diff --git a/.bazelrc b/.bazelrc new file mode 100644 index 0000000000000000000000000000000000000000..c70c57136102b483a4332ca22f775d7a2c5b849e --- /dev/null +++ b/.bazelrc @@ -0,0 +1,104 @@ +# Android configs. Bazel needs to have --cpu and --fat_apk_cpu both set to the +# target CPU to build transient dependencies correctly. See +# https://docs.bazel.build/versions/master/user-manual.html#flag--fat_apk_cpu +build:android --crosstool_top=//external:android/crosstool +build:android --host_crosstool_top=@bazel_tools//tools/cpp:toolchain +build:android_arm --config=android +build:android_arm --cpu=armeabi-v7a +build:android_arm --fat_apk_cpu=armeabi-v7a +build:android_arm64 --config=android +build:android_arm64 --cpu=arm64-v8a +build:android_arm64 --fat_apk_cpu=arm64-v8a + +# Config to use a mostly-static build and disable modular op registration +# support (this will revert to loading TensorFlow with RTLD_GLOBAL in Python). +# By default, TensorFlow will build with a dependence on +# //tensorflow:libtensorflow_framework.so. +build:monolithic --define framework_shared_object=false + +# For projects which use TensorFlow as part of a Bazel build process, putting +# nothing in a bazelrc will default to a monolithic build. The following line +# opts in to modular op registration support by default. +build --define framework_shared_object=true + +# Please note that MKL on MacOS or windows is still not supported. +# 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 +# Instruct clang to use LLD for linking. +# This only works with GPU builds currently, since Bazel sets -B/usr/bin in +# auto-generated CPU crosstool, forcing /usr/bin/ld.lld to be preferred over +# the downloaded one. +build:download_clang_use_lld --linkopt='-fuse-ld=lld' + +build:cuda --crosstool_top=@local_config_cuda//crosstool:toolchain +build:cuda --define=using_cuda=true --define=using_cuda_nvcc=true + +build:rocm --crosstool_top=@local_config_rocm//crosstool:toolchain +build:rocm --define=using_rocm=true --define=using_rocm_hipcc=true + +build:cuda_clang --crosstool_top=@local_config_cuda//crosstool:toolchain +build:cuda_clang --define=using_cuda=true --define=using_cuda_clang=true --define=using_clang=true + +build:sycl --crosstool_top=@local_config_sycl//crosstool:toolchain +build:sycl --define=using_sycl=true --define=using_trisycl=false + +build:sycl_nodouble --crosstool_top=@local_config_sycl//crosstool:toolchain +build:sycl_nodouble --define=using_sycl=true --cxxopt -DTENSORFLOW_SYCL_NO_DOUBLE + +build:sycl_asan --crosstool_top=@local_config_sycl//crosstool:toolchain +build:sycl_asan --define=using_sycl=true --define=using_trisycl=false --copt -fno-omit-frame-pointer --copt -fsanitize-coverage=3 --copt -DGPR_NO_DIRECT_SYSCALLS --linkopt -fPIC --linkopt -fsanitize=address + +build:sycl_trisycl --crosstool_top=@local_config_sycl//crosstool:toolchain +build:sycl_trisycl --define=using_sycl=true --define=using_trisycl=true + +# Options extracted from configure script +build:gdr --define=with_gdr_support=true +build:ngraph --define=with_ngraph_support=true +build:verbs --define=with_verbs_support=true + +# Options to disable default on features +build:noaws --define=no_aws_support=true +build:nogcp --define=no_gcp_support=true +build:nohdfs --define=no_hdfs_support=true +build:nokafka --define=no_kafka_support=true +build:noignite --define=no_ignite_support=true +build:nonccl --define=no_nccl_support=true + +build --define=use_fast_cpp_protos=true +build --define=allow_oversize_protos=true + +build --spawn_strategy=standalone +build --strategy=Genrule=standalone +build -c opt + +# Other build flags. +build --define=grpc_no_ares=true + +# Modular TF build options +build:dynamic_kernels --define=dynamic_loaded_kernels=true +build:dynamic_kernels --copt=-DAUTOLOAD_DYNAMIC_KERNELS + +# Default paths for TF_SYSTEM_LIBS +build --define=PREFIX=/usr +build --define=LIBDIR=$(PREFIX)/lib +build --define=INCLUDEDIR=$(PREFIX)/include + +# Default options should come above this line + +# Options from ./configure +try-import %workspace%/.tf_configure.bazelrc + +# Put user-specific options in .bazelrc.user +try-import %workspace%/.bazelrc.user diff --git a/.github/ISSUE_TEMPLATE/bug-performance-issue.md b/.github/ISSUE_TEMPLATE/00-bug-performance-issue.md similarity index 100% rename from .github/ISSUE_TEMPLATE/bug-performance-issue.md rename to .github/ISSUE_TEMPLATE/00-bug-performance-issue.md diff --git a/.github/ISSUE_TEMPLATE/build-installation-issue.md b/.github/ISSUE_TEMPLATE/10-build-installation-issue.md similarity index 100% rename from .github/ISSUE_TEMPLATE/build-installation-issue.md rename to .github/ISSUE_TEMPLATE/10-build-installation-issue.md diff --git a/.github/ISSUE_TEMPLATE/documentation-issue.md b/.github/ISSUE_TEMPLATE/20-documentation-issue.md similarity index 100% rename from .github/ISSUE_TEMPLATE/documentation-issue.md rename to .github/ISSUE_TEMPLATE/20-documentation-issue.md diff --git a/.github/ISSUE_TEMPLATE/feature-request.md b/.github/ISSUE_TEMPLATE/30-feature-request.md similarity index 100% rename from .github/ISSUE_TEMPLATE/feature-request.md rename to .github/ISSUE_TEMPLATE/30-feature-request.md diff --git a/.github/ISSUE_TEMPLATE/40-tflite-op-request.md b/.github/ISSUE_TEMPLATE/40-tflite-op-request.md new file mode 100644 index 0000000000000000000000000000000000000000..7b391279e479ade4ed5327728f19be8752e11507 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/40-tflite-op-request.md @@ -0,0 +1,24 @@ +--- +name: TensorFlow Lite Op Request +about: Use this template for reporting ops you are using or missing. + +--- + + +**System information** +- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): +- TensorFlow installed from (source or binary): +- TensorFlow version (or github SHA if from source): + + +**Provide the text output from tflite_convert** + +``` +# Copy and paste here +``` + +Also, please include a link to a GraphDef or the model if possible. + +**Any other info / logs** + +Include any logs or source code that would be helpful to diagnose the problem. If including tracebacks, please include the full traceback. Large logs and files should be attached. diff --git a/.github/ISSUE_TEMPLATE/other-issues.md b/.github/ISSUE_TEMPLATE/50-other-issues.md similarity index 100% rename from .github/ISSUE_TEMPLATE/other-issues.md rename to .github/ISSUE_TEMPLATE/50-other-issues.md diff --git a/.gitignore b/.gitignore index 1ef4c297ee4f369775c13b32a46a55887de719e7..e1d352c238a1b2d4febe0f5d4a30cfa0c942f7e7 100644 --- a/.gitignore +++ b/.gitignore @@ -1,7 +1,7 @@ .DS_Store .ipynb_checkpoints node_modules -/.bazelrc +/.bazelrc.user /.tf_configure.bazelrc /bazel-* /bazel_pip @@ -24,10 +24,10 @@ Pods Podfile.lock *.pbxproj *.xcworkspacedata -/tensorflow/contrib/lite/downloads/** -/tensorflow/contrib/lite/gen/** -/tensorflow/contrib/lite/examples/ios/simple/data/*.txt -/tensorflow/contrib/lite/examples/ios/simple/data/*.tflite +/tensorflow/lite/tools/make/downloads/** +/tensorflow/lite/gen/** +/tensorflow/lite/examples/ios/simple/data/*.txt +/tensorflow/lite/examples/ios/simple/data/*.tflite xcuserdata/** /api_init_files_list.txt /estimator_api_init_files_list.txt diff --git a/BUILD b/BUILD index 4bf647e47aa56cff0b3fd5af7d5df99d8b70549b..1200cf5f7103cad12ab9693c339c372f4f3bc0fb 100644 --- a/BUILD +++ b/BUILD @@ -2,5 +2,7 @@ exports_files( [ "LICENSE", "ACKNOWLEDGEMENTS", + "configure", + "configure.py", ], ) diff --git a/CODEOWNERS b/CODEOWNERS index 94cc865479cd6ab5cdb589490d3a2d650f06b160..cb3fa2312405ce44d5dfc30ea4164740f436e07e 100644 --- a/CODEOWNERS +++ b/CODEOWNERS @@ -1,6 +1,7 @@ # Where component owners are known, add them here. /tenosrflow/core/debug @caisq +/tensorflow/core/nccl/ @azaks2 @chsigg /tensorflow/core/platform/windows/ @mrry /tensorflow/core/platform/s3 @yongtang /tensorflow/go @asimshankar @@ -46,18 +47,17 @@ /tensorflow/contrib/losses/ @alextp @ispirmustafa /tensorflow/contrib/makefile/ @petewarden @satok16 @wolffg /tensorflow/contrib/metrics/ @alextp @honkentuber @ispirmustafa -/tensorflow/contrib/nccl/ @cwhipkey @zheng-xq /tensorflow/contrib/opt/ @strategist333 @alextp /tensorflow/contrib/pi_examples/ @maciekcc /tensorflow/contrib/quantization/ @petewarden /tensorflow/contrib/rnn/ @ebrevdo @scottzhu -/tensorflow/contrib/saved_model/ @nfiedel @sukritiramesh @allenl +/tensorflow/contrib/saved_model/ @nfiedel @sukritiramesh @allenlavoie /tensorflow/contrib/seq2seq/ @ebrevdo @lmthang /tensorflow/contrib/session_bundle/ @nfiedel @sukritiramesh /tensorflow/contrib/slim/ @sguada @thenbasilmanran /tensorflow/contrib/stateless/ @girving @alextp /tensorflow/contrib/tensor_forest/ @gilberthendry @thomascolthurst @yupbank -/tensorflow/contrib/tensorrt/ @aaroey +/tensorflow/contrib/tensorrt/ @aaroey @smit-hinsu @azaks2 # NEED OWNER: /tensorflow/contrib/testing/ /tensorflow/contrib/timeseries/ @allenlavoie /tensorflow/contrib/tpu/ @frankchn @saeta @jhseu @sourabhbajaj diff --git a/CODE_OF_CONDUCT.md b/CODE_OF_CONDUCT.md index 20601eaf611d98f78382a7d260629e72e24a07c0..a4647020ff76830badd75f3d3f76a41a637159bb 100644 --- a/CODE_OF_CONDUCT.md +++ b/CODE_OF_CONDUCT.md @@ -51,10 +51,12 @@ However, for the vast majority of issues, we aim to empower individuals to first If you are experiencing or witnessing conflict, we ask you to use the following escalation strategy to address the conflict: -1. Address the perceived conflict directly with those involved, preferably in a real-time medium. -2. If this fails, get a third party (e.g. a mutual friend, and/or someone with background on the issue, but not involved in conflict) to intercede. -3. If you are still unable to resolve the conflict, and you believe it rises to harassment or another code of conduct violation, report it. - +1. Address the perceived conflict directly with those involved, preferably in a + real-time medium. +2. If this fails, get a third party (e.g. a mutual friend, and/or someone with + background on the issue, but not involved in the conflict) to intercede. +3. If you are still unable to resolve the conflict, and you believe it rises to + harassment or another code of conduct violation, report it. ## Reporting Violations 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 8af5370befbb090966a8b3af54d80c84a969aaa5..4e37b239b16e6eeefc587aeb242a03e1f88eddbd 100644 --- a/README.md +++ b/README.md @@ -9,12 +9,14 @@ |-----------------| | [![Documentation](https://img.shields.io/badge/api-reference-blue.svg)](https://www.tensorflow.org/api_docs/) | -**TensorFlow** is an open source software library for numerical computation using -data flow graphs. The graph nodes represent mathematical operations, while +**TensorFlow** is an open source software library for numerical computation +using data flow graphs. The graph nodes represent mathematical operations, while the graph edges represent the multidimensional data arrays (tensors) that flow -between them. This flexible architecture enables you to deploy computation to one -or more CPUs or GPUs in a desktop, server, or mobile device without rewriting -code. TensorFlow also includes [TensorBoard](https://www.tensorflow.org/guide/summaries_and_tensorboard), a data visualization toolkit. +between them. This flexible architecture enables you to deploy computation to +one or more CPUs or GPUs in a desktop, server, or mobile device without +rewriting code. TensorFlow also includes +[TensorBoard](https://github.com/tensorflow/tensorboard), a data visualization +toolkit. TensorFlow was originally developed by researchers and engineers working on the Google Brain team within Google's Machine Intelligence Research @@ -55,21 +57,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 @@ -111,22 +116,25 @@ The TensorFlow project strives to abide by generally accepted best practices in Build Type | Status | Artifacts ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | --------- **IBM s390x** | [![Build Status](http://ibmz-ci.osuosl.org/job/TensorFlow_IBMZ_CI/badge/icon)](http://ibmz-ci.osuosl.org/job/TensorFlow_IBMZ_CI/) | TBA -**IBM ppc64le CPU** | [![Build Status](http://powerci.osuosl.org/job/TensorFlow_Ubuntu_16.04_CPU/badge/icon)](http://powerci.osuosl.org/job/TensorFlow_Ubuntu_16.04_CPU/) | TBA -**IBM ppc64le GPU** Nightly | [![Build Status](https://powerci.osuosl.org/job/TensorFlow_PPC64LE_GPU_Nightly_Artifact/badge/icon)](https://powerci.osuosl.org/job/TensorFlow_PPC64LE_GPU_Nightly_Artifact/) | [Nightly](https://powerci.osuosl.org/job/TensorFlow_PPC64LE_GPU_Nightly_Artifact/) -**IBM ppc64le GPU** Stable Release | [![Build Status](https://powerci.osuosl.org/job/TensorFlow_PPC64LE_GPU_Release_Build/badge/icon)](https://powerci.osuosl.org/job/TensorFlow_PPC64LE_GPU_Release_Build/) | [Release](https://powerci.osuosl.org/job/TensorFlow_PPC64LE_GPU_Release_Build/) +**Linux ppc64le CPU** Nightly | [![Build Status](https://powerci.osuosl.org/job/TensorFlow_PPC64LE_CPU_Build/badge/icon)](https://powerci.osuosl.org/job/TensorFlow_PPC64LE_CPU_Build/) | [Nightly](https://powerci.osuosl.org/job/TensorFlow_PPC64LE_CPU_Nightly_Artifact/) +**Linux ppc64le CPU** Stable Release | [![Build Status](https://powerci.osuosl.org/job/TensorFlow_PPC64LE_CPU_Release_Build/badge/icon)](https://powerci.osuosl.org/job/TensorFlow_PPC64LE_CPU_Release_Build/) | [Release](https://powerci.osuosl.org/job/TensorFlow_PPC64LE_CPU_Release_Build/) +**Linux ppc64le GPU** Nightly | [![Build Status](https://powerci.osuosl.org/job/TensorFlow_PPC64LE_GPU_Build/badge/icon)](https://powerci.osuosl.org/job/TensorFlow_PPC64LE_GPU_Build/) | [Nightly](https://powerci.osuosl.org/job/TensorFlow_PPC64LE_GPU_Nightly_Artifact/) +**Linux ppc64le GPU** Stable Release | [![Build Status](https://powerci.osuosl.org/job/TensorFlow_PPC64LE_GPU_Release_Build/badge/icon)](https://powerci.osuosl.org/job/TensorFlow_PPC64LE_GPU_Release_Build/) | [Release](https://powerci.osuosl.org/job/TensorFlow_PPC64LE_GPU_Release_Build/) **Linux CPU with Intel® MKL-DNN** Nightly | [![Build Status](https://tensorflow-ci.intel.com/job/tensorflow-mkl-linux-cpu/badge/icon)](https://tensorflow-ci.intel.com/job/tensorflow-mkl-linux-cpu/) | [Nightly](https://tensorflow-ci.intel.com/job/tensorflow-mkl-build-whl-nightly/) -**Linux CPU with Intel® MKL-DNN** Python 2.7
**Linux CPU with Intel® MKL-DNN** Python 3.4
**Linux CPU with Intel® MKL-DNN** Python 3.5
**Linux CPU with Intel® MKL-DNN** Python 3.6 | [![Build Status](https://tensorflow-ci.intel.com/job/tensorflow-mkl-build-release-whl/badge/icon)](https://tensorflow-ci.intel.com/job/tensorflow-mkl-build-release-whl/lastStableBuild) | [1.11.0 py2.7](https://storage.googleapis.com/intel-optimized-tensorflow/tensorflow-1.11.0-cp27-cp27mu-linux_x86_64.whl)
[1.11.0 py3.4](https://storage.googleapis.com/intel-optimized-tensorflow/tensorflow-1.11.0-cp34-cp34m-linux_x86_64.whl)
[1.11.0 py3.5](https://storage.googleapis.com/intel-optimized-tensorflow/tensorflow-1.11.0-cp35-cp35m-linux_x86_64.whl)
[1.11.0 py3.6](https://storage.googleapis.com/intel-optimized-tensorflow/tensorflow-1.11.0-cp36-cp36m-linux_x86_64.whl) +**Linux CPU with Intel® MKL-DNN** Python 2.7
**Linux CPU with Intel® MKL-DNN** Python 3.4
**Linux CPU with Intel® MKL-DNN** Python 3.5
**Linux CPU with Intel® MKL-DNN** Python 3.6 | [![Build Status](https://tensorflow-ci.intel.com/job/tensorflow-mkl-build-release-whl/badge/icon)](https://tensorflow-ci.intel.com/job/tensorflow-mkl-build-release-whl/lastStableBuild) | [1.12.0 py2.7](https://storage.googleapis.com/intel-optimized-tensorflow/tensorflow-1.12.0-cp27-cp27mu-linux_x86_64.whl)
[1.12.0 py3.4](https://storage.googleapis.com/intel-optimized-tensorflow/tensorflow-1.12.0-cp34-cp34m-linux_x86_64.whl)
[1.12.0 py3.5](https://storage.googleapis.com/intel-optimized-tensorflow/tensorflow-1.12.0-cp35-cp35m-linux_x86_64.whl)
[1.12.0 py3.6](https://storage.googleapis.com/intel-optimized-tensorflow/tensorflow-1.12.0-cp36-cp36m-linux_x86_64.whl) ## For more information -* [TensorFlow Website](https://www.tensorflow.org) -* [TensorFlow Tutorials](https://www.tensorflow.org/tutorials/) -* [TensorFlow Model Zoo](https://github.com/tensorflow/models) -* [TensorFlow Twitter](https://twitter.com/tensorflow) -* [TensorFlow Blog](https://medium.com/tensorflow) -* [TensorFlow Course at Stanford](https://web.stanford.edu/class/cs20si) -* [TensorFlow Roadmap](https://www.tensorflow.org/community/roadmap) -* [TensorFlow White Papers](https://www.tensorflow.org/about/bib) -* [TensorFlow YouTube Channel](https://www.youtube.com/channel/UC0rqucBdTuFTjJiefW5t-IQ) + +* [TensorFlow Website](https://www.tensorflow.org) +* [TensorFlow Tutorials](https://www.tensorflow.org/tutorials/) +* [TensorFlow Model Zoo](https://github.com/tensorflow/models) +* [TensorFlow Twitter](https://twitter.com/tensorflow) +* [TensorFlow Blog](https://medium.com/tensorflow) +* [TensorFlow Course at Stanford](https://web.stanford.edu/class/cs20si) +* [TensorFlow Roadmap](https://www.tensorflow.org/community/roadmap) +* [TensorFlow White Papers](https://www.tensorflow.org/about/bib) +* [TensorFlow YouTube Channel](https://www.youtube.com/channel/UC0rqucBdTuFTjJiefW5t-IQ) +* [TensorFlow Visualization Toolkit](https://github.com/tensorflow/tensorboard) Learn more about the TensorFlow community at the [community page of tensorflow.org](https://www.tensorflow.org/community) for a few ways to participate. diff --git a/RELEASE.md b/RELEASE.md index 2b00d06580d925a4afed5753afb8f51f0ebac99f..0a56e6909870e398c9d6349576cd2f8e6734f072 100644 --- a/RELEASE.md +++ b/RELEASE.md @@ -7,6 +7,8 @@ Serving. * Keras models now support evaluating with a `tf.data.Dataset`. * TensorFlow binaries are built with XLA support linked in by default. +* Ignite Dataset added to contrib/ignite that allows to work with Apache + Ignite. ## Bug Fixes and Other Changes @@ -258,8 +260,8 @@ Ag Ramesh, Alex Wiltschko, Alexander Pantyukhin, Amogh Mannekote, An Jiaoyang, A * Update `tf.keras` to the Keras 2.1.6 API. * Added [`tf.keras.layers.CuDNNGRU`](https://www.tensorflow.org/versions/r1.9/api_docs/python/tf/keras/layers/CuDNNGRU) and [`tf.keras.layers.CuDNNLSTM`](https://www.tensorflow.org/versions/r1.9/api_docs/python/tf/keras/layers/CuDNNLSTM) layers. [Try it](https://colab.sandbox.google.com/github/tensorflow/tensorflow/blob/master/tensorflow/contrib/eager/python/examples/nmt_with_attention/nmt_with_attention.ipynb?linkId=53292082). * Adding support of core [feature columns](https://www.tensorflow.org/get_started/feature_columns) and [losses](https://www.tensorflow.org/api_docs/python/tf/losses) to [gradient boosted trees estimators](https://github.com/tensorflow/models/tree/master/official/boosted_trees). -* The [python interface](https://www.tensorflow.org/versions/r1.9/api_docs/python/tf/contrib/lite) - for the [TFLite Optimizing Converter](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/lite/toco/README.md) +* The [python interface](https://www.tensorflow.org/versions/r1.9/api_docs/python/tf/lite) + for the [TFLite Optimizing Converter](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/toco/README.md) has been expanded, and the command line interface (AKA: `toco`, `tflite_convert`) is once again included in the standard `pip` installation. * Improved data-loading and text processing with: @@ -280,50 +282,76 @@ Ag Ramesh, Alex Wiltschko, Alexander Pantyukhin, Amogh Mannekote, An Jiaoyang, A ## Bug Fixes and Other Changes -* `tfe.Network` is deprecated. Please inherit from `tf.keras.Model`. -* Layered variable names have changed in the following conditions: - * Using `tf.keras.layers` with custom variable scopes. - * Using `tf.layers` in a subclassed `tf.keras.Model` class. See - [here](https://www.tensorflow.org/versions/r1.9/api_docs/python/tf/layers) for more details -* `tf.data`: - * `Dataset.from_generator()` now accepts an `args` list, in order to create nested generators. - * `Dataset.list_files()` now produces determinstic results when `shuffle=False` or a `seed` is passed. - * `tf.contrib.data.sample_from_datasets()` and `tf.contrib.data.choose_from_datasets()` make it easier to sample or deterministically choose elements from multiple datasets. - * `tf.contrib.data.make_csv_dataset()` now supports line breaks in quoted strings, and two infrequently used arguments removed. - * (C++) `DatasetBase::DebugString()` is now `const`. - * (C++) `DatasetBase::MakeIterator()` has been renamed to `DatasetBase::MakeIteratorInternal()`. - * (C++) `IteratorBase::Initialize()` method was added to support raising errors during iterator construction. -* Eager Execution: - * Added the ability to pause recording operations for gradient computation via `tf.GradientTape.stop_recording`. - * Updated documentation, introductory notebooks. -* `tf.keras`: - * Move Keras code out of _impl folder and remove API files. - * `tf.keras.Model.save_weights` now saves in TensorFlow format by default. - * Enable dataset iterators to be passed to `tf.keras.Model` training/eval methods. -* TensorFlow Debugger (tfdbg) CLI: fix an issue in which the TensorBoard Debugger Plugin could not handle total source file size exceeding gRPC message size limit (4 MB). -* `tf.contrib`: - * `tf.contrib.framework.zero_initializer` supports ResourceVariable. - * Adding "constrained_optimization" to tensorflow/contrib. -* Other: - * Add GCS Configuration Ops. - * Changing signature of `MakeIterator` to enable propagating error status. - * KL divergence for two Dirichlet distributions. - * More consistent GcsFileSystem behavior for certain reads past EOF. - * Update benchmark for tf.scan to match ranges across eager and graph modes. - * Fixed bug in `tf.reduce_prod gradient` for complex dtypes. - * Allow the use of '.' in variables (e.g. "hparams.parse('a.b=1.0')"), which would previously raise an error. This will correspond to an attribute name with an embedded '.' symbol (e.g. 'a.b'), which can only be accessed indirectly (e.g. through getattr and setattr). To set this up the user will first need to explicitly add the variable to the hparam object (e.g. "hparams.add_hparam(name='a.b', value=0.0)"). - * Benchmark for tf.scan in graph and eager modes. - * Added complex128 support to FFT, FFT2D, FFT3D, IFFT, IFFT2D, and IFFT3D. - * Making ids unique in `nn.embedding_lookup_sparse`. This helps to reduce RPC calls for looking up the embeddings when there are repeated ids in the batch. - * Support indicator column in boosted trees. - * Prevent `tf.gradients()` from backpropagating through integer tensors. - * LinearOperator[1D,2D,3D]Circulant added to `tensorflow.linalg`. - * Conv3D, Conv3DBackpropInput, Conv3DBackpropFilter now supports arbitrary. - * Added `tf.train.Checkpoint` for reading/writing object-based checkpoints. - * Added LinearOperatorKronecker, a dense-free implementation of the Kronecker Product. - * Allow LinearOperator to broadcast. - * SavedModelBuilder will now deduplicate asset names that point to files with the same basename and the same contents. Note that this may result in new asset files included in SavedModels in cases where assets with the same name but different contents were previously overwriting each other. - +* `tfe.Network` is deprecated. Please inherit from `tf.keras.Model`. +* Layered variable names have changed in the following conditions: + * Using `tf.keras.layers` with custom variable scopes. + * Using `tf.layers` in a subclassed `tf.keras.Model` class. See + [here](https://www.tensorflow.org/versions/r1.9/api_docs/python/tf/layers) + for more details +* `tf.data`: + * `Dataset.from_generator()` now accepts an `args` list, in order to + create nested generators. + * `Dataset.list_files()` now produces deterministic results when + `shuffle=False` or a `seed` is passed. + * `tf.contrib.data.sample_from_datasets()` and + `tf.contrib.data.choose_from_datasets()` make it easier to sample or + deterministically choose elements from multiple datasets. + * `tf.contrib.data.make_csv_dataset()` now supports line breaks in quoted + strings, and two infrequently used arguments removed. + * (C++) `DatasetBase::DebugString()` is now `const`. + * (C++) `DatasetBase::MakeIterator()` has been renamed to + `DatasetBase::MakeIteratorInternal()`. + * (C++) `IteratorBase::Initialize()` method was added to support raising + errors during iterator construction. +* Eager Execution: + * Added the ability to pause recording operations for gradient computation + via `tf.GradientTape.stop_recording`. + * Updated documentation, introductory notebooks. +* `tf.keras`: + * Move Keras code out of _impl folder and remove API files. + * `tf.keras.Model.save_weights` now saves in TensorFlow format by default. + * Enable dataset iterators to be passed to `tf.keras.Model` training/eval + methods. +* TensorFlow Debugger (tfdbg) CLI: fix an issue in which the TensorBoard + Debugger Plugin could not handle total source file size exceeding gRPC + message size limit (4 MB). +* `tf.contrib`: + * `tf.contrib.framework.zero_initializer` supports ResourceVariable. + * Adding "constrained_optimization" to tensorflow/contrib. +* Other: + * Add GCS Configuration Ops. + * Changing signature of `MakeIterator` to enable propagating error status. + * KL divergence for two Dirichlet distributions. + * More consistent GcsFileSystem behavior for certain reads past EOF. + * Update benchmark for tf.scan to match ranges across eager and graph + modes. + * Fixed bug in `tf.reduce_prod gradient` for complex dtypes. + * Allow the use of '.' in variables (e.g. "hparams.parse('a.b=1.0')"), + which would previously raise an error. This will correspond to an + attribute name with an embedded '.' symbol (e.g. 'a.b'), which can only + be accessed indirectly (e.g. through getattr and setattr). To set this + up the user will first need to explicitly add the variable to the hparam + object (e.g. "hparams.add_hparam(name='a.b', value=0.0)"). + * Benchmark for tf.scan in graph and eager modes. + * Added complex128 support to FFT, FFT2D, FFT3D, IFFT, IFFT2D, and IFFT3D. + * Making ids unique in `nn.embedding_lookup_sparse`. This helps to reduce + RPC calls for looking up the embeddings when there are repeated ids in + the batch. + * Support indicator column in boosted trees. + * Prevent `tf.gradients()` from backpropagating through integer tensors. + * LinearOperator[1D,2D,3D]Circulant added to `tensorflow.linalg`. + * Conv3D, Conv3DBackpropInput, Conv3DBackpropFilter now supports + arbitrary. + * Added `tf.train.Checkpoint` for reading/writing object-based + checkpoints. + * Added LinearOperatorKronecker, a dense-free implementation of the + Kronecker Product. + * Allow LinearOperator to broadcast. + * SavedModelBuilder will now deduplicate asset names that point to files + with the same basename and the same contents. Note that this may result + in new asset files included in SavedModels in cases where assets with + the same name but different contents were previously overwriting each + other. ## Thanks to our Contributors @@ -562,7 +590,7 @@ Yoni Tsafir, yordun, Yuan (Terry) Tang, Yuxin Wu, zhengdi, Zhengsheng Wei, 田 ## Major Features And Improvements * [Eager execution](https://github.com/tensorflow/tensorflow/tree/r1.5/tensorflow/contrib/eager) preview version is now available. -* [TensorFlow Lite](https://github.com/tensorflow/tensorflow/tree/r1.5/tensorflow/contrib/lite) +* [TensorFlow Lite](https://github.com/tensorflow/tensorflow/tree/r1.5/tensorflow/lite) dev preview is now available. * CUDA 9.0 and cuDNN 7 support. * Accelerated Linear Algebra (XLA): @@ -821,7 +849,7 @@ answered questions, and were part of inspiring discussions. * Remove `tf.contrib.data.Iterator.from_dataset()` method. Use `Dataset.make_initializable_iterator()` instead. * Remove seldom used and unnecessary `tf.contrib.data.Iterator.dispose_op()`. -* Reorder some TFGAN loss functions in a non-backwards compatible way. +* Reorder some TF-GAN loss functions in a non-backwards compatible way. ## Known Issues * In Python 3, `Dataset.from_generator()` does not support Unicode strings. @@ -909,7 +937,7 @@ See also [TensorBoard 0.1.4](https://github.com/tensorflow/tensorboard/releases/ * Adds tf.contrib.nn.rank_sampled_softmax_loss, a sampled-softmax variant that can improve rank loss. * `tf.contrib.metrics`.{streaming_covariance,streaming_pearson_correlation} modified to return nan when they have seen less or equal to 1 unit of weight. * Adds time series models to contrib. See contrib/timeseries/README.md for details. -* Adds FULLY_CONNECTED Op to tensorflow/contrib/lite/schema.fbs +* Adds FULLY_CONNECTED Op to tensorflow/lite/schema.fbs ## Known Issues * Tensorflow_gpu compilation fails with Bazel 0.5.3. diff --git a/WORKSPACE b/WORKSPACE index 17961829a605c2d1f2d2ba86a7c30c47618c139b..957b8d8528dc9b5e2ea134921b28601aa6fed2d1 100644 --- a/WORKSPACE +++ b/WORKSPACE @@ -1,12 +1,14 @@ workspace(name = "org_tensorflow") +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 ], ) @@ -14,11 +16,64 @@ load("@io_bazel_rules_closure//closure:defs.bzl", "closure_repositories") closure_repositories() +load("//third_party/toolchains/preconfig/generate:archives.bzl", + "bazel_toolchains_archive") + +bazel_toolchains_archive() + +load( + "@bazel_toolchains//repositories:repositories.bzl", + bazel_toolchains_repositories = "repositories", +) + +bazel_toolchains_repositories() + +load( + "@io_bazel_rules_docker//container:container.bzl", + container_repositories = "repositories", +) + +container_repositories() + +load("//third_party/toolchains/preconfig/generate:workspace.bzl", + "remote_config_workspace") + +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"], +) +http_file( + name = "xctestrunner", + executable = 1, + urls = ["https://github.com/google/xctestrunner/releases/download/0.2.5/ios_test_runner.par"], +) +load("@build_bazel_rules_apple//apple:repositories.bzl", "apple_rules_dependencies") +apple_rules_dependencies(ignore_version_differences = True) +load("@build_bazel_rules_swift//swift:repositories.bzl", "swift_rules_dependencies") +swift_rules_dependencies() + # We must check the bazel version before trying to parse any other BUILD # 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.15.0") +check_bazel_version_at_least("0.19.0") load("//tensorflow:workspace.bzl", "tf_workspace") @@ -30,9 +85,9 @@ android_workspace() # Please add all new TensorFlow dependencies in workspace.bzl. tf_workspace() -new_http_archive( +http_archive( name = "inception_v1", - build_file = "models.BUILD", + build_file = "//:models.BUILD", sha256 = "7efe12a8363f09bc24d7b7a450304a15655a57a7751929b2c1593a71183bb105", urls = [ "http://storage.googleapis.com/download.tensorflow.org/models/inception_v1.zip", @@ -40,9 +95,9 @@ new_http_archive( ], ) -new_http_archive( +http_archive( name = "mobile_ssd", - build_file = "models.BUILD", + build_file = "//:models.BUILD", sha256 = "bddd81ea5c80a97adfac1c9f770e6f55cbafd7cce4d3bbe15fbeb041e6b8f3e8", urls = [ "http://storage.googleapis.com/download.tensorflow.org/models/object_detection/ssd_mobilenet_v1_android_export.zip", @@ -50,9 +105,9 @@ new_http_archive( ], ) -new_http_archive( +http_archive( name = "mobile_multibox", - build_file = "models.BUILD", + build_file = "//:models.BUILD", sha256 = "859edcddf84dddb974c36c36cfc1f74555148e9c9213dedacf1d6b613ad52b96", urls = [ "http://storage.googleapis.com/download.tensorflow.org/models/mobile_multibox_v1a.zip", @@ -60,9 +115,9 @@ new_http_archive( ], ) -new_http_archive( +http_archive( name = "stylize", - build_file = "models.BUILD", + build_file = "//:models.BUILD", sha256 = "3d374a730aef330424a356a8d4f04d8a54277c425e274ecb7d9c83aa912c6bfa", urls = [ "http://storage.googleapis.com/download.tensorflow.org/models/stylize_v1.zip", @@ -70,12 +125,13 @@ new_http_archive( ], ) -new_http_archive( +http_archive( name = "speech_commands", - build_file = "models.BUILD", + build_file = "//:models.BUILD", sha256 = "c3ec4fea3158eb111f1d932336351edfe8bd515bb6e87aad4f25dbad0a600d0c", urls = [ "http://storage.googleapis.com/download.tensorflow.org/models/speech_commands_v0.01.zip", "http://download.tensorflow.org/models/speech_commands_v0.01.zip", ], ) + diff --git a/configure.py b/configure.py index ceaae6399343b6d3690b37ae9624776b7276c93e..adc9ef9caca8c0128c63896fdebbbadf7f86da81 100644 --- a/configure.py +++ b/configure.py @@ -33,7 +33,7 @@ except ImportError: from distutils.spawn import find_executable as which # pylint: enable=g-import-not-at-top -_DEFAULT_CUDA_VERSION = '9.0' +_DEFAULT_CUDA_VERSION = '10.0' _DEFAULT_CUDNN_VERSION = '7' _DEFAULT_CUDA_COMPUTE_CAPABILITIES = '3.5,7.0' _DEFAULT_CUDA_PATH = '/usr/local/cuda' @@ -43,7 +43,7 @@ _DEFAULT_CUDA_PATH_WIN = ('C:/Program Files/NVIDIA GPU Computing ' _TF_OPENCL_VERSION = '1.2' _DEFAULT_COMPUTECPP_TOOLKIT_PATH = '/usr/local/computecpp' _DEFAULT_TRISYCL_INCLUDE_DIR = '/usr/local/triSYCL/include' -_SUPPORTED_ANDROID_NDK_VERSIONS = [10, 11, 12, 13, 14, 15, 16] +_SUPPORTED_ANDROID_NDK_VERSIONS = [10, 11, 12, 13, 14, 15, 16, 17, 18] _DEFAULT_PROMPT_ASK_ATTEMPTS = 10 @@ -238,6 +238,13 @@ def setup_python(environ_cp): write_to_bazelrc('build --python_path=\"%s"' % python_bin_path) environ_cp['PYTHON_BIN_PATH'] = python_bin_path + # If choosen python_lib_path is from a path specified in the PYTHONPATH + # variable, need to tell bazel to include PYTHONPATH + if environ_cp.get('PYTHONPATH'): + python_paths = environ_cp.get('PYTHONPATH').split(':') + if python_lib_path in python_paths: + write_action_env_to_bazelrc('PYTHONPATH', environ_cp.get('PYTHONPATH')) + # Write tools/python_bin_path.sh with open( os.path.join(_TF_WORKSPACE_ROOT, 'tools', 'python_bin_path.sh'), @@ -248,18 +255,6 @@ def setup_python(environ_cp): def reset_tf_configure_bazelrc(): """Reset file that contains customized config settings.""" open(_TF_BAZELRC, 'w').close() - bazelrc_path = os.path.join(_TF_WORKSPACE_ROOT, '.bazelrc') - - data = [] - if os.path.exists(bazelrc_path): - with open(bazelrc_path, 'r') as f: - data = f.read().splitlines() - with open(bazelrc_path, 'w') as f: - for l in data: - if _TF_BAZELRC_FILENAME in l: - continue - f.write('%s\n' % l) - f.write('import %%workspace%%/%s\n' % _TF_BAZELRC_FILENAME) def cleanup_makefile(): """Delete any leftover BUILD files from the Makefile build. @@ -445,11 +440,12 @@ def convert_version_to_int(version): return int(version_str) -def check_bazel_version(min_version): - """Check installed bazel version is at least min_version. +def check_bazel_version(min_version, max_version): + """Check installed bazel version is between min_version and max_version. Args: min_version: string for minimum bazel version. + max_version: string for maximum bazel version. Returns: The bazel version detected. @@ -467,6 +463,7 @@ def check_bazel_version(min_version): min_version_int = convert_version_to_int(min_version) curr_version_int = convert_version_to_int(curr_version) + max_version_int = convert_version_to_int(max_version) # Check if current bazel version can be detected properly. if not curr_version_int: @@ -479,7 +476,14 @@ def check_bazel_version(min_version): if curr_version_int < min_version_int: print('Please upgrade your bazel installation to version %s or higher to ' 'build TensorFlow!' % min_version) - sys.exit(0) + sys.exit(1) + if (curr_version_int > max_version_int and + 'TF_IGNORE_MAX_BAZEL_VERSION' not in os.environ): + print('Please downgrade your bazel installation to version %s or lower to ' + 'build TensorFlow! To downgrade: download the installer for the old ' + 'version (from https://github.com/bazelbuild/bazel/releases) then ' + 'run the installer.' % max_version) + sys.exit(1) return curr_version @@ -859,7 +863,7 @@ def set_tf_cuda_version(environ_cp): cuda_toolkit_paths_full = [ os.path.join(cuda_toolkit_path, x) for x in cuda_rt_lib_paths ] - if any([os.path.exists(x) for x in cuda_toolkit_paths_full]): + if any(os.path.exists(x) for x in cuda_toolkit_paths_full): break # Reset and retry @@ -1182,6 +1186,7 @@ def set_tf_nccl_install_path(environ_cp): if is_windows() or is_cygwin(): nccl_install_path = cygpath(nccl_install_path) + nccl_lib_path = '' if is_windows(): nccl_lib_path = 'lib/x64/nccl.lib' elif is_linux(): @@ -1417,11 +1422,16 @@ def set_mpi_home(environ_cp): def valid_mpi_path(mpi_home): exists = ( os.path.exists(os.path.join(mpi_home, 'include')) and - os.path.exists(os.path.join(mpi_home, 'lib'))) + (os.path.exists(os.path.join(mpi_home, 'lib')) or + os.path.exists(os.path.join(mpi_home, 'lib64')) or + os.path.exists(os.path.join(mpi_home, 'lib32')))) if not exists: - print('Invalid path to the MPI Toolkit. %s or %s cannot be found' % - (os.path.join(mpi_home, 'include'), - os.path.exists(os.path.join(mpi_home, 'lib')))) + print( + 'Invalid path to the MPI Toolkit. %s or %s or %s or %s cannot be found' + % (os.path.join(mpi_home, 'include'), + os.path.exists(os.path.join(mpi_home, 'lib')), + os.path.exists(os.path.join(mpi_home, 'lib64')), + os.path.exists(os.path.join(mpi_home, 'lib32')))) return exists _ = prompt_loop_or_load_from_env( @@ -1462,8 +1472,17 @@ def set_other_mpi_vars(environ_cp): if os.path.exists(os.path.join(mpi_home, 'lib/libmpi.so')): symlink_force( os.path.join(mpi_home, 'lib/libmpi.so'), 'third_party/mpi/libmpi.so') + elif os.path.exists(os.path.join(mpi_home, 'lib64/libmpi.so')): + symlink_force( + os.path.join(mpi_home, 'lib64/libmpi.so'), 'third_party/mpi/libmpi.so') + elif os.path.exists(os.path.join(mpi_home, 'lib32/libmpi.so')): + symlink_force( + os.path.join(mpi_home, 'lib32/libmpi.so'), 'third_party/mpi/libmpi.so') + else: - raise ValueError('Cannot find the MPI library file in %s/lib' % mpi_home) + raise ValueError( + 'Cannot find the MPI library file in %s/lib or %s/lib64 or %s/lib32' % + mpi_home, mpi_home, mpi_home) def set_system_libs_flag(environ_cp): @@ -1537,9 +1556,10 @@ def main(): # environment variables. environ_cp = dict(os.environ) - check_bazel_version('0.15.0') + check_bazel_version('0.19.0', '0.21.0') reset_tf_configure_bazelrc() + cleanup_makefile() setup_python(environ_cp) @@ -1674,10 +1694,10 @@ def main(): 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.') if __name__ == '__main__': main() - diff --git a/tensorflow/BUILD b/tensorflow/BUILD index 82526cead476bdc2eb9a5c5d53922d3a3d3ba5ae..0be7920d1193f2709a4d09d2d5e51daba35a5ff4 100644 --- a/tensorflow/BUILD +++ b/tensorflow/BUILD @@ -43,6 +43,11 @@ TENSORFLOW_API_INIT_FILES_V2 = ( TENSORFLOW_API_INIT_FILES + 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) +) + # Config setting used when building for products # which requires restricted licenses to be avoided. config_setting( @@ -197,6 +202,12 @@ config_setting( visibility = ["//visibility:public"], ) +config_setting( + name = "arm", + values = {"cpu": "arm"}, + visibility = ["//visibility:public"], +) + config_setting( name = "freebsd", values = {"cpu": "freebsd"}, @@ -213,31 +224,37 @@ config_setting( # config_setting( name = "no_aws_support", - define_values = {"no_aws_support": "false"}, + define_values = {"no_aws_support": "true"}, visibility = ["//visibility:public"], ) config_setting( name = "no_gcp_support", - define_values = {"no_gcp_support": "false"}, + define_values = {"no_gcp_support": "true"}, visibility = ["//visibility:public"], ) config_setting( name = "no_hdfs_support", - define_values = {"no_hdfs_support": "false"}, + define_values = {"no_hdfs_support": "true"}, visibility = ["//visibility:public"], ) config_setting( name = "no_ignite_support", - define_values = {"no_ignite_support": "false"}, + define_values = {"no_ignite_support": "true"}, visibility = ["//visibility:public"], ) config_setting( name = "no_kafka_support", - define_values = {"no_kafka_support": "false"}, + define_values = {"no_kafka_support": "true"}, + visibility = ["//visibility:public"], +) + +config_setting( + name = "no_nccl_support", + define_values = {"no_nccl_support": "true"}, visibility = ["//visibility:public"], ) @@ -256,6 +273,15 @@ config_setting( visibility = ["//visibility:public"], ) +# By default, XLA GPU is compiled into tensorflow when building with +# --config=cuda even when `with_xla_support` is false. The config setting +# here allows us to override the behavior if needed. +config_setting( + name = "no_xla_deps_in_cuda", + define_values = {"no_xla_deps_in_cuda": "true"}, + visibility = ["//visibility:public"], +) + config_setting( name = "with_gdr_support", define_values = {"with_gdr_support": "true"}, @@ -344,14 +370,26 @@ config_setting( define_values = {"tf_api_version": "2"}, ) +# This flag is defined for select statements that match both +# on 'windows' and 'api_version_2'. In this case, bazel requires +# having a flag which is a superset of these two. +config_setting( + name = "windows_and_api_version_2", + define_values = {"tf_api_version": "2"}, + values = {"cpu": "x64_windows"}, +) + package_group( name = "internal", packages = [ "-//third_party/tensorflow/python/estimator", + "//learning/deepmind/...", "//learning/meta_rank/...", + "//platforms/performance/autograppler/...", "//tensorflow/...", - "//tensorflow_estimator/...", + "//tensorflow_estimator/contrib/...", "//tensorflow_fold/llgtm/...", + "//tensorflow_text/...", "//third_party/py/tensor2tensor/...", ], ) @@ -553,29 +591,40 @@ genrule( }), outs = ["__init__.py"], cmd = select({ - "api_version_2": "cp $(@D)/_api/v2/__init__.py $(OUTS)", - "//conditions:default": "cp $(@D)/_api/v1/__init__.py $(OUTS)", + "api_version_2": "cp $(@D)/_api/v2/v2.py $(OUTS)", + "//conditions:default": "cp $(@D)/_api/v1/v1.py $(OUTS)", }), ) gen_api_init_files( name = "tf_python_api_gen_v1", - srcs = ["api_template.__init__.py"], + srcs = [ + "api_template_v1.__init__.py", + "compat_template_v1.__init__.py", + ], api_version = 1, + compat_api_versions = [1], + compat_init_templates = ["compat_template_v1.__init__.py"], output_dir = "_api/v1/", - output_files = TENSORFLOW_API_INIT_FILES_V1, + output_files = TENSORFLOW_API_INIT_FILES_V1_WITH_COMPAT, output_package = "tensorflow._api.v1", - root_init_template = "api_template.__init__.py", + root_file_name = "v1.py", + root_init_template = "api_template_v1.__init__.py", ) gen_api_init_files( name = "tf_python_api_gen_v2", - srcs = ["api_template.__init__.py"], + srcs = [ + "api_template.__init__.py", + "compat_template_v1.__init__.py", + ], api_version = 2, compat_api_versions = [1], + compat_init_templates = ["compat_template_v1.__init__.py"], output_dir = "_api/v2/", output_files = TENSORFLOW_API_INIT_FILES_V2, output_package = "tensorflow._api.v2", + root_file_name = "v2.py", root_init_template = "api_template.__init__.py", ) @@ -583,9 +632,11 @@ py_library( name = "tensorflow_py", srcs_version = "PY2AND3", visibility = ["//visibility:public"], - deps = [ + deps = select({ + "api_version_2": [], + "//conditions:default": ["//tensorflow/contrib:contrib_py"], + }) + [ ":tensorflow_py_no_contrib", - "//tensorflow/contrib:contrib_py", "//tensorflow/python/estimator:estimator_py", ], ) @@ -595,7 +646,11 @@ py_library( srcs = select({ "api_version_2": [":tf_python_api_gen_v2"], "//conditions:default": [":tf_python_api_gen_v1"], - }) + [":root_init_gen"], + }) + [":root_init_gen"] + [ + "//tensorflow/python/keras/api:keras_python_api_gen", + "//tensorflow/python/keras/api:keras_python_api_gen_compat_v1", + "//tensorflow/python/keras/api:keras_python_api_gen_compat_v2", + ], srcs_version = "PY2AND3", visibility = ["//visibility:public"], deps = ["//tensorflow/python:no_contrib"], diff --git a/tensorflow/api_template.__init__.py b/tensorflow/api_template.__init__.py index 65172fd74a1660adc021ae97f769b05483bc0ba0..a93799bfe84b0f9c4743e1ad0effd6e69ad7f3f2 100644 --- a/tensorflow/api_template.__init__.py +++ b/tensorflow/api_template.__init__.py @@ -18,36 +18,77 @@ from __future__ import absolute_import as _absolute_import from __future__ import division as _division from __future__ import print_function as _print_function +import distutils as _distutils +import inspect as _inspect import os as _os +import site as _site +import sys as _sys + +# API IMPORTS PLACEHOLDER # pylint: disable=g-bad-import-order -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._v2.estimator')) -from tensorflow.python.tools import component_api_helper -component_api_helper.package_hook( +_current_module = _sys.modules[__name__] +if not hasattr(_current_module, 'estimator'): + _component_api_helper.package_hook( + parent_package_str=__name__, + child_package_str=( + 'tensorflow_estimator.python.estimator.api.estimator')) +_component_api_helper.package_hook( parent_package_str=__name__, - child_package_str=('tensorflow_estimator.python.estimator.api.estimator')) -del component_api_helper + child_package_str=('tensorflow.python.keras.api._v2.keras')) +# Make sure directory containing top level submodules is in +# the __path__ so that "from tensorflow.foo import bar" works. +# We're using bitwise, but there's nothing special about that. +_tf_api_dir = _os.path.dirname(_os.path.dirname(bitwise.__file__)) # pylint: disable=undefined-variable +if not hasattr(_current_module, '__path__'): + __path__ = [_tf_api_dir] +elif _tf_api_dir not in __path__: + __path__.append(_tf_api_dir) -# API IMPORTS PLACEHOLDER +# Enable TF2 behaviors +from tensorflow.python.compat import v2_compat as _compat # pylint: disable=g-import-not-at-top +_compat.enable_v2_behavior() -from tensorflow.python.util.lazy_loader import LazyLoader # pylint: disable=g-import-not-at-top -contrib = LazyLoader('contrib', globals(), 'tensorflow.contrib') -del LazyLoader -# The templated code that replaces the placeholder above sometimes -# sets the __all__ variable. If it does, we have to be sure to add -# "contrib". -if '__all__' in vars(): - vars()['__all__'].append('contrib') -from tensorflow.python.platform import flags # pylint: disable=g-import-not-at-top -app.flags = flags # pylint: disable=undefined-variable +# Load all plugin libraries from site-packages/tensorflow-plugins if we are +# running under pip. +# TODO(gunan): Enable setting an environment variable to define arbitrary plugin +# directories. +# TODO(gunan): Find a better location for this code snippet. +from tensorflow.python.framework import load_library as _ll +from tensorflow.python.lib.io import file_io as _fi -# 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 -if _tf_api_dir not in __path__: - __path__.append(_tf_api_dir) +# Get sitepackages directories for the python installation. +_site_packages_dirs = [] +_site_packages_dirs += [_site.USER_SITE] +_site_packages_dirs += [_p for _p in _sys.path if 'site-packages' in _p] +if 'getsitepackages' in dir(_site): + _site_packages_dirs += _site.getsitepackages() + +if 'sysconfig' in dir(_distutils): + _site_packages_dirs += [_distutils.sysconfig.get_python_lib()] + +_site_packages_dirs = list(set(_site_packages_dirs)) + +# Find the location of this exact file. +_current_file_location = _inspect.getfile(_inspect.currentframe()) + +def _running_from_pip_package(): + return any( + _current_file_location.startswith(dir_) for dir_ in _site_packages_dirs) + +if _running_from_pip_package(): + for s in _site_packages_dirs: + # TODO(gunan): Add sanity checks to loaded modules here. + plugin_dir = _os.path.join(s, 'tensorflow-plugins') + if _fi.file_exists(plugin_dir): + _ll.load_library(plugin_dir) # These symbols appear because we import the python package which # in turn imports from tensorflow.core and tensorflow.python. They @@ -59,7 +100,16 @@ try: del core except NameError: # Don't fail if these modules are not available. - # For e.g. we are using this file for compat.v1 module as well and - # 'python', 'core' directories are not under compat/v1. + # For e.g. this file will be originally placed under tensorflow/_api/v1 which + # does not have 'python', 'core' directories. Then, it will be copied + # to tensorflow/ which does have these two directories. pass +# Similarly for compiler. Do it separately to make sure we do this even if the +# others don't exist. +try: + del compiler +except NameError: + pass + + # pylint: enable=undefined-variable diff --git a/tensorflow/api_template_v1.__init__.py b/tensorflow/api_template_v1.__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..eeca8f0d566a6401cb64e4fe3f0ee3c5aeb4ece2 --- /dev/null +++ b/tensorflow/api_template_v1.__init__.py @@ -0,0 +1,133 @@ +# 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 distutils as _distutils +import inspect as _inspect +import os as _os +import site as _site +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._v1.estimator')) + +_current_module = _sys.modules[__name__] +if not hasattr(_current_module, 'estimator'): + _component_api_helper.package_hook( + parent_package_str=__name__, + child_package_str=( + 'tensorflow_estimator.python.estimator.api.estimator')) +_component_api_helper.package_hook( + parent_package_str=__name__, + child_package_str=('tensorflow.python.keras.api._v1.keras')) +from tensorflow.python.util.lazy_loader import LazyLoader # pylint: disable=g-import-not-at-top +_CONTRIB_WARNING = """ +WARNING: The TensorFlow contrib module will not be included in TensorFlow 2.0. +For more information, please see: + * https://github.com/tensorflow/community/blob/master/rfcs/20180907-contrib-sunset.md + * https://github.com/tensorflow/addons +If you depend on functionality not listed there, please file an issue. +""" +contrib = LazyLoader('contrib', globals(), 'tensorflow.contrib', + _CONTRIB_WARNING) +del LazyLoader +# The templated code that replaces the placeholder above sometimes +# sets the __all__ variable. If it does, we have to be sure to add +# "contrib". +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(_API_MODULE.__file__)) # pylint: disable=undefined-variable +if not hasattr(_current_module, '__path__'): + __path__ = [_tf_api_dir] +elif _tf_api_dir not in __path__: + __path__.append(_tf_api_dir) + +# Load all plugin libraries from site-packages/tensorflow-plugins if we are +# running under pip. +# TODO(gunan): Enable setting an environment variable to define arbitrary plugin +# directories. +# TODO(gunan): Find a better location for this code snippet. +from tensorflow.python.framework import load_library as _ll +from tensorflow.python.lib.io import file_io as _fi + +# Get sitepackages directories for the python installation. +_site_packages_dirs = [] +_site_packages_dirs += [_site.USER_SITE] +_site_packages_dirs += [_p for _p in _sys.path if 'site-packages' in _p] +if 'getsitepackages' in dir(_site): + _site_packages_dirs += _site.getsitepackages() + +if 'sysconfig' in dir(_distutils): + _site_packages_dirs += [_distutils.sysconfig.get_python_lib()] + +_site_packages_dirs = list(set(_site_packages_dirs)) + +# Find the location of this exact file. +_current_file_location = _inspect.getfile(_inspect.currentframe()) + +def _running_from_pip_package(): + return any( + _current_file_location.startswith(dir_) for dir_ in _site_packages_dirs) + +if _running_from_pip_package(): + for s in _site_packages_dirs: + # TODO(gunan): Add sanity checks to loaded modules here. + plugin_dir = _os.path.join(s, 'tensorflow-plugins') + if _fi.file_exists(plugin_dir): + _ll.load_library(plugin_dir) + +# These symbols appear because we import the python package which +# in turn imports from tensorflow.core and tensorflow.python. They +# must come from this module. So python adds these symbols for the +# resolution to succeed. +# pylint: disable=undefined-variable +try: + del python + del core +except NameError: + # Don't fail if these modules are not available. + # For e.g. this file will be originally placed under tensorflow/_api/v1 which + # does not have 'python', 'core' directories. Then, it will be copied + # to tensorflow/ which does have these two directories. + pass +# Similarly for compiler. Do it separately to make sure we do this even if the +# others don't exist. +try: + del compiler +except NameError: + pass +# pylint: enable=undefined-variable diff --git a/tensorflow/c/BUILD b/tensorflow/c/BUILD index 56f5e6767ac68b1008c786e3b5a47b9b173ab9cb..6e50a09bfc5ed3a8f2f7e05e6a6a151525e8dfce 100644 --- a/tensorflow/c/BUILD +++ b/tensorflow/c/BUILD @@ -60,6 +60,7 @@ tf_cuda_library( "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:op_gen_lib", + "//tensorflow/core/distributed_runtime:server_lib", ], }), ) @@ -82,7 +83,7 @@ tf_cuda_library( ], "//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", @@ -95,6 +96,7 @@ tf_cuda_library( "//tensorflow/core:protos_all_cc", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", + "//tensorflow/core/distributed_runtime:server_lib", ], }) + select({ "//tensorflow:with_xla_support": [ @@ -119,13 +121,15 @@ tf_cuda_library( ":c_api", ":c_api_internal", "//tensorflow/c/eager:c_api", - "//tensorflow/compiler/jit/legacy_flags:mark_for_compilation_pass_flags", - "//tensorflow/contrib/tpu:all_ops", + "//tensorflow/c/eager:c_api_internal", + "//tensorflow/compiler/jit:flags", "//tensorflow/core:core_cpu", "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:lib_platform", "//tensorflow/core:protos_all_cc", + "//tensorflow/core/common_runtime/eager:attr_builder", + "@com_google_absl//absl/strings", ], ) @@ -171,6 +175,58 @@ tf_cuda_library( ], ) +tf_cuda_library( + name = "env", + srcs = [ + "env.cc", + ], + hdrs = [ + "env.h", + ], + copts = tf_copts(), + visibility = ["//visibility:public"], + deps = select({ + "//tensorflow:android": [ + ":c_api", + ":tf_status_helper", + "//tensorflow/core:android_tensorflow_lib_lite", + "//tensorflow/core:lib", + ], + "//conditions:default": [ + ":c_api", + ":tf_status_helper", + "//tensorflow/core:framework", + "//tensorflow/core:lib", + ], + }) + [":c_api_internal"], +) + +tf_cuda_library( + name = "kernels", + srcs = [ + "kernels.cc", + ], + hdrs = [ + "kernels.h", + ], + copts = tf_copts(), + visibility = ["//visibility:public"], + deps = select({ + "//tensorflow:android": [ + ":c_api", + ":c_api_internal", + ":tf_status_helper", + "//tensorflow/core:android_tensorflow_lib_lite", + ], + "//conditions:default": [ + ":c_api", + ":c_api_internal", + ":tf_status_helper", + "//tensorflow/core:framework", + ], + }), +) + # ----------------------------------------------------------------------------- # Tests @@ -193,11 +249,24 @@ tf_cuda_library( ], ) +tf_cc_test( + name = "c_test", + srcs = ["c_test.c"], + extra_copts = ["-std=c11"], + deps = [ + ":c_api", + ":c_api_experimental", + ":env", + ":kernels", + ], +) + tf_cuda_cc_test( name = "c_api_test", size = "small", srcs = ["c_api_test.cc"], data = [ + ":test_op1.so", "//tensorflow/cc/saved_model:saved_model_half_plus_two", ], kernels = [":test_op_kernel"], @@ -205,7 +274,10 @@ tf_cuda_cc_test( "//tensorflow:darwin": ["-headerpad_max_install_names"], "//conditions:default": [], }), - tags = ["noasan"], + tags = [ + "no_oss", # http://b/119522529 + "noasan", + ], # We must ensure that the dependencies can be dynamically linked since # the shared library must be able to use core:framework. # linkstatic = tf_kernel_tests_linkstatic(), @@ -216,13 +288,21 @@ tf_cuda_cc_test( "//tensorflow/cc:grad_ops", "//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:proto_text", "//tensorflow/core:protos_all_cc", + "//tensorflow/core:spectral_ops_op_lib", "//tensorflow/core:test", "//tensorflow/core:test_main", "//tensorflow/core/kernels:array", @@ -233,7 +313,7 @@ tf_cuda_cc_test( tf_cc_test( name = "c_api_experimental_test", - size = "small", + size = "medium", srcs = ["c_api_experimental_test.cc"], data = ["testdata/tf_record"], linkopts = select({ @@ -244,8 +324,12 @@ tf_cc_test( # the shared library must be able to use core:framework. # linkstatic = tf_kernel_tests_linkstatic(), deps = [ + ":c_api", ":c_api_experimental", + ":c_api_internal", ":c_test_util", + "//tensorflow/c/eager:c_api", + "//tensorflow/c/eager:c_api_test_util", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", "//tensorflow/core:test", @@ -282,8 +366,8 @@ tf_cc_test( ) tf_custom_op_library( - name = "test_op.so", - srcs = ["test_op.cc"], + name = "test_op1.so", + srcs = ["test_op1.cc"], ) tf_kernel_library( @@ -296,6 +380,51 @@ tf_kernel_library( alwayslink = 1, ) +tf_cuda_cc_test( + name = "env_test", + size = "small", + srcs = ["env_test.cc"], + linkopts = select({ + "//tensorflow:darwin": ["-headerpad_max_install_names"], + "//conditions:default": [], + }), + tags = ["noasan"], + # We must ensure that the dependencies can be dynamically linked since + # the shared library must be able to use core:framework. + # linkstatic = tf_kernel_tests_linkstatic(), + deps = [ + ":c_api", + ":env", + "//tensorflow/core:lib", + "//tensorflow/core:test", + "//tensorflow/core:test_main", + ], +) + +tf_cuda_cc_test( + name = "kernels_test", + size = "small", + srcs = ["kernels_test.cc"], + linkopts = select({ + "//tensorflow:darwin": ["-headerpad_max_install_names"], + "//conditions:default": [], + }), + tags = ["noasan"], + # We must ensure that the dependencies can be dynamically linked since + # the shared library must be able to use core:framework. + # linkstatic = tf_kernel_tests_linkstatic(), + deps = [ + ":c_api", + ":kernels", + "//tensorflow/core:framework", + "//tensorflow/core:lib", + "//tensorflow/core:proto_text", + "//tensorflow/core:protos_all_cc", + "//tensorflow/core:test", + "//tensorflow/core:test_main", + ], +) + # ----------------------------------------------------------------------------- # Python API target diff --git a/tensorflow/c/README.md b/tensorflow/c/README.md new file mode 100644 index 0000000000000000000000000000000000000000..b386998ceaf3e91daba04125fe83e2f3bdd508e5 --- /dev/null +++ b/tensorflow/c/README.md @@ -0,0 +1,7 @@ +# TensorFlow C API + +- See [www.tensorflow.org/install/lang_c](https://www.tensorflow.org/install/lang_c) +- Nightly builds: + - [Linux CPU-only](https://storage.googleapis.com/tensorflow-nightly/github/tensorflow/lib_package/libtensorflow-cpu-linux-x86_64.tar.gz) + - [Linux GPU](https://storage.googleapis.com/tensorflow-nightly/github/tensorflow/lib_package/libtensorflow-gpu-linux-x86_64.tar.gz) + - [MacOS CPU-only](https://storage.googleapis.com/tensorflow-nightly/github/tensorflow/lib_package/libtensorflow-cpu-darwin-x86_64.tar.gz) diff --git a/tensorflow/c/c_api.cc b/tensorflow/c/c_api.cc index 1726db12fa62c5a3665de9fc306da38c1b7f0f9c..94d9f4a6fa2f14cb3343bdd51b7e4d61944444d0 100644 --- a/tensorflow/c/c_api.cc +++ b/tensorflow/c/c_api.cc @@ -136,16 +136,22 @@ const char* TF_Message(const TF_Status* s) { namespace { class TF_ManagedBuffer : public TensorBuffer { public: - void* data_; - size_t len_; - void (*deallocator_)(void* data, size_t len, void* arg); - void* deallocator_arg_; + TF_ManagedBuffer(void* data, size_t len, + void (*deallocator)(void* data, size_t len, void* arg), + void* deallocator_arg) + : TensorBuffer(data), + len_(len), + deallocator_(deallocator), + deallocator_arg_(deallocator_arg) {} + + const size_t len_; + void (*const deallocator_)(void* data, size_t len, void* arg); + void* const deallocator_arg_; ~TF_ManagedBuffer() override { - (*deallocator_)(data_, len_, deallocator_arg_); + (*deallocator_)(data(), len_, deallocator_arg_); } - void* data() const override { return data_; } size_t size() const override { return len_; } TensorBuffer* root_buffer() override { return this; } void FillAllocationDescription(AllocationDescription* proto) const override { @@ -199,8 +205,7 @@ TF_Tensor* TF_NewTensor(TF_DataType dtype, const int64_t* dims, int num_dims, dimvec[i] = static_cast(dims[i]); } - TF_ManagedBuffer* buf = new TF_ManagedBuffer; - buf->len_ = len; + TF_ManagedBuffer* buf = nullptr; if (dtype != TF_STRING && dtype != TF_RESOURCE && tensorflow::DataTypeCanUseMemcpy(static_cast(dtype)) && reinterpret_cast(data) % std::max(1, EIGEN_MAX_ALIGN_BYTES) != @@ -212,17 +217,15 @@ TF_Tensor* TF_NewTensor(TF_DataType dtype, const int64_t* dims, int num_dims, // // Other types have the same representation, so copy only if it is safe to // do so. - buf->data_ = allocate_tensor("TF_NewTensor", len); - std::memcpy(buf->data_, data, len); - buf->deallocator_ = deallocate_buffer; - buf->deallocator_arg_ = nullptr; + buf = new TF_ManagedBuffer(allocate_tensor("TF_NewTensor", len), len, + deallocate_buffer, nullptr); + std::memcpy(buf->data(), data, len); // Free the original buffer. deallocator(data, len, deallocator_arg); } else { - buf->data_ = data; - buf->deallocator_ = deallocator; - buf->deallocator_arg_ = deallocator_arg; + buf = new TF_ManagedBuffer(data, len, deallocator, deallocator_arg); } + TF_Tensor* ret = new TF_Tensor{dtype, TensorShape(dimvec), buf}; size_t elem_size = TF_DataTypeSize(dtype); if (elem_size > 0 && len < (elem_size * ret->shape.num_elements())) { @@ -254,6 +257,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) { @@ -477,14 +548,15 @@ static TF_Tensor* EmptyTensor(TF_DataType dtype, const TensorShape& shape) { CHECK_EQ(nelems, 0); static_assert(sizeof(int64_t) == sizeof(tensorflow::int64), "64-bit int types should match in size"); - return TF_NewTensor(dtype, reinterpret_cast(dims.data()), - shape.dims(), reinterpret_cast(&empty), 0, - [](void*, size_t, void*) {}, nullptr); + return TF_NewTensor( + dtype, reinterpret_cast(dims.data()), shape.dims(), + reinterpret_cast(&empty), 0, [](void*, size_t, void*) {}, nullptr); } // Non-static for testing. TF_Tensor* TF_TensorFromTensor(const tensorflow::Tensor& src, TF_Status* status) { + TF_SetStatus(status, TF_OK, ""); if (!src.IsInitialized()) { status->status = FailedPrecondition( "attempt to use a tensor with an uninitialized value"); @@ -1592,18 +1664,20 @@ TF_AttrMetadata TF_OperationGetAttrMetadata(TF_Operation* oper, break; \ } - LIST_CASE(s, TF_ATTR_STRING, metadata.total_size = 0; - for (int i = 0; i < attr->list().s_size(); - ++i) { metadata.total_size += attr->list().s(i).size(); }); + LIST_CASE( + s, TF_ATTR_STRING, metadata.total_size = 0; + for (int i = 0; i < attr->list().s_size(); + ++i) { metadata.total_size += attr->list().s(i).size(); }); LIST_CASE(i, TF_ATTR_INT); LIST_CASE(f, TF_ATTR_FLOAT); LIST_CASE(b, TF_ATTR_BOOL); LIST_CASE(type, TF_ATTR_TYPE); - LIST_CASE(shape, TF_ATTR_SHAPE, metadata.total_size = 0; - for (int i = 0; i < attr->list().shape_size(); ++i) { - const auto& s = attr->list().shape(i); - metadata.total_size += s.unknown_rank() ? 0 : s.dim_size(); - }); + LIST_CASE( + shape, TF_ATTR_SHAPE, metadata.total_size = 0; + for (int i = 0; i < attr->list().shape_size(); ++i) { + const auto& s = attr->list().shape(i); + metadata.total_size += s.unknown_rank() ? 0 : s.dim_size(); + }); LIST_CASE(tensor, TF_ATTR_TENSOR); LIST_CASE(tensor, TF_ATTR_FUNC); #undef LIST_CASE @@ -1942,6 +2016,10 @@ void TF_ImportGraphDefOptionsSetPrefix(TF_ImportGraphDefOptions* opts, const char* prefix) { opts->opts.prefix = prefix; } +void TF_ImportGraphDefOptionsSetDefaultDevice(TF_ImportGraphDefOptions* opts, + const char* device) { + opts->opts.default_device = device; +} void TF_ImportGraphDefOptionsSetUniquifyNames(TF_ImportGraphDefOptions* opts, unsigned char uniquify_names) { @@ -2806,4 +2884,74 @@ TF_Buffer* TF_GetRegisteredKernelsForOp(const char* name, TF_Status* status) { } return ret; } + +// TF_Server functions ---------------------------------------------- + +#ifndef __ANDROID__ +TF_Server::TF_Server(std::unique_ptr server) + : target(server->target()), server(std::move(server)) {} +#endif // __ANDROID__ + +TF_Server* TF_NewServer(const void* proto, size_t proto_len, + TF_Status* status) { +#ifdef __ANDROID__ + status->status = tensorflow::errors::Unimplemented( + "Server functionality is not supported in Android"); + return nullptr; +#else + tensorflow::ServerDef server_def; + if (!server_def.ParseFromArray(proto, static_cast(proto_len))) { + status->status = InvalidArgument( + "Could not parse provided bytes into a ServerDef protocol buffer"); + return nullptr; + } + + std::unique_ptr out_server; + status->status = tensorflow::NewServer(server_def, &out_server); + if (!status->status.ok()) return nullptr; + + return new TF_Server(std::move(out_server)); +#endif +} + +void TF_ServerStart(TF_Server* server, TF_Status* status) { +#ifdef __ANDROID__ + status->status = tensorflow::errors::Unimplemented( + "Server functionality is not supported in Android"); +#else + status->status = server->server->Start(); +#endif +} + +void TF_ServerStop(TF_Server* server, TF_Status* status) { +#ifdef __ANDROID__ + status->status = tensorflow::errors::Unimplemented( + "Server functionality is not supported in Android"); +#else + status->status = server->server->Stop(); +#endif +} + +void TF_ServerJoin(TF_Server* server, TF_Status* status) { +#ifdef __ANDROID__ + status->status = tensorflow::errors::Unimplemented( + "Server functionality is not supported in Android"); +#else + status->status = server->server->Join(); +#endif +} + +const char* TF_ServerTarget(TF_Server* server) { +#ifdef __ANDROID__ + return nullptr; +#else + return server->target.c_str(); +#endif +} + +void TF_DeleteServer(TF_Server* server) { +#ifndef __ANDROID__ + delete server; +#endif +} } // end extern "C" diff --git a/tensorflow/c/c_api.h b/tensorflow/c/c_api.h index 850f6ecd637d768bca99720e0add07680829e17a..8031928dac4de2391f0aec46e69d61a137606e4d 100644 --- a/tensorflow/c/c_api.h +++ b/tensorflow/c/c_api.h @@ -91,7 +91,7 @@ extern "C" { // -------------------------------------------------------------------------- // TF_Version returns a string describing version information of the // TensorFlow library. TensorFlow using semantic versioning. -TF_CAPI_EXPORT extern const char* TF_Version(); +TF_CAPI_EXPORT extern const char* TF_Version(void); // -------------------------------------------------------------------------- // TF_DataType holds the type for a scalar value. E.g., one slot in a tensor. @@ -157,7 +157,7 @@ typedef enum TF_Code { typedef struct TF_Status TF_Status; // Return a new status object. -TF_CAPI_EXPORT extern TF_Status* TF_NewStatus(); +TF_CAPI_EXPORT extern TF_Status* TF_NewStatus(void); // Delete a previously created status object. TF_CAPI_EXPORT extern void TF_DeleteStatus(TF_Status*); @@ -196,7 +196,7 @@ TF_CAPI_EXPORT extern TF_Buffer* TF_NewBufferFromString(const void* proto, size_t proto_len); // Useful for passing *out* a protobuf. -TF_CAPI_EXPORT extern TF_Buffer* TF_NewBuffer(); +TF_CAPI_EXPORT extern TF_Buffer* TF_NewBuffer(void); TF_CAPI_EXPORT extern void TF_DeleteBuffer(TF_Buffer*); @@ -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` @@ -305,7 +338,7 @@ TF_CAPI_EXPORT extern size_t TF_StringEncodedSize(size_t len); typedef struct TF_SessionOptions TF_SessionOptions; // Return a new options object. -TF_CAPI_EXPORT extern TF_SessionOptions* TF_NewSessionOptions(); +TF_CAPI_EXPORT extern TF_SessionOptions* TF_NewSessionOptions(void); // Set the target in TF_SessionOptions.options. // target can be empty, a single entry, or a comma separated list of entries. @@ -338,7 +371,7 @@ TF_CAPI_EXPORT extern void TF_DeleteSessionOptions(TF_SessionOptions*); typedef struct TF_Graph TF_Graph; // Return a new graph object. -TF_CAPI_EXPORT extern TF_Graph* TF_NewGraph(); +TF_CAPI_EXPORT extern TF_Graph* TF_NewGraph(void); // Destroy an options object. Graph will be deleted once no more // TFSession's are referencing it. @@ -890,7 +923,8 @@ TF_CAPI_EXPORT extern void TF_GraphVersions(TF_Graph* graph, // TF_GraphImportGraphDef. typedef struct TF_ImportGraphDefOptions TF_ImportGraphDefOptions; -TF_CAPI_EXPORT extern TF_ImportGraphDefOptions* TF_NewImportGraphDefOptions(); +TF_CAPI_EXPORT extern TF_ImportGraphDefOptions* TF_NewImportGraphDefOptions( + void); TF_CAPI_EXPORT extern void TF_DeleteImportGraphDefOptions( TF_ImportGraphDefOptions* opts); @@ -900,6 +934,12 @@ TF_CAPI_EXPORT extern void TF_DeleteImportGraphDefOptions( TF_CAPI_EXPORT extern void TF_ImportGraphDefOptionsSetPrefix( TF_ImportGraphDefOptions* opts, const char* prefix); +// Set the execution device for nodes in `graph_def`. +// Only applies to nodes where a device was not already explicitly specified. +// `device` is copied and has no lifetime requirements. +TF_CAPI_EXPORT extern void TF_ImportGraphDefOptionsSetDefaultDevice( + TF_ImportGraphDefOptions* opts, const char* device); + // Set whether to uniquify imported operation names. If true, imported operation // names will be modified if their name already exists in the graph. If false, // conflicting names will be treated as an error. Note that this option has no @@ -1605,7 +1645,7 @@ TF_CAPI_EXPORT extern void TF_DeleteLibraryHandle(TF_Library* lib_handle); // // The data in the buffer will be the serialized OpList proto for ops registered // in this address space. -TF_CAPI_EXPORT extern TF_Buffer* TF_GetAllOpList(); +TF_CAPI_EXPORT extern TF_Buffer* TF_GetAllOpList(void); // TF_ApiDefMap encapsulates a collection of API definitions for an operation. // @@ -1662,6 +1702,47 @@ TF_CAPI_EXPORT extern TF_Buffer* TF_GetAllRegisteredKernels(TF_Status* status); TF_CAPI_EXPORT extern TF_Buffer* TF_GetRegisteredKernelsForOp( const char* name, TF_Status* status); +// -------------------------------------------------------------------------- +// In-process TensorFlow server functionality, for use in distributed training. +// A Server instance encapsulates a set of devices and a Session target that +// can participate in distributed training. A server belongs to a cluster +// (specified by a ClusterSpec), and corresponds to a particular task in a +// named job. The server can communicate with any other server in the same +// cluster. + +// In-process TensorFlow server. +typedef struct TF_Server TF_Server; + +// Creates a new in-process TensorFlow server configured using a serialized +// ServerDef protocol buffer provided via `proto` and `proto_len`. +// +// The server will not serve any requests until TF_ServerStart is invoked. +// The server will stop serving requests once TF_ServerStop or +// TF_DeleteServer is invoked. +TF_CAPI_EXPORT extern TF_Server* TF_NewServer(const void* proto, + size_t proto_len, + TF_Status* status); + +// Starts an in-process TensorFlow server. +TF_CAPI_EXPORT extern void TF_ServerStart(TF_Server* server, TF_Status* status); + +// Stops an in-process TensorFlow server. +TF_CAPI_EXPORT extern void TF_ServerStop(TF_Server* server, TF_Status* status); + +// Blocks until the server has been successfully stopped (via TF_ServerStop or +// TF_ServerClose). +TF_CAPI_EXPORT extern void TF_ServerJoin(TF_Server* server, TF_Status* status); + +// Returns the target string that can be provided to TF_SetTarget() to connect +// a TF_Session to `server`. +// +// The returned string is valid only until TF_DeleteServer is invoked. +TF_CAPI_EXPORT extern const char* TF_ServerTarget(TF_Server* server); + +// Destroy an in-process TensorFlow server, frees memory. If server is running +// it will be stopped and joined. +TF_CAPI_EXPORT extern void TF_DeleteServer(TF_Server* server); + #ifdef __cplusplus } /* end extern "C" */ #endif diff --git a/tensorflow/c/c_api_experimental.cc b/tensorflow/c/c_api_experimental.cc index d4b78138e93624a7e41e917f8210281b500661bc..6cc74cfb3246e9526e862f363590ce43e390ffaa 100644 --- a/tensorflow/c/c_api_experimental.cc +++ b/tensorflow/c/c_api_experimental.cc @@ -15,12 +15,19 @@ 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/compiler/jit/legacy_flags/mark_for_compilation_pass_flags.h" +#include "tensorflow/c/eager/c_api.h" +#include "tensorflow/c/eager/c_api_internal.h" +#include "tensorflow/compiler/jit/flags.h" +#include "tensorflow/core/common_runtime/eager/attr_builder.h" #include "tensorflow/core/framework/tensor.pb.h" #include "tensorflow/core/graph/graph.h" #include "tensorflow/core/graph/node_builder.h" #include "tensorflow/core/lib/strings/strcat.h" +#include "tensorflow/core/platform/init_main.h" +#include "tensorflow/core/platform/net.h" #include "tensorflow/core/platform/platform.h" #include "tensorflow/core/protobuf/config.pb.h" #include "tensorflow/core/protobuf/tensorflow_server.pb.h" @@ -50,8 +57,8 @@ void TF_EnableXLACompilation(TF_SessionOptions* options, unsigned char enable) { // These XLA flags are needed to trigger XLA properly from C (more generally // non-Python) clients. If this API is called again with `enable` set to // false, it is safe to keep these flag values as is. - tensorflow::legacy_flags::MarkForCompilationPassFlags* flags = - tensorflow::legacy_flags::GetMarkForCompilationPassFlags(); + tensorflow::MarkForCompilationPassFlags* flags = + tensorflow::GetMarkForCompilationPassFlags(); flags->tf_xla_cpu_global_jit = true; flags->tf_xla_min_cluster_size = 1; } else { @@ -60,7 +67,8 @@ void TF_EnableXLACompilation(TF_SessionOptions* options, unsigned char enable) { } TF_Buffer* TF_CreateConfig(unsigned char enable_xla_compilation, - unsigned char gpu_memory_allow_growth) { + unsigned char gpu_memory_allow_growth, + unsigned int num_cpu_devices) { tensorflow::ConfigProto config; auto* optimizer_options = config.mutable_graph_options()->mutable_optimizer_options(); @@ -70,8 +78,8 @@ TF_Buffer* TF_CreateConfig(unsigned char enable_xla_compilation, // These XLA flags are needed to trigger XLA properly from C (more generally // non-Python) clients. If this API is called again with `enable` set to // false, it is safe to keep these flag values as is. - tensorflow::legacy_flags::MarkForCompilationPassFlags* flags = - tensorflow::legacy_flags::GetMarkForCompilationPassFlags(); + tensorflow::MarkForCompilationPassFlags* flags = + tensorflow::GetMarkForCompilationPassFlags(); flags->tf_xla_cpu_global_jit = true; flags->tf_xla_min_cluster_size = 1; } else { @@ -81,6 +89,8 @@ TF_Buffer* TF_CreateConfig(unsigned char enable_xla_compilation, auto* gpu_options = config.mutable_gpu_options(); gpu_options->set_allow_growth(gpu_memory_allow_growth); + (*config.mutable_device_count())["CPU"] = num_cpu_devices; + // TODO(b/113217601): This is needed for EagerContext::runner_ to use a // threadpool, so that we avoid the possibility of running the runner_ in the // threadpool of GPU event mgr, as that can trigger more callbacks to be @@ -119,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. // @@ -6524,7 +6542,7 @@ library { } } node_def { - name: "ParallelInterleaveDataset/cycle_length" + name: "ExperimentalParallelInterleaveDataset/cycle_length" op: "Const" attr { key: "dtype" @@ -6545,7 +6563,7 @@ library { } } node_def { - name: "ParallelInterleaveDataset/block_length" + name: "ExperimentalParallelInterleaveDataset/block_length" op: "Const" attr { key: "dtype" @@ -6566,7 +6584,7 @@ library { } } node_def { - name: "ParallelInterleaveDataset/sloppy" + name: "ExperimentalParallelInterleaveDataset/sloppy" op: "Const" attr { key: "dtype" @@ -6587,7 +6605,7 @@ library { } } node_def { - name: "ParallelInterleaveDataset/buffer_output_elements" + name: "ExperimentalParallelInterleaveDataset/buffer_output_elements" op: "Const" attr { key: "dtype" @@ -6608,7 +6626,7 @@ library { } } node_def { - name: "ParallelInterleaveDataset/prefetch_input_elements" + name: "ExperimentalParallelInterleaveDataset/prefetch_input_elements" op: "Const" attr { key: "dtype" @@ -6629,14 +6647,14 @@ library { } } node_def { - name: "ParallelInterleaveDataset" - op: "ParallelInterleaveDataset" + name: "ExperimentalParallelInterleaveDataset" + op: "ExperimentalParallelInterleaveDataset" input: "RepeatDataset:handle:0" - input: "ParallelInterleaveDataset/cycle_length:output:0" - input: "ParallelInterleaveDataset/block_length:output:0" - input: "ParallelInterleaveDataset/sloppy:output:0" - input: "ParallelInterleaveDataset/buffer_output_elements:output:0" - input: "ParallelInterleaveDataset/prefetch_input_elements:output:0" + input: "ExperimentalParallelInterleaveDataset/cycle_length:output:0" + input: "ExperimentalParallelInterleaveDataset/block_length:output:0" + input: "ExperimentalParallelInterleaveDataset/sloppy:output:0" + input: "ExperimentalParallelInterleaveDataset/buffer_output_elements:output:0" + input: "ExperimentalParallelInterleaveDataset/prefetch_input_elements:output:0" attr { key: "Targuments" value { @@ -6736,7 +6754,7 @@ library { node_def { name: "ShuffleDataset_2" op: "ShuffleDataset" - input: "ParallelInterleaveDataset:handle:0" + input: "ExperimentalParallelInterleaveDataset:handle:0" input: "ShuffleDataset_2/buffer_size_1:output:0" input: "ShuffleDataset_2/seed_2:output:0" input: "ShuffleDataset_2/seed2_2:output:0" @@ -8529,8 +8547,9 @@ TFE_Context* TFE_CreateContextFromSession(TF_Session* session, // Reduce GPU memory allocation, and set appropriate config options for TFE // context. - auto* config = - TF_CreateConfig(/*xla*/ false, /* gpu_memory_allow_growth */ true); + auto* config = TF_CreateConfig( + /*xla*/ false, /* gpu_memory_allow_growth */ true, /* num_cpu_devices */ + 10); TFE_ContextOptionsSetConfig(opts, config->data, config->length, status); if (!status->status.ok()) { CHECK(!config); @@ -8727,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); @@ -8738,7 +8763,359 @@ void TFE_TensorHandlePrintDebugString(TFE_TensorHandle* handle) { TF_DeleteStatus(status); } -TF_CAPI_EXPORT extern void TF_MakeInternalErrorStatus(TF_Status* status, - const char* errMsg) { +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; + std::unique_ptr thread; + std::unique_ptr status; +}; + +TFE_ExecuteOpNotification* TFE_ExecuteOpInNewThread(TFE_Op* op, + TFE_TensorHandle** retvals, + int* num_retvals, + TF_Status* status) { + TFE_ExecuteOpNotification* n = new TFE_ExecuteOpNotification; + + n->thread.reset(op->operation.EagerContext()->TFEnv()->StartThread( + tensorflow::ThreadOptions(), "ExecuteOpThread", + [op, retvals, num_retvals, n]() { + TFE_Execute(op, retvals, num_retvals, n->status.get()); + n->n.Notify(); + })); + + return n; +} + +void TFE_ExecuteOpNotificationWaitAndDelete( + TFE_ExecuteOpNotification* notification, TF_Status* status) { + if (notification == nullptr) { + status->status = tensorflow::errors::InvalidArgument( + "Passed in notification is a nullptr."); + + return; + } + if (notification->thread == nullptr) { + status->status = tensorflow::errors::InvalidArgument( + "Passed in notification didn't start a thread correctly. Cleaning up " + "this notification. Please re-execute the operation to get a new " + "notification."); + + delete notification; + return; + } + + notification->n.WaitForNotification(); + + status->status = notification->status->status; + + delete notification; +} + +void TF_MakeInternalErrorStatus(TF_Status* status, const char* errMsg) { status->status = tensorflow::errors::Internal(errMsg); } + +// This builder is used in the eager API to build a NodeDef. +struct TF_AttrBuilder : public tensorflow::AttrBuilder { + using tensorflow::AttrBuilder::AttrBuilder; + // The string buffers to make sure that any `attr_name` we pass into + // `builder->Set()` will outlive the subsequent + // `TF_AttrBuilderCheckCanRunOnDevice()` call(s) on the same `builder`. + std::set attr_names; +}; + +TF_AttrBuilder* TF_NewAttrBuilder(const char* op_name) { + return new TF_AttrBuilder(op_name); +} + +void TF_DeleteAttrBuilder(TF_AttrBuilder* builder) { delete builder; } + +void TF_AttrBuilderSetType(TF_AttrBuilder* builder, const char* attr_name, + TF_DataType value) { + auto iter = builder->attr_names.insert(attr_name).first; + builder->Set((*iter).c_str(), static_cast(value)); +} + +void TF_AttrBuilderSetTypeList(TF_AttrBuilder* builder, const char* attr_name, + const TF_DataType* values, int num_values) { + auto iter = builder->attr_names.insert(attr_name).first; + builder->Set( + (*iter).c_str(), + tensorflow::gtl::ArraySlice( + reinterpret_cast(values), num_values)); +} + +void TF_AttrBuilderCheckCanRunOnDevice(TF_AttrBuilder* builder, + const char* device_type, + TF_Status* status) { + status->status = tensorflow::FindKernelDef( + tensorflow::DeviceType(device_type), builder->BuildNodeDef(), + /* def = */ nullptr, /* kernel_class_name = */ nullptr); +} + +const char* TF_GetNumberAttrForOpListInput(const char* op_name, int input_index, + TF_Status* status) { + const tensorflow::OpDef* op_def = nullptr; + status->status = + tensorflow::OpRegistry::Global()->LookUpOpDef(op_name, &op_def); + if (!status->status.ok()) return nullptr; + + if (input_index >= op_def->input_arg_size() || input_index < 0) { + status->status = tensorflow::errors::InvalidArgument( + input_index, " out of range for ", op_name); + return nullptr; + } + + const tensorflow::OpDef_ArgDef& input_arg = op_def->input_arg()[input_index]; + + if (input_arg.number_attr().empty()) { + status->status = tensorflow::errors::NotFound( + op_name, " does not have number_attr() defined."); + return nullptr; + } + + // The returned string is owned by OpRegistry, so liveness is not a concern. + return input_arg.number_attr().c_str(); +} + +int TF_OpIsStateful(const char* op_type, TF_Status* status) { + const tensorflow::OpRegistrationData* op_reg_data; + status->status = + tensorflow::OpRegistry::Global()->LookUp(op_type, &op_reg_data); + if (!status->status.ok()) { + return 0; + } + return op_reg_data->op_def.is_stateful(); +} + +void TF_InitMain(const char* usage, int* argc, char*** argv) { + tensorflow::port::InitMain(usage, argc, argv); +} + +int TF_PickUnusedPortOrDie() { + return tensorflow::internal::PickUnusedPortOrDie(); +} + +TFE_TensorHandle* TFE_NewTensorHandleFromScalar(TF_DataType dtype_arg, + void* data, size_t len) { + auto dtype = static_cast(dtype_arg); + DCHECK(tensorflow::DataTypeCanUseMemcpy(dtype)); + + tensorflow::Tensor tensor(dtype, tensorflow::TensorShape({})); + std::memcpy(tensorflow::TensorCApi::Buffer(tensor)->data(), data, len); + return new TFE_TensorHandle(tensor, nullptr, nullptr); +} + +namespace { +tensorflow::Status EnableCollectiveOps(const tensorflow::ServerDef& server_def, + TFE_Context* ctx) { + // We don't use the TF_RETURN_IF_ERROR macro directly since that destroys the + // server object (which currently CHECK-fails) and we miss the error, instead, + // we log the error, and then return to allow the user to see the error + // message. +#define LOG_AND_RETURN_IF_ERROR(...) \ + do { \ + const ::tensorflow::Status _status = (__VA_ARGS__); \ + if (TF_PREDICT_FALSE(!_status.ok())) { \ + LOG(ERROR) << _status.error_message(); \ + return _status; \ + } \ + } while (0); + + std::unique_ptr server; + LOG_AND_RETURN_IF_ERROR(tensorflow::NewServer(server_def, &server)); + + tensorflow::GrpcServer* grpc_server = + dynamic_cast(server.get()); + if (grpc_server == nullptr) { + LOG_AND_RETURN_IF_ERROR(tensorflow::errors::Internal( + "Currently, TFE_NewContext only supports tensorflow::GrpcServer.")); + } + + LOG_AND_RETURN_IF_ERROR(grpc_server->Start()); + + LOG_AND_RETURN_IF_ERROR(ctx->context.StoreCollectiveOpsServer( + std::move(server), grpc_server->worker_env()->device_mgr, + grpc_server->worker_env()->collective_executor_mgr)); + + return tensorflow::Status::OK(); +#undef LOG_AND_RETURN_IF_ERROR +} +} // namespace + +// Set server_def on the context, possibly updating it. +TF_CAPI_EXPORT extern void TFE_EnableCollectiveOps(TFE_Context* ctx, + const void* proto, + size_t proto_len, + TF_Status* status) { + tensorflow::ServerDef server_def; + if (!server_def.ParseFromArray(proto, proto_len)) { + status->status = tensorflow::errors::InvalidArgument( + "Invalid tensorflow.ServerDef protocol buffer"); + return; + } + 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; +} + +void 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; + TF_AddInput(desc, symbolic_input); + } + + VLOG(1) << "Adding attrs."; + // TODO(hongm): add attrs + + auto* graph_op = TF_FinishOperation(desc, status); + if (!status->status.ok()) return; + + 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); + } +} + +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 d98d532e32e891e21f5b7ba360c74c3256fb1947..48ea0ec1ed78a071b7bf7c858881d943a3ff3acd 100644 --- a/tensorflow/c/c_api_experimental.h +++ b/tensorflow/c/c_api_experimental.h @@ -67,9 +67,10 @@ TF_CAPI_EXPORT extern void TF_EnableXLACompilation(TF_SessionOptions* options, // a) ConfigProto.optimizer_options.global_jit_level is set to to ON_1 if // `enable_xla_compilation` is non-zero, and OFF otherwise. // b) ConfigProto.gpu_options.allow_growth is set to `gpu_memory_allow_growth`. +// c) ConfigProto.device_count is set to `num_cpu_devices`. TF_CAPI_EXPORT extern TF_Buffer* TF_CreateConfig( - unsigned char enable_xla_compilation, - unsigned char gpu_memory_allow_growth); + unsigned char enable_xla_compilation, unsigned char gpu_memory_allow_growth, + unsigned int num_cpu_devices); // Create a serialized tensorflow.RunOptions proto, where RunOptions.trace_level // is set to FULL_TRACE if `enable_full_trace` is non-zero, and NO_TRACE @@ -83,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. @@ -180,9 +190,131 @@ 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 +// notification that can be waited upon. This always executes the kernel in a +// new thread. +// 1. `retvals` and `num_retvals` can only be consumed after +// `TFE_ExecuteOp` returns successfully. They shouldn't be used +// if the return is unsuccessful +// 2. These new APIs cannot be used together with the TFE context level async +// support. +TF_CAPI_EXPORT extern TFE_ExecuteOpNotification* TFE_ExecuteOpInNewThread( + TFE_Op* op, TFE_TensorHandle** retvals, int* num_retvals, + TF_Status* status); + +// Waits to complete the op execution, and cleans up the notification. +// Errors reported by op execution are set in `status`. +TF_CAPI_EXPORT extern void TFE_ExecuteOpNotificationWaitAndDelete( + TFE_ExecuteOpNotification* notification, TF_Status* status); + TF_CAPI_EXPORT extern void TF_MakeInternalErrorStatus(TF_Status* status, const char* errMsg); +// TF_NewAttrBuilder() returns an object that you can set attributes on as +// though it were an op. This allows querying properties of that op for +// type-checking purposes like if the op will run on a particular device type. +typedef struct TF_AttrBuilder TF_AttrBuilder; +TF_CAPI_EXPORT extern TF_AttrBuilder* TF_NewAttrBuilder(const char* op_name); +TF_CAPI_EXPORT extern void TF_DeleteAttrBuilder(TF_AttrBuilder* builder); +TF_CAPI_EXPORT extern void TF_AttrBuilderSetType(TF_AttrBuilder* builder, + const char* attr_name, + TF_DataType value); +TF_CAPI_EXPORT extern void TF_AttrBuilderSetTypeList(TF_AttrBuilder* builder, + const char* attr_name, + const TF_DataType* values, + int num_values); + +// Checks the tensorflow::NodeDef built via the methods above to see if it can +// run on device_type. +TF_CAPI_EXPORT extern void TF_AttrBuilderCheckCanRunOnDevice( + TF_AttrBuilder* builder, const char* device_type, TF_Status* status); + +// For argument number input_index, fetch the corresponding number_attr that +// needs to be updated with the argument length of the input list. +// Returns nullptr if there is any problem like op_name is not found, or the +// argument does not support this attribute type. +TF_CAPI_EXPORT extern const char* TF_GetNumberAttrForOpListInput( + const char* op_name, int input_index, TF_Status* status); + +// Returns 1 if the op is stateful, 0 otherwise. The return value is undefined +// if the status is not ok. +TF_CAPI_EXPORT extern int TF_OpIsStateful(const char* op_type, + TF_Status* status); + +// Platform specific initialization routine. Very few platforms actually require +// this to be called. +TF_CAPI_EXPORT void TF_InitMain(const char* usage, int* argc, char*** argv); + +// Platform-specific implementation to return an unused port. (This should used +// in tests only.) +TF_CAPI_EXPORT int TF_PickUnusedPortOrDie(void); + +// Fast path method that makes constructing a single scalar tensor require less +// overhead and copies. +TF_CAPI_EXPORT extern TFE_TensorHandle* TFE_NewTensorHandleFromScalar( + TF_DataType dtype, void* scalar, size_t len); + +// Specify the server_def that enables collective ops. +// This is different to the above function in that it doesn't create remote +// contexts, and remotely executing ops is not possible. It just enables +// communication for collective ops. +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`. +TF_CAPI_EXPORT extern void 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 c6effd39697e0397278770b53e98508074f99862..4cfcf2ef3b2ccd9d8aedaf8efa4a31ac12d91c1b 100644 --- a/tensorflow/c/c_api_experimental_test.cc +++ b/tensorflow/c/c_api_experimental_test.cc @@ -14,7 +14,10 @@ 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" #include "tensorflow/core/lib/io/path.h" #include "tensorflow/core/platform/env.h" #include "tensorflow/core/platform/logging.h" @@ -162,5 +165,205 @@ protocol: "grpc" TF_DeleteStatus(status); } +TEST(CAPI_EXPERIMENTAL, IsStateful) { + std::unique_ptr status( + TF_NewStatus(), TF_DeleteStatus); + int assign = TF_OpIsStateful("AssignAddVariableOp", status.get()); + ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); + EXPECT_EQ(assign, 1); + int id = TF_OpIsStateful("Identity", status.get()); + ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); + EXPECT_EQ(id, 0); +} + +TEST(CAPI_EXPERIMENTAL, TFE_ExecuteOpInNewThreadTest_Simple) { + TF_Status* status = TF_NewStatus(); + TFE_ContextOptions* opts = TFE_NewContextOptions(); + TFE_Context* ctx = TFE_NewContext(opts, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_DeleteContextOptions(opts); + + TFE_TensorHandle* m = TestMatrixTensorHandle(); + + TFE_Op* matmul_op = MatMulOp(ctx, m, m); + + TFE_TensorHandle* retvals[1] = {nullptr}; + int num_retvals = 1; + + auto* r = + TFE_ExecuteOpInNewThread(matmul_op, &retvals[0], &num_retvals, status); + + TFE_ExecuteOpNotificationWaitAndDelete(r, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + + TF_Tensor* t = TFE_TensorHandleResolve(retvals[0], status); + 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]); + + TFE_DeleteOp(matmul_op); + TFE_DeleteTensorHandle(m); + + TFE_DeleteTensorHandle(retvals[0]); + TFE_DeleteContext(ctx); + TF_DeleteStatus(status); +} + +// Perform a send/recv test. Recv blocks, so they need to be executed +// asynchronously. +TEST(CAPI_EXPERIMENTAL, TFE_ExecuteOpInNewThreadTest_Blocking) { + TF_Status* status = TF_NewStatus(); + TFE_ContextOptions* opts = TFE_NewContextOptions(); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_Context* ctx = TFE_NewContext(opts, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_DeleteContextOptions(opts); + + // Returns a 2x2 float32 Tensor on the CPU, with data 1., 2., 3., 4. + TFE_TensorHandle* m = TestMatrixTensorHandle(); + + // Build a send op. + TFE_Op* send_op = TFE_NewOp(ctx, "_Send", status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_OpAddInput(send_op, m, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + + string tensor_name = "Tensor"; + TFE_OpSetAttrType(send_op, "T", TF_FLOAT); + TFE_OpSetAttrString(send_op, "tensor_name", tensor_name.c_str(), + tensor_name.size()); + string send_device = "/job:localhost/replica:0/task:0/device:CPU:0"; + TFE_OpSetAttrString(send_op, "send_device", send_device.c_str(), + send_device.size()); + TFE_OpSetAttrInt(send_op, "send_device_incarnation", 1234); + string recv_device = "/job:localhost/replica:0/task:0/device:CPU:0"; + TFE_OpSetAttrString(send_op, "recv_device", recv_device.c_str(), + recv_device.size()); + TFE_OpSetAttrBool(send_op, "client_terminated", true); + + // Build a recv op. + TFE_Op* recv_op = TFE_NewOp(ctx, "_Recv", status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + + TFE_OpSetAttrType(recv_op, "tensor_type", TF_FLOAT); + TFE_OpSetAttrString(recv_op, "tensor_name", tensor_name.c_str(), + tensor_name.size()); + TFE_OpSetAttrString(recv_op, "send_device", send_device.c_str(), + send_device.size()); + TFE_OpSetAttrInt(recv_op, "send_device_incarnation", 1234); + TFE_OpSetAttrString(recv_op, "recv_device", recv_device.c_str(), + recv_device.size()); + TFE_OpSetAttrBool(recv_op, "client_terminated", true); + + TFE_TensorHandle* send_retvals; + int send_num_retvals = 0; + auto* send_result = TFE_ExecuteOpInNewThread(send_op, &send_retvals, + &send_num_retvals, status); + + TFE_TensorHandle* recv_retvals[1] = {nullptr}; + int recv_num_retvals = 1; + auto* recv_result = TFE_ExecuteOpInNewThread(recv_op, &recv_retvals[0], + &recv_num_retvals, status); + + TFE_ExecuteOpNotificationWaitAndDelete(send_result, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_ExecuteOpNotificationWaitAndDelete(recv_result, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + + TF_Tensor* t = TFE_TensorHandleResolve(recv_retvals[0], status); + 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(1, product[0]); + EXPECT_EQ(2, product[1]); + EXPECT_EQ(3, product[2]); + EXPECT_EQ(4, product[3]); + + TFE_DeleteOp(send_op); + TFE_DeleteOp(recv_op); + TFE_DeleteTensorHandle(m); + + TFE_DeleteTensorHandle(recv_retvals[0]); + TFE_DeleteContext(ctx); + 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); +} + +TEST(CAPI_EXPERIMENTAL, DebugPrintAndSymbolicExecution) { + TF_Status* status = TF_NewStatus(); + TFE_ContextOptions* opts = TFE_NewContextOptions(); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_Context* ctx = TFE_NewContext(opts, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_DeleteContextOptions(opts); + + TFE_TensorHandle* m = TestMatrixTensorHandle(); + TFE_Op* op = MatMulOp(ctx, m, m); + + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_OpPrintDebugString(op); + + auto* graph = TF_NewGraph(); + auto* trace_ctx = TFE_NewTraceContext(graph); + 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_DeleteTraceContext(trace_ctx); + TF_DeleteGraph(graph); + + TFE_DeleteTensorHandle(m); + TFE_DeleteOp(op); + TFE_DeleteContext(ctx); + TF_DeleteStatus(status); +} + } // namespace } // namespace tensorflow diff --git a/tensorflow/c/c_api_function.cc b/tensorflow/c/c_api_function.cc index f68f8a3e90a971b5e4a024feaf26ba498afc48da..28b9f8df9c873ee394eb6a241dd9ac06ba6c8796 100644 --- a/tensorflow/c/c_api_function.cc +++ b/tensorflow/c/c_api_function.cc @@ -392,26 +392,26 @@ Status ProcessInputs( EXCLUSIVE_LOCKS_REQUIRED(fn_body->mu) { input_tensors->reserve(ninputs); for (int i = 0; i < ninputs; ++i) { - const Node& node = inputs[i].oper->node; + Node* node = &inputs[i].oper->node; int idx = inputs[i].index; TF_RETURN_WITH_CONTEXT_IF_ERROR( - fn_body->graph.IsValidOutputTensor(&node, idx), + fn_body->graph.IsValidOutputTensor(node, idx), "Encountered while processing input ", i, " into function '", fn_name, "'"); - TF_RETURN_WITH_CONTEXT_IF_ERROR(ValidateNonRefOutput(&node, idx), + TF_RETURN_WITH_CONTEXT_IF_ERROR(ValidateNonRefOutput(node, idx), "Encountered while processing input ", i, " into function '", fn_name, "'"); - input_tensors->emplace_back(&node, idx); + input_tensors->emplace_back(node, idx); - const auto& iter = input_nodes->find(&node); + const auto& iter = input_nodes->find(node); if (iter == input_nodes->end()) { - input_nodes->insert({&node, {idx}}); + input_nodes->insert({node, {idx}}); } else { auto& indices = iter->second; if (std::find(indices.begin(), indices.end(), idx) != indices.end()) { - return InvalidArgument("TF_Output ", node.name(), ":", idx, + return InvalidArgument("TF_Output ", node->name(), ":", idx, " appears more than once in the input list"); } indices.push_back(idx); @@ -428,16 +428,16 @@ Status ProcessOutputs(const TF_Graph* fn_body, const char* fn_name, EXCLUSIVE_LOCKS_REQUIRED(fn_body->mu) { output_tensors->reserve(noutputs); for (int i = 0; i < noutputs; ++i) { - const Node& node = outputs[i].oper->node; + Node* node = &outputs[i].oper->node; int idx = outputs[i].index; TF_RETURN_WITH_CONTEXT_IF_ERROR( - fn_body->graph.IsValidOutputTensor(&node, idx), + fn_body->graph.IsValidOutputTensor(node, idx), "Encountered while processing output ", i, " from function '", fn_name, "'"); - TF_RETURN_WITH_CONTEXT_IF_ERROR(ValidateNonRefOutput(&node, idx), + TF_RETURN_WITH_CONTEXT_IF_ERROR(ValidateNonRefOutput(node, idx), "Encountered while creating function '", fn_name, "'"); - output_tensors->emplace_back(&node, idx); + output_tensors->emplace_back(node, idx); } return Status::OK(); } diff --git a/tensorflow/c/c_api_internal.h b/tensorflow/c/c_api_internal.h index 95652a11378d6276b5ba6540a07baa15aa77cc1c..73283d775639b297857b2a50007dc7c28b1f39a3 100644 --- a/tensorflow/c/c_api_internal.h +++ b/tensorflow/c/c_api_internal.h @@ -25,6 +25,7 @@ limitations under the License. #include #ifndef __ANDROID__ +#include "tensorflow/core/distributed_runtime/server_lib.h" #include "tensorflow/core/framework/op_gen_lib.h" #endif #include "tensorflow/core/common_runtime/shape_refiner.h" @@ -179,6 +180,15 @@ struct TF_ApiDefMap { tensorflow::mutex lock; }; +#ifndef __ANDROID__ +struct TF_Server { + TF_Server(std::unique_ptr server); + + const tensorflow::string target; + std::unique_ptr server; +}; +#endif + namespace tensorflow { class TensorCApi { @@ -218,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 c4746b4990bc3bf80b749428f803056e552421c3..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); } @@ -187,15 +188,26 @@ TEST(CAPI, LibraryLoadFunctions) { // tf_cuda_cc_test() bazel rule and remove the next line. if (!GPUDeviceName().empty()) return; - // Load the library. - TF_Status* status = TF_NewStatus(); - TF_Library* lib = - TF_LoadLibrary("tensorflow/c/test_op.so", status); - TF_Code code = TF_GetCode(status); - string status_msg(TF_Message(status)); - TF_DeleteStatus(status); - ASSERT_EQ(TF_OK, code) << status_msg; +#if !defined(TENSORFLOW_NO_SHARED_OBJECTS) + { + // Load the library. + TF_Status* status = TF_NewStatus(); + TF_Library* lib = + TF_LoadLibrary("tensorflow/c/test_op1.so", status); + TF_Code code = TF_GetCode(status); + string status_msg(TF_Message(status)); + TF_DeleteStatus(status); + ASSERT_EQ(TF_OK, code) << status_msg; + // Test op list. + TF_Buffer op_list_buf = TF_GetOpList(lib); + tensorflow::OpList op_list; + EXPECT_TRUE(op_list.ParseFromArray(op_list_buf.data, op_list_buf.length)); + ASSERT_EQ(op_list.op_size(), 1); + EXPECT_EQ("TestCApi1", op_list.op(0).name()); + TF_DeleteLibraryHandle(lib); + } +#endif // !defined(TENSORFLOW_NO_SHARED_OBJECTS) { TF_Buffer* op_list_buffer = TF_GetAllOpList(); tensorflow::OpList op_list; @@ -210,19 +222,6 @@ TEST(CAPI, LibraryLoadFunctions) { EXPECT_TRUE(found); TF_DeleteBuffer(op_list_buffer); } - -#if !defined(TENSORFLOW_NO_SHARED_OBJECTS) - { - // Test op list. - TF_Buffer op_list_buf = TF_GetOpList(lib); - tensorflow::OpList op_list; - EXPECT_TRUE(op_list.ParseFromArray(op_list_buf.data, op_list_buf.length)); - ASSERT_EQ(op_list.op_size(), 1); - EXPECT_EQ("TestCApi", op_list.op(0).name()); - } -#endif // !defined(TENSORFLOW_NO_SHARED_OBJECTS) - - TF_DeleteLibraryHandle(lib); } void TestEncodeDecode(int line, const std::vector& data) { @@ -1469,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") @@ -2349,14 +2383,8 @@ TEST(TestApiDef, TestCreateApiDef) { // tf_cuda_cc_test() bazel rule and remove the next line. if (!GPUDeviceName().empty()) return; - TF_Status* status = TF_NewStatus(); - TF_Library* lib = - TF_LoadLibrary("tensorflow/c/test_op.so", status); - EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); - TF_DeleteStatus(status); - TF_Buffer* op_list_buf = TF_GetAllOpList(); - status = TF_NewStatus(); + TF_Status* status = TF_NewStatus(); auto* api_def_map = TF_NewApiDefMap(op_list_buf, status); EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); TF_DeleteStatus(status); @@ -2376,7 +2404,6 @@ TEST(TestApiDef, TestCreateApiDef) { TF_DeleteBuffer(api_def_buf); TF_DeleteApiDefMap(api_def_map); TF_DeleteBuffer(op_list_buf); - TF_DeleteLibraryHandle(lib); } TEST(TestApiDef, TestCreateApiDefWithOverwrites) { @@ -2384,14 +2411,8 @@ TEST(TestApiDef, TestCreateApiDefWithOverwrites) { // tf_cuda_cc_test() bazel rule and remove the next line. if (!GPUDeviceName().empty()) return; - TF_Status* status = TF_NewStatus(); - TF_Library* lib = - TF_LoadLibrary("tensorflow/c/test_op.so", status); - EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); - TF_DeleteStatus(status); - TF_Buffer* op_list_buf = TF_GetAllOpList(); - status = TF_NewStatus(); + TF_Status* status = TF_NewStatus(); auto* api_def_map = TF_NewApiDefMap(op_list_buf, status); EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); TF_DeleteStatus(status); @@ -2422,7 +2443,6 @@ TEST(TestApiDef, TestCreateApiDefWithOverwrites) { TF_DeleteBuffer(api_def_buf); TF_DeleteApiDefMap(api_def_map); TF_DeleteBuffer(op_list_buf); - TF_DeleteLibraryHandle(lib); } class DummyKernel : public tensorflow::OpKernel { diff --git a/tensorflow/c/c_test.c b/tensorflow/c/c_test.c new file mode 100644 index 0000000000000000000000000000000000000000..b86d8eb8e300e02a3871ecd5f424a82c521b18fc --- /dev/null +++ b/tensorflow/c/c_test.c @@ -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. +==============================================================================*/ + +#include +#include +#include +#include +#include +#include + +#include "tensorflow/c/c_api.h" +#include "tensorflow/c/c_api_experimental.h" +#include "tensorflow/c/env.h" +#include "tensorflow/c/kernels.h" + +// 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) { + TF_Tensor* input; + TF_Status* s = TF_NewStatus(); + TF_GetInput(ctx, 0, &input, s); + TF_DeleteTensor(input); + + TF_DataType type; + TF_OpKernelContext_GetAttrType(ctx, "foobar", &type, s); + + TF_DeleteStatus(s); + +} + +// Exercises tensorflow's C API. +int main(int argc, char** argv) { + TF_InitMain(argv[0], &argc, &argv); + + struct TF_StringStream* s = TF_GetLocalTempDirectories(); + const char* path; + + if (!TF_StringStreamNext(s, &path)) { + fprintf(stderr, "TF_GetLocalTempDirectories returned no results\n"); + return 1; + } + + char file_name[100]; + struct timeval t; + if (gettimeofday(&t, NULL)) { + perror("gettimeofday failed"); + return 1; + } + snprintf(file_name, sizeof(file_name), "test-%d-%ld.txt", getpid(), t.tv_sec); + + size_t length = 2 + strlen(path) + strlen(file_name); + char* full_path = malloc(length); + snprintf(full_path, length, "%s/%s", path, file_name); + + TF_WritableFileHandle* h; + TF_Status* status = TF_NewStatus(); + TF_NewWritableFile(full_path, &h, status); + if (TF_GetCode(status) != TF_OK) { + fprintf(stderr, "TF_NewWritableFile failed: %s\n", TF_Message(status)); + return 1; + } + fprintf(stderr, "wrote %s\n", full_path); + free(full_path); + TF_CloseWritableFile(h, status); + if (TF_GetCode(status) != TF_OK) { + fprintf(stderr, "TF_CloseWritableFile failed: %s\n", TF_Message(status)); + } + TF_StringStreamDone(s); + + TF_KernelBuilder* b = + TF_NewKernelBuilder("SomeOp", "SomeDevice", NULL, &compute, NULL); + TF_RegisterKernelBuilder("someKernel", b, status); + + TF_DeleteStatus(status); + return 0; +} diff --git a/tensorflow/c/eager/BUILD b/tensorflow/c/eager/BUILD index 3ee31a6a7ac641bbd3fc4c05568b61e433a1d523..04dfefa6da28429b73856d392d94fa402ecda580 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( @@ -50,6 +58,7 @@ tf_cuda_library( ], "//conditions:default": [], }) + [ + "@com_google_absl//absl/memory", "//tensorflow/core/common_runtime/eager:eager_operation", "//tensorflow/core/distributed_runtime/eager:eager_client", "//tensorflow/core/distributed_runtime/rpc/eager:grpc_eager_client", @@ -61,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:eager_profiler", "//tensorflow/core:gpu_runtime", ], ) @@ -69,7 +79,7 @@ tf_cuda_library( name = "c_api_internal", hdrs = ["c_api_internal.h"], visibility = [ - "//learning/deepmind/courier:__pkg__", + "//learning/deepmind/courier:__subpackages__", "//tensorflow:internal", ], deps = [ @@ -100,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:eager_profiler", ], ) @@ -143,6 +154,89 @@ tf_cuda_cc_test( "//tensorflow/core:test", "//tensorflow/core:test_main", "//tensorflow/core/distributed_runtime/rpc:grpc_server_lib", + "@com_google_absl//absl/strings", + ], +) + +tf_cuda_library( + name = "c_api_experimental", + srcs = [ + "c_api_experimental.cc", + ], + hdrs = ["c_api_experimental.h"], + copts = tf_copts() + tfe_xla_copts(), + visibility = ["//visibility:public"], + deps = select({ + "//tensorflow:android": [ + "//tensorflow/core:android_tensorflow_lib_lite", + ], + "//conditions:default": [ + ":c_api", + ":c_api_internal", + "//tensorflow/c:c_api", + "//tensorflow/c:c_api_internal", + "//tensorflow/core:core_cpu", + "//tensorflow/core/common_runtime/eager:attr_builder", + "//tensorflow/core/common_runtime/eager:context", + "//tensorflow/core/common_runtime/eager:eager_executor", + "//tensorflow/core/common_runtime/eager:execute", + "//tensorflow/core/common_runtime/eager:kernel_and_device", + "//tensorflow/core/common_runtime/eager:tensor_handle", + "//tensorflow/core/common_runtime/eager:copy_to_device_node", + "//tensorflow/core:core_cpu_internal", + "//tensorflow/core:framework", + "//tensorflow/core:framework_internal", + "//tensorflow/core:lib", + "//tensorflow/core:lib_internal", + "//tensorflow/core:protos_all_cc", + ], + }) + select({ + "//tensorflow:with_xla_support": [ + "//tensorflow/compiler/tf2xla:xla_compiler", + "//tensorflow/compiler/jit", + "//tensorflow/compiler/jit:xla_device", + ], + "//conditions:default": [], + }) + [ + "@com_google_absl//absl/memory", + "//tensorflow/core/common_runtime/eager:eager_operation", + "//tensorflow/core/distributed_runtime/eager:eager_client", + "//tensorflow/core/distributed_runtime/rpc/eager:grpc_eager_client", + "//tensorflow/core/distributed_runtime/rpc:grpc_channel", + "//tensorflow/core/distributed_runtime/rpc:grpc_server_lib", + "//tensorflow/core/distributed_runtime/rpc:grpc_worker_cache", + "//tensorflow/core/distributed_runtime/rpc:grpc_worker_service", + "//tensorflow/core/distributed_runtime/rpc:rpc_rendezvous_mgr", + "//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", ], ) diff --git a/tensorflow/c/eager/c_api.cc b/tensorflow/c/eager/c_api.cc index 3554ec0bf3202b54bfc38d67e51b89df19832302..af13f487af91594fedd4d5f77592682a6f98c34f 100755 --- a/tensorflow/c/eager/c_api.cc +++ b/tensorflow/c/eager/c_api.cc @@ -21,9 +21,11 @@ limitations under the License. #include #include +#include "absl/memory/memory.h" #include "tensorflow/c/c_api.h" #include "tensorflow/c/c_api_internal.h" #include "tensorflow/c/eager/c_api_internal.h" +#include "tensorflow/core/platform/host_info.h" #ifdef TENSORFLOW_EAGER_USE_XLA #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #endif // TENSORFLOW_EAGER_USE_XLA @@ -79,7 +81,7 @@ tensorflow::Status GetAllRemoteDevices( const std::vector& remote_workers, tensorflow::WorkerCacheInterface* worker_cache, std::unique_ptr* device_mgr) { - std::vector remote_devices; + std::vector> remote_devices; tensorflow::Status status; // TODO(nareshmodi) do this in parallel instead of serially. for (const string& remote_worker : remote_workers) { @@ -92,7 +94,7 @@ tensorflow::Status GetAllRemoteDevices( status = s; if (s.ok()) { for (tensorflow::Device* d : *devices) { - remote_devices.push_back(d); + remote_devices.emplace_back(d); } } n.Notify(); @@ -100,7 +102,7 @@ tensorflow::Status GetAllRemoteDevices( n.WaitForNotification(); } std::unique_ptr remote_device_mgr( - new tensorflow::DeviceMgr(remote_devices)); + new tensorflow::DeviceMgr(std::move(remote_devices))); TF_RETURN_IF_ERROR(status); @@ -261,13 +263,13 @@ TF_CAPI_EXPORT extern void TFE_ContextSetAsyncForThread(TFE_Context* ctx, void TFE_DeleteContextOptions(TFE_ContextOptions* options) { delete options; } TFE_Context* TFE_NewContext(const TFE_ContextOptions* opts, TF_Status* status) { - std::vector devices; + std::vector> devices; status->status = tensorflow::DeviceFactory::AddDevices( opts->session_options.options, "/job:localhost/replica:0/task:0", &devices); if (!status->status.ok()) return nullptr; std::unique_ptr device_mgr( - new tensorflow::DeviceMgr(devices)); + new tensorflow::DeviceMgr(std::move(devices))); tensorflow::Rendezvous* r = new tensorflow::IntraProcessRendezvous(device_mgr.get()); @@ -354,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(); } @@ -404,8 +408,19 @@ const char* TFE_TensorHandleDeviceName(TFE_TensorHandle* h, TF_Status* status) { "The passed in handle is a nullptr"); return nullptr; } - tensorflow::Device* d = nullptr; - status->status = h->handle->OpDevice(&d); + tensorflow::Device* d = h->handle->op_device(); + return (d == nullptr) ? "/job:localhost/replica:0/task:0/device:CPU:0" + : d->name().c_str(); +} + +const char* TFE_TensorHandleBackingDeviceName(TFE_TensorHandle* h, + TF_Status* status) { + if (h == nullptr || h->handle == nullptr) { + status->status = tensorflow::errors::InvalidArgument( + "The passed in handle is a nullptr"); + return nullptr; + } + tensorflow::Device* d = h->handle->device(); return (d == nullptr) ? "/job:localhost/replica:0/task:0/device:CPU:0" : d->name().c_str(); } @@ -430,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; } @@ -447,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) { @@ -459,13 +490,20 @@ TFE_Op* TFE_NewOp(TFE_Context* ctx, const char* op_or_function_name, TF_Status* status) { const char* name = op_or_function_name; // Shorthand const tensorflow::AttrTypeMap* types; - status->status = tensorflow::AttrTypeMapForOp(name, &types); - if (status->status.ok()) return new TFE_Op(ctx, name, types); - if (TF_GetCode(status) == TF_NOT_FOUND) { - if (ctx->context.FindFunctionByName(name)) { - status->status = tensorflow::Status::OK(); - return new TFE_Op(ctx, name, nullptr); + bool is_function = false; + status->status = tensorflow::AttrTypeMapForOp(name, &types, &is_function); + if (status->status.ok()) { + if (is_function && !ctx->context.FindFunctionByName(name)) { + status->status = tensorflow::errors::NotFound( + "'", name, + "' is neither a type of a primitive operation nor a name " + "of a function registered in binary running on ", + tensorflow::port::Hostname(), + ". Make sure the operation or function is " + "registered in the binary running in this process."); + return nullptr; } + return new TFE_Op(ctx, name, is_function, types); } return nullptr; } @@ -498,12 +536,6 @@ void TFE_OpAddInput(TFE_Op* op, TFE_TensorHandle* h, TF_Status* status) { TF_AttrType TFE_OpGetAttrType(TFE_Op* op, const char* attr_name, unsigned char* is_list, TF_Status* status) { TF_AttrType ret; - if (op->operation.is_function()) { - status->status = tensorflow::errors::Unimplemented( - "TODO(apassos): Support for attributes for TensorFlow functions is not " - "ready yet."); - return TF_ATTR_INT; // The compiler requires that we return something. - } status->status = tensorflow::AttrTypeByName(*op->operation.AttrTypes(), attr_name, &ret, is_list); return ret; @@ -682,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 = @@ -724,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); } @@ -760,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 b2454d872207e26feb3764671474a5d87c01f84d..044dfb7415b027b707af05a197fdb41fe1f6d2e5 100755 --- a/tensorflow/c/eager/c_api.h +++ b/tensorflow/c/eager/c_api.h @@ -48,7 +48,7 @@ extern "C" { typedef struct TFE_ContextOptions TFE_ContextOptions; // Return a new options object. -TF_CAPI_EXPORT extern TFE_ContextOptions* TFE_NewContextOptions(); +TF_CAPI_EXPORT extern TFE_ContextOptions* TFE_NewContextOptions(void); // Set the config in TF_ContextOptions.options. // config should be a serialized tensorflow.ConfigProto proto. @@ -169,10 +169,21 @@ TF_CAPI_EXPORT extern int64_t TFE_TensorHandleNumElements(TFE_TensorHandle* h, TF_CAPI_EXPORT extern int64_t TFE_TensorHandleDim(TFE_TensorHandle* h, int dim_index, TF_Status* status); -// This function will block till the operation that produces `h` has completed. + +// Returns the device of the operation that produced `h`. If `h` was produced by +// a copy, returns the destination device of the copy. Note that the returned +// device name is not always the device holding the tensor handle's memory. If +// you want the latter, use TFE_TensorHandleBackingDeviceName. This function +// will block till the operation that produces `h` has completed. TF_CAPI_EXPORT extern const char* TFE_TensorHandleDeviceName( TFE_TensorHandle* h, TF_Status* status); +// Returns the name of the device in whose memory `h` resides. +// +// This function will block till the operation that produces `h` has completed. +TF_CAPI_EXPORT extern const char* TFE_TensorHandleBackingDeviceName( + TFE_TensorHandle* h, TF_Status* status); + // Return a pointer to a new TFE_TensorHandle that shares the underlying tensor // with `h`. On success, `status` is set to OK. On failure, `status` reflects // the error and a nullptr is returned. @@ -382,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 5006b76f1981d068e99a2c081115ebb3a66d8c7f..ffcd5ace0b98597363abe63201bf6c328a03212f 100644 --- a/tensorflow/c/eager/c_api_debug.cc +++ b/tensorflow/c/eager/c_api_debug.cc @@ -57,13 +57,9 @@ TF_CAPI_EXPORT extern TFE_TensorDebugInfo* TFE_TensorHandleTensorDebugInfo( return nullptr; } - tensorflow::Device* device; - status->status = handle->handle->Device(&device); - if (!status->status.ok()) { - return nullptr; - } - #ifdef TENSORFLOW_EAGER_USE_XLA + tensorflow::Device* device = handle->handle->device(); + // If tensor resides on an XLA device, use XLA device's PaddedShapeFn. tensorflow::XlaDevice* xla_device = dynamic_cast(device); @@ -87,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 @@ -103,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 new file mode 100644 index 0000000000000000000000000000000000000000..dab17505643e791e6294a64247898ae23769a055 --- /dev/null +++ b/tensorflow/c/eager/c_api_experimental.cc @@ -0,0 +1,52 @@ +/* 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/c/eager/c_api_experimental.h" + +#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_Context* ctx) { + return new TFE_Profiler(ctx); +} + +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); + }; +} + +void TFE_StartProfilerServer(TFE_Context* ctx, int port) { + auto server_thread = tensorflow::StartProfilerServer(&ctx->context, port); + ctx->context.AddChildThread(std::move(server_thread)); +} diff --git a/tensorflow/c/eager/c_api_experimental.h b/tensorflow/c/eager/c_api_experimental.h new file mode 100644 index 0000000000000000000000000000000000000000..8c85d0e51695fde09cf0e2bb3930f9173e6cfb54 --- /dev/null +++ b/tensorflow/c/eager/c_api_experimental.h @@ -0,0 +1,58 @@ +/* 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_C_EAGER_C_API_EXPERIMENTAL_H_ +#define TENSORFLOW_C_EAGER_C_API_EXPERIMENTAL_H_ + +#include "tensorflow/c/c_api.h" +#include "tensorflow/c/eager/c_api.h" + +#ifdef __cplusplus +extern "C" { +#endif + +TF_CAPI_EXPORT extern void TFE_OpConsumeInput(TFE_Op* op, TFE_TensorHandle* h, + TF_Status* status); + +// 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_Context* ctx); +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); + +// Start a profiler grpc server which listens to specified port. It will start +// the server on its own thread. It can be shutdown by destructing TFE_Context. +// 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_Context* ctx, int port); + +#ifdef __cplusplus +} /* end extern "C" */ +#endif + +#endif // TENSORFLOW_C_EAGER_C_API_EXPERIMENTAL_H_ 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..af55fee66e8708e39626da3b10b6dd2f73af92bb --- /dev/null +++ b/tensorflow/c/eager/c_api_experimental_test.cc @@ -0,0 +1,104 @@ +/* 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_Profiler* profiler = TFE_NewProfiler(ctx); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_DeleteContextOptions(opts); + + 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); } + +} // namespace +} // namespace tensorflow diff --git a/tensorflow/c/eager/c_api_internal.h b/tensorflow/c/eager/c_api_internal.h index 104d52430cf7aa14d4d2a335a1b96e667f21ce87..3b9e681194b7cebc61d9028525d200c692bbd529 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/eager_profiler.h" #include "tensorflow/core/public/version.h" struct TFE_ContextOptions { @@ -79,13 +80,15 @@ struct TFE_TensorHandle { tensorflow::Device* op_device) : handle(new tensorflow::TensorHandle(t, d, op_device, nullptr)) {} - TFE_TensorHandle(tensorflow::uint64 node_id, tensorflow::DataType dtype, - tensorflow::EagerContext* ctx) - : handle(new tensorflow::TensorHandle(node_id, dtype, ctx)) {} - 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 { @@ -97,14 +100,20 @@ struct TFE_TensorDebugInfo { }; struct TFE_Op { - // t is NULL iff the TFE_Op corresponds to a TensorFlow function instead of a - // primitive operation. - TFE_Op(TFE_Context* ctx, const char* op, const tensorflow::AttrTypeMap* t) - : operation(&ctx->context, op, t) {} + TFE_Op(TFE_Context* ctx, const char* op, bool is_function, + const tensorflow::AttrTypeMap* t) + : operation(&ctx->context, op, is_function, t) {} tensorflow::EagerOperation operation; }; +struct TFE_Profiler { + TFE_Profiler(TFE_Context* ctx) + : profiler(tensorflow::EagerProfiler::Create(&ctx->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, @@ -112,4 +121,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 55331022b9dbd0696928fa44430f340f371432ac..3d1ca4fb4b561a03ea9d879b1876fb1fd08a3139 100644 --- a/tensorflow/c/eager/c_api_test.cc +++ b/tensorflow/c/eager/c_api_test.cc @@ -16,6 +16,7 @@ limitations under the License. #include "tensorflow/c/eager/c_api.h" #include +#include "absl/strings/match.h" #include "tensorflow/c/eager/c_api_test_util.h" #include "tensorflow/core/distributed_runtime/rpc/grpc_server_lib.h" #include "tensorflow/core/framework/function.pb.h" @@ -174,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)); @@ -589,9 +585,22 @@ void TensorHandleCopyBetweenTwoGPUDevices(bool async) { TF_DeviceList* devices = TFE_ContextListDevices(ctx, status.get()); ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); const int num_devices = TF_DeviceListCount(devices); + bool has_gpu0 = false; + bool has_gpu1 = false; + for (int i = 0; i < num_devices; ++i) { + const char* dev = TF_DeviceListName(devices, i, status.get()); + ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); + string device_name(dev); + if (device_name.find("GPU:0") != string::npos) { + has_gpu0 = true; + } + if (device_name.find("GPU:1") != string::npos) { + has_gpu1 = true; + } + } const char* kCPUDevice = "CPU:0"; - if (num_devices < 3) { + if (!has_gpu0 || !has_gpu1) { TF_DeleteDeviceList(devices); TF_DeleteTensor(t); TFE_DeleteTensorHandle(hcpu); @@ -781,6 +790,14 @@ TEST(CAPI, TensorHandleNullptr) { TF_SetStatus(status.get(), TF_OK, ""); + device_name = TFE_TensorHandleBackingDeviceName(h, status.get()); + ASSERT_EQ(TF_INVALID_ARGUMENT, TF_GetCode(status.get())); + ASSERT_EQ(device_name, nullptr); + ASSERT_EQ("The passed in handle is a nullptr", + string(TF_Message(status.get()))); + + TF_SetStatus(status.get(), TF_OK, ""); + int num_dims = TFE_TensorHandleNumDims(h, status.get()); ASSERT_EQ(TF_INVALID_ARGUMENT, TF_GetCode(status.get())); ASSERT_EQ(num_dims, -1); @@ -796,6 +813,62 @@ TEST(CAPI, TensorHandleNullptr) { string(TF_Message(status.get()))); } +TEST(CAPI, TensorHandleDevices) { + std::unique_ptr status( + TF_NewStatus(), TF_DeleteStatus); + TFE_ContextOptions* opts = TFE_NewContextOptions(); + TFE_Context* ctx = TFE_NewContext(opts, status.get()); + TFE_DeleteContextOptions(opts); + ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); + + TFE_TensorHandle* hcpu = TestMatrixTensorHandle(); + const char* device_name = TFE_TensorHandleDeviceName(hcpu, status.get()); + ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); + ASSERT_TRUE(absl::StrContains(device_name, "CPU:0")) << device_name; + const char* backing_device_name = + TFE_TensorHandleBackingDeviceName(hcpu, status.get()); + ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); + ASSERT_TRUE(absl::StrContains(backing_device_name, "CPU:0")) + << backing_device_name; + + // Disable the test if no GPU is present. + string gpu_device_name; + if (GetDeviceName(ctx, &gpu_device_name, "GPU")) { + TFE_TensorHandle* hgpu = TFE_TensorHandleCopyToDevice( + hcpu, ctx, gpu_device_name.c_str(), status.get()); + ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get()); + + TFE_Op* shape_op = ShapeOp(ctx, hgpu); + TFE_OpSetDevice(shape_op, gpu_device_name.c_str(), status.get()); + ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get()); + TFE_TensorHandle* retvals[1]; + int num_retvals = 1; + TFE_Execute(shape_op, &retvals[0], &num_retvals, status.get()); + ASSERT_TRUE(TF_GetCode(status.get()) == TF_OK) << TF_Message(status.get()); + + // .device of shape is GPU since the op is executed on GPU + device_name = TFE_TensorHandleDeviceName(retvals[0], status.get()); + ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); + ASSERT_TRUE(absl::StrContains(device_name, "GPU:0")) << device_name; + + // .backing_device of shape is CPU since the tensor is backed by CPU + backing_device_name = + TFE_TensorHandleBackingDeviceName(retvals[0], status.get()); + ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); + ASSERT_TRUE(absl::StrContains(backing_device_name, "CPU:0")) + << backing_device_name; + + TFE_DeleteOp(shape_op); + TFE_DeleteTensorHandle(retvals[0]); + TFE_DeleteTensorHandle(hgpu); + } + + TFE_DeleteTensorHandle(hcpu); + TFE_ContextAsyncWait(ctx, status.get()); + EXPECT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); + TFE_DeleteContext(ctx); +} + void Execute_MatMul_CPU(bool async) { TF_Status* status = TF_NewStatus(); TFE_ContextOptions* opts = TFE_NewContextOptions(); diff --git a/tensorflow/c/eager/c_api_test_util.cc b/tensorflow/c/eager/c_api_test_util.cc index 008f088c2dcdd7d9114103516a4702e47a55c6de..bd38127d50c171af801dd1b937acefdba491b4a6 100644 --- a/tensorflow/c/eager/c_api_test_util.cc +++ b/tensorflow/c/eager/c_api_test_util.cc @@ -104,6 +104,19 @@ TFE_Op* MatMulOp(TFE_Context* ctx, TFE_TensorHandle* a, TFE_TensorHandle* b) { return op; } +TFE_Op* ShapeOp(TFE_Context* ctx, TFE_TensorHandle* a) { + TF_Status* status = TF_NewStatus(); + + TFE_Op* op = TFE_NewOp(ctx, "Shape", status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_OpAddInput(op, a, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TF_DeleteStatus(status); + TFE_OpSetAttrType(op, "T", TFE_TensorHandleDataType(a)); + + return op; +} + TFE_TensorHandle* TestAxisTensorHandle() { int64_t dims[] = {1}; int data[] = {1}; diff --git a/tensorflow/c/eager/c_api_test_util.h b/tensorflow/c/eager/c_api_test_util.h index 474cae67c89249af3a62707f0db00ba458ca8f31..75ef9459e93b4f2ed471c423a34565594efc1714 100644 --- a/tensorflow/c/eager/c_api_test_util.h +++ b/tensorflow/c/eager/c_api_test_util.h @@ -37,6 +37,9 @@ TFE_TensorHandle* TestMatrixTensorHandle3X2(); // Return a matmul op multiplying `a` by `b`. TFE_Op* MatMulOp(TFE_Context* ctx, TFE_TensorHandle* a, TFE_TensorHandle* b); +// Return a shape op fetching the shape of `a`. +TFE_Op* ShapeOp(TFE_Context* ctx, TFE_TensorHandle* a); + // Return an 1-D INT32 tensor containing a single value 1. TFE_TensorHandle* TestAxisTensorHandle(); diff --git a/tensorflow/c/eager/tape.h b/tensorflow/c/eager/tape.h index 5ba55a203ff70cc64c07e96b5a869a1f11c9334e..5c11f51e8749de84547ae873f5f55ebd42bc4b3d 100644 --- a/tensorflow/c/eager/tape.h +++ b/tensorflow/c/eager/tape.h @@ -141,8 +141,9 @@ class GradientTape { // null. The result is populated with one tensor per target element. Status ComputeGradient( const VSpace& vspace, - gtl::ArraySlice target_tensor_ids, - gtl::ArraySlice source_tensor_id, + const gtl::ArraySlice target_tensor_ids, + const gtl::ArraySlice source_tensor_ids, + const gtl::FlatMap sources_that_are_targets, gtl::ArraySlice output_gradients, std::vector* result); @@ -396,6 +397,7 @@ template Status InitialGradients( const VSpace& vspace, gtl::ArraySlice target_tensor_ids, + gtl::FlatMap sources_that_are_targets, gtl::ArraySlice output_gradients, const TensorTape& tensor_tape, const OpTape& op_tape, gtl::FlatMap>* result) { @@ -425,8 +427,13 @@ Status InitialGradients( "none of operations outputs match expected tensor"); } } else { - // No record of the target tensor found on the tape, so no gradient - // needs to be computed from it. Do nothing. + // This target tensor was not generated by any operation recorded on + // the tape, so no gradient needs to be computed from it unless this + // target is also a source. + auto source_tensor = sources_that_are_targets.find(id); + if (source_tensor != sources_that_are_targets.end()) { + (*result)[id].push_back(vspace.Ones(source_tensor->second)); + } } } else { (*result)[id].push_back(output_gradients[i]); @@ -467,8 +474,9 @@ constexpr int kMinAggregateBytes = 128 * 1024 * 1024; template Status GradientTape::ComputeGradient( const VSpace& vspace, - gtl::ArraySlice target_tensor_ids, - gtl::ArraySlice source_tensor_ids, + const gtl::ArraySlice target_tensor_ids, + const gtl::ArraySlice source_tensor_ids, + const gtl::FlatMap sources_that_are_targets, gtl::ArraySlice output_gradients, std::vector* result) { gtl::FlatSet sources_set(source_tensor_ids.begin(), @@ -478,7 +486,8 @@ Status GradientTape::ComputeGradient( std::vector op_stack = InitialStack(state.op_tape, state.op_missing_tensor); gtl::FlatMap> gradients; - Status s = InitialGradients(vspace, target_tensor_ids, output_gradients, + Status s = InitialGradients(vspace, target_tensor_ids, + sources_that_are_targets, output_gradients, tensor_tape_, state.op_tape, &gradients); auto cleanup = [this, &state]() { if (!persistent_) { diff --git a/tensorflow/c/env.cc b/tensorflow/c/env.cc new file mode 100644 index 0000000000000000000000000000000000000000..1c35ff9001d0ee1ab0fbae9e1bcc07116fab1065 --- /dev/null +++ b/tensorflow/c/env.cc @@ -0,0 +1,183 @@ +/* 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/c/env.h" + +#include "tensorflow/c/c_api_internal.h" +#include "tensorflow/c/tf_status_helper.h" +#include "tensorflow/core/platform/env.h" +#include "tensorflow/core/platform/types.h" + +struct TF_StringStream { + std::vector<::tensorflow::string>* list; + size_t position; +}; + +void TF_CreateDir(const char* dirname, TF_Status* status) { + TF_SetStatus(status, TF_OK, ""); + ::tensorflow::Set_TF_Status_from_Status( + status, ::tensorflow::Env::Default()->CreateDir(dirname)); +} + +void TF_DeleteDir(const char* dirname, TF_Status* status) { + TF_SetStatus(status, TF_OK, ""); + ::tensorflow::Set_TF_Status_from_Status( + status, ::tensorflow::Env::Default()->DeleteDir(dirname)); +} + +void TF_DeleteRecursively(const char* dirname, int64_t* undeleted_file_count, + int64_t* undeleted_dir_count, TF_Status* status) { + ::tensorflow::int64 f, d; + + TF_SetStatus(status, TF_OK, ""); + ::tensorflow::Set_TF_Status_from_Status( + status, ::tensorflow::Env::Default()->DeleteRecursively(dirname, &f, &d)); + *undeleted_file_count = f; + *undeleted_dir_count = d; +} + +void TF_FileStat(const char* filename, TF_FileStatistics* stats, + TF_Status* status) { + ::tensorflow::FileStatistics cc_stats; + TF_SetStatus(status, TF_OK, ""); + ::tensorflow::Status s = + ::tensorflow::Env::Default()->Stat(filename, &cc_stats); + ::tensorflow::Set_TF_Status_from_Status(status, s); + if (s.ok()) { + stats->length = cc_stats.length; + stats->mtime_nsec = cc_stats.mtime_nsec; + stats->is_directory = cc_stats.is_directory; + } +} + +void TF_NewWritableFile(const char* filename, TF_WritableFileHandle** handle, + TF_Status* status) { + std::unique_ptr<::tensorflow::WritableFile> f; + TF_SetStatus(status, TF_OK, ""); + ::tensorflow::Status s = + ::tensorflow::Env::Default()->NewWritableFile(filename, &f); + ::tensorflow::Set_TF_Status_from_Status(status, s); + + if (s.ok()) { + *handle = reinterpret_cast(f.release()); + } +} + +void TF_CloseWritableFile(TF_WritableFileHandle* handle, TF_Status* status) { + auto* cc_file = reinterpret_cast<::tensorflow::WritableFile*>(handle); + TF_SetStatus(status, TF_OK, ""); + ::tensorflow::Set_TF_Status_from_Status(status, cc_file->Close()); + delete cc_file; +} + +void TF_SyncWritableFile(TF_WritableFileHandle* handle, TF_Status* status) { + auto* cc_file = reinterpret_cast<::tensorflow::WritableFile*>(handle); + TF_SetStatus(status, TF_OK, ""); + ::tensorflow::Set_TF_Status_from_Status(status, cc_file->Sync()); +} + +void TF_FlushWritableFile(TF_WritableFileHandle* handle, TF_Status* status) { + auto* cc_file = reinterpret_cast<::tensorflow::WritableFile*>(handle); + TF_SetStatus(status, TF_OK, ""); + ::tensorflow::Set_TF_Status_from_Status(status, cc_file->Flush()); +} + +void TF_AppendWritableFile(TF_WritableFileHandle* handle, const char* data, + size_t length, TF_Status* status) { + auto* cc_file = reinterpret_cast<::tensorflow::WritableFile*>(handle); + TF_SetStatus(status, TF_OK, ""); + ::tensorflow::Set_TF_Status_from_Status( + status, cc_file->Append(::tensorflow::StringPiece{data, length})); +} + +void TF_DeleteFile(const char* filename, TF_Status* status) { + TF_SetStatus(status, TF_OK, ""); + ::tensorflow::Set_TF_Status_from_Status( + status, ::tensorflow::Env::Default()->DeleteFile(filename)); +} + +bool TF_StringStreamNext(TF_StringStream* list, const char** result) { + if (list->position >= list->list->size()) { + *result = nullptr; + return false; + } + + *result = list->list->at(list->position++).c_str(); + return true; +} + +void TF_StringStreamDone(TF_StringStream* list) { + delete list->list; + delete list; +} +TF_StringStream* TF_GetChildren(const char* dirname, TF_Status* status) { + auto* children = new std::vector<::tensorflow::string>; + + TF_SetStatus(status, TF_OK, ""); + ::tensorflow::Set_TF_Status_from_Status( + status, ::tensorflow::Env::Default()->GetChildren(dirname, children)); + + auto* list = new TF_StringStream; + list->list = children; + list->position = 0; + return list; +} + +TF_StringStream* TF_GetLocalTempDirectories() { + auto* tmpdirs = new std::vector<::tensorflow::string>; + + ::tensorflow::Env::Default()->GetLocalTempDirectories(tmpdirs); + + auto* list = new TF_StringStream; + list->list = tmpdirs; + list->position = 0; + return list; +} + +TF_CAPI_EXPORT extern uint64_t TF_NowNanos(void) { + return ::tensorflow::Env::Default()->NowNanos(); +} + +// Returns the number of microseconds since the Unix epoch. +TF_CAPI_EXPORT extern uint64_t TF_NowMicros(void) { + return ::tensorflow::Env::Default()->NowMicros(); +} + +// Returns the number of seconds since the Unix epoch. +TF_CAPI_EXPORT extern uint64_t TF_NowSeconds(void) { + return ::tensorflow::Env::Default()->NowSeconds(); +} + +void TF_DefaultThreadOptions(TF_ThreadOptions* options) { + options->stack_size = 0; + options->guard_size = 0; + options->numa_node = -1; +} + +TF_Thread* TF_StartThread(const TF_ThreadOptions* options, + const char* thread_name, void (*work_func)(void*), + void* param) { + ::tensorflow::ThreadOptions cc_options; + cc_options.stack_size = options->stack_size; + cc_options.guard_size = options->guard_size; + cc_options.numa_node = options->numa_node; + return reinterpret_cast(::tensorflow::Env::Default()->StartThread( + cc_options, thread_name, [=]() { (*work_func)(param); })); +} + +void TF_JoinThread(TF_Thread* thread) { + // ::tensorflow::Thread joins on destruction + delete reinterpret_cast<::tensorflow::Thread*>(thread); +} diff --git a/tensorflow/c/env.h b/tensorflow/c/env.h new file mode 100644 index 0000000000000000000000000000000000000000..73078fcbbc5ae4c042f4a992655072a838e42915 --- /dev/null +++ b/tensorflow/c/env.h @@ -0,0 +1,195 @@ +/* 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 + +#ifndef TENSORFLOW_C_ENV_H_ +#define TENSORFLOW_C_ENV_H_ + +#include "tensorflow/c/c_api.h" + +// -------------------------------------------------------------------------- +// C API for tensorflow::Env. + +#ifdef __cplusplus +extern "C" { +#endif + +typedef struct TF_WritableFileHandle TF_WritableFileHandle; +typedef struct TF_StringStream TF_StringStream; +typedef struct TF_Thread TF_Thread; + +typedef struct TF_FileStatistics { + // The length of the file in bytes. + int64_t length; + // The last modified time in nanoseconds. + int64_t mtime_nsec; + // Whether the name refers to a directory. + bool is_directory; +} TF_FileStatistics; + +typedef struct TF_ThreadOptions { + // Thread stack size to use (in bytes), zero implies that the system default + // will be used. + size_t stack_size; + + // Guard area size to use near thread stacks to use (in bytes), zero implies + // that the system default will be used. + size_t guard_size; + + // The NUMA node to use, -1 implies that there should be no NUMA affinity for + // this thread. + int numa_node; +} TF_ThreadOptions; + +// Creates the specified directory. Typical status code are: +// * TF_OK - successfully created the directory +// * TF_ALREADY_EXISTS - directory already exists +// * TF_PERMISSION_DENIED - dirname is not writable +TF_CAPI_EXPORT extern void TF_CreateDir(const char* dirname, TF_Status* status); + +// Deletes the specified directory. Typical status codes are: +// * TF_OK - successfully deleted the directory +// * TF_FAILED_PRECONDITION - the directory is not empty +TF_CAPI_EXPORT extern void TF_DeleteDir(const char* dirname, TF_Status* status); + +// Deletes the specified directory and all subdirectories and files underneath +// it. This is accomplished by traversing the directory tree rooted at dirname +// and deleting entries as they are encountered. +// +// If dirname itself is not readable or does not exist, *undeleted_dir_count is +// set to 1, *undeleted_file_count is set to 0 and an appropriate status (e.g. +// TF_NOT_FOUND) is returned. +// +// If dirname and all its descendants were successfully deleted, TF_OK is +// returned and both error counters are set to zero. +// +// Otherwise, while traversing the tree, undeleted_file_count and +// undeleted_dir_count are updated if an entry of the corresponding type could +// not be deleted. The returned error status represents the reason that any one +// of these entries could not be deleted. +// +// Typical status codes: +// * TF_OK - dirname exists and we were able to delete everything underneath +// * TF_NOT_FOUND - dirname doesn't exist +// * TF_PERMISSION_DENIED - dirname or some descendant is not writable +// * TF_UNIMPLEMENTED - some underlying functions (like Delete) are not +// implemented +TF_CAPI_EXPORT extern void TF_DeleteRecursively(const char* dirname, + int64_t* undeleted_file_count, + int64_t* undeleted_dir_count, + TF_Status* status); + +// Obtains statistics for the given path. If status is TF_OK, *stats is +// updated, otherwise it is not touched. +TF_CAPI_EXPORT extern void TF_FileStat(const char* filename, + TF_FileStatistics* stats, + TF_Status* status); + +// Creates or truncates the given filename and returns a handle to be used for +// appending data to the file. If status is TF_OK, *handle is updated and the +// caller is responsible for freeing it (see TF_CloseWritableFile). +TF_CAPI_EXPORT extern void TF_NewWritableFile(const char* filename, + TF_WritableFileHandle** handle, + TF_Status* status); + +// Closes the given handle and frees its memory. If there was a problem closing +// the file, it is indicated by status. Memory is freed in any case. +TF_CAPI_EXPORT extern void TF_CloseWritableFile(TF_WritableFileHandle* handle, + TF_Status* status); + +// Syncs content of the handle to the filesystem. Blocks waiting for the +// filesystem to indicate that the content has been persisted. +TF_CAPI_EXPORT extern void TF_SyncWritableFile(TF_WritableFileHandle* handle, + TF_Status* status); + +// Flush local buffers to the filesystem. If the process terminates after a +// successful flush, the contents may still be persisted, since the underlying +// filesystem may eventually flush the contents. If the OS or machine crashes +// after a successful flush, the contents may or may not be persisted, depending +// on the implementation. +TF_CAPI_EXPORT extern void TF_FlushWritableFile(TF_WritableFileHandle* handle, + TF_Status* status); + +// Appends the given bytes to the file. Any failure to do so is indicated in +// status. +TF_CAPI_EXPORT extern void TF_AppendWritableFile(TF_WritableFileHandle* handle, + const char* data, + size_t length, + TF_Status* status); + +// Deletes the named file and indicates whether successful in *status. +TF_CAPI_EXPORT extern void TF_DeleteFile(const char* filename, + TF_Status* status); + +// Retrieves the next item from the given TF_StringStream and places a pointer +// to it in *result. If no more items are in the list, *result is set to NULL +// and false is returned. +// +// Ownership of the items retrieved with this function remains with the library. +// Item points are invalidated after a call to TF_StringStreamDone. +TF_CAPI_EXPORT extern bool TF_StringStreamNext(TF_StringStream* list, + const char** result); + +// Frees the resources associated with given string list. All pointers returned +// by TF_StringStreamNext are invalid after this call. +TF_CAPI_EXPORT extern void TF_StringStreamDone(TF_StringStream* list); + +// Retrieves the list of children of the given directory. You can iterate +// through the list with TF_StringStreamNext. The caller is responsible for +// freeing the list (see TF_StringStreamDone). +TF_CAPI_EXPORT extern TF_StringStream* TF_GetChildren(const char* filename, + TF_Status* status); + +// Retrieves a list of directory names on the local machine that may be used for +// temporary storage. You can iterate through the list with TF_StringStreamNext. +// The caller is responsible for freeing the list (see TF_StringStreamDone). +TF_CAPI_EXPORT extern TF_StringStream* TF_GetLocalTempDirectories(void); + +// Returns the number of nanoseconds since the Unix epoch. +TF_CAPI_EXPORT extern uint64_t TF_NowNanos(void); + +// Returns the number of microseconds since the Unix epoch. +TF_CAPI_EXPORT extern uint64_t TF_NowMicros(void); + +// Returns the number of seconds since the Unix epoch. +TF_CAPI_EXPORT extern uint64_t TF_NowSeconds(void); + +// Populates a TF_ThreadOptions struct with system-default values. +TF_CAPI_EXPORT extern void TF_DefaultThreadOptions(TF_ThreadOptions* options); + +// Returns a new thread that is running work_func and is identified +// (for debugging/performance-analysis) by thread_name. +// +// The given param (which may be null) is passed to work_func when the thread +// starts. In this way, data may be passed from the thread back to the caller. +// +// Caller takes ownership of the result and must call TF_JoinThread on it +// eventually. +TF_CAPI_EXPORT extern TF_Thread* TF_StartThread(const TF_ThreadOptions* options, + const char* thread_name, + void (*work_func)(void*), + void* param); + +// Waits for the given thread to finish execution, then deletes it. +TF_CAPI_EXPORT extern void TF_JoinThread(TF_Thread* thread); + +#ifdef __cplusplus +} +#endif + +#endif // TENSORFLOW_C_ENV_H_ diff --git a/tensorflow/c/env_test.cc b/tensorflow/c/env_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..687ad024137352662759ec1f43df87e89faca353 --- /dev/null +++ b/tensorflow/c/env_test.cc @@ -0,0 +1,127 @@ +/* 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/c/env.h" + +#include "tensorflow/c/c_api.h" +#include "tensorflow/core/lib/io/path.h" +#include "tensorflow/core/platform/mutex.h" +#include "tensorflow/core/platform/test.h" +#include "tensorflow/core/platform/types.h" + +#define ASSERT_TF_OK(x) ASSERT_EQ(TF_OK, TF_GetCode(x)) + +TEST(TestEnv, TestDirHandling) { + TF_StringStream* tempdirs = TF_GetLocalTempDirectories(); + const char* tempdir; + bool found = false; + while (TF_StringStreamNext(tempdirs, &tempdir)) { + found = true; + + TF_Status* s = TF_NewStatus(); + + ::tensorflow::string dirpath = + ::tensorflow::io::JoinPath(tempdir, "somedir"); + TF_CreateDir(dirpath.c_str(), s); + ASSERT_TF_OK(s) << "TF_CreateDir failed for " << dirpath << ": " + << TF_Message(s); + + ::tensorflow::string filepath = + ::tensorflow::io::JoinPath(dirpath, "somefile.txt"); + TF_WritableFileHandle* handle; + TF_NewWritableFile(filepath.c_str(), &handle, s); + ASSERT_TF_OK(s) << "NewWritableFile failed for " << filepath << ": " + << TF_Message(s); + + const char* data = "Hello, world!\n"; + TF_AppendWritableFile(handle, data, strlen(data), s); + ASSERT_TF_OK(s) << "TF_AppendWritableFile failed to append data to file at " + << filepath << ": " << TF_Message(s); + + TF_CloseWritableFile(handle, s); + ASSERT_TF_OK(s) << "TF_CloseWritableFile failed to close handle to " + << filepath << ": " << TF_Message(s); + + TF_StringStream* children = TF_GetChildren(dirpath.c_str(), s); + ASSERT_TF_OK(s) << "TF_GetChildren failed for " << dirpath; + const char* childpath; + ASSERT_TRUE(TF_StringStreamNext(children, &childpath)); + ASSERT_EQ(::tensorflow::string(childpath), "somefile.txt"); + // There should only be one file in this directory. + ASSERT_FALSE(TF_StringStreamNext(children, &childpath)); + ASSERT_EQ(childpath, nullptr); + TF_StringStreamDone(children); + + TF_FileStatistics stats; + TF_FileStat(filepath.c_str(), &stats, s); + ASSERT_EQ(stats.length, strlen(data)); + ASSERT_FALSE(stats.is_directory); + ASSERT_GT(stats.mtime_nsec, 0); + + // Trying to delete a non-empty directory should fail. + TF_DeleteDir(dirpath.c_str(), s); + ASSERT_NE(TF_OK, TF_GetCode(s)) + << "TF_DeleteDir unexpectedly succeeded with a non-empty directory " + << dirpath; + + TF_DeleteFile(filepath.c_str(), s); + ASSERT_TF_OK(s) << "TF_DeleteFile failed for " << filepath << ": " + << TF_Message(s); + + // Now deleting the directory should work. + TF_DeleteDir(dirpath.c_str(), s); + ASSERT_TF_OK(s) << "TF_DeleteDir failed for " << dirpath << ": " + << TF_Message(s); + + TF_DeleteStatus(s); + break; + } + + ASSERT_TRUE(found) << "expected at least one temp dir"; + + TF_StringStreamDone(tempdirs); +} + +TEST(TestEnv, TestTimeFunctions) { + ASSERT_GE(TF_NowSeconds(), 946684800); // Midnight Jan 1, 2000 + ASSERT_GE(TF_NowMicros(), 946684800 * 1e6); + ASSERT_GE(TF_NowNanos(), 946684800 * 1e9); +} + +namespace { + +struct SomeThreadData { + ::tensorflow::mutex mu; + bool did_work = false; +}; + +void SomeThreadFunc(void* data) { + auto* real_data = static_cast(data); + ::tensorflow::mutex_lock l(real_data->mu); + real_data->did_work = true; +} + +} // namespace + +TEST(TestEnv, TestThreads) { + TF_ThreadOptions options; + TF_DefaultThreadOptions(&options); + SomeThreadData data; + TF_Thread* thread = + TF_StartThread(&options, "SomeThreadName", &SomeThreadFunc, &data); + TF_JoinThread(thread); + ::tensorflow::mutex_lock l(data.mu); + ASSERT_TRUE(data.did_work); +} diff --git a/tensorflow/c/kernels.cc b/tensorflow/c/kernels.cc new file mode 100644 index 0000000000000000000000000000000000000000..9505bf9dda32b9a338b574f1d31ec555a5628c6a --- /dev/null +++ b/tensorflow/c/kernels.cc @@ -0,0 +1,202 @@ +/* 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 + +#include "tensorflow/c/c_api_internal.h" +#include "tensorflow/c/kernels.h" +#include "tensorflow/c/tf_status_helper.h" +#include "tensorflow/core/framework/kernel_def_builder.h" +#include "tensorflow/core/framework/op_kernel.h" + +// This file forms the basis of a stable ABI for third-party kernel +// implementations. It is crucial that changes to this file are made cautiously +// and with a focus on maintaining both source and binary compatibility. + +struct TF_KernelBuilder { + ::tensorflow::KernelDefBuilder* cc_builder; + + void* (*create_function)(TF_OpKernelConstruction*); + void (*compute_function)(void*, TF_OpKernelContext*); + void (*delete_function)(void*); +}; + +TF_KernelBuilder* TF_NewKernelBuilder( + const char* op_name, const char* device_name, + void* (*create_func)(TF_OpKernelConstruction*), + void (*compute_func)(void*, TF_OpKernelContext*), + void (*delete_func)(void*)) { + TF_KernelBuilder* result = new TF_KernelBuilder; + result->cc_builder = new ::tensorflow::KernelDefBuilder(op_name); + result->cc_builder->Device(device_name); + result->create_function = create_func; + result->compute_function = compute_func; + result->delete_function = delete_func; + return result; +} + +void TF_DeleteKernelBuilder(TF_KernelBuilder* builder) { + if (builder != nullptr) { + delete builder->cc_builder; + delete builder; + } +} + +namespace tensorflow { +namespace { + +// An OpKernel whose methods delegate to C function pointers. +class COpKernel : public OpKernel { + public: + explicit COpKernel(OpKernelConstruction* ctx, + void* (*create_func)(TF_OpKernelConstruction*), + void (*compute_func)(void*, TF_OpKernelContext*), + void (*delete_func)(void*)) + : OpKernel(ctx), compute_func_(compute_func), delete_func_(delete_func) { + if (create_func != nullptr) { + c_kernel_ = + (*create_func)(reinterpret_cast(ctx)); + } else { + c_kernel_ = nullptr; + } + } + + void Compute(OpKernelContext* ctx) override { + (*compute_func_)(c_kernel_, reinterpret_cast(ctx)); + } + + ~COpKernel() override { + if (delete_func_ != nullptr) { + (*delete_func_)(c_kernel_); + } + } + + private: + void (*compute_func_)(void*, TF_OpKernelContext* context); + void (*delete_func_)(void*); + void* c_kernel_; +}; + +// A KernelFactory that returns COpKernel instances. +class KernelBuilderFactory + : public ::tensorflow::kernel_factory::OpKernelFactory { + public: + explicit KernelBuilderFactory(TF_KernelBuilder* builder) + : builder_(builder) {} + ::tensorflow::OpKernel* Create( + ::tensorflow::OpKernelConstruction* context) override { + return new ::tensorflow::COpKernel(context, builder_->create_function, + builder_->compute_function, + builder_->delete_function); + } + ~KernelBuilderFactory() override { TF_DeleteKernelBuilder(builder_); } + + private: + TF_KernelBuilder* builder_; +}; +} // namespace +} // namespace tensorflow + +void TF_RegisterKernelBuilder(const char* name, TF_KernelBuilder* builder, + TF_Status* status) { + using tensorflow::register_kernel::Name; + + tensorflow::kernel_factory::OpKernelRegistrar( + builder->cc_builder->Build(), name, + absl::make_unique(builder)); + + TF_SetStatus(status, TF_OK, ""); +} + +int TF_NumInputs(TF_OpKernelContext* ctx) { + auto* cc_ctx = reinterpret_cast<::tensorflow::OpKernelContext*>(ctx); + return cc_ctx->num_inputs(); +} + +int TF_NumOutputs(TF_OpKernelContext* ctx) { + auto* cc_ctx = reinterpret_cast<::tensorflow::OpKernelContext*>(ctx); + return cc_ctx->num_outputs(); +} + +void TF_GetInput(TF_OpKernelContext* ctx, int i, TF_Tensor** tensor, + TF_Status* status) { + auto* cc_ctx = reinterpret_cast<::tensorflow::OpKernelContext*>(ctx); + if (i < 0 || i >= cc_ctx->num_inputs()) { + TF_SetStatus(status, TF_OUT_OF_RANGE, "input index out of range"); + return; + } + const ::tensorflow::Tensor& cc_tensor(cc_ctx->input(i)); + TF_Tensor* result = ::tensorflow::TF_TensorFromTensor(cc_tensor, status); + if (TF_GetCode(status) == TF_OK) { + *tensor = result; + } +} + +void TF_SetOutput(TF_OpKernelContext* ctx, int i, const TF_Tensor* tensor, + TF_Status* status) { + auto* cc_ctx = reinterpret_cast<::tensorflow::OpKernelContext*>(ctx); + if (i < 0 || i >= cc_ctx->num_inputs()) { + TF_SetStatus(status, TF_OUT_OF_RANGE, "input index out of range"); + return; + } + ::tensorflow::Tensor cc_tensor; + ::tensorflow::Status s = ::tensorflow::TF_TensorToTensor(tensor, &cc_tensor); + TF_SetStatus(status, TF_OK, ""); + ::tensorflow::Set_TF_Status_from_Status(status, s); + if (s.ok()) { + 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_(struct_name, func, c_type, cc_type) \ + void struct_name##_GetAttr##func(struct_name* 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 DEFINE_TF_GETATTR(func, c_type, cc_type) \ + DEFINE_TF_GETATTR_(TF_OpKernelConstruction, func, c_type, cc_type) \ + DEFINE_TF_GETATTR_(TF_OpKernelContext, func, c_type, cc_type) + +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 new file mode 100644 index 0000000000000000000000000000000000000000..b015d0103969355e8566242bfcc007f697c6ae18 --- /dev/null +++ b/tensorflow/c/kernels.h @@ -0,0 +1,153 @@ +/* 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_C_KERNELS_H_ +#define TENSORFLOW_C_KERNELS_H_ + +#include "tensorflow/c/c_api.h" + +#ifdef __cplusplus +extern "C" { +#endif + +// -------------------------------------------------------------------------- +// C API for TensorFlow Kernels. +// +// This API allows developers to register custom kernel implementations for +// TensorFlow. +// +// See c_api.h header comments for a discussion about API conventions. +// +// Users wishing to extend TensorFlow with new kernels will call +// `TF_NewKernelBuilder`. The resulting kernel builder can be registered with +// `TF_RegisterKernelBuilder`, which will allow TF to construct user-provided +// kernels when necessary. + +typedef struct TF_KernelBuilder TF_KernelBuilder; +typedef struct TF_OpKernelConstruction TF_OpKernelConstruction; +typedef struct TF_OpKernelContext TF_OpKernelContext; + +// Allocates a new kernel builder and returns a pointer to it. +// +// If non-null, TensorFlow will call create_func when it needs to instantiate +// the kernel. The pointer returned by create_func will be passed to +// compute_func and delete_func, thereby functioning as a "this" pointer for +// referring to kernel instances. +// +// The TF_OpKernelConstruction pointer passed to create_func is owned by +// TensorFlow and will be deleted once create_func returns. It must not be used +// after this. +// +// When TensorFlow needs to perform a computation with this kernel, it will +// call compute_func. This function will receive the pointer returned by +// create_func (or null if no create_func was provided), along with the inputs +// to the computation. +// +// The TF_OpKernelContext pointer received by compute_func is owned by +// TensorFlow and will be deleted once compute_func returns. It must not be used +// after this. +// +// Finally, when TensorFlow no longer needs the kernel, it will call +// delete_func if one is provided. This function will receive the pointer +// returned in `create_func` or nullptr if no `create_func` was provided. +// +// The caller should pass the result of this function to +// TF_RegisterKernelBuilder, which will take ownership of the pointer. If, for +// some reason, the kernel builder will not be registered, the caller should +// delete it with TF_DeleteKernelBuilder. +TF_CAPI_EXPORT extern TF_KernelBuilder* TF_NewKernelBuilder( + const char* op_name, const char* device_name, + void* (*create_func)(TF_OpKernelConstruction*), + void (*compute_func)(void*, TF_OpKernelContext*), + void (*delete_func)(void*)); + +// Register the given kernel builder with the TensorFlow runtime. If +// registration fails, the given status will be populated. +// +// This call takes ownership of the `builder` pointer. +TF_CAPI_EXPORT extern void TF_RegisterKernelBuilder(const char* kernel_name, + TF_KernelBuilder* builder, + TF_Status* status); + +// Deletes the given TF_KernelBuilder. This should be called only if the kernel +// builder is not registered with TensorFlow via TF_RegisterKernelBuilder. +TF_CAPI_EXPORT extern void TF_DeleteKernelBuilder(TF_KernelBuilder* builder); + +// -------------------------------------------------------------------------- +// OpKernelContext routines + +// TF_NumInputs returns the number of inputs available in ctx. +TF_CAPI_EXPORT extern int TF_NumInputs(TF_OpKernelContext* ctx); + +// TF_NumOutputs returns the number of outputs to be placed in *ctx by the +// kernel. +TF_CAPI_EXPORT extern int TF_NumOutputs(TF_OpKernelContext* ctx); + +// Retrieves the ith input from ctx. If TF_GetCode(status) is TF_OK, *tensor is +// populated and its ownership is passed to the caller. In any other case, +// *tensor is not modified. +// +// If i < 0 or i >= TF_NumInputs(ctx), *status is set to TF_OUT_OF_RANGE. +TF_CAPI_EXPORT extern void TF_GetInput(TF_OpKernelContext* ctx, int i, + TF_Tensor** tensor, TF_Status* status); + +// Sets the ith output of ctx to tensor. If TF_GetCode(status) is anything but +// TF_OK, ctx is left unmodified. +// +// If i < 0 or i >= TF_NumOutputs(ctx), *status is set to TF_OUT_OF_RANGE. +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); + +// Interprets the named kernel context 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_OpKernelContext_GetAttrType( + TF_OpKernelContext* ctx, const char* attr_name, TF_DataType* val, + TF_Status* status); + +#ifdef __cplusplus +} /* end extern "C" */ +#endif + +#endif // TENSORFLOW_C_KERNELS_H_ diff --git a/tensorflow/c/kernels_test.cc b/tensorflow/c/kernels_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..0d2954717e7a83c102a35815809a554e3a917e07 --- /dev/null +++ b/tensorflow/c/kernels_test.cc @@ -0,0 +1,231 @@ +/* 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/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" +#include "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/framework/types.h" +#include "tensorflow/core/lib/core/status_test_util.h" +#include "tensorflow/core/platform/test.h" + +struct MyCustomKernel { + bool created; + bool compute_called; +}; + +static bool delete_called = false; + +static void* MyCreateFunc(TF_OpKernelConstruction* ctx) { + struct MyCustomKernel* s = new struct MyCustomKernel; + s->created = true; + s->compute_called = false; + return s; +} + +static void MyComputeFunc(void* kernel, TF_OpKernelContext* ctx) { + struct MyCustomKernel* s = static_cast(kernel); + s->compute_called = true; + if (ctx != nullptr) { + TF_Status* status = TF_NewStatus(); + + EXPECT_EQ(43, TF_StepId(ctx)); + + // Exercise attribute reads. + TF_DataType type; + TF_OpKernelContext_GetAttrType(ctx, "SomeDataTypeAttr", &type, status); + EXPECT_EQ(TF_OK, TF_GetCode(status)); + EXPECT_EQ(TF_FLOAT, type); + + TF_DeleteStatus(status); + } +} + +static void MyDeleteFunc(void* kernel) { + struct MyCustomKernel* s = static_cast(kernel); + EXPECT_TRUE(s->created); + EXPECT_TRUE(s->compute_called); + delete_called = true; + delete s; +} + +namespace tensorflow { + +static std::unique_ptr GetFakeKernel(const char* device_name, + const char* op_name, + Status* status) { + NodeDef def; + def.set_op(op_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); +} + +// Tests registration of a single C kernel and checks that calls through the +// C/C++ boundary are being made. +TEST(TestKernel, TestRegisterKernelBuilder) { + const char* kernel_name = "SomeKernelName"; + const char* op_name = "FooOp"; + const char* device_name = "FakeDeviceName1"; + + REGISTER_OP(op_name) + .Input("input1: double") + .Input("input2: uint8") + .Output("output1: uint8") + .Attr("SomeDataTypeAttr: type"); + + TF_KernelBuilder* builder = TF_NewKernelBuilder( + op_name, device_name, &MyCreateFunc, &MyComputeFunc, &MyDeleteFunc); + + { + TF_Status* status = TF_NewStatus(); + TF_RegisterKernelBuilder(kernel_name, builder, status); + EXPECT_EQ(TF_OK, TF_GetCode(status)); + TF_Buffer* buf = TF_GetRegisteredKernelsForOp(op_name, status); + EXPECT_EQ(TF_OK, TF_GetCode(status)); + KernelList list; + list.ParseFromArray(buf->data, buf->length); + ASSERT_EQ(1, list.kernel_size()); + ASSERT_EQ(device_name, list.kernel(0).device_type()); + TF_DeleteBuffer(buf); + TF_DeleteStatus(status); + } + + { + Status status; + std::unique_ptr kernel = + GetFakeKernel(device_name, op_name, &status); + TF_EXPECT_OK(status); + ASSERT_NE(nullptr, kernel.get()); + kernel->Compute(nullptr); + } + + ASSERT_TRUE(delete_called); +} + +class DummyDevice : public DeviceBase { + public: + DummyDevice(Env* env, bool save) : DeviceBase(env), save_(save) {} + bool RequiresRecordingAccessedTensors() const override { return save_; } + Allocator* GetAllocator(AllocatorAttributes /*attr*/) override { + return cpu_allocator(); + } + + private: + bool save_; +}; + +TEST(TestKernel, TestInputAndOutputCount) { + const char* kernel_name = "InputOutputCounterKernel"; + const char* op_name = "BarOp"; + const char* device_name = "FakeDeviceName2"; + + REGISTER_OP(op_name) + .Input("input1: double") + .Input("input2: uint8") + .Output("output1: uint8") + .Attr("SomeDataTypeAttr: type"); + + static int num_inputs = 0; + static int num_outputs = 0; + + // A kernel whose Compute function has a side-effect of updating num_inputs + // and num_outputs. Various functions on TF_OpKernelContext are also + // exercised. + auto my_compute_func = [](void* kernel, TF_OpKernelContext* ctx) { + num_inputs = TF_NumInputs(ctx); + num_outputs = TF_NumOutputs(ctx); + + TF_Tensor* input = nullptr; + TF_Status* s = TF_NewStatus(); + TF_GetInput(ctx, 0, &input, s); + EXPECT_EQ(TF_OK, TF_GetCode(s)) << "Failed to get input: " << TF_Message(s); + EXPECT_EQ(123, *static_cast(TF_TensorData(input))); + TF_GetInput(ctx, -1, &input, s); + EXPECT_EQ(TF_OUT_OF_RANGE, TF_GetCode(s)); + TF_GetInput(ctx, 3, &input, s); + EXPECT_EQ(TF_OUT_OF_RANGE, TF_GetCode(s)); + + // Copy the input tensor to output. + TF_SetOutput(ctx, 0, input, s); + EXPECT_EQ(TF_OK, TF_GetCode(s)); + + 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); + } + }; + + TF_KernelBuilder* builder = TF_NewKernelBuilder(op_name, device_name, nullptr, + my_compute_func, nullptr); + + { + TF_Status* status = TF_NewStatus(); + TF_RegisterKernelBuilder(kernel_name, builder, status); + EXPECT_EQ(TF_OK, TF_GetCode(status)); + TF_DeleteStatus(status); + } + + { + OpKernelContext::Params p; + DummyDevice dummy_device(nullptr, false); + p.device = &dummy_device; + p.step_id = 43; + + Tensor t(tensorflow::uint8(123)); + + gtl::InlinedVector inputs; + // Simulate 2 inputs + inputs.emplace_back(&t); + inputs.emplace_back(); + p.inputs = &inputs; + + Status status; + std::unique_ptr kernel = + GetFakeKernel(device_name, op_name, &status); + TF_EXPECT_OK(status); + ASSERT_NE(nullptr, kernel.get()); + + p.op_kernel = kernel.get(); + OpKernelContext ctx(&p); + kernel->Compute(&ctx); + + ASSERT_EQ(2, num_inputs); + ASSERT_EQ(1, num_outputs); + ASSERT_EQ(123, ctx.mutable_output(0)->scalar()()); + } +} + +TEST(TestKernel, DeleteKernelBuilderIsOkOnNull) { + TF_DeleteKernelBuilder(nullptr); +} + +} // namespace tensorflow diff --git a/tensorflow/c/python_api.cc b/tensorflow/c/python_api.cc index 247236b760dd8c07bbb08426100b6a4d34296d2e..98d8393332269ae349cf8aa5c0b612c6f17172e6 100644 --- a/tensorflow/c/python_api.cc +++ b/tensorflow/c/python_api.cc @@ -160,4 +160,17 @@ void SetHandleShapeAndType(TF_Graph* graph, TF_Output output, const void* proto, ic->set_output_handle_shapes_and_types(output.index, shapes_and_types); } +void AddWhileInputHack(TF_Graph* graph, TF_Output new_src, TF_Operation* dst, + TF_Status* status) { + mutex_lock l(graph->mu); + status->status = graph->graph.AddWhileInputHack(&new_src.oper->node, + new_src.index, &dst->node); + if (status->status.ok()) { + // This modification only updates the destination node for + // the purposes of running this graph in a session. Thus, we don't + // record the source node as being modified. + RecordMutation(graph, *dst, "adding input tensor"); + } +} + } // namespace tensorflow diff --git a/tensorflow/c/python_api.h b/tensorflow/c/python_api.h index 5cce84020bc68d912d259f51512341eb5f464a2c..44779ca656165dd65590cb5e9ea3ccf71165ed63 100644 --- a/tensorflow/c/python_api.h +++ b/tensorflow/c/python_api.h @@ -34,6 +34,7 @@ void SetAttr(TF_Graph* graph, TF_Operation* op, const char* attr_name, void SetRequestedDevice(TF_Graph* graph, TF_Operation* op, const char* device); +// Updates 'dst' to consume 'new_src'. void UpdateEdge(TF_Graph* graph, TF_Output new_src, TF_Input dst, TF_Status* status); @@ -65,6 +66,13 @@ std::string GetHandleShapeAndType(TF_Graph* graph, TF_Output output); // because I couldn't get SWIG to work otherwise. void SetHandleShapeAndType(TF_Graph* graph, TF_Output output, const void* proto, size_t proto_len, TF_Status* status); + +// This method is used to add a new input edge to 'dst', which must be a While +// op. The While op's "T" attribute must have already been updated to include +// the new edge. This is used to construct tf.while_loop gradients. +void AddWhileInputHack(TF_Graph* graph, TF_Output new_src, TF_Operation* dst, + TF_Status* status); + } // namespace tensorflow #endif // TENSORFLOW_C_PYTHON_API_H_ diff --git a/tensorflow/c/test_op1.cc b/tensorflow/c/test_op1.cc new file mode 100644 index 0000000000000000000000000000000000000000..b22cc9aef2b344282f45340ff12ee849935a26f9 --- /dev/null +++ b/tensorflow/c/test_op1.cc @@ -0,0 +1,23 @@ +/* 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/op.h" +#include "tensorflow/core/framework/op_kernel.h" + +namespace tensorflow { + +REGISTER_OP("TestCApi1").Doc(R"doc(Used to test C API)doc"); + +} // namespace tensorflow diff --git a/tensorflow/cc/BUILD b/tensorflow/cc/BUILD index c18b07603ae3841d3581741ab5a43f2e8b628356..a09becc49b10d2c58f98fbcc11df5190f794c1d4 100644 --- a/tensorflow/cc/BUILD +++ b/tensorflow/cc/BUILD @@ -170,6 +170,7 @@ cc_library_with_android_deps( "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:protos_all_cc", + "@com_google_absl//absl/strings", ], ) @@ -488,6 +489,7 @@ tf_gen_op_wrappers_cc( "image_ops", "io_ops", "linalg_ops", + "list_ops", "logging_ops", "lookup_ops", "manip_ops", @@ -516,6 +518,8 @@ tf_gen_op_wrappers_cc( ":array_ops", ":const_op", ":math_ops", + "//tensorflow/cc:ops", + "//tensorflow/cc:scope", ], ) diff --git a/tensorflow/cc/framework/scope.cc b/tensorflow/cc/framework/scope.cc index 6abc9e268e3ac97379954a34017ddffa010db67f..81785b2d89b3d36b46992b7ae376b5175a806027 100644 --- a/tensorflow/cc/framework/scope.cc +++ b/tensorflow/cc/framework/scope.cc @@ -95,6 +95,7 @@ Scope::Impl::Impl(const Scope& other, Tags::ScopeName, const string& name, kernel_label_(other.impl()->kernel_label_), device_(other.impl()->device_), assigned_device_(other.impl()->assigned_device_), + xla_cluster_(other.impl()->xla_cluster_), colocation_constraints_(other.impl()->colocation_constraints_), disable_shape_inference_(other.impl()->disable_shape_inference_) {} @@ -112,6 +113,7 @@ Scope::Impl::Impl(const Scope& other, Tags::OpName, const string& name, kernel_label_(other.impl()->kernel_label_), device_(other.impl()->device_), assigned_device_(other.impl()->assigned_device_), + xla_cluster_(other.impl()->xla_cluster_), colocation_constraints_(other.impl()->colocation_constraints_), disable_shape_inference_(other.impl()->disable_shape_inference_) {} @@ -135,6 +137,7 @@ Scope::Impl::Impl(const Scope& other, Tags::ControlDeps, kernel_label_(other.impl()->kernel_label_), device_(other.impl()->device_), assigned_device_(other.impl()->assigned_device_), + xla_cluster_(other.impl()->xla_cluster_), colocation_constraints_(other.impl()->colocation_constraints_), disable_shape_inference_(other.impl()->disable_shape_inference_) {} @@ -167,6 +170,7 @@ Scope::Impl::Impl(const Scope& other, Tags::SingleUseScope, kernel_label_(other.impl()->kernel_label_), device_(other.impl()->device_), assigned_device_(other.impl()->assigned_device_), + xla_cluster_(other.impl()->xla_cluster_), colocation_constraints_(other.impl()->colocation_constraints_), disable_shape_inference_(other.impl()->disable_shape_inference_) {} @@ -183,6 +187,7 @@ Scope::Impl::Impl(const Scope& other, Tags::ExitOnError) kernel_label_(other.impl()->kernel_label_), device_(other.impl()->device_), assigned_device_(other.impl()->assigned_device_), + xla_cluster_(other.impl()->xla_cluster_), colocation_constraints_(other.impl()->colocation_constraints_), disable_shape_inference_(other.impl()->disable_shape_inference_) {} @@ -200,6 +205,7 @@ Scope::Impl::Impl(const Scope& other, Tags::KernelLabel, kernel_label_(kernel_label), device_(other.impl()->device_), assigned_device_(other.impl()->assigned_device_), + xla_cluster_(other.impl()->xla_cluster_), colocation_constraints_(other.impl()->colocation_constraints_), disable_shape_inference_(other.impl()->disable_shape_inference_) {} @@ -217,6 +223,7 @@ Scope::Impl::Impl(const Scope& other, Tags::Colocate, kernel_label_(other.impl()->kernel_label_), device_(other.impl()->device_), assigned_device_(other.impl()->assigned_device_), + xla_cluster_(other.impl()->xla_cluster_), colocation_constraints_( clear_colocations ? std::unordered_set() @@ -237,6 +244,25 @@ Scope::Impl::Impl(const Scope& other, Tags::AssignedDevice, kernel_label_(other.impl()->kernel_label_), device_(other.impl()->device_), assigned_device_(assigned_device), + xla_cluster_(other.impl()->xla_cluster_), + colocation_constraints_(other.impl()->colocation_constraints_), + disable_shape_inference_(other.impl()->disable_shape_inference_) {} + +Scope::Impl::Impl(const Scope& other, Tags::XlaCluster, + const string& xla_cluster) + : graph_(other.impl()->graph_), + status_(other.impl()->status_), + name_map_(other.impl()->name_map_), + refiner_(other.impl()->refiner_), + scope_used_(other.impl()->scope_used_), + control_deps_(other.impl()->control_deps_), + name_(other.impl()->name_), + op_name_(other.impl()->op_name_), + exit_on_error_(other.impl()->exit_on_error_), + kernel_label_(other.impl()->kernel_label_), + device_(other.impl()->device_), + assigned_device_(other.impl()->assigned_device_), + xla_cluster_(xla_cluster), colocation_constraints_(other.impl()->colocation_constraints_), disable_shape_inference_(other.impl()->disable_shape_inference_) {} @@ -326,6 +352,9 @@ void Scope::UpdateBuilder(NodeBuilder* builder) const { if (!impl()->assigned_device_.empty()) { builder->AssignedDevice(impl()->assigned_device_); } + if (!impl()->xla_cluster_.empty()) { + builder->XlaCluster(impl()->xla_cluster_); + } } string Scope::Impl::GetUniqueName(const string& prefix, @@ -388,7 +417,7 @@ Scope Scope::NewSubScope(const string& child_scope_name) const { false /* copy_names */)); } -Scope Scope::WithOpName(const string& op_name) const { +Scope Scope::WithOpNameImpl(const string& op_name) const { if (impl()->single_use_scope()) { UpdateStatus(errors::InvalidArgument("Cannot set op name ", op_name, " on this scope")); @@ -425,6 +454,10 @@ Scope Scope::WithAssignedDevice(const string& assigned_device) const { return Scope(new Impl(*this, Impl::Tags::AssignedDevice(), assigned_device)); } +Scope Scope::WithXlaCluster(const string& xla_cluster) const { + return Scope(new Impl(*this, Impl::Tags::XlaCluster(), xla_cluster)); +} + Scope Scope::ColocateWith(const Operation& op) const { return Scope(new Impl(*this, Impl::Tags::Colocate(), op, /* clear_colocations */ false)); diff --git a/tensorflow/cc/framework/scope.h b/tensorflow/cc/framework/scope.h index e307d8989b6647dfac8d2691ed2171c86b7f3a7c..0a75f23725c143e6b22ee6dffae1428ed8209fe8 100644 --- a/tensorflow/cc/framework/scope.h +++ b/tensorflow/cc/framework/scope.h @@ -22,6 +22,7 @@ limitations under the License. #include #include +#include "absl/strings/str_cat.h" #include "tensorflow/cc/framework/ops.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/lib/gtl/array_slice.h" @@ -69,8 +70,9 @@ struct CompositeOpScopes; /// // W will be named "linear/W" /// auto W = Variable(linear.WithOpName("W"), /// {2, 2}, DT_FLOAT); -/// // b will be named "linear/b" -/// auto b = Variable(linear.WithOpName("b"), +/// // b will be named "linear/b_3" +/// int idx = 3; +/// auto b = Variable(linear.WithOpName("b_", idx), /// {2}, DT_FLOAT); /// auto x = Const(linear, {...}); // name: "linear/Const" /// auto m = MatMul(linear, x, W); // name: "linear/MatMul" @@ -113,8 +115,11 @@ class Scope { Scope NewSubScope(const string& child_scope_name) const; /// Return a new scope. All ops created within the returned scope will have - /// names of the form `name/op_name[_suffix]`. - Scope WithOpName(const string& op_name) const; + /// names of the form `name/StrCat(fragments...)[_suffix]` + template + Scope WithOpName(Ty... fragments) const { + return WithOpNameImpl(absl::StrCat(fragments...)); + } /// Return a new scope. All ops created within the returned scope will have as /// control dependencies the union of operations in the control_deps vector @@ -137,6 +142,10 @@ class Scope { /// their assigned device set to `assigned_device`. Scope WithAssignedDevice(const string& assigned_device) const; + /// Returns a new scope. All ops created within the returned scope will have + /// their _XlaCluster attribute set to `xla_cluster`. + Scope WithXlaCluster(const string& xla_cluster) const; + /// Return a new scope. All ops created within the returned scope will be /// co-located on the device where op is placed. /// NOTE: This function is intended to be use internal libraries only for @@ -227,6 +236,8 @@ class Scope { // END_SKIP_DOXYGEN private: + Scope WithOpNameImpl(const string& op_name) const; + friend class InternalScope; std::unique_ptr impl_; explicit Scope(Impl*); diff --git a/tensorflow/cc/framework/scope_internal.h b/tensorflow/cc/framework/scope_internal.h index 514e02e84146b6d95147d83182e5d9a07509cfa1..5db7eab2b819c2c5d8fc358953d4607848f1cba5 100644 --- a/tensorflow/cc/framework/scope_internal.h +++ b/tensorflow/cc/framework/scope_internal.h @@ -61,6 +61,7 @@ class Scope::Impl { enum class KernelLabel; enum class Colocate; enum class AssignedDevice; + enum class XlaCluster; }; Impl(Graph* graph, Status* status, NameMap* name_map, ShapeRefiner* refiner, @@ -78,6 +79,7 @@ class Scope::Impl { Impl(const Scope& other, Tags::Colocate, const Operation& colocate_with_op, bool clear_colocations); Impl(const Scope& other, Tags::AssignedDevice, const string& assigned_device); + Impl(const Scope& other, Tags::XlaCluster, const string& xla_cluster); std::unordered_set GetColocationConstraints( const Operation& colocate_with_op) const; @@ -112,6 +114,7 @@ class Scope::Impl { const string kernel_label_ = ""; const string device_ = ""; const string assigned_device_ = ""; + const string xla_cluster_ = ""; const std::unordered_set colocation_constraints_; // If true, Scope::DoShapeInference() always returns Status:OK(). diff --git a/tensorflow/cc/gradients/image_grad.cc b/tensorflow/cc/gradients/image_grad.cc index 882709e1e2817431a32c453fe0f35f2b2e6c69b0..05c287bdc62cdb8be7208ce3975f280aaa816766 100644 --- a/tensorflow/cc/gradients/image_grad.cc +++ b/tensorflow/cc/gradients/image_grad.cc @@ -69,6 +69,23 @@ Status ResizeBicubicGradHelper(const Scope& scope, const Operation& op, } REGISTER_GRADIENT_OP("ResizeBicubic", ResizeBicubicGradHelper); +Status ScaleAndTranslateGradHelper(const Scope& scope, const Operation& op, + const std::vector& grad_inputs, + std::vector* grad_outputs) { + string kernel_type; + TF_RETURN_IF_ERROR( + GetNodeAttr(op.node()->attrs(), "kernel_type", &kernel_type)); + grad_outputs->push_back(internal::ScaleAndTranslateGrad( + scope, grad_inputs[0], op.input(0), op.input(2), op.input(3), + internal::ScaleAndTranslateGrad::KernelType(kernel_type))); + + grad_outputs->push_back(NoGradient()); + grad_outputs->push_back(NoGradient()); + grad_outputs->push_back(NoGradient()); + return scope.status(); +} +REGISTER_GRADIENT_OP("ScaleAndTranslate", ScaleAndTranslateGradHelper); + } // anonymous namespace } // namespace ops } // namespace tensorflow diff --git a/tensorflow/cc/gradients/image_grad_test.cc b/tensorflow/cc/gradients/image_grad_test.cc index 2e55c7561b030c50bd67bd53fd0d55710085c5d2..1d150226538093467e092e02f38090a327f9c9b6 100644 --- a/tensorflow/cc/gradients/image_grad_test.cc +++ b/tensorflow/cc/gradients/image_grad_test.cc @@ -30,6 +30,7 @@ using ops::Const; using ops::ResizeBicubic; using ops::ResizeBilinear; using ops::ResizeNearestNeighbor; +using ops::ScaleAndTranslate; class ImageGradTest : public ::testing::Test { protected: @@ -153,5 +154,45 @@ TEST_F(ImageGradTest, TestBicubic) { TestResize(RESIZE_BICUBIC); } +class ScaleAndTranslateGradTest : public ::testing::Test { + protected: + ScaleAndTranslateGradTest() : scope_(Scope::NewRootScope()) {} + + template + Tensor MakeData(const TensorShape& data_shape) { + DataType data_type = DataTypeToEnum::v(); + Tensor data(data_type, data_shape); + auto data_flat = data.flat(); + for (int i = 0; i < data_flat.size(); ++i) { + data_flat(i) = T(i); + } + return data; + } + + template + void MakeOp(const Tensor& x_data, const Input& y_shape, Output* x, + Output* y) { + *x = Const(scope_, x_data); + *y = ScaleAndTranslate(scope_, *x, y_shape, {1.8f, 2.1f}, {0.5f, 0.7f}); + TF_ASSERT_OK(scope_.status()); + } + + template + void TestResize() { + TensorShape x_shape({1, 2, 3, 1}); + Tensor x_data = MakeData(x_shape); + Output x, y; + MakeOp(x_data, {4, 6}, &x, &y); + JAC_T max_error; + TF_ASSERT_OK((ComputeGradientError( + scope_, x, x_data, y, {1, 4, 6, 1}, &max_error))); + EXPECT_LT(max_error, 1e-3); + } + + Scope scope_; +}; + +TEST_F(ScaleAndTranslateGradTest, Works) { TestResize(); } + } // namespace } // namespace tensorflow 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/BUILD b/tensorflow/cc/saved_model/BUILD index 3d3895c8fa82c3c0e2974228e9cad767d0e00df4..52345a376cc29ee47ccb9888c9bb26292468b5a9 100644 --- a/tensorflow/cc/saved_model/BUILD +++ b/tensorflow/cc/saved_model/BUILD @@ -133,5 +133,6 @@ filegroup( "testdata/half_plus_two_pbtxt/**", "testdata/half_plus_two_main_op/**", "testdata/half_plus_two/**", + "testdata/half_plus_two_v2/**", ]), ) diff --git a/tensorflow/cc/saved_model/constants.h b/tensorflow/cc/saved_model/constants.h index 645a3f101d1ae7dda88ec4ca622c694dc5a7a919..6f00dc324bd7054b28de2c35023581e1666bfa01 100644 --- a/tensorflow/cc/saved_model/constants.h +++ b/tensorflow/cc/saved_model/constants.h @@ -33,10 +33,10 @@ constexpr char kSavedModelFilenamePb[] = "saved_model.pb"; /// SavedModel text format proto filename. constexpr char kSavedModelFilenamePbTxt[] = "saved_model.pbtxt"; -/// SavedModel legacy init op key. +/// SavedModel legacy init op collection key. Used in v1 SavedModels. constexpr char kSavedModelLegacyInitOpKey[] = "legacy_init_op"; -/// SavedModel main op key. +/// SavedModel main op collection key. Used in v1 SavedModels. constexpr char kSavedModelMainOpKey[] = "saved_model_main_op"; /// Directory in which to save the SavedModel variables. @@ -45,6 +45,11 @@ constexpr char kSavedModelVariablesDirectory[] = "variables"; /// SavedModel variables filename. constexpr char kSavedModelVariablesFilename[] = "variables"; +/// SavedModel SignatureDef keys for the initialization and train ops. Used in +/// V2 SavedModels. +constexpr char kSavedModelInitOpSignatureKey[] = "__saved_model_init_op"; +constexpr char kSavedModelTrainOpSignatureKey[] = "__saved_model_train_op"; + } // namespace tensorflow #endif // TENSORFLOW_CC_SAVED_MODEL_CONSTANTS_H_ diff --git a/tensorflow/cc/saved_model/loader.cc b/tensorflow/cc/saved_model/loader.cc index c6abe2f41b9b5ec2faee6f65b429ff606f8ac08e..10f7abf09e925c0c31cfd595ecee4605f189476f 100644 --- a/tensorflow/cc/saved_model/loader.cc +++ b/tensorflow/cc/saved_model/loader.cc @@ -21,6 +21,7 @@ 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" @@ -42,9 +43,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) { @@ -122,34 +142,54 @@ Status RunOnce(const RunOptions& run_options, return run_status; } -bool HasMainOp(const MetaGraphDef& meta_graph_def) { +// RunInitOp will return OK if the initialization op was run successfully. +// An empty init_op_name indicates that there are no init ops to run. +Status RunInitOp(const RunOptions& run_options, const string& export_dir, + const MetaGraphDef& meta_graph_def, + const std::vector& asset_file_defs, + Session* session, const string& init_op_name) { + if (!init_op_name.empty()) { + LOG(INFO) << "Running initialization op on SavedModel bundle."; + std::vector> inputs; + AddAssetsTensorsToInputs(export_dir, asset_file_defs, &inputs); + RunMetadata run_metadata; + return RunOnce(run_options, inputs, {}, {init_op_name}, + nullptr /* outputs */, &run_metadata, session); + } + return Status::OK(); +} + +// A SavedModel may store the name of the initialization op to run in the +// in the SignatureDef (v2) or a collection (v1). If an init_op collection +// exists, then the collection must contain exactly one op. +Status GetInitOp(const string& export_dir, const MetaGraphDef& meta_graph_def, + string* init_op_name) { + const auto& sig_def_map = meta_graph_def.signature_def(); + const auto& init_op_sig_it = + meta_graph_def.signature_def().find(kSavedModelInitOpSignatureKey); + if (init_op_sig_it != sig_def_map.end()) { + *init_op_name = init_op_sig_it->second.outputs() + .find(kSavedModelInitOpSignatureKey) + ->second.name(); + return Status::OK(); + } + const auto& collection_def_map = meta_graph_def.collection_def(); + string init_op_collection_key; if (collection_def_map.find(kSavedModelMainOpKey) != collection_def_map.end()) { - return true; + init_op_collection_key = kSavedModelMainOpKey; + } else { + init_op_collection_key = kSavedModelLegacyInitOpKey; } - return false; -} -Status RunMainOp(const RunOptions& run_options, const string& export_dir, - const MetaGraphDef& meta_graph_def, - const std::vector& asset_file_defs, - Session* session, const string& main_op_key) { - LOG(INFO) << "Running MainOp with key " << main_op_key - << " on SavedModel bundle."; - const auto& collection_def_map = meta_graph_def.collection_def(); - const auto main_op_it = collection_def_map.find(main_op_key); - if (main_op_it != collection_def_map.end()) { - if (main_op_it->second.node_list().value_size() != 1) { + const auto init_op_it = collection_def_map.find(init_op_collection_key); + if (init_op_it != collection_def_map.end()) { + if (init_op_it->second.node_list().value_size() != 1) { return errors::FailedPrecondition( strings::StrCat("Expected exactly one main op in : ", export_dir)); } - std::vector> inputs; - AddAssetsTensorsToInputs(export_dir, asset_file_defs, &inputs); - RunMetadata run_metadata; - const StringPiece main_op_name = main_op_it->second.node_list().value(0); - return RunOnce(run_options, inputs, {}, {string(main_op_name)}, - nullptr /* outputs */, &run_metadata, session); + *init_op_name = init_op_it->second.node_list().value(0); } return Status::OK(); } @@ -193,6 +233,15 @@ Status RunRestore(const RunOptions& run_options, const string& export_dir, Status GetAssetFileDefs(const MetaGraphDef& meta_graph_def, std::vector* asset_file_defs) { + // With SavedModel v2, we write asset file def into metagraph instead of + // collection, so read from metagraph first. + if (meta_graph_def.asset_file_def_size() > 0) { + for (const auto& asset : meta_graph_def.asset_file_def()) { + asset_file_defs->push_back(asset); + } + return Status::OK(); + } + // Fall back to read from collection to be backward compatible with v1. const auto& collection_def_map = meta_graph_def.collection_def(); const auto assets_it = collection_def_map.find(kSavedModelAssetsKey); if (assets_it == collection_def_map.end()) { @@ -213,6 +262,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)); @@ -227,15 +277,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())); - if (HasMainOp(bundle->meta_graph_def)) { - TF_RETURN_IF_ERROR(RunMainOp(run_options, export_dir, - bundle->meta_graph_def, asset_file_defs, - bundle->session.get(), kSavedModelMainOpKey)); - } else { - TF_RETURN_IF_ERROR(RunMainOp( - run_options, export_dir, bundle->meta_graph_def, asset_file_defs, - bundle->session.get(), kSavedModelLegacyInitOpKey)); - } + // 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(); } @@ -249,16 +307,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()) { @@ -266,7 +318,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/saved_model/loader_test.cc b/tensorflow/cc/saved_model/loader_test.cc index 72b8bc18710b0ee77cb01ed3ad0c2abb5183efb2..597e42bb65ab5536664089f7e65ec52d77fc8f23 100644 --- a/tensorflow/cc/saved_model/loader_test.cc +++ b/tensorflow/cc/saved_model/loader_test.cc @@ -36,6 +36,8 @@ constexpr char kTestDataMainOp[] = "cc/saved_model/testdata/half_plus_two_main_op/00000123"; constexpr char kTestDataSharded[] = "cc/saved_model/testdata/half_plus_two/00000123"; +constexpr char kTestDataInitOpV2[] = + "cc/saved_model/testdata/half_plus_two_v2/00000123"; class LoaderTest : public ::testing::Test { protected: @@ -227,5 +229,17 @@ TEST_F(LoaderTest, MaybeSavedModelDirectory) { EXPECT_FALSE(MaybeSavedModelDirectory(invalid_export_dir)); } +TEST_F(LoaderTest, SavedModelInitOpV2Format) { + SavedModelBundle bundle; + SessionOptions session_options; + RunOptions run_options; + + const string export_dir = + io::JoinPath(testing::TensorFlowSrcRoot(), kTestDataInitOpV2); + TF_ASSERT_OK(LoadSavedModel(session_options, run_options, export_dir, + {kSavedModelTagServe}, &bundle)); + CheckSavedModelBundle(export_dir, bundle); +} + } // namespace } // namespace tensorflow diff --git a/tensorflow/cc/saved_model/testdata/half_plus_two_v2/00000123/assets/foo.txt b/tensorflow/cc/saved_model/testdata/half_plus_two_v2/00000123/assets/foo.txt new file mode 100644 index 0000000000000000000000000000000000000000..f9ff036688007836524129e23f5cf82edd1e8910 --- /dev/null +++ b/tensorflow/cc/saved_model/testdata/half_plus_two_v2/00000123/assets/foo.txt @@ -0,0 +1 @@ +asset-file-contents \ No newline at end of file diff --git a/tensorflow/cc/saved_model/testdata/half_plus_two_v2/00000123/saved_model.pb b/tensorflow/cc/saved_model/testdata/half_plus_two_v2/00000123/saved_model.pb new file mode 100644 index 0000000000000000000000000000000000000000..a10bbf8fb6bca0fcee6414b2927d2f706de85ebc Binary files /dev/null and b/tensorflow/cc/saved_model/testdata/half_plus_two_v2/00000123/saved_model.pb differ diff --git a/tensorflow/cc/saved_model/testdata/half_plus_two_v2/00000123/variables/variables.data-00000-of-00001 b/tensorflow/cc/saved_model/testdata/half_plus_two_v2/00000123/variables/variables.data-00000-of-00001 new file mode 100644 index 0000000000000000000000000000000000000000..15b75d6ef6bffc336d138d923badb3928b8c4c13 Binary files /dev/null and b/tensorflow/cc/saved_model/testdata/half_plus_two_v2/00000123/variables/variables.data-00000-of-00001 differ diff --git a/tensorflow/cc/saved_model/testdata/half_plus_two_v2/00000123/variables/variables.index b/tensorflow/cc/saved_model/testdata/half_plus_two_v2/00000123/variables/variables.index new file mode 100644 index 0000000000000000000000000000000000000000..7ec9fb4fe2dd21d0a6c324aecd7658fc37cf2326 Binary files /dev/null and b/tensorflow/cc/saved_model/testdata/half_plus_two_v2/00000123/variables/variables.index differ diff --git a/tensorflow/compat_template_v1.__init__.py b/tensorflow/compat_template_v1.__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9549a71c41a0ba2aac58abd8cfb182aa4eaf3b4f --- /dev/null +++ b/tensorflow/compat_template_v1.__init__.py @@ -0,0 +1,37 @@ +# 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 + +# 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._v1.estimator')) +_component_api_helper.package_hook( + parent_package_str=__name__, + child_package_str=('tensorflow.python.keras.api._v1.keras')) +from tensorflow.python.platform import flags # pylint: disable=g-import-not-at-top +app.flags = flags # pylint: disable=undefined-variable diff --git a/tensorflow/compiler/aot/BUILD b/tensorflow/compiler/aot/BUILD index 6c29f09cde7ee17c11cb44ce48d8e9128daae4d0..16151e77737429f4fbf690fc34b12a70bacebdc4 100644 --- a/tensorflow/compiler/aot/BUILD +++ b/tensorflow/compiler/aot/BUILD @@ -93,7 +93,7 @@ cc_library( ":tfcompile_lib", "//tensorflow/compiler/tf2xla:tf2xla_proto", "//tensorflow/compiler/tf2xla:tf2xla_util", - "//tensorflow/compiler/xla/legacy_flags:debug_options_flags", + "//tensorflow/compiler/xla:debug_options_flags", "//tensorflow/compiler/xla/service:compiler", "//tensorflow/core:core_cpu", "//tensorflow/core:core_cpu_internal", diff --git a/tensorflow/compiler/aot/codegen.cc b/tensorflow/compiler/aot/codegen.cc index b17bc658fa06b9feb7edb292bd89ef31e6309169..d016632da2a9d7c2c2f81c02dd573787a0502923 100644 --- a/tensorflow/compiler/aot/codegen.cc +++ b/tensorflow/compiler/aot/codegen.cc @@ -129,7 +129,7 @@ Status AddRewritesForShape(int i, const xla::Shape& shape, TF_RETURN_IF_ERROR(XLATypeToCpp(shape.element_type(), &type)); std::vector dim_vars; string dim_sizes, indices; - if (xla::ShapeUtil::Rank(shape) == 0 || + if (shape.rank() == 0 || (shape.dimensions_size() == 1 && shape.dimensions(0) == 1)) { dim_sizes = "[1]"; indices = "[0]"; @@ -164,7 +164,8 @@ string RewriteWithName(const string& name, string code, } // Generate methods for args (inputs). -Status GenArgMethods(const tf2xla::Config& config, const xla::ProgramShape& ps, +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) { @@ -174,9 +175,10 @@ Status GenArgMethods(const tf2xla::Config& config, const xla::ProgramShape& ps, } for (int i = 0; i < num_args; ++i) { std::vector> rewrites; - TF_RETURN_IF_ERROR(AddRewritesForShape(i, ps.parameters(i), &rewrites)); + TF_RETURN_IF_ERROR( + AddRewritesForShape(i, xla::Shape(ps.parameters(i)), &rewrites)); const string code = R"( - void set_arg{{NAME}}_data(void* data) { + void set_arg{{NAME}}_data(const void* data) { set_arg_data({{I}}, data); } {{TYPE}}* arg{{NAME}}_data() { @@ -204,7 +206,7 @@ Status GenArgMethods(const tf2xla::Config& config, const xla::ProgramShape& ps, // Generate methods for results (outputs). Status GenResultMethods(const tf2xla::Config& config, - const xla::ProgramShape& ps, string* methods) { + const xla::ProgramShapeProto& ps, string* methods) { if (ps.result().element_type() != xla::TUPLE) { // The XlaCompiler we use to build the xla computation always generates a // tuple result, and we rely on this to simplify code generation. @@ -217,8 +219,8 @@ Status GenResultMethods(const tf2xla::Config& config, } for (int i = 0; i < ps.result().tuple_shapes_size(); ++i) { std::vector> rewrites; - TF_RETURN_IF_ERROR( - AddRewritesForShape(i, ps.result().tuple_shapes(i), &rewrites)); + TF_RETURN_IF_ERROR(AddRewritesForShape( + i, xla::Shape(ps.result().tuple_shapes(i)), &rewrites)); string code = R"( {{TYPE}}* result{{NAME}}_data() { return static_cast<{{TYPE}}*>(result_data({{I}})); @@ -336,7 +338,7 @@ Status GenerateHeader(const CodegenOpts& opts, const tf2xla::Config& config, ExtractEntryParamBufferInfos(buffer_infos); std::vector buffer_infos_for_temps = ExtractTempBufferInfos(buffer_infos); - const xla::ProgramShape& ps = compile_result.program_shape; + const xla::ProgramShapeProto& ps = compile_result.program_shape; string methods_arg, methods_result; TF_RETURN_IF_ERROR(GenArgMethods(config, ps, compile_result, &methods_arg)); TF_RETURN_IF_ERROR(GenResultMethods(config, ps, &methods_result)); @@ -382,8 +384,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 @@ -447,7 +450,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}}; @@ -462,16 +465,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; }(); @@ -548,8 +552,8 @@ class {{CLASS}} : public tensorflow::XlaCompiledCpuFunction { static const char** StaticResultNames() {{RESULT_NAMES_CODE}} // Shape of the args and results. - static const xla::ProgramShape* StaticProgramShape() { - static const xla::ProgramShape* kShape = {{PROGRAM_SHAPE_SHIM_EXPRESSION}}; + static const xla::ProgramShapeProto* StaticProgramShape() { + static const xla::ProgramShapeProto* kShape = {{PROGRAM_SHAPE_SHIM_EXPRESSION}}; return kShape; } @@ -587,7 +591,7 @@ class {{CLASS}} : public tensorflow::XlaCompiledCpuFunction { {"{{METHODS_RESULT}}\n", methods_result}, {"{{NS_END}}\n", ns_end}, {"{{NS_START}}\n", ns_start}, - {"{{PROGRAM_SHAPE}}", xla::ShapeUtil::HumanString(ps)}, + {"{{PROGRAM_SHAPE}}", xla::ShapeUtil::HumanString(xla::ProgramShape(ps))}, {"{{PROGRAM_SHAPE_SHIM_EXPRESSION}}", metadata_result.program_shape_access_shim}, {"{{RESULT_INDEX}}", absl::StrCat(result_index)}, @@ -615,11 +619,11 @@ static string CreateUniqueIdentifier(const CodegenOpts& opts, Status GenerateMetadata(const CodegenOpts& opts, const CompileResult& compile_result, MetadataResult* metadata_result) { - std::unique_ptr program_shape; + std::unique_ptr program_shape; if (opts.gen_program_shape) { program_shape = - absl::make_unique(compile_result.program_shape); + absl::make_unique(compile_result.program_shape); // The parameter names are currently meaningless, and redundant with the // rest of our metadata, so clear them out to avoid confusion and save @@ -631,8 +635,8 @@ Status GenerateMetadata(const CodegenOpts& opts, // a shim that evaluates to nullptr, which is what we want. ProtobufToEmbed program_shape_protobuf{ - CreateUniqueIdentifier(opts, "ProgramShape"), "xla::ProgramShape", - program_shape.get()}; + CreateUniqueIdentifier(opts, "ProgramShapeProto"), + "xla::ProgramShapeProto", program_shape.get()}; ProtobufToEmbed hlo_profile_printer_data_protobuf{ CreateUniqueIdentifier(opts, "HloProfilePrinterData"), diff --git a/tensorflow/compiler/aot/codegen.h b/tensorflow/compiler/aot/codegen.h index 90410c46a8e36e44454f1219ad76d0fb0937070d..9485e86b10e225a3c9c12eafd9905bdf7c15c9fa 100644 --- a/tensorflow/compiler/aot/codegen.h +++ b/tensorflow/compiler/aot/codegen.h @@ -57,7 +57,7 @@ struct MetadataResult { std::vector header_variable_decls; // program_shape_access_shim is a C++ expression that constructs the - // xla::ProgramShape instance for the CompileResult passed to + // xla::ProgramShapeProto instance for the CompileResult passed to // GenerateMetadata. string program_shape_access_shim; diff --git a/tensorflow/compiler/aot/codegen_test.cc b/tensorflow/compiler/aot/codegen_test.cc index bb288d23000527be74f01630d20bbf82e50007ce..c1788ca32a1d099284eeb870f9513891051fd29e 100644 --- a/tensorflow/compiler/aot/codegen_test.cc +++ b/tensorflow/compiler/aot/codegen_test.cc @@ -181,13 +181,15 @@ TEST(CodegenTest, Golden) { BufferInfo::MakeEntryParameter(/*size=*/96, /*param_number=*/1), BufferInfo::MakeTempBuffer(3), BufferInfo::MakeTempBuffer(120)}, 5, {})); - compile_result.program_shape = xla::ShapeUtil::MakeProgramShape( - { - xla::ShapeUtil::MakeShape(xla::F32, {1, 2}), - xla::ShapeUtil::MakeShape(xla::S64, {3, 4}), - }, - xla::ShapeUtil::MakeTupleShape( - {xla::ShapeUtil::MakeShape(xla::U32, {5, 6})})); + compile_result.program_shape = + xla::ShapeUtil::MakeProgramShape( + { + xla::ShapeUtil::MakeShape(xla::F32, {1, 2}), + xla::ShapeUtil::MakeShape(xla::S64, {3, 4}), + }, + xla::ShapeUtil::MakeTupleShape( + {xla::ShapeUtil::MakeShape(xla::U32, {5, 6})})) + .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 e4d8a02877c75fa72c5747650ab9c7ac229955b3..35994fc785d3e1d5e883c49bec96de315e189d2e 100644 --- a/tensorflow/compiler/aot/codegen_test_h.golden +++ b/tensorflow/compiler/aot/codegen_test_h.golden @@ -22,7 +22,7 @@ extern "C" void entry_point( void* result, const xla::ExecutableRunOptions* run_options, const void** args, void** temps, tensorflow::int64* profile_counters); -extern "C" char __tfcompile_foo_bar_MyClass_ProgramShape_protobuf_array_contents[]; +extern "C" char __tfcompile_foo_bar_MyClass_ProgramShapeProto_protobuf_array_contents[]; namespace foo { @@ -59,7 +59,7 @@ namespace bar { // 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; }(); @@ -114,7 +115,7 @@ class MyClass : public tensorflow::XlaCompiledCpuFunction { // with dim indices specifying which value. No bounds checking is performed // on dim indices. - void set_arg0_data(void* data) { + void set_arg0_data(const void* data) { set_arg_data(0, data); } float* arg0_data() { @@ -132,7 +133,7 @@ class MyClass : public tensorflow::XlaCompiledCpuFunction { arg_data(0)))[dim0][dim1]; } - void set_arg_myfeed_data(void* data) { + void set_arg_myfeed_data(const void* data) { set_arg_data(0, data); } float* arg_myfeed_data() { @@ -150,7 +151,7 @@ class MyClass : public tensorflow::XlaCompiledCpuFunction { arg_data(0)))[dim0][dim1]; } - void set_arg1_data(void* data) { + void set_arg1_data(const void* data) { set_arg_data(1, data); } tensorflow::int64* arg1_data() { @@ -253,10 +254,10 @@ class MyClass : public tensorflow::XlaCompiledCpuFunction { } // Shape of the args and results. - static const xla::ProgramShape* StaticProgramShape() { - static const xla::ProgramShape* kShape = []() { - xla::ProgramShape* proto = new xla::ProgramShape; - proto->ParseFromArray(&__tfcompile_foo_bar_MyClass_ProgramShape_protobuf_array_contents[0], 52); + 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); return proto; }(); return kShape; diff --git a/tensorflow/compiler/aot/codegen_test_o.golden b/tensorflow/compiler/aot/codegen_test_o.golden index eb001c5d45bdfefc76629d7303d89f5480432235..7f7b96428572705f30144e6c95cd4cf9c44ce2a3 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 2b5f97b34cd928d32eb220536342c715d91d45bb..9fc223bdc7c0e207ce2005cb86250aa77e709df8 100644 --- a/tensorflow/compiler/aot/compile.cc +++ b/tensorflow/compiler/aot/compile.cc @@ -56,17 +56,23 @@ Status CompileXla(xla::CompileOnlyClient* client, return errors::Unknown("Couldn't get XLA program shape: ", pshape_or.status().error_message()); } - compile_result->program_shape = *pshape_or.ValueOrDie(); - xla::ProgramShape* pshape = &compile_result->program_shape; - std::vector arg_layouts; - arg_layouts.reserve(pshape->parameters_size()); + compile_result->program_shape = pshape_or.ValueOrDie()->ToProto(); + xla::ProgramShapeProto* pshape = &compile_result->program_shape; + + // AotXlaComputationInstance::argument_layouts is a vector of Shape + // pointers. Accumulate the Shape objects themselves in a separate vector + // while building the vector of pointers. + std::vector arg_layout_ptrs(pshape->parameters_size()); + std::vector arg_layouts(pshape->parameters_size()); for (int i = 0; i < pshape->parameters_size(); ++i) { - arg_layouts.push_back(pshape->mutable_parameters(i)); + arg_layouts[i] = xla::Shape(*pshape->mutable_parameters(i)); + arg_layout_ptrs[i] = &arg_layouts[i]; } xla::CompileOnlyClient::AotXlaComputationInstance instance; instance.computation = &computation; - instance.argument_layouts = std::move(arg_layouts); - instance.result_layout = &pshape->result(); + instance.argument_layouts = std::move(arg_layout_ptrs); + xla::Shape result_shape(pshape->result()); + instance.result_layout = &result_shape; xla::StatusOr>> aot_or = client->CompileAheadOfTime({instance}, aot_opts); if (!aot_or.ok()) { diff --git a/tensorflow/compiler/aot/compile.h b/tensorflow/compiler/aot/compile.h index e03c5b1aa77c1262ed903aae3072ef65f34d80a2..ee7bb26fabd2d897b85b62f38778ecbfe2238eb6 100644 --- a/tensorflow/compiler/aot/compile.h +++ b/tensorflow/compiler/aot/compile.h @@ -33,9 +33,9 @@ namespace tfcompile { struct CompileResult { // Contains object file and meta-info. std::unique_ptr aot; - xla::ProgramShape program_shape; // Static shape of args and results. - string entry_point; // Name of generated function. - int pointer_size = 0; // Size of a pointer in bytes. + xla::ProgramShapeProto program_shape; // Static shape of args and results. + string entry_point; // Name of generated function. + int pointer_size = 0; // Size of a pointer in bytes. }; // CompileGraph compiles the graph_def into an object file containing a function diff --git a/tensorflow/compiler/aot/tests/make_test_graphs.py b/tensorflow/compiler/aot/tests/make_test_graphs.py index 64b861a73091642b03573543a5c55618bf33915d..7bac79ec062af7e790134286e34eda4e123e138a 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) diff --git a/tensorflow/compiler/aot/tests/tfcompile_test.cc b/tensorflow/compiler/aot/tests/tfcompile_test.cc index f10852c7850f61bfd8b99fa9f1648202d182085e..4dd79e5882d7da61be029735ef2b165908c599f9 100644 --- a/tensorflow/compiler/aot/tests/tfcompile_test.cc +++ b/tensorflow/compiler/aot/tests/tfcompile_test.cc @@ -526,13 +526,15 @@ TEST(TFCompileTest, ProgramShape) { // muladd has the program shape defined. MatMulAndAddComp muladd; - const xla::ProgramShape* muladd_shape = muladd.ProgramShape(); + const xla::ProgramShapeProto* muladd_shape = muladd.ProgramShape(); ASSERT_TRUE(muladd_shape != nullptr); ASSERT_EQ(muladd_shape->parameters_size(), 2); - EXPECT_TRUE(ShapeUtil::Compatible(muladd_shape->parameters(0), f32_2x2)); - EXPECT_TRUE(ShapeUtil::Compatible(muladd_shape->parameters(1), f32_2x2)); + EXPECT_TRUE( + ShapeUtil::Compatible(xla::Shape(muladd_shape->parameters(0)), f32_2x2)); + EXPECT_TRUE( + ShapeUtil::Compatible(xla::Shape(muladd_shape->parameters(1)), f32_2x2)); - const xla::Shape& muladd_result = muladd_shape->result(); + const xla::Shape muladd_result(muladd_shape->result()); ASSERT_EQ(muladd_result.element_type(), xla::TUPLE); ASSERT_EQ(ShapeUtil::TupleElementCount(muladd_result), 2); const xla::Shape& muladd_result0 = diff --git a/tensorflow/compiler/aot/tfcompile.bzl b/tensorflow/compiler/aot/tfcompile.bzl index 2dc3e8c9113b37bf9d575ad66783f4ab49478af4..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 @@ -283,7 +283,7 @@ def tf_library( ) # Variables used for gen_test and gen_benchmark. - cpp_class_split = cpp_class.rsplit("::", maxsplit = 2) + cpp_class_split = cpp_class.rsplit("::", 2) if len(cpp_class_split) == 1: no_ns_name = cpp_class_split[0] else: diff --git a/tensorflow/compiler/aot/tfcompile_main.cc b/tensorflow/compiler/aot/tfcompile_main.cc index b95b063348c5cdfdcaed635ba527e9f0bfd6092d..0b6ab7e723d6e3a55da2f1c30b75f44cbdaa75bb 100644 --- a/tensorflow/compiler/aot/tfcompile_main.cc +++ b/tensorflow/compiler/aot/tfcompile_main.cc @@ -26,7 +26,7 @@ limitations under the License. #include "tensorflow/compiler/aot/flags.h" #include "tensorflow/compiler/tf2xla/tf2xla.pb.h" #include "tensorflow/compiler/tf2xla/tf2xla_util.h" -#include "tensorflow/compiler/xla/legacy_flags/debug_options_flags.h" +#include "tensorflow/compiler/xla/debug_options_flags.h" #include "tensorflow/compiler/xla/service/compiler.h" #include "tensorflow/core/framework/function.h" #include "tensorflow/core/framework/graph.pb.h" @@ -103,7 +103,7 @@ Status Main(const MainFlags& flags) { return errors::InvalidArgument("Must specify --cpp_class"); } codegen_opts.gen_hlo_profile_printer_data = - xla::legacy_flags::GetDebugOptionsFromFlags().xla_hlo_profile(); + xla::GetDebugOptionsFromFlags().xla_hlo_profile(); TF_RETURN_IF_ERROR(ParseCppClass(flags.cpp_class, &codegen_opts.class_name, &codegen_opts.namespaces)); @@ -132,10 +132,14 @@ int main(int argc, char** argv) { std::vector flag_list; AppendMainFlags(&flag_list, &flags); - xla::legacy_flags::AppendDebugOptionsFlags(&flag_list); + xla::AppendDebugOptionsFlags(&flag_list); 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 77cae5668b6732eb1eda68daac017845477b7680..55e2e6d11f7a0984d2e1a40990c3d3db85bd1ff4 100644 --- a/tensorflow/compiler/jit/BUILD +++ b/tensorflow/compiler/jit/BUILD @@ -21,10 +21,8 @@ package( ) load("//tensorflow:tensorflow.bzl", "cc_header_only_library") -load("//tensorflow:tensorflow.bzl", "tf_kernel_library") load("//tensorflow:tensorflow.bzl", "tf_cc_test") load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda") -load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda_is_configured") load("//tensorflow:tensorflow.bzl", "tf_cuda_cc_test") load("//tensorflow:tensorflow.bzl", "tf_custom_op_py_library") @@ -39,7 +37,7 @@ cc_library( ":xla_cpu_device", ":xla_cpu_jit", "//tensorflow/compiler/plugin", - ] + if_cuda_is_configured([ + ] + if_cuda([ ":xla_gpu_device", ":xla_gpu_jit", ]), @@ -52,6 +50,8 @@ cc_library( deps = [ ":jit_compilation_passes", "//tensorflow/compiler/jit/kernels:xla_ops", + "//tensorflow/compiler/tf2xla/kernels:xla_cpu_only_ops", + "//tensorflow/compiler/tf2xla/kernels:xla_dummy_ops", "//tensorflow/compiler/tf2xla/kernels:xla_ops", "//tensorflow/compiler/xla/service:cpu_plugin", ], @@ -65,6 +65,7 @@ cc_library( ":jit_compilation_passes", "//tensorflow/compiler/jit/kernels:xla_ops", "//tensorflow/compiler/tf2xla/kernels:xla_ops", + "//tensorflow/compiler/tf2xla/kernels:xla_dummy_ops", "//tensorflow/compiler/xla/service:gpu_plugin", ]), alwayslink = 1, @@ -75,10 +76,11 @@ cc_library( srcs = ["xla_cpu_device.cc"], visibility = [":friends"], deps = [ + ":create_xla_launch_op", # buildcleaner: keep + ":flags", ":jit_compilation_passes", ":xla_device", "//tensorflow/compiler/jit/kernels:xla_ops", - "//tensorflow/compiler/jit/legacy_flags:xla_device_flags", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/compiler/tf2xla/kernels:xla_ops", "//tensorflow/compiler/xla/service:cpu_plugin", # buildcleaner: keep @@ -94,6 +96,7 @@ cc_library( srcs = ["xla_gpu_device.cc"], visibility = [":friends"], deps = [ + ":create_xla_launch_op", # buildcleaner: keep ":jit_compilation_passes", ":xla_device", "//tensorflow/compiler/jit/kernels:xla_ops", @@ -103,6 +106,7 @@ cc_library( "//tensorflow/core:core_cpu_internal", "//tensorflow/core:lib", "@com_google_absl//absl/memory", + "@com_google_absl//absl/strings", ], alwayslink = 1, ) @@ -171,12 +175,18 @@ 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: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:protos_all_cc", + "//tensorflow/core:resource_variable_ops_op_lib", "//tensorflow/core:stream_executor_no_cuda", "//tensorflow/core/kernels:cast_op", "//tensorflow/core/kernels:constant_op", @@ -190,11 +200,14 @@ cc_library( "//tensorflow/core/kernels:resource_variable_ops", "//tensorflow/core/kernels:sendrecv_ops", "//tensorflow/core/kernels:shape_ops", + "//tensorflow/core/kernels:stack", "//tensorflow/core/kernels:variable_ops", "//tensorflow/core/kernels/data:generator_dataset_op", "//tensorflow/core/kernels/data:iterator_ops", "//tensorflow/core/kernels/data:prefetch_dataset_op", "@com_google_absl//absl/memory", + "@com_google_absl//absl/synchronization", + "@com_google_absl//absl/types:optional", ], ) @@ -207,6 +220,18 @@ cc_library( # Internal targets below this point. +cc_library( + name = "flags", + srcs = ["flags.cc"], + hdrs = ["flags.h"], + visibility = [":friends"], + deps = [ + "//tensorflow/compiler/xla:parse_flags_from_env", + "//tensorflow/core:framework_internal", + "//tensorflow/core:lib", + ], +) + cc_library( name = "common", srcs = [ @@ -239,6 +264,8 @@ cc_library( "//tensorflow/core:lib_internal", "//tensorflow/core:protos_all_cc", "//tensorflow/core/kernels:variable_ops", + "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/memory", ], ) @@ -251,6 +278,8 @@ 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", "//tensorflow/core:core_cpu", @@ -261,6 +290,22 @@ cc_library( "//tensorflow/core:protos_all_cc", "//tensorflow/core/kernels:variable_ops", "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/types:optional", + "@com_google_absl//absl/types:span", + ], +) + +tf_cc_test( + name = "xla_compilation_cache_test", + srcs = [ + "xla_compilation_cache_test.cc", + ], + deps = [ + ":xla_compilation_cache", + "//tensorflow/compiler/tf2xla:common", + "//tensorflow/core:test", + "//tensorflow/core:test_main", ], ) @@ -447,6 +492,7 @@ cc_library( "encapsulate_subgraphs_pass.cc", "encapsulate_xla_computations_pass.cc", "extract_outside_compilation_pass.cc", + "increase_dynamism_for_auto_jit_pass.cc", "mark_for_compilation_pass.cc", "mark_for_compilation_pass_test_helper.cc", "partially_decluster_pass.cc", @@ -457,6 +503,7 @@ cc_library( "encapsulate_subgraphs_pass.h", "encapsulate_xla_computations_pass.h", "extract_outside_compilation_pass.h", + "increase_dynamism_for_auto_jit_pass.h", "mark_for_compilation_pass.h", "mark_for_compilation_pass_test_helper.h", "partially_decluster_pass.h", @@ -464,6 +511,7 @@ cc_library( deps = [ ":common", ":encapsulate_util", + ":flags", ":shape_inference_helpers", ":union_find", ":xla_cluster_util", @@ -471,14 +519,14 @@ cc_library( "//tensorflow/cc:ops", "//tensorflow/cc:scope_internal", "//tensorflow/compiler/jit/graphcycles", - "//tensorflow/compiler/jit/legacy_flags:build_xla_ops_pass_flags", - "//tensorflow/compiler/jit/legacy_flags:mark_for_compilation_pass_flags", "//tensorflow/compiler/jit/ops:xla_ops", "//tensorflow/compiler/tf2xla:dump_graph", "//tensorflow/compiler/tf2xla:resource_operation_table", + "//tensorflow/compiler/tf2xla:side_effect_util", "//tensorflow/compiler/tf2xla:tf2xla_util", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/compiler/tf2xla/cc:xla_jit_ops", + "//tensorflow/compiler/tf2xla/cc:xla_ops", "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:util", "//tensorflow/core:core_cpu", @@ -492,8 +540,10 @@ cc_library( "@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", + "@com_google_absl//absl/types:optional", ], ) @@ -518,25 +568,6 @@ cc_library( hdrs = ["union_find.h"], ) -cc_library( - name = "producer_consumer_queue", - hdrs = ["producer_consumer_queue.h"], - deps = ["//tensorflow/core:lib"], -) - -tf_cc_test( - name = "producer_consumer_queue_test", - size = "small", - srcs = ["producer_consumer_queue_test.cc"], - deps = [ - ":producer_consumer_queue", - "//tensorflow/core:lib", - "//tensorflow/core:test", - "//tensorflow/core:test_main", - "//tensorflow/core:testlib", - ], -) - tf_cc_test( name = "deadness_analysis_test", size = "small", @@ -575,6 +606,7 @@ tf_cc_test( "encapsulate_subgraphs_pass_test.cc", "encapsulate_xla_computations_pass_test.cc", "extract_outside_compilation_pass_test.cc", + "increase_dynamism_for_auto_jit_pass_test.cc", "mark_for_compilation_pass_test.cc", "partially_decluster_pass_test.cc", ], @@ -589,25 +621,29 @@ tf_cc_test( "//tensorflow/cc:cc_ops", "//tensorflow/cc:cc_ops_internal", "//tensorflow/cc:function_ops", + "//tensorflow/cc:functional_ops", "//tensorflow/cc:ops", "//tensorflow/cc:resource_variable_ops", "//tensorflow/cc:scope", "//tensorflow/cc:sendrecv_ops", "//tensorflow/compiler/jit/kernels:xla_ops", + "//tensorflow/compiler/tf2xla:side_effect_util", "//tensorflow/compiler/tf2xla:test_util", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/compiler/tf2xla/cc:xla_jit_ops", "//tensorflow/compiler/tf2xla/cc:xla_ops", + "//tensorflow/compiler/tf2xla/kernels:xla_dummy_ops", "//tensorflow/compiler/tf2xla/kernels:xla_ops", + "//tensorflow/compiler/xla:test", "//tensorflow/core:core_cpu", "//tensorflow/core:framework", "//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", @@ -640,31 +676,6 @@ tf_cc_test( ], ) -tf_cc_test( - name = "xla_launch_util_test", - size = "small", - srcs = ["xla_launch_util_test.cc"], - deps = [ - ":common", - ":xla_compilation_cache", - ":xla_launch_util", - ":xla_tensor", - "//tensorflow/compiler/tf2xla:common", - "//tensorflow/compiler/tf2xla:xla_compiler", - "//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:gpu_runtime", - "//tensorflow/core:lib", - "//tensorflow/core:lib_internal", - "//tensorflow/core:protos_all_cc", - "//tensorflow/core:test", - "//tensorflow/core/kernels:variable_ops", - ], -) - cc_library( name = "xla_fusion_optimizer", srcs = ["xla_fusion_optimizer.cc"], @@ -739,7 +750,10 @@ tf_custom_op_py_library( visibility = [ ":friends", ], - deps = ["//tensorflow/compiler/jit/ops:xla_ops_wrapper_py"], + deps = [ + "//tensorflow/compiler/jit/ops:xla_ops_grad", + "//tensorflow/compiler/jit/ops:xla_ops_wrapper_py", + ], ) # This target can be used by XLA device plugins to prevent circular dependencies, and provides access to all of the required headers for building a device library. diff --git a/tensorflow/compiler/jit/build_xla_ops_pass.cc b/tensorflow/compiler/jit/build_xla_ops_pass.cc index 054f31ba3352b2215e6b0448c8ec8a70cb98b8e5..9f4042630edaec1b9519b6434d859a48372e8b15 100644 --- a/tensorflow/compiler/jit/build_xla_ops_pass.cc +++ b/tensorflow/compiler/jit/build_xla_ops_pass.cc @@ -23,7 +23,7 @@ limitations under the License. #include "tensorflow/cc/ops/control_flow_ops.h" #include "tensorflow/compiler/jit/defs.h" #include "tensorflow/compiler/jit/encapsulate_subgraphs_pass.h" -#include "tensorflow/compiler/jit/legacy_flags/build_xla_ops_pass_flags.h" +#include "tensorflow/compiler/jit/flags.h" #include "tensorflow/compiler/jit/xla_cluster_util.h" #include "tensorflow/compiler/tf2xla/cc/ops/xla_jit_ops.h" #include "tensorflow/compiler/tf2xla/dump_graph.h" @@ -214,7 +214,8 @@ Status NodeRequiresCompilation(Node* n, bool* result) { return errors::Internal("Could not find compilation device ", device_type.type()); } - *result = registration->requires_compilation; + *result = registration->autoclustering_policy == + XlaOpRegistry::AutoclusteringPolicy::kAlways; return Status::OK(); } @@ -319,10 +320,10 @@ Status BuildXlaOpsPass::Run(const GraphOptimizationPassOptions& options) { return IsXlaCompiledKernel(*n); }); - bool lazy_compilation_enabled = enable_lazy_compilation_ - ? *enable_lazy_compilation_ - : legacy_flags::GetBuildXlaOpsPassFlags() - .tf_xla_enable_lazy_compilation; + bool lazy_compilation_enabled = + enable_lazy_compilation_ + ? *enable_lazy_compilation_ + : GetBuildXlaOpsPassFlags().tf_xla_enable_lazy_compilation; for (Node* n : xla_compiled_kernels) { TF_RETURN_IF_ERROR(ReplaceNodeWithXlaCompileAndXlaRun( diff --git a/tensorflow/compiler/jit/build_xla_ops_pass_test.cc b/tensorflow/compiler/jit/build_xla_ops_pass_test.cc index 11df946cc186660242574c2644463a26ead44f1f..390ffa694b6f127544d92f3024a02d877556aacd 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" @@ -42,14 +41,8 @@ class BuildXlaOpsTest : public ::testing::Test { .ok()); } - void TearDown() override { - for (Device* device : devices_) { - delete device; - } - } - private: - std::vector devices_; + std::vector> devices_; }; using ::tensorflow::testing::FindNodeByName; diff --git a/tensorflow/compiler/jit/create_xla_launch_op_test.cc b/tensorflow/compiler/jit/create_xla_launch_op_test.cc index 73866607621cd745f6e640a14405daebf0dd9985..0f872a480f4d4843217f1df3452c4dc62531264e 100644 --- a/tensorflow/compiler/jit/create_xla_launch_op_test.cc +++ b/tensorflow/compiler/jit/create_xla_launch_op_test.cc @@ -59,8 +59,9 @@ class CreateXlaLaunchOpTest : public ::testing::Test { SessionOptions options; auto* device_count = options.config.mutable_device_count(); device_count->insert({"CPU", 1}); + std::vector> devices; TF_CHECK_OK(DeviceFactory::AddDevices( - options, "/job:localhost/replica:0/task:0", &devices_)); + options, "/job:localhost/replica:0/task:0", &devices)); FunctionDefLibrary proto; for (const auto& fdef : flib) { @@ -69,7 +70,7 @@ class CreateXlaLaunchOpTest : public ::testing::Test { lib_def_ = absl::make_unique( OpRegistry::Global(), proto); OptimizerOptions opts; - device_mgr_ = absl::make_unique(devices_); + device_mgr_ = absl::make_unique(std::move(devices)); pflr_ = absl::make_unique( device_mgr_.get(), Env::Default(), TF_GRAPH_DEF_VERSION, lib_def_.get(), opts, /*default_thread_pool=*/nullptr, /*cluster_flr=*/nullptr); @@ -77,7 +78,6 @@ class CreateXlaLaunchOpTest : public ::testing::Test { } FunctionLibraryRuntime* flr_; - std::vector devices_; std::unique_ptr device_mgr_; std::unique_ptr lib_def_; std::unique_ptr pflr_; diff --git a/tensorflow/compiler/jit/deadness_analysis.cc b/tensorflow/compiler/jit/deadness_analysis.cc index e29da8500f9ce0f668373414597753d241f7d973..08fabb3e3085cdfe87125cf54980810d48783ea0 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" @@ -222,29 +225,40 @@ 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. // -// An AndRecurrence with start = S and step = X is printed as {S,&,X} and stands -// for the list of predicates: +// More concretely, an and recurrence {S,&,X} represents the liveness of V +// in the following graph: // -// S, S & GenSym(X,1), S & GenSym(X,1) & GenSym(X,2), ... +// V = Merge(S', V_NextIt) +// V = Op(V, X') +// V_NextIt = NextIteration(V) // -// 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". +// where Predicate(S') = S and Predicate(X') = X. +// +// `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(Predicate* start, Predicate* step, + std::vector frame) + : Predicate(Hash(start, step, frame)), + 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 +269,17 @@ class AndRecurrencePredicate : public Predicate { private: std::array operands_; + std::vector frame_; + + static int64 Hash(Predicate* start, Predicate* step, + const std::vector& frame) { + uint64 frame_hash = 0; + for (const string& sub_frame : frame) { + frame_hash = Hash64Combine(Hash64(sub_frame), frame_hash); + } + return Hash64Combine( + HashPredicateSequence(Kind::kAndRecurrence, {start, step}), frame_hash); + } }; // Represents an uninterpreted symbol in a logical predicate. @@ -281,7 +306,7 @@ 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_; } @@ -333,34 +358,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,16 +418,63 @@ 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) { return std::unique_ptr( @@ -402,7 +498,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 +519,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,6 +559,7 @@ class PredicateFactory { absl::flat_hash_map, HashSignatureForSymbol> interned_symbol_instances_; + int stack_depth_ = 0; }; Predicate* PredicateFactory::MakeInternedAndOr( @@ -466,6 +594,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 +629,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 @@ -525,7 +675,6 @@ Predicate* PredicateFactory::MakeAndOrImpl( op->GetOperands().begin(), op->GetOperands().end()); } else { - std::vector sub_ops_intersection; common_inner_operands.clear(); absl::c_copy_if(op->GetOperands(), std::back_inserter(common_inner_operands), @@ -620,6 +769,7 @@ class DeadnessAnalysisImpl : public DeadnessAnalysis { const Graph& graph_; absl::flat_hash_map predicate_map_; PredicateFactory predicate_factory_; + std::vector control_flow_info_; bool vlog_; }; @@ -662,9 +812,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. @@ -762,6 +915,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, @@ -784,8 +954,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(); @@ -826,8 +998,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(); } @@ -842,8 +1016,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); @@ -893,6 +1069,25 @@ Status DeadnessAnalysisImpl::Populate() { Status DeadnessAnalysisImpl::PopulateWithReversePostOrder( absl::Span rpo) { + std::vector unreachable_nodes; + // Compute the loop structure of the graph. + std::vector control_flow_info; + 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 617e31488c7daeb714c0ff7056b786e4eaf7873f..16ee8f86d55c72785368ac2fd67635eba2fa7cd7 100644 --- a/tensorflow/compiler/jit/deadness_analysis_test.cc +++ b/tensorflow/compiler/jit/deadness_analysis_test.cc @@ -123,11 +123,11 @@ 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::internal::Exit exit(root.WithOpName(prefix + "/exit"), iv.output); + 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"), latch.output_true, increment_by); Output next_iteration = @@ -139,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, @@ -191,7 +191,8 @@ DependentInductionVar CreateDependentLoopInvariantValue( value, frame_name); ops::Merge iv(root.WithOpName(prefix + "/iv"), {enter_value, enter_value}); ops::Switch latch(root.WithOpName(prefix + "/latch"), iv.output, loop_cond); - ops::internal::Exit exit(root.WithOpName(prefix + "/exit"), iv.output); + ops::internal::Exit exit(root.WithOpName(prefix + "/exit"), + latch.output_false); Output next_iteration = ops::NextIteration( root.WithOpName(prefix + "/next_iteration"), latch.output_true); CHECK(root.graph() @@ -513,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,&,*iv1/cond:0} & {#true,&,*iv0/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); @@ -547,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/cond:0 & iv0/iv:0)}"); EXPECT_EQ(predicate_map[ControlOutputFor(dependent_iv1)], - "{#true,&,(*iv0/cond:0 & iv0/iv:0)}"); + "{#true,&,(*iv0/cond:0 & iv0/iv:0)}"); EXPECT_EQ(predicate_map[ControlOutputFor(add0)], - "{#true,&,(*iv0/cond:0 & iv0/iv:0)}"); + "{#true,&,(*iv0/cond:0 & iv0/iv:0)}"); } } @@ -593,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); @@ -636,46 +641,50 @@ 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}"); + "{({#true,&,*iv_outer/cond:0} & " + "*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/cond:0 & " + "iv_inner/iv: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/cond:0 & " + "iv_inner/iv: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/cond:0 & " + "iv_inner/iv: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()); @@ -690,21 +699,76 @@ 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(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})"); + EXPECT_EQ(predicate_map[ControlOutputFor(outer_iv[0])], + "{#true,&,*iv_outer/cond:0}"); + EXPECT_EQ(predicate_map[ControlOutputFor(inner_iv[0])], + "{({#true,&,*iv_outer/cond:0} & " + "*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])], + "{({#true,&,*iv_outer/cond_1:0} & " + "*iv_outer/cond_1:0),&,*iv_inner/cond_1:0}"); + EXPECT_EQ( + predicate_map[ControlOutputFor(add0)], + "({({#true,&,*iv_outer/cond:0} & " + "*iv_outer/cond:0),&,*iv_inner/cond:0} & " + "{({#true,&,*iv_outer/cond_1:0} & " + "*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])], ""); } } @@ -816,5 +880,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 da030b3bcc7aacae2306bec30f4b8927aa042d7c..d0d7a3f3785469acd79a83b6897668f94fc6ea2e 100644 --- a/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc +++ b/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc @@ -779,7 +779,8 @@ Status Encapsulator::Subgraph::RecordArg( if (inserted) { NodeDef arg_def; NodeDefBuilder builder( - absl::StrCat(src_node->name(), "_", src_slot, "_arg"), kArgOp); + absl::StrCat(src_node->name(), "_", src_slot, "_arg"), kArgOp, + NodeDebugInfo(src_node->def())); DataType dtype = edge->dst()->input_type(edge->dst_input()); builder.Attr("T", dtype); builder.Attr("index", arg_index); @@ -814,7 +815,8 @@ Status Encapsulator::Subgraph::RecordResult( if (inserted) { NodeDef ret_def; NodeDefBuilder builder( - absl::StrCat(src_node->name(), "_", src_slot, "_retval"), kRetValOp); + absl::StrCat(src_node->name(), "_", src_slot, "_retval"), kRetValOp, + NodeDebugInfo(src_node->def())); DataType dtype = src_node->output_type(src_slot); builder.Attr("T", dtype); builder.Attr("index", ret_index); @@ -974,6 +976,7 @@ Status Encapsulator::Subgraph::AddHostComputes( } NodeDef host_compute_def; + // TODO(shikharagarwal): What source node should we use for errors? NodeDefBuilder builder(absl::StrCat("outside_compilation_", oc_subgraph_name, "_host_compute"), kHostComputeOp); @@ -1005,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. @@ -1028,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); } } } @@ -1040,6 +1046,7 @@ Status Encapsulator::Subgraph::MakeSequencingNode(const string& subgraph_name, Graph* graph_out) { if (sequencer_ == nullptr) { NodeDef seq_def; + // TODO(shikharagarwal): What source node should we use for errors? NodeDefBuilder builder(absl::StrCat(subgraph_name, "_sequencer"), "NoOp"); builder.Attr(kXlaHostTransferSequencerAttr, subgraph_name); builder.Device(device_); @@ -1055,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); } } @@ -1122,8 +1130,11 @@ Status Encapsulator::Subgraph::BuildFunctionDef( fdef); } - if (!reuse_existing_functions || library->Find(name) == nullptr) { + const FunctionDef* original_fdef = library->Find(name); + if (!reuse_existing_functions || original_fdef == nullptr) { TF_RETURN_IF_ERROR(library->AddFunctionDef(fdef)); + } else if (!FunctionDefsEqual(*original_fdef, fdef)) { + TF_RETURN_IF_ERROR(library->ReplaceFunction(name, fdef)); } return Status::OK(); } @@ -1211,7 +1222,8 @@ Status Encapsulator::Subgraph::AddHostComputeKeyPlaceholder( GraphDefBuilder::Options options(graph_out, /*status=*/nullptr); NodeDef key_def; NodeDefBuilder builder( - absl::StrCat(call_node_def_.name(), "_key_placeholder"), "Placeholder"); + absl::StrCat(call_node_def_.name(), "_key_placeholder"), "Placeholder", + NodeDebugInfo(call_node_def_)); builder.Attr("dtype", DT_STRING); builder.Attr("shape", shape_proto); builder.Attr("_host_compute_call_node", call_node_def_.name()); @@ -1245,6 +1257,7 @@ Status Encapsulator::Subgraph::AddRecvAtHostNode( } NodeDef recv_def; + // TODO(shikharagarwal): What source node should we use for errors? NodeDefBuilder builder(absl::StrCat("outside_compilation_", subgraph_name, "_", oc_subgraph_name, "_recv"), kRecvAtHostOp); @@ -1270,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(); } @@ -1300,6 +1314,7 @@ Status Encapsulator::Subgraph::AddSendFromHostNode( } NodeDef send_def; + // TODO(shikharagarwal): What source node should we use for errors? NodeDefBuilder builder(absl::StrCat("outside_compilation_", subgraph_name, "_", oc_subgraph_name, "_send"), kSendFromHostOp); @@ -1326,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(); } @@ -1436,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()); } @@ -1722,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(); @@ -1751,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(); @@ -1830,8 +1849,9 @@ Node* AddDummyShapedNode(const Node* src_node, int src_port, // Add any Enter nodes required to bring the constant to the correct control // flow frame. while (!control_flow_info[src_node->id()].frame_name.empty()) { + NodeDebugInfo debug_info(*src_node); NodeBuilder enter_builder(options.GetNameForOp("Enter"), "Enter", - options.op_registry()); + options.op_registry(), &debug_info); enter_builder.Attr("frame_name", control_flow_info[src_node->id()].frame_name); enter_builder.Attr("is_constant", true); @@ -2015,7 +2035,8 @@ Status Encapsulator::DoStaticShapeInferenceForOutsideCompilationSend( return errors::InvalidArgument( "Shape inference is not possible for outside_compilation " "SendFromHost node ", - send_node->name(), " because shape of node ", n->name(), + send_node->name(), " because shape of node ", + FormatNodeForError(*n), " will not be known at compilation time."); } } @@ -2044,8 +2065,7 @@ Status Encapsulator::DoStaticShapeInferenceForOutsideCompilationSend( return errors::Internal( "Internal assumption failed while rewriting an outside_compilation " "cluster that contains a while loop. Logic assumes back-edge is to " - "port 1 of a 2-input " - "Merge node."); + "port 1 of a 2-input Merge node."); } // Connect the existing edge to both inputs of the Merge node so that the // graph will be well-formed. @@ -2118,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(); } @@ -2162,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); @@ -2524,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 49958093b8dcf35e8adcdfd2f7dfce8558d5db6f..1f8ec09e19c01d0a8b2a3761135ed53dfb2ad3b0 100644 --- a/tensorflow/compiler/jit/encapsulate_subgraphs_pass_test.cc +++ b/tensorflow/compiler/jit/encapsulate_subgraphs_pass_test.cc @@ -16,18 +16,26 @@ limitations under the License. #include #include -#include "absl/strings/str_cat.h" #include "tensorflow/compiler/jit/encapsulate_subgraphs_pass.h" #include "absl/strings/match.h" +#include "absl/strings/str_cat.h" #include "tensorflow/cc/framework/ops.h" #include "tensorflow/cc/ops/standard_ops.h" +#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" #include "tensorflow/core/graph/graph_def_builder.h" +#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 { @@ -295,7 +303,7 @@ REGISTER_OP("XlaHostCompute") .Attr("Toutputs: list(type) >= 0") .Attr("ancestors: list(string) >= 0") .Attr("key: string") - .Attr("shape_inference_graph: string = ''") + .Attr("shape_inference_graph: func") .Attr("shapes: list(shape) >= 0") .SetShapeFn(::tensorflow::shape_inference::UnknownShape); @@ -406,8 +414,8 @@ Node* KeyPlaceholderShape(const GraphDefBuilder::Options& opts) { Node* KeyPlaceholder(const string& call_node, const GraphDefBuilder::Options& opts) { if (opts.HaveError()) return nullptr; - NodeBuilder node_builder(opts.GetNameForOp("Placeholder"), "Placeholder", - opts.op_registry()); + NodeBuilder node_builder(absl::StrCat(call_node, "_key_placeholder"), + "Placeholder", opts.op_registry()); TensorShapeProto shape; shape.add_dim()->set_size(2); return opts.WithAttr("shape", shape) @@ -494,7 +502,8 @@ Node* RetOp(int index, ops::NodeOut a, const GraphDefBuilder::Options& opts) { return opts.FinalizeBuilder(&node_builder); } -Status Encapsulate(GraphDef* graphdef, FunctionDefLibrary* library) { +Status Encapsulate(GraphDef* graphdef, FunctionDefLibrary* library, + const std::vector& encapsulated_functions) { Status s; // Convert the GraphDef to a Graph std::unique_ptr lib_def( @@ -505,11 +514,47 @@ Status Encapsulate(GraphDef* graphdef, FunctionDefLibrary* library) { s = ConvertGraphDefToGraph(options, *graphdef, graph.get()); if (!s.ok()) return s; + 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", *graph, - /*rewrite_subgraph_fn=*/{}, - /*reuse_existing_functions=*/false, - &graph_out, lib_def.get()); + s = EncapsulateSubgraphsInFunctions( + "_encapsulate", /*outside_compilation_attribute=*/"", *graph, + /*rewrite_subgraph_fn=*/{}, + /*reuse_existing_functions=*/false, &graph_out, lib_def.get()); + if (!s.ok()) return s; + + std::unordered_map clusters; + for (const auto& func : encapsulated_functions) { + Node* xla_computation_node; + for (Node* n : graph_out->nodes()) { + if (n->name() == func) { + xla_computation_node = n; + } + } + if (!xla_computation_node) { + return errors::Internal("Cannot find node ", func); + } + NameAttrList func_name_attrs; + func_name_attrs.set_name(func); + clusters.emplace(func, + XlaClusterInfo{func, func_name_attrs, xla_computation_node, + std::map{}}); + } + s = ExtractOutsideCompilation("_encapsulate", "_outside", clusters, + graph_out.get(), flr, lib_def.get()); if (!s.ok()) return s; GraphDef graphdef_out; @@ -517,9 +562,22 @@ Status Encapsulate(GraphDef* graphdef, FunctionDefLibrary* library) { graphdef->Swap(&graphdef_out); *library = lib_def->ToProto(); + // Remove "_xla_inferred_shapes" attr. They are added by + // `PerformStaticShapeInferenceBeforeEncapsulation`. + for (FunctionDef& fdef : *library->mutable_function()) { + for (NodeDef& node_def : *fdef.mutable_node_def()) { + node_def.mutable_attr()->erase("_xla_inferred_shapes"); + } + } + return s; } +Status Encapsulate(GraphDef* graphdef, FunctionDefLibrary* library) { + std::vector encapsulated_functions; + return Encapsulate(graphdef, library, encapsulated_functions); +} + // If there are no marked nodes, funcification should be a no-op. TEST(EncapsulateSubgraphsTest, NoFunctions) { GraphDefBuilder builder(GraphDefBuilder::kFailImmediately); @@ -703,7 +761,7 @@ TEST(EncapsulateSubgraphsTest, InputDeduplication) { FunctionLibraryDefinition library(OpRegistry::Global(), {}); std::unique_ptr graph; TF_ASSERT_OK(EncapsulateSubgraphsInFunctions( - "_cluster", "_outside", graph_before_encapsulation, + "_cluster", "", graph_before_encapsulation, /*rewrite_subgraph_fn=*/{}, /*reuse_existing_functions=*/false, &graph, &library)); @@ -755,7 +813,7 @@ TEST(EncapsulateSubgraphsWithGuaranteeConstOpTest, Simple) { FunctionLibraryDefinition library(OpRegistry::Global(), {}); int guaranteed_consts = 0; TF_ASSERT_OK(EncapsulateSubgraphsInFunctions( - "_encapsulate", "_outside", graph_before, + "_encapsulate", "", graph_before, /*rewrite_subgraph_fn=*/ [&guaranteed_consts](const std::vector& arg_source_tensors, std::unique_ptr* graph_ptr, @@ -800,7 +858,7 @@ TEST(EncapsulateSubgraphsWithGuaranteeConstOpTest, Add) { FunctionLibraryDefinition library(OpRegistry::Global(), {}); int guaranteed_consts = 0; TF_ASSERT_OK(EncapsulateSubgraphsInFunctions( - "_encapsulate", "_outside", graph_before, + "_encapsulate", "", graph_before, /*rewrite_subgraph_fn=*/ [&guaranteed_consts](const std::vector& arg_source_tensors, std::unique_ptr* graph_ptr, @@ -854,30 +912,34 @@ TEST(EncapsulateSubgraphsTest, OneFunctionOneOutside) { TF_EXPECT_OK(b1.ToGraphDef(&graphdef)); } - TF_EXPECT_OK(Encapsulate(&graphdef, &library)); + std::vector encapsulated_functions{"F1"}; + TF_EXPECT_OK(Encapsulate(&graphdef, &library, encapsulated_functions)); FunctionDefLibrary library_expected; GraphDef graphdef_expected; { GraphDefBuilder shape(GraphDefBuilder::kFailImmediately); - Node* key_constant = - KeyPlaceholderShape(shape.opts().WithName("KnownShape/_0")); - Node* recv = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", - {DT_FLOAT, DT_FLOAT}, shape.opts()); + Node* key_constant = KeyPlaceholder("F1", shape.opts()); + Node* recv = RecvAtHost( + ops::NodeOut(key_constant, 0), "F1", "O1", {DT_FLOAT, DT_FLOAT}, + shape.opts().WithAttr(kXlaHasHostTransferAttrName, true)); Node* e = Binary(ops::NodeOut(recv, 0), ops::NodeOut(recv, 1), shape.opts() .WithName("E") .WithAttr("_encapsulate", "F1") .WithAttr("_outside", "O1")); - SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, shape.opts()); + SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, + shape.opts().WithAttr(kXlaHasHostTransferAttrName, true)); TF_EXPECT_OK( AddGraphDefToFunctionLibrary(shape, "F1_O1", &library_expected)); } + NameAttrList shape_inference_graph; + shape_inference_graph.set_name("_outside_compilation_shape_inference_F1_O1"); *library_expected.add_function() = test::function::XTimesTwo(); *library_expected.add_function() = FunctionDefHelper::Create( - "F1", {"a_0_arg:float", "b_0_arg:float"}, {"f_0_retval:float"}, {}, + "F1", {"a_0_arg:float", "b_0_arg:float"}, {"f_0_retval_retval:float"}, {}, { {{"C"}, "UnaryTest", {"a_0_arg"}}, {{"c"}, "BinaryTest", {"b_0_arg", "C:o:0"}, {}, {"C"}}, @@ -893,13 +955,14 @@ TEST(EncapsulateSubgraphsTest, OneFunctionOneOutside) { {"Toutputs", absl::Span({DT_FLOAT})}, {"ancestors", absl::Span({})}, {"key", "host_compute_channel_F1_O1"}, - {"shape_inference_graph", - "_outside_compilation_shape_inference_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", "F:o:0"}}); + {{"f_0_retval_retval", "F:o:0"}}); { std::unique_ptr lib_def( @@ -910,16 +973,18 @@ TEST(EncapsulateSubgraphsTest, OneFunctionOneOutside) { Node* key_constant = KeyPlaceholder("F1", b2.opts().WithName("F1_key_placeholder")); - Node* recv = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", - {DT_FLOAT, DT_FLOAT}, b2.opts()); + Node* recv = RecvAtHost( + ops::NodeOut(key_constant, 0), "F1", "O1", {DT_FLOAT, DT_FLOAT}, + b2.opts().WithAttr(kXlaHasHostTransferAttrName, true)); Node* e = Binary(ops::NodeOut(recv, 0), ops::NodeOut(recv, 1), b2.opts() .WithName("E") - .WithControlInputs({recv, b}) + .WithControlInputs({recv}) .WithAttr("_encapsulate", "F1") .WithAttr("_outside", "O1")); Node* send = SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, - b2.opts().WithControlInput(e)); + b2.opts().WithControlInput(e).WithAttr( + kXlaHasHostTransferAttrName, true)); Node* s = Sequencer( b2.opts().WithName("F1_sequencer").WithControlInputs({recv, send}), @@ -928,9 +993,9 @@ TEST(EncapsulateSubgraphsTest, OneFunctionOneOutside) { NodeBuilder node_builder("F1", "F1", lib_def.get()); node_builder.Input(a).Input(b); Node* call = - b2.opts().WithControlInputs({s}).FinalizeBuilder(&node_builder); + b2.opts().WithControlInputs({s, b}).FinalizeBuilder(&node_builder); - Binary(a, call, b2.opts().WithName("G").WithControlInputs({e})); + Binary(a, call, b2.opts().WithName("G").WithControlInputs({call})); TF_EXPECT_OK(b2.ToGraphDef(&graphdef_expected)); } @@ -975,58 +1040,71 @@ TEST(EncapsulateSubgraphsTest, OneFunctionTwoOutside) { TF_EXPECT_OK(b1.ToGraphDef(&graphdef)); } - TF_EXPECT_OK(Encapsulate(&graphdef, &library)); + std::vector encapsulated_functions{"F1"}; + TF_EXPECT_OK(Encapsulate(&graphdef, &library, encapsulated_functions)); FunctionDefLibrary library_expected; GraphDef graphdef_expected; { GraphDefBuilder shape1(GraphDefBuilder::kFailImmediately); - Node* key_constant = - KeyPlaceholderShape(shape1.opts().WithName("KnownShape/_0")); - Node* recv = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", - {DT_FLOAT, DT_FLOAT}, shape1.opts()); + Node* key_constant = KeyPlaceholder("F1", shape1.opts()); + Node* recv = RecvAtHost( + ops::NodeOut(key_constant, 0), "F1", "O1", {DT_FLOAT, DT_FLOAT}, + shape1.opts().WithAttr(kXlaHasHostTransferAttrName, true)); Node* e = Binary(ops::NodeOut(recv, 0), ops::NodeOut(recv, 1), shape1.opts() .WithName("E") .WithAttr("_encapsulate", "F1") .WithAttr("_outside", "O1")); - SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, shape1.opts()); + SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, + shape1.opts().WithAttr(kXlaHasHostTransferAttrName, true)); TF_EXPECT_OK( AddGraphDefToFunctionLibrary(shape1, "F1_O1", &library_expected)); } { GraphDefBuilder shape2(GraphDefBuilder::kFailImmediately); - Node* key_constant = - KeyPlaceholderShape(shape2.opts().WithName("KnownShape/_0")); - Node* recv1 = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", - {DT_FLOAT, DT_FLOAT}, shape2.opts()); + Node* key_constant = KeyPlaceholder("F1", shape2.opts()); + Node* recv1 = RecvAtHost( + ops::NodeOut(key_constant, 0), "F1", "O1", {DT_FLOAT, DT_FLOAT}, + shape2.opts().WithAttr(kXlaHasHostTransferAttrName, true)); Node* e = Binary(ops::NodeOut(recv1, 0), ops::NodeOut(recv1, 1), shape2.opts() .WithName("E") .WithAttr("_encapsulate", "F1") .WithAttr("_outside", "O1")); - Node* recv2 = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O2", - {DT_FLOAT, DT_FLOAT}, shape2.opts()); + Node* recv2 = RecvAtHost( + ops::NodeOut(key_constant, 0), "F1", "O2", {DT_FLOAT, DT_FLOAT}, + shape2.opts().WithAttr(kXlaHasHostTransferAttrName, true)); + Node* g = Binary(e, ops::NodeOut(recv2, 0), + shape2.opts() + .WithName("G") + .WithAttr("_encapsulate", "F1") + .WithAttr("_outside", "O2")); Node* h = Binary(ops::NodeOut(recv2, 1), e, shape2.opts() .WithName("H") .WithAttr("_encapsulate", "F1") .WithAttr("_outside", "O2")); - SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O2", {h}, shape2.opts()); + SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O2", {g, h}, + shape2.opts().WithAttr(kXlaHasHostTransferAttrName, true)); TF_EXPECT_OK( AddGraphDefToFunctionLibrary(shape2, "F1_O2", &library_expected)); } + NameAttrList shape_inference_graph1, shape_inference_graph2; + shape_inference_graph1.set_name("_outside_compilation_shape_inference_F1_O1"); + shape_inference_graph2.set_name("_outside_compilation_shape_inference_F1_O2"); *library_expected.add_function() = FunctionDefHelper::Create( - "F1", {"a_0_arg:float", "b_0_arg:float"}, {"i_0_retval:float"}, {}, + "F1", {"a_0_arg:float", "b_0_arg:float"}, + {"g_0_retval_retval:float", "i_0_retval_retval:float"}, {}, { {{"C"}, "UnaryTest", {"a_0_arg"}}, {{"D"}, "BinaryTest", {"b_0_arg", "C:o:0"}, {}}, {{"I"}, "UnaryTest", - {"outside_compilation_O2_host_compute:outputs:0"}}, + {"outside_compilation_O2_host_compute:outputs:1"}}, {{"F"}, "BinaryTest", {"C:o:0", "outside_compilation_O1_host_compute:outputs:0"}, @@ -1036,15 +1114,15 @@ TEST(EncapsulateSubgraphsTest, OneFunctionTwoOutside) { "XlaHostCompute", {"F:o:0", "D:o:0"}, {{"Tinputs", absl::Span({DT_FLOAT, DT_FLOAT})}, - {"Toutputs", absl::Span({DT_FLOAT})}, - {"ancestors", - absl::Span({"outside_compilation_O1_host_compute"})}, + {"Toutputs", absl::Span({DT_FLOAT, DT_FLOAT})}, + {"ancestors", absl::Span({})}, {"key", "host_compute_channel_F1_O2"}, - {"shape_inference_graph", - "_outside_compilation_shape_inference_F1_O2"}, + {"shape_inference_graph", shape_inference_graph2}, {"shapes", absl::Span({})}, - {"_outside_compilation_subgraph", "O2"}}, - {"F", "outside_compilation_O1_host_compute"}}, + {"_outside_compilation_subgraph", "O2"}, + {"_xla_token_input_nodes", + absl::Span({"_xla_token_arg_node"})}}, + {"F"}}, {{"outside_compilation_O1_host_compute"}, "XlaHostCompute", {"C:o:0", "D:o:0"}, @@ -1052,13 +1130,15 @@ TEST(EncapsulateSubgraphsTest, OneFunctionTwoOutside) { {"Toutputs", absl::Span({DT_FLOAT})}, {"ancestors", absl::Span({})}, {"key", "host_compute_channel_F1_O1"}, - {"shape_inference_graph", - "_outside_compilation_shape_inference_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"}}, }, - {{"i_0_retval", "I:o:0"}}); + {{"g_0_retval_retval", "outside_compilation_O2_host_compute:outputs:0"}, + {"i_0_retval_retval", "I:o:0"}}); { std::unique_ptr lib_def( @@ -1069,19 +1149,22 @@ TEST(EncapsulateSubgraphsTest, OneFunctionTwoOutside) { Node* key_constant = KeyPlaceholder("F1", b2.opts().WithName("F1_key_placeholder")); - Node* recv1 = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", - {DT_FLOAT, DT_FLOAT}, b2.opts()); + Node* recv1 = RecvAtHost( + ops::NodeOut(key_constant, 0), "F1", "O1", {DT_FLOAT, DT_FLOAT}, + b2.opts().WithAttr(kXlaHasHostTransferAttrName, true)); Node* e = Binary(ops::NodeOut(recv1, 0), ops::NodeOut(recv1, 1), b2.opts() .WithName("E") - .WithControlInputs({recv1, b}) + .WithControlInputs({recv1}) .WithAttr("_encapsulate", "F1") .WithAttr("_outside", "O1")); Node* send1 = SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, - b2.opts().WithControlInput(e)); + b2.opts().WithControlInput(e).WithAttr( + kXlaHasHostTransferAttrName, true)); - Node* recv2 = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O2", - {DT_FLOAT, DT_FLOAT}, b2.opts()); + Node* recv2 = RecvAtHost( + ops::NodeOut(key_constant, 0), "F1", "O2", {DT_FLOAT, DT_FLOAT}, + b2.opts().WithAttr(kXlaHasHostTransferAttrName, true)); Node* g = Binary(e, ops::NodeOut(recv2, 0), b2.opts() .WithName("G") @@ -1094,7 +1177,8 @@ TEST(EncapsulateSubgraphsTest, OneFunctionTwoOutside) { .WithAttr("_encapsulate", "F1") .WithAttr("_outside", "O2")); Node* send2 = - SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O2", {h}, b2.opts()); + SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O2", {g, h}, + b2.opts().WithAttr(kXlaHasHostTransferAttrName, true)); Node* s = Sequencer(b2.opts() .WithName("F1_sequencer") @@ -1103,12 +1187,13 @@ TEST(EncapsulateSubgraphsTest, OneFunctionTwoOutside) { NodeBuilder node_builder("F1", "F1", lib_def.get()); node_builder.Input(a).Input(b); - Node* call = b2.opts().WithControlInput(s).FinalizeBuilder(&node_builder); + Node* call = + b2.opts().WithControlInputs({s, b}).FinalizeBuilder(&node_builder); - Binary(g, call, b2.opts().WithName("J")); + Binary(ops::NodeOut(call, 0), ops::NodeOut(call, 1), + b2.opts().WithName("J")); TF_EXPECT_OK(b2.ToGraphDef(&graphdef_expected)); } - TF_EXPECT_GRAPH_EQ(graphdef_expected, graphdef); TF_EXPECT_FUNCTIONDEFLIBRARY_EQ(library_expected, library); } @@ -1149,33 +1234,20 @@ TEST(EncapsulateSubgraphsTest, TwoFunctionsTwoOutside) { TF_EXPECT_OK(b1.ToGraphDef(&graphdef)); } - TF_EXPECT_OK(Encapsulate(&graphdef, &library)); + std::vector encapsulated_functions{"F1", "F2"}; + TF_EXPECT_OK(Encapsulate(&graphdef, &library, encapsulated_functions)); FunctionDefLibrary library_expected; GraphDef graphdef_expected; - { - GraphDefBuilder shape(GraphDefBuilder::kFailImmediately); - Node* key_constant = - KeyPlaceholderShape(shape.opts().WithName("KnownShape/_0")); - Node* recv = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", - {DT_FLOAT, DT_FLOAT}, shape.opts()); - Node* e = Binary(ops::NodeOut(recv, 0), ops::NodeOut(recv, 1), - shape.opts() - .WithName("E") - .WithAttr("_encapsulate", "F1") - .WithAttr("_outside", "O1")); - SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, shape.opts()); - TF_EXPECT_OK( - AddGraphDefToFunctionLibrary(shape, "F1_O1", &library_expected)); - } - TensorShapeProto shape_proto_expected; shape_proto_expected.add_dim()->set_size(2); *library_expected.add_function() = FunctionDefHelper::Create( "F1", {"a_0_arg:float", "b_0_arg:float"}, - {"f_0_retval:float", "d_0_retval:float"}, {}, + {"e_0_retval_retval:float", "f_0_retval_retval:float", + "d_0_retval_retval:float"}, + {}, { {{"C"}, "UnaryTest", {"a_0_arg"}}, {{"D"}, "BinaryTest", {"b_0_arg", "C:o:0"}}, @@ -1191,17 +1263,21 @@ TEST(EncapsulateSubgraphsTest, TwoFunctionsTwoOutside) { {"Toutputs", absl::Span({DT_FLOAT})}, {"ancestors", absl::Span({})}, {"key", "host_compute_channel_F1_O1"}, - {"shape_inference_graph", - "_outside_compilation_shape_inference_F1_O1"}, - {"shapes", absl::Span({})}, - {"_outside_compilation_subgraph", "O1"}}, + {"shape_inference_graph", NameAttrList()}, + {"shapes", + absl::Span({shape_proto_expected})}, + {"_outside_compilation_subgraph", "O1"}, + {"_xla_token_input_nodes", + absl::Span({"_xla_token_arg_node"})}}, {"D"}}, }, - {{"d_0_retval", "D:o:0"}, {"f_0_retval", "F:o:0"}}); + {{"e_0_retval_retval", "outside_compilation_O1_host_compute:outputs:0"}, + {"d_0_retval_retval", "D:o:0"}, + {"f_0_retval_retval", "F:o:0"}}); *library_expected.add_function() = FunctionDefHelper::Create( - "F2", {"e_0_arg:float", "f_0_arg:float"}, - {"g_0_retval:float", "i_0_retval:float"}, {}, + "F2", {"e_0_arg:float", "f_0_arg:float", "d_0_arg:float"}, + {"g_0_retval_retval:float", "i_0_retval_retval:float"}, {}, { {{"G"}, "BinaryTest", {"e_0_arg", "f_0_arg"}}, {{"I"}, @@ -1209,17 +1285,19 @@ TEST(EncapsulateSubgraphsTest, TwoFunctionsTwoOutside) { {"f_0_arg", "outside_compilation_O1_host_compute:outputs:0"}}, {{"outside_compilation_O1_host_compute"}, "XlaHostCompute", - {"G:o:0"}, - {{"Tinputs", absl::Span({DT_FLOAT})}, + {"d_0_arg", "G:o:0"}, + {{"Tinputs", absl::Span({DT_FLOAT, DT_FLOAT})}, {"Toutputs", absl::Span({DT_FLOAT})}, {"ancestors", absl::Span({})}, {"key", "host_compute_channel_F2_O1"}, - {"shape_inference_graph", ""}, + {"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", "G:o:0"}, {"i_0_retval", "I:o:0"}}); + {{"g_0_retval_retval", "G:o:0"}, {"i_0_retval_retval", "I:o:0"}}); { std::unique_ptr lib_def( @@ -1230,16 +1308,18 @@ TEST(EncapsulateSubgraphsTest, TwoFunctionsTwoOutside) { Node* key_constant1 = KeyPlaceholder("F1", b2.opts().WithName("F1_key_placeholder")); - Node* recv1 = RecvAtHost(ops::NodeOut(key_constant1, 0), "F1", "O1", - {DT_FLOAT, DT_FLOAT}, b2.opts()); + Node* recv1 = RecvAtHost( + ops::NodeOut(key_constant1, 0), "F1", "O1", {DT_FLOAT, DT_FLOAT}, + b2.opts().WithAttr(kXlaHasHostTransferAttrName, true)); Node* e = Binary(ops::NodeOut(recv1, 0), ops::NodeOut(recv1, 1), b2.opts() .WithName("E") - .WithControlInputs({recv1, b}) + .WithControlInputs({recv1}) .WithAttr("_encapsulate", "F1") .WithAttr("_outside", "O1")); Node* send1 = SendFromHost(ops::NodeOut(key_constant1, 0), "F1", "O1", {e}, - b2.opts().WithControlInput(e)); + b2.opts().WithControlInput(e).WithAttr( + kXlaHasHostTransferAttrName, true)); Node* s1 = Sequencer( b2.opts().WithName("F1_sequencer").WithControlInputs({recv1, send1}), "F1"); @@ -1247,27 +1327,31 @@ TEST(EncapsulateSubgraphsTest, TwoFunctionsTwoOutside) { NodeBuilder node_builder1("F1", "F1", lib_def.get()); node_builder1.Input(a).Input(b); Node* call1 = - b2.opts().WithControlInput(s1).FinalizeBuilder(&node_builder1); + b2.opts().WithControlInputs({s1, b}).FinalizeBuilder(&node_builder1); Node* key_constant2 = KeyPlaceholder("F2", b2.opts().WithName("F2_key_placeholder")); - Node* recv2 = RecvAtHost(ops::NodeOut(key_constant2, 0), "F2", "O1", - {DT_FLOAT}, b2.opts()); - Node* h = Binary(ops::NodeOut(call1, 1), recv2, + Node* recv2 = RecvAtHost( + ops::NodeOut(key_constant2, 0), "F2", "O1", {DT_FLOAT, DT_FLOAT}, + b2.opts().WithAttr(kXlaHasHostTransferAttrName, true)); + Node* h = Binary(recv2, ops::NodeOut(recv2, 1), b2.opts() .WithName("H") .WithAttr("_encapsulate", "F2") .WithAttr("_outside", "O1")); - Node* send2 = SendFromHost(ops::NodeOut(key_constant2, 0), "F2", "O1", {h}, - b2.opts()); + Node* send2 = + SendFromHost(ops::NodeOut(key_constant2, 0), "F2", "O1", {h}, + b2.opts().WithAttr(kXlaHasHostTransferAttrName, true)); Node* s2 = Sequencer( b2.opts().WithName("F2_sequencer").WithControlInputs({recv2, send2}), "F2"); NodeBuilder node_builder2("F2", "F2", lib_def.get()); - node_builder2.Input(e).Input(call1); + node_builder2.Input(call1) + .Input(ops::NodeOut(call1, 1)) + .Input(ops::NodeOut(call1, 2)); Node* call2 = b2.opts() - .WithControlInputs({s2, e, call1}) + .WithControlInputs({s2, call1}) .FinalizeBuilder(&node_builder2); Binary(call2, ops::NodeOut(call2, 1), b2.opts().WithName("J")); TF_EXPECT_OK(b2.ToGraphDef(&graphdef_expected)); @@ -1305,51 +1389,22 @@ TEST(EncapsulateSubgraphsTest, TwoFunctionsTwoOutsideDependencyFromOutside) { Node* h = Unary(g, b1.opts() .WithName("H") .WithAttr("_encapsulate", "F2") - .WithAttr("_outside", "O1") - .WithControlInput(e)); + .WithAttr("_outside", "O1")); Node* i = Unary(h, b1.opts().WithName("I").WithAttr("_encapsulate", "F2")); Binary(f, i, b1.opts().WithName("J")); TF_EXPECT_OK(b1.ToGraphDef(&graphdef)); } - TF_EXPECT_OK(Encapsulate(&graphdef, &library)); + std::vector encapsulated_functions{"F1", "F2"}; + TF_EXPECT_OK(Encapsulate(&graphdef, &library, encapsulated_functions)); FunctionDefLibrary library_expected; GraphDef graphdef_expected; - - { - GraphDefBuilder shape(GraphDefBuilder::kFailImmediately); - Node* key_constant = - KeyPlaceholderShape(shape.opts().WithName("KnownShape/_0")); - Node* recv = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", - {DT_FLOAT, DT_FLOAT}, shape.opts()); - Node* e = Binary(ops::NodeOut(recv, 0), ops::NodeOut(recv, 1), - shape.opts() - .WithName("E") - .WithAttr("_encapsulate", "F1") - .WithAttr("_outside", "O1")); - SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, shape.opts()); - TF_EXPECT_OK( - AddGraphDefToFunctionLibrary(shape, "F1_O1", &library_expected)); - } - - { - GraphDefBuilder shape(GraphDefBuilder::kFailImmediately); - Node* key_constant = - KeyPlaceholderShape(shape.opts().WithName("KnownShape/_0")); - Node* recv = RecvAtHost(ops::NodeOut(key_constant, 0), "F2", "O1", - {DT_FLOAT}, shape.opts()); - Node* h = Unary(recv, shape.opts() - .WithName("H") - .WithAttr("_encapsulate", "F2") - .WithAttr("_outside", "O1")); - SendFromHost(ops::NodeOut(key_constant, 0), "F2", "O1", {h}, shape.opts()); - TF_EXPECT_OK( - AddGraphDefToFunctionLibrary(shape, "F2_O1", &library_expected)); - } + TensorShapeProto shape_proto_expected; + shape_proto_expected.add_dim()->set_size(2); *library_expected.add_function() = FunctionDefHelper::Create( - "F1", {"a_0_arg:float", "b_0_arg:float"}, {"f_0_retval:float"}, {}, + "F1", {"a_0_arg:float", "b_0_arg:float"}, {"f_0_retval_retval:float"}, {}, { {{"C"}, "UnaryTest", {"a_0_arg"}}, {{"D"}, "BinaryTest", {"b_0_arg", "C:o:0"}}, @@ -1365,16 +1420,18 @@ TEST(EncapsulateSubgraphsTest, TwoFunctionsTwoOutsideDependencyFromOutside) { {"Toutputs", absl::Span({DT_FLOAT})}, {"ancestors", absl::Span({})}, {"key", "host_compute_channel_F1_O1"}, - {"shape_inference_graph", - "_outside_compilation_shape_inference_F1_O1"}, - {"shapes", absl::Span({})}, - {"_outside_compilation_subgraph", "O1"}}, + {"shape_inference_graph", NameAttrList()}, + {"shapes", + absl::Span({shape_proto_expected})}, + {"_outside_compilation_subgraph", "O1"}, + {"_xla_token_input_nodes", + absl::Span({"_xla_token_arg_node"})}}, {"D"}}, }, - {{"f_0_retval", "F:o:0"}}); + {{"f_0_retval_retval", "F:o:0"}}); *library_expected.add_function() = FunctionDefHelper::Create( - "F2", {"a_0_arg:float", "b_0_arg:float"}, {"i_0_retval:float"}, {}, + "F2", {"a_0_arg:float", "b_0_arg:float"}, {"i_0_retval_retval:float"}, {}, { {{"G"}, "BinaryTest", {"a_0_arg", "b_0_arg"}}, {{"I"}, @@ -1387,12 +1444,14 @@ TEST(EncapsulateSubgraphsTest, TwoFunctionsTwoOutsideDependencyFromOutside) { {"Toutputs", absl::Span({DT_FLOAT})}, {"ancestors", absl::Span({})}, {"key", "host_compute_channel_F2_O1"}, - {"shape_inference_graph", - "_outside_compilation_shape_inference_F2_O1"}, - {"shapes", absl::Span({})}, - {"_outside_compilation_subgraph", "O1"}}}, + {"shape_inference_graph", NameAttrList()}, + {"shapes", + absl::Span({shape_proto_expected})}, + {"_outside_compilation_subgraph", "O1"}, + {"_xla_token_input_nodes", + absl::Span({"_xla_token_arg_node"})}}}, }, - {{"i_0_retval", "I:o:0"}}); + {{"i_0_retval_retval", "I:o:0"}}); { std::unique_ptr lib_def( @@ -1408,7 +1467,7 @@ TEST(EncapsulateSubgraphsTest, TwoFunctionsTwoOutsideDependencyFromOutside) { Node* e = Binary(ops::NodeOut(recv1, 0), ops::NodeOut(recv1, 1), b2.opts() .WithName("E") - .WithControlInputs({recv1, b}) + .WithControlInputs({recv1}) .WithAttr("_encapsulate", "F1") .WithAttr("_outside", "O1")); Node* send1 = SendFromHost(ops::NodeOut(key_constant1, 0), "F1", "O1", {e}, @@ -1420,7 +1479,7 @@ TEST(EncapsulateSubgraphsTest, TwoFunctionsTwoOutsideDependencyFromOutside) { NodeBuilder node_builder1("F1", "F1", lib_def.get()); node_builder1.Input(a).Input(b); Node* call1 = - b2.opts().WithControlInput(s1).FinalizeBuilder(&node_builder1); + b2.opts().WithControlInputs({s1, b}).FinalizeBuilder(&node_builder1); Node* key_constant2 = KeyPlaceholder("F2", b2.opts().WithName("F2_key_placeholder")); @@ -1429,8 +1488,7 @@ TEST(EncapsulateSubgraphsTest, TwoFunctionsTwoOutsideDependencyFromOutside) { Node* h = Unary(recv2, b2.opts() .WithName("H") .WithAttr("_encapsulate", "F2") - .WithAttr("_outside", "O1") - .WithControlInput(e)); + .WithAttr("_outside", "O1")); Node* send2 = SendFromHost(ops::NodeOut(key_constant2, 0), "F2", "O1", {h}, b2.opts()); @@ -1439,9 +1497,8 @@ TEST(EncapsulateSubgraphsTest, TwoFunctionsTwoOutsideDependencyFromOutside) { "F2"); NodeBuilder node_builder2("F2", "F2", lib_def.get()); node_builder2.Input(a).Input(b); - Node* call2 = b2.opts() - .WithControlInputs({s2, call1}) - .FinalizeBuilder(&node_builder2); + Node* call2 = + b2.opts().WithControlInputs({s2}).FinalizeBuilder(&node_builder2); Binary(call1, call2, b2.opts().WithName("J")); TF_EXPECT_OK(b2.ToGraphDef(&graphdef_expected)); } @@ -1473,7 +1530,8 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationNoInputs) { TF_EXPECT_OK(b1.ToGraphDef(&graphdef)); } - TF_EXPECT_OK(Encapsulate(&graphdef, &library)); + std::vector encapsulated_functions{"F1"}; + TF_EXPECT_OK(Encapsulate(&graphdef, &library, encapsulated_functions)); FunctionDefLibrary library_expected; GraphDef graphdef_expected; @@ -1482,7 +1540,7 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationNoInputs) { shape_proto_expected.add_dim()->set_size(2); *library_expected.add_function() = FunctionDefHelper::Create( - "F1", {"a_0_arg:float", "b_0_arg:float"}, {"f_0_retval:float"}, {}, + "F1", {"a_0_arg:float", "b_0_arg:float"}, {"f_0_retval_retval:float"}, {}, { {{"C"}, "UnaryTest", {"a_0_arg"}}, {{"D"}, "BinaryTest", {"b_0_arg", "C:o:0"}}, @@ -1491,17 +1549,19 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationNoInputs) { {"D:o:0", "outside_compilation_O1_host_compute:outputs:0"}}, {{"outside_compilation_O1_host_compute"}, "XlaHostCompute", - {}, - {{"Tinputs", absl::Span({})}, + {"a_0_arg"}, + {{"Tinputs", absl::Span({DT_FLOAT})}, {"Toutputs", absl::Span({DT_FLOAT})}, {"ancestors", absl::Span({})}, {"key", "host_compute_channel_F1_O1"}, - {"shape_inference_graph", ""}, + {"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", "F:o:0"}}); + {{"f_0_retval_retval", "F:o:0"}}); { std::unique_ptr lib_def( @@ -1510,16 +1570,19 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationNoInputs) { Node* a = InputShaped(b2.opts().WithName("A")); Node* b = Input(b2.opts().WithName("B")); - Node* e = Unary(a, b2.opts() - .WithName("E") - .WithAttr("_encapsulate", "F1") - .WithAttr("_outside", "O1")); Node* key_constant = KeyPlaceholder("F1", b2.opts().WithName("F1_key_placeholder")); + Node* recv1 = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", + {DT_FLOAT}, b2.opts()); + Node* e = Unary(recv1, b2.opts() + .WithName("E") + .WithAttr("_encapsulate", "F1") + .WithAttr("_outside", "O1")); Node* send1 = SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, b2.opts()); Node* s1 = Sequencer( - b2.opts().WithName("F1_sequencer").WithControlInput(send1), "F1"); + b2.opts().WithName("F1_sequencer").WithControlInputs({send1, recv1}), + "F1"); NodeBuilder node_builder1("F1", "F1", lib_def.get()); node_builder1.Input(a).Input(b); Node* call1 = @@ -1557,7 +1620,8 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationControlInput) { TF_EXPECT_OK(b1.ToGraphDef(&graphdef)); } - TF_EXPECT_OK(Encapsulate(&graphdef, &library)); + std::vector encapsulated_functions{"F1"}; + TF_EXPECT_OK(Encapsulate(&graphdef, &library, encapsulated_functions)); FunctionDefLibrary library_expected; GraphDef graphdef_expected; @@ -1566,7 +1630,7 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationControlInput) { shape_proto_expected.add_dim()->set_size(2); *library_expected.add_function() = FunctionDefHelper::Create( - "F1", {"a_0_arg:float", "b_0_arg:float"}, {"f_0_retval:float"}, {}, + "F1", {"a_0_arg:float", "b_0_arg:float"}, {"f_0_retval_retval:float"}, {}, { {{"C"}, "UnaryTest", {"a_0_arg"}}, {{"D"}, "BinaryTest", {"b_0_arg", "C:o:0"}}, @@ -1575,18 +1639,20 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationControlInput) { {"D:o:0", "outside_compilation_O1_host_compute:outputs:0"}}, {{"outside_compilation_O1_host_compute"}, "XlaHostCompute", - {}, - {{"Tinputs", absl::Span({})}, + {"a_0_arg"}, + {{"Tinputs", absl::Span({DT_FLOAT})}, {"Toutputs", absl::Span({DT_FLOAT})}, {"ancestors", absl::Span({})}, {"key", "host_compute_channel_F1_O1"}, - {"shape_inference_graph", ""}, + {"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", "F:o:0"}}); + {{"f_0_retval_retval", "F:o:0"}}); { std::unique_ptr lib_def( @@ -1597,13 +1663,13 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationControlInput) { Node* key_constant = KeyPlaceholder("F1", b2.opts().WithName("F1_key_placeholder")); - Node* recv1 = - RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", {}, b2.opts()); - Node* e = Unary(a, b2.opts() - .WithName("E") - .WithControlInput(recv1) - .WithAttr("_encapsulate", "F1") - .WithAttr("_outside", "O1")); + Node* recv1 = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", + {DT_FLOAT}, b2.opts()); + Node* e = Unary(recv1, b2.opts() + .WithName("E") + .WithControlInput(recv1) + .WithAttr("_encapsulate", "F1") + .WithAttr("_outside", "O1")); Node* send1 = SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, b2.opts()); Node* s1 = Sequencer( @@ -1644,13 +1710,33 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationNoOutputs) { TF_EXPECT_OK(b1.ToGraphDef(&graphdef)); } - TF_EXPECT_OK(Encapsulate(&graphdef, &library)); + std::vector encapsulated_functions{"F1"}; + TF_EXPECT_OK(Encapsulate(&graphdef, &library, encapsulated_functions)); FunctionDefLibrary library_expected; GraphDef graphdef_expected; + { + GraphDefBuilder shape1(GraphDefBuilder::kFailImmediately); + Node* key_constant = KeyPlaceholder("F1", shape1.opts()); + Node* recv1 = + RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", {DT_FLOAT}, + shape1.opts().WithAttr(kXlaHasHostTransferAttrName, true)); + Node* e = Unary(ops::NodeOut(recv1, 0), shape1.opts() + .WithName("E") + .WithAttr("_encapsulate", "F1") + .WithAttr("_outside", "O1")); + SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, + shape1.opts().WithAttr(kXlaHasHostTransferAttrName, true)); + TF_EXPECT_OK( + AddGraphDefToFunctionLibrary(shape1, "F1_O1", &library_expected)); + } + + NameAttrList shape_inference_graph; + shape_inference_graph.set_name("_outside_compilation_shape_inference_F1_O1"); *library_expected.add_function() = FunctionDefHelper::Create( - "F1", {"a_0_arg:float", "b_0_arg:float"}, {"f_0_retval:float"}, {}, + "F1", {"a_0_arg:float", "b_0_arg:float"}, + {"e_0_retval_retval:float", "f_0_retval_retval:float"}, {}, { {{"C"}, "UnaryTest", {"a_0_arg"}}, {{"D"}, "BinaryTest", {"b_0_arg", "C:o:0"}}, @@ -1659,14 +1745,17 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationNoOutputs) { "XlaHostCompute", {"D:o:0"}, {{"Tinputs", absl::Span({DT_FLOAT})}, - {"Toutputs", absl::Span({})}, + {"Toutputs", absl::Span({DT_FLOAT})}, {"ancestors", absl::Span({})}, {"key", "host_compute_channel_F1_O1"}, - {"shape_inference_graph", ""}, + {"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"})}}}, }, - {{"f_0_retval", "F:o:0"}}); + {{"e_0_retval_retval", "outside_compilation_O1_host_compute:outputs:0"}, + {"f_0_retval_retval", "F:o:0"}}); { std::unique_ptr lib_def( @@ -1683,14 +1772,17 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationNoOutputs) { .WithName("E") .WithAttr("_encapsulate", "F1") .WithAttr("_outside", "O1")); + Node* send1 = + SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, b2.opts()); Node* s1 = Sequencer( - b2.opts().WithName("F1_sequencer").WithControlInput(recv1), "F1"); + b2.opts().WithName("F1_sequencer").WithControlInputs({recv1, send1}), + "F1"); NodeBuilder node_builder1("F1", "F1", lib_def.get()); node_builder1.Input(a).Input(b); Node* call1 = b2.opts().WithControlInput(s1).FinalizeBuilder(&node_builder1); - Binary(e, call1, b2.opts().WithName("G")); + Binary(call1, ops::NodeOut(call1, 1), b2.opts().WithName("G")); TF_EXPECT_OK(b2.ToGraphDef(&graphdef_expected)); } @@ -1721,13 +1813,33 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationControlOutput) { TF_EXPECT_OK(b1.ToGraphDef(&graphdef)); } - TF_EXPECT_OK(Encapsulate(&graphdef, &library)); + std::vector encapsulated_functions{"F1"}; + TF_EXPECT_OK(Encapsulate(&graphdef, &library, encapsulated_functions)); FunctionDefLibrary library_expected; GraphDef graphdef_expected; + { + GraphDefBuilder shape1(GraphDefBuilder::kFailImmediately); + Node* key_constant = KeyPlaceholder("F1", shape1.opts()); + Node* recv1 = + RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", {DT_FLOAT}, + shape1.opts().WithAttr(kXlaHasHostTransferAttrName, true)); + Node* e = Unary(ops::NodeOut(recv1, 0), shape1.opts() + .WithName("E") + .WithAttr("_encapsulate", "F1") + .WithAttr("_outside", "O1")); + SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, + shape1.opts().WithAttr(kXlaHasHostTransferAttrName, true)); + TF_EXPECT_OK( + AddGraphDefToFunctionLibrary(shape1, "F1_O1", &library_expected)); + } + + NameAttrList shape_inference_graph; + shape_inference_graph.set_name("_outside_compilation_shape_inference_F1_O1"); *library_expected.add_function() = FunctionDefHelper::Create( - "F1", {"a_0_arg:float", "b_0_arg:float"}, {"f_0_retval:float"}, {}, + "F1", {"a_0_arg:float", "b_0_arg:float"}, + {"e_0_retval_retval:float", "f_0_retval_retval:float"}, {}, { {{"C"}, "UnaryTest", {"a_0_arg"}}, {{"D"}, "BinaryTest", {"b_0_arg", "C:o:0"}}, @@ -1740,14 +1852,17 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationControlOutput) { "XlaHostCompute", {"D:o:0"}, {{"Tinputs", absl::Span({DT_FLOAT})}, - {"Toutputs", absl::Span({})}, + {"Toutputs", absl::Span({DT_FLOAT})}, {"ancestors", absl::Span({})}, {"key", "host_compute_channel_F1_O1"}, - {"shape_inference_graph", ""}, + {"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"})}}}, }, - {{"f_0_retval", "F:o:0"}}); + {{"e_0_retval_retval", "outside_compilation_O1_host_compute:outputs:0"}, + {"f_0_retval_retval", "F:o:0"}}); { std::unique_ptr lib_def( @@ -1764,7 +1879,7 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationControlOutput) { .WithName("E") .WithAttr("_encapsulate", "F1") .WithAttr("_outside", "O1")); - Node* send1 = SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {}, + Node* send1 = SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, b2.opts().WithControlInput(e)); Node* s1 = Sequencer( b2.opts().WithName("F1_sequencer").WithControlInputs({recv1, send1}), @@ -1774,7 +1889,7 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationControlOutput) { Node* call1 = b2.opts().WithControlInput(s1).FinalizeBuilder(&node_builder1); - Binary(e, call1, b2.opts().WithName("G")); + Binary(call1, ops::NodeOut(call1, 1), b2.opts().WithName("G")); TF_EXPECT_OK(b2.ToGraphDef(&graphdef_expected)); } @@ -1811,28 +1926,51 @@ TEST(EncapsulateSubgraphsTest, TF_EXPECT_OK(b1.ToGraphDef(&graphdef)); } - TF_EXPECT_OK(Encapsulate(&graphdef, &library)); + std::vector encapsulated_functions{"F1"}; + TF_EXPECT_OK(Encapsulate(&graphdef, &library, encapsulated_functions)); FunctionDefLibrary library_expected; GraphDef graphdef_expected; + { + GraphDefBuilder shape1(GraphDefBuilder::kFailImmediately); + Node* key_constant = KeyPlaceholder("F1", shape1.opts()); + Node* recv1 = + RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", {DT_FLOAT}, + shape1.opts().WithAttr(kXlaHasHostTransferAttrName, true)); + Node* e = Unary(ops::NodeOut(recv1, 0), shape1.opts() + .WithName("E") + .WithAttr("_encapsulate", "F1") + .WithAttr("_outside", "O1")); + SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, + shape1.opts().WithAttr(kXlaHasHostTransferAttrName, true)); + TF_EXPECT_OK( + AddGraphDefToFunctionLibrary(shape1, "F1_O1", &library_expected)); + } + { GraphDefBuilder shape2(GraphDefBuilder::kFailImmediately); - Node* key_constant = - KeyPlaceholderShape(shape2.opts().WithName("KnownShape/_0")); - Node* recv2 = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O2", - {DT_FLOAT}, shape2.opts()); + Node* key_constant = KeyPlaceholder("F1", shape2.opts()); + Node* recv2 = + RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O2", {DT_FLOAT}, + shape2.opts().WithAttr(kXlaHasHostTransferAttrName, true)); Node* g = Unary(ops::NodeOut(recv2, 0), shape2.opts() .WithName("G") .WithAttr("_encapsulate", "F1") .WithAttr("_outside", "O2")); - SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O2", {g}, shape2.opts()); + SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O2", {g}, + shape2.opts().WithAttr(kXlaHasHostTransferAttrName, true)); TF_EXPECT_OK( AddGraphDefToFunctionLibrary(shape2, "F1_O2", &library_expected)); } + NameAttrList shape_inference_graph1; + shape_inference_graph1.set_name("_outside_compilation_shape_inference_F1_O1"); + NameAttrList shape_inference_graph2; + shape_inference_graph2.set_name("_outside_compilation_shape_inference_F1_O2"); *library_expected.add_function() = FunctionDefHelper::Create( - "F1", {"a_0_arg:float", "b_0_arg:float"}, {"h_0_retval:float"}, {}, + "F1", {"a_0_arg:float", "b_0_arg:float"}, + {"e_0_retval_retval:float", "h_0_retval_retval:float"}, {}, { {{"C"}, "UnaryTest", {"a_0_arg"}}, {{"D"}, "BinaryTest", {"b_0_arg", "C:o:0"}}, @@ -1840,6 +1978,18 @@ TEST(EncapsulateSubgraphsTest, {{"H"}, "UnaryTest", {"outside_compilation_O2_host_compute:outputs:0"}}, + {{"outside_compilation_O1_host_compute"}, + "XlaHostCompute", + {"a_0_arg"}, + {{"Tinputs", absl::Span({DT_FLOAT})}, + {"Toutputs", absl::Span({DT_FLOAT})}, + {"ancestors", absl::Span({})}, + {"key", "host_compute_channel_F1_O1"}, + {"shape_inference_graph", shape_inference_graph1}, + {"shapes", absl::Span({})}, + {"_outside_compilation_subgraph", "O1"}, + {"_xla_token_input_nodes", + absl::Span({"_xla_token_arg_node"})}}}, {{"outside_compilation_O2_host_compute"}, "XlaHostCompute", {"F:o:0"}, @@ -1847,12 +1997,14 @@ TEST(EncapsulateSubgraphsTest, {"Toutputs", absl::Span({DT_FLOAT})}, {"ancestors", absl::Span({})}, {"key", "host_compute_channel_F1_O2"}, - {"shape_inference_graph", - "_outside_compilation_shape_inference_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"})}}}, }, - {{"h_0_retval", "H:o:0"}}); + {{"e_0_retval_retval", "outside_compilation_O1_host_compute:outputs:0"}, + {"h_0_retval_retval", "H:o:0"}}); { std::unique_ptr lib_def( @@ -1860,30 +2012,39 @@ TEST(EncapsulateSubgraphsTest, GraphDefBuilder b2(GraphDefBuilder::kFailImmediately, lib_def.get()); Node* a = Input(b2.opts().WithName("A")); Node* b = Input(b2.opts().WithName("B")); - - Node* e = Unary(a, b2.opts() - .WithName("E") - .WithAttr("_encapsulate", "F1") - .WithAttr("_outside", "O1")); Node* key_constant = KeyPlaceholder("F1", b2.opts().WithName("F1_key_placeholder")); - Node* recv = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O2", - {DT_FLOAT}, b2.opts()); - Node* g = Unary(recv, b2.opts() - .WithName("G") - .WithAttr("_encapsulate", "F1") - .WithAttr("_outside", "O2") - .WithControlInput(e)); - Node* send = - SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O2", {g}, b2.opts()); - Node* s1 = Sequencer( - b2.opts().WithName("F1_sequencer").WithControlInputs({recv, send}), - "F1"); + Node* recv1 = + RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", {DT_FLOAT}, + b2.opts().WithAttr(kXlaHasHostTransferAttrName, true)); + + Node* e = Unary(recv1, b2.opts() + .WithName("E") + .WithAttr("_encapsulate", "F1") + .WithAttr("_outside", "O1")); + Node* send1 = + SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, + b2.opts().WithAttr(kXlaHasHostTransferAttrName, true)); + Node* recv2 = + RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O2", {DT_FLOAT}, + b2.opts().WithAttr(kXlaHasHostTransferAttrName, true)); + Node* g = Unary(recv2, b2.opts() + .WithName("G") + .WithAttr("_encapsulate", "F1") + .WithAttr("_outside", "O2") + .WithControlInput(e)); + Node* send2 = + SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O2", {g}, + b2.opts().WithAttr(kXlaHasHostTransferAttrName, true)); + Node* s1 = Sequencer(b2.opts() + .WithName("F1_sequencer") + .WithControlInputs({recv1, send1, recv2, send2}), + "F1"); NodeBuilder node_builder1("F1", "F1", lib_def.get()); node_builder1.Input(a).Input(b).ControlInput(s1); Node* call1 = b2.opts().FinalizeBuilder(&node_builder1); - Binary(e, call1, b2.opts().WithName("I")); + Binary(call1, ops::NodeOut(call1, 1), b2.opts().WithName("I")); TF_EXPECT_OK(b2.ToGraphDef(&graphdef_expected)); } @@ -1920,28 +2081,33 @@ TEST(EncapsulateSubgraphsTest, TF_EXPECT_OK(b1.ToGraphDef(&graphdef)); } - TF_EXPECT_OK(Encapsulate(&graphdef, &library)); + std::vector encapsulated_functions{"F1"}; + TF_EXPECT_OK(Encapsulate(&graphdef, &library, encapsulated_functions)); FunctionDefLibrary library_expected; GraphDef graphdef_expected; { GraphDefBuilder shape1(GraphDefBuilder::kFailImmediately); - Node* key_constant = - KeyPlaceholderShape(shape1.opts().WithName("KnownShape/_0")); - Node* recv2 = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", - {DT_FLOAT}, shape1.opts()); + Node* key_constant = KeyPlaceholder("F1", shape1.opts()); + Node* recv2 = + RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", {DT_FLOAT}, + shape1.opts().WithAttr(kXlaHasHostTransferAttrName, true)); Node* e = Unary(ops::NodeOut(recv2, 0), shape1.opts() .WithName("E") .WithAttr("_encapsulate", "F1") .WithAttr("_outside", "O1")); - SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, shape1.opts()); + SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, + shape1.opts().WithAttr(kXlaHasHostTransferAttrName, true)); TF_EXPECT_OK( AddGraphDefToFunctionLibrary(shape1, "F1_O1", &library_expected)); } + NameAttrList shape_inference_graph; + shape_inference_graph.set_name("_outside_compilation_shape_inference_F1_O1"); *library_expected.add_function() = FunctionDefHelper::Create( - "F1", {"a_0_arg:float", "b_0_arg:float"}, {"h_0_retval:float"}, {}, + "F1", {"a_0_arg:float", "b_0_arg:float"}, + {"e_0_retval_retval:float", "h_0_retval_retval:float"}, {}, { {{"C"}, "UnaryTest", {"a_0_arg"}}, {{"D"}, "BinaryTest", {"b_0_arg", "C:o:0"}}, @@ -1949,6 +2115,18 @@ TEST(EncapsulateSubgraphsTest, "UnaryTest", {"outside_compilation_O1_host_compute:outputs:0"}}, {{"H"}, "UnaryTest", {"F:o:0"}}, + {{"outside_compilation_O2_host_compute"}, + "XlaHostCompute", + {"a_0_arg"}, + {{"Tinputs", absl::Span({DT_FLOAT})}, + {"Toutputs", absl::Span({})}, + {"ancestors", absl::Span({})}, + {"key", "host_compute_channel_F1_O2"}, + {"shape_inference_graph", NameAttrList()}, + {"shapes", absl::Span({})}, + {"_outside_compilation_subgraph", "O2"}, + {"_xla_token_input_nodes", + absl::Span({"_xla_token_arg_node"})}}}, {{"outside_compilation_O1_host_compute"}, "XlaHostCompute", {"D:o:0"}, @@ -1956,12 +2134,14 @@ TEST(EncapsulateSubgraphsTest, {"Toutputs", absl::Span({DT_FLOAT})}, {"ancestors", absl::Span({})}, {"key", "host_compute_channel_F1_O1"}, - {"shape_inference_graph", - "_outside_compilation_shape_inference_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"})}}}, }, - {{"h_0_retval", "H:o:0"}}); + {{"e_0_retval_retval", "outside_compilation_O1_host_compute:outputs:0"}, + {"h_0_retval_retval", "H:o:0"}}); { std::unique_ptr lib_def( @@ -1972,27 +2152,33 @@ TEST(EncapsulateSubgraphsTest, Node* key_constant = KeyPlaceholder("F1", b2.opts().WithName("F1_key_placeholder")); - Node* recv = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", - {DT_FLOAT}, b2.opts()); - Node* e = Unary(recv, b2.opts() - .WithName("E") - .WithAttr("_encapsulate", "F1") - .WithAttr("_outside", "O1")); + Node* recv1 = + RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", {DT_FLOAT}, + b2.opts().WithAttr(kXlaHasHostTransferAttrName, true)); + Node* e = Unary(recv1, b2.opts() + .WithName("E") + .WithAttr("_encapsulate", "F1") + .WithAttr("_outside", "O1")); Node* send = - SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, b2.opts()); - /*Node* g =*/Unary(a, b2.opts() - .WithName("G") - .WithAttr("_encapsulate", "F1") - .WithAttr("_outside", "O2") - .WithControlInput(e)); - Node* s1 = Sequencer( - b2.opts().WithName("F1_sequencer").WithControlInputs({recv, send}), - "F1"); + SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, + b2.opts().WithAttr(kXlaHasHostTransferAttrName, true)); + Node* recv2 = + RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O2", {DT_FLOAT}, + b2.opts().WithAttr(kXlaHasHostTransferAttrName, true)); + /*Node* g =*/Unary(recv2, b2.opts() + .WithName("G") + .WithAttr("_encapsulate", "F1") + .WithAttr("_outside", "O2") + .WithControlInput(e)); + Node* s1 = Sequencer(b2.opts() + .WithName("F1_sequencer") + .WithControlInputs({recv1, recv2, send}), + "F1"); NodeBuilder node_builder1("F1", "F1", lib_def.get()); node_builder1.Input(a).Input(b).ControlInput(s1); Node* call1 = b2.opts().FinalizeBuilder(&node_builder1); - Binary(e, call1, b2.opts().WithName("I")); + Binary(call1, ops::NodeOut(call1, 1), b2.opts().WithName("I")); TF_EXPECT_OK(b2.ToGraphDef(&graphdef_expected)); } @@ -2034,28 +2220,33 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationClusterDependency) { TF_EXPECT_OK(b1.ToGraphDef(&graphdef)); } - TF_EXPECT_OK(Encapsulate(&graphdef, &library)); + std::vector encapsulated_functions{"F1"}; + TF_EXPECT_OK(Encapsulate(&graphdef, &library, encapsulated_functions)); FunctionDefLibrary library_expected; GraphDef graphdef_expected; { GraphDefBuilder shape1(GraphDefBuilder::kFailImmediately); - Node* key_constant = - KeyPlaceholderShape(shape1.opts().WithName("KnownShape/_0")); - Node* recv2 = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", - {DT_FLOAT}, shape1.opts()); + Node* key_constant = KeyPlaceholder("F1", shape1.opts()); + Node* recv2 = + RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", {DT_FLOAT}, + shape1.opts().WithAttr(kXlaHasHostTransferAttrName, true)); Node* e = Unary(ops::NodeOut(recv2, 0), shape1.opts() .WithName("E") .WithAttr("_encapsulate", "F1") .WithAttr("_outside", "O1")); - SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, shape1.opts()); + SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, + shape1.opts().WithAttr(kXlaHasHostTransferAttrName, true)); TF_EXPECT_OK( AddGraphDefToFunctionLibrary(shape1, "F1_O1", &library_expected)); } + NameAttrList shape_inference_graph; + shape_inference_graph.set_name("_outside_compilation_shape_inference_F1_O1"); *library_expected.add_function() = FunctionDefHelper::Create( - "F1", {"a_0_arg:float", "b_0_arg:float"}, {"h_0_retval:float"}, {}, + "F1", {"a_0_arg:float", "b_0_arg:float"}, + {"e_0_retval_retval:float", "h_0_retval_retval:float"}, {}, {{{"C"}, "UnaryTest", {"a_0_arg"}}, {{"D"}, "BinaryTest", {"b_0_arg", "C:o:0"}}, {{"F"}, "UnaryTest", {"outside_compilation_O1_host_compute:outputs:0"}}, @@ -2067,37 +2258,39 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationClusterDependency) { {"Toutputs", absl::Span({DT_FLOAT})}, {"ancestors", absl::Span({})}, {"key", "host_compute_channel_F1_O1"}, - {"shape_inference_graph", - "_outside_compilation_shape_inference_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"}, {{"Tinputs", absl::Span({DT_FLOAT})}, {"Toutputs", absl::Span({})}, - {"ancestors", - absl::Span({"outside_compilation_O1_host_compute"})}, + {"ancestors", absl::Span({})}, {"key", "host_compute_channel_F1_O2"}, - {"shape_inference_graph", ""}, + {"shape_inference_graph", NameAttrList()}, {"shapes", absl::Span({})}, - {"_outside_compilation_subgraph", "O2"}}, - {"outside_compilation_O1_host_compute"}}, + {"_outside_compilation_subgraph", "O2"}, + {"_xla_token_input_nodes", + absl::Span({"_xla_token_arg_node"})}}, + {}}, {{"outside_compilation_O3_host_compute"}, "XlaHostCompute", {"D:o:0"}, {{"Tinputs", absl::Span({DT_FLOAT})}, {"Toutputs", absl::Span({})}, - {"ancestors", - absl::Span({"outside_compilation_O1_host_compute", - "outside_compilation_O2_host_compute"})}, + {"ancestors", absl::Span({})}, {"key", "host_compute_channel_F1_O3"}, - {"shape_inference_graph", ""}, + {"shape_inference_graph", NameAttrList()}, {"shapes", absl::Span({})}, - {"_outside_compilation_subgraph", "O3"}}, - {"outside_compilation_O1_host_compute", - "outside_compilation_O2_host_compute"}}}, - {{"h_0_retval", "H:o:0"}}); + {"_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"}}); { std::unique_ptr lib_def( @@ -2108,23 +2301,27 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationClusterDependency) { Node* key_constant = KeyPlaceholder("F1", b2.opts().WithName("F1_key_placeholder")); - Node* recv1 = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", - {DT_FLOAT}, b2.opts()); + Node* recv1 = + RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", {DT_FLOAT}, + b2.opts().WithAttr(kXlaHasHostTransferAttrName, true)); Node* e = Unary(recv1, b2.opts() .WithName("E") .WithAttr("_encapsulate", "F1") .WithAttr("_outside", "O1")); Node* send = - SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, b2.opts()); - Node* recv2 = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O2", - {DT_FLOAT}, b2.opts()); + SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, + b2.opts().WithAttr(kXlaHasHostTransferAttrName, true)); + Node* recv2 = + RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O2", {DT_FLOAT}, + b2.opts().WithAttr(kXlaHasHostTransferAttrName, true)); Node* g = Unary(recv2, b2.opts() .WithName("G") .WithAttr("_encapsulate", "F1") .WithAttr("_outside", "O2") .WithControlInput(e)); - Node* recv3 = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O3", - {DT_FLOAT}, b2.opts()); + Node* recv3 = + RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O3", {DT_FLOAT}, + b2.opts().WithAttr(kXlaHasHostTransferAttrName, true)); /*Node* i =*/Binary(recv3, e, b2.opts() .WithName("I") @@ -2139,7 +2336,7 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationClusterDependency) { node_builder1.Input(a).Input(b).ControlInput(s1); Node* call1 = b2.opts().FinalizeBuilder(&node_builder1); - Binary(e, call1, b2.opts().WithName("J")); + Binary(call1, ops::NodeOut(call1, 1), b2.opts().WithName("J")); TF_EXPECT_OK(b2.ToGraphDef(&graphdef_expected)); } @@ -2169,19 +2366,52 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationNoInputsOrOutputs) { TF_EXPECT_OK(b1.ToGraphDef(&graphdef)); } - TF_EXPECT_OK(Encapsulate(&graphdef, &library)); + std::vector encapsulated_functions{"F1"}; + TF_EXPECT_OK(Encapsulate(&graphdef, &library, encapsulated_functions)); FunctionDefLibrary library_expected; GraphDef graphdef_expected; + { + GraphDefBuilder shape1(GraphDefBuilder::kFailImmediately); + Node* key_constant = KeyPlaceholder("F1", shape1.opts()); + Node* recv2 = + RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", {DT_FLOAT}, + shape1.opts().WithAttr(kXlaHasHostTransferAttrName, true)); + Node* e = Unary(ops::NodeOut(recv2, 0), shape1.opts() + .WithName("E") + .WithAttr("_encapsulate", "F1") + .WithAttr("_outside", "O1")); + SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, + shape1.opts().WithAttr(kXlaHasHostTransferAttrName, true)); + TF_EXPECT_OK( + AddGraphDefToFunctionLibrary(shape1, "F1_O1", &library_expected)); + } + + NameAttrList shape_inference_graph; + shape_inference_graph.set_name("_outside_compilation_shape_inference_F1_O1"); *library_expected.add_function() = FunctionDefHelper::Create( - "F1", {"a_0_arg:float", "b_0_arg:float"}, {"f_0_retval:float"}, {}, + "F1", {"a_0_arg:float", "b_0_arg:float"}, + {"e_0_retval_retval:float", "f_0_retval_retval:float"}, {}, { {{"C"}, "UnaryTest", {"a_0_arg"}}, {{"D"}, "BinaryTest", {"b_0_arg", "C:o:0"}}, {{"F"}, "UnaryTest", {"D:o:0"}}, + {{"outside_compilation_O1_host_compute"}, + "XlaHostCompute", + {"a_0_arg"}, + {{"Tinputs", absl::Span({DT_FLOAT})}, + {"Toutputs", absl::Span({DT_FLOAT})}, + {"ancestors", absl::Span({})}, + {"key", "host_compute_channel_F1_O1"}, + {"shape_inference_graph", shape_inference_graph}, + {"shapes", absl::Span({})}, + {"_outside_compilation_subgraph", "O1"}, + {"_xla_token_input_nodes", + absl::Span({"_xla_token_arg_node"})}}}, }, - {{"f_0_retval", "F:o:0"}}); + {{"e_0_retval_retval", "outside_compilation_O1_host_compute:outputs:0"}, + {"f_0_retval_retval", "F:o:0"}}); { std::unique_ptr lib_def( @@ -2190,15 +2420,26 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationNoInputsOrOutputs) { Node* a = Input(b2.opts().WithName("A")); Node* b = Input(b2.opts().WithName("B")); - Node* e = Unary(a, b2.opts() - .WithName("E") - .WithAttr("_encapsulate", "F1") - .WithAttr("_outside", "O1")); + Node* key_constant = + KeyPlaceholder("F1", b2.opts().WithName("F1_key_placeholder")); + Node* recv = + RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", {DT_FLOAT}, + b2.opts().WithAttr(kXlaHasHostTransferAttrName, true)); + Node* e = Unary(recv, b2.opts() + .WithName("E") + .WithAttr("_encapsulate", "F1") + .WithAttr("_outside", "O1")); + Node* send = + SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, + b2.opts().WithAttr(kXlaHasHostTransferAttrName, true)); + Node* s = Sequencer( + b2.opts().WithName("F1_sequencer").WithControlInputs({recv, send}), + "F1"); NodeBuilder node_builder1("F1", "F1", lib_def.get()); - node_builder1.Input(a).Input(b); + node_builder1.Input(a).Input(b).ControlInput(s); Node* call1 = b2.opts().FinalizeBuilder(&node_builder1); - Binary(e, call1, b2.opts().WithName("G")); + Binary(call1, ops::NodeOut(call1, 1), b2.opts().WithName("G")); TF_EXPECT_OK(b2.ToGraphDef(&graphdef_expected)); } @@ -2234,31 +2475,34 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationShapeInference) { TF_EXPECT_OK(b1.ToGraphDef(&graphdef)); } - TF_EXPECT_OK(Encapsulate(&graphdef, &library)); + std::vector encapsulated_functions{"F1"}; + TF_EXPECT_OK(Encapsulate(&graphdef, &library, encapsulated_functions)); FunctionDefLibrary library_expected; GraphDef graphdef_expected; { GraphDefBuilder shape(GraphDefBuilder::kFailImmediately); - Node* key_constant = - KeyPlaceholderShape(shape.opts().WithName("KnownShape/_0")); - Node* known = KnownShape({2}, shape.opts().WithName("KnownShape/_1")); - Node* recv = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", - {DT_FLOAT}, shape.opts()); - Node* e = BinaryUnknownShape(known, recv, + Node* key_constant = KeyPlaceholder("F1", shape.opts()); + Node* recv = RecvAtHost( + ops::NodeOut(key_constant, 0), "F1", "O1", {DT_FLOAT, DT_FLOAT}, + shape.opts().WithAttr(kXlaHasHostTransferAttrName, true)); + Node* e = BinaryUnknownShape(recv, ops::NodeOut(recv, 1), shape.opts() .WithName("E") .WithAttr("_encapsulate", "F1") .WithAttr("_outside", "O1")); - SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, shape.opts()); + SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, + shape.opts().WithAttr(kXlaHasHostTransferAttrName, true)); TF_EXPECT_OK( AddGraphDefToFunctionLibrary(shape, "F1_O1", &library_expected)); } + NameAttrList shape_inference_graph; + shape_inference_graph.set_name("_outside_compilation_shape_inference_F1_O1"); *library_expected.add_function() = test::function::XTimesTwo(); *library_expected.add_function() = FunctionDefHelper::Create( - "F1", {"b_0_arg:float", "c_0_arg:float"}, {"f_0_retval:float"}, {}, + "F1", {"b_0_arg:float", "c_0_arg:float"}, {"f_0_retval_retval:float"}, {}, { {{"c"}, "UnaryTest", {"b_0_arg"}, {}, {}}, {{"F"}, @@ -2268,18 +2512,19 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationShapeInference) { {"outside_compilation_O1_host_compute"}}, {{"outside_compilation_O1_host_compute"}, "XlaHostCompute", - {"c:o:0"}, - {{"Tinputs", absl::Span({DT_FLOAT})}, + {"c_0_arg", "c:o:0"}, + {{"Tinputs", absl::Span({DT_FLOAT, DT_FLOAT})}, {"Toutputs", absl::Span({DT_FLOAT})}, {"ancestors", absl::Span({})}, {"key", "host_compute_channel_F1_O1"}, - {"shape_inference_graph", - "_outside_compilation_shape_inference_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", "F:o:0"}}); + {{"f_0_retval_retval", "F:o:0"}}); { std::unique_ptr lib_def( @@ -2291,16 +2536,18 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationShapeInference) { Node* key_constant = KeyPlaceholder("F1", b2.opts().WithName("F1_key_placeholder")); - Node* recv = RecvAtHost(ops::NodeOut(key_constant, 0), "F1", "O1", - {DT_FLOAT}, b2.opts()); - Node* e = BinaryUnknownShape(c, ops::NodeOut(recv, 0), + Node* recv = RecvAtHost( + ops::NodeOut(key_constant, 0), "F1", "O1", {DT_FLOAT, DT_FLOAT}, + b2.opts().WithAttr(kXlaHasHostTransferAttrName, true)); + Node* e = BinaryUnknownShape(recv, ops::NodeOut(recv, 1), b2.opts() .WithName("E") - .WithControlInputs({recv, b}) + .WithControlInputs({recv}) .WithAttr("_encapsulate", "F1") .WithAttr("_outside", "O1")); Node* send = SendFromHost(ops::NodeOut(key_constant, 0), "F1", "O1", {e}, - b2.opts().WithControlInput(e)); + b2.opts().WithControlInput(e).WithAttr( + kXlaHasHostTransferAttrName, true)); Node* s = Sequencer( b2.opts().WithName("F1_sequencer").WithControlInputs({recv, send}), @@ -2309,9 +2556,9 @@ TEST(EncapsulateSubgraphsTest, OutsideCompilationShapeInference) { NodeBuilder node_builder("F1", "F1", lib_def.get()); node_builder.Input(b).Input(c); Node* call = - b2.opts().WithControlInputs({s, c}).FinalizeBuilder(&node_builder); + b2.opts().WithControlInputs({s, b, c}).FinalizeBuilder(&node_builder); - Binary(a, call, b2.opts().WithName("G").WithControlInputs({e})); + Binary(a, call, b2.opts().WithName("G").WithControlInputs({call})); TF_EXPECT_OK(b2.ToGraphDef(&graphdef_expected)); } diff --git a/tensorflow/compiler/jit/encapsulate_util.cc b/tensorflow/compiler/jit/encapsulate_util.cc index a3581b4fa00494ab3e98d19f97e1fcff6534b89a..2264806d6bdabd9f26d9f83b681524399f996317 100644 --- a/tensorflow/compiler/jit/encapsulate_util.cc +++ b/tensorflow/compiler/jit/encapsulate_util.cc @@ -62,10 +62,10 @@ void ReplaceAttr(Node* n, const string& attr_name, const T& value) { n->AddAttr(attr_name, value); } -// Step 1a ~ 1d for PreprocessForEncapsulation(). See comments of -// PreprocessForEncapsulation() for details. -Status ProcessControlEdges(Graph* g, const string& xla_computation_attr_name, - const string& outside_compilation_attr_name) { +// Step 1 for `PreprocessEdgesBetweenOutsideCompilations`. See comments of +// `PreprocessEdgesBetweenOutsideCompilations` for details. +Status PreprocessControlEdgesBetweenOutsideCompilations( + Graph* g, const string& outside_compilation_attr_name) { // Gather edges to remove. We should not remove the edge while iterating. std::vector edges_to_remove; for (const Edge* e : g->edges()) { @@ -73,67 +73,26 @@ Status ProcessControlEdges(Graph* g, const string& xla_computation_attr_name, continue; } - auto src_xla_computation = - GetStringAttr(*e->src(), xla_computation_attr_name); - auto dst_xla_computation = - GetStringAttr(*e->dst(), xla_computation_attr_name); auto src_outside_compilation = GetStringAttr(*e->src(), outside_compilation_attr_name); auto dst_outside_compilation = GetStringAttr(*e->dst(), outside_compilation_attr_name); - if (!src_xla_computation && !dst_xla_computation) { - continue; - } else if (src_xla_computation && !dst_xla_computation) { - if (src_outside_compilation) { - // Case 1d: outside compilation to host computation control edge. - TF_RETURN_IF_ERROR(AppendToListAttr( - e->dst(), kXlaControlDependenciesAttrName, e->src()->name())); - } - } else if (!src_xla_computation && dst_xla_computation) { - if (dst_outside_compilation) { - // Case 1d: host computation control to outside compilation edge. + if (src_outside_compilation && dst_outside_compilation) { + if (*src_outside_compilation != *dst_outside_compilation) { + // Case 1a: outside compilation to outside compilation control edge. + edges_to_remove.push_back(e); + TF_RETURN_IF_ERROR(AppendToListAttr( - e->dst(), kXlaControlDependenciesAttrName, e->src()->name())); - } - } else { // src_xla_computation && dst_xla_computation - if (*src_xla_computation != *dst_xla_computation) { - if (src_outside_compilation && dst_outside_compilation) { - // Case 1c: outside compilation to outside compilation control edge. - edges_to_remove.push_back(e); - - TF_RETURN_IF_ERROR(AppendToListAttr( - e->dst(), kXlaControlDependenciesAttrName, e->src()->name())); - } else if (src_outside_compilation && !dst_outside_compilation) { - // Case 1b: outside compilation to another XLA computaition control - // edge. - TF_RETURN_IF_ERROR(AppendToListAttr( - e->src(), kXlaConnectedToOtherXlaComputationAttrName, - *dst_xla_computation)); - } else if (!src_outside_compilation && dst_outside_compilation) { - // Case 1b: another XLA computaition to outside compilation control - // edge. - TF_RETURN_IF_ERROR(AppendToListAttr( - e->dst(), kXlaConnectedFromOtherXlaComputationAttrName, - *src_xla_computation)); - } - } else { // *src_xla_computation == *dst_xla_computation - if (src_outside_compilation && dst_outside_compilation) { - if (*src_outside_compilation != *dst_outside_compilation) { - // Case 1c: outside compilation to outside compilation control edge. - edges_to_remove.push_back(e); - - TF_RETURN_IF_ERROR(AppendToListAttr( - e->dst(), kXlaControlDependenciesAttrName, e->src()->name())); - } - } else if (src_outside_compilation && !dst_outside_compilation) { - // Case 1a: outside compilation to its XLA computation control edge. - ReplaceAttr(e->src(), kXlaConnectedToXlaComputationAttrName, true); - } else if (!src_outside_compilation && dst_outside_compilation) { - // Case 1a: XLA computation to outside compilation in it control edge. - ReplaceAttr(e->dst(), kXlaConnectedFromXlaComputationAttrName, true); - } + e->dst(), kXlaControlDependenciesWithinXlaClusterAttrName, + e->src()->name())); } + } else if (src_outside_compilation && !dst_outside_compilation) { + // Case 1b: outside compilation to its XLA computation control edge. + ReplaceAttr(e->src(), kXlaConnectedToXlaComputationAttrName, true); + } else if (!src_outside_compilation && dst_outside_compilation) { + // Case 1b: XLA computation to outside compilation in it control edge. + ReplaceAttr(e->dst(), kXlaConnectedFromXlaComputationAttrName, true); } } @@ -143,14 +102,13 @@ Status ProcessControlEdges(Graph* g, const string& xla_computation_attr_name, return Status::OK(); } -// Step 2 for PreprocessForEncapsulation(). See comments of -// PreprocessForEncapsulation() for details. -Status ProcessXlaToXlaDataEdges(Graph* g, - const string& xla_computation_attr_name, - const string& outside_compilation_attr_name) { - // Gather edges between XLA computations. Notice that we do not store `Edge*` - // directly because we remove some nodes while adding Identity nodes, and - // those Edge pointers might be invalidated. +// Step 2 for `PreprocessEdgesBetweenOutsideCompilations`. See comments of +// `PreprocessEdgesBetweenOutsideCompilations` for details. +Status PreprocessDataEdgesBetweenOutsideCompilations( + Graph* g, const string& outside_compilation_attr_name) { + // Gather edges between outside compilation and host computation. Notice that + // we do not store `Edge*` directly because we remove some nodes while adding + // Identity nodes, and those Edge pointers might be invalidated. struct EdgeInfo { int dst_input, dst_node_id; }; @@ -160,106 +118,21 @@ Status ProcessXlaToXlaDataEdges(Graph* g, continue; } - auto src_xla_computation = - GetStringAttr(*e->src(), xla_computation_attr_name); - auto dst_xla_computation = - GetStringAttr(*e->dst(), xla_computation_attr_name); auto src_outside_compilation = GetStringAttr(*e->src(), outside_compilation_attr_name); auto dst_outside_compilation = GetStringAttr(*e->dst(), outside_compilation_attr_name); - if (!src_xla_computation || !dst_xla_computation) { - continue; - } - if (*src_xla_computation != *dst_xla_computation) { - if (src_outside_compilation || dst_outside_compilation) { - edges.push_back(EdgeInfo{e->dst_input(), e->dst()->id()}); - VLOG(4) << "XLA -> XLA edge: " << e->DebugString(); - } - } else { // *src_xla_computation == *dst_xla_computation - if (src_outside_compilation && dst_outside_compilation && - *src_outside_compilation != *dst_outside_compilation) { - edges.push_back(EdgeInfo{e->dst_input(), e->dst()->id()}); - VLOG(4) << "XLA -> XLA edge: " << e->DebugString(); - } - } - } - - // For each XLA -> XLA edge, add an Identity node between src and dst. - for (int i = 0; i < edges.size(); i++) { - Node* dst = g->FindNodeId(edges[i].dst_node_id); - const Edge* e; - TF_RETURN_IF_ERROR(dst->input_edge(edges[i].dst_input, &e)); - Node* src = e->src(); - int src_output = e->src_output(), dst_input = e->dst_input(); - g->RemoveEdge(e); - - // Create Identity node, and connect it between `src` and `dst`. - string identity_node_name = - absl::StrCat("bridge_", src->name(), "_", dst->name()); - DataType dtype = src->output_type(src_output); - TF_ASSIGN_OR_RETURN(Node * identity_node, - BuildIdentityNode(g, identity_node_name, dtype, src, - /*requested_device=*/absl::nullopt)); - identity_node->AddAttr(kBridgeSourceNodeAttrName, src->name()); - g->AddEdge(src, src_output, identity_node, 0); - g->AddEdge(identity_node, 0, dst, dst_input); - - // Replace `e->dst()` because its input node changed. - NodeDef new_def = dst->def(); - *new_def.mutable_input(dst_input) = identity_node->name(); - TF_ASSIGN_OR_RETURN(Node * dst_replace_node, ReplaceNode(g, dst, new_def)); - - // Other edge in `edges` might have `e->dst()` as src or dst - // node. Before removing `e->dst()`, replace those edges with corresponding - // edges for `dst_replace_node`. - for (int j = i + 1; j < edges.size(); j++) { - if (edges[j].dst_node_id == edges[i].dst_node_id) { - edges[j].dst_node_id = dst_replace_node->id(); - } - } - } - return Status::OK(); -} - -// Step 3 for PreprocessForEncapsulation(). See comments of -// PreprocessForEncapsulation() for details. -Status ProcessDataEdgeBetweenOutsideCompilationAndHostComputation( - Graph* g, const string& xla_computation_attr_name, - const string& outside_compilation_attr_name) { - // Gather edges between outside compilation and host computation. Notice that - // we do not store `Edge*` directly because we remove some nodes while adding - // Identity nodes, and those Edge pointers might be invalidated. - struct EdgeInfo { - int dst_input, dst_node_id; - bool is_host_to_outside_compilation; - }; - std::vector edges; - for (const Edge* e : g->edges()) { - if (e->IsControlEdge()) { - continue; - } - - if (e->src()->attrs().Find(xla_computation_attr_name) == nullptr && - e->dst()->attrs().Find(xla_computation_attr_name) != nullptr && - e->dst()->attrs().Find(outside_compilation_attr_name) != nullptr) { - edges.push_back(EdgeInfo{e->dst_input(), e->dst()->id(), - /*is_host_to_outside_compilation=*/true}); - VLOG(4) << "Host -> oc edge: " << e->DebugString(); - } else if (e->dst()->attrs().Find(xla_computation_attr_name) == nullptr && - e->src()->attrs().Find(xla_computation_attr_name) != nullptr && - e->src()->attrs().Find(outside_compilation_attr_name) != - nullptr) { - edges.push_back(EdgeInfo{e->dst_input(), e->dst()->id(), - /*is_host_to_outside_compilation=*/false}); - VLOG(4) << "Oc -> host edge: " << e->DebugString(); + if (src_outside_compilation && dst_outside_compilation && + *src_outside_compilation != *dst_outside_compilation) { + edges.push_back(EdgeInfo{e->dst_input(), e->dst()->id()}); + VLOG(4) << "Oc -> oc edge: " << e->DebugString(); } } // Remove the edge from host to outside compilation. Add a placeholder as // outside compilation node input. - std::map placeholders; + std::map, Node*> placeholders; for (int i = 0; i < edges.size(); i++) { Node* dst = g->FindNodeId(edges[i].dst_node_id); const Edge* e; @@ -270,48 +143,33 @@ Status ProcessDataEdgeBetweenOutsideCompilationAndHostComputation( // Find or create placeholder node. string new_name = - edges[i].is_host_to_outside_compilation - ? absl::StrCat(src->name(), "_host_to_oc_placeholder") - : absl::StrCat(src->name(), "_oc_to_host_placeholder"); - auto iter = placeholders.find(new_name); + absl::StrCat(src->name(), "_oc_to_oc_placeholder_", src_output); + auto placeholder_index = std::make_pair(src->name(), src_output); + auto iter = placeholders.find(placeholder_index); Node* placeholder_node; if (iter == placeholders.end()) { NodeDefBuilder placeholder_builder(new_name, "Placeholder"); placeholder_builder.Attr("dtype", src->output_type(src_output)); - if (edges[i].is_host_to_outside_compilation) { - placeholder_builder.Attr(kHostToOutsideCompilationOriginalNodeAttrName, - src->name()); - placeholder_builder.Attr(kHostToOutsideCompilationSrcOutputAttrName, - src_output); - // If this placeholder node is in outside compilation, we need to set - // `xla_computation_attr_name` and `outside_compilation_attr_name`. - string xla_computation_attr, outside_compilation_attr; - TF_RETURN_IF_ERROR(GetNodeAttr(dst->attrs(), xla_computation_attr_name, - &xla_computation_attr)); - TF_RETURN_IF_ERROR(GetNodeAttr(dst->attrs(), - outside_compilation_attr_name, - &outside_compilation_attr)); - placeholder_builder.Attr(xla_computation_attr_name, - xla_computation_attr); - placeholder_builder.Attr(outside_compilation_attr_name, - outside_compilation_attr); - } else { - placeholder_builder.Attr(kOutsideCompilationToHostOriginalNodeAttrName, - src->name()); - placeholder_builder.Attr(kOutsideCompilationToHostSrcOutputAttrName, - src_output); - } + string outside_compilation_attr; + TF_RETURN_IF_ERROR(GetNodeAttr(dst->attrs(), + outside_compilation_attr_name, + &outside_compilation_attr)); + placeholder_builder.Attr(outside_compilation_attr_name, + outside_compilation_attr); + placeholder_builder.Attr(kOutsideCompilationOriginalNodeAttrName, + src->name()); + placeholder_builder.Attr(kOutsideCompilationSrcOutputAttrName, + src_output); NodeDef placeholder_def; TF_RETURN_IF_ERROR(placeholder_builder.Finalize(&placeholder_def)); Status s; placeholder_node = g->AddNode(placeholder_def, &s); TF_RETURN_IF_ERROR(s); - placeholders[new_name] = placeholder_node; + placeholders[placeholder_index] = placeholder_node; } else { placeholder_node = iter->second; } g->AddEdge(placeholder_node, 0, dst, dst_input); - g->RemoveEdge(e); // Replace `e->dst()` because its input node changed. NodeDef new_def = dst->def(); @@ -319,8 +177,8 @@ Status ProcessDataEdgeBetweenOutsideCompilationAndHostComputation( TF_ASSIGN_OR_RETURN(Node * dst_replace_node, ReplaceNode(g, dst, new_def)); // Other edge in `edges` might have `e->dst()` as src or dst - // node. Before removing `e->dst()`, replace those edges with corresponding - // edges for `dst_replace_node`. + // node. Before removing `e->dst()`, replace those edges with + // corresponding edges for `dst_replace_node`. for (int j = i + 1; j < edges.size(); j++) { if (edges[j].dst_node_id == edges[i].dst_node_id) { edges[j].dst_node_id = dst_replace_node->id(); @@ -330,55 +188,129 @@ Status ProcessDataEdgeBetweenOutsideCompilationAndHostComputation( return Status::OK(); } -} // namespace +// Step 1 for `PostprocessEdgesBetweenOutsideCompilations`. See comments of +// `PostprocessEdgesBetweenOutsideCompilations` for details. +Status PostprocessDataEdgesBetweenOutsideCompilations( + Graph* g, const string& outside_compilation_attr_name) { + // Gather all outside compilation to outside compilation nodes. + std::vector placeholder_nodes; + for (Node* n : g->nodes()) { + if (n->type_string() == "Placeholder" && + HasNodeAttr(n->def(), kOutsideCompilationOriginalNodeAttrName)) { + placeholder_nodes.push_back(n); + } + } -const char kXlaInferredShapesAttrName[] = "_xla_inferred_shapes"; + // Remove the placeholder nodes, and reconnect original edge. + auto node_name_index = g->BuildNodeNameIndex(); + for (auto n : placeholder_nodes) { + string node_name; + int node_src_output; + TF_RETURN_IF_ERROR(GetNodeAttr( + n->attrs(), kOutsideCompilationOriginalNodeAttrName, &node_name)); + TF_RETURN_IF_ERROR(GetNodeAttr( + n->attrs(), kOutsideCompilationSrcOutputAttrName, &node_src_output)); + auto iter = node_name_index.find(node_name); + if (iter == node_name_index.end()) { + return errors::Internal( + "Cannot find original node for oc -> host placeholder node ", + node_name); + } -const char kXlaConnectedToXlaComputationAttrName[] = - "_xla_connected_to_xla_computation"; -const char kXlaConnectedFromXlaComputationAttrName[] = - "_xla_connected_from_xla_computation"; -const char kXlaConnectedToOtherXlaComputationAttrName[] = - "_xla_connected_to_other_xla_computation"; -const char kXlaConnectedFromOtherXlaComputationAttrName[] = - "_xla_connected_from_other_xla_computation"; -const char kXlaControlDependenciesAttrName[] = "_xla_control_dependencies"; -const char kBridgeSourceNodeAttrName[] = "_xla_bridge_src"; -const char kOutsideCompilationToHostOriginalNodeAttrName[] = - "_xla_oc_to_host_node_name"; -const char kOutsideCompilationToHostSrcOutputAttrName[] = - "_xla_oc_to_host_src_output"; -const char kHostToOutsideCompilationOriginalNodeAttrName[] = - "_xla_host_to_oc_node_name"; -const char kHostToOutsideCompilationSrcOutputAttrName[] = - "_xla_host_to_oc_src_output"; - -Status PerformStaticShapeInferenceBeforeEncapsulation( - Graph* g, const string& xla_computation_attr_name, - const string& outside_compilation_attr_name) { - // Find all outside compilation to XLA computation data edges. - std::unordered_set outside_compilation_send_nodes; - for (auto e : g->edges()) { - if (e->IsControlEdge()) { - continue; + // Change all usage node to use the original node instead. + Node* original_node = iter->second; + std::vector control_edges; + std::vector data_edges; + for (auto e : n->out_edges()) { + if (e->IsControlEdge()) { + control_edges.push_back(e); + } else { + data_edges.push_back({e->dst(), e->src_output(), e->dst_input()}); + } + } + for (const Edge* e : control_edges) { + g->AddControlEdge(original_node, e->dst()); + g->RemoveEdge(e); } + for (int i = 0; i < data_edges.size(); i++) { + Node* dst = data_edges[i].dst; + NodeDef new_def = dst->def(); + int dst_input = data_edges[i].dst_input; + *new_def.mutable_input(dst_input) = + absl::StrCat(original_node->name(), ":", node_src_output); + TF_ASSIGN_OR_RETURN(Node * replace_node, ReplaceNode(g, dst, new_def)); + + const Edge* edge_to_replace = nullptr; + TF_RETURN_IF_ERROR(replace_node->input_edge(dst_input, &edge_to_replace)); + g->RemoveEdge(edge_to_replace); + g->AddEdge(original_node, node_src_output, replace_node, dst_input); + + // Other edges might have `dst` as dst node. Update those edges with + // `replace_node`. + for (int j = i + 1; j < data_edges.size(); j++) { + if (data_edges[j].dst == dst) { + data_edges[j].dst = replace_node; + } + } - auto src_computation = GetStringAttr(*e->src(), xla_computation_attr_name); - auto dst_computation = GetStringAttr(*e->dst(), xla_computation_attr_name); - if (!src_computation || !dst_computation || - *src_computation != *dst_computation) { - continue; + // Other placeholder node might have `dst` as original node. Update + // `node_name_index` with `replace_node`. + node_name_index[replace_node->name()] = replace_node; } - auto src_outside_compilation = - GetStringAttr(*e->src(), outside_compilation_attr_name); - auto dst_outside_compilation = - GetStringAttr(*e->dst(), outside_compilation_attr_name); - if (src_outside_compilation && !dst_outside_compilation) { - outside_compilation_send_nodes.insert(e->src()); + // Remove placeholder node. + g->RemoveNode(n); + } + return Status::OK(); +} + +// Step 2 for `PostprocessEdgesBetweenOutsideCompilations`. See comments of +// `PostprocessEdgesBetweenOutsideCompilations` for details. +Status PostprocessControlEdgesBetweenOutsideCompilations( + Graph* g, const string& outside_compilation_attr_name) { + auto node_name_index = g->BuildNodeNameIndex(); + + // Reconnect outside compilation to outside compilation control edge. + for (Node* n : g->nodes()) { + std::vector control_deps; + Status s = + GetNodeAttr(n->attrs(), kXlaControlDependenciesWithinXlaClusterAttrName, + &control_deps); + if (!s.ok()) { + if (s.code() != error::NOT_FOUND) { + return s; + } else { + continue; + } + } else { + n->ClearAttr(kXlaControlDependenciesWithinXlaClusterAttrName); + for (const string& control_input : control_deps) { + auto iter = node_name_index.find(control_input); + if (iter == node_name_index.end()) { + return errors::Internal("Cannot find original node for ", + control_input); + } + g->AddControlEdge(iter->second, n); + } } } + return Status::OK(); +} +} // namespace + +const char kXlaInferredShapesAttrName[] = "_xla_inferred_shapes"; +const char kXlaConnectedToXlaComputationAttrName[] = + "_xla_connected_to_xla_computation"; +const char kXlaConnectedFromXlaComputationAttrName[] = + "_xla_connected_from_xla_computation"; +const char kOutsideCompilationOriginalNodeAttrName[] = + "_xla_oc_to_oc_node_name"; +const char kOutsideCompilationSrcOutputAttrName[] = "_xla_oc_to_oc_src_output"; +const char kXlaControlDependenciesWithinXlaClusterAttrName[] = + "_xla_control_dependencies_within_xla_cluster"; + +Status PerformStaticShapeInferenceBeforeEncapsulation(Graph* g) { // Perform shape inference. std::map arg_shapes; GraphShapeInfo shape_info; @@ -386,33 +318,53 @@ Status PerformStaticShapeInferenceBeforeEncapsulation( InferShapes(g, arg_shapes, /*fnlib_def=*/nullptr, &shape_info)); // Add attribute for output shapes. - for (Node* n : outside_compilation_send_nodes) { - auto iter = shape_info.find(n->name()); - if (iter == shape_info.end()) { - continue; - } - + auto node_name_index = g->BuildNodeNameIndex(); + for (auto iter : shape_info) { std::vector output_shapes; - std::transform(iter->second.begin(), iter->second.end(), + std::transform(iter.second.begin(), iter.second.end(), std::back_inserter(output_shapes), [](const InferredShape& inferred_shape) { return inferred_shape.shape; }); + Node* n = node_name_index[iter.first]; n->AddAttr(kXlaInferredShapesAttrName, output_shapes); } return Status::OK(); } -Status PreprocessForEncapsulation(Graph* g, - const string& xla_computation_attr_name, - const string& outside_compilation_attr_name) { - TF_RETURN_IF_ERROR(ProcessControlEdges(g, xla_computation_attr_name, - outside_compilation_attr_name)); - TF_RETURN_IF_ERROR(ProcessXlaToXlaDataEdges(g, xla_computation_attr_name, - outside_compilation_attr_name)); - TF_RETURN_IF_ERROR(ProcessDataEdgeBetweenOutsideCompilationAndHostComputation( - g, xla_computation_attr_name, outside_compilation_attr_name)); +Status PreprocessEdgesBetweenOutsideCompilations( + Graph* g, const string& outside_compilation_attr_name) { + // Remove edges from source node to outside compilation nodes, and edges + // from outside compilation nodes to sink node. + std::vector edges_to_remove; + for (const Edge* e : g->source_node()->out_edges()) { + if (HasNodeAttr(e->dst()->def(), outside_compilation_attr_name)) { + edges_to_remove.push_back(e); + } + } + for (const Edge* e : g->sink_node()->in_edges()) { + if (HasNodeAttr(e->src()->def(), outside_compilation_attr_name)) { + edges_to_remove.push_back(e); + } + } + for (auto e : edges_to_remove) { + g->RemoveEdge(e); + } + + TF_RETURN_IF_ERROR(PreprocessControlEdgesBetweenOutsideCompilations( + g, outside_compilation_attr_name)); + TF_RETURN_IF_ERROR(PreprocessDataEdgesBetweenOutsideCompilations( + g, outside_compilation_attr_name)); + return Status::OK(); +} + +Status PostprocessEdgesBetweenOutsideCompilations( + Graph* g, const string& outside_compilation_attr_name) { + TF_RETURN_IF_ERROR(PostprocessDataEdgesBetweenOutsideCompilations( + g, outside_compilation_attr_name)); + TF_RETURN_IF_ERROR(PostprocessControlEdgesBetweenOutsideCompilations( + g, outside_compilation_attr_name)); return Status::OK(); } diff --git a/tensorflow/compiler/jit/encapsulate_util.h b/tensorflow/compiler/jit/encapsulate_util.h index bd76c844c491de0e20bfd83eb3d3ea8d38e4e60c..c9f16d14168163e11bb19092f566f1de8724aca3 100644 --- a/tensorflow/compiler/jit/encapsulate_util.h +++ b/tensorflow/compiler/jit/encapsulate_util.h @@ -27,22 +27,13 @@ namespace tensorflow { // a list of PartialTensorShape objects. extern const char kXlaInferredShapesAttrName[]; -// Infer output shapes for outside compilation nodes which have output data -// edges to XLA computation nodes. These shapes will be used later by XLA -// compiler as output shapes of the outside compilation's XlaHostCompute op. -// XLA computation nodes will be mark by attr `xla_computation_attr_name`; -// outside compilation nodes will be marked by both attr -// `xla_computation_attr_name` and `outside_compilation_attr_name`. -// -// Those outside compilation nodes will be marked with attribute -// `kXlaInferredShapesAttrName`. +// Infers output shapes for all nodes in graph `g`. The output shapes will be +// stored in node attribute `kXlaInferredShapesAttrName`. // // We have to perform shape inference before encapsulation because after // encapsulation, some nodes will be encapsulated into function call, and shape // inference does not handle function call at the moment. -Status PerformStaticShapeInferenceBeforeEncapsulation( - Graph* g, const string& xla_computation_attr_name, - const string& outside_compilation_attr_name); +Status PerformStaticShapeInferenceBeforeEncapsulation(Graph* g); // Attribute indicating that some ops in this node's XLA computation has control // dependency on this node. Attribute value will always be "true". @@ -52,69 +43,81 @@ extern const char kXlaConnectedToXlaComputationAttrName[]; // this node's XLA computation. Attribute value will always be "true". extern const char kXlaConnectedFromXlaComputationAttrName[]; -// Attribute indicating that some ops in other XLA computation has control -// dependency on this node. Attribute value will be a list of string (XLA -// computation names). -extern const char kXlaConnectedToOtherXlaComputationAttrName[]; - -// Attribute indicating that this node has control dependency on some ops in -// other XLA computation. Attribute value will be a list of string (XLA -// computation names). -extern const char kXlaConnectedFromOtherXlaComputationAttrName[]; - -// Attribute indicating that this node has control dependencies on some other -// nodes. Attribute value will be a list of string (node names). -extern const char kXlaControlDependenciesAttrName[]; - -// Attribute indicating that this is an Identity node added to act as a bridge -// between different XLA computations. Attribute value will be string (source -// node name). -extern const char kBridgeSourceNodeAttrName[]; - // Attribute indicating that this is an Placeholder node added to act as a // temporary input node for an outside compilation node. Attribute value will be // string (original input node name). -extern const char kOutsideCompilationToHostOriginalNodeAttrName[]; +extern const char kOutsideCompilationOriginalNodeAttrName[]; // Attribute indicating that this is an Placeholder node added to act as a // temporary input node for an outside compilation node. Attribute value will be // int (src_output for original edge). -extern const char kOutsideCompilationToHostSrcOutputAttrName[]; - -// Attribute indicating that this is an Placeholder node added to act as a -// temporary input node for an host node. Attribute value will be string -// (original input node name). -extern const char kHostToOutsideCompilationOriginalNodeAttrName[]; - -// Attribute indicating that this is an Placeholder node added to act as a -// temporary input node for a host node. Attribute value will be int (src_output -// for original edge). -extern const char kHostToOutsideCompilationSrcOutputAttrName[]; +extern const char kOutsideCompilationSrcOutputAttrName[]; -// Preprocesses the graph for encapsulation. It will perform the following +// Attribute indicating that this node has control dependencies on some other +// nodes within the same XLA cluster. Attribute value will be a list of string +// (node names). +extern const char kXlaControlDependenciesWithinXlaClusterAttrName[]; + +// Information for XLA computation. +struct XlaClusterInfo { + // Add an explicitly-defined default constructor for this class. + // + // The compiler may delete the default constructor here because + // host_compute_core is a const member whose type (std::map) doesn't + // necessarily have a user provided constructor -- while libc++ and + // libstdc++ 4.8 provide a user defined default constructor, libstdc++ at + // least >= 7.3 does not. See also c++11 [class.ctor] p5. + // + // TODO(klimek): In c++17 we'll be able to initialize host_compute_core + // without losing aggregate initialization, which allows us to get rid of + // the constructor definitions again. + XlaClusterInfo() {} + XlaClusterInfo(const string& cluster_name, + const NameAttrList& func_name_attrs, Node* node, + const std::map& host_compute_core) + : cluster_name(cluster_name), + func_name_attrs(func_name_attrs), + node(node), + host_compute_core(host_compute_core) {} + // XLA cluster name. It might be different from `func_name`. + const string cluster_name; + // Name and attributes of XLA computation function. + const NameAttrList func_name_attrs; + // The XLA computation node in the graph. + Node* node; + // A mapping from outside compilation cluster name to its device assignment. + const std::map host_compute_core; +}; + +// Preprocesses edges within the same XLA cluster. It will perform the following // operations in order: // -// 1a. For control edges between outside compilation and its XLA computation, +// 0. Remove edges from source node to outside compilation nodes, and edges +// from outside compilation nodes to sink node. +// 1a. For edges between different outside compilation clusters, remove the edge +// and add attr "kXlaControlDependenciesWithinXlaClusterAttrName = src node +// name" to dst node. +// 1b. For control edges between outside compilation and its XLA computation, // add attr "kXlaConnected{From, To}XlaComputationAttrName = true" to the // outside compilation node. -// 1b. For control edges between outside compilation and another XLA -// computation, add attr "kXlaConnected{From, To}OtherXlaComputationAttrName -// = XLA computation node name" to the outside compilation node. -// 1c. For control edges between different outside compilations, remove the edge -// and add attr "kXlaControlDependenciesAttrName = src node name" to dst -// node. -// 1d. For control edges between outside compilation and host computation, -// remove the edge and add attr "kXlaControlDependenciesAttrName = src node -// name" to dst node. -// 2. For data edges between different XLA computations, if either src or dst -// is outside compilation, add an Identity node in between the edge. The -// identity node will have attr kBridgeSourceNodeAttrName. -// 3. For data edges between outside compilation and host computation, remove -// the edge and create a Placeholder node as dst node's input. -Status PreprocessForEncapsulation(Graph* g, - const string& xla_computation_attr_name, - const string& outside_compilation_attr_name); - +// 2. For data edges between different outside compilations, remove the edge +// and create a Placeholder node as dst node's input. +Status PreprocessEdgesBetweenOutsideCompilations( + Graph* g, const string& outside_compilation_attr_name); + +// Postprocesses edges within the same XLA cluster. This function reverts what +// `PreprocessEdgesBetweenOutsideCompilations` did. It will perform the +// following operations in order: +// +// 1. Remove Placeholder nodes between different outside compilations (created +// in `PreprocessEdgesBetweenOutsideCompilations` step 2). +// 2a. Reconnect control edges between different outside compilations (marked by +// `PreprocessEdgesBetweenOutsideCompilations` step 1a). +// Notice that control edges marked by +// `PreprocessEdgesBetweenOutsideCompilations` step 1b are not handled here. +// They are handled in `RewriteOutsideCompilationSubgraphFn`. +Status PostprocessEdgesBetweenOutsideCompilations( + Graph* g, const string& outside_compilation_attr_name); } // namespace tensorflow #endif // TENSORFLOW_COMPILER_JIT_ENCAPSULATE_UTIL_H_ diff --git a/tensorflow/compiler/jit/encapsulate_util_test.cc b/tensorflow/compiler/jit/encapsulate_util_test.cc index 18422f73b5bc4dcd1af7fffe89211b93f7b578e8..3bb979e0698d2d6be42ed5bae66c25267928192c 100644 --- a/tensorflow/compiler/jit/encapsulate_util_test.cc +++ b/tensorflow/compiler/jit/encapsulate_util_test.cc @@ -21,6 +21,7 @@ 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 { @@ -37,24 +38,11 @@ TEST(PerformStaticShapeInferenceBeforeEncapsulationTest, Basic) { Graph g(OpRegistry::Global()); TF_CHECK_OK(s.ToGraph(&g)); - // "add" node is outside compilation node, "identity" node is XLA node. - auto node_index = g.BuildNodeNameIndex(); - Node *add_node = node_index["add"], *identity_node = node_index["identity"]; - add_node->AddAttr("_xla", "cluster"); - add_node->AddAttr("_oc", "cluster"); - identity_node->AddAttr("_xla", "cluster"); - TF_CHECK_OK( - PerformStaticShapeInferenceBeforeEncapsulation(&g, "_xla", "_oc")); + TF_CHECK_OK(PerformStaticShapeInferenceBeforeEncapsulation(&g)); - // Check that only "add" node now has _xla_inferred_shapes attr. - std::vector nodes_with_inferred_shape; - for (Node *n : g.nodes()) { - if (HasNodeAttr(n->def(), kXlaInferredShapesAttrName)) { - nodes_with_inferred_shape.push_back(n); - } - } - EXPECT_EQ(nodes_with_inferred_shape.size(), 1); - EXPECT_EQ(nodes_with_inferred_shape[0], add_node); + // Check that "add" node now has _xla_inferred_shapes attr. + auto node_index = g.BuildNodeNameIndex(); + Node *add_node = node_index["add"]; std::vector output_shapes; TF_CHECK_OK(GetNodeAttr(add_node->attrs(), kXlaInferredShapesAttrName, &output_shapes)); @@ -65,175 +53,4 @@ TEST(PerformStaticShapeInferenceBeforeEncapsulationTest, Basic) { EXPECT_EQ(shape_proto.dim(0).size(), 2); } -TEST(PreprocessForEncapsulationTest, ControlEdges) { - // Build the graph: - // "const_0" and "const_1" in host computation - // "add" = "const_0" + "const_1" in XLA computation 0 - // "identity0" = "add" in XLA computation 0 & outside compilation 0 - // "identity1" = "identity0" in XLA computation 0 - // "identity2" = "identity1" in host computation - // "identity3" = "identity2" in XLA computation 1 - // "identity4" = "identity3" in XLA computation 1 & outside compilation 1 - // "identity5" = "identity4" in XLA computation 1 - // "identity6" = "identity5" in host computation - tensorflow::Scope s = tensorflow::Scope::NewRootScope(); - Output const_0 = ops::Const(s.WithOpName("const_0"), 1, {}); - Output const_1 = ops::Const(s.WithOpName("const_1"), 2, {}); - Output add = ops::Add(s.WithOpName("add"), const_0, const_1); - Output identity0 = ops::Identity(s.WithOpName("identity0"), add); - Output identity1 = ops::Identity(s.WithOpName("identity1"), identity0); - Output identity2 = ops::Identity(s.WithOpName("identity2"), identity1); - Output identity3 = ops::Identity(s.WithOpName("identity3"), identity2); - Output identity4 = ops::Identity(s.WithOpName("identity4"), identity3); - Output identity5 = ops::Identity(s.WithOpName("identity5"), identity4); - Graph g(OpRegistry::Global()); - TF_CHECK_OK(s.ToGraph(&g)); - auto node_index = g.BuildNodeNameIndex(); - - // Set XLA computation/outside compilation attr, and add control edges. - Node *const0_node = node_index["const_0"], *add_node = node_index["add"], - *identity0_node = node_index["identity0"], - *identity1_node = node_index["identity1"], - *identity2_node = node_index["identity2"], - *identity3_node = node_index["identity3"], - *identity4_node = node_index["identity4"], - *identity5_node = node_index["identity5"]; - add_node->AddAttr("_xla", "0"); - identity0_node->AddAttr("_xla", "0"); - identity0_node->AddAttr("_oc", "0"); - identity1_node->AddAttr("_xla", "0"); - identity3_node->AddAttr("_xla", "1"); - identity4_node->AddAttr("_xla", "1"); - identity4_node->AddAttr("_oc", "0"); - identity5_node->AddAttr("_xla", "1"); - // Case 1a: control edges between outside compilation and its XLA computation. - g.AddControlEdge(add_node, identity0_node); - g.AddControlEdge(identity0_node, identity1_node); - // Case 1b: control edges between outside compilation and another XLA - // computation. - g.AddControlEdge(identity0_node, identity3_node); - g.AddControlEdge(identity1_node, identity4_node); - // Case 1c: control edges between different outside compilations. - g.AddControlEdge(identity0_node, identity4_node); - // Case 1d: control edges between outside compilation and host computation. - g.AddControlEdge(const0_node, identity0_node); - g.AddControlEdge(identity0_node, identity2_node); - - TF_CHECK_OK(PreprocessForEncapsulation(&g, "_xla", "_oc")); - - // Case 1a: add attr "_xla_connected_{from/to}_xla_computation = true" to the - // outside compilation node. - EXPECT_TRUE(HasNodeAttr(identity0_node->def(), - kXlaConnectedFromXlaComputationAttrName)); - EXPECT_TRUE(HasNodeAttr(identity0_node->def(), - kXlaConnectedToXlaComputationAttrName)); - // Case 1b: add attr "_xla_control_deps_{from/to} = XLA computation node name" - // to the outside compilation node. - std::vector attr; - TF_CHECK_OK(GetNodeAttr(identity0_node->def(), - kXlaConnectedToOtherXlaComputationAttrName, &attr)); - EXPECT_EQ(attr.size(), 1); - EXPECT_EQ(attr[0], "1"); - attr.clear(); - TF_CHECK_OK(GetNodeAttr(identity4_node->def(), - kXlaConnectedFromOtherXlaComputationAttrName, &attr)); - EXPECT_EQ(attr.size(), 1); - EXPECT_EQ(attr[0], "0"); - // Case 1c: add attr "_xla_control_deps = src node name" to dst node. - attr.clear(); - TF_CHECK_OK(GetNodeAttr(identity4_node->def(), - kXlaControlDependenciesAttrName, &attr)); - EXPECT_EQ(attr.size(), 1); - EXPECT_EQ(attr[0], "identity0"); - // Case 1d: add attr "_xla_control_deps = src node name" to dst node. - attr.clear(); - TF_CHECK_OK(GetNodeAttr(identity0_node->def(), - kXlaControlDependenciesAttrName, &attr)); - EXPECT_EQ(attr.size(), 1); - EXPECT_EQ(attr[0], "const_0"); - attr.clear(); - TF_CHECK_OK(GetNodeAttr(identity2_node->def(), - kXlaControlDependenciesAttrName, &attr)); - EXPECT_EQ(attr.size(), 1); - EXPECT_EQ(attr[0], "identity0"); -} - -TEST(PreprocessForEncapsulationTest, DataEdges) { - // Build the graph: - // "const_0" and "const_1" in host computation - // "add0" = "const_0" + "const_1" in XLA computation 0 - // "add1" = "add0" + "const_0" in XLA computation 0 & outside compilation 0 - // "identity0" = "add1" in XLA computation 0 - // "add2" = "add1" + "identity0" in host computation - // "add3" = "add1" + "add2" in XLA computation 1 - // "add4" = "identity0" + "add2" in XLA computation 1 & outside compilation 1 - // "identity1" = "add4" in XLA computation 1 - // "identity2" = "identity1" in host computation - tensorflow::Scope s = tensorflow::Scope::NewRootScope(); - Output const_0 = ops::Const(s.WithOpName("const_0"), 1, {}); - Output const_1 = ops::Const(s.WithOpName("const_1"), 2, {}); - Output add0 = ops::Add(s.WithOpName("add0"), const_0, const_1); - Output add1 = ops::Add(s.WithOpName("add1"), add0, const_0); - Output identity0 = ops::Identity(s.WithOpName("identity0"), add1); - Output add2 = ops::Add(s.WithOpName("add2"), add1, identity0); - Output add3 = ops::Add(s.WithOpName("add3"), add1, add2); - Output add4 = ops::Add(s.WithOpName("add4"), identity0, add2); - Output identity1 = ops::Identity(s.WithOpName("identity1"), add4); - Output identity2 = ops::Identity(s.WithOpName("identity2"), add4); - Graph g(OpRegistry::Global()); - TF_CHECK_OK(s.ToGraph(&g)); - auto node_index = g.BuildNodeNameIndex(); - - // Set XLA computation/outside compilation attr. - Node *add0_node = node_index["add0"], *add1_node = node_index["add1"], - *identity0_node = node_index["identity0"], - *add3_node = node_index["add3"], *add4_node = node_index["add4"], - *identity1_node = node_index["identity1"]; - add0_node->AddAttr("_xla", "0"); - add1_node->AddAttr("_xla", "0"); - add1_node->AddAttr("_oc", "0"); - identity0_node->AddAttr("_xla", "0"); - add3_node->AddAttr("_xla", "1"); - add4_node->AddAttr("_xla", "1"); - add4_node->AddAttr("_oc", "0"); - identity1_node->AddAttr("_xla", "1"); - - TF_CHECK_OK(PreprocessForEncapsulation(&g, "_xla", "_oc")); - - // Check input nodes for related data edges. - node_index = g.BuildNodeNameIndex(); - // Step 2: add an Identity node between different XLA computations. - Node *bridge_add1_add3 = node_index["bridge_add1_add3"]; - EXPECT_NE(bridge_add1_add3, nullptr); - string str; - TF_CHECK_OK( - GetNodeAttr(bridge_add1_add3->attrs(), kBridgeSourceNodeAttrName, &str)); - EXPECT_EQ(str, "add1"); - Node *bridge_identity0_add4 = node_index["bridge_identity0_add4"]; - EXPECT_NE(bridge_identity0_add4, nullptr); - // Step 3: add placeholder for edges between host computation and outside - // compilation. - EXPECT_EQ(bridge_add1_add3->def().input(0), "add1_oc_to_host_placeholder"); - Node *add1_oc_to_host_placeholder = node_index["add1_oc_to_host_placeholder"]; - TF_CHECK_OK(GetNodeAttr(add1_oc_to_host_placeholder->attrs(), - kOutsideCompilationToHostOriginalNodeAttrName, &str)); - EXPECT_EQ(str, "add1"); - int i; - TF_CHECK_OK(GetNodeAttr(add1_oc_to_host_placeholder->attrs(), - kOutsideCompilationToHostSrcOutputAttrName, &i)); - EXPECT_EQ(i, 0); - add4_node = node_index["add4"]; - ASSERT_NE(add4_node, nullptr); - EXPECT_EQ(add4_node->def().input(0), - "bridge_identity0_add4_host_to_oc_placeholder"); - Node *identity0_host_to_oc_placeholder = - node_index["bridge_identity0_add4_host_to_oc_placeholder"]; - TF_CHECK_OK(GetNodeAttr(identity0_host_to_oc_placeholder->attrs(), - kHostToOutsideCompilationOriginalNodeAttrName, &str)); - EXPECT_EQ(str, "bridge_identity0_add4"); - TF_CHECK_OK(GetNodeAttr(identity0_host_to_oc_placeholder->attrs(), - kHostToOutsideCompilationSrcOutputAttrName, &i)); - EXPECT_EQ(i, 0); -} - } // namespace tensorflow diff --git a/tensorflow/compiler/jit/encapsulate_xla_computations_pass.cc b/tensorflow/compiler/jit/encapsulate_xla_computations_pass.cc index 2ce6fa73fc448ca83fa392aa909cb385453eb8b6..ec745cdbb7e237f8b4935dd41e9791fc75f5355d 100644 --- a/tensorflow/compiler/jit/encapsulate_xla_computations_pass.cc +++ b/tensorflow/compiler/jit/encapsulate_xla_computations_pass.cc @@ -195,8 +195,11 @@ Status RewriteSubgraph(const std::vector& arg_source_tensors, e->dst()->attrs().Find(kXlaClusterAttr) == nullptr && e->dst()->type_string() != kXlaClusterOutput) { return errors::InvalidArgument( - "Undeclared output of XLA computation. A common cause of this error " - "is variable initializers that depend on the XLA computation. Edge: ", + "Undeclared output of XLA computation. Some common causes of this " + "error are: 1) variable initializers that depend on the XLA " + "computation; 2) gradient computations that depend on the XLA " + "computation, which can be mitigated by moving gradient computations " + "inside XLA computation. Offending edge: ", e->src()->name(), ":", e->src_output(), " -> ", e->dst()->name(), ":", e->dst_input()); } @@ -294,6 +297,7 @@ Status RewriteSubgraph(const std::vector& arg_source_tensors, NodeDef def; def.set_name(launch->name()); + MergeDebugInfo(NodeDebugInfo(launch->def()), &def); // Target the XLA CPU/GPU backends. VLOG(2) << "Replacing with XlaLaunch"; diff --git a/tensorflow/compiler/jit/extract_outside_compilation_pass.cc b/tensorflow/compiler/jit/extract_outside_compilation_pass.cc index 3963fb012e029358298c3cb721a89dbb7e758ea3..2a770c527b2fae91352fd17dacb13495a3a73f34 100644 --- a/tensorflow/compiler/jit/extract_outside_compilation_pass.cc +++ b/tensorflow/compiler/jit/extract_outside_compilation_pass.cc @@ -17,11 +17,20 @@ limitations under the License. #include "absl/strings/match.h" #include "absl/strings/str_cat.h" +#include "tensorflow/compiler/jit/encapsulate_subgraphs_pass.h" #include "tensorflow/compiler/jit/encapsulate_util.h" +#include "tensorflow/compiler/tf2xla/dump_graph.h" +#include "tensorflow/compiler/tf2xla/side_effect_util.h" #include "tensorflow/compiler/tf2xla/tf2xla_util.h" +#include "tensorflow/core/common_runtime/function.h" +#include "tensorflow/core/framework/function.h" +#include "tensorflow/core/framework/graph_to_functiondef.h" #include "tensorflow/core/framework/node_def_builder.h" +#include "tensorflow/core/framework/node_def_util.h" #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 { @@ -45,6 +54,12 @@ xla::StatusOr AddHostComputeKeyPlaceholder( return n; } +// Returns if the node is a XLA computation key placeholder. +bool IsKeyPlaceholderNode(const Node& n) { + return n.type_string() == "Placeholder" && + absl::EndsWith(n.name(), "_key_placeholder"); +} + // Returns nodes with given type. std::vector GatherNodesWithType(const Graph& g, const string& type) { std::vector result; @@ -86,9 +101,12 @@ xla::StatusOr BuildRecvAtHostNode( recv_at_host_builder.Attr("Toutputs", recv_at_host_dtypes); // The correct device_ordinal will be inserted during replication in a // subsequent rewrite. - recv_at_host_builder.Attr("device_ordinal", 0); + AttrValue device_ordinal_value; + device_ordinal_value.set_placeholder("device_ordinal"); + recv_at_host_builder.Attr("device_ordinal", device_ordinal_value); recv_at_host_builder.Attr( "key", absl::StrCat("host_compute_channel_", oc_cluster_name)); + recv_at_host_builder.Attr(kXlaHasHostTransferAttrName, true); recv_at_host_builder.Input(key_placeholder->name(), 0, DT_STRING); TF_RETURN_IF_ERROR(recv_at_host_builder.Finalize(&recv_at_host_def)); Status s; @@ -101,6 +119,8 @@ xla::StatusOr BuildRecvAtHostNode( xla::StatusOr ReplaceArgNodesWithRecvAtHostNode( Graph* g, const string& oc_cluster_name, std::vector* recv_at_host_dtypes, Node* key_placeholder) { + // TODO(b/77601805): use out nodes for source node, instead of traversing all + // nodes. std::vector arg_nodes = GatherNodesWithType(*g, "_Arg"); TF_RETURN_IF_ERROR(GetArgDataTypes(arg_nodes, recv_at_host_dtypes)); TF_ASSIGN_OR_RETURN( @@ -183,9 +203,12 @@ xla::StatusOr BuildSendFromHostNode( send_from_host_builder.Attr("Tinputs", send_from_host_dtypes); // The correct device_ordinal will be inserted during replication in a // subsequent rewrite. - send_from_host_builder.Attr("device_ordinal", 0); + AttrValue device_ordinal_value; + device_ordinal_value.set_placeholder("device_ordinal"); + send_from_host_builder.Attr("device_ordinal", device_ordinal_value); send_from_host_builder.Attr( "key", absl::StrCat("host_compute_channel_", oc_cluster_name)); + send_from_host_builder.Attr(kXlaHasHostTransferAttrName, true); std::vector inputs(send_from_host_dtypes.size()); for (auto* n : ret_nodes) { int index; @@ -212,6 +235,8 @@ xla::StatusOr BuildSendFromHostNode( xla::StatusOr ReplaceRetNodesWithSendFromHostNode( Graph* g, const string& oc_cluster_name, std::vector* send_from_host_dtypes, Node* key_placeholder) { + // TODO(b/77601805): use in nodes for sink node, instead of traversing all + // nodes. std::vector ret_nodes = GatherNodesWithType(*g, "_Retval"); TF_RETURN_IF_ERROR(GetRetDataTypes(ret_nodes, send_from_host_dtypes)); TF_ASSIGN_OR_RETURN( @@ -252,7 +277,7 @@ absl::optional> GetInferredInputShapes( return absl::nullopt; } - const PartialTensorShape shape = shapes[e->dst_input()]; + const PartialTensorShape shape = shapes[e->src_output()]; if (!shape.IsFullyDefined()) { return absl::nullopt; } @@ -262,6 +287,1196 @@ absl::optional> GetInferredInputShapes( return results; } +// Builds XlaHostCompute NodeDef from the outside compilation call node. +xla::StatusOr BuildXlaHostComputeNodeDef( + const Node* call_node, const std::map& host_compute_core) { + string original_oc_name; + TF_RETURN_IF_ERROR(GetNodeAttr( + call_node->attrs(), "_outside_compilation_subgraph", &original_oc_name)); + NodeDefBuilder host_compute_builder( + absl::StrCat("outside_compilation_", original_oc_name, "_host_compute"), + "XlaHostCompute"); + + // Copy all attributes. + for (auto attr : call_node->attrs()) { + host_compute_builder.Attr(attr.first, attr.second); + } + + // Populate tpu_core assignment. + const auto iter = host_compute_core.find(original_oc_name); + if (iter != host_compute_core.end()) { + int core = iter->second; + 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)); + std::vector inputs(input_dtypes.size()); + for (auto e : call_node->in_edges()) { + if (e->IsControlEdge()) { + continue; + } + + if (e->dst_input() < 0 || e->dst_input() >= input_dtypes.size()) { + return errors::Internal("Invalid dst_input: ", e->dst_input()); + } + inputs[e->dst_input()] = NodeDefBuilder::NodeOut{ + e->src()->name(), e->src_output(), input_dtypes[e->dst_input()]}; + } + host_compute_builder.Input(inputs); + + NodeDef new_def; + TF_RETURN_IF_ERROR(host_compute_builder.Finalize(&new_def)); + return new_def; +} + +Status ValidateOutsideCompilationCallNode(Node* call_node) { + // DT_INT64 as input/output for outside compilation is not supported yet: + // b/120809951. + for (const Edge* e : call_node->in_edges()) { + if (e->IsControlEdge()) { + continue; + } + DataType dtype = e->src()->output_type(e->src_output()); + if (dtype == DT_INT64) { + return errors::Unimplemented( + "int64 input for outside compilation is not supported yet: " + "b/120809951. Please cast output of node ", + e->src()->DebugString(), + " to int32 before feeding it into outside compilation."); + } + } + for (const Edge* e : call_node->out_edges()) { + if (e->IsControlEdge()) { + continue; + } + DataType dtype = e->dst()->input_type(e->dst_input()); + if (dtype == DT_INT64) { + return errors::Unimplemented( + "int64 output for outside compilation is not supported yet: " + "b/120809951. Please cast input of node ", + e->dst()->DebugString(), + " to int32 before returning it from outside compilation."); + } + } + return Status::OK(); +} + +// Replace outside compilation function call node with XlaHostCompute node. +// If the function call node has no input/output edges, we will just remove it +// and not create a XlaHostCompute node. +Status ReplaceOrRemoveOutsideCompilationCallNode( + Graph* g, Node* call_node, const std::map& host_compute_core) { + // If the function call node has no input/output edges, just remove it. + bool has_edge = false; + for (auto e : call_node->in_edges()) { + if (!e->IsControlEdge() || e->src() != g->source_node()) { + has_edge = true; + break; + } + } + for (auto e : call_node->out_edges()) { + if (!e->IsControlEdge() || e->dst() != g->sink_node()) { + has_edge = true; + break; + } + } + if (!has_edge) { + VLOG(4) << "Did not add HostCompute node for " << call_node->DebugString(); + g->RemoveNode(call_node); + return Status::OK(); + } + + // Build XlaHostCompute NodeDef. + TF_ASSIGN_OR_RETURN(NodeDef node_def, + BuildXlaHostComputeNodeDef(call_node, host_compute_core)); + TF_ASSIGN_OR_RETURN(Node * host_compute_node, + ReplaceNode(g, call_node, node_def)); + VLOG(4) << "Added HostCompute node: " << host_compute_node->DebugString(); + + return Status::OK(); +} + +// Resets "device_ordinal" attr to placeholder value for related nodes +// (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"); + for (Node* n : g->nodes()) { + if (!HasNodeAttr(n->def(), kXlaHasHostTransferAttrName)) { + continue; + } + + if (n->type_string() == "_XlaRecvAtHost" || + n->type_string() == "_XlaSendFromHost") { + n->ClearAttr("device_ordinal"); + n->AddAttr("device_ordinal", device_ordinal_value); + } else if (n->type_string() == "If") { + for (const string& attr_name : + std::vector{"then_branch", "else_branch"}) { + NameAttrList branch_func; + TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), attr_name, &branch_func)); + (*branch_func.mutable_attr())["device_ordinal"] = device_ordinal_value; + n->ClearAttr(attr_name); + n->AddAttr(attr_name, branch_func); + } + } else if (n->type_string() == "While") { + for (const string& attr_name : std::vector{"cond", "body"}) { + NameAttrList branch_func; + TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), attr_name, &branch_func)); + (*branch_func.mutable_attr())["device_ordinal"] = device_ordinal_value; + 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, ": ", + n->DebugString()); + } + } + return Status::OK(); +} + +// For an XLA computation, builds host side graph given all outside compilation +// graphs inside it. The host side graph contains: +// 1) a "sequencer" node (we will add control edge between XlaRecvAtHost and +// XlaSendFromHost to this sequencer node, so all outside compilation nodes +// will be executed *before* this sequencer). +// 2) a "key placeholder" node. Later in ExpandHostGraphIntoMainGraph(), we will +// replace this node with compilation result node. +// 3) all outside compilation graphs. +Status ConstructHostGraph( + const string& xla_cluster_name, const string& outside_compilation_attr_name, + const std::vector& outside_compilation_host_graphs, + FunctionLibraryDefinition* fld, const string& host_graph_func_name) { + Graph host_graph(fld); + + // Create sequencer node in host graph. + NodeDefBuilder sequencer_builder(absl::StrCat(xla_cluster_name, "_sequencer"), + "NoOp"); + sequencer_builder.Attr("_xla_host_transfer_sequencer", xla_cluster_name); + NodeDef sequencer_def; + TF_RETURN_IF_ERROR(sequencer_builder.Finalize(&sequencer_def)); + Status s; + Node* sequencer = host_graph.AddNode(sequencer_def, &s); + TF_RETURN_IF_ERROR(s); + + // Create key placeholder in host graph. + TF_ASSIGN_OR_RETURN( + Node * key_placeholder, + AddHostComputeKeyPlaceholder(xla_cluster_name, &host_graph)); + + // For each outside compilation graph, copy them to host graph with the + // following changes: + // a) Use key_placeholder in host graph instead of its own. + // b) Add control edge from host transfer nodes (XlaRecvAtHost, + // XlaSendFromHost, If/While nodes containing + // XlaRecvAtHost/XlaSendFromHost) to sequencer node. + // c) Clear node_def.device(), so device placer won't get confused. + for (const string& host_func : outside_compilation_host_graphs) { + VLOG(4) << "Expanding host graph " << host_func; + // Temporarily use "0" as "device_ordinal". It will be reset to placeholder + // value after we expanded all host graphs. We cannot just use placeholder + // value here because FunctionDef instantiation does not allow placeholder + // value for attributes. + AttrValue device_ordinal_attr; + device_ordinal_attr.set_i(0); + protobuf::Map attrs; + attrs["device_ordinal"] = device_ordinal_attr; + FunctionBody* host_fbody = nullptr; + TF_RETURN_IF_ERROR(FunctionDefToBodyHelper( + *fld->Find(host_func), AttrSlice(&attrs), fld, + [&](const string& op, const OpDef** sig) { + return fld->LookUpOpDef(op, sig); + }, + &host_fbody)); + std::unique_ptr host_fbody_deleter(host_fbody); + + // We use ReverseDFS() to copy nodes. Make sure all nodes are reverse + // reachable from sink node so all nodes will be copied. + // TODO(b/77601805): consolidate copy graph functions. + FixupSourceAndSinkEdges(host_fbody->graph); + + std::map node_map; + node_map[host_fbody->graph->source_node()] = host_graph.source_node(); + node_map[host_fbody->graph->sink_node()] = host_graph.sink_node(); + Status s; + ReverseDFS( + *host_fbody->graph, /*enter=*/nullptr, + [&](const Node* n) { + if (!s.ok()) { + return; + } + + Node* copy; + if (node_map.find(n) != node_map.end()) { + // Already copied this node. + copy = node_map.at(n); + } else if (IsKeyPlaceholderNode(*n)) { + // Change a). + copy = key_placeholder; + node_map[n] = copy; + } else { + // Copy the node. + NodeDef copy_def = n->def(); + // Change c). + copy_def.clear_device(); + copy = host_graph.AddNode(copy_def, &s); + if (!s.ok()) { + return; + } + node_map[n] = copy; + } + + // Only handle input edges. Output edges will be added later as + // its output nodes' input edges. + for (auto e : n->in_edges()) { + if (node_map.find(e->src()) == node_map.end()) { + s = errors::Internal("Cannot find node image for ", + e->src()->DebugString()); + return; + } + host_graph.AddEdge(node_map[e->src()], e->src_output(), copy, + e->dst_input()); + } + + // Change b). + if (HasNodeAttr(copy->def(), kXlaHasHostTransferAttrName)) { + host_graph.AddControlEdge(copy, sequencer); + } + }, + NodeComparatorID()); + + if (!s.ok()) { + return s; + } + } + // Reset "device_ordinal" to placeholder value. + TF_RETURN_IF_ERROR(ResetDeviceOrdinalToPlaceholderValue(&host_graph)); + + // sequencer and key_placeholder might be dead nodes. Prune them if necessary. + // - sequencer should be pruned iff it has no input control edges from + // RecvAtHost/SendFromHost. If it has input control edge, we connect it to + // sink node so it won't be pruned. + // - key_placeholder should be pruned iff there's no RecvAtHost/SendFromHost. + // We don't need to do anything special. + if (!sequencer->in_edges().empty()) { + host_graph.AddControlEdge(sequencer, host_graph.sink_node()); + } + PruneForReverseReachability( + &host_graph, std::unordered_set{host_graph.sink_node()}); + + // Postprocess edges between different outside compilations. + TF_RETURN_IF_ERROR(PostprocessEdgesBetweenOutsideCompilations( + &host_graph, outside_compilation_attr_name)); + + if (VLOG_IS_ON(4)) { + dump_graph::DumpGraphToFile( + absl::StrCat("extract_outside_compilation_host_graph_for_", + xla_cluster_name), + host_graph, fld); + } + + FunctionDef host_graph_fdef; + TF_RETURN_IF_ERROR( + GraphToFunctionDef(host_graph, host_graph_func_name, &host_graph_fdef)); + if (fld->Find(host_graph_func_name)) { + TF_RETURN_IF_ERROR( + fld->ReplaceFunction(host_graph_func_name, host_graph_fdef)); + } else { + TF_RETURN_IF_ERROR(fld->AddFunctionDef(host_graph_fdef)); + } + + return Status::OK(); +} + +// Expand XLA computation's outside compilation host side graph into main graph. +// Add a control edge between sequencer node and the XLA computation node. +Status ExpandHostGraphIntoMainGraph(Graph* main_graph, + FunctionLibraryDefinition* fld, + const string& host_graph_func_name, + Node* xla_computation_node) { + // Temporarily use "0" as "device_ordinal". It will be rewritten with the + // correct value in a later pass. We cannot just use placeholder value here + // because FunctionDef instantiation does not allow placeholder value for + // attributes. + AttrValue device_ordinal_attr; + device_ordinal_attr.set_i(0); + protobuf::Map attrs; + attrs["device_ordinal"] = device_ordinal_attr; + FunctionBody* fbody = nullptr; + TF_RETURN_IF_ERROR(FunctionDefToBodyHelper( + *fld->Find(host_graph_func_name), AttrSlice(&attrs), fld, + [&](const string& op, const OpDef** sig) { + return fld->LookUpOpDef(op, sig); + }, + &fbody)); + std::unique_ptr fbody_deleter(fbody); + Graph* host_graph = fbody->graph; + + // We use ReverseDFS() to copy nodes. Make sure all nodes are reverse + // reachable from sink node so all nodes will be copied. + // TODO(b/77601805): consolidate copy graph functions. + FixupSourceAndSinkEdges(host_graph); + + // Copy all nodes. + std::map node_map; + node_map[host_graph->source_node()] = main_graph->source_node(); + node_map[host_graph->sink_node()] = main_graph->sink_node(); + Status s = Status::OK(); + auto copy_node_fn = [&](const Node* n) { + if (!s.ok()) { + return; + } + + Node* copy; + if (node_map.find(n) != node_map.end()) { + // Already copied this node. + copy = node_map.at(n); + } else { + // Copy the node. + NodeDef copy_def = n->def(); + copy = main_graph->AddNode(copy_def, &s); + if (!s.ok()) { + return; + } + node_map[n] = copy; + } + + // Only handle input edges. Output edges will be added later as its output + // nodes' input edges. + for (auto e : n->in_edges()) { + if (node_map.find(e->src()) == node_map.end()) { + s = errors::Internal("Cannot find node image for ", + e->src()->DebugString()); + return; + } + main_graph->AddEdge(node_map[e->src()], e->src_output(), copy, + e->dst_input()); + } + + // Add control edge from sequencer to XLA computation node. + if (copy->type_string() == "NoOp" && + HasNodeAttr(copy->def(), "_xla_host_transfer_sequencer")) { + main_graph->AddControlEdge(copy, xla_computation_node); + } + }; + ReverseDFS(*host_graph, /*enter=*/nullptr, copy_node_fn, NodeComparatorID()); + return s; +} + +// Rewrites shape inference graph for outside compilation: +// 1) If XlaSendFromHost also exists in `host_graph`, copy nodes from +// `host_graph`. Because we might still have outside compilation to outside +// compilation placeholder nodes in shape inference graph, which will prevent +// us from inferring XlaSendFromHost shape. But in `host_graph`, we already +// removed those placeholder nodes. +// 2) Remove control edges. +// 3) Prune nodes that are not useful for shape inference. +Status RewriteShapeInferenceGraph(const string& shape_inference_graph_name, + Graph* host_graph, + FunctionLibraryDefinition* fld) { + // Use "0" as "device_ordinal". It does not matter for shape inference. + AttrValue device_ordinal_attr; + device_ordinal_attr.set_i(0); + protobuf::Map attrs; + attrs["device_ordinal"] = device_ordinal_attr; + FunctionBody* fbody = nullptr; + TF_RETURN_IF_ERROR(FunctionDefToBodyHelper( + *fld->Find(shape_inference_graph_name), AttrSlice(&attrs), fld, + [&](const string& op, const OpDef** sig) { + return fld->LookUpOpDef(op, sig); + }, + &fbody)); + std::unique_ptr fbody_deleter(fbody); + Graph* g = fbody->graph; + + // Find SendFromHost node. + Node* send_from_host = nullptr; + for (Node* n : g->nodes()) { + if (n->type_string() == "_XlaSendFromHost") { + send_from_host = n; + break; + } + } + if (!send_from_host) { + return errors::Internal("Shape inference graph ", + shape_inference_graph_name, + " does not have _XlaSendFromHost node."); + } + + // See if the SendFromHost node exists in `host_graph`. + Node* send_from_host_main_graph = nullptr; + for (Node* n : host_graph->nodes()) { + if (n->name() == send_from_host->name()) { + send_from_host_main_graph = n; + break; + } + } + if (send_from_host_main_graph) { + // This is an "top-level" outside compilation. Clear the graph, and copy + // SendFromHost and all its predecessors from `host_graph`. + std::vector nodes; + for (Node* n : g->op_nodes()) { + nodes.push_back(n); + } + for (Node* n : nodes) { + g->RemoveNode(n); + } + + std::map node_map; + node_map[host_graph->source_node()] = g->source_node(); + Status s; + auto copy_node_fn = [&](const Node* n) { + if (!s.ok()) { + return; + } + + if (node_map.find(n) != node_map.end()) { + return; + } + + NodeDef copy_def = n->def(); + Node* copy = g->AddNode(copy_def, &s); + if (!s.ok()) { + return; + } + for (auto e : n->in_edges()) { + if (node_map.find(e->src()) == node_map.end()) { + s = errors::Internal("Cannot find node image for ", + e->src()->DebugString()); + return; + } + g->AddEdge(node_map[e->src()], e->src_output(), copy, e->dst_input()); + } + + node_map[n] = copy; + }; + // TODO(b/77601805): consolidate copy graph functions. + ReverseDFSFrom(*host_graph, + std::vector{send_from_host_main_graph}, + /*enter=*/nullptr, copy_node_fn, NodeComparatorID()); + if (!s.ok()) { + return s; + } + + send_from_host = node_map[send_from_host_main_graph]; + } else { + // This is an outside compilation embedded in If/While/gradient/etc. + // It will be enough for shape inference. Leave `g` unchanged. + } + + // Control edges are not useful for shape inference. Remove them. + for (auto e : g->edges()) { + if (e->IsControlEdge()) { + g->RemoveEdge(e); + } + } + + // Nodes that are not reverse reachable from SendFromHost are not useful for + // shape inference. Prune them. + PruneForReverseReachability(g, + std::unordered_set{send_from_host}); + + if (VLOG_IS_ON(4)) { + dump_graph::DumpGraphToFile(shape_inference_graph_name, *g, fld); + } + + // Replace original shape inference graph. + FunctionDef fdef_replace; + TF_RETURN_IF_ERROR( + GraphToFunctionDef(*g, shape_inference_graph_name, &fdef_replace)); + TF_RETURN_IF_ERROR( + fld->ReplaceFunction(shape_inference_graph_name, fdef_replace)); + + return Status::OK(); +} + +// Builds XlaSendToHost node which sends cond predicate to host. +xla::StatusOr BuildSendIfPredNode(const string& name, + const string& host_transfer_key, + Node* pred_node, Graph* g) { + NodeDefBuilder send_pred_builder(name, "XlaSendToHost"); + send_pred_builder.Attr("Tinput", DT_BOOL); + send_pred_builder.Attr("key", absl::StrCat(host_transfer_key, "_dtoh_0")); + send_pred_builder.Attr(kXlaTokenInputNodesAttrName, + std::vector{kXlaTokenArgNodeName}); + send_pred_builder.Input(pred_node->name(), 0, DT_BOOL); + NodeDef send_pred_def; + TF_RETURN_IF_ERROR(send_pred_builder.Finalize(&send_pred_def)); + Status s; + Node* send_pred_node = g->AddNode(send_pred_def, &s); + TF_RETURN_IF_ERROR(s); + g->AddEdge(pred_node, 0, send_pred_node, 0); + return send_pred_node; +} + +// Replaces key placeholder node with an _Arg node. +Status ReplaceKeyPlaceholderWithArgNode(const string& xla_cluster_name, + const string& func_name, + FunctionLibraryDefinition* fld) { + // Temporarily use "0" as "device_ordinal". It will be reset to placeholder + // value after rewriting. + AttrValue device_ordinal_attr; + device_ordinal_attr.set_i(0); + protobuf::Map attrs; + attrs["device_ordinal"] = device_ordinal_attr; + FunctionBody* fbody = nullptr; + TF_RETURN_IF_ERROR(FunctionDefToBodyHelper( + *fld->Find(func_name), AttrSlice(&attrs), fld, + [&](const string& op, const OpDef** sig) { + return fld->LookUpOpDef(op, sig); + }, + &fbody)); + std::unique_ptr fbody_deleter(fbody); + Graph* g = fbody->graph; + + // Find or create the key placeholder node. + Node* key_placeholder = nullptr; + for (Node* n : g->nodes()) { + if (IsKeyPlaceholderNode(*n)) { + key_placeholder = n; + break; + } + } + if (!key_placeholder) { + TF_ASSIGN_OR_RETURN(key_placeholder, + AddHostComputeKeyPlaceholder(xla_cluster_name, g)); + } + + // Build the _Arg node, and replace key placeholder node with it. + NodeDefBuilder arg_builder("key_arg", FunctionLibraryDefinition::kArgOp); + arg_builder.Attr("T", DT_STRING); + arg_builder.Attr("index", 0); + NodeDef arg_def; + TF_RETURN_IF_ERROR(arg_builder.Finalize(&arg_def)); + TF_RETURN_IF_ERROR(ReplaceNode(g, key_placeholder, arg_def).status()); + + // Reset "device_ordinal" to placeholder value. + TF_RETURN_IF_ERROR(ResetDeviceOrdinalToPlaceholderValue(g)); + + FunctionDef replace_fdef; + TF_RETURN_IF_ERROR(GraphToFunctionDef(*g, func_name, &replace_fdef)); + TF_RETURN_IF_ERROR(fld->ReplaceFunction(func_name, replace_fdef)); + return Status::OK(); +} + +// Builds host side graph for If node. +Status BuildHostGraphForIfNode(const string& xla_cluster_attr_name, + const string& outside_compilation_attr_name, + const string& xla_cluster_name, + const string& if_node_name, + const string& host_transfer_key, + const string& host_graph_func_name, + FunctionLibraryDefinition* fld, + const string& then_branch_host_func_name, + const string& else_branch_host_func_name) { + Graph host_graph(fld); + string outside_compilation_name = absl::StrCat("oc_if_", if_node_name); + 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: build XlaRecvAtHost node to recv predicate. + NodeDefBuilder recv_pred_builder( + absl::StrCat("recv_oc_if_pred_", if_node_name), "_XlaRecvAtHost"); + recv_pred_builder.Attr("Toutputs", std::vector{DT_BOOL}); + recv_pred_builder.Attr("key", host_transfer_key); + recv_pred_builder.Attr("device_ordinal", device_ordinal_value); + recv_pred_builder.Attr(xla_cluster_attr_name, xla_cluster_name); + recv_pred_builder.Attr(outside_compilation_attr_name, + outside_compilation_name); + recv_pred_builder.Attr(kXlaHasHostTransferAttrName, true); + recv_pred_builder.Input(key_placeholder->name(), 0, DT_STRING); + NodeDef recv_pred_def; + TF_RETURN_IF_ERROR(recv_pred_builder.Finalize(&recv_pred_def)); + Status s; + Node* recv_pred_node = host_graph.AddNode(recv_pred_def, &s); + TF_RETURN_IF_ERROR(s); + host_graph.AddEdge(key_placeholder, 0, recv_pred_node, 0); + + // Step 3: rewrite `{then, else}_branch_host_func_name`, replace key + // placeholder with an _Arg node. + TF_RETURN_IF_ERROR(ReplaceKeyPlaceholderWithArgNode( + xla_cluster_name, then_branch_host_func_name, fld)); + TF_RETURN_IF_ERROR(ReplaceKeyPlaceholderWithArgNode( + xla_cluster_name, else_branch_host_func_name, fld)); + + // Step 4: build If node to choose between `{then, else}_branch_host_graph`. + NodeDefBuilder if_builder(absl::StrCat("oc_if_", if_node_name), "If"); + if_builder.Attr("Tcond", DT_BOOL); + if_builder.Attr("Tin", std::vector{DT_STRING}); + if_builder.Attr("Tout", std::vector{}); + NameAttrList host_then_branch, host_else_branch; + host_then_branch.set_name(then_branch_host_func_name); + (*host_then_branch.mutable_attr())["device_ordinal"] = device_ordinal_value; + host_else_branch.set_name(else_branch_host_func_name); + (*host_else_branch.mutable_attr())["device_ordinal"] = device_ordinal_value; + if_builder.Attr("then_branch", host_then_branch); + if_builder.Attr("else_branch", host_else_branch); + if_builder.Attr(kXlaHasHostTransferAttrName, true); + if_builder.Attr(xla_cluster_attr_name, xla_cluster_name); + if_builder.Attr(outside_compilation_attr_name, outside_compilation_name); + if_builder.Input(recv_pred_node->name(), 0, DT_BOOL); + std::vector if_inputs{ + {key_placeholder->name(), 0, DT_STRING}}; + if_builder.Input(if_inputs); + NodeDef if_def; + TF_RETURN_IF_ERROR(if_builder.Finalize(&if_def)); + Node* if_node = host_graph.AddNode(if_def, &s); + TF_RETURN_IF_ERROR(s); + host_graph.AddEdge(recv_pred_node, 0, if_node, 0); + host_graph.AddEdge(key_placeholder, 0, if_node, 1); + + // 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(); +} + +// Rewrites loop cond to add a node which sends loop cond to host. +Status AddSendLoopPredToLoopCond(FunctionLibraryDefinition* fld, + const NameAttrList& loop_cond_func, + const string& while_node_name, + const string& host_transfer_key) { + // Instantiate the loop cond function. + FunctionBody* fbody = nullptr; + TF_RETURN_IF_ERROR(FunctionDefToBodyHelper( + *fld->Find(loop_cond_func.name()), AttrSlice(&loop_cond_func.attr()), fld, + [&](const string& op, const OpDef** sig) { + return fld->LookUpOpDef(op, sig); + }, + &fbody)); + std::unique_ptr fbody_deleter(fbody); + Graph* g = fbody->graph; + + // Find the _Retval node and the loop cond node. + Node* ret_node = nullptr; + for (Node* n : g->nodes()) { + if (n->type_string() == "_Retval") { + if (ret_node) { + return errors::Internal("Multiple return node for loop cond function ", + loop_cond_func.name(), ": ", + ret_node->DebugString(), " and ", + n->DebugString()); + } else { + ret_node = n; + } + } + } + if (!ret_node) { + return errors::Internal("No _Retval node for loop cond function ", + loop_cond_func.name()); + } + Node* loop_cond; + TF_RETURN_IF_ERROR(ret_node->input_node(0, &loop_cond)); + + // Build the XlaSendToHost node. + NodeDefBuilder send_loop_cond_builder( + absl::StrCat("send_oc_while_cond_", while_node_name), "XlaSendToHost"); + send_loop_cond_builder.Attr("Tinput", DT_BOOL); + send_loop_cond_builder.Attr("key", + absl::StrCat(host_transfer_key, "_dtoh_0")); + send_loop_cond_builder.Attr(kXlaTokenInputNodesAttrName, + std::vector{kXlaTokenArgNodeName}); + send_loop_cond_builder.Input(loop_cond->name(), 0, DT_BOOL); + NodeDef send_loop_cond_def; + TF_RETURN_IF_ERROR(send_loop_cond_builder.Finalize(&send_loop_cond_def)); + Status s; + Node* send_loop_cond_node = g->AddNode(send_loop_cond_def, &s); + TF_RETURN_IF_ERROR(s); + g->AddEdge(loop_cond, 0, send_loop_cond_node, 0); + + // Replace original function. + FunctionDef replace_fdef; + TF_RETURN_IF_ERROR( + GraphToFunctionDef(*g, loop_cond_func.name(), &replace_fdef)); + TF_RETURN_IF_ERROR(fld->ReplaceFunction(loop_cond_func.name(), replace_fdef)); + + return Status::OK(); +} + +// Rewrites while loop cond function for host. +Status RewriteHostWhileLoopCond( + const string& cond_host_func_name, const string& while_node_name, + const string& host_transfer_key, const string& xla_cluster_attr_name, + const string& xla_cluster_name, const string& outside_compilation_attr_name, + const string& outside_compilation_name, FunctionLibraryDefinition* fld) { + // Replace key placeholder node with _Arg node. + TF_RETURN_IF_ERROR(ReplaceKeyPlaceholderWithArgNode( + xla_cluster_name, cond_host_func_name, fld)); + + // Instantiate cond function. + AttrValue device_ordinal_temp_value; + device_ordinal_temp_value.set_i(0); + protobuf::Map attrs; + attrs["device_ordinal"] = device_ordinal_temp_value; + FunctionBody* cond_fbody = nullptr; + TF_RETURN_IF_ERROR(FunctionDefToBodyHelper( + *fld->Find(cond_host_func_name), AttrSlice(&attrs), fld, + [&](const string& op, const OpDef** sig) { + return fld->LookUpOpDef(op, sig); + }, + &cond_fbody)); + std::unique_ptr cond_fbody_deleter(cond_fbody); + Graph* cond_graph = cond_fbody->graph; + Node* key_arg = nullptr; + for (Node* n : cond_graph->nodes()) { + if (n->type_string() == "_Arg") { + key_arg = n; + } + } + if (!key_arg) { + return errors::Internal( + "No _Arg node found for host compute key in function ", + cond_host_func_name); + } + + // Add an XlaRecvAtHost node to use as cond function return value. + // We don't need to set kXlaHasHostTransferAttrName for this node, because + // it's already added for the "While" node on the host. + NodeDefBuilder recv_pred_builder( + absl::StrCat("recv_oc_while_cond_", while_node_name), "_XlaRecvAtHost"); + recv_pred_builder.Attr("Toutputs", std::vector{DT_BOOL}); + recv_pred_builder.Attr("key", host_transfer_key); + AttrValue device_ordinal_value; + device_ordinal_value.set_placeholder("device_ordinal"); + recv_pred_builder.Attr("device_ordinal", device_ordinal_value); + recv_pred_builder.Attr(xla_cluster_attr_name, xla_cluster_name); + recv_pred_builder.Attr(outside_compilation_attr_name, + outside_compilation_name); + recv_pred_builder.Input(key_arg->name(), 0, DT_STRING); + NodeDef recv_pred_def; + TF_RETURN_IF_ERROR(recv_pred_builder.Finalize(&recv_pred_def)); + Status s; + Node* recv_pred_node = cond_graph->AddNode(recv_pred_def, &s); + TF_RETURN_IF_ERROR(s); + cond_graph->AddEdge(key_arg, 0, recv_pred_node, 0); + NodeDefBuilder ret_builder( + absl::StrCat("recv_oc_while_cond_ret_", while_node_name), "_Retval"); + ret_builder.Attr("T", DT_BOOL); + ret_builder.Attr("index", 0); + ret_builder.Input(recv_pred_node->name(), 0, DT_BOOL); + NodeDef ret_def; + TF_RETURN_IF_ERROR(ret_builder.Finalize(&ret_def)); + Node* ret_node = cond_graph->AddNode(ret_def, &s); + TF_RETURN_IF_ERROR(s); + cond_graph->AddEdge(recv_pred_node, 0, ret_node, 0); + + // Reset device_ordinal to placeholder value. + TF_RETURN_IF_ERROR(ResetDeviceOrdinalToPlaceholderValue(cond_graph)); + + // Replace original function. + FunctionDef cond_replace_fdef; + TF_RETURN_IF_ERROR( + GraphToFunctionDef(*cond_graph, cond_host_func_name, &cond_replace_fdef)); + TF_RETURN_IF_ERROR( + fld->ReplaceFunction(cond_host_func_name, cond_replace_fdef)); + + return Status::OK(); +} + +// Rewrites while loop body function for host. +Status RewriteHostWhileLoopBody( + const string& body_host_func_name, const string& while_node_name, + const string& host_transfer_key, const string& xla_cluster_attr_name, + const string& xla_cluster_name, const string& outside_compilation_attr_name, + const string& outside_compilation_name, FunctionLibraryDefinition* fld) { + // Replace key placeholder node with _Arg node. + TF_RETURN_IF_ERROR(ReplaceKeyPlaceholderWithArgNode( + xla_cluster_name, body_host_func_name, fld)); + + // Instantiate body function. + AttrValue device_ordinal_temp_value; + device_ordinal_temp_value.set_i(0); + protobuf::Map attrs; + attrs["device_ordinal"] = device_ordinal_temp_value; + FunctionBody* body_fbody = nullptr; + TF_RETURN_IF_ERROR(FunctionDefToBodyHelper( + *fld->Find(body_host_func_name), AttrSlice(&attrs), fld, + [&](const string& op, const OpDef** sig) { + return fld->LookUpOpDef(op, sig); + }, + &body_fbody)); + std::unique_ptr body_fbody_deleter(body_fbody); + Graph* body_graph = body_fbody->graph; + Node* key_arg = nullptr; + for (Node* n : body_graph->nodes()) { + if (n->type_string() == "_Arg") { + key_arg = n; + } + } + if (!key_arg) { + return errors::Internal( + "No _Arg node found for host compute key in function ", + body_host_func_name); + } + + // Add a _Retval node to loop body. + NodeDefBuilder ret_builder( + absl::StrCat("recv_oc_while_body_ret_", while_node_name), "_Retval"); + ret_builder.Attr("T", DT_STRING); + ret_builder.Attr("index", 0); + ret_builder.Input(key_arg->name(), 0, DT_STRING); + NodeDef ret_def; + TF_RETURN_IF_ERROR(ret_builder.Finalize(&ret_def)); + Status s; + Node* ret_node = body_graph->AddNode(ret_def, &s); + TF_RETURN_IF_ERROR(s); + body_graph->AddEdge(key_arg, 0, ret_node, 0); + + // Reset device_ordinal to placeholder value. + TF_RETURN_IF_ERROR(ResetDeviceOrdinalToPlaceholderValue(body_graph)); + + // Replace original function. + FunctionDef body_replace_fdef; + TF_RETURN_IF_ERROR( + GraphToFunctionDef(*body_graph, body_host_func_name, &body_replace_fdef)); + TF_RETURN_IF_ERROR( + fld->ReplaceFunction(body_host_func_name, body_replace_fdef)); + + return Status::OK(); +} + +// Builds host side graph for while node. +Status BuildHostGraphForWhileNode( + const string& xla_cluster_attr_name, + const string& outside_compilation_attr_name, const string& xla_cluster_name, + const string& while_node_name, const string& host_transfer_key, + const string& host_graph_func_name, FunctionLibraryDefinition* fld, + const string& cond_host_func_name, const string& body_host_func_name) { + Graph host_graph(fld); + string outside_compilation_name = absl::StrCat("oc_while_", while_node_name); + + // Step 1: add key placeholder node. + TF_ASSIGN_OR_RETURN( + Node * key_placeholder, + AddHostComputeKeyPlaceholder(xla_cluster_name, &host_graph)); + + // Step 2: rewrite cond function. + TF_RETURN_IF_ERROR(RewriteHostWhileLoopCond( + cond_host_func_name, while_node_name, host_transfer_key, + xla_cluster_attr_name, xla_cluster_name, outside_compilation_attr_name, + outside_compilation_name, fld)); + + // Step 3: rewrite body function. + TF_RETURN_IF_ERROR(RewriteHostWhileLoopBody( + body_host_func_name, while_node_name, host_transfer_key, + xla_cluster_attr_name, xla_cluster_name, outside_compilation_attr_name, + outside_compilation_name, fld)); + + // Step 4: build While node. + NodeDefBuilder while_builder(absl::StrCat("oc_while_", while_node_name), + "While"); + while_builder.Attr("T", std::vector{DT_STRING}); + NameAttrList func; + AttrValue device_ordinal_value; + device_ordinal_value.set_placeholder("device_ordinal"); + (*func.mutable_attr())["device_ordinal"] = device_ordinal_value; + func.set_name(cond_host_func_name); + while_builder.Attr("cond", func); + func.set_name(body_host_func_name); + while_builder.Attr("body", func); + while_builder.Attr(kXlaHasHostTransferAttrName, true); + while_builder.Attr(xla_cluster_attr_name, xla_cluster_name); + while_builder.Attr(outside_compilation_attr_name, outside_compilation_name); + std::vector while_inputs{ + {key_placeholder->name(), 0, DT_STRING}}; + while_builder.Input(while_inputs); + NodeDef while_def; + TF_RETURN_IF_ERROR(while_builder.Finalize(&while_def)); + Status s; + Node* while_node = host_graph.AddNode(while_def, &s); + TF_RETURN_IF_ERROR(s); + host_graph.AddEdge(key_placeholder, 0, while_node, 0); + + // Convert `host_graph` to function. + 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(); +} + +// 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, FunctionLibraryRuntime* flr, + FunctionLibraryDefinition* fld, std::vector* host_graphs, + std::vector* shape_inference_graphs, + bool* has_outside_compilation) { + 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) { + // Instantiate "then_branch" and "else_branch". + NameAttrList then_branch, else_branch; + TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), "then_branch", &then_branch)); + TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), "else_branch", &else_branch)); + + // Extract outside compilation for then_branch and else_branch. + bool then_branch_has_outside_compilation = false; + bool else_branch_has_outside_compilation = false; + string then_branch_host_func_name = + absl::StrCat("oc_then_branch_host_if_", n->name()), + else_branch_host_func_name = + absl::StrCat("oc_else_branch_host_if_", n->name()); + string then_branch_xla_func_name = absl::StrCat(then_branch.name(), "_oc"), + else_branch_xla_func_name = absl::StrCat(else_branch.name(), "_oc"); + 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, 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, flr, fld, shape_inference_graphs, + &else_branch_has_outside_compilation)); + + // If then/else branch do not have outside compilation, nothing to do. + if (!then_branch_has_outside_compilation && + !else_branch_has_outside_compilation) { + continue; + } + + *has_outside_compilation = true; + + // Change If node to call the new functions. + then_branch.set_name(then_branch_xla_func_name); + n->ClearAttr("then_branch"); + n->AddAttr("then_branch", then_branch); + else_branch.set_name(else_branch_xla_func_name); + n->ClearAttr("else_branch"); + n->AddAttr("else_branch", else_branch); + + string host_transfer_key = absl::StrCat("oc_if_pred_", n->name()); + + // XLA computation: add a SendToHost node to send cond predicate. + Node* pred_node; + TF_RETURN_IF_ERROR(n->input_node(0, &pred_node)); + TF_ASSIGN_OR_RETURN( + Node * send_pred_node, + BuildSendIfPredNode(absl::StrCat("send_oc_if_pred_", n->name()), + host_transfer_key, pred_node, g)); + n->AddAttr(kXlaTokenInputNodesAttrName, + std::vector{send_pred_node->name()}); + + // Add a control edge from `send_pred_node` to If node, so XlaCompiler will + // visit If node after `send_pred_node`, thus the token output for + // `send_pred_node` has been generated. + g->AddControlEdge(send_pred_node, n); + + // Build host side graph for the "If" node. + string oc_host_graph_name = absl::StrCat("oc_if_host_graph_", n->name()); + TF_RETURN_IF_ERROR(BuildHostGraphForIfNode( + xla_cluster_attr_name, outside_compilation_attr_name, xla_cluster_name, + n->name(), host_transfer_key, oc_host_graph_name, fld, + then_branch_host_func_name, else_branch_host_func_name)); + host_graphs->push_back(oc_host_graph_name); + } + + for (Node* n : while_nodes) { + // Instantiate "cond" and "body". + NameAttrList cond, body; + TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), "cond", &cond)); + TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), "body", &body)); + + // Extract outside compilation for cond and body. + bool cond_has_outside_compilation = false; + bool body_has_outside_compilation = false; + string cond_host_func_name = absl::StrCat("oc_cond_host_while_", n->name()), + body_host_func_name = absl::StrCat("oc_body_host_while_", n->name()); + string cond_xla_func_name = absl::StrCat(cond.name(), "_oc"), + 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, 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, 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) { + continue; + } + + *has_outside_compilation = true; + + // Change While node to call the new functions. + cond.set_name(cond_xla_func_name); + n->ClearAttr("cond"); + n->AddAttr("cond", cond); + body.set_name(body_xla_func_name); + n->ClearAttr("body"); + n->AddAttr("body", body); + + string host_transfer_key = absl::StrCat("oc_while_pred_", n->name()); + + // XLA computation: rewrite cond function to add a SendToHost node to send + // loop predicate. + TF_RETURN_IF_ERROR( + AddSendLoopPredToLoopCond(fld, cond, n->name(), host_transfer_key)); + n->AddAttr(kXlaTokenInputNodesAttrName, + std::vector{kXlaTokenArgNodeName}); + + // Build host side graph for the "While" node. + string oc_host_graph_name = absl::StrCat("oc_while_host_graph_", n->name()); + TF_RETURN_IF_ERROR(BuildHostGraphForWhileNode( + xla_cluster_attr_name, outside_compilation_attr_name, xla_cluster_name, + n->name(), host_transfer_key, oc_host_graph_name, fld, + cond_host_func_name, body_host_func_name)); + host_graphs->push_back(oc_host_graph_name); + } + + return Status::OK(); +} + } // namespace Status RewriteOutsideCompilationSubgraphFn::operator()( @@ -298,8 +1513,7 @@ Status RewriteOutsideCompilationSubgraphFn::operator()( // Step 4: add XLA cluster and outside compilation attr. for (Node* n : (*graph)->nodes()) { - if (n->type_string() == "Placeholder" && - absl::EndsWith(n->name(), "_key_placeholder")) { + if (IsKeyPlaceholderNode(*n)) { continue; } @@ -312,6 +1526,9 @@ Status RewriteOutsideCompilationSubgraphFn::operator()( // shape inference graph and set `shape_inference_graph` for the call node. absl::optional> shapes = GetInferredInputShapes(send_from_host_dtypes.size(), send_from_host_node); + for (Node* n : (*graph)->nodes()) { + n->ClearAttr(kXlaInferredShapesAttrName); + } // Step 5: add control edges for originally XLA <-> outside compilation // control edges. @@ -346,12 +1563,15 @@ Status RewriteOutsideCompilationSubgraphFn::operator()( // it with HostCompute node later. AddNodeAttr("_outside_compilation_subgraph", old_name, node_def); if (shapes) { - AddNodeAttr("shape_inference_graph", "", node_def); + NameAttrList shape_inference_graph; + AddNodeAttr("shape_inference_graph", shape_inference_graph, node_def); AddNodeAttr("shapes", *shapes, node_def); } else { string shape_inference_func_name = absl::StrCat("_outside_compilation_shape_inference_", new_name); - AddNodeAttr("shape_inference_graph", shape_inference_func_name, node_def); + NameAttrList shape_inference_graph; + shape_inference_graph.set_name(shape_inference_func_name); + AddNodeAttr("shape_inference_graph", shape_inference_graph, node_def); AddNodeAttr("shapes", std::vector{}, node_def); } AddNodeAttr("ancestors", std::vector{}, node_def); @@ -362,4 +1582,174 @@ Status RewriteOutsideCompilationSubgraphFn::operator()( return Status::OK(); } +Status ExtractOutsideCompilationForFunction( + 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, 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(); + 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 (Node* n : fbody->graph->nodes()) { + if (HasNodeAttr(n->def(), outside_compilation_attr_name)) { + *has_outside_compilation = true; + break; + } + } + // We cannot early return here, because we might have outside compilation in + // If/While function body. + + // Preprocess edges between different outside compilations. They will be + // restored in `ConstructHostGraph()`. + TF_RETURN_IF_ERROR(PreprocessEdgesBetweenOutsideCompilations( + fbody->graph, outside_compilation_attr_name)); + if (VLOG_IS_ON(4)) { + dump_graph::DumpGraphToFile( + absl::StrCat("extract_outside_compilation_for_func_before_", func_name), + *fbody->graph, fld); + } + + // Encapsulate outside_compilation cluster into function call node. + std::unique_ptr graph_out; + RewriteOutsideCompilationSubgraphFn rewrite_fn( + xla_cluster_attr_name, outside_compilation_attr_name, xla_cluster_name); + TF_RETURN_IF_ERROR(EncapsulateSubgraphsInFunctions( + outside_compilation_attr_name, "", *fbody->graph, rewrite_fn, + /*reuse_existing_functions=*/true, &graph_out, fld)); + + // Replace outside_compilation function nodes with HostCompute ops. + std::vector outside_compilation_nodes; + std::vector outside_compilation_host_graphs; + for (Node* n : graph_out->nodes()) { + if (HasNodeAttr(n->def(), "_outside_compilation_subgraph")) { + outside_compilation_nodes.push_back(n); + outside_compilation_host_graphs.push_back(n->name()); + + // If we could not infer shapes for XlaSendFromHost inputs statically, we + // will set the "shape_inference_graph" attribute. In that case, copy + // outside compilation subgraph as shape inference graph in `fld`. + NameAttrList shape_inference_graph; + TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), "shape_inference_graph", + &shape_inference_graph)); + if (!shape_inference_graph.name().empty()) { + shape_inference_graphs->push_back(shape_inference_graph.name()); + + const FunctionDef* xla_fdef = fld->Find(n->name()); + if (!xla_fdef) { + return errors::Internal("Cannot find XLA function ", n->name()); + } + FunctionDef shape_inference_fdef = *xla_fdef; + shape_inference_fdef.mutable_signature()->set_name( + shape_inference_graph.name()); + if (fld->Find(shape_inference_graph.name())) { + TF_RETURN_IF_ERROR(fld->ReplaceFunction(shape_inference_graph.name(), + shape_inference_fdef)); + } else { + TF_RETURN_IF_ERROR(fld->AddFunctionDef(shape_inference_fdef)); + } + } + } + } + for (Node* n : outside_compilation_nodes) { + TF_RETURN_IF_ERROR(ValidateOutsideCompilationCallNode(n)); + TF_RETURN_IF_ERROR(ReplaceOrRemoveOutsideCompilationCallNode( + graph_out.get(), n, host_compute_core)); + } + + // 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, flr, fld, + &outside_compilation_host_graphs, shape_inference_graphs, + has_outside_compilation)); + + // Construct host graph. + TF_RETURN_IF_ERROR(ConstructHostGraph( + xla_cluster_name, outside_compilation_attr_name, + outside_compilation_host_graphs, fld, host_graph_func_name)); + + // Remove the outside compilation graphs from function library. + for (const string& func : outside_compilation_host_graphs) { + TF_RETURN_IF_ERROR(fld->RemoveFunction(func)); + } + + // Replace original function. + 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 ret_status; +} + +Status ExtractOutsideCompilation( + const string& xla_cluster_attr_name, + const string& outside_compilation_attr_name, + const std::unordered_map& clusters, Graph* g, + FunctionLibraryRuntime* flr, FunctionLibraryDefinition* fld) { + if (VLOG_IS_ON(4)) { + dump_graph::DumpGraphToFile("extract_outside_compilation_before", *g, fld); + } + + std::vector shape_inference_graphs; + for (auto& iter : clusters) { + string xla_cluster_name = iter.first; + Node* n = iter.second.node; + auto const& func_name_attrs = iter.second.func_name_attrs; + auto const& host_compute_core = iter.second.host_compute_core; + + bool has_outside_compilation; + string host_graph_func_name = absl::StrCat("oc_host_graph_", n->name()); + 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, flr, fld, &shape_inference_graphs, + &has_outside_compilation)); + TF_RETURN_IF_ERROR( + ExpandHostGraphIntoMainGraph(g, fld, host_graph_func_name, n)); + TF_RETURN_IF_ERROR(fld->RemoveFunction(host_graph_func_name)); + } + + for (auto shape_inference_graph_name : shape_inference_graphs) { + TF_RETURN_IF_ERROR( + RewriteShapeInferenceGraph(shape_inference_graph_name, g, fld)); + } + + if (VLOG_IS_ON(4)) { + dump_graph::DumpGraphToFile("extract_outside_compilation_after", *g, fld); + } + return Status::OK(); +} + } // namespace tensorflow diff --git a/tensorflow/compiler/jit/extract_outside_compilation_pass.h b/tensorflow/compiler/jit/extract_outside_compilation_pass.h index 4aa34d76c686996fb10dd8940f7813af9e526292..d64cc2a103ed040cbf413ac736f97f84459e869b 100644 --- a/tensorflow/compiler/jit/extract_outside_compilation_pass.h +++ b/tensorflow/compiler/jit/extract_outside_compilation_pass.h @@ -17,6 +17,7 @@ limitations under the License. #define TENSORFLOW_COMPILER_JIT_EXTRACT_OUTSIDE_COMPILATION_PASS_H_ #include "absl/types/optional.h" +#include "tensorflow/compiler/jit/encapsulate_util.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/core/graph/graph.h" @@ -61,6 +62,47 @@ class RewriteOutsideCompilationSubgraphFn { string xla_cluster_name_; }; +// For an XLA computation function, replace all outside compilations with +// XlaHostCompute nodes. Each outside compilation subgraph will be rewritten by +// `RewriteOutsideCompilationSubgraphFn`, and they will be merged into one +// single host side graph (`host_graph`). +// +// xla_cluster_attr_name and outside_compilation_attr_name: attr name for XLA +// computation and outside compilation. Required for +// `RewriteOutsideCompilationSubgraphFn`. +// xla_cluster_name: XLA cluster name for this XLA computation. We need it +// because XLA cluster name might be different from `func_name`. +// func_name_attrs: they will be used to instantiate the XLA computation func. +// new_func_name: new function name for rewritten XLA computation func. +// host_compute_core: mapping from outside compilation cluster name to XLA +// device assignment. +// fld: FunctionLibraryDefinition object. +// host_graph: Graph object to store host side graph for all outside +// compilations within this XLA computation func. If there is no outside +// compilation, it will be empty. +// shape_inference_graphs: a list of outside compilation shape inference +// function names. These functions need to be rewritten later. +// has_outside_compilation: a bool indicating whether this function has any +// outside compilation nodes. +Status ExtractOutsideCompilationForFunction( + 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, FunctionLibraryRuntime* flr, + FunctionLibraryDefinition* fld, std::vector* shape_inference_graphs, + bool* has_outside_compilation); + +// Rewrites XLA computation in `clusters` to replace outside compilation nodes +// with XlaHostCompute, and moves those outside compilations into `g`. If shapes +// of outside compilation outputs cannot be determined now, we will store shape +// inference graph into `fld`. +Status ExtractOutsideCompilation( + const string& xla_cluster_attr_name, + const string& outside_compilation_attr_name, + const std::unordered_map& clusters, Graph* g, + FunctionLibraryRuntime* flr, FunctionLibraryDefinition* fld); + } // namespace tensorflow #endif // TENSORFLOW_COMPILER_JIT_EXTRACT_OUTSIDE_COMPILATION_PASS_H_ diff --git a/tensorflow/compiler/jit/extract_outside_compilation_pass_test.cc b/tensorflow/compiler/jit/extract_outside_compilation_pass_test.cc index 64913f5aab68e07bf91e877d05f56eed147cda26..7c3a24feff81b21a5d2347d21fb80988bc3e6065 100644 --- a/tensorflow/compiler/jit/extract_outside_compilation_pass_test.cc +++ b/tensorflow/compiler/jit/extract_outside_compilation_pass_test.cc @@ -15,16 +15,25 @@ limitations under the License. #include "tensorflow/compiler/jit/extract_outside_compilation_pass.h" +#include "absl/strings/match.h" #include "tensorflow/cc/framework/scope.h" #include "tensorflow/cc/ops/array_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/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" +#include "tensorflow/core/framework/graph_to_functiondef.h" #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/platform/test.h" +#include "tensorflow/core/public/session_options.h" +#include "tensorflow/core/public/version.h" namespace tensorflow { @@ -105,10 +114,10 @@ TEST(RewriteOutsideCompilationSubgraphFnTest, Basic) { } EXPECT_TRUE(has_control_edge_to_send_from_host); // Verify step 7: necessary attrs added to call_node_def. - string shape_inference_graph; + NameAttrList shape_inference_graph; TF_CHECK_OK(GetNodeAttr(AttrSlice(&call_node_def.attr()), "shape_inference_graph", &shape_inference_graph)); - EXPECT_EQ(shape_inference_graph, + EXPECT_EQ(shape_inference_graph.name(), "_outside_compilation_shape_inference_cluster_0"); } @@ -216,4 +225,757 @@ TEST(RewriteOutsideCompilationSubgraphFnTest, ShapesInferred) { EXPECT_EQ(shapes[0].dim_size(), 1); } +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") + // "identity1" = "identity0" (outside compilation cluster "1") + // "identity2" = "identity1" + FunctionDefLibrary fdl; + { + tensorflow::Scope s = tensorflow::Scope::NewRootScope(); + Output const0 = ops::Const(s.WithOpName("const0"), 1, {2}); + Output identity0 = ops::Identity(s.WithOpName("identity0"), const0); + Output identity1 = ops::Identity(s.WithOpName("identity1"), identity0); + Output identity2 = ops::Identity(s.WithOpName("identity2"), identity1); + std::unique_ptr g(new Graph(OpRegistry::Global())); + TF_CHECK_OK(s.ToGraph(g.get())); + auto node_name_image = g->BuildNodeNameIndex(); + node_name_image["identity0"]->AddAttr("_oc", "0"); + node_name_image["identity1"]->AddAttr("_oc", "1"); + PartialTensorShape shape({2}); + node_name_image["identity1"]->AddAttr( + kXlaInferredShapesAttrName, std::vector{shape}); + + FunctionDef *xla_fdef = fdl.add_function(); + TF_CHECK_OK(GraphToFunctionDef(*g, "cluster", xla_fdef)); + } + FunctionLibraryDefinition fld(OpRegistry::Global(), fdl); + + protobuf::Map attrs; + std::map host_compute_core = {{"0", 1}, {"1", 0}}; + 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)); + + // Get rewritten XLA computation function. + 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); + auto node_name_index = xla_fbody->graph->BuildNodeNameIndex(); + + // Check XlaHostCompute nodes. + Node *host_compute_0 = node_name_index["outside_compilation_0_host_compute"]; + EXPECT_NE(host_compute_0, nullptr); + Node *host_compute_1 = node_name_index["outside_compilation_1_host_compute"]; + EXPECT_NE(host_compute_1, nullptr); + // Check XlaHostCompute nodes' "tpu_core" attr. + int tpu_core; + TF_CHECK_OK(GetNodeAttr(host_compute_0->attrs(), "tpu_core", &tpu_core)); + EXPECT_EQ(tpu_core, 1); + TF_CHECK_OK(GetNodeAttr(host_compute_1->attrs(), "tpu_core", &tpu_core)); + EXPECT_EQ(tpu_core, 0); + // Check XlaHostCompute nodes' "shapes" attr. "0" should not have shapes, and + // "1" should have shapes. + std::vector shapes; + TF_CHECK_OK(GetNodeAttr(host_compute_0->attrs(), "shapes", &shapes)); + EXPECT_EQ(shapes.size(), 0); + TF_CHECK_OK(GetNodeAttr(host_compute_1->attrs(), "shapes", &shapes)); + EXPECT_EQ(shapes.size(), 1); + EXPECT_EQ(shapes[0].dim_size(), 1); + // Check XlaHostCompute nodes' "shape_inference_graph" attr. Both should have + // empty values. + NameAttrList shape_inference_graph; + TF_CHECK_OK(GetNodeAttr(host_compute_0->attrs(), "shape_inference_graph", + &shape_inference_graph)); + EXPECT_EQ(shape_inference_graph.name(), ""); + TF_CHECK_OK(GetNodeAttr(host_compute_1->attrs(), "shape_inference_graph", + &shape_inference_graph)); + EXPECT_EQ(shape_inference_graph.name(), ""); + + // Check `shape_inference_graphs`. + EXPECT_EQ(shape_inference_graphs.size(), 0); + + // Check host graph: verify we have key placeholder and sequencer. + 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; + Node *key_placeholder = nullptr, *sequencer = nullptr; + for (Node *n : host_graph->nodes()) { + if (n->type_string() == "Placeholder" && + absl::EndsWith(n->name(), "_key_placeholder")) { + EXPECT_EQ(key_placeholder, nullptr); + key_placeholder = n; + } else if (HasNodeAttr(n->def(), "_xla_host_transfer_sequencer")) { + EXPECT_EQ(sequencer, nullptr); + sequencer = n; + } + } + EXPECT_NE(key_placeholder, nullptr); + EXPECT_NE(sequencer, nullptr); + // Check SendFromHost and RecvAtHost has key placeholder as input, and have + // control edge to sequencer. + int num_send_from_host = 0, num_recv_at_host = 0; + std::vector send_recv_nodes; + for (Node *n : host_graph->nodes()) { + if (n->type_string() == "_XlaSendFromHost") { + num_send_from_host++; + send_recv_nodes.push_back(n); + } else if (n->type_string() == "_XlaRecvAtHost") { + num_recv_at_host++; + send_recv_nodes.push_back(n); + } + } + EXPECT_EQ(num_send_from_host, 1); + EXPECT_EQ(num_recv_at_host, 1); + for (Node *n : send_recv_nodes) { + Node *input_node; + TF_CHECK_OK(n->input_node(n->num_inputs() - 1, &input_node)); + EXPECT_EQ(input_node, key_placeholder); + + bool has_control_edge_to_sequencer = false; + for (const Edge *e : n->out_edges()) { + if (e->IsControlEdge() && e->dst() == sequencer) { + has_control_edge_to_sequencer = true; + break; + } + } + EXPECT_TRUE(has_control_edge_to_sequencer); + } +} + +TEST_F(ExtractOutsideCompilationForFunctionTest, NoHostGraph) { + // Build the XLA computation func. + // "const0" + FunctionDefLibrary fdl; + { + tensorflow::Scope s = tensorflow::Scope::NewRootScope(); + Output const0 = ops::Const(s.WithOpName("const0"), 1, {2}); + std::unique_ptr g(new Graph(OpRegistry::Global())); + TF_CHECK_OK(s.ToGraph(g.get())); + + FunctionDef *xla_fdef = fdl.add_function(); + TF_CHECK_OK(GraphToFunctionDef(*g, "cluster", xla_fdef)); + } + FunctionLibraryDefinition fld(OpRegistry::Global(), fdl); + + protobuf::Map attrs; + std::map host_compute_core = {{"0", 1}, {"1", 0}}; + 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 is empty. + 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; + EXPECT_EQ(host_graph->num_nodes(), 2); +} + +TEST_F(ExtractOutsideCompilationForFunctionTest, XlaHostComputeRemoved) { + // Build the XLA computation func. + // "const0" + // "const1" (outside compilation cluster "0") + FunctionDefLibrary fdl; + { + tensorflow::Scope s = tensorflow::Scope::NewRootScope(); + Output const0 = ops::Const(s.WithOpName("const0"), 1, {2}); + Output const1 = ops::Const(s.WithOpName("const1"), 1, {2}); + std::unique_ptr g(new Graph(OpRegistry::Global())); + TF_CHECK_OK(s.ToGraph(g.get())); + auto node_name_image = g->BuildNodeNameIndex(); + node_name_image["const1"]->AddAttr("_oc", "0"); + + FunctionDef *xla_fdef = fdl.add_function(); + TF_CHECK_OK(GraphToFunctionDef(*g, "cluster", xla_fdef)); + } + FunctionLibraryDefinition fld(OpRegistry::Global(), fdl); + + protobuf::Map attrs; + std::map host_compute_core = {{"0", 1}, {"1", 0}}; + 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 rewritten XLA graph: verify that we have no XlaHostCompute. + 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); + for (Node *n : xla_fbody->graph->nodes()) { + EXPECT_NE(n->type_string(), "XlaHostCompute"); + } + + // Check host graph: verify we have no placeholder, but we have "const1". + 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; + int num_key_placeholders = 0; + for (Node *n : host_graph->nodes()) { + if (n->type_string() == "Placeholder" && + absl::EndsWith(n->name(), "_key_placeholder")) { + num_key_placeholders++; + } + } + EXPECT_EQ(num_key_placeholders, 0); + auto node_name_index = host_graph->BuildNodeNameIndex(); + EXPECT_NE(node_name_index.find("const1"), node_name_index.end()); +} + +REGISTER_OP("XlaSendToHost") + .Input("input: Tinput") + .Attr("Tinput: type") + .Attr("key: string") + .SetIsStateful(); + +REGISTER_OP("XlaRecvFromHost") + .Output("output: Toutput") + .Attr("Toutput: type") + .Attr("shape: shape") + .Attr("key: string") + .SetIsStateful(); + +TEST_F(ExtractOutsideCompilationForFunctionTest, OutsideCompilationInIf) { + // Build the XLA computation func. + // "const0" (bool) + // "const1" (int32) + // "if0" (pred = "const0", input = "const1", then_branch = "true_fn", + // else_branch = "false_fn") + 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_true_fn"), 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_true_fn"]->AddAttr("_oc", "0"); + PartialTensorShape shape({2}); + node_name_image["identity_true_fn"]->AddAttr( + kXlaInferredShapesAttrName, std::vector{shape}); + + FunctionDef *true_fn_fdef = fdl.add_function(); + TF_CHECK_OK(GraphToFunctionDef(*g, "true_fn", true_fn_fdef)); + } + { + tensorflow::Scope s = tensorflow::Scope::NewRootScope(); + Output arg = ops::_Arg(s.WithOpName("arg"), DT_INT32, 0); + Output identity = ops::Identity(s.WithOpName("identity_false_fn"), 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_false_fn"]->AddAttr("_oc", "0"); + PartialTensorShape shape({2}); + node_name_image["identity_false_fn"]->AddAttr( + kXlaInferredShapesAttrName, std::vector{shape}); + + FunctionDef *false_fn_fdef = fdl.add_function(); + TF_CHECK_OK(GraphToFunctionDef(*g, "false_fn", false_fn_fdef)); + } + { + tensorflow::Scope s = tensorflow::Scope::NewRootScope(); + Output cond = ops::Const(s.WithOpName("const0"), true, {2}); + Output input = ops::Const(s.WithOpName("const1"), 1, {2}); + NameAttrList true_fn; + true_fn.set_name("true_fn"); + NameAttrList false_fn; + false_fn.set_name("false_fn"); + auto if_op = ops::If(s.WithOpName("if"), cond, + std::initializer_list{cond, input}, {DT_INT32}, + true_fn, false_fn); + ops::_Retval retval(s.WithOpName("retval"), if_op.output[0], 0); + std::unique_ptr g(new Graph(OpRegistry::Global())); + TF_CHECK_OK(s.ToGraph(g.get())); + + FunctionDef *xla_fdef = fdl.add_function(); + TF_CHECK_OK(GraphToFunctionDef(*g, "cluster", xla_fdef)); + } + FunctionLibraryDefinition fld(OpRegistry::Global(), fdl); + + 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 XlaRecvAtHost to receive "If" predicate. + Node *recv_if_pred_node = node_name_index["recv_oc_if_pred_if"]; + EXPECT_NE(recv_if_pred_node, nullptr); + + // Verify we have an "If" to choose outside compilation between then_branch + // and else_branch, and it has `recv_if_pred_node` as cond input. + Node *if_oc_node = node_name_index["oc_if_if"]; + EXPECT_NE(if_oc_node, nullptr); + Node *if_oc_node_cond_input; + TF_CHECK_OK(if_oc_node->input_node(0, &if_oc_node_cond_input)); + EXPECT_EQ(if_oc_node_cond_input, recv_if_pred_node); + + // Check that then_branch outside compilation has node "identity_true_fn". + const FunctionDef *true_def = fld.Find("oc_then_branch_host_if_if"); + EXPECT_NE(true_def, nullptr); + bool has_identity_true_fn_node = false; + for (const auto &node_def : true_def->node_def()) { + if (node_def.name() == "identity_true_fn") { + has_identity_true_fn_node = true; + break; + } + } + EXPECT_TRUE(has_identity_true_fn_node); + + // Check that else_branch outside compilation has node "identity_false_fn". + const FunctionDef *false_def = fld.Find("oc_else_branch_host_if_if"); + EXPECT_NE(false_def, nullptr); + bool has_identity_false_fn_node = false; + for (const auto &node_def : false_def->node_def()) { + if (node_def.name() == "identity_false_fn") { + has_identity_false_fn_node = true; + break; + } + } + EXPECT_TRUE(has_identity_false_fn_node); + } + + // 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 XlaSendToHost to send cond predicate to host, and + // there is a control edge to If node. + Node *send_if_pred_node = node_name_index["send_oc_if_pred_if"]; + EXPECT_NE(send_if_pred_node, nullptr); + bool has_control_edge_to_if = false; + for (const Edge *e : send_if_pred_node->out_edges()) { + if (e->IsControlEdge() && e->dst()->name() == "if") { + has_control_edge_to_if = true; + break; + } + } + EXPECT_TRUE(has_control_edge_to_if); + + // Check that the "If" node now has `send_if_pred_node` as attribute + // _xla_token_input_nodes. + Node *if_node = node_name_index["if"]; + EXPECT_NE(if_node, nullptr); + std::vector token_inputs; + TF_CHECK_OK( + GetNodeAttr(if_node->def(), "_xla_token_input_nodes", &token_inputs)); + EXPECT_THAT(token_inputs, ::testing::ElementsAre("send_oc_if_pred_if")); + } +} + +TEST_F(ExtractOutsideCompilationForFunctionTest, OutsideCompilationInWhile) { + // Build the XLA computation func. + // "const0" (bool) + // "while0" (input = "const0", cond = "cond_fn", body = "body_fn") + FunctionDefLibrary fdl; + { + tensorflow::Scope s = tensorflow::Scope::NewRootScope(); + Output arg = ops::_Arg(s.WithOpName("arg"), DT_BOOL, 0); + Output identity = ops::Identity(s.WithOpName("identity_cond_fn"), 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_cond_fn"]->AddAttr("_oc", "0"); + PartialTensorShape shape({2}); + node_name_image["identity_cond_fn"]->AddAttr( + kXlaInferredShapesAttrName, std::vector{shape}); + + FunctionDef *cond_fn_fdef = fdl.add_function(); + TF_CHECK_OK(GraphToFunctionDef(*g, "cond_fn", cond_fn_fdef)); + } + { + tensorflow::Scope s = tensorflow::Scope::NewRootScope(); + Output arg = ops::_Arg(s.WithOpName("arg"), DT_BOOL, 0); + Output identity = ops::Identity(s.WithOpName("identity_body_fn"), 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_body_fn"]->AddAttr("_oc", "0"); + PartialTensorShape shape({2}); + node_name_image["identity_body_fn"]->AddAttr( + kXlaInferredShapesAttrName, std::vector{shape}); + + FunctionDef *body_fn_fdef = fdl.add_function(); + TF_CHECK_OK(GraphToFunctionDef(*g, "body_fn", body_fn_fdef)); + } + { + tensorflow::Scope s = tensorflow::Scope::NewRootScope(); + Output input = ops::Const(s.WithOpName("const0"), true, {2}); + NameAttrList cond_fn; + cond_fn.set_name("cond_fn"); + NameAttrList body_fn; + body_fn.set_name("body_fn"); + auto while_op = + ops::While(s.WithOpName("while"), std::initializer_list{input}, + cond_fn, body_fn); + ops::_Retval retval(s.WithOpName("retval"), while_op.output[0], 0); + std::unique_ptr g(new Graph(OpRegistry::Global())); + TF_CHECK_OK(s.ToGraph(g.get())); + + FunctionDef *xla_fdef = fdl.add_function(); + TF_CHECK_OK(GraphToFunctionDef(*g, "cluster", xla_fdef)); + } + FunctionLibraryDefinition fld(OpRegistry::Global(), fdl); + + 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 an "While" to execute outside compilation. + Node *while_oc_node = node_name_index["oc_while_while"]; + EXPECT_NE(while_oc_node, nullptr); + + // Check that cond outside compilation has node "identity_cond_fn". + const FunctionDef *cond_def = fld.Find("oc_cond_host_while_while"); + EXPECT_NE(cond_def, nullptr); + bool has_identity_cond_fn_node = false; + for (const auto &node_def : cond_def->node_def()) { + if (node_def.name() == "identity_cond_fn") { + has_identity_cond_fn_node = true; + break; + } + } + EXPECT_TRUE(has_identity_cond_fn_node); + + // Check that body outside compilation has node "identity_body_fn". + const FunctionDef *body_def = fld.Find("oc_body_host_while_while"); + EXPECT_NE(body_def, nullptr); + bool has_identity_body_fn_node = false; + for (const auto &node_def : body_def->node_def()) { + if (node_def.name() == "identity_body_fn") { + has_identity_body_fn_node = true; + break; + } + } + EXPECT_TRUE(has_identity_body_fn_node); + } + + // Check XLA graph. + { + // Verify that rewritten cond fn has XlaSendToHost to send loop predicate to + // host. + const FunctionDef *cond_def = fld.Find("cond_fn_oc"); + EXPECT_NE(cond_def, nullptr); + bool has_send_oc_while_cond_node = false; + for (const auto &node_def : cond_def->node_def()) { + if (node_def.name() == "send_oc_while_cond_while") { + has_send_oc_while_cond_node = true; + break; + } + } + EXPECT_TRUE(has_send_oc_while_cond_node); + } +} + +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 new file mode 100644 index 0000000000000000000000000000000000000000..98e344b3a080aa8aab27cd41564a90427bac151e --- /dev/null +++ b/tensorflow/compiler/jit/flags.cc @@ -0,0 +1,152 @@ +/* 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 // NOLINT + +#include "tensorflow/compiler/jit/flags.h" +#include "tensorflow/compiler/xla/parse_flags_from_env.h" +#include "tensorflow/core/util/command_line_flags.h" + +namespace tensorflow { +namespace { + +BuildXlaOpsPassFlags* build_ops_flags; +DumpGraphFlags* dump_graph_flags; +MarkForCompilationPassFlags* mark_for_compilation_flags; +XlaDeviceFlags* device_flags; +XlaOpsCommonFlags* ops_flags; + +std::vector* flag_list; +std::once_flag flags_init; + +void AppendDumpGraphFlagsInternal(std::vector* flag_list) { + std::vector new_flags = { + Flag("tf_dump_graph_prefix", &dump_graph_flags->tf_dump_graph_prefix, + "Path prefix to which graphs dumped during debugging should be " + "written."), + }; + flag_list->insert(flag_list->end(), new_flags.begin(), new_flags.end()); +} + +void AppendMarkForCompilationPassFlagsInternal(std::vector* flag_list) { + std::vector new_flags = { + Flag("tf_xla_auto_jit", &mark_for_compilation_flags->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."), + Flag("tf_xla_min_cluster_size", + &mark_for_compilation_flags->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."), + Flag("tf_xla_max_cluster_size", + &mark_for_compilation_flags->tf_xla_max_cluster_size, + "Maximum number of operators in an XLA compilation."), + Flag("tf_xla_clustering_debug", + &mark_for_compilation_flags->tf_xla_clustering_debug, + "Dump graphs during XLA compilation."), + Flag("tf_xla_cpu_global_jit", + &mark_for_compilation_flags->tf_xla_cpu_global_jit, + "Enables global JIT compilation for CPU via SessionOptions."), + Flag("tf_xla_clustering_fuel", + &mark_for_compilation_flags->tf_xla_clustering_fuel, + "Places an artificial limit on the number of ops marked as " + "eligible for clustering."), + 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*.")}; + flag_list->insert(flag_list->end(), new_flags.begin(), new_flags.end()); +} + +void AllocateAndParseFlags() { + build_ops_flags = new BuildXlaOpsPassFlags; + build_ops_flags->tf_xla_enable_lazy_compilation = true; + + dump_graph_flags = new DumpGraphFlags; + dump_graph_flags->tf_dump_graph_prefix = "/tmp/"; + + mark_for_compilation_flags = new MarkForCompilationPassFlags; + mark_for_compilation_flags->tf_xla_auto_jit = 0; + mark_for_compilation_flags->tf_xla_min_cluster_size = 2; + mark_for_compilation_flags->tf_xla_max_cluster_size = + std::numeric_limits::max(); + mark_for_compilation_flags->tf_xla_clustering_debug = false; + mark_for_compilation_flags->tf_xla_cpu_global_jit = false; + mark_for_compilation_flags->tf_xla_clustering_fuel = + std::numeric_limits::max(); + mark_for_compilation_flags->tf_xla_fusion_only = false; + + device_flags = new XlaDeviceFlags; + device_flags->tf_xla_compile_on_demand = false; + + ops_flags = new XlaOpsCommonFlags; + ops_flags->tf_xla_always_defer_compilation = false; + + flag_list = new std::vector({ + Flag("tf_xla_enable_lazy_compilation", + &build_ops_flags->tf_xla_enable_lazy_compilation, ""), + + Flag("tf_xla_compile_on_demand", &device_flags->tf_xla_compile_on_demand, + "Switch a device into 'on-demand' mode, where instead of " + "autoclustering ops are compiled one by one just-in-time."), + + Flag("tf_xla_always_defer_compilation", + &ops_flags->tf_xla_always_defer_compilation, ""), + }); + AppendDumpGraphFlagsInternal(flag_list); + AppendMarkForCompilationPassFlagsInternal(flag_list); + xla::ParseFlagsFromEnvAndDieIfUnknown("TF_XLA_FLAGS", *flag_list); +} + +} // namespace + +const BuildXlaOpsPassFlags& GetBuildXlaOpsPassFlags() { + std::call_once(flags_init, &AllocateAndParseFlags); + return *build_ops_flags; +} + +DumpGraphFlags* GetDumpGraphFlags() { + std::call_once(flags_init, &AllocateAndParseFlags); + return dump_graph_flags; +} + +MarkForCompilationPassFlags* GetMarkForCompilationPassFlags() { + std::call_once(flags_init, &AllocateAndParseFlags); + return mark_for_compilation_flags; +} + +XlaDeviceFlags* GetXlaDeviceFlags() { + std::call_once(flags_init, &AllocateAndParseFlags); + return device_flags; +} + +const XlaOpsCommonFlags& GetXlaOpsCommonFlags() { + std::call_once(flags_init, &AllocateAndParseFlags); + return *ops_flags; +} + +void AppendMarkForCompilationPassFlags(std::vector* flag_list) { + std::call_once(flags_init, &AllocateAndParseFlags); + AppendMarkForCompilationPassFlagsInternal(flag_list); +} + +void AppendDumpGraphFlags(std::vector* flag_list) { + std::call_once(flags_init, &AllocateAndParseFlags); + AppendDumpGraphFlagsInternal(flag_list); +} + +} // namespace tensorflow diff --git a/tensorflow/compiler/jit/flags.h b/tensorflow/compiler/jit/flags.h new file mode 100644 index 0000000000000000000000000000000000000000..5ddea588eef5270880d91623dc05893da265960a --- /dev/null +++ b/tensorflow/compiler/jit/flags.h @@ -0,0 +1,103 @@ +/* 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_JIT_FLAGS_H_ +#define TENSORFLOW_COMPILER_JIT_FLAGS_H_ + +#include + +#include "tensorflow/core/platform/types.h" +#include "tensorflow/core/util/command_line_flags.h" + +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. +}; + +// Flags associated with the XLA bridge's xla_device module. +struct XlaDeviceFlags { + // Switch the CPU device into "on-demand" mode, where instead of + // autoclustering ops are compiled one by one just-in-time. + // Enabling this mode by a legacy flag is a temporary mechanism. When this + // feature is battle-tested, we will switch this to be a session option. + bool tf_xla_compile_on_demand; +}; + +// Flags common to the _Xla* ops and their kernels. +struct XlaOpsCommonFlags { + // If true, _XlaCompile always refuses to compile the cluster, which means the + // XLA clusters always run in the TF executor. Defaults to false. + bool tf_xla_always_defer_compilation; +}; + +// Flags for the build_xla_ops pass. +struct BuildXlaOpsPassFlags { + // Enables lazy compilation for TF/XLA (only when auto-clustering) if true. + // Defaults to true. + bool tf_xla_enable_lazy_compilation; +}; + +// Flags for the XLA bridge's dump_graph module. +struct DumpGraphFlags { + // Path prefix to which graphs dumped during debugging should be written. + string tf_dump_graph_prefix; +}; + +// Return a pointer to the DumpGraphFlags struct; +// repeated calls return the same pointer. +// This should be called only after Flags::Parse() has returned. + +// Getters for flags structs defined above. The first call to any of these +// parses TF_XLA_FLAGS for all of them. Those functions which return a pointer +// always return the same pointer. +MarkForCompilationPassFlags* GetMarkForCompilationPassFlags(); +const BuildXlaOpsPassFlags& GetBuildXlaOpsPassFlags(); +XlaDeviceFlags* GetXlaDeviceFlags(); +const XlaOpsCommonFlags& GetXlaOpsCommonFlags(); +DumpGraphFlags* GetDumpGraphFlags(); + +// Appends the flag definitions associated with +// MarkForCompilationPassFlags/DumpGraphFlags to `flag_list`. +// +// Has the side-effect of parsing TF_XLA_FLAGS if that hasn't happened yet. +void AppendMarkForCompilationPassFlags( + std::vector* flag_list); +void AppendDumpGraphFlags(std::vector* flag_list); + +} // namespace tensorflow + +#endif // TENSORFLOW_COMPILER_JIT_FLAGS_H_ diff --git a/tensorflow/compiler/jit/increase_dynamism_for_auto_jit_pass.cc b/tensorflow/compiler/jit/increase_dynamism_for_auto_jit_pass.cc new file mode 100644 index 0000000000000000000000000000000000000000..ce53f70b79d97ab087fefe542920b33f883632a2 --- /dev/null +++ b/tensorflow/compiler/jit/increase_dynamism_for_auto_jit_pass.cc @@ -0,0 +1,364 @@ +/* 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/jit/increase_dynamism_for_auto_jit_pass.h" +#include "absl/algorithm/container.h" +#include "absl/container/inlined_vector.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/str_replace.h" +#include "absl/types/optional.h" +#include "tensorflow/cc/framework/scope_internal.h" +#include "tensorflow/cc/ops/array_ops.h" +#include "tensorflow/cc/ops/const_op.h" +#include "tensorflow/cc/ops/math_ops.h" +#include "tensorflow/compiler/jit/flags.h" +#include "tensorflow/compiler/jit/xla_cluster_util.h" +#include "tensorflow/compiler/tf2xla/cc/ops/xla_ops.h" +#include "tensorflow/compiler/tf2xla/dump_graph.h" +#include "tensorflow/compiler/xla/status_macros.h" +#include "tensorflow/core/common_runtime/shape_refiner.h" +#include "tensorflow/core/graph/algorithm.h" +#include "tensorflow/core/public/session_options.h" +#include "tensorflow/core/util/device_name_utils.h" + +namespace tensorflow { +namespace { + +// StatusOrOptional instances hold +// +// - A non-OK Status to indicate an error that needs to be propagated out of +// this pass (e.g. the Graph is malformed). +// +// - A nullopt to indicate the function that created the instance failed to do +// what it set out to do but this is not actually an error +// (e.g. TryToGetTensorFromConstOp was passed a non-Const node). +// +// - A T to indicate a successful operation. +template +using StatusOrOptional = xla::StatusOr>; + +StatusOrOptional TryToGetTensorFromConstOp(Node* n) { + if (n->type_string() != "Const") { + return {absl::nullopt}; + } + + const TensorProto* proto = nullptr; + TF_RETURN_IF_ERROR(GetNodeAttr(n->def(), "value", &proto)); + Tensor tensor(proto->dtype()); + TF_RET_CHECK(tensor.FromProto(*proto)); + return {tensor}; +} + +struct SliceInputs { + Output slice_op; + Output input; + Output begin; + Output size; + + // The size of the TF slice operation as a std::vector. We can always compute + // this because we only manipulate slices with a Const size. + std::vector size_as_vector; +}; + +std::vector IntTensorAsVector(const Tensor& t) { + DCHECK(t.dtype() == DT_INT32 || t.dtype() == DT_INT64); + std::vector result; + result.reserve(t.NumElements()); + for (int i = 0; i < t.NumElements(); i++) { + int64 element = t.dtype() == DT_INT32 + ? static_cast(t.flat()(i)) + : t.flat()(i); + result.push_back(element); + } + return result; +} + +// Packages up the inputs to a Slice operation into an instance of +// `SliceInputs`. +StatusOrOptional GetSliceInputs(Node* slice) { + const int kSliceInputIndex = 0; + const int kSliceBeginIndex = 1; + const int kSliceSizeIndex = 2; + + const Edge* slice_input_edge; + TF_RETURN_IF_ERROR(slice->input_edge(kSliceInputIndex, &slice_input_edge)); + const Edge* slice_size_edge; + TF_RETURN_IF_ERROR(slice->input_edge(kSliceSizeIndex, &slice_size_edge)); + const Edge* slice_begin_edge; + TF_RETURN_IF_ERROR(slice->input_edge(kSliceBeginIndex, &slice_begin_edge)); + + SliceInputs slice_inputs; + slice_inputs.input = + Output(slice_input_edge->src(), slice_input_edge->src_output()); + slice_inputs.begin = + Output(slice_begin_edge->src(), slice_begin_edge->src_output()); + slice_inputs.size = + Output(slice_size_edge->src(), slice_size_edge->src_output()); + + TF_ASSIGN_OR_RETURN(absl::optional tf_slice_size, + TryToGetTensorFromConstOp(slice_inputs.size.node())); + if (!tf_slice_size.has_value()) { + return {absl::nullopt}; + } + + if (tf_slice_size->dims() != 1) { + return {absl::nullopt}; + } + + slice_inputs.size_as_vector = IntTensorAsVector(*tf_slice_size); + return {slice_inputs}; +} + +// Casts `x` to a DT_INT64 if it isn't one already. +Output MakeInt64(const Scope& host_scope, absl::string_view name, + const Output& x) { + return x.type() == DT_INT64 + ? x + : ops::Cast(host_scope.WithOpName(name, "_s64"), x, DT_INT64); +} + +// Returns `slice_inputs` with the index and size inputs cast to DT_INT64. +SliceInputs MakeSliceIndexAndSizeInt64(const Scope& host_scope, + const SliceInputs& slice_inputs) { + SliceInputs result; + result.input = slice_inputs.input; + result.begin = MakeInt64(host_scope, "begin", slice_inputs.begin); + result.size = MakeInt64(host_scope, "size", slice_inputs.size); + result.size_as_vector = slice_inputs.size_as_vector; + return result; +} + +// This class caches emitted constants to avoid creating multiple nodes for the +// same constant value. This helps make the generated GraphDef more readable. +class ConstantCache { + public: + explicit ConstantCache(const Scope& s) : scope_(s) {} + + Output Get1DHostConstant(int64 constant) { + auto it = cache_.find(constant); + if (it == cache_.end()) { + Output new_const = + ops::Const(scope_.WithOpName("const_", constant), {constant}); + it = cache_.insert({constant, new_const}).first; + } + return it->second; + } + + private: + Scope scope_; + std::unordered_map cache_; +}; + +// 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) { + // If slice_size[i] >= 0 then slice_size[i] = slice_size[i]. + // + // If slice_size[i] == -1 then slice_size[i] = input_size[i] - + // begin[i]. + // + // If slice_size[i] < -1 then executing the slice will throw an error, and we + // don't do anything here. We've already filtered these cases out in + // IsRewritableSlice. + + if (absl::c_all_of(slice_inputs.size_as_vector, + [](int64 i) { return i >= 0; })) { + *size = slice_inputs.size; + return Status::OK(); + } + + Output input_shape = + ops::Shape(host_scope.WithOpName("input_shape"), slice_inputs.input, + ops::Shape::OutType(DT_INT64)); + + ConstantCache constant_pool(host_scope); + + std::vector slice_size; + for (int i = 0; i < slice_inputs.size_as_vector.size(); i++) { + if (slice_inputs.size_as_vector[i] >= 0) { + slice_size.push_back( + constant_pool.Get1DHostConstant(slice_inputs.size_as_vector[i])); + continue; + } + + DCHECK_EQ(slice_inputs.size_as_vector[i], -1); + + Output begin_i = ops::Slice( + host_scope.WithOpName("begin_", i), slice_inputs.begin, + constant_pool.Get1DHostConstant(i), constant_pool.Get1DHostConstant(1)); + + Output input_shape_i = ops::Slice( + host_scope.WithOpName("input_shape_", i), input_shape, + constant_pool.Get1DHostConstant(i), constant_pool.Get1DHostConstant(1)); + + slice_size.push_back(ops::Sub(host_scope.WithOpName("slice_size_", i), + input_shape_i, begin_i)); + DCHECK_EQ(slice_size.back().type(), DT_INT64); + } + + // 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)); + return Status::OK(); +} + +// Terminology: "static sized" slice is a slice with the +// _XlaCompileTimeConstantInputs attribute set to {2}. The output shape of +// these slices can be solely determined by their "size" input. +Status ConvertTensorFlowSliceToStaticShapedSlice( + Graph* g, Node* slice, const SliceInputs& slice_inputs, + absl::string_view cluster_name, Node** result) { + string host_name; + TF_RETURN_IF_ERROR(DeviceNameUtils::DeviceNameToCpuDeviceName( + slice->assigned_device_name(), &host_name)); + + Status status; + Scope main_scope = + NewInternalScope(g, &status, /*refiner=*/nullptr) + .WithXlaCluster(string(cluster_name)) + .NewSubScope(absl::StrCat(slice->name(), "/static_shaped_slice")); + Scope host_scope = main_scope.WithAssignedDevice(host_name); + + SliceInputs slice_inputs_int64 = + MakeSliceIndexAndSizeInt64(host_scope, slice_inputs); + + Output slice_size; + TF_RETURN_IF_ERROR( + ComputeSliceSize(host_scope, slice_inputs_int64, &slice_size)); + + *result = + ops::Slice(main_scope.WithAssignedDevice(slice->assigned_device_name()) + .WithOpName("static_shaped_slice"), + slice_inputs_int64.input, slice_inputs_int64.begin, slice_size) + .node(); + + TF_RETURN_IF_ERROR(main_scope.status()); + + std::vector compile_time_const_inputs; + compile_time_const_inputs.push_back("size"); + (*result)->AddAttr(kXlaCompileTimeConstantInputsAttr, + compile_time_const_inputs); + return status; +} + +void ReplaceTensorFlowSliceWithStaticShapedSlice(Graph* g, Node* slice, + Node* static_shaped_slice) { + absl::InlinedVector edges_to_remove; + std::vector slice_out_edges; + absl::c_copy(slice->out_edges(), std::back_inserter(slice_out_edges)); + for (const Edge* e : slice_out_edges) { + DCHECK(e->src_output() == 0 || e->src_output() == Graph::kControlSlot); + + int src_output = e->src_output(); + int dst_input = e->dst_input(); + Node* dst = e->dst(); + g->RemoveEdge(e); + g->AddEdge(static_shaped_slice, src_output, dst, dst_input); + } + + for (const Edge* e : slice->in_edges()) { + if (e->IsControlEdge()) { + g->AddControlEdge(e->src(), static_shaped_slice); + } + } + + g->RemoveNode(slice); +} + +Status RewriteSlice(Graph* g, Node* slice, const SliceInputs& slice_inputs, + absl::string_view cluster_name) { + VLOG(3) << "Rewriting slice " << slice->name() + << " to a \"static shaped\" Slice"; + Node* static_shaped_slice; + TF_RETURN_IF_ERROR(ConvertTensorFlowSliceToStaticShapedSlice( + g, slice, slice_inputs, cluster_name, &static_shaped_slice)); + ReplaceTensorFlowSliceWithStaticShapedSlice(g, slice, static_shaped_slice); + return Status::OK(); +} + +// Return true if `n` is a slice we can rewrite to have a static shape +// (i.e. have the output shape only depend on the "size" input). +xla::StatusOr IsRewritableSlice(Node* n) { + if (n->type_string() != "Slice") { + return false; + } + + if (!GetXlaClusterForNode(*n).has_value()) { + // There is no need to change slice ops outside XLA clusters. + return false; + } + + TF_ASSIGN_OR_RETURN(absl::optional slice_inputs, + GetSliceInputs(n)); + if (!slice_inputs.has_value()) { + return false; + } + + // 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; }); +} + +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)); + if (is_rewritable) { + slices_to_rewrite.push_back(n); + } + } + + for (Node* n : slices_to_rewrite) { + TF_ASSIGN_OR_RETURN(absl::optional slice_inputs, + GetSliceInputs(n)); + TF_RET_CHECK(slice_inputs.has_value()); + TF_RETURN_IF_ERROR( + RewriteSlice(g, n, *slice_inputs, *GetXlaClusterForNode(*n))); + } + + if (!slices_to_rewrite.empty()) { + // We've added constants to the graph; hook them up to _SOURCE. + FixupSourceAndSinkEdges(g); + } + + *changed = !slices_to_rewrite.empty(); + + return Status::OK(); +} +} // namespace + +Status IncreaseDynamismForAutoJitPass::Run( + const GraphOptimizationPassOptions& options) { + MarkForCompilationPassFlags* flags = GetMarkForCompilationPassFlags(); + if (flags->tf_xla_clustering_debug) { + dump_graph::DumpGraphToFile("before_increase_dynamism_for_auto_jit_pass", + **options.graph, options.flib_def); + } + + bool changed; + TF_RETURN_IF_ERROR(FindAndRewriteSlices(options.graph->get(), &changed)); + if (changed && flags->tf_xla_clustering_debug) { + dump_graph::DumpGraphToFile("increase_dynamism_for_auto_jit_pass", + **options.graph, options.flib_def); + } + + return Status::OK(); +} + +} // namespace tensorflow diff --git a/tensorflow/compiler/jit/increase_dynamism_for_auto_jit_pass.h b/tensorflow/compiler/jit/increase_dynamism_for_auto_jit_pass.h new file mode 100644 index 0000000000000000000000000000000000000000..818ca948d64b0353b08f393c3bd7d874c9b2480b --- /dev/null +++ b/tensorflow/compiler/jit/increase_dynamism_for_auto_jit_pass.h @@ -0,0 +1,57 @@ +/* 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_JIT_INCREASE_DYNAMISM_FOR_AUTO_JIT_PASS_H_ +#define TENSORFLOW_COMPILER_JIT_INCREASE_DYNAMISM_FOR_AUTO_JIT_PASS_H_ + +#include "tensorflow/core/common_runtime/optimization_registry.h" + +namespace tensorflow { + +// Increases the amount of "dynamism" representable by XLA clusters by rewriting +// the TensorFlow graph. This pass does the following rewrites: +// +// Slice +// ----- +// +// Slice(op, begin, size ) => +// Slice(op, begin, actual_size(op.shape(), size, begin)); +// _XlaCompileTimeConstantInputs={2} +// +// where +// +// actual_size(op_shape, size, begin)[i] = +// size[i] == -1 ? (op_shape[i] - size[i]) +// : size[i] +// +// This pass, combined with jit/partially_decluster_pass, reduces the number of +// unnecessary cluster recompilations in some common cases. After the rewrite +// shown above jit/partially_decluster_pass extracts the actual_size(...) +// computation to outside the XLA cluster, causing the cluster to be versioned +// only on the actual size of the XlaDynamicSlice. This avoids recompilation +// due to superficial changes that don't affect tensor shapes. +// +// Future Work TODO(b/111210515) +// ----------------------------- +// +// In the future we will also translate StridedSlice and Pad a similar way. +class IncreaseDynamismForAutoJitPass : public GraphOptimizationPass { + public: + Status Run(const GraphOptimizationPassOptions& options) override; +}; + +} // namespace tensorflow + +#endif // TENSORFLOW_COMPILER_JIT_INCREASE_DYNAMISM_FOR_AUTO_JIT_PASS_H_ 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 new file mode 100644 index 0000000000000000000000000000000000000000..a2f1b831ad7605237e23c15cc43b337e06265553 --- /dev/null +++ b/tensorflow/compiler/jit/increase_dynamism_for_auto_jit_pass_test.cc @@ -0,0 +1,405 @@ +/* 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/jit/increase_dynamism_for_auto_jit_pass.h" + +#include "tensorflow/cc/framework/ops.h" +#include "tensorflow/cc/ops/array_ops.h" +#include "tensorflow/cc/ops/const_op.h" +#include "tensorflow/compiler/jit/node_matchers.h" +#include "tensorflow/compiler/jit/xla_cluster_util.h" +#include "tensorflow/core/lib/core/status_test_util.h" +#include "tensorflow/core/platform/test.h" +#include "tensorflow/core/public/session_options.h" + +namespace tensorflow { +namespace { + +using ::testing::_; +using testing::matchers::AssignedDevice; +using testing::matchers::Attr; +using testing::matchers::Const; +using testing::matchers::CtrlDeps; +using testing::matchers::Inputs; +using testing::matchers::Name; +using testing::matchers::NodeWith; +using testing::matchers::Op; +using testing::matchers::Out; + +// A fake device used to populate a DeviceSet. +class FakeDevice : public Device { + public: + explicit FakeDevice(const DeviceAttributes& device_attributes) + : Device(nullptr, device_attributes) {} + + Status Sync() override { return errors::Unimplemented("FakeDevice::Sync()"); } + + Allocator* GetAllocator(AllocatorAttributes attr) override { return nullptr; } + + static std::unique_ptr Make(const string& name, const string& type) { + DeviceAttributes device_attributes; + device_attributes.set_name(name); + device_attributes.set_device_type(DeviceType(type).type()); + return absl::make_unique(device_attributes); + } +}; + +const char* kHostName = "/job:worker/replica:0/task:0/device:CPU:0"; +const char* kDeviceName = "/job:worker/replica:0/task:0/device:GPU:0"; + +Status IncreaseDynamismForAutoJit(const Scope& s, + std::unique_ptr* result) { + std::vector> devices; + devices.push_back(FakeDevice::Make(kDeviceName, DEVICE_GPU)); + devices.push_back(FakeDevice::Make(kHostName, DEVICE_CPU)); + + std::unique_ptr device_set(new DeviceSet()); + for (auto& device : devices) { + device_set->AddDevice(device.get()); + } + + auto graph = absl::make_unique(OpRegistry::Global()); + SessionOptions session_options; + session_options.config.mutable_graph_options() + ->mutable_optimizer_options() + ->set_global_jit_level(OptimizerOptions::ON_2); + GraphOptimizationPassOptions options; + options.graph = &graph; + options.device_set = device_set.get(); + options.session_options = &session_options; + + // Scope::ToGraph seems to drop assigned devices, probably because it goes + // through a GraphDef. So explicitly maintain the device assignment. + std::unordered_map assigned_device_names; + for (Node* n : s.graph()->nodes()) { + assigned_device_names[n->name()] = n->assigned_device_name(); + } + TF_RETURN_IF_ERROR(s.ToGraph(graph.get())); + for (Node* n : graph->nodes()) { + n->set_assigned_device_name(assigned_device_names[n->name()]); + } + + IncreaseDynamismForAutoJitPass rewriter; + TF_RETURN_IF_ERROR(rewriter.Run(options)); + *result = std::move(graph); + return Status::OK(); +} + +TEST(SliceToDynamicSliceRewriteTest, Basic) { + 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, 500}); + Output slice = ops::Slice(root.WithOpName("slice"), input, begin, size); + + std::unique_ptr result; + TF_ASSERT_OK(IncreaseDynamismForAutoJit(root, &result)); + + const int64 zero_64 = 0; + const int32 zero_32 = 0; + const int64 one_64 = 1; + + auto m_input = Out(NodeWith(Op("Placeholder"), Name("input"))); + auto m_begin_s64 = Out(NodeWith( + Op("Cast"), Inputs(Out(NodeWith(Op("Placeholder"), Name("begin")))))); + auto m_input_shape = Out(NodeWith(Op("Shape"), Inputs(m_input))); + auto m_slice_size_0 = Out(NodeWith( + Op("Sub"), AssignedDevice(kHostName), + Inputs( + Out(NodeWith(Op("Slice"), AssignedDevice(kHostName), + Inputs(m_input_shape, Const(zero_64), Const(one_64)))), + Out(NodeWith(Op("Slice"), AssignedDevice(kHostName), + Inputs(m_begin_s64, Const(zero_64), Const(one_64))))))); + auto m_dynamic_slice_size = Out(NodeWith( + Op("ConcatV2"), AssignedDevice(kHostName), + Inputs(m_slice_size_0, Const(static_cast(500)), Const(zero_32)))); + + std::vector compile_time_constant_inputs; + compile_time_constant_inputs.push_back("size"); + auto m_dynamic_slice = NodeWith( + Op("Slice"), AssignedDevice(kDeviceName), + Attr(kXlaCompileTimeConstantInputsAttr, compile_time_constant_inputs), + Inputs(m_input, m_begin_s64, m_dynamic_slice_size)); + + Node* static_shaped_slice = testing::FindNodeByName( + result.get(), "slice/static_shaped_slice/static_shaped_slice"); + ASSERT_NE(static_shaped_slice, nullptr); + EXPECT_THAT(static_shaped_slice, m_dynamic_slice); +} + +TEST(SliceToDynamicSliceRewriteTest, SliceFromVector) { + 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); + + std::unique_ptr result; + TF_ASSERT_OK(IncreaseDynamismForAutoJit(root, &result)); + + Node* static_shaped_slice = testing::FindNodeByName( + result.get(), "slice/static_shaped_slice/static_shaped_slice"); + EXPECT_NE(static_shaped_slice, nullptr); + EXPECT_THAT(result->nodes(), Not(Contains(NodeWith(Op("ConcatV2"))))); +} + +TEST(SliceToDynamicSliceRewriteTest, ControlDependencePreserved) { + 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, 500}); + Output control_pred = ops::Placeholder(root.WithOpName("control"), DT_BOOL); + Output slice = ops::Slice(root.WithOpName("slice"), input, begin, size); + root.graph()->AddControlEdge(control_pred.node(), slice.node()); + + std::unique_ptr result; + TF_ASSERT_OK(IncreaseDynamismForAutoJit(root, &result)); + + Node* static_shaped_slice = testing::FindNodeByName( + result.get(), "slice/static_shaped_slice/static_shaped_slice"); + ASSERT_NE(static_shaped_slice, nullptr); + EXPECT_THAT(static_shaped_slice, + NodeWith(Op("Slice"), + CtrlDeps(NodeWith(Op("Placeholder"), Name("control"))))); +} + +int64 ToInt64(int v) { return static_cast(v); } + +TEST(SliceToDynamicSliceRewriteTest, Int64Indices) { + 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_INT64); + Output size = + ops::Const(root.WithOpName("size"), {ToInt64(-1), ToInt64(500)}); + Output slice = ops::Slice(root.WithOpName("slice"), input, begin, size); + + std::unique_ptr result; + TF_ASSERT_OK(IncreaseDynamismForAutoJit(root, &result)); + + EXPECT_THAT(result->nodes(), Not(Contains(NodeWith(Op("Cast"))))); +} + +TEST(SliceToDynamicSliceRewriteTest, DontRewriteInvalidSlice) { + 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); + + // The shape refiner throws an error if we use a bogus constant value for + // size. So we first use a Placeholder to placate the shape refiner, and + // later replace it with a bogus constant. + Output size_placeholder = + ops::Placeholder(root.WithOpName("size_placeholder"), DT_INT32); + Output slice = + ops::Slice(root.WithOpName("slice"), input, begin, size_placeholder); + + Output size = ops::Const(root.WithOpName("size"), {-8, 500}); + TF_ASSERT_OK(root.graph()->UpdateEdge(/*new_src=*/size.node(), + /*new_src_index=*/0, + /*dst=*/slice.node(), /*dst_index=*/2)); + + std::unique_ptr result; + TF_ASSERT_OK(IncreaseDynamismForAutoJit(root, &result)); + + EXPECT_THAT(result->nodes(), + Not(Contains(NodeWith(Op("Slice"), + Attr(kXlaCompileTimeConstantInputsAttr))))); +} + +TEST(SliceToDynamicSliceRewriteTest, DontRewriteUnclusteredSlice) { + Scope root = + Scope::NewRootScope().ExitOnError().WithAssignedDevice(kDeviceName); + + 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, 500}); + Output slice = ops::Slice(root.WithOpName("slice"), input, begin, size); + + std::unique_ptr result; + TF_ASSERT_OK(IncreaseDynamismForAutoJit(root, &result)); + + EXPECT_THAT(result->nodes(), + Not(Contains(NodeWith(Op("Slice"), + Attr(kXlaCompileTimeConstantInputsAttr))))); +} + +TEST(SliceToDynamicSliceRewriteTest, DontRewriteSliceWithNonConstSize) { + 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_INT64); + Output size = ops::Placeholder(root.WithOpName("size"), DT_INT64); + Output slice = ops::Slice(root.WithOpName("slice"), input, begin, size); + + std::unique_ptr result; + TF_ASSERT_OK(IncreaseDynamismForAutoJit(root, &result)); + + EXPECT_THAT(result->nodes(), + Not(Contains(NodeWith(Op("Slice"), + Attr(kXlaCompileTimeConstantInputsAttr))))); +} + +TEST(SliceToDynamicSliceRewriteTest, ScalarSlice) { + 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_INT64); + Output size = ops::Const(root.WithOpName("size"), {}); + Output slice = ops::Slice(root.WithOpName("slice"), input, begin, size); + + std::unique_ptr result; + TF_ASSERT_OK(IncreaseDynamismForAutoJit(root, &result)); + + Node* static_shaped_slice = testing::FindNodeByName( + result.get(), "slice/static_shaped_slice/static_shaped_slice"); + ASSERT_NE(static_shaped_slice, nullptr); + EXPECT_THAT(static_shaped_slice, + NodeWith(Op("Slice"), Attr(kXlaCompileTimeConstantInputsAttr), + Inputs(_, _, Out(NodeWith(Name(size.node()->name())))))); +} + +TEST(SliceToDynamicSliceRewriteTest, IndicesNotVector) { + Scope root = Scope::NewRootScope() + .ExitOnError() + .WithAssignedDevice(kDeviceName) + .WithXlaCluster("cluster_0"); + + auto ToInt64 = [](int v) { return static_cast(v); }; + + Output input = ops::Placeholder(root.WithOpName("input"), DT_FLOAT); + Output begin = ops::Placeholder(root.WithOpName("begin"), DT_INT64); + + // The C++ node bindings immediately error out when we try construct a bogus + // slice so we first use a placeholder to construct the Slice and then replace + // the input. + Output size_placeholder = ops::Placeholder(root.WithOpName("size"), DT_INT64); + Output slice = + ops::Slice(root.WithOpName("slice"), input, begin, size_placeholder); + + Output size = + ops::Const(root.WithOpName("size"), {{ToInt64(-1)}, {ToInt64(500)}}); + TF_ASSERT_OK(root.graph()->UpdateEdge(size.node(), 0, slice.node(), 2)); + + std::unique_ptr result; + TF_ASSERT_OK(IncreaseDynamismForAutoJit(root, &result)); + + EXPECT_THAT(result->nodes(), + Not(Contains(NodeWith(Op("Slice"), + Attr(kXlaCompileTimeConstantInputsAttr))))); +} + +TEST(SliceToDynamicSliceRewriteTest, SliceWithSliceInput) { + 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_a = ops::Const(root.WithOpName("size_a"), {-1, 500}); + Output slice = ops::Slice(root.WithOpName("slice"), input, begin, size_a); + + Output size_b = ops::Const(root.WithOpName("size_a"), {-1, 200}); + Output slice_with_slice_input = ops::Slice( + root.WithOpName("slice_with_slice_input"), slice, begin, size_b); + + std::unique_ptr result; + TF_ASSERT_OK(IncreaseDynamismForAutoJit(root, &result)); + + Node* static_shaped_slice = testing::FindNodeByName( + result.get(), + "slice_with_slice_input/static_shaped_slice/static_shaped_slice"); + ASSERT_NE(static_shaped_slice, nullptr); + EXPECT_EQ(static_shaped_slice->output_type(0), DT_FLOAT) + << "Expected DT_FLOAT, was " + << DataType_Name(static_shaped_slice->output_type(0)); + EXPECT_THAT( + static_shaped_slice, + NodeWith( + Op("Slice"), + Inputs(Out(NodeWith( + Op("Slice"), + Name("slice/static_shaped_slice/static_shaped_slice"))), + _, _))); +} + +TEST(SliceToDynamicSliceRewriteTest, SliceWithSliceBegin) { + Scope root = Scope::NewRootScope() + .ExitOnError() + .WithAssignedDevice(kDeviceName) + .WithXlaCluster("cluster_0"); + + Output input_float = + ops::Placeholder(root.WithOpName("input_float"), DT_FLOAT); + Output input_i64 = ops::Placeholder(root.WithOpName("input_i64"), DT_INT64); + + Output begin_begin = + ops::Placeholder(root.WithOpName("begin_begin"), DT_INT32); + Output begin_size = ops::Const(root.WithOpName("begin_size"), {-1}); + Output begin = + ops::Slice(root.WithOpName("begin"), input_i64, begin_begin, begin_size); + + Output size = + ops::Const(root.WithOpName("size"), {ToInt64(-1), ToInt64(200)}); + Output slice_with_slice_begin = ops::Slice( + root.WithOpName("slice_with_slice_begin"), input_float, begin, size); + + std::unique_ptr result; + TF_ASSERT_OK(IncreaseDynamismForAutoJit(root, &result)); + + Node* static_shaped_slice = testing::FindNodeByName( + result.get(), + "slice_with_slice_begin/static_shaped_slice/static_shaped_slice"); + ASSERT_NE(static_shaped_slice, nullptr); + EXPECT_EQ(static_shaped_slice->output_type(0), DT_FLOAT) + << "Expected DT_FLOAT, was " + << DataType_Name(static_shaped_slice->output_type(0)); + EXPECT_THAT( + static_shaped_slice, + NodeWith( + Op("Slice"), + Inputs(_, + Out(NodeWith( + Op("Slice"), + Name("begin/static_shaped_slice/static_shaped_slice"))), + _))); +} +} // namespace +} // namespace tensorflow diff --git a/tensorflow/compiler/jit/jit_compilation_pass_registration.cc b/tensorflow/compiler/jit/jit_compilation_pass_registration.cc index 085c0e5adbb270e71ff3447a936555c99904e26c..f79bdc1e2e8d82c9144d1bb9923ad36d8541cbdb 100644 --- a/tensorflow/compiler/jit/jit_compilation_pass_registration.cc +++ b/tensorflow/compiler/jit/jit_compilation_pass_registration.cc @@ -16,6 +16,7 @@ limitations under the License. #include "tensorflow/compiler/jit/build_xla_ops_pass.h" #include "tensorflow/compiler/jit/encapsulate_subgraphs_pass.h" #include "tensorflow/compiler/jit/encapsulate_xla_computations_pass.h" +#include "tensorflow/compiler/jit/increase_dynamism_for_auto_jit_pass.h" #include "tensorflow/compiler/jit/mark_for_compilation_pass.h" #include "tensorflow/compiler/jit/partially_decluster_pass.h" #include "tensorflow/core/common_runtime/optimization_registry.h" @@ -44,17 +45,20 @@ REGISTER_OPTIMIZATION(OptimizationPassRegistry::POST_REWRITE_FOR_EXEC, 10, MarkForCompilationPass); REGISTER_OPTIMIZATION(OptimizationPassRegistry::POST_REWRITE_FOR_EXEC, 20, + IncreaseDynamismForAutoJitPass); + +REGISTER_OPTIMIZATION(OptimizationPassRegistry::POST_REWRITE_FOR_EXEC, 30, PartiallyDeclusterPass); // The EncapsulateSubgraphs pass must run after the MarkForCompilationPass. We // also need to run it after the graph been rewritten to have _Send nodes added // for fetches. Before the _Send nodes are added, fetch nodes are identified by // name, and encapsulation might remove that node from the graph. -REGISTER_OPTIMIZATION(OptimizationPassRegistry::POST_REWRITE_FOR_EXEC, 30, +REGISTER_OPTIMIZATION(OptimizationPassRegistry::POST_REWRITE_FOR_EXEC, 40, EncapsulateSubgraphsPass); // Must run after EncapsulateSubgraphsPass. -REGISTER_OPTIMIZATION(OptimizationPassRegistry::POST_REWRITE_FOR_EXEC, 40, +REGISTER_OPTIMIZATION(OptimizationPassRegistry::POST_REWRITE_FOR_EXEC, 50, BuildXlaOpsPass); } // namespace tensorflow diff --git a/tensorflow/compiler/jit/kernels/BUILD b/tensorflow/compiler/jit/kernels/BUILD index 26cb3af9d69ba1877c67853cde28d2477d394efc..0583774714c6db7a2fa515fc8a0d304e1898db97 100644 --- a/tensorflow/compiler/jit/kernels/BUILD +++ b/tensorflow/compiler/jit/kernels/BUILD @@ -12,6 +12,7 @@ cc_library( hdrs = ["xla_ops.h"], deps = [ "//tensorflow/compiler/jit:common", + "//tensorflow/compiler/jit:flags", "//tensorflow/compiler/jit:xla_compilation_cache", "//tensorflow/compiler/jit:xla_device", "//tensorflow/compiler/jit:xla_launch_util", diff --git a/tensorflow/compiler/jit/kernels/xla_ops.cc b/tensorflow/compiler/jit/kernels/xla_ops.cc index 2268d9042860f6556cb69469ee52ad7cbbb81954..ad71df5a694a5f8da94675049df1062a7edb6253 100644 --- a/tensorflow/compiler/jit/kernels/xla_ops.cc +++ b/tensorflow/compiler/jit/kernels/xla_ops.cc @@ -18,6 +18,7 @@ limitations under the License. #include "absl/container/flat_hash_map.h" #include "absl/memory/memory.h" #include "tensorflow/compiler/jit/defs.h" +#include "tensorflow/compiler/jit/flags.h" #include "tensorflow/compiler/tf2xla/shape_util.h" #include "tensorflow/compiler/tf2xla/tf2xla_util.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" @@ -38,12 +39,22 @@ limitations under the License. #include "tensorflow/core/platform/stream_executor_no_cuda.h" #include "tensorflow/core/util/stream_executor_util.h" +// OP_REQUIRES_OK_RETURN is the same as OP_REQUIRES_OK except that +// in error case, it returns RET instead of void. +#define OP_REQUIRES_OK_RETURN(CTX, RET, ...) \ + do { \ + ::tensorflow::Status _s(__VA_ARGS__); \ + if (!TF_PREDICT_TRUE(_s.ok())) { \ + (CTX)->CtxFailureWithWarning(__FILE__, __LINE__, _s); \ + return RET; \ + } \ + } while (0) + namespace tensorflow { namespace { -Status PlatformInfoFromContext(OpKernelConstruction* ctx, - XlaPlatformInfo* result) { +XlaPlatformInfo PlatformInfoFromContext(OpKernelConstruction* ctx) { DeviceType device_type = ctx->device_type(); se::Platform::Id platform_id = nullptr; const XlaDevice::Metadata* xla_device_metadata = nullptr; @@ -75,16 +86,16 @@ Status PlatformInfoFromContext(OpKernelConstruction* ctx, } if (!device_allocator) { - TF_ASSIGN_OR_RETURN(se::Platform* const platform, - se::MultiPlatformManager::PlatformWithId(platform_id)); + xla::StatusOr maybe_platform = + se::MultiPlatformManager::PlatformWithId(platform_id); + OP_REQUIRES_OK_RETURN(ctx, XlaPlatformInfo(), maybe_platform.status()); + xla_allocator = absl::make_unique( - platform, ctx->device()->GetAllocator({})); + maybe_platform.ValueOrDie(), ctx->device()->GetAllocator({})); } - *result = XlaPlatformInfo(device_type, platform_id, xla_device_metadata, - std::move(xla_allocator), device_allocator); - - return Status::OK(); + return XlaPlatformInfo(device_type, platform_id, xla_device_metadata, + std::move(xla_allocator), device_allocator); } // A closure describing how to run a compiled version of a TensorFlow function. @@ -178,9 +189,8 @@ XlaLocalLaunchBase::XlaLocalLaunchBase(OpKernelConstruction* ctx, : OpKernel(ctx), constants_(constants), resources_(resources), - function_(function) { - OP_REQUIRES_OK(ctx, PlatformInfoFromContext(ctx, &platform_info_)); -} + function_(function), + platform_info_(PlatformInfoFromContext(ctx)) {} static Status BuildCompilationCache(OpKernelContext* ctx, const XlaPlatformInfo& platform_info, @@ -241,7 +251,7 @@ static Status CompileToLocalExecutable( // this is more obviously correct.) core::ScopedUnref cache_ref(cache); - *variables = SnapshotResourceVariables(ctx, resources); + TF_RETURN_IF_ERROR(SnapshotResourceVariables(ctx, resources, variables)); *client = static_cast(cache->client()); XlaCompiler::Options options; @@ -276,8 +286,10 @@ static Status CompileToLocalExecutable( // rather than a one-element tuple. compile_options.always_return_tuple = false; - return cache->Compile(options, function, constant_args, *variables, ctx, - compile_options, + std::vector args; + TF_RETURN_IF_ERROR(XlaComputationLaunchContext::BuildXlaCompilerArguments( + constant_args, *variables, ctx, &args)); + return cache->Compile(options, function, args, compile_options, lazy ? XlaCompilationCache::CompileMode::kLazy : XlaCompilationCache::CompileMode::kStrict, kernel, executable); @@ -332,18 +344,6 @@ void XlaLocalLaunchBase::Compute(OpKernelContext* ctx) { } namespace { - -// OP_REQUIRES_OK_RETURN is the same as OP_REQUIRES_OK except that -// in error case, it returns RET instead of void. -#define OP_REQUIRES_OK_RETURN(CTX, RET, ...) \ - do { \ - ::tensorflow::Status _s(__VA_ARGS__); \ - if (!TF_PREDICT_TRUE(_s.ok())) { \ - (CTX)->CtxFailureWithWarning(__FILE__, __LINE__, _s); \ - return RET; \ - } \ - } while (0) - // Helper static functions to construct parameters for // XlaLocalLaunchBase constructor from OpKernelConstruction. std::vector ConstantsVector(OpKernelConstruction* ctx) { @@ -380,7 +380,12 @@ NameAttrList FunctionAttr(OpKernelConstruction* ctx) { return *func; } -#undef OP_REQUIRES_OK_RETURN +bool MustCompileAttr(OpKernelConstruction* ctx) { + bool must_compile; + OP_REQUIRES_OK_RETURN(ctx, false, + ctx->GetAttr("must_compile", &must_compile)); + return must_compile; +} } // namespace XlaLocalLaunchOp::XlaLocalLaunchOp(OpKernelConstruction* ctx) @@ -395,10 +400,9 @@ XlaCompileOp::XlaCompileOp(OpKernelConstruction* ctx) : OpKernel(ctx), constants_(ConstantsVector(ctx)), resources_(ResourcesVector(ctx)), - function_(FunctionAttr(ctx)) { - OP_REQUIRES_OK(ctx, PlatformInfoFromContext(ctx, &platform_info_)); - OP_REQUIRES_OK(ctx, ctx->GetAttr("must_compile", &must_compile_)); -} + function_(FunctionAttr(ctx)), + platform_info_(PlatformInfoFromContext(ctx)), + must_compile_(MustCompileAttr(ctx)) {} void XlaCompileOp::Compute(OpKernelContext* ctx) { VLOG(3) << "XlaCompileOp " << def().name() @@ -408,10 +412,31 @@ void XlaCompileOp::Compute(OpKernelContext* ctx) { xla::LocalExecutable* executable; std::map variables; - OP_REQUIRES_OK( - ctx, CompileToLocalExecutable(ctx, function_, platform_info_, resources_, - constants_, /*lazy=*/!must_compile_, - &client, &variables, &kernel, &executable)); + bool cannot_compile_cluster; + { + mutex_lock guard(cannot_compile_cluster_mu_); + cannot_compile_cluster = cannot_compile_cluster_; + } + + if (GetXlaOpsCommonFlags().tf_xla_always_defer_compilation || + cannot_compile_cluster) { + executable = nullptr; + } else { + Status status = CompileToLocalExecutable( + ctx, function_, platform_info_, resources_, constants_, + /*lazy=*/!must_compile_, &client, &variables, &kernel, &executable); + if (must_compile_ || status.code() != error::UNIMPLEMENTED) { + OP_REQUIRES_OK(ctx, status); + } + + if (status.code() == error::UNIMPLEMENTED) { + LOG(WARNING) << "Compilation failed:" << status.ToString() + << ". Falling back to TF function call."; + executable = nullptr; + mutex_lock guard(cannot_compile_cluster_mu_); + cannot_compile_cluster_ = true; + } + } AllocatorAttributes host_alloc_attrs; host_alloc_attrs.set_gpu_compatible(true); @@ -447,9 +472,8 @@ void XlaCompileOp::Compute(OpKernelContext* ctx) { ctx->set_output(1, compilation_successful); } -XlaRunOp::XlaRunOp(OpKernelConstruction* ctx) : OpKernel(ctx) { - OP_REQUIRES_OK(ctx, PlatformInfoFromContext(ctx, &platform_info_)); -} +XlaRunOp::XlaRunOp(OpKernelConstruction* ctx) + : OpKernel(ctx), platform_info_(PlatformInfoFromContext(ctx)) {} void XlaRunOp::Compute(OpKernelContext* ctx) { VLOG(3) << "XlaRunOp " << def().name(); diff --git a/tensorflow/compiler/jit/kernels/xla_ops.h b/tensorflow/compiler/jit/kernels/xla_ops.h index ac90837e0d90943b93e2cdb01a30fa0837ba94df..7b4d4b5b4737784d4fe277d5bbe9cab79cfaf4c9 100644 --- a/tensorflow/compiler/jit/kernels/xla_ops.h +++ b/tensorflow/compiler/jit/kernels/xla_ops.h @@ -16,6 +16,8 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_JIT_KERNELS_XLA_OPS_H_ #define TENSORFLOW_COMPILER_JIT_KERNELS_XLA_OPS_H_ +#include + #include "tensorflow/compiler/jit/xla_compilation_cache.h" #include "tensorflow/compiler/jit/xla_device.h" #include "tensorflow/compiler/jit/xla_launch_util.h" @@ -33,6 +35,7 @@ namespace tensorflow { class XlaPlatformInfo { public: XlaPlatformInfo() : device_type_("") {} + XlaPlatformInfo(XlaPlatformInfo&&) = default; explicit XlaPlatformInfo(const DeviceType device_type, se::Platform::Id platform_id, const XlaDevice::Metadata* xla_device_metadata, @@ -110,12 +113,12 @@ class XlaLocalLaunchBase : public OpKernel { protected: // Indexes of compile-time constant inputs - std::vector constants_; + const std::vector constants_; // Indexes of resource inputs - std::vector resources_; + const std::vector resources_; - NameAttrList function_; - XlaPlatformInfo platform_info_; + const NameAttrList function_; + const XlaPlatformInfo platform_info_; }; // XlaLocalLaunchOp is used to replace a region of the TensorFlow graph @@ -144,15 +147,23 @@ class XlaCompileOp : public OpKernel { private: // Indexes of compile-time constant inputs - std::vector constants_; + const std::vector constants_; // Indexes of resource inputs - std::vector resources_; + const std::vector resources_; - NameAttrList function_; + const NameAttrList function_; XlaPlatformInfo platform_info_; - bool must_compile_; + const bool must_compile_; + + // cannot_compile_cluster_ is set to true if XLA returns an Unimplemented + // error when compiling the cluster this _XlaCompile is supposed to compile. + // If `cannot_compile_cluster_` is true then we avoid compiling this cluster + // on any future calls to _XlaCompile. + bool cannot_compile_cluster_ GUARDED_BY(cannot_compile_cluster_mu_) = false; + + mutex cannot_compile_cluster_mu_; }; class XlaRunOp : public OpKernel { @@ -162,7 +173,7 @@ class XlaRunOp : public OpKernel { void Compute(OpKernelContext* ctx) override; private: - XlaPlatformInfo platform_info_; + const XlaPlatformInfo platform_info_; }; } // namespace tensorflow diff --git a/tensorflow/compiler/jit/legacy_flags/BUILD b/tensorflow/compiler/jit/legacy_flags/BUILD deleted file mode 100644 index d8fe4026f51d8aa4b027aeedf0795ad30e28d986..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/jit/legacy_flags/BUILD +++ /dev/null @@ -1,53 +0,0 @@ -# Legacy command line flags for the XLA bridge libraries. - -# Please do not add more flags to this package. - -# The XLA bridge libraries were written in an environment that allowed -# command-line flags to be scattered freely throughout the libraries. This -# model, while initially convenient, leads to a proliferation in unused command -# line flags in tests and binaries, and serious problems in servers, where one -# might wish parameters to be different in independent RPC calls to the same -# routine. -# -# Please don't add more flags. If you're a library author, pass options and -# parameters explicitly through the library's interface. - -licenses(["notice"]) # Apache 2.0 - -package(default_visibility = ["//tensorflow:internal"]) - -cc_library( - name = "mark_for_compilation_pass_flags", - srcs = ["mark_for_compilation_pass_flags.cc"], - hdrs = ["mark_for_compilation_pass_flags.h"], - deps = - [ - "//tensorflow/compiler/xla/legacy_flags:parse_flags_from_env", - "//tensorflow/core:framework_internal", - "//tensorflow/core:lib", - ], -) - -cc_library( - name = "xla_device_flags", - srcs = ["xla_device_flags.cc"], - hdrs = ["xla_device_flags.h"], - deps = - [ - "//tensorflow/compiler/xla/legacy_flags:parse_flags_from_env", - "//tensorflow/core:framework_internal", - "//tensorflow/core:lib", - ], -) - -cc_library( - name = "build_xla_ops_pass_flags", - srcs = ["build_xla_ops_pass_flags.cc"], - hdrs = ["build_xla_ops_pass_flags.h"], - deps = - [ - "//tensorflow/compiler/xla/legacy_flags:parse_flags_from_env", - "//tensorflow/core:framework_internal", - "//tensorflow/core:lib", - ], -) diff --git a/tensorflow/compiler/jit/legacy_flags/build_xla_ops_pass_flags.cc b/tensorflow/compiler/jit/legacy_flags/build_xla_ops_pass_flags.cc deleted file mode 100644 index 58157d2b9800a2e8269533607c2ea688ff4e7766..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/jit/legacy_flags/build_xla_ops_pass_flags.cc +++ /dev/null @@ -1,47 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include // NOLINT - -#include "tensorflow/compiler/jit/legacy_flags/build_xla_ops_pass_flags.h" -#include "tensorflow/compiler/xla/legacy_flags/parse_flags_from_env.h" -#include "tensorflow/core/util/command_line_flags.h" - -namespace tensorflow { -namespace legacy_flags { -namespace { - -BuildXlaOpsPassFlags* flags; -std::vector* flag_list; -std::once_flag flags_init; - -void AllocateAndParseFlags() { - flags = new BuildXlaOpsPassFlags; - flags->tf_xla_enable_lazy_compilation = false; - flag_list = new std::vector({ - Flag("tf_xla_enable_lazy_compilation", - &flags->tf_xla_enable_lazy_compilation, ""), - }); - xla::legacy_flags::ParseFlagsFromEnv(*flag_list); -} - -} // namespace - -const BuildXlaOpsPassFlags& GetBuildXlaOpsPassFlags() { - std::call_once(flags_init, &AllocateAndParseFlags); - return *flags; -} -} // namespace legacy_flags -} // namespace tensorflow diff --git a/tensorflow/compiler/jit/legacy_flags/build_xla_ops_pass_flags.h b/tensorflow/compiler/jit/legacy_flags/build_xla_ops_pass_flags.h deleted file mode 100644 index 539314cbf72d38ed973b8a526aa6424b19ef344d..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/jit/legacy_flags/build_xla_ops_pass_flags.h +++ /dev/null @@ -1,37 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#ifndef TENSORFLOW_COMPILER_JIT_LEGACY_FLAGS_BUILD_XLA_OPS_PASS_FLAGS_H_ -#define TENSORFLOW_COMPILER_JIT_LEGACY_FLAGS_BUILD_XLA_OPS_PASS_FLAGS_H_ - -namespace tensorflow { -namespace legacy_flags { - -// Flags for the build_xla_ops pass. -struct BuildXlaOpsPassFlags { - // Enables lazy compilation for TF/XLA (only when auto-clustering) if true. - // Defaults to false. - bool tf_xla_enable_lazy_compilation; -}; - -// Parses the flags in BuildXlaOpsPassFlags from the TF_XLA_FLAGS environment -// variable and returns a reference to the parsed copy. Parses TF_XLA_FLAGS -// only the first time this routine is called. -const BuildXlaOpsPassFlags& GetBuildXlaOpsPassFlags(); - -} // namespace legacy_flags -} // namespace tensorflow - -#endif // TENSORFLOW_COMPILER_JIT_LEGACY_FLAGS_BUILD_XLA_OPS_PASS_FLAGS_H_ diff --git a/tensorflow/compiler/jit/legacy_flags/mark_for_compilation_pass_flags.cc b/tensorflow/compiler/jit/legacy_flags/mark_for_compilation_pass_flags.cc deleted file mode 100644 index 7277a1d1f8ad5fa045645ead839ab9efa01e89c7..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/jit/legacy_flags/mark_for_compilation_pass_flags.cc +++ /dev/null @@ -1,86 +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. -==============================================================================*/ - -// Legacy flags for the XLA bridge's mark_for_compilation_pass module. - -#include -#include - -#include "tensorflow/compiler/jit/legacy_flags/mark_for_compilation_pass_flags.h" -#include "tensorflow/compiler/xla/legacy_flags/parse_flags_from_env.h" -#include "tensorflow/core/platform/types.h" -#include "tensorflow/core/util/command_line_flags.h" - -namespace tensorflow { -namespace legacy_flags { - -// Pointers to the parsed value of the flags and flag descriptors, initialized -// via flags_init. -static MarkForCompilationPassFlags* flags; -static std::vector* flag_list; -static std::once_flag flags_init; - -// Allocate *flags. Called via call_once(&flags_init,...). -static void AllocateFlags() { - flags = new MarkForCompilationPassFlags; - flags->tf_xla_auto_jit = 0; - flags->tf_xla_min_cluster_size = 2; - flags->tf_xla_max_cluster_size = std::numeric_limits::max(); - flags->tf_xla_clustering_debug = false; - flags->tf_xla_cpu_global_jit = false; - flags->tf_xla_clustering_fuel = std::numeric_limits::max(); - flags->tf_xla_fusion_only = false; - flag_list = new std::vector( - {Flag("tf_xla_auto_jit", &flags->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."), - Flag("tf_xla_min_cluster_size", &flags->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."), - Flag("tf_xla_max_cluster_size", &flags->tf_xla_max_cluster_size, - "Maximum number of operators in an XLA compilation."), - Flag("tf_xla_clustering_debug", &flags->tf_xla_clustering_debug, - "Dump graphs during XLA compilation."), - Flag("tf_xla_cpu_global_jit", &flags->tf_xla_cpu_global_jit, - "Enables global JIT compilation for CPU via SessionOptions."), - Flag("tf_xla_clustering_fuel", &flags->tf_xla_clustering_fuel, - "Places an artificial limit on the number of ops marked as " - "eligible for clustering."), - Flag("tf_xla_fusion_only", &flags->tf_xla_fusion_only, - "enable fusion of element-wise operations only using XLA when " - "global_jit_level is ON*.")}); - xla::legacy_flags::ParseFlagsFromEnv(*flag_list); -} - -// Append to *append_to flag definitions associated with the XLA bridge's -// mark_for_compilation_pass module. -void AppendMarkForCompilationPassFlags(std::vector* append_to) { - std::call_once(flags_init, &AllocateFlags); - append_to->insert(append_to->end(), flag_list->begin(), flag_list->end()); -} - -// Return a pointer to the MarkForCompilationPassFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -MarkForCompilationPassFlags* GetMarkForCompilationPassFlags() { - std::call_once(flags_init, &AllocateFlags); - return flags; -} - -} // namespace legacy_flags -} // namespace tensorflow diff --git a/tensorflow/compiler/jit/legacy_flags/mark_for_compilation_pass_flags.h b/tensorflow/compiler/jit/legacy_flags/mark_for_compilation_pass_flags.h deleted file mode 100644 index 2affda6ab4e0fbad32a246744fa5b38aeb629c1b..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/jit/legacy_flags/mark_for_compilation_pass_flags.h +++ /dev/null @@ -1,68 +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_JIT_LEGACY_FLAGS_MARK_FOR_COMPILATION_PASS_FLAGS_H_ -#define TENSORFLOW_COMPILER_JIT_LEGACY_FLAGS_MARK_FOR_COMPILATION_PASS_FLAGS_H_ - -// Legacy flags for the XLA bridge's mark_for_compilation_pass module. - -#include - -#include "tensorflow/core/platform/types.h" -#include "tensorflow/core/util/command_line_flags.h" - -namespace tensorflow { -namespace legacy_flags { - -// Append to *flag_list flag definitions associated with the XLA bridge's -// mark_for_compilation_pass module. -void AppendMarkForCompilationPassFlags( - std::vector* flag_list); - -// The values of flags associated with the XLA bridge's -// mark_for_compilation_pass module. -typedef struct { - 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. -} MarkForCompilationPassFlags; - -// Return a pointer to the MarkForCompilationPassFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -MarkForCompilationPassFlags* GetMarkForCompilationPassFlags(); - -} // namespace legacy_flags -} // namespace tensorflow - -#endif // TENSORFLOW_COMPILER_JIT_LEGACY_FLAGS_MARK_FOR_COMPILATION_PASS_FLAGS_H_ diff --git a/tensorflow/compiler/jit/legacy_flags/xla_device_flags.cc b/tensorflow/compiler/jit/legacy_flags/xla_device_flags.cc deleted file mode 100644 index 1bb2fce2dbad5bffce2e33b665b7222090d0855a..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/jit/legacy_flags/xla_device_flags.cc +++ /dev/null @@ -1,56 +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. -==============================================================================*/ - -// Legacy flags for the XLA bridge's xla_device module. - -#include -#include - -#include "tensorflow/compiler/jit/legacy_flags/xla_device_flags.h" -#include "tensorflow/compiler/xla/legacy_flags/parse_flags_from_env.h" -#include "tensorflow/core/platform/types.h" -#include "tensorflow/core/util/command_line_flags.h" - -namespace tensorflow { -namespace legacy_flags { - -// Pointers to the parsed value of the flags and flag descriptors, initialized -// via flags_init. -static XlaDeviceFlags* flags; -static std::vector* flag_list; -static std::once_flag flags_init; - -// Allocate *flags. Called via call_once(&flags_init,...). -static void AllocateFlags() { - flags = new XlaDeviceFlags; - flags->tf_xla_compile_on_demand = false; - flag_list = new std::vector({ - Flag("tf_xla_compile_on_demand", &flags->tf_xla_compile_on_demand, - "Switch a device into 'on-demand' mode, where instead of " - "autoclustering ops are compiled one by one just-in-time."), - }); - xla::legacy_flags::ParseFlagsFromEnv(*flag_list); -} - -// Return a pointer to the XlaDeviceFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -XlaDeviceFlags* GetXlaDeviceFlags() { - std::call_once(flags_init, &AllocateFlags); - return flags; -} - -} // namespace legacy_flags -} // namespace tensorflow diff --git a/tensorflow/compiler/jit/legacy_flags/xla_device_flags.h b/tensorflow/compiler/jit/legacy_flags/xla_device_flags.h deleted file mode 100644 index 27b22121ac1e089bd5d5a494e1e3fb60b05bc76d..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/jit/legacy_flags/xla_device_flags.h +++ /dev/null @@ -1,47 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#ifndef TENSORFLOW_COMPILER_JIT_LEGACY_FLAGS_XLA_DEVICE_FLAGS_H_ -#define TENSORFLOW_COMPILER_JIT_LEGACY_FLAGS_XLA_DEVICE_FLAGS_H_ - -// Legacy flags for the XLA bridge's xla_device module. - -#include - -#include "tensorflow/core/platform/types.h" -#include "tensorflow/core/util/command_line_flags.h" - -namespace tensorflow { -namespace legacy_flags { - -// The values of flags associated with the XLA bridge's -// xla_device module. -typedef struct { - // Switch the CPU device into "on-demand" mode, where instead of - // autoclustering ops are compiled one by one just-in-time. - // Enabling this mode by a legacy flag is a temporary mechanism. When this - // feature is battle-tested, we will switch this to be a session option. - bool tf_xla_compile_on_demand; -} XlaDeviceFlags; - -// Return a pointer to the XlaDeviceFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -XlaDeviceFlags* GetXlaDeviceFlags(); - -} // namespace legacy_flags -} // namespace tensorflow - -#endif // TENSORFLOW_COMPILER_JIT_LEGACY_FLAGS_XLA_DEVICE_FLAGS_H_ diff --git a/tensorflow/compiler/jit/mark_for_compilation_pass.cc b/tensorflow/compiler/jit/mark_for_compilation_pass.cc index 4f0c370e65159c89c91ea58733f20f852d9acc99..234e8052d16572fef803c7334afb2ab8e689a72c 100644 --- a/tensorflow/compiler/jit/mark_for_compilation_pass.cc +++ b/tensorflow/compiler/jit/mark_for_compilation_pass.cc @@ -24,8 +24,8 @@ limitations under the License. #include "absl/container/flat_hash_set.h" #include "tensorflow/compiler/jit/deadness_analysis.h" #include "tensorflow/compiler/jit/defs.h" +#include "tensorflow/compiler/jit/flags.h" #include "tensorflow/compiler/jit/graphcycles/graphcycles.h" -#include "tensorflow/compiler/jit/legacy_flags/mark_for_compilation_pass_flags.h" #include "tensorflow/compiler/jit/union_find.h" #include "tensorflow/compiler/jit/xla_cluster_util.h" #include "tensorflow/compiler/tf2xla/const_analysis.h" @@ -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" @@ -49,6 +50,51 @@ limitations under the License. namespace tensorflow { namespace { +// Aggregates information about what kinds of ops are allowed. +struct OperationFilter { + // Whether resource variable ops are allowed. We do not allow resource + // variable ops in called functions (either as direct TF calls or as higher + // order control flow ops) because we do not yet model their memory effects in + // jit/resource_variable_safety_analysis. + bool allow_resource_ops; + + // Whether stateful RNG ops are allowed. XLA's RNG does not have the same + // seeding behavior as TensorFlow's RNG (b/34749654). So we avoid + // auto-clustering stateful RNG ops. + bool allow_stateful_rng_ops; + + // TODO(b/118970344): Whether ControlTrigger ops are allowed. It is unsound + // to cluster ControlTrigger because of how we use deadness analysis. + bool allow_control_trigger; + + // Whether ops with dummy implementations are allowed. We avoid + // auto-clustering these ops so that the user is not surprised when XLA is + // implicitly enabled. If the user explicitly specifies to use XLA, it is fine + // to resort to a dummy implementation. Currently Assert and CheckNumerics ops + // have dummy XLA implementations. + bool allow_dummy_ops; + + // Whether ops that produce or consume DT_VARIANT values are allowed. We + // don't auto-cluster these ops because we don't yet support live-in or + // live-out DT_VARIANT values. + bool allow_ops_producing_or_consuming_variant; +}; + +bool IsDummyImplOp(absl::string_view op_name) { + return op_name == "Assert" || op_name == "CheckNumerics"; +} + +bool IsStatefulRandomOp(absl::string_view op_name) { + return op_name == "RandomUniform" || op_name == "RandomShuffle" || + op_name == "RandomUniformInt" || op_name == "RandomStandardNormal" || + op_name == "TruncatedNormal" || op_name == "Multinomial"; +} + +bool OpProducesOrConsumesVariant(const Node& node) { + auto is_variant = [](DataType dtype) { return dtype == DT_VARIANT; }; + return absl::c_any_of(node.input_types(), is_variant) || + absl::c_any_of(node.output_types(), is_variant); +} bool HasXLAKernel(const Node& node, const DeviceType& jit_device_type) { // There is a SymbolicGradient kernel on the XLA_JIT device, but the gradient @@ -101,7 +147,7 @@ const int kMaxRecursionDepth = 10; bool IsCompilableCall(const NodeDef& call_def, const DeviceType& jit_device_type, - bool allow_resource_ops, int depth, + const OperationFilter& op_filter, int depth, FunctionLibraryRuntime* lib_runtime); // Tests whether 'while_node' is a completely compilable loop. @@ -109,7 +155,7 @@ bool IsCompilableCall(const NodeDef& call_def, // while loop to be compilable. bool IsCompilableWhile(const Node& while_node, const DeviceType& jit_device_type, - bool allow_resource_ops, int depth, + const OperationFilter& op_filter, int depth, FunctionLibraryRuntime* lib_runtime) { const NameAttrList* name_attr; NodeDef call; @@ -124,7 +170,7 @@ bool IsCompilableWhile(const Node& while_node, call.set_name("while_cond"); call.set_op(cond_func); *call.mutable_attr() = name_attr->attr(); - if (!IsCompilableCall(call, jit_device_type, allow_resource_ops, depth + 1, + if (!IsCompilableCall(call, jit_device_type, op_filter, depth + 1, lib_runtime)) { VLOG(2) << "Rejecting While " << while_node.name() << ": can't compile loop condition: " << cond_func; @@ -140,7 +186,7 @@ bool IsCompilableWhile(const Node& while_node, call.set_name("while_body"); call.set_op(body_func); *call.mutable_attr() = name_attr->attr(); - if (!IsCompilableCall(call, jit_device_type, allow_resource_ops, depth + 1, + if (!IsCompilableCall(call, jit_device_type, op_filter, depth + 1, lib_runtime)) { VLOG(2) << "Rejecting While " << while_node.name() << ": can't compile loop body: " << body_func; @@ -154,7 +200,7 @@ bool IsCompilableWhile(const Node& while_node, // compilable. bool IsCompilableCall(const NodeDef& call_def, const DeviceType& jit_device_type, - bool allow_resource_ops, int depth, + const OperationFilter& op_filter, int depth, FunctionLibraryRuntime* lib_runtime) { if (depth > kMaxRecursionDepth) { VLOG(2) << "Rejecting " << call_def.op() @@ -195,16 +241,30 @@ bool IsCompilableCall(const NodeDef& call_def, continue; if (node->type_string() == "While") { // Handle functional While loop. - return IsCompilableWhile(*node, jit_device_type, allow_resource_ops, - depth + 1, lib_runtime); + return IsCompilableWhile(*node, jit_device_type, op_filter, depth + 1, + lib_runtime); } - if (!allow_resource_ops && + if (!op_filter.allow_resource_ops && (HasResourceInput(*node) || HasResourceOutput(*node))) { return false; } + if (!op_filter.allow_stateful_rng_ops && + IsStatefulRandomOp(node->type_string())) { + return false; + } + if (!op_filter.allow_control_trigger && node->IsControlTrigger()) { + return false; + } + if (!op_filter.allow_dummy_ops && IsDummyImplOp(node->type_string())) { + return false; + } + if (!op_filter.allow_ops_producing_or_consuming_variant && + OpProducesOrConsumesVariant(*node)) { + return false; + } if (!HasXLAKernel(*node, jit_device_type) && - !IsCompilableCall(node->def(), jit_device_type, allow_resource_ops, - depth + 1, lib_runtime)) { + !IsCompilableCall(node->def(), jit_device_type, op_filter, depth + 1, + lib_runtime)) { VLOG(2) << "Rejecting " << call_def.op() << ": unsupported op " << node->name() << ": " << node->def().ShortDebugString(); return false; @@ -383,8 +443,7 @@ Status FindCompilationCandidates( BackwardsConstAnalysis(graph, /*compile_time_const_arg_indices=*/nullptr, &compile_time_const_nodes)); - int64& fuel = - legacy_flags::GetMarkForCompilationPassFlags()->tf_xla_clustering_fuel; + int64& fuel = GetMarkForCompilationPassFlags()->tf_xla_clustering_fuel; // Iterate over nodes in sorted order so that compiler fuel is deterministic. // We can't simply pass op_nodes().begin() and op_nodes().end to the @@ -426,14 +485,47 @@ Status FindCompilationCandidates( CHECK( XlaOpRegistry::GetCompilationDevice(device_type.type(), ®istration)); DeviceType jit_device_type(registration->compilation_device_name); + + bool always_auto_cluster = registration->autoclustering_policy == + XlaOpRegistry::AutoclusteringPolicy::kAlways; + + OperationFilter op_filter; + op_filter.allow_resource_ops = registration->compile_resource_ops; + op_filter.allow_stateful_rng_ops = always_auto_cluster; + op_filter.allow_control_trigger = always_auto_cluster; + op_filter.allow_dummy_ops = always_auto_cluster; + op_filter.allow_ops_producing_or_consuming_variant = always_auto_cluster; + if (!HasXLAKernel(*node, jit_device_type) && - !IsCompilableCall(node->def(), jit_device_type, - registration->compile_resource_ops, 0, lib_runtime)) { + !IsCompilableCall(node->def(), jit_device_type, op_filter, 0, + lib_runtime)) { VLOG(2) << "Rejecting " << node->name() << ": unsupported op " << node->type_string(); continue; } - if (!registration->compile_resource_ops && + + if (!op_filter.allow_stateful_rng_ops && + IsStatefulRandomOp(node->type_string())) { + VLOG(2) << "Rejecting " << node->name() << ": stateful random operation"; + continue; + } + if (!op_filter.allow_control_trigger && node->IsControlTrigger()) { + VLOG(2) << "Rejecting " << node->name() << ": is a control trigger op"; + continue; + } + if (!op_filter.allow_dummy_ops && IsDummyImplOp(node->type_string())) { + VLOG(2) << "Rejecting " << node->name() << ": dummy op (" + << node->type_string() << ")"; + continue; + } + if (!op_filter.allow_ops_producing_or_consuming_variant && + OpProducesOrConsumesVariant(*node)) { + VLOG(2) << "Rejecting " << node->name() + << ": produces or consumes DT_VARIANT"; + continue; + } + + if (!op_filter.allow_resource_ops && (HasResourceOutput(*node) || IsNonResourceVarResourceOp(*node))) { // We don't have a way of returning values of type DT_RESOURCE from XLA // computations so we avoid auto-clustering nodes producing DT_RESOURCE. @@ -444,6 +536,7 @@ Status FindCompilationCandidates( << node->type_string(); continue; } + if (compile_time_const_nodes[node->id()]) { const OpDef* op_def; TF_RETURN_IF_ERROR( @@ -501,9 +594,7 @@ Status FindCompilationCandidates( // registration->compile_resource_ops is true for XLA_CPU/XLA_GPU but not // for CPU/GPU. if (node->type_string() == "While" && - !IsCompilableWhile(*node, jit_device_type, - registration->compile_resource_ops, 0, - lib_runtime)) { + !IsCompilableWhile(*node, jit_device_type, op_filter, 0, lib_runtime)) { continue; } // _Arg nodes in a top-level function represent feeds. @@ -536,8 +627,7 @@ OptimizerOptions::GlobalJitLevel GetGlobalJitLevel( // To set compilation to be on by default, change the following line. global_jit_level = OptimizerOptions::OFF; } - legacy_flags::MarkForCompilationPassFlags* flags = - legacy_flags::GetMarkForCompilationPassFlags(); + MarkForCompilationPassFlags* flags = GetMarkForCompilationPassFlags(); if (flags->tf_xla_auto_jit == -1 || (1 <= flags->tf_xla_auto_jit && flags->tf_xla_auto_jit <= 2)) { // If the flag tf_xla_auto_jit is a valid, non-zero setting, it overrides @@ -563,10 +653,16 @@ bool IsCompilable(FunctionLibraryRuntime* flr, const NodeDef& ndef) { ®istration)); DeviceType jit_device_type(registration->compilation_device_name); - // We can always *compile* resource operations, even if we are sometimes - // unable to auto-cluster them. - const bool compile_resource_ops = true; - return IsCompilableCall(ndef, jit_device_type, compile_resource_ops, 0, flr); + // We can always *compile* resource operations, stateful RNGs and dummy ops, + // even if we are sometimes unable to auto-cluster them. + OperationFilter op_filter; + op_filter.allow_resource_ops = true; + op_filter.allow_stateful_rng_ops = true; + op_filter.allow_control_trigger = true; + op_filter.allow_dummy_ops = true; + op_filter.allow_ops_producing_or_consuming_variant = true; + + return IsCompilableCall(ndef, jit_device_type, op_filter, 0, flr); } Status MarkForCompilationPass::Run( @@ -575,16 +671,18 @@ Status MarkForCompilationPass::Run( // device ahead of time. OptimizerOptions::GlobalJitLevel global_jit_level = GetGlobalJitLevel(options); - legacy_flags::MarkForCompilationPassFlags* flags = - legacy_flags::GetMarkForCompilationPassFlags(); - bool cpu_global_jit = flags->tf_xla_cpu_global_jit; + MarkForCompilationPassFlags* flags = GetMarkForCompilationPassFlags(); bool fusion_only = flags->tf_xla_fusion_only; - VLOG(1) << "flags->tf_xla_cpu_global_jit = " << flags->tf_xla_cpu_global_jit; VLOG(1) << "flags->tf_xla_fusion_only = " << flags->tf_xla_fusion_only; 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); @@ -599,9 +697,6 @@ Status MarkForCompilationPass::Run( return false; } - // If this device requires a JIT, we must say yes. - if (registration->requires_compilation) return true; - // If there is a _XlaCompile annotation, use its value. bool compile = false; Status status = GetNodeAttr(node->attrs(), kXlaCompileAttr, &compile); @@ -638,18 +733,21 @@ Status MarkForCompilationPass::Run( return false; } - // Otherwise use the value of global_jit_level. - // Ignore enable_jit_by_default if global jit compilation for CPU - // is explicitly requested via tf_xla_cpu_global_jit flag - bool ignore_registration = cpu_global_jit && device_type == DEVICE_CPU; + // Otherwise use the value of global_jit_level and the device's + // autoclustering policy. bool should_compile = - (ignore_registration || registration->enable_jit_by_default) && - global_jit_level != OptimizerOptions::OFF; + registration->autoclustering_policy == + XlaOpRegistry::AutoclusteringPolicy::kAlways || + (registration->autoclustering_policy == + XlaOpRegistry::AutoclusteringPolicy::kIfEnabledGlobally && + global_jit_level != OptimizerOptions::OFF); if (!should_compile) { if (global_jit_level == OptimizerOptions::OFF) { VLOG(2) << "Rejecting " << node->name() << ": global jit disabled."; } else { - VLOG(2) << "Rejecting " << node->name() << ": JIT for device disabled."; + VLOG(2) + << "Rejecting " << node->name() + << ": autoclustering for device only when requested explicitly."; } } return should_compile; @@ -879,8 +977,7 @@ Status MarkForCompilationPass::RunImpl( OptimizerOptions::GlobalJitLevel global_jit_level = GetGlobalJitLevel(options); - legacy_flags::MarkForCompilationPassFlags* flags = - legacy_flags::GetMarkForCompilationPassFlags(); + MarkForCompilationPassFlags* flags = GetMarkForCompilationPassFlags(); // Repeatedly contract edges between clusters that are on the same device, // provided the contraction would not create a cycle. @@ -952,6 +1049,28 @@ Status MarkForCompilationPass::RunImpl( continue; } + // If any of the consumer's producers are on a different device, do not + // cluster these nodes. This prevents other work on this device from being + // delayed by work on other devices. We consider predecessors of the + // entire cluster rather than just the inputs to the node to prevent the + // cluster still being combined in cases where the 'to' cluster has + // multiple dependencies on the 'from' cluster and another dependency + // leads to a merging of the clusters. + // + // TODO(b/117085735): We probably want to handle the reciprocal of this + // case where a cluster is producing data for multiple devices. + bool found_split = false; + for (const auto& in_id : cycles.Predecessors(to)) { + if (in_id >= graph->num_node_ids()) continue; + + Node* in = graph->FindNodeId(in_id); + if (compilation_candidates.find(in) != compilation_candidates.cend() && + in->assigned_device_name() != node_to->assigned_device_name()) { + found_split = true; + } + } + if (found_split) continue; + // If contracting the edge would create a cycle, bail out. // However, just because we can't merge the clusters now does not mean // we won't be able to merge them in the future. @@ -1015,12 +1134,10 @@ Status MarkForCompilationPass::RunImpl( XlaOpRegistry::GetCompilationDevice(device_type.type(), ®istration); // Compile if this is a cluster of >= min_cluster_size compilable operators. - // Also, always compile if the operator is placed on a device that requires - // compilation, or if it contains at least one op that is marked for + // Also, always compile if it contains at least one op that is marked for // compilation that is not an Identity op. if (effective_cluster_sizes[cluster] >= min_cluster_size || - (effective_cluster_sizes[cluster] > 0 && marked_for_compilation) || - registration->requires_compilation) { + (effective_cluster_sizes[cluster] > 0 && marked_for_compilation)) { string& name = cluster_names[cluster]; if (name.empty()) { @@ -1034,6 +1151,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 2a80c745e3fcebf97bcccb03551feb3d6fb9f831..c2b6250f738fafa35b2c5f79e97cf1281b50a316 100644 --- a/tensorflow/compiler/jit/mark_for_compilation_pass_test.cc +++ b/tensorflow/compiler/jit/mark_for_compilation_pass_test.cc @@ -22,6 +22,7 @@ limitations under the License. #include "tensorflow/cc/ops/array_ops.h" #include "tensorflow/cc/ops/control_flow_ops_internal.h" #include "tensorflow/cc/ops/function_ops.h" +#include "tensorflow/cc/ops/list_ops.h" #include "tensorflow/cc/ops/resource_variable_ops.h" #include "tensorflow/cc/ops/sendrecv_ops.h" #include "tensorflow/cc/ops/standard_ops.h" @@ -150,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; { @@ -158,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())); } @@ -817,14 +818,10 @@ TEST(XlaCompilationTest, ClusterControlTrigger) { std::unordered_map clusters = GetClusters(*graph); - ASSERT_FALSE(clusters.empty()); - string cluster_name = clusters.begin()->second; - - // ctrl_trigger_a has inputs with mismatching deadness so it won't be - // clustered. ctrl_trigger_b is okay to cluster. - std::unordered_map expected_clusters( - {{"const_a", cluster_name}, {"ctrl_trigger_b", cluster_name}}); - EXPECT_EQ(clusters, expected_clusters); + // TODO(b/118970344): ctrl_trigger_a has inputs with mismatching deadness so + // it won't be clustered. ctrl_trigger_b is okay to cluster but we don't + // cluster it because of b/118970344. + EXPECT_TRUE(clusters.empty()); } TEST(XlaCompilationTest, RandomShape) { @@ -923,9 +920,8 @@ TEST(XlaCompilationTest, RandomShapeOnXlaDevice) { TF_ASSERT_OK(MarkForCompilationPassTestHelper::MarkForCompilation(&graph)); std::unordered_map clusters = GetClusters(*graph); - EXPECT_NE(clusters["test/shape_rng"], ""); - EXPECT_NE(clusters["test/reshape"], ""); - EXPECT_NE(clusters["test/shape_rng"], clusters["test/reshape"]); + EXPECT_EQ(clusters["test/shape_rng"], ""); + EXPECT_EQ(clusters["test/reshape"], ""); } TEST(XlaCompilationTest, TensorArrayShapeOnXlaDevice) { @@ -961,5 +957,271 @@ TEST(XlaCompilationTest, TensorArrayShapeOnXlaDevice) { EXPECT_EQ(clusters["test/read"], clusters["test/reshape"]); } +TEST(XlaCompilationTest, DontClusterMergingNodes) { + // MatMulCombined below takes data from nodes on GPU0 and GPU1 and is placed + // on GPU1. However, it should not be clustered with the previous node on + // GPU1, because that will serialize production of its inputs that should be + // done in parallel. + // + // This graph is: + // (Const0, Const0) -> MatMul0 + // (Const1, Const1) -> MatMul1 + // (MatMul0, MatMul1) -> MatMulCombined + // + // Device0: [Const0, Const0, MatMul0] + // Device1: [Const1, Const1, MatMul1, MatMulCombined] + // + // Cluster0: [Const0, Const0, MatMul0] + // Cluster1: [Const1, Const1, MatMul1] + // Cluster2: [MatMulCombined] + Scope root = Scope::NewRootScope().ExitOnError(); + absl::string_view xla_gpu_dev0 = + "/job:worker/replica:0/task:0/device:XLA_GPU:0"; + absl::string_view xla_gpu_dev1 = + "/job:worker/replica:0/task:0/device:XLA_GPU:1"; + std::unique_ptr graph(new Graph(OpRegistry::Global())); + Output a = ops::Const(root.WithOpName("A_dev0"), 1.0f, {2, 2}); + Output b = ops::Const(root.WithOpName("B_dev1"), 1.0f, {2, 2}); + Output matmul0 = ops::MatMul(root.WithOpName("MatMul0_dev0"), a, a); + Output matmul1 = ops::MatMul(root.WithOpName("MatMul1_dev1"), b, b); + + Output combined = + ops::MatMul(root.WithOpName("MatMulCombined_dev1"), matmul0, matmul1); + TF_ASSERT_OK(root.ToGraph(graph.get())); + + for (Node* n : graph->nodes()) { + if (absl::EndsWith(n->name(), /*suffix=*/"dev0")) { + n->set_assigned_device_name(string(xla_gpu_dev0)); + } else if (absl::EndsWith(n->name(), /*suffix=*/"dev1")) { + n->set_assigned_device_name(string(xla_gpu_dev1)); + } + } + TF_ASSERT_OK(MarkForCompilationPassTestHelper::MarkForCompilation(&graph)); + + // Each of the MatMuls should be in a separate cluster. + std::unordered_map clusters = GetClusters(*graph); + EXPECT_NE(clusters["MatMul0_dev0"], clusters["MatMul1_dev1"]); + EXPECT_NE(clusters["MatMulCombined_dev1"], clusters["MatMul0_dev0"]); + EXPECT_NE(clusters["MatMulCombined_dev1"], clusters["MatMul1_dev1"]); + EXPECT_EQ(clusters["A_dev0"], clusters["MatMul0_dev0"]); + EXPECT_EQ(clusters["B_dev1"], clusters["MatMul1_dev1"]); +} + +// TODO(b/117085735): This form of clustering should be prevented. +TEST(XlaCompilationTest, NOT_DontClusterSpreadingNodes) { + // MatMulSource below creates data for nodes on GPU0 and GPU1 and is placed + // on GPU0. However, it should not be clustered with the next node on + // GPU0, because that will prevent the node on GPU1 from beginning its work as + // soon as the data has been produced. + // + // This graph is: + // (Const0, Const0) -> MatMulSource + // MatMulSource -> (MatMul0, MatMul1) + // + // Device0: [Const0, Const1, MatMulSource, MatMul0] + // Device1: [MatMul1] + // + // Cluster0: [Const0, Const1, MatMulSource] + // Cluster1: [MatMul0] + // Cluster2: [MatMul1] + Scope root = Scope::NewRootScope().ExitOnError(); + absl::string_view xla_gpu_dev0 = + "/job:worker/replica:0/task:0/device:XLA_GPU:0"; + absl::string_view xla_gpu_dev1 = + "/job:worker/replica:0/task:0/device:XLA_GPU:1"; + std::unique_ptr graph(new Graph(OpRegistry::Global())); + Output a = ops::Const(root.WithOpName("A_dev0"), 1.0f, {2, 2}); + Output matmul_source = + ops::MatMul(root.WithOpName("MatMulSource_dev0"), a, a); + + Output matmul0 = ops::MatMul(root.WithOpName("MatMul0_dev0"), matmul_source, + matmul_source); + Output matmul1 = ops::MatMul(root.WithOpName("MatMul1_dev1"), matmul_source, + matmul_source); + + TF_ASSERT_OK(root.ToGraph(graph.get())); + for (Node* n : graph->nodes()) { + if (absl::EndsWith(n->name(), /*suffix=*/"dev0")) { + n->set_assigned_device_name(string(xla_gpu_dev0)); + } else if (absl::EndsWith(n->name(), /*suffix=*/"dev1")) { + n->set_assigned_device_name(string(xla_gpu_dev1)); + } + } + TF_ASSERT_OK(MarkForCompilationPassTestHelper::MarkForCompilation(&graph)); + + std::unordered_map clusters = GetClusters(*graph); + EXPECT_EQ(clusters["A_dev0"], clusters["MatMulSource_dev0"]); + EXPECT_NE(clusters["MatMul0_dev0"], clusters["MatMul1_dev1"]); + EXPECT_NE(clusters["MatMulSource_dev0"], clusters["MatMul1_dev1"]); + + // Improved Heuristics should prevent this probably. + EXPECT_EQ(clusters["MatMulSource_dev0"], clusters["MatMul0_dev0"]); +} + +TEST(XlaCompilationTest, ClusterStatefulRandomOpOnXlaDevice) { + absl::string_view xla_cpu_device = + "/job:worker/replica:0/task:0/device:XLA_CPU:0"; + + Scope root = Scope::NewRootScope().ExitOnError(); + Output shape = ops::Const(root.WithOpName("test/shape_shape"), {200, 200}); + Output a = ops::RandomUniform(root.WithOpName("test/a"), shape, DT_FLOAT); + Output b = ops::RandomUniform(root.WithOpName("test/b"), shape, DT_FLOAT); + Output c = ops::Add(root.WithOpName("test/c"), a, b); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(root.ToGraph(graph.get())); + + for (Node* n : graph->nodes()) { + if (absl::StartsWith(n->name(), /*prefix=*/"test/")) { + n->set_assigned_device_name(string(xla_cpu_device)); + } + } + TF_ASSERT_OK(MarkForCompilationPassTestHelper::MarkForCompilation(&graph)); + + std::unordered_map clusters = GetClusters(*graph); + EXPECT_NE(clusters["test/a"], ""); + EXPECT_NE(clusters["test/b"], ""); + EXPECT_NE(clusters["test/c"], ""); +} + +TEST(XlaCompilationTest, DontAutoClusterStatefulRandomOp) { + Scope root = Scope::NewRootScope().ExitOnError(); + Output shape = ops::Const(root.WithOpName("test/shape_shape"), {200, 200}); + Output a = ops::RandomUniform(root.WithOpName("test/a"), shape, DT_FLOAT); + Output b = ops::RandomUniform(root.WithOpName("test/b"), shape, DT_FLOAT); + Output c = ops::Add(root.WithOpName("test/c"), a, b); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(root.ToGraph(graph.get())); + + TF_ASSERT_OK(MarkForCompilationPassTestHelper::MarkForCompilation(&graph)); + + std::unordered_map clusters = GetClusters(*graph); + EXPECT_EQ(clusters["test/a"], ""); + EXPECT_EQ(clusters["test/b"], ""); +} + +TEST(XlaCompilationTest, ClusterDummyOpsOnXlaDevice) { + absl::string_view xla_cpu_device = + "/job:worker/replica:0/task:0/device:XLA_CPU:0"; + + Scope root = Scope::NewRootScope().ExitOnError(); + Output a = ops::Placeholder(root.WithOpName("test/a"), DT_FLOAT); + Output b = ops::Placeholder(root.WithOpName("test/b"), DT_FLOAT); + Output check = + ops::CheckNumerics(root.WithOpName("test/check"), a, "test/check"); + Output ge = ops::GreaterEqual(root.WithOpName("test/greaterequal"), check, b); + Operation assert = ops::Assert(root.WithOpName("test/assert"), ge, {a, b}); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(root.ToGraph(graph.get())); + + for (Node* n : graph->nodes()) { + if (absl::StartsWith(n->name(), /*prefix=*/"test/")) { + n->set_assigned_device_name(string(xla_cpu_device)); + } + } + TF_ASSERT_OK(MarkForCompilationPassTestHelper::MarkForCompilation(&graph)); + + std::unordered_map clusters = GetClusters(*graph); + EXPECT_NE(clusters["test/check"], ""); + EXPECT_NE(clusters["test/greaterequal"], ""); + EXPECT_NE(clusters["test/assert"], ""); +} + +TEST(XlaCompilationTest, DontAutoClusterDummyOps) { + Scope root = Scope::NewRootScope().ExitOnError(); + Output a = ops::Placeholder(root.WithOpName("test/a"), DT_FLOAT); + Output b = ops::Placeholder(root.WithOpName("test/b"), DT_FLOAT); + Output check = + ops::CheckNumerics(root.WithOpName("test/check"), a, "test/check"); + Output ge = ops::GreaterEqual(root.WithOpName("test/greaterequal"), check, b); + Operation assert = ops::Assert(root.WithOpName("test/assert"), ge, {a, b}); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(root.ToGraph(graph.get())); + + TF_ASSERT_OK(MarkForCompilationPassTestHelper::MarkForCompilation(&graph)); + + std::unordered_map clusters = GetClusters(*graph); + EXPECT_EQ(clusters["test/assert"], ""); + EXPECT_EQ(clusters["test/check"], ""); +} + +TEST(XlaCompilationTest, DontAutoClusterOpsProducingVariant) { + Scope root = Scope::NewRootScope().ExitOnError(); + Output a = ops::Placeholder(root.WithOpName("test/a"), DT_INT64); + Output b = ops::Placeholder(root.WithOpName("test/b"), DT_INT64); + + Output cast_a = ops::Cast(root.WithOpName("test/cast_a"), a, DT_INT32); + Output cast_b = ops::Cast(root.WithOpName("test/cast_b"), b, DT_INT32); + + Output tensor_list_reserve = ops::TensorListReserve( + root.WithOpName("test/tensor_list_reserve"), cast_a, cast_b, DT_FLOAT); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(root.ToGraph(graph.get())); + + TF_ASSERT_OK(MarkForCompilationPassTestHelper::MarkForCompilation(&graph)); + + std::unordered_map clusters = GetClusters(*graph); + EXPECT_EQ(clusters["test/tensor_list_reserve"], ""); +} + +TEST(XlaCompilationTest, DontAutoClusterOpsConsumingVariant) { + Scope root = Scope::NewRootScope().ExitOnError(); + Output dummy_input = + ops::Placeholder(root.WithOpName("test/dummy_input"), DT_INT64); + Output variant_input = + ops::Placeholder(root.WithOpName("test/variant_input"), DT_VARIANT); + + // Create one more node so that we don't avoid creating a cluster solely + // because it would be trivial. + Output dummy_cast = + ops::Cast(root.WithOpName("test/dummy_cast"), dummy_input, DT_INT32); + + Output tensor_list_element_shape = ops::TensorListElementShape( + root.WithOpName("test/tensor_list_element_shape"), variant_input, + DT_INT32); + + root.graph()->AddControlEdge(dummy_cast.node(), + tensor_list_element_shape.node()); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(root.ToGraph(graph.get())); + + TF_ASSERT_OK(MarkForCompilationPassTestHelper::MarkForCompilation(&graph)); + + std::unordered_map clusters = GetClusters(*graph); + EXPECT_EQ(clusters["test/tensor_list_element_shape"], ""); +} + +TEST(XlaCompilationTest, ClusterOpsProducingVariantIfOnXlaDevice) { + Scope root = Scope::NewRootScope().ExitOnError(); + Output a = ops::Placeholder(root.WithOpName("test/a"), DT_INT64); + Output b = ops::Placeholder(root.WithOpName("test/b"), DT_INT64); + + Output cast_a = ops::Cast(root.WithOpName("test/cast_a"), a, DT_INT32); + Output cast_b = ops::Cast(root.WithOpName("test/cast_b"), b, DT_INT32); + + Output tensor_list_reserve = ops::TensorListReserve( + root.WithOpName("test/tensor_list_reserve"), cast_a, cast_b, DT_FLOAT); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(root.ToGraph(graph.get())); + + string xla_cpu_device = "/job:worker/replica:0/task:0/device:XLA_CPU:0"; + for (Node* n : graph->nodes()) { + if (absl::StartsWith(n->name(), /*prefix=*/"test/")) { + n->set_assigned_device_name(xla_cpu_device); + } + } + + TF_ASSERT_OK(MarkForCompilationPassTestHelper::MarkForCompilation(&graph)); + + std::unordered_map clusters = GetClusters(*graph); + EXPECT_NE(clusters["test/tensor_list_reserve"], ""); +} + } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/jit/mark_for_compilation_pass_test_helper.cc b/tensorflow/compiler/jit/mark_for_compilation_pass_test_helper.cc index d56d0f8ccfcdab40003be38059228cb255921b64..64a3301745790132fe3149bf8fb52d6c45ecc3c1 100644 --- a/tensorflow/compiler/jit/mark_for_compilation_pass_test_helper.cc +++ b/tensorflow/compiler/jit/mark_for_compilation_pass_test_helper.cc @@ -34,15 +34,9 @@ namespace tensorflow { // // It may be worth refactoring out XlaOpRegistry::RegisterCompilationDevice to // make this more direct, but probably not worth it solely for this test. - std::vector devices; + std::vector> devices; TF_RETURN_IF_ERROR(DeviceFactory::AddDevices(*session_options, "", &devices)); - auto delete_devices = gtl::MakeCleanup([&] { - for (Device* d : devices) { - delete d; - } - }); - GraphOptimizationPassOptions opt_options; opt_options.graph = graph; opt_options.session_options = session_options; diff --git a/tensorflow/compiler/jit/node_matchers.cc b/tensorflow/compiler/jit/node_matchers.cc index a09a6eb1553cb4bcf5587a7602097a40b64cfcdf..c788091724e443ba1e3bcd60515d68e71e2e0824 100644 --- a/tensorflow/compiler/jit/node_matchers.cc +++ b/tensorflow/compiler/jit/node_matchers.cc @@ -204,11 +204,12 @@ struct NodeMatcher : public ::testing::MatcherInterface { } return false; } - if (!AreAttrValuesEqual(it->second, attr_kv_pair.second)) { + if (attr_kv_pair.second && + !AreAttrValuesEqual(it->second, *attr_kv_pair.second)) { if (listener->IsInterested()) { *listener << "attribute named " << attr_kv_pair.first << " does not match value; expected: \"" - << SummarizeAttrValue(attr_kv_pair.second) + << SummarizeAttrValue(*attr_kv_pair.second) << "\", found: \"" << SummarizeAttrValue(it->second) << "\""; } @@ -278,12 +279,14 @@ struct NodeMatcher : public ::testing::MatcherInterface { if (!attrs.empty()) { printed_something = true; std::vector attrs_str; - absl::c_transform(attrs, std::back_inserter(attrs_str), - [](const std::pair& attr_kv_pair) { - return absl::StrCat( - attr_kv_pair.first, "->", - SummarizeAttrValue(attr_kv_pair.second)); - }); + absl::c_transform( + attrs, std::back_inserter(attrs_str), + [](const std::pair>& attr_kv_pair) { + return absl::StrCat(attr_kv_pair.first, "->", + attr_kv_pair.second + ? SummarizeAttrValue(*attr_kv_pair.second) + : "*"); + }); *os << " and attr values matching [" << absl::StrJoin(attrs_str, ", ") << "]"; } @@ -327,7 +330,7 @@ struct NodeMatcher : public ::testing::MatcherInterface { absl::optional>> input_matchers; absl::optional<::testing::Matcher>> control_dep_set; - std::map attrs; + std::map> attrs; }; // Matches a dst and dst_output on an input edge. Today we only use this with @@ -472,12 +475,38 @@ std::pair impl::AttrLiteralHelper( return {bool_attr.first, attr_value}; } +std::pair impl::AttrLiteralHelper( + const std::pair>& int_list_attr) { + AttrValue attr_value; + AttrValue::ListValue* list = attr_value.mutable_list(); + for (int i : int_list_attr.second) { + list->add_i(i); + } + return {int_list_attr.first, attr_value}; +} + +std::pair impl::AttrLiteralHelper( + const std::pair>& string_list_attr) { + AttrValue attr_value; + AttrValue::ListValue* list = attr_value.mutable_list(); + for (string s : string_list_attr.second) { + list->add_s(s); + } + return {string_list_attr.first, attr_value}; +} + impl::NodeMatcherProperties impl::Attr(std::pair attr) { impl::NodeMatcherProperties props; props.set_attr(std::move(attr)); return props; } +impl::NodeMatcherProperties impl::Attr(string name) { + impl::NodeMatcherProperties props; + props.set_attr({std::move(name), absl::nullopt}); + return props; +} + NodeMatcherProperties ConstantValue( const ::tensorflow::Input::Initializer& val) { TF_CHECK_OK(val.status); @@ -486,9 +515,9 @@ NodeMatcherProperties ConstantValue( return props; } -::testing::Matcher Const( +::testing::Matcher Const( const ::tensorflow::Input::Initializer& val) { - return NodeWith(ConstantValue(val)); + return Out(NodeWith(ConstantValue(val))); } ::testing::Matcher Out( int oidx, ::testing::Matcher node_matcher) { diff --git a/tensorflow/compiler/jit/node_matchers.h b/tensorflow/compiler/jit/node_matchers.h index 35c2f5fd7b533d0e8716dc6c70c21afe9a32c9c8..0d4f02c236bba353799f75ee91cf03235b424b29 100644 --- a/tensorflow/compiler/jit/node_matchers.h +++ b/tensorflow/compiler/jit/node_matchers.h @@ -84,7 +84,7 @@ class NodeMatcherProperties { public: using NodeSeqMatcher = std::vector<::testing::Matcher>; using InputSeqMatcher = std::vector<::testing::Matcher>; - using AttrKeyValuePair = std::pair; + using AttrKeyValuePair = std::pair>; const absl::optional& name() const { return name_; } const absl::optional& op() const { return op_; } @@ -163,9 +163,16 @@ impl::NodeMatcherProperties CtrlDeps( absl::Span> control_deps); impl::NodeMatcherProperties Attr(std::pair attrs); +impl::NodeMatcherProperties Attr(string name); std::pair AttrLiteralHelper( const std::pair& bool_attr); + +std::pair AttrLiteralHelper( + const std::pair>& int_list_attr); + +std::pair AttrLiteralHelper( + const std::pair>& string_list_attr); } // namespace impl // ----------------------------------------------------------------------------- @@ -187,6 +194,10 @@ impl::NodeMatcherProperties Attr(const string& name, ValueTy value) { return impl::Attr({impl::AttrLiteralHelper({name, value})}); } +inline impl::NodeMatcherProperties Attr(const string& name) { + return impl::Attr(name); +} + // Matches a node with inputs `inputs`. // // `inputs` are ordered; `inputs`[i] must match input i. @@ -200,7 +211,8 @@ impl::NodeMatcherProperties Inputs(Ts... inputs) { ::testing::Matcher node); // Matches the first output of a node that matches `node`. -::testing::Matcher Out(::testing::Matcher node) { +inline ::testing::Matcher Out( + ::testing::Matcher node) { return Out(0, node); } @@ -224,7 +236,7 @@ template return impl::NodeWith(array); } -::testing::Matcher Const( +::testing::Matcher Const( const ::tensorflow::Input::Initializer& val); } // namespace matchers diff --git a/tensorflow/compiler/jit/ops/BUILD b/tensorflow/compiler/jit/ops/BUILD index f72224545b25bc7100e0b6788e6fbf0a7ca63dad..64409d9334751e0edfce9091a4e5697dd2c712c5 100644 --- a/tensorflow/compiler/jit/ops/BUILD +++ b/tensorflow/compiler/jit/ops/BUILD @@ -18,3 +18,9 @@ tf_gen_op_wrapper_py( out = "xla_ops.py", deps = ["//tensorflow/compiler/jit/ops:xla_ops"], ) + +py_library( + name = "xla_ops_grad", + srcs = ["xla_ops_grad.py"], + deps = ["//tensorflow/python:framework_ops"], +) diff --git a/tensorflow/compiler/jit/ops/xla_ops_grad.py b/tensorflow/compiler/jit/ops/xla_ops_grad.py new file mode 100644 index 0000000000000000000000000000000000000000..2d31d8dc714307a48932d061fb1af643940a0872 --- /dev/null +++ b/tensorflow/compiler/jit/ops/xla_ops_grad.py @@ -0,0 +1,29 @@ +"""Gradients for XLA ops.""" +# 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. +# ============================================================================== + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.python.framework import ops + + +@ops.RegisterGradient("XlaClusterOutput") +def _XlaClusterOutputGrad(_, grad): + del grad # unused + raise RuntimeError("Gradient computation of graph in xla.compile() is " + "prohibited because it can cause performance degradation." + "Please move gradient computation inside xla.compile().") diff --git a/tensorflow/compiler/jit/partially_decluster_pass.cc b/tensorflow/compiler/jit/partially_decluster_pass.cc index 5b9610322336acbcede0bef0538043b8ff917c16..e1fd2aaee2822daeffb415d053c9c4f56002a856 100644 --- a/tensorflow/compiler/jit/partially_decluster_pass.cc +++ b/tensorflow/compiler/jit/partially_decluster_pass.cc @@ -26,6 +26,10 @@ limitations under the License. namespace tensorflow { namespace { + +bool NotBackedge(const Edge& edge) { return !edge.src()->IsNextIteration(); } + +namespace reduce_device_to_host_copies { Status FindNodesToDecluster(const Graph& graph, absl::flat_hash_set* result, absl::Span post_order) { @@ -116,6 +120,7 @@ Status PartiallyDeclusterNode(Graph* graph, Node* n) { NodeDef ndef = n->def(); ndef.set_name(absl::StrCat(n->name(), "/declustered")); + MergeDebugInfo(NodeDebugInfo(n->def()), &ndef); RemoveFromXlaCluster(&ndef); Status s; Node* cloned_node = graph->AddNode(ndef, &s); @@ -133,11 +138,13 @@ Status PartiallyDeclusterNode(Graph* graph, Node* n) { graph->RemoveEdge(out_edge_to_clone); } + if (n->out_edges().empty()) { + graph->RemoveNode(n); + } + return Status::OK(); } -bool NotBackedge(const Edge& edge) { return !edge.src()->IsNextIteration(); } - // Clones nodes to outside their cluster to avoid device-to-host copies. For // instance, converts this: // @@ -164,7 +171,7 @@ bool NotBackedge(const Edge& edge) { return !edge.src()->IsNextIteration(); } // where the ===> arrow has a hostmem source and destination and would entail a // device to host copy if the source and destination were not in the same XLA // cluster. -Status PartiallyDeclusterToRemoveDeviceToHostCopies(Graph* graph) { +Status PartiallyDeclusterGraph(Graph* graph) { // When deciding whether to decluster a particular node, we base our decision // on if we've decided that some of its consumers have to be declustered too. // Iterating the graph in post-order guarantees that consumers have been @@ -191,6 +198,10 @@ Status PartiallyDeclusterToRemoveDeviceToHostCopies(Graph* graph) { } } + // Recompute post order since PartiallyDeclusterNode may have deleted nodes. + post_order.clear(); + GetPostOrder(*graph, &post_order, /*stable_comparator=*/NodeComparatorName(), + /*edge_filter=*/NotBackedge); nodes_to_partially_decluster.clear(); TF_RETURN_IF_ERROR( FindNodesToDecluster(*graph, &nodes_to_partially_decluster, post_order)); @@ -198,7 +209,9 @@ Status PartiallyDeclusterToRemoveDeviceToHostCopies(Graph* graph) { return Status::OK(); } +} // namespace reduce_device_to_host_copies +namespace reduce_recompilation { bool IsIntraClusterEdge(const Edge& edge) { absl::optional src_cluster_name = GetXlaClusterForNode(*edge.src()); @@ -210,7 +223,8 @@ bool IsIntraClusterEdge(const Edge& edge) { bool IsMustCompileDevice(const DeviceType& device_type) { const XlaOpRegistry::DeviceRegistration* registration; if (XlaOpRegistry::GetCompilationDevice(device_type.type(), ®istration)) { - return registration->requires_compilation; + return registration->autoclustering_policy == + XlaOpRegistry::AutoclusteringPolicy::kAlways; } return false; @@ -260,7 +274,7 @@ Status MustCompileNode(const Node* n, bool* must_compile) { // regress performance in any significant manner. We will have to revisit this // algorith with a more complex cost model if this assumption turns out to be // incorrect. -Status DeclusterNodesToReduceRecompilations(Graph* graph) { +Status PartiallyDeclusterGraph(Graph* graph) { std::vector compile_time_const_nodes(graph->num_node_ids()); TF_RETURN_IF_ERROR(BackwardsConstAnalysis( *graph, nullptr, &compile_time_const_nodes, IsIntraClusterEdge)); @@ -313,7 +327,7 @@ Status DeclusterNodesToReduceRecompilations(Graph* graph) { return Status::OK(); } - +} // namespace reduce_recompilation } // namespace Status PartiallyDeclusterPass::Run( @@ -325,8 +339,9 @@ Status PartiallyDeclusterPass::Run( Graph* graph = options.graph->get(); - TF_RETURN_IF_ERROR(PartiallyDeclusterToRemoveDeviceToHostCopies(graph)); - TF_RETURN_IF_ERROR(DeclusterNodesToReduceRecompilations(graph)); + TF_RETURN_IF_ERROR( + reduce_device_to_host_copies::PartiallyDeclusterGraph(graph)); + TF_RETURN_IF_ERROR(reduce_recompilation::PartiallyDeclusterGraph(graph)); return Status::OK(); } diff --git a/tensorflow/compiler/jit/partially_decluster_pass_test.cc b/tensorflow/compiler/jit/partially_decluster_pass_test.cc index 74d5ef57184197ad6e9e5048722e84863756a3f5..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" @@ -386,7 +385,7 @@ TEST(PartiallyDeclusterPassTest, DontDeclusterXlaDeviceOps) { TF_ASSERT_OK(s.ToGraph(graph.get())); // This is needed to register the XLA_GPU device. - std::vector devices; + std::vector> devices; TF_ASSERT_OK(DeviceFactory::AddDevices( SessionOptions(), "/job:localhost/replica:0/task:0", &devices)); @@ -400,10 +399,6 @@ TEST(PartiallyDeclusterPassTest, DontDeclusterXlaDeviceOps) { TF_ASSERT_OK(PartiallyDecluster(&graph)); EXPECT_EQ(GetXlaClusterForNode(*n), "cluster_0"); - - for (Device* d : devices) { - delete d; - } } TEST(PartiallyDeclusterPassTest, DontDeclusterNonTensorFlowOps) { @@ -437,5 +432,32 @@ TEST(PartiallyDeclusterPassTest, DontDeclusterNonTensorFlowOps) { EXPECT_EQ(GetXlaClusterForNode(*n), "cluster_0"); } +TEST(PartiallyDeclusterPassTest, EliminatedUnusedNodes) { + const char* const kClusteredProducer0Name = "ClusteredProducer0"; + const char* const kClusteredProducer1Name = "ClusteredProducer1"; + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + { + GraphDefBuilder builder(GraphDefBuilder::kFailImmediately); + Node* input = + ops::SourceOp("FakeNullary", builder.opts().WithName("Input")); + Node* clustered_producer_0 = + ops::BinaryOp("FakeBinary", input, input, + builder.opts().WithName(kClusteredProducer0Name)); + Node* clustered_producer_1 = + ops::BinaryOp("FakeBinary", clustered_producer_0, input, + builder.opts().WithName(kClusteredProducer1Name)); + ops::BinaryOp("FakeBinary", clustered_producer_1, input, + builder.opts().WithName("UnclusteredConsumer")); + clustered_producer_0->AddAttr(kXlaClusterAttr, "cluster_0"); + clustered_producer_1->AddAttr(kXlaClusterAttr, "cluster_0"); + TF_EXPECT_OK(GraphDefBuilderToGraph(builder, graph.get())); + } + + TF_ASSERT_OK(PartiallyDecluster(&graph)); + EXPECT_EQ(FindNodeByName(*graph, kClusteredProducer0Name), nullptr); + EXPECT_EQ(FindNodeByName(*graph, kClusteredProducer1Name), nullptr); +} + } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/jit/producer_consumer_queue.h b/tensorflow/compiler/jit/producer_consumer_queue.h deleted file mode 100644 index 7c8c04152d2f3a0fd46711df24756b7e68b967ea..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/jit/producer_consumer_queue.h +++ /dev/null @@ -1,132 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#ifndef TENSORFLOW_COMPILER_JIT_PRODUCER_CONSUMER_QUEUE_H_ -#define TENSORFLOW_COMPILER_JIT_PRODUCER_CONSUMER_QUEUE_H_ - -#include -#include "tensorflow/core/platform/logging.h" -#include "tensorflow/core/platform/mutex.h" - -namespace tensorflow { - -// A thread-safe, first-in-first-out queue. -template -class ProducerConsumerQueue { - public: - ProducerConsumerQueue() - : capacity_(std::numeric_limits::max()) {} - ~ProducerConsumerQueue() = default; - - // Wait until the queue is non-full, then append a copy of v. - void Put(const T &v); - - // Wait until the queue is non-empty, then remove and return the head value. - T Get(); - - // If the queue is non-empty, remove the head value, placing it in *pv, and - // return true; otherwise return false. - bool TryGet(T *pv); - - // Set the capacity of the queue; the queue is full whenever count() >= - // capacity(). The initial value is the maximum size_t. Requires size > 0. - void set_capacity(std::size_t size); - - // Return the capacity of the queue. - std::size_t capacity() const; - - // Return the number of elements in the queue. - std::size_t count() const; - - // Implementation details follow. Clients should ignore. - private: - mutable tensorflow::mutex mu_; // protects all fields below - tensorflow::condition_variable non_empty_ GUARDED_BY(mu_); - tensorflow::condition_variable non_full_ GUARDED_BY(mu_); - std::size_t capacity_ GUARDED_BY(mu_); - std::deque queue_ GUARDED_BY(mu_); - - TF_DISALLOW_COPY_AND_ASSIGN(ProducerConsumerQueue); -}; - -// ------------------------------------------------------ -// Implementation details follow. Clients should ignore. - -// Wait until the queue is non-full, then append a copy of v. -template -void ProducerConsumerQueue::Put(const T &v) { - mutex_lock lock(mu_); - while (queue_.size() >= capacity_) { - non_full_.wait(lock); - } - queue_.push_back(v); - non_empty_.notify_one(); -} - -// Wait until the queue is non-empty, then remove and return the head value. -template -T ProducerConsumerQueue::Get() { - mutex_lock lock(mu_); - while (queue_.empty()) { - non_empty_.wait(lock); - } - non_full_.notify_one(); - T result_value = queue_.front(); - queue_.pop_front(); - return result_value; -} - -// If the queue is non-empty, remove the head value, placing it in *pv, and -// return true; otherwise return false. -template -bool ProducerConsumerQueue::TryGet(T *pv) { - mutex_lock lock(mu_); - bool got_element = !queue_.empty(); - if (got_element) { - non_full_.notify_one(); - *pv = queue_.front(); - queue_.pop_front(); - } - return got_element; -} - -// Set the capacity of the queue; the queue is full whenever count() >= -// capacity(). The initial value is the maximum size_t. Requires size > 0. -template -void ProducerConsumerQueue::set_capacity(std::size_t size) { - mutex_lock lock(mu_); - CHECK_NE(size, 0); - capacity_ = size; - non_full_.notify_all(); -} - -// Return the capacity of the queue. -template -std::size_t ProducerConsumerQueue::capacity() const { - mutex_lock lock(mu_); - std::size_t max_elements = capacity_; - return max_elements; -} - -// Return the number of elements in the queue. -template -std::size_t ProducerConsumerQueue::count() const { - mutex_lock lock(mu_); - std::size_t num_elements = queue_.size(); - return num_elements; -} -} // namespace tensorflow - -#endif // TENSORFLOW_COMPILER_JIT_PRODUCER_CONSUMER_QUEUE_H_ diff --git a/tensorflow/compiler/jit/producer_consumer_queue_test.cc b/tensorflow/compiler/jit/producer_consumer_queue_test.cc deleted file mode 100644 index f61260c6e52756ee039829afdc7452f5f760c221..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/jit/producer_consumer_queue_test.cc +++ /dev/null @@ -1,139 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/compiler/jit/producer_consumer_queue.h" - -#include "tensorflow/core/lib/core/threadpool.h" -#include "tensorflow/core/platform/env.h" -#include "tensorflow/core/platform/mutex.h" -#include "tensorflow/core/platform/test.h" - -namespace tensorflow { -namespace { - -typedef ProducerConsumerQueue IntQueue; - -// Insert integers between low inclusive and high exclusive into q. -void PushRange(IntQueue *q, int low, int high) { - while (low != high) { - q->Put(low); - VLOG(2) << "Pushing " << low; - ++low; - } -} - -// Push the numbers between 0 and 999 inclusive from several threads in the -// pool. -void PushRanges(IntQueue *queue, thread::ThreadPool *pool) { - VLOG(1) << "Adding 20-36"; - pool->Schedule([queue] { PushRange(queue, 20, 36); }); - VLOG(1) << "Adding 7-20"; - pool->Schedule([queue] { PushRange(queue, 7, 20); }); - VLOG(1) << "Adding 36-501"; - pool->Schedule([queue] { PushRange(queue, 36, 501); }); - VLOG(1) << "Adding 501-1000"; - pool->Schedule([queue] { PushRange(queue, 501, 1000); }); - VLOG(1) << "Adding 0-5"; - pool->Schedule([queue] { PushRange(queue, 0, 5); }); - VLOG(1) << "Adding 5-7"; - pool->Schedule([queue] { PushRange(queue, 5, 7); }); -} - -// Pop elements from queue using Get(). Make sure that exactly elements -// were present and their values are all integers between 0 and high-1 -// inclusive. -void GetRange(IntQueue *queue, int high) { - VLOG(1) << "Testing Wait"; - std::vector results; - for (int i = 0; i != high; ++i) { - int r = queue->Get(); - VLOG(2) << "Waited and got " << r; - results.push_back(r); - } - CHECK_EQ(queue->count(), 0); - std::sort(results.begin(), results.end()); - for (int i = 0; i != high; ++i) { - CHECK(results[i] == i); - } -} - -// Pop elements from queue using TryGet(). Make sure that exactly -// elements were present and their values are all integers between 0 and high-1 -// inclusive. -void TryGetRange(IntQueue *queue, int high) { - std::vector results; - // Give up if we don't get all the elements back from the queue - // in 10 seconds. - int timeout = 10; - int r; - for (int i = 0; i != high; ++i) { - while (!queue->TryGet(&r)) { - if (!timeout--) { - LOG(FATAL) << "Can't find all elements in the queue"; - } - VLOG(1) << "Sleeping for a second..."; - sleep(1); - } - VLOG(2) << "Popped " << r; - results.push_back(r); - } - CHECK_EQ(queue->count(), 0); - CHECK(!queue->TryGet(&r)); - std::sort(results.begin(), results.end()); - for (int i = 0; i != high; ++i) { - CHECK_EQ(i, results[i]); - } -} - -const int kNumThreads = 15; - -TEST(ProducerConsumerQueue, GetRange) { - IntQueue queue; - { - thread::ThreadPool pool(Env::Default(), "test", kNumThreads); - PushRanges(&queue, &pool); - } - GetRange(&queue, 1000); -} - -TEST(ProducerConsumerQueue, TryGetRange) { - IntQueue queue; - { - thread::ThreadPool pool(Env::Default(), "test", kNumThreads); - PushRanges(&queue, &pool); - } - TryGetRange(&queue, 1000); -} - -TEST(ProducerConsumerQueue, ParallelGetRange) { - IntQueue queue; - { - thread::ThreadPool pool(Env::Default(), "test", kNumThreads); - pool.Schedule([&queue] { GetRange(&queue, 1000); }); - PushRanges(&queue, &pool); - } -} - -TEST(ProducerConsumerQueue, ParallelTryGetRange) { - IntQueue queue; - { - thread::ThreadPool pool(Env::Default(), "test", kNumThreads); - pool.Schedule([&queue] { TryGetRange(&queue, 1000); }); - PushRanges(&queue, &pool); - } -} - -} // namespace -} // namespace tensorflow diff --git a/tensorflow/compiler/jit/resource_operation_safety_analysis.cc b/tensorflow/compiler/jit/resource_operation_safety_analysis.cc index e039d46ec863920eb7deb5bc20525fdab866415c..c0897217bcbd895003ce3018835da93a779a51a2 100644 --- a/tensorflow/compiler/jit/resource_operation_safety_analysis.cc +++ b/tensorflow/compiler/jit/resource_operation_safety_analysis.cc @@ -39,8 +39,7 @@ limitations under the License. // resource variables). // // The result is incorrect around loops because we ignore edges from -// NextIteration to Merge, but that should be fine because we don't cluster -// these edges. For instance, in: +// NextIteration to Merge. For instance, in: // // Init -----> Merge <-------+ // | | @@ -55,21 +54,20 @@ limitations under the License. // // we won't put (Read, Write) in the returned set. This is fine if // auto-clustering can only cluster the Read->Write edge, but it is a problem if -// it clusters the Write->NextIteration->Merge->Read edges instead. The same -// problem is present for the functional version of the loop above. We rely on -// auto-clustering to not cluster control flow edges like NextIteration->Merge. -// This is enough to avoid the explicit-control-flow problem shown above. One -// way to think about this is that we only care about cases where two nodes, A -// and B, would normally have been put in the same cluster but cannot legally be -// in the same cluster because of resourcevar-dependencies. If A and B would +// it clusters the Write->NextIteration->Merge->Read edges instead. So we rely +// on auto-clustering to not cluster NextIteration->Merge edges. The same +// problem is present for the functional version of the loop above and we also +// rely on auto-clustering not clustering functional while loops containing +// resource operations. +// +// One way to think about this is that we only care about cases where two nodes, +// A and B, would normally have been put in the same cluster but cannot legally +// be in the same cluster because of resourcevar-dependencies. If A and B would // normally have been put in the same cluster then all paths between A and B // would have to be clusterable (otherwise we'd have introduced a cycle). Ergo // there could not have been a NextIteration->Merge edge between A and B since // we don't cluster these edges. // -// We also rely on auto-clustering to not cluster functional control flow nodes -// that contain resource operations. -// // IMPLEMENTATION // -------------- // @@ -152,13 +150,12 @@ Status XlaResourceOpKindForNode( // can be represented by an XLA cluster and needs no special handling around // auto-jit. bool IsEdgeSafe(XlaResourceOpKind from, XlaResourceOpKind to) { - // XLA clusters forces all reads to happen before all writes, which means the - // kinds of edges it can faithfully represent are: Read->Write, Read->Modify, - // Modify->Write, Read->Read, Write->Write. - // - // TODO(b/112856632): We can, in theory, support Read->Read and Write->Write - // dependencies. - return from == XlaResourceOpKind::kRead && to == XlaResourceOpKind::kWrite; + // XLA clusters force all reads to happen before all writes. Moreover the set + // of reads are executed as one atomic operation, and the set of writes are as + // another atomic operation. This means we can faithfully represent the + // following edges: Read->*, *->Write. + + return from == XlaResourceOpKind::kRead || to == XlaResourceOpKind::kWrite; } using ResourceOp = std::pair; diff --git a/tensorflow/compiler/jit/resource_operation_safety_analysis_test.cc b/tensorflow/compiler/jit/resource_operation_safety_analysis_test.cc index e54b547abcfea698fe79e81dce547ea7858ff829..67304412fd384edde931fa2c5efb05f49e10411f 100644 --- a/tensorflow/compiler/jit/resource_operation_safety_analysis_test.cc +++ b/tensorflow/compiler/jit/resource_operation_safety_analysis_test.cc @@ -130,9 +130,7 @@ TEST(ResourceOperationSafetyAnalysisTest, ReadModify) { std::vector> incompatible_pairs; TF_ASSERT_OK(ComputeIncompatiblePairs(root.graph(), &incompatible_pairs)); - EXPECT_EQ(incompatible_pairs.size(), 1); - std::pair read_modify_pair = {read->id(), modify->id()}; - EXPECT_EQ(incompatible_pairs[0], read_modify_pair); + EXPECT_EQ(incompatible_pairs.size(), 0); } TEST(ResourceOperationSafetyAnalysisTest, ModifyRead) { @@ -162,9 +160,7 @@ TEST(ResourceOperationSafetyAnalysisTest, ModifyWrite) { std::vector> incompatible_pairs; TF_ASSERT_OK(ComputeIncompatiblePairs(root.graph(), &incompatible_pairs)); - EXPECT_EQ(incompatible_pairs.size(), 1); - std::pair modify_write_pair = {modify->id(), write->id()}; - EXPECT_EQ(incompatible_pairs[0], modify_write_pair); + EXPECT_EQ(incompatible_pairs.size(), 0); } TEST(ResourceOperationSafetyAnalysisTest, WriteModify) { @@ -196,11 +192,7 @@ TEST(ResourceOperationSafetyAnalysisTest, ReadModifyWrite) { std::vector> incompatible_pairs; TF_ASSERT_OK(ComputeIncompatiblePairs(root.graph(), &incompatible_pairs)); - EXPECT_EQ(incompatible_pairs.size(), 2); - std::pair modify_write_pair = {modify->id(), write->id()}; - std::pair read_modify_pair = {read->id(), modify->id()}; - EXPECT_EQ(incompatible_pairs[0], read_modify_pair); - EXPECT_EQ(incompatible_pairs[1], modify_write_pair); + EXPECT_EQ(incompatible_pairs.size(), 0); } TEST(ResourceOperationSafetyAnalysisTest, WriteModifyRead) { @@ -239,14 +231,12 @@ TEST(ResourceOperationSafetyAnalysisTest, WriteReadModify) { std::vector> incompatible_pairs; TF_ASSERT_OK(ComputeIncompatiblePairs(root.graph(), &incompatible_pairs)); - ASSERT_EQ(incompatible_pairs.size(), 3); + ASSERT_EQ(incompatible_pairs.size(), 2); std::pair write_modify_pair = {write->id(), modify->id()}; std::pair write_read_pair = {write->id(), read->id()}; - std::pair read_modify_pair = {read->id(), modify->id()}; - EXPECT_EQ(incompatible_pairs[0], read_modify_pair); - EXPECT_EQ(incompatible_pairs[1], write_read_pair); - EXPECT_EQ(incompatible_pairs[2], write_modify_pair); + EXPECT_EQ(incompatible_pairs[0], write_read_pair); + EXPECT_EQ(incompatible_pairs[1], write_modify_pair); } FunctionDefLibrary CreateFunctionDefLibWithConstFunction(const string& name) { @@ -307,9 +297,7 @@ TEST(ResourceOperationSafetyAnalysisTest, ReadCall) { std::vector> incompatible_pairs; TF_ASSERT_OK(ComputeIncompatiblePairs(root.graph(), &incompatible_pairs)); - ASSERT_EQ(incompatible_pairs.size(), 1); - std::pair read_call_edge = {read->id(), call->id()}; - EXPECT_EQ(incompatible_pairs[0], read_call_edge); + EXPECT_EQ(incompatible_pairs.size(), 0); } TEST(ResourceOperationSafetyAnalysisTest, CallWrite) { @@ -329,9 +317,7 @@ TEST(ResourceOperationSafetyAnalysisTest, CallWrite) { std::vector> incompatible_pairs; TF_ASSERT_OK(ComputeIncompatiblePairs(root.graph(), &incompatible_pairs)); - ASSERT_EQ(incompatible_pairs.size(), 1); - std::pair call_write_edge = {call->id(), write->id()}; - EXPECT_EQ(incompatible_pairs[0], call_write_edge); + EXPECT_EQ(incompatible_pairs.size(), 0); } TEST(ResourceOperationSafetyAnalysisTest, WriteCall) { @@ -429,18 +415,14 @@ TEST(ResourceOperationSafetyAnalysisTest, ChainOfOps) { std::vector> incompatible_pairs; TF_ASSERT_OK(ComputeIncompatiblePairs(root.graph(), &incompatible_pairs)); - ASSERT_EQ(incompatible_pairs.size(), 5); + ASSERT_EQ(incompatible_pairs.size(), 3); std::pair write_0_read_0_pair = {write_0->id(), read_0->id()}; std::pair write_0_read_1_pair = {write_0->id(), read_1->id()}; std::pair write_1_read_1_pair = {write_1->id(), read_1->id()}; - std::pair write_0_write_1_pair = {write_0->id(), write_1->id()}; - std::pair read_0_read_1_pair = {read_0->id(), read_1->id()}; EXPECT_EQ(incompatible_pairs[0], write_0_read_0_pair); - EXPECT_EQ(incompatible_pairs[1], write_0_write_1_pair); - EXPECT_EQ(incompatible_pairs[2], write_0_read_1_pair); - EXPECT_EQ(incompatible_pairs[3], read_0_read_1_pair); - EXPECT_EQ(incompatible_pairs[4], write_1_read_1_pair); + EXPECT_EQ(incompatible_pairs[1], write_0_read_1_pair); + EXPECT_EQ(incompatible_pairs[2], write_1_read_1_pair); } TEST(ResourceOperationSafetyAnalysisTest, DagOfOps) { diff --git a/tensorflow/compiler/jit/shape_inference.cc b/tensorflow/compiler/jit/shape_inference.cc index 80c691fe490c1092315708a2da754d367d585300..a27e0d9f2a6ecddfdbdb29be673084d77a178d8a 100644 --- a/tensorflow/compiler/jit/shape_inference.cc +++ b/tensorflow/compiler/jit/shape_inference.cc @@ -53,7 +53,15 @@ Status PropagateShapes(const Graph& graph, // shapes, even if no shape function is registered for a node. Status status = shape_refiner->AddNode(n); if (!status.ok()) { - VLOG(1) << "Shape inference failed for node: " << status; + VLOG(1) << "Shape inference failed for node " << n->name() << ": " + << status; + } else { + shape_inference::InferenceContext* context = shape_refiner->GetContext(n); + for (int i = 0; i < n->num_outputs(); i++) { + shape_inference::ShapeHandle handle = context->output(i); + VLOG(4) << "Output " << i << " for node " << n->name() << ": " + << context->DebugString(handle); + } } if (n->type_string() == "_Arg") { diff --git a/tensorflow/compiler/jit/xla_cluster_util.cc b/tensorflow/compiler/jit/xla_cluster_util.cc index f85121ca27ad3da918315f93b28e9000dfd65e67..80993861abba050fa3d6a133023d3c99f41f73e3 100644 --- a/tensorflow/compiler/jit/xla_cluster_util.cc +++ b/tensorflow/compiler/jit/xla_cluster_util.cc @@ -19,15 +19,17 @@ 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 { const char* const kXlaClusterAttr = "_XlaCluster"; const char* const kXlaOutsideCompilationAttr = "_XlaOutsideCompilation"; +const char* const kXlaCompileTimeConstantInputsAttr = + "_XlaCompileTimeConstantInputs"; namespace { // Returns a string describing how an edge from src to dst would diff --git a/tensorflow/compiler/jit/xla_cluster_util.h b/tensorflow/compiler/jit/xla_cluster_util.h index ba218f3315d2607c47342fdade0403678faa2362..fa6eaab3900b37baf7271c8c431c8384ceeda59f 100644 --- a/tensorflow/compiler/jit/xla_cluster_util.h +++ b/tensorflow/compiler/jit/xla_cluster_util.h @@ -32,6 +32,15 @@ extern const char* const kXlaClusterAttr; // compilation by the encapsulate subgraphs pass. extern const char* const kXlaOutsideCompilationAttr; +// The attribute that marks certain inputs to a Node as required to be a +// constant at compile time. If this attribute is present then the +// CompileTimeConstantInput information in the corresponding XlaOpKernel is +// ignored. +// +// The value for this attribute, if present, has to be a list of strings naming +// the inputs to the node that must be constant. +extern const char* const kXlaCompileTimeConstantInputsAttr; + using OrderedNodeSet = std::set; // Returns the DeviceType corresponding to 'device'. diff --git a/tensorflow/compiler/jit/xla_compilation_cache.cc b/tensorflow/compiler/jit/xla_compilation_cache.cc index 826e98b96620165604594a22b81cd02422605c12..611515cf33bc1abe21e06eb7f1513800276e095b 100644 --- a/tensorflow/compiler/jit/xla_compilation_cache.cc +++ b/tensorflow/compiler/jit/xla_compilation_cache.cc @@ -17,6 +17,8 @@ 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" @@ -37,6 +39,8 @@ limitations under the License. namespace tensorflow { +constexpr int64 XlaCompilationCache::kDefaultCompilationThreshold; + XlaCompilationCache::XlaCompilationCache(xla::LocalClient* client, DeviceType device_type) : client_(client), device_type_(std::move(device_type)) {} @@ -59,20 +63,20 @@ XlaCompilationCache::~XlaCompilationCache() { // about? } -string XlaCompilationCache::DebugString() { +string XlaCompilationCache::DebugString() const { return "XLA JIT compilation cache"; } // Compute a string signature which encodes the shapes of the // arguments in the supplied list. -string XlaCompilationCache::SignatureDebugString(const Signature& sig) { - string result = sig.name; - for (const auto& a : sig.arg_types) { - absl::StrAppend(&result, ",", DataTypeString(a.first), - a.second.DebugString()); +string XlaCompilationCache::Signature::HumanString() const { + string result = name; + for (const auto& a : arg_shapes) { + absl::StrAppend(&result, ",", DataTypeString(a.first)); + absl::StrAppend(&result, " [", absl::StrJoin(a.second, ","), "]"); } - for (const auto& v : sig.arg_values) { + for (const auto& v : arg_values) { absl::StrAppend(&result, "; ", v.DebugString()); } return result; @@ -80,11 +84,13 @@ string XlaCompilationCache::SignatureDebugString(const Signature& sig) { 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) { - if (arg_values[i].tensor_data() != other.arg_values[i].tensor_data()) { + if (arg_values[i].dtype() != other.arg_values[i].dtype() || + arg_values[i].shape() != other.arg_values[i].shape() || + arg_values[i].tensor_data() != other.arg_values[i].tensor_data()) { return false; } } @@ -94,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)); } } @@ -108,96 +114,30 @@ uint64 XlaCompilationCache::Signature::Hash::operator()( return h; } -Status XlaCompilationCache::BuildSignature( - const NameAttrList& function, const std::map& constant_args, - const std::map& variable_args, OpKernelContext* ctx, - Signature* signature) { - signature->name = Canonicalize(function.name(), AttrSlice(&function.attr())); - signature->arg_values.reserve(constant_args.size()); - - signature->arg_types.reserve(ctx->num_inputs() - constant_args.size()); - - for (int i = 0; i < ctx->num_inputs(); ++i) { - if (constant_args.count(i) > 0) { - // Use the values of compile time constants in the signature. - signature->arg_values.push_back(constant_args.at(i)); - } else if (variable_args.count(i) > 0) { - const OptionalTensor& variable = variable_args.at(i); - if (variable.present) { - signature->arg_types.emplace_back(variable.value.dtype(), - variable.value.shape()); - } else { - signature->arg_types.emplace_back(DT_INVALID, TensorShape()); - } - } else { - signature->arg_types.emplace_back(ctx->input_dtype(i), - ctx->input(i).shape()); +xla::StatusOr +XlaCompilationCache::BuildSignature( + const NameAttrList& function, + absl::Span args) { + Signature signature; + signature.name = Canonicalize(function.name(), AttrSlice(&function.attr())); + for (const XlaCompiler::Argument& arg : args) { + switch (arg.kind) { + case XlaCompiler::Argument::kConstant: + signature.arg_values.push_back(arg.constant_value); + break; + case XlaCompiler::Argument::kParameter: + case XlaCompiler::Argument::kResource: + signature.arg_shapes.emplace_back(arg.type, arg.DimensionSizes()); + break; + default: + return errors::InvalidArgument( + "Unhandled argument kind in XlaCompilationCache: ", + arg.HumanString()); } } - return Status::OK(); + return std::move(signature); } -namespace { - -// Builds a XlaCompiler::Argument vector from the arguments to the XlaLaunch op. -Status BuildArguments(const std::map& constant_args, - const std::map& variable_args, - OpKernelContext* ctx, - std::vector* args) { - args->resize(ctx->num_inputs()); - - for (int64 input_num = 0; input_num < ctx->num_inputs(); ++input_num) { - XlaCompiler::Argument& arg = (*args)[input_num]; - if (constant_args.count(input_num) > 0) { - // Handles compile-time constants. - const Tensor& input = constant_args.at(input_num); - TF_RET_CHECK(input.dtype() != DT_RESOURCE); - arg.kind = XlaCompiler::Argument::kConstant; - arg.type = input.dtype(); - arg.shape = input.shape(); - arg.constant_value = input; - } else if (variable_args.count(input_num) == 0) { - // Handles the non-constant arguments. - const Tensor& input = ctx->input(input_num); - TF_RET_CHECK(input.dtype() != DT_RESOURCE); - if (input.NumElements() > 0) { - arg.kind = XlaCompiler::Argument::kParameter; - } else { - arg.kind = XlaCompiler::Argument::kConstant; - arg.constant_value = input; - } - arg.type = input.dtype(); - arg.shape = input.shape(); - } else { - // Handles resource variables. - const Tensor& input = ctx->input(input_num); - TF_RET_CHECK(input.dtype() == DT_RESOURCE); - const OptionalTensor& variable = variable_args.at(input_num); - arg.name = variable.name; - arg.kind = XlaCompiler::Argument::kResource; - arg.resource_kind = XlaResource::kVariable; - if (variable.present) { - const Tensor& value = variable.value; - arg.type = value.dtype(); - arg.shape = value.shape(); - arg.initialized = true; - } else { - // The values of uninitialized variables are not passed as inputs, since - // they are meaningless. However, it is legal to assign to a resource - // variable for the first time inside the XLA computation, so we do - // permit uninitialized variables. - arg.initialized = false; - arg.type = DT_INVALID; - arg.shape = TensorShape(); - } - } - } - - return Status::OK(); -} - -} // namespace - Status XlaCompilationCache::BuildExecutable( const XlaCompiler::Options& options, const XlaCompiler::CompilationResult& result, @@ -227,25 +167,38 @@ Status XlaCompilationCache::BuildExecutable( Status XlaCompilationCache::Compile( const XlaCompiler::Options& options, const NameAttrList& function, - const std::map& constant_args, - const std::map& variable_args, OpKernelContext* ctx, + absl::Span args, const XlaCompiler::CompileOptions& compile_options, CompileMode compile_mode, const XlaCompiler::CompilationResult** out_compilation_result, xla::LocalExecutable** out_executable) { - // Set the compile threshold to 1 to implement CompileMode::kStrict. - int64 compile_threshold = - compile_mode == CompileMode::kLazy ? kDefaultCompilationThreshold : 1; - return CompileImpl(options, function, constant_args, variable_args, ctx, - compile_options, /*compile_single_op=*/false, + absl::optional compile_threshold; + if (compile_mode == CompileMode::kLazy) { + compile_threshold = kDefaultCompilationThreshold; + } + auto compile_fn = [&](XlaCompiler* compiler, + XlaCompiler::CompilationResult* result) { + return compiler->CompileFunction(compile_options, function, args, result); + }; + return CompileImpl(options, function, args, compile_fn, /*compile_threshold=*/compile_threshold, out_compilation_result, out_executable); } +static bool IsMegamorphic(int64 compile_count, int64 execution_count) { + const int64 kCompileThreshold = 10; + const int64 kMinExecutionsPerCompile = 50; + + // This heuristic is trying to capture the following property: have we sunk a + // certain minimum amount of compile time into the cluster that didn't quite + // "pay off"? + return compile_count > kCompileThreshold && + execution_count < kMinExecutionsPerCompile * compile_count; +} + Status XlaCompilationCache::CompileSingleOp( const XlaCompiler::Options& options, - const std::map& constant_args, - const std::map& variable_args, OpKernelContext* ctx, + absl::Span args, OpKernelContext* ctx, const XlaCompiler::CompileOptions& compile_options, const XlaCompiler::CompilationResult** out_compilation_result, xla::LocalExecutable** out_executable) { @@ -253,54 +206,41 @@ Status XlaCompilationCache::CompileSingleOp( NameAttrList name; name.set_name(def.op()); *name.mutable_attr() = def.attr(); - return CompileImpl(options, name, constant_args, variable_args, ctx, - compile_options, - /*compile_single_op=*/true, /*compile_threshold=*/1, + auto compile_op = [&](XlaCompiler* compiler, + XlaCompiler::CompilationResult* result) { + std::vector result_dtypes(ctx->num_outputs()); + for (int i = 0; i < result_dtypes.size(); ++i) { + result_dtypes[i] = ctx->expected_output_dtype(i); + } + return compiler->CompileSingleOp(compile_options, ctx->op_kernel().def(), + args, result_dtypes, result); + }; + return CompileImpl(options, name, args, compile_op, + /*compile_threshold=*/absl::nullopt, out_compilation_result, out_executable); } Status XlaCompilationCache::CompileImpl( const XlaCompiler::Options& options, const NameAttrList& function, - const std::map& constant_args, - const std::map& variable_args, OpKernelContext* ctx, - const XlaCompiler::CompileOptions& compile_options, bool compile_single_op, - int64 compile_threshold, + absl::Span args, + const std::function& compile_fn, + absl::optional compile_threshold, const XlaCompiler::CompilationResult** out_compilation_result, xla::LocalExecutable** out_executable) { DCHECK_NE(out_executable, nullptr); VLOG(2) << "XlaCompilationCache::Compile " << DebugString(); if (VLOG_IS_ON(2)) { - VLOG(2) << "num_inputs=" << ctx->num_inputs() - << " num_constant_args=" << constant_args.size() - << " num_variable_args=" << variable_args.size(); - for (int i = 0; i < ctx->num_inputs(); i++) { - TensorShape shape = ctx->input(i).shape(); - VLOG(2) << i << ": dtype=" << DataTypeString(ctx->input_dtype(i)) - << " present=" << ctx->has_input(i) - << " shape=" << shape.DebugString(); - } - for (auto& iterator : variable_args) { - const OptionalTensor& variable = iterator.second; - VLOG(2) << "variable present=" << variable.present - << " type=" << DataTypeString(variable.value.dtype()) - << " shape=" << variable.value.shape().DebugString() - << " TF arg= " << iterator.first; - } - VLOG(2) << "num_outputs = " << ctx->num_outputs(); - for (int i = 0; i < ctx->num_outputs(); i++) { - VLOG(2) << i << ": dtype=" << ctx->expected_output_dtype(i); + VLOG(2) << "num_inputs=" << args.size(); + for (int i = 0; i < args.size(); i++) { + VLOG(2) << i << ": " << args[i].HumanString(); } } - TF_RET_CHECK(constant_args.size() + variable_args.size() <= - ctx->num_inputs()); - - Signature signature; - TF_RETURN_IF_ERROR( - BuildSignature(function, constant_args, variable_args, ctx, &signature)); + TF_ASSIGN_OR_RETURN(Signature signature, BuildSignature(function, args)); + VLOG(2) << "Signature: " << signature.HumanString(); - VLOG(2) << "Signature: " << SignatureDebugString(signature); // The outer lock protects the existence of the cache entry. It does not // protect the contents of the cache entry. Entry* entry; @@ -314,17 +254,72 @@ Status XlaCompilationCache::CompileImpl( entry = e.get(); } + // We always compile a cluster the very first time it is executed. This is an + // optimistic guess that pays off for statically shaped TensorFlow graphs + // (since they get the benefit of XLA right away without waiting for warmup) + // and doesn't hurt much for dynamically shaped TensorFlow graphs (we "pay" at + // most one cluster-compilation's worth of compile time). + bool is_first_execution; + + // We avoid compiling clusters that have "gone megamorphic" i.e. have an + // excessive amount of shape dynamism. + bool is_megamorphic; + + { + mutex_lock lock(cluster_compile_stats_mu_); + auto it = + cluster_compile_stats_.emplace(function.name(), ClusterCompileStats{}) + .first; + is_first_execution = it->second.execution_count++ == 0; + + // The is_megamorphic bit is "sticky". We assume clusters that have been + // observed to be megamorphic once stay megamorphic forever. + it->second.is_megamorphic |= + IsMegamorphic(/*compile_count=*/it->second.compile_count, + /*execution_count=*/it->second.execution_count); + is_megamorphic = it->second.is_megamorphic; + } + // Acquire the cache entry lock and compile, if necessary. // TODO(phawkins): this locking will need to be restructured when we implement // cache eviction. mutex_lock entry_lock(entry->mu); int64 current_request_count = ++entry->request_count; + VLOG(2) << "Compilation cache entry hit: " << entry->compiled + << " signature: " << signature.HumanString() << " with request count " + << current_request_count << " and compile threshold " + << compile_threshold.value_or(0); if (!entry->compiled) { - VLOG(2) << "Compilation cache miss for signature: " - << SignatureDebugString(signature) << " with request count " - << current_request_count << " and compile threshold " - << compile_threshold; - if (current_request_count < compile_threshold) { + const bool should_compile = [&] { + if (!compile_threshold.has_value()) { + // Lazy compilation is disabled. + return true; + } + + if (is_megamorphic) { + VLOG(3) << "Not compiling cluster " << function.name() + << " because it is megamorphic."; + return false; + } + + if (is_first_execution) { + return true; + } + + bool reached_compile_threshold = + current_request_count >= *compile_threshold; + if (!reached_compile_threshold) { + VLOG(3) + << "Not compiling cluster " << function.name() + << " because it has not reached compile threshold; threshold is " + << *compile_threshold << " execution count " + << current_request_count << "."; + } + return reached_compile_threshold; + }(); + + if (!should_compile) { + VLOG(2) << "Not compiling for signature: " << signature.HumanString(); *out_compilation_result = nullptr; *out_executable = nullptr; return Status::OK(); @@ -334,21 +329,12 @@ Status XlaCompilationCache::CompileImpl( const uint64 compile_start_us = env->NowMicros(); // Do the actual JIT compilation without holding the lock (it can take // a long time.) - std::vector args; - TF_RETURN_IF_ERROR( - BuildArguments(constant_args, variable_args, ctx, &args)); XlaCompiler compiler(options); entry->compiled = true; - if (compile_single_op) { - entry->compilation_status = - compiler.CompileSingleOp(compile_options, signature.name, ctx, args, - &entry->compilation_result); - } else { - entry->compilation_status = compiler.CompileFunction( - compile_options, function, args, &entry->compilation_result); - } + entry->compilation_status = + compile_fn(&compiler, &entry->compilation_result); TF_RETURN_IF_ERROR(entry->compilation_status); CHECK_EQ(entry->executable.get(), nullptr); entry->compilation_status = @@ -357,8 +343,8 @@ Status XlaCompilationCache::CompileImpl( const uint64 compile_end_us = env->NowMicros(); const uint64 compile_time_us = compile_end_us - compile_start_us; { - mutex_lock lock(compile_stats_mu_); - auto it = compile_stats_.emplace(function.name(), CompileStats{}).first; + mutex_lock lock(cluster_compile_stats_mu_); + auto it = cluster_compile_stats_.find(function.name()); it->second.compile_count++; it->second.cumulative_compile_time_us += compile_time_us; VLOG(1) << "compiled " << function.name() << " " diff --git a/tensorflow/compiler/jit/xla_compilation_cache.h b/tensorflow/compiler/jit/xla_compilation_cache.h index f06a991818db53adb3e5c0cc483c6180128a87e7..7748b4700f39da4f952278ca6c6d2cadff4d3fb8 100644 --- a/tensorflow/compiler/jit/xla_compilation_cache.h +++ b/tensorflow/compiler/jit/xla_compilation_cache.h @@ -17,9 +17,12 @@ limitations under the License. #define TENSORFLOW_COMPILER_JIT_XLA_COMPILATION_CACHE_H_ #include "absl/container/flat_hash_map.h" +#include "absl/types/optional.h" +#include "absl/types/span.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" #include "tensorflow/compiler/tf2xla/xla_context.h" #include "tensorflow/compiler/xla/client/local_client.h" +#include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/core/common_runtime/device.h" #include "tensorflow/core/common_runtime/device_mgr.h" #include "tensorflow/core/framework/graph.pb.h" @@ -30,13 +33,6 @@ limitations under the License. namespace tensorflow { -// Struct that represents a possibly-absent Tensor. -struct OptionalTensor { - string name; // A descriptive name - bool present = false; // Is the tensor present? - Tensor value; // If present, what is the Tensor's value? -}; - // The XlaCompilationCache class caches the results of the XlaCompiler class, // which converts a Tensorflow graph into a compiled XLA compilation. // @@ -58,11 +54,7 @@ class XlaCompilationCache : public ResourceBase { // Compiles a function into a XlaCompiler::CompilationResult that can be used // to execute an XLA Computation. Compilation results are cached. // `function` is the name of a Tensorflow function to compile. - // `constant_args` is a map of tensorflow argument number to its constant - // value. - // `variable_args` is a snapshot of the current values of the - // resource variable arguments to `function`; uninitialized variables are - // represented by an absent OptionalTensor. + // `args` is a description of the arguments to the computation. // // `compile_mode` controls the behavior of the compilation cache on a cache // miss. If `compile_mode` is `kLazy` then, based on some profitability @@ -78,9 +70,7 @@ class XlaCompilationCache : public ResourceBase { // outputs. Status Compile(const XlaCompiler::Options& options, const NameAttrList& function, - const std::map& constant_args, - const std::map& variable_args, - OpKernelContext* ctx, + absl::Span args, const XlaCompiler::CompileOptions& compile_options, CompileMode compile_mode, const XlaCompiler::CompilationResult** out_compilation_result, @@ -90,8 +80,7 @@ class XlaCompilationCache : public ResourceBase { // XlaCompiler::CompileFunction. Status CompileSingleOp( const XlaCompiler::Options& options, - const std::map& constant_args, - const std::map& variable_args, OpKernelContext* ctx, + absl::Span args, OpKernelContext* ctx, const XlaCompiler::CompileOptions& compile_options, const XlaCompiler::CompilationResult** out_compilation_result, xla::LocalExecutable** out_executable); @@ -99,34 +88,16 @@ class XlaCompilationCache : public ResourceBase { xla::LocalClient* client() const { return client_; } const DeviceType& device_type() const { return device_type_; } - string DebugString() override; - - private: - // Common implementation of Compile and CompileSingleOp. - Status CompileImpl( - const XlaCompiler::Options& options, const NameAttrList& function, - const std::map& constant_args, - const std::map& variable_args, OpKernelContext* ctx, - const XlaCompiler::CompileOptions& compile_options, - bool compile_single_op, int64 compile_threshold, - const XlaCompiler::CompilationResult** out_compilation_result, - xla::LocalExecutable** out_executable); - - // Takes `result` which has been compiled from a Tensorflow subgraph to a - // XLA computation already, and generates an XLA LocalExecutable `executable`. - Status BuildExecutable(const XlaCompiler::Options& options, - const XlaCompiler::CompilationResult& result, - std::unique_ptr* executable); - - xla::LocalClient* const client_; - const DeviceType device_type_; + 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. @@ -137,14 +108,35 @@ class XlaCompilationCache : public ResourceBase { struct Hash { uint64 operator()(const Signature& signature) const; }; + + // Returns a human-readable description of the signature. + string HumanString() const; }; - static string SignatureDebugString(const Signature& sig); // Builds the signature for a compilation. - Status BuildSignature(const NameAttrList& function, - const std::map& constant_args, - const std::map& variable_args, - OpKernelContext* ctx, Signature* signature); + static xla::StatusOr BuildSignature( + const NameAttrList& function, + absl::Span args); + + private: + // Common implementation of Compile and CompileSingleOp. + Status CompileImpl( + const XlaCompiler::Options& options, const NameAttrList& function, + absl::Span args, + const std::function& compile_fn, + absl::optional compile_threshold, + const XlaCompiler::CompilationResult** out_compilation_result, + xla::LocalExecutable** out_executable); + + // Takes `result` which has been compiled from a Tensorflow subgraph to a + // XLA computation already, and generates an XLA LocalExecutable `executable`. + Status BuildExecutable(const XlaCompiler::Options& options, + const XlaCompiler::CompilationResult& result, + std::unique_ptr* executable); + + xla::LocalClient* const client_; + const DeviceType device_type_; // The value associated with a cache entry. struct Entry { @@ -171,18 +163,27 @@ class XlaCompilationCache : public ResourceBase { absl::flat_hash_map, Signature::Hash> cache_ GUARDED_BY(compile_cache_mu_); - struct CompileStats { + struct ClusterCompileStats { // Number of times the cluster has been (re-)compiled. int64 compile_count = 0; + // The number of times this cluster has been executed. + int64 execution_count = 0; + // Cumulative time spent compiling the cluster. int64 cumulative_compile_time_us = 0; + + // True if we have decided that this cluster is too dynamic (i.e. its shapes + // change too frequently) to profitably JIT compile. Once a cluster is + // tagged megamorphic, it stays megamorphic forever. + bool is_megamorphic = false; }; - mutex compile_stats_mu_; + + mutex cluster_compile_stats_mu_; // Maps cluster names to compilation statistics for said cluster. - absl::flat_hash_map compile_stats_ - GUARDED_BY(compile_stats_mu_); + absl::flat_hash_map cluster_compile_stats_ + GUARDED_BY(cluster_compile_stats_mu_); // The number of times a lazy compilation must be requested for a specific // signature before we attempt to compile it. diff --git a/tensorflow/compiler/jit/xla_compilation_cache_test.cc b/tensorflow/compiler/jit/xla_compilation_cache_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..018c7c219f445bdca17f4f8b060e3678fe1be9ee --- /dev/null +++ b/tensorflow/compiler/jit/xla_compilation_cache_test.cc @@ -0,0 +1,54 @@ +/* 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/jit/xla_compilation_cache.h" +#include "tensorflow/compiler/tf2xla/shape_util.h" +#include "tensorflow/core/platform/test.h" + +namespace tensorflow { +namespace { + +TEST(XlaCompilationCacheTest, SignatureEquality) { + NameAttrList fn; + fn.set_name("afunction"); + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kConstant; + args[0].type = DT_INT32; + args[0].shape = TensorShape({4, 0}); + args[0].constant_value = Tensor(DT_INT32, {4, 0}); + TF_ASSERT_OK_AND_ASSIGN(XlaCompilationCache::Signature s1, + XlaCompilationCache::BuildSignature(fn, args)); + + args[0].type = DT_FLOAT; + args[0].constant_value = Tensor(DT_FLOAT, {4, 0}); + TF_ASSERT_OK_AND_ASSIGN(XlaCompilationCache::Signature s2, + XlaCompilationCache::BuildSignature(fn, args)); + + args[0].shape = TensorShape({0, 4}); + args[0].constant_value = Tensor(DT_FLOAT, {0, 4}); + TF_ASSERT_OK_AND_ASSIGN(XlaCompilationCache::Signature s3, + XlaCompilationCache::BuildSignature(fn, args)); + + std::vector signatures = {s1, s2, s3}; + for (int i = 0; i < signatures.size(); ++i) { + for (int j = 0; j < signatures.size(); ++j) { + EXPECT_EQ(i == j, signatures[i] == signatures[j]) + << signatures[i].HumanString() << " " << signatures[j].HumanString(); + } + } +} + +} // namespace +} // namespace tensorflow diff --git a/tensorflow/compiler/jit/xla_compile_on_demand_op.cc b/tensorflow/compiler/jit/xla_compile_on_demand_op.cc index 129528bb4428564a130f1eaa30f5d887ce0349dc..c7e8d61d280a33a83c3386d8ef801018634d31ec 100644 --- a/tensorflow/compiler/jit/xla_compile_on_demand_op.cc +++ b/tensorflow/compiler/jit/xla_compile_on_demand_op.cc @@ -88,29 +88,26 @@ Status XlaCompileOnDemandOp::Run(OpKernelContext* ctx, return Status::OK(); } -bool XlaCompileOnDemandOp::MustArgumentBeConstant(const OpKernel* op_kernel, - int64 argument_idx) { +Status XlaCompileOnDemandOp::MustArgumentBeConstant(const OpKernel* op_kernel, + int64 argument_idx, + bool* result) { + *result = false; + // TODO(jmolloy): This could be expensive, so memoize. - auto* constant_inputs = tensorflow::XlaOpRegistry::CompileTimeConstantInputs( - op_kernel->def().op()); - CHECK(constant_inputs); - std::set constant_input_indices; - for (const auto& name : *constant_inputs) { - int start, stop; - TF_CHECK_OK(op_kernel->InputRange(name, &start, &stop)); - for (int i = start; i < stop; ++i) { - constant_input_indices.insert(i); - } - } - return constant_input_indices.count(argument_idx) > 0; + std::vector constant_input_indices; + TF_RETURN_IF_ERROR(XlaOpRegistry::CompileTimeConstantInputs( + *op_kernel, &constant_input_indices)); + *result = absl::c_binary_search(constant_input_indices, argument_idx); + return Status::OK(); } -bool XlaCompileOnDemandOp::ShouldArgumentBeConstant(const OpKernel* op_kernel, - int64 argument_idx) { +Status XlaCompileOnDemandOp::ShouldArgumentBeConstant(const OpKernel* op_kernel, + int64 argument_idx, + bool* result) { // Right now we only create kConstant arguments when absolutely required, but // there may be benefit in eagerly constant-folding a larger subset of // arguments in the future. - return MustArgumentBeConstant(op_kernel, argument_idx); + return MustArgumentBeConstant(op_kernel, argument_idx, result); } Status XlaCompileOnDemandOp::Compile( @@ -121,27 +118,48 @@ Status XlaCompileOnDemandOp::Compile( for (int64 i = 0; i < ctx->num_inputs(); ++i) { const Tensor& device_tensor = ctx->input(i); if (const XlaTensor* xla_tensor = XlaTensor::FromTensor(&device_tensor)) { - if (xla_tensor->has_host_tensor() && - ShouldArgumentBeConstant(&ctx->op_kernel(), i)) { - constant_arguments[i] = xla_tensor->host_tensor(); + if (xla_tensor->has_host_tensor()) { + bool should_arg_be_const; + TF_RETURN_IF_ERROR(ShouldArgumentBeConstant(&ctx->op_kernel(), i, + &should_arg_be_const)); + if (should_arg_be_const) { + constant_arguments[i] = xla_tensor->host_tensor(); + } } } - if (constant_arguments.count(i) == 0 && - MustArgumentBeConstant(&ctx->op_kernel(), i)) { - // Slow path; the argument is not available as a host constant so we must - // fetch it synchronously. - Tensor host_tensor; - AllocatorAttributes attrs; - attrs.set_on_host(true); - TF_RETURN_IF_ERROR(ctx->allocate_temp( - device_tensor.dtype(), device_tensor.shape(), &host_tensor, attrs)); - Notification n; - ctx->op_device_context()->CopyDeviceTensorToCPU( - &device_tensor, "ConstantArgument", - reinterpret_cast(ctx->device()), &host_tensor, - [&](Status status) { n.Notify(); }); - n.WaitForNotification(); - constant_arguments[i] = host_tensor; + + if (constant_arguments.count(i) == 0) { + bool must_argument_be_const; + TF_RETURN_IF_ERROR(MustArgumentBeConstant(&ctx->op_kernel(), i, + &must_argument_be_const)); + + if (must_argument_be_const) { + // Slow path; the argument is not available as a host constant so we + // must fetch it synchronously. + Tensor host_tensor; + AllocatorAttributes attrs; + attrs.set_on_host(true); + TF_RETURN_IF_ERROR(ctx->allocate_temp( + device_tensor.dtype(), device_tensor.shape(), &host_tensor, attrs)); + Notification n; + Status status; + ctx->op_device_context()->CopyDeviceTensorToCPU( + &device_tensor, "ConstantArgument", + reinterpret_cast(ctx->device()), &host_tensor, + [&](Status s) { + status = s; + n.Notify(); + }); + n.WaitForNotification(); + if (!status.ok()) { + LOG(ERROR) << "Copying tensor of shape " + << device_tensor.shape().DebugString() << " from " + << ctx->device()->name() << "to CPU failed with " + << status.ToString(); + return status; + } + constant_arguments[i] = host_tensor; + } } } @@ -166,9 +184,7 @@ Status XlaCompileOnDemandOp::Compile( XlaCompiler::Options options; options.device_type = metadata.jit_device_type(); options.client = metadata.client(); - auto flib_def = absl::make_unique( - OpRegistry::Global(), FunctionDefLibrary{}); - options.flib_def = flib_def.get(); + options.flib_def = ctx->function_library()->GetFunctionLibraryDefinition(); options.shape_representation_fn = metadata.shape_representation_fn(); XlaCompiler::CompileOptions compile_options; @@ -182,8 +198,14 @@ Status XlaCompileOnDemandOp::Compile( compile_options.always_return_tuple = false; std::map variable_args = GetVariables(ctx); - return cache->CompileSingleOp(options, constant_arguments, variable_args, ctx, - compile_options, result, executable); + + std::vector args; + + TF_RETURN_IF_ERROR(XlaComputationLaunchContext::BuildXlaCompilerArguments( + constant_arguments, variable_args, ctx, &args)); + + return cache->CompileSingleOp(options, args, ctx, compile_options, result, + executable); } void XlaCompileOnDemandOp::Compute(OpKernelContext* ctx) { diff --git a/tensorflow/compiler/jit/xla_compile_on_demand_op.h b/tensorflow/compiler/jit/xla_compile_on_demand_op.h index 7cc3d0e007ba2974fbfbe6fbabc4aa08f9fa910f..b93bb15ce34688f26316e22bf59f448e787df9fc 100644 --- a/tensorflow/compiler/jit/xla_compile_on_demand_op.h +++ b/tensorflow/compiler/jit/xla_compile_on_demand_op.h @@ -38,8 +38,10 @@ class XlaCompileOnDemandOp : public OpKernel { private: XlaCompiler::Argument CreateCompilerArgument(OpKernelContext* ctx, int64 i); - bool ShouldArgumentBeConstant(const OpKernel* op_kernel, int64 argument_idx); - bool MustArgumentBeConstant(const OpKernel* op_kernel, int64 argument_idx); + Status ShouldArgumentBeConstant(const OpKernel* op_kernel, int64 argument_idx, + bool* result); + Status MustArgumentBeConstant(const OpKernel* op_kernel, int64 argument_idx, + bool* result); Status Compile(OpKernelContext* ctx, const XlaDevice::Metadata& metadata, const XlaCompiler::CompilationResult** result, xla::LocalExecutable** executable); diff --git a/tensorflow/compiler/jit/xla_cpu_device.cc b/tensorflow/compiler/jit/xla_cpu_device.cc index f2cea3d000e239ce537cc79aca80f579ba299d1d..94dc61d55fb047c0ea81d98fde24cb55387c27d7 100644 --- a/tensorflow/compiler/jit/xla_cpu_device.cc +++ b/tensorflow/compiler/jit/xla_cpu_device.cc @@ -17,8 +17,8 @@ limitations under the License. // operators using XLA via the XLA "Host" (CPU) backend. #include "absl/memory/memory.h" +#include "tensorflow/compiler/jit/flags.h" #include "tensorflow/compiler/jit/kernels/xla_ops.h" -#include "tensorflow/compiler/jit/legacy_flags/xla_device_flags.h" #include "tensorflow/compiler/jit/xla_compile_on_demand_op.h" #include "tensorflow/compiler/jit/xla_device.h" #include "tensorflow/compiler/jit/xla_device_ops.h" @@ -31,19 +31,21 @@ namespace tensorflow { class XlaCpuDeviceFactory : public DeviceFactory { public: Status CreateDevices(const SessionOptions& options, const string& name_prefix, - std::vector* devices) override; + std::vector>* devices) override; }; -Status XlaCpuDeviceFactory::CreateDevices(const SessionOptions& session_options, - const string& name_prefix, - std::vector* devices) { - legacy_flags::XlaDeviceFlags* flags = legacy_flags::GetXlaDeviceFlags(); +Status XlaCpuDeviceFactory::CreateDevices( + const SessionOptions& session_options, const string& name_prefix, + std::vector>* devices) { + XlaDeviceFlags* flags = GetXlaDeviceFlags(); bool compile_on_demand = flags->tf_xla_compile_on_demand; XlaOpRegistry::DeviceRegistration registration; registration.compilation_device_name = DEVICE_CPU_XLA_JIT; - registration.requires_compilation = !compile_on_demand; - registration.enable_jit_by_default = false; + registration.autoclustering_policy = + compile_on_demand + ? XlaOpRegistry::AutoclusteringPolicy::kIfExplicitlyRequested + : XlaOpRegistry::AutoclusteringPolicy::kAlways; registration.compile_resource_ops = true; XlaOpRegistry::RegisterCompilationDevice(DEVICE_XLA_CPU, registration); @@ -60,10 +62,20 @@ Status XlaCpuDeviceFactory::CreateDevices(const SessionOptions& session_options, options.device_name = DEVICE_XLA_CPU; options.device_ordinal = 0; options.compilation_device_name = DEVICE_CPU_XLA_JIT; - options.transfer_as_literal = false; options.use_multiple_streams = false; auto device = absl::make_unique(session_options, options); - devices->push_back(device.release()); + + // Setting GpuDeviceInfo because eager runtime relies on the device + // context in tensorflow_gpu_device_info(). Also, + // tensorflow_gpu_device_info() == nullptr is used as an IsCPU test. + // We need XlaCpuDevice to be treated not as CPU because it allocates + // XlaTensors, not regular Tensors. + Status status = device->UseGpuDeviceInfo(); + if (!status.ok()) { + errors::AppendToMessage(&status, "while setting up ", DEVICE_GPU_XLA_JIT); + return status; + } + devices->push_back(std::move(device)); return Status::OK(); } @@ -71,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 17353456eb5d18a818e8cef196151ab85980bb44..e2397f6fcb8677f4bd5151646f9ebacd3e23af5b 100644 --- a/tensorflow/compiler/jit/xla_device.cc +++ b/tensorflow/compiler/jit/xla_device.cc @@ -201,12 +201,19 @@ XlaDevice::XlaDevice(const SessionOptions& session_options, jit_device_name_(options.compilation_device_name), platform_(options.platform), use_multiple_streams_(options.use_multiple_streams), - transfer_as_literal_(options.transfer_as_literal), - shape_representation_fn_(options.shape_representation_fn) { + shape_representation_fn_(options.shape_representation_fn), + allowed_devices_(options.allowed_devices) { VLOG(1) << "Created XLA device " << options.compilation_device_name << " " << this; thread_pool_.reset(new thread::ThreadPool(session_options.env, "xla_device", /*num_threads=*/1)); + + // We have multiple device to device streams to allow for some concurrency + // between transfers. The particular value of '4' is chosen fairly + // arbitrarily. It may be necessary to make this tunable via + // XlaDevice::Options. + static constexpr int kNumDeviceToDeviceStreams = 4; + device_to_device_streams_.resize(kNumDeviceToDeviceStreams); } XlaDevice::~XlaDevice() { @@ -225,7 +232,8 @@ xla::LocalClient* XlaDevice::client() const { // TODO(b/78468222): This can fail, at least when the backend is GPU and // there is no GPU on the host. - return xla::ClientLibrary::GetOrCreateLocalClient(platform_).ValueOrDie(); + return xla::ClientLibrary::GetOrCreateLocalClient(platform_, allowed_devices_) + .ValueOrDie(); } Allocator* XlaDevice::GetAllocator(AllocatorAttributes attr) { @@ -274,8 +282,9 @@ xla::StatusOr XlaDevice::GetDeviceContextLocked() { TF_RETURN_IF_ERROR(EnsureStreamOkLocked(backend, "stream", &stream_, &need_new_device_context)); - std::shared_ptr host_to_device_stream = stream_; - std::shared_ptr device_to_host_stream = stream_; + std::shared_ptr host_to_device_stream; + std::shared_ptr device_to_host_stream; + std::vector> device_to_device_streams; if (use_multiple_streams_) { TF_RETURN_IF_ERROR(EnsureStreamOkLocked(backend, "host_to_device_stream", &host_to_device_stream_, @@ -283,8 +292,18 @@ xla::StatusOr XlaDevice::GetDeviceContextLocked() { TF_RETURN_IF_ERROR(EnsureStreamOkLocked(backend, "device_to_host_stream", &device_to_host_stream_, &need_new_device_context)); + for (std::shared_ptr& stream : device_to_device_streams_) { + TF_RETURN_IF_ERROR( + EnsureStreamOkLocked(backend, "device_to_device_stream", &stream, + &need_new_device_context)); + } host_to_device_stream = host_to_device_stream_; device_to_host_stream = device_to_host_stream_; + device_to_device_streams = device_to_device_streams_; + } else { + host_to_device_stream = stream_; + device_to_host_stream = stream_; + device_to_device_streams = {stream_}; } if (!need_new_device_context) { @@ -302,8 +321,9 @@ xla::StatusOr XlaDevice::GetDeviceContextLocked() { // ensures that the streams remain live for the duration of a run, even if // an error is encountered and the streams are replaced with new ones. device_context_ = new XlaDeviceContext( - stream_, host_to_device_stream, device_to_host_stream, client(), - transfer_as_literal_, shape_representation_fn_, thread_pool_.get()); + stream_, std::move(host_to_device_stream), + std::move(device_to_host_stream), std::move(device_to_device_streams), + client(), shape_representation_fn_, thread_pool_.get()); VLOG(1) << "XlaDevice " << this << " new XlaDeviceContext " << device_context_; @@ -366,6 +386,7 @@ void XlaDevice::ComputeAsync(AsyncOpKernel* op_kernel, OpKernelContext* context, Status XlaDevice::Sync() { VLOG(1) << "XlaDevice::Sync"; + tracing::ScopedActivity activity("XlaDevice::Sync", /*is_expensive=*/true); std::shared_ptr stream; { mutex_lock lock(mu_); @@ -373,13 +394,48 @@ Status XlaDevice::Sync() { } if (!stream) return Status::OK(); - if (!stream->parent()->SynchronizeAllActivity() || !stream->ok()) { + Status status = stream->BlockHostUntilDone(); + TF_RETURN_IF_ERROR(status); + if (!stream->ok()) { return errors::Internal("XlaDevice::Sync() failed."); } VLOG(1) << "XlaDevice::Sync completed"; 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; + { + mutex_lock lock(mu_); + stream = stream_; + } + if (!stream) { + done(Status::OK()); + 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*) { + tracing::ScopedActivity activity("XlaDevice::Sync::Callback", + /*is_expensive=*/true); + done(stream->ok() ? Status::OK() + : errors::Internal("XlaDevice::Sync() failed.")); + }); +} + Status XlaDevice::MakeTensorFromProto(const TensorProto& tensor_proto, const AllocatorAttributes alloc_attrs, Tensor* tensor) { @@ -413,22 +469,34 @@ 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_; } +Status XlaDevice::CurrentStatus() { + std::shared_ptr stream; + { + mutex_lock lock(mu_); + stream = stream_; + } + if (!stream) { + return Status::OK(); + } + return stream->ok() ? Status::OK() : errors::Internal("XlaDevice is not OK."); +} + XlaDeviceOpRegistrations* RegisterXlaDeviceKernels(const char* device, const char* jit_device) { // Any op assigned to the device that isn't rewritten by the graph rewriter // gets executed by a n XlaCompileOnDemandOp, which compiles it and executes // it just-in-time. - kernel_factory::OpKernelRegistrar::Factory factory = + OpKernel* (*factory)(OpKernelConstruction*) = [](OpKernelConstruction* context) -> OpKernel* { return new XlaCompileOnDemandOp(context); }; diff --git a/tensorflow/compiler/jit/xla_device.h b/tensorflow/compiler/jit/xla_device.h index 223f0f6649f0544db242c59e85208196d9513978..e35a1c7d29514dc5777bdbd3858c56401d7b9044 100644 --- a/tensorflow/compiler/jit/xla_device.h +++ b/tensorflow/compiler/jit/xla_device.h @@ -24,7 +24,9 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_JIT_XLA_DEVICE_H_ #define TENSORFLOW_COMPILER_JIT_XLA_DEVICE_H_ +#include +#include "absl/types/optional.h" #include "tensorflow/compiler/jit/xla_device_context.h" #include "tensorflow/compiler/jit/xla_tensor.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" @@ -108,20 +110,26 @@ class XlaDevice : public LocalDevice { // The name of the compilation device (e.g., "XLA_CPU_JIT"); string compilation_device_name; - // 'transfer_as_literal' is true if device<->host transfers must be done - // using XLA's TransferLiteral{To,From}Device interface. If false, we can - // use ThenMemcpy instead. - bool transfer_as_literal = false; - // If 'use_multiple_streams' is true, we create separate streams for // compute, host-to-device, and device-to-host communication. bool use_multiple_streams = false; + // A function that describes how the on-host shapes of + // a) argument and return value, for entry computations + // b) variables, for all computations, + // should be represented in XLA. Parameters/return values will be shaped + // according to this function, and reshaped back to/from their declared + // shapes for computations. Must be non-null. XlaCompiler::ShapeRepresentationFn shape_representation_fn; // If padded_shape_fn is empty, a default implementation that returns // the logical on-device shape without padding is used. PaddedShapeFn padded_shape_fn; + + // Set of devices to use. This controls which of the devices on the given + // platform will have resources allocated. For GPUs this will be + // filled from visible_gpu_devices list from session configuration. + absl::optional> allowed_devices; }; // Creates a new XLA Device. @@ -134,6 +142,7 @@ class XlaDevice : public LocalDevice { void ComputeAsync(AsyncOpKernel* op_kernel, OpKernelContext* context, AsyncOpKernel::DoneCallback done) override; Status Sync() override; + void Sync(const DoneCallback& done) override; Status FillContextMap(const Graph* graph, DeviceContextMap* device_context_map) override @@ -158,10 +167,12 @@ 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_); + Status CurrentStatus() override LOCKS_EXCLUDED(mu_); private: xla::LocalClient* client() const; @@ -188,6 +199,7 @@ class XlaDevice : public LocalDevice { se::Platform* const platform_; // Not owned. // Memory allocator associated with this device. Allocator* xla_allocator_ GUARDED_BY(mu_) = nullptr; // Not owned. + // Stream associated with this device. Operations enqueued on this // stream are executed on the device. Operations include data // copying back and forth between CPU and the device, and @@ -203,9 +215,11 @@ class XlaDevice : public LocalDevice { // If use_multiple_streams_, device to host transfers are performed using this // stream. std::shared_ptr device_to_host_stream_ GUARDED_BY(mu_); - // Must we use XLA's transfer manager for correct host<->device transfers? if - // false, we can use ThenMemcpy() instead. - const bool transfer_as_literal_; + // If use_multiple_streams_, transfers between different devices are performed + // using these streams. + std::vector> device_to_device_streams_ + GUARDED_BY(mu_); + const XlaCompiler::ShapeRepresentationFn shape_representation_fn_; // The device context accessed by all users of the XlaDevice, set by calls to @@ -220,9 +234,14 @@ 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; + 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 + // filled from visible_gpu_devices list from session configuration. + absl::optional> allowed_devices_; }; // Builds OpKernel registrations on 'device' for the JIT operators diff --git a/tensorflow/compiler/jit/xla_device_context.cc b/tensorflow/compiler/jit/xla_device_context.cc index 090021093d67384521f5fad43b226b5263829c99..28681bb8b03dbf97e8145972f9a04b5855fafdae 100644 --- a/tensorflow/compiler/jit/xla_device_context.cc +++ b/tensorflow/compiler/jit/xla_device_context.cc @@ -53,85 +53,37 @@ void XlaDeviceAllocator::GetStats(AllocatorStats* stats) { stats->Clear(); } XlaDeviceContext::XlaDeviceContext( std::shared_ptr compute_stream, std::shared_ptr host_to_device_stream, - std::shared_ptr device_to_host_stream, xla::LocalClient* client, - bool transfer_as_literal, + std::shared_ptr device_to_host_stream, + std::vector> device_to_device_streams, + xla::LocalClient* client, XlaCompiler::ShapeRepresentationFn shape_representation_fn, thread::ThreadPool* thread_pool) : stream_(std::move(compute_stream)), host_to_device_stream_(std::move(host_to_device_stream)), device_to_host_stream_(std::move(device_to_host_stream)), + device_to_device_streams_(std::move(device_to_device_streams)), client_(client), transfer_manager_(client->backend().transfer_manager()), - transfer_as_literal_(transfer_as_literal), shape_representation_fn_(std::move(shape_representation_fn)), thread_pool_(thread_pool) { CHECK(host_to_device_stream_ != nullptr); CHECK(device_to_host_stream_ != nullptr); CHECK(stream_ != nullptr); if (!shape_representation_fn_) { - shape_representation_fn_ = - [](const TensorShape& shape, - DataType dtype) -> xla::StatusOr { return shape; }; + shape_representation_fn_ = [](const TensorShape& shape, + DataType dtype) -> xla::StatusOr { + xla::Shape xla_shape; + TF_RETURN_IF_ERROR(TensorShapeToXLAShape(dtype, shape, &xla_shape)); + return xla_shape; + }; } } -Status XlaDeviceContext::TransferLiteralToDevice(const Tensor& host_tensor, - Tensor* device_tensor) const { - xla::Shape xla_shape; - TF_RETURN_IF_ERROR(TensorShapeToXLAShape(host_tensor.dtype(), - host_tensor.shape(), &xla_shape)); - // Create a reference to hold onto host_tensor until after the literal has - // been transferred. Also make sure the literal exists until the function - // asynchronously completes, as it will be wrapped in an xla::LiteralSlice. - TensorReference ref(host_tensor); - auto literal = std::make_shared( - static_cast(DMAHelper::base(&host_tensor)), xla_shape); - - XlaTensor* xla_tensor = XlaTensor::FromTensor(device_tensor); - const xla::ShapedBuffer& shaped_buffer = xla_tensor->shaped_buffer(); - VLOG(1) << "Transfer to device as literal: " << literal->ToString() << " " - << shaped_buffer.ToString(); - if (UseMultipleStreams() && !transfer_manager_->CanShapedBufferBeAccessedNow( - stream_->parent(), shaped_buffer)) { - // Initially wait for the compute stream so that memory allocations are - // synchronized. - host_to_device_stream_->ThenWaitFor(stream_.get()); - } - TF_RETURN_IF_ERROR(transfer_manager_->TransferLiteralToDeviceAsync( - host_to_device_stream_.get(), *literal, shaped_buffer)); - if (UseMultipleStreams()) { - auto event = std::make_shared(stream_->parent()); - TF_RET_CHECK(event->Init()) << "Event failed to initialize!"; - host_to_device_stream_->ThenRecordEvent(event.get()); - xla_tensor->SetDefinedOn(host_to_device_stream_.get(), std::move(event)); - } - // Unref the host tensor, and capture the literal shared_ptr too so it goes - // out of scope when the lambda completes. - host_to_device_stream_->ThenDoHostCallback([ref, literal]() { ref.Unref(); }); - - return Status::OK(); -} - -void XlaDeviceContext::TransferLiteralFromDevice( - Tensor* host_tensor, const Tensor& device_tensor, - const StatusCallback& done) const { - xla::MutableBorrowingLiteral literal; - TF_CHECK_OK(HostTensorToMutableBorrowingLiteral(host_tensor, &literal)); - - const xla::ShapedBuffer& shaped_buffer = - XlaTensor::FromTensor(&device_tensor)->shaped_buffer(); - - TensorReference ref(device_tensor); - transfer_manager_->TransferLiteralFromDevice( - device_to_host_stream_.get(), shaped_buffer, literal, - [=, &shaped_buffer](xla::Status status) { - ref.Unref(); - done([&]() -> Status { - VLOG(1) << "Transfer from device as literal: " - << shaped_buffer.ToString(); - return status; - }()); - }); +void XlaDeviceContext::CopyTensorInSameDevice(const Tensor* input_tensor, + Device* device, + Tensor* output_tensor, + StatusCallback done) const { + done(errors::Unimplemented("XLA->XLA same-device copies not implemented.")); } void XlaDeviceContext::CopyCPUTensorToDevice(const Tensor* cpu_tensor, @@ -152,54 +104,79 @@ void XlaDeviceContext::CopyCPUTensorToDevice(const Tensor* cpu_tensor, << cpu_tensor->shape().DebugString() << " " << device_tensor->shape().DebugString(); - void* src_ptr = const_cast(DMAHelper::base(cpu_tensor)); - const int64 total_bytes = cpu_tensor->TotalBytes(); XlaTensor* xla_tensor = XlaTensor::FromTensor(device_tensor); CHECK(xla_tensor); - xla::StatusOr shape_or_status = - shape_representation_fn_(device_tensor->shape(), device_tensor->dtype()); - if (!shape_or_status.ok()) { - done(shape_or_status.status()); - return; - } - TensorShape shape = shape_or_status.ValueOrDie(); - if (!xla_tensor->has_shaped_buffer()) { - Status s = + Status status = [&]() -> Status { + TF_ASSIGN_OR_RETURN(xla::Shape shape, + shape_representation_fn_(device_tensor->shape(), + device_tensor->dtype())); + + // The device tensor should always be fresh. + TF_RET_CHECK(!xla_tensor->has_shaped_buffer()); + + xla_tensor->set_host_tensor(*cpu_tensor); + TF_RETURN_IF_ERROR( xla_tensor->AllocateShapedBuffer(device_tensor->dtype(), shape, client_, - stream_->parent()->device_ordinal()); - if (!s.ok()) { - done(s); - return; + stream_->parent()->device_ordinal())); + + // The cpu_tensor and literal that we created here hold the data of host + // tensor in descending layout. The layout could be different from layout in + // device_tensor (but the logical shape has to be the same). The + // transfer_manager is responsible to do corresponding transposing when + // transferring the data to device. + xla::BorrowingLiteral literal( + static_cast(DMAHelper::base(cpu_tensor)), + xla::ShapeUtil::MakeShape(shape.element_type(), + xla::AsInt64Slice(shape.dimensions()))); + + VLOG(2) << "Transfer to device as literal: " << literal.ToString() << " " + << xla_tensor->shaped_buffer().ToString(); + if (UseMultipleStreams() && + !transfer_manager_->CanShapedBufferBeAccessedNow( + stream_->parent(), xla_tensor->shaped_buffer())) { + // Initially wait for the compute stream so that memory allocations are + // synchronized. + host_to_device_stream_->ThenWaitFor(stream_.get()); } - } - Status status; - if (transfer_as_literal_) { - Tensor reshaped_cpu_tensor; - if (!reshaped_cpu_tensor.CopyFrom(*cpu_tensor, shape)) { - done(errors::Internal( - "Tensor::CopyFrom failed when copying from CPU to XLA device")); - return; - } - status = TransferLiteralToDevice(reshaped_cpu_tensor, device_tensor); - } else { - se::DeviceMemoryBase dev_dst_ptr = - XlaTensor::DeviceMemoryFromTensor(*device_tensor); - host_to_device_stream_->ThenMemcpy(&dev_dst_ptr, src_ptr, total_bytes); - // TODO(hpucha): Make this asynchronous. - Status block_status = host_to_device_stream_->BlockHostUntilDone(); - if (!block_status.ok()) { - status = xla::InternalError( - "Failed to complete data transfer on stream %p: %s", - host_to_device_stream_.get(), block_status.error_message().c_str()); + TF_RETURN_IF_ERROR(transfer_manager_->TransferLiteralToDeviceAsync( + host_to_device_stream_.get(), literal, xla_tensor->shaped_buffer())); + + if (UseMultipleStreams()) { + auto event = std::make_shared(stream_->parent()); + TF_RET_CHECK(event->Init()) << "Event failed to initialize!"; + host_to_device_stream_->ThenRecordEvent(event.get()); + xla_tensor->ResetDefinitionEvent(std::move(event), + host_to_device_stream_.get()); } + + return Status::OK(); + }(); + if (!status.ok()) { + done(status); + return; } - if (status.ok()) { - xla_tensor->set_host_tensor(*cpu_tensor); + + // Create a reference to hold onto cpu_tensor until after the literal has + // been transferred + TensorReference ref(*cpu_tensor); + if (UseMultipleStreams()) { + // Unref the host tensor when the transfer completes. + // We don't defer the call to done() onto the stream here, and the reasons + // why this is correct are subtle. We assume that: + // a) all consumers of the device tensor will wait for its definition event. + // b) if the tensor is destroyed, then the memory allocator will not hand + // out the same buffers until the transfer has completed. + host_to_device_stream_->ThenDoHostCallback([ref]() { ref.Unref(); }); + done(status); + } else { + host_to_device_stream_->ThenDoHostCallback([ref, done]() { + ref.Unref(); + done(Status::OK()); + }); } - done(status); } void XlaDeviceContext::CopyDeviceTensorToCPU(const Tensor* device_tensor, @@ -219,104 +196,38 @@ void XlaDeviceContext::CopyDeviceTensorToCPU(const Tensor* device_tensor, << cpu_tensor->shape().DebugString() << " " << device_tensor->shape().DebugString(); - const int64 total_bytes = cpu_tensor->TotalBytes(); - se::DeviceMemoryBase dev_src_ptr = - XlaTensor::DeviceMemoryFromTensor(*device_tensor); - void* dst_ptr = DMAHelper::base(cpu_tensor); XlaTensor* xla_tensor = XlaTensor::FromTensor(device_tensor); + xla_tensor->WaitForDefinitionEventOnStream(device_to_host_stream_.get()); - if (se::Event* event = - xla_tensor->GetDefinitionEvent(device_to_host_stream_.get())) { - device_to_host_stream_->ThenWaitFor(event); - xla_tensor->SetDefinedOn(device_to_host_stream_.get()); - } - - Status status; - if (transfer_as_literal_) { - TransferLiteralFromDevice(cpu_tensor, *device_tensor, done); - return; - } else { - device_to_host_stream_->ThenMemcpy(dst_ptr, dev_src_ptr, total_bytes); - // TODO(hpucha): Make this asynchronous. - Status block_status = device_to_host_stream_->BlockHostUntilDone(); - if (!block_status.ok()) { - status = xla::InternalError( - "Failed to complete data transfer on stream %p: %s", stream_.get(), - block_status.error_message().c_str()); - } - } + // Transfer manager requires the shape of the shaped buffer to be the same as + // literal shape except for the layout. Set the literal to use xla_tensor's + // shape as it is derived from the cpu_tensor's shape using + // shape_representation_fn_. + xla::MutableBorrowingLiteral literal; + TF_CHECK_OK(HostTensorToMutableBorrowingLiteral( + xla::LayoutUtil::GetWithDefaultLayout( + xla_tensor->shaped_buffer().on_host_shape()), + cpu_tensor, &literal)); - done(status); + TensorReference ref(*device_tensor); + transfer_manager_->TransferLiteralFromDevice( + device_to_host_stream_.get(), xla_tensor->shaped_buffer(), literal, + [ref, xla_tensor, done](xla::Status status) { + done([&]() -> Status { + VLOG(2) << "Transfer from device as literal: " + << xla_tensor->shaped_buffer().ToString(); + return status; + }()); + ref.Unref(); + }); } -void XlaDeviceContext::CopyDeviceTensorToDevice(const Tensor& src_tensor, - Tensor* dst_tensor, - const StatusCallback& done) { - VLOG(2) << "CopyDeviceTensorToDevice " - << reinterpret_cast(src_tensor.tensor_data().data()) - << " " - << reinterpret_cast(dst_tensor->tensor_data().data()); - // Perform memory allocation now, and enqueue the device-to-device transfer. - Status status = [&]() -> Status { - if (src_tensor.NumElements() == 0) { - return Status::OK(); - } - // TODO(jmolloy): We co-opt the device_to_host stream for device to device - // transfers; perhaps we should have a dedicated device to device stream? or - // one per device? - auto device_to_device_stream = stream_; - XlaTensor* xla_src = XlaTensor::FromTensor(&src_tensor); - XlaTensor* xla_dst = XlaTensor::FromTensor(dst_tensor); - CHECK(xla_src && xla_dst) - << "Missing destination tensor for device-to-device copy"; - if (!xla_dst->has_shaped_buffer()) { - TF_ASSIGN_OR_RETURN( - TensorShape shape, - shape_representation_fn_(src_tensor.shape(), src_tensor.dtype())); - TF_RETURN_IF_ERROR( - xla_dst->AllocateShapedBuffer(src_tensor.dtype(), shape, client_, - stream_->parent()->device_ordinal())); - if (stream_ != device_to_device_stream) { - // Initially wait for the compute stream so that memory allocations are - // synchronized. - device_to_device_stream->ThenWaitFor(stream_.get()); - } - } - - if (se::Event* event = - xla_src->GetDefinitionEvent(device_to_device_stream.get())) { - device_to_device_stream->ThenWaitFor(event); - xla_src->SetDefinedOn(device_to_device_stream.get()); - } - - auto from_iter = xla_src->shaped_buffer().buffers().begin(); - auto to_iter = xla_dst->shaped_buffer().buffers().begin(); - for (auto end_iter = xla_src->shaped_buffer().buffers().end(); - from_iter != end_iter; ++from_iter, ++to_iter) { - device_to_device_stream->ThenMemcpyD2D( - &to_iter->second, from_iter->second, to_iter->second.size()); - } - - if (UseMultipleStreams()) { - auto event = std::make_shared(stream_->parent()); - TF_RET_CHECK(event->Init()) << "Event failed to initialize"; - device_to_device_stream->ThenRecordEvent(event.get()); - xla_dst->SetDefinedOn(device_to_device_stream.get(), std::move(event)); - } - return Status::OK(); - }(); - if (!status.ok()) { - return done(status); - } else { - stream_->ThenDoHostCallback([this, done]() { - // We must not call the done closure directly from DoHostCallback to avoid - // a deadlock. If done() is the callback that ends an Executor's run, the - // Executor may call XlaDevice::Sync() inside the callback. This - // deadlocks, because XlaDevice::Sync() waits for all stream activity to - // complete. - thread_pool_->Schedule([done]() { done(Status::OK()); }); - }); - } +se::Stream* XlaDeviceContext::GetDeviceToDeviceStream() { + DCHECK_GT(device_to_device_streams_.size(), 0); + absl::MutexLock lock(&mu_); + int stream = next_stream_; + next_stream_ = (next_stream_ + 1) % device_to_device_streams_.size(); + return device_to_device_stream(stream); } } // namespace tensorflow diff --git a/tensorflow/compiler/jit/xla_device_context.h b/tensorflow/compiler/jit/xla_device_context.h index babb60acb5ca547d47825022003b296b1e5d0324..e45db989fac720df6c3458c93a6b8dbb0919f930 100644 --- a/tensorflow/compiler/jit/xla_device_context.h +++ b/tensorflow/compiler/jit/xla_device_context.h @@ -18,6 +18,7 @@ limitations under the License. #include +#include "absl/synchronization/mutex.h" #include "tensorflow/compiler/jit/xla_tensor.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" #include "tensorflow/compiler/xla/client/global_data.h" @@ -50,7 +51,8 @@ class XlaDeviceContext : public DeviceContext { std::shared_ptr compute_stream, std::shared_ptr host_to_device_stream, std::shared_ptr device_to_host_stream, - xla::LocalClient* client, bool transfer_as_literal, + std::vector> device_to_device_streams, + xla::LocalClient* client, XlaCompiler::ShapeRepresentationFn shape_representation_fn, thread::ThreadPool* thread_pool); @@ -60,18 +62,30 @@ class XlaDeviceContext : public DeviceContext { void CopyDeviceTensorToCPU(const Tensor* device_tensor, absl::string_view tensor_name, Device* device, Tensor* cpu_tensor, StatusCallback done) override; + void CopyTensorInSameDevice(const Tensor* input_tensor, Device* device, + Tensor* output_tensor, + StatusCallback done) const override; - void CopyDeviceTensorToDevice(const Tensor& src_tensor, Tensor* dst_tensor, - const StatusCallback& done); - + xla::LocalClient* client() const { return client_; } se::Stream* stream() const { return stream_.get(); } + se::Stream* host_to_device_stream() const { + return host_to_device_stream_.get(); + } + se::Stream* device_to_host_stream() const { + return device_to_host_stream_.get(); + } + se::Stream* device_to_device_stream(int index) const { + return device_to_device_streams_.at(index).get(); + } + xla::TransferManager* transfer_manager() const { return transfer_manager_; } + const XlaCompiler::ShapeRepresentationFn& shape_representation_fn() const { + return shape_representation_fn_; + } + + // Returns a device-to-device stream, in round-robin fashion. + se::Stream* GetDeviceToDeviceStream(); private: - Status TransferLiteralToDevice(const Tensor& host_tensor, - Tensor* device_tensor) const; - void TransferLiteralFromDevice(Tensor* host_tensor, - const Tensor& device_tensor, - const StatusCallback& done) const; bool UseMultipleStreams() const { return stream_ != host_to_device_stream_; } // The main compute stream of the device, used to synchronize the transfer @@ -83,16 +97,22 @@ class XlaDeviceContext : public DeviceContext { // The stream to use for transferring data from device to host. Can be // idential to stream_, but must not be nullptr. std::shared_ptr device_to_host_stream_; + // Streams to use for transferring data directly between different devices, + // e.g., over NVLINK. + std::vector> device_to_device_streams_; + // For the underlying memory allocator and XLA's TransferManager. xla::LocalClient* client_; // Transfer manager, for marshalling data to and from the device. xla::TransferManager* transfer_manager_; - // True if we must use XLA's TransferManager for correct device transfers. - const bool transfer_as_literal_; + XlaCompiler::ShapeRepresentationFn shape_representation_fn_; // Thread pool used for running closures thread::ThreadPool* thread_pool_; + + absl::Mutex mu_; + int next_stream_ GUARDED_BY(mu_) = 0; }; } // namespace tensorflow diff --git a/tensorflow/compiler/jit/xla_device_ops.cc b/tensorflow/compiler/jit/xla_device_ops.cc index 5ecb1afa7bcec910ca843ccd3a782745f2bb6ca8..f56c26ba0103fed152322f0c8971a449610cdc2b 100644 --- a/tensorflow/compiler/jit/xla_device_ops.cc +++ b/tensorflow/compiler/jit/xla_device_ops.cc @@ -30,81 +30,43 @@ void XlaDeviceDummyOp::Compute(OpKernelContext* ctx) { } XlaAssignVariableOp::XlaAssignVariableOp(OpKernelConstruction* c) - : AsyncOpKernel(c) { + : OpKernel(c) { OP_REQUIRES_OK(c, c->GetAttr("dtype", &dtype_)); } -void XlaAssignVariableOp::ComputeAsync(OpKernelContext* context, - DoneCallback done) { - OP_REQUIRES_ASYNC(context, dtype_ == context->input(1).dtype(), - errors::InvalidArgument( - "Variable and value dtypes don't match; respectively, ", - dtype_, " and ", context->input(1).dtype()), - done); +void XlaAssignVariableOp::Compute(OpKernelContext* context) { + OP_REQUIRES(context, dtype_ == context->input(1).dtype(), + errors::InvalidArgument( + "Variable and value dtypes don't match; respectively, ", + DataTypeString(dtype_), " and ", + DataTypeString(context->input(1).dtype()))); Var* variable = nullptr; - OP_REQUIRES_OK_ASYNC( - context, - LookupOrCreateResource( - context, HandleFromInput(context, 0), &variable, - [this, context](Var** ptr) { - *ptr = new Var(dtype_); - PersistentTensor unused; - Tensor* tmp; - AllocatorAttributes attr; - TF_RETURN_IF_ERROR(context->allocate_persistent( - dtype_, context->input(1).shape(), &unused, &tmp, attr)); - *(*ptr)->tensor() = *tmp; - return Status::OK(); - }), - done); - core::ScopedUnref s(variable); - - OP_REQUIRES_ASYNC(context, variable->tensor()->dtype() == dtype_, - errors::InvalidArgument( - "Trying to assign variable with wrong dtype. Expected ", - DataTypeString(variable->tensor()->dtype()), " got ", - DataTypeString(dtype_)), - done); - const Tensor& value = context->input(1); - AllocatorAttributes attr; - - // Copying is unnecessary if we are the last user of the value tensor, we can - // just adopt the input tensor's buffer instead. - std::unique_ptr input_alias = context->forward_input( - 1, /*output_index=*/OpKernelContext::Params::kNoReservation, dtype_, - value.shape(), DEVICE_MEMORY, attr); + // Note: every resource-variable-manipulating op assumes copy-on-write + // semantics, and creates a copy of the variable's Tensor if its refcount is + // bigger than 1 when we try to modify it. This means we never need to copy + // the original tensor for AssignVariableOp; even if there are other live + // users of it we know none can modify it so this is always safe (even in + // esoteric cases where the same tensor is used to initialize multiple + // variables or the tensor is a constant this is safe, as future writes will + // trigger copies). + OP_REQUIRES_OK(context, LookupOrCreateResource( + context, HandleFromInput(context, 0), &variable, + [this, &value](Var** ptr) { + *ptr = new Var(dtype_); + *(*ptr)->tensor() = value; + (*ptr)->is_initialized = true; + return Status::OK(); + })); + core::ScopedUnref s(variable); mutex_lock ml(*variable->mu()); + OP_REQUIRES(context, variable->tensor()->dtype() == dtype_, + errors::InvalidArgument( + "Trying to assign variable with wrong dtype. Expected ", + DataTypeString(variable->tensor()->dtype()), " got ", + DataTypeString(dtype_))); variable->is_initialized = true; - if (input_alias) { - *variable->tensor() = *input_alias; - done(); - return; - } - - // Need to copy, but maybe we can re-use variable's buffer? - if (!XlaTensor::RefCountIsOne(*variable->tensor()) || - !variable->tensor()->shape().IsSameSize(value.shape())) { - // Copy to new buffer - PersistentTensor unused; - Tensor* tmp; - OP_REQUIRES_OK_ASYNC(context, - context->allocate_persistent(dtype_, value.shape(), - &unused, &tmp, attr), - done); - *variable->tensor() = *tmp; - } - - XlaDeviceContext* device_context = - static_cast(context->op_device_context()); - - variable->Ref(); - device_context->CopyDeviceTensorToDevice( - value, variable->tensor(), [context, variable, done](Status status) { - variable->Unref(); - context->SetStatus(status); - done(); - }); + *variable->tensor() = value; } } // namespace tensorflow diff --git a/tensorflow/compiler/jit/xla_device_ops.h b/tensorflow/compiler/jit/xla_device_ops.h index 6a1c43aa96026a991c8a8d016d67b5ca048c293c..927f983ba9ef23c8509523f42366c0c89c29db9f 100644 --- a/tensorflow/compiler/jit/xla_device_ops.h +++ b/tensorflow/compiler/jit/xla_device_ops.h @@ -35,6 +35,7 @@ limitations under the License. #include "tensorflow/core/kernels/resource_variable_ops.h" #include "tensorflow/core/kernels/sendrecv_ops.h" #include "tensorflow/core/kernels/shape_ops.h" +#include "tensorflow/core/kernels/stack.h" #include "tensorflow/core/kernels/variable_ops.h" namespace tensorflow { @@ -49,10 +50,10 @@ class XlaDeviceDummyOp : public OpKernel { void Compute(OpKernelContext* ctx) override; }; -class XlaAssignVariableOp : public AsyncOpKernel { +class XlaAssignVariableOp : public OpKernel { public: explicit XlaAssignVariableOp(OpKernelConstruction* c); - void ComputeAsync(OpKernelContext* context, DoneCallback done) override; + void Compute(OpKernelContext* context) override; private: DataType dtype_; @@ -93,6 +94,9 @@ class XlaAssignVariableOp : public AsyncOpKernel { ConstantOp); \ REGISTER_KERNEL_BUILDER( \ Name("Identity").Device(DEVICE).TypeConstraint("T", TYPES), IdentityOp); \ + REGISTER_KERNEL_BUILDER( \ + Name("Identity").Device(DEVICE).TypeConstraint("T", DT_STRING), \ + IdentityOp); \ REGISTER_KERNEL_BUILDER(Name("IdentityN").Device(DEVICE), IdentityNOp); \ REGISTER_KERNEL_BUILDER(Name("Placeholder").Device(DEVICE), PlaceholderOp); \ REGISTER_KERNEL_BUILDER(Name("PlaceholderV2").Device(DEVICE), \ @@ -199,6 +203,8 @@ class XlaAssignVariableOp : public AsyncOpKernel { .HostMemory("output") \ .TypeConstraint("T"), \ ArgOp); \ + REGISTER_KERNEL_BUILDER( \ + Name(kArgOp).Device(DEVICE).TypeConstraint("T"), ArgOp); \ \ REGISTER_KERNEL_BUILDER(Name(kRetOp) \ .Device(DEVICE) \ @@ -254,9 +260,27 @@ class XlaAssignVariableOp : public AsyncOpKernel { .Device(DEVICE) \ .TypeConstraint("T") \ .HostMemory("input"), \ - RetvalOp); + RetvalOp); \ + \ + REGISTER_KERNEL_BUILDER(Name("StackV2") \ + .Device(DEVICE) \ + .HostMemory("max_size") \ + .HostMemory("handle"), \ + StackOp); \ + REGISTER_KERNEL_BUILDER(Name("StackPushV2") \ + .Device(DEVICE) \ + .HostMemory("handle") \ + .TypeConstraint("T", TYPES), \ + TemplatedStackPushOp); \ + REGISTER_KERNEL_BUILDER(Name("StackPopV2") \ + .Device(DEVICE) \ + .HostMemory("handle") \ + .TypeConstraint("elem_type", TYPES), \ + StackPopOp); \ + REGISTER_KERNEL_BUILDER( \ + Name("StackCloseV2").Device(DEVICE).HostMemory("handle"), StackCloseOp); -// TODO(phawkins): currently we do not register the QueueEnqueueMany, +// TODO(b/118881356): currently we do not register the QueueEnqueueMany, // QueueDequeueMany, or QueueDequeueUpTo kernels because they attempt to read // and write the tensors they access in order to concatenate them into a batch. // We would need either to call out to an XLA computation to perform the diff --git a/tensorflow/compiler/jit/xla_gpu_device.cc b/tensorflow/compiler/jit/xla_gpu_device.cc index d9021fb001abda9559fbc337c75e8f5ab3972197..b29f6a009b9e9fdba76ac55386a4bec2f339cc0e 100644 --- a/tensorflow/compiler/jit/xla_gpu_device.cc +++ b/tensorflow/compiler/jit/xla_gpu_device.cc @@ -16,7 +16,10 @@ limitations under the License. // Registers the XLA_GPU device, which is an XlaDevice instantiation that runs // operators using XLA via the XLA "CUDA" (GPU) backend. +#include #include "absl/memory/memory.h" +#include "absl/strings/numbers.h" +#include "absl/strings/str_split.h" #include "tensorflow/compiler/jit/kernels/xla_ops.h" #include "tensorflow/compiler/jit/xla_device.h" #include "tensorflow/compiler/jit/xla_device_ops.h" @@ -26,19 +29,43 @@ limitations under the License. namespace tensorflow { +// Returns a set containing the device ids contained in visible_device_list or +// nullopt if it is empty. It returns error in case of malformed configuration +// string. +static xla::StatusOr>> ParseVisibleDeviceList( + const string& visible_device_list) { + std::set gpu_ids; + if (visible_device_list.empty()) { + return {{absl::nullopt}}; + } + const std::vector visible_devices = + absl::StrSplit(visible_device_list, ','); + for (const string& platform_gpu_id_str : visible_devices) { + int32 platform_gpu_id; + if (!absl::SimpleAtoi(platform_gpu_id_str, &platform_gpu_id)) { + return errors::InvalidArgument( + "Could not parse entry in 'visible_device_list': '", + platform_gpu_id_str, + "'. visible_device_list = ", visible_device_list); + } + gpu_ids.insert(platform_gpu_id); + } + return {{gpu_ids}}; +} + class XlaGpuDeviceFactory : public DeviceFactory { public: Status CreateDevices(const SessionOptions& options, const string& name_prefix, - std::vector* devices) override; + std::vector>* devices) override; }; -Status XlaGpuDeviceFactory::CreateDevices(const SessionOptions& session_options, - const string& name_prefix, - std::vector* devices) { +Status XlaGpuDeviceFactory::CreateDevices( + const SessionOptions& session_options, const string& name_prefix, + std::vector>* devices) { XlaOpRegistry::DeviceRegistration registration; registration.compilation_device_name = DEVICE_GPU_XLA_JIT; - registration.requires_compilation = true; - registration.enable_jit_by_default = false; + registration.autoclustering_policy = + XlaOpRegistry::AutoclusteringPolicy::kAlways; registration.compile_resource_ops = true; XlaOpRegistry::RegisterCompilationDevice(DEVICE_XLA_GPU, registration); @@ -52,26 +79,37 @@ Status XlaGpuDeviceFactory::CreateDevices(const SessionOptions& session_options, VLOG(1) << "Failed to create XLA_GPU device: " << platform.status(); return Status::OK(); } - - XlaDevice::Options options; - options.platform = platform.ValueOrDie(); - options.device_name_prefix = name_prefix; - options.device_name = DEVICE_XLA_GPU; - options.device_ordinal = 0; - options.compilation_device_name = DEVICE_GPU_XLA_JIT; - options.transfer_as_literal = false; - options.use_multiple_streams = false; - auto device = absl::make_unique(session_options, options); - - // TODO(b/78468222): Uncomment after fixing this bug - // status = device->UseGpuDeviceInfo(); - // if (!status.ok()) { - // errors::AppendToMessage(&status, "while setting up ", DEVICE_GPU_XLA_JIT, - // " device"); - // return status; - // } - - devices->push_back(device.release()); + string allowed_gpus = + session_options.config.gpu_options().visible_device_list(); + absl::optional> gpu_ids = + ParseVisibleDeviceList(allowed_gpus).ValueOrDie(); + if (!gpu_ids) { + gpu_ids.emplace(); + // Fill the gpu_ids set with all devices if config string is empty. + for (int i = 0; i < platform.ValueOrDie()->VisibleDeviceCount(); ++i) { + gpu_ids->insert(i); + } + } + for (int i : *gpu_ids) { + XlaDevice::Options options; + options.platform = platform.ValueOrDie(); + options.device_name_prefix = name_prefix; + options.device_name = DEVICE_XLA_GPU; + options.device_ordinal = i; + options.compilation_device_name = DEVICE_GPU_XLA_JIT; + options.use_multiple_streams = true; + options.allowed_devices = gpu_ids; + auto device = absl::make_unique(session_options, options); + + Status status = device->UseGpuDeviceInfo(); + if (!status.ok()) { + errors::AppendToMessage(&status, "while setting up ", DEVICE_GPU_XLA_JIT, + " device number ", i); + return status; + } + + devices->push_back(std::move(device)); + } return Status::OK(); } diff --git a/tensorflow/compiler/jit/xla_interpreter_device.cc b/tensorflow/compiler/jit/xla_interpreter_device.cc index aee3b58c997ff032527340332fbc109965f71889..e1a582406153d2af447fa9d4ebcaf0bf0842b132 100644 --- a/tensorflow/compiler/jit/xla_interpreter_device.cc +++ b/tensorflow/compiler/jit/xla_interpreter_device.cc @@ -26,27 +26,27 @@ 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: Status CreateDevices(const SessionOptions& options, const string& name_prefix, - std::vector* devices) override; + std::vector>* devices) override; }; Status XlaInterpreterDeviceFactory::CreateDevices( const SessionOptions& session_options, const string& name_prefix, - std::vector* devices) { + std::vector>* devices) { static XlaDeviceOpRegistrations* registrations = RegisterXlaDeviceKernels( DEVICE_XLA_INTERPRETER, DEVICE_INTERPRETER_XLA_JIT); (void)registrations; XlaOpRegistry::DeviceRegistration registration; registration.compilation_device_name = DEVICE_INTERPRETER_XLA_JIT; - registration.requires_compilation = true; - registration.enable_jit_by_default = false; + registration.autoclustering_policy = + XlaOpRegistry::AutoclusteringPolicy::kAlways; registration.compile_resource_ops = true; XlaOpRegistry::RegisterCompilationDevice(DEVICE_XLA_INTERPRETER, registration); @@ -60,10 +60,8 @@ Status XlaInterpreterDeviceFactory::CreateDevices( options.device_name = DEVICE_XLA_INTERPRETER; options.device_ordinal = 0; options.compilation_device_name = DEVICE_INTERPRETER_XLA_JIT; - options.transfer_as_literal = false; options.use_multiple_streams = false; - auto device = absl::make_unique(session_options, options); - devices->push_back(device.release()); + devices->push_back(absl::make_unique(session_options, options)); return Status::OK(); } diff --git a/tensorflow/compiler/jit/xla_launch_util.cc b/tensorflow/compiler/jit/xla_launch_util.cc index 0e8ee56ed8979111b66e3886f07994c8b665c388..c64981053fad2dbf1e8bcd623a940ded8b4d9150 100644 --- a/tensorflow/compiler/jit/xla_launch_util.cc +++ b/tensorflow/compiler/jit/xla_launch_util.cc @@ -17,6 +17,7 @@ limitations under the License. #include +#include "absl/algorithm/container.h" #include "absl/memory/memory.h" #include "tensorflow/compiler/jit/defs.h" #include "tensorflow/compiler/tf2xla/shape_util.h" @@ -41,22 +42,127 @@ using xla::ScopedShapedBuffer; using xla::ShapedBuffer; } // anonymous namespace -std::map SnapshotResourceVariables( - OpKernelContext* ctx, absl::Span variables) { - std::map snapshot; - for (int i : variables) { - Var* variable = nullptr; - ResourceHandle handle = HandleFromInput(ctx, i); - OptionalTensor& tensor = snapshot[i]; - if (LookupResource(ctx, handle, &variable).ok()) { - core::ScopedUnref scoped_unref(variable); - tf_shared_lock lock(*variable->mu()); - tensor.name = handle.name(); +VariableInfo::VariableInfo(int index, Var* var) : index_(index), var_(var) {} +VariableInfo::VariableInfo(VariableInfo&& other) + : index_(other.index_), var_(other.var_), lock_held_(other.lock_held_) { + other.index_ = -1; + other.var_ = nullptr; +} + +VariableInfo& VariableInfo::operator=(VariableInfo&& other) { + index_ = other.index_; + var_ = other.var_; + lock_held_ = other.lock_held_; + + other.index_ = -1; + other.var_ = nullptr; + + return *this; +} + +VariableInfo::~VariableInfo() { + // Release the variable's lock if we hold it. Ensures that the lock is + // released even on error. It does not matter in what order we release the + // locks. + if (var()) { + if (lock_held()) { + var()->mu()->unlock(); + } + + // Unref the variable so it can be released by ResourceManager. + var()->Unref(); + } +} + +// Returns a vector of VaribleInfo instances for the resource variable inputs to +// the kernel with context `ctx`. The input indices for the resource variable +// inputs are in `variable_indices`. +static Status GetVariableInfosFromCtxInputs( + OpKernelContext* ctx, absl::Span variable_indices, + std::vector* result) { + std::vector resource_handles; + absl::c_transform( + variable_indices, std::back_inserter(resource_handles), + [&](int variable_idx) { return &HandleFromInput(ctx, variable_idx); }); + + std::vector> variables; + TF_RETURN_IF_ERROR(LookupResources(ctx, resource_handles, &variables)); + + result->clear(); + result->reserve(variable_indices.size()); + for (int i = 0; i < variable_indices.size(); i++) { + // *Release* the variable because we're going to unref it later in + // ~VariableInfo. + Var* variable = variables[i].release(); + result->emplace_back(variable_indices[i], variable); + } + + return Status::OK(); +} + +Status LockVariables(absl::Span variables) { + std::vector lock_order(variables.size()); + std::iota(lock_order.begin(), lock_order.end(), 0); + + // VariableInfoComparator orders all empty VariableInfo instances as + // equivalent so it looks like we may want to stable sort these to maintain a + // deterministic order between the empty VariableInfo instances. However + // since we're sorting by pointer value the sort is pretty non-deterministic + // anyway so we don't bother using std::stable_sort for now. + absl::c_sort(lock_order, [&](int a, int b) { + if (variables[a].var() && variables[b].var()) { + return variables[a].var()->mu() < variables[b].var()->mu(); + } + + // Move all the empty VariableInfo instances to the end. + return variables[a].var() != nullptr; + }); + + mutex* prev = nullptr; + for (int i : lock_order) { + Var* variable = variables[i].var(); + if (variable == nullptr) { + // All empty VariableInfo instances are at the end of the order + // so we're done. + break; + } + mutex* mu = variable->mu(); + if (prev == mu) { + // It is an error to pass the same variable handle twice to the same XLA + // cluster because we would not handle variable updates correctly. Any + // locks we have already acquired will be released when the VariableInfo + // objects are destroyed. + return errors::Internal("Duplicate variable passed to XLA cluster"); + } + VLOG(4) << "Acquiring lock for variable " + << reinterpret_cast(variable); + mu->lock(); + variables[i].set_lock_held(); + prev = mu; + } + VLOG(4) << "Finished acquiring variable locks."; + return Status::OK(); +} + +Status SnapshotResourceVariables(OpKernelContext* ctx, + absl::Span variable_indices, + std::map* result) { + std::vector variable_infos; + TF_RETURN_IF_ERROR( + GetVariableInfosFromCtxInputs(ctx, variable_indices, &variable_infos)); + TF_RETURN_IF_ERROR(LockVariables(absl::MakeSpan(variable_infos))); + + for (int i = 0; i < variable_indices.size(); i++) { + if (variable_infos[i].var()) { + OptionalTensor& tensor = (*result)[variable_indices[i]]; + tensor.name = HandleFromInput(ctx, variable_indices[i]).name(); tensor.present = true; - tensor.value = *variable->tensor(); + tensor.value = *variable_infos[i].var()->tensor(); + } else { + (*result)[variable_indices[i]] = OptionalTensor(); } } - return snapshot; + return Status::OK(); } XlaAllocator::XlaAllocator(const se::Platform* platform, Allocator* wrapped) @@ -85,40 +191,6 @@ Status XlaAllocator::Deallocate(int device_ordinal, se::DeviceMemoryBase mem) { return Status::OK(); } -namespace internal { -// Return the 'index''th subtree of the given ShapedBuffer as a -// ScopedShapedBuffer. The returned ScopedShapedBuffer takes ownership of the -// subtree, and sets the input's buffer pointers to nullptr for the subtree. -ScopedShapedBuffer ExtractSubShapedBuffer( - ShapedBuffer* shaped_buffer, int index, - xla::DeviceMemoryAllocator* allocator) { - const xla::Shape& on_host_shape = xla::ShapeUtil::GetTupleElementShape( - shaped_buffer->on_host_shape(), index); - const xla::Shape& on_device_shape = xla::ShapeUtil::GetTupleElementShape( - shaped_buffer->on_device_shape(), index); - - ShapedBuffer sub_shaped_buffer(on_host_shape, on_device_shape, - shaped_buffer->platform(), - shaped_buffer->device_ordinal()); - - auto& shape_tree = shaped_buffer->buffers(); - auto& sub_shape_tree = sub_shaped_buffer.buffers(); - sub_shape_tree.CopySubtreeFrom(shape_tree, - /*source_base_index=*/{index}, - /*target_base_index=*/{}); - shape_tree.ForEachMutableElement( - [index](const xla::ShapeIndex& shape_index, - tensorflow::se::DeviceMemoryBase* data) { - // shape_index is empty for the root node. Ignore that. - if (!shape_index.empty() && shape_index[0] == index) { - *data = tensorflow::se::DeviceMemoryBase(nullptr, 0); - } - }); - return ScopedShapedBuffer(std::move(sub_shaped_buffer), allocator); -} -} // namespace internal -using internal::ExtractSubShapedBuffer; - XlaComputationLaunchContext::XlaComputationLaunchContext( xla::LocalClient* client, xla::DeviceMemoryAllocator* xla_allocator, bool allocate_xla_tensors, bool use_multiple_streams) @@ -160,15 +232,12 @@ void XlaComputationLaunchContext::PopulateInputs( CHECK(stream) << "Must have a stream available when using XLA tensors!"; XlaTensor* xla_tensor = XlaTensor::FromTensor(t); CHECK(xla_tensor); - if (se::Event* event = xla_tensor->GetDefinitionEvent(stream)) { - stream->ThenWaitFor(event); - xla_tensor->SetDefinedOn(stream); - } + xla_tensor->WaitForDefinitionEventOnStream(stream); } const xla::Shape on_device_shape = client_->backend().transfer_manager()->HostShapeToDeviceShape(shape); - if (xla::ShapeUtil::IsTuple(on_device_shape)) { + if (on_device_shape.IsTuple()) { const XlaTensor* xla_tensor = XlaTensor::FromTensor(t); CHECK(xla_tensor && xla_tensor->has_shaped_buffer()); arg_ptrs_[i] = const_cast(&xla_tensor->shaped_buffer()); @@ -205,7 +274,7 @@ Status XlaComputationLaunchContext::PopulateOutputs( // If the on-host-shape isn't a tuple, create a new single-element tuple // buffer with a nullptr root index table. This allows the code below to treat // output as a tuple unconditionally. - if (!xla::ShapeUtil::IsTuple(output.on_host_shape())) { + if (!output.on_host_shape().IsTuple()) { ShapedBuffer nontuple_buffer = output.release(); ShapedBuffer buffer( xla::ShapeUtil::MakeTupleShape({nontuple_buffer.on_host_shape()}), @@ -239,7 +308,7 @@ Status XlaComputationLaunchContext::PopulateOutputs( // Copy host -> device. (Empty tensors don't have backing buffers.) // Manually allocate memory using an XlaTensorBuffer so we can allocate // as much memory as the device requires (as given by - // GetByteSizeRequirement). This avoids XlaDeviceContext having to + // GetByteSizeRequirement). This avoids XlaTransferManager having to // reallocate the device buffer later. VLOG(1) << "Constant output tensor on device"; @@ -288,10 +357,9 @@ Status XlaComputationLaunchContext::PopulateOutputs( TF_RETURN_IF_ERROR(ctx->allocate_output(i, shape, &output_tensor)); XlaTensor* xla_tensor = XlaTensor::FromTensor(output_tensor); if (xla_tensor) { - xla_tensor->set_shaped_buffer(ScopedShapedBuffer( - ExtractSubShapedBuffer(&output, output_num, xla_allocator_))); + xla_tensor->set_shaped_buffer(output.TakeSubTree({output_num})); if (use_multiple_streams_) { - xla_tensor->SetDefinedOn(stream, definition_event); + xla_tensor->ResetDefinitionEvent(definition_event, stream); } } else { // xla_tensor wasn't valid, which must mean this is a zero-element @@ -309,36 +377,41 @@ Status XlaComputationLaunchContext::PopulateOutputs( } if (VLOG_IS_ON(3)) { - VLOG(3) << ctx->mutable_output(i)->DebugString(); + VLOG(3) << ctx->mutable_output(i)->DeviceSafeDebugString(); } } // Apply variable updates, if any. VLOG(2) << "Applying variable updates"; + std::vector variable_infos; + variable_infos.reserve(kernel->resource_updates.size()); + for (int i = 0; i < kernel->resource_updates.size(); ++i) { - Allocator* allocator = ctx->device()->GetAllocator({}); const XlaCompiler::ResourceUpdate& write = kernel->resource_updates[i]; int actual_input_index = write.input_index - missing_ctx_input_prefix; if (actual_input_index < 0 || actual_input_index >= ctx->num_inputs()) { return errors::Internal("Invalid input index for variable write."); } - se::DeviceMemoryBase buffer = output.buffer({output_num}); - - Var* variable = nullptr; // TODO(b/35625933): tensorflow::Var should contain a PersistentTensor, // not a Tensor. + Var* variable = nullptr; TF_RETURN_IF_ERROR(LookupOrCreateResource( ctx, HandleFromInput(ctx, actual_input_index), &variable, [&write](Var** ptr) { *ptr = new Var(write.type); return Status::OK(); })); + variable_infos.emplace_back(actual_input_index, variable); + } - core::ScopedUnref s(variable); + TF_RETURN_IF_ERROR(LockVariables(absl::MakeSpan(variable_infos))); - mutex_lock ml(*variable->mu()); - if (variable->tensor()->dtype() != write.type) { + for (int i = 0; i < kernel->resource_updates.size(); ++i) { + Allocator* allocator = ctx->device()->GetAllocator({}); + const XlaCompiler::ResourceUpdate& write = kernel->resource_updates[i]; + + if (variable_infos[i].var()->tensor()->dtype() != write.type) { return errors::Internal("Mismatched type in variable write"); } @@ -346,23 +419,81 @@ Status XlaComputationLaunchContext::PopulateOutputs( Tensor output_tensor; TF_RETURN_IF_ERROR( ctx->allocate_temp(write.type, write.shape, &output_tensor)); - XlaTensor* xla_tensor = XlaTensor::FromTensor(&output_tensor); - CHECK(xla_tensor); - xla_tensor->set_shaped_buffer( - ExtractSubShapedBuffer(&output, output_num, xla_allocator_)); - if (use_multiple_streams_) { - xla_tensor->SetDefinedOn(stream, definition_event); + if (write.shape.num_elements() > 0) { + XlaTensor* xla_tensor = XlaTensor::FromTensor(&output_tensor); + CHECK(xla_tensor); + xla_tensor->set_shaped_buffer(output.TakeSubTree({output_num})); + if (use_multiple_streams_) { + xla_tensor->ResetDefinitionEvent(definition_event, stream); + } } - *variable->tensor() = output_tensor; + *variable_infos[i].var()->tensor() = output_tensor; } else { + se::DeviceMemoryBase buffer = output.buffer({output_num}); + output.set_buffer(xla::OwningDeviceMemory(), {output_num}); Tensor output_tensor = XlaTensorBuffer::MakeTensor( write.type, write.shape, buffer, allocator); - output.set_buffer(xla::OwningDeviceMemory(), {output_num}); - *variable->tensor() = output_tensor; + *variable_infos[i].var()->tensor() = output_tensor; } ++output_num; } return Status::OK(); } +Status XlaComputationLaunchContext::BuildXlaCompilerArguments( + const std::map& constant_args, + const std::map& variable_args, OpKernelContext* ctx, + std::vector* args) { + args->resize(ctx->num_inputs()); + + for (int64 input_num = 0; input_num < ctx->num_inputs(); ++input_num) { + XlaCompiler::Argument& arg = (*args)[input_num]; + if (constant_args.count(input_num) > 0) { + // Handles compile-time constants. + const Tensor& input = constant_args.at(input_num); + TF_RET_CHECK(input.dtype() != DT_RESOURCE); + arg.kind = XlaCompiler::Argument::kConstant; + arg.type = input.dtype(); + arg.shape = input.shape(); + arg.constant_value = input; + } else if (variable_args.count(input_num) == 0) { + // Handles the non-constant arguments. + const Tensor& input = ctx->input(input_num); + TF_RET_CHECK(input.dtype() != DT_RESOURCE); + if (input.NumElements() > 0) { + arg.kind = XlaCompiler::Argument::kParameter; + } else { + arg.kind = XlaCompiler::Argument::kConstant; + arg.constant_value = input; + } + arg.type = input.dtype(); + arg.shape = input.shape(); + } else { + // Handles resource variables. + const Tensor& input = ctx->input(input_num); + TF_RET_CHECK(input.dtype() == DT_RESOURCE); + const OptionalTensor& variable = variable_args.at(input_num); + arg.name = variable.name; + arg.kind = XlaCompiler::Argument::kResource; + arg.resource_kind = XlaResource::kVariable; + if (variable.present) { + const Tensor& value = variable.value; + arg.type = value.dtype(); + arg.shape = value.shape(); + arg.initialized = true; + } else { + // The values of uninitialized variables are not passed as inputs, since + // they are meaningless. However, it is legal to assign to a resource + // variable for the first time inside the XLA computation, so we do + // permit uninitialized variables. + arg.initialized = false; + arg.type = DT_INVALID; + arg.shape = TensorShape(); + } + } + } + + return Status::OK(); +} + } // namespace tensorflow diff --git a/tensorflow/compiler/jit/xla_launch_util.h b/tensorflow/compiler/jit/xla_launch_util.h index 326d70a027564343408df356833c97e131495da0..554227f09de0ab4d9e07f199b957657f3121ff06 100644 --- a/tensorflow/compiler/jit/xla_launch_util.h +++ b/tensorflow/compiler/jit/xla_launch_util.h @@ -18,6 +18,7 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_JIT_XLA_LAUNCH_UTIL_H_ #define TENSORFLOW_COMPILER_JIT_XLA_LAUNCH_UTIL_H_ +#include "absl/base/thread_annotations.h" #include "tensorflow/compiler/jit/xla_compilation_cache.h" #include "tensorflow/compiler/jit/xla_tensor.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" @@ -34,17 +35,75 @@ limitations under the License. namespace tensorflow { class XlaAllocator; -// Takes a snapshot of the values of resource variable arguments, whose -// indices are specified in `variables` argument. We snapshot tensors that back +// Struct that represents a possibly-absent Tensor. +struct OptionalTensor { + string name; // A descriptive name + bool present = false; // Is the tensor present? + Tensor value; // If present, what is the Tensor's value? +}; + +// Takes a snapshot of the values of resource variable arguments, whose indices +// are specified in `variable_indices` argument. We snapshot tensors that back // resource variables since concurrent updates may modify the shape, and it is // important that the shapes used for compilation match the true shapes of the // buffers. // +// We snapshot the entire set of resource variables as one atomic operation. +// This models Read->* dependencies between resource variable operations. See +// jit/resource_operation_safety_analysis for details. +// // Returns a map of TensorFlow argument index to resource variable. If a // resource variable is not initialized, the corresponding OptionalTensor // will have its `present` field set to false. -std::map SnapshotResourceVariables( - OpKernelContext* ctx, absl::Span variables); +Status SnapshotResourceVariables(OpKernelContext* ctx, + absl::Span variable_indices, + std::map* result); + +// Information about the state of a variable passed as input to the _XlaCompile +// and _XlaRun operators. Unlocks the resource variable and decrements its +// refcount on destruction. +class VariableInfo { + public: + explicit VariableInfo(int index, Var* var); + VariableInfo(VariableInfo&& other); + + VariableInfo& operator=(VariableInfo&& other); + + VariableInfo(const VariableInfo&) = delete; + VariableInfo& operator=(const VariableInfo&) = delete; + + // The index of the DT_RESOURCE input to the _XlaCompile/_XlaRun operator. + // Note that the indices can be different between _XlaCompile and _XlaRun. + int index() const { return index_; } + + // A pointer to the resource variable. May be null if this VariableInfo is + // "empty", i.e. it does not track a resource variable. + Var* var() const { return var_; } + + // Returns true if the resource variable lock was successfully acquired by + // this thread. + bool lock_held() const { return lock_held_; } + void set_lock_held() { lock_held_ = true; } + + ~VariableInfo(); + + private: + int index_; + Var* var_; + + // We can't use a optional here because it confuses the compiler's + // thread safety analysis. Instead we use a boolean flag and release the lock + // in the VariableInfo destructor. + bool lock_held_ = false; +}; + +// Acquires the mutexes for all the variables in `variables` using a +// deadlock-safe protocol (acquire the mutexes in increasing-address order). +// +// `variables` is allowed to contain instances that don't track a resource +// variable (i.e. variables[i].var() can be null for some i). +Status LockVariables(absl::Span variables) + EXCLUSIVE_LOCK_FUNCTION(); // Adapter class that wraps a Tensorflow allocator as an XLA allocator. // Assumes that the Tensorflow allocator permits asynchronous deallocation: @@ -87,6 +146,13 @@ class XlaComputationLaunchContext { bool allocate_xla_tensors, bool use_multiple_streams); + // Builds a XlaCompiler::Argument vector from the arguments to an XlaLaunch + // op. + static Status BuildXlaCompilerArguments( + const std::map& constant_args, + const std::map& variable_args, OpKernelContext* ctx, + std::vector* args); + // Add all inputs within `ctx` as XLA arguments (returned by arguments()). // `variables` is a map from TensorFlow argument number to resource variable. // @@ -99,7 +165,13 @@ class XlaComputationLaunchContext { const std::map& variables, int missing_ctx_input_prefix); - // Given the XLA output in `output`, populate all outputs of `ctx`. + // Given the XLA output in `output`, populate all outputs of `ctx`. Also + // writes out the resource variable updates. + // + // Updates to all resource variables are written in a single atomic operation. + // This models *->Write dependencies between resource variable operations. + // See jit/resource_operation_safety_analysis for details. + // // // Assumes that the first `missing_ctx_input_prefix` inputs to the kernel are // missing and adjusts input indices accordingly. @@ -127,19 +199,17 @@ class XlaTensorBuffer : public TensorBuffer { public: XlaTensorBuffer(const void* ptr, size_t expected_size, size_t actual_size, Allocator* allocator) - : expected_size_(expected_size), + : TensorBuffer(const_cast(ptr)), + expected_size_(expected_size), actual_size_(actual_size), - allocator_(allocator) { - data_ = const_cast(ptr); - } + allocator_(allocator) {} ~XlaTensorBuffer() override { - if (data_) { - allocator_->DeallocateRaw(data_); + if (data()) { + allocator_->DeallocateRaw(data()); } } - void* data() const override { return data_; } size_t size() const override { return expected_size_; } TensorBuffer* root_buffer() override { return this; } @@ -159,23 +229,11 @@ class XlaTensorBuffer : public TensorBuffer { } private: - void* data_; size_t expected_size_; size_t actual_size_; Allocator* allocator_; }; -// Exposed in this header file for microbenchmarking purposes, but this is an -// internal implementation detail. -namespace internal { -// Return the 'index''th subtree of the given ShapedBuffer as a -// ScopedShapedBuffer. The returned ScopedShapedBuffer takes ownership of the -// subtree, and sets the input's buffer pointers to nullptr for the subtree. -xla::ScopedShapedBuffer ExtractSubShapedBuffer( - xla::ShapedBuffer* shaped_buffer, int index, - xla::DeviceMemoryAllocator* allocator); -} // namespace internal - } // namespace tensorflow #endif // TENSORFLOW_COMPILER_JIT_XLA_LAUNCH_UTIL_H_ diff --git a/tensorflow/compiler/jit/xla_launch_util_test.cc b/tensorflow/compiler/jit/xla_launch_util_test.cc deleted file mode 100644 index a45932403ec1760d6b985d5357fd6d84fbf257a2..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/jit/xla_launch_util_test.cc +++ /dev/null @@ -1,64 +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. -==============================================================================*/ - -// Contains microbenchmarks for performance critical functions in -// xla_launch_util.cc. - -#include "tensorflow/compiler/jit/xla_launch_util.h" -#include "tensorflow/compiler/tf2xla/shape_util.h" -#include "tensorflow/core/platform/test.h" -#include "tensorflow/core/platform/test_benchmark.h" - -// Test ExtractSubBuffer with different depths (depth of ShapeTree) and fan-outs -// (cardinality of each non-leaf node's children). -void BM_ExtractSubBuffer(int iters, int depth, int fan_out) { - tensorflow::testing::StopTiming(); - xla::Shape shape = xla::ShapeUtil::MakeShape(xla::F32, {32, 64, 128}); - for (int i = 0; i < depth; ++i) { - std::vector shapes(fan_out, shape); - shape = xla::ShapeUtil::MakeTupleShape(shapes); - } - xla::ShapedBuffer shaped_buffer(shape, shape, /*platform=*/nullptr, - /*device_ordinal=*/0); - tensorflow::testing::StartTiming(); - for (int i = 0; i < iters; ++i) { - // Extract a buffer from approximately the middle of the first level of the - // tree. - (void)tensorflow::internal::ExtractSubShapedBuffer(&shaped_buffer, - /*index=*/fan_out / 2, - /*allocator=*/nullptr) - .release(); - } -} - -BENCHMARK(BM_ExtractSubBuffer) - ->ArgPair(1, 4) - ->ArgPair(1, 8) - ->ArgPair(1, 32) - ->ArgPair(1, 64) - ->ArgPair(1, 128) - ->ArgPair(1, 256) - ->ArgPair(1, 512) - ->ArgPair(2, 4) - ->ArgPair(2, 8) - ->ArgPair(2, 32) - ->ArgPair(2, 64) - ->ArgPair(2, 128); - -int main(int argc, char** argv) { - testing::InitGoogleTest(&argc, argv); - tensorflow::testing::RunBenchmarks(); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/jit/xla_tensor.cc b/tensorflow/compiler/jit/xla_tensor.cc index 92ba7de1b7d32fcf693cd12a380d7a1e0d861d71..d1f7f754c8338487557eda512c56be34c9e958b7 100644 --- a/tensorflow/compiler/jit/xla_tensor.cc +++ b/tensorflow/compiler/jit/xla_tensor.cc @@ -43,11 +43,10 @@ namespace tensorflow { } } -Status XlaTensor::AllocateShapedBuffer(DataType dtype, const TensorShape& shape, +Status XlaTensor::AllocateShapedBuffer(DataType dtype, + const xla::Shape& on_host_shape, xla::LocalClient* client, int device_ordinal) { - xla::Shape on_host_shape; - TF_RETURN_IF_ERROR(TensorShapeToXLAShape(dtype, shape, &on_host_shape)); xla::Shape on_device_shape = client->backend().transfer_manager()->HostShapeToDeviceShape( on_host_shape); @@ -73,10 +72,10 @@ Status XlaTensor::AllocateShapedBuffer(DataType dtype, const TensorShape& shape, return Status::OK(); } -se::Event* XlaTensor::GetDefinitionEvent(se::Stream* stream) { +void XlaTensor::WaitForDefinitionEventOnStream(se::Stream* stream) { mutex_lock lock(mu_); if (!definition_event_) { - return nullptr; + return; } // The set of defined streams is expected to be very small indeed (usually @@ -84,24 +83,20 @@ se::Event* XlaTensor::GetDefinitionEvent(se::Stream* stream) { if (std::find(streams_defined_on_.begin(), streams_defined_on_.end(), stream) != streams_defined_on_.end()) { // stream is in streams_defined_on_; it doesn't need to be waited on. - return nullptr; + return; } - return definition_event_.get(); + stream->ThenWaitFor(definition_event_.get()); + streams_defined_on_.push_back(stream); } -void XlaTensor::SetDefinedOn(se::Stream* stream, - std::shared_ptr event) { +void XlaTensor::ResetDefinitionEvent(std::shared_ptr event, + se::Stream* stream) { mutex_lock lock(mu_); definition_event_ = std::move(event); streams_defined_on_ = {stream}; } -void XlaTensor::SetDefinedOn(se::Stream* stream) { - mutex_lock lock(mu_); - streams_defined_on_.push_back(stream); -} - // The pointer tag, OR-ed into the XlaTensor's address to distinguish it from // device-side tensors, which are either CPU or GPU memory pointers. This works // because we're guaranteed that CPU and GPU pointers are aligned to > 1 bits. diff --git a/tensorflow/compiler/jit/xla_tensor.h b/tensorflow/compiler/jit/xla_tensor.h index d95da63405889dfd0c279b17789a2195072c7277..77e80aa2527ecc2221ac61f7b7e6ebcce0982931 100644 --- a/tensorflow/compiler/jit/xla_tensor.h +++ b/tensorflow/compiler/jit/xla_tensor.h @@ -50,7 +50,7 @@ class XlaTensor { // Assign the internal ShapedBuffer to new memory for the given dtype and // shape. If a ShapedBuffer exists already (has_shaped_buffer() == true), it // is replaced and the managed memory deallocated. - Status AllocateShapedBuffer(DataType dtype, const TensorShape& shape, + Status AllocateShapedBuffer(DataType dtype, const xla::Shape& on_host_shape, xla::LocalClient* client, int device_ordinal); // Some Tensors can have complex on-device shapes, including tuple shapes. To @@ -88,23 +88,19 @@ class XlaTensor { host_tensor_.reset(new Tensor(tensor)); } - // If the tensor's content is not yet defined on 'stream', and there exists an - // se::Event declaring when the tensor's content is defined, return it. - // Otherwise, return nullptr. If this function returns nullptr then the - // tensor's content can be read on 'stream' without additional - // synchronization. - se::Event* GetDefinitionEvent(se::Stream* stream); - - // Assert that the tensor's content is defined on 'stream' by the time 'event' - // triggers. - void SetDefinedOn(se::Stream* stream, std::shared_ptr event); - - // Assert that the tensor's content is defined on 'stream'. This version does - // not provide an event, and must be called *after* SetDefinedOn(Stream, - // Event). This call can be read as an assertion that the definition event has - // been waited on by 'stream', so further calls to GetDefinitionEvent(stream) - // do not need to also wait on the event. - void SetDefinedOn(se::Stream* stream); + // Adds synchronization events to 'stream' that wait for this tensor to be + // defined on 'stream'. Does nothing if the tensor is already defined on that + // stream. + void WaitForDefinitionEventOnStream(se::Stream* stream); + + // (Re)sets the definition event of the tensor to 'event', and promises that + // the tensor has already been defined on stream. Removes any previous + // definition event or any previous promises about the tensor being defined on + // streams. + // It is legal to reset the definition event of a tensor when overwriting the + // tensor's value (at which point, it is effectively a new tensor once again.) + void ResetDefinitionEvent(std::shared_ptr event, + se::Stream* stream); // Convert from a raw pointer to an XlaTensor, removing the pointer tag. static XlaTensor* FromOpaquePointer(void* ptr); diff --git a/tensorflow/compiler/tests/BUILD b/tensorflow/compiler/tests/BUILD index d6e3f0817edbc21a1dfdccfd9d075c12f7010d97..f80cb1812f00d36ddb7c28ae0e77c58498058ef3 100644 --- a/tensorflow/compiler/tests/BUILD +++ b/tensorflow/compiler/tests/BUILD @@ -230,6 +230,7 @@ tf_xla_py_test( "//tensorflow/python:framework", "//tensorflow/python:platform_test", "//tensorflow/python:random_ops", + "//tensorflow/python:standard_ops", ], ) @@ -294,33 +295,6 @@ tf_xla_py_test( ], ) -tf_xla_py_test( - name = "oom_test", - size = "medium", - srcs = ["oom_test.py"], - # TODO(b/80081500): Re-enable on GPU. Disabled on 2018-05-21. - disabled_backends = [ - "cpu", - "cpu_ondemand", - "gpu", - ], - tags = [ - # Allocates very large amounts of memory and does not work under TSAN. - "notsan", - "optonly", # Times out frequently in fastbuild. - ], - deps = [ - ":xla_test", - "//tensorflow/python:array_ops", - "//tensorflow/python:array_ops_gen", - "//tensorflow/python:framework", - "//tensorflow/python:gradient_checker", - "//tensorflow/python:gradients", - "//tensorflow/python:math_ops", - "//tensorflow/python:platform_test", - ], -) - tf_xla_py_test( name = "conv2d_test", size = "medium", @@ -435,13 +409,6 @@ tf_xla_py_test( name = "eager_test", size = "large", srcs = ["eager_test.py"], - disabled_backends = [ - # TODO(b/78199195) Support XLA CPU devices in eager runtime - "cpu", - "cpu_ondemand", - # TODO(b/78468222) Enable GPU backend - "gpu", - ], deps = [ ":xla_test", "//tensorflow/python:array_ops", @@ -476,12 +443,11 @@ tf_xla_py_test( tags = ["optonly"], deps = [ ":xla_test", - "//tensorflow/contrib/signal:signal_py", "//tensorflow/python:array_ops", "//tensorflow/python:extra_py_tests_deps", "//tensorflow/python:framework", "//tensorflow/python:platform_test", - "//tensorflow/python:spectral_ops", + "//tensorflow/python/ops/signal", ], ) @@ -516,8 +482,6 @@ tf_xla_py_test( name = "function_test", size = "small", srcs = ["function_test.py"], - # Functions are not implemented in the on-demand compilation model yet. - disabled_backends = "cpu_ondemand", deps = [ ":xla_test", "//tensorflow/python:array_ops", @@ -707,9 +671,6 @@ tf_xla_py_test( name = "random_ops_test", size = "small", srcs = ["random_ops_test.py"], - disabled_backends = [ - "cpu_ondemand", - ], deps = [ ":xla_test", "//tensorflow/python:array_ops", @@ -717,6 +678,7 @@ tf_xla_py_test( "//tensorflow/python:math_ops", "//tensorflow/python:platform_test", "//tensorflow/python:random_ops", + "//tensorflow/python:standard_ops", ], ) @@ -740,7 +702,6 @@ tf_xla_py_test( name = "reduce_window_test", size = "small", srcs = ["reduce_window_test.py"], - disabled_backends = ["cpu_ondemand"], deps = [ ":xla_test", "//tensorflow/compiler/tf2xla/python:xla", @@ -849,8 +810,6 @@ tf_xla_py_test( name = "stack_ops_test", size = "small", srcs = ["stack_ops_test.py"], - # Stack ops are not implemented in the on-demand compilation model yet. - disabled_backends = "cpu_ondemand", deps = [ ":xla_test", "//tensorflow/python:array_ops", @@ -869,6 +828,7 @@ tf_xla_py_test( ":xla_test", "//tensorflow/python:framework", "//tensorflow/python:platform_test", + "//tensorflow/python:standard_ops", "//tensorflow/python:stateless_random_ops", ], ) @@ -878,7 +838,7 @@ tf_xla_py_test( size = "small", srcs = ["tensor_array_ops_test.py"], # TensorArray ops are not implemented in the on-demand compilation model yet. - disabled_backends = "cpu_ondemand", + disabled_backends = ["cpu_ondemand"], deps = [ ":xla_test", "//tensorflow/python:array_ops", @@ -899,7 +859,7 @@ tf_xla_py_test( size = "small", srcs = ["tensor_list_ops_test.py"], # TensorList ops are not implemented in the on-demand compilation model yet. - disabled_backends = "cpu_ondemand", + disabled_backends = ["cpu_ondemand"], deps = [ ":xla_test", "//tensorflow/python:array_ops", @@ -979,7 +939,6 @@ tf_xla_py_test( name = "while_test", size = "small", srcs = ["while_test.py"], - disabled_backends = ["cpu_ondemand"], deps = [ ":xla_test", "//tensorflow/compiler/tf2xla/python:xla", @@ -1136,6 +1095,7 @@ cc_library( "//tensorflow/core:test", "//tensorflow/core:testlib", "//tensorflow/core/kernels:ops_util", + "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/strings", ], @@ -1231,11 +1191,18 @@ tf_xla_py_test( tf_xla_py_test( name = "quantized_ops_test", - size = "small", + size = "medium", srcs = ["quantized_ops_test.py"], + disabled_backends = [ + "cpu", + "cpu_ondemand", + ], deps = [ ":xla_test", + "//tensorflow/compiler/tf2xla/python:xla", "//tensorflow/python:array_ops", + "//tensorflow/python:bitwise_ops", + "//tensorflow/python:constant_op", "//tensorflow/python:dtypes", "//tensorflow/python:math_ops", "//tensorflow/python:platform_test", @@ -1246,7 +1213,6 @@ tf_xla_py_test( name = "xla_ops_test", size = "medium", srcs = ["xla_ops_test.py"], - disabled_backends = ["cpu_ondemand"], deps = [ ":xla_test", "//tensorflow/compiler/tf2xla/python:xla", diff --git a/tensorflow/compiler/tests/adagrad_da_test.py b/tensorflow/compiler/tests/adagrad_da_test.py index 69fb3ec2964a09508e612515b9e291fc14121d68..e9c2d363acab96c0fb968cb7f901ce105ea8703e 100644 --- a/tensorflow/compiler/tests/adagrad_da_test.py +++ b/tensorflow/compiler/tests/adagrad_da_test.py @@ -50,8 +50,8 @@ class AdagradDAOptimizerTest(xla_test.XLATestCase): zip([grads0, grads1], [var0, var1]), global_step=global_step) variables.global_variables_initializer().run() - self.assertAllClose([0.0, 0.0], var0.eval()) - self.assertAllClose([0.0, 0.0], var1.eval()) + self.assertAllClose([0.0, 0.0], self.evaluate(var0)) + self.assertAllClose([0.0, 0.0], self.evaluate(var1)) # Run a step of AdagradDA update.run() @@ -63,9 +63,9 @@ class AdagradDAOptimizerTest(xla_test.XLATestCase): # For -0.1*3.0*(0.1 - 0)/(0 + sqrt(0.1 + 0.1*0.1)) = -0.904534 # similarly for others. self.assertAllCloseAccordingToType( - np.array([-0.904534, -1.603567]), var0.eval()) + np.array([-0.904534, -1.603567]), self.evaluate(var0)) self.assertAllCloseAccordingToType( - np.array([-0.094821, -0.189358]), var1.eval()) + np.array([-0.094821, -0.189358]), self.evaluate(var1)) def testAdagradDAwithoutRegularizationBasic2(self): for dtype in self.float_types: @@ -87,16 +87,16 @@ class AdagradDAOptimizerTest(xla_test.XLATestCase): zip([grads0, grads1], [var0, var1]), global_step=global_step) variables.global_variables_initializer().run() - self.assertAllCloseAccordingToType([1.0, 2.0], var0.eval()) - self.assertAllCloseAccordingToType([4.0, 3.0], var1.eval()) + self.assertAllCloseAccordingToType([1.0, 2.0], self.evaluate(var0)) + self.assertAllCloseAccordingToType([4.0, 3.0], self.evaluate(var1)) # Run a step of AdagradDA update.run() self.assertAllCloseAccordingToType( - np.array([-0.904534, -1.603567]), var0.eval()) + np.array([-0.904534, -1.603567]), self.evaluate(var0)) self.assertAllCloseAccordingToType( - np.array([-0.094821, -0.189358]), var1.eval()) + np.array([-0.094821, -0.189358]), self.evaluate(var1)) def testAdagradDAWithL1(self): for dtype in self.float_types: @@ -118,16 +118,16 @@ class AdagradDAOptimizerTest(xla_test.XLATestCase): zip([grads0, grads1], [var0, var1]), global_step=global_step) variables.global_variables_initializer().run() - self.assertAllCloseAccordingToType([1.0, 2.0], var0.eval()) - self.assertAllCloseAccordingToType([4.0, 3.0], var1.eval()) + self.assertAllCloseAccordingToType([1.0, 2.0], self.evaluate(var0)) + self.assertAllCloseAccordingToType([4.0, 3.0], self.evaluate(var1)) # Run a step of AdagradDA update.run() self.assertAllCloseAccordingToType( - np.array([-0.895489, -1.59555]), var0.eval()) + np.array([-0.895489, -1.59555]), self.evaluate(var0)) self.assertAllCloseAccordingToType( - np.array([-0.085339, -0.17989]), var1.eval()) + np.array([-0.085339, -0.17989]), self.evaluate(var1)) def testAdagradDAWithL1_L2(self): for dtype in self.float_types: @@ -149,16 +149,16 @@ class AdagradDAOptimizerTest(xla_test.XLATestCase): zip([grads0, grads1], [var0, var1]), global_step=global_step) variables.global_variables_initializer().run() - self.assertAllCloseAccordingToType([1.0, 2.0], var0.eval()) - self.assertAllCloseAccordingToType([4.0, 3.0], var1.eval()) + self.assertAllCloseAccordingToType([1.0, 2.0], self.evaluate(var0)) + self.assertAllCloseAccordingToType([4.0, 3.0], self.evaluate(var1)) # Run a step of AdagradDA update.run() self.assertAllCloseAccordingToType( - np.array([-0.046907, -0.093659]), var0.eval()) + np.array([-0.046907, -0.093659]), self.evaluate(var0)) self.assertAllCloseAccordingToType( - np.array([-0.004275, -0.009023]), var1.eval()) + np.array([-0.004275, -0.009023]), self.evaluate(var1)) if __name__ == "__main__": diff --git a/tensorflow/compiler/tests/adagrad_test.py b/tensorflow/compiler/tests/adagrad_test.py index ab69319c59fb07e7ce56c3c287a50a6290effdfd..e26483303c3934fd51675cb1fbc998b276caf527 100644 --- a/tensorflow/compiler/tests/adagrad_test.py +++ b/tensorflow/compiler/tests/adagrad_test.py @@ -42,17 +42,19 @@ class AdagradOptimizerTest(xla_test.XLATestCase): zip([grads0, grads1], [var0, var1])) variables.global_variables_initializer().run() # Fetch params to validate initial values - self.assertAllClose([1.0, 2.0], var0.eval()) - self.assertAllClose([3.0, 4.0], var1.eval()) + self.assertAllClose([1.0, 2.0], self.evaluate(var0)) + self.assertAllClose([3.0, 4.0], self.evaluate(var1)) # Run 3 steps of adagrad for _ in range(3): ada_update.run() # Validate updated params self.assertAllCloseAccordingToType( - np.array([-1.6026098728179932, -0.6026098728179932]), var0.eval(), + np.array([-1.6026098728179932, -0.6026098728179932]), + self.evaluate(var0), float_rtol=1e-5) self.assertAllCloseAccordingToType( - np.array([2.715679168701172, 3.715679168701172]), var1.eval(), + np.array([2.715679168701172, 3.715679168701172]), + self.evaluate(var1), float_rtol=1e-5) def testTensorLearningRate(self): @@ -68,17 +70,19 @@ class AdagradOptimizerTest(xla_test.XLATestCase): zip([grads0, grads1], [var0, var1])) variables.global_variables_initializer().run() # Fetch params to validate initial values - self.assertAllClose([1.0, 2.0], var0.eval()) - self.assertAllClose([3.0, 4.0], var1.eval()) + self.assertAllClose([1.0, 2.0], self.evaluate(var0)) + self.assertAllClose([3.0, 4.0], self.evaluate(var1)) # Run 3 steps of adagrad for _ in range(3): ada_update.run() # Validate updated params self.assertAllCloseAccordingToType( - np.array([-1.6026098728179932, -0.6026098728179932]), var0.eval(), + np.array([-1.6026098728179932, -0.6026098728179932]), + self.evaluate(var0), float_rtol=1e-5) self.assertAllCloseAccordingToType( - np.array([2.715679168701172, 3.715679168701172]), var1.eval(), + np.array([2.715679168701172, 3.715679168701172]), + self.evaluate(var1), float_rtol=1e-5) def testSharing(self): @@ -103,18 +107,20 @@ class AdagradOptimizerTest(xla_test.XLATestCase): variables.global_variables_initializer().run() # Fetch params to validate initial values. - self.assertAllClose([1.0, 2.0], var0.eval()) - self.assertAllClose([3.0, 4.0], var1.eval()) + self.assertAllClose([1.0, 2.0], self.evaluate(var0)) + self.assertAllClose([3.0, 4.0], self.evaluate(var1)) # Mix the first and the second adagrad for 3 steps. ada_update1.run() ada_update2.run() ada_update1.run() # Validate updated params (the same as with only 1 Adagrad). self.assertAllCloseAccordingToType( - np.array([-1.6026098728179932, -0.6026098728179932]), var0.eval(), + np.array([-1.6026098728179932, -0.6026098728179932]), + self.evaluate(var0), float_rtol=1e-5) self.assertAllCloseAccordingToType( - np.array([2.715679168701172, 3.715679168701172]), var1.eval(), + np.array([2.715679168701172, 3.715679168701172]), + self.evaluate(var1), float_rtol=1e-5) diff --git a/tensorflow/compiler/tests/adam_test.py b/tensorflow/compiler/tests/adam_test.py index 058576b3d4b695209952158769162bb24e7ccfce..8bcff9d379d34f8a6bb8b0fdc60b7588c6d80be9 100644 --- a/tensorflow/compiler/tests/adam_test.py +++ b/tensorflow/compiler/tests/adam_test.py @@ -75,23 +75,24 @@ class AdamOptimizerTest(xla_test.XLATestCase): variables.global_variables_initializer().run() # Fetch params to validate initial values - self.assertAllClose([1.0, 2.0], var0.eval()) - self.assertAllClose([3.0, 4.0], var1.eval()) + self.assertAllClose([1.0, 2.0], self.evaluate(var0)) + self.assertAllClose([3.0, 4.0], self.evaluate(var1)) beta1_power, beta2_power = opt._get_beta_accumulators() # Run 3 steps of Adam for t in range(1, 4): - self.assertAllCloseAccordingToType(0.9**t, beta1_power.eval()) - self.assertAllCloseAccordingToType(0.999**t, beta2_power.eval()) + self.assertAllCloseAccordingToType(0.9**t, self.evaluate(beta1_power)) + self.assertAllCloseAccordingToType(0.999**t, + self.evaluate(beta2_power)) update.run(feed_dict={grads0: grads0_np, grads1: grads1_np}) var0_np, m0, v0 = adam_update_numpy(var0_np, grads0_np, t, m0, v0) var1_np, m1, v1 = adam_update_numpy(var1_np, grads1_np, t, m1, v1) # Validate updated params - self.assertAllCloseAccordingToType(var0_np, var0.eval()) - self.assertAllCloseAccordingToType(var1_np, var1.eval()) + self.assertAllCloseAccordingToType(var0_np, self.evaluate(var0)) + self.assertAllCloseAccordingToType(var1_np, self.evaluate(var1)) def testTensorLearningRate(self): for dtype in self.float_types: @@ -117,23 +118,24 @@ class AdamOptimizerTest(xla_test.XLATestCase): variables.global_variables_initializer().run() # Fetch params to validate initial values - self.assertAllClose([1.0, 2.0], var0.eval()) - self.assertAllClose([3.0, 4.0], var1.eval()) + self.assertAllClose([1.0, 2.0], self.evaluate(var0)) + self.assertAllClose([3.0, 4.0], self.evaluate(var1)) beta1_power, beta2_power = opt._get_beta_accumulators() # Run 3 steps of Adam for t in range(1, 4): - self.assertAllCloseAccordingToType(0.9**t, beta1_power.eval()) - self.assertAllCloseAccordingToType(0.999**t, beta2_power.eval()) + self.assertAllCloseAccordingToType(0.9**t, self.evaluate(beta1_power)) + self.assertAllCloseAccordingToType(0.999**t, + self.evaluate(beta2_power)) update.run(feed_dict={grads0: grads0_np, grads1: grads1_np}) var0_np, m0, v0 = adam_update_numpy(var0_np, grads0_np, t, m0, v0) var1_np, m1, v1 = adam_update_numpy(var1_np, grads1_np, t, m1, v1) # Validate updated params - self.assertAllCloseAccordingToType(var0_np, var0.eval()) - self.assertAllCloseAccordingToType(var1_np, var1.eval()) + self.assertAllCloseAccordingToType(var0_np, self.evaluate(var0)) + self.assertAllCloseAccordingToType(var1_np, self.evaluate(var1)) def testSharing(self): for dtype in self.float_types: @@ -162,13 +164,14 @@ class AdamOptimizerTest(xla_test.XLATestCase): beta1_power, beta2_power = opt._get_beta_accumulators() # Fetch params to validate initial values - self.assertAllClose([1.0, 2.0], var0.eval()) - self.assertAllClose([3.0, 4.0], var1.eval()) + self.assertAllClose([1.0, 2.0], self.evaluate(var0)) + self.assertAllClose([3.0, 4.0], self.evaluate(var1)) # Run 3 steps of intertwined Adam1 and Adam2. for t in range(1, 4): - self.assertAllCloseAccordingToType(0.9**t, beta1_power.eval()) - self.assertAllCloseAccordingToType(0.999**t, beta2_power.eval()) + self.assertAllCloseAccordingToType(0.9**t, self.evaluate(beta1_power)) + self.assertAllCloseAccordingToType(0.999**t, + self.evaluate(beta2_power)) if t % 2 == 0: update1.run(feed_dict={grads0: grads0_np, grads1: grads1_np}) else: @@ -178,8 +181,8 @@ class AdamOptimizerTest(xla_test.XLATestCase): var1_np, m1, v1 = adam_update_numpy(var1_np, grads1_np, t, m1, v1) # Validate updated params - self.assertAllCloseAccordingToType(var0_np, var0.eval()) - self.assertAllCloseAccordingToType(var1_np, var1.eval()) + self.assertAllCloseAccordingToType(var0_np, self.evaluate(var0)) + self.assertAllCloseAccordingToType(var1_np, self.evaluate(var1)) if __name__ == "__main__": diff --git a/tensorflow/compiler/tests/adamax_test.py b/tensorflow/compiler/tests/adamax_test.py index 3ed1d41b7121f44dd7470f61180f7a7055369174..961b46375c941bdc3922e460a2f58345086dbceb 100644 --- a/tensorflow/compiler/tests/adamax_test.py +++ b/tensorflow/compiler/tests/adamax_test.py @@ -78,8 +78,8 @@ class AdaMaxOptimizerTest(xla_test.XLATestCase): variables.global_variables_initializer().run() # Fetch params to validate initial values - self.assertAllClose([1.0, 2.0], var0.eval()) - self.assertAllClose([3.0, 4.0], var1.eval()) + self.assertAllClose([1.0, 2.0], self.evaluate(var0)) + self.assertAllClose([3.0, 4.0], self.evaluate(var1)) beta1_power = opt._get_beta_accumulators() @@ -87,14 +87,17 @@ class AdaMaxOptimizerTest(xla_test.XLATestCase): for t in range(1, 4): update.run() - self.assertAllCloseAccordingToType(0.9**(t + 1), beta1_power.eval()) + self.assertAllCloseAccordingToType(0.9**(t + 1), + self.evaluate(beta1_power)) var0_np, m0, v0 = adamax_update_numpy(var0_np, grads0_np, t, m0, v0) var1_np, m1, v1 = adamax_update_numpy(var1_np, grads1_np, t, m1, v1) # Validate updated params - self.assertAllCloseAccordingToType(var0_np, var0.eval(), rtol=1e-2) - self.assertAllCloseAccordingToType(var1_np, var1.eval(), rtol=1e-2) + self.assertAllCloseAccordingToType( + var0_np, self.evaluate(var0), rtol=1e-2) + self.assertAllCloseAccordingToType( + var1_np, self.evaluate(var1), rtol=1e-2) self.assertEqual("var0_%d/AdaMax:0" % (i,), opt.get_slot(var=var0, name="m").name) @@ -118,22 +121,23 @@ class AdaMaxOptimizerTest(xla_test.XLATestCase): variables.global_variables_initializer().run() # Fetch params to validate initial values - self.assertAllClose([1.0, 2.0], var0.eval()) - self.assertAllClose([3.0, 4.0], var1.eval()) + self.assertAllClose([1.0, 2.0], self.evaluate(var0)) + self.assertAllClose([3.0, 4.0], self.evaluate(var1)) beta1_power = opt._get_beta_accumulators() # Run 3 steps of AdaMax for t in range(1, 4): - self.assertAllCloseAccordingToType(0.9**t, beta1_power.eval()) + self.assertAllCloseAccordingToType(0.9**t, self.evaluate(beta1_power)) update.run() var0_np, m0, v0 = adamax_update_numpy(var0_np, grads0_np, t, m0, v0) var1_np, m1, v1 = adamax_update_numpy(var1_np, grads1_np, t, m1, v1) # Validate updated params - self.assertAllCloseAccordingToType(var0_np, var0.eval()) - self.assertAllCloseAccordingToType(var1_np, var1.eval()) + self.assertAllCloseAccordingToType(var0_np, self.evaluate(var0)) + self.assertAllCloseAccordingToType(var1_np, self.evaluate(var1)) + if __name__ == "__main__": test.main() diff --git a/tensorflow/compiler/tests/addsign_test.py b/tensorflow/compiler/tests/addsign_test.py index 1bc07ace23ccdc83103abe71ee11b72994c75a6d..a37c97e6d374440aeb860b9d02f2d5dd95c91f62 100644 --- a/tensorflow/compiler/tests/addsign_test.py +++ b/tensorflow/compiler/tests/addsign_test.py @@ -90,8 +90,8 @@ class AddSignTest(xla_test.XLATestCase): variables.global_variables_initializer().run() # Fetch params to validate initial values - self.assertAllClose([1.0, 2.0], var0.eval()) - self.assertAllClose([3.0, 4.0], var1.eval()) + self.assertAllClose([1.0, 2.0], self.evaluate(var0)) + self.assertAllClose([3.0, 4.0], self.evaluate(var1)) # Run 7 steps of AddSign # first 4 steps with positive gradient @@ -125,8 +125,8 @@ class AddSignTest(xla_test.XLATestCase): # Validate updated params self.assertAllCloseAccordingToType( - var0_np, var0.eval(), half_rtol=1e-2) - self.assertAllCloseAccordingToType(var1_np, var1.eval()) + var0_np, self.evaluate(var0), half_rtol=1e-2) + self.assertAllCloseAccordingToType(var1_np, self.evaluate(var1)) def testDense(self): decay_steps = 10 diff --git a/tensorflow/compiler/tests/binary_ops_test.py b/tensorflow/compiler/tests/binary_ops_test.py index 1b39d53dc0908e1fa05f766ca1e601731b26846d..c829c50b5518b29c96c0b0117a6cd143911bd1fc 100644 --- a/tensorflow/compiler/tests/binary_ops_test.py +++ b/tensorflow/compiler/tests/binary_ops_test.py @@ -18,6 +18,8 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import itertools + import numpy as np from tensorflow.compiler.tests import xla_test @@ -178,6 +180,13 @@ class BinaryOpsTest(xla_test.XLATestCase): [0, 0, 0, 0, 0, 0.1, 0.3, 0.5, 0.7, 0.9, 6.1, 10.0], dtype=dtype), expected=np.array([0, 0, 0, 0, 0, 6, 7, 8, 9, 10, 0, 0], dtype=dtype)) + self._testBinary( + gen_nn_ops.leaky_relu_grad, + np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10], dtype=dtype), + np.array([0, 0, 0, 0, 0, 0.1, 0.3, 0.5, 0.7, 0.9], dtype=dtype), + expected=np.array([0.2, 0.4, 0.6, 0.8, 1, 6, 7, 8, 9, 10], + dtype=dtype)) + self._testBinary( gen_nn_ops.softmax_cross_entropy_with_logits, np.array([[1, 2, 3, 4], [5, 6, 7, 8]], dtype=dtype), @@ -209,6 +218,21 @@ class BinaryOpsTest(xla_test.XLATestCase): ], equality_test=self.ListsAreClose) + # TF doesn't define these for bf16. + if dtype != dtypes.bfloat16.as_numpy_dtype: + self._testBinary( + gen_math_ops.xdivy, + np.array([0, 4, 3, 2, 1, 0], dtype=dtype), + np.array([0, 5, 6, 7, 8, float("NaN")], dtype=dtype), + expected=np.array([0, 0.8, 0.5, 0.285714, 0.125, 0], dtype=dtype)) + + self._testBinary( + gen_math_ops.xlogy, + np.array([0, 4, 3, 2, 1, 0], dtype=dtype), + np.array([0, 5, 6, 7, 8, float("NaN")], dtype=dtype), + expected=np.array([0, 6.437752, 5.375278, 3.89182, 2.079442, 0], + dtype=dtype)) + def testIntOps(self): for dtype in self.signed_int_types: self._testBinary( @@ -287,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( @@ -376,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]), @@ -960,7 +1008,7 @@ class BinaryOpsTest(xla_test.XLATestCase): self._testBinary( array_ops.expand_dims, np.array([42], dtype=dtype), - np.int32(0), + np.array([0], dtype=np.int64), expected=np.array([[42]], dtype=dtype)) self._testBinary( array_ops.expand_dims, @@ -987,15 +1035,21 @@ class BinaryOpsTest(xla_test.XLATestCase): np.array([[[1, 2], [3, 4]]], dtype=dtype), np.int32(3), expected=np.array([[[[1], [2]], [[3], [4]]]], dtype=dtype)) + self._testBinary( + array_ops.expand_dims, + np.array([[[1, 2], [3, 4]]], dtype=dtype), + np.array([2], dtype=np.int64), + expected=np.array([[[[1, 2]], [[3, 4]]]], dtype=dtype)) def testPad(self): - for dtype in self.numeric_types: + for dtype, pad_type in itertools.product( + self.numeric_types, [np.int32, np.int64]): self._testBinary( array_ops.pad, np.array( [[1, 2, 3], [4, 5, 6]], dtype=dtype), np.array( - [[1, 2], [2, 1]], dtype=np.int32), + [[1, 2], [2, 1]], dtype=pad_type), expected=np.array( [[0, 0, 0, 0, 0, 0], [0, 0, 1, 2, 3, 0], @@ -1009,7 +1063,7 @@ class BinaryOpsTest(xla_test.XLATestCase): np.array( [[1, 2, 3], [4, 5, 6]], dtype=dtype), np.array( - [[0, 3], [2, 1]], dtype=np.int32), + [[0, 3], [2, 1]], dtype=pad_type), expected=np.array( [[7, 7, 1, 2, 3, 7], [7, 7, 4, 5, 6, 7], diff --git a/tensorflow/compiler/tests/build_defs.bzl b/tensorflow/compiler/tests/build_defs.bzl index 1d3979b21bfd915a641fabe1ef40301b3e5a17b4..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(): @@ -50,6 +51,8 @@ def tf_xla_py_test( """ if disabled_backends == None: disabled_backends = [] + if type(disabled_backends) != "list": + fail("disabled_backends must be a list of strings", "disabled_backends") enabled_backends = [b for b in all_backends() if b not in disabled_backends] test_names = [] @@ -62,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 += [ @@ -82,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, @@ -90,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 a57d1dc81ea2c9c188b0a3005904738aa8156bf3..5d5e486f616937601214aa169a4c329ab78932c8 100644 --- a/tensorflow/compiler/tests/categorical_op_test.py +++ b/tensorflow/compiler/tests/categorical_op_test.py @@ -27,6 +27,7 @@ from tensorflow.python.framework import dtypes from tensorflow.python.framework import random_seed from tensorflow.python.ops import array_ops from tensorflow.python.ops import random_ops +from tensorflow.python.ops import stateless_random_ops from tensorflow.python.platform import googletest @@ -56,11 +57,11 @@ class CategoricalTest(xla_test.XLATestCase): Returns: Frequencies from sampled classes; shape [batch_size, num_classes]. """ - with self.cached_session() as sess, self.test_scope(): + with self.cached_session(), self.test_scope(): random_seed.set_random_seed(1618) op = random_ops.multinomial(logits, num_samples, output_dtype=dtypes.int32) - d = sess.run(op) + d = self.evaluate(op) batch_size, num_classes = logits.shape freqs_mat = [] @@ -79,15 +80,15 @@ class CategoricalTest(xla_test.XLATestCase): def _testRngIsNotConstant(self, rng, dtype, output_dtype): # Tests that 'rng' does not always return the same value. - with self.cached_session() as sess: + with self.cached_session(): with self.test_scope(): x = rng(dtype, output_dtype) # The random-number generator, if working correctly, should produce the # same output multiple times with low probability. - y = sess.run(x) - z = sess.run(x) - w = sess.run(x) + y = self.evaluate(x) + z = self.evaluate(x) + w = self.evaluate(x) # We use exact equality here. If the random-number generator is producing # deterministic output, all three outputs will be bitwise identical. @@ -107,12 +108,12 @@ class CategoricalTest(xla_test.XLATestCase): def testCategoricalIsInRange(self): for dtype in self.float_types: for output_dtype in self.output_dtypes(): - with self.cached_session() as sess: + with self.cached_session(): with self.test_scope(): x = random_ops.multinomial( array_ops.ones(shape=[1, 20], dtype=dtype), 1000, output_dtype=output_dtype) - y = sess.run(x) + y = self.evaluate(x) self.assertTrue((y >= 0).sum() == 1000) self.assertTrue((y < 20).sum() == 1000) @@ -138,6 +139,57 @@ class CategoricalTest(xla_test.XLATestCase): chi2 = self._chi2(probs, freqs) self.assertLess(chi2, 1e-3) + def testStatelessMultinomialIsInRange(self): + for dtype in self.float_types: + for output_dtype in self.output_dtypes(): + with self.cached_session() as sess: + with self.test_scope(): + seed_t = array_ops.placeholder(dtypes.int32, shape=[2]) + x = stateless_random_ops.stateless_multinomial( + array_ops.ones(shape=[1, 20], dtype=dtype), + 1000, + seed_t, + output_dtype=output_dtype) + y = sess.run(x, {seed_t: [0x12345678, 0xabcdef12]}) + self.assertTrue((y >= 0).sum() == 1000) + self.assertTrue((y < 20).sum() == 1000) + + def testDeterminismMultinomial(self): + # Stateless values should be equal iff the seeds are equal (roughly) + num_samples = 10 + with self.cached_session(), self.test_scope(): + seed_t = array_ops.placeholder(dtypes.int32, shape=[2]) + seeds = [(x, y) for x in range(5) for y in range(5)] * 3 + for logits in ([[0.1, 0.25, 0.5, 0.15]], [[0.5, 0.5], [0.8, 0.2], + [0.25, 0.75]]): + pure = stateless_random_ops.stateless_multinomial( + logits, num_samples, seed=seed_t) + values = [(seed, pure.eval(feed_dict={seed_t: seed})) for seed in seeds] + for s0, v0 in values: + for s1, v1 in values: + self.assertEqual(s0 == s1, np.all(v0 == v1)) + + def testEmpty(self): + with self.cached_session(): + with self.test_scope(): + x = random_ops.multinomial( + array_ops.zeros([42, 40]), 0, output_dtype=dtypes.int32) + y = self.evaluate(x) + self.assertEqual(y.shape, (42, 0)) + + def testEmptyStateless(self): + with self.cached_session() as sess: + with self.test_scope(): + seed_t = array_ops.placeholder(dtypes.int32, shape=[2]) + x = stateless_random_ops.stateless_multinomial( + array_ops.zeros([42, 40]), + 0, + seed=seed_t, + output_dtype=dtypes.int32) + y = sess.run(x, {seed_t: [0x12345678, 0xabcdef12]}) + self.assertEqual(y.shape, (42, 0)) + + if __name__ == '__main__': googletest.main() diff --git a/tensorflow/compiler/tests/clustering_test.py b/tensorflow/compiler/tests/clustering_test.py index 88bd58b2da6b2892f898ad10f3467d8ce39d6388..ef2d7af69deeebd5f4c4c7225d7027f8f76bf861 100644 --- a/tensorflow/compiler/tests/clustering_test.py +++ b/tensorflow/compiler/tests/clustering_test.py @@ -43,7 +43,7 @@ class ClusteringTest(xla_test.XLATestCase): input1 = constant_op.constant(val1, name="const1") input2 = constant_op.constant(val2, name="const2") output = math_ops.add(input1, input2) - result = output.eval() + result = self.evaluate(output) self.assertAllClose(result, expected, rtol=1e-3) def testAddFromCpuMultiple(self): @@ -57,7 +57,7 @@ class ClusteringTest(xla_test.XLATestCase): with self.test_scope(): output = math_ops.add(input1, input2) for _ in xrange(10): - result = output.eval() + result = self.evaluate(output) self.assertAllClose(result, expected, rtol=1e-3) def testDeadlock(self): diff --git a/tensorflow/compiler/tests/concat_ops_test.py b/tensorflow/compiler/tests/concat_ops_test.py index 2d225ad226cac368042b95eae8fc29e6fd8e82e0..2187f57960f80300d631bdc7eb8fe5e9c8dddeea 100644 --- a/tensorflow/compiler/tests/concat_ops_test.py +++ b/tensorflow/compiler/tests/concat_ops_test.py @@ -72,7 +72,7 @@ class ConcatTest(xla_test.XLATestCase): x2 = constant_op.constant(p2) with self.test_scope(): c = array_ops.concat([x1, x2], 0) - result = c.eval() + result = self.evaluate(c) self.assertAllEqual(result[:2, :], p1) self.assertAllEqual(result[2:, :], p2) @@ -150,7 +150,7 @@ class ConcatTest(xla_test.XLATestCase): [float(x) for x in grad_inp.flatten()], shape=output_shape) grad = gradients_impl.gradients([c], inp_tensors, [grad_tensor]) concated_grad = array_ops.concat(grad, 1) - result = concated_grad.eval() + result = self.evaluate(concated_grad) self.assertAllEqual(result, grad_inp) def testGradientsSimpleAll(self): @@ -177,7 +177,7 @@ class ConcatTest(xla_test.XLATestCase): [float(x) for x in grad_inp.flatten()], shape=output_shape) grad = gradients_impl.gradients([c], inp_tensors, [grad_tensor]) concated_grad = array_ops.concat(grad, 0) - result = concated_grad.eval() + result = self.evaluate(concated_grad) self.assertAllEqual(result, grad_inp) @@ -205,7 +205,7 @@ class ConcatTest(xla_test.XLATestCase): [float(x) for x in grad_inp.flatten()], shape=output_shape) grad = gradients_impl.gradients([c], inp_tensors, [grad_tensor]) concated_grad = array_ops.concat(grad, 2) - result = concated_grad.eval() + result = self.evaluate(concated_grad) self.assertAllEqual(result, grad_inp) @@ -242,7 +242,7 @@ class ConcatTest(xla_test.XLATestCase): [float(x) for x in grad_inp.flatten()], shape=output_shape) grad = gradients_impl.gradients([c], inp_tensors, [grad_tensor]) concated_grad = array_ops.concat(grad, concat_dim) - result = concated_grad.eval() + result = self.evaluate(concated_grad) self.assertAllEqual(result, grad_inp) @@ -254,7 +254,7 @@ class ConcatTest(xla_test.XLATestCase): def DISABLED_testZeroSize(self): # Verify that concat doesn't crash and burn for zero size inputs np.random.seed(7) - with self.cached_session() as sess: + with self.cached_session(): with self.test_scope(): for shape0 in (), (2,): axis = len(shape0) @@ -270,7 +270,7 @@ class ConcatTest(xla_test.XLATestCase): self.assertAllEqual(c.eval(), correct) # Check gradients dc = np.random.randn(*c.get_shape().as_list()) - dxs = sess.run(gradients_impl.gradients(c, xs, dc)) + dxs = self.evaluate(gradients_impl.gradients(c, xs, dc)) self.assertAllEqual(dc, np.concatenate(dxs, axis=axis)) def testConcatTuple(self): @@ -280,7 +280,7 @@ class ConcatTest(xla_test.XLATestCase): with self.test_scope(): concat_list_t = array_ops.concat([c1, c2], 0) concat_tuple_t = array_ops.concat((c1, c2), 0) - self.assertAllEqual(concat_list_t.eval(), concat_tuple_t.eval()) + self.assertAllEqual(concat_list_t.eval(), self.evaluate(concat_tuple_t)) def testConcatNoScalars(self): with self.cached_session(): @@ -330,47 +330,47 @@ class ConcatTest(xla_test.XLATestCase): class ConcatOffsetTest(xla_test.XLATestCase): def testBasic(self): - with self.cached_session() as sess: + with self.cached_session(): with self.test_scope(): cdim = constant_op.constant(1, dtypes.int32) s0 = constant_op.constant([2, 3, 5], dtypes.int32) s1 = constant_op.constant([2, 7, 5], dtypes.int32) s2 = constant_op.constant([2, 20, 5], dtypes.int32) off = gen_array_ops.concat_offset(cdim, [s0, s1, s2]) - ans = sess.run(off) + ans = self.evaluate(off) self.assertAllEqual(ans, [[0, 0, 0], [0, 3, 0], [0, 10, 0]]) class PackTest(xla_test.XLATestCase): def testBasic(self): - with self.cached_session() as sess: + with self.cached_session(): with self.test_scope(): s0 = constant_op.constant([2, 3, 5], dtypes.int32) s1 = constant_op.constant([2, 7, 5], dtypes.int32) s2 = constant_op.constant([2, 20, 5], dtypes.int32) packed = array_ops.stack([s0, s1, s2]) - ans = sess.run(packed) + ans = self.evaluate(packed) self.assertAllEqual(ans, [[2, 3, 5], [2, 7, 5], [2, 20, 5]]) def testScalars(self): - with self.cached_session() as sess: + with self.cached_session(): with self.test_scope(): s0 = constant_op.constant(2, dtypes.int32) s1 = constant_op.constant(3, dtypes.int32) s2 = constant_op.constant(5, dtypes.int32) packed = array_ops.stack([s0, s1, s2]) - ans = sess.run(packed) + ans = self.evaluate(packed) self.assertAllEqual(ans, [2, 3, 5]) def testEmpty(self): - with self.cached_session() as sess: + with self.cached_session(): with self.test_scope(): s0 = constant_op.constant([[]], dtypes.int32) s1 = constant_op.constant([[]], dtypes.int32) s2 = constant_op.constant([[]], dtypes.int32) packed = array_ops.stack([s0, s1, s2]) - ans = sess.run(packed) + ans = self.evaluate(packed) self.assertAllEqual(ans, [[[]], [[]], [[]]]) diff --git a/tensorflow/compiler/tests/conv3d_test.py b/tensorflow/compiler/tests/conv3d_test.py index 33fd983b5485e503c2fcc96db2dfdecfc41e309f..01cc1b6392845be2418c50d55be97487eb290843 100644 --- a/tensorflow/compiler/tests/conv3d_test.py +++ b/tensorflow/compiler/tests/conv3d_test.py @@ -85,7 +85,7 @@ class Conv3DTransposeTest(xla_test.XLATestCase): 1.0, shape=f_shape, name="filter", dtype=dtypes.float32) output = nn_ops.conv3d_transpose( x, f, y_shape, strides=strides, padding="SAME") - value = output.eval() + value = self.evaluate(output) # We count the number of cells being added at the locations in the output. # At the center, #cells = kernel_depth * kernel_height * kernel_width @@ -135,7 +135,7 @@ class Conv3DTransposeTest(xla_test.XLATestCase): 1.0, shape=f_shape, name="filter", dtype=dtypes.float32) output = nn_ops.conv3d_transpose( x, f, y_shape, strides=strides, padding="SAME") - value = output.eval() + value = self.evaluate(output) for n in xrange(x_shape[0]): for k in xrange(f_shape[3]): @@ -173,7 +173,7 @@ class Conv3DTransposeTest(xla_test.XLATestCase): 1.0, shape=f_shape, name="filter", dtype=dtypes.float32) output = nn_ops.conv3d_transpose( x, f, y_shape, strides=strides, padding="VALID") - value = output.eval() + value = self.evaluate(output) cache_values = np.zeros(y_shape, dtype=np.float32) @@ -225,7 +225,7 @@ class Conv3DTransposeTest(xla_test.XLATestCase): err = gradient_checker.compute_gradient_error([x, f], [x_shape, f_shape], output, y_shape) print("conv3d_transpose gradient err = %g " % err) - err_tolerance = 0.0005 + err_tolerance = 0.001 self.assertLess(err, err_tolerance) diff --git a/tensorflow/compiler/tests/dense_layer_test.py b/tensorflow/compiler/tests/dense_layer_test.py index d1b90f098d7d6574999ba0af44b285f5ad5e4f8d..b7d08df9f7d144b71fd0b09535e10b8f596ea6ca 100644 --- a/tensorflow/compiler/tests/dense_layer_test.py +++ b/tensorflow/compiler/tests/dense_layer_test.py @@ -42,7 +42,7 @@ def GetRunMetadataLabels(run_metadata): def InLabels(labels, substr): """Returns true iff one of the labels contains substr.""" - return any([substr in x for x in labels]) + return any(substr in x for x in labels) class DenseLayerTest(test.TestCase): @@ -72,7 +72,7 @@ class DenseLayerTest(test.TestCase): x = array_ops.placeholder(shape=[None, None, 3], dtype=np.float32) y = layers.dense(x, 3) - sess.run(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) - sess.run(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) - sess.run(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/depthwise_conv_op_test.py b/tensorflow/compiler/tests/depthwise_conv_op_test.py index 6ef8a68ca5d35d3d2f78f0cb491e7bb98ff97ac9..90146e6b27ca31304a2549ec247412341efe390c 100644 --- a/tensorflow/compiler/tests/depthwise_conv_op_test.py +++ b/tensorflow/compiler/tests/depthwise_conv_op_test.py @@ -255,7 +255,7 @@ class DepthwiseConv2DTest(xla_test.XLATestCase): t1, t2, strides=[1, stride, stride, 1], padding=padding) value = sess.run(conv, {t1: x1, t2: x2}) print("value = ", value) - self.assertArrayNear(expected, np.ravel(value), 1e-5) + self.assertArrayNear(expected, np.ravel(value), 1e-4) self.assertShapeEqual(value, conv) def testConv2D2x2Filter(self): @@ -350,8 +350,13 @@ class DepthwiseConv2DTest(xla_test.XLATestCase): self._CompareBackpropInput(input_size, filter_size, output_size, stride, padding) - def _CompareBackpropFilter(self, input_sizes, filter_sizes, output_sizes, - stride, padding): + def _CompareBackpropFilter(self, + input_sizes, + filter_sizes, + output_sizes, + stride, + padding, + data_format="NHWC"): x0 = np.random.rand(*input_sizes).astype(np.float32) x2 = np.random.rand(*output_sizes).astype(np.float32) @@ -360,13 +365,30 @@ class DepthwiseConv2DTest(xla_test.XLATestCase): t0 = array_ops.placeholder(np.float32, shape=input_sizes) t1 = constant_op.constant(filter_sizes, shape=[len(filter_sizes)]) t2 = array_ops.placeholder(np.float32, shape=output_sizes) + native_t0 = t0 + native_t2 = t2 + strides = [1, stride, stride, 1] + if use_xla: + if data_format == "NCHW": + # Transpose from NWHC input to NCHW + # Ex. [4, 5, 5, 48] to [4, 48, 5, 5] + native_t0 = array_ops.transpose(t0, [0, 3, 1, 2]) + native_t2 = array_ops.transpose(t2, [0, 3, 1, 2]) + strides = [1, 1, stride, stride] with self.test_scope(): backprop = nn_ops.depthwise_conv2d_native_backprop_filter( - t0, t1, t2, strides=[1, stride, stride, 1], padding=padding) + native_t0, + t1, + native_t2, + strides=strides, + padding=padding, + data_format=data_format) else: + # For CPU, the format NCHW is not supported. Therefore we always use + # NHWC here. backprop = nn_ops.depthwise_conv2d_native_backprop_filter( - t0, t1, t2, strides=[1, stride, stride, 1], padding=padding) + native_t0, t1, native_t2, strides=strides, padding=padding) ret = backprop.eval({t0: x0, t2: x2}) self.assertShapeEqual(ret, backprop) return ret @@ -379,11 +401,24 @@ class DepthwiseConv2DTest(xla_test.XLATestCase): for index, (input_size, filter_size, output_size, stride, padding) in enumerate(ConfigsToTest()): print("Testing DepthwiseConv2DFilterGradCompare,", index, "th config:", - input_size, "*", filter_size, "stride:", stride, "padding:", - padding) + input_size, "*", filter_size, "producing output", output_size, + "stride:", stride, "padding:", padding) self._CompareBackpropFilter(input_size, filter_size, output_size, stride, padding) + def testDepthwiseConv2DFilterGradFormatNCHWCompare(self): + for index, (input_size, filter_size, output_size, stride, + padding) in enumerate(ConfigsToTest()): + print("Testing DepthwiseConv2DFilterGradFormatNCHWCompare,", index, + "th config:", input_size, "*", filter_size, "producing output", + output_size, "stride:", stride, "padding:", padding) + self._CompareBackpropFilter( + input_size, + filter_size, + output_size, + stride, + padding, + data_format="NCHW") if __name__ == "__main__": test.main() diff --git a/tensorflow/compiler/tests/dynamic_stitch_test.py b/tensorflow/compiler/tests/dynamic_stitch_test.py index 50b04daa6b9f4159a3c4bdeecaf900a5b35a833c..e89cf975f5d889091ce92a35165aef55ee5ad4b0 100644 --- a/tensorflow/compiler/tests/dynamic_stitch_test.py +++ b/tensorflow/compiler/tests/dynamic_stitch_test.py @@ -58,6 +58,15 @@ class DynamicStitchTest(xla_test.XLATestCase): [idx1, idx2], [val1, val2], expected=np.array([[], [], [], []], np.int32)) + def testEmptyIndex(self): + idx1 = np.array([], dtype=np.int32) + idx2 = np.array([[], []], dtype=np.int32) + val1 = np.ndarray(shape=(0, 9), dtype=np.int32) + val2 = np.ndarray(shape=(2, 0, 9), dtype=np.int32) + self._AssertDynamicStitchResultIs([idx1, idx2], [val1, val2], + expected=np.ndarray( + shape=(0, 9), dtype=np.int32)) + def testSimple1D(self): val1 = np.array([0, 4, 7], dtype=np.int32) val2 = np.array([1, 6, 2, 3, 5], dtype=np.int32) diff --git a/tensorflow/compiler/tests/eager_test.py b/tensorflow/compiler/tests/eager_test.py index 63cee550fde9d9d4314b1541fba191df776a4da2..c9fce39f6c5111f93a54708b59b4c42c3ba844b6 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,6 +32,7 @@ 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 gen_random_ops from tensorflow.python.ops import init_ops @@ -101,12 +103,12 @@ class EagerTest(xla_test.XLATestCase): self.assertAllEqual(15, product) # Run some ops graphly - with context.graph_mode(), self.cached_session() as sess: + with context.graph_mode(), self.cached_session(): with self.test_scope(): three = constant_op.constant(3) five = constant_op.constant(5) product = three * five - self.assertAllEqual(15, sess.run(product)) + self.assertAllEqual(15, self.evaluate(product)) def testDegenerateSlices(self): with self.test_scope(): @@ -463,7 +465,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 +481,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 +508,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 +557,56 @@ 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()) + class ExcessivePaddingTest(xla_test.XLATestCase): """Test that eager execution works with TPU flattened tensors. diff --git a/tensorflow/compiler/tests/fft_test.py b/tensorflow/compiler/tests/fft_test.py index b3e13fbaa6b33bdaa1be123be558059e96de282e..0edd0c35aa2d417a3ed24decbaa0b5d62d35bb62 100644 --- a/tensorflow/compiler/tests/fft_test.py +++ b/tensorflow/compiler/tests/fft_test.py @@ -24,11 +24,10 @@ import numpy as np import scipy.signal as sps from tensorflow.compiler.tests import xla_test -from tensorflow.contrib.signal.python.ops import spectral_ops as signal from tensorflow.python.framework import dtypes from tensorflow.python.ops import array_ops from tensorflow.python.ops import gradients_impl -from tensorflow.python.ops import spectral_ops +from tensorflow.python.ops.signal import signal from tensorflow.python.platform import googletest BATCH_DIMS = (3, 5) @@ -107,39 +106,39 @@ class FFTTest(xla_test.XLATestCase): def testFFT(self): self._VerifyFftMethod(INNER_DIMS_1D, lambda x: x, np.fft.fft, - spectral_ops.fft) + signal.fft) def testFFT2D(self): self._VerifyFftMethod(INNER_DIMS_2D, lambda x: x, np.fft.fft2, - spectral_ops.fft2d) + signal.fft2d) def testFFT3D(self): self._VerifyFftMethod(INNER_DIMS_3D, lambda x: x, lambda x: np.fft.fftn(x, axes=(-3, -2, -1)), - spectral_ops.fft3d) + signal.fft3d) def testIFFT(self): self._VerifyFftMethod(INNER_DIMS_1D, lambda x: x, np.fft.ifft, - spectral_ops.ifft) + signal.ifft) def testIFFT2D(self): self._VerifyFftMethod(INNER_DIMS_2D, lambda x: x, np.fft.ifft2, - spectral_ops.ifft2d) + signal.ifft2d) def testIFFT3D(self): self._VerifyFftMethod(INNER_DIMS_3D, lambda x: x, lambda x: np.fft.ifftn(x, axes=(-3, -2, -1)), - spectral_ops.ifft3d) + signal.ifft3d) def testRFFT(self): self._VerifyFftMethod( INNER_DIMS_1D, np.real, lambda x: np.fft.rfft(x, n=x.shape[-1]), - lambda x: spectral_ops.rfft(x, fft_length=[x.shape[-1].value])) + lambda x: signal.rfft(x, fft_length=[x.shape[-1].value])) def testRFFT2D(self): def _tf_fn(x): - return spectral_ops.rfft2d( + return signal.rfft2d( x, fft_length=[x.shape[-2].value, x.shape[-1].value]) self._VerifyFftMethod( @@ -153,16 +152,33 @@ class FFTTest(xla_test.XLATestCase): x, axes=(-3, -2, -1), s=[x.shape[-3], x.shape[-2], x.shape[-1]]) def _tf_fn(x): - return spectral_ops.rfft3d( + return signal.rfft3d( x, fft_length=[x.shape[-3].value, x.shape[-2].value, x.shape[-1].value]) self._VerifyFftMethod(INNER_DIMS_3D, np.real, _to_expected, _tf_fn) + def testRFFT3DMismatchedSize(self): + + def _to_expected(x): + return np.fft.rfftn( + x, + axes=(-3, -2, -1), + s=[x.shape[-3] // 2, x.shape[-2], x.shape[-1] * 2]) + + def _tf_fn(x): + return signal.rfft3d( + x, + fft_length=[ + x.shape[-3].value // 2, x.shape[-2].value, x.shape[-1].value * 2 + ]) + + self._VerifyFftMethod(INNER_DIMS_3D, np.real, _to_expected, _tf_fn) + def testIRFFT(self): def _tf_fn(x): - return spectral_ops.irfft(x, fft_length=[2 * (x.shape[-1].value - 1)]) + return signal.irfft(x, fft_length=[2 * (x.shape[-1].value - 1)]) self._VerifyFftMethod( INNER_DIMS_1D, lambda x: np.fft.rfft(np.real(x), n=x.shape[-1]), @@ -171,7 +187,7 @@ class FFTTest(xla_test.XLATestCase): def testIRFFT2D(self): def _tf_fn(x): - return spectral_ops.irfft2d( + return signal.irfft2d( x, fft_length=[x.shape[-2].value, 2 * (x.shape[-1].value - 1)]) self._VerifyFftMethod( @@ -195,7 +211,7 @@ class FFTTest(xla_test.XLATestCase): s=[x.shape[-3], x.shape[-2], 2 * (x.shape[-1] - 1)]) def _tf_fn(x): - return spectral_ops.irfft3d( + return signal.irfft3d( x, fft_length=[ x.shape[-3].value, x.shape[-2].value, 2 * (x.shape[-1].value - 1) @@ -203,6 +219,30 @@ class FFTTest(xla_test.XLATestCase): self._VerifyFftMethod(INNER_DIMS_3D, _to_input, _to_expected, _tf_fn) + def testIRFFT3DMismatchedSize(self): + + def _to_input(x): + return np.fft.rfftn( + np.real(x), + axes=(-3, -2, -1), + s=[x.shape[-3] // 2, x.shape[-2], x.shape[-1] * 2]) + + def _to_expected(x): + return np.fft.irfftn( + x, + axes=(-3, -2, -1), + s=[x.shape[-3] // 2, x.shape[-2], x.shape[-1] * 2]) + + def _tf_fn(x): + return signal.irfft3d( + x, + fft_length=[ + x.shape[-3].value // 2, x.shape[-2].value, x.shape[-1].value * 2 + ]) + + self._VerifyFftMethod(INNER_DIMS_3D, _to_input, _to_expected, _tf_fn) + + if __name__ == "__main__": googletest.main() diff --git a/tensorflow/compiler/tests/fifo_queue_test.py b/tensorflow/compiler/tests/fifo_queue_test.py index 8c7edfd277c992c35a81dd5f261256a86352254e..91d77d2f791834346f43aecb60d116ddbf2faa6e 100644 --- a/tensorflow/compiler/tests/fifo_queue_test.py +++ b/tensorflow/compiler/tests/fifo_queue_test.py @@ -129,7 +129,7 @@ class FIFOQueueTest(xla_test.XLATestCase): enqueue_op.run() for i in xrange(len(elems)): - vals = dequeued_t.eval() + vals = self.evaluate(dequeued_t) self.assertEqual([elems[i]], vals) def testEnqueueAndBlockingDequeue(self): @@ -192,9 +192,9 @@ class FIFOQueueTest(xla_test.XLATestCase): self.assertEqual([], size.get_shape()) enqueue_op.run() - self.assertEqual(1, size.eval()) + self.assertEqual(1, self.evaluate(size)) dequeued_t.op.run() - self.assertEqual(0, size.eval()) + self.assertEqual(0, self.evaluate(size)) if __name__ == "__main__": diff --git a/tensorflow/compiler/tests/ftrl_test.py b/tensorflow/compiler/tests/ftrl_test.py index f1b87a5ffb73bed62a80abaa152d335f64d970c5..b078053cdbd6d129645734492d34dd25d28ab3ef 100644 --- a/tensorflow/compiler/tests/ftrl_test.py +++ b/tensorflow/compiler/tests/ftrl_test.py @@ -50,14 +50,14 @@ class FtrlOptimizerTest(xla_test.XLATestCase): ftrl_update = opt.apply_gradients(zip([grads0, grads1], [var0, var1])) variables.global_variables_initializer().run() # Fetch params to validate initial values - self.assertAllClose([0.0, 0.0], var0.eval()) - self.assertAllClose([0.0, 0.0], var1.eval()) + self.assertAllClose([0.0, 0.0], self.evaluate(var0)) + self.assertAllClose([0.0, 0.0], self.evaluate(var1)) # Run Ftrl for a few steps for _ in range(steps): ftrl_update.run() - return var0.eval(), var1.eval() + return self.evaluate(var0), self.evaluate(var1) def equivAdagradTest_AdagradPart(self, steps, dtype): var0, var1, grads0, grads1 = self.initVariableAndGradient(dtype) @@ -65,14 +65,14 @@ class FtrlOptimizerTest(xla_test.XLATestCase): adagrad_update = opt.apply_gradients(zip([grads0, grads1], [var0, var1])) variables.global_variables_initializer().run() # Fetch params to validate initial values - self.assertAllClose([0.0, 0.0], var0.eval()) - self.assertAllClose([0.0, 0.0], var1.eval()) + self.assertAllClose([0.0, 0.0], self.evaluate(var0)) + self.assertAllClose([0.0, 0.0], self.evaluate(var1)) # Run Adagrad for a few steps for _ in range(steps): adagrad_update.run() - return var0.eval(), var1.eval() + return self.evaluate(var0), self.evaluate(var1) def equivGradientDescentTest_FtrlPart(self, steps, dtype): var0, var1, grads0, grads1 = self.initVariableAndGradient(dtype) @@ -85,14 +85,14 @@ class FtrlOptimizerTest(xla_test.XLATestCase): ftrl_update = opt.apply_gradients(zip([grads0, grads1], [var0, var1])) variables.global_variables_initializer().run() # Fetch params to validate initial values - self.assertAllClose([0.0, 0.0], var0.eval()) - self.assertAllClose([0.0, 0.0], var1.eval()) + self.assertAllClose([0.0, 0.0], self.evaluate(var0)) + self.assertAllClose([0.0, 0.0], self.evaluate(var1)) # Run Ftrl for a few steps for _ in range(steps): ftrl_update.run() - return var0.eval(), var1.eval() + return self.evaluate(var0), self.evaluate(var1) def equivGradientDescentTest_GradientDescentPart(self, steps, dtype): var0, var1, grads0, grads1 = self.initVariableAndGradient(dtype) @@ -100,14 +100,14 @@ class FtrlOptimizerTest(xla_test.XLATestCase): sgd_update = opt.apply_gradients(zip([grads0, grads1], [var0, var1])) variables.global_variables_initializer().run() # Fetch params to validate initial values - self.assertAllClose([0.0, 0.0], var0.eval()) - self.assertAllClose([0.0, 0.0], var1.eval()) + self.assertAllClose([0.0, 0.0], self.evaluate(var0)) + self.assertAllClose([0.0, 0.0], self.evaluate(var1)) # Run GradientDescent for a few steps for _ in range(steps): sgd_update.run() - return var0.eval(), var1.eval() + return self.evaluate(var0), self.evaluate(var1) def testFtrlwithoutRegularization(self): for dtype in self.float_types: @@ -124,8 +124,8 @@ class FtrlOptimizerTest(xla_test.XLATestCase): ftrl_update = opt.apply_gradients(zip([grads0, grads1], [var0, var1])) variables.global_variables_initializer().run() # Fetch params to validate initial values - self.assertAllClose([0.0, 0.0], var0.eval()) - self.assertAllClose([0.0, 0.0], var1.eval()) + self.assertAllClose([0.0, 0.0], self.evaluate(var0)) + self.assertAllClose([0.0, 0.0], self.evaluate(var1)) # Run 3 steps FTRL for _ in range(3): @@ -134,12 +134,12 @@ class FtrlOptimizerTest(xla_test.XLATestCase): # Validate updated params self.assertAllCloseAccordingToType( np.array([-2.60260963, -4.29698515]), - var0.eval(), - float_rtol=1e-5, + self.evaluate(var0), + float_rtol=1e-4, half_rtol=1e-2) self.assertAllCloseAccordingToType( np.array([-0.28432083, -0.56694895]), - var1.eval(), + self.evaluate(var1), float_rtol=1e-5, half_rtol=1e-2) @@ -158,8 +158,8 @@ class FtrlOptimizerTest(xla_test.XLATestCase): ftrl_update = opt.apply_gradients(zip([grads0, grads1], [var0, var1])) variables.global_variables_initializer().run() # Fetch params to validate initial values - self.assertAllClose([1.0, 2.0], var0.eval()) - self.assertAllClose([4.0, 3.0], var1.eval()) + self.assertAllClose([1.0, 2.0], self.evaluate(var0)) + self.assertAllClose([4.0, 3.0], self.evaluate(var1)) # Run 3 steps FTRL for _ in range(3): @@ -167,9 +167,14 @@ class FtrlOptimizerTest(xla_test.XLATestCase): # Validate updated params self.assertAllCloseAccordingToType( - np.array([-2.55607247, -3.98729396]), var0.eval(), 1e-5, 1e-5) + np.array([-2.55607247, -3.98729396]), + self.evaluate(var0), + 1e-5, + 1e-5, + float_rtol=1e-4) self.assertAllCloseAccordingToType( - np.array([-0.28232238, -0.56096673]), var1.eval(), 1e-5, 1e-5) + np.array([-0.28232238, -0.56096673]), self.evaluate(var1), 1e-5, + 1e-5) def testFtrlWithL1(self): for dtype in self.float_types: @@ -186,8 +191,8 @@ class FtrlOptimizerTest(xla_test.XLATestCase): ftrl_update = opt.apply_gradients(zip([grads0, grads1], [var0, var1])) variables.global_variables_initializer().run() # Fetch params to validate initial values - self.assertAllClose([1.0, 2.0], var0.eval()) - self.assertAllClose([4.0, 3.0], var1.eval()) + self.assertAllClose([1.0, 2.0], self.evaluate(var0)) + self.assertAllClose([4.0, 3.0], self.evaluate(var1)) # Run 10 steps FTRL for _ in range(10): @@ -196,12 +201,14 @@ class FtrlOptimizerTest(xla_test.XLATestCase): # Validate updated params self.assertAllCloseAccordingToType( np.array([-7.66718769, -10.91273689]), - var0.eval(), + self.evaluate(var0), rtol=1e-4, bfloat16_rtol=1e-1, bfloat16_atol=1e-1) self.assertAllCloseAccordingToType( - np.array([-0.93460727, -1.86147261]), var1.eval(), rtol=1e-4) + np.array([-0.93460727, -1.86147261]), + self.evaluate(var1), + rtol=1e-4) def testFtrlWithL1_L2(self): for dtype in self.float_types: @@ -218,8 +225,8 @@ class FtrlOptimizerTest(xla_test.XLATestCase): ftrl_update = opt.apply_gradients(zip([grads0, grads1], [var0, var1])) variables.global_variables_initializer().run() # Fetch params to validate initial values - self.assertAllClose([1.0, 2.0], var0.eval()) - self.assertAllClose([4.0, 3.0], var1.eval()) + self.assertAllClose([1.0, 2.0], self.evaluate(var0)) + self.assertAllClose([4.0, 3.0], self.evaluate(var1)) # Run 10 steps FTRL for _ in range(10): @@ -227,9 +234,13 @@ class FtrlOptimizerTest(xla_test.XLATestCase): # Validate updated params self.assertAllCloseAccordingToType( - np.array([-0.24059935, -0.46829352]), var0.eval(), rtol=1e-5) + np.array([-0.24059935, -0.46829352]), + self.evaluate(var0), + rtol=1e-5) self.assertAllCloseAccordingToType( - np.array([-0.02406147, -0.04830509]), var1.eval(), rtol=1e-5) + np.array([-0.02406147, -0.04830509]), + self.evaluate(var1), + rtol=1e-5) def testFtrlWithL1_L2_L2Shrinkage(self): """Test the new FTRL op with support for l2 shrinkage. @@ -253,8 +264,8 @@ class FtrlOptimizerTest(xla_test.XLATestCase): ftrl_update = opt.apply_gradients(zip([grads0, grads1], [var0, var1])) variables.global_variables_initializer().run() # Fetch params to validate initial values - self.assertAllCloseAccordingToType([1.0, 2.0], var0.eval()) - self.assertAllCloseAccordingToType([4.0, 3.0], var1.eval()) + self.assertAllCloseAccordingToType([1.0, 2.0], self.evaluate(var0)) + self.assertAllCloseAccordingToType([4.0, 3.0], self.evaluate(var1)) # Run 10 steps FTRL for _ in range(10): @@ -262,9 +273,13 @@ class FtrlOptimizerTest(xla_test.XLATestCase): # Validate updated params self.assertAllCloseAccordingToType( - np.array([-0.22578996, -0.44345799]), var0.eval(), rtol=1e-4) + np.array([-0.22578996, -0.44345799]), + self.evaluate(var0), + rtol=1e-4) self.assertAllCloseAccordingToType( - np.array([-0.14378493, -0.13229476]), var1.eval(), rtol=1e-4) + np.array([-0.14378493, -0.13229476]), + self.evaluate(var1), + rtol=1e-4) def testFtrlWithL2ShrinkageDoesNotChangeLrSchedule(self): """Verifies that l2 shrinkage in FTRL does not change lr schedule.""" @@ -290,8 +305,8 @@ class FtrlOptimizerTest(xla_test.XLATestCase): update1 = opt1.apply_gradients([(grads1, var1)]) variables.global_variables_initializer().run() - self.assertAllCloseAccordingToType([1.0, 2.0], var0.eval()) - self.assertAllCloseAccordingToType([1.0, 2.0], var1.eval()) + self.assertAllCloseAccordingToType([1.0, 2.0], self.evaluate(var0)) + self.assertAllCloseAccordingToType([1.0, 2.0], self.evaluate(var1)) # Run 10 steps FTRL for _ in range(10): @@ -300,7 +315,7 @@ class FtrlOptimizerTest(xla_test.XLATestCase): # var0 is experiencing L2 shrinkage so it should be smaller than var1 # in magnitude. - self.assertTrue((var0.eval()**2 < var1.eval()**2).all()) + self.assertTrue((var0.eval()**2 < self.evaluate(var1)**2).all()) accum0 = list(opt0._slots["accum"].values())[0].eval() accum1 = list(opt1._slots["accum"].values())[0].eval() # L2 shrinkage should not change how we update grad accumulator. diff --git a/tensorflow/compiler/tests/function_test.py b/tensorflow/compiler/tests/function_test.py index b1891b918c6584abce9da382088ed0037f5319fb..a61827c2ae44de117abad5b7db5c6bcd78fa171e 100644 --- a/tensorflow/compiler/tests/function_test.py +++ b/tensorflow/compiler/tests/function_test.py @@ -40,7 +40,7 @@ class FunctionTest(xla_test.XLATestCase): bval = np.array([5, 6, 7, 8]).reshape([2, 2]).astype(np.float32) expected = APlus2B(aval, bval) - with self.cached_session() as sess: + with self.cached_session(): @function.Defun(dtypes.float32, dtypes.float32) def Foo(a, b): @@ -50,7 +50,7 @@ class FunctionTest(xla_test.XLATestCase): b = constant_op.constant(bval, name="b") with self.test_scope(): call_f = Foo(a, b) - result = sess.run(call_f) + result = self.evaluate(call_f) self.assertAllClose(result, expected, rtol=1e-3) def testNestedFunctions(self): @@ -66,7 +66,7 @@ class FunctionTest(xla_test.XLATestCase): bval = np.array([4, 3, 2, 1]).reshape([2, 2]).astype(np.float32) expected = APlus2B(aval, bval) - with self.cached_session() as sess: + with self.cached_session(): @function.Defun(dtypes.float32, dtypes.float32) def Foo(a, b): @@ -76,7 +76,7 @@ class FunctionTest(xla_test.XLATestCase): b = constant_op.constant(bval, name="b") with self.test_scope(): call_g = Foo(a, b) - result = sess.run(call_g) + result = self.evaluate(call_g) self.assertAllClose(result, expected, rtol=1e-3) def testFunctionMultipleRetvals(self): @@ -90,7 +90,7 @@ class FunctionTest(xla_test.XLATestCase): bval = np.array([5, 6, 7, 8]).reshape([2, 2]).astype(np.float32) expected = Func(aval, bval) - with self.cached_session() as sess: + with self.cached_session(): @function.Defun(dtypes.float32, dtypes.float32) def Foo(a, b): @@ -100,7 +100,7 @@ class FunctionTest(xla_test.XLATestCase): b = constant_op.constant(bval, name="b") with self.test_scope(): call_f = Foo(a, b) - result = sess.run(call_f) + result = self.evaluate(call_f) self.assertAllClose(result, expected, rtol=1e-3) def testCompileTimeConstantsInDefun(self): diff --git a/tensorflow/compiler/tests/image_ops_test.py b/tensorflow/compiler/tests/image_ops_test.py index d67b16f8e9e7320d5717b0203be340a2356e53d0..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)) + 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), @@ -477,8 +588,8 @@ class ResizeBilinearTest(xla_test.XLATestCase): for dtype in self.float_types: self._assertForwardOpMatchesExpected( np.array([[1, 2], [3, 4]], dtype=dtype), [3, 3], - expected=np.array( - [[1, 1.5, 2], [2, 2.5, 3], [3, 3.5, 4]], dtype=np.float32)) + expected=np.array([[1, 1.5, 2], [2, 2.5, 3], [3, 3.5, 4]], + dtype=np.float32)) def testAlignCorners2x2To3x3Grad(self): self._assertBackwardOpMatchesExpected( @@ -498,8 +609,8 @@ class ResizeBilinearTest(xla_test.XLATestCase): np.array([[7, 13], [22, 4]], dtype=np.float32), input_shape=[3, 3], dtype=dtype, - expected=np.array( - [[7, 0, 13], [0, 0, 0], [22, 0, 4]], dtype=np.float32)) + expected=np.array([[7, 0, 13], [0, 0, 0], [22, 0, 4]], + dtype=np.float32)) def testAlignCorners4x4To3x3(self): for dtype in self.float_types: @@ -507,8 +618,8 @@ class ResizeBilinearTest(xla_test.XLATestCase): np.array( [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]], dtype=dtype), [3, 3], - expected=np.array( - [[1, 2.5, 4], [7, 8.5, 10], [13, 14.5, 16]], dtype=np.float32)) + expected=np.array([[1, 2.5, 4], [7, 8.5, 10], [13, 14.5, 16]], + dtype=np.float32)) def testAlignCorners4x4To3x3Grad(self): for dtype in self.float_types: @@ -516,41 +627,39 @@ class ResizeBilinearTest(xla_test.XLATestCase): np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=np.float32), input_shape=[4, 4], dtype=dtype, - expected=np.array( - [[1, 1, 1, 3], [2, 1.25, 1.25, 3], [2, 1.25, 1.25, 3], - [7, 4, 4, 9]], - dtype=np.float32)) + expected=np.array([[1, 1, 1, 3], [2, 1.25, 1.25, 3], + [2, 1.25, 1.25, 3], [7, 4, 4, 9]], + dtype=np.float32)) def testAlignCorners3x3To9x9(self): for dtype in self.float_types: self._assertForwardOpMatchesExpected( np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=dtype), [9, 9], expected=np.array( - [[1.0, 1.25, 1.50, 1.75, 2.00, 2.25, 2.50, 2.75, 3.00], [ - 1.75, 2.00, 2.25, 2.50, 2.75, 3.00, 3.25, 3.50, 3.75 - ], [2.50, 2.75, 3.00, 3.25, 3.50, 3.75, 4.00, 4.25, 4.50], [ - 3.25, 3.50, 3.75, 4.00, 4.25, 4.50, 4.75, 5.00, 5.25 - ], [4.00, 4.25, 4.50, 4.75, 5.00, 5.25, 5.50, 5.75, 6.00], [ - 4.75, 5.00, 5.25, 5.50, 5.75, 6.00, 6.25, 6.50, 6.75 - ], [5.50, 5.75, 6.00, 6.25, 6.50, 6.75, 7.00, 7.25, 7.50], [ - 6.25, 6.50, 6.75, 7.00, 7.25, 7.50, 7.75, 8.00, 8.25 - ], [7.00, 7.25, 7.50, 7.75, 8.00, 8.25, 8.50, 8.75, 9.00]], + [[1.0, 1.25, 1.50, 1.75, 2.00, 2.25, 2.50, 2.75, 3.00], + [1.75, 2.00, 2.25, 2.50, 2.75, 3.00, 3.25, 3.50, 3.75], + [2.50, 2.75, 3.00, 3.25, 3.50, 3.75, 4.00, 4.25, 4.50], + [3.25, 3.50, 3.75, 4.00, 4.25, 4.50, 4.75, 5.00, 5.25], + [4.00, 4.25, 4.50, 4.75, 5.00, 5.25, 5.50, 5.75, 6.00], + [4.75, 5.00, 5.25, 5.50, 5.75, 6.00, 6.25, 6.50, 6.75], + [5.50, 5.75, 6.00, 6.25, 6.50, 6.75, 7.00, 7.25, 7.50], + [6.25, 6.50, 6.75, 7.00, 7.25, 7.50, 7.75, 8.00, 8.25], + [7.00, 7.25, 7.50, 7.75, 8.00, 8.25, 8.50, 8.75, 9.00]], dtype=np.float32)) def testAlignCorners3x3To9x9Grad(self): for dtype in self.float_types: self._assertBackwardOpMatchesExpected( - np.array( - [[1.00, 1.25, 1.50, 1.75, 2.00, 2.25, 2.50, 2.75, 3.00], [ - 1.75, 2.00, 2.25, 2.50, 2.75, 3.00, 3.25, 3.50, 3.75 - ], [2.50, 2.75, 3.00, 3.25, 3.50, 3.75, 4.00, 4.25, 4.50], [ - 3.25, 3.50, 3.75, 4.00, 4.25, 4.50, 4.75, 5.00, 5.25 - ], [4.00, 4.25, 4.50, 4.75, 5.00, 5.25, 5.50, 5.75, 6.00], [ - 4.75, 5.00, 5.25, 5.50, 5.75, 6.00, 6.25, 6.50, 6.75 - ], [5.50, 5.75, 6.00, 6.25, 6.50, 6.75, 7.00, 7.25, 7.50], [ - 6.25, 6.50, 6.75, 7.00, 7.25, 7.50, 7.75, 8.00, 8.25 - ], [7.00, 7.25, 7.50, 7.75, 8.00, 8.25, 8.50, 8.75, 9.00]], - dtype=np.float32), + np.array([[1.00, 1.25, 1.50, 1.75, 2.00, 2.25, 2.50, 2.75, 3.00], + [1.75, 2.00, 2.25, 2.50, 2.75, 3.00, 3.25, 3.50, 3.75], + [2.50, 2.75, 3.00, 3.25, 3.50, 3.75, 4.00, 4.25, 4.50], + [3.25, 3.50, 3.75, 4.00, 4.25, 4.50, 4.75, 5.00, 5.25], + [4.00, 4.25, 4.50, 4.75, 5.00, 5.25, 5.50, 5.75, 6.00], + [4.75, 5.00, 5.25, 5.50, 5.75, 6.00, 6.25, 6.50, 6.75], + [5.50, 5.75, 6.00, 6.25, 6.50, 6.75, 7.00, 7.25, 7.50], + [6.25, 6.50, 6.75, 7.00, 7.25, 7.50, 7.75, 8.00, 8.25], + [7.00, 7.25, 7.50, 7.75, 8.00, 8.25, 8.50, 8.75, 9.00]], + dtype=np.float32), input_shape=[3, 3], dtype=dtype, expected=np.array( @@ -571,12 +680,12 @@ class ResizeBilinearTest(xla_test.XLATestCase): (np.array([[0, 1, 2, 3, 4, 5, 6, 7]], dtype=np.float32) + np.array( [[0], [1], [2], [3], [4], [5], [6], [7]], dtype=np.float32)) * 15.0, [16, 16], - expected=7 * (np.array( - [[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]], - dtype=np.float32) + np.array( - [[0], [1], [2], [3], [4], [5], [6], [7], [8], [9], [10], [11], - [12], [13], [14], [15]], - dtype=np.float32)), + expected=7 * + (np.array([[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]], + dtype=np.float32) + + np.array([[0], [1], [2], [3], [4], [5], [6], [7], [8], [9], [10], [11], + [12], [13], [14], [15]], + dtype=np.float32)), large_tolerance=True) def testNonAlignCorners3x2To6x4(self): @@ -600,6 +709,26 @@ class ResizeBilinearTest(xla_test.XLATestCase): expected=np.array(expected_data, dtype=dtype), align_corners=False) + def testNonAlignCorners3x2To6x4Batch2(self): + input_data = [[[64, 32], [32, 64], [50, 100]], [[32, 16], [16, 32], + [25, 50]]] + expected_data = [[[64.0, 48.0, 32.0, 32.0], [48.0, 48.0, 48.0, 48.0], + [32.0, 48.0, 64.0, 64.0], [41.0, 61.5, 82.0, 82.0], + [50.0, 75.0, 100.0, 100.0], [50.0, 75.0, 100.0, 100.0]], + [[32.0, 24.0, 16.0, 16.0], [24.0, 24.0, 24.0, 24.0], + [16.0, 24.0, 32.0, 32.0], [20.5, 30.75, 41.0, 41.0], + [25.0, 37.5, 50.0, 50.0], [25.0, 37.5, 50.0, 50.0]]] + + for dtype in self.float_types: + input_image = np.array(input_data, dtype=dtype) + expected = np.array(expected_data, dtype=dtype) + with self.cached_session() as sess, self.test_scope(): + image = array_ops.placeholder(input_image.dtype) + resized = gen_image_ops.resize_bilinear( + image, [6, 4], align_corners=False) + out = sess.run(resized, {image: input_image[:, :, :, np.newaxis]}) + self.assertAllClose(expected[:, :, :, np.newaxis], out) + class NonMaxSuppressionTest(xla_test.XLATestCase): @@ -804,5 +933,6 @@ class NonMaxSuppressionTest(xla_test.XLATestCase): self.assertEqual(num_valid, 3) self.assertAllClose(indices_tf[:num_valid], [0, 2, 4]) + if __name__ == "__main__": test.main() diff --git a/tensorflow/compiler/tests/jit_test.py b/tensorflow/compiler/tests/jit_test.py index 8778b54dfaf35003c83cf2ab03e9e218c60c98ed..dbea9849e217519874352b789588a2af62f1c826 100644 --- a/tensorflow/compiler/tests/jit_test.py +++ b/tensorflow/compiler/tests/jit_test.py @@ -75,7 +75,7 @@ def RunMetadataLabels(run_metadata): def InLabels(labels, substr): """Returns true iff one of the labels contains substr.""" - return any([substr in x for x in labels]) + return any(substr in x for x in labels) def MetadataHasXlaRunOp(run_metadata): @@ -539,6 +539,20 @@ class LazyCompilationTest(test.TestCase): x = array_ops.placeholder(dtypes.float32) y = CompiledFunction(x) + # The very first run of the cluster is always compiled (non-lazily). + run_metadata_for_first_run = config_pb2.RunMetadata() + sess.run( + y, + feed_dict={x: [2., 10., 19., 77., 100.]}, + run_metadata=run_metadata_for_first_run, + options=config_pb2.RunOptions( + trace_level=config_pb2.RunOptions.FULL_TRACE)) + self.assertTrue( + InLabels( + RunMetadataLabels(run_metadata_for_first_run), "_XlaCompile")) + self.assertTrue( + InLabels(RunMetadataLabels(run_metadata_for_first_run), "_XlaRun")) + run_metadata_before_warmup = config_pb2.RunMetadata() sess.run( y, @@ -579,6 +593,67 @@ class LazyCompilationTest(test.TestCase): self.assertFalse( InLabels(RunMetadataLabels(run_metadata_for_new_shape), "_XlaRun")) + def testIsMegamorphic(self): + + @function.Defun(compiled=True) + def CompiledFunction(x): + return math_ops.log(x) + + with session_lib.Session(config=NoRewriteSessionConfig()) as sess: + x = array_ops.placeholder(dtypes.float32) + y = CompiledFunction(x) + + # Make the cluster go megamorphic by running it with lots of shape + # signatures where the cluster is executed with each signature only a few + # times. Then check that we don't compile the cluster ever again. + + for shape in range(10, 50): + for _ in range(0, 49): + sess.run(y, feed_dict={x: [0.] * shape}) + + for _ in range(0, 50): + run_metadata = config_pb2.RunMetadata() + sess.run( + y, + feed_dict={x: [0.] * 60}, + run_metadata=run_metadata, + options=config_pb2.RunOptions( + trace_level=config_pb2.RunOptions.FULL_TRACE)) + self.assertTrue( + InLabels(RunMetadataLabels(run_metadata), "_XlaCompile")) + self.assertFalse(InLabels(RunMetadataLabels(run_metadata), "_XlaRun")) + + def testIsNotMegamorphic(self): + + @function.Defun(compiled=True) + def CompiledFunction(x): + return math_ops.log(x) + + with session_lib.Session(config=NoRewriteSessionConfig()) as sess: + x = array_ops.placeholder(dtypes.float32) + y = CompiledFunction(x) + + # Run the cluster with lots of shape signatures, but in a way that it + # isn't megamorphic (i.e. each shape signature sees a lot of executions). + # Then check that the cluster has not been marked as megamorphic. + + for shape in range(10, 50): + for _ in range(0, 1000): + sess.run(y, feed_dict={x: [0.] * shape}) + + for _ in range(0, 10): + sess.run(y, feed_dict={x: [0.] * 60}) + + run_metadata = config_pb2.RunMetadata() + sess.run( + y, + feed_dict={x: [0.] * 60}, + run_metadata=run_metadata, + options=config_pb2.RunOptions( + trace_level=config_pb2.RunOptions.FULL_TRACE)) + self.assertTrue(InLabels(RunMetadataLabels(run_metadata), "_XlaCompile")) + self.assertTrue(InLabels(RunMetadataLabels(run_metadata), "_XlaRun")) + if __name__ == "__main__": os.environ["TF_XLA_FLAGS"] = ("--tf_xla_enable_lazy_compilation=true " + diff --git a/tensorflow/compiler/tests/listdiff_op_test.py b/tensorflow/compiler/tests/listdiff_op_test.py index 58622114e4f552fb71db9b040a39b57d7da0037c..0210201fa71a6e790e94667073ab4dba542537a5 100644 --- a/tensorflow/compiler/tests/listdiff_op_test.py +++ b/tensorflow/compiler/tests/listdiff_op_test.py @@ -33,13 +33,13 @@ class ListDiffTest(xla_test.XLATestCase): def _testListDiff(self, x, y, out, idx): for dtype in [dtypes.int32, dtypes.int64]: for index_dtype in [dtypes.int32, dtypes.int64]: - with self.cached_session() as sess: + with self.cached_session(): x_tensor = ops.convert_to_tensor(x, dtype=dtype) y_tensor = ops.convert_to_tensor(y, dtype=dtype) with self.test_scope(): out_tensor, idx_tensor = array_ops.listdiff( x_tensor, y_tensor, out_idx=index_dtype) - tf_out, tf_idx = sess.run([out_tensor, idx_tensor]) + tf_out, tf_idx = self.evaluate([out_tensor, idx_tensor]) self.assertAllEqual(out, tf_out) self.assertAllEqual(idx, tf_idx) self.assertEqual(1, out_tensor.get_shape().ndims) diff --git a/tensorflow/compiler/tests/lrn_ops_test.py b/tensorflow/compiler/tests/lrn_ops_test.py index c6ad67993e8bc196a74c9a328df8c9200c92c575..5dddf6ae4e8c8a3d5e9eb7b2c62298df02a0093c 100644 --- a/tensorflow/compiler/tests/lrn_ops_test.py +++ b/tensorflow/compiler/tests/lrn_ops_test.py @@ -120,8 +120,8 @@ class LRNTest(xla_test.XLATestCase): with self.test_scope(): actual = gen_nn_ops.lrn_grad(out_grads, in_image, out_image, depth_radius, bias, alpha, beta) - expected_val = expected.eval() - actual_val = actual.eval() + expected_val = self.evaluate(expected) + actual_val = self.evaluate(actual) self.assertAllClose(actual_val, expected_val, rtol=1e-3) diff --git a/tensorflow/compiler/tests/lstm_test.py b/tensorflow/compiler/tests/lstm_test.py index 265c0b6d1412de7be3a5bf5e79129cb330ceb162..776ed899e68ddd3893b8bb30b7c8034297aa6515 100644 --- a/tensorflow/compiler/tests/lstm_test.py +++ b/tensorflow/compiler/tests/lstm_test.py @@ -88,8 +88,8 @@ class LSTMTest(test.TestCase): (basename, m_prev_scalar, c_prev_scalar, pad_scalar)) # Initialize variables and run the unrolled LSTM step. - sess.run(variables.global_variables_initializer()) - return sess.run([m, c]) + self.evaluate(variables.global_variables_initializer()) + return self.evaluate([m, c]) def testLSTMCell(self): # Run with all-0 weights, no padding. @@ -173,8 +173,8 @@ class LSTMTest(test.TestCase): (basename, m_init_scalar, c_init_scalar, pad_scalar)) # Initialize variables and run the unrolled LSTM layer. - sess.run(variables.global_variables_initializer()) - return sess.run(out_seq) + self.evaluate(variables.global_variables_initializer()) + return self.evaluate(out_seq) def testLSTMLayer(self): # Run with all-0 weights, no padding. diff --git a/tensorflow/compiler/tests/momentum_test.py b/tensorflow/compiler/tests/momentum_test.py index f77521a7c49dba39849869ddceb7c0e885147722..3416f7dbd6bdd264bf79785084f981f5b07cb8a9 100644 --- a/tensorflow/compiler/tests/momentum_test.py +++ b/tensorflow/compiler/tests/momentum_test.py @@ -61,37 +61,43 @@ class MomentumOptimizerTest(xla_test.XLATestCase): self.assertFalse(slot1 in variables.trainable_variables()) # Fetch params to validate initial values - self.assertAllClose([1.0, 2.0], var0.eval()) - self.assertAllClose([3.0, 4.0], var1.eval()) + self.assertAllClose([1.0, 2.0], self.evaluate(var0)) + self.assertAllClose([3.0, 4.0], self.evaluate(var1)) # Step 1: the momentum accumulators where 0. So we should see a normal # update: v -= grad * learning_rate mom_update.run() # Check that the momentum accumulators have been updated. - self.assertAllCloseAccordingToType(np.array([0.1, 0.1]), slot0.eval()) - self.assertAllCloseAccordingToType(np.array([0.01, 0.01]), slot1.eval()) + self.assertAllCloseAccordingToType( + np.array([0.1, 0.1]), self.evaluate(slot0)) + self.assertAllCloseAccordingToType( + np.array([0.01, 0.01]), self.evaluate(slot1)) # Check that the parameters have been updated. self.assertAllCloseAccordingToType( - np.array([1.0 - (0.1 * 2.0), 2.0 - (0.1 * 2.0)]), var0.eval()) + np.array([1.0 - (0.1 * 2.0), 2.0 - (0.1 * 2.0)]), + self.evaluate(var0)) self.assertAllCloseAccordingToType( - np.array([3.0 - (0.01 * 2.0), 4.0 - (0.01 * 2.0)]), var1.eval()) + np.array([3.0 - (0.01 * 2.0), 4.0 - (0.01 * 2.0)]), + self.evaluate(var1)) # Step 2: the momentum accumulators contain the previous update. mom_update.run() # Check that the momentum accumulators have been updated. self.assertAllCloseAccordingToType( - np.array([(0.9 * 0.1 + 0.1), (0.9 * 0.1 + 0.1)]), slot0.eval()) + np.array([(0.9 * 0.1 + 0.1), (0.9 * 0.1 + 0.1)]), + self.evaluate(slot0)) self.assertAllCloseAccordingToType( - np.array([(0.9 * 0.01 + 0.01), (0.9 * 0.01 + 0.01)]), slot1.eval()) + np.array([(0.9 * 0.01 + 0.01), (0.9 * 0.01 + 0.01)]), + self.evaluate(slot1)) # Check that the parameters have been updated. self.assertAllCloseAccordingToType( np.array([ 1.0 - (0.1 * 2.0) - ((0.9 * 0.1 + 0.1) * 2.0), 2.0 - (0.1 * 2.0) - ((0.9 * 0.1 + 0.1) * 2.0) - ]), var0.eval()) + ]), self.evaluate(var0)) self.assertAllCloseAccordingToType( np.array([ - 2.98 - ((0.9 * 0.01 + 0.01) * 2.0), 3.98 - ( - (0.9 * 0.01 + 0.01) * 2.0) - ]), var1.eval()) + 2.98 - ((0.9 * 0.01 + 0.01) * 2.0), + 3.98 - ((0.9 * 0.01 + 0.01) * 2.0) + ]), self.evaluate(var1)) def testNesterovMomentum(self): for dtype in self.float_types: @@ -115,8 +121,8 @@ class MomentumOptimizerTest(xla_test.XLATestCase): var0_np, accum0_np, var0_np * 0.8, 0.1, 0.9) var1_np, accum1_np = self._update_nesterov_momentum_numpy( var1_np, accum1_np, 0.9, 0.1, 0.9) - self.assertAllCloseAccordingToType(var0_np, var0.eval()) - self.assertAllCloseAccordingToType(var1_np, var1.eval()) + self.assertAllCloseAccordingToType(var0_np, self.evaluate(var0)) + self.assertAllCloseAccordingToType(var1_np, self.evaluate(var1)) def testTensorLearningRateAndMomentum(self): for dtype in self.float_types: @@ -141,37 +147,43 @@ class MomentumOptimizerTest(xla_test.XLATestCase): self.assertFalse(slot1 in variables.trainable_variables()) # Fetch params to validate initial values - self.assertAllClose([1.0, 2.0], var0.eval()) - self.assertAllClose([3.0, 4.0], var1.eval()) + self.assertAllClose([1.0, 2.0], self.evaluate(var0)) + self.assertAllClose([3.0, 4.0], self.evaluate(var1)) # Step 1: the momentum accumulators where 0. So we should see a normal # update: v -= grad * learning_rate mom_update.run() # Check that the momentum accumulators have been updated. - self.assertAllCloseAccordingToType(np.array([0.1, 0.1]), slot0.eval()) - self.assertAllCloseAccordingToType(np.array([0.01, 0.01]), slot1.eval()) + self.assertAllCloseAccordingToType( + np.array([0.1, 0.1]), self.evaluate(slot0)) + self.assertAllCloseAccordingToType( + np.array([0.01, 0.01]), self.evaluate(slot1)) # Check that the parameters have been updated. self.assertAllCloseAccordingToType( - np.array([1.0 - (0.1 * 2.0), 2.0 - (0.1 * 2.0)]), var0.eval()) + np.array([1.0 - (0.1 * 2.0), 2.0 - (0.1 * 2.0)]), + self.evaluate(var0)) self.assertAllCloseAccordingToType( - np.array([3.0 - (0.01 * 2.0), 4.0 - (0.01 * 2.0)]), var1.eval()) + np.array([3.0 - (0.01 * 2.0), 4.0 - (0.01 * 2.0)]), + self.evaluate(var1)) # Step 2: the momentum accumulators contain the previous update. mom_update.run() # Check that the momentum accumulators have been updated. self.assertAllCloseAccordingToType( - np.array([(0.9 * 0.1 + 0.1), (0.9 * 0.1 + 0.1)]), slot0.eval()) + np.array([(0.9 * 0.1 + 0.1), (0.9 * 0.1 + 0.1)]), + self.evaluate(slot0)) self.assertAllCloseAccordingToType( - np.array([(0.9 * 0.01 + 0.01), (0.9 * 0.01 + 0.01)]), slot1.eval()) + np.array([(0.9 * 0.01 + 0.01), (0.9 * 0.01 + 0.01)]), + self.evaluate(slot1)) # Check that the parameters have been updated. self.assertAllCloseAccordingToType( np.array([ 1.0 - (0.1 * 2.0) - ((0.9 * 0.1 + 0.1) * 2.0), 2.0 - (0.1 * 2.0) - ((0.9 * 0.1 + 0.1) * 2.0) - ]), var0.eval()) + ]), self.evaluate(var0)) self.assertAllCloseAccordingToType( np.array([ - 2.98 - ((0.9 * 0.01 + 0.01) * 2.0), 3.98 - ( - (0.9 * 0.01 + 0.01) * 2.0) - ]), var1.eval()) + 2.98 - ((0.9 * 0.01 + 0.01) * 2.0), + 3.98 - ((0.9 * 0.01 + 0.01) * 2.0) + ]), self.evaluate(var1)) if __name__ == "__main__": diff --git a/tensorflow/compiler/tests/oom_test.py b/tensorflow/compiler/tests/oom_test.py deleted file mode 100644 index 7635f89249b7b71e5353e0b7cb1cea5c1f7bca1d..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/tests/oom_test.py +++ /dev/null @@ -1,76 +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. -# ============================================================================== -"""Functional tests for out-of-memory conditions.""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -from tensorflow.compiler.tests import xla_test -from tensorflow.python.framework import dtypes -from tensorflow.python.framework import errors -from tensorflow.python.ops import array_ops -from tensorflow.python.ops import math_ops -from tensorflow.python.ops import nn_ops -from tensorflow.python.platform import googletest - - -class OutOfMemoryTest(xla_test.XLATestCase): - - def testOutputOutOfMemory(self): - """Allocates tensors until out of memory. - - Generates a large rank-1 tensor. The tensor is an output of an XLA - computation, not constant. - - Check that a ResourceExhaustedError is raised and can be caught. - - We spin in a loop generating larger and larger tensors until an OOM event - happens. We may be running sandboxed, so have a small host memory limit, so - any hardcoded value is unlikely to land in the sweet spot between device - memory size and host memory size with stability. - """ - - def test_loop(): - size = int(2e8) - while True: - with self.cached_session(): - # Force the compiled code to not be constant by feeding in a - # parameter. - p = array_ops.placeholder(dtypes.float32, shape=[2, 1, 1]) - with self.test_scope(): - # Create a computation that produces a large R1 tensor as an - # intermediate result. Reduce it down so that if this file was - # compiled without --config=cuda, we don't force a D2H copy of a - # large tensor and potentially OOM the host. - # - # This is a bit tricky because XLA:GPU doesn't currently support RNG - # ops. Here we rely on the fact that XLA doesn't do algebraic - # simplifications on conv(, ). - c = math_ops.reduce_sum( - nn_ops.convolution( - array_ops.ones([1, size, 1]), - p, - padding='SAME', - data_format='NWC')) - - c.eval(feed_dict={p: [[[1.0]], [[2.0]]]}) - size *= 2 - - self.assertRaises(errors.ResourceExhaustedError, test_loop) - - -if __name__ == '__main__': - googletest.main() diff --git a/tensorflow/compiler/tests/placeholder_test.py b/tensorflow/compiler/tests/placeholder_test.py index 77bb839409f0c323ff6ed2c8d6bd105d3003b398..9671ae0ae973ff82d22744a1feb9b4293d94bbdd 100644 --- a/tensorflow/compiler/tests/placeholder_test.py +++ b/tensorflow/compiler/tests/placeholder_test.py @@ -33,7 +33,7 @@ class PlaceholderTest(xla_test.XLATestCase): ph = array_ops.placeholder_with_default(v, shape=[]) out = ph * 2 sess.run(variables.variables_initializer([v])) - self.assertEqual(8.0, sess.run(out)) + self.assertEqual(8.0, self.evaluate(out)) def test_placeholder_with_default_fed(self): with self.cached_session() as sess, self.test_scope(): diff --git a/tensorflow/compiler/tests/powersign_test.py b/tensorflow/compiler/tests/powersign_test.py index 86536da7fed0e2309beb32fee9c7c605491592ed..5b35c20027700b34500a31e174061d7087094b61 100644 --- a/tensorflow/compiler/tests/powersign_test.py +++ b/tensorflow/compiler/tests/powersign_test.py @@ -91,8 +91,8 @@ class PowerSignTest(xla_test.XLATestCase): variables.global_variables_initializer().run() # Fetch params to validate initial values - self.assertAllClose([1.0, 2.0], var0.eval()) - self.assertAllClose([3.0, 4.0], var1.eval()) + self.assertAllClose([1.0, 2.0], self.evaluate(var0)) + self.assertAllClose([3.0, 4.0], self.evaluate(var1)) # Run 7 steps of powersign # first 4 steps with positive gradient @@ -125,8 +125,8 @@ class PowerSignTest(xla_test.XLATestCase): ) # Validate updated params - self.assertAllCloseAccordingToType(var0_np, var0.eval()) - self.assertAllCloseAccordingToType(var1_np, var1.eval()) + self.assertAllCloseAccordingToType(var0_np, self.evaluate(var0)) + self.assertAllCloseAccordingToType(var1_np, self.evaluate(var1)) def testDense(self): decay_steps = 10 diff --git a/tensorflow/compiler/tests/proximal_adagrad_test.py b/tensorflow/compiler/tests/proximal_adagrad_test.py index c41b4171e26af4f7ad0237d7407a5b3691299595..63cc51a470164915b2614a06d18ca1850bb64a3c 100644 --- a/tensorflow/compiler/tests/proximal_adagrad_test.py +++ b/tensorflow/compiler/tests/proximal_adagrad_test.py @@ -45,15 +45,17 @@ class ProximalAdagradOptimizerTest(xla_test.XLATestCase): update = opt.apply_gradients(zip([grads0, grads1], [var0, var1])) variables.global_variables_initializer().run() - self.assertAllClose([0.0, 0.0], var0.eval()) - self.assertAllClose([0.0, 0.0], var1.eval()) + self.assertAllClose([0.0, 0.0], self.evaluate(var0)) + self.assertAllClose([0.0, 0.0], self.evaluate(var1)) # Run 3 steps Proximal Adagrad. for _ in range(3): update.run() - self.assertAllClose(np.array([-2.60260963, -4.29698515]), var0.eval()) - self.assertAllClose(np.array([-0.28432083, -0.56694895]), var1.eval()) + self.assertAllClose( + np.array([-2.60260963, -4.29698515]), self.evaluate(var0)) + self.assertAllClose( + np.array([-0.28432083, -0.56694895]), self.evaluate(var1)) opt_vars = opt.variables() self.assertStartsWith(opt_vars[0].name, var0._shared_name) self.assertStartsWith(opt_vars[1].name, var1._shared_name) @@ -74,14 +76,14 @@ class ProximalAdagradOptimizerTest(xla_test.XLATestCase): update = opt.apply_gradients(zip([grads0, grads1], [var0, var1])) variables.global_variables_initializer().run() - self.assertAllClose([1.0, 2.0], var0.eval()) - self.assertAllClose([4.0, 3.0], var1.eval()) + self.assertAllClose([1.0, 2.0], self.evaluate(var0)) + self.assertAllClose([4.0, 3.0], self.evaluate(var1)) # Run 3 steps Proximal Adagrad. for _ in range(3): update.run() - self.assertAllClose(np.array([-1.60261, -2.296985]), var0.eval()) - self.assertAllClose(np.array([3.715679, 2.433051]), var1.eval()) + self.assertAllClose(np.array([-1.60261, -2.296985]), self.evaluate(var0)) + self.assertAllClose(np.array([3.715679, 2.433051]), self.evaluate(var1)) def testProximalAdagradWithL1(self): with self.cached_session(), self.test_scope(): @@ -98,14 +100,14 @@ class ProximalAdagradOptimizerTest(xla_test.XLATestCase): update = opt.apply_gradients(zip([grads0, grads1], [var0, var1])) variables.global_variables_initializer().run() - self.assertAllClose([1.0, 2.0], var0.eval()) - self.assertAllClose([4.0, 3.0], var1.eval()) + self.assertAllClose([1.0, 2.0], self.evaluate(var0)) + self.assertAllClose([4.0, 3.0], self.evaluate(var1)) # Run 10 steps Proximal Adagrad for _ in range(10): update.run() - self.assertAllClose(np.array([-6.663634, -9.190331]), var0.eval()) - self.assertAllClose(np.array([2.959304, 1.029232]), var1.eval()) + self.assertAllClose(np.array([-6.663634, -9.190331]), self.evaluate(var0)) + self.assertAllClose(np.array([2.959304, 1.029232]), self.evaluate(var1)) def testProximalAdagradWithL1_L2(self): with self.cached_session(), self.test_scope(): @@ -122,15 +124,15 @@ class ProximalAdagradOptimizerTest(xla_test.XLATestCase): update = opt.apply_gradients(zip([grads0, grads1], [var0, var1])) variables.global_variables_initializer().run() - self.assertAllClose([1.0, 2.0], var0.eval()) - self.assertAllClose([4.0, 3.0], var1.eval()) + self.assertAllClose([1.0, 2.0], self.evaluate(var0)) + self.assertAllClose([4.0, 3.0], self.evaluate(var1)) # Run 10 steps Proximal Adagrad. for _ in range(10): update.run() - self.assertAllClose(np.array([-0.0495, -0.0995]), var0.eval()) - self.assertAllClose(np.array([-0.0045, -0.0095]), var1.eval()) + self.assertAllClose(np.array([-0.0495, -0.0995]), self.evaluate(var0)) + self.assertAllClose(np.array([-0.0045, -0.0095]), self.evaluate(var1)) def applyOptimizer(self, opt, steps=5): var0 = resource_variable_ops.ResourceVariable([1.0, 2.0]) @@ -141,14 +143,14 @@ class ProximalAdagradOptimizerTest(xla_test.XLATestCase): update = opt.apply_gradients(zip([grads0, grads1], [var0, var1])) variables.global_variables_initializer().run() - self.assertAllClose([1.0, 2.0], var0.eval()) - self.assertAllClose([3.0, 4.0], var1.eval()) + self.assertAllClose([1.0, 2.0], self.evaluate(var0)) + self.assertAllClose([3.0, 4.0], self.evaluate(var1)) # Run ProximalAdagrad for a few steps for _ in range(steps): update.run() - return var0.eval(), var1.eval() + return self.evaluate(var0), self.evaluate(var1) def testEquivAdagradwithoutRegularization(self): with self.cached_session(), self.test_scope(): diff --git a/tensorflow/compiler/tests/proximal_gradient_descent_test.py b/tensorflow/compiler/tests/proximal_gradient_descent_test.py index 3d808e6b8a71ef9fa60b671d07bfd907e9f58efc..5aec433be765dd0a04bd7ab10d5c39a5a7f48c5c 100644 --- a/tensorflow/compiler/tests/proximal_gradient_descent_test.py +++ b/tensorflow/compiler/tests/proximal_gradient_descent_test.py @@ -42,15 +42,15 @@ class ProximalGradientDescentOptimizerTest(xla_test.XLATestCase): update = opt.apply_gradients(zip([grads0, grads1], [var0, var1])) variables.global_variables_initializer().run() - self.assertAllClose([0.0, 0.0], var0.eval()) - self.assertAllClose([0.0, 0.0], var1.eval()) + self.assertAllClose([0.0, 0.0], self.evaluate(var0)) + self.assertAllClose([0.0, 0.0], self.evaluate(var1)) # Run 3 steps Proximal Gradient Descent. for _ in range(3): update.run() - self.assertAllClose(np.array([-0.9, -1.8]), var0.eval()) - self.assertAllClose(np.array([-0.09, -0.18]), var1.eval()) + self.assertAllClose(np.array([-0.9, -1.8]), self.evaluate(var0)) + self.assertAllClose(np.array([-0.09, -0.18]), self.evaluate(var1)) def testProximalGradientDescentwithoutRegularization2(self): with self.cached_session(), self.test_scope(): @@ -64,15 +64,15 @@ class ProximalGradientDescentOptimizerTest(xla_test.XLATestCase): update = opt.apply_gradients(zip([grads0, grads1], [var0, var1])) variables.global_variables_initializer().run() - self.assertAllClose([1.0, 2.0], var0.eval()) - self.assertAllClose([4.0, 3.0], var1.eval()) + self.assertAllClose([1.0, 2.0], self.evaluate(var0)) + self.assertAllClose([4.0, 3.0], self.evaluate(var1)) # Run 3 steps Proximal Gradient Descent for _ in range(3): update.run() - self.assertAllClose(np.array([0.1, 0.2]), var0.eval()) - self.assertAllClose(np.array([3.91, 2.82]), var1.eval()) + self.assertAllClose(np.array([0.1, 0.2]), self.evaluate(var0)) + self.assertAllClose(np.array([3.91, 2.82]), self.evaluate(var1)) def testProximalGradientDescentWithL1(self): with self.cached_session(), self.test_scope(): @@ -86,15 +86,15 @@ class ProximalGradientDescentOptimizerTest(xla_test.XLATestCase): update = opt.apply_gradients(zip([grads0, grads1], [var0, var1])) variables.global_variables_initializer().run() - self.assertAllClose([1.0, 2.0], var0.eval()) - self.assertAllClose([4.0, 3.0], var1.eval()) + self.assertAllClose([1.0, 2.0], self.evaluate(var0)) + self.assertAllClose([4.0, 3.0], self.evaluate(var1)) # Run 10 steps proximal gradient descent. for _ in range(10): update.run() - self.assertAllClose(np.array([-1.988, -3.988001]), var0.eval()) - self.assertAllClose(np.array([3.67, 2.37]), var1.eval()) + self.assertAllClose(np.array([-1.988, -3.988001]), self.evaluate(var0)) + self.assertAllClose(np.array([3.67, 2.37]), self.evaluate(var1)) def testProximalGradientDescentWithL1_L2(self): with self.cached_session(), self.test_scope(): @@ -108,15 +108,15 @@ class ProximalGradientDescentOptimizerTest(xla_test.XLATestCase): update = opt.apply_gradients(zip([grads0, grads1], [var0, var1])) variables.global_variables_initializer().run() - self.assertAllClose([1.0, 2.0], var0.eval()) - self.assertAllClose([4.0, 3.0], var1.eval()) + self.assertAllClose([1.0, 2.0], self.evaluate(var0)) + self.assertAllClose([4.0, 3.0], self.evaluate(var1)) # Run 10 steps Proximal Gradient Descent for _ in range(10): update.run() - self.assertAllClose(np.array([-0.0495, -0.0995]), var0.eval()) - self.assertAllClose(np.array([-0.0045, -0.0095]), var1.eval()) + self.assertAllClose(np.array([-0.0495, -0.0995]), self.evaluate(var0)) + self.assertAllClose(np.array([-0.0045, -0.0095]), self.evaluate(var1)) def applyOptimizer(self, opt, steps=5): var0 = resource_variable_ops.ResourceVariable([1.0, 2.0]) @@ -127,14 +127,14 @@ class ProximalGradientDescentOptimizerTest(xla_test.XLATestCase): update = opt.apply_gradients(zip([grads0, grads1], [var0, var1])) variables.global_variables_initializer().run() - self.assertAllClose([1.0, 2.0], var0.eval()) - self.assertAllClose([3.0, 4.0], var1.eval()) + self.assertAllClose([1.0, 2.0], self.evaluate(var0)) + self.assertAllClose([3.0, 4.0], self.evaluate(var1)) # Run ProximalAdagrad for a few steps for _ in range(steps): update.run() - return var0.eval(), var1.eval() + return self.evaluate(var0), self.evaluate(var1) def testEquivGradientDescentwithoutRegularization(self): with self.cached_session(), self.test_scope(): diff --git a/tensorflow/compiler/tests/qr_op_test.py b/tensorflow/compiler/tests/qr_op_test.py index 236b1b881dcaffc1a5b0c6395f0605c1d7ef0269..b4d4193e35f9e0e3b23d0242ed076dd811f4ee2b 100644 --- a/tensorflow/compiler/tests/qr_op_test.py +++ b/tensorflow/compiler/tests/qr_op_test.py @@ -63,7 +63,7 @@ class QrOpTest(xla_test.XLATestCase, parameterized.TestCase): # Tests that x[...,:,:]^H * x[...,:,:] is close to the identity. xx = math_ops.matmul(x, x, adjoint_a=True) identity = array_ops.matrix_band_part(array_ops.ones_like(xx), 0, 0) - precision = self.AdjustedNorm(xx.eval() - identity.eval()) + precision = self.AdjustedNorm(xx.eval() - self.evaluate(identity)) self.assertTrue(np.all(precision < 5.0)) def _test(self, dtype, shape, full_matrices): diff --git a/tensorflow/compiler/tests/quantized_ops_test.py b/tensorflow/compiler/tests/quantized_ops_test.py index 80c338513bc9ff6b8e56c5ad6b904af9e06a3715..cd9b728ab314d29e4eb585e00a9131024ea3a207 100644 --- a/tensorflow/compiler/tests/quantized_ops_test.py +++ b/tensorflow/compiler/tests/quantized_ops_test.py @@ -18,11 +18,16 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import math import numpy as np from tensorflow.compiler.tests import xla_test +from tensorflow.compiler.tf2xla.python import xla +from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes +from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops +from tensorflow.python.ops import bitwise_ops from tensorflow.python.ops import math_ops from tensorflow.python.platform import googletest @@ -44,5 +49,55 @@ class QuantizedOpsTest(xla_test.XLATestCase): self.assertAllEqual(value, expected) +class DeuantizedOpsTest(xla_test.XLATestCase): + + def pack_uint8_r2_to_uint32(self, test_input): + num_rows, num_columns = test_input.get_shape().as_list() + num_output_columns = int(math.ceil(num_columns / 4.0)) + padding_input = array_ops.pad( + math_ops.cast(test_input, dtype=dtypes.uint8), + constant_op.constant([[ + 0, + 0, + ], [0, num_output_columns * 4 - num_columns]])) + output = array_ops.zeros([num_rows, num_output_columns], + dtype=dtypes.uint32) + num_elements_per_pack = 4 + shift_bits = 8 + + iota_r1 = math_ops.range(num_output_columns * num_elements_per_pack) + + for p in range(num_elements_per_pack): + selected_index = math_ops.equal( + math_ops.mod(iota_r1, num_elements_per_pack), p) + gather_index = array_ops.boolean_mask(iota_r1, selected_index) + gathered_input = array_ops.gather(padding_input, gather_index, axis=1) + total_shift_bits = shift_bits * (num_elements_per_pack - p - 1) + left_shift_input = bitwise_ops.left_shift( + math_ops.cast(gathered_input, dtype=dtypes.uint32), total_shift_bits) + output = bitwise_ops.bitwise_or(output, left_shift_input) + return output + + def testDequantizeQuint8(self): + num_rows = 100 + num_columns = 3547 + random_input = np.random.normal(128.0, 10.0, [num_rows, num_columns]) + with self.cached_session() as session: + with ops.device("CPU"): + test_input = ops.convert_to_tensor(random_input, dtype=dtypes.float32) + transposed_input = array_ops.transpose(test_input, [1, 0]) + quantized_input = array_ops.quantize(transposed_input, 0.0, 255.0, + dtypes.quint8) + packed_input = self.pack_uint8_r2_to_uint32(quantized_input.output) + with self.test_scope(): + transposed_quantized_output = xla.dequantize(packed_input, 0.0, 255.0, + "MIN_COMBINED", True) + quantized_output = array_ops.slice(transposed_quantized_output, [0, 0], + [num_rows, num_columns]) + + value = session.run(quantized_output) + self.assertAllClose(value, random_input, 1.0) + + if __name__ == "__main__": googletest.main() diff --git a/tensorflow/compiler/tests/random_ops_test.py b/tensorflow/compiler/tests/random_ops_test.py index 36ef6ed5fee78bad10bb1ee0bf3eb7824d05c206..97ffad34c00b8ec16eb1ec109ba5d980e0ce673d 100644 --- a/tensorflow/compiler/tests/random_ops_test.py +++ b/tensorflow/compiler/tests/random_ops_test.py @@ -46,9 +46,9 @@ class RandomOpsTest(xla_test.XLATestCase): # The random-number generator, if working correctly, should produce the # same output multiple times with low probability. - y = sess.run(x) - z = sess.run(x) - w = sess.run(x) + y = self.evaluate(x) + z = self.evaluate(x) + w = self.evaluate(x) # We use exact equality here. If the random-number generator is producing # deterministic output, all three outputs will be bitwise identical. @@ -83,7 +83,7 @@ class RandomOpsTest(xla_test.XLATestCase): with self.test_scope(): x = random_ops.random_uniform( shape=[1000], dtype=dtype, minval=-2, maxval=33) - y = sess.run(x) + y = self.evaluate(x) self.assertTrue((y >= -2).sum() == 1000) self.assertTrue((y < 33).sum() == 1000) @@ -102,7 +102,7 @@ class RandomOpsTest(xla_test.XLATestCase): with self.cached_session() as sess: with self.test_scope(): x = random_ops.truncated_normal(shape=[count], dtype=dtype) - y = sess.run(x) + y = self.evaluate(x) def normal_cdf(x): return .5 * math.erfc(-x / math.sqrt(2)) @@ -111,7 +111,7 @@ class RandomOpsTest(xla_test.XLATestCase): return math.exp(-(x**2) / 2.) / math.sqrt(2 * math.pi) def probit(x, sess=sess): - return sess.run(special_math.ndtri(x)) + return self.evaluate(special_math.ndtri(x)) a = -2. b = 2. @@ -148,7 +148,7 @@ class RandomOpsTest(xla_test.XLATestCase): with self.test_scope(): x = math_ops.range(1 << 16) shuffle = random_ops.random_shuffle(x) - result = sess.run(shuffle) + result = self.evaluate(shuffle) expected = range(1 << 16) # Compare sets to avoid randomness behavior changes but make sure still # have all the values. @@ -159,7 +159,7 @@ class RandomOpsTest(xla_test.XLATestCase): with self.test_scope(): x = array_ops.diag(math_ops.range(20)) shuffle = random_ops.random_shuffle(x) - result = sess.run(shuffle) + result = self.evaluate(shuffle) expected = np.diag(range(20)).flatten() # Compare sets to avoid randomness behavior changes but make sure still # have all the values. diff --git a/tensorflow/compiler/tests/randomized_tests.cc b/tensorflow/compiler/tests/randomized_tests.cc index dc119fb0f8a41a3772a8c9508bf2db657f57de88..1521cc760b85b176acb27c1489640e92ef90e247 100644 --- a/tensorflow/compiler/tests/randomized_tests.cc +++ b/tensorflow/compiler/tests/randomized_tests.cc @@ -45,6 +45,7 @@ limitations under the License. #include #include +#include "absl/algorithm/container.h" #include "absl/container/flat_hash_set.h" #include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" @@ -62,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" @@ -79,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) { @@ -320,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; @@ -452,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)); @@ -759,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(); @@ -796,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()); } @@ -828,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; @@ -844,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; @@ -1370,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( @@ -2465,20 +2505,21 @@ TEST_F(OpTest, Pack) { }); } -// TODO(b/31741898): crashes on GPU. TEST_F(OpTest, Pad) { Repeatedly([this]() { auto type = Choose(kAllXlaTypes); std::vector t_dims = RandomDims(); - // TODO(b/31741996): re-enable DT_INT64 when bug is fixed. - // DataType tpaddings = Choose({DT_INT32, DT_INT64}); - DataType tpaddings = DT_INT32; + DataType tpaddings = Choose({DT_INT32, DT_INT64}); std::vector paddings_vec; - std::uniform_int_distribution distribution(0, 7); for (int i = 0; i < t_dims.size(); ++i) { - paddings_vec.push_back(distribution(generator())); - paddings_vec.push_back(distribution(generator())); + std::uniform_int_distribution pad_distribution(0, t_dims[i]); + int pad_size = pad_distribution(generator()); + std::uniform_int_distribution lower_distribution(0, pad_size); + int low_pad_size = lower_distribution(generator()); + paddings_vec.push_back(low_pad_size); + paddings_vec.push_back(pad_size - low_pad_size); + t_dims[i] -= pad_size; } Tensor paddings; CHECK( @@ -2687,6 +2728,37 @@ TEST_F(OpTest, Reverse) { }); } +TEST_F(OpTest, ReverseSequence) { + Repeatedly([this]() { + std::vector dims = RandomDims(/*min_rank=*/2); + auto type = Choose(kAllXlaTypes); + int64 rank = dims.size(); + + // Choose random batch and sequence dimensions. + std::vector shuffled_dim_ids(rank); + absl::c_iota(shuffled_dim_ids, 0); + absl::c_shuffle(shuffled_dim_ids, generator()); + shuffled_dim_ids.resize(2); + int batch_dim = shuffled_dim_ids[0]; + int seq_dim = shuffled_dim_ids[1]; + + int batch_size = dims[batch_dim]; + int max_seq_len = dims[seq_dim]; + std::vector seq_lens(batch_size); + std::uniform_int_distribution d(0, max_seq_len); + absl::c_generate(seq_lens, [&]() { return d(generator()); }); + + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("ReverseSequence") + .RandomInput(type, dims) + .Input(test::AsTensor(seq_lens)) + .Attr("seq_dim", seq_dim) + .Attr("batch_dim", batch_dim) + .Attr("T", type) + .Attr("Tlen", DT_INT32)); + }); +} + TEST_F(OpTest, ReverseV2) { Repeatedly([this]() { auto type = Choose(kAllXlaTypes); @@ -3313,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, @@ -3333,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"), }; @@ -3349,10 +3454,10 @@ int main(int argc, char** argv) { } // XLA devices register kernels at construction time; create all known devices // to make sure the kernels are registered. - std::vector devices; + std::vector> devices; TF_CHECK_OK(tensorflow::DeviceFactory::AddDevices( tensorflow::SessionOptions(), "", &devices)); - tensorflow::DeviceMgr device_mgr(devices); + tensorflow::DeviceMgr device_mgr(std::move(devices)); tensorflow::Device* ignored; TF_QCHECK_OK( diff --git a/tensorflow/compiler/tests/reduce_ops_test.py b/tensorflow/compiler/tests/reduce_ops_test.py index 132c59c32c9db0c8759bdbb31f8613c3ef88b485..e8fc81bbb5472669c408b8bbdbcdfcdcf461131f 100644 --- a/tensorflow/compiler/tests/reduce_ops_test.py +++ b/tensorflow/compiler/tests/reduce_ops_test.py @@ -91,6 +91,7 @@ class ReduceOpsTest(xla_test.XLATestCase, parameterized.TestCase): np.array([], dtype=np.bool).reshape(0, 3), np.array([[False, True, False], [True, True, False]]), ] + ONES = [np.ones([34000, 2])] def testReduceSumF32(self, index_dtype): self._testReduction(math_ops.reduce_sum, np.sum, np.float32, self.REAL_DATA, @@ -149,6 +150,11 @@ class ReduceOpsTest(xla_test.XLATestCase, parameterized.TestCase): self._testReduction(math_ops.reduce_mean, np.mean, np.float32, self.NONEMPTY_REAL_DATA, index_dtype) + def testReduceMeanF16(self, index_dtype): + if np.float16 in self.all_types: + self._testReduction(math_ops.reduce_mean, np.mean, np.float16, self.ONES, + index_dtype) + def testReduceMeanC64(self, index_dtype): self._testReduction(math_ops.reduce_mean, np.mean, np.complex64, self.NONEMPTY_COMPLEX_DATA, index_dtype) diff --git a/tensorflow/compiler/tests/rmsprop_test.py b/tensorflow/compiler/tests/rmsprop_test.py index 8840a1329a907bddc6ef1cb6dd1c2a6d234def5c..dc3e90b4afa41c08d899ee195d42fb91678bad1c 100644 --- a/tensorflow/compiler/tests/rmsprop_test.py +++ b/tensorflow/compiler/tests/rmsprop_test.py @@ -76,7 +76,7 @@ class RmspropTest(xla_test.XLATestCase): rms_opt = rmsprop.RMSPropOptimizer(learning_rate, centered=centered) rms_update = rms_opt.apply_gradients( zip([grads0, grads1], [var0, var1])) - variables.global_variables_initializer().run() + self.evaluate(variables.global_variables_initializer()) mg0 = rms_opt.get_slot(var0, "mg") self.assertEqual(mg0 is not None, centered) @@ -92,12 +92,12 @@ class RmspropTest(xla_test.XLATestCase): self.assertTrue(mom1 is not None) # Fetch params to validate initial values - self.assertAllClose([1.0, 2.0], var0.eval()) - self.assertAllClose([3.0, 4.0], var1.eval()) + self.assertAllClose([1.0, 2.0], self.evaluate(var0)) + self.assertAllClose([3.0, 4.0], self.evaluate(var1)) # Run 3 steps of RMSProp for _ in range(3): - rms_update.run() + self.evaluate(rms_update) var0_np, mg0_np, rms0_np, mom0_np = self._rmsprop_update_numpy( var0_np, @@ -118,14 +118,14 @@ class RmspropTest(xla_test.XLATestCase): # Validate updated params if centered: - self.assertAllCloseAccordingToType(mg0_np, mg0.eval()) - self.assertAllCloseAccordingToType(mg1_np, mg1.eval()) - self.assertAllCloseAccordingToType(rms0_np, rms0.eval()) - self.assertAllCloseAccordingToType(rms1_np, rms1.eval()) - self.assertAllCloseAccordingToType(mom0_np, mom0.eval()) - self.assertAllCloseAccordingToType(mom1_np, mom1.eval()) - self.assertAllCloseAccordingToType(var0_np, var0.eval()) - self.assertAllCloseAccordingToType(var1_np, var1.eval()) + self.assertAllCloseAccordingToType(mg0_np, self.evaluate(mg0)) + self.assertAllCloseAccordingToType(mg1_np, self.evaluate(mg1)) + self.assertAllCloseAccordingToType(rms0_np, self.evaluate(rms0)) + self.assertAllCloseAccordingToType(rms1_np, self.evaluate(rms1)) + self.assertAllCloseAccordingToType(mom0_np, self.evaluate(mom0)) + self.assertAllCloseAccordingToType(mom1_np, self.evaluate(mom1)) + self.assertAllCloseAccordingToType(var0_np, self.evaluate(var0)) + self.assertAllCloseAccordingToType(var1_np, self.evaluate(var1)) if __name__ == "__main__": diff --git a/tensorflow/compiler/tests/scan_ops_test.py b/tensorflow/compiler/tests/scan_ops_test.py index 897db384b7e8067b0460b5f344201f101a4d8479..17639bd8a755b9e9f5acc77979ac7a4149f112db 100644 --- a/tensorflow/compiler/tests/scan_ops_test.py +++ b/tensorflow/compiler/tests/scan_ops_test.py @@ -71,7 +71,7 @@ def handle_options(func, x, axis, exclusive, reverse): class CumsumTest(xla_test.XLATestCase): - valid_dtypes = [np.float32] + valid_dtypes = [np.float32, np.int32] def axis_dtypes(self): return set(self.int_types).intersection([np.int32, np.int64]) @@ -149,7 +149,7 @@ class CumsumTest(xla_test.XLATestCase): class CumprodTest(xla_test.XLATestCase): - valid_dtypes = [np.float32] + valid_dtypes = [np.float32, np.int32] def axis_dtypes(self): return set(self.int_types).intersection([np.int32, np.int64]) diff --git a/tensorflow/compiler/tests/stateless_random_ops_test.py b/tensorflow/compiler/tests/stateless_random_ops_test.py index b7747414ead7599d885b319b758976328aaf788b..ee7ca7e6f196e114ff18e2597145e5c198980b08 100644 --- a/tensorflow/compiler/tests/stateless_random_ops_test.py +++ b/tensorflow/compiler/tests/stateless_random_ops_test.py @@ -33,8 +33,11 @@ from tensorflow.python.platform import test class StatelessRandomOpsTest(xla_test.XLATestCase): """Test cases for stateless random-number generator operators.""" - def _random_types(self): - return self.float_types & {dtypes.float32, dtypes.float64} + def _random_types(self, include_int=False): + allowed_types = {dtypes.float32, dtypes.float64, dtypes.bfloat16} + if include_int: + allowed_types.update({dtypes.int32, dtypes.int64}) + return self.all_tf_types & allowed_types def testDeterminism(self): # Stateless values should be equal iff the seeds are equal (roughly) @@ -46,6 +49,11 @@ class StatelessRandomOpsTest(xla_test.XLATestCase): ]: for shape in (), (3,), (2, 5): for dtype in self._random_types(): + # Skip bfloat16. The result of bfloat16 is truncated from 32-bit + # result. With different seeds, the 32-bit results are different, + # but the truncated 16-bit results might be the same. + if dtype == dtypes.bfloat16: + continue pure = stateless_op(shape, seed=seed_t, dtype=dtype) values = [(seed, pure.eval(feed_dict={ seed_t: seed @@ -56,13 +64,16 @@ class StatelessRandomOpsTest(xla_test.XLATestCase): def testRandomUniformIsInRange(self): with self.cached_session() as sess, self.test_scope(): - for dtype in self._random_types(): + for dtype in self._random_types(include_int=True): + maxval = 1 + if dtype.is_integer: + maxval = 100 seed_t = array_ops.placeholder(dtypes.int32, shape=[2]) x = stateless.stateless_random_uniform( - shape=[1000], seed=seed_t, dtype=dtype) + shape=[1000], seed=seed_t, maxval=maxval, dtype=dtype) y = sess.run(x, {seed_t: [0x12345678, 0xabcdef12]}) self.assertTrue(np.all(y >= 0)) - self.assertTrue(np.all(y < 1)) + self.assertTrue(np.all(y < maxval)) def _chi_squared(self, x, bins): """Pearson's Chi-squared test.""" @@ -75,12 +86,18 @@ class StatelessRandomOpsTest(xla_test.XLATestCase): def testDistributionOfStatelessRandomUniform(self): """Use Pearson's Chi-squared test to test for uniformity.""" with self.cached_session() as sess, self.test_scope(): - for dtype in self._random_types(): + for dtype in self._random_types(include_int=True): seed_t = array_ops.placeholder(dtypes.int32, shape=[2]) n = 1000 + maxval = 1 + if dtype.is_integer: + maxval = 100 x = stateless.stateless_random_uniform( - shape=[n], seed=seed_t, dtype=dtype) + shape=[n], seed=seed_t, maxval=maxval, dtype=dtype) y = sess.run(x, {seed_t: [565656, 121212]}) + if maxval > 1: + # Normalize y to range [0, 1). + y = y.astype(float) / maxval # Tests that the values are distributed amongst 10 bins with equal # probability. 16.92 is the Chi^2 value for 9 degrees of freedom with # p=0.05. This test is probabilistic and would be flaky if the random @@ -121,7 +138,7 @@ class StatelessRandomOpsTest(xla_test.XLATestCase): # The constant 2.492 is the 5% critical value for the Anderson-Darling # test where the mean and variance are known. This test is probabilistic # so to avoid flakiness the seed is fixed. - self.assertTrue(self._anderson_darling(y) < 2.492) + self.assertTrue(self._anderson_darling(y.astype(float)) < 2.492) def testTruncatedNormalIsInRange(self): for dtype in self._random_types(): @@ -139,7 +156,7 @@ class StatelessRandomOpsTest(xla_test.XLATestCase): return math.exp(-(x**2) / 2.) / math.sqrt(2 * math.pi) def probit(x, sess=sess): - return sess.run(special_math.ndtri(x)) + return self.evaluate(special_math.ndtri(x)) a = -2. b = 2. @@ -157,6 +174,7 @@ class StatelessRandomOpsTest(xla_test.XLATestCase): # Burkardt, John. "The Truncated Normal Distribution". # Department of Scientific Computing website. Florida State University. expected_mean = mu + (normal_pdf(alpha) - normal_pdf(beta)) / z * sigma + y = y.astype(float) actual_mean = np.mean(y) self.assertAllClose(actual_mean, expected_mean, atol=5e-4) @@ -169,8 +187,8 @@ class StatelessRandomOpsTest(xla_test.XLATestCase): (alpha * normal_pdf(alpha) - beta * normal_pdf(beta)) / z) - ( (normal_pdf(alpha) - normal_pdf(beta)) / z)**2) actual_variance = np.var(y) - self.assertAllClose(actual_variance, expected_variance, rtol=1e-3) - + self.assertAllClose(actual_variance, expected_variance, + rtol=5e-3 if dtype == dtypes.bfloat16 else 1e-3) if __name__ == '__main__': diff --git a/tensorflow/compiler/tests/tensor_array_ops_test.py b/tensorflow/compiler/tests/tensor_array_ops_test.py index 46ca371c8abf1cb4710717a183ee12820c4c4ca0..d7e26d79c4c054860ade5c8960a3bca984e020b0 100644 --- a/tensorflow/compiler/tests/tensor_array_ops_test.py +++ b/tensorflow/compiler/tests/tensor_array_ops_test.py @@ -79,7 +79,8 @@ class TensorArrayTest(xla_test.XLATestCase): c0 = w2.stack() self.assertAllEqual( - convert([[[4.0, 5.0]], [[6.0, 7.0]], [[8.0, 9.0]]]), c0.eval()) + convert([[[4.0, 5.0]], [[6.0, 7.0]], [[8.0, 9.0]]]), + self.evaluate(c0)) def testTensorArrayWritePack(self): for dtype in self.numeric_tf_types: @@ -97,7 +98,7 @@ class TensorArrayTest(xla_test.XLATestCase): c0 = w2.stack() - self.assertAllEqual([3, 0, 1], c0.eval().shape) + self.assertAllEqual([3, 0, 1], self.evaluate(c0).shape) def _testTensorArrayWriteConcat(self, tf_dtype): with self.cached_session(), self.test_scope(): @@ -113,8 +114,8 @@ class TensorArrayTest(xla_test.XLATestCase): c0 = w2.concat() self.assertAllEqual( - convert([[4.0, 5.0], [104.0, 105.0], [6.0, 7.0], - [106.0, 107.0], [8.0, 9.0], [204.0, 205.0]]), c0.eval()) + convert([[4.0, 5.0], [104.0, 105.0], [6.0, 7.0], [106.0, 107.0], + [8.0, 9.0], [204.0, 205.0]]), self.evaluate(c0)) def testTensorArrayWriteConcat(self): for dtype in self.numeric_tf_types: @@ -341,7 +342,7 @@ class TensorArrayTest(xla_test.XLATestCase): r0_bad = gen_data_flow_ops.tensor_array_read_v3( handle=w0.handle, index=0, dtype=dtype2, flow_in=w0.flow) with self.assertRaisesOpError("TensorArray dtype is "): - r0_bad.eval() + self.evaluate(r0_bad) # Test reading from a different index than the one we wrote to w0.read(1) @@ -422,7 +423,7 @@ class TensorArrayTest(xla_test.XLATestCase): w2 = h2.write(0, 5.0) r2 = w2.read(0) r = r1 + r2 - self.assertAllClose(9.0, r.eval()) + self.assertAllClose(9.0, self.evaluate(r)) def _testTensorArrayGradientWriteReadType(self, dtype): with self.cached_session() as session, self.test_scope(): @@ -504,7 +505,7 @@ class TensorArrayTest(xla_test.XLATestCase): [-0.5, 1.5], # read(0) gradient [20.0, 30.0, 40.0, 50.0], # concat gradient ]) - grad_vals = sess.run(grad_r) # 2 + 2 entries + grad_vals = self.evaluate(grad_r) # 2 + 2 entries self.assertAllClose([2.0 - 0.5 + 20.0, 3.0 + 1.5 + 30.0], grad_vals[0]) self.assertAllEqual([4.0 + 40.0, 5.0 + 50.0], grad_vals[1]) @@ -526,7 +527,7 @@ class TensorArrayTest(xla_test.XLATestCase): with ops.control_dependencies([r0_readtwice]): r1_readtwice = w_readtwice.read(0) - self.assertAllEqual([1.0, -1.0], r1_readtwice.eval()) + self.assertAllEqual([1.0, -1.0], self.evaluate(r1_readtwice)) def _testTensorArrayGradientUnpackRead(self): with self.cached_session() as session, self.test_scope(): @@ -592,7 +593,7 @@ class TensorArrayTest(xla_test.XLATestCase): ta = tensor_array_ops.TensorArray( dtype=dtypes.float32, tensor_array_name="foo", size=3) s = ta.size() - self.assertAllEqual(3, s.eval()) + self.assertAllEqual(3, self.evaluate(s)) def testWriteCloseTensorArray(self): with self.cached_session(), self.test_scope(): @@ -722,7 +723,7 @@ class TensorArrayTest(xla_test.XLATestCase): # r = acc2.stack() # grad = gradients_impl.gradients(r, [x])[0] - # self.assertAllClose(31.0, grad.eval()) + # self.assertAllClose(31.0, self.evaluate(grad)) def testSumOfTwoReadVariablesWithoutRepeatGrad(self): with self.cached_session() as session, self.test_scope(): @@ -912,7 +913,7 @@ class TensorArrayTest(xla_test.XLATestCase): self.assertEqual(0, ta.size().eval()) ta = ta.unstack(array_ops.zeros([0, 3, 5])) packed = ta.stack() - self.assertAllEqual([0, 3, 5], packed.eval().shape) + self.assertAllEqual([0, 3, 5], self.evaluate(packed).shape) # Concatenating zero tensors along their first dimension gives a # first dimension of zero self.assertAllEqual([0, 5], ta.concat().eval().shape) @@ -1041,8 +1042,8 @@ class TensorArrayTest(xla_test.XLATestCase): (read0, read1, size0, size1)) # Tests that the control dependencies was added and executed. - self.assertEqual(1, v0.eval()) - self.assertEqual(1, v1.eval()) + self.assertEqual(1, self.evaluate(v0)) + self.assertEqual(1, self.evaluate(v1)) # Tests correct TensorArray. self.assertEqual(read0_v, 0) diff --git a/tensorflow/compiler/tests/tensor_list_ops_test.py b/tensorflow/compiler/tests/tensor_list_ops_test.py index 5c079d595c440cac644f5461154509abe7b1d1ed..1ecdb22cd0bc7e42d7ff67d20544fd26a65f6204 100644 --- a/tensorflow/compiler/tests/tensor_list_ops_test.py +++ b/tensorflow/compiler/tests/tensor_list_ops_test.py @@ -48,24 +48,39 @@ 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) + element_shape=(7, 15), num_elements=10, element_dtype=dtypes.float32) 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.tensor_list_reserve( + element_shape=(7, 15), num_elements=10, element_dtype=dtypes.float32) + # 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, + num_elements=20, element_dtype=dtypes.float32) 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)) @@ -75,21 +90,29 @@ 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))) if __name__ == "__main__": diff --git a/tensorflow/compiler/tests/unary_ops_test.py b/tensorflow/compiler/tests/unary_ops_test.py index 77f6eee0cf8ddc9b76f150e1038bf66da34c5218..3c2875ba477fa71e9e56a18d10efe0808533dd03 100644 --- a/tensorflow/compiler/tests/unary_ops_test.py +++ b/tensorflow/compiler/tests/unary_ops_test.py @@ -358,6 +358,11 @@ class UnaryOpsTest(xla_test.XLATestCase): np.array([[-0.05, 6.05, 5]], dtype=dtype), expected=np.array([[0, 6, 5]], dtype=dtype)) + self._assertOpOutputMatchesExpected( + nn_ops.leaky_relu, + np.array([[-2, -1, 0, 1, 2]], dtype=dtype), + expected=np.array([[-0.4, -0.2, 0.0, 1.0, 2.0]], dtype=dtype)) + self._assertOpOutputMatchesExpected( nn_ops.softmax, np.array([1, 2, 3, 4], dtype=dtype), @@ -476,6 +481,72 @@ class UnaryOpsTest(xla_test.XLATestCase): np.array([-1, -0.5, 0, 0.3], dtype=dtype), expected=np.array([-1., -0.5, 0., 0.296875], dtype=dtype)) + def quantize_and_dequantize_v2_round_half_up(x): + return array_ops.quantize_and_dequantize_v2( + x, + -1, + 1.0, + signed_input=True, + num_bits=8, + range_given=True, + round_mode="HALF_UP") + + self._assertOpOutputMatchesExpected( + quantize_and_dequantize_v2_round_half_up, + np.array([-0.8, -0.5, 0, 0.3, 0.8, -2, 33], dtype=dtype), + expected=np.array([ + -102.0 / 127, + -63.0 / 127, + 0, + 38.0 / 127, + 102.0 / 127, + -128.0 / 127, + 1, + ], + dtype=dtype)) + + def quantize_and_dequantize_v2_round_half_to_even(x): + return array_ops.quantize_and_dequantize_v2( + x, + -1.0, + 1.0, + signed_input=True, + num_bits=8, + range_given=True, + round_mode="HALF_TO_EVEN") + + self._assertOpOutputMatchesExpected( + quantize_and_dequantize_v2_round_half_to_even, + np.array( + [ + -0.8, + # The -0.5 should become -63.5 after scaling and with + # rounding this should become -64. But with the test + # unary_ops_test_cpu_ondemand, this fails as the result + # before scaling becomes -63.499996 and gets rounded to -63. + # TODO(sreenik): Some one more familiar with this test needs + # to take a look and resolve this. This works on all other + # variations of the platform like cpu, and gpu. + # -0.5, + 0, + 0.3, + 0.8, + -2, + 33 + ], + dtype=dtype), + expected=np.array( + [ + -102.0 / 127, + # -64.0 / 127, + 0, + 38.0 / 127, + 102.0 / 127, + -128.0 / 127, + 1, + ], + dtype=dtype)) + def quantize_and_dequantize_v3(x): return array_ops.quantize_and_dequantize_v3( x, -127, 127, num_bits=8, signed_input=True, range_given=False) @@ -576,7 +647,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), @@ -724,6 +795,15 @@ class UnaryOpsTest(xla_test.XLATestCase): lambda x: array_ops.bitcast(x, dtypes.int32), np.array([1e-45, 1.0], np.float32), expected=np.array([1, 0x3f800000], np.int32)) + if np.int64 in self.numeric_types: + self._assertOpOutputMatchesExpected( + lambda x: array_ops.bitcast(x, dtypes.int64), + np.array([1, 0x100000003f800000], np.uint64), + expected=np.array([1, 0x100000003f800000], np.int64)) + self._assertOpOutputMatchesExpected( + lambda x: array_ops.bitcast(x, dtypes.uint64), + np.array([1, 0x100000003f800000], np.int64), + expected=np.array([1, 0x100000003f800000], np.uint64)) def testInvertPermutation(self): self._assertOpOutputMatchesExpected( diff --git a/tensorflow/compiler/tests/variable_ops_test.py b/tensorflow/compiler/tests/variable_ops_test.py index dd2c252d383bca9c59033ac07e442b487e4975a6..fcd7ac5ba1ca5049246e93e6f5f76746fb28c6b8 100644 --- a/tensorflow/compiler/tests/variable_ops_test.py +++ b/tensorflow/compiler/tests/variable_ops_test.py @@ -40,6 +40,19 @@ from tensorflow.python.training.gradient_descent import GradientDescentOptimizer class VariableOpsTest(xla_test.XLATestCase): """Test cases for resource variable operators.""" + def testWriteEmptyShape(self): + # Verifies that we can pass an uninitialized variable with an empty shape, + # assign it a value, and successfully return it. + for dtype in self.numeric_types: + with self.test_session() as sess, self.test_scope(): + zeros = np.zeros([3, 0], dtype=dtype) + v = resource_variable_ops.ResourceVariable(zeros) + p = array_ops.placeholder(dtype) + x = v.assign(p) + with ops.control_dependencies([x]): + y = v.read_value() + self.assertAllClose(zeros, sess.run(y, {p: zeros})) + def testOneWriteOneOutput(self): # Regression test for a bug where computations with one non-constant # output and one variable update were mishandled. @@ -64,7 +77,7 @@ class VariableOpsTest(xla_test.XLATestCase): sess.run(variables.variables_initializer([v])) x = v.sparse_read(2) self.assertAllClose( - np.array([8j, 9, 10, 11]).astype(dtype), sess.run(x)) + np.array([8j, 9, 10, 11]).astype(dtype), self.evaluate(x)) def testSparseRead1DIndices(self): for dtype in self.numeric_types: @@ -76,7 +89,7 @@ class VariableOpsTest(xla_test.XLATestCase): x = v.sparse_read([2, 1]) self.assertAllClose( np.array([[8, 9, 10, 11], [4, 5, 6j, 7]]).astype(dtype), - sess.run(x)) + self.evaluate(x)) def testSparseRead2DIndices(self): for dtype in self.numeric_types: @@ -89,7 +102,7 @@ class VariableOpsTest(xla_test.XLATestCase): self.assertAllClose( np.array([[[8, 9, 10, 11], [4, 5, 6, 7]], [[0, 1, 2j, 3], [8, 9, 10, 11]]]).astype(dtype), - sess.run(x)) + self.evaluate(x)) def testSparseRead2DIndices3DTensor(self): for dtype in self.numeric_types: @@ -102,9 +115,9 @@ class VariableOpsTest(xla_test.XLATestCase): x = v.sparse_read([[2, 1], [3, 0]]) self.assertAllClose( np.array( - [[[[20, 21, 22], [23, 24j, 25]], [[10, 11, 12], [13, 14, 15]] - ], [[[30, 31, 32], [33, 34, 35]], [[0, 1, 2], [3, 4, 5]]] - ],).astype(dtype), sess.run(x)) + [[[[20, 21, 22], [23, 24j, 25]], [[10, 11, 12], [13, 14, 15]]], + [[[30, 31, 32], [33, 34, 35]], [[0, 1, 2], [3, 4, 5]]] + ],).astype(dtype), self.evaluate(x)) def testShape(self): for dtype in self.numeric_types: @@ -216,7 +229,7 @@ class VariableOpsTest(xla_test.XLATestCase): resource_variable_ops.resource_scatter_add( handle, [0], constant_op.constant([[2]], dtype=dtypes.int32))) read = resource_variable_ops.read_variable_op(handle, dtype=dtypes.int32) - self.assertAllEqual(sess.run(read), [[3], [7]]) + self.assertAllEqual(self.evaluate(read), [[3], [7]]) def testScatterSub(self): with self.test_session() as sess, self.test_scope(): @@ -229,7 +242,7 @@ class VariableOpsTest(xla_test.XLATestCase): resource_variable_ops.resource_scatter_sub( handle, [1], constant_op.constant([[2]], dtype=dtypes.int32))) read = resource_variable_ops.read_variable_op(handle, dtype=dtypes.int32) - self.assertAllEqual(sess.run(read), [[4], [-1]]) + self.assertAllEqual(self.evaluate(read), [[4], [-1]]) def testScatterMul(self): with self.test_session() as sess, self.test_scope(): @@ -242,7 +255,7 @@ class VariableOpsTest(xla_test.XLATestCase): resource_variable_ops.resource_scatter_mul( handle, [0], constant_op.constant([[5]], dtype=dtypes.int32))) read = resource_variable_ops.read_variable_op(handle, dtype=dtypes.int32) - self.assertEqual(sess.run(read), [[5]]) + self.assertEqual(self.evaluate(read), [[5]]) def testScatterDiv(self): with self.test_session() as sess, self.test_scope(): @@ -255,7 +268,7 @@ class VariableOpsTest(xla_test.XLATestCase): resource_variable_ops.resource_scatter_div( handle, [0], constant_op.constant([[3]], dtype=dtypes.int32))) read = resource_variable_ops.read_variable_op(handle, dtype=dtypes.int32) - self.assertAllEqual(sess.run(read), [[2]]) + self.assertAllEqual(self.evaluate(read), [[2]]) def testScatterMin(self): with self.test_session() as sess, self.test_scope(): @@ -268,7 +281,7 @@ class VariableOpsTest(xla_test.XLATestCase): resource_variable_ops.resource_scatter_min( handle, [0], constant_op.constant([[3]], dtype=dtypes.int32))) read = resource_variable_ops.read_variable_op(handle, dtype=dtypes.int32) - self.assertEqual(sess.run(read), [[3]]) + self.assertEqual(self.evaluate(read), [[3]]) def testScatterMax(self): with self.test_session() as sess, self.test_scope(): @@ -281,7 +294,7 @@ class VariableOpsTest(xla_test.XLATestCase): resource_variable_ops.resource_scatter_max( handle, [0], constant_op.constant([[3]], dtype=dtypes.int32))) read = resource_variable_ops.read_variable_op(handle, dtype=dtypes.int32) - self.assertEqual(sess.run(read), [[6]]) + self.assertEqual(self.evaluate(read), [[6]]) def testScatterUpdate(self): with self.test_session() as sess, self.test_scope(): @@ -294,7 +307,7 @@ class VariableOpsTest(xla_test.XLATestCase): resource_variable_ops.resource_scatter_update( handle, [0], constant_op.constant([[3]], dtype=dtypes.int32))) read = resource_variable_ops.read_variable_op(handle, dtype=dtypes.int32) - self.assertEqual(sess.run(read), [[3]]) + self.assertEqual(self.evaluate(read), [[3]]) def testScatterAddScalar(self): with self.test_session() as sess, self.test_scope(): @@ -307,7 +320,7 @@ class VariableOpsTest(xla_test.XLATestCase): resource_variable_ops.resource_scatter_add( handle, [0], constant_op.constant(2, dtype=dtypes.int32))) read = resource_variable_ops.read_variable_op(handle, dtype=dtypes.int32) - self.assertEqual(sess.run(read), [[3]]) + self.assertEqual(self.evaluate(read), [[3]]) def testScatterSubScalar(self): with self.test_session() as sess, self.test_scope(): @@ -320,7 +333,7 @@ class VariableOpsTest(xla_test.XLATestCase): resource_variable_ops.resource_scatter_sub( handle, [0], constant_op.constant(2, dtype=dtypes.int32))) read = resource_variable_ops.read_variable_op(handle, dtype=dtypes.int32) - self.assertEqual(sess.run(read), [[-1]]) + self.assertEqual(self.evaluate(read), [[-1]]) def testScatterMulScalar(self): with self.test_session() as sess, self.test_scope(): @@ -333,7 +346,7 @@ class VariableOpsTest(xla_test.XLATestCase): resource_variable_ops.resource_scatter_mul( handle, [0], constant_op.constant(5, dtype=dtypes.int32))) read = resource_variable_ops.read_variable_op(handle, dtype=dtypes.int32) - self.assertEqual(sess.run(read), [[5]]) + self.assertEqual(self.evaluate(read), [[5]]) def testScatterDivScalar(self): with self.test_session() as sess, self.test_scope(): @@ -346,7 +359,7 @@ class VariableOpsTest(xla_test.XLATestCase): resource_variable_ops.resource_scatter_div( handle, [0], constant_op.constant(3, dtype=dtypes.int32))) read = resource_variable_ops.read_variable_op(handle, dtype=dtypes.int32) - self.assertEqual(sess.run(read), [[2]]) + self.assertEqual(self.evaluate(read), [[2]]) def testScatterMinScalar(self): with self.test_session() as sess, self.test_scope(): @@ -359,7 +372,7 @@ class VariableOpsTest(xla_test.XLATestCase): resource_variable_ops.resource_scatter_min( handle, [0], constant_op.constant(3, dtype=dtypes.int32))) read = resource_variable_ops.read_variable_op(handle, dtype=dtypes.int32) - self.assertEqual(sess.run(read), [[3]]) + self.assertEqual(self.evaluate(read), [[3]]) def testScatterMaxScalar(self): with self.test_session() as sess, self.test_scope(): @@ -372,7 +385,7 @@ class VariableOpsTest(xla_test.XLATestCase): resource_variable_ops.resource_scatter_max( handle, [0], constant_op.constant(3, dtype=dtypes.int32))) read = resource_variable_ops.read_variable_op(handle, dtype=dtypes.int32) - self.assertEqual(sess.run(read), [[6]]) + self.assertEqual(self.evaluate(read), [[6]]) def testScatterNdAddOps(self): with self.test_session() as sess, self.test_scope(): @@ -387,7 +400,7 @@ class VariableOpsTest(xla_test.XLATestCase): sess.run(gen_state_ops.resource_scatter_nd_add(handle, indices, updates)) read = resource_variable_ops.read_variable_op( handle, dtype=dtypes.float32) - self.assertAllClose(expected, sess.run(read)) + self.assertAllClose(expected, self.evaluate(read)) def testScatterNdUpdateAddOps(self): with self.test_session() as sess, self.test_scope(): @@ -403,7 +416,7 @@ class VariableOpsTest(xla_test.XLATestCase): gen_state_ops.resource_scatter_nd_update(handle, indices, updates)) read = resource_variable_ops.read_variable_op( handle, dtype=dtypes.float32) - self.assertAllClose(expected, sess.run(read)) + self.assertAllClose(expected, self.evaluate(read)) class StridedSliceAssignChecker(object): diff --git a/tensorflow/compiler/tests/xla_device_test.py b/tensorflow/compiler/tests/xla_device_test.py index 28d61fb07dcb665fa0dbe3f3e566e291e24fa662..ef55292b1be91a731ec556d7efa9cdf1a696e5cc 100644 --- a/tensorflow/compiler/tests/xla_device_test.py +++ b/tensorflow/compiler/tests/xla_device_test.py @@ -81,7 +81,7 @@ class XlaDeviceTest(xla_test.XLATestCase): with self.cached_session() as sess: with self.test_scope(): x = gen_control_flow_ops.control_trigger() - sess.run(x) + self.evaluate(x) if __name__ == "__main__": diff --git a/tensorflow/compiler/tests/xla_ops_test.py b/tensorflow/compiler/tests/xla_ops_test.py index 4cf88fc523735cc2d22e085afb83790c7ebb48e4..28274ff799de2c85e1e80512cadbe0206cb640a4 100644 --- a/tensorflow/compiler/tests/xla_ops_test.py +++ b/tensorflow/compiler/tests/xla_ops_test.py @@ -319,7 +319,7 @@ class XlaOpsTest(xla_test.XLATestCase, parameterized.TestCase): session.run(output) self.assertRegexpMatches( invalid_arg_error.exception.message, - (r'^start_indices must be a vector with length equal to input rank, ' + (r'start_indices must be a vector with length equal to input rank, ' r'but input rank is 3 and start_indices has shape \[2\].*')) def testDynamicSliceWithIncorrectSizeIndicesShape(self): @@ -332,7 +332,7 @@ class XlaOpsTest(xla_test.XLATestCase, parameterized.TestCase): session.run(output) self.assertRegexpMatches( invalid_arg_error.exception.message, - (r'^size_indices must be a vector with length equal to input rank, ' + (r'size_indices must be a vector with length equal to input rank, ' r'but input rank is 3 and size_indices has shape \[2\].*')) diff --git a/tensorflow/compiler/tf2tensorrt/BUILD b/tensorflow/compiler/tf2tensorrt/BUILD new file mode 100644 index 0000000000000000000000000000000000000000..a67e511826ae161e78d504c1513934065cbfd19f --- /dev/null +++ b/tensorflow/compiler/tf2tensorrt/BUILD @@ -0,0 +1,440 @@ +# 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", + ], + 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", + "@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(), + # TODO(laigd): fix this by merging header file in cc file. + alwayslink = 1, # buildozer: disable=alwayslink-with-hdrs +) + +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 = [ + ":get_serialized_resource_op_op_lib", + ":trt_op_kernels", + ":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", + ], +) + +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 = [ + ":get_serialized_resource_op_op_lib", + ":trt_engine_op_op_lib", + ":trt_logging", + ], +) + +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_engine_op_op_lib", + ":get_serialized_resource_op_op_lib", + ], + srcs_version = "PY2AND3", + deps = [ + "//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", + "//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//: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", + ], + 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 78% rename from tensorflow/contrib/tensorrt/convert/convert_graph.cc rename to tensorflow/compiler/tf2tensorrt/convert/convert_graph.cc index 070dc294a784b9978b8e2e1416e892b7a04c5e2d..1fdb099cc1d658b4259177e357b639ea72d636d0 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() { @@ -81,59 +82,96 @@ std::vector GetLoadedTensorRTVersion() { return {ver_major, ver_minor, ver_patch}; } -namespace { +TrtCandidateSelector::TrtCandidateSelector( + const grappler::GraphProperties& graph_properties, int precision_mode) + : graph_properties_(graph_properties), precision_mode_(precision_mode) {} -bool IsTensorRTCandidate(const tensorflow::Node* node) { +Status TrtCandidateSelector::IsTensorRTCandidate(const tensorflow::Node* node) { + // TODO(laigd): move this set to TrtNodeValidator where it should belong. // LINT.IfChange - // TODO(jie): Segmentation shouldn't associated with op name. - // Split it into a registration for each kernel. static const std::set candidate_ops = { - "Identity", - "Snapshot", - "Const", - "Conv2D", - "MaxPool", - "BiasAdd", - "Relu", - "Add", - "Mul", - "Sub", - "Rsqrt", - "Pad", - "Mean", - "AvgPool", - "ConcatV2", - "DepthwiseConv2dNative", - "FusedBatchNorm", - "FusedBatchNormV2", - "Div", - "RealDiv", - "Rsqrt", - "Reciprocal", - "Exp", - "Log", - "Sqrt", - "Abs", - "Neg", - "Transpose", - "Reshape", - "MatMul", - "BatchMatMul", - "Softmax", - "Minimum", - "Maximum", - "TopKV2", - "Sum", - "Prod", - "Max", - "Min", - // TODO(ben,jie): ... + "Abs", + "Add", + "AvgPool", + "BatchMatMul", + "BiasAdd", + "ConcatV2", + "Const", + "Conv2D", + "DepthwiseConv2dNative", + "Div", + "Exp", + "ExpandDims", + "FusedBatchNorm", + "FusedBatchNormV2", + "Identity", + "Log", + "MatMul", + "Max", + "MaxPool", + "Maximum", + "Mean", + "Min", + "Minimum", + "Mul", + "Neg", + "Pad", + "Prod", + "RealDiv", + "Reciprocal", + "Relu", + "Relu6", + "Reshape", + "Rsqrt", + "Rsqrt", + "Sigmoid", + "Snapshot", + "Softmax", + "Sqrt", + "Square", + "Squeeze", + "StridedSlice", + "Sub", + "Sum", + "Tanh", + "TopKV2", + "Transpose", + }; + bool is_supported_op_type = + (candidate_ops.count(node->type_string()) || + PluginFactoryTensorRT::GetInstance()->IsPlugin(node->type_string())); + static const std::set quantize_ops = { + "QuantizeAndDequantizeV2", + "QuantizeAndDequantizeV3", + "FakeQuantWithMinMaxVars", + "FakeQuantWithMinMaxArgs", }; - // LINT.ThenChange(//tensorflow/contrib/tensorrt/convert/convert_nodes.cc) - return (candidate_ops.count(node->type_string()) || - PluginFactoryTensorRT::GetInstance()->IsPlugin(node->type_string())); + // 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())) { + is_supported_op_type = true; + } + // LINT.ThenChange(//tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc) + if (!is_supported_op_type) { + return errors::Unimplemented("Op type ", node->type_string(), + " is not supported"); + } + + std::vector input_edges; + TF_RETURN_IF_ERROR(node->input_edges(&input_edges)); + std::vector> input_node_and_ports; + input_node_and_ports.reserve(input_edges.size()); + for (const Edge* input_edge : input_edges) { + input_node_and_ports.emplace_back(&input_edge->src()->def(), + input_edge->src_output()); + } + return validator_.ValidateNode(node->def(), input_node_and_ports, + graph_properties_); } +namespace { + tensorflow::Status BuildNodeMap( const tensorflow::Graph& graph, std::unordered_map* node_map) { @@ -152,7 +190,7 @@ tensorflow::Status BuildNodeMap( tensorflow::Status ConvertCalibGraphToInferGraph( const tensorflow::GraphDef& graph_def, tensorflow::GraphDef* infer_graph, bool is_dyn_op) { - VLOG(0) << "Starting Calib Conversion"; + LOG(INFO) << "Starting Calib Conversion"; infer_graph->CopyFrom(graph_def); auto trt_rm = TRTResourceManager::instance(); auto calib_rm = trt_rm->getManager("TRTCalibration"); @@ -202,18 +240,19 @@ tensorflow::Status ConvertGraphDefToTensorRT( 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, - int max_cached_engines, std::vector cached_engine_batches) { + int max_cached_engines, std::vector cached_engine_batches, + bool use_calibration) { // Create GrapplerItem. tensorflow::grappler::GrapplerItem item; item.fetch = output_names; item.graph = graph_def; - // TODO(aaroey): we should have used single machine cluster like the - // following, but the problem is then wrap_conversion will depend on - // direct_session and cause double linking problems. To fix this we need to - // fix or get rid of the swig dependency. Here we use VirtualCluster - // as a work around, and we need to create a session to initialize the - // underlying device before calling this method. +// TODO(aaroey): we should have used single machine cluster like the +// following, but the problem is then wrap_conversion will depend on +// direct_session and cause double linking problems. To fix this we need to +// fix or get rid of the swig dependency. Here we use VirtualCluster +// as a work around, and we need to create a session to initialize the +// underlying device before calling this method. #if 0 // Create single machine cluster. Note that this will create a session and // initialize the gpu devices. @@ -246,7 +285,9 @@ tensorflow::Status ConvertGraphDefToTensorRT( #endif // Create RewriterConfig. - tensorflow::RewriterConfig rw_cfg; + tensorflow::ConfigProto config_proto; + auto& rw_cfg = + *config_proto.mutable_graph_options()->mutable_rewrite_options(); // TODO(aaroey): use only const folding and layout for the time being since // new optimizers break the graph for trt. rw_cfg.add_optimizers("constfold"); @@ -267,9 +308,10 @@ tensorflow::Status ConvertGraphDefToTensorRT( list->add_i(batch); } } + parameters["use_calibration"].set_b(use_calibration); // Run optimizer. - tensorflow::grappler::MetaOptimizer meta_opt(nullptr, rw_cfg); + tensorflow::grappler::MetaOptimizer meta_opt(nullptr, config_proto); TF_RETURN_IF_ERROR(meta_opt.Optimize(cluster.get(), item, new_graph_def)); if (VLOG_IS_ON(5)) { @@ -282,17 +324,23 @@ tensorflow::Status ConvertGraphDefToTensorRT( return Status::OK(); } +struct EdgePtrCompare { + bool operator()(const tensorflow::Edge* lhs, + const tensorflow::Edge* rhs) const { + return lhs->id() < rhs->id(); + } +}; + // Function to get subsegment information structure. tensorflow::Status GetEngineInfo( const tensorflow::Graph* g, const tensorflow::grappler::GraphProperties& graph_properties, - const std::set& segment_nodes, + const std::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 @@ -304,26 +352,45 @@ 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); - // Create input connections. - for (const auto edge : node->in_edges()) { + 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(), + node->in_edges().end()); + 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()) { @@ -340,12 +407,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. @@ -369,10 +435,14 @@ tensorflow::Status GetEngineInfo( node_id, edge->dst_input(), /*input_edge=*/true, port); } } - // Create output connections. - for (const auto edge : node->out_edges()) { + // Create output connections. Sort edges first to make determnistic since + // out_edges is a set of pointers. + std::vector out_edges(node->out_edges().begin(), + node->out_edges().end()); + std::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()) { @@ -400,12 +470,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(); @@ -415,7 +484,8 @@ tensorflow::Status GetEngineInfo( << "but this shouldn't have happened"; info->device = *segment_devices.begin(); } else { - LOG(ERROR) << "Can't find a device placement for the op!"; + VLOG(1) << "No device is assigned to the segment. " + << "A device will be assigned during graph execution (inference)."; } return Status::OK(); } @@ -525,6 +595,18 @@ tensorflow::Status CreateTRTNode(const std::vector& infos, int pos, } input_shape_protos.at(conn.port_number) = in_shape; input_shapes.at(conn.port_number) = conn.outside_shape; + // Shape must be fully defined (excluding batch dimension) for static + // mode. + if (info.engine_type == EngineInfo::EngineType::TRTStatic) { + for (int i = 1; i < conn.outside_shape.dims(); i++) { + if (conn.outside_shape.dim_size(i) <= 0) { + return tensorflow::errors::Internal( + "Input shapes must be fully defined when in static mode. " + "Please try is_dynamic_op=True (shape was ", + conn.outside_shape.DebugString(), ")"); + } + } + } // Rewrire data input if it's not found in original graph. tensorflow::Node* input_node = graph->FindNodeId(conn.outside_id); @@ -546,27 +628,38 @@ tensorflow::Status CreateTRTNode(const std::vector& infos, int pos, } } } + // We don't support segments with no inputs. Fall back to native TF here to + // avoid crash later. Constant folding should've folded the ops that make up + // these segments. + if (inputs.empty()) { + return tensorflow::errors::Internal( + "Segment has no inputs (possible " + "constfold failure)"); + } + + const bool calibrate_int8 = + (info.precision_mode == INT8MODE && info.use_calibration); + // Build the engine and get its serialized representation. string segment_string; - if (info.engine_type == EngineInfo::EngineType::TRTStatic || - info.precision_mode == INT8MODE) { + if (info.engine_type == EngineInfo::EngineType::TRTStatic || calibrate_int8) { // Create static engine for fp32/fp16 mode, and test validity of the engine - // for int8 mode. We don't want engine to fail at the calibration time. - // So we are constructing a FP32 engine here to check its validity, and if - // it is a valid engine then we put the serialized graphdef to the op. - // Otherwise we skip node creation for this engine. + // for int8 calibration mode. We don't want engine to fail at the + // calibration time. So we are constructing a FP32 engine here to check its + // validity, and if it is a valid engine then we put the serialized graphdef + // to the op. Otherwise we skip node creation for this engine. Logger trt_logger; TrtUniquePtrType engine; // TODO(sami): What happens if 1st dim is not batch? TF_RETURN_IF_ERROR(ConvertGraphDefToEngine( - info.segment_graph_def, - info.precision_mode == INT8MODE ? FP32MODE : info.precision_mode, + info.segment_graph_def, calibrate_int8 ? FP32MODE : 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()); - if (info.precision_mode == INT8MODE) { + if (calibrate_int8) { // See above comment about why not putting this inside the 'else' branch. segment_string = info.segment_graph_def.SerializeAsString(); } @@ -574,14 +667,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 && - !TRTResourceManager::instance()->getManager("TRTCalibration")) { - LOG(ERROR) << "Failed to construct calibration storage"; - } tensorflow::NodeDefBuilder node_builder(info.engine_name, "TRTEngineOp"); if (!info.device.empty()) node_builder.Device(info.device); if (VLOG_IS_ON(1)) { @@ -597,7 +684,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; @@ -611,9 +698,9 @@ 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) .Attr("OutT", out_types) .Finalize(&trt_node); if (!status.ok()) { @@ -846,13 +933,18 @@ tensorflow::Status ConvertAfterShapes(ConversionParams& params) { } segment_options.minimum_segment_size = params.minimum_segment_size; tensorflow::tensorrt::segment::SegmentNodesVector initial_segments; + TrtCandidateSelector candidate_selector(*params.graph_properties, + params.precision_mode); TF_RETURN_IF_ERROR(tensorrt::segment::SegmentGraph( - &graph, IsTensorRTCandidate, InputEdgeValidator(*params.graph_properties), - OutputEdgeValidator(), segment_options, &initial_segments)); - if (initial_segments.size() > 1) { - VLOG(0) << "MULTIPLE tensorrt candidate conversion: " + &graph, + std::bind(&TrtCandidateSelector::IsTensorRTCandidate, &candidate_selector, + std::placeholders::_1), + // Input validation is already done by TrtCandidateSelector, so we don't + // need to check the input edges. + [](const Edge* edge) { return true; }, OutputEdgeValidator(), + segment_options, &initial_segments)); + LOG(INFO) << "Number of TensorRT candidate segments: " << initial_segments.size(); - } // Get the EngineInfo for each segment. std::unordered_map node_map; @@ -878,13 +970,17 @@ tensorflow::Status ConvertAfterShapes(ConversionParams& params) { continue; } curr_engine.precision_mode = params.precision_mode; - curr_engine.engine_type = - (params.is_dyn_op || params.precision_mode == INT8MODE - ? EngineInfo::EngineType::TRTDynamic - : EngineInfo::EngineType::TRTStatic); + if (params.use_calibration && params.precision_mode != INT8MODE) { + return errors::InvalidArgument( + "Calibration with FP32 or FP16 is not supported."); + } + curr_engine.engine_type = ((params.is_dyn_op || params.use_calibration) + ? EngineInfo::EngineType::TRTDynamic + : EngineInfo::EngineType::TRTStatic); + curr_engine.use_calibration = params.use_calibration; curr_engine.cached_engine_batches = params.cached_engine_batches; curr_engine.maximum_cached_engines = params.max_cached_engines; - StrAppend(&curr_engine.engine_name, "my_trt_op_", t); + StrAppend(&curr_engine.engine_name, "TRTEngineOp_", t); status = RegisterSegmentFunctionToFunctionLibrary( &graph, curr_engine.segment_graph_def, curr_engine.engine_name); if (!status.ok()) { @@ -900,7 +996,7 @@ tensorflow::Status ConvertAfterShapes(ConversionParams& params) { converted_segments.push_back(std::move(curr_segment)); if (VLOG_IS_ON(8)) { - string fname = curr_engine.engine_name; + string fname = engine_segments.back().engine_name; StrAppend(&fname, ".pb"); std::fstream f; f.open(fname.c_str(), std::fstream::out | std::fstream::binary); @@ -943,19 +1039,30 @@ 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. - const string msg = StrCat("Engine ", engine.engine_name, - " creation for segment ", i, ", composed of ", - converted_segments.at(i).first.size(), " nodes"); + + 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 << ". Skipping..."; + LOG(WARNING) << msg << " failed: " << status << ". Fallback to TF..."; + } + if (VLOG_IS_ON(1)) { + msg = "Segment consists of nodes: "; + 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); diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_graph.h b/tensorflow/compiler/tf2tensorrt/convert/convert_graph.h new file mode 100644 index 0000000000000000000000000000000000000000..fb82a430c632781047487a280e23e7da4c385929 --- /dev/null +++ b/tensorflow/compiler/tf2tensorrt/convert/convert_graph.h @@ -0,0 +1,126 @@ +/* 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_CONVERT_CONVERT_GRAPH_H_ +#define TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_CONVERT_GRAPH_H_ + +#include + +#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" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/platform/types.h" + +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT + +namespace tensorflow { +namespace tensorrt { +namespace convert { + +// Helper class for the segmenter to determine whether given TF node is +// supported by TRT. +class TrtCandidateSelector { + public: + TrtCandidateSelector(const grappler::GraphProperties& graph_properties, + int 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. + Status IsTensorRTCandidate(const tensorflow::Node* node); + + private: + // The TF-TRT node converter used to verify whether individual node is + // supported. It will operate in validation-only mode. + TrtNodeValidator validator_; + + // GraphProperties of the graph whose nodes are to be validated by + // IsTensorRTCandidate(). + const grappler::GraphProperties& graph_properties_; + + // Quantization ops are only converted when using quantized precisions. + const int precision_mode_; +}; + +struct ConversionParams { + ConversionParams() + : input_graph_def(nullptr), + max_batch_size(1), + max_workspace_size_bytes(1 << 30), + output_graph_def(nullptr), + precision_mode(1), + minimum_segment_size(3), + graph_properties(nullptr), + cluster(nullptr), + is_dyn_op(false), + fixed_input_size(true), + use_calibration(true), + max_cached_engines(1) {} + const tensorflow::GraphDef* input_graph_def; + const std::vector* output_names; + size_t max_batch_size; + size_t max_workspace_size_bytes; + tensorflow::GraphDef* output_graph_def; + int precision_mode; + int minimum_segment_size; + const tensorflow::grappler::GraphProperties* graph_properties; + const tensorflow::grappler::Cluster* cluster; + bool is_dyn_op; // Whether to create engine on conversion or execution time + bool fixed_input_size; // Assume non-batch ranks of input tensors are fixed + int max_cached_engines; // maximum number of cached engines + bool use_calibration; + std::vector cached_engine_batches; // list of cached engines +}; + +// This method extracts calibration information from the resource managers +// and puts them in to engine nodedefs. +tensorflow::Status ConvertCalibGraphToInferGraph( + const tensorflow::GraphDef& graph_def, tensorflow::GraphDef* new_graph_def, + bool is_dyn_op); + +// - max_batch_size: maximum batch size which can be used for inference for +// optimization targets inference run with max batch size. +// - max_workspace_size_bytes: The upper bound of memory allowance for engine +// building. +tensorflow::Status ConvertGraphDefToTensorRT( + const tensorflow::GraphDef& graph_def, + const std::vector& output_names, size_t max_batch_size, + size_t max_workspace_size_bytes, tensorflow::GraphDef* new_graph_def, + 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); + +// Method to call from optimization pass +tensorflow::Status ConvertAfterShapes(ConversionParams& params); + +// Return compile time TensorRT library version information. +std::vector GetLinkedTensorRTVersion(); + +// Return runtime time TensorRT library version information. +std::vector GetLoadedTensorRTVersion(); + +// Helper method for the conversion, expose for testing. +std::pair GetDeviceAndAllocator( + const ConversionParams& params, const EngineInfo& engine); + +} // namespace convert +} // namespace tensorrt +} // namespace tensorflow + +#endif // GOOGLE_TENSORRT +#endif // GOOGLE_CUDA + +#endif // TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_CONVERT_GRAPH_H_ diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_graph_test.cc b/tensorflow/compiler/tf2tensorrt/convert/convert_graph_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..a3c3a8ac6561259c974aebb6c6eeac05c71b7161 --- /dev/null +++ b/tensorflow/compiler/tf2tensorrt/convert/convert_graph_test.cc @@ -0,0 +1,229 @@ +/* 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/convert/convert_graph.h" + +#include +#include +#include "tensorflow/cc/framework/scope.h" +#include "tensorflow/cc/ops/standard_ops.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" +#include "tensorflow/core/grappler/clusters/cluster.h" +#include "tensorflow/core/grappler/grappler_item.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/lib/core/status_test_util.h" +#include "tensorflow/core/platform/test.h" +#include "tensorflow/core/protobuf/config.pb.h" // NOLINT +#include "tensorflow/core/public/session.h" + +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT + +namespace tensorflow { +namespace tensorrt { +namespace convert { + +// TODO(laigd): put this into some test utils file. +void ExpectStatus(Status status, error::Code code = error::OK, + const char* substr = nullptr) { + EXPECT_EQ(code, status.code()) + << status << " vs expected error code \"" << error::Code_Name(code) + << "\" and message \"" << substr << "\""; + if (substr) { + EXPECT_THAT(status.error_message(), ::testing::HasSubstr(substr)) << status; + } +} + +TEST(TrtCandidateSelector, Basics) { + // Create a graph containing both TRT-compatible and TRT-incompatible nodes + // and use it to test TrtCandidateSelector::IsTensorRTCandidate(). + const std::vector input_shape_array{2, 2}; + TensorShape input_shape; + TF_EXPECT_OK(TensorShapeUtils::MakeShape(input_shape_array, &input_shape)); + + Scope s = Scope::NewRootScope(); + ops::Placeholder::Attrs feed_attrs; + TF_EXPECT_OK( + TensorShapeUtils::MakeShape(input_shape_array, &feed_attrs.shape_)); + + // Compatible input. + auto feed = ops::Placeholder(s.WithOpName("feed"), DT_FLOAT, feed_attrs); + auto const_1 = ops::Const(s.WithOpName("const_1"), 1.0f, input_shape); + + // Compatible MatMul. + auto matmul = ops::MatMul(s.WithOpName("matmul"), feed, const_1); + + // Incompatible MatMul. + ops::MatMul::Attrs matmul_attrs; + matmul_attrs.transpose_a_ = true; + auto incompatible_matmul = ops::MatMul(s.WithOpName("incompatible_matmul"), + feed, const_1, matmul_attrs); + + // Unsupported op. + auto unsupported_op = ops::Sin(s.WithOpName("sin"), feed); + + // Incompatible input. + auto incompatible_feed = ops::Placeholder(s.WithOpName("feed"), DT_DOUBLE); + auto const_2 = ops::Const(s.WithOpName("const_2"), 1.0, input_shape); + // Compatible op with incompatible input. + auto matmul_with_incompatible_input = + ops::MatMul(s.WithOpName("matmul_with_incompatible_input"), + incompatible_feed, const_2); + + // Quantize ops. + auto quantize_attrs = ops::FakeQuantWithMinMaxArgs::Min(-6.0f).Max(6.0f); + auto quantize = ops::FakeQuantWithMinMaxArgs(s.WithOpName("quantize"), feed, + quantize_attrs); + + // Get GrapplerItem and GraphProperties. + grappler::GrapplerItem item; + TF_EXPECT_OK(s.ToGraphDef(&item.graph)); + Tensor feed_tensor(DT_FLOAT, input_shape); + item.feed.push_back(std::make_pair("feed", feed_tensor)); + grappler::GraphProperties graph_properties(item); + TF_EXPECT_OK(graph_properties.InferStatically(true)); + + for (const int precision_mode : {FP32MODE, INT8MODE}) { + TrtCandidateSelector selector(graph_properties, precision_mode); + TF_EXPECT_OK(selector.IsTensorRTCandidate(matmul.operation.node())); + ExpectStatus( + selector.IsTensorRTCandidate(incompatible_matmul.operation.node()), + error::INVALID_ARGUMENT, + "transpose_a is not supported for TensorRT FullyConnected " + "(op: MatMul), at: incompatible_matmul"); + ExpectStatus(selector.IsTensorRTCandidate(unsupported_op.operation.node()), + error::UNIMPLEMENTED, "Op type Sin is not supported"); + ExpectStatus( + selector.IsTensorRTCandidate( + matmul_with_incompatible_input.operation.node()), + error::INTERNAL, + "Failed to convert input with index 0 to a TRT_TensorOrWeights"); + if (precision_mode == INT8MODE) { + TF_EXPECT_OK(selector.IsTensorRTCandidate(quantize.operation.node())); + } else { + ExpectStatus(selector.IsTensorRTCandidate(quantize.operation.node()), + error::UNIMPLEMENTED, + "Op type FakeQuantWithMinMaxArgs is not supported"); + } + } +} + +class FakeCluster : public grappler::Cluster { + public: + FakeCluster() : Cluster(0) {} + + void SetDeviceSet(const DeviceSet* device_set) { device_set_ = device_set; } + + const DeviceSet* GetDeviceSet() const override { return device_set_; } + + string type() const override { return ""; } + Status Provision() override { return Status::OK(); } + Status Initialize(const grappler::GrapplerItem& item) override { + return Status::OK(); + } + Status Run(const GraphDef& graph_def, + const std::vector>& feed, + const std::vector& fetch, RunMetadata* metadata) override { + return Status::OK(); + } + + private: + const DeviceSet* device_set_; +}; + +TEST(ConvertGraphTest, GetDeviceAndAllocator) { + ConversionParams params; + EngineInfo engine_info; + { + // params.cluster is not set, and no gpu device is available. + auto result = GetDeviceAndAllocator(params, engine_info); + EXPECT_EQ(-1, result.first); + EXPECT_EQ(nullptr, result.second); + } + + // Create a session with two (virtual) gpu device. + SessionOptions options; + ConfigProto* config = &options.config; + GPUOptions* gpu_options = config->mutable_gpu_options(); + auto virtual_devices = + gpu_options->mutable_experimental()->add_virtual_devices(); + virtual_devices->add_memory_limit_mb(200); + virtual_devices->add_memory_limit_mb(200); + std::unique_ptr session(NewSession(options)); + + { + // params.cluster is not set, should find and return first gpu id and + // corresponding allocator. + auto result = GetDeviceAndAllocator(params, engine_info); + EXPECT_EQ(0, result.first); + EXPECT_NE(nullptr, result.second); + EXPECT_EQ("GPU_0_bfc", result.second->Name()); + } + + FakeCluster cluster; + params.cluster = &cluster; + { + // params.cluster->GetDeviceSet() returns null, should find and return first + // gpu id and corresponding allocator. + auto result = GetDeviceAndAllocator(params, engine_info); + EXPECT_EQ(0, result.first); + EXPECT_NE(nullptr, result.second); + EXPECT_EQ("GPU_0_bfc", result.second->Name()); + } + + // Build the DeviceSet. + DeviceSet device_set; + const DeviceMgr* device_mgr = nullptr; + TF_ASSERT_OK(session->LocalDeviceManager(&device_mgr)); + for (auto d : device_mgr->ListDevices()) { + device_set.AddDevice(d); + } + cluster.SetDeviceSet(&device_set); + { + // engine_info.device is not set, should find and return first gpu id and + // corresponding allocator. + auto result = GetDeviceAndAllocator(params, engine_info); + EXPECT_EQ(0, result.first); + EXPECT_NE(nullptr, result.second); + EXPECT_EQ("GPU_0_bfc", result.second->Name()); + } + + engine_info.device = "/GPU:1"; + { + // Set to use second device. + auto result = GetDeviceAndAllocator(params, engine_info); + EXPECT_EQ(0, result.first); + EXPECT_NE(nullptr, result.second); + EXPECT_EQ("GPU_1_bfc", result.second->Name()); + } + + engine_info.device = "/GPU:3"; + { + // Set to use nonexistent device. + auto result = GetDeviceAndAllocator(params, engine_info); + EXPECT_EQ(-1, result.first); + EXPECT_EQ(nullptr, result.second); + } +} + +} // namespace convert +} // namespace tensorrt +} // namespace tensorflow + +#endif // GOOGLE_TENSORRT +#endif // GOOGLE_CUDA diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc new file mode 100644 index 0000000000000000000000000000000000000000..c08582a42e24fd55e785ad045725e06f1d414bfd --- /dev/null +++ b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc @@ -0,0 +1,3916 @@ +/* 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/convert/convert_nodes.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "absl/strings/str_cat.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 +#include "tensorflow/core/framework/tensor_shape.pb.h" // NOLINT +#include "tensorflow/core/framework/types.h" +#include "tensorflow/core/graph/algorithm.h" +#include "tensorflow/core/graph/graph.h" +#include "tensorflow/core/graph/graph_constructor.h" +#include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/lib/strings/numbers.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/protobuf.h" +#include "tensorflow/core/platform/tensor_coding.h" +#include "tensorflow/core/platform/types.h" + +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT +#include "tensorrt/include/NvInfer.h" + +// Check if the types are equal. Cast to int first so that failure log message +// would work! +#define TFTRT_CHECK_EQ_TYPE(val1, val2) CHECK_EQ((int)val1, (int)val2) + +#define TFTRT_INTERNAL_ERROR_AT_NODE(node) \ + do { \ + return tensorflow::errors::Internal( \ + "TFTRT::", __FUNCTION__, " failed to add TRT layer, at: ", node); \ + } while (0) + +#define TFTRT_RETURN_ERROR_IF_FALSE(status, node) \ + do { \ + if (status == false) { \ + TFTRT_INTERNAL_ERROR_AT_NODE(node); \ + } \ + } while (0) + +#define TFTRT_RETURN_ERROR_IF_NULLPTR(ptr, node) \ + do { \ + if (ptr == nullptr) { \ + TFTRT_INTERNAL_ERROR_AT_NODE(node); \ + } \ + } while (0) + +namespace tensorflow { +namespace tensorrt { +// TODO(aaroey): put these constants into some class. +const char* const kInputPHName = "TensorRTInputPH_"; +const char* const kOutputPHName = "TensorRTOutputPH_"; + +namespace convert { +using absl::StrAppend; +using absl::StrCat; +using ::tensorflow::str_util::Split; + +inline tensorflow::Status ConvertDType(tensorflow::DataType tf_dtype, + nvinfer1::DataType* trt_dtype) { + switch (tf_dtype) { + case tensorflow::DataType::DT_FLOAT: + *trt_dtype = nvinfer1::DataType::kFLOAT; + break; + // TODO(aaroey): this should be DT_QINT8 which is not a well supported type. + case tensorflow::DataType::DT_INT8: + *trt_dtype = nvinfer1::DataType::kINT8; + break; + case tensorflow::DataType::DT_HALF: + *trt_dtype = nvinfer1::DataType::kHALF; + break; + case tensorflow::DataType::DT_INT32: + *trt_dtype = nvinfer1::DataType::kINT32; + break; + default: + return tensorflow::errors::InvalidArgument( + "Unsupported data type ", tensorflow::DataTypeString(tf_dtype)); + } + return tensorflow::Status::OK(); +} + +template +inline nvinfer1::Dims TensorShapeToTrtDims(const TensorShapeType& shape, + bool ignore_first_dim) { + nvinfer1::Dims trt_dims; + const int offset = (ignore_first_dim ? 1 : 0); + for (int i = offset; i < shape.dims(); i++) { + trt_dims.d[i - offset] = shape.dim_size(i); + } + trt_dims.nbDims = shape.dims() - offset; + return trt_dims; +} + +Status TensorShapeArrayToTrtDims(const std::vector& shape, + nvinfer1::Dims* out, + bool ignore_first_dim = false) { + PartialTensorShape tensor_shape; + TF_RETURN_IF_ERROR(TensorShapeUtils::MakeShape(shape, &tensor_shape)); + *out = TensorShapeToTrtDims(tensor_shape, ignore_first_dim); + return tensorflow::Status::OK(); +} + +void GetOutputProperties(const grappler::GraphProperties& graph_properties, + const Node* node, const int out_port, + PartialTensorShape* shape, + tensorflow::DataType* dtype) { + if (graph_properties.HasOutputProperties(node->name())) { + auto output_params = graph_properties.GetOutputProperties(node->name()); + auto out_shape = output_params.at(out_port); + *dtype = out_shape.dtype(); + *shape = out_shape.shape(); + } else { + LOG(INFO) << "Unknown output shape" << node->name(); + *dtype = node->output_type(out_port); + } +} + +void GetInputProperties(const grappler::GraphProperties& graph_properties, + const Node* node, const int in_port, + PartialTensorShape* shape, + tensorflow::DataType* dtype) { + if (graph_properties.HasInputProperties(node->name())) { + auto input_params = graph_properties.GetInputProperties(node->name()); + auto in_shape = input_params.at(in_port); + *dtype = in_shape.dtype(); + *shape = in_shape.shape(); + } else { + *dtype = node->input_type(in_port); + } +} + +Status ValidateTensorProperties(const string& producer_node_type, + const tensorflow::DataType dtype, + const PartialTensorShape& shape, + bool validation_only, + nvinfer1::DataType* trt_dtype, + nvinfer1::Dims* trt_dims, int* batch_size) { + // Convert data type. + TF_RETURN_IF_ERROR(ConvertDType(dtype, trt_dtype)); + + // Convert shape. + if (shape.dims() < 0) { + return errors::InvalidArgument("Input tensor rank is unknown."); + } + if (shape.dims() > nvinfer1::Dims::MAX_DIMS + 1) { // +1 for batch dim + return errors::OutOfRange("Input tensor rank is greater than ", + nvinfer1::Dims::MAX_DIMS + 1); + } + if (producer_node_type != "Const" && shape.dims() < 2) { + return errors::InvalidArgument( + "Input tensor with rank<2 is not supported since the first dimension " + "is treated as batch dimension by TRT"); + } + *trt_dims = TensorShapeToTrtDims(shape, /*ignore_first_dim=*/true); + *batch_size = shape.dim_size(0); + + if (validation_only) return Status::OK(); + // Following are validations at runtime. + + for (int d = 1; d < shape.dims(); ++d) { + if (shape.dim_size(d) < 0) { + return errors::InvalidArgument( + "Input tensor with shape ", shape.DebugString(), + " has an unknown non-batch dimension at dim ", d); + } + } + return Status::OK(); +} + +string DebugString(const nvinfer1::DimensionType type) { + switch (type) { + case nvinfer1::DimensionType::kSPATIAL: + return "kSPATIAL"; + case nvinfer1::DimensionType::kCHANNEL: + return "kCHANNEL"; + case nvinfer1::DimensionType::kINDEX: + return "kINDEX"; + case nvinfer1::DimensionType::kSEQUENCE: + return "kSEQUENCE"; + default: + return StrCat(static_cast(type), "=unknown"); + } +} + +string DebugString(const nvinfer1::DataType trt_dtype) { + switch (trt_dtype) { + case nvinfer1::DataType::kFLOAT: + return "kFLOAT"; + case nvinfer1::DataType::kHALF: + return "kHALF"; + case nvinfer1::DataType::kINT8: + return "kINT8"; + case nvinfer1::DataType::kINT32: + return "kINT32"; + default: + return "Invalid TRT data type"; + } +} + +string DebugString(const nvinfer1::Dims& dims) { + string out = StrCat("nvinfer1::Dims(nbDims=", dims.nbDims, ", d="); + for (int i = 0; i < dims.nbDims; ++i) { + StrAppend(&out, dims.d[i], "[", DebugString(dims.type[i]), "],"); + } + StrAppend(&out, ")"); + return out; +} + +string DebugString(const nvinfer1::Permutation& permutation, int len) { + string out = "nvinfer1::Permutation("; + for (int i = 0; i < len; ++i) { + StrAppend(&out, permutation.order[i], ","); + } + StrAppend(&out, ")"); + return out; +} + +string DebugString(const nvinfer1::ITensor& tensor) { + return StrCat("nvinfer1::ITensor(@", reinterpret_cast(&tensor), + ", name=", tensor.getName(), + ", dtype=", DebugString(tensor.getType()), + ", dims=", DebugString(tensor.getDimensions()), ")"); +} + +Status Converter::GetTrtBroadcastShape( + const TRT_TensorOrWeights& operand_l, const TRT_TensorOrWeights& operand_r, + nvinfer1::Dims* operand_l_new_dims, + nvinfer1::Dims* operand_r_new_dims) const { + // *************************************************************************** + // TensorRT Elementwise op supports broadcast but requires both tensor to be + // of Identical rank + // + // We consider case of: + // 1. operand_l to be a Tensor & operand_r to be a Const; + // 2. operand_l to be a Tensor & operand_r to be a Tensor; + // note: const op const (constant folding) should fallback to TensorFlow + // + // broadcast scheme: + // T: 1 3 5 (tensor would not have batch dimension) + // W: 1 1 3 1 (weight would have all explicit dimensions) + // i. fill in explicit dimensions + // -> T: -1 1 3 5 (we put a -1 for batch dimension) + // -> W: 1 1 3 1 + // ii. compare broadcast feasibility + // + // We cannot support the following since TensorRT does not allow manipulation + // on batch dimension, we cannot generate output with proper shape + // T: 3 5 1 + // W: 1 1 1 1 3 5 1 + // -> T: 1 1 1 -1 3 5 1 + // -> W: 1 1 1 1 3 5 1 + // *************************************************************************** + if (!operand_l.is_tensor() && !operand_r.is_tensor()) { + return errors::InvalidArgument( + "Broadcasting requires at least one of the operands be tensors"); + } + + 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 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, + output_dims_array + broadcast_num_dims - input_dims.nbDims); + if (input.is_tensor()) { + const int true_input_dims = input_dims.nbDims + 1; + if (true_input_dims < broadcast_num_dims) { + return errors::InvalidArgument( + "Broadcasting beyond batch dimension is not supported ", + "(tensor #dims ", true_input_dims, " vs broadcast #dims ", + broadcast_num_dims, ")"); + } + // Set the batch dimension to -1, since batch size is not supposed to + // be broadcasted. + output_dims_array[0] = -1; + } + // Copy to output dimensions (stripping the batch dimension). + output_dims->nbDims = broadcast_num_dims - 1; + std::copy(output_dims_array + 1, output_dims_array + broadcast_num_dims, + output_dims->d); + return Status::OK(); + }; + + // Compute the output dimensions. + const int broadcast_num_dims = + std::max(operand_l.GetTrtDims().nbDims + (operand_l.is_tensor() ? 1 : 0), + operand_r.GetTrtDims().nbDims + (operand_r.is_tensor() ? 1 : 0)); + int output_l[max_nb_dims], output_r[max_nb_dims]; + TF_RETURN_IF_ERROR(compute_output_dims(operand_l, broadcast_num_dims, + output_l, operand_l_new_dims)); + TF_RETURN_IF_ERROR(compute_output_dims(operand_r, broadcast_num_dims, + output_r, operand_r_new_dims)); + + // Compare broadcast feasibility + for (int i = 0; i < broadcast_num_dims; ++i) { + if ((output_l[i] != output_r[i]) && (output_l[i] != 1) && + (output_r[i] != 1)) { + return errors::InvalidArgument( + "Infeasible broadcast scheme (", "batch_dim: ", output_l[0], ", ", + DebugString(*operand_l_new_dims), " vs ", "batch_dim: ", output_r[0], + ", ", DebugString(*operand_r_new_dims), ")"); + } + } + 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; +} + +inline bool DimsEqual(const nvinfer1::Dims& dim_l, + const nvinfer1::Dims& dim_r) { + if (dim_l.nbDims != dim_r.nbDims) { + return false; + } + for (int i = 0; i < dim_l.nbDims; i++) { + if (dim_l.d[i] != dim_r.d[i]) { + return false; + } + } + return true; +} + +inline nvinfer1::Dims GetTrtDimsForTensor(const tensorflow::Tensor& tensor) { + nvinfer1::Dims dims; + dims.nbDims = tensor.dims(); + for (int i = 0; i < dims.nbDims; i++) { + dims.d[i] = tensor.dim_size(i); + } + return dims; +} + +inline bool HasStaticShape(const nvinfer1::Dims& dims) { + if (dims.nbDims < 0) return false; + for (int d = 0; d < dims.nbDims; ++d) { + if (dims.d[d] < 0) return false; + } + return true; +} + +// Returns total number of elements in dims. Returning 0 means either some dim +// is 0 or the number of dims is 0. +// Note that for TF scalar constant, we always convert to dims [1]. +int64_t TrtDimsNumElements(const nvinfer1::Dims& dims) { + if (dims.nbDims == 0) return 0; + int64_t count = 1; + for (int d = 0; d < dims.nbDims; ++d) { + count *= dims.d[d]; + } + return count; +} + +static std::vector> CreateSamePadding( + const nvinfer1::DimsHW& stride, const nvinfer1::DimsHW& kernel, + const std::vector& input_dims) { + std::vector> padding(input_dims.size()); + CHECK_EQ(stride.nbDims, input_dims.size()); // TODO(jie): N+C? NC+? + + for (size_t i = 0; i < input_dims.size(); ++i) { + // Formula to calculate the padding + int p = ((input_dims[i] - 1) / stride.d[i]) * stride.d[i] + kernel.d[i] - + input_dims[i]; + p = (p > 0) ? p : 0; + + // Right precedence padding, like in TensorFlow + int left = p / 2; + int right = p - left; + + VLOG(2) << "PADDING_" << i << " pre: " << left << ", post: " << right + << "paras: " << input_dims[i] << ", " << stride.d[i] << ", " + << "kernel: " << kernel.d[i]; + padding[i] = {left, right}; + } + return padding; +} + +string GetCommonNameScope(const string& op_name_a, const string& op_name_b) { + size_t last_scope_separator = 0; + const size_t min_size = std::min(op_name_a.size(), op_name_b.size()); + for (size_t i = 0; i < min_size; ++i) { + if (op_name_a[i] != op_name_b[i]) break; + if (op_name_a[i] == '/') last_scope_separator = i + 1; + } + return op_name_a.substr(0, last_scope_separator); +} + +TRT_ShapedWeights::TRT_ShapedWeights(DataType type) : type_(type) { + shape_.nbDims = 0; +} + +TRT_ShapedWeights::TRT_ShapedWeights(DataType type, nvinfer1::Dims dims, + Tensor tensor) + : shape_(dims), type_(type), tensor_(tensor) {} + +TRT_ShapedWeights::TRT_ShapedWeights(const TRT_ShapedWeights& rhs) + : shape_(rhs.shape_), type_(rhs.type_), tensor_(rhs.tensor_) {} + +int64_t TRT_ShapedWeights::count() const { return TrtDimsNumElements(shape_); } + +nvinfer1::Weights TRT_ShapedWeights::GetTrtWeights() const { + nvinfer1::DataType trt_type(nvinfer1::DataType::kFLOAT); + TF_CHECK_OK(ConvertDType(type_, &trt_type)); + return nvinfer1::Weights{trt_type, GetValues(), count()}; +} + +size_t TRT_ShapedWeights::size_bytes() const { + return this->count() * tensorflow::DataTypeSize(this->type_); +} + +string TRT_ShapedWeights::DebugString() const { + return StrCat("TRT_ShapedWeights(shape=", convert::DebugString(shape_), + ", type=", DataTypeString(type_), + ", values=", reinterpret_cast(GetValues()), ")"); +} + +// A fake ITensor implementation used to check whether the TF-TRT converter can +// handle specific node. We only need shape and type information, and the +// converter won't (and shouldn't) use this to build the TRT network. +class TRT_TensorOrWeights::SimpleITensor : public nvinfer1::ITensor { + public: + SimpleITensor(nvinfer1::DataType trt_dtype, const nvinfer1::Dims& trt_dims) + : trt_dtype_(trt_dtype), trt_dims_(trt_dims) {} + + void setName(const char* name) override {} + + const char* getName() const override { return ""; } + + void setDimensions(nvinfer1::Dims dimensions) override { + trt_dims_ = dimensions; + } + + nvinfer1::Dims getDimensions() const override { return trt_dims_; } + + void setType(nvinfer1::DataType trt_dtype) override { + trt_dtype_ = trt_dtype; + } + + nvinfer1::DataType getType() const override { return trt_dtype_; } + + bool isNetworkInput() const override { return false; } + + bool isNetworkOutput() const override { return false; } + + void setBroadcastAcrossBatch(bool broadcastAcrossBatch) override {} + + bool getBroadcastAcrossBatch() const override { return false; } + + nvinfer1::TensorLocation getLocation() const override { + // This is arbitrary, since we don't use it. + return nvinfer1::TensorLocation::kDEVICE; + } + + void setLocation(nvinfer1::TensorLocation location) override {} + +#if NV_TENSORRT_MAJOR >= 5 + bool setDynamicRange(float min, float max) override { return true; } + + float getDynamicRange() const override { return 0; } +#endif + + private: + nvinfer1::DataType trt_dtype_; + nvinfer1::Dims trt_dims_; +}; + +TRT_TensorOrWeights::TRT_TensorOrWeights(nvinfer1::ITensor* tensor, + int batch_size) + : tensor_(tensor), + batch_size_(batch_size), + initialized_(true), + is_tensor_(true) {} + +TRT_TensorOrWeights::TRT_TensorOrWeights(nvinfer1::DataType trt_dtype, + const nvinfer1::Dims& trt_dims, + int batch_size) + : simple_itensor_(new SimpleITensor(trt_dtype, trt_dims)), + batch_size_(batch_size), + initialized_(true), + is_tensor_(true) {} + +TRT_TensorOrWeights::TRT_TensorOrWeights(const TRT_ShapedWeights& weights) + : weights_(weights), initialized_(true), is_tensor_(false) {} + +TRT_TensorOrWeights::TRT_TensorOrWeights(const TRT_TensorOrWeights& rhs) + : tensor_(rhs.tensor_), + simple_itensor_(rhs.simple_itensor_), + batch_size_(rhs.batch_size_), + weights_(rhs.weights_), + initialized_(rhs.initialized_), + is_tensor_(rhs.is_tensor_) {} + +void TRT_TensorOrWeights::operator=(const TRT_TensorOrWeights& rhs) { + tensor_ = rhs.tensor_; + simple_itensor_ = rhs.simple_itensor_; + batch_size_ = rhs.batch_size_; + weights_ = rhs.weights_; + initialized_ = rhs.initialized_; + is_tensor_ = rhs.is_tensor_; +} + +nvinfer1::ITensor* TRT_TensorOrWeights::tensor() { + CHECK(is_tensor()); + return tensor_ == nullptr ? simple_itensor_.get() : tensor_; +} + +const nvinfer1::ITensor* TRT_TensorOrWeights::tensor() const { + CHECK(is_tensor()); + return tensor_ == nullptr ? simple_itensor_.get() : tensor_; +} + +nvinfer1::Dims TRT_TensorOrWeights::GetTrtDims() const { + if (is_tensor()) { + return tensor()->getDimensions(); + } else { + return weights().shape_; + } +} + +string TRT_TensorOrWeights::DebugString() const { + string output = "TRT_TensorOrWeights(type="; + if (is_tensor()) { + StrAppend(&output, "tensor=", convert::DebugString(*tensor()), + ", batch_size=", batch_size_); + } else { + StrAppend(&output, "weights=", weights_.DebugString()); + } + StrAppend(&output, ")"); + return output; +} + +class TFAttrs { + public: + explicit TFAttrs(const tensorflow::NodeDef& tf_node) { + for (const auto& attr : tf_node.attr()) { + attrs_.insert({attr.first, &attr.second}); + } + } + + bool count(const string& key) const { return attrs_.count(key); } + + tensorflow::AttrValue const* at(const string& key) const { + if (!attrs_.count(key)) { + LOG(FATAL) << "Attribute not found: " << key; + } + return attrs_.at(key); + } + + template + T get(const string& key) const; + + template + T get(const string& key, const T& default_value) const { + return attrs_.count(key) ? this->get(key) : default_value; + } + + std::vector GetAllAttrKeys() const { + std::vector attr_list; + for (const auto& attr_item : attrs_) { + attr_list.emplace_back(attr_item.first); + } + return attr_list; + } + + private: + typedef std::map AttrMap; + AttrMap attrs_; +}; + +template <> +string TFAttrs::get(const string& key) const { + return this->at(key)->s(); +} + +template <> +std::vector TFAttrs::get>(const string& key) const { + auto attr = this->at(key)->list().i(); + return std::vector(attr.begin(), attr.end()); +} + +template <> +std::vector TFAttrs::get>(const string& key) const { + auto attr = this->at(key)->list().f(); + return std::vector(attr.begin(), attr.end()); +} + +template <> +nvinfer1::DataType TFAttrs::get(const string& key) const { + nvinfer1::DataType trt_dtype(nvinfer1::DataType::kFLOAT); + TF_CHECK_OK(ConvertDType(this->at(key)->type(), &trt_dtype)); + return trt_dtype; +} + +template <> +tensorflow::DataType TFAttrs::get( + const string& key) const { + return this->at(key)->type(); +} + +template <> +float TFAttrs::get(const string& key) const { + return this->at(key)->f(); +} + +template <> +bool TFAttrs::get(const string& key) const { + return this->at(key)->b(); +} + +template <> +int TFAttrs::get(const string& key) const { + return this->at(key)->i(); +} + +// TODO(jie): reorder4 & reorder2 should be merged? +// TODO(aaroey): fix the order of parameters. +template +void Reorder4(const nvinfer1::DimsNCHW& shape, const T* idata, + const nvinfer1::DimsNCHW& istrides, T* odata, + const nvinfer1::DimsNCHW& ostrides) { + for (int n = 0; n < shape.n(); ++n) { + for (int c = 0; c < shape.c(); ++c) { + for (int h = 0; h < shape.h(); ++h) { + for (int w = 0; w < shape.w(); ++w) { + odata[n * ostrides.n() + c * ostrides.c() + h * ostrides.h() + + w * ostrides.w()] = idata[n * istrides.n() + c * istrides.c() + + h * istrides.h() + w * istrides.w()]; + } + } + } + } +} + +template +void Reorder2(const nvinfer1::DimsHW& shape, const T* idata, + const nvinfer1::DimsHW& istrides, T* odata, + const nvinfer1::DimsHW& ostrides) { + for (int h = 0; h < shape.h(); ++h) { + for (int w = 0; w < shape.w(); ++w) { + odata[h * ostrides.h() + w * ostrides.w()] = + idata[h * istrides.h() + w * istrides.w()]; + } + } +} + +// TODO(jie): fallback to tensorflow!! +void ReorderCKtoKC(const TRT_ShapedWeights& iweights, + TRT_ShapedWeights* oweights) { + const int c = iweights.shape_.d[0]; + const int k = iweights.shape_.d[1]; + oweights->shape_.d[0] = k; + oweights->shape_.d[1] = c; + const nvinfer1::DimsHW istrides = {1, k}; + const nvinfer1::DimsHW ostrides = {c, 1}; + switch (iweights.type_) { + case tensorflow::DataType::DT_FLOAT: { + Reorder2({k, c}, static_cast(iweights.GetValues()), + istrides, + // TODO(aaroey): get rid of all the const_cast like this. + static_cast(const_cast(oweights->GetValues())), + ostrides); + break; + } + case tensorflow::DataType::DT_HALF: { + Reorder2( + {k, c}, static_cast(iweights.GetValues()), + istrides, + static_cast(const_cast(oweights->GetValues())), + ostrides); + break; + } + default: + LOG(FATAL) << "Unsupported type in reorder expected fp32 or fp16 but got " + << DataTypeString(iweights.type_); + } +} + +void ReorderRSCKToKCRS(const TRT_ShapedWeights& iweights, + TRT_ShapedWeights* oweights, const int num_groups) { + CHECK_EQ(iweights.type_, oweights->type_); + CHECK_EQ(iweights.size_bytes(), oweights->size_bytes()); + // K indexes over output channels, C over input channels, and R and S over the + // height and width of the convolution + const int r = iweights.shape_.d[0]; + const int s = iweights.shape_.d[1]; + // TRT requires GKcRS, while TF depthwise has RSCK where c=1, C=G + const int c = iweights.shape_.d[2] / num_groups; + const int k = iweights.shape_.d[3] * num_groups; + VLOG(2) << "num_groups: " << num_groups << "c" << iweights.shape_.d[2] + << " then " << c << "k" << iweights.shape_.d[3] << " then " << k + << "r" << iweights.shape_.d[0] << " then " << r << "s" + << iweights.shape_.d[1] << " then " << s; + oweights->shape_.d[0] = k / num_groups; + oweights->shape_.d[1] = c * num_groups; + oweights->shape_.d[2] = r; + oweights->shape_.d[3] = s; + const nvinfer1::DimsNCHW istrides = {1, k, s * k * c, c * k}; + const nvinfer1::DimsNCHW ostrides = {c * r * s, r * s, s, 1}; + switch (iweights.type_) { + case tensorflow::DataType::DT_FLOAT: { + Reorder4({k, c, r, s}, static_cast(iweights.GetValues()), + istrides, + static_cast(const_cast(oweights->GetValues())), + ostrides); + break; + } + case tensorflow::DataType::DT_HALF: { + Reorder4( + {k, c, r, s}, static_cast(iweights.GetValues()), + istrides, + static_cast(const_cast(oweights->GetValues())), + ostrides); + break; + } + + default: + LOG(FATAL) << "Unsupported type, expected fp32 or fp16 but got " + << DataTypeString(iweights.type_); + } +} + +TRT_ShapedWeights TrtWeightStore::GetTempWeights(tensorflow::DataType type, + const nvinfer1::Dims& dims) { + TensorShape shape; + // TODO(laigd): make it return a status. + TF_CHECK_OK(TensorShapeUtils::MakeShape(dims.d, dims.nbDims, &shape)); + // TODO(jie): check weights size_bytes. 0 means type error + Tensor tensor(type, shape); + TRT_ShapedWeights weights(type, dims, tensor); + store_.emplace_back(std::move(tensor)); + return weights; +} + +TrtNodeValidator::TrtNodeValidator() { RegisterOpValidators(); } + +Status TrtNodeValidator::ConvertToTensorOrWeights( + const NodeDef& node_def, int output_port, + const grappler::GraphProperties& graph_properties, + TRT_TensorOrWeights* tensor_or_weights) { + if (node_def.op() == "Const") { + if (output_port != 0) { + return errors::InvalidArgument("Const node should only have one output."); + } + // The output of the conversion will be used as input to other nodes to + // determine whether TRT supports those nodes. If it cannot convert the + // Const, it's very likely we cannot treat it as a tensor and make it an + // input to the TRT network, since TRT removes the first dimension and + // treats it as batch size. Also, it's not likely that the converter can + // support the op, and performance may suffer even if it can, so we just + // simply return error if the conversion fails. + std::vector inputs; + return ConvertConstToWeights(node_def, inputs, tensor_or_weights); + } + if (!graph_properties.HasOutputProperties(node_def.name())) { + return errors::InvalidArgument("Shape and data type are unknown"); + } + + // Validate and convert shape and dtype. + const auto& output_params = + graph_properties.GetOutputProperties(node_def.name()); + const auto& tensor_properties = output_params.at(output_port); + const DataType dtype = tensor_properties.dtype(); + const PartialTensorShape shape = tensor_properties.shape(); + nvinfer1::DataType trt_dtype; + nvinfer1::Dims trt_dims; + int batch_size = -1; + TF_RETURN_IF_ERROR(ValidateTensorProperties( + node_def.op(), dtype, shape, /*validation_only_=*/true, &trt_dtype, + &trt_dims, &batch_size)); + + // Adds a fake ITensor. This is fine since op converter operates in + // validation-only mode and it won't (and shouldn't) use the tensor to do + // any TRT network operations. + *tensor_or_weights = TRT_TensorOrWeights(trt_dtype, trt_dims, batch_size); + return Status::OK(); +} + +Status TrtNodeValidator::ValidateNode( + const tensorflow::NodeDef& node_def, + const std::vector>& input_node_and_ports, + const grappler::GraphProperties& graph_properties) { + // Convert input NodeDef and corresponding output ports to + // TRT_TensorOrWeights. + std::vector inputs; + for (int i = 0; i < input_node_and_ports.size(); ++i) { + const auto& pair = input_node_and_ports[i]; + TRT_TensorOrWeights tensor_or_weights; + Status status = ConvertToTensorOrWeights( + *pair.first, pair.second, graph_properties, &tensor_or_weights); + if (!status.ok()) { + return errors::Internal( + "Failed to convert input with index ", i, + " to a TRT_TensorOrWeights: ", status.error_message()); + } + inputs.push_back(tensor_or_weights); + } + + // Validate the node. + const auto iter = op_validators_.find(node_def.op()); + if (iter == op_validators_.end()) { + // If validator is not registered, it means no validation is needed. + return Status::OK(); + } + + OpConverter validator = iter->second; + OpConverterParams params( + /*arg_converter=*/nullptr, node_def, inputs, /*arg_outputs=*/nullptr, + /*arg_validation_only=*/true, &weight_store_); + return validator(¶ms); +} + +Status TrtNodeValidator::ConvertConstToWeights( + const NodeDef& const_node_def, + const std::vector& inputs, + TRT_TensorOrWeights* output) { + std::vector outputs; + OpConverterParams params( + /*arg_converter=*/nullptr, const_node_def, inputs, &outputs, + /*arg_validation_only=*/true, &weight_store_); + Status status = op_validators_["Const"](¶ms); + if (status.ok() && output) *output = outputs[0]; + return status; +} + +Converter::Converter(nvinfer1::INetworkDefinition* trt_network, + int precision_mode, bool use_calibration) + : trt_network_(trt_network), + precision_mode_(precision_mode), + use_calibration_(use_calibration) { + this->RegisterOpConverters(); +} + +Status Converter::ConvertNode(const NodeDef& node_def) { + std::vector inputs, outputs; + TF_RETURN_IF_ERROR(this->GetInputs(node_def, &inputs)); + + OpConverterParams params(this, node_def, inputs, &outputs, + /*arg_validation_only=*/false, &weight_store_); + const string& op = node_def.op(); + if (PluginFactoryTensorRT::GetInstance()->IsPlugin(op)) { + TF_RETURN_IF_ERROR(plugin_converter_(¶ms)); + } else { + if (!op_registry_.count(op)) { + return errors::Unimplemented("No converter registered for op: " + op); + } + OpConverter op_converter = op_registry_.at(op); + TF_RETURN_IF_ERROR(op_converter(¶ms)); + } + + 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); + // 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)) { + // 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 + // not inputs or outputs of a node will not be renamed. This is a + // potential cause of confusion if an error message or warning + // mentions the unnamed tensor. + output.tensor()->setName(output_name.c_str()); + } + } + VLOG(2) << "Adding out tensor " << output_name << ": " + << output.DebugString(); + Status status = AddTensorOrWeights(output_name, output); + if (!status.ok()) { + return Status(status.code(), + StrCat("Failed to add output for node ", node_def.name(), + ": ", status.error_message())); + } + } + return Status::OK(); +} + +Status Converter::AddInputTensor(const string& name, nvinfer1::DataType dtype, + const nvinfer1::Dims& dims, int batch_size) { + // We verify the batch size only for the input nodes, and rely on individual + // op converter to ensure the batch size of the outputs is not changed. + // TODO(laigd): we need to test this properties. + Status status = MaybeUpdateBatchSize(batch_size); + if (!status.ok()) { + return Status(status.code(), StrCat("Batch size doesn't match for tensor ", + name, ": ", status.error_message())); + } + nvinfer1::ITensor* tensor = network()->addInput(name.c_str(), dtype, dims); + if (tensor == nullptr) { + return errors::InvalidArgument("Failed to create Input layer tensor ", name, + " rank=", dims.nbDims); + } + status = AddTensorOrWeights(name, TRT_TensorOrWeights(tensor)); + if (!status.ok()) { + return Status(status.code(), StrCat("Failed to add input tensor ", name, + ": ", status.error_message())); + } + return Status::OK(); +} + +Status Converter::RenameAndMarkOutputTensors( + 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)); + if (!tensor_or_weights.is_tensor()) { + return errors::InvalidArgument("Output ", output.first, + " is weights not tensor"); + } + nvinfer1::ITensor* tensor = tensor_or_weights.tensor(); + if (tensor == nullptr) { + return errors::NotFound("Output tensor not found: ", output.first); + } + // Check if this tensor has already been marked as an 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 (tensorflow::str_util::StartsWith(tensor->getName(), kOutputPHName)) { + // 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; + network()->markOutput(*tensor); + } + return Status::OK(); +} + +Status Converter::MaybeUpdateBatchSize(int batch_size) { + // OK iff either is unknown or they equal to each other. + if (this->batch_size_ < 0 || batch_size < 0 || + this->batch_size_ == batch_size) { + if (this->batch_size_ < 0 && batch_size >= 0) { + this->batch_size_ = batch_size; + } + return Status::OK(); + } + return errors::InvalidArgument( + "Provided batch size does not match converter batch size: ", batch_size, + " vs ", batch_size_); +} + +Status Converter::AddTensorOrWeights(const string& name, + TRT_TensorOrWeights input) { + // Set the batch size of the tensor, using batch size collected from the + // input tensors to the TRT subgraph at the beginning of the conversion. + // We rely on the individual op converter to understand the semantics of the + // TF node, and make sure it doesn't change the batch size nor introduce + // intra-element dependency inside the batch. + if (input.is_tensor()) input.set_batch_size(batch_size_); + if (trt_tensors_.insert({name, std::move(input)}).second) return Status::OK(); + return errors::AlreadyExists("tensor/weights ", name, " already exist."); +} + +Status Converter::GetTensorOrWeights(const string& name, + TRT_TensorOrWeights* output) { + if (!trt_tensors_.count(name)) { + return errors::NotFound("Tensor or weights with name ", name, + " could not be found."); + } + *output = trt_tensors_.at(name); + return Status::OK(); +} + +Status Converter::TransposeTensor(nvinfer1::ITensor* input_tensor, + const std::vector& order_with_batch_dim, + const nvinfer1::ITensor** output_tensor) { + const auto dims = input_tensor->getDimensions(); + + if (order_with_batch_dim.size() - 1 != size_t(dims.nbDims)) { + return tensorflow::errors::InvalidArgument( + "Rank of perm for transpose does not match with that of the input."); + } + if (order_with_batch_dim[0] != 0) { + return tensorflow::errors::Unimplemented( + "Transpose at batch dimension is not supported."); + } + + nvinfer1::IShuffleLayer* layer = this->network()->addShuffle(*input_tensor); + TFTRT_RETURN_ERROR_IF_NULLPTR(layer, "TF-TRT Internal Transpose"); + MarkQuantizationRangesAsInferrable(input_tensor, layer->getOutput(0)); + + nvinfer1::Permutation permutation; + for (int32_t i = 0; i < dims.nbDims; ++i) { + permutation.order[i] = order_with_batch_dim[i + 1] - 1; + } + VLOG(1) << "TransposeTensor permutation: " + << DebugString(permutation, dims.nbDims); + layer->setFirstTranspose(permutation); + + nvinfer1::Dims reshape_dims; + reshape_dims.nbDims = dims.nbDims; + for (int32_t i = 0; i < reshape_dims.nbDims; ++i) { + reshape_dims.d[i] = 0; + // TODO(aaroey): why not transposing the types as well? + reshape_dims.type[i] = dims.type[i]; + } + layer->setReshapeDimensions(reshape_dims); + + *output_tensor = layer->getOutput(0); + return tensorflow::Status::OK(); +} + +Status Converter::GetWeightRange(const TRT_ShapedWeights& weights, + float* out_min, float* out_max) const { + switch (weights.type_) { + case DataType::DT_FLOAT: { + auto inp = static_cast(weights.GetValues()); + auto result = std::minmax_element(inp, inp + weights.count()); + *out_min = *result.first; + *out_max = *result.second; + break; + } + case DataType::DT_HALF: { + auto inp = static_cast(weights.GetValues()); + auto result = std::minmax_element(inp, inp + weights.count()); + *out_min = Eigen::half_impl::half_to_float(*result.first); + *out_max = Eigen::half_impl::half_to_float(*result.second); + break; + } + case DataType::DT_INT32: { + auto inp = static_cast(weights.GetValues()); + auto result = std::minmax_element(inp, inp + weights.count()); + *out_min = static_cast(*result.first); + *out_max = static_cast(*result.second); + break; + } + default: + return errors::Unimplemented( + "Data type not supported for GetWeightRange: ", + DataTypeString(weights.type_)); + } + return Status::OK(); +} + +Status Converter::PrepareTensorForShape(const TRT_TensorOrWeights& input, + const nvinfer1::Dims& dims, + const nvinfer1::ITensor** tensor) { + // If -1 is not used for one of the dims, we can check if the shapes are + // compatible. + bool can_check_shapes = true; + for (int i = 0; i < dims.nbDims; i++) { + if (dims.d[i] == -1) { + can_check_shapes = false; + break; + } + } + if (can_check_shapes && + TrtDimsNumElements(input.GetTrtDims()) != TrtDimsNumElements(dims)) { + return errors::InvalidArgument("Reshape shapes are not compatible (", + DebugString(input.GetTrtDims()), " vs ", + DebugString(dims), ")"); + } + + if (input.is_tensor()) { + if (DimsEqual(input.GetTrtDims(), dims)) { + *tensor = input.tensor(); + } else { + nvinfer1::IShuffleLayer* layer = this->network()->addShuffle( + *const_cast(input.tensor())); + TFTRT_RETURN_ERROR_IF_NULLPTR(layer, "TF-TRT Internal Reshape"); + layer->setReshapeDimensions(dims); + MarkQuantizationRangesAsInferrable( + const_cast(input.tensor()), layer->getOutput(0)); + *tensor = layer->getOutput(0); + } + } else { + *tensor = CreateConstantLayer(input.weights(), dims); + TFTRT_RETURN_ERROR_IF_NULLPTR(*tensor, "TF-TRT Internal Reshape"); + if (precision_mode() == INT8MODE && !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)]. + float min_range = 0.0f; + float max_range = 0.0f; + TF_RETURN_IF_ERROR( + GetWeightRange(input.weights(), &min_range, &max_range)); + // Avoid setting range to 0 because TRT will throw an error. If the + // weights are zero then the range doesn't matter: using 127.0f should + // ensure the quantized weight will be exactly zero. + if (min_range == 0.0f && max_range == 0.0f) { + min_range = -127.0f; + max_range = 127.0f; + } + ProvideQuantizationRange(const_cast(*tensor), + min_range, max_range); + } + } + return tensorflow::Status::OK(); +} + +void Converter::MarkQuantizationRangesAsInferrable(nvinfer1::ITensor* input, + nvinfer1::ITensor* output) { + quantization_infer_.push_back({input, output}); + quantization_infer_.push_back({output, input}); +} + +void Converter::ProvideQuantizationRange(nvinfer1::ITensor* tensor, + float min_range, float max_range) { + float symmetric_range = std::max(std::abs(min_range), std::abs(max_range)); + quantization_ranges_[tensor] = symmetric_range; +} + +void Converter::MaybeApplyQuantizationRanges() { + if (precision_mode() != INT8MODE) return; + + // Infer ranges across marked ops. + PropagateQuantizationRanges(); + // Apply ranges. +#if NV_TENSORRT_MAJOR >= 5 + for (auto pair : quantization_ranges_) { + nvinfer1::ITensor* tensor = pair.first; + const float range = pair.second; + VLOG(1) << "Setting range for: " << tensor->getName() << ": " << range; + // TODO(laigd): if 'tensor' already has a range set which doesn't match + // 'range', it should report error. + tensor->setDynamicRange(-range, range); + } +#endif + + // Warn user about tensors that are missing ranges. If TRT fuses some layers + // then these tensors may not actually be required, which is why this is + // just a warning. If we are still missing ranges even after fusion, + // Builder::buildCudaEngine() will return nullptr and we will catch the + // error at that point. + if (!use_calibration()) { + // Get all tensors from network + std::set all_tensors; + for (int i = 0; i < this->network()->getNbLayers(); i++) { + nvinfer1::ILayer* layer = this->network()->getLayer(i); + for (int j = 0; j < layer->getNbInputs(); j++) { + all_tensors.insert(layer->getInput(j)); + } + for (int j = 0; j < layer->getNbOutputs(); j++) { + all_tensors.insert(layer->getOutput(j)); + } + } + // Find tensors with no ranges + for (auto tensor : all_tensors) { + if (!quantization_ranges_.count(tensor)) { + // Note: there may be some warnings for "(Unnamed ITensor* N)". These + // are tensors which are created internally by TF-TRT. The ranges for + // these unnamed ITensors are always inferred from user provided ranges, + // thus there will also be a warning for the range(s) the user missed. + LOG(WARNING) << "Quantization range was not found for " + << tensor->getName() << ". " + << "This is okay if TensorRT does not need the range " + << "(e.g. due to node fusion)."; + } + } + } +} + +void Converter::PropagateQuantizationRanges() { + // Propagate ranges across edges in quantization_infer_ until no new + // information is added. + // Note: this function modifies quantization_infer_, it might be better to + // modify a copy instead if we for some reason need quantization_infer_ + // later. + bool information_added = true; + while (information_added) { + information_added = false; + for (auto it = quantization_infer_.begin(); + it != quantization_infer_.end();) { + auto input_tensor_range = quantization_ranges_.find(it->first); + auto output_tensor_range = quantization_ranges_.find(it->second); + if (input_tensor_range != quantization_ranges_.end() && + output_tensor_range == quantization_ranges_.end()) { + // Input has range but output doesn't: copy range + // TODO(laigd): consider reporting error if it a different range is + // already set. + quantization_ranges_[it->second] = input_tensor_range->second; + information_added = true; + VLOG(1) << "Copy quantization range: " << it->first->getName() << " -> " + << it->second->getName(); + } + // We can remove edges when the output range is known + if (quantization_ranges_.find(it->second) != quantization_ranges_.end()) { + it = quantization_infer_.erase(it); + } else { + ++it; + } + } + } +} + +Status Converter::GetInputs(const tensorflow::NodeDef& node_def, + std::vector* inputs) const { + for (auto const& input_name : node_def.input()) { + /************************************************************************* + * TODO(jie): handle case 1) here. + * Normalizes the inputs and extracts associated metadata: + * 1) Inputs can contain a colon followed by a suffix of characters. + * That suffix may be a single number (e.g. inputName:1) or several + * word characters separated from a number by a colon + * (e.g. inputName:foo:1). The + * latter case is used to denote inputs and outputs of functions. + * 2) Control dependency inputs contain caret at the beginning and we + * remove this and annotate the edge as a control dependency. + ************************************************************************/ + // skip control nodes + if (input_name[0] == '^') continue; + string name = input_name; + auto last = name.find_last_of(':'); + // TODO(aaroey): use TensorId + if (last != string::npos && last + 2 == name.size() && + name[last + 1] == '0') { + name.erase(last); + } + + if (trt_tensors_.count(name)) { + TRT_TensorOrWeights input = trt_tensors_.at(name); + inputs->push_back(input); + VLOG(2) << "Retrieved input " << name << ": " << input.DebugString(); + } else { + // TODO(aaroey): this should not happen, make it a CHECK. + // TODO(aaroey): use StrCat for pattern like this. + string msg("Node "); + StrAppend(&msg, node_def.name(), " should have an input named '", name, + "' but it is not available"); + LOG(ERROR) << msg; + return tensorflow::errors::InvalidArgument(msg); + } + } + return tensorflow::Status::OK(); +} + +TRT_ShapedWeights ConvertFP32ToFP16(TrtWeightStore* store, + const TRT_ShapedWeights& weights_src) { + auto dtype_new = tensorflow::DataType::DT_HALF; + TRT_ShapedWeights weights = + store->GetTempWeights(dtype_new, weights_src.shape_); + const float* src = static_cast(weights_src.GetValues()); + Eigen::half* dst = const_cast( + static_cast(weights.GetValues())); + for (int64_t i = 0; i < weights_src.count(); i++) { + dst[i] = Eigen::half_impl::float_to_half_rtne(src[i]); + } + return weights; +} + +// **************************************************************************** +// Constant folding functions for weights. +// TODO(laigd): we should probably use eigen directly. +// ***************************************************************************** +struct LambdaFactory { + enum class OP_CATEGORY : int { RSQRT = 0, NEG, RECIP }; + OP_CATEGORY op; + + template + std::function unary() { + switch (op) { + case OP_CATEGORY::RSQRT: { + VLOG(2) << "RSQRT GETS DONE"; + return [](T t) -> T { return 1.0 / sqrt(t); }; + } + case OP_CATEGORY::NEG: + return [](T t) -> T { return -t; }; + case OP_CATEGORY::RECIP: + return [](T t) -> T { return 1.0 / t; }; + default: + LOG(ERROR) << "Not supported op for unary: " << static_cast(op); + return nullptr; + } + } +}; + +template <> +std::function LambdaFactory::unary() { + switch (op) { + case OP_CATEGORY::RSQRT: { + VLOG(2) << "RSQRT GETS DONE"; + return [](Eigen::half t) { + return Eigen::half(1.0 / sqrt(static_cast(t))); + }; + } + case OP_CATEGORY::NEG: + return [](Eigen::half t) { return -t; }; + case OP_CATEGORY::RECIP: + return [](Eigen::half t) { + return Eigen::half(1.0 / static_cast(t)); + }; + default: + LOG(ERROR) << "Not supported op for unary: " << static_cast(op); + return nullptr; + } +} + +tensorflow::Status UnaryCompute(const TRT_ShapedWeights& iweights, + TRT_ShapedWeights* oweights, + LambdaFactory unary_op) { + CHECK_EQ(iweights.type_, oweights->type_); + switch (iweights.type_) { + case tensorflow::DataType::DT_FLOAT: { + auto inp = static_cast(iweights.GetValues()); + auto oup = static_cast(const_cast(oweights->GetValues())); + std::transform(inp, inp + iweights.count(), oup, unary_op.unary()); + break; + } + case tensorflow::DataType::DT_HALF: { + auto inp = static_cast(iweights.GetValues()); + auto oup = + static_cast(const_cast(oweights->GetValues())); + std::transform(inp, inp + iweights.count(), oup, + unary_op.unary()); + break; + } + default: + return tensorflow::errors::Unimplemented( + "Data type not supported: " + + tensorflow::DataTypeString(iweights.type_)); + } + return tensorflow::Status::OK(); +} + +// If swapped_inputs is false, 'tensor' is the left operand and 'weights' is the +// right operand. If swapped_inputs is true, those two are swapped. +// +// TODO(jie): broadcast is needed yet not implemented. +// Only implemented channel wise for the time being. +Status BinaryTensorOpWeight(OpConverterParams* params, + const nvinfer1::ITensor* tensor, + TRT_ShapedWeights weights, bool swapped_inputs) { + static const std::unordered_set supported_ops = {"Sub", "Add", "Mul", + "Div", "RealDiv"}; + const auto& node_def = params->node_def; + if (!supported_ops.count(node_def.op())) { + return errors::Unimplemented(node_def.op(), " is not supported, at ", + node_def.name()); + } + + // Check type consistency. + nvinfer1::DataType trt_dtype; + TF_RETURN_IF_ERROR(ConvertDType(weights.type_, &trt_dtype)); + + // Check scale mode. + auto dims_w = weights.shape_; + const auto dims_t = tensor->getDimensions(); + + // TODO(jie): addScale checks for input tensor dimension + if (dims_t.nbDims != 3) { + return errors::InvalidArgument("addScale requires tensor with rank 3, at ", + node_def.name()); + } + + // Default to element-wise + auto scale_mode = nvinfer1::ScaleMode::kELEMENTWISE; + + // TODO(jie): maybe use a permutation instead to support more cases; + bool need_to_permute = false; + + if (weights.count() == 1) { + scale_mode = nvinfer1::ScaleMode::kUNIFORM; + } else { + VLOG(2) << "weights dims: " << DebugString(dims_w) + << "; tensor dims: " << DebugString(dims_t); + // Make sure no broadcasting on batch dimension. + if (dims_w.nbDims == dims_t.nbDims + 1) { + if (dims_w.d[0] == 1) { + for (int i = 1; i < dims_w.nbDims; i++) { + dims_w.d[i - 1] = dims_w.d[i]; + } + dims_w.nbDims--; + } else { + return errors::InvalidArgument("Binary op cannot operate on batch, at ", + node_def.name()); + } + } + + if (dims_w.nbDims == dims_t.nbDims && dims_w.d[0] == dims_t.d[0]) { + scale_mode = nvinfer1::ScaleMode::kELEMENTWISE; + // Default is element-wise + for (int i = 1; i < dims_w.nbDims; i++) { + if (dims_w.d[i] != dims_t.d[i]) { + // If dimension does not match, switch back to per-channel + scale_mode = nvinfer1::ScaleMode::kCHANNEL; + break; + } + } + // If the mode is per-channel, since channel dimension is assumed to be + // the third to last dimension, we need to make sure all other dimensions + // have size 1. + if (scale_mode == nvinfer1::ScaleMode::kCHANNEL) { + for (int i = 1; i < dims_w.nbDims; i++) { + if (dims_w.d[i] != 1) + return errors::InvalidArgument( + "Weight dims not compatible for channel-wise broadcast at ", + node_def.name()); + } + } + } else if (dims_w.nbDims == 1 && + dims_w.d[0] == dims_t.d[dims_t.nbDims - 1]) { + // Channel wise and broadcast required. We compare the last dimension of + // the tensor shape because of tensorflow default broadcasting rules. + need_to_permute = true; + scale_mode = nvinfer1::ScaleMode::kCHANNEL; + } else { + return errors::InvalidArgument("Weight dims not compatible at ", + node_def.name()); + } + } + // TODO(laigd): we should add validation_only support in TransposeTensor() and + // PrepareTensorForShape(). + if (params->validation_only) return Status::OK(); + + // Transpose last dimension. + std::vector permutation(dims_t.nbDims + 1); + if (need_to_permute) { + // We swap the last dimension into channel for trt, because of tensorflow + // default broadcasting rules. + for (int i = 0; i < static_cast(permutation.size()); i++) { + permutation[i] = i; + } + permutation[1] = dims_t.nbDims; + permutation[dims_t.nbDims] = 1; + TF_RETURN_IF_ERROR(params->converter->TransposeTensor( + const_cast(tensor), permutation, &tensor)); + } + + if (params->converter->precision_mode() == FP16MODE) { + weights = ConvertFP32ToFP16(params->weight_store, weights); + } + + // Prepare weights + TRT_ShapedWeights shift_weights(weights.type_); + TRT_ShapedWeights scale_weights(weights.type_); + TRT_ShapedWeights power_weights(weights.type_); + + if (node_def.op() == "Sub") { + if (swapped_inputs) { + shift_weights = weights; + nvinfer1::IUnaryLayer* layer = params->converter->network()->addUnary( + *const_cast(tensor), + nvinfer1::UnaryOperation::kNEG); + TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); + // Since quantization ranges are symmetric, the same range as the input + // will work for the negation of the input. + params->converter->MarkQuantizationRangesAsInferrable( + const_cast(tensor), layer->getOutput(0)); + tensor = layer->getOutput(0); + } else { + TRT_ShapedWeights neg_weights = + params->weight_store->GetTempWeights(weights); + LambdaFactory unary_op; + unary_op.op = LambdaFactory::OP_CATEGORY::NEG; + TF_RETURN_IF_ERROR(UnaryCompute(weights, &neg_weights, unary_op)); + shift_weights = neg_weights; + } + } else if (node_def.op() == "Div" || node_def.op() == "RealDiv") { + if (swapped_inputs) { + // We need to infer the quantization range for this intermediate tensor. + // + // x -> [Recip] -> 1/x -> [Scale] -> s/x + // ^ + // need range for this + // + // We have the quantization scales for x and s/x - can we divide the scale + // for s/x by s? Only if it is a scalar. + // + // 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 && + !params->converter->use_calibration()) { + return errors::Unimplemented( + "Intermediate quantization range cannot be determined without" + " calibration. Falling back to BinaryTensorOpTensor for ", + node_def.op(), ", at ", node_def.name()); + } + scale_weights = weights; + nvinfer1::IUnaryLayer* layer = params->converter->network()->addUnary( + *const_cast(tensor), + nvinfer1::UnaryOperation::kRECIP); + TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); + tensor = layer->getOutput(0); + } else { + TRT_ShapedWeights recip_weights = + params->weight_store->GetTempWeights(weights); + LambdaFactory unary_op; + unary_op.op = LambdaFactory::OP_CATEGORY::RECIP; + TF_RETURN_IF_ERROR(UnaryCompute(weights, &recip_weights, unary_op)); + scale_weights = recip_weights; + } + } else if (node_def.op() == "Mul") { + scale_weights = weights; + } else if (node_def.op() == "Add") { + shift_weights = weights; + } else { + // This should not happen. + return errors::Unimplemented("Binary op not supported at ", node_def.op()); + } + + nvinfer1::IScaleLayer* layer = params->converter->network()->addScale( + *const_cast(tensor), scale_mode, + shift_weights.GetTrtWeights(), scale_weights.GetTrtWeights(), + power_weights.GetTrtWeights()); + TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); + + const nvinfer1::ITensor* output_tensor = layer->getOutput(0); + // Transpose back dimension + if (need_to_permute) { + TF_RETURN_IF_ERROR(params->converter->TransposeTensor( + const_cast(output_tensor), permutation, + &output_tensor)); + } + + // Pass the output + params->outputs->push_back( + TRT_TensorOrWeights(const_cast(output_tensor))); + return tensorflow::Status::OK(); +} + +enum class ConvolutionType { DEFAULT, DEPTHWISE_CONV }; + +tensorflow::Status ConvertConv2DHelper(OpConverterParams* params, int group) { + const auto& inputs = params->inputs; + const auto& node_def = params->node_def; + if (inputs.size() != 2) { + return tensorflow::errors::InvalidArgument("Two inputs are expected for ", + node_def.op(), ", at ", + node_def.name()); + } + if (inputs.at(0).is_weights()) { + return tensorflow::errors::Unimplemented( + node_def.op(), " is only implemented for tensors, not weights, 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()); + } + 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( + "Dilation rate must be 1 for batch and channel dimensions, at ", + node_def.name()); + } + const nvinfer1::DimsHW dilation(tf_dilations[h_index], tf_dilations[w_index]); + + 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(); + + // 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)); + } + // Dimensions of transposed tensor. + const auto tensor_dim = tensor->getDimensions(); + + // For depthwise convolution, group will be 0 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()); + } + 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; + nvinfer1::DimsHW kernel_size; + kernel_size.h() = weights.shape_.d[2]; + kernel_size.w() = weights.shape_.d[3]; + + // Add padding. + std::vector> padding; + if (attrs.get("padding") == "SAME") { + 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); + padding = CreateSamePadding( + stride, effective_kernel_size, + {static_cast(tensor_dim.d[1]), static_cast(tensor_dim.d[2])}); + } else { + padding = {{0, 0}, {0, 0}}; + } + if (padding[0].first != padding[0].second || + padding[1].first != padding[1].second) { + // Handle asymmetric padding. + auto pad_layer = params->converter->network()->addPadding( + *const_cast(tensor), + nvinfer1::DimsHW(padding[0].first, padding[1].first), + nvinfer1::DimsHW(padding[0].second, padding[1].second)); + TFTRT_RETURN_ERROR_IF_NULLPTR(pad_layer, node_def.name()); + params->converter->MarkQuantizationRangesAsInferrable( + const_cast(tensor), pad_layer->getOutput(0)); + padding = {{0, 0}, {0, 0}}; + tensor = pad_layer->getOutput(0); + } + + // Add convolution. + 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); + const nvinfer1::ITensor* output_tensor = layer->getOutput(0); + + // Restore transpose. + if (need_transpose) { + TF_RETURN_IF_ERROR(params->converter->TransposeTensor( + const_cast(output_tensor), {0, 2, 3, 1}, + &output_tensor)); + } + params->outputs->push_back( + TRT_TensorOrWeights(const_cast(output_tensor))); + 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) { + const auto& node_def = params->node_def; + static const std::unordered_map ops{ + {"Add", nvinfer1::ElementWiseOperation::kSUM}, + {"Mul", nvinfer1::ElementWiseOperation::kPROD}, + {"Sub", nvinfer1::ElementWiseOperation::kSUB}, + {"Div", nvinfer1::ElementWiseOperation::kDIV}, + {"RealDiv", nvinfer1::ElementWiseOperation::kDIV}, + {"Minimum", nvinfer1::ElementWiseOperation::kMIN}, + {"Maximum", nvinfer1::ElementWiseOperation::kMAX}, + }; + auto op_pair = ops.find(node_def.op()); + if (op_pair == ops.end()) { + return errors::Unimplemented("Binary op ", node_def.op(), + " not supported at: ", node_def.name()); + } + + nvinfer1::Dims broadcasted_dims_l, broadcasted_dims_r; + Status status = params->converter->GetTrtBroadcastShape( + operand_l, operand_r, &broadcasted_dims_l, &broadcasted_dims_r); + if (!status.ok()) { + return errors::InvalidArgument( + "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; + const nvinfer1::ITensor* tensor_r = nullptr; + status = params->converter->PrepareTensorForShape( + operand_l, broadcasted_dims_l, &tensor_l); + if (status.ok()) { + status = params->converter->PrepareTensorForShape( + operand_r, broadcasted_dims_r, &tensor_r); + } + if (!status.ok()) { + return errors::Internal("Failed to convert binary op ", node_def.name(), + ": ", status.error_message()); + } + + // Check type consistency. + TFTRT_CHECK_EQ_TYPE(tensor_l->getType(), dtype) + << DebugString(tensor_l->getType()) << " vs " << DebugString(dtype); + TFTRT_CHECK_EQ_TYPE(tensor_r->getType(), dtype) + << DebugString(tensor_r->getType()) << " vs " << DebugString(dtype); + + // Add ElementWise layer. + nvinfer1::IElementWiseLayer* layer = + params->converter->network()->addElementWise( + *const_cast(tensor_l), + *const_cast(tensor_r), op_pair->second); + TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); + nvinfer1::ITensor* output_tensor = layer->getOutput(0); + + // Pass the output + params->outputs->push_back(TRT_TensorOrWeights(output_tensor)); + return tensorflow::Status::OK(); +} + +tensorflow::Status ConvertPlugin(OpConverterParams* params) { + const auto& inputs = params->inputs; + const auto& node_def = params->node_def; + // prepare input + std::vector all_inputs; + all_inputs.reserve(inputs.size()); + for (auto input : inputs) { + all_inputs.emplace_back(const_cast(input.tensor())); + } + + // plugin is owned by PluginFactory + // TODO(jie): destroy plugins later (resource management) + PluginTensorRT* plugin = + PluginFactoryTensorRT::GetInstance()->CreatePlugin(node_def.op()); + + // passing attributes + // TODO(jie): support more general attribute + TFAttrs attrs(node_def); + auto attr_key_vector = attrs.GetAllAttrKeys(); + for (auto attr_key : attr_key_vector) { + // TODO(jie): support only list of float for toy example here. + auto data = attrs.get>(attr_key); + size_t size_data = data.size() * sizeof(float); + if (!plugin->SetAttribute(attr_key, static_cast(data.data()), + size_data)) { + return tensorflow::errors::InvalidArgument("plugin SetAttribute failed"); + } + } + + nvinfer1::IPluginLayer* layer = params->converter->network()->addPlugin( + &all_inputs[0], static_cast(inputs.size()), *plugin); + + for (int i = 0; i < layer->getNbOutputs(); i++) { + nvinfer1::ITensor* output_tensor = layer->getOutput(i); + params->outputs->push_back(TRT_TensorOrWeights(output_tensor)); + } + return tensorflow::Status::OK(); +} + +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()); + } + + // Get the permutation from weights. + TRT_ShapedWeights weights = inputs.at(1).weights(); + const int* weights_ptr = + static_cast(const_cast(weights.GetValues())); + std::vector perm(weights_ptr, weights_ptr + weights.count()); + + // Verify the permutation. + nvinfer1::ITensor* input_tensor = + const_cast(inputs.at(0).tensor()); + if (perm.size() - 1 != size_t(input_tensor->getDimensions().nbDims)) { + return errors::InvalidArgument( + "Rank of perm for transpose does not match with that of the input."); + } + if (perm[0] != 0) { + return errors::Unimplemented( + "Transpose at batch dimension is not supported."); + } + + if (params->validation_only) return Status::OK(); + + // Start conversion. + const nvinfer1::ITensor* output_tensor = nullptr; + TF_RETURN_IF_ERROR( + params->converter->TransposeTensor(input_tensor, perm, &output_tensor)); + params->outputs->push_back( + TRT_TensorOrWeights(const_cast(output_tensor))); + return tensorflow::Status::OK(); +} + +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()); + } + + TRT_TensorOrWeights input_tensor = inputs.at(0); + TRT_ShapedWeights weights = inputs.at(1).weights(); + if (weights.count() == 0) { + return tensorflow::errors::Unimplemented( + "Reshape to shape=[] is not supported, at ", node_def.name()); + } + + const int* weights_ptr = + static_cast(const_cast(weights.GetValues())); + + // Check that it doesn't change the batch dimension. This check is + // conservative, for example, when the first dim of the shape is -1 and input + // tensor shape is not fixed, it is still possible that the reshape doesn't + // change the batch dim, but as long as there is a possibility that it could + // change the batch dim, it reject the conversion. The parameters are: + // + // * reshape_batch_dim: the value of the first dim of the input shape constant + // * reshape_dims: all other dims of the input shape constant + // * input_batch_dim: the value of the first dim of the input tensor to + // reshape + // * input_dims: all other dims of the input tensor to reshape + // + // The validation logic is: + // + // if input_batch_dim is fixed: + // if reshape_batch_dim == input_batch_dim: + // ok + // elif reshape_batch_dim == -1 (meaning reshape_dims are fixed) and + // input_dims are fixed and + // prod(input_dims) == prod(reshape_dims) + // ok + // else: + // not ok + // elif input_dims are fixed: + // if reshape_dims are fixed and + // prod(input_dims) == prod(reshape_dims): + // ok + // else: + // not ok + // else: + // not ok + + const int input_batch_dim = input_tensor.batch_size(); + const int reshape_batch_dim = weights_ptr[0]; + const nvinfer1::Dims input_dims = input_tensor.GetTrtDims(); + + nvinfer1::Dims reshape_dims; + reshape_dims.nbDims = weights.count() - 1; + for (int i = 1; i < weights.count(); i++) { + reshape_dims.d[i - 1] = weights_ptr[i]; + } + + // Check that it doesn't change the batch dimension according to the logic + // mentioned above. + bool reshape_may_change_batch_dim = false; + if (input_batch_dim > 0) { // Batch size is fixed. + if (reshape_batch_dim == -1) { // Other dims of the shape must be fixed. + if (!HasStaticShape(input_dims) || + TrtDimsNumElements(reshape_dims) != TrtDimsNumElements(input_dims)) { + reshape_may_change_batch_dim = true; + } + } else if (reshape_batch_dim != input_batch_dim) { + reshape_may_change_batch_dim = true; + } + } else if (HasStaticShape(input_dims)) { + if (!HasStaticShape(reshape_dims) || + TrtDimsNumElements(reshape_dims) != TrtDimsNumElements(input_dims)) { + reshape_may_change_batch_dim = true; + } + } else { + reshape_may_change_batch_dim = true; + } + VLOG(1) << "input_batch_dim=" << input_batch_dim + << ", input_dims=" << DebugString(input_dims) + << "\nreshape_batch_dim=" << reshape_batch_dim + << ", reshape_dims=" << DebugString(reshape_dims); + if (reshape_may_change_batch_dim) { + const string msg = StrCat( + "Reshape on batch dimension is not supported, at ", node_def.name()); + return errors::Unimplemented(msg); + } + if (params->validation_only) return Status::OK(); + + // Start conversion. + const nvinfer1::ITensor* output_tensor = nullptr; + TF_RETURN_IF_ERROR(params->converter->PrepareTensorForShape( + input_tensor, reshape_dims, &output_tensor)); + params->outputs->push_back( + TRT_TensorOrWeights(const_cast(output_tensor))); + return tensorflow::Status::OK(); +} + +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()); + } + // Get input shape as vector. + TRT_TensorOrWeights input_tensor = inputs.at(0); + const nvinfer1::Dims dims = input_tensor.GetTrtDims(); + std::vector input_dims(dims.d, dims.d + dims.nbDims); + // Add batch dim back. + input_dims.insert(input_dims.begin(), -1); + const int input_rank = input_dims.size(); + // Get axis to expand on. + TRT_ShapedWeights weights = inputs.at(1).weights(); + if (weights.count() != 1) { + return tensorflow::errors::InvalidArgument( + "ExpandDims axis must be a scalar, at ", node_def.name()); + } + const int* weights_ptr = + static_cast(const_cast(weights.GetValues())); + int axis = weights_ptr[0]; + // Make sure axis is valid. + if ((axis < (-input_rank - 1)) || (axis > input_rank)) { + return tensorflow::errors::InvalidArgument( + "Axis for ExpandDims is invalid, must be in the range " + "[-rank(input) - 1, rank(input)], at ", + node_def.name()); + } + // Convert negative axis to corresponding positive axis. + if (axis < 0) axis += input_rank + 1; + if (axis == 0) { + return tensorflow::errors::Unimplemented( + "Modifying batch dimension is not supported for ExpandDims, at ", + node_def.name()); + } + if (params->validation_only) return Status::OK(); + + // ExpandDims: Insert new dim of size 1. + input_dims.insert(input_dims.begin() + axis, 1); + // Reshape tensor. + nvinfer1::Dims new_dims; + TF_RETURN_IF_ERROR(TensorShapeArrayToTrtDims(input_dims, &new_dims, + /*ignore_first_dim=*/true)); + const nvinfer1::ITensor* output_tensor = nullptr; + TF_RETURN_IF_ERROR(params->converter->PrepareTensorForShape( + input_tensor, new_dims, &output_tensor)); + params->outputs->push_back( + TRT_TensorOrWeights(const_cast(output_tensor))); + return tensorflow::Status::OK(); +} + +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()); + } + // Get input shape. + TRT_TensorOrWeights input_tensor = inputs.at(0); + const nvinfer1::Dims dims = input_tensor.GetTrtDims(); + std::vector input_dims(dims.d, dims.d + dims.nbDims); + // Add batch dim back. + input_dims.insert(input_dims.begin(), -1); + const int input_rank = input_dims.size(); + // Mark axes to remove by setting them to 0. + TFAttrs attrs(node_def); + auto squeeze_dims = attrs.get>("squeeze_dims"); + if (squeeze_dims.size() == 0) { + return tensorflow::errors::Unimplemented( + "Squeeze is only implemented for explicit dims, at ", node_def.name()); + } + for (int axis : squeeze_dims) { + // Make sure axis is valid. + if ((axis < -input_rank) || (axis >= input_rank)) { + return tensorflow::errors::InvalidArgument( + "Axis for Squeeze is invalid, must be in the range " + "[-rank(input), rank(input)), at ", + node_def.name()); + } + // Convert negative axis to corresponding positive axis. + if (axis < 0) axis += input_rank; + // Don't squeeze batch dim. + if (axis == 0) { + return tensorflow::errors::Unimplemented( + "Cannot squeeze batch dimension, at ", node_def.name()); + } + // Make sure target dimension is size 1. + if (input_dims[axis] != 1) { + return tensorflow::errors::InvalidArgument( + "Cannot squeeze a dimension which isn't size 1, at ", + node_def.name()); + } + // Mark dim for removal by setting to 0. + input_dims[axis] = 0; + } + if (params->validation_only) return Status::OK(); + + // Remove all dims which are equal to 0. + input_dims.erase(std::remove(input_dims.begin(), input_dims.end(), 0), + input_dims.end()); + // Reshape tensor. + nvinfer1::Dims new_dims; + TF_RETURN_IF_ERROR(TensorShapeArrayToTrtDims(input_dims, &new_dims, + /*ignore_first_dim=*/true)); + const nvinfer1::ITensor* output_tensor = nullptr; + TF_RETURN_IF_ERROR(params->converter->PrepareTensorForShape( + input_tensor, new_dims, &output_tensor)); + params->outputs->push_back( + TRT_TensorOrWeights(const_cast(output_tensor))); + return tensorflow::Status::OK(); +} + +// Gets the bounds (start or end) from the weights of a StridedSlice op. +tensorflow::Status GetStridedSliceBound(const std::vector& input_dims, + const TRT_ShapedWeights& bound_weights, + int mask, bool begin, string node_name, + std::vector* output_bound) { + const string bound_name = (begin) ? "begin" : "end"; + const int* weights_ptr = static_cast(bound_weights.GetValues()); + *output_bound = + std::vector(weights_ptr, weights_ptr + bound_weights.count()); + if (output_bound->size() != input_dims.size()) { + return tensorflow::errors::InvalidArgument( + "StridedSlice \"", bound_name, "\" specified ", + std::to_string(output_bound->size()), " dimensions, but input rank is ", + std::to_string(input_dims.size()), ", at ", node_name); + } + for (int i = 0; i < output_bound->size(); i++) { + if ((1 << i) & mask) { + // Apply mask. + (*output_bound)[i] = (begin) ? 0 : input_dims[i]; + // Masked bound will always result in a valid, non-negative bound, so we + // don't need the following checks. For the common case of using masks on + // a undefined batch dim (-1), we specifically don't want to do the + // following checks because they will erroneously detect an out of range + // bound or try to correct the negative value. + continue; + } + // Make sure bound is valid. + if (((*output_bound)[i] < -input_dims[i]) || + ((*output_bound)[i] > input_dims[i])) { + return tensorflow::errors::InvalidArgument( + bound_name, " value of ", std::to_string((*output_bound)[i]), + " for StridedSlice is invalid, must be in the range " + "[-dim_size(i), dim_size(i)], at ", + node_name); + } + // Convert negative values to their positive equivalent. + if ((*output_bound)[i] < 0) { + (*output_bound)[i] += input_dims[i]; + } + } + return tensorflow::Status::OK(); +} + +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()); + } + // Get input dims. + nvinfer1::Dims dims = inputs.at(0).GetTrtDims(); + std::vector input_dims(dims.d, dims.d + dims.nbDims); + if (inputs.at(0).is_tensor()) { + // Temporarily add batch dimension so that indexes line up properly. + input_dims.insert(input_dims.begin(), inputs.at(0).batch_size()); + } + if (input_dims.size() > 4) { + return tensorflow::errors::Unimplemented( + "StridedSlice is not implemented for tensors with rank > 4, at ", + node_def.name()); + } + TFAttrs attrs(node_def); + // Get begin and end bounds per axis. + std::vector begin, end; + TF_RETURN_IF_ERROR(GetStridedSliceBound(input_dims, inputs.at(1).weights(), + attrs.get("begin_mask"), true, + node_def.name(), &begin)); + TF_RETURN_IF_ERROR(GetStridedSliceBound(input_dims, inputs.at(2).weights(), + attrs.get("end_mask"), false, + node_def.name(), &end)); + // Get strides per axis (must all be 1). + TRT_ShapedWeights stride_weights = inputs.at(3).weights(); + const int* stride_weights_ptr = static_cast(stride_weights.GetValues()); + std::vector strides(stride_weights_ptr, + stride_weights_ptr + stride_weights.count()); + for (int x : strides) { + if (x != 1) { + return tensorflow::errors::Unimplemented( + "StridedSlice is only implemented for stride of 1, at ", + node_def.name()); + } + } + // Unsupported mask options. + for (const string& attr : + {"ellipsis_mask", "new_axis_mask", "shrink_axis_mask"}) { + int attr_val = attrs.get(attr); + if (attr_val != 0) { + return tensorflow::errors::Unimplemented( + attr, " is not supported for StridedSlice, at ", node_def.name()); + } + } + + nvinfer1::ITensor* tensor = + const_cast(inputs.at(0).tensor()); + // Reshape if necessary to 4-D, since IPaddingLayer requires a 4-D input. + const bool need_reshape = (input_dims.size() != 4); + int reshape_dims_added = 0; + nvinfer1::Dims reshape_dims; + if (need_reshape) { + // Add new dims after batch dim until tensor is 4D. + while (input_dims.size() < 4) { + input_dims.insert(input_dims.begin() + 1, 1); + begin.insert(begin.begin() + 1, 0); + end.insert(end.begin() + 1, 1); + reshape_dims_added++; + } + TF_RETURN_IF_ERROR(TensorShapeArrayToTrtDims(input_dims, &reshape_dims, + /*ignore_first_dim=*/true)); + } + // Find dimensions which need to be sliced. + std::vector pad_dims; + for (int i = 0; i < input_dims.size(); i++) { + if ((begin[i] != 0) || (end[i] != input_dims[i])) { + if (i == 0) { + return tensorflow::errors::Unimplemented( + "StridedSlice can't modify batch dim, at ", node_def.name()); + } else if ((end[i] - begin[i]) < 0) { + return tensorflow::errors::InvalidArgument( + "New size of sliced dimension is negative, at ", node_def.name()); + } + pad_dims.push_back(i); + } + } + if (pad_dims.size() == 0) { + // No dimensions are changed. We could create a padding layer anyway with + // values of 0. + if (params->validation_only) return Status::OK(); + params->outputs->push_back(inputs.at(0)); + return tensorflow::Status::OK(); + } else if (pad_dims.size() == 1) { + // Only one dim is modified but we have to have 2, mark a second dim which + // will have padding of 0. The dim we add is chosen to avoid an unecessary + // transpose. + if (pad_dims[0] != 2) { + pad_dims.push_back(2); + } else { + pad_dims.push_back(3); + } + } else if (pad_dims.size() > 2) { + return tensorflow::errors::Unimplemented( + "StridedSlice can only modify 2 dimensions, at ", node_def.name()); + } + std::sort(pad_dims.begin(), pad_dims.end()); + // Convert to pre/post padding values. Since TRT does not have a StridedSlice + // or Slice layer, we instead create an IPaddingLayer with negative padding. + nvinfer1::DimsHW pre_padding, post_padding; + for (int i = 0; i < pad_dims.size(); i++) { + const int axis = pad_dims[i]; + pre_padding.d[i] = -begin[axis]; + post_padding.d[i] = end[axis] - input_dims[axis]; + } + + // IPaddingLayer will always apply the padding to dims 2,3 (input format is + // NCHW). + const bool need_transpose = !(pad_dims[0] == 2 && pad_dims[1] == 3); + std::vector transpose_order(input_dims.size()); + std::vector inv_transpose_order(input_dims.size()); + if (need_transpose) { + if (pad_dims[0] == 1 && pad_dims[1] == 3) { + transpose_order = {0, 2, 1, 3}; + inv_transpose_order = {0, 2, 1, 3}; + } else if (pad_dims[0] == 1 && pad_dims[1] == 2) { + transpose_order = {0, 3, 1, 2}; + inv_transpose_order = {0, 2, 3, 1}; + } + } + if (params->validation_only) return Status::OK(); + + // Start conversion. + if (need_reshape) { + const nvinfer1::ITensor* output_tensor = nullptr; + TF_RETURN_IF_ERROR(params->converter->PrepareTensorForShape( + inputs.at(0), reshape_dims, &output_tensor)); + tensor = const_cast(output_tensor); + } + if (need_transpose) { + const nvinfer1::ITensor* output_tensor = nullptr; + TF_RETURN_IF_ERROR(params->converter->TransposeTensor( + tensor, transpose_order, &output_tensor)); + tensor = const_cast(output_tensor); + } + + // Add padding layer + nvinfer1::IPaddingLayer* layer = params->converter->network()->addPadding( + *const_cast(tensor), pre_padding, post_padding); + TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); + params->converter->MarkQuantizationRangesAsInferrable(tensor, + layer->getOutput(0)); + tensor = layer->getOutput(0); + + // Restore transpose + if (need_transpose) { + const nvinfer1::ITensor* output_tensor = nullptr; + TF_RETURN_IF_ERROR(params->converter->TransposeTensor( + tensor, inv_transpose_order, &output_tensor)); + tensor = const_cast(output_tensor); + } + // Restore reshape + if (need_reshape) { + // Calculate output dimensions + for (int i = 0; i < pad_dims.size(); i++) { + const int axis = pad_dims[i]; + input_dims[axis] = end[axis] - begin[axis]; + } + // Remove added 1 dimensions + for (int i = 0; i < reshape_dims_added; i++) { + int value = input_dims[1]; + if (value != 1) { + return tensorflow::errors::Internal( + "StridedSlice error when reshaping, at ", node_def.name()); + } + input_dims.erase(input_dims.begin() + 1); + } + + nvinfer1::Dims new_dims; + TF_RETURN_IF_ERROR(TensorShapeArrayToTrtDims(input_dims, &new_dims, + /*ignore_first_dim=*/true)); + const nvinfer1::ITensor* output_tensor = nullptr; + TF_RETURN_IF_ERROR(params->converter->PrepareTensorForShape( + TRT_TensorOrWeights(tensor), new_dims, &output_tensor)); + tensor = const_cast(output_tensor); + } + + params->outputs->push_back( + TRT_TensorOrWeights(const_cast(tensor))); + return tensorflow::Status::OK(); +} + +tensorflow::Status ConvertConv2D(OpConverterParams* params) { + return ConvertConv2DHelper(params, ConvolutionType::DEFAULT); +} + +tensorflow::Status ConvertConv2DDepthwise(OpConverterParams* params) { + return ConvertConv2DHelper(params, ConvolutionType::DEPTHWISE_CONV); +} + +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()); + } + nvinfer1::PoolingType type; + if (node_def.op() == "MaxPool") { + type = nvinfer1::PoolingType::kMAX; + } else if (node_def.op() == "AvgPool") { + type = nvinfer1::PoolingType::kAVERAGE; + } else { + return tensorflow::errors::Unimplemented( + "Unsupported pooling type: ", node_def.op(), ", at ", node_def.name()); + } + TFAttrs attrs(node_def); + const string padding_type = attrs.get("padding"); + if ((padding_type != "SAME") && (padding_type != "VALID")) { + return tensorflow::errors::Unimplemented( + "Unsupported padding type: ", padding_type, ", at ", node_def.name()); + } + if (params->validation_only) return Status::OK(); + + const nvinfer1::ITensor* tensor = inputs.at(0).tensor(); + int h_index = 2; + int w_index = 3; + const auto data_format = attrs.get("data_format"); + if (data_format == "NHWC") { + h_index = 1; + w_index = 2; + TF_RETURN_IF_ERROR(params->converter->TransposeTensor( + const_cast(tensor), {0, 3, 1, 2}, &tensor)); + } + + const auto tf_stride = attrs.get>("strides"); + const nvinfer1::DimsHW stride(tf_stride[h_index], tf_stride[w_index]); + + const auto tf_kernel = attrs.get>("ksize"); + const nvinfer1::DimsHW ksize(tf_kernel[h_index], tf_kernel[w_index]); + + auto tensor_dim = tensor->getDimensions(); + std::vector> padding; + if (padding_type == "SAME") { + // This is NCHW tensor with no batch dimension. + // 1 -> h + // 2 -> w + padding = CreateSamePadding( + stride, ksize, + {static_cast(tensor_dim.d[1]), static_cast(tensor_dim.d[2])}); + } else if (padding_type == "VALID") { + padding = {{0, 0}, {0, 0}}; + } + + if (padding[0].first != padding[0].second || + padding[1].first != padding[1].second) { + VLOG(2) << "Padding!!!: " << padding[0].first << padding[0].second + << padding[1].first << padding[1].second; + auto pad_layer = params->converter->network()->addPadding( + *const_cast(tensor), + nvinfer1::DimsHW(padding[0].first, padding[1].first), + nvinfer1::DimsHW(padding[0].second, padding[1].second)); + TFTRT_RETURN_ERROR_IF_NULLPTR(pad_layer, node_def.name()); + params->converter->MarkQuantizationRangesAsInferrable( + const_cast(tensor), pad_layer->getOutput(0)); + padding = {{0, 0}, {0, 0}}; + tensor = pad_layer->getOutput(0); + } + + nvinfer1::IPoolingLayer* layer = params->converter->network()->addPooling( + *const_cast(tensor), type, ksize); + TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); + // TODO(tmorris): Average pooling may not be entirely safe to infer + // quantization range through (at least forwards - backwards should be fine). + // Max pooling is okay. + params->converter->MarkQuantizationRangesAsInferrable( + const_cast(tensor), layer->getOutput(0)); + + layer->setStride(stride); + layer->setPadding({padding[0].first, padding[1].first}); + layer->setName(node_def.name().c_str()); + const nvinfer1::ITensor* output_tensor = layer->getOutput(0); + + if (data_format == "NHWC") { + TF_RETURN_IF_ERROR(params->converter->TransposeTensor( + const_cast(output_tensor), {0, 2, 3, 1}, + &output_tensor)); + } + params->outputs->push_back( + TRT_TensorOrWeights(const_cast(output_tensor))); + return tensorflow::Status::OK(); +} + +tensorflow::Status ConvertActivation(OpConverterParams* params) { + const auto& inputs = params->inputs; + const auto& node_def = params->node_def; + if (inputs.size() != 1) { + return tensorflow::errors::InvalidArgument( + node_def.op(), " expects one input, at ", node_def.name()); + } + if (!inputs.at(0).is_tensor()) { + return tensorflow::errors::Unimplemented( + node_def.op(), " is only implemented for tensors, at ", + node_def.name()); + } + static const std::unordered_map ops{ + {"Relu", nvinfer1::ActivationType::kRELU}, + {"Sigmoid", nvinfer1::ActivationType::kSIGMOID}, + {"Tanh", nvinfer1::ActivationType::kTANH}, + }; + auto op_pair = ops.find(node_def.op()); + if (op_pair == ops.end()) { + return tensorflow::errors::Unimplemented( + "Activation op: ", node_def.op(), + " not supported at: ", node_def.name()); + } + if (params->validation_only) return tensorflow::Status::OK(); + + // Start conversion. + const nvinfer1::ITensor* tensor = inputs.at(0).tensor(); + nvinfer1::IActivationLayer* layer = + params->converter->network()->addActivation( + *const_cast(tensor), op_pair->second); + TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); + nvinfer1::ITensor* output_tensor = layer->getOutput(0); + // Set quantization range for output of Sigmoid, Tanh. + if (node_def.op() == "Sigmoid") { + params->converter->ProvideQuantizationRange(output_tensor, 0.0f, 1.0f); + } else if (node_def.op() == "Tanh") { + params->converter->ProvideQuantizationRange(output_tensor, -1.0f, 1.0f); + } + params->outputs->push_back(TRT_TensorOrWeights(output_tensor)); + return tensorflow::Status::OK(); +} + +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(); + } + float min_range = 0.0f; + float max_range = 0.0f; + if (node_def.op() == "FakeQuantWithMinMaxArgs") { + // Get ranges via node attributes. + TFAttrs attrs(node_def); + if (attrs.count("min") == 0 || attrs.count("max") == 0) { + return errors::InvalidArgument("Min or max attribute not found for ", + node_def.op(), " at ", node_def.name()); + } + min_range = attrs.get("min"); + max_range = attrs.get("max"); + } else if (node_def.op() == "FakeQuantWithMinMaxVars" || + 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())); + return raw_weights[0]; + }; + min_range = get_weights_value(1); + max_range = get_weights_value(2); + } else { + return errors::InvalidArgument("Unknown quantization op ", node_def.op(), + ", at ", node_def.name()); + } + if (params->validation_only) return Status::OK(); + + // Store ranges for tensor + params->converter->ProvideQuantizationRange( + const_cast(inputs.at(0).tensor()), min_range, + max_range); + // Sometimes, TRT may not quantize a tensor, either because it chooses to + // execute a higher precision kernel or because of op fusion. In these cases, + // accuracy will suffer if the model was trained to expect quantization at + // that tensor. We should consider adding a clip(tensor, min_range, max_range) + // operation here to ensure that any arbitrarily placed quantize node will + // execute as expected. However, this will negatively affect performance. If + // users train their models in a way which models inference as close as + // possible (i.e. not quantizing in place where fusion will occur), then there + // is no problem with the current implementation. + params->outputs->push_back(inputs.at(0)); + return Status::OK(); +} + +// TODO(pdavoodi): we should update relu6 implementation once TensorRT supports +// Relu6 natively. +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()); + } + if (params->validation_only) return Status::OK(); + // *************************************************************************** + // TensorRT does not implement Relu6 natively. This function converts Relu6 op + // to available TensorRT ops: Relu6(x) = min(Relu(x), 6) + // *************************************************************************** + + // Input Tensor + const nvinfer1::ITensor* tensor = inputs.at(0).tensor(); + + // Relu operation i.e. Relu(x) = max(0, x) + nvinfer1::IActivationLayer* relu_layer = + params->converter->network()->addActivation( + *const_cast(tensor), + nvinfer1::ActivationType::kRELU); + TFTRT_RETURN_ERROR_IF_NULLPTR(relu_layer, node_def.name()); + + // Large range of relu is problematic during quantization in INT8 precision + // mode. Setting dynamic range of relu = [0.f, 6.0f] helps with quantization. + // TRT only uses dynamic ranges in INT8 precision mode, + // and this does not affect the FP32 path. + 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::ITensor* const6_tensor = + params->converter->CreateConstantLayer(weights, dims); + TFTRT_RETURN_ERROR_IF_NULLPTR(const6_tensor, node_def.name()); + params->converter->ProvideQuantizationRange(const6_tensor, 0.0f, 6.0f); + + // 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]. + nvinfer1::IElementWiseLayer* relu6_layer = + params->converter->network()->addElementWise( + *const_cast(relu_layer->getOutput(0)), + *const6_tensor, nvinfer1::ElementWiseOperation::kMIN); + TFTRT_RETURN_ERROR_IF_NULLPTR(relu6_layer, node_def.name()); + nvinfer1::ITensor* output_tensor = relu6_layer->getOutput(0); + params->converter->ProvideQuantizationRange(output_tensor, 0.0f, 6.0f); + + params->outputs->push_back(TRT_TensorOrWeights(output_tensor)); + return Status::OK(); +} + +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()); + } + 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(); + const string data_format = attrs.get("data_format"); + const int channel_index = + (data_format == "NHWC" ? original_dims.nbDims - 1 : 0); + + nvinfer1::Permutation permutation; + if (channel_index != 0) { + // Permute the dimensions so that the channel dimension is the first + // dimension. + for (int i = 0; i < original_dims.nbDims; ++i) { + permutation.order[i] = i; + } + permutation.order[0] = channel_index; + permutation.order[channel_index] = 0; + VLOG(1) << "ConvertBiasAdd permutation: " + << DebugString(permutation, original_dims.nbDims); + } + + // TensorRT addScale requires input to be of rank 3, we need to apply + // transpose as well as reshape. + // TODO(laigd): this doesn't match what the TRT doc says, fix the doc? + if (channel_index != 0 || original_dims.nbDims != 3) { + nvinfer1::IShuffleLayer* shuffle_layer = + params->converter->network()->addShuffle(*tensor); + TFTRT_RETURN_ERROR_IF_NULLPTR(shuffle_layer, node_def.name()); + params->converter->MarkQuantizationRangesAsInferrable( + tensor, shuffle_layer->getOutput(0)); + + // NOTE(laigd): for some reason we need to apply the reshape + // unconditionally. The default shape has nbDims==-1 and it seems the + // behavior is undefined in some cases. + nvinfer1::Dims reshape_dims; + reshape_dims.nbDims = 3; + // 0 means copying from input; -1 means inferring from the rest. + reshape_dims.d[0] = 0; + reshape_dims.d[1] = original_dims.nbDims >= 2 ? 0 : 1; + reshape_dims.d[2] = original_dims.nbDims >= 3 ? -1 : 1; + shuffle_layer->setReshapeDimensions(reshape_dims); + + if (channel_index != 0) { + shuffle_layer->setFirstTranspose(permutation); + } + tensor = shuffle_layer->getOutput(0); + } + + TRT_ShapedWeights weights = inputs.at(1).weights(); + if (params->converter->precision_mode() == FP16MODE) { + weights = ConvertFP32ToFP16(params->weight_store, weights); + } + nvinfer1::ScaleMode mode = nvinfer1::ScaleMode::kCHANNEL; + if (weights.shape_.d[0] == 1) { + mode = nvinfer1::ScaleMode::kUNIFORM; + } + + TRT_ShapedWeights empty_weights(weights.type_); + nvinfer1::IScaleLayer* layer = params->converter->network()->addScale( + *tensor, mode, weights.GetTrtWeights(), empty_weights.GetTrtWeights(), + empty_weights.GetTrtWeights()); + TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); + + nvinfer1::ITensor* output_tensor = layer->getOutput(0); + + // Restore transpose & reshape. + if (channel_index != 0 || original_dims.nbDims != 3) { + nvinfer1::IShuffleLayer* shuffle_layer = + params->converter->network()->addShuffle(*output_tensor); + TFTRT_RETURN_ERROR_IF_NULLPTR(shuffle_layer, node_def.name()); + // NOTE: for same reason as mentioned above we need to apply the reshape + // unconditionally. + nvinfer1::Dims reshape_dims = original_dims; + if (channel_index != 0) { + // NOTE: according to NVIDIA dimension types are deprecated, so we don't + // need to copy them back. + reshape_dims.d[channel_index] = original_dims.d[0]; + reshape_dims.d[0] = original_dims.d[channel_index]; + } + shuffle_layer->setReshapeDimensions(reshape_dims); + + if (channel_index != 0) { + shuffle_layer->setSecondTranspose(permutation); + } + params->converter->MarkQuantizationRangesAsInferrable( + output_tensor, shuffle_layer->getOutput(0)); + output_tensor = shuffle_layer->getOutput(0); + } + + params->outputs->push_back(TRT_TensorOrWeights(output_tensor)); + return Status::OK(); +} + +void GetTensorDimsWithProtoShape(const Tensor& tensor, nvinfer1::Dims* dims) { + if (tensor.dims() > 0) { + *dims = GetTrtDimsForTensor(tensor); + } else { + dims->nbDims = 1; + // No dimension provided. Flatten it. + 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; + } + } +} + +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; + 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 { + // 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(); +} + +// Convert a Const NodeDef to TRT_ShapedWeights. This is a special converter, it +// always ignores the params->validation_only parameter but adds the converted +// weights to params->outputs. We did this since TrtNodeValidator needs the +// weights as input to other nodes, and use it to determine whether those nodes +// are supported by TRT. +tensorflow::Status ConvertConst(OpConverterParams* params) { + const auto& inputs = params->inputs; + const auto& node_def = params->node_def; + if (!inputs.empty()) { + return errors::InvalidArgument( + "Constant node is expected to have empty input list: ", + node_def.name()); + } + + // Create shaped weights as output + const auto& tensor_proto = node_def.attr().at("value").tensor(); + tensorflow::Tensor tensor; + if (!tensor.FromProto(tensor_proto)) { + return tensorflow::errors::Internal("Cannot parse weight tensor proto: ", + 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)); + } + return Status::OK(); +} + +tensorflow::Status ConvertIdentity(OpConverterParams* params) { + // TODO(tmorris): TRT's Identity layer does not get optimized away as of TRT + // 5.0, however once we know that it does it would be nice to use that + // instead. + params->outputs->push_back(params->inputs.at(0)); + return tensorflow::Status::OK(); +} + +Status ConvertBinary(OpConverterParams* params) { + const auto& inputs = params->inputs; + const auto& node_def = params->node_def; + if (inputs.size() != 2) { + return errors::InvalidArgument("Binary ops require two inputs, at ", + node_def.name()); + } + + // Constant folding should have been done by TensorFlow + if (inputs.at(0).is_weights() && inputs.at(1).is_weights()) { + return errors::Unimplemented( + "Constant folding is falled back to TensorFlow, binary op received " + "both input as constant at: ", + node_def.name()); + } + + // TODO(tmorris): TRT plans to deprecate IScaleLayer and will replace it with + // IElementwiseLayer. At that point, we can remove BinaryTensorOpWeight. For + // now, the performance will be slightly better with IScaleLayer because it + // can be fused in more situations. However, most of the benefits of + // IScaleLayer are when the layer performs both a shift and a scale, which we + // don't do except for convolutions. + // + // Try to convert into Scale layer first (for better performance). + // Since scale layer supports restricted broadcast policy and op types, we + // allow failure and try to handle it through Elementwise op + // (BinaryTensorOpTensor). + Status status = Status::OK(); + if (inputs.at(0).is_tensor() && inputs.at(1).is_weights()) { + status = BinaryTensorOpWeight(params, inputs.at(0).tensor(), + inputs.at(1).weights(), false); + } else if (inputs.at(0).is_weights() && inputs.at(1).is_tensor()) { + status = BinaryTensorOpWeight(params, inputs.at(1).tensor(), + inputs.at(0).weights(), true); + } + // If both input are tensors, or one of them is weights but the conversion + // above failed, try the conversion using BinaryTensorOpTensor. + if ((inputs.at(0).is_tensor() && inputs.at(1).is_tensor()) || !status.ok()) { + if (!status.ok()) VLOG(1) << status; + status = BinaryTensorOpTensor(params, inputs.at(0), inputs.at(1)); + } + return status; +} + +tensorflow::Status ConvertUnary(OpConverterParams* params) { + const auto& inputs = params->inputs; + const auto& node_def = params->node_def; + static const std::unordered_map ops{ + {"Neg", nvinfer1::UnaryOperation::kNEG}, + {"Exp", nvinfer1::UnaryOperation::kEXP}, + {"Log", nvinfer1::UnaryOperation::kLOG}, + {"Sqrt", nvinfer1::UnaryOperation::kSQRT}, + {"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()); + } + + // TODO(jie): check type + const nvinfer1::ITensor* tensor = nullptr; + TF_RETURN_IF_ERROR(params->converter->PrepareTensorForShape( + inputs.at(0), inputs.at(0).GetTrtDims(), &tensor)); + + nvinfer1::IUnaryLayer* layer; + if (node_def.op() == "Rsqrt") { + // We will need a quantization range for intermediate tensor if not using + // calibration. + // + // x -> [Sqrt] -> sqrt(x) -> [Recip] -> 1/sqrt(x) + // ^ + // need range here + if (params->converter->precision_mode() == INT8MODE && + !params->converter->use_calibration()) { + return errors::Unimplemented( + "Intermediate quantization range cannot be determined without" + " calibration for Rsqrt, consider replacing with " + "Sqrt -> FakeQuant -> Reciprocal ops, at ", + node_def.name()); + } + layer = params->converter->network()->addUnary( + *const_cast(tensor), + nvinfer1::UnaryOperation::kSQRT); + TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); + tensor = layer->getOutput(0); + layer = params->converter->network()->addUnary( + *const_cast(tensor), + nvinfer1::UnaryOperation::kRECIP); + } else if (ops.count(node_def.op()) != 0) { + layer = params->converter->network()->addUnary( + *const_cast(tensor), ops.at(node_def.op())); + } else { + return tensorflow::errors::InvalidArgument( + "Binary op: ", node_def.op(), " not supported, at ", node_def.name()); + } + + TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); + nvinfer1::ITensor* output_tensor = layer->getOutput(0); + params->outputs->push_back( + TRT_TensorOrWeights(const_cast(output_tensor))); + return tensorflow::Status::OK(); +} + +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()); + } + if (params->validation_only) return Status::OK(); + + // Constant 2 with same rank as input + nvinfer1::Dims dims = inputs.at(0).GetTrtDims(); + 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] = 2.f; + nvinfer1::ITensor* const2_tensor = + params->converter->CreateConstantLayer(weights, dims); + TFTRT_RETURN_ERROR_IF_NULLPTR(const2_tensor, node_def.name()); + + // ElementWise Pow Operation + nvinfer1::IElementWiseLayer* layer = + params->converter->network()->addElementWise( + *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); + + params->outputs->push_back(TRT_TensorOrWeights(output_tensor)); + return tensorflow::Status::OK(); +} + +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()); + } + + const nvinfer1::ITensor* tensor = inputs.at(0).tensor(); + TRT_ShapedWeights index_list = inputs.at(1).weights(); + + TFAttrs attrs(node_def); + auto index_type = attrs.get("Tidx"); + + // Only expect to handle INT32 as attributes for now + if (index_type != tensorflow::DataType::DT_INT32) { + return tensorflow::errors::Unimplemented("Tidx supports only DT_INT32"); + } + + int axes = 0; + if (index_list.count() == 0) { + return tensorflow::errors::InvalidArgument( + "TRT cannot support reduce on all (batch) dimensions, at", + node_def.name()); + } else { + auto index_list_data = + static_cast(const_cast(index_list.GetValues())); + for (int i = 0; i < index_list.count(); i++) { + int axis = index_list_data[i]; + if (axis < 0) axis += tensor->getDimensions().nbDims + 1; + if (axis == 0) { + return tensorflow::errors::InvalidArgument( + "TRT cannot reduce at batch dimension, at", node_def.name()); + } + axes |= (1 << (axis - 1)); + } + } + + nvinfer1::ReduceOperation reduce_operation; + if (node_def.op() == "Sum") { + reduce_operation = nvinfer1::ReduceOperation::kSUM; + } else if (node_def.op() == "Prod") { + reduce_operation = nvinfer1::ReduceOperation::kPROD; + } else if (node_def.op() == "Max") { + reduce_operation = nvinfer1::ReduceOperation::kMAX; + } else if (node_def.op() == "Min") { + reduce_operation = nvinfer1::ReduceOperation::kMIN; + } else if (node_def.op() == "Mean") { + reduce_operation = nvinfer1::ReduceOperation::kAVG; + } else { + return tensorflow::errors::Unimplemented("Op not supported ", node_def.op(), + " , at ", node_def.name()); + } + + const auto keep_dims = attrs.get("keep_dims"); + nvinfer1::ILayer* layer = params->converter->network()->addReduce( + *const_cast(tensor), reduce_operation, axes, + keep_dims); + TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); + + params->outputs->push_back(TRT_TensorOrWeights(layer->getOutput(0))); + return tensorflow::Status::OK(); +} + +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()); + } + + // Implement tensor binaryOp weight [channel wise] for now; + const nvinfer1::ITensor* tensor = inputs.at(0).tensor(); + const auto dims = tensor->getDimensions(); + // Restore implicit batch dimension + const int nb_dims = dims.nbDims + 1; + + TRT_ShapedWeights pads = inputs.at(1).weights(); + + TFAttrs attrs(node_def); + // Padding type here is done through TF type + // so I can leverage their EnumToDataType for my cast + auto padding_type = attrs.get("Tpaddings"); + // TODO(jie): handle data type conversion for TRT? + + if (pads.shape_.d[0] != nb_dims || pads.shape_.d[1] != 2) { + return tensorflow::errors::InvalidArgument( + "Pad only supports explicit padding on 4 dimensional tensor, at ", + node_def.name()); + } + + // Only expect to handle INT32 as attributes for now + if (padding_type != tensorflow::DataType::DT_INT32) { + return tensorflow::errors::Unimplemented( + "Tpaddings supports only DT_INT32"); + } + auto pad_data = static_cast(const_cast(pads.GetValues())); + + std::vector pad_index; + for (int i = 0; i < nb_dims; i++) { + if (pad_data[2 * i] != 0 || pad_data[2 * i + 1] != 0) { + pad_index.push_back(i); + } + } + + // No padding at all, we should exit + if (pad_index.size() == 0) { + params->outputs->push_back(inputs.at(0)); + return tensorflow::Status::OK(); + } + + // Only supports padding on less than 2 axis GIE-2579 + if (pad_index.size() > 2) { + return tensorflow::errors::InvalidArgument( + "Padding layer does not support padding on > 2"); + } + + // Padding on batch dimension is not supported + if (pad_index[0] == 0) { + return tensorflow::errors::InvalidArgument( + "Padding layer does not support padding on batch dimension"); + } + + // Not doing the legit thing here. ignoring padding on dim 1 and 3; + // TODO(jie): implement pad as uff parser + if (pad_index.size() == 2 && pad_index[0] == 0 && pad_index[1] == 3) { + return tensorflow::errors::Unimplemented( + "Padding layer does not support padding on dimension 1 and 3 yet"); + } + if (params->validation_only) return Status::OK(); + + bool legit_pad = true; + nvinfer1::DimsHW pre_padding(0, 0); + nvinfer1::DimsHW post_padding(0, 0); + + std::vector permuted_pad_index(pad_index); + if (pad_index[0] == 1) { + legit_pad = false; + TF_RETURN_IF_ERROR(params->converter->TransposeTensor( + const_cast(tensor), {0, 3, 2, 1}, &tensor)); + permuted_pad_index[0] = 3; + } + + for (size_t i = 0; i < pad_index.size(); i++) { + int index = pad_index[i]; + if (permuted_pad_index[i] == 2) { + pre_padding.h() = pad_data[index * 2]; + post_padding.h() = pad_data[index * 2 + 1]; + } else if (permuted_pad_index[i] == 3) { + pre_padding.w() = pad_data[index * 2]; + post_padding.w() = pad_data[index * 2 + 1]; + } + } + + nvinfer1::IPaddingLayer* layer = params->converter->network()->addPadding( + *const_cast(tensor), pre_padding, post_padding); + TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); + const nvinfer1::ITensor* output_tensor = layer->getOutput(0); + + if (!legit_pad) { + TF_RETURN_IF_ERROR(params->converter->TransposeTensor( + const_cast(output_tensor), {0, 3, 2, 1}, + &output_tensor)); + } + + params->outputs->push_back( + TRT_TensorOrWeights(const_cast(output_tensor))); + return tensorflow::Status::OK(); +} + +tensorflow::Status ConvertConcat(OpConverterParams* params) { + const auto& inputs = params->inputs; + const auto& node_def = params->node_def; + // not including the last input (axis) here + int input_size = static_cast(inputs.size()) - 1; + + if (!inputs.at(0).is_tensor()) { + return tensorflow::errors::InvalidArgument( + "Concat in TRT support only Tensor input, at ", node_def.name()); + } + + // We are retrieving the axis + TRT_ShapedWeights axis = inputs.at(input_size).weights(); + + TFAttrs attrs(node_def); + auto index_type = attrs.get("Tidx"); + + // TODO(jie): handle data type + // Only expect to handle INT32 as index attributes for now + if (index_type != tensorflow::DataType::DT_INT32) + return tensorflow::errors::Unimplemented("Tidx supports only DT_INT32, at ", + node_def.name()); + + int index = *(static_cast(const_cast(axis.GetValues()))); + + // TODO(jie): early termination with no-op (attr_size==1) + + auto dim = inputs.at(0).tensor()->getDimensions(); + // dimension check + if (index > dim.nbDims + 1) { + return tensorflow::errors::InvalidArgument( + "Concatenate on axis out of dimension range, at ", node_def.name()); + } + if (index == 0) { + return tensorflow::errors::InvalidArgument( + "Concatenate on batch dimension not supported, at ", node_def.name()); + } + if (index < 0) { + index = dim.nbDims + index + 1; + } + + std::vector inputs_vec; + // Shap chack (all input tensor should have same shape) + // starting from 0 since we are probably also doing transpose here; + for (int i = 0; i < input_size; i++) { + auto tensor_i = inputs.at(i).tensor(); + auto dim_i = tensor_i->getDimensions(); + if (dim_i.nbDims != dim.nbDims) { + return tensorflow::errors::InvalidArgument( + "Concatenate receives inputs with inconsistent dimensions, at ", + node_def.name()); + } + for (int j = 0; j < dim.nbDims; j++) { + // check dimension consistency on non-concatenate axis + if (j != index - 1 && dim_i.d[j] != dim.d[j]) { + return tensorflow::errors::InvalidArgument( + "Concatenate receives inputs with inconsistent shape, at", + node_def.name()); + } + } + + inputs_vec.push_back(tensor_i); + } + if (params->validation_only) return tensorflow::Status::OK(); + + // nvinfer1::ITensor const* tensor = inputs.at(0).tensor(); + nvinfer1::IConcatenationLayer* layer = + params->converter->network()->addConcatenation( + const_cast(inputs_vec.data()), + inputs_vec.size()); + TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); + layer->setAxis(index - 1); + nvinfer1::ITensor* output_tensor = layer->getOutput(0); + params->outputs->push_back(TRT_TensorOrWeights(output_tensor)); + return tensorflow::Status::OK(); +} + +tensorflow::Status ConvertFusedBatchNorm(OpConverterParams* params) { + const auto& inputs = params->inputs; + const auto& node_def = params->node_def; + TFAttrs attrs(node_def); + float epsilon = attrs.get("epsilon"); + auto data_format = attrs.get("data_format"); + if (data_format != "NCHW") { + return tensorflow::errors::Unimplemented( + node_def.op(), " only supports data_format=NCHW, at ", node_def.name()); + } + bool is_training = attrs.get("is_training"); + if (is_training) { + // Trying to use batchnorm in training mode is a very common problem. + // Because the error message will only be printed in VLOG(1) by the + // segmenter, we issue a special warning so that users will actually see it. + LOG(WARNING) << node_def.op() << " only supports is_training=false. If you " + << "are using Keras, please call " + << "keras.backend.set_learning_phase(0) before constructing " + << "your model. At " << node_def.name(); + return tensorflow::errors::Unimplemented( + node_def.op(), " only supports is_training=false, at ", + node_def.name()); + } + 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 + auto parameter_type = inputs.at(1).weights().type_; + if ((parameter_type != tensorflow::DataType::DT_FLOAT) && + (parameter_type != tensorflow::DataType::DT_HALF)) { + return tensorflow::errors::Unimplemented( + "only float32 or float16 weight data type is supported, for node " + + node_def.name() + " got " + tensorflow::DataTypeString(parameter_type)); + } + for (int i = 1; i < 5; i++) { + if (inputs.at(i).weights().type_ != parameter_type) { + return tensorflow::errors::Unimplemented( + "Inconsistent parameter type for batchnorm is not supported, at: " + + node_def.name()); + } + } + + 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()); + } + TRT_ShapedWeights* ptr_shape_weights = nullptr; + for (int i = 1; i < 5; i++) { + if (inputs.at(i).weights().count() == nweight) { + ptr_shape_weights = + const_cast(&(inputs.at(i).weights())); + } else if (inputs.at(i).weights().count() != 1) { + return tensorflow::errors::InvalidArgument( + "Inconsistent batchnorm parameter count, at: " + node_def.name()); + } + } + if (params->validation_only) return Status::OK(); + + // We could technically have two weights with different shape. + // that requires two addScale op, arguably less performant + TRT_ShapedWeights combined_scale_weights = + params->weight_store->GetTempWeights(*ptr_shape_weights); + TRT_ShapedWeights combined_offset_weights = + params->weight_store->GetTempWeights(*ptr_shape_weights); + + const Eigen::half* cast_vals_array[4]; + const float* vals_array[4]; + for (int j = 0; j < 4; j++) { + cast_vals_array[j] = + static_cast(inputs.at(j + 1).weights().GetValues()); + vals_array[j] = + static_cast(inputs.at(j + 1).weights().GetValues()); + } + Eigen::half* cast_combined_scale_vals = const_cast( + static_cast(combined_scale_weights.GetValues())); + Eigen::half* cast_combined_offset_vals = const_cast( + static_cast(combined_offset_weights.GetValues())); + float* combined_scale_vals = const_cast( + static_cast(combined_scale_weights.GetValues())); + float* combined_offset_vals = const_cast( + static_cast(combined_offset_weights.GetValues())); + + for (size_t i = 0; i < nweight; ++i) { + float batchnorm_data[4]; + for (int j = 0; j < 4; j++) { + if (inputs.at(j + 1).weights().count() != 1) { + if (parameter_type == tensorflow::DT_FLOAT) { + batchnorm_data[j] = vals_array[j][i]; + } else if (parameter_type == tensorflow::DT_HALF) { + batchnorm_data[j] = + Eigen::half_impl::half_to_float(cast_vals_array[j][i]); + } + } else { + if (parameter_type == tensorflow::DT_FLOAT) { + batchnorm_data[j] = vals_array[j][0]; + } else if (parameter_type == tensorflow::DT_HALF) { + batchnorm_data[j] = + Eigen::half_impl::half_to_float(cast_vals_array[j][0]); + } + } + } + float scale = batchnorm_data[0]; + float offset = batchnorm_data[1]; + float mean = batchnorm_data[2]; + float variance = batchnorm_data[3]; + float combined_scale_val = scale / sqrtf(variance + epsilon); + float combined_offset_val = offset - mean * combined_scale_val; + if (parameter_type == tensorflow::DT_FLOAT) { + combined_scale_vals[i] = combined_scale_val; + combined_offset_vals[i] = combined_offset_val; + } else if (parameter_type == tensorflow::DT_HALF) { + cast_combined_scale_vals[i] = Eigen::half(combined_scale_val); + cast_combined_offset_vals[i] = Eigen::half(combined_offset_val); + } + } + + nvinfer1::ScaleMode mode = nweight == 1 ? nvinfer1::ScaleMode::kUNIFORM + : nvinfer1::ScaleMode::kCHANNEL; + nvinfer1::IScaleLayer* layer = params->converter->network()->addScale( + *const_cast(tensor), mode, + combined_offset_weights.GetTrtWeights(), + combined_scale_weights.GetTrtWeights(), + dummy_power_weights.GetTrtWeights()); + TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); + nvinfer1::ITensor* output_tensor = layer->getOutput(0); + params->outputs->push_back(TRT_TensorOrWeights(output_tensor)); + return tensorflow::Status::OK(); +} + +tensorflow::Status ConvertMatMulHelper(OpConverterParams* params, + TRT_TensorOrWeights tensor_input, + TRT_ShapedWeights weights_raw, + bool transpose_weight, + string node_name) { + nvinfer1::ITensor* output_tensor; + if (!tensor_input.is_tensor()) { + return tensorflow::errors::InvalidArgument("Input 0 expects tensor"); + } + const nvinfer1::ITensor* tensor = tensor_input.tensor(); + + TRT_ShapedWeights weights(weights_raw.type_); + if (transpose_weight) { + weights = weights_raw; + } else { + weights = params->weight_store->GetTempWeights(weights_raw); + ReorderCKtoKC(weights_raw, &weights); + } + TRT_ShapedWeights biases(weights.type_); + + int noutput = weights.shape_.d[0]; + + auto input_dim = tensor->getDimensions(); + while (input_dim.nbDims != 3) { + input_dim.d[input_dim.nbDims++] = 1; + } + TF_RETURN_IF_ERROR(params->converter->PrepareTensorForShape( + tensor_input, input_dim, &tensor)); + + nvinfer1::IFullyConnectedLayer* layer = + params->converter->network()->addFullyConnected( + *const_cast(tensor), noutput, + weights.GetTrtWeights(), biases.GetTrtWeights()); + TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_name); + output_tensor = layer->getOutput(0); + + const nvinfer1::ITensor* temp_tensor = nullptr; + auto output_dim = output_tensor->getDimensions(); + output_dim.nbDims = 1; + TF_RETURN_IF_ERROR(params->converter->PrepareTensorForShape( + TRT_TensorOrWeights(output_tensor), output_dim, &temp_tensor)); + output_tensor = const_cast(temp_tensor); + params->outputs->push_back(TRT_TensorOrWeights(output_tensor)); + return tensorflow::Status::OK(); +} + +// inputs are both two dimensional (tensorflow::ops::MatMul) +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()); + } + + 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)); + } + bool transpose_a = attrs.get("transpose_a"); + bool transpose_b = attrs.get("transpose_b"); + + // FullyConnected: + if (transpose_a) { + return errors::InvalidArgument( + "transpose_a is not supported for TensorRT FullyConnected (op: ", + node_def.op(), "), at: ", node_def.name()); + } + if (params->validation_only) return Status::OK(); + return ConvertMatMulHelper(params, inputs.at(0), inputs.at(1).weights(), + transpose_b, node_def.name()); +} + +tensorflow::Status ConvertBatchMatMul(OpConverterParams* params) { + const auto& inputs = params->inputs; + const auto& node_def = params->node_def; + TFAttrs attrs(node_def); + + tensorflow::DataType tf_dtype = attrs.get("T"); + if (tf_dtype != tensorflow::DataType::DT_FLOAT && + tf_dtype != tensorflow::DataType::DT_HALF) { + return tensorflow::errors::Unimplemented( + "data type is not supported, for node " + node_def.name() + " got " + + tensorflow::DataTypeString(tf_dtype)); + } + + bool transpose_a = attrs.get("adj_x"); + bool transpose_b = attrs.get("adj_y"); + + auto dims = inputs.at(0).GetTrtDims(); + if (dims.nbDims == 1) { // NC * CK is only supported through fully connected + if (transpose_a == false && inputs.at(0).is_tensor() && + inputs.at(1).is_weights()) { + return ConvertMatMulHelper(params, inputs.at(0), inputs.at(1).weights(), + transpose_b, node_def.name()); + } else { + return tensorflow::errors::InvalidArgument( + "Invalid configuration for MatMul, at: " + node_def.name()); + } + } + + const nvinfer1::ITensor* tensor_l; + const nvinfer1::ITensor* tensor_r; + auto dims_l = inputs.at(0).GetTrtDims(); + auto dims_r = inputs.at(1).GetTrtDims(); + if (inputs.at(0).is_weights()) { + if (inputs.at(0).GetTrtDims().d[0] != 1) { + return tensorflow::errors::InvalidArgument( + "Input 0 as weight assumes broadcast across batch for MatMul, at: " + + node_def.name()); + } else { + for (int i = 0; i < dims_l.nbDims - 1; i++) { + dims_l.d[i] = dims_l.d[i + 1]; + } + dims_l.nbDims--; + } + } + if (inputs.at(1).is_weights()) { + if (inputs.at(1).GetTrtDims().d[0] != 1) { + return tensorflow::errors::InvalidArgument( + "Input 1 as weight assumes broadcast across batch for MatMul, at: " + + node_def.name()); + } else { + for (int i = 0; i < dims_r.nbDims - 1; i++) { + dims_r.d[i] = dims_r.d[i + 1]; + } + dims_r.nbDims--; + } + } + TF_RETURN_IF_ERROR(params->converter->PrepareTensorForShape( + inputs.at(0), dims_l, &tensor_l)); + TF_RETURN_IF_ERROR(params->converter->PrepareTensorForShape( + inputs.at(1), dims_r, &tensor_r)); + + nvinfer1::IMatrixMultiplyLayer* layer = + params->converter->network()->addMatrixMultiply( + *const_cast(tensor_l), transpose_a, + *const_cast(tensor_r), transpose_b); + TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); + nvinfer1::ITensor* output_tensor = layer->getOutput(0); + params->outputs->push_back(TRT_TensorOrWeights(output_tensor)); + return tensorflow::Status::OK(); +} + +tensorflow::Status ConvertSoftmax(OpConverterParams* params) { + const auto& inputs = params->inputs; + const auto& node_def = params->node_def; + const nvinfer1::ITensor* tensor = inputs.at(0).tensor(); + + int nbDims = tensor->getDimensions().nbDims; + if (nbDims == 0) { + return tensorflow::errors::InvalidArgument( + "TensorRT Softmax cannot apply on batch dimension, at" + + node_def.name()); + } + nvinfer1::ISoftMaxLayer* layer = params->converter->network()->addSoftMax( + *const_cast(tensor)); + TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); + // Tensorflow SoftMax assumes applying softmax on the last dimension. + layer->setAxes(1 << (nbDims - 1)); + + nvinfer1::ITensor* output_tensor = layer->getOutput(0); + // Quantization range for SoftMax is always (0, 1) + params->converter->ProvideQuantizationRange(output_tensor, 0.0f, 1.0f); + params->outputs->push_back(TRT_TensorOrWeights(output_tensor)); + return tensorflow::Status::OK(); +} + +tensorflow::Status ConvertTopK(OpConverterParams* params) { + const auto& inputs = params->inputs; + const auto& node_def = params->node_def; + 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()); + } + + 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()); + } + + nvinfer1::ITopKLayer* layer = params->converter->network()->addTopK( + *const_cast(tensor), op, k, reducedAxes); + TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); + + nvinfer1::ITensor* output_value_tensor = layer->getOutput(0); + nvinfer1::ITensor* output_indices_tensor = layer->getOutput(1); + // Tensor type for network output is not inferred. Indices should be INT32 + // (default is float). + output_indices_tensor->setType(nvinfer1::DataType::kINT32); + params->outputs->push_back(TRT_TensorOrWeights(output_value_tensor)); + params->outputs->push_back(TRT_TensorOrWeights(output_indices_tensor)); + return tensorflow::Status::OK(); +} + +static void RegisterValidatableOpConverters( + std::unordered_map* registration) { + // TODO(laigd): support all op types. + (*registration)["BiasAdd"] = ConvertBiasAdd; + (*registration)["ConcatV2"] = ConvertConcat; + (*registration)["Const"] = ConvertConst; + (*registration)["Conv2D"] = ConvertConv2D; + (*registration)["DepthwiseConv2dNative"] = ConvertConv2DDepthwise; + (*registration)["ExpandDims"] = ConvertExpandDims; + (*registration)["MatMul"] = ConvertMatMul; + (*registration)["Pad"] = ConvertPad; + (*registration)["Relu6"] = ConvertRelu6; + (*registration)["Reshape"] = ConvertReshape; + (*registration)["Square"] = ConvertSquare; + (*registration)["Squeeze"] = ConvertSqueeze; + (*registration)["StridedSlice"] = ConvertStridedSlice; + (*registration)["Transpose"] = ConvertTranspose; + + for (auto quantization_op_type : + {"QuantizeAndDequantizeV2", "QuantizeAndDequantizeV3", + "FakeQuantWithMinMaxVars", "FakeQuantWithMinMaxArgs"}) { + (*registration)[quantization_op_type] = ConvertQuantize; + } + for (auto binary_op_type : + {"Add", "Mul", "Sub", "Div", "RealDiv", "Maximum", "Minimum"}) { + (*registration)[binary_op_type] = ConvertBinary; + } + for (auto activation_op_type : {"Relu", "Sigmoid", "Tanh"}) { + (*registration)[activation_op_type] = ConvertActivation; + } + for (auto pool_op_type : {"AvgPool", "MaxPool"}) { + (*registration)[pool_op_type] = ConvertPool; + } + for (auto normalization_op_type : {"FusedBatchNorm", "FusedBatchNormV2"}) { + (*registration)[normalization_op_type] = ConvertFusedBatchNorm; + } +} + +void TrtNodeValidator::RegisterOpValidators() { + RegisterValidatableOpConverters(&op_validators_); +} + +void Converter::RegisterOpConverters() { + RegisterValidatableOpConverters(&op_registry_); + // TODO(ben,jie): this is a temp hack. + op_registry_["Identity"] = ConvertIdentity; // Identity should be removed + op_registry_["Snapshot"] = ConvertIdentity; // Snapshot should be removed + + op_registry_["Rsqrt"] = ConvertUnary; + op_registry_["Reciprocal"] = ConvertUnary; + op_registry_["Exp"] = ConvertUnary; + op_registry_["Log"] = ConvertUnary; + op_registry_["Sqrt"] = ConvertUnary; + op_registry_["Abs"] = ConvertUnary; + op_registry_["Neg"] = ConvertUnary; + + op_registry_["Sum"] = ConvertReduce; + op_registry_["Prod"] = ConvertReduce; + op_registry_["Max"] = ConvertReduce; + op_registry_["Min"] = ConvertReduce; + op_registry_["Mean"] = ConvertReduce; + op_registry_["Softmax"] = ConvertSoftmax; + op_registry_["BatchMatMul"] = ConvertBatchMatMul; + 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 std::vector& input_shapes, + Logger* logger, nvinfer1::IGpuAllocator* allocator, + TRTInt8Calibrator* calibrator, + TrtUniquePtrType* engine, bool use_calibration, + bool* convert_successfully) { + engine->reset(); + if (convert_successfully) *convert_successfully = false; + + // Create the builder. + TrtUniquePtrType builder( + nvinfer1::createInferBuilder(*logger)); + builder->setMaxBatchSize(max_batch_size); + builder->setMaxWorkspaceSize(max_workspace_size_bytes); + builder->setGpuAllocator(allocator); + if (precision_mode == FP16MODE) { + builder->setHalf2Mode(true); + } else if (precision_mode == INT8MODE) { + builder->setInt8Mode(true); + if (use_calibration) { + builder->setInt8Calibrator(calibrator); + } else { + builder->setInt8Calibrator(nullptr); + } + } + + // Create the network. + auto trt_network = + TrtUniquePtrType(builder->createNetwork()); + if (!trt_network) { + return tensorflow::errors::Internal( + "Failed to create TensorRT network object"); + } + + // Build the network + VLOG(1) << "Starting engine conversion "; + Converter converter(trt_network.get(), precision_mode, use_calibration); + 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")) { + int32 slot_number = -1; + 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); + } + nvinfer1::DataType trt_dtype; + nvinfer1::Dims trt_dims; + int batch_size = -1; + auto shape = input_shapes.at(slot_number); + auto status = ValidateTensorProperties( + node_def.op(), node_def.attr().at("dtype").type(), shape, + /*validation_only=*/false, &trt_dtype, &trt_dims, &batch_size); + if (!status.ok()) { + const string error_message = + StrCat("Validation failed for ", node_name, " and input slot ", + slot_number, ": ", status.error_message()); + LOG(WARNING) << error_message; + return Status(status.code(), error_message); + } + VLOG(2) << "Adding engine input tensor " << node_name << " with shape " + << DebugString(trt_dims); + // TODO(laigd): the conversion should always happen at runtime where all + // the shapes are known, and we can provide a mode to generate the + // 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")) { + int32 slot_number = -1; + 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); + } + if (output_tensors.size() <= slot_number) { + output_tensors.resize(slot_number + 1); + } + output_tensors.at(slot_number) = {node_def.input(0), node_name}; + } else { + VLOG(2) << "Converting node: " << node_def.name() << " , " + << node_def.op(); + TF_RETURN_IF_ERROR(converter.ConvertNode(node_def)); + } + } + TF_RETURN_IF_ERROR(converter.RenameAndMarkOutputTensors(output_tensors)); + if (convert_successfully) *convert_successfully = true; + + // Apply user provided quantization ranges to tensors + converter.MaybeApplyQuantizationRanges(); + + // Build the engine. + VLOG(1) << "Starting engine creation"; + engine->reset(builder->buildCudaEngine(*converter.network())); + if (engine->get() == nullptr) { + return tensorflow::errors::Internal("Failed to build TensorRT engine"); + } + VLOG(1) << "Finished conversion"; + return tensorflow::Status::OK(); +} + +tensorflow::Status ConvertSegmentToGraphDef( + const tensorflow::Graph* graph, + const tensorflow::grappler::GraphProperties& graph_properties, + const std::vector& subgraph_nodes, // In topological order + std::vector* connections, + tensorflow::GraphDef* segment_def, string* common_scope) { + std::set marker_nodes; + // Update connection shapes/data types and add corresponding input/output + // nodes in the segment graphdef. + for (size_t i = 0; i < connections->size(); ++i) { + auto& connection = connections->at(i); + if (connection.is_control_edge()) continue; + auto outside_node = graph->FindNodeId(connection.outside_id); + if (!outside_node) { + // This should never happen, unless the original graph is problematic. + return tensorflow::errors::NotFound( + "Cannot find node with id ", connection.outside_id, " in the graph."); + } + // Updates the shape and data types of input/output connections. + tensorflow::DataType dtype; + tensorflow::PartialTensorShape partial_shape; + if (connection.is_input_edge) { + GetOutputProperties(graph_properties, + graph->FindNodeId(connection.outside_id), + connection.outside_port, &partial_shape, &dtype); + connection.outside_shape = partial_shape; + } else { + GetInputProperties(graph_properties, + graph->FindNodeId(connection.outside_id), + connection.outside_port, &partial_shape, &dtype); + connection.inside_shape = partial_shape; + } + connection.connection_type = dtype; + + // Add dummy input/output nodes to the segment graphdef. + if (connection.is_input_edge) { + const string node_name = StrCat(kInputPHName, connection.port_number); + if (marker_nodes.count(node_name)) { + VLOG(1) << "Reusing input " << node_name << " for the edge " + << connection.outside_node_name << ":" + << connection.outside_port << " -> " + << connection.inside_node_name << ":" << connection.inside_port; + continue; + } + marker_nodes.insert(node_name); + auto seg_node = segment_def->add_node(); + tensorflow::NodeDefBuilder builder(node_name, "Placeholder"); + auto status = builder.Attr("shape", partial_shape) + .Attr("dtype", dtype) + .Finalize(seg_node); + VLOG(1) << "Constructing input " << node_name << " for the edge " + << connection.outside_node_name << ":" << connection.outside_port + << " -> " << connection.inside_node_name << ":" + << connection.inside_port; + } else { + const string node_name = StrCat(kOutputPHName, connection.port_number); + if (marker_nodes.count(node_name)) { + VLOG(1) << "Reusing output " << node_name << " for the edge " + << connection.inside_node_name << ":" << connection.inside_port + << " -> " << connection.outside_node_name << ":" + << connection.outside_port; + continue; + } + 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, 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 << ":" + << connection.outside_port; + } + } // for each connection. + + std::unordered_map old_to_new_id_map; + // Copy internal nodes to new graphdef + 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(); + auto snode = segment_def->add_node(); + snode->CopyFrom(node->def()); + VLOG(2) << "Copying " << snode->name() << " to subgraph"; + } + // Update the inputs of the new input nodes to point to placeholder nodes. + for (int i = 0; i < connections->size(); ++i) { + auto& connection = connections->at(i); + if (connection.is_control_edge() || !connection.is_input_edge) continue; + auto snode = + segment_def->mutable_node(old_to_new_id_map[connection.inside_id]); + const string placeholder_name = + StrCat(kInputPHName, connection.port_number); + VLOG(1) << "Updating " << snode->name() << ":" << connection.inside_port + << " from " << snode->input(connection.inside_port) << " to " + << 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); + const int input_size = snode->input_size(); + int input_idx = 0; + int actual_input_idx = 0; + while (input_idx < input_size) { + 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)) { + if (input.second == Graph::kControlSlot) { + VLOG(1) << "... removing control inputs " << input.first + << " from subgraph."; + ++input_idx; + continue; + } else { + return tensorflow::errors::InvalidArgument( + "Found non control input outside the segment that is not an " + "engine connection to ", + snode->name(), ": ", input.first); + } + } + if (actual_input_idx != input_idx) { + snode->set_input(actual_input_idx, snode->input(input_idx)); + } + ++input_idx; + ++actual_input_idx; + } + for (int remove = input_size - actual_input_idx; remove > 0; --remove) { + snode->mutable_input()->RemoveLast(); + } + } + *common_scope = local_scope; + VLOG(1) << "Converted TensorRT candidate segment @scope '" << local_scope + << "' to a GraphDef"; + return tensorflow::Status::OK(); +} + +bool OutputEdgeValidator::operator()(const tensorflow::Edge* out_edge) const { + if (out_edge->IsControlEdge()) return true; + if (out_edge->src()->type_string() == "Const") { + VLOG(1) << "--> Need to remove output node " << out_edge->src()->name() + << " which is a Const."; + return false; + } + return true; +} + +} // namespace convert +} // namespace tensorrt +} // namespace tensorflow + +#endif // GOOGLE_TENSORRT +#endif // GOOGLE_CUDA diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h new file mode 100644 index 0000000000000000000000000000000000000000..aebc0ca38de449dd716b3948f9a0b2e581fc8c80 --- /dev/null +++ b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h @@ -0,0 +1,558 @@ +/* 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_CONVERT_CONVERT_NODES_H_ +#define TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_CONVERT_NODES_H_ + +#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/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" +#include "tensorflow/core/lib/core/status.h" + +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT +#include "tensorrt/include/NvInfer.h" + +namespace tensorflow { +namespace tensorrt { +extern const char* const kInputPHName; +extern const char* const kOutputPHName; + +namespace convert { + +struct EngineConnection { + // Constructs a non-control edge. + EngineConnection(const string& outside, int out_id, int out_port, + const string& inside, int in_id, int in_port, + bool input_edge, int port) + : outside_node_name(outside), + outside_id(out_id), + outside_port(out_port), + inside_node_name(inside), + inside_id(in_id), + inside_port(in_port), + is_input_edge(input_edge), + port_number(port) {} + + // Constructs a control edge. + EngineConnection(const string& outside, int out_id, const string& inside, + int in_id, bool input_edge) + : outside_node_name(outside), + outside_id(out_id), + outside_port(Graph::kControlSlot), + inside_node_name(inside), + inside_id(in_id), + inside_port(Graph::kControlSlot), + is_input_edge(input_edge), + port_number(Graph::kControlSlot) {} + + bool is_control_edge() const { return port_number == Graph::kControlSlot; } + + const string outside_node_name; + const int outside_id; + const int outside_port; + tensorflow::PartialTensorShape outside_shape; // Only set for input edge. + + const string inside_node_name; + const int inside_id; + const int inside_port; + tensorflow::PartialTensorShape inside_shape; // Only set for output edge. + + tensorflow::DataType connection_type; + const bool is_input_edge; + + // The port number of the TRT node connected with this edge. + const int port_number; +}; + +struct EngineInfo { + EngineInfo() + : engine_type(EngineType::TRTStatic), + max_workspace_size_bytes(0), + precision_mode(FP32MODE), + use_calibration(true) {} + + string engine_name; + string device; + tensorflow::GraphDef segment_graph_def; + + // Non-control input connections inside this vector are sorted in a way such + // that, the segment nodes connecting to them are topological sorted. + // In addition, for non-control connections, there must be no duplicates. + std::vector connections; + + enum class EngineType { TRTStatic = 0, TRTDynamic = 1 }; + EngineType engine_type; + int64 max_workspace_size_bytes; + int maximum_cached_engines; + std::vector cached_engine_batches; + int precision_mode; + bool use_calibration; +}; + +// Constructs a graphdef from the segment in the given graph. Adds placeholder +// nodes for input edges (InputPH_*) and identity nodes for output edges +// (OutputPH_*). This function needs to be called before TensorRT nodes +// inserted in order to correctly get sizes from the original graph. +// +// - subgraph_node_names: the node names of the subgraph. +// - subgraph_node_ids: the node ids of the subgraph, must be sorted in +// topological order. +// - segment_def: the output GraphDef, whose non-input/output nodedefs will be +// sorted in topological order. +// +// TODO(aaroey): add tests to validate these properties. +tensorflow::Status ConvertSegmentToGraphDef( + const tensorflow::Graph* graph, + const tensorflow::grappler::GraphProperties& graph_properties, + const std::vector& subgraph_nodes, + std::vector* connections, + tensorflow::GraphDef* segment_def, string* common_scope); + +// Converts given subgraph to a TRT engine saved in 'engine'. Returns ok iff +// 'builder' successfully build the engine. If the result is not ok, 'engine' +// will be set to nullptr +// Once returned, 'builder' is not needed any more and can be safely detroyed. +// +// - convert_successfully: indicates whether the converson to TensorRT network +// is successful. This is different than successfully building the engine: +// building can still fail afterwards. +tensorflow::Status ConvertGraphDefToEngine( + const tensorflow::GraphDef& gdef, int precision_mode, int max_batch_size, + size_t max_workspace_size_bytes, + const std::vector& input_shapes, + Logger* logger, nvinfer1::IGpuAllocator* allocator, + TRTInt8Calibrator* calibrator, + TrtUniquePtrType* engine, bool use_calibration, + bool* convert_successfully); + +// Helper class for the segmenter to determine whether an output edge from the +// TRT segment is valid. +class OutputEdgeValidator { + public: + // Return true if the specified edge is eligible to be an output edge of the + // TRT segment. + bool operator()(const tensorflow::Edge* out_edge) const; +}; + +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); + +// Class to convert TF compile-time constants (e.g. Const nodes) to TRT weight. +class TRT_ShapedWeights { + public: + explicit TRT_ShapedWeights(DataType type = DT_FLOAT); + + // Copy from another weights. + // + // NOTE: this does not copy the underlying buffer but only increase its + // reference count. + TRT_ShapedWeights(const TRT_ShapedWeights& rhs); + + nvinfer1::Weights GetTrtWeights() const; + + void* GetValues() const { + return const_cast(tensor_.tensor_data().data()); + } + + int64_t count() const; + + size_t size_bytes() const; + + string DebugString() const; + + // TODO(aaroey): make these private. + nvinfer1::Dims shape_; // Note: shape.type[] is not used. + tensorflow::DataType type_; + + private: + // This constructor is only used by TrtWeightStore, which creates the + // 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; +}; + +// Container for TRT_ShapedWeights. We need this container because, TRT doesn't +// manage the lifetime of the weights buffer, it only keeps a pointer to it and +// requires that the data referenced by the pointer be available until the +// building of engine is complete. For more information see +// https://docs.nvidia.com/deeplearning/sdk/tensorrt-api/c_api/classnvinfer1_1_1_weights.html +// +// TODO(laigd): consider adding garbage collection to the unused weights. +class TrtWeightStore { + public: + // Get a TRT_ShapedWeights with 'type' and 'dims'. + TRT_ShapedWeights GetTempWeights(tensorflow::DataType type, + const nvinfer1::Dims& dims); + + // Get a TRT_ShapedWeights with the same data type and dimensions as + // 'weights'. + TRT_ShapedWeights GetTempWeights(const TRT_ShapedWeights& weights) { + return GetTempWeights(weights.type_, weights.shape_); + } + + private: + // The backend storage of the TRT_ShapedWeights. + std::vector store_; +}; + +// Represents a TRT-style input to a TF node, it can be either a +// nvinfer1::ITensor, or TRT_ShapedWeights which is compile-time constant. +// +// TODO(laigd): maybe rename it to TrtArgument, or mimic XlaCompiler::Argument. +class TRT_TensorOrWeights { + public: + TRT_TensorOrWeights() {} + + // Constructor that makes it an ITensor, doesn't take ownership of 'tensor'. + // This is used by Converter when building the TRT network, where the ITensor + // is owned by the TRT network being built. See comment for 'tensor_' below. + explicit TRT_TensorOrWeights(nvinfer1::ITensor* tensor, int batch_size = -1); + + // Constructor that makes it an ITensor by creating one using provided data + // type and shape, and takes ownership of the created ITensor. This is used by + // TrtNodeValidator to encapsulate the type and shape information for + // validation of graph nodes, and the created ITensor is fake and temporary, + // and should not be used to build any TRT network. See comment for + // 'simple_itensor_' below. + explicit TRT_TensorOrWeights(nvinfer1::DataType trt_dtype, + const nvinfer1::Dims& trt_dims, int batch_size); + + // Constructor that makes it a TRT_TensorOrWeights. + explicit TRT_TensorOrWeights(const TRT_ShapedWeights& weights); + + TRT_TensorOrWeights(const TRT_TensorOrWeights& rhs); + + void operator=(const TRT_TensorOrWeights& rhs); + + bool is_tensor() const { return initialized_ && is_tensor_; } + bool is_weights() const { return initialized_ && !is_tensor_; } + + nvinfer1::ITensor* tensor(); + + const nvinfer1::ITensor* tensor() const; + + TRT_ShapedWeights& weights() { + CHECK(is_weights()); + return weights_; + } + + const TRT_ShapedWeights& weights() const { + CHECK(is_weights()); + return weights_; + } + + nvinfer1::Dims GetTrtDims() const; + + int batch_size() const { return batch_size_; } + + string DebugString() const; + + private: + class SimpleITensor; + + void set_batch_size(int batch_size) { batch_size_ = batch_size; } + + // When it represents an ITensor, the ITensor can be either passed by the + // caller via the constructor that takes an ITensor* as parameter, or be + // created as a SimpleITensor. + // + // In the first case, the ITensor pointer is stored in 'tensor_' below, and + // the ITensor itself is not owned by this class. This method is used by + // Converter (e.g. AddInputTensor) and op converters during TRT network + // construction, where the TRT network owns the ITensor. + // + // In the second case, the created SimpleITensor is stored in + // 'simple_itensor_' below and is owned by this class. SimpleITensor is a fake + // implementation of ITensor and is used only by TrtNodeValidator to validate + // the graph nodes. + nvinfer1::ITensor* tensor_ = nullptr; // Not owned. + std::shared_ptr simple_itensor_ = nullptr; + + // First dimension of the TF tensor (NOT tensor_) that is represented by + // tensor_ is treated as the "batch dimension" by TRT, and tensor_'s + // dimensions (obtained via tensor_->getDimensions()) do not contain the batch + // dimension. For example, when a TF tensor with shape (A,B,C) is represented + // in TRT, tensor_->getDimensions() will be (B,C) and batch_size_ will be A. + // + // This requires that all tensors in the subgraph that is converted to a TRT + // engine have the same batch size are represented by the first dimension of + // their shape, and Converter will verify this during conversion. The drawback + // is that currently it cannot convert a graph that doesn't have the batch + // size represented in the shapes or the batch sizes are different. See + // b/118387490 for more details. + int batch_size_ = -1; + + TRT_ShapedWeights weights_; + bool initialized_ = false; + bool is_tensor_ = false; + + friend class Converter; +}; + +class Converter; + +// Parameters for each op converter. +struct OpConverterParams { + OpConverterParams(Converter* arg_converter, + const tensorflow::NodeDef& arg_node_def, + const std::vector& arg_inputs, + std::vector* arg_outputs, + bool arg_validation_only, TrtWeightStore* arg_weight_store) + : converter(arg_converter), + node_def(arg_node_def), + inputs(arg_inputs), + outputs(arg_outputs), + validation_only(arg_validation_only), + weight_store(arg_weight_store) {} + + Converter* converter; + const tensorflow::NodeDef& node_def; + const std::vector& inputs; + std::vector* outputs; + const bool validation_only; + TrtWeightStore* weight_store; +}; + +using OpConverter = std::function; + +// Class to verify if specific TF node is supported by TRT. +class TrtNodeValidator { + public: + TrtNodeValidator(); + + // Validate the node, and return ok if it's supported by TRT. + // + // - 'node_def' is the node to validate. + // - 'input_node_and_ports' are the input NodeDefs and their output ports that + // are connected to 'node_def' in the TF graph. + // - 'graph_properties' is the GraphProperties of the graph where 'node_def' + // belongs. It is used to get the shape and data type information of a + // tensor for validation purpose. + Status ValidateNode( + const NodeDef& node_def, + const std::vector>& input_node_and_ports, + const grappler::GraphProperties& graph_properties); + + private: + void RegisterOpValidators(); + + // Convert a Const node to a TRT_TensorOrWeights. + Status ConvertConstToWeights(const NodeDef& const_node_def, + const std::vector& inputs, + TRT_TensorOrWeights* output); + + // Convert the output tensor at 'output_port' of 'node_def' to a + // TRT_TensorOrWeights which will be later used as an input to other nodes and + // passed to ValidateNode() below. + Status ConvertToTensorOrWeights( + const NodeDef& node_def, int output_port, + const grappler::GraphProperties& graph_properties, + TRT_TensorOrWeights* tensor_or_weights); + + // Stores all the validators by op type. If no validator is registered for + // specific op, it means no validation is needed and ValidateNode() will + // return OK. + std::unordered_map op_validators_; + + // Store the weights added during validation. Some validations (e.g. + // validation for Const node) may produce weights. + TrtWeightStore weight_store_; + + friend class ValidatorTest; + friend class OpConverterTest; +}; + +// Class to convert TF nodes to TRT network. +class Converter { + public: + Converter(nvinfer1::INetworkDefinition* trt_network, int precision_mode, + bool use_calibration); + + ////////////////////////////////////////////////////////////////////////////// + // Methods used by the TRT engine builder to build a TRT network from a TF + // function/subgraph. + + // Convert the node to TRT network. + Status ConvertNode(const tensorflow::NodeDef& node_def); + + // Add input tensor to the TRT network with given 'name', 'dtype', 'dims' and + // 'batch_size'. + 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). + Status RenameAndMarkOutputTensors( + const std::vector>& output_tensors); + + ////////////////////////////////////////////////////////////////////////////// + // Methods used by op converters to convert individual TF node and add layers + // to the TRT network. + + // Op converters (e.g. ConvertReshape) need to access the TRT network in order + // to add TRT layers. + nvinfer1::INetworkDefinition* network() { return trt_network_; } + + // What precision are we targeting? + int precision_mode() const { return precision_mode_; } + + // Calibration will be or was previously performed on this network? + bool use_calibration() const { return use_calibration_; } + + // This should be called on the inputs and outputs of any layer we create + // where we know that the quantization range does not change during that + // operation. (e.g. Reshape, Transpose, Identity, MaxPool). + void MarkQuantizationRangesAsInferrable(nvinfer1::ITensor* input, + nvinfer1::ITensor* output); + + // This function should be called when we know the quantization range of a + // tensor, either from a quantize/dequantize node or when the output is a + // fixed range (e.g. SoftMax, Relu6, Sigmoid). + void ProvideQuantizationRange(nvinfer1::ITensor* tensor, float min_range, + float max_range); + + // Should be called when full TRT network has been constructed and before + // building the engine. + void MaybeApplyQuantizationRanges(); + + // Below are helper methods for op converters to add different layers to the + // TRT network. + + // Transpose 'input_tensor' with given permutation 'order_with_batch_dim' to + // 'output_tensor'. The permutation 'order_with_batch_dim' contains the batch + // dimension which should always be 0. + Status TransposeTensor(nvinfer1::ITensor* input_tensor, + const std::vector& order_with_batch_dim, + const nvinfer1::ITensor** output_tensor); + + // Converts 'input' into 'tensor' with shape specified by 'dims'. + Status PrepareTensorForShape(const TRT_TensorOrWeights& input, + const nvinfer1::Dims& dims, + const nvinfer1::ITensor** tensor); + + // Return OK if the broadcast scheme is supported and compute the shapes after + // broadcasting. + Status GetTrtBroadcastShape(const TRT_TensorOrWeights& operand_l, + const TRT_TensorOrWeights& operand_r, + 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. + Status MaybeUpdateBatchSize(int batch_size); + + // Add the provided tensor/weights to the map trt_tensors_. + Status AddTensorOrWeights(const string& name, TRT_TensorOrWeights input); + + // Get the tensor/weights from trt_tensors_ by 'name'. + Status GetTensorOrWeights(const string& name, TRT_TensorOrWeights* output); + + // Get the inputs of 'node_def' from trt_tensors_. + Status GetInputs(const tensorflow::NodeDef& node_def, + std::vector* inputs) const; + + void RegisterOpConverters(); + + void PropagateQuantizationRanges(); + + // Gets the min and max value in a TRT_ShapedWeights + Status GetWeightRange(const TRT_ShapedWeights& weights, float* out_min, + float* out_max) const; + + // Registered op converters by op type. + std::unordered_map op_registry_; + + // Tensors/weights added during construction of trt_network_. + std::unordered_map trt_tensors_; + + // Special op converter for custom plugins. + OpConverter plugin_converter_; + + // The TRT networking being built. + nvinfer1::INetworkDefinition* trt_network_; + + // Store the weights added during construction of trt_network_. + TrtWeightStore weight_store_; + + // During conversion, this table is populated with quantization ranges per + // tensor. MaybeApplyQuantizationRanges() will use this table to set the TRT + // quantization ranges. Since TRT only supports symmetric ranges, we will + // store the range as a single float = max(abs(min_range), abs(max_range)). + // Range refers to the floating point values, e.g. min_range = 0.0f, max_range + // = 6.0f for Relu6. + std::unordered_map quantization_ranges_; + + // Edges where quantization ranges can be inferred (copied) across ops - from + // first tensor to second tensor. PropagateQuantizationRanges() will propagate + // known ranges from quantization_ranges_ across these edges, adding the new + // ranges to quantization_ranges_ so that they can be applied in + // MaybeApplyQuantizationRanges(). + std::vector> + quantization_infer_; + + const int precision_mode_; + + const bool use_calibration_; + + // Batch size of inputs to trt_network_ added by AddInputTensor(). During + // network construction it will update this, use it to verify the batch + // size of all inputs are compatible, and make sure individual TF node is + // acceptable by TRT. + int batch_size_ = -1; + + friend class ConverterTest; + friend class OpConverterTest; +}; + +} // namespace convert +} // namespace tensorrt +} // namespace tensorflow + +#endif // GOOGLE_TENSORRT +#endif // GOOGLE_CUDA + +#endif // TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_CONVERT_NODES_H_ diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..3a70423d12b35e46d2709dcdc25920a3143f41c4 --- /dev/null +++ b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc @@ -0,0 +1,2927 @@ +/* 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/convert/convert_nodes.h" + +#include +#include +#include + +#include +#include +#include "absl/strings/str_cat.h" +#include "tensorflow/cc/framework/ops.h" +#include "tensorflow/cc/framework/scope.h" +#include "tensorflow/cc/ops/standard_ops.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 +#include "tensorflow/core/framework/tensor_shape.h" +#include "tensorflow/core/framework/tensor_testutil.h" +#include "tensorflow/core/framework/types.h" +#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/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" + +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT +#include "cuda/include/cuda.h" +#include "cuda/include/cuda_runtime_api.h" +#include "tensorrt/include/NvInfer.h" + +namespace tensorflow { +namespace tensorrt { +namespace convert { + +using absl::StrCat; +using ::testing::ElementsAre; +using ::testing::ElementsAreArray; + +// TODO(laigd): put this into some test utils file. +void ExpectStatus(Status status, error::Code code = error::OK, + const char* substr = nullptr) { + EXPECT_EQ(code, status.code()) + << status << " vs expected error code \"" << error::Code_Name(code) + << "\" and message \"" << substr << "\""; + if (substr) { + EXPECT_THAT(status.error_message(), ::testing::HasSubstr(substr)) << status; + } +} + +nvinfer1::Dims GetTestDims(const std::vector& d) { + nvinfer1::Dims dims; + dims.nbDims = d.size(); + for (int i = 0; i < d.size(); ++i) { + dims.d[i] = d[i]; + } + return dims; +} + +nvinfer1::DataType TfDataTypeToTrt(DataType tf_dtype) { + switch (tf_dtype) { + case DT_FLOAT: + return nvinfer1::DataType::kFLOAT; + case DT_HALF: + return nvinfer1::DataType::kHALF; + case DT_INT32: + return nvinfer1::DataType::kINT32; + default: + QCHECK(false) << "Unexpected data type " << DataTypeString(tf_dtype); + } +} + +DataType TrtDataTypeToTf(nvinfer1::DataType trt_dtype) { + switch (trt_dtype) { + case nvinfer1::DataType::kFLOAT: + return DT_FLOAT; + case nvinfer1::DataType::kHALF: + return DT_HALF; + case nvinfer1::DataType::kINT32: + return DT_INT32; + default: + QCHECK(false) << "Unexpected data type " << static_cast(trt_dtype); + } +} + +NodeDef MakeNodeDef(const string& name, const string& op, + const std::vector& inputs) { + NodeDef node_def; + node_def.set_name(name); + node_def.set_op(op); + for (const string& input : inputs) { + node_def.add_input(input); + } + return node_def; +} + +template +NodeDef MakeConstNodeDef(const string& name, const std::vector& vals, + const TensorShape& shape) { + Scope s = Scope::NewRootScope(); + Tensor t = ::tensorflow::test::AsTensor(vals, shape); + auto const_op = ops::Const(s.WithOpName(name), t); + return const_op.node()->def(); +} + +template +NodeDef MakeConstNodeDef(const string& name, const std::vector& vals) { + TensorShape shape; + const std::vector shape_dims = {static_cast(vals.size())}; + TF_EXPECT_OK(TensorShapeUtils::MakeShape(shape_dims, &shape)); + return MakeConstNodeDef(name, vals, shape); +} + +bool TrtDimsEquals(const nvinfer1::Dims& lhs, const nvinfer1::Dims& rhs) { + if (lhs.nbDims != rhs.nbDims) return false; + for (int i = 0; i < lhs.nbDims; ++i) { + if (lhs.d[i] != rhs.d[i]) return false; + // We don't check the types in the tests. + } + return true; +} + +bool TrtDimsEqualsArray(const std::vector& lhs, + const nvinfer1::Dims& rhs) { + return TrtDimsEquals(GetTestDims(lhs), rhs); +} + +// TODO(laigd): define a parameterized matcher that can compare against the +// vector. +void ExpectTrtDimsEqualsArray(const std::vector& lhs, + const nvinfer1::Dims& rhs) { + EXPECT_TRUE(TrtDimsEqualsArray(lhs, rhs)) + << "expected: " << DebugString(GetTestDims(lhs)) << "\n" + << " actual: " << DebugString(rhs); +} + +template +void ExpectArrayNear(const std::vector& lhs, const std::vector& rhs) { + ASSERT_EQ(lhs.size(), rhs.size()); + for (int i = 0; i < lhs.size(); i++) { + EXPECT_FLOAT_EQ(lhs[i], rhs[i]); + } +} + +// Eigen::half cannot implicitly convert to float which is required for +// EXPECT_FLOAT_EQ. +template <> +void ExpectArrayNear(const std::vector& lhs, + const std::vector& 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]), + Eigen::half_impl::half_to_float(rhs[i])); + } +} + +bool TrtShapedWeightsEquals(const TRT_ShapedWeights& lhs, + const TRT_ShapedWeights& rhs) { + return TrtDimsEquals(lhs.shape_, rhs.shape_) && lhs.type_ == rhs.type_ && + lhs.GetValues() == rhs.GetValues(); +} + +template +void ValidateWeights(const TRT_ShapedWeights& weights, + const std::vector& expected_dims, + const std::vector& expected_value) { + ExpectTrtDimsEqualsArray(expected_dims, weights.shape_); + ASSERT_EQ(expected_value.size(), weights.count()) << weights.DebugString(); + const T* actual_values = static_cast(weights.GetValues()); + for (int i = 0; i < expected_value.size(); ++i) { + EXPECT_EQ(expected_value[i], actual_values[i]); + } +} + +// Fake ITensor implementation for testing purposes. +class FakeITensor : public nvinfer1::ITensor { + public: + FakeITensor() : dynamic_range_(0.0f) {} + + FakeITensor(const nvinfer1::Dims& dims) : dims_(dims), dynamic_range_(0.0f) {} + + FakeITensor(const std::vector& dims) + : dims_(GetTestDims(dims)), dynamic_range_(0.0f) {} + + void setName(const char* name) override { name_ = name; } + + const char* getName() const override { return name_.c_str(); } + + void setDimensions(nvinfer1::Dims dimensions) override { dims_ = dimensions; } + + nvinfer1::Dims getDimensions() const override { return dims_; } + + void setType(nvinfer1::DataType type) override { type_ = type; } + + nvinfer1::DataType getType() const override { return type_; } + + bool isNetworkInput() const override { return false; } + + bool isNetworkOutput() const override { return false; } + + void setBroadcastAcrossBatch(bool broadcastAcrossBatch) override {} + + bool getBroadcastAcrossBatch() const override { return false; } + + nvinfer1::TensorLocation getLocation() const override { return location_; } + + void setLocation(nvinfer1::TensorLocation location) override { + location_ = location; + } + +#if NV_TENSORRT_MAJOR >= 5 + bool setDynamicRange(float min, float max) override { + dynamic_range_ = std::max(std::abs(min), std::abs(max)); + return true; + } + + float getDynamicRange() const override { return dynamic_range_; } +#endif + + private: + string name_; + nvinfer1::Dims dims_; + nvinfer1::DataType type_; + nvinfer1::TensorLocation location_; + float dynamic_range_; +}; + +TEST(TRT_ShapedWeights_Test, Basic) { + // Test constructor with no arguments. + { + TRT_ShapedWeights weights; + TRT_ShapedWeights copy(weights); + for (auto ptr : {&weights, ©}) { + nvinfer1::Weights trt_weights = ptr->GetTrtWeights(); + EXPECT_EQ(nvinfer1::DataType::kFLOAT, trt_weights.type); + EXPECT_EQ(nullptr, trt_weights.values); + EXPECT_EQ(0, trt_weights.count); + + EXPECT_EQ(nullptr, ptr->GetValues()); + EXPECT_EQ(0, ptr->count()); + EXPECT_EQ(0, ptr->size_bytes()); + } + } + // Test constructor with DataType argument. + { + TRT_ShapedWeights weights(DT_FLOAT); + TRT_ShapedWeights copy(weights); + for (auto ptr : {&weights, ©}) { + nvinfer1::Weights trt_weights = ptr->GetTrtWeights(); + EXPECT_EQ(nvinfer1::DataType::kFLOAT, trt_weights.type); + EXPECT_EQ(nullptr, trt_weights.values); + EXPECT_EQ(0, trt_weights.count); + + EXPECT_EQ(nullptr, ptr->GetValues()); + EXPECT_EQ(0, ptr->count()); + EXPECT_EQ(0, ptr->size_bytes()); + } + } + // Test constructor with DataType and nvinfer1::Dims arguments. + { + TrtWeightStore store; + TRT_ShapedWeights weights = + store.GetTempWeights(DT_FLOAT, GetTestDims({2, 5})); + TRT_ShapedWeights copy(weights); + for (auto ptr : {&weights, ©}) { + nvinfer1::Weights trt_weights = ptr->GetTrtWeights(); + EXPECT_EQ(nvinfer1::DataType::kFLOAT, trt_weights.type); + EXPECT_NE(nullptr, trt_weights.values); + EXPECT_EQ(10, trt_weights.count); + + EXPECT_EQ(trt_weights.values, ptr->GetValues()); + EXPECT_EQ(10, ptr->count()); + EXPECT_EQ(40, ptr->size_bytes()); + } + // Test that it doesn't copy the underlying buffer. + EXPECT_EQ(weights.GetValues(), copy.GetValues()); + } +} + +TEST(TRT_TensorOrWeights_Test, Basic) { + // Test constructor with no arguments. + { + TRT_TensorOrWeights tw; + TRT_TensorOrWeights copy(tw); + TRT_TensorOrWeights assigned; + assigned = tw; + for (auto ptr : {&tw, ©, &assigned}) { + EXPECT_EQ(false, ptr->is_tensor()); + EXPECT_EQ(false, ptr->is_weights()); + EXPECT_EQ(-1, ptr->batch_size()); + } + } + + // Test constructor with ITensor and batch size argument. + { + nvinfer1::Dims dims; + dims.nbDims = 1; + dims.d[0] = 1; + FakeITensor itensor(dims); + TRT_TensorOrWeights tw(&itensor); + TRT_TensorOrWeights tw1(&itensor, /*batch_size=*/1); + + for (auto original_ptr : {&tw, &tw1}) { + TRT_TensorOrWeights copy(*original_ptr); + TRT_TensorOrWeights assigned; + assigned = *original_ptr; + + for (auto ptr : {original_ptr, ©, &assigned}) { + EXPECT_EQ(true, ptr->is_tensor()); + EXPECT_EQ(false, ptr->is_weights()); + if (original_ptr == &tw) { + EXPECT_EQ(-1, ptr->batch_size()); + } else { + EXPECT_EQ(1, ptr->batch_size()); + } + EXPECT_EQ(&itensor, ptr->tensor()); + ExpectTrtDimsEqualsArray({1}, ptr->GetTrtDims()); + } + } + } + // Test constructor which creates and owns an ITensor. + { + nvinfer1::Dims dims; + dims.nbDims = 1; + dims.d[0] = 1; + TRT_TensorOrWeights tw(nvinfer1::DataType::kFLOAT, dims, /*batch_size=*/1); + TRT_TensorOrWeights copy(tw); + TRT_TensorOrWeights assigned; + assigned = tw; + + for (auto ptr : {&tw, ©, &assigned}) { + EXPECT_EQ(true, ptr->is_tensor()); + EXPECT_EQ(false, ptr->is_weights()); + EXPECT_EQ(1, ptr->batch_size()); + EXPECT_NE(nullptr, ptr->tensor()); + ExpectTrtDimsEqualsArray({1}, ptr->GetTrtDims()); + } + } + // Test constructor with TRT_ShapedWeights argument. + { + TRT_ShapedWeights weights; + TRT_TensorOrWeights tw(weights); + TRT_TensorOrWeights copy(tw); + TRT_TensorOrWeights assigned; + assigned = tw; + for (auto ptr : {&tw, ©, &assigned}) { + EXPECT_EQ(false, ptr->is_tensor()); + EXPECT_EQ(true, ptr->is_weights()); + EXPECT_TRUE(TrtShapedWeightsEquals(weights, ptr->weights())); + ExpectTrtDimsEqualsArray({}, ptr->GetTrtDims()); + } + } +} + +class ValidatorTest : public ::testing::Test { + public: + void AddOpValidator(const string& op_name, OpConverter op_validator) { + validator_.op_validators_[op_name] = op_validator; + } + + Status ConvertToTensorOrWeights( + const NodeDef& node_def, int output_port, + const grappler::GraphProperties& graph_properties, + TRT_TensorOrWeights* tensor_or_weights) { + return validator_.ConvertToTensorOrWeights( + node_def, output_port, graph_properties, tensor_or_weights); + } + + protected: + TrtNodeValidator validator_; +}; + +TEST_F(ValidatorTest, ConvertToTensorOrWeights) { + // Convert Const. + { + NodeDef node_def = MakeConstNodeDef("my_const", {1.0f, 2.0f}); + TRT_TensorOrWeights output; + grappler::GrapplerItem item; + grappler::GraphProperties graph_properties(item); + ExpectStatus(ConvertToTensorOrWeights(node_def, /*output_port=*/0, + graph_properties, &output)); + ValidateWeights(output.weights(), {2}, {1.0, 2.0}); + } + + // Helper method to run ConvertToTensorOrWeights() with predefined parameters. + auto convert_to_tensor_or_weights = [this](const std::vector& dims, + TRT_TensorOrWeights* output) { + Scope s = Scope::NewRootScope(); + const auto attrs = ops::Placeholder::Shape(PartialTensorShape{dims}); + auto feed = ops::Placeholder(s.WithOpName("feed"), DT_FLOAT, attrs); + auto add = ops::Add(s.WithOpName("add"), feed, feed); + + grappler::GrapplerItem item; + TF_EXPECT_OK(s.ToGraphDef(&item.graph)); + grappler::GraphProperties graph_properties(item); + TF_EXPECT_OK(graph_properties.InferStatically(true)); + const NodeDef& node_def = add.operation.node()->def(); + return this->ConvertToTensorOrWeights(node_def, /*output_port=*/0, + graph_properties, output); + }; + // Convert non-Const with #dims > nvinfer1::Dims::MAX_DIMS+1. + { + TRT_TensorOrWeights output; + ExpectStatus( + convert_to_tensor_or_weights( + std::vector(nvinfer1::Dims::MAX_DIMS + 2, 1), &output), + error::OUT_OF_RANGE, "Input tensor rank is greater than 9"); + } + // Convert non-Const with #dims < 2. + { + TRT_TensorOrWeights output; + ExpectStatus( + convert_to_tensor_or_weights({1}, &output), error::INVALID_ARGUMENT, + "Input tensor with rank<2 is not supported since the first dimension " + "is treated as batch dimension by TRT"); + } + // Convert non-Const. We test the case where the non-batch dimemsion is + // unknown as well, to make sure the validator allows that. + for (const int32 non_batch_dim : {-1, 2}) { + const int32 batch_size = 12; + TRT_TensorOrWeights output; + ExpectStatus( + convert_to_tensor_or_weights({batch_size, non_batch_dim}, &output)); + EXPECT_EQ(true, output.is_tensor()); + EXPECT_EQ(batch_size, output.batch_size()); + EXPECT_NE(nullptr, output.tensor()); + ExpectTrtDimsEqualsArray({non_batch_dim}, output.GetTrtDims()); + } +} + +TEST_F(ValidatorTest, ValidateNode) { + grappler::GrapplerItem item; + grappler::GraphProperties graph_properties(item); + + bool start_conversion = false; + bool should_fail = false; + auto op_converter = [&start_conversion, + &should_fail](OpConverterParams* params) -> Status { + if (should_fail) return errors::InvalidArgument(""); + if (!params->validation_only) start_conversion = true; + return Status::OK(); + }; + NodeDef node_def = MakeNodeDef("my_op", "MyOp", {}); + + // Validator not registered, validation should pass. + TF_EXPECT_OK(validator_.ValidateNode(node_def, {}, graph_properties)); + + // Register validator. + AddOpValidator("MyOp", op_converter); + TF_EXPECT_OK(validator_.ValidateNode(node_def, {}, graph_properties)); + EXPECT_EQ(false, start_conversion); + + // Let the converter return error. + should_fail = true; + ExpectStatus(validator_.ValidateNode(node_def, {}, graph_properties), + error::INVALID_ARGUMENT); +} + +class ConverterTest : public ::testing::Test { + public: + ConverterTest() { + builder_.reset(nvinfer1::createInferBuilder(logger_)); + network_.reset(builder_->createNetwork()); + converter_.reset(new Converter(network_.get(), + /*precision_mode=*/FP32MODE, + /*use_calibration=*/false)); + weight_store_ = &converter_->weight_store_; + } + + void AddOpConverter(const string& op_name, OpConverter op_converter) { + converter_->op_registry_[op_name] = op_converter; + } + + // Below we expose private methods of Converter for testing. + + Status MaybeUpdateBatchSize(int batch_size) { + return converter_->MaybeUpdateBatchSize(batch_size); + } + + Status AddTensorOrWeights(const string& name, TRT_TensorOrWeights input) { + return converter_->AddTensorOrWeights(name, input); + } + + Status GetTensorOrWeights(const string& name, TRT_TensorOrWeights* output) { + return converter_->GetTensorOrWeights(name, output); + } + + Status GetInputs(const NodeDef& node_def, + std::vector* inputs) const { + return converter_->GetInputs(node_def, inputs); + } + + Status GetWeightRange(const TRT_ShapedWeights& weights, float* out_min, + float* out_max) const { + return converter_->GetWeightRange(weights, out_min, out_max); + } + + void PropagateQuantizationRanges() { + converter_->PropagateQuantizationRanges(); + } + + int batch_size() const { return converter_->batch_size_; } + + std::unordered_map& quantization_ranges() { + return converter_->quantization_ranges_; + } + + private: + Logger logger_; + // These members are ordered in a way such that the destruction order is: + // converter_ -> network_ -> builder_ + TrtUniquePtrType builder_; + TrtUniquePtrType network_; + + protected: + std::unique_ptr converter_; + TrtWeightStore* weight_store_; +}; + +TEST_F(ConverterTest, ConvertNode) { + FakeITensor output_tensors[2]; + auto op_converter = [&output_tensors](OpConverterParams* params) -> Status { + nvinfer1::Dims dims = params->inputs[0].tensor()->getDimensions(); + for (int i = 0; i < 2; ++i) { + dims.d[0] += 1; + output_tensors[i].setDimensions(dims); + params->outputs->push_back(TRT_TensorOrWeights(&output_tensors[i])); + } + return Status::OK(); + }; + NodeDef node_def = MakeNodeDef("my_op", "MyOp", {"my_input"}); + TF_EXPECT_OK(converter_->AddInputTensor( + "my_input", nvinfer1::DataType::kFLOAT, GetTestDims({123}), 1)); + + // Converter not registered. + ExpectStatus(converter_->ConvertNode(node_def), error::UNIMPLEMENTED, + "No converter registered for op: MyOp"); + + // Register the converter and retry. + AddOpConverter("MyOp", op_converter); + TF_EXPECT_OK(converter_->ConvertNode(node_def)); + + TRT_TensorOrWeights actual_output_1; + TF_EXPECT_OK(GetTensorOrWeights("my_op", &actual_output_1)); + EXPECT_EQ(&output_tensors[0], actual_output_1.tensor()); + EXPECT_EQ(124, actual_output_1.tensor()->getDimensions().d[0]); + + TRT_TensorOrWeights actual_output_2; + TF_EXPECT_OK(GetTensorOrWeights("my_op:1", &actual_output_2)); + EXPECT_EQ(&output_tensors[1], actual_output_2.tensor()); + EXPECT_EQ(125, actual_output_2.tensor()->getDimensions().d[0]); +} + +TEST_F(ConverterTest, AddAndGetInputs) { + NodeDef node_def; + node_def.add_input("^control_input"); + node_def.add_input("input"); + node_def.add_input("input:0"); + node_def.add_input("input:1"); + node_def.add_input("weird_input:2:3:4:0"); + + TF_EXPECT_OK(converter_->AddInputTensor("input", nvinfer1::DataType::kFLOAT, + GetTestDims({1}), 1)); + TF_EXPECT_OK(converter_->AddInputTensor("input:1", nvinfer1::DataType::kINT32, + GetTestDims({2, 3}), 1)); + TF_EXPECT_OK(converter_->AddInputTensor( + "weird_input:2:3:4", nvinfer1::DataType::kHALF, GetTestDims({5, 3}), 1)); + + std::vector inputs; + TF_EXPECT_OK(GetInputs(node_def, &inputs)); + + EXPECT_EQ(4, inputs.size()); + EXPECT_EQ(inputs[0].tensor(), inputs[1].tensor()); + + EXPECT_EQ(nvinfer1::DataType::kFLOAT, inputs[0].tensor()->getType()); + EXPECT_EQ(nvinfer1::DataType::kINT32, inputs[2].tensor()->getType()); + EXPECT_EQ(nvinfer1::DataType::kHALF, inputs[3].tensor()->getType()); + ExpectTrtDimsEqualsArray({1}, inputs[0].tensor()->getDimensions()); + ExpectTrtDimsEqualsArray({2, 3}, inputs[2].tensor()->getDimensions()); + ExpectTrtDimsEqualsArray({5, 3}, inputs[3].tensor()->getDimensions()); +} + +TEST_F(ConverterTest, RenameAndMarkOutputTensors) { + // Test that the tensor are actually named and marked as output after + // Converter::RenameAndMarkOutputTensors() is called. + + // Register a custom converter which shuffles the input. We use it to build a + // TRT network whose output will be later marked. + std::vector output_tensors; + auto op_converter = [&output_tensors](OpConverterParams* params) -> Status { + nvinfer1::Permutation perm; + perm.order[0] = 1; + perm.order[1] = 0; + for (int i = 0; i < 2; ++i) { + nvinfer1::ITensor* input_tensor = + const_cast(params->inputs[0].tensor()); + nvinfer1::IShuffleLayer* layer = + params->converter->network()->addShuffle(*input_tensor); + layer->setFirstTranspose(perm); + nvinfer1::ITensor* output_tensor = layer->getOutput(0); + params->outputs->emplace_back(output_tensor); + output_tensors.push_back(output_tensor); + } + TRT_ShapedWeights output_weights(DT_FLOAT); + params->outputs->emplace_back(output_weights); + return Status::OK(); + }; + AddOpConverter("MyOp", op_converter); + + // Run the conversion. + NodeDef node_def = MakeNodeDef("my_op", "MyOp", {"my_input"}); + TF_EXPECT_OK(converter_->AddInputTensor( + "my_input", nvinfer1::DataType::kFLOAT, GetTestDims({1, 2}), 1)); + TF_EXPECT_OK(converter_->ConvertNode(node_def)); + + // Mark a weight as output, should fail. + ExpectStatus( + converter_->RenameAndMarkOutputTensors({{"my_op:2", "my_output"}}), + error::INVALID_ARGUMENT, "Output my_op:2 is weights not tensor"); + + // Mark tensors as output, should pass. + TF_EXPECT_OK(converter_->RenameAndMarkOutputTensors( + {{"my_op", "my_output"}, {"my_op:1", "my_output_1"}})); + EXPECT_EQ(2, output_tensors.size()); + for (auto output_tensor : output_tensors) { + ExpectTrtDimsEqualsArray({2, 1}, output_tensor->getDimensions()); + } + EXPECT_EQ("my_output", string(output_tensors[0]->getName())); + EXPECT_EQ("my_output_1", string(output_tensors[1]->getName())); +} + +TEST_F(ConverterTest, TransposeTensor) { + nvinfer1::ITensor* input_tensor = converter_->network()->addInput( + "", nvinfer1::DataType::kFLOAT, GetTestDims({2, 3, 5})); + const nvinfer1::ITensor* output_tensor = nullptr; + + // Rank doesn't match. + ExpectStatus( + converter_->TransposeTensor(input_tensor, {0, 1}, &output_tensor), + error::INVALID_ARGUMENT, + "Rank of perm for transpose does not match with that of the input"); + + // Transpose at batch dimension. + ExpectStatus( + converter_->TransposeTensor(input_tensor, {1, 0, 2, 3}, &output_tensor), + error::UNIMPLEMENTED, "Transpose at batch dimension is not supported."); + + // OK. + TF_EXPECT_OK( + converter_->TransposeTensor(input_tensor, {0, 3, 1, 2}, &output_tensor)); + ExpectTrtDimsEqualsArray({5, 2, 3}, output_tensor->getDimensions()); +} + +TEST_F(ConverterTest, PrepareTensorForShape_Tensor) { + nvinfer1::ITensor* input_tensor = converter_->network()->addInput( + "", nvinfer1::DataType::kFLOAT, GetTestDims({2, 3, 5})); + TRT_TensorOrWeights tw(input_tensor); + const nvinfer1::ITensor* output_tensor = nullptr; + + // Shape size doesn't match. + ExpectStatus(converter_->PrepareTensorForShape(tw, GetTestDims({2, 3, 6}), + &output_tensor), + error::INVALID_ARGUMENT, "Reshape shapes are not compatible"); + + // TODO(aaroey): we should check the case where uninferred dimensions are not + // an exact divisor of input dim ensions, e.g. for dims {-1, 7}. + + // Infer shape, ok. + TF_EXPECT_OK(converter_->PrepareTensorForShape(tw, GetTestDims({-1, 2}), + &output_tensor)); + ExpectTrtDimsEqualsArray({15, 2}, output_tensor->getDimensions()); + + // Regular shape. + TF_EXPECT_OK(converter_->PrepareTensorForShape(tw, GetTestDims({10, 3}), + &output_tensor)); + ExpectTrtDimsEqualsArray({10, 3}, output_tensor->getDimensions()); +} + +TEST_F(ConverterTest, PrepareTensorForShape_Weights) { + TRT_ShapedWeights weights = + weight_store_->GetTempWeights(DT_FLOAT, GetTestDims({2, 3, 5})); + TRT_TensorOrWeights tw(weights); + const nvinfer1::ITensor* output_tensor = nullptr; + TF_EXPECT_OK(converter_->PrepareTensorForShape(tw, GetTestDims({10, 3}), + &output_tensor)); + ExpectTrtDimsEqualsArray({10, 3}, output_tensor->getDimensions()); +} + +TEST_F(ConverterTest, MaybeUpdateBatchSize) { + EXPECT_EQ(-1, batch_size()); + + TF_EXPECT_OK(MaybeUpdateBatchSize(-1)); + EXPECT_EQ(-1, batch_size()); + + TF_EXPECT_OK(MaybeUpdateBatchSize(123)); + EXPECT_EQ(123, batch_size()); + + TF_EXPECT_OK(MaybeUpdateBatchSize(123)); + EXPECT_EQ(123, batch_size()); + + TF_EXPECT_OK(MaybeUpdateBatchSize(-1)); + EXPECT_EQ(123, batch_size()); + + ExpectStatus(MaybeUpdateBatchSize(124), error::INVALID_ARGUMENT, + "Provided batch size does not match converter batch size"); +} + +TEST_F(ConverterTest, AddAndGetTensorOrWeights) { + // Add a tensor. + FakeITensor fake_tensor; + TRT_TensorOrWeights tensor(&fake_tensor); + EXPECT_EQ(-1, tensor.batch_size()); + TF_EXPECT_OK(MaybeUpdateBatchSize(123)); + TF_EXPECT_OK(AddTensorOrWeights("my_tensor", tensor)); + + // Get the added tensor. + TRT_TensorOrWeights added_tensor; + TF_EXPECT_OK(GetTensorOrWeights("my_tensor", &added_tensor)); + EXPECT_EQ(123, added_tensor.batch_size()); + + // Add the same tensor again. + ExpectStatus(AddTensorOrWeights("my_tensor", tensor), error::ALREADY_EXISTS, + "tensor/weights my_tensor already exist"); +} + +template +void TestGetWeightRange(ConverterTest* test, TrtWeightStore* weight_store) { + TRT_ShapedWeights weights = + weight_store->GetTempWeights(DataTypeToEnum::v(), GetTestDims({2, 3})); + const std::vector values = {T(3), T(1), T(2), T(6), T(5), T(4)}; + memcpy(const_cast(weights.GetValues()), values.data(), + weights.size_bytes()); + + float out_min = 0.0f; + float out_max = 0.0f; + TF_EXPECT_OK(test->GetWeightRange(weights, &out_min, &out_max)); + EXPECT_EQ(1.0f, out_min); + EXPECT_EQ(6.0f, out_max); +} + +TEST_F(ConverterTest, GetWeightRange) { + TestGetWeightRange(this, weight_store_); + TestGetWeightRange(this, weight_store_); + TestGetWeightRange(this, weight_store_); +} + +TEST_F(ConverterTest, ProvideQuantizationRange) { + FakeITensor fake_tensor; + // Assymetric range + converter_->ProvideQuantizationRange(&fake_tensor, 0.0f, 6.0f); + EXPECT_EQ(6.0f, quantization_ranges()[&fake_tensor]); + converter_->ProvideQuantizationRange(&fake_tensor, 1.0f, 6.0f); + EXPECT_EQ(6.0f, quantization_ranges()[&fake_tensor]); + converter_->ProvideQuantizationRange(&fake_tensor, -8.0f, 6.0f); + EXPECT_EQ(8.0f, quantization_ranges()[&fake_tensor]); + converter_->ProvideQuantizationRange(&fake_tensor, -8.123f, -6.123f); + EXPECT_EQ(8.123f, quantization_ranges()[&fake_tensor]); + // Symmetric range + converter_->ProvideQuantizationRange(&fake_tensor, -6.123f, 6.123f); + EXPECT_EQ(6.123f, quantization_ranges()[&fake_tensor]); +} + +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, + /*use_calibration=*/true); + int8_converter.ProvideQuantizationRange(&input, -5.0f, 5.0f); + int8_converter.ProvideQuantizationRange(¬_infer, -100.0f, 100.0f); + int8_converter.MarkQuantizationRangesAsInferrable(&input, &infer_1); + int8_converter.MarkQuantizationRangesAsInferrable(&infer_1, &infer_2); + int8_converter.MarkQuantizationRangesAsInferrable(&infer_2, &infer_3); + + // Input range should be inferred along the chain and applied to tensors. + int8_converter.MaybeApplyQuantizationRanges(); +#if NV_TENSORRT_MAJOR >= 5 + EXPECT_EQ(input.getDynamicRange(), 5.0f); + EXPECT_EQ(infer_1.getDynamicRange(), 5.0f); + EXPECT_EQ(infer_2.getDynamicRange(), 5.0f); + EXPECT_EQ(infer_3.getDynamicRange(), 5.0f); + EXPECT_EQ(not_infer.getDynamicRange(), 100.0f); +#endif +} + +TEST_F(ConverterTest, PropagateQuantizationRanges) { + // infer0 <-> infer1 <-> infer2 <-> infer3 + // | + // infer4 <-> infer5 + FakeITensor infer[6]; + FakeITensor not_infer; + converter_->ProvideQuantizationRange(&infer[4], -5.0f, 5.0f); + converter_->MarkQuantizationRangesAsInferrable(&infer[0], &infer[1]); + converter_->MarkQuantizationRangesAsInferrable(&infer[1], &infer[2]); + converter_->MarkQuantizationRangesAsInferrable(&infer[3], &infer[2]); + converter_->MarkQuantizationRangesAsInferrable(&infer[4], &infer[1]); + converter_->MarkQuantizationRangesAsInferrable(&infer[4], &infer[5]); + + // Input range should be inferred along the chain. + PropagateQuantizationRanges(); + auto ranges = quantization_ranges(); + for (int i = 0; i < 6; ++i) { + EXPECT_EQ(5.0f, ranges[&infer[i]]); + } + EXPECT_EQ(ranges.count(¬_infer), 0); +} + +TEST_F(ConverterTest, GetTrtBroadcastShape) { + const bool kIsTensor = true; + const bool kIsNotTensor = false; + auto symmetric_test = [this](const std::vector& operand_1_shape, + const std::vector& operand_2_shape, + const bool operand_1_is_tensor, + const bool operand_2_is_tensor, + const std::vector& expected_operand_1_shape, + const std::vector& expected_operand_2_shape, + error::Code expected_code = error::OK, + const char* expected_error_msg_substr = nullptr, + const int operand_1_batch_size = -1, + const int operand_2_batch_size = -1) { + auto create_tensor_or_weights = [](const std::vector& shape, + bool is_tensor, int batch_size = -1) { + if (is_tensor) { + return TRT_TensorOrWeights{nvinfer1::DataType::kFLOAT, + GetTestDims(shape), batch_size}; + } + TRT_ShapedWeights weights; + weights.shape_ = GetTestDims(shape); + return TRT_TensorOrWeights(weights); + }; + + nvinfer1::Dims operand_1_new_dims, operand_2_new_dims; + TRT_TensorOrWeights operand_1 = create_tensor_or_weights( + operand_1_shape, operand_1_is_tensor, operand_1_batch_size); + TRT_TensorOrWeights operand_2 = create_tensor_or_weights( + operand_2_shape, operand_2_is_tensor, operand_2_batch_size); + + // operand_1 broadcast operand_2 + ExpectStatus( + this->converter_->GetTrtBroadcastShape( + operand_1, operand_2, &operand_1_new_dims, &operand_2_new_dims), + expected_code, expected_error_msg_substr); + if (expected_code == error::OK) { + ExpectTrtDimsEqualsArray(expected_operand_1_shape, operand_1_new_dims); + ExpectTrtDimsEqualsArray(expected_operand_2_shape, operand_2_new_dims); + } + // operand_2 broadcast operand_1 + ExpectStatus( + this->converter_->GetTrtBroadcastShape( + operand_2, operand_1, &operand_2_new_dims, &operand_1_new_dims), + expected_code, expected_error_msg_substr); + if (expected_code == error::OK) { + ExpectTrtDimsEqualsArray(expected_operand_1_shape, operand_1_new_dims); + ExpectTrtDimsEqualsArray(expected_operand_2_shape, operand_2_new_dims); + } + }; + + // Both inputs are weights. + symmetric_test( + {1}, {1}, kIsNotTensor, kIsNotTensor, {}, {}, error::INVALID_ARGUMENT, + "Broadcasting requires at least one of the operands be tensors"); + + // One tensor and one weights. + symmetric_test({1, 1, 1}, {2}, kIsTensor, kIsNotTensor, {1, 1, 1}, {1, 1, 2}); + symmetric_test({1, 1, 2}, {2}, kIsTensor, kIsNotTensor, {1, 1, 2}, {1, 1, 2}); + symmetric_test({1, 3, 2}, {1}, kIsTensor, kIsNotTensor, {1, 3, 2}, {1, 1, 1}); + symmetric_test({1, 1, 1}, {2, 3}, kIsTensor, kIsNotTensor, {1, 1, 1}, + {1, 2, 3}); + symmetric_test({1, 1, 1}, {2, 3, 4}, kIsTensor, kIsNotTensor, {1, 1, 1}, + {2, 3, 4}); + symmetric_test({1, 1, 1}, {1, 2, 3, 4}, kIsTensor, kIsNotTensor, {1, 1, 1}, + {2, 3, 4}); + symmetric_test({1, 3, 4}, {1, 2, 1, 4}, kIsTensor, kIsNotTensor, {1, 3, 4}, + {2, 1, 4}); + symmetric_test({1, 1, 1}, {2, 1, 1, 1}, kIsTensor, kIsNotTensor, {}, {}, + error::INVALID_ARGUMENT, "Infeasible broadcast scheme"); + symmetric_test({1, 1, 1}, {2, 1, 1, 1}, kIsTensor, kIsNotTensor, {}, {}, + error::INVALID_ARGUMENT, "Infeasible broadcast scheme", + /*operand_1_batch_size=*/2); + symmetric_test({1, 1, 1}, {1, 1, 1, 1, 1}, kIsTensor, kIsNotTensor, {}, {}, + error::INVALID_ARGUMENT, + "Broadcasting beyond batch dimension is not supported " + "(tensor #dims 4 vs broadcast #dims 5)"); + + // Both inputs are tensors. + symmetric_test({1, 1, 1}, {1, 1}, kIsTensor, kIsTensor, {}, {}, + error::INVALID_ARGUMENT, + "Broadcasting beyond batch dimension is not supported " + "(tensor #dims 3 vs broadcast #dims 4)"); + symmetric_test({1, 3, 4}, {2, 1, 4}, kIsTensor, kIsTensor, {1, 3, 4}, + {2, 1, 4}); + symmetric_test({1, 1, 1}, {1, 1, 1, 1}, kIsTensor, kIsTensor, {}, {}, + error::INVALID_ARGUMENT, + "Broadcasting beyond batch dimension is not supported " + "(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 to test various op converters, using both a TrtNodeValidator and +// Converter. +class OpConverterTest : public ::testing::Test { + public: + OpConverterTest() : scope_(Scope::NewRootScope()) { + QCHECK_EQ(0, cudaStreamCreate(&stream_)); + Reset(); + } + + ~OpConverterTest() override { QCHECK_EQ(0, cudaStreamDestroy(stream_)); } + + Status GetTensorOrWeights(const string& name, TRT_TensorOrWeights* output) { + return converter_->GetTensorOrWeights(name, output); + } + + void Reset() { + validator_.reset(nullptr); + converter_.reset(nullptr); + + // Reset the INetworkDefinition. + engine_.reset(nullptr); + network_.reset(nullptr); + builder_.reset(nvinfer1::createInferBuilder(logger_)); + network_.reset(builder_->createNetwork()); + builder_->setMaxBatchSize(1); + + // Reset the validator and converter. + validator_.reset(new TrtNodeValidator); + converter_.reset(new Converter(network_.get(), + /*precision_mode=*/FP32MODE, + /*use_calibration=*/false)); + + // Reset other related artifacts. + scope_ = Scope::NewRootScope(); + validator_inputs_.clear(); + } + + // TODO(laigd): test fp16 and int8 support. + template + void BuildAndRun( + const std::vector>>& + input_data, + const char* output_name, std::vector* output_data) { + // Mark the output tensor as TRT engine output. + TF_EXPECT_OK(converter_->RenameAndMarkOutputTensors( + {{string(output_name), string(output_name)}})); + + // Build the TRT engine. + ASSERT_EQ(nullptr, engine_.get()); + engine_.reset(builder_->buildCudaEngine(*converter_->network())); + 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 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); + + 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_)); + cudaStreamSynchronize(stream_); + + for (int i = 0; i < input_data.size() + 1; ++i) { + ASSERT_EQ(0, cudaFree(buffers[i])); + } + } + + bool HasStaticShape(const nvinfer1::Dims& dims) const { + if (dims.nbDims < 0) return false; + for (int i = 0; i < dims.nbDims; ++i) { + if (dims.d[i] < 0) return false; + } + return true; + } + + // Add ITensor for both validation and conversion. + void AddTestTensor( + const char* name, const std::vector& dims, int batch_size = 1, + nvinfer1::DataType trt_dtype = nvinfer1::DataType::kFLOAT) { + DataType tf_dtype = TrtDataTypeToTf(trt_dtype); + ops::Placeholder::Attrs attrs; + TF_EXPECT_OK(TensorShapeUtils::MakeShape(dims, &attrs.shape_)); + attrs.shape_.InsertDim(0, batch_size); + auto input = ops::Placeholder(scope_.WithOpName(name), tf_dtype, attrs); + validator_inputs_[name] = input.operation.node()->def(); + + // Add a real ITensor for conversion conditionally. + const nvinfer1::Dims trt_dims = GetTestDims(dims); + if (HasStaticShape(trt_dims)) { + TF_EXPECT_OK( + converter_->AddInputTensor(name, trt_dtype, trt_dims, batch_size)); + ASSERT_EQ(batch_size, converter_->batch_size_); + } + } + + // Add weights for both validation and conversion. + template + void AddTestWeights(const char* name, const std::vector& dims, + const std::vector& values) { + const DataType dtype = DataTypeToEnum::v(); + const nvinfer1::Dims trt_dims = GetTestDims(dims); + const int64_t num_elements = TrtDimsNumElements(trt_dims); + QCHECK_EQ(num_elements, values.size()) + << num_elements << " vs " << values.size(); + TRT_ShapedWeights weights(dtype); + if (num_elements) { + weights = converter_->weight_store_.GetTempWeights(dtype, trt_dims); + QCHECK_EQ(weights.size_bytes(), sizeof(T) * values.size()) + << weights.size_bytes() << " vs " << sizeof(T) * values.size(); + memcpy(const_cast(weights.GetValues()), values.data(), + weights.size_bytes()); + } + // Add weights for validation. + TensorShape shape; + TF_EXPECT_OK(TensorShapeUtils::MakeShape(dims, &shape)); + validator_inputs_[name] = MakeConstNodeDef(name, values, shape); + // Add weights for conversion. + TF_EXPECT_OK( + converter_->AddTensorOrWeights(name, TRT_TensorOrWeights{weights})); + } + + // Test validation in validation-only mode. + void RunValidation(const NodeDef& node_def, + error::Code expected_code = error::OK, + const char* expected_msg_substr = nullptr) { + std::vector> input_node_and_ports; + for (const string& input : node_def.input()) { + input_node_and_ports.emplace_back(&validator_inputs_[input], 0); + } + grappler::GrapplerItem item; + TF_EXPECT_OK(scope_.ToGraphDef(&item.graph)); + grappler::GraphProperties graph_properties(item); + TF_EXPECT_OK(graph_properties.InferStatically(true)); + + ExpectStatus(validator_->ValidateNode(node_def, input_node_and_ports, + graph_properties), + expected_code, expected_msg_substr); + } + + void RunConversion(const NodeDef& node_def, + error::Code expected_code = error::OK, + const char* expected_msg_substr = nullptr) { + ExpectStatus(converter_->ConvertNode(node_def), expected_code, + expected_msg_substr); + } + + // Helper method to run both validation and conversion, when the expected + // output are same. + void RunValidationAndConversion(const NodeDef& node_def, + error::Code expected_code = error::OK, + const char* expected_msg_substr = nullptr, + bool should_run_conversion = true) { + RunValidation(node_def, expected_code, expected_msg_substr); + if (should_run_conversion) { + RunConversion(node_def, expected_code, expected_msg_substr); + } + } + + // Expose quantization_ranges_ for tests + std::unordered_map& quantization_ranges() { + return converter_->quantization_ranges_; + } + + std::unique_ptr converter_; + std::unique_ptr validator_; + + private: + Logger logger_; + TrtUniquePtrType builder_; + TrtUniquePtrType network_; + TrtUniquePtrType engine_; + cudaStream_t stream_; + // Used to create placeholders with shape and data type information. The + // created placeholders will be used as inputs to the node to be verified, + // thus we need the shape and data type information to get a non-empty + // GraphProperties. + // TODO(laigd): consider use this Scope to create the NodeDef to verify. + Scope scope_; + 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; + node_def.set_name("my_const"); + node_def.set_op("Const"); + + auto reset_and_test = [&node_def, test]( + const Tensor& tensor, const bool as_tensor_content, + const std::vector& expected_dims, + const std::vector& expected_value) { + test->Reset(); + + TensorProto* tensor_attr = + (*node_def.mutable_attr())["value"].mutable_tensor(); + tensor_attr->Clear(); + + if (as_tensor_content) { + tensor.AsProtoTensorContent(tensor_attr); + } else { + 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; + TF_EXPECT_OK(test->GetTensorOrWeights("my_const", &output)); + ValidateWeights(output.weights(), expected_dims, expected_value); + }; + + auto& attr = *node_def.mutable_attr(); + attr["dtype"].set_type(dtype); + { + // By default empty tensor will pick DT_FLOAT as data type and we fix it + // here. + Tensor t(dtype); // Empty tensor. + reset_and_test(t, false, {}, {}); + } + { + Tensor t = ::tensorflow::test::AsScalar(12); + reset_and_test(t, false, {1}, {12}); + reset_and_test(t, true, {1}, {12}); + } + { + Tensor t = ::tensorflow::test::AsTensor({1, 2}); + reset_and_test(t, false, {2}, {1, 2}); + reset_and_test(t, true, {2}, {1, 2}); + } + { + Tensor t = ::tensorflow::test::AsTensor({1, 2, 3, 4, 5, 6}, + TensorShape({2, 3})); + 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) { + { + Reset(); + NodeDef node_def = MakeNodeDef("my_const", "Const", {"input"}); + AddTestTensor("input", {1}); + RunValidationAndConversion( + node_def, error::INVALID_ARGUMENT, + "Constant node is expected to have empty input list: my_const"); + } + { + Reset(); + NodeDef node_def = MakeConstNodeDef("my_const", {}); + RunValidationAndConversion(node_def, error::INVALID_ARGUMENT, + "Unsupported data type double"); + } + + TestConvertConst(this); + TestConvertConst(this); + TestConvertConst(this); +} + +TEST_F(OpConverterTest, ConvertTranspose) { + { + // Input list is empty, should fail. + NodeDef node_def = MakeNodeDef("my_transpose", "Transpose", {}); + RunValidationAndConversion( + node_def, error::INVALID_ARGUMENT, + "Input expects tensor and weights, at my_transpose"); + } + + // Get the NodeDef for Transpose. + Scope s = Scope::NewRootScope(); + auto input = ops::Placeholder(s.WithOpName("input"), DT_FLOAT); + auto weights = ops::Placeholder(s.WithOpName("weights"), DT_INT32); + auto transpose = ops::Transpose(s.WithOpName("my_transpose"), input, weights); + const NodeDef& node_def = transpose.operation.node()->def(); + + { + // Permutation is a tensor, should fail. + Reset(); + AddTestTensor("input", {1, 2, 3}); + AddTestTensor("weights", {3}); + RunValidationAndConversion( + node_def, error::INVALID_ARGUMENT, + "Input expects tensor and weights, at my_transpose"); + } + { + // Transpose at batch dimension, should fail. + Reset(); + AddTestTensor("input", {1, 2, 3}); + AddTestWeights("weights", {4}, {1, 0, 2, 3}); + RunValidationAndConversion(node_def, error::UNIMPLEMENTED, + "Transpose at batch dimension is not supported"); + } + { + // Permutation rank doesn't match, should fail. + Reset(); + AddTestTensor("input", {1, 2, 3}); + AddTestWeights("weights", {3}, {0, 1, 2}); + RunValidationAndConversion( + node_def, error::INVALID_ARGUMENT, + "Rank of perm for transpose does not match with that of the input."); + } + { + // Ok. + Reset(); + AddTestTensor("input", {1, 2, 3}); + AddTestWeights("weights", {4}, {0, 3, 1, 2}); + RunValidationAndConversion(node_def); + TRT_TensorOrWeights output; + TF_EXPECT_OK(GetTensorOrWeights("my_transpose", &output)); + 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)); + } +} + +TEST_F(OpConverterTest, ConvertReshape) { + { + // Input list is empty, should fail. + NodeDef node_def = MakeNodeDef("my_reshape", "Reshape", {}); + RunValidationAndConversion( + node_def, error::INVALID_ARGUMENT, + "Input expects weights for shape, at my_reshape"); + } + + // Get the NodeDef for Reshape. + Scope s = Scope::NewRootScope(); + auto input = ops::Placeholder(s.WithOpName("input"), DT_FLOAT); + auto weights = ops::Placeholder(s.WithOpName("weights"), DT_INT32); + auto reshape = ops::Reshape(s.WithOpName("my_reshape"), input, weights); + const NodeDef& node_def = reshape.operation.node()->def(); + + { + // Shape is a tensor, should fail. + Reset(); + AddTestTensor("input", {1, 2, 3}); + AddTestTensor("weights", {3}); + RunValidationAndConversion( + node_def, error::INVALID_ARGUMENT, + "Input expects weights for shape, at my_reshape"); + } + { + // Reshape to scalar, should fail. + Reset(); + AddTestTensor("input", {1, 2, 3}); + AddTestWeights("weights", {0}, {}); + RunValidationAndConversion( + node_def, error::UNIMPLEMENTED, + "Reshape to shape=[] is not supported, at my_reshape"); + } + + 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; + }; + + // Reshape at batch dimension, should fail. + const int kReshapeBatchDimsCases = 5; + TestParams params[kReshapeBatchDimsCases] = { + TestParams{1, {1, 2, 3}, {3, 1, 1, 2}}, + TestParams{1, {1, 2, -1}, {-1, 1, 1, 2}}, + TestParams{1, {1, 2, 3}, {-1, 1, 1, 2}}, + TestParams{-1, {1, 2, 3}, {1, 1, 1, 2}}, + TestParams{-1, {-1, 2, 3}, {1, 1, 1, 6}}, // TODO(laigd): it should pass. + }; + for (int i = 0; i < kReshapeBatchDimsCases; ++i) { + Reset(); + const std::vector& dims = params[i].tensor_dims; + AddTestTensor("input", dims, params[i].batch_size); + AddTestWeights("weights", {4}, params[i].shape); + RunValidationAndConversion( + node_def, error::UNIMPLEMENTED, + "Reshape on batch dimension is not supported, at my_reshape", + /*should_run_conversion=*/(dims[0] > 0 && dims[1] > 0 && dims[2] > 0)); + } + + // Reshape on non batch dimensions, ok. + const int kReshapeOKCases = 3; + TestParams ok_params[kReshapeOKCases] = { + TestParams{-1, {1, 2, 3}, {-1, 1, 3, 2}}, + TestParams{1, {1, 2, 3}, {-1, 1, 3, 2}}, + TestParams{1, {1, 2, 3}, {1, 1, 3, 2}}, + }; + for (int i = 0; i < kReshapeOKCases; ++i) { + Reset(); + AddTestTensor("input", ok_params[i].tensor_dims, ok_params[i].batch_size); + AddTestWeights("weights", {4}, ok_params[i].shape); + RunValidationAndConversion(node_def); + TRT_TensorOrWeights output; + TF_EXPECT_OK(GetTensorOrWeights("my_reshape", &output)); + 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)); + } +} + +TEST_F(OpConverterTest, ConvertMatMul) { + { + // Input list is empty, should fail. + NodeDef node_def = MakeNodeDef("my_matmul", "MatMul", {}); + RunValidationAndConversion( + node_def, error::INVALID_ARGUMENT, + "Input expects tensor and weights, at my_matmul"); + } + + // Get the NodeDef for MatMul. + auto get_matmul_nodedef = [](DataType dtype, bool transpose_a, + bool transpose_b) -> NodeDef { + Scope s = Scope::NewRootScope(); + auto input = ops::Placeholder(s.WithOpName("input"), dtype); + auto weights = ops::Placeholder(s.WithOpName("weights"), dtype); + const auto matmul_attrs = + ops::MatMul::TransposeA(transpose_a).TransposeB(transpose_b); + auto matmul = + ops::MatMul(s.WithOpName("my_matmul"), input, weights, matmul_attrs); + return matmul.operation.node()->def(); + }; + + { + // Unsupported data type. + Reset(); + NodeDef node_def = get_matmul_nodedef(DT_INT32, false, false); + AddTestTensor("input", {2}, /*batch_size=*/1, nvinfer1::DataType::kINT32); + AddTestWeights("weights", {2, 1}, {3, 5}); + RunValidationAndConversion( + node_def, error::UNIMPLEMENTED, + "Data type is not supported, for node my_matmul got int32"); + } + // transpose_a is set. + for (bool transpose_b : {false, true}) { + Reset(); + NodeDef node_def = + get_matmul_nodedef(DT_FLOAT, /*transpose_a=*/true, transpose_b); + AddTestTensor("input", {2}, /*batch_size=*/1); + AddTestWeights("weights", {2, 2}, {0, 1, 2, 3}); + RunValidationAndConversion( + node_def, error::INVALID_ARGUMENT, + "transpose_a is not supported for TensorRT FullyConnected"); + } + // OK. + for (bool transpose_b : {false, true}) { + Reset(); + NodeDef node_def = + get_matmul_nodedef(DT_FLOAT, /*transpose_a=*/false, transpose_b); + AddTestTensor("input", {2}, /*batch_size=*/1); + AddTestWeights("weights", {2, 2}, {0, 1, 2, 3}); + RunValidationAndConversion(node_def); + TRT_TensorOrWeights output; + TF_EXPECT_OK(GetTensorOrWeights("my_matmul", &output)); + EXPECT_TRUE(output.is_tensor()); + ExpectTrtDimsEqualsArray({2}, output.tensor()->getDimensions()); + + std::vector output_data(2); + BuildAndRun({{"input", {0, 1}}}, "my_matmul", &output_data); + if (transpose_b) { + EXPECT_THAT(output_data, ElementsAre(1, 3)); + } else { + EXPECT_THAT(output_data, ElementsAre(2, 3)); + } + } +} + +template +void TestConvertBiasAdd(OpConverterTest* test) { + // Get the NodeDef for BiasAdd. + auto get_biasadd_nodedef = [](const string& data_format) -> NodeDef { + Scope s = Scope::NewRootScope(); + auto input = ops::Placeholder(s.WithOpName("input"), dtype); + auto weights = ops::Placeholder(s.WithOpName("weights"), dtype); + const auto biasadd_attrs = ops::BiasAdd::DataFormat(data_format); + auto biasadd = + ops::BiasAdd(s.WithOpName("my_biasadd"), input, weights, biasadd_attrs); + return biasadd.operation.node()->def(); + }; + + typedef typename EnumToDataType::Type CType; + for (const string& data_format : {"NHWC", "NCHW"}) { + for (const int trt_input_rank : {1, 2, 3, 4}) { + test->Reset(); + NodeDef node_def = get_biasadd_nodedef(data_format); + + // Add input, dims_array will be like {2, 1, ..., 1, 3} + std::vector dims_array(trt_input_rank, 1); + if (trt_input_rank == 1) { + dims_array[0] = (data_format == "NHWC" ? 3 : 2); + } else { + dims_array[0] = 2; + dims_array[trt_input_rank - 1] = 3; + } + test->AddTestTensor("input", dims_array, /*batch_size=*/1, + TfDataTypeToTrt(dtype)); + + // Add bias weights. + const int channel_size = (data_format == "NHWC" ? 3 : 2); + std::vector bias(channel_size); + for (int i = 0; i < channel_size; ++i) { + bias[i] = CType(i + 1); // bias will be {1, 2, 3, ...} + } + test->AddTestWeights("weights", {channel_size}, bias); + + // Run the conversion. + test->RunValidationAndConversion(node_def); + TRT_TensorOrWeights output; + TF_EXPECT_OK(test->GetTensorOrWeights("my_biasadd", &output)); + EXPECT_TRUE(output.is_tensor()); + ExpectTrtDimsEqualsArray(dims_array, output.tensor()->getDimensions()); + + // Build and run the engine. + 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); + if (trt_input_rank == 1) { + if (data_format == "NHWC") { + EXPECT_THAT(output_data, ElementsAre(CType(1), CType(2), CType(3))); + } else { + EXPECT_THAT(output_data, 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))); + } else { + EXPECT_THAT(output_data, ElementsAre(CType(1), CType(1), CType(1), + CType(2), CType(2), CType(2))); + } + } + } + } +} + +TEST_F(OpConverterTest, ConvertBiasAdd) { + { + // Input list is empty, should fail. + NodeDef node_def = MakeNodeDef("my_biasadd", "BiasAdd", {}); + RunValidationAndConversion( + node_def, error::INVALID_ARGUMENT, + "Input expects tensor and weights, at my_biasadd"); + } + + // OK. Note that kINT32 is not supported by IScaleLayer, so we don't test + // DT_INT32 type here. + TestConvertBiasAdd(this); + TestConvertBiasAdd(this); +} + +template +NodeDef GetBinaryOpNodeDef(const string& input_name_l, + const string& input_name_r, DataType dtype) { + Scope s = Scope::NewRootScope(); + auto input_l = ops::Placeholder(s.WithOpName(input_name_l), dtype); + auto input_r = ops::Placeholder(s.WithOpName(input_name_r), dtype); + auto op = OpType(s.WithOpName("my_binary"), input_l, input_r); + return op.operation.node()->def(); +} + +void CheckAddedLayers(OpConverterTest* test, bool expect_scale_layer) { + bool element_wise_layer_found = false; + bool scale_layer_found = false; + for (int i = 0; i < test->converter_->network()->getNbLayers(); i++) { + nvinfer1::ILayer* layer = test->converter_->network()->getLayer(i); + if (dynamic_cast(layer)) { + scale_layer_found = true; + } else if (dynamic_cast(layer)) { + element_wise_layer_found = true; + } + } + EXPECT_EQ(expect_scale_layer, scale_layer_found); + EXPECT_NE(expect_scale_layer, element_wise_layer_found); +} + +template +void TestBinaryTensorOpWeightNoBroadcast(OpConverterTest* test) { + typedef typename EnumToDataType::Type CType; + for (auto swap_inputs : {false, true}) { + test->Reset(); + NodeDef node_def; + if (swap_inputs) { + node_def = GetBinaryOpNodeDef("weights", "input", dtype); + } else { + node_def = GetBinaryOpNodeDef("input", "weights", dtype); + } + + const std::vector operand1{CType(3), CType(7.5)}; + const std::vector operand2{CType(2), CType(3)}; + + // It requires the dims to be at least of rank 3 to apply an IScaleLayer. + test->AddTestTensor("input", /*dims=*/{1, 1, 2}, /*batch_size=*/1, + TfDataTypeToTrt(dtype)); + test->AddTestWeights("weights", /*dims=*/{1, 1, 2}, + /*values=*/swap_inputs ? operand1 : operand2); + test->RunValidationAndConversion(node_def); + + // Make sure it does use BinaryTensorOpWeight, not BinaryTensorOpTensor. + CheckAddedLayers(test, /*expect_scale_layer=*/true); + + // Check the dims of the output ITensor. + TRT_TensorOrWeights output; + TF_EXPECT_OK(test->GetTensorOrWeights("my_binary", &output)); + 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); + if (node_def.op() == "Add") { + EXPECT_THAT(output_data, ElementsAre(CType(5), CType(10.5))); + } else if (node_def.op() == "Sub") { + EXPECT_THAT(output_data, ElementsAre(CType(1), CType(4.5))); + } else if (node_def.op() == "Mul") { + EXPECT_THAT(output_data, ElementsAre(CType(6), CType(22.5))); + } else if (node_def.op() == "Div") { + EXPECT_THAT(output_data, ElementsAre(CType(1.5), CType(2.5))); + } else if (node_def.op() == "RealDiv") { + EXPECT_THAT(output_data, ElementsAre(CType(1.5), CType(2.5))); + } else { + ASSERT_TRUE(false); + } + } +} + +template +void TestBinaryTensorOpWeightWithChannelWiseBroadcast(OpConverterTest* test) { + typedef typename EnumToDataType::Type CType; + const NodeDef node_def = + GetBinaryOpNodeDef("input", "weights", dtype); + const std::vector input{CType(1), CType(2), CType(3), CType(4)}; + const std::vector weights{CType(10), CType(20)}; + // There are two types of valid dim pairs which requires channel-wise + // broadcasting: + // - input dims (X Y Z) vs weights dims (X 1 1) + // - input dims (X Y Z) vs weights dims (Z) + // Here X=Z=2 and Y=1. + for (auto weights_dims : std::vector>{{2, 1, 1}, {2}}) { + test->Reset(); + test->AddTestTensor("input", /*dims=*/{2, 1, 2}, /*batch_size=*/1, + TfDataTypeToTrt(dtype)); + test->AddTestWeights("weights", weights_dims, weights); + test->RunValidationAndConversion(node_def); + + // Make sure it does use BinaryTensorOpWeight, not BinaryTensorOpTensor. + CheckAddedLayers(test, /*expect_scale_layer=*/true); + + // Check the dims of the output ITensor. + TRT_TensorOrWeights output; + TF_EXPECT_OK(test->GetTensorOrWeights("my_binary", &output)); + 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); + if (weights_dims.size() == 1) { + EXPECT_THAT(output_data, + ElementsAre(CType(11), CType(22), CType(13), CType(24))); + } else { + EXPECT_THAT(output_data, + ElementsAre(CType(11), CType(12), CType(23), CType(24))); + } + } +} + +template +void TestBinaryTensorOpWeightWithUniformlyBroadcast(OpConverterTest* test) { + typedef typename EnumToDataType::Type CType; + const NodeDef node_def = + GetBinaryOpNodeDef("input", "weights", dtype); + const std::vector input{CType(1), CType(2), CType(3), CType(4)}; + const std::vector weights{CType(10)}; + test->Reset(); + test->AddTestTensor("input", /*dims=*/{2, 1, 2}, /*batch_size=*/1, + TfDataTypeToTrt(dtype)); + test->AddTestWeights("weights", {1, 1, 1, 1}, weights); + test->RunValidationAndConversion(node_def); + + // Make sure it does use BinaryTensorOpWeight, not BinaryTensorOpTensor. + CheckAddedLayers(test, /*expect_scale_layer=*/true); + + // Check the dims of the output ITensor. + TRT_TensorOrWeights output; + TF_EXPECT_OK(test->GetTensorOrWeights("my_binary", &output)); + 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, + ElementsAre(CType(11), CType(12), CType(13), CType(14))); +} + +template +void TestBinaryTensorOpWeightFallback(OpConverterTest* test, + const std::vector& input_dims, + const std::vector& weights_dims, + error::Code code = error::OK, + const char* error_msg_substr = nullptr, + const int input_batch_size = 1) { + const DataType dtype = DT_FLOAT; + typedef typename EnumToDataType::Type CType; + const size_t num_inputs = TrtDimsNumElements(GetTestDims(input_dims)); + const size_t num_weights = TrtDimsNumElements(GetTestDims(weights_dims)); + + test->Reset(); + const NodeDef node_def = + GetBinaryOpNodeDef("input", "weights", dtype); + test->AddTestTensor("input", /*dims=*/input_dims, input_batch_size, + TfDataTypeToTrt(dtype)); + test->AddTestWeights( + "weights", /*dims=*/weights_dims, + /*values=*/std::vector(num_weights, CType(1))); + test->RunValidationAndConversion(node_def, code, error_msg_substr); + if (code != error::OK) return; + + // Make sure it does use BinaryTensorOpTensor, not BinaryTensorOpWeight. + CheckAddedLayers(test, /*expect_scale_layer=*/false); + + TRT_TensorOrWeights output; + TF_EXPECT_OK(test->GetTensorOrWeights("my_binary", &output)); + EXPECT_TRUE(output.is_tensor()); + + // Check the dims of the output ITensor. + std::vector expected_output_dims = input_dims; + for (int i = expected_output_dims.size() - 1, j = weights_dims.size() - 1; + i >= 0 && j >= 0; --i, --j) { + if (expected_output_dims[i] == 1) { + expected_output_dims[i] = weights_dims[j]; + } + } + ExpectTrtDimsEqualsArray(expected_output_dims, + output.tensor()->getDimensions()); + + // 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); + if (node_def.op() == "Add") { + EXPECT_THAT(output_data, 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)))); + } else { + ASSERT_TRUE(false); + } +} + +template +void TestBinaryTensorOpTensor(OpConverterTest* test) { + typedef typename EnumToDataType::Type CType; + test->Reset(); + const NodeDef node_def = + GetBinaryOpNodeDef("input1", "input2", dtype); + test->AddTestTensor("input1", /*dims=*/{1, 2}, /*batch_size=*/1, + TfDataTypeToTrt(dtype)); + test->AddTestTensor("input2", /*dims=*/{2, 1}, /*batch_size=*/1, + TfDataTypeToTrt(dtype)); + test->RunValidationAndConversion(node_def); + + // Make sure it does use BinaryTensorOpTensor, not BinaryTensorOpWeight. + CheckAddedLayers(test, /*expect_scale_layer=*/false); + + // Check output dims. + TRT_TensorOrWeights output; + TF_EXPECT_OK(test->GetTensorOrWeights("my_binary", &output)); + EXPECT_TRUE(output.is_tensor()); + ExpectTrtDimsEqualsArray({2, 2}, output.tensor()->getDimensions()); + + std::vector output_data(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); + if (node_def.op() == "Add") { + EXPECT_THAT(output_data, + ElementsAre(CType(5), CType(8), CType(6), CType(9))); + } else if (node_def.op() == "Sub") { + EXPECT_THAT(output_data, + ElementsAre(CType(1), CType(4), CType(0), CType(3))); + } else if (node_def.op() == "Mul") { + EXPECT_THAT(output_data, + ElementsAre(CType(6), CType(12), CType(9), CType(18))); + } else if (node_def.op() == "Div") { + EXPECT_THAT(output_data, + ElementsAre(CType(1.5), CType(3), CType(1), CType(2))); + } else if (node_def.op() == "RealDiv") { + EXPECT_THAT(output_data, + ElementsAre(CType(1.5), CType(3), CType(1), CType(2))); + } else if (node_def.op() == "Minimum") { + EXPECT_THAT(output_data, + ElementsAre(CType(2), CType(2), CType(3), CType(3))); + } else if (node_def.op() == "Maximum") { + EXPECT_THAT(output_data, + ElementsAre(CType(3), CType(6), CType(3), CType(6))); + } else { + ASSERT_TRUE(false); + } +} + +TEST_F(OpConverterTest, ConvertBinary) { + // Input size doesn't match, should fail. + for (size_t num_inputs = 0; num_inputs < 2; ++num_inputs) { + Reset(); + 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"); + } + { + // Both inputs are weights. + Reset(); + NodeDef node_def = MakeNodeDef("my_add", "Add", {"weights1", "weights2"}); + AddTestWeights("weights1", {1}, {1}); + AddTestWeights("weights2", {1}, {1}); + RunValidationAndConversion( + node_def, error::UNIMPLEMENTED, + "Constant folding is falled back to TensorFlow, binary op received " + "both input as constant at: my_add"); + } + + // Test BinaryTensorOpWeight() without broadcasting. + TestBinaryTensorOpWeightNoBroadcast(this); + TestBinaryTensorOpWeightNoBroadcast(this); + TestBinaryTensorOpWeightNoBroadcast(this); + TestBinaryTensorOpWeightNoBroadcast(this); + TestBinaryTensorOpWeightNoBroadcast(this); +#if 0 + // TODO(b/119560144): it doesn't support FP16 constants and the following test + // will fail. + TestBinaryTensorOpWeightNoBroadcast(this); + TestBinaryTensorOpWeightNoBroadcast(this); + TestBinaryTensorOpWeightNoBroadcast(this); + TestBinaryTensorOpWeightNoBroadcast(this); + TestBinaryTensorOpWeightNoBroadcast(this); +#endif + + // Test BinaryTensorOpWeight() with channel-wise broadcasting. + TestBinaryTensorOpWeightWithChannelWiseBroadcast(this); + + // Test BinaryTensorOpWeight() with uniformly broadcasting. + TestBinaryTensorOpWeightWithUniformlyBroadcast(this); + + // Test BinaryTensorOpWeight() falling back to BinaryTensorOpTensor(). + // Unsupported op. + TestBinaryTensorOpWeightFallback(this, {1, 1, 1}, {1}); + // Rank of input tensor dimension <3. + TestBinaryTensorOpWeightFallback(this, {1, 1}, {1}); + // Broadcast on batch dimension, should fail. + TestBinaryTensorOpWeightFallback( + this, {1, 1, 1}, {2, 1, 1, 1}, error::INVALID_ARGUMENT, + "Unsupported binary op broadcast scheme for op my_binary", + /*input_batch_size=*/2); + // Incompatible dims with per-channel mode. + TestBinaryTensorOpWeightFallback(this, {1, 1, 1}, {1, 2, 1}); + // Incompatible dims. + TestBinaryTensorOpWeightFallback(this, {1, 2, 1}, {2}); + + // Test BinaryTensorOpTensor() with broadcasting. + TestBinaryTensorOpTensor(this); + TestBinaryTensorOpTensor(this); + TestBinaryTensorOpTensor(this); + TestBinaryTensorOpTensor(this); + TestBinaryTensorOpTensor(this); + TestBinaryTensorOpTensor(this); + TestBinaryTensorOpTensor(this); + + TestBinaryTensorOpTensor(this); + TestBinaryTensorOpTensor(this); + TestBinaryTensorOpTensor(this); + TestBinaryTensorOpTensor(this); + TestBinaryTensorOpTensor(this); + TestBinaryTensorOpTensor(this); + TestBinaryTensorOpTensor(this); +} + +TEST_F(OpConverterTest, ConvertQuantize) { + for (const string& op : + {"FakeQuantWithMinMaxArgs", "FakeQuantWithMinMaxVars", + "QuantizeAndDequantizeV2", "QuantizeAndDequantizeV3"}) { + // Input list is empty, should fail. + NodeDef node_def = MakeNodeDef("my_quantize", op, {}); + RunValidationAndConversion( + node_def, error::INVALID_ARGUMENT, + StrCat("Invalid number of inputs for ", op, ", at my_quantize") + .c_str()); + } + { + // FakeQuantWithMinMaxArgs attributes are empty, should fail. + NodeDef node_def = + MakeNodeDef("my_quantize", "FakeQuantWithMinMaxArgs", {"input"}); + AddTestTensor("input", {1, 2, 3}); + RunValidationAndConversion( + node_def, error::INVALID_ARGUMENT, + "Min or max attribute not found for FakeQuantWithMinMaxArgs " + "at my_quantize"); + } + { + // FakeQuantWithMinMaxArgs ranges set via attributes, ok. + Reset(); + Scope s = Scope::NewRootScope(); + auto input = ops::Placeholder(s.WithOpName("input"), DT_FLOAT); + auto quantize_attrs = ops::FakeQuantWithMinMaxArgs::Min(-6.0f).Max(6.0f); + auto quantize = ops::FakeQuantWithMinMaxArgs(s.WithOpName("my_quantize"), + input, quantize_attrs); + const NodeDef& node_def = quantize.operation.node()->def(); + AddTestTensor("input", {1, 2, 3}); + RunValidationAndConversion(node_def); + TRT_TensorOrWeights output; + TF_EXPECT_OK(GetTensorOrWeights("my_quantize", &output)); + EXPECT_TRUE(output.is_tensor()); + auto ranges = quantization_ranges(); + EXPECT_EQ(1, ranges.count(output.tensor())); + EXPECT_EQ(6.0f, ranges[output.tensor()]); + } + { + // FakeQuantWithMinMaxVars ranges set via inputs, ok. + Reset(); + Scope s = Scope::NewRootScope(); + auto input = ops::Placeholder(s.WithOpName("input"), DT_FLOAT); + auto weights_min = ops::Placeholder(s.WithOpName("weights_min"), DT_FLOAT); + auto weights_max = ops::Placeholder(s.WithOpName("weights_max"), DT_FLOAT); + auto quantize = ops::FakeQuantWithMinMaxVars( + s.WithOpName("my_quantize"), input, weights_min, weights_max); + const NodeDef& node_def = quantize.operation.node()->def(); + AddTestTensor("input", {1, 2, 3}); + AddTestWeights("weights_min", {1}, {-6.0f}); + AddTestWeights("weights_max", {1}, {6.0f}); + RunValidationAndConversion(node_def); + TRT_TensorOrWeights output; + TF_EXPECT_OK(GetTensorOrWeights("my_quantize", &output)); + EXPECT_TRUE(output.is_tensor()); + auto ranges = quantization_ranges(); + EXPECT_EQ(1, ranges.count(output.tensor())); + EXPECT_EQ(6.0f, ranges[output.tensor()]); + } + { + // QuantizeAndDequantizeV2 ranges set via inputs, ok. + Reset(); + Scope s = Scope::NewRootScope(); + auto input = ops::Placeholder(s.WithOpName("input"), DT_FLOAT); + auto weights_min = ops::Placeholder(s.WithOpName("weights_min"), DT_FLOAT); + auto weights_max = ops::Placeholder(s.WithOpName("weights_max"), DT_FLOAT); + auto quantize = ops::QuantizeAndDequantizeV2( + s.WithOpName("my_quantize"), input, weights_min, weights_max); + const NodeDef& node_def = quantize.operation.node()->def(); + AddTestTensor("input", {1, 2, 3}); + AddTestWeights("weights_min", {1}, {-6.0f}); + AddTestWeights("weights_max", {1}, {6.0f}); + RunValidationAndConversion(node_def); + TRT_TensorOrWeights output; + TF_EXPECT_OK(GetTensorOrWeights("my_quantize", &output)); + EXPECT_TRUE(output.is_tensor()); + auto ranges = quantization_ranges(); + EXPECT_EQ(1, ranges.count(output.tensor())); + EXPECT_EQ(6.0f, ranges[output.tensor()]); + } + { + // QuantizeAndDequantizeV2 Range inputs are tensors, should fail. + Reset(); + Scope s = Scope::NewRootScope(); + auto input = ops::Placeholder(s.WithOpName("input"), DT_FLOAT); + auto weights_min = ops::Placeholder(s.WithOpName("weights_min"), DT_FLOAT); + auto weights_max = ops::Placeholder(s.WithOpName("weights_max"), DT_FLOAT); + auto quantize = ops::QuantizeAndDequantizeV2( + s.WithOpName("my_quantize"), input, weights_min, weights_max); + const NodeDef& node_def = quantize.operation.node()->def(); + AddTestTensor("input", {1, 2, 3}); + 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"); + } + { + // QuantizeAndDequantizeV3 ranges set via inputs, ok. + Reset(); + Scope s = Scope::NewRootScope(); + auto input = ops::Placeholder(s.WithOpName("input"), DT_FLOAT); + auto weights_min = ops::Placeholder(s.WithOpName("weights_min"), DT_FLOAT); + auto weights_max = ops::Placeholder(s.WithOpName("weights_max"), DT_FLOAT); + auto num_bits = ops::Placeholder(s.WithOpName("num_bits"), DT_INT32); + auto quantize = ops::QuantizeAndDequantizeV3( + s.WithOpName("my_quantize"), input, weights_min, weights_max, num_bits); + const NodeDef& node_def = quantize.operation.node()->def(); + AddTestTensor("input", {1, 2, 3}); + AddTestWeights("weights_min", {1}, {-6.0f}); + AddTestWeights("weights_max", {1}, {6.0f}); + AddTestWeights("num_bits", {1}, {8}); + RunValidationAndConversion(node_def); + TRT_TensorOrWeights output; + TF_EXPECT_OK(GetTensorOrWeights("my_quantize", &output)); + EXPECT_TRUE(output.is_tensor()); + auto ranges = quantization_ranges(); + EXPECT_EQ(1, ranges.count(output.tensor())); + EXPECT_EQ(6.0f, ranges[output.tensor()]); + } +} + +TEST_F(OpConverterTest, ConvertRelu6) { + { + // Input list is empty, should fail. + NodeDef node_def = MakeNodeDef("my_relu6", "Relu6", {}); + RunValidationAndConversion( + node_def, error::INVALID_ARGUMENT, + "Invalid number of inputs for Relu6, at my_relu6"); + } + + // Get the NodeDef for Relu6. + Scope s = Scope::NewRootScope(); + auto input = ops::Placeholder(s.WithOpName("input"), DT_FLOAT); + auto relu6 = ops::Relu6(s.WithOpName("my_relu6"), input); + const NodeDef node_def = relu6.operation.node()->def(); + { + // Input is weights, should fail. + Reset(); + AddTestWeights("input", {1}, {1.0f}); + RunValidationAndConversion( + node_def, error::UNIMPLEMENTED, + "Relu6 is only implemented for tensors, not weights, at my_relu6"); + } + { + // Clip tensor values and set quantization ranges, ok. + Reset(); + AddTestTensor("input", {1, 2, 3}); + RunValidationAndConversion(node_def); + TRT_TensorOrWeights output; + TF_EXPECT_OK(GetTensorOrWeights("my_relu6", &output)); + EXPECT_TRUE(output.is_tensor()); + auto ranges = quantization_ranges(); + EXPECT_EQ(ranges[output.tensor()], 6.0f); + + 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)); + } +} + +template +void TestConvertSquare(OpConverterTest* test) { + test->Reset(); + typedef typename EnumToDataType::Type CType; + + Scope s = Scope::NewRootScope(); + auto input = ops::Placeholder(s.WithOpName("input"), dtype); + auto square = ops::Square(s.WithOpName("my_square"), input); + NodeDef node_def = square.operation.node()->def(); + + test->AddTestTensor("input", {1, 20}); + test->RunValidationAndConversion(node_def); + TRT_TensorOrWeights output; + TF_EXPECT_OK(test->GetTensorOrWeights("my_square", &output)); + EXPECT_TRUE(output.is_tensor()); + ExpectTrtDimsEqualsArray({1, 20}, output.tensor()->getDimensions()); + + const int num_inputs = 20; + std::vector input_data(num_inputs); + std::vector expected_output_data(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; + } + std::vector output_data(num_inputs); + test->BuildAndRun({{"input", input_data}}, "my_square", &output_data); + ExpectArrayNear(expected_output_data, output_data); +} + +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"); + } + { + // Input is weights, should fail. + Reset(); + Scope s = Scope::NewRootScope(); + auto input = ops::Placeholder(s.WithOpName("input"), DT_FLOAT); + auto square = ops::Square(s.WithOpName("my_square"), input); + NodeDef node_def = square.operation.node()->def(); + 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"); + } + + // OK. Note that kINT32 is not supported by IElementWiseLayer, so we don't + // test DT_INT32 type here. + TestConvertSquare(this); + // TODO(tmorris): Looks like there may be a bug with this layer for FP16 + // inputs. Disabling for now. + // TestConvertSquare(this); +} + +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"); + } + { + // Input is weights, should fail. + Reset(); + Scope s = Scope::NewRootScope(); + auto input = ops::Placeholder(s.WithOpName("input"), DT_FLOAT); + auto relu = ops::Relu(s.WithOpName("my_act"), input); + const NodeDef& node_def = relu.operation.node()->def(); + 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"); + } + + // Get nodedef for activation layer. + auto get_act_nodedef = [](string op_name) -> NodeDef { + Scope s = Scope::NewRootScope(); + auto input = ops::Placeholder(s.WithOpName("input"), DT_FLOAT); + if (op_name == "Relu") { + auto act = ops::Relu(s.WithOpName("my_act"), input); + return act.operation.node()->def(); + } else if (op_name == "Sigmoid") { + auto act = ops::Sigmoid(s.WithOpName("my_act"), input); + return act.operation.node()->def(); + } else if (op_name == "Tanh") { + auto act = ops::Tanh(s.WithOpName("my_act"), input); + return act.operation.node()->def(); + } + EXPECT_TRUE(false); + return NodeDef(); + }; + // Get expected output for activation layer. + auto get_act_output = [](string op_name, float input) -> float { + if (op_name == "Relu") { + return (input > 0.0f) ? input : 0.0f; + } else if (op_name == "Sigmoid") { + return 1.0f / (1.0f + std::exp(-input)); + } else if (op_name == "Tanh") { + return std::tanh(input); + } + EXPECT_TRUE(false); + return 0; + }; + + // Ok. + for (string op_name : {"Relu", "Sigmoid", "Tanh"}) { + Reset(); + NodeDef node_def = get_act_nodedef(op_name); + AddTestTensor("input", {1, 2, 3}); + RunValidationAndConversion(node_def); + TRT_TensorOrWeights output; + TF_EXPECT_OK(GetTensorOrWeights("my_act", &output)); + 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); + } + } +} + +TEST_F(OpConverterTest, ConvertExpandDims) { + { + // Input list is empty, should fail. + NodeDef node_def = MakeNodeDef("my_expanddims", "ExpandDims", {}); + RunValidationAndConversion( + node_def, error::INVALID_ARGUMENT, + "Two inputs expected for ExpandDims, at my_expanddims"); + } + + // Get the NodeDef for ExpandDims. + Scope s = Scope::NewRootScope(); + auto input = ops::Placeholder(s.WithOpName("input"), DT_FLOAT); + auto weights = ops::Placeholder(s.WithOpName("weights"), DT_INT32); + auto expanddims = + ops::ExpandDims(s.WithOpName("my_expanddims"), input, weights); + const NodeDef& node_def = expanddims.operation.node()->def(); + { + // Input is weights, should fail. + 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"); + } + { + // 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"); + } + { + // Add dim at batch dimension, should fail. + Reset(); + AddTestTensor("input", {1, 2, 3}); + AddTestWeights("weights", {1}, {0}); + RunValidationAndConversion( + node_def, error::UNIMPLEMENTED, + "Modifying batch dimension is not supported for ExpandDims, at " + "my_expanddims"); + } + { + // Add dim at batch dimension via negative axis, should fail. + Reset(); + AddTestTensor("input", {1, 2, 3}); + // Input is rank 4 (batch dim included) + AddTestWeights("weights", {1}, {-5}); + RunValidationAndConversion( + node_def, error::UNIMPLEMENTED, + "Modifying batch dimension is not supported for ExpandDims, at " + "my_expanddims"); + } + { + // Axis > rank(input), should fail. + Reset(); + AddTestTensor("input", {1, 2, 3}); + // Input is rank 4 (batch dim included) + AddTestWeights("weights", {1}, {5}); + RunValidationAndConversion( + node_def, error::INVALID_ARGUMENT, + "Axis for ExpandDims is invalid, must be in the range " + "[-rank(input) - 1, rank(input)], at my_expanddims"); + } + { + // Axis < -rank(input)-1, should fail. + Reset(); + AddTestTensor("input", {1, 2, 3}); + // Input is rank 4 (batch dim included) + AddTestWeights("weights", {1}, {-6}); + RunValidationAndConversion( + node_def, error::INVALID_ARGUMENT, + "Axis for ExpandDims is invalid, must be in the range " + "[-rank(input) - 1, rank(input)], at my_expanddims"); + } + + 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; + }; + + // Ok. + const int kExpandDimsOKCases = 8; + TestParams ok_params[kExpandDimsOKCases] = { + TestParams{{2, 3}, 1, {1, 2, 3}}, TestParams{{2, 3}, -3, {1, 2, 3}}, + TestParams{{2, 3}, 3, {2, 3, 1}}, TestParams{{2, 3}, -1, {2, 3, 1}}, + TestParams{{2, 3}, 2, {2, 1, 3}}, TestParams{{2, 3}, -2, {2, 1, 3}}, + TestParams{{6}, 1, {1, 6}}, TestParams{{6}, -1, {6, 1}}, + }; + for (int i = 0; i < kExpandDimsOKCases; ++i) { + Reset(); + AddTestTensor("input", ok_params[i].input_dims); + AddTestWeights("weights", {1}, {ok_params[i].axis}); + RunValidationAndConversion(node_def); + TRT_TensorOrWeights output; + TF_EXPECT_OK(GetTensorOrWeights("my_expanddims", &output)); + EXPECT_TRUE(output.is_tensor()); + 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)); + } +} + +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"); + } + { + // No attrs, should fail. + Reset(); + Scope s = Scope::NewRootScope(); + auto input = ops::Placeholder(s.WithOpName("input"), DT_FLOAT); + auto squeeze = ops::Squeeze(s.WithOpName("my_squeeze"), input); + const NodeDef& node_def = squeeze.operation.node()->def(); + AddTestTensor("input", {1, 2, 3}); + RunValidationAndConversion( + node_def, error::UNIMPLEMENTED, + "Squeeze is only implemented for explicit dims, at my_squeeze"); + } + + // Get the NodeDef for Squeeze. + auto get_squeeze_nodedef = [](std::vector axis) -> NodeDef { + Scope s = Scope::NewRootScope(); + auto input = ops::Placeholder(s.WithOpName("input"), DT_FLOAT); + ops::Squeeze::Attrs squeeze_attrs; + 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(); + }; + + { + // Input is weights, should fail. + Reset(); + NodeDef node_def = get_squeeze_nodedef({0}); + AddTestWeights("input", {1, 2, 3}, {1, 2, 3, 4, 5, 6}); + RunValidationAndConversion( + node_def, error::UNIMPLEMENTED, + "Squeeze expects tensor for input, at my_squeeze"); + } + { + // Squeeze batch dim, should fail. + Reset(); + NodeDef node_def = get_squeeze_nodedef({0}); + AddTestTensor("input", {1, 2, 3}); + RunValidationAndConversion(node_def, error::UNIMPLEMENTED, + "Cannot squeeze batch dimension, at my_squeeze"); + } + { + // Squeeze batch dim via negative axis, should fail. + Reset(); + NodeDef node_def = get_squeeze_nodedef({-4}); + AddTestTensor("input", {1, 2, 3}); + RunValidationAndConversion(node_def, error::UNIMPLEMENTED, + "Cannot squeeze batch dimension, at my_squeeze"); + } + { + // Squeeze >= rank(input), should fail. + Reset(); + NodeDef node_def = get_squeeze_nodedef({4}); + AddTestTensor("input", {1, 2, 3}); + RunValidationAndConversion( + node_def, error::INVALID_ARGUMENT, + "Axis for Squeeze is invalid, must be in the range " + "[-rank(input), rank(input)), at my_squeeze"); + } + { + // Squeeze < -rank(input), should fail. + Reset(); + NodeDef node_def = get_squeeze_nodedef({-5}); + AddTestTensor("input", {1, 2, 3}); + RunValidationAndConversion( + node_def, error::INVALID_ARGUMENT, + "Axis for Squeeze is invalid, must be in the range " + "[-rank(input), rank(input)), at my_squeeze"); + } + + 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; + }; + + // Ok. + const int kSqueezeOKCases = 10; + TestParams ok_params[kSqueezeOKCases] = { + TestParams{{1, 2, 3}, {1}, {2, 3}}, + TestParams{{1, 2, 3}, {-3}, {2, 3}}, + TestParams{{2, 3, 1}, {3}, {2, 3}}, + TestParams{{2, 3, 1}, {-1}, {2, 3}}, + TestParams{{1, 2, 1, 3, 1}, {1, 3, 5}, {2, 3}}, + TestParams{{1, 2, 1, 3, 1}, {3, 1, 5}, {2, 3}}, + TestParams{{1, 2, 1, 3, 1}, {-1, -3, -5}, {2, 3}}, + TestParams{{1, 2, 1, 3, 1}, {1, -3, 5}, {2, 3}}, + TestParams{{1, 6}, {1}, {6}}, + TestParams{{6, 1}, {2}, {6}}, + }; + for (int i = 0; i < kSqueezeOKCases; ++i) { + Reset(); + NodeDef node_def = get_squeeze_nodedef(ok_params[i].axis); + AddTestTensor("input", ok_params[i].input_dims); + RunValidationAndConversion(node_def); + TRT_TensorOrWeights output; + TF_EXPECT_OK(GetTensorOrWeights("my_squeeze", &output)); + EXPECT_TRUE(output.is_tensor()); + 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)); + } +} + +TEST_F(OpConverterTest, ConvertStridedSlice) { + { + // Input list is empty, should fail. + NodeDef node_def = MakeNodeDef("my_strided_slice", "StridedSlice", {}); + RunValidationAndConversion( + node_def, error::INVALID_ARGUMENT, + "StridedSlice expects 4 inputs, at my_strided_slice"); + } + + // Get nodedef for StridedSlice layer. + auto get_strided_slice_nodedef = + [](int begin_mask = 0, int end_mask = 0, int ellipsis_mask = 0, + int new_axis_mask = 0, int shrink_axis_mask = 0) -> NodeDef { + Scope s = Scope::NewRootScope(); + auto input = ops::Placeholder(s.WithOpName("input"), DT_FLOAT); + auto begin = ops::Placeholder(s.WithOpName("begin"), DT_INT32); + auto end = ops::Placeholder(s.WithOpName("end"), DT_INT32); + auto strides = ops::Placeholder(s.WithOpName("strides"), DT_INT32); + ops::StridedSlice::Attrs attrs = ops::StridedSlice::Attrs() + .BeginMask(begin_mask) + .EndMask(end_mask) + .EllipsisMask(ellipsis_mask) + .NewAxisMask(new_axis_mask) + .ShrinkAxisMask(shrink_axis_mask); + auto strided_slice = ops::StridedSlice(s.WithOpName("my_strided_slice"), + input, begin, end, strides, attrs); + return strided_slice.operation.node()->def(); + }; + + { + // 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"); + } + { + // Begin, end, strides are tensors, should fail. + Reset(); + NodeDef node_def = get_strided_slice_nodedef(); + AddTestTensor("input", {1, 2, 3}); + AddTestTensor("begin", {4}); + AddTestTensor("end", {4}); + AddTestTensor("strides", {4}); + RunValidationAndConversion( + node_def, error::INVALID_ARGUMENT, + "StridedSlice expects weights for begin, end, and strides, at " + "my_strided_slice"); + } + { + // Non-zero ellipsis_mask, should fail. + Reset(); + NodeDef node_def = get_strided_slice_nodedef( + /*begin_mask=*/0, /*end_mask=*/0, /*ellipsis_mask=*/2, + /*new_axis_mask=*/0, /*shrink_axis_mask=*/0); + AddTestTensor("input", {1, 2, 3}); + 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, + "ellipsis_mask is not supported for StridedSlice, at " + "my_strided_slice"); + } + { + // Modify batch dim, should fail. + Reset(); + NodeDef node_def = get_strided_slice_nodedef(); + AddTestTensor("input", {1, 2, 3}); + AddTestWeights("begin", {4}, {0, 0, 0, 0}); + AddTestWeights("end", {4}, {0, 1, 2, 3}); + AddTestWeights("strides", {4}, {1, 1, 1, 1}); + RunValidationAndConversion( + node_def, error::UNIMPLEMENTED, + "StridedSlice can't modify batch dim, at my_strided_slice"); + } + { + // Stride is not 1, should fail. + Reset(); + NodeDef node_def = get_strided_slice_nodedef(); + AddTestTensor("input", {1, 2, 3}); + AddTestWeights("begin", {4}, {0, 0, 0, 0}); + AddTestWeights("end", {4}, {1, 1, 2, 3}); + AddTestWeights("strides", {4}, {1, 2, -1, 3}); + RunValidationAndConversion(node_def, error::UNIMPLEMENTED, + "StridedSlice is only implemented for stride of " + "1, at my_strided_slice"); + } + { + // Begin out of bounds, should fail. + Reset(); + NodeDef node_def = get_strided_slice_nodedef(); + AddTestTensor("input", {1, 2, 3}); + AddTestWeights("begin", {4}, {1, 2, 3, 4}); + AddTestWeights("end", {4}, {0, 1, 2, 3}); + AddTestWeights("strides", {4}, {1, 1, 1, 1}); + RunValidationAndConversion( + node_def, error::INVALID_ARGUMENT, + "begin value of 2 for StridedSlice is invalid, must be in the range " + "[-dim_size(i), dim_size(i)], at my_strided_slice"); + } + { + // End out of bounds, should fail. + Reset(); + NodeDef node_def = get_strided_slice_nodedef(); + AddTestTensor("input", {1, 2, 3}); + AddTestWeights("begin", {4}, {0, 0, 0, 0}); + AddTestWeights("end", {4}, {1, 2, 3, 4}); + AddTestWeights("strides", {4}, {1, 1, 1, 1}); + RunValidationAndConversion( + node_def, error::INVALID_ARGUMENT, + "end value of 2 for StridedSlice is invalid, must be in the range " + "[-dim_size(i), dim_size(i)], at my_strided_slice"); + } + { + // Size of sliced dim is negative, should fail. + Reset(); + NodeDef node_def = get_strided_slice_nodedef(); + AddTestTensor("input", {1, 2, 3}); + AddTestWeights("begin", {4}, {0, 0, 2, 0}); + AddTestWeights("end", {4}, {1, 1, 0, 3}); + AddTestWeights("strides", {4}, {1, 1, 1, 1}); + RunValidationAndConversion( + node_def, error::INVALID_ARGUMENT, + "New size of sliced dimension is negative, at my_strided_slice"); + } + + 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; + std::vector end; + int begin_mask; + int end_mask; + std::vector expected_output; + }; + + // 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}, + /*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}, + /*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}, + /*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}, + /*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}, + /*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}, + /*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}, + /*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}, + /*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}, + /*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}, + /*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}, + /*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}, + /*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}, + /*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}, + /*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}, + /*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}, + /*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}, + /*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}, + /*expected_output=*/{1, 2, 3, 4, 5}}, + }; + + for (int i = 0; i < kStridedSliceOKCases; i++) { + Reset(); + NodeDef node_def = get_strided_slice_nodedef(ok_params[i].begin_mask, + ok_params[i].end_mask); + AddTestTensor("input", ok_params[i].input_dims); + AddTestWeights("begin", + {static_cast(ok_params[i].begin.size())}, + ok_params[i].begin); + AddTestWeights("end", {static_cast(ok_params[i].end.size())}, + ok_params[i].end); + std::vector strides(ok_params[i].input_dims.size(), 1); + AddTestWeights("strides", {static_cast(strides.size())}, + strides); + RunValidationAndConversion(node_def); + + 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)); + } +} + +TEST_F(OpConverterTest, ConvertConv2D) { + { + // Input list is empty, should fail. + NodeDef node_def = MakeNodeDef("my_conv2d", "Conv2D", {}); + RunValidationAndConversion( + node_def, error::INVALID_ARGUMENT, + "Two inputs are expected for Conv2D, 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}) -> NodeDef { + Scope s = Scope::NewRootScope(); + auto input = ops::Placeholder(s.WithOpName("input"), DT_FLOAT); + auto filter = ops::Placeholder(s.WithOpName("weights"), DT_FLOAT); + 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, + "Conv2D is only implemented for tensors, not weights, 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, + "Kernel for Conv2D must be constant weights, 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"); + } + { + // 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 { + TestParams(const std::vector& input_dims, + const std::vector& input, + const std::vector& filter_dims, + const std::vector& filter, + const std::vector& strides, const string& padding, + const string& data_format, const std::vector& dilations, + const std::vector& expected_output_dims, + const std::vector& expected_output) + : input_dims(input_dims), + input(input), + filter_dims(filter_dims), + filter(filter), + strides(strides), + padding(padding), + data_format(data_format), + dilations(dilations), + expected_output_dims(expected_output_dims), + expected_output(expected_output) {} + + 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; + std::vector expected_output_dims; + std::vector expected_output; + }; + + // Ok. + const int kConv2DOKCases = 6; + 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}, + /*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}, + /*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}, + /*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}, + /*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}, + /*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}, + /*expected_output_dims=*/{1, 2, 2}, + /*expected_output=*/{1, 0, 1, 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); + AddTestTensor("input", ok_params[i].input_dims); + AddTestWeights("weights", ok_params[i].filter_dims, + ok_params[i].filter); + 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()); + std::vector output_data(ok_params[i].expected_output.size()); + BuildAndRun({{"input", ok_params[i].input}}, "my_conv2d", + &output_data); + EXPECT_THAT(output_data, ElementsAreArray(ok_params[i].expected_output)); + } +} + +} // namespace convert +} // namespace tensorrt +} // namespace tensorflow + +#endif // GOOGLE_TENSORRT +#endif // GOOGLE_CUDA diff --git a/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.cc b/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.cc new file mode 100644 index 0000000000000000000000000000000000000000..ebf8df1349363e9986020ea705b32edfef43bc93 --- /dev/null +++ b/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.cc @@ -0,0 +1,302 @@ +/* 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/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" +#include "tensorflow/core/lib/strings/numbers.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/stacktrace.h" + +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT +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; + +tensorflow::Status TRTOptimizationPass::Init( + const tensorflow::RewriterConfig_CustomGraphOptimizer* config) { + VLOG(1) << "Called INIT for " << name_ << " with config = " << config; + if (config == nullptr) { + return tensorflow::Status::OK(); + } + const auto params = config->parameter_map(); + if (params.count("minimum_segment_size")) { + minimum_segment_size_ = params.at("minimum_segment_size").i(); + } + if (params.count("max_batch_size")) { + maximum_batch_size_ = params.at("max_batch_size").i(); + } + if (params.count("is_dynamic_op")) { + is_dynamic_op_ = params.at("is_dynamic_op").b(); + } + if (params.count("cached_engine_batches")) { + auto batch_vec = params.at("cached_engine_batches").list(); + batches_.reserve(batch_vec.i_size()); + for (const auto i : batch_vec.i()) { + batches_.push_back(i); + } + } + if (params.count("maximum_cached_engines")) { + max_cached_batches_ = params.at("maximum_cached_engines").i(); + } + if (params.count("max_workspace_size_bytes")) { + max_workspace_size_bytes_ = params.at("max_workspace_size_bytes").i(); + } + if (params.count("precision_mode")) { + TF_RETURN_IF_ERROR(GetPrecisionMode( + Uppercase(params.at("precision_mode").s()), &precision_mode_)); + } + if (params.count("use_calibration")) { + use_calibration_ = params.at("use_calibration").b(); + } + return tensorflow::Status::OK(); +} + +void TRTOptimizationPass::PrintDebugInfo( + tensorflow::grappler::Cluster* cluster, + const tensorflow::grappler::GrapplerItem& item) { + LOG(INFO) << "Cluster = " << cluster; + string offset(" "); + string offset2 = StrCat(offset, offset); + string offset3 = StrCat(offset2, offset); + string offset4 = StrCat(offset2, offset2); + if (cluster) { + LOG(INFO) << offset << "type = " << cluster->type(); + LOG(INFO) << offset << "num warmup steps = " << cluster->NumWarmupSteps(); + const auto dev_names = cluster->GetDeviceNames(); + if (dev_names.size()) { + LOG(INFO) << offset << " Device names:"; + for (const auto s : dev_names) { + LOG(INFO) << offset2 << s; + } + } + std::unordered_map peak_mem; + auto status = cluster->GetPeakMemoryUsage(&peak_mem); + if (status == tensorflow::Status::OK()) { + LOG(INFO) << offset << "Peak Memory Usage :"; + for (auto s : peak_mem) { + LOG(INFO) << offset2 << s.first << " = " << s.second; + } + } + + const auto dev_props = cluster->GetDevices(); + if (dev_props.size()) { + LOG(INFO) << offset << "Device properties:"; + for (auto k : dev_props) { + LOG(INFO) << offset2 << k.first; + const auto& dt = k.second; + LOG(INFO) << offset3 << "type = " << dt.type(); + LOG(INFO) << offset3 << "vendor = " << dt.vendor(); + LOG(INFO) << offset3 << "model = " << dt.model(); + LOG(INFO) << offset3 << "frequency = " << dt.frequency(); + LOG(INFO) << offset3 << "num cores = " << dt.num_cores(); + LOG(INFO) << offset3 << "num registers = " << dt.num_registers(); + LOG(INFO) << offset3 << "L1 cache size = " << dt.l1_cache_size(); + LOG(INFO) << offset3 << "L2 cache size = " << dt.l2_cache_size(); + LOG(INFO) << offset3 << "L3 cache size = " << dt.l3_cache_size(); + LOG(INFO) << offset3 << "SHMem per SMP = " + << dt.shared_memory_size_per_multiprocessor(); + LOG(INFO) << offset3 << "memory size = " << dt.memory_size(); + LOG(INFO) << offset3 << "bandwidth = " << dt.bandwidth(); + if (dt.environment_size()) { + LOG(INFO) << offset3 << "environment :"; + for (const auto e : dt.environment()) { + LOG(INFO) << offset4 << e.first << " = " << e.second; + } + } + } + } + } + LOG(INFO) << "item: " << item.id; + if (item.feed.size()) { + LOG(INFO) << offset << "Feeds :"; + for (const auto& f : item.feed) { + const auto& shape = f.second.shape(); + LOG(INFO) << offset2 << f.first << " = shaped " << shape.DebugString(); + } + } else { + LOG(INFO) << offset << "No Feeds"; + } + if (item.fetch.size()) { + LOG(INFO) << offset << "Fetches :"; + for (const auto& f : item.fetch) { + LOG(INFO) << offset2 << f; + } + } else { + LOG(INFO) << offset << "No Fetches"; + } + + if (item.init_ops.size()) { + LOG(INFO) << offset << "init ops :"; + for (const auto& f : item.init_ops) { + LOG(INFO) << offset2 << f; + } + } else { + LOG(INFO) << offset << "No init ops"; + } + 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()) { + LOG(INFO) << offset << "keep ops :"; + for (const auto& f : item.keep_ops) { + LOG(INFO) << offset2 << f; + } + } else { + LOG(INFO) << offset << "No keep ops"; + } + for (const auto dev : cluster->GetDeviceSet()->devices()) { + const auto& pname = dev->parsed_name(); + LOG(INFO) << "Device name= " << dev->name() + << " parsedname job= " << pname.job << " id= " << pname.id + << " has_id: " << pname.has_id << " has_job: " << pname.has_job + << "has_type: " << pname.has_type << " type =" << pname.type; + } +} + +tensorflow::Status TRTOptimizationPass::Optimize( + tensorflow::grappler::Cluster* cluster, + const tensorflow::grappler::GrapplerItem& item, GraphDef* optimized_graph) { + VLOG(1) << "Called TRTOptimization Pass " << name_; + // This is a hack to workaround optimizer issue. MetaOptimizer calls + // optimization passes on function objects as well, we should not modify + // generated funcdefs! This is fragile but we don't have any other option + // until framework fixes it. + if (item.id != "tf_graph") { + LOG(WARNING) << name_ + << " is probably called on funcdef! This optimizer must *NOT* " + "be called on function objects."; + *optimized_graph = item.graph; + return tensorflow::Status::OK(); + } + if (VLOG_IS_ON(3)) { + LOG(INFO) << CurrentStackTrace(); + PrintDebugInfo(cluster, item); + } + int max_dim = -1; + if (item.feed.size()) { + for (const auto& f : item.feed) { + const auto& shape = f.second.shape(); + if (shape.dims() > 0) { + if (shape.dim_size(0) > max_dim) max_dim = shape.dim_size(0); + } + } + } + if (maximum_batch_size_ < 0) { // automatic batch size from input + if (max_dim > 0) { + maximum_batch_size_ = max_dim; + VLOG(1) << "Setting maximum batch size to " << max_dim; + } else { + maximum_batch_size_ = 128; + LOG(WARNING) << "Maximum batch size is not set" + " and can't be deduced from inputs setting it to" + << maximum_batch_size_ + << ". Suggest configuring it from configuration parameters"; + } + } else { + if (max_dim > maximum_batch_size_) { + LOG(WARNING) << "Configured batch size " << maximum_batch_size_ + << " is less than input batch size " << max_dim + << " adjusting maximum batch size to match input batch size"; + } + } + tensorflow::grappler::GraphProperties static_graph_properties(item); + TF_RETURN_IF_ERROR(static_graph_properties.InferStatically(true)); + tensorflow::tensorrt::convert::ConversionParams cp; + + if (use_calibration_ && precision_mode_ != INT8MODE) { + 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."; + use_calibration_ = false; + } + + std::vector nodes_to_preserve; + for (const auto& n : item.NodesToPreserve()) { + auto tokens = str_util::Split(n, ":"); + string s = tokens.at(0); + for (int i = 1; i < tokens.size() - 1; ++i) { + StrAppend(&s, ":", tokens.at(i)); + } + int dumm_port = -1; + // 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)) { // non-absl ok + StrAppend(&s, ":", tokens.back()); + } + nodes_to_preserve.push_back(s); + } + cp.input_graph_def = &item.graph; + cp.output_names = &nodes_to_preserve; + cp.max_batch_size = maximum_batch_size_; + cp.max_workspace_size_bytes = max_workspace_size_bytes_; + cp.output_graph_def = optimized_graph; + cp.precision_mode = precision_mode_; + cp.minimum_segment_size = minimum_segment_size_; + cp.graph_properties = &static_graph_properties; + cp.cluster = cluster; + cp.is_dyn_op = is_dynamic_op_; + cp.cached_engine_batches = batches_; + cp.max_cached_engines = max_cached_batches_; + cp.use_calibration = use_calibration_; + auto status = tensorflow::tensorrt::convert::ConvertAfterShapes(cp); + VLOG(1) << "Returning from " << name_; + return status; +} + +void TRTOptimizationPass::Feedback( + tensorflow::grappler::Cluster* cluster, + const tensorflow::grappler::GrapplerItem& item, + const GraphDef& optimized_graph, double result) {} + +} // namespace convert +} // namespace tensorrt +} // namespace tensorflow + +class VerboseCustomGraphOptimizerRegistrar + : public tensorflow::grappler::CustomGraphOptimizerRegistrar { + public: + VerboseCustomGraphOptimizerRegistrar( + const tensorflow::grappler::CustomGraphOptimizerRegistry::Creator& cr, + const tensorflow::string& name) + : tensorflow::grappler::CustomGraphOptimizerRegistrar(cr, name) { + VLOG(1) << "Constructing a CustomOptimizationPass registration object for " + << name; + } +}; + +static VerboseCustomGraphOptimizerRegistrar TRTOptimizationPass_Registrar( + []() { + VLOG(1) + << "Instantiating CustomOptimizationPass object TensorRTOptimizer"; + return new tensorflow::tensorrt::convert::TRTOptimizationPass( + "TensorRTOptimizer"); + }, + ("TensorRTOptimizer")); + +#endif +#endif diff --git a/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.h b/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.h similarity index 88% rename from tensorflow/contrib/tensorrt/convert/trt_optimization_pass.h rename to tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.h index 71b51d13681cb3f75dad034f3fb0f73dea2bacc1..bd6c6dbce1ddb8757227a1c71408770ee8be48d8 100644 --- a/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.h +++ b/tensorflow/compiler/tf2tensorrt/convert/trt_optimization_pass.h @@ -13,8 +13,8 @@ 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 @@ -38,7 +38,8 @@ class TRTOptimizationPass : public tensorflow::grappler::CustomGraphOptimizer { maximum_batch_size_(-1), is_dynamic_op_(false), max_cached_batches_(1), - max_workspace_size_bytes_(256LL << 20) { + max_workspace_size_bytes_(256LL << 20), + use_calibration_(true) { VLOG(1) << "Constructing " << name_; } @@ -67,6 +68,7 @@ class TRTOptimizationPass : public tensorflow::grappler::CustomGraphOptimizer { std::vector batches_; int max_cached_batches_; int64_t max_workspace_size_bytes_; + bool use_calibration_; }; } // namespace convert @@ -75,4 +77,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/compiler/tf2tensorrt/convert/utils.cc b/tensorflow/compiler/tf2tensorrt/convert/utils.cc new file mode 100644 index 0000000000000000000000000000000000000000..62a0f62ad6657f2d1551cd093f4f2d93c25f4cae --- /dev/null +++ b/tensorflow/compiler/tf2tensorrt/convert/utils.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 "tensorflow/compiler/tf2tensorrt/convert/utils.h" + +#include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/lib/core/status.h" + +namespace tensorflow { +namespace tensorrt { + +bool IsGoogleTensorRTEnabled() { + // TODO(laigd): consider also checking if tensorrt shared libraries are + // accessible. We can then direct users to this function to make sure they can + // safely write code that uses tensorrt conditionally. E.g. if it does not + // check for for tensorrt, and user mistakenly uses tensorrt, they will just + // crash and burn. +#if GOOGLE_CUDA && GOOGLE_TENSORRT + return true; +#else + return false; +#endif +} + +Status GetPrecisionModeName(const int precision_mode, string* name) { + switch (precision_mode) { + case FP32MODE: + *name = "FP32"; + break; + case FP16MODE: + *name = "FP16"; + break; + case INT8MODE: + *name = "INT8"; + break; + default: + return tensorflow::errors::OutOfRange("Unknown precision mode"); + } + return Status::OK(); +} + +Status GetPrecisionMode(const string& name, int* precision_mode) { + if (name == "FP32") { + *precision_mode = FP32MODE; + } else if (name == "FP16") { + *precision_mode = FP16MODE; + } else if (name == "INT8") { + *precision_mode = INT8MODE; + } else { + return tensorflow::errors::InvalidArgument("Invalid precision mode name: ", + name); + } + return Status::OK(); +} + +} // namespace tensorrt +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2tensorrt/convert/utils.h b/tensorflow/compiler/tf2tensorrt/convert/utils.h new file mode 100644 index 0000000000000000000000000000000000000000..9f9ee59087d461bdc825346d9adc976e42f47c5e --- /dev/null +++ b/tensorflow/compiler/tf2tensorrt/convert/utils.h @@ -0,0 +1,50 @@ +/* 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_CONVERT_UTILS_H_ +#define TENSORFLOW_COMPILER_TF2TENSORRT_CONVERT_UTILS_H_ + +#include + +#include "tensorflow/core/lib/core/status.h" + +namespace tensorflow { +namespace tensorrt { + +template +struct TrtDestroyer { + void operator()(T* t) { + if (t) t->destroy(); + } +}; + +template +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; + +Status GetPrecisionModeName(const int precision_mode, string* name); + +Status GetPrecisionMode(const string& name, int* precision_mode); + +} // namespace tensorrt +} // namespace tensorflow + +#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..eae1f8e7525f1816d1c50072ebe4ba6713c96e47 --- /dev/null +++ b/tensorflow/compiler/tf2tensorrt/kernels/get_serialized_resource_op.cc @@ -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_KERNELS_GET_SERIALIZED_RESOURCE_OP_H_ +#define TENSORFLOW_COMPILER_TF2TENSORRT_KERNELS_GET_SERIALIZED_RESOURCE_OP_H_ + +#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 +#endif // TENSORFLOW_COMPILER_TF2TENSORRT_KERNELS_GET_SERIALIZED_RESOURCE_OP_H_ 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/compiler/tf2tensorrt/kernels/trt_engine_op.cc b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc new file mode 100644 index 0000000000000000000000000000000000000000..198d68b60985d2b3f2ef958c4f13f94054d4875a --- /dev/null +++ b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc @@ -0,0 +1,651 @@ +/* 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/kernels/trt_engine_op.h" + +#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_logger.h" +#include "tensorflow/compiler/tf2tensorrt/utils/trt_resource_manager.h" +#include "tensorflow/compiler/tf2tensorrt/utils/trt_resources.h" +#include "tensorflow/core/framework/graph_to_functiondef.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/stream_executor.h" +#include "tensorflow/core/platform/types.h" + +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT +#include "cuda/include/cuda_runtime_api.h" + +namespace tensorflow { +namespace tensorrt { +static Logger logger; +using absl::StrAppend; +using absl::StrCat; +using ::nvinfer1::IRuntime; + +// A helper class to call done() when destructed for asynchronous execution. +// Helps simultaneous execution of native and TRT engines. +class AsyncHelper : public tensorflow::core::RefCounted { + public: + AsyncHelper(AsyncOpKernel::DoneCallback done) { done_ = done; } + ~AsyncHelper() override { done_(); } + + private: + AsyncOpKernel::DoneCallback done_; +}; + +#define TYPECASE(dt, X, Y) \ + case dt: { \ + return (void*)X->flat::Type>().data(); \ + } + +void* GetTensorAddress(const Tensor* tensor_ptr) { + auto tensor_type = tensor_ptr->dtype(); + switch (tensor_type) { + TYPECASE(tensorflow::DT_FLOAT, tensor_ptr, dest_ptr); + TYPECASE(tensorflow::DT_HALF, tensor_ptr, dest_ptr); + TYPECASE(tensorflow::DT_INT8, tensor_ptr, dest_ptr); + default: { + LOG(ERROR) << "Unsupported Data type " + << tensorflow::DataTypeString(tensor_type); + return nullptr; + } + } +} + +tensorflow::Status TRTEngineOp::ConstructFunctionHandle(OpKernelContext* ctx) { + VLOG(1) << "Constructing function handle"; + auto lib = ctx->function_library(); + if (lib == nullptr) { + return tensorflow::errors::Internal("Context function library is null"); + } + auto fdef = lib->GetFunctionLibraryDefinition()->Find(funcdef_name_); + if (fdef == nullptr) { + return tensorflow::errors::Internal("Native FunctionDef ", funcdef_name_, + " can't be found in function library"); + } + tensorflow::FunctionLibraryRuntime::InstantiateOptions inst_ops; + inst_ops.overlay_lib = nullptr; + inst_ops.state_handle = ""; + inst_ops.target = ctx->device()->name(); + native_func_ = 0; + auto status = lib->Instantiate(funcdef_name_, AttrSlice(&fdef->attr()), + inst_ops, &native_func_); + if (!status.ok()) { + LOG(ERROR) << " Instantiating native function " << funcdef_name_ + << " failed!"; + } + return status; +} + +TRTEngineOp::TRTEngineOp(OpKernelConstruction* context) + : AsyncOpKernel(context) { + // read serialized_engine + OP_REQUIRES_OK(context, + context->GetAttr("serialized_segment", &serialized_segment_)); + OP_REQUIRES_OK(context, + context->GetAttr("workspace_size_bytes", &workspace_size_)); + OP_REQUIRES_OK(context, context->GetAttr("static_engine", &static_engine_)); + if (!static_engine_) { + if (!segment_graph_.ParseFromString(serialized_segment_)) { + LOG(ERROR) << "Parsing segment graph failed!"; + context->SetStatus(tensorflow::errors::InvalidArgument( + "Failed to parse segment graphdef!")); + return; + } + serialized_segment_.resize(0); + } + VLOG(1) << "Constructing " << name(); + string precision_string; + OP_REQUIRES_OK(context, + context->GetAttr("precision_mode", &precision_string)); + string calibration_data; + OP_REQUIRES_OK(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, + context->GetAttr("use_calibration", &use_calibration_)); + calibration_mode_ = (use_calibration_ && precision_mode_ == INT8MODE && + calibration_data.size() == 0); + if (calibration_data.size()) { + 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("cached_engine_batches", + &cached_engine_batches_)); + std::sort(cached_engine_batches_.begin(), cached_engine_batches_.end()); + if (VLOG_IS_ON(1)) { + string s("Engine Batches= "); + for (auto i : cached_engine_batches_) { + StrAppend(&s, i, " "); + } + VLOG(1) << s; + } +} + +void TRTEngineOp::ExecuteNativeSegment(OpKernelContext* ctx, + AsyncHelper* helper) { + std::vector inputs; + std::vector* outputs = new std::vector(); + if (native_func_ == tensorflow::kInvalidHandle) { + auto status = ConstructFunctionHandle(ctx); + if (!status.ok()) { + LOG(ERROR) << "Couldn't construct function handle " << funcdef_name_; + ctx->SetStatus(status); + return; + } + } + auto lib = ctx->function_library(); + tensorflow::FunctionLibraryRuntime::Options opts; + opts.step_id = ctx->step_id(); + opts.rendezvous = ctx->rendezvous(); + opts.cancellation_manager = ctx->cancellation_manager(); + opts.runner = ctx->runner(); + for (int i = 0; i < ctx->num_inputs(); i++) { + inputs.push_back(ctx->input(i)); + } + helper->Ref(); // Increment count for calculating native graph + VLOG(1) << "Executing native segment: " << name(); + lib->Run(opts, native_func_, inputs, outputs, + [this, ctx, outputs, helper](const tensorflow::Status& s) { + tensorflow::core::ScopedUnref sc(helper); + 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)); + } + test::AddTestValue(StrCat(this->name(), ":ExecuteNativeSegment"), + "done"); + delete outputs; + }); +} + +void TRTEngineOp::ExecuteCalibration(OpKernelContext* ctx, + AsyncHelper* helper) { + VLOG(1) << "Executing TRT calibration: " << name(); + helper->Ref(); + tensorflow::core::ScopedUnref sc(helper); + auto res_mgr = ctx->resource_manager(); + TRTCalibrationResource* calib_res = nullptr; + OP_REQUIRES_OK( + ctx, + res_mgr->LookupOrCreate( + "TF_TRT_Calibration", name(), + reinterpret_cast(&calib_res), + {[ctx, this](SerializableResourceBase** cr) -> tensorflow::Status { + return this->AllocateCalibrationResources(ctx, cr); + }})); + tensorflow::core::ScopedUnref calib_sc(calib_res); + // 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 + std::unordered_map input_data; + for (int i = 0; i < num_inputs; i++) { + const Tensor& t = ctx->input(i); + void* data_address = GetTensorAddress(&t); + if (data_address == nullptr) { + ctx->SetStatus(tensorflow::errors::InvalidArgument( + "Unsupported data type encountered in input ", i)); + return; + } + // Check the allocated buffer is sufficient for input + 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); + } + VLOG(2) << "Filled map for sending"; + // copied from cuda_kernel_helper since it seems only valid in *.cu.cc files + const cudaStream_t* stream = CHECK_NOTNULL( + reinterpret_cast(ctx->op_device_context() + ->stream() + ->implementation() + ->GpuStreamMemberHack())); + calib_res->calibrator_->setBatch(input_data, *stream); + test::AddTestValue(StrCat(name(), ":ExecuteCalibration"), "done"); + VLOG(2) << "Passed calibration data"; + ExecuteNativeSegment(ctx, helper); +} + +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_batch_size != -1); +} + +void TRTEngineOp::ComputeAsync(OpKernelContext* ctx, + AsyncOpKernel::DoneCallback done) { + auto helper = new AsyncHelper(done); + tensorflow::core::ScopedUnref sc(helper); + if (calibration_mode_) { + ExecuteCalibration(ctx, helper); + return; + } + // Get shapes of inputs to engine. + std::vector input_shapes; + for (int i = 0; i < ctx->num_inputs(); ++i) { + input_shapes.emplace_back(ctx->input(i).shape()); + } + 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, engine_context); + if (retry) { + LOG(WARNING) << "Failed to execute engine, " + << "retrying with native segment for " << name(); + ExecuteNativeSegment(ctx, helper); + return; + } +} + +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 = cuda_engine->getBindingIndex(input_name.c_str()); + if (binding_index == -1) { + LOG(ERROR) << "Input node not found, at " << input_name; + return kRetry; + } + + const Tensor& input_tensor = ctx->input(i); + const TensorShape& input_shape = input_tensor.shape(); + if (num_batch != input_shape.dim_size(0)) { + LOG(ERROR) << "Input data has inconsistent batch size: " << num_batch + << " vs " << input_shape.dim_size(0); + return kRetry; + } + auto dtype = cuda_engine->getBindingDataType(binding_index); + switch (dtype) { + case nvinfer1::DataType::kFLOAT: + buffers[binding_index] = (void*)(input_tensor.flat().data()); + break; + case nvinfer1::DataType::kHALF: + LOG(ERROR) << "FP16 inputs are not supported yet!"; + return kRetry; + case nvinfer1::DataType::kINT8: + LOG(ERROR) << "INT8 inputs are not supported yet!"; + return kRetry; + case nvinfer1::DataType::kINT32: + buffers[binding_index] = (void*)(input_tensor.flat().data()); + break; + default: + LOG(ERROR) << "Unknown TRT data type: " << int(dtype); + return kRetry; + } + } + + for (int i = 0; i < ctx->num_outputs(); i++) { + // Create an output tensor + const string output_name = StrCat(kOutputPHName, i); + const int binding_index = cuda_engine->getBindingIndex(output_name.c_str()); + Tensor* output_tensor = nullptr; + + TensorShape output_shape; + if (binding_index != -1) { + 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]; + auto status = TensorShapeUtils::MakeShape( + trt_shape.data(), trt_shape.size(), &output_shape); + if (!status.ok()) { + LOG(ERROR) << "Failed to get output shape: " << status; + return kRetry; + } + } else { + LOG(ERROR) << "Output node not found, at " << output_name; + return kRetry; + } + auto status = ctx->allocate_output(i, output_shape, &output_tensor); + if (!status.ok()) { + LOG(ERROR) << "Allocating output failed with " << status; + ctx->SetStatus(status); + // Do not retry since we cannot allocate the same output twice. + // TODO(aaroey): ideally we should retry, fix this. + return !kRetry; + } + auto dtype = cuda_engine->getBindingDataType(binding_index); + switch (dtype) { + case nvinfer1::DataType::kFLOAT: + buffers[binding_index] = + reinterpret_cast(output_tensor->flat().data()); + break; + case nvinfer1::DataType::kHALF: + LOG(WARNING) << "half size is not supported yet!"; + return kRetry; + case nvinfer1::DataType::kINT8: + LOG(WARNING) << "int8 is not supported yet!"; + return kRetry; + case nvinfer1::DataType::kINT32: + buffers[binding_index] = + reinterpret_cast(output_tensor->flat().data()); + break; + default: + LOG(WARNING) << "Unknown TRT data type: " << static_cast(dtype); + return kRetry; + } + } + // Copied from cuda_kernel_helper since it seems only valid in *.cu.cc files + const cudaStream_t* stream = CHECK_NOTNULL( + reinterpret_cast(ctx->op_device_context() + ->stream() + ->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 = 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; + } + test::AddTestValue(StrCat(name(), ":ExecuteTrtEngine"), "done"); + // Synchronization will be done by TF. + return !kRetry; +} + +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; + } + tensorflow::core::ScopedUnref sc(cache_res); + auto& cache = cache_res->cache_; + auto allocator = cache_res->allocator_.get(); + if (allocator == nullptr) { + return &empty_context; + } + + // Handle the static engine case. For static engines, the cache will have a + // single element containing the only engine. + if (static_engine_) { + 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 &empty_context; + } + + TrtUniquePtrType infer(nvinfer1::createInferRuntime(logger)); + infer->setGpuAllocator(allocator); + TrtUniquePtrType static_engine( + infer->deserializeCudaEngine(serialized_segment_.c_str(), + serialized_segment_.size(), + PluginFactoryTensorRT::GetInstance())); + auto raw_static_engine = static_engine.get(); + const auto max_batch_size = raw_static_engine->getMaxBatchSize(); + // 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 &empty_context; + } + return cache.at(engine_input_shapes).get(); + } // static_engine_ + + // Handle the dynamic engine case. + // See if there is a compatible engine cached. The batch size should be <= the + // cached batch size. + std::vector engine_input_shapes; + 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() + << " input shapes: " + << TensorShapeUtils::ShapeListString(engine_input_shapes); + // Convert to partial shapes + std::vector partial_shapes; + for (int i = 0; i < engine_input_shapes.size(); i++) { + partial_shapes.emplace_back(engine_input_shapes[i]); + } + // 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_, + 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. + cache.emplace(engine_input_shapes, absl::make_unique()); + } + LOG(WARNING) << "Engine creation for batch size " << batch_size + << " failed " << status; + return &empty_context; + } + VLOG(1) << "Conversion is done"; + TrtUniquePtrType exec_context( + engine->createExecutionContext()); + cache.emplace(engine_input_shapes, + absl::make_unique(std::move(engine), + std::move(exec_context))); + } + return cache.at(engine_input_shapes).get(); +} + +tensorflow::Status TRTEngineOp::AllocateCalibrationResources( + OpKernelContext* ctx, SerializableResourceBase** cr) { + auto cres = new TRTCalibrationResource(); + *cr = cres; + // Get the allocator. + auto alloc = ctx->device()->GetAllocator(tensorflow::AllocatorAttributes()); + if (!alloc) { + LOG(WARNING) << "Can't get device allocator will not be able to " + "allocate memory from TensorFlow memory pool"; + cres->allocator_.reset(new TRTCudaAllocator); + } else { + cres->allocator_.reset(new TRTDeviceAllocator(alloc)); + } + // Get the input shapes. + const int batch_size = ctx->input(0).dim_size(0); + const int num_inputs = ctx->num_inputs(); + std::vector shapes; + cres->device_tensors_.resize(num_inputs); + VLOG(1) << " Constructing calibrator"; + for (int i = 0; i < num_inputs; i++) { + // allocate workspace on device for inputs + const tensorflow::Tensor& t = ctx->input(i); + shapes.emplace_back(t.shape()); + Tensor* device_tensor; + TF_RETURN_IF_ERROR(ctx->allocate_persistent( + 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); + } + cres->device_buffers_.emplace( + StrCat(kInputPHName, i), + std::pair(device_address, device_tensor->TotalBytes())); + } + cres->calibrator_.reset( + new TRTInt8Calibrator(cres->device_buffers_, batch_size, name())); + const string label(name()); + auto segment_graph = &segment_graph_; + const int platform_gpu_id = + ctx->device()->tensorflow_gpu_device_info()->gpu_id; + if (platform_gpu_id < 0) { + LOG(ERROR) << "Can't get gpu_device_info from context->device()"; + return tensorflow::errors::InvalidArgument( + "Context->device doesn't contain device info!"); + } + const int64 workspace_size_bytes = workspace_size_; + cres->thr_.reset(new std::thread([cres, label, segment_graph, shapes, + platform_gpu_id, workspace_size_bytes]() { + LOG(INFO) << "Starting calibration thread on device " << platform_gpu_id + << ", Calibration Resource @ " << cres; + auto err = cudaSetDevice(platform_gpu_id); + if (err != cudaSuccess) { + // TODO(aaroey): should return error here. + LOG(ERROR) << "Couldn't set cuda device to " << platform_gpu_id + << " in calibration thread"; + } + // ConvertGraphDefToEngine() will try to build the engine. This thread + // will loop inside buildCudaEngine() consuming the calibration data + // that is set by the TF op, and drive the builder until calibrator returns + // false. Engine is discarded after calibration table is generated + // + // 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_, + /*use_calibration=*/true, + /*convert_successfully=*/nullptr); + if (!s.ok()) { + LOG(ERROR) << "Calibration failed: " << s; + cres->calibrator_->setDone(); // Ignore further pushes + } + VLOG(1) << "Calibration loop terminated " << label; + })); + VLOG(1) << "initialized calibrator resource"; + return tensorflow::Status::OK(); +} + +REGISTER_KERNEL_BUILDER(Name("TRTEngineOp").Device(DEVICE_GPU), TRTEngineOp); + +} // namespace tensorrt +} // namespace tensorflow + +#endif // GOOGLE_TENSORRT +#endif // GOOGLE_CUDA diff --git a/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.h b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.h new file mode 100644 index 0000000000000000000000000000000000000000..64f8c97a74092ac075de9cc7993283e3ce1e27cf --- /dev/null +++ b/tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.h @@ -0,0 +1,130 @@ +/* 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_KERNELS_TRT_ENGINE_OP_H_ +#define TENSORFLOW_COMPILER_TF2TENSORRT_KERNELS_TRT_ENGINE_OP_H_ + +#include +#include + +#include "tensorflow/compiler/tf2tensorrt/convert/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_resources.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" +#include "tensorflow/core/platform/thread_annotations.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 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; + + private: + // TODO(samikama): context should go to a resource manager! + + // 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. + 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_; + + // 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_COMPILER_TF2TENSORRT_KERNELS_TRT_ENGINE_OP_H_ diff --git a/tensorflow/compiler/tf2tensorrt/ops/get_serialized_resource_op.cc b/tensorflow/compiler/tf2tensorrt/ops/get_serialized_resource_op.cc new file mode 100644 index 0000000000000000000000000000000000000000..59da73f5efc8eedc20c35cf35cb1eae6cda136c9 --- /dev/null +++ b/tensorflow/compiler/tf2tensorrt/ops/get_serialized_resource_op.cc @@ -0,0 +1,40 @@ +/* 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. +==============================================================================*/ + +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT + +#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 { + +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"); + +} // namespace tensorflow + +#endif // GOOGLE_TENSORRT +#endif // GOOGLE_CUDA diff --git a/tensorflow/compiler/tf2tensorrt/ops/trt_engine_op.cc b/tensorflow/compiler/tf2tensorrt/ops/trt_engine_op.cc new file mode 100644 index 0000000000000000000000000000000000000000..b84d2fe0b8cef3475f2a7d0f5383d5e11cde099a --- /dev/null +++ b/tensorflow/compiler/tf2tensorrt/ops/trt_engine_op.cc @@ -0,0 +1,65 @@ +/* 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. +==============================================================================*/ + +#if GOOGLE_CUDA +#if GOOGLE_TENSORRT + +#include "tensorflow/core/framework/common_shape_fns.h" +#include "tensorflow/core/framework/op.h" +#include "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/framework/shape_inference.h" +#include "tensorflow/core/framework/tensor_shape.h" + +namespace tensorflow { + +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,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) >= 0 = []") + .Attr("max_cached_engines_count: int = 1") + .Attr("workspace_size_bytes: int") + .Attr("precision_mode: {'FP32', 'FP16', 'INT8'}") + .Attr("calibration_data: string = ''") + .Attr("use_calibration: bool = true") + .Input("in_tensor: InT") + .Output("out_tensor: OutT") + // TODO(jie): TF requires concrete output shape for concrete input shapes. + // This is tricky for batch dimension, since we cannot ensure which input + // would carry the correct batch dimension (for the current stage of the + // implementation, we do require all input tensor to carry the same batch + // size, but this could change in the future). Hence we disable shape + // inference function as a workaround. + // .SetShapeFn(shape_inference::TRTEngineOpShapeInference); + .SetShapeFn(shape_inference::UnknownShape); +} // namespace tensorflow + +#endif // GOOGLE_TENSORRT +#endif // GOOGLE_CUDA 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..86bfabf99e08a8e447a28504c72eebca4d3a582c --- /dev/null +++ b/tensorflow/compiler/tf2tensorrt/python/ops/trt_ops.py @@ -0,0 +1,34 @@ +# 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 platform + +if platform.system() != "Windows": + # pylint: disable=wildcard-import,unused-import,g-import-not-at-top + from tensorflow.compiler.tf2tensorrt.ops.gen_trt_ops import * + + from tensorflow.python.framework import load_library + from tensorflow.python.platform import resource_loader + # pylint: enable=wildcard-import,unused-import,g-import-not-at-top + + _trt_ops = load_library.load_op_library( + resource_loader.get_path_to_datafile("_trt_ops.so")) +else: + raise RuntimeError("Windows platforms are not supported") diff --git a/tensorflow/contrib/tensorrt/segment/segment.cc b/tensorflow/compiler/tf2tensorrt/segment/segment.cc similarity index 88% rename from tensorflow/contrib/tensorrt/segment/segment.cc rename to tensorflow/compiler/tf2tensorrt/segment/segment.cc index c82d4a018392be19a0bae5893158c7180f15acc3..4a8a4ac7589a4b68b129e8e88ee999e8a2495728 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" @@ -32,7 +33,8 @@ limitations under the License. namespace tensorflow { namespace tensorrt { namespace segment { -using ::tensorflow::strings::StrAppend; +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 @@ -224,6 +226,24 @@ SimpleGraph::~SimpleGraph() { for (auto x : edges_) delete x; } +// Define comparison functions for std::set with pointer keys so that behavior +// is deterministic. When using std::set with pointer key types, the items are +// sorted by pointer address which is non-deterministic. This can cause issues +// for INT8 mode because the graph is converted twice and non-determinism may +// cause a mismatch between the calibration tables of the conversions. +struct SimpleEdgePtrCompare { + bool operator()(const SimpleEdge* lhs, const SimpleEdge* rhs) const { + return lhs->id() < rhs->id(); + } +}; + +struct NodePtrCompare { + bool operator()(const tensorflow::Node* lhs, + const tensorflow::Node* rhs) const { + return lhs->name() < rhs->name(); + } +}; + namespace { // Copied from TF ReverseDFS, which only works for tensorflow::Graph. @@ -389,7 +409,7 @@ void ContractEdge(SimpleEdge* edge, SimpleGraph* graph, tensorflow::Status SegmentGraph( const tensorflow::Graph* tf_graph, - const std::function& candidate_fn, + const std::function& candidate_fn, const std::function& input_candidate_fn, const std::function& output_candidate_fn, const SegmentOptions& options, SegmentNodesVector* segments) { @@ -406,15 +426,42 @@ tensorflow::Status SegmentGraph( // Use a union-find to collect the nodes that belong to the same // segment. A node value of nullptr indicates that the node is not a candidate // for TRT. + std::unordered_set unsupported_ops; + int num_unsupported_ops = 0; std::vector> node_segments; for (int i = 0; i < graph->num_node_ids(); ++i) { SimpleNode* node = graph->FindNodeId(i); - if (options.exclude_node_list.count(node->name()) != 0 || - !candidate_fn(node->tf_node())) { + if (options.exclude_node_list.count(node->name()) != 0) { + VLOG(1) << "Not a TF-TRT candidate, " + << "(Op type: " << node->tf_node()->type_string() << "), " + << "(Op name: " << node->name() << "), " + << "(Reason: excluded by segmenter option)"; + unsupported_ops.emplace(node->tf_node()->type_string()); + num_unsupported_ops++; node = nullptr; + } else { + const Status status = candidate_fn(node->tf_node()); + if (!status.ok()) { + VLOG(1) << "Not a TF-TRT candidate, " + << "(Op type: " << node->tf_node()->type_string() << "), " + << "(Op name: " << node->name() << "), " + << "(Reason: " << status << ")"; + unsupported_ops.emplace(node->tf_node()->type_string()); + num_unsupported_ops++; + node = nullptr; + } } node_segments.emplace_back(node); } + string msg = StrCat( + "There are ", num_unsupported_ops, " ops of ", unsupported_ops.size(), + " different types in the graph that", " are not converted to TensorRT: "); + for (const auto& elem : unsupported_ops) { + StrAppend(&msg, elem, ", "); + } + LOG(INFO) << msg << "(For more information see " + << "https://docs.nvidia.com/deeplearning" + << "/dgx/integrate-tf-trt/index.html#support-ops)."; // The segmentation algorithm below visits nodes in reverse topological order // and attempts to merge nodes along output edges. That means that subgraphs @@ -448,7 +495,7 @@ tensorflow::Status SegmentGraph( // nodes. Iterate since combining two nodes may unblock other // combining. while (true) { - std::set contract_edges; + std::set contract_edges; for (const SimpleEdge* out_edge : node->out_edges()) { VLOG(3) << "... out node " << out_edge->dst()->name() << " ( " << out_edge->dst()->id() << " <- " << node->id() << " )"; @@ -502,7 +549,7 @@ tensorflow::Status SegmentGraph( // A map from the segment identifier (currently the name of the root node of // the segment tree) to the segment nodes set. - std::map> sg_map; + std::map> sg_map; // A map from the segment identifier (currently the name of the root node of // the segment tree) to the device names that the nodes in the segment are @@ -538,7 +585,8 @@ tensorflow::Status SegmentGraph( // --------------------------------- Step 2 --------------------------------- // Remove ineligible input/output nodes. for (auto& itr : sg_map) { - std::set& segment_nodes = itr.second; + std::set& segment_nodes = + itr.second; VLOG(1) << "Segment original size: " << segment_nodes.size(); while (true) { std::deque in_nodes_que, out_nodes_que; @@ -590,8 +638,9 @@ tensorflow::Status SegmentGraph( bool is_input_nodes, std::deque* que) { // Run a BFS on the queue to find all the input/output nodes. - std::set visited; - std::set logged(que->begin(), que->end()); + std::set visited; + std::set logged(que->begin(), + que->end()); while (!que->empty()) { auto node = que->front(); que->pop_front(); @@ -625,9 +674,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; } @@ -640,12 +691,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: "); @@ -654,10 +703,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)) { diff --git a/tensorflow/contrib/tensorrt/segment/segment.h b/tensorflow/compiler/tf2tensorrt/segment/segment.h similarity index 77% rename from tensorflow/contrib/tensorrt/segment/segment.h rename to tensorflow/compiler/tf2tensorrt/segment/segment.h index 8c44eb782aa37052680d0e06023f29dc65e327c6..9a0ccc9aef475edfb0ffb83a2be21d4d4ca0e028 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 @@ -29,10 +29,10 @@ 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. @@ -43,7 +43,7 @@ struct SegmentOptions { // Get the subgraphs of a graph that can be handled by TensorRT. // // @param graph tensorflow::Graph of the network -// @param candidate_fn A function that returns true for a Node* if +// @param candidate_fn A function that returns OK for a Node* if // that node can be handled by TensorRT. // @param segments Returns the TensorRT segments/subgraphs. Each entry // in the vector describes a subgraph by giving a set of the names of @@ -51,7 +51,7 @@ struct SegmentOptions { // @return the status. tensorflow::Status SegmentGraph( const tensorflow::Graph* tf_graph, - const std::function& candidate_fn, + const std::function& candidate_fn, const std::function& input_candidate_fn, const std::function& output_candidate_fn, const SegmentOptions& options, SegmentNodesVector* segments); @@ -60,4 +60,4 @@ tensorflow::Status SegmentGraph( } // namespace tensorrt } // namespace tensorflow -#endif // TENSORFLOW_CONTRIB_TENSORRT_SEGMENT_SEGMENT_H_ +#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 95% rename from tensorflow/contrib/tensorrt/segment/segment_test.cc rename to tensorflow/compiler/tf2tensorrt/segment/segment_test.cc index 5937fa8259a39339e92b150862d195ee1f23f70a..58512d3b09d7c6f523710bc09843c628a5838b53 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" @@ -34,10 +34,13 @@ namespace ops = ::tensorflow::ops; class SegmentTest : public ::testing::Test { protected: - std::function MakeCandidateFn( + std::function MakeCandidateFn( const std::set& node_names) { - return [node_names](const tensorflow::Node* node) -> bool { - return node_names.find(node->name()) != node_names.end(); + return [node_names](const tensorflow::Node* node) -> Status { + if (node_names.find(node->name()) != node_names.end()) { + return Status::OK(); + } + return errors::NotFound(""); }; } @@ -72,7 +75,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)) 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/compiler/tf2tensorrt/utils/test_utils.cc b/tensorflow/compiler/tf2tensorrt/utils/test_utils.cc new file mode 100644 index 0000000000000000000000000000000000000000..3bcca99afbff8b84d2dd628ae9211ee94e86af2a --- /dev/null +++ b/tensorflow/compiler/tf2tensorrt/utils/test_utils.cc @@ -0,0 +1,101 @@ +/* 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/test_utils.h" + +#include +#include + +#include "re2/re2.h" +#include "tensorflow/core/platform/macros.h" + +namespace tensorflow { +namespace tensorrt { +namespace test { + +// TODO(aaroey): make this class thread-safe. +class TestValueManager { + public: + static TestValueManager* singleton() { + static TestValueManager* manager = new TestValueManager(); + return manager; + } + + void Enable() { + VLOG(1) << "Enabling test value"; + enabled_ = true; + } + + void Add(const string& label, const string& value) { + if (TF_PREDICT_FALSE(enabled_)) { + QCHECK_NE("", value); + VLOG(1) << "Adding test value: " << label << " -> " << value; + values_.insert({label, value}); + } + } + + string Get(const string& label) { + if (TF_PREDICT_FALSE(enabled_)) { + VLOG(1) << "Getting test value by " << label; + auto itr = values_.find(label); + if (itr == values_.end()) return ""; + return itr->second; + } + return ""; + } + + void Clear(const string& pattern) { + if (TF_PREDICT_FALSE(enabled_)) { + VLOG(1) << "Clearing test values"; + if (pattern.empty()) { + values_.clear(); + return; + } + std::vector keys_to_clear; + for (const auto& kv : values_) { + if (RE2::FullMatch(kv.first, pattern)) { + keys_to_clear.push_back(kv.first); + } + } + for (const string& key : keys_to_clear) { + values_.erase(key); + } + } + } + + private: + TestValueManager() : enabled_(false) {} + + bool enabled_; + std::unordered_map values_; +}; + +void EnableTestValue() { TestValueManager::singleton()->Enable(); } + +void ClearTestValues(const string& pattern) { + TestValueManager::singleton()->Clear(pattern); +} + +void AddTestValue(const string& label, const string& value) { + TestValueManager::singleton()->Add(label, value); +} + +string GetTestValue(const string& label) { + return TestValueManager::singleton()->Get(label); +} + +} // namespace test +} // namespace tensorrt +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2tensorrt/utils/test_utils.h b/tensorflow/compiler/tf2tensorrt/utils/test_utils.h new file mode 100644 index 0000000000000000000000000000000000000000..bcd628b62f0320f7ce9dfe6240316d876f1d5a20 --- /dev/null +++ b/tensorflow/compiler/tf2tensorrt/utils/test_utils.h @@ -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. +==============================================================================*/ + +#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" + +namespace tensorflow { +namespace tensorrt { +namespace test { + +// Helper methods to inject values used by testing tools. +void EnableTestValue(); +void ClearTestValues(const string& pattern); +void AddTestValue(const string& label, const string& value); +string GetTestValue(const string& label); + +#define TRT_RETURN_IF_TEST_VALUE(label, value_to_return) \ + do { \ + if (::tensorflow::tensorrt::test::GetTestValue(label) == \ + value_to_return) { \ + return errors::Internal("Injected manually"); \ + } \ + } while (0) + +} // namespace test +} // namespace tensorrt +} // namespace tensorflow + +#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 80% rename from tensorflow/contrib/tensorrt/resources/trt_allocator_test.cc rename to tensorflow/compiler/tf2tensorrt/utils/trt_allocator_test.cc index ad6b1d7d4c57d696d3dee3b479733e152e669211..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" @@ -48,11 +48,14 @@ TEST(TRTAllocatorTest, Align) { 513ul, 700ul, 12345ul, 1ul << 32}) { for (uint64_t alignment = 1; alignment <= space * 4; alignment *= 2) { for (const uintptr_t ptr_val : - {1ul, alignment == 1 ? 1ul : alignment - 1, alignment, alignment + 1, - alignment + (alignment / 2)}) { + {static_cast(1), + alignment == 1 ? static_cast(1) : alignment - 1, + alignment, alignment + 1, alignment + (alignment / 2)}) { if (ptr_val % alignment == 0) { for (const uint64_t size : - {1ul, space == 1 ? 1ul : space - 1, space, space + 1}) { + {static_cast(1), + space == 1 ? static_cast(1) : space - 1, space, + space + 1}) { EXPECT_EQ(space >= size, RunTest(alignment, size, ptr_val, space)); } } else { @@ -62,8 +65,10 @@ TEST(TRTAllocatorTest, Align) { EXPECT_TRUE( RunTest(alignment, space - diff, ptr_val + diff, space - diff)); for (const uint64_t size : - {1ul, space - diff > 1 ? space - diff - 1 : 1ul, space - diff, - space - diff + 1, space - 1}) { + {static_cast(1), + space - diff > 1 ? space - diff - 1 + : static_cast(1), + space - diff, space - diff + 1, space - 1}) { EXPECT_EQ(space - diff >= size, RunTest(alignment, size, ptr_val, space)); } diff --git a/tensorflow/contrib/tensorrt/resources/trt_int8_calibrator.cc b/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.cc similarity index 98% rename from tensorflow/contrib/tensorrt/resources/trt_int8_calibrator.cc rename to tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.cc index dab1dd9343be7d5b033a3e04bf0b49fbbf37e9e5..bf111d3a2ee2fbec9151d12bbb6ff7181761c2aa 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 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 96% rename from tensorflow/contrib/tensorrt/log/trt_logger.cc rename to tensorflow/compiler/tf2tensorrt/utils/trt_logger.cc index dda0dc9e712eb726800abfb6084f4f708d04825b..c48bd6bf7747d1646c4e450b780822728e8573f1 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 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..37f7fe99fbb2b9e121953fc0de211db1bbf34b7a --- /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->size()) { + 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 f0e7791e9811533502fae0d4dea5a2e1ca2cf33c..92ba474fbcd085e3e33ceea4395cca4034969bd9 100644 --- a/tensorflow/compiler/tf2xla/BUILD +++ b/tensorflow/compiler/tf2xla/BUILD @@ -9,6 +9,7 @@ package_group( "//tensorflow/compiler/jit/...", "//tensorflow/compiler/tests/...", "//tensorflow/compiler/tf2xla/...", + "//tensorflow/contrib/compiler/...", ], ) @@ -166,6 +167,7 @@ cc_library( "xla_compilation_device.cc", "xla_compiler.cc", "xla_context.cc", + "xla_expression.cc", "xla_helpers.cc", "xla_op_kernel.cc", "xla_op_registry.cc", @@ -180,6 +182,7 @@ cc_library( "xla_compilation_device.h", "xla_compiler.h", "xla_context.h", + "xla_expression.h", "xla_helpers.h", "xla_op_kernel.h", "xla_op_registry.h", @@ -193,20 +196,23 @@ cc_library( ":sharding_util", ":side_effect_util", ":tf2xla_util", + "//tensorflow/compiler/jit:flags", + "//tensorflow/compiler/jit:xla_cluster_util", "//tensorflow/compiler/tf2xla/lib:util", "//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", + "//tensorflow/compiler/xla/client", "//tensorflow/compiler/xla/client:client_library", "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/client:xla_builder", "//tensorflow/compiler/xla/client:xla_computation", "//tensorflow/compiler/xla/client/lib:arithmetic", "//tensorflow/compiler/xla/client/lib:constants", - "//tensorflow/compiler/xla/client/lib:numeric", "//tensorflow/core:core_cpu", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", @@ -214,8 +220,12 @@ cc_library( "//tensorflow/core:lib_internal", "//tensorflow/core:protos_all_cc", "//tensorflow/core:stream_executor_no_cuda", + "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/memory", + "@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, ) @@ -236,6 +246,7 @@ cc_library( deps = [ "//tensorflow/compiler/xla:literal", "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", @@ -360,8 +371,12 @@ tf_cc_test( tf_cc_test( name = "xla_compiler_test", - srcs = ["xla_compiler_test.cc"], + srcs = [ + "xla_compiler_test.cc", + "xla_expression_test.cc", + ], deps = [ + ":common", ":side_effect_util", ":xla_compiler", "//tensorflow/cc:cc_ops", @@ -384,6 +399,7 @@ tf_cc_test( "//tensorflow/core:test", "//tensorflow/core:test_main", "//tensorflow/core:testlib", + "@com_google_absl//absl/memory", "@com_google_absl//absl/strings", ], ) @@ -413,6 +429,7 @@ tf_cc_test( "//tensorflow/cc:cc_ops", "//tensorflow/cc:function_ops", "//tensorflow/cc:ops", + "//tensorflow/compiler/jit:xla_cluster_util", "//tensorflow/compiler/tf2xla/kernels:xla_ops", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:ops", @@ -425,21 +442,15 @@ cc_library( name = "dump_graph", srcs = [ "dump_graph.cc", - "dump_graph_flags.cc", - "dump_graph_flags.h", ], hdrs = [ "dump_graph.h", ], deps = [ - "//tensorflow/compiler/xla/legacy_flags:parse_flags_from_env", - "//tensorflow/core:core_cpu", - "//tensorflow/core:core_cpu_internal", + "//tensorflow/compiler/jit:flags", "//tensorflow/core:framework", - "//tensorflow/core:framework_internal", - "//tensorflow/core:lib", + "//tensorflow/core:graph", "//tensorflow/core:protos_all_cc", - "@com_google_absl//absl/strings", ], ) @@ -660,6 +671,7 @@ 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", diff --git a/tensorflow/compiler/tf2xla/const_analysis.cc b/tensorflow/compiler/tf2xla/const_analysis.cc index 027ca6d2d2f616177d91d9d57d1ff373bab2a754..a57095f91e43f6b31b58e5a5f36331241451b545 100644 --- a/tensorflow/compiler/tf2xla/const_analysis.cc +++ b/tensorflow/compiler/tf2xla/const_analysis.cc @@ -18,6 +18,7 @@ limitations under the License. #include #include +#include "absl/algorithm/container.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "tensorflow/core/framework/node_def_util.h" #include "tensorflow/core/graph/algorithm.h" @@ -67,25 +68,18 @@ Status BackwardsConstAnalysis(const Graph& g, } // Mark any compile-time constant operator arguments as const. - const std::unordered_set* const_inputs = - XlaOpRegistry::CompileTimeConstantInputs(node->type_string()); - if (!const_inputs || const_inputs->empty()) return; + std::vector const_input_idxs; + status = XlaOpRegistry::CompileTimeConstantInputs( + node->def(), node->op_def(), &const_input_idxs); - NameRangeMap input_name_ranges; - status = - NameRangesForNode(*node, node->op_def(), &input_name_ranges, nullptr); - if (!status.ok()) return; - - for (const string& input : *const_inputs) { - auto name_range = input_name_ranges.find(input); - if (name_range == input_name_ranges.end()) continue; + if (!status.ok()) { + return; + } - for (Edge const* edge : node->in_edges()) { - if (edge->dst_input() >= name_range->second.first && - edge->dst_input() < name_range->second.second && - edge_filter(*edge)) { - (*compile_time_const_nodes)[edge->src()->id()] = true; - } + for (Edge const* edge : node->in_edges()) { + if (absl::c_binary_search(const_input_idxs, edge->dst_input()) && + edge_filter(*edge)) { + (*compile_time_const_nodes)[edge->src()->id()] = true; } } }; diff --git a/tensorflow/compiler/tf2xla/const_analysis_test.cc b/tensorflow/compiler/tf2xla/const_analysis_test.cc index 56065be894697bc72ecc0089c665c19aafee7bf8..40c6d0e01701d9104a200d9ea27706a0a7c12146 100644 --- a/tensorflow/compiler/tf2xla/const_analysis_test.cc +++ b/tensorflow/compiler/tf2xla/const_analysis_test.cc @@ -20,6 +20,7 @@ limitations under the License. #include "tensorflow/cc/framework/ops.h" #include "tensorflow/cc/ops/function_ops.h" #include "tensorflow/cc/ops/standard_ops.h" +#include "tensorflow/compiler/jit/xla_cluster_util.h" #include "tensorflow/core/graph/algorithm.h" #include "tensorflow/core/lib/core/status_test_util.h" #include "tensorflow/core/platform/test.h" @@ -107,5 +108,54 @@ TEST(ConstAnalysisTest, DontFollowControlDependencies) { EXPECT_EQ(const_args, std::vector({false, true})); } +TEST(ConstAnalysisTest, RespectExplicitAttr_0) { + Scope root = Scope::NewRootScope(); + + Output arg0 = ops::_Arg(root.WithOpName("Arg0"), DT_INT32, 0); + Output arg1 = ops::_Arg(root.WithOpName("Arg1"), DT_INT32, 1); + Output c1 = + ops::Const(root.WithOpName("c1").WithControlDependencies(arg0), 1, {1}); + Output add = ops::Add(root, arg1, c1); + + // Force const analysis to pretend that the shape argument to `reshape` does + // not need to be a constant. + Output reshape = ops::Reshape(root, arg1, add); + reshape.node()->AddAttr(kXlaCompileTimeConstantInputsAttr, + std::vector()); + + Graph graph(OpRegistry::Global()); + TF_ASSERT_OK(root.ToGraph(&graph)); + + std::vector const_args(2, false); + TF_ASSERT_OK(BackwardsConstAnalysis(graph, &const_args, + /*compile_time_const_nodes=*/nullptr)); + + EXPECT_EQ(const_args, std::vector({false, false})); +} + +TEST(ConstAnalysisTest, RespectExplicitAttr_1) { + Scope root = Scope::NewRootScope(); + + Output arg0 = ops::_Arg(root.WithOpName("Arg0"), DT_INT32, 0); + Output c1 = + ops::Const(root.WithOpName("c1").WithControlDependencies(arg0), 1, {1}); + Output add = ops::Add(root, arg0, c1); + + // Force const analysis to pretend that the first argument to `add` needs to + // be a constant. + std::vector add_constant_inputs; + add_constant_inputs.push_back("x"); + add.node()->AddAttr(kXlaCompileTimeConstantInputsAttr, add_constant_inputs); + + Graph graph(OpRegistry::Global()); + TF_ASSERT_OK(root.ToGraph(&graph)); + + std::vector const_args(1, false); + TF_ASSERT_OK(BackwardsConstAnalysis(graph, &const_args, + /*compile_time_const_nodes=*/nullptr)); + + EXPECT_EQ(const_args, std::vector({true})); +} + } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/cpu_function_runtime.h b/tensorflow/compiler/tf2xla/cpu_function_runtime.h index dfc1e8b8aebcf3142e9f61f60171c6b58634c71d..78970fb39bae7067c7668baa2aec65732b5b2352 100644 --- a/tensorflow/compiler/tf2xla/cpu_function_runtime.h +++ b/tensorflow/compiler/tf2xla/cpu_function_runtime.h @@ -104,7 +104,7 @@ class BufferInfo { private: BufferInfo() = default; - enum class Kind : unsigned { + enum class Kind : uint64 { kConstant, kTempBuffer, kEntryParameter, diff --git a/tensorflow/compiler/tf2xla/dump_graph.cc b/tensorflow/compiler/tf2xla/dump_graph.cc index 380c6a7e23da92d949b26876836b999bf6406c6c..64fdbbebc65bff4ed0b965fcdd534cc9696472b6 100644 --- a/tensorflow/compiler/tf2xla/dump_graph.cc +++ b/tensorflow/compiler/tf2xla/dump_graph.cc @@ -18,87 +18,26 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/dump_graph.h" -#include "absl/strings/str_cat.h" -#include "tensorflow/compiler/tf2xla/dump_graph_flags.h" -#include "tensorflow/core/platform/env.h" -#include "tensorflow/core/platform/mutex.h" +#include "tensorflow/compiler/jit/flags.h" +#include "tensorflow/core/util/dump_graph.h" namespace tensorflow { namespace dump_graph { -namespace { - -struct NameCounts { - mutex counts_mutex; - std::unordered_map counts; -}; - -string MakeUniqueFilename(string name) { - static NameCounts& instance = *new NameCounts; - - // Remove illegal characters from `name`. - for (int i = 0; i < name.size(); ++i) { - char ch = name[i]; - if (ch == '/' || ch == '[' || ch == ']' || ch == '*' || ch == '?') { - name[i] = '_'; - } - } - - int count; - { - mutex_lock lock(instance.counts_mutex); - count = instance.counts[name]++; - } - - string filename = name; - if (count > 0) { - absl::StrAppend(&filename, "_", count); - } - absl::StrAppend(&filename, ".pbtxt"); - return filename; -} - -string WriteTextProtoToUniqueFile( - Env* env, const string& name, const char* proto_type, - const ::tensorflow::protobuf::Message& proto) { - const string& dirname = - legacy_flags::GetDumpGraphFlags()->tf_dump_graph_prefix; - Status status = env->RecursivelyCreateDir(dirname); - if (!status.ok()) { - LOG(WARNING) << "Failed to create " << dirname << " for dumping " - << proto_type << ": " << status; - return "(unavailable)"; - } - string filepath = absl::StrCat(dirname, "/", MakeUniqueFilename(name)); - status = WriteTextProto(Env::Default(), filepath, proto); - if (!status.ok()) { - LOG(WARNING) << "Failed to dump " << proto_type << " to file: " << filepath - << " : " << status; - return "(unavailable)"; - } - LOG(INFO) << "Dumped " << proto_type << " to " << filepath; - return filepath; -} - -} // anonymous namespace - string DumpGraphDefToFile(const string& name, GraphDef const& graph_def) { - return WriteTextProtoToUniqueFile(Env::Default(), name, "GraphDef", - graph_def); + return tensorflow::DumpGraphDefToFile( + name, graph_def, GetDumpGraphFlags()->tf_dump_graph_prefix); } string DumpGraphToFile(const string& name, Graph const& graph, const FunctionLibraryDefinition* flib_def) { - GraphDef graph_def; - graph.ToGraphDef(&graph_def); - if (flib_def) { - *graph_def.mutable_library() = flib_def->ToProto(); - } - return DumpGraphDefToFile(name, graph_def); + return tensorflow::DumpGraphToFile(name, graph, flib_def, + GetDumpGraphFlags()->tf_dump_graph_prefix); } string DumpFunctionDefToFile(const string& name, FunctionDef const& fdef) { - return WriteTextProtoToUniqueFile(Env::Default(), name, "FunctionDef", fdef); + return tensorflow::DumpFunctionDefToFile( + name, fdef, GetDumpGraphFlags()->tf_dump_graph_prefix); } } // namespace dump_graph diff --git a/tensorflow/compiler/tf2xla/dump_graph_flags.cc b/tensorflow/compiler/tf2xla/dump_graph_flags.cc deleted file mode 100644 index a6c908ba011afb90fabacc855df8c6afbb35d254..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/tf2xla/dump_graph_flags.cc +++ /dev/null @@ -1,63 +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. -==============================================================================*/ - -// Legacy flags for the XLA bridge's dump_graph module. - -#include -#include - -#include "tensorflow/compiler/tf2xla/dump_graph_flags.h" -#include "tensorflow/compiler/xla/legacy_flags/parse_flags_from_env.h" -#include "tensorflow/core/platform/types.h" -#include "tensorflow/core/util/command_line_flags.h" - -namespace tensorflow { -namespace legacy_flags { - -// Pointers to the parsed value of the flags and flag descriptors, initialized -// via flags_init. -static DumpGraphFlags* flags; -static std::vector* flag_list; -static std::once_flag flags_init; - -// Allocate *flags. Called via call_once(&flags_init,...). -static void AllocateFlags() { - flags = new DumpGraphFlags; - flags->tf_dump_graph_prefix = "/tmp/"; - flag_list = new std::vector({ - Flag("tf_dump_graph_prefix", &flags->tf_dump_graph_prefix, - "Path prefix to which graphs dumped during debugging should be " - "written."), - }); - xla::legacy_flags::ParseFlagsFromEnv(*flag_list); -} - -// Append to *append_to flag definitions associated with the XLA bridge's -// dump_graph module. -void AppendDumpGraphFlags(std::vector* append_to) { - std::call_once(flags_init, &AllocateFlags); - append_to->insert(append_to->end(), flag_list->begin(), flag_list->end()); -} - -// Return a pointer to the DumpGraphFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -DumpGraphFlags* GetDumpGraphFlags() { - std::call_once(flags_init, &AllocateFlags); - return flags; -} - -} // namespace legacy_flags -} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/dump_graph_flags.h b/tensorflow/compiler/tf2xla/dump_graph_flags.h deleted file mode 100644 index 80a3307d920f2cc3d668d507786a02e43589f86f..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/tf2xla/dump_graph_flags.h +++ /dev/null @@ -1,48 +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_TF2XLA_DUMP_GRAPH_FLAGS_H_ -#define TENSORFLOW_COMPILER_TF2XLA_DUMP_GRAPH_FLAGS_H_ - -// Legacy flags for the XLA bridge's dump_graph module. - -#include - -#include "tensorflow/core/platform/types.h" -#include "tensorflow/core/util/command_line_flags.h" - -namespace tensorflow { -namespace legacy_flags { - -// Append to *flag_list flag definitions associated with the XLA bridge's -// dump_graph module. -void AppendDumpGraphFlags(std::vector* flag_list); - -// The values of flags associated with the XLA bridge's -// dump_graph module. -typedef struct { - string tf_dump_graph_prefix; // Path prefix to which graphs dumped during - // debugging should be written. -} DumpGraphFlags; - -// Return a pointer to the DumpGraphFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -DumpGraphFlags* GetDumpGraphFlags(); - -} // namespace legacy_flags -} // namespace tensorflow - -#endif // TENSORFLOW_COMPILER_TF2XLA_DUMP_GRAPH_FLAGS_H_ diff --git a/tensorflow/compiler/tf2xla/functionalize_cond.cc b/tensorflow/compiler/tf2xla/functionalize_cond.cc index 46649b8cc43016d4a62f49e20256c77ca8accc79..7ae96e1d484900e28e8c23c3bb2232401144ad82 100644 --- a/tensorflow/compiler/tf2xla/functionalize_cond.cc +++ b/tensorflow/compiler/tf2xla/functionalize_cond.cc @@ -339,6 +339,7 @@ Status Conditional::AddSwitch(Node* s) { DebugString(switch_predicate_), " vs ", DebugString(predicate), ")."); } switches_.insert(s); + parent_->AddSwitchId(s->id()); return Status::OK(); } @@ -639,7 +640,8 @@ Status Conditional::ExtractBodies(Graph* graph) { Status Conditional::BuildIfNode(Graph* graph, FunctionLibraryDefinition* library) { VLOG(2) << "Build cond function for " << name(); - NodeDefBuilder builder(name(), "If", library); + NodeDebugInfo debug_info((*merges_.begin())->def()); + NodeDefBuilder builder(name(), "If", library, &debug_info); const string branch_name[] = {"else_branch", "then_branch"}; for (auto branch : {BranchType::kElseBranch, BranchType::kThenBranch}) { int branch_index = static_cast(branch); @@ -1185,7 +1187,7 @@ Status FunctionalizeCond::DetermineAncestorState(Node* dst) { } void FunctionalizeCond::DeleteReachableAndDeadNodes( - const std::vector& switch_ids, const std::vector& merge_order) { + const std::vector& merge_order) { // Delete all nodes that have been extracted or are reachable from // deleted/dead nodes. The input and outgoing edges should have already been // removed. @@ -1197,7 +1199,7 @@ void FunctionalizeCond::DeleteReachableAndDeadNodes( // All remaining Switch nodes are not reachable from a Merge node and // removed. This is to account for dead Switch nodes. - for (int s_id : switch_ids) { + for (int s_id : switch_ids_) { Node* s = graph_->FindNodeId(s_id); if (s == nullptr) continue; for (const Edge* e : s->out_edges()) { @@ -1288,11 +1290,10 @@ Status FunctionalizeCond::FunctionalizeInternal() { // reverse topological sorting); // * Record reverse topological for merge and switch nodes; std::vector rev_topo_order; - std::vector switch_ids; std::vector merge_order; DFS(*graph_, nullptr, [&](Node* n) { if (IsSwitch(n)) { - switch_ids.push_back(n->id()); + AddSwitchId(n->id()); } if (IsMerge(n)) { merge_order.push_back(n); @@ -1306,9 +1307,7 @@ Status FunctionalizeCond::FunctionalizeInternal() { if (merge_order.empty()) { // No merges mean no switch values consumed (as only considering values // fetchable as output of merge); - for (auto it = switch_ids.begin(); it != switch_ids.end(); ++it) { - graph_->RemoveNode(graph_->FindNodeId(*it)); - } + DeleteReachableAndDeadNodes(merge_order); return Status::OK(); } @@ -1351,7 +1350,7 @@ Status FunctionalizeCond::FunctionalizeInternal() { if (VLOG_IS_ON(4)) DumpGraphWithCondState("after_extract"); } - DeleteReachableAndDeadNodes(switch_ids, merge_order); + DeleteReachableAndDeadNodes(merge_order); return Status::OK(); } @@ -1371,6 +1370,10 @@ void FunctionalizeCond::DumpGraphWithCondState(const string& name) { library_); } +void FunctionalizeCond::AddSwitchId(int switch_id) { + switch_ids_.push_back(switch_id); +} + Status FunctionalizeCond::Functionalize(Graph* graph, FunctionLibraryDefinition* library) { VLOG(1) << "FunctionalizeCond::Functionalize"; diff --git a/tensorflow/compiler/tf2xla/functionalize_cond.h b/tensorflow/compiler/tf2xla/functionalize_cond.h index 189980894073b1da1a12d1c284536336eb920900..8525d7af61b4471e53a9ae16b081060bfd234c9c 100644 --- a/tensorflow/compiler/tf2xla/functionalize_cond.h +++ b/tensorflow/compiler/tf2xla/functionalize_cond.h @@ -166,6 +166,9 @@ class FunctionalizeCond { // Dump graph with the CondState annotated. void DumpGraphWithCondState(const string& name); + // Adds `switch_id` to the list of Switch node ids. + void AddSwitchId(int switch_id); + private: FunctionalizeCond(Graph* graph, FunctionLibraryDefinition* library); @@ -219,8 +222,7 @@ class FunctionalizeCond { // Deletes all nodes in/consumers reachable from switch/merge nodes that were // extracted. - void DeleteReachableAndDeadNodes(const std::vector& switch_ids, - const std::vector& merge_order); + void DeleteReachableAndDeadNodes(const std::vector& merge_order); // Member used to unique the CondState to a unique CondId (AncestorState to a // unique AncestorId) and keep track of CondState/CondId @@ -234,6 +236,8 @@ class FunctionalizeCond { Graph* graph_; friend class FunctionalizeCondTest; + + std::vector switch_ids_; }; } // namespace functionalize_cond diff --git a/tensorflow/compiler/tf2xla/functionalize_control_flow.cc b/tensorflow/compiler/tf2xla/functionalize_control_flow.cc index f818d80022da0bad851c896f2714c15b20b22195..3dfd3f854c8646ebbf06d3378201d22e8741b7eb 100644 --- a/tensorflow/compiler/tf2xla/functionalize_control_flow.cc +++ b/tensorflow/compiler/tf2xla/functionalize_control_flow.cc @@ -75,6 +75,25 @@ Status FunctionalizeControlFlow(Graph* graph, return FunctionalizeControlFlow(/*lookup_library=*/nullptr, graph, library); } +Status FunctionalizeControlFlowForGraphDef(GraphDef* graph_def, + FunctionLibraryDefinition* library) { + return FunctionalizeControlFlowForGraphDef(/*lookup_library=*/nullptr, + graph_def, library); +} + +Status FunctionalizeControlFlowForGraphDef( + const FunctionLibraryDefinition* lookup_library, GraphDef* graph_def, + FunctionLibraryDefinition* library) { + FunctionDefLibrary function_lib = graph_def->library(); + Graph graph(OpRegistry::Global()); + + TF_RETURN_IF_ERROR(ConvertGraphDefToGraph({}, *graph_def, &graph)); + TF_RETURN_IF_ERROR(FunctionalizeControlFlow(lookup_library, &graph, library)); + graph.ToGraphDef(graph_def); + std::swap(*graph_def->mutable_library(), function_lib); + return Status::OK(); +} + Status FunctionalizeControlFlowForFunction( const string& func_name, const string& new_func_name, const protobuf::Map& attrs, @@ -242,23 +261,20 @@ Status FunctionalizeControlFlowPass::Run( continue; } const string func_attr = it->second; - if (kNodeTypeToFunctionAttrMapping->find(n->type_string()) != - kNodeTypeToFunctionAttrMapping->end()) { - NameAttrList func; - TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), func_attr, &func)); - VLOG(2) << "Graph has node " << n->type_string() - << ". Corresponding function: " << func.name(); - string new_func_name = options.flib_def->UniqueFunctionName( - absl::StrCat(func.name(), "_f15n_")); - bool modified; - TF_RETURN_IF_ERROR(FunctionalizeControlFlowForFunction( - func.name(), new_func_name, func.attr(), options.flib_def, flr, - &canonicalized_name_to_new_name, &modified)); - if (modified) { - n->ClearAttr(func_attr); - func.set_name(new_func_name); - n->AddAttr(func_attr, func); - } + NameAttrList func; + TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), func_attr, &func)); + VLOG(2) << "Graph has node " << n->type_string() + << ". Corresponding function: " << func.name(); + string new_func_name = options.flib_def->UniqueFunctionName( + absl::StrCat(func.name(), "_f15n_")); + bool modified; + TF_RETURN_IF_ERROR(FunctionalizeControlFlowForFunction( + func.name(), new_func_name, func.attr(), options.flib_def, flr, + &canonicalized_name_to_new_name, &modified)); + if (modified) { + n->ClearAttr(func_attr); + func.set_name(new_func_name); + n->AddAttr(func_attr, func); } } diff --git a/tensorflow/compiler/tf2xla/functionalize_control_flow.h b/tensorflow/compiler/tf2xla/functionalize_control_flow.h index ba99205640ccdc83a3a4d50e3ec474907894a835..91d33fa405834d7f1f8f66180583580f4f2e448a 100644 --- a/tensorflow/compiler/tf2xla/functionalize_control_flow.h +++ b/tensorflow/compiler/tf2xla/functionalize_control_flow.h @@ -33,6 +33,12 @@ Status FunctionalizeControlFlow(const FunctionLibraryDefinition* lookup_library, Graph* graph, FunctionLibraryDefinition* library); +Status FunctionalizeControlFlowForGraphDef(GraphDef* graph_def, + FunctionLibraryDefinition* library); +Status FunctionalizeControlFlowForGraphDef( + const FunctionLibraryDefinition* lookup_library, GraphDef* graph_def, + FunctionLibraryDefinition* library); + // This pass looks at the graph and all associated FunctionDefs, and turns // traditional control flow structure (Switch/Merge/etc.) into functional // control flow structure (If/While). diff --git a/tensorflow/compiler/tf2xla/functionalize_control_flow_test.cc b/tensorflow/compiler/tf2xla/functionalize_control_flow_test.cc index c3841f996f801e855da75b23f01d41674ec51c4d..9784985af83a18619d837528f99a60b98a501ec5 100644 --- a/tensorflow/compiler/tf2xla/functionalize_control_flow_test.cc +++ b/tensorflow/compiler/tf2xla/functionalize_control_flow_test.cc @@ -95,77 +95,87 @@ TEST(FunctionalizeControlFlow, Conditional) { } FunctionLibraryDefinition library(OpRegistry::Global(), {}); + GraphDef optimized_graph_def; + graph.ToGraphDef(&optimized_graph_def); + TF_ASSERT_OK( + FunctionalizeControlFlowForGraphDef(&optimized_graph_def, &library)); TF_ASSERT_OK(FunctionalizeControlFlow(&graph, &library)); + GraphDef converted_graph_def; + graph.ToGraphDef(&converted_graph_def); + + for (const GraphDef& graph_def : {optimized_graph_def, converted_graph_def}) { + string op_name; + NameAttrList then_fn; + NameAttrList else_fn; + TF_EXPECT_OK(FindIfThenAndElse(graph_def, &op_name, &then_fn, &else_fn)); + InstantiationResultForTest else_result; + TF_EXPECT_OK( + InstantiateFunctionForTest(else_fn.name(), library, &else_result)); + + // Outer graph + { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto y = ops::Placeholder(scope.WithOpName("y"), DT_INT32); + auto x = ops::Placeholder(scope.WithOpName("x"), DT_INT32); + auto less = ops::Less(scope.WithOpName("cond/Less"), y, x); + auto if_op = ops::If(scope.WithOpName(op_name), less, + std::initializer_list{less, y, x}, {DT_INT32}, + then_fn, else_fn); + auto id = ops::Identity(scope.WithOpName("cond/Merge"), if_op.output[0]); + GraphDef expected; + TF_EXPECT_OK(scope.ToGraphDef(&expected)); + TF_EXPECT_GRAPH_EQ(expected, graph_def); + } - GraphDef graph_def; - graph.ToGraphDef(&graph_def); - string op_name; - NameAttrList then_fn; - NameAttrList else_fn; - TF_EXPECT_OK(FindIfThenAndElse(graph_def, &op_name, &then_fn, &else_fn)); - InstantiationResultForTest else_result; - TF_EXPECT_OK( - InstantiateFunctionForTest(else_fn.name(), library, &else_result)); - - // Outer graph - { - Scope scope = Scope::NewRootScope().ExitOnError(); - auto y = ops::Placeholder(scope.WithOpName("y"), DT_INT32); - auto x = ops::Placeholder(scope.WithOpName("x"), DT_INT32); - auto less = ops::Less(scope.WithOpName("cond/Less"), y, x); - auto if_op = ops::If(scope.WithOpName(op_name), less, - std::initializer_list{less, y, x}, {DT_INT32}, - then_fn, else_fn); - auto id = ops::Identity(scope.WithOpName("cond/Merge"), if_op.output[0]); - GraphDef expected; - TF_EXPECT_OK(scope.ToGraphDef(&expected)); - TF_EXPECT_GRAPH_EQ(expected, graph_def); - } - - // then body. - { - Scope scope = Scope::NewRootScope().ExitOnError(); - auto arg_0 = ops::_Arg(scope.WithOpName("_arg0"), DT_BOOL, 0); - auto arg_1 = ops::_Arg(scope.WithOpName("_arg1"), DT_INT32, 1); - auto arg_2 = ops::_Arg(scope.WithOpName("_arg2"), DT_INT32, 2); - auto identity = ops::Identity(scope.WithOpName("cond/Identity"), arg_0); - auto cond = ops::Const( - scope.WithOpName("cond").WithControlDependencies(identity), 17); - auto mul = ops::Mul(scope.WithOpName("cond/Mul"), arg_1, cond); - auto retval0 = ops::_Retval(scope.WithOpName("_retval0_RetVal"), mul, 0); - - GraphDef expected; - TF_EXPECT_OK(scope.ToGraphDef(&expected)); - - InstantiationResultForTest result; - TF_EXPECT_OK(InstantiateFunctionForTest(then_fn.name(), library, &result)); - - EXPECT_EQ(DataTypeVector{DT_INT32}, result.ret_types); - EXPECT_EQ((DataTypeVector{DT_BOOL, DT_INT32, DT_INT32}), result.arg_types); - TF_EXPECT_GRAPH_EQ(expected, result.gdef); - } + // then body. + { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto arg_0 = ops::_Arg(scope.WithOpName("_arg0"), DT_BOOL, 0); + auto arg_1 = ops::_Arg(scope.WithOpName("_arg1"), DT_INT32, 1); + auto arg_2 = ops::_Arg(scope.WithOpName("_arg2"), DT_INT32, 2); + auto identity = ops::Identity(scope.WithOpName("cond/Identity"), arg_0); + auto cond = ops::Const( + scope.WithOpName("cond").WithControlDependencies(identity), 17); + auto mul = ops::Mul(scope.WithOpName("cond/Mul"), arg_1, cond); + auto retval0 = ops::_Retval(scope.WithOpName("_retval0_RetVal"), mul, 0); + + GraphDef expected; + TF_EXPECT_OK(scope.ToGraphDef(&expected)); + + InstantiationResultForTest result; + TF_EXPECT_OK( + InstantiateFunctionForTest(then_fn.name(), library, &result)); + + EXPECT_EQ(DataTypeVector{DT_INT32}, result.ret_types); + EXPECT_EQ((DataTypeVector{DT_BOOL, DT_INT32, DT_INT32}), + result.arg_types); + TF_EXPECT_GRAPH_EQ(expected, result.gdef); + } - // else body. - { - Scope scope = Scope::NewRootScope().ExitOnError(); - auto arg_0 = ops::_Arg(scope.WithOpName("_arg0"), DT_BOOL, 0); - auto arg_1 = ops::_Arg(scope.WithOpName("_arg1"), DT_INT32, 1); - auto arg_2 = ops::_Arg(scope.WithOpName("_arg2"), DT_INT32, 2); - auto identity = ops::Identity(scope.WithOpName("cond/Identity_1"), arg_0); - auto cond_1 = ops::Const( - scope.WithOpName("cond_1").WithControlDependencies(identity), 23); - auto add = ops::Add(scope.WithOpName("cond/false/add"), arg_2, cond_1); - auto retval0 = ops::_Retval(scope.WithOpName("_retval0_RetVal"), add, 0); - - GraphDef expected; - TF_EXPECT_OK(scope.ToGraphDef(&expected)); - - InstantiationResultForTest result; - TF_EXPECT_OK(InstantiateFunctionForTest(else_fn.name(), library, &result)); - - EXPECT_EQ(DataTypeVector{DT_INT32}, result.ret_types); - EXPECT_EQ((DataTypeVector{DT_BOOL, DT_INT32, DT_INT32}), result.arg_types); - TF_EXPECT_GRAPH_EQ(expected, result.gdef); + // else body. + { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto arg_0 = ops::_Arg(scope.WithOpName("_arg0"), DT_BOOL, 0); + auto arg_1 = ops::_Arg(scope.WithOpName("_arg1"), DT_INT32, 1); + auto arg_2 = ops::_Arg(scope.WithOpName("_arg2"), DT_INT32, 2); + auto identity = ops::Identity(scope.WithOpName("cond/Identity_1"), arg_0); + auto cond_1 = ops::Const( + scope.WithOpName("cond_1").WithControlDependencies(identity), 23); + auto add = ops::Add(scope.WithOpName("cond/false/add"), arg_2, cond_1); + auto retval0 = ops::_Retval(scope.WithOpName("_retval0_RetVal"), add, 0); + + GraphDef expected; + TF_EXPECT_OK(scope.ToGraphDef(&expected)); + + InstantiationResultForTest result; + TF_EXPECT_OK( + InstantiateFunctionForTest(else_fn.name(), library, &result)); + + EXPECT_EQ(DataTypeVector{DT_INT32}, result.ret_types); + EXPECT_EQ((DataTypeVector{DT_BOOL, DT_INT32, DT_INT32}), + result.arg_types); + TF_EXPECT_GRAPH_EQ(expected, result.gdef); + } } } @@ -239,75 +249,77 @@ TEST(FunctionalizeControlFlow, OneLoopVar) { } FunctionLibraryDefinition library(OpRegistry::Global(), {}); + GraphDef optimized_graph_def; + graph.ToGraphDef(&optimized_graph_def); + TF_ASSERT_OK( + FunctionalizeControlFlowForGraphDef(&optimized_graph_def, &library)); TF_ASSERT_OK(FunctionalizeControlFlow(&graph, &library)); + GraphDef converted_graph_def; + graph.ToGraphDef(&converted_graph_def); + + for (const GraphDef& graph_def : {optimized_graph_def, converted_graph_def}) { + NameAttrList cond_fn, body_fn; + TF_EXPECT_OK(FindWhileCondAndBody(graph_def, &cond_fn, &body_fn)); + + // Outer graph + { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto source = ops::Placeholder(scope.WithOpName("source"), DT_INT32); + auto while_op = + ops::While(scope.WithOpName("while/LoopCond"), + std::initializer_list{source}, cond_fn, body_fn); + auto sink = ops::Identity(scope.WithOpName("sink"), while_op[0]); + GraphDef expected; + TF_EXPECT_OK(scope.ToGraphDef(&expected)); + TF_EXPECT_GRAPH_EQ(expected, graph_def); + } - GraphDef graph_def; - graph.ToGraphDef(&graph_def); - - NameAttrList cond_fn, body_fn; - TF_EXPECT_OK(FindWhileCondAndBody(graph_def, &cond_fn, &body_fn)); - - // Outer graph - { - Scope scope = Scope::NewRootScope().ExitOnError(); - auto source = ops::Placeholder(scope.WithOpName("source"), DT_INT32); - auto while_op = - ops::While(scope.WithOpName("while/LoopCond"), - std::initializer_list{source}, cond_fn, body_fn); - auto sink = ops::Identity(scope.WithOpName("sink"), while_op[0]); - GraphDef expected; - TF_EXPECT_OK(scope.ToGraphDef(&expected)); - TF_EXPECT_GRAPH_EQ(expected, graph_def); - } - - // Condition graph - { - Scope scope = Scope::NewRootScope().ExitOnError(); - auto arg = ops::_Arg(scope.WithOpName("_arg0"), DT_INT32, 0); - auto ten = ops::Const( - scope.WithOpName("while/Less/y").WithControlDependencies(arg), 10); - auto less = ops::Less(scope.WithOpName("while/Less"), arg, ten); - auto retval = ops::_Retval(scope.WithOpName("_retval0_RetVal"), less, 0); - - GraphDef expected; - TF_EXPECT_OK(scope.ToGraphDef(&expected)); - - InstantiationResultForTest result; - TF_EXPECT_OK(InstantiateFunctionForTest(cond_fn.name(), library, &result)); - - EXPECT_EQ(DataTypeVector{DT_INT32}, result.arg_types); - EXPECT_EQ(DataTypeVector{DT_BOOL}, result.ret_types); - TF_EXPECT_GRAPH_EQ(expected, result.gdef); - } - - // Body graph. - { - Scope scope = Scope::NewRootScope().ExitOnError(); - auto arg = ops::_Arg(scope.WithOpName("_arg0"), DT_INT32, 0); - auto identity = ops::Identity(scope.WithOpName("while/Identity"), arg); - auto one = ops::Const( - scope.WithOpName("while/add/y").WithControlDependencies(identity), 1); - auto add = ops::Add(scope.WithOpName("while/add"), identity, one); - auto retval = ops::_Retval(scope.WithOpName("_retval0_RetVal"), add, 0); - - GraphDef expected; - TF_EXPECT_OK(scope.ToGraphDef(&expected)); - - InstantiationResultForTest result; - TF_EXPECT_OK(InstantiateFunctionForTest(body_fn.name(), library, &result)); + // Condition graph + { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto arg = ops::_Arg(scope.WithOpName("_arg0"), DT_INT32, 0); + auto ten = ops::Const( + scope.WithOpName("while/Less/y").WithControlDependencies(arg), 10); + auto less = ops::Less(scope.WithOpName("while/Less"), arg, ten); + auto retval = ops::_Retval(scope.WithOpName("_retval0_RetVal"), less, 0); + + GraphDef expected; + TF_EXPECT_OK(scope.ToGraphDef(&expected)); + + InstantiationResultForTest result; + TF_EXPECT_OK( + InstantiateFunctionForTest(cond_fn.name(), library, &result)); + + EXPECT_EQ(DataTypeVector{DT_INT32}, result.arg_types); + EXPECT_EQ(DataTypeVector{DT_BOOL}, result.ret_types); + TF_EXPECT_GRAPH_EQ(expected, result.gdef); + } - EXPECT_EQ(DataTypeVector{DT_INT32}, result.arg_types); - EXPECT_EQ(DataTypeVector{DT_INT32}, result.ret_types); - TF_EXPECT_GRAPH_EQ(expected, result.gdef); + // Body graph. + { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto arg = ops::_Arg(scope.WithOpName("_arg0"), DT_INT32, 0); + auto identity = ops::Identity(scope.WithOpName("while/Identity"), arg); + auto one = ops::Const( + scope.WithOpName("while/add/y").WithControlDependencies(identity), 1); + auto add = ops::Add(scope.WithOpName("while/add"), identity, one); + auto retval = ops::_Retval(scope.WithOpName("_retval0_RetVal"), add, 0); + + GraphDef expected; + TF_EXPECT_OK(scope.ToGraphDef(&expected)); + + InstantiationResultForTest result; + TF_EXPECT_OK( + InstantiateFunctionForTest(body_fn.name(), library, &result)); + + EXPECT_EQ(DataTypeVector{DT_INT32}, result.arg_types); + EXPECT_EQ(DataTypeVector{DT_INT32}, result.ret_types); + TF_EXPECT_GRAPH_EQ(expected, result.gdef); + } } } -// @function.Defun(noinline=True) -// def increment_fn(x): -// return [x + 1] -// Define the above function, and add it to the given graph. It's used as the -// while loop body in NoinlineLoopBody test. -Status AddNoinlineFunctionToGraph(const string& node_name, Graph* graph) { +FunctionDef GetNoinlineFunctionDef() { FunctionDef fdef = FunctionDefHelper::Create( "increment_fn", {"x:int32"}, {"add:int32"}, {}, { @@ -316,8 +328,17 @@ Status AddNoinlineFunctionToGraph(const string& node_name, Graph* graph) { }, {{"add", "add_0:z:0"}}); (*fdef.mutable_attr())["_noinline"].set_b(true); + return fdef; +} + +// @function.Defun(noinline=True) +// def increment_fn(x): +// return [x + 1] +// Define the above function, and add it to the given graph. It's used as the +// while loop body in NoinlineLoopBody test. +Status AddNoinlineFunctionToGraph(const string& node_name, Graph* graph) { FunctionDefLibrary fdef_lib; - *(fdef_lib.add_function()) = fdef; + *(fdef_lib.add_function()) = GetNoinlineFunctionDef(); TF_RETURN_IF_ERROR(graph->AddFunctionLibrary(fdef_lib)); NodeDef increment_fn; increment_fn.set_name(node_name); @@ -376,55 +397,88 @@ TEST(FunctionalizeControlFlow, NoinlineLoopBody) { FunctionLibraryDefinition lookup_lib(graph.flib_def()); FunctionLibraryDefinition library(OpRegistry::Global(), {}); // Function increment_fn will be copied from lookup_lib to library. - TF_ASSERT_OK(FunctionalizeControlFlow(&lookup_lib, &graph, &library)); + GraphDef optimized_graph_def; + graph.ToGraphDef(&optimized_graph_def); - GraphDef graph_def; - graph.ToGraphDef(&graph_def); + *(optimized_graph_def.mutable_library()->add_function()) = + GetNoinlineFunctionDef(); - NameAttrList cond_fn, body_fn; - TF_ASSERT_OK(FindWhileCondAndBody(graph_def, &cond_fn, &body_fn)); + TF_ASSERT_OK(FunctionalizeControlFlowForGraphDef( + &lookup_lib, &optimized_graph_def, &library)); + TF_ASSERT_OK(FunctionalizeControlFlow(&lookup_lib, &graph, &library)); + GraphDef converted_graph_def; + graph.ToGraphDef(&converted_graph_def); + + for (const GraphDef& graph_def : {optimized_graph_def, converted_graph_def}) { + NameAttrList cond_fn, body_fn; + TF_ASSERT_OK(FindWhileCondAndBody(graph_def, &cond_fn, &body_fn)); + + // Outer graph + { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto source = ops::Placeholder(scope.WithOpName("source"), DT_INT32); + auto while_op = + ops::While(scope.WithOpName("while/LoopCond"), + std::initializer_list{source}, cond_fn, body_fn); + GraphDef expected; + TF_ASSERT_OK(scope.ToGraphDef(&expected)); + TF_EXPECT_GRAPH_EQ(expected, graph_def); + } - // Outer graph - { - Scope scope = Scope::NewRootScope().ExitOnError(); - auto source = ops::Placeholder(scope.WithOpName("source"), DT_INT32); - auto while_op = - ops::While(scope.WithOpName("while/LoopCond"), - std::initializer_list{source}, cond_fn, body_fn); - GraphDef expected; - TF_ASSERT_OK(scope.ToGraphDef(&expected)); - TF_EXPECT_GRAPH_EQ(expected, graph_def); + // Body graph. + { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto arg = ops::_Arg(scope.WithOpName("_arg0"), DT_INT32, 0); + TF_ASSERT_OK( + AddNoinlineFunctionToGraph(noinline_node_name, scope.graph())); + auto identity = ops::Identity(scope.WithOpName("while/Identity"), arg); + NodeDef retval; + retval.set_name("_retval0_RetVal"); + retval.set_op(FunctionLibraryDefinition::kRetOp); + *retval.add_input() = noinline_node_name; + (*retval.mutable_attr())["T"].set_type(DT_INT32); + (*retval.mutable_attr())["index"].set_i(0); + Status status; + scope.graph()->AddNode(retval, &status); + TF_ASSERT_OK(status); + + GraphDef expected; + TF_ASSERT_OK(scope.ToGraphDef(&expected)); + + InstantiationResultForTest result; + // Verify that increment_fn has been copied to library. + TF_EXPECT_OK( + InstantiateFunctionForTest(body_fn.name(), library, &result)); + + EXPECT_EQ(DataTypeVector{DT_INT32}, result.arg_types); + EXPECT_EQ(DataTypeVector{DT_INT32}, result.ret_types); + // Ignore the function library when comparing the graphs. + expected.clear_library(); + TF_EXPECT_GRAPH_EQ(expected, result.gdef); + } } +} - // Body graph. +TEST(FunctionalizeControlFlow, MissingFunctionDefInLibrary) { + const string& noinline_node_name = "while/increment_fn"; + Graph graph(OpRegistry::Global()); { Scope scope = Scope::NewRootScope().ExitOnError(); - auto arg = ops::_Arg(scope.WithOpName("_arg0"), DT_INT32, 0); + auto source = ops::Placeholder(scope.WithOpName("source"), DT_INT32); + auto identity = ops::Identity(scope.WithOpName("while/Identity"), source); TF_ASSERT_OK(AddNoinlineFunctionToGraph(noinline_node_name, scope.graph())); - auto identity = ops::Identity(scope.WithOpName("while/Identity"), arg); - NodeDef retval; - retval.set_name("_retval0_RetVal"); - retval.set_op(FunctionLibraryDefinition::kRetOp); - *retval.add_input() = noinline_node_name; - (*retval.mutable_attr())["T"].set_type(DT_INT32); - (*retval.mutable_attr())["index"].set_i(0); - Status status; - scope.graph()->AddNode(retval, &status); - TF_ASSERT_OK(status); - - GraphDef expected; - TF_ASSERT_OK(scope.ToGraphDef(&expected)); + TF_ASSERT_OK(scope.ToGraph(&graph)); + } - InstantiationResultForTest result; - // Verify that increment_fn has been copied to library. - TF_EXPECT_OK(InstantiateFunctionForTest(body_fn.name(), library, &result)); + FunctionLibraryDefinition lookup_lib(graph.flib_def()); + FunctionLibraryDefinition library(OpRegistry::Global(), {}); + GraphDef graph_def; + graph.ToGraphDef(&graph_def); + graph_def.clear_library(); - EXPECT_EQ(DataTypeVector{DT_INT32}, result.arg_types); - EXPECT_EQ(DataTypeVector{DT_INT32}, result.ret_types); - // Ignore the function library when comparing the graphs. - expected.clear_library(); - TF_EXPECT_GRAPH_EQ(expected, result.gdef); - } + Status status = + FunctionalizeControlFlowForGraphDef(&lookup_lib, &graph_def, &library); + EXPECT_EQ(tensorflow::error::NOT_FOUND, status.code()); } // Tests functionalizing OneLoopVar where the loop value is not used post the @@ -467,65 +521,72 @@ TEST(FunctionalizeControlFlow, OneLoopVarWithoutExit) { } FunctionLibraryDefinition library(OpRegistry::Global(), {}); + GraphDef optimized_graph_def; + graph.ToGraphDef(&optimized_graph_def); + TF_ASSERT_OK( + FunctionalizeControlFlowForGraphDef(&optimized_graph_def, &library)); TF_ASSERT_OK(FunctionalizeControlFlow(&graph, &library)); + GraphDef converted_graph_def; + graph.ToGraphDef(&converted_graph_def); + + for (const GraphDef& graph_def : {optimized_graph_def, converted_graph_def}) { + NameAttrList cond_fn, body_fn; + TF_EXPECT_OK(FindWhileCondAndBody(graph_def, &cond_fn, &body_fn)); + + // Outer graph + { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto source = ops::Placeholder(scope.WithOpName("source"), DT_INT32); + auto while_op = + ops::While(scope.WithOpName("while/LoopCond"), + std::initializer_list{source}, cond_fn, body_fn); + GraphDef expected; + TF_EXPECT_OK(scope.ToGraphDef(&expected)); + TF_EXPECT_GRAPH_EQ(expected, graph_def); + } - GraphDef graph_def; - graph.ToGraphDef(&graph_def); - - NameAttrList cond_fn, body_fn; - TF_EXPECT_OK(FindWhileCondAndBody(graph_def, &cond_fn, &body_fn)); - - // Outer graph - { - Scope scope = Scope::NewRootScope().ExitOnError(); - auto source = ops::Placeholder(scope.WithOpName("source"), DT_INT32); - auto while_op = - ops::While(scope.WithOpName("while/LoopCond"), - std::initializer_list{source}, cond_fn, body_fn); - GraphDef expected; - TF_EXPECT_OK(scope.ToGraphDef(&expected)); - TF_EXPECT_GRAPH_EQ(expected, graph_def); - } - - // Condition graph - { - Scope scope = Scope::NewRootScope().ExitOnError(); - auto arg = ops::_Arg(scope.WithOpName("_arg0"), DT_INT32, 0); - auto ten = ops::Const( - scope.WithOpName("while/Less/y").WithControlDependencies(arg), 10); - auto less = ops::Less(scope.WithOpName("while/Less"), arg, ten); - auto retval = ops::_Retval(scope.WithOpName("_retval0_RetVal"), less, 0); - - GraphDef expected; - TF_EXPECT_OK(scope.ToGraphDef(&expected)); - - InstantiationResultForTest result; - TF_EXPECT_OK(InstantiateFunctionForTest(cond_fn.name(), library, &result)); - - EXPECT_EQ(DataTypeVector{DT_INT32}, result.arg_types); - EXPECT_EQ(DataTypeVector{DT_BOOL}, result.ret_types); - TF_EXPECT_GRAPH_EQ(expected, result.gdef); - } - - // Body graph. - { - Scope scope = Scope::NewRootScope().ExitOnError(); - auto arg = ops::_Arg(scope.WithOpName("_arg0"), DT_INT32, 0); - auto identity = ops::Identity(scope.WithOpName("while/Identity"), arg); - auto one = ops::Const( - scope.WithOpName("while/add/y").WithControlDependencies(identity), 1); - auto add = ops::Add(scope.WithOpName("while/add"), identity, one); - auto retval = ops::_Retval(scope.WithOpName("_retval0_RetVal"), add, 0); - - GraphDef expected; - TF_EXPECT_OK(scope.ToGraphDef(&expected)); - - InstantiationResultForTest result; - TF_EXPECT_OK(InstantiateFunctionForTest(body_fn.name(), library, &result)); + // Condition graph + { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto arg = ops::_Arg(scope.WithOpName("_arg0"), DT_INT32, 0); + auto ten = ops::Const( + scope.WithOpName("while/Less/y").WithControlDependencies(arg), 10); + auto less = ops::Less(scope.WithOpName("while/Less"), arg, ten); + auto retval = ops::_Retval(scope.WithOpName("_retval0_RetVal"), less, 0); + + GraphDef expected; + TF_EXPECT_OK(scope.ToGraphDef(&expected)); + + InstantiationResultForTest result; + TF_EXPECT_OK( + InstantiateFunctionForTest(cond_fn.name(), library, &result)); + + EXPECT_EQ(DataTypeVector{DT_INT32}, result.arg_types); + EXPECT_EQ(DataTypeVector{DT_BOOL}, result.ret_types); + TF_EXPECT_GRAPH_EQ(expected, result.gdef); + } - EXPECT_EQ(DataTypeVector{DT_INT32}, result.arg_types); - EXPECT_EQ(DataTypeVector{DT_INT32}, result.ret_types); - TF_EXPECT_GRAPH_EQ(expected, result.gdef); + // Body graph. + { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto arg = ops::_Arg(scope.WithOpName("_arg0"), DT_INT32, 0); + auto identity = ops::Identity(scope.WithOpName("while/Identity"), arg); + auto one = ops::Const( + scope.WithOpName("while/add/y").WithControlDependencies(identity), 1); + auto add = ops::Add(scope.WithOpName("while/add"), identity, one); + auto retval = ops::_Retval(scope.WithOpName("_retval0_RetVal"), add, 0); + + GraphDef expected; + TF_EXPECT_OK(scope.ToGraphDef(&expected)); + + InstantiationResultForTest result; + TF_EXPECT_OK( + InstantiateFunctionForTest(body_fn.name(), library, &result)); + + EXPECT_EQ(DataTypeVector{DT_INT32}, result.arg_types); + EXPECT_EQ(DataTypeVector{DT_INT32}, result.ret_types); + TF_EXPECT_GRAPH_EQ(expected, result.gdef); + } } } @@ -608,86 +669,95 @@ TEST(FunctionalizeControlFlow, TwoLoopVars) { } FunctionLibraryDefinition library(OpRegistry::Global(), {}); + GraphDef optimized_graph_def; + graph.ToGraphDef(&optimized_graph_def); + TF_ASSERT_OK( + FunctionalizeControlFlowForGraphDef(&optimized_graph_def, &library)); TF_ASSERT_OK(FunctionalizeControlFlow(&graph, &library)); + GraphDef converted_graph_def; + graph.ToGraphDef(&converted_graph_def); + + for (const GraphDef& graph_def : {optimized_graph_def, converted_graph_def}) { + NameAttrList cond_fn, body_fn; + TF_EXPECT_OK(FindWhileCondAndBody(graph_def, &cond_fn, &body_fn)); + + // Outer graph. + { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto x = ops::Placeholder(scope.WithOpName("Placeholder/x"), DT_INT32); + auto y = ops::Placeholder(scope.WithOpName("Placeholder/y"), DT_INT32); + auto while_op = + ops::While(scope.WithOpName("while/LoopCond"), + std::initializer_list{x, y}, cond_fn, body_fn); + auto sink_x = ops::Identity(scope.WithOpName("sink_x"), while_op[0]); + auto sink_y = ops::Identity(scope.WithOpName("sink_y"), while_op[1]); + GraphDef expected; + TF_EXPECT_OK(scope.ToGraphDef(&expected)); + TF_EXPECT_GRAPH_EQ(expected, graph_def); + } - GraphDef graph_def; - graph.ToGraphDef(&graph_def); - - NameAttrList cond_fn, body_fn; - TF_EXPECT_OK(FindWhileCondAndBody(graph_def, &cond_fn, &body_fn)); - - // Outer graph. - { - Scope scope = Scope::NewRootScope().ExitOnError(); - auto x = ops::Placeholder(scope.WithOpName("Placeholder/x"), DT_INT32); - auto y = ops::Placeholder(scope.WithOpName("Placeholder/y"), DT_INT32); - auto while_op = - ops::While(scope.WithOpName("while/LoopCond"), - std::initializer_list{x, y}, cond_fn, body_fn); - auto sink_x = ops::Identity(scope.WithOpName("sink_x"), while_op[0]); - auto sink_y = ops::Identity(scope.WithOpName("sink_y"), while_op[1]); - GraphDef expected; - TF_EXPECT_OK(scope.ToGraphDef(&expected)); - TF_EXPECT_GRAPH_EQ(expected, graph_def); - } - - // Condition graph. - { - Scope scope = Scope::NewRootScope().ExitOnError(); - auto arg0 = ops::_Arg(scope.WithOpName("_arg0"), DT_INT32, 0); - auto arg1 = ops::_Arg(scope.WithOpName("_arg1"), DT_INT32, 1); - auto three = ops::Const(scope.WithOpName("while/cond/three") + // Condition graph. + { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto arg0 = ops::_Arg(scope.WithOpName("_arg0"), DT_INT32, 0); + auto arg1 = ops::_Arg(scope.WithOpName("_arg1"), DT_INT32, 1); + auto three = ops::Const(scope.WithOpName("while/cond/three") + .WithControlDependencies(arg0.output), + 3); + auto cond_add = + ops::Add(scope.WithOpName("while/cond/Add"), arg0.output, three); + auto ten = ops::Const(scope.WithOpName("while/cond/ten") .WithControlDependencies(arg0.output), - 3); - auto cond_add = - ops::Add(scope.WithOpName("while/cond/Add"), arg0.output, three); - auto ten = ops::Const( - scope.WithOpName("while/cond/ten").WithControlDependencies(arg0.output), - 10); - auto less = ops::Less(scope.WithOpName("while/cond/Less"), cond_add, ten); - auto retval = ops::_Retval(scope.WithOpName("_retval0_RetVal"), less, 0); - - GraphDef expected; - TF_EXPECT_OK(scope.ToGraphDef(&expected)); - - InstantiationResultForTest result; - TF_EXPECT_OK(InstantiateFunctionForTest(cond_fn.name(), library, &result)); - - EXPECT_EQ((DataTypeVector{DT_INT32, DT_INT32}), result.arg_types); - EXPECT_EQ(DataTypeVector{DT_BOOL}, result.ret_types); - TF_EXPECT_GRAPH_EQ(expected, result.gdef); - } - - // Body graph. - { - Scope scope = Scope::NewRootScope().ExitOnError(); - auto arg0 = ops::_Arg(scope.WithOpName("_arg0"), DT_INT32, 0); - auto arg1 = ops::_Arg(scope.WithOpName("_arg1"), DT_INT32, 1); - - auto identity_x = ops::Identity(scope.WithOpName("while/Identity/x"), arg0); - auto identity_y = ops::Identity(scope.WithOpName("while/Identity/y"), arg1); - - auto one = ops::Const( - scope.WithOpName("while/add/one").WithControlDependencies(identity_x), - 1); - auto two = ops::Const( - scope.WithOpName("while/mul/two").WithControlDependencies(identity_x), - 2); + 10); + auto less = ops::Less(scope.WithOpName("while/cond/Less"), cond_add, ten); + auto retval = ops::_Retval(scope.WithOpName("_retval0_RetVal"), less, 0); - auto add = ops::Add(scope.WithOpName("while/add"), identity_x, one); - auto mul = ops::Add(scope.WithOpName("while/mul"), identity_y, two); - auto retval0 = ops::_Retval(scope.WithOpName("_retval0_RetVal"), add, 0); - auto retval1 = ops::_Retval(scope.WithOpName("_retval1_RetVal"), mul, 1); + GraphDef expected; + TF_EXPECT_OK(scope.ToGraphDef(&expected)); - GraphDef expected; - TF_EXPECT_OK(scope.ToGraphDef(&expected)); + InstantiationResultForTest result; + TF_EXPECT_OK( + InstantiateFunctionForTest(cond_fn.name(), library, &result)); - InstantiationResultForTest result; - TF_EXPECT_OK(InstantiateFunctionForTest(body_fn.name(), library, &result)); + EXPECT_EQ((DataTypeVector{DT_INT32, DT_INT32}), result.arg_types); + EXPECT_EQ(DataTypeVector{DT_BOOL}, result.ret_types); + TF_EXPECT_GRAPH_EQ(expected, result.gdef); + } - EXPECT_EQ((DataTypeVector{DT_INT32, DT_INT32}), result.arg_types); - EXPECT_EQ((DataTypeVector{DT_INT32, DT_INT32}), result.ret_types); - TF_EXPECT_GRAPH_EQ(expected, result.gdef); + // Body graph. + { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto arg0 = ops::_Arg(scope.WithOpName("_arg0"), DT_INT32, 0); + auto arg1 = ops::_Arg(scope.WithOpName("_arg1"), DT_INT32, 1); + + auto identity_x = + ops::Identity(scope.WithOpName("while/Identity/x"), arg0); + auto identity_y = + ops::Identity(scope.WithOpName("while/Identity/y"), arg1); + + auto one = ops::Const( + scope.WithOpName("while/add/one").WithControlDependencies(identity_x), + 1); + auto two = ops::Const( + scope.WithOpName("while/mul/two").WithControlDependencies(identity_x), + 2); + + auto add = ops::Add(scope.WithOpName("while/add"), identity_x, one); + auto mul = ops::Add(scope.WithOpName("while/mul"), identity_y, two); + auto retval0 = ops::_Retval(scope.WithOpName("_retval0_RetVal"), add, 0); + auto retval1 = ops::_Retval(scope.WithOpName("_retval1_RetVal"), mul, 1); + + GraphDef expected; + TF_EXPECT_OK(scope.ToGraphDef(&expected)); + + InstantiationResultForTest result; + TF_EXPECT_OK( + InstantiateFunctionForTest(body_fn.name(), library, &result)); + + EXPECT_EQ((DataTypeVector{DT_INT32, DT_INT32}), result.arg_types); + EXPECT_EQ((DataTypeVector{DT_INT32, DT_INT32}), result.ret_types); + TF_EXPECT_GRAPH_EQ(expected, result.gdef); + } } } @@ -841,177 +911,192 @@ TEST(FunctionalizeControlFlow, Complex) { } FunctionLibraryDefinition library(OpRegistry::Global(), {}); + GraphDef optimized_graph_def; + graph.ToGraphDef(&optimized_graph_def); + TF_ASSERT_OK( + FunctionalizeControlFlowForGraphDef(&optimized_graph_def, &library)); TF_ASSERT_OK(FunctionalizeControlFlow(&graph, &library)); + GraphDef converted_graph_def; + graph.ToGraphDef(&converted_graph_def); - GraphDef graph_def; - graph.ToGraphDef(&graph_def); - - NameAttrList outer_cond_fn, outer_body_fn; - TF_EXPECT_OK(FindWhileCondAndBody(graph_def, &outer_cond_fn, &outer_body_fn)); - - // Outer graph. - { - Scope scope = Scope::NewRootScope().ExitOnError(); - auto x = ops::Placeholder(scope.WithOpName("x"), DT_INT32); - auto three = ops::Const(scope.WithOpName("three"), 3); - auto y = ops::Add(scope.WithOpName("y"), x, three); - - auto var = ops::VarHandleOp(scope.WithOpName("Variable"), DT_INT32, - TensorShape({})); - - auto zero = ops::Const(scope.WithOpName("outer/Const"), 0); - - auto while_op = ops::While(scope.WithOpName("outer/LoopCond"), - std::initializer_list{zero, y, x, var}, - outer_cond_fn, outer_body_fn); - auto sink = ops::Identity(scope.WithOpName("sink"), while_op[0]); - GraphDef expected; - TF_EXPECT_OK(scope.ToGraphDef(&expected)); - TF_EXPECT_GRAPH_EQ(expected, graph_def); - } - - // Outer condition graph. - { - Scope scope = Scope::NewRootScope().ExitOnError(); - auto arg0 = ops::_Arg(scope.WithOpName("_arg0"), DT_INT32, 0); - auto arg1 = ops::_Arg(scope.WithOpName("_arg1"), DT_INT32, 1); - auto arg2 = ops::_Arg(scope.WithOpName("_arg2"), DT_INT32, 2); - auto arg3 = ops::_Arg(scope.WithOpName("_arg3"), DT_RESOURCE, 3); - - auto ten = ops::Const( - scope.WithOpName("outer/Less/y").WithControlDependencies(arg0.output), - 10); - auto less = ops::Less(scope.WithOpName("outer/Less_i"), arg0, ten); - auto retval = ops::_Retval(scope.WithOpName("_retval0_RetVal"), less, 0); - - GraphDef expected; - TF_EXPECT_OK(scope.ToGraphDef(&expected)); - - InstantiationResultForTest result; - TF_EXPECT_OK( - InstantiateFunctionForTest(outer_cond_fn.name(), library, &result)); - - EXPECT_EQ((DataTypeVector{DT_INT32, DT_INT32, DT_INT32, DT_RESOURCE}), - result.arg_types); - EXPECT_EQ(DataTypeVector{DT_BOOL}, result.ret_types); - TF_EXPECT_GRAPH_EQ(expected, result.gdef); - } - - // Outer body graph. - NameAttrList inner_cond_fn, inner_body_fn; - { - InstantiationResultForTest result; - TF_EXPECT_OK( - InstantiateFunctionForTest(outer_body_fn.name(), library, &result)); - - // Find the inner condition and body names. - TF_EXPECT_OK( - FindWhileCondAndBody(result.gdef, &inner_cond_fn, &inner_body_fn)); - - Scope scope = Scope::NewRootScope().ExitOnError(); - auto arg0 = ops::_Arg(scope.WithOpName("_arg0"), DT_INT32, 0); - auto arg1 = ops::_Arg(scope.WithOpName("_arg1"), DT_INT32, 1); - auto arg2 = ops::_Arg(scope.WithOpName("_arg2"), DT_INT32, 2); - auto arg3 = ops::_Arg(scope.WithOpName("_arg3"), DT_RESOURCE, 3); - - auto identity_i = ops::Identity(scope.WithOpName("outer/Identity"), arg0); - auto one_j = ops::Const( - scope.WithOpName("outer/j").WithControlDependencies(identity_i), 1); - auto while_op = - ops::While(scope.WithOpName("outer/LoopCond_1"), - std::initializer_list{one_j, arg1, arg2, arg3}, - inner_cond_fn, inner_body_fn); - - auto one_outer = ops::Const( - scope.WithOpName("outer/add/y").WithControlDependencies(identity_i), 1); - auto add_i = - ops::Add(scope.WithOpName("outer/add") - .WithControlDependencies(absl::Span{ - while_op[0].op(), while_op[1].op()}), - identity_i, one_outer); - - auto retval0 = ops::_Retval(scope.WithOpName("_retval0_RetVal"), add_i, 0); - auto retval1 = ops::_Retval(scope.WithOpName("_retval1_RetVal"), arg1, 1); - auto retval2 = ops::_Retval(scope.WithOpName("_retval2_RetVal"), arg2, 2); - - GraphDef expected; - TF_EXPECT_OK(scope.ToGraphDef(&expected)); - - EXPECT_EQ((DataTypeVector{DT_INT32, DT_INT32, DT_INT32, DT_RESOURCE}), - result.arg_types); - EXPECT_EQ((DataTypeVector{DT_INT32, DT_INT32, DT_INT32}), result.ret_types); - TF_EXPECT_GRAPH_EQ(expected, result.gdef); - } - - // Inner condition graph. - { - Scope scope = Scope::NewRootScope().ExitOnError(); - auto arg0 = ops::_Arg(scope.WithOpName("_arg0"), DT_INT32, 0); - auto arg1 = ops::_Arg(scope.WithOpName("_arg1"), DT_INT32, 1); - auto arg2 = ops::_Arg(scope.WithOpName("_arg2"), DT_INT32, 2); - auto arg3 = ops::_Arg(scope.WithOpName("_arg3"), DT_RESOURCE, 3); - - auto five = ops::Const( - scope.WithOpName("outer/inner/Five").WithControlDependencies(arg0), 5); - auto less_j = ops::Less(scope.WithOpName("outer/inner/Less_j"), arg0, five); - auto retval = ops::_Retval(scope.WithOpName("_retval0_RetVal"), less_j, 0); - - GraphDef expected; - TF_EXPECT_OK(scope.ToGraphDef(&expected)); - - InstantiationResultForTest result; + for (const GraphDef& graph_def : {optimized_graph_def, converted_graph_def}) { + NameAttrList outer_cond_fn, outer_body_fn; TF_EXPECT_OK( - InstantiateFunctionForTest(inner_cond_fn.name(), library, &result)); - - EXPECT_EQ((DataTypeVector{DT_INT32, DT_INT32, DT_INT32, DT_RESOURCE}), - result.arg_types); - EXPECT_EQ(DataTypeVector{DT_BOOL}, result.ret_types); - TF_EXPECT_GRAPH_EQ(expected, result.gdef); - } - - // Inner body graph. - { - Scope scope = Scope::NewRootScope().ExitOnError(); - auto arg0 = ops::_Arg(scope.WithOpName("_arg0"), DT_INT32, 0); - auto arg1 = ops::_Arg(scope.WithOpName("_arg1"), DT_INT32, 1); - auto arg2 = ops::_Arg(scope.WithOpName("_arg2"), DT_INT32, 2); - auto arg3 = ops::_Arg(scope.WithOpName("_arg3"), DT_RESOURCE, 3); - - auto identity_j = - ops::Identity(scope.WithOpName("outer/inner/Identity_j"), arg0); - auto identity_k = - ops::Identity(scope.WithOpName("outer/inner/Identity_k"), arg1); - - auto mul_jk = - ops::Mul(scope.WithOpName("outer/inner/mul"), identity_j, identity_k); - auto add_jkx = ops::Add(scope.WithOpName("outer/inner/add"), mul_jk, arg2); - auto assign = ops::AssignAddVariableOp( - scope.WithOpName("outer/inner/assign_add"), arg3, add_jkx); - - auto one = ops::Const( - scope.WithOpName("outer/inner/One") - .WithControlDependencies( - absl::Span{assign.operation}), - 1); - auto add_j = - ops::Add(scope.WithOpName("outer/inner/add_j"), identity_j, one); + FindWhileCondAndBody(graph_def, &outer_cond_fn, &outer_body_fn)); + + // Outer graph. + { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto x = ops::Placeholder(scope.WithOpName("x"), DT_INT32); + auto three = ops::Const(scope.WithOpName("three"), 3); + auto y = ops::Add(scope.WithOpName("y"), x, three); + + auto var = ops::VarHandleOp(scope.WithOpName("Variable"), DT_INT32, + TensorShape({})); + + auto zero = ops::Const(scope.WithOpName("outer/Const"), 0); + + auto while_op = ops::While(scope.WithOpName("outer/LoopCond"), + std::initializer_list{zero, y, x, var}, + outer_cond_fn, outer_body_fn); + auto sink = ops::Identity(scope.WithOpName("sink"), while_op[0]); + GraphDef expected; + TF_EXPECT_OK(scope.ToGraphDef(&expected)); + TF_EXPECT_GRAPH_EQ(expected, graph_def); + } - auto retval0 = ops::_Retval(scope.WithOpName("_retval0_RetVal"), add_j, 0); - auto retval1 = - ops::_Retval(scope.WithOpName("_retval1_RetVal"), identity_k, 1); - auto retval2 = ops::_Retval(scope.WithOpName("_retval2_RetVal"), arg2, 2); + // Outer condition graph. + { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto arg0 = ops::_Arg(scope.WithOpName("_arg0"), DT_INT32, 0); + auto arg1 = ops::_Arg(scope.WithOpName("_arg1"), DT_INT32, 1); + auto arg2 = ops::_Arg(scope.WithOpName("_arg2"), DT_INT32, 2); + auto arg3 = ops::_Arg(scope.WithOpName("_arg3"), DT_RESOURCE, 3); + + auto ten = ops::Const( + scope.WithOpName("outer/Less/y").WithControlDependencies(arg0.output), + 10); + auto less = ops::Less(scope.WithOpName("outer/Less_i"), arg0, ten); + auto retval = ops::_Retval(scope.WithOpName("_retval0_RetVal"), less, 0); + + GraphDef expected; + TF_EXPECT_OK(scope.ToGraphDef(&expected)); + + InstantiationResultForTest result; + TF_EXPECT_OK( + InstantiateFunctionForTest(outer_cond_fn.name(), library, &result)); + + EXPECT_EQ((DataTypeVector{DT_INT32, DT_INT32, DT_INT32, DT_RESOURCE}), + result.arg_types); + EXPECT_EQ(DataTypeVector{DT_BOOL}, result.ret_types); + TF_EXPECT_GRAPH_EQ(expected, result.gdef); + } - GraphDef expected; - TF_EXPECT_OK(scope.ToGraphDef(&expected)); + // Outer body graph. + NameAttrList inner_cond_fn, inner_body_fn; + { + InstantiationResultForTest result; + TF_EXPECT_OK( + InstantiateFunctionForTest(outer_body_fn.name(), library, &result)); + + // Find the inner condition and body names. + TF_EXPECT_OK( + FindWhileCondAndBody(result.gdef, &inner_cond_fn, &inner_body_fn)); + + Scope scope = Scope::NewRootScope().ExitOnError(); + auto arg0 = ops::_Arg(scope.WithOpName("_arg0"), DT_INT32, 0); + auto arg1 = ops::_Arg(scope.WithOpName("_arg1"), DT_INT32, 1); + auto arg2 = ops::_Arg(scope.WithOpName("_arg2"), DT_INT32, 2); + auto arg3 = ops::_Arg(scope.WithOpName("_arg3"), DT_RESOURCE, 3); + + auto identity_i = ops::Identity(scope.WithOpName("outer/Identity"), arg0); + auto one_j = ops::Const( + scope.WithOpName("outer/j").WithControlDependencies(identity_i), 1); + auto while_op = + ops::While(scope.WithOpName("outer/LoopCond_1"), + std::initializer_list{one_j, arg1, arg2, arg3}, + inner_cond_fn, inner_body_fn); + + auto one_outer = ops::Const( + scope.WithOpName("outer/add/y").WithControlDependencies(identity_i), + 1); + auto add_i = + ops::Add(scope.WithOpName("outer/add") + .WithControlDependencies(absl::Span{ + while_op[0].op(), while_op[1].op()}), + identity_i, one_outer); + + auto retval0 = + ops::_Retval(scope.WithOpName("_retval0_RetVal"), add_i, 0); + auto retval1 = ops::_Retval(scope.WithOpName("_retval1_RetVal"), arg1, 1); + auto retval2 = ops::_Retval(scope.WithOpName("_retval2_RetVal"), arg2, 2); + + GraphDef expected; + TF_EXPECT_OK(scope.ToGraphDef(&expected)); + + EXPECT_EQ((DataTypeVector{DT_INT32, DT_INT32, DT_INT32, DT_RESOURCE}), + result.arg_types); + EXPECT_EQ((DataTypeVector{DT_INT32, DT_INT32, DT_INT32}), + result.ret_types); + TF_EXPECT_GRAPH_EQ(expected, result.gdef); + } - InstantiationResultForTest result; - TF_EXPECT_OK( - InstantiateFunctionForTest(inner_body_fn.name(), library, &result)); + // Inner condition graph. + { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto arg0 = ops::_Arg(scope.WithOpName("_arg0"), DT_INT32, 0); + auto arg1 = ops::_Arg(scope.WithOpName("_arg1"), DT_INT32, 1); + auto arg2 = ops::_Arg(scope.WithOpName("_arg2"), DT_INT32, 2); + auto arg3 = ops::_Arg(scope.WithOpName("_arg3"), DT_RESOURCE, 3); + + auto five = ops::Const( + scope.WithOpName("outer/inner/Five").WithControlDependencies(arg0), + 5); + auto less_j = + ops::Less(scope.WithOpName("outer/inner/Less_j"), arg0, five); + auto retval = + ops::_Retval(scope.WithOpName("_retval0_RetVal"), less_j, 0); + + GraphDef expected; + TF_EXPECT_OK(scope.ToGraphDef(&expected)); + + InstantiationResultForTest result; + TF_EXPECT_OK( + InstantiateFunctionForTest(inner_cond_fn.name(), library, &result)); + + EXPECT_EQ((DataTypeVector{DT_INT32, DT_INT32, DT_INT32, DT_RESOURCE}), + result.arg_types); + EXPECT_EQ(DataTypeVector{DT_BOOL}, result.ret_types); + TF_EXPECT_GRAPH_EQ(expected, result.gdef); + } - EXPECT_EQ((DataTypeVector{DT_INT32, DT_INT32, DT_INT32, DT_RESOURCE}), - result.arg_types); - EXPECT_EQ((DataTypeVector{DT_INT32, DT_INT32, DT_INT32}), result.ret_types); - TF_EXPECT_GRAPH_EQ(expected, result.gdef); + // Inner body graph. + { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto arg0 = ops::_Arg(scope.WithOpName("_arg0"), DT_INT32, 0); + auto arg1 = ops::_Arg(scope.WithOpName("_arg1"), DT_INT32, 1); + auto arg2 = ops::_Arg(scope.WithOpName("_arg2"), DT_INT32, 2); + auto arg3 = ops::_Arg(scope.WithOpName("_arg3"), DT_RESOURCE, 3); + + auto identity_j = + ops::Identity(scope.WithOpName("outer/inner/Identity_j"), arg0); + auto identity_k = + ops::Identity(scope.WithOpName("outer/inner/Identity_k"), arg1); + + auto mul_jk = + ops::Mul(scope.WithOpName("outer/inner/mul"), identity_j, identity_k); + auto add_jkx = + ops::Add(scope.WithOpName("outer/inner/add"), mul_jk, arg2); + auto assign = ops::AssignAddVariableOp( + scope.WithOpName("outer/inner/assign_add"), arg3, add_jkx); + + auto one = ops::Const( + scope.WithOpName("outer/inner/One") + .WithControlDependencies( + absl::Span{assign.operation}), + 1); + auto add_j = + ops::Add(scope.WithOpName("outer/inner/add_j"), identity_j, one); + + auto retval0 = + ops::_Retval(scope.WithOpName("_retval0_RetVal"), add_j, 0); + auto retval1 = + ops::_Retval(scope.WithOpName("_retval1_RetVal"), identity_k, 1); + auto retval2 = ops::_Retval(scope.WithOpName("_retval2_RetVal"), arg2, 2); + + GraphDef expected; + TF_EXPECT_OK(scope.ToGraphDef(&expected)); + + InstantiationResultForTest result; + TF_EXPECT_OK( + InstantiateFunctionForTest(inner_body_fn.name(), library, &result)); + + EXPECT_EQ((DataTypeVector{DT_INT32, DT_INT32, DT_INT32, DT_RESOURCE}), + result.arg_types); + EXPECT_EQ((DataTypeVector{DT_INT32, DT_INT32, DT_INT32}), + result.ret_types); + TF_EXPECT_GRAPH_EQ(expected, result.gdef); + } } } diff --git a/tensorflow/compiler/tf2xla/graph_compiler.cc b/tensorflow/compiler/tf2xla/graph_compiler.cc index 706ed4f5bbfac60de4653cc8c326214cd4d8d886..5e4699bbb6218089d2e76a36c7351bf7fbd23264 100644 --- a/tensorflow/compiler/tf2xla/graph_compiler.cc +++ b/tensorflow/compiler/tf2xla/graph_compiler.cc @@ -22,10 +22,11 @@ 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_compilation_device.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" #include "tensorflow/compiler/tf2xla/xla_context.h" +#include "tensorflow/compiler/tf2xla/xla_expression.h" #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" #include "tensorflow/compiler/xla/client/client_library.h" #include "tensorflow/compiler/xla/client/xla_builder.h" @@ -40,6 +41,7 @@ limitations under the License. #include "tensorflow/core/graph/graph_constructor.h" #include "tensorflow/core/graph/node_builder.h" #include "tensorflow/core/graph/validate.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" @@ -51,12 +53,11 @@ namespace { Status PrepareArguments(XlaOpKernelContext* ctx, Graph* graph, const std::vector& expressions, std::vector* args) { - auto builder = ctx->builder(); auto client = ctx->compiler()->client(); - std::vector compile_time_constant_flags(expressions.size()); + std::vector arg_must_be_compile_time_constant(expressions.size()); TF_RETURN_IF_ERROR( - BackwardsConstAnalysis(*graph, &compile_time_constant_flags, + BackwardsConstAnalysis(*graph, &arg_must_be_compile_time_constant, /*compile_time_const_nodes=*/nullptr)); args->resize(expressions.size()); @@ -65,24 +66,34 @@ Status PrepareArguments(XlaOpKernelContext* ctx, Graph* graph, arg.type = ctx->input_type(i); arg.shape = ctx->InputShape(i); - if (arg.type == DT_RESOURCE) { - return errors::InvalidArgument( - "Resource as function argument is not yet implemented."); - } else if (expressions[i]->has_constant_value()) { - arg.kind = XlaCompiler::Argument::kConstant; - arg.constant_value = expressions[i]->constant_value(); - } else if (compile_time_constant_flags[i]) { - arg.kind = XlaCompiler::Argument::kConstant; - TF_RET_CHECK(expressions[i]->resource() == nullptr) - << "Input with resource is not yet implemented."; - TF_ASSIGN_OR_RETURN(auto constant_graph, builder->BuildConstantSubGraph( - expressions[i]->handle())); - TF_ASSIGN_OR_RETURN(auto literal, - client->ComputeConstant(constant_graph)); - TF_RETURN_IF_ERROR( - LiteralToHostTensor(literal, arg.type, &arg.constant_value)); - } else { - arg.kind = XlaCompiler::Argument::kParameter; + switch (expressions[i]->kind()) { + case XlaExpression::Kind::kConstant: + arg.kind = XlaCompiler::Argument::kConstant; + arg.constant_value = expressions[i]->constant_value(); + break; + case XlaExpression::Kind::kXlaOp: + if (arg_must_be_compile_time_constant[i]) { + TF_ASSIGN_OR_RETURN(absl::optional value, + expressions[i]->ResolveConstant(client)); + if (!value.has_value()) { + return errors::InvalidArgument( + "Argument to function must be a compile-time constant, but " + "unable to resolve argument value to a constant."); + } + arg.kind = XlaCompiler::Argument::kConstant; + arg.constant_value = *value; + } else { + arg.kind = XlaCompiler::Argument::kParameter; + } + break; + 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"); } } return Status::OK(); @@ -184,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; @@ -212,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)); @@ -227,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 @@ -244,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 9ee4178f5c213e919255bb33e9b15800a77256e6..52d2901e73d16f71ecbf633ede0d2cf553b6e521 100644 --- a/tensorflow/compiler/tf2xla/kernels/BUILD +++ b/tensorflow/compiler/tf2xla/kernels/BUILD @@ -1,16 +1,11 @@ +load("//tensorflow:tensorflow.bzl", "tf_copts", "tf_kernel_library") + licenses(["notice"]) # Apache 2.0 package( default_visibility = ["//tensorflow/compiler/tf2xla:internal"], ) -load("//tensorflow:tensorflow.bzl", "tf_copts") -load("//tensorflow:tensorflow.bzl", "tf_kernel_library") -load( - "//third_party/mkl:build_defs.bzl", - "if_mkl", -) - tf_kernel_library( name = "xla_ops", srcs = [ @@ -106,6 +101,7 @@ tf_kernel_library( "variable_ops.cc", "xla_broadcast_helper_op.cc", "xla_conv_op.cc", + "xla_dequantize_op.cc", "xla_dot_op.cc", "xla_pad_op.cc", "xla_reduce_op.cc", @@ -121,15 +117,10 @@ tf_kernel_library( ":while_op", "//tensorflow/compiler/tf2xla:common", "//tensorflow/compiler/tf2xla:xla_compiler", - "//tensorflow/compiler/tf2xla/lib:batch_dot", "//tensorflow/compiler/tf2xla/lib:broadcast", - "//tensorflow/compiler/tf2xla/lib:cholesky", - "//tensorflow/compiler/tf2xla/lib:qr", "//tensorflow/compiler/tf2xla/lib:random", "//tensorflow/compiler/tf2xla/lib:scatter", - "//tensorflow/compiler/tf2xla/lib:triangular_solve", "//tensorflow/compiler/tf2xla/lib:util", - "//tensorflow/compiler/tf2xla/lib:while_loop", "//tensorflow/compiler/tf2xla/ops:xla_ops", "//tensorflow/compiler/xla:array4d", "//tensorflow/compiler/xla:literal", @@ -142,19 +133,33 @@ tf_kernel_library( "//tensorflow/compiler/xla/client:xla_builder", "//tensorflow/compiler/xla/client:xla_computation", "//tensorflow/compiler/xla/client/lib:arithmetic", + "//tensorflow/compiler/xla/client/lib:cholesky", "//tensorflow/compiler/xla/client/lib:constants", + "//tensorflow/compiler/xla/client/lib:loops", "//tensorflow/compiler/xla/client/lib:math", - "//tensorflow/compiler/xla/client/lib:numeric", + "//tensorflow/compiler/xla/client/lib:matrix", "//tensorflow/compiler/xla/client/lib:pooling", "//tensorflow/compiler/xla/client/lib:prng", + "//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: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:protos_all_cc", + "//tensorflow/core:random_ops_op_lib", + "//tensorflow/core:resource_variable_ops_op_lib", "//tensorflow/core:spectral_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", @@ -178,6 +183,31 @@ tf_kernel_library( ], ) +# A separate cc_library for resampler_ops is needed because resampler is in +# contrib/, and thus the declaration of resampler cannot be pulled into the deps +# of xla_ops. Therefore, resampler_ops is its own cc_library target, and its +# corresponding tf_kernel_library is defined in contrib/resampler/BUILD. +cc_library( + name = "resampler_ops", + srcs = ["resampler_ops.cc"], + visibility = ["//visibility:public"], + deps = [ + "//tensorflow/compiler/tf2xla:common", + "//tensorflow/compiler/tf2xla:xla_compiler", + "//tensorflow/compiler/xla:array4d", + "//tensorflow/compiler/xla:literal", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla/client:xla_builder", + "//tensorflow/compiler/xla/client/lib:arithmetic", + "//tensorflow/compiler/xla/client/lib:constants", + "//tensorflow/core:framework", + "//tensorflow/core:lib", + "//tensorflow/core:protos_all_cc", + ], + alwayslink = 1, +) + cc_library( name = "conv_op_helpers", srcs = ["conv_op_helpers.cc"], @@ -190,7 +220,6 @@ cc_library( "//tensorflow/compiler/xla/client:xla_builder", "//tensorflow/compiler/xla/client/lib:arithmetic", "//tensorflow/compiler/xla/client/lib:constants", - "//tensorflow/compiler/xla/client/lib:numeric", "//tensorflow/core:framework", "//tensorflow/core/kernels:bounds_check", "//tensorflow/core/kernels:conv_ops", diff --git a/tensorflow/compiler/tf2xla/kernels/arg_op.cc b/tensorflow/compiler/tf2xla/kernels/arg_op.cc index 276d744c096f8996c774964204feaa3762bdb844..795ea09831e183a26fb3498b9bbaf9c3adaef9ed 100644 --- a/tensorflow/compiler/tf2xla/kernels/arg_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/arg_op.cc @@ -14,11 +14,13 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/compiler/tf2xla/type_util.h" +#include "tensorflow/compiler/tf2xla/xla_compilation_device.h" #include "tensorflow/compiler/tf2xla/xla_compiler.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/core/framework/kernel_def_builder.h" +#include "tensorflow/core/lib/core/errors.h" namespace tensorflow { @@ -48,14 +50,10 @@ class XlaArgOp : public XlaOpKernel { return; } - const XlaExpression& arg = XlaContext::Get(ctx).args()[index_]; - if (arg.resource() != nullptr) { - ctx->SetResourceOutput(0, arg.resource()); - } else if (arg.has_constant_value()) { - ctx->SetConstantOutput(0, arg.constant_value()); - } else { - ctx->SetOutput(0, arg.handle()); - } + 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); } private: diff --git a/tensorflow/compiler/tf2xla/kernels/batch_matmul_op.cc b/tensorflow/compiler/tf2xla/kernels/batch_matmul_op.cc index 4cfe946b2e6146f034867c06e996ffae42b90705..1b254e328a8c71bd81a0ec700e2af1d81a5fa67a 100644 --- a/tensorflow/compiler/tf2xla/kernels/batch_matmul_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/batch_matmul_op.cc @@ -13,9 +13,11 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/compiler/tf2xla/lib/batch_dot.h" +#include "tensorflow/compiler/tf2xla/lib/util.h" #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" +#include "tensorflow/compiler/xla/client/lib/math.h" +#include "tensorflow/compiler/xla/client/lib/matrix.h" namespace tensorflow { namespace { @@ -28,9 +30,11 @@ class BatchMatMulOp : public XlaOpKernel { } void Compile(XlaOpKernelContext* ctx) override { - auto result = BatchDot(ctx->Input(0), ctx->Input(1), - /*transpose_x=*/adj_x_, /*transpose_y=*/adj_y_, - /*conjugate_x=*/adj_x_, /*conjugate_y=*/adj_y_); + auto result = + xla::BatchDot(MaybeTransposeInMinorDims( + MaybeConjugate(ctx->Input(0), adj_x_), adj_x_), + MaybeTransposeInMinorDims( + MaybeConjugate(ctx->Input(1), adj_y_), adj_y_)); ctx->SetOutput(0, result); } diff --git a/tensorflow/compiler/tf2xla/kernels/batch_norm_op.cc b/tensorflow/compiler/tf2xla/kernels/batch_norm_op.cc index a267c0c72fce67d7c22c55a57f8d5ac4ffd2b7e2..0e2f335f3354e3ae6008bdc0ac0b80683fe479c1 100644 --- a/tensorflow/compiler/tf2xla/kernels/batch_norm_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/batch_norm_op.cc @@ -115,9 +115,9 @@ class FusedBatchNormGradOp : public XlaOpKernel { // operators. For now, cast everything to the statistics type (which // may be more precise than the input type). auto grad_backprop = - XlaHelpers::ConvertElementType(b, ctx->Input(0), scale_dtype); + XlaHelpers::ConvertElementType(ctx->Input(0), scale_dtype); auto activations = - XlaHelpers::ConvertElementType(b, ctx->Input(1), scale_dtype); + XlaHelpers::ConvertElementType(ctx->Input(1), scale_dtype); auto scale = ctx->Input(2); auto mean = ctx->Input(3); auto var = ctx->Input(4); @@ -151,11 +151,11 @@ class FusedBatchNormGradOp : public XlaOpKernel { const DataType accumulation_type = XlaHelpers::SumAccumulationType(scale_dtype); auto converted = - XlaHelpers::ConvertElementType(b, grad_backprop, accumulation_type); + XlaHelpers::ConvertElementType(grad_backprop, accumulation_type); auto reduce = xla::Reduce(converted, XlaHelpers::Zero(b, accumulation_type), *ctx->GetOrCreateAdd(accumulation_type), reduction_dims); - offset_backprop = XlaHelpers::ConvertElementType(b, reduce, scale_dtype); + offset_backprop = XlaHelpers::ConvertElementType(reduce, scale_dtype); // scratch1 = rsqrt(pop_var + epsilon) auto neg_half = XlaHelpers::FloatLiteral(b, scale_dtype, -0.5); @@ -165,19 +165,18 @@ class FusedBatchNormGradOp : public XlaOpKernel { // scratch2 = sum(y_backprop * (x - mean)) auto mul = xla::Mul(grad_backprop, xla::Sub(activations, mean, {feature_index})); - converted = XlaHelpers::ConvertElementType(b, mul, accumulation_type); + converted = XlaHelpers::ConvertElementType(mul, accumulation_type); reduce = xla::Reduce(converted, XlaHelpers::Zero(b, accumulation_type), *ctx->GetOrCreateAdd(accumulation_type), reduction_dims); - auto scratch2 = XlaHelpers::ConvertElementType(b, reduce, scale_dtype); + auto scratch2 = XlaHelpers::ConvertElementType(reduce, scale_dtype); x_backprop = xla::Mul(grad_backprop, xla::Mul(scratch1, scale), {feature_index}); scale_backprop = xla::Mul(scratch1, scratch2); } - ctx->SetOutput(0, - XlaHelpers::ConvertElementType(b, x_backprop, input_dtype)); + ctx->SetOutput(0, XlaHelpers::ConvertElementType(x_backprop, input_dtype)); ctx->SetOutput(1, scale_backprop); ctx->SetOutput(2, offset_backprop); ctx->SetConstantOutput(3, Tensor()); diff --git a/tensorflow/compiler/tf2xla/kernels/batchtospace_op.cc b/tensorflow/compiler/tf2xla/kernels/batchtospace_op.cc index a18e04995b5e1e0b0374f7b0edd6f5e114cf994a..6b675fa8a94e0bc932baaa359565cbc8e4614ee5 100644 --- a/tensorflow/compiler/tf2xla/kernels/batchtospace_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/batchtospace_op.cc @@ -39,7 +39,7 @@ void BatchToSpace(XlaOpKernelContext* ctx, const xla::XlaOp& input, OP_REQUIRES( ctx, - xla::ShapeUtil::Rank(crops.shape()) == 2 && + crops.shape().rank() == 2 && block_rank == xla::ShapeUtil::GetDimension(crops.shape(), 0) && 2 == xla::ShapeUtil::GetDimension(crops.shape(), 1), errors::InvalidArgument("crops should have shape [", block_rank, @@ -159,8 +159,8 @@ class BatchToSpaceNDOp : public XlaOpKernel { } }; REGISTER_XLA_OP(Name("BatchToSpaceND") - .CompileTimeConstInput("block_shape") - .CompileTimeConstInput("crops"), + .CompileTimeConstantInput("block_shape") + .CompileTimeConstantInput("crops"), BatchToSpaceNDOp); class BatchToSpaceOp : public XlaOpKernel { @@ -183,7 +183,7 @@ class BatchToSpaceOp : public XlaOpKernel { private: int block_size_; }; -REGISTER_XLA_OP(Name("BatchToSpace").CompileTimeConstInput("crops"), +REGISTER_XLA_OP(Name("BatchToSpace").CompileTimeConstantInput("crops"), BatchToSpaceOp); } // namespace diff --git a/tensorflow/compiler/tf2xla/kernels/bcast_ops.cc b/tensorflow/compiler/tf2xla/kernels/bcast_ops.cc index 182f7c99344845964f7010127718f876ab6e8a44..c022284fec6bc91951170e243ea3609c8d5d0c43 100644 --- a/tensorflow/compiler/tf2xla/kernels/bcast_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/bcast_ops.cc @@ -67,8 +67,8 @@ class BCastArgsOp : public XlaOpKernel { TF_DISALLOW_COPY_AND_ASSIGN(BCastArgsOp); }; REGISTER_XLA_OP(Name("BroadcastArgs") - .CompileTimeConstInput("s0") - .CompileTimeConstInput("s1"), + .CompileTimeConstantInput("s0") + .CompileTimeConstantInput("s1"), BCastArgsOp); // Given shapes of two tensors, computes the reduction indices for the @@ -94,14 +94,10 @@ class BCastGradArgsOp : public XlaOpKernel { OP_REQUIRES(ctx, TensorShapeUtils::IsVector(in_shape), errors::InvalidArgument("In[", i, "] must be a vector.", in_shape.DebugString())); - xla::Literal literal; - OP_REQUIRES_OK(ctx, ctx->ConstantInput(i, &literal)); - - BCast::Vec vec; - for (int64 i = 0; i < in_shape.num_elements(); ++i) { - vec.push_back(literal.Get({i})); - } - shapes.push_back(vec); + std::vector vec; + OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector(i, &vec)); + + shapes.push_back(BCast::Vec(vec.begin(), vec.end())); } BCast bcast(shapes[0], shapes[1]); OP_REQUIRES(ctx, bcast.IsValid(), @@ -126,8 +122,8 @@ class BCastGradArgsOp : public XlaOpKernel { }; REGISTER_XLA_OP(Name("BroadcastGradientArgs") - .CompileTimeConstInput("s0") - .CompileTimeConstInput("s1"), + .CompileTimeConstantInput("s0") + .CompileTimeConstantInput("s1"), BCastGradArgsOp); } // namespace diff --git a/tensorflow/compiler/tf2xla/kernels/bias_ops.cc b/tensorflow/compiler/tf2xla/kernels/bias_ops.cc index 41f540506ba41fbe7f91393e7b8e26a89e72ef0a..e7f369b761f36a717ea5fb536780af91a8955b1e 100644 --- a/tensorflow/compiler/tf2xla/kernels/bias_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/bias_ops.cc @@ -107,11 +107,11 @@ class BiasAddGradOp : public XlaOpKernel { const DataType accumulation_type = XlaHelpers::SumAccumulationType(input_type(0)); auto converted = - XlaHelpers::ConvertElementType(b, ctx->Input(0), accumulation_type); + XlaHelpers::ConvertElementType(ctx->Input(0), accumulation_type); auto reduce = xla::Reduce(converted, XlaHelpers::Zero(b, accumulation_type), *ctx->GetOrCreateAdd(accumulation_type), reduce_dims); - ctx->SetOutput(0, XlaHelpers::ConvertElementType(b, reduce, input_type(0))); + ctx->SetOutput(0, XlaHelpers::ConvertElementType(reduce, input_type(0))); } private: diff --git a/tensorflow/compiler/tf2xla/kernels/binary_ops.cc b/tensorflow/compiler/tf2xla/kernels/binary_ops.cc index 47e517a6576d3a848bc41ceb703df2bd778c4a35..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" @@ -43,6 +45,9 @@ namespace { const std::vector& extend_dimensions) override { \ xla::XlaBuilder* b = ctx->builder(); \ (void)b; \ + (void)lhs_shape; \ + (void)rhs_shape; \ + (void)extend_dimensions; \ return HLO; \ } \ }; \ @@ -103,23 +108,23 @@ static xla::XlaOp FloorDivImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x, XLA_MAKE_BINARY(FloorDiv, FloorDivImpl(b, input_type(0), lhs, rhs, broadcast_helper)); -static xla::XlaOp XlogyImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x, - xla::XlaOp y, const BCast& broadcast_helper) { +xla::XlaOp XlogyImpl(xla::XlaOp x, xla::XlaOp y, + const BCast& broadcast_helper) { std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper); - auto zero = XlaHelpers::Zero(b, dtype); + auto zero = xla::ZerosLike(x); auto is_zero = xla::Eq(x, zero); return xla::Select(is_zero, zero, xla::Mul(x, xla::Log(y))); } -XLA_MAKE_BINARY(Xlogy, XlogyImpl(b, input_type(0), lhs, rhs, broadcast_helper)); +XLA_MAKE_BINARY(Xlogy, XlogyImpl(lhs, rhs, broadcast_helper)); -static xla::XlaOp XdivyImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x, - xla::XlaOp y, const BCast& broadcast_helper) { +xla::XlaOp XdivyImpl(xla::XlaOp x, xla::XlaOp y, + const BCast& broadcast_helper) { std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper); - auto zero = XlaHelpers::Zero(b, dtype); + auto zero = xla::ZerosLike(x); auto is_zero = xla::Eq(x, zero); return xla::Select(is_zero, zero, xla::Div(x, y)); } -XLA_MAKE_BINARY(Xdivy, XdivyImpl(b, input_type(0), lhs, rhs, broadcast_helper)); +XLA_MAKE_BINARY(Xdivy, XdivyImpl(lhs, rhs, broadcast_helper)); // Implementation of FloorMod. Pseudo-code: // T trunc_mod = std::fmod(x, y); @@ -162,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)); @@ -192,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)), @@ -201,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/broadcast_to_op.cc b/tensorflow/compiler/tf2xla/kernels/broadcast_to_op.cc index 9bb11fb67e3e4ddc48d68631c60f96c60b921094..d7a8e67dd33aab5c32b7465ce505b745b5c1ca2f 100644 --- a/tensorflow/compiler/tf2xla/kernels/broadcast_to_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/broadcast_to_op.cc @@ -38,7 +38,7 @@ class BroadcastToOp : public XlaOpKernel { } }; -REGISTER_XLA_OP(Name("BroadcastTo").CompileTimeConstInput("shape"), +REGISTER_XLA_OP(Name("BroadcastTo").CompileTimeConstantInput("shape"), BroadcastToOp); } // namespace 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 e7fef77edcba0ea5a521956a704225ac4f7fcb22..c2b4c28d1566f5429c5d8109db94af0c3762b131 100644 --- a/tensorflow/compiler/tf2xla/kernels/categorical_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/categorical_op.cc @@ -21,10 +21,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/client/lib/arithmetic.h" +#include "tensorflow/compiler/xla/client/lib/prng.h" #include "tensorflow/compiler/xla/client/xla_builder.h" +#include "tensorflow/compiler/xla/xla_data.pb.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.pb.h" namespace tensorflow { namespace { @@ -57,42 +60,114 @@ class CategoricalOp : public XlaOpKernel { const int64 batch_size = logits_shape.dim_size(0); const int64 num_classes = logits_shape.dim_size(1); - xla::XlaBuilder* builder = ctx->builder(); - - std::array uniform_shape_array = { - {batch_size, num_samples, num_classes}}; - xla::PrimitiveType uniform_xla_type; - OP_REQUIRES_OK(ctx, - DataTypeToPrimitiveType(input_type(0), &uniform_xla_type)); - xla::Shape uniform_shape = - xla::ShapeUtil::MakeShape(uniform_xla_type, uniform_shape_array); - auto uniforms = - xla::RngUniform(XlaHelpers::Zero(builder, input_type(0)), - XlaHelpers::One(builder, input_type(0)), uniform_shape); + xla::Shape uniform_shape; + int class_dimension; + if (num_samples != 1) { + std::array uniform_shape_array = { + {batch_size, num_samples, num_classes}}; + xla::PrimitiveType uniform_xla_type; + OP_REQUIRES_OK(ctx, + DataTypeToPrimitiveType(input_type(0), &uniform_xla_type)); + uniform_shape = + xla::ShapeUtil::MakeShape(uniform_xla_type, uniform_shape_array); + class_dimension = 2; + } else { + // Have a special case for when we only need one sample, because + // dimensions may be padded on architectures with tiled memory layouts, so + // if the num_classes or batch size is large then this can lead to + // expensive wasted memory. + std::array uniform_shape_array = {{batch_size, num_classes}}; + xla::PrimitiveType uniform_xla_type; + OP_REQUIRES_OK(ctx, + DataTypeToPrimitiveType(input_type(0), &uniform_xla_type)); + uniform_shape = + xla::ShapeUtil::MakeShape(uniform_xla_type, uniform_shape_array); + class_dimension = 1; + } + xla::PrimitiveType type; + OP_REQUIRES_OK(ctx, DataTypeToPrimitiveType(input_type(0), &type)); + xla::XlaOp log_uniforms = GetLogUniforms(uniform_shape, type, ctx); // Use Gumbel softmax trick to generate categorical samples. // See: // https://hips.seas.harvard.edu/blog/2013/04/06/the-gumbel-max-trick-for-discrete-distributions/ // TODO(b/68769470): Switch to using a cumulative sum approach. - auto softmax_entries = xla::Sub(logits, xla::Log(-xla::Log(uniforms)), - /*broadcast_dimensions=*/{0, 2}); + auto softmax_entries = + xla::Sub(logits, log_uniforms, + /*broadcast_dimensions=*/{0, class_dimension}); 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=*/2); + xla::XlaOp argmax = xla::ArgMax(softmax_entries, xla_output_type, + /*axis=*/class_dimension); + if (num_samples == 1) { + argmax = xla::Reshape(argmax, {batch_size, 1}); + } ctx->SetOutput(0, argmax); } + virtual xla::XlaOp GetLogUniforms(xla::Shape uniform_shape, + xla::PrimitiveType type, + XlaOpKernelContext* ctx) { + xla::XlaBuilder* builder = ctx->builder(); + auto uniforms = + xla::RngUniform(XlaHelpers::Zero(builder, input_type(0)), + XlaHelpers::One(builder, input_type(0)), uniform_shape); + return xla::Log(-xla::Log(uniforms)); + } + private: TF_DISALLOW_COPY_AND_ASSIGN(CategoricalOp); }; // TODO(b/68769717): Rename this sampler to Categorical. -REGISTER_XLA_OP(Name("Multinomial").CompileTimeConstInput("num_samples"), +REGISTER_XLA_OP(Name("Multinomial").CompileTimeConstantInput("num_samples"), CategoricalOp); +class StatelessCategoricalOp : public CategoricalOp { + public: + explicit StatelessCategoricalOp(OpKernelConstruction* ctx) + : CategoricalOp(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("T", &dtype_)); + } + + xla::XlaOp GetLogUniforms(xla::Shape uniform_shape, xla::PrimitiveType type, + XlaOpKernelContext* ctx) override { + xla::XlaOp seed = ctx->Input(2); + auto seed0 = xla::Reshape(xla::Slice(seed, {0}, {1}, {1}), {}); + auto seed1 = xla::Reshape(xla::Slice(seed, {1}, {2}, {1}), {}); + + xla::XlaBuilder* builder = ctx->builder(); + if (uniform_shape.element_type() == xla::BF16) { + uniform_shape.set_element_type(xla::F32); + } + auto uniforms = xla::StatelessRngUniform( + {seed0, seed1}, uniform_shape, XlaHelpers::Zero(builder, DT_FLOAT), + XlaHelpers::One(builder, DT_FLOAT)); + return xla::ConvertElementType(xla::Log(-xla::Log(uniforms)), type); + } + + void Compile(XlaOpKernelContext* ctx) override { + TensorShape seed_shape = ctx->InputShape(2); + OP_REQUIRES(ctx, seed_shape.dims() == 1 && seed_shape.dim_size(0) == 2, + errors::InvalidArgument("seed must have shape [2], not ", + seed_shape.DebugString())); + CategoricalOp::Compile(ctx); + } + + private: + DataType dtype_; + + TF_DISALLOW_COPY_AND_ASSIGN(StatelessCategoricalOp); +}; + +REGISTER_XLA_OP(Name("StatelessMultinomial") + .CompileTimeConstantInput("num_samples") + .TypeConstraint("T", {DT_FLOAT, DT_BFLOAT16}) + .TypeConstraint("Tseed", DT_INT32), + StatelessCategoricalOp); + } // anonymous namespace } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/cholesky_op.cc b/tensorflow/compiler/tf2xla/kernels/cholesky_op.cc index 9fcbc86adc0967cbb7fb73da8bdabc58b60953da..0ed3044efa5b1060d2b0ad2d5563b0e02ebf66ec 100644 --- a/tensorflow/compiler/tf2xla/kernels/cholesky_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/cholesky_op.cc @@ -13,9 +13,9 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/compiler/tf2xla/lib/cholesky.h" #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" +#include "tensorflow/compiler/xla/client/lib/cholesky.h" namespace tensorflow { namespace { @@ -24,7 +24,7 @@ class CholeskyOp : public XlaOpKernel { public: explicit CholeskyOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} void Compile(XlaOpKernelContext* ctx) override { - ctx->SetOutput(0, Cholesky(ctx->Input(0))); + ctx->SetOutput(0, xla::Cholesky(ctx->Input(0))); } }; diff --git a/tensorflow/compiler/tf2xla/kernels/concat_op.cc b/tensorflow/compiler/tf2xla/kernels/concat_op.cc index 0ae23aa6dfe49048ac5cb8ae00c12432b2e2a2fe..91e4d9cea7cbf6075e30250587044174c4b8e7f4 100644 --- a/tensorflow/compiler/tf2xla/kernels/concat_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/concat_op.cc @@ -24,12 +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" @@ -37,16 +38,6 @@ limitations under the License. namespace tensorflow { namespace { -// Used to determine the number of Tensors allowed in a Concat op to prevent -// going over the max gpu parameter memory size. This is an issue because concat -// is variadic and can have an unlimited number of arguments when called. -// Concat ops with more Tensors than this will be split into multiple concat -// ops. -// -// TODO(b/112613927): Remove the logic here and put it properly in an HLO pass -// along with boxing large numbers of parameters. -constexpr int64 kMaxConcatArgsPerOp = 500; - // -------------------------------------------------------------------------- class ConcatBaseOp : public XlaOpKernel { public: @@ -55,15 +46,13 @@ class ConcatBaseOp : public XlaOpKernel { void Compile(XlaOpKernelContext* ctx) override { const TensorShape concat_dim_tensor_shape = ctx->InputShape(axis_index_); - OP_REQUIRES( - ctx, IsLegacyScalar(concat_dim_tensor_shape), - errors::InvalidArgument( - "Concat dim tensor should be a scalar integer, but got shape ", - concat_dim_tensor_shape.DebugString())); - xla::Literal literal; - OP_REQUIRES_OK(ctx, ctx->ConstantInput(axis_index_, &literal)); - // TODO(annarev): add a helper to support int64 input. - const int32 concat_dim = literal.Get({}); + OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(concat_dim_tensor_shape), + errors::InvalidArgument( + "Concat dim tensor should be a scalar, but got shape ", + concat_dim_tensor_shape.DebugString())); + int64 concat_dim; + OP_REQUIRES_OK(ctx, + ctx->ConstantInputAsIntScalar(axis_index_, &concat_dim)); std::vector values; std::vector shapes; @@ -73,9 +62,7 @@ class ConcatBaseOp : public XlaOpKernel { const TensorShape& input_shape = shapes[0]; int32 axis = concat_dim < 0 ? concat_dim + input_dims : concat_dim; - OP_REQUIRES(ctx, - (0 <= axis && axis < input_dims) || - (allow_legacy_scalars() && concat_dim == 0), + OP_REQUIRES(ctx, 0 <= axis && axis < input_dims, errors::InvalidArgument( "ConcatOp : Expected concatenating dimensions in the range " "[", @@ -84,16 +71,12 @@ class ConcatBaseOp : public XlaOpKernel { // Make a vector holding the XlaOp for each of the inputs that has non-zero // elements. std::vector input_data; - std::vector partial_concats; int output_concat_dim = 0; - const bool input_is_scalar = IsLegacyScalar(input_shape); for (int i = 0; i < N; ++i) { xla::XlaOp handle = values[i]; const TensorShape& in_shape = shapes[i]; - const bool in_is_scalar = IsLegacyScalar(in_shape); OP_REQUIRES( - ctx, - in_shape.dims() == input_dims || (input_is_scalar && in_is_scalar), + ctx, in_shape.dims() == input_dims, errors::InvalidArgument( "ConcatOp : Ranks of all input tensors should match: shape[0] = ", input_shape.DebugString(), " vs. shape[", i, @@ -105,30 +88,10 @@ class ConcatBaseOp : public XlaOpKernel { input_data.push_back(handle); } output_concat_dim += in_shape.dims() > 0 ? in_shape.dim_size(axis) : 1; - - // Concat is associative, so it can be split into many operations when too - // many arguments are in a single op. This is a temporary workaround for - // b/112613927 where too many parameters in an XlaLaunchOp later result in - // too many parameters to a single GPU kernel. - if (i && i % kMaxConcatArgsPerOp == 0) { - partial_concats.push_back( - xla::ConcatInDim(ctx->builder(), input_data, axis)); - input_data.clear(); - } } - // Add any inputs that have not been put into another concat yet. - partial_concats.insert(partial_concats.end(), input_data.begin(), - input_data.end()); VLOG(1) << "Concat dim " << concat_dim << " equivalent to " << axis; - // Don't add an additional "identity" concatenate for better readibility of - // IR. - if (partial_concats.size() == 1) { - ctx->SetOutput(0, partial_concats.front()); - } else { - ctx->SetOutput(0, - xla::ConcatInDim(ctx->builder(), partial_concats, axis)); - } + ctx->SetOutput(0, xla::ConcatInDim(ctx->builder(), input_data, axis)); } private: @@ -149,10 +112,11 @@ class ConcatV2Op : public ConcatBaseOp { : ConcatBaseOp(c, /* axis_index */ c->num_inputs() - 1) {} }; -REGISTER_XLA_OP(Name("Concat").CompileTimeConstInput("concat_dim"), ConcatOp); +REGISTER_XLA_OP(Name("Concat").CompileTimeConstantInput("concat_dim"), + ConcatOp); REGISTER_XLA_OP(Name("ConcatV2") .TypeConstraint("Tidx", DT_INT32) - .CompileTimeConstInput("axis"), + .CompileTimeConstantInput("axis"), ConcatV2Op); class ConcatOffsetOp : public XlaOpKernel { @@ -161,11 +125,10 @@ class ConcatOffsetOp : public XlaOpKernel { void Compile(XlaOpKernelContext* ctx) override { const TensorShape concat_dim_shape = ctx->InputShape(0); - OP_REQUIRES( - ctx, IsLegacyScalar(concat_dim_shape), - errors::InvalidArgument( - "Concat dim tensor should be a scalar integer, but got shape ", - concat_dim_shape.DebugString())); + OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(concat_dim_shape), + errors::InvalidArgument( + "Concat dim tensor should be a scalar, but got shape ", + concat_dim_shape.DebugString())); for (int i = 1; i < ctx->num_inputs(); ++i) { OP_REQUIRES(ctx, TensorShapeUtils::IsVector(ctx->InputShape(i)), errors::InvalidArgument("input ", i, @@ -192,39 +155,38 @@ class ConcatOffsetOp : public XlaOpKernel { // [0, 5, 0, 0] const int32 N = ctx->num_inputs() - 1; const TensorShape inp0_shape = ctx->InputShape(1); - xla::Literal inp0_literal; - OP_REQUIRES_OK(ctx, ctx->ConstantInput(1, &inp0_literal)); - const int64 dims = inp0_shape.num_elements(); + std::vector inp0_dims; + OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector(1, &inp0_dims)); + const int64 inp0_rank = inp0_shape.num_elements(); - xla::Literal concat_dim_literal; - OP_REQUIRES_OK(ctx, ctx->ConstantInput(0, &concat_dim_literal)); - const int64 cdim = concat_dim_literal.Get({}); + int64 cdim; + OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntScalar(0, &cdim)); - VLOG(1) << "ConcatOffset " << cdim << "," << dims; - int32 axis = cdim < 0 ? cdim + dims : cdim; - OP_REQUIRES(ctx, FastBoundsCheck(axis, dims), + VLOG(1) << "ConcatOffset " << cdim << "," << inp0_rank; + int32 axis = cdim < 0 ? cdim + inp0_rank : cdim; + OP_REQUIRES(ctx, FastBoundsCheck(axis, inp0_rank), errors::InvalidArgument("Concat dim is out of range: ", axis, - " vs. ", dims)); + " vs. ", inp0_rank)); int32 offset = 0; for (int i = 0; i < N; ++i) { const TensorShape inp_shape = ctx->InputShape(1 + i); - OP_REQUIRES(ctx, dims == inp_shape.num_elements(), - errors::InvalidArgument("input ", i, " should contain ", dims, - " elements, but got ", + OP_REQUIRES(ctx, inp0_rank == inp_shape.num_elements(), + errors::InvalidArgument("input ", i, " should contain ", + inp0_rank, " elements, but got ", inp_shape.num_elements())); - xla::Literal inp_literal; - OP_REQUIRES_OK(ctx, ctx->ConstantInput(1 + i, &inp_literal)); + std::vector inp_dims; + OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector(1 + i, &inp_dims)); - Tensor out_constant(DT_INT32, TensorShape({dims})); + Tensor out_constant(DT_INT32, TensorShape({inp0_rank})); auto out_vec = out_constant.vec(); - for (int64 j = 0; j < dims; ++j) { + for (int64 j = 0; j < inp0_rank; ++j) { if (j == axis) { out_vec(j) = offset; - offset += inp_literal.Get({j}); + offset += inp_dims[j]; } else { - const int32 inp0_element = inp0_literal.Get({j}); - const int32 inp_element = inp_literal.Get({j}); - OP_REQUIRES(ctx, (inp0_element == inp_element), + const int32 inp0_element = inp0_dims[j]; + const int32 inp_element = inp_dims[j]; + OP_REQUIRES(ctx, inp0_element == inp_element, errors::InvalidArgument("input[", i, ",", j, "] mismatch: ", inp0_element, " vs. ", inp_element)); @@ -238,8 +200,8 @@ class ConcatOffsetOp : public XlaOpKernel { }; REGISTER_XLA_OP(Name("ConcatOffset") - .CompileTimeConstInput("concat_dim") - .CompileTimeConstInput("shape"), + .CompileTimeConstantInput("concat_dim") + .CompileTimeConstantInput("shape"), ConcatOffsetOp); } // namespace diff --git a/tensorflow/compiler/tf2xla/kernels/const_op.cc b/tensorflow/compiler/tf2xla/kernels/const_op.cc index 2628ef8e2454976aeff3859fa5dc1d8e106f32e1..ff6c54e47c62f0555ef045e25051f6ec5a3c1d39 100644 --- a/tensorflow/compiler/tf2xla/kernels/const_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/const_op.cc @@ -42,11 +42,6 @@ class ConstOp : public XlaOpKernel { void Compile(XlaOpKernelContext* ctx) override { TensorShape shape(proto_.tensor_shape()); - if (proto_.dtype() == DT_STRING) { - LOG(WARNING) << "Not computing Const of type DT_STRING"; - ctx->SetInvalidOutput(0); - return; - } xla::XlaBuilder* b = ctx->builder(); // To avoid blowups for large constants filled with the same value, @@ -88,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 c9a1be494066e4f935a1d818bc86c86333e34fae..5f99b24e221ba6c926032ef7a1b4bf1e92df7a68 100644 --- a/tensorflow/compiler/tf2xla/kernels/conv_op_helpers.cc +++ b/tensorflow/compiler/tf2xla/kernels/conv_op_helpers.cc @@ -24,16 +24,15 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "tensorflow/compiler/xla/client/lib/arithmetic.h" #include "tensorflow/compiler/xla/client/lib/constants.h" -#include "tensorflow/compiler/xla/client/lib/numeric.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" @@ -65,60 +64,63 @@ xla::Shape ExpandedFilterShapeForDepthwiseConvolution(const xla::Shape& shape) { // 0 0 1 1 0 0 0 0 1 1 0 0 // 0 0 0 0 1 1 0 0 0 0 1 1 // -// The first step is to create a one tensor, A, that is [3] -// 0 1 2 +// The first step is to create a iota A with iota_dimension = 2 +// 0 0 0 0 0 0 0 0 0 0 0 0 +// 1 1 1 1 1 1 1 1 1 1 1 1 +// 2 2 2 2 2 2 2 2 2 2 2 2 // -// and another tensor, B, that is [3 * 2] -// 0 1 2 3 4 5 +// 0 0 0 0 0 0 0 0 0 0 0 0 +// 1 1 1 1 1 1 1 1 1 1 1 1 +// 2 2 2 2 2 2 2 2 2 2 2 2 // -// and divide B it by 2 to get -// 0 0 1 1 2 2 +// and another iota B with iota_dimension = 3 +// 0 1 2 3 4 5 0 1 2 3 4 5 +// 0 1 2 3 4 5 0 1 2 3 4 5 +// 0 1 2 3 4 5 0 1 2 3 4 5 // -// then we broadcast the B to [2, 2, 3, 3 * 2] -// 0 0 1 1 2 2 0 0 1 1 2 2 -// 0 0 1 1 2 2 0 0 1 1 2 2 -// 0 0 1 1 2 2 0 0 1 1 2 2 +// 0 1 2 3 4 5 0 1 2 3 4 5 +// 0 1 2 3 4 5 0 1 2 3 4 5 +// 0 1 2 3 4 5 0 1 2 3 4 5 // -// 0 0 1 1 2 2 0 0 1 1 2 2 -// 0 0 1 1 2 2 0 0 1 1 2 2 -// 0 0 1 1 2 2 0 0 1 1 2 2 +// and divide B by 2 to get +// 0 0 1 1 2 2 0 0 1 1 2 2 +// 0 0 1 1 2 2 0 0 1 1 2 2 +// 0 0 1 1 2 2 0 0 1 1 2 2 // -// Finally compare A and broadcasted B in dimension 2 amd return the result at -// the beginning of the comment. +// 0 0 1 1 2 2 0 0 1 1 2 2 +// 0 0 1 1 2 2 0 0 1 1 2 2 +// 0 0 1 1 2 2 0 0 1 1 2 2 +// +// Finally compare A and B and return the result at the beginning of the +// comment. xla::XlaOp CreateExpandedFilterMask(const xla::Shape& filter_shape, xla::XlaBuilder* builder) { xla::Shape expanded_filter_shape = ExpandedFilterShapeForDepthwiseConvolution(filter_shape); int64 depthwise_multiplier = filter_shape.dimensions(filter_shape.dimensions_size() - 1); - int64 input_feature = - filter_shape.dimensions(filter_shape.dimensions_size() - 2); - - // Create a M sized linspace and an M*N sized linspace that will be - // broadcasted into perpendicular dimensions and compared. - xla::XlaOp input_feature_iota = xla::Iota(builder, xla::S32, input_feature); - xla::XlaOp expanded_feature_iota = - xla::Iota(builder, xla::S32, input_feature * depthwise_multiplier); - // Divide the M*N sized linspace by the depthwise_multiplier to create - // [0 0 1 1 2 2] in the example in the function comment. + // Create two iotas with the shape of the expanded filter, one of them with + // the iota dimension chosen as the feature dimension, and the other a iota + // with the iota dimension chosen as the expanded output feature dimension. + std::vector iota_dimensions(expanded_filter_shape.dimensions().begin(), + expanded_filter_shape.dimensions().end()); + xla::Shape iota_shape = xla::ShapeUtil::MakeShape(xla::S32, iota_dimensions); + xla::XlaOp input_feature_iota = xla::Iota( + builder, iota_shape, /*iota_dimension=*/iota_dimensions.size() - 2); + xla::XlaOp expanded_feature_iota = xla::Iota( + builder, iota_shape, /*iota_dimension=*/iota_dimensions.size() - 1); + + // Divide 'expanded_feature_iota' by the depthwise_multiplier to create + // [0 0 1 1 2 2] ... in the example in the function comment. expanded_feature_iota = xla::Div(expanded_feature_iota, XlaHelpers::IntegerLiteral(builder, DataType::DT_INT32, depthwise_multiplier)); - // Broadcast the N*M linspace to [H, W, ..., M, M*N]. - std::vector expanded_feature_broadcast_dims( - expanded_filter_shape.dimensions().begin(), - expanded_filter_shape.dimensions().end()); - expanded_feature_broadcast_dims.pop_back(); - auto broadcasted_expanded_feature_iota = - xla::Broadcast(expanded_feature_iota, expanded_feature_broadcast_dims); - - // Compare the broadcasted linspace to the input feature linspace in the - // input feature dimension to create a diagonal predicate. - return xla::Eq(broadcasted_expanded_feature_iota, input_feature_iota, - {expanded_filter_shape.dimensions_size() - 2}); + // Compare 'input_feature_iota' with 'expanded_feature_iota' to create a + // diagonal predicate. + return xla::Eq(expanded_feature_iota, input_feature_iota); } // Reshapes a filter of shape [H, W, ..., M, N] to [H, W, ..., 1, M*N]. Used to @@ -210,8 +212,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 @@ -225,6 +227,11 @@ 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)); + // TODO(reedwm): Support explicit padding. + if (attrs.padding == EXPLICIT) { + return errors::Unimplemented( + "XLA does not yet support Conv2D with explicit padding."); + } string data_format; TF_RETURN_IF_ERROR(ctx->GetAttr("data_format", &data_format)); @@ -390,23 +397,31 @@ xla::StatusOr MakeXlaBackpropFilterConvOp( builder->GetShape(activations)); TF_ASSIGN_OR_RETURN(xla::Shape out_backprop_shape, builder->GetShape(gradients)); + xla::XlaOp filter_backprop; + + xla::Shape input_shape = activations_shape; + xla::Shape output_shape = out_backprop_shape; + + TensorShape input_tensor_shape, filter_tensor_shape, output_tensor_shape; + TF_RETURN_IF_ERROR(XLAShapeToTensorShape(filter_shape, &filter_tensor_shape)); + TF_RETURN_IF_ERROR(XLAShapeToTensorShape(input_shape, &input_tensor_shape)); + TF_RETURN_IF_ERROR(XLAShapeToTensorShape(output_shape, &output_tensor_shape)); + const xla::Shape expanded_filter_shape = attrs.depthwise ? ExpandedFilterShapeForDepthwiseConvolution(filter_shape) : filter_shape; - // Reuse dimension computation logic from conv_grad_ops.cc. ConvBackpropDimensions dims; - 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)); - // The filter gradients are computed by a convolution of the input // activations and the output gradients, with some appropriate padding. // See the comment at the top of conv_grad_ops.h for details. - xla::ConvolutionDimensionNumbers dnums; + 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)); + // 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 @@ -418,6 +433,14 @@ xla::StatusOr MakeXlaBackpropFilterConvOp( int n_dim = GetTensorBatchDimIndex(num_dims, attrs.data_format); int c_dim = GetTensorFeatureDimIndex(num_dims, attrs.data_format); + bool use_batch_group_count = + 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); + // Swap n_dim and c_dim in the activations. dnums.set_input_batch_dimension(c_dim); dnums.set_input_feature_dimension(n_dim); @@ -428,19 +451,21 @@ xla::StatusOr MakeXlaBackpropFilterConvOp( dnums.set_kernel_input_feature_dimension(n_dim); dnums.set_kernel_output_feature_dimension(c_dim); - 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 dimension swap below is needed because filter shape is KH,KW,F,DM. + if (use_batch_group_count) { + dnums.set_output_batch_dimension(attrs.num_spatial_dims + 1); + dnums.set_output_feature_dimension(attrs.num_spatial_dims); + } else { + dnums.set_output_batch_dimension(attrs.num_spatial_dims); + dnums.set_output_feature_dimension(attrs.num_spatial_dims + 1); + } // Tensorflow filter shape is [ H, W, ..., inC, outC ]. for (int i = 0; i < attrs.num_spatial_dims; ++i) { dnums.add_output_spatial_dimensions(i); } - dnums.set_output_batch_dimension(attrs.num_spatial_dims); - dnums.set_output_feature_dimension(attrs.num_spatial_dims + 1); - for (int i = 0; i < attrs.num_spatial_dims; ++i) { + for (int64 i = 0; i < attrs.num_spatial_dims; ++i) { int64 dim = GetTensorSpatialDimIndex(num_dims, attrs.data_format, i); dnums.add_input_spatial_dimensions(dim); dnums.add_kernel_spatial_dimensions(dim); @@ -449,7 +474,7 @@ xla::StatusOr MakeXlaBackpropFilterConvOp( // 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]; @@ -494,11 +519,14 @@ xla::StatusOr MakeXlaBackpropFilterConvOp( // // This is done by specifying the window dilation factors in the // convolution HLO below. - auto filter_backprop = - xla::ConvGeneralDilated(activations, gradients, window_strides, padding, - /*lhs_dilation=*/ones, rhs_dilation, dnums); - if (attrs.depthwise) { + filter_backprop = xla::ConvGeneralDilated( + activations, gradients, window_strides, padding, /*lhs_dilation=*/ones, + rhs_dilation, dnums, + /*feature_group_count=*/1, + /*batch_group_count=*/use_batch_group_count ? dims.in_depth : 1); + + if (!use_batch_group_count && attrs.depthwise) { filter_backprop = ContractFilterForDepthwiseBackprop( filter_shape, filter_backprop, activations.builder()); } diff --git a/tensorflow/compiler/tf2xla/kernels/conv_ops.cc b/tensorflow/compiler/tf2xla/kernels/conv_ops.cc index cd7c820be0b6029514ff74288e7bdd3f75b5d6b1..52c3c2c4a903a8c51f6b511774bc0312d39df826 100644 --- a/tensorflow/compiler/tf2xla/kernels/conv_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/conv_ops.cc @@ -22,16 +22,16 @@ 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/constants.h" -#include "tensorflow/compiler/xla/client/lib/numeric.h" +#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" @@ -124,7 +124,7 @@ class Conv2DBackpropInputOp : public ConvBackpropInputOp { : ConvBackpropInputOp(ctx, /*num_spatial_dims=*/2, /*depthwise=*/false) {} }; REGISTER_XLA_OP( - Name("Conv2DBackpropInput").CompileTimeConstInput("input_sizes"), + Name("Conv2DBackpropInput").CompileTimeConstantInput("input_sizes"), Conv2DBackpropInputOp); class Conv3DBackpropInputOp : public ConvBackpropInputOp { @@ -133,7 +133,7 @@ class Conv3DBackpropInputOp : public ConvBackpropInputOp { : ConvBackpropInputOp(ctx, /*num_spatial_dims=*/3, /*depthwise=*/false) {} }; REGISTER_XLA_OP( - Name("Conv3DBackpropInputV2").CompileTimeConstInput("input_sizes"), + Name("Conv3DBackpropInputV2").CompileTimeConstantInput("input_sizes"), Conv3DBackpropInputOp); class DepthwiseConv2DBackpropInputOp : public ConvBackpropInputOp { @@ -142,7 +142,7 @@ class DepthwiseConv2DBackpropInputOp : public ConvBackpropInputOp { : ConvBackpropInputOp(ctx, /*num_spatial_dims=*/2, /*depthwise=*/true) {} }; REGISTER_XLA_OP(Name("DepthwiseConv2dNativeBackpropInput") - .CompileTimeConstInput("input_sizes"), + .CompileTimeConstantInput("input_sizes"), DepthwiseConv2DBackpropInputOp); class ConvBackpropFilterOp : public XlaOpKernel { @@ -183,7 +183,7 @@ class Conv2DBackpropFilterOp : public ConvBackpropFilterOp { } }; REGISTER_XLA_OP( - Name("Conv2DBackpropFilter").CompileTimeConstInput("filter_sizes"), + Name("Conv2DBackpropFilter").CompileTimeConstantInput("filter_sizes"), Conv2DBackpropFilterOp); class Conv3DBackpropFilterOp : public ConvBackpropFilterOp { @@ -193,7 +193,7 @@ class Conv3DBackpropFilterOp : public ConvBackpropFilterOp { } }; REGISTER_XLA_OP( - Name("Conv3DBackpropFilterV2").CompileTimeConstInput("filter_sizes"), + Name("Conv3DBackpropFilterV2").CompileTimeConstantInput("filter_sizes"), Conv3DBackpropFilterOp); class DepthwiseConv2DBackpropFilterOp : public ConvBackpropFilterOp { @@ -202,7 +202,7 @@ class DepthwiseConv2DBackpropFilterOp : public ConvBackpropFilterOp { : ConvBackpropFilterOp(ctx, /*num_spatial_dims=*/2, /*depthwise=*/true) {} }; REGISTER_XLA_OP(Name("DepthwiseConv2dNativeBackpropFilter") - .CompileTimeConstInput("filter_sizes"), + .CompileTimeConstantInput("filter_sizes"), DepthwiseConv2DBackpropFilterOp); } // namespace diff --git a/tensorflow/compiler/tf2xla/kernels/diag_op.cc b/tensorflow/compiler/tf2xla/kernels/diag_op.cc index 49c12fc232092873b69961644a059abc6035f64f..ee79cbc70da269be7586c47b4fd33c901f4fd581 100644 --- a/tensorflow/compiler/tf2xla/kernels/diag_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/diag_op.cc @@ -19,7 +19,7 @@ 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/constants.h" -#include "tensorflow/compiler/xla/client/lib/numeric.h" +#include "tensorflow/compiler/xla/client/lib/matrix.h" #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/compiler/xla/util.h" #include "tensorflow/core/framework/op_kernel.h" diff --git a/tensorflow/compiler/tf2xla/kernels/dynamic_slice_ops.cc b/tensorflow/compiler/tf2xla/kernels/dynamic_slice_ops.cc index 4af1e8b44cbbd02d8e3ea5e42d841c92288b5d56..bb2c0d9ddb8504a1156a74b6ece5d41b620803c7 100644 --- a/tensorflow/compiler/tf2xla/kernels/dynamic_slice_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/dynamic_slice_ops.cc @@ -102,8 +102,9 @@ class DynamicSliceOp : public XlaOpKernel { } }; -REGISTER_XLA_OP(Name("XlaDynamicSlice").CompileTimeConstInput("size_indices"), - DynamicSliceOp); +REGISTER_XLA_OP( + Name("XlaDynamicSlice").CompileTimeConstantInput("size_indices"), + DynamicSliceOp); } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/dynamic_stitch_op.cc b/tensorflow/compiler/tf2xla/kernels/dynamic_stitch_op.cc index cb73053666d4c32bc0a2ef19b174aee1a29f101e..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 { @@ -113,8 +113,20 @@ class DynamicStitchOp : public XlaOpKernel { } } int number_of_indices = max_index + 1; - OP_REQUIRES(ctx, number_of_indices > 0, - errors::InvalidArgument("no indices supplied")); + int64 result_rank = 1 + data0_shape.dims() - indices0_shape.dims(); + if (number_of_indices == 0) { + std::vector result_shape(result_rank); + for (int d = indices0_shape.dims(); d < data0_shape.dims(); d++) { + result_shape[d - indices0_shape.dims() + 1] = data0_shape.dim_size(d); + } + xla::PrimitiveType element_type = + ctx->input_xla_type(ctx->num_inputs() - 1); + xla::Literal empty_literal = xla::Literal::CreateFromShape( + xla::ShapeUtil::MakeShape(element_type, result_shape)); + ctx->SetOutput(0, xla::ConstantLiteral(ctx->builder(), empty_literal)); + return; + } + // Construct the reverse mapping, for each index, of which slice of which // input it comes from. std::vector src_input_vector(number_of_indices); @@ -157,12 +169,9 @@ class DynamicStitchOp : public XlaOpKernel { // Set up the vectors for slicing: the first dimension will vary // slice by slice, and the rest take the full common extra shape. - std::vector slice_start(1 + data0_shape.dims() - - indices0_shape.dims()); - std::vector slice_limit(1 + data0_shape.dims() - - indices0_shape.dims()); - std::vector stride(1 + data0_shape.dims() - indices0_shape.dims(), - 1); + std::vector slice_start(result_rank); + std::vector slice_limit(result_rank); + std::vector stride(result_rank, 1); for (int d = indices0_shape.dims(); d < data0_shape.dims(); d++) { slice_limit[1 + d - indices0_shape.dims()] = data0_shape.dim_size(d); } @@ -200,10 +209,11 @@ class DynamicStitchOp : public XlaOpKernel { } }; -REGISTER_XLA_OP(Name("DynamicStitch").CompileTimeConstInput("indices"), - DynamicStitchOp); -REGISTER_XLA_OP(Name("ParallelDynamicStitch").CompileTimeConstInput("indices"), +REGISTER_XLA_OP(Name("DynamicStitch").CompileTimeConstantInput("indices"), DynamicStitchOp); +REGISTER_XLA_OP( + Name("ParallelDynamicStitch").CompileTimeConstantInput("indices"), + DynamicStitchOp); } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/extract_image_patches_op.cc b/tensorflow/compiler/tf2xla/kernels/extract_image_patches_op.cc index c68b0bfd7961892294c2931e5c4c44de534a7740..29687c7b82f92d9f336854c4575746589c63b64f 100644 --- a/tensorflow/compiler/tf2xla/kernels/extract_image_patches_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/extract_image_patches_op.cc @@ -17,7 +17,6 @@ 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/numeric.h" #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/core/util/tensor_format.h" diff --git a/tensorflow/compiler/tf2xla/kernels/fake_quantize_ops.cc b/tensorflow/compiler/tf2xla/kernels/fake_quantize_ops.cc index cdba6680dee3fade5bdf0c453ed672b653072b0d..142be030f737f105980ab9c80a5a849e1ca6eb47 100644 --- a/tensorflow/compiler/tf2xla/kernels/fake_quantize_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/fake_quantize_ops.cc @@ -260,19 +260,19 @@ class FakeQuantWithMinMaxVarsGradOp : public XlaOpKernel { xla::XlaOp below_min = xla::Lt(input, nudged_input_min); xla::XlaOp select1 = xla::Select(below_min, gradient, zeroes); xla::XlaOp reduce1 = xla::ReduceAll( - XlaHelpers::ConvertElementType(b, select1, accumulation_type), + XlaHelpers::ConvertElementType(select1, accumulation_type), XlaHelpers::Zero(b, accumulation_type), *ctx->GetOrCreateAdd(accumulation_type)); - xla::XlaOp output1 = XlaHelpers::ConvertElementType(b, reduce1, data_type); + xla::XlaOp output1 = XlaHelpers::ConvertElementType(reduce1, data_type); ctx->SetOutput(1, output1); xla::XlaOp above_max = xla::Gt(input, nudged_input_max); xla::XlaOp select2 = xla::Select(above_max, gradient, zeroes); xla::XlaOp reduce2 = xla::ReduceAll( - XlaHelpers::ConvertElementType(b, select2, accumulation_type), + XlaHelpers::ConvertElementType(select2, accumulation_type), XlaHelpers::Zero(b, accumulation_type), *ctx->GetOrCreateAdd(accumulation_type)); - xla::XlaOp output2 = XlaHelpers::ConvertElementType(b, reduce2, data_type); + xla::XlaOp output2 = XlaHelpers::ConvertElementType(reduce2, data_type); ctx->SetOutput(2, output2); } diff --git a/tensorflow/compiler/tf2xla/kernels/fft_ops.cc b/tensorflow/compiler/tf2xla/kernels/fft_ops.cc index 80bcef966360ec9a1ca63a02741108ce41b31846..a623585aad3b1b8f1f096ca527e7694d74f1ba46 100644 --- a/tensorflow/compiler/tf2xla/kernels/fft_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/fft_ops.cc @@ -20,12 +20,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/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" @@ -50,11 +51,36 @@ class GenericFftOp : public XlaOpKernel { errors::InvalidArgument("input must be at least 1 dimensional")); std::vector fft_length; + xla::XlaOp input = ctx->Input(0); if (fft_type_ == FftType::RFFT || fft_type_ == FftType::IRFFT) { OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector(1, &fft_length)); OP_REQUIRES(ctx, fft_length.size() == fft_rank_, errors::InvalidArgument("fft_length must be length ", fft_rank_, " vector")); + + // Zero pad or truncate the axes we're doing FFT on. + absl::InlinedVector slice_sizes = input_shape.dim_sizes(); + std::vector> padding_sizes(slice_sizes.size()); + std::vector expected_sizes = fft_length; + // IRFFT wants the innermost axis to be n / 2 + 1. + if (fft_type_ == FftType::IRFFT) { + expected_sizes[fft_rank_ - 1] = fft_length[fft_rank_ - 1] / 2 + 1; + } + for (int i = 0; i < fft_rank_; i++) { + int index = input_shape.dims() - fft_rank_ + i; + if (input_shape.dim_size(index) > expected_sizes[i]) { + slice_sizes[index] = expected_sizes[i]; + } else { + padding_sizes[index].second = + expected_sizes[i] - input_shape.dim_size(index); + } + } + + std::vector start_indices(input_shape.dims(), 0); + std::vector strides(input_shape.dims(), 1); + input = xla::Pad(xla::Slice(input, start_indices, slice_sizes, strides), + XlaHelpers::Zero(ctx->builder(), ctx->input_type(0)), + xla::MakeEdgePaddingConfig(padding_sizes)); } else { // Innermost axis provides the FFT length. for (int i = 0; i < fft_rank_; i++) { @@ -63,7 +89,7 @@ class GenericFftOp : public XlaOpKernel { } } - xla::XlaOp fft = xla::Fft(ctx->Input(0), fft_type_, fft_length); + xla::XlaOp fft = xla::Fft(input, fft_type_, fft_length); ctx->SetOutput(0, fft); } @@ -106,9 +132,11 @@ class RFFTOp : public GenericFftOp { explicit RFFTOp(OpKernelConstruction* ctx) : GenericFftOp(ctx, /*fft_type=*/FftType::RFFT, /*fft_rank=*/FFTRank) {} }; -REGISTER_XLA_OP(Name("RFFT").CompileTimeConstInput("fft_length"), RFFTOp<1>); -REGISTER_XLA_OP(Name("RFFT2D").CompileTimeConstInput("fft_length"), RFFTOp<2>); -REGISTER_XLA_OP(Name("RFFT3D").CompileTimeConstInput("fft_length"), RFFTOp<3>); +REGISTER_XLA_OP(Name("RFFT").CompileTimeConstantInput("fft_length"), RFFTOp<1>); +REGISTER_XLA_OP(Name("RFFT2D").CompileTimeConstantInput("fft_length"), + RFFTOp<2>); +REGISTER_XLA_OP(Name("RFFT3D").CompileTimeConstantInput("fft_length"), + RFFTOp<3>); template class IRFFTOp : public GenericFftOp { @@ -116,10 +144,11 @@ class IRFFTOp : public GenericFftOp { explicit IRFFTOp(OpKernelConstruction* ctx) : GenericFftOp(ctx, /*fft_type=*/FftType::IRFFT, /*fft_rank=*/FFTRank) {} }; -REGISTER_XLA_OP(Name("IRFFT").CompileTimeConstInput("fft_length"), IRFFTOp<1>); -REGISTER_XLA_OP(Name("IRFFT2D").CompileTimeConstInput("fft_length"), +REGISTER_XLA_OP(Name("IRFFT").CompileTimeConstantInput("fft_length"), + IRFFTOp<1>); +REGISTER_XLA_OP(Name("IRFFT2D").CompileTimeConstantInput("fft_length"), IRFFTOp<2>); -REGISTER_XLA_OP(Name("IRFFT3D").CompileTimeConstInput("fft_length"), +REGISTER_XLA_OP(Name("IRFFT3D").CompileTimeConstantInput("fft_length"), IRFFTOp<3>); } // namespace diff --git a/tensorflow/compiler/tf2xla/kernels/fill_op.cc b/tensorflow/compiler/tf2xla/kernels/fill_op.cc index 54b21a278229024e3e54e9135548be6b69b077e1..35e0625dbb0d4c696d36cce642d6f50f1d220c45 100644 --- a/tensorflow/compiler/tf2xla/kernels/fill_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/fill_op.cc @@ -22,6 +22,7 @@ limitations under the License. #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 { @@ -33,44 +34,25 @@ class FillOp : public XlaOpKernel { void Compile(XlaOpKernelContext* ctx) override { // The output of this Op is a tensor of shape 'dims_shape' with each // element set to the scalar 'dims_literal'. - const TensorShape dims_shape = ctx->InputShape(0); - const TensorShape value_shape = ctx->InputShape(1); + const TensorShape dims_shape = ctx->InputShape("dims"); + const TensorShape value_shape = ctx->InputShape("value"); OP_REQUIRES( - ctx, IsLegacyVector(dims_shape), + ctx, TensorShapeUtils::IsVector(dims_shape), errors::InvalidArgument("dims must be a vector of int32, got shape ", dims_shape.DebugString())); - OP_REQUIRES(ctx, IsLegacyScalar(value_shape), + OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(value_shape), errors::InvalidArgument("value must be a scalar, got shape ", value_shape.DebugString())); - // Evaluate the 'dims' constant input, reshaping to a vector if it - // was a 'legacy' vector (secretly a scalar). - xla::Literal dims_literal; - OP_REQUIRES_OK(ctx, ctx->ConstantInputReshaped( - 0, {dims_shape.num_elements()}, &dims_literal)); - // Convert the dims literal into a vector that we can pass to - // XlaBuilder. - std::vector broadcast; - broadcast.reserve(dims_literal.shape().dimensions(0)); - for (int i = 0; i < dims_literal.shape().dimensions(0); ++i) { - broadcast.push_back(dims_literal.Get({i})); - } - // Look up the value input, reshaping to a scalar if it was a - // 'legacy' scalar (secretly a vector). - xla::XlaOp data = ctx->Input(1); - if (value_shape.dims() > 0) { - CHECK_EQ(value_shape.dims(), 1); - data = xla::Reshape(data, {}); - } - // Emit the actual computation, which broadcasts the scalar to the - // desired shape. - auto result = xla::Broadcast(data, broadcast); + std::vector dims; + OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector("dims", &dims)); + auto result = xla::Broadcast(ctx->Input("value"), dims); ctx->SetOutput(0, result); } }; -REGISTER_XLA_OP(Name("Fill").CompileTimeConstInput("dims"), FillOp); +REGISTER_XLA_OP(Name("Fill").CompileTimeConstantInput("dims"), FillOp); } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/gather_op.cc b/tensorflow/compiler/tf2xla/kernels/gather_op.cc index 44140304fdf5cdf60d8ad8b85c532fcadff8ba86..6472045265e4d930a5da770a68f5c502192201ae 100644 --- a/tensorflow/compiler/tf2xla/kernels/gather_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/gather_op.cc @@ -14,7 +14,6 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/compiler/tf2xla/kernels/gather_op_helpers.h" -#include "tensorflow/compiler/tf2xla/lib/while_loop.h" #include "tensorflow/compiler/tf2xla/shape_util.h" #include "tensorflow/compiler/tf2xla/type_util.h" #include "tensorflow/compiler/tf2xla/xla_context.h" @@ -168,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); @@ -194,7 +193,7 @@ class GatherOp : public XlaOpKernel { }; REGISTER_XLA_OP(Name("Gather"), GatherOp); -REGISTER_XLA_OP(Name("GatherV2").CompileTimeConstInput("axis"), GatherOp); +REGISTER_XLA_OP(Name("GatherV2").CompileTimeConstantInput("axis"), GatherOp); class GatherNdOp : public XlaOpKernel { public: diff --git a/tensorflow/compiler/tf2xla/kernels/if_op.cc b/tensorflow/compiler/tf2xla/kernels/if_op.cc index 56da50f140893c68c8a1556853884720b21c7229..aa5637e2669555da17af8bb05ab08beeba6a89c3 100644 --- a/tensorflow/compiler/tf2xla/kernels/if_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/if_op.cc @@ -56,6 +56,7 @@ void XlaIfOp::Compile(XlaOpKernelContext* ctx) { VLOG(1) << "Building If: " << input_types_.size() << " inputs"; std::vector arguments(input_types_.size()); + int num_resource_args = 0; for (int i = 0; i < input_types_.size(); ++i) { XlaCompiler::Argument& arg = arguments[i]; DataType type = ctx->input_type(i + 1); @@ -72,21 +73,23 @@ void XlaIfOp::Compile(XlaOpKernelContext* ctx) { arg.shape = resource->shape(); OP_REQUIRES(ctx, arg.initialized, errors::Unimplemented("Uninitialized arguments: ", arg.name)); - arg.tensor_array_size = resource->tensor_array_size(); + arg.max_array_size = resource->max_array_size(); for (const auto& gradient : resource->tensor_array_gradients()) { arg.tensor_array_gradients.insert(gradient.first); } arg.name = resource->name(); VLOG(2) << "Resource " << resource->name() << " type: " << DataTypeString(arg.type) - << " shape: " << arg.shape.DebugString() + << " shape: " << arg.HumanString() << " initialized: " << arg.initialized; + + num_resource_args++; } else { arg.kind = XlaCompiler::Argument::kParameter; arg.type = input_types_[i]; arg.shape = ctx->InputShape(i + 1); VLOG(2) << "Arg type: " << DataTypeString(arg.type) - << " shape: " << arg.shape.DebugString(); + << " shape: " << arg.HumanString(); } } @@ -147,12 +150,12 @@ void XlaIfOp::Compile(XlaOpKernelContext* ctx) { OP_REQUIRES(ctx, then_result.xla_input_shapes.size() == 1, errors::FailedPrecondition("Expected one input shape")); xla::Shape then_input_shape = then_result.xla_input_shapes[0]; - OP_REQUIRES(ctx, xla::ShapeUtil::IsTuple(then_input_shape), + OP_REQUIRES(ctx, then_input_shape.IsTuple(), errors::FailedPrecondition("Expected tuple shape")); OP_REQUIRES(ctx, else_result.xla_input_shapes.size() == 1, errors::FailedPrecondition("Expected one input shape")); xla::Shape else_input_shape = else_result.xla_input_shapes[0]; - OP_REQUIRES(ctx, xla::ShapeUtil::IsTuple(else_input_shape), + OP_REQUIRES(ctx, else_input_shape.IsTuple(), errors::FailedPrecondition("Expected tuple shape")); OP_REQUIRES(ctx, xla::ShapeUtil::Compatible(then_input_shape, else_input_shape), @@ -236,12 +239,16 @@ void XlaIfOp::Compile(XlaOpKernelContext* ctx) { ctx->SetOutput(i, output_handle); } if (has_token_input_output_) { - // Set token output for this "if" op. + // Set token output for this "If" op. Token output is the last output of + // XLA computation, which comes after all "normal" TF outputs and resource + // updates. For "If" node, num of resource updates equals to number of + // resource args because we set `return_updated_values_for_all_resources` + // to true in XlaCompiler option. xla::XlaOp token_output = - xla::GetTupleElement(outputs, output_types_.size()); + xla::GetTupleElement(outputs, output_types_.size() + num_resource_args); auto shape_or = b->GetShape(token_output); OP_REQUIRES_OK(ctx, shape_or.status()); - OP_REQUIRES(ctx, xla::ShapeUtil::IsToken(shape_or.ValueOrDie()), + OP_REQUIRES(ctx, shape_or.ValueOrDie().IsToken(), errors::FailedPrecondition( "Token output is not token type: ", xla::ShapeUtil::HumanString(shape_or.ValueOrDie()))); diff --git a/tensorflow/compiler/tf2xla/kernels/image_ops.cc b/tensorflow/compiler/tf2xla/kernels/image_ops.cc index 6713d6bc921b24b25baddfb3fd7296fffcc3d6ea..92b20fe0ba5611ca5314cd954026f7b71ea75f84 100644 --- a/tensorflow/compiler/tf2xla/kernels/image_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/image_ops.cc @@ -13,18 +13,20 @@ 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/lib/while_loop.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/client/lib/constants.h" +#include "tensorflow/compiler/xla/client/lib/loops.h" #include "tensorflow/compiler/xla/client/lib/sorting.h" #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,20 +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(b, input, accumulation_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(b, 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); @@ -234,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); @@ -250,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); @@ -284,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); @@ -300,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); @@ -352,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); @@ -506,9 +518,9 @@ class NonMaxSuppressionOp : public XlaOpKernel { init_values.push_back(included_iou); auto suppress_loop_result = - XlaWhileLoop(WhileCondFn(num_boxes, output_size), - SuppressBodyFn(num_boxes), init_values, "suppress_loop", - builder) + xla::WhileLoopHelper(WhileCondFn(num_boxes, output_size), + SuppressBodyFn(num_boxes), init_values, + "suppress_loop", builder) .ValueOrDie(); xla::XlaOp included_score = @@ -563,7 +575,7 @@ class NonMaxSuppressionOp : public XlaOpKernel { }; REGISTER_XLA_OP( - Name("NonMaxSuppressionV4").CompileTimeConstInput("max_output_size"), + Name("NonMaxSuppressionV4").CompileTimeConstantInput("max_output_size"), NonMaxSuppressionOp); } // namespace diff --git a/tensorflow/compiler/tf2xla/kernels/image_resize_ops.cc b/tensorflow/compiler/tf2xla/kernels/image_resize_ops.cc index 7b2bb4a7c50fc954237e09a32f71009f790b60d0..b96d45316f626e678a64392a4315979eeeb6e83c 100644 --- a/tensorflow/compiler/tf2xla/kernels/image_resize_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/image_resize_ops.cc @@ -19,7 +19,6 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "tensorflow/compiler/xla/array4d.h" #include "tensorflow/compiler/xla/client/lib/constants.h" -#include "tensorflow/compiler/xla/client/lib/numeric.h" #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/core/framework/kernel_def_builder.h" #include "tensorflow/core/framework/register_types.h" @@ -73,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, @@ -118,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: @@ -133,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; @@ -143,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 @@ -231,20 +253,22 @@ xla::XlaOp ResizeUsingDilationAndConvolution(xla::XlaBuilder* builder, num_extended[0] = upper_padding[0] / (dims.kernel_size[0]); num_extended[1] = upper_padding[1] / (dims.kernel_size[1]); + const int64 batch_dim_size = + builder->GetShape(input).ValueOrDie().dimensions(0); if (num_extended[0] > 0) { - auto slice = - xla::Slice(input_data, {0, in_size[0] - 1, 0, 0}, - {1, in_size[0], in_size[1], channels}, {1, 1, 1, 1}); + auto slice = xla::Slice( + input_data, {0, in_size[0] - 1, 0, 0}, + {batch_dim_size, in_size[0], in_size[1], channels}, {1, 1, 1, 1}); for (int i = 0; i < num_extended[0]; i++) { input_data = xla::ConcatInDim(builder, {input_data, slice}, 1); } } if (num_extended[1] > 0) { - auto slice = - xla::Slice(input_data, {0, 0, in_size[1] - 1, 0}, - {1, in_size[0] + num_extended[0], in_size[1], channels}, - {1, 1, 1, 1}); + auto slice = xla::Slice( + input_data, {0, 0, in_size[1] - 1, 0}, + {batch_dim_size, in_size[0] + num_extended[0], in_size[1], channels}, + {1, 1, 1, 1}); for (int i = 0; i < num_extended[1]; i++) { input_data = xla::ConcatInDim(builder, {input_data, slice}, 2); } @@ -263,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=*/ @@ -274,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=*/ @@ -283,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=*/ @@ -305,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); @@ -331,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 @@ -354,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), @@ -406,112 +428,142 @@ 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").CompileTimeConstInput("size"), +REGISTER_XLA_OP(Name("ResizeBilinear").CompileTimeConstantInput("size"), ResizeBilinearOp); class ResizeBilinearGradOp : public XlaOpKernel { @@ -580,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 f3964748587c1b31cf8b1b76643ff19a9044bf44..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); @@ -78,7 +77,7 @@ XlaArgMaxOp::XlaArgMaxOp(OpKernelConstruction* ctx) : XlaArgMinMaxOp(ctx, /*is_min=*/false) {} REGISTER_XLA_OP(Name("ArgMax") .Device(DEVICE_GPU_XLA_JIT) - .CompileTimeConstInput("dimension"), + .CompileTimeConstantInput("dimension"), XlaArgMaxOp); namespace { @@ -89,7 +88,8 @@ class XlaArgMinOp : public XlaArgMinMaxOp { }; XlaArgMinOp::XlaArgMinOp(OpKernelConstruction* ctx) : XlaArgMinMaxOp(ctx, /*is_min=*/true) {} -REGISTER_XLA_OP(Name("ArgMin").CompileTimeConstInput("dimension"), XlaArgMinOp); +REGISTER_XLA_OP(Name("ArgMin").CompileTimeConstantInput("dimension"), + XlaArgMinOp); } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/index_ops_cpu.cc b/tensorflow/compiler/tf2xla/kernels/index_ops_cpu.cc index f210bfbd886e48b8d7972393ed1899491486646c..e4bbdef6480104a1051acfc647644deb65c80171 100644 --- a/tensorflow/compiler/tf2xla/kernels/index_ops_cpu.cc +++ b/tensorflow/compiler/tf2xla/kernels/index_ops_cpu.cc @@ -16,21 +16,23 @@ 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 { -// The logic below uses a custom-call to implement argmax. +// The logic below uses a custom-call to implement argmax when possible. When +// custom-call is not allowed or input shapes are not supported, this kernel +// falls back to using XLA HLO native ArgMax. // // Also see b/29507024 for first-class XLA support for indexing ops. class ArgMaxCustomCallOp : public XlaOpKernel { @@ -48,30 +50,42 @@ class ArgMaxCustomCallOp : public XlaOpKernel { // We require that the dimension argument is a constant, since it lets us // dispatch to a specialized custom-call function without any run-time // overhead, when compiling ahead-of-time. - xla::Literal literal; - OP_REQUIRES_OK(ctx, ctx->ConstantInput(1, &literal)); - const int32 dim = literal.Get({}); - OP_REQUIRES(ctx, dim >= 0, errors::InvalidArgument("dim must be >= 0")); - OP_REQUIRES( - ctx, dim < input_shape.dims(), - errors::InvalidArgument("dim must be < input rank (", - input_shape.dims(), "), but got: ", dim)); - const int64 dim_size = input_shape.dim_size(dim); - OP_REQUIRES(ctx, dim_size > 0, + int64 dim; + OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntScalar(1, &dim)); + + const int input_dims = input_shape.dims(); + const int axis = dim < 0 ? dim + input_dims : dim; + OP_REQUIRES(ctx, axis >= 0 && axis < input_dims, + errors::InvalidArgument("Expected dimension in the range [", + -input_dims, ", ", input_dims, + "), but got ", dim)); + + const int64 axis_size = input_shape.dim_size(axis); + OP_REQUIRES(ctx, axis_size > 0, errors::InvalidArgument( "Reduction axis ", dim, " is empty in shape: ", input_shape.DebugString())); - // The output shape is the input shape contracted along dim. + const DataType dtype = output_type(0); + xla::PrimitiveType output_type; + OP_REQUIRES_OK(ctx, DataTypeToPrimitiveType(dtype, &output_type)); + + // Fall back to XLA ArgMax HLO when CustomCall is not allowed or when input + // shape isn't supported. + if (!ctx->compiler()->options().allow_cpu_custom_calls || + (input_dims != 1 && input_dims != 2)) { + xla::XlaOp output = xla::ArgMax(ctx->Input(0), output_type, axis); + ctx->SetOutput(0, output); + return; + } + + xla::XlaOp output; + // The output shape is the input shape contracted along axis. TensorShape output_shape; for (int d = 0; d < input_shape.dims() - 1; ++d) { - output_shape.AddDim(input_shape.dim_size((d < dim) ? d : d + 1)); + output_shape.AddDim(input_shape.dim_size((d < axis) ? d : d + 1)); } - // For now we use a custom-call, only for the 1d and 2d cases. - OP_REQUIRES(ctx, XlaContext::Get(ctx).allow_cpu_custom_calls(), - errors::InvalidArgument( - "ArgMax implementation requires a CustomCall on CPU")); xla::XlaBuilder& b = *ctx->builder(); // XLA passes to the function, so it is not included here. @@ -85,7 +99,7 @@ class ArgMaxCustomCallOp : public XlaOpKernel { args.push_back(xla::ConstantLiteral( &b, xla::LiteralUtil::CreateR1(output_shape.dim_sizes()))); args.push_back( - xla::ConstantLiteral(&b, xla::LiteralUtil::CreateR0(dim))); + xla::ConstantLiteral(&b, xla::LiteralUtil::CreateR0(axis))); } // The argmax function expects row-major layout. @@ -96,30 +110,21 @@ class ArgMaxCustomCallOp : public XlaOpKernel { auto shape_status = b.GetShape(arg); OP_REQUIRES_OK(ctx, shape_status.status()); xla::Shape arg_shape = shape_status.ConsumeValueOrDie(); - *arg_shape.mutable_layout() = xla::LayoutUtil::MakeDescendingLayout( - xla::ShapeUtil::Rank(arg_shape)); + *arg_shape.mutable_layout() = + xla::LayoutUtil::MakeDescendingLayout(arg_shape.rank()); arg_shapes.push_back(std::move(arg_shape)); } // Tell XLA to call the custom code, defined in - // index_ops_kernel_argmax_float_1d.cc. - xla::XlaOp output; - switch (input_shape.dims()) { - case 1: - output = xla::CustomCallWithLayout(&b, "argmax_float_1d_xla_impl", args, - xla_shape, arg_shapes); - break; - case 2: - output = xla::CustomCallWithLayout(&b, "argmax_float_2d_xla_impl", args, - xla_shape, arg_shapes); - break; - default: - OP_REQUIRES(ctx, false, - errors::Unimplemented( - "Argmax is only implemented for 1d and 2d tensors" - ", but got shape: ", - input_shape.DebugString())); + // index_ops_kernel_argmax_float_{1, 2}d.cc. + if (input_dims == 1) { + output = xla::CustomCallWithLayout(&b, "argmax_float_1d_xla_impl", args, + xla_shape, arg_shapes); + } else { + output = xla::CustomCallWithLayout(&b, "argmax_float_2d_xla_impl", args, + xla_shape, arg_shapes); } + output = xla::ConvertElementType(output, output_type); ctx->SetOutput(0, output); } @@ -130,7 +135,7 @@ class ArgMaxCustomCallOp : public XlaOpKernel { REGISTER_XLA_OP(Name("ArgMax") .TypeConstraint("T", DT_FLOAT) .Device(DEVICE_CPU_XLA_JIT) - .CompileTimeConstInput("dimension"), + .CompileTimeConstantInput("dimension"), ArgMaxCustomCallOp); } // namespace diff --git a/tensorflow/compiler/tf2xla/kernels/l2loss_op.cc b/tensorflow/compiler/tf2xla/kernels/l2loss_op.cc index f028e361bccd51de0bd69a1d2227c7afaed53455..93f029731c34e84000a3dc00df8af05654cccf2d 100644 --- a/tensorflow/compiler/tf2xla/kernels/l2loss_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/l2loss_op.cc @@ -37,12 +37,11 @@ class L2LossOp : public XlaOpKernel { // output = sum(t ** 2) / 2 const DataType accumulation_type = XlaHelpers::SumAccumulationType(dtype); - auto t = - XlaHelpers::ConvertElementType(b, ctx->Input(0), accumulation_type); + auto t = XlaHelpers::ConvertElementType(ctx->Input(0), accumulation_type); auto square = xla::Mul(t, t); auto reduce = xla::Reduce(square, XlaHelpers::Zero(b, accumulation_type), *ctx->GetOrCreateAdd(accumulation_type), dims); - auto deconverted = XlaHelpers::ConvertElementType(b, reduce, dtype); + auto deconverted = XlaHelpers::ConvertElementType(reduce, dtype); auto two = XlaHelpers::IntegerLiteral(b, dtype, 2); ctx->SetOutput(0, xla::Div(deconverted, two)); } diff --git a/tensorflow/compiler/tf2xla/kernels/listdiff_op.cc b/tensorflow/compiler/tf2xla/kernels/listdiff_op.cc index a11bbe918f7f8eb050aaa40d4344f9cc9e9a10a4..e46f4e72dc9cb245916b138d5365ee42371f0e4c 100644 --- a/tensorflow/compiler/tf2xla/kernels/listdiff_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/listdiff_op.cc @@ -115,8 +115,8 @@ class ListDiffOp : public XlaOpKernel { REGISTER_XLA_OP(Name("ListDiff") .TypeConstraint("T", kListDiffTypes) - .CompileTimeConstInput("x") - .CompileTimeConstInput("y"), + .CompileTimeConstantInput("x") + .CompileTimeConstantInput("y"), ListDiffOp); } // namespace diff --git a/tensorflow/compiler/tf2xla/kernels/lrn_ops.cc b/tensorflow/compiler/tf2xla/kernels/lrn_ops.cc index 87ee2d3aede50eb24e65570f106d49030e1d4236..987901d82b3f3798dd52f18ef2497b8f0cf80b11 100644 --- a/tensorflow/compiler/tf2xla/kernels/lrn_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/lrn_ops.cc @@ -49,16 +49,14 @@ class LRNOp : public XlaOpKernel { // We use a window of depth_radius_ * 2 + 1, to account for the current // element and a depth_radius_ on either side. auto accumulation_type = XlaHelpers::SumAccumulationType(input_type(0)); - auto converted = - XlaHelpers::ConvertElementType(builder, input, accumulation_type); + auto converted = XlaHelpers::ConvertElementType(input, accumulation_type); auto squared = xla::Mul(converted, converted); auto reduce = xla::ReduceWindow( squared, XlaHelpers::Zero(builder, accumulation_type), *ctx->GetOrCreateAdd(accumulation_type), /* window_dimensions = */ {1, 1, 1, depth_radius_ * 2 + 1}, /* window_strides = */ {1, 1, 1, 1}, xla::Padding::kSame); - auto sqr_sum = - XlaHelpers::ConvertElementType(builder, reduce, input_type(0)); + auto sqr_sum = XlaHelpers::ConvertElementType(reduce, input_type(0)); auto scale = xla::Pow( xla::Add(xla::ConstantR0(builder, bias_), @@ -138,15 +136,14 @@ class LRNGradOp : public XlaOpKernel { auto accumulation_type = XlaHelpers::SumAccumulationType(input_type(0)); auto converted = - XlaHelpers::ConvertElementType(builder, in_image, accumulation_type); + XlaHelpers::ConvertElementType(in_image, accumulation_type); auto squared = xla::Mul(converted, converted); auto reduce = xla::ReduceWindow( squared, XlaHelpers::Zero(builder, accumulation_type), *ctx->GetOrCreateAdd(accumulation_type), /* window_dimensions = */ {1, 1, 1, depth_radius_ * 2 + 1}, /* window_strides = */ {1, 1, 1, 1}, xla::Padding::kSame); - auto sqr_sum = - XlaHelpers::ConvertElementType(builder, reduce, input_type(0)); + auto sqr_sum = XlaHelpers::ConvertElementType(reduce, input_type(0)); auto norm = xla::Add(xla::ConstantR0(builder, bias_), @@ -157,15 +154,13 @@ class LRNGradOp : public XlaOpKernel { xla::Div(out_image, norm)), in_grads); - auto converted_dy = - XlaHelpers::ConvertElementType(builder, dy, accumulation_type); + auto converted_dy = XlaHelpers::ConvertElementType(dy, accumulation_type); auto dy_reduce = xla::ReduceWindow( converted_dy, XlaHelpers::Zero(builder, accumulation_type), *ctx->GetOrCreateAdd(accumulation_type), /* window_dimensions = */ {1, 1, 1, depth_radius_ * 2 + 1}, /* window_strides = */ {1, 1, 1, 1}, xla::Padding::kSame); - auto dy_reduced = - XlaHelpers::ConvertElementType(builder, dy_reduce, input_type(0)); + auto dy_reduced = XlaHelpers::ConvertElementType(dy_reduce, input_type(0)); xla::XlaOp gradients = xla::Add( xla::Mul(in_image, dy_reduced), 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_band_part_op.cc b/tensorflow/compiler/tf2xla/kernels/matrix_band_part_op.cc index 8dfd7de591c4a3c4768dd60b41e03d294ad49397..2dd0a710e47ec8cad6153402fdb3be59f5868205 100644 --- a/tensorflow/compiler/tf2xla/kernels/matrix_band_part_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/matrix_band_part_op.cc @@ -16,8 +16,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/numeric.h" #include "tensorflow/compiler/xla/client/xla_builder.h" +#include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/core/framework/tensor_shape.h" namespace tensorflow { @@ -61,11 +61,11 @@ class MatrixBandPartOp : public XlaOpKernel { // Compute 'offset', which is how many diagonals we are above/below the // diagonal. - xla::XlaOp iota_m = xla::Iota(builder, index_xla_type, m); - xla::XlaOp iota_n = xla::Iota(builder, index_xla_type, n); + xla::Shape iota_shape = xla::ShapeUtil::MakeShape(index_xla_type, {m, n}); + xla::XlaOp iota_m = xla::Iota(builder, iota_shape, /*iota_dimension=*/0); + xla::XlaOp iota_n = xla::Iota(builder, iota_shape, /*iota_dimension=*/1); - auto offset = xla::Sub(xla::Broadcast(iota_n, {m}), iota_m, - /*broadcast_dimensions=*/{0}); + auto offset = xla::Sub(iota_n, iota_m); // If num_lower or num_upper are negative, include all lower/upper // diagonals. diff --git a/tensorflow/compiler/tf2xla/kernels/matrix_set_diag_op.cc b/tensorflow/compiler/tf2xla/kernels/matrix_set_diag_op.cc index c0ca881ff82cee04e0c5e35f9a2d5732fabdd8a6..4f980b6d14ed667bdf4756ed740894098cae5919 100644 --- a/tensorflow/compiler/tf2xla/kernels/matrix_set_diag_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/matrix_set_diag_op.cc @@ -16,7 +16,6 @@ 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/numeric.h" #include "tensorflow/compiler/xla/client/xla_builder.h" namespace tensorflow { diff --git a/tensorflow/compiler/tf2xla/kernels/matrix_triangular_solve_op.cc b/tensorflow/compiler/tf2xla/kernels/matrix_triangular_solve_op.cc index f4def11d08c31513aec5aad15187016a7294c2fd..90c0ebefb24ec2c4378782e9b15d3f57c33032a4 100644 --- a/tensorflow/compiler/tf2xla/kernels/matrix_triangular_solve_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/matrix_triangular_solve_op.cc @@ -13,9 +13,9 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/compiler/tf2xla/lib/triangular_solve.h" #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" namespace tensorflow { namespace { @@ -29,7 +29,7 @@ class MatrixTriangularSolveOp : public XlaOpKernel { } void Compile(XlaOpKernelContext* ctx) override { - auto result = TriangularSolve( + auto result = xla::TriangularSolve( ctx->Input(0), ctx->Input(1), /*left_side=*/true, /*lower=*/lower_, /*transpose_a=*/adjoint_, /*conjugate_a=*/adjoint_); ctx->SetOutput(0, result); diff --git a/tensorflow/compiler/tf2xla/kernels/mirror_pad_op.cc b/tensorflow/compiler/tf2xla/kernels/mirror_pad_op.cc index 2a42eeaf76ab3aa88ff3a93ef7eb7ab217964bb6..656f9b898f32dfc05215014f51c2bbaf07580836 100644 --- a/tensorflow/compiler/tf2xla/kernels/mirror_pad_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/mirror_pad_op.cc @@ -38,13 +38,10 @@ class MirrorPadOp : public XlaOpKernel { // - [1, 2, 3, 3, 2] in symmetric mode. int64 excluded_edges = mode == MirrorPadMode::REFLECT ? 1 : 0; xla::XlaOp accum = t; - for (int64 dimno = xla::ShapeUtil::Rank(original_shape) - 1; dimno >= 0; - --dimno) { + for (int64 dimno = original_shape.rank() - 1; dimno >= 0; --dimno) { auto t_rev = xla::Rev(accum, {dimno}); - TF_ASSIGN_OR_RETURN(int64 lhs_padding, - pad_literal.GetIntegralAsS64({dimno, 0})); - TF_ASSIGN_OR_RETURN(int64 rhs_padding, - pad_literal.GetIntegralAsS64({dimno, 1})); + int64 lhs_padding = pad_literal.Get({dimno, 0}); + int64 rhs_padding = pad_literal.Get({dimno, 1}); int64 dim_size = original_shape.dimensions(dimno); // Padding amounts on each side must be no more than the size of the @@ -65,8 +62,8 @@ class MirrorPadOp : public XlaOpKernel { } void Compile(XlaOpKernelContext* ctx) override { - const TensorShape input_shape = ctx->InputShape(0); - const TensorShape pad_shape = ctx->InputShape(1); + const TensorShape input_shape = ctx->InputShape("input"); + const TensorShape pad_shape = ctx->InputShape("paddings"); MirrorPadMode mode; OP_REQUIRES_OK(ctx, GetNodeAttr(def(), "mode", &mode)); @@ -81,23 +78,19 @@ class MirrorPadOp : public XlaOpKernel { TensorShapeUtils::IsMatrix(pad_shape) && pad_shape.dim_size(1) == 2, errors::InvalidArgument("paddings must be a matrix with 2 columns: ", pad_shape.DebugString())); - const int fixed_dims = - (allow_legacy_scalars() && dims == 0 && pad_shape.dim_size(0) == 1) - ? 1 - : dims; OP_REQUIRES( - ctx, fixed_dims == pad_shape.dim_size(0), + ctx, dims == pad_shape.dim_size(0), errors::InvalidArgument( "The first dimension of paddings must be the rank of inputs", pad_shape.DebugString(), " ", input_shape.DebugString())); // Evaluate the 'padding' constant input, reshaping to a matrix. xla::Literal pad_literal; - OP_REQUIRES_OK( - ctx, ctx->ConstantInputReshaped(1, {fixed_dims, 2}, &pad_literal)); + OP_REQUIRES_OK(ctx, + ctx->ConstantInputAsInt64Literal("paddings", &pad_literal)); xla::XlaBuilder* b = ctx->builder(); - auto in0 = ctx->Input(0); + auto in0 = ctx->Input("input"); xla::StatusOr in0_shape = b->GetShape(in0); OP_REQUIRES(ctx, in0_shape.ok(), in0_shape.status()); xla::StatusOr accum_status = @@ -112,7 +105,7 @@ class MirrorPadOp : public XlaOpKernel { TF_DISALLOW_COPY_AND_ASSIGN(MirrorPadOp); }; -REGISTER_XLA_OP(Name("MirrorPad").CompileTimeConstInput("paddings"), +REGISTER_XLA_OP(Name("MirrorPad").CompileTimeConstantInput("paddings"), MirrorPadOp); } // namespace diff --git a/tensorflow/compiler/tf2xla/kernels/one_hot_op.cc b/tensorflow/compiler/tf2xla/kernels/one_hot_op.cc index cac2eea96eeed723b2a63bc9193070cad04b005d..aba54578d97c1e455f67efa2877ddc25dab68ac0 100644 --- a/tensorflow/compiler/tf2xla/kernels/one_hot_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/one_hot_op.cc @@ -76,7 +76,7 @@ class OneHotOp : public XlaOpKernel { TF_DISALLOW_COPY_AND_ASSIGN(OneHotOp); }; -REGISTER_XLA_OP(Name("OneHot").CompileTimeConstInput("depth"), OneHotOp); +REGISTER_XLA_OP(Name("OneHot").CompileTimeConstantInput("depth"), OneHotOp); } // namespace } // namespace tensorflow 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/pad_op.cc b/tensorflow/compiler/tf2xla/kernels/pad_op.cc index e5937b56c17d01892928b073da09f38941ea1bbb..36ea70ac392ff18fb52d400efa886533f8335eba 100644 --- a/tensorflow/compiler/tf2xla/kernels/pad_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/pad_op.cc @@ -20,6 +20,7 @@ limitations under the License. #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 { @@ -29,40 +30,36 @@ class PadOp : public XlaOpKernel { explicit PadOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} void Compile(XlaOpKernelContext* ctx) override { - const TensorShape input_shape = ctx->InputShape(0); - const TensorShape pad_shape = ctx->InputShape(1); + const TensorShape input_shape = ctx->InputShape("input"); + const TensorShape pad_shape = ctx->InputShape("paddings"); const int dims = input_shape.dims(); OP_REQUIRES( ctx, TensorShapeUtils::IsMatrix(pad_shape) && pad_shape.dim_size(1) == 2, errors::InvalidArgument("paddings must be a matrix with 2 columns: ", pad_shape.DebugString())); - const int fixed_dims = - (allow_legacy_scalars() && dims == 0 && pad_shape.dim_size(0) == 1) - ? 1 - : dims; OP_REQUIRES( - ctx, fixed_dims == pad_shape.dim_size(0), + ctx, dims == pad_shape.dim_size(0), errors::InvalidArgument( "The first dimension of paddings must be the rank of inputs", pad_shape.DebugString(), " ", input_shape.DebugString())); - if (fixed_dims == 0) { + xla::XlaOp input = ctx->Input("input"); + if (dims == 0) { // Tensor is rank 0. Return it unchanged. - ctx->SetOutput(0, ctx->Input(0)); + ctx->SetOutput(0, input); return; } - // Evaluate the 'padding' constant input, reshaping to a matrix. xla::Literal pad_literal; - OP_REQUIRES_OK( - ctx, ctx->ConstantInputReshaped(1, {fixed_dims, 2}, &pad_literal)); + OP_REQUIRES_OK(ctx, + ctx->ConstantInputAsInt64Literal("paddings", &pad_literal)); xla::PaddingConfig config; - for (int i = 0; i < fixed_dims; ++i) { + for (int i = 0; i < dims; ++i) { auto* dim = config.add_dimensions(); - int before = pad_literal.Get({i, 0}); - int after = pad_literal.Get({i, 1}); + int before = pad_literal.Get({i, 0}); + int after = pad_literal.Get({i, 1}); OP_REQUIRES(ctx, before >= 0 && after >= 0, errors::InvalidArgument( "Paddings must be non-negative: ", before, " ", after)); @@ -73,18 +70,19 @@ class PadOp : public XlaOpKernel { // PadV2 added a "constant_values" input that indicates the pad value. xla::XlaOp constant_values; if (ctx->num_inputs() == 3) { - OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(ctx->InputShape(2)), - errors::InvalidArgument("constant_values must be a scalar.")); - ctx->SetOutput(0, xla::Pad(ctx->Input(0), ctx->Input(2), config)); + OP_REQUIRES( + ctx, TensorShapeUtils::IsScalar(ctx->InputShape("constant_values")), + errors::InvalidArgument("constant_values must be a scalar.")); + ctx->SetOutput(0, xla::Pad(input, ctx->Input("constant_values"), config)); } else { auto zero = XlaHelpers::Zero(ctx->builder(), input_type(0)); - ctx->SetOutput(0, xla::Pad(ctx->Input(0), zero, config)); + ctx->SetOutput(0, xla::Pad(input, zero, config)); } } }; -REGISTER_XLA_OP(Name("Pad").CompileTimeConstInput("paddings"), PadOp); -REGISTER_XLA_OP(Name("PadV2").CompileTimeConstInput("paddings"), PadOp); +REGISTER_XLA_OP(Name("Pad").CompileTimeConstantInput("paddings"), PadOp); +REGISTER_XLA_OP(Name("PadV2").CompileTimeConstantInput("paddings"), PadOp); } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/permute_op.cc b/tensorflow/compiler/tf2xla/kernels/permute_op.cc index 94b51e1a586c6cf623c181abf200b91851c7ba05..71920bf5c1e6aa5981aafa8b611cc01c0917e02b 100644 --- a/tensorflow/compiler/tf2xla/kernels/permute_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/permute_op.cc @@ -75,8 +75,7 @@ class DataFormatVecPermuteOp : public XlaOpKernel { } auto keys = xla::ConstantR1(builder, absl::Span(dst_indices)); if (input_rank == 2) { - keys = xla::BroadcastInDim( - keys, xla::ShapeUtil::MakeShape(xla::S32, {4, 2}), {0}); + keys = xla::BroadcastInDim(keys, {4, 2}, {0}); } auto sorted = xla::Sort(keys, {ctx->Input(0)}, 0); auto output = xla::GetTupleElement(sorted, 1); diff --git a/tensorflow/compiler/tf2xla/kernels/pooling_ops.cc b/tensorflow/compiler/tf2xla/kernels/pooling_ops.cc index 27690c156e4da129ad139f3880bba3a208b5606d..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" @@ -152,7 +152,12 @@ class MaxPoolOp : public PoolingOp { public: MaxPoolOp(OpKernelConstruction* ctx, int num_spatial_dims) : PoolingOp(ctx, /*num_spatial_dims=*/num_spatial_dims, - /*reduction_type=*/ctx->input_type(0)) {} + /*reduction_type=*/ctx->input_type(0)) { + string data_format_str; + OP_REQUIRES_OK(ctx, ctx->GetAttr("data_format", &data_format_str)); + OP_REQUIRES(ctx, FormatFromString(data_format_str, &data_format_), + errors::InvalidArgument("Invalid data format")); + } void Compile(XlaOpKernelContext* ctx) override { auto ksize_or_error = GetKernelSize(ctx); @@ -180,16 +185,12 @@ class MaxPool2DOp : public MaxPoolOp { public: explicit MaxPool2DOp(OpKernelConstruction* ctx) : MaxPoolOp(ctx, /*num_spatial_dims=*/2) { - string data_format_str; - OP_REQUIRES_OK(ctx, ctx->GetAttr("data_format", &data_format_str)); - OP_REQUIRES(ctx, FormatFromString(data_format_str, &data_format_), - errors::InvalidArgument("Invalid data format")); } }; REGISTER_XLA_OP(Name("MaxPool"), MaxPool2DOp); REGISTER_XLA_OP(Name("MaxPoolV2") - .CompileTimeConstInput("ksize") - .CompileTimeConstInput("strides"), + .CompileTimeConstantInput("ksize") + .CompileTimeConstantInput("strides"), MaxPool2DOp); class MaxPool3DOp : public MaxPoolOp { @@ -204,7 +205,12 @@ class AvgPoolOp : public PoolingOp { AvgPoolOp(OpKernelConstruction* ctx, int num_spatial_dims) : PoolingOp(ctx, /*num_spatial_dims=*/num_spatial_dims, /*reduction_type=*/ - XlaHelpers::SumAccumulationType(ctx->input_type(0))) {} + XlaHelpers::SumAccumulationType(ctx->input_type(0))) { + string data_format_str; + OP_REQUIRES_OK(ctx, ctx->GetAttr("data_format", &data_format_str)); + OP_REQUIRES(ctx, FormatFromString(data_format_str, &data_format_), + errors::InvalidArgument("Invalid data format")); + } void Compile(XlaOpKernelContext* ctx) override { auto ksize_or_error = GetKernelSize(ctx); @@ -241,10 +247,6 @@ class AvgPool2DOp : public AvgPoolOp { public: explicit AvgPool2DOp(OpKernelConstruction* ctx) : AvgPoolOp(ctx, /*num_spatial_dims=*/2) { - string data_format_str; - OP_REQUIRES_OK(ctx, ctx->GetAttr("data_format", &data_format_str)); - OP_REQUIRES(ctx, FormatFromString(data_format_str, &data_format_), - errors::InvalidArgument("Invalid data format")); } }; REGISTER_XLA_OP(Name("AvgPool"), AvgPool2DOp); @@ -360,8 +362,8 @@ class MaxPool2DGradOp : public MaxPoolGradOp { }; REGISTER_XLA_OP(Name("MaxPoolGrad"), MaxPool2DGradOp); REGISTER_XLA_OP(Name("MaxPoolGradV2") - .CompileTimeConstInput("ksize") - .CompileTimeConstInput("strides"), + .CompileTimeConstantInput("ksize") + .CompileTimeConstantInput("strides"), MaxPool2DGradOp); class MaxPool3DGradOp : public MaxPoolGradOp { @@ -390,6 +392,11 @@ class AvgPoolGradOp : public XlaOpKernel { OP_REQUIRES(ctx, ksize_[0] == 1 && stride_[0] == 1, errors::Unimplemented( "Pooling is not yet supported on the batch dimension.")); + + string data_format; + OP_REQUIRES_OK(ctx, ctx->GetAttr("data_format", &data_format)); + OP_REQUIRES(ctx, FormatFromString(data_format, &data_format_), + errors::InvalidArgument("Invalid data format")); } int num_dims() const { return num_spatial_dims_ + 2; } @@ -449,22 +456,20 @@ class AvgPool2DGradOp : public AvgPoolGradOp { public: explicit AvgPool2DGradOp(OpKernelConstruction* ctx) : AvgPoolGradOp(ctx, /*num_spatial_dims=*/2) { - string data_format; - OP_REQUIRES_OK(ctx, ctx->GetAttr("data_format", &data_format)); - OP_REQUIRES(ctx, FormatFromString(data_format, &data_format_), - errors::InvalidArgument("Invalid data format")); } }; -REGISTER_XLA_OP(Name("AvgPoolGrad").CompileTimeConstInput("orig_input_shape"), - AvgPool2DGradOp); +REGISTER_XLA_OP( + Name("AvgPoolGrad").CompileTimeConstantInput("orig_input_shape"), + AvgPool2DGradOp); class AvgPool3DGradOp : public AvgPoolGradOp { public: explicit AvgPool3DGradOp(OpKernelConstruction* ctx) : AvgPoolGradOp(ctx, /*num_spatial_dims=*/3) {} }; -REGISTER_XLA_OP(Name("AvgPool3DGrad").CompileTimeConstInput("orig_input_shape"), - AvgPool3DGradOp); +REGISTER_XLA_OP( + Name("AvgPool3DGrad").CompileTimeConstantInput("orig_input_shape"), + AvgPool3DGradOp); class MaxPoolGradGradOp : public XlaOpKernel { public: @@ -632,8 +637,8 @@ REGISTER_XLA_OP(Name("MaxPoolGradGrad").TypeConstraint("T", DT_FLOAT), MaxPool2DGradGradOp); REGISTER_XLA_OP(Name("MaxPoolGradGradV2") .TypeConstraint("T", DT_FLOAT) - .CompileTimeConstInput("ksize") - .CompileTimeConstInput("strides"), + .CompileTimeConstantInput("ksize") + .CompileTimeConstantInput("strides"), MaxPool2DGradGradOp); class MaxPool3DGradGradOp : public MaxPoolGradGradOp { diff --git a/tensorflow/compiler/tf2xla/kernels/qr_op.cc b/tensorflow/compiler/tf2xla/kernels/qr_op.cc index 7ea0afc1f53cbe4cfcc3f6121a4ecd55864c1b52..66ec40a946b8a063d84acd33daf81f52ea2c35ed 100644 --- a/tensorflow/compiler/tf2xla/kernels/qr_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/qr_op.cc @@ -13,9 +13,9 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/compiler/tf2xla/lib/qr.h" #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" +#include "tensorflow/compiler/xla/client/lib/qr.h" namespace tensorflow { namespace { @@ -26,7 +26,7 @@ class QROp : public XlaOpKernel { OP_REQUIRES_OK(ctx, ctx->GetAttr("full_matrices", &full_matrices_)); } void Compile(XlaOpKernelContext* ctx) override { - auto result = QRDecomposition(ctx->Input(0), full_matrices_); + auto result = xla::QRDecomposition(ctx->Input(0), full_matrices_); if (!result.ok()) { ctx->SetStatus(result.status()); return; diff --git a/tensorflow/compiler/tf2xla/kernels/quantize_and_dequantize_op.cc b/tensorflow/compiler/tf2xla/kernels/quantize_and_dequantize_op.cc index 6f4ed496a1774dde68dd9d5fbd37995d615b678c..7fe102428db1cc5ce16037f56fa301d1941da8e3 100644 --- a/tensorflow/compiler/tf2xla/kernels/quantize_and_dequantize_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/quantize_and_dequantize_op.cc @@ -19,6 +19,7 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "tensorflow/compiler/xla/client/lib/arithmetic.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/client/xla_computation.h" #include "tensorflow/core/platform/macros.h" @@ -26,12 +27,26 @@ limitations under the License. namespace tensorflow { namespace { +enum QuantizerRoundMode { + // Round half up: if the fraction of y is exactly 0.5, then + // round(y) = y + 0.5 + // E.g., -5.5 gets rounded to -5, -5.4 goes to -5, + // 5.4 goes to 5, and 5.5 goes to 6. + ROUND_HALF_UP, + // Round half to even: if the fraction of y is exactly 0.5, then round(y) is + // the nearest even integer to y. + // E.g., 23.5 gets rounded to 24, 24.5 gets rounded to 24, while -23.5 becomes + // -24, and -24.5 gets rounded to 24. + ROUND_HALF_TO_EVEN, +}; + class QuantizeAndDequantizeOp : public XlaOpKernel { public: explicit QuantizeAndDequantizeOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { OP_REQUIRES_OK(ctx, ctx->GetAttr("signed_input", &signed_input_)); OP_REQUIRES_OK(ctx, ctx->GetAttr("range_given", &range_given_)); + round_mode_ = ROUND_HALF_TO_EVEN; } void Compile(XlaOpKernelContext* ctx) override { @@ -117,8 +132,17 @@ class QuantizeAndDequantizeOp : public XlaOpKernel { // in that case they were measured from the tensor. input = Clamp(min_range, input, max_range); } - xla::XlaOp result = - Floor((input - min_range) * scale + half) * inverse_scale + min_range; + xla::XlaOp result; + switch (round_mode_) { + case ROUND_HALF_TO_EVEN: { + result = xla::RoundToEven(input * scale) * inverse_scale; + break; + } + case ROUND_HALF_UP: { + result = Floor(input * scale + half) * inverse_scale; + break; + } + } ctx->SetOutput(0, result); } @@ -126,6 +150,7 @@ class QuantizeAndDequantizeOp : public XlaOpKernel { int64 num_bits_ = -1; bool signed_input_; bool range_given_; + QuantizerRoundMode round_mode_; }; class QuantizeAndDequantizeV2Op : public QuantizeAndDequantizeOp { @@ -136,6 +161,20 @@ class QuantizeAndDequantizeV2Op : public QuantizeAndDequantizeOp { OP_REQUIRES(ctx, num_bits_ > 0 && num_bits_ < (signed_input_ ? 62 : 63), errors::InvalidArgument("num_bits is out of range: ", num_bits_, " with signed_input_ ", signed_input_)); + string round_mode_string; + OP_REQUIRES_OK(ctx, ctx->GetAttr("round_mode", &round_mode_string)); + OP_REQUIRES( + ctx, + (round_mode_string == "HALF_UP" || round_mode_string == "HALF_TO_EVEN"), + errors::InvalidArgument("Round mode string must be " + "'HALF_UP' or " + "'HALF_TO_EVEN', is '" + + round_mode_string + "'")); + if (round_mode_string == "HALF_UP") { + round_mode_ = ROUND_HALF_UP; + } else if (round_mode_string == "HALF_TO_EVEN") { + round_mode_ = ROUND_HALF_TO_EVEN; + } } }; diff --git a/tensorflow/compiler/tf2xla/kernels/random_ops.cc b/tensorflow/compiler/tf2xla/kernels/random_ops.cc index 7ef6fa305b7f5b5aae187808f856a9273f101e14..01b047f732f0e9fb3b45b272e7886e2f8cf4fff4 100644 --- a/tensorflow/compiler/tf2xla/kernels/random_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/random_ops.cc @@ -20,13 +20,12 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/kernels/gather_op_helpers.h" #include "tensorflow/compiler/tf2xla/lib/random.h" #include "tensorflow/compiler/tf2xla/lib/util.h" -#include "tensorflow/compiler/tf2xla/lib/while_loop.h" #include "tensorflow/compiler/tf2xla/shape_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/client/lib/numeric.h" +#include "tensorflow/compiler/xla/client/lib/loops.h" #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/tensor.h" @@ -58,7 +57,7 @@ class RandomUniformOp : public XlaOpKernel { TF_DISALLOW_COPY_AND_ASSIGN(RandomUniformOp); }; -REGISTER_XLA_OP(Name("RandomUniform").CompileTimeConstInput("shape"), +REGISTER_XLA_OP(Name("RandomUniform").CompileTimeConstantInput("shape"), RandomUniformOp); class RandomShuffleOp : public XlaOpKernel { @@ -161,23 +160,30 @@ 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): auto swap_loop_result = - XlaForEachIndex(n, xla::S32, swap_body_fn, {swaps, indices}, - "indices_swap_loop", builder) + xla::ForEachIndex(n, xla::S32, swap_body_fn, {swaps, indices}, + "indices_swap_loop", builder) .ValueOrDie(); auto swapped_indices = swap_loop_result[1]; @@ -227,7 +233,7 @@ class RandomUniformIntOp : public XlaOpKernel { TF_DISALLOW_COPY_AND_ASSIGN(RandomUniformIntOp); }; -REGISTER_XLA_OP(Name("RandomUniformInt").CompileTimeConstInput("shape"), +REGISTER_XLA_OP(Name("RandomUniformInt").CompileTimeConstantInput("shape"), RandomUniformIntOp); class RandomStandardNormalOp : public XlaOpKernel { @@ -256,7 +262,7 @@ class RandomStandardNormalOp : public XlaOpKernel { TF_DISALLOW_COPY_AND_ASSIGN(RandomStandardNormalOp); }; -REGISTER_XLA_OP(Name("RandomStandardNormal").CompileTimeConstInput("shape"), +REGISTER_XLA_OP(Name("RandomStandardNormal").CompileTimeConstantInput("shape"), RandomStandardNormalOp); class TruncatedNormalOp : public XlaOpKernel { @@ -282,7 +288,7 @@ class TruncatedNormalOp : public XlaOpKernel { }; REGISTER_XLA_OP(Name("TruncatedNormal") - .CompileTimeConstInput("shape") + .CompileTimeConstantInput("shape") .TypeConstraint("dtype", DT_FLOAT), TruncatedNormalOp); diff --git a/tensorflow/compiler/tf2xla/kernels/reduce_window_op.cc b/tensorflow/compiler/tf2xla/kernels/reduce_window_op.cc index 8eee5b12991fb377203d780cecd8916952bd699a..dacdbc88e005304bc64ea35c8985711afca41eae 100644 --- a/tensorflow/compiler/tf2xla/kernels/reduce_window_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/reduce_window_op.cc @@ -130,11 +130,11 @@ class ReduceWindowOp : public XlaOpKernel { }; REGISTER_XLA_OP(Name("XlaReduceWindow") - .CompileTimeConstInput("window_dimensions") - .CompileTimeConstInput("window_strides") - .CompileTimeConstInput("base_dilations") - .CompileTimeConstInput("window_dilations") - .CompileTimeConstInput("padding"), + .CompileTimeConstantInput("window_dimensions") + .CompileTimeConstantInput("window_strides") + .CompileTimeConstantInput("base_dilations") + .CompileTimeConstantInput("window_dilations") + .CompileTimeConstantInput("padding"), ReduceWindowOp); } // namespace diff --git a/tensorflow/compiler/tf2xla/kernels/reduction_ops.cc b/tensorflow/compiler/tf2xla/kernels/reduction_ops.cc index 0d260fa8fcaa513d7854c1e9215952404d555c70..65e158d64fdd7df62d50b81c9e488b2d03476fb7 100644 --- a/tensorflow/compiler/tf2xla/kernels/reduction_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/reduction_ops.cc @@ -41,7 +41,8 @@ class SumOp : public XlaReductionOp { } }; -REGISTER_XLA_OP(Name("Sum").CompileTimeConstInput("reduction_indices"), SumOp); +REGISTER_XLA_OP(Name("Sum").CompileTimeConstantInput("reduction_indices"), + SumOp); class ProdOp : public XlaReductionOp { public: @@ -59,7 +60,7 @@ class ProdOp : public XlaReductionOp { } }; -REGISTER_XLA_OP(Name("Prod").CompileTimeConstInput("reduction_indices"), +REGISTER_XLA_OP(Name("Prod").CompileTimeConstantInput("reduction_indices"), ProdOp); class MinOp : public XlaReductionOp { @@ -77,7 +78,8 @@ class MinOp : public XlaReductionOp { } }; -REGISTER_XLA_OP(Name("Min").CompileTimeConstInput("reduction_indices"), MinOp); +REGISTER_XLA_OP(Name("Min").CompileTimeConstantInput("reduction_indices"), + MinOp); class MaxOp : public XlaReductionOp { public: @@ -94,7 +96,8 @@ class MaxOp : public XlaReductionOp { } }; -REGISTER_XLA_OP(Name("Max").CompileTimeConstInput("reduction_indices"), MaxOp); +REGISTER_XLA_OP(Name("Max").CompileTimeConstantInput("reduction_indices"), + MaxOp); class MeanOp : public XlaReductionOp { public: @@ -110,16 +113,25 @@ class MeanOp : public XlaReductionOp { xla::Add(scalar_lhs, scalar_rhs); } - xla::XlaOp BuildFinalizer(xla::XlaBuilder* builder, - const xla::XlaOp& reduce_output, - int64 num_elements_reduced) override { - auto divisor = XlaHelpers::IntegerLiteral(builder, input_type(0), - num_elements_reduced); - return reduce_output / divisor; + xla::XlaOp BuildFinalizer( + xla::XlaBuilder* /*builder*/, const xla::XlaOp& input, + const xla::XlaOp& reduce_output, + const std::vector& dimensions_to_reduce) override { + if (dimensions_to_reduce.empty()) { + return reduce_output; + } + auto divisor = xla::GetDimensionSize(input, dimensions_to_reduce[0]); + for (int i = 1; i < dimensions_to_reduce.size(); i++) { + auto size = xla::GetDimensionSize(input, dimensions_to_reduce[i]); + divisor = xla::Mul(divisor, size); + } + divisor = xla::ConvertElementType(divisor, xla_reduction_type_); + return XlaHelpers::ConvertElementType(reduce_output / divisor, + input_type(0)); } }; -REGISTER_XLA_OP(Name("Mean").CompileTimeConstInput("reduction_indices"), +REGISTER_XLA_OP(Name("Mean").CompileTimeConstantInput("reduction_indices"), MeanOp); class AllOp : public XlaReductionOp { @@ -137,7 +149,8 @@ class AllOp : public XlaReductionOp { } }; -REGISTER_XLA_OP(Name("All").CompileTimeConstInput("reduction_indices"), AllOp); +REGISTER_XLA_OP(Name("All").CompileTimeConstantInput("reduction_indices"), + AllOp); class AnyOp : public XlaReductionOp { public: @@ -154,7 +167,8 @@ class AnyOp : public XlaReductionOp { } }; -REGISTER_XLA_OP(Name("Any").CompileTimeConstInput("reduction_indices"), AnyOp); +REGISTER_XLA_OP(Name("Any").CompileTimeConstantInput("reduction_indices"), + AnyOp); } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/reduction_ops.h b/tensorflow/compiler/tf2xla/kernels/reduction_ops.h index 466e79828d111ee7cadcf713703e8f252c63e62c..af716eab79886791e8507a84984b7ca60865d00e 100644 --- a/tensorflow/compiler/tf2xla/kernels/reduction_ops.h +++ b/tensorflow/compiler/tf2xla/kernels/reduction_ops.h @@ -48,13 +48,14 @@ class XlaReductionOp : public XlaOpKernel { const xla::XlaOp& scalar_rhs) = 0; // Applies a transformation to the output of the reduction. The desired - // computation should be added to 'builder'. Argument 'reduce_output' is the - // output of the reduction. 'num_elements_reduced' is the number of elements - // that contributed to the reduction. Returns the transformed reduction - // output, Defaults to returning 'reduce_output' unchanged. - virtual xla::XlaOp BuildFinalizer(xla::XlaBuilder* builder, - const xla::XlaOp& reduce_output, - int64 num_elements_reduced); + // computation should be added to 'builder'. Argument 'input' is the original + // input of the reduction; 'reduce_output' is the output of the reduction. + // Returns the transformed reduction output. Defaults to returning + // 'reduce_output' converted to the input type. + virtual xla::XlaOp BuildFinalizer( + xla::XlaBuilder* builder, const xla::XlaOp& input, + const xla::XlaOp& reduce_output, + const std::vector& dimensions_to_reduce); void Compile(XlaOpKernelContext* ctx) override; diff --git a/tensorflow/compiler/tf2xla/kernels/reduction_ops_common.cc b/tensorflow/compiler/tf2xla/kernels/reduction_ops_common.cc index 118f2798d559f43acb7f6394a7337426164325ef..2ca2a85244b4edfe75db3d4fff6c2058adc2bf71 100644 --- a/tensorflow/compiler/tf2xla/kernels/reduction_ops_common.cc +++ b/tensorflow/compiler/tf2xla/kernels/reduction_ops_common.cc @@ -35,12 +35,13 @@ XlaReductionOp::XlaReductionOp(OpKernelConstruction* ctx, ctx, DataTypeToPrimitiveType(reduction_type_, &xla_reduction_type_)); } -// Unless BuildFinalizer is overridden the reduction has no -// finalizer. -xla::XlaOp XlaReductionOp::BuildFinalizer(xla::XlaBuilder* builder, - const xla::XlaOp& reduce_output, - int64 num_elements_reduced) { - return reduce_output; +// The default finalizer converts the results back into the input type. This can +// be overridden. +xla::XlaOp XlaReductionOp::BuildFinalizer( + xla::XlaBuilder* /*builder*/, const xla::XlaOp& /*input*/, + const xla::XlaOp& reduce_output, + const std::vector& /*dimensions_to_reduce*/) { + return XlaHelpers::ConvertElementType(reduce_output, input_type(0)); } void XlaReductionOp::Compile(XlaOpKernelContext* ctx) { @@ -71,7 +72,6 @@ void XlaReductionOp::Compile(XlaOpKernelContext* ctx) { absl::InlinedVector bitmap(data_shape.dims(), false); std::vector xla_axes; - int64 num_elements_reduced = 1LL; for (int64 i = 0; i < axes_tensor_shape.num_elements(); ++i) { int64 index = axes[i]; OP_REQUIRES(ctx, @@ -82,7 +82,6 @@ void XlaReductionOp::Compile(XlaOpKernelContext* ctx) { index = (index + data_shape.dims()) % data_shape.dims(); bitmap[index] = true; xla_axes.push_back(index); - num_elements_reduced *= data_shape.dim_size(index); } std::vector final_shape; @@ -118,8 +117,7 @@ void XlaReductionOp::Compile(XlaOpKernelContext* ctx) { xla::XlaComputation reduction_computation = r.Build().ConsumeValueOrDie(); auto reduce = xla::Reduce(data, initial, reduction_computation, xla_axes); - auto deconverted = XlaHelpers::ConvertElementType(b, reduce, input_type(0)); - auto finalized = BuildFinalizer(b, deconverted, num_elements_reduced); + auto finalized = BuildFinalizer(b, data, reduce, xla_axes); auto result = keep_dims_ ? xla::Reshape(finalized, final_shape) : finalized; ctx->SetOutput(0, result); } diff --git a/tensorflow/compiler/tf2xla/kernels/relu_op.cc b/tensorflow/compiler/tf2xla/kernels/relu_op.cc index d35777ccb1271ec6a7c9972c714d06b2415d9c34..a8e230ba107ce8a73f3e80f0e55fa27eea31338f 100644 --- a/tensorflow/compiler/tf2xla/kernels/relu_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/relu_op.cc @@ -15,14 +15,12 @@ limitations under the License. // Native XLA implementations of XLA Relu Ops -#include "tensorflow/compiler/tf2xla/kernels/cwise_ops.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/compiler/xla/literal.h" -#include "tensorflow/core/framework/kernel_def_builder.h" -#include "tensorflow/core/framework/types.h" -#include "tensorflow/core/kernels/no_op.h" namespace tensorflow { namespace { @@ -37,6 +35,7 @@ class ReluOp : public XlaOpKernel { ctx->SetOutput(0, xla::Max(zero, ctx->Input(0))); } }; +REGISTER_XLA_OP(Name("Relu"), ReluOp); class Relu6Op : public XlaOpKernel { public: @@ -49,6 +48,22 @@ class Relu6Op : public XlaOpKernel { ctx->SetOutput(0, xla::Clamp(zero, ctx->Input(0), six)); } }; +REGISTER_XLA_OP(Name("Relu6"), Relu6Op); + +class LeakyReluOp : public XlaOpKernel { + public: + explicit LeakyReluOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("alpha", &alpha_)); + } + void Compile(XlaOpKernelContext* ctx) override { + auto features = ctx->Input("features"); + auto output = + xla::Max(features, features * xla::ScalarLike(features, alpha_)); + ctx->SetOutput(0, output); + } + float alpha_; +}; +REGISTER_XLA_OP(Name("LeakyRelu"), LeakyReluOp); class ReluGradOp : public XlaOpKernel { public: @@ -64,6 +79,7 @@ class ReluGradOp : public XlaOpKernel { ctx->SetOutput(0, xla::Select(pred, ctx->Input(0), zero)); } }; +REGISTER_XLA_OP(Name("ReluGrad"), ReluGradOp); class Relu6GradOp : public XlaOpKernel { public: @@ -83,11 +99,24 @@ class Relu6GradOp : public XlaOpKernel { ctx->SetOutput(0, out); } }; - -REGISTER_XLA_OP(Name("Relu"), ReluOp); -REGISTER_XLA_OP(Name("Relu6"), Relu6Op); -REGISTER_XLA_OP(Name("ReluGrad"), ReluGradOp); REGISTER_XLA_OP(Name("Relu6Grad"), Relu6GradOp); +class LeakyReluGradOp : public XlaOpKernel { + public: + explicit LeakyReluGradOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("alpha", &alpha_)); + } + void Compile(XlaOpKernelContext* ctx) override { + auto gradients = ctx->Input("gradients"); + auto features = ctx->Input("features"); + auto output = + xla::Select(xla::Gt(features, xla::ScalarLike(features, 0)), gradients, + gradients * xla::ScalarLike(gradients, alpha_)); + ctx->SetOutput(0, output); + } + float alpha_; +}; +REGISTER_XLA_OP(Name("LeakyReluGrad"), LeakyReluGradOp); + } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/resampler_ops.cc b/tensorflow/compiler/tf2xla/kernels/resampler_ops.cc new file mode 100644 index 0000000000000000000000000000000000000000..f9985d526033ca675c701a508a3d1576e46bc5f7 --- /dev/null +++ b/tensorflow/compiler/tf2xla/kernels/resampler_ops.cc @@ -0,0 +1,682 @@ +/* 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/tf2xla/shape_util.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/array4d.h" +#include "tensorflow/compiler/xla/client/lib/arithmetic.h" +#include "tensorflow/compiler/xla/client/lib/constants.h" +#include "tensorflow/compiler/xla/client/xla_builder.h" +#include "tensorflow/compiler/xla/literal.h" +#include "tensorflow/compiler/xla/shape_util.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/register_types.h" +#include "tensorflow/core/framework/tensor_shape.h" +#include "tensorflow/core/framework/types.pb.h" +#include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/lib/gtl/inlined_vector.h" +#include "tensorflow/core/lib/math/math_util.h" +#include "tensorflow/core/platform/types.h" + +namespace tensorflow { +namespace { + +using xla::XlaOp; + +// Calculates the bilinear weight tensor, given basis ratio (px, py) of the +// sampling position: +// W = [(1-px)*(1-py), px*(1-py), (1-px)*py, px*py] +// 'ratio' tensor has dimensions [batch, dim_0, ...dim_n, 2]. +// +// The returned tensor has dimensions [batch, dim_0, ... dim_n, 4]. +XlaOp BilinearWeights(XlaOpKernelContext* ctx, XlaOp ratio, + const TensorShape warp_shape, + xla::PrimitiveType xla_type) { + auto first_term = xla::ConstantR2( + ctx->builder(), {{1.0, 1.0}, {0.0, 1.0}, {1.0, 0.0}, {0.0, 0.0}}); + first_term = xla::ConvertElementType(first_term, xla_type); + + auto warp_dims = warp_shape.dim_sizes(); + std::vector broadcast_dims(warp_dims.begin(), warp_dims.end() - 1); + broadcast_dims.push_back(4); + broadcast_dims.push_back(2); + + const int64 broadcast_dims_size = broadcast_dims.size(); + + std::vector last_two_dims_indices = {(broadcast_dims_size - 2), + (broadcast_dims_size - 1)}; + + auto broadcast_first_term = + xla::BroadcastInDim(first_term, broadcast_dims, last_two_dims_indices); + + // Ratio is of the same dimension as warp, which is [batch, dim_0,... dim_n, + // 2], we broadcast ratio tensor to 'broadcast_dim' by keeping the + // [batch, dim_0,...dim_n] dimensions and the [2] dimension as the last + // dimension. + std::vector ratio_broadcast_indices(broadcast_dims.size()); + std::iota(ratio_broadcast_indices.begin(), ratio_broadcast_indices.end(), 0); + ratio_broadcast_indices.erase(ratio_broadcast_indices.end() - 2); + + auto broadcast_ratio = + xla::BroadcastInDim(ratio, broadcast_dims, ratio_broadcast_indices); + + auto first_term_subtract_weights = broadcast_first_term - broadcast_ratio; + + // Now we have [(1-px, 1-py), (-px, 1-py), (1-px, -py), (px, py)], need to + // flip the signs of the second and the third term. + auto sign_change = xla::ConstantR2( + ctx->builder(), {{1.0, 1.0}, {-1.0, 1.0}, {-1.0, 1.0}, {1.0, 1.0}}); + sign_change = xla::ConvertElementType(sign_change, xla_type); + + auto broadcast_sign_change = + xla::BroadcastInDim(sign_change, broadcast_dims, last_two_dims_indices); + + auto flipped = first_term_subtract_weights * broadcast_sign_change; + + // Build up the final bilinear weight tensor by multiply reduction, which + // gives: + // [(1-px)*(1-py), px*(1-py), (1-px)*py, px*py] + // for each 4 neighboring pixels where px and py are the weight of the target + // pixel we are sampling from. + return xla::Reduce( + flipped, xla::One(ctx->builder(), xla_type), + xla::CreateScalarMultiplyComputation(xla_type, ctx->builder()), + {broadcast_dims_size - 1}); +} + +// Concatenates the batch indices to the (x, y) coordinate indices. +// This is done by first creating an Iota tensor that represents the current +// batch it is in, then concatenate with the givin (coordinate) indices. +// +// The resulting tensor has dimension (batch, dim_0, ... dim_n, 3) where +// the last dimension of size 3 in turn is [batch_number, x, y]. +// The [batch_number, x, y] dimension is needed because the indices +// [x,y] alone cannot allow the xla::Gather operation to gather from the input +// data, which is of dimension [batch, height(y), width(x), channel] with +// 'batch' being the first dimension. +XlaOp ConcatenateIota(xla::XlaBuilder* b, XlaOp indices, + const TensorShape& warp_shape) { + // We need to create an iota tensor with the same batch dimension. + std::vector dimensions; + for (auto dim : warp_shape) { + dimensions.push_back(dim.size); + } + // Except the last dimension, which is of size 1. + dimensions.back() = 1; + + auto batch_indices = + xla::Iota(b, xla::ShapeUtil::MakeShape(xla::S32, dimensions), + /*iota_dimension=*/0); + + return xla::ConcatInDim(b, {batch_indices, indices}, dimensions.size() - 1); +} + +// Gathers the 2x2 neighbors of the input starting_indices, and return a +// tensor of dimension [batch, dim_0, ... dim_n, 4, data_channels]. +// 'gather_indices' is of dimension [batch, dim_0, ..., dim_n, 3] where the last +// dimension of size 3 is (batch_no, x, y). +XlaOp Gather2by2Neighbors(xla::XlaBuilder* b, XlaOp data, XlaOp gather_indices, + int64 data_channels, int warp_dims) { + xla::GatherDimensionNumbers gather_dim_numbers; + const int64 neighbor_data_dimensions = warp_dims + 2; + // Since the Gather output dimensions are [batch, dim_0, ... dim_n, 2, 2, + // data_channels], the offset dimensions for Gather is the last 3 dimensions. + gather_dim_numbers.add_offset_dims(neighbor_data_dimensions - 3); + gather_dim_numbers.add_offset_dims(neighbor_data_dimensions - 2); + gather_dim_numbers.add_offset_dims(neighbor_data_dimensions - 1); + // The last dimension of 'gather_indices' is the starting indices for gather. + gather_dim_numbers.set_index_vector_dim(warp_dims - 1); + gather_dim_numbers.add_collapsed_slice_dims(0); + gather_dim_numbers.add_start_index_map(0); + // Since input is of dimension [batch, height(y), width(x), channel], and warp + // is of dimension [batch, x, y], the ordering of x, y here needs to be + // swapped when gathering. + gather_dim_numbers.add_start_index_map(2); + gather_dim_numbers.add_start_index_map(1); + // Data dimensions are [batch, x, y, channel]. + // Output dimensions are [batch, dim_0, ... dim_n, 2, 2, data_channels]. + auto neighbors_data = xla::Gather(data, gather_indices, gather_dim_numbers, + /*slice_sizes=*/{1, 2, 2, data_channels}); + // Collapse the ...,2,2,... dimensions into ...,4,... + return xla::Collapse(neighbors_data, {warp_dims - 1, warp_dims}); +} + +// Scatter 'updates' tensor to 'grad_data' based on 'indices'. Returns the +// resulting tensor of dimension: [batch, dim_0, ...dim_n, 2, 2, data_channels]. +// This function can also be seen as the inverse of 'Gather2by2Neighbors'. +XlaOp ScatterToGradData(XlaOpKernelContext* ctx, XlaOp grad_data, XlaOp indices, + XlaOp updates, int64 warp_dims, + xla::PrimitiveType xla_type) { + xla::ScatterDimensionNumbers scatter_dim_numbers; + const int64 neighbor_data_dimensions = warp_dims + 2; + // Since the Scatter output dimensions are [batch, dim_0, ... dim_n, 2, 2, + // data_channels], the update window dimensions is the last 3 dimensions. + scatter_dim_numbers.add_update_window_dims(neighbor_data_dimensions - 3); + scatter_dim_numbers.add_update_window_dims(neighbor_data_dimensions - 2); + scatter_dim_numbers.add_update_window_dims(neighbor_data_dimensions - 1); + scatter_dim_numbers.set_index_vector_dim(warp_dims - 1); + + scatter_dim_numbers.add_inserted_window_dims(0); + scatter_dim_numbers.add_scatter_dims_to_operand_dims(0); + // Since input is of dimension [batch, height(y), width(x), channel], and warp + // is of dimension [batch, x, y], the ordering of x, y here needs to be + // swapped when scattering. + scatter_dim_numbers.add_scatter_dims_to_operand_dims(2); + scatter_dim_numbers.add_scatter_dims_to_operand_dims(1); + + return xla::Scatter(grad_data, indices, updates, + xla::CreateScalarAddComputation(xla_type, ctx->builder()), + scatter_dim_numbers); +} + +// Bounds samples to 0 if the warp image indices are out of the (-1, image_size) +// bound. +// The resulting dimension is given by 'result_dims'. +XlaOp BoundSamples(XlaOpKernelContext* ctx, XlaOp warp, + xla::PrimitiveType warp_type, TensorShape warp_shape, + std::vector result_dims, + std::vector broadcasted_dims, int64 last_warp_dim, + xla::Shape data_shape, XlaOp sample) { + auto is_gt_minus_one = + xla::Gt(warp, + xla::ConvertElementType( + xla::ConstantR1(ctx->builder(), {-1, -1}), warp_type), + /*broadcast_dimensions=*/{warp_shape.dims() - 1}); + auto is_lt_image_size = xla::Lt( + warp, + xla::ConvertElementType( + xla::ConstantR1( + ctx->builder(), + {/*width=*/static_cast(data_shape.dimensions(2)), + /*height=*/static_cast(data_shape.dimensions(1))}), + warp_type), + /*broadcast_dimensions=*/{warp_shape.dims() - 1}); + + auto is_in_bound_padded_x_y = xla::And(is_gt_minus_one, is_lt_image_size); + // Reduce along last dimension. The resulting dimension is: + // [batch, dim_0, ...dim_n]. + auto is_in_bound = xla::Reduce( + is_in_bound_padded_x_y, xla::ConstantR0(ctx->builder(), true), + xla::CreateScalarAndComputation(xla::PrimitiveType::PRED, ctx->builder()), + {last_warp_dim}); + + // Broadcast 'is_in_bound' to the same dimension as 'result_dims'. + auto broadcasted_is_in_bound = + xla::BroadcastInDim(is_in_bound, result_dims, broadcasted_dims); + + // Set out of bound samples to zero. + auto zeros = + xla::Broadcast(xla::Zero(ctx->builder(), warp_type), result_dims); + return xla::Select(broadcasted_is_in_bound, sample, zeros); +} + +// Build computation the backprop into input 'data'. +// Where input: +// grad_output is of dimension [batch, dim_0, ...dim_n, channel] +// ratio is of dimension [batch, dim_0, ...dim_n, 2] +// gather_indices is of dimension [batch, dim_0, ...dim_n, 3] +// data_shape is of dimension [batch, x(width), y(height), channel] +// +// Output: +// scatter-add to each 2x2 grad_data neighbor: +// grad_data[fx, fy, chan] += output_grad * dx * dy +// grad_data[cx, fy, chan] += output_grad * (1 - dx) * dy +// grad_data[fx, cy, chan] += output_grad * dx * (1 - dy) +// grad_data[cx, cy, chan] += output_grad * (1 - dx) * (1 - dy) +// where (dx, dy) is (1 - ratio). If (dx, dy) is out of bound, then the their +// contribution is 0 to 'grad_data'. +XlaOp CalculateGradData(XlaOpKernelContext* ctx, XlaOp grad_output, XlaOp ratio, + XlaOp gather_indices, XlaOp warp, + xla::PrimitiveType warp_type, TensorShape warp_shape, + int64 last_warp_dim, int64 data_channels, + xla::Shape data_shape) { + // Weights tensor has dimension [batch, dim_0, ... dim_n, 4]. + auto weights = BilinearWeights(ctx, ratio, warp_shape, warp_type); + + auto warp_dims = warp_shape.dim_sizes(); + std::vector warp_dims_without_last_dims(warp_dims.begin(), + warp_dims.end() - 1); + + std::vector reshaped_weights_dims = warp_dims_without_last_dims; + // Reshape the last dimension of size 4 to two dimensions [2, 2]. + reshaped_weights_dims.push_back(2); + reshaped_weights_dims.push_back(2); + std::vector reshape_dims(warp_shape.dims()); + std::iota(reshape_dims.begin(), reshape_dims.end(), 0); + // The dimension is [batch, dim_0,..., dim_n, 2, 2]. + auto reshaped_weights = xla::Reshape(weights, /*dimensions=*/reshape_dims, + /*new_sizes=*/reshaped_weights_dims); + + std::vector weights_with_channels_dims = reshaped_weights_dims; + weights_with_channels_dims.push_back(data_channels); + std::vector reshaped_weights_indices(reshaped_weights_dims.size()); + std::iota(reshaped_weights_indices.begin(), reshaped_weights_indices.end(), + 0); + + // Set out of bound weights to 0. + // The dimension of the reshaped_weight: [batch, dim_0, ...dim_n, 2, 2]. + std::vector reshaped_result_dims(warp_dims.begin(), + warp_dims.end() - 1); + reshaped_result_dims.push_back(2); + reshaped_result_dims.push_back(2); + std::vector broadcasted_dims(warp_dims.size() - 1); + std::iota(broadcasted_dims.begin(), broadcasted_dims.end(), 0); + reshaped_weights = BoundSamples(ctx, warp, warp_type, warp_shape, + reshaped_result_dims, broadcasted_dims, + last_warp_dim, data_shape, reshaped_weights); + + // The dimension is [batch, dim_0, ..., dim_n, 2, 2, data_channel]. + auto broadcast_reshaped_weights = xla::BroadcastInDim( + reshaped_weights, weights_with_channels_dims, reshaped_weights_indices); + + std::vector grad_output_indices(warp_dims_without_last_dims.size()); + std::iota(grad_output_indices.begin(), grad_output_indices.end(), 0); + grad_output_indices.push_back(weights_with_channels_dims.size() - 1); + XlaOp broadcast_grad_output = xla::BroadcastInDim( + grad_output, weights_with_channels_dims, grad_output_indices); + + auto grad_output_multiply_weights = + broadcast_grad_output * broadcast_reshaped_weights; + + auto grad_data = xla::ConstantLiteral( + ctx->builder(), xla::Literal::CreateFromShape(data_shape)); + + // Pad grad data then slice it back. + // + // After left and right column 0-padding, the new dimension of padded data + // will be [batch, x+2, y+2, channel]. + auto padded_grad_data = + xla::Pad(grad_data, xla::Zero(ctx->builder(), warp_type), + xla::MakeEdgePaddingConfig({{0, 0}, {1, 1}, {1, 1}, {0, 0}})); + + auto shifting_value = xla::ConstantR1( + ctx->builder(), {/*batch=*/0, /*x(width)=*/1, /*y(height)=*/1}); + auto shifted_gather_indices = + xla::Add(gather_indices, shifting_value, {last_warp_dim}); + + auto updated_grad_data = ScatterToGradData( + ctx, padded_grad_data, shifted_gather_indices, + grad_output_multiply_weights, warp_shape.dims(), warp_type); + + const int64 batch_size = data_shape.dimensions(0); + const int64 width = data_shape.dimensions(1); + const int64 height = data_shape.dimensions(2); + // Slice out the result accounting for the padding. + return xla::Slice( + updated_grad_data, /*start_indices=*/{0, 1, 1, 0}, + /*limit_indices=*/{batch_size, width + 1, height + 1, data_channels}, + /*strides=*/{1, 1, 1, 1}); +} + +// Build computation for the backprop into input 'warp'. +// Where input: +// warp is of dimension [batch, dim_0, ...dim_n, 2] +// grad_output is of dimension [batch, dim_0, ...dim_n, channel] +// ratio is of dimension [batch, dim_0, ...dim_n, 2] +// gather_indices is of dimension [batch, dim_0, ...dim_n, 3] where the last +// dimension of size 3 is for {batch, x(width), y(height)}. +// data is of dimension [batch, x, y, channel] +// +// Output (simplified by ignoring the batch dimensions): +// Since the forward path has: +// output = dot(weights * neighbors) +// The backprop into warp will therefore be: +// grad_warp = output_grad * d_output / d_warp +// = output_grad * (d_weights / d_warp * neighbors + d_neighbors / +// d_warp * weight) +// Where: +// d_weights / d_warp_x = [-(1 - py), (1 - py), -py, py] +// d_weights / d_warp_y = [-(1 - px), -px, (1-px), px] +// and +// d_neighbors / d_warp_x = 0 +// +// Therefore: +// grad_warp_x = py * (img_cxcy - img_fxcy) + (1-py) * (img_cxfy-img_fxfy) +// grad_warp_y = px * (img_cxcy - img_cxfy) + (1-px) * (img_fxcy-img_fxfy) +// +// where (px, py) is warp, (fx, fy) is the top left corner and (cx, cy) is the +// bottom right corner in a 2x2 neighborhood. +XlaOp CalculateGradWarp(XlaOpKernelContext* ctx, XlaOp grad_output, XlaOp ratio, + XlaOp gather_indices, XlaOp data, + TensorShape warp_shape, int64 data_channels, + xla::PrimitiveType data_type, xla::Shape data_shape) { + auto warp_dims = warp_shape.dim_sizes(); + std::vector warp_dims_without_last_dims(warp_dims.begin(), + warp_dims.end() - 1); + + // With dimension [batch, dim_0, ...dim_n, 4] + std::vector neighbor_broadcast_dims = warp_dims_without_last_dims; + neighbor_broadcast_dims.push_back(4); + + // With dimension [batch, dim_0, ...dim_n, 4] + auto neighbor_broadcast_shape = + xla::ShapeUtil::MakeShape(data_type, neighbor_broadcast_dims); + + const int64 last_warp_dim = warp_shape.dims() - 1; + + // Pad data with 0, before gathering such that 0 will be returned for samples + // in the range of (-1, 0) or (image_dimension-1, image_dimension). + // After left and right column 0-padding, the new dimension of padded data + // will be [batch, x+2, y+2, channel]. + auto padded_data = + xla::Pad(data, xla::Zero(ctx->builder(), data_type), + xla::MakeEdgePaddingConfig({{0, 0}, {1, 1}, {1, 1}, {0, 0}})); + + auto shifting_value = xla::ConstantR1( + ctx->builder(), {/*batch=*/0, /*x(width)=*/1, /*y(height)=*/1}); + auto shifted_gather_indices = + xla::Add(gather_indices, shifting_value, {last_warp_dim}); + + // The dimension is [batch, dim_0, ... dim_n, 4, data_channels] + auto neighbors_data = + Gather2by2Neighbors(ctx->builder(), padded_data, shifted_gather_indices, + data_channels, warp_shape.dims()); + + // Since we will be creating the dot product of: + // lhs: [batch, dim_0, ...dim_n, 4] + // and + // rhs: [batch, dim_0, ...dim_n, 4, data_channels] + // we choose the last dimension of lhs and the second last dimension of rhs, + // with size 4, as the contracting dimension. + xla::DotDimensionNumbers dot_dims; + for (int i = 0; i < warp_shape.dims() - 1; ++i) { + dot_dims.add_lhs_batch_dimensions(i); + dot_dims.add_rhs_batch_dimensions(i); + } + dot_dims.add_lhs_contracting_dimensions(warp_shape.dims() - 1); + dot_dims.add_rhs_contracting_dimensions(warp_shape.dims() - 1); + + // img_cxcy - img_fxcy + auto bottom_right_minus_bottom_left = xla::DotGeneral( + xla::BroadcastInDim( + xla::ConvertElementType( + xla::ConstantR1(ctx->builder(), {0, 0, -1, 1}), data_type), + neighbor_broadcast_dims, {last_warp_dim}), + neighbors_data, dot_dims, /*precision_config=*/nullptr); + + // img_cxfy - img_fxfy + auto top_right_minus_top_left = xla::DotGeneral( + xla::BroadcastInDim( + xla::ConvertElementType( + xla::ConstantR1(ctx->builder(), {-1, 1, 0, 0}), data_type), + neighbor_broadcast_dims, {last_warp_dim}), + neighbors_data, dot_dims, /*precision_config=*/nullptr); + + // img_cxcy - img_cxfy + auto bottom_right_minus_top_right = xla::DotGeneral( + xla::BroadcastInDim( + xla::ConvertElementType( + xla::ConstantR1(ctx->builder(), {0, -1, 0, 1}), data_type), + neighbor_broadcast_dims, {last_warp_dim}), + neighbors_data, dot_dims, /*precision_config=*/nullptr); + + // img_fxcy - img_fxfy + auto bottom_left_minus_top_left = xla::DotGeneral( + xla::BroadcastInDim( + xla::ConvertElementType( + xla::ConstantR1(ctx->builder(), {-1, 0, 1, 0}), data_type), + neighbor_broadcast_dims, {last_warp_dim}), + neighbors_data, dot_dims, /*precision_config=*/nullptr); + + // Slice out x and y. + auto weight_x = xla::SliceInDim(ratio, /*start_index=*/0, /*limit_index=*/1, + /*stride=*/1, /*dimno=*/last_warp_dim); + auto weight_y = xla::SliceInDim(ratio, /*start_index=*/1, /*limit_index=*/2, + /*stride=*/1, /*dimno=*/last_warp_dim); + + // Build 1 - y and 1 - x. + auto one_minus_y = xla::One(ctx->builder(), data_type) - weight_y; + auto one_minus_x = xla::One(ctx->builder(), data_type) - weight_x; + + auto x_before_reduce = + grad_output * weight_y * bottom_right_minus_bottom_left + + one_minus_y * top_right_minus_top_left; + + std::vector reshaped_sizes = warp_dims_without_last_dims; + reshaped_sizes.push_back(1); + + std::vector reshaped_dims(warp_dims_without_last_dims.size()); + std::iota(reshaped_dims.begin(), reshaped_dims.end(), 0); + + // Reduce-add along the channel dimension. + auto x_result = + xla::Reduce(x_before_reduce, xla::Zero(ctx->builder(), data_type), + xla::CreateScalarAddComputation(data_type, ctx->builder()), + {last_warp_dim}); + // Reshape before concatenating with y values. + XlaOp reshaped_x = xla::Reshape(x_result, reshaped_dims, reshaped_sizes); + + auto y_before_reduce = grad_output * weight_x * bottom_right_minus_top_right + + one_minus_x * bottom_left_minus_top_left; + // Reduce-add along the channel dimension. + auto y_result = + xla::Reduce(y_before_reduce, xla::Zero(ctx->builder(), data_type), + + xla::CreateScalarAddComputation(data_type, ctx->builder()), + {last_warp_dim}); + XlaOp reshaped_y = xla::Reshape(y_result, reshaped_dims, reshaped_sizes); + + return xla::ConcatInDim(ctx->builder(), {reshaped_x, reshaped_y}, + last_warp_dim); +} + +class ResamplerOp : public XlaOpKernel { + public: + explicit ResamplerOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} + + void Compile(XlaOpKernelContext* ctx) override { + TensorShape data_shape = ctx->InputShape("data"); + OP_REQUIRES(ctx, data_shape.dims() == 4, + errors::InvalidArgument("data must be 4-dimensional", + data_shape.DebugString())); + const int64 data_channels = data_shape.dim_size(3); + xla::PrimitiveType data_type = ctx->input_xla_type(0); + + TensorShape warp_shape = ctx->InputShape("warp"); + OP_REQUIRES(ctx, warp_shape.dims() >= 2, + errors::InvalidArgument("warp must be at least 2-dimensional", + warp_shape.DebugString())); + for (int size : warp_shape.dim_sizes()) { + OP_REQUIRES(ctx, size > 0, + errors::InvalidArgument("warp sizes must be positive, got [", + size, "]")); + } + const int64 last_warp_dim = warp_shape.dims() - 1; + // Last dimension of warp shape must be of size 2. + OP_REQUIRES(ctx, warp_shape.dim_size(last_warp_dim) == 2, + errors::InvalidArgument( + "the last dimension of warp must be exactly size 2.")); + xla::PrimitiveType warp_type = ctx->input_xla_type(1); + + XlaOp data = ctx->Input("data"); + XlaOp warp = ctx->Input("warp"); + + // Find the coordinates of the top left corner for the 2x2 region to be + // sampled from. The dimensions are [batch, dim_0, ... dim_n, 2] where the + // last dimension of size 2 in turn is [x, y]. + XlaOp top_left = xla::ConvertElementType(warp, xla::S32); + + auto gather_indices = ConcatenateIota(ctx->builder(), top_left, warp_shape); + + // The dimension is [batch, dim_0, ... dim_n, 4, data_channels] + auto neighbors_data = Gather2by2Neighbors( + ctx->builder(), data, gather_indices, data_channels, warp_shape.dims()); + + // Dimensions are [batch, dim_0, ... dim_n, 2]. + XlaOp ratio = warp - xla::ConvertElementType(top_left, data_type); + + // Obtain the bilinear blending weights, the dimension is [batch, dim_0, + // ...dim_n, 4]. + auto weights = BilinearWeights(ctx, ratio, warp_shape, data_type); + + // Since we will be creating the dot product of: + // lhs: [batch, dim_0, ...dim_n, 4] + // and + // rhs: [batch, dim_0, ...dim_n, 4, data_channels] + // we choose the last dimension of lhs and the second last dimension of rhs, + // with size 4, as the contracting dimension. + xla::DotDimensionNumbers dot_dims; + for (int i = 0; i < warp_shape.dims() - 1; ++i) { + dot_dims.add_lhs_batch_dimensions(i); + dot_dims.add_rhs_batch_dimensions(i); + } + dot_dims.add_lhs_contracting_dimensions(warp_shape.dims() - 1); + dot_dims.add_rhs_contracting_dimensions(warp_shape.dims() - 1); + + // The dimension is [batch, dim_0, ...dim_n, data_channels]. + auto blended_pixels = xla::DotGeneral(weights, neighbors_data, dot_dims, + /*precision_config=*/nullptr); + + // Handle out of boundary cases by constructing a predicate mask array based + // on the in-bound condition, and output 0 for the blended pixel value if + // out-bound. The dimension is the same as top_left: [batch, dim_0, + // ...dim_n, 2] where the last dimension of size 2 is the [x, y] coordinate. + + auto is_ge_zero = xla::Ge(warp, xla::ZerosLike(warp)); + + auto is_lt_image_size = xla::Lt( + warp, + xla::ConvertElementType( + xla::ConstantR1( + ctx->builder(), + {/*width=*/static_cast(data_shape.dim_size(2) - 1), + /*height=*/static_cast(data_shape.dim_size(1) - 1)}), + warp_type), + /*broadcast_dimensions=*/{warp_shape.dims() - 1}); + + auto is_in_bound_x_y = xla::And(is_ge_zero, is_lt_image_size); + // Reduce along last dimension. The resulting dimension is: + // [batch, dim_0, ...dim_n]. + auto is_in_bound = xla::Reduce( + is_in_bound_x_y, xla::ConstantR0(ctx->builder(), true), + xla::CreateScalarAndComputation(xla::PrimitiveType::PRED, + ctx->builder()), + {last_warp_dim}); + + // Broadcast 'is_in_bound' to the same dimension as 'blended_pixels', which + // is the dimension of the result: + // [batch, dim_0, ...dim_n, data_channels]. + auto warp_dims = warp_shape.dim_sizes(); + std::vector result_dims(warp_dims.begin(), warp_dims.end() - 1); + result_dims.push_back(data_channels); + + std::vector broadcasted_dims(warp_dims.size() - 1); + std::iota(broadcasted_dims.begin(), broadcasted_dims.end(), 0); + auto broadcasted_is_in_bound = + xla::BroadcastInDim(is_in_bound, result_dims, broadcasted_dims); + + // Set out of bound samples to zero. + auto zeros = + xla::Broadcast(xla::Zero(ctx->builder(), data_type), result_dims); + auto result = xla::Select(broadcasted_is_in_bound, blended_pixels, zeros); + + ctx->SetOutput(0, result); + } +}; + +REGISTER_XLA_OP(Name("Resampler"), ResamplerOp); + +class ResamplerGradOp : public XlaOpKernel { + public: + explicit ResamplerGradOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + DataType output_dtype; + OP_REQUIRES_OK(ctx, ctx->GetAttr("T", &output_dtype)); + } + + // TODO(b/112295522): note that sampling from image boundary is not currently + // being handled properly. + void Compile(XlaOpKernelContext* ctx) override { + TensorShape data_shape_tf = ctx->InputShape("data"); + OP_REQUIRES(ctx, data_shape_tf.dims() == 4, + errors::InvalidArgument("data must be 4-dimensional", + data_shape_tf.DebugString())); + const int64 data_channels = data_shape_tf.dim_size(3); + xla::PrimitiveType data_type = ctx->input_xla_type(0); + + TensorShape warp_shape = ctx->InputShape("warp"); + OP_REQUIRES(ctx, warp_shape.dims() >= 2, + errors::InvalidArgument("warp must be at least 2-dimensional", + warp_shape.DebugString())); + for (int size : warp_shape.dim_sizes()) { + OP_REQUIRES(ctx, size > 0, + errors::InvalidArgument("warp sizes must be positive, got [", + size, "]")); + } + // Last dimension of warp shape must be of size 2. + const int64 last_warp_dim = warp_shape.dims() - 1; + OP_REQUIRES(ctx, warp_shape.dim_size(last_warp_dim) == 2, + errors::InvalidArgument( + "the last dimension of warp must be exactly size 2.")); + xla::PrimitiveType warp_type = ctx->input_xla_type(1); + + TensorShape output_grad_shape = ctx->InputShape("grad_output"); + OP_REQUIRES( + ctx, output_grad_shape.dims() >= 2, + errors::InvalidArgument("output_grad must be at least 2-dimensional", + output_grad_shape.DebugString())); + + // Dimensions are [batch, x, y, channel]. + XlaOp data = ctx->Input("data"); + xla::Shape data_shape = TensorShapeToXLAShape(data_type, data_shape_tf); + + // Dimensions are [batch, dim_0, ...dim_n, 2]. + XlaOp warp = ctx->Input("warp"); + // Dimensions are [batch, dim_0, ...dim_n, channel]. + XlaOp grad_output = ctx->Input("grad_output"); + + // Find the top left corner coordinate for the region to be sampled from. + // The dimensions are [batch, dim_0, ... dim_n, 2] where the last dimension + // of size 2 in turn is [x, y]. + XlaOp top_left = xla::ConvertElementType(xla::Floor(warp), xla::S32); + + // Dimensions are [batch, dim_0, ... dim_n, 2]. + XlaOp ratio = warp - xla::ConvertElementType(top_left, warp_type); + + // Indices for gathering neighboring pixels. + auto gather_indices = ConcatenateIota(ctx->builder(), top_left, warp_shape); + + auto grad_data = CalculateGradData( + ctx, grad_output, ratio, gather_indices, warp, warp_type, warp_shape, + last_warp_dim, data_channels, data_shape); + + auto grad_warp = + CalculateGradWarp(ctx, grad_output, ratio, gather_indices, data, + warp_shape, data_channels, data_type, data_shape); + auto warp_dims = warp_shape.dim_sizes(); + std::vector result_dims(warp_dims.begin(), warp_dims.end() - 1); + result_dims.push_back(2); + std::vector broadcasted_dims(warp_dims.size() - 1); + std::iota(broadcasted_dims.begin(), broadcasted_dims.end(), 0); + auto grad_warp_bounded = + BoundSamples(ctx, warp, warp_type, warp_shape, result_dims, + broadcasted_dims, last_warp_dim, data_shape, grad_warp); + + ctx->SetOutput(0, grad_data); + ctx->SetOutput(1, grad_warp_bounded); + } +}; + +REGISTER_XLA_OP(Name("ResamplerGrad"), ResamplerGradOp); + +} // namespace +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/reshape_op.cc b/tensorflow/compiler/tf2xla/kernels/reshape_op.cc index 366ce42866e9f1375ee0ff6f4985c8f461fc0885..fa1b6b91710f5507f41f3f69b0715398ae879aaf 100644 --- a/tensorflow/compiler/tf2xla/kernels/reshape_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/reshape_op.cc @@ -24,6 +24,7 @@ limitations under the License. #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" namespace tensorflow { namespace { @@ -36,7 +37,7 @@ class ReshapeOp : public XlaOpKernel { const TensorShape input_shape = ctx->InputShape(0); const TensorShape sizes_shape = ctx->InputShape(1); // Preliminary validation of sizes. - OP_REQUIRES(ctx, IsLegacyVector(sizes_shape), + OP_REQUIRES(ctx, TensorShapeUtils::IsVector(sizes_shape), errors::InvalidArgument("sizes input must be 1-D, not shape ", sizes_shape.DebugString())); const int64 num_dims = sizes_shape.num_elements(); @@ -95,7 +96,7 @@ class ReshapeOp : public XlaOpKernel { } }; -REGISTER_XLA_OP(Name("Reshape").CompileTimeConstInput("shape"), ReshapeOp); +REGISTER_XLA_OP(Name("Reshape").CompileTimeConstantInput("shape"), ReshapeOp); } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/retval_op.cc b/tensorflow/compiler/tf2xla/kernels/retval_op.cc index e172c649325adb6f7761ce0be141f21e8d545bc1..e4046c795577983bff1a8053743bf4d3a258e583 100644 --- a/tensorflow/compiler/tf2xla/kernels/retval_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/retval_op.cc @@ -13,6 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include "tensorflow/compiler/tf2xla/shape_util.h" #include "tensorflow/compiler/tf2xla/xla_context.h" #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" @@ -46,61 +47,7 @@ class RetvalOp : public XlaOpKernel { // compilation. OP_REQUIRES_OK(ctx, frame->SetRetval(index_, input)); } else { - xla::XlaOp input = ctx->Input(0); - const TensorShape input_shape = ctx->InputShape(0); - DataType input_type = ctx->input_type(0); - XlaContext& tc = XlaContext::Get(ctx); - - if (input_type == DT_RESOURCE) { - XlaResource* resource; - OP_REQUIRES_OK(ctx, ctx->GetResourceInput(0, &resource)); - ctx->SetStatus(tc.AddResourceRetval(index_, resource)); - return; - } - - auto is_constant = ctx->builder()->IsConstant(input); - if (!is_constant.ok()) { - ctx->SetStatus(is_constant.status()); - return; - } - - if (tc.resolve_compile_time_constants() && - (input_shape.num_elements() == 0 || is_constant.ValueOrDie())) { - xla::Literal literal; - OP_REQUIRES_OK(ctx, ctx->ConstantInput(0, &literal)); - OP_REQUIRES_OK(ctx, tc.AddConstRetval(index_, dtype_, literal)); - } else { - TensorShape shape = ctx->InputShape(0); - ctx->SetStatus(is_constant.status()); - TensorShape representation_shape; - if (tc.is_entry_computation()) { - xla::StatusOr shape_or_status = - tc.RepresentationShape(shape, ctx->input_type(0)); - if (!shape_or_status.ok()) { - ctx->SetStatus(shape_or_status.status()); - return; - } else { - representation_shape = shape_or_status.ValueOrDie(); - } - } else { - representation_shape = shape; - } - - xla::XlaOp output = input; - if (tc.is_entry_computation()) { - output = xla::Reshape(input, representation_shape.dim_sizes()); - } else { - // The core from which a return value is returned depends on the - // device assignment of the input to the retval. Since we can't change - // the device assignment of "input" at this point, we must always - // introduce an operator here, even if the shape does not change. - // TODO(b/76097077): propagate device assignments onto arguments and - // return values of functions, and then reshape unconditionally. - output = - xla::GetTupleElement(xla::Tuple(ctx->builder(), {output}), 0); - } - tc.AddRetval(index_, dtype_, shape, output); - } + ctx->xla_context()->SetRetval(index_, ctx->InputExpression(0)); } } diff --git a/tensorflow/compiler/tf2xla/kernels/reverse_op.cc b/tensorflow/compiler/tf2xla/kernels/reverse_op.cc index 8494864b33a44b03a07e3fea7766285f54074e7d..2ceadaf79c5cef35ad50aa84a0d66a46527a6458 100644 --- a/tensorflow/compiler/tf2xla/kernels/reverse_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/reverse_op.cc @@ -51,14 +51,11 @@ class ReverseOp : public XlaOpKernel { } // XlaBuilder::Rev() requires concrete values for dimensions arg. xla::Literal lax; - OP_REQUIRES_OK(ctx, ctx->ConstantInputReshaped(1, {x_shape.dims()}, &lax)); - std::vector revdims(x_shape.dims()); - std::copy(lax.data().begin(), lax.data().end(), - revdims.begin()); - std::vector dimensions; + OP_REQUIRES_OK(ctx, ctx->ConstantInput(1, &lax)); + std::vector dimensions; for (int d = 0; d < x_shape.dims(); ++d) { - if (revdims[d]) { + if (lax.Get({d})) { dimensions.push_back(d); } } @@ -67,7 +64,7 @@ class ReverseOp : public XlaOpKernel { } }; -REGISTER_XLA_OP(Name("Reverse").CompileTimeConstInput("dims"), ReverseOp); +REGISTER_XLA_OP(Name("Reverse").CompileTimeConstantInput("dims"), ReverseOp); class ReverseV2Op : public XlaOpKernel { public: @@ -119,7 +116,8 @@ class ReverseV2Op : public XlaOpKernel { } }; -REGISTER_XLA_OP(Name("ReverseV2").CompileTimeConstInput("axis"), ReverseV2Op); +REGISTER_XLA_OP(Name("ReverseV2").CompileTimeConstantInput("axis"), + ReverseV2Op); } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/reverse_sequence_op.cc b/tensorflow/compiler/tf2xla/kernels/reverse_sequence_op.cc index 03a50ef8a059e5a005c4cc2e5e98acedfea8619a..d7b38e86cc985d608116488f9e76756a8e904f9c 100644 --- a/tensorflow/compiler/tf2xla/kernels/reverse_sequence_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/reverse_sequence_op.cc @@ -17,8 +17,9 @@ 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/numeric.h" +#include "tensorflow/compiler/xla/client/lib/constants.h" #include "tensorflow/compiler/xla/client/xla_builder.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/framework/tensor_shape.h" namespace tensorflow { @@ -61,113 +62,79 @@ class ReverseSequenceOp : public XlaOpKernel { const auto seq_lens = context->Input(1); const int64 batch_size = input_shape.dim_size(batch_dim_); + if (batch_size == 0) { + context->SetOutput(0, input); + return; + } - const DataType input_type = context->input_type(0); - const DataType seq_lens_type = context->input_type(1); + // Given the input + // + // 012345 + // 6789AB + // + // and sequence lens {2, 3} we: + // + // 1. Reverse and pad each row to get + // + // 543210XXXXXX + // BA9876XXXXXX + // + // 2. Gather out the suffix from each row to get + // + // 10XXXX + // 876XXX + // + // 3. Select from the input and the array created by (2) to get the result. + // + // 102345 + // 8769AB + const xla::PrimitiveType input_type = context->input_xla_type(0); + const xla::PrimitiveType seq_lens_type = context->input_xla_type(1); const int64 max_seq_len = input_shape.dim_size(seq_dim_); - xla::Shape input_xla_shape; - OP_REQUIRES_OK(context, TensorShapeToXLAShape(input_type, input_shape, - &input_xla_shape)); - xla::Shape seq_lens_xla_shape; - OP_REQUIRES_OK(context, TensorShapeToXLAShape(seq_lens_type, seq_lens_shape, - &seq_lens_xla_shape)); - - const auto tuple_shape = xla::ShapeUtil::MakeTupleShape({ - xla::ShapeUtil::MakeShape(seq_lens_xla_shape.element_type(), {}), - seq_lens_xla_shape, - input_xla_shape, - }); - - // For each entry in the batch, reverse the sequence. - // TODO(b/65689298): generalize the Map() operator to non-scalar cases and - // use it here, instead of a While loop. - - // Condition: lambda (i, _, _): i < batch_size - auto condition_builder = - builder->CreateSubBuilder("reverse_sequence_condition"); - { - auto param = - xla::Parameter(condition_builder.get(), 0, tuple_shape, "param"); - auto i = xla::GetTupleElement(param, 0); - xla::Lt(i, XlaHelpers::IntegerLiteral(condition_builder.get(), - seq_lens_type, batch_size)); - } - auto condition = condition_builder->Build(); - OP_REQUIRES_OK(context, condition.status()); - - auto body_builder = builder->CreateSubBuilder("reverse_sequence_body"); - { - auto param = xla::Parameter(body_builder.get(), 0, tuple_shape, "param"); - auto i = xla::GetTupleElement(param, 0); - auto seq_lens = xla::GetTupleElement(param, 1); - auto output = xla::GetTupleElement(param, 2); - - // seq_len is the sequence length of the current batch element (rank 1) - auto seq_len = xla::DynamicSlice(seq_lens, xla::Reshape(i, {1}), {1}); - - // Indices is the offset of the batch element in the input. - auto batch_element_indices = - xla::Broadcast(XlaHelpers::Zero(body_builder.get(), seq_lens_type), - {input_shape.dims()}); - batch_element_indices = xla::DynamicUpdateSlice( - batch_element_indices, xla::Reshape(i, {1}), - xla::Reshape(XlaHelpers::IntegerLiteral(body_builder.get(), - seq_lens_type, batch_dim_), - {1})); - - // Slice out the current batch element and pad it out in the sequence - // dimension. - TensorShape slice_shape = input_shape; - slice_shape.set_dim(batch_dim_, 1); - slice_shape.set_dim(seq_dim_, max_seq_len); - auto slice = xla::DynamicSlice(output, batch_element_indices, - slice_shape.dim_sizes()); - auto padding_config = xla::MakeNoPaddingConfig(slice_shape.dims()); - padding_config.mutable_dimensions(seq_dim_)->set_edge_padding_high( - slice_shape.dim_size(seq_dim_)); - slice = xla::Pad(slice, XlaHelpers::Zero(body_builder.get(), input_type), - padding_config); - - // Now slice out the reversed sequence from its actual start. - // sequence_start_indices is the offset of the start of the reversed - // sequence in the input. The slice will go into the padding, however, we - // will mask off these elements and replace them with elements from the - // original input so their values do not matter. - auto sequence_start_indices = - xla::Broadcast(XlaHelpers::Zero(body_builder.get(), seq_lens_type), - {slice_shape.dims()}); - sequence_start_indices = xla::DynamicUpdateSlice( - sequence_start_indices, - xla::Sub(XlaHelpers::IntegerLiteral(body_builder.get(), seq_lens_type, - max_seq_len), - seq_len), - xla::Reshape(XlaHelpers::IntegerLiteral(body_builder.get(), - seq_lens_type, seq_dim_), - {1})); - slice = xla::DynamicSlice(slice, sequence_start_indices, - slice_shape.dim_sizes()); - - // Shift the reversed sequence to the left. - output = xla::DynamicUpdateSlice(output, slice, batch_element_indices); - - xla::Tuple( - body_builder.get(), - {xla::Add(i, XlaHelpers::One(body_builder.get(), seq_lens_type)), - seq_lens, output}); + xla::XlaOp rev = xla::Rev(input, {seq_dim_}); + + auto padding_config = xla::MakeNoPaddingConfig(input_shape.dims()); + padding_config.mutable_dimensions(seq_dim_)->set_edge_padding_high( + max_seq_len); + xla::XlaOp padded = + xla::Pad(rev, xla::Zero(builder, input_type), padding_config); + + // Form a start indices tensor with shape [2, batch_size]. For each batch + // entry we have a (batch offset, seq offset) pair. + xla::XlaOp start_indices = xla::ConcatInDim( + builder, + { + xla::Iota(builder, + xla::ShapeUtil::MakeShape(seq_lens_type, {1, batch_size}), + /*iota_dimension=*/1), + xla::Reshape(xla::ScalarLike(seq_lens, max_seq_len) - seq_lens, + {1, batch_size}), + }, + /*dimension=*/0); + + xla::GatherDimensionNumbers dnums; + // The first dimension of start_indices contains the batch/seq dim choice. + dnums.set_index_vector_dim(0); + dnums.add_start_index_map(batch_dim_); + dnums.add_start_index_map(seq_dim_); + + // All other dimensions other than the batch dim are offset dimensions. + for (int i = 0; i < input_shape.dims(); ++i) { + if (i != batch_dim_) { + dnums.add_offset_dims(i); + } } - auto body = body_builder->Build(); - OP_REQUIRES_OK(context, body.status()); - - auto loop_output = xla::While( - condition.ValueOrDie(), body.ValueOrDie(), - xla::Tuple(builder, {XlaHelpers::Zero(builder, seq_lens_type), seq_lens, - xla::Rev(input, {seq_dim_})})); - auto output = xla::GetTupleElement(loop_output, 2); - - // Mask out elements after the sequence length. - xla::XlaOp iota = - xla::Iota(builder, seq_lens_xla_shape.element_type(), max_seq_len); + dnums.add_collapsed_slice_dims(batch_dim_); + + auto slice_sizes = input_shape.dim_sizes(); + slice_sizes[batch_dim_] = 1; + + xla::XlaOp output = xla::Gather(padded, start_indices, dnums, slice_sizes); + + // Mask out elements after the sequence length, and copy the corresponding + // elements from the input. + xla::XlaOp iota = xla::Iota(builder, seq_lens_type, max_seq_len); std::vector dims(input_shape.dims(), 1); dims[batch_dim_] = batch_size; auto mask = xla::Lt(iota, xla::Reshape(seq_lens, dims), {seq_dim_}); diff --git a/tensorflow/compiler/tf2xla/kernels/scan_ops.cc b/tensorflow/compiler/tf2xla/kernels/scan_ops.cc index 57afd608de820573821d605cadcc8779474b5fd6..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" @@ -39,8 +39,8 @@ namespace { // TODO(phawkins): implement double-sized windowed reductions in XLA and remove // the type constraint. -constexpr std::array kScanOpTypes = { - {DT_HALF, DT_BFLOAT16, DT_FLOAT}}; +constexpr std::array kScanOpTypes = { + {DT_HALF, DT_BFLOAT16, DT_FLOAT, DT_INT32}}; class ScanOp : public XlaOpKernel { public: @@ -103,11 +103,10 @@ class ScanOp : public XlaOpKernel { reducer = ctx->GetOrCreateMul(dtype); } auto output = xla::ReduceWindowWithGeneralPadding( - XlaHelpers::ConvertElementType(builder, ctx->Input(0), dtype), init, - *reducer, window_dims, window_strides, + XlaHelpers::ConvertElementType(ctx->Input(0), dtype), init, *reducer, + window_dims, window_strides, /*base_dilations=*/{}, /*window_dilations=*/{}, padding); - output = - XlaHelpers::ConvertElementType(builder, output, ctx->input_type(0)); + output = XlaHelpers::ConvertElementType(output, ctx->input_type(0)); // In exclusive mode, we have computed an extra element containing the sum // of all the input elements. Slice off this extra "last" element. @@ -136,7 +135,7 @@ class CumsumOp : public ScanOp { }; REGISTER_XLA_OP(Name("Cumsum") .TypeConstraint("T", kScanOpTypes) - .CompileTimeConstInput("axis"), + .CompileTimeConstantInput("axis"), CumsumOp); class CumprodOp : public ScanOp { @@ -145,7 +144,7 @@ class CumprodOp : public ScanOp { }; REGISTER_XLA_OP(Name("Cumprod") .TypeConstraint("T", kScanOpTypes) - .CompileTimeConstInput("axis"), + .CompileTimeConstantInput("axis"), CumprodOp); } // anonymous namespace diff --git a/tensorflow/compiler/tf2xla/kernels/scatter_nd_op.cc b/tensorflow/compiler/tf2xla/kernels/scatter_nd_op.cc index f1f32699fee5f03f603f830722fe65622dee5d3e..a95e7adacf194ba6eb33cbeb56abe1a5a2479337 100644 --- a/tensorflow/compiler/tf2xla/kernels/scatter_nd_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/scatter_nd_op.cc @@ -116,7 +116,8 @@ class ScatterNdOp : public XlaOpKernel { } }; -REGISTER_XLA_OP(Name("ScatterNd").CompileTimeConstInput("shape"), ScatterNdOp); +REGISTER_XLA_OP(Name("ScatterNd").CompileTimeConstantInput("shape"), + ScatterNdOp); } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/segment_reduction_ops.cc b/tensorflow/compiler/tf2xla/kernels/segment_reduction_ops.cc index b22ecb7c6dbb42a33a4f4d90b18b20816df16a50..97359f81eee4aa0b46f03941ab6ca3ea3d468f1f 100644 --- a/tensorflow/compiler/tf2xla/kernels/segment_reduction_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/segment_reduction_ops.cc @@ -105,7 +105,7 @@ class UnsortedSegmentSum : public UnsortedSegmentReduce { }; REGISTER_XLA_OP( - Name("UnsortedSegmentSum").CompileTimeConstInput("num_segments"), + Name("UnsortedSegmentSum").CompileTimeConstantInput("num_segments"), UnsortedSegmentSum); class UnsortedSegmentProd : public UnsortedSegmentReduce { @@ -120,7 +120,7 @@ class UnsortedSegmentProd : public UnsortedSegmentReduce { }; REGISTER_XLA_OP( - Name("UnsortedSegmentProd").CompileTimeConstInput("num_segments"), + Name("UnsortedSegmentProd").CompileTimeConstantInput("num_segments"), UnsortedSegmentProd); class UnsortedSegmentMin : public UnsortedSegmentReduce { @@ -137,7 +137,7 @@ class UnsortedSegmentMin : public UnsortedSegmentReduce { }; REGISTER_XLA_OP( - Name("UnsortedSegmentMin").CompileTimeConstInput("num_segments"), + Name("UnsortedSegmentMin").CompileTimeConstantInput("num_segments"), UnsortedSegmentMin); class UnsortedSegmentMax : public UnsortedSegmentReduce { @@ -154,7 +154,7 @@ class UnsortedSegmentMax : public UnsortedSegmentReduce { }; REGISTER_XLA_OP( - Name("UnsortedSegmentMax").CompileTimeConstInput("num_segments"), + Name("UnsortedSegmentMax").CompileTimeConstantInput("num_segments"), UnsortedSegmentMax); } // namespace 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/sendrecv_ops.cc b/tensorflow/compiler/tf2xla/kernels/sendrecv_ops.cc index a7f5a8f1698b9d02560de427d356e9e6be5caa7c..84470b230d421658e0d79dcecb175a24155f49b7 100644 --- a/tensorflow/compiler/tf2xla/kernels/sendrecv_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/sendrecv_ops.cc @@ -42,7 +42,7 @@ SendOp::SendOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { } void SendOp::Compile(XlaOpKernelContext* ctx) { - XlaCompiler* compiler = XlaContext::Get(ctx).compiler(); + XlaCompiler* compiler = ctx->compiler(); xla::ChannelHandle channel; OP_REQUIRES_OK(ctx, compiler->GetChannelHandle(tensor_name_, &channel)); xla::Send(ctx->Input(0), channel); @@ -73,7 +73,7 @@ RecvOp::RecvOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { } void RecvOp::Compile(XlaOpKernelContext* ctx) { - XlaCompiler* compiler = XlaContext::Get(ctx).compiler(); + XlaCompiler* compiler = ctx->compiler(); xla::ChannelHandle channel; OP_REQUIRES_OK(ctx, compiler->GetChannelHandle(tensor_name_, &channel)); ctx->SetOutput(0, xla::Recv(ctx->builder(), shape_, channel)); diff --git a/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc b/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc index 0c32b8def0f7b741c93e803f8359b6504087e257..7a620d2a6518f8686ef570b33aac971d1dccb6c1 100644 --- a/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc @@ -18,7 +18,7 @@ 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/numeric.h" +#include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/compiler/xla/literal.h" #include "tensorflow/compiler/xla/primitive_util.h" #include "tensorflow/core/framework/op_kernel.h" @@ -30,31 +30,6 @@ limitations under the License. namespace tensorflow { namespace { -template -Status GetValue(int index, XlaOpKernelContext* ctx, T* value) { - xla::Literal literal; - TF_RETURN_IF_ERROR(ctx->ConstantInput(index, &literal)); - *value = literal.Get({}); - return Status::OK(); -} - -Status GetIntValue(int index, XlaOpKernelContext* ctx, int64* value) { - xla::Literal literal; - TF_RETURN_IF_ERROR(ctx->ConstantInput(index, &literal)); - switch (literal.shape().element_type()) { - case xla::S32: - *value = literal.Get({}); - break; - case xla::S64: - *value = literal.Get({}); - break; - default: - return errors::InvalidArgument("Invalid argument type for argument", - index); - } - return Status::OK(); -} - // The type-specific part of the implementation of Range. template xla::StatusOr CreateRangeTensor( @@ -98,13 +73,13 @@ class RangeOp : public XlaOpKernel { const TensorShape start_in_shape = ctx->InputShape(0); const TensorShape limit_in_shape = ctx->InputShape(1); const TensorShape delta_in_shape = ctx->InputShape(2); - OP_REQUIRES(ctx, IsLegacyScalar(start_in_shape), + OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(start_in_shape), errors::InvalidArgument("start must be a scalar, not shape ", start_in_shape.DebugString())); - OP_REQUIRES(ctx, IsLegacyScalar(limit_in_shape), + OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(limit_in_shape), errors::InvalidArgument("limit must be a scalar, not shape ", limit_in_shape.DebugString())); - OP_REQUIRES(ctx, IsLegacyScalar(delta_in_shape), + OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(delta_in_shape), errors::InvalidArgument("delta must be a scalar, not shape ", delta_in_shape.DebugString())); xla::Literal start, limit, delta; @@ -137,9 +112,9 @@ class RangeOp : public XlaOpKernel { }; REGISTER_XLA_OP(Name("Range") - .CompileTimeConstInput("start") - .CompileTimeConstInput("limit") - .CompileTimeConstInput("delta"), + .CompileTimeConstantInput("start") + .CompileTimeConstantInput("limit") + .CompileTimeConstantInput("delta"), RangeOp); class LinSpaceOp : public XlaOpKernel { @@ -147,9 +122,9 @@ class LinSpaceOp : public XlaOpKernel { explicit LinSpaceOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} void Compile(XlaOpKernelContext* ctx) override { - const TensorShape start_in_shape = ctx->InputShape(0); - const TensorShape stop_in_shape = ctx->InputShape(1); - const TensorShape num_in_shape = ctx->InputShape(2); + const TensorShape start_in_shape = ctx->InputShape("start"); + const TensorShape stop_in_shape = ctx->InputShape("stop"); + const TensorShape num_in_shape = ctx->InputShape("num"); OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(start_in_shape), errors::InvalidArgument("start must be a scalar, not shape ", start_in_shape.DebugString())); @@ -163,39 +138,46 @@ class LinSpaceOp : public XlaOpKernel { DataType type = ctx->input_type(0); int64 num; - OP_REQUIRES_OK(ctx, GetIntValue(2, ctx, &num)); + OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntScalar("num", &num)); OP_REQUIRES(ctx, num > 0, errors::InvalidArgument("Requires num > 0: ", num)); Tensor out_constant(type, TensorShape({num})); + xla::Literal start_literal; + OP_REQUIRES_OK(ctx, ctx->ConstantInput("start", &start_literal)); + xla::Literal stop_literal; + OP_REQUIRES_OK(ctx, ctx->ConstantInput("stop", &stop_literal)); + switch (type) { case DT_FLOAT: { - float start, stop; - OP_REQUIRES_OK(ctx, GetValue(0, ctx, &start)); - OP_REQUIRES_OK(ctx, GetValue(1, ctx, &stop)); + float start = start_literal.GetFirstElement(); + float stop = stop_literal.GetFirstElement(); auto flat = out_constant.flat(); if (num == 1) { 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; } case DT_DOUBLE: { - double start, stop; - OP_REQUIRES_OK(ctx, GetValue(0, ctx, &start)); - OP_REQUIRES_OK(ctx, GetValue(1, ctx, &stop)); + double start = start_literal.GetFirstElement(); + double stop = stop_literal.GetFirstElement(); auto flat = out_constant.flat(); if (num == 1) { 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; } @@ -210,9 +192,9 @@ class LinSpaceOp : public XlaOpKernel { }; REGISTER_XLA_OP(Name("LinSpace") - .CompileTimeConstInput("start") - .CompileTimeConstInput("stop") - .CompileTimeConstInput("num"), + .CompileTimeConstantInput("start") + .CompileTimeConstantInput("stop") + .CompileTimeConstantInput("num"), LinSpaceOp); } // namespace diff --git a/tensorflow/compiler/tf2xla/kernels/shape_op.cc b/tensorflow/compiler/tf2xla/kernels/shape_op.cc index c8a0f31a0375abacaca26688a23f4835e11c692e..31d4cc131600f360c764ffa02831046c85d846e5 100644 --- a/tensorflow/compiler/tf2xla/kernels/shape_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/shape_op.cc @@ -20,9 +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/kernels/bounds_check.h" +#include "tensorflow/core/framework/tensor_shape.h" namespace tensorflow { namespace { @@ -90,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); } }; @@ -108,21 +116,16 @@ class ExpandDimsOp : public XlaOpKernel { explicit ExpandDimsOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} void Compile(XlaOpKernelContext* ctx) override { - const TensorShape input_shape = ctx->InputShape(0); - const TensorShape dim_shape = ctx->InputShape(1); + const TensorShape input_shape = ctx->InputShape("input"); + const TensorShape dim_shape = ctx->InputShape("dim"); - // TODO(phawkins): the standard implementation of ExpandDimsOp seems to - // accept legacy scalars, even when they should be forbidden by the graphdef - // version. - OP_REQUIRES(ctx, dim_shape.num_elements() == 1, + std::vector dims; + OP_REQUIRES_OK(ctx, ctx->ConstantInputReshapedToIntVector("dim", &dims)); + OP_REQUIRES(ctx, dims.size() == 1, errors::InvalidArgument(absl::StrCat( "dim input to ExpandDims must be a scalar; got ", dim_shape.DebugString()))); - - xla::Literal literal; - OP_REQUIRES_OK(ctx, ctx->ConstantInputReshaped(1, {1}, &literal)); - - int dim = literal.data()[0]; + int dim = dims[0]; OP_REQUIRES(ctx, (dim >= -1 - input_shape.dims() && dim <= input_shape.dims()), @@ -148,10 +151,11 @@ class ExpandDimsOp : public XlaOpKernel { dim = std::min(dim, existing_dims_size); new_shape.emplace(new_shape.begin() + dim, 1); - ctx->SetOutput(0, xla::Reshape(ctx->Input(0), new_shape)); + ctx->SetOutput(0, xla::Reshape(ctx->Input("input"), new_shape)); } }; -REGISTER_XLA_OP(Name("ExpandDims").CompileTimeConstInput("dim"), ExpandDimsOp); +REGISTER_XLA_OP(Name("ExpandDims").CompileTimeConstantInput("dim"), + ExpandDimsOp); class SqueezeOp : public XlaOpKernel { public: 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/slice_op.cc b/tensorflow/compiler/tf2xla/kernels/slice_op.cc index 537b71f3c0cf3622a8a45a717ac406da69f5c3c7..88da64e5a217a0c026106f03cb26958f6738446c 100644 --- a/tensorflow/compiler/tf2xla/kernels/slice_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/slice_op.cc @@ -24,6 +24,7 @@ limitations under the License. #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/ops_util.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/platform/mem.h" @@ -42,8 +43,8 @@ class SliceOp : public XlaOpKernel { OP_REQUIRES( ctx, - IsLegacyVector(begin_tensor_shape) && - IsLegacyVector(size_tensor_shape) && + TensorShapeUtils::IsVector(begin_tensor_shape) && + TensorShapeUtils::IsVector(size_tensor_shape) && begin_tensor_shape.num_elements() == input_shape.dims() && size_tensor_shape.num_elements() == input_shape.dims(), errors::InvalidArgument( @@ -111,9 +112,10 @@ class SliceOp : public XlaOpKernel { } }; -REGISTER_XLA_OP( - Name("Slice").CompileTimeConstInput("begin").CompileTimeConstInput("size"), - SliceOp); +REGISTER_XLA_OP(Name("Slice") + .CompileTimeConstantInput("begin") + .CompileTimeConstantInput("size"), + SliceOp); } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/softmax_op.cc b/tensorflow/compiler/tf2xla/kernels/softmax_op.cc index d6bd927135c013ac1ec3f6547aef358dc2741896..20da8033536e3af3da0fcb216db45f808cacc1d5 100644 --- a/tensorflow/compiler/tf2xla/kernels/softmax_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/softmax_op.cc @@ -71,7 +71,7 @@ class SoftmaxOp : public XlaOpKernel { auto reduce = xla::Reduce(converted, xla::Zero(b, xla_accumulation_type), *ctx->GetOrCreateAdd(accumulation_type), {kClassDim}); - auto sum = XlaHelpers::ConvertElementType(b, reduce, type); + auto sum = XlaHelpers::ConvertElementType(reduce, type); auto softmax = log_ // softmax = shifted_logits - log(sum(exp(shifted_logits))) @@ -111,11 +111,11 @@ std::pair CrossEntropyWithLogits( // sum_{class} (exp(logits - max_logits)) const DataType accumulation_type = XlaHelpers::SumAccumulationType(type); auto converted = - XlaHelpers::ConvertElementType(b, exp_shifted_logits, accumulation_type); + XlaHelpers::ConvertElementType(exp_shifted_logits, accumulation_type); auto reduce = xla::Reduce(converted, XlaHelpers::Zero(b, accumulation_type), *ctx->GetOrCreateAdd(accumulation_type), {kClassDim}); - auto sum_exp = XlaHelpers::ConvertElementType(b, reduce, type); + auto sum_exp = XlaHelpers::ConvertElementType(reduce, type); // log(sum(exp(logits - max_logits))) auto log_sum_exp = xla::Log(sum_exp); @@ -126,11 +126,10 @@ std::pair CrossEntropyWithLogits( // (The subtraction broadcasts along the batch dimension.) auto sub = xla::Sub(shifted_logits, log_sum_exp, {kBatchDim}); auto mul = xla::Mul(xla::Neg(labels), sub); - auto sum = - xla::Reduce(XlaHelpers::ConvertElementType(b, mul, accumulation_type), - XlaHelpers::Zero(b, accumulation_type), - *ctx->GetOrCreateAdd(accumulation_type), {kClassDim}); - auto loss = XlaHelpers::ConvertElementType(b, sum, type); + auto sum = xla::Reduce(XlaHelpers::ConvertElementType(mul, accumulation_type), + XlaHelpers::Zero(b, accumulation_type), + *ctx->GetOrCreateAdd(accumulation_type), {kClassDim}); + auto loss = XlaHelpers::ConvertElementType(sum, type); // backprop: prob - labels, where // prob = exp(logits - max_logits) / sum(exp(logits - max_logits)) diff --git a/tensorflow/compiler/tf2xla/kernels/spacetobatch_op.cc b/tensorflow/compiler/tf2xla/kernels/spacetobatch_op.cc index 76b79be6f6f6b5ecbe9edcffb81f2834fdac9a56..52bed2670b4b8408e3b2f72b64bf370aea5325f6 100644 --- a/tensorflow/compiler/tf2xla/kernels/spacetobatch_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/spacetobatch_op.cc @@ -39,7 +39,7 @@ void SpaceToBatch(XlaOpKernelContext* ctx, const xla::XlaOp& input, OP_REQUIRES( ctx, - xla::ShapeUtil::Rank(paddings.shape()) == 2 && + paddings.shape().rank() == 2 && block_rank == xla::ShapeUtil::GetDimension(paddings.shape(), 0) && 2 == xla::ShapeUtil::GetDimension(paddings.shape(), 1), errors::InvalidArgument("paddings should have shape [", block_rank, @@ -161,8 +161,8 @@ class SpaceToBatchNDOp : public XlaOpKernel { } }; REGISTER_XLA_OP(Name("SpaceToBatchND") - .CompileTimeConstInput("paddings") - .CompileTimeConstInput("block_shape"), + .CompileTimeConstantInput("paddings") + .CompileTimeConstantInput("block_shape"), SpaceToBatchNDOp); class SpaceToBatchOp : public XlaOpKernel { @@ -185,7 +185,7 @@ class SpaceToBatchOp : public XlaOpKernel { private: int block_size_; }; -REGISTER_XLA_OP(Name("SpaceToBatch").CompileTimeConstInput("paddings"), +REGISTER_XLA_OP(Name("SpaceToBatch").CompileTimeConstantInput("paddings"), SpaceToBatchOp); } // namespace diff --git a/tensorflow/compiler/tf2xla/kernels/sparse_to_dense_op.cc b/tensorflow/compiler/tf2xla/kernels/sparse_to_dense_op.cc index e831dc30a9d3c27ec3b1494e7d8a6de836ff2a11..def3c147bf3fc619784044357e95bf32b404954b 100644 --- a/tensorflow/compiler/tf2xla/kernels/sparse_to_dense_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/sparse_to_dense_op.cc @@ -80,7 +80,7 @@ class SparseToDenseOp : public XlaOpKernel { } }; -REGISTER_XLA_OP(Name("SparseToDense").CompileTimeConstInput("output_shape"), +REGISTER_XLA_OP(Name("SparseToDense").CompileTimeConstantInput("output_shape"), SparseToDenseOp); } // namespace diff --git a/tensorflow/compiler/tf2xla/kernels/split_op.cc b/tensorflow/compiler/tf2xla/kernels/split_op.cc index 93fc14e9efca868e84444dd0e07d7f0dfa84c042..7a0e240400b344ab25743997ce3baad81bd5f476 100644 --- a/tensorflow/compiler/tf2xla/kernels/split_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/split_op.cc @@ -35,26 +35,16 @@ class SplitOp : public XlaOpKernel { void Compile(XlaOpKernelContext* ctx) override { const int32 num_split = num_outputs(); - const TensorShape index_shape = ctx->InputShape(0); + const TensorShape split_dim_shape = ctx->InputShape("split_dim"); const TensorShape input_shape = ctx->InputShape(1); - xla::Literal literal_index; - OP_REQUIRES_OK(ctx, ctx->ConstantInput(0, &literal_index)); - - int32 split_dim_orig; - if (index_shape.dims() == 0) { - split_dim_orig = literal_index.Get({}); - } else { - OP_REQUIRES( - ctx, index_shape.dims() == 1, - errors::InvalidArgument("split_index input to Split Op must be a " - "scalar or a vector with 1 element")); - OP_REQUIRES( - ctx, index_shape.dim_size(0) == 1, - errors::InvalidArgument("split_index input to Split Op must be a " - "scalar or a vector with 1 element")); - split_dim_orig = literal_index.Get({0}); - } + OP_REQUIRES( + ctx, TensorShapeUtils::IsScalar(split_dim_shape), + errors::InvalidArgument("split_dim must be a scalar but has rank ", + split_dim_shape.dims())); + int64 split_dim_orig; + OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntScalar(0, &split_dim_orig)); + int32 split_dim = split_dim_orig < 0 ? split_dim_orig + input_shape.dims() : split_dim_orig; OP_REQUIRES(ctx, 0 <= split_dim && split_dim < input_shape.dims(), @@ -104,7 +94,7 @@ class SplitOp : public XlaOpKernel { } }; -REGISTER_XLA_OP(Name("Split").CompileTimeConstInput("split_dim"), SplitOp); +REGISTER_XLA_OP(Name("Split").CompileTimeConstantInput("split_dim"), SplitOp); class SplitVOp : public XlaOpKernel { public: @@ -138,7 +128,6 @@ class SplitVOp : public XlaOpKernel { // Check that sizes are correct. int total_split_size = 0; int neg_one_dim = -1; - std::vector split_sizes_vec(num_split, -1); const TensorShape split_size_shape = ctx->InputShape(1); OP_REQUIRES(ctx, split_size_shape.dims() == 1 && @@ -150,12 +139,11 @@ class SplitVOp : public XlaOpKernel { split_size_shape.dims(), "-D and ", split_size_shape.num_elements(), " elements")); // Get the dimension of this split. - xla::Literal split_size_literal; - OP_REQUIRES_OK(ctx, ctx->ConstantInput(1, &split_size_literal)); + std::vector split_sizes; + OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector(1, &split_sizes)); for (int i = 0; i < num_split; ++i) { - int slice_size; - slice_size = split_size_literal.Get({i}); + int64 slice_size = split_sizes[i]; if (slice_size == -1) { OP_REQUIRES( ctx, neg_one_dim == -1, @@ -164,7 +152,6 @@ class SplitVOp : public XlaOpKernel { i)); neg_one_dim = i; } else { - split_sizes_vec[i] = slice_size; total_split_size += slice_size; } } @@ -183,7 +170,7 @@ class SplitVOp : public XlaOpKernel { total_split_size)); if (neg_one_dim >= 0) { - split_sizes_vec[neg_one_dim] = + split_sizes[neg_one_dim] = input_shape.dim_size(split_dim) - total_split_size; } @@ -195,7 +182,7 @@ class SplitVOp : public XlaOpKernel { std::vector strides(input_shape.dims(), 1); for (int i = 0; i < num_split; ++i) { TensorShape output_shape(input_shape); - int slice_size = split_sizes_vec[i]; + int slice_size = split_sizes[i]; output_shape.set_dim(split_dim, slice_size); // Slice out the ith split from the split dimension. @@ -207,8 +194,8 @@ class SplitVOp : public XlaOpKernel { }; REGISTER_XLA_OP(Name("SplitV") - .CompileTimeConstInput("split_dim") - .CompileTimeConstInput("size_splits"), + .CompileTimeConstantInput("split_dim") + .CompileTimeConstantInput("size_splits"), SplitVOp); } // namespace diff --git a/tensorflow/compiler/tf2xla/kernels/stack_ops.cc b/tensorflow/compiler/tf2xla/kernels/stack_ops.cc index ee70f508a9586d5f47bd7bb7670506d4c92b369f..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" @@ -45,7 +45,7 @@ Status GetStackShape(xla::XlaBuilder* builder, XlaResource* resource, return shape_or_status.status(); } xla::Shape shape = shape_or_status.ValueOrDie(); - TF_RET_CHECK(xla::ShapeUtil::IsTuple(shape)); + TF_RET_CHECK(shape.IsTuple()); return XLAShapeToTensorShape(xla::ShapeUtil::GetTupleElementShape(shape, 0), stack_shape); } @@ -69,7 +69,7 @@ Status MaybeInitializeStack(xla::XlaBuilder* builder, XlaResource* resource, } TensorShape stack_shape; - stack_shape.AddDim(resource->tensor_array_size()); + stack_shape.AddDim(resource->max_array_size()); stack_shape.AppendShape(elem_shape); if (!resource->initialized()) { @@ -97,10 +97,10 @@ class StackOp : public XlaOpKernel { } void Compile(XlaOpKernelContext* ctx) override { - int64 size; - OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntScalar(0, &size)); + int64 max_size; + OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntScalar(0, &max_size)); OP_REQUIRES( - ctx, size >= 0, + ctx, max_size >= 0, errors::InvalidArgument( "XLA compilation requires a fixed stack size upper bound. If " "you are using tf.while_loop, set the maximum_iterations parameter " @@ -108,14 +108,9 @@ class StackOp : public XlaOpKernel { // We defer initializing the Stack resource until we see the first push. // Otherwise we do not know the shape of the stack elements. - xla::XlaOp value; - XlaContext& xc = XlaContext::Get(ctx); - XlaResource* resource; - string name = absl::StrCat("Stack: ", stack_name_); - OP_REQUIRES_OK( - ctx, xc.CreateResource(XlaResource::kStack, -1, std::move(name), dtype_, - TensorShape(), value, /*tensor_array_size=*/size, - /*tensor_array_gradients=*/{}, &resource)); + XlaResource* resource = + ctx->xla_context()->AddResource(XlaResource::CreateStack( + /*name=*/absl::StrCat("Stack: ", stack_name_), dtype_, max_size)); ctx->SetResourceOutput(0, resource); } @@ -126,7 +121,9 @@ class StackOp : public XlaOpKernel { TF_DISALLOW_COPY_AND_ASSIGN(StackOp); }; -REGISTER_XLA_OP(Name("StackV2").CompileTimeConstInput("max_size"), StackOp); +REGISTER_XLA_OP( + Name("StackV2").CompileTimeConstantInput("max_size").CompilationOnly(), + StackOp); class StackPushOp : public XlaOpKernel { public: @@ -149,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); @@ -173,7 +170,7 @@ class StackPushOp : public XlaOpKernel { TF_DISALLOW_COPY_AND_ASSIGN(StackPushOp); }; -REGISTER_XLA_OP(Name("StackPushV2"), StackPushOp); +REGISTER_XLA_OP(Name("StackPushV2").CompilationOnly(), StackPushOp); class StackPopOp : public XlaOpKernel { public: @@ -205,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; @@ -227,7 +224,7 @@ class StackPopOp : public XlaOpKernel { TF_DISALLOW_COPY_AND_ASSIGN(StackPopOp); }; -REGISTER_XLA_OP(Name("StackPopV2"), StackPopOp); +REGISTER_XLA_OP(Name("StackPopV2").CompilationOnly(), StackPopOp); class StackCloseOp : public XlaOpKernel { public: @@ -241,7 +238,7 @@ class StackCloseOp : public XlaOpKernel { TF_DISALLOW_COPY_AND_ASSIGN(StackCloseOp); }; -REGISTER_XLA_OP(Name("StackCloseV2"), StackCloseOp); +REGISTER_XLA_OP(Name("StackCloseV2").CompilationOnly(), StackCloseOp); } // anonymous namespace } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/stateless_random_ops.cc b/tensorflow/compiler/tf2xla/kernels/stateless_random_ops.cc index 5412e135478361d08965e4621ec52cfb4a792f1d..50653d7b3973b73d580cdeec5d71943b575d7cc9 100644 --- a/tensorflow/compiler/tf2xla/kernels/stateless_random_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/stateless_random_ops.cc @@ -17,27 +17,43 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/lib/random.h" #include "tensorflow/compiler/tf2xla/shape_util.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/constants.h" #include "tensorflow/compiler/xla/client/lib/math.h" -#include "tensorflow/compiler/xla/client/lib/numeric.h" #include "tensorflow/compiler/xla/client/lib/prng.h" #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_shape.h" -#include "tensorflow/core/lib/core/casts.h" #include "tensorflow/core/lib/math/math_util.h" namespace tensorflow { namespace { +xla::XlaOp MaybeConvertF32ToBF16(xla::XlaOp input, DataType dtype) { + // Mask the last 16 bit. With normal rounding, values near "maxval" would be + // converted to "maxval" which is out of range ["minval", "maxval"). In + // addition, the distribution near the limit is not uniform. + if (dtype == DT_BFLOAT16) { + xla::XlaBuilder* builder = input.builder(); + auto output = xla::BitcastConvertType(input, xla::U32) & + xla::ConstantR0(builder, 0xFFFF0000); + return xla::ConvertElementType(xla::BitcastConvertType(output, xla::F32), + xla::BF16); + } else { + return input; + } +} + class StatelessRandomUniformOp : public XlaOpKernel { public: explicit StatelessRandomUniformOp(OpKernelConstruction* ctx) - : XlaOpKernel(ctx) {} + : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("dtype", &dtype_)); + } void Compile(XlaOpKernelContext* ctx) override { xla::XlaBuilder* builder = ctx->builder(); @@ -60,24 +76,81 @@ class StatelessRandomUniformOp : public XlaOpKernel { auto uniform = xla::StatelessRngUniform( {seed0, seed1}, xla_shape, xla::ConstantR0(builder, 0.0), xla::ConstantR0(builder, 1.0)); + uniform = MaybeConvertF32ToBF16(uniform, dtype_); ctx->SetOutput(0, uniform); } private: + DataType dtype_; + TF_DISALLOW_COPY_AND_ASSIGN(StatelessRandomUniformOp); }; // TODO(phawkins): generalize to non-float, non-int32 seed types. REGISTER_XLA_OP(Name("StatelessRandomUniform") - .CompileTimeConstInput("shape") - .TypeConstraint("dtype", DT_FLOAT) + .CompileTimeConstantInput("shape") + .TypeConstraint("dtype", {DT_FLOAT, DT_BFLOAT16}) .TypeConstraint("Tseed", DT_INT32), StatelessRandomUniformOp); +class StatelessRandomUniformIntOp : public XlaOpKernel { + public: + explicit StatelessRandomUniformIntOp(OpKernelConstruction* ctx) + : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("dtype", &dtype_)); + } + + void Compile(XlaOpKernelContext* ctx) override { + TensorShape shape; + OP_REQUIRES_OK(ctx, ctx->ConstantInputAsShape(0, &shape)); + + TensorShape seed_shape = ctx->InputShape(1); + OP_REQUIRES(ctx, seed_shape.dims() == 1 && seed_shape.dim_size(0) == 2, + errors::InvalidArgument("seed must have shape [2], not ", + seed_shape.DebugString())); + TensorShape minval_shape = ctx->InputShape(2); + OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(minval_shape), + errors::InvalidArgument("minval must be scalar, got shape ", + minval_shape.DebugString())); + TensorShape maxval_shape = ctx->InputShape(3); + OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(maxval_shape), + errors::InvalidArgument("minval must be scalar, got shape ", + maxval_shape.DebugString())); + + xla::XlaOp seed = ctx->Input(1); + xla::XlaOp minval = ctx->Input(2); + xla::XlaOp maxval = ctx->Input(3); + + xla::Shape xla_shape; + OP_REQUIRES_OK(ctx, TensorShapeToXLAShape(dtype_, shape, &xla_shape)); + + auto seed0 = xla::Reshape(xla::Slice(seed, {0}, {1}, {1}), {}); + auto seed1 = xla::Reshape(xla::Slice(seed, {1}, {2}, {1}), {}); + + auto uniform = + xla::StatelessRngUniform({seed0, seed1}, xla_shape, minval, maxval); + ctx->SetOutput(0, uniform); + } + + private: + DataType dtype_; + + TF_DISALLOW_COPY_AND_ASSIGN(StatelessRandomUniformIntOp); +}; + +// TODO(phawkins): generalize to non-int32 seed types. +REGISTER_XLA_OP(Name("StatelessRandomUniformInt") + .CompileTimeConstantInput("shape") + .TypeConstraint("dtype", {DT_INT32, DT_INT64}) + .TypeConstraint("Tseed", DT_INT32), + StatelessRandomUniformIntOp); + class StatelessRandomNormalOp : public XlaOpKernel { public: explicit StatelessRandomNormalOp(OpKernelConstruction* ctx) - : XlaOpKernel(ctx) {} + : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("dtype", &dtype_)); + } void Compile(XlaOpKernelContext* ctx) override { TensorShape shape; @@ -103,24 +176,29 @@ class StatelessRandomNormalOp : public XlaOpKernel { // sqrt(2) * erfinv(x) auto normal = xla::ScalarLike(uniform, std::sqrt(2.0)) * xla::ErfInv(uniform); + normal = MaybeConvertF32ToBF16(normal, dtype_); ctx->SetOutput(0, normal); } private: + DataType dtype_; + TF_DISALLOW_COPY_AND_ASSIGN(StatelessRandomNormalOp); }; // TODO(phawkins): generalize to non-float, non-int32 seed types. REGISTER_XLA_OP(Name("StatelessRandomNormal") - .CompileTimeConstInput("shape") - .TypeConstraint("dtype", DT_FLOAT) + .CompileTimeConstantInput("shape") + .TypeConstraint("dtype", {DT_FLOAT, DT_BFLOAT16}) .TypeConstraint("Tseed", DT_INT32), StatelessRandomNormalOp); class StatelessTruncatedNormalOp : public XlaOpKernel { public: explicit StatelessTruncatedNormalOp(OpKernelConstruction* ctx) - : XlaOpKernel(ctx) {} + : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("dtype", &dtype_)); + } void Compile(XlaOpKernelContext* ctx) override { TensorShape shape; @@ -142,17 +220,20 @@ class StatelessTruncatedNormalOp : public XlaOpKernel { {seed0, seed1}, xla_shape, xla::ConstantR0(builder, std::numeric_limits::min()), xla::ConstantR0(builder, 1.0)); - - ctx->SetOutput(0, TruncatedNormal(uniform)); + auto output = TruncatedNormal(uniform); + output = MaybeConvertF32ToBF16(output, dtype_); + ctx->SetOutput(0, output); } private: + DataType dtype_; + TF_DISALLOW_COPY_AND_ASSIGN(StatelessTruncatedNormalOp); }; REGISTER_XLA_OP(Name("StatelessTruncatedNormal") - .CompileTimeConstInput("shape") - .TypeConstraint("dtype", DT_FLOAT) + .CompileTimeConstantInput("shape") + .TypeConstraint("dtype", {DT_FLOAT, DT_BFLOAT16}) .TypeConstraint("Tseed", DT_INT32), StatelessTruncatedNormalOp); diff --git a/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc b/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc index 2b2e3de64fd0db9d99efa46ecaf7a0fefbae6645..2273b592466431f59abcc43fcac4c37eecd53bff 100644 --- a/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc @@ -109,9 +109,9 @@ class StridedSliceOp : public XlaOpKernel { }; REGISTER_XLA_OP(Name("StridedSlice") - .CompileTimeConstInput("begin") - .CompileTimeConstInput("end") - .CompileTimeConstInput("strides"), + .CompileTimeConstantInput("begin") + .CompileTimeConstantInput("end") + .CompileTimeConstantInput("strides"), StridedSliceOp); class StridedSliceGradOp : public XlaOpKernel { @@ -218,10 +218,10 @@ class StridedSliceGradOp : public XlaOpKernel { }; REGISTER_XLA_OP(Name("StridedSliceGrad") - .CompileTimeConstInput("shape") - .CompileTimeConstInput("begin") - .CompileTimeConstInput("end") - .CompileTimeConstInput("strides"), + .CompileTimeConstantInput("shape") + .CompileTimeConstantInput("begin") + .CompileTimeConstantInput("end") + .CompileTimeConstantInput("strides"), StridedSliceGradOp); class StridedSliceAssignOp : public XlaOpKernel { @@ -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 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)); } @@ -331,9 +326,9 @@ class StridedSliceAssignOp : public XlaOpKernel { }; REGISTER_XLA_OP(Name("ResourceStridedSliceAssign") - .CompileTimeConstInput("begin") - .CompileTimeConstInput("end") - .CompileTimeConstInput("strides"), + .CompileTimeConstantInput("begin") + .CompileTimeConstantInput("end") + .CompileTimeConstantInput("strides"), StridedSliceAssignOp); } // namespace diff --git a/tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc b/tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc index 06a560d9471c352065ef7e9f6903ebdca542f5b1..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" @@ -61,8 +61,8 @@ Status MaybeInitializeTensorArray(xla::XlaBuilder* builder, " but op has dtype ", DataTypeString(dtype), "."); } - TF_RET_CHECK(resource->tensor_array_size() >= 0) - << resource->name() << " size " << resource->tensor_array_size(); + TF_RET_CHECK(resource->max_array_size() >= 0) + << resource->name() << " size " << resource->max_array_size(); if (!resource->initialized()) { TF_RETURN_IF_ERROR(resource->SetTypeAndShape(dtype, elem_shape)); @@ -78,7 +78,7 @@ Status MaybeInitializeTensorArray(xla::XlaBuilder* builder, XLAShapeToTensorShape(shape_or_status.ValueOrDie(), &shape)); TensorShape ta_shape; - ta_shape.AddDim(resource->tensor_array_size()); + ta_shape.AddDim(resource->max_array_size()); ta_shape.AppendShape(elem_shape); if (ta_shape != shape) { return errors::InvalidArgument( @@ -114,7 +114,7 @@ Status CheckTensorArrayIsInitialized(const string& op_name, Status GetTensorArrayShape(const XlaResource* resource, xla::XlaBuilder* builder, TensorShape* shape) { *shape = resource->shape(); - shape->InsertDim(0, resource->tensor_array_size()); + shape->InsertDim(0, resource->max_array_size()); return Status::OK(); } @@ -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); @@ -166,13 +167,10 @@ class TensorArrayOp : public XlaOpKernel { value = xla::Broadcast(zero, ta_shape.dim_sizes()); } - XlaContext& xc = XlaContext::Get(ctx); - XlaResource* var; - string name = absl::StrCat("TensorArray: ", tensor_array_name_); - OP_REQUIRES_OK( - ctx, xc.CreateResource(XlaResource::kTensorArray, -1, std::move(name), - dtype_, shape, value, /*tensor_array_size=*/size, - /*tensor_array_gradients=*/{}, &var)); + XlaResource* var = + ctx->xla_context()->AddResource(XlaResource::CreateTensorArray( + /*name=*/absl::StrCat("TensorArray: ", tensor_array_name_), dtype_, + shape, /*initial_value=*/value, /*max_array_size=*/size)); ctx->SetResourceOutput(0, var); Tensor flow(DT_FLOAT, TensorShape({})); @@ -188,7 +186,7 @@ class TensorArrayOp : public XlaOpKernel { TF_DISALLOW_COPY_AND_ASSIGN(TensorArrayOp); }; -REGISTER_XLA_OP(Name("TensorArrayV3").CompileTimeConstInput("size"), +REGISTER_XLA_OP(Name("TensorArrayV3").CompileTimeConstantInput("size"), TensorArrayOp); class TensorArrayWriteOp : public XlaOpKernel { @@ -215,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); @@ -266,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; @@ -422,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_); } } @@ -517,14 +515,13 @@ class TensorArraySplitOp : public XlaOpKernel { xla::XlaOp ta = resource->value(); TensorShape ta_shape; - ta_shape.AddDim(resource->tensor_array_size()); + ta_shape.AddDim(resource->max_array_size()); ta_shape.AppendShape(elem_shape); - OP_REQUIRES( - ctx, lengths.size() == resource->tensor_array_size(), - errors::InvalidArgument( - "TensorArray's size is not equal to the size of lengths (", - lengths.size(), " vs. ", resource->tensor_array_size(), ")")); + OP_REQUIRES(ctx, lengths.size() == resource->max_array_size(), + errors::InvalidArgument( + "TensorArray's size is not equal to the size of lengths (", + lengths.size(), " vs. ", resource->max_array_size(), ")")); const xla::XlaOp value = ctx->Input(1); const xla::XlaOp flow = ctx->Input(3); @@ -551,7 +548,7 @@ class TensorArraySplitOp : public XlaOpKernel { TF_DISALLOW_COPY_AND_ASSIGN(TensorArraySplitOp); }; -REGISTER_XLA_OP(Name("TensorArraySplitV3").CompileTimeConstInput("lengths"), +REGISTER_XLA_OP(Name("TensorArraySplitV3").CompileTimeConstantInput("lengths"), TensorArraySplitOp); class TensorArraySizeOp : public XlaOpKernel { @@ -562,8 +559,7 @@ class TensorArraySizeOp : public XlaOpKernel { XlaResource* var; OP_REQUIRES_OK(ctx, ctx->GetResourceInput(0, &var)); Tensor size_tensor(DT_INT32, {}); - size_tensor.scalar()() = - static_cast(var->tensor_array_size()); + size_tensor.scalar()() = static_cast(var->max_array_size()); ctx->SetConstantOutput(0, size_tensor); } diff --git a/tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc b/tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc index 74d4fcc425bdadb70a7bedf2487deaf6c4a4f7b9..6a68e0b75ea10651252a0b61f968a16875539e5f 100644 --- a/tensorflow/compiler/tf2xla/kernels/tensor_list_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/tensor_list_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" @@ -45,7 +45,7 @@ Status GetTensorListShape(xla::XlaBuilder* builder, xla::XlaOp op, return shape_or_status.status(); } xla::Shape shape = shape_or_status.ValueOrDie(); - TF_RET_CHECK(xla::ShapeUtil::IsTuple(shape)); + TF_RET_CHECK(shape.IsTuple()); return XLAShapeToTensorShape(xla::ShapeUtil::GetTupleElementShape(shape, 0), tensor_list_shape); } @@ -67,9 +67,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, 0)})); } private: @@ -79,25 +80,47 @@ class TensorListReserveOp : public XlaOpKernel { }; REGISTER_XLA_OP(Name("TensorListReserve") - .CompileTimeConstInput("element_shape") - .CompileTimeConstInput("num_elements"), + .CompileTimeConstantInput("element_shape") + .CompileTimeConstantInput("num_elements"), TensorListReserveOp); 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: @@ -147,17 +170,17 @@ 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); @@ -165,7 +188,7 @@ class TensorListPushBackOp : public XlaOpKernel { // 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,10 +220,9 @@ 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; @@ -210,7 +232,7 @@ class TensorListPopBackOp : public XlaOpKernel { // 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/tile_ops.cc b/tensorflow/compiler/tf2xla/kernels/tile_ops.cc index 52f2b36e19edd96f491f6706d1872e0d3af2df3b..e1c764f3d5c28cf0d812519e4a16786e1f2d3a3a 100644 --- a/tensorflow/compiler/tf2xla/kernels/tile_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/tile_ops.cc @@ -16,7 +16,9 @@ limitations under the License. // XLA-specific Tile Op. #include +#include "absl/algorithm/container.h" #include "absl/types/span.h" +#include "tensorflow/compiler/tf2xla/lib/broadcast.h" #include "tensorflow/compiler/tf2xla/type_util.h" #include "tensorflow/compiler/tf2xla/xla_helpers.h" #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" @@ -25,6 +27,7 @@ limitations under the License. #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/type_index.h" #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/platform/macros.h" @@ -38,11 +41,11 @@ class TileOp : public XlaOpKernel { explicit TileOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} void Compile(XlaOpKernelContext* ctx) override { - const TensorShape input_shape = ctx->InputShape(0); - const TensorShape multiples_shape = ctx->InputShape(1); + const TensorShape input_shape = ctx->InputShape("input"); + const TensorShape multiples_shape = ctx->InputShape("multiples"); OP_REQUIRES( - ctx, IsLegacyVector(multiples_shape), + ctx, TensorShapeUtils::IsVector(multiples_shape), errors::InvalidArgument("Expected multiples to be 1-D, but got shape ", multiples_shape.DebugString())); OP_REQUIRES(ctx, input_shape.dims() == multiples_shape.num_elements(), @@ -51,51 +54,46 @@ class TileOp : public XlaOpKernel { input_shape.dims(), " but got length ", multiples_shape.dim_size(0))); const int input_dims = input_shape.dims(); - + auto input = ctx->Input(0); // If input is a scalar then multiples has 0 elements and this is // a NoOp. if (input_dims == 0) { - ctx->SetOutput(0, ctx->Input(0)); + ctx->SetOutput(0, input); return; } - xla::Literal literal; - OP_REQUIRES_OK(ctx, ctx->ConstantInput(1, &literal)); - - // zero_element_result is true if the final shape has 0 elements, - // i.e. if any of the input dimensions or multiples is zero. - std::vector multiples_array(input_dims); - std::vector output_shape; - bool all_multiples_are_one = true; - bool one_dimension_is_broadcasted_without_multiple = true; - for (int i = 0; i < input_dims; ++i) { - int multiple = literal.Get({i}); - OP_REQUIRES(ctx, multiple >= 0, + std::vector multiples; + OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector("multiples", &multiples)); + std::vector output_dims(input_shape.dims()); + for (int64 i = 0; i < input_shape.dims(); ++i) { + OP_REQUIRES(ctx, multiples[i] >= 0, errors::InvalidArgument("Expected multiples[", i, - "] >= 0, but got ", multiple)); - int64 new_dim = input_shape.dim_size(i) * multiple; - output_shape.push_back(new_dim); - multiples_array[i] = multiple; - all_multiples_are_one = all_multiples_are_one && multiple == 1; - // If the multiple of a non-one dimensions is not one, then binary - // operation broadcast semantics will not be sufficient to implement the - // tile operation. - one_dimension_is_broadcasted_without_multiple = - one_dimension_is_broadcasted_without_multiple && - ((input_shape.dim_size(i) > 1 && multiple == 1) || - input_shape.dim_size(i) == 1); + "] >= 0, but got ", output_dims[i])); + output_dims[i] = input_shape.dim_size(i) * multiples[i]; } - auto input = ctx->Input(0); + // If all multiples are 1, than the input is the same as the output. - if (all_multiples_are_one) { + if (absl::c_all_of(multiples, + [](int64 multiple) { return multiple == 1; })) { ctx->SetOutput(0, input); return; } - if (one_dimension_is_broadcasted_without_multiple) { + + bool can_tile_with_implicit_broadcast = true; + for (int i = 0; i < input_dims; ++i) { + int64 multiple = multiples[i]; + // If the multiple and input dimension are not 1, then tile cannot be + // implemented with a single hlo broadcast. + if (multiple != 1 && input_shape.dim_size(i) != 1) { + can_tile_with_implicit_broadcast = false; + } + } + + if (can_tile_with_implicit_broadcast) { // Create a constant Zero the size of the output shape to leverage binary // operation broadcast semantics. auto broadcasted_zero = xla::Broadcast( - XlaHelpers::Zero(ctx->builder(), ctx->input_type(0)), output_shape); + XlaHelpers::Zero(ctx->builder(), ctx->input_type(0)), output_dims); if (ctx->input_type(0) == DT_BOOL) { ctx->SetOutput(0, xla::Or(broadcasted_zero, input)); } else { @@ -104,29 +102,16 @@ class TileOp : public XlaOpKernel { return; } - // First broadcast the requisite number of multiples along each - // dimension. This prepends the broadcasted dimensions, so an - // input of shape [2,3,1] broadcast with multiples [5,4,3] will - // end up with shape [5,4,3,2,3,1]. - auto broadcasted = xla::Broadcast(input, multiples_array); - // Now flatten and reshape. The broadcasted dimensions are - // paired with the original dimensions so in the above example - // we flatten [0,3,1,4,2,5] then reshape to [10,12,3]. - std::vector flattened; - for (int i = 0; i < output_shape.size(); ++i) { - flattened.push_back(i); - flattened.push_back(i + output_shape.size()); - } - xla::XlaOp output = xla::Reshape(broadcasted, flattened, output_shape); - - ctx->SetOutput(0, output); + auto result = BroadcastTo(ctx->Input("input"), output_dims); + OP_REQUIRES_OK(ctx, result.status()); + ctx->SetOutput(0, result.ValueOrDie()); } private: TF_DISALLOW_COPY_AND_ASSIGN(TileOp); }; -REGISTER_XLA_OP(Name("Tile").CompileTimeConstInput("multiples"), TileOp); +REGISTER_XLA_OP(Name("Tile").CompileTimeConstantInput("multiples"), TileOp); } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/topk_op.cc b/tensorflow/compiler/tf2xla/kernels/topk_op.cc index 183879c7602ccbbd74fca6cb9fa3fc94c066c37d..ee3bdf3394e37c757f31724e73e95417becaa534 100644 --- a/tensorflow/compiler/tf2xla/kernels/topk_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/topk_op.cc @@ -15,7 +15,6 @@ 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/numeric.h" #include "tensorflow/compiler/xla/client/lib/sorting.h" #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/compiler/xla/literal.h" @@ -59,7 +58,7 @@ class TopKOp : public XlaOpKernel { bool sorted_; }; -REGISTER_XLA_OP(Name("TopKV2").CompileTimeConstInput("k").TypeConstraint( +REGISTER_XLA_OP(Name("TopKV2").CompileTimeConstantInput("k").TypeConstraint( "T", {DT_UINT32, DT_INT32, DT_FLOAT, DT_BFLOAT16}), TopKOp); diff --git a/tensorflow/compiler/tf2xla/kernels/training_ops.cc b/tensorflow/compiler/tf2xla/kernels/training_ops.cc index 7077c2e3a546e198bdb4ff944ea531f3158810f2..26d4214099d1d07c1b2e275d783654d9cd948e28 100644 --- a/tensorflow/compiler/tf2xla/kernels/training_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/training_ops.cc @@ -172,6 +172,65 @@ class ResourceApplyMomentum : public XlaOpKernel { REGISTER_XLA_OP(Name("ResourceApplyMomentum").TypeConstraint("T", kFloatTypes), ResourceApplyMomentum); +class ResourceApplyKerasMomentum : public XlaOpKernel { + public: + explicit ResourceApplyKerasMomentum(OpKernelConstruction* ctx) + : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("use_nesterov", &use_nesterov_)); + } + + void Compile(XlaOpKernelContext* ctx) override { + DataType type = ctx->input_type(2); + + TensorShape var_shape, accum_shape; + xla::XlaOp var, accum; + OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(0, type, &var_shape, &var)); + OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(1, type, &accum_shape, &accum)); + + OP_REQUIRES(ctx, var_shape.IsSameSize(accum_shape), + errors::InvalidArgument( + "var and accum do not have the same shape", + var_shape.DebugString(), " ", accum_shape.DebugString())); + + TensorShape lr_shape = ctx->InputShape(2); + OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(lr_shape), + errors::InvalidArgument("lr is not a scalar: ", + lr_shape.DebugString())); + + TensorShape grad_shape = ctx->InputShape(3); + OP_REQUIRES(ctx, var_shape.IsSameSize(grad_shape), + errors::InvalidArgument( + "var and grad do not have the same shape", + var_shape.DebugString(), " ", grad_shape.DebugString())); + + TensorShape momentum_shape = ctx->InputShape(4); + OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(momentum_shape), + errors::InvalidArgument("momentum is not a scalar: ", + momentum_shape.DebugString())); + + xla::XlaOp lr = ctx->Input(2); + xla::XlaOp grad = ctx->Input(3); + xla::XlaOp momentum = ctx->Input(4); + + accum = accum * momentum - grad * lr; + if (use_nesterov_) { + // See https://github.com/tensorflow/tensorflow/pull/2798 for an + // explanation of the reparameterization used here. + var = var + accum * momentum - grad * lr; + } else { + var = var + accum; + } + OP_REQUIRES_OK(ctx, ctx->AssignVariable(0, type, var)); + OP_REQUIRES_OK(ctx, ctx->AssignVariable(1, type, accum)); + } + + private: + bool use_nesterov_; +}; +REGISTER_XLA_OP( + Name("ResourceApplyKerasMomentum").TypeConstraint("T", kFloatTypes), + ResourceApplyKerasMomentum); + class ResourceApplyAdagrad : public XlaOpKernel { public: explicit ResourceApplyAdagrad(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} @@ -320,9 +379,8 @@ class ResourceApplyAdagradDA : public XlaOpKernel { xla::XlaOp lr = ctx->Input(4); xla::XlaOp l1 = ctx->Input(5); xla::XlaOp l2 = ctx->Input(6); - xla::XlaBuilder* const b = ctx->builder(); xla::XlaOp global_step = - XlaHelpers::ConvertElementType(b, ctx->Input(7), dtype_); + XlaHelpers::ConvertElementType(ctx->Input(7), dtype_); accum = accum + grad; squared_accum = squared_accum + xla::Square(grad); diff --git a/tensorflow/compiler/tf2xla/kernels/transpose_op.cc b/tensorflow/compiler/tf2xla/kernels/transpose_op.cc index 6b303b31d43ce2249a87f25723caf34f84c8387d..76793d677ba45f8e863e684a149da684c8ce8787 100644 --- a/tensorflow/compiler/tf2xla/kernels/transpose_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/transpose_op.cc @@ -24,9 +24,9 @@ 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/framework/register_types.h" -#include "tensorflow/core/kernels/bounds_check.h" namespace tensorflow { namespace { @@ -37,8 +37,8 @@ class TransposeOp : public XlaOpKernel { : XlaOpKernel(ctx), conjugate_(conjugate) {} void Compile(XlaOpKernelContext* ctx) override { - const TensorShape input_shape = ctx->InputShape(0); - const TensorShape perm_tensor_shape = ctx->InputShape(1); + const TensorShape input_shape = ctx->InputShape("x"); + const TensorShape perm_tensor_shape = ctx->InputShape("perm"); // Preliminary validation of sizes. OP_REQUIRES(ctx, TensorShapeUtils::IsVector(perm_tensor_shape), @@ -52,19 +52,15 @@ class TransposeOp : public XlaOpKernel { ". But input(1) is a vector of size ", perm_tensor_shape.num_elements())); - xla::Literal literal; - OP_REQUIRES_OK(ctx, ctx->ConstantInputReshaped(1, {dims}, &literal)); - - std::vector perm(dims); - std::copy(literal.data().begin(), literal.data().end(), - perm.begin()); + std::vector perm; + OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector("perm", &perm)); std::vector transposed_order; // Check whether permutation is a permutation of integers of [0 .. dims). absl::InlinedVector bits(dims); bool is_identity = true; for (int i = 0; i < dims; ++i) { - const int32 d = perm[i]; + const int64 d = perm[i]; OP_REQUIRES( ctx, 0 <= d && d < dims, errors::InvalidArgument(d, " is out of range [0 .. ", dims, ")")); @@ -83,9 +79,9 @@ class TransposeOp : public XlaOpKernel { xla::XlaOp transposed; // 0-D, 1-D, and identity transposes do nothing. if (dims <= 1 || is_identity) { - transposed = ctx->Input(0); + transposed = ctx->Input("x"); } else { - transposed = xla::Transpose(ctx->Input(0), transposed_order); + transposed = xla::Transpose(ctx->Input("x"), transposed_order); } // Conjugate the transposed result if this is ConjugateTransposeOp. @@ -106,9 +102,10 @@ class ConjugateTransposeOp : public TransposeOp { : TransposeOp(ctx, /*conjugate=*/true) {} }; -REGISTER_XLA_OP(Name("Transpose").CompileTimeConstInput("perm"), TransposeOp); +REGISTER_XLA_OP(Name("Transpose").CompileTimeConstantInput("perm"), + TransposeOp); -REGISTER_XLA_OP(Name("ConjugateTranspose").CompileTimeConstInput("perm"), +REGISTER_XLA_OP(Name("ConjugateTranspose").CompileTimeConstantInput("perm"), ConjugateTransposeOp); // InvertPermutation frequently forms part of the gradient of Transpose. @@ -153,7 +150,7 @@ class InvertPermutationOp : public XlaOpKernel { REGISTER_XLA_OP(Name("InvertPermutation") .TypeConstraint("T", DT_INT32) - .CompileTimeConstInput("x"), + .CompileTimeConstantInput("x"), InvertPermutationOp); } // namespace diff --git a/tensorflow/compiler/tf2xla/kernels/unary_ops.cc b/tensorflow/compiler/tf2xla/kernels/unary_ops.cc index 0bdfc05726105e2d18362a691cbe2aab00bf77f3..a0ea6422d732b00fc1b8cf855d9c9ad603b87c82 100644 --- a/tensorflow/compiler/tf2xla/kernels/unary_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/unary_ops.cc @@ -80,24 +80,8 @@ XLAJIT_MAKE_UNARY(Invert, xla::Not(x)); XLAJIT_MAKE_UNARY(LogicalNot, xla::Not(x)); XLAJIT_MAKE_UNARY(Neg, -x); -// Implements Banker's rounding: numbers that are equidistant between two -// integers are rounded towards even. -xla::XlaOp RoundToEven(xla::XlaOp x) { - auto half = xla::ScalarLike(x, 0.5); - auto one = xla::ScalarLike(x, 1.0); - auto two = xla::ScalarLike(x, 2.0); - - auto round_val = xla::Floor(x); - auto fraction = x - round_val; - auto nearest_even_int = round_val - two * xla::Floor(half * x); - auto is_odd = xla::Eq(nearest_even_int, one); - return xla::Select(xla::Or(xla::Gt(fraction, half), - xla::And(xla::Eq(fraction, half), is_odd)), - round_val + one, round_val); -} - -XLAJIT_MAKE_UNARY(Rint, RoundToEven(x)); -XLAJIT_MAKE_UNARY(Round, RoundToEven(x)); +XLAJIT_MAKE_UNARY(Rint, xla::RoundToEven(x)); +XLAJIT_MAKE_UNARY(Round, xla::RoundToEven(x)); XLAJIT_MAKE_UNARY(Rsqrt, xla::Rsqrt(x)); 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/variable_ops.cc b/tensorflow/compiler/tf2xla/kernels/variable_ops.cc index 2c92a585f5679242d672d0402e617ff199b94f17..dfa09b16081e93ba843a1858e68e6ff756de20c1 100644 --- a/tensorflow/compiler/tf2xla/kernels/variable_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/variable_ops.cc @@ -291,5 +291,19 @@ class ResourceScatterNdAddOp : public ResourceScatterOp { }; REGISTER_XLA_OP(Name("ResourceScatterNdAdd"), ResourceScatterNdAddOp); +class ResourceScatterNdSubOp : public ResourceScatterOp { + public: + explicit ResourceScatterNdSubOp(OpKernelConstruction* context) + : ResourceScatterOp(context, /*indices_are_vectors=*/true, + /*combiner=*/Combine) {} + + private: + static xla::XlaOp Combine(const xla::XlaOp& x, const xla::XlaOp& y, + xla::XlaBuilder* builder) { + return xla::Sub(x, y); + } +}; +REGISTER_XLA_OP(Name("ResourceScatterNdSub"), ResourceScatterNdSubOp); + } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/while_op.cc b/tensorflow/compiler/tf2xla/kernels/while_op.cc index 559414eeaa5fec75e5a9d1866baaf738c024cd15..fd5ff10ae0a8cb39075fa6c594707dbc833f5f16 100644 --- a/tensorflow/compiler/tf2xla/kernels/while_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/while_op.cc @@ -41,8 +41,7 @@ Status MakeXlaCompilerArgumentsFromInputs( *has_uninitialized_vars = false; *has_tensor_arrays = false; for (int i = 0; i < ctx->num_inputs(); ++i) { - VLOG(2) << " Input " << i - << " type: " << DataTypeString(ctx->input_type(i)) + VLOG(2) << " Input " << i << " type: " << DataTypeString(ctx->input_type(i)) << " shape: " << ctx->InputShape(i).DebugString(); XlaCompiler::Argument& arg = (*args)[i]; DataType type = ctx->input_type(i); @@ -64,20 +63,27 @@ Status MakeXlaCompilerArgumentsFromInputs( if (!arg.initialized) { *has_uninitialized_vars = true; } - arg.tensor_array_size = resource->tensor_array_size(); + arg.max_array_size = resource->max_array_size(); for (const auto& gradient : resource->tensor_array_gradients()) { arg.tensor_array_gradients.insert(gradient.first); } arg.name = resource->name(); VLOG(2) << " resource " << resource->name() << " type: " << DataTypeString(arg.type) - << " shape: " << arg.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(); @@ -207,12 +213,12 @@ void XlaWhileOp::Compile(XlaOpKernelContext* ctx) { OP_REQUIRES(ctx, body.xla_input_shapes.size() == 1, errors::FailedPrecondition("Expected one input shape")); xla::Shape body_input_shape = body.xla_input_shapes[0]; - OP_REQUIRES(ctx, xla::ShapeUtil::IsTuple(body_input_shape), + OP_REQUIRES(ctx, body_input_shape.IsTuple(), errors::FailedPrecondition("Expected tuple shape")); OP_REQUIRES(ctx, cond.xla_input_shapes.size() == 1, errors::FailedPrecondition("Expected one input shape")); xla::Shape cond_input_shape = cond.xla_input_shapes[0]; - OP_REQUIRES(ctx, xla::ShapeUtil::IsTuple(cond_input_shape), + OP_REQUIRES(ctx, cond_input_shape.IsTuple(), errors::FailedPrecondition("Expected tuple shape")); VLOG(2) << "Body shape: " << xla::ShapeUtil::HumanString(body_input_shape) @@ -233,13 +239,22 @@ void XlaWhileOp::Compile(XlaOpKernelContext* ctx) { xla::ShapeUtil::HumanString(body_input_shape), " vs. ", xla::ShapeUtil::HumanString(body.xla_output_shape))); - xla::Shape expected_cond_output_shape = xla::ShapeUtil::MakeTupleShape( - {xla::ShapeUtil::MakeShape(xla::PRED, {})}); + xla::Shape expected_cond_output_shape_without_side_effect = + xla::ShapeUtil::MakeTupleShape( + {xla::ShapeUtil::MakeShape(xla::PRED, {})}); + xla::Shape expected_cond_output_shape_with_side_effect = + xla::ShapeUtil::MakeTupleShape({xla::ShapeUtil::MakeShape(xla::PRED, {}), + xla::ShapeUtil::MakeTokenShape()}); OP_REQUIRES(ctx, - xla::ShapeUtil::Compatible(cond.xla_output_shape, - expected_cond_output_shape), + xla::ShapeUtil::Compatible( + cond.xla_output_shape, + expected_cond_output_shape_without_side_effect) || + xla::ShapeUtil::Compatible( + cond.xla_output_shape, + expected_cond_output_shape_with_side_effect), errors::InvalidArgument( - "Output shape of loop condition should be (pred[]), got: ", + "Output shape of loop condition should be (pred[]) or " + "(pred[], token[]), got: ", xla::ShapeUtil::HumanString(cond.xla_output_shape))); int num_inputs = body.input_mapping.size(); @@ -283,11 +298,15 @@ void XlaWhileOp::Compile(XlaOpKernelContext* ctx) { xla::XlaOp while_result = xla::While(cond_wrapper, *body.computation, init); - // Sets non-variable outputs. + // Sets non-variable outputs and determine when resource variables start. + int resource_index = 0; for (int i = 0; i < ctx->num_outputs(); ++i) { if (ctx->input_type(i) != DT_RESOURCE) { ctx->SetOutput(body.input_mapping[i], xla::GetTupleElement(while_result, i)); + ++resource_index; + } else { + break; } } if (has_token_input_output_) { @@ -296,7 +315,7 @@ void XlaWhileOp::Compile(XlaOpKernelContext* ctx) { xla::GetTupleElement(while_result, ctx->num_outputs()); auto shape_or = builder->GetShape(token_output); OP_REQUIRES_OK(ctx, shape_or.status()); - OP_REQUIRES(ctx, xla::ShapeUtil::IsToken(shape_or.ValueOrDie()), + OP_REQUIRES(ctx, shape_or.ValueOrDie().IsToken(), errors::FailedPrecondition( "Token output is not token type: ", xla::ShapeUtil::HumanString(shape_or.ValueOrDie()))); @@ -309,7 +328,7 @@ void XlaWhileOp::Compile(XlaOpKernelContext* ctx) { XlaResource* resource; OP_REQUIRES_OK(ctx, ctx->GetResourceInput(update.input_index, &resource)); if (update.modified) { - int pos = body.outputs.size() + i; + int pos = resource_index + i; OP_REQUIRES_OK(ctx, resource->SetFromPack( arguments[update.input_index].tensor_array_gradients, diff --git a/tensorflow/compiler/tf2xla/kernels/xla_broadcast_helper_op.cc b/tensorflow/compiler/tf2xla/kernels/xla_broadcast_helper_op.cc index 412afeaaad96842521fbd306f5b666e837e675fd..ad8e707e1116d01d492575986a7ab9586022f6b3 100644 --- a/tensorflow/compiler/tf2xla/kernels/xla_broadcast_helper_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/xla_broadcast_helper_op.cc @@ -89,13 +89,10 @@ class XlaBroadcastHelperOp : public XlaOpKernel { lhs_shape.DebugString(), " and ", rhs_shape.DebugString())); broadcast_shape[dim] = min_rank_shape->dim_size(i); } - xla::PrimitiveType type = context->input_xla_type(0); - xla::Shape broadcast_xla_shape = - xla::ShapeUtil::MakeShape(type, broadcast_shape); if (broadcast_lhs) { - lhs = xla::BroadcastInDim(lhs, broadcast_xla_shape, broadcast_dims); + lhs = xla::BroadcastInDim(lhs, broadcast_shape, broadcast_dims); } else { - rhs = xla::BroadcastInDim(rhs, broadcast_xla_shape, broadcast_dims); + rhs = xla::BroadcastInDim(rhs, broadcast_shape, broadcast_dims); } context->SetOutput(0, lhs); context->SetOutput(1, rhs); @@ -108,7 +105,7 @@ class XlaBroadcastHelperOp : public XlaOpKernel { }; REGISTER_XLA_OP( - Name("XlaBroadcastHelper").CompileTimeConstInput("broadcast_dims"), + Name("XlaBroadcastHelper").CompileTimeConstantInput("broadcast_dims"), XlaBroadcastHelperOp); } // namespace diff --git a/tensorflow/compiler/tf2xla/kernels/xla_conv_op.cc b/tensorflow/compiler/tf2xla/kernels/xla_conv_op.cc index fecc7c556eb4121b912796e5811632c46769b479..b20adc592a0d3d2129c897218ddbfc891b4cd40a 100644 --- a/tensorflow/compiler/tf2xla/kernels/xla_conv_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/xla_conv_op.cc @@ -78,7 +78,7 @@ class XlaConvOp : public XlaOpKernel { xla::XlaOp output = xla::ConvGeneralDilated( context->Input(0), context->Input(1), window_strides, padding, lhs_dilation, rhs_dilation, dnums_, feature_group_count, - &precision_config_); + /*batch_group_count=*/1, &precision_config_); context->SetOutput(0, output); } @@ -90,11 +90,11 @@ class XlaConvOp : public XlaOpKernel { }; REGISTER_XLA_OP(Name("XlaConv") - .CompileTimeConstInput("window_strides") - .CompileTimeConstInput("lhs_dilation") - .CompileTimeConstInput("rhs_dilation") - .CompileTimeConstInput("feature_group_count") - .CompileTimeConstInput("padding"), + .CompileTimeConstantInput("window_strides") + .CompileTimeConstantInput("lhs_dilation") + .CompileTimeConstantInput("rhs_dilation") + .CompileTimeConstantInput("feature_group_count") + .CompileTimeConstantInput("padding"), XlaConvOp); } // namespace diff --git a/tensorflow/compiler/tf2xla/kernels/xla_dequantize_op.cc b/tensorflow/compiler/tf2xla/kernels/xla_dequantize_op.cc new file mode 100644 index 0000000000000000000000000000000000000000..a30b4861f6b3a964c0c874a3affab7d6198264d7 --- /dev/null +++ b/tensorflow/compiler/tf2xla/kernels/xla_dequantize_op.cc @@ -0,0 +1,60 @@ +/* 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/tf2xla/shape_util.h" +#include "tensorflow/compiler/tf2xla/xla_compiler.h" +#include "tensorflow/compiler/tf2xla/xla_op_kernel.h" +#include "tensorflow/compiler/tf2xla/xla_op_registry.h" +#include "tensorflow/compiler/xla/client/lib/quantize.h" +#include "tensorflow/compiler/xla/client/xla_builder.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/framework/op_kernel.h" + +namespace tensorflow { +namespace { + +class XlaDequantizeOp : public XlaOpKernel { + public: + explicit XlaDequantizeOp(OpKernelConstruction* context) + : XlaOpKernel(context) { + OP_REQUIRES_OK(context, context->GetAttr("min_range", &min_range_)); + OP_REQUIRES_OK(context, context->GetAttr("max_range", &max_range_)); + OP_REQUIRES_OK(context, context->GetAttr("mode", &mode_)); + OP_REQUIRES_OK(context, + context->GetAttr("transpose_output", &transpose_output_)); + } + + void Compile(XlaOpKernelContext* context) override { + const xla::XlaOp& input = context->Input(0); + + xla::QuantizedRange range(min_range_, max_range_); + + xla::XlaOp output = + xla::Dequantize(input, range, mode_, transpose_output_); + context->SetOutput(0, output); + } + + private: + float min_range_; + float max_range_; + bool transpose_output_; + string mode_; + TF_DISALLOW_COPY_AND_ASSIGN(XlaDequantizeOp); +}; + +REGISTER_XLA_OP(Name("XlaDequantize"), XlaDequantizeOp); + +} // namespace +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/xla_pad_op.cc b/tensorflow/compiler/tf2xla/kernels/xla_pad_op.cc index 59502d83c7338bd1b05b3323a97761fff2da186a..a3c2eef993c80e43e7cf9e1f6147e5b337c41cfe 100644 --- a/tensorflow/compiler/tf2xla/kernels/xla_pad_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/xla_pad_op.cc @@ -96,9 +96,9 @@ class XlaPadOp : public XlaOpKernel { }; REGISTER_XLA_OP(Name("XlaPad") - .CompileTimeConstInput("padding_low") - .CompileTimeConstInput("padding_high") - .CompileTimeConstInput("padding_interior"), + .CompileTimeConstantInput("padding_low") + .CompileTimeConstantInput("padding_high") + .CompileTimeConstantInput("padding_interior"), XlaPadOp); } // namespace diff --git a/tensorflow/compiler/tf2xla/kernels/xla_select_and_scatter_op.cc b/tensorflow/compiler/tf2xla/kernels/xla_select_and_scatter_op.cc index 089776fcf74fcf6b363dfff5de8d86d7449eacd6..9043af995386a179f74d95bbc6c17a1cac7881cd 100644 --- a/tensorflow/compiler/tf2xla/kernels/xla_select_and_scatter_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/xla_select_and_scatter_op.cc @@ -138,9 +138,9 @@ class XlaSelectAndScatterOp : public XlaOpKernel { }; REGISTER_XLA_OP(Name("XlaSelectAndScatter") - .CompileTimeConstInput("window_dimensions") - .CompileTimeConstInput("window_strides") - .CompileTimeConstInput("padding"), + .CompileTimeConstantInput("window_dimensions") + .CompileTimeConstantInput("window_strides") + .CompileTimeConstantInput("padding"), XlaSelectAndScatterOp); } // namespace diff --git a/tensorflow/compiler/tf2xla/lib/BUILD b/tensorflow/compiler/tf2xla/lib/BUILD index 1ce3930fd1cd91f8e8dfb765b49be2dc969d1bd7..3d7b0bc959f9dbf3c1b9749379e2ea0d285b302b 100644 --- a/tensorflow/compiler/tf2xla/lib/BUILD +++ b/tensorflow/compiler/tf2xla/lib/BUILD @@ -15,22 +15,6 @@ filegroup( ]), ) -load("//tensorflow/compiler/xla/tests:build_defs.bzl", "xla_test") - -cc_library( - name = "batch_dot", - srcs = ["batch_dot.cc"], - hdrs = ["batch_dot.h"], - deps = [ - "//tensorflow/compiler/xla:shape_util", - "//tensorflow/compiler/xla:status_macros", - "//tensorflow/compiler/xla:statusor", - "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/client:xla_builder", - "//tensorflow/core:lib", - ], -) - cc_library( name = "broadcast", srcs = ["broadcast.cc"], @@ -47,26 +31,6 @@ cc_library( ], ) -cc_library( - name = "cholesky", - srcs = ["cholesky.cc"], - hdrs = ["cholesky.h"], - deps = [ - ":batch_dot", - ":triangular_solve", - ":util", - ":while_loop", - "//tensorflow/compiler/xla:literal", - "//tensorflow/compiler/xla:shape_util", - "//tensorflow/compiler/xla:status_macros", - "//tensorflow/compiler/xla:statusor", - "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/client:xla_builder", - "//tensorflow/compiler/xla/client/lib:constants", - "//tensorflow/core:lib", - ], -) - cc_library( name = "random", srcs = ["random.cc"], @@ -82,35 +46,12 @@ cc_library( ], ) -cc_library( - name = "qr", - srcs = ["qr.cc"], - hdrs = ["qr.h"], - deps = [ - ":batch_dot", - ":util", - ":while_loop", - "//tensorflow/compiler/xla:literal_util", - "//tensorflow/compiler/xla:shape_util", - "//tensorflow/compiler/xla:status_macros", - "//tensorflow/compiler/xla:statusor", - "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/client:xla_builder", - "//tensorflow/compiler/xla/client/lib:arithmetic", - "//tensorflow/compiler/xla/client/lib:constants", - "//tensorflow/compiler/xla/client/lib:math", - "//tensorflow/compiler/xla/client/lib:numeric", - "//tensorflow/core:lib", - ], -) - cc_library( name = "scatter", srcs = ["scatter.cc"], hdrs = ["scatter.h"], deps = [ ":util", - ":while_loop", "//tensorflow/compiler/xla:literal", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:status_macros", @@ -124,51 +65,6 @@ cc_library( ], ) -cc_library( - name = "triangular_solve", - srcs = ["triangular_solve.cc"], - hdrs = ["triangular_solve.h"], - deps = [ - ":batch_dot", - ":util", - "//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: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:numeric", - "//tensorflow/core:lib", - ], -) - -xla_test( - name = "triangular_solve_test", - srcs = ["triangular_solve_test.cc"], - tags = ["noasan"], # sometimes times out, http://b/78650012 - deps = [ - ":triangular_solve", - "//tensorflow/compiler/xla:array2d", - "//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", - ], -) - cc_library( name = "util", srcs = ["util.cc"], @@ -186,42 +82,3 @@ cc_library( "@com_google_absl//absl/types:span", ], ) - -xla_test( - name = "util_test", - srcs = ["util_test.cc"], - deps = [ - ":batch_dot", - ":util", - "//tensorflow/compiler/xla:array2d", - "//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/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", - ], -) - -cc_library( - name = "while_loop", - srcs = ["while_loop.cc"], - hdrs = ["while_loop.h"], - deps = [ - ":util", - "//tensorflow/compiler/xla:shape_util", - "//tensorflow/compiler/xla:status_macros", - "//tensorflow/compiler/xla:statusor", - "//tensorflow/compiler/xla/client:xla_builder", - "//tensorflow/compiler/xla/client:xla_computation", - "@com_google_absl//absl/strings", - "@com_google_absl//absl/types:span", - ], -) diff --git a/tensorflow/compiler/tf2xla/lib/batch_dot.cc b/tensorflow/compiler/tf2xla/lib/batch_dot.cc deleted file mode 100644 index 5400e8834cb9807f6dd71abe7789b2672e29e905..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/tf2xla/lib/batch_dot.cc +++ /dev/null @@ -1,115 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/compiler/tf2xla/lib/batch_dot.h" - -#include -#include - -#include "tensorflow/compiler/xla/client/xla_builder.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 tensorflow { - -xla::XlaOp BatchDot(xla::XlaOp x, xla::XlaOp y, bool transpose_x, - bool transpose_y, bool conjugate_x, bool conjugate_y, - xla::PrecisionConfig::Precision precision) { - xla::XlaBuilder* builder = x.builder(); - return builder->ReportErrorOrReturn([&]() -> xla::StatusOr { - TF_ASSIGN_OR_RETURN(xla::Shape x_shape, builder->GetShape(x)); - TF_ASSIGN_OR_RETURN(xla::Shape y_shape, builder->GetShape(y)); - - // Check that both tensors have the same number of dimensions. There must be - // at least two (the batch dimensions can be empty). - if (xla::ShapeUtil::Rank(x_shape) != xla::ShapeUtil::Rank(y_shape)) { - return errors::InvalidArgument( - "Arguments to BatchedDot have different ranks: ", - xla::ShapeUtil::HumanString(x_shape), " vs. ", - xla::ShapeUtil::HumanString(y_shape)); - } - const int ndims = xla::ShapeUtil::Rank(x_shape); - if (ndims < 2) { - return errors::InvalidArgument( - "Arguments to BatchedDot must have rank >= 2: ", ndims); - } - - // 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 errors::InvalidArgument( - "Dimension ", i, " of inputs to BatchedDot must be equal: ", - xla::ShapeUtil::HumanString(x_shape), " vs ", - xla::ShapeUtil::HumanString(y_shape)); - } - batch_dimension_numbers.push_back(i); - } - - int x_inner_dim = transpose_x ? (ndims - 2) : (ndims - 1); - int y_inner_dim = transpose_y ? (ndims - 1) : (ndims - 2); - if (x_shape.dimensions(x_inner_dim) != y_shape.dimensions(y_inner_dim)) { - return errors::InvalidArgument( - "Dimensions ", x_inner_dim, " and ", y_inner_dim, - " of arguments to BatchedDot must be equal: ", - xla::ShapeUtil::HumanString(x_shape), " transpose: ", transpose_x, - " vs. ", xla::ShapeUtil::HumanString(y_shape), - " transpose: ", transpose_y); - } - - // Check for zero lhs/rhs dim size. - if (xla::ShapeUtil::IsZeroElementArray(x_shape) || - xla::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]); - } - int x_outer_dim = transpose_x ? (ndims - 1) : (ndims - 2); - int y_outer_dim = transpose_y ? (ndims - 2) : (ndims - 1); - dimensions.push_back(x_shape.dimensions(x_outer_dim)); - dimensions.push_back(y_shape.dimensions(y_outer_dim)); - return xla::Broadcast( - xla::ConstantLiteral(builder, - xla::LiteralUtil::Zero(x_shape.element_type())), - dimensions); - } - - if (x_shape.element_type() == xla::C64 && conjugate_x) { - x = xla::Conj(x); - } - if (y_shape.element_type() == xla::C64 && conjugate_y) { - y = xla::Conj(y); - } - - xla::PrecisionConfig precision_proto; - precision_proto.add_operand_precision(precision); - precision_proto.add_operand_precision(precision); - - xla::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); - } - - return xla::DotGeneral(x, y, dot_dnums, &precision_proto); - }); -} - -} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/lib/batch_dot.h b/tensorflow/compiler/tf2xla/lib/batch_dot.h deleted file mode 100644 index 6edd63a4d3b66c21aa4cce8c9f36eef0dc363cd8..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/tf2xla/lib/batch_dot.h +++ /dev/null @@ -1,54 +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_TF2XLA_LIB_BATCH_DOT_H_ -#define TENSORFLOW_COMPILER_TF2XLA_LIB_BATCH_DOT_H_ - -#include "tensorflow/compiler/xla/client/xla_builder.h" -#include "tensorflow/compiler/xla/xla_data.pb.h" - -namespace tensorflow { - -// Multiplies slices of two tensors in batches. - -// Multiplies all slices of `Tensor` `x` and `y` (each slice can be -// viewed as an element of a batch), and arranges the individual results -// in a single output tensor of the same batch size. Each of the -// individual slices can optionally be transposed before multiplication by -// setting the `transpose_x` or `transpose_y` flag to `true`. Similarly, each -// can be elementwise-complex-conjugated by setting the `conjugate_x` or -// `conjugate_y` flag to `true`. To apply a Hermitian adjoint to `x`, set both -// `transpose_x` and `conjugate_x` to `true`, and analogously for `y`. -// -// The input tensors `x` and `y` are 2-D or higher with shape `[..., r_x, c_x]` -// and `[..., r_y, c_y]`. -// -// The output tensor is 2-D or higher with shape `[..., r_o, c_o]`, where: -// -// r_o = c_x if transpose_x else r_x -// c_o = r_y if transpose_y else c_y -// -// It is computed as: -// -// output[..., :, :] = matrix(x[..., :, :]) * matrix(y[..., :, :]) -xla::XlaOp BatchDot( - xla::XlaOp x, xla::XlaOp y, bool transpose_x = false, - bool transpose_y = false, bool conjugate_x = false, - bool conjugate_y = false, - xla::PrecisionConfig::Precision precision = xla::PrecisionConfig::DEFAULT); - -} // namespace tensorflow - -#endif // TENSORFLOW_COMPILER_TF2XLA_LIB_BATCH_DOT_H_ diff --git a/tensorflow/compiler/tf2xla/lib/broadcast.cc b/tensorflow/compiler/tf2xla/lib/broadcast.cc index 3e402ef855cd7c114332d84032bc869232404fc8..be31f116686a2e302ece730e9d03312a45888a61 100644 --- a/tensorflow/compiler/tf2xla/lib/broadcast.cc +++ b/tensorflow/compiler/tf2xla/lib/broadcast.cc @@ -80,10 +80,8 @@ xla::StatusOr BroadcastTo(xla::XlaOp input, broadcast_dim = broadcast_shape_size - broadcast_dim - 1; } absl::c_reverse(broadcast_shape); - xla::XlaOp output = xla::BroadcastInDim( - input, - xla::ShapeUtil::MakeShape(input_shape.element_type(), broadcast_shape), - broadcast_dims); + xla::XlaOp output = + xla::BroadcastInDim(input, broadcast_shape, broadcast_dims); if (broadcast_shape != output_dims) { output = xla::Reshape(output, output_dims); } diff --git a/tensorflow/compiler/tf2xla/lib/cholesky.cc b/tensorflow/compiler/tf2xla/lib/cholesky.cc deleted file mode 100644 index ab3d0a566839343828d176d9a46672824e425613..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/tf2xla/lib/cholesky.cc +++ /dev/null @@ -1,217 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/compiler/tf2xla/lib/cholesky.h" - -#include -#include - -#include "tensorflow/compiler/tf2xla/lib/batch_dot.h" -#include "tensorflow/compiler/tf2xla/lib/triangular_solve.h" -#include "tensorflow/compiler/tf2xla/lib/util.h" -#include "tensorflow/compiler/tf2xla/lib/while_loop.h" -#include "tensorflow/compiler/xla/client/lib/constants.h" -#include "tensorflow/compiler/xla/client/xla_builder.h" -#include "tensorflow/compiler/xla/literal.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 tensorflow { - -namespace { - -// The Cholesky–Banachiewicz algorithm. See -// https://en.wikipedia.org/wiki/Cholesky_decomposition#The_Cholesky–Banachiewicz_and_Cholesky–Crout_algorithms -// for a description. -// -// def cholesky_unblocked(a): -// assert len(a.shape) == 2 and a.shape[-2] == a.shape[-1] -// n = a.shape[-2] -// l = np.zeros_like(a) -// for j in xrange(n): -// row = l[..., j, :j] -// row_t = np.swapaxes(row, -1, -2) -// l[..., j, j] = np.sqrt(a[..., j, j] - np.dot(row, row_t)) -// l[..., j+1:, j] = (a[..., j+1:, j] - np.dot(l[..., j+1:, :j], row_t)) / -// l[..., j, j] -// return l -xla::XlaOp CholeskyUnblocked(xla::XlaOp a, - xla::PrecisionConfig::Precision precision) { - xla::XlaBuilder* builder = a.builder(); - return builder->ReportErrorOrReturn([&]() -> xla::StatusOr { - TF_ASSIGN_OR_RETURN(xla::Shape a_shape, builder->GetShape(a)); - const int n_dims = xla::ShapeUtil::Rank(a_shape); - const int64 n = xla::ShapeUtil::GetDimension(a_shape, -1); - auto major_dims = xla::AsInt64Slice(a_shape.dimensions()) - .subspan( - /*pos=*/0, - /*len=*/n_dims - 2); - - xla::XlaOp l = xla::ZerosLike(a); - - // Construct the for loop body to iterate over rows. - auto body_fn = [&](xla::XlaOp i, absl::Span loop_vars, - xla::XlaBuilder* body_builder) - -> xla::StatusOr> { - xla::Shape col_shape; - xla::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 = xla::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 = xla::Zeros(body_builder, col_shape); - - std::vector mask_vector(n); - std::iota(mask_vector.begin(), mask_vector.end(), 0); - auto mask_range = xla::ConstantR1(body_builder, mask_vector); - auto mask_range_row = - xla::Broadcast(xla::Reshape(mask_range, {0}, {1, n}), major_dims); - auto mask_range_col = - xla::Broadcast(xla::Reshape(mask_range, {0}, {n, 1}), major_dims); - auto body_a = loop_vars[0]; - auto body_l = loop_vars[1]; - - // row = l[..., i, :i] - // select the whole i-th row, then mask out all columns past i-1 - auto zero = xla::ConstantR0(body_builder, 0); - auto l_i = DynamicSliceInMinorDims(body_l, {i, zero}, {1, n}); - auto row = xla::Select(xla::Ge(mask_range_row, i), mask_zeros_row, l_i); - // a[..., i, i] - auto a_ii = DynamicSliceInMinorDims(body_a, {i, i}, {1, 1}); - // np.dot(row, np.swapaxes(row, -1, -2)) - auto diag_dot = BatchDot(row, row, - /*transpose_x=*/false, - /*transpose_y=*/true, /*conjugate_x=*/false, - /*conjugate_y=*/false, precision); - // l[..., i, i] = np.sqrt(a[..., i, i] - np.dot(row, - // np.swapaxes(row, -1, -2))) - auto l_ii = - xla::Pow(a_ii - diag_dot, - FloatLiteral(body_builder, a_shape.element_type(), 0.5)); - - // a[..., i+1:, i] - // select the whole i-th column, then mask out all rows above i+1 - auto a_0i = DynamicSliceInMinorDims(body_a, {i}, {1}); - auto a_ip1i = - xla::Select(xla::Le(mask_range_col, i), mask_zeros_col, a_0i); - - // l[..., i+1:, i] = (a[..., i+1:, i] - np.dot(l[..., i+1:, :i], r.T)) / - // l[..., i, i] - // The columns in [i, n] are zeroed out in `row`, so we just have to - // zero out rows above i+1 after the BatchDot. np.dot(l[..., :, :i], - // r.T) - auto dot = BatchDot(body_l, row, - /*transpose_x=*/false, - /*transpose_y=*/true, /*conjugate_x=*/false, - /*conjugate_y=*/false, precision); - // np.dot(l[..., i+1:, :i], r.T) - auto dot_ip1 = - xla::Select(xla::Le(mask_range_col, i), mask_zeros_col, dot); - - body_l = - DynamicUpdateSliceInMinorDims(body_l, (a_ip1i - dot_ip1) / l_ii, {i}); - // Assign the diagonal after the rest of the column because otherwise the - // column assign will wrap around and overwrite the diagonal assign. - body_l = DynamicUpdateSliceInMinorDims(body_l, l_ii, {i, i}); - - return std::vector{body_a, body_l}; - }; - - TF_ASSIGN_OR_RETURN( - auto cholesky_while, - XlaForEachIndex(n, xla::S32, body_fn, {a, l}, "unblocked", builder)); - - return cholesky_while[1]; - }); -} - -} // namespace - -xla::XlaOp Cholesky(xla::XlaOp a, int64 block_size, - xla::PrecisionConfig::Precision precision) { - xla::XlaBuilder* builder = a.builder(); - return builder->ReportErrorOrReturn([&]() -> xla::StatusOr { - TF_ASSIGN_OR_RETURN(xla::Shape a_shape, builder->GetShape(a)); - const int ndims = xla::ShapeUtil::Rank(a_shape); - if (ndims < 2) { - return errors::InvalidArgument( - "Arguments to Cholesky must have rank >= 2: ", ndims); - } - - const int64 n = xla::ShapeUtil::GetDimension(a_shape, -1); - if (n != xla::ShapeUtil::GetDimension(a_shape, -2)) { - return errors::InvalidArgument( - "Arguments to Cholesky must be square matrices: ", - xla::ShapeUtil::HumanString(a_shape)); - } - - if (block_size < 1) { - return errors::InvalidArgument( - "block_size argument to Cholesky must be >= 1; got ", block_size); - } - - // Blocked left-looking Cholesky factorization. - // Algorithm 1 from - // Haidar, Azzam, et al. "High-performance Cholesky factorization for - // GPU-only execution." Proceedings of General Purpose GPUs. ACM, 2017. - xla::XlaOp l = xla::ZerosLike(a); - for (int64 i = 0; i < n; i += block_size) { - int64 k = std::min(block_size, n - i); - if (i > 0) { - // TODO(phawkins): consider implementing SYRK for the diagonal part of - // the panel. - // a[i:, i:i+k] -= np.dot(l[i:, :i], np.transpose(l[i:i+k, :i])) - auto lhs = SliceInMinorDims(l, {i, 0}, {n, i}); - auto rhs = SliceInMinorDims(l, {i, 0}, {i + k, i}); - auto delta = BatchDot(lhs, rhs, /*transpose_x=*/false, - /*transpose_y=*/true, /*conjugate_x=*/false, - /*conjugate_y=*/false, precision); - auto before = SliceInMinorDims(a, {i, i}, {n, i + k}); - a = UpdateSliceInMinorDims(a, before - delta, {i, i}); - } - - // l[i:i+k, i:i+k] = cholesky_unblocked(a[i:i+k, i:i+k]) - auto x = SliceInMinorDims(a, {i, i}, {i + k, i + k}); - auto factorized = CholeskyUnblocked(x, precision); - l = UpdateSliceInMinorDims(l, factorized, {i, i}); - - if (i + k < n) { - // 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); - l = UpdateSliceInMinorDims(l, update, {i + k, i}); - } - } - return l; - }); -} - -} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/lib/qr.cc b/tensorflow/compiler/tf2xla/lib/qr.cc deleted file mode 100644 index 6b3f2b6e065b5c99e2d0248237369ecc30188aa5..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/tf2xla/lib/qr.cc +++ /dev/null @@ -1,411 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/compiler/tf2xla/lib/qr.h" - -#include -#include - -#include "tensorflow/compiler/tf2xla/lib/batch_dot.h" -#include "tensorflow/compiler/tf2xla/lib/util.h" -#include "tensorflow/compiler/tf2xla/lib/while_loop.h" -#include "tensorflow/compiler/xla/client/lib/arithmetic.h" -#include "tensorflow/compiler/xla/client/lib/constants.h" -#include "tensorflow/compiler/xla/client/lib/math.h" -#include "tensorflow/compiler/xla/client/lib/numeric.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 tensorflow { - -namespace { - -// Computes a Householder reflection of the form: -// H = I - tau v v.T. -// such that -// H . ( x1 ) = ( x1 ) -// ( x2 ) = ( x2 ) -// ( ... ) = ( ... ) -// ( xk ) = ( beta ) -// ( ... ) ( 0 ) -// ( ... ) ( 0 ) -// Unlike the usual formulation, we allow the caller to supply 'k' rather than -// only providing the relevant part of 'x' to maintain XLA's static shape -// invariant. In addition, the implementation supports batching. -// Pseudo-code, without batching: -// alpha = x[k] -// x_copy = np.copy(x) -// x_copy[:k+1] = 0 -// xnorm = norm2(x_copy) -// if xnorm == 0: -// beta = alpha -// tau = 0 -// v = np.zeros_like(x) -// else: -// beta = - np.sign(alpha) * dlapy2(alpha, xnorm) -// tau = (beta - alpha) / beta -// v = x / (alpha - beta) -// v[k] = 1 -// return (v, tau, beta) -// TODO(phawkins): LAPACK's xLARFG implementation has code for handling -// overflows in the norm/beta calculations. Perhaps do the same here. -xla::Status House(xla::XlaOp x, xla::XlaOp k, - absl::Span batch_dims, const int64 m, - xla::XlaOp* v, xla::XlaOp* tau, xla::XlaOp* beta) { - xla::XlaBuilder* const builder = x.builder(); - TF_ASSIGN_OR_RETURN(xla::Shape x_shape, builder->GetShape(x)); - const xla::PrimitiveType type = x_shape.element_type(); - - std::vector batch_dim_ids(batch_dims.size()); - std::iota(batch_dim_ids.begin(), batch_dim_ids.end(), 0); - const int64 minor_dim = batch_dims.size(); - - xla::XlaOp zero = xla::ScalarLike(x, 0.0); - xla::XlaOp one = xla::ScalarLike(x, 1.0); - - // alpha = x[k] - xla::XlaOp alpha = - xla::Reshape(DynamicSliceInMinorDims(x, {k}, {1}), batch_dims); - - // Compute x[k+1:] (padded with zeros in elements 0..k) - xla::XlaOp iota = xla::Iota(builder, xla::S32, m); - xla::XlaOp x_after_k = - xla::Mul(x, xla::ConvertElementType(xla::Gt(iota, k), type), - /*broadcast_dimensions=*/{minor_dim}); - - // sigma = np.dot(x[k+1:], x[k+1:]) - auto sigma = - xla::Reduce(x_after_k * x_after_k, zero, - xla::CreateScalarAddComputation(type, builder), {minor_dim}); - // mu = np.sqrt(x[k]*x[k] + sigma) - auto mu = xla::Sqrt(xla::Square(alpha) + sigma); - - auto sigma_is_zero = xla::Eq(sigma, zero); - - *beta = xla::Select(sigma_is_zero, alpha, -xla::Sign(alpha) * mu); - *tau = xla::Select(sigma_is_zero, xla::Broadcast(zero, batch_dims), - (*beta - alpha) / *beta); - auto divisor = xla::Select(sigma_is_zero, xla::Broadcast(one, batch_dims), - alpha - *beta); - - auto e_k = xla::Broadcast(xla::ConvertElementType(xla::Eq(iota, k), type), - std::vector(batch_dims.size(), 1)); - - // Form v as [0, 0, ..., 1] ++ x[k+1:] / divisor - // If sigma is zero, x[k+1:] is zero, so use any non-zero divisor. - *v = e_k + - xla::Div(x_after_k, divisor, /*broadcast_dimensions=*/batch_dim_ids); - return Status::OK(); -} - -// Householder QR decomposition. Algorithm 5.2.1 from Golub and Van -// Loan "Matrix Computations", 4th Edition. This is an unblocked implementation -// used as an inner routine of the blocked implementation. -// Algorithm is adapted slightly so the shapes inside the loop are static, at -// the cost of some redundant computation. Since this is used as an inner block -// kernel, accumulates the Householder transformations (vs, taus) rather than -// the matrix q. -// Equivalent Python code, without batching: -// def qr(a): -// m = a.shape[0] -// n = a.shape[1] -// vs = np.zeros([m, n]) -// taus = np.zeros([n]) -// for j in xrange(min(m, n)): -// v, tau, beta = house(a[:, j], j) -// # Unusually, we apply the Householder transformation to the entirety of -// # a, wasting FLOPs to maintain the static shape invariant that XLA -// # requires. For columns that precede j this has no effect. -// a[:, :] -= tau * np.dot(v[:, np.newaxis], -// np.dot(v[np.newaxis, :], a[:, :])) -// # Form column j explicitly rather than relying on the precision of the -// # Householder update. -// a[j, j] = beta -// a[j+1:, j] = np.zeros([m - j - 1], dtype=a.dtype) -// vs[:, j] = v -// taus[j] = tau -// return (q, vs, taus) -struct QRBlockResult { - // The factored R value - xla::XlaOp r; - - // Representation of the Householder matrices I - beta v v.T - xla::XlaOp taus; // Shape: [..., n] - xla::XlaOp vs; // Shape: [..., m, n] -}; -xla::StatusOr QRBlock( - xla::XlaOp a, xla::PrecisionConfig::Precision precision) { - xla::XlaBuilder* builder = a.builder(); - TF_ASSIGN_OR_RETURN(xla::Shape a_shape, builder->GetShape(a)); - const int num_dims = xla::ShapeUtil::Rank(a_shape); - if (num_dims < 2) { - return errors::InvalidArgument("Arguments to QR must have rank >= 2: ", - num_dims); - } - xla::PrimitiveType type = a_shape.element_type(); - - const int64 m = xla::ShapeUtil::GetDimension(a_shape, -2); - const int64 n = xla::ShapeUtil::GetDimension(a_shape, -1); - - 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] = xla::ShapeUtil::GetDimension(a_shape, i); - } - - std::vector batch_dim_indices(num_batch_dims); - std::iota(batch_dim_indices.begin(), batch_dim_indices.end(), 0); - - auto qr_body_fn = - [&](xla::XlaOp j, absl::Span values, - xla::XlaBuilder* builder) -> xla::StatusOr> { - auto a = values[0]; - auto vs = values[1]; - auto taus = values[2]; - - // v, beta = house(a[:, j], j) - auto x = DynamicSliceInMinorDims(a, {j}, {1}); - xla::XlaOp v, tau, beta; - TF_RETURN_IF_ERROR(House(xla::Collapse(x, {num_dims - 2, num_dims - 1}), j, - batch_dims, m, &v, &tau, &beta)); - - std::vector shape = batch_dims; - shape.push_back(1); - shape.push_back(m); - auto v_broadcast = xla::Reshape(v, shape); - // a[:, :] -= tau * np.dot(v[:, np.newaxis], - // np.dot(v[np.newaxis, :], a[:, :])) - auto vva = - BatchDot(v_broadcast, a, /*transpose_x=*/false, /*transpose_y=*/false, - /*conjugate_x=*/false, /*conjugate_y=*/false, precision); - vva = - BatchDot(v_broadcast, vva, /*transpose_x=*/true, /*transpose_y=*/false, - /*conjugate_x=*/false, /*conjugate_y=*/false, precision); - a = a - xla::Mul(tau, vva, - /*broadcast_dimensions=*/batch_dim_indices); - - // It is more precise to populate column 'k' explicitly, rather than - // computing it implicitly by applying the Householder transformation. - // a[k,k] = beta - // a[k+1:,k] = np.zeros([m-k-1], dtype=a.dtype) - auto iota = xla::Reshape(xla::Iota(a.builder(), xla::S32, m), {m, 1}); - auto predecessor_mask = xla::ConvertElementType(xla::Lt(iota, j), type); - auto mask = xla::Broadcast(xla::ConvertElementType(xla::Eq(iota, j), type), - std::vector(batch_dims.size(), 1)); - auto new_x = - xla::Mul(x, predecessor_mask, - /*broadcast_dimensions=*/{num_dims - 2, num_dims - 1}) + - xla::Mul(beta, mask, /*broadcast_dimensions=*/batch_dim_indices); - a = DynamicUpdateSliceInMinorDims(a, new_x, {j}); - - // vs[:, j] = v - vs = DynamicUpdateSliceInMinorDims( - vs, xla::Reshape(v, ConcatVectors(batch_dims, {m, 1})), {j}); - // taus[j] = tau - taus = DynamicUpdateSliceInMinorDims( - taus, xla::Reshape(tau, ConcatVectors(batch_dims, {1})), {j}); - return std::vector{a, vs, taus}; - }; - - auto vs = xla::Zeros(builder, xla::ShapeUtil::MakeShape( - type, ConcatVectors(batch_dims, {m, n}))); - auto taus = xla::Zeros( - builder, xla::ShapeUtil::MakeShape(type, ConcatVectors(batch_dims, {n}))); - - TF_ASSIGN_OR_RETURN(auto values, - XlaForEachIndex(std::min(m, n), xla::S32, qr_body_fn, - {a, vs, taus}, "qr", builder)); - - QRBlockResult result; - result.r = values[0]; - result.vs = values[1]; - result.taus = values[2]; - return result; -} - -// Computes W and Y such that I-WY is equivalent to the sequence of Householder -// transformations given by vs and taus. -// Golub and van Loan, "Matrix Computations", algorithm 5.1.2. -// Y = np.zeros([m, n]) -// W = np.zeros([m, n]) -// Y[:, 0] = vs[:, 0] -// W[:, 0] = -taus[0] * vs[:, 0] -// for j in xrange(1, n): -// v = vs[:, j] -// z = -taus[j] * v - taus[j] * np.dot(W, np.dot(Y.T, v)) -// W[:, j] = z -// Y[:, j] = v -// return W -// There is no need to return Y since at termination of the loop it is equal to -// vs. -xla::StatusOr ComputeWYRepresentation( - xla::PrimitiveType type, absl::Span batch_dims, xla::XlaOp vs, - xla::XlaOp taus, int64 m, int64 n, - xla::PrecisionConfig::Precision precision) { - std::vector batch_dim_indices(batch_dims.size()); - std::iota(batch_dim_indices.begin(), batch_dim_indices.end(), 0); - int64 n_index = batch_dims.size() + 1; - - auto body_fn = - [&](xla::XlaOp j, absl::Span values, - xla::XlaBuilder* builder) -> xla::StatusOr> { - auto w = values[0]; - auto y = values[1]; - const auto vs = values[2]; - const auto taus = values[3]; - - // Want j values in range [1, ... n). - j = j + xla::ConstantR0(builder, 1); - // vs has shape [..., m, 1] - auto v = DynamicSliceInMinorDims(vs, {j}, {1}); - // beta has shape [..., 1] - auto beta = DynamicSliceInMinorDims(taus, {j}, {1}); - - // yv has shape [..., n, 1] - auto yv = BatchDot(y, v, /*transpose_x=*/true, /*transpose_y=*/false, - /*conjugate_x=*/false, /*conjugate_y=*/false, precision); - // wyv has shape [..., m, 1] - auto wyv = - BatchDot(w, yv, /*transpose_x=*/false, /*transpose_y=*/false, - /*conjugate_x=*/false, /*conjugate_y=*/false, precision); - - auto z = xla::Mul( - -beta, v + wyv, - /*broadcast_dimensions=*/ConcatVectors(batch_dim_indices, {n_index})); - - w = DynamicUpdateSliceInMinorDims(w, z, {j}); - y = DynamicUpdateSliceInMinorDims(y, v, {j}); - - return std::vector{w, y, vs, taus}; - }; - - xla::XlaBuilder* builder = vs.builder(); - auto w = xla::Zeros(builder, xla::ShapeUtil::MakeShape( - type, ConcatVectors(batch_dims, {m, n}))); - auto y = w; - auto v = SliceInMinorDims(vs, {0}, {1}); - auto beta = SliceInMinorDims(taus, {0}, {1}); - y = UpdateSliceInMinorDims(y, v, {0}); - auto bv = xla::Mul( - -beta, v, - /*broadcast_dimensions=*/ConcatVectors(batch_dim_indices, {n_index})); - w = UpdateSliceInMinorDims(w, bv, {0}); - - TF_ASSIGN_OR_RETURN( - auto values, XlaForEachIndex(n - 1, xla::S32, body_fn, {w, y, vs, taus}, - "wy", builder)); - return values[0]; -} - -} // namespace - -// Block Householder QR Factorization. Algorithm 5.2.2 of Golub and van Loan. -// def qr_blocked(a, block_size): -// m = a.shape[0] -// n = a.shape[1] -// q = np.eye(m) -// for i in xrange(0, min(m, n), block_size): -// k = min(block_size, min(m, n) - s) -// (a, vs, taus) = qr(a[i:, i:i+k]) -// y = vs -// w = ComputeWYRepresentation(vs, taus, m-i, k) -// a[i:, i+r:] += np.dot(y, np.dot(w.T, a[i:, i+k:])) -// q[:, i:] += np.dot(q[:, i:], np.dot(w, y.T)) -// return (q, a) -// TODO(phawkins): consider using UT transformations (in the form I - V U V') -// rather than WY transformations. -xla::StatusOr QRDecomposition( - xla::XlaOp a, bool full_matrices, int64 block_size, - xla::PrecisionConfig::Precision precision) { - xla::XlaBuilder* builder = a.builder(); - TF_ASSIGN_OR_RETURN(xla::Shape a_shape, builder->GetShape(a)); - const int num_dims = xla::ShapeUtil::Rank(a_shape); - if (num_dims < 2) { - return errors::InvalidArgument("Arguments to QR must have rank >= 2: ", - num_dims); - } - xla::PrimitiveType type = a_shape.element_type(); - - const int64 m = xla::ShapeUtil::GetDimension(a_shape, -2); - const int64 n = xla::ShapeUtil::GetDimension(a_shape, -1); - const int64 p = std::min(m, n); - - if (block_size < 1) { - return errors::InvalidArgument( - "block_size argument to QR must be >= 1; got ", block_size); - } - - 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] = xla::ShapeUtil::GetDimension(a_shape, i); - } - - auto q = xla::Broadcast(xla::IdentityMatrix(builder, type, m, m), batch_dims); - for (int64 i = 0; i < p; i += block_size) { - int64 k = std::min(block_size, p - i); - - auto a_block = SliceInMinorDims(a, {i, i}, {m, i + k}); - TF_ASSIGN_OR_RETURN(auto qr_block, QRBlock(a_block, precision)); - - a = UpdateSliceInMinorDims(a, qr_block.r, {i, i}); - - // Compute the I-WY block representation of a product of Householder - // matrices. - TF_ASSIGN_OR_RETURN( - auto w, ComputeWYRepresentation(type, batch_dims, qr_block.vs, - qr_block.taus, m - i, k, precision)); - auto y = qr_block.vs; - - // a[i:, i+k:] += np.dot(Y, np.dot(W.T, a[i:, i+k:])) - auto a_panel = SliceInMinorDims(a, {i, i + k}, {m, n}); - auto a_update = - BatchDot(w, a_panel, /*transpose_x=*/true, /*transpose_y=*/false, - /*conjugate_x=*/false, /*conjugate_y=*/false, precision); - a_update = - BatchDot(y, a_update, /*transpose_x=*/false, /*transpose_y=*/false, - /*conjugate_x=*/false, /*conjugate_y=*/false, precision); - a_panel = a_panel + a_update; - a = UpdateSliceInMinorDims(a, a_panel, {i, i + k}); - - // q[:, i:] += np.dot(np.dot(q[:, i:], W), Y.T)) - auto q_panel = SliceInMinorDims(q, {0, i}, {m, m}); - auto q_update = - BatchDot(q_panel, w, /*transpose_x=*/false, /*transpose_y=*/false, - /*conjugate_x=*/false, /*conjugate_y=*/false, precision); - q_update = BatchDot(q_update, y, /*transpose_x=*/false, - /*transpose_y=*/true, /*conjugate_x=*/false, - /*conjugate_y=*/false, precision); - q_panel = q_panel + q_update; - q = UpdateSliceInMinorDims(q, q_panel, {0, i}); - } - QRDecompositionResult result; - - // full_matrices is false when only a partial result in needed. Slice to the - // needed dimensions here. - if (!full_matrices) { - q = SliceInMinorDims(q, {0, 0}, {m, p}); - a = SliceInMinorDims(a, {0, 0}, {p, n}); - } - result.q = q; - result.r = a; - return result; -} - -} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/lib/qr.h b/tensorflow/compiler/tf2xla/lib/qr.h deleted file mode 100644 index 24b537ac8b63b93e734c3d0e335ea455f7d51a54..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/tf2xla/lib/qr.h +++ /dev/null @@ -1,42 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#ifndef TENSORFLOW_COMPILER_TF2XLA_LIB_QR_H_ -#define TENSORFLOW_COMPILER_TF2XLA_LIB_QR_H_ - -#include "tensorflow/compiler/xla/client/xla_builder.h" -#include "tensorflow/compiler/xla/xla_data.pb.h" - -namespace tensorflow { - -// Computes the QR decompositions of a batch of matrices. That is, -// given a (batched) matrix a, computes an orthonormal matrix Q and an -// upper-triangular matrix R such that a = QR. -// `a` must be a (batched) matrix of size [..., m, n]. -// The algorithm implements a blocked QR decomposition; `block_size` is -// the block size to use. -// TODO(phawkins): handle the complex case. -struct QRDecompositionResult { - xla::XlaOp q; - xla::XlaOp r; -}; - -xla::StatusOr QRDecomposition( - xla::XlaOp a, bool full_matrices, int64 block_size = 128, - xla::PrecisionConfig::Precision precision = xla::PrecisionConfig::HIGHEST); - -} // namespace tensorflow - -#endif // TENSORFLOW_COMPILER_TF2XLA_LIB_QR_H_ diff --git a/tensorflow/compiler/tf2xla/lib/scatter.cc b/tensorflow/compiler/tf2xla/lib/scatter.cc index 2b1c2ced925d9fee7392986015a6e716a94d356f..1cd5a79171dccd57fc1b7941cdf16417301ff7f8 100644 --- a/tensorflow/compiler/tf2xla/lib/scatter.cc +++ b/tensorflow/compiler/tf2xla/lib/scatter.cc @@ -20,7 +20,6 @@ limitations under the License. #include "absl/types/span.h" #include "tensorflow/compiler/tf2xla/lib/util.h" -#include "tensorflow/compiler/tf2xla/lib/while_loop.h" #include "tensorflow/compiler/xla/client/lib/arithmetic.h" #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/compiler/xla/literal.h" @@ -49,7 +48,7 @@ xla::StatusOr XlaScatter( if (indices_are_vectors) { TF_RET_CHECK(!indices_dims.empty()); num_index_dims = indices_dims.back(); - if (num_index_dims > xla::ShapeUtil::Rank(buffer_shape)) { + if (num_index_dims > buffer_shape.rank()) { return errors::InvalidArgument( "The size of the minor dimension of the indices (shape: ", xla::ShapeUtil::HumanString(indices_shape), @@ -141,8 +140,8 @@ xla::StatusOr XlaScatter( ? indices_shape.dimensions_size() - 1 : indices_shape.dimensions_size()); - int64 updates_rank = xla::ShapeUtil::Rank(updates_shape); - int64 buffer_rank = xla::ShapeUtil::Rank(buffer_shape); + int64 updates_rank = updates_shape.rank(); + int64 buffer_rank = buffer_shape.rank(); int64 num_window_dims_in_updates = buffer_rank - num_index_dims; // If the rank of `updates` is 0 and does not match the expected rank of @@ -157,7 +156,7 @@ xla::StatusOr XlaScatter( if (updates_rank == 0 && expected_updates_rank != 0) { new_updates = xla::Broadcast(updates, expected_updates_dims); TF_ASSIGN_OR_RETURN(updates_shape, builder->GetShape(new_updates)); - updates_rank = xla::ShapeUtil::Rank(updates_shape); + updates_rank = updates_shape.rank(); } if (updates_rank > 0) { diff --git a/tensorflow/compiler/tf2xla/lib/triangular_solve.cc b/tensorflow/compiler/tf2xla/lib/triangular_solve.cc deleted file mode 100644 index 6524c2a9b1ada632d80edd234272760c2b545cc4..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/tf2xla/lib/triangular_solve.cc +++ /dev/null @@ -1,416 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/compiler/tf2xla/lib/triangular_solve.h" - -#include -#include - -#include "tensorflow/compiler/tf2xla/lib/batch_dot.h" -#include "tensorflow/compiler/tf2xla/lib/util.h" -#include "tensorflow/compiler/xla/client/lib/constants.h" -#include "tensorflow/compiler/xla/client/lib/numeric.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/shape_util.h" -#include "tensorflow/compiler/xla/status_macros.h" -#include "tensorflow/compiler/xla/statusor.h" -#include "tensorflow/compiler/xla/util.h" -#include "tensorflow/core/lib/core/errors.h" -#include "tensorflow/core/lib/math/math_util.h" - -namespace tensorflow { - -// Get the diagonal blocks of the coefficient matrix -xla::XlaOp DiagonalBlocks(xla::XlaOp a, int64 block_size) { - xla::XlaBuilder* builder = a.builder(); - return builder->ReportErrorOrReturn([&]() -> xla::StatusOr { - TF_ASSIGN_OR_RETURN(xla::Shape shape, builder->GetShape(a)); - int ndims = xla::ShapeUtil::Rank(shape); - int64 n = xla::ShapeUtil::GetDimension(shape, -1); - int64 num_blocks = n / block_size; - - xla::XlaOp diag_blocks; - - // If the coefficient matrix is exactly the block size, we just add a - // singleton dimension i.e. [..., n, n] -> [..., 1, n, n] - if (n == block_size) { - std::vector permutation(ndims); - std::iota(permutation.begin(), permutation.end(), 1); - permutation.insert(permutation.end() - 2, 0); - return Transpose(Broadcast(a, /*broadcast_sizes=*/{1}), permutation); - } - - // We can grab entire blocks using gather - if (n > block_size) { - // Construct the starting indices of the diagonal blocks - auto start_indices = - Transpose(Broadcast(Mul(Iota(builder, xla::S32, num_blocks), - xla::ConstantR0(builder, block_size)), - /*broadcast_sizes=*/{2}), - /*permutation=*/{1, 0}); - - // Gather the diagonal blocks - xla::GatherDimensionNumbers dim_numbers; - dim_numbers.add_offset_dims(ndims - 1); - dim_numbers.add_offset_dims(ndims); - dim_numbers.add_start_index_map(ndims - 2); - dim_numbers.add_start_index_map(ndims - 1); - dim_numbers.set_index_vector_dim(1); - diag_blocks = Gather(a, start_indices, dim_numbers, - /*slice_sizes=*/{block_size, block_size}); - } - - // The last block might be smaller than the block size, - // so we will need to pad it - if (n % block_size != 0) { - // Pad with zeros - auto last_blocks = - SliceInMinorDims(a, {n - n % block_size, n - n % block_size}, {n, n}); - xla::PaddingConfig config = xla::MakeNoPaddingConfig(ndims); - int64 padding = block_size - n % block_size; - config.mutable_dimensions(ndims - 1)->set_edge_padding_high(padding); - config.mutable_dimensions(ndims - 2)->set_edge_padding_high(padding); - last_blocks = - Pad(last_blocks, Zero(builder, shape.element_type()), config); - - // Add a singleton dimension - // i.e. [..., block_size, block_size] -> [..., 1, block_size, block_size] - TF_ASSIGN_OR_RETURN(xla::Shape blocks_shape, - builder->GetShape(last_blocks)); - auto shape_dims = xla::AsInt64Slice(blocks_shape.dimensions()); - auto last_blocks_dims = std::vector(ndims); - std::copy(shape_dims.begin(), shape_dims.end(), last_blocks_dims.begin()); - last_blocks_dims.insert(last_blocks_dims.end() - 2, 1); - last_blocks = Reshape(last_blocks, last_blocks_dims); - - // Concatenate with the other blocks if necessary - if (n > block_size) { - diag_blocks = - xla::ConcatInDim(builder, {diag_blocks, last_blocks}, ndims - 2); - } else { - diag_blocks = last_blocks; - } - } - - return diag_blocks; - }); -} - -xla::XlaOp InvertDiagonalBlocks(xla::XlaOp diag_blocks, bool lower, - bool transpose_a, bool conjugate_a, - xla::PrecisionConfig::Precision precision) { - xla::XlaBuilder* builder = diag_blocks.builder(); - return builder->ReportErrorOrReturn([&]() -> xla::StatusOr { - // Input is a batch of square lower triangular square matrices. Its shape is - // (..., size, size). We resize this to (num_blocks, size, size). - TF_ASSIGN_OR_RETURN(xla::Shape shape, builder->GetShape(diag_blocks)); - int64 block_size = xla::ShapeUtil::GetDimension(shape, -1); - int64 num_blocks = xla::ShapeUtil::ElementsIn(shape) / - tensorflow::MathUtil::IPow(block_size, 2); - diag_blocks = Reshape(diag_blocks, {num_blocks, block_size, block_size}); - - // The input must be triangular because we rely on that when doing - // multiplications later on - diag_blocks = Triangle(diag_blocks, /*lower=*/lower); - - // Rescale blocks to be unit triangular, but avoid dividing by - // 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(xla::Shape diags_shape, builder->GetShape(diags)); - auto one = ScalarLike(diags, 1); - auto ones = Broadcast(one, xla::AsInt64Slice(diags_shape.dimensions())); - diags = Select(Eq(diags, Zero(builder, shape.element_type())), ones, diags); - auto scaled_diag_blocks = Div(diag_blocks, diags, {0, 2}); - - // We can now use the fact that for an upper triangular matrix - // [[L11, 0], [L21, L22]], given the inverses L11' and L22', we have - // L22' = -L22' * L21 * L11'. In our case, L21 is a vector and our blocks - // have been rescaled to be unit triangular, so L22 = L22' = 1. - - // Initialize the output matrix with -1s on the diagonal. We use -1 instead - // of 1 because we cannot do matrix-vector multiplies with variable shapes - // inside of a loop, or do irregularly shaped in-place updates. Hence, - // L21 <- -L22 * L21 * L11 cannot be done naively. Instead, we update the - // entire row i.e. we calculate - // [L21 L22 0] <- -[L21 L22 0] @ diag_blocks([L11', -I, -I]) - // which means [L21 L22 0] <- [-L21 * L11', L22, 0]. - auto identity = - IdentityMatrix(builder, shape.element_type(), block_size, block_size); - auto neg_identity = -identity; - - // 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=*/xla::ConstantR1(builder, 2, start_index)); - - // Broadcast diag([1, -1, -1, ...]) to every block - xla::XlaOp output = Broadcast(output_block, - /*broadcast_sizes=*/{num_blocks}); - - // Now we construct a loop that performs matrix-vector multiplications - // inverting the blocks one row at a time - std::vector tuple_shapes = { - // The loop iteration counter is a scalar, incremented each iteration. - xla::ShapeUtil::MakeShape(xla::S32, {}), - // The output has the shape of A, with one row updated each iteration. - xla::ShapeUtil::MakeShape(shape.element_type(), - {num_blocks, block_size, block_size}), - // The input is a loop invariant. - xla::ShapeUtil::MakeShape(shape.element_type(), - {num_blocks, block_size, block_size})}; - xla::Shape tuple_shape = xla::ShapeUtil::MakeTupleShape(tuple_shapes); - - auto init_i = One(builder, xla::S32); - auto init = xla::Tuple(builder, {init_i, output, scaled_diag_blocks}); - - // Construct the loop condition function. - std::unique_ptr condb = - builder->CreateSubBuilder("InvertDiagCond"); - { - auto i = GetTupleElement( - Parameter(condb.get(), 0, tuple_shape, "InvertDiagCondTuple"), 0); - Lt(i, xla::ConstantR0(condb.get(), block_size)); - } - TF_ASSIGN_OR_RETURN(auto cond, condb->Build()); - - // Construct the loop body function. - std::unique_ptr bodyb = - builder->CreateSubBuilder("InvertDiagBody"); - { - auto input_tuple = - Parameter(bodyb.get(), 0, tuple_shape, "InvertDiagBodyTuple"); - - auto i = GetTupleElement(input_tuple, 0); - auto body_out = GetTupleElement(input_tuple, 1); - auto body_input = GetTupleElement(input_tuple, 2); - - auto zero = xla::ConstantR1(bodyb.get(), 1, 0); - auto j = (lower) ? i : ScalarLike(i, block_size - 1) - i; - auto start_indices = - xla::ConcatInDim(bodyb.get(), {zero, Reshape(j, {1}), zero}, 0); - auto input_row = - DynamicSlice(body_input, start_indices, - /*slice_sizes=*/{num_blocks, 1, block_size}); - - // We want -L21 L11^{-1} - xla::DotDimensionNumbers dnums; - dnums.add_lhs_batch_dimensions(0); - dnums.add_rhs_batch_dimensions(0); - dnums.add_lhs_contracting_dimensions(2); - dnums.add_rhs_contracting_dimensions(1); - xla::PrecisionConfig precision_proto; - precision_proto.add_operand_precision(precision); - precision_proto.add_operand_precision(precision); - auto update = -DotGeneral(input_row, body_out, dnums, &precision_proto); - - body_out = DynamicUpdateSlice(body_out, update, start_indices); - - auto next_i = i + ScalarLike(i, 1); - xla::Tuple(bodyb.get(), {next_i, body_out, body_input}); - } - TF_ASSIGN_OR_RETURN(auto body, bodyb->Build()); - - // Construct the While loop and return the result, - // return while_loop(cond_fun, body_fun, init)[1] - auto invert_while = While(cond, body, init); - auto inv_diag_blocks = GetTupleElement(invert_while, 1); - - // Undo the scaling - inv_diag_blocks = Div(inv_diag_blocks, diags, - /*broadcast_dimensions=*/{0, 1}); - - // Reshape back to original batch major dimensions - return Reshape(inv_diag_blocks, xla::AsInt64Slice(shape.dimensions())); - }); -} - -xla::XlaOp SolveWithInvertedDiagonalBlocks( - xla::XlaOp a, xla::XlaOp b, xla::XlaOp inv_diag_blocks, bool left_side, - bool lower, bool transpose_a, bool conjugate_a, - xla::PrecisionConfig::Precision precision) { - xla::XlaBuilder* builder = a.builder(); - return builder->ReportErrorOrReturn([&]() -> xla::StatusOr { - TF_ASSIGN_OR_RETURN(xla::Shape blocks_shape, - builder->GetShape(inv_diag_blocks)); - TF_ASSIGN_OR_RETURN(xla::Shape b_shape, builder->GetShape(b)); - int64 block_size = xla::ShapeUtil::GetDimension(blocks_shape, -1); - - TF_ASSIGN_OR_RETURN(xla::Shape a_shape, builder->GetShape(a)); - int64 ndims = xla::ShapeUtil::Rank(a_shape); - int64 n = xla::ShapeUtil::GetDimension(a_shape, -1); - int64 num_blocks = n / block_size + (n % block_size != 0); - int64 m_dim = (left_side) ? -1 : -2; - int64 m = xla::ShapeUtil::GetDimension(b_shape, m_dim); - - // Initialize the solution - auto x = ZerosLike(b); - - // This loop is unrolled for performance reasons, but it could be expressed - // rolled as well since the matrices are of the same size each iteration - for (int i = 0; i < num_blocks; i++) { - // High-level intuition: We have B[i] = L[i] @ X. Since L is upper - // triangular this means B[i] = L[i, :i + 1] @ X[:i + 1]. We can split - // this into two parts: B[i] = L[i, :i] @ X[:i] + L[i, i] @ X[i] which - // can be solved for X[i] as X[i] = inv(L[i, i]) @ B[i] - L[i, :i] @ X[:i] - - // Decide whether we go from first block to last or vice versa - auto j = (left_side ^ lower ^ transpose_a) ? num_blocks - 1 - i : i; - - // Get the size of the inverse blocks (the last one might be smaller) - int64 block = (n % block_size != 0 && j + 1 == num_blocks) - ? n % block_size - : block_size; - auto inv_block = - MaybeConjugate(Collapse(SliceInMinorDims(inv_diag_blocks, {j, 0, 0}, - {j + 1, block, block}), - /*dimensions=*/{ndims - 2, ndims - 1}), - conjugate_a); - - // Get the corresponding row of B - int64 k = std::min((j + 1) * block_size, n); - std::vector start = {j * block_size, 0}; - std::vector end = {k, m}; - if (!left_side) { - std::swap(start[0], start[1]); - std::swap(end[0], end[1]); - } - auto b_row = SliceInMinorDims(b, start, end); - - xla::XlaOp remainder; - if (i == 0) { - remainder = b_row; - } else { - // This matrix multiply involves a lot of multiplying with zero (namely, - // X[i * block_size:] = 0), but this is faster than slicing... - end = {k, n}; - if (!left_side) { - std::swap(end[0], end[1]); - } - if (transpose_a) { - std::swap(start[0], start[1]); - std::swap(end[0], end[1]); - } - auto a_row = - MaybeConjugate(SliceInMinorDims(a, start, end), conjugate_a); - if (left_side) { - remainder = b_row - BatchDot(a_row, x, transpose_a, false, - /*conjugate_x=*/false, - /*conjugate_y=*/false, precision); - } else { - remainder = b_row - BatchDot(x, a_row, false, transpose_a, - /*conjugate_x=*/false, - /*conjugate_y=*/false, precision); - } - } - - xla::XlaOp x_update; - auto zero = Zero(builder, xla::S32); - auto start_index = - xla::ConstantR0WithType(builder, xla::S32, j * block_size); - std::vector update_starts = {start_index, zero}; - if (left_side) { - x_update = - BatchDot(inv_block, remainder, transpose_a, false, - /*conjugate_x=*/false, /*conjugate_y=*/false, precision); - } else { - x_update = - BatchDot(remainder, inv_block, false, transpose_a, - /*conjugate_x=*/false, /*conjugate_y=*/false, precision); - std::swap(update_starts[0], update_starts[1]); - } - x = DynamicUpdateSliceInMinorDims(x, x_update, /*starts=*/update_starts); - } - - return x; - }); -} - -xla::XlaOp TriangularSolve(xla::XlaOp a, xla::XlaOp b, bool left_side, - bool lower, bool transpose_a, bool conjugate_a, - int64 block_size, - xla::PrecisionConfig::Precision precision) { - xla::XlaBuilder* builder = a.builder(); - return builder->ReportErrorOrReturn([&]() -> xla::StatusOr { - TF_ASSIGN_OR_RETURN(xla::Shape a_shape, builder->GetShape(a)); - TF_ASSIGN_OR_RETURN(xla::Shape b_shape, builder->GetShape(b)); - if (xla::ShapeUtil::Rank(a_shape) != xla::ShapeUtil::Rank(b_shape)) { - return errors::InvalidArgument( - "Arguments to TriangularSolve have different ranks: ", - xla::ShapeUtil::HumanString(a_shape), " vs. ", - xla::ShapeUtil::HumanString(b_shape)); - } - const int64 ndims = xla::ShapeUtil::Rank(a_shape); - if (ndims < 2) { - return errors::InvalidArgument( - "Arguments to TriangularSolve must have rank >= 2: ", ndims); - } - // The batch dimensions must be equal. - std::vector batch_dimensions; - for (int i = 0; i < ndims - 2; ++i) { - int64 a_size = a_shape.dimensions(i); - int64 b_size = b_shape.dimensions(i); - if (a_size != b_size) { - return errors::InvalidArgument( - "Batch dimensions of arguments to TriangularSolve must be equal: ", - xla::ShapeUtil::HumanString(a_shape), " vs ", - xla::ShapeUtil::HumanString(b_shape)); - } - batch_dimensions.push_back(a_size); - } - - if (xla::ShapeUtil::GetDimension(a_shape, -1) != - xla::ShapeUtil::GetDimension(a_shape, -2)) { - return errors::InvalidArgument( - "The 'a' arguments to TriangularSolve must be square matrices: ", - xla::ShapeUtil::HumanString(a_shape)); - } - const int64 m = xla::ShapeUtil::GetDimension(b_shape, -2); - const int64 n = xla::ShapeUtil::GetDimension(b_shape, -1); - if ((left_side ? m : n) != xla::ShapeUtil::GetDimension(a_shape, -1)) { - return errors::InvalidArgument( - "Arguments to TriangularSolve have incompatible matrix shapes: ", - xla::ShapeUtil::HumanString(a_shape), " vs ", - xla::ShapeUtil::HumanString(b_shape)); - } - - if (block_size < 1) { - return errors::InvalidArgument( - "block_size argument to TriangularSolve must be >= 1; got ", - block_size); - } - - // We find the diagonal blocks of the coefficient matrix - auto diag_blocks = DiagonalBlocks(a, block_size); - - // We invert these blocks in parallel using batched matrix-vector products - auto inv_diag_blocks = InvertDiagonalBlocks(diag_blocks, lower, transpose_a, - conjugate_a, precision); - - // We now find the solution using GEMMs - auto x = - SolveWithInvertedDiagonalBlocks(a, b, inv_diag_blocks, left_side, lower, - transpose_a, conjugate_a, precision); - - return x; - }); -} - -} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/lib/triangular_solve_test.cc b/tensorflow/compiler/tf2xla/lib/triangular_solve_test.cc deleted file mode 100644 index aeebf16028d40189203cdfd815f06a339ee72902..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/tf2xla/lib/triangular_solve_test.cc +++ /dev/null @@ -1,333 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/compiler/tf2xla/lib/triangular_solve.h" - -#include -#include -#include - -#include "tensorflow/compiler/xla/array2d.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/types.h" -#include "tensorflow/core/lib/core/status_test_util.h" - -namespace tensorflow { -namespace { - -using TriangularSolveTest = xla::ClientLibraryTestBase; -using TriangularSolveLeftLookingTest = xla::ClientLibraryTestBase; -using complex64 = xla::complex64; - -xla::Array2D AValsLower() { - return {{2, 0, 0, 0}, {3, 6, 0, 0}, {4, 7, 9, 0}, {5, 8, 10, 11}}; -} - -xla::Array2D AValsUpper() { - return {{2, 3, 4, 5}, {0, 6, 7, 8}, {0, 0, 9, 10}, {0, 0, 0, 11}}; -} - -xla::Array2D BValsRight() { - return {{1, 2, 3, 4}, {5, 6, 7, 8}, {9, 10, 11, 12}}; -} - -xla::Array2D BValsLeft() { - return {{1, 2, 3}, {4, 5, 6}, {7, 8, 9}, {10, 11, 12}}; -} - -xla::Array2D AValsLowerComplex() { - return {{2, 0, 0, 0}, - {complex64(3, 1), 6, 0, 0}, - {4, complex64(7, 2), 9, 0}, - {5, 8, complex64(10, 3), 11}}; -} - -xla::Array2D AValsUpperComplex() { - return {{2, 3, complex64(4, 3), 5}, - {0, 6, complex64(7, 2), 8}, - {0, 0, complex64(9, 1), 10}, - {0, 0, 0, 11}}; -} - -xla::Array2D BValsRightComplex() { - return {{1, 2, 3, 4}, {5, 6, 7, 8}, {9, 10, 11, 12}}; -} - -xla::Array2D BValsLeftComplex() { - return {{1, 2, 3}, {4, 5, 6}, {7, 8, 9}, {10, 11, 12}}; -} - -xla::Array2D AValsFull() { - return {{2, 0, 1, 2}, {3, 6, 0, 1}, {4, 7, 9, 0}, {5, 8, 10, 11}}; -} - -XLA_TEST_F(TriangularSolveTest, SimpleRightLowerTranspose) { - xla::XlaBuilder builder(TestName()); - - xla::XlaOp a, b; - auto a_data = CreateR2Parameter(AValsLower(), 0, "a", &builder, &a); - 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); - - xla::Array2D expected({ - {0.5, 0.08333334, 0.04629629, 0.03367003}, - {2.5, -0.25, -0.1388889, -0.1010101}, - {4.5, -0.58333331, -0.32407406, -0.23569024}, - }); - - ComputeAndCompareR2(&builder, expected, {a_data.get(), b_data.get()}, - xla::ErrorSpec(1e-2, 1e-2)); -} - -XLA_TEST_F(TriangularSolveTest, SimpleRightLowerNotranspose) { - xla::XlaBuilder builder(TestName()); - - xla::XlaOp a, b; - auto a_data = CreateR2Parameter(AValsLower(), 0, "a", &builder, &a); - 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); - - xla::Array2D expected({ - {-0.16414141, -0.06902357, -0.07070707, 0.36363636}, - {0.64393939, 0.06565657, -0.03030303, 0.72727273}, - {1.4520202, 0.2003367, 0.01010101, 1.09090909}, - }); - - ComputeAndCompareR2(&builder, expected, {a_data.get(), b_data.get()}, - xla::ErrorSpec(1e-2, 1e-2)); -} - -XLA_TEST_F(TriangularSolveTest, SimpleRightUpperTranspose) { - xla::XlaBuilder builder(TestName()); - - xla::XlaOp a, b; - auto a_data = CreateR2Parameter(AValsUpper(), 0, "a", &builder, &a); - 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); - - xla::Array2D expected({ - {-0.16414141, -0.06902357, -0.07070707, 0.36363636}, - {0.64393939, 0.06565657, -0.03030303, 0.72727273}, - {1.4520202, 0.2003367, 0.01010101, 1.09090909}, - }); - - ComputeAndCompareR2(&builder, expected, {a_data.get(), b_data.get()}, - xla::ErrorSpec(1e-2, 1e-2)); -} - -XLA_TEST_F(TriangularSolveTest, SimpleRightUpperNotranspose) { - xla::XlaBuilder builder(TestName()); - - xla::XlaOp a, b; - auto a_data = CreateR2Parameter(AValsUpper(), 0, "a", &builder, &a); - 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); - - xla::Array2D expected({ - {0.5, 0.08333334, 0.04629629, 0.03367003}, - {2.5, -0.25, -0.1388889, -0.1010101}, - {4.5, -0.58333331, -0.32407406, -0.23569024}, - }); - - ComputeAndCompareR2(&builder, expected, {a_data.get(), b_data.get()}, - xla::ErrorSpec(1e-2, 1e-2)); -} - -XLA_TEST_F(TriangularSolveTest, SimpleLeftLowerTranspose) { - xla::XlaBuilder builder(TestName()); - - xla::XlaOp a, b; - auto a_data = CreateR2Parameter(AValsLower(), 0, "a", &builder, &a); - 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); - - xla::Array2D expected({ - {-0.89646465, -0.69444444, -0.49242424}, - {-0.27441077, -0.24074074, -0.20707071}, - {-0.23232323, -0.22222222, -0.21212121}, - {0.90909091, 1., 1.09090909}, - }); - - ComputeAndCompareR2(&builder, expected, {a_data.get(), b_data.get()}, - xla::ErrorSpec(1e-2, 1e-2)); -} - -XLA_TEST_F(TriangularSolveTest, SimpleLeftLowerNotranspose) { - xla::XlaBuilder builder(TestName()); - - xla::XlaOp a, b; - auto a_data = CreateR2Parameter(AValsLower(), 0, "a", &builder, &a); - 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); - - xla::Array2D expected({ - {0.5, 1.0, 1.5}, - {0.41666667, 0.33333333, 0.25}, - {0.23148148, 0.18518519, 0.13888889}, - {0.16835017, 0.13468013, 0.1010101}, - }); - - ComputeAndCompareR2(&builder, expected, {a_data.get(), b_data.get()}, - xla::ErrorSpec(1e-2, 1e-2)); -} - -XLA_TEST_F(TriangularSolveTest, SimpleLeftLowerNotransposeIrregularblock) { - xla::XlaBuilder builder(TestName()); - - xla::XlaOp a, b; - auto a_data = CreateR2Parameter(AValsLower(), 0, "a", &builder, &a); - 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); - - xla::Array2D expected({ - {0.5, 1.0, 1.5}, - {0.41666667, 0.33333333, 0.25}, - {0.23148148, 0.18518519, 0.13888889}, - {0.16835017, 0.13468013, 0.1010101}, - }); - - ComputeAndCompareR2(&builder, expected, {a_data.get(), b_data.get()}, - xla::ErrorSpec(1e-2, 1e-2)); -} - -XLA_TEST_F(TriangularSolveTest, SimpleLeftUpperTranspose) { - xla::XlaBuilder builder(TestName()); - - xla::XlaOp a, b; - auto a_data = CreateR2Parameter(AValsUpper(), 0, "a", &builder, &a); - 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); - - xla::Array2D expected({ - {0.5, 1.0, 1.5}, - {0.41666667, 0.33333333, 0.25}, - {0.23148148, 0.18518519, 0.13888889}, - {0.16835017, 0.13468013, 0.1010101}, - }); - - ComputeAndCompareR2(&builder, expected, {a_data.get(), b_data.get()}, - xla::ErrorSpec(1e-2, 1e-2)); -} - -XLA_TEST_F(TriangularSolveTest, SimpleLeftUpperNotranspose) { - xla::XlaBuilder builder(TestName()); - - xla::XlaOp a, b; - auto a_data = CreateR2Parameter(AValsUpper(), 0, "a", &builder, &a); - 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); - - xla::Array2D expected({ - {-0.89646465, -0.69444444, -0.49242424}, - {-0.27441077, -0.24074074, -0.20707071}, - {-0.23232323, -0.22222222, -0.21212121}, - {0.90909091, 1., 1.09090909}, - }); - - ComputeAndCompareR2(&builder, expected, {a_data.get(), b_data.get()}, - xla::ErrorSpec(1e-2, 1e-2)); -} - -XLA_TEST_F(TriangularSolveTest, SimpleRightLowerTransposeConjugate) { - xla::XlaBuilder builder(TestName()); - - xla::XlaOp a, b; - auto a_data = - CreateR2Parameter(AValsLowerComplex(), 0, "a", &builder, &a); - auto b_data = - CreateR2Parameter(BValsRightComplex(), 1, "b", &builder, &b); - TriangularSolve(a, b, - /*left_side=*/false, /*lower=*/true, - /*transpose_a=*/true, /*conjugate_a=*/true, - /*block_size=*/2); - - xla::Array2D expected({ - {0.5, complex64(0.08333333, 0.08333333), - complex64(0.02777778, -0.0462963), complex64(0.06313131, -0.01094276)}, - {2.5, complex64(-0.25, 0.41666667), complex64(-0.23148148, -0.37962963), - complex64(0.08670034, -0.02104377)}, - {4.5, complex64(-0.58333333, 0.75), complex64(-0.49074074, -0.71296296), - complex64(0.11026936, -0.03114478)}, - }); - - ComputeAndCompareR2(&builder, expected, - {a_data.get(), b_data.get()}, - xla::ErrorSpec(1e-2, 1e-2)); -} - -XLA_TEST_F(TriangularSolveTest, SimpleLeftUpperTransposeNoconjugate) { - xla::XlaBuilder builder(TestName()); - - xla::XlaOp a, b; - auto a_data = - CreateR2Parameter(AValsUpperComplex(), 0, "a", &builder, &a); - auto b_data = - CreateR2Parameter(BValsLeftComplex(), 1, "b", &builder, &b); - TriangularSolve(a, b, - /*left_side=*/true, /*lower=*/false, - /*transpose_a=*/true, /*conjugate_a=*/false, - /*block_size=*/2); - - xla::Array2D expected({ - {0.5, 1., 1.5}, - {0.41666667, 0.33333333, 0.25}, - {complex64(0.20020325, -2.81504065e-01), - complex64(0.13821138, -4.22764228e-01), - complex64(0.07621951, -5.64024390e-01)}, - {complex64(0.19678492, 2.55912786e-01), - complex64(0.17738359, 3.84331116e-01), - complex64(0.15798226, 5.12749446e-01)}, - }); - - ComputeAndCompareR2(&builder, expected, - {a_data.get(), b_data.get()}, - xla::ErrorSpec(1e-2, 1e-2)); -} - -} // namespace -} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/lib/util.cc b/tensorflow/compiler/tf2xla/lib/util.cc index 804671fbc75b0a5a6e04b204822b6f084013cd8b..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: @@ -113,36 +119,6 @@ xla::XlaOp IntegerLiteral(xla::XlaBuilder* builder, xla::PrimitiveType type, return xla::ConstantLiteral(builder, literal); } -xla::XlaOp SliceInMinorDims(xla::XlaOp x, absl::Span start, - absl::Span end) { - xla::XlaBuilder* builder = x.builder(); - return builder->ReportErrorOrReturn([&]() -> xla::StatusOr { - TF_RET_CHECK(start.size() == end.size()); - int64 n_minor_dims = start.size(); - - TF_ASSIGN_OR_RETURN(xla::Shape shape, builder->GetShape(x)); - - const int64 n_dims = xla::ShapeUtil::Rank(shape); - TF_RET_CHECK(n_minor_dims <= n_dims); - auto major_dims = xla::AsInt64Slice(shape.dimensions()) - .subspan( - /*pos=*/0, - /*len=*/n_dims - n_minor_dims); - - // Prepends 0s in the major dim - std::vector padded_start(n_dims, 0); - std::copy(start.begin(), start.end(), - padded_start.begin() + major_dims.size()); - - // Prepends the shape of the major dims. - std::vector padded_end(n_dims); - std::copy(major_dims.begin(), major_dims.end(), padded_end.begin()); - std::copy(end.begin(), end.end(), padded_end.begin() + major_dims.size()); - - std::vector strides(n_dims, 1); - return xla::Slice(x, padded_start, padded_end, strides); - }); -} std::vector ConcatVectors(absl::Span xs, absl::Span ys) { @@ -152,100 +128,4 @@ std::vector ConcatVectors(absl::Span xs, return output; } -xla::XlaOp DynamicSliceInMinorDims(xla::XlaOp x, - absl::Span starts, - absl::Span sizes) { - xla::XlaBuilder* builder = x.builder(); - return builder->ReportErrorOrReturn([&]() -> xla::StatusOr { - TF_ASSIGN_OR_RETURN(xla::Shape shape, builder->GetShape(x)); - const int64 n_dims = xla::ShapeUtil::Rank(shape); - int64 n_minor_dims = starts.size(); - TF_RET_CHECK(n_minor_dims == sizes.size()); - TF_RET_CHECK(n_minor_dims <= n_dims); - auto major_dims = xla::AsInt64Slice(shape.dimensions()) - .subspan( - /*pos=*/0, - /*len=*/n_dims - sizes.size()); - auto padded_starts = PrependZerosInMajorDims(x, starts); - auto padded_sizes = ConcatVectors(major_dims, sizes); - return xla::DynamicSlice(x, padded_starts, padded_sizes); - }); -} - -xla::XlaOp UpdateSlice(xla::XlaOp x, xla::XlaOp update, - absl::Span start) { - xla::XlaBuilder* builder = x.builder(); - return builder->ReportErrorOrReturn([&]() -> xla::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 = xla::ConstantR1(builder, start_as_int32); - TF_ASSIGN_OR_RETURN(xla::Shape shape, builder->GetShape(x)); - const int64 n_dims = xla::ShapeUtil::Rank(shape); - TF_ASSIGN_OR_RETURN(xla::Shape start_constant_shape, - builder->GetShape(start_constant)); - const int64 start_length = - xla::ShapeUtil::GetDimension(start_constant_shape, -1); - TF_RET_CHECK(start_length == n_dims); - return xla::DynamicUpdateSlice(x, update, start_constant); - }); -} - -xla::XlaOp UpdateSliceInMinorDims(xla::XlaOp x, xla::XlaOp update, - absl::Span start) { - xla::XlaBuilder* builder = x.builder(); - return builder->ReportErrorOrReturn([&]() -> xla::StatusOr { - TF_ASSIGN_OR_RETURN(xla::Shape shape, builder->GetShape(x)); - const int64 n_dims = xla::ShapeUtil::Rank(shape); - const int64 n_minor_dims = start.size(); - TF_RET_CHECK(n_minor_dims <= n_dims); - std::vector padded_start(n_dims, 0); - std::copy(start.begin(), start.end(), - padded_start.begin() + (n_dims - n_minor_dims)); - return UpdateSlice(x, update, padded_start); - }); -} - -xla::XlaOp DynamicUpdateSliceInMinorDims(xla::XlaOp x, xla::XlaOp update, - absl::Span starts) { - auto padded_starts = PrependZerosInMajorDims(x, starts); - return xla::DynamicUpdateSlice(x, update, padded_starts); -} - -xla::XlaOp PrependZerosInMajorDims(xla::XlaOp x, - absl::Span starts) { - xla::XlaBuilder* builder = x.builder(); - return builder->ReportErrorOrReturn([&]() -> xla::StatusOr { - TF_ASSIGN_OR_RETURN(xla::Shape shape, builder->GetShape(x)); - const int64 n_dims = xla::ShapeUtil::Rank(shape); - auto zero = xla::Reshape(xla::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] = xla::Reshape(starts[i], {1}); - } - return xla::ConcatInDim(builder, padded_starts, 0); - }); -} - -xla::XlaOp TransposeInMinorDims(xla::XlaOp x) { - xla::XlaBuilder* builder = x.builder(); - return builder->ReportErrorOrReturn([&]() -> xla::StatusOr { - TF_ASSIGN_OR_RETURN(xla::Shape shape, builder->GetShape(x)); - const int64 n_dims = xla::ShapeUtil::Rank(shape); - TF_RET_CHECK(n_dims >= 2); - std::vector permutation(n_dims); - std::iota(permutation.begin(), permutation.end(), 0); - std::swap(permutation[n_dims - 1], permutation[n_dims - 2]); - return xla::Transpose(x, permutation); - }); -} - -xla::XlaOp MaybeConjugate(xla::XlaOp x, bool conjugate) { - xla::XlaBuilder* builder = x.builder(); - return builder->ReportErrorOrReturn([&]() -> xla::StatusOr { - TF_ASSIGN_OR_RETURN(xla::Shape shape, builder->GetShape(x)); - auto perform_conj = shape.element_type() == xla::C64 && conjugate; - return perform_conj ? xla::Conj(x) : x; - }); -} - } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/lib/util.h b/tensorflow/compiler/tf2xla/lib/util.h index 80e9e5b002d49581209e608b98606e02709c5876..aec8061cb4322b8d315b6cdc80c7fff1e0cb4cb1 100644 --- a/tensorflow/compiler/tf2xla/lib/util.h +++ b/tensorflow/compiler/tf2xla/lib/util.h @@ -38,44 +38,10 @@ xla::XlaOp PrependZerosInMajorDims(xla::XlaOp x, xla::XlaOp IntegerLiteral(xla::XlaBuilder* builder, xla::PrimitiveType type, int64 value); -// Builds a vector of zeros of length rank(x) with the last values being -// those in `starts`. -xla::XlaOp PrependZerosInMajorDims(xla::XlaOp x, - absl::Span starts); - -// Performs a slice in the minor dimensions of a Tensor. -xla::XlaOp SliceInMinorDims(xla::XlaOp x, absl::Span start, - absl::Span end); - // Returns the concatenation of `xs` and `ys`. std::vector ConcatVectors(absl::Span xs, absl::Span ys); -// Performs a dynamic slice in the minor dimensions of a Tensor. -xla::XlaOp DynamicSliceInMinorDims(xla::XlaOp x, - absl::Span starts, - absl::Span sizes); - -// Updates a slice of 'x', i.e., -// x[start[0], ..., start[n]] = update -xla::XlaOp UpdateSlice(xla::XlaOp x, xla::XlaOp update, - absl::Span start); - -// Updates a slice of 'x', where 'start' contains a list of minor dimensions: -// x[..., start[0], ..., start[n]] = update -xla::XlaOp UpdateSliceInMinorDims(xla::XlaOp x, xla::XlaOp update, - absl::Span start); - -xla::XlaOp DynamicUpdateSliceInMinorDims(xla::XlaOp x, xla::XlaOp update, - absl::Span starts); - -// Transposes a stack of matrices `x` by swapping the last two dimensions. -xla::XlaOp TransposeInMinorDims(xla::XlaOp x); - -// Applies a complex conjugation operation if `a` is complex and `conjugate_a` -// is true, otherwise returns its argument. -xla::XlaOp MaybeConjugate(xla::XlaOp x, bool conjugate); - } // namespace tensorflow #endif // TENSORFLOW_COMPILER_TF2XLA_LIB_UTIL_H_ diff --git a/tensorflow/compiler/tf2xla/lib/util_test.cc b/tensorflow/compiler/tf2xla/lib/util_test.cc deleted file mode 100644 index 442fe92c34ca26cb1a854cc90da8dc034bca79bb..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/tf2xla/lib/util_test.cc +++ /dev/null @@ -1,136 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/compiler/tf2xla/lib/util.h" - -#include -#include -#include - -#include "tensorflow/compiler/tf2xla/lib/batch_dot.h" -#include "tensorflow/compiler/xla/array2d.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/types.h" -#include "tensorflow/core/lib/core/status_test_util.h" - -namespace tensorflow { -namespace { - -using UtilTest = xla::ClientLibraryTestBase; -using UtilLeftLookingTest = xla::ClientLibraryTestBase; - -xla::Array2D BValsRight() { - return {{1, 2, 3, 4}, {5, 6, 7, 8}, {9, 10, 11, 12}}; -} - -xla::Array2D BValsLeft() { - return {{1, 2, 3}, {4, 5, 6}, {7, 8, 9}, {10, 11, 12}}; -} - -xla::Array2D AValsFull() { - return {{2, 0, 1, 2}, {3, 6, 0, 1}, {4, 7, 9, 0}, {5, 8, 10, 11}}; -} - -xla::Array3D BatchedAValsFull() { - return {{ - {2, 0, 1, 2}, - {3, 6, 0, 1}, - {4, 7, 9, 0}, - {5, 8, 10, 11}, - }, - { - {16, 24, 8, 12}, - {24, 61, 82, 48}, - {8, 82, 456, 106}, - {12, 48, 106, 62}, - }}; -} - -XLA_TEST_F(UtilTest, Simple2dLookup) { - xla::XlaBuilder builder(TestName()); - - xla::XlaOp a, x, y; - auto a_data = CreateR2Parameter(BValsRight(), 0, "a", &builder, &a); - auto x_data = CreateR0Parameter(2, 1, "x", &builder, &x); - auto y_data = CreateR0Parameter(1, 2, "y", &builder, &y); - DynamicSliceInMinorDims(a, {x, y}, {1, 1}); - - ComputeAndCompareR2(&builder, {{10}}, - {a_data.get(), x_data.get(), y_data.get()}, - xla::ErrorSpec(1e-2, 1e-2)); -} - -XLA_TEST_F(UtilTest, Simple3dLookup) { - xla::XlaBuilder builder(TestName()); - - xla::XlaOp a, index; - auto a_data = - CreateR3Parameter(BatchedAValsFull(), 0, "a", &builder, &a); - auto index_data = CreateR0Parameter(1, 1, "index", &builder, &index); - - DynamicSliceInMinorDims(a, {index, xla::ConstantR0(&builder, 0)}, - {1, 4}); - - ComputeAndCompareR3(&builder, {{{3, 6, 0, 1}}, {{24, 61, 82, 48}}}, - {a_data.get(), index_data.get()}); -} - -XLA_TEST_F(UtilTest, SimpleSliceUpdate) { - xla::XlaBuilder builder(TestName()); - - xla::XlaOp a, b, x, y; - auto a_data = CreateR2Parameter(AValsFull(), 0, "a", &builder, &a); - auto b_data = CreateR2Parameter({{9, 1, -10}}, 1, "b", &builder, &b); - auto x_data = CreateR0Parameter(2, 2, "x", &builder, &x); - auto y_data = CreateR0Parameter(1, 3, "y", &builder, &y); - - DynamicUpdateSliceInMinorDims(a, b, {x, y}); - - xla::Array2D expected( - {{{2, 0, 1, 2}, {3, 6, 0, 1}, {4, 9, 1, -10}, {5, 8, 10, 11}}}); - - ComputeAndCompareR2( - &builder, expected, - {a_data.get(), b_data.get(), x_data.get(), y_data.get()}); -} - -XLA_TEST_F(UtilTest, RowBatchDot) { - xla::XlaBuilder builder(TestName()); - - int n = 4; - - xla::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, xla::ConstantR0(&builder, 0)}, {1, n}); - BatchDot(l_index, row, /*transpose_x=*/false, /*transpose_y=*/true); - - ComputeAndCompareR3(&builder, {{{33}}, {{292}}}, - {a_data.get(), row_data.get(), index_data.get()}); -} - -} // namespace -} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/lib/while_loop.cc b/tensorflow/compiler/tf2xla/lib/while_loop.cc deleted file mode 100644 index 594ab1dfd0700f47501712183f6efe62d17e15e7..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/tf2xla/lib/while_loop.cc +++ /dev/null @@ -1,127 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/compiler/tf2xla/lib/while_loop.h" -#include "tensorflow/compiler/tf2xla/lib/util.h" -#include "tensorflow/compiler/xla/client/xla_builder.h" -#include "tensorflow/compiler/xla/shape_util.h" -#include "tensorflow/compiler/xla/status_macros.h" - -namespace tensorflow { - -xla::StatusOr> XlaWhileLoop( - const LoopConditionFunction& condition_function, - const LoopBodyFunction& body_function, - absl::Span initial_values, absl::string_view name, - xla::XlaBuilder* builder) { - int arity = initial_values.size(); - std::vector var_shapes; - var_shapes.reserve(arity); - for (const xla::XlaOp& input : initial_values) { - TF_ASSIGN_OR_RETURN(auto shape, builder->GetShape(input)); - var_shapes.push_back(std::move(shape)); - } - xla::Shape tuple_shape = xla::ShapeUtil::MakeTupleShape(var_shapes); - - // Unpacks a tuple into its component parts. - auto unpack_tuple = [](xla::XlaOp tuple, int arity, - xla::XlaBuilder* builder) { - std::vector elements(arity); - for (int i = 0; i < arity; ++i) { - elements[i] = xla::GetTupleElement(tuple, i); - } - return elements; - }; - - // Build the condition. - std::unique_ptr cond_builder = - builder->CreateSubBuilder(absl::StrCat(name, "_condition")); - { - auto parameter = - xla::Parameter(cond_builder.get(), 0, tuple_shape, "parameter"); - - TF_RETURN_IF_ERROR( - condition_function(unpack_tuple(parameter, arity, cond_builder.get()), - cond_builder.get()) - .status()); - } - TF_ASSIGN_OR_RETURN(auto cond, cond_builder->Build()); - - // Build the body. - std::unique_ptr body_builder = - builder->CreateSubBuilder(absl::StrCat(name, "_body")); - { - auto parameter = - xla::Parameter(body_builder.get(), 0, tuple_shape, "parameter"); - - TF_ASSIGN_OR_RETURN( - auto result, - body_function(unpack_tuple(parameter, arity, body_builder.get()), - body_builder.get())); - - TF_RET_CHECK(result.size() == initial_values.size()); - xla::Tuple(body_builder.get(), result); - } - TF_ASSIGN_OR_RETURN(auto body, body_builder->Build()); - - auto outputs = xla::While(cond, body, xla::Tuple(builder, initial_values)); - - return unpack_tuple(outputs, arity, builder); -} - -xla::StatusOr> XlaForEachIndex( - int64 num_iterations, xla::PrimitiveType num_iterations_type, - const ForEachIndexBodyFunction& body_function, - absl::Span initial_values, absl::string_view name, - xla::XlaBuilder* builder) { - auto while_cond_fn = - [&](absl::Span values, - xla::XlaBuilder* cond_builder) -> xla::StatusOr { - return xla::Lt(values[0], IntegerLiteral(cond_builder, num_iterations_type, - num_iterations)); - }; - auto while_body_fn = [&](absl::Span values, - xla::XlaBuilder* body_builder) - -> xla::StatusOr> { - xla::XlaOp iteration = values[0]; - - std::vector updated_values; - updated_values.reserve(values.size()); - updated_values.push_back(xla::Add( - iteration, - xla::ConstantLiteral(body_builder, - xla::LiteralUtil::One(num_iterations_type)))); - - values.remove_prefix(1); - TF_ASSIGN_OR_RETURN(std::vector body_outputs, - body_function(iteration, values, body_builder)); - updated_values.insert(updated_values.end(), body_outputs.begin(), - body_outputs.end()); - return updated_values; - }; - - std::vector values; - values.reserve(initial_values.size() + 1); - values.push_back(xla::ConstantLiteral( - builder, xla::LiteralUtil::Zero(num_iterations_type))); - values.insert(values.end(), initial_values.begin(), initial_values.end()); - - TF_ASSIGN_OR_RETURN(values, XlaWhileLoop(while_cond_fn, while_body_fn, values, - name, builder)); - values.erase(values.begin(), values.begin() + 1); - return values; -} - -} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/lib/while_loop.h b/tensorflow/compiler/tf2xla/lib/while_loop.h deleted file mode 100644 index f2134bb4495a12b8342961d96f70e7737f816c7d..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/tf2xla/lib/while_loop.h +++ /dev/null @@ -1,73 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#ifndef TENSORFLOW_COMPILER_TF2XLA_LIB_WHILE_LOOP_H_ -#define TENSORFLOW_COMPILER_TF2XLA_LIB_WHILE_LOOP_H_ - -#include -#include - -#include "absl/strings/string_view.h" -#include "absl/types/span.h" -#include "tensorflow/compiler/xla/client/xla_builder.h" -#include "tensorflow/compiler/xla/client/xla_computation.h" -#include "tensorflow/compiler/xla/statusor.h" - -namespace tensorflow { - -// Function that builds a loop condition. Takes as input a sequence of input -// values, and returns a boolean value representing if the condition succeeds. -typedef std::function(absl::Span, - xla::XlaBuilder*)> - LoopConditionFunction; - -// Function that builds a loop body. Takes as input a sequence of input values -// and returns a sequence of output values. -typedef std::function>( - absl::Span, xla::XlaBuilder*)> - LoopBodyFunction; - -// Helper function for building an XLA while loop, where the values carried by -// the loop are a tuple of values, e.g., (a, b, c): -// while( -// condition: (a, b, c) -> bool, -// body: (a, b, c) -> (a, b, c) -// init: (a, b, c) -// ) -// 'name' is a descriptive name for the loop. -xla::StatusOr> XlaWhileLoop( - const LoopConditionFunction& condition_function, - const LoopBodyFunction& body_function, - absl::Span initial_values, absl::string_view name, - xla::XlaBuilder* builder); - -// Builds an XLA loop that repeats a computation `num_iterations` times. -// -// The body function (ForEachIndexBodyFunction) takes as input a pair of -// (current iteration number, loop-carried values), and returns an updated -// vector of the loop-carried values. -typedef std::function>( - xla::XlaOp, absl::Span, xla::XlaBuilder*)> - ForEachIndexBodyFunction; - -xla::StatusOr> XlaForEachIndex( - int64 num_iterations, xla::PrimitiveType num_iterations_type, - const ForEachIndexBodyFunction& body_function, - absl::Span initial_values, absl::string_view name, - xla::XlaBuilder* builder); - -} // namespace tensorflow - -#endif // TENSORFLOW_COMPILER_TF2XLA_LIB_WHILE_LOOP_H_ diff --git a/tensorflow/compiler/tf2xla/literal_util.cc b/tensorflow/compiler/tf2xla/literal_util.cc index 20103ec3ae00b57723e05326dbbb1b0f6e1a671a..749a7c3054a65d6ec9f9dc13f6f4a713ac9d3d5a 100644 --- a/tensorflow/compiler/tf2xla/literal_util.cc +++ b/tensorflow/compiler/tf2xla/literal_util.cc @@ -32,6 +32,12 @@ Status HostTensorToBorrowingLiteral(const Tensor& host_tensor, return Status::OK(); } +xla::StatusOr HostTensorToLiteral(const Tensor& host_tensor) { + xla::BorrowingLiteral literal; + TF_RETURN_IF_ERROR(HostTensorToBorrowingLiteral(host_tensor, &literal)); + return literal.Clone(); +} + Status HostTensorToMutableBorrowingLiteral( Tensor* host_tensor, xla::MutableBorrowingLiteral* literal) { xla::Shape xla_shape; @@ -71,7 +77,7 @@ Status HostTensorsToBorrowingLiteralTuple(absl::Span host_tensors, Status CopyLiteralToHostTensor(const xla::LiteralSlice& literal, Tensor* host_tensor) { - TF_RET_CHECK(xla::ShapeUtil::IsArray(literal.shape()) && + TF_RET_CHECK(literal.shape().IsArray() && xla::ShapeUtil::ElementsIn(literal.shape()) == host_tensor->NumElements()); xla::PrimitiveType primitive_type; diff --git a/tensorflow/compiler/tf2xla/literal_util.h b/tensorflow/compiler/tf2xla/literal_util.h index 1db7470ee2a839099454b772d4833492e033bc92..a153dddee6127ff9c0858220f2d8a735ab3f0e19 100644 --- a/tensorflow/compiler/tf2xla/literal_util.h +++ b/tensorflow/compiler/tf2xla/literal_util.h @@ -30,6 +30,11 @@ namespace tensorflow { // 'host_tensor'. Status HostTensorToBorrowingLiteral(const Tensor& host_tensor, xla::BorrowingLiteral* literal); + +// Returns a Literal with the contents of 'host_tensor', backed by its own +// storage (i.e., not reusing 'host_tensor's buffers.) +xla::StatusOr HostTensorToLiteral(const Tensor& host_tensor); + // Returns a MutableBorrowingLiteral that utilizes the same underlying buffer // owned by 'host_tensor', but is mutable via the xla::Literal methods. Status HostTensorToMutableBorrowingLiteral( 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 bd2c0a5ee88869ba60701c0a7ace05857452eed9..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 @@ -409,5 +413,29 @@ body: A function that takes a list of tensors and returns another list of tensors. Both lists have the same types as specified by T. )doc"); +REGISTER_OP("XlaDequantize") + .Input("input: uint32") + .Output("output: bfloat16") + .Attr("min_range: float") + .Attr("max_range: float") + .Attr("mode: string") + .Attr("transpose_output: bool") + .SetIsStateful() + .SetShapeFn(shape_inference::UnknownShape) + .Doc(R"doc( +Takes the packed uint32 input and unpacks the input to uint8 to do +Dequantization on deivce. + +input: Input tensors whose types is uint32, shape is [d0, ..., dn]. +output: Output tensors whose types is bloat16. If transpose_output is true, + output shape is [dn * 4, dn-1, ..., d1, d0]. If transpose_output + is false, output shape is [d0,..., dn * 4]. +min_range: The minimum scalar value possibly produced for the input. +max_range: The maximum scalar value possibly produced for the input. +mode: String to determine the dequantize mode in {"MIN_COMBINED", "MIN_FIRST", "SCALED"}. +transpose_output: Boolean to determine if output is transposed. transpose_output + is faster when input is large and rank of input is higher than 1. +)doc"); + } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/python/BUILD b/tensorflow/compiler/tf2xla/python/BUILD index 69ca39436013ec5cf09ba502a1540d5df322e213..9abdb04d7736e8ff5225688af4759a522d3e7fc7 100644 --- a/tensorflow/compiler/tf2xla/python/BUILD +++ b/tensorflow/compiler/tf2xla/python/BUILD @@ -1,9 +1,13 @@ licenses(["notice"]) # Apache 2.0 +package_group( + name = "friends", + includes = ["//tensorflow:internal"], +) + package( default_visibility = [ - "//learning/tfx:__subpackages__", - "//tensorflow:internal", + ":friends", ], ) @@ -11,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", @@ -23,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/python/xla.py b/tensorflow/compiler/tf2xla/python/xla.py index 5e86b5d8ec0a2690f004bc67decea09185d9cbb6..345193c936a885e5a9e468979c4b73b5b0c9e5c2 100644 --- a/tensorflow/compiler/tf2xla/python/xla.py +++ b/tensorflow/compiler/tf2xla/python/xla.py @@ -250,7 +250,7 @@ def conv(lhs, rhs_dilation: dilation to apply between kernel elements dimension_numbers: a `ConvolutionDimensionNumbers` proto. feature_group_count: number of feature groups for grouped convolution. - precision_config: a `PrecisionConfigProto` proto. + precision_config: a `xla.PrecisionConfig` proto. name: an optional name for the operator Returns: @@ -386,3 +386,4 @@ def slice(x, start_dims, limit_dims, strides): sort = gen_xla_ops.xla_sort key_value_sort = gen_xla_ops.xla_key_value_sort while_loop = gen_xla_ops.xla_while +dequantize = gen_xla_ops.xla_dequantize diff --git a/tensorflow/compiler/tf2xla/resource_operation_table.cc b/tensorflow/compiler/tf2xla/resource_operation_table.cc index 72b240996fb4d9dcb5f5dfd919da618cbae08c16..c20d6a5fd1f3bd7dad30cb3359d13ed4609a2250 100644 --- a/tensorflow/compiler/tf2xla/resource_operation_table.cc +++ b/tensorflow/compiler/tf2xla/resource_operation_table.cc @@ -65,6 +65,7 @@ CreateResourceOpInfoMap() { add("ResourceApplyFtrlV2" , kReadWrite, kVariable); add("ResourceApplyGradientDescent" , kReadWrite, kVariable); add("ResourceApplyMomentum" , kReadWrite, kVariable); + add("ResourceApplyKerasMomentum" , kReadWrite, kVariable); add("ResourceApplyPowerSign" , kReadWrite, kVariable); add("ResourceApplyProximalAdagrad" , kReadWrite, kVariable); add("ResourceApplyProximalGradientDescent" , kReadWrite, kVariable); @@ -76,6 +77,7 @@ CreateResourceOpInfoMap() { add("ResourceScatterMin" , kReadWrite, kVariable); add("ResourceScatterMul" , kReadWrite, kVariable); add("ResourceScatterNdAdd" , kReadWrite, kVariable); + add("ResourceScatterNdSub" , kReadWrite, kVariable); add("ResourceScatterNdUpdate" , kReadWrite, kVariable); add("ResourceScatterSub" , kReadWrite, kVariable); add("ResourceScatterUpdate" , kReadWrite, kVariable); diff --git a/tensorflow/compiler/tf2xla/shape_util.cc b/tensorflow/compiler/tf2xla/shape_util.cc index b589512dcdfa32050281120aba6a5ae89a980c2f..8997b2f5c68da480e9d4cb1f7ff8776690363392 100644 --- a/tensorflow/compiler/tf2xla/shape_util.cc +++ b/tensorflow/compiler/tf2xla/shape_util.cc @@ -18,21 +18,81 @@ limitations under the License. #include #include "tensorflow/compiler/tf2xla/type_util.h" +#include "tensorflow/compiler/xla/layout_util.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/core/lib/core/status.h" namespace tensorflow { +namespace { + +Status PopulateInfeedLayoutVector(const xla::Shape& shape, + std::vector* layouts) { + if (shape.IsTuple()) { + int64 tuple_elements = xla::ShapeUtil::TupleElementCount(shape); + for (int64 i = 0; i < tuple_elements; ++i) { + const xla::Shape& subshape = + xla::ShapeUtil::GetTupleElementShape(shape, i); + TF_RETURN_IF_ERROR(PopulateInfeedLayoutVector(subshape, layouts)); + } + } else if (xla::LayoutUtil::HasLayout(shape)) { + for (auto dim : xla::LayoutUtil::MinorToMajor(shape)) { + layouts->push_back(dim); + } + } else { + layouts->insert(layouts->end(), shape.rank(), -1); + } + 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. Status XLAShapeToTensorShape(const xla::Shape& shape, TensorShape* tensor_shape) { - if (xla::ShapeUtil::IsTuple(shape)) { + if (shape.IsTuple()) { return errors::InvalidArgument("XLA shape ", xla::ShapeUtil::HumanString(shape), " cannot be converted to a TensorShape"); } *tensor_shape = TensorShape(); - for (int i = 0; i < xla::ShapeUtil::Rank(shape); ++i) { + for (int i = 0; i < shape.rank(); ++i) { tensor_shape->AddDim(shape.dimensions(i)); } return Status::OK(); @@ -61,4 +121,64 @@ xla::Shape TensorShapeToXLAShape(xla::PrimitiveType type, return xla::ShapeUtil::MakeShapeWithLayout(type, dimensions, layout); } +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 f7e34a5b40c2f9244c029ed325a76322b8cf54dd..e775c4462c3dc15cf4b8d9e8d8e7d9a61e024cd0 100644 --- a/tensorflow/compiler/tf2xla/shape_util.h +++ b/tensorflow/compiler/tf2xla/shape_util.h @@ -18,6 +18,10 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_TF2XLA_SHAPE_UTIL_H_ #define TENSORFLOW_COMPILER_TF2XLA_SHAPE_UTIL_H_ +#include + +#include "tensorflow/compiler/xla/shape.h" +#include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/framework/tensor_shape.h" #include "tensorflow/core/framework/types.pb.h" @@ -40,6 +44,25 @@ Status TensorShapeToXLAShape(DataType dtype, const TensorShape& tensor_shape, 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 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. +// 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 #endif // TENSORFLOW_COMPILER_TF2XLA_SHAPE_UTIL_H_ diff --git a/tensorflow/compiler/tf2xla/side_effect_util.cc b/tensorflow/compiler/tf2xla/side_effect_util.cc index b233e6b2c28e1968bb74901fc684e808ae45ab60..412f31adbb7df52b2d6933be054cc6d40947dc44 100644 --- a/tensorflow/compiler/tf2xla/side_effect_util.cc +++ b/tensorflow/compiler/tf2xla/side_effect_util.cc @@ -24,6 +24,51 @@ const char kXlaTokenInputNodesAttrName[] = "_xla_token_input_nodes"; 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 f22ddb2f58e1fa5c10ca0fdb956d9136942388b7..75e1f253fb08ae61b0336a8783b7449c69197dd1 100644 --- a/tensorflow/compiler/tf2xla/side_effect_util.h +++ b/tensorflow/compiler/tf2xla/side_effect_util.h @@ -35,6 +35,13 @@ extern const char kXlaTokenInputNodesAttrName[]; // node has side-effect dependency on current graph's token input. 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..cf48576ec2746fb29779633275eac4c638b91e45 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); } diff --git a/tensorflow/compiler/tf2xla/tf2xla_test.cc b/tensorflow/compiler/tf2xla/tf2xla_test.cc index ab26d939ccba75ce58609ffd71c7ccadbe90cfa8..24afe595b18b823818bd8fe65bc599af8bce040a 100644 --- a/tensorflow/compiler/tf2xla/tf2xla_test.cc +++ b/tensorflow/compiler/tf2xla/tf2xla_test.cc @@ -91,7 +91,7 @@ TEST(ConvertGraphDefToXla, Sum) { client->ExecuteAndTransfer(computation, {x_global.get(), y_global.get()}); TF_EXPECT_OK(result_or.status()); xla::Literal result = std::move(result_or.ValueOrDie()); - EXPECT_EQ("(s32[]) (\n42\n)", result.ToString()); + EXPECT_EQ("(\ns32[] 42\n)", result.ToString()); config.mutable_feed(0)->mutable_id()->set_output_index( 123); /* invalid output_index */ diff --git a/tensorflow/compiler/tf2xla/tf2xla_util.cc b/tensorflow/compiler/tf2xla/tf2xla_util.cc index 0394b6b533ff971786dd18c76123cbf4bb29c43e..18d87727c500619bf386be7d8c7085724f44aba3 100644 --- a/tensorflow/compiler/tf2xla/tf2xla_util.cc +++ b/tensorflow/compiler/tf2xla/tf2xla_util.cc @@ -30,6 +30,7 @@ limitations under the License. #include "tensorflow/core/framework/graph_to_functiondef.h" #include "tensorflow/core/framework/node_def.pb.h" #include "tensorflow/core/framework/node_def_builder.h" +#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/versions.pb.h" @@ -75,6 +76,222 @@ Status CheckFeedFetchNameConflicts(const string& kind, return Status::OK(); } +// For graph `g`, copy all function call nodes' FunctionDef from `lookup_fld` to +// `fld`. This is to ensure that `fld` can instantiate FunctionDef of graph `g`. +Status CopyAssociatedFunctions(Graph* g, + const FunctionLibraryDefinition* lookup_fld, + FunctionLibraryDefinition* fld) { + for (Node* n : g->op_nodes()) { + for (const auto& associated_function : + GetAssociatedFunctions(*n, lookup_fld)) { + switch (associated_function.type()) { + case AssociatedFunctionInfo::kFunctionCallNode: { + const FunctionDef* fdef = + lookup_fld->Find(associated_function.func_name()); + if (!fdef) { + return errors::Internal( + "Cannot find function ", associated_function.func_name(), + " for function call node ", n->DebugString()); + } + TF_RETURN_IF_ERROR(fld->AddFunctionDef(*fdef)); + break; + } + case AssociatedFunctionInfo::kSymbolicGradient: + case AssociatedFunctionInfo::kFunctionAttr: + break; + } + } + } + return Status::OK(); +} + +// For graph `g`, replaces _Arg nodes whose "index" attribute is in +// `const_input_index_to_node` with Const nodes. +Status ReplaceArgUsageWithConstNode( + Graph* g, + const std::unordered_map& const_input_index_to_node) { + // Collect all _Arg nodes. + std::unordered_map arg_nodes; + for (Node* n : g->op_nodes()) { + if (n->type_string() == FunctionLibraryDefinition::kArgOp) { + int index; + TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), "index", &index)); + arg_nodes[index] = n; + } + } + + for (const auto& iter : const_input_index_to_node) { + int arg_index = iter.first; + Node* const_node = g->CopyNode(iter.second); + Node* arg_node = arg_nodes[arg_index]; + + // Collect all usages of the _Arg node. + struct OutEdgeInfo { + int dst_node_id, dst_input; + }; + std::vector usages; + for (const Edge* e : arg_node->out_edges()) { + if (e->IsControlEdge()) { + continue; + } + usages.push_back({e->dst()->id(), e->dst_input()}); + } + + for (int i = 0; i < usages.size(); i++) { + // Make a copy of `usage_node`, and change its input to const node. + Node* usage_node = g->FindNodeId(usages[i].dst_node_id); + NodeDef replace_def = usage_node->def(); + *replace_def.mutable_input(usages[i].dst_input) = const_node->name(); + TF_ASSIGN_OR_RETURN(Node * replace_node, + ReplaceNode(g, usage_node, replace_def)); + const Edge* usage_edge; + TF_RETURN_IF_ERROR( + replace_node->input_edge(usages[i].dst_input, &usage_edge)); + g->RemoveEdge(usage_edge); + g->AddEdge(const_node, 0, replace_node, usages[i].dst_input); + + // Later entries in `usages` might have `usage_node` as dst node, but + // `usage_node` is removed. Replace such entries with `replace_node`. + for (int j = i + 1; j < usages.size(); j++) { + if (usages[j].dst_node_id == usages[i].dst_node_id) { + usages[j].dst_node_id = replace_node->id(); + } + } + } + } + return Status::OK(); +} + +// For a node's function attr (e.g. then/else branch for "If" nodes), rewrites +// the function to replace _Arg nodes in `const_input_index_to_node` with Const +// inputs. +Status PropagateConstIntoFuncAttr( + Node* n, const string& attr_name, + const std::unordered_map& const_input_index_to_node, + const FunctionLibraryDefinition* lookup_fld, + FunctionLibraryDefinition* fld) { + // Instantiate the function. + NameAttrList func_attr; + TF_RETURN_IF_ERROR(GetNodeAttr(n->def(), attr_name, &func_attr)); + const FunctionDef* fdef = lookup_fld->Find(func_attr.name()); + if (!fdef) { + return errors::Internal("Cannot find function ", func_attr.name(), + " for node ", n->name()); + } + FunctionBody* fbody; + TF_RETURN_IF_ERROR(FunctionDefToBodyHelper( + *fdef, AttrSlice(&func_attr.attr()), lookup_fld, + [lookup_fld](const string& op, const OpDef** sig) { + return lookup_fld->LookUpOpDef(op, sig); + }, + &fbody)); + std::unique_ptr fbody_deleter(fbody); + + // Rewrite _Arg usages with Const node. + Graph* func_graph = fbody->graph; + TF_RETURN_IF_ERROR( + ReplaceArgUsageWithConstNode(func_graph, const_input_index_to_node)); + + // Save rewritten function. + FunctionDef replace_fdef; + string new_func_name = + fld->UniqueFunctionName(absl::StrCat(func_attr.name(), "_const_")); + TF_RETURN_IF_ERROR( + GraphToFunctionDef(*func_graph, new_func_name, &replace_fdef)); + TF_RETURN_IF_ERROR(fld->AddFunctionDef(replace_fdef)); + + // Change the node to use rewritten function. + func_attr.set_name(new_func_name); + n->ClearAttr(attr_name); + n->AddAttr(attr_name, func_attr); + + // Copy associated functions. + TF_RETURN_IF_ERROR(CopyAssociatedFunctions(func_graph, lookup_fld, fld)); + + return Status::OK(); +} + +// For an "If" node in graph `g`, if it has Const node inputs, rewrite its +// then/else branch function to replace _Arg nodes with those Const inputs. +Status PropagateConstIntoIfNode(Graph* g, Node* if_node, + const FunctionLibraryDefinition* lookup_fld, + FunctionLibraryDefinition* fld) { + // Notice that first input for If node is predicate; other inputs are function + // inputs. + std::unordered_map const_input_index_to_node; + for (int i = 1; i < if_node->num_inputs(); i++) { + const Node* input_node; + TF_RETURN_IF_ERROR(if_node->input_node(i, &input_node)); + if (input_node->type_string() == "Const") { + const_input_index_to_node[i - 1] = input_node; + } + } + if (const_input_index_to_node.empty()) { + return Status::OK(); + } + + // Rewrite "then_branch" and "else_branch" function, replace usage of those + // _Arg nodes with corresponding const node. + for (const auto& attr_name : + std::vector{"then_branch", "else_branch"}) { + TF_RETURN_IF_ERROR(PropagateConstIntoFuncAttr( + if_node, attr_name, const_input_index_to_node, lookup_fld, fld)); + } + + return Status::OK(); +} + +// For a "While" node in graph `g`, if it has Const node inputs, rewrite its +// cond/body function to replace _Arg nodes with those Const inputs. +Status PropagateConstIntoWhileNode(Graph* g, Node* while_node, + const FunctionLibraryDefinition* lookup_fld, + FunctionLibraryDefinition* fld) { + // For "While" node, we should only replace _Arg nodes which are loop + // invariants. For such _Arg nodes, the return value's input will come + // directly from the corresponding arg. + std::unordered_map const_input_index_to_node; + NameAttrList body_attr; + TF_RETURN_IF_ERROR(GetNodeAttr(while_node->def(), "body", &body_attr)); + const FunctionDef* body_func = lookup_fld->Find(body_attr.name()); + if (!body_func) { + return errors::Internal("Cannot find body function ", body_attr.name(), + " for While node ", while_node->name()); + } + for (int i = 0; i < while_node->num_inputs(); i++) { + const Node* input_node; + TF_RETURN_IF_ERROR(while_node->input_node(i, &input_node)); + if (input_node->type_string() != "Const") { + continue; + } + + // Check if i-th retval's input comes from i-th arg directly. + 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()) { + return errors::Internal("Cannot find input for output arg ", + output_arg.name(), " in function ", + body_attr.name()); + } + const OpDef_ArgDef& input_arg = body_func->signature().input_arg(i); + if (output_arg_input->second != input_arg.name()) { + continue; + } + + const_input_index_to_node[i] = input_node; + } + if (const_input_index_to_node.empty()) { + return Status::OK(); + } + + // Rewrite "cond" and "body" function, replace usage of those _Arg nodes with + // corresponding const node. + for (const auto& attr_name : std::vector{"cond", "body"}) { + TF_RETURN_IF_ERROR(PropagateConstIntoFuncAttr( + while_node, attr_name, const_input_index_to_node, lookup_fld, fld)); + } + return Status::OK(); +} + } // namespace const char kXlaOutsideCompilationAttrName[] = "_xla_outside_compilation"; @@ -147,6 +364,7 @@ Status AddPlaceholdersForFeeds( GraphDef gd; *gd.mutable_versions() = graph_def->versions(); *gd.add_node() = *existing; + MergeDebugInfo(NodeDebugInfo(*existing), gd.mutable_node(0)); TF_RETURN_IF_ERROR( AddDefaultAttrsToGraphDef(&gd, *op_registry, 0 /*node_offset*/)); @@ -173,6 +391,7 @@ Status AddPlaceholdersForFeeds( // in this code. for (auto it = placeholder_info.begin(); it != placeholder_info.end(); ++it) { const PlaceholderInfo& info = it->second; + // TODO(shikharagarwal): Add original node information. NodeDef* d = graph_def->add_node(); d->set_name(info.placeholder_name); d->set_op("PlaceholderV2"); @@ -340,6 +559,12 @@ bool HasAssociatedFunction(const NodeDef& node_def, return true; } + if (node_def.op() == "XlaHostCompute") { + // XlaHostCompute has "shape_inference_graph" func attr, but that's not + // related to graph execution. + return false; + } + for (const auto& iter : node_def.attr()) { if (iter.second.has_func()) { return true; @@ -361,6 +586,9 @@ std::vector GetAssociatedFunctions( // This is a SymbolicGradient op. AttrValueMap attrs(node.attrs().begin(), node.attrs().end()); results.emplace_back(AssociatedFunctionInfo::SymbolicGradient(op, attrs)); + } else if (node.type_string() == "XlaHostCompute") { + // XlaHostCompute has "shape_inference_graph" func attr, but that's not + // related to graph execution. } else { // Collect all function attrs for the node. for (auto& iter : node.attrs()) { @@ -382,7 +610,9 @@ Status RewriteAssociatedFunction( switch (associated_function.type()) { case AssociatedFunctionInfo::kFunctionCallNode: { // Change this node to call the new function. - NodeDefBuilder builder(node->name(), rewritten_function_name, fld); + NodeDebugInfo debug_info(*node); + NodeDefBuilder builder(node->name(), rewritten_function_name, fld, + &debug_info); for (auto attr : node->attrs()) { builder.Attr(attr.first, attr.second); } @@ -520,4 +750,17 @@ xla::StatusOr BuildIdentityNode( return id_node; } +Status PropagateConstIntoFunctionalNodes( + Graph* g, const FunctionLibraryDefinition* lookup_fld, + FunctionLibraryDefinition* fld) { + for (Node* n : g->op_nodes()) { + if (n->type_string() == "If") { + TF_RETURN_IF_ERROR(PropagateConstIntoIfNode(g, n, lookup_fld, fld)); + } else if (n->type_string() == "While") { + TF_RETURN_IF_ERROR(PropagateConstIntoWhileNode(g, n, lookup_fld, fld)); + } + } + return Status::OK(); +} + } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/tf2xla_util.h b/tensorflow/compiler/tf2xla/tf2xla_util.h index 1232ed8c676ff3f32cce6ce81a79536cafd920a4..cf3aa2f847c5ada8897110c7735b207f388f88d4 100644 --- a/tensorflow/compiler/tf2xla/tf2xla_util.h +++ b/tensorflow/compiler/tf2xla/tf2xla_util.h @@ -183,6 +183,20 @@ xla::StatusOr BuildIdentityNode(Graph* graph, const string& node_name, DataType dtype, const Node* input, absl::optional requested_device); +// For "If"/"While" nodes, if some of their inputs are Const nodes, rewrite +// body functions to use the Const nodes instead of original _Arg nodes. +// +// For example, say we have the following computation: +// shape = constant_op.constant([1]) +// return tf.cond(pred, lambda: tf.ones(shape), lambda: tf.zeros(shape)) +// If we do not rewrite then/else function, they will use _Arg node as shape +// input for tf.ones/tf.zeros. But XLA requires that shape input to be compile +// time constant, so XLA compilation will fail. This rewriting process will +// change the shape input to Const node. +Status PropagateConstIntoFunctionalNodes( + Graph* g, const FunctionLibraryDefinition* lookup_fld, + FunctionLibraryDefinition* fld); + } // namespace tensorflow #endif // TENSORFLOW_COMPILER_TF2XLA_TF2XLA_UTIL_H_ 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_compilation_device.cc b/tensorflow/compiler/tf2xla/xla_compilation_device.cc index cb7843850c352eee2e55baf52a0c4445dc861d7b..ddb284966eeb97cc7c9d3ed77fb313e567975e59 100644 --- a/tensorflow/compiler/tf2xla/xla_compilation_device.cc +++ b/tensorflow/compiler/tf2xla/xla_compilation_device.cc @@ -124,13 +124,4 @@ Status XlaCompilationDevice::MakeTensorFromProto( "XLACompilationDevice::MakeTensorFromProto should not be called"); } -XlaExpression::XlaExpression() = default; - -void XlaExpression::set_handle(const xla::XlaOp& h) { handle_ = h; } - -void XlaExpression::set_constant_value(Tensor value) { - has_constant_value_ = true; - constant_value_ = std::move(value); -} - } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/xla_compilation_device.h b/tensorflow/compiler/tf2xla/xla_compilation_device.h index a6e78825334fec748be5fee80669649df699d2fb..de6a3356e05d8ab45c269d7c6c653853d2c63a79 100644 --- a/tensorflow/compiler/tf2xla/xla_compilation_device.h +++ b/tensorflow/compiler/tf2xla/xla_compilation_device.h @@ -18,9 +18,6 @@ limitations under the License. #include -#include "tensorflow/compiler/tf2xla/xla_resource.h" -#include "tensorflow/compiler/xla/client/xla_builder.h" -#include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/common_runtime/local_device.h" #include "tensorflow/core/framework/device_base.h" #include "tensorflow/core/framework/tensor.h" @@ -38,8 +35,8 @@ class XlaCompilationAllocator; // This is a 'dummy' TensorFlow device that is only used to execute a // subgraph of XLA compilation Ops to construct a compiled version // of the subgraph's computation. It has a 'dummy' allocator that -// backs each Tensor with metadata indicating the computation the -// Tensor represents. +// backs each Tensor with an XlaExpression. The shape of the Tensor +// matches the shape of XlaExpression. // // We deliberately don't register a device factory because we *never* // want placement to put Ops on a compilation device. The device is created @@ -67,40 +64,6 @@ class XlaCompilationDevice : public LocalDevice { std::unique_ptr allocator_; }; -// A XlaExpression wraps an XLA computation. Each Tensor on an -// XlaCompilationDevice contains an XlaExpression, and the shape of the Tensor -// matches the shape of the subcomputation in the XlaOp. Each -// expression is either a constant, or a function of previously-compiled -// expressions. -class XlaExpression { - public: - XlaExpression(); - - // handle() stores the XLA handle of the computation that the - // expression represents. - void set_handle(const xla::XlaOp& h); - const xla::XlaOp& handle() const { return handle_; } - - void set_constant_value(Tensor value); - bool has_constant_value() const { return has_constant_value_; } - const Tensor& constant_value() const { return constant_value_; } - - void set_resource(XlaResource* resource) { resource_ = resource; } - XlaResource* resource() const { return resource_; } - - private: - // The XLA handle of the expression's computation. - xla::XlaOp handle_; - - // If this expression is a constant with a known value, 'constant_value' is a - // host-memory Tensor containing the value. Used to avoid invoking XLA for - // expressions that are trivially constant. - bool has_constant_value_ = false; - Tensor constant_value_; - - XlaResource* resource_ = nullptr; // Not owned. -}; - } // namespace tensorflow #endif // TENSORFLOW_COMPILER_TF2XLA_XLA_COMPILATION_DEVICE_H_ diff --git a/tensorflow/compiler/tf2xla/xla_compiled_cpu_function.h b/tensorflow/compiler/tf2xla/xla_compiled_cpu_function.h index 425e769346ffcbc548495d93cb7adc779f860110..de2e485a47c18ae8e58a06aba408dbb61a30d00a 100644 --- a/tensorflow/compiler/tf2xla/xla_compiled_cpu_function.h +++ b/tensorflow/compiler/tf2xla/xla_compiled_cpu_function.h @@ -26,7 +26,7 @@ limitations under the License. // Forward-declare, rather than include, to reduce code size for users that // never use this functionality. namespace xla { -class ProgramShape; +class ProgramShapeProto; class HloProfilePrinterData; } @@ -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::ProgramShape* 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_; @@ -122,7 +85,7 @@ class XlaCompiledCpuFunction { const char** result_names_ = nullptr; // [Optional] Arg and result shapes. - const xla::ProgramShape* program_shape_ = nullptr; + const xla::ProgramShapeProto* program_shape_ = nullptr; // [Optional] Profile printer data. Null if profiling is disabled. const xla::HloProfilePrinterData* hlo_profile_printer_data_ = nullptr; @@ -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(); @@ -206,8 +170,14 @@ class XlaCompiledCpuFunction { // // Aliasing of argument and result buffers is not allowed, and results in // undefined behavior. - void set_arg_data(size_t index, void* data) { - buffer_table_[arg_index_table_[index]] = data; + void set_arg_data(size_t index, const void* data) { + // The const_cast is safe because the generated code does not write to arg + // buffers. + // + // buffer_table_ contains pointers to buffers that _will_ be written to by + // generated code so it would be misleading to make buffer_table_ a `const + // void**`. + buffer_table_[arg_index_table_[index]] = const_cast(data); } // ------------------------------ @@ -264,7 +234,7 @@ class XlaCompiledCpuFunction { // Returns the shape of the args and results. May return nullptr if the // program shape isn't available. - const xla::ProgramShape* ProgramShape() const { return program_shape_; } + const xla::ProgramShapeProto* ProgramShape() const { return program_shape_; } bool hlo_profiling_enabled() const { return hlo_profile_printer_data_ != nullptr; @@ -274,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_; @@ -287,11 +327,6 @@ class XlaCompiledCpuFunction { // Argument i needs to be placed in buffer_table_[arg_index_to_temp_index_[i]] // for XLA generated code to be able to find it. - // - // For now we need to keep around the args_ array because there is code that - // depends on args() returning a void**. However, in the future we may remove - // args_ in favor of using buffer_table_ as the sole storage for the - // arguments. const int32* const arg_index_table_; // The number of incoming arguments. @@ -310,8 +345,12 @@ class XlaCompiledCpuFunction { // Optional metadata. const char** arg_names_ = nullptr; const char** result_names_ = nullptr; - const xla::ProgramShape* program_shape_ = 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 b2c57e88803e0661a9a514f844dff97ff9edf2ea..1f9cfcdd246f36bd7e0325bca34c7480d4ce2843 100644 --- a/tensorflow/compiler/tf2xla/xla_compiler.cc +++ b/tensorflow/compiler/tf2xla/xla_compiler.cc @@ -31,15 +31,19 @@ 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" #include "tensorflow/core/common_runtime/graph_optimizer.h" #include "tensorflow/core/framework/attr_value_util.h" +#include "tensorflow/core/framework/function.h" #include "tensorflow/core/framework/node_def_util.h" +#include "tensorflow/core/framework/types.h" #include "tensorflow/core/graph/algorithm.h" #include "tensorflow/core/graph/graph_constructor.h" #include "tensorflow/core/graph/node_builder.h" +#include "tensorflow/core/lib/gtl/cleanup.h" #include "tensorflow/core/lib/hash/hash.h" #include "tensorflow/core/platform/logging.h" @@ -48,7 +52,7 @@ namespace { // Checks that arguments `args` match types `types`. Status CheckSignature(const DataTypeVector& types, - const std::vector& args) { + absl::Span args) { if (args.size() != types.size()) { return errors::Internal("Compilation arguments have ", args.size(), " elements while function has ", types.size()); @@ -63,19 +67,290 @@ Status CheckSignature(const DataTypeVector& types, return Status::OK(); } +// Uses the _Arg and _Retval nodes in the graph to determine a core assignment +// for each argument and return value. +xla::StatusOr, std::map>> +ComputeArgAndRetvalCores(const Graph& graph) { + auto get_sharding_for_node = [](const Node* n) -> xla::StatusOr { + TF_ASSIGN_OR_RETURN( + auto sharding, + ParseShardingFromDevice(*n, std::numeric_limits::max())); + if (sharding.has_value()) { + TF_RET_CHECK(sharding.value().type() == + xla::OpSharding::Type::OpSharding_Type_MAXIMAL); + return sharding.value().tile_assignment_devices(0); + } else { + return -1; + } + }; + std::map arg_cores; + std::map retval_cores; + for (const Node* n : graph.nodes()) { + if (n->type_string() == FunctionLibraryDefinition::kArgOp) { + TF_ASSIGN_OR_RETURN(int core, get_sharding_for_node(n)); + if (core < 0) continue; + int index; + TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), "index", &index)); + TF_RET_CHECK(index >= 0) << "Negative _Arg index"; + arg_cores[index] = core; + } else if (n->type_string() == FunctionLibraryDefinition::kRetOp) { + TF_ASSIGN_OR_RETURN(int core, get_sharding_for_node(n)); + if (core < 0) continue; + int index; + TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), "index", &index)); + TF_RET_CHECK(index >= 0) << "Negative _Retval index"; + TF_ASSIGN_OR_RETURN(retval_cores[index], get_sharding_for_node(n)); + retval_cores[index] = core; + } + } + return std::make_pair(std::move(arg_cores), std::move(retval_cores)); +} + +Status ExecuteGraph(XlaContext* xla_context, std::unique_ptr graph, + XlaCompilationDevice* device, FunctionLibraryRuntime* flib, + int64 step_id) { + // Resource cleanup is a bit messy. XlaContext is a ref-countd resource; the + // resource manager takes ownership via Create, and unrefs via Cleanup. We + // explicitly add a reference to ensure the refcount at entry is maintained at + // all exit points; Create and Cleanup are always called in this function. + // + // The Executor requires us to use ScopedStepContainer. We wrap it in a + // unique_ptr so we can capture the cleanup status in the end. + xla_context->Ref(); + Status status; + auto step_container = absl::make_unique( + step_id, [&status, device](const string& name) { + status = device->resource_manager()->Cleanup(name); + }); + TF_RETURN_IF_ERROR(device->resource_manager()->Create( + step_container->name(), XlaContext::kXlaContextResourceName, + xla_context)); + + GraphCompiler graph_compiler(device, graph.get(), flib, step_container.get()); + TF_RETURN_IF_ERROR(graph_compiler.Compile()); + // Explicitly clean up the step container, to capture the cleanup status. + step_container.reset(); + return Status::OK(); +} + +// Builds the XLA computation. +// - `args` is the list of input arguments +// - `retvals` is the list of retvals produced by _Retval operators, in index +// order. +// - `args_core` and `retval_cores` are mapping from arg/return indices to core +// assignments. +// - If `return_updated_values_for_all_resources` is true, all resources will be +// included in `resource_updates`, regardless of whether their value changed. +// - Sets `*num_nonconst_outputs` to the number of outputs of the `computation`. +// - Sets `*resource_updates` to a description of resources whose values are +// written by the computation; the variable writes are the last +// - `resource_updates.size()` return values from the computation. Each entry in +// `resource_updates` is a ResourceUpdate, whose `index` is the index of a +// resource variable argument to the computation to be updated, and `type` is +// the type of the final output. +Status BuildComputation( + const std::vector& args, + const std::vector& retvals, + const std::map& arg_cores, const std::map& retval_cores, + const std::vector>& resources, + std::unique_ptr token_output, + const XlaCompiler::ShapeRepresentationFn& shape_representation_fn, + bool return_updated_values_for_all_resources, bool always_return_tuple, + xla::XlaBuilder* builder, xla::XlaComputation* computation, + int* num_computation_outputs, int* num_nonconst_outputs, + std::vector* outputs, + std::vector* resource_updates, + xla::Shape* output_shape) { + // Attach a common operator name as metadata. This has no semantic effect — it + // merely makes the HLO graph more readable when visualized via TensorBoard, + // since TensorBoard forms groups out of operators with similar names. + xla::OpMetadata retval_metadata; + retval_metadata.set_op_name("XLA_Retvals"); + builder->SetOpMetadata(retval_metadata); + auto cleanup = gtl::MakeCleanup([builder]() { builder->ClearOpMetadata(); }); + + // Builds a no-op XLA computation. We need to set the sharding of outputs, but + // cannot change the sharding of the existing output op. To do this, we build + // a new identity op to which shardings can be applied. + auto identity_op = [builder](xla::XlaOp op) { + return xla::GetTupleElement(xla::Tuple(builder, {op}), 0); + }; + + std::vector elems; + elems.reserve(retvals.size()); + + // Keeps track of which retvals have layout to update. The first element is + // the output index, second element is the new layout. + std::vector> retval_to_update_layout; + for (int i = 0; i < retvals.size(); ++i) { + XlaCompiler::OutputDescription& output = (*outputs)[i]; + const XlaExpression& retval = retvals[i]; + output.type = retval.dtype(); + switch (retval.kind()) { + case XlaExpression::Kind::kConstant: + output.is_constant = true; + output.constant_value = retval.constant_value(); + 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()); + xla::XlaOp value = retval.handle(); + auto it = retval_cores.find(i); + xla::XlaScopedShardingAssignment assign_sharding( + builder, it == retval_cores.end() + ? absl::optional() + : xla::sharding_builder::AssignDevice(it->second)); + if (shape_representation_fn) { + // If there is a shape representation function, reshape the output + // tensor to the shape given by the representation shape function. + TF_ASSIGN_OR_RETURN(xla::Shape shape, shape_representation_fn( + output.shape, output.type)); + value = xla::Reshape(value, xla::AsInt64Slice(shape.dimensions())); + retval_to_update_layout.emplace_back(elems.size(), shape.layout()); + } else if (it != retval_cores.end()) { + // Apply the sharding to the output, if there is a core assignment. + value = identity_op(value); + } + + elems.push_back(value); + break; + } + + case XlaExpression::Kind::kResource: + output.is_constant = false; + output.input_index = retval.resource()->arg_num(); + output.shape = retval.resource()->shape(); + break; + + case XlaExpression::Kind::kInvalid: + return errors::InvalidArgument( + "Invalid expression returned by computation. " + "This probably means a return value was not set."); + } + } + *num_nonconst_outputs = elems.size(); + + // Add return values for resources whose values have changed. + std::vector arg_resources; + arg_resources.reserve(resources.size()); + for (const auto& resource : resources) { + if (resource->arg_num() >= 0) { + arg_resources.push_back(resource.get()); + } + } + std::sort(arg_resources.begin(), arg_resources.end(), + [](const XlaResource* a, const XlaResource* b) { + return a->arg_num() < b->arg_num(); + }); + + for (const XlaResource* resource : arg_resources) { + DCHECK_LT(resource->arg_num(), args.size()); + const XlaCompiler::Argument& arg = args[resource->arg_num()]; + auto it = arg_cores.find(resource->arg_num()); + const int core = it == arg_cores.end() ? -1 : it->second; + bool modified = !resource->value().IsIdenticalTo(resource->initial_value()); + // TensorArray gradients were modified if their values changed or there are + // any newly created gradients. + for (const auto& grad : resource->tensor_array_gradients()) { + modified = + modified || + !grad.second->value().IsIdenticalTo(grad.second->initial_value()) || + arg.tensor_array_gradients.count(grad.first) == 0; + } + if (return_updated_values_for_all_resources || modified) { + resource_updates->emplace_back(); + XlaCompiler::ResourceUpdate& update = resource_updates->back(); + update.input_index = resource->arg_num(); + update.type = resource->type(); + update.shape = resource->shape(); + update.modified = modified; + for (const auto& grad : resource->tensor_array_gradients()) { + update.tensor_array_gradients_accessed.insert(grad.first); + } + + // Request that the value be returned on a specific core. + xla::XlaScopedShardingAssignment assign_sharding( + builder, core == -1 ? absl::optional() + : xla::sharding_builder::AssignDevice(core)); + + xla::XlaOp handle; + TF_RETURN_IF_ERROR(resource->Pack(&handle, builder)); + + // Ensures the correct sharding is applied to the output. + handle = identity_op(handle); + + elems.push_back(handle); + } + } + + // If we have token output, append it as the last one. + if (token_output) { + elems.push_back(*token_output); + } + + *num_computation_outputs = elems.size(); + + // Builds the XLA computation. We *always* form a tuple here to ensure that + // the output value is the last thing added into the XLA computation, even + // if there is only one output value. + auto tuple = xla::Tuple(builder, elems); + if (!always_return_tuple && elems.size() == 1) { + xla::GetTupleElement(tuple, 0); + } + + xla::StatusOr computation_status = builder->Build(); + if (!computation_status.ok()) { + return computation_status.status(); + } + *computation = computation_status.ConsumeValueOrDie(); + + TF_ASSIGN_OR_RETURN(const auto& program_shape, + computation->GetProgramShape()); + *output_shape = program_shape.result(); + // Update the output layout to the layout of retval. + for (auto& update : retval_to_update_layout) { + if (!always_return_tuple && elems.size() == 1) { + *output_shape->mutable_layout() = update.second; + continue; + } + + xla::Shape* output_sub_shape = + xla::ShapeUtil::GetMutableSubshape(output_shape, {update.first}); + *output_sub_shape->mutable_layout() = update.second; + } + return Status::OK(); +} + } // namespace bool XlaCompiler::Argument::operator==( const XlaCompiler::Argument& other) const { - if (std::tie(kind, resource_kind, type, name, initialized, tensor_array_size, + if (std::tie(kind, resource_kind, type, name, initialized, max_array_size, tensor_array_gradients) != std::tie(other.kind, other.resource_kind, other.type, other.name, - other.initialized, other.tensor_array_size, + other.initialized, other.max_array_size, 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; @@ -83,12 +358,62 @@ bool XlaCompiler::Argument::operator==( return constant_value.tensor_data() == other.constant_value.tensor_data(); } +string XlaCompiler::Argument::HumanString() const { + string common; + if (!name.empty()) { + common = absl::StrCat(" name=", name); + } + absl::StrAppend(&common, " type=", DataTypeString(type), + " shape=", ShapeHumanString()); + switch (kind) { + case kInvalid: + return "invalid"; + case kConstant: + return absl::StrCat("kind=constant", common, + " value=", constant_value.DebugString()); + case kResource: { + string output = absl::StrCat("kind=resource", common, " resource_kind=", + XlaResource::KindToString(resource_kind), + " initialized=", initialized); + if (max_array_size >= 0) { + absl::StrAppend(&output, " max_array_size=", max_array_size); + } + if (!tensor_array_gradients.empty()) { + absl::StrAppend(&output, " tensor_array_gradients=", + absl::StrJoin(tensor_array_gradients, ",")); + } + return output; + } + case kParameter: + return absl::StrCat("kind=parameter", common); + case kToken: + return absl::StrCat("token", common); + } +} + +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()), next_step_id_(1), device_(new XlaCompilationDevice(SessionOptions(), options_.device_type)), - device_mgr_({device_}) { + device_mgr_(absl::WrapUnique(device_)) { CHECK(!options_.device_type.type_string().empty()); if (options_.populate_resource_manager) { initialization_status_ = @@ -110,8 +435,13 @@ XlaCompiler::XlaCompiler(XlaCompiler::Options options) // The default shape representation function is the identity. if (!options_.shape_representation_fn) { - options_.shape_representation_fn = [](const TensorShape& shape, - DataType type) { return shape; }; + options_.shape_representation_fn = + [](const TensorShape& shape, + DataType dtype) -> xla::StatusOr { + xla::Shape xla_shape; + TF_RETURN_IF_ERROR(TensorShapeToXLAShape(dtype, shape, &xla_shape)); + return xla_shape; + }; } } @@ -165,21 +495,48 @@ 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; } -Status XlaCompiler::CompileFunction(const XlaCompiler::CompileOptions& options, - const NameAttrList& function, - std::vector args, - XlaCompiler::CompilationResult* result) { +Status XlaCompiler::CompileFunction( + const XlaCompiler::CompileOptions& options, const NameAttrList& function, + absl::Span args, + XlaCompiler::CompilationResult* result) { const string function_id = Canonicalize(function.name(), AttrSlice(&function.attr())); VLOG(1) << "XlaCompiler::CompileFunction " << function_id; - auto it = cache_.find({function_id, args}); + const std::vector arg_vector(args.begin(), args.end()); + auto it = cache_.find({function_id, arg_vector}); if (it != cache_.end()) { *result = it->second; return Status::OK(); @@ -212,14 +569,16 @@ Status XlaCompiler::CompileFunction(const XlaCompiler::CompileOptions& options, // lowest-numbered core that consumes the argument. We choose the // lowest-numbered core so the assignment is deterministic. for (Node* n : graph->nodes()) { - if (absl::string_view(n->type_string()) == "_Arg") { + if (absl::string_view(n->type_string()) == + FunctionLibraryDefinition::kArgOp) { TF_RETURN_IF_ERROR(SetNodeShardingFromNeighbors(n, /*out_edges=*/true)); } } // Do _Retval as a second loop, in case the retval's input is an _Arg (which // may have gotten a device assignment from the first loop). for (Node* n : graph->nodes()) { - if (absl::string_view(n->type_string()) == "_Retval") { + if (absl::string_view(n->type_string()) == + FunctionLibraryDefinition::kRetOp) { TF_RETURN_IF_ERROR(SetNodeShardingFromNeighbors(n, /*out_edges=*/false)); } } @@ -235,7 +594,7 @@ Status XlaCompiler::CompileFunction(const XlaCompiler::CompileOptions& options, CompileGraph(options, function_id, std::move(graph), args, result)); VLOG(1) << "===================================================="; - cache_[{function_id, args}] = *result; + cache_[{function_id, arg_vector}] = *result; return Status::OK(); } @@ -247,34 +606,47 @@ Status XlaCompiler::XLAShapeForArgument(const XlaCompiler::Argument& arg, case XlaCompiler::Argument::kConstant: LOG(FATAL) << "Unreachable case"; case XlaCompiler::Argument::kParameter: { - TensorShape shape; if (is_entry_computation) { - TF_ASSIGN_OR_RETURN( - 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 { - shape = arg.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 TensorShapeToXLAShape(arg.type, shape, xla_shape); + return Status::OK(); } case XlaCompiler::Argument::kResource: { TF_RET_CHECK(arg.initialized); switch (arg.resource_kind) { case XlaResource::kVariable: { - TF_ASSIGN_OR_RETURN( - TensorShape representation_shape, - options_.shape_representation_fn(arg.shape, arg.type)); - return TensorShapeToXLAShape(arg.type, representation_shape, - xla_shape); + 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(); } case XlaResource::kTensorArray: { - if (arg.tensor_array_size < 0) { + if (arg.max_array_size < 0) { return errors::InvalidArgument( - "Negative tensor_array_size in XLAShapeForArgument"); + "Negative max_array_size in XLAShapeForArgument"); } + TF_RET_CHECK(absl::holds_alternative(arg.shape)); TensorShape shape; - shape.AddDim(arg.tensor_array_size); - shape.AppendShape(arg.shape); + shape.AddDim(arg.max_array_size); + shape.AppendShape(absl::get(arg.shape)); TF_RETURN_IF_ERROR(TensorShapeToXLAShape(arg.type, shape, xla_shape)); if (!arg.tensor_array_gradients.empty()) { @@ -285,13 +657,14 @@ Status XlaCompiler::XLAShapeForArgument(const XlaCompiler::Argument& arg, return Status::OK(); } case XlaResource::kStack: { - if (arg.tensor_array_size < 0) { + if (arg.max_array_size < 0) { return errors::InvalidArgument( - "Negative tensor_array_size in XLAShapeForArgument"); + "Negative max_array_size in XLAShapeForArgument"); } + TF_RET_CHECK(absl::holds_alternative(arg.shape)); TensorShape shape; - shape.AddDim(arg.tensor_array_size); - shape.AppendShape(arg.shape); + shape.AddDim(arg.max_array_size); + shape.AppendShape(absl::get(arg.shape)); xla::Shape buffer_shape; TF_RETURN_IF_ERROR( TensorShapeToXLAShape(arg.type, shape, &buffer_shape)); @@ -314,174 +687,22 @@ Status XlaCompiler::XLAShapeForArgument(const XlaCompiler::Argument& arg, } } -namespace { - -Status ExecuteGraph(XlaContext* xla_context, std::unique_ptr graph, - XlaCompilationDevice* device, FunctionLibraryRuntime* flib, - int64 step_id) { - // Resource cleanup is a bit messy. XlaContext is a ref-countd resource; the - // resource manager takes ownership via Create, and unrefs via Cleanup. We - // explicitly add a reference to ensure the refcount at entry is maintained at - // all exit points; Create and Cleanup are always called in this function. - // - // The Executor requires us to use ScopedStepContainer. We wrap it in a - // unique_ptr so we can capture the cleanup status in the end. - xla_context->Ref(); - Status status; - auto step_container = absl::make_unique( - step_id, [&status, device](const string& name) { - status = device->resource_manager()->Cleanup(name); - }); - TF_RETURN_IF_ERROR(device->resource_manager()->Create( - step_container->name(), XlaContext::kXlaContextResourceName, - xla_context)); - - GraphCompiler graph_compiler(device, graph.get(), flib, step_container.get()); - TF_RETURN_IF_ERROR(graph_compiler.Compile()); - // Explicitly clean up the step container, to capture the cleanup status. - step_container.reset(); - return Status::OK(); -} - -// Builds the XLA computation. -// `args` is the list of input arguments, `retvals` is the list of retvals -// produced by _Retval operators, in index order. -// If `return_updated_values_for_all_resources` is true, all resources will be -// included in `resource_updates`, regardless of whether their value changed. -// Sets `*num_nonconst_outputs` to the number of outputs of the `computation`. -// Sets `*resource_updates` to a description of resources whose values are -// written by the computation; the variable writes are the last -// `resource_updates.size()` return values from the computation. Each entry in -// `resource_updates` is a (input_index, type) pair, where `input_index` is the -// index of a resource variable argument to the computation, and `type` is the -// type of the final output. -Status BuildComputation( - const std::vector& args, - const std::vector& arg_cores, - const std::vector& retvals, - const std::vector>& resources, - bool return_updated_values_for_all_resources, bool always_return_tuple, - xla::XlaBuilder* builder, xla::XlaComputation* computation, - int* num_computation_outputs, int* num_nonconst_outputs, - std::vector* outputs, - std::vector* resource_updates) { - std::vector elems; - elems.reserve(retvals.size()); - for (int i = 0; i < retvals.size(); ++i) { - XlaCompiler::OutputDescription& output = (*outputs)[i]; - output.type = retvals[i].type; - output.shape = retvals[i].shape; - const XlaExpression& retval = retvals[i].expression; - if (retval.has_constant_value()) { - output.is_constant = true; - output.constant_value = retval.constant_value(); - } else if (retval.resource() != nullptr) { - output.is_constant = false; - output.input_index = retval.resource()->arg_num(); - } else { - output.is_constant = false; - elems.push_back(retval.handle()); - } - } - *num_nonconst_outputs = elems.size(); - - // Add return values for resources whose values have changed. - std::vector arg_resources; - arg_resources.reserve(resources.size()); - for (const auto& resource : resources) { - if (resource->arg_num() >= 0) { - arg_resources.push_back(resource.get()); - } - } - std::sort(arg_resources.begin(), arg_resources.end(), - [](const XlaResource* a, const XlaResource* b) { - return a->arg_num() < b->arg_num(); - }); - - // Attach a common operator name as metadata. This has no semantic effect — it - // merely makes the HLO graph more readable when visualized via TensorBoard, - // since TensorBoard forms groups out of operators with similar names. - xla::OpMetadata retval_metadata; - retval_metadata.set_op_name("XLA_Retvals"); - builder->SetOpMetadata(retval_metadata); - - for (const XlaResource* resource : arg_resources) { - const XlaCompiler::Argument& arg = args[resource->arg_num()]; - const int core = arg_cores[resource->arg_num()]; - DCHECK_LT(resource->arg_num(), arg_cores.size()); - bool modified = !resource->value().IsIdenticalTo(resource->initial_value()); - // TensorArray gradients were modified if their values changed or there are - // any newly created gradients. - for (const auto& grad : resource->tensor_array_gradients()) { - modified = - modified || - !grad.second->value().IsIdenticalTo(grad.second->initial_value()) || - arg.tensor_array_gradients.count(grad.first) == 0; - } - if (return_updated_values_for_all_resources || modified) { - resource_updates->emplace_back(); - XlaCompiler::ResourceUpdate& update = resource_updates->back(); - update.input_index = resource->arg_num(); - update.type = resource->type(); - update.shape = resource->shape(); - update.modified = modified; - for (const auto& grad : resource->tensor_array_gradients()) { - update.tensor_array_gradients_accessed.insert(grad.first); - } - - // Request that the value be returned on a specific core. - xla::XlaScopedShardingAssignment assign_sharding( - builder, core == -1 ? absl::optional() - : xla::sharding_builder::AssignDevice(core)); - - xla::XlaOp handle; - TF_RETURN_IF_ERROR(resource->Pack(&handle, builder)); - - // Since we can't change the sharding metadata of as this point, - // create a tuple/get-tuple-element combination so that sharding - // assignment will be placed on this value, which will cause the resource - // update to be returned from the same device that provided the resource. - handle = xla::GetTupleElement(xla::Tuple(builder, {handle}), 0); - elems.push_back(handle); - } - } - - *num_computation_outputs = elems.size(); - - // Builds the XLA computation. We *always* form a tuple here to ensure that - // the output value is the last thing added into the XLA computation, even - // if there is only one output value. - auto tuple = xla::Tuple(builder, elems); - if (!always_return_tuple && elems.size() == 1) { - xla::GetTupleElement(tuple, 0); - } - builder->ClearOpMetadata(); - - xla::StatusOr computation_status = builder->Build(); - if (!computation_status.ok()) { - return computation_status.status(); - } - *computation = computation_status.ConsumeValueOrDie(); - return Status::OK(); -} - -} // namespace - // Builds XLA computations for each of the arguments to the computation. // `args` are the arguments to the computation. Status XlaCompiler::BuildArguments( const Graph& graph, const std::vector& args, bool use_tuple_arg, xla::XlaBuilder* builder, XlaContext* context, - std::vector* arg_cores, std::vector* arg_expressions, - std::vector* input_mapping, std::vector* input_shapes, + const std::map& arg_cores, + std::vector* arg_expressions, + std::vector* input_to_args, std::vector* input_shapes, bool is_entry_computation) { arg_expressions->resize(args.size()); - *arg_cores = std::vector(args.size(), -1); // 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(); @@ -489,28 +710,31 @@ Status XlaCompiler::BuildArguments( const XlaCompiler::Argument& arg = args[i]; XlaExpression& arg_expression = (*arg_expressions)[i]; switch (arg.kind) { - case XlaCompiler::Argument::kResource: + 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; - TF_RETURN_IF_ERROR(context->CreateResource( - arg.resource_kind, i, arg.name, arg.type, arg.shape, xla::XlaOp(), - /*tensor_array_size=*/arg.tensor_array_size, - /*tensor_array_gradients=*/arg.tensor_array_gradients, &resource)); - arg_expression.set_resource(resource); + XlaResource* resource = + context->AddResource(absl::make_unique( + 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: - arg_expression.set_constant_value(arg.constant_value); + arg_expression = XlaExpression::Constant(arg.constant_value); break; case XlaCompiler::Argument::kInvalid: return errors::Internal( @@ -518,15 +742,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) { @@ -535,26 +767,6 @@ Status XlaCompiler::BuildArguments( *input_shapes = arg_shapes; } - // Use the _Arg nodes in the graph to resolve core assignments. - for (const Node* n : graph.nodes()) { - if (absl::string_view(n->type_string()) != "_Arg") continue; - int index; - TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), "index", &index)); - TF_RET_CHECK(index >= 0 && index < args.size()) - << "_Arg out of bounds: " << index << " vs " << args.size(); - TF_ASSIGN_OR_RETURN( - auto sharding, - ParseShardingFromDevice(*n, std::numeric_limits::max())); - if (sharding.has_value()) { - TF_RET_CHECK(sharding.value().type() == - xla::OpSharding::Type::OpSharding_Type_MAXIMAL); - const int core = sharding.value().tile_assignment_devices(0); - if ((*arg_cores)[index] == -1 || core < (*arg_cores)[index]) { - (*arg_cores)[index] = core; - } - } - } - // Attach a common operator name as metadata. This has no semantic effect — it // merely makes the HLO graph more readable when visualized via TensorBoard, // since TensorBoard forms groups out of operators with similar names. @@ -563,18 +775,17 @@ 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) { - const int core = (*arg_cores)[parameter]; - const int root_device = 0; + 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() = - core == -1 ? xla::sharding_builder::AssignDevice(root_device) - : xla::sharding_builder::AssignDevice(core); + xla::sharding_builder::AssignDevice(core); } xla::XlaScopedShardingAssignment assign_tuple_sharding(builder, tuple_sharding); @@ -582,22 +793,47 @@ 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) { - const int core = (*arg_cores)[input_mapping->at(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( builder, core == -1 ? absl::optional() : xla::sharding_builder::AssignDevice(core)); arg_handles[i] = xla::GetTupleElement(tuple, i); } } else { - for (std::vector::size_type i = 0; i < input_mapping->size(); ++i) { - const int core = (*arg_cores)[input_mapping->at(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( builder, core == -1 ? absl::optional() : xla::sharding_builder::AssignDevice(core)); 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(); @@ -605,12 +841,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); @@ -626,14 +862,14 @@ Status XlaCompiler::BuildArguments( // TODO(b/76097077): propagate device assignments onto arguments and // return values of functions, and then reshape unconditionally. if (is_entry_computation) { - arg_expression.set_handle( - xla::Reshape(arg_handles[i], arg.shape.dim_sizes())); + arg_expression = XlaExpression::XlaOp( + xla::Reshape(arg_handles[i], arg.DimensionSizes()), arg.type); } else { - arg_expression.set_handle(arg_handles[i]); + arg_expression = XlaExpression::XlaOp(arg_handles[i], arg.type); } break; case XlaCompiler::Argument::kToken: { - arg_expression.set_handle(arg_handles[i]); + arg_expression = XlaExpression::XlaOp(arg_handles[i], arg.type); break; } case XlaCompiler::Argument::kConstant: @@ -647,46 +883,48 @@ Status XlaCompiler::BuildArguments( } Status XlaCompiler::CompileSingleOp( - const XlaCompiler::CompileOptions& options, string const& name, - OpKernelContext* ctx, const std::vector& args, - CompilationResult* result) { + const XlaCompiler::CompileOptions& options, const NodeDef& node_def, + absl::Span args, + absl::Span result_types, CompilationResult* result) { // TODO(b/74182462): We implement this by creating a new dummy Graph including // _Arg nodes, and let CompileGraph walk it. This could be optimized. std::unique_ptr graph(new Graph(OpRegistry::Global())); Status status; // First create the actual node we care about computing. - Node* main_node = graph->AddNode(ctx->op_kernel().def(), &status); + Node* main_node = graph->AddNode(node_def, &status); TF_RETURN_IF_ERROR(status); // Create dummy _Arg nodes. Link these to `node` and also via a control // dependency edge to the _SOURCE node. - for (int64 i = 0; i < ctx->num_inputs(); ++i) { + for (int64 i = 0; i < args.size(); ++i) { Node* node; - string name = absl::StrCat(ctx->op_kernel().name(), "_", i, "_arg"); - Status status = NodeBuilder(name, "_Arg") - .ControlInput(graph->source_node()) - .Attr("T", ctx->input_dtype(i)) - .Attr("index", i) - .Finalize(graph.get(), &node); + string arg_name = absl::StrCat("_arg", i); + Status status = + NodeBuilder(arg_name, FunctionLibraryDefinition::kArgOp) + .ControlInput(graph->source_node()) + .Attr("T", args[i].kind == Argument::kResource ? DT_RESOURCE + : args[i].type) + .Attr("index", i) + .Finalize(graph.get(), &node); TF_RETURN_IF_ERROR(status); graph->AddEdge(node, 0, main_node, i); } // Similarly with return values, create dummy _Retval nodes fed by `node`. - for (int64 i = 0; i < ctx->num_outputs(); ++i) { + for (int64 i = 0; i < result_types.size(); ++i) { Node* node; - string name = absl::StrCat(ctx->op_kernel().name(), "_", i, "_retval"); - Status status = NodeBuilder(name, "_Retval") + string retval_name = absl::StrCat("_retval", i); + Status status = NodeBuilder(retval_name, FunctionLibraryDefinition::kRetOp) .Input(main_node, i) - .Attr("T", ctx->expected_output_dtype(i)) + .Attr("T", result_types[i]) .Attr("index", i) .Finalize(graph.get(), &node); TF_RETURN_IF_ERROR(status); } FixupSourceAndSinkEdges(graph.get()); - return CompileGraph(options, name, std::move(graph), args, result); + return CompileGraph(options, node_def.name(), std::move(graph), args, result); } namespace { @@ -741,15 +979,43 @@ Status ValidateGraph(const Graph* graph, return Status::OK(); } +// Converts the value of any expressions whose values are known at compile-time +// to constants. +Status ResolveConstantExpressionsToConstants( + xla::Client* client, absl::Span expressions) { + for (XlaExpression& expression : expressions) { + if (expression.kind() == XlaExpression::Kind::kXlaOp) { + TF_ASSIGN_OR_RETURN(absl::optional constant, + expression.ResolveConstant(client)); + if (constant.has_value()) { + expression = XlaExpression::Constant(*constant); + } + } + } + return Status::OK(); +} + +void ConvertConstantsToExpressions(xla::XlaBuilder* builder, + absl::Span expressions) { + for (XlaExpression& expression : expressions) { + if (expression.kind() == XlaExpression::Kind::kConstant) { + expression = + XlaExpression::XlaOp(expression.AsXlaOp(builder), expression.dtype()); + } + } +} + } // namespace Status XlaCompiler::CompileGraph(const XlaCompiler::CompileOptions& options, string const& name, std::unique_ptr graph, - const std::vector& args, + absl::Span args, CompilationResult* result) { VLOG(1) << "Executing graph symbolically to populate XlaBuilder."; + TF_RETURN_IF_ERROR(PropagateConstIntoFunctionalNodes( + graph.get(), options_.flib_def, local_flib_def_.get())); if (VLOG_IS_ON(2)) { VLOG(2) << "XlaCompiler::CompileGraph: " << dump_graph::DumpGraphToFile( @@ -766,14 +1032,12 @@ Status XlaCompiler::CompileGraph(const XlaCompiler::CompileOptions& options, options_.device_type, name)); xla::XlaBuilder builder(name); - XlaContext* context = new XlaContext( - this, &builder, options_.allow_cpu_custom_calls, - options.resolve_compile_time_constants, options.is_entry_computation, - &options_.shape_representation_fn); + XlaContext* context = new XlaContext(this, &builder); core::ScopedUnref context_unref(context); - std::vector real_args(args); + std::vector real_args(args.begin(), args.end()); int token_input_index = -1; + std::unique_ptr token_output; if (options.add_token_input_output) { // Add extra token input. token_input_index = real_args.size(); @@ -783,10 +1047,14 @@ Status XlaCompiler::CompileGraph(const XlaCompiler::CompileOptions& options, real_args.push_back(token_arg); } + std::map arg_cores; + std::map retval_cores; + TF_ASSIGN_OR_RETURN(std::tie(arg_cores, retval_cores), + ComputeArgAndRetvalCores(*graph)); + std::vector arg_expressions; - std::vector arg_cores; TF_RETURN_IF_ERROR(BuildArguments( - *graph, real_args, options.use_tuple_arg, &builder, context, &arg_cores, + *graph, real_args, options.use_tuple_arg, &builder, context, arg_cores, &arg_expressions, &result->input_mapping, &result->xla_input_shapes, options.is_entry_computation)); context->set_args(std::move(arg_expressions)); @@ -826,8 +1094,7 @@ Status XlaCompiler::CompileGraph(const XlaCompiler::CompileOptions& options, TF_RETURN_IF_ERROR(token_or.status()); token_inputs.push_back(token_or.ValueOrDie()); } - TF_RETURN_IF_ERROR( - context->AppendTokenRetval(xla::AfterAll(&builder, token_inputs))); + token_output.reset(new xla::XlaOp(xla::AfterAll(&builder, token_inputs))); } TF_RETURN_IF_ERROR(PopNodeTokenMapping()); @@ -835,28 +1102,27 @@ Status XlaCompiler::CompileGraph(const XlaCompiler::CompileOptions& options, int num_computation_outputs; result->computation = std::make_shared(); 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))); + } else { + ConvertConstantsToExpressions(&builder, absl::Span(retvals)); + } TF_RETURN_IF_ERROR(BuildComputation( - real_args, arg_cores, context->retvals(), context->resources(), + real_args, retvals, arg_cores, retval_cores, context->resources(), + std::move(token_output), + options.is_entry_computation ? options_.shape_representation_fn + : ShapeRepresentationFn{}, options.return_updated_values_for_all_resources, options.always_return_tuple, &builder, result->computation.get(), &num_computation_outputs, &num_nonconst_outputs, &result->outputs, - &result->resource_updates)); + &result->resource_updates, &result->xla_output_shape)); VLOG(2) << "Outputs: total: " << context->retvals().size() << " nonconstant: " << num_nonconst_outputs; - - // Compute the XLA output shape, if there is a computation with non-constant - // outputs. - TF_ASSIGN_OR_RETURN(std::unique_ptr computation_shape, - client()->GetComputationShape(*result->computation)); - - result->xla_output_shape.Swap(computation_shape->mutable_result()); VLOG(2) << "XLA output shape: " - << xla::ShapeUtil::HumanString(result->xla_output_shape); - - // Tensorflow expects a major-to-minor order of results. - xla::LayoutUtil::SetToDefaultLayout(&result->xla_output_shape); - + << xla::ShapeUtil::HumanStringWithLayout(result->xla_output_shape); return Status::OK(); } diff --git a/tensorflow/compiler/tf2xla/xla_compiler.h b/tensorflow/compiler/tf2xla/xla_compiler.h index 2cc603a58016a509fafdf6f95423dd6c0864cce3..ad3144b41bdf3fc8b75ab5230e8e128df2962884 100644 --- a/tensorflow/compiler/tf2xla/xla_compiler.h +++ b/tensorflow/compiler/tf2xla/xla_compiler.h @@ -18,10 +18,14 @@ 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" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "tensorflow/compiler/xla/client/local_client.h" +#include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/compiler/xla/client/xla_computation.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/core/common_runtime/device.h" @@ -118,10 +122,11 @@ class XlaCompiler { // The type of the argument. If the argument is a resource, this // is the type of the variable's value, not DT_RESOURCE. - DataType type; + 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 @@ -130,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. @@ -147,14 +152,27 @@ class XlaCompiler { // For a TensorArray or Stack resource, what is the array's declared size? // (Used for lazy initialization.) - int64 tensor_array_size = -1; + int64 max_array_size = -1; // TensorArray resource parameters are passed as (array, gradient array 0, // ..., gradient array k), where the gradient arrays are in the same order // 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 @@ -259,8 +277,7 @@ class XlaCompiler { std::shared_ptr computation; }; - typedef std::function(const TensorShape&, - DataType)> + typedef std::function(const TensorShape&, DataType)> ShapeRepresentationFn; struct Options { // Name of the compilation device to use. It must be set by the caller. @@ -316,22 +333,23 @@ class XlaCompiler { Status CompileFunction(const CompileOptions& options, const NameAttrList& fn_name_attrs, - std::vector args, CompilationResult* result); + absl::Span args, + CompilationResult* result); // Compiles a tensorflow::Graph into an xla::XlaComputation. // Similar to CompileFunction, but takes a Graph as input rather than a // function. Status CompileGraph(const CompileOptions& options, string const& name, std::unique_ptr graph, - const std::vector& args, + absl::Span args, CompilationResult* result); - // Compiles a single Op, given by an OpKernelContext, into an + // Compiles a single Op, given by `node_def`, into an // xla::XlaComputation. Similar to CompileFunction but takes a single Op as // input. - Status CompileSingleOp(const CompileOptions& options, string const& name, - OpKernelContext* ctx, - const std::vector& args, + Status CompileSingleOp(const CompileOptions& options, const NodeDef& node_def, + absl::Span args, + absl::Span result_types, CompilationResult* result); // Returns the shape of the XLA parameter for an argument 'arg'. @@ -411,9 +429,10 @@ class XlaCompiler { Status BuildArguments(const Graph& graph, const std::vector& args, bool use_tuple_arg, xla::XlaBuilder* builder, - XlaContext* context, std::vector* arg_cores, + 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 2cb8b3331cf7a26657739dac4dfb85fd84d501ca..492010f7317d32a8a620147cd2cd9356d4f13fde 100644 --- a/tensorflow/compiler/tf2xla/xla_compiler_test.cc +++ b/tensorflow/compiler/tf2xla/xla_compiler_test.cc @@ -20,7 +20,9 @@ limitations under the License. #include "tensorflow/cc/ops/function_ops.h" #include "tensorflow/cc/ops/resource_variable_ops.h" #include "tensorflow/cc/ops/standard_ops.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_op_kernel.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "tensorflow/compiler/xla/client/client_library.h" @@ -80,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_; } @@ -648,7 +650,7 @@ TEST_F(XlaCompilerTest, CanPassTensorArraysToAndFromComputation) { args[0].initialized = true; args[0].type = DT_INT32; args[0].shape = TensorShape({}); - args[0].tensor_array_size = 2; + args[0].max_array_size = 2; args[0].tensor_array_gradients = {"grad2"}; // Compiles the graph. @@ -707,7 +709,7 @@ TEST_F(XlaCompilerTest, UnwrittenTensorArrayGradientsAreNotComputationOutputs) { args[0].initialized = true; args[0].type = DT_INT32; args[0].shape = TensorShape({}); - args[0].tensor_array_size = 2; + args[0].max_array_size = 2; args[0].tensor_array_gradients = {"grad1"}; // Compiles the graph. @@ -739,7 +741,7 @@ TEST_F(XlaCompilerTest, NewTensorArrayGradientsAreComputationOutputs) { args[0].initialized = true; args[0].type = DT_INT32; args[0].shape = TensorShape({}); - args[0].tensor_array_size = 2; + args[0].max_array_size = 2; args[0].tensor_array_gradients = {"grad1"}; // Compiles the graph. @@ -909,6 +911,82 @@ TEST_F(XlaCompilerTest, Variables) { RunAndCheckVariablesComputation(client_, result); } +TEST_F(XlaCompilerTest, ResultLayoutSingle) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto a = ops::_Arg(scope.WithOpName("A"), DT_INT32, 0); + auto b = ops::_Retval(scope.WithOpName("RET"), a, 0); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + // Builds a description of the arguments. + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = TensorShape({2, 3}); + + auto options = DefaultOptions(); + // Sets the representation function to return a non-default layout. + options.shape_representation_fn = + [](const TensorShape& shape, DataType type) -> xla::StatusOr { + xla::Shape xla_shape; + TF_RETURN_IF_ERROR(TensorShapeToXLAShape(type, shape, &xla_shape)); + *xla_shape.mutable_layout() = xla::LayoutUtil::MakeLayout({0, 1}); + return xla_shape; + }; + + // Compiles the graph. + XlaCompiler compiler(options); + + XlaCompiler::CompilationResult result; + auto compile_options = XlaCompiler::CompileOptions(); + compile_options.always_return_tuple = false; + TF_ASSERT_OK(compiler.CompileGraph(compile_options, "id", std::move(graph), + args, &result)); + EXPECT_TRUE(xla::ShapeUtil::Equal( + result.xla_output_shape, + xla::ShapeUtil::MakeShapeWithLayout(xla::S32, {2, 3}, {0, 1}))); +} + +TEST_F(XlaCompilerTest, ResultLayoutMultiple) { + Scope scope = Scope::NewRootScope().ExitOnError(); + auto a = ops::_Arg(scope.WithOpName("A"), DT_INT32, 0); + auto b = ops::_Retval(scope.WithOpName("RET1"), a, 0); + auto c = ops::_Retval(scope.WithOpName("RET2"), a, 1); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + TF_ASSERT_OK(scope.ToGraph(graph.get())); + + // Builds a description of the arguments. + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kParameter; + args[0].type = DT_INT32; + args[0].shape = TensorShape({2, 3}); + + auto options = DefaultOptions(); + // Sets the representation function to return a non-default layout. + options.shape_representation_fn = + [](const TensorShape& shape, DataType type) -> xla::StatusOr { + xla::Shape xla_shape; + TF_RETURN_IF_ERROR(TensorShapeToXLAShape(type, shape, &xla_shape)); + *xla_shape.mutable_layout() = xla::LayoutUtil::MakeLayout({0, 1}); + return xla_shape; + }; + + // Compiles the graph. + XlaCompiler compiler(options); + + XlaCompiler::CompilationResult result; + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "id", + std::move(graph), args, &result)); + xla::Shape result_shape = + xla::ShapeUtil::MakeShapeWithLayout(xla::S32, {2, 3}, {0, 1}); + + EXPECT_TRUE(xla::ShapeUtil::Equal( + result.xla_output_shape, + xla::ShapeUtil::MakeTupleShape({result_shape, result_shape}))); +} + // Tests a simple graph that reads and writes a variable. TEST_F(XlaCompilerTest, ReturnResourceHandleOnly) { Scope scope = Scope::NewRootScope().ExitOnError(); @@ -1018,9 +1096,11 @@ TEST_F(XlaCompilerTest, VariableRepresentationShapeFunction) { // Compiles the graph. XlaCompiler::Options options = DefaultOptions(); - options.shape_representation_fn = [](const TensorShape& shape, - DataType type) { - return TensorShape({shape.num_elements()}); + options.shape_representation_fn = + [](const TensorShape& shape, DataType type) -> xla::StatusOr { + xla::PrimitiveType ptype; + TF_RETURN_IF_ERROR(DataTypeToPrimitiveType(type, &ptype)); + return xla::ShapeUtil::MakeShape(ptype, {shape.num_elements()}); }; XlaCompiler compiler(options); @@ -1086,9 +1166,11 @@ TEST_F(XlaCompilerTest, ArgRetvalShapeRepresentationFunction) { // Compiles the graph. XlaCompiler::Options options = DefaultOptions(); - options.shape_representation_fn = [](const TensorShape& shape, - DataType type) { - return TensorShape({shape.num_elements()}); + options.shape_representation_fn = + [](const TensorShape& shape, DataType type) -> xla::StatusOr { + xla::PrimitiveType ptype; + TF_RETURN_IF_ERROR(DataTypeToPrimitiveType(type, &ptype)); + return xla::ShapeUtil::MakeShape(ptype, {shape.num_elements()}); }; XlaCompiler compiler(options); @@ -1258,23 +1340,30 @@ TEST_F(XlaCompilerTest, TokenInputAndOutput) { TF_ASSERT_OK(status); EXPECT_TRUE(FixupSourceAndSinkEdges(graph.get())); - const std::vector empty_args; + std::vector args(1); + args[0].kind = XlaCompiler::Argument::kResource; + args[0].resource_kind = XlaResource::kVariable; + args[0].initialized = true; + args[0].type = DT_INT32; + args[0].shape = TensorShape({2, 2}); + { // The case for entry computation: we don't add token input/output. Instead, // we use CreateToken HLO to create the entry token. XlaCompiler::CompileOptions options; options.is_entry_computation = true; options.add_token_input_output = false; + options.return_updated_values_for_all_resources = true; XlaCompiler compiler(DefaultOptions()); std::unique_ptr graph_copy(new Graph(OpRegistry::Global())); CopyGraph(*graph, graph_copy.get()); XlaCompiler::CompilationResult result; TF_ASSERT_OK(compiler.CompileGraph(options, "NoOp", std::move(graph_copy), - empty_args, &result)); - EXPECT_EQ(result.xla_input_shapes.size(), 0); - EXPECT_TRUE(xla::ShapeUtil::IsTuple(result.xla_output_shape)); - EXPECT_EQ(xla::ShapeUtil::TupleElementCount(result.xla_output_shape), 0); + args, &result)); + EXPECT_EQ(result.xla_input_shapes.size(), 1); + EXPECT_TRUE(result.xla_output_shape.IsTuple()); + EXPECT_EQ(xla::ShapeUtil::TupleElementCount(result.xla_output_shape), 1); } { // The case for non-entry computation (e.g. while loop body). We add token @@ -1282,19 +1371,20 @@ TEST_F(XlaCompilerTest, TokenInputAndOutput) { XlaCompiler::CompileOptions options; options.is_entry_computation = false; options.add_token_input_output = true; + options.return_updated_values_for_all_resources = true; XlaCompiler compiler(DefaultOptions()); std::unique_ptr graph_copy(new Graph(OpRegistry::Global())); CopyGraph(*graph, graph_copy.get()); XlaCompiler::CompilationResult result; TF_ASSERT_OK(compiler.CompileGraph(options, "NoOp", std::move(graph_copy), - empty_args, &result)); - EXPECT_EQ(result.xla_input_shapes.size(), 1); - EXPECT_TRUE(xla::ShapeUtil::IsToken(result.xla_input_shapes[0])); - EXPECT_TRUE(xla::ShapeUtil::IsTuple(result.xla_output_shape)); - EXPECT_EQ(xla::ShapeUtil::TupleElementCount(result.xla_output_shape), 1); - EXPECT_TRUE(xla::ShapeUtil::IsToken( - xla::ShapeUtil::GetTupleElementShape(result.xla_output_shape, 0))); + args, &result)); + EXPECT_EQ(result.xla_input_shapes.size(), 2); + EXPECT_TRUE(result.xla_input_shapes[1].IsToken()); + EXPECT_TRUE(result.xla_output_shape.IsTuple()); + EXPECT_EQ(xla::ShapeUtil::TupleElementCount(result.xla_output_shape), 2); + EXPECT_TRUE(xla::ShapeUtil::GetTupleElementShape(result.xla_output_shape, 1) + .IsToken()); } } diff --git a/tensorflow/compiler/tf2xla/xla_context.cc b/tensorflow/compiler/tf2xla/xla_context.cc index 2095a6b8099f48a867ec2c7c7d6e84d8f2426dce..6139bf3cea0790c2697130a993e92be96c81848b 100644 --- a/tensorflow/compiler/tf2xla/xla_context.cc +++ b/tensorflow/compiler/tf2xla/xla_context.cc @@ -54,99 +54,25 @@ const char XlaContext::kXlaContextResourceName[] = "_xla_context"; return *context; } -/* static */ XlaContext& XlaContext::Get(const XlaOpKernelContext* ctx) { - return Get(ctx->op_kernel_context()); -} - void XlaContext::set_args(std::vector args) { args_ = std::move(args); } -XlaContext::XlaContext( - XlaCompiler* compiler, xla::XlaBuilder* builder, - bool allow_cpu_custom_calls, bool resolve_compile_time_constants, - bool is_entry_computation, - const std::function( - const TensorShape&, DataType)>* shape_representation_fn) - : compiler_(compiler), - builder_(builder), - allow_cpu_custom_calls_(allow_cpu_custom_calls), - resolve_compile_time_constants_(resolve_compile_time_constants), - is_entry_computation_(is_entry_computation), - shape_representation_fn_(shape_representation_fn) {} - -string XlaContext::DebugString() { return "TLA JIT context"; } - -// This is called by the Retval Op to associate a computed value -// with a specific return value of the subgraph. -void XlaContext::AddRetval(int retval_index, DataType type, - const TensorShape& shape, const xla::XlaOp& handle) { - VLOG(1) << "Added retval index " << retval_index << " to XLA computation"; - // Add the return value to the list being built up. - if (retvals_.size() <= retval_index) { - retvals_.resize(retval_index + 1); - } - XlaExpression e; - e.set_handle(handle); - retvals_[retval_index] = Retval{type, shape, e}; -} +XlaContext::XlaContext(XlaCompiler* compiler, xla::XlaBuilder* builder) + : compiler_(compiler), builder_(builder) {} -Status XlaContext::AddConstRetval(int retval_index, DataType dtype, - const xla::LiteralSlice& literal) { - VLOG(1) << "Adding retval index " << retval_index - << " with non-data-dependent tensor to XLA computation"; - if (retvals_.size() <= retval_index) { - retvals_.resize(retval_index + 1); - } - Tensor value; - TF_RETURN_IF_ERROR(LiteralToHostTensor(literal, dtype, &value)); - XlaExpression e; - e.set_constant_value(value); - retvals_[retval_index] = Retval{dtype, value.shape(), e}; - return Status::OK(); -} +string XlaContext::DebugString() const { return "XLA JIT context"; } -Status XlaContext::AddResourceRetval(int retval_index, XlaResource* resource) { - VLOG(1) << "Adding retval index " << retval_index << " with resource " - << resource->name() << ":" << resource->shape().DebugString() - << " to XLA computation"; - if (retvals_.size() <= retval_index) { - retvals_.resize(retval_index + 1); +void XlaContext::SetRetval(int index, const XlaExpression& expression) { + if (retvals_.size() <= index) { + retvals_.resize(index + 1); } - XlaExpression e; - e.set_resource(resource); - retvals_[retval_index] = Retval{DT_RESOURCE, resource->shape(), e}; - return Status::OK(); -} - -Status XlaContext::AppendTokenRetval(const xla::XlaOp& token) { - VLOG(1) << "Adding retval index " << retvals_.size() - << " with token to XLA computation"; - XlaExpression e; - e.set_handle(token); - // We use DT_INVALID because there is no TF DataType which corresponds to XLA - // token. XlaCompiler handles this case separately, so putting it here is OK. - retvals_.push_back(Retval{DT_INVALID, TensorShape(), e}); - return Status::OK(); -} - -xla::XlaBuilder* XlaContext::builder() { return builder_; } - -Status XlaContext::CreateResource( - XlaResource::Kind kind, int arg_num, string name, DataType type, - TensorShape shape, const xla::XlaOp& handle, int64 tensor_array_size, - const std::set& tensor_array_gradients, XlaResource** resource) { - resources_.emplace_back( - new XlaResource(kind, arg_num, std::move(name), type, std::move(shape), - handle, tensor_array_size, tensor_array_gradients, - /*tensor_array_multiple_writes_aggregate=*/false)); - *resource = resources_.back().get(); - return Status::OK(); + retvals_[index] = expression; } -xla::StatusOr XlaContext::RepresentationShape( - const TensorShape& shape, DataType type) const { - return (*shape_representation_fn_)(shape, type); +XlaResource* XlaContext::AddResource(std::unique_ptr resource) { + resources_.push_back(std::move(resource)); + return resources_.back().get(); } const xla::XlaComputation* XlaContext::GetOrCreateMax(const DataType type) { diff --git a/tensorflow/compiler/tf2xla/xla_context.h b/tensorflow/compiler/tf2xla/xla_context.h index d7dbdc957f0e7969db5098b815381866cdc71ab6..eb4ad3fe6a14b42a4df2c73c71cb6df1331fd796 100644 --- a/tensorflow/compiler/tf2xla/xla_context.h +++ b/tensorflow/compiler/tf2xla/xla_context.h @@ -20,8 +20,8 @@ limitations under the License. #include -#include "tensorflow/compiler/tf2xla/xla_compilation_device.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" +#include "tensorflow/compiler/tf2xla/xla_expression.h" #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/compiler/xla/client/xla_computation.h" #include "tensorflow/compiler/xla/status_macros.h" @@ -41,76 +41,36 @@ class XlaContext : public ResourceBase { public: // Retrieves the XlaContext of the current compilation. static XlaContext& Get(const OpKernelContext* ctx); - static XlaContext& Get(const XlaOpKernelContext* ctx); // Creates a new XlaContext. See the documentation on the class data fields // for descriptions of the arguments. - XlaContext(XlaCompiler* compiler, xla::XlaBuilder* builder, - bool allow_cpu_custom_calls, bool resolve_compile_time_constants, - bool is_entry_computation, - const std::function( - const TensorShape&, DataType)>* shape_representation_fn); + XlaContext(XlaCompiler* compiler, xla::XlaBuilder* builder); // Virtual method defined by ResourceBase. - string DebugString() override; + string DebugString() const override; XlaCompiler* compiler() const { return compiler_; } // Returns the XlaBuilder that Ops use for compiling new expressions. - xla::XlaBuilder* builder(); - - bool allow_cpu_custom_calls() const { return allow_cpu_custom_calls_; } - - bool resolve_compile_time_constants() const { - return resolve_compile_time_constants_; - } - bool is_entry_computation() const { return is_entry_computation_; } + xla::XlaBuilder* builder() { return builder_; } const std::vector& args() const { return args_; } void set_args(std::vector args); - struct Retval { - DataType type; - TensorShape shape; - // An XlaExpression representing the Retval's value. - XlaExpression expression; - }; - const std::vector& retvals() { return retvals_; } - - // This is called by the Retval Op to associate a computed value - // with a specific return value of the subgraph. - void AddRetval(int retval_index, DataType type, const TensorShape& shape, - const xla::XlaOp& handle); - - // As for Retval, but for return values that are compile-time constants. - Status AddConstRetval(int retval_index, DataType dtype, - const xla::LiteralSlice& literal); - - // As for Retval, but for return values that are resource handles. - Status AddResourceRetval(int retval_index, XlaResource* resource); - - // As for Retval, but for return values that are XLA tokens. - Status AppendTokenRetval(const xla::XlaOp& token); - - // Creates a resource with resource `kind` and initial value `handle`. `name` - // is a descriptive name for use in error messages. See the `XlaResource` - // constructor for a description of the remaining arguments. - // Fails if the resource already exists. - Status CreateResource(XlaResource::Kind kind, int arg_num, string name, - DataType type, TensorShape shape, - const xla::XlaOp& handle, int64 tensor_array_size, - const std::set& tensor_array_gradients, - XlaResource** resource); + const std::vector& retvals() { return retvals_; } + + // Sets a return value. + // Since we do not always know in advance how many return values there are, + // grows the return values vector to size index+1 if it is smaller. + void SetRetval(int index, const XlaExpression& expression); + + // Adds 'resource' to the set of resources owned by the context. + XlaResource* AddResource(std::unique_ptr resource); const std::vector>& resources() { return resources_; } - // Returns the XLA shape to be used to represent a variable of TF `shape` - // and `type`, or of an argument or return value of a top-level computation. - xla::StatusOr RepresentationShape(const TensorShape& shape, - DataType type) const; - // Get an XLA lambda to compute Max. This is cached in the // XlaContext since it may be used by multiple Ops. There is a // separate specialization of the computation for each DataType. @@ -140,36 +100,16 @@ class XlaContext : public ResourceBase { // The XlaBuilder used to construct the subgraph's compiled representation. xla::XlaBuilder* builder_; - // Allow ops to emit CustomCall operations for CPU. - const bool allow_cpu_custom_calls_; - - // If true, constant return values are returned as Tensors instead of - // run-time computation outputs. - const bool resolve_compile_time_constants_; - // Arguments to the Tensorflow graph, indexed by _Arg index. // Includes both compile-time constant arguments and runtime parameters. std::vector args_; // Return values of the Tensorflow graph, indexed by _Retval index. - std::vector retvals_; + std::vector retvals_; // Holds ownership of resources. The resources are not ordered. std::vector> resources_; - // Is this a top-level computation, or an inner computation (e.g., a while - // body)? - const bool is_entry_computation_; - - // A function that describes how the shapes of - // a) argument and return value, for entry computations - // b) variables, for all computations, - // should be represented in XLA. Parameters/return values will be shaped - // according to this function, and reshaped back to/from their declared shapes - // for computations. Must be non-null. - const std::function(const TensorShape&, DataType)>* - shape_representation_fn_; - // Cache of prebuilt computations indexed by their type. using ComputationMap = std::map; diff --git a/tensorflow/compiler/tf2xla/xla_expression.cc b/tensorflow/compiler/tf2xla/xla_expression.cc new file mode 100644 index 0000000000000000000000000000000000000000..3d228c92adcbe3d093a4fe70d157e57ab3e80c80 --- /dev/null +++ b/tensorflow/compiler/tf2xla/xla_expression.cc @@ -0,0 +1,162 @@ +/* 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/compiler/tf2xla/xla_expression.h" + +#include "tensorflow/compiler/tf2xla/literal_util.h" +#include "tensorflow/compiler/tf2xla/shape_util.h" +#include "tensorflow/core/framework/types.pb.h" +#include "tensorflow/core/lib/core/errors.h" + +namespace tensorflow { + +XlaExpression::XlaExpression() = default; + +XlaExpression XlaExpression::Invalid() { + XlaExpression e; + e.kind_ = Kind::kInvalid; + return e; +} + +XlaExpression XlaExpression::Constant(Tensor value) { + XlaExpression e; + e.kind_ = Kind::kConstant; + e.dtype_ = value.dtype(); + e.constant_value_ = value; + return e; +} + +XlaExpression XlaExpression::XlaOp(xla::XlaOp value, DataType dtype) { + XlaExpression e; + e.kind_ = Kind::kXlaOp; + e.dtype_ = dtype; + e.handle_ = value; + 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; + e.dtype_ = DT_RESOURCE; + e.resource_ = resource; + return e; +} + +string XlaExpression::HumanString() const { + switch (kind_) { + case Kind::kInvalid: + return "invalid"; + case Kind::kConstant: + return "constant"; + case Kind::kXlaOp: + return "xla_op"; + case Kind::kResource: + return "resource"; + case Kind::kTensorList: + return "tensor_list"; + } +} + +xla::XlaOp XlaExpression::AsXlaOp(xla::XlaBuilder* builder) const { + return builder->ReportErrorOrReturn([&]() -> xla::StatusOr { + switch (kind_) { + case Kind::kConstant: { + xla::BorrowingLiteral literal; + TF_RETURN_IF_ERROR( + 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( + "Mismatched builders in XlaExpression::AsXlaOp"); + } + return handle_; + default: + return errors::InvalidArgument("AsXlaOp called on XlaExpression: ", + HumanString()); + } + }); +} + +xla::StatusOr> XlaExpression::ResolveConstant( + xla::Client* client) const { + switch (kind()) { + case Kind::kConstant: + 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()); + } + + TF_ASSIGN_OR_RETURN(bool is_constant, + handle().builder()->IsConstant(handle())); + if (!is_constant) return {absl::nullopt}; + + TF_ASSIGN_OR_RETURN(xla::XlaComputation constant_graph, + handle().builder()->BuildConstantSubGraph(handle())); + + TF_ASSIGN_OR_RETURN(TensorShape shape, GetShape()); + + // The XLA layout is specified minor to major, and TensorFlow uses a major to + // minor order. + std::vector layout_indices(shape.dims()); + std::iota(layout_indices.rbegin(), layout_indices.rend(), 0); + xla::Layout layout = xla::LayoutUtil::MakeLayout(layout_indices); + TF_ASSIGN_OR_RETURN(xla::Literal literal, + client->ComputeConstant(constant_graph, &layout)); + Tensor tensor; + TF_RETURN_IF_ERROR(LiteralToHostTensor(literal, dtype(), &tensor)); + return {tensor}; +} + +xla::StatusOr XlaExpression::GetShape() const { + switch (kind_) { + case Kind::kConstant: + return constant_value().shape(); + case Kind::kXlaOp: { + TF_ASSIGN_OR_RETURN(xla::Shape xla_shape, + handle().builder()->GetShape(handle())); + TensorShape shape; + TF_RETURN_IF_ERROR(XLAShapeToTensorShape(xla_shape, &shape)); + return shape; + } + case Kind::kTensorList: + return TensorShape({}); + case Kind::kResource: + return TensorShape({}); + case Kind::kInvalid: + return errors::InvalidArgument( + "GetShape() called on invalid XlaExpression"); + } +} + +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/xla_expression.h b/tensorflow/compiler/tf2xla/xla_expression.h new file mode 100644 index 0000000000000000000000000000000000000000..ac0232d8924cf2c9e35ad3f0772a3a2adc18af87 --- /dev/null +++ b/tensorflow/compiler/tf2xla/xla_expression.h @@ -0,0 +1,125 @@ +/* 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_TF2XLA_XLA_EXPRESSION_H_ +#define TENSORFLOW_COMPILER_TF2XLA_XLA_EXPRESSION_H_ + +#include "absl/types/optional.h" +#include "tensorflow/compiler/tf2xla/xla_resource.h" +#include "tensorflow/compiler/xla/client/client.h" +#include "tensorflow/compiler/xla/client/xla_builder.h" +#include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/core/framework/tensor.h" +#include "tensorflow/core/lib/core/status.h" + +namespace tensorflow { + +// A XlaExpression represents a symbolic TensorFlow value in a TF->XLA +// compilation. +// An expression is one of: +// * 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 { + kInvalid, + kConstant, + kXlaOp, + kResource, + kTensorList, + }; + + XlaExpression(); + XlaExpression(const XlaExpression&) = default; + XlaExpression& operator=(const XlaExpression&) = default; + + // Builds an invalid expression. (Same as the default constructor, but makes + // the intent clearer.) + static XlaExpression Invalid(); + + // Builds a constant XLA expression. + static XlaExpression Constant(Tensor value); + + // Builds a XlaOp expression. Since the mapping from TF data types to XLA + // types is not 1-1, the TF type must also be provided; in general it cannot + // 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); + + Kind kind() const { return kind_; } + + DataType dtype() const { return dtype_; } + + // handle() returns the XlaOp that backs a kXlaOp expression. + const xla::XlaOp& handle() const { return handle_; } + + const Tensor& constant_value() const { return constant_value_; } + + XlaResource* resource() const { return resource_; } + + // Returns a human-readable summary of the expression. + string HumanString() const; + + // Returns the value of a kConstant or kXlaOp as an xla::XlaOp. Returns + // an erroneous XlaOp if the expression is not a constant or an expression. + xla::XlaOp AsXlaOp(xla::XlaBuilder* builder) const; + + // If a kXlaOp or kConstant expression can be resolved to a compile-time + // constant, returns the value as a host-memory Tensor. Returns an empty + // optional if it cannot be resolved. Returns an error if passed a resource + // expression. + xla::StatusOr> ResolveConstant( + xla::Client* client) const; + + // Returns the shape of the tensor. + // The shape of a resource is the shape of a resource handle (i.e., a scalar), + // not the shape of the resource's value. + xla::StatusOr GetShape() const; + + private: + Kind kind_ = Kind::kInvalid; + + DataType dtype_ = DT_INVALID; + + // 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. + Tensor constant_value_; + + // The resource, if kind_ == kResource. Not owned. + XlaResource* resource_ = nullptr; +}; + +} // namespace tensorflow + +#endif // TENSORFLOW_COMPILER_TF2XLA_XLA_EXPRESSION_H_ diff --git a/tensorflow/compiler/tf2xla/xla_expression_test.cc b/tensorflow/compiler/tf2xla/xla_expression_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..84202c931390f2d68f6d381aef0752bfff00a53d --- /dev/null +++ b/tensorflow/compiler/tf2xla/xla_expression_test.cc @@ -0,0 +1,135 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include + +#include "absl/memory/memory.h" +#include "tensorflow/compiler/tf2xla/xla_expression.h" +#include "tensorflow/compiler/tf2xla/xla_resource.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" +#include "tensorflow/compiler/xla/literal.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/status_macros.h" +#include "tensorflow/compiler/xla/tests/literal_test_util.h" +#include "tensorflow/core/framework/tensor_testutil.h" +#include "tensorflow/core/lib/core/status_test_util.h" +#include "tensorflow/core/platform/test.h" + +namespace tensorflow { +namespace { + +class XlaExpressionTest : public ::testing::Test { + protected: + void SetUp() override { + client_ = xla::ClientLibrary::LocalClientOrDie(); + builder_ = absl::make_unique("acomputation"); + constant_ = test::AsScalar(42); + op_ = xla::ConstantR0(builder_.get(), 7); + non_constant_op_ = xla::Parameter( + builder_.get(), 0, xla::ShapeUtil::MakeShape(xla::F32, {}), "x"); + resource_ = absl::make_unique( + XlaResource::kVariable, /*arg_num=*/0, /*name=*/string("avariable"), + DT_INT32, TensorShape({17, 3}), op_, /*tensor_array_size=*/-1, + /*tensor_array_gradients=*/std::set(), + /*tensor_array_multiple_writes_aggregate=*/false); + } + + xla::Client* client_; + std::unique_ptr builder_; + Tensor constant_; + xla::XlaOp op_; + xla::XlaOp non_constant_op_; + std::unique_ptr resource_; +}; + +TEST_F(XlaExpressionTest, Kind) { + EXPECT_TRUE(XlaExpression::Kind::kInvalid == XlaExpression().kind()); + EXPECT_TRUE(XlaExpression::Kind::kInvalid == XlaExpression::Invalid().kind()); + EXPECT_TRUE(XlaExpression::Kind::kConstant == + XlaExpression::Constant(constant_).kind()); + EXPECT_TRUE(XlaExpression::Kind::kXlaOp == + XlaExpression::XlaOp(op_, DT_INT32).kind()); + EXPECT_TRUE(XlaExpression::Kind::kResource == + XlaExpression::Resource(resource_.get()).kind()); +} + +TEST_F(XlaExpressionTest, HumanString) { + EXPECT_EQ("invalid", XlaExpression().HumanString()); + EXPECT_EQ("invalid", XlaExpression::Invalid().HumanString()); + EXPECT_EQ("constant", XlaExpression::Constant(constant_).HumanString()); + EXPECT_EQ("xla_op", XlaExpression::XlaOp(op_, DT_INT32).HumanString()); + EXPECT_EQ("resource", XlaExpression::Resource(resource_.get()).HumanString()); +} + +TEST_F(XlaExpressionTest, AsXlaOp) { + xla::XlaOp op_as_op = + XlaExpression::XlaOp(op_, DT_INT32).AsXlaOp(builder_.get()); + EXPECT_TRUE(op_.IsIdenticalTo(op_as_op)); + + xla::XlaOp const_as_op = + XlaExpression::Constant(constant_).AsXlaOp(builder_.get()); + TF_ASSERT_OK_AND_ASSIGN(xla::XlaComputation computation, + builder_->BuildConstantSubGraph(const_as_op)); + TF_ASSERT_OK_AND_ASSIGN(xla::Literal value, + client_->ComputeConstant(computation)); + EXPECT_TRUE(xla::LiteralTestUtil::Equal(xla::LiteralUtil::CreateR0(42), + value)); +} + +TEST_F(XlaExpressionTest, GetShape) { + EXPECT_FALSE(XlaExpression().GetShape().ok()); + EXPECT_FALSE(XlaExpression::Invalid().GetShape().ok()); + + TF_ASSERT_OK_AND_ASSIGN(TensorShape resource_shape, + XlaExpression::Resource(resource_.get()).GetShape()); + EXPECT_EQ(TensorShape({}), resource_shape); + + TF_ASSERT_OK_AND_ASSIGN(TensorShape op_shape, + XlaExpression::XlaOp(op_, DT_INT32).GetShape()); + EXPECT_EQ(TensorShape({}), op_shape); + + TF_ASSERT_OK_AND_ASSIGN(TensorShape constant_shape, + XlaExpression::Constant(constant_).GetShape()); + EXPECT_EQ(TensorShape({}), constant_shape); +} + +TEST_F(XlaExpressionTest, ResolveConstant) { + EXPECT_FALSE(XlaExpression().ResolveConstant(client_).ok()); + EXPECT_FALSE(XlaExpression::Invalid().ResolveConstant(client_).ok()); + EXPECT_FALSE( + XlaExpression::Resource(resource_.get()).ResolveConstant(client_).ok()); + + TF_ASSERT_OK_AND_ASSIGN( + absl::optional op_constant, + XlaExpression::XlaOp(op_, DT_INT32).ResolveConstant(client_)); + ASSERT_TRUE(op_constant.has_value()); + test::ExpectTensorEqual(test::AsScalar(7), *op_constant); + + TF_ASSERT_OK_AND_ASSIGN(absl::optional op_nonconstant, + XlaExpression::XlaOp(non_constant_op_, DT_FLOAT) + .ResolveConstant(client_)); + EXPECT_FALSE(op_nonconstant.has_value()); + + TF_ASSERT_OK_AND_ASSIGN( + absl::optional constant_constant, + XlaExpression::Constant(constant_).ResolveConstant(client_)); + ASSERT_TRUE(constant_constant.has_value()); + test::ExpectTensorEqual(constant_, *constant_constant); +} + +} // namespace +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/xla_helpers.cc b/tensorflow/compiler/tf2xla/xla_helpers.cc index 9a34cd8c6ae2dc6d52a3cc69168df96f5322c6da..04a5d934064a9083a41cc210b48df65bbc862fff 100644 --- a/tensorflow/compiler/tf2xla/xla_helpers.cc +++ b/tensorflow/compiler/tf2xla/xla_helpers.cc @@ -26,7 +26,6 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" #include "tensorflow/compiler/xla/client/lib/arithmetic.h" #include "tensorflow/compiler/xla/client/lib/constants.h" -#include "tensorflow/compiler/xla/client/lib/numeric.h" #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/compiler/xla/client/xla_computation.h" #include "tensorflow/compiler/xla/types.h" @@ -35,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(xla::ShapeUtil::Rank(input_shape) - 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)); @@ -121,7 +63,7 @@ xla::XlaOp XlaHelpers::FloatLiteral(xla::XlaBuilder* b, DataType data_type, /* static */ Status XlaHelpers::ReshapeLiteral( const xla::Literal& input, absl::Span dimensions, xla::Literal* output) { - if (xla::ShapeUtil::IsTuple(input.shape())) { + if (input.shape().IsTuple()) { return errors::InvalidArgument("ReshapeLiteral does not support tuples."); } xla::Shape shape = @@ -149,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, @@ -216,8 +148,7 @@ DataType XlaHelpers::SumAccumulationType(const DataType& dtype) { return dtype; } -xla::XlaOp XlaHelpers::ConvertElementType(xla::XlaBuilder* const builder, - const xla::XlaOp& operand, +xla::XlaOp XlaHelpers::ConvertElementType(const xla::XlaOp& operand, const DataType new_element_type) { xla::PrimitiveType convert_to; TF_CHECK_OK(DataTypeToPrimitiveType(new_element_type, &convert_to)); diff --git a/tensorflow/compiler/tf2xla/xla_helpers.h b/tensorflow/compiler/tf2xla/xla_helpers.h index 39578144caaadf293d24ea91aa874e56e27ecc01..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 @@ -80,8 +70,7 @@ class XlaHelpers { // A helper for creating a ConvertElementType xla op given a DataType rather // than the xla::PrimitiveType. - static xla::XlaOp ConvertElementType(xla::XlaBuilder* const builder, - const xla::XlaOp& operand, + static xla::XlaOp ConvertElementType(const xla::XlaOp& operand, const DataType new_element_type); }; diff --git a/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.cc b/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.cc index 86a78ee429e8913edb4a948727fa692083c472f4..884dc45cb11b18ae557c3da3f4192b3805cb7980 100644 --- a/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.cc +++ b/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.cc @@ -133,25 +133,36 @@ XlaJitCompiledCpuFunction::Compile( jit->executable_ = std::move(executable); jit->buffer_infos_ = std::move(buffer_infos); jit->arg_index_table_ = std::move(arg_index_table); - jit->program_shape_ = std::move(program_shape); - 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); + jit->program_shape_ = + absl::make_unique(program_shape->ToProto()); + 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_jit_compiled_cpu_function.h b/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.h index d3c8f22a8078d03d15447ed200c914390f40b04f..a5392057177e983e11787c31bb496a8947add1e6 100644 --- a/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.h +++ b/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.h @@ -80,8 +80,10 @@ class XlaJitCompiledCpuFunction { std::vector arg_names_; std::vector result_names_; - // The backing data for the program shape. - std::unique_ptr program_shape_; + // The backing data for the program shape. The proto form of program shape is + // used because the program shape is serialized and embedded in the object + // file. + std::unique_ptr program_shape_; }; } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function_test.cc b/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function_test.cc index 6d49298a6f3e8a726695fafc42f3c5341fe98b5f..8846088678b53f6b9ecff0de732d6b5c82392b5a 100644 --- a/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function_test.cc +++ b/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function_test.cc @@ -116,13 +116,13 @@ TEST(XlaJitCompiledCpuFunction, Sum) { // Check program shape. using xla::ShapeUtil; const xla::Shape s32 = ShapeUtil::MakeShape(xla::S32, {}); - const xla::ProgramShape* program_shape = function.ProgramShape(); - ASSERT_TRUE(program_shape != nullptr); - ASSERT_EQ(program_shape->parameters_size(), 2); - EXPECT_TRUE(ShapeUtil::Compatible(program_shape->parameters(0), s32)); - EXPECT_TRUE(ShapeUtil::Compatible(program_shape->parameters(1), s32)); + ASSERT_TRUE(function.ProgramShape() != nullptr); + const xla::ProgramShape program_shape(*function.ProgramShape()); + ASSERT_EQ(program_shape.parameters_size(), 2); + EXPECT_TRUE(ShapeUtil::Compatible(program_shape.parameters(0), s32)); + EXPECT_TRUE(ShapeUtil::Compatible(program_shape.parameters(1), s32)); - const xla::Shape& result = program_shape->result(); + const xla::Shape& result = program_shape.result(); ASSERT_EQ(result.element_type(), xla::TUPLE); ASSERT_EQ(ShapeUtil::TupleElementCount(result), 1); const xla::Shape& result0 = ShapeUtil::GetTupleElementShape(result, 0); diff --git a/tensorflow/compiler/tf2xla/xla_op_kernel.cc b/tensorflow/compiler/tf2xla/xla_op_kernel.cc index dd3498ef7aa242d3ad946cae5f60bc2c8853a342..1147d0c726fd4fc4b368941eca461520e36e2c24 100644 --- a/tensorflow/compiler/tf2xla/xla_op_kernel.cc +++ b/tensorflow/compiler/tf2xla/xla_op_kernel.cc @@ -20,6 +20,7 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/literal_util.h" #include "tensorflow/compiler/tf2xla/shape_util.h" #include "tensorflow/compiler/tf2xla/type_util.h" +#include "tensorflow/compiler/tf2xla/xla_compilation_device.h" #include "tensorflow/compiler/tf2xla/xla_context.h" #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/compiler/xla/client/xla_computation.h" @@ -35,40 +36,52 @@ bool XlaOpKernelContext::ValidateInputsAreSameShape(OpKernel* op) { return context_->ValidateInputsAreSameShape(op); } +XlaContext* XlaOpKernelContext::xla_context() const { + return &XlaContext::Get(context_); +} + xla::XlaBuilder* XlaOpKernelContext::builder() const { - return XlaContext::Get(this).builder(); + return xla_context()->builder(); +} + +XlaCompiler* XlaOpKernelContext::compiler() const { + return xla_context()->compiler(); } // Retrieves an XlaExpression that was allocated by a previous Op. static const XlaExpression* CastExpressionFromTensor(const Tensor& tensor) { const XlaExpression* expression = reinterpret_cast(tensor.tensor_data().data()); - CHECK(expression->handle().valid() || expression->resource() != nullptr); - VLOG(1) << "Fetched T" << expression->handle(); + CHECK(expression->kind() != XlaExpression::Kind::kInvalid) + << expression->HumanString(); return expression; } -// Retrieves an uninitialized XlaExpression from a newly-allocated tensor. -static XlaExpression* CastExpressionFromUninitializedTensor(Tensor* tensor) { +// Assigns an XlaExpression to a tensor on an XLA compilation device. +static void AssignExpressionToTensor(Tensor* tensor, + const XlaExpression& value) { const XlaExpression* expression = reinterpret_cast(tensor->tensor_data().data()); - CHECK(!expression->handle().valid()); - return const_cast(expression); + CHECK(expression->kind() == XlaExpression::Kind::kInvalid) + << expression->HumanString(); + *const_cast(expression) = value; +} + +const XlaExpression& XlaOpKernelContext::InputExpression(int index) { + return *CastExpressionFromTensor(context_->input(index)); } -// Retrieves the XlaOp from an input Tensor to an Op. This computation was -// constructed by an Op that executed previously and created the output Tensor -// using CreateOutputTensorFromComputation or CreateConstantOutputTensor. -static const xla::XlaOp& GetComputationFromTensor(const Tensor& tensor) { - return CastExpressionFromTensor(tensor)->handle(); +const XlaExpression& XlaOpKernelContext::InputExpression( + absl::string_view name) { + return *CastExpressionFromTensor(GetInputTensorByName(name)); } -const xla::XlaOp& XlaOpKernelContext::Input(int index) { - return GetComputationFromTensor(context_->input(index)); +xla::XlaOp XlaOpKernelContext::Input(int index) { + return InputExpression(index).AsXlaOp(builder()); } -const xla::XlaOp& XlaOpKernelContext::Input(absl::string_view name) { - return GetComputationFromTensor(GetInputTensorByName(name)); +xla::XlaOp XlaOpKernelContext::Input(absl::string_view name) { + return InputExpression(name).AsXlaOp(builder()); } TensorShape XlaOpKernelContext::InputShape(int index) { @@ -80,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) { @@ -125,77 +138,18 @@ Status XlaOpKernelContext::ConstantInput(absl::string_view name, Status XlaOpKernelContext::ConstantInputReshaped( int index, absl::Span new_dims, xla::Literal* constant_literal) { - const Tensor& tensor = context_->input(index); - TensorShape new_shape(new_dims); - if (tensor.NumElements() != new_shape.num_elements()) { - return errors::InvalidArgument( - context_->op_kernel().name(), " input ", index, " has shape ", - tensor.shape().DebugString(), - " but was asked to be reshaped to incompatible shape ", - new_shape.DebugString()); - } - const XlaExpression* expression = CastExpressionFromTensor(tensor); - - auto copy_tensor_to_literal = [](const Tensor& tensor, - xla::Literal* literal) { - xla::Shape literal_shape; - TF_RETURN_IF_ERROR( - TensorShapeToXLAShape(tensor.dtype(), tensor.shape(), &literal_shape)); - - *literal = xla::Literal(literal_shape); - - // memcpy over the payload ... - // TODO(phawkins): handle string types. - size_t total_bytes = tensor.TotalBytes(); - if (total_bytes > 0) { - void* dst_ptr = literal->untyped_data(); - const void* src_ptr = DMAHelper::base(&tensor); - memcpy(dst_ptr, src_ptr, total_bytes); - } - return Status::OK(); - }; - - // If the tensor has a known constant value, there is no need to invoke XLA. - if (expression->has_constant_value()) { - Tensor temp(tensor.dtype()); - if (!temp.CopyFrom(expression->constant_value(), new_shape)) { - // This should never happen. The constant should have a shape compatible - // with the enclosing Tensor. - return errors::Internal("Incompatible shapes in ConstantInputReshaped."); - } - - return copy_tensor_to_literal(temp, constant_literal); - } - - // Make sure we treat zero-element tensors as constant. - if (new_shape.num_elements() == 0) { - Tensor temp(tensor.dtype(), new_shape); - - return copy_tensor_to_literal(temp, constant_literal); - } - - xla::XlaOp handle = expression->handle(); - if (new_shape != tensor.shape()) { - // Reshape the handle to the desired shape. - handle = xla::Reshape(handle, new_shape.dim_sizes()); - } - - // The XLA layout is specified minor to major, and TensorFlow's minor - // dimension is the last one. - std::vector layout_indices(new_shape.dims()); - std::iota(layout_indices.rbegin(), layout_indices.rend(), 0); - xla::Layout layout = xla::LayoutUtil::MakeLayout(layout_indices); - - xla::StatusOr is_constant = builder()->IsConstant(handle); - if (!is_constant.ok()) { - Status status = is_constant.status(); + XlaExpression e = InputExpression(index); + xla::StatusOr> constant_or_status = + e.ResolveConstant(compiler()->client()); + if (!constant_or_status.ok()) { + Status status = constant_or_status.status(); errors::AppendToMessage(&status, "while evaluating input ", index, " of ", context_->op_kernel().type_string(), " operator as a compile-time constant."); return status; } - - if (!is_constant.ValueOrDie()) { + absl::optional constant = constant_or_status.ValueOrDie(); + if (!constant.has_value()) { return errors::InvalidArgument( "Input ", index, " to ", context_->op_kernel().type_string(), " operator must be a compile-time constant.\n" @@ -208,32 +162,23 @@ Status XlaOpKernelContext::ConstantInputReshaped( "stateful operation such as a random number generator."); } - // Ask the XLA compiler to evaluate the data handle to a literal. - xla::StatusOr constant_graph = - builder()->BuildConstantSubGraph(handle); - if (!constant_graph.ok()) { - return errors::Internal( - "Error getting a compile-time constant graph for ", - context_->op_kernel().name(), " input ", index, - ".\nError: ", constant_graph.status().error_message()); - } - xla::StatusOr computed = compiler()->client()->ComputeConstant( - constant_graph.ValueOrDie(), &layout); - if (!computed.ok()) { - return errors::Internal("Error evaluating ", context_->op_kernel().name(), - " input ", index, - " as a compile-time constant.\nError: ", - computed.status().error_message()); + Tensor temp(constant->dtype()); + if (!temp.CopyFrom(*constant, TensorShape(new_dims))) { + return errors::InvalidArgument( + context_->op_kernel().name(), " input ", index, " has shape ", + constant->shape().DebugString(), + " but was asked to be reshaped to incompatible shape ", + TensorShape(new_dims).DebugString()); } - *constant_literal = std::move(computed).ValueOrDie(); + TF_ASSIGN_OR_RETURN(*constant_literal, HostTensorToLiteral(temp)); return Status::OK(); } // Converts an int32 or int64 scalar literal to an int64. static Status LiteralToInt64Scalar(const xla::LiteralSlice& literal, int64* out) { - if (xla::ShapeUtil::Rank(literal.shape()) != 0) { + if (literal.shape().rank() != 0) { return errors::InvalidArgument("value is not a scalar"); } if (literal.shape().element_type() == xla::S32) { @@ -249,7 +194,7 @@ static Status LiteralToInt64Scalar(const xla::LiteralSlice& literal, // Converts an float32 or float64 scalar literal to a float64. static Status LiteralToFloat64Scalar(const xla::LiteralSlice& literal, double* out) { - if (xla::ShapeUtil::Rank(literal.shape()) != 0) { + if (literal.shape().rank() != 0) { return errors::InvalidArgument("value is not a scalar"); } if (literal.shape().element_type() == xla::F32) { @@ -283,8 +228,9 @@ Status XlaOpKernelContext::ConstantInputAsFloatScalar(int index, double* out) { // Converts an int32 or int64 1D literal to an int64 vector. static Status LiteralToInt64Vector(const xla::LiteralSlice& literal, std::vector* out) { - if (xla::ShapeUtil::Rank(literal.shape()) != 1) { - return errors::InvalidArgument("value is not 1D"); + if (literal.shape().rank() != 1) { + 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) { @@ -322,6 +268,15 @@ Status XlaOpKernelContext::ConstantInputReshapedToIntVector( return LiteralToInt64Vector(literal, out); } +Status XlaOpKernelContext::ConstantInputReshapedToIntVector( + absl::string_view name, std::vector* out) { + TF_ASSIGN_OR_RETURN(int index, InputIndex(this, name)); + xla::Literal literal; + TF_RETURN_IF_ERROR(ConstantInputReshaped( + index, {InputShape(index).num_elements()}, &literal)); + return LiteralToInt64Vector(literal, out); +} + Status XlaOpKernelContext::ConstantInputAsInt64Literal(int index, xla::Literal* out) { xla::Literal literal; @@ -372,7 +327,7 @@ Status XlaOpKernelContext::InputList(absl::string_view name, handles->clear(); shapes->clear(); for (const Tensor& input : inputs) { - handles->push_back(GetComputationFromTensor(input)); + handles->push_back(CastExpressionFromTensor(input)->AsXlaOp(builder())); shapes->push_back(input.shape()); } return Status::OK(); @@ -392,8 +347,8 @@ Status XlaOpKernelContext::ConstantInputList( namespace { Status ReadVariableInputTensor(const Tensor& tensor, DataType type, - const OpKernelContext* ctx, TensorShape* shape, - xla::XlaOp* value) { + const XlaOpKernelContext* ctx, + TensorShape* shape, xla::XlaOp* value) { const XlaExpression* expression = CastExpressionFromTensor(tensor); XlaResource* variable = expression->resource(); TF_RET_CHECK(variable != nullptr); @@ -411,11 +366,13 @@ Status ReadVariableInputTensor(const Tensor& tensor, DataType type, *shape = variable->shape(); } - XlaContext& xla_context = XlaContext::Get(ctx); - TF_ASSIGN_OR_RETURN( - TensorShape representation_shape, - xla_context.RepresentationShape(variable->shape(), variable->type())); - if (representation_shape == variable->shape()) { + TF_ASSIGN_OR_RETURN(xla::Shape representation_shape, + ctx->compiler()->options().shape_representation_fn( + variable->shape(), variable->type())); + xla::Shape xla_shape; + TF_RETURN_IF_ERROR( + TensorShapeToXLAShape(variable->type(), variable->shape(), &xla_shape)); + if (xla::ShapeUtil::Compatible(xla_shape, representation_shape)) { *value = variable->value(); } else { *value = xla::Reshape(variable->value(), variable->shape().dim_sizes()); @@ -428,15 +385,15 @@ Status ReadVariableInputTensor(const Tensor& tensor, DataType type, Status XlaOpKernelContext::ReadVariableInput(int index, DataType type, TensorShape* shape, xla::XlaOp* value) { - return ReadVariableInputTensor(context_->input(index), type, context_, shape, + return ReadVariableInputTensor(context_->input(index), type, this, shape, value); } Status XlaOpKernelContext::ReadVariableInput(absl::string_view name, DataType type, TensorShape* shape, xla::XlaOp* value) { - return ReadVariableInputTensor(GetInputTensorByName(name), type, context_, - shape, value); + return ReadVariableInputTensor(GetInputTensorByName(name), type, this, shape, + value); } Status XlaOpKernelContext::GetVariableTypeAndShape(int index, DataType* type, @@ -455,90 +412,58 @@ Status XlaOpKernelContext::GetVariableTypeAndShape(int index, DataType* type, return Status::OK(); } -Status XlaOpKernelContext::allocate_output(int index, const xla::Shape& shape, - Tensor** output) { - // The step's default allocator is the dummy XlaCompilationAllocator which - // simply allocates a metadata buffer to hold the expression to which it - // corresponds. - if (expected_output_dtype(index) == DT_VARIANT) { - // tensor_data() is not supported for variant Tensor (i.e., - // DataTypeCanUseMemcpy is false for DT_VARIANT), and so storing the - // XlaExpression inside the Tensor's tensor_data() does not work for - // variant. Instead construct a uint8 tensor and store the expression in its - // value. - // TODO(jpienaar): This should be refactored to stop masquerading - // XlaExpressions as Tensors. - *output = new Tensor(); - TensorShape tensor_shape; - TF_RETURN_IF_ERROR( - context_->allocate_temp(DT_UINT8, tensor_shape, *output)); - context_->set_output(index, **output); - } else { - TensorShape tensor_shape; - TF_RETURN_IF_ERROR(XLAShapeToTensorShape(shape, &tensor_shape)); - TF_RETURN_IF_ERROR(context_->allocate_output(index, tensor_shape, output)); +void XlaOpKernelContext::SetOutputExpression(int index, + const XlaExpression& expression) { + Status status = [&] { + // The step's default allocator is the dummy XlaCompilationAllocator which + // simply allocates a metadata buffer to hold the expression to which it + // corresponds. + Tensor* output = nullptr; + // Provides a special behavior for DT_VARIANT: a variant is treated as + // DT_UINT8 scalar as the type to allow mapping for variant to more generic + // types. + if (expression.dtype() == DT_VARIANT) { + // tensor_data() is not supported for variant Tensor (i.e., + // DataTypeCanUseMemcpy is false for DT_VARIANT), and so storing the + // XlaExpression inside the Tensor's tensor_data() does not work for + // variant. Instead construct a uint8 tensor and store the expression in + // its value. + // TODO(jpienaar): This should be refactored to stop masquerading + // XlaExpressions as Tensors. + output = new Tensor(); + TensorShape tensor_shape; + TF_RETURN_IF_ERROR( + context_->allocate_temp(DT_UINT8, tensor_shape, output)); + context_->set_output(index, *output); + } else { + TF_ASSIGN_OR_RETURN(TensorShape shape, expression.GetShape()); + TF_RETURN_IF_ERROR(context_->allocate_output(index, shape, &output)); + } + AssignExpressionToTensor(output, expression); + return Status::OK(); + }(); + if (!status.ok()) { + SetStatus(status); } - return Status::OK(); } void XlaOpKernelContext::SetOutput(int index, const xla::XlaOp& handle) { - // Makes the host Tensor that will refer to the expression. - Tensor* output = nullptr; - auto shape_or = builder()->GetShape(handle); - if (!shape_or.ok()) { - SetStatus(shape_or.status()); - return; - } - - OP_REQUIRES_OK(context_, - allocate_output(index, shape_or.ValueOrDie(), &output)); - - // The expression is stored in the tensor's data buffer. Fill in the - // fields now. - XlaExpression* expression = CastExpressionFromUninitializedTensor(output); - expression->set_handle(handle); + SetOutputExpression( + index, + XlaExpression::XlaOp(handle, context_->expected_output_dtype(index))); } void XlaOpKernelContext::SetConstantOutput(int index, const Tensor& constant) { - const TensorShape& shape = constant.shape(); - - xla::BorrowingLiteral literal; - OP_REQUIRES_OK(context_, HostTensorToBorrowingLiteral(constant, &literal)); - - xla::XlaOp handle = xla::ConstantLiteral(builder(), literal); - CHECK(handle.valid()); - - // Make the Tensor that will refer to the expression. - Tensor* output = nullptr; - // The step's default allocator is the dummy XlaCompilationAllocator which - // simply allocates a metadata buffer to hold the expression to which it - // corresponds. - OP_REQUIRES_OK(context_, context_->allocate_output(index, shape, &output)); - - // The expression is stored in the tensor's data buffer. Fill in the - // fields now. - XlaExpression* expression = CastExpressionFromUninitializedTensor(output); - expression->set_handle(handle); - expression->set_constant_value(constant); + SetOutputExpression(index, XlaExpression::Constant(constant)); } -void XlaOpKernelContext::SetInvalidOutput(int index) { - Tensor* output = nullptr; - OP_REQUIRES_OK(context_, - context_->allocate_output(index, TensorShape({}), &output)); - XlaExpression* expression = CastExpressionFromUninitializedTensor(output); - xla::XlaOp handle; - expression->set_handle(handle); +void XlaOpKernelContext::SetTensorListOutput(int index, + const xla::XlaOp& handle) { + SetOutputExpression(index, XlaExpression::TensorList(handle)); } void XlaOpKernelContext::SetResourceOutput(int index, XlaResource* resource) { - Tensor* output = nullptr; - // The shape of the output tensor is the shape of the resource itself - // (i.e., a scalar), not the shape of the resource's value. - OP_REQUIRES_OK(context_, - context_->allocate_output(index, TensorShape(), &output)); - XlaExpression* expression = CastExpressionFromUninitializedTensor(output); - expression->set_resource(resource); + SetOutputExpression(index, XlaExpression::Resource(resource)); } Status XlaOpKernelContext::GetResourceInput(int index, XlaResource** resource) { @@ -552,7 +477,7 @@ Status XlaOpKernelContext::GetResourceInput(int index, XlaResource** resource) { namespace { Status AssignVariableTensor(const Tensor& tensor, DataType type, - const OpKernelContext* ctx, xla::XlaOp handle, + const XlaOpKernelContext* ctx, xla::XlaOp handle, xla::XlaBuilder* builder) { const XlaExpression* expression = CastExpressionFromTensor(tensor); XlaResource* variable = expression->resource(); @@ -569,11 +494,14 @@ Status AssignVariableTensor(const Tensor& tensor, DataType type, TF_RETURN_IF_ERROR(variable->SetTypeAndShape(type, shape)); - XlaContext& xla_context = XlaContext::Get(ctx); - TF_ASSIGN_OR_RETURN(TensorShape representation_shape, - xla_context.RepresentationShape(shape, type)); - if (shape != representation_shape) { - handle = xla::Reshape(handle, representation_shape.dim_sizes()); + TF_ASSIGN_OR_RETURN( + xla::Shape representation_shape, + ctx->compiler()->options().shape_representation_fn(shape, type)); + xla::Shape xla_shape; + TF_RETURN_IF_ERROR(TensorShapeToXLAShape(type, shape, &xla_shape)); + if (!xla::ShapeUtil::Compatible(xla_shape, representation_shape)) { + handle = xla::Reshape(handle, + xla::AsInt64Slice(representation_shape.dimensions())); } return variable->SetValue(handle); } @@ -583,19 +511,15 @@ Status AssignVariableTensor(const Tensor& tensor, DataType type, Status XlaOpKernelContext::AssignVariable(int input_index, DataType type, xla::XlaOp handle) { TF_RET_CHECK(handle.valid()); - return AssignVariableTensor(context_->input(input_index), type, context_, - handle, builder()); + return AssignVariableTensor(context_->input(input_index), type, this, handle, + builder()); } Status XlaOpKernelContext::AssignVariable(absl::string_view name, DataType type, xla::XlaOp handle) { TF_RET_CHECK(handle.valid()); - return AssignVariableTensor(GetInputTensorByName(name), type, context_, - handle, builder()); -} - -XlaCompiler* XlaOpKernelContext::compiler() const { - return XlaContext::Get(context_).compiler(); + return AssignVariableTensor(GetInputTensorByName(name), type, this, handle, + builder()); } void XlaOpKernelContext::CtxFailure(const Status& s) { @@ -615,22 +539,22 @@ void XlaOpKernelContext::CtxFailureWithWarning(const char* file, int line, const xla::XlaComputation* XlaOpKernelContext::GetOrCreateMax( const DataType type) { - return XlaContext::Get(context_).GetOrCreateMax(type); + return xla_context()->GetOrCreateMax(type); } const xla::XlaComputation* XlaOpKernelContext::GetOrCreateMin( const DataType type) { - return XlaContext::Get(context_).GetOrCreateMin(type); + return xla_context()->GetOrCreateMin(type); } const xla::XlaComputation* XlaOpKernelContext::GetOrCreateAdd( const DataType type) { - return XlaContext::Get(context_).GetOrCreateAdd(type); + return xla_context()->GetOrCreateAdd(type); } const xla::XlaComputation* XlaOpKernelContext::GetOrCreateMul( const DataType type) { - return XlaContext::Get(context_).GetOrCreateMul(type); + return xla_context()->GetOrCreateMul(type); } const Tensor& XlaOpKernelContext::GetInputTensorByName(absl::string_view name) { diff --git a/tensorflow/compiler/tf2xla/xla_op_kernel.h b/tensorflow/compiler/tf2xla/xla_op_kernel.h index aa00a454968ad29495e34dc080e55b62bb0b5f7b..e44415f60bff82fb92d0cf4ec81935564a2f083a 100644 --- a/tensorflow/compiler/tf2xla/xla_op_kernel.h +++ b/tensorflow/compiler/tf2xla/xla_op_kernel.h @@ -60,6 +60,8 @@ class XlaOpKernelContext { public: explicit XlaOpKernelContext(OpKernelContext* context); + XlaContext* xla_context() const; + // Returns the XLA XlaBuilder containing the output of compilation. xla::XlaBuilder* builder() const; @@ -88,9 +90,9 @@ class XlaOpKernelContext { // Returns input `index` as a XlaOp. Unlike // OpKernelContext::Input returns a symbolic value rather than a concrete // Tensor. - const xla::XlaOp& Input(int index); + xla::XlaOp Input(int index); // Returns input `name` as a XlaOp. - const xla::XlaOp& Input(absl::string_view name); + xla::XlaOp Input(absl::string_view name); // Returns true if all inputs are the same shape, otherwise sets the // status to a non-OK value and returns false. @@ -111,14 +113,6 @@ class XlaOpKernelContext { Status ConstantInput(int index, xla::Literal* constant_literal); Status ConstantInput(absl::string_view name, xla::Literal* constant_literal); - // Evaluates input `index`, reshapes it to `new_shape` if new_shape != - // InputShape(index), and stores it in `*constant_literal`. If the input - // cannot be evaluated, e.g., because it depends on unbound parameters, - // returns a non-Ok status. If InputShape(index).num_elements() != - // new_shape.num_elements(), returns an error status. - Status ConstantInputReshaped(int index, absl::Span new_dims, - xla::Literal* constant_literal); - // Converts a constant scalar int32 or int64 tensor into an int64. Status ConstantInputAsIntScalar(int index, int64* out); Status ConstantInputAsIntScalar(absl::string_view name, int64* out); @@ -134,6 +128,8 @@ class XlaOpKernelContext { // Reshapes and converts a constant int32 or int64 tensor into a vector of // int64s. Status ConstantInputReshapedToIntVector(int index, std::vector* out); + Status ConstantInputReshapedToIntVector(absl::string_view name, + std::vector* out); // Converts a constant int32 or int64 Tensor into an xla int64 Literal. Status ConstantInputAsInt64Literal(int index, xla::Literal* out); @@ -148,6 +144,10 @@ class XlaOpKernelContext { Status ConstantInputList(absl::string_view name, std::vector* literals); + // Returns an XlaExpression describing the value of 'index'. + const XlaExpression& InputExpression(int index); + const XlaExpression& InputExpression(absl::string_view name); + // Outputs int num_outputs() const { return context_->num_outputs(); } @@ -165,9 +165,11 @@ class XlaOpKernelContext { // SetConstantOutput where possible. void SetConstantOutput(int index, const Tensor& host_tensor); - // Sets output `index` to an invalid value. - // Any subsequent attempt to consume this output will cause an error. - void SetInvalidOutput(int index); + // 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); } @@ -255,10 +257,13 @@ class XlaOpKernelContext { // Returns the tensor of input `name`. const Tensor& GetInputTensorByName(absl::string_view name); - // Wraps OpKernelContext's allocate_output method while providing special - // behavior for DT_VARIANT: a variant is treated as DT_UINT8 scalar as the - // type to allow mapping for variant to more generic types. - Status allocate_output(int index, const xla::Shape& shape, Tensor** output); + // Evaluates input `index`, reshapes it to `new_shape` if new_shape != + // InputShape(index), and stores it in `*constant_literal`. If the input + // cannot be evaluated, e.g., because it depends on unbound parameters, + // returns a non-Ok status. If InputShape(index).num_elements() != + // new_shape.num_elements(), returns an error status. + Status ConstantInputReshaped(int index, absl::Span new_dims, + xla::Literal* constant_literal); OpKernelContext* const context_; }; diff --git a/tensorflow/compiler/tf2xla/xla_op_registry.cc b/tensorflow/compiler/tf2xla/xla_op_registry.cc index 91d48125f1d21092db7e5f9307e44af9c16e4e2b..14237df69081016817fbd1a5332f22996e7f264d 100644 --- a/tensorflow/compiler/tf2xla/xla_op_registry.cc +++ b/tensorflow/compiler/tf2xla/xla_op_registry.cc @@ -18,6 +18,8 @@ limitations under the License. #include #include +#include "tensorflow/compiler/jit/flags.h" +#include "tensorflow/compiler/jit/xla_cluster_util.h" #include "tensorflow/compiler/tf2xla/type_util.h" #include "tensorflow/compiler/tf2xla/xla_context.h" #include "tensorflow/compiler/xla/client/client_library.h" @@ -128,21 +130,26 @@ XlaOpRegistry::~XlaOpRegistry() = default; // Lazily register the CPU and GPU JIT devices the first time // GetCompilationDevice is called. static void* registration_init = [®istry]() { + MarkForCompilationPassFlags* flags = GetMarkForCompilationPassFlags(); + bool cpu_global_jit = flags->tf_xla_cpu_global_jit; + mutex_lock lock(registry.mutex_); if (LaunchOpHasKernelForDevice(DeviceType(DEVICE_CPU)).ok()) { DeviceRegistration& registration = registry.compilation_devices_[DEVICE_CPU]; registration.compilation_device_name = DEVICE_CPU_XLA_JIT; - registration.requires_compilation = false; - registration.enable_jit_by_default = false; + registration.autoclustering_policy = + cpu_global_jit + ? XlaOpRegistry::AutoclusteringPolicy::kIfEnabledGlobally + : XlaOpRegistry::AutoclusteringPolicy::kIfExplicitlyRequested; registration.compile_resource_ops = false; } if (LaunchOpHasKernelForDevice(DeviceType(DEVICE_GPU)).ok()) { DeviceRegistration& registration = registry.compilation_devices_[DEVICE_GPU]; registration.compilation_device_name = DEVICE_GPU_XLA_JIT; - registration.requires_compilation = false; - registration.enable_jit_by_default = true; + registration.autoclustering_policy = + XlaOpRegistry::AutoclusteringPolicy::kIfEnabledGlobally; registration.compile_resource_ops = false; } return nullptr; @@ -341,18 +348,69 @@ std::vector XlaOpRegistry::DeviceKernels( return ops; } -/* static */ const std::unordered_set* -XlaOpRegistry::CompileTimeConstantInputs(const string& op) { - XlaOpRegistry& registry = Instance(); - mutex_lock lock(registry.mutex_); - auto it = registry.ops_.find(op); - if (it == registry.ops_.end() || it->second.empty()) { - return nullptr; +/* static */ Status XlaOpRegistry::CompileTimeConstantInputs( + const NodeDef& node_def, const OpKernel* op_kernel, const OpDef* op_def, + std::vector* result) { + result->clear(); + + DCHECK(op_def != nullptr || op_kernel != nullptr); + + std::unordered_set compile_time_constant_inputs_from_attr; + std::vector compile_time_constant_inputs_vect_from_attr; + + const std::unordered_set* compile_time_constant_inputs; + + if (GetNodeAttr(node_def, kXlaCompileTimeConstantInputsAttr, + &compile_time_constant_inputs_vect_from_attr) + .ok()) { + absl::c_copy(compile_time_constant_inputs_vect_from_attr, + std::inserter(compile_time_constant_inputs_from_attr, + compile_time_constant_inputs_from_attr.end())); + compile_time_constant_inputs = &compile_time_constant_inputs_from_attr; + } else { + const string& op = node_def.op(); + + XlaOpRegistry& registry = Instance(); + mutex_lock lock(registry.mutex_); + auto it = registry.ops_.find(op); + if (it == registry.ops_.end() || it->second.empty()) { + return Status::OK(); + } else { + // The test in IsCompatible ensures that if there are multiple matching + // registrations for this op name, they all have the same value of + // compile_time_constant_inputs, so only the first match is returned. + // + // TODO(sanjoy): This can probably be a std::vector. + compile_time_constant_inputs = + &it->second.front()->compile_time_constant_inputs; + } } - // The test in IsCompatible ensures that if there are multiple matching - // registrations for this op name, they all have the same value of - // compile_time_constant_inputs, so only the first match is returned. - return &it->second.front()->compile_time_constant_inputs; + + for (const string& input : *compile_time_constant_inputs) { + if (op_def) { + NameRangeMap input_name_ranges; + TF_RETURN_IF_ERROR( + NameRangesForNode(node_def, *op_def, &input_name_ranges, nullptr)); + auto name_range = input_name_ranges.find(input); + if (name_range == input_name_ranges.end()) { + continue; + } + + for (int i = name_range->second.first; i < name_range->second.second; + i++) { + result->push_back(i); + } + } else { + int start, stop; + TF_CHECK_OK(op_kernel->InputRange(input, &start, &stop)); + for (int i = start; i < stop; ++i) { + result->push_back(i); + } + } + } + + absl::c_sort(*result); + return Status::OK(); } /*static*/ bool XlaOpRegistry::IsMetadataOp(const string& op) { @@ -445,7 +503,7 @@ XlaOpRegistrationBuilder& XlaOpRegistrationBuilder::TypeConstraint( return *this; } -XlaOpRegistrationBuilder& XlaOpRegistrationBuilder::CompileTimeConstInput( +XlaOpRegistrationBuilder& XlaOpRegistrationBuilder::CompileTimeConstantInput( absl::string_view input_name) { registration_->compile_time_constant_inputs.emplace(input_name); return *this; diff --git a/tensorflow/compiler/tf2xla/xla_op_registry.h b/tensorflow/compiler/tf2xla/xla_op_registry.h index 413e1eeaff0c64a4976a81fe0aedfe56f776ad10..ce3b6b298c6dc5a08e7b794bbab3a28575967d28 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, @@ -66,19 +67,26 @@ class XlaOpRegistry { public: typedef OpKernel* (*Factory)(OpKernelConstruction*); + enum class AutoclusteringPolicy { + // Enable autoclustering if the user requests it, e.g., via + // experimental_jit_scope. Does not autocluster if the JIT is enabled + // globally (e.g., via the OptimizerOptions in the TF session + // configuration.) + kIfExplicitlyRequested, + // Enable autoclustering if explicitly requested, or if the JIT is enabled + // globally in the session options, or via TF_XLA_FLAGS=--tf_xla_auto_jit=N. + kIfEnabledGlobally, + // Always try to autocluster ops placed on this device. + kAlways, + }; + // Describes how to compile operators assigned to a device. struct DeviceRegistration { // The name of the an XLA compilation device to use to compile code. string compilation_device_name; - // Do operators assigned to this device require compilation? - bool requires_compilation; - - // If !requires_compilation, should we try to JIT operators on this device - // when XLA JIT compilation is enabled globally via the SessionOptions? - // (It is still possible to explicitly mark operators to JIT compile, even - // if enable_jit_by_default is false.) - bool enable_jit_by_default; + // When should we autocluster operators assigned to this device? + AutoclusteringPolicy autoclustering_policy; // Enable compilation of operators that use DT_RESOURCE types? bool compile_resource_ops = false; @@ -133,10 +141,27 @@ class XlaOpRegistry { // Returns all operations for which there are XLA kernels on any device. static std::vector GetAllRegisteredOps(); - // Returns the set of compile-time constant inputs to 'op'. Returns nullptr - // if the op is not registered. - static const std::unordered_set* CompileTimeConstantInputs( - const string& op); + // Returns (via `result`) the indices of inputs to `node_def` that must be + // compile-time constants. Returns an empty vector if the op is not + // registered. + // + // `result` is sorted. + static Status CompileTimeConstantInputs(const NodeDef& node_def, + const OpDef& op_def, + std::vector* result) { + return CompileTimeConstantInputs(node_def, /*op_kernel=*/nullptr, &op_def, + result); + } + + // Returns (via `result`) the indices of inputs to `op_kernel` that must be + // compile-time constants. + // + // `result` is sorted. + static Status CompileTimeConstantInputs(const OpKernel& op_kernel, + std::vector* result) { + return CompileTimeConstantInputs(op_kernel.def(), /*op_kernel=*/&op_kernel, + /*op_def=*/nullptr, result); + } // Returns true if `op` is a "metadata" op, one that only looks at the shapes // of its operands and not their values. @@ -213,6 +238,11 @@ class XlaOpRegistry { // whitelists must not intersect. static bool IsCompatible(const OpRegistration& x, const OpRegistration& y); + static Status CompileTimeConstantInputs(const NodeDef& node_def, + const OpKernel* op_kernel, + const OpDef* op_def, + std::vector* result); + // Map from operator name to OpRegistrations, populated by REGISTER_XLA_OP. // Registrations present under the same key must satisfy IsCompatible above, // and this is checked during registration. @@ -264,7 +294,8 @@ class XlaOpRegistrationBuilder { XlaOpRegistrationBuilder& AllowResourceTypes(); // Mark 'input_name' as an argument whose value must be known at compile-time. - XlaOpRegistrationBuilder& CompileTimeConstInput(absl::string_view input_name); + XlaOpRegistrationBuilder& CompileTimeConstantInput( + absl::string_view input_name); // Mark this op as a "metadata" op, one that only looks at the shapes of its // operands and not their values. diff --git a/tensorflow/compiler/tf2xla/xla_resource.cc b/tensorflow/compiler/tf2xla/xla_resource.cc index 63b09c8f02a60e91576544d13227d29f56d3e88c..48a3c012727acd8472d3d5d4072ae700f5497d96 100644 --- a/tensorflow/compiler/tf2xla/xla_resource.cc +++ b/tensorflow/compiler/tf2xla/xla_resource.cc @@ -18,6 +18,7 @@ limitations under the License. #include #include +#include "absl/memory/memory.h" #include "tensorflow/compiler/tf2xla/shape_util.h" #include "tensorflow/compiler/tf2xla/sharding_util.h" #include "tensorflow/compiler/tf2xla/xla_context.h" @@ -26,9 +27,42 @@ limitations under the License. namespace tensorflow { +/*static*/ absl::string_view XlaResource::KindToString(XlaResource::Kind kind) { + switch (kind) { + case XlaResource::kInvalid: + return "invalid"; + case XlaResource::kVariable: + return "variable"; + case XlaResource::kStack: + return "stack"; + case XlaResource::kTensorArray: + return "tensorarray"; + } +} + +/*static*/ std::unique_ptr XlaResource::CreateStack( + string name, DataType type, int64 max_size) { + return absl::make_unique( + XlaResource::kStack, /*arg_num=*/-1, std::move(name), type, TensorShape(), + /*initial_value=*/xla::XlaOp(), + /*max_array_size=*/max_size, + /*tensor_array_gradients=*/std::set{}, + /*tensor_array_multiple_writes_aggregate=*/false); +} + +/*static*/ std::unique_ptr XlaResource::CreateTensorArray( + string name, DataType type, TensorShape shape, xla::XlaOp initial_value, + int64 max_array_size) { + return absl::make_unique( + XlaResource::kTensorArray, /*arg_num=*/-1, std::move(name), type, shape, + initial_value, max_array_size, + /*tensor_array_gradients=*/std::set{}, + /*tensor_array_multiple_writes_aggregate=*/false); +} + XlaResource::XlaResource(Kind kind, int arg_num, string name, DataType type, TensorShape shape, const xla::XlaOp& initial_value, - int64 tensor_array_size, + int64 max_array_size, const std::set& tensor_array_gradients, bool tensor_array_multiple_writes_aggregate) : kind_(kind), @@ -38,7 +72,7 @@ XlaResource::XlaResource(Kind kind, int arg_num, string name, DataType type, shape_(std::move(shape)), value_(initial_value), initial_value_(initial_value), - tensor_array_size_(tensor_array_size), + max_array_size_(max_array_size), tensor_array_multiple_writes_aggregate_( tensor_array_multiple_writes_aggregate) { CHECK(kind_ != kInvalid); @@ -47,7 +81,7 @@ XlaResource::XlaResource(Kind kind, int arg_num, string name, DataType type, tensor_array_gradients_[gradient].reset(new XlaResource( /*kind=*/kTensorArray, /*arg_num=*/-1, /*name=*/absl::StrCat("TensorArrayGrad: ", name_), type_, shape_, - xla::XlaOp(), tensor_array_size_, /*tensor_array_gradients=*/{}, + xla::XlaOp(), max_array_size_, /*tensor_array_gradients=*/{}, /*tensor_array_multiple_writes_aggregate=*/true)); } } @@ -100,7 +134,7 @@ Status XlaResource::SetZeroValue(xla::XlaBuilder* builder) { } case kTensorArray: { TensorShape ta_shape; - ta_shape.AddDim(tensor_array_size_); + ta_shape.AddDim(max_array_size_); ta_shape.AppendShape(shape_); value_ = xla::Broadcast(XlaHelpers::Zero(builder, type_), ta_shape.dim_sizes()); @@ -108,7 +142,7 @@ Status XlaResource::SetZeroValue(xla::XlaBuilder* builder) { } case kStack: { TensorShape ta_shape; - ta_shape.AddDim(tensor_array_size_); + ta_shape.AddDim(max_array_size_); ta_shape.AppendShape(shape_); value_ = xla::Tuple(builder, {xla::Broadcast(XlaHelpers::Zero(builder, type_), @@ -133,14 +167,14 @@ Status XlaResource::GetOrCreateTensorArrayGradient(const string& source, std::unique_ptr& gradient = tensor_array_gradients_[source]; if (!gradient) { TensorShape ta_shape; - ta_shape.AddDim(tensor_array_size_); + ta_shape.AddDim(max_array_size_); ta_shape.AppendShape(shape_); xla::XlaOp gradient_value = xla::Broadcast(XlaHelpers::Zero(builder, type_), ta_shape.dim_sizes()); gradient.reset( new XlaResource(/*kind=*/kTensorArray, /*arg_num=*/-1, /*name=*/absl::StrCat("TensorArrayGrad: ", name_), - type_, shape_, gradient_value, tensor_array_size_, + type_, shape_, gradient_value, max_array_size_, /*tensor_array_gradients=*/{}, /*tensor_array_multiple_writes_aggregate=*/true)); } diff --git a/tensorflow/compiler/tf2xla/xla_resource.h b/tensorflow/compiler/tf2xla/xla_resource.h index aa9ce1b171f11ea0de4db0123098729c1c97f93a..736588bb8b89ba756cdce77eeebff8d1fcf4774c 100644 --- a/tensorflow/compiler/tf2xla/xla_resource.h +++ b/tensorflow/compiler/tf2xla/xla_resource.h @@ -18,6 +18,7 @@ limitations under the License. #include +#include "absl/strings/string_view.h" #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/framework/tensor_shape.h" @@ -35,10 +36,20 @@ class XlaResource { kTensorArray, kStack, }; + static absl::string_view KindToString(Kind kind); + + // Creates a new Stack resource. + static std::unique_ptr CreateStack(string name, DataType type, + int64 max_size); + + // Creates a new TensorArray resource. + static std::unique_ptr CreateTensorArray( + string name, DataType type, TensorShape shape, xla::XlaOp initial_value, + int64 max_array_size); XlaResource(Kind kind, int arg_num, string name, DataType type, TensorShape shape, const xla::XlaOp& initial_value, - int64 tensor_array_size, + int64 max_array_size, const std::set& tensor_array_gradients, bool tensor_array_multiple_writes_aggregate); @@ -117,12 +128,12 @@ class XlaResource { // TODO(phawkins): refactor this code to use subclasses, rather than putting // kind-specific fields in XlaResource. - // 'tensor_array_size' stores the expected size of the TensorArray or Stack. + // 'max_array_size' stores the expected size of the TensorArray or Stack. // We need to store this since sometimes TensorArrays must be initialized // lazily since we do not know the element shape at construction time. // Used by both TensorArrays and Stacks. - int64 tensor_array_size() const { return tensor_array_size_; } - void set_tensor_array_size(int64 size) { tensor_array_size_ = size; } + int64 max_array_size() const { return max_array_size_; } + void set_max_array_size(int64 size) { max_array_size_ = size; } bool tensor_array_multiple_writes_aggregate() const { return tensor_array_multiple_writes_aggregate_; @@ -149,7 +160,7 @@ class XlaResource { xla::XlaOp value_; xla::XlaOp initial_value_; - int64 tensor_array_size_ = -1; + int64 max_array_size_ = -1; bool tensor_array_multiple_writes_aggregate_ = false; std::map> tensor_array_gradients_; diff --git a/tensorflow/compiler/xla/BUILD b/tensorflow/compiler/xla/BUILD index cc7390c6e60375b4c31c38f9f7dee25730f8f51e..636e5ef721f58c009566c10a653d09a7667619c0 100644 --- a/tensorflow/compiler/xla/BUILD +++ b/tensorflow/compiler/xla/BUILD @@ -7,6 +7,7 @@ package_group( packages = [ "//tensorflow/compiler/...", "//tensorflow/contrib/tpu/...", + "//third_party/py/jax/...", ], ) @@ -67,7 +68,7 @@ cc_library( visibility = [":friends"], deps = [ ":xla_proto", - "//tensorflow/compiler/xla/legacy_flags:debug_options_flags", + "//tensorflow/compiler/xla:debug_options_flags", ], ) @@ -108,7 +109,7 @@ cc_library( name = "status_macros", srcs = ["status_macros.cc"], hdrs = ["status_macros.h"], - visibility = [":friends"], + visibility = ["//visibility:public"], deps = [ ":statusor", ":types", @@ -151,7 +152,7 @@ cc_library( ":status", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", - "//tensorflow/stream_executor", + "//tensorflow/stream_executor/lib", ], ) @@ -223,14 +224,18 @@ cc_library( name = "shape_util", srcs = [ "index_util.cc", + "layout.cc", "layout_util.cc", "primitive_util.cc", + "shape.cc", "shape_util.cc", ], hdrs = [ "index_util.h", + "layout.h", "layout_util.h", "primitive_util.h", + "shape.h", "shape_util.h", ], visibility = ["//visibility:public"], @@ -253,6 +258,23 @@ cc_library( ], ) +tf_cc_test( + name = "shape_test", + srcs = ["shape_test.cc"], + deps = [ + ":shape_util", + ":status_macros", + ":test", + ":test_helpers", + ":types", + ":util", + ":xla_data_proto", + "//tensorflow/core:lib", + "//tensorflow/core:test_main", + "@com_google_absl//absl/strings", + ], +) + tf_cc_test( name = "shape_util_test", srcs = ["shape_util_test.cc"], @@ -270,6 +292,22 @@ tf_cc_test( ], ) +tf_cc_test( + name = "primitive_util_test", + srcs = ["primitive_util_test.cc"], + deps = [ + ":shape_util", + ":status_macros", + ":test", + ":test_helpers", + ":types", + ":util", + ":xla_data_proto", + "//tensorflow/core:lib", + "//tensorflow/core:test_main", + ], +) + tf_cc_test( name = "layout_util_test", srcs = ["layout_util_test.cc"], @@ -281,6 +319,22 @@ tf_cc_test( ], ) +tf_cc_test( + name = "layout_test", + srcs = ["layout_test.cc"], + deps = [ + ":shape_util", + ":status_macros", + ":test", + ":test_helpers", + ":types", + ":util", + ":xla_data_proto", + "//tensorflow/core:test_main", + "@com_google_absl//absl/strings", + ], +) + tf_cc_test( name = "index_util_test", srcs = ["index_util_test.cc"], @@ -308,6 +362,7 @@ cc_library( ":util", ":xla_data_proto", "//tensorflow/core:lib", + "@com_google_absl//absl/base", "@com_google_absl//absl/memory", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", @@ -330,6 +385,7 @@ tf_cc_test( "//tensorflow/core:lib", "//tensorflow/core:test", "//tensorflow/core:test_main", + "@com_google_absl//absl/base", "@com_google_absl//absl/memory", "@com_google_absl//absl/strings", ], @@ -373,6 +429,7 @@ cc_library( ":literal_util", ":util", "//tensorflow/core:lib", + "@com_google_absl//absl/base", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", ], @@ -552,6 +609,7 @@ cc_library( ":types", ":util", ":xla_data_proto", + "//tensorflow/compiler/xla/service:hlo_parser", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", "@com_google_absl//absl/memory", @@ -659,6 +717,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", ], @@ -682,8 +741,8 @@ cc_library( "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/service:hlo_evaluator", "//tensorflow/compiler/xla/service:shape_inference", - "//tensorflow/compiler/xla/service/cpu:runtime_single_threaded_matmul", "//tensorflow/core:lib", + "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/memory", "@com_google_absl//absl/types:span", ], @@ -731,6 +790,73 @@ tf_cc_test( ], ) +cc_library( + name = "parse_flags_from_env", + srcs = ["parse_flags_from_env.cc"], + hdrs = ["parse_flags_from_env.h"], + deps = + [ + "//tensorflow/compiler/xla:types", + "//tensorflow/core:framework_internal", + "//tensorflow/core:lib", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/strings:str_format", + "@com_google_absl//absl/types:span", + ], +) + +tf_cc_test( + name = "parse_flags_from_env_test", + srcs = ["parse_flags_from_env_test.cc"], + deps = + [ + ":parse_flags_from_env", + "//tensorflow/compiler/xla:types", + "//tensorflow/core:framework_internal", + "//tensorflow/core:lib", + "//tensorflow/core:test", + "@com_google_absl//absl/strings:str_format", + ], +) + +cc_library( + name = "debug_options_flags", + srcs = [ + "debug_options_flags.cc", + "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", + ], +) + +tf_cc_test( + name = "debug_options_parsers_test", + size = "small", + srcs = [ + "debug_options_parsers.h", + "debug_options_parsers_test.cc", + ], + deps = + [ + "//tensorflow/compiler/xla:xla_proto", + "//tensorflow/compiler/xla/service:hlo", + "//tensorflow/core:framework_internal", + "//tensorflow/core:lib", + "//tensorflow/core:test", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/strings:str_format", + ], +) + # ----------------------------------------------------------------------------- # This is a headers target that extra XLA devices can use to prevent circular dependencies. Devices that are compiled as separate shared objects can also use it to prevent linking of library code. diff --git a/tensorflow/compiler/xla/README.md b/tensorflow/compiler/xla/README.md index 39f8caaa961dc7b57d2b45f974fc6ecf89cf6748..f9c93707f7af30a0fa0c4224240dc40848a24f66 100644 --- a/tensorflow/compiler/xla/README.md +++ b/tensorflow/compiler/xla/README.md @@ -1,7 +1,6 @@

- +

XLA (Accelerated Linear Algebra) is a domain-specific compiler for linear -algebra that optimizes TensorFlow computations. See the -[documentation](https://www.tensorflow.org/performance/xla/) for more details. +algebra that optimizes TensorFlow computations. See the [documentation](./g3doc/overview.md). 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/array2d.h b/tensorflow/compiler/xla/array2d.h index 782c966b4c57672d137569a318fb20ace14d493b..e4aca98f67d50287a83afc6f41a59458f3df2da2 100644 --- a/tensorflow/compiler/xla/array2d.h +++ b/tensorflow/compiler/xla/array2d.h @@ -104,7 +104,7 @@ std::unique_ptr> MakeLinspaceArray2D(double from, double to, int64 count = n1 * n2; NativeT step = static_cast((count > 1) ? (to - from) / (count - 1) : 0); - auto set = [&array, n1, n2](int64 index, NativeT value) { + auto set = [&array, n2](int64 index, NativeT value) { (*array)(index / n2, index % n2) = value; }; for (int64 i = 0; i < count - 1; ++i) { diff --git a/tensorflow/compiler/xla/client/BUILD b/tensorflow/compiler/xla/client/BUILD index dc097f3696e22d75d7dc72ec4877a9c8b5dda059..f5d56e8a9e1f3a05e1039f7cc90194407200f1ab 100644 --- a/tensorflow/compiler/xla/client/BUILD +++ b/tensorflow/compiler/xla/client/BUILD @@ -3,7 +3,7 @@ licenses(["notice"]) # Apache 2.0 -package(default_visibility = [":friends"]) +package(default_visibility = ["//visibility:public"]) package_group( name = "friends", @@ -33,6 +33,8 @@ cc_library( "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla:xla_proto", "//tensorflow/core:lib", + "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/types:span", ], ) @@ -66,6 +68,7 @@ cc_library( deps = [ ":global_data", ":xla_computation", + "//tensorflow/compiler/xla:debug_options_flags", "//tensorflow/compiler/xla:execution_options_util", "//tensorflow/compiler/xla:literal", "//tensorflow/compiler/xla:service_interface", @@ -74,11 +77,11 @@ cc_library( "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla:xla_proto", - "//tensorflow/compiler/xla/legacy_flags:debug_options_flags", "//tensorflow/compiler/xla/service:hlo_proto", "//tensorflow/core:lib", "@com_google_absl//absl/memory", "@com_google_absl//absl/strings", + "@com_google_absl//absl/types:optional", "@com_google_absl//absl/types:span", ], ) @@ -88,11 +91,12 @@ cc_library( srcs = ["executable_build_options.cc"], hdrs = ["executable_build_options.h"], deps = [ + "//tensorflow/compiler/xla:debug_options_flags", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla:xla_proto", "//tensorflow/compiler/xla/service:device_memory_allocator", - "//tensorflow/core:lib", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", "@com_google_absl//absl/types:optional", @@ -166,6 +170,7 @@ cc_library( "//tensorflow/core:lib", "//tensorflow/core:stream_executor_no_cuda", "@com_google_absl//absl/memory", + "@com_google_absl//absl/types:optional", ], ) @@ -189,6 +194,7 @@ cc_library( hdrs = ["xla_computation.h"], visibility = ["//visibility:public"], deps = [ + "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", @@ -234,13 +240,14 @@ tf_cc_test( deps = [ ":xla_builder", ":xla_computation", + "//tensorflow/compiler/xla:debug_options_flags", "//tensorflow/compiler/xla:literal", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:test_helpers", + "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/legacy_flags:debug_options_flags", "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/service:hlo_matchers", "//tensorflow/core:test", diff --git a/tensorflow/compiler/xla/client/client.cc b/tensorflow/compiler/xla/client/client.cc index 5dde5b432f136c16d4e3795569499ee5de709763..4f020bcec2756a328755d86ab04154d54f532465 100644 --- a/tensorflow/compiler/xla/client/client.cc +++ b/tensorflow/compiler/xla/client/client.cc @@ -20,9 +20,10 @@ limitations under the License. #include "absl/memory/memory.h" #include "absl/strings/str_cat.h" +#include "absl/types/optional.h" #include "tensorflow/compiler/xla/client/xla_computation.h" +#include "tensorflow/compiler/xla/debug_options_flags.h" #include "tensorflow/compiler/xla/execution_options_util.h" -#include "tensorflow/compiler/xla/legacy_flags/debug_options_flags.h" #include "tensorflow/compiler/xla/literal.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/types.h" @@ -42,7 +43,7 @@ StatusOr Client::Transfer(const GlobalData& data, TransferToClientRequest request; *request.mutable_data() = data.handle(); if (shape_with_layout != nullptr) { - *request.mutable_shape_with_layout() = *shape_with_layout; + *request.mutable_shape_with_layout() = shape_with_layout->ToProto(); } TransferToClientResponse response; @@ -123,7 +124,7 @@ StatusOr Client::TransferFromOutfeed( } request.set_replica_id(replica_id); if (shape_with_layout != nullptr) { - *request.mutable_shape_with_layout() = *shape_with_layout; + *request.mutable_shape_with_layout() = shape_with_layout->ToProto(); } TransferFromOutfeedResponse response; @@ -170,11 +171,14 @@ StatusOr Client::ExecuteAndTransfer( std::unique_ptr data, Execute(computation, arguments, execution_options, execution_profile)); - const Shape* shape_with_output_layout = nullptr; + absl::optional shape_with_output_layout; if (execution_options && execution_options->has_shape_with_output_layout()) { - shape_with_output_layout = &execution_options->shape_with_output_layout(); + shape_with_output_layout = + Shape(execution_options->shape_with_output_layout()); } - return Transfer(*data, shape_with_output_layout); + return Transfer(*data, shape_with_output_layout.has_value() + ? &(*shape_with_output_layout) + : nullptr); } StatusOr Client::ComputeConstant(const XlaComputation& computation, @@ -182,7 +186,7 @@ StatusOr Client::ComputeConstant(const XlaComputation& computation, ComputeConstantGraphRequest request; *request.mutable_computation() = computation.proto(); if (output_layout != nullptr) { - *request.mutable_output_layout() = *output_layout; + *request.mutable_output_layout() = output_layout->ToProto(); } ComputeConstantResponse response; @@ -210,11 +214,10 @@ StatusOr Client::LoadSnapshot(const HloSnapshot& module) { return XlaComputation(module.hlo().hlo_module()); } -StatusOr> Client::Execute( - const XlaComputation& computation, absl::Span arguments, - const ExecutionOptions* execution_options, - ExecutionProfile* execution_profile) { - ExecuteGraphRequest request; +StatusOr Client::Compile( + const XlaComputation& computation, absl::Span argument_shapes, + const ExecutionOptions* execution_options) { + CompileRequest request; *request.mutable_computation() = computation.proto(); if (execution_options == nullptr) { @@ -222,6 +225,34 @@ StatusOr> Client::Execute( } else { *request.mutable_execution_options() = *execution_options; } + if (request.execution_options().device_handles_size() > 1) { + return InvalidArgument( + "Compiling with multiple device handles is not supported. Use " + "'Execute' instead."); + } + + // The argument shapes affect how the computation is compiled. + for (const auto& arg_shape : argument_shapes) { + *request.add_input_shape_with_layout() = arg_shape.ToProto(); + } + + CompileResponse response; + VLOG(1) << "making compile request: " << request.ShortDebugString(); + Status s = stub_->Compile(&request, &response); + VLOG(1) << "done with request"; + + if (!s.ok()) { + return s; + } + TF_RET_CHECK(response.has_handle()); + return response.handle(); +} + +StatusOr> Client::Execute( + const ExecutionHandle& handle, absl::Span arguments, + ExecutionProfile* execution_profile) { + ExecuteRequest request; + *request.mutable_handle() = handle; for (GlobalData* argument : arguments) { CHECK(argument != nullptr) << "Argument pointers must not be null."; *request.add_arguments() = argument->handle(); @@ -229,7 +260,7 @@ StatusOr> Client::Execute( ExecuteResponse response; VLOG(1) << "making execute request: " << request.ShortDebugString(); - Status s = stub_->ExecuteGraph(&request, &response); + Status s = stub_->Execute(&request, &response); VLOG(1) << "done with request"; if (!s.ok()) { @@ -238,17 +269,62 @@ StatusOr> Client::Execute( if (execution_profile != nullptr) { *execution_profile = response.profile(); - if (VLOG_IS_ON(1)) { - TF_ASSIGN_OR_RETURN( - auto execution_stats, - ExecutionStatsAsString(computation, response.profile())); - VLOG(1) << execution_stats; - } } return absl::make_unique(stub_, response.output()); } +StatusOr> Client::Execute( + const XlaComputation& computation, absl::Span arguments, + const ExecutionOptions* execution_options, + ExecutionProfile* execution_profile) { + // 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()); + } + 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]); + } + } + TF_RET_CHECK(!results.empty()); + VLOG(1) << "Defaulting to device 0 result"; + return std::move(results[0]); +} + StatusOr>> Client::ExecuteParallel( absl::Span computations) { ExecuteGraphParallelRequest request; @@ -274,10 +350,11 @@ StatusOr>> Client::ExecuteParallel( } std::vector> outputs; - for (size_t i = 0; i < computations.size(); ++i) { + for (size_t i = 0; i < response.responses_size(); ++i) { outputs.push_back( absl::make_unique(stub_, response.responses(i).output())); - if (computations[i].execution_profile != nullptr) { + if (i < computations.size() && + computations[i].execution_profile != nullptr) { *computations[i].execution_profile = response.responses(i).profile(); } } @@ -312,7 +389,7 @@ StatusOr> Client::GetDeviceHandles( Status Client::Unregister(const GlobalData& data) { UnregisterRequest request; - *request.mutable_data() = data.handle(); + *request.add_data() = data.handle(); UnregisterResponse response; VLOG(1) << "making unregister request"; @@ -383,15 +460,14 @@ StatusOr Client::GetShape(const GlobalData& data) { return s; } - return response.shape(); + return Shape(response.shape()); } StatusOr Client::ExecutionStatsAsString( const XlaComputation& computation, const ExecutionProfile& profile) { TF_ASSIGN_OR_RETURN( auto computation_stats, - GetComputationStats(computation, - legacy_flags::GetDebugOptionsFromFlags())); + GetComputationStats(computation, GetDebugOptionsFromFlags())); int64 total_flops = computation_stats.flop_count() + computation_stats.transcendental_count(); if (profile.compute_time_ns() > 0) { diff --git a/tensorflow/compiler/xla/client/client.h b/tensorflow/compiler/xla/client/client.h index 6f4d33c469f1f885cfeef546e3981dc3417ef71f..eff8713ac340e82ee7633f1f078334ba73b67b2f 100644 --- a/tensorflow/compiler/xla/client/client.h +++ b/tensorflow/compiler/xla/client/client.h @@ -40,6 +40,37 @@ class Client { explicit Client(ServiceInterface* stub); virtual ~Client(); + // Compile the computation with the given argument shapes and returns the + // handle to the compiled executable. The compiled executable is cached on the + // service, and the returned handle can be used for exection without + // re-compile. + // * The shape and layout of the arguments being executed with will affect how + // the computation is compiled. If argument_shapes is empty, the parameters' + // shape and layout will be used in the compilation. + // * If execution_options is not nullptr, these options are passed to the + // service to affect how it compiles our computation. (The pointer does not + // 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, + const ExecutionOptions* execution_options = nullptr); + + // Executes the compiled executable for the given handle with the given + // arguments and returns the global data that was produced from the execution. + // * If execution_profile is not nullptr then the pointed-to ExecutionProfile + // will be filled with profile data from the execution. + StatusOr> Execute( + const ExecutionHandle& handle, absl::Span arguments, + ExecutionProfile* execution_profile = nullptr); + // Executes the computation with the given arguments and returns the global // data that was produced from the execution. // * If execution_options is not nullptr, these options are passed to the @@ -51,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/client_library.cc b/tensorflow/compiler/xla/client/client_library.cc index 27b7fa7b29206affa9f9c2e4becd9e4ea66484ab..42aae026229a49fd801cc90562fa51f604336148 100644 --- a/tensorflow/compiler/xla/client/client_library.cc +++ b/tensorflow/compiler/xla/client/client_library.cc @@ -24,12 +24,14 @@ limitations under the License. namespace xla { -LocalClientOptions::LocalClientOptions(se::Platform* platform, - int number_of_replicas, - int intra_op_parallelism_threads) +LocalClientOptions::LocalClientOptions( + se::Platform* platform, int number_of_replicas, + int intra_op_parallelism_threads, + const absl::optional>& allowed_devices) : platform_(platform), number_of_replicas_(number_of_replicas), - intra_op_parallelism_threads_(intra_op_parallelism_threads) {} + intra_op_parallelism_threads_(intra_op_parallelism_threads), + allowed_devices_(allowed_devices) {} LocalClientOptions& LocalClientOptions::set_platform(se::Platform* platform) { platform_ = platform; @@ -58,6 +60,17 @@ int LocalClientOptions::intra_op_parallelism_threads() const { return intra_op_parallelism_threads_; } +LocalClientOptions& LocalClientOptions::set_allowed_devices( + const absl::optional>& allowed_devices) { + allowed_devices_ = allowed_devices; + return *this; +} + +const absl::optional>& LocalClientOptions::allowed_devices() + const { + return allowed_devices_; +} + /* static */ ClientLibrary& ClientLibrary::Singleton() { static ClientLibrary* c = new ClientLibrary; return *c; @@ -67,9 +80,10 @@ ClientLibrary::ClientLibrary() = default; ClientLibrary::~ClientLibrary() = default; /* static */ StatusOr ClientLibrary::GetOrCreateLocalClient( - se::Platform* platform) { + se::Platform* platform, const absl::optional>& device_set) { LocalClientOptions default_options; default_options.set_platform(platform); + default_options.set_allowed_devices(device_set); return GetOrCreateLocalClient(default_options); } @@ -94,7 +108,7 @@ ClientLibrary::~ClientLibrary() = default; service_options.set_number_of_replicas(replica_count); service_options.set_intra_op_parallelism_threads( options.intra_op_parallelism_threads()); - + service_options.set_allowed_devices(options.allowed_devices()); auto instance = absl::make_unique(); TF_ASSIGN_OR_RETURN(instance->service, LocalService::NewService(service_options)); diff --git a/tensorflow/compiler/xla/client/client_library.h b/tensorflow/compiler/xla/client/client_library.h index 3ad558fa532931937fab898f7b855f0a3370eaec..62d225c6c298b26bbbd248fc1f4be64fc8efcf6b 100644 --- a/tensorflow/compiler/xla/client/client_library.h +++ b/tensorflow/compiler/xla/client/client_library.h @@ -23,9 +23,11 @@ limitations under the License. #include #include +#include #include #include +#include "absl/types/optional.h" #include "tensorflow/compiler/xla/client/compile_only_client.h" #include "tensorflow/compiler/xla/client/local_client.h" #include "tensorflow/compiler/xla/service/compile_only_service.h" @@ -43,9 +45,10 @@ namespace xla { // Options to configure the local client when it is created. class LocalClientOptions { public: - LocalClientOptions(se::Platform* platform = nullptr, - int number_of_replicas = 1, - int intra_op_parallelism_threads = -1); + LocalClientOptions( + se::Platform* platform = nullptr, int number_of_replicas = 1, + int intra_op_parallelism_threads = -1, + const absl::optional>& allowed_devices = absl::nullopt); // Set the platform backing the service, or nullptr for the default platform. LocalClientOptions& set_platform(se::Platform* platform); @@ -60,10 +63,17 @@ class LocalClientOptions { LocalClientOptions& set_intra_op_parallelism_threads(int num_threads); int intra_op_parallelism_threads() const; + // Sets the allowed_devices set for selectively constructing stream executors + // on the platform. + LocalClientOptions& set_allowed_devices( + const absl::optional>& allowed_devices); + const absl::optional>& allowed_devices() const; + private: se::Platform* platform_; int number_of_replicas_; int intra_op_parallelism_threads_; + absl::optional> allowed_devices_; }; class ClientLibrary { @@ -73,8 +83,11 @@ class ClientLibrary { // // platform : The platform the underlying XLA service should target. If // null then default platform is used. + // device_set: Set of device IDs for which the stream executor will be + // created, for the given platform. static StatusOr GetOrCreateLocalClient( - se::Platform* platform = nullptr); + se::Platform* platform = nullptr, + const absl::optional>& allowed_devices = absl::nullopt); static StatusOr GetOrCreateLocalClient( const LocalClientOptions& options); diff --git a/tensorflow/compiler/xla/client/executable_build_options.cc b/tensorflow/compiler/xla/client/executable_build_options.cc index 0f1745366b7c33e573aff2e66d85431b01488c49..ec0e08975926f36c36c854f83a40b374b12a09a4 100644 --- a/tensorflow/compiler/xla/client/executable_build_options.cc +++ b/tensorflow/compiler/xla/client/executable_build_options.cc @@ -16,6 +16,7 @@ limitations under the License. #include "tensorflow/compiler/xla/client/executable_build_options.h" #include "absl/strings/str_format.h" +#include "tensorflow/compiler/xla/debug_options_flags.h" #include "tensorflow/compiler/xla/shape_util.h" namespace xla { @@ -39,6 +40,13 @@ ExecutableBuildOptions& ExecutableBuildOptions::set_device_ordinal( int ExecutableBuildOptions::device_ordinal() const { return device_ordinal_; } +DebugOptions* ExecutableBuildOptions::mutable_debug_options() { + if (!has_debug_options()) { + debug_options_ = GetDebugOptionsFromFlags(); + } + return &debug_options_.value(); +} + ExecutableBuildOptions& ExecutableBuildOptions::set_result_layout( const Shape& shape_with_layout) { result_layout_set_ = true; @@ -50,73 +58,22 @@ 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_) { result_layout = ShapeUtil::HumanStringWithLayout(result_layout_); } - string generate_hlo_graph = "nullopt"; - if (generate_hlo_graph_.has_value()) { - generate_hlo_graph = generate_hlo_graph_.value(); - } return absl::StrFormat( "ExecutableBuildOptions{device_ordinal=%d, result_layout=%s, " - "generate_hlo_graph=%s}", - device_ordinal_, result_layout, generate_hlo_graph); -} - -ExecutableBuildOptions& ExecutableBuildOptions::set_generate_hlo_graph( - string regex) { - generate_hlo_graph_ = std::move(regex); - return *this; -} - -const absl::optional& ExecutableBuildOptions::generate_hlo_graph() - const { - return generate_hlo_graph_; -} - -ExecutableBuildOptions& ExecutableBuildOptions::set_dump_optimized_hlo_proto_to( - absl::string_view dirpath) { - dump_optimized_hlo_proto_to_ = string(dirpath); - return *this; -} - -const absl::optional& -ExecutableBuildOptions::dump_optimized_hlo_proto_to() const { - return dump_optimized_hlo_proto_to_; -} - -ExecutableBuildOptions& -ExecutableBuildOptions::set_dump_unoptimized_hlo_proto_to( - absl::string_view dirpath) { - dump_unoptimized_hlo_proto_to_ = string(dirpath); - return *this; -} - -const absl::optional& -ExecutableBuildOptions::dump_unoptimized_hlo_proto_to() const { - return dump_unoptimized_hlo_proto_to_; -} - -ExecutableBuildOptions& ExecutableBuildOptions::set_dump_per_pass_hlo_proto_to( - absl::string_view dirpath) { - dump_per_pass_hlo_proto_to_ = string(dirpath); - return *this; -} - -const absl::optional& -ExecutableBuildOptions::dump_per_pass_hlo_proto_to() const { - return dump_per_pass_hlo_proto_to_; -} - -ExecutableBuildOptions& ExecutableBuildOptions::set_hlo_profile(bool enabled) { - hlo_profile_ = enabled; - return *this; -} - -absl::optional ExecutableBuildOptions::hlo_profile() const { - return hlo_profile_; + "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 93334db88bc24f2ffbf3c7a57ee45ef238286739..1d85fb34304b95d1fccdb0b0d6a7a65e739fae18 100644 --- a/tensorflow/compiler/xla/client/executable_build_options.h +++ b/tensorflow/compiler/xla/client/executable_build_options.h @@ -19,7 +19,9 @@ limitations under the License. #include "absl/strings/string_view.h" #include "absl/types/optional.h" #include "tensorflow/compiler/xla/service/device_memory_allocator.h" +#include "tensorflow/compiler/xla/shape.h" #include "tensorflow/compiler/xla/util.h" +#include "tensorflow/compiler/xla/xla.pb.h" #include "tensorflow/compiler/xla/xla_data.pb.h" namespace xla { @@ -44,6 +46,12 @@ class ExecutableBuildOptions { ExecutableBuildOptions& set_result_layout(const Shape& shape_with_layout); const Shape* result_layout() const; + // Expose access to the XLA debug options which will be passed to the + // compilation process. + bool has_debug_options() const { return debug_options_.has_value(); } + const DebugOptions& debug_options() const { return *debug_options_; } + DebugOptions* mutable_debug_options(); + // If set, this specifies an allocator that can be used to allocate temporary // space on the device during compilation. For example, the compiler might // want to run various algorithms on the device and pick the fastest one -- it @@ -55,56 +63,22 @@ class ExecutableBuildOptions { DeviceMemoryAllocator* allocator); DeviceMemoryAllocator* device_allocator() const; - // If set, specifies a regexp of HLO graphs to dump (as in DebugOptions). - ExecutableBuildOptions& set_generate_hlo_graph(string regex); - const absl::optional& generate_hlo_graph() const; - - // If set, specifies a dirpath to dump the end-of-optimization-pipeline HLO - // protobuf to (as in DebugOptions). - ExecutableBuildOptions& set_dump_optimized_hlo_proto_to( - absl::string_view dirpath); - const absl::optional& dump_optimized_hlo_proto_to() const; - - // If set, specifies a dirpath to dump the start-of-optimization-pipeline HLO - // protobuf to (as in DebugOptions). - ExecutableBuildOptions& set_dump_unoptimized_hlo_proto_to( - absl::string_view dirpath); - const absl::optional& dump_unoptimized_hlo_proto_to() const; - - // If set, specifies a dirpath to dump the per-pass-in-pipeline HLO protobufs - // to (as in DebugOptions). - ExecutableBuildOptions& set_dump_per_pass_hlo_proto_to( - absl::string_view dirpath); - const absl::optional& dump_per_pass_hlo_proto_to() const; - - // If true, specifies that we should record an HLO profile during execution - // and log it after execution (as in DebugOptions). If nullopt the default is - // used. - ExecutableBuildOptions& set_hlo_profile(bool enabled); - absl::optional hlo_profile() const; - - void add_disabled_hlo_pass(absl::string_view pass_name) { - disabled_hlo_passes_.push_back(std::string(pass_name)); - } - const absl::Span disabled_hlo_passes() const { - return disabled_hlo_passes_; - } - // Returns a string representation of the build options, suitable for // 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: - absl::optional hlo_profile_; int device_ordinal_ = -1; Shape result_layout_; bool result_layout_set_ = false; - absl::optional generate_hlo_graph_; - absl::optional dump_optimized_hlo_proto_to_; - absl::optional dump_unoptimized_hlo_proto_to_; - absl::optional dump_per_pass_hlo_proto_to_; + absl::optional debug_options_; DeviceMemoryAllocator* device_allocator_ = nullptr; - std::vector disabled_hlo_passes_; + int num_replicas_ = 1; }; } // namespace xla diff --git a/tensorflow/compiler/xla/client/global_data.cc b/tensorflow/compiler/xla/client/global_data.cc index 2986d4060013703873b2cffb6aacbb012606d16f..f1fa13d95c035d182746d3ce5400178890aa42b1 100644 --- a/tensorflow/compiler/xla/client/global_data.cc +++ b/tensorflow/compiler/xla/client/global_data.cc @@ -18,25 +18,53 @@ limitations under the License. #include #include +#include "absl/container/flat_hash_map.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/core/platform/logging.h" namespace xla { +namespace { + +// Releases a set of global data handles owned by the parent service +// interface. +void ReleaseHandles(ServiceInterface* parent, + const absl::Span handles) { + UnregisterRequest request; + for (auto& handle : handles) { + VLOG(1) << "Requesting to unregister " << handle.ShortDebugString(); + *request.add_data() = handle; + } + UnregisterResponse response; + Status status = parent->Unregister(&request, &response); + VLOG(1) << "Done with request"; + if (!status.ok()) { + LOG(WARNING) << "Failed to unregister handles: " << status + << "; continuing anyway..."; + } +} + +} // namespace GlobalData::GlobalData(ServiceInterface* parent, GlobalDataHandle handle) : handle_(std::move(handle)), parent_(parent) {} GlobalData::~GlobalData() { - UnregisterRequest request; - *request.mutable_data() = handle_; - UnregisterResponse response; - VLOG(1) << "requesting to unregister " << handle_.ShortDebugString(); - Status s = parent_->Unregister(&request, &response); - VLOG(1) << "done with request"; + if (parent_ != nullptr) { + ReleaseHandles(parent_, {handle_}); + } +} - if (!s.ok()) { - LOG(WARNING) << "failed to unregister " << handle_.ShortDebugString() - << "; continuing anyway..."; +/* static */ void GlobalData::Release( + std::vector> instances) { + absl::flat_hash_map> + parent_handles_map; + for (auto& instance : instances) { + if (instance->parent_ != nullptr) { + parent_handles_map[instance->parent_].push_back(instance->Release()); + } + } + for (auto& parent_handles : parent_handles_map) { + ReleaseHandles(parent_handles.first, parent_handles.second); } } diff --git a/tensorflow/compiler/xla/client/global_data.h b/tensorflow/compiler/xla/client/global_data.h index b7929357d06032b55c04bf0391f7fa703ee15f17..4d48d2c53fc6171fe1940924598a4d48519c5adf 100644 --- a/tensorflow/compiler/xla/client/global_data.h +++ b/tensorflow/compiler/xla/client/global_data.h @@ -16,6 +16,10 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_CLIENT_GLOBAL_DATA_H_ #define TENSORFLOW_COMPILER_XLA_CLIENT_GLOBAL_DATA_H_ +#include +#include + +#include "absl/types/span.h" #include "tensorflow/compiler/xla/service_interface.h" #include "tensorflow/compiler/xla/xla.pb.h" #include "tensorflow/compiler/xla/xla_data.pb.h" @@ -36,7 +40,18 @@ class GlobalData { const GlobalDataHandle& handle() const { return handle_; } + // Releases a set of GlobalData handles. A single RPC will be issued + // per unique ServiceInterface of the given GlobalData objects. + static void Release(std::vector> instances); + private: + // Detaches the global data handle from the object, such that the destructor + // will not try to release it. + GlobalDataHandle Release() { + parent_ = nullptr; + return handle_; + } + GlobalDataHandle handle_; // Handle being wrapped. ServiceInterface* parent_; // Service used to unregister handle_. diff --git a/tensorflow/compiler/xla/client/lib/BUILD b/tensorflow/compiler/xla/client/lib/BUILD index a18c94c4e695a6cdcb9dcc60b64b617cecd276d8..26c5e8eb73f0908cdc2d7df65936fadeda627423 100644 --- a/tensorflow/compiler/xla/client/lib/BUILD +++ b/tensorflow/compiler/xla/client/lib/BUILD @@ -1,5 +1,7 @@ # Common computation builders for XLA. +load("//tensorflow/compiler/xla/tests:build_defs.bzl", "generate_backend_suites", "xla_test") + licenses(["notice"]) # Apache 2.0 package(default_visibility = ["//tensorflow/compiler/xla/client:friends"]) @@ -13,9 +15,6 @@ filegroup( ]), ) -load("//tensorflow/compiler/xla/tests:build_defs.bzl", "xla_test") -load("//tensorflow/compiler/xla/tests:build_defs.bzl", "generate_backend_suites") - # Generate test_suites for all backends, named "${backend}_tests". generate_backend_suites() @@ -35,6 +34,96 @@ 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"], + hdrs = ["cholesky.h"], + deps = [ + ":math", + "//tensorflow/compiler/xla:literal", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:status_macros", + "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla/client:xla_builder", + "//tensorflow/compiler/xla/client/lib:constants", + "//tensorflow/compiler/xla/client/lib:loops", + "//tensorflow/compiler/xla/client/lib:matrix", + "//tensorflow/compiler/xla/client/lib:slicing", + "//tensorflow/compiler/xla/client/lib:triangular_solve", + "//tensorflow/core:lib", + ], +) + +xla_test( + name = "cholesky_test", + srcs = ["cholesky_test.cc"], + tags = ["optonly"], + deps = [ + ":arithmetic", + ":cholesky", + ":matrix", + "//tensorflow/compiler/xla:array2d", + "//tensorflow/compiler/xla:literal", + "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:test", + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla/client:xla_builder", + "//tensorflow/compiler/xla/tests:client_library_test_base", + "//tensorflow/compiler/xla/tests:literal_test_util", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", + "//tensorflow/core:test", + ], +) + +cc_library( + name = "comparators", + srcs = ["comparators.cc"], + 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"], @@ -52,7 +141,6 @@ cc_library( xla_test( name = "constants_test", srcs = ["constants_test.cc"], - tags = ["enable_for_xla_interpreter"], deps = [ ":constants", "//tensorflow/compiler/xla:test", @@ -75,6 +163,22 @@ cc_library( ], ) +cc_library( + name = "loops", + srcs = ["loops.cc"], + hdrs = ["loops.h"], + deps = [ + ":constants", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:status_macros", + "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla/client:xla_builder", + "//tensorflow/compiler/xla/client:xla_computation", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/types:span", + ], +) + cc_library( name = "math", srcs = ["math.cc"], @@ -90,7 +194,22 @@ 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: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", @@ -104,31 +223,43 @@ xla_test( ) cc_library( - name = "numeric", - srcs = ["numeric.cc"], - hdrs = ["numeric.h"], + name = "matrix", + srcs = ["matrix.cc"], + hdrs = ["matrix.h"], deps = [ ":arithmetic", ":constants", + "//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", ], ) xla_test( - name = "numeric_test", - srcs = ["numeric_test.cc"], - tags = ["enable_for_xla_interpreter"], + name = "matrix_test", + srcs = ["matrix_test.cc"], deps = [ - ":numeric", + ":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", ], ) @@ -164,21 +295,88 @@ cc_library( deps = [ ":constants", ":math", - ":numeric", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/client:xla_builder", + "@com_google_absl//absl/base", + ], +) + +cc_library( + name = "qr", + srcs = ["qr.cc"], + hdrs = ["qr.h"], + deps = [ + ":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:xla_data_proto", + "//tensorflow/compiler/xla/client:xla_builder", "//tensorflow/core:lib", ], ) +xla_test( + name = "qr_test", + srcs = ["qr_test.cc"], + tags = ["optonly"], + deps = [ + ":matrix", + ":qr", + "//tensorflow/compiler/xla:array2d", + "//tensorflow/compiler/xla:array3d", + "//tensorflow/compiler/xla:literal", + "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:test", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla/client:xla_builder", + "//tensorflow/compiler/xla/tests:client_library_test_base", + "//tensorflow/compiler/xla/tests:literal_test_util", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", + "//tensorflow/core:test", + ], +) + +cc_library( + name = "slicing", + srcs = ["slicing.cc"], + hdrs = ["slicing.h"], + deps = [ + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla/client:xla_builder", + "@com_google_absl//absl/types:span", + ], +) + +xla_test( + name = "slicing_test", + srcs = ["slicing_test.cc"], + deps = [ + ":slicing", + "//tensorflow/compiler/xla:literal_util", + "//tensorflow/compiler/xla:test", + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla/client:xla_builder", + "//tensorflow/compiler/xla/tests:client_library_test_base", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", + ], +) + cc_library( name = "sorting", srcs = ["sorting.cc"], hdrs = ["sorting.h"], deps = [ - ":numeric", + "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/client:xla_builder", ], @@ -187,17 +385,42 @@ cc_library( xla_test( name = "sorting_test", srcs = ["sorting_test.cc"], - blacklisted_backends = [ - "cpu", - "gpu", - ], - tags = ["enable_for_xla_interpreter"], deps = [ ":sorting", "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla/client:xla_builder", + "//tensorflow/compiler/xla/tests:client_library_test_base", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", + ], +) + +cc_library( + name = "quantize", + hdrs = ["quantize.h"], + deps = [ + ":constants", + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/client:xla_builder", + "//tensorflow/core:lib", + ], +) + +xla_test( + name = "quantize_test", + srcs = ["quantize_test.cc"], + # TODO(b/122119490): re-enable TAP after fixing. + tags = [ + "notap", + ], + deps = [ + ":quantize", + "//tensorflow/compiler/xla:test", + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla:util", + "//tensorflow/compiler/xla/client:xla_builder", "//tensorflow/compiler/xla/tests:client_library_test_base", "//tensorflow/compiler/xla/tests:xla_internal_test_main", ], @@ -224,3 +447,48 @@ cc_library( "@com_google_absl//absl/strings", ], ) + +cc_library( + name = "triangular_solve", + srcs = ["triangular_solve.cc"], + hdrs = ["triangular_solve.h"], + deps = [ + "//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: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 + ], + deps = [ + ":math", + ":matrix", + ":triangular_solve", + "//tensorflow/compiler/xla:array2d", + "//tensorflow/compiler/xla:literal", + "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:test", + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla/client:xla_builder", + "//tensorflow/compiler/xla/tests:client_library_test_base", + "//tensorflow/compiler/xla/tests:literal_test_util", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", + "//tensorflow/core:test", + ], +) diff --git a/tensorflow/compiler/xla/client/lib/arithmetic.cc b/tensorflow/compiler/xla/client/lib/arithmetic.cc index e86c10f030f3990d67e5a6638100640f73c82307..3b875135af29f142463ffd783bfeaadc61ada1af 100644 --- a/tensorflow/compiler/xla/client/lib/arithmetic.cc +++ b/tensorflow/compiler/xla/client/lib/arithmetic.cc @@ -117,10 +117,70 @@ XlaOp Any(XlaOp predicates) { XlaComputation logical_or = CreateScalarOrComputation(PRED, builder); TF_ASSIGN_OR_RETURN(const Shape& predicates_shape, builder->GetShape(predicates)); - std::vector all_dimensions(ShapeUtil::Rank(predicates_shape)); + std::vector all_dimensions(predicates_shape.rank()); std::iota(all_dimensions.begin(), all_dimensions.end(), 0); return Reduce(predicates, f, logical_or, all_dimensions); }); } +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 new file mode 100644 index 0000000000000000000000000000000000000000..414bd1494cd32f32a5c37e84119de930678a776b --- /dev/null +++ b/tensorflow/compiler/xla/client/lib/cholesky.cc @@ -0,0 +1,210 @@ +/* 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/client/lib/cholesky.h" + +#include +#include + +#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/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" +#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 { + +// The Cholesky–Banachiewicz algorithm. See +// https://en.wikipedia.org/wiki/Cholesky_decomposition#The_Cholesky–Banachiewicz_and_Cholesky–Crout_algorithms +// for a description. +// +// def cholesky_unblocked(a): +// assert len(a.shape) == 2 and a.shape[-2] == a.shape[-1] +// n = a.shape[-2] +// l = np.zeros_like(a) +// for j in xrange(n): +// row = l[..., j, :j] +// row_t = np.swapaxes(row, -1, -2) +// l[..., j, j] = np.sqrt(a[..., j, j] - np.dot(row, row_t)) +// l[..., j+1:, j] = (a[..., j+1:, j] - np.dot(l[..., j+1:, :j], row_t)) / +// l[..., j, j] +// return l +XlaOp CholeskyUnblocked(XlaOp a, PrecisionConfig::Precision precision) { + XlaBuilder* builder = a.builder(); + return builder->ReportErrorOrReturn([&]() -> StatusOr { + TF_ASSIGN_OR_RETURN(Shape a_shape, builder->GetShape(a)); + const int n_dims = a_shape.rank(); + const int64 n = ShapeUtil::GetDimension(a_shape, -1); + auto major_dims = AsInt64Slice(a_shape.dimensions()) + .subspan( + /*pos=*/0, + /*len=*/n_dims - 2); + + XlaOp l = ZerosLike(a); + + // Construct the for loop body to iterate over rows. + auto body_fn = + [&](XlaOp i, absl::Span loop_vars, + XlaBuilder* body_builder) -> StatusOr> { + 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 = + Iota(body_builder, ShapeUtil::MakeShape(S32, row_shape_dims), + /*iota_dimension=*/n_dims - 1); + auto mask_range_col = + 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]; + + // row = l[..., i, :i] + // select the whole i-th row, then mask out all columns past i-1 + auto zero = ConstantR0(body_builder, 0); + auto l_i = DynamicSliceInMinorDims(body_l, {i, zero}, {1, n}); + auto row = Select(Ge(mask_range_row, i), mask_zeros_row, l_i); + // a[..., i, i] + auto a_ii = DynamicSliceInMinorDims(body_a, {i, i}, {1, 1}); + // np.dot(row, np.swapaxes(row, -1, -2)) + auto diag_dot = BatchDot(row, TransposeInMinorDims(row), precision); + // l[..., i, i] = np.sqrt(a[..., i, i] - np.dot(row, + // np.swapaxes(row, -1, -2))) + auto l_ii = Sqrt(a_ii - diag_dot); + + // a[..., i+1:, i] + // select the whole i-th column, then mask out all rows above i+1 + auto a_0i = DynamicSliceInMinorDims(body_a, {i}, {1}); + auto a_ip1i = Select(Le(mask_range_col, i), mask_zeros_col, a_0i); + + // l[..., i+1:, i] = (a[..., i+1:, i] - np.dot(l[..., i+1:, :i], r.T)) / + // l[..., i, i] + // The columns in [i, n] are zeroed out in `row`, so we just have to + // zero out rows above i+1 after the BatchDot. np.dot(l[..., :, :i], + // r.T) + auto dot = BatchDot(body_l, TransposeInMinorDims(row), precision); + // np.dot(l[..., i+1:, :i], r.T) + auto dot_ip1 = Select(Le(mask_range_col, i), mask_zeros_col, dot); + + body_l = + DynamicUpdateSliceInMinorDims(body_l, (a_ip1i - dot_ip1) / l_ii, {i}); + // Assign the diagonal after the rest of the column because otherwise the + // column assign will wrap around and overwrite the diagonal assign. + body_l = DynamicUpdateSliceInMinorDims(body_l, l_ii, {i, i}); + + return std::vector{body_a, body_l}; + }; + + TF_ASSIGN_OR_RETURN( + auto cholesky_while, + ForEachIndex(n, S32, body_fn, {a, l}, "unblocked", builder)); + + return cholesky_while[1]; + }); +} + +} // namespace + +XlaOp Cholesky(XlaOp a, int64 block_size, + PrecisionConfig::Precision precision) { + XlaBuilder* builder = a.builder(); + return builder->ReportErrorOrReturn([&]() -> StatusOr { + TF_ASSIGN_OR_RETURN(Shape a_shape, builder->GetShape(a)); + const int ndims = a_shape.rank(); + if (ndims < 2) { + return InvalidArgument( + "Argument to Cholesky must have rank >= 2; shape was %s", + a_shape.ToString()); + } + + const int64 n = ShapeUtil::GetDimension(a_shape, -1); + if (n != ShapeUtil::GetDimension(a_shape, -2)) { + return InvalidArgument( + "Argument to Cholesky must be batched square matrices; got shape %s", + ShapeUtil::HumanString(a_shape)); + } + + if (primitive_util::IsComplexType(a_shape.element_type())) { + return Unimplemented( + "Complex types are not implemented in Cholesky; got shape %s", + ShapeUtil::HumanString(a_shape)); + } + + if (block_size < 1) { + return InvalidArgument( + "block_size argument to Cholesky must be >= 1; got %d", block_size); + } + + // Blocked left-looking Cholesky factorization. + // Algorithm 1 from + // Haidar, Azzam, et al. "High-performance Cholesky factorization for + // GPU-only execution." Proceedings of General Purpose GPUs. ACM, 2017. + XlaOp l = ZerosLike(a); + for (int64 i = 0; i < n; i += block_size) { + int64 k = std::min(block_size, n - i); + if (i > 0) { + // TODO(phawkins): consider implementing SYRK for the diagonal part of + // the panel. + // a[i:, i:i+k] -= np.dot(l[i:, :i], np.transpose(l[i:i+k, :i])) + auto lhs = SliceInMinorDims(l, {i, 0}, {n, i}); + auto rhs = SliceInMinorDims(l, {i, 0}, {i + k, i}); + auto delta = BatchDot(lhs, TransposeInMinorDims(rhs), precision); + auto before = SliceInMinorDims(a, {i, i}, {n, i + k}); + a = UpdateSliceInMinorDims(a, before - delta, {i, i}); + } + + // l[i:i+k, i:i+k] = cholesky_unblocked(a[i:i+k, i:i+k]) + auto x = SliceInMinorDims(a, {i, i}, {i + k, i + k}); + auto factorized = CholeskyUnblocked(x, precision); + l = UpdateSliceInMinorDims(l, factorized, {i, i}); + + if (i + k < n) { + // 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); + l = UpdateSliceInMinorDims(l, update, {i + k, i}); + } + } + return l; + }); +} + +} // namespace xla diff --git a/tensorflow/compiler/tf2xla/lib/cholesky.h b/tensorflow/compiler/xla/client/lib/cholesky.h similarity index 87% rename from tensorflow/compiler/tf2xla/lib/cholesky.h rename to tensorflow/compiler/xla/client/lib/cholesky.h index 9a561c34b92ee45059f2a05336e682838f8e36e2..0bae26837c0f14dd0cfab82cf426becc787ec11c 100644 --- a/tensorflow/compiler/tf2xla/lib/cholesky.h +++ b/tensorflow/compiler/xla/client/lib/cholesky.h @@ -13,13 +13,13 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef TENSORFLOW_COMPILER_TF2XLA_LIB_CHOLESKY_H_ -#define TENSORFLOW_COMPILER_TF2XLA_LIB_CHOLESKY_H_ +#ifndef TENSORFLOW_COMPILER_XLA_CLIENT_LIB_CHOLESKY_H_ +#define TENSORFLOW_COMPILER_XLA_CLIENT_LIB_CHOLESKY_H_ #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/compiler/xla/xla_data.pb.h" -namespace tensorflow { +namespace xla { // Computes the Cholesky decompositions of a batch of symmetric positive // definite matrices. @@ -34,6 +34,6 @@ xla::XlaOp Cholesky( xla::XlaOp a, int64 block_size = 256, xla::PrecisionConfig::Precision precision = xla::PrecisionConfig::HIGHEST); -} // namespace tensorflow +} // namespace xla -#endif // TENSORFLOW_COMPILER_TF2XLA_LIB_CHOLESKY_H_ +#endif // TENSORFLOW_COMPILER_XLA_CLIENT_LIB_CHOLESKY_H_ diff --git a/tensorflow/compiler/xla/client/lib/cholesky_test.cc b/tensorflow/compiler/xla/client/lib/cholesky_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..095dd4fbf8b7c90047c4428b50c626c16e9c1e94 --- /dev/null +++ b/tensorflow/compiler/xla/client/lib/cholesky_test.cc @@ -0,0 +1,166 @@ +/* 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/client/lib/cholesky.h" + +#include +#include +#include + +#include "tensorflow/compiler/xla/array2d.h" +#include "tensorflow/compiler/xla/client/lib/arithmetic.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/types.h" +#include "tensorflow/core/lib/core/status_test_util.h" + +namespace { + +using xla::int64; + +using CholeskyTest = xla::ClientLibraryTestBase; + +XLA_TEST_F(CholeskyTest, Simple) { + xla::XlaBuilder builder(TestName()); + + xla::Array2D a_vals({ + {4, 6, 8, 10}, + {6, 45, 54, 63}, + {8, 54, 146, 166}, + {10, 63, 166, 310}, + }); + + xla::XlaOp a; + auto a_data = CreateR2Parameter(a_vals, 0, "a", &builder, &a); + xla::Cholesky(a, /*block_size=*/2); + + xla::Array2D expected({ + {2, 0, 0, 0}, + {3, 6, 0, 0}, + {4, 7, 9, 0}, + {5, 8, 10, 11}, + }); + + ComputeAndCompareR2(&builder, expected, {a_data.get()}, + xla::ErrorSpec(1e-4, 1e-4)); +} + +XLA_TEST_F(CholeskyTest, Simple2) { + xla::XlaBuilder builder(TestName()); + + xla::Array2D a_vals({ + {16, 24, 8, 12}, + {24, 61, 82, 48}, + {8, 82, 456, 106}, + {12, 48, 106, 62}, + }); + + xla::XlaOp a; + auto a_data = CreateR2Parameter(a_vals, 0, "a", &builder, &a); + xla::Cholesky(a); + + xla::Array2D expected( + {{4, 0, 0, 0}, {6, 5, 0, 0}, {2, 14, 16, 0}, {3, 6, 1, 4}}); + + ComputeAndCompareR2(&builder, expected, {a_data.get()}, + xla::ErrorSpec(1e-4, 1e-4)); +} + +XLA_TEST_F(CholeskyTest, SimpleBatched) { + xla::XlaBuilder builder(TestName()); + + xla::Array3D a_vals({ + { + {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, 456, 106}, + {12, 48, 106, 62}, + }, + }); + + xla::XlaOp a; + auto a_data = CreateR3Parameter(a_vals, 0, "a", &builder, &a); + xla::Cholesky(a); + + xla::Array3D expected({ + { + {2, 0, 0, 0}, + {3, 6, 0, 0}, + {4, 7, 9, 0}, + {5, 8, 10, 11}, + }, + {{4, 0, 0, 0}, {6, 5, 0, 0}, {2, 14, 16, 0}, {3, 6, 1, 4}}, + }); + + ComputeAndCompareR3(&builder, expected, {a_data.get()}, + xla::ErrorSpec(1e-4, 1e-4)); +} + +using CholeskyTestCase = std::tuple; + +class RandomCholeskyTest + : public xla::ClientLibraryTestBase, + public ::testing::WithParamInterface {}; + +XLA_TEST_P(RandomCholeskyTest, Random) { + xla::XlaBuilder builder(TestName()); + + auto test_params = GetParam(); + std::vector dimensions = {std::get<0>(test_params), + std::get<1>(test_params), + std::get<1>(test_params)}; + xla::Shape shape = xla::ShapeUtil::MakeShape(xla::F32, dimensions); + TF_ASSERT_OK_AND_ASSIGN( + auto literal, + xla::LiteralUtil::CreateRandomLiteral(shape, 0.0, 1.0)); + + auto input = xla::Parameter(&builder, 0, shape, "input"); + // Form a random positive definite matrix. + auto matrix = xla::BatchDot(input, TransposeInMinorDims(input), + xla::PrecisionConfig::HIGHEST); + + auto cholesky = xla::Cholesky(matrix, /*block_size=*/4); + + // Verify that ||matrix - cholesky * cholesky_t||_2 ~= 0 + auto verification = xla::BatchDot(cholesky, TransposeInMinorDims(cholesky), + xla::PrecisionConfig::HIGHEST); + auto delta = matrix - verification; + xla::Reduce(delta * delta, xla::ConstantR0(&builder, 0.0), + CreateScalarAddComputation(xla::F32, &builder), {0, 1, 2}); + + TF_ASSERT_OK_AND_ASSIGN(auto input_data, client_->TransferToServer(literal)); + ComputeAndCompareR0(&builder, 0.0, {input_data.get()}, + xla::ErrorSpec(1e-4, 1e-4)); +} + +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.h b/tensorflow/compiler/xla/client/lib/constants.h index 81624614c1e3599dfe116eb61d9e2edcd5230684..4e5310a380e8bda15348dae2cbb0ea9e2c381bcb 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); diff --git a/tensorflow/compiler/xla/client/lib/loops.cc b/tensorflow/compiler/xla/client/lib/loops.cc new file mode 100644 index 0000000000000000000000000000000000000000..721f987628a8ac7da3f3f872939c3f0457d6bbe2 --- /dev/null +++ b/tensorflow/compiler/xla/client/lib/loops.cc @@ -0,0 +1,123 @@ +/* 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/client/lib/loops.h" + +#include "tensorflow/compiler/xla/client/lib/constants.h" +#include "tensorflow/compiler/xla/client/xla_builder.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/status_macros.h" + +namespace xla { + +StatusOr> WhileLoopHelper( + const WhileLoopHelperConditionFunction& condition_function, + const WhileLoopHelperBodyFunction& body_function, + absl::Span initial_values, absl::string_view name, + XlaBuilder* builder) { + int arity = initial_values.size(); + std::vector var_shapes; + var_shapes.reserve(arity); + for (const XlaOp& input : initial_values) { + TF_ASSIGN_OR_RETURN(auto shape, builder->GetShape(input)); + var_shapes.push_back(std::move(shape)); + } + Shape tuple_shape = ShapeUtil::MakeTupleShape(var_shapes); + + // Unpacks a tuple into its component parts. + auto unpack_tuple = [](XlaOp tuple, int arity, XlaBuilder* builder) { + std::vector elements(arity); + for (int i = 0; i < arity; ++i) { + elements[i] = GetTupleElement(tuple, i); + } + return elements; + }; + + // Build the condition. + std::unique_ptr cond_builder = + builder->CreateSubBuilder(absl::StrCat(name, "_condition")); + { + auto parameter = Parameter(cond_builder.get(), 0, tuple_shape, "parameter"); + + TF_RETURN_IF_ERROR( + condition_function(unpack_tuple(parameter, arity, cond_builder.get()), + cond_builder.get()) + .status()); + } + TF_ASSIGN_OR_RETURN(auto cond, cond_builder->Build()); + + // Build the body. + std::unique_ptr body_builder = + builder->CreateSubBuilder(absl::StrCat(name, "_body")); + { + auto parameter = Parameter(body_builder.get(), 0, tuple_shape, "parameter"); + + TF_ASSIGN_OR_RETURN( + auto result, + body_function(unpack_tuple(parameter, arity, body_builder.get()), + body_builder.get())); + + TF_RET_CHECK(result.size() == initial_values.size()); + Tuple(body_builder.get(), result); + } + TF_ASSIGN_OR_RETURN(auto body, body_builder->Build()); + + auto outputs = While(cond, body, Tuple(builder, initial_values)); + + return unpack_tuple(outputs, arity, builder); +} + +StatusOr> ForEachIndex( + int64 num_iterations, PrimitiveType num_iterations_type, + const ForEachIndexBodyFunction& body_function, + absl::Span initial_values, absl::string_view name, + XlaBuilder* builder) { + auto while_cond_fn = [&](absl::Span values, + XlaBuilder* cond_builder) -> StatusOr { + return Lt(values[0], ConstantR0WithType(cond_builder, num_iterations_type, + num_iterations)); + }; + auto while_body_fn = + [&](absl::Span values, + XlaBuilder* body_builder) -> StatusOr> { + XlaOp iteration = values[0]; + + std::vector updated_values; + updated_values.reserve(values.size()); + updated_values.push_back(Add( + iteration, + ConstantLiteral(body_builder, LiteralUtil::One(num_iterations_type)))); + + values.remove_prefix(1); + TF_ASSIGN_OR_RETURN(std::vector body_outputs, + body_function(iteration, values, body_builder)); + updated_values.insert(updated_values.end(), body_outputs.begin(), + body_outputs.end()); + return updated_values; + }; + + std::vector values; + values.reserve(initial_values.size() + 1); + values.push_back( + ConstantLiteral(builder, LiteralUtil::Zero(num_iterations_type))); + values.insert(values.end(), initial_values.begin(), initial_values.end()); + + TF_ASSIGN_OR_RETURN(values, WhileLoopHelper(while_cond_fn, while_body_fn, + values, name, builder)); + values.erase(values.begin(), values.begin() + 1); + return values; +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/client/lib/loops.h b/tensorflow/compiler/xla/client/lib/loops.h new file mode 100644 index 0000000000000000000000000000000000000000..e11de59493e9c1de51fbdb6c45dab6d82b85a62a --- /dev/null +++ b/tensorflow/compiler/xla/client/lib/loops.h @@ -0,0 +1,72 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef TENSORFLOW_COMPILER_XLA_CLIENT_LIB_LOOPS_H_ +#define TENSORFLOW_COMPILER_XLA_CLIENT_LIB_LOOPS_H_ + +#include +#include + +#include "absl/strings/string_view.h" +#include "absl/types/span.h" +#include "tensorflow/compiler/xla/client/xla_builder.h" +#include "tensorflow/compiler/xla/client/xla_computation.h" +#include "tensorflow/compiler/xla/statusor.h" + +namespace xla { + +// Function that builds a loop condition. Takes as input a sequence of input +// values, and returns a boolean value representing if the condition succeeds. +typedef std::function(absl::Span, XlaBuilder*)> + WhileLoopHelperConditionFunction; + +// Function that builds a loop body. Takes as input a sequence of input values +// and returns a sequence of output values. +typedef std::function>(absl::Span, + XlaBuilder*)> + WhileLoopHelperBodyFunction; + +// Helper function for building an XLA while loop, where the values carried by +// the loop are a tuple of values, e.g., (a, b, c): +// while( +// condition: (a, b, c) -> bool, +// body: (a, b, c) -> (a, b, c) +// init: (a, b, c) +// ) +// 'name' is a descriptive name for the loop. +StatusOr> WhileLoopHelper( + const WhileLoopHelperConditionFunction& condition_function, + const WhileLoopHelperBodyFunction& body_function, + absl::Span initial_values, absl::string_view name, + XlaBuilder* builder); + +// Builds an XLA loop that repeats a computation `num_iterations` times. +// +// The body function (ForEachIndexBodyFunction) takes as input a pair of +// (current iteration number, loop-carried values), and returns an updated +// vector of the loop-carried values. +typedef std::function>( + XlaOp, absl::Span, XlaBuilder*)> + ForEachIndexBodyFunction; + +StatusOr> ForEachIndex( + int64 num_iterations, PrimitiveType num_iterations_type, + const ForEachIndexBodyFunction& body_function, + absl::Span initial_values, absl::string_view name, + XlaBuilder* builder); + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_CLIENT_LIB_LOOPS_H_ diff --git a/tensorflow/compiler/xla/client/lib/math.cc b/tensorflow/compiler/xla/client/lib/math.cc index d3d7edb42a38595bbf9fdb36e0dd946ae5df51f9..f7e29db690f5feb7fc4c939f3b51546ff15dafe7 100644 --- a/tensorflow/compiler/xla/client/lib/math.cc +++ b/tensorflow/compiler/xla/client/lib/math.cc @@ -16,6 +16,7 @@ limitations under the License. #include "tensorflow/compiler/xla/client/lib/math.h" #include "tensorflow/compiler/xla/client/lib/constants.h" +#include "tensorflow/compiler/xla/primitive_util.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/status_macros.h" @@ -79,25 +80,47 @@ XlaOp EvaluatePolynomial(XlaOp x, absl::Span coefficients) { // Compute an approximation of the error function complement (1 - erf(x)). XlaOp Erfc(XlaOp x) { - XlaOp abs_x = Abs(x); - XlaOp z = Exp(-x * x); + auto& b = *x.builder(); + return b.ReportErrorOrReturn([&]() -> StatusOr { + // Reject non-real non-fp inputs. (We could extend erfc to accept complex + // types, but it doesn't seem necessary at this point.) + TF_ASSIGN_OR_RETURN(auto shape, b.GetShape(x)); + if (!ShapeUtil::ElementIsFloating(shape)) { + return InvalidArgument( + "erfc only accepts real floating-point arrays or scalars, but got %s", + shape.ToString()); + } + XlaOp abs_x = Abs(x); + XlaOp z = Exp(-x * x); - 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 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); + XlaOp y = Select(Lt(abs_x, ScalarLike(x, 8.0)), z * pp / pq, z * pr / ps); - return Select(Lt(x, ScalarLike(x, 0.0)), ScalarLike(x, 2.0) - y, y); + return Select(Lt(x, ScalarLike(x, 0.0)), ScalarLike(x, 2.0) - y, y); + }); } // 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 { + // Reject non-real non-fp inputs. (We could extend erf to accept complex + // types, but it doesn't seem necessary at this point.) + TF_ASSIGN_OR_RETURN(auto shape, b.GetShape(x)); + if (!ShapeUtil::ElementIsFloating(shape)) { + return InvalidArgument( + "erf only accepts real floating-point arrays or scalars, but got %s", + shape.ToString()); + } + XlaOp z = x * x; + XlaOp pt = EvaluatePolynomial(z, kErfTCoefficient); + XlaOp pu = EvaluatePolynomial(z, kErfUCoefficient); + return x * pt / pu; + }); } // Approximation for the inverse error function from @@ -113,37 +136,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}; - - 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; - }); + 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, 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 +186,94 @@ 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); + auto& b = *input.builder(); + return b.ReportErrorOrReturn([&]() -> StatusOr { + // Reject non-real non-fp inputs. (We could extend lgamma to accept complex + // types, but it doesn't seem necessary at this point.) + TF_ASSIGN_OR_RETURN(auto shape, b.GetShape(input)); + if (!ShapeUtil::ElementIsFloating(shape)) { + return InvalidArgument( + "lgamma only accepts real floating-point arrays or scalars, but got " + "%s", + shape.ToString()); + } - 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 one_half = ScalarLike(input, 0.5); + XlaOp one = ScalarLike(input, 1); - 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 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 base_lanczos_coeff = ScalarLike(input, kBaseLanczosCoeff); + 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)); - // 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 base_lanczos_coeff = ScalarLike(input, kBaseLanczosCoeff); - 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); - } + // 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); - // 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; + 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); + + // 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,54 +284,96 @@ 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 { + // Reject non-real non-fp inputs. (We could extend digamma to accept + // complex types, but it doesn't seem necessary at this point.) + TF_ASSIGN_OR_RETURN(auto shape, b.GetShape(input)); + if (!ShapeUtil::ElementIsFloating(shape)) { + return InvalidArgument( + "digamma only accepts real floating-point arrays or scalars, but got " + "%s", + shape.ToString()); + } + + XlaOp zero = ScalarLike(input, 0); + XlaOp one_half = ScalarLike(input, 0.5); + XlaOp one = ScalarLike(input, 1); + + XlaOp pi = ScalarLike(input, M_PI); - // 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; + 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); + 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& 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_ASSIGN_OR_RETURN(auto shape, b.GetShape(x)); + if (ShapeUtil::ElementIsComplex(shape)) { + return InvalidArgument( + "RoundToEven doesn't accept complex inputs, but got %s", + shape.ToString()); + } + 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))) @@ -304,4 +407,92 @@ XlaOp Cosh(XlaOp x) { return (Exp(x) + Exp(-x)) * ScalarLike(x, 0.5); } XlaOp Sinh(XlaOp x) { return (Exp(x) - Exp(-x)) * ScalarLike(x, 0.5); } +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 = + 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 a6cafd42077367bf23ffa1f45eab31c01dc31b16..583481c7f329fec9b7c5262e820b6796654cb7a2 100644 --- a/tensorflow/compiler/xla/client/lib/math.h +++ b/tensorflow/compiler/xla/client/lib/math.h @@ -32,7 +32,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); @@ -51,6 +51,10 @@ XlaOp Lgamma(XlaOp input); // Computes an approximation of the digamma function. XlaOp Digamma(XlaOp input); +// Rounds the given number to even when the number is equidistant between two +// integers. +XlaOp RoundToEven(XlaOp x); + // Trigonometric functions // Computes the arc cosine of 'x'. @@ -82,6 +86,14 @@ 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' +// is true, otherwise returns its argument. +xla::XlaOp MaybeConjugate(xla::XlaOp x, bool conjugate); + +// 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); + } // namespace xla #endif // TENSORFLOW_COMPILER_XLA_CLIENT_LIB_MATH_H_ 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..0408ed1ba23ffdb00d61d393de29475d0f06be65 --- /dev/null +++ b/tensorflow/compiler/xla/client/lib/math_exhaustive_test.cc @@ -0,0 +1,176 @@ +/* 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}; + + // 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 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{"erf", Erf, std::erf}, +// Testcase{"erfc", Erfc, std::erfc}.set_skip_infs(), +// 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}, +// +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{"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 14c259a7fa2a47642663b65d2785e5bbdc040cfd..c2e1251fc2fa09956b9b60d4e3e13a5d0cb61d2b 100644 --- a/tensorflow/compiler/xla/client/lib/math_test.cc +++ b/tensorflow/compiler/xla/client/lib/math_test.cc @@ -30,6 +30,45 @@ 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_); + } +}; + +// 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_TEST_F(MathTest, SqrtF32) { XlaBuilder builder(TestName()); Literal zero_literal = LiteralUtil::Zero(PrimitiveType::F32); @@ -106,6 +145,28 @@ XLA_TEST_F(MathTest, Lgamma) { ComputeAndCompareR1(&builder, expected, {}, error_spec_); } +// TODO(jlebar): Fails on interpreter due to unimplemented operation. +XLA_TEST_F(MathTest, DISABLED_ON_INTERPRETER(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, @@ -136,5 +197,52 @@ XLA_TEST_F(MathTest, Digamma) { ComputeAndCompareR1(&builder, expected, {}, error_spec_); } +XLA_TEST_F(MathTest, RoundToEven) { + XlaBuilder builder(TestName()); + auto x = ConstantR1( + &builder, {-1.4, -1.5, -2.5, -0.5, 0, 0.5, 1.5, 2.5, 3.5, 4.5}); + RoundToEven(x); + + std::vector expected = {-1.0, -2.0, -2.0, -0.0, 0, + 0.0, 2.0, 2.0, 4.0, 4.0}; + + 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 new file mode 100644 index 0000000000000000000000000000000000000000..a5aea96090c59c78d20cfc10a4bd6b312be592c1 --- /dev/null +++ b/tensorflow/compiler/xla/client/lib/matrix.cc @@ -0,0 +1,339 @@ +/* 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/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/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 { + +XlaOp IdentityMatrix(XlaBuilder* builder, PrimitiveType type, int64 m, + int64 n) { + auto a = Iota(builder, U32, m); + auto b = Iota(builder, U32, n); + auto indicator = Eq(a, Broadcast(b, {m}), /*broadcast_dimensions=*/{0}); + return ConvertElementType(indicator, type); +} + +XlaOp GetMatrixDiagonal(XlaOp x) { + XlaBuilder* builder = x.builder(); + return builder->ReportErrorOrReturn([&]() -> StatusOr { + TF_ASSIGN_OR_RETURN(Shape shape, builder->GetShape(x)); + const int64 n_dims = shape.rank(); + TF_RET_CHECK(n_dims >= 2); + const int64 m = shape.dimensions(n_dims - 2); + const int64 n = shape.dimensions(n_dims - 1); + 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 indicator = Eq(b, Broadcast(a, {m}), /*broadcast_dimensions=*/{0}); + auto mask = Broadcast(indicator, major_dims); + + // TPUs don't support S64 add reduction at the moment. But fortunately + // OR-reductions work just as well for integers. + XlaComputation reducer = + 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}); + }); +} + +XlaOp TriangleMask(XlaOp x, int diagonal) { + XlaBuilder* builder = x.builder(); + return builder->ReportErrorOrReturn([&]() -> StatusOr { + TF_ASSIGN_OR_RETURN(Shape shape, builder->GetShape(x)); + const int64 n_dims = shape.rank(); + TF_RET_CHECK(n_dims >= 2); + const int64 m = shape.dimensions(n_dims - 2); + const int64 n = shape.dimensions(n_dims - 1); + absl::Span major_dims = + AsInt64Slice(shape.dimensions()).subspan(/*pos=*/0, /*len=*/n_dims - 2); + auto a = Iota(builder, S32, n); + auto b = Iota(builder, S32, m) + ConstantR0(builder, diagonal); + XlaOp indicator; + indicator = Ge(b, Broadcast(a, {m}), /*broadcast_dimensions=*/{0}); + return Broadcast(indicator, major_dims); + }); +} + +XlaOp Triangle(XlaOp x, bool lower) { + return lower ? Select(TriangleMask(x, 0), x, ZerosLike(x)) + : Select(TriangleMask(x, -1), ZerosLike(x), x); +} + +XlaOp UpperTriangle(XlaOp x) { return Triangle(x, false); } + +XlaOp LowerTriangle(XlaOp x) { return Triangle(x, true); } + +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_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; + }; + + 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); + } + + for (auto d : y_config) { + insert(y_map, d); + } + + for (auto d : output_config) { + insert(output_map, d); + } + + 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); + } + } + + 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); + } + } + + 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, ','); + + 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; +} + +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); + }); +} + +XlaOp TransposeInMinorDims(XlaOp x) { + XlaBuilder* builder = x.builder(); + return builder->ReportErrorOrReturn([&]() -> StatusOr { + TF_ASSIGN_OR_RETURN(Shape shape, builder->GetShape(x)); + const int64 n_dims = shape.rank(); + TF_RET_CHECK(n_dims >= 2); + std::vector permutation(n_dims); + std::iota(permutation.begin(), permutation.end(), 0); + std::swap(permutation[n_dims - 1], permutation[n_dims - 2]); + return Transpose(x, permutation); + }); +} + +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 new file mode 100644 index 0000000000000000000000000000000000000000..491f1eab4cbffbbf9df70d4c35a61351df3e98aa --- /dev/null +++ b/tensorflow/compiler/xla/client/lib/matrix.h @@ -0,0 +1,115 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef TENSORFLOW_COMPILER_XLA_CLIENT_LIB_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" + +namespace xla { + +// Returns an m x n matrix with 1s on the diagonal elements, zeros everywhere +// 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); + +// Returns a lower-triangular mask, i.e., true below the `diagonal`-th diagonal +// and false above that diagonal. +XlaOp TriangleMask(XlaOp x, int diagonal); + +// Get the upper or lower triangle part of the last two dimensions +XlaOp Triangle(XlaOp x, bool lower); + +// Get the upper triangle part of the last two dimensions +XlaOp UpperTriangle(XlaOp x); + +// Get the lower triangle part of the last two dimensions +XlaOp LowerTriangle(XlaOp x); + +// Multiplies slices of two tensors in batches. + +// Multiplies all slices of `Tensor` `x` and `y` (each slice can be +// viewed as an element of a batch), and arranges the individual results +// in a single output tensor of the same batch size. +// +// The input tensors `x` and `y` are 2-D or higher with shape `[..., r_x, c_x]` +// and `[..., r_y, c_y]`. +// +// The output tensor is 2-D or higher with shape `[..., r_o, c_o]`, where: +// +// r_o = c_x if transpose_x else r_x +// c_o = r_y if transpose_y else c_y +// +// It is computed as: +// +// output[..., :, :] = matrix(x[..., :, :]) * matrix(y[..., :, :]) +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); + +// Transposes `x` in its minor dimensions if `transpose` is true, otherwise +// returns `x` unchanged. +xla::XlaOp MaybeTransposeInMinorDims(xla::XlaOp x, bool transpose); + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_CLIENT_LIB_MATRIX_H_ diff --git a/tensorflow/compiler/xla/client/lib/matrix_test.cc b/tensorflow/compiler/xla/client/lib/matrix_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..79cf529ee94b044ee0af788522200cd28c778997 --- /dev/null +++ b/tensorflow/compiler/xla/client/lib/matrix_test.cc @@ -0,0 +1,181 @@ +/* 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/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 { + +class MatrixTest : public ClientLibraryTestBase { + protected: + template + void TestMatrixDiagonal(); +}; + +XLA_TEST_F(MatrixTest, Triangle) { + XlaBuilder builder(TestName()); + Array3D input(2, 3, 4); + input.FillIota(0); + + XlaOp a; + auto a_data = CreateR3Parameter(input, 0, "a", &builder, &a); + LowerTriangle(a); + Array3D expected({{{0, 0, 0, 0}, {4, 5, 0, 0}, {8, 9, 10, 0}}, + {{12, 0, 0, 0}, {16, 17, 0, 0}, {20, 21, 22, 0}}}); + + ComputeAndCompareR3(&builder, expected, {a_data.get()}); +} + +template +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()}); +} + +XLA_TEST_F(MatrixTest, GetMatrixDiagonal_S32) { TestMatrixDiagonal(); } + +XLA_TEST_F(MatrixTest, GetMatrixDiagonal_S64) { TestMatrixDiagonal(); } + +XLA_TEST_F(MatrixTest, GetMatrixDiagonal_F32) { TestMatrixDiagonal(); } + +Array3D BatchedAValsFull() { + return {{ + {2, 0, 1, 2}, + {3, 6, 0, 1}, + {4, 7, 9, 0}, + {5, 8, 10, 11}, + }, + { + {16, 24, 8, 12}, + {24, 61, 82, 48}, + {8, 82, 456, 106}, + {12, 48, 106, 62}, + }}; +} + +XLA_TEST_F(MatrixTest, RowBatchDot) { + 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}); + BatchDot(l_index, TransposeInMinorDims(row)); + + 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/numeric.cc b/tensorflow/compiler/xla/client/lib/numeric.cc deleted file mode 100644 index 377654220b5df4487e9e194361473d54ff46a54e..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/client/lib/numeric.cc +++ /dev/null @@ -1,89 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include -#include - -#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/numeric.h" - -namespace xla { - -XlaOp IdentityMatrix(XlaBuilder* builder, PrimitiveType type, int64 m, - int64 n) { - auto a = Iota(builder, type, m); - auto b = Iota(builder, type, n); - auto indicator = Eq(a, Broadcast(b, {m}), /*broadcast_dimensions=*/{0}); - return ConvertElementType(indicator, type); -} - -XlaOp GetMatrixDiagonal(XlaOp x) { - XlaBuilder* builder = x.builder(); - return builder->ReportErrorOrReturn([&]() -> StatusOr { - TF_ASSIGN_OR_RETURN(Shape shape, builder->GetShape(x)); - const int64 n_dims = ShapeUtil::Rank(shape); - TF_RET_CHECK(n_dims >= 2); - const int64 m = shape.dimensions(n_dims - 2); - const int64 n = shape.dimensions(n_dims - 1); - 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 indicator = Eq(b, Broadcast(a, {m}), /*broadcast_dimensions=*/{0}); - auto mask = Broadcast(indicator, major_dims); - - // TPUs don't support S64 add reduction at the moment. But fortunately - // OR-reductions work just as well for integers. - XlaComputation reducer = - 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}); - }); -} - -XlaOp Triangle(XlaOp x, bool lower) { - XlaBuilder* builder = x.builder(); - return builder->ReportErrorOrReturn([&]() -> StatusOr { - TF_ASSIGN_OR_RETURN(Shape shape, builder->GetShape(x)); - const int64 n_dims = ShapeUtil::Rank(shape); - TF_RET_CHECK(n_dims >= 2); - const int64 m = shape.dimensions(n_dims - 2); - const int64 n = shape.dimensions(n_dims - 1); - 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); - xla::XlaOp indicator; - if (lower) { - indicator = Ge(b, Broadcast(a, {m}), /*broadcast_dimensions=*/{0}); - } else { - indicator = Le(b, Broadcast(a, {m}), /*broadcast_dimensions=*/{0}); - } - auto mask = Broadcast(indicator, major_dims); - - return Select(mask, x, Zeros(builder, shape)); - }); -} - -XlaOp UpperTriangle(XlaOp x) { return Triangle(x, false); } - -XlaOp LowerTriangle(XlaOp x) { return Triangle(x, true); } - -} // namespace xla diff --git a/tensorflow/compiler/xla/client/lib/numeric.h b/tensorflow/compiler/xla/client/lib/numeric.h deleted file mode 100644 index efd8cdc25724198633e0bf1c48c4e7d9e4b4c9e1..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/client/lib/numeric.h +++ /dev/null @@ -1,48 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#ifndef TENSORFLOW_COMPILER_XLA_CLIENT_LIB_NUMERIC_H_ -#define TENSORFLOW_COMPILER_XLA_CLIENT_LIB_NUMERIC_H_ - -#include "tensorflow/compiler/xla/client/xla_builder.h" -#include "tensorflow/compiler/xla/types.h" -#include "tensorflow/compiler/xla/xla_data.pb.h" - -namespace xla { - -// Returns a rank 1 tensor of `type` containing values [0, 1, 2, ...]. -XlaOp Iota(XlaBuilder* builder, PrimitiveType type, int64 size); - -// Returns an m x n matrix with 1s on the diagonal elements, zeros everywhere -// 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 upper or lower triangle part of the last two dimensions -XlaOp Triangle(XlaOp x, bool lower); - -// Get the upper triangle part of the last two dimensions -XlaOp UpperTriangle(XlaOp x); - -// Get the lower triangle part of the last two dimensions -XlaOp LowerTriangle(XlaOp x); - -} // namespace xla - -#endif // TENSORFLOW_COMPILER_XLA_CLIENT_LIB_NUMERIC_H_ diff --git a/tensorflow/compiler/xla/client/lib/numeric_test.cc b/tensorflow/compiler/xla/client/lib/numeric_test.cc deleted file mode 100644 index 7d6aedd49462bd4f075f90d0b0f85c40f1191aa1..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/client/lib/numeric_test.cc +++ /dev/null @@ -1,68 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/compiler/xla/client/lib/numeric.h" -#include "tensorflow/compiler/xla/client/xla_builder.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 NumericTest : public ClientLibraryTestBase { - protected: - template - void TestMatrixDiagonal(); -}; - -XLA_TEST_F(NumericTest, Triangle) { - XlaBuilder builder(TestName()); - Array3D input(2, 3, 4); - input.FillIota(0); - - XlaOp a; - auto a_data = CreateR3Parameter(input, 0, "a", &builder, &a); - LowerTriangle(a); - Array3D expected({{{0, 0, 0, 0}, {4, 5, 0, 0}, {8, 9, 10, 0}}, - {{12, 0, 0, 0}, {16, 17, 0, 0}, {20, 21, 22, 0}}}); - - ComputeAndCompareR3(&builder, expected, {a_data.get()}); -} - -template -void NumericTest::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()}); -} - -XLA_TEST_F(NumericTest, GetMatrixDiagonal_S32) { TestMatrixDiagonal(); } - -XLA_TEST_F(NumericTest, GetMatrixDiagonal_S64) { TestMatrixDiagonal(); } - -XLA_TEST_F(NumericTest, GetMatrixDiagonal_F32) { TestMatrixDiagonal(); } - -} // namespace -} // namespace xla diff --git a/tensorflow/compiler/xla/client/lib/prng.cc b/tensorflow/compiler/xla/client/lib/prng.cc index 6ef81689489d8117d5951bcb75693c2e3413e4d6..85b9e1827dcef5ed907d893277deb5a52f8f30e9 100644 --- a/tensorflow/compiler/xla/client/lib/prng.cc +++ b/tensorflow/compiler/xla/client/lib/prng.cc @@ -15,20 +15,19 @@ limitations under the License. #include +#include "absl/base/casts.h" #include "tensorflow/compiler/xla/client/lib/constants.h" #include "tensorflow/compiler/xla/client/lib/math.h" -#include "tensorflow/compiler/xla/client/lib/numeric.h" #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/compiler/xla/util.h" -#include "tensorflow/core/lib/core/casts.h" namespace xla { namespace { // Rotates a 32-bit integer 'v' left by 'distance' bits. -XlaOp RotateLeftS32(XlaOp v, int distance) { - return (v << ConstantR0(v.builder(), distance)) | - ShiftRightLogical(v, ConstantR0(v.builder(), 32 - distance)); +XlaOp RotateLeftU32(XlaOp v, int distance) { + return (v << ConstantR0(v.builder(), distance)) | + ShiftRightLogical(v, ConstantR0(v.builder(), 32 - distance)); } using ThreeFry2x32State = std::array; @@ -38,13 +37,16 @@ using ThreeFry2x32State = std::array; // http://www.thesalmons.org/john/random123/papers/random123sc11.pdf ThreeFry2x32State ThreeFry2x32(ThreeFry2x32State input, ThreeFry2x32State key) { XlaBuilder* builder = input[0].builder(); + key[0] = BitcastConvertType(key[0], U32); + key[1] = BitcastConvertType(key[1], U32); + // Rotation distances specified by the Threefry2x32 algorithm. constexpr std::array rotations = {13, 15, 26, 6, 17, 29, 16, 24}; ThreeFry2x32State x; std::array ks; // 0x1BD11BDA is a parity constant specified by the ThreeFry2x32 algorithm. - ks[2] = ConstantR0(builder, 0x1BD11BDA); + ks[2] = ConstantR0(builder, 0x1BD11BDA); for (int i = 0; i < 2; ++i) { ks[i] = key[i]; x[i] = input[i]; @@ -58,7 +60,7 @@ ThreeFry2x32State ThreeFry2x32(ThreeFry2x32State input, ThreeFry2x32State key) { // amount 'rotation'. auto round = [](ThreeFry2x32State v, int rotation) { v[0] = v[0] + v[1]; - v[1] = RotateLeftS32(v[1], rotation); + v[1] = RotateLeftU32(v[1], rotation); v[1] = v[0] ^ v[1]; return v; }; @@ -70,74 +72,83 @@ ThreeFry2x32State ThreeFry2x32(ThreeFry2x32State input, ThreeFry2x32State key) { x = round(x, rotations[2]); x = round(x, rotations[3]); x[0] = x[0] + ks[1]; - x[1] = x[1] + ks[2] + ConstantR0(builder, 1); + x[1] = x[1] + ks[2] + ConstantR0(builder, 1); x = round(x, rotations[4]); x = round(x, rotations[5]); x = round(x, rotations[6]); x = round(x, rotations[7]); x[0] = x[0] + ks[2]; - x[1] = x[1] + ks[0] + ConstantR0(builder, 2); + x[1] = x[1] + ks[0] + ConstantR0(builder, 2); x = round(x, rotations[0]); x = round(x, rotations[1]); x = round(x, rotations[2]); x = round(x, rotations[3]); x[0] = x[0] + ks[0]; - x[1] = x[1] + ks[1] + ConstantR0(builder, 3); + x[1] = x[1] + ks[1] + ConstantR0(builder, 3); x = round(x, rotations[4]); x = round(x, rotations[5]); x = round(x, rotations[6]); x = round(x, rotations[7]); x[0] = x[0] + ks[1]; - x[1] = x[1] + ks[2] + ConstantR0(builder, 4); + x[1] = x[1] + ks[2] + ConstantR0(builder, 4); x = round(x, rotations[0]); x = round(x, rotations[1]); x = round(x, rotations[2]); x = round(x, rotations[3]); x[0] = x[0] + ks[2]; - x[1] = x[1] + ks[0] + ConstantR0(builder, 5); + x[1] = x[1] + ks[0] + ConstantR0(builder, 5); return x; } -} // namespace +// Returns the inputs with unique counter values for ThreeFry2x32. +ThreeFry2x32State GetInputs(const int64 size, XlaBuilder* builder) { + ThreeFry2x32State inputs; + inputs[0] = Iota(builder, U32, size); + inputs[1] = inputs[0] + ConstantR0(builder, size); + return inputs; +} -XlaOp StatelessRngUniform(std::array seeds, const Shape& shape, - XlaOp minval, XlaOp maxval) { - XlaBuilder* builder = seeds[0].builder(); - if (shape.element_type() != F32) { - return builder->ReportError(Unimplemented( - "Types other than F32 are not implemented by StatelessRngUniform.")); - } - ThreeFry2x32State key = seeds; +XlaOp StatelessRngUniformU32(std::array key, const Shape& shape) { + XlaBuilder* builder = key[0].builder(); const int64 size = ShapeUtil::ElementsIn(shape); - const int64 half_size = CeilOfRatio(size, 2); const bool size_is_odd = (half_size * 2 != size); - - // Fill the generator inputs with unique counter values. - ThreeFry2x32State inputs; - inputs[0] = Iota(builder, S32, half_size); - inputs[1] = inputs[0] + ConstantR0(builder, half_size); + ThreeFry2x32State inputs = GetInputs(half_size, builder); ThreeFry2x32State outputs = ThreeFry2x32(inputs, key); - if (size_is_odd) { outputs[1] = Slice(outputs[1], {0}, {half_size - 1}, {1}); } + auto result = ConcatInDim(builder, outputs, 0); + return Reshape(result, AsInt64Slice(shape.dimensions())); +} - auto bits = Reshape(ConcatInDim(builder, outputs, 0), - AsInt64Slice(shape.dimensions())); +XlaOp StatelessRngUniformU64(std::array key, const Shape& shape) { + XlaBuilder* builder = key[0].builder(); + const int64 size = ShapeUtil::ElementsIn(shape); + ThreeFry2x32State inputs = GetInputs(size, builder); + ThreeFry2x32State outputs = ThreeFry2x32(inputs, key); + // low 32 bit: outputs[0], high 32 bit: outputs[1] + auto result = ConvertElementType(outputs[0], U64) | + ShiftLeft(ConvertElementType(outputs[1], U64), + ConstantR0WithType(builder, U64, 32)); + return Reshape(result, AsInt64Slice(shape.dimensions())); +} + +XlaOp StatelessRngUniformF32(XlaOp bits, XlaOp minval, XlaOp maxval) { + XlaBuilder* builder = bits.builder(); // Form 23 random mantissa bits, with a leading 1 bit. The leading 1 bit // forces the random bits into the mantissa. constexpr int kFloatBits = 32; constexpr int kMantissaBits = 23; bits = ShiftRightLogical( - bits, ConstantR0(builder, kFloatBits - kMantissaBits)) | - ConstantR0(builder, tensorflow::bit_cast(1.0f)); + bits, ConstantR0(builder, kFloatBits - kMantissaBits)) | + ConstantR0(builder, absl::bit_cast(1.0f)); auto floats = BitcastConvertType(bits, F32); // We have a floating point number in the range [1.0, 2.0). @@ -147,4 +158,47 @@ XlaOp StatelessRngUniform(std::array seeds, const Shape& shape, return floats * (maxval - minval) + minval; } +XlaOp StatelessRngUniformInt(XlaOp bits, XlaOp minval, XlaOp maxval, + PrimitiveType type, PrimitiveType unsigned_type) { + XlaBuilder* builder = bits.builder(); + // TODO(b/72573764): Generate real uniform integer distribution. + // The following algorithm is the same one that TF uses right now, but it's + // uniform only when maxval - minval is a divisor of the range that bits is + // generated from. + auto range = BitcastConvertType(maxval, unsigned_type) - + BitcastConvertType(minval, unsigned_type); + auto dist = Rem(bits, range); + auto dist_div_2 = + ShiftRightLogical(dist, ConstantR0WithType(builder, unsigned_type, 1)); + + return minval + BitcastConvertType(dist_div_2, type) + + BitcastConvertType(dist - dist_div_2, type); +} + +} // namespace + +XlaOp StatelessRngUniform(std::array seeds, const Shape& shape, + XlaOp minval, XlaOp maxval) { + XlaBuilder* builder = seeds[0].builder(); + PrimitiveType type = shape.element_type(); + switch (type) { + case F32: { + auto bits = StatelessRngUniformU32(seeds, shape); + return StatelessRngUniformF32(bits, minval, maxval); + } + case S32: { + auto bits = StatelessRngUniformU32(seeds, shape); + return StatelessRngUniformInt(bits, minval, maxval, type, U32); + } + case S64: { + auto bits = StatelessRngUniformU64(seeds, shape); + return StatelessRngUniformInt(bits, minval, maxval, type, U64); + } + default: + return builder->ReportError(Unimplemented( + "Types other than F32, S32 and S64 are not implemented by " + "StatelessRngUniform.")); + } +} + } // namespace xla diff --git a/tensorflow/compiler/xla/client/lib/prng.h b/tensorflow/compiler/xla/client/lib/prng.h index ad000b1fa1d0655c8fccc0bb33379f2499b77f26..2603818de26888566a533334e49b039b126db66e 100644 --- a/tensorflow/compiler/xla/client/lib/prng.h +++ b/tensorflow/compiler/xla/client/lib/prng.h @@ -25,7 +25,7 @@ namespace xla { // Returns a tensor containing 'shape' random values uniformly distributed in // the range [minval, maxval). Requires 2 32-bit integer seeds. -// Currently only 'shape's of type F32 are implemented. +// Currently only 'shape's of type F32, S32 and S64 are implemented. XlaOp StatelessRngUniform(std::array seeds, const Shape& shape, XlaOp minval, XlaOp maxval); diff --git a/tensorflow/compiler/xla/client/lib/qr.cc b/tensorflow/compiler/xla/client/lib/qr.cc new file mode 100644 index 0000000000000000000000000000000000000000..640412ec8bcffd2565b11ba25b87f6bf6438d848 --- /dev/null +++ b/tensorflow/compiler/xla/client/lib/qr.cc @@ -0,0 +1,393 @@ +/* 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/client/lib/qr.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 { + +std::vector ConcatVectors(absl::Span xs, + absl::Span ys) { + std::vector output(xs.size() + ys.size()); + std::copy(xs.begin(), xs.end(), output.begin()); + std::copy(ys.begin(), ys.end(), output.begin() + xs.size()); + return output; +} + +// Computes a Householder reflection of the form: +// H = I - tau v v.T. +// such that +// H . ( x1 ) = ( x1 ) +// ( x2 ) = ( x2 ) +// ( ... ) = ( ... ) +// ( xk ) = ( beta ) +// ( ... ) ( 0 ) +// ( ... ) ( 0 ) +// Unlike the usual formulation, we allow the caller to supply 'k' rather than +// only providing the relevant part of 'x' to maintain XLA's static shape +// invariant. In addition, the implementation supports batching. +// Pseudo-code, without batching: +// alpha = x[k] +// x_copy = np.copy(x) +// x_copy[:k+1] = 0 +// xnorm = norm2(x_copy) +// if xnorm == 0: +// beta = alpha +// tau = 0 +// v = np.zeros_like(x) +// else: +// beta = - np.sign(alpha) * dlapy2(alpha, xnorm) +// tau = (beta - alpha) / beta +// v = x / (alpha - beta) +// v[k] = 1 +// return (v, tau, beta) +// TODO(phawkins): LAPACK's xLARFG implementation has code for handling +// overflows in the norm/beta calculations. Perhaps do the same here. +Status House(XlaOp x, XlaOp k, absl::Span batch_dims, + const int64 m, XlaOp* v, XlaOp* tau, XlaOp* beta) { + XlaBuilder* const builder = x.builder(); + TF_ASSIGN_OR_RETURN(Shape x_shape, builder->GetShape(x)); + const PrimitiveType type = x_shape.element_type(); + + std::vector batch_dim_ids(batch_dims.size()); + std::iota(batch_dim_ids.begin(), batch_dim_ids.end(), 0); + const int64 minor_dim = batch_dims.size(); + + XlaOp zero = ScalarLike(x, 0.0); + XlaOp one = ScalarLike(x, 1.0); + + // alpha = x[k] + XlaOp alpha = Reshape(DynamicSliceInMinorDims(x, {k}, {1}), batch_dims); + + // Compute x[k+1:] (padded with zeros in elements 0..k) + XlaOp iota = Iota(builder, S32, m); + XlaOp x_after_k = Mul(x, ConvertElementType(Gt(iota, k), type), + /*broadcast_dimensions=*/{minor_dim}); + + // sigma = np.dot(x[k+1:], x[k+1:]) + auto sigma = Reduce(x_after_k * x_after_k, zero, + CreateScalarAddComputation(type, builder), {minor_dim}); + // mu = np.sqrt(x[k]*x[k] + sigma) + auto mu = Sqrt(Square(alpha) + sigma); + + auto sigma_is_zero = Eq(sigma, zero); + + *beta = Select(sigma_is_zero, alpha, -Sign(alpha) * mu); + *tau = Select(sigma_is_zero, Broadcast(zero, batch_dims), + (*beta - alpha) / *beta); + auto divisor = + Select(sigma_is_zero, Broadcast(one, batch_dims), alpha - *beta); + + auto e_k = Broadcast(ConvertElementType(Eq(iota, k), type), + std::vector(batch_dims.size(), 1)); + + // Form v as [0, 0, ..., 1] ++ x[k+1:] / divisor + // If sigma is zero, x[k+1:] is zero, so use any non-zero divisor. + *v = e_k + Div(x_after_k, divisor, /*broadcast_dimensions=*/batch_dim_ids); + return Status::OK(); +} + +// Householder QR decomposition. Algorithm 5.2.1 from Golub and Van +// Loan "Matrix Computations", 4th Edition. This is an unblocked implementation +// used as an inner routine of the blocked implementation. +// Algorithm is adapted slightly so the shapes inside the loop are static, at +// the cost of some redundant computation. Since this is used as an inner block +// kernel, accumulates the Householder transformations (vs, taus) rather than +// the matrix q. +// Equivalent Python code, without batching: +// def qr(a): +// m = a.shape[0] +// n = a.shape[1] +// vs = np.zeros([m, n]) +// taus = np.zeros([n]) +// for j in xrange(min(m, n)): +// v, tau, beta = house(a[:, j], j) +// # Unusually, we apply the Householder transformation to the entirety of +// # a, wasting FLOPs to maintain the static shape invariant that XLA +// # requires. For columns that precede j this has no effect. +// a[:, :] -= tau * np.dot(v[:, np.newaxis], +// np.dot(v[np.newaxis, :], a[:, :])) +// # Form column j explicitly rather than relying on the precision of the +// # Householder update. +// a[j, j] = beta +// a[j+1:, j] = np.zeros([m - j - 1], dtype=a.dtype) +// vs[:, j] = v +// taus[j] = tau +// return (q, vs, taus) +struct QRBlockResult { + // The factored R value + XlaOp r; + + // Representation of the Householder matrices I - beta v v.T + XlaOp taus; // Shape: [..., n] + XlaOp vs; // Shape: [..., m, n] +}; +StatusOr QRBlock(XlaOp a, PrecisionConfig::Precision precision) { + XlaBuilder* builder = a.builder(); + TF_ASSIGN_OR_RETURN(Shape a_shape, builder->GetShape(a)); + const int num_dims = a_shape.rank(); + if (num_dims < 2) { + return InvalidArgument("Argument to QR must have rank >= 2; got shape %s", + a_shape.ToString()); + } + PrimitiveType type = a_shape.element_type(); + + const int64 m = ShapeUtil::GetDimension(a_shape, -2); + const int64 n = ShapeUtil::GetDimension(a_shape, -1); + + 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); + } + + std::vector batch_dim_indices(num_batch_dims); + std::iota(batch_dim_indices.begin(), batch_dim_indices.end(), 0); + + auto qr_body_fn = [&](XlaOp j, absl::Span values, + XlaBuilder* builder) -> StatusOr> { + auto a = values[0]; + auto vs = values[1]; + auto taus = values[2]; + + // v, beta = house(a[:, j], j) + auto x = DynamicSliceInMinorDims(a, {j}, {1}); + XlaOp v, tau, beta; + TF_RETURN_IF_ERROR(House(Collapse(x, {num_dims - 2, num_dims - 1}), j, + batch_dims, m, &v, &tau, &beta)); + + std::vector shape = batch_dims; + shape.push_back(1); + shape.push_back(m); + auto v_broadcast = Reshape(v, shape); + // a[:, :] -= tau * np.dot(v[:, np.newaxis], + // np.dot(v[np.newaxis, :], a[:, :])) + auto vva = BatchDot(v_broadcast, a, precision); + vva = BatchDot(TransposeInMinorDims(v_broadcast), vva, precision); + a = a - Mul(tau, vva, + /*broadcast_dimensions=*/batch_dim_indices); + + // It is more precise to populate column 'k' explicitly, rather than + // computing it implicitly by applying the Householder transformation. + // a[k,k] = beta + // a[k+1:,k] = np.zeros([m-k-1], dtype=a.dtype) + auto iota = Reshape(Iota(a.builder(), S32, m), {m, 1}); + auto predecessor_mask = ConvertElementType(Lt(iota, j), type); + auto mask = Broadcast(ConvertElementType(Eq(iota, j), type), + std::vector(batch_dims.size(), 1)); + auto new_x = Mul(x, predecessor_mask, + /*broadcast_dimensions=*/{num_dims - 2, num_dims - 1}) + + Mul(beta, mask, /*broadcast_dimensions=*/batch_dim_indices); + a = DynamicUpdateSliceInMinorDims(a, new_x, {j}); + + // vs[:, j] = v + vs = DynamicUpdateSliceInMinorDims( + vs, Reshape(v, ConcatVectors(batch_dims, {m, 1})), {j}); + // taus[j] = tau + taus = DynamicUpdateSliceInMinorDims( + taus, Reshape(tau, ConcatVectors(batch_dims, {1})), {j}); + return std::vector{a, vs, taus}; + }; + + auto vs = Zeros( + builder, ShapeUtil::MakeShape(type, ConcatVectors(batch_dims, {m, n}))); + auto taus = Zeros(builder, + ShapeUtil::MakeShape(type, ConcatVectors(batch_dims, {n}))); + + TF_ASSIGN_OR_RETURN(auto values, ForEachIndex(std::min(m, n), S32, qr_body_fn, + {a, vs, taus}, "qr", builder)); + + QRBlockResult result; + result.r = values[0]; + result.vs = values[1]; + result.taus = values[2]; + return result; +} + +// Computes W and Y such that I-WY is equivalent to the sequence of Householder +// transformations given by vs and taus. +// Golub and van Loan, "Matrix Computations", algorithm 5.1.2. +// Y = np.zeros([m, n]) +// W = np.zeros([m, n]) +// Y[:, 0] = vs[:, 0] +// W[:, 0] = -taus[0] * vs[:, 0] +// for j in xrange(1, n): +// v = vs[:, j] +// z = -taus[j] * v - taus[j] * np.dot(W, np.dot(Y.T, v)) +// W[:, j] = z +// Y[:, j] = v +// return W +// There is no need to return Y since at termination of the loop it is equal to +// vs. +StatusOr ComputeWYRepresentation(PrimitiveType type, + absl::Span batch_dims, + XlaOp vs, XlaOp taus, int64 m, int64 n, + PrecisionConfig::Precision precision) { + std::vector batch_dim_indices(batch_dims.size()); + std::iota(batch_dim_indices.begin(), batch_dim_indices.end(), 0); + int64 n_index = batch_dims.size() + 1; + + auto body_fn = [&](XlaOp j, absl::Span values, + XlaBuilder* builder) -> StatusOr> { + auto w = values[0]; + auto y = values[1]; + const auto vs = values[2]; + const auto taus = values[3]; + + // Want j values in range [1, ... n). + j = j + ConstantR0(builder, 1); + // vs has shape [..., m, 1] + auto v = DynamicSliceInMinorDims(vs, {j}, {1}); + // beta has shape [..., 1] + auto beta = DynamicSliceInMinorDims(taus, {j}, {1}); + + // yv has shape [..., n, 1] + auto yv = BatchDot(TransposeInMinorDims(y), v, precision); + // wyv has shape [..., m, 1] + auto wyv = BatchDot(w, yv, precision); + + auto z = Mul( + -beta, v + wyv, + /*broadcast_dimensions=*/ConcatVectors(batch_dim_indices, {n_index})); + + w = DynamicUpdateSliceInMinorDims(w, z, {j}); + y = DynamicUpdateSliceInMinorDims(y, v, {j}); + + return std::vector{w, y, vs, taus}; + }; + + XlaBuilder* builder = vs.builder(); + auto w = Zeros(builder, + ShapeUtil::MakeShape(type, ConcatVectors(batch_dims, {m, n}))); + auto y = w; + auto v = SliceInMinorDims(vs, {0}, {1}); + auto beta = SliceInMinorDims(taus, {0}, {1}); + y = UpdateSliceInMinorDims(y, v, {0}); + auto bv = + Mul(-beta, v, + /*broadcast_dimensions=*/ConcatVectors(batch_dim_indices, {n_index})); + w = UpdateSliceInMinorDims(w, bv, {0}); + + TF_ASSIGN_OR_RETURN( + auto values, + ForEachIndex(n - 1, S32, body_fn, {w, y, vs, taus}, "wy", builder)); + return values[0]; +} + +} // namespace + +// Block Householder QR Factorization. Algorithm 5.2.2 of Golub and van Loan. +// def qr_blocked(a, block_size): +// m = a.shape[0] +// n = a.shape[1] +// q = np.eye(m) +// for i in xrange(0, min(m, n), block_size): +// k = min(block_size, min(m, n) - s) +// (a, vs, taus) = qr(a[i:, i:i+k]) +// y = vs +// w = ComputeWYRepresentation(vs, taus, m-i, k) +// a[i:, i+r:] += np.dot(y, np.dot(w.T, a[i:, i+k:])) +// q[:, i:] += np.dot(q[:, i:], np.dot(w, y.T)) +// return (q, a) +// TODO(phawkins): consider using UT transformations (in the form I - V U V') +// rather than WY transformations. +StatusOr QRDecomposition( + XlaOp a, bool full_matrices, int64 block_size, + PrecisionConfig::Precision precision) { + XlaBuilder* builder = a.builder(); + TF_ASSIGN_OR_RETURN(Shape a_shape, builder->GetShape(a)); + const int num_dims = a_shape.rank(); + if (num_dims < 2) { + return InvalidArgument("Arguments to QR must have rank >= 2: got shape %s", + a_shape.ToString()); + } + PrimitiveType type = a_shape.element_type(); + + const int64 m = ShapeUtil::GetDimension(a_shape, -2); + const int64 n = ShapeUtil::GetDimension(a_shape, -1); + const int64 p = std::min(m, n); + + if (block_size < 1) { + return InvalidArgument("block_size argument to QR must be >= 1; got %d", + block_size); + } + + 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 q = Broadcast(IdentityMatrix(builder, type, m, m), batch_dims); + for (int64 i = 0; i < p; i += block_size) { + int64 k = std::min(block_size, p - i); + + auto a_block = SliceInMinorDims(a, {i, i}, {m, i + k}); + TF_ASSIGN_OR_RETURN(auto qr_block, QRBlock(a_block, precision)); + + a = UpdateSliceInMinorDims(a, qr_block.r, {i, i}); + + // Compute the I-WY block representation of a product of Householder + // matrices. + TF_ASSIGN_OR_RETURN( + auto w, ComputeWYRepresentation(type, batch_dims, qr_block.vs, + qr_block.taus, m - i, k, precision)); + auto y = qr_block.vs; + + // a[i:, i+k:] += np.dot(Y, np.dot(W.T, a[i:, i+k:])) + auto a_panel = SliceInMinorDims(a, {i, i + k}, {m, n}); + auto a_update = BatchDot(TransposeInMinorDims(w), a_panel, precision); + a_update = BatchDot(y, a_update, precision); + a_panel = a_panel + a_update; + a = UpdateSliceInMinorDims(a, a_panel, {i, i + k}); + + // q[:, i:] += np.dot(np.dot(q[:, i:], W), Y.T)) + auto q_panel = SliceInMinorDims(q, {0, i}, {m, m}); + auto q_update = BatchDot(q_panel, w, precision); + q_update = BatchDot(q_update, TransposeInMinorDims(y), precision); + q_panel = q_panel + q_update; + q = UpdateSliceInMinorDims(q, q_panel, {0, i}); + } + QRDecompositionResult result; + + // full_matrices is false when only a partial result in needed. Slice to the + // needed dimensions here. + if (!full_matrices) { + q = SliceInMinorDims(q, {0, 0}, {m, p}); + a = SliceInMinorDims(a, {0, 0}, {p, n}); + } + result.q = q; + result.r = a; + return result; +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/client/lib/qr.h b/tensorflow/compiler/xla/client/lib/qr.h new file mode 100644 index 0000000000000000000000000000000000000000..827c8eeca05ef09a0d77363eb3c40961b95813d8 --- /dev/null +++ b/tensorflow/compiler/xla/client/lib/qr.h @@ -0,0 +1,42 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef TENSORFLOW_COMPILER_XLA_CLIENT_LIB_QR_H_ +#define TENSORFLOW_COMPILER_XLA_CLIENT_LIB_QR_H_ + +#include "tensorflow/compiler/xla/client/xla_builder.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" + +namespace xla { + +// Computes the QR decompositions of a batch of matrices. That is, +// given a (batched) matrix a, computes an orthonormal matrix Q and an +// upper-triangular matrix R such that a = QR. +// `a` must be a (batched) matrix of size [..., m, n]. +// The algorithm implements a blocked QR decomposition; `block_size` is +// the block size to use. +// TODO(phawkins): handle the complex case. +struct QRDecompositionResult { + XlaOp q; + XlaOp r; +}; + +StatusOr QRDecomposition( + XlaOp a, bool full_matrices, int64 block_size = 128, + PrecisionConfig::Precision precision = PrecisionConfig::HIGHEST); + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_CLIENT_LIB_QR_H_ diff --git a/tensorflow/compiler/xla/client/lib/qr_test.cc b/tensorflow/compiler/xla/client/lib/qr_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..b27d364b62444d6d5fb1278b6e6461affc15b2e6 --- /dev/null +++ b/tensorflow/compiler/xla/client/lib/qr_test.cc @@ -0,0 +1,93 @@ +/* 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/client/lib/qr.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/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 { + +using QrTest = xla::ClientLibraryTestBase; + +XLA_TEST_F(QrTest, Simple) { + xla::XlaBuilder builder(TestName()); + + xla::Array2D a_vals({ + {4, 6, 8, 10}, + {6, 45, 54, 63}, + {8, 54, 146, 166}, + {10, 63, 166, 310}, + }); + + xla::XlaOp a; + auto a_data = CreateR2Parameter(a_vals, 0, "a", &builder, &a); + TF_ASSERT_OK_AND_ASSIGN( + auto result, + xla::QRDecomposition(a, /*full_matrices=*/true, /*block_size=*/2)); + + // Verifies that the decomposition composes back to the original matrix. + // + // This isn't a terribly demanding test, (e.g., we should verify that Q is + // orthonormal and R is upper-triangular) but it's awkward to write such tests + // without more linear algebra libraries. It's easier to test the numerics + // from Python, anyway, where we have access to numpy and scipy. + xla::BatchDot(result.q, result.r, xla::PrecisionConfig::HIGHEST); + + ComputeAndCompareR2(&builder, a_vals, {a_data.get()}, + xla::ErrorSpec(1e-4, 1e-4)); +} + +XLA_TEST_F(QrTest, SimpleBatched) { + xla::XlaBuilder builder(TestName()); + + xla::Array3D a_vals({ + { + {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, 456, 106}, + {12, 48, 106, 62}, + }, + }); + + xla::XlaOp a; + auto a_data = CreateR3Parameter(a_vals, 0, "a", &builder, &a); + TF_ASSERT_OK_AND_ASSIGN( + auto result, + xla::QRDecomposition(a, /*full_matrices=*/true, /*block_size=*/2)); + + xla::BatchDot(result.q, result.r, xla::PrecisionConfig::HIGHEST); + + ComputeAndCompareR3(&builder, a_vals, {a_data.get()}, + xla::ErrorSpec(1e-4, 1e-4)); +} + +} // namespace diff --git a/tensorflow/compiler/xla/client/lib/quantize.h b/tensorflow/compiler/xla/client/lib/quantize.h new file mode 100644 index 0000000000000000000000000000000000000000..26dbbd5b00bd1a29f4047c9a4294fcac7340cf6c --- /dev/null +++ b/tensorflow/compiler/xla/client/lib/quantize.h @@ -0,0 +1,186 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef TENSORFLOW_COMPILER_XLA_CLIENT_LIB_QUANTIZE_H_ +#define TENSORFLOW_COMPILER_XLA_CLIENT_LIB_QUANTIZE_H_ + +#include +#include +#include + +#include "tensorflow/compiler/xla/client/lib/constants.h" +#include "tensorflow/compiler/xla/client/xla_builder.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/bfloat16/bfloat16.h" + +namespace xla { + +constexpr int64 kBitsOfByte = 8; + +// Represents the range used for quantization +struct QuantizedRange { + QuantizedRange() = default; + QuantizedRange(float min_in, float max_in) : min(min_in), max(max_in) {} + + bool operator==(const QuantizedRange& rhs) const { + return this->min == rhs.min && this->max == rhs.max; + } + + bool operator!=(const QuantizedRange& rhs) const { return !(*this == rhs); } + + tensorflow::bfloat16 min = tensorflow::bfloat16(0.0f); + tensorflow::bfloat16 max = tensorflow::bfloat16(0.0f); +}; + +template +inline std::vector PackToUint32(absl::Span input) { + const int64 kElementsPerPack = sizeof(uint32) / sizeof(T); + const int64 input_size = input.size(); + const int64 output_size = CeilOfRatio(input_size, kElementsPerPack); + + std::vector output_vec; + constexpr int64 kShiftBits = sizeof(T) / sizeof(uint8) * kBitsOfByte; + + for (int64 i = 0; i < output_size; i++) { + uint32 result = 0; + for (int64 p = 0; p < kElementsPerPack; p++) { + int64 index = i * kElementsPerPack + p; + if (index < input_size) { + int64 total_shift_bits = kShiftBits * (kElementsPerPack - p - 1); + result |= (input[index] << total_shift_bits); + } + } + output_vec.push_back(result); + } + + return output_vec; +} + +// Dequantize the quantized input of packed uint32 to bfloat16. +// Only uint8 or uint16 is supported for the original unpacked input. +// Returns a tensor of shape [d0,..., dn * unpack_size] if +// input shape is [d0, ..., dn], where unpack_size = sizeof(unit32) / sizeof(T). +// If transpose_output is true, will return a tensor of shape +// [dn * unpack_size, dn-1, ..., d1, d0]. transpose_output is faster when +// input's rank higher than 1. The input needs to be transposed to use +// transpose_output feature. +template +inline XlaOp Dequantize(XlaOp input, const QuantizedRange& range, + absl::string_view mode_string = "MIN_COMBINED", + bool transpose_output = false) { + XlaBuilder* const builder = input.builder(); + return builder->ReportErrorOrReturn([&]() -> StatusOr { + float half_range = + !std::is_signed::value + ? 0.0f + : (static_cast(std::numeric_limits::max()) - + std::numeric_limits::min() + 1) / + 2.0f; + const int64 unpack_size = sizeof(uint32) / sizeof(T); + TF_ASSIGN_OR_RETURN(Shape shape, builder->GetShape(input)); + + auto element_type = shape.element_type(); + if (element_type != U32) { + return InvalidArgument( + "Only U32 is supported for input type of xla::Dequantize Op."); + } + + // Broadcast the input to [unpack_size, d0, ..., dn] if input size is + // [d0, ..., dn]. + auto broadcast_input = Broadcast(input, {unpack_size}); + + XlaOp iota_r1 = Iota(builder, U32, unpack_size); + // Highest significant bytes needs to shift more bytes than lower + // significant bytes. + XlaOp shift_bytes = + xla::ConstantR0(builder, unpack_size - 1) - iota_r1; + + const int bytes_of_type = sizeof(T) / sizeof(uint8); + std::vector shift_vec(unpack_size, kBitsOfByte * bytes_of_type); + XlaOp shift_bits = + shift_bytes * xla::ConstantR1(builder, shift_vec); + + // Make bit_mask for different data type T. + uint32 bit_mask = 0x00000000; + for (int i = 0; i < bytes_of_type; i++) { + bit_mask <<= kBitsOfByte; + bit_mask |= 0x000000ff; + } + + std::vector shift_transpose_dimensions(shape.dimensions_size()); + std::iota(shift_transpose_dimensions.begin(), + shift_transpose_dimensions.end(), 0); + shift_transpose_dimensions.insert(shift_transpose_dimensions.begin(), 1, + shape.dimensions_size()); + + // Shift the input by sizeof(T) bytes and apply bit_mask to unpack. + XlaOp shifted_input = ShiftRightLogical( + broadcast_input, Transpose(Broadcast(shift_bits, shape.dimensions()), + shift_transpose_dimensions)); + XlaOp unpack_input = + And(shifted_input, xla::ConstantR0(builder, bit_mask)); + + XlaOp result; + + if (mode_string == "MIN_COMBINED") { + const tensorflow::bfloat16 scale_factor = + (range.max - range.min) / + (static_cast(std::numeric_limits::max() - + std::numeric_limits::min())); + // result = bfloat16(input + half_range) * scale_factor + range.min + XlaOp unpack_input_bf16 = ConvertElementType(unpack_input, BF16); + XlaOp half_range_bf16 = xla::ConstantR0( + builder, static_cast(half_range)); + XlaOp sum = unpack_input_bf16 + half_range_bf16; + + result = + sum * xla::ConstantR0(builder, scale_factor) + + xla::ConstantR0(builder, range.min); + } else { + // TODO(wangtao): support other modes. + return InvalidArgument( + "Only MIN_COMBINED mode is supported in xla::Dequantize Op."); + } + + std::vector transpose_dimensions(shape.dimensions_size()); + std::iota(transpose_dimensions.begin(), transpose_dimensions.end(), 1); + std::reverse(transpose_dimensions.begin(), transpose_dimensions.end()); + transpose_dimensions.insert(transpose_dimensions.begin() + 1, 1, 0); + + // Transpose the result to be [dn, unpack_size, dn-1, ..., d1, d0]. + XlaOp transposed_result = Transpose(result, transpose_dimensions); + + // Reshape to be [dn * unpack_size, dn-1, ..., d1, d0]. + XlaOp reshaped_result = Collapse(transposed_result, {0, 1}); + + // Return the transpose result if transpose_output is true. + if (transpose_output) { + return reshaped_result; + } + + // Transpose the result to be [d0, d1, ..., dn-1, dn * unpack_size]. + std::vector result_dimensions(shape.dimensions_size()); + std::iota(result_dimensions.begin(), result_dimensions.end(), 0); + std::reverse(result_dimensions.begin(), result_dimensions.end()); + + return Transpose(reshaped_result, result_dimensions); + }); +} + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_CLIENT_LIB_QUANTIZE_H_ diff --git a/tensorflow/compiler/xla/client/lib/quantize_test.cc b/tensorflow/compiler/xla/client/lib/quantize_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..be3603d9e11670913c21a834d2216a999306d582 --- /dev/null +++ b/tensorflow/compiler/xla/client/lib/quantize_test.cc @@ -0,0 +1,337 @@ +/* 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/client/lib/quantize.h" + +#include + +#include "tensorflow/compiler/xla/client/xla_builder.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/util.h" + +namespace xla { +namespace { + +using bfloat16 = tensorflow::bfloat16; + +template +std::vector GenerateInput() { + std::vector input; + + for (int64 i = std::numeric_limits::min(); + i < std::numeric_limits::max(); ++i) { + input.push_back(static_cast(i)); + } + + return input; +} + +template +Array2D GenerateLargeSizeInput(int num_columns, int num_rows) { + Array2D input(num_columns, num_rows); + + input.FillRandom(6, 128); + + return input; +} + +template +Array2D PackLargeInput(Array2D &input) { + const int64 size_per_pack = sizeof(uint32) / sizeof(NativeT); + int64 width = input.width(); + + int64 padded_output_width = CeilOfRatio(width, size_per_pack); + + Array2D pack_input(input.height(), padded_output_width); + + for (int h = 0; h < input.height(); h++) { + std::vector input_row; + for (int w = 0; w < width; w++) { + input_row.push_back(input({h, w})); + } + + auto pack_input_vec = PackToUint32(input_row); + + for (int w = 0; w < padded_output_width; w++) { + pack_input(h, w) = pack_input_vec[w]; + } + } + + return pack_input; +} + +template +Array2D GenerateLargeSizeMinCombinedOutput( + Array2D &input, const QuantizedRange &range, + bool transpose_output = false) { + const int64 size_per_pack = sizeof(uint32) / sizeof(NativeT); + int64 width = input.width(); + + int64 padded_output_width = CeilOfRatio(width, size_per_pack) * size_per_pack; + + int64 output_height; + int64 output_width; + + if (transpose_output) { + output_height = padded_output_width; + output_width = input.height(); + } else { + output_height = input.height(); + output_width = padded_output_width; + } + + Array2D output(output_height, output_width, bfloat16(0.0)); + + float half_range = + !std::is_signed::value + ? 0.0f + : (static_cast(std::numeric_limits::max() - + std::numeric_limits::min() + 1)) / + 2.0f; + const bfloat16 scale_factor = + (range.max - range.min) / + (static_cast(std::numeric_limits::max() - + std::numeric_limits::min())); + + for (int h = 0; h < input.height(); h++) { + std::vector input_row; + for (int w = 0; w < width; w++) { + bfloat16 result = + static_cast(input(h, w) + half_range) * scale_factor + + range.min; + if (transpose_output) { + output(w, h) = result; + } else { + output(h, w) = result; + } + } + } + + return output; +} + +template +std::vector GenerateMinCombinedOutput(const QuantizedRange &range) { + float half_range = + !std::is_signed::value + ? 0.0f + : (static_cast(std::numeric_limits::max() - + std::numeric_limits::min() + 1)) / + 2.0f; + const bfloat16 scale_factor = + (range.max - range.min) / + (static_cast(std::numeric_limits::max() - + std::numeric_limits::min())); + std::vector output; + for (int64 i = std::numeric_limits::min(); + i < std::numeric_limits::max(); ++i) { + bfloat16 result = + static_cast(i + half_range) * scale_factor + range.min; + output.push_back(result); + } + + const int64 pack_size = sizeof(uint32) / sizeof(NativeT); + const int64 output_size = output.size(); + + int64 num_tailing_zeros = + CeilOfRatio(output_size, pack_size) * pack_size - output_size; + + output.insert(output.end(), num_tailing_zeros, bfloat16(0.0)); + return output; +} + +// TODO(wangtao): add a test to make sure this op is the inverse of the existing +// TF quantize op defined in: third_party/tensorflow/core/kernels/quantize_op.cc + +using DequantizeTest = ClientLibraryTestBase; + +TEST(PackTest, PackUint8ToUint32) { + std::vector input = {0xAB, 0x0B, 0x00, 0xF0, 0x01}; + auto output = PackToUint32(input); + EXPECT_THAT(output, ::testing::ElementsAre(0xAB0B00F0, 0x01000000)); +} + +TEST(PackTest, PackInt8ToUint32) { + std::vector input = {static_cast(0x81), 0x0B, 0x00, 0x20, + 0x01}; + auto output = PackToUint32(input); + EXPECT_THAT(output, ::testing::ElementsAre(0x810B0020, 0x01000000)); +} + +TEST(PackTest, PackUint8ToUint32PerfectSize) { + std::vector input = {3, 2, 1, 0}; + auto output = PackToUint32(input); + EXPECT_THAT(output, ::testing::ElementsAre(0x03020100)); +} + +XLA_TEST_F(DequantizeTest, MinCombinedUint16R1) { + XlaBuilder builder(TestName()); + auto input = GenerateInput(); + auto x = ConstantR1(&builder, PackToUint32(input)); + QuantizedRange range(0, 255.0f); + xla::Dequantize(x, range, "MIN_COMBINED"); + auto expected = GenerateMinCombinedOutput(range); + ComputeAndCompareR1(&builder, expected, {}); +} + +XLA_TEST_F(DequantizeTest, MinCombinedUint8R1) { + XlaBuilder builder(TestName()); + auto input = GenerateInput(); + auto x = ConstantR1(&builder, PackToUint32(input)); + QuantizedRange range(0, 127.0f); + xla::Dequantize(x, range, "MIN_COMBINED"); + auto expected = GenerateMinCombinedOutput(range); + ComputeAndCompareR1(&builder, expected, {}); +} + +XLA_TEST_F(DequantizeTest, MinCombinedUint8R2) { + XlaBuilder builder(TestName()); + std::vector> input = { + {0, 1, 2, 3}, + {4, 5, 6, 7}, + {8, 9, 10, 11}, + {12, 13, 16, 15}, + }; + auto x = ConstantR2(&builder, {{PackToUint32(input[0])[0]}, + {PackToUint32(input[1])[0]}, + {PackToUint32(input[2])[0]}, + {PackToUint32(input[3])[0]}}); + QuantizedRange range(0, 255.0f); + xla::Dequantize(x, range, "MIN_COMBINED"); + const Array2D expected = { + {bfloat16(0.0), bfloat16(1.0), bfloat16(2.0), bfloat16(3.0)}, + {bfloat16(4.0), bfloat16(5.0), bfloat16(6.0), bfloat16(7.0)}, + {bfloat16(8.0), bfloat16(9.0), bfloat16(10.0), bfloat16(11.0)}, + {bfloat16(12.0), bfloat16(13.0), bfloat16(16.0), bfloat16(15.0)}, + }; + ComputeAndCompareR2(&builder, expected, {}); +} + +XLA_TEST_F(DequantizeTest, MinCombinedUint8R2TransposeOutput) { + XlaBuilder builder(TestName()); + std::vector> input = { + {0, 1, 2, 3}, + {4, 5, 6, 7}, + {8, 9, 10, 11}, + {12, 13, 16, 15}, + }; + auto x = ConstantR2(&builder, {{PackToUint32(input[0])[0]}, + {PackToUint32(input[1])[0]}, + {PackToUint32(input[2])[0]}, + {PackToUint32(input[3])[0]}}); + QuantizedRange range(0, 255.0f); + xla::Dequantize(x, range, "MIN_COMBINED", /*transpose_output=*/true); + const Array2D expected = { + {bfloat16(0.0), bfloat16(4.0), bfloat16(8.0), bfloat16(12.0)}, + {bfloat16(1.0), bfloat16(5.0), bfloat16(9.0), bfloat16(13.0)}, + {bfloat16(2.0), bfloat16(6.0), bfloat16(10.0), bfloat16(16.0)}, + {bfloat16(3.0), bfloat16(7.0), bfloat16(11.0), bfloat16(15.0)}, + }; + ComputeAndCompareR2(&builder, expected, {}); +} + +XLA_TEST_F(DequantizeTest, MinCombinedUint8R2TailingZero) { + XlaBuilder builder(TestName()); + std::vector> input = { + {0, 1, 2, 3, 16}, + {4, 5, 6, 7, 17}, + {8, 9, 10, 11, 18}, + {12, 13, 16, 15, 19}, + }; + auto x = ConstantR2( + &builder, + {{PackToUint32(input[0])[0], PackToUint32(input[0])[1]}, + {PackToUint32(input[1])[0], PackToUint32(input[1])[1]}, + {PackToUint32(input[2])[0], PackToUint32(input[2])[1]}, + {PackToUint32(input[3])[0], PackToUint32(input[3])[1]}}); + QuantizedRange range(0, 255.0f); + xla::Dequantize(x, range, "MIN_COMBINED"); + + const Array2D expected = { + {bfloat16(0.0), bfloat16(1.0), bfloat16(2.0), bfloat16(3.0), + bfloat16(16.0), bfloat16(0.0), bfloat16(0.0), bfloat16(0.0)}, + {bfloat16(4.0), bfloat16(5.0), bfloat16(6.0), bfloat16(7.0), + bfloat16(17.0), bfloat16(0.0), bfloat16(0.0), bfloat16(0.0)}, + {bfloat16(8.0), bfloat16(9.0), bfloat16(10.0), bfloat16(11.0), + bfloat16(18.0), bfloat16(0.0), bfloat16(0.0), bfloat16(0.0)}, + {bfloat16(12.0), bfloat16(13.0), bfloat16(16.0), bfloat16(15.0), + bfloat16(19.0), bfloat16(0.0), bfloat16(0.0), bfloat16(0.0)}, + }; + ComputeAndCompareR2(&builder, expected, {}); +} + +XLA_TEST_F(DequantizeTest, MinCombinedUint8R2TailingZeroTransposeOutput) { + XlaBuilder builder(TestName()); + std::vector> input = { + {0, 1, 2, 3, 16}, + {4, 5, 6, 7, 17}, + {8, 9, 10, 11, 18}, + {12, 13, 16, 15, 19}, + }; + auto x = ConstantR2( + &builder, + {{PackToUint32(input[0])[0], PackToUint32(input[0])[1]}, + {PackToUint32(input[1])[0], PackToUint32(input[1])[1]}, + {PackToUint32(input[2])[0], PackToUint32(input[2])[1]}, + {PackToUint32(input[3])[0], PackToUint32(input[3])[1]}}); + QuantizedRange range(0, 255.0f); + xla::Dequantize(x, range, "MIN_COMBINED", /*transpose_output=*/true); + + const Array2D expected = { + {bfloat16(0.0), bfloat16(4.0), bfloat16(8.0), bfloat16(12.0)}, + {bfloat16(1.0), bfloat16(5.0), bfloat16(9.0), bfloat16(13.0)}, + {bfloat16(2.0), bfloat16(6.0), bfloat16(10.0), bfloat16(16.0)}, + {bfloat16(3.0), bfloat16(7.0), bfloat16(11.0), bfloat16(15.0)}, + {bfloat16(16.0), bfloat16(17.0), bfloat16(18.0), bfloat16(19.0)}, + {bfloat16(0.0), bfloat16(0.0), bfloat16(0.0), bfloat16(0.0)}, + {bfloat16(0.0), bfloat16(0.0), bfloat16(0.0), bfloat16(0.0)}, + {bfloat16(0.0), bfloat16(0.0), bfloat16(0.0), bfloat16(0.0)}, + }; + ComputeAndCompareR2(&builder, expected, {}); +} + +XLA_TEST_F(DequantizeTest, MinCombinedUint8LargeSizeTest) { + XlaBuilder builder(TestName()); + Array2D input = GenerateLargeSizeInput(500, 3547); + Array2D input_packed = PackLargeInput(input); + + auto x = ConstantR2FromArray2D(&builder, input_packed); + QuantizedRange range(0, 255.0f); + xla::Dequantize(x, range, "MIN_COMBINED"); + + const Array2D expected = + GenerateLargeSizeMinCombinedOutput(input, range); + ComputeAndCompareR2(&builder, expected, {}); +} + +XLA_TEST_F(DequantizeTest, MinCombinedUint8LargeSizeTestTransposeOutput) { + XlaBuilder builder(TestName()); + Array2D input = GenerateLargeSizeInput(500, 3547); + Array2D input_packed = PackLargeInput(input); + + auto x = ConstantR2FromArray2D(&builder, input_packed); + QuantizedRange range(0, 255.0f); + xla::Dequantize(x, range, "MIN_COMBINED", /*transpose_output=*/true); + + const Array2D expected = GenerateLargeSizeMinCombinedOutput( + input, range, /*transpose_output=*/true); + ComputeAndCompareR2(&builder, expected, {}); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/client/lib/slicing.cc b/tensorflow/compiler/xla/client/lib/slicing.cc new file mode 100644 index 0000000000000000000000000000000000000000..77145ba7d4c72435450d3e33d57b2507eb84d2fc --- /dev/null +++ b/tensorflow/compiler/xla/client/lib/slicing.cc @@ -0,0 +1,137 @@ +/* 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/client/lib/slicing.h" +#include "tensorflow/compiler/xla/client/xla_builder.h" + +namespace xla { + +XlaOp SliceInMinorDims(XlaOp x, absl::Span start, + absl::Span end) { + XlaBuilder* builder = x.builder(); + return builder->ReportErrorOrReturn([&]() -> StatusOr { + TF_RET_CHECK(start.size() == end.size()); + int64 n_minor_dims = start.size(); + + TF_ASSIGN_OR_RETURN(Shape shape, builder->GetShape(x)); + + const int64 n_dims = shape.rank(); + TF_RET_CHECK(n_minor_dims <= n_dims); + auto major_dims = AsInt64Slice(shape.dimensions()) + .subspan( + /*pos=*/0, + /*len=*/n_dims - n_minor_dims); + + // Prepends 0s in the major dim + std::vector padded_start(n_dims, 0); + std::copy(start.begin(), start.end(), + padded_start.begin() + major_dims.size()); + + // Prepends the shape of the major dims. + std::vector padded_end(n_dims); + std::copy(major_dims.begin(), major_dims.end(), padded_end.begin()); + std::copy(end.begin(), end.end(), padded_end.begin() + major_dims.size()); + + std::vector strides(n_dims, 1); + return Slice(x, padded_start, padded_end, strides); + }); +} + +XlaOp UpdateSlice(XlaOp x, XlaOp update, absl::Span start) { + XlaBuilder* builder = x.builder(); + return builder->ReportErrorOrReturn([&]() -> StatusOr { + TF_ASSIGN_OR_RETURN(Shape shape, builder->GetShape(x)); + const int64 n_dims = shape.rank(); + 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); + }); +} + +XlaOp UpdateSliceInMinorDims(XlaOp x, XlaOp update, + absl::Span start) { + XlaBuilder* builder = x.builder(); + return builder->ReportErrorOrReturn([&]() -> StatusOr { + TF_ASSIGN_OR_RETURN(Shape shape, builder->GetShape(x)); + const int64 n_dims = shape.rank(); + const int64 n_minor_dims = start.size(); + TF_RET_CHECK(n_minor_dims <= n_dims); + std::vector padded_start(n_dims, 0); + std::copy(start.begin(), start.end(), + padded_start.begin() + (n_dims - n_minor_dims)); + return UpdateSlice(x, update, padded_start); + }); +} + +namespace { + +std::vector ConcatVectors(absl::Span xs, + absl::Span ys) { + std::vector output(xs.size() + ys.size()); + std::copy(xs.begin(), xs.end(), output.begin()); + std::copy(ys.begin(), ys.end(), output.begin() + xs.size()); + return output; +} + +StatusOr> PrependZerosInMajorDims( + XlaOp x, absl::Span starts) { + XlaBuilder* builder = x.builder(); + 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 + +XlaOp DynamicSliceInMinorDims(XlaOp x, absl::Span starts, + absl::Span sizes) { + XlaBuilder* builder = x.builder(); + return builder->ReportErrorOrReturn([&]() -> StatusOr { + TF_ASSIGN_OR_RETURN(Shape shape, builder->GetShape(x)); + const int64 n_dims = shape.rank(); + int64 n_minor_dims = starts.size(); + TF_RET_CHECK(n_minor_dims == sizes.size()); + TF_RET_CHECK(n_minor_dims <= n_dims); + auto major_dims = AsInt64Slice(shape.dimensions()) + .subspan( + /*pos=*/0, + /*len=*/n_dims - sizes.size()); + TF_ASSIGN_OR_RETURN(auto padded_starts, PrependZerosInMajorDims(x, starts)); + auto padded_sizes = ConcatVectors(major_dims, sizes); + return DynamicSlice(x, padded_starts, padded_sizes); + }); +} + +XlaOp DynamicUpdateSliceInMinorDims(XlaOp x, XlaOp update, + absl::Span 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/slicing.h b/tensorflow/compiler/xla/client/lib/slicing.h new file mode 100644 index 0000000000000000000000000000000000000000..6c482a38b5489c9fb17c3dca9ee3d2a1b8fd1890 --- /dev/null +++ b/tensorflow/compiler/xla/client/lib/slicing.h @@ -0,0 +1,48 @@ +/* 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 "absl/types/span.h" +#include "tensorflow/compiler/xla/client/xla_builder.h" +#include "tensorflow/compiler/xla/types.h" + +#ifndef TENSORFLOW_COMPILER_XLA_CLIENT_LIB_SLICING_H_ +#define TENSORFLOW_COMPILER_XLA_CLIENT_LIB_SLICING_H_ + +namespace xla { + +// Updates a slice of 'x', i.e., +// x[start[0], ..., start[n]] = update +XlaOp UpdateSlice(XlaOp x, XlaOp update, absl::Span start); + +// Performs a slice in the minor dimensions of a tensor. +// x[..., start[0]:end[0], ..., start[n]:end[n]] +XlaOp SliceInMinorDims(XlaOp x, absl::Span start, + absl::Span end); + +// Updates a slice of 'x', where 'start' contains a list of minor dimensions: +// x[..., start[0]:..., ..., start[n]:...] = update +XlaOp UpdateSliceInMinorDims(XlaOp x, XlaOp update, + absl::Span start); + +// Performs a dynamic slice in the minor dimensions of a tensor. +XlaOp DynamicSliceInMinorDims(XlaOp x, absl::Span starts, + absl::Span sizes); + +XlaOp DynamicUpdateSliceInMinorDims(XlaOp x, XlaOp update, + absl::Span starts); + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_CLIENT_LIB_SLICING_H_ diff --git a/tensorflow/compiler/xla/client/lib/slicing_test.cc b/tensorflow/compiler/xla/client/lib/slicing_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..8d362119e01006555db0f82d02626175936e1d05 --- /dev/null +++ b/tensorflow/compiler/xla/client/lib/slicing_test.cc @@ -0,0 +1,106 @@ +/* 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/client/lib/slicing.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 SlicingTest = xla::ClientLibraryTestBase; + +xla::Array2D BValsRight() { + return {{1, 2, 3, 4}, {5, 6, 7, 8}, {9, 10, 11, 12}}; +} + +xla::Array2D BValsLeft() { + return {{1, 2, 3}, {4, 5, 6}, {7, 8, 9}, {10, 11, 12}}; +} + +xla::Array2D AValsFull() { + return {{2, 0, 1, 2}, {3, 6, 0, 1}, {4, 7, 9, 0}, {5, 8, 10, 11}}; +} + +xla::Array3D BatchedAValsFull() { + return {{ + {2, 0, 1, 2}, + {3, 6, 0, 1}, + {4, 7, 9, 0}, + {5, 8, 10, 11}, + }, + { + {16, 24, 8, 12}, + {24, 61, 82, 48}, + {8, 82, 456, 106}, + {12, 48, 106, 62}, + }}; +} + +XLA_TEST_F(SlicingTest, Simple2dLookup) { + xla::XlaBuilder builder(TestName()); + + xla::XlaOp a, x, y; + auto a_data = CreateR2Parameter(BValsRight(), 0, "a", &builder, &a); + auto x_data = CreateR0Parameter(2, 1, "x", &builder, &x); + auto y_data = CreateR0Parameter(1, 2, "y", &builder, &y); + DynamicSliceInMinorDims(a, {x, y}, {1, 1}); + + ComputeAndCompareR2(&builder, {{10}}, + {a_data.get(), x_data.get(), y_data.get()}, + xla::ErrorSpec(1e-2, 1e-2)); +} + +XLA_TEST_F(SlicingTest, Simple3dLookup) { + xla::XlaBuilder builder(TestName()); + + xla::XlaOp a, index; + auto a_data = + CreateR3Parameter(BatchedAValsFull(), 0, "a", &builder, &a); + auto index_data = CreateR0Parameter(1, 1, "index", &builder, &index); + + DynamicSliceInMinorDims(a, {index, xla::ConstantR0(&builder, 0)}, + {1, 4}); + + ComputeAndCompareR3(&builder, {{{3, 6, 0, 1}}, {{24, 61, 82, 48}}}, + {a_data.get(), index_data.get()}); +} + +XLA_TEST_F(SlicingTest, SimpleSliceUpdate) { + xla::XlaBuilder builder(TestName()); + + xla::XlaOp a, b, x, y; + auto a_data = CreateR2Parameter(AValsFull(), 0, "a", &builder, &a); + auto b_data = CreateR2Parameter({{9, 1, -10}}, 1, "b", &builder, &b); + auto x_data = CreateR0Parameter(2, 2, "x", &builder, &x); + auto y_data = CreateR0Parameter(1, 3, "y", &builder, &y); + + DynamicUpdateSliceInMinorDims(a, b, {x, y}); + + xla::Array2D expected( + {{{2, 0, 1, 2}, {3, 6, 0, 1}, {4, 9, 1, -10}, {5, 8, 10, 11}}}); + + ComputeAndCompareR2( + &builder, expected, + {a_data.get(), b_data.get(), x_data.get(), y_data.get()}); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/client/lib/sorting.cc b/tensorflow/compiler/xla/client/lib/sorting.cc index 0475fd9c94f6e390b5169cfe2cbba8eae28ddc18..e8553a08bb014e790822a14e128686b60b8d6b7c 100644 --- a/tensorflow/compiler/xla/client/lib/sorting.cc +++ b/tensorflow/compiler/xla/client/lib/sorting.cc @@ -14,7 +14,9 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/compiler/xla/client/lib/sorting.h" -#include "tensorflow/compiler/xla/client/lib/numeric.h" +#include "tensorflow/compiler/xla/client/xla_builder.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/util.h" namespace xla { @@ -23,13 +25,12 @@ XlaOp TopK(XlaOp input, int64 k) { return builder->ReportErrorOrReturn([&]() -> StatusOr { TF_ASSIGN_OR_RETURN(Shape input_shape, builder->GetShape(input)); int last_dim = input_shape.dimensions_size() - 1; - int last_dim_size = input_shape.dimensions(last_dim); - XlaOp iota_s32 = Iota(builder, S32, last_dim_size); + Shape iota_shape = + ShapeUtil::MakeShape(S32, AsInt64Slice(input_shape.dimensions())); + XlaOp iota_s32 = Iota(builder, iota_shape, last_dim); auto input_dims = input_shape.dimensions(); - std::vector broadcast_dims(input_dims.begin(), input_dims.end() - 1); - XlaOp broadcast_s32 = Broadcast(iota_s32, broadcast_dims); - XlaOp sort_result = Sort(Neg(input), {broadcast_s32}); + XlaOp sort_result = Sort(Neg(input), {iota_s32}); 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 fef98c9923096e21a755c6d730de2c7c10852b2d..0fbd138aca1e86f219d0459086fc09d20844f135 100644 --- a/tensorflow/compiler/xla/client/lib/sorting_test.cc +++ b/tensorflow/compiler/xla/client/lib/sorting_test.cc @@ -14,6 +14,9 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/compiler/xla/client/lib/sorting.h" + +#include + #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/compiler/xla/test.h" #include "tensorflow/compiler/xla/tests/client_library_test_base.h" @@ -41,6 +44,28 @@ XLA_TEST_F(SortingTest, TopK3From8Indices) { ComputeAndCompareR1(&builder, {0, 1, 2}, {}); } +// TODO(b/119930279): enable this test. +XLA_TEST_F(SortingTest, DISABLED_TopKFullSortMinInt) { + XlaBuilder builder(TestName()); + auto x_rev = ConstantR1(&builder, {std::numeric_limits::min(), + std::numeric_limits::min() + 1, + std::numeric_limits::max()}); + xla::GetTupleElement(xla::TopK(x_rev, 3), 1); + ComputeAndCompareR1(&builder, {2, 1, 0}, {}); +} + +XLA_TEST_F(SortingTest, NOT_TopKFullSortMinInt) { + XlaBuilder builder(TestName()); + auto x_rev = ConstantR1(&builder, {std::numeric_limits::min(), + std::numeric_limits::min() + 1, + std::numeric_limits::max()}); + xla::GetTupleElement(xla::TopK(x_rev, 3), 1); + // TopK currently negates the keys, which doesn't work correctly for + // std::numeric_limits::min(). Therefore, it will sort this key to the + // front instead of to the back. + ComputeAndCompareR1(&builder, {0, 2, 1}, {}); +} + XLA_TEST_F(SortingTest, TopKFullSort) { XlaBuilder builder(TestName()); const int kSize = 16; @@ -52,9 +77,17 @@ 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) { + XlaBuilder builder(TestName()); + XlaOp a; + auto a_data = CreateR1Parameter({1, 1, 2, 2, 1}, 0, "a", &builder, &a); + xla::GetTupleElement(xla::TopK(a, 5), 1); + ComputeAndCompareR1(&builder, {2, 3, 0, 1, 4}, {a_data.get()}); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/client/lib/testing.cc b/tensorflow/compiler/xla/client/lib/testing.cc index a44681f586278bf03f3fb2b8c812936cbf3ad47b..9f520bcdadfabc8ca9f9ee82b20804fd2c50d1db 100644 --- a/tensorflow/compiler/xla/client/lib/testing.cc +++ b/tensorflow/compiler/xla/client/lib/testing.cc @@ -34,7 +34,7 @@ namespace { // specified shape. In case of a (nested) tuple shape this is the total byte // size of all sub-shapes within the tuple. int64 DataSizeOfShape(const Shape& shape) { - if (ShapeUtil::IsArray(shape)) { + if (shape.IsArray()) { return ShapeUtil::ByteSizeOf(shape); } @@ -47,7 +47,7 @@ int64 DataSizeOfShape(const Shape& shape) { // Creates a XlaOp for an op what generates fake data with the given shape. XlaOp BuildFakeDataOpOnDevice(const Shape& shape, XlaBuilder* builder) { - if (ShapeUtil::IsArray(shape)) { + if (shape.IsArray()) { return Broadcast( ConstantLiteral(builder, LiteralUtil::One(shape.element_type())), AsInt64Slice(shape.dimensions())); @@ -59,22 +59,25 @@ XlaOp BuildFakeDataOpOnDevice(const Shape& shape, XlaBuilder* builder) { return Tuple(builder, parts); } -std::unique_ptr MakeFakeDataViaDeviceOrDie(const Shape& shape, - Client* client) { +std::unique_ptr MakeFakeDataViaDeviceOrDie( + const Shape& shape, Client* client, DebugOptions* debug_opts) { XlaBuilder b(absl::StrCat("make_fake_", ShapeUtil::HumanString(shape))); BuildFakeDataOpOnDevice(shape, &b); XlaComputation computation = b.Build().ConsumeValueOrDie(); auto execution_options = CreateDefaultExecutionOptions(); - *execution_options.mutable_shape_with_output_layout() = shape; + *execution_options.mutable_shape_with_output_layout() = shape.ToProto(); + if (debug_opts) { + *execution_options.mutable_debug_options() = *debug_opts; + } return client->Execute(computation, /*arguments=*/{}, &execution_options) .ConsumeValueOrDie(); } } // namespace -std::unique_ptr MakeFakeDataOrDie(const Shape& shape, - Client* client) { +std::unique_ptr MakeFakeDataOrDie( + const Shape& shape, Client* client, DebugOptions* debug_opts /*=nullptr*/) { if (DataSizeOfShape(shape) < (1LL << 20)) { StatusOr literal_status = MakeFakeLiteral(shape); if (!literal_status.ok()) { @@ -82,24 +85,25 @@ std::unique_ptr MakeFakeDataOrDie(const Shape& shape, // an on-device computation. CHECK_EQ(literal_status.status().code(), tensorflow::error::UNIMPLEMENTED); - return MakeFakeDataViaDeviceOrDie(shape, client); + return MakeFakeDataViaDeviceOrDie(shape, client, debug_opts); } return client->TransferToServer(literal_status.ValueOrDie()).ValueOrDie(); } // If the data is large, generate it on-device. - return MakeFakeDataViaDeviceOrDie(shape, client); + return MakeFakeDataViaDeviceOrDie(shape, client, debug_opts); } std::vector> MakeFakeArgumentsOrDie( - const XlaComputation& computation, Client* client) { + const XlaComputation& computation, Client* client, + DebugOptions* debug_opts /*=nullptr*/) { CHECK(computation.proto().has_host_program_shape()) << "Computation should have progran shape."; auto program_shape = computation.proto().host_program_shape(); std::vector> results; - for (const Shape& shape : program_shape.parameters()) { - results.push_back(MakeFakeDataOrDie(shape, client)); + for (const ShapeProto& shape : program_shape.parameters()) { + results.push_back(MakeFakeDataOrDie(Shape(shape), client, debug_opts)); } return results; } diff --git a/tensorflow/compiler/xla/client/lib/testing.h b/tensorflow/compiler/xla/client/lib/testing.h index 03695ce2a339735e3e49522f4fe1bbf2d83a3834..428fa3e93d1b46983aae60176e7c2242d2552fdb 100644 --- a/tensorflow/compiler/xla/client/lib/testing.h +++ b/tensorflow/compiler/xla/client/lib/testing.h @@ -29,14 +29,19 @@ namespace xla { // Generates fake data of the given shape on the device or dies. The fake data // is created by performing a computation on the device rather than transferring // data from the host to the device. -std::unique_ptr MakeFakeDataOrDie(const Shape& shape, - Client* client); +// +// The optional DebugOptions are used when generating fake data on the device. +std::unique_ptr MakeFakeDataOrDie( + const Shape& shape, Client* client, DebugOptions* debug_opts = nullptr); // Returns vector of GlobalData handles of fake data (created using // MakeFakeDataOrDie) that are correctly shaped arguments for the given // xla computation. +// +// The optional DebugOptions are used when generating fake data on the device. std::vector> MakeFakeArgumentsOrDie( - const XlaComputation& computation, Client* client); + const XlaComputation& computation, Client* client, + DebugOptions* debug_opts = nullptr); } // namespace xla diff --git a/tensorflow/compiler/xla/client/lib/triangular_solve.cc b/tensorflow/compiler/xla/client/lib/triangular_solve.cc new file mode 100644 index 0000000000000000000000000000000000000000..ba7fde118fde990fbb4aa9a34dd0f0e67ff5a93b --- /dev/null +++ b/tensorflow/compiler/xla/client/lib/triangular_solve.cc @@ -0,0 +1,430 @@ +/* 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/client/lib/triangular_solve.h" + +#include +#include + +#include "tensorflow/compiler/xla/client/lib/constants.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/client/xla_computation.h" +#include "tensorflow/compiler/xla/literal.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/util.h" +#include "tensorflow/core/lib/math/math_util.h" + +namespace xla { + +// Get the diagonal blocks of the coefficient matrix +XlaOp DiagonalBlocks(XlaOp a, int64 block_size) { + XlaBuilder* builder = a.builder(); + return builder->ReportErrorOrReturn([&]() -> StatusOr { + TF_ASSIGN_OR_RETURN(Shape shape, builder->GetShape(a)); + int ndims = shape.rank(); + int64 n = ShapeUtil::GetDimension(shape, -1); + int64 num_blocks = n / block_size; + + XlaOp diag_blocks; + + // If the coefficient matrix is exactly the block size, we just add a + // singleton dimension i.e. [..., n, n] -> [..., 1, n, n] + if (n == block_size) { + std::vector permutation(ndims); + std::iota(permutation.begin(), permutation.end(), 1); + permutation.insert(permutation.end() - 2, 0); + return Transpose(Broadcast(a, /*broadcast_sizes=*/{1}), permutation); + } + + // We can grab entire blocks using gather + if (n > block_size) { + // Construct the starting indices of the diagonal blocks + auto start_indices = + Transpose(Broadcast(Mul(Iota(builder, S32, num_blocks), + ConstantR0(builder, block_size)), + /*broadcast_sizes=*/{2}), + /*permutation=*/{1, 0}); + + PaddingConfig padding_config = + MakeEdgePaddingConfig({{0, 0}, {ndims - 2, 0}}); + start_indices = + Pad(start_indices, ConstantR0(builder, 0), padding_config); + + // Gather the diagonal blocks + std::vector slice_sizes(ndims); + GatherDimensionNumbers dim_numbers; + for (int i = 0; i < ndims - 2; ++i) { + dim_numbers.add_offset_dims(i); + dim_numbers.add_start_index_map(i); + slice_sizes[i] = ShapeUtil::GetDimension(shape, i); + } + slice_sizes[ndims - 2] = slice_sizes[ndims - 1] = block_size; + dim_numbers.add_offset_dims(ndims - 1); + dim_numbers.add_offset_dims(ndims); + dim_numbers.add_start_index_map(ndims - 2); + dim_numbers.add_start_index_map(ndims - 1); + dim_numbers.set_index_vector_dim(1); + diag_blocks = Gather(a, start_indices, dim_numbers, slice_sizes); + } + + // The last block might be smaller than the block size, + // so we will need to pad it + if (n % block_size != 0) { + // Pad with zeros + auto last_blocks = + SliceInMinorDims(a, {n - n % block_size, n - n % block_size}, {n, n}); + PaddingConfig config = MakeNoPaddingConfig(ndims); + int64 padding = block_size - n % block_size; + config.mutable_dimensions(ndims - 1)->set_edge_padding_high(padding); + config.mutable_dimensions(ndims - 2)->set_edge_padding_high(padding); + last_blocks = + Pad(last_blocks, Zero(builder, shape.element_type()), config); + + // Add a singleton dimension + // i.e. [..., block_size, block_size] -> [..., 1, block_size, block_size] + TF_ASSIGN_OR_RETURN(Shape blocks_shape, builder->GetShape(last_blocks)); + auto shape_dims = AsInt64Slice(blocks_shape.dimensions()); + auto last_blocks_dims = std::vector(ndims); + std::copy(shape_dims.begin(), shape_dims.end(), last_blocks_dims.begin()); + last_blocks_dims.insert(last_blocks_dims.end() - 2, 1); + last_blocks = Reshape(last_blocks, last_blocks_dims); + + // Concatenate with the other blocks if necessary + if (n > block_size) { + diag_blocks = + ConcatInDim(builder, {diag_blocks, last_blocks}, ndims - 2); + } else { + diag_blocks = last_blocks; + } + } + + return diag_blocks; + }); +} + +XlaOp InvertDiagonalBlocks(XlaOp diag_blocks, bool lower, bool transpose_a, + bool conjugate_a, + PrecisionConfig::Precision precision) { + XlaBuilder* builder = diag_blocks.builder(); + return builder->ReportErrorOrReturn([&]() -> StatusOr { + // Input is a batch of square lower triangular square matrices. Its shape is + // (..., size, size). We resize this to (num_blocks, size, size). + TF_ASSIGN_OR_RETURN(Shape shape, builder->GetShape(diag_blocks)); + int64 block_size = ShapeUtil::GetDimension(shape, -1); + int64 num_blocks = ShapeUtil::ElementsIn(shape) / + tensorflow::MathUtil::IPow(block_size, 2); + diag_blocks = Reshape(diag_blocks, {num_blocks, block_size, block_size}); + + // The input must be triangular because we rely on that when doing + // multiplications later on + diag_blocks = Triangle(diag_blocks, /*lower=*/lower); + + // Rescale blocks to be unit triangular, but avoid dividing by + // zero (which can happen if the last block was padded) otherwise it will + // introduce nans which will propagate + auto diags = GetMatrixDiagonal(diag_blocks); + 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}); + + // We can now use the fact that for an upper triangular matrix + // [[L11, 0], [L21, L22]], given the inverses L11' and L22', we have + // L22' = -L22' * L21 * L11'. In our case, L21 is a vector and our blocks + // have been rescaled to be unit triangular, so L22 = L22' = 1. + + // Initialize the output matrix with -1s on the diagonal. We use -1 instead + // of 1 because we cannot do matrix-vector multiplies with variable shapes + // inside of a loop, or do irregularly shaped in-place updates. Hence, + // L21 <- -L22 * L21 * L11 cannot be done naively. Instead, we update the + // entire row i.e. we calculate + // [L21 L22 0] <- -[L21 L22 0] @ diag_blocks([L11', -I, -I]) + // which means [L21 L22 0] <- [-L21 * L11', L22, 0]. + auto identity = + IdentityMatrix(builder, shape.element_type(), block_size, block_size); + auto neg_identity = -identity; + + // 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 = 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, + /*broadcast_sizes=*/{num_blocks}); + + // Now we construct a loop that performs matrix-vector multiplications + // inverting the blocks one row at a time + std::vector tuple_shapes = { + // The loop iteration counter is a scalar, incremented each iteration. + ShapeUtil::MakeShape(S32, {}), + // The output has the shape of A, with one row updated each iteration. + ShapeUtil::MakeShape(shape.element_type(), + {num_blocks, block_size, block_size}), + // The input is a loop invariant. + ShapeUtil::MakeShape(shape.element_type(), + {num_blocks, block_size, block_size})}; + Shape tuple_shape = ShapeUtil::MakeTupleShape(tuple_shapes); + + auto init_i = One(builder, S32); + auto init = Tuple(builder, {init_i, output, scaled_diag_blocks}); + + // Construct the loop condition function. + std::unique_ptr condb = + builder->CreateSubBuilder("InvertDiagCond"); + { + auto i = GetTupleElement( + Parameter(condb.get(), 0, tuple_shape, "InvertDiagCondTuple"), 0); + Lt(i, ConstantR0(condb.get(), block_size)); + } + TF_ASSIGN_OR_RETURN(auto cond, condb->Build()); + + // Construct the loop body function. + std::unique_ptr bodyb = + builder->CreateSubBuilder("InvertDiagBody"); + { + auto input_tuple = + Parameter(bodyb.get(), 0, tuple_shape, "InvertDiagBodyTuple"); + + auto i = GetTupleElement(input_tuple, 0); + auto body_out = GetTupleElement(input_tuple, 1); + auto body_input = GetTupleElement(input_tuple, 2); + + auto zero = ConstantR0(bodyb.get(), 0); + auto j = (lower) ? i : ScalarLike(i, block_size - 1) - i; + auto input_row = + DynamicSlice(body_input, {zero, j, zero}, + /*slice_sizes=*/{num_blocks, 1, block_size}); + + // We want -L21 L11^{-1} + DotDimensionNumbers dnums; + dnums.add_lhs_batch_dimensions(0); + dnums.add_rhs_batch_dimensions(0); + dnums.add_lhs_contracting_dimensions(2); + dnums.add_rhs_contracting_dimensions(1); + PrecisionConfig precision_proto; + precision_proto.add_operand_precision(precision); + precision_proto.add_operand_precision(precision); + auto update = -DotGeneral(input_row, body_out, dnums, &precision_proto); + + 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}); + } + TF_ASSIGN_OR_RETURN(auto body, bodyb->Build()); + + // Construct the While loop and return the result, + // return while_loop(cond_fun, body_fun, init)[1] + auto invert_while = While(cond, body, init); + auto inv_diag_blocks = GetTupleElement(invert_while, 1); + + // Undo the scaling + inv_diag_blocks = Div(inv_diag_blocks, diags, + /*broadcast_dimensions=*/{0, 1}); + + // Reshape back to original batch major dimensions + return Reshape(inv_diag_blocks, AsInt64Slice(shape.dimensions())); + }); +} + +XlaOp SolveWithInvertedDiagonalBlocks(XlaOp a, XlaOp b, XlaOp inv_diag_blocks, + bool left_side, bool lower, + bool transpose_a, bool conjugate_a, + PrecisionConfig::Precision precision) { + XlaBuilder* builder = a.builder(); + return builder->ReportErrorOrReturn([&]() -> StatusOr { + TF_ASSIGN_OR_RETURN(Shape blocks_shape, builder->GetShape(inv_diag_blocks)); + TF_ASSIGN_OR_RETURN(Shape b_shape, builder->GetShape(b)); + int64 block_size = ShapeUtil::GetDimension(blocks_shape, -1); + + TF_ASSIGN_OR_RETURN(Shape a_shape, builder->GetShape(a)); + int64 ndims = a_shape.rank(); + int64 n = ShapeUtil::GetDimension(a_shape, -1); + int64 num_blocks = n / block_size + (n % block_size != 0); + int64 m_dim = (left_side) ? -1 : -2; + int64 m = ShapeUtil::GetDimension(b_shape, m_dim); + + // Initialize the solution + auto x = ZerosLike(b); + + // This loop is unrolled for performance reasons, but it could be expressed + // rolled as well since the matrices are of the same size each iteration + for (int i = 0; i < num_blocks; i++) { + // High-level intuition: We have B[i] = L[i] @ X. Since L is upper + // triangular this means B[i] = L[i, :i + 1] @ X[:i + 1]. We can split + // this into two parts: B[i] = L[i, :i] @ X[:i] + L[i, i] @ X[i] which + // can be solved for X[i] as X[i] = inv(L[i, i]) @ B[i] - L[i, :i] @ X[:i] + + // Decide whether we go from first block to last or vice versa + auto j = (left_side ^ lower ^ transpose_a) ? num_blocks - 1 - i : i; + + // Get the size of the inverse blocks (the last one might be smaller) + int64 block = (n % block_size != 0 && j + 1 == num_blocks) + ? n % block_size + : block_size; + auto inv_block = + MaybeConjugate(Collapse(SliceInMinorDims(inv_diag_blocks, {j, 0, 0}, + {j + 1, block, block}), + /*dimensions=*/{ndims - 2, ndims - 1}), + conjugate_a); + + // Get the corresponding row of B + int64 k = std::min((j + 1) * block_size, n); + std::vector start = {j * block_size, 0}; + std::vector end = {k, m}; + if (!left_side) { + std::swap(start[0], start[1]); + std::swap(end[0], end[1]); + } + auto b_row = SliceInMinorDims(b, start, end); + + XlaOp remainder; + if (i == 0) { + remainder = b_row; + } else { + // This matrix multiply involves a lot of multiplying with zero (namely, + // X[i * block_size:] = 0), but this is faster than slicing... + end = {k, n}; + if (!left_side) { + std::swap(end[0], end[1]); + } + if (transpose_a) { + std::swap(start[0], start[1]); + std::swap(end[0], end[1]); + } + auto a_row = + MaybeConjugate(SliceInMinorDims(a, start, end), conjugate_a); + if (left_side) { + remainder = + b_row - BatchDot(MaybeTransposeInMinorDims(a_row, transpose_a), x, + precision); + } else { + remainder = + b_row - BatchDot(x, MaybeTransposeInMinorDims(a_row, transpose_a), + precision); + } + } + + XlaOp x_update; + auto zero = Zero(builder, S32); + auto start_index = ConstantR0WithType(builder, S32, j * block_size); + std::vector update_starts = {start_index, zero}; + if (left_side) { + x_update = BatchDot(MaybeTransposeInMinorDims(inv_block, transpose_a), + remainder, precision); + } else { + x_update = BatchDot(remainder, + MaybeTransposeInMinorDims(inv_block, transpose_a), + precision); + std::swap(update_starts[0], update_starts[1]); + } + x = DynamicUpdateSliceInMinorDims(x, x_update, /*starts=*/update_starts); + } + + return x; + }); +} + +XlaOp TriangularSolve(XlaOp a, XlaOp b, bool left_side, bool lower, + bool transpose_a, bool conjugate_a, int64 block_size, + PrecisionConfig::Precision precision) { + XlaBuilder* builder = a.builder(); + return builder->ReportErrorOrReturn([&]() -> StatusOr { + TF_ASSIGN_OR_RETURN(Shape a_shape, builder->GetShape(a)); + TF_ASSIGN_OR_RETURN(Shape b_shape, builder->GetShape(b)); + if (a_shape.rank() != b_shape.rank()) { + return InvalidArgument( + "Arguments to TriangularSolve have shapes with different ranks: " + "%s vs. %s", + ShapeUtil::HumanString(a_shape), ShapeUtil::HumanString(b_shape)); + } + const int64 ndims = a_shape.rank(); + if (ndims < 2) { + return InvalidArgument( + "Arguments to TriangularSolve was rank %d but must have rank >= 2.", + ndims); + } + // The batch dimensions must be equal. + std::vector batch_dimensions; + for (int i = 0; i < ndims - 2; ++i) { + int64 a_size = a_shape.dimensions(i); + int64 b_size = b_shape.dimensions(i); + if (a_size != b_size) { + return InvalidArgument( + "Batch dimensions of arguments to TriangularSolve must be equal; " + "shapes were %s and %s.", + ShapeUtil::HumanString(a_shape), ShapeUtil::HumanString(b_shape)); + } + batch_dimensions.push_back(a_size); + } + + if (ShapeUtil::GetDimension(a_shape, -1) != + ShapeUtil::GetDimension(a_shape, -2)) { + return InvalidArgument( + "The 'a' argument to TriangularSolve must be a batched square matrix;" + " shape was: %s", + ShapeUtil::HumanString(a_shape)); + } + const int64 m = ShapeUtil::GetDimension(b_shape, -2); + const int64 n = ShapeUtil::GetDimension(b_shape, -1); + if ((left_side ? m : n) != ShapeUtil::GetDimension(a_shape, -1)) { + return InvalidArgument( + "Arguments to TriangularSolve have incompatible matrix shapes %s and " + "%s", + ShapeUtil::HumanString(a_shape), ShapeUtil::HumanString(b_shape)); + } + + if (block_size < 1) { + return InvalidArgument( + "block_size argument to TriangularSolve must be >= 1; got %d", + block_size); + } + + if (ShapeUtil::IsZeroElementArray(b_shape)) { + // The output has the same shape as 'b', and since the output has zero + // elements, any such array will do. + return b; + } + + // We find the diagonal blocks of the coefficient matrix + auto diag_blocks = DiagonalBlocks(a, block_size); + + // We invert these blocks in parallel using batched matrix-vector products + 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, + transpose_a, conjugate_a, precision); + + return x; + }); +} + +} // namespace xla diff --git a/tensorflow/compiler/tf2xla/lib/triangular_solve.h b/tensorflow/compiler/xla/client/lib/triangular_solve.h similarity index 88% rename from tensorflow/compiler/tf2xla/lib/triangular_solve.h rename to tensorflow/compiler/xla/client/lib/triangular_solve.h index 2303234f361e54cd2a0ad495cb03b371bed76877..50a3b30ebd1c15eb6d2ace4e351cb41f21db7093 100644 --- a/tensorflow/compiler/tf2xla/lib/triangular_solve.h +++ b/tensorflow/compiler/xla/client/lib/triangular_solve.h @@ -13,13 +13,13 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef TENSORFLOW_COMPILER_TF2XLA_LIB_TRIANGULAR_SOLVE_H_ -#define TENSORFLOW_COMPILER_TF2XLA_LIB_TRIANGULAR_SOLVE_H_ +#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 tensorflow { +namespace xla { // Solves systems of linear equations with lower or upper triangular coefficient // matrices by forward- or back-substitution. Broadcasting along leading @@ -57,11 +57,11 @@ namespace tensorflow { // // Uses a blocked algorithm if `block_size` is > 1; if block_size == 1 then no // blocking is used. -xla::XlaOp TriangularSolve( - xla::XlaOp a, xla::XlaOp b, bool left_side, bool lower, bool transpose_a, +XlaOp TriangularSolve( + XlaOp a, XlaOp b, bool left_side, bool lower, bool transpose_a, bool conjugate_a, int64 block_size = 128, - xla::PrecisionConfig::Precision precision = xla::PrecisionConfig::HIGHEST); + PrecisionConfig::Precision precision = PrecisionConfig::HIGHEST); -} // namespace tensorflow +} // namespace xla -#endif // TENSORFLOW_COMPILER_TF2XLA_LIB_TRIANGULAR_SOLVE_H_ +#endif // TENSORFLOW_COMPILER_XLA_CLIENT_LIB_TRIANGULAR_SOLVE_H_ diff --git a/tensorflow/compiler/xla/client/lib/triangular_solve_test.cc b/tensorflow/compiler/xla/client/lib/triangular_solve_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..284a2e9d183a6a7923fb59ac134ce3b3a3a96e35 --- /dev/null +++ b/tensorflow/compiler/xla/client/lib/triangular_solve_test.cc @@ -0,0 +1,447 @@ +/* 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/compiler/xla/client/lib/triangular_solve.h" + +#include +#include +#include + +#include "tensorflow/compiler/xla/array2d.h" +#include "tensorflow/compiler/xla/client/lib/math.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/types.h" +#include "tensorflow/core/lib/core/status_test_util.h" + +namespace xla { +namespace { + +using TriangularSolveTest = ClientLibraryTestBase; +using TriangularSolveLeftLookingTest = ClientLibraryTestBase; + +static constexpr float kNan = std::numeric_limits::quiet_NaN(); + +Array2D AValsLower() { + return {{2, kNan, kNan, kNan}, + {3, 6, kNan, kNan}, + {4, 7, 9, kNan}, + {5, 8, 10, 11}}; +} + +Array2D AValsUpper() { + return {{2, 3, 4, 5}, + {kNan, 6, 7, 8}, + {kNan, kNan, 9, 10}, + {kNan, kNan, kNan, 11}}; +} + +Array2D BValsRight() { + return {{1, 2, 3, 4}, {5, 6, 7, 8}, {9, 10, 11, 12}}; +} + +Array2D BValsLeft() { + return {{1, 2, 3}, {4, 5, 6}, {7, 8, 9}, {10, 11, 12}}; +} + +static constexpr complex64 kNanC64 = complex64(kNan, kNan); + +Array2D AValsLowerComplex() { + return {{2, kNanC64, kNanC64, kNanC64}, + {complex64(3, 1), 6, kNanC64, kNanC64}, + {4, complex64(7, 2), 9, kNanC64}, + {5, 8, complex64(10, 3), 11}}; +} + +Array2D AValsUpperComplex() { + return {{2, 3, complex64(4, 3), 5}, + {kNanC64, 6, complex64(7, 2), 8}, + {kNanC64, kNanC64, complex64(9, 1), 10}, + {kNanC64, kNanC64, kNanC64, 11}}; +} + +Array2D BValsRightComplex() { + return {{1, 2, 3, 4}, {5, 6, 7, 8}, {9, 10, 11, 12}}; +} + +Array2D BValsLeftComplex() { + return {{1, 2, 3}, {4, 5, 6}, {7, 8, 9}, {10, 11, 12}}; +} + +XLA_TEST_F(TriangularSolveTest, EmptyArrays) { + XlaBuilder builder(TestName()); + + XlaOp a, b; + auto a_data = + CreateR2Parameter(Array2D(0, 0), 0, "a", &builder, &a); + auto b_data = + 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); + + ComputeAndCompareR2(&builder, Array2D(0, 10), + {a_data.get(), b_data.get()}); +} + +XLA_TEST_F(TriangularSolveTest, SimpleRightLowerTranspose) { + XlaBuilder builder(TestName()); + + XlaOp a, b; + auto a_data = CreateR2Parameter(AValsLower(), 0, "a", &builder, &a); + 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); + + Array2D expected({ + {0.5, 0.08333334, 0.04629629, 0.03367003}, + {2.5, -0.25, -0.1388889, -0.1010101}, + {4.5, -0.58333331, -0.32407406, -0.23569024}, + }); + + ComputeAndCompareR2(&builder, expected, {a_data.get(), b_data.get()}, + ErrorSpec(1e-2, 1e-2)); +} + +XLA_TEST_F(TriangularSolveTest, SimpleRightLowerNotranspose) { + XlaBuilder builder(TestName()); + + XlaOp a, b; + auto a_data = CreateR2Parameter(AValsLower(), 0, "a", &builder, &a); + 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); + + Array2D expected({ + {-0.16414141, -0.06902357, -0.07070707, 0.36363636}, + {0.64393939, 0.06565657, -0.03030303, 0.72727273}, + {1.4520202, 0.2003367, 0.01010101, 1.09090909}, + }); + + ComputeAndCompareR2(&builder, expected, {a_data.get(), b_data.get()}, + ErrorSpec(1e-2, 1e-2)); +} + +XLA_TEST_F(TriangularSolveTest, SimpleRightUpperTranspose) { + XlaBuilder builder(TestName()); + + XlaOp a, b; + auto a_data = CreateR2Parameter(AValsUpper(), 0, "a", &builder, &a); + 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); + + Array2D expected({ + {-0.16414141, -0.06902357, -0.07070707, 0.36363636}, + {0.64393939, 0.06565657, -0.03030303, 0.72727273}, + {1.4520202, 0.2003367, 0.01010101, 1.09090909}, + }); + + ComputeAndCompareR2(&builder, expected, {a_data.get(), b_data.get()}, + ErrorSpec(1e-2, 1e-2)); +} + +XLA_TEST_F(TriangularSolveTest, SimpleRightUpperNotranspose) { + XlaBuilder builder(TestName()); + + XlaOp a, b; + auto a_data = CreateR2Parameter(AValsUpper(), 0, "a", &builder, &a); + 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); + + Array2D expected({ + {0.5, 0.08333334, 0.04629629, 0.03367003}, + {2.5, -0.25, -0.1388889, -0.1010101}, + {4.5, -0.58333331, -0.32407406, -0.23569024}, + }); + + ComputeAndCompareR2(&builder, expected, {a_data.get(), b_data.get()}, + ErrorSpec(1e-2, 1e-2)); +} + +XLA_TEST_F(TriangularSolveTest, SimpleLeftLowerTranspose) { + XlaBuilder builder(TestName()); + + XlaOp a, b; + auto a_data = CreateR2Parameter(AValsLower(), 0, "a", &builder, &a); + 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); + + Array2D expected({ + {-0.89646465, -0.69444444, -0.49242424}, + {-0.27441077, -0.24074074, -0.20707071}, + {-0.23232323, -0.22222222, -0.21212121}, + {0.90909091, 1., 1.09090909}, + }); + + ComputeAndCompareR2(&builder, expected, {a_data.get(), b_data.get()}, + ErrorSpec(1e-2, 1e-2)); +} + +XLA_TEST_F(TriangularSolveTest, SimpleLeftLowerNotranspose) { + XlaBuilder builder(TestName()); + + XlaOp a, b; + auto a_data = CreateR2Parameter(AValsLower(), 0, "a", &builder, &a); + 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); + + Array2D expected({ + {0.5, 1.0, 1.5}, + {0.41666667, 0.33333333, 0.25}, + {0.23148148, 0.18518519, 0.13888889}, + {0.16835017, 0.13468013, 0.1010101}, + }); + + ComputeAndCompareR2(&builder, expected, {a_data.get(), b_data.get()}, + ErrorSpec(1e-2, 1e-2)); +} + +XLA_TEST_F(TriangularSolveTest, SimpleLeftLowerNotransposeIrregularblock) { + XlaBuilder builder(TestName()); + + XlaOp a, b; + auto a_data = CreateR2Parameter(AValsLower(), 0, "a", &builder, &a); + 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); + + Array2D expected({ + {0.5, 1.0, 1.5}, + {0.41666667, 0.33333333, 0.25}, + {0.23148148, 0.18518519, 0.13888889}, + {0.16835017, 0.13468013, 0.1010101}, + }); + + ComputeAndCompareR2(&builder, expected, {a_data.get(), b_data.get()}, + ErrorSpec(1e-2, 1e-2)); +} + +XLA_TEST_F(TriangularSolveTest, SimpleLeftUpperTranspose) { + XlaBuilder builder(TestName()); + + XlaOp a, b; + auto a_data = CreateR2Parameter(AValsUpper(), 0, "a", &builder, &a); + 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); + + Array2D expected({ + {0.5, 1.0, 1.5}, + {0.41666667, 0.33333333, 0.25}, + {0.23148148, 0.18518519, 0.13888889}, + {0.16835017, 0.13468013, 0.1010101}, + }); + + ComputeAndCompareR2(&builder, expected, {a_data.get(), b_data.get()}, + ErrorSpec(1e-2, 1e-2)); +} + +XLA_TEST_F(TriangularSolveTest, SimpleLeftUpperNotranspose) { + XlaBuilder builder(TestName()); + + XlaOp a, b; + auto a_data = CreateR2Parameter(AValsUpper(), 0, "a", &builder, &a); + 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); + + Array2D expected({ + {-0.89646465, -0.69444444, -0.49242424}, + {-0.27441077, -0.24074074, -0.20707071}, + {-0.23232323, -0.22222222, -0.21212121}, + {0.90909091, 1., 1.09090909}, + }); + + ComputeAndCompareR2(&builder, expected, {a_data.get(), b_data.get()}, + ErrorSpec(1e-2, 1e-2)); +} + +XLA_TEST_F(TriangularSolveTest, SimpleRightLowerTransposeConjugate) { + XlaBuilder builder(TestName()); + + XlaOp a, b; + auto a_data = + CreateR2Parameter(AValsLowerComplex(), 0, "a", &builder, &a); + auto b_data = + CreateR2Parameter(BValsRightComplex(), 1, "b", &builder, &b); + TriangularSolve(a, b, + /*left_side=*/false, /*lower=*/true, + /*transpose_a=*/true, /*conjugate_a=*/true, + /*block_size=*/2); + + Array2D expected({ + {0.5, complex64(0.08333333, 0.08333333), + complex64(0.02777778, -0.0462963), complex64(0.06313131, -0.01094276)}, + {2.5, complex64(-0.25, 0.41666667), complex64(-0.23148148, -0.37962963), + complex64(0.08670034, -0.02104377)}, + {4.5, complex64(-0.58333333, 0.75), complex64(-0.49074074, -0.71296296), + complex64(0.11026936, -0.03114478)}, + }); + + ComputeAndCompareR2( + &builder, expected, {a_data.get(), b_data.get()}, ErrorSpec(1e-2, 1e-2)); +} + +XLA_TEST_F(TriangularSolveTest, SimpleLeftUpperTransposeNoconjugate) { + XlaBuilder builder(TestName()); + + XlaOp a, b; + auto a_data = + CreateR2Parameter(AValsUpperComplex(), 0, "a", &builder, &a); + auto b_data = + CreateR2Parameter(BValsLeftComplex(), 1, "b", &builder, &b); + TriangularSolve(a, b, + /*left_side=*/true, /*lower=*/false, + /*transpose_a=*/true, /*conjugate_a=*/false, + /*block_size=*/2); + + Array2D expected({ + {0.5, 1., 1.5}, + {0.41666667, 0.33333333, 0.25}, + {complex64(0.20020325, -2.81504065e-01), + complex64(0.13821138, -4.22764228e-01), + complex64(0.07621951, -5.64024390e-01)}, + {complex64(0.19678492, 2.55912786e-01), + complex64(0.17738359, 3.84331116e-01), + complex64(0.15798226, 5.12749446e-01)}, + }); + + ComputeAndCompareR2( + &builder, expected, {a_data.get(), b_data.get()}, ErrorSpec(1e-2, 1e-2)); +} + +XLA_TEST_F(TriangularSolveTest, BatchedLeftUpper) { + XlaBuilder builder(TestName()); + + Array3D bvals(7, 5, 5); + bvals.FillIota(1.); + + // Set avals to the upper triangle of bvals. + Array3D avals = bvals; + avals.Each([](absl::Span indices, float* value) { + if (indices[1] > indices[2]) { + *value = 0; + } + }); + + 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)); + + ComputeAndCompareR3(&builder, bvals, {a_data.get(), b_data.get()}, + ErrorSpec(1e-2, 1e-2)); +} + +struct TriangularSolveTestSpec { + int m, n; // A is mxm, B is mxn + bool left_side; + bool lower; + bool transpose_a; +}; + +class TriangularSolveParametricTest + : public ClientLibraryTestBase, + public ::testing::WithParamInterface {}; + +XLA_TEST_P(TriangularSolveParametricTest, Random) { + TriangularSolveTestSpec spec = GetParam(); + + XlaBuilder builder(TestName()); + + Array2D avals(spec.m, spec.m); + avals.FillRandom(1.0); + for (int i = 0; i < spec.m; ++i) { + avals(i, i) += 10; + } + + std::pair bdims = spec.left_side ? std::make_pair(spec.m, spec.n) + : std::make_pair(spec.n, spec.m); + Array2D bvals(bdims.first, bdims.second); + bvals.FillRandom(1.0); + + 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 a_tri = Triangle(a, spec.lower); + a_tri = MaybeTransposeInMinorDims(a_tri, spec.transpose_a); + if (spec.left_side) { + BatchDot(a_tri, x); + } else { + BatchDot(x, a_tri); + } + + ComputeAndCompareR2(&builder, bvals, {a_data.get(), b_data.get()}, + ErrorSpec(1e-2, 1e-2)); +} + +std::vector TriangularSolveTests() { + std::vector specs; + for (int m : {5, 10}) { + for (int n : {5, 10}) { + for (bool left_side : {false, true}) { + for (bool lower : {false, true}) { + for (bool transpose_a : {false, true}) { + specs.push_back({m, n, left_side, lower, transpose_a}); + } + } + } + } + } + return specs; +} + +INSTANTIATE_TEST_SUITE_P(TriangularSolveParametricTestInstantiation, + TriangularSolveParametricTest, + ::testing::ValuesIn(TriangularSolveTests())); + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/client/local_client.cc b/tensorflow/compiler/xla/client/local_client.cc index f96b6c9c261a9686fb647e3da0dcc933cd1f70df..48b5f94538f453785194bc434a91ee0a10c020c2 100644 --- a/tensorflow/compiler/xla/client/local_client.cc +++ b/tensorflow/compiler/xla/client/local_client.cc @@ -71,9 +71,9 @@ Status LocalExecutable::ValidateExecutionOptions( "parameter " "%d: want %s, got %s", i, - ShapeUtil::HumanString( + ShapeUtil::HumanStringWithLayout( computation_layout.parameter_layout(i).shape()), - ShapeUtil::HumanString(arguments[i]->on_host_shape())); + ShapeUtil::HumanStringWithLayout(arguments[i]->on_host_shape())); } } @@ -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); @@ -310,4 +309,28 @@ StatusOr LocalClient::ReplicaNumberToDeviceOrdinal(int replica_number) { return local_service_->ReplicaNumberToDeviceOrdinal(replica_number); } +StatusOr LocalClient::TransferToLocalServer( + const ::xla::BorrowingLiteral& literal, int device_oridinal) { + const ::xla::Shape& shape = literal.shape(); + + TF_ASSIGN_OR_RETURN( + ::xla::ScopedShapedBuffer shaped_buffer, + backend().transfer_manager()->AllocateScopedShapedBuffer( + shape, backend().memory_allocator(), device_oridinal)); + TF_ASSIGN_OR_RETURN(auto stream, + mutable_backend()->BorrowStream(device_oridinal)); + TF_RETURN_IF_ERROR(backend().transfer_manager()->TransferLiteralToDevice( + stream.get(), literal, shaped_buffer)); + std::vector<::xla::ScopedShapedBuffer> replicated_buffer; + replicated_buffer.emplace_back(std::move(shaped_buffer)); + ::xla::TransferToServerResponse result; + TF_ASSIGN_OR_RETURN(*result.mutable_data(), + local_service_->RegisterReplicatedBuffers( + std::move(replicated_buffer), + absl::StrCat("TransferToServer literal of shape ", + ::xla::ShapeUtil::HumanString(shape)))); + + return result; +} + } // namespace xla diff --git a/tensorflow/compiler/xla/client/local_client.h b/tensorflow/compiler/xla/client/local_client.h index feb2f8ec9dab5bf13afdc866d10ccbe74f8edcb9..ddb36680e8b185b053368baffa6f1d5cac50dc07 100644 --- a/tensorflow/compiler/xla/client/local_client.h +++ b/tensorflow/compiler/xla/client/local_client.h @@ -60,8 +60,8 @@ class LocalExecutable { // Validates that the given arguments and options satisfy various constraints // of the computation. // - // The given ExecutableRunOptions override any values from legacy_flags - // (TF_XLA_FLAGS environment variable). + // The given ExecutableRunOptions override any values from TF_XLA_FLAGS + // environment variable. Status ValidateExecutionOptions( const absl::Span arguments, const ExecutableRunOptions& run_options, const Backend& backend); @@ -69,8 +69,8 @@ class LocalExecutable { // Records the computation in a SessionModule proto with the arguments used to // invoke it, and the result. Enabled by flag: --tla_dump_executions_to. // - // The given ServiceExecutableRunOptions override any values from legacy_flags - // (TF_XLA_FLAGS environment variable). + // The given ServiceExecutableRunOptions override any values from TF_XLA_FLAGS + // environment variable. StatusOr ExecuteAndDump( const ServiceExecutableRunOptions* run_options, const absl::Span arguments); @@ -114,8 +114,8 @@ 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 legacy_flags - // (TF_XLA_FLAGS environment variable). + // The given ExecutableBuildOptions override any values from TF_XLA_FLAGS + // environment variable. StatusOr> Compile( const XlaComputation& computation, const absl::Span argument_layouts, @@ -129,6 +129,10 @@ class LocalClient : public Client { const Literal& literal, int device_ordinal, DeviceMemoryAllocator* allocator = nullptr); + // Transfer the BorrowingLiteral to the device with the given ordinal. + StatusOr TransferToLocalServer( + const ::xla::BorrowingLiteral& literal, int device_oridinal); + // Copy the data from the device contained in the given ShapedBuffer and // return as a Literal. StatusOr ShapedBufferToLiteral(const ShapedBuffer& shaped_buffer); diff --git a/tensorflow/compiler/xla/client/sharding_builder.cc b/tensorflow/compiler/xla/client/sharding_builder.cc index 176802b33ef824a1f898255a19e44def3c1fc982..b9bff06cbdbc3525eb19d5df885952c3971d9d6a 100644 --- a/tensorflow/compiler/xla/client/sharding_builder.cc +++ b/tensorflow/compiler/xla/client/sharding_builder.cc @@ -36,7 +36,7 @@ OpSharding Tile(const Shape& tile_shape, const TileAssignment& tile_assignment) { OpSharding result; result.set_type(OpSharding::Type::OpSharding_Type_OTHER); - *result.mutable_tile_shape() = tile_shape; + *result.mutable_tile_shape() = tile_shape.ToProto(); for (int64 dim : tile_assignment.dimensions()) { result.add_tile_assignment_dimensions(dim); } @@ -50,9 +50,9 @@ OpSharding Tile1D(const Shape& tile_shape, int64 num_tiles) { OpSharding result; result.set_type(OpSharding::Type::OpSharding_Type_OTHER); - CHECK_EQ(ShapeUtil::Rank(tile_shape), 1); + CHECK_EQ(tile_shape.rank(), 1); std::vector dimensions(1, num_tiles); - *result.mutable_tile_shape() = tile_shape; + *result.mutable_tile_shape() = tile_shape.ToProto(); auto& tile_dimension = (*result.mutable_tile_shape()->mutable_dimensions())[0]; tile_dimension = CeilOfRatio(static_cast(tile_dimension), num_tiles); diff --git a/tensorflow/compiler/xla/client/xla_builder.cc b/tensorflow/compiler/xla/client/xla_builder.cc index 7d081b27222bd31ddbe7c64b4dea8a4d5a371acb..5c9f9f708883f458b67205058fc7c1e1e2ad02f5 100644 --- a/tensorflow/compiler/xla/client/xla_builder.cc +++ b/tensorflow/compiler/xla/client/xla_builder.cc @@ -22,6 +22,7 @@ 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" @@ -30,11 +31,11 @@ limitations under the License. #include "tensorflow/compiler/xla/client/sharding_builder.h" #include "tensorflow/compiler/xla/client/xla_computation.h" #include "tensorflow/compiler/xla/execution_options_util.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_opcode.h" #include "tensorflow/compiler/xla/service/shape_inference.h" #include "tensorflow/compiler/xla/util.h" -#include "tensorflow/core/platform/mutex.h" namespace xla { @@ -42,12 +43,30 @@ using absl::StrCat; namespace { -int64 GetUniqueId() { - static tensorflow::mutex mu(tensorflow::LINKER_INITIALIZED); - static int64 built_counter = 0; - tensorflow::mutex_lock loc(mu); - const int64 id = built_counter++; - return id; +static const char kNameSeparator = '.'; + +// Retrieves the base name of an instruction or computation fully qualified +// name, using separator as boundary between the initial base name part, and +// the numeric identification. +string GetBaseName(const string& name, char separator) { + auto pos = name.rfind(separator); + CHECK_NE(pos, string::npos) << name; + return name.substr(0, pos); +} + +// Generates a fully qualified computation/instruction name. +string GetFullName(const string& base_name, char separator, int64 id) { + const char separator_str[] = {separator, '\0'}; + return StrCat(base_name, separator_str, id); +} + +// Common function to standardize setting name and IDs on computation and +// instruction proto entities. +template +void SetProtoIdAndName(T* entry, const string& base_name, char separator, + int64 id) { + entry->set_id(id); + entry->set_name(GetFullName(base_name, separator, id)); } } // namespace @@ -86,7 +105,7 @@ StatusOr XlaBuilder::GetShape(const XlaOp& op) const { TF_RETURN_IF_ERROR(first_error_); TF_ASSIGN_OR_RETURN(auto instr, LookUpInstruction(op)); - return instr->shape(); + return Shape(instr->shape()); } StatusOr> XlaBuilder::GetOperandShapes( @@ -139,7 +158,7 @@ StatusOr XlaBuilder::GetProgramShape(int64 root_id) const { ProgramShape program_shape; - *program_shape.mutable_result() = root_proto->shape(); + *program_shape.mutable_result() = Shape(root_proto->shape()); // Check that the parameter numbers are continuous from 0, and add parameter // shapes and names to the program shape. @@ -156,7 +175,7 @@ StatusOr XlaBuilder::GetProgramShape(int64 root_id) const { const int64 index = instr.parameter_number(); TF_RET_CHECK(index >= 0 && index < param_count) << "invalid parameter number: " << index; - *program_shape.mutable_parameters(index) = instr.shape(); + *program_shape.mutable_parameters(index) = Shape(instr.shape()); *program_shape.mutable_parameter_names(index) = instr.name(); } } @@ -176,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; } @@ -192,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::kCrossReplicaSum: - // TODO(b/33009255): Implmement constant folding for cross replica sum. + case HloOpcode::kAllReduce: + // TODO(b/33009255): Implement constant folding for cross replica sum. case HloOpcode::kInfeed: case HloOpcode::kOutfeed: case HloOpcode::kCall: @@ -223,6 +252,42 @@ void XlaBuilder::IsConstantVisitor(const int64 op_handle, visited->insert(op_handle); } +Status XlaBuilder::SetDynamicBinding(int64 dynamic_size_param_num, + ShapeIndex dynamic_size_param_index, + 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}, + DynamicParameterBinding::DynamicDimension{ + target_param_num, target_param_index, target_dim_num})); + return Status::OK(); +} + XlaComputation XlaBuilder::BuildAndNoteError() { DCHECK(parent_builder_ != nullptr); auto build_status = Build(); @@ -234,41 +299,62 @@ 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); } - HloComputationProto entry; - entry.set_id(GetUniqueId()); // Give the computation a global unique id. - entry.set_name(StrCat(name_, entry.id())); // Ensure that the name is unique. + // 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); + } + }; - TF_ASSIGN_OR_RETURN(*entry.mutable_program_shape(), GetProgramShape(root_id)); + 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)); + *entry.mutable_program_shape() = program_shape.ToProto(); entry.set_root_id(root_id); for (auto& instruction : instructions_) { // Ensures that the instruction names are unique among the whole graph. - const string& new_name = - StrCat(instruction.name(), ".", entry.id(), ".", instruction.id()); - instruction.set_name(new_name); + instruction.set_name( + GetFullName(instruction.name(), kNameSeparator, instruction.id())); entry.add_instructions()->Swap(&instruction); } @@ -283,6 +369,12 @@ StatusOr XlaBuilder::Build(int64 root_id) { module->add_computations()->Swap(&e.second); } module->add_computations()->Swap(&entry); + if (!input_output_aliases_.empty()) { + TF_RETURN_IF_ERROR( + PopulateInputOutputAlias(module, program_shape, input_output_aliases_)); + } + *(module->mutable_dynamic_parameter_binding()) = + dynamic_parameter_binding_.ToProto(); // Clear data held by this builder. this->instructions_.clear(); @@ -293,13 +385,42 @@ StatusOr XlaBuilder::Build(int64 root_id) { return std::move(computation); } +/* static */ Status XlaBuilder::PopulateInputOutputAlias( + HloModuleProto* module, const ProgramShape& program_shape, + const std::vector& input_output_aliases) { + HloInputOutputAliasConfig config(program_shape.result()); + for (auto& alias : input_output_aliases) { + // The HloInputOutputAliasConfig does not do parameter validation as it only + // carries the result shape. Maybe it should be constructed with a + // ProgramShape to allow full validation. We will still get an error when + // trying to compile the HLO module, but would be better to have validation + // at this stage. + if (alias.param_number >= program_shape.parameters_size()) { + return InvalidArgument("Invalid parameter number %ld (total %ld)", + alias.param_number, + program_shape.parameters_size()); + } + const Shape& parameter_shape = program_shape.parameters(alias.param_number); + if (!ShapeUtil::IndexIsValid(parameter_shape, alias.param_index)) { + return InvalidArgument("Invalid parameter %ld index: %s", + alias.param_number, + alias.param_index.ToString().c_str()); + } + 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(); +} + StatusOr XlaBuilder::InDimBroadcast( const Shape& shape, const XlaOp& operand, absl::Span broadcast_dimensions) { TF_RETURN_IF_ERROR(first_error_); HloInstructionProto instr; - *instr.mutable_shape() = shape; + *instr.mutable_shape() = shape.ToProto(); for (int64 dim : broadcast_dimensions) { instr.add_dimensions(dim); } @@ -313,7 +434,7 @@ StatusOr XlaBuilder::AddBroadcastSequence(const Shape& output_shape, TF_ASSIGN_OR_RETURN(const Shape& operand_shape, GetShape(operand)); CHECK(ShapeUtil::IsScalar(operand_shape) || - ShapeUtil::Rank(operand_shape) == ShapeUtil::Rank(output_shape)); + operand_shape.rank() == output_shape.rank()); Shape broadcast_shape = ShapeUtil::ChangeElementType(output_shape, operand_shape.element_type()); @@ -325,7 +446,7 @@ StatusOr XlaBuilder::AddBroadcastSequence(const Shape& output_shape, // Do explicit broadcast for degenerate broadcast. std::vector broadcast_dimensions; std::vector reshaped_dimensions; - for (int i = 0; i < ShapeUtil::Rank(operand_shape); i++) { + for (int i = 0; i < operand_shape.rank(); i++) { if (operand_shape.dimensions(i) == output_shape.dimensions(i)) { broadcast_dimensions.push_back(i); reshaped_dimensions.push_back(operand_shape.dimensions(i)); @@ -350,8 +471,9 @@ XlaOp XlaBuilder::UnaryOp(HloOpcode unop, const XlaOp& operand) { return ReportErrorOrReturn([&]() -> StatusOr { HloInstructionProto instr; TF_ASSIGN_OR_RETURN(const Shape& operand_shape, GetShape(operand)); - TF_ASSIGN_OR_RETURN(*instr.mutable_shape(), + TF_ASSIGN_OR_RETURN(Shape shape, ShapeInference::InferUnaryOpShape(unop, operand_shape)); + *instr.mutable_shape() = shape.ToProto(); return AddInstruction(std::move(instr), unop, {operand}); }); } @@ -362,12 +484,13 @@ XlaOp XlaBuilder::BinaryOp(HloOpcode binop, 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)); - TF_ASSIGN_OR_RETURN(*instr.mutable_shape(), + TF_ASSIGN_OR_RETURN(Shape shape, ShapeInference::InferBinaryOpShape( binop, lhs_shape, rhs_shape, broadcast_dimensions)); + *instr.mutable_shape() = shape.ToProto(); - const int64 lhs_rank = ShapeUtil::Rank(lhs_shape); - const int64 rhs_rank = ShapeUtil::Rank(rhs_shape); + const int64 lhs_rank = lhs_shape.rank(); + const int64 rhs_rank = rhs_shape.rank(); XlaOp updated_lhs = lhs; XlaOp updated_rhs = rhs; @@ -378,17 +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 : instr.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 < ShapeUtil::Rank(from_shape); - from_dim++) { + 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)); @@ -398,14 +523,14 @@ XlaOp XlaBuilder::BinaryOp(HloOpcode binop, const XlaOp& lhs, const XlaOp& rhs, } TF_ASSIGN_OR_RETURN(Shape updated_lhs_shape, GetShape(updated_lhs)); - if (!ShapeUtil::SameDimensions(instr.shape(), updated_lhs_shape)) { + if (!ShapeUtil::SameDimensions(shape, updated_lhs_shape)) { TF_ASSIGN_OR_RETURN(updated_lhs, - AddBroadcastSequence(instr.shape(), updated_lhs)); + AddBroadcastSequence(shape, updated_lhs)); } TF_ASSIGN_OR_RETURN(Shape updated_rhs_shape, GetShape(updated_rhs)); - if (!ShapeUtil::SameDimensions(instr.shape(), updated_rhs_shape)) { + if (!ShapeUtil::SameDimensions(shape, updated_rhs_shape)) { TF_ASSIGN_OR_RETURN(updated_rhs, - AddBroadcastSequence(instr.shape(), updated_rhs)); + AddBroadcastSequence(shape, updated_rhs)); } return AddInstruction(std::move(instr), binop, {updated_lhs, updated_rhs}); @@ -419,30 +544,28 @@ XlaOp XlaBuilder::TernaryOp(HloOpcode triop, const XlaOp& lhs, const XlaOp& rhs, TF_ASSIGN_OR_RETURN(const Shape& lhs_shape, GetShape(lhs)); TF_ASSIGN_OR_RETURN(const Shape& rhs_shape, GetShape(rhs)); TF_ASSIGN_OR_RETURN(const Shape& ehs_shape, GetShape(ehs)); - TF_ASSIGN_OR_RETURN(*instr.mutable_shape(), - ShapeInference::InferTernaryOpShape( - triop, lhs_shape, rhs_shape, ehs_shape)); + TF_ASSIGN_OR_RETURN( + Shape shape, ShapeInference::InferTernaryOpShape(triop, lhs_shape, + rhs_shape, ehs_shape)); + *instr.mutable_shape() = shape.ToProto(); XlaOp updated_lhs = lhs; XlaOp updated_rhs = rhs; XlaOp updated_ehs = ehs; - if (!ShapeUtil::IsTuple(instr.shape())) { - if (!ShapeUtil::IsTuple(lhs_shape) && - !ShapeUtil::SameDimensions(instr.shape(), lhs_shape)) { + if (!shape.IsTuple()) { + if (!lhs_shape.IsTuple() && + !ShapeUtil::SameDimensions(shape, lhs_shape)) { // lhs is being implicitly broadcasted. Change to explicit. - TF_ASSIGN_OR_RETURN(updated_lhs, - AddBroadcastSequence(instr.shape(), lhs)); + TF_ASSIGN_OR_RETURN(updated_lhs, AddBroadcastSequence(shape, lhs)); } - if (!ShapeUtil::IsTuple(rhs_shape) && - !ShapeUtil::SameDimensions(instr.shape(), rhs_shape)) { + if (!rhs_shape.IsTuple() && + !ShapeUtil::SameDimensions(shape, rhs_shape)) { // rhs is being implicitly broadcasted. Change to explicit. - TF_ASSIGN_OR_RETURN(updated_rhs, - AddBroadcastSequence(instr.shape(), rhs)); + TF_ASSIGN_OR_RETURN(updated_rhs, AddBroadcastSequence(shape, rhs)); } - if (!ShapeUtil::IsTuple(ehs_shape) && - !ShapeUtil::SameDimensions(instr.shape(), ehs_shape)) { + if (!ehs_shape.IsTuple() && + !ShapeUtil::SameDimensions(shape, ehs_shape)) { // ehs is being implicitly broadcasted. Change to explicit. - TF_ASSIGN_OR_RETURN(updated_ehs, - AddBroadcastSequence(instr.shape(), ehs)); + TF_ASSIGN_OR_RETURN(updated_ehs, AddBroadcastSequence(shape, ehs)); } } return AddInstruction(std::move(instr), triop, @@ -463,7 +586,7 @@ XlaOp XlaBuilder::Mul(const XlaOp& lhs, const XlaOp& rhs, XlaOp XlaBuilder::ConstantLiteral(const LiteralSlice& literal) { return ReportErrorOrReturn([&]() -> StatusOr { HloInstructionProto instr; - *instr.mutable_shape() = literal.shape(); + *instr.mutable_shape() = literal.shape().ToProto(); *instr.mutable_literal() = literal.ToProto(); return AddInstruction(std::move(instr), HloOpcode::kConstant); }); @@ -472,7 +595,7 @@ XlaOp XlaBuilder::ConstantLiteral(const LiteralSlice& literal) { XlaOp XlaBuilder::Iota(const Shape& shape, int64 iota_dimension) { return ReportErrorOrReturn([&]() -> StatusOr { HloInstructionProto instr; - *instr.mutable_shape() = shape; + *instr.mutable_shape() = shape.ToProto(); instr.add_dimensions(iota_dimension); return AddInstruction(std::move(instr), HloOpcode::kIota); }); @@ -492,10 +615,10 @@ XlaOp XlaBuilder::Call(const XlaComputation& computation, [](const Shape& shape) { return &shape; }); TF_ASSIGN_OR_RETURN(const ProgramShape& called_program_shape, computation.GetProgramShape()); - TF_ASSIGN_OR_RETURN( - *instr.mutable_shape(), - ShapeInference::InferCallShape(operand_shape_ptrs, - /*to_apply=*/called_program_shape)); + TF_ASSIGN_OR_RETURN(Shape shape, ShapeInference::InferCallShape( + operand_shape_ptrs, + /*to_apply=*/called_program_shape)); + *instr.mutable_shape() = shape.ToProto(); AddCalledComputation(computation, &instr); @@ -513,7 +636,7 @@ XlaOp XlaBuilder::Parameter(int64 parameter_number, const Shape& shape, } instr.set_parameter_number(parameter_number); instr.set_name(name); - *instr.mutable_shape() = shape; + *instr.mutable_shape() = shape.ToProto(); return AddInstruction(std::move(instr), HloOpcode::kParameter); }); } @@ -533,20 +656,54 @@ XlaOp XlaBuilder::Broadcast(const XlaOp& operand, // output, so to append dimensions on the left the instruction's dimensions // should just be the n highest dimension numbers of the output shape where // n is the number of input dimensions. - const int64 operand_rank = ShapeUtil::Rank(operand_shape); + const int64 operand_rank = operand_shape.rank(); std::vector dimensions(operand_rank); for (int i = 0; i < operand_rank; ++i) { - dimensions[i] = i + ShapeUtil::Rank(shape) - operand_rank; + dimensions[i] = i + shape.rank() - operand_rank; } return InDimBroadcast(shape, operand, dimensions); }); } XlaOp XlaBuilder::BroadcastInDim( - const XlaOp& operand, const Shape& shape, + const XlaOp& operand, const absl::Span out_dim_size, const absl::Span broadcast_dimensions) { return ReportErrorOrReturn([&]() -> StatusOr { - return InDimBroadcast(shape, operand, broadcast_dimensions); + 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. + 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) + .status()); + std::vector in_dim_size(out_dim_size.begin(), out_dim_size.end()); + for (int i = 0; i < broadcast_dimensions.size(); i++) { + in_dim_size[broadcast_dimensions[i]] = operand_shape.dimensions(i); + } + const auto& in_dim_shape = + ShapeUtil::MakeShape(operand_shape.element_type(), in_dim_size); + TF_ASSIGN_OR_RETURN( + XlaOp in_dim_broadcast, + InDimBroadcast(in_dim_shape, operand, broadcast_dimensions)); + + // If broadcast is not degenerate, return broadcasted result. + if (ShapeUtil::Equal(in_dim_shape, output_shape)) { + return in_dim_broadcast; + } + + // Otherwise handle degenerate broadcast case. + return AddBroadcastSequence(output_shape, in_dim_broadcast); }); } @@ -554,7 +711,7 @@ StatusOr XlaBuilder::Reshape(const Shape& shape, const XlaOp& operand) { TF_RETURN_IF_ERROR(first_error_); HloInstructionProto instr; - *instr.mutable_shape() = shape; + *instr.mutable_shape() = shape.ToProto(); return AddInstruction(std::move(instr), HloOpcode::kReshape, {operand}); } @@ -566,9 +723,9 @@ XlaOp XlaBuilder::Slice(const XlaOp& operand, HloInstructionProto instr; TF_ASSIGN_OR_RETURN(const Shape& operand_shape, GetShape(operand)); TF_ASSIGN_OR_RETURN( - *instr.mutable_shape(), - ShapeInference::InferSliceShape(operand_shape, start_indices, - limit_indices, strides)); + Shape shape, ShapeInference::InferSliceShape( + operand_shape, start_indices, limit_indices, strides)); + *instr.mutable_shape() = shape.ToProto(); for (int i = 0; i < start_indices.size(); i++) { auto* slice_config = instr.add_slice_dimensions(); slice_config->set_start(start_indices[i]); @@ -584,10 +741,10 @@ XlaOp XlaBuilder::SliceInDim(const XlaOp& operand, int64 start_index, int64 limit_index, int64 stride, int64 dimno) { return ReportErrorOrReturn([&]() -> StatusOr { TF_ASSIGN_OR_RETURN(const Shape& shape, GetShape(operand)); - std::vector starts(ShapeUtil::Rank(shape), 0); + std::vector starts(shape.rank(), 0); std::vector limits(shape.dimensions().begin(), shape.dimensions().end()); - std::vector strides(ShapeUtil::Rank(shape), 1); + std::vector strides(shape.rank(), 1); starts[dimno] = start_index; limits[dimno] = limit_index; strides[dimno] = stride; @@ -603,9 +760,10 @@ XlaOp XlaBuilder::DynamicSlice(const XlaOp& operand, const XlaOp& start_indices, TF_ASSIGN_OR_RETURN(const Shape& operand_shape, GetShape(operand)); TF_ASSIGN_OR_RETURN(const Shape& start_indices_shape, GetShape(start_indices)); - TF_ASSIGN_OR_RETURN(*instr.mutable_shape(), + TF_ASSIGN_OR_RETURN(Shape shape, ShapeInference::InferDynamicSliceShape( - operand_shape, start_indices_shape, slice_sizes)); + operand_shape, {start_indices_shape}, slice_sizes)); + *instr.mutable_shape() = shape.ToProto(); for (int64 size : slice_sizes) { instr.add_dynamic_slice_sizes(size); @@ -616,6 +774,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 { @@ -625,15 +811,41 @@ XlaOp XlaBuilder::DynamicUpdateSlice(const XlaOp& operand, const XlaOp& update, TF_ASSIGN_OR_RETURN(const Shape& update_shape, GetShape(update)); TF_ASSIGN_OR_RETURN(const Shape& start_indices_shape, GetShape(start_indices)); - TF_ASSIGN_OR_RETURN(*instr.mutable_shape(), - ShapeInference::InferDynamicUpdateSliceShape( - operand_shape, update_shape, start_indices_shape)); + TF_ASSIGN_OR_RETURN( + Shape shape, ShapeInference::InferDynamicUpdateSliceShape( + operand_shape, update_shape, {start_indices_shape})); + *instr.mutable_shape() = shape.ToProto(); return AddInstruction(std::move(instr), HloOpcode::kDynamicUpdateSlice, {operand, update, start_indices}); }); } +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 { @@ -643,9 +855,9 @@ XlaOp XlaBuilder::ConcatInDim(absl::Span operands, TF_ASSIGN_OR_RETURN(const auto& operand_shapes, GetOperandShapes(operands)); absl::c_transform(operand_shapes, std::back_inserter(operand_shape_ptrs), [](const Shape& shape) { return &shape; }); - TF_ASSIGN_OR_RETURN( - *instr.mutable_shape(), - ShapeInference::InferConcatOpShape(operand_shape_ptrs, dimension)); + TF_ASSIGN_OR_RETURN(Shape shape, ShapeInference::InferConcatOpShape( + operand_shape_ptrs, dimension)); + *instr.mutable_shape() = shape.ToProto(); instr.add_dimensions(dimension); @@ -662,10 +874,9 @@ XlaOp XlaBuilder::Pad(const XlaOp& operand, const XlaOp& padding_value, TF_ASSIGN_OR_RETURN(const Shape& padding_value_shape, GetShape(padding_value)); TF_ASSIGN_OR_RETURN( - *instr.mutable_shape(), - ShapeInference::InferPadShape(operand_shape, padding_value_shape, - padding_config)); - + Shape shape, ShapeInference::InferPadShape( + operand_shape, padding_value_shape, padding_config)); + *instr.mutable_shape() = shape.ToProto(); *instr.mutable_padding_config() = padding_config; return AddInstruction(std::move(instr), HloOpcode::kPad, @@ -678,7 +889,7 @@ XlaOp XlaBuilder::Reshape(const XlaOp& operand, absl::Span new_sizes) { return ReportErrorOrReturn([&]() -> StatusOr { TF_ASSIGN_OR_RETURN(const Shape& operand_shape, GetShape(operand)); - TF_ASSIGN_OR_RETURN(const Shape& shape, + TF_ASSIGN_OR_RETURN(const Shape shape, ShapeInference::InferReshapeShape( operand_shape, dimensions, new_sizes)); XlaOp transposed = IsIdentityPermutation(dimensions) @@ -691,7 +902,7 @@ XlaOp XlaBuilder::Reshape(const XlaOp& operand, XlaOp XlaBuilder::Reshape(const XlaOp& operand, absl::Span new_sizes) { return ReportErrorOrReturn([&]() -> StatusOr { - TF_ASSIGN_OR_RETURN(auto shape, GetShape(operand)); + TF_ASSIGN_OR_RETURN(Shape shape, GetShape(operand)); std::vector dimensions(shape.dimensions_size()); std::iota(dimensions.begin(), dimensions.end(), 0); return Reshape(operand, dimensions, new_sizes); @@ -724,7 +935,7 @@ XlaOp XlaBuilder::Collapse(const XlaOp& operand, VLOG(3) << "dims to collapse: " << absl::StrJoin(dimensions, ","); std::vector new_sizes; - for (int i = 0; i < ShapeUtil::Rank(original_shape); ++i) { + for (int i = 0; i < original_shape.rank(); ++i) { if (i <= dimensions.front() || i > dimensions.back()) { new_sizes.push_back(original_shape.dimensions(i)); } else { @@ -741,7 +952,7 @@ XlaOp XlaBuilder::Collapse(const XlaOp& operand, void XlaBuilder::Trace(const string& tag, const XlaOp& operand) { ReportErrorOrReturn([&]() -> StatusOr { HloInstructionProto instr; - *instr.mutable_shape() = ShapeUtil::MakeNil(); + *instr.mutable_shape() = ShapeUtil::MakeNil().ToProto(); *instr.mutable_literal() = LiteralUtil::CreateR1U8(tag).ToProto(); return AddInstruction(std::move(instr), HloOpcode::kTrace, {operand}); }); @@ -752,10 +963,9 @@ XlaOp XlaBuilder::Select(const XlaOp& pred, const XlaOp& on_true, return ReportErrorOrReturn([&]() -> StatusOr { TF_ASSIGN_OR_RETURN(const Shape& true_shape, GetShape(on_true)); TF_ASSIGN_OR_RETURN(const Shape& false_shape, GetShape(on_false)); - TF_RET_CHECK(ShapeUtil::IsTuple(true_shape) == - ShapeUtil::IsTuple(false_shape)); - HloOpcode opcode = ShapeUtil::IsTuple(true_shape) ? HloOpcode::kTupleSelect - : HloOpcode::kSelect; + TF_RET_CHECK(true_shape.IsTuple() == false_shape.IsTuple()); + HloOpcode opcode = + true_shape.IsTuple() ? HloOpcode::kTupleSelect : HloOpcode::kSelect; return TernaryOp(opcode, pred, on_true, on_false); }); } @@ -767,9 +977,10 @@ XlaOp XlaBuilder::Tuple(absl::Span elements) { TF_ASSIGN_OR_RETURN(const auto& operand_shapes, GetOperandShapes(elements)); absl::c_transform(operand_shapes, std::back_inserter(operand_shape_ptrs), [](const Shape& shape) { return &shape; }); - TF_ASSIGN_OR_RETURN(*instr.mutable_shape(), + TF_ASSIGN_OR_RETURN(const Shape shape, ShapeInference::InferVariadicOpShape( HloOpcode::kTuple, operand_shape_ptrs)); + *instr.mutable_shape() = shape.ToProto(); return AddInstruction(std::move(instr), HloOpcode::kTuple, elements); }); } @@ -778,13 +989,13 @@ XlaOp XlaBuilder::GetTupleElement(const XlaOp& tuple_data, int64 index) { return ReportErrorOrReturn([&]() -> StatusOr { HloInstructionProto instr; TF_ASSIGN_OR_RETURN(const Shape& tuple_shape, GetShape(tuple_data)); - if (!ShapeUtil::IsTuple(tuple_shape)) { + if (!tuple_shape.IsTuple()) { return InvalidArgument( "Operand to GetTupleElement() is not a tuple; got %s", ShapeUtil::HumanString(tuple_shape)); } *instr.mutable_shape() = - ShapeUtil::GetTupleElementShape(tuple_shape, index); + ShapeUtil::GetTupleElementShape(tuple_shape, index).ToProto(); instr.set_tuple_index(index); @@ -843,9 +1054,10 @@ 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)); - TF_ASSIGN_OR_RETURN(*instr.mutable_shape(), + TF_ASSIGN_OR_RETURN(Shape shape, ShapeInference::InferDotOpShape(lhs_shape, rhs_shape, dimension_numbers)); + *instr.mutable_shape() = shape.ToProto(); *instr.mutable_dot_dimension_numbers() = dimension_numbers; if (precision_config != nullptr) { *instr.mutable_precision_config() = *precision_config; @@ -857,13 +1069,13 @@ XlaOp XlaBuilder::DotGeneral(const XlaOp& lhs, const XlaOp& rhs, Status XlaBuilder::VerifyConvolution( const Shape& lhs_shape, const Shape& rhs_shape, const ConvolutionDimensionNumbers& dimension_numbers) const { - if (ShapeUtil::Rank(lhs_shape) != ShapeUtil::Rank(rhs_shape)) { + if (lhs_shape.rank() != rhs_shape.rank()) { return InvalidArgument( "Convolution arguments must have same number of " "dimensions. Got: %s and %s", ShapeUtil::HumanString(lhs_shape), ShapeUtil::HumanString(rhs_shape)); } - int num_dims = ShapeUtil::Rank(lhs_shape); + int num_dims = lhs_shape.rank(); if (num_dims < 2) { return InvalidArgument( "Convolution expects argument arrays with >= 3 dimensions. " @@ -901,27 +1113,29 @@ Status XlaBuilder::VerifyConvolution( XlaOp XlaBuilder::Conv(const XlaOp& lhs, const XlaOp& rhs, absl::Span window_strides, Padding padding, - int64 feature_group_count, + int64 feature_group_count, int64 batch_group_count, const PrecisionConfig* precision_config) { return ConvWithGeneralDimensions( lhs, rhs, window_strides, padding, CreateDefaultConvDimensionNumbers(window_strides.size()), - feature_group_count, precision_config); + feature_group_count, batch_group_count, precision_config); } XlaOp XlaBuilder::ConvWithGeneralPadding( const XlaOp& lhs, const XlaOp& rhs, absl::Span window_strides, absl::Span> padding, - int64 feature_group_count, const PrecisionConfig* precision_config) { + int64 feature_group_count, int64 batch_group_count, + const PrecisionConfig* precision_config) { return ConvGeneral(lhs, rhs, window_strides, padding, CreateDefaultConvDimensionNumbers(window_strides.size()), - feature_group_count, precision_config); + feature_group_count, batch_group_count, precision_config); } XlaOp XlaBuilder::ConvWithGeneralDimensions( const XlaOp& lhs, const XlaOp& rhs, absl::Span window_strides, Padding padding, const ConvolutionDimensionNumbers& dimension_numbers, - int64 feature_group_count, const PrecisionConfig* precision_config) { + int64 feature_group_count, int64 batch_group_count, + const PrecisionConfig* precision_config) { return ReportErrorOrReturn([&]() -> StatusOr { TF_ASSIGN_OR_RETURN(const Shape& lhs_shape, GetShape(lhs)); TF_ASSIGN_OR_RETURN(const Shape& rhs_shape, GetShape(rhs)); @@ -949,7 +1163,7 @@ XlaOp XlaBuilder::ConvWithGeneralDimensions( MakePadding(base_area_dimensions, window_dimensions, window_strides, padding), dimension_numbers, feature_group_count, - precision_config); + batch_group_count, precision_config); }); } @@ -957,10 +1171,11 @@ XlaOp XlaBuilder::ConvGeneral( const XlaOp& lhs, const XlaOp& rhs, absl::Span window_strides, absl::Span> padding, const ConvolutionDimensionNumbers& dimension_numbers, - int64 feature_group_count, const PrecisionConfig* precision_config) { + int64 feature_group_count, int64 batch_group_count, + const PrecisionConfig* precision_config) { return ConvGeneralDilated(lhs, rhs, window_strides, padding, {}, {}, dimension_numbers, feature_group_count, - precision_config); + batch_group_count, precision_config); } XlaOp XlaBuilder::ConvGeneralDilated( @@ -968,7 +1183,8 @@ XlaOp XlaBuilder::ConvGeneralDilated( absl::Span> padding, absl::Span lhs_dilation, absl::Span rhs_dilation, const ConvolutionDimensionNumbers& dimension_numbers, - int64 feature_group_count, const PrecisionConfig* precision_config) { + int64 feature_group_count, int64 batch_group_count, + const PrecisionConfig* precision_config) { return ReportErrorOrReturn([&]() -> StatusOr { HloInstructionProto instr; TF_ASSIGN_OR_RETURN(const Shape& lhs_shape, GetShape(lhs)); @@ -987,13 +1203,15 @@ XlaOp XlaBuilder::ConvGeneralDilated( MakeWindow(window_dimensions, window_strides, padding, lhs_dilation, rhs_dilation)); - TF_ASSIGN_OR_RETURN(*instr.mutable_shape(), - ShapeInference::InferConvolveShape( - lhs_shape, rhs_shape, feature_group_count, - instr.window(), dimension_numbers)); + TF_ASSIGN_OR_RETURN( + Shape shape, ShapeInference::InferConvolveShape( + lhs_shape, rhs_shape, feature_group_count, + batch_group_count, instr.window(), dimension_numbers)); + *instr.mutable_shape() = shape.ToProto(); *instr.mutable_convolution_dimension_numbers() = dimension_numbers; instr.set_feature_group_count(feature_group_count); + instr.set_batch_group_count(batch_group_count); if (precision_config != nullptr) { *instr.mutable_precision_config() = *precision_config; @@ -1063,10 +1281,9 @@ XlaOp XlaBuilder::Fft(const XlaOp& operand, const FftType fft_type, return ReportErrorOrReturn([&]() -> StatusOr { HloInstructionProto instr; TF_ASSIGN_OR_RETURN(const Shape& operand_shape, GetShape(operand)); - TF_ASSIGN_OR_RETURN( - *instr.mutable_shape(), - ShapeInference::InferFftShape(operand_shape, fft_type, fft_length)); - + TF_ASSIGN_OR_RETURN(Shape shape, ShapeInference::InferFftShape( + operand_shape, fft_type, fft_length)); + *instr.mutable_shape() = shape.ToProto(); instr.set_fft_type(fft_type); for (int64 i : fft_length) { instr.add_fft_length(i); @@ -1084,10 +1301,10 @@ XlaOp XlaBuilder::Infeed(const Shape& shape, const string& config) { } const Shape infeed_instruction_shape = ShapeUtil::MakeTupleShape({shape, ShapeUtil::MakeTokenShape()}); - *instr.mutable_shape() = infeed_instruction_shape; + *instr.mutable_shape() = infeed_instruction_shape.ToProto(); instr.set_infeed_config(config); - if (ShapeUtil::IsArray(shape) && sharding() && + if (shape.IsArray() && sharding() && sharding()->type() == OpSharding::Type::OpSharding_Type_OTHER) { // TODO(b/110793772): Support tiled array-shaped infeeds. return InvalidArgument( @@ -1105,7 +1322,7 @@ XlaOp XlaBuilder::Infeed(const Shape& shape, const string& config) { XlaOp token; auto make_token = [&]() { HloInstructionProto token_instr; - *token_instr.mutable_shape() = ShapeUtil::MakeTokenShape(); + *token_instr.mutable_shape() = ShapeUtil::MakeTokenShape().ToProto(); return AddInstruction(std::move(token_instr), HloOpcode::kAfterAll, {}); }; if (sharding()) { @@ -1144,7 +1361,7 @@ XlaOp XlaBuilder::Infeed(const Shape& shape, const string& config) { // TODO(b/80000000): Remove this when clients have been updated to handle // tokens. HloInstructionProto infeed_data; - *infeed_data.mutable_shape() = shape; + *infeed_data.mutable_shape() = shape.ToProto(); infeed_data.set_tuple_index(0); return AddInstruction(std::move(infeed_data), HloOpcode::kGetTupleElement, {infeed}); @@ -1160,10 +1377,10 @@ XlaOp XlaBuilder::InfeedWithToken(const XlaOp& token, const Shape& shape, } const Shape infeed_instruction_shape = ShapeUtil::MakeTupleShape({shape, ShapeUtil::MakeTokenShape()}); - *instr.mutable_shape() = infeed_instruction_shape; + *instr.mutable_shape() = infeed_instruction_shape.ToProto(); instr.set_infeed_config(config); - if (ShapeUtil::IsArray(shape) && sharding() && + if (shape.IsArray() && sharding() && sharding()->type() == OpSharding::Type::OpSharding_Type_OTHER) { // TODO(b/110793772): Support tiled array-shaped infeeds. return InvalidArgument( @@ -1185,7 +1402,7 @@ void XlaBuilder::Outfeed(const XlaOp& operand, const Shape& shape_with_layout, ReportErrorOrReturn([&]() -> StatusOr { HloInstructionProto instr; - *instr.mutable_shape() = ShapeUtil::MakeTokenShape(); + *instr.mutable_shape() = ShapeUtil::MakeTokenShape().ToProto(); // Check and set outfeed shape. if (!LayoutUtil::HasLayout(shape_with_layout)) { @@ -1198,14 +1415,14 @@ void XlaBuilder::Outfeed(const XlaOp& operand, const Shape& shape_with_layout, ShapeUtil::HumanStringWithLayout(shape_with_layout), ShapeUtil::HumanStringWithLayout(operand_shape)); } - *instr.mutable_outfeed_shape() = shape_with_layout; + *instr.mutable_outfeed_shape() = shape_with_layout.ToProto(); instr.set_outfeed_config(outfeed_config); // Outfeed takes a token as its second operand. Generate the token to pass // to the outfeed. HloInstructionProto token_instr; - *token_instr.mutable_shape() = ShapeUtil::MakeTokenShape(); + *token_instr.mutable_shape() = ShapeUtil::MakeTokenShape().ToProto(); TF_ASSIGN_OR_RETURN(XlaOp token, AddInstruction(std::move(token_instr), HloOpcode::kAfterAll, {})); @@ -1219,7 +1436,7 @@ void XlaBuilder::Outfeed(const XlaOp& operand, const Shape& shape_with_layout, // TODO(b/80000000): Remove this when clients have been updated to handle // tokens. HloInstructionProto tuple_instr; - *tuple_instr.mutable_shape() = ShapeUtil::MakeNil(); + *tuple_instr.mutable_shape() = ShapeUtil::MakeNil().ToProto(); // The dummy tuple should have no sharding. { @@ -1238,7 +1455,7 @@ XlaOp XlaBuilder::OutfeedWithToken(const XlaOp& operand, const XlaOp& token, return ReportErrorOrReturn([&]() -> StatusOr { HloInstructionProto instr; - *instr.mutable_shape() = ShapeUtil::MakeTokenShape(); + *instr.mutable_shape() = ShapeUtil::MakeTokenShape().ToProto(); // Check and set outfeed shape. if (!LayoutUtil::HasLayout(shape_with_layout)) { @@ -1251,7 +1468,7 @@ XlaOp XlaBuilder::OutfeedWithToken(const XlaOp& operand, const XlaOp& token, ShapeUtil::HumanStringWithLayout(shape_with_layout), ShapeUtil::HumanStringWithLayout(operand_shape)); } - *instr.mutable_outfeed_shape() = shape_with_layout; + *instr.mutable_outfeed_shape() = shape_with_layout.ToProto(); instr.set_outfeed_config(outfeed_config); @@ -1263,7 +1480,7 @@ XlaOp XlaBuilder::OutfeedWithToken(const XlaOp& operand, const XlaOp& token, XlaOp XlaBuilder::CreateToken() { return ReportErrorOrReturn([&]() -> StatusOr { HloInstructionProto instr; - *instr.mutable_shape() = ShapeUtil::MakeTokenShape(); + *instr.mutable_shape() = ShapeUtil::MakeTokenShape().ToProto(); return AddInstruction(std::move(instr), HloOpcode::kAfterAll); }); } @@ -1273,8 +1490,17 @@ XlaOp XlaBuilder::AfterAll(absl::Span tokens) { if (tokens.empty()) { return InvalidArgument("AfterAll requires at least one operand"); } + for (int i = 0; i < tokens.size(); ++i) { + const XlaOp& operand = tokens[i]; + TF_ASSIGN_OR_RETURN(const Shape& operand_shape, GetShape(operand)); + if (!operand_shape.IsToken()) { + return InvalidArgument( + "All operands to AfterAll must be tokens; operand %d has shape %s", + i, ShapeUtil::HumanString(operand_shape)); + } + } HloInstructionProto instr; - *instr.mutable_shape() = ShapeUtil::MakeTokenShape(); + *instr.mutable_shape() = ShapeUtil::MakeTokenShape().ToProto(); return AddInstruction(std::move(instr), HloOpcode::kAfterAll, tokens); }); } @@ -1291,7 +1517,7 @@ XlaOp XlaBuilder::CustomCall( "are reserved for internal use.", call_target_name); } - *instr.mutable_shape() = shape; + *instr.mutable_shape() = shape.ToProto(); instr.set_custom_call_target(call_target_name); instr.set_custom_call_opaque(opaque); if (operand_shapes_with_layout.has_value()) { @@ -1315,7 +1541,7 @@ XlaOp XlaBuilder::CustomCall( "constrained layout.", operand_num); } - *instr.add_operand_shapes_with_layout() = operand_shape; + *instr.add_operand_shapes_with_layout() = operand_shape.ToProto(); ++operand_num; } } @@ -1469,9 +1695,9 @@ XlaOp XlaBuilder::Transpose(const XlaOp& operand, return ReportErrorOrReturn([&]() -> StatusOr { HloInstructionProto instr; TF_ASSIGN_OR_RETURN(const Shape& operand_shape, GetShape(operand)); - TF_ASSIGN_OR_RETURN( - *instr.mutable_shape(), - ShapeInference::InferTransposeShape(operand_shape, permutation)); + TF_ASSIGN_OR_RETURN(Shape shape, ShapeInference::InferTransposeShape( + operand_shape, permutation)); + *instr.mutable_shape() = shape.ToProto(); for (int64 dim : permutation) { instr.add_dimensions(dim); } @@ -1484,9 +1710,9 @@ XlaOp XlaBuilder::Rev(const XlaOp& operand, return ReportErrorOrReturn([&]() -> StatusOr { HloInstructionProto instr; TF_ASSIGN_OR_RETURN(const Shape& operand_shape, GetShape(operand)); - TF_ASSIGN_OR_RETURN( - *instr.mutable_shape(), - ShapeInference::InferReverseShape(operand_shape, dimensions)); + TF_ASSIGN_OR_RETURN(Shape shape, ShapeInference::InferReverseShape( + operand_shape, dimensions)); + *instr.mutable_shape() = shape.ToProto(); for (int64 dim : dimensions) { instr.add_dimensions(dim); } @@ -1505,12 +1731,12 @@ XlaOp XlaBuilder::Sort(const XlaOp& keys, absl::Span values, GetOperandShapes(values)); absl::c_transform(values_shapes, std::back_inserter(operand_shape_ptrs), [](const Shape& shape) { return &shape; }); - TF_ASSIGN_OR_RETURN(*instr.mutable_shape(), - ShapeInference::InferVariadicOpShape( - HloOpcode::kSort, operand_shape_ptrs)); + 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)); - dimension = ShapeUtil::Rank(keys_shape) - 1; + dimension = keys_shape.rank() - 1; } instr.add_dimensions(dimension); std::vector operands{keys}; @@ -1529,9 +1755,9 @@ XlaOp XlaBuilder::ConvertElementType(const XlaOp& operand, return ReportErrorOrReturn([&]() -> StatusOr { HloInstructionProto instr; TF_ASSIGN_OR_RETURN(const Shape& operand_shape, GetShape(operand)); - TF_ASSIGN_OR_RETURN( - *instr.mutable_shape(), - ShapeInference::InferConvertShape(operand_shape, new_element_type)); + TF_ASSIGN_OR_RETURN(Shape shape, ShapeInference::InferConvertShape( + operand_shape, new_element_type)); + *instr.mutable_shape() = shape.ToProto(); return AddInstruction(std::move(instr), HloOpcode::kConvert, {operand}); }); } @@ -1541,9 +1767,9 @@ XlaOp XlaBuilder::BitcastConvertType(const XlaOp& operand, return ReportErrorOrReturn([&]() -> StatusOr { HloInstructionProto instr; TF_ASSIGN_OR_RETURN(const Shape& operand_shape, GetShape(operand)); - TF_ASSIGN_OR_RETURN( - *instr.mutable_shape(), - ShapeInference::InferConvertShape(operand_shape, new_element_type)); + TF_ASSIGN_OR_RETURN(Shape shape, ShapeInference::InferConvertShape( + operand_shape, new_element_type)); + *instr.mutable_shape() = shape.ToProto(); return AddInstruction(std::move(instr), HloOpcode::kBitcastConvert, {operand}); }); @@ -1575,17 +1801,17 @@ XlaOp XlaBuilder::Map(absl::Span operands, TF_ASSIGN_OR_RETURN(const ProgramShape& called_program_shape, computation.GetProgramShape()); TF_ASSIGN_OR_RETURN( - *instr.mutable_shape(), - ShapeInference::InferMapShape(operand_shape_ptrs, called_program_shape, - dimensions)); + Shape shape, ShapeInference::InferMapShape( + operand_shape_ptrs, called_program_shape, dimensions)); + *instr.mutable_shape() = shape.ToProto(); - const Shape& output_shape = instr.shape(); - const int64 output_rank = ShapeUtil::Rank(output_shape); + Shape output_shape(instr.shape()); + const int64 output_rank = output_shape.rank(); AddCalledComputation(computation, &instr); std::vector new_operands(operands.begin(), operands.end()); for (XlaOp& new_operand : new_operands) { TF_ASSIGN_OR_RETURN(Shape shape, GetShape(new_operand)); - const int64 rank = ShapeUtil::Rank(shape); + const int64 rank = shape.rank(); if (rank != output_rank) { TF_ASSIGN_OR_RETURN(new_operand, InDimBroadcast(output_shape, new_operand, {})); @@ -1622,7 +1848,7 @@ XlaOp XlaBuilder::RngOp(RandomDistribution distribution, } TF_RETURN_IF_ERROR(ShapeUtil::ValidateShapeWithOptionalLayout(shape)); - *instr.mutable_shape() = shape; + *instr.mutable_shape() = shape.ToProto(); instr.set_distribution(distribution); @@ -1650,10 +1876,10 @@ XlaOp XlaBuilder::While(const XlaComputation& condition, TF_ASSIGN_OR_RETURN(const auto& condition_program_shape, condition.GetProgramShape()); TF_ASSIGN_OR_RETURN(const Shape& init_shape, GetShape(init)); - TF_ASSIGN_OR_RETURN( - *instr.mutable_shape(), - ShapeInference::InferWhileShape(condition_program_shape, - body_program_shape, init_shape)); + TF_ASSIGN_OR_RETURN(Shape shape, ShapeInference::InferWhileShape( + condition_program_shape, + body_program_shape, init_shape)); + *instr.mutable_shape() = shape.ToProto(); // Body comes before condition computation in the vector. AddCalledComputation(body, &instr); AddCalledComputation(condition, &instr); @@ -1670,10 +1896,10 @@ XlaOp XlaBuilder::Gather(const XlaOp& input, const XlaOp& start_indices, TF_ASSIGN_OR_RETURN(const Shape& input_shape, GetShape(input)); TF_ASSIGN_OR_RETURN(const Shape& start_indices_shape, GetShape(start_indices)); - TF_ASSIGN_OR_RETURN( - *instr.mutable_shape(), - ShapeInference::InferGatherShape(input_shape, start_indices_shape, + TF_ASSIGN_OR_RETURN(Shape shape, ShapeInference::InferGatherShape( + input_shape, start_indices_shape, dimension_numbers, slice_sizes)); + *instr.mutable_shape() = shape.ToProto(); *instr.mutable_gather_dimension_numbers() = dimension_numbers; for (int64 bound : slice_sizes) { @@ -1698,10 +1924,11 @@ XlaOp XlaBuilder::Scatter(const XlaOp& input, const XlaOp& scatter_indices, TF_ASSIGN_OR_RETURN(const Shape& updates_shape, GetShape(updates)); TF_ASSIGN_OR_RETURN(const ProgramShape& to_apply_shape, update_computation.GetProgramShape()); - TF_ASSIGN_OR_RETURN(*instr.mutable_shape(), + TF_ASSIGN_OR_RETURN(Shape shape, ShapeInference::InferScatterShape( input_shape, scatter_indices_shape, updates_shape, to_apply_shape, dimension_numbers)); + *instr.mutable_shape() = shape.ToProto(); *instr.mutable_scatter_dimension_numbers() = dimension_numbers; @@ -1728,10 +1955,11 @@ XlaOp XlaBuilder::Conditional(const XlaOp& predicate, const XlaOp& true_operand, TF_ASSIGN_OR_RETURN(const ProgramShape& false_computation_shape, false_computation.GetProgramShape()); TF_ASSIGN_OR_RETURN( - *instr.mutable_shape(), + Shape shape, ShapeInference::InferConditionalShape( predicate_shape, true_operand_shape, false_operand_shape, true_computation_shape, false_computation_shape)); + *instr.mutable_shape() = shape.ToProto(); // The index of true_computation must be 0 and that of false computation // must be 1. @@ -1773,9 +2001,10 @@ XlaOp XlaBuilder::Reduce(absl::Span operands, [](const Shape& shape) { return &shape; }); TF_ASSIGN_OR_RETURN( - *instr.mutable_shape(), + Shape shape, ShapeInference::InferReduceShape( operand_shape_ptrs, dimensions_to_reduce, called_program_shape)); + *instr.mutable_shape() = shape.ToProto(); for (int64 dim : dimensions_to_reduce) { instr.add_dimensions(dim); @@ -1791,7 +2020,7 @@ XlaOp XlaBuilder::ReduceAll(const XlaOp& operand, const XlaOp& init_value, const XlaComputation& computation) { return ReportErrorOrReturn([&]() -> StatusOr { TF_ASSIGN_OR_RETURN(const Shape& operand_shape, GetShape(operand)); - std::vector all_dimnos(ShapeUtil::Rank(operand_shape)); + std::vector all_dimnos(operand_shape.rank()); std::iota(all_dimnos.begin(), all_dimnos.end(), 0); return Reduce(operand, init_value, computation, all_dimnos); }); @@ -1838,10 +2067,10 @@ XlaOp XlaBuilder::ReduceWindowWithGeneralPadding( MakeWindow(window_dimensions, window_strides, padding, /*lhs_dilation=*/base_dilations, /*rhs_dilation=*/window_dilations)); - TF_ASSIGN_OR_RETURN( - *instr.mutable_shape(), - ShapeInference::InferReduceWindowShape(operand_shape, init_shape, - instr.window(), to_apply_shape)); + TF_ASSIGN_OR_RETURN(Shape shape, ShapeInference::InferReduceWindowShape( + operand_shape, init_shape, + instr.window(), to_apply_shape)); + *instr.mutable_shape() = shape.ToProto(); AddCalledComputation(computation, &instr); return AddInstruction(std::move(instr), HloOpcode::kReduceWindow, @@ -1859,9 +2088,10 @@ XlaOp XlaBuilder::BatchNormTraining(const XlaOp& operand, const XlaOp& scale, TF_ASSIGN_OR_RETURN(const Shape& scale_shape, GetShape(scale)); TF_ASSIGN_OR_RETURN(const Shape& offset_shape, GetShape(offset)); TF_ASSIGN_OR_RETURN( - *instr.mutable_shape(), + Shape shape, ShapeInference::InferBatchNormTrainingShape( operand_shape, scale_shape, offset_shape, feature_index)); + *instr.mutable_shape() = shape.ToProto(); instr.set_epsilon(epsilon); instr.set_feature_index(feature_index); @@ -1883,10 +2113,11 @@ XlaOp XlaBuilder::BatchNormInference(const XlaOp& operand, const XlaOp& scale, TF_ASSIGN_OR_RETURN(const Shape& offset_shape, GetShape(offset)); TF_ASSIGN_OR_RETURN(const Shape& mean_shape, GetShape(mean)); TF_ASSIGN_OR_RETURN(const Shape& variance_shape, GetShape(variance)); - TF_ASSIGN_OR_RETURN(*instr.mutable_shape(), - ShapeInference::InferBatchNormInferenceShape( - operand_shape, scale_shape, offset_shape, - mean_shape, variance_shape, feature_index)); + TF_ASSIGN_OR_RETURN( + Shape shape, ShapeInference::InferBatchNormInferenceShape( + operand_shape, scale_shape, offset_shape, mean_shape, + variance_shape, feature_index)); + *instr.mutable_shape() = shape.ToProto(); instr.set_epsilon(epsilon); instr.set_feature_index(feature_index); @@ -1908,10 +2139,11 @@ XlaOp XlaBuilder::BatchNormGrad(const XlaOp& operand, const XlaOp& scale, TF_ASSIGN_OR_RETURN(const Shape& batch_mean_shape, GetShape(batch_mean)); TF_ASSIGN_OR_RETURN(const Shape& batch_var_shape, GetShape(batch_var)); TF_ASSIGN_OR_RETURN(const Shape& grad_output_shape, GetShape(grad_output)); - TF_ASSIGN_OR_RETURN(*instr.mutable_shape(), + TF_ASSIGN_OR_RETURN(Shape shape, ShapeInference::InferBatchNormGradShape( operand_shape, scale_shape, batch_mean_shape, batch_var_shape, grad_output_shape, feature_index)); + *instr.mutable_shape() = shape.ToProto(); instr.set_epsilon(epsilon); instr.set_feature_index(feature_index); @@ -1942,9 +2174,9 @@ XlaOp XlaBuilder::CrossReplicaSum( return ReportErrorOrReturn([&]() -> StatusOr { HloInstructionProto instr; TF_ASSIGN_OR_RETURN(const Shape& operand_shape, GetShape(operand)); - TF_ASSIGN_OR_RETURN( - *instr.mutable_shape(), - ShapeInference::InferCrossReplicaSumShape({&operand_shape})); + TF_ASSIGN_OR_RETURN(Shape shape, + ShapeInference::InferAllReduceShape({&operand_shape})); + *instr.mutable_shape() = shape.ToProto(); for (const ReplicaGroup& group : replica_groups) { *instr.add_replica_groups() = group; @@ -1956,8 +2188,7 @@ XlaOp XlaBuilder::CrossReplicaSum( AddCalledComputation(computation, &instr); - return AddInstruction(std::move(instr), HloOpcode::kCrossReplicaSum, - {operand}); + return AddInstruction(std::move(instr), HloOpcode::kAllReduce, {operand}); }); } @@ -1997,8 +2228,8 @@ XlaOp XlaBuilder::AllToAll(const XlaOp& operand, int64 split_dimension, absl::c_transform(slice_shapes, std::back_inserter(slice_shape_ptrs), [](const Shape& shape) { return &shape; }); TF_ASSIGN_OR_RETURN( - *instr.mutable_shape(), - ShapeInference::InferAllToAllTupleShape(slice_shape_ptrs)); + Shape shape, ShapeInference::InferAllToAllTupleShape(slice_shape_ptrs)); + *instr.mutable_shape() = shape.ToProto(); for (const ReplicaGroup& group : replica_groups) { *instr.add_replica_groups() = group; } @@ -2023,8 +2254,9 @@ XlaOp XlaBuilder::CollectivePermute( TF_ASSIGN_OR_RETURN(const Shape& operand_shape, GetShape(operand)); HloInstructionProto instr; TF_ASSIGN_OR_RETURN( - *instr.mutable_shape(), + Shape shape, ShapeInference::InferCollectivePermuteShape(operand_shape)); + *instr.mutable_shape() = shape.ToProto(); for (const auto& pair : source_target_pairs) { auto* proto_pair = instr.add_source_target_pairs(); @@ -2073,10 +2305,11 @@ XlaOp XlaBuilder::SelectAndScatterWithGeneralPadding( TF_ASSIGN_OR_RETURN(*instr.mutable_window(), MakeWindow(window_dimensions, window_strides, padding, /*lhs_dilation=*/{}, /*rhs_dilation=*/{})); - TF_ASSIGN_OR_RETURN(*instr.mutable_shape(), + TF_ASSIGN_OR_RETURN(Shape shape, ShapeInference::InferSelectAndScatterShape( operand_shape, select_shape, instr.window(), source_shape, init_shape, scatter_shape)); + *instr.mutable_shape() = shape.ToProto(); AddCalledComputation(select, &instr); AddCalledComputation(scatter, &instr); @@ -2091,9 +2324,10 @@ XlaOp XlaBuilder::ReducePrecision(const XlaOp& operand, const int exponent_bits, return ReportErrorOrReturn([&]() -> StatusOr { HloInstructionProto instr; TF_ASSIGN_OR_RETURN(const Shape& operand_shape, GetShape(operand)); - TF_ASSIGN_OR_RETURN(*instr.mutable_shape(), + TF_ASSIGN_OR_RETURN(Shape shape, ShapeInference::InferReducePrecisionShape( operand_shape, exponent_bits, mantissa_bits)); + *instr.mutable_shape() = shape.ToProto(); instr.set_exponent_bits(exponent_bits); instr.set_mantissa_bits(mantissa_bits); return AddInstruction(std::move(instr), HloOpcode::kReducePrecision, @@ -2108,7 +2342,7 @@ void XlaBuilder::Send(const XlaOp& operand, const ChannelHandle& handle) { // TODO(b/80000000): Remove this when clients have been updated to handle // tokens. HloInstructionProto token_instr; - *token_instr.mutable_shape() = ShapeUtil::MakeTokenShape(); + *token_instr.mutable_shape() = ShapeUtil::MakeTokenShape().ToProto(); TF_ASSIGN_OR_RETURN(XlaOp token, AddInstruction(std::move(token_instr), HloOpcode::kAfterAll, {})); @@ -2127,15 +2361,17 @@ XlaOp XlaBuilder::SendWithToken(const XlaOp& operand, const XlaOp& token, // token}. HloInstructionProto send_instr; TF_ASSIGN_OR_RETURN(const Shape& shape, GetShape(operand)); - *send_instr.mutable_shape() = ShapeUtil::MakeTupleShape( - {shape, ShapeUtil::MakeShape(U32, {}), ShapeUtil::MakeTokenShape()}); + *send_instr.mutable_shape() = + ShapeUtil::MakeTupleShape( + {shape, ShapeUtil::MakeShape(U32, {}), ShapeUtil::MakeTokenShape()}) + .ToProto(); send_instr.set_channel_id(handle.handle()); TF_ASSIGN_OR_RETURN(XlaOp send, AddInstruction(std::move(send_instr), HloOpcode::kSend, {operand, token})); HloInstructionProto send_done_instr; - *send_done_instr.mutable_shape() = ShapeUtil::MakeTokenShape(); + *send_done_instr.mutable_shape() = ShapeUtil::MakeTokenShape().ToProto(); send_done_instr.set_channel_id(handle.handle()); return AddInstruction(std::move(send_done_instr), HloOpcode::kSendDone, {send}); @@ -2149,7 +2385,7 @@ XlaOp XlaBuilder::Recv(const Shape& shape, const ChannelHandle& handle) { // TODO(b/80000000): Remove this when clients have been updated to handle // tokens. HloInstructionProto token_instr; - *token_instr.mutable_shape() = ShapeUtil::MakeTokenShape(); + *token_instr.mutable_shape() = ShapeUtil::MakeTokenShape().ToProto(); TF_ASSIGN_OR_RETURN(XlaOp token, AddInstruction(std::move(token_instr), HloOpcode::kAfterAll, {})); @@ -2160,7 +2396,7 @@ XlaOp XlaBuilder::Recv(const Shape& shape, const ChannelHandle& handle) { // TODO(b/80000000): Remove this when clients have been updated to handle // tokens. HloInstructionProto recv_data; - *recv_data.mutable_shape() = shape; + *recv_data.mutable_shape() = shape.ToProto(); recv_data.set_tuple_index(0); return AddInstruction(std::move(recv_data), HloOpcode::kGetTupleElement, {recv}); @@ -2177,15 +2413,18 @@ XlaOp XlaBuilder::RecvWithToken(const XlaOp& token, const Shape& shape, // Recv instruction produces a tuple of {receive buffer, U32 context, // token}. HloInstructionProto recv_instr; - *recv_instr.mutable_shape() = ShapeUtil::MakeTupleShape( - {shape, ShapeUtil::MakeShape(U32, {}), ShapeUtil::MakeTokenShape()}); + *recv_instr.mutable_shape() = + ShapeUtil::MakeTupleShape( + {shape, ShapeUtil::MakeShape(U32, {}), ShapeUtil::MakeTokenShape()}) + .ToProto(); recv_instr.set_channel_id(handle.handle()); TF_ASSIGN_OR_RETURN(XlaOp recv, AddInstruction(std::move(recv_instr), HloOpcode::kRecv, {token})); HloInstructionProto recv_done_instr; *recv_done_instr.mutable_shape() = - ShapeUtil::MakeTupleShape({shape, ShapeUtil::MakeTokenShape()}); + ShapeUtil::MakeTupleShape({shape, ShapeUtil::MakeTokenShape()}) + .ToProto(); recv_done_instr.set_channel_id(handle.handle()); return AddInstruction(std::move(recv_done_instr), HloOpcode::kRecvDone, {recv}); @@ -2207,7 +2446,7 @@ XlaOp XlaBuilder::SendToHost(const XlaOp& operand, const XlaOp& token, ShapeUtil::HumanStringWithLayout(operand_shape)); } // TODO(b/111544877): Support tuple shapes. - if (!ShapeUtil::IsArray(operand_shape)) { + if (!operand_shape.IsArray()) { return InvalidArgument("SendToHost only supports array shapes, shape: %s", ShapeUtil::HumanString(operand_shape)); } @@ -2219,9 +2458,11 @@ XlaOp XlaBuilder::SendToHost(const XlaOp& operand, const XlaOp& token, // Send instruction produces a tuple of {aliased operand, U32 context, // token}. HloInstructionProto send_instr; - *send_instr.mutable_shape() = ShapeUtil::MakeTupleShape( - {shape_with_layout, ShapeUtil::MakeShape(U32, {}), - ShapeUtil::MakeTokenShape()}); + *send_instr.mutable_shape() = + ShapeUtil::MakeTupleShape({shape_with_layout, + ShapeUtil::MakeShape(U32, {}), + ShapeUtil::MakeTokenShape()}) + .ToProto(); send_instr.set_channel_id(handle.handle()); send_instr.set_is_host_transfer(true); TF_ASSIGN_OR_RETURN(XlaOp send, @@ -2229,7 +2470,7 @@ XlaOp XlaBuilder::SendToHost(const XlaOp& operand, const XlaOp& token, {operand, token})); HloInstructionProto send_done_instr; - *send_done_instr.mutable_shape() = ShapeUtil::MakeTokenShape(); + *send_done_instr.mutable_shape() = ShapeUtil::MakeTokenShape().ToProto(); send_done_instr.set_channel_id(handle.handle()); send_done_instr.set_is_host_transfer(true); return AddInstruction(std::move(send_done_instr), HloOpcode::kSendDone, @@ -2245,7 +2486,7 @@ XlaOp XlaBuilder::RecvFromHost(const XlaOp& token, const Shape& shape, } // TODO(b/111544877): Support tuple shapes. - if (!ShapeUtil::IsArray(shape)) { + if (!shape.IsArray()) { return InvalidArgument( "RecvFromHost only supports array shapes, shape: %s", ShapeUtil::HumanString(shape)); @@ -2258,8 +2499,10 @@ XlaOp XlaBuilder::RecvFromHost(const XlaOp& token, const Shape& shape, // Recv instruction produces a tuple of {receive buffer, U32 context, // token}. HloInstructionProto recv_instr; - *recv_instr.mutable_shape() = ShapeUtil::MakeTupleShape( - {shape, ShapeUtil::MakeShape(U32, {}), ShapeUtil::MakeTokenShape()}); + *recv_instr.mutable_shape() = + ShapeUtil::MakeTupleShape( + {shape, ShapeUtil::MakeShape(U32, {}), ShapeUtil::MakeTokenShape()}) + .ToProto(); recv_instr.set_channel_id(handle.handle()); recv_instr.set_is_host_transfer(true); TF_ASSIGN_OR_RETURN(XlaOp recv, AddInstruction(std::move(recv_instr), @@ -2267,7 +2510,8 @@ XlaOp XlaBuilder::RecvFromHost(const XlaOp& token, const Shape& shape, HloInstructionProto recv_done_instr; *recv_done_instr.mutable_shape() = - ShapeUtil::MakeTupleShape({shape, ShapeUtil::MakeTokenShape()}); + ShapeUtil::MakeTupleShape({shape, ShapeUtil::MakeTokenShape()}) + .ToProto(); recv_done_instr.set_channel_id(handle.handle()); recv_done_instr.set_is_host_transfer(true); return AddInstruction(std::move(recv_done_instr), HloOpcode::kRecvDone, @@ -2275,6 +2519,19 @@ XlaOp XlaBuilder::RecvFromHost(const XlaOp& token, const Shape& shape, }); } +XlaOp XlaBuilder::GetDimensionSize(const XlaOp& operand, int64 dimension) { + return ReportErrorOrReturn([&]() -> StatusOr { + HloInstructionProto instr; + TF_ASSIGN_OR_RETURN(const auto& operand_shape, GetShape(operand)); + TF_ASSIGN_OR_RETURN(Shape shape, ShapeInference::InferGetDimensionSizeShape( + operand_shape, dimension)); + *instr.mutable_shape() = shape.ToProto(); + instr.add_dimensions(dimension); + return AddInstruction(std::move(instr), HloOpcode::kGetDimensionSize, + {operand}); + }); +} + StatusOr XlaBuilder::IsConstant(const XlaOp& operand) const { TF_RETURN_IF_ERROR(first_error_); @@ -2282,13 +2539,13 @@ 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; } StatusOr XlaBuilder::BuildConstantSubGraph( - const XlaOp& root_op) const { + const XlaOp& root_op) { TF_ASSIGN_OR_RETURN(bool is_constant, IsConstant(root_op)); if (!is_constant) { auto op_status = LookUpInstruction(root_op); @@ -2310,10 +2567,10 @@ StatusOr XlaBuilder::BuildConstantSubGraph( LookUpInstruction(root_op)); HloComputationProto entry; - entry.set_id(GetUniqueId()); // Give the computation a global unique id. - entry.set_name(StrCat(name_, entry.id(), "_compute_constant")); + SetProtoIdAndName(&entry, StrCat(name_, "_compute_constant"), kNameSeparator, + GetNextId()); entry.set_root_id(root->id()); - ProgramShape* program_shape = entry.mutable_program_shape(); + ProgramShapeProto* program_shape = entry.mutable_program_shape(); *program_shape->mutable_result() = root->shape(); // We use std::set to keep the instruction ids in ascending order (which is @@ -2329,21 +2586,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 = @@ -2451,7 +2745,7 @@ StatusOr XlaBuilder::AddInstruction(HloInstructionProto&& instr, absl::Span operands) { TF_RETURN_IF_ERROR(first_error_); - const int64 handle = GetUniqueId(); + const int64 handle = GetNextId(); instr.set_id(handle); instr.set_opcode(HloOpcodeString(opcode)); if (instr.name().empty()) { @@ -2482,9 +2776,50 @@ StatusOr XlaBuilder::AddInstruction(HloInstructionProto&& instr, void XlaBuilder::AddCalledComputation(const XlaComputation& computation, HloInstructionProto* instr) { - instr->add_called_computation_ids(computation.proto().entry_computation_id()); + absl::flat_hash_map remapped_ids; + std::vector imported_computations; + imported_computations.reserve(computation.proto().computations_size()); + // Before we import the computations by remapping IDs, and capturing the + // old->new mappings in remapped_ids. for (const HloComputationProto& e : computation.proto().computations()) { - embedded_.insert({e.id(), e}); + HloComputationProto new_computation(e); + int64 computation_id = GetNextId(); + remapped_ids[new_computation.id()] = computation_id; + SetProtoIdAndName(&new_computation, + GetBaseName(new_computation.name(), kNameSeparator), + kNameSeparator, computation_id); + for (auto& instruction : *new_computation.mutable_instructions()) { + int64 instruction_id = GetNextId(); + remapped_ids[instruction.id()] = instruction_id; + SetProtoIdAndName(&instruction, + GetBaseName(instruction.name(), kNameSeparator), + kNameSeparator, instruction_id); + } + new_computation.set_root_id(remapped_ids.at(new_computation.root_id())); + + imported_computations.push_back(std::move(new_computation)); + } + // Once we have imported all the computations, and captured all the ID + // mappings, we go back and fixup the IDs in the imported computations. + instr->add_called_computation_ids( + remapped_ids.at(computation.proto().entry_computation_id())); + for (auto& imported_computation : imported_computations) { + for (auto& instruction : *imported_computation.mutable_instructions()) { + for (auto& operand_id : *instruction.mutable_operand_ids()) { + operand_id = remapped_ids.at(operand_id); + } + for (auto& control_predecessor_id : + *instruction.mutable_control_predecessor_ids()) { + control_predecessor_id = remapped_ids.at(control_predecessor_id); + } + for (auto& called_computation_id : + *instruction.mutable_called_computation_ids()) { + called_computation_id = remapped_ids.at(called_computation_id); + } + } + + int64 computation_id = imported_computation.id(); + embedded_.insert({computation_id, std::move(imported_computation)}); } } @@ -2533,9 +2868,10 @@ XlaOp Broadcast(const XlaOp& operand, absl::Span broadcast_sizes) { return operand.builder()->Broadcast(operand, broadcast_sizes); } -XlaOp BroadcastInDim(const XlaOp& operand, const Shape& shape, +XlaOp BroadcastInDim(const XlaOp& operand, + const absl::Span out_dim_size, const absl::Span broadcast_dimensions) { - return operand.builder()->BroadcastInDim(operand, shape, + return operand.builder()->BroadcastInDim(operand, out_dim_size, broadcast_dimensions); } @@ -2574,12 +2910,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); @@ -2645,38 +2990,42 @@ XlaOp DotGeneral(const XlaOp& lhs, const XlaOp& rhs, XlaOp Conv(const XlaOp& lhs, const XlaOp& rhs, absl::Span window_strides, Padding padding, - int64 feature_group_count, const PrecisionConfig* precision_config) { + int64 feature_group_count, int64 batch_group_count, + const PrecisionConfig* precision_config) { return lhs.builder()->Conv(lhs, rhs, window_strides, padding, - feature_group_count, precision_config); + feature_group_count, batch_group_count, + precision_config); } XlaOp ConvWithGeneralPadding(const XlaOp& lhs, const XlaOp& rhs, absl::Span window_strides, absl::Span> padding, - int64 feature_group_count, + int64 feature_group_count, int64 batch_group_count, const PrecisionConfig* precision_config) { return lhs.builder()->ConvWithGeneralPadding( - lhs, rhs, window_strides, padding, feature_group_count, precision_config); + lhs, rhs, window_strides, padding, feature_group_count, batch_group_count, + precision_config); } XlaOp ConvWithGeneralDimensions( const XlaOp& lhs, const XlaOp& rhs, absl::Span window_strides, Padding padding, const ConvolutionDimensionNumbers& dimension_numbers, - int64 feature_group_count, const PrecisionConfig* precision_config) { + int64 feature_group_count, int64 batch_group_count, + const PrecisionConfig* precision_config) { return lhs.builder()->ConvWithGeneralDimensions( lhs, rhs, window_strides, padding, dimension_numbers, feature_group_count, - precision_config); + batch_group_count, precision_config); } XlaOp ConvGeneral(const XlaOp& lhs, const XlaOp& rhs, absl::Span window_strides, absl::Span> padding, const ConvolutionDimensionNumbers& dimension_numbers, - int64 feature_group_count, + int64 feature_group_count, int64 batch_group_count, const PrecisionConfig* precision_config) { return lhs.builder()->ConvGeneral(lhs, rhs, window_strides, padding, dimension_numbers, feature_group_count, - precision_config); + batch_group_count, precision_config); } XlaOp ConvGeneralDilated(const XlaOp& lhs, const XlaOp& rhs, @@ -2685,11 +3034,12 @@ XlaOp ConvGeneralDilated(const XlaOp& lhs, const XlaOp& rhs, absl::Span lhs_dilation, absl::Span rhs_dilation, const ConvolutionDimensionNumbers& dimension_numbers, - int64 feature_group_count, + int64 feature_group_count, int64 batch_group_count, const PrecisionConfig* precision_config) { return lhs.builder()->ConvGeneralDilated( lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, - dimension_numbers, feature_group_count, precision_config); + dimension_numbers, feature_group_count, batch_group_count, + precision_config); } XlaOp Fft(const XlaOp& operand, FftType fft_type, @@ -3087,4 +3437,8 @@ XlaOp Iota(XlaBuilder* builder, const Shape& shape, int64 iota_dimension) { return builder->Iota(shape, iota_dimension); } +XlaOp GetDimensionSize(const XlaOp& operand, int64 dimension) { + return operand.builder()->GetDimensionSize(operand, dimension); +} + } // namespace xla diff --git a/tensorflow/compiler/xla/client/xla_builder.h b/tensorflow/compiler/xla/client/xla_builder.h index 5747661c34b411bbf22575f9c1d9fe09aa32911f..3bd6d42363664721ee4c15c8dc4fc75a42d0591b 100644 --- a/tensorflow/compiler/xla/client/xla_builder.h +++ b/tensorflow/compiler/xla/client/xla_builder.h @@ -29,6 +29,7 @@ limitations under the License. #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/dynamic_parameter_binding.h" #include "tensorflow/compiler/xla/service/hlo.pb.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" #include "tensorflow/compiler/xla/shape_util.h" @@ -196,11 +197,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. @@ -216,7 +225,7 @@ class XlaBuilder { // compile-time constant (see `IsConstant`), returns an error. // // This will copy the needed ops/computations to the subgraph. - StatusOr BuildConstantSubGraph(const XlaOp& root_op) const; + StatusOr BuildConstantSubGraph(const XlaOp& root_op); // Returns the first error that was encountered while building the // computation. When an error is encountered, by default we return a vacuous @@ -263,35 +272,49 @@ class XlaBuilder { // evaluating the computation. StatusOr IsConstant(const XlaOp& operand) const; + // Sets up binding which indicates that the `target_dim_num` in the subshape + // `target_param_index` of parameter `target_param_num` is a dynamic dimension + // 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, + ShapeIndex dynamic_size_param_index, + int64 target_param_num, + ShapeIndex target_param_index, int64 target_dim_num); + + // Adds a new input/output alias. Since the input/ouput shape information are + // not available until the computation is built, and eventual error in the + // arguments of this API will be detected only at computation Build() time. + void SetUpAlias(const ShapeIndex& output_index, int64 param_number, + const ShapeIndex& param_index) { + input_output_aliases_.push_back({output_index, param_number, param_index}); + } + private: + // Describes an input/output alias as inserted by the SetUpAlias() API. + struct InputOutputAlias { + ShapeIndex output_index; + int64 param_number; + ShapeIndex param_index; + }; + // 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. - // Enqueues a "retrieve parameter value" instruction for a parameter that was - // passed to the computation. XlaOp Parameter(int64 parameter_number, const Shape& shape, const string& name); - // Enqueues a constant with the value of the given literal onto the - // computation. XlaOp ConstantLiteral(const LiteralSlice& literal); - // Enqueues a constant onto the computation. Methods are templated on the - // native host type (NativeT) which corresponds to a specific XLA - // PrimitiveType as given in the following table: - // - // Native Type PrimitiveType - // ----------------------------- - // bool PRED - // int32 S32 - // int64 S64 - // uint32 U32 - // uint64 U64 - // float F32 - // double F64 - // - // Note: not all primitive types defined in xla_data.proto have a - // corresponding native type yet. template XlaOp ConstantR0(NativeT value); template @@ -321,225 +344,107 @@ class XlaBuilder { template XlaOp ConstantR4FromArray4D(const Array4D& values); - // Enqueues a rank one constant (vector) onto the computation. The vector has - // size 'length' and every element has the value 'value'. template XlaOp ConstantR1(int64 length, NativeT value); - // Adds dimensions to an array by duplicating the data in the array. - // - // The new dimensions are inserted on the left, i.e. if - // broadcast_sizes has values {a0, ..., aN} and the operand shape - // has dimensions {b0, ..., bM} then the shape of the output has - // dimensions {a0, ..., aN, b0, ..., bM}. - // - // The new dimensions index into copies of the operand, i.e. - // - // output[i0, ..., iN, j0, ..., jM] = operand[j0, ..., jM] XlaOp Broadcast(const XlaOp& operand, absl::Span broadcast_sizes); - // Performs in-dimension-style broadcast. - // - // Operand specifies the input to be broadcast. "shape" is expected output - // shape. "broadcast_dimensions" are the dimensions to be broadcasting into. - // Dimension numbers in broadcast_dimensions map to individual dimensions - // of the operand, and specify what dimension of the output shape they - // should be broadcast. - // e.g. - // Say operand = [1, 2], i.e., a 1D tensor with 2 elements. - // and dimension of shape is [2,2]. - // Specifying {1} as brodcast_dimension will generate output - // [1 , 2] - // [1 , 2] - // On the other hand, specifying {0} as broadcast_dimension - // will generate output - // [1 , 1] - // [2 , 2] - XlaOp BroadcastInDim(const XlaOp& operand, const Shape& shape, + XlaOp BroadcastInDim(const XlaOp& operand, + const absl::Span out_dim_size, const absl::Span broadcast_dimensions); - // Enqueues a pad operation onto the computation that pads the given value on - // the edges as well as between the elements of the input. padding_config - // specifies the padding amount for each dimension. XlaOp Pad(const XlaOp& operand, const XlaOp& padding_value, const PaddingConfig& padding_config); - // Enqueues an operation onto the computation that flattens the operand based - // on the dimension order (major/slowest-varying to minor/fastest-varying) - // given, followed by reshaping it into the shape with the given dimension - // sizes (also major to minor). Conceptually, this is a limited form of - // "shape casting". XlaOp Reshape(const XlaOp& operand, absl::Span dimensions, absl::Span new_sizes); - // Enqueues an operation onto the computation that collapses the operand, from - // first to last dimension (C order), then reshapes it to the given dimension - // sizes. Conceptually, this is a limited form of "shape casting". XlaOp Reshape(const XlaOp& operand, absl::Span new_sizes); - // Wrapper for Reshape. - // Enqueues an operation to collapse the provided dimensions; e.g. an - // operand with dimensions {x=256, y=2, z=2, p=32} can be collapsed to - // {x=1024, y=32} by collapsing dims {0, 1, 2}. Collapsing dimensions must - // be a consecutive, in-order subsequence of the operand dimensions. - // - // Note that collapsing a single dimension does nothing: - // - // {256} collapsing {0} => {256} - // {1} collapsing {0} => {1} - // - // Collapsing multiple dimensions produces a single result dimension: - // - // {256, 2} collapsing {0,1} => {512} - // {256, 2, 3} collapsing {0,1} => {512, 3} - // - // This could potentially cause data to be moved -- it provides a more - // structured form of reshaping than an arbitrary Reshape operation. XlaOp Collapse(const XlaOp& operand, absl::Span dimensions); - // Enqueues a slice operation onto the computation that slices the operand - // from the start indices to the limit indices; e.g. - // - // x - // [ 0 1 2 3 ] - // y [ 4 5 6 7 ] => slice(start={1, 1}, limit={2, 3}) => [ 5 6 ] - // [ 8 9 a b ] - // - // Note that "limit" means up-to-but-not-including; i.e. [start, limit) in 1D - // range notation. - // The strides parameter determines the stride over the slice XlaOp Slice(const XlaOp& operand, absl::Span start_indices, absl::Span limit_indices, absl::Span strides); - // Enqueues a slice operation in a given dimension, taking all other - // dimensions as they are; e.g. if dimno is 1 from start_index 2 to - // limit_index 4 by 1, and the shape is f32[7,8,9], this call is short-hand - // for: - // - // array[:, 2:4:1, :] XlaOp SliceInDim(const XlaOp& operand, int64 start_index, int64 limit_index, int64 stride, int64 dimno); - // Enqueues a slice operation onto the computation that slices the 'operand' - // from dynamic start indices which are passed in 'start_indices'. - // 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'. - // Slice index calculations are computed modulo input dimension sizes to - // prevent dynamic start indices from generating out-of-bound array accesses. + 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); - // Enqueues a dynamic update slice operation onto the computation, which - // updates a slice of 'operand' with 'update' at dynamic 'start_indices'. - // The shape of 'update' determines the shape of the slice of 'operand' - // which is updated. - // The indices specified in 'start_indices' specify the offset of the slice - // of 'operand' which is updated. - // - // update = {10, 11} // calculated at runtime. - // [1 2 3] start = {1, 1} // calculated at runtime. [1 2 3 ] - // [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'. - // Slice index calculations are computed modulo update dimension sizes to - // prevent dynamic start indices from generating out-of-bound array accesses. + 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); - // Enqueues a concatenate instruction onto the computation. 'operands' must - // have >= 1 entry. XlaOp ConcatInDim(absl::Span operands, int64 dimension); - // Enqueue a tracing operation onto the computation; the computation will emit - // a logging message with the operand. void Trace(const string& tag, const XlaOp& operand); - // Enqueues a conditional-move-like select operation onto the computation; - // predicated on pred, selects between on_true and on_false. XlaOp Select(const XlaOp& pred, const XlaOp& on_true, const XlaOp& on_false); - // Enqueues a tuple-creation instruction onto the computation. XlaOp Tuple(absl::Span elements); - // Enqueues a tuple-element-get instruction onto the computation. XlaOp GetTupleElement(const XlaOp& tuple_data, int64 index); - // Enqueues an equal-to comparison instruction onto the computation. XlaOp Eq(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions = {}); - // Enqueues a not-equal comparison instruction onto the computation. XlaOp Ne(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions = {}); - // Enqueues a greater-or-equal comparison instruction onto the computation. XlaOp Ge(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions = {}); - // Enqueues a greater-than comparison instruction onto the computation. XlaOp Gt(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions = {}); - // Enqueues a less-than comparison instruction onto the computation. XlaOp Lt(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions = {}); - // Enqueues a less-or-equal comparison instruction onto the computation. XlaOp Le(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions = {}); - // Enqueues a dot instruction onto the computation. XlaOp Dot(const XlaOp& lhs, const XlaOp& rhs, const PrecisionConfig* precision_config = nullptr); - // Enqueues a general dot instruction onto the computation. XlaOp DotGeneral(const XlaOp& lhs, const XlaOp& rhs, const DotDimensionNumbers& dimension_numbers, const PrecisionConfig* precision_config = nullptr); - // Enqueues a convolution instruction onto the computation, which uses the - // default convolution dimension numbers. XlaOp Conv(const XlaOp& lhs, const XlaOp& rhs, absl::Span window_strides, Padding padding, - int64 feature_group_count = 1, + int64 feature_group_count = 1, int64 batch_group_count = 1, const PrecisionConfig* precision_config = nullptr); - // Enqueues a convolution instruction onto the computation, with the caller - // provided padding configuration in the format returned by MakePadding(). XlaOp ConvWithGeneralPadding( const XlaOp& lhs, const XlaOp& rhs, absl::Span window_strides, absl::Span> padding, - int64 feature_group_count = 1, + int64 feature_group_count = 1, int64 batch_group_count = 1, const PrecisionConfig* precision_config = nullptr); - // Enqueues a convolution instruction onto the computation, with the caller - // provided dimension numbers configuration. XlaOp ConvWithGeneralDimensions( const XlaOp& lhs, const XlaOp& rhs, absl::Span window_strides, Padding padding, const ConvolutionDimensionNumbers& dimension_numbers, - int64 feature_group_count = 1, + int64 feature_group_count = 1, int64 batch_group_count = 1, const PrecisionConfig* precision_config = nullptr); - // Enqueues a convolution instruction onto the computation, with the caller - // provided padding configuration as well as the dimension numbers. XlaOp ConvGeneral(const XlaOp& lhs, const XlaOp& rhs, absl::Span window_strides, absl::Span> padding, const ConvolutionDimensionNumbers& dimension_numbers, - int64 feature_group_count = 1, + int64 feature_group_count = 1, int64 batch_group_count = 1, const PrecisionConfig* precision_config = nullptr); - // Enqueues a convolution instruction onto the computation, with the caller - // provided padding configuration, dilation factors and dimension numbers. XlaOp ConvGeneralDilated(const XlaOp& lhs, const XlaOp& rhs, absl::Span window_strides, absl::Span> padding, @@ -547,82 +452,56 @@ class XlaBuilder { absl::Span rhs_dilation, const ConvolutionDimensionNumbers& dimension_numbers, int64 feature_group_count = 1, + int64 batch_group_count = 1, const PrecisionConfig* precision_config = nullptr); - // Enqueues an FFT instruction onto the computation, of the given type and - // with the given FFT length. XlaOp Fft(const XlaOp& operand, FftType fft_type, absl::Span fft_length); - // Enqueues an infeed instruction onto the computation, which writes data of - // the given shape to the infeed buffer of the device. XlaOp Infeed(const Shape& shape, const string& config = ""); XlaOp InfeedWithToken(const XlaOp& token, const Shape& shape, const string& config = ""); - // Enqueues an outfeed instruction onto the computation. This instruction - // generates outgoing data transfers for the given data. - // - // shape_with_layout communicates the laid out shape that we want to outfeed - // -- if !ShapeUtil::Compatible(GetShape(operand), shape_with_layout) an error - // will occur. void Outfeed(const XlaOp& operand, const Shape& shape_with_layout, const string& outfeed_config); XlaOp OutfeedWithToken(const XlaOp& operand, const XlaOp& token, const Shape& shape_with_layout, const string& outfeed_config); - // Enqueues a call instruction onto the computation. XlaOp Call(const XlaComputation& computation, absl::Span operands); - // Enqueues a custom call instruction onto the computation. XlaOp CustomCall( const string& call_target_name, absl::Span operands, const Shape& shape_with_layout, const string& opaque, absl::optional> operand_shapes_with_layout); - // The following methods enqueue element-wise binary arithmetic operations - // onto the computation. The shapes of the operands have to match unless one - // of the operands is a scalar, or an explicit broadcast dimension is given - // (see g3doc for more details). - - // Enqueues a complex compose instruction onto the computation. XlaOp Complex(const XlaOp& real, const XlaOp& imag, absl::Span broadcast_dimensions = {}); - // Enqueues a complex conjugate instruction onto the computation. XlaOp Conj(const XlaOp& operand); - // Enqueues an add instruction onto the computation. XlaOp Add(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions = {}); - // Enqueues a subtract instruction onto the computation. XlaOp Sub(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions = {}); - // Enqueues a multiply instruction onto the computation. XlaOp Mul(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions = {}); - // Enqueues a divide instruction onto the computation. XlaOp Div(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions = {}); - // Enqueues a remainder instruction onto the computation. XlaOp Rem(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions = {}); - // Enqueues a max instruction onto the computation. XlaOp Max(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions = {}); - // Enqueues a min instruction onto the computation. XlaOp Min(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions = {}); - // Element-wise logical operators XlaOp And(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions = {}); @@ -641,32 +520,23 @@ class XlaBuilder { XlaOp ShiftRightLogical(const XlaOp& lhs, const XlaOp& rhs, absl::Span broadcast_dimensions = {}); - // Reduces an array among the provided dimensions, given "computation" as a - // reduction operator. XlaOp Reduce(const XlaOp& operand, const XlaOp& init_value, const XlaComputation& computation, absl::Span dimensions_to_reduce); - // Reduces several arrays simultaneously among the provided dimensions, given - // "computation" as a reduction operator. XlaOp Reduce(absl::Span operands, absl::Span init_values, const XlaComputation& computation, absl::Span dimensions_to_reduce); - // Convenience wrapper around the above that reduces all the dimensions in the - // operand shape. XlaOp ReduceAll(const XlaOp& operand, const XlaOp& init_value, const XlaComputation& computation); - // Enqueues a windowed reduce instruction onto the computation. XlaOp ReduceWindow(const XlaOp& operand, const XlaOp& init_value, const XlaComputation& computation, absl::Span window_dimensions, absl::Span window_strides, Padding padding); - // As ReduceWindow(), but the padding is given in the format - // returned by MakePadding(). XlaOp ReduceWindowWithGeneralPadding( const XlaOp& operand, const XlaOp& init_value, const XlaComputation& computation, @@ -676,48 +546,22 @@ class XlaBuilder { absl::Span window_dilations, absl::Span> padding); - // Returns the sum of the operand value within each subgroup of replicas. All - // replicas supply one input to the sum and all replicas receive the resulting - // sum for each subgroup. XlaOp CrossReplicaSum(const XlaOp& operand, absl::Span replica_groups = {}); - // Enqueues an operation that do an AllReduce of the operand cross cores. Here - // AllReduce means doing a reduction on the input operand cross cores and then - // broadcasting the reduction result to those cores. The reduction function is - // defined by `computation`, which should be a commutative computation on - // scalars, e.g., add, min, or max. The way that AllReduce is applied is - // configured by: - // - // - `replica_groups`: each ReplicaGroup contains a list of replica id. If - // empty, all replicas belong to one group. Allreduce will be applied within - // subgroups. For example, we have 4 replicas, then - // replica_groups={{0,2},{1,3}} means, replica 0 and 2 are in subgroup 0, - // replica 1 and 3 are in subgroup 1. - // - // - `channel_id`: for Allreduce nodes from different modules, if they have - // the same channel_id, they will be 'Allreduce'd. If empty, Allreduce will - // not be applied cross modules. - // - // TODO(b/117564385): Rename this to AllReduce when it's ready to use. XlaOp CrossReplicaSum( const XlaOp& operand, const XlaComputation& computation, absl::Span replica_groups = {}, const absl::optional& channel_id = absl::nullopt); - // Enqueues an operation that do an Alltoall of the operand cross cores. XlaOp AllToAll(const XlaOp& operand, int64 split_dimension, int64 concat_dimension, int64 split_count, const std::vector& replica_groups); - // Enqueues an operation that do an CollectivePermute of the operand cross - // cores. XlaOp CollectivePermute( const XlaOp& operand, const std::vector>& source_target_pairs); - // Enqueues an operation that scatters the `source` array to the selected - // indices of each window. XlaOp SelectAndScatter(const XlaOp& operand, const XlaComputation& select, absl::Span window_dimensions, absl::Span window_strides, @@ -725,8 +569,6 @@ class XlaBuilder { const XlaOp& init_value, const XlaComputation& scatter); - // As SelectAndScatter(), but the padding is given in the format - // returned by MakePadding(). XlaOp SelectAndScatterWithGeneralPadding( const XlaOp& operand, const XlaComputation& select, absl::Span window_dimensions, @@ -734,222 +576,126 @@ class XlaBuilder { absl::Span> padding, const XlaOp& source, const XlaOp& init_value, const XlaComputation& scatter); - // Enqueues an abs instruction onto the computation. XlaOp Abs(const XlaOp& operand); - // Enqueues a atan2 instruction onto the computation. XlaOp Atan2(const XlaOp& y, const XlaOp& x, absl::Span broadcast_dimensions = {}); - // Enqueues an exp instruction onto the computation. XlaOp Exp(const XlaOp& operand); - // Enqueues an expm1 instruction onto the computation. XlaOp Expm1(const XlaOp& operand); - // Enqueues a floor instruction onto the computation. XlaOp Floor(const XlaOp& operand); - // Enqueues a ceil instruction onto the computation. XlaOp Ceil(const XlaOp& operand); - // Enqueues a round instruction onto the computation, rounding to nearest even - // with half-way cases rounding away from zero. XlaOp Round(const XlaOp& operand); - // Enqueues an log instruction (natural logarithm) onto the computation. XlaOp Log(const XlaOp& operand); - // Enqueues an log1p instruction (log(x+1)) onto the computation. XlaOp Log1p(const XlaOp& operand); - // Enqueues a sign instruction onto the computation. XlaOp Sign(const XlaOp& operand); - // Enqueues a count leading zeros instruction onto the computation. XlaOp Clz(const XlaOp& operand); - // Enqueues a cosine instruction onto the computation. XlaOp Cos(const XlaOp& operand); - // Enqueues a sine instruction onto the computation. XlaOp Sin(const XlaOp& operand); - // Enqueues a tanh instruction onto the computation. XlaOp Tanh(const XlaOp& operand); - // Enqueues a real-part instruction onto the computation. XlaOp Real(const XlaOp& operand); - // Enqueues an imaginary-part instruction onto the computation. XlaOp Imag(const XlaOp& operand); - // Enqueues a lhs^rhs computation onto the computation. 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. XlaOp IsFinite(const XlaOp& operand); - // Enqueues an iota operation onto the computation. XlaOp Iota(const Shape& shape, int64 iota_dimension); - // Enqueues a rank-1 iota operation onto the computation. XlaOp Iota(PrimitiveType type, int64 size); - // Enqueues a convert instruction onto the computation that changes the - // element type of the operand array to primitive_type. XlaOp ConvertElementType(const XlaOp& operand, PrimitiveType new_element_type); - // Enqueues a no-op instruction onto the computation that changes - // the element type of the operand array to primitive_type. The - // bit-widths of the source and destination element types must be - // identical. XlaOp BitcastConvertType(const XlaOp& operand, PrimitiveType new_element_type); - // Enqueues a negate instruction onto the computation. XlaOp Neg(const XlaOp& operand); - // Enqueues a transpose instruction onto the computation. XlaOp Transpose(const XlaOp& operand, absl::Span permutation); - // Enqueues a reverse instruction onto the computation. The order of the - // elements in the given dimensions is reversed (i.e., the element at index i - // is moved to index dimension_size - 1 - i). XlaOp Rev(const XlaOp& operand, absl::Span dimensions); - // Enqueues a sort (as increasing order) instruction onto the computation. - // If only keys are provided: - // * If the keys are an rank-1 tensor (an array), the result is a sorted array - // 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 - // will independently sort each row. If no dimension number is provided, then - // the last dimension is chosen by default. - // - // If both keys and values are provided: - // * The keys and all values must be tensors with the same dimensions. The - // element types of the tensors may be different. - // * 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. XlaOp Sort(const XlaOp& keys, absl::Span values = {}, int64 dimension = -1); - // Enqueues a clamp instruction onto the computation. XlaOp Clamp(const XlaOp& min, const XlaOp& operand, const XlaOp& max); - // Enqueues a map instruction onto the computation. XlaOp Map(absl::Span operands, const XlaComputation& computation, absl::Span dimensions, absl::Span static_operands = {}); - // Enqueues a N(mu, sigma) random number generation instruction onto the - // computation. XlaOp RngNormal(const XlaOp& mu, const XlaOp& sigma, const Shape& shape); - // Enqueues a U(a, b) random number generation instruction onto the - // computation. Returns values in the semi-open interval [a, b). XlaOp RngUniform(const XlaOp& a, const XlaOp& b, const Shape& shape); - // Enqueues a while node onto the computation. XlaOp While(const XlaComputation& condition, const XlaComputation& body, const XlaOp& init); - // Enqueues a conditional node onto the computation. XlaOp Conditional(const XlaOp& predicate, const XlaOp& true_operand, const XlaComputation& true_computation, const XlaOp& false_operand, const XlaComputation& false_computation); - // Enqueues a ReducePrecision node onto the computation. XlaOp ReducePrecision(const XlaOp& operand, const int exponent_bits, const int mantissa_bits); - // Enqueues a Gather node onto the computation. XlaOp Gather(const XlaOp& input, const XlaOp& start_indices, const GatherDimensionNumbers& dimension_numbers, absl::Span slice_sizes); - // Enqueues a Scatter node onto the computation. XlaOp Scatter(const XlaOp& input, const XlaOp& scatter_indices, const XlaOp& updates, const XlaComputation& update_computation, const ScatterDimensionNumbers& dimension_numbers); - // Enqueues a Send node onto the computation for device-to-device - // communication, to send the given operand to a Recv instruction that shares - // the same channel handle. void Send(const XlaOp& operand, const ChannelHandle& handle); XlaOp SendWithToken(const XlaOp& operand, const XlaOp& token, const ChannelHandle& handle); - // Enqueues a Send node which sends data to the host. XlaOp SendToHost(const XlaOp& operand, const XlaOp& token, const Shape& shape_with_layout, const ChannelHandle& handle); - // Enqueues a Recv node which receives data from the host. XlaOp RecvFromHost(const XlaOp& token, const Shape& shape, const ChannelHandle& handle); - // Enqueues an AfterAll operation with no operands producing a token-shaped - // value. XlaOp CreateToken(); - // Enqueues an AfterAll operation with no operands producing a token-shaped - // value. XlaOp AfterAll(absl::Span tokens); - // Enqueues a Recv node onto the computation. The data comes from a Send - // instruction that shares the same channel handle and its shape must - // be the same as the given shape. XlaOp Recv(const Shape& shape, const ChannelHandle& handle); XlaOp RecvWithToken(const XlaOp& token, const Shape& shape, const ChannelHandle& handle); - // Normalizes operand across spatial and batch dimensions for each feature. - // - // Returns a tuple (normalized, batch_mean, batch_var) where `normalized` - // is the normalized result and batch_mean and batch_var are the mean and - // variance, respectively, across batch for the operand. XlaOp BatchNormTraining(const XlaOp& operand, const XlaOp& scale, const XlaOp& offset, float epsilon, int64 feature_index); - // Normalizes operand across spatial and batch dimensions for each feature. - // - // `BatchNormInference` is equivalent to calling `BatchNormTraining` without - // computing `mean` and `variance` for each batch inside the operation. It - // uses the input `mean` and `variance` instead as estimated values. The - // purpose of this op is to reduce latency in inference, hence the name - // `BatchNormInference`. - // - // The output has the same shape as `operand`, and contains the normalized - // values for each batch. XlaOp BatchNormInference(const XlaOp& operand, const XlaOp& scale, const XlaOp& offset, const XlaOp& mean, const XlaOp& variance, float epsilon, int64 feature_index); - // Calculates the gradients of a batch norm op. - // - // The inputs `batch_mean` and `batch_var` represent the mean and variance - // across the batch. - // - // Returns a tuple of three elements: - // - grad_operand: Gradient with respect to input `operand` - // - grad_offset: Gradient with respect to input `offset` - // - grad_scale: Gradient with respect to input `scale` XlaOp BatchNormGrad(const XlaOp& operand, const XlaOp& scale, const XlaOp& batch_mean, const XlaOp& batch_var, const XlaOp& grad_output, float epsilon, int64 feature_index); + XlaOp GetDimensionSize(const XlaOp& operand, int64 dimension); + StatusOr AddInstruction(HloInstructionProto&& instr, HloOpcode opcode, absl::Span operands = {}); @@ -1000,7 +746,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. @@ -1016,8 +763,20 @@ class XlaBuilder { absl::Span lhs_dilation, absl::Span rhs_dilation) const; + int64 GetNextId() { return ++next_id_; } + + // Populates the module with the input/output alias information stored within + // the input_output_aliases vector. + static Status PopulateInputOutputAlias( + HloModuleProto* module, const ProgramShape& program_shape, + const std::vector& input_output_aliases); + string name_; // Name to use for the built computation. + // The next sequential ID for every instruction/computation contained within + // this computation. + int64 next_id_ = 0; + // The first error encountered while building the computation. // This is OK until the first error is encountered. Status first_error_; @@ -1028,6 +787,12 @@ class XlaBuilder { // The instructions of this computation. std::vector instructions_; + // Dynamic parameter configuration of this computation. + DynamicParameterBinding dynamic_parameter_binding_; + + // Holds the input/output alias information populated by the SetUpAlias() API. + std::vector input_output_aliases_; + // A map from XlaOp::Handle to the index in the instructions_ vector where the // instruction is held. absl::flat_hash_map handle_to_index_; @@ -1105,7 +870,7 @@ class XlaBuilder { absl::Span broadcast_sizes); friend XlaOp BroadcastInDim( - const XlaOp& operand, const Shape& shape, + const XlaOp& operand, const absl::Span out_dim_size, const absl::Span broadcast_dimensions); friend XlaOp Pad(const XlaOp& operand, const XlaOp& padding_value, @@ -1129,9 +894,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); @@ -1161,23 +931,25 @@ class XlaBuilder { const PrecisionConfig* precision_config); friend XlaOp Conv(const XlaOp& lhs, const XlaOp& rhs, absl::Span window_strides, Padding padding, - int64 feature_group_count, + int64 feature_group_count, int64 batch_group_count, const PrecisionConfig* precision_config); friend XlaOp ConvWithGeneralPadding( const XlaOp& lhs, const XlaOp& rhs, absl::Span window_strides, absl::Span> padding, - int64 feature_group_count, const PrecisionConfig* precision_config); + int64 feature_group_count, int64 batch_group_count, + const PrecisionConfig* precision_config); friend XlaOp ConvWithGeneralDimensions( const XlaOp& lhs, const XlaOp& rhs, absl::Span window_strides, Padding padding, const ConvolutionDimensionNumbers& dimension_numbers, - int64 feature_group_count, const PrecisionConfig* precision_config); + int64 feature_group_count, int64 batch_group_count, + const PrecisionConfig* precision_config); friend XlaOp ConvGeneral(const XlaOp& lhs, const XlaOp& rhs, absl::Span window_strides, absl::Span> padding, const ConvolutionDimensionNumbers& dimension_numbers, - int64 feature_group_count, + int64 feature_group_count, int64 batch_group_count, const PrecisionConfig* precision_config); friend XlaOp ConvGeneralDilated( const XlaOp& lhs, const XlaOp& rhs, @@ -1186,7 +958,8 @@ class XlaBuilder { absl::Span lhs_dilation, absl::Span rhs_dilation, const ConvolutionDimensionNumbers& dimension_numbers, - int64 feature_group_count, const PrecisionConfig* precision_config); + int64 feature_group_count, int64 batch_group_count, + const PrecisionConfig* precision_config); friend XlaOp Fft(const XlaOp& operand, FftType fft_type, absl::Span fft_length); friend XlaOp Infeed(XlaBuilder* builder, const Shape& shape, @@ -1366,6 +1139,8 @@ class XlaBuilder { const string& outfeed_config); friend XlaOp CreateToken(XlaBuilder* builder); friend XlaOp AfterAll(XlaBuilder* builder, absl::Span tokens); + + friend XlaOp GetDimensionSize(const XlaOp& operand, int64 dimension); }; // RAII-style object: sets the current sharding assignment in builder on @@ -1400,6 +1175,7 @@ class XlaScopedShardingAssignment { // Free functions for building XlaOps. The intention is that these will // become the public API for building XlaOps rather than calling methods on // XlaBuilder directly. +// // Enqueues a "retrieve parameter value" instruction for a parameter that was // passed to the computation. @@ -1480,24 +1256,23 @@ XlaOp ConstantR1(XlaBuilder* builder, int64 length, NativeT value); // output[i0, ..., iN, j0, ..., jM] = operand[j0, ..., jM] XlaOp Broadcast(const XlaOp& operand, absl::Span broadcast_sizes); -// Performs in-dimension-style broadcast. +// This op broadcasts the `operand` to an output with the given `shape`. +// `broadcast_dimensions` are the dimensions to be broadcasting into, i.e., the +// i'th dimension of the operand is mapped to the broadcast_dimensions[i]'th +// dimension of the output. This also requires that the i'th input dimension is +// either 1 or is the same as the output dimension it's broadcasting into. // -// Operand specifies the input to be broadcast. "shape" is expected output -// shape. "broadcast_dimensions" are the dimensions to be broadcasting into. -// Dimension numbers in broadcast_dimensions map to individual dimensions -// of the operand, and specify what dimension of the output shape they -// should be broadcast. -// e.g. -// Say operand = [1, 2], i.e., a 1D tensor with 2 elements. -// and dimension of shape is [2,2]. -// Specifying {1} as brodcast_dimension will generate output -// [1 , 2] -// [1 , 2] -// On the other hand, specifying {0} as broadcast_dimension -// will generate output -// [1 , 1] -// [2 , 2] -XlaOp BroadcastInDim(const XlaOp& operand, const Shape& shape, +// For example, say operand = {1, 2}, i.e., a 1D tensor in shape s32[2]; the +// output shape is s32[2,2]: +// - Specifying {1} as brodcast_dimension will generate output +// {{1, 2}, +// {1, 2}} +// - On the other hand, specifying {0} as broadcast_dimension +// will generate output +// {{1 , 1}, +// {2 , 2}} +XlaOp BroadcastInDim(const XlaOp& operand, + const absl::Span out_dim_size, const absl::Span broadcast_dimensions); // Enqueues a pad operation onto the computation that pads the given value on @@ -1568,10 +1343,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); @@ -1587,10 +1367,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); @@ -1650,7 +1435,7 @@ XlaOp DotGeneral(const XlaOp& lhs, const XlaOp& rhs, // default convolution dimension numbers. XlaOp Conv(const XlaOp& lhs, const XlaOp& rhs, absl::Span window_strides, Padding padding, - int64 feature_group_count = 1, + int64 feature_group_count = 1, int64 batch_group_count = 1, const PrecisionConfig* precision_config = nullptr); // Enqueues a convolution instruction onto the computation, with the caller @@ -1659,6 +1444,7 @@ XlaOp ConvWithGeneralPadding(const XlaOp& lhs, const XlaOp& rhs, absl::Span window_strides, absl::Span> padding, int64 feature_group_count = 1, + int64 batch_group_count = 1, const PrecisionConfig* precision_config = nullptr); // Enqueues a convolution instruction onto the computation, with the caller @@ -1666,7 +1452,7 @@ XlaOp ConvWithGeneralPadding(const XlaOp& lhs, const XlaOp& rhs, XlaOp ConvWithGeneralDimensions( const XlaOp& lhs, const XlaOp& rhs, absl::Span window_strides, Padding padding, const ConvolutionDimensionNumbers& dimension_numbers, - int64 feature_group_count = 1, + int64 feature_group_count = 1, int64 batch_group_count = 1, const PrecisionConfig* precision_config = nullptr); // Enqueues a convolution instruction onto the computation, with the caller @@ -1675,7 +1461,7 @@ XlaOp ConvGeneral(const XlaOp& lhs, const XlaOp& rhs, absl::Span window_strides, absl::Span> padding, const ConvolutionDimensionNumbers& dimension_numbers, - int64 feature_group_count = 1, + int64 feature_group_count = 1, int64 batch_group_count = 1, const PrecisionConfig* precision_config = nullptr); // Enqueues a convolution instruction onto the computation, with the caller @@ -1687,6 +1473,7 @@ XlaOp ConvGeneralDilated(const XlaOp& lhs, const XlaOp& rhs, absl::Span rhs_dilation, const ConvolutionDimensionNumbers& dimension_numbers, int64 feature_group_count = 1, + int64 batch_group_count = 1, const PrecisionConfig* precision_config = nullptr); // Enqueues an FFT instruction onto the computation, of the given type and @@ -2142,7 +1929,12 @@ XlaOp BatchNormGrad(const XlaOp& operand, const XlaOp& scale, const XlaOp& grad_output, float epsilon, int64 feature_index); +// Returns the size of the given dimension of the operand. The operand must be +// array shaped. +XlaOp GetDimensionSize(const XlaOp& operand, int64 dimension); + // Implementation details below this point. +// template XlaOp XlaBuilder::ConstantR0(NativeT value) { diff --git a/tensorflow/compiler/xla/client/xla_builder_test.cc b/tensorflow/compiler/xla/client/xla_builder_test.cc index 7c37ed00cd3dcc214fb0b36c0161d3c39a5bf8c8..098165000a29cb28cb0ef906dbdb1ff9ae2f24e8 100644 --- a/tensorflow/compiler/xla/client/xla_builder_test.cc +++ b/tensorflow/compiler/xla/client/xla_builder_test.cc @@ -18,13 +18,14 @@ limitations under the License. #include #include "tensorflow/compiler/xla/client/xla_computation.h" -#include "tensorflow/compiler/xla/legacy_flags/debug_options_flags.h" +#include "tensorflow/compiler/xla/debug_options_flags.h" #include "tensorflow/compiler/xla/service/hlo_matchers.h" #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/test.h" #include "tensorflow/compiler/xla/test_helpers.h" +#include "tensorflow/compiler/xla/util.h" #include "tensorflow/compiler/xla/xla_data.pb.h" namespace xla { @@ -39,22 +40,24 @@ 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( - proto, legacy_flags::GetDebugOptionsFromFlags())); + proto, GetDebugOptionsFromFlags())); return HloModule::CreateFromProto(proto, config); } // 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( - proto, legacy_flags::GetDebugOptionsFromFlags())); + proto, GetDebugOptionsFromFlags())); return HloModule::CreateFromProto(proto, config); } @@ -264,6 +267,26 @@ TEST_F(XlaBuilderTest, BinopHasInDimAndDegenerateBroadcast) { op::Broadcast(op::Reshape(op::Parameter(1))))); } +TEST_F(XlaBuilderTest, BroadcastInDim) { + XlaBuilder b(TestName()); + auto x = Parameter(&b, 0, ShapeUtil::MakeShape(F32, {2, 3}), "x"); + BroadcastInDim(x, {2, 4, 3}, + /*broadcast_dimensions=*/{0, 2}); + TF_ASSERT_OK_AND_ASSIGN(auto module, BuildHloModule(&b)); + auto root = module->entry_computation()->root_instruction(); + EXPECT_THAT(root, op::Broadcast()); +} + +TEST_F(XlaBuilderTest, BroadcastInDimWithDegeneratedDim) { + XlaBuilder b(TestName()); + auto x = Parameter(&b, 0, ShapeUtil::MakeShape(F32, {2, 1, 4}), "x"); + BroadcastInDim(x, {2, 3, 4}, + /*broadcast_dimensions=*/{0, 1, 2}); + TF_ASSERT_OK_AND_ASSIGN(auto module, BuildHloModule(&b)); + EXPECT_THAT(module->entry_computation()->root_instruction(), + op::Broadcast(op::Reshape(op::Broadcast()))); +} + TEST_F(XlaBuilderTest, OperandFromWrongBuilder) { XlaBuilder b1("b1"); auto p0 = Parameter(&b1, 0, ShapeUtil::MakeShape(F32, {}), "p0"); @@ -329,6 +352,15 @@ TEST_F(XlaBuilderTest, CollectivePermute) { EXPECT_EQ(root->opcode(), HloOpcode::kCollectivePermute); } +TEST_F(XlaBuilderTest, GetDimensionSize) { + XlaBuilder b(TestName()); + auto x = Parameter(&b, 0, ShapeUtil::MakeShape(F32, {5, 7}), "x"); + GetDimensionSize(x, 1); + TF_ASSERT_OK_AND_ASSIGN(auto module, BuildHloModule(&b)); + auto root = module->entry_computation()->root_instruction(); + EXPECT_EQ(root->opcode(), HloOpcode::kGetDimensionSize); +} + TEST_F(XlaBuilderTest, ReportError) { XlaBuilder b(TestName()); auto x = Parameter(&b, 0, ShapeUtil::MakeShape(F32, {5, 7}), "x"); @@ -396,5 +428,472 @@ TEST_F(XlaBuilderTest, BuildWithSpecificRootWithWrongBuilder) { ::testing::HasSubstr("root operation is not in this computation")); } +TEST_F(XlaBuilderTest, ProtoMatches) { + std::vector computations; + for (int i = 0; i < 2; ++i) { + XlaBuilder b_call("the_only_to_apply"); + auto p0 = Parameter(&b_call, 0, ShapeUtil::MakeShape(F32, {}), "p0"); + auto p1 = Parameter(&b_call, 1, ShapeUtil::MakeShape(F32, {}), "p1"); + Add(p0, Add(p1, p0)); + TF_ASSERT_OK_AND_ASSIGN(auto call, b_call.Build()); + XlaBuilder b(TestName()); + auto x = Parameter(&b, 0, ShapeUtil::MakeShape(F32, {}), "x"); + auto y = Parameter(&b, 1, ShapeUtil::MakeShape(F32, {}), "y"); + auto one = ConstantR0(&b, 1); + auto two = ConstantR0(&b, 2); + Add(Call(&b, call, {x, y}), Call(&b, call, {one, two})); + computations.push_back(b.Build().ValueOrDie()); + } + auto c0_string = computations[0].proto().SerializeAsString(); + auto c1_string = computations[1].proto().SerializeAsString(); + 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, 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)}); + Status status = b.Build().status(); + ASSERT_IS_NOT_OK(status); + EXPECT_THAT(status.error_message(), + ::testing::HasSubstr("All operands to AfterAll must be tokens")); +} + +TEST_F(XlaBuilderTest, CheckInputOutputAlias) { + XlaBuilder b(TestName()); + auto p0 = Parameter(&b, 0, ShapeUtil::MakeShape(F32, {8, 4}), "p0"); + auto p1 = Parameter(&b, 1, ShapeUtil::MakeShape(F32, {8, 4}), "p1"); + auto add = Add(p0, p1); + auto sub = Sub(p0, p1); + auto root = Tuple(&b, {add, sub}); + + b.SetUpAlias({1}, 0, {}); + b.SetUpAlias({0}, 1, {}); + + TF_ASSERT_OK_AND_ASSIGN(auto module, BuildHloModule(&b, root)); + + const HloInputOutputAliasConfig& config = module->input_output_alias_config(); + EXPECT_TRUE(config.ParameterHasAlias(0, {})); + EXPECT_TRUE(config.ParameterHasAlias(1, {})); + + auto alias_p0 = config.GetAliasedOutput(0, {}); + ASSERT_TRUE(alias_p0.has_value()); + EXPECT_EQ(*alias_p0, ShapeIndex({1})); + + auto alias_p1 = config.GetAliasedOutput(1, {}); + ASSERT_TRUE(alias_p1.has_value()); + EXPECT_EQ(*alias_p1, ShapeIndex({0})); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/client/xla_computation.cc b/tensorflow/compiler/xla/client/xla_computation.cc index c9870b65b91c1ebd7d44143faf215a2d5c2a2fc5..f317892c12529b2ee8a81788f6bbcae3b3d6489d 100644 --- a/tensorflow/compiler/xla/client/xla_computation.cc +++ b/tensorflow/compiler/xla/client/xla_computation.cc @@ -25,7 +25,7 @@ namespace xla { StatusOr XlaComputation::GetProgramShape() const { TF_RET_CHECK(proto_.has_host_program_shape()); - return proto_.host_program_shape(); + return ProgramShape(proto_.host_program_shape()); } StatusOr> XlaComputation::Snapshot() const { diff --git a/tensorflow/compiler/xla/client/xla_computation.h b/tensorflow/compiler/xla/client/xla_computation.h index 71598ef8b296a760b0ee818fce0a59aed5cfc6b4..3ccbfb28bd0c5939ee40878e9cc298688882ac62 100644 --- a/tensorflow/compiler/xla/client/xla_computation.h +++ b/tensorflow/compiler/xla/client/xla_computation.h @@ -19,6 +19,7 @@ limitations under the License. #include #include "tensorflow/compiler/xla/service/hlo.pb.h" +#include "tensorflow/compiler/xla/shape.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/xla_data.pb.h" diff --git a/tensorflow/compiler/xla/legacy_flags/debug_options_flags.cc b/tensorflow/compiler/xla/debug_options_flags.cc similarity index 83% rename from tensorflow/compiler/xla/legacy_flags/debug_options_flags.cc rename to tensorflow/compiler/xla/debug_options_flags.cc index 3ed3afcfcede20fbf5c7d4f004378817febeb4c7..a9a91648ac377987e7f226116e11c9c697ace103 100644 --- a/tensorflow/compiler/xla/legacy_flags/debug_options_flags.cc +++ b/tensorflow/compiler/xla/debug_options_flags.cc @@ -13,60 +13,58 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/compiler/xla/legacy_flags/debug_options_flags.h" +#include "tensorflow/compiler/xla/debug_options_flags.h" #include // NOLINT(build/c++11): only using std::call_once, not mutex. #include #include "absl/strings/str_split.h" -#include "tensorflow/compiler/xla/legacy_flags/debug_options_parsers.h" -#include "tensorflow/compiler/xla/legacy_flags/parse_flags_from_env.h" +#include "tensorflow/compiler/xla/debug_options_parsers.h" +#include "tensorflow/compiler/xla/parse_flags_from_env.h" namespace xla { -namespace legacy_flags { -namespace { - -DebugOptions* flag_values; -std::vector* flag_objects; -std::once_flag flags_init; - -void SetDebugOptionsDefaults(DebugOptions* flags) { - flags->set_xla_llvm_enable_alias_scope_metadata(true); - flags->set_xla_llvm_enable_noalias_metadata(true); - flags->set_xla_llvm_enable_invariant_load_metadata(true); - flags->set_xla_llvm_disable_expensive_passes(false); - flags->set_xla_backend_optimization_level(3); - flags->set_xla_cpu_multi_thread_eigen(true); - flags->set_xla_gpu_cuda_data_dir("./cuda_sdk_lib"); - flags->set_xla_eliminate_hlo_implicit_broadcast(true); +DebugOptions DefaultDebugOptionsIgnoringFlags() { + DebugOptions opts; + opts.set_xla_llvm_enable_alias_scope_metadata(true); + opts.set_xla_llvm_enable_noalias_metadata(true); + opts.set_xla_llvm_enable_invariant_load_metadata(true); + opts.set_xla_llvm_disable_expensive_passes(false); + opts.set_xla_backend_optimization_level(3); + opts.set_xla_cpu_multi_thread_eigen(true); + opts.set_xla_gpu_cuda_data_dir("./cuda_sdk_lib"); + opts.set_xla_eliminate_hlo_implicit_broadcast(true); + opts.set_xla_hlo_dump_as_html(false); #ifdef INTEL_MKL - flags->set_xla_cpu_use_mkl_dnn(true); + opts.set_xla_cpu_use_mkl_dnn(true); #endif // INTEL_MKL - flags->set_xla_gpu_max_kernel_unroll_factor(4); + opts.set_xla_gpu_max_kernel_unroll_factor(4); // Set cudnn batchnorm off by default; it does not provide a performance win // on average. - flags->set_xla_gpu_use_cudnn_batchnorm(false); + opts.set_xla_gpu_use_cudnn_batchnorm(false); // Run all GPU work on one stream by default. Using multiple streams // increases memory usage and we lack strong motivating benchmarks for tuning // the heuristics needed to decide when to run on multiple streams. See // b/77879207. - flags->set_xla_gpu_disable_multi_streaming(true); + opts.set_xla_gpu_disable_multi_streaming(true); // TODO(jlebar): Disable fastmath once doing so is not a performance // regression. - flags->set_xla_cpu_enable_fast_math(true); - flags->set_xla_gpu_enable_fast_math(true); + opts.set_xla_cpu_enable_fast_math(true); + opts.set_xla_gpu_enable_fast_min_max(true); - flags->set_xla_force_host_platform_device_count(1); + opts.set_xla_force_host_platform_device_count(1); + return opts; } +static DebugOptions* flag_values; +static std::vector* flag_objects; +static std::once_flag flags_init; + // Allocates flag_values and flag_objects; this function must not be called more // than once - its call done via call_once. -void AllocateFlags() { - flag_values = new DebugOptions; - - SetDebugOptionsDefaults(flag_values); +static void AllocateFlags() { + flag_values = new DebugOptions(DefaultDebugOptionsIgnoringFlags()); // Returns a lambda that calls "member_setter" on "flag_values" with the // argument passed in to the lambda. @@ -101,8 +99,8 @@ void AllocateFlags() { [](string comma_separated_values) { auto* extra_options_map = flag_values->mutable_xla_backend_extra_options(); - impl::parse_xla_backend_extra_options(extra_options_map, - comma_separated_values); + parse_xla_backend_extra_options(extra_options_map, + comma_separated_values); return true; }; @@ -111,8 +109,8 @@ void AllocateFlags() { [](string reduce_precision_option_value) { HloReducePrecisionOptions* option_proto = flag_values->add_hlo_reduce_precision_options(); - return impl::parse_xla_reduce_precision_option( - option_proto, reduce_precision_option_value); + return parse_xla_reduce_precision_option(option_proto, + reduce_precision_option_value); }; flag_objects = new std::vector({ @@ -135,6 +133,11 @@ void AllocateFlags() { bool_setter_for(&DebugOptions::set_xla_hlo_dump_as_graphdef), flag_values->xla_hlo_dump_as_graphdef(), "Dump HLO graphs as TensorFlow GraphDefs."), + tensorflow::Flag("xla_hlo_dump_as_html", + bool_setter_for(&DebugOptions::set_xla_hlo_dump_as_html), + flag_values->xla_hlo_dump_as_html(), + "Dump HLO graphs as an HTML (DOT rendered into SVG " + "inlined in HTML)."), tensorflow::Flag( "xla_hlo_graph_sharding_color", bool_setter_for(&DebugOptions::set_xla_hlo_graph_sharding_color), @@ -162,11 +165,11 @@ void AllocateFlags() { "Enable unsafe fast-math optimizations in the CPU compiler; " "this may produce faster code at the expense of some accuracy."), tensorflow::Flag( - "xla_gpu_enable_fast_math", - bool_setter_for(&DebugOptions::set_xla_cpu_enable_fast_math), - flag_values->xla_cpu_enable_fast_math(), - "Enable unsafe fast-math optimizations in the GPU compiler; " - "this may produce faster code at the expense of some accuracy."), + "xla_gpu_enable_fast_min_max", + bool_setter_for(&DebugOptions::set_xla_gpu_enable_fast_min_max), + flag_values->xla_gpu_enable_fast_min_max(), + "Enable fast floating point min/max lowering that does not propagate " + "NaNs."), tensorflow::Flag( "xla_llvm_enable_alias_scope_metadata", bool_setter_for( @@ -204,6 +207,16 @@ void AllocateFlags() { "Comma-separated list of hlo passes to be disabled. These names " "must exactly match the passes' names; no whitespace around " "commas."), + tensorflow::Flag( + "xla_disable_all_hlo_passes", + bool_setter_for(&DebugOptions::set_xla_disable_all_hlo_passes), false, + "Disables all HLO passes. Notes that some passes are necessary for " + "correctness and the invariants that must be satisfied by 'fully " + "optimized' HLO are different for different devices and may change " + "over time. The only 'guarantee', such as it is, is that if you " + "compile XLA and dump the optimized HLO for some graph, you should " + "be able to run it again on the same device with the same build of " + "XLA."), tensorflow::Flag( "xla_embed_ir_in_executable", bool_setter_for(&DebugOptions::set_xla_embed_ir_in_executable), @@ -336,12 +349,16 @@ void AllocateFlags() { "overhead from context switching but we let the user override this " "behavior to help run tests on the host that run models in parallel " "across multiple devices."), + tensorflow::Flag( + "xla_gpu_disable_ptxas_optimizations", + bool_setter_for( + &DebugOptions::set_xla_gpu_disable_ptxas_optimizations), + flag_values->xla_gpu_disable_ptxas_optimizations(), + "In XLA:GPU run ptxas in -O0 (default is -O3)."), }); - ParseFlagsFromEnv(*flag_objects); + ParseFlagsFromEnvAndDieIfUnknown("XLA_FLAGS", *flag_objects); } -} // namespace - void AppendDebugOptionsFlags(std::vector* flag_list) { std::call_once(flags_init, &AllocateFlags); flag_list->insert(flag_list->end(), flag_objects->begin(), @@ -353,5 +370,4 @@ xla::DebugOptions GetDebugOptionsFromFlags() { return *flag_values; } -} // namespace legacy_flags } // namespace xla diff --git a/tensorflow/compiler/xla/legacy_flags/debug_options_flags.h b/tensorflow/compiler/xla/debug_options_flags.h similarity index 76% rename from tensorflow/compiler/xla/legacy_flags/debug_options_flags.h rename to tensorflow/compiler/xla/debug_options_flags.h index b53157f59c61cf4e0850e006ad3656f4be63a936..dbf86a40f052af09c61da0e1abb3116ef5214357 100644 --- a/tensorflow/compiler/xla/legacy_flags/debug_options_flags.h +++ b/tensorflow/compiler/xla/debug_options_flags.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_DEBUG_OPTIONS_FLAGS_H_ -#define TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_DEBUG_OPTIONS_FLAGS_H_ +#ifndef TENSORFLOW_COMPILER_XLA_DEBUG_OPTIONS_FLAGS_H_ +#define TENSORFLOW_COMPILER_XLA_DEBUG_OPTIONS_FLAGS_H_ #include @@ -22,7 +22,6 @@ limitations under the License. #include "tensorflow/core/util/command_line_flags.h" namespace xla { -namespace legacy_flags { // Appends flag definitions for debug options to flag_list. void AppendDebugOptionsFlags(std::vector* flag_list); @@ -30,9 +29,11 @@ void AppendDebugOptionsFlags(std::vector* flag_list); // Fetches a DebugOptions proto message from flags provided to the program. // Flags must be registered with the flags parser using AppendDebugOptionsFlags // first. -xla::DebugOptions GetDebugOptionsFromFlags(); +DebugOptions GetDebugOptionsFromFlags(); + +// Gets a DebugOptions proto that reflects the defaults as if no flags were set. +DebugOptions DefaultDebugOptionsIgnoringFlags(); -} // namespace legacy_flags } // namespace xla -#endif // TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_DEBUG_OPTIONS_FLAGS_H_ +#endif // TENSORFLOW_COMPILER_XLA_DEBUG_OPTIONS_FLAGS_H_ diff --git a/tensorflow/compiler/xla/legacy_flags/debug_options_parsers.h b/tensorflow/compiler/xla/debug_options_parsers.h similarity index 94% rename from tensorflow/compiler/xla/legacy_flags/debug_options_parsers.h rename to tensorflow/compiler/xla/debug_options_parsers.h index ee7eb019c07cf898e48886955b18710146644cac..80aadfd5ece0e768afaf1842d2b6c5b11c288b55 100644 --- a/tensorflow/compiler/xla/legacy_flags/debug_options_parsers.h +++ b/tensorflow/compiler/xla/debug_options_parsers.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_DEBUG_OPTIONS_PARSERS_H_ -#define TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_DEBUG_OPTIONS_PARSERS_H_ +#ifndef TENSORFLOW_COMPILER_XLA_DEBUG_OPTIONS_PARSERS_H_ +#define TENSORFLOW_COMPILER_XLA_DEBUG_OPTIONS_PARSERS_H_ #include #include "absl/strings/numbers.h" @@ -23,8 +23,6 @@ limitations under the License. #include "tensorflow/compiler/xla/xla.pb.h" namespace xla { -namespace legacy_flags { -namespace impl { template void parse_xla_backend_extra_options(T* extra_options_map, @@ -140,8 +138,6 @@ inline bool parse_xla_reduce_precision_option( return true; } -} // namespace impl -} // namespace legacy_flags } // namespace xla -#endif // TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_DEBUG_OPTIONS_PARSERS_H_ +#endif // TENSORFLOW_COMPILER_XLA_DEBUG_OPTIONS_PARSERS_H_ diff --git a/tensorflow/compiler/xla/legacy_flags/debug_options_parsers_test.cc b/tensorflow/compiler/xla/debug_options_parsers_test.cc similarity index 88% rename from tensorflow/compiler/xla/legacy_flags/debug_options_parsers_test.cc rename to tensorflow/compiler/xla/debug_options_parsers_test.cc index 6f197aec53c7596e84437a03affa9118f22f5a1d..8003c3496d5df9be2ff8a99bc171972c8e090c43 100644 --- a/tensorflow/compiler/xla/legacy_flags/debug_options_parsers_test.cc +++ b/tensorflow/compiler/xla/debug_options_parsers_test.cc @@ -15,7 +15,7 @@ limitations under the License. // Test for parse_flags_from_env.cc -#include "tensorflow/compiler/xla/legacy_flags/debug_options_parsers.h" +#include "tensorflow/compiler/xla/debug_options_parsers.h" #include #include @@ -23,13 +23,12 @@ limitations under the License. #include "tensorflow/core/platform/test.h" namespace xla { -namespace legacy_flags { // Test that the xla_backend_extra_options flag is parsed correctly. TEST(DebugOptionsFlags, ParseXlaBackendExtraOptions) { std::unordered_map test_map; string test_string = "aa=bb,cc,dd=,ee=ff=gg"; - impl::parse_xla_backend_extra_options(&test_map, test_string); + parse_xla_backend_extra_options(&test_map, test_string); EXPECT_EQ(test_map.size(), 4); EXPECT_EQ(test_map.at("aa"), "bb"); EXPECT_EQ(test_map.at("cc"), ""); @@ -41,7 +40,7 @@ TEST(DebugOptionsFlags, ParseXlaBackendExtraOptions) { TEST(DebugOptionsFlags, ParseXlaReducePrecisionOptionNoStrings) { HloReducePrecisionOptions proto; string test_string = "OP_OUTPUTS=5,10:add,dot"; - EXPECT_TRUE(impl::parse_xla_reduce_precision_option(&proto, test_string)); + EXPECT_TRUE(parse_xla_reduce_precision_option(&proto, test_string)); EXPECT_EQ(proto.location(), HloReducePrecisionOptions::OP_OUTPUTS); EXPECT_EQ(proto.exponent_bits(), 5); EXPECT_EQ(proto.mantissa_bits(), 10); @@ -56,7 +55,7 @@ TEST(DebugOptionsFlags, ParseXlaReducePrecisionOptionNoStrings) { TEST(DebugOptionsFlags, ParseXlaReducePrecisionOptionNoStringsSemicolon) { HloReducePrecisionOptions proto; string test_string = "OP_OUTPUTS=5,10:add,dot;"; - EXPECT_TRUE(impl::parse_xla_reduce_precision_option(&proto, test_string)); + EXPECT_TRUE(parse_xla_reduce_precision_option(&proto, test_string)); EXPECT_EQ(proto.location(), HloReducePrecisionOptions::OP_OUTPUTS); EXPECT_EQ(proto.exponent_bits(), 5); EXPECT_EQ(proto.mantissa_bits(), 10); @@ -71,7 +70,7 @@ TEST(DebugOptionsFlags, ParseXlaReducePrecisionOptionNoStringsSemicolon) { TEST(DebugOptionsFlags, ParseXlaReducePrecisionOptionNoOpcodes) { HloReducePrecisionOptions proto; string test_string = "UNFUSED_OP_OUTPUTS=5,10:;foo,bar/baz"; - EXPECT_TRUE(impl::parse_xla_reduce_precision_option(&proto, test_string)); + EXPECT_TRUE(parse_xla_reduce_precision_option(&proto, test_string)); EXPECT_EQ(proto.location(), HloReducePrecisionOptions::UNFUSED_OP_OUTPUTS); EXPECT_EQ(proto.exponent_bits(), 5); EXPECT_EQ(proto.mantissa_bits(), 10); @@ -84,7 +83,7 @@ TEST(DebugOptionsFlags, ParseXlaReducePrecisionOptionNoOpcodes) { TEST(DebugOptionsFlags, ParseXlaReducePrecisionOptionBoth) { HloReducePrecisionOptions proto; string test_string = "UNFUSED_OP_OUTPUTS=5,10:subtract;foo,bar/baz"; - EXPECT_TRUE(impl::parse_xla_reduce_precision_option(&proto, test_string)); + EXPECT_TRUE(parse_xla_reduce_precision_option(&proto, test_string)); EXPECT_EQ(proto.location(), HloReducePrecisionOptions::UNFUSED_OP_OUTPUTS); EXPECT_EQ(proto.exponent_bits(), 5); EXPECT_EQ(proto.mantissa_bits(), 10); @@ -96,7 +95,6 @@ TEST(DebugOptionsFlags, ParseXlaReducePrecisionOptionBoth) { EXPECT_EQ(proto.opname_substrings_to_suffix(1), "bar/baz"); } -} // namespace legacy_flags } // namespace xla int main(int argc, char* argv[]) { 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/execution_options_util.cc b/tensorflow/compiler/xla/execution_options_util.cc index e83ff7cddd675197c7f6d7018257edb4c25b6228..cf569863bbe1c92bdcafb133d49dcf5ae8890ffe 100644 --- a/tensorflow/compiler/xla/execution_options_util.cc +++ b/tensorflow/compiler/xla/execution_options_util.cc @@ -13,14 +13,13 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ #include "tensorflow/compiler/xla/execution_options_util.h" -#include "tensorflow/compiler/xla/legacy_flags/debug_options_flags.h" +#include "tensorflow/compiler/xla/debug_options_flags.h" namespace xla { ExecutionOptions CreateDefaultExecutionOptions() { ExecutionOptions execution_options; - *(execution_options.mutable_debug_options()) = - legacy_flags::GetDebugOptionsFromFlags(); + *(execution_options.mutable_debug_options()) = GetDebugOptionsFromFlags(); return execution_options; } diff --git a/tensorflow/compiler/xla/experimental/xla_sharding/xla_sharding.py b/tensorflow/compiler/xla/experimental/xla_sharding/xla_sharding.py index fb135f5ceda67ce6c001de15b8f3f084ca164826..c34e84efc80ba970624d80802841d6ec534b6fd0 100644 --- a/tensorflow/compiler/xla/experimental/xla_sharding/xla_sharding.py +++ b/tensorflow/compiler/xla/experimental/xla_sharding/xla_sharding.py @@ -18,12 +18,9 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import math - import numpy as _np # Avoids becoming a part of public Tensorflow API. from tensorflow.compiler.xla import xla_data_pb2 -from tensorflow.compiler.xla.python_api import xla_shape from tensorflow.core.framework import attr_value_pb2 @@ -64,22 +61,18 @@ class Sharding(object): tile_assignment_devices=[core])) @classmethod - def tile(cls, tile_shape, tile_assignment): + def tile(cls, tile_assignment): """Returns a Tiled sharding attribute. This causes an op to be partially computed on multiple cores in the XLA device. Args: - tile_shape: A xla_shape.Shape describing the tile shape that each core - will compute. - The tile shape does not need to be divisible by the tile assignment. tile_assignment: An np.ndarray describing the topology of the tiling and which device will compute which part of the topology. Raises: - TypeError: tile_assignment was not of np.array type or tile_shape was - not of xla_shape.Shape type. + TypeError: tile_assignment was not of np.array type. TODO(jmolloy): This concept is nefarious and is not something we really want to expose to users (especially as the @@ -87,14 +80,11 @@ class Sharding(object): """ if not isinstance(tile_assignment, _np.ndarray): raise TypeError('Tile assignment must be of type np.ndarray') - if not isinstance(tile_shape, xla_shape.Shape): - raise TypeError('Tile shape must be of type xla_shape.Shape') dims = list(tile_assignment.shape) flattened_devices = tile_assignment.reshape(-1, order='C') return Sharding( proto=xla_data_pb2.OpSharding( type=xla_data_pb2.OpSharding.OTHER, - tile_shape=tile_shape.message, tile_assignment_dimensions=dims, tile_assignment_devices=list(flattened_devices))) @@ -114,18 +104,12 @@ class Sharding(object): ValueError: The tensor to split was smaller in the split dimension than the number of devices to split over. """ - tensor.shape.assert_is_fully_defined() shape = tensor.shape.as_list() - if shape[split_dimension] < num_devices: + if (shape[split_dimension] is not None and + shape[split_dimension] < num_devices): raise ValueError('Split dimension was smaller than the required number ' - 'of splits: shape=%r, dimension=%r, num_devices=%r', - shape, split_dimension, num_devices) - - tile_shape = shape - tile_shape[split_dimension] = int( - math.ceil(tile_shape[split_dimension] / num_devices)) - tile_shape_proto = xla_data_pb2.Shape( - element_type=xla_data_pb2.F32, dimensions=tile_shape) + 'of splits: shape=%r, dimension=%r, num_devices=%r' % + (shape, split_dimension, num_devices)) tile_assignment_dims = [1] * len(shape) tile_assignment_dims[split_dimension] = num_devices @@ -133,7 +117,6 @@ class Sharding(object): return Sharding( proto=xla_data_pb2.OpSharding( type=xla_data_pb2.OpSharding.OTHER, - tile_shape=tile_shape_proto, tile_assignment_dimensions=tile_assignment_dims, tile_assignment_devices=range(num_devices))) @@ -149,7 +132,6 @@ class Sharding(object): type=xla_data_pb2.OpSharding.TUPLE, tuple_shardings=tuple_shardings) else: proto = self._proto - attr_value = attr_value_pb2.AttrValue(s=proto.SerializeToString()) # TODO(jmolloy): This need to be seriously revisited before declaring this # API available for public use. @@ -194,8 +176,8 @@ def assign_device(tensor, device): return tensor -def tile(tensor, tile_shape, tile_assignment): - Sharding.tile(tile_shape, tile_assignment).apply_to_tensor(tensor) +def tile(tensor, tile_assignment): + Sharding.tile(tile_assignment).apply_to_tensor(tensor) return tensor diff --git a/tensorflow/compiler/xla/g3doc/README.md b/tensorflow/compiler/xla/g3doc/README.md new file mode 100644 index 0000000000000000000000000000000000000000..6643bf0aab3078ff24c86b81de69216355da69a1 --- /dev/null +++ b/tensorflow/compiler/xla/g3doc/README.md @@ -0,0 +1,3 @@ +# XLA: Accelerated Linear Algebra + +These are the docs for: https://www.tensorflow.org/xla diff --git a/tensorflow/compiler/xla/g3doc/_book.yaml b/tensorflow/compiler/xla/g3doc/_book.yaml new file mode 100644 index 0000000000000000000000000000000000000000..267701e9c0e42a21d2cda6238520f6a9692e7e76 --- /dev/null +++ b/tensorflow/compiler/xla/g3doc/_book.yaml @@ -0,0 +1,35 @@ +upper_tabs: +# Tabs left of dropdown menu +- include: /_upper_tabs_left.yaml +- include: /api_docs/_upper_tabs_api.yaml +# Dropdown menu +- name: Resources + path: /resources + is_default: true + menu: + - include: /resources/_menu_toc.yaml + lower_tabs: + # Subsite tabs + other: + - name: Guide & Tutorials + contents: + - title: XLA overview + path: /xla/overview + - title: Broadcasting semantics + path: /xla/broadcasting + - title: Developing a new backend for XLA + path: /xla/developing_new_backend + - title: Using JIT compilation + path: /xla/jit + - title: Operation semantics + path: /xla/operation_semantics + - title: Shapes and layout + path: /xla/shapes + - title: Using AOT compilation + path: /xla/tfcompile + - heading: Tutorials + - title: XLA compile API + path: /xla/tutorials/xla_compile + status: experimental + +- include: /_upper_tabs_right.yaml diff --git a/tensorflow/compiler/xla/g3doc/_index.yaml b/tensorflow/compiler/xla/g3doc/_index.yaml new file mode 100644 index 0000000000000000000000000000000000000000..858de427119bfcfa82d0b1158776bf269129fd92 --- /dev/null +++ b/tensorflow/compiler/xla/g3doc/_index.yaml @@ -0,0 +1,35 @@ +book_path: /xla/_book.yaml +project_path: /xla/_project.yaml +description: +landing_page: + custom_css_path: /site-assets/css/style.css + rows: + - heading: XLA is a compiler that optimizes TensorFlow computations. + items: + - classname: devsite-landing-row-50 + description: > + XLA (Accelerated Linear Algebra) is a domain-specific compiler for linear + algebra that optimizes TensorFlow computations. The results are + improvements in speed, memory usage, and portability on server and mobile + platforms. The XLA framework is experimental and in active development. + For details, read the XLA guide. + + - classname: devsite-landing-row-cards + items: + - heading: XLA - TensorFlow, compiled + image_path: /resources/images/tf-logo-card-16x9.png + path: https://developers.googleblog.com/2017/03/xla-tensorflow-compiled.html + buttons: + - label: Read on Google Developers blog + path: https://developers.googleblog.com/2017/03/xla-tensorflow-compiled.html + - heading: XLA at the Dev Summit + youtube_id: kAOanJczHA0 + buttons: + - label: Watch the video + path: https://www.youtube.com/watch?v=kAOanJczHA0 + - heading: XLA on GitHub + image_path: /resources/images/github-card-16x9.png + path: https://github.com/tensorflow/tensorflow/tree/master/tensorflow/compiler/xla + buttons: + - label: View on GitHub + path: https://github.com/tensorflow/tensorflow/tree/master/tensorflow/compiler/xla diff --git a/tensorflow/compiler/xla/g3doc/_project.yaml b/tensorflow/compiler/xla/g3doc/_project.yaml new file mode 100644 index 0000000000000000000000000000000000000000..33d8bdb27a664d9e282d1d65c007ebf5838b196a --- /dev/null +++ b/tensorflow/compiler/xla/g3doc/_project.yaml @@ -0,0 +1,10 @@ +name: XLA +breadcrumb_name: XLA +home_url: /xla/ +parent_project_metadata_path: /_project.yaml +description: > + XLA is a compiler-based linear algebra execution engine. +use_site_branding: true +hide_from_products_list: true +content_license: cc3-apache2 +buganizer_id: 171704 diff --git a/tensorflow/compiler/xla/g3doc/broadcasting.md b/tensorflow/compiler/xla/g3doc/broadcasting.md new file mode 100644 index 0000000000000000000000000000000000000000..5c0525c1e9adf9f37d945170d05e7c18fa3d8852 --- /dev/null +++ b/tensorflow/compiler/xla/g3doc/broadcasting.md @@ -0,0 +1,204 @@ +# Broadcasting semantics + +This document describes how the broadcasting semantics in XLA work. + +## What is broadcasting? + +Broadcasting is the process of making arrays with different shapes have +compatible shapes for arithmetic operations. The terminology is borrowed from +Numpy +[broadcasting](http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html). + +Broadcasting may be required for operations between multi-dimensional arrays of +different ranks, or between multi-dimensional arrays with different but +compatible shapes. Consider the addition `X+v` where `X` is a matrix (an array +of rank 2) and `v` is a vector (an array of rank 1). To perform element-wise +addition, XLA needs to "broadcast" the vector `v` to the same rank as the +matrix `X`, by replicating `v` a certain number of times. The vector's length +has to match at least one of the dimensions of the matrix. + +For example: + + |1 2 3| + |7 8 9| + |4 5 6| + +The matrix's dimensions are (2,3), the vector's are (3). The vector is broadcast +by replicating it over rows to get: + + |1 2 3| + |7 8 9| = |8 10 12| + |4 5 6| |7 8 9| |11 13 15| + +In Numpy, this is called +[broadcasting](http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html). + +## Principles + +The XLA language is as strict and explicit as possible, avoiding implicit and +"magical" features. Such features may make some computations slightly easier to +define, at the cost of more assumptions baked into user code that will be +difficult to change in the long term. If necessary, implicit and magical +features can be added in client-level wrappers. + +In regards to broadcasting, explicit broadcasting specifications on operations +between arrays of different ranks is required. This is different from Numpy, +which infers the specification when possible. + +## Broadcasting a lower-rank array onto a higher-rank array + +*Scalars* can always be broadcast over arrays without an explicit specification +of broadcasting dimensions. An element-wise binary operation between a scalar +and an array means applying the operation with the scalar for each element in +the array. For example, adding a scalar to a matrix means producing a matrix +each element of which is a sum of the scalar with the corresponding input +matrix's element. + + |1 2 3| + 7 = |8 9 10| + |4 5 6| |11 12 13| + +Most broadcasting needs can be captured by using a tuple of dimensions on a +binary operation. When the inputs to the operation have different ranks, this +broadcasting tuple specifies which dimension(s) in the **higher-rank** array to +match with the **lower-rank** array. + +Consider the previous example, instead of adding a scalar to a (2,3) matrix, add +a vector of dimension (3) to a matrix of dimensions (2,3). *Without specifying +broadcasting, this operation is invalid.* To correctly request matrix-vector +addition, specify the broadcasting dimension to be (1), meaning the vector's +dimension is matched to dimension 1 of the matrix. In 2D, if dimension 0 is +considered as rows and dimension 1 as columns, this means that each element of +the vector becomes a column of a size matching the number of rows in the matrix: + + |7 8 9| ==> |7 8 9| + |7 8 9| + +As a more complex example, consider adding a 3-element vector (dimension (3)) to +a 3x3 matrix (dimensions (3,3)). There are two ways broadcasting can happen for +this example: + +(1) A broadcasting dimension of 1 can be used. Each vector element becomes a +column and the vector is duplicated for each row in the matrix. + + |7 8 9| ==> |7 8 9| + |7 8 9| + |7 8 9| + +(2) A broadcasting dimension of 0 can be used. Each vector element becomes a row +and the vector is duplicated for each column in the matrix. + + |7| ==> |7 7 7| + |8| |8 8 8| + |9| |9 9 9| + +> Note: when adding a 2x3 matrix to a 3-element vector, a broadcasting dimension +> of 0 is invalid. + +The broadcasting dimensions can be a tuple that describes how a smaller rank +shape is broadcast into a larger rank shape. For example, given a 2x3x4 cuboid +and a 3x4 matrix, a broadcasting tuple (1,2) means matching the matrix to +dimensions 1 and 2 of the cuboid. + +This type of broadcast is used in the binary ops in `XlaBuilder`, if the +`broadcast_dimensions` argument is given. For example, see +[XlaBuilder::Add](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.cc). +In the XLA source code, this type of broadcasting is sometimes called "InDim" +broadcasting. + +### Formal definition + +The broadcasting attribute allows matching a lower-rank array to a higher-rank +array, by specifying which dimensions of the higher-rank array to match. For +example, for an array with dimensions MxNxPxQ, a vector with dimension T can be +matched as follows: + + MxNxPxQ + + dim 3: T + dim 2: T + dim 1: T + dim 0: T + +In each case, T has to be equal to the matching dimension of the higher-rank +array. The vector's values are then broadcast from the matched dimension to all +the other dimensions. + +To match a TxV matrix onto the MxNxPxQ array, a pair of broadcasting dimensions +are used: + + MxNxPxQ + dim 2,3: T V + dim 1,2: T V + dim 0,3: T V + etc... + +The order of dimensions in the broadcasting tuple has to be the order in which +the lower-rank array's dimensions are expected to match the higher-rank array's +dimensions. The first element in the tuple says which dimension in the +higher-rank array has to match dimension 0 in the lower-rank array. The second +element for dimension 1, and so on. The order of broadcast dimensions has to be +strictly increasing. For example, in the previous example it is illegal to match +V to N and T to P; it is also illegal to match V to both P and N. + +## Broadcasting similar-rank arrays with degenerate dimensions + +A related broadcasting problem is broadcasting two arrays that have the same +rank but different dimension sizes. Similarly to Numpy's rules, this is only +possible when the arrays are *compatible*. Two arrays are compatible when all +their dimensions are compatible. Two dimensions are compatible if: + +* They are equal, or +* One of them is 1 (a "degenerate" dimension) + +When two compatible arrays are encountered, the result shape has the maximum +among the two inputs at every dimension index. + +Examples: + +1. (2,1) and (2,3) broadcast to (2,3). +2. (1,2,5) and (7,2,5) broadcast to (7,2,5) +3. (7,2,5) and (7,1,5) broadcast to (7,2,5) +4. (7,2,5) and (7,2,6) are incompatible and cannot be broadcast. + +A special case arises, and is also supported, where each of the input arrays has +a degenerate dimension at a different index. In this case, the result is an +"outer operation": (2,1) and (1,3) broadcast to (2,3). For more examples, +consult the +[Numpy documentation on broadcasting](http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html). + +## Broadcast composition + +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 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. + +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 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: + + |1 1| + [5 6] + |2 2| + |3 3| + |4 4| + +Then "degenerate dimension broadcasting" broadcasts dimension zero of the 1x2 +matrix to match the corresponding dimension size of the right hand side: + + |1 1| + |5 6| |6 7| + |2 2| + |5 6| = |7 8| + |3 3| + |5 6| |8 9| + |4 4| + |5 6| |9 10| + +A more complicated example is a matrix of size 1x2 added to an array of size +4x3x1 using broadcast dimensions of (1, 2). First the 1x2 matrix is broadcast up +to rank 3 using the broadcast dimensions to produces an intermediate Mx1x2 array +where the dimension size M is determined by the size of the larger operand (the +4x3x1 array) producing a 4x1x2 intermediate array. The M is at dimension 0 +(left-most dimension) because the dimensions 1 and 2 are mapped to the +dimensions of the original 1x2 matrix as the broadcast dimension are (1, 2). +This intermediate array can be added to the 4x3x1 matrix using broadcasting of +degenerate dimensions to produce a 4x3x2 array result. diff --git a/tensorflow/compiler/xla/g3doc/developing_new_backend.md b/tensorflow/compiler/xla/g3doc/developing_new_backend.md new file mode 100644 index 0000000000000000000000000000000000000000..5ede7f523131cf715575074b8e27487be5ea77c6 --- /dev/null +++ b/tensorflow/compiler/xla/g3doc/developing_new_backend.md @@ -0,0 +1,76 @@ +# Developing a new backend for XLA + +This preliminary guide is for early adopters that want to easily retarget +TensorFlow to their hardware in an efficient manner. The guide is not +step-by-step and assumes knowledge of [LLVM](http://llvm.org), +[Bazel](https://bazel.build/), and TensorFlow. + +XLA provides an abstract interface that a new architecture or accelerator can +implement to create a backend to run TensorFlow graphs. Retargeting XLA should +be significantly simpler and scalable than implementing every existing +TensorFlow Op for new hardware. + +Most implementations will fall into one of the following scenarios: + +1. Existing CPU architecture not yet officially supported by XLA, with or + without an existing [LLVM](http://llvm.org) backend. +2. Non-CPU-like hardware with an existing LLVM backend. +3. Non-CPU-like hardware without an existing LLVM backend. + +> Note: An LLVM backend can mean either one of the officially released LLVM +> backends or a custom LLVM backend developed in-house. + +## Scenario 1: Existing CPU architecture not yet officially supported by XLA + +In this scenario, start by looking at the existing +[XLA CPU backend](https://www.tensorflow.org/code/tensorflow/compiler/xla/service/cpu/). +XLA makes it easy to retarget TensorFlow to different CPUs by using LLVM, since +the main difference between XLA backends for CPUs is the code generated by LLVM. +Google tests XLA for x64 and ARM64 architectures. + +If the hardware vendor has an LLVM backend for their hardware, it is simple to +link the backend with the LLVM built with XLA. In JIT mode, the XLA CPU backend +emits code for the host CPU. For ahead-of-time compilation, +[`xla::AotCompilationOptions`](https://www.tensorflow.org/code/tensorflow/compiler/xla/service/compiler.h) +can provide an LLVM triple to configure the target architecture. + +If there is no existing LLVM backend but another kind of code generator exists, +it should be possible to reuse most of the existing CPU backend. + +## Scenario 2: Non-CPU-like hardware with an existing LLVM backend + +It is possible to model a new +[`xla::Compiler`](https://www.tensorflow.org/code/tensorflow/compiler/xla/service/compiler.h) +implementation on the existing +[`xla::CPUCompiler`](https://www.tensorflow.org/code/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc) +and [`xla::GPUCompiler`](https://www.tensorflow.org/code/tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc) +classes, since these already emit LLVM IR. Depending on the nature of the +hardware, it is possible that many of the LLVM IR generation aspects will have +to be changed, but a lot of code can be shared with the existing backends. + +A good example to follow is the +[GPU backend](https://www.tensorflow.org/code/tensorflow/compiler/xla/service/gpu/) +of XLA. The GPU backend targets a non-CPU-like ISA, and therefore some aspects +of its code generation are unique to the GPU domain. Other kinds of hardware, +e.g. DSPs like Hexagon (which has an upstream LLVM backend), can reuse parts of +the LLVM IR emission logic, but other parts will be unique. + +## Scenario 3: Non-CPU-like hardware without an existing LLVM backend + +If it is not possible to utilize LLVM, then the best option is to implement a +new backend for XLA for the desired hardware. This option requires the most +effort. The classes that need to be implemented are as follows: + +* [`StreamExecutor`](https://www.tensorflow.org/code/tensorflow/stream_executor/stream_executor.h): + For many devices not all methods of `StreamExecutor` are needed. See + existing `StreamExecutor` implementations for details. +* [`xla::Compiler`](https://www.tensorflow.org/code/tensorflow/compiler/xla/service/compiler.h): + This class encapsulates the compilation of an HLO computation into an + `xla::Executable`. +* [`xla::Executable`](https://www.tensorflow.org/code/tensorflow/compiler/xla/service/executable.h): + This class is used to launch a compiled computation on the platform. +* [`xla::TransferManager`](https://www.tensorflow.org/code/tensorflow/compiler/xla/service/transfer_manager.h): + This class enables backends to provide platform-specific mechanisms for + constructing XLA literal data from given device memory handles. In other + words, it helps encapsulate the transfer of data from the host to the device + and back. diff --git a/tensorflow/compiler/xla/g3doc/images/how-does-xla-work.png b/tensorflow/compiler/xla/g3doc/images/how-does-xla-work.png new file mode 100644 index 0000000000000000000000000000000000000000..15f86c3221d3637f2087a2db9f4cb008fe2690fa Binary files /dev/null and b/tensorflow/compiler/xla/g3doc/images/how-does-xla-work.png differ diff --git a/tensorflow/compiler/xla/g3doc/images/jit_cpu_xla_graph.png b/tensorflow/compiler/xla/g3doc/images/jit_cpu_xla_graph.png new file mode 100644 index 0000000000000000000000000000000000000000..4e2dc091fee1d13ae659988b1a68505e9ff77b27 Binary files /dev/null and b/tensorflow/compiler/xla/g3doc/images/jit_cpu_xla_graph.png differ diff --git a/tensorflow/compiler/xla/g3doc/images/jit_gpu_xla_graph.png b/tensorflow/compiler/xla/g3doc/images/jit_gpu_xla_graph.png new file mode 100644 index 0000000000000000000000000000000000000000..39d7c90c4fc3d707df062562fcf9ebdc37344af0 Binary files /dev/null and b/tensorflow/compiler/xla/g3doc/images/jit_gpu_xla_graph.png differ diff --git a/tensorflow/compiler/xla/g3doc/images/jit_timeline_cpu.png b/tensorflow/compiler/xla/g3doc/images/jit_timeline_cpu.png new file mode 100644 index 0000000000000000000000000000000000000000..a38f636983b527b678f17d3b0c92646ac1485f86 Binary files /dev/null and b/tensorflow/compiler/xla/g3doc/images/jit_timeline_cpu.png differ diff --git a/tensorflow/compiler/xla/g3doc/images/jit_timeline_cpu_xla.png b/tensorflow/compiler/xla/g3doc/images/jit_timeline_cpu_xla.png new file mode 100644 index 0000000000000000000000000000000000000000..285c3a96d5aa33605cab2486522a5e815901a2fc Binary files /dev/null and b/tensorflow/compiler/xla/g3doc/images/jit_timeline_cpu_xla.png differ diff --git a/tensorflow/compiler/xla/g3doc/images/jit_timeline_gpu.png b/tensorflow/compiler/xla/g3doc/images/jit_timeline_gpu.png new file mode 100644 index 0000000000000000000000000000000000000000..488fc2c2f1009706b7e2c5ded154f47e2b7f4bcb Binary files /dev/null and b/tensorflow/compiler/xla/g3doc/images/jit_timeline_gpu.png differ diff --git a/tensorflow/compiler/xla/g3doc/images/jit_timeline_gpu_xla.png b/tensorflow/compiler/xla/g3doc/images/jit_timeline_gpu_xla.png new file mode 100644 index 0000000000000000000000000000000000000000..d0df38cf18197f89224cc0f5ff643dd537d03fcc Binary files /dev/null and b/tensorflow/compiler/xla/g3doc/images/jit_timeline_gpu_xla.png differ diff --git a/tensorflow/compiler/xla/g3doc/images/ops_2d_matrix.png b/tensorflow/compiler/xla/g3doc/images/ops_2d_matrix.png new file mode 100644 index 0000000000000000000000000000000000000000..4846d1700607ced60dd3b8038996894d4dd0f8af Binary files /dev/null and b/tensorflow/compiler/xla/g3doc/images/ops_2d_matrix.png differ diff --git a/tensorflow/compiler/xla/g3doc/images/ops_alltoall.png b/tensorflow/compiler/xla/g3doc/images/ops_alltoall.png new file mode 100644 index 0000000000000000000000000000000000000000..c8150bda5bd6fb5723832a5e42e71c12cee3d399 Binary files /dev/null and b/tensorflow/compiler/xla/g3doc/images/ops_alltoall.png differ diff --git a/tensorflow/compiler/xla/g3doc/images/ops_concatenate.png b/tensorflow/compiler/xla/g3doc/images/ops_concatenate.png new file mode 100644 index 0000000000000000000000000000000000000000..26ded3d88c07205dd6eceef2d2ee151b4e390977 Binary files /dev/null and b/tensorflow/compiler/xla/g3doc/images/ops_concatenate.png differ diff --git a/tensorflow/compiler/xla/g3doc/images/ops_pad.png b/tensorflow/compiler/xla/g3doc/images/ops_pad.png new file mode 100644 index 0000000000000000000000000000000000000000..dc1948a627a88721d44bd22027ab75540f61feda Binary files /dev/null and b/tensorflow/compiler/xla/g3doc/images/ops_pad.png differ diff --git a/tensorflow/compiler/xla/g3doc/images/ops_reduce_from_2d_matrix.png b/tensorflow/compiler/xla/g3doc/images/ops_reduce_from_2d_matrix.png new file mode 100644 index 0000000000000000000000000000000000000000..c2ff037ab5c6ad7b2b2157339f189cff3b16df09 Binary files /dev/null and b/tensorflow/compiler/xla/g3doc/images/ops_reduce_from_2d_matrix.png differ diff --git a/tensorflow/compiler/xla/g3doc/images/ops_reduce_from_3d_matrix.png b/tensorflow/compiler/xla/g3doc/images/ops_reduce_from_3d_matrix.png new file mode 100644 index 0000000000000000000000000000000000000000..ebeeca093b2dda7fc5871e53302bce0e73e670be Binary files /dev/null and b/tensorflow/compiler/xla/g3doc/images/ops_reduce_from_3d_matrix.png differ diff --git a/tensorflow/compiler/xla/g3doc/images/ops_reduce_window.png b/tensorflow/compiler/xla/g3doc/images/ops_reduce_window.png new file mode 100644 index 0000000000000000000000000000000000000000..e9cdc3d148ab4ebb46bef8af84724134eae75d55 Binary files /dev/null and b/tensorflow/compiler/xla/g3doc/images/ops_reduce_window.png differ diff --git a/tensorflow/compiler/xla/g3doc/images/ops_reduce_window_stride.png b/tensorflow/compiler/xla/g3doc/images/ops_reduce_window_stride.png new file mode 100644 index 0000000000000000000000000000000000000000..f1ef5270dbac9f4ca1eb884e5cb27fd57a02ba8e Binary files /dev/null and b/tensorflow/compiler/xla/g3doc/images/ops_reduce_window_stride.png differ diff --git a/tensorflow/compiler/xla/g3doc/images/ops_scatter_to_selected_window_element.png b/tensorflow/compiler/xla/g3doc/images/ops_scatter_to_selected_window_element.png new file mode 100644 index 0000000000000000000000000000000000000000..4a82afaefab42d837a46178ba9aef3a1b6ddc434 Binary files /dev/null and b/tensorflow/compiler/xla/g3doc/images/ops_scatter_to_selected_window_element.png differ diff --git a/tensorflow/compiler/xla/g3doc/images/ops_while.png b/tensorflow/compiler/xla/g3doc/images/ops_while.png new file mode 100644 index 0000000000000000000000000000000000000000..da32b553eb0226bfb1122c236dfefe151758b9fa Binary files /dev/null and b/tensorflow/compiler/xla/g3doc/images/ops_while.png differ diff --git a/tensorflow/compiler/xla/g3doc/images/ops_xla_gather_0.svg b/tensorflow/compiler/xla/g3doc/images/ops_xla_gather_0.svg new file mode 100644 index 0000000000000000000000000000000000000000..7d324aa35bd92aeef7bc2987eaf346f1c3aa0966 --- /dev/null +++ b/tensorflow/compiler/xla/g3doc/images/ops_xla_gather_0.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/tensorflow/compiler/xla/g3doc/images/ops_xla_gather_1.svg b/tensorflow/compiler/xla/g3doc/images/ops_xla_gather_1.svg new file mode 100644 index 0000000000000000000000000000000000000000..f460b923f0efa5594a25251ba308f0fe4b9bf786 --- /dev/null +++ b/tensorflow/compiler/xla/g3doc/images/ops_xla_gather_1.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/tensorflow/compiler/xla/g3doc/images/ops_xla_gather_2.svg b/tensorflow/compiler/xla/g3doc/images/ops_xla_gather_2.svg new file mode 100644 index 0000000000000000000000000000000000000000..d9c35e972d152c63df44d3c9be65ec3a840d5544 --- /dev/null +++ b/tensorflow/compiler/xla/g3doc/images/ops_xla_gather_2.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/tensorflow/compiler/xla/g3doc/images/send_recv_order.png b/tensorflow/compiler/xla/g3doc/images/send_recv_order.png new file mode 100644 index 0000000000000000000000000000000000000000..721200e3cb0af984f58cb8594607e5c0a39ddd18 Binary files /dev/null and b/tensorflow/compiler/xla/g3doc/images/send_recv_order.png differ diff --git a/tensorflow/compiler/xla/g3doc/images/send_recv_schedule.png b/tensorflow/compiler/xla/g3doc/images/send_recv_schedule.png new file mode 100644 index 0000000000000000000000000000000000000000..c830f987ab9b7e53730555d5734ce37bd1854211 Binary files /dev/null and b/tensorflow/compiler/xla/g3doc/images/send_recv_schedule.png differ diff --git a/tensorflow/compiler/xla/g3doc/images/xla_array_layout_figure1.png b/tensorflow/compiler/xla/g3doc/images/xla_array_layout_figure1.png new file mode 100644 index 0000000000000000000000000000000000000000..00cefe4c7806c1c09dd51499375e720bfb0baac6 Binary files /dev/null and b/tensorflow/compiler/xla/g3doc/images/xla_array_layout_figure1.png differ diff --git a/tensorflow/compiler/xla/g3doc/images/xla_array_layout_figure2.png b/tensorflow/compiler/xla/g3doc/images/xla_array_layout_figure2.png new file mode 100644 index 0000000000000000000000000000000000000000..6439c6e40272ae6b2954e9d7f3de2df470a2b36d Binary files /dev/null and b/tensorflow/compiler/xla/g3doc/images/xla_array_layout_figure2.png differ diff --git a/tensorflow/compiler/xla/xlalogo.png b/tensorflow/compiler/xla/g3doc/images/xlalogo.png similarity index 100% rename from tensorflow/compiler/xla/xlalogo.png rename to tensorflow/compiler/xla/g3doc/images/xlalogo.png diff --git a/tensorflow/compiler/xla/g3doc/jit.md b/tensorflow/compiler/xla/g3doc/jit.md new file mode 100644 index 0000000000000000000000000000000000000000..85fa16ccc7f48a3dce840564e79097c9e136767f --- /dev/null +++ b/tensorflow/compiler/xla/g3doc/jit.md @@ -0,0 +1,180 @@ +# Using JIT Compilation + +> Note: TensorFlow must be compiled from source to include XLA. + +## Why use just-in-time (JIT) compilation? + +The TensorFlow/XLA JIT compiler compiles and runs parts of TensorFlow graphs via +XLA. The benefit of this over the standard TensorFlow implementation is that XLA +can fuse multiple operators (kernel fusion) into a small number of compiled +kernels. Fusing operators can reduce memory bandwidth requirements and improve +performance compared to executing operators one-at-a-time, as the TensorFlow +executor does. + +## Running TensorFlow graphs via XLA + +There are two ways to run TensorFlow computations via XLA, either by +JIT-compiling operators placed on a CPU or GPU device, or by placing operators +on the `XLA_CPU` or `XLA_GPU` TensorFlow devices. Placing operators directly on +a TensorFlow XLA device forces the operator to run on that device and is mainly +used for testing. + +> Note: The XLA CPU backend supports intra-op parallelism (i.e. it can shard a +> single operation across multiple cores) but it does not support inter-op +> parallelism (i.e. it cannot execute independent operations concurrently across +> multiple cores). The XLA GPU backend is competitive with the standard +> TensorFlow implementation, sometimes faster, sometimes slower. + +### Turning on JIT compilation + +JIT compilation can be turned on at the session level or manually for select +operations. Both of these approaches are zero-copy --- data does not need to be +copied when passing data between a compiled XLA kernel and a TensorFlow operator +placed on the same device. + +#### Session + +Turning on JIT compilation at the session level will result in all possible +operators being greedily compiled into XLA computations. Each XLA computation +will be compiled into one or more kernels for the underlying device. + +Subject to a few constraints, if there are two adjacent operators in the graph +that both have XLA implementations, then they will be compiled into a single XLA +computation. + +JIT compilation is turned on at the session level by setting the +`global_jit_level` config to `tf.OptimizerOptions.ON_1` and passing the config +during session initialization. + +```python +# Config to turn on JIT compilation +config = tf.ConfigProto() +config.graph_options.optimizer_options.global_jit_level = tf.OptimizerOptions.ON_1 + +sess = tf.Session(config=config) +``` + +> Note: Turning on JIT at the session level will not result in operations being +> compiled for the CPU. JIT compilation for CPU operations must be done via +> the manual method documented below. + +#### Manual with experimental_jit_scope() + +JIT compilation can also be turned on manually for one or more operators. This +is done by tagging the operators to compile with the attribute +`_XlaCompile=true`. The simplest way to do this is via the +`tf.contrib.compiler.jit.experimental_jit_scope()` scope defined in +[`tensorflow/contrib/compiler/jit.py`](https://www.tensorflow.org/code/tensorflow/contrib/compiler/jit.py). +Example usage: + +```python + jit_scope = tf.contrib.compiler.jit.experimental_jit_scope + + x = tf.placeholder(np.float32) + with jit_scope(): + y = tf.add(x, x) # The "add" will be compiled with XLA. +``` + +The `_XlaCompile` attribute is currently supported on a best-effort basis. If an +operator cannot be compiled, TensorFlow will silently fall back to the normal +implementation. + +#### Manual with xla.compile() + +Unlike experimental_jit_scope() which silently falls back to normal Tensorflow +on uncompilable operator, xla.compile() returns an explicit error. This is +useful if you want more predictable behaviors from XLA compilation. + +Please see +[xla.compile() tutorial Colab](./tutorials/xla_compile.ipynb) +for how to use it. + +### Placing operators on XLA devices + +Another way to run computations via XLA is to place an operator on a specific +XLA device. This method is normally only used for testing. Valid targets are +`XLA_CPU` or `XLA_GPU`. + +```python +with tf.device("/job:localhost/replica:0/task:0/device:XLA_GPU:0"): + output = tf.add(input1, input2) +``` + +Unlike JIT compilation on the standard CPU and GPU devices, these devices make a +copy of data when it is transferred on and off the device. The extra copy makes +it expensive to mix XLA and TensorFlow operators in the same graph. + +## Tutorial + +This tutorial covers training a simple version of MNIST softmax with JIT turned +on. Currently JIT at the session level, which is what is used for the tutorial, +only supports GPU. + +Before starting the tutorial verify that the LD_LIBRARY environment variable or +ldconfig contains `$CUDA_ROOT/extras/CUPTI/lib64`, which contains libraries for +the CUDA Profiling Tools Interface +[(CUPTI)](http://docs.nvidia.com/cuda/cupti/index.html). TensorFlow uses CUPTI +to pull tracing information from the GPU. + +### Step #1: Prepare sample script + +Download or move +[mnist_softmax_xla.py](https://www.tensorflow.org/code/tensorflow/examples/tutorials/mnist/mnist_softmax_xla.py) +into a folder outside of the TensorFlow source tree. + +### Step #2: Run without XLA + +Execute the python script to train the model without XLA. + +```shell +python mnist_softmax_xla.py --xla='' +``` + +Using the Chrome Trace Event Profiler (browse to chrome://tracing), +open the timeline file created when the script finishes: `timeline.ctf.json`. +The rendered timeline should look similar to the picture below with multiple +green boxes labeled `MatMul`, possibly across multiple CPUs. +
+ +
+ +### Step #3 Run with XLA + +Execute the python script to train the model with XLA and turn on a debugging +feature of XLA via an environmental variable that outputs the XLA graph. + +```shell +XLA_FLAGS="--xla_hlo_graph_path=/tmp --xla_generate_hlo_graph=.*" python mnist_softmax_xla.py +``` + +Open the timeline file created (`timeline.ctf.json`). The rendered timeline +should look similar to the picture below with one long bar labeled `XlaLaunch`. +
+ +
+ +To understand what is happening in `XlaLaunch`, look at the console output for +statements similar to the following: + +```shell +computation cluster_0[_XlaCompiledKernel=true,_XlaNumConstantArgs=1].v82 [CPU: +pipeline start, before inline]: /tmp/hlo_graph_0.dot + +``` + +The console statements point to the location of `hlo_graph_xx.dot` files that +contain information about the graph created by XLA. The process that XLA takes +to fuse Ops is visible by starting at `hlo_graph_0.dot` and viewing each diagram +in succession. + +To Render the .dot file into a png, install +[GraphViz](https://www.graphviz.org/download/) and run: + +```shell +dot -Tpng hlo_graph_80.dot -o hlo_graph_80.png +``` + +The result will look like the following: +
+ +
diff --git a/tensorflow/compiler/xla/g3doc/layout_with_tiling.md b/tensorflow/compiler/xla/g3doc/layout_with_tiling.md new file mode 100644 index 0000000000000000000000000000000000000000..5e990851af7495ebd4417e44f1d955fcc14dadf1 --- /dev/null +++ b/tensorflow/compiler/xla/g3doc/layout_with_tiling.md @@ -0,0 +1,159 @@ +# 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.* + +
![](images/xla_array_layout_figure1.png) + +Figure 1
+ +Figure 1 shows how an array F32[3,5] is laid out in memory with 2x2 tiling. A +shape with this layout is written as F32[3,5]{1,0:(2,2)}, where 1,0 relates to +the physical order of dimensions (minor_to_major field in Layout) while (2,2) +after the colon indicates tiling of the physical dimensions by a 2x2 tile. + +Intuitively tiles are laid out to cover the shape and then within each tile, +elements are then laid out without tiling, as in the example above, where the +right part of the example shows the layout in memory, including the white +padding elements that are added in order to have complete 2x2 tiles even though +the original array bounds are not even. + +The extra elements in the padding are not required to contain any particular +value. + +## Linear index formulas for tiling given a shape and a tile + +Without tiling, an element e=(en, en-1, ... , +e1) in an array with array bounds d=(dn, dn-1, +... , d1) (d1 is the most minor dimension) is laid out by major to +minor order at position: + +   linear_index(e, d) \ += linear_index((en, en-1, ... , e1), +(dn, dn-1, ... , d1)) \ += endn-1...d1 + +en-1dn-2...d1 + ... + e1 + +For simplicity of notation in this document we assume a tile has the same number +of dimensions as the array. In XLA's implementation of tiling, this is +generalized to tilings with fewer dimensions by leaving the initial most-major +dimensions unchanged and applying the tiling only to the most minor dimensions, +so that the tiling that is specified mentions a suffix of the physical +dimensions of the shape being tiled. + +When tiling of size (tn, tn-1, ... , t1) is +used, an element in the array with indices (en, en-1, ... +, e1) is mapped to this position in the final layout: + +   linear_index_with_tile(e, d, t) \ += linear_index((⌊e/t⌋, e mod t), (⌈d/t⌉, t))     (arithmetic is +elementwise, (a,b) is concatenation) \ += linear_index((⌊en/tn⌋, ... , +⌊e1/t1⌋, en mod tn, ... , +e1 mod t1), (⌈dn/tn⌉, ... , +⌈d1/t1⌉, tn, tn-1, ... , +t1)) \ += linear_index((⌊en/tn⌋, ... , +⌊e1/t1⌋), (⌈dn/tn⌉, ... , +⌈d1/t1⌉))∙tntn-1...t1 + +linear_index((en mod tn, ... , e1 mod +t1), (tn, tn-1, ... , t1)) + +The layout can be thought of as having two parts: +(⌊en/tn⌋, ... , ⌊e1/t1⌋), which +corresponds to a tile index in an array of tiles of size +(⌈dn/tn⌉, ... , ⌈d1/t1⌉), and +(en mod tn, ... , e1 mod t1), which +corresponds to a within-tile index. The ceil function appears in +⌈di/ti⌉ because if tiles overrun the bounds of the larger +array, padding is inserted as in Figure 1. Both the tiles and elements within +tiles are laid out recursively without tiling. + +For the example in Figure 1, element (2,3) has tile index (1,1), and within-tile +index (0,1), for a combined coordinate vector of (1, 1, 0, 1). The tile indices +have bounds (2, 3) and the tile itself is (2, 2) for a combined vector of (2, 3, +2, 2). The linear index with tile for the element with index (2, 3) in the +logical shape is then + +   linear_index_with_tile((2,3), (3,5), (2,2)) \ += linear_index((1,1,0,1), (2,3,2,2)) \ += linear_index((1,1), (2,3)) ∙ 2 ∙ 2 + linear_index((0,1), (2,2)) \ += (1 ∙ 3 + 1) ∙ 2 ∙ 2 + (0 ∙ 2 + 1) \ += 17. + +# Tiling as pad-reshape-transpose + +Tiling-based layout operates as follows: \ +Consider an array of dimensions (dn, dn-1, ... , d1) (d1 +is the most minor dimension). When it’s laid out with tiling of size +(tn, tn-1, ... , t1) (t1 is the most +minor dimension), that tiling can be described in terms of pad-reshape-transpose +in the following way. + +1. The array is padded to (⌈dn/tn⌉∙tn, ... , + ⌈d1/t1⌉∙t1). +2. Each dimension i is broken into (⌈di/ti⌉, + ti), i.e. the array is reshaped to \ +     (⌈dn/tn⌉, tn, ... , + ⌈d1/t1⌉, t1). \ + There is no physical layout change in this reshape by itself, so this + reshape is a bitcast. If one is not explicitly thinking of a tiling, this + reshape could express any shape with the same number of elements as the + padded shape - the example here is of how to express a tile in this way. +3. A transpose happens by moving tn, ... , t1 to the most + minor dimensions while keeping their relative order, so that the order of + dimensions from most major to most minor becomes \ +     (⌈dn/tn⌉, ... , + ⌈d1/t1⌉, tn, ... , t1). + +The final shape has the prefix \ +    (⌈dn/tn⌉, ... , +⌈d1/t1⌉), which describes the number of tiles in each +dimension. An element in the array (en, ... , e1) is +mapped to this element in the final shape: \ +    (⌊en/tn⌋, ... , +⌊e0/t0⌋, en mod tn, ... , +e1 mod t1). It is easy to see that the linear index of the +element follows the formula above as expected. + +# Repeated tiling + +XLA's tiling becomes even more flexible by applying it repeatedly. + +
![](images/xla_array_layout_figure2.png) + +Figure 2
+ +Figure 2 shows how an array of size 4x8 is tiled by two levels of tiling (first +2x4 then 2x1). We represent this repeated tiling as (2,4)(2,1). Each color +indicates a 2x4 tile and each red border box is a 2x1 tile. The numbers +indicates the linear index in memory of that element in the tiled format. This +format matches the format used for BF16 on TPU, except that the initial tile is +bigger, namely the tiling is (8,128)(2,1), where the purpose of the second +tiling by 2x1 is to collect together two 16 bit values to form one 32 bit value +in a way that aligns with the architecture of a TPU. + +Note that a second or later tile can refer to both the minor within-tile +dimensions, which just rearranges data within the tile, as in this example with +(8,128)(2,1), but can also refer to the major cross-tile dimensions from the +prior tiling. + +# Combining dimensions using tiles + +XLA's tiling also supports combining dimensions. For example, it can combine +dimensions in F32[2,7,8,11,10]{4,3,2,1,0} into F32[112,110]{1,0} first before +tiling it with (2,3). The tile used is (∗,∗,2,∗,3). Here an +asterisk in a tile implies taking that dimension and combining it with the next +more minor dimension. Multiple adjacent dimensions can be subsumed together into +one dimension. A subsumed dimension is represented by a tile value of -1 in that +dimension of the tile, which is not otherwise valid in a tile as a dimension +size. + +More precisely, if dimension i of the shape is eliminated via an asterisk in the +tile, then before the prior definition of tiling is applied, that dimension is +removed from both the shape being tiled and the tile vector, and what was +dimension i-1 of the shape has its array bound increased from di-1 to +didi-1. This step is repeated for each asterisk in the +tile vector. diff --git a/tensorflow/compiler/xla/g3doc/operation_semantics.md b/tensorflow/compiler/xla/g3doc/operation_semantics.md new file mode 100644 index 0000000000000000000000000000000000000000..c5f9377f98868cdf6d5c711cf80ede5d41fd8305 --- /dev/null +++ b/tensorflow/compiler/xla/g3doc/operation_semantics.md @@ -0,0 +1,2498 @@ +# Operation Semantics + +The following describes the semantics of operations defined in the +[`XlaBuilder`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h) +interface. Typically, these operations map one-to-one to operations defined in +the RPC interface in +[`xla_data.proto`](https://www.tensorflow.org/code/tensorflow/compiler/xla/xla_data.proto). + +A note on nomenclature: the generalized data type XLA deals with is an +N-dimensional array holding elements of some uniform type (such as 32-bit +float). Throughout the documentation, *array* is used to denote an +arbitrary-dimensional array. For convenience, special cases have more specific +and familiar names; for example a *vector* is a 1-dimensional array and a +*matrix* is a 2-dimensional array. + +## AfterAll + +See also +[`XlaBuilder::AfterAll`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + +AfterAll takes a variadic number of tokens and produces a single token. Tokens +are primitive types which can be threaded between side-effecting operations to +enforce ordering. `AfterAll` can be used as a join of tokens for ordering a +operation after a set operations. + + `AfterAll(operands)` + +Arguments | Type | Semantics +---------- | ------- | ------------------------- +`operands` | `XlaOp` | variadic number of tokens + +## AllToAll + +See also +[`XlaBuilder::AllToAll`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + +Alltoall is a collective operation that sends data from all cores to all cores. +It has two phases: + +1. the scatter phase. On each core, the operand is split into `split_count` +number of blocks along the `split_dimensions`, and the blocks are scattered +to all cores, e.g., the ith block is send to the ith core. +2. the gather phase. Each core concatenates the received blocks along the +`concat_dimension`. + +The participating cores can be configured by: + +- `replica_groups`: each ReplicaGroup contains a list of replica id. If empty, +all replicas belong to one group in the order of 0 - (n-1). Alltoall will be +applied within subgroups in the specified order. For example, replica +groups = {{1,2,3},{4,5,0}} means, an Alltoall will be applied within replica +1, 2, 3, and in the gather phase, the received blocks will be concatenated +in the order of 1, 2, 3; another Alltoall will be applied within replica 4, +5, 0, and the concatenation order is 4, 5, 0. + +Prerequisites: + +- The dimension size of the operand on the split_dimension is divisible by +split_count. +- The operand's shape is not tuple. + + `AllToAll(operand, split_dimension, concat_dimension, split_count, +replica_groups)` + + +| Arguments | Type | Semantics | +| ------------------ | --------------------- | ------------------------------- | +| `operand` | `XlaOp` | n dimensional input array | +| `split_dimension` | `int64` | A value in the interval `[0, | +: : : n)` that names the dimension : +: : : along which the operand is : +: : : split : +| `concat_dimension` | `int64` | a value in the interval `[0, | +: : : n)` that names the dimension : +: : : along which the split blocks : +: : : are concatenated : +| `split_count` | `int64` | the number of cores that | +: : : participate this operation. If : +: : : `replica_groups` is empty, this : +: : : should be the number of : +: : : replicas; otherwise, this : +: : : should be equal to the number : +: : : of replicas in each group. : +| `replica_groups` | `ReplicaGroup` vector | each group contains a list of | +: : : replica id. : + +Below shows an example of Alltoall. + +``` +XlaBuilder b("alltoall"); +auto x = Parameter(&b, 0, ShapeUtil::MakeShape(F32, {4, 16}), "x"); +AllToAll(x, /*split_dimension=*/1, /*concat_dimension=*/0, /*split_count=*/4); +``` + +
+ +
+ +In this example, there are 4 cores participating the Alltoall. On each core, the +operand is split into 4 parts along dimension 0, so each part has shape +f32[4,4]. The 4 parts are scattered to all cores. Then each core concatenates +the received parts along dimension 1, in the order or core 0-4. So the output on +each core has shape f32[16,4]. + +## BatchNormGrad + +See also +[`XlaBuilder::BatchNormGrad`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h) +and [the original batch normalization paper](https://arxiv.org/abs/1502.03167) +for a detailed description of the algorithm. + +Calculates gradients of batch norm. + + `BatchNormGrad(operand, scale, mean, variance, grad_output, epsilon, feature_index)` + +| Arguments | Type | Semantics | +| --------------- | ----------------------- | -------------------------------- | +| `operand` | `XlaOp` | n dimensional array to be | +: : : normalized (x) : +| `scale` | `XlaOp` | 1 dimensional array | +: : : (\\(\gamma\\)) : +| `mean` | `XlaOp` | 1 dimensional array (\\(\mu\\)) | +| `variance` | `XlaOp` | 1 dimensional array | +: : : (\\(\sigma^2\\)) : +| `grad_output` | `XlaOp` | Gradients passed to | +: : : `BatchNormTraining` : +: : : (\\( \nabla y\\)) : +| `epsilon` | `float` | Epsilon value (\\(\epsilon\\)) | +| `feature_index` | `int64` | Index to feature dimension in | +: : : `operand` : + +For each feature in the feature dimension (`feature_index` is the index for the +feature dimension in `operand`), the operation calculates the gradients with +respect to `operand`, `offset` and `scale` across all the other dimensions. The +`feature_index` must be a valid index for the feature dimension in `operand`. + +The three gradients are defined by the following formulas (assuming a +4-dimensional array as `operand` and with feature dimension index $$l$$, batch +size `m` and spatial sizes `w` and `h`): + +\\[ \begin{split} c_l&= +\frac{1}{mwh}\sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h +\left( \nabla y_{ijkl} \frac{x_{ijkl} - \mu_l}{\sigma^2_l+\epsilon} \right) +\\\\ +\nabla x_{ijkl} &= \frac{\gamma_{l}}{\sqrt{\sigma^2_{l}+\epsilon}} +\left( \nabla y_{ijkl} - \mathrm{mean}(\nabla y) - c_l (x_{ijkl} - \mu_{l}) +\right) +\\\\ +\nabla \gamma_l &= \sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h \left( \nabla y_{ijkl} +\frac{x_{ijkl} - \mu_l}{\sqrt{\sigma^2_{l}+\epsilon}} \right) +\\\\\ +\nabla \beta_l &= \sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h \nabla y_{ijkl} +\end{split} \\] + +The inputs `mean` and `variance` represent moments value +across batch and spatial dimensions. + +The output type is a tuple of three handles: + +| Outputs | Type | Semantics | +| ------------- | ----------------------- | --------------------------------- | +| `grad_operand` | `XlaOp` | gradient with respect to input | +: : : `operand` (\\( \nabla x\\)) : +| `grad_scale` | `XlaOp` | gradient with respect to input | +: : : `scale` (\\( \nabla \gamma\\)) : +| `grad_offset` | `XlaOp` | gradient with respect to input | +: : : `offset`(\\( \nabla \beta\\)) : + +## BatchNormInference + +See also +[`XlaBuilder::BatchNormInference`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h) +and [the original batch normalization paper](https://arxiv.org/abs/1502.03167) +for a detailed description of the algorithm. + +Normalizes an array across batch and spatial dimensions. + + `BatchNormInference(operand, scale, offset, mean, variance, epsilon, feature_index)` + +Arguments | Type | Semantics +--------------- | ------- | --------------------------------------- +`operand` | `XlaOp` | n dimensional array to be normalized +`scale` | `XlaOp` | 1 dimensional array +`offset` | `XlaOp` | 1 dimensional array +`mean` | `XlaOp` | 1 dimensional array +`variance` | `XlaOp` | 1 dimensional array +`epsilon` | `float` | Epsilon value +`feature_index` | `int64` | Index to feature dimension in `operand` + +For each feature in the feature dimension (`feature_index` is the index for the +feature dimension in `operand`), the operation calculates the mean and variance +across all the other dimensions and uses the mean and variance to normalize each +element in `operand`. The `feature_index` must be a valid index for the feature +dimension in `operand`. + +`BatchNormInference` is equivalent to calling `BatchNormTraining` without +computing `mean` and `variance` for each batch. It uses the input `mean` and +`variance` instead as estimated values. The purpose of this op is to reduce +latency in inference, hence the name `BatchNormInference`. + +The output is an n-dimensional, normalized array with the same shape as input +`operand`. + +## BatchNormTraining + +See also +[`XlaBuilder::BatchNormTraining`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h) +and [`the original batch normalization paper`](https://arxiv.org/abs/1502.03167) +for a detailed description of the algorithm. + +Normalizes an array across batch and spatial dimensions. + + `BatchNormTraining(operand, scale, offset, epsilon, feature_index)` + +Arguments | Type | Semantics +--------------- | ------- | ---------------------------------------- +`operand` | `XlaOp` | n dimensional array to be normalized (x) +`scale` | `XlaOp` | 1 dimensional array (\\(\gamma\\)) +`offset` | `XlaOp` | 1 dimensional array (\\(\beta\\)) +`epsilon` | `float` | Epsilon value (\\(\epsilon\\)) +`feature_index` | `int64` | Index to feature dimension in `operand` + +For each feature in the feature dimension (`feature_index` is the index for the +feature dimension in `operand`), the operation calculates the mean and variance +across all the other dimensions and uses the mean and variance to normalize each +element in `operand`. The `feature_index` must be a valid index for the feature +dimension in `operand`. + +The algorithm goes as follows for each batch in `operand` \\(x\\) that +contains `m` elements with `w` and `h` as the size of spatial dimensions +(assuming `operand` is an 4 dimensional array): + +- Calculates batch mean \\(\mu_l\\) for each feature `l` in feature dimension: +\\(\mu_l=\frac{1}{mwh}\sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h x_{ijkl}\\) + +- Calculates batch variance \\(\sigma^2_l\\): +\\(\sigma^2_l=\frac{1}{mwh}\sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h (x_{ijkl} - \mu_l)^2\\) + +- Normalizes, scales and shifts: +\\(y_{ijkl}=\frac{\gamma_l(x_{ijkl}-\mu_l)}{\sqrt[2]{\sigma^2_l+\epsilon}}+\beta_l\\) + +The epsilon value, usually a small number, is added to avoid divide-by-zero errors. + +The output type is a tuple of three `XlaOp`s: + +| Outputs | Type | Semantics | +| ------------ | ----------------------- | -------------------------------------| +| `output` | `XlaOp` | n dimensional array with the same | +: : : shape as input `operand` (y) : +| `batch_mean` | `XlaOp` | 1 dimensional array (\\(\mu\\)) | +| `batch_var` | `XlaOp` | 1 dimensional array (\\(\sigma^2\\)) | + +The `batch_mean` and `batch_var` are moments calculated across the batch and +spatial dimensions using the formulas above. + +## BitcastConvertType + +See also +[`XlaBuilder::BitcastConvertType`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + +Similar to a `tf.bitcast` in TensorFlow, performs an element-wise bitcast +operation from a data shape to a target shape. The dimensions must match, and +the conversion is an element-wise one; e.g. `s32` elements become `f32` elements +via bitcast routine. Bitcast is implemented as a low-level cast, so machines +with different floating-point representations will give different results. + + `BitcastConvertType(operand, new_element_type)` + +Arguments | Type | Semantics +------------------ | --------------- | --------------------------- +`operand` | `XlaOp` | array of type T with dims D +`new_element_type` | `PrimitiveType` | type U + +The dimensions of the operand and the target shape must match. The bit-width of +the source and destination element types must be equal. The source +and destination element types must not be tuples. + +## Broadcast + +See also +[`XlaBuilder::Broadcast`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + +Adds dimensions to an array by duplicating the data in the array. + + `Broadcast(operand, broadcast_sizes)` + +Arguments | Type | Semantics +----------------- | ------------------- | ------------------------------- +`operand` | `XlaOp` | The array to duplicate +`broadcast_sizes` | `ArraySlice` | The sizes of the new dimensions + +The new dimensions are inserted on the left, i.e. if `broadcast_sizes` has +values `{a0, ..., aN}` and the operand shape has dimensions `{b0, ..., bM}` then +the shape of the output has dimensions `{a0, ..., aN, b0, ..., bM}`. + +The new dimensions index into copies of the operand, i.e. + +``` +output[i0, ..., iN, j0, ..., jM] = operand[j0, ..., jM] +``` + +For example, if `operand` is a scalar `f32` with value `2.0f`, and +`broadcast_sizes` is `{2, 3}`, then the result will be an array with shape +`f32[2, 3]` and all the values in the result will be `2.0f`. + +## Call + +See also +[`XlaBuilder::Call`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + +Invokes a computation with the given arguments. + + `Call(computation, args...)` + +| Arguments | Type | Semantics | +| ------------- | ---------------------- | ----------------------------------- | +| `computation` | `XlaComputation` | computation of type `T_0, T_1, ..., | +: : : T_N -> S` with N parameters of : +: : : arbitrary type : +| `args` | sequence of N `XlaOp`s | N arguments of arbitrary type | + +The arity and types of the `args` must match the parameters of the +`computation`. It is allowed to have no `args`. + +## Clamp + +See also +[`XlaBuilder::Clamp`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + +Clamps an operand to within the range between a minimum and maximum value. + + `Clamp(min, operand, max)` + +Arguments | Type | Semantics +--------- | ------- | --------------- +`min` | `XlaOp` | array of type T +`operand` | `XlaOp` | array of type T +`max` | `XlaOp` | array of type T + +Given an operand and minimum and maximum values, returns the operand if it is in +the range between the minimum and maximum, else returns the minimum value if the +operand is below this range or the maximum value if the operand is above this +range. That is, `clamp(a, x, b) = min(max(a, x), b)`. + +All three arrays must be the same shape. Alternatively, as a restricted form of +[broadcasting](broadcasting.md), `min` and/or `max` can be a scalar of type `T`. + +Example with scalar `min` and `max`: + +``` +let operand: s32[3] = {-1, 5, 9}; +let min: s32 = 0; +let max: s32 = 6; +==> +Clamp(min, operand, max) = s32[3]{0, 5, 6}; +``` + +## Collapse + +See also +[`XlaBuilder::Collapse`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h) +and the `tf.reshape` operation. + +Collapses dimensions of an array into one dimension. + + `Collapse(operand, dimensions)` + +Arguments | Type | Semantics +------------ | -------------- | ----------------------------------------------- +`operand` | `XlaOp` | array of type T +`dimensions` | `int64` vector | in-order, consecutive subset of T's dimensions. + +Collapse replaces the given subset of the operand's dimensions by a single +dimension. The input arguments are an arbitrary array of type T and a +compile-time-constant vector of dimension indices. The dimension indices must be +an in-order (low to high dimension numbers), consecutive subset of T's +dimensions. Thus, {0, 1, 2}, {0, 1}, or {1, 2} are all valid dimension sets, but +{1, 0} or {0, 2} are not. They are replaced by a single new dimension, in the +same position in the dimension sequence as those they replace, with the new +dimension size equal to the product of original dimension sizes. The lowest +dimension number in `dimensions` is the slowest varying dimension (most major) +in the loop nest which collapses these dimension, and the highest dimension +number is fastest varying (most minor). See the `tf.reshape` operator +if more general collapse ordering is needed. + +For example, let v be an array of 24 elements: + +``` +let v = f32[4x2x3] {{{10, 11, 12}, {15, 16, 17}}, +{{20, 21, 22}, {25, 26, 27}}, +{{30, 31, 32}, {35, 36, 37}}, +{{40, 41, 42}, {45, 46, 47}}}; + +// Collapse to a single dimension, leaving one dimension. +let v012 = Collapse(v, {0,1,2}); +then v012 == f32[24] {10, 11, 12, 15, 16, 17, +20, 21, 22, 25, 26, 27, +30, 31, 32, 35, 36, 37, +40, 41, 42, 45, 46, 47}; + +// Collapse the two lower dimensions, leaving two dimensions. +let v01 = Collapse(v, {0,1}); +then v01 == f32[4x6] {{10, 11, 12, 15, 16, 17}, +{20, 21, 22, 25, 26, 27}, +{30, 31, 32, 35, 36, 37}, +{40, 41, 42, 45, 46, 47}}; + +// Collapse the two higher dimensions, leaving two dimensions. +let v12 = Collapse(v, {1,2}); +then v12 == f32[8x3] {{10, 11, 12}, +{15, 16, 17}, +{20, 21, 22}, +{25, 26, 27}, +{30, 31, 32}, +{35, 36, 37}, +{40, 41, 42}, +{45, 46, 47}}; + +``` + +## CollectivePermute + +See also +[`XlaBuilder::CollectivePermute`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + +CollectivePermute is a collective operation that sends and receives data cross +replicas. + + `CollectivePermute(operand, source_target_pairs)` + +| Arguments | Type | Semantics | +| --------------------- | ----------------------- | -------------------------- | +| `operand` | `XlaOp` | n dimensional input array | +| `source_target_pairs` | `` vector | A list of | +: : : (source_replica_id, : +: : : target_replica_id) pairs. : +: : : For each pair, the operand : +: : : is sent from source : +: : : replica to target replica. : + +Note that there are the following restrictions on the `source_target_pair`: + +- Any two pairs should not have the same target replica id, and they should +not have the same source replica id. +- If a replica id is not a target in any pair, then the output on that replica +is a tensor consists of 0(s) with the same shape as the input. + +## Concatenate + +See also +[`XlaBuilder::ConcatInDim`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + +Concatenate composes an array from multiple array operands. The array is of the +same rank as each of the input array operands (which must be of the same rank as +each other) and contains the arguments in the order that they were specified. + + `Concatenate(operands..., dimension)` + +| Arguments | Type | Semantics | +| ----------- | --------------------- | -------------------------------------- | +| `operands` | sequence of N `XlaOp` | N arrays of type T with dimensions | +: : : [L0, L1, ...]. Requires N >= 1. : +| `dimension` | `int64` | A value in the interval `[0, N)` that | +: : : names the dimension to be concatenated : +: : : between the `operands`. : + +With the exception of `dimension` all dimensions must be the same. This is +because XLA does not support "ragged" arrays. Also note that rank-0 values +cannot be concatenated (as it's impossible to name the dimension along which the +concatenation occurs). + +1-dimensional example: + +``` +Concat({{2, 3}, {4, 5}, {6, 7}}, 0) +>>> {2, 3, 4, 5, 6, 7} +``` + +2-dimensional example: + +``` +let a = { +{1, 2}, +{3, 4}, +{5, 6}, +}; +let b = { +{7, 8}, +}; +Concat({a, b}, 0) +>>> { +{1, 2}, +{3, 4}, +{5, 6}, +{7, 8}, +} +``` + +Diagram: +
+ +
+ +## Conditional + +See also +[`XlaBuilder::Conditional`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + + `Conditional(pred, true_operand, true_computation, false_operand, +false_computation)` + +Arguments | Type | Semantics +------------------- | ---------------- | --------------------------------- +`pred` | `XlaOp` | Scalar of type `PRED` +`true_operand` | `XlaOp` | Argument of type `T_0` +`true_computation` | `XlaComputation` | XlaComputation of type `T_0 -> S` +`false_operand` | `XlaOp` | Argument of type `T_1` +`false_computation` | `XlaComputation` | XlaComputation of type `T_1 -> S` + +Executes `true_computation` if `pred` is `true`, `false_computation` if `pred` +is `false`, and returns the result. + +The `true_computation` must take in a single argument of type `T_0` and will be +invoked with `true_operand` which must be of the same type. The +`false_computation` must take in a single argument of type `T_1` and will be +invoked with `false_operand` which must be of the same type. The type of the +returned value of `true_computation` and `false_computation` must be the same. + +Note that only one of `true_computation` and `false_computation` will be +executed depending on the value of `pred`. + +## Conv (convolution) + +See also +[`XlaBuilder::Conv`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + +As ConvWithGeneralPadding, but the padding is specified in a short-hand way as +either SAME or VALID. SAME padding pads the input (`lhs`) with zeroes so that +the output has the same shape as the input when not taking striding into +account. VALID padding simply means no padding. + +## ConvWithGeneralPadding (convolution) + +See also +[`XlaBuilder::ConvWithGeneralPadding`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + +Computes a convolution of the kind used in neural networks. Here, a convolution +can be thought of as a n-dimensional window moving across a n-dimensional base +area and a computation is performed for each possible position of the window. + +| Arguments | Type | Semantics | +| --------------------- | ------------------------ | ------------------------ | +| `lhs` | `XlaOp` | rank n+2 array of inputs | +| `rhs` | `XlaOp` | rank n+2 array of kernel | +: : : weights : +| `window_strides` | `ArraySlice` | n-d array of kernel | +: : : strides : +| `padding` | `ArraySlice< pair>` : padding : +| `lhs_dilation` | `ArraySlice` | n-d lhs dilation factor | +: : : array : +| `rhs_dilation` | `ArraySlice` | n-d rhs dilation factor | +: : : array : +| `feature_group_count` | int64 | the number of feature | +: : : groups : +| `batch_group_count` | int64 | the number of batch | +: : : groups : + +Let n be the number of spatial dimensions. The `lhs` argument is a rank n+2 +array describing the base area. This is called the input, even though of course +the rhs is also an input. In a neural network, these are the input activations. +The n+2 dimensions are, in this order: + +* `batch`: Each coordinate in this dimension represents an independent input +for which convolution is carried out. +* `z/depth/features`: Each (y,x) position in the base area has a vector +associated to it, which goes into this dimension. +* `spatial_dims`: Describes the `n` spatial dimensions that define the base +area that the window moves across. + +The `rhs` argument is a rank n+2 array describing the convolutional +filter/kernel/window. The dimensions are, in this order: + +* `output-z`: The `z` dimension of the output. +* `input-z`: The size of this dimension times `feature_group_count` should +equal the size of the `z` dimension in lhs. +* `spatial_dims`: Describes the `n` spatial dimensions that define the n-d +window that moves across the base area. + +The `window_strides` argument specifies the stride of the convolutional window +in the spatial dimensions. For example, if the stride in the first spatial +dimension is 3, then the window can only be placed at coordinates where the +first spatial index is divisible by 3. + +The `padding` argument specifies the amount of zero padding to be applied to the +base area. The amount of padding can be negative -- the absolute value of +negative padding indicates the number of elements to remove from the specified +dimension before doing the convolution. `padding[0]` specifies the padding for +dimension `y` and `padding[1]` specifies the padding for dimension `x`. Each +pair has the low padding as the first element and the high padding as the second +element. The low padding is applied in the direction of lower indices while the +high padding is applied in the direction of higher indices. For example, if +`padding[1]` is `(2,3)` then there will be a padding by 2 zeroes on the left and +by 3 zeroes on the right in the second spatial dimension. Using padding is +equivalent to inserting those same zero values into the input (`lhs`) before +doing the convolution. + +The `lhs_dilation` and `rhs_dilation` arguments specify the dilation factor to +be applied to the lhs and rhs, respectively, in each spatial dimension. If the +dilation factor in a spatial dimension is d, then d-1 holes are implicitly +placed between each of the entries in that dimension, increasing the size of the +array. The holes are filled with a no-op value, which for convolution means +zeroes. + +Dilation of the rhs is also called atrous convolution. For more details, see +`tf.nn.atrous_conv2d`. Dilation of the lhs is also called transposed +convolution. For more details, see `tf.nn.conv2d_transpose`. + +The `feature_group_count` argument (default value 1) can be used for grouped +convolutions. `feature_group_count` needs to be a divisor of both the input and +the output feature dimension. If `feature_group_count` is greater than 1, it +means that conceptually the input and output feature dimension and the `rhs` +output feature dimension are split evenly into `feature_group_count` many +groups, each group consisting of a consecutive subsequence of features. The +input feature dimension of `rhs` needs to be equal to the `lhs` input feature +dimension divided by `feature_group_count` (so it already has the size of a +group of input features). The i-th groups are used together to compute +`feature_group_count` many separate convolutions. The results of these +convolutions are concatenated together in the output feature dimension. + +For depthwise convolution the `feature_group_count` argument would be set to the +input feature dimension, and the filter would be reshaped from +`[filter_height, filter_width, in_channels, channel_multiplier]` to +`[filter_height, filter_width, 1, in_channels * channel_multiplier]`. For more +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` (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: + +* `batch`: The size of this dimension times `batch_group_count` should equal + the size of the `batch` dimension in lhs. +* `z`: Same size as `output-z` on the kernel (`rhs`). +* `spatial_dims`: One value for each valid placement of the convolutional + window. + +The valid placements of the convolutional window are determined by the strides +and the size of the base area after padding. + +To describe what a convolution does, consider a 2d convolution, and pick some +fixed `batch`, `z`, `y`, `x` coordinates in the output. Then `(y,x)` is a +position of a corner of the window within the base area (e.g. the upper left +corner, depending on how you interpret the spatial dimensions). We now have a 2d +window, taken from the base area, where each 2d point is associated to a 1d +vector, so we get a 3d box. From the convolutional kernel, since we fixed the +output coordinate `z`, we also have a 3d box. The two boxes have the same +dimensions, so we can take the sum of the element-wise products between the two +boxes (similar to a dot product). That is the output value. + +Note that if `output-z` is e.g., 5, then each position of the window produces 5 +values in the output into the `z` dimension of the output. These values differ +in what part of the convolutional kernel is used - there is a separate 3d box of +values used for each `output-z` coordinate. So you could think of it as 5 +separate convolutions with a different filter for each of them. + +Here is pseudo-code for a 2d convolution with padding and striding: + +``` +for (b, oz, oy, ox) { // output coordinates +value = 0; +for (iz, ky, kx) { // kernel coordinates and input z +iy = oy*stride_y + ky - pad_low_y; +ix = ox*stride_x + kx - pad_low_x; +if ((iy, ix) inside the base area considered without padding) { +value += input(b, iz, iy, ix) * kernel(oz, iz, ky, kx); +} +} +output(b, oz, oy, ox) = value; +} +``` + +## ConvertElementType + +See also +[`XlaBuilder::ConvertElementType`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + +Similar to an element-wise `static_cast` in C++, performs an element-wise +conversion operation from a data shape to a target shape. The dimensions must +match, and the conversion is an element-wise one; e.g. `s32` elements become +`f32` elements via an `s32`-to-`f32` conversion routine. + + `ConvertElementType(operand, new_element_type)` + +Arguments | Type | Semantics +------------------ | --------------- | --------------------------- +`operand` | `XlaOp` | array of type T with dims D +`new_element_type` | `PrimitiveType` | type U + +The dimensions of the operand and the target shape must match. The source and +destination element types must not be tuples. + +A conversion such as `T=s32` to `U=f32` will perform a normalizing int-to-float +conversion routine such as round-to-nearest-even. + +> Note: The precise float-to-int and visa-versa conversions are currently +> unspecified, but may become additional arguments to the convert operation in +> the future. Not all possible conversions have been implemented for all +>targets. + +``` +let a: s32[3] = {0, 1, 2}; +let b: f32[3] = convert(a, f32); +then b == f32[3]{0.0, 1.0, 2.0} +``` + +## CrossReplicaSum + +See also +[`XlaBuilder::CrossReplicaSum`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + +Computes a sum across replicas. + + `CrossReplicaSum(operand)` + +Arguments | Type | Semantics +--------- | ------- | ----------------------------- +`operand` | `XlaOp` | Array to sum across replicas. +| `replica_group_ids` | `int64` vector | Group ID for each replica. | + +The output shape is the same as the input shape. For example, if there are two +replicas and the operand has the value `(1.0, 2.5)` and `(3.0, 5.25)` +respectively on the two replicas, then the output value from this op will be +`(4.0, 7.75)` on both replicas. + +`replica_group_ids` identifies the group ID of each replica. The group ID must +either be empty (all replicas belong to a single group), or contain the same +number of elements as the number of replicas. For example, if +`replica_group_ids` = {0, 1, 2, 3, 0, 1, 2, 3} has eight replicas, there are +four subgroups of replica IDs: {0, 4}, {1, 5}, {2, 6}, and {3, 7}. The size of +each subgroup *must* be identical, so, for example, using: +`replica_group_ids` = {0, 1, 2, 0} for four replicas is invalid. + +Computing the result of CrossReplicaSum requires having one input from each +replica, so if one replica executes a CrossReplicaSum node more times than +another, then the former replica will wait forever. Since the replicas are all +running the same program, there are not a lot of ways for that to happen, but it +is possible when a while loop's condition depends on data from infeed and the +data that is infed causes the while loop to iterate more times on one replica +than another. + +## CustomCall + +See also +[`XlaBuilder::CustomCall`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + +Call a user-provided function within a computation. + + `CustomCall(target_name, args..., shape)` + +| Arguments | Type | Semantics | +| ------------- | ---------------------- | --------------------------------- | +| `target_name` | `string` | Name of the function. A call | +: : : instruction will be emitted which : +: : : targets this symbol name. : +| `args` | sequence of N `XlaOp`s | N arguments of arbitrary type, | +: : : which will be passed to the : +: : : function. : +| `shape` | `Shape` | Output shape of the function | + +The function signature is the same, regardless of the arity or type of args: + +``` +extern "C" void target_name(void* out, void** in); +``` + +For example, if CustomCall is used as follows: + +``` +let x = f32[2] {1,2}; +let y = f32[2x3] {{10, 20, 30}, {40, 50, 60}}; + +CustomCall("myfunc", {x, y}, f32[3x3]) +``` + +Here is an example of an implementation of `myfunc`: + +``` +extern "C" void myfunc(void* out, void** in) { +float (&x)[2] = *static_cast(in[0]); +float (&y)[2][3] = *static_cast(in[1]); +EXPECT_EQ(1, x[0]); +EXPECT_EQ(2, x[1]); +EXPECT_EQ(10, y[0][0]); +EXPECT_EQ(20, y[0][1]); +EXPECT_EQ(30, y[0][2]); +EXPECT_EQ(40, y[1][0]); +EXPECT_EQ(50, y[1][1]); +EXPECT_EQ(60, y[1][2]); +float (&z)[3][3] = *static_cast(out); +z[0][0] = x[1] + y[1][0]; +// ... +} +``` + +The user-provided function must not have side-effects and its execution must be +idempotent. + +> Note: The opaque nature of the user-provided function restricts optimization +> opportunities for the compiler. Try to express your computation in terms of +> native XLA ops whenever possible; only use CustomCall as a last resort. + +## Dot + +See also +[`XlaBuilder::Dot`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + + `Dot(lhs, rhs)` + +Arguments | Type | Semantics +--------- | ------- | --------------- +`lhs` | `XlaOp` | array of type T +`rhs` | `XlaOp` | array of type T + +The exact semantics of this operation depend on the ranks of the operands: + +| Input | Output | Semantics | +| ----------------------- | --------------------- | ----------------------- | +| vector [n] `dot` vector | scalar | vector dot product | +: [n] : : : +| matrix [m x k] `dot` | vector [m] | matrix-vector | +: vector [k] : : multiplication : +| matrix [m x k] `dot` | matrix [m x n] | matrix-matrix | +: matrix [k x n] : : multiplication : + +The operation performs sum of products over the last dimension of `lhs` and the +one-before-last dimension of `rhs`. These are the "contracted" dimensions. The +contracted dimensions of `lhs` and `rhs` must be of the same size. In practice, +it can be used to perform dot products between vectors, vector/matrix +multiplications or matrix/matrix multiplications. + +## DotGeneral + +See also +[`XlaBuilder::DotGeneral`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + + `DotGeneral(lhs, rhs, dimension_numbers)` + +Arguments | Type | Semantics +------------------- | --------------------- | --------------- +`lhs` | `XlaOp` | array of type T +`rhs` | `XlaOp` | array of type T +`dimension_numbers` | `DotDimensionNumbers` | array of type T + +As Dot, but allows contracting and batch dimension numbers to be specified for +both the 'lhs' and 'rhs'. + +| DotDimensionNumbers Fields | Type | Semantics +| --------- | ----------------------- | --------------- +| 'lhs_contracting_dimensions' | repeated int64 | 'lhs' contracting dimension numbers | +| 'rhs_contracting_dimensions' | repeated int64 | 'rhs' contracting dimension numbers | +| 'lhs_batch_dimensions' | repeated int64 | 'lhs' batch dimension numbers | +| 'rhs_batch_dimensions' | repeated int64 | 'rhs' batch dimension numbers | + +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 and but must have the same dimension sizes. + +Example with contracting dimension numbers: + +``` +lhs = { {1.0, 2.0, 3.0}, +{4.0, 5.0, 6.0} } + +rhs = { {1.0, 1.0, 1.0}, +{2.0, 2.0, 2.0} } + +DotDimensionNumbers dnums; +dnums.add_lhs_contracting_dimensions(1); +dnums.add_rhs_contracting_dimensions(1); + +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 sizes. + +Example with batch dimension numbers (batch size 2, 2x2 matrices): + +``` +lhs = { { {1.0, 2.0}, +{3.0, 4.0} }, +{ {5.0, 6.0}, +{7.0, 8.0} } } + +rhs = { { {1.0, 0.0}, +{0.0, 1.0} }, +{ {1.0, 0.0}, +{0.0, 1.0} } } + +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) -> { { {1.0, 2.0}, +{3.0, 4.0} }, +{ {5.0, 6.0}, +{7.0, 8.0} } } +``` + +| Input | Output | Semantics | +| ----------------------------------- | ----------------- | ---------------- | +| [b0, m, k] `dot` [b0, k, n] | [b0, m, n] | batch matmul | +| [b0, b1, m, k] `dot` [b0, b1, k, n] | [b0, b1, m, n] | batch matmul | + +It follows that the resulting dimension number starts with the batch dimension, +then the 'lhs' non-contracting/non-batch dimension, and finally the 'rhs' +non-contracting/non-batch dimension. + +## DynamicSlice + +See also +[`XlaBuilder::DynamicSlice`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + +DynamicSlice extracts a sub-array from the input array at dynamic +`start_indices`. The size of the slice in each dimension is passed in +`size_indices`, 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 `operand`. + + `DynamicSlice(operand, start_indices, size_indices)` + +| 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: + +``` +start_indices[i] = clamp(start_indices[i], 0, operand.dimension_size[i] - size_indices[i]) +``` + +This ensures that the extracted slice is always in-bounds with respect to the +operand array. If the slice is in-bounds before the transformation is applied, +the transformation has no effect. + +1-dimensional example: + +``` +let a = {0.0, 1.0, 2.0, 3.0, 4.0} +let s = {2} + +DynamicSlice(a, s, {2}) produces: +{2.0, 3.0} +``` + +2-dimensional example: + +``` +let b = +{ {0.0, 1.0, 2.0}, +{3.0, 4.0, 5.0}, +{6.0, 7.0, 8.0}, +{9.0, 10.0, 11.0} } +let s = {2, 1} + +DynamicSlice(b, s, {2, 2}) produces: +{ { 7.0, 8.0}, +{10.0, 11.0} } +``` +## DynamicUpdateSlice + +See also +[`XlaBuilder::DynamicUpdateSlice`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + +DynamicUpdateSlice generates a result which is the value of the input array +`operand`, with a slice `update` overwritten at `start_indices`. +The shape of `update` determines the shape of the sub-array of the result which +is updated. +The shape of `start_indices` must be rank == 1, with dimension size equal to +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` | 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: + +``` +start_indices[i] = clamp(start_indices[i], 0, operand.dimension_size[i] - update.dimension_size[i]) +``` + +This ensures that the updated slice is always in-bounds with respect to the +operand array. If the slice is in-bounds before the transformation is applied, +the transformation has no effect. + +1-dimensional example: + +``` +let a = {0.0, 1.0, 2.0, 3.0, 4.0} +let u = {5.0, 6.0} +let s = {2} + +DynamicUpdateSlice(a, u, s) produces: +{0.0, 1.0, 5.0, 6.0, 4.0} +``` + +2-dimensional example: + +``` +let b = +{ {0.0, 1.0, 2.0}, +{3.0, 4.0, 5.0}, +{6.0, 7.0, 8.0}, +{9.0, 10.0, 11.0} } +let u = +{ {12.0, 13.0}, +{14.0, 15.0}, +{16.0, 17.0} } + +let s = {1, 1} + +DynamicUpdateSlice(b, u, s) produces: +{ {0.0, 1.0, 2.0}, +{3.0, 12.0, 13.0}, +{6.0, 14.0, 15.0}, +{9.0, 16.0, 17.0} } +``` + +## Element-wise binary arithmetic operations + +See also +[`XlaBuilder::Add`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + +A set of element-wise binary arithmetic operations is supported. + + `Op(lhs, rhs)` + +Where `Op` is one of `Add` (addition), `Sub` (subtraction), `Mul` +(multiplication), `Div` (division), `Rem` (remainder), `Max` (maximum), `Min` +(minimum), `LogicalAnd` (logical AND), or `LogicalOr` (logical OR). + +Arguments | Type | Semantics +--------- | ------- | ---------------------------------------- +`lhs` | `XlaOp` | left-hand-side operand: array of type T +`rhs` | `XlaOp` | right-hand-side operand: array of type T + +The arguments' shapes have to be either similar or compatible. See the +[broadcasting](broadcasting.md) documentation about what it means for shapes to +be compatible. The result of an operation has a shape which is the result of +broadcasting the two input arrays. In this variant, operations between arrays of +different ranks are *not* supported, unless one of the operands is a scalar. + +When `Op` is `Rem`, the sign of the result is taken from the dividend, and the +absolute value of the result is always less than the divisor's absolute value. + +Integer division overflow (signed/unsigned division/remainder by zero or signed +division/remainder of `INT_SMIN` with `-1`) produces an implementation defined +value. + +An alternative variant with different-rank broadcasting support exists for these +operations: + + `Op(lhs, rhs, broadcast_dimensions)` + +Where `Op` is the same as above. This variant of the operation should be used +for arithmetic operations between arrays of different ranks (such as adding a +matrix to a vector). + +The additional `broadcast_dimensions` operand is a slice of integers used to +expand the rank of the lower-rank operand up to the rank of the higher-rank +operand. `broadcast_dimensions` maps the dimensions of the lower-rank shape to +the dimensions of the higher-rank shape. The unmapped dimensions of the expanded +shape are filled with dimensions of size one. Degenerate-dimension broadcasting +then broadcasts the shapes along these degenerate dimensions to equalize the +shapes of both operands. The semantics are described in detail on the +[broadcasting page](broadcasting.md). + +## Element-wise comparison operations + +See also +[`XlaBuilder::Eq`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + +A set of standard element-wise binary comparison operations is supported. Note +that standard IEEE 754 floating-point comparison semantics apply when comparing +floating-point types. + + `Op(lhs, rhs)` + +Where `Op` is one of `Eq` (equal-to), `Ne` (not equal-to), `Ge` +(greater-or-equal-than), `Gt` (greater-than), `Le` (less-or-equal-than), `Lt` +(less-than). + +Arguments | Type | Semantics +--------- | ------- | ---------------------------------------- +`lhs` | `XlaOp` | left-hand-side operand: array of type T +`rhs` | `XlaOp` | right-hand-side operand: array of type T + +The arguments' shapes have to be either similar or compatible. See the +[broadcasting](broadcasting.md) documentation about what it means for shapes to +be compatible. The result of an operation has a shape which is the result of +broadcasting the two input arrays with the element type `PRED`. In this variant, +operations between arrays of different ranks are *not* supported, unless one of +the operands is a scalar. + +An alternative variant with different-rank broadcasting support exists for these +operations: + + `Op(lhs, rhs, broadcast_dimensions)` + +Where `Op` is the same as above. This variant of the operation should be used +for comparison operations between arrays of different ranks (such as adding a +matrix to a vector). + +The additional `broadcast_dimensions` operand is a slice of integers specifying +the dimensions to use for broadcasting the operands. The semantics are described +in detail on the [broadcasting page](broadcasting.md). + +## Element-wise unary functions + +XlaBuilder supports these element-wise unary functions: + +`Abs(operand)` Element-wise abs `x -> |x|`. + +`Ceil(operand)` Element-wise ceil `x -> ⌈x⌉`. + +`Cos(operand)` Element-wise cosine `x -> cos(x)`. + +`Exp(operand)` Element-wise natural exponential `x -> e^x`. + +`Floor(operand)` Element-wise floor `x -> ⌊x⌋`. + +`IsFinite(operand)` Tests whether each element of `operand` is finite, +i.e., is not positive or negative infinity, and is not `NaN`. Returns an array +of `PRED` values with the same shape as the input, where each element is `true` +if and only if the corresponding input element is finite. + +`Log(operand)` Element-wise natural logarithm `x -> ln(x)`. + +`LogicalNot(operand)` Element-wise logical not `x -> !(x)`. + +`Neg(operand)` Element-wise negation `x -> -x`. + +`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}$$ + +using the comparison operator of the element type of `operand`. + +`Tanh(operand)` Element-wise hyperbolic tangent `x -> tanh(x)`. + + +Arguments | Type | Semantics +--------- | ------- | --------------------------- +`operand` | `XlaOp` | The operand to the function + +The function is applied to each element in the `operand` array, resulting in an +array with the same shape. It is allowed for `operand` to be a scalar (rank 0). + +## Gather + +The XLA gather operation stitches together several slices (each slice at a +potentially different runtime offset) of an input array. + +### General Semantics + +See also +[`XlaBuilder::Gather`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). +For a more intuitive description, see the "Informal Description" section below. + + `gather(operand, start_indices, offset_dims, collapsed_slice_dims, slice_sizes, start_index_map)` + +| Arguments | Type | Semantics | +| ---------------------- | ------------------- | ----------------------------- | +| `operand` | `XlaOp` | The array we’re gathering | +: : : from. : +| `start_indices` | `XlaOp` | Array containing the starting | +: : : indices of the slices we : +: : : gather. : +| `index_vector_dim` | `int64` | The dimension in | +: : : `start_indices` that : +: : : "contains" the starting : +: : : indices. See below for a : +: : : detailed description. : +| `offset_dims` | `ArraySlice` | The set of dimensions in the | +: : : output shape that offset into : +: : : a array sliced from operand. : +| `slice_sizes` | `ArraySlice` | `slice_sizes[i]` is the | +: : : bounds for the slice on : +: : : dimension `i`. : +| `collapsed_slice_dims` | `ArraySlice` | The set of dimensions in each | +: : : \: slice that are collapsed : +: : : away. These dimensions must : +: : : have size 1. : +| `start_index_map` | `ArraySlice` | A map that describes how to | +: : : map indices in : +: : : `start_indices` to legal : +: : : indices into operand. : + +For convenience, we label dimensions in the output array not in `offset_dims` +as `batch_dims`. + +The output is an array of rank `batch_dims.size` + `operand.rank` - +`collapsed_slice_dims`.size. + +If `index_vector_dim` is equal to `start_indices.rank` we implicitly consider +`start_indices` to have a trailing `1` dimension (i.e. if `start_indices` was of +shape `[6,7]` and `index_vector_dim` is `2` then we implicitly consider the +shape of `start_indices` to be `[6,7,1]`). + +The bounds for the output array along dimension `i` is computed as follows: + +1. If `i` is present in `batch_dims` (i.e. is equal to `batch_dims[k]` for +some `k`) then we pick the corresponding dimension bounds out of +`start_indices.shape`, skipping `index_vector_dim` (i.e. pick +`start_indices.shape.dims`[`k`] if `k` < `index_vector_dim` and +`start_indices.shape.dims`[`k`+`1`] otherwise). + +2. If `i` is present in `offset_dims` (i.e. equal to `offset_dims`[`k`] for +some `k`) then we pick the corresponding bound out of `slice_sizes` after +accounting for `collapsed_slice_dims` (i.e. we pick +`adjusted_slice_sizes`[`k`] where `adjusted_slice_sizes` is `slice_sizes` +with the bounds at indices `collapsed_slice_dims` removed). + +Formally, the operand index `In` corresponding to an output index `Out` is +computed as follows: + +1. Let `G` = { `Out`[`k`] for `k` in `batch_dims` }. Use `G` to slice out +vector `S` such that `S`[`i`] = `start_indices`[Combine(`G`, `i`)] where +Combine(A, b) inserts b at position `index_vector_dim` into A. Note that +this is well defined even if `G` is empty -- if `G` is empty then `S` = +`start_indices`. + +2. Create a starting index, `S``in`, into `operand` using `S` by +scattering `S` using `start_index_map`. More precisely: +1. `S``in`[`start_index_map`[`k`]] = `S`[`k`] if `k` < +`start_index_map.size`. +2. `S``in`[`_`] = `0` otherwise. + +3. Create an index `O``in` into `operand` by scattering the indices +at the offset dimensions in `Out` according to the `collapsed_slice_dims` +set. More precisely: +1. `O``in`[`expand_offset_dims`(`k`)] = +`Out`[`offset_dims`[`k`]] if `k` < `offset_dims.size` +(`expand_offset_dims` is defined below). +2. `O``in`[`_`] = `0` otherwise. +4. `In` is `O``in` + `S``in` where + is element-wise +addition. + +`expand_offset_dims` is the monotonic function with domain [`0`, `offset.size`) +and range [`0`, `operand.rank`) \ `collapsed_slice_dims`. So if, e.g., +`offset.size` is `4`, `operand.rank` is `6` and `collapsed_slice_dims` is {`0`, +`2`} then `expand_offset_dims` is {`0`→`1`, `1`→`3`, `2`→`4`, `3`→`5`}. + +### Informal Description and Examples + +Informally, every index `Out` in the output array corresponds to an element `E` +in the operand array, computed as follows: + +- We use the batch dimensions in `Out` to look up a starting index from +`start_indices`. + +- We use `start_index_map` to map the starting index (which may have size less +than operand.rank) to a "full" starting index into operand. + +- We dynamic-slice out a slice with size `slice_sizes` using the full starting +index. + +- We reshape the slice by collapsing the `collapsed_slice_dims` dimensions. +Since all collapsed slice dimensions have to have bound 1 this reshape is +always legal. + +- We use the offset dimensions in `Out` to index into this slice to get the +input element, `E`, corresponding to output index `Out`. + +`index_vector_dim` is set to `start_indices.rank` - `1` in all of the +examples that follow. More interesting values for `index_vector_dim` does not +change the operation fundamentally, but makes the visual representation more +cumbersome. + +To get an intuition on how all of the above fits together, let's look at an +example that gathers 5 slices of shape `[8,6]` from a `[16,11]` array. The +position of a slice into the `[16,11]` array can be represented as an index +vector of shape `S64[2]`, so the set of 5 positions can be represented as a +`S64[5,2]` array. + +The behavior of the gather operation can then be depicted as an index +transformation that takes [`G`,`O``0`,`O``1`], an index in +the output shape, and maps it to an element in the input array in the following +way: + +
+ +
+ +We first select an (`X`,`Y`) vector from the gather indices array using `G`. +The element in the output array at index +[`G`,`O``0`,`O``1`] is then the element in the input +array at index [`X`+`O``0`,`Y`+`O``1`]. + +`slice_sizes` is `[8,6]`, which decides the range of W`0` and +W`1`, and this in turn decides the bounds of the slice. + +This gather operation acts as a batch dynamic slice with `G` as the batch +dimension. + +The gather indices may be multidimensional. For instance, a more general +version of the example above using a "gather indices" array of shape `[4,5,2]` +would translate indices like this: + +
+ +
+ +Again, this acts as a batch dynamic slice `G``0` and +`G``1` as the batch dimensions. The slice size is still `[8,6]`. + +The gather operation in XLA generalizes the informal semantics outlined above in +the following ways: + +1. We can configure which dimensions in the output shape are the offset +dimensions (dimensions containing `O``0`, `O``1` in +the last example). The output batch dimensions (dimensions containing +`G``0`, `G``1` in the last example) are defined to be +the output dimensions that are not offset dimensions. + +2. The number of output offset dimensions explicitly present in the output +shape may be smaller than the input rank. These "missing" dimensions, which +are listed explicitly as `collapsed_slice_dims`, must have a slice size of +`1`. Since they have a slice size of `1` the only valid index for them is +`0` and eliding them does not introduce ambiguity. + +3. The slice extracted from the "Gather Indices" array ((`X`, `Y`) in the last +example) may have fewer elements than the input array rank, and an explicit +mapping dictates how the index should be expanded to have the same rank as +the input. + +As a final example, we use (2) and (3) to implement `tf.gather_nd`: + +
+ +
+ +`G``0` and `G``1` are used to slice out a starting index +from the gather indices array as usual, except the starting index has only one +element, `X`. Similarly, there is only one output offset index with the value +`O``0`. However, before being used as indices into the input array, +these are expanded in accordance to "Gather Index Mapping" (`start_index_map` in +the formal description) and "Offset Mapping" (`expand_offset_dims` in the formal +description) into [`X`,`0`] and [`0`,`O``0`] respectively, adding up +to [`X`,`O``0`]. In other words, the output index +[`G``0`,`G``1`,`O``0`] maps to the input index +[`GatherIndices`[`G``0`,`G``1`,`0`],`X`] which gives us +the semantics for `tf.gather_nd`. + +`slice_sizes` for this case is `[1,11]`. Intuitively this means that every +index `X` in the gather indices array picks an entire row and the result is the +concatenation of all these rows. + +## GetDimensionSize + +See also +[`XlaBuilder::GetDimensionSize`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + +Returns the size of the given dimension of the operand. The operand must be +array shaped. + + `GetDimensionSize(operand, dimension)` + +| Arguments | Type | Semantics | +| ----------- | ------- | --------------------------------------------------- | +| `operand` | `XlaOp` | n dimensional input array | +| `dimension` | `int64` | A value in the interval `[0, n)` that specifies the | +: : : dimension : + +## GetTupleElement + +See also +[`XlaBuilder::GetTupleElement`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + +Indexes into a tuple with a compile-time-constant value. + +The value must be a compile-time-constant so that shape inference can determine +the type of the resulting value. + +This is analogous to `std::get(t)` in C++. Conceptually: + +``` +let v: f32[10] = f32[10]{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}; +let s: s32 = 5; +let t: (f32[10], s32) = tuple(v, s); +let element_1: s32 = gettupleelement(t, 1); // Inferred shape matches s32. +``` + +See also `tf.tuple`. + +## Infeed + +See also +[`XlaBuilder::Infeed`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + + `Infeed(shape)` + +| Argument | Type | Semantics | +| -------- | ------- | ----------------------------------------------------- | +| `shape` | `Shape` | Shape of the data read from the Infeed interface. The | +: : : layout field of the shape must be set to match the : +: : : layout of the data sent to the device; otherwise its : +: : : behavior is undefined. : + +Reads a single data item from the implicit Infeed streaming interface of the +device, interpreting the data as the given shape and its layout, and returns a +`XlaOp` of the data. Multiple Infeed operations are allowed in a +computation, but there must be a total order among the Infeed operations. For +example, two Infeeds in the code below have a total order since there is a +dependency between the while loops. + +``` +result1 = while (condition, init = init_value) { +Infeed(shape) +} + +result2 = while (condition, init = result1) { +Infeed(shape) +} +``` + +Nested tuple shapes are not supported. For an empty tuple shape, the Infeed +operation is effectively a no-op and proceeds without reading any data from the +Infeed of the device. + +> Note: We plan to allow multiple Infeed operations without a total order, in +> which case the compiler will provide information about how the Infeed +> operations are serialized in the compiled program. + +## Iota + + `Iota()` + +Builds a constant literal on device rather than a potentially large host +transfer. Creates a rank 1 array of values starting at zero and incrementing by +one. For floating-point types, the produced array is equivalent to +`ConvertElementType(Iota(...))` where the `Iota` is of integral type and the +conversion is to the floating-point type. + +Arguments | Type | Semantics +---------------- | --------------- | ------------------------------------ +`type` | `PrimitiveType` | type U +`size` | `int64` | The number of elements in the array. +`iota_dimension` | `int64` | The dimension to increment along. + +## Map + +See also +[`XlaBuilder::Map`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + + `Map(operands..., computation)` + +| Arguments | Type | Semantics | +| ----------------- | ---------------------- | ------------------------------ | +| `operands` | sequence of N `XlaOp`s | N arrays of types T_0..T_{N-1} | +| `computation` | `XlaComputation` | computation of type `T_0, T_1, | +: : : ..., T_{N + M -1} -> S` with N : +: : : parameters of type T and M of : +: : : arbitrary type : +| `dimensions` | `int64` array | array of map dimensions | + +Applies a scalar function over the given `operands` arrays, producing an array +of the same dimensions where each element is the result of the mapped function +applied to the corresponding elements in the input arrays. + +The mapped function is an arbitrary computation with the restriction that it has +N inputs of scalar type `T` and a single output with type `S`. The output has +the same dimensions as the operands except that the element type T is replaced +with S. + +For example: `Map(op1, op2, op3, computation, par1)` maps `elem_out <- +computation(elem1, elem2, elem3, par1)` at each (multi-dimensional) index in the +input arrays to produce the output array. + +## Pad + +See also +[`XlaBuilder::Pad`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + + `Pad(operand, padding_value, padding_config)` + +| Arguments | Type | Semantics | +| ---------------- | --------------- | --------------------------------------- | +| `operand` | `XlaOp` | array of type `T` | +| `padding_value` | `XlaOp` | scalar of type `T` to fill in the added | +: : : padding : +| `padding_config` | `PaddingConfig` | padding amount on both edges (low, | +: : : high) and between the elements of each : +: : : dimension : + +Expands the given `operand` array by padding around the array as well as between +the elements of the array with the given `padding_value`. `padding_config` +specifies the amount of edge padding and the interior padding for each +dimension. + +`PaddingConfig` is a repeated field of `PaddingConfigDimension`, which contains +three fields for each dimension: `edge_padding_low`, `edge_padding_high`, and +`interior_padding`. + +`edge_padding_low` and `edge_padding_high` specify the amount of padding added +at the low-end (next to index 0) and the high-end (next to the highest index) of +each dimension respectively. The amount of edge padding can be negative -- the +absolute value of negative padding indicates the number of elements to remove +from the specified dimension. + +`interior_padding` specifies the amount of padding added between any two +elements in each dimension; it may not be negative. Interior padding occurs +logically before edge padding, so in the case of negative edge padding, elements +are removed from the interior-padded operand. + +This operation is a no-op if the edge padding pairs are all (0, 0) and the +interior padding values are all 0. The figure below shows examples of different +`edge_padding` and `interior_padding` values for a two-dimensional array. + +
+ +
+ +## Recv + +See also +[`XlaBuilder::Recv`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + + `Recv(shape, channel_handle)` + +| Arguments | Type | Semantics | +| ---------------- | --------------- | ------------------------------------ | +| `shape` | `Shape` | shape of the data to receive | +| `channel_handle` | `ChannelHandle` | unique identifier for each send/recv pair | + +Receives data of the given shape from a `Send` instruction in another +computation that shares the same channel handle. Returns a +XlaOp for the received data. + +The client API of `Recv` operation represents synchronous communication. +However, the instruction is internally decomposed into 2 HLO instructions +(`Recv` and `RecvDone`) to enable asynchronous data transfers. See also +[`HloInstruction::CreateRecv` and `HloInstruction::CreateRecvDone`](https://www.tensorflow.org/code/tensorflow/compiler/xla/service/hlo_instruction.h). + +`Recv(const Shape& shape, int64 channel_id)` + +Allocates resources required to receive data from a `Send` instruction with the +same channel_id. Returns a context for the allocated resources, which is used +by a following `RecvDone` instruction to wait for the completion of the data +transfer. The context is a tuple of {receive buffer (shape), request identifier +(U32)} and it can only be used by a `RecvDone` instruction. + + `RecvDone(HloInstruction context)` + +Given a context created by a `Recv` instruction, waits for the data transfer to +complete and returns the received data. + +## Reduce + +See also +[`XlaBuilder::Reduce`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + +Applies a reduction function to one or more arrays in parallel. + + `Reduce(operands..., init_values..., computation, dimensions)` + +Arguments | Type | Semantics +------------- | --------------------- | --------------------------------------- +`operands` | Sequence of N `XlaOp` | N arrays of types `T_0, ..., T_N`. +`init_values` | Sequence of N `XlaOp` | N scalars of types `T_0, ..., T_N`. +`computation` | `XlaComputation` | computation of type + : : `T_0, ..., T_N, T_0, ..., T_N -> Collate(T_0, ..., T_N)` +`dimensions` | `int64` array | unordered array of dimensions to reduce + +Where: +* N is required to be greater or equal to 1. +* All input arrays must have the same dimensions. +* If `N = 1`, `Collate(T)` is `T`. +* If `N > 1`, `Collate(T_0, ..., T_N)` is a tuple of `N` elements of type `T`. + +The output of the op is `Collate(Q_0, ..., Q_N)` where `Q_i` is an array of type +`T_i`, the dimensions of which are described below. + +This operation reduces one or more dimensions of each input array into scalars. +The rank of each returned array is `rank(operand) - len(dimensions)`. +`init_value` is the initial value used for every reduction and may be inserted +anywhere during computation by the back-end. In most cases, `init_value` is an +identity of the reduction function (for example, 0 for addition). The applied +`computation` is always passed the `init_value` on the left-hand side. + +The evaluation order of the reduction function is arbitrary and may be +non-deterministic. Therefore, the reduction function should not be overly +sensitive to reassociation. + +Some reduction functions like addition are not strictly associative for floats. +However, if the range of the data is limited, floating-point addition is close +enough to being associative for most practical uses. It is possible to conceive +of some completely non-associative reductions, however, and these will produce +incorrect or unpredictable results in XLA reductions. + +As an example, when reducing across one dimension in a single 1D array with +values [10, 11, 12, 13], with reduction function `f` (this is `computation`) +then that could be computed as + +`f(10, f(11, f(12, f(init_value, 13)))` + +but there are also many other possibilities, e.g. + +`f(init_value, f(f(10, f(init_value, 11)), f(f(init_value, 12), f(init_value, 13))))` + +The following is a rough pseudo-code example of how reduction could be +implemented, using summation as the reduction computation with an initial value +of 0. + +```python +result_shape <- remove all dims in dimensions from operand_shape + +# Iterate over all elements in result_shape. The number of r's here is equal +# to the rank of the result +for r0 in range(result_shape[0]), r1 in range(result_shape[1]), ...: + # Initialize this result element + result[r0, r1...] <- 0 + + # Iterate over all the reduction dimensions + for d0 in range(dimensions[0]), d1 in range(dimensions[1]), ...: + # Increment the result element with the value of the operand's element. + # The index of the operand's element is constructed from all ri's and di's + # in the right order (by construction ri's and di's together index over the + # whole operand shape). + result[r0, r1...] += operand[ri... di] +``` + +Here's an example of reducing a 2D array (matrix). The shape has rank 2, +dimension 0 of size 2 and dimension 1 of size 3: + +
+ +
+ +Results of reducing dimensions 0 or 1 with an "add" function: + +
+ +
+ +Note that both reduction results are 1D arrays. The diagram shows one as column +and another as row just for visual convenience. + +For a more complex example, here is a 3D array. Its rank is 3, dimension 0 of +size 4, dimension 1 of size 2 and dimension 2 of size 3. For simplicity, the +values 1 to 6 are replicated across dimension 0. + +
+ +
+ +Similarly to the 2D example, we can reduce just one dimension. If we reduce +dimension 0, for example, we get a rank-2 array where all values across +dimension 0 were folded into a scalar: + +```text +| 4 8 12 | +| 16 20 24 | +``` + +If we reduce dimension 2, we also get a rank-2 array where all values across +dimension 2 were folded into a scalar: + +```text +| 6 15 | +| 6 15 | +| 6 15 | +| 6 15 | +``` + +Note that the relative order between the remaining dimensions in the input is +preserved in the output, but some dimensions may get assigned new numbers (since +the rank changes). + +We can also reduce multiple dimensions. Add-reducing dimensions 0 and 1 produces +the 1D array `| 20 28 36 |`. + +Reducing the 3D array over all its dimensions produces the scalar `84`. + +When `N > 1`, reduce function application is slightly more complex, as it is +applied simultaneously to all inputs. For example, consider the following +reduction function, which can be used to compute the max and the argmax of a a +1-D array in parallel: + +``` +f: (Float, Int, Float, Int) -> Float, Int +f(max, argmax, value, index): + if value >= argmax: + return (value, index) + else: + return (max, argmax) +``` + +For 1-D Input arrays `V = Float[N], K = Int[N]`, and init values +`I_V = Float, I_K = Int`, the result `f_(N-1)` of reducing across the only +input dimension is equivalent to the following recursive application: +``` +f_0 = f(I_V, I_K, V_0, K_0) +f_1 = f(f_0.first, f_0.second, V_1, K_1) +... +f_(N-1) = f(f_(N-2).first, f_(N-2).second, V_(N-1), K_(N-1)) +``` + +Applying this reduction to an array of values, and an array of sequential +indices (i.e. iota), will co-iterate over the arrays, and return a tuple +containing the maximal value and the matching index. + +## ReducePrecision + +See also +[`XlaBuilder::ReducePrecision`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + +Models the effect of converting floating-point values to a lower-precision +format (such as IEEE-FP16) and back to the original format. The number of +exponent and mantissa bits in the lower-precision format can be specified +arbitrarily, although all bit sizes may not be supported on all hardware +implementations. + + `ReducePrecision(operand, mantissa_bits, exponent_bits)` + +Arguments | Type | Semantics +--------------- | ------- | ------------------------------------------------- +`operand` | `XlaOp` | array of floating-point type `T`. +`exponent_bits` | `int32` | number of exponent bits in lower-precision format +`mantissa_bits` | `int32` | number of mantissa bits in lower-precision format + +The result is an array of type `T`. The input values are rounded to the nearest +value representable with the given number of mantissa bits (using "ties to even" +semantics), and any values that exceed the range specified by the number of +exponent bits are clamped to positive or negative infinity. `NaN` values are +retained, although they may be converted to canonical `NaN` values. + +The lower-precision format must have at least one exponent bit (in order to +distinguish a zero value from an infinity, since both have a zero mantissa), and +must have a non-negative number of mantissa bits. The number of exponent or +mantissa bits may exceed the corresponding value for type `T`; the corresponding +portion of the conversion is then simply a no-op. + +## ReduceWindow + +See also +[`XlaBuilder::ReduceWindow`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + +Applies a reduction function to all elements in each window of the input +multi-dimensional array, producing an output multi-dimensional array with the +same number of elements as the number of valid positions of the window. A +pooling layer can be expressed as a `ReduceWindow`. Similar to +[`Reduce`](#reduce), the applied `computation` is always passed the `init_value` +on the left-hand side. + + `ReduceWindow(operand, init_value, computation, window_dimensions, +window_strides, padding)` + +| Arguments | Type | Semantics | +| ------------------- | ------------------- | -------------------------------- | +| `operand` | `XlaOp` | N dimensional array containing | +: : : elements of type T. This is the : +: : : base area on which the window is : +: : : placed. : +| `init_value` | `XlaOp` | Starting value for the | +: : : reduction. See [Reduce](#reduce) : +: : : for details. : +| `computation` | `XlaComputation` | Reduction function of type `T, T | +: : : -> T`, to apply to all elements : +: : : in each window : +| `window_dimensions` | `ArraySlice` | array of integers for window | +: : : dimension values : +| `window_strides` | `ArraySlice` | array of integers for window | +: : : stride values : +| `base_dilations` | `ArraySlice` | array of integers for base | +: : : dilation values : +| `window_dilations` | `ArraySlice` | array of integers for window | +: : : dilation values : +| `padding` | `Padding` | padding type for window | +: : : (Padding\:\:kSame or : +: : : Padding\:\:kValid) : + +Below code and figure shows an example of using `ReduceWindow`. Input is a +matrix of size [4x6] and both window_dimensions and window_stride_dimensions are +[2x3]. + +``` +// Create a computation for the reduction (maximum). +XlaComputation max; +{ + XlaBuilder builder(client_, "max"); + auto y = builder.Parameter(0, ShapeUtil::MakeShape(F32, {}), "y"); + auto x = builder.Parameter(1, ShapeUtil::MakeShape(F32, {}), "x"); + builder.Max(y, x); + max = builder.Build().ConsumeValueOrDie(); +} + +// Create a ReduceWindow computation with the max reduction computation. +XlaBuilder builder(client_, "reduce_window_2x3"); +auto shape = ShapeUtil::MakeShape(F32, {4, 6}); +auto input = builder.Parameter(0, shape, "input"); +builder.ReduceWindow( + input, + /*init_val=*/builder.ConstantLiteral(LiteralUtil::MinValue(F32)), + *max, + /*window_dimensions=*/{2, 3}, + /*window_stride_dimensions=*/{2, 3}, + Padding::kValid); +``` + +
+ +
+ +Stride of 1 in a dimension specifies that the position of a window in the +dimension is 1 element away from its adjacent window. In order to specify that +no windows overlap with each other, window_stride_dimensions should be equal to +window_dimensions. The figure below illustrates the use of two different stride +values. Padding is applied to each dimension of the input and the calculations +are the same as though the input came in with the dimensions it has after +padding. + +
+ +
+ +The evaluation order of the reduction function is arbitrary and may be +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. + +## Reshape + +See also +[`XlaBuilder::Reshape`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h) +and the [`Collapse`](#collapse) operation. + +Reshapes the dimensions of an array into a new configuration. + + `Reshape(operand, new_sizes)` + `Reshape(operand, dimensions, new_sizes)` + +Arguments | Type | Semantics +------------ | -------------- | --------------------------------------- +`operand` | `XlaOp` | array of type T +`dimensions` | `int64` vector | order in which dimensions are collapsed +`new_sizes` | `int64` vector | vector of sizes of new dimensions + +Conceptually, reshape first flattens an array into a one-dimensional vector of +data values, and then refines this vector into a new shape. The input arguments +are an arbitrary array of type T, a compile-time-constant vector of dimension +indices, and a compile-time-constant vector of dimension sizes for the result. +The values in the `dimension` vector, if given, must be a permutation of all of +T's dimensions; the default if not given is `{0, ..., rank - 1}`. The order of +the dimensions in `dimensions` is from slowest-varying dimension (most major) to +fastest-varying dimension (most minor) in the loop nest which collapses the +input array into a single dimension. The `new_sizes` vector determines the size +of the output array. The value at index 0 in `new_sizes` is the size of +dimension 0, the value at index 1 is the size of dimension 1, and so on. The +product of the `new_size` dimensions must equal the product of the operand's +dimension sizes. When refining the collapsed array into the multidimensional +array defined by `new_sizes`, the dimensions in `new_sizes` are ordered from +slowest varying (most major) and to fastest varying (most minor). + +For example, let v be an array of 24 elements: + +``` +let v = f32[4x2x3] {{{10, 11, 12}, {15, 16, 17}}, + {{20, 21, 22}, {25, 26, 27}}, + {{30, 31, 32}, {35, 36, 37}}, + {{40, 41, 42}, {45, 46, 47}}}; + +In-order collapse: +let v012_24 = Reshape(v, {0,1,2}, {24}); +then v012_24 == f32[24] {10, 11, 12, 15, 16, 17, 20, 21, 22, 25, 26, 27, + 30, 31, 32, 35, 36, 37, 40, 41, 42, 45, 46, 47}; + +let v012_83 = Reshape(v, {0,1,2}, {8,3}); +then v012_83 == f32[8x3] {{10, 11, 12}, {15, 16, 17}, + {20, 21, 22}, {25, 26, 27}, + {30, 31, 32}, {35, 36, 37}, + {40, 41, 42}, {45, 46, 47}}; + +Out-of-order collapse: +let v021_24 = Reshape(v, {1,2,0}, {24}); +then v012_24 == f32[24] {10, 20, 30, 40, 11, 21, 31, 41, 12, 22, 32, 42, + 15, 25, 35, 45, 16, 26, 36, 46, 17, 27, 37, 47}; + +let v021_83 = Reshape(v, {1,2,0}, {8,3}); +then v021_83 == f32[8x3] {{10, 20, 30}, {40, 11, 21}, + {31, 41, 12}, {22, 32, 42}, + {15, 25, 35}, {45, 16, 26}, + {36, 46, 17}, {27, 37, 47}}; + + +let v021_262 = Reshape(v, {1,2,0}, {2,6,2}); +then v021_262 == f32[2x6x2] {{{10, 20}, {30, 40}, + {11, 21}, {31, 41}, + {12, 22}, {32, 42}}, + {{15, 25}, {35, 45}, + {16, 26}, {36, 46}, + {17, 27}, {37, 47}}}; +``` + +As a special case, reshape can transform a single-element array to a scalar and +vice versa. For example, + +``` +Reshape(f32[1x1] {{5}}, {0,1}, {}) == 5; +Reshape(5, {}, {1,1}) == f32[1x1] {{5}}; +``` + +## Rev (reverse) + +See also +[`XlaBuilder::Rev`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + +`Rev(operand, dimensions)` + +Arguments | Type | Semantics +------------ | ------------------- | --------------------- +`operand` | `XlaOp` | array of type T +`dimensions` | `ArraySlice` | dimensions to reverse + +Reverses the order of elements in the `operand` array along the specified +`dimensions`, generating an output array of the same shape. Each element of the +operand array at a multidimensional index is stored into the output array at a +transformed index. The multidimensional index is transformed by reversing the +index in each dimension to be reversed (i.e., if a dimension of size N is one of +the reversing dimensions, its index i is transformed into N - 1 - i). + +One use for the `Rev` operation is to reverse the convolution weight array along +the two window dimensions during the gradient computation in neural networks. + +## RngNormal + +See also +[`XlaBuilder::RngNormal`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + +Constructs an output of a given shape with random numbers generated following +the $$N(\mu, \sigma)$$ normal distribution. The parameters $$\mu$$ and +$$\sigma$$, and output shape have to have a floating point elemental type. The +parameters furthermore have to be scalar valued. + +`RngNormal(mu, sigma, shape)` + +| Arguments | Type | Semantics | +| --------- | ------- | --------------------------------------------------- | +| `mu` | `XlaOp` | Scalar of type T specifying mean of generated | +: : : numbers : +| `sigma` | `XlaOp` | Scalar of type T specifying standard deviation of | +: : : generated numbers : +| `shape` | `Shape` | Output shape of type T | + +## RngUniform + +See also +[`XlaBuilder::RngUniform`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + +Constructs an output of a given shape with random numbers generated following +the uniform distribution over the interval $$[a,b)$$. The parameters and output +element type have to be a boolean type, an integral type or a floating point +types, and the types have to be consistent. The CPU and GPU backends currently +only support F64, F32, F16, BF16, S64, U64, S32 and U32. Furthermore, the +parameters need to be scalar valued. If $$b <= a$$ the result is +implementation-defined. + +`RngUniform(a, b, shape)` + +| Arguments | Type | Semantics | +| --------- | ----------------------- | --------------------------------- | +| `a` | `XlaOp` | Scalar of type T specifying lower | +: : : limit of interval : +| `b` | `XlaOp` | Scalar of type T specifying upper | +: : : limit of interval : +| `shape` | `Shape` | Output shape of type T | + +## Scatter + +The XLA scatter operation generates a result which is the value of the input +array `operand`, with several slices (at indices specified by `scatter_indices`) +updated with the values in `updates` using `update_computation`. + +See also +[`XlaBuilder::Scatter`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + + `scatter(operand, scatter_indices, updates, update_computation, index_vector_dim, update_window_dims, inserted_window_dims, scatter_dims_to_operand_dims)` + +Arguments | Type | Semantics +------------------------------ | ------------------- | --------- +`operand` | `XlaOp` | Array to be scattered into. +`scatter_indices` | `XlaOp` | Array containing the starting indices of the slices that must be scattered to. +`updates` | `XlaOp` | Array containing the values that must be used for scattering. +`update_computation` | `XlaComputation` | Computation to be used for combining the existing values in the input array and the updates during scatter. This computation should be of type `T, T -> T`. +`index_vector_dim` | `int64` | The dimension in `scatter_indices` that contains the starting indices. +`update_window_dims` | `ArraySlice` | The set of dimensions in `updates` shape that are _window dimensions_. +`inserted_window_dims` | `ArraySlice` | The set of _window dimensions_ that must be inserted into `updates` shape. +`scatter_dims_to_operand_dims` | `ArraySlice` | A dimensions map from the scatter indices to the operand index space. This array is interpreted as mapping `i` to `scatter_dims_to_operand_dims[i]` . It has to be one-to-one and total. + +If `index_vector_dim` is equal to `scatter_indices.rank` we implicitly consider +`scatter_indices` to have a trailing `1` dimension. + +We define `update_scatter_dims` of type `ArraySlice` as the set of +dimensions in `updates` shape that are not in `update_window_dims`, in ascending +order. + +The arguments of scatter should follow these constraints: + +- `updates` array must be of rank `update_window_dims.size + + scatter_indices.rank - 1`. + +- Bounds of dimension `i` in `updates` must conform to the following: + + - If `i` is present in `update_window_dims` (i.e. equal to + `update_window_dims`[`k`] for some `k`), then the bound of dimension `i` + in `updates` must not exceed the corresponding bound of `operand` after + accounting for the `inserted_window_dims` (i.e. + `adjusted_window_bounds`[`k`], where `adjusted_window_bounds` contains + the bounds of `operand` with the bounds at indices + `inserted_window_dims` removed). + - If `i` is present in `update_scatter_dims` (i.e. equal to + `update_scatter_dims`[`k`] for some `k`), then the bound of dimension + `i` in `updates` must be equal to the corresponding bound of + `scatter_indices`, skipping `index_vector_dim` (i.e. + `scatter_indices.shape.dims`[`k`], if `k` < `index_vector_dim` and + `scatter_indices.shape.dims`[`k+1`] otherwise). + +- `update_window_dims` must be in ascending order, not have any repeating + dimension numbers, and be in the range `[0, updates.rank)`. + +- `inserted_window_dims` must be in ascending order, not have any repeating + dimension numbers, and be in the range `[0, operand.rank)`. + +- `scatter_dims_to_operand_dims.size` must be equal to + `scatter_indices`[`index_vector_dim`], and its values must be in the range + `[0, operand.rank)`. + +For a given index `U` in the `updates` array, the corresponding index `I` in the +`operand` array into which this update has to be applied is computed as follows: + +1. Let `G` = { `U`[`k`] for `k` in `update_scatter_dims` }. Use `G` to look up + an index vector `S` in the `scatter_indices` array such that `S`[`i`] = + `scatter_indices`[Combine(`G`, `i`)] where Combine(A, b) inserts b at + positions `index_vector_dim` into A. +2. Create an index `S``in` into `operand` using `S` by scattering + `S` using the `scatter_dims_to_operand_dims` map. More formally: + 1. `S``in`[`scatter_dims_to_operand_dims`[`k`]] = `S`[`k`] if + `k` < `scatter_dims_to_operand_dims.size`. + 2. `S``in`[`_`] = `0` otherwise. +3. Create an index `W``in` into `operand` by scattering the indices + at `update_window_dims` in `U` according to `inserted_window_dims`. More + formally: + 1. `W``in`[`window_dims_to_operand_dims`(`k`)] = `U`[`k`] if `k` + < `update_window_dims.size`, where `window_dims_to_operand_dims` is the + monotonic function with domain [`0`, `update_window_dims.size`) and + range [`0`, `operand.rank`) \\ `inserted_window_dims`. (For example, if + `update_window_dims.size` is `4`, `operand.rank` is `6`, and + `inserted_window_dims` is {`0`, `2`} then `window_dims_to_operand_dims` + is {`0`→`1`, `1`→`3`, `2`→`4`, `3`→`5`}). + 2. `W``in`[`_`] = `0` otherwise. +4. `I` is `W``in` + `S``in` where + is element-wise + addition. + +In summary, the scatter operation can be defined as follows. + +- Initialize `output` with `operand`, i.e. for all indices `O` in the + `operand` array: \ + `output`[`O`] = `operand`[`O`] +- For every index `U` in the `updates` array and the corresponding index `O` + in the `operand` array: \ + `output`[`O`] = `update_computation`(`output`[`O`], `updates`[`U`]) + +The order in which updates are applied is non-deterministic. So, when multiple +indices in `updates` refer to the same index in `operand`, the corresponding +value in `output` will be non-deterministic. + +Note that the first parameter that is passed into the `update_computation` will +always be the current value from the `output` array and the second parameter +will always be the value from the `updates` array. This is important +specifically for cases when the `update_computation` is _not commutative_. + +Informally, the scatter op can be viewed as an _inverse_ of the gather op, i.e. +the scatter op updates the elements in the input that are extracted by the +corresponding gather op. + +For a detailed informal description and examples, refer to the +"Informal Description" section under `Gather`. + +## Select + +See also +[`XlaBuilder::Select`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + +Constructs an output array from elements of two input arrays, based on the +values of a predicate array. + + `Select(pred, on_true, on_false)` + +Arguments | Type | Semantics +---------- | ------- | ------------------ +`pred` | `XlaOp` | array of type PRED +`on_true` | `XlaOp` | array of type T +`on_false` | `XlaOp` | array of type T + +The arrays `on_true` and `on_false` must have the same shape. This is also the +shape of the output array. The array `pred` must have the same dimensionality as +`on_true` and `on_false`, with the `PRED` element type. + +For each element `P` of `pred`, the corresponding element of the output array is +taken from `on_true` if the value of `P` is `true`, and from `on_false` if the +value of `P` is `false`. As a restricted form of [broadcasting](broadcasting.md), +`pred` can be a scalar of type `PRED`. In this case, the output array is taken +wholly from `on_true` if `pred` is `true`, and from `on_false` if `pred` is `false`. + +Example with non-scalar `pred`: + +``` +let pred: PRED[4] = {true, false, false, true}; +let v1: s32[4] = {1, 2, 3, 4}; +let v2: s32[4] = {100, 200, 300, 400}; +==> +Select(pred, v1, v2) = s32[4]{1, 200, 300, 4}; +``` + +Example with scalar `pred`: + +``` +let pred: PRED = true; +let v1: s32[4] = {1, 2, 3, 4}; +let v2: s32[4] = {100, 200, 300, 400}; +==> +Select(pred, v1, v2) = s32[4]{1, 2, 3, 4}; +``` + +Selections between tuples are supported. Tuples are considered to be scalar +types for this purpose. If `on_true` and `on_false` are tuples (which must have +the same shape!) then `pred` has to be a scalar of type `PRED`. + +## SelectAndScatter + +See also +[`XlaBuilder::SelectAndScatter`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + +This operation can be considered as a composite operation that first computes +`ReduceWindow` on the `operand` array to select an element from each window, and +then scatters the `source` array to the indices of the selected elements to +construct an output array with the same shape as the operand array. The binary +`select` function is used to select an element from each window by applying it +across each window, and it is called with the property that the first +parameter's index vector is lexicographically less than the second parameter's +index vector. The `select` function returns `true` if the first parameter is +selected and returns `false` if the second parameter is selected, and the +function must hold transitivity (i.e., if `select(a, b)` and `select(b, c)` are +`true`, then `select(a, c)` is also `true`) so that the selected element does +not depend on the order of the elements traversed for a given window. + +The function `scatter` is applied at each selected index in the output array. It +takes two scalar parameters: + +1. Current value at the selected index in the output array +2. The scatter value from `source` that applies to the selected index + +It combines the two parameters and returns a scalar value that's used to update +the value at the selected index in the output array. Initially, all indices of +the output array are set to `init_value`. + +The output array has the same shape as the `operand` array and the `source` +array must have the same shape as the result of applying a `ReduceWindow` +operation on the `operand` array. `SelectAndScatter` can be used to +backpropagate the gradient values for a pooling layer in a neural network. + +`SelectAndScatter(operand, select, window_dimensions, window_strides, +padding, source, init_value, scatter)` + +| Arguments | Type | Semantics | +| ------------------- | ------------------- | -------------------------------- | +| `operand` | `XlaOp` | array of type T over which the | +: : : windows slide : +| `select` | `XlaComputation` | binary computation of type `T, T | +: : : -> PRED`, to apply to all : +: : : elements in each window; returns : +: : : `true` if the first parameter is : +: : : selected and returns `false` if : +: : : the second parameter is selected : +| `window_dimensions` | `ArraySlice` | array of integers for window | +: : : dimension values : +| `window_strides` | `ArraySlice` | array of integers for window | +: : : stride values : +| `padding` | `Padding` | padding type for window | +: : : (Padding\:\:kSame or : +: : : Padding\:\:kValid) : +| `source` | `XlaOp` | array of type T with the values | +: : : to scatter : +| `init_value` | `XlaOp` | scalar value of type T for the | +: : : initial value of the output : +: : : array : +| `scatter` | `XlaComputation` | binary computation of type `T, T | +: : : -> T`, to apply each scatter : +: : : source element with its : +: : : destination element : + +The figure below shows examples of using `SelectAndScatter`, with the `select` +function computing the maximal value among its parameters. Note that when the +windows overlap, as in the figure (2) below, an index of the `operand` array may +be selected multiple times by different windows. In the figure, the element of +value 9 is selected by both of the top windows (blue and red) and the binary +addition `scatter` function produces the output element of value 8 (2 + 6). + +
+ +
+ +The evaluation order of the `scatter` function is arbitrary and may be +non-deterministic. Therefore, the `scatter` function should not be overly +sensitive to reassociation. See the discussion about associativity in the +context of [`Reduce`](#reduce) for more details. + +## Send + +See also +[`XlaBuilder::Send`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + + `Send(operand, channel_handle)` + +Arguments | Type | Semantics +---------------- | --------------- | ----------------------------------------- +`operand` | `XlaOp` | data to send (array of type T) +`channel_handle` | `ChannelHandle` | unique identifier for each send/recv pair + +Sends the given operand data to a `Recv` instruction in another computation +that shares the same channel handle. Does not return any data. + +Similar to the `Recv` operation, the client API of `Send` operation represents +synchronous communication, and is internally decomposed into 2 HLO instructions +(`Send` and `SendDone`) to enable asynchronous data transfers. See also +[`HloInstruction::CreateSend` and `HloInstruction::CreateSendDone`](https://www.tensorflow.org/code/tensorflow/compiler/xla/service/hlo_instruction.h). + +`Send(HloInstruction operand, int64 channel_id)` + +Initiates an asynchronous transfer of the operand to the resources allocated by +the `Recv` instruction with the same channel id. Returns a context, which is +used by a following `SendDone` instruction to wait for the completion of the +data transfer. The context is a tuple of {operand (shape), request identifier +(U32)} and it can only be used by a `SendDone` instruction. + + `SendDone(HloInstruction context)` + +Given a context created by a `Send` instruction, waits for the data transfer to +complete. The instruction does not return any data. + + Scheduling of channel instructions + +The execution order of the 4 instructions for each channel (`Recv`, `RecvDone`, +`Send`, `SendDone`) is as below. + +
+ +
+ +* `Recv` happens before `Send` +* `Send` happens before `RecvDone` +* `Recv` happens before `RecvDone` +* `Send` happens before `SendDone` + +When the backend compilers generate a linear schedule for each computation that +communicates via channel instructions, there must not be cycles across the +computations. For example, below schedules lead to deadlocks. + +
+ +
+ +## Slice + +See also +[`XlaBuilder::Slice`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + +Slicing extracts a sub-array from the input array. The sub-array is of the same +rank as the input and contains the values inside a bounding box within the input +array where the dimensions and indices of the bounding box are given as +arguments to the slice operation. + + `Slice(operand, start_indices, limit_indices)` + +| Arguments | Type | Semantics | +| --------------- | ------------------- | ------------------------------------ | +| `operand` | `XlaOp` | N dimensional array of type T | +| `start_indices` | `ArraySlice` | List of N integers containing the | +: : : starting indices of the slice for : +: : : each dimension. Values must be : +: : : greater than or equal to zero. : +| `limit_indices` | `ArraySlice` | List of N integers containing the | +: : : ending indices (exclusive) for the : +: : : slice for each dimension. Each value : +: : : must be greater than or equal to the : +: : : respective `start_indices` value for : +: : : the dimension and less than or equal : +: : : to the size of the dimension. : + +1-dimensional example: + +``` +let a = {0.0, 1.0, 2.0, 3.0, 4.0} +Slice(a, {2}, {4}) produces: + {2.0, 3.0} +``` + +2-dimensional example: + +``` +let b = + { {0.0, 1.0, 2.0}, + {3.0, 4.0, 5.0}, + {6.0, 7.0, 8.0}, + {9.0, 10.0, 11.0} } + +Slice(b, {2, 1}, {4, 3}) produces: + { { 7.0, 8.0}, + {10.0, 11.0} } +``` + +## Sort + +See also +[`XlaBuilder::Sort`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + +There are two versions of the Sort instruction: a single-operand and a +multi-operand version. + +`Sort(operand, dimension)` + +Arguments | Type | Semantics +----------- | ------- | -------------------- +`operand` | `XlaOp` | The operand to sort. +`dimension` | `int64` | The dimension along which to sort. + +Sorts the elements in the operand in ascending order along the provided +dimension. For example, for a rank-2 (matrix) operand, a `dimension` value of 0 +will sort each column independently, and a `dimension` value of 1 will sort each +row independently. If the operand's elements have floating point type, and the +operand contains NaN elements, the order of elements in the output is +implementation-defined. + +`Sort(keys, values, ... values, dimension)` + +Sorts both the key and one or more value operands. The keys are sorted as in the +single-operand version. Each of the values inputs is sorted according to the +order of the corresponding keys. For example, if the three inputs are `keys = +[3, 1]`, `values0 = [42, 50]`, `values1 = [-3.0, 1.1]`, then the output of the +sort is the tuple `{[1, 3], [50, 42], [1.1, -3.0]}`. + +The sort is not guaranteed to be stable, that is, if the keys array contains +duplicates, the order of values corresponding to these keys may not be +preserved. + +Arguments | Type | Semantics +----------- | ---------------------- | ---------------------------------- +`keys` | `XlaOp` | The sort keys. +`values` | Sequence of N `XlaOp`s | The values to sort. +`dimension` | `int64` | The dimension along which to sort. + +The `keys` and each of the `values` inputs must have the same dimensions, but +may have different element types. + +## Transpose + +See also the `tf.reshape` operation. + +`Transpose(operand)` + +Arguments | Type | Semantics +------------- | ------------------- | ------------------------------ +`operand` | `XlaOp` | The operand to transpose. +`permutation` | `ArraySlice` | How to permute the dimensions. + + +Permutes the operand dimensions with the given permutation, so +`∀ i . 0 ≤ i < rank ⇒ input_dimensions[permutation[i]] = output_dimensions[i]`. + +This is the same as Reshape(operand, permutation, + Permute(permutation, operand.shape.dimensions)). + +## Tuple + +See also +[`XlaBuilder::Tuple`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + +A tuple containing a variable number of data handles, each of which has its own +shape. + +This is analogous to `std::tuple` in C++. Conceptually: + +``` +let v: f32[10] = f32[10]{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}; +let s: s32 = 5; +let t: (f32[10], s32) = tuple(v, s); +``` + +Tuples can be deconstructed (accessed) via the [`GetTupleElement`] +(#gettupleelement) operation. + +## While + +See also +[`XlaBuilder::While`](https://www.tensorflow.org/code/tensorflow/compiler/xla/client/xla_builder.h). + + `While(condition, body, init)` + +| Arguments | Type | Semantics | +| ----------- | ---------------- | ---------------------------------------- | +| `condition` | `XlaComputation` | XlaComputation of type `T -> PRED` which | +: : : defines the termination condition of the : +: : : loop. : +| `body` | `XlaComputation` | XlaComputation of type `T -> T` which | +: : : defines the body of the loop. : +| `init` | `T` | Initial value for the parameter of | +: : : `condition` and `body`. : + +Sequentially executes the `body` until the `condition` fails. This is similar to +a typical while loop in many other languages except for the differences and +restrictions listed below. + +* A `While` node returns a value of type `T`, which is the result from the + last execution of the `body`. +* The shape of the type `T` is statically determined and must be the same + across all iterations. + +The T parameters of the computations are initialized with the `init` value in +the first iteration and are automatically updated to the new result from `body` +in each subsequent iteration. + +One main use case of the `While` node is to implement the repeated execution of +training in neural networks. Simplified pseudocode is shown below with a graph +that represents the computation. The code can be found in +[`while_test.cc`](https://www.tensorflow.org/code/tensorflow/compiler/xla/tests/while_test.cc). +The type `T` in this example is a `Tuple` consisting of an `int32` for the +iteration count and a `vector[10]` for the accumulator. For 1000 iterations, the +loop keeps adding a constant vector to the accumulator. + +``` +// Pseudocode for the computation. +init = {0, zero_vector[10]} // Tuple of int32 and float[10]. +result = init; +while (result(0) < 1000) { + iteration = result(0) + 1; + new_vector = result(1) + constant_vector[10]; + result = {iteration, new_vector}; +} +``` + +
+ +
diff --git a/tensorflow/compiler/xla/g3doc/overview.md b/tensorflow/compiler/xla/g3doc/overview.md new file mode 100644 index 0000000000000000000000000000000000000000..d3428b7276131e8f406f60cfea9a9346c5478433 --- /dev/null +++ b/tensorflow/compiler/xla/g3doc/overview.md @@ -0,0 +1,98 @@ +# XLA Overview + +
+ +
+ +> Note: XLA is still under development. Some use cases will not +> see improvements in speed or decreased memory usage. + +XLA (Accelerated Linear Algebra) is a domain-specific compiler for linear +algebra that optimizes TensorFlow computations. The results are improvements in +speed, memory usage, and portability on server and mobile platforms. Initially, +most users will not see large benefits from XLA, but are welcome to experiment +by using XLA via [just-in-time (JIT) compilation](./jit.md) or +[ahead-of-time (AOT) compilation](./tfcompile.md). Developers targeting new +hardware accelerators are especially encouraged to try out XLA. + +The XLA framework is experimental and in active development. In particular, +while it is unlikely that the semantics of existing operations will change, it +is expected that more operations will be added to cover important use cases. The +team welcomes feedback from the community about missing functionality and +community contributions via GitHub. + +## Why did we build XLA? + +We had several objectives for XLA to work with TensorFlow: + +* *Improve execution speed.* Compile subgraphs to reduce the execution time of + short-lived Ops to eliminate overhead from the TensorFlow runtime, fuse + pipelined operations to reduce memory overhead, and specialize to known + tensor shapes to allow for more aggressive constant propagation. + +* *Improve memory usage.* Analyze and schedule memory usage, in principle + eliminating many intermediate storage buffers. + +* *Reduce reliance on custom Ops.* Remove the need for many custom Ops by + improving the performance of automatically fused low-level Ops to match the + performance of custom Ops that were fused by hand. + +* *Reduce mobile footprint.* Eliminate the TensorFlow runtime by ahead-of-time + compiling the subgraph and emitting an object/header file pair that can be + linked directly into another application. The results can reduce the + footprint for mobile inference by several orders of magnitude. + +* *Improve portability.* Make it relatively easy to write a new backend for + novel hardware, at which point a large fraction of TensorFlow programs will + run unmodified on that hardware. This is in contrast with the approach of + specializing individual monolithic Ops for new hardware, which requires + TensorFlow programs to be rewritten to make use of those Ops. + +## How does XLA work? + +The input language to XLA is called "HLO IR", or just HLO (High Level +Optimizer). The semantics of HLO are described on the +[Operation Semantics](./operation_semantics.md) page. It +is most convenient to think of HLO as a +[compiler IR](https://en.wikipedia.org/wiki/Intermediate_representation). + +XLA takes graphs ("computations") defined in HLO and compiles them into machine +instructions for various architectures. XLA is modular in the sense that it is +easy to slot in an alternative backend to +[target some novel HW architecture](./developing_new_backend.md). +The CPU backend for x64 and ARM64 as well as the NVIDIA GPU backend are in the +TensorFlow source tree. + +The following diagram shows the compilation process in XLA: + +
+ +
+ +XLA comes with several optimizations and analysis passes that are +target-independent, such as +[CSE](https://en.wikipedia.org/wiki/Common_subexpression_elimination), +target-independent operation fusion, and buffer analysis for allocating runtime +memory for the computation. + +After the target-independent step, XLA sends the HLO computation to a backend. +The backend can perform further HLO-level optimizations, this time with target +specific information and needs in mind. For example, the XLA GPU backend may +perform operation fusion beneficial specifically for the GPU programming model +and determine how to partition the computation into streams. At this stage, +backends may also pattern-match certain operations or combinations thereof to +optimized library calls. + +The next step is target-specific code generation. The CPU and GPU backends +included with XLA use [LLVM](http://llvm.org) for low-level IR, optimization, +and code-generation. These backends emit the LLVM IR necessary to represent the +XLA HLO computation in an efficient manner, and then invoke LLVM to emit native +code from this LLVM IR. + +The GPU backend currently supports NVIDIA GPUs via the LLVM NVPTX backend; the +CPU backend supports multiple CPU ISAs. + +## Supported Platforms + +XLA currently supports [JIT compilation](./jit.md) on x86-64 and NVIDIA GPUs; and +[AOT compilation](./tfcompile.md) for x86-64 and ARM. diff --git a/tensorflow/compiler/xla/g3doc/shapes.md b/tensorflow/compiler/xla/g3doc/shapes.md new file mode 100644 index 0000000000000000000000000000000000000000..39e74ff307cde49ef378a1201cb074dce4ababf0 --- /dev/null +++ b/tensorflow/compiler/xla/g3doc/shapes.md @@ -0,0 +1,150 @@ +# Shapes and Layout + +The XLA `Shape` proto +([xla_data.proto](https://www.tensorflow.org/code/tensorflow/compiler/xla/xla_data.proto)) +describes the rank, size, and data type of an N-dimensional array (*array* in +short). + +## Terminology, Notation, and Conventions + +* The rank of an array is equal to the number of dimensions. The *true rank* + of an array is the number of dimensions which have a size greater than 1. + +* Dimensions are numbered from `0` up to `N-1` for an `N` dimensional array. + The dimension numbers are arbitrary labels for convenience. The order of + these dimension numbers does not imply a particular minor/major ordering in + the layout of the shape. The layout is determined by the `Layout` proto. + +* By convention, dimensions are listed in increasing order of dimension + number. For example, for a 3-dimensional array of size `[A x B x C]`, + dimension 0 has size `A`, dimension 1 has size `B` and dimension 2 has size + `C`. + + Some utilities in XLA also support negative indexing, similarly to Python; + dimension -1 is the last dimension (equivalent to `N-1` for an `N` + dimensional array). For example, for the 3-dimensional array described + above, dimension -1 has size `C`, dimension -2 has size `B` and so on. + +* Two, three, and four dimensional arrays often have specific letters + associated with dimensions. For example, for a 2D array: + + * dimension 0: `y` + * dimension 1: `x` + + For a 3D array: + + * dimension 0: `z` + * dimension 1: `y` + * dimension 2: `x` + + For a 4D array: + + * dimension 0: `p` + * dimension 1: `z` + * dimension 2: `y` + * dimension 3: `x` + +* Functions in the XLA API which take dimensions do so in increasing order of + dimension number. This matches the ordering used when passing dimensions as + an `initializer_list`; e.g. + + `ShapeUtil::MakeShape(F32, {A, B, C, D})` + + Will create a shape whose dimension size array consists of the sequence + `[A, B, C, D]`. + +## Layout + +The `Layout` proto describes how an array is represented in memory. The `Layout` +proto includes the following fields: + +``` +message Layout { + repeated int64 minor_to_major = 1; + repeated int64 padded_dimensions = 2; + optional PaddingValue padding_value = 3; +} +``` + +### Minor-to-major dimension ordering + +The only required field is `minor_to_major`. This field describes the +minor-to-major ordering of the dimensions within a shape. Values in +`minor_to_major` are an ordering of the dimensions of the array (`0` to `N-1` +for an `N` dimensional array) with the first value being the most-minor +dimension up to the last value which is the most-major dimension. The most-minor +dimension is the dimension which changes most rapidly when stepping through the +elements of the array laid out in linear memory. + +For example, consider the following 2D array of size `[2 x 3]`: + +``` +a b c +d e f +``` + +Here dimension `0` is size 2, and dimension `1` is size 3. If the +`minor_to_major` field in the layout is `[0, 1]` then dimension `0` is the +most-minor dimension and dimension `1` is the most-major dimension. This +corresponds to the following layout in linear memory: + +``` +a d b e c f +``` + +This minor-to-major dimension order of `0` up to `N-1` is akin to *column-major* +(at rank 2). Assuming a monotonic ordering of dimensions, another name we may +use to refer to this layout in the code is simply "dim 0 is minor". + +On the other hand, if the `minor_to_major` field in the layout is `[1, 0]` then +the layout in linear memory is: + +``` +a b c d e f +``` + +A minor-to-major dimension order of `N-1` down to `0` for an `N` dimensional +array is akin to *row-major* (at rank 2). Assuming a monotonic ordering of +dimensions, another name we may use to refer to this layout in the code is +simply "dim 0 is major". + +#### Default minor-to-major ordering + +The default layout for newly created Shapes is "dimension order is +major-to-minor" (akin to row-major at rank 2). + +### Padding + +Padding is defined in the optional `padded_dimensions` and `padding_value` +fields. The field `padded_dimensions` describes the sizes (widths) to which each +dimension is padded. If present, the number of elements in `padded_dimensions` +must equal the rank of the shape. + +For example, given the `[2 x 3]` array defined above, if `padded_dimension` is +`[3, 5]` then dimension 0 is padded to a width of 3 and dimension 1 is padded to +a width of 5. The layout in linear memory (assuming a padding value of 0 and +column-major layout) is: + +``` +a d 0 b e 0 c f 0 0 0 0 0 0 0 +``` + +This is equivalent to the layout of the following array with the same +minor-to-major dimension order: + +``` +a b c 0 0 +d e f 0 0 +0 0 0 0 0 +``` + +### Indexing into arrays + +The class `IndexUtil` in +[index_util.h](https://www.tensorflow.org/code/tensorflow/compiler/xla/index_util.h) +provides utilities for converting between multidimensional indices and linear +indices given a shape and layout. Multidimensional indices include a `int64` +index for each dimension. Linear indices are a single `int64` value which +indexes into the buffer holding the array. See `shape_util.h` and +`layout_util.h` in the same directory for utilities that simplify creation and +manipulation of shapes and layouts. diff --git a/tensorflow/compiler/xla/g3doc/tfcompile.md b/tensorflow/compiler/xla/g3doc/tfcompile.md new file mode 100644 index 0000000000000000000000000000000000000000..5ee09fd302ba0edf84a7c99bb369586067141bef --- /dev/null +++ b/tensorflow/compiler/xla/g3doc/tfcompile.md @@ -0,0 +1,281 @@ +# Using AOT compilation + +## What is tfcompile? + +`tfcompile` is a standalone tool that ahead-of-time (AOT) compiles TensorFlow +graphs into executable code. It can reduce total binary size, and also avoid +some runtime overheads. A typical use-case of `tfcompile` is to compile an +inference graph into executable code for mobile devices. + +The TensorFlow graph is normally executed by the TensorFlow runtime. This incurs +some runtime overhead for execution of each node in the graph. This also leads +to a larger total binary size, since the code for the TensorFlow runtime needs +to be available, in addition to the graph itself. The executable code produced +by `tfcompile` does not use the TensorFlow runtime, and only has dependencies on +kernels that are actually used in the computation. + +The compiler is built on top of the XLA framework. The code bridging TensorFlow +to the XLA framework resides under +[tensorflow/compiler](https://www.tensorflow.org/code/tensorflow/compiler/), +which also includes support for [just-in-time (JIT) compilation](jit.md) of +TensorFlow graphs. + +## What does tfcompile do? + +`tfcompile` takes a subgraph, identified by the TensorFlow concepts of +feeds and fetches, and generates a function that implements that subgraph. +The `feeds` are the input arguments for the function, and the `fetches` are the +output arguments for the function. All inputs must be fully specified by the +feeds; the resulting pruned subgraph cannot contain Placeholder or Variable +nodes. It is common to specify all Placeholders and Variables as feeds, which +ensures the resulting subgraph no longer contains these nodes. The generated +function is packaged as a `cc_library`, with a header file exporting the +function signature, and an object file containing the implementation. The user +writes code to invoke the generated function as appropriate. + +## Using tfcompile + +This section details high level steps for generating an executable binary with +`tfcompile` from a TensorFlow subgraph. The steps are: + +* Step 1: Configure the subgraph to compile +* Step 2: Use the `tf_library` build macro to compile the subgraph +* Step 3: Write code to invoke the subgraph +* Step 4: Create the final binary + +### Step 1: Configure the subgraph to compile + +Identify the feeds and fetches that correspond to the input and output +arguments for the generated function. Then configure the `feeds` and `fetches` +in a [`tensorflow.tf2xla.Config`](https://www.tensorflow.org/code/tensorflow/compiler/tf2xla/tf2xla.proto) +proto. + +```textproto +# Each feed is a positional input argument for the generated function. The order +# of each entry matches the order of each input argument. Here “x_hold” and “y_hold” +# refer to the names of placeholder nodes defined in the graph. +feed { + id { node_name: "x_hold" } + shape { + dim { size: 2 } + dim { size: 3 } + } +} +feed { + id { node_name: "y_hold" } + shape { + dim { size: 3 } + dim { size: 2 } + } +} + +# Each fetch is a positional output argument for the generated function. The order +# of each entry matches the order of each output argument. Here “x_y_prod” +# refers to the name of a matmul node defined in the graph. +fetch { + id { node_name: "x_y_prod" } +} +``` + +### Step 2: Use tf_library build macro to compile the subgraph + +This step converts the graph into a `cc_library` using the `tf_library` build +macro. The `cc_library` consists of an object file containing the code generated +from the graph, along with a header file that gives access to the generated +code. `tf_library` utilizes `tfcompile` to compile the TensorFlow graph into +executable code. + +```build +load("//tensorflow/compiler/aot:tfcompile.bzl", "tf_library") + +# Use the tf_library macro to compile your graph into executable code. +tf_library( + # name is used to generate the following underlying build rules: + # : cc_library packaging the generated header and object files + # _test : cc_test containing a simple test and benchmark + # _benchmark : cc_binary containing a stand-alone benchmark with minimal deps; + # can be run on a mobile device + name = "test_graph_tfmatmul", + # cpp_class specifies the name of the generated C++ class, with namespaces allowed. + # The class will be generated in the given namespace(s), or if no namespaces are + # given, within the global namespace. + cpp_class = "foo::bar::MatMulComp", + # graph is the input GraphDef proto, by default expected in binary format. To + # use the text format instead, just use the ‘.pbtxt’ suffix. A subgraph will be + # created from this input graph, with feeds as inputs and fetches as outputs. + # No Placeholder or Variable ops may exist in this subgraph. + graph = "test_graph_tfmatmul.pb", + # config is the input Config proto, by default expected in binary format. To + # use the text format instead, use the ‘.pbtxt’ suffix. This is where the + # feeds and fetches were specified above, in the previous step. + config = "test_graph_tfmatmul.config.pbtxt", +) +``` + +> To generate the GraphDef proto (test_graph_tfmatmul.pb) for this example, run +> [make_test_graphs.py](https://www.tensorflow.org/code/tensorflow/compiler/aot/tests/make_test_graphs.py) +> and specify the output location with the --out_dir flag. + +Typical graphs contain [`Variables`](https://www.tensorflow.org/guide/variables) +representing the weights that are learned via training, but `tfcompile` cannot +compile a subgraph that contain `Variables`. The +[freeze_graph.py](https://www.tensorflow.org/code/tensorflow/python/tools/freeze_graph.py) +tool converts variables into constants, using values stored in a checkpoint +file. As a convenience, the `tf_library` macro supports the `freeze_checkpoint` +argument, which runs the tool. For more examples see +[tensorflow/compiler/aot/tests/BUILD](https://www.tensorflow.org/code/tensorflow/compiler/aot/tests/BUILD). + +> Constants that show up in the compiled subgraph are compiled directly into the +> generated code. To pass the constants into the generated function, rather than +> having them compiled-in, simply pass them in as feeds. + +For details on the `tf_library` build macro, see +[tfcompile.bzl](https://www.tensorflow.org/code/tensorflow/compiler/aot/tfcompile.bzl). + +For details on the underlying `tfcompile` tool, see +[tfcompile_main.cc](https://www.tensorflow.org/code/tensorflow/compiler/aot/tfcompile_main.cc). + +### Step 3: Write code to invoke the subgraph + +This step uses the header file (`test_graph_tfmatmul.h`) generated by the +`tf_library` build macro in the previous step to invoke the generated code. The +header file is located in the `bazel-genfiles` directory corresponding to the +build package, and is named based on the name attribute set on the `tf_library` +build macro. For example, the header generated for `test_graph_tfmatmul` would +be `test_graph_tfmatmul.h`. Below is an abbreviated version of what is +generated. The generated file, in `bazel-genfiles`, contains additional useful +comments. + +```c++ +namespace foo { +namespace bar { + +// MatMulComp represents a computation previously specified in a +// TensorFlow graph, now compiled into executable code. +class MatMulComp { + public: + // AllocMode controls the buffer allocation mode. + enum class AllocMode { + ARGS_RESULTS_AND_TEMPS, // Allocate arg, result and temp buffers + RESULTS_AND_TEMPS_ONLY, // Only allocate result and temp buffers + }; + + MatMulComp(AllocMode mode = AllocMode::ARGS_RESULTS_AND_TEMPS); + ~MatMulComp(); + + // Runs the computation, with inputs read from arg buffers, and outputs + // written to result buffers. Returns true on success and false on failure. + bool Run(); + + // Arg methods for managing input buffers. Buffers are in row-major order. + // There is a set of methods for each positional argument. + void** args(); + + void set_arg0_data(float* data); + float* arg0_data(); + float& arg0(size_t dim0, size_t dim1); + + void set_arg1_data(float* data); + float* arg1_data(); + float& arg1(size_t dim0, size_t dim1); + + // Result methods for managing output buffers. Buffers are in row-major order. + // Must only be called after a successful Run call. There is a set of methods + // for each positional result. + void** results(); + + + float* result0_data(); + float& result0(size_t dim0, size_t dim1); +}; + +} // end namespace bar +} // end namespace foo +``` + +The generated C++ class is called `MatMulComp` in the `foo::bar` namespace, +because that was the `cpp_class` specified in the `tf_library` macro. All +generated classes have a similar API, with the only difference being the methods +to handle arg and result buffers. Those methods differ based on the number and +types of the buffers, which were specified by the `feed` and `fetch` arguments +to the `tf_library` macro. + +There are three types of buffers managed within the generated class: `args` +representing the inputs, `results` representing the outputs, and `temps` +representing temporary buffers used internally to perform the computation. By +default, each instance of the generated class allocates and manages all of these +buffers for you. The `AllocMode` constructor argument may be used to change this +behavior. All buffers are aligned to 64-byte boundaries. + +The generated C++ class is just a wrapper around the low-level code generated by +XLA. + +Example of invoking the generated function based on +[`tfcompile_test.cc`](https://www.tensorflow.org/code/tensorflow/compiler/aot/tests/tfcompile_test.cc): + +```c++ +#define EIGEN_USE_THREADS +#define EIGEN_USE_CUSTOM_THREAD_POOL + +#include +#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" +#include "tensorflow/compiler/aot/tests/test_graph_tfmatmul.h" // generated + +int main(int argc, char** argv) { + Eigen::ThreadPool tp(2); // Size the thread pool as appropriate. + Eigen::ThreadPoolDevice device(&tp, tp.NumThreads()); + + + foo::bar::MatMulComp matmul; + matmul.set_thread_pool(&device); + + // Set up args and run the computation. + const float args[12] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}; + std::copy(args + 0, args + 6, matmul.arg0_data()); + std::copy(args + 6, args + 12, matmul.arg1_data()); + matmul.Run(); + + // Check result + if (matmul.result0(0, 0) == 58) { + std::cout << "Success" << std::endl; + } else { + std::cout << "Failed. Expected value 58 at 0,0. Got:" + << matmul.result0(0, 0) << std::endl; + } + + return 0; +} +``` + +### Step 4: Create the final binary + +This step combines the library generated by `tf_library` in step 2 and the code +written in step 3 to create a final binary. Below is an example `bazel` BUILD +file. + +```build +# Example of linking your binary +# Also see //tensorflow/compiler/aot/tests/BUILD +load("//tensorflow/compiler/aot:tfcompile.bzl", "tf_library") + +# The same tf_library call from step 2 above. +tf_library( + name = "test_graph_tfmatmul", + ... +) + +# The executable code generated by tf_library can then be linked into your code. +cc_binary( + name = "my_binary", + srcs = [ + "my_code.cc", # include test_graph_tfmatmul.h to access the generated header + ], + deps = [ + ":test_graph_tfmatmul", # link in the generated object file + "//third_party/eigen3", + ], + linkopts = [ + "-lpthread", + ] +) +``` diff --git a/tensorflow/compiler/xla/g3doc/tutorials/xla_compile.ipynb b/tensorflow/compiler/xla/g3doc/tutorials/xla_compile.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..2a83092805be5efdd7b9ab54449b2bcc6a2ec481 --- /dev/null +++ b/tensorflow/compiler/xla/g3doc/tutorials/xla_compile.ipynb @@ -0,0 +1,373 @@ +{ + "nbformat": 4, + "nbformat_minor": 0, + "metadata": { + "colab": { + "name": "The XLA compile API", + "version": "0.3.2", + "provenance": [], + "collapsed_sections": [], + "toc_visible": true + }, + "kernelspec": { + "name": "python3", + "display_name": "Python 3" + } + }, + "cells": [ + { + "metadata": { + "colab_type": "text", + "id": "f4TSNCvpENrW" + }, + "cell_type": "markdown", + "source": [ + "##### Copyright 2018 The TensorFlow Authors." + ] + }, + { + "metadata": { + "cellView": "form", + "colab_type": "code", + "id": "vamNSA0vEP-m", + "colab": {} + }, + "cell_type": "code", + "source": [ + "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n", + "# you may not use this file except in compliance with the License.\n", + "# You may obtain a copy of the License at\n", + "#\n", + "# https://www.apache.org/licenses/LICENSE-2.0\n", + "#\n", + "# Unless required by applicable law or agreed to in writing, software\n", + "# distributed under the License is distributed on an \"AS IS\" BASIS,\n", + "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", + "# See the License for the specific language governing permissions and\n", + "# limitations under the License." + ], + "execution_count": 0, + "outputs": [] + }, + { + "metadata": { + "colab_type": "text", + "id": "e1oSi4lHFt3z" + }, + "cell_type": "markdown", + "source": [ + "# The XLA compile API" + ] + }, + { + "metadata": { + "colab_type": "text", + "id": "b7noD9NjFRL-" + }, + "cell_type": "markdown", + "source": [ + "\n", + " \n", + " \n", + " \n", + "
\n", + " View on TensorFlow.org\n", + " \n", + " Run in Google Colab\n", + " \n", + " View source on GitHub\n", + "
" + ] + }, + { + "metadata": { + "colab_type": "text", + "id": "v9YbsuLZaBXy" + }, + "cell_type": "markdown", + "source": [ + "\n", + "\n", + "Import TensorFlow and the XLA library. XLA contains `xla.compile()`, an experimental API that compiles part or all of a model with [XLA](https://www.tensorflow.org/extend/xla/)." + ] + }, + { + "metadata": { + "colab_type": "code", + "id": "45kUPj5ZFrRa", + "colab": {} + }, + "cell_type": "code", + "source": [ + "import tensorflow as tf\n", + "\n", + "from tensorflow.contrib.compiler import xla" + ], + "execution_count": 0, + "outputs": [] + }, + { + "metadata": { + "colab_type": "text", + "id": "GZVNiRmTDV-5" + }, + "cell_type": "markdown", + "source": [ + "Define some necessary constants and prepare the MNIST dataset." + ] + }, + { + "metadata": { + "colab_type": "code", + "id": "f37TSEGvGX4_", + "colab": {} + }, + "cell_type": "code", + "source": [ + "# Size of each input image, 28 x 28 pixels\n", + "IMAGE_SIZE = 28 * 28\n", + "# Number of distinct number labels, [0..9]\n", + "NUM_CLASSES = 10\n", + "# Number of examples in each training batch (step)\n", + "TRAIN_BATCH_SIZE = 100\n", + "# Number of training steps to run\n", + "TRAIN_STEPS = 1000" + ], + "execution_count": 0, + "outputs": [] + }, + { + "metadata": { + "colab_type": "code", + "id": "TiVXchblG5hK", + "colab": {} + }, + "cell_type": "code", + "source": [ + "# Loads MNIST dataset.\n", + "train, test = tf.keras.datasets.mnist.load_data()\n", + "train_ds = tf.data.Dataset.from_tensor_slices(train).batch(TRAIN_BATCH_SIZE).repeat()\n", + "test_ds = tf.data.Dataset.from_tensor_slices(test).batch(TRAIN_BATCH_SIZE)\n", + "\n", + "iterator = tf.data.Iterator.from_structure(train_ds.output_types, train_ds.output_shapes)\n", + "images, labels = iterator.get_next()\n", + "images = tf.reshape(images, [-1, IMAGE_SIZE])\n", + "images, labels = tf.cast(images, tf.float32), tf.cast(labels, tf.int64)" + ], + "execution_count": 0, + "outputs": [] + }, + { + "metadata": { + "colab_type": "text", + "id": "x_ZehpZP-SfS" + }, + "cell_type": "markdown", + "source": [ + "# Define the model constructing function\n", + "\n", + "Following code block contains a function that constructs a simple model with one dense layer, including both forward and backward propagation.\n", + "\n", + "When called, it returns two values. `y` is a `tf.Tensor` representing predicted probability of each target class, `train_step` is a `tf.Operation` that increments `global_step` and applies variable update." + ] + }, + { + "metadata": { + "colab_type": "code", + "id": "ZbhJl_WvGa3g", + "colab": {} + }, + "cell_type": "code", + "source": [ + "def build_mnist_model(x, y_):\n", + " y = tf.keras.layers.Dense(NUM_CLASSES).apply(x)\n", + "\n", + " cross_entropy = tf.losses.sparse_softmax_cross_entropy(labels=y_, logits=y)\n", + " train_step = tf.train.GradientDescentOptimizer(0.5).minimize(cross_entropy)\n", + "\n", + " return y, train_step" + ], + "execution_count": 0, + "outputs": [] + }, + { + "metadata": { + "colab_type": "text", + "id": "7Jh3lyQHDfM9" + }, + "cell_type": "markdown", + "source": [ + "# Enable XLA\n", + "\n", + "Use `xla.compile` with the `build_mnist_model` function to enable XLA. Following code block wraps the model with `xla.compile()`, which allows the target function with provided inputs to be executed by XLA." + ] + }, + { + "metadata": { + "colab_type": "code", + "id": "kYpCXCdRHNuN", + "colab": {} + }, + "cell_type": "code", + "source": [ + "[y] = xla.compile(build_mnist_model, inputs=[images, labels])" + ], + "execution_count": 0, + "outputs": [] + }, + { + "metadata": { + "colab_type": "text", + "id": "4giQh62IrZGF" + }, + "cell_type": "markdown", + "source": [ + "When compiling the graph, XLA replaces all the graph nodes constructed in the target function with a few XLA ops.\n", + "\n", + "xla.compile does not return any\n", + "`tf.Operation` nodes that can be executed independently from the generated XLA ops. Instead, returned `tf.Operation` nodes from the target function are added as control dependencies of all returned `tf.Tensor` values. This triggers execution of the `tf.Operation` nodes when the returned tensors are evaluated.\n", + "\n", + "In pseudo-code, xla.compile's implementation looks as follows:\n", + "\n", + "---\n", + "```\n", + "# Ask Tensorflow to execute code in XLA-friendly manner\n", + "\n", + "y, train_step = build_mnist_model(images, labels)\n", + "with tf.control_dependencies([train_step]):\n", + " y = tf.identity(y)\n", + "\n", + "# Ask Tensorflow to STOP executing code in XLA-friendly manner\n", + "```\n", + "---\n", + "\n", + "xla.compile() always returns a list of `tf.Tensor`'s (even if there is only one-element)." + ] + }, + { + "metadata": { + "colab_type": "text", + "id": "TPGas4jjFLZl" + }, + "cell_type": "markdown", + "source": [ + "If you were to print the constructed graph now, you will see that it is not much different from a normal Tensorflow graph and you won't be able to find XLA ops mentioned before. This is because the actual compilation happens later when you try to execute the graph with `sess.run()`. At that time, Tensorflow triggers a series of graph rewrite passes that actually generate XLA ops, which compiles and executes computation when all inputs are ready." + ] + }, + { + "metadata": { + "colab_type": "text", + "id": "EZD1m_n1DxAF" + }, + "cell_type": "markdown", + "source": [ + "# Train and test the model" + ] + }, + { + "metadata": { + "colab_type": "code", + "id": "qe28bAHNHUG2", + "colab": {} + }, + "cell_type": "code", + "source": [ + "# Creates session and initialize all variables.\n", + "# xla.compile() doesn't work with Keras model.fit() API or TF eager mode yet.\n", + "sess = tf.Session()\n", + "sess.run(tf.global_variables_initializer())" + ], + "execution_count": 0, + "outputs": [] + }, + { + "metadata": { + "colab_type": "text", + "id": "qgsKmz3n2UiW" + }, + "cell_type": "markdown", + "source": [ + "Following code block trains model. Evaluating `y` also triggers its control dependency node `train_step`, which updates model variables." + ] + }, + { + "metadata": { + "colab_type": "code", + "id": "_GxF6jTRHVuA", + "colab": { + "base_uri": "https://localhost:8080/", + "height": 34 + }, + "outputId": "fbf299ca-02d5-4e95-f9fe-8f3c0432d132" + }, + "cell_type": "code", + "source": [ + "# Feeds training dataset\n", + "sess.run(iterator.make_initializer(train_ds))\n", + "\n", + "# Runs TRAIN_STEPS steps\n", + "for i in range(TRAIN_STEPS):\n", + " sess.run(y)\n", + "\n", + "print(\"Model trained for %s steps.\" % TRAIN_STEPS)" + ], + "execution_count": 21, + "outputs": [ + { + "output_type": "stream", + "text": [ + "Model trained for 1000 steps.\n" + ], + "name": "stdout" + } + ] + }, + { + "metadata": { + "colab_type": "code", + "id": "dHlQlRSRHXD1", + "colab": { + "base_uri": "https://localhost:8080/", + "height": 34 + }, + "outputId": "9c3677a2-ec84-406f-9d2c-d722844f3093" + }, + "cell_type": "code", + "source": [ + "# Tests trained model\n", + "\n", + "# Feeds testing dataset\n", + "sess.run(iterator.make_initializer(test_ds))\n", + "\n", + "# Calculates accuracy\n", + "correct_prediction = tf.equal(tf.argmax(y, 1), labels)\n", + "accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))\n", + "print(\"Prediction accuracy after training: %s\" % sess.run(accuracy))" + ], + "execution_count": 22, + "outputs": [ + { + "output_type": "stream", + "text": [ + "Prediction accuracy after training: 0.91\n" + ], + "name": "stdout" + } + ] + }, + { + "metadata": { + "colab_type": "code", + "id": "ynJQIuzjHYOb", + "colab": {} + }, + "cell_type": "code", + "source": [ + "# Cleans up session\n", + "sess.close()" + ], + "execution_count": 0, + "outputs": [] + } + ] +} \ No newline at end of file diff --git a/tensorflow/compiler/xla/index_util.cc b/tensorflow/compiler/xla/index_util.cc index 3fadabcf5207097aa875d654320b930b1ed94ad3..7e22a32e545e4155545ffcfb9582187eadec3a82 100644 --- a/tensorflow/compiler/xla/index_util.cc +++ b/tensorflow/compiler/xla/index_util.cc @@ -29,8 +29,6 @@ namespace xla { /* static */ int64 IndexUtil::MultidimensionalIndexToLinearIndex( const Shape& shape, absl::Span multi_index) { DCHECK_EQ(shape.dimensions_size(), multi_index.size()); - // Padding and nested layouts not supported yet. - DCHECK_EQ(0, shape.layout().padded_dimensions_size()); for (size_t i = 0; i < multi_index.size(); ++i) { DCHECK_GE(multi_index[i], 0); @@ -94,8 +92,6 @@ namespace xla { /* static */ std::vector IndexUtil::LinearIndexToMultidimensionalIndex( const Shape& shape, int64 linear_index) { - // Padding and nested layouts not supported yet. - DCHECK_EQ(0, shape.layout().padded_dimensions_size()); DCHECK_GE(linear_index, 0); DCHECK_LT(linear_index, ShapeUtil::ElementsIn(shape)); @@ -133,25 +129,19 @@ namespace xla { /* static */ int64 IndexUtil::GetDimensionStride(const Shape& shape, int64 dimension) { - int64 pdim_size = LayoutUtil::PaddedDimensions(shape).size(); int64 stride = 1; - DCHECK(pdim_size == 0 || pdim_size == shape.dimensions_size()); for (auto dim : LayoutUtil::MinorToMajor(shape)) { if (dim == dimension) { break; } - if (pdim_size == 0) { - stride *= shape.dimensions(dim); - } else { - stride *= LayoutUtil::PaddedDimension(shape, dim); - } + stride *= shape.dimensions()[dim]; } return stride; } /* static */ bool IndexUtil::IndexInBounds(const Shape& shape, absl::Span index) { - int64 rank = ShapeUtil::Rank(shape); + int64 rank = shape.rank(); if (rank != index.size()) { return false; } diff --git a/tensorflow/compiler/xla/index_util.h b/tensorflow/compiler/xla/index_util.h index 2979cf87dde92893ce2151cb09b46c8db8473b31..d76f61eb62c0fc89d6bc3ca2033e8c7170f30e78 100644 --- a/tensorflow/compiler/xla/index_util.h +++ b/tensorflow/compiler/xla/index_util.h @@ -21,6 +21,7 @@ limitations under the License. #include #include "absl/types/span.h" +#include "tensorflow/compiler/xla/shape.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/platform/macros.h" @@ -61,8 +62,7 @@ class IndexUtil { static bool BumpIndices(const Shape& shape, absl::Span indices); // Calculates the stride size (in number of elements, not byte size) of a - // given logical shape dimension (from 0 to rank-1). If available, padded - // dimensions are used. + // given logical shape dimension (from 0 to rank-1). // Example: // GetDimensionStride(F32[5,8,10,4]{3,2,1,0}, 1) == // sizeof(dimension(3)) * sizeof(dimension(2)) == 4 * 10 diff --git a/tensorflow/compiler/xla/index_util_test.cc b/tensorflow/compiler/xla/index_util_test.cc index 93522d2ca87a7eba8d3c7533785c54e63ce507b0..fa94d0afb4c9280b8f8fa9642c1b0ab7285ee6f3 100644 --- a/tensorflow/compiler/xla/index_util_test.cc +++ b/tensorflow/compiler/xla/index_util_test.cc @@ -24,8 +24,7 @@ limitations under the License. namespace xla { namespace { -void SetMinorToMajorLayout(Shape* shape, - std::initializer_list dimensions) { +void SetMinorToMajorLayout(Shape* shape, std::vector dimensions) { shape->mutable_layout()->clear_minor_to_major(); for (auto dimension : dimensions) { shape->mutable_layout()->add_minor_to_major(dimension); @@ -122,7 +121,7 @@ TEST(IndexUtilTest, LinearToMultiToLinear) { std::vector linear_indexes = {0, 1439999999, 1145567336, 43883404, 617295214, 1117613654}; - std::vector> minor_to_major_orders; + std::vector> minor_to_major_orders; minor_to_major_orders.push_back({6, 5, 4, 3, 2, 1, 0}); minor_to_major_orders.push_back({0, 1, 2, 3, 4, 5, 6}); minor_to_major_orders.push_back({4, 5, 1, 2, 6, 0, 3}); diff --git a/tensorflow/compiler/xla/layout.cc b/tensorflow/compiler/xla/layout.cc new file mode 100644 index 0000000000000000000000000000000000000000..e3b5fcd5274881cec31ecf906e3461685f82a1f4 --- /dev/null +++ b/tensorflow/compiler/xla/layout.cc @@ -0,0 +1,96 @@ +/* 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/layout.h" + +#include "absl/strings/str_cat.h" +#include "absl/strings/str_join.h" +#include "tensorflow/compiler/xla/layout_util.h" + +namespace xla { + +TileProto Tile::ToProto() const { + TileProto tile_proto; + for (int64 i : dimensions()) { + tile_proto.add_dimensions(i); + } + return tile_proto; +} + +string Tile::ToString() const { + return absl::StrCat("(", absl::StrJoin(dimensions(), ","), ")"); +} + +/* static */ Layout Layout::CreateFromProto(const LayoutProto& proto) { + Layout layout; + layout.set_format(proto.format()); + layout.minor_to_major_.reserve(proto.minor_to_major_size()); + for (const int64 dimension : proto.minor_to_major()) { + layout.add_minor_to_major(dimension); + } + layout.set_max_sparse_elements(proto.max_sparse_elements()); + for (const TileProto& tile_proto : proto.tiles()) { + *layout.add_tiles() = Tile::CreateFromProto(tile_proto); + } + layout.set_element_size_in_bits(proto.element_size_in_bits()); + return layout; +} + +LayoutProto Layout::ToProto() const { + LayoutProto proto; + proto.set_format(format_); + proto.mutable_minor_to_major()->Reserve(minor_to_major_size()); + for (const int64 dimension : minor_to_major()) { + proto.add_minor_to_major(dimension); + } + proto.set_max_sparse_elements(max_sparse_elements_); + for (const Tile& tile : tiles()) { + *proto.add_tiles() = tile.ToProto(); + } + proto.set_element_size_in_bits(element_size_in_bits()); + return proto; +} + +string Layout::ToString() const { + // TODO(b/119839262): Emit tiles in string. + if (format() == SPARSE) { + return absl::StrCat("sparse{", max_sparse_elements(), "}"); + } else if (format() == DENSE) { + return absl::StrCat("{", absl::StrJoin(minor_to_major(), ","), "}"); + } else { + CHECK_EQ(format(), INVALID_FORMAT); + return "invalid{}"; + } +} + +bool Layout::operator==(const Layout& other) const { + return (other.format() == format() && + other.minor_to_major() == minor_to_major() && + other.element_size_in_bits() == element_size_in_bits() && + other.max_sparse_elements() == max_sparse_elements() && + other.tiles() == tiles()); +} + +std::ostream& operator<<(std::ostream& out, const Tile& tile) { + out << tile.ToString(); + return out; +} + +std::ostream& operator<<(std::ostream& out, const Layout& layout) { + out << layout.ToString(); + return out; +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/layout.h b/tensorflow/compiler/xla/layout.h new file mode 100644 index 0000000000000000000000000000000000000000..313368c39e4c976fc481941eb17325101f2ba69a --- /dev/null +++ b/tensorflow/compiler/xla/layout.h @@ -0,0 +1,187 @@ +/* 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_LAYOUT_H_ +#define TENSORFLOW_COMPILER_XLA_LAYOUT_H_ + +#include + +#include "absl/types/span.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/types.h" + +namespace xla { + +// Describes a tile used in tiling-based layout. Refer to +// g3doc/third_party/tensorflow/compiler/xla/g3doc/layout_with_tiling.md for +// details. +class Tile { + public: + Tile() = default; + explicit Tile(absl::Span dimensions) + : dimensions_(dimensions.begin(), dimensions.end()) {} + + // De/Serialize a Tile to and from a TileProto. + static Tile CreateFromProto(const TileProto& tile_proto) { + return Tile(AsInt64Slice(tile_proto.dimensions())); + } + TileProto ToProto() const; + + bool operator==(const Tile& other) const { + return dimensions() == other.dimensions(); + } + bool operator!=(const Tile& other) const { return !(*this == other); } + + string ToString() const; + + // Returns the bound of the tile in the given dimension index. + int64 dimension(int i) const { return dimensions_.at(i); } + + // Returns the dimensions of the tile. + const std::vector& dimensions() const { return dimensions_; } + + private: + // The bounds of the tile. + std::vector dimensions_; +}; + +class Layout { + public: + Layout() = default; + + // Constructs a dense layout with the given minor-to-major order. + explicit Layout(absl::Span minor_to_major) + : format_(DENSE), + minor_to_major_(minor_to_major.begin(), minor_to_major.end()) {} + + // Constructs a dense tiled layout with the given minor-to-major order and + // tiles. + Layout(absl::Span minor_to_major, absl::Span tiles) + : format_(DENSE), + minor_to_major_(minor_to_major.begin(), minor_to_major.end()), + tiles_(tiles.begin(), tiles.end()) {} + + // Construct a shape from a LayoutProto. + static Layout CreateFromProto(const LayoutProto& proto); + + // Returns a LayoutProto representation of the Layout. + LayoutProto ToProto() const; + + // Returns a human-readable string that represents this layout. + string ToString() const; + + bool operator==(const Layout& other) const; + bool operator!=(const Layout& other) const { return !(*this == other); } + + // The following methods mirror the protobuf generated code interface for the + // message LayoutProto. This enabled easy migration of this data structure + // from a proto to a proper C++ class. + // + // TODO(b/29771030): Replace or augment these methods with a more ergonomic + // interface. + + // Methods for accessing the format. + Format format() const { return format_; } + Layout& set_format(Format value) { + format_ = value; + return *this; + } + + // Methods for accessing the minor-to-major array. + int minor_to_major_size() const { return minor_to_major_.size(); } + int64 minor_to_major(int index) const { return minor_to_major_.at(index); } + Layout& set_minor_to_major(int index, int64 value) { + minor_to_major_.at(index) = value; + return *this; + } + Layout& add_minor_to_major(int64 value) { + minor_to_major_.push_back(value); + return *this; + } + Layout& clear_minor_to_major() { + minor_to_major_.clear(); + return *this; + } + const std::vector& minor_to_major() const { return minor_to_major_; } + std::vector* mutable_minor_to_major() { return &minor_to_major_; } + + // Methods for accessing the tile field. + int tiles_size() const { return tiles_.size(); } + const Tile& tiles(int index) const { return tiles_.at(index); } + Tile* mutable_tiles(int index) { return &tiles_.at(index); } + Tile* add_tiles() { + tiles_.push_back(Tile()); + return &tiles_.back(); + } + Layout& clear_tiles() { + tiles_.clear(); + return *this; + } + const std::vector& tiles() const { return tiles_; } + std::vector* mutable_tiles() { return &tiles_; } + + // Methods for accessing the int64 fields. + int64 max_sparse_elements() const { return max_sparse_elements_; } + Layout& set_max_sparse_elements(int64 value) { + max_sparse_elements_ = value; + return *this; + } + int64 element_size_in_bits() const { return element_size_in_bits_; } + Layout& set_element_size_in_bits(int64 value) { + element_size_in_bits_ = value; + return *this; + } + + void Swap(Layout* other) { + using std::swap; + swap(*this, *other); + } + + void Clear() { + format_ = INVALID_FORMAT; + minor_to_major_.clear(); + max_sparse_elements_ = 0; + element_size_in_bits_ = 0; + } + + public: + // The format of this layout. + Format format_ = INVALID_FORMAT; + + // Sequence of dimension numbers, from minor (fastest varying index) to major + // (slowest varying index). + std::vector minor_to_major_; + + // The maximum number of elements that can be stored for SPARSE formats. This + // can be used to determine the maximum size in bytes of arrays stored in + // memory. This field must be zero unless the format is SPARSE. + int64 max_sparse_elements_ = 0; + + // The number of bits used to store an individual array element. + int64 element_size_in_bits_ = 0; + + // The tiles used in tiling-based layout. + std::vector tiles_; +}; + +std::ostream& operator<<(std::ostream& out, const Tile& Tile); +std::ostream& operator<<(std::ostream& out, const Layout& layout); + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_LAYOUT_H_ diff --git a/tensorflow/compiler/xla/layout_test.cc b/tensorflow/compiler/xla/layout_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..fb6abd3f6523b978e72b21ec082ae06973e86243 --- /dev/null +++ b/tensorflow/compiler/xla/layout_test.cc @@ -0,0 +1,104 @@ +/* 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/layout.h" + +#include "absl/strings/str_cat.h" +#include "absl/strings/str_join.h" +#include "tensorflow/compiler/xla/layout_util.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/status_macros.h" +#include "tensorflow/compiler/xla/test.h" +#include "tensorflow/compiler/xla/test_helpers.h" +#include "tensorflow/compiler/xla/types.h" +#include "tensorflow/compiler/xla/util.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" + +namespace xla { +namespace { + +class LayoutTest : public ::testing::Test {}; + +TEST_F(LayoutTest, ToString) { + EXPECT_EQ(Layout().ToString(), "invalid{}"); + EXPECT_EQ(Layout({4, 5, 6}).ToString(), "{4,5,6}"); + EXPECT_EQ(Layout().set_format(SPARSE).set_max_sparse_elements(123).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}"); + EXPECT_EQ( + Layout({1, 0}, {Tile({2, 55})}).set_element_size_in_bits(42).ToString(), + "{1,0}"); +} + +TEST_F(LayoutTest, StreamOut) { + { + std::ostringstream oss; + oss << Tile({7, 8}); + EXPECT_EQ(oss.str(), "(7,8)"); + } + + { + std::ostringstream oss; + oss << Layout({0, 1, 2}); + EXPECT_EQ(oss.str(), "{0,1,2}"); + } +} + +TEST_F(LayoutTest, SparseLayoutMaxElements) { + EXPECT_EQ(LayoutUtil::MaxSparseElements(LayoutUtil::MakeSparseLayout(101)), + 101); +} + +TEST_F(LayoutTest, Equality) { + EXPECT_EQ(Layout(), Layout()); + const std::vector empty_dims; + EXPECT_EQ(Layout(empty_dims), Layout(empty_dims)); + EXPECT_NE(Layout(), Layout(empty_dims)); + EXPECT_EQ(Layout({0, 1, 2, 3}), Layout({0, 1, 2, 3})); + EXPECT_NE(Layout({0, 1, 2, 3}), Layout({0, 1, 2})); + EXPECT_EQ(Layout({0, 1, 2}, {Tile({42, 44})}), + Layout({0, 1, 2}, {Tile({42, 44})})); + EXPECT_NE(Layout({0, 1, 2}, {Tile({42, 44})}), + Layout({0, 1, 2}, {Tile({42, 45})})); + EXPECT_NE(Layout({0, 1, 2}, {Tile({42, 44})}), Layout({0, 1, 2, 3})); + EXPECT_EQ(Layout({0, 1, 2}).set_element_size_in_bits(33), + Layout({0, 1, 2}).set_element_size_in_bits(33)); + EXPECT_NE(Layout({0, 1, 2}).set_element_size_in_bits(33), + Layout({0, 1, 2}).set_element_size_in_bits(7)); + EXPECT_EQ(Layout().set_format(SPARSE), Layout().set_format(SPARSE)); + EXPECT_EQ(Layout().set_format(SPARSE).set_max_sparse_elements(42), + Layout().set_format(SPARSE).set_max_sparse_elements(42)); + EXPECT_NE(Layout().set_format(SPARSE).set_max_sparse_elements(42), + Layout().set_format(SPARSE).set_max_sparse_elements(24)); +} + +TEST_F(LayoutTest, LayoutToFromProto) { + // Round-trips a Layout through proto de/serialization. + auto expect_unchanged = [](const Layout& layout) { + EXPECT_EQ(layout, Layout::CreateFromProto(layout.ToProto())); + }; + + expect_unchanged(Layout()); + expect_unchanged(Layout({1, 3, 2, 0})); + expect_unchanged(Layout().set_format(SPARSE)); + expect_unchanged(Layout().set_format(SPARSE).set_max_sparse_elements(123)); + expect_unchanged(Layout({0, 1}).set_element_size_in_bits(42)); + expect_unchanged(Layout({3, 2, 1, 0}, {Tile({42, 123}), Tile({4, 5})})); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/layout_util.cc b/tensorflow/compiler/xla/layout_util.cc index 66af644cf78f3cc3ebecfaba67cf7d023b0360d5..2fe9b56c6bdffb931726f60ab75081361b43ebb4 100644 --- a/tensorflow/compiler/xla/layout_util.cc +++ b/tensorflow/compiler/xla/layout_util.cc @@ -41,15 +41,13 @@ namespace { // Internal helper for GetDefaultLayoutForShape and SetToDefaultLayout. Sets // minor_to_major to the value that represents the default layout. -void SetDefaultLayoutToContainer( - tensorflow::protobuf::RepeatedField* - minor_to_major) { +void SetDefaultLayoutToContainer(std::vector* minor_to_major) { // The default XLA layout is major-to-minor (dim 0 is major). // For more information on XLA layouts, see: // https://www.tensorflow.org/performance/xla/shapes const int64 size = minor_to_major->size(); for (int64 i = 0; i < size; ++i) { - minor_to_major->Set(i, size - 1 - i); + (*minor_to_major)[i] = size - 1 - i; } } @@ -94,9 +92,8 @@ namespace { Layout CreateDefaultLayoutForRank(int64 rank) { Layout layout; layout.set_format(DENSE); - tensorflow::protobuf::RepeatedField* - minor_to_major = layout.mutable_minor_to_major(); - minor_to_major->Resize(rank, 0); + std::vector* minor_to_major = layout.mutable_minor_to_major(); + minor_to_major->resize(rank, 0); SetDefaultLayoutToContainer(minor_to_major); return layout; } @@ -104,13 +101,13 @@ Layout CreateDefaultLayoutForRank(int64 rank) { } // namespace /* static */ Layout LayoutUtil::GetDefaultLayoutForShape(const Shape& shape) { - if (ShapeUtil::IsOpaque(shape) || ShapeUtil::IsToken(shape)) { + if (shape.IsOpaque() || shape.IsToken()) { // Opaque and token types have empty layouts. return Layout(); } // A Layout proto corresponds to a single array, not a tuple. - CHECK(ShapeUtil::IsArray(shape)); + CHECK(shape.IsArray()); return CreateDefaultLayoutForRank(shape.dimensions_size()); } @@ -131,17 +128,16 @@ Layout CreateDefaultLayoutForRank(int64 rank) { } /* static */ void LayoutUtil::SetToDefaultLayout(Shape* shape) { - if (ShapeUtil::IsTuple(*shape)) { + if (shape->IsTuple()) { // Tuple shape. for (auto& element_shape : *shape->mutable_tuple_shapes()) { SetToDefaultLayout(&element_shape); } shape->clear_layout(); - } else if (ShapeUtil::IsArray(*shape)) { + } else if (shape->IsArray()) { shape->mutable_layout()->set_format(DENSE); - tensorflow::protobuf::RepeatedField* - minor_to_major = shape->mutable_layout()->mutable_minor_to_major(); - minor_to_major->Resize(shape->dimensions_size(), 0); + auto* minor_to_major = shape->mutable_layout()->mutable_minor_to_major(); + minor_to_major->resize(shape->dimensions_size(), 0); SetDefaultLayoutToContainer(minor_to_major); } else { // Opaque, token types etc. have no layout. @@ -164,7 +160,7 @@ Layout CreateDefaultLayoutForRank(int64 rank) { /* static */ Status LayoutUtil::ValidateLayoutInShape( const Shape& shape, bool allow_missing_layouts) { - if (ShapeUtil::IsTuple(shape)) { + if (shape.IsTuple()) { // Tuple shape. if (shape.has_layout()) { return InvalidArgument("tuple should not have a layout field"); @@ -174,7 +170,7 @@ Layout CreateDefaultLayoutForRank(int64 rank) { ValidateLayoutInShape(element_shape, allow_missing_layouts)); } return Status::OK(); - } else if (ShapeUtil::IsArray(shape)) { + } else if (shape.IsArray()) { if (!shape.has_layout()) { if (allow_missing_layouts) { return Status::OK(); @@ -196,13 +192,12 @@ Layout CreateDefaultLayoutForRank(int64 rank) { /* static */ Status LayoutUtil::ValidateLayoutForShape(const Layout& layout, const Shape& shape) { - if (ShapeUtil::IsTuple(shape)) { + if (shape.IsTuple()) { return InvalidArgument("a single Layout is not valid for tuple shapes"); } - if (!ShapeUtil::IsArray(shape)) { - if (layout.minor_to_major_size() != 0 || - layout.padded_dimensions_size() != 0) { + if (!shape.IsArray()) { + if (layout.minor_to_major_size() != 0) { return InvalidArgument( "shape of primitive type %s should not have a non-trivial layout", PrimitiveType_Name(shape.element_type())); @@ -211,25 +206,24 @@ Layout CreateDefaultLayoutForRank(int64 rank) { } if (layout.format() == INVALID_FORMAT || !Format_IsValid(layout.format())) { - return InvalidArgument( - "Layout has an invalid format (%d) in layout {%s}, shape {%s}", - layout.format(), layout.ShortDebugString(), shape.ShortDebugString()); + return InvalidArgument("Layout has an invalid format (%d)", + layout.format()); } if (layout.format() == DENSE) { - if (layout.minor_to_major_size() != ShapeUtil::Rank(shape)) { + if (layout.minor_to_major_size() != shape.rank()) { return InvalidArgument( "layout minor_to_major field contains %d elements, " "but shape is rank %d: {%s}; shape: %s", - layout.minor_to_major_size(), ShapeUtil::Rank(shape), + layout.minor_to_major_size(), shape.rank(), absl::StrJoin(layout.minor_to_major(), ", "), shape.ShortDebugString()); } - std::vector dimensions_in_layout(ShapeUtil::Rank(shape), false); - for (int64 i = 0; i < ShapeUtil::Rank(shape); ++i) { + std::vector dimensions_in_layout(shape.rank(), false); + for (int64 i = 0; i < shape.rank(); ++i) { int64 dim = layout.minor_to_major(i); - if (dim < 0 || dim >= ShapeUtil::Rank(shape)) { + if (dim < 0 || dim >= shape.rank()) { return InvalidArgument( "layout minor_to_major field has out-of-bounds value: %s", HumanString(layout)); @@ -241,28 +235,6 @@ Layout CreateDefaultLayoutForRank(int64 rank) { } dimensions_in_layout[dim] = true; } - - if (layout.padded_dimensions_size() > 0) { - if (layout.padded_dimensions_size() != ShapeUtil::Rank(shape)) { - return InvalidArgument( - "layout has %d padded dimensions, but shape is rank %d", - layout.padded_dimensions_size(), ShapeUtil::Rank(shape)); - } - for (int i = 0; i < layout.padded_dimensions_size(); ++i) { - if (layout.padded_dimensions(i) < shape.dimensions(i)) { - return InvalidArgument( - "for dimension %d, dimension padding (%d) is smaller than " - "the dimension size (%d) of the shape", - i, layout.padded_dimensions(i), shape.dimensions(i)); - } - } - } - } - - if (layout.format() == SPARSE) { - if (!layout.padded_dimensions().empty()) { - return InvalidArgument("Sparse layout has padded dimensions"); - } } return Status::OK(); @@ -283,8 +255,7 @@ Layout CreateDefaultLayoutForRank(int64 rank) { } /* static */ bool LayoutUtil::IsDenseArray(const Shape& shape) { - return ShapeUtil::IsArray(shape) && shape.has_layout() && - IsDense(shape.layout()); + return shape.IsArray() && shape.has_layout() && IsDense(shape.layout()); } /* static */ bool LayoutUtil::IsDense(const Layout& layout) { @@ -303,41 +274,8 @@ Layout CreateDefaultLayoutForRank(int64 rank) { layout.minor_to_major().end(), std::greater()); } -/* static */ bool LayoutUtil::IsPadded(const Shape& shape) { - if (!ShapeUtil::IsArray(shape) || !HasLayout(shape) || - shape.layout().padded_dimensions_size() == 0) { - return false; - } - CHECK(IsDenseArray(shape)) << shape.ShortDebugString(); - CHECK_EQ(shape.dimensions_size(), shape.layout().padded_dimensions_size()); - for (int64 i = 0; i < shape.dimensions_size(); ++i) { - if (shape.layout().padded_dimensions(i) > shape.dimensions(i)) { - return true; - } - } - return false; -} - -/* static */ absl::Span LayoutUtil::PaddedDimensions( - const Shape& shape) { - CHECK(IsDenseArray(shape)); - return AsInt64Slice(shape.layout().padded_dimensions()); -} - -/* static */ int64 LayoutUtil::PaddedDimension(const Shape& shape, - int64 index) { - CHECK(IsDenseArray(shape)); - return shape.layout().padded_dimensions(index); -} - -/* static */ PaddingValue LayoutUtil::GetPaddingValue(const Shape& shape) { - CHECK(IsDenseArray(shape)); - return shape.layout().padding_value(); -} - /* static */ bool LayoutUtil::IsSparseArray(const Shape& shape) { - return ShapeUtil::IsArray(shape) && shape.has_layout() && - IsSparse(shape.layout()); + return shape.IsArray() && shape.has_layout() && IsSparse(shape.layout()); } /* static */ bool LayoutUtil::IsSparse(const Layout& layout) { @@ -350,11 +288,11 @@ Layout CreateDefaultLayoutForRank(int64 rank) { } /* static */ bool LayoutUtil::HasLayout(const Shape& shape) { - if (ShapeUtil::IsTuple(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); }); - } else if (!ShapeUtil::IsArray(shape)) { + 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; } @@ -371,7 +309,7 @@ Layout CreateDefaultLayoutForRank(int64 rank) { } /* static */ bool LayoutUtil::Equal(const Layout& lhs, const Layout& rhs) { - return protobuf_util::ProtobufEquals(lhs, rhs); + return lhs == rhs; } /* static */ absl::Span LayoutUtil::MinorToMajor( @@ -413,22 +351,18 @@ Layout CreateDefaultLayoutForRank(int64 rank) { } /* static */ string LayoutUtil::HumanString(const Layout& layout) { - if (IsSparse(layout)) { - return absl::StrCat("sparse{", layout.max_sparse_elements(), "}"); - } - CHECK(IsDense(layout)); - return absl::StrCat("{", absl::StrJoin(layout.minor_to_major(), ","), "}"); + return layout.ToString(); } namespace { // Internal helper for recursively copying layouts. Status CopyLayoutInternal(const Shape& src, Shape* dst) { - if (ShapeUtil::IsTuple(src) != ShapeUtil::IsTuple(*dst)) { + if (src.IsTuple() != dst->IsTuple()) { return InvalidArgument( "cannot copy layout from shape: shape structure differs"); } - if (ShapeUtil::IsTuple(src)) { + if (src.IsTuple()) { if (ShapeUtil::TupleElementCount(src) != ShapeUtil::TupleElementCount(*dst)) { return InvalidArgument( @@ -440,7 +374,7 @@ Status CopyLayoutInternal(const Shape& src, Shape* dst) { } } else { if (src.has_layout()) { - if (ShapeUtil::Rank(src) != ShapeUtil::Rank(*dst)) { + if (src.rank() != dst->rank()) { return InvalidArgument("cannot copy layout from shape: ranks differs"); } TF_RETURN_IF_ERROR( @@ -462,9 +396,9 @@ Status LayoutUtil::CopyLayoutBetweenShapes(const Shape& src, Shape* dst) { /* static */ bool LayoutUtil::LayoutsInShapesEqual(const Shape& lhs, const Shape& rhs) { - if (ShapeUtil::IsTuple(lhs)) { - if (!ShapeUtil::IsTuple(rhs) || ShapeUtil::TupleElementCount(lhs) != - ShapeUtil::TupleElementCount(rhs)) { + if (lhs.IsTuple()) { + if (!rhs.IsTuple() || ShapeUtil::TupleElementCount(lhs) != + ShapeUtil::TupleElementCount(rhs)) { return false; } for (int i = 0; i < ShapeUtil::TupleElementCount(lhs); ++i) { @@ -473,8 +407,8 @@ Status LayoutUtil::CopyLayoutBetweenShapes(const Shape& src, Shape* dst) { } } return true; - } else if (ShapeUtil::IsArray(lhs)) { - return ShapeUtil::Rank(lhs) == ShapeUtil::Rank(rhs) && + } else if (lhs.IsArray()) { + return lhs.rank() == rhs.rank() && LayoutUtil::Equal(lhs.layout(), rhs.layout()); } else { // Layouts of non-array and non-tuple shapes is ignored. @@ -490,7 +424,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; @@ -499,11 +433,6 @@ Status LayoutUtil::CopyLayoutBetweenShapes(const Shape& src, Shape* dst) { return true; } -std::ostream& operator<<(std::ostream& out, const Layout& layout) { - out << LayoutUtil::HumanString(layout); - return out; -} - /*static*/ size_t LayoutUtil::Hash(const Layout& layout) { using tensorflow::hash; using tensorflow::Hash64Combine; @@ -513,14 +442,14 @@ std::ostream& operator<<(std::ostream& out, const Layout& layout) { for (int64 minor_to_major : layout.minor_to_major()) { hash_value = Hash64Combine(hash_value, hash()(minor_to_major)); } + hash_value = Hash64Combine(hash_value, layout.max_sparse_elements()); - for (int64 padded_dim : layout.padded_dimensions()) { - hash_value = Hash64Combine(hash_value, hash()(padded_dim)); + for (Tile tile : layout.tiles()) { + for (int64 tile_dim : tile.dimensions()) { + hash_value = Hash64Combine(hash_value, hash()(tile_dim)); + } } - - hash_value = - Hash64Combine(hash_value, hash()(layout.padding_value())); - hash_value = Hash64Combine(hash_value, layout.max_sparse_elements()); + hash_value = Hash64Combine(hash_value, layout.element_size_in_bits()); return hash_value; } diff --git a/tensorflow/compiler/xla/layout_util.h b/tensorflow/compiler/xla/layout_util.h index 97806d7e3311141920551a17d56d8ae9a1fe4af9..609dba67bcdbcb11be0906b7d87a52a17ba0dfbd 100644 --- a/tensorflow/compiler/xla/layout_util.h +++ b/tensorflow/compiler/xla/layout_util.h @@ -21,6 +21,8 @@ limitations under the License. #include #include "absl/types/span.h" +#include "tensorflow/compiler/xla/layout.h" +#include "tensorflow/compiler/xla/shape.h" #include "tensorflow/compiler/xla/status.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" @@ -104,23 +106,6 @@ class LayoutUtil { // more minor, and so on until dimension N-1 which is the minor. static bool IsMonotonicWithDim0Major(const Layout& layout); - // Returns whether the layout of the given shape has padding (a - // padded_dimension value in Layout is greater than the corresponding - // dimension size). - static bool IsPadded(const Shape& shape); - - // Returns the padded_dimensions array for the given Shape. Requires that the - // shape is an array and has a dense layout. - static absl::Span PaddedDimensions(const Shape& shape); - - // Returns the given index of the padded_dimensions array for the given Shape. - // Requires that the shape is an array and has a dense layout. - static int64 PaddedDimension(const Shape& shape, int64 index); - - // Returns the padding_value for the given Shape. Requires that the shape is - // an array and has a dense layout. - static PaddingValue GetPaddingValue(const Shape& shape); - // Returns whether the given Shape is an array (i.e. not a tuple) and has a // sparse format layout. static bool IsSparseArray(const Shape& shape); @@ -211,8 +196,6 @@ class LayoutUtil { TF_DISALLOW_COPY_AND_ASSIGN(LayoutUtil); }; -std::ostream& operator<<(std::ostream& out, const Layout& layout); - } // namespace xla #endif // TENSORFLOW_COMPILER_XLA_LAYOUT_UTIL_H_ diff --git a/tensorflow/compiler/xla/layout_util_test.cc b/tensorflow/compiler/xla/layout_util_test.cc index a50d53eaeb15daa9f7a98a816e180d3a55568bb8..4cc94c270cd64eb19761cc1044861c7d185b7888 100644 --- a/tensorflow/compiler/xla/layout_util_test.cc +++ b/tensorflow/compiler/xla/layout_util_test.cc @@ -304,30 +304,6 @@ TEST_F(LayoutUtilTest, SetToDefaultLayoutTuple) { shape.tuple_shapes(1).layout())); } -TEST_F(LayoutUtilTest, IsPadded) { - Shape shape_without_layout = ShapeUtil::MakeShape(F32, {2, 3, 4}); - LayoutUtil::ClearLayout(&shape_without_layout); - EXPECT_FALSE(LayoutUtil::IsPadded(shape_without_layout)); - - Shape shape_with_layout = ShapeUtil::MakeShape(F32, {2, 3, 4}); - LayoutUtil::SetToDefaultLayout(&shape_with_layout); - EXPECT_FALSE(LayoutUtil::IsPadded(shape_with_layout)); - - // Add padding equal to the dimension sizes. In this case the padding is a - // nop. - Shape shape_with_degenerate_padding = ShapeUtil::MakeShape(F32, {2, 3, 4}); - shape_with_degenerate_padding.mutable_layout()->add_padded_dimensions(2); - shape_with_degenerate_padding.mutable_layout()->add_padded_dimensions(3); - shape_with_degenerate_padding.mutable_layout()->add_padded_dimensions(4); - EXPECT_FALSE(LayoutUtil::IsPadded(shape_with_degenerate_padding)); - - Shape shape_with_padding = ShapeUtil::MakeShape(F32, {2, 3, 4}); - shape_with_padding.mutable_layout()->add_padded_dimensions(2); - shape_with_padding.mutable_layout()->add_padded_dimensions(14); - shape_with_padding.mutable_layout()->add_padded_dimensions(42); - EXPECT_TRUE(LayoutUtil::IsPadded(shape_with_padding)); -} - TEST_F(LayoutUtilTest, DefaultLayoutGettersMajorToMinor) { EXPECT_TRUE(LayoutUtil::Equal(LayoutUtil::MakeLayout({1, 0}), LayoutUtil::GetDefaultLayoutForR2())); @@ -341,17 +317,6 @@ TEST_F(LayoutUtilTest, DefaultLayoutGettersMajorToMinor) { ShapeUtil::MakeShape(F32, {10, 20, 30, 15, 25})))); } -TEST_F(LayoutUtilTest, SparseLayoutMaxElements) { - EXPECT_EQ(LayoutUtil::MaxSparseElements(LayoutUtil::MakeSparseLayout(101)), - 101); -} - -TEST_F(LayoutUtilTest, StreamOut) { - std::ostringstream oss; - oss << LayoutUtil::MakeLayout({0, 1, 2}); - EXPECT_EQ(oss.str(), "{0,1,2}"); -} - TEST_F(LayoutUtilTest, ValidateLayout_ValidArrayLayout) { Shape shape = ShapeUtil::MakeShapeWithLayout(F32, {2, 3}, {0, 1}); auto status = diff --git a/tensorflow/compiler/xla/legacy_flags/BUILD b/tensorflow/compiler/xla/legacy_flags/BUILD deleted file mode 100644 index 3e79129aafd234e5eab05d205f2017b54057795e..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/legacy_flags/BUILD +++ /dev/null @@ -1,82 +0,0 @@ -# Legacy command-line flags for the XLA libraries. - -# Please do not add more flags to this package. - -# The XLA libraries were written in an environment that allowed command-line -# flags to be scattered freely throughout the libraries. This model, while -# initially convenient, leads to a proliferation in unused command-line flags -# in tests and binaries, and serious problems in servers, where one might wish -# parameters to be different in independent RPC calls to the same routine. -# -# Please don't add more flags. If you're a library author, pass options and -# parameters explicitly through the library's interface. - -package(default_visibility = ["//tensorflow:internal"]) - -licenses(["notice"]) # Apache 2.0 - -load("//tensorflow:tensorflow.bzl", "tf_cc_test") - -cc_library( - name = "parse_flags_from_env", - srcs = ["parse_flags_from_env.cc"], - hdrs = ["parse_flags_from_env.h"], - deps = - [ - "//tensorflow/compiler/xla:types", - "//tensorflow/core:framework_internal", - "//tensorflow/core:lib", - "@com_google_absl//absl/strings", - ], -) - -tf_cc_test( - name = "parse_flags_from_env_test", - srcs = ["parse_flags_from_env_test.cc"], - deps = - [ - ":parse_flags_from_env", - "//tensorflow/compiler/xla:types", - "//tensorflow/core:framework_internal", - "//tensorflow/core:lib", - "//tensorflow/core:test", - "@com_google_absl//absl/strings:str_format", - ], -) - -cc_library( - name = "debug_options_flags", - srcs = [ - "debug_options_flags.cc", - "debug_options_parsers.h", - ], - hdrs = ["debug_options_flags.h"], - 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", - ], -) - -tf_cc_test( - name = "debug_options_parsers_test", - size = "small", - srcs = [ - "debug_options_parsers.h", - "debug_options_parsers_test.cc", - ], - deps = - [ - "//tensorflow/compiler/xla:xla_proto", - "//tensorflow/compiler/xla/service:hlo", - "//tensorflow/core:framework_internal", - "//tensorflow/core:lib", - "//tensorflow/core:test", - "@com_google_absl//absl/strings", - "@com_google_absl//absl/strings:str_format", - ], -) diff --git a/tensorflow/compiler/xla/legacy_flags/parse_flags_from_env.cc b/tensorflow/compiler/xla/legacy_flags/parse_flags_from_env.cc deleted file mode 100644 index 2a4e49b05aa0d1eed2197095694cfc6aa8814983..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/legacy_flags/parse_flags_from_env.cc +++ /dev/null @@ -1,206 +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. -==============================================================================*/ - -// This module exports ParseFlagsFromEnv(), which allows other modules to parse -// flags from an environtment variable, or a file named by the environment -// variable. - -#include -#include -#include -#include - -#include "tensorflow/compiler/xla/legacy_flags/parse_flags_from_env.h" -#include "tensorflow/compiler/xla/types.h" -#include "tensorflow/core/platform/logging.h" -#include "tensorflow/core/platform/macros.h" -#include "tensorflow/core/platform/mutex.h" -#include "tensorflow/core/platform/types.h" -#include "tensorflow/core/util/command_line_flags.h" - -namespace xla { -namespace legacy_flags { - -static const char kEnvVar[] = "TF_XLA_FLAGS"; // environment variable queried -static const char kWS[] = " \t\r\n"; // whitespace - -// The following struct represents an argv[]-style array, parsed -// from data gleaned from the environment. -// -// As usual, an anonymous namespace is advisable to avoid -// constructor/destructor collisions with other "private" types -// in the same named namespace. -namespace { -struct EnvArgv { - EnvArgv() : initialized(false), argc(0) {} - bool initialized; // whether the other fields have been set. - int argc; // elements used in argv[] - std::vector argv; // flag arguments parsed from environment string. - std::vector argv_save; // saved values from argv[] to avoid leaks -}; -} // anonymous namespace - -// Append the string s0[0, .., s0len-1] concatenated with s1[0, .., s1len-1] as -// a newly allocated nul-terminated string to the array *a. If s0==nullptr, a -// nullptr is appended without increasing a->argc. -static void AppendToEnvArgv(const char* s0, size_t s0len, const char* s1, - size_t s1len, EnvArgv* a) { - if (s0 == nullptr) { - a->argv.push_back(nullptr); - a->argv_save.push_back(nullptr); - } else { - string s = string(s0, s0len) + string(s1, s1len); - char* str = strdup(s.c_str()); - a->argv.push_back(str); - a->argv_save.push_back(str); - a->argc++; - } -} - -// Like s.find_first_of(x, pos), but return s.size() when find_first_of() would -// return string::npos. This avoids if-statements elsewhere. -static size_t FindFirstOf(const string& s, const char* x, size_t pos) { - size_t result = s.find_first_of(x, pos); - return result == string::npos ? s.size() : result; -} - -// Like s.find_first_not_of(x, pos), but return s.size() when -// find_first_not_of() would return string::npos. This avoids if-statements -// elsewhere. -static size_t FindFirstNotOf(const string& s, const char* x, size_t pos) { - size_t result = s.find_first_not_of(x, pos); - return result == string::npos ? s.size() : result; -} - -// Given a string containing flags, parse them into the XLA command line flags. -// The parse is best effort, and gives up on the first syntax error. -static void ParseArgvFromString(const string& flag_str, EnvArgv* a) { - size_t b = FindFirstNotOf(flag_str, kWS, 0); - while (b != flag_str.size() && flag_str[b] == '-') { - // b is the index of the start of a flag. - // 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]))) { - e++; - } - if (e != flag_str.size() && flag_str[e] == '=' && - e + 1 != flag_str.size() && strchr("'\"", flag_str[e + 1]) != nullptr) { - // A flag of the form --flag="something in double or single quotes" - int c; - e++; // point just past '=' - size_t eflag = e; - char quote = flag_str[e]; - e++; // point just past quote - // Put in value the string with quotes removed. - string value; - for (; e != flag_str.size() && (c = flag_str[e]) != quote; e++) { - if (quote == '"' && c == '\\' && e + 1 != flag_str.size()) { - // Handle backslash in double quoted strings. They are literal in - // single-quoted strings. - e++; - c = flag_str[e]; - } - value += c; - } - if (e != flag_str.size()) { // skip final " or ' - e++; - } - AppendToEnvArgv(flag_str.data() + b, eflag - b, value.data(), - value.size(), a); - } else { // A flag without a quoted value. - e = FindFirstOf(flag_str, kWS, e); - AppendToEnvArgv(flag_str.data() + b, e - b, "", 0, a); - } - b = FindFirstNotOf(flag_str, kWS, e); - } -} - -// Call ParseArgvFromString(..., a) on a string derived from the setting of an -// environment variable kEnvVar, or a file it points to. -static void SetArgvFromEnv(EnvArgv* a) { - if (!a->initialized) { - static const char kDummyArgv[] = ""; - AppendToEnvArgv(kDummyArgv, strlen(kDummyArgv), nullptr, 0, - a); // dummy argv[0] - const char* env = getenv(kEnvVar); - if (env == nullptr || env[0] == '\0') { - // nothing - } else if (env[strspn(env, kWS)] == '-') { // flags in env var value - ParseArgvFromString(env, a); - } else { // assume it's a file name - FILE* fp = fopen(env, "r"); - if (fp != nullptr) { - string str; - char buf[512]; - int n; - while ((n = fread(buf, 1, sizeof(buf), fp)) > 0) { - str.append(buf, n); - } - fclose(fp); - ParseArgvFromString(str, a); - } - } - AppendToEnvArgv(nullptr, 0, nullptr, 0, a); // add trailing nullptr to *a. - a->initialized = true; - } -} - -// The simulated argv[] parsed from the environment. -static EnvArgv* env_argv; - -// Used to protect accesses to env_argv. -static tensorflow::mutex env_argv_mu(tensorflow::LINKER_INITIALIZED); - -// Call Flags::Parse(argc, argv, flag_list) against any as yet unrecognized -// flags passed in from the environment. -bool ParseFlagsFromEnv(const std::vector& flag_list) { - env_argv_mu.lock(); - if (env_argv == nullptr) { - env_argv = new EnvArgv; - } - SetArgvFromEnv(env_argv); // a no-op if already initialized - bool result = - tensorflow::Flags::Parse(&env_argv->argc, &env_argv->argv[0], flag_list); - env_argv_mu.unlock(); - return result; -} - -// Testing only. -// Reset the env_argv struct so that subsequent calls to ParseFlagsFromEnv() -// will parse the environment variable (or the file it points to) anew, and set -// *pargc, and *pargv to point to the internal locations of the argc and argv -// constructed from the environment. -void ResetFlagsFromEnvForTesting(int** pargc, std::vector** pargv) { - env_argv_mu.lock(); - if (env_argv == nullptr) { - env_argv = new EnvArgv; - } - if (!env_argv->argv_save.empty()) { - for (int i = 0; env_argv->argv_save[i] != nullptr; i++) { - free(env_argv->argv_save[i]); - } - } - env_argv->initialized = false; - env_argv->argc = 0; - env_argv->argv.clear(); - env_argv->argv_save.clear(); - env_argv_mu.unlock(); - *pargc = &env_argv->argc; - *pargv = &env_argv->argv; -} - -} // namespace legacy_flags -} // namespace xla diff --git a/tensorflow/compiler/xla/legacy_flags/parse_flags_from_env.h b/tensorflow/compiler/xla/legacy_flags/parse_flags_from_env.h deleted file mode 100644 index b54482ad2ba2224c781861341a80ceb878ffd343..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/legacy_flags/parse_flags_from_env.h +++ /dev/null @@ -1,66 +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_LEGACY_FLAGS_PARSE_FLAGS_FROM_ENV_H_ -#define TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_PARSE_FLAGS_FROM_ENV_H_ - -// This module exports ParseFlagsFromEnv(), which allows other modules to parse -// flags from the environtment variable TF_XLA_FLAGS, or (if the first -// non-whitespace in the variable value is not '-'), a file named by that -// environment variable. The accepted syntax is that flags arguments are of -// the form --flag=value or (for boolean flags) --flag, and are whitespace -// separated. The may be one of: -// - -// in which case the effective value is the string itself -// - in which case the effective value is the -// string with the single-quotes removed -// - in which case the effective value if the -// string with the double-quotes removed, and escaped sequences of -// replaced by . -// -// Flags values inconsistent with the type of the flag will be rejected by the -// flag parser. -// -// Examples: -// TF_XLA_FLAGS="--foo=bar --wombat='value with a space'" -// -// TF_XLA_FLAGS=/tmp/flagfile -// where /tmp/flagfile might contain -// --some_flag="This is a string containing a \" and a '." -// --another_flag=wombats - -#include - -#include "tensorflow/compiler/xla/types.h" -#include "tensorflow/core/platform/types.h" -#include "tensorflow/core/util/command_line_flags.h" - -namespace xla { -namespace legacy_flags { - -// Call tensorflow::Flags::Parse(argc, argv, flag_list) against any as yet -// unrecognized flags passed in from the environment, and return its -// return value. -bool ParseFlagsFromEnv(const std::vector& flag_list); - -// Used only for testing. Not to be used by clients. -void ResetFlagsFromEnvForTesting(int** pargc, std::vector** pargv); - -} // namespace legacy_flags -} // namespace xla - -#endif // TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_PARSE_FLAGS_FROM_ENV_H_ diff --git a/tensorflow/compiler/xla/literal.cc b/tensorflow/compiler/xla/literal.cc index 8a8f49ccd041ea9454eb20e216d533b43965651b..8600e8752cfbe072407391559d210d0b49bea511 100644 --- a/tensorflow/compiler/xla/literal.cc +++ b/tensorflow/compiler/xla/literal.cc @@ -22,16 +22,19 @@ limitations under the License. #include #include +#include "absl/base/casts.h" #include "absl/memory/memory.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_format.h" #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/core/lib/core/casts.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" @@ -62,6 +65,14 @@ void ConvertEndianShort(char* bytes, int64 size) { } } +// Since Eigen::half doesn't satisfy the absl::bit_cast contract, we need to be +// able to transparently access the raw 16-bit value contained within. +template +T GetRawValue(T val) { + return val; +} +uint16 GetRawValue(Eigen::half val) { return val.x; } + } // namespace LiteralBase::~LiteralBase() {} @@ -98,7 +109,7 @@ Literal::Literal(const Shape& shape) : Literal(shape, /*allocate_arrays=*/true) {} void Literal::SetPiece(const Shape& shape, Piece* piece, bool allocate_arrays) { - if (ShapeUtil::IsTuple(shape)) { + if (shape.IsTuple()) { for (int i = 0; i < ShapeUtil::TupleElementCount(shape); ++i) { const Shape& subshape = shape.tuple_shapes(i); @@ -109,7 +120,7 @@ void Literal::SetPiece(const Shape& shape, Piece* piece, bool allocate_arrays) { piece->emplace_back(std::move(child_piece)); } - } else if (ShapeUtil::IsArray(shape)) { + } else if (shape.IsArray()) { if (allocate_arrays) { if (LayoutUtil::IsSparseArray(shape)) { // For sparse arrays, the buffer must be of the size of the maximum @@ -120,7 +131,7 @@ void Literal::SetPiece(const Shape& shape, Piece* piece, bool allocate_arrays) { new char[max_sparse_elements * ShapeUtil::ByteSizeOfPrimitiveType(shape.element_type())]); piece->set_sparse_indices( - new SparseIndexArray(max_sparse_elements, ShapeUtil::Rank(shape))); + new SparseIndexArray(max_sparse_elements, shape.rank())); } else { piece->set_buffer(new char[piece->size_bytes()]); } @@ -178,7 +189,7 @@ Literal LiteralBase::CreateFromShape(const Shape& shape) { Literal literal(shape); literal.root_piece_->ForEachMutableSubpiece( [&](const ShapeIndex& index, Piece* piece) { - if (ShapeUtil::IsArray(piece->subshape())) { + if (piece->subshape().IsArray()) { memset(piece->untyped_data(), 0, piece->size_bytes()); } }); @@ -199,16 +210,15 @@ template Status MutableLiteralBase::CopySliceFromInternal( const LiteralBase& src_literal, absl::Span src_base, absl::Span dest_base, absl::Span copy_size) { - TF_RET_CHECK(ShapeUtil::Rank(src_literal.shape()) == src_base.size()); - TF_RET_CHECK(ShapeUtil::Rank(shape()) == dest_base.size()); + TF_RET_CHECK(src_literal.shape().rank() == src_base.size()); + TF_RET_CHECK(shape().rank() == dest_base.size()); auto linear_index = [](const Shape& shape, absl::Span multi_index) { return IndexUtil::MultidimensionalIndexToLinearIndex(shape, multi_index); }; - if (ShapeUtil::Rank(src_literal.shape()) == 0 || - ShapeUtil::Rank(shape()) == 0) { + if (src_literal.shape().rank() == 0 || shape().rank() == 0) { // If any of the two shapes are scalars, we can just call the StridedCopy() // directly, and we know we will be copying only one value. TF_RET_CHECK(copy_size.empty()); @@ -283,16 +293,17 @@ Status MutableLiteralBase::CopyElementFrom(const LiteralSlice& src_literal, if (!proto.has_shape()) { return InvalidArgument("LiteralProto has no shape"); } - if (ShapeUtil::HasPrimitiveType(proto.shape(), OPAQUE)) { + Shape shape(proto.shape()); + if (ShapeUtil::HasPrimitiveType(shape, OPAQUE)) { return InvalidArgument("Literal shape cannot include OPAQUE sub-shape"); } - if (!LayoutUtil::HasLayout(proto.shape())) { + if (!LayoutUtil::HasLayout(shape)) { return InvalidArgument("LiteralProto has no layout"); } - TF_RETURN_IF_ERROR(ShapeUtil::ValidateShapeWithOptionalLayout(proto.shape())); + TF_RETURN_IF_ERROR(ShapeUtil::ValidateShapeWithOptionalLayout(shape)); - Literal literal(proto.shape()); + Literal literal(shape); TF_RETURN_IF_ERROR(literal.root_piece_->ForEachMutableSubpieceWithStatus( [&](const ShapeIndex& index, Piece* piece) { @@ -302,7 +313,7 @@ Status MutableLiteralBase::CopyElementFrom(const LiteralSlice& src_literal, proto_element = &proto_element->tuple_literals(i); } - if (ShapeUtil::IsTuple(piece->subshape())) { + if (piece->subshape().IsTuple()) { if (proto_element->tuple_literals_size() != ShapeUtil::TupleElementCount(piece->subshape())) { return InvalidArgument( @@ -316,7 +327,7 @@ Status MutableLiteralBase::CopyElementFrom(const LiteralSlice& src_literal, return Status::OK(); } - CHECK(ShapeUtil::IsArray(piece->subshape())); + CHECK(piece->subshape().IsArray()); TF_RETURN_IF_ERROR(piece->CopyFromProto(*proto_element)); return Status::OK(); @@ -326,7 +337,7 @@ Status MutableLiteralBase::CopyElementFrom(const LiteralSlice& src_literal, } std::vector Literal::DecomposeTuple() { - CHECK(ShapeUtil::IsTuple(shape())); + CHECK(shape().IsTuple()); std::vector elements; for (int i = 0; i < ShapeUtil::TupleElementCount(shape()); ++i) { elements.push_back(Literal(ShapeUtil::GetSubshape(shape(), {i}), @@ -365,7 +376,7 @@ void CopyElementsBetween(absl::Span dest, if (ShapeUtil::IsZeroElementArray(dest_shape)) { return; } - std::vector index(ShapeUtil::Rank(dest_shape)); + std::vector index(dest_shape.rank()); do { dest[IndexUtil::MultidimensionalIndexToLinearIndex(dest_shape, index)] = src[IndexUtil::MultidimensionalIndexToLinearIndex(src_shape, index)]; @@ -382,7 +393,7 @@ Status LiteralBase::Piece::CopyFrom(const LiteralBase::Piece& src) { memcpy(buffer(), src.buffer(), src.size_bytes()); } else { TF_RET_CHECK(ShapeUtil::Compatible(src.subshape(), subshape())); - std::vector origin(ShapeUtil::Rank(subshape()), 0); + std::vector origin(subshape().rank(), 0); switch (subshape().element_type()) { #define COPY_ELEMENTS(XLA_T, NATIVE_T) \ case (XLA_T): \ @@ -402,6 +413,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: @@ -428,7 +440,7 @@ Status MutableLiteralBase::CopyFrom(const LiteralSlice& src_literal, } return root_piece_->ForEachMutableSubpieceWithStatus( [&](const ShapeIndex& index, Piece* piece) { - if (!ShapeUtil::IsArray(piece->subshape())) { + if (!piece->subshape().IsArray()) { return Status::OK(); } @@ -467,7 +479,7 @@ Status Literal::MoveFrom(Literal&& src_literal, src_literal.root_piece_->ForEachSubpiece( [&](const ShapeIndex& src_index, const Piece& src_piece) { - if (!ShapeUtil::IsArray(src_piece.subshape())) { + if (!src_piece.subshape().IsArray()) { return; } @@ -494,8 +506,8 @@ Status MutableLiteralBase::CopySliceFrom(const LiteralSlice& src_literal, absl::Span src_base, absl::Span dest_base, absl::Span copy_size) { - TF_RET_CHECK(ShapeUtil::IsArray(shape())) << ShapeUtil::HumanString(shape()); - TF_RET_CHECK(ShapeUtil::IsArray(src_literal.shape())) + TF_RET_CHECK(shape().IsArray()) << ShapeUtil::HumanString(shape()); + TF_RET_CHECK(src_literal.shape().IsArray()) << ShapeUtil::HumanString(src_literal.shape()); TF_RET_CHECK(ShapeUtil::SameElementType(src_literal.shape(), shape())); @@ -539,6 +551,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); @@ -552,8 +567,8 @@ Status MutableLiteralBase::CopySliceFrom(const LiteralSlice& src_literal, } void MutableLiteralBase::PopulateR1(const tensorflow::core::Bitmap& values) { - CHECK(ShapeUtil::IsArray(shape())); - CHECK_EQ(ShapeUtil::Rank(shape()), 1); + CHECK(shape().IsArray()); + CHECK_EQ(shape().rank(), 1); CHECK_EQ(element_count(), values.bits()); CHECK_EQ(shape().element_type(), PRED); for (int64 i = 0; i < static_cast(values.bits()); ++i) { @@ -582,7 +597,7 @@ Literal LiteralBase::Relayout(const Shape& shape_with_layout) const { ShapeUtil::ForEachSubshape( result.shape(), [this, &result](const Shape& subshape, const ShapeIndex& index) { - if (ShapeUtil::IsArray(subshape)) { + if (subshape.IsArray()) { TF_CHECK_OK(result.CopyFrom(*this, /*dest_shape_index=*/index, /*src_shape_index=*/index)); @@ -593,7 +608,7 @@ Literal LiteralBase::Relayout(const Shape& shape_with_layout) const { StatusOr LiteralBase::Broadcast( const Shape& result_shape, absl::Span dimensions) const { - if (!ShapeUtil::IsArray(shape())) { + if (!shape().IsArray()) { return InvalidArgument("Broadcast only supports arrays."); } @@ -633,13 +648,12 @@ StatusOr LiteralBase::Broadcast( StatusOr LiteralBase::Reshape( absl::Span dimensions) const { - if (!ShapeUtil::IsArray(shape())) { + if (!shape().IsArray()) { return InvalidArgument("Reshape does not support tuples."); } Literal output; if (!LayoutUtil::IsMonotonicWithDim0Major(shape().layout())) { - output = - Relayout(LayoutUtil::GetDefaultLayoutForRank(ShapeUtil::Rank(shape()))); + output = Relayout(LayoutUtil::GetDefaultLayoutForRank(shape().rank())); } else { output = Clone(); } @@ -661,8 +675,8 @@ StatusOr LiteralBase::Reshape( } Literal LiteralBase::Transpose(absl::Span permutation) const { - CHECK(ShapeUtil::IsArray(shape())) << "Tuple is not supported for transpose"; - CHECK(IsPermutation(permutation, ShapeUtil::Rank(shape()))) + CHECK(shape().IsArray()) << "Tuple is not supported for transpose"; + CHECK(IsPermutation(permutation, shape().rank())) << "Given permutation is not a permutation of dimension numbers"; // To transpose the array, we just permute the dimensions and layout, and // do a straight memory copy of the raw data set. @@ -701,10 +715,10 @@ template Literal LiteralBase::SliceInternal( const Shape& result_shape, absl::Span start_indices) const { Literal result_literal(result_shape); - DimensionVector new_indices(ShapeUtil::Rank(result_shape)); + DimensionVector new_indices(result_shape.rank()); result_literal.EachCell( [&](absl::Span indices, NativeT /*value*/) { - for (int64 i = 0; i < ShapeUtil::Rank(result_shape); ++i) { + for (int64 i = 0; i < result_shape.rank(); ++i) { new_indices[i] = indices[i] + start_indices[i]; } NativeT value = Get(new_indices); @@ -715,10 +729,10 @@ Literal LiteralBase::SliceInternal( Literal LiteralBase::Slice(absl::Span start_indices, absl::Span limit_indices) const { - CHECK(ShapeUtil::IsArray(shape())) << "tuple is not supported for slice"; + CHECK(shape().IsArray()) << "tuple is not supported for slice"; DimensionVector result_dimensions; - for (int64 dnum = 0; dnum < ShapeUtil::Rank(shape()); ++dnum) { + for (int64 dnum = 0; dnum < shape().rank(); ++dnum) { CHECK_GE(start_indices[dnum], 0); CHECK_LE(limit_indices[dnum], shape().dimensions(dnum)) << "dnum = " << dnum; @@ -758,6 +772,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()); @@ -806,6 +822,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()); } @@ -860,6 +880,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()); @@ -896,7 +921,7 @@ size_t LiteralBase::Hash() const { ShapeUtil::ForEachSubshape( shape(), [&](const Shape& subshape, const ShapeIndex& index) { - if (!ShapeUtil::IsArray(subshape)) { + if (!subshape.IsArray()) { return; } @@ -988,6 +1013,9 @@ void LiteralBase::Piece::SortSparseElements() { case C64: SortSparseElementsInternal(); break; + case C128: + SortSparseElementsInternal(); + break; case F16: SortSparseElementsInternal(); break; @@ -1012,167 +1040,150 @@ void LiteralBase::Piece::SortSparseElementsInternal() { namespace { -void ToStringHelper(const LiteralBase& literal, const ShapeIndex& shape_index, - bool print_layout, std::vector* pieces) { - const Shape& subshape = ShapeUtil::GetSubshape(literal.shape(), shape_index); - CHECK(LayoutUtil::HasLayout(literal.shape())); - CHECK(LayoutUtil::HasLayout(subshape)); +string ShapeToString(bool print_layout, const Shape& shape) { + return print_layout ? ShapeUtil::HumanStringWithLayout(shape) + : ShapeUtil::HumanString(shape); +} - auto shape_to_string = [print_layout](const Shape& shape) { - if (print_layout) { - return ShapeUtil::HumanStringWithLayout(shape); - } else { - return ShapeUtil::HumanString(shape); - } - }; +void ToStringHelper(const LiteralBase& literal, const ShapeIndex& shape_index, + bool print_shape, bool print_layout, + std::vector* pieces); - // TODO(b/32894291): refactor this code to reduce code duplication. - if (ShapeUtil::IsTuple(subshape)) { - pieces->push_back(shape_to_string(subshape)); - pieces->push_back(" (\n"); - std::vector tuple_pieces; - for (int i = 0; i < ShapeUtil::TupleElementCount(subshape); ++i) { - ShapeIndex element_index = shape_index; - element_index.push_back(i); - std::vector element_pieces; - ToStringHelper(literal, element_index, print_layout, &element_pieces); - tuple_pieces.push_back(absl::StrJoin(element_pieces, "")); +void TupleToStringHelper(const LiteralBase& literal, + const ShapeIndex& shape_index, bool print_shape, + bool print_layout, std::vector* pieces) { + const Shape& subshape = ShapeUtil::GetSubshape(literal.shape(), shape_index); + pieces->push_back("(\n"); + std::vector tuple_pieces; + for (int i = 0; i < ShapeUtil::TupleElementCount(subshape); ++i) { + ShapeIndex element_index = shape_index; + element_index.push_back(i); + std::vector element_pieces; + ToStringHelper(literal, element_index, print_shape, print_layout, + &element_pieces); + tuple_pieces.push_back(absl::StrJoin(element_pieces, "")); + } + pieces->push_back(absl::StrJoin(tuple_pieces, ",\n")); + pieces->push_back("\n)"); +} + +void SparseArrayToStringHelper(const LiteralBase& literal, + const Shape& subshape, bool print_shape, + bool print_layout, std::vector* pieces) { + if (print_shape) { + pieces->push_back(ShapeToString(print_layout, subshape)); + } + pieces->push_back("{"); + int64 rank = subshape.rank(); + int64 num_elements = literal.sparse_element_count(); + for (int64 i = 0; i < num_elements; ++i) { + if (i > 0) { + pieces->push_back(", "); } - pieces->push_back(absl::StrJoin(tuple_pieces, ",\n")); - pieces->push_back("\n)"); - return; - } - - if (ShapeUtil::IsToken(subshape)) { - pieces->push_back("token"); - return; - } - - if (LayoutUtil::IsSparseArray(subshape)) { - pieces->push_back(shape_to_string(subshape)); - pieces->push_back("{"); - int64 rank = ShapeUtil::Rank(subshape); - int64 num_elements = literal.sparse_element_count(); - for (int64 i = 0; i < num_elements; ++i) { - if (i > 0) { - pieces->push_back(", "); - } - if (rank == 1) { - pieces->push_back(StrCat(literal.GetSparseIndex(i)[0])); - pieces->push_back(": "); - } else { - pieces->push_back("["); - pieces->push_back(absl::StrJoin(literal.GetSparseIndex(i), ", ")); - pieces->push_back("]: "); - } - pieces->push_back(literal.GetSparseElementAsString(i)); + if (rank == 1) { + pieces->push_back(StrCat(literal.GetSparseIndex(i)[0])); + pieces->push_back(": "); + } else { + pieces->push_back("["); + pieces->push_back(absl::StrJoin(literal.GetSparseIndex(i), ", ")); + pieces->push_back("]: "); } - pieces->push_back("}"); - return; + pieces->push_back(literal.GetSparseElementAsString(i)); } + pieces->push_back("}"); +} - CHECK(LayoutUtil::IsDenseArray(subshape)); - - auto element_to_string = [&](absl::Span indices) -> string { - PrimitiveType element_type = subshape.element_type(); - if (element_type == PRED) { - // We display predicates in a densely packed form. - return literal.Get(indices, shape_index) ? "1" : "0"; - } - return ((!indices.empty() && indices.back() > 0) ? ", " : "") + - literal.GetAsString(indices, shape_index); - }; +void DenseArrayToStringHelper(const LiteralBase& literal, + const ShapeIndex& shape_index, bool print_shape, + bool print_layout, std::vector* pieces) { + const Shape& subshape = ShapeUtil::GetSubshape(literal.shape(), shape_index); + int64 rank = subshape.rank(); + + std::function dimensions, std::vector*)> + to_string_recursive = [&](absl::Span dimensions, + std::vector* accum_indices) { + // dimensions.size() decreases by 1 at each recursive call, + // and accum_indices->size() increases by 1. + // Their sum is equal to the rank of the tensor. + CHECK_EQ(rank, dimensions.size() + accum_indices->size()); + + auto brace_to_string = [&](string brace) -> string { + // Handle 1D tensor + if (rank == 1) { + return brace; + } + // Handle the innermost tensor of a 2D+ tensor. + if (dimensions.size() == 1 && brace == "{") { + return StrCat(" ", brace, dimensions[0] <= 1 ? "" : " "); + } + if (dimensions.size() == 1 && brace == "}") { + return StrCat(dimensions[0] <= 1 ? "" : " ", brace); + } + // Handle the non-innermost tensors of a 2D+ tensor. + if (brace == "{") { + if (rank > 3 && !accum_indices->empty() && + accum_indices->size() < rank) { + int index = accum_indices->size() - 1; + int value = accum_indices->back(); + return StrCat(brace, " /*i", index, "=", value, "*/\n"); + } + return StrCat(brace, "\n"); + } + return StrCat("\n", brace); + }; - if (ShapeUtil::Rank(subshape) == 0) { - pieces->push_back(literal.GetAsString({}, shape_index)); - } else if (ShapeUtil::Rank(subshape) == 1) { - pieces->push_back("{"); - for (int64 i0 = 0; i0 < subshape.dimensions(0); ++i0) { - pieces->push_back(element_to_string({i0})); - } - pieces->push_back("}"); - } else if (ShapeUtil::Rank(subshape) == 2) { - pieces->push_back(shape_to_string(subshape)); - pieces->push_back(" {\n"); - for (int64 i0 = 0; i0 < subshape.dimensions(0); ++i0) { - pieces->push_back(" { "); - for (int64 i1 = 0; i1 < subshape.dimensions(1); ++i1) { - pieces->push_back(element_to_string({i0, i1})); - } - pieces->push_back(" "); - pieces->push_back(i0 == subshape.dimensions(0) - 1 ? "}\n" : "},\n"); - } - pieces->push_back("}"); - } else if (ShapeUtil::Rank(subshape) == 3) { - pieces->push_back(shape_to_string(subshape)); - pieces->push_back(" {\n"); - for (int64 i0 = 0; i0 < subshape.dimensions(0); ++i0) { - pieces->push_back(i0 > 0 ? ",\n{" : "{"); - for (int64 i1 = 0; i1 < subshape.dimensions(1); ++i1) { - pieces->push_back(i1 > 0 ? ",\n { " : " { "); - for (int64 i2 = 0; i2 < subshape.dimensions(2); ++i2) { - pieces->push_back(element_to_string({i0, i1, i2})); - } - pieces->push_back(" }"); - } - pieces->push_back(" }"); - } - pieces->push_back("\n}"); - } else if (ShapeUtil::Rank(subshape) == 4) { - pieces->push_back(shape_to_string(subshape)); - pieces->push_back(" {\n"); - for (int64 i0 = 0; i0 < subshape.dimensions(0); ++i0) { - pieces->push_back(StrFormat(" { /*i0=%d*/\n", i0)); - for (int64 i1 = 0; i1 < subshape.dimensions(1); ++i1) { - pieces->push_back(StrFormat(" { /*i1=%d*/\n", i1)); - for (int64 i2 = 0; i2 < subshape.dimensions(2); ++i2) { - pieces->push_back(" {"); - for (int64 i3 = 0; i3 < subshape.dimensions(3); ++i3) { - pieces->push_back(element_to_string({i0, i1, i2, i3})); + if (dimensions.empty()) { + // Display predicates as 0s and 1s so that the string is more dense. + string elem; + if (subshape.element_type() == PRED && rank > 0) { + elem = literal.Get(*accum_indices, shape_index) ? "1" : "0"; + } else { + elem = literal.GetAsString(*accum_indices, shape_index); } - pieces->push_back(i2 == subshape.dimensions(2) - 1 ? "}\n" : "},\n"); - } - pieces->push_back(i1 == subshape.dimensions(1) - 1 ? " }\n" - : " },\n"); - } - pieces->push_back(i0 == subshape.dimensions(0) - 1 ? " }\n" : " },\n"); - } - pieces->push_back("}"); - } else if (ShapeUtil::Rank(subshape) == 5) { - pieces->push_back(shape_to_string(subshape)); - pieces->push_back(" {\n"); - for (int64 i0 = 0; i0 < subshape.dimensions(0); ++i0) { - pieces->push_back(StrFormat(" { /*i0=%d*/\n", i0)); - for (int64 i1 = 0; i1 < subshape.dimensions(1); ++i1) { - pieces->push_back(StrFormat(" { /*i1=%d*/\n", i1)); - for (int64 i2 = 0; i2 < subshape.dimensions(2); ++i2) { - pieces->push_back(StrFormat(" { /*i2=%d*/\n", i2)); - for (int64 i3 = 0; i3 < subshape.dimensions(3); ++i3) { - pieces->push_back(" {"); - for (int64 i4 = 0; i4 < subshape.dimensions(4); ++i4) { - pieces->push_back(element_to_string({i0, i1, i2, i3, i4})); + pieces->push_back(elem); + } else { + pieces->push_back(brace_to_string("{")); + for (int i = 0; i < dimensions[0]; ++i) { + std::vector cloned_indices(*accum_indices); + cloned_indices.push_back(i); + to_string_recursive(dimensions.subspan(1), &cloned_indices); + if (i < dimensions[0] - 1) { + pieces->push_back(","); + pieces->push_back(dimensions.size() > 1 ? "\n" : " "); } - pieces->push_back(i3 == subshape.dimensions(3) - 1 ? "}\n" - : "},\n"); } - pieces->push_back(i2 == subshape.dimensions(2) - 1 ? " }\n" - : " },\n"); + pieces->push_back(brace_to_string("}")); } - pieces->push_back(i1 == subshape.dimensions(1) - 1 ? " }\n" - : " },\n"); - } - pieces->push_back(i0 == subshape.dimensions(0) - 1 ? " }\n" : " },\n"); - } - pieces->push_back("}"); + }; + + if (print_shape) { + pieces->push_back(ShapeToString(print_layout, subshape)); + pieces->push_back(" "); + } + std::vector indices = {}; + std::vector dimensions(subshape.dimensions().begin(), + subshape.dimensions().end()); + to_string_recursive(dimensions, &indices); +} + +void ToStringHelper(const LiteralBase& literal, const ShapeIndex& shape_index, + bool print_shape, bool print_layout, + std::vector* pieces) { + const Shape& subshape = ShapeUtil::GetSubshape(literal.shape(), shape_index); + CHECK(LayoutUtil::HasLayout(literal.shape())); + CHECK(LayoutUtil::HasLayout(subshape)); + if (subshape.IsTuple()) { + TupleToStringHelper(literal, shape_index, print_shape, print_layout, + pieces); + } else if (subshape.IsToken()) { + pieces->push_back("token"); + } else if (LayoutUtil::IsSparseArray(subshape)) { + SparseArrayToStringHelper(literal, subshape, print_shape, print_layout, + pieces); } else { - pieces->push_back(shape_to_string(subshape)); - pieces->push_back(" {"); - literal.EachCellAsString( - [&](absl::Span indices, const string& value) { - pieces->push_back(" "); - pieces->push_back(value); - }); - pieces->push_back("}"); + CHECK(LayoutUtil::IsDenseArray(subshape)); + DenseArrayToStringHelper(literal, shape_index, print_shape, print_layout, + pieces); } } @@ -1183,10 +1194,27 @@ int64 LiteralBase::sparse_element_count() const { return sparse_indices()->index_count(); } -string LiteralBase::ToString(bool print_layout) const { +string LiteralBase::ToString() const { + std::vector pieces; + CHECK(LayoutUtil::HasLayout(this->shape())); + ToStringHelper(*this, {}, /*print_shape=*/true, + /*print_layout=*/false, &pieces); + return absl::StrJoin(pieces, ""); +} + +string LiteralBase::ToStringWithoutShape() const { std::vector pieces; CHECK(LayoutUtil::HasLayout(this->shape())); - ToStringHelper(*this, {}, print_layout, &pieces); + ToStringHelper(*this, {}, /*print_shape=*/false, + /*print_layout=*/false, &pieces); + return absl::StrJoin(pieces, ""); +} + +string LiteralBase::ToStringWithLayout() const { + std::vector pieces; + CHECK(LayoutUtil::HasLayout(this->shape())); + ToStringHelper(*this, {}, /*print_shape=*/true, + /*print_layout=*/true, &pieces); return absl::StrJoin(pieces, ""); } @@ -1207,7 +1235,7 @@ namespace { template Literal ConvertBetweenNativeTypesWithConverter(const LiteralBase& src_literal, const ConverterType& converter) { - CHECK(ShapeUtil::IsArray(src_literal.shape())); + CHECK(src_literal.shape().IsArray()); Literal result_literal(ShapeUtil::ChangeElementType( src_literal.shape(), primitive_util::NativeToPrimitiveType())); @@ -1222,23 +1250,56 @@ 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); } template -typename std::enable_if<(sizeof(NativeSrcT) == sizeof(NativeDestT)), +typename std::enable_if<(sizeof(NativeSrcT) == sizeof(NativeDestT) && + !std::is_same::value), Literal>::type BitcastBetweenNativeTypes(const LiteralBase& src_literal) { auto converter = [](NativeSrcT src) { - return tensorflow::bit_cast(src); + return absl::bit_cast(GetRawValue(src)); }; return ConvertBetweenNativeTypesWithConverter( src_literal, converter); } +template +typename std::enable_if<(sizeof(NativeSrcT) == sizeof(Eigen::half) && + std::is_same::value), + Literal>::type +BitcastBetweenNativeTypes(const LiteralBase& src_literal) { + // Eigen::half doesn't satisfy the absl::bit_cast contract, so explicitly + // cast to unsigned short and then use raw_uint16_to_half. + auto converter = [](NativeSrcT src) { + return Eigen::half_impl::raw_uint16_to_half( + absl::bit_cast(GetRawValue(src))); + }; + return ConvertBetweenNativeTypesWithConverter( + src_literal, converter); +} + // This template specialization is here to make the compiler happy. bit_cast has // a static check that the types are the same size. This specialization should // never be used because the source and destination types are checked for @@ -1250,22 +1311,6 @@ BitcastBetweenNativeTypes(const LiteralBase& src_literal) { LOG(FATAL) << "Invalid bitcast between types of different sizes."; } -template -Literal ConvertToC64(const LiteralBase& src_literal) { - CHECK(ShapeUtil::IsArray(src_literal.shape())); - 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()); @@ -1295,9 +1340,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) @@ -1306,10 +1353,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; @@ -1322,7 +1374,7 @@ StatusOr ConvertIfDestTypeMatches(const LiteralBase& src_literal, StatusOr ConvertSwitch(const LiteralBase& literal, PrimitiveType primitive_dest_type, bool bitcast) { - TF_RET_CHECK(ShapeUtil::IsArray(literal.shape())); + TF_RET_CHECK(literal.shape().IsArray()); if (literal.shape().element_type() == primitive_dest_type) { return literal.Clone(); } @@ -1333,9 +1385,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) @@ -1375,7 +1429,7 @@ StatusOr LiteralBase::BitcastConvert( } StatusOr LiteralBase::ConvertToShape(const Shape& dest_shape) const { - if (!ShapeUtil::IsTuple(dest_shape)) { + if (!dest_shape.IsTuple()) { return Convert(dest_shape.element_type()); } std::vector elements; @@ -1407,7 +1461,7 @@ StatusOr LiteralBase::ConvertToShape(const Shape& dest_shape) const { template bool LiteralBase::Piece::EqualElementsInternal( const LiteralBase::Piece& other, std::vector* multi_index) const { - if (multi_index->size() == ShapeUtil::Rank(subshape())) { + if (multi_index->size() == subshape().rank()) { return (Get(*multi_index) == other.Get(*multi_index)); } for (int64 i = 0; i < subshape().dimensions(multi_index->size()); ++i) { @@ -1435,10 +1489,14 @@ bool LiteralBase::Piece::EqualElements(const LiteralBase::Piece& other) const { return EqualElementsInternal(other, &multi_index); case U8: return EqualElementsInternal(other, &multi_index); + case S16: + return EqualElementsInternal(other, &multi_index); case S32: return EqualElementsInternal(other, &multi_index); case S64: return EqualElementsInternal(other, &multi_index); + case U16: + return EqualElementsInternal(other, &multi_index); case U32: return EqualElementsInternal(other, &multi_index); case U64: @@ -1453,6 +1511,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()); @@ -1466,7 +1526,7 @@ bool LiteralBase::operator==(const LiteralBase& other) const { return root_piece().ForEachSubpieceWithBool( [&](const ShapeIndex& index, const Piece& piece) { - if (!ShapeUtil::IsArray(piece.subshape())) { + if (!piece.subshape().IsArray()) { return true; } @@ -1496,7 +1556,7 @@ static bool AllElementsEqualValue(absl::Span data, bool LiteralBase::IsAll(int8 value) const { return root_piece().ForEachSubpieceWithBool([&](const ShapeIndex& index, const Piece& piece) { - if (!ShapeUtil::IsArray(piece.subshape())) { + if (!piece.subshape().IsArray()) { return true; } @@ -1507,6 +1567,11 @@ bool LiteralBase::IsAll(int8 value) const { return AllElementsEqualValue(piece.data(), value); } return false; + case U16: + if (value >= 0) { + return AllElementsEqualValue(piece.data(), value); + } + return false; case U32: if (value >= 0) { return AllElementsEqualValue(piece.data(), value); @@ -1519,6 +1584,8 @@ bool LiteralBase::IsAll(int8 value) const { return false; case S8: return AllElementsEqualValue(piece.data(), value); + case S16: + return AllElementsEqualValue(piece.data(), value); case S32: return AllElementsEqualValue(piece.data(), value); case S64: @@ -1557,7 +1624,7 @@ bool LiteralBase::IsAll(int8 value) const { bool LiteralBase::IsAllFloat(float value) const { return root_piece().ForEachSubpieceWithBool( [&](const ShapeIndex& index, const Piece& piece) { - if (!ShapeUtil::IsArray(piece.subshape())) { + if (!piece.subshape().IsArray()) { return true; } @@ -1589,6 +1656,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; } @@ -1597,7 +1667,7 @@ bool LiteralBase::IsAllComplex(complex64 value) const { bool LiteralBase::IsAllFirst() const { return root_piece().ForEachSubpieceWithBool( [&](const ShapeIndex& index, const Piece& piece) { - if (!ShapeUtil::IsArray(piece.subshape())) { + if (!piece.subshape().IsArray()) { return true; } @@ -1668,6 +1738,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; } @@ -1681,11 +1756,11 @@ bool LiteralBase::IsAllFirst() const { } bool LiteralBase::IsR1Iota() const { - if (!ShapeUtil::IsArray(shape())) { + if (!shape().IsArray()) { return false; } - if (ShapeUtil::Rank(shape()) != 1) { + if (shape().rank() != 1) { return false; } @@ -1717,6 +1792,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. @@ -1736,16 +1813,20 @@ bool LiteralBase::IsR1Iota() const { } bool LiteralBase::IsZero(absl::Span indices) const { - CHECK(ShapeUtil::IsArray(shape())); + CHECK(shape().IsArray()); switch (shape().element_type()) { case U8: return Get(indices) == 0; + case U16: + return Get(indices) == 0; case U32: return Get(indices) == 0; case U64: return Get(indices) == 0; case S8: return Get(indices) == 0; + case S16: + return Get(indices) == 0; case S32: return Get(indices) == 0; case S64: @@ -1756,6 +1837,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: @@ -1778,7 +1861,7 @@ void CopyToRepeatedField(RepeatedFieldT* dest, } // namespace void LiteralBase::Piece::WriteToProto(LiteralProto* proto) const { - *proto->mutable_shape() = subshape(); + *proto->mutable_shape() = subshape().ToProto(); switch (subshape().element_type()) { case PRED: CopyToRepeatedField(proto->mutable_preds(), data()); @@ -1803,6 +1886,20 @@ void LiteralBase::Piece::WriteToProto(LiteralProto* proto) const { case S64: CopyToRepeatedField(proto->mutable_s64s(), data()); break; + case U16: + *proto->mutable_u16s() = string( + reinterpret_cast(data().data()), size_bytes()); + if (!kLittleEndian) { + ConvertEndianShort(proto->mutable_u16s()); + } + break; + case S16: + *proto->mutable_s16s() = string( + reinterpret_cast(data().data()), size_bytes()); + if (!kLittleEndian) { + ConvertEndianShort(proto->mutable_s16s()); + } + break; case F16: *proto->mutable_f16s() = string( reinterpret_cast(data().data()), size_bytes()); @@ -1829,6 +1926,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. @@ -1841,12 +1944,12 @@ void LiteralBase::Piece::WriteToProto(LiteralProto* proto) const { } const void* LiteralBase::Piece::untyped_data() const { - CHECK(ShapeUtil::IsArray(subshape())) << ShapeUtil::HumanString(subshape()); + CHECK(subshape().IsArray()) << ShapeUtil::HumanString(subshape()); return buffer(); } void* LiteralBase::Piece::untyped_data() { - CHECK(ShapeUtil::IsArray(subshape())) << ShapeUtil::HumanString(subshape()); + CHECK(subshape().IsArray()) << ShapeUtil::HumanString(subshape()); return buffer(); } @@ -1870,20 +1973,19 @@ Status LiteralBase::Piece::CopyFromProto(const LiteralProto& proto) { // These conditions should have been checked in // MutableLiteralBase::CreateFromProto. TF_RET_CHECK(proto.has_shape()); - TF_RET_CHECK(LayoutUtil::HasLayout(proto.shape())); - TF_RET_CHECK(ShapeUtil::Equal(proto.shape(), subshape())); + Shape shape(proto.shape()); + TF_RET_CHECK(LayoutUtil::HasLayout(shape)); + TF_RET_CHECK(ShapeUtil::Equal(shape, subshape())); if (LayoutUtil::IsSparseArray(subshape())) { // Compute the number of elements (indices) in the sparse shape and reserve // the necessary space in spare_indices. - TF_RET_CHECK(ShapeUtil::Rank(subshape()) != 0) - << "Scalar shapes cannot be sparse"; - TF_RET_CHECK(proto.sparse_indices_size() % ShapeUtil::Rank(subshape()) == 0) + TF_RET_CHECK(subshape().rank() != 0) << "Scalar shapes cannot be sparse"; + TF_RET_CHECK(proto.sparse_indices_size() % subshape().rank() == 0) << "Unexpected number of indices in proto (" << proto.sparse_indices_size() << ") for shape of rank " - << ShapeUtil::Rank(subshape()); - const int64 index_count = - proto.sparse_indices_size() / ShapeUtil::Rank(subshape()); + << subshape().rank(); + const int64 index_count = proto.sparse_indices_size() / subshape().rank(); sparse_indices()->Resize(index_count); // Copy the indices from the proto into the SparseIndexArray object. @@ -1917,6 +2019,22 @@ Status LiteralBase::Piece::CopyFromProto(const LiteralProto& proto) { case U64: TF_RETURN_IF_ERROR(CopyFromRepeatedField(data(), proto.u64s())); break; + case S16: { + const string& s(proto.s16s()); + TF_RET_CHECK(data().size() * sizeof(int16_t) == s.size()); + memcpy(untyped_data(), s.data(), s.size()); + if (!kLittleEndian) { + ConvertEndianShort(reinterpret_cast(untyped_data()), s.size()); + } + } break; + case U16: { + const string& s(proto.u16s()); + TF_RET_CHECK(data().size() * sizeof(uint16_t) == s.size()); + memcpy(untyped_data(), s.data(), s.size()); + if (!kLittleEndian) { + ConvertEndianShort(reinterpret_cast(untyped_data()), s.size()); + } + } break; case F16: { const string& s(proto.f16s()); TF_RET_CHECK(data().size() * sizeof(half) == s.size()); @@ -1946,7 +2064,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())); @@ -1992,10 +2120,10 @@ int64 LiteralBase::size_bytes(const ShapeIndex& shape_index) const { } string LiteralBase::GetR1U8AsString() const { - CHECK(ShapeUtil::IsArray(shape())); - CHECK_EQ(ShapeUtil::Rank(shape()), 1); + CHECK(shape().IsArray()); + CHECK_EQ(shape().rank(), 1); CHECK_EQ(shape().element_type(), U8); - return string(tensorflow::bit_cast(data().data()), + return string(absl::bit_cast(data().data()), ShapeUtil::ElementsIn(shape())); } @@ -2007,7 +2135,7 @@ void MutableBorrowingLiteral::CopyPieceSubtree(const Shape& shape, << ShapeUtil::HumanString(src_piece->subshape()) << "dest_piece has shape: " << ShapeUtil::HumanString(dest_piece->subshape()); - if (ShapeUtil::IsTuple(shape)) { + if (shape.IsTuple()) { for (int i = 0; i < ShapeUtil::TupleElementCount(shape); ++i) { const Shape& subshape = shape.tuple_shapes(i); @@ -2018,7 +2146,7 @@ void MutableBorrowingLiteral::CopyPieceSubtree(const Shape& shape, dest_piece->emplace_back(std::move(child_piece)); } - } else if (ShapeUtil::IsArray(shape)) { + } else if (shape.IsArray()) { dest_piece->set_buffer(src_piece->buffer()); } else { // If the shape is neither an array nor tuple, then it must be @@ -2094,7 +2222,7 @@ MutableBorrowingLiteral::MutableBorrowingLiteral(const char* src_buf_ptr, : MutableLiteralBase() { shape_ = absl::make_unique(shape); CHECK(LayoutUtil::HasLayout(*shape_)); - CHECK(!ShapeUtil::IsTuple(*shape_)); + CHECK(!shape_->IsTuple()); root_piece_ = new Piece(); root_piece_->set_buffer(const_cast(src_buf_ptr)); @@ -2121,14 +2249,14 @@ LiteralSlice::LiteralSlice(const LiteralBase& literal, : LiteralBase(), root_piece_(&literal.piece(view_root)) {} void BorrowingLiteral::BuildPieceSubtree(const Shape& shape, Piece* piece) { - CHECK(ShapeUtil::IsTuple(shape)); + CHECK(shape.IsTuple()); for (int i = 0; i < ShapeUtil::TupleElementCount(shape); ++i) { const Shape& subshape = shape.tuple_shapes(i); auto child_piece = Piece(); child_piece.set_subshape(&subshape); - if (ShapeUtil::IsTuple(subshape)) { + if (subshape.IsTuple()) { BuildPieceSubtree(subshape, &child_piece); } @@ -2138,7 +2266,7 @@ void BorrowingLiteral::BuildPieceSubtree(const Shape& shape, Piece* piece) { BorrowingLiteral::BorrowingLiteral(const char* src_buf_ptr, const Shape& shape) : LiteralBase(), shape_(absl::make_unique(shape)) { - CHECK(ShapeUtil::IsArray(*shape_)); + CHECK(shape_->IsArray()); CHECK(LayoutUtil::HasLayout(*shape_)); root_piece_ = Piece(); @@ -2149,7 +2277,7 @@ BorrowingLiteral::BorrowingLiteral(const char* src_buf_ptr, const Shape& shape) BorrowingLiteral::BorrowingLiteral(absl::Span src_buf_ptrs, const Shape& shape) : LiteralBase(), shape_(absl::make_unique(shape)) { - CHECK(ShapeUtil::IsTuple(*shape_)); + CHECK(shape_->IsTuple()); CHECK(!ShapeUtil::IsNestedTuple(*shape_)); CHECK_EQ(src_buf_ptrs.size(), ShapeUtil::TupleElementCount(*shape_)); root_piece_ = Piece(); @@ -2158,7 +2286,7 @@ BorrowingLiteral::BorrowingLiteral(absl::Span src_buf_ptrs, for (int i = 0; i < src_buf_ptrs.size(); ++i) { const auto& src_shape = shape_->tuple_shapes(i); - CHECK(ShapeUtil::IsArray(src_shape)); + CHECK(src_shape.IsArray()); root_piece_.child(i).set_buffer(const_cast(src_buf_ptrs[i])); } } diff --git a/tensorflow/compiler/xla/literal.h b/tensorflow/compiler/xla/literal.h index e791048b4d9f5dcf877e05e3b5cf16eb37c07dbc..041151fda1280d6ae7b35d5857ca79788d4f7203 100644 --- a/tensorflow/compiler/xla/literal.h +++ b/tensorflow/compiler/xla/literal.h @@ -92,9 +92,20 @@ class LiteralBase { // array. string GetR1U8AsString() const; - // Returns a string representation of the literal value. - // Warning: this function can take minutes for multi-million element Literals. - string ToString(bool print_layout = false) const; + // Returns a string representation of the literal value. The Shape of the + // literal is a prefix of the literal value in the string. + + // Warning: this function can take minutes for multi-million + // element Literals. + string ToString() const; + + // Returns a string representation of the literal value which does *not* + // include the shape string. + string ToStringWithoutShape() const; + + // Returns a string representation of the literal value which includes the + // shape string with its layout.does *not* include the shape string. + string ToStringWithLayout() const; // Gets an element in the literal at the given index. The multi_index is // CHECKed against the dimension sizes. @@ -301,7 +312,7 @@ class LiteralBase { // // Note: It's an antipattern to use this method then immediately call // MutableLiteralBase::Populate on the result (since that results in zero - // initialization, then reinitialization. Conside if a call to + // initialization, then reinitialization. Consider if a call to // absl::make_unique(shape), followed by the call to // MutableLiteralBase::Populate can be used instead. static Literal CreateFromShape(const Shape& shape); @@ -856,7 +867,7 @@ class BorrowingLiteral : public LiteralBase { template absl::Span LiteralBase::Piece::data() const { - DCHECK(ShapeUtil::IsArray(subshape())) << ShapeUtil::HumanString(subshape()); + DCHECK(subshape().IsArray()) << ShapeUtil::HumanString(subshape()); DCHECK_EQ(subshape().element_type(), primitive_util::NativeToPrimitiveType()) << "Attempting to access " @@ -869,7 +880,7 @@ absl::Span LiteralBase::Piece::data() const { template absl::Span LiteralBase::Piece::data() { - DCHECK(ShapeUtil::IsArray(subshape())) << ShapeUtil::HumanString(subshape()); + DCHECK(subshape().IsArray()) << ShapeUtil::HumanString(subshape()); DCHECK_EQ(subshape().element_type(), primitive_util::NativeToPrimitiveType()) << "Attempting to access " @@ -950,7 +961,7 @@ void MutableLiteralBase::AppendSparseElement( Piece& p = piece(shape_index); const Shape& subshape = p.subshape(); CHECK(LayoutUtil::IsSparseArray(subshape)); - int64 rank = ShapeUtil::Rank(subshape); + int64 rank = subshape.rank(); CHECK_EQ(multi_index.size(), rank); int64 last_element = p.sparse_indices()->index_count(); CHECK_LT(last_element, LayoutUtil::MaxSparseElements(subshape.layout())); @@ -966,7 +977,7 @@ void LiteralBase::EachCell( if (ShapeUtil::IsZeroElementArray(shape())) { return; } - std::vector indices(ShapeUtil::Rank(shape()), 0); + std::vector indices(shape().rank(), 0); do { per_cell(indices, Get(indices)); } while (IndexUtil::BumpIndices(shape(), absl::MakeSpan(indices))); @@ -974,8 +985,8 @@ void LiteralBase::EachCell( template inline void MutableLiteralBase::PopulateR1(absl::Span values) { - CHECK(ShapeUtil::IsArray(shape())); - CHECK_EQ(ShapeUtil::Rank(shape()), 1); + CHECK(shape().IsArray()); + CHECK_EQ(shape().rank(), 1); CHECK_EQ(ShapeUtil::ElementsIn(shape()), values.size()); CHECK_EQ(shape().element_type(), primitive_util::NativeToPrimitiveType()); @@ -986,8 +997,8 @@ inline void MutableLiteralBase::PopulateR1(absl::Span values) { template void MutableLiteralBase::PopulateR2( std::initializer_list> values) { - CHECK(ShapeUtil::IsArray(shape())); - CHECK_EQ(ShapeUtil::Rank(shape()), 2); + CHECK(shape().IsArray()); + CHECK_EQ(shape().rank(), 2); CHECK_EQ(shape().element_type(), primitive_util::NativeToPrimitiveType()); @@ -1010,10 +1021,10 @@ void MutableLiteralBase::PopulateR2( template void MutableLiteralBase::PopulateFromArray(const Array& values) { - CHECK(ShapeUtil::IsArray(shape())); + CHECK(shape().IsArray()); CHECK_EQ(shape().element_type(), primitive_util::NativeToPrimitiveType()); - CHECK_EQ(ShapeUtil::Rank(shape()), values.num_dimensions()); + CHECK_EQ(shape().rank(), values.num_dimensions()); for (int dim = 0; dim < values.num_dimensions(); ++dim) { CHECK_EQ(values.dim(dim), shape().dimensions(dim)); } @@ -1042,7 +1053,7 @@ void MutableLiteralBase::PopulateSparse(SparseIndexArray indices, absl::Span values, bool sort) { CHECK(LayoutUtil::IsSparseArray(shape())); - int rank = ShapeUtil::Rank(shape()); + int rank = shape().rank(); CHECK_EQ(indices.rank(), rank); int64 max_elements = LayoutUtil::MaxSparseElements(shape().layout()); CHECK_LE(indices.max_indices(), max_elements); @@ -1066,7 +1077,7 @@ template Status MutableLiteralBase::PopulateInternal(const FnType& generator, bool parallel) { const Shape& this_shape = shape(); - const int64 rank = ShapeUtil::Rank(this_shape); + const int64 rank = this_shape.rank(); TF_RET_CHECK(LayoutUtil::IsDenseArray(this_shape)); TF_RET_CHECK(this_shape.element_type() == primitive_util::NativeToPrimitiveType()); @@ -1118,7 +1129,7 @@ Status MutableLiteralBase::PopulateParallel(const FnType& generator) { template void MutableLiteralBase::PopulateWithValue(NativeT value) { - CHECK(ShapeUtil::IsArray(shape())); + CHECK(shape().IsArray()); CHECK_EQ(shape().element_type(), primitive_util::NativeToPrimitiveType()); for (NativeT& element : data()) { diff --git a/tensorflow/compiler/xla/literal_comparison.cc b/tensorflow/compiler/xla/literal_comparison.cc index 3d8725ed7051cafc97987f25a96004fa876dfdd3..9b3de75dd4e9d495778af86fb8fc07909ab4ba81 100644 --- a/tensorflow/compiler/xla/literal_comparison.cc +++ b/tensorflow/compiler/xla/literal_comparison.cc @@ -19,11 +19,11 @@ limitations under the License. #include #include +#include "absl/base/casts.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_format.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/util.h" -#include "tensorflow/core/lib/core/casts.h" #include "tensorflow/core/platform/env.h" using absl::StrAppend; @@ -34,72 +34,128 @@ namespace xla { namespace literal_comparison { namespace { +// Since Eigen::half doesn't satisfy the absl::bit_cast contract, we need to be +// able to transparently access the raw 16-bit value contained within. +template +T GetRawValue(T val) { + return val; +} +uint16 GetRawValue(Eigen::half val) { return val.x; } + // Helper function for comparing a floating point type, FloatT, bitwise equal // between the left-hand-side and right-hand-side, by bit-casting to UnsignedT // -- on miscompare, a nice error message is given in the AssertionFailure. template -Status CompareFloatsBitwiseEqual(FloatT lhs, FloatT rhs, - absl::Span multi_index) { - auto ulhs = tensorflow::bit_cast(lhs); - auto urhs = tensorflow::bit_cast(rhs); +bool CompareFloatsBitwiseEqual(FloatT lhs, FloatT rhs, + absl::Span multi_index) { + auto ulhs = absl::bit_cast(GetRawValue(lhs)); + auto urhs = absl::bit_cast(GetRawValue(rhs)); + return ulhs == urhs; +} + +// Templated comparator that specializes for float equality comparison with the +// bitwise helper above (this is the un-specialized fallback, to just use the +// default gunit implementation). +template +bool CompareEqual(NativeT lhs, NativeT rhs, + absl::Span multi_index) { + return lhs == rhs; +} + +// Specializations for floating types that do bitwise comparisons when equality +// comparison is requested. +template <> +bool CompareEqual(bfloat16 lhs, bfloat16 rhs, + absl::Span multi_index) { + return CompareFloatsBitwiseEqual(lhs, rhs, multi_index); +} +template <> +bool CompareEqual(Eigen::half lhs, Eigen::half rhs, + absl::Span multi_index) { + return CompareFloatsBitwiseEqual(lhs, rhs, multi_index); +} +template <> +bool CompareEqual(float lhs, float rhs, + absl::Span multi_index) { + return CompareFloatsBitwiseEqual(lhs, rhs, multi_index); +} +template <> +bool CompareEqual(double lhs, double rhs, + absl::Span multi_index) { + return CompareFloatsBitwiseEqual(lhs, rhs, multi_index); +} +template <> +bool CompareEqual(complex64 lhs, complex64 rhs, + absl::Span multi_index) { + 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, + absl::Span multi_index) { + auto ulhs = absl::bit_cast(GetRawValue(lhs)); + auto urhs = absl::bit_cast(GetRawValue(rhs)); auto lhs_double = static_cast(lhs); auto rhs_double = static_cast(rhs); - if (ulhs != urhs) { return InvalidArgument( "floating values are not bitwise-equal; and equality testing " "was requested: %s=%g=%a vs %s=%g=%a at array index %s", StrCat(absl::Hex(ulhs)), lhs_double, lhs_double, StrCat(absl::Hex(urhs)), rhs_double, rhs_double, LiteralUtil::MultiIndexAsString(multi_index)); - } - return Status::OK(); } -// Templated comparator that specializes for float equality comparison with the -// bitwise helper above (this is the un-specialized fallback, to just use the -// default gunit implementation). template -Status CompareEqual(NativeT lhs, NativeT rhs, - absl::Span multi_index) { - if (lhs == rhs) { - return Status::OK(); - } +Status MakeErrorStatus(NativeT lhs, NativeT rhs, + absl::Span multi_index) { return InvalidArgument( "first mismatch at array index %s:\n expected value: %s\n actual " "value: %s", LiteralUtil::MultiIndexAsString(multi_index), StrCat(lhs), StrCat(rhs)); } -// Specializations for floating types that do bitwise comparisons when equality -// comparison is requested. template <> -Status CompareEqual(bfloat16 lhs, bfloat16 rhs, - absl::Span multi_index) { - return CompareFloatsBitwiseEqual(lhs, rhs, multi_index); +Status MakeErrorStatus(bfloat16 lhs, bfloat16 rhs, + absl::Span multi_index) { + return MakeBitwiseErrorStatus(lhs, rhs, multi_index); } template <> -Status CompareEqual(Eigen::half lhs, Eigen::half rhs, - absl::Span multi_index) { - return CompareFloatsBitwiseEqual(lhs, rhs, multi_index); +Status MakeErrorStatus(Eigen::half lhs, Eigen::half rhs, + absl::Span multi_index) { + return MakeBitwiseErrorStatus(lhs, rhs, multi_index); } template <> -Status CompareEqual(float lhs, float rhs, - absl::Span multi_index) { - return CompareFloatsBitwiseEqual(lhs, rhs, multi_index); +Status MakeErrorStatus(float lhs, float rhs, + absl::Span multi_index) { + return MakeBitwiseErrorStatus(lhs, rhs, multi_index); } template <> -Status CompareEqual(double lhs, double rhs, - absl::Span multi_index) { - return CompareFloatsBitwiseEqual(lhs, rhs, multi_index); +Status MakeErrorStatus(double lhs, double rhs, + absl::Span multi_index) { + return MakeBitwiseErrorStatus(lhs, rhs, multi_index); } template <> -Status CompareEqual(complex64 lhs, complex64 rhs, - absl::Span multi_index) { - auto res = CompareEqual(lhs.real(), rhs.real(), multi_index); - if (!res.ok()) { - return res; +Status MakeErrorStatus(complex64 lhs, complex64 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); +} +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 CompareEqual(lhs.imag(), rhs.imag(), multi_index); + return MakeErrorStatus(lhs.imag(), rhs.imag(), multi_index); } // A recursive function which iterates through every index of expected and @@ -111,7 +167,11 @@ Status Equal(LiteralSlice expected, LiteralSlice actual, if (dimension == expected.shape().dimensions_size()) { NativeT expected_value = expected.Get(multi_index); NativeT actual_value = actual.Get(multi_index); - return CompareEqual(expected_value, actual_value, multi_index); + bool result = + CompareEqual(expected_value, actual_value, multi_index); + return result ? Status::OK() + : MakeErrorStatus(expected_value, actual_value, + multi_index); } Status result; @@ -126,42 +186,20 @@ Status Equal(LiteralSlice expected, LiteralSlice actual, // Gets the total element count. For tuples, this is not the count of tuple // elements, but the sum of elements of each tuple element. int64 RecursiveElementCount(const Shape& shape) { - if (ShapeUtil::IsTuple(shape)) { + if (shape.IsTuple()) { const int64 tuple_elements = ShapeUtil::TupleElementCount(shape); int64 total = 0; for (int64 i = 0; i < tuple_elements; ++i) { total += RecursiveElementCount(ShapeUtil::GetTupleElementShape(shape, i)); } return total; - } else { + } else if (shape.IsArray()) { return ShapeUtil::ElementsIn(shape); - } -} - -// Returns whether the actual and expected values are mismatched with respect to -// nans. 'relaxed_nans' is interpreted as in xla::ErrorSpec. -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); + return 0; } } -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); -} - -template <> -bool NanMismatch(half expected, half actual, bool relaxed_nans) { - return NanMismatch(static_cast(expected), - static_cast(actual), relaxed_nans); -} - // Returns whether the given value is infinity. template bool IsInf(NativeT val) { @@ -173,6 +211,17 @@ bool IsInf(half val) { return std::isinf(static_cast(val)); } +// Returns whether the given value is nan. +template +float IsNan(NativeT value) { + return std::isnan(value); +} + +template <> +float IsNan(half value) { + return IsNan(static_cast(value)); +} + // Converts the given floating-point value to a string. template string FpValueToString(NativeT value) { @@ -184,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. @@ -263,7 +317,7 @@ class NearComparator { // If the shapes mismatch, we simply fail the expectation instead of // printing out data, as it's a type error rather than a value error. TF_RETURN_IF_ERROR(EqualShapes(expected_.shape(), actual_.shape())); - if (!ShapeUtil::IsArray(expected_.shape())) { + if (!expected_.shape().IsArray()) { return InvalidArgument("Expected array shape; got %s.", ShapeUtil::HumanString(expected_.shape())); } @@ -316,35 +370,59 @@ 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}).ok()) { + 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. - CHECK(!CompareEqual(expected, actual, {linear_index}).ok()); + CHECK(!CompareEqual(expected, actual, {linear_index})); abs_error = std::numeric_limits::infinity(); rel_error = std::numeric_limits::infinity(); } else { abs_error = FpAbsoluteValue(actual - expected); - rel_error = abs_error / FpAbsoluteValue(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(); + } } 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. @@ -379,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); @@ -402,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. @@ -415,7 +516,7 @@ class NearComparator { } return; } - std::vector multi_index(ShapeUtil::Rank(actual_.shape()), 0); + std::vector multi_index(actual_.shape().rank(), 0); CompareLiteralsSlow(0, &multi_index); } @@ -610,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}), @@ -632,12 +736,12 @@ 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())); - if (ShapeUtil::IsTuple(expected.shape())) { + if (expected.shape().IsTuple()) { Status return_status; for (int64 i = 0; i < ShapeUtil::TupleElementCount(expected.shape()); ++i) { const auto expected_element = LiteralSlice(expected, {i}); @@ -673,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: " @@ -713,7 +823,7 @@ Status EqualShapes(const Shape& expected, const Shape& actual) { ShapeUtil::HumanString(expected), ShapeUtil::HumanString(actual)); } - if (ShapeUtil::IsTuple(expected)) { + if (expected.IsTuple()) { if (ShapeUtil::TupleElementCount(expected) != ShapeUtil::TupleElementCount(actual)) { return InvalidArgument( @@ -728,8 +838,8 @@ Status EqualShapes(const Shape& expected, const Shape& actual) { return AppendStatus(result, StrCat("mismatch in tuple index", i)); } } - } else if (ShapeUtil::IsArray(expected)) { - if (ShapeUtil::Rank(expected) != ShapeUtil::Rank(actual)) { + } else if (expected.IsArray()) { + if (expected.rank() != actual.rank()) { return InvalidArgument("want rank of %s got rank of %s", ShapeUtil::HumanString(expected), ShapeUtil::HumanString(actual)); @@ -783,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 dd5b54e4c99998f676419cf98a3da16593338829..b54a71ae68218ef578535a913f5867d843236e32 100644 --- a/tensorflow/compiler/xla/literal_test.cc +++ b/tensorflow/compiler/xla/literal_test.cc @@ -17,6 +17,7 @@ limitations under the License. #include +#include "absl/base/casts.h" #include "absl/memory/memory.h" #include "absl/strings/match.h" #include "absl/strings/str_cat.h" @@ -28,7 +29,6 @@ limitations under the License. #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/test.h" #include "tensorflow/compiler/xla/types.h" -#include "tensorflow/core/lib/core/casts.h" #include "tensorflow/core/lib/core/status_test_util.h" #include "tensorflow/core/platform/macros.h" #include "tensorflow/core/platform/types.h" @@ -98,42 +98,45 @@ class LiteralUtilTest : public ::testing::Test { TEST_F(LiteralUtilTest, LiteralScalarToString) { auto true_lit = LiteralUtil::CreateR0(true); - EXPECT_EQ("true", true_lit.ToString()); + EXPECT_EQ("pred[] true", true_lit.ToString()); auto false_lit = LiteralUtil::CreateR0(false); - EXPECT_EQ("false", false_lit.ToString()); + EXPECT_EQ("pred[] false", false_lit.ToString()); auto u32_lit = LiteralUtil::CreateR0(42); - EXPECT_EQ("42", u32_lit.ToString()); + EXPECT_EQ("u32[] 42", u32_lit.ToString()); auto s32_lit = LiteralUtil::CreateR0(-999); - EXPECT_EQ("-999", s32_lit.ToString()); + EXPECT_EQ("s32[] -999", s32_lit.ToString()); auto f32_lit = LiteralUtil::CreateR0(3.14f); - EXPECT_EQ("3.14", f32_lit.ToString()); + EXPECT_EQ("f32[] 3.14", f32_lit.ToString()); auto f16_lit = LiteralUtil::CreateR0(static_cast(0.5f)); - EXPECT_EQ("0.5", f16_lit.ToString()); + EXPECT_EQ("f16[] 0.5", f16_lit.ToString()); auto c64_lit = LiteralUtil::CreateR0({3.14f, 2.78f}); - EXPECT_EQ("(3.14, 2.78)", c64_lit.ToString()); + 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("0.5", bf16_lit.ToString()); + EXPECT_EQ("bf16[] 0.5", bf16_lit.ToString()); // 3.14 will be rounded to 3.14062 in bfloat16 format. auto bf16_lit_truncated = LiteralUtil::CreateR0(static_cast(3.14f)); - ASSERT_EQ("3.14062", bf16_lit_truncated.ToString()); + ASSERT_EQ("bf16[] 3.14062", bf16_lit_truncated.ToString()); auto bf16_lit_truncated2 = LiteralUtil::CreateR0(static_cast(9.001f)); - EXPECT_EQ("9", bf16_lit_truncated2.ToString()); + EXPECT_EQ("bf16[] 9", bf16_lit_truncated2.ToString()); } TEST_F(LiteralUtilTest, LiteralVectorToString) { auto pred_vec = LiteralUtil::CreateR1({true, false, true}); - EXPECT_EQ("{101}", pred_vec.ToString()); + EXPECT_EQ("pred[3] {1, 0, 1}", pred_vec.ToString()); } TEST_F(LiteralUtilTest, R2ToString) { @@ -150,12 +153,58 @@ TEST_F(LiteralUtilTest, R3ToString) { const auto literal = LiteralUtil::CreateR3({{{1}, {2}}, {{3}, {4}}, {{5}, {6}}}); const string expected = R"(s32[3,2,1] { -{ { 1 }, - { 2 } }, -{ { 3 }, - { 4 } }, -{ { 5 }, - { 6 } } +{ + {1}, + {2} +}, +{ + {3}, + {4} +}, +{ + {5}, + {6} +} +})"; + EXPECT_EQ(expected, literal.ToString()); +} + +TEST_F(LiteralUtilTest, R6ToString) { + const auto literal = + LiteralUtil::CreateFromDimensions(S32, {2, 2, 1, 1, 1, 2}); + const string expected = R"(s32[2,2,1,1,1,2] { +{ /*i0=0*/ +{ /*i1=0*/ +{ /*i2=0*/ +{ /*i3=0*/ + { 0, 0 } +} +} +}, +{ /*i1=1*/ +{ /*i2=0*/ +{ /*i3=0*/ + { 0, 0 } +} +} +} +}, +{ /*i0=1*/ +{ /*i1=0*/ +{ /*i2=0*/ +{ /*i3=0*/ + { 0, 0 } +} +} +}, +{ /*i1=1*/ +{ /*i2=0*/ +{ /*i3=0*/ + { 0, 0 } +} +} +} +} })"; EXPECT_EQ(expected, literal.ToString()); } @@ -164,8 +213,8 @@ TEST_F(LiteralUtilTest, TupleToString) { auto scalar = LiteralUtil::CreateR0(1.0); auto matrix = LiteralUtil::CreateR2({{1.0, 2.0}, {3.0, 4.0}}); auto tuple = LiteralUtil::MakeTuple({&scalar, &matrix}); - const string expected = R"((f32[], f32[2,2]) ( -1, + const string expected = R"(( +f32[] 1, f32[2,2] { { 1, 2 }, { 3, 4 } @@ -190,12 +239,16 @@ TEST_F(LiteralUtilTest, CreateR3FromArray3d) { EXPECT_THAT(literal.shape().dimensions(), ElementsAre(2, 3, 2)); string result = literal.ToString(); const string expected = R"(f32[2,3,2] { -{ { 1, 2 }, +{ + { 1, 2 }, { 3, 4 }, - { 5, 6 } }, -{ { 7, 8 }, + { 5, 6 } +}, +{ + { 7, 8 }, { 9, 10 }, - { 11, 12 } } + { 11, 12 } +} })"; EXPECT_EQ(expected, result); } @@ -247,18 +300,18 @@ TEST_F(LiteralUtilTest, LiteralR4F32ProjectedStringifies) { EXPECT_THAT(literal.shape().dimensions(), ElementsAre(1, 2, 3, 2)); string result = literal.ToString(); const string expected = R"(f32[1,2,3,2] { - { /*i0=0*/ - { /*i1=0*/ - {1, 2}, - {1001, 1002}, - {2001, 2002} - }, - { /*i1=1*/ - {1, 2}, - {1001, 1002}, - {2001, 2002} - } - } +{ /*i0=0*/ +{ /*i1=0*/ + { 1, 2 }, + { 1001, 1002 }, + { 2001, 2002 } +}, +{ /*i1=1*/ + { 1, 2 }, + { 1001, 1002 }, + { 2001, 2002 } +} +} })"; EXPECT_EQ(expected, result); } @@ -268,30 +321,30 @@ TEST_F(LiteralUtilTest, LiteralR4F32Stringifies) { ElementsAre(2, 2, 3, 3)); string result = literal_r4_2x2x3x3_dim0major_.ToString(); const string expected = R"(f32[2,2,3,3] { - { /*i0=0*/ - { /*i1=0*/ - {1, 2, 3}, - {4, 5, 6}, - {7, 8, 9} - }, - { /*i1=1*/ - {11, 12, 13}, - {14, 15, 16}, - {17, 18, 19} - } - }, - { /*i0=1*/ - { /*i1=0*/ - {101, 102, 103}, - {104, 105, 106}, - {107, 108, 109} - }, - { /*i1=1*/ - {201, 202, 203}, - {204, 205, 206}, - {207, 208, 209} - } - } +{ /*i0=0*/ +{ /*i1=0*/ + { 1, 2, 3 }, + { 4, 5, 6 }, + { 7, 8, 9 } +}, +{ /*i1=1*/ + { 11, 12, 13 }, + { 14, 15, 16 }, + { 17, 18, 19 } +} +}, +{ /*i0=1*/ +{ /*i1=0*/ + { 101, 102, 103 }, + { 104, 105, 106 }, + { 107, 108, 109 } +}, +{ /*i1=1*/ + { 201, 202, 203 }, + { 204, 205, 206 }, + { 207, 208, 209 } +} +} })"; EXPECT_EQ(expected, result); } @@ -419,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}}); @@ -573,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. @@ -786,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}}}); @@ -847,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); @@ -1187,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}}, @@ -1248,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); @@ -1302,21 +1405,34 @@ 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) { auto original = LiteralUtil::CreateR1( - {tensorflow::bit_cast(2.5f), - tensorflow::bit_cast(-42.25f), - tensorflow::bit_cast(100.f), 0xbeef}); + {absl::bit_cast(2.5f), absl::bit_cast(-42.25f), + absl::bit_cast(100.f), 0xbeef}); auto expected = LiteralUtil::CreateR1( - {2.5f, -42.25f, 100.0f, tensorflow::bit_cast(0xbeef)}); + {2.5f, -42.25f, 100.0f, absl::bit_cast(0xbeef)}); TF_ASSERT_OK_AND_ASSIGN(Literal converted, original.BitcastConvert(F32)); } @@ -1328,13 +1444,26 @@ TEST_F(LiteralUtilTest, BitcastConvertBetweenInvalidTypes) { absl::StrContains(status.error_message(), "bit widths are different")); } +// Sets the layout of the given ShapeProto to the default. +void SetDefaultLayoutOnProto(ShapeProto* shape_proto) { + CHECK(ShapeUtil::IsArrayPrimitiveType(shape_proto->element_type())); + shape_proto->mutable_layout()->set_format(DENSE); + auto* minor_to_major = + shape_proto->mutable_layout()->mutable_minor_to_major(); + minor_to_major->Resize(shape_proto->dimensions_size(), 0); + const int64 size = minor_to_major->size(); + for (int64 i = 0; i < size; ++i) { + minor_to_major->Set(i, size - 1 - i); + } +} + TEST_F(LiteralUtilTest, CopyFromProto_Bool) { LiteralProto p; p.mutable_shape()->set_element_type(PRED); for (int len = 0; len < 25; ++len) { p.mutable_shape()->clear_dimensions(); p.mutable_shape()->add_dimensions(len); - LayoutUtil::SetToDefaultLayout(p.mutable_shape()); + SetDefaultLayoutOnProto(p.mutable_shape()); p.clear_preds(); for (int i = 0; i < len; ++i) { p.add_preds((i % 2) == (len % 2)); @@ -1360,7 +1489,7 @@ TEST_F(LiteralUtilTest, ToProto_f16) { EXPECT_EQ(4, m.data().size()); LiteralProto p = m.ToProto(); - EXPECT_EQ(4, ShapeUtil::ElementsIn(p.shape())); + EXPECT_EQ(4, ShapeUtil::ElementsIn(Shape(p.shape()))); EXPECT_EQ(8, p.f16s().size()); const char* d = p.f16s().data(); EXPECT_EQ(d[0], 0); @@ -1383,7 +1512,7 @@ TEST_F(LiteralUtilTest, CopyFromProto_f16) { p.mutable_shape()->set_element_type(F16); p.mutable_shape()->clear_dimensions(); p.mutable_shape()->add_dimensions(4); - LayoutUtil::SetToDefaultLayout(p.mutable_shape()); + SetDefaultLayoutOnProto(p.mutable_shape()); p.clear_f16s(); p.set_f16s(half_vals, 8); TF_ASSERT_OK_AND_ASSIGN(Literal literal, Literal::CreateFromProto(p)); @@ -1395,6 +1524,28 @@ TEST_F(LiteralUtilTest, CopyFromProto_f16) { EXPECT_EQ(h1, r[3]); } +TEST_F(LiteralUtilTest, CopyFromProto_u16) { + uint16 u1(0xabcd); + uint16 u2(0x1234); + + const unsigned char uint16_vals[8] = {0xcd, 0xab, 0x34, 0x12, + 0x34, 0x12, 0xcd, 0xab}; + LiteralProto p; + p.mutable_shape()->set_element_type(U16); + p.mutable_shape()->clear_dimensions(); + p.mutable_shape()->add_dimensions(4); + SetDefaultLayoutOnProto(p.mutable_shape()); + p.clear_u16s(); + p.set_u16s(uint16_vals, 8); + TF_ASSERT_OK_AND_ASSIGN(Literal literal, Literal::CreateFromProto(p)); + auto r = literal.data(); + ASSERT_EQ(4, r.size()); + EXPECT_EQ(u1, r[0]); + EXPECT_EQ(u2, r[1]); + EXPECT_EQ(u2, r[2]); + EXPECT_EQ(u1, r[3]); +} + TEST_F(LiteralUtilTest, LiteralSliceTest) { auto scalar = LiteralUtil::CreateR0(1.0); auto matrix = LiteralUtil::CreateR2({{1.0, 2.0}, {3.0, 4.0}}); @@ -1516,9 +1667,9 @@ TEST_F(LiteralUtilTest, DecomposeTuple) { Literal nested_tuple = LiteralUtil::MakeTuple( {&tuple_elements[0], &tuple_elements[1], &nil_literal}); - EXPECT_FALSE(ShapeUtil::IsNil(nested_tuple.shape())); + EXPECT_FALSE(ShapeUtil::IsEmptyTuple(nested_tuple.shape())); std::vector elements = nested_tuple.DecomposeTuple(); - EXPECT_TRUE(ShapeUtil::IsNil(nested_tuple.shape())); + EXPECT_TRUE(ShapeUtil::IsEmptyTuple(nested_tuple.shape())); ASSERT_EQ(elements.size(), 3); @@ -1558,7 +1709,7 @@ TEST_F(LiteralUtilTest, MoveIntoTuple) { LiteralUtil::MakeTuple({&inner_elements[0], &inner_elements[1]})); Literal literal = Literal::MoveIntoTuple(absl::MakeSpan(elements)); - ASSERT_TRUE(ShapeUtil::IsTuple(literal.shape())); + ASSERT_TRUE(literal.shape().IsTuple()); ASSERT_EQ(ShapeUtil::TupleElementCount(literal.shape()), 3); EXPECT_EQ(literal.Get({}, /*shape_index=*/{0}), 1.0); @@ -1569,13 +1720,13 @@ TEST_F(LiteralUtilTest, MoveIntoTuple) { EXPECT_EQ(literal.Get({1}, /*shape_index=*/{2, 1}), 44.0); for (const Literal& element : elements) { - EXPECT_TRUE(ShapeUtil::IsNil(element.shape())); + EXPECT_TRUE(ShapeUtil::IsEmptyTuple(element.shape())); } } TEST_F(LiteralUtilTest, MoveIntoEmptyTuple) { Literal literal = Literal::MoveIntoTuple({}); - ASSERT_TRUE(ShapeUtil::IsTuple(literal.shape())); + ASSERT_TRUE(literal.shape().IsTuple()); EXPECT_EQ(ShapeUtil::TupleElementCount(literal.shape()), 0); } @@ -1635,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); @@ -1643,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) { @@ -1652,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 = @@ -1672,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)); @@ -1685,7 +1841,7 @@ TEST_F(LiteralUtilTest, ProtoRoundTrip) { TEST_F(LiteralUtilTest, InvalidProtoNoValues) { // Proto contains a shape, but no values. LiteralProto proto; - *proto.mutable_shape() = ShapeUtil::MakeShape(F32, {3}); + *proto.mutable_shape() = ShapeUtil::MakeShape(F32, {3}).ToProto(); Status status = Literal::CreateFromProto(proto).status(); ASSERT_FALSE(status.ok()); EXPECT_THAT(status.error_message(), @@ -1706,7 +1862,7 @@ TEST_F(LiteralUtilTest, InvalidProtoNoShape) { TEST_F(LiteralUtilTest, InvalidProtoWrongContainer) { // Proto contains values in wrong container. LiteralProto proto; - *proto.mutable_shape() = ShapeUtil::MakeShape(F32, {3}); + *proto.mutable_shape() = ShapeUtil::MakeShape(F32, {3}).ToProto(); proto.add_preds(false); proto.add_preds(true); proto.add_preds(false); @@ -1719,7 +1875,7 @@ TEST_F(LiteralUtilTest, InvalidProtoWrongContainer) { TEST_F(LiteralUtilTest, InvalidProtoTooFewValues) { // Proto contains too few values. LiteralProto proto; - *proto.mutable_shape() = ShapeUtil::MakeShape(F32, {42, 2}); + *proto.mutable_shape() = ShapeUtil::MakeShape(F32, {42, 2}).ToProto(); proto.add_f32s(1.0); proto.add_f32s(2.0); proto.add_f32s(3.0); @@ -1732,7 +1888,7 @@ TEST_F(LiteralUtilTest, InvalidProtoTooFewValues) { TEST_F(LiteralUtilTest, InvalidProtoTooManyValues) { // Proto contains too many values. LiteralProto proto; - *proto.mutable_shape() = ShapeUtil::MakeShape(S32, {2}); + *proto.mutable_shape() = ShapeUtil::MakeShape(S32, {2}).ToProto(); proto.add_s32s(42); proto.add_s32s(-10); proto.add_s32s(100); @@ -1745,8 +1901,8 @@ TEST_F(LiteralUtilTest, InvalidProtoTooManyValues) { TEST_F(LiteralUtilTest, InvalidProtoMissingLayout) { // Proto shape missing layout. LiteralProto proto; - *proto.mutable_shape() = ShapeUtil::MakeShape(PRED, {2, 2}); - LayoutUtil::ClearLayout(proto.mutable_shape()); + *proto.mutable_shape() = ShapeUtil::MakeShape(PRED, {2, 2}).ToProto(); + proto.mutable_shape()->clear_layout(); proto.add_preds(true); proto.add_preds(false); proto.add_preds(true); @@ -1759,11 +1915,13 @@ TEST_F(LiteralUtilTest, InvalidProtoMissingLayout) { TEST_F(LiteralUtilTest, InvalidProtoTooFewTupleElements) { // Proto has the too few tuple elements. LiteralProto proto; - *proto.mutable_shape() = ShapeUtil::MakeTupleShape( - {ShapeUtil::MakeShape(PRED, {2}), ShapeUtil::MakeShape(F32, {})}); + *proto.mutable_shape() = + ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeShape(PRED, {2}), ShapeUtil::MakeShape(F32, {})}) + .ToProto(); LiteralProto* element0 = proto.add_tuple_literals(); *element0->mutable_shape() = - ShapeUtil::GetTupleElementShape(proto.shape(), 0); + ShapeUtil::GetTupleElementShape(Shape(proto.shape()), 0).ToProto(); element0->add_preds(false); element0->add_preds(true); @@ -1775,19 +1933,21 @@ TEST_F(LiteralUtilTest, InvalidProtoTooFewTupleElements) { TEST_F(LiteralUtilTest, InvalidProtoTooManyTupleElements) { // Proto has the too many tuple elements. LiteralProto proto; - *proto.mutable_shape() = ShapeUtil::MakeTupleShape( - {ShapeUtil::MakeShape(PRED, {2}), ShapeUtil::MakeShape(F32, {})}); + *proto.mutable_shape() = + ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeShape(PRED, {2}), ShapeUtil::MakeShape(F32, {})}) + .ToProto(); LiteralProto* element0 = proto.add_tuple_literals(); *element0->mutable_shape() = - ShapeUtil::GetTupleElementShape(proto.shape(), 0); + ShapeUtil::GetTupleElementShape(Shape(proto.shape()), 0).ToProto(); element0->add_preds(false); element0->add_preds(true); LiteralProto* element1 = proto.add_tuple_literals(); *element1->mutable_shape() = - ShapeUtil::GetTupleElementShape(proto.shape(), 1); + ShapeUtil::GetTupleElementShape(Shape(proto.shape()), 1).ToProto(); element1->add_f32s(42.0); LiteralProto* element2 = proto.add_tuple_literals(); - *element2->mutable_shape() = ShapeUtil::MakeShape(F32, {}); + *element2->mutable_shape() = ShapeUtil::MakeShape(F32, {}).ToProto(); element2->add_f32s(123.0); Status status = Literal::CreateFromProto(proto).status(); @@ -1802,7 +1962,7 @@ TEST_F(LiteralUtilTest, SortSparseElements) { literal.AppendSparseElement({3, 4, 5}, 3.0); literal.AppendSparseElement({1, 2, 3}, 1.0); literal.SortSparseElements(); - EXPECT_EQ(literal.ToString(false), + EXPECT_EQ(literal.ToString(), "f32[10,10,10]{[1, 2, 3]: 1, [2, 3, 4]: 2, [3, 4, 5]: 3}"); } diff --git a/tensorflow/compiler/xla/literal_util.cc b/tensorflow/compiler/xla/literal_util.cc index 0cb1ae35f4ad31f091063d78ed32c1463be8ee0a..26b029c8d0c52e38510f9279def7c4af2904931d 100644 --- a/tensorflow/compiler/xla/literal_util.cc +++ b/tensorflow/compiler/xla/literal_util.cc @@ -30,7 +30,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/core/lib/core/casts.h" #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/hash/hash.h" #include "tensorflow/core/platform/logging.h" @@ -63,7 +62,7 @@ Literal ConvertType(LiteralSlice literal) { ShapeUtil::ForEachSubshape( literal.shape(), [&](const Shape& subshape, const ShapeIndex& shape_index) { - if (ShapeUtil::IsArray(subshape)) { + if (subshape.IsArray()) { if (subshape.element_type() == primitive_util::NativeToPrimitiveType()) { auto src = literal.data(shape_index); @@ -107,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: @@ -127,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: @@ -165,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: @@ -201,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: @@ -345,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()); @@ -356,7 +366,7 @@ Literal ConvertType(LiteralSlice literal) { /* static */ Literal LiteralUtil::GetFirstScalarLiteral( const LiteralSlice& literal) { - CHECK(ShapeUtil::IsArray(literal.shape())); + CHECK(literal.shape().IsArray()); CHECK_GT(ShapeUtil::ElementsIn(literal.shape()), 0); switch (literal.shape().element_type()) { case PRED: @@ -393,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/packed_literal_reader.cc b/tensorflow/compiler/xla/packed_literal_reader.cc index 0f86f9f35e105713aa3072a9ebf572d33d35d66d..339660cf44fd64fc5859e72255d63762fcf20efe 100644 --- a/tensorflow/compiler/xla/packed_literal_reader.cc +++ b/tensorflow/compiler/xla/packed_literal_reader.cc @@ -42,8 +42,7 @@ PackedLiteralReader::~PackedLiteralReader() { delete file_; } StatusOr PackedLiteralReader::Read(const Shape& shape, const Layout* layout) { VLOG(3) << "reading shape from file: " << ShapeUtil::HumanString(shape) - << " layout: " - << (layout == nullptr ? "" : layout->ShortDebugString()); + << " layout: " << (layout == nullptr ? "" : layout->ToString()); Shape literal_shape = shape; if (layout != nullptr) { TF_RETURN_IF_ERROR( diff --git a/tensorflow/compiler/xla/parse_flags_from_env.cc b/tensorflow/compiler/xla/parse_flags_from_env.cc new file mode 100644 index 0000000000000000000000000000000000000000..91e71f5d1d02d135158d0dffc140c21cf8ea5e3a --- /dev/null +++ b/tensorflow/compiler/xla/parse_flags_from_env.cc @@ -0,0 +1,236 @@ +/* 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. +==============================================================================*/ + +// This module exports ParseFlagsFromEnvAndDieIfUnknown(), which allows other +// modules to parse flags from an environtment variable, or a file named by the +// environment variable. + +#include +#include +#include +#include +#include +#include + +#include "absl/strings/ascii.h" +#include "absl/strings/str_format.h" +#include "absl/strings/str_join.h" +#include "absl/types/span.h" +#include "tensorflow/compiler/xla/parse_flags_from_env.h" +#include "tensorflow/compiler/xla/types.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/macros.h" +#include "tensorflow/core/platform/mutex.h" +#include "tensorflow/core/platform/types.h" +#include "tensorflow/core/util/command_line_flags.h" + +namespace xla { + +static const char kWS[] = " \t\r\n"; // whitespace + +// The following struct represents an argv[]-style array, parsed +// from data gleaned from the environment. +// +// As usual, an anonymous namespace is advisable to avoid +// constructor/destructor collisions with other "private" types +// in the same named namespace. +namespace { + +// Functor which deletes objects by calling `free`. Necessary to free strdup'ed +// strings created by AppendToEnvArgv. +struct FreeDeleter { + void operator()(char* ptr) { free(ptr); } +}; + +struct EnvArgv { + EnvArgv() : initialized(false), argc(0) {} + bool initialized; // whether the other fields have been set. + int argc; // elements used in argv[] + std::vector argv; // flag arguments parsed from environment string. + // saved values from argv[] to avoid leaks + std::vector> argv_save; +}; +} // anonymous namespace + +// Append the string s0[0, .., s0len-1] concatenated with s1[0, .., s1len-1] as +// a newly allocated nul-terminated string to the array *a. If s0==nullptr, a +// nullptr is appended without increasing a->argc. +static void AppendToEnvArgv(const char* s0, size_t s0len, const char* s1, + size_t s1len, EnvArgv* a) { + if (s0 == nullptr) { + a->argv.push_back(nullptr); + a->argv_save.push_back(nullptr); + } else { + string s = string(s0, s0len) + string(s1, s1len); + char* str = strdup(s.c_str()); + a->argv.push_back(str); + a->argv_save.emplace_back(str); + a->argc++; + } +} + +// Like s.find_first_of(x, pos), but return s.size() when find_first_of() would +// return string::npos. This avoids if-statements elsewhere. +static size_t FindFirstOf(const string& s, const char* x, size_t pos) { + size_t result = s.find_first_of(x, pos); + return result == string::npos ? s.size() : result; +} + +// Like s.find_first_not_of(x, pos), but return s.size() when +// find_first_not_of() would return string::npos. This avoids if-statements +// elsewhere. +static size_t FindFirstNotOf(const string& s, const char* x, size_t pos) { + size_t result = s.find_first_not_of(x, pos); + return result == string::npos ? s.size() : result; +} + +// Given a string containing flags, parse them into the XLA command line flags. +// The parse is best effort, and gives up on the first syntax error. +static void ParseArgvFromString(const string& flag_str, EnvArgv* a) { + size_t b = FindFirstNotOf(flag_str, kWS, 0); + while (b != flag_str.size() && flag_str[b] == '-') { + // b is the index of the start of a flag. + // 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 || + absl::ascii_isalnum(flag_str[e]))) { + e++; + } + if (e != flag_str.size() && flag_str[e] == '=' && + e + 1 != flag_str.size() && strchr("'\"", flag_str[e + 1]) != nullptr) { + // A flag of the form --flag="something in double or single quotes" + int c; + e++; // point just past '=' + size_t eflag = e; + char quote = flag_str[e]; + e++; // point just past quote + // Put in value the string with quotes removed. + string value; + for (; e != flag_str.size() && (c = flag_str[e]) != quote; e++) { + if (quote == '"' && c == '\\' && e + 1 != flag_str.size()) { + // Handle backslash in double quoted strings. They are literal in + // single-quoted strings. + e++; + c = flag_str[e]; + } + value += c; + } + if (e != flag_str.size()) { // skip final " or ' + e++; + } + AppendToEnvArgv(flag_str.data() + b, eflag - b, value.data(), + value.size(), a); + } else { // A flag without a quoted value. + e = FindFirstOf(flag_str, kWS, e); + AppendToEnvArgv(flag_str.data() + b, e - b, "", 0, a); + } + b = FindFirstNotOf(flag_str, kWS, e); + } +} + +// Call ParseArgvFromString(..., a) on a string derived from the setting of the +// environment variable `envvar`, or a file it points to. +static void SetArgvFromEnv(absl::string_view envvar, EnvArgv* a) { + if (!a->initialized) { + static const char kDummyArgv[] = ""; + AppendToEnvArgv(kDummyArgv, strlen(kDummyArgv), nullptr, 0, + a); // dummy argv[0] + const char* env = getenv(string(envvar).c_str()); + if (env == nullptr || env[0] == '\0') { + // nothing + } else if (env[strspn(env, kWS)] == '-') { // flags in env var value + ParseArgvFromString(env, a); + } else { // assume it's a file name + FILE* fp = fopen(env, "r"); + if (fp != nullptr) { + string str; + char buf[512]; + int n; + while ((n = fread(buf, 1, sizeof(buf), fp)) > 0) { + str.append(buf, n); + } + fclose(fp); + ParseArgvFromString(str, a); + } + } + AppendToEnvArgv(nullptr, 0, nullptr, 0, a); // add trailing nullptr to *a. + a->initialized = true; + } +} + +// The simulated argv[] parsed from the environment, one for each different +// environment variable we've seen. +static std::unordered_map& EnvArgvs() { + static auto* env_argvs = new std::unordered_map(); + return *env_argvs; +} + +// Used to protect accesses to env_argvs. +static tensorflow::mutex env_argv_mu(tensorflow::LINKER_INITIALIZED); + +bool ParseFlagsFromEnvAndDieIfUnknown( + absl::string_view envvar, const std::vector& flag_list) { + tensorflow::mutex_lock lock(env_argv_mu); + auto* env_argv = &EnvArgvs()[string(envvar)]; + SetArgvFromEnv(envvar, env_argv); // a no-op if already initialized + bool result = + tensorflow::Flags::Parse(&env_argv->argc, &env_argv->argv[0], flag_list); + + // There's always at least one unparsed argc, namely the fake argv[0]. + if (result && env_argv->argc != 1) { + // Skip the first argv, which is the fake argv[0]. + auto unknown_flags = absl::MakeSpan(env_argv->argv); + unknown_flags.remove_prefix(1); + + // Some flags are set on XLA_FLAGS, others on TF_XLA_FLAGS. If we find an + // unrecognized flag, suggest the alternative. + string alternate_envvar; + if (envvar == "TF_XLA_FLAGS") { + alternate_envvar = "XLA_FLAGS"; + } else if (envvar == "XLA_FLAGS") { + alternate_envvar = "TF_XLA_FLAGS"; + } + string did_you_mean; + if (!alternate_envvar.empty()) { + did_you_mean = absl::StrFormat( + "\nPerhaps you meant to specify these on the %s envvar?", + alternate_envvar); + } + + LOG(FATAL) << "Unknown flag" << (unknown_flags.size() > 1 ? "s" : "") + << " in " << envvar << ": " << absl::StrJoin(unknown_flags, " ") + << did_you_mean; + return false; + } + return result; +} + +// Testing only. +// +// Resets the env_argv struct so that subsequent calls to +// ParseFlagsFromEnvAndDieIfUnknown() will parse the environment variable (or +// the file it points to) anew, and set *pargc, and *pargv to point to the +// internal locations of the argc and argv constructed from the environment. +void ResetFlagsFromEnvForTesting(absl::string_view envvar, int** pargc, + std::vector** pargv) { + tensorflow::mutex_lock lock(env_argv_mu); + EnvArgvs().erase(string(envvar)); + auto& env_argv = EnvArgvs()[string(envvar)]; + *pargc = &env_argv.argc; + *pargv = &env_argv.argv; +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/parse_flags_from_env.h b/tensorflow/compiler/xla/parse_flags_from_env.h new file mode 100644 index 0000000000000000000000000000000000000000..76940a4299ac50138222333ff250a264cc941288 --- /dev/null +++ b/tensorflow/compiler/xla/parse_flags_from_env.h @@ -0,0 +1,74 @@ +/* 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_PARSE_FLAGS_FROM_ENV_H_ +#define TENSORFLOW_COMPILER_XLA_PARSE_FLAGS_FROM_ENV_H_ + +// This module exports ParseFlagsFromEnvAndDieIfUnknown(), which allows other +// modules to parse flags from an environtment variable, or (if the first +// non-whitespace in the variable value is not '-'), a file named by that +// environment variable. +// +// The accepted syntax is that flags arguments are of the form --flag=value or +// (for boolean flags) --flag, and are whitespace separated. The may be +// one of: +// +// - +// in which case the effective value is the string itself +// - in which case the effective value is the +// string with the single-quotes removed +// - in which case the effective value if the +// string with the double-quotes removed, and escaped sequences of +// replaced by . +// +// Flags values inconsistent with the type of the flag will be rejected by the +// flag parser. +// +// Examples: +// +// - TF_XLA_FLAGS="--foo=bar --wombat='value with a space'" +// - TF_XLA_FLAGS=/tmp/flagfile +// +// where /tmp/flagfile might contain +// +// --some_flag="This is a string containing a \" and a '." +// --another_flag=wombats + +#include + +#include "absl/strings/string_view.h" +#include "tensorflow/compiler/xla/types.h" +#include "tensorflow/core/platform/types.h" +#include "tensorflow/core/util/command_line_flags.h" + +namespace xla { + +// Calls tensorflow::Flags::Parse(argc, argv, flag_list) against any as yet +// unrecognized flags passed in the environment variable `envvar`, and returns +// its return value. +// +// Raises a fatal error if any flags in `envvar` were not recognized. +bool ParseFlagsFromEnvAndDieIfUnknown( + absl::string_view envvar, const std::vector& flag_list); + +// Used only for testing. Not to be used by clients. +void ResetFlagsFromEnvForTesting(absl::string_view envvar, int** pargc, + std::vector** pargv); + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_PARSE_FLAGS_FROM_ENV_H_ diff --git a/tensorflow/compiler/xla/legacy_flags/parse_flags_from_env_test.cc b/tensorflow/compiler/xla/parse_flags_from_env_test.cc similarity index 89% rename from tensorflow/compiler/xla/legacy_flags/parse_flags_from_env_test.cc rename to tensorflow/compiler/xla/parse_flags_from_env_test.cc index 138c0c852e2bb0527d171f25b4d96cedc5671516..3465552ebbf52140fb954b247d99d3c6afe7fcde 100644 --- a/tensorflow/compiler/xla/legacy_flags/parse_flags_from_env_test.cc +++ b/tensorflow/compiler/xla/parse_flags_from_env_test.cc @@ -15,7 +15,7 @@ limitations under the License. // Test for parse_flags_from_env.cc -#include "tensorflow/compiler/xla/legacy_flags/parse_flags_from_env.h" +#include "tensorflow/compiler/xla/parse_flags_from_env.h" #include #include @@ -30,7 +30,6 @@ limitations under the License. #include "tensorflow/core/util/command_line_flags.h" namespace xla { -namespace legacy_flags { // Test that XLA flags can be set from the environment. // Failure messages are accompanied by the text in msg[]. @@ -38,20 +37,7 @@ static void TestParseFlagsFromEnv(const char* msg) { // Initialize module under test. int* pargc; std::vector* pargv; - ResetFlagsFromEnvForTesting(&pargc, &pargv); - - // Ensure that environment variable can be parsed when - // no flags are expected. - std::vector empty_flag_list; - bool parsed_ok = ParseFlagsFromEnv(empty_flag_list); - CHECK(parsed_ok) << msg; - const std::vector& argv_first = *pargv; - CHECK_NE(argv_first[0], nullptr) << msg; - int i = 0; - while (argv_first[i] != nullptr) { - i++; - } - CHECK_EQ(i, *pargc) << msg; + ResetFlagsFromEnvForTesting("TF_XLA_FLAGS", &pargc, &pargv); // Check that actual flags can be parsed. bool simple = false; @@ -66,7 +52,7 @@ static void TestParseFlagsFromEnv(const char* msg) { tensorflow::Flag("single_quoted", &single_quoted, ""), tensorflow::Flag("double_quoted", &double_quoted, ""), }; - parsed_ok = ParseFlagsFromEnv(flag_list); + bool parsed_ok = ParseFlagsFromEnvAndDieIfUnknown("TF_XLA_FLAGS", flag_list); CHECK_EQ(*pargc, 1) << msg; const std::vector& argv_second = *pargv; CHECK_NE(argv_second[0], nullptr) << msg; @@ -159,12 +145,11 @@ TEST(ParseFlagsFromEnv, EnvAndFlag) { } } -} // namespace legacy_flags } // namespace xla int main(int argc, char* argv[]) { // Save name of binary so that it may invoke itself. - xla::legacy_flags::binary_name = argv[0]; + xla::binary_name = argv[0]; bool recursing = false; xla::int32 int_flag = 1; const std::vector flag_list = { @@ -173,7 +158,8 @@ int main(int argc, char* argv[]) { tensorflow::Flag("int_flag", &int_flag, "An integer flag to test with"), }; xla::string usage = tensorflow::Flags::Usage(argv[0], flag_list); - bool parse_ok = xla::legacy_flags::ParseFlagsFromEnv(flag_list); + bool parse_ok = + xla::ParseFlagsFromEnvAndDieIfUnknown("TF_XLA_FLAGS", flag_list); if (!parse_ok) { LOG(QFATAL) << "can't parse from environment\n" << usage; } diff --git a/tensorflow/compiler/xla/primitive_util.cc b/tensorflow/compiler/xla/primitive_util.cc index b16147e3be71771269d8b7a18528bef3a8c72d99..1eedddf72c1d393cb1b88e589881e24de02ad802 100644 --- a/tensorflow/compiler/xla/primitive_util.cc +++ b/tensorflow/compiler/xla/primitive_util.cc @@ -15,16 +15,35 @@ limitations under the License. #include "tensorflow/compiler/xla/primitive_util.h" +#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; @@ -64,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"; @@ -75,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); @@ -90,5 +129,65 @@ bool IsArrayType(PrimitiveType primitive_type) { primitive_type != OPAQUE && primitive_type != TOKEN; } +// Class to memoize the computation of +// absl::AsciiStrToLower(PrimitiveType_Name(p)) +// for all PrimitiveType values "p" +class PrimitiveTypeNameGenerator { + public: + PrimitiveTypeNameGenerator() { + for (int i = 0; i < PrimitiveType_ARRAYSIZE; i++) { + if (PrimitiveType_IsValid(i)) { + lowercase_name_[i] = absl::AsciiStrToLower( + PrimitiveType_Name(static_cast(i))); + } + } + } + const string& LowercaseName(PrimitiveType t) { + return lowercase_name_[static_cast(t)]; + } + + private: + string lowercase_name_[PrimitiveType_ARRAYSIZE]; +}; + +const string& LowercasePrimitiveTypeName(PrimitiveType s) { + static auto* gen = new PrimitiveTypeNameGenerator(); + return gen->LowercaseName(s); +} + +namespace { + +// Returns a map from lower-case primitive type name to primitive type. +const std::unordered_map& GetPrimitiveTypeStringMap() { + static std::unordered_map* name_to_type = [] { + static auto* map = new std::unordered_map; + for (int i = 0; i < PrimitiveType_ARRAYSIZE; i++) { + if (PrimitiveType_IsValid(i) && i != PRIMITIVE_TYPE_INVALID) { + auto value = static_cast(i); + (*map)[LowercasePrimitiveTypeName(value)] = value; + } + } + return map; + }(); + return *name_to_type; +} + +} // namespace + +StatusOr StringToPrimitiveType(absl::string_view name) { + const auto& map = GetPrimitiveTypeStringMap(); + auto found = map.find(string(name)); + if (found == map.end()) { + return InvalidArgument("Invalid element type string: \"%s\".", name); + } + return found->second; +} + +bool IsPrimitiveTypeName(absl::string_view name) { + const auto& map = GetPrimitiveTypeStringMap(); + auto found = map.find(string(name)); + return found != map.end(); +} + } // namespace primitive_util } // namespace xla diff --git a/tensorflow/compiler/xla/primitive_util.h b/tensorflow/compiler/xla/primitive_util.h index 889e9a1ceca675689406d255d348c82c398563aa..295d353003276b4c1731f7d6a378fd1ae0288d3c 100644 --- a/tensorflow/compiler/xla/primitive_util.h +++ b/tensorflow/compiler/xla/primitive_util.h @@ -20,12 +20,19 @@ limitations under the License. #include +#include "absl/strings/string_view.h" +#include "tensorflow/compiler/xla/status_macros.h" +#include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" 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; @@ -123,6 +130,11 @@ inline PrimitiveType NativeToPrimitiveType() { return C64; } +template <> +inline PrimitiveType NativeToPrimitiveType() { + return C128; +} + bool IsFloatingPointType(PrimitiveType type); bool IsComplexType(PrimitiveType type); @@ -139,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); @@ -221,6 +235,22 @@ template <> 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); + +// Returns the PrimitiveType matching the given name. The given name is expected +// to be lower-case. +StatusOr StringToPrimitiveType(absl::string_view name); + +// Returns true if the given name is a primitive type string (lower-case). +bool IsPrimitiveTypeName(absl::string_view name); + } // namespace primitive_util } // namespace xla diff --git a/tensorflow/compiler/xla/primitive_util_test.cc b/tensorflow/compiler/xla/primitive_util_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..1f765d6da9ef65849fe8ede56ced7597d623cb59 --- /dev/null +++ b/tensorflow/compiler/xla/primitive_util_test.cc @@ -0,0 +1,46 @@ +/* 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/primitive_util.h" + +#include +#include "tensorflow/compiler/xla/status_macros.h" +#include "tensorflow/compiler/xla/test.h" +#include "tensorflow/compiler/xla/test_helpers.h" +#include "tensorflow/compiler/xla/types.h" +#include "tensorflow/compiler/xla/util.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" + +namespace xla { +namespace { + +TEST(PrimitiveUtilTest, StringToPrimitiveType) { + auto expect_ok_and_equal = [](const string& str, PrimitiveType expected) { + TF_ASSERT_OK_AND_ASSIGN(PrimitiveType actual, + primitive_util::StringToPrimitiveType(str)); + EXPECT_EQ(expected, actual); + }; + expect_ok_and_equal("f32", F32); + expect_ok_and_equal("tuple", TUPLE); + expect_ok_and_equal("pred", PRED); + expect_ok_and_equal("s32", S32); + + EXPECT_IS_NOT_OK(primitive_util::StringToPrimitiveType("F32").status()); + EXPECT_IS_NOT_OK(primitive_util::StringToPrimitiveType("Pred").status()); + EXPECT_IS_NOT_OK(primitive_util::StringToPrimitiveType("preD").status()); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/protobuf_util.cc b/tensorflow/compiler/xla/protobuf_util.cc index b507a2ef79f1d7e9ae632744675dddf574490805..ac342bf40fbc0052acbb09a346b9d062561ed06b 100644 --- a/tensorflow/compiler/xla/protobuf_util.cc +++ b/tensorflow/compiler/xla/protobuf_util.cc @@ -40,16 +40,6 @@ bool ProtobufEquals(const tensorflow::protobuf::Message& m1, namespace { -string SanitizeFilename(const string& file_name) { - string safe_file_name = file_name; - for (char& c : safe_file_name) { - if (c == '/' || c == '\\') { - c = '_'; - } - } - return safe_file_name; -} - std::pair>*> GetDirectoryExpanders() { static auto* mutex = new tensorflow::mutex; diff --git a/tensorflow/compiler/xla/python/BUILD b/tensorflow/compiler/xla/python/BUILD index f0d84646b9f01ad3ad209073f13b7b3ec21635d1..8276ee5a32e390edee6103544bfcfd951f461187 100644 --- a/tensorflow/compiler/xla/python/BUILD +++ b/tensorflow/compiler/xla/python/BUILD @@ -3,6 +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") py_library( name = "xla_client", @@ -16,6 +18,12 @@ py_library( ], ) +pyx_library( + name = "custom_call_for_test", + testonly = True, + srcs = ["custom_call_for_test.pyx"], +) + py_test( name = "xla_client_test", srcs = ["xla_client_test.py"], @@ -23,6 +31,7 @@ py_test( srcs_version = "PY2AND3", tags = ["no_oss"], deps = [ + ":custom_call_for_test", ":xla_client", "//tensorflow/python:platform_test", ], @@ -50,7 +59,14 @@ cc_library( srcs = ["local_computation_builder.cc"], hdrs = ["local_computation_builder.h"], deps = [ + "//tensorflow/cc:cc_ops", + "//tensorflow/cc:client_session", + "//tensorflow/cc:ops", + "//tensorflow/cc:scope", "//tensorflow/compiler/xla:executable_run_options", + "//tensorflow/compiler/xla:literal", + "//tensorflow/compiler/xla:literal_util", + "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/client:client_library", @@ -58,10 +74,18 @@ cc_library( "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/client:xla_builder", "//tensorflow/compiler/xla/client:xla_computation", + "//tensorflow/compiler/xla/client/lib:cholesky", "//tensorflow/compiler/xla/client/lib:math", + "//tensorflow/compiler/xla/client/lib:qr", + "//tensorflow/compiler/xla/client/lib:triangular_solve", + "//tensorflow/compiler/xla/service:platform_util", "//tensorflow/compiler/xla/service:shaped_buffer", - "//tensorflow/core:framework_lite", + "//tensorflow/compiler/xla/service/cpu:custom_call_target_registry", + "//tensorflow/compiler/xrt:xrt_proto", + "//tensorflow/compiler/xrt/cc:xrt_ops", + "//tensorflow/core:framework", "//tensorflow/core:lib", + "//third_party/python_runtime:headers", "@com_google_absl//absl/memory", "@com_google_absl//absl/types:span", ], @@ -72,7 +96,13 @@ tf_py_wrap_cc( srcs = ["xla.i"], swig_includes = [ "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", @@ -80,5 +110,7 @@ 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", + ]), ) diff --git a/tensorflow/compiler/xla/python/custom_call_for_test.pyx b/tensorflow/compiler/xla/python/custom_call_for_test.pyx new file mode 100644 index 0000000000000000000000000000000000000000..530dffd1755d8438f52569c223525000c97df6ea --- /dev/null +++ b/tensorflow/compiler/xla/python/custom_call_for_test.pyx @@ -0,0 +1,21 @@ +# distutils: language = c++ + +# Test case for defining a XLA custom call target in Cython, and registering +# it via the xla_client SWIG API. + +from cpython.pycapsule cimport PyCapsule_New + +cdef void test_subtract_f32(void* out_ptr, void** data_ptr) nogil: + cdef float a = ((data_ptr[0]))[0] + cdef float b = ((data_ptr[1]))[0] + cdef float* out = (out_ptr) + out[0] = a - b + + +cpu_custom_call_targets = {} + +cdef register_custom_call_target(fn_name, void* fn): + cdef const char* name = "xla._CPU_CUSTOM_CALL_TARGET" + cpu_custom_call_targets[fn_name] = PyCapsule_New(fn, name, NULL) + +register_custom_call_target(b"test_subtract_f32", (test_subtract_f32)) diff --git a/tensorflow/compiler/xla/python/local_computation_builder.cc b/tensorflow/compiler/xla/python/local_computation_builder.cc index 92df404b8ec0aed4899906877a4dd41102bdf7a0..0e898d494e044509a41209891c28d929dff11b9a 100644 --- a/tensorflow/compiler/xla/python/local_computation_builder.cc +++ b/tensorflow/compiler/xla/python/local_computation_builder.cc @@ -14,16 +14,46 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/compiler/xla/python/local_computation_builder.h" + +#include +#include +#include + #include "absl/memory/memory.h" +#include "tensorflow/cc/client/client_session.h" +#include "tensorflow/cc/framework/ops.h" +#include "tensorflow/cc/framework/scope.h" +#include "tensorflow/cc/ops/standard_ops.h" +#include "tensorflow/compiler/xla/client/lib/cholesky.h" #include "tensorflow/compiler/xla/client/lib/math.h" +#include "tensorflow/compiler/xla/client/lib/qr.h" +#include "tensorflow/compiler/xla/client/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" +#include "tensorflow/compiler/xla/literal.h" +#include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/compiler/xla/service/cpu/custom_call_target_registry.h" +#include "tensorflow/compiler/xla/service/platform_util.h" +#include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/util.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/compiler/xrt/cc/ops/xrt_compile_ops.h" +#include "tensorflow/compiler/xrt/cc/ops/xrt_execute_op.h" +#include "tensorflow/compiler/xrt/cc/ops/xrt_state_ops.h" +#include "tensorflow/compiler/xrt/xrt.pb.h" +#include "tensorflow/core/framework/tensor.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/lib/gtl/array_slice.h" #include "tensorflow/core/platform/thread_annotations.h" +#include "tensorflow/core/platform/types.h" namespace xla { namespace swig { +// TODO(b/118641336): Factor out XRT parts into a small c++ library of their +// own. + // TODO(b/34473877) Ideally XLA would support AllReduce among arbitrary sets of // device handles instead of needing to set the number of replicas at XLA // service initialization time. @@ -31,6 +61,12 @@ tensorflow::mutex g_local_client_mutex(tensorflow::LINKER_INITIALIZED); int g_replica_count GUARDED_BY(g_local_client_mutex) = 1; LocalClient* g_local_client GUARDED_BY(g_local_client_mutex) = nullptr; +string* GetPlatformNameString() { + static string* platform_name_string PT_GUARDED_BY(g_local_client_mutex) = + new string("Host"); + return platform_name_string; +} + Status InitializeReplicaCount(int replica_count) { if (replica_count < 1) { return InvalidArgument("Replica count must be >= 1; got %d.", @@ -47,27 +83,58 @@ Status InitializeReplicaCount(int replica_count) { return Status::OK(); } +Status InitializePlatformName(const string& platform_name) { + string* g_platform_name = GetPlatformNameString(); + tensorflow::mutex_lock lock(g_local_client_mutex); + if (g_local_client != nullptr) { + return FailedPrecondition( + "Attempted to set the platform name to %s, but a local XLA service was " + "previously created with a platform name of %s.", + platform_name, *g_platform_name); + } + TF_RETURN_IF_ERROR(PlatformUtil::GetPlatform(platform_name).status()); + *g_platform_name = platform_name; + return Status::OK(); +} + int GetReplicaCount() { tensorflow::mutex_lock lock(g_local_client_mutex); 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) { return g_local_client; } 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; } +Status RegisterCpuCustomCallTarget(const string& fn_name, PyObject* capsule) { + const char* name = "xla._CPU_CUSTOM_CALL_TARGET"; + if (!PyCapsule_IsValid(capsule, name)) { + return InvalidArgument( + "Argument to RegisterCpuCustomCallTargetRegistry was not a " + "xla._CPU_CUSTOM_CALL_TARGET capsule."); + } + void* fn_ptr = PyCapsule_GetPointer(capsule, name); + CHECK(fn_ptr != nullptr); + cpu::CustomCallTargetRegistry::Global()->Register( + std::string(fn_name.begin(), fn_name.end()), fn_ptr); + return Status::OK(); +} + 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); } @@ -75,7 +142,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); @@ -85,12 +152,39 @@ 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); } +static StatusOr ToBuffer(LocalClient* client, + int device_ordinal, + const Literal& arg) { + return client->LiteralToShapedBuffer(arg, device_ordinal, + client->backend().memory_allocator()); +} + +/* static */ +StatusOr LocalShapedBuffer::FromLiteral( + const Literal& argument, const absl::optional& shape_with_layout, + int replica_number) { + 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: " + << replica_number << "/" << device_ordinal; + StatusOr buf = [&] { + if (shape_with_layout) { + Literal relaid = argument.Relayout(shape_with_layout.value()); + return ToBuffer(client, device_ordinal, relaid); + } + return ToBuffer(client, device_ordinal, argument); + }(); + TF_RETURN_IF_ERROR(buf.status()); + return new LocalShapedBuffer(std::move(buf).ValueOrDie()); +} + LocalShapedBuffer::LocalShapedBuffer(ScopedShapedBuffer shaped_buffer) : shaped_buffer_(std::move(shaped_buffer)) {} @@ -100,11 +194,20 @@ const ScopedShapedBuffer* LocalShapedBuffer::shaped_buffer() const { ShapedBuffer LocalShapedBuffer::Release() { return shaped_buffer_.release(); } +const Shape& LocalShapedBuffer::shape() const { + return shaped_buffer()->on_device_shape(); +} + +StatusOr LocalShapedBuffer::ToLiteral() const { + TF_ASSIGN_OR_RETURN(LocalClient * client, GetOrCreateLocalClient()); + return client->ShapedBufferToLiteral(*shaped_buffer()); +} + LocalShapedBufferTuple::LocalShapedBufferTuple( std::vector elements) : elements_(std::move(elements)) { for (auto* element : elements_) { - DCHECK(element != nullptr); + CHECK(element != nullptr); } } @@ -126,157 +229,307 @@ StatusOr LocalShapedBufferTuple::Release(int i) { return element; } -int LocalShapedBufferTuple::size() const { return elements_.size(); } +int64 LocalShapedBufferTuple::size() const { return elements_.size(); } + +XrtAllocation::XrtAllocation(int64 handle, Shape shape, + const string& session_target) + : handle_(handle), shape_(shape), session_target_(session_target) {} + +XrtAllocation::~XrtAllocation() { + tensorflow::Scope root = tensorflow::Scope::NewRootScope(); + auto allocation_handle = + tensorflow::ops::Placeholder(root, tensorflow::DT_INT64); + auto release = + tensorflow::ops::XRTReleaseAllocationHandle(root, allocation_handle); + if (!root.status().ok()) { + LOG(ERROR) << root.status(); + return; + } -static StatusOr ToBuffer(LocalClient* client, - int device_ordinal, - const Literal& arg) { - return client->LiteralToShapedBuffer(arg, device_ordinal, - client->backend().memory_allocator()); + tensorflow::ClientSession session(root, session_target_); + tensorflow::ClientSession::FeedType inputs; + inputs.insert({allocation_handle, handle()}); + std::vector outputs; + auto status = session.Run(inputs, {}, {release}, &outputs); + if (!status.ok()) { + LOG(ERROR) << status; + return; + } } /* static */ -StatusOr LocalShapedBuffer::FromLiteral( - const Literal& argument, const absl::optional& shape_with_layout) { - LocalClient* client = GetOrCreateLocalClient(); - StatusOr buf = [&] { - if (shape_with_layout) { - Literal relaid = argument.Relayout(shape_with_layout.value()); - return ToBuffer(client, /*device_ordinal=*/0, relaid); +StatusOr XrtAllocation::FromLiteral( + const Literal& argument, const string& session_target) { + xrt::XLAAllocation alloc; + *alloc.mutable_value() = argument.ToProto(); + + tensorflow::Scope root = tensorflow::Scope::NewRootScope(); + auto literal_string = + tensorflow::ops::Placeholder(root, tensorflow::DT_STRING); + auto literal_handle = tensorflow::ops::XRTAllocate(root, literal_string); + TF_RETURN_IF_ERROR(root.status()); + + tensorflow::ClientSession session(root, session_target); + tensorflow::ClientSession::FeedType inputs; + inputs.insert({literal_string, alloc.SerializeAsString()}); + std::vector outputs; + TF_RETURN_IF_ERROR(session.Run(inputs, {literal_handle}, &outputs)); + + int64 handle = outputs[0].scalar()(); + return new XrtAllocation(handle, argument.shape(), session_target); +} + +const int64 XrtAllocation::handle() const { return handle_; } + +const Shape& XrtAllocation::shape() const { return shape_; } + +StatusOr XrtAllocation::ToLiteral() const { + tensorflow::Scope root = tensorflow::Scope::NewRootScope(); + auto allocation_handle = + tensorflow::ops::Placeholder(root, tensorflow::DT_INT64); + auto read_literal = tensorflow::ops::XRTReadLiteral(root, allocation_handle); + TF_RETURN_IF_ERROR(root.status()); + + tensorflow::ClientSession session(root, session_target_); + tensorflow::ClientSession::FeedType inputs; + inputs.insert({allocation_handle, handle()}); + std::vector outputs; + TF_RETURN_IF_ERROR(session.Run(inputs, {read_literal}, &outputs)); + + xla::LiteralProto response; + TF_RET_CHECK(response.ParseFromString(outputs[0].scalar()())); + return Literal::CreateFromProto(response); +} + +XrtAllocationTuple::XrtAllocationTuple(std::vector elements) + : elements_(std::move(elements)) { + for (auto* element : elements_) { + CHECK(element != nullptr); + } +} + +XrtAllocationTuple::~XrtAllocationTuple() { + for (XrtAllocation* element : elements_) { + if (element != nullptr) { + delete element; } - return ToBuffer(client, /*device_ordinal=*/0, argument); - }(); - TF_RETURN_IF_ERROR(buf.status()); - return new LocalShapedBuffer(std::move(buf).ValueOrDie()); + } } -StatusOr LocalShapedBuffer::ToLiteral() const { - LocalClient* client = GetOrCreateLocalClient(); - return client->ShapedBufferToLiteral(*shaped_buffer()); +StatusOr XrtAllocationTuple::Release(int i) { + XrtAllocation* element = elements_[i]; + if (element == nullptr) { + return InvalidArgument("Attempted to release already-released element %d.", + i); + } + elements_[i] = nullptr; + return element; } +int64 XrtAllocationTuple::size() const { return elements_.size(); } + CompiledLocalComputation::CompiledLocalComputation( std::unique_ptr executable) : executable_(std::move(executable)) {} -StatusOr CompiledLocalComputation::Execute( - const std::vector& arguments, - const std::vector>& shapes_with_layout) { - LocalClient* client = GetOrCreateLocalClient(); +StatusOr CompiledLocalComputation::Execute( + absl::Span argument_handles) { + 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; + const int device_ordinal = device_assignment(0, 0); + VLOG(3) << "Replica 0 mapped to device ordinal for execution: " + << device_ordinal; - VLOG(1) << "Execution requested with " << GetReplicaCount() << " replicas."; + std::vector argument_buffers; + argument_buffers.reserve(argument_handles.size()); + for (auto& handle : argument_handles) { + argument_buffers.push_back(handle->shaped_buffer()); + } - // Each replica populates a StatusOr result, but only replica zero actually - // retrieves its literal value. - std::vector> results(GetReplicaCount()); - { + 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); + + if (!result_buffer_status.ok()) { + return InternalError( + "Failed running replica 0 (other replicas may have failed as well): " + "%s.", + result_buffer_status.status().ToString()); + } + return new LocalShapedBuffer(std::move(result_buffer_status).ValueOrDie()); +} + +StatusOr CompiledLocalComputation::ExecutePerReplica( + absl::Span> argument_handles) { + TF_ASSIGN_OR_RETURN(LocalClient * client, GetOrCreateLocalClient()); + const int num_devices = client->device_count(); + + if (argument_handles.size() != num_replicas()) { + return InvalidArgument( + "Attempted to execute with %d replicas when replica count is %d", + 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."; + + 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) { + const int device_ordinal = device_assignment(replica, 0); + VLOG(3) << "Replica " << replica + << " mapped to device ordinal for execution: " << device_ordinal; + + std::vector argument_buffers; + argument_buffers.reserve(argument_handles[replica].size()); + for (auto& handle : argument_handles[replica]) { + 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); + StatusOr result_buffer_status = + executable_->Run(argument_buffers, options); + + results[replica] = std::move(result_buffer_status); + }; + + 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", - GetReplicaCount()); - - for (int replica = 0; replica < GetReplicaCount(); ++replica) { - pool.Schedule( - [this, client, replica, &arguments, &shapes_with_layout, &results] { - 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(); - VLOG(3) << "Replica " << replica - << " mapped to device ordinal for execution: " - << device_ordinal; - - // Transfer arguments in - std::vector scoped_buffers; - scoped_buffers.reserve(arguments.size()); - for (int i = 0; i < arguments.size(); ++i) { - const Literal& argument = arguments[i]; - const absl::optional& shape_with_layout = - shapes_with_layout[i]; - - StatusOr pushed; - if (shape_with_layout) { - Literal relaid = argument.Relayout(shape_with_layout.value()); - pushed = ToBuffer(client, device_ordinal, relaid); - } else { - pushed = ToBuffer(client, device_ordinal, argument); - } - if (!pushed.ok()) { - results[replica] = pushed.status(); - return; - } - - scoped_buffers.push_back(std::move(pushed).ValueOrDie()); - } - - // Execute - std::vector argument_buffers; - argument_buffers.reserve(scoped_buffers.size()); - for (auto& buffer : scoped_buffers) { - argument_buffers.push_back(&buffer); - } - - DeviceAssignment device_assignment = - client->backend() - .computation_placer() - ->AssignDevices(GetReplicaCount(), /*computation_count=*/1) - .ConsumeValueOrDie(); - - 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); - StatusOr result_buffer_status = - executable_->Run(argument_buffers, options); - if (!result_buffer_status.ok()) { - results[replica] = result_buffer_status.status(); - return; - } - - // Transfer result out - results[replica] = client->ShapedBufferToLiteral( - std::move(result_buffer_status).ValueOrDie()); - }); + num_replicas() - 1); + + for (int replica = 0; replica < num_replicas() - 1; ++replica) { + pool.Schedule([&execute, replica] { execute(replica); }); } + execute(num_replicas() - 1); } - for (int replica = 0; replica < GetReplicaCount(); ++replica) { - const auto& statusor = results[replica]; + std::vector wrapped_results(num_replicas()); + for (int replica = 0; replica < num_replicas(); ++replica) { + auto& statusor = results[replica]; if (!statusor.ok()) { return InternalError( "Failed running replica %d (other replicas may have failed as well): " "%s.", replica, statusor.status().ToString()); } + wrapped_results[replica] = + new LocalShapedBuffer(std::move(statusor).ValueOrDie()); } - return std::move(results[0]); + return new LocalShapedBufferTuple(std::move(wrapped_results)); } -LocalShapedBuffer* CompiledLocalComputation::ExecuteWithShapedBuffers( - absl::Span argument_handles) { - LocalClient* client = GetOrCreateLocalClient(); +static StatusOr GetReturnValueShape(const XlaComputation& computation) { + TF_ASSIGN_OR_RETURN(ProgramShape program_shape, + computation.GetProgramShape()); + return std::move(*program_shape.mutable_result()); +} - std::vector argument_buffers; - argument_buffers.reserve(argument_handles.size()); - for (auto& handle : argument_handles) { - argument_buffers.push_back(handle->shaped_buffer()); +CompiledXrtComputation::CompiledXrtComputation( + const ProgramShape& program_shape, int64 handle, + const string& session_target) + : program_shape_(program_shape), + handle_(handle), + session_target_(session_target) {} + +CompiledXrtComputation::~CompiledXrtComputation() { + tensorflow::Scope root = tensorflow::Scope::NewRootScope(); + auto computation_handle = + tensorflow::ops::Placeholder(root, tensorflow::DT_INT64); + auto release = + tensorflow::ops::XRTReleaseCompilationHandle(root, computation_handle); + if (!root.status().ok()) { + LOG(ERROR) << root.status(); + return; } - // Execute - ExecutableRunOptions options; - options.set_allocator(client->backend().memory_allocator()); - options.set_intra_op_thread_pool( - client->backend().eigen_intra_op_thread_pool_device()); - ScopedShapedBuffer result_buffer = - executable_->Run(argument_buffers, options).ConsumeValueOrDie(); + tensorflow::ClientSession session(root, session_target_); + tensorflow::ClientSession::FeedType inputs; + inputs.insert({computation_handle, handle()}); + std::vector outputs; + auto status = session.Run(inputs, {}, {release}, &outputs); + if (!status.ok()) { + LOG(ERROR) << status; + return; + } +} - return new LocalShapedBuffer(std::move(result_buffer)); +StatusOr CompiledXrtComputation::Execute( + absl::Span argument_handles) { + const int num_expected_arguments = program_shape().parameters().size(); + + tensorflow::Scope root = tensorflow::Scope::NewRootScope(); + std::vector arguments; + arguments.reserve(num_expected_arguments); + for (int i = 0; i < num_expected_arguments; ++i) { + arguments.push_back( + tensorflow::ops::Placeholder(root, tensorflow::DT_INT64)); + } + auto computation_handle = + tensorflow::ops::Placeholder(root, tensorflow::DT_INT64); + auto execution_config = + tensorflow::ops::Placeholder(root, tensorflow::DT_STRING); + auto execute = tensorflow::ops::XRTExecute(root, computation_handle, + execution_config, arguments); + TF_RETURN_IF_ERROR(root.status()); + + TF_RET_CHECK(argument_handles.size() == arguments.size()); + + xrt::XRTExecutionConfig e; + e.set_release_input_handles(false); + e.set_release_compilation_handle(false); + + tensorflow::ClientSession session(root, session_target_); + tensorflow::ClientSession::FeedType inputs; + for (int i = 0; i < arguments.size(); ++i) { + inputs.insert({arguments[i], argument_handles[i]->handle()}); + } + inputs.insert({computation_handle, handle()}); + inputs.insert({execution_config, e.SerializeAsString()}); + std::vector outputs; + TF_RETURN_IF_ERROR(session.Run(inputs, {execute}, &outputs)); + + int64 output = outputs[0].scalar()(); + return new XrtAllocation(output, program_shape().result(), session_target_); +} + +const ProgramShape& CompiledXrtComputation::program_shape() const { + return program_shape_; } +int64 CompiledXrtComputation::handle() const { return handle_; } + LocalComputation::LocalComputation(XlaComputation computation) : computation_(std::move(computation)) {} @@ -289,7 +542,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; @@ -300,6 +553,37 @@ StatusOr LocalComputation::Compile( return new CompiledLocalComputation(std::move(local_executable)); } +StatusOr LocalComputation::CompileForXrt( + const std::vector& argument_shapes, const string& session_target) { + tensorflow::Scope root = tensorflow::Scope::NewRootScope(); + auto program = tensorflow::ops::Placeholder(root, tensorflow::DT_STRING); + auto compile = tensorflow::ops::XRTCompile(root, program); + TF_RETURN_IF_ERROR(root.status()); + + xrt::XLAComputation c; + auto config = c.mutable_config(); + ProgramShape shapes; + for (auto& shape : argument_shapes) { + *shapes.add_parameters() = shape; + } + TF_ASSIGN_OR_RETURN(*shapes.mutable_result(), GetReturnValueShape()); + LayoutUtil::SetToDefaultLayout(&shapes); + *config->mutable_program_shape() = shapes.ToProto(); + auto snapshot = computation().Snapshot().ValueOrDie(); + *c.mutable_hlo_snapshot() = *snapshot; + + tensorflow::ClientSession session(root, session_target); + tensorflow::ClientSession::FeedType inputs; + inputs.insert({program, c.SerializeAsString()}); + std::vector outputs; + TF_RETURN_IF_ERROR(session.Run(inputs, {compile.handle}, &outputs)); + + TF_ASSIGN_OR_RETURN(ProgramShape program_shape, + computation().GetProgramShape()); + int64 handle = outputs[0].scalar()(); + return new CompiledXrtComputation(program_shape, handle, session_target); +} + const XlaComputation& LocalComputation::computation() const { return computation_; } @@ -314,9 +598,7 @@ string LocalComputation::GetSerializedProto() const { } StatusOr LocalComputation::GetReturnValueShape() const { - TF_ASSIGN_OR_RETURN(ProgramShape program_shape, - computation_.GetProgramShape()); - return std::move(*program_shape.mutable_result()); + return swig::GetReturnValueShape(computation_); } LocalOp::LocalOp(const XlaOp& op) : op_(op) {} @@ -343,6 +625,12 @@ LocalOp LocalComputationBuilder::Parameter(int64 parameter_number, return xla::Parameter(&builder_, parameter_number, shape, name); } +StatusOr LocalComputationBuilder::BuildWithRoot( + const LocalOp& root) { + TF_ASSIGN_OR_RETURN(XlaComputation computation, builder_.Build(root.op())); + return new LocalComputation(std::move(computation)); +} + StatusOr LocalComputationBuilder::GetShape(const LocalOp& operand) { return builder_.GetShape(operand.op()); } @@ -366,11 +654,26 @@ LocalOp LocalComputationBuilder::ConstantLiteral(const Literal& literal) { return xla::ConstantLiteral(&builder_, literal); } +LocalOp LocalComputationBuilder::Iota(PrimitiveType element_type, int64 size) { + return xla::Iota(&builder_, element_type, size); +} + +LocalOp LocalComputationBuilder::BroadcastedIota(const Shape& shape, + int64 dimension) { + return xla::Iota(&builder_, shape, dimension); +} + LocalOp LocalComputationBuilder::Broadcast( const LocalOp& operand, absl::Span broadcast_sizes) { return xla::Broadcast(operand.op(), broadcast_sizes); } +LocalOp LocalComputationBuilder::BroadcastInDim( + const LocalOp& operand, absl::Span out_dim_sizes, + absl::Span broadcast_dimensions) { + return xla::BroadcastInDim(operand.op(), out_dim_sizes, broadcast_dimensions); +} + LocalOp LocalComputationBuilder::Pad(const LocalOp& operand, const LocalOp& padding_value, const PaddingConfig& padding_config) { @@ -388,8 +691,9 @@ 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::CrossReplicaSum( + const LocalOp& operand, absl::Span replica_groups) { + return xla::CrossReplicaSum(operand.op(), replica_groups); } LocalOp LocalComputationBuilder::Slice(const LocalOp& operand, @@ -496,6 +800,21 @@ LocalOp LocalComputationBuilder::Call(const LocalComputation& local_computation, return xla::Call(&builder_, local_computation.computation(), xla_ops); } +LocalOp LocalComputationBuilder::CustomCall( + const string& call_target_name, absl::Span operands, + const Shape& shape_with_layout, + const std::vector& operand_shapes_with_layout, + const string& opaque) { + std::vector xla_ops; + xla_ops.reserve(operands.size()); + for (const auto& op : operands) { + xla_ops.push_back(op.op()); + } + return xla::CustomCallWithLayout(&builder_, call_target_name, xla_ops, + shape_with_layout, + operand_shapes_with_layout, opaque); +} + LocalOp LocalComputationBuilder::Transpose( const LocalOp& operand, absl::Span permutation) { return xla::Transpose(operand.op(), permutation); @@ -581,6 +900,43 @@ LocalOp LocalComputationBuilder::SortKeyVal(const LocalOp& keys, return xla::Sort(keys.op(), {values.op()}, dimension); } +LocalOp LocalComputationBuilder::Cholesky(const LocalOp& a) { + return xla::Cholesky(a.op()); +} + +LocalOp LocalComputationBuilder::QR(const LocalOp& a, bool full_matrices) { + XlaBuilder* builder = a.op().builder(); + return builder->ReportErrorOrReturn([&]() -> StatusOr { + TF_ASSIGN_OR_RETURN(auto qr, xla::QRDecomposition(a.op(), full_matrices)); + return xla::Tuple(builder, {qr.q, qr.r}); + }); +} + +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); +} + +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( const LocalOp& operand) { TF_ASSIGN_OR_RETURN(XlaComputation computation, @@ -677,23 +1033,29 @@ void DeleteLocalShapedBuffer(LocalShapedBuffer* local_shaped_buffer) { delete local_shaped_buffer; } +void DeleteXrtAllocation(XrtAllocation* allocation) { delete allocation; } + void DeleteCompiledLocalComputation(CompiledLocalComputation* computation) { delete computation; } +void DeleteCompiledXrtComputation(CompiledXrtComputation* computation) { + delete computation; +} + void DeleteLocalComputation(LocalComputation* computation) { delete computation; } StatusOr DestructureLocalShapedBufferTuple( LocalShapedBuffer* local_shaped_buffer) { - if (!ShapeUtil::IsTuple( - local_shaped_buffer->shaped_buffer()->on_device_shape())) { + const Shape tuple_shape = local_shaped_buffer->shape(); + + if (!tuple_shape.IsTuple()) { return InvalidArgument( "Attemped to destructure a LocalShapedBuffer that did not have a tuple " "shape; shape: %s", - ShapeUtil::HumanString( - local_shaped_buffer->shaped_buffer()->on_device_shape())); + ShapeUtil::HumanString(tuple_shape)); } DeviceMemoryAllocator* allocator = @@ -705,7 +1067,6 @@ StatusOr DestructureLocalShapedBufferTuple( int device_ordinal = tuple_buffer.device_ordinal(); ShapeTree& shape_tree = tuple_buffer.buffers(); - const Shape& tuple_shape = tuple_buffer.on_device_shape(); std::vector results; for (int64 i = 0; i < ShapeUtil::TupleElementCount(tuple_shape); ++i) { // Create a shaped buffer for this destructured tuple element. @@ -733,5 +1094,47 @@ StatusOr DestructureLocalShapedBufferTuple( return new LocalShapedBufferTuple(std::move(results)); } +StatusOr DestructureXrtAllocationTuple( + XrtAllocation* allocation, const string& session_target) { + const Shape& tuple_shape = allocation->shape(); + + if (!tuple_shape.IsTuple()) { + return InvalidArgument( + "Attemped to destructure a LocalShapedBuffer that did not have a tuple " + "shape; shape: %s", + ShapeUtil::HumanString(tuple_shape)); + } + + tensorflow::Scope root = tensorflow::Scope::NewRootScope(); + auto base_handle = tensorflow::ops::Placeholder(root, tensorflow::DT_INT64); + auto shape_index = tensorflow::ops::Placeholder(root, tensorflow::DT_INT32); + auto subtuple = tensorflow::ops::XRTSubTuple(root, base_handle, shape_index); + TF_RETURN_IF_ERROR(root.status()); + + tensorflow::ClientSession session(root, session_target); + tensorflow::ClientSession::FeedType inputs; + std::vector results; + for (int32 i = 0; i < ShapeUtil::TupleElementCount(tuple_shape); ++i) { + inputs.clear(); + inputs.insert({base_handle, allocation->handle()}); + inputs.insert({shape_index, {i}}); + std::vector outputs; + auto status = session.Run(inputs, {subtuple}, &outputs); + if (!status.ok()) { + // Clean up before returning non-ok status. + for (int j = 0; j < results.size(); ++j) { + delete results[j]; + } + return status; + } + const int64 subtuple_handle = outputs[0].scalar()(); + const Shape& subtuple_shape = + ShapeUtil::GetTupleElementShape(tuple_shape, i); + results.push_back( + new XrtAllocation(subtuple_handle, subtuple_shape, session_target)); + } + return new XrtAllocationTuple(std::move(results)); +} + } // namespace swig } // namespace xla diff --git a/tensorflow/compiler/xla/python/local_computation_builder.h b/tensorflow/compiler/xla/python/local_computation_builder.h index 43332e0abd410c08dc5a40f7de39dbc96d34a72c..6170567f9ff8f5a062f47d148900fe3676a74542 100644 --- a/tensorflow/compiler/xla/python/local_computation_builder.h +++ b/tensorflow/compiler/xla/python/local_computation_builder.h @@ -16,7 +16,14 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_PYTHON_LOCAL_COMPUTATION_BUILDER_H_ #define TENSORFLOW_COMPILER_XLA_PYTHON_LOCAL_COMPUTATION_BUILDER_H_ +#include +#include + +#include + #include "absl/types/span.h" +#include "tensorflow/cc/framework/ops.h" +#include "tensorflow/cc/framework/scope.h" #include "tensorflow/compiler/xla/client/client_library.h" #include "tensorflow/compiler/xla/client/executable_build_options.h" #include "tensorflow/compiler/xla/client/local_client.h" @@ -34,10 +41,21 @@ namespace swig { // returned. Status InitializeReplicaCount(int replica_count); +// Initializes the platform name that XLA will be initialized with (when +// first obtaining a handle to the local XLA service). If this is called after +// the handle to the local XLA service has been established, then an error is +// returned. +Status InitializePlatformName(const string& platform_name); + // Returns the replica count that is currently set, regardless of whether the // local XLA service has been instantiated yet or not. int GetReplicaCount(); +// Registers a 'fn_capsule' as a CPU custom call target. +// 'fn_capsule' is a void* pointer encapsulated in a PyCapsule object, with name +// "xla._CPU_CUSTOM_CALL_TARGET". +Status RegisterCpuCustomCallTarget(const string& name, PyObject* fn_capsule); + // Wraps the local client's infeed-transfer function. // // The default device ordinal (0) is used. @@ -54,18 +72,19 @@ Status TransferToInfeedLocalReplica(const Literal& literal, int replica_number); StatusOr TransferFromOutfeedLocalReplica(const Shape& shape, int replica_number); -// Wraps a ScopedShapedBuffer produced by copying a literal "to -// device," i.e. copying a literal to a scoped buffer via the local -// client. +// Represents a reference to literals that live in a device-allocated buffer via +// XLA. Specifically, wraps a ScopedShapedBuffer produced by transferring a +// literal to device via the local client. class LocalShapedBuffer { public: static StatusOr FromLiteral( - const Literal& argument, const absl::optional& shape_with_layout); + const Literal& argument, const absl::optional& shape_with_layout, + int replica_number); LocalShapedBuffer(ScopedShapedBuffer shaped_buffer); - const ScopedShapedBuffer* shaped_buffer() const; - StatusOr ToLiteral() const; + const Shape& shape() const; + const ScopedShapedBuffer* shaped_buffer() const; // Transfers ownership of the encapsulated ShapedBuffer to the caller, // analogous to std::unique_ptr::release(). @@ -92,7 +111,7 @@ class LocalShapedBufferTuple { StatusOr Release(int i); // Returns the number of elements in the destructured tuple. - int size() const; + int64 size() const; private: std::vector elements_; @@ -103,31 +122,103 @@ class LocalShapedBufferTuple { StatusOr DestructureLocalShapedBufferTuple( LocalShapedBuffer* local_shaped_buffer); -// Wraps a LocalExecutable produced by compiling a -// LocalComputation. The Execute method forwards to that of the -// underlying LocalExecutable, and additionally handles tranferring -// arguments and return values in and back out of the client library's -// local client. This class is intended to be made available to Python -// via SWIG. +// Represents a reference to literals that live in a device-allocated buffer via +// XRT. Specifically, wraps an int64 handle produced by running the allocation +// graph, and an XLA shape to track the referent's shape. +class XrtAllocation { + public: + // Accepts a `session_target` argument, used in constructing the + // `tensorflow::ClientSession` instance in which allocation and deallocation + // graphs are run. + static StatusOr FromLiteral(const Literal& argument, + const string& session_target); + + XrtAllocation(int64 handle, Shape shape, const string& session_target); + ~XrtAllocation(); + StatusOr ToLiteral() const; + const Shape& shape() const; + const int64 handle() const; + + private: + const int64 handle_; + const Shape shape_; + const string session_target_; +}; + +// Result of a tuple destructuring operation on an XrtAllocation. +class XrtAllocationTuple { + public: + // Note: any XrtAllocation elements that are not Release()'d will be + // deallocated in the destructor. + explicit XrtAllocationTuple(std::vector elements); + + ~XrtAllocationTuple(); + + // Releases the ith element to the caller. Further attempts to release the ith + // element will return an invalid argument error. + StatusOr Release(int i); + + // Returns the number of elements in the destructured tuple. + int64 size() const; + + private: + std::vector elements_; +}; + +// Destructures a tuple-valued XrtAllocation into its constitutent elements +// in XrtAllocationTuple form. +// +// Accepts a `session_target` argument, used in constructing the +// `tensorflow::ClientSession` instance in which the sub-tupling graph is run, +// and passed along in constructing each constituent XrtAllocation. +StatusOr DestructureXrtAllocationTuple( + XrtAllocation* allocation, const string& session_target); + +// Represents a compiled computation that can be executed given handles to +// device-allocated literals. Specifically, wraps an XLA LocalExecutable. class CompiledLocalComputation { public: CompiledLocalComputation(std::unique_ptr executable); - // Execute the computation with the given argument literals, and - // with optionally-specified argument layouts. The literals will be - // re-laid out according to the corresponding elements of - // shapes_with_layout. - StatusOr Execute( - const std::vector& arguments, - const std::vector >& shapes_with_layout); + int num_replicas() const { + return executable_->build_options().num_replicas(); + } - LocalShapedBuffer* ExecuteWithShapedBuffers( + StatusOr Execute( absl::Span argument_handles); + // Execute on many replicas. Takes a sequence of argument lists (one argument + // list per replica) and returns a tuple of results (one result per replica). + // The number of argument lists must be equal to the replica count. + StatusOr ExecutePerReplica( + absl::Span > argument_handles); + private: std::unique_ptr executable_; }; +// Represents a compiled computation that can be executed given handles to +// device-allocated literals. Specifically, wraps an XRT computation handle. +class CompiledXrtComputation { + public: + // Accepts a `session_target` argument, used in constructing the + // `tensorflow::ClientSession` instance in which the execution graph is run. + CompiledXrtComputation(const ProgramShape& program_shape, int64 handle, + const string& session_target); + ~CompiledXrtComputation(); + + StatusOr Execute( + absl::Span argument_handles); + + const ProgramShape& program_shape() const; + int64 handle() const; + + private: + const ProgramShape program_shape_; + const int64 handle_; + const string session_target_; +}; + // Wraps a XlaComputation produced by a LocalComputationBuilder. The // Compile method compiles the computation to a (local) executable via // the client library's local client. This class is intended to be @@ -140,6 +231,11 @@ class LocalComputation { const std::vector& argument_shapes, const ExecutableBuildOptions* build_options); + // Accepts a `session_target` argument, used in constructing the + // `tensorflow::ClientSession` instance in which the compilation graph is run. + StatusOr CompileForXrt( + const std::vector& argument_shapes, const string& session_target); + const XlaComputation& computation() const; // Returns the HloModuleProto contained in the XlaComputation in the @@ -183,6 +279,9 @@ class LocalComputationBuilder { // Returns an owned LocalComputation to the caller on success. StatusOr Build(); + // Returns an owned LocalComputation to the caller on success with given root. + StatusOr BuildWithRoot(const LocalOp& root); + LocalOp Parameter(int64 parameter_number, const Shape& shape, const string& name); @@ -198,9 +297,17 @@ class LocalComputationBuilder { LocalOp ConstantLiteral(const Literal& literal); + LocalOp Iota(PrimitiveType element_type, int64 size); + + LocalOp BroadcastedIota(const Shape& shape, int64 dimension); + LocalOp Broadcast(const LocalOp& operand, absl::Span broadcast_sizes); + LocalOp BroadcastInDim(const LocalOp& operand, + absl::Span out_dim_sizes, + absl::Span broadcast_dimensions); + LocalOp Pad(const LocalOp& operand, const LocalOp& padding_value, const PaddingConfig& padding_config); @@ -209,7 +316,8 @@ class LocalComputationBuilder { LocalOp Collapse(const LocalOp& operand, absl::Span dimensions); - LocalOp CrossReplicaSum(const LocalOp& operand); + LocalOp CrossReplicaSum(const LocalOp& operand, + absl::Span replica_groups); LocalOp Slice(const LocalOp& operand, absl::Span start_indices, absl::Span limit_indices, @@ -260,6 +368,12 @@ class LocalComputationBuilder { LocalOp Call(const LocalComputation& local_computation, absl::Span operands); + LocalOp CustomCall(const string& call_target_name, + absl::Span operands, + const Shape& shape_with_layout, + const std::vector& operand_shapes_with_layout, + const string& opaque); + LocalOp Transpose(const LocalOp& operand, absl::Span permutation); @@ -302,6 +416,22 @@ class LocalComputationBuilder { LocalOp SortKeyVal(const LocalOp& keys, const LocalOp& values, int64 dimension); + LocalOp QR(const LocalOp& a, bool full_matrices); + + LocalOp Cholesky(const LocalOp& a); + + LocalOp TriangularSolve(const LocalOp& a, const LocalOp& b, bool left_side, + bool lower, bool transpose_a, bool conjugate_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); #define _FORWARD(method_name, return_sig, args_sig) \ @@ -391,7 +521,9 @@ class LocalComputationBuilder { // Functions for freeing resources from the Python side. void DeleteLocalShapedBuffer(LocalShapedBuffer* local_shaped_buffer); +void DeleteXrtAllocation(XrtAllocation* allocation); void DeleteCompiledLocalComputation(CompiledLocalComputation* computation); +void DeleteCompiledXrtComputation(CompiledXrtComputation* computation); void DeleteLocalComputation(LocalComputation* computation); } // namespace swig diff --git a/tensorflow/compiler/xla/python/local_computation_builder.i b/tensorflow/compiler/xla/python/local_computation_builder.i index 521490e76c138553c5cc6895412eadb35a939881..6a85ed62dea3dbdbb25a990e6d774a0152439673 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. @@ -176,6 +212,81 @@ bool HandleStringAttribute(PyObject* o, tensorflow::ImportNumpy(); %} +// Basic types + +%typemap(out) StatusOr { + if ($1.ok()) { + $result = PyBool_FromLong($1.ConsumeValueOrDie()); + } else { + PyErr_SetString(PyExc_RuntimeError, $1.status().ToString().c_str()); + SWIG_fail; + } +} + +%typemap(out) Status { + if (!$1.ok()) { + PyErr_SetString( + PyExc_RuntimeError, $1.ToString().c_str()); + SWIG_fail; + } + Py_INCREF(Py_None); + $result = Py_None; +} + +%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.resize(size); + for (int i = 0; i < size; ++i) { + PyObject* o = PySequence_GetItem($input, i); + PyObject* py_int = numpy::PyNumberToPyInt(o); + if (!py_int) { + PyErr_SetString( + PyExc_TypeError, + "Argument sequence element cannot be converted to int"); + Py_DECREF(o); + SWIG_fail; + } + temps[i] = numpy::PyIntOrPyLongToLong(py_int); + if (temps[i] == -1 && PyErr_Occurred()) { + Py_DECREF(py_int); + Py_DECREF(o); + SWIG_fail; + } + Py_DECREF(py_int); + Py_DECREF(o); + } + $1 = temps; +} + +// Computation builder types + +%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); + for (int i = 0; i < size; ++i) { + PyObject* o = PySequence_GetItem($input, i); + LocalOp* op; + if ((SWIG_ConvertPtr(o, (void**)&op, $descriptor(xla::swig::LocalOp*), + SWIG_POINTER_EXCEPTION)) == -1) { + SWIG_fail; + } + temps.push_back(*op); + Py_DECREF(o); + } + $1 = temps; +} + +// Computation and buffer/allocation types + %typemap(out) StatusOr { if ($1.ok()) { auto* value = $1.ValueOrDie(); @@ -189,12 +300,12 @@ tensorflow::ImportNumpy(); } } -%typemap(out) StatusOr { +%typemap(out) StatusOr { if ($1.ok()) { auto* value = $1.ValueOrDie(); { auto* $1 = value; - $typemap(out, xla::swig::LocalShapedBuffer*) + $typemap(out, xla::swig::CompiledXrtComputation*) } } else { PyErr_SetString(PyExc_RuntimeError, $1.status().ToString().c_str()); @@ -202,12 +313,12 @@ tensorflow::ImportNumpy(); } } -%typemap(out) StatusOr { +%typemap(out) StatusOr { if ($1.ok()) { auto* value = $1.ValueOrDie(); { auto* $1 = value; - $typemap(out, xla::swig::LocalShapedBufferTuple*) + $typemap(out, xla::swig::LocalShapedBuffer*) } } else { PyErr_SetString(PyExc_RuntimeError, $1.status().ToString().c_str()); @@ -215,23 +326,25 @@ tensorflow::ImportNumpy(); } } - -%typemap(out) StatusOr { +%typemap(out) StatusOr { if ($1.ok()) { - Literal value = $1.ConsumeValueOrDie(); - $result = numpy::PyObjectFromXlaLiteral(*value); + auto* value = $1.ValueOrDie(); + { + auto* $1 = value; + $typemap(out, xla::swig::LocalShapedBufferTuple*) + } } else { PyErr_SetString(PyExc_RuntimeError, $1.status().ToString().c_str()); SWIG_fail; } } -%typemap(out) StatusOr { +%typemap(out) StatusOr { if ($1.ok()) { auto* value = $1.ValueOrDie(); { auto* $1 = value; - $typemap(out, xla::swig::LocalComputation*) + $typemap(out, xla::swig::XrtAllocation*) } } else { PyErr_SetString(PyExc_RuntimeError, $1.status().ToString().c_str()); @@ -239,92 +352,86 @@ tensorflow::ImportNumpy(); } } -%typemap(out) StatusOr { +%typemap(out) StatusOr { if ($1.ok()) { - $result = numpy::PyShapeInfoFromXlaShape($1.ConsumeValueOrDie()); + auto* value = $1.ValueOrDie(); + { + auto* $1 = value; + $typemap(out, xla::swig::XrtAllocationTuple*) + } } else { PyErr_SetString(PyExc_RuntimeError, $1.status().ToString().c_str()); SWIG_fail; } } -%typemap(out) StatusOr { +%typemap(out) StatusOr { if ($1.ok()) { - $result = PyBool_FromLong($1.ConsumeValueOrDie()); + auto* value = $1.ValueOrDie(); + { + auto* $1 = value; + $typemap(out, xla::swig::LocalComputation*) + } } else { PyErr_SetString(PyExc_RuntimeError, $1.status().ToString().c_str()); SWIG_fail; } } -%typemap(out) Status { - if (!$1.ok()) { - PyErr_SetString( - PyExc_RuntimeError, $1.ToString().c_str()); - SWIG_fail; - } - Py_INCREF(Py_None); - $result = Py_None; -} - -// Span - -%typemap(in) absl::Span - (std::vector temps) { +%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.resize(size); + temps.reserve(size); for (int i = 0; i < size; ++i) { PyObject* o = PySequence_GetItem($input, i); - PyObject* py_int = numpy::PyNumberToPyInt(o); - if (!py_int) { - PyErr_SetString( - PyExc_TypeError, - "Argument sequence element cannot be converted to int"); - Py_DECREF(o); - SWIG_fail; - } - temps[i] = numpy::PyIntOrPyLongToLong(py_int); - if (temps[i] == -1 && PyErr_Occurred()) { - Py_DECREF(py_int); - Py_DECREF(o); + LocalShapedBuffer* lsbp; + if ((SWIG_ConvertPtr(o, (void**) &lsbp, $descriptor(xla::swig::LocalShapedBuffer*), + SWIG_POINTER_EXCEPTION)) == -1) { SWIG_fail; } - Py_DECREF(py_int); + temps.push_back(lsbp); Py_DECREF(o); } $1 = temps; } -// Span - -%typemap(in) absl::Span( - std::vector temps) { +%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); - LocalOp* op; - if ((SWIG_ConvertPtr(o, (void**)&op, $descriptor(xla::swig::LocalOp*), - SWIG_POINTER_EXCEPTION)) == -1) { - SWIG_fail; + std::vector vec; + const int vec_size = PySequence_Size(o); + vec.reserve(vec_size); + for (int j = 0; j < vec_size; ++j) { + PyObject* vec_elt = PySequence_GetItem(o, j); + LocalShapedBuffer* lsbp; + if ((SWIG_ConvertPtr(vec_elt, (void**) &lsbp, $descriptor(xla::swig::LocalShapedBuffer*), + SWIG_POINTER_EXCEPTION)) == -1) { + Py_DECREF(vec_elt); + Py_DECREF(o); + SWIG_fail; + } + vec.push_back(lsbp); + Py_DECREF(vec_elt); } - temps.push_back(*op); + temps.push_back(vec); Py_DECREF(o); } $1 = temps; } -// LocalShapedBuffer* - -%typemap(in) absl::Span - (std::vector temps) { +%typemap(in) absl::Span + (std::vector temps) { if (!PySequence_Check($input)) { PyErr_SetString(PyExc_TypeError, "Argument is not a sequence"); SWIG_fail; @@ -333,12 +440,12 @@ tensorflow::ImportNumpy(); temps.reserve(size); for (int i = 0; i < size; ++i) { PyObject* o = PySequence_GetItem($input, i); - LocalShapedBuffer* lsbp; - if ((SWIG_ConvertPtr(o, (void**) &lsbp, $descriptor(xla::swig::LocalShapedBuffer*), + XrtAllocation* xrta; + if ((SWIG_ConvertPtr(o, (void**) &xrta, $descriptor(xla::swig::XrtAllocation*), SWIG_POINTER_EXCEPTION)) == -1) { SWIG_fail; } - temps.push_back(lsbp); + temps.push_back(xrta); Py_DECREF(o); } $1 = temps; @@ -346,6 +453,16 @@ 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()) { @@ -401,6 +518,19 @@ tensorflow::ImportNumpy(); // Shape +%typemap(out) const Shape& { + $result = numpy::PyShapeInfoFromXlaShape(*$1); +} + +%typemap(out) StatusOr { + if ($1.ok()) { + $result = numpy::PyShapeInfoFromXlaShape($1.ConsumeValueOrDie()); + } else { + PyErr_SetString(PyExc_RuntimeError, $1.status().ToString().c_str()); + SWIG_fail; + } +} + %typemap(in) const Shape& (Shape temp) { StatusOr statusor = numpy::XlaShapeFromPyShape($input); if (!statusor.ok()) { @@ -563,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) { + if (!HandleRepeatedInt64Attribute( + $input, "lhs_contracting_dimensions", + dimension_numbers.mutable_lhs_contracting_dimensions())) { SWIG_fail; } - - length = PySequence_Size(lhs_contracting_dimensions); - if (length == -1) { - Py_DECREF(lhs_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(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_batch_dimensions", + dimension_numbers.mutable_lhs_batch_dimensions())) { SWIG_fail; } - - length = PySequence_Size(rhs_contracting_dimensions); - if (length == -1) { - Py_DECREF(rhs_contracting_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(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) { - SWIG_fail; - } - - length = PySequence_Size(lhs_batch_dimensions); - if (length == -1) { - Py_DECREF(lhs_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; } @@ -766,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* @@ -858,22 +905,22 @@ tensorflow::ImportNumpy(); $1 = NULL; } else { if (!HandleStringAttribute($input, "generate_hlo_graph", [&](string s) { - build_options.set_generate_hlo_graph(std::move(s)); + build_options.mutable_debug_options()->set_xla_generate_hlo_graph(std::move(s)); })) { return nullptr; } if (!HandleStringAttribute($input, "dump_optimized_hlo_proto_to", [&](string s) { - build_options.set_dump_optimized_hlo_proto_to(std::move(s)); + build_options.mutable_debug_options()->set_xla_dump_optimized_hlo_proto_to(std::move(s)); })) { return nullptr; } if (!HandleStringAttribute($input, "dump_unoptimized_hlo_proto_to", [&](string s) { - build_options.set_dump_unoptimized_hlo_proto_to(std::move(s)); + build_options.mutable_debug_options()->set_xla_dump_unoptimized_hlo_proto_to(std::move(s)); })) { return nullptr; } if (!HandleStringAttribute($input, "dump_per_pass_hlo_proto_to", [&](string s) { - build_options.set_dump_per_pass_hlo_proto_to(std::move(s)); + build_options.mutable_debug_options()->set_xla_dump_per_pass_hlo_proto_to(std::move(s)); })) { return nullptr; } @@ -887,7 +934,7 @@ tensorflow::ImportNumpy(); PyErr_SetString(PyExc_TypeError, "ExecutableBuildOptions.hlo_profile must be a bool or None."); SWIG_fail; } - build_options.set_hlo_profile(o == Py_True); + build_options.mutable_debug_options()->set_xla_hlo_profile(o == Py_True); } Py_DECREF(o); @@ -906,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; } } @@ -914,27 +967,41 @@ tensorflow::ImportNumpy(); %unignore xla; %unignore xla::swig; %unignore xla::swig::InitializeReplicaCount; +%unignore xla::swig::InitializePlatformName; %unignore xla::swig::GetReplicaCount; +%unignore xla::swig::RegisterCpuCustomCallTarget; %unignore xla::swig::TransferToInfeedLocal; %unignore xla::swig::TransferToInfeedLocalReplica; %unignore xla::swig::TransferFromOutfeedLocalReplica; %unignore xla::swig::LocalShapedBuffer; %unignore xla::swig::LocalShapedBuffer::FromLiteral; %unignore xla::swig::LocalShapedBuffer::ToLiteral; +%unignore xla::swig::LocalShapedBuffer::shape; %unignore xla::swig::LocalShapedBufferTuple; %unignore xla::swig::LocalShapedBufferTuple::Release; %unignore xla::swig::LocalShapedBufferTuple::size; +%unignore xla::swig::XrtAllocation; +%unignore xla::swig::XrtAllocation::FromLiteral; +%unignore xla::swig::XrtAllocation::ToLiteral; +%unignore xla::swig::XrtAllocation::shape; +%unignore xla::swig::XrtAllocationTuple; +%unignore xla::swig::XrtAllocationTuple::Release; +%unignore xla::swig::XrtAllocationTuple::size; %unignore xla::swig::CompiledLocalComputation; %unignore xla::swig::CompiledLocalComputation::Execute; -%unignore xla::swig::CompiledLocalComputation::ExecuteWithShapedBuffers; +%unignore xla::swig::CompiledLocalComputation::ExecutePerReplica; +%unignore xla::swig::CompiledXrtComputation; +%unignore xla::swig::CompiledXrtComputation::Execute; %unignore xla::swig::LocalComputation; %unignore xla::swig::LocalComputation::Compile; +%unignore xla::swig::LocalComputation::CompileForXrt; %unignore xla::swig::LocalComputation::GetReturnValueShape; %unignore xla::swig::LocalComputation::GetSerializedProto; %unignore xla::swig::LocalOp; %unignore xla::swig::LocalComputationBuilder; %unignore xla::swig::LocalComputationBuilder::LocalComputationBuilder; %unignore xla::swig::LocalComputationBuilder::Build; +%unignore xla::swig::LocalComputationBuilder::BuildWithRoot; %unignore xla::swig::LocalComputationBuilder::SetOpMetadata; %unignore xla::swig::LocalComputationBuilder::ClearOpMetadata; %unignore xla::swig::LocalComputationBuilder::Parameter; @@ -944,7 +1011,10 @@ tensorflow::ImportNumpy(); %unignore xla::swig::LocalComputationBuilder::Outfeed; %unignore xla::swig::LocalComputationBuilder::ConstantLiteral; %unignore xla::swig::LocalComputationBuilder::ConstantR0; +%unignore xla::swig::LocalComputationBuilder::Iota; +%unignore xla::swig::LocalComputationBuilder::BroadcastedIota; %unignore xla::swig::LocalComputationBuilder::Broadcast; +%unignore xla::swig::LocalComputationBuilder::BroadcastInDim; %unignore xla::swig::LocalComputationBuilder::Pad; %unignore xla::swig::LocalComputationBuilder::Reshape; %unignore xla::swig::LocalComputationBuilder::Collapse; @@ -1036,10 +1106,19 @@ tensorflow::ImportNumpy(); %unignore xla::swig::LocalComputationBuilder::Imag; %unignore xla::swig::LocalComputationBuilder::Conj; %unignore xla::swig::LocalComputationBuilder::Complex; +%unignore xla::swig::LocalComputationBuilder::Cholesky; +%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; %unignore xla::swig::DeleteLocalShapedBuffer; -%unignore xla::swig::DeleteLocalComputation; +%unignore xla::swig::DeleteXrtAllocation; %unignore xla::swig::DeleteCompiledLocalComputation; +%unignore xla::swig::DeleteCompiledXrtComputation; %thread; %include "tensorflow/compiler/xla/python/local_computation_builder.h" diff --git a/tensorflow/compiler/xla/python/numpy_bridge.cc b/tensorflow/compiler/xla/python/numpy_bridge.cc index b0aa024c7474cf8e6934432b2f364be464714999..52c5c621f7294c5da341879d15b77559fe870551 100644 --- a/tensorflow/compiler/xla/python/numpy_bridge.cc +++ b/tensorflow/compiler/xla/python/numpy_bridge.cc @@ -54,6 +54,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 +91,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 +115,7 @@ bool NumpyTypeIsValid(int np_type) { case NPY_FLOAT32: case NPY_FLOAT64: case NPY_COMPLEX64: + case NPY_COMPLEX128: case NPY_OBJECT: return true; default: @@ -123,7 +128,7 @@ PyObject* PyShapeInfoFromXlaShape(const Shape& shape) { PyArray_Descr* np_dtype = PyArray_DescrFromType(np_typenum); PyObject* dimensions; - if (ShapeUtil::IsTuple(shape)) { + if (shape.IsTuple()) { int num_elements = ShapeUtil::TupleElementCount(shape); dimensions = PyTuple_New(ShapeUtil::TupleElementCount(shape)); for (int i = 0; i < num_elements; ++i) { @@ -132,7 +137,7 @@ PyObject* PyShapeInfoFromXlaShape(const Shape& shape) { PyShapeInfoFromXlaShape(ShapeUtil::GetTupleElementShape(shape, i))); } } else { - int rank = ShapeUtil::Rank(shape); + int rank = shape.rank(); dimensions = PyTuple_New(rank); for (int i = 0; i < rank; ++i) { PyTuple_SET_ITEM(dimensions, i, @@ -345,7 +350,7 @@ StatusOr OpMetadataFromPyObject(PyObject* o) { } PyObject* PyObjectFromXlaLiteral(const LiteralSlice& literal) { - if (ShapeUtil::IsTuple(literal.shape())) { + if (literal.shape().IsTuple()) { int num_elements = ShapeUtil::TupleElementCount(literal.shape()); PyObject* tuple = PyTuple_New(num_elements); for (int i = 0; i < num_elements; i++) { @@ -354,7 +359,7 @@ PyObject* PyObjectFromXlaLiteral(const LiteralSlice& literal) { } return tuple; } else { - int rank = ShapeUtil::Rank(literal.shape()); + int rank = literal.shape().rank(); std::vector dimensions(rank); // NOLINT - PyArray requires a long* for (int i = 0; i < rank; i++) { dimensions[i] = ShapeUtil::GetDimension(literal.shape(), i); @@ -430,6 +435,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); @@ -470,6 +478,9 @@ 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; } 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..bce6c1acf8a1cc0005ca93e0466c5a0e29d880de --- /dev/null +++ b/tensorflow/compiler/xla/python/pywrap_xla_exported_symbols.lds @@ -0,0 +1 @@ +_PyInit__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 f8197488fb3bacb312cc7fbf149b773851992b8a..8964b158292371d662368cfb0b644667985f719e 100644 --- a/tensorflow/compiler/xla/python/xla_client.py +++ b/tensorflow/compiler/xla/python/xla_client.py @@ -26,6 +26,9 @@ import os import numpy as np +import six +from six.moves import xrange + from tensorflow.compiler.xla import xla_data_pb2 from tensorflow.compiler.xla.python import pywrap_xla as c_api from tensorflow.compiler.xla.service import hlo_pb2 @@ -46,6 +49,15 @@ _OP_METADATA_FIELDS = [ OpMetadata = collections.namedtuple('OpMetadata', _OP_METADATA_FIELDS) +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 OpMetadataToProto(pyobj): proto = xla_data_pb2.OpMetadata() for field in _OP_METADATA_FIELDS: @@ -66,6 +78,13 @@ def CurrentSourceInfoMetadata(op_type=None, op_name=None, skip_frames=1): source_line=lineno) +def _maybe_encode_string(s): + if six.PY3: + return s.encode('utf-8') + else: + return s + + class PaddingType(enum.Enum): VALID = 1 SAME = 2 @@ -73,16 +92,32 @@ class PaddingType(enum.Enum): def _convert_padding_type_to_pad_values(padding_type, lhs_dims, rhs_dims, window_strides): - """Maps PaddingType (VALID or SAME) to pad values (list of pairs of ints).""" + """Maps PaddingType or string to pad values (list of pairs of ints).""" + if not isinstance(padding_type, (str, PaddingType)): + msg = 'padding_type must be str or PaddingType, got {}.' + raise TypeError(msg.format(type(padding_type))) + + if isinstance(padding_type, str): + if padding_type.upper() == 'VALID': + padding_type = PaddingType.VALID + elif padding_type.upper() == 'SAME': + padding_type = PaddingType.SAME + else: + msg = 'Unknown padding type string: expected "VALID" or "SAME", got {}.' + raise ValueError(msg.format(padding_type)) + if padding_type == PaddingType.VALID: return [(0, 0)] * len(window_strides) - - out_shape = np.ceil(np.true_divide(lhs_dims, window_strides)).astype(int) - pad_sizes = [max((out_size - 1) * stride + filter_size - in_size, 0) - for out_size, stride, filter_size, in_size - in zip(out_shape, window_strides, rhs_dims, lhs_dims)] - return [(pad_size // 2, pad_size - pad_size // 2) - for pad_size in pad_sizes] + elif padding_type == PaddingType.SAME: + out_shape = np.ceil(np.true_divide(lhs_dims, window_strides)).astype(int) + pad_sizes = [max((out_size - 1) * stride + filter_size - in_size, 0) + for out_size, stride, filter_size, in_size + in zip(out_shape, window_strides, rhs_dims, lhs_dims)] + return [(pad_size // 2, pad_size - pad_size // 2) + for pad_size in pad_sizes] + else: + msg = 'Unexpected PaddingType value: {}' + raise ValueError(msg.format(padding_type)) _UNARY_OPS = [ @@ -164,6 +199,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), } @@ -187,38 +223,66 @@ class LocalBuffer(object): means the referent is in device memory. """ - def __init__(self, c_local_shaped_buffer): - self.c_local_shaped_buffer = c_local_shaped_buffer - self._delete = c_api.DeleteLocalShapedBuffer + def __init__(self, c_buffer, backend, replica): + 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, layout_fn=None): + def from_pyval(pyval, replica=0, backend=XLA_LOCAL_BACKEND): + """Allocate and copy to XLA the given python value.""" pyval = require_numpy_array_layout(pyval) - if layout_fn: - shape = Shape.from_pyval(pyval) - shape = shape.map_leaves(layout_fn) + num_replicas = get_replica_count() + if not 0 <= replica < num_replicas: + 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: - shape = None - return LocalBuffer(c_api.LocalShapedBuffer.FromLiteral(pyval, shape)) + cbuf = c_api.LocalShapedBuffer.FromLiteral(pyval, None, replica) + return LocalBuffer(cbuf, backend, replica) def to_py(self): - return self.c_local_shaped_buffer.ToLiteral() + return self.c_buffer.ToLiteral() + + def shape(self): + return _wrap_shape(self.c_buffer.shape()) + + def replica(self): + return self._replica def delete(self): - if self.c_local_shaped_buffer is not None: - self._delete(self.c_local_shaped_buffer) - self.c_local_shaped_buffer = None + if self.c_buffer is not None: + self._delete(self.c_buffer) + self.c_buffer = None def destructure(self): - assert self.c_local_shaped_buffer is not None - result = c_api.DestructureLocalShapedBufferTuple(self.c_local_shaped_buffer) - self.c_local_shaped_buffer = None + """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) + self.delete() size = result.size() - destructured = tuple(LocalBuffer(result.Release(i)) for i in xrange(size)) + destructured = tuple( + LocalBuffer( + result.Release(i), replica=self._replica, backend=self._backend) + for i in xrange(size)) return destructured def is_deleted(self): - return self.c_local_shaped_buffer is None + return self.c_buffer is None def __del__(self): self.delete() @@ -283,6 +347,9 @@ class Shape(object): def __ne__(self, other): return not self == other + def __hash__(self): + return hash((self._dtype, self._dimensions, self._minor_to_major)) + def __repr__(self): return ('xla_client.Shape(_dtype={!r}, _dimensions={!r}, ' '_is_tuple={!r}, _minor_to_major={!r})').format( @@ -349,7 +416,7 @@ class Shape(object): assert mtm is None, self if mtm is not None: assert self.rank() == len(mtm), self - assert sorted(mtm) == range(len(mtm)), self + assert sorted(mtm) == list(range(len(mtm))), self def update_minor_to_major(self, minor_to_major): if not self.is_array(): @@ -392,6 +459,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): @@ -436,26 +504,37 @@ class LocalComputation(object): ComputationBuilder methods. """ - def __init__(self, c_local_computation, is_compiled): - self.c_local_computation = c_local_computation - self.is_compiled = is_compiled + def __init__(self, c_computation, is_compiled, backend=XLA_LOCAL_BACKEND): + self._c_computation = c_computation + self._backend = backend + self._is_compiled = is_compiled # Ensure a reference to C-based destructor for use in __del__. if is_compiled: - assert isinstance(c_local_computation, c_api.CompiledLocalComputation) - self._delete = c_api.DeleteCompiledLocalComputation + 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_local_computation, c_api.LocalComputation) + assert isinstance(c_computation, c_api.LocalComputation) self._delete = c_api.DeleteLocalComputation + @property + def computation(self): + if self._is_compiled: + raise ValueError( + 'Attempt to read the XLA computation of a compiled LocalComputation.') + return self._c_computation + def GetProto(self): """Get the HloModuleProto proto object in this local computation. Returns: An HloModuleProto proto object that has the whole-graph information. """ - - serialized = self.c_local_computation.GetSerializedProto() + serialized = self.computation.GetSerializedProto() proto = hlo_pb2.HloModuleProto.FromString(serialized) return proto @@ -480,10 +559,10 @@ class LocalComputation(object): Returns: A newly *compiled* local computation instance. """ - if self.is_compiled: + if self._is_compiled: raise ValueError('Attempt to compile a compiled local XLA computation.') - result_shape = _wrap_shape(self.c_local_computation.GetReturnValueShape()) + result_shape = _wrap_shape(self.computation.GetReturnValueShape()) if layout_fn: argument_shapes = [ @@ -491,11 +570,16 @@ class LocalComputation(object): ] result_shape = result_shape.map_leaves(layout_fn) + argument_shapes = list(argument_shapes) + compile_options = compile_options or CompileOptions() compile_options.result_shape = result_shape - return LocalComputation( - self.c_local_computation.Compile(argument_shapes, compile_options), - is_compiled=True) + 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) + return LocalComputation(c, is_compiled=True, backend=self._backend) def CompileWithExampleArguments(self, arguments=(), @@ -506,33 +590,89 @@ class LocalComputation(object): compile_options=compile_options, layout_fn=layout_fn) - def Execute(self, arguments=(), layout_fn=None): - """Execute with Python values as arguments and return value.""" - if not self.is_compiled: + def GetReturnValueShape(self): + return _wrap_shape(self._c_computation.GetReturnValueShape()) + + def Execute(self, arguments=(), check_for_deleted_args=True): + """Execute on one replica with LocalBuffer arguments and return value.""" + 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) + return LocalBuffer(output_buffer, backend=self._backend, replica=0) + + def ExecutePerReplica(self, arguments=None): + """Execute on many replicas with LocalBuffer arguments and return value. + + Args: + arguments: A sequence of sequences of LocalBuffers. The i'th inner + sequence comprises the arguments for execution on the i'th replica. + + Returns: + A list of the computation's outputs on each replica, as a LocalBuffer. If + a shallow sequence of arguments was passed in for `arguments`, then the + sole, zero'th replica's output is returned instead, as a LocalBuffer. + """ + if not self._is_compiled: raise ValueError('Cannot execute an uncompiled local XLA computation.') - argument_shapes = [Shape.from_pyval(arg) for arg in arguments] - if layout_fn: - argument_shapes = [ - shape.map_leaves(layout_fn) for shape in argument_shapes - ] + if arguments is None: + arguments = ((),) * get_replica_count() else: - argument_shapes = [None for shape in argument_shapes] - arguments = tuple(map(require_numpy_array_layout, arguments)) - return self.c_local_computation.Execute(arguments, argument_shapes) + arguments = [list(replica_args) for replica_args in arguments] + + # Check arguments + for replica, replica_args in enumerate(arguments): + for arg in replica_args: + if arg.is_deleted(): + raise ValueError('Executing with deleted local buffer argument') + if arg.replica() != replica: + raise ValueError( + 'Executing on replica {} with argument from replica {}'.format( + replica, arg.replica())) + + # Pull out argument buffer handles + stripped_args = [ + [arg.c_buffer for arg in replica_args] for replica_args in arguments + ] + + # 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)] - def ExecuteWithLocalBuffers(self, arguments=()): - """Execute with LocalBuffer arguments and return value.""" - if not self.is_compiled: - raise ValueError('Cannot execute an uncompiled local XLA computation.') - arguments = tuple(arguments) - if any(arg.is_deleted() for arg in arguments): - raise ValueError('Executing with deleted local buffer argument') - return LocalBuffer( - self.c_local_computation.ExecuteWithShapedBuffers( - [arg.c_local_shaped_buffer for arg in arguments])) + # Wrap output handles in LocalBuffer instances + return tuple( + LocalBuffer(output_buffer, backend=self._backend, replica=replica) + for replica, output_buffer in enumerate(output_buffers)) + + def ExecuteWithPythonValues(self, arguments=()): + """Execute on one replica with Python values as arguments and output.""" + + def put(arg): + return LocalBuffer.from_pyval(arg, backend=self._backend) + + arguments = [put(arg) for arg in arguments] + return self.Execute(arguments).to_py() + + def ExecuteWithPythonValuesPerReplica(self, arguments): + """Execute on many replicas with Python values as arguments and output.""" + + def put(arg, replica): + return LocalBuffer.from_pyval(arg, replica, backend=self._backend) + + arguments = [[put(arg, replica) + for arg in replica_args] + for replica, replica_args in enumerate(arguments)] + return [out.to_py() for out in self.ExecutePerReplica(arguments)] def __del__(self): - self._delete(self.c_local_computation) + self._delete(self._c_computation) class ComputationBuilder(object): @@ -554,8 +694,13 @@ class ComputationBuilder(object): self._client = c_api.LocalComputationBuilder(name.encode('utf8')) self._parameter_numbering = itertools.count() - def Build(self): - return LocalComputation(self._client.Build(), is_compiled=False) + def Build(self, root=None, backend=XLA_LOCAL_BACKEND): + if root is not None: + return LocalComputation( + self._client.BuildWithRoot(root), is_compiled=False, backend=backend) + else: + return LocalComputation( + self._client.Build(), is_compiled=False, backend=backend) def SetOpMetadata(self, op_metadata): """Set metadata for operations that are about to be enqueued.""" @@ -688,6 +833,33 @@ class ComputationBuilder(object): return self.ParameterWithShape( Shape.from_pyval(value), name=name, parameter_num=parameter_num) + def Iota(self, dtype, size): + """Enqueues an iota constant onto the computation. + + Args: + dtype: expected numpy dtype of the output. + size: integer, the number of elements in the array. + + Returns: + A LocalOp representing the added iota constant. + """ + element_type = DTYPE_TO_XLA_ELEMENT_TYPE[str(np.dtype(dtype))] + return self._client.Iota(element_type, size) + + def BroadcastedIota(self, dtype, shape, dimension): + """Enqueues a broadcasted iota constant onto the computation. + + Args: + dtype: expected numpy dtype of the output. + shape: tuple of integers, the expected output shape (dimensions). + dimension: positive integer, dimension along which to increment values. + + Returns: + A LocalOp representing the added broadcasted iota constant. + """ + xla_shape = Shape.array_shape(dtype, shape) + return self._client.BroadcastedIota(xla_shape, dimension) + def Broadcast(self, operand, sizes): """Enqueues a broadcast operation onto the computation. @@ -700,6 +872,20 @@ class ComputationBuilder(object): """ return self._client.Broadcast(operand, sizes) + def BroadcastInDim(self, operand, shape, broadcast_dimensions): + """Enqueues a broadcast-in-dimensions operation onto the computation. + + Args: + operand: the operand LocalOp to broadcast. + shape: tuple of integers, the expected output shape. + broadcast_dimensions: tuple of integers identifying which dimensions + of the output are to be broadcast into. + + Returns: + A LocalOp representing the added broadcast-in-dimensions op. + """ + return self._client.BroadcastInDim(operand, shape, broadcast_dimensions) + def Concatenate(self, operands, dimension): """Enqueues a concatenate operation onto the computation. @@ -779,16 +965,30 @@ class ComputationBuilder(object): dimensions = tuple(range(ndim)) return self._client.Reshape(operand, dimensions, new_sizes) - def CrossReplicaSum(self, operand): + 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) + + def make_proto(replica_group): + replica_group_proto = xla_data_pb2.ReplicaGroup() + replica_group_proto.replica_ids.extend(replica_group) + return replica_group_proto + + if replica_groups is None: + replica_groups = [] # special value for XLA API + else: + replica_groups = [make_proto(group) for group in replica_groups] + return self._client.CrossReplicaSum(operand, replica_groups) def Collapse(self, operand, dimensions): """Collapse op.""" @@ -834,8 +1034,8 @@ class ComputationBuilder(object): padding, self.GetShape(operand).dimensions(), window_dimensions, window_strides) return self._client.SelectAndScatterWithGeneralPadding( - operand, select.c_local_computation, window_dimensions, window_strides, - pads, source, init_value, scatter.c_local_computation) + operand, select.computation, window_dimensions, window_strides, pads, + source, init_value, scatter.computation) def Select(self, pred, on_true, on_false): """Element-wise selection op. @@ -943,7 +1143,32 @@ class ComputationBuilder(object): Returns: A LocalOp representing the added call op. """ - return self._client.Call(computation_to_apply.c_local_computation, operands) + return self._client.Call(computation_to_apply.computation, operands) + + def CustomCall(self, + call_target_name, + operands, + shape_with_layout, + operand_shapes_with_layout, + opaque=None): + """Enqueues a custom call operation onto the computation. + + Args: + call_target_name: the name of the function to call. + operands: an iterable of LocalOp. The number and types of operands must + match the arity of `operand_shapes_with_layout`. + shape_with_layout: the shape of the operator's output, with layout. + operand_shapes_with_layout: the shapes of `operands`, including the + expected layouts. + opaque: an opaque string passed to the backend. + + Returns: + A LocalOp representing the added custom call op. + """ + opaque = opaque or b'' + return self._client.CustomCall(call_target_name, operands, + shape_with_layout, + operand_shapes_with_layout, opaque) def Map(self, operands, computation_to_apply, dimensions): """Enqueues a map operation onto the computation. @@ -956,7 +1181,7 @@ class ComputationBuilder(object): Returns: A LocalOp representing the added Map op. """ - return self._client.Map(operands, computation_to_apply.c_local_computation, + return self._client.Map(operands, computation_to_apply.computation, dimensions) def Reduce(self, operand, init_value, computation_to_apply, dimensions): @@ -972,8 +1197,7 @@ class ComputationBuilder(object): A LocalOp representing the added Reduce op. """ return self._client.Reduce(operand, init_value, - computation_to_apply.c_local_computation, - dimensions) + computation_to_apply.computation, dimensions) def ReduceWindow(self, operand, init_value, computation_to_apply, window_dimensions, window_strides, padding): @@ -994,7 +1218,7 @@ class ComputationBuilder(object): padding, self.GetShape(operand).dimensions(), window_dimensions, window_strides) return self._client.ReduceWindowWithGeneralPadding( - operand, init_value, computation_to_apply.c_local_computation, + operand, init_value, computation_to_apply.computation, window_dimensions, window_strides, (), (), pads) def ReduceWindowWithGeneralPadding( @@ -1016,7 +1240,7 @@ class ComputationBuilder(object): A LocalOp representing the added ReduceWindow op. """ return self._client.ReduceWindowWithGeneralPadding( - operand, init_value, computation_to_apply.c_local_computation, + operand, init_value, computation_to_apply.computation, window_dimensions, window_strides, base_dilations, window_dilations, padding) @@ -1062,8 +1286,7 @@ class ComputationBuilder(object): Returns: a LocalOp representing the While operation. """ - return self._client.While(cond.c_local_computation, - body.c_local_computation, init) + return self._client.While(cond.computation, body.computation, init) def Conditional(self, pred, true_operand, true_computation, false_operand, false_computation): @@ -1079,8 +1302,8 @@ class ComputationBuilder(object): Returns: a LocalOp representing the Conditional operation. """ return self._client.Conditional( - pred, true_operand, true_computation.c_local_computation, false_operand, - false_computation.c_local_computation) + pred, true_operand, true_computation.computation, false_operand, + false_computation.computation) def IsConstant(self, operand): """Checks whether the given operand is a compile-time constant. @@ -1147,10 +1370,9 @@ class ComputationBuilder(object): pads = _convert_padding_type_to_pad_values( padding, self.GetShape(lhs).dimensions()[2:], self.GetShape(rhs).dimensions()[2:], window_strides) - dimension_numbers = self._GetConvDimensionNumbers(len(window_strides)) - return self._client.ConvGeneralDilated(lhs, rhs, window_strides, pads, (), - (), dimension_numbers, - feature_group_count) + return self.ConvGeneralDilated( + lhs, rhs, window_strides, pads, (), (), + dimension_numbers=None, feature_group_count=feature_group_count) def ConvWithGeneralPadding(self, lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, feature_group_count=1): @@ -1168,11 +1390,9 @@ class ComputationBuilder(object): Returns: A ComputationdataHandle representing the added ConvWithGeneralPadding op. """ - dimension_numbers = self._GetConvDimensionNumbers(len(window_strides)) - return self._client.ConvGeneralDilated(lhs, rhs, window_strides, padding, - lhs_dilation, rhs_dilation, - dimension_numbers, - feature_group_count) + return self.ConvGeneralDilated( + lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, + dimension_numbers=None, feature_group_count=feature_group_count) def _GetConvDimensionNumbers(self, num_spatial_dims): """Create ConvolutionDimensionNumbers proto for convolutions.""" @@ -1190,7 +1410,7 @@ class ComputationBuilder(object): return dimension_numbers def ConvGeneralDilated(self, lhs, rhs, window_strides, padding, lhs_dilation, - rhs_dilation, dimension_numbers, + rhs_dilation, dimension_numbers=None, feature_group_count=1): """Enqueues a ConvGeneralDilated operation onto the computation. @@ -1201,10 +1421,11 @@ class ComputationBuilder(object): padding: length-N array-like of pairs of integers of (low, high) padding. lhs_dilation: length-N array-like of integer dilation factors. rhs_dilation: length-N array-like of integer dilation factors. - dimension_numbers: either an xla_data_pb2.ConvolutionDimensionNumbers or a - triple (lhs_spec, rhs_spec, out_spec) where each element is a string of - length N+2 identifying by position (1) batch dimensions in lhs, rhs, and - the output with the character 'N', (2) feature dimensions in lhs and the + dimension_numbers: optional, either an + xla_data_pb2.ConvolutionDimensionNumbers proto instance or a tuple + (lhs_spec, rhs_spec, out_spec) where each element is a string of length + N+2 identifying by position (1) batch dimensions in lhs, rhs, and the + output with the character 'N', (2) feature dimensions in lhs and the output with the character 'C', (3) input and output feature dimensions in rhs with the characters 'I' and 'O' respectively, and (4) spatial dimension correspondences between lhs, rhs, and the output using any @@ -1217,13 +1438,16 @@ class ComputationBuilder(object): spatial dimension character labels according to the order in which the labels appear in the rhs_spec string, so that window_strides[0] is matched with the dimension corresponding to the first character - appearing in rhs_spec that is not 'I' or 'O'. + appearing in rhs_spec that is not 'I' or 'O'. By default, use the same + dimension numbering as Conv and ConvWithGeneralPadding. feature_group_count: number of feature groups for grouped convolution. Returns: a LocalOp representing the ConvGenralDilated operation. """ - if not isinstance(dimension_numbers, - xla_data_pb2.ConvolutionDimensionNumbers): + if dimension_numbers is None: + dimension_numbers = self._GetConvDimensionNumbers(len(window_strides)) + elif not isinstance(dimension_numbers, + xla_data_pb2.ConvolutionDimensionNumbers): lhs_spec, rhs_spec, out_spec = dimension_numbers dimension_numbers = xla_data_pb2.ConvolutionDimensionNumbers() @@ -1255,6 +1479,32 @@ class ComputationBuilder(object): """Enqueues a key-value sort operation onto the computation.""" return self._client.SortKeyVal(keys, values, dimension) + def Cholesky(self, a): + """Enqueues a Cholesky decomposition onto the computation.""" + return self._client.Cholesky(a) + + def QR(self, a, full_matrices=True): + """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): + """Enqueues a triangular-solve operation onto the computation.""" + return self._client.TriangularSolve( + a, b, left_side, lower, transpose_a, conjugate_a) + + 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(): """Forward remaining ComputationBuilder methods to the C API. @@ -1308,6 +1558,19 @@ def initialize_replica_count(replica_count): c_api.InitializeReplicaCount(replica_count) +def initialize_platform_name(platform_name): + """Initializes the desired platform name to use on XLA service init. + + Args: + platform_name: string name of platform. + + Raises: + A runtime exception if the XLA service has already been initialized. + """ + platform_name = _maybe_encode_string(platform_name) + c_api.InitializePlatformName(platform_name) + + def get_replica_count(): """Returns the current replica count used for the XLA service. @@ -1317,6 +1580,16 @@ def get_replica_count(): return c_api.GetReplicaCount() +def register_cpu_custom_call_target(name, fn): + """Registers a CPU custom call target. + + Args: + name: bytes containing the name of the function. + fn: a PyCapsule object containing the function pointer. + """ + c_api.RegisterCpuCustomCallTarget(name, fn) + + def GetPaddingConfigFromTriples(triples): """Create PaddingConfig proto from list of triples of integers.""" padding_config = xla_data_pb2.PaddingConfig() diff --git a/tensorflow/compiler/xla/python/xla_client_test.py b/tensorflow/compiler/xla/python/xla_client_test.py index 82103f03132e45ff822ce1ebcc2be47b24f5869f..54c76241b9929fb39a6d63648f8ff35d78534b28 100644 --- a/tensorflow/compiler/xla/python/xla_client_test.py +++ b/tensorflow/compiler/xla/python/xla_client_test.py @@ -18,11 +18,13 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import functools import itertools import threading import numpy as np +from tensorflow.compiler.xla.python import custom_call_for_test from tensorflow.compiler.xla.python import xla_client import unittest @@ -37,7 +39,7 @@ class LocalComputationTest(unittest.TestCase): def _Execute(self, c, arguments): compiled_c = c.Build().CompileWithExampleArguments(arguments) - return compiled_c.Execute(arguments) + return compiled_c.ExecuteWithPythonValues(arguments) def _ExecuteAndAssertWith(self, assert_func, c, arguments, expected): assert expected is not None @@ -51,9 +53,11 @@ class LocalComputationTest(unittest.TestCase): def _ExecuteAndCompareExact(self, c, arguments=(), expected=None): self._ExecuteAndAssertWith(np.testing.assert_equal, c, arguments, expected) - def _ExecuteAndCompareClose(self, c, arguments=(), expected=None): - self._ExecuteAndAssertWith(np.testing.assert_allclose, c, arguments, - expected) + def _ExecuteAndCompareClose(self, c, arguments=(), expected=None, rtol=1e-7, + atol=0): + self._ExecuteAndAssertWith( + functools.partial(np.testing.assert_allclose, rtol=rtol, atol=atol), + c, arguments, expected) def NumpyArrayF32(*args, **kwargs): @@ -143,6 +147,17 @@ class ComputationsWithConstantsTest(LocalComputationTest): c.Pow(c.Constant(NumpyArrayF64([1.5, 2.5, 3.0])), c.ConstantF64Scalar(2.)) self._ExecuteAndCompareClose(c, expected=[2.25, 6.25, 9.]) + def testIota(self): + c = self._NewComputation() + c.Iota(np.float32, 10) + self._ExecuteAndCompareExact(c, expected=np.arange(10, dtype=np.float32)) + + def testBroadcastedIota(self): + c = self._NewComputation() + c.BroadcastedIota(np.int64, (2, 3), 1) + expected = np.array([[0, 1, 2], [0, 1, 2]], dtype=np.int64) + self._ExecuteAndCompareExact(c, expected=expected) + def testBooleanAnd(self): c = self._NewComputation() c.And( @@ -268,6 +283,20 @@ class ComputationsWithConstantsTest(LocalComputationTest): c.Constant(NumpyArrayF64([100, -100, 200, -200]))) self._ExecuteAndCompareClose(c, expected=[104.4, -93.4, 208.8, -189]) + def testCustomCall(self): + c = self._NewComputation() + for name, fn in custom_call_for_test.cpu_custom_call_targets.items(): + xla_client.register_cpu_custom_call_target(name, fn) + c.CustomCall( + b"test_subtract_f32", + operands=(c.ConstantF32Scalar(1.25), c.ConstantF32Scalar(0.5)), + shape_with_layout=xla_client.Shape.array_shape(np.float32, (), ()), + operand_shapes_with_layout=( + xla_client.Shape.array_shape(np.float32, (), ()), + xla_client.Shape.array_shape(np.float32, (), ()), + )) + self._ExecuteAndCompareClose(c, expected=0.75) + class ParametersTest(LocalComputationTest): """Tests focusing on Parameter ops and argument-passing.""" @@ -355,7 +384,7 @@ class LocalBufferTest(LocalComputationTest): def _Execute(self, c, arguments): compiled_c = c.Build().CompileWithExampleArguments(arguments) arg_buffers = [xla_client.LocalBuffer.from_pyval(arg) for arg in arguments] - result_buffer = compiled_c.ExecuteWithLocalBuffers(arg_buffers) + result_buffer = compiled_c.Execute(arg_buffers) return result_buffer.to_py() def testConstantSum(self): @@ -388,7 +417,7 @@ class LocalBufferTest(LocalComputationTest): arg_buffer = xla_client.LocalBuffer.from_pyval(arg) arg_buffer.delete() with self.assertRaises(ValueError): - compiled_c.ExecuteWithLocalBuffers([arg_buffer]) + compiled_c.Execute([arg_buffer]) def testDestructureTupleEmpty(self): t = () @@ -439,6 +468,13 @@ class LocalBufferTest(LocalComputationTest): np.testing.assert_equal(NumpyArrayF32([1.0, 2.0]), got[0]) np.testing.assert_equal(NumpyArrayS32([3, 4]), got[1]) + def testShape(self): + pyval = np.array([[1., 2.]], np.float32) + local_buffer = xla_client.LocalBuffer.from_pyval(pyval) + xla_shape = local_buffer.shape() + self.assertEqual(xla_shape.dimensions(), (1, 2,)) + self.assertEqual(np.dtype(xla_shape.element_type()), np.dtype(np.float32)) + class SingleOpTest(LocalComputationTest): """Tests for single ops. @@ -478,7 +514,7 @@ class SingleOpTest(LocalComputationTest): x = c.Constant(np.array(template, dtype=src_dtype)) c.ConvertElementType(x, xla_types[dst_dtype]) - result = c.Build().Compile().Execute() + result = c.Build().Compile().ExecuteWithPythonValues() expected = np.array(template, dtype=dst_dtype) self.assertEqual(result.shape, expected.shape) @@ -505,7 +541,7 @@ class SingleOpTest(LocalComputationTest): x = c.Constant(np.array(template, dtype=src_dtype)) c.BitcastConvertType(x, dst_etype) - result = c.Build().Compile().Execute() + result = c.Build().Compile().ExecuteWithPythonValues() expected = np.array(template, src_dtype).view(dst_dtype) self.assertEqual(result.shape, expected.shape) @@ -529,6 +565,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]]) @@ -987,7 +1035,7 @@ class SingleOpTest(LocalComputationTest): c.Tuple( c.ConstantS32Scalar(42), c.Constant(NumpyArrayF32([1.0, 2.0])), c.Constant(NumpyArrayBool([True, False, False, True]))) - result = c.Build().Compile().Execute() + result = c.Build().Compile().ExecuteWithPythonValues() self.assertIsInstance(result, tuple) np.testing.assert_equal(result[0], 42) np.testing.assert_allclose(result[1], [1.0, 2.0]) @@ -1007,12 +1055,19 @@ class SingleOpTest(LocalComputationTest): self._ExecuteAndCompareExact( c, expected=[[10, 20, 30, 40], [10, 20, 30, 40], [10, 20, 30, 40]]) + def testBroadcastInDim(self): + c = self._NewComputation() + c.BroadcastInDim(c.Constant(NumpyArrayS32([1, 2])), [2, 2], [0]) + self._ExecuteAndCompareExact(c, expected=[[1, 1], [2, 2]]) + c.BroadcastInDim(c.Constant(NumpyArrayS32([1, 2])), [2, 2], [1]) + self._ExecuteAndCompareExact(c, expected=[[1, 2], [1, 2]]) + def testRngNormal(self): shape = (2, 3) c = self._NewComputation() c.RngNormal(c.Constant(NumpyArrayF32(0.)), c.Constant(NumpyArrayF32(1.)), dims=shape) - result = c.Build().Compile().Execute() + result = c.Build().Compile().ExecuteWithPythonValues() # since the result is random, we just check shape and uniqueness self.assertEqual(result.shape, shape) self.assertEqual(len(np.unique(result)), np.prod(shape)) @@ -1023,7 +1078,7 @@ class SingleOpTest(LocalComputationTest): c = self._NewComputation() c.RngUniform(c.Constant(NumpyArrayF32(lo)), c.Constant(NumpyArrayF32(hi)), dims=shape) - result = c.Build().Compile().Execute() + result = c.Build().Compile().ExecuteWithPythonValues() # since the result is random, we just check shape, uniqueness, and range self.assertEqual(result.shape, shape) self.assertEqual(len(np.unique(result)), np.prod(shape)) @@ -1036,13 +1091,45 @@ class SingleOpTest(LocalComputationTest): c = self._NewComputation() c.RngUniform(c.Constant(NumpyArrayS32(lo)), c.Constant(NumpyArrayS32(hi)), dims=shape) - result = c.Build().Compile().Execute() + result = c.Build().Compile().ExecuteWithPythonValues() # since the result is random, we just check shape, integrality, and range self.assertEqual(result.shape, shape) self.assertEqual(result.dtype, np.int32) self.assertTrue(np.all(lo <= result)) self.assertTrue(np.all(result < hi)) + def testCholesky(self): + l = np.array([[4, 0, 0, 0], [6, 5, 0, 0], [2, 14, 16, 0], [3, 6, 1, 4]], + dtype=np.float32) + c = self._NewComputation() + c.Cholesky(c.Constant(np.dot(l, l.T))) + self._ExecuteAndCompareClose(c, expected=l, rtol=1e-4) + + def testQR(self): + a = np.array( + [[4, 6, 8, 10], [6, 45, 54, 63], [8, 54, 146, 166], [10, 63, 166, 310]], + dtype=np.float32) + c = self._NewComputation() + c.QR(c.Constant(a), full_matrices=True) + q, r = self._Execute(c, ()) + np.testing.assert_allclose(np.dot(q, r), a, rtol=1e-4) + + def testTriangularSolve(self): + a_vals = np.array( + [[2, 0, 0, 0], [3, 6, 0, 0], [4, 7, 9, 0], [5, 8, 10, 11]], + dtype=np.float32) + b_vals = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]], + dtype=np.float32) + + c = self._NewComputation() + c.TriangularSolve(c.Constant(a_vals), c.Constant(b_vals), left_side=False, + lower=True, transpose_a=True) + self._ExecuteAndCompareClose(c, expected=np.array([ + [0.5, 0.08333334, 0.04629629, 0.03367003], + [2.5, -0.25, -0.1388889, -0.1010101], + [4.5, -0.58333331, -0.32407406, -0.23569024], + ], dtype=np.float32), rtol=1e-4) + def testIsConstant(self): c = self._NewComputation() a = c.ConstantS32Scalar(3) @@ -1054,6 +1141,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).""" @@ -1111,6 +1213,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") @@ -1473,7 +1583,7 @@ class EmbeddedComputationsTest(LocalComputationTest): xla_client.transfer_to_infeed(item) for item in to_infeed: - result = compiled_c.Execute() + result = compiled_c.ExecuteWithPythonValues() self.assertEqual(result, item) def testInfeedThenOutfeedS32(self): @@ -1493,6 +1603,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): @@ -1511,5 +1638,20 @@ class ErrorTest(LocalComputationTest): lambda: c.Build().CompileWithExampleArguments([self.f32_scalar_2])) +class ComputationRootTest(LocalComputationTest): + """Tests related to setting the root of the computation.""" + + def testComputationRootDifferentFromLastOp(self): + c = self._NewComputation() + x = c.ParameterFromNumpy(NumpyArrayF32(2.0)) + result = c.Add(x, c.ConstantF32Scalar(3.14)) + extra = c.Add(result, c.ConstantF32Scalar(1.618)) # pylint: disable=unused-variable + + arg = NumpyArrayF32(1.0) + compiled_c = c.Build(result).CompileWithExampleArguments([arg]) + ans = compiled_c.ExecuteWithPythonValues([arg]) + np.testing.assert_allclose(ans, 4.14) + + if __name__ == "__main__": unittest.main() 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/python_api/xla_shape.py b/tensorflow/compiler/xla/python_api/xla_shape.py index f158f6b2410352432445f669155aff0af5526abf..bdcd4abd6cc708795416b15412f37dde10d7fe97 100644 --- a/tensorflow/compiler/xla/python_api/xla_shape.py +++ b/tensorflow/compiler/xla/python_api/xla_shape.py @@ -20,14 +20,17 @@ from __future__ import print_function import numpy as _np # Avoids becoming a part of public Tensorflow API. +from six.moves import xrange + from tensorflow.compiler.xla import xla_data_pb2 from tensorflow.compiler.xla.python_api import types class Shape(object): - """Wraps a xla_data_pb2.Shape message with a convenient Python type. + """Wraps a xla_data_pb2.ShapeProto message with a convenient Python type. - Provides direct access to the underlying xla_data_pb2.Shape message in the + Provides direct access to the underlying xla_data_pb2.ShapeProto message in + the message attribute, along with accessor wrappers to the message's fields. Avoid direct access to .message unless interacting directly with protobuf APIs like CopyFrom. In other words, prefer hauling the shape around in a Shape, and @@ -48,7 +51,7 @@ class Shape(object): Raises: ValueError: if element_type is TUPLE but dimensions are not Shape objects. """ - self.message = xla_data_pb2.Shape() + self.message = xla_data_pb2.ShapeProto() self.message.element_type = element_type if element_type == xla_data_pb2.TUPLE: if not all(isinstance(subshape, Shape) for subshape in dimensions): diff --git a/tensorflow/compiler/xla/reference_util.cc b/tensorflow/compiler/xla/reference_util.cc index ceb5e74db7c3b9305e9d77068df9ae0a3690af8a..08b78ee244844f41d551d7e249cec0cbf157d639 100644 --- a/tensorflow/compiler/xla/reference_util.cc +++ b/tensorflow/compiler/xla/reference_util.cc @@ -18,10 +18,10 @@ 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" -#include "tensorflow/compiler/xla/service/cpu/runtime_single_threaded_matmul.h" #include "tensorflow/compiler/xla/service/hlo_evaluator.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/shape_inference.h" @@ -32,48 +32,19 @@ limitations under the License. namespace xla { -namespace { - -template -std::unique_ptr> MatmulArray2DImpl( - const Array2D& lhs, const Array2D& rhs, - const std::function& impl_fn) { - CHECK_EQ(lhs.width(), rhs.height()); - int m = lhs.height(); - int n = rhs.width(); - int k = lhs.width(); - auto result = absl::make_unique>(m, n); - // Because Eigen is a header-oriented library, make sure that the Eigen code - // is the same as the code used by the CPU backend (otherwise the linker will - // randomly pick *some* definition). - impl_fn( - /*run_options_ptr=*/nullptr, result->data(), rhs.data(), lhs.data(), n, m, - k, - /*transpose_lhs=*/0, - /*transpose_rhs=*/0); - return result; -} - -} // namespace - /* static */ std::unique_ptr> ReferenceUtil::MatmulArray2D( const Array2D& lhs, const Array2D& rhs) { - return MatmulArray2DImpl( - lhs, rhs, __xla_cpu_runtime_EigenSingleThreadedMatMulF16); + return HloEvaluator::MatmulArray2D(lhs, rhs); } /* static */ std::unique_ptr> ReferenceUtil::MatmulArray2D( const Array2D& lhs, const Array2D& rhs) { - return MatmulArray2DImpl( - lhs, rhs, __xla_cpu_runtime_EigenSingleThreadedMatMulF32); + return HloEvaluator::MatmulArray2D(lhs, rhs); } /* static */ std::unique_ptr> ReferenceUtil::MatmulArray2D( const Array2D& lhs, const Array2D& rhs) { - return MatmulArray2DImpl( - lhs, rhs, __xla_cpu_runtime_EigenSingleThreadedMatMulF64); + return HloEvaluator::MatmulArray2D(lhs, rhs); } /* static */ std::unique_ptr> ReferenceUtil::Array2DF32ToF64( @@ -557,10 +528,11 @@ ReferenceUtil::ConvArray4DGeneralDimensionsDilated( dim2.set_base_dilation(lhs_dilation.second); *window.add_dimensions() = dim2; - const Shape& shape = ShapeInference::InferConvolveShape( - lhs_literal.shape(), rhs_literal.shape(), - /*feature_group_count=*/1, window, dnums) - .ConsumeValueOrDie(); + const Shape& shape = + ShapeInference::InferConvolveShape( + lhs_literal.shape(), rhs_literal.shape(), + /*feature_group_count=*/1, /*batch_group_count=*/1, window, dnums) + .ConsumeValueOrDie(); HloInstruction* lhs_instruction = b.AddInstruction(HloInstruction::CreateConstant(std::move(lhs_literal))); @@ -572,16 +544,16 @@ ReferenceUtil::ConvArray4DGeneralDimensionsDilated( /*new_size=*/2, PrecisionConfig::DEFAULT); b.AddInstruction(HloInstruction::CreateConvolve( shape, lhs_instruction, rhs_instruction, /*feature_group_count=*/1, - window, dnums, precision_config)); + /*batch_group_count=*/1, window, dnums, precision_config)); HloModuleConfig config; HloModule module("ReferenceUtil", config); auto computation = module.AddEntryComputation(b.Build()); HloEvaluator evaluator; Literal result_literal = - evaluator.Evaluate(*computation, {}).ConsumeValueOrDie(); + evaluator.Evaluate(*computation, {}).ConsumeValueOrDie(); - CHECK_EQ(ShapeUtil::Rank(result_literal.shape()), 4); + CHECK_EQ(result_literal.shape().rank(), 4); auto result = absl::make_unique>(result_literal.shape().dimensions(0), result_literal.shape().dimensions(1), @@ -634,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/BUILD b/tensorflow/compiler/xla/rpc/BUILD index 3abb3855a42b8b5222115262448d359da3a80e87..26affbcceb33110baf41d507173e56f8b1c8c9eb 100644 --- a/tensorflow/compiler/xla/rpc/BUILD +++ b/tensorflow/compiler/xla/rpc/BUILD @@ -16,7 +16,6 @@ xla_proto_library( use_grpc_plugin = True, visibility = ["//visibility:public"], deps = [ - "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla:xla_proto", ], ) diff --git a/tensorflow/compiler/xla/rpc/grpc_service.cc b/tensorflow/compiler/xla/rpc/grpc_service.cc index 4e1435fa30a24c320ddbedb84d37b369a3158a54..22b4218fbd5e9bc59a0de22735eb51db46670f09 100644 --- a/tensorflow/compiler/xla/rpc/grpc_service.cc +++ b/tensorflow/compiler/xla/rpc/grpc_service.cc @@ -47,11 +47,34 @@ namespace xla { }); } -::grpc::Status GRPCService::ExecuteGraph(::grpc::ServerContext* /*context*/, - const ExecuteGraphRequest* arg, - ExecuteResponse* result) { +::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) { return DelegateRPC( - [this, arg, result]() { return service_->ExecuteGraph(arg, result); }); + [this, arg, result]() { return service_->Compile(arg, result); }); +} + +::grpc::Status GRPCService::Execute(::grpc::ServerContext* /*context*/, + const ExecuteRequest* arg, + ExecuteResponse* result) { + return DelegateRPC( + [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, diff --git a/tensorflow/compiler/xla/rpc/grpc_service.h b/tensorflow/compiler/xla/rpc/grpc_service.h index ca1b09b648013ad45d806040c5ddcf11d9e5604e..b546704f73e34941cbf7bc2fe08062aa438039f7 100644 --- a/tensorflow/compiler/xla/rpc/grpc_service.h +++ b/tensorflow/compiler/xla/rpc/grpc_service.h @@ -39,9 +39,20 @@ class GRPCService : public grpc::XlaService::Service { const DeconstructTupleRequest* arg, DeconstructTupleResponse* result) override; - ::grpc::Status ExecuteGraph(::grpc::ServerContext* context, - const ExecuteGraphRequest* arg, - ExecuteResponse* 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; + + ::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/rpc/grpc_stub.cc b/tensorflow/compiler/xla/rpc/grpc_stub.cc index 7b8ab158e1396d7087a407be180ab44d2e16e121..66abf66cfd6c2f753c5507aa373452ac880e9a29 100644 --- a/tensorflow/compiler/xla/rpc/grpc_stub.cc +++ b/tensorflow/compiler/xla/rpc/grpc_stub.cc @@ -62,10 +62,17 @@ Status GRPCStub::ResetDevice(const ResetDeviceRequest* request, }); } -Status GRPCStub::ExecuteGraph(const ExecuteGraphRequest* request, - ExecuteResponse* response) { +Status GRPCStub::Compile(const CompileRequest* request, + CompileResponse* response) { return MakeRPC([this, request, response](::grpc::ClientContext* context) { - return grpc_stub_->ExecuteGraph(context, *request, response); + return grpc_stub_->Compile(context, *request, response); + }); +} + +Status GRPCStub::Execute(const ExecuteRequest* request, + ExecuteResponse* response) { + return MakeRPC([this, request, response](::grpc::ClientContext* context) { + return grpc_stub_->Execute(context, *request, response); }); } diff --git a/tensorflow/compiler/xla/rpc/grpc_stub.h b/tensorflow/compiler/xla/rpc/grpc_stub.h index 8dfcb761387d608abbb1f62974f49b976a7ff7ff..f02b401399f3e895153f0b08e325bc9c2c2336ec 100644 --- a/tensorflow/compiler/xla/rpc/grpc_stub.h +++ b/tensorflow/compiler/xla/rpc/grpc_stub.h @@ -43,8 +43,11 @@ class GRPCStub : public ServiceInterface { Status ResetDevice(const ResetDeviceRequest* arg, ResetDeviceResponse* result) override; - Status ExecuteGraph(const ExecuteGraphRequest* request, - ExecuteResponse* response) override; + Status Compile(const CompileRequest* request, + CompileResponse* response) override; + + Status Execute(const ExecuteRequest* request, + ExecuteResponse* response) override; Status ExecuteGraphParallel(const ExecuteGraphParallelRequest* request, ExecuteParallelResponse* response) override; diff --git a/tensorflow/compiler/xla/rpc/xla_service.proto b/tensorflow/compiler/xla/rpc/xla_service.proto index 551ae895e05586daec0ffcd425f4950f76bdd50d..0ff8adc2acbe5fd21e85027dd63bfb14f5672a7d 100644 --- a/tensorflow/compiler/xla/rpc/xla_service.proto +++ b/tensorflow/compiler/xla/rpc/xla_service.proto @@ -43,7 +43,6 @@ limitations under the License. syntax = "proto3"; import "tensorflow/compiler/xla/xla.proto"; -import "tensorflow/compiler/xla/xla_data.proto"; package xla; @@ -128,11 +127,14 @@ service XlaService { returns (CreateChannelHandleResponse) { } - // Invokes the provided computation with the provided global data passed as - // immutable arguments. The request contains the whole computation graph. + // Compiles the provided computation into executable. Returns the handle of + // the executable. + rpc Compile(CompileRequest) returns (CompileResponse) {} + + // Invokes the provided executable with the provided global data passed as + // immutable arguments. The request contains the handle to the executable. // Returns global data output and execution timing. - rpc ExecuteGraph(ExecuteGraphRequest) returns (ExecuteResponse) { - } + rpc Execute(ExecuteRequest) returns (ExecuteResponse) {} // Invokes the provided list of computations in parallel with the provided // global data for each computation. Returns a list of global data output and diff --git a/tensorflow/compiler/xla/service/BUILD b/tensorflow/compiler/xla/service/BUILD index 3a716c385b2bced7b36e65012d2ff6888525524a..3c123fd5c6b88872ce1d9633893e4fedaf3dd328 100644 --- a/tensorflow/compiler/xla/service/BUILD +++ b/tensorflow/compiler/xla/service/BUILD @@ -1,6 +1,14 @@ # Description: # XLA service implementation. +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.bzl", + "tf_proto_library_py", +) +load("//tensorflow:tensorflow.bzl", "tf_cc_binary", "tf_cc_test") + licenses(["notice"]) # Apache 2.0 package(default_visibility = [":friends"]) @@ -12,15 +20,6 @@ package_group( ], ) -load("//tensorflow/compiler/xla/tests:build_defs.bzl", "xla_test") -load("//tensorflow/compiler/xla:xla.bzl", "xla_proto_library") -load("//tensorflow:tensorflow.bzl", "tf_cc_test") -load("//tensorflow:tensorflow.bzl", "tf_cc_binary") -load( - "//tensorflow/core:platform/default/build_config.bzl", - "tf_proto_library_py", -) - xla_proto_library( name = "hlo_proto", srcs = ["hlo.proto"], @@ -87,7 +86,6 @@ tf_cc_test( "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:lib", ], @@ -124,7 +122,6 @@ tf_cc_test( "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:lib", ], @@ -158,12 +155,12 @@ tf_cc_test( ":bfloat16_propagation", ":bfloat16_support", ":hlo", + "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:test_helpers", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", "//tensorflow/compiler/xla/tests:literal_test_util", "//tensorflow/compiler/xla/tests:xla_internal_test_main", # fixdeps: keep ], @@ -226,23 +223,28 @@ 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", ":shape_inference", + "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:literal", "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", @@ -251,11 +253,14 @@ cc_library( "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:window_util", "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla/service/cpu:runtime_single_threaded_matmul", "//tensorflow/core:lib", "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/base", "@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", @@ -280,12 +285,14 @@ tf_cc_test( "//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_verified_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", # fixdeps: keep "//tensorflow/core:lib", "//tensorflow/core:test", "@com_google_absl//absl/memory", + "@com_google_absl//absl/strings:str_format", ], ) @@ -293,6 +300,7 @@ cc_library( name = "hlo", srcs = [ "dfs_hlo_visitor.cc", + "dynamic_parameter_binding.cc", "hlo_computation.cc", "hlo_input_output_alias_config.cc", "hlo_instruction.cc", @@ -306,6 +314,7 @@ cc_library( hdrs = [ "dfs_hlo_visitor.h", "dfs_hlo_visitor_with_default.h", + "dynamic_parameter_binding.h", "hlo_clone_context.h", "hlo_computation.h", "hlo_domain_metadata.h", @@ -322,7 +331,6 @@ cc_library( ":hlo_casting_utils", ":hlo_module_config", ":hlo_proto", - ":hlo_reachability", ":name_uniquer", "//tensorflow/compiler/xla:array", "//tensorflow/compiler/xla:literal", @@ -352,6 +360,25 @@ cc_library( ], ) +tf_cc_test( + name = "dynamic_parameter_binding_test", + srcs = ["dynamic_parameter_binding_test.cc"], + deps = [ + ":hlo", + ":hlo_dce", + ":hlo_memory_scheduler", + ":hlo_ordering", + ":hlo_parser", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla/tests:hlo_test_base", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", + "//tensorflow/core:test", + "@com_google_absl//absl/algorithm:container", + ], +) + tf_cc_test( name = "dfs_hlo_visitor_with_default_test", srcs = ["dfs_hlo_visitor_with_default_test.cc"], @@ -364,7 +391,6 @@ tf_cc_test( "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:test", ], @@ -390,9 +416,36 @@ tf_cc_test( ":hlo", ":pattern_matcher", "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla/service:hlo_parser", "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:test", + "@com_google_absl//absl/strings", + ], +) + +cc_library( + name = "pattern_matcher_gmock", + testonly = 1, + hdrs = ["pattern_matcher_gmock.h"], + deps = [ + ":pattern_matcher", + "//tensorflow/compiler/xla:test", + "//tensorflow/core:test", + ], +) + +tf_cc_test( + name = "pattern_matcher_gmock_test", + srcs = ["pattern_matcher_gmock_test.cc"], + deps = [ + ":hlo", + ":pattern_matcher", + ":pattern_matcher_gmock", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:test", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", + "//tensorflow/core:test", ], ) @@ -401,10 +454,12 @@ cc_library( srcs = ["hlo_reachability.cc"], hdrs = ["hlo_reachability.h"], deps = [ + ":hlo", "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", + "@com_google_absl//absl/base", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/types:span", ], @@ -419,7 +474,6 @@ tf_cc_test( "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:test_helpers", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", "//tensorflow/compiler/xla/tests:xla_internal_test_main", ], ) @@ -465,8 +519,8 @@ tf_cc_test( "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:window_util", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", "//tensorflow/compiler/xla/tests:xla_internal_test_main", + "@com_google_absl//absl/container:flat_hash_map", ], ) @@ -518,7 +572,6 @@ tf_cc_test( "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:test", ], @@ -567,7 +620,6 @@ tf_cc_test( "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:lib", "//tensorflow/core:test", @@ -590,7 +642,6 @@ tf_cc_test( "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:test", ], @@ -602,11 +653,11 @@ cc_library( hdrs = ["platform_util.h"], deps = [ ":compiler", + "//tensorflow/compiler/xla:debug_options_flags", "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", - "//tensorflow/compiler/xla/legacy_flags:debug_options_flags", "//tensorflow/core:lib", "//tensorflow/core:stream_executor_no_cuda", "@com_google_absl//absl/strings", @@ -632,6 +683,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", @@ -646,9 +698,11 @@ cc_library( ":allocation_tracker", ":backend", ":channel_tracker", + ":compilation_cache", ":compiler", ":computation_layout", ":device_memory_allocator", + ":dynamic_dimension_inference", ":executable", ":execution_tracker", ":hlo", @@ -661,6 +715,7 @@ cc_library( ":source_map_util", ":stream_pool", ":transfer_manager", + "//tensorflow/compiler/xla:debug_options_flags", "//tensorflow/compiler/xla:executable_run_options", "//tensorflow/compiler/xla:execution_options_util", "//tensorflow/compiler/xla:service_interface", @@ -672,7 +727,6 @@ cc_library( "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla:xla_proto", - "//tensorflow/compiler/xla/legacy_flags:debug_options_flags", "//tensorflow/core:lib", "//tensorflow/core:ptr_util", "//tensorflow/core:stream_executor_no_cuda", @@ -729,12 +783,12 @@ cc_library( ":computation_layout", ":platform_util", ":service", + "//tensorflow/compiler/xla:debug_options_flags", "//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/legacy_flags:debug_options_flags", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", "//tensorflow/core:stream_executor_no_cuda", @@ -810,6 +864,7 @@ tf_cc_test( "//tensorflow/compiler/xla:test_helpers", "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:ptr_util", + "//tensorflow/core:stream_executor_no_cuda", "//tensorflow/core:test", "@com_google_absl//absl/memory", ], @@ -832,6 +887,7 @@ cc_library( ":maybe_owning_device_memory", ":shaped_buffer", ":stream_pool", + "//tensorflow/compiler/xla:debug_options_flags", "//tensorflow/compiler/xla:executable_run_options", "//tensorflow/compiler/xla:shape_tree", "//tensorflow/compiler/xla:status", @@ -839,7 +895,6 @@ cc_library( "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/legacy_flags:debug_options_flags", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", "//tensorflow/core:stream_executor_no_cuda", @@ -955,6 +1010,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", @@ -966,6 +1022,7 @@ cc_library( srcs = ["name_uniquer.cc"], hdrs = ["name_uniquer.h"], deps = [ + "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:types", "//tensorflow/core:lib", "@com_google_absl//absl/container:flat_hash_map", @@ -1005,7 +1062,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", @@ -1043,7 +1099,6 @@ cc_library( ":buffer_value_containers", ":heap_simulator", ":hlo", - ":hlo_memory_scheduler", ":hlo_proto", ":logical_buffer", ":tuple_points_to_analysis", @@ -1085,10 +1140,10 @@ tf_cc_test( "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/service:hlo_parser", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", "//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", ], ) @@ -1102,6 +1157,7 @@ cc_library( ":hlo", ":hlo_dataflow_analysis", ":hlo_proto", + ":hlo_reachability", ":hlo_value", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:status_macros", @@ -1167,7 +1223,6 @@ tf_cc_test( "//tensorflow/compiler/xla:literal", "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:lib", "//tensorflow/core:test", @@ -1183,7 +1238,6 @@ cc_library( deps = [ ":hlo", ":hlo_proto", - "//tensorflow/compiler/xla:statusor", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:span", ], @@ -1342,6 +1396,7 @@ cc_library( ":hlo", "//tensorflow/compiler/xla:literal", "//tensorflow/compiler/xla:shape_util", + "@com_google_absl//absl/container:flat_hash_set", ], ) @@ -1361,10 +1416,12 @@ cc_library( ":fusion_queue", ":hlo", ":hlo_pass", + ":hlo_reachability", "//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", "@com_google_absl//absl/memory", ], ) @@ -1386,6 +1443,7 @@ cc_library( srcs = ["multi_output_fusion.cc"], hdrs = ["multi_output_fusion.h"], deps = [ + ":hlo_reachability", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla/service:hlo", @@ -1426,7 +1484,6 @@ tf_cc_test( "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:test", "@com_google_absl//absl/memory", @@ -1448,7 +1505,6 @@ 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", ], @@ -1502,7 +1558,6 @@ tf_cc_test( "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:lib", "@com_google_absl//absl/memory", @@ -1530,6 +1585,9 @@ cc_library( "//tensorflow/compiler/xla:xla_data_proto", "//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", "@com_google_absl//absl/types:optional", @@ -1545,7 +1603,10 @@ tf_cc_test( ":hlo", ":hlo_casting_utils", ":hlo_matchers", + ":hlo_parser", ":hlo_pass", + ":pattern_matcher", + ":pattern_matcher_gmock", "//tensorflow/compiler/xla:literal", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:test", @@ -1554,7 +1615,6 @@ tf_cc_test( "//tensorflow/compiler/xla:window_util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", "//tensorflow/compiler/xla/tests:xla_internal_test_main", # fixdeps: keep "//tensorflow/core:lib", "//tensorflow/core:test", @@ -1591,7 +1651,6 @@ tf_cc_test( "//tensorflow/compiler/xla:window_util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", "//tensorflow/compiler/xla/tests:xla_internal_test_main", # fixdeps: keep "//tensorflow/core:lib", "//tensorflow/core:test", @@ -1641,16 +1700,16 @@ tf_cc_test( "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", + "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/core:lib", "//tensorflow/core:test", ], ) cc_library( - name = "convolution_feature_group_converter", - srcs = ["convolution_feature_group_converter.cc"], - hdrs = ["convolution_feature_group_converter.h"], + name = "convolution_group_converter", + srcs = ["convolution_group_converter.cc"], + hdrs = ["convolution_group_converter.h"], deps = [ ":hlo", ":hlo_pass", @@ -1668,11 +1727,11 @@ 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_feature_group_converter", + ":convolution_group_converter", ":hlo", ":hlo_matchers", ":hlo_parser", @@ -1693,6 +1752,19 @@ cc_library( ], ) +tf_cc_test( + name = "while_loop_analysis_test", + srcs = ["while_loop_analysis_test.cc"], + deps = [ + ":while_loop_analysis", + "//tensorflow/compiler/xla:test", + "//tensorflow/compiler/xla/service:hlo_parser", + "//tensorflow/compiler/xla/tests:hlo_test_base", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", + "//tensorflow/core:test", + ], +) + cc_library( name = "while_loop_simplifier", srcs = ["while_loop_simplifier.cc"], @@ -1701,9 +1773,11 @@ cc_library( ":call_inliner", ":hlo", ":hlo_pass", + ":hlo_query", + ":pattern_matcher", ":while_loop_analysis", + "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:statusor", - "//tensorflow/core:lib", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/strings", @@ -1715,10 +1789,18 @@ tf_cc_test( name = "while_loop_simplifier_test", srcs = ["while_loop_simplifier_test.cc"], deps = [ + ":algebraic_simplifier", + ":hlo", + ":hlo_cse", + ":hlo_dce", ":hlo_matchers", + ":hlo_parser", + ":hlo_pass", + ":hlo_pass_pipeline", + ":tuple_simplifier", ":while_loop_simplifier", "//tensorflow/compiler/xla:test", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", + "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/core:lib", "//tensorflow/core:test", "@com_google_absl//absl/strings", @@ -1749,7 +1831,7 @@ tf_cc_test( ":hlo_matchers", "//tensorflow/compiler/xla:literal", "//tensorflow/compiler/xla:shape_util", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", + "//tensorflow/compiler/xla/tests:hlo_test_base", ], ) @@ -1777,7 +1859,7 @@ tf_cc_test( ":implicit_broadcast_remover", "//tensorflow/compiler/xla:literal", "//tensorflow/compiler/xla:shape_util", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", + "//tensorflow/compiler/xla/tests:hlo_test_base", ], ) @@ -1791,8 +1873,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", ], ) @@ -1822,7 +1905,6 @@ tf_cc_test( "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", "//tensorflow/core:test", ], ) @@ -1842,6 +1924,81 @@ cc_library( ], ) +cc_library( + name = "dynamic_dimension_inference", + srcs = ["dynamic_dimension_inference.cc"], + hdrs = ["dynamic_dimension_inference.h"], + deps = [ + ":hlo", + "//tensorflow/compiler/xla:status", + "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:types", + "//tensorflow/core:lib", + "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/types:span", + ], +) + +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"], + deps = [ + ":dynamic_dimension_inference", + "//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 = "reshape_mover_test", srcs = ["reshape_mover_test.cc"], @@ -1856,7 +2013,7 @@ tf_cc_test( "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", + "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:lib", "@com_google_absl//absl/memory", @@ -1908,7 +2065,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", @@ -1949,6 +2105,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", @@ -1999,13 +2156,15 @@ tf_cc_test( srcs = ["hlo_computation_test.cc"], deps = [ ":hlo", - ":hlo_matchers", + ":pattern_matcher", + ":pattern_matcher_gmock", "//tensorflow/compiler/xla:literal", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla: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_set", ], ) @@ -2178,6 +2337,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", ], @@ -2262,7 +2422,6 @@ tf_cc_test( "//tensorflow/compiler/xla:test_helpers", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:lib", "//tensorflow/core:test", @@ -2325,13 +2484,27 @@ tf_cc_test( "//tensorflow/compiler/xla:test_helpers", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:lib", "//tensorflow/core:test", ], ) +cc_library( + name = "compilation_cache", + srcs = ["compilation_cache.cc"], + hdrs = ["compilation_cache.h"], + deps = [ + ":executable", + ":hlo_module_config", + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla:util", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/core:lib", + "@com_google_absl//absl/container:flat_hash_map", + ], +) + cc_library( name = "layout_assignment", srcs = [ @@ -2401,14 +2574,13 @@ tf_cc_test( ":hlo_graph_dumper", ":hlo_matchers", ":hlo_runner", + "//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/legacy_flags:debug_options_flags", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", "//tensorflow/core:test", ], ) @@ -2426,6 +2598,7 @@ cc_library( "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", "//tensorflow/core:lib", + "@com_google_absl//absl/container:flat_hash_set", ], ) @@ -2470,6 +2643,7 @@ tf_cc_test( srcs = ["hlo_verifier_test.cc"], deps = [ ":hlo", + ":hlo_module_config", ":hlo_parser", ":hlo_verifier", ":layout_assignment", @@ -2477,6 +2651,7 @@ tf_cc_test( "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla:xla_proto", "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:test", @@ -2526,7 +2701,6 @@ tf_cc_test( "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:test", ], @@ -2582,8 +2756,9 @@ tf_cc_test( ":algebraic_simplifier", ":computation_layout", ":hlo", - ":hlo_matchers", ":layout_assignment", + ":pattern_matcher", + ":pattern_matcher_gmock", "//tensorflow/compiler/xla:literal", "//tensorflow/compiler/xla:shape_layout", "//tensorflow/compiler/xla:shape_util", @@ -2593,8 +2768,8 @@ tf_cc_test( "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/service:hlo_parser", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/compiler/xla/tests:hlo_verified_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/types:span", @@ -2655,7 +2830,7 @@ tf_cc_test( "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", + "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/compiler/xla/tests:test_utils", "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:lib", @@ -2696,7 +2871,6 @@ tf_cc_test( "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/service:hlo_parser", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", "//tensorflow/compiler/xla/tests:literal_test_util", "//tensorflow/compiler/xla/tests:test_utils", "//tensorflow/core:lib", @@ -2730,12 +2904,13 @@ tf_cc_test( ":hlo_matchers", ":hlo_parser", ":hlo_pass", + ":pattern_matcher", + ":pattern_matcher_gmock", "//tensorflow/compiler/xla:literal", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", "//tensorflow/compiler/xla/tests:literal_test_util", "//tensorflow/compiler/xla/tests:xla_internal_test_main", ], @@ -2807,10 +2982,9 @@ tf_cc_test( ":hlo_domain_isolator", ":hlo_domain_remover", ":hlo_parser", + "//tensorflow/compiler/xla:debug_options_flags", "//tensorflow/compiler/xla:test", - "//tensorflow/compiler/xla/legacy_flags:debug_options_flags", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:test", "@com_google_absl//absl/memory", @@ -2843,6 +3017,44 @@ tf_cc_test( ], ) +cc_library( + name = "hlo_get_dimension_size_rewriter", + 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_util", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/core:lib", + "@com_google_absl//absl/algorithm:container", + ], +) + +tf_cc_test( + name = "hlo_get_dimension_size_rewriter_test", + srcs = ["hlo_get_dimension_size_rewriter_test.cc"], + deps = [ + ":hlo", + ":hlo_get_dimension_size_rewriter", + ":hlo_matchers", + ":hlo_parser", + "//tensorflow/compiler/xla:literal", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla:util", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla/tests:hlo_test_base", + "//tensorflow/compiler/xla/tests:literal_test_util", + "//tensorflow/compiler/xla/tests:test_utils", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", + "//tensorflow/core:lib", + "//tensorflow/core:test", + ], +) + cc_library( name = "device_memory_allocator", srcs = [ @@ -2901,6 +3113,7 @@ cc_library( "//tensorflow/core:lib", "//tensorflow/core:lib_internal", "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/strings", "@llvm//:core", "@llvm//:transform_utils", @@ -2998,7 +3211,6 @@ tf_cc_test( deps = [ ":hlo_tfgraph_builder", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:protos_all_cc", ], @@ -3008,6 +3220,7 @@ cc_library( name = "hlo_graph_dumper", srcs = [ "hlo_graph_dumper.cc", + "hlo_graph_html_renderer.cc", ], hdrs = ["hlo_graph_dumper.h"], deps = [ @@ -3015,6 +3228,7 @@ cc_library( ":hlo_casting_utils", ":hlo_execution_profile", ":hlo_tfgraph_builder", + ":pattern_matcher", "//tensorflow/compiler/xla:literal", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:types", @@ -3023,6 +3237,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", @@ -3183,7 +3398,6 @@ cc_library( ":hlo_pass_pipeline", "//tensorflow/compiler/xla:shape_util", "//tensorflow/core:lib", - "@com_google_absl//absl/container:flat_hash_map", ], ) @@ -3240,6 +3454,36 @@ 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", ], ) @@ -3276,6 +3520,8 @@ cc_library( ":tuple_util", "//tensorflow/compiler/xla:literal_util", "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/container:inlined_vector", "@com_google_absl//absl/strings", ], ) @@ -3302,10 +3548,11 @@ cc_library( ":hlo", ":hlo_pass", ":tuple_util", + ":while_loop_analysis", ":while_util", + "//tensorflow/compiler/xla:shape_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:flat_hash_set", @@ -3321,7 +3568,7 @@ tf_cc_test( ":while_loop_invariant_code_motion", "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla/service:hlo_parser", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", + "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/core:test", ], ) @@ -3338,7 +3585,6 @@ cc_library( "//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", ], ) @@ -3351,7 +3597,7 @@ tf_cc_test( ":while_loop_constant_sinking", "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla/service:hlo_parser", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", + "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/core:test", ], ) @@ -3364,6 +3610,7 @@ cc_library( ":bfloat16_normalization", ":defuser", ":hlo", + ":hlo_memory_scheduler", ":hlo_pass", ":hlo_pass_pipeline", ":implicit_broadcast_remover", @@ -3406,14 +3653,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_verified_test_base", + "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/compiler/xla/tests:test_utils", "//tensorflow/core:test", + "@com_google_absl//absl/strings", ], ) @@ -3447,6 +3696,9 @@ tf_cc_test( ":hlo_casting_utils", ":hlo_matchers", ":hlo_parser", + ":pattern_matcher", + ":pattern_matcher_gmock", + "//tensorflow/compiler/xla:test_helpers", "//tensorflow/compiler/xla:window_util", "//tensorflow/core:lib", "//tensorflow/core:test", @@ -3460,7 +3712,6 @@ cc_library( srcs = ["hlo_lexer.cc"], hdrs = [ "hlo_lexer.h", - "hlo_token.h", ], deps = [ "//tensorflow/compiler/xla:shape_util", @@ -3496,6 +3747,74 @@ cc_library( ], ) +cc_library( + name = "ar_crs_combiner", + srcs = ["ar_crs_combiner.cc"], + hdrs = ["ar_crs_combiner.h"], + deps = [ + ":call_graph", + ":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", + "@com_google_absl//absl/strings", + ], +) + +cc_library( + name = "dynamic_index_splitter", + srcs = ["dynamic_index_splitter.cc"], + hdrs = ["dynamic_index_splitter.h"], + deps = [ + ":hlo_casting_utils", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla/service:hlo", + "//tensorflow/compiler/xla/service:hlo_pass", + "@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/strings", + ], +) + +tf_cc_test( + name = "dynamic_index_splitter_test", + srcs = ["dynamic_index_splitter_test.cc"], + deps = [ + ":dynamic_index_splitter", + ":hlo", + ":hlo_matchers", + "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:test", + "//tensorflow/compiler/xla:test_helpers", + "//tensorflow/compiler/xla/tests:hlo_test_base", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", + ], +) + +tf_cc_test( + name = "ar_crs_combiner_test", + srcs = ["ar_crs_combiner_test.cc"], + deps = [ + ":ar_crs_combiner", + ":hlo", + ":hlo_matchers", + "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla/tests:hlo_test_base", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", + "//tensorflow/core:lib", + "//tensorflow/core:test", + ], +) + tf_cc_test( name = "map_inliner_test", srcs = ["map_inliner_test.cc"], @@ -3507,7 +3826,7 @@ tf_cc_test( "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", + "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/compiler/xla/tests:literal_test_util", "//tensorflow/compiler/xla/tests:xla_internal_test_main", # fixdeps: keep "@com_google_absl//absl/memory", diff --git a/tensorflow/compiler/xla/service/algebraic_simplifier.cc b/tensorflow/compiler/xla/service/algebraic_simplifier.cc index 72ed5ca48217298cab6dc63b1f2dd30a0730817d..da15ff7d7a2bee8f142bacc996f7fcd063598f77 100644 --- a/tensorflow/compiler/xla/service/algebraic_simplifier.cc +++ b/tensorflow/compiler/xla/service/algebraic_simplifier.cc @@ -16,6 +16,9 @@ limitations under the License. #include "tensorflow/compiler/xla/service/algebraic_simplifier.h" #include +#include +#include +#include #include #include #include @@ -23,6 +26,9 @@ 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/container/inlined_vector.h" #include "absl/memory/memory.h" #include "absl/strings/str_cat.h" #include "absl/types/optional.h" @@ -30,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" @@ -39,12 +46,14 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_opcode.h" #include "tensorflow/compiler/xla/service/hlo_query.h" #include "tensorflow/compiler/xla/service/pattern_matcher.h" +#include "tensorflow/compiler/xla/shape.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/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" @@ -67,6 +76,45 @@ bool IsAll(const HloInstruction* op, int8 value) { } } +// Checks whether `op` is a floating-point constant or broadcast of a constant +// of the form +/- 2^k for some integer k positive, negative, or zero. Such +// values are interesting because multiplying by a power of 2 just moves the +// exponent. +bool IsAllFpConstantPowerOf2(const HloInstruction* op) { + // Unwrap the broadcast if necessary. + const HloInstruction* c; + if (!Match(op, m::ConstantEffectiveScalar(&c)) && + !Match(op, m::Broadcast(m::Constant(&c).WithShape( + m::Shape().IsEffectiveScalar())))) { + return false; + } + auto val = [&]() -> absl::optional { + switch (c->shape().element_type()) { + case BF16: + return static_cast(c->literal().GetFirstElement()); + case F16: + return static_cast(c->literal().GetFirstElement()); + case F32: + return c->literal().GetFirstElement(); + case F64: + return c->literal().GetFirstElement(); + default: + // Cowardly refuse to consider complex types. + return absl::nullopt; + } + }(); + if (!val) { + return false; + } + + int exp; + double mantissa = std::frexp(*val, &exp); + // frexp returns a value in the range (-1; -0.5] U [0.5, 1). A return value + // of +/-0.5 therefore indicates that the floating point value is a power of + // 2. + return mantissa == 0.5 || mantissa == -0.5; +} + // Returns whether the given transpose produces a result which is bit-wise // identical to its operand and thus may be replaced with a bitcast. bool TransposeIsBitcast(const HloInstruction* transpose) { @@ -76,22 +124,42 @@ 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 AlgebraicSimplifier::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()); + return BitcastingOperandOfReshapeOrCopyChainHelper( + instr, instr->mutable_operand(0), options); +} - 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()); +bool IsUnstridedSlice(const HloInstruction* hlo) { + return absl::c_all_of(hlo->slice_strides(), + [](int64 stride) { return stride == 1; }); } // AlgebraicSimplifierVisitor traverses the HLO computation and reduces certain @@ -107,6 +175,8 @@ class AlgebraicSimplifierVisitor : public DfsHloVisitorWithDefault { Status HandleAdd(HloInstruction* add) override; + Status HandleAnd(HloInstruction* logical_and) override; + Status HandleBitcast(HloInstruction* bitcast) override; Status HandleBitcastConvert(HloInstruction* bitcast) override; @@ -141,10 +211,18 @@ class AlgebraicSimplifierVisitor : public DfsHloVisitorWithDefault { Status HandleMultiply(HloInstruction* multiply) override; + Status HandleNegate(HloInstruction* negate) override; + + Status HandleNot(HloInstruction* logical_not) override; + + Status HandleOr(HloInstruction* logical_or) override; + Status HandlePad(HloInstruction* pad) override; Status HandlePower(HloInstruction* power) override; + Status HandleRemainder(HloInstruction* remainder) override; + Status HandleReshape(HloInstruction* reshape) override; Status HandleReduce(HloInstruction* reduce) override; @@ -171,30 +249,29 @@ class AlgebraicSimplifierVisitor : public DfsHloVisitorWithDefault { const bool changed() const { return changed_; } // Runs the visitor on a computation. - static bool Run( - HloComputation* computation, bool is_layout_sensitive, - AlgebraicSimplifier::ValidBitcastCallback valid_bitcast_callback, - bool enable_dot_strength_reduction, bool enable_conv_simplification); + static bool Run(HloComputation* computation, + const AlgebraicSimplifierOptions& options); private: - explicit AlgebraicSimplifierVisitor( - HloComputation* computation, bool is_layout_sensitive, - AlgebraicSimplifier::ValidBitcastCallback valid_bitcast_callback, - bool enable_dot_strength_reduction, bool enable_conv_simplification) - : computation_(computation), - is_layout_sensitive_(is_layout_sensitive), - valid_bitcast_callback_(std::move(valid_bitcast_callback)), - enable_dot_strength_reduction_(enable_dot_strength_reduction), - enable_conv_simplification_(enable_conv_simplification) {} + explicit AlgebraicSimplifierVisitor(HloComputation* computation, + const AlgebraicSimplifierOptions& options) + : computation_(computation), options_(options) {} // Transforms Dots where at least one input is a vector or has a degenerate // dimension and converts it into a multiply and reduce. This should enable // more fusion than leaving the nodes as Dot operations. StatusOr HandleDotStrengthReduction(HloInstruction* dot); + // Removes dimension dim from hlo. + HloInstruction* StripDim(HloInstruction* hlo, int64 dim) { + CHECK_EQ(hlo->shape().dimensions(dim), 1); + return computation_->AddInstruction(HloInstruction::CreateReshape( + ShapeUtil::DeleteDimension(dim, hlo->shape()), hlo)); + } + // Reshapes an instruction to rank 1 if it is not already rank 1. HloInstruction* Flatten(HloInstruction* hlo) { - if (ShapeUtil::Rank(hlo->shape()) == 1) { + if (hlo->shape().rank() == 1) { return hlo; } return computation_->AddInstruction(HloInstruction::CreateReshape( @@ -214,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 @@ -224,10 +304,10 @@ class AlgebraicSimplifierVisitor : public DfsHloVisitorWithDefault { HloInstruction* new_instruction); // Returns whether the shape of the output of the given instructions are the - // same for the purposes of simplification. If is_layout_sensitive_ is true, - // then this tests shape equality including layout (ShapeUtil::Equal). If - // is_layout_sensitive_ is false, then the tests shape compatibility - // (ShapeUtil::Compatible). + // same for the purposes of simplification. If options_.is_layout_sensitive() + // is true, then this tests shape equality including layout + // (ShapeUtil::Equal). If options_.is_layout_sensitive() is false, then the + // tests shape compatibility (ShapeUtil::Compatible). bool SameShape(const HloInstruction* lhs, const HloInstruction* rhs) const; // Returns whether it was possible to transform `root` to a clamp instruction. @@ -306,26 +386,22 @@ class AlgebraicSimplifierVisitor : public DfsHloVisitorWithDefault { // Tries to use a kDot in place of the given convolution. StatusOr SimplifyConvToDot(HloInstruction* convolution); + // Tries to simplify a slice where the result of the slice is a scalar. + StatusOr TrySimplifyScalarSlice(HloInstruction* slice); + + // Tries to convert slice(reshape(X)) into reshape(slice(X)) + StatusOr TryToReorderSliceAndReshape(HloInstruction* slice); + // Current HloComputation instance the AlgebraicSimplifierVisitor is // traversing. HloComputation* computation_; + // The backend-specific options selected for the algebraic simplifier. + const AlgebraicSimplifierOptions& options_; + // Whether algebraic simplification has occurred. bool changed_ = false; - // Whether layout is considered during transformation. - bool is_layout_sensitive_; - - // Callback used to determine if a bitcast is possible. - AlgebraicSimplifier::ValidBitcastCallback valid_bitcast_callback_; - - // Disable dot strength reduction on platforms where it causes a slowdown. - bool enable_dot_strength_reduction_; - - // Disable convolution -> dot simplification on platforms where it causes a - // slowdown. - bool enable_conv_simplification_; - // Cached computation for adding two scalar F32. HloComputation* scalar_add_computation_ = nullptr; }; @@ -333,36 +409,34 @@ class AlgebraicSimplifierVisitor : public DfsHloVisitorWithDefault { } // namespace bool AlgebraicSimplifierVisitor::Run( - HloComputation* computation, bool is_layout_sensitive, - AlgebraicSimplifier::ValidBitcastCallback valid_bitcast_callback, - bool enable_dot_strength_reduction, bool enable_conv_simplification) { - AlgebraicSimplifierVisitor visitor( - computation, is_layout_sensitive, std::move(valid_bitcast_callback), - enable_dot_strength_reduction, enable_conv_simplification); + HloComputation* computation, const AlgebraicSimplifierOptions& options) { + AlgebraicSimplifierVisitor visitor(computation, options); TF_CHECK_OK(computation->Accept(&visitor)); return visitor.changed_; } bool AlgebraicSimplifierVisitor::SameShape(const HloInstruction* lhs, const HloInstruction* rhs) const { - if (is_layout_sensitive_) { + if (options_.is_layout_sensitive()) { return ShapeUtil::Equal(lhs->shape(), rhs->shape()); } else { return ShapeUtil::Compatible(lhs->shape(), rhs->shape()); } } -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)); } @@ -416,6 +490,77 @@ Status AlgebraicSimplifierVisitor::HandleAdd(HloInstruction* add) { sum_of_constants)); } + // A*C + B*C => (A+B)*C + // + // - If A, B, and C are integers, do this unconditionally. Proof of + // correctness: https://rise4fun.com/Alive/u9X. + // + // - If A, B, and C are floating point, do this if C is a scalar constant or + // broadcast of scalar constant and is equal to +/- 2^k for some (possibly + // negative) integer k. + // + // Multiplying by a power of 2 just moves the exponent, so our answer is + // exact modulo rounding of intermediate results so long as + // + // - none of the three products has an exponent which underflows (so the + // result is 0 or denormal), and + // - none of the three products overflows to inf. + // + // Proof: See algebraic_simplifier_proof_distributive_property.py. + // + // We deem these differences in rounding, underflow, and overflow + // acceptable in the ML context. + HloInstruction *b, *c; + if (((Match(lhs, m::Multiply(m::Op(&a), m::Op(&c))) && + Match(rhs, m::MultiplyAnyOrder(m::Op().Is(c), m::Op(&b)))) || + (Match(lhs, m::Multiply(m::Op(&c), m::Op(&a))) && + Match(rhs, m::MultiplyAnyOrder(m::Op().Is(c), m::Op(&b))))) && + (ShapeUtil::ElementIsIntegral(add->shape()) || + IsAllFpConstantPowerOf2(c))) { + return ReplaceWithNewInstruction( + add, HloInstruction::CreateBinary( + add->shape(), HloOpcode::kMultiply, + computation_->AddInstruction(HloInstruction::CreateBinary( + add->shape(), HloOpcode::kAdd, a, b)), + c)); + } + return Status::OK(); +} + +Status AlgebraicSimplifierVisitor::HandleAnd(HloInstruction* logical_and) { + HloInstruction *lhs, *rhs; + CHECK(Match(logical_and, m::And(m::Op(&lhs), m::Op(&rhs)))); + // Simplify logical and + if (ShapeUtil::HasPrimitiveType(lhs->shape(), xla::PRED) && + ShapeUtil::HasPrimitiveType(rhs->shape(), xla::PRED)) { + // A && True => A + VLOG(10) << "trying transform [A && True => A]: " + << logical_and->ToString(); + if (IsAll(rhs, 1) && ReplaceInstructionIfSameShape(logical_and, lhs)) { + return Status::OK(); + } + // True && A => A + VLOG(10) << "trying transform [True && A => A]: " + << logical_and->ToString(); + if (IsAll(lhs, 1) && ReplaceInstructionIfSameShape(logical_and, rhs)) { + return Status::OK(); + } + + // A && False => False + VLOG(10) << "trying transform [A && False => False]: " + << logical_and->ToString(); + if (IsAll(rhs, 0) && ReplaceInstructionIfSameShape(logical_and, rhs)) { + return Status::OK(); + } + + // False && A => False + VLOG(10) << "trying transform [False && A => False]: " + << logical_and->ToString(); + if (IsAll(lhs, 0) && ReplaceInstructionIfSameShape(logical_and, lhs)) { + return Status::OK(); + } + } + return Status::OK(); } @@ -452,9 +597,9 @@ Status AlgebraicSimplifierVisitor::HandleCopy(HloInstruction* copy) { return Status::OK(); } - if (is_layout_sensitive_ && - ReshapeOrCopyIsBitcast(copy, valid_bitcast_callback_)) { - ReplaceWithBitcast(copy); + if (HloInstruction* bitcast_operand = + BitcastingOperandOfReshapeOrCopyChain(copy, options_)) { + ReplaceWithBitcast(copy, bitcast_operand); } return Status::OK(); @@ -489,7 +634,74 @@ Status AlgebraicSimplifierVisitor::HandleConcatenate( VLOG(10) << "trying to replace " << concatenate->ToString() << " with " << replacement->ToString(); ReplaceInstructionIfSameShape(concatenate, replacement); - } else if (operands.size() == 2) { + return Status::OK(); + } + + // Check if we can merge "adjacent" slice operands which take slices from the + // same other op. For simplicity we only merge unstrided slices. + int64 concatenate_dimension = concatenate->concatenate_dimension(); + for (int64 i = 0; i < operands.size(); ++i) { + if (operands[i]->opcode() != HloOpcode::kSlice || + !IsUnstridedSlice(operands[i])) { + continue; + } + int64 slice_end = operands[i]->slice_limits(concatenate_dimension); + HloInstruction* slice_operand = operands[i]->mutable_operand(0); + int64 j = i + 1; + while (j < operands.size() && operands[j]->opcode() == HloOpcode::kSlice && + IsUnstridedSlice(operands[j]) && + operands[j]->operand(0) == slice_operand && + operands[j]->slice_starts(concatenate_dimension) == slice_end) { + // Check that all the slice_start values are the same in all other + // dimensions. This implies that the slice_limit values are also the same, + // because operands of concatenate need to have the same shape, and we + // already checked that the slices are unstrided. + bool same_other_starts = true; + for (int64 k = 0; k < operands[j]->slice_starts().size(); ++k) { + if (k == concatenate_dimension) { + continue; + } + if (operands[i]->slice_starts(k) != operands[j]->slice_starts(k)) { + same_other_starts = false; + break; + } + } + if (!same_other_starts) { + break; + } + slice_end = operands[j]->slice_limits(concatenate_dimension); + ++j; + } + if (j - i > 1) { + Shape new_slice_shape = operands[i]->shape(); + new_slice_shape.set_dimensions( + concatenate_dimension, + slice_end - operands[i]->slice_starts(concatenate_dimension)); + auto new_limit_indices = operands[i]->slice_limits(); + new_limit_indices[concatenate_dimension] = slice_end; + auto new_slice_op = + computation_->AddInstruction(HloInstruction::CreateSlice( + new_slice_shape, slice_operand, + /*start_indices=*/operands[i]->slice_starts(), + /*limit_indices=*/new_limit_indices, + /*strides=*/operands[i]->slice_strides())); + std::vector new_operands; + for (int64 k = 0; k < i; ++k) { + new_operands.push_back(operands[k]); + } + new_operands.push_back(new_slice_op); + for (int64 k = j; k < operands.size(); ++k) { + new_operands.push_back(operands[k]); + } + auto replacement = + computation_->AddInstruction(concatenate->CloneWithNewOperands( + concatenate->shape(), new_operands)); + ReplaceInstructionIfSameShape(concatenate, replacement); + return Status::OK(); + } + } + + if (operands.size() == 2) { // A binary concat with a broadcasted scalar as an operand can be converted // into a pad which is simpler to fold into other operations. bool is_effective_low_pad = Match( @@ -500,12 +712,12 @@ Status AlgebraicSimplifierVisitor::HandleConcatenate( return Status::OK(); } PaddingConfig padding_config; - for (int64 dim = 0; dim < ShapeUtil::Rank(operands[0]->shape()); ++dim) { + for (int64 dim = 0; dim < operands[0]->shape().rank(); ++dim) { auto padding_config_dim = padding_config.add_dimensions(); padding_config_dim->set_edge_padding_high(0); padding_config_dim->set_edge_padding_low(0); padding_config_dim->set_interior_padding(0); - if (dim == concatenate->concatenate_dimension()) { + if (dim == concatenate_dimension) { if (is_effective_low_pad) { padding_config_dim->set_edge_padding_low( operands[0]->shape().dimensions(dim)); @@ -528,7 +740,7 @@ Status AlgebraicSimplifierVisitor::HandleConcatenate( static HloInstruction* BuildTupleConstant(HloComputation* computation, const LiteralSlice& literal) { - if (ShapeUtil::IsTuple(literal.shape())) { + if (literal.shape().IsTuple()) { std::vector elems; elems.reserve(ShapeUtil::TupleElementCount(literal.shape())); for (int i = 0; i < ShapeUtil::TupleElementCount(literal.shape()); ++i) { @@ -545,7 +757,7 @@ static HloInstruction* BuildTupleConstant(HloComputation* computation, Status AlgebraicSimplifierVisitor::HandleConstant(HloInstruction* constant) { // Tuple constants aren't directly supported by any backend. Expand them into // explicit Tuple instructions. - if (ShapeUtil::IsTuple(constant->shape())) { + if (constant->shape().IsTuple()) { return ReplaceInstruction( constant, BuildTupleConstant(computation_, constant->literal())); } @@ -567,7 +779,7 @@ Status AlgebraicSimplifierVisitor::HandleConstant(HloInstruction* constant) { } // If a literal is an increasing sequence from zero, replace it with an iota. - if (ShapeUtil::Rank(constant->shape()) == 1 && + if (constant->shape().rank() == 1 && ShapeUtil::ElementsIn(constant->shape()) > 1 && constant->literal().IsR1Iota()) { return ReplaceWithNewInstruction( @@ -604,6 +816,79 @@ 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) { @@ -616,6 +901,60 @@ 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; + } + // 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)))) { @@ -683,6 +1022,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(); } @@ -731,21 +1073,51 @@ StatusOr AlgebraicSimplifierVisitor::HandleDotStrengthReduction( HloInstruction* dot) { HloInstruction *lhs, *rhs; CHECK(Match(dot, m::Dot(m::Op(&lhs), m::Op(&rhs)))); - int64 lhs_collapsing_dim = - dot->dot_dimension_numbers().lhs_contracting_dimensions(0); + + const auto kept_dim = [](int64 rank, int64 contracting_dimension, + absl::Span batch_dimensions) -> int64 { + for (int64 i = 0; i < rank; ++i) { + if (i != contracting_dimension && + !absl::c_linear_search(batch_dimensions, i)) { + return i; + } + } + return -1; + }; + + const int64 dot_rank = dot->shape().rank(); + 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) { + return false; + } + if (dot_rank > 2 && (lhs_rank != rhs_rank || lhs_rank != dot_rank)) { + return false; + } + int64 lhs_collapsing_dim = dnums.lhs_contracting_dimensions(0); + int64 lhs_kept_dim = kept_dim(lhs_rank, lhs_collapsing_dim, + AsInt64Slice(dnums.lhs_batch_dimensions())); + // If there is no non-contracting dimension in rank 2, do not strength reduce. + if (lhs_kept_dim == -1 && lhs_rank > 1) { + return false; + } if (lhs->IsRank2Transpose()) { lhs = lhs->mutable_operand(0); - lhs_collapsing_dim = 1 - lhs_collapsing_dim; + std::swap(lhs_collapsing_dim, lhs_kept_dim); } - const int64 lhs_kept_dim = 1 - lhs_collapsing_dim; - int64 rhs_collapsing_dim = - dot->dot_dimension_numbers().rhs_contracting_dimensions(0); + int64 rhs_collapsing_dim = dnums.rhs_contracting_dimensions(0); + int64 rhs_kept_dim = kept_dim(rhs_rank, rhs_collapsing_dim, + AsInt64Slice(dnums.rhs_batch_dimensions())); + // If there is no non-contracting dimension in rank 2, do not strength reduce. + if (rhs_kept_dim == -1 && rhs_rank > 1) { + return false; + } if (rhs->IsRank2Transpose()) { rhs = rhs->mutable_operand(0); - rhs_collapsing_dim = 1 - rhs_collapsing_dim; + std::swap(rhs_collapsing_dim, rhs_kept_dim); } - const int64 rhs_kept_dim = 1 - rhs_collapsing_dim; auto as_type = [&](HloInstruction* hlo, const PrimitiveType element_type) { if (hlo->shape().element_type() == element_type) { @@ -768,10 +1140,15 @@ StatusOr AlgebraicSimplifierVisitor::HandleDotStrengthReduction( return AddReduce(as_type(hlo, F32), dim); }; + auto broadcast = [&](HloInstruction* hlo, const Shape& shape, + absl::Span dims) { + return computation_->AddInstruction( + HloInstruction::CreateBroadcast(shape, hlo, dims)); + }; + auto broadcast_to_dim = [&](HloInstruction* hlo, const Shape& shape, int64 dim) { - return computation_->AddInstruction( - HloInstruction::CreateBroadcast(shape, hlo, {dim})); + return broadcast(hlo, shape, {dim}); }; auto multiply = [&](HloInstruction* local_lhs, HloInstruction* local_rhs) { @@ -782,11 +1159,9 @@ StatusOr AlgebraicSimplifierVisitor::HandleDotStrengthReduction( // Strength reduce dot(a[K] , b[K]) = // reshape(result.shape, // reduce_sum(multiply(a, b), {0})) - if (ShapeUtil::Rank(rhs->shape()) == 1 && - ShapeUtil::Rank(lhs->shape()) == 1) { - TF_RETURN_IF_ERROR( - ReplaceInstruction(dot, reshape_if_necessary(add_reduce_in_f32( - multiply(Flatten(lhs), Flatten(rhs)), 0)))); + if (rhs_rank == 1 && lhs_rank == 1) { + TF_RETURN_IF_ERROR(ReplaceInstruction( + dot, reshape_if_necessary(add_reduce_in_f32(multiply(lhs, rhs), 0)))); return true; } @@ -800,8 +1175,7 @@ StatusOr AlgebraicSimplifierVisitor::HandleDotStrengthReduction( // Simplify outer product into multiply with implicit broadcasting. // // A dot(a[M, 1], b[1, N]) = multiply(a [M,1], b [1, N]) - if (ShapeUtil::Rank(rhs->shape()) == 2 && - rhs->shape().dimensions(rhs_collapsing_dim) == 1) { + if (rhs_rank == 2 && rhs->shape().dimensions(rhs_collapsing_dim) == 1) { TF_RETURN_IF_ERROR(ReplaceInstruction( dot, multiply(broadcast_to_dim(Flatten(lhs), dot->shape(), 0), broadcast_to_dim(Flatten(rhs), dot->shape(), 1)))); @@ -815,10 +1189,9 @@ StatusOr AlgebraicSimplifierVisitor::HandleDotStrengthReduction( // {0}) // ) // ) - if (ShapeUtil::Rank(lhs->shape()) == 1 || - (ShapeUtil::Rank(lhs->shape()) == 2 && - lhs->shape().dimensions(lhs_kept_dim) == 1)) { - if (ShapeUtil::Rank(rhs->shape()) == 1) { + if (lhs_rank == 1 || + (lhs_rank == 2 && lhs->shape().dimensions(lhs_kept_dim) == 1)) { + if (rhs->shape().rank() == 1) { TF_RETURN_IF_ERROR( ReplaceInstruction(dot, reshape_if_necessary(add_reduce_in_f32( multiply(Flatten(lhs), rhs), 0)))); @@ -837,9 +1210,8 @@ StatusOr AlgebraicSimplifierVisitor::HandleDotStrengthReduction( // reshape(result.shape, // reduce_sum(multiply(a, broadcast(reshape([K],b), {1})), {0}) // ) - if (ShapeUtil::Rank(rhs->shape()) == 1 || - (ShapeUtil::Rank(rhs->shape()) == 2 && - rhs->shape().dimensions(rhs_kept_dim) == 1)) { + if (rhs_rank == 1 || + (rhs_rank == 2 && rhs->shape().dimensions(rhs_kept_dim) == 1)) { TF_RETURN_IF_ERROR(ReplaceInstruction( dot, reshape_if_necessary(add_reduce_in_f32( multiply(lhs, broadcast_to_dim(Flatten(rhs), lhs->shape(), @@ -847,6 +1219,97 @@ StatusOr AlgebraicSimplifierVisitor::HandleDotStrengthReduction( lhs_collapsing_dim)))); return true; } + + // Only consider kDot with batch dimension. + if (dot_rank <= 2) { + return false; + } + + CHECK_EQ(rhs_rank, lhs_rank); + CHECK_EQ(dot_rank, lhs_rank); + // If there is more than one non-contracting dimension or the batch dimensions + // are not equal, bail out since transposes may be required to do a strength + // reduction. + if (dnums.rhs_batch_dimensions_size() + 2 != dot_rank || + !absl::c_equal(dnums.lhs_batch_dimensions(), + dnums.rhs_batch_dimensions())) { + return false; + } + + auto broadcast_dims = [](int64 rank, int64 non_broadcast_dim) { + absl::InlinedVector dims; + for (int64 i = 0; i < rank; ++i) { + if (i != non_broadcast_dim) { + dims.push_back(i); + } + } + return dims; + }; + + // If the contracting dimension is 1, remove the degnerate dimnesions from the + // lhs and rhs, broadcast each to the result shape and multiply. + if (lhs->shape().dimensions(lhs_collapsing_dim) == 1 && + (rhs_kept_dim == rhs_rank - 1 || + (rhs_collapsing_dim == rhs_rank - 1 && rhs_kept_dim == rhs_rank - 2))) { + CHECK_EQ(rhs->shape().dimensions(rhs_collapsing_dim), 1); + const int64 lhs_kept_dim_in_output = + lhs_kept_dim > lhs_collapsing_dim ? (lhs_kept_dim - 1) : lhs_kept_dim; + absl::InlinedVector lhs_broadcast_dims; + for (const int64 dim : dnums.lhs_batch_dimensions()) { + lhs_broadcast_dims.push_back(dim > lhs_collapsing_dim ? (dim - 1) : dim); + } + absl::InlinedVector rhs_broadcast_dims = lhs_broadcast_dims; + lhs_broadcast_dims.push_back(lhs_kept_dim_in_output); + absl::c_sort(lhs_broadcast_dims); + rhs_broadcast_dims.push_back(dot_rank - 1); + absl::c_sort(rhs_broadcast_dims); + TF_RETURN_IF_ERROR(ReplaceInstruction( + dot, reshape_if_necessary( + multiply(broadcast(StripDim(lhs, lhs_collapsing_dim), + dot->shape(), lhs_broadcast_dims), + broadcast(StripDim(rhs, rhs_collapsing_dim), + dot->shape(), rhs_broadcast_dims))))); + return true; + } + + // If the lhs and rhs non-contracting dimensions are both one, strip each one, + // multiply and then reduce the collapsing dimension + if (lhs->shape().dimensions(lhs_kept_dim) == 1 && + rhs->shape().dimensions(rhs_kept_dim) == 1 && + lhs_kept_dim == rhs_kept_dim) { + auto new_lhs = StripDim(lhs, lhs_kept_dim); + auto new_rhs = StripDim(rhs, rhs_kept_dim); + const int64 reduce_dim = rhs_kept_dim < rhs_collapsing_dim + ? (rhs_collapsing_dim - 1) + : rhs_collapsing_dim; + TF_RETURN_IF_ERROR( + ReplaceInstruction(dot, reshape_if_necessary(add_reduce_in_f32( + multiply(new_lhs, new_rhs), reduce_dim)))); + return true; + } + + // If the lhs non-contracting dimensions is one, strip the one, brodcast to + // the rhs shape, multiply and then reduce the collapsing dimension + if (lhs->shape().dimensions(lhs_kept_dim) == 1) { + auto new_lhs = broadcast(StripDim(lhs, lhs_kept_dim), rhs->shape(), + broadcast_dims(rhs_rank, rhs_kept_dim)); + TF_RETURN_IF_ERROR(ReplaceInstruction( + dot, reshape_if_necessary(add_reduce_in_f32(multiply(new_lhs, rhs), + rhs_collapsing_dim)))); + return true; + } + + // If the rhs non-contracting dimensions is one, strip the one, brodcast to + // the lhs shape, multiply and then reduce the collapsing dimension + if (rhs->shape().dimensions(rhs_kept_dim) == 1) { + auto new_rhs = broadcast(StripDim(rhs, rhs_kept_dim), lhs->shape(), + broadcast_dims(lhs_rank, lhs_kept_dim)); + TF_RETURN_IF_ERROR(ReplaceInstruction( + dot, reshape_if_necessary(add_reduce_in_f32(multiply(lhs, new_rhs), + lhs_collapsing_dim)))); + return true; + } + return false; } @@ -1065,6 +1528,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 = @@ -1082,8 +1548,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 @@ -1092,23 +1556,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; @@ -1125,25 +1585,31 @@ Status AlgebraicSimplifierVisitor::HandleDot(HloInstruction* dot) { HloInstruction *lhs, *rhs; CHECK(Match(dot, m::Dot(m::Op(&lhs), m::Op(&rhs)))); - // Only optimize F32 or BF16 dot operations where the dot, rhs and lhs are - // rank 2 or below. - if ((dot->shape().element_type() != F32 && - dot->shape().element_type() != BF16) || - ShapeUtil::Rank(lhs->shape()) > 2 || ShapeUtil::Rank(rhs->shape()) > 2 || - ShapeUtil::Rank(dot->shape()) > 2) { - return Status::OK(); - } - // Replace a zero element dot with a broadcast of the constant 0. if (ShapeUtil::IsZeroElementArray(dot->shape()) || ShapeUtil::IsZeroElementArray(lhs->shape()) || ShapeUtil::IsZeroElementArray(rhs->shape())) { - auto zero = computation_->AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(0.0f))); + auto zero = computation_->AddInstruction(HloInstruction::CreateConstant( + LiteralUtil::Zero(dot->shape().element_type()))); return ReplaceWithNewInstruction( dot, HloInstruction::CreateBroadcast(dot->shape(), zero, {})); } + // Only optimize F32 or BF16 dot operations where the dot, rhs and lhs are + // rank 2 or below. + if (dot->shape().element_type() != F32 && + dot->shape().element_type() != BF16) { + return Status::OK(); + } + if (lhs->shape().rank() > 2 || rhs->shape().rank() > 2 || + dot->shape().rank() > 2) { + if (options_.enable_dot_strength_reduction() && + !options_.is_layout_sensitive()) { + TF_RETURN_IF_ERROR(HandleDotStrengthReduction(dot).status()); + } + return Status::OK(); + } + TF_ASSIGN_OR_RETURN(HloInstruction * dot_of_concat_optimized, OptimizeDotOfConcat(dot)); if (dot_of_concat_optimized) { @@ -1163,7 +1629,8 @@ Status AlgebraicSimplifierVisitor::HandleDot(HloInstruction* dot) { return ReplaceInstruction(dot, dot_of_gather_optimized); } - if (enable_dot_strength_reduction_ && !is_layout_sensitive_) { + if (options_.enable_dot_strength_reduction() && + !options_.is_layout_sensitive()) { TF_ASSIGN_OR_RETURN(bool did_strength_reduction, HandleDotStrengthReduction(dot)); if (did_strength_reduction) { @@ -1225,6 +1692,64 @@ Status AlgebraicSimplifierVisitor::HandleMultiply(HloInstruction* multiply) { return Status::OK(); } +Status AlgebraicSimplifierVisitor::HandleNegate(HloInstruction* negate) { + // negate(negate(x)) => x + HloInstruction* x; + if (Match(negate, m::Negate(m::Negate(m::Op(&x)))) && + ReplaceInstructionIfSameShape(negate, x)) { + return Status::OK(); + } + return Status::OK(); +} + +Status AlgebraicSimplifierVisitor::HandleNot(HloInstruction* logical_not) { + // not(not(x)) => x + HloInstruction* x; + if (Match(logical_not, m::Not(m::Not(m::Op(&x)))) && + ReplaceInstructionIfSameShape(logical_not, x)) { + return Status::OK(); + } + return Status::OK(); +} + +Status AlgebraicSimplifierVisitor::HandleOr(HloInstruction* logical_or) { + HloInstruction *lhs, *rhs; + CHECK(Match(logical_or, m::Or(m::Op(&lhs), m::Op(&rhs)))); + + // Simplify logical or + if (ShapeUtil::HasPrimitiveType(lhs->shape(), xla::PRED) && + ShapeUtil::HasPrimitiveType(rhs->shape(), xla::PRED)) { + // A || True => True + VLOG(10) << "trying transform [A || True => True]: " + << logical_or->ToString(); + if (IsAll(rhs, 1) && ReplaceInstructionIfSameShape(logical_or, rhs)) { + return Status::OK(); + } + // True || A => True + VLOG(10) << "trying transform [True || A => True]: " + << logical_or->ToString(); + if (IsAll(lhs, 1) && ReplaceInstructionIfSameShape(logical_or, lhs)) { + return Status::OK(); + } + + // A || False => A + VLOG(10) << "trying transform [A || False => A]: " + << logical_or->ToString(); + if (IsAll(rhs, 0) && ReplaceInstructionIfSameShape(logical_or, lhs)) { + return Status::OK(); + } + + // False || A => A + VLOG(10) << "trying transform [False || A => A]: " + << logical_or->ToString(); + if (IsAll(lhs, 0) && ReplaceInstructionIfSameShape(logical_or, rhs)) { + return Status::OK(); + } + } + + return Status::OK(); +} + Status AlgebraicSimplifierVisitor::HandleLog(HloInstruction* log) { // ln(exp(A)) => A VLOG(10) << "trying transform [ln(exp(A)) => A]: " << log->ToString(); @@ -1313,7 +1838,7 @@ bool OutputIsPermutationOfOperandElements(HloInstruction* instruction, case HloOpcode::kTranspose: return true; case HloOpcode::kSort: - return (!ShapeUtil::IsTuple(instruction->shape())); + return (!instruction->shape().IsTuple()); default: return false; } @@ -1359,8 +1884,7 @@ Status AlgebraicSimplifierVisitor::HandleBroadcast(HloInstruction* broadcast) { // A degenerate broadcast that has the same input and output rank can be // converted into a transpose. - if (ShapeUtil::Rank(broadcast->shape()) == - ShapeUtil::Rank(operand->shape()) && + if (broadcast->shape().rank() == operand->shape().rank() && ShapeUtil::ElementsIn(broadcast->shape()) == ShapeUtil::ElementsIn(operand->shape())) { VLOG(10) << "transform broadcast(X) -> transpose(X) where " @@ -1509,6 +2033,27 @@ Status AlgebraicSimplifierVisitor::HandlePad(HloInstruction* pad) { pad, HloInstruction::CreateBroadcast(pad->shape(), pad->mutable_operand(1), {})); } + + // Interior padding on one sized dimensions have no effect. As a result it + // makes other simplifications possible if there is no interior padding. + if (HasInteriorPadding(pad->padding_config())) { + PaddingConfig padding_config = pad->padding_config(); + bool cleared_interior_padding = false; + for (int64 i = 0; i < pad->shape().rank(); ++i) { + if (padding_config.dimensions(i).interior_padding() > 0 && + pad->operand(0)->shape().dimensions(i) == 1) { + cleared_interior_padding = true; + padding_config.mutable_dimensions(i)->set_interior_padding(0); + } + } + if (cleared_interior_padding) { + return ReplaceWithNewInstruction( + pad, + HloInstruction::CreatePad(pad->shape(), pad->mutable_operand(0), + pad->mutable_operand(1), padding_config)); + } + } + // Eliminate nop pads (padding all zero), and replace a pad with negative // padding with a pad with non-negative padding followed by a slice. bool all_zero = true; @@ -1745,6 +2290,137 @@ 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); @@ -1769,6 +2445,7 @@ Status AlgebraicSimplifierVisitor::HandleReshape(HloInstruction* reshape) { reshape, HloInstruction::CreateReshape(reshape->shape(), operand->mutable_operand(0))); } + if (operand->opcode() == HloOpcode::kRng && operand->user_count() == 1) { *operand->mutable_shape() = reshape->shape(); return ReplaceInstruction(reshape, operand); @@ -1800,12 +2477,10 @@ Status AlgebraicSimplifierVisitor::HandleReshape(HloInstruction* reshape) { } // Make this a bitcast if possible. - if (is_layout_sensitive_ && - ReshapeOrCopyIsBitcast(reshape, valid_bitcast_callback_)) { - ReplaceWithBitcast(reshape); - return Status::OK(); + if (HloInstruction* bitcast_operand = + BitcastingOperandOfReshapeOrCopyChain(reshape, options_)) { + ReplaceWithBitcast(reshape, bitcast_operand); } - return Status::OK(); } @@ -1815,25 +2490,171 @@ 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(); } +StatusOr AlgebraicSimplifierVisitor::TrySimplifyScalarSlice( + HloInstruction* slice) { + // Only try to do this for effective scalars. We could do the same for slicing + // out larger pieces of padding (replacing with a broadcast of the padding + // value), but this is probably not worth it. + if (!ShapeUtil::IsEffectiveScalar(slice->shape())) { + return false; + } + + if (slice->operand(0)->opcode() == HloOpcode::kPad) { + VLOG(10) << "Trying to simplify scalar slice of pad"; + // Check there's no internal padding. Again, we could handle that too, since + // everything is statically known, but it's not worth it. + auto pad = Cast(slice->mutable_operand(0)); + auto padding_config = pad->padding_config(); + int64 rank = padding_config.dimensions_size(); + if (HasInteriorPadding(padding_config)) { + VLOG(10) << "Not folding scalar slice of pad, pad has interior padding"; + return false; + } + + // Check whether the scalar we're slicing out falls into the padding. + bool in_padding = [&]() { + for (int64 i = 0; i < rank; ++i) { + int64 start = slice->slice_starts(i); + int64 low = padding_config.dimensions(i).edge_padding_low(); + int64 data = pad->operand(0)->shape().dimensions(i); + if (start >= low && start < low + data) { + return false; + } + } + return true; + }(); + + if (in_padding) { + VLOG(10) << "Folding scalar slice of pad into padding value"; + TF_RETURN_IF_ERROR(ReplaceWithNewInstruction( + slice, HloInstruction::CreateReshape(slice->shape(), + pad->mutable_padding_value()))); + return true; + } else { + // We already know the output of the slice is scalar. If the padded + // value is scalar, and it's not in the padding, then it's exactly the + // output value. + bool replaced = + ReplaceInstructionIfSameShape(slice, pad->mutable_operand(0)); + if (replaced) { + VLOG(10) << "Folding scalar slice of pad into padded value"; + } else { + VLOG(10) << "Not folding scalar slice of pad into padded value as they " + "have different shapes."; + } + return replaced; + } + } + + if (slice->operand(0)->opcode() == HloOpcode::kConcatenate) { + VLOG(10) << "Trying to simplify scalar slice of concat"; + // Only do this for R1, there's no chance of this being useful otherwise. + if (slice->shape().rank() != 1) { + VLOG(10) << "Not folding, slice is not rank 1"; + return false; + } + HloConcatenateInstruction* concat = + Cast(slice->mutable_operand(0)); + int64 operand_start = 0; + int64 operand_num = 0; + // Weird loop structure to avoid annoying off-by-one errors. + while (true) { + TF_RET_CHECK(operand_num < concat->operand_count()); + const HloInstruction* operand = concat->operand(operand_num); + int64 next_operand_start = operand_start + operand->shape().dimensions(0); + if (next_operand_start > slice->slice_starts(0)) { + break; + } + operand_start = next_operand_start; + operand_num++; + } + + bool replaced = ReplaceInstructionIfSameShape( + slice, concat->mutable_operand(operand_num)); + if (replaced) { + VLOG(10) << "Folding scalar slice of concat into concat operand"; + } else { + VLOG(10) << "Folding scalar slice of concat into slice of concat operand"; + TF_RETURN_IF_ERROR(ReplaceWithNewInstruction( + slice, HloInstruction::CreateSlice( + slice->shape(), concat->mutable_operand(operand_num), + {slice->slice_starts(0) - operand_start}, + {slice->slice_starts(0) - operand_start + 1}, + slice->slice_strides()))); + } + return true; + } + + return false; +} + +StatusOr AlgebraicSimplifierVisitor::TryToReorderSliceAndReshape( + HloInstruction* slice) { + CHECK_EQ(slice->opcode(), HloOpcode::kSlice); + if (!IsUnstridedSlice(slice)) { + return false; + } + HloInstruction* reshape = slice->mutable_operand(0); + if (reshape->opcode() != HloOpcode::kReshape) { + return false; + } + HloInstruction* new_slice_operand = reshape->mutable_operand(0); + int64 slice_rank = slice->shape().rank(); + std::vector sliced_dims; + for (int64 i = 0; i < slice_rank; ++i) { + if (slice->slice_starts(i) != 0 || + slice->slice_limits(i) != reshape->shape().dimensions(i)) { + sliced_dims.push_back(i); + } + } + + if (sliced_dims.size() == 1 && sliced_dims[0] == 0 && + slice->slice_starts(0) == 0) { + const Shape& new_slice_shape = new_slice_operand->shape(); + const int64 rank = new_slice_shape.rank(); + std::vector new_slice_starts(rank, 0); + std::vector new_slice_stides(rank, 1); + std::vector new_slice_limits(new_slice_shape.dimensions().begin(), + new_slice_shape.dimensions().end()); + int64 slice_elements = ShapeUtil::ElementsIn(slice->shape()); + for (int64 i = rank - 1; i >= 0; --i) { + if (slice_elements >= new_slice_limits[i]) { + if (slice_elements % new_slice_limits[i] != 0) { + return false; + } + slice_elements /= new_slice_limits[i]; + } else { + new_slice_limits[i] = slice_elements; + slice_elements = 1; + } + } + HloInstruction* new_slice = + computation_->AddInstruction(HloInstruction::CreateSlice( + ShapeUtil::MakeShape(new_slice_shape.element_type(), + new_slice_limits), + new_slice_operand, new_slice_starts, new_slice_limits, + new_slice_stides)); + TF_RETURN_IF_ERROR(ReplaceWithNewInstruction( + slice, HloInstruction::CreateReshape(slice->shape(), new_slice))); + return true; + } + return false; +} + Status AlgebraicSimplifierVisitor::HandleSlice(HloInstruction* slice) { // Delete no-op slices, i.e. where shape = operand shape. if (ReplaceInstructionIfSameShape(slice, slice->mutable_operand(0))) { return Status::OK(); } - auto is_unstrided_slice = [](const HloInstruction* hlo) { - return absl::c_all_of(hlo->slice_strides(), - [](int64 stride) { return stride == 1; }); - }; if (slice->operand(0)->opcode() == HloOpcode::kSlice && - is_unstrided_slice(slice) && is_unstrided_slice(slice->operand(0))) { + IsUnstridedSlice(slice) && IsUnstridedSlice(slice->operand(0))) { HloInstruction* operand_slice = slice->mutable_operand(0); std::vector new_slice_starts = slice->slice_starts(); std::vector new_slice_limits = slice->slice_limits(); @@ -1846,6 +2667,16 @@ Status AlgebraicSimplifierVisitor::HandleSlice(HloInstruction* slice) { slice->shape(), operand_slice->mutable_operand(0), new_slice_starts, new_slice_limits, slice->slice_strides())); } + + TF_ASSIGN_OR_RETURN(bool replaced, TrySimplifyScalarSlice(slice)); + if (replaced) { + return Status::OK(); + } + + TF_ASSIGN_OR_RETURN(replaced, TryToReorderSliceAndReshape(slice)); + if (replaced) { + return Status::OK(); + } return Status::OK(); } @@ -1886,7 +2717,7 @@ Status AlgebraicSimplifierVisitor::HandleDynamicUpdateSlice( Status AlgebraicSimplifierVisitor::HandleReduce(HloInstruction* reduce) { // TODO(b/112040122): Most of those optimizations can be done for multi-output // reduces. - if (ShapeUtil::IsTuple(reduce->shape())) { + if (reduce->shape().IsTuple()) { return Status::OK(); } @@ -1904,8 +2735,7 @@ Status AlgebraicSimplifierVisitor::HandleReduce(HloInstruction* reduce) { // 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 (ShapeUtil::Rank(reduce->shape()) <= 1 && - arg->opcode() == HloOpcode::kTranspose) { + if (reduce->shape().rank() <= 1 && arg->opcode() == HloOpcode::kTranspose) { auto transpose_dimensions = arg->dimensions(); std::vector new_reduce_dimensions; for (auto dim : dimensions) { @@ -1935,9 +2765,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) { @@ -1964,8 +2794,8 @@ Status AlgebraicSimplifierVisitor::HandleReduce(HloInstruction* reduce) { std::vector> unmodified_dims = ShapeUtil::DimensionsUnmodifiedByReshape(arg->operand(0)->shape(), arg->shape()); - std::vector arg_dim_in_output(ShapeUtil::Rank(arg->shape()), true); - std::vector arg_dim_unmodified(ShapeUtil::Rank(arg->shape()), false); + std::vector arg_dim_in_output(arg->shape().rank(), true); + std::vector arg_dim_unmodified(arg->shape().rank(), false); for (auto dim : dimensions) { arg_dim_in_output[dim] = false; } @@ -1983,15 +2813,15 @@ 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); } } std::vector new_reduce_dimensions; - for (int64 i = 0; i < ShapeUtil::Rank(arg->operand(0)->shape()); ++i) { - if (dimensions_not_to_reduce.count(i) == 0) { + for (int64 i = 0; i < arg->operand(0)->shape().rank(); ++i) { + if (!dimensions_not_to_reduce.contains(i)) { new_reduce_dimensions.push_back(i); } } @@ -2045,6 +2875,55 @@ Status AlgebraicSimplifierVisitor::HandleReduceWindow( function)); } + if (options_.enable_window_reduce_to_reduce_replacement()) { + // A reduce window can be expressed as a reduce and a reshape if all + // dimensions either have a window size of one or the entire dimension. If + // there is no stride, dilation, or padding, this is as easy as checking the + // size of the output shape and window dimension. + // + // The reshape is a bitcast since it adds one-sized dimensions. Often these + // ones are immediately removed as well with another reshape. The + // implementation of reduce tends to be slightly more efficient at reducing + // entire dimensions compared to reduce window. + auto effective_reduce_dims = [&] { + if (window_util::HasStride(window) || window_util::HasDilation(window) || + window_util::HasPadding(window)) { + return absl::InlinedVector{}; + } + absl::InlinedVector reduce_dims; + for (int64 i = 0; i < window.dimensions_size(); ++i) { + if (window.dimensions(i).size() == 1) { + continue; + } else if (reduce_window->shape().dimensions(i) == 1) { + reduce_dims.push_back(i); + } else { + return absl::InlinedVector{}; + } + } + return reduce_dims; + }(); + + // If a reduce window can be expressed as a reduce, do so and reshape the + // output. + if (!effective_reduce_dims.empty()) { + Shape reduce_shape = ShapeUtil::FilterDimensions( + [&](int64 dim) { + return !absl::c_linear_search(effective_reduce_dims, dim); + }, + reduce_window->shape()); + HloInstruction* reduce = + computation_->AddInstruction(HloInstruction::CreateReduce( + /*shape=*/reduce_shape, + /*operand=*/operand, + /*init_value=*/reduce_window->mutable_operand(1), + /*dimensions_to_reduce=*/effective_reduce_dims, + /*reduce_computation=*/function)); + return ReplaceWithNewInstruction( + reduce_window, + HloInstruction::CreateReshape(reduce_window->shape(), reduce)); + } + } + // This optimization folds a pad op into reduce_window. HloInstruction* pad; const HloInstruction* convert = nullptr; @@ -2180,7 +3059,7 @@ Status AlgebraicSimplifierVisitor::HandleReduceWindow( // Carry out the folding of the pad into reduce_window. VLOG(10) << "Folding pad into reduce-window."; Window new_window = window; - const int64 rank = ShapeUtil::Rank(reduce_window->shape()); + const int64 rank = reduce_window->shape().rank(); TF_RET_CHECK(pad_config.dimensions_size() == rank); TF_RET_CHECK(window.dimensions_size() == rank); for (int64 i = 0; i < rank; ++i) { @@ -2229,9 +3108,128 @@ Status AlgebraicSimplifierVisitor::HandleSort(HloInstruction* sort) { return ReplaceWithNewInstruction( sort, HloInstruction::CreateTuple(sort->operands())); } + + 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(); } +namespace { +bool OnlyPermutesMoreThanOneDegenerateDim(const Shape& shape, + absl::Span perm) { + std::vector new_permutation; + int64 degenerate_count = 0; + for (int64 i = 0; i < perm.size(); ++i) { + if (shape.dimensions(i) != 1) { + new_permutation.push_back(perm[i]); + } else { + ++degenerate_count; + } + } + return degenerate_count > 1 && absl::c_is_sorted(new_permutation); +} +} // namespace + Status AlgebraicSimplifierVisitor::HandleTranspose(HloInstruction* transpose) { auto operand = transpose->mutable_operand(0); if (std::is_sorted(transpose->dimensions().begin(), @@ -2248,12 +3246,21 @@ Status AlgebraicSimplifierVisitor::HandleTranspose(HloInstruction* transpose) { transpose->dimensions()))); } + // Replace transpose with a reshape if more than one degenerate method is + // permuted. + if (OnlyPermutesMoreThanOneDegenerateDim(transpose->shape(), + transpose->dimensions())) { + return ReplaceWithNewInstruction( + transpose, HloInstruction::CreateReshape( + transpose->shape(), transpose->mutable_operand(0))); + } + if (operand->opcode() == HloOpcode::kRng && operand->user_count() == 1) { *operand->mutable_shape() = transpose->shape(); return ReplaceInstruction(transpose, operand); } - if (is_layout_sensitive_ && TransposeIsBitcast(transpose)) { + if (options_.is_layout_sensitive() && TransposeIsBitcast(transpose)) { ReplaceWithBitcast(transpose); return Status::OK(); } @@ -2402,13 +3409,13 @@ StatusOr AlgebraicSimplifierVisitor::SimplifyConvToDot( const ConvolutionDimensionNumbers& dnums = convolution->convolution_dimension_numbers(); - if (!enable_conv_simplification_) { + if (!options_.enable_conv_simplification()) { return false; } // TODO(b/31337498): For now, we cowardly refuse to do this optimization in // layout-insensitive mode, for fear of adding nontrivial reshapes. - if (!is_layout_sensitive_) { + if (!options_.is_layout_sensitive()) { return false; } @@ -2495,15 +3502,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 (!valid_bitcast_callback_(input_shape, new_input_shape) || - !valid_bitcast_callback_(filter_shape, new_filter_shape) || - !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; @@ -2606,9 +3604,7 @@ StatusOr AlgebraicSimplifier::Run(HloModule* module) { "AlgebraicSimplifier::Run(), before:\n" + module->ToString()); bool changed = false; for (auto* comp : module->MakeNonfusionComputations()) { - if (AlgebraicSimplifierVisitor::Run( - comp, is_layout_sensitive_, valid_bitcast_callback_, - enable_dot_strength_reduction_, enable_conv_simplification_)) { + if (AlgebraicSimplifierVisitor::Run(comp, options_)) { changed = true; } } diff --git a/tensorflow/compiler/xla/service/algebraic_simplifier.h b/tensorflow/compiler/xla/service/algebraic_simplifier.h index 9f8d0ee88bdebcf17310cd0407b1b99e4b0a7b5f..ff3f638b22e290f6f6237a5a72a257aa23ecd78b 100644 --- a/tensorflow/compiler/xla/service/algebraic_simplifier.h +++ b/tensorflow/compiler/xla/service/algebraic_simplifier.h @@ -23,29 +23,93 @@ limitations under the License. namespace xla { -// A pass which performs algebraic simplifications. -class AlgebraicSimplifier : public HloModulePass { +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( + 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 - // transformation. Otherwise, layout is ignored. If valid_bitcast_callback - // returns true, then the pass will replace reshapes and transposes with - // bitcasts. - AlgebraicSimplifier(bool is_layout_sensitive, - ValidBitcastCallback valid_bitcast_callback, - bool enable_dot_strength_reduction = true, - bool enable_conv_simplification = true) - : is_layout_sensitive_(is_layout_sensitive), - valid_bitcast_callback_(std::move(valid_bitcast_callback)), - enable_dot_strength_reduction_(enable_dot_strength_reduction), - enable_conv_simplification_(enable_conv_simplification) {} + // transformation. Otherwise, layout is ignored. + 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_; + } + + // Enable convolution simplification on platforms where it is profitable. + void set_enable_conv_simplification(bool enable_conv_simplification) { + enable_conv_simplification_ = enable_conv_simplification; + } + bool enable_conv_simplification() const { + 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( + bool enable_window_reduce_to_reduce_replacement) { + 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: + 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}; +}; + +// A pass which performs algebraic simplifications. +class AlgebraicSimplifier : public HloModulePass { + public: + // If is_layout_sensitive is true, then the simplifier preserves layout during + // transformation. Otherwise, layout is ignored. + explicit AlgebraicSimplifier(const AlgebraicSimplifierOptions& options) + : options_(options) {} ~AlgebraicSimplifier() override = default; absl::string_view name() const override { return "algsimp"; } @@ -54,14 +118,7 @@ class AlgebraicSimplifier : public HloModulePass { StatusOr Run(HloModule* module) override; private: - bool is_layout_sensitive_; - ValidBitcastCallback valid_bitcast_callback_; - - // Enable dot simplification on platforms where it is profitable. - bool enable_dot_strength_reduction_; - - // Enable convolution simplification on platforms where it is profitable. - bool enable_conv_simplification_; + AlgebraicSimplifierOptions options_; }; } // namespace xla diff --git a/tensorflow/compiler/xla/service/algebraic_simplifier_proof_distributive_property.py b/tensorflow/compiler/xla/service/algebraic_simplifier_proof_distributive_property.py new file mode 100644 index 0000000000000000000000000000000000000000..5da13da041b4ded813876af7ca379025187545ab --- /dev/null +++ b/tensorflow/compiler/xla/service/algebraic_simplifier_proof_distributive_property.py @@ -0,0 +1,82 @@ +# 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. +# ============================================================================== +"""Proof that transforming (A*C)+(B*C) <=> (A+B)*C is "safe" if C=2^k. + +Specifically, for all floating-point values A, B, and C, if + + - C is equal to +/- 2^k for some (possibly negative) integer k, and + - A, B, C, A*C, B*C, and A+B are not subnormal, zero, or inf, + +then there exists a rounding mode rm in [RTZ, RNE] such that + + (A*C) + (B*C) == (A+B) * C (computed with rounding mode rm). + +Informally, this means that the equivalence holds for powers of 2 C, modulo +flushing to zero or inf, and modulo rounding of intermediate results. + +Requires z3 python bindings; try `pip install z3-solver`. +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function +import z3 + +# We do float16 because it lets the solver run much faster. These results +# should generalize to fp32 and fp64, and you can verify this by changing the +# value of FLOAT_TY (and then waiting a while). +FLOAT_TY = z3.Float16 + +a = z3.FP("a", FLOAT_TY()) +b = z3.FP("b", FLOAT_TY()) +c = z3.FP("c", FLOAT_TY()) + +s = z3.Solver() + +# C must be a power of 2, i.e. significand bits must all be 0. +s.add(z3.Extract(FLOAT_TY().sbits() - 1, 0, z3.fpToIEEEBV(c)) == 0) + +for rm in [z3.RTZ(), z3.RNE()]: + z3.set_default_rounding_mode(rm) + before = a * c + b * c + after = (a + b) * c + + # Check that before == after, allowing that 0 == -0. + s.add( + z3.Not( + z3.Or( + before == after, # + z3.And(z3.fpIsZero(before), z3.fpIsZero(after))))) + + for x in [ + (a * c), + (b * c), + (a + b), + ]: + s.add(z3.Not(z3.fpIsSubnormal(x))) + s.add(z3.Not(z3.fpIsZero(x))) + s.add(z3.Not(z3.fpIsInf(x))) + +if s.check() == z3.sat: + m = s.model() + print("Counterexample found!") + print(m) + print("a*c: ", z3.simplify(m[a] * m[c])) + print("b*c: ", z3.simplify(m[b] * m[c])) + print("a+b: ", z3.simplify(m[a] + m[b])) + print("a*c + b*c: ", z3.simplify(m[a] * m[c] + m[b] * m[c])) + print("(a+b) * c: ", z3.simplify((m[a] + m[b]) * m[c])) +else: + print("Proved!") diff --git a/tensorflow/compiler/xla/service/algebraic_simplifier_test.cc b/tensorflow/compiler/xla/service/algebraic_simplifier_test.cc index c79c518700b63be2fda8a415b38fe246689ab7c6..3602ab82b248bb3d7cd8203ed7664e3c460374d2 100644 --- a/tensorflow/compiler/xla/service/algebraic_simplifier_test.cc +++ b/tensorflow/compiler/xla/service/algebraic_simplifier_test.cc @@ -27,13 +27,14 @@ 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_instructions.h" -#include "tensorflow/compiler/xla/service/hlo_matchers.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" +#include "tensorflow/compiler/xla/service/hlo_parser.h" #include "tensorflow/compiler/xla/service/hlo_pass_fix.h" +#include "tensorflow/compiler/xla/service/pattern_matcher.h" +#include "tensorflow/compiler/xla/service/pattern_matcher_gmock.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/test.h" #include "tensorflow/compiler/xla/tests/hlo_test_base.h" -#include "tensorflow/compiler/xla/tests/hlo_verified_test_base.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/window_util.h" #include "tensorflow/compiler/xla/xla_data.pb.h" @@ -43,21 +44,16 @@ namespace xla { namespace { using ::testing::ElementsAre; +namespace m = match; -namespace op = xla::testing::opcode_matchers; - -AlgebraicSimplifier::ValidBitcastCallback bitcasting_callback() { - return [](const Shape&, const Shape&) { return true; }; -} - -AlgebraicSimplifier::ValidBitcastCallback non_bitcasting_callback() { - return [](const Shape&, const Shape&) { return false; }; -} - -class AlgebraicSimplifierTest : public HloVerifiedTestBase {}; +class AlgebraicSimplifierTest : public HloTestBase { + protected: + AlgebraicSimplifierOptions default_options_; +}; // Test that A + 0 is simplified to A TEST_F(AlgebraicSimplifierTest, AddZero) { + auto m = CreateNewVerifiedModule(); Shape r0f32 = ShapeUtil::MakeShape(F32, {}); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction( @@ -67,18 +63,220 @@ TEST_F(AlgebraicSimplifierTest, AddZero) { builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kAdd, param0, zero)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root->opcode(), HloOpcode::kAdd); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); root = computation->root_instruction(); EXPECT_EQ(root, param0); } +TEST_F(AlgebraicSimplifierTest, FactorIntegerAddition) { + const char* kModuleStr = R"( + HloModule m + test { + p0 = s32[8] parameter(0) + p1 = s32[8] parameter(1) + p2 = s32[8] parameter(2) + x = s32[8] multiply(p0, p2) + y = s32[8] multiply(p1, p2) + ROOT sum = s32[8] add(x, y) + } + )"; + TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr)); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); + EXPECT_THAT( + m->entry_computation()->root_instruction(), + GmockMatch(m::MultiplyAnyOrder( + m::AddAnyOrder(m::Parameter(0), m::Parameter(1)), m::Parameter(2)))); +} + +// A*C + B*C => (A+B)*C if C is a floating-point power of 2. +TEST_F(AlgebraicSimplifierTest, FactorFpAddition) { + const char* kModuleStr = R"( + HloModule m + test { + p0 = f32[] parameter(0) + p1 = f32[] parameter(1) + c = f32[] constant(0.125) + x = f32[] multiply(p0, c) + y = f32[] multiply(p1, c) + ROOT sum = f32[] add(x, y) + } + )"; + TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr)); + ASSERT_TRUE(AlgebraicSimplifier(default_options_).Run(m.get()).ValueOrDie()); + EXPECT_THAT(m->entry_computation()->root_instruction(), + GmockMatch(m::MultiplyAnyOrder( + m::AddAnyOrder(m::Parameter(0), m::Parameter(1)), + m::ConstantScalar(0.125)))); +} + +// A*C + B*C => (A+B)*C if C is a broadcast of a floating-point power of 2. +TEST_F(AlgebraicSimplifierTest, FactorFpAdditionWithBroadcast) { + const char* kModuleStr = R"( + HloModule m + test { + p0 = f32[4] parameter(0) + p1 = f32[4] parameter(1) + c = f32[] constant(0.125) + b = f32[4] broadcast(c), dimensions={} + x = f32[4] multiply(p0, b) + y = f32[4] multiply(p1, b) + ROOT sum = f32[4] add(x, y) + } + )"; + TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr)); + ASSERT_TRUE(AlgebraicSimplifier(default_options_).Run(m.get()).ValueOrDie()); + EXPECT_THAT(m->entry_computation()->root_instruction(), + GmockMatch(m::MultiplyAnyOrder( + m::AddAnyOrder(m::Parameter(0), m::Parameter(1)), + m::Broadcast(m::ConstantScalar(0.125))))); +} + +// A*C + B*C => (A+B)*C simplification should not happen if C is not a +// floating-point power of 2. +TEST_F(AlgebraicSimplifierTest, FactorFpAdditionNotPowerOf2) { + const char* kModuleStr = R"( + HloModule m + test { + p0 = f32[] parameter(0) + p1 = f32[] parameter(1) + c = f32[] constant(0.3) + x = f32[] multiply(p0, c) + y = f32[] multiply(p1, c) + ROOT sum = f32[] add(x, y) + } + )"; + TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr)); + EXPECT_FALSE(AlgebraicSimplifier(default_options_).Run(m.get()).ValueOrDie()); +} + +// A*C + B*C => (A+B)*C simplification should not happen if A, B, and C are +// complex numbers. +TEST_F(AlgebraicSimplifierTest, FactorFpAdditionComplex) { + const char* kModuleStr = R"( + HloModule m + test { + p0 = c64[8] parameter(0) + p1 = c64[8] parameter(1) + p2 = c64[8] parameter(2) + x = c64[8] multiply(p0, p2) + y = c64[8] multiply(p1, p2) + ROOT sum = c64[8] add(x, y) + } + )"; + TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr)); + EXPECT_FALSE(AlgebraicSimplifier(default_options_).Run(m.get()).ValueOrDie()); +} + +// A*C + B*C => (A+B)*C simplification is OK if A, B, and C are complex. +TEST_F(AlgebraicSimplifierTest, FactorFpAdditionBfloat16) { + const char* kModuleStr = R"( + HloModule m + test { + p0 = bf16[4] parameter(0) + p1 = bf16[4] parameter(1) + c = bf16[] constant(0.125) + b = bf16[4] broadcast(c), dimensions={} + x = bf16[4] multiply(p0, b) + y = bf16[4] multiply(p1, b) + ROOT sum = bf16[4] add(x, y) + } + )"; + TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr)); + ASSERT_TRUE(AlgebraicSimplifier(default_options_).Run(m.get()).ValueOrDie()); + EXPECT_THAT(m->entry_computation()->root_instruction(), + GmockMatch(m::MultiplyAnyOrder( + m::AddAnyOrder(m::Parameter(0), m::Parameter(1)), + 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(); Shape r0s32 = ShapeUtil::MakeShape(S32, {}); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction( @@ -88,12 +286,11 @@ TEST_F(AlgebraicSimplifierTest, MulZero) { builder.AddInstruction( HloInstruction::CreateBinary(r0s32, HloOpcode::kMultiply, param0, zero)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root->opcode(), HloOpcode::kMultiply); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); EXPECT_EQ(computation->root_instruction(), zero); } @@ -114,8 +311,7 @@ TEST_F(AlgebraicSimplifierTest, SelectTrue) { auto computation = module->AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root->opcode(), HloOpcode::kSelect); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); + AlgebraicSimplifier simplifier(default_options_); ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); EXPECT_EQ(computation->root_instruction(), param0); } @@ -137,8 +333,7 @@ TEST_F(AlgebraicSimplifierTest, SelectFalse) { auto computation = module->AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root->opcode(), HloOpcode::kSelect); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); + AlgebraicSimplifier simplifier(default_options_); ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); EXPECT_EQ(computation->root_instruction(), param1); } @@ -158,14 +353,14 @@ TEST_F(AlgebraicSimplifierTest, SelectIdentical) { auto computation = module->AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root->opcode(), HloOpcode::kSelect); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); + AlgebraicSimplifier simplifier(default_options_); ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); EXPECT_EQ(computation->root_instruction(), param1); } // Test that Reduce(Reduce(A)) -> Reduce(A) TEST_F(AlgebraicSimplifierTest, TwoReducesToOne) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); // Create add computation. HloInstruction* zero = builder.AddInstruction( @@ -180,7 +375,7 @@ TEST_F(AlgebraicSimplifierTest, TwoReducesToOne) { HloInstruction::CreateParameter(1, scalar_shape, "p1")); builder.AddInstruction( HloInstruction::CreateBinary(scalar_shape, HloOpcode::kAdd, p0, p1)); - add_computation = module().AddEmbeddedComputation(builder.Build()); + add_computation = m->AddEmbeddedComputation(builder.Build()); } Shape r4f32 = ShapeUtil::MakeShape(F32, {4, 5, 6, 7}); HloInstruction* param = builder.AddInstruction( @@ -193,17 +388,17 @@ TEST_F(AlgebraicSimplifierTest, TwoReducesToOne) { Shape r1f32 = ShapeUtil::MakeShape(F32, {5}); builder.AddInstruction(HloInstruction::CreateReduce(r1f32, reduce0, zero, dims1, add_computation)); - module().AddEntryComputation(builder.Build()); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); - HloInstruction* root = module().entry_computation()->root_instruction(); - EXPECT_THAT(root, op::Reduce(param, zero)); + m->AddEntryComputation(builder.Build()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); + HloInstruction* root = m->entry_computation()->root_instruction(); + EXPECT_THAT(root, GmockMatch(m::Reduce(m::Parameter(0), m::Op().Is(zero)))); EXPECT_EQ(root->dimensions(), std::vector({0, 2, 3})); } // Test that Const + A is canonicalized to A + Const. TEST_F(AlgebraicSimplifierTest, AddConstOnLHS) { + auto m = CreateNewVerifiedModule(); Shape r0f32 = ShapeUtil::MakeShape(F32, {}); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction( @@ -213,18 +408,18 @@ TEST_F(AlgebraicSimplifierTest, AddConstOnLHS) { builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kAdd, constant, param0)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root->opcode(), HloOpcode::kAdd); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); root = computation->root_instruction(); - EXPECT_THAT(root, op::Add(param0, op::Constant())); + EXPECT_THAT(root, GmockMatch(m::Add(m::Parameter(0), m::Constant()))); } // Test that [(A + C1) + C2] => [A + (C1 + C2)] for constants C1 and C2. TEST_F(AlgebraicSimplifierTest, AddReassociateMergeConstants) { + auto m = CreateNewVerifiedModule(); Shape r0f32 = ShapeUtil::MakeShape(F32, {}); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction( @@ -239,17 +434,19 @@ TEST_F(AlgebraicSimplifierTest, AddReassociateMergeConstants) { builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kAdd, add1, constant2)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root->opcode(), HloOpcode::kAdd); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); root = computation->root_instruction(); - EXPECT_THAT(root, op::Add(param0, op::Add(constant1, constant2))); + EXPECT_THAT(root, GmockMatch(m::Add( + m::Op().Is(param0), + m::Add(m::Op().Is(constant1), m::Op().Is(constant2))))); } TEST_F(AlgebraicSimplifierTest, AddBroadcastZeroR0Operand) { + auto m = CreateNewVerifiedModule(); Shape r2f32 = ShapeUtil::MakeShape(F32, {3, 2}); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction( @@ -261,17 +458,17 @@ TEST_F(AlgebraicSimplifierTest, AddBroadcastZeroR0Operand) { builder.AddInstruction( HloInstruction::CreateBinary(r2f32, HloOpcode::kAdd, bcast, param0)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root->opcode(), HloOpcode::kAdd); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); root = computation->root_instruction(); EXPECT_EQ(root, param0); } TEST_F(AlgebraicSimplifierTest, InlineTrivialMap) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); // Create add computation. HloComputation* add_computation = nullptr; @@ -284,7 +481,7 @@ TEST_F(AlgebraicSimplifierTest, InlineTrivialMap) { HloInstruction::CreateParameter(1, scalar_shape, "p1")); builder.AddInstruction( HloInstruction::CreateBinary(scalar_shape, HloOpcode::kAdd, p0, p1)); - add_computation = module().AddEmbeddedComputation(builder.Build()); + add_computation = m->AddEmbeddedComputation(builder.Build()); } Shape r2f32 = ShapeUtil::MakeShape(F32, {32, 1}); HloInstruction* param0 = builder.AddInstruction( @@ -297,17 +494,18 @@ TEST_F(AlgebraicSimplifierTest, InlineTrivialMap) { HloInstruction::CreateBroadcast(r2f32, zero, {}))}, add_computation)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root->opcode(), HloOpcode::kMap); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); root = computation->root_instruction(); - EXPECT_THAT(root, op::Add(param0, op::Broadcast(zero))); + EXPECT_THAT(root, GmockMatch(m::Add(m::Parameter(0), + m::Broadcast(m::Op().Is(zero))))); } TEST_F(AlgebraicSimplifierTest, AddBroadcastZeroR1Operand) { + auto m = CreateNewVerifiedModule(); Shape r2f32 = ShapeUtil::MakeShape(F32, {3, 2}); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction( @@ -319,64 +517,64 @@ TEST_F(AlgebraicSimplifierTest, AddBroadcastZeroR1Operand) { builder.AddInstruction( HloInstruction::CreateBinary(r2f32, HloOpcode::kAdd, bcast, param0)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root->opcode(), HloOpcode::kAdd); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); root = computation->root_instruction(); EXPECT_EQ(root, param0); } TEST_F(AlgebraicSimplifierTest, ConstantToBroadcast) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); builder.AddInstruction(HloInstruction::CreateConstant( LiteralUtil::CreateR1({3.14f, 3.14f, 3.14f}))); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); - EXPECT_THAT(root, op::Constant()); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + EXPECT_THAT(root, GmockMatch(m::Constant())); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); root = computation->root_instruction(); - EXPECT_THAT(root, op::Broadcast(op::Constant())); + EXPECT_THAT(root, GmockMatch(m::Broadcast(m::Constant()))); EXPECT_EQ(3.14f, root->operand(0)->literal().GetFirstElement()); } TEST_F(AlgebraicSimplifierTest, ConstantNotToBroadcast) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); builder.AddInstruction(HloInstruction::CreateConstant( LiteralUtil::CreateR1({3.14, 3.14, 4}))); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); - EXPECT_THAT(root, op::Constant()); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_FALSE(simplifier.Run(&module()).ValueOrDie()); + EXPECT_THAT(root, GmockMatch(m::Constant())); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_FALSE(simplifier.Run(m.get()).ValueOrDie()); root = computation->root_instruction(); - EXPECT_THAT(root, op::Constant()); + EXPECT_THAT(root, GmockMatch(m::Constant())); } TEST_F(AlgebraicSimplifierTest, IotaToBroadcast) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); builder.AddInstruction(HloInstruction::CreateConstant( LiteralUtil::CreateR1({0.0f, 1.0f, 2.0f}))); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); - EXPECT_THAT(root, op::Constant()); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + EXPECT_THAT(root, GmockMatch(m::Constant())); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); root = computation->root_instruction(); - EXPECT_THAT(root, op::Iota()); + EXPECT_THAT(root, GmockMatch(m::Iota())); } // Test that A - 0 is simplified to A TEST_F(AlgebraicSimplifierTest, SubZero) { + auto m = CreateNewVerifiedModule(); Shape r0f32 = ShapeUtil::MakeShape(F32, {}); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction( @@ -386,18 +584,18 @@ TEST_F(AlgebraicSimplifierTest, SubZero) { builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kSubtract, param0, zero)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root->opcode(), HloOpcode::kSubtract); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); root = computation->root_instruction(); EXPECT_EQ(root, param0); } // Test that A - Const is canonicalized to A + (-Const). TEST_F(AlgebraicSimplifierTest, SubConstCanonicalization) { + auto m = CreateNewVerifiedModule(); Shape r0f32 = ShapeUtil::MakeShape(F32, {}); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction( @@ -407,18 +605,19 @@ TEST_F(AlgebraicSimplifierTest, SubConstCanonicalization) { builder.AddInstruction(HloInstruction::CreateBinary( r0f32, HloOpcode::kSubtract, param0, constant)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root->opcode(), HloOpcode::kSubtract); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); root = computation->root_instruction(); - EXPECT_THAT(root, op::Add(param0, op::Negate(constant))); + EXPECT_THAT(root, GmockMatch(m::Add(m::Parameter(0), + m::Negate(m::Op().Is(constant))))); } // Test that (A/B)/C is simplified to A/(B*C). TEST_F(AlgebraicSimplifierTest, LhsDivOfDiv) { + auto m = CreateNewVerifiedModule(); Shape r0f32 = ShapeUtil::MakeShape(F32, {}); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction( @@ -432,21 +631,24 @@ TEST_F(AlgebraicSimplifierTest, LhsDivOfDiv) { builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kDivide, div, param2)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), - op::Divide(op::Divide(param0, param1), param2)); + GmockMatch(m::Divide(m::Divide(m::Parameter(0), m::Parameter(1)), + m::Parameter(2)))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); - EXPECT_THAT(computation->root_instruction(), - op::Divide(param0, op::Multiply(param1, param2))); + EXPECT_THAT( + computation->root_instruction(), + GmockMatch(m::Divide(m::Parameter(0), + m::Multiply(m::Parameter(1), m::Parameter(2))))); } // Test that A/(B/C) is simplified to (A*C)/B. TEST_F(AlgebraicSimplifierTest, RhsDivOfDiv) { + auto m = CreateNewVerifiedModule(); Shape r0f32 = ShapeUtil::MakeShape(F32, {}); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction( @@ -460,21 +662,25 @@ TEST_F(AlgebraicSimplifierTest, RhsDivOfDiv) { builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kDivide, param0, div)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); - EXPECT_THAT(computation->root_instruction(), - op::Divide(param0, op::Divide(param1, param2))); + EXPECT_THAT( + computation->root_instruction(), + GmockMatch(m::Divide(m::Parameter(0), + m::Divide(m::Parameter(1), m::Parameter(2))))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); - EXPECT_THAT(computation->root_instruction(), - op::Divide(op::Multiply(param0, param2), param1)); + EXPECT_THAT( + computation->root_instruction(), + GmockMatch(m::Divide(m::Multiply(m::Parameter(0), m::Parameter(2)), + m::Parameter(1)))); } // Test that (A/B)/(C/D) is simplified to (A*D)/(B*C). TEST_F(AlgebraicSimplifierTest, DivOfDivAndDiv) { + auto m = CreateNewVerifiedModule(); Shape r2f32 = ShapeUtil::MakeShape(F32, {42, 123}); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction( @@ -492,23 +698,25 @@ TEST_F(AlgebraicSimplifierTest, DivOfDivAndDiv) { builder.AddInstruction( HloInstruction::CreateBinary(r2f32, HloOpcode::kDivide, div0, div1)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); EXPECT_THAT( computation->root_instruction(), - op::Divide(op::Divide(param0, param1), op::Divide(param2, param3))); + GmockMatch(m::Divide(m::Divide(m::Parameter(0), m::Parameter(1)), + m::Divide(m::Parameter(2), m::Parameter(3))))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); EXPECT_THAT( computation->root_instruction(), - op::Divide(op::Multiply(param0, param3), op::Multiply(param1, param2))); + GmockMatch(m::Divide(m::Multiply(m::Parameter(0), m::Parameter(3)), + m::Multiply(m::Parameter(1), m::Parameter(2))))); } // Test that A/exp(B) is simplified to A*exp(-B). TEST_F(AlgebraicSimplifierTest, DivOfExp) { + auto m = CreateNewVerifiedModule(); Shape r0f32 = ShapeUtil::MakeShape(F32, {}); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction( @@ -520,21 +728,22 @@ TEST_F(AlgebraicSimplifierTest, DivOfExp) { builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kDivide, param0, exp)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), - op::Divide(param0, op::Exp(param1))); + GmockMatch(m::Divide(m::Parameter(0), m::Exp(m::Parameter(1))))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), - op::Multiply(param0, op::Exp(op::Negate(param1)))); + GmockMatch(m::Multiply(m::Parameter(0), + m::Exp(m::Negate(m::Parameter(1)))))); } // Test that A/pow(B,C) is simplified to A*pow(B,-C). TEST_F(AlgebraicSimplifierTest, DivOfPower) { + auto m = CreateNewVerifiedModule(); Shape r0f32 = ShapeUtil::MakeShape(F32, {}); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction( @@ -548,22 +757,26 @@ TEST_F(AlgebraicSimplifierTest, DivOfPower) { builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kDivide, param0, power)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); - EXPECT_THAT(computation->root_instruction(), - op::Divide(param0, op::Power(param1, param2))); + EXPECT_THAT( + computation->root_instruction(), + GmockMatch(m::Divide(m::Parameter(0), + m::Power(m::Parameter(1), m::Parameter(2))))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), - op::Multiply(param0, op::Power(param1, op::Negate(param2)))); + GmockMatch(m::Multiply( + m::Parameter(0), + m::Power(m::Parameter(1), m::Negate(m::Parameter(2)))))); } // Test that broadcasting is done on the right step when simplifying A/pow(B,C) // to A*pow(B,-C). TEST_F(AlgebraicSimplifierTest, DivOfBroadcastingPower) { + auto m = CreateNewVerifiedModule(); Shape r1f32 = ShapeUtil::MakeShape(F32, {7}); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction( @@ -577,21 +790,25 @@ TEST_F(AlgebraicSimplifierTest, DivOfBroadcastingPower) { builder.AddInstruction( HloInstruction::CreateBinary(r1f32, HloOpcode::kDivide, param0, power)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); - EXPECT_THAT(computation->root_instruction(), - op::Divide(param0, op::Power(param1, param2))); + EXPECT_THAT( + computation->root_instruction(), + GmockMatch(m::Divide(m::Parameter(0), + m::Power(m::Parameter(1), m::Parameter(2))))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); ASSERT_THAT(computation->root_instruction(), - op::Multiply(param0, op::Power(param1, op::Negate(param2)))); + GmockMatch(m::Multiply( + m::Parameter(0), + m::Power(m::Parameter(1), m::Negate(m::Parameter(2)))))); } // A / Const => A * InvertedConst TEST_F(AlgebraicSimplifierTest, DivideByConstant) { + auto m = CreateNewVerifiedModule(); Shape r1f32 = ShapeUtil::MakeShape(F32, {3}); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction( @@ -602,18 +819,18 @@ TEST_F(AlgebraicSimplifierTest, DivideByConstant) { builder.AddInstruction(HloInstruction::CreateBinary(r1f32, HloOpcode::kDivide, param0, constant)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), - op::Multiply(param0, op::Constant())); + GmockMatch(m::Multiply(m::Parameter(0), m::Constant()))); } // pow(pow(A, X), Y) => pow(A, X*Y) TEST_F(AlgebraicSimplifierTest, PowerOfPower) { + auto m = CreateNewVerifiedModule(); Shape r1f32 = ShapeUtil::MakeShape(F32, {7}); HloComputation::Builder builder(TestName()); HloInstruction* base = builder.AddInstruction( @@ -627,17 +844,19 @@ TEST_F(AlgebraicSimplifierTest, PowerOfPower) { builder.AddInstruction(HloInstruction::CreateBinary(r1f32, HloOpcode::kPower, inner_power, exp2)); - auto computation = module().AddEntryComputation(builder.Build()); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); - EXPECT_THAT(computation->root_instruction(), - op::Power(base, op::Multiply(exp1, exp2))); + auto computation = m->AddEntryComputation(builder.Build()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); + EXPECT_THAT( + computation->root_instruction(), + GmockMatch(m::Power(m::Op().Is(base), + m::Multiply(m::Op().Is(exp1), m::Op().Is(exp2))))); } // Don't simplify pow(pow(A, X), Y) => pow(A, X*Y) if X and Y are complex // numbers. TEST_F(AlgebraicSimplifierTest, PowerOfPowerComplex) { + auto m = CreateNewVerifiedModule(); Shape r1c64 = ShapeUtil::MakeShape(C64, {7}); HloComputation::Builder builder(TestName()); HloInstruction* base = builder.AddInstruction( @@ -651,14 +870,14 @@ TEST_F(AlgebraicSimplifierTest, PowerOfPowerComplex) { builder.AddInstruction(HloInstruction::CreateBinary(r1c64, HloOpcode::kPower, inner_power, exp2)); - module().AddEntryComputation(builder.Build()); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_FALSE(simplifier.Run(&module()).ValueOrDie()); + m->AddEntryComputation(builder.Build()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_FALSE(simplifier.Run(m.get()).ValueOrDie()); } // Test that A/1 is simplified to A for a scalar. TEST_F(AlgebraicSimplifierTest, DivOneScalar) { + auto m = CreateNewVerifiedModule(); Shape r0f32 = ShapeUtil::MakeShape(F32, {}); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction( @@ -668,18 +887,18 @@ TEST_F(AlgebraicSimplifierTest, DivOneScalar) { HloInstruction* div = builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kDivide, param0, one)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root, div); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); root = computation->root_instruction(); EXPECT_EQ(root, param0); } // Test that A/1 is simplified to A for an array. TEST_F(AlgebraicSimplifierTest, DivOneArray) { + auto m = CreateNewVerifiedModule(); Shape r2f32 = ShapeUtil::MakeShape(F32, {2, 2}); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction( @@ -689,18 +908,18 @@ TEST_F(AlgebraicSimplifierTest, DivOneArray) { HloInstruction* div = builder.AddInstruction( HloInstruction::CreateBinary(r2f32, HloOpcode::kDivide, param0, one)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root, div); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); root = computation->root_instruction(); EXPECT_EQ(root, param0); } // Test that complex(real(c), imag(c)) is simplified to c. TEST_F(AlgebraicSimplifierTest, ComplexOfRealImagC) { + auto m = CreateNewVerifiedModule(); Shape r2f32 = ShapeUtil::MakeShape(F32, {2, 2}); Shape r2c64 = ShapeUtil::MakeShape(C64, {2, 2}); HloComputation::Builder builder(TestName()); @@ -713,18 +932,18 @@ TEST_F(AlgebraicSimplifierTest, ComplexOfRealImagC) { HloInstruction* cplx = builder.AddInstruction( HloInstruction::CreateBinary(r2c64, HloOpcode::kComplex, real, imag)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root, cplx); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); root = computation->root_instruction(); EXPECT_EQ(root, param0); } // Test that real(complex(r,i)) is simplified to r. TEST_F(AlgebraicSimplifierTest, RealOfComplex) { + auto m = CreateNewVerifiedModule(); Shape r2f32 = ShapeUtil::MakeShape(F32, {2, 2}); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction( @@ -737,18 +956,18 @@ TEST_F(AlgebraicSimplifierTest, RealOfComplex) { HloInstruction* real = builder.AddInstruction( HloInstruction::CreateUnary(r2f32, HloOpcode::kReal, cplx)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root, real); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); root = computation->root_instruction(); EXPECT_EQ(root, param0); } // Test that imag(complex(r,i)) is simplified to i. TEST_F(AlgebraicSimplifierTest, ImagOfComplex) { + auto m = CreateNewVerifiedModule(); Shape r2f32 = ShapeUtil::MakeShape(F32, {2, 2}); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction( @@ -761,18 +980,18 @@ TEST_F(AlgebraicSimplifierTest, ImagOfComplex) { HloInstruction* imag = builder.AddInstruction( HloInstruction::CreateUnary(r2f32, HloOpcode::kImag, cplx)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root, imag); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); root = computation->root_instruction(); EXPECT_EQ(root, param1); } // Test that get_element(make_tuple({A,B}),1) is simplified to B TEST_F(AlgebraicSimplifierTest, SelectMakeTuple) { + auto m = CreateNewVerifiedModule(); Shape r0f32 = ShapeUtil::MakeShape(F32, {}); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction( @@ -788,18 +1007,18 @@ TEST_F(AlgebraicSimplifierTest, SelectMakeTuple) { HloInstruction* add = builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kAdd, get, param2)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root, add); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); root = computation->root_instruction(); - EXPECT_THAT(root, op::Add(param1, param2)); + EXPECT_THAT(root, GmockMatch(m::Add(m::Parameter(1), m::Parameter(2)))); } // Test that exp(A)/exp(B) is simplified to exp(A-B) TEST_F(AlgebraicSimplifierTest, ExpDiv) { + auto m = CreateNewVerifiedModule(); Shape r0f32 = ShapeUtil::MakeShape(F32, {}); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction( @@ -813,21 +1032,23 @@ TEST_F(AlgebraicSimplifierTest, ExpDiv) { builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kDivide, exp0, exp1)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); - EXPECT_THAT(computation->root_instruction(), - op::Divide(op::Exp(param0), op::Exp(param1))); + EXPECT_THAT( + computation->root_instruction(), + GmockMatch(m::Divide(m::Exp(m::Parameter(0)), m::Exp(m::Parameter(1))))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); - EXPECT_THAT(computation->root_instruction(), - op::Exp(op::Subtract(param0, param1))); + EXPECT_THAT( + computation->root_instruction(), + GmockMatch(m::Exp(m::Subtract(m::Parameter(0), m::Parameter(1))))); } // Test that exp(A)*exp(B) is simplified to exp(A+B) TEST_F(AlgebraicSimplifierTest, ExpMul) { + auto m = CreateNewVerifiedModule(); Shape r0f32 = ShapeUtil::MakeShape(F32, {}); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction( @@ -841,21 +1062,22 @@ TEST_F(AlgebraicSimplifierTest, ExpMul) { builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kMultiply, exp0, exp1)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), - op::Multiply(op::Exp(param0), op::Exp(param1))); + GmockMatch(m::Multiply(m::Exp(m::Parameter(0)), + m::Exp(m::Parameter(1))))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), - op::Exp(op::Add(param0, param1))); + GmockMatch(m::Exp(m::Add(m::Parameter(0), m::Parameter(1))))); } // Test that pow(exp(A), B) is simplified to exp(A*B) TEST_F(AlgebraicSimplifierTest, PowExp) { + auto m = CreateNewVerifiedModule(); Shape r0f32 = ShapeUtil::MakeShape(F32, {}); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction( @@ -867,21 +1089,22 @@ TEST_F(AlgebraicSimplifierTest, PowExp) { builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kPower, exp0, param1)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), - op::Power(op::Exp(param0), param1)); + GmockMatch(m::Power(m::Exp(m::Parameter(0)), m::Parameter(1)))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); - EXPECT_THAT(computation->root_instruction(), - op::Exp(op::Multiply(param0, param1))); + EXPECT_THAT( + computation->root_instruction(), + GmockMatch(m::Exp(m::Multiply(m::Parameter(0), m::Parameter(1))))); } // Test that ln(pow(A, B)) is simplified to ln(A)*B TEST_F(AlgebraicSimplifierTest, LnPow) { + auto m = CreateNewVerifiedModule(); Shape r0f32 = ShapeUtil::MakeShape(F32, {}); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction( @@ -893,21 +1116,22 @@ TEST_F(AlgebraicSimplifierTest, LnPow) { builder.AddInstruction( HloInstruction::CreateUnary(r0f32, HloOpcode::kLog, pow)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), - op::Log(op::Power(param0, param1))); + GmockMatch(m::Log(m::Power(m::Parameter(0), m::Parameter(1))))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); - EXPECT_THAT(computation->root_instruction(), - op::Multiply(op::Log(param0), param1)); + EXPECT_THAT( + computation->root_instruction(), + GmockMatch(m::Multiply(m::Log(m::Parameter(0)), m::Parameter(1)))); } // Test that ln(exp(A)) is simplified to A TEST_F(AlgebraicSimplifierTest, LnExp) { + auto m = CreateNewVerifiedModule(); Shape r0f32 = ShapeUtil::MakeShape(F32, {}); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction( @@ -917,19 +1141,20 @@ TEST_F(AlgebraicSimplifierTest, LnExp) { builder.AddInstruction( HloInstruction::CreateUnary(r0f32, HloOpcode::kLog, exp0)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); - EXPECT_THAT(computation->root_instruction(), op::Log(op::Exp(param0))); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Log(m::Exp(m::Parameter(0))))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); EXPECT_EQ(computation->root_instruction(), param0); } // Test that ln(exp(A)/exp(B)) is simplified to A-B TEST_F(AlgebraicSimplifierTest, LnExpDiv) { + auto m = CreateNewVerifiedModule(); Shape r0f32 = ShapeUtil::MakeShape(F32, {}); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction( @@ -945,21 +1170,23 @@ TEST_F(AlgebraicSimplifierTest, LnExpDiv) { builder.AddInstruction( HloInstruction::CreateUnary(r0f32, HloOpcode::kLog, div)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), - op::Log(op::Divide(op::Exp(param0), op::Exp(param1)))); + GmockMatch(m::Log(m::Divide(m::Exp(m::Parameter(0)), + m::Exp(m::Parameter(1)))))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); - EXPECT_THAT(computation->root_instruction(), op::Subtract(param0, param1)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Subtract(m::Parameter(0), m::Parameter(1)))); } // Test that pow(A, 0) where A is a scalar is simplified to the scalar // constant 1. TEST_F(AlgebraicSimplifierTest, Pow0Scalar) { + auto m = CreateNewVerifiedModule(); Shape r0f32 = ShapeUtil::MakeShape(F32, {}); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction( @@ -969,21 +1196,22 @@ TEST_F(AlgebraicSimplifierTest, Pow0Scalar) { builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kPower, param0, zero)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); - EXPECT_THAT(computation->root_instruction(), op::Power(param0, zero)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Power(m::Parameter(0), m::Op().Is(zero)))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); HloInstruction* root = computation->root_instruction(); - EXPECT_THAT(root, op::Constant()); + EXPECT_THAT(root, GmockMatch(m::Constant())); EXPECT_EQ(root->literal().GetFirstElement(), 1); } // Test that pow(A, 0) where A is not a scalar is simplified to broadcast(1). TEST_F(AlgebraicSimplifierTest, Pow0Vector) { + auto m = CreateNewVerifiedModule(); Shape r1f32 = ShapeUtil::MakeShape(F32, {42}); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction( @@ -993,16 +1221,16 @@ TEST_F(AlgebraicSimplifierTest, Pow0Vector) { builder.AddInstruction( HloInstruction::CreateBinary(r1f32, HloOpcode::kPower, param0, zero)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); - EXPECT_THAT(computation->root_instruction(), op::Power(param0, zero)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Power(m::Parameter(0), m::Op().Is(zero)))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); HloInstruction* root = computation->root_instruction(); - EXPECT_THAT(root, op::Broadcast()); + EXPECT_THAT(root, GmockMatch(m::Broadcast())); EXPECT_TRUE(ShapeUtil::Equal(root->shape(), r1f32)) << ShapeUtil::HumanString(root->shape()); EXPECT_EQ(root->dimensions().size(), 0); @@ -1012,6 +1240,7 @@ TEST_F(AlgebraicSimplifierTest, Pow0Vector) { // Test that pow(A, 1) is simplified to A. TEST_F(AlgebraicSimplifierTest, Pow1) { + auto m = CreateNewVerifiedModule(); Shape r0f32 = ShapeUtil::MakeShape(F32, {}); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction( @@ -1021,19 +1250,20 @@ TEST_F(AlgebraicSimplifierTest, Pow1) { builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kPower, param0, one)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); - EXPECT_THAT(computation->root_instruction(), op::Power(param0, one)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Power(m::Parameter(0), m::Op().Is(one)))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); EXPECT_EQ(computation->root_instruction(), param0); } // Test that pow(A, 2) is simplified to A*A. TEST_F(AlgebraicSimplifierTest, Pow2) { + auto m = CreateNewVerifiedModule(); Shape r0f32 = ShapeUtil::MakeShape(F32, {}); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction( @@ -1043,19 +1273,21 @@ TEST_F(AlgebraicSimplifierTest, Pow2) { builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kPower, param0, two)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); - EXPECT_THAT(computation->root_instruction(), op::Power(param0, two)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Power(m::Parameter(0), m::Op().Is(two)))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); - EXPECT_THAT(computation->root_instruction(), op::Multiply(param0, param0)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Multiply(m::Parameter(0), m::Parameter(0)))); } // Test that pow(A, -1) is simplified to 1/A. TEST_F(AlgebraicSimplifierTest, PowNegative1) { + auto m = CreateNewVerifiedModule(); Shape r0f32 = ShapeUtil::MakeShape(F32, {}); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction( @@ -1065,22 +1297,23 @@ TEST_F(AlgebraicSimplifierTest, PowNegative1) { builder.AddInstruction(HloInstruction::CreateBinary(r0f32, HloOpcode::kPower, param0, negative_one)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); - EXPECT_THAT(computation->root_instruction(), op::Power(param0, negative_one)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Power(m::Parameter(0), m::Op().Is(negative_one)))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); HloInstruction* root = computation->root_instruction(); - EXPECT_THAT(root, op::Divide(op::Broadcast(), param0)); + EXPECT_THAT(root, GmockMatch(m::Divide(m::Broadcast(), m::Parameter(0)))); EXPECT_EQ(root->operand(0)->opcode(), HloOpcode::kBroadcast); EXPECT_EQ(root->operand(0)->operand(0)->literal().GetFirstElement(), 1); } TEST_F(AlgebraicSimplifierTest, ZeroSizedConvolution) { + auto m = CreateNewVerifiedModule(); auto builder = HloComputation::Builder(TestName()); HloInstruction* lhs = builder.AddInstruction(HloInstruction::CreateParameter( 0, ShapeUtil::MakeShape(F32, {3, 3, 0}), "lhs")); @@ -1112,18 +1345,63 @@ TEST_F(AlgebraicSimplifierTest, ZeroSizedConvolution) { // Create add computation. builder.AddInstruction(HloInstruction::CreateConvolve( ShapeUtil::MakeShape(F32, {3, 3, 3}), lhs, rhs, /*feature_group_count=*/1, - window, dnums, DefaultPrecisionConfig(2))); - module().AddEntryComputation(builder.Build()); - HloPassFix simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - EXPECT_THAT(module().entry_computation()->root_instruction(), - op::Convolution(lhs, rhs)); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); - EXPECT_THAT(module().entry_computation()->root_instruction(), - op::Broadcast(op::Constant())); + /*batch_group_count=*/1, window, dnums, DefaultPrecisionConfig(2))); + m->AddEntryComputation(builder.Build()); + HloPassFix simplifier(default_options_); + EXPECT_THAT(m->entry_computation()->root_instruction(), + GmockMatch(m::Convolution(m::Op().Is(lhs), m::Op().Is(rhs)))); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); + EXPECT_THAT(m->entry_computation()->root_instruction(), + GmockMatch(m::Broadcast(m::Constant()))); +} + +TEST_F(AlgebraicSimplifierTest, ReduceWindowIsReduceAndReshape) { + auto m = CreateNewVerifiedModule(); + auto builder = HloComputation::Builder(TestName()); + HloInstruction* param = + builder.AddInstruction(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShape(F32, {1, 2, 3, 4}), "param")); + Window window; + for (int64 i = 0; i < 4; ++i) { + WindowDimension* dim = window.add_dimensions(); + // Makes 1x2x3x1 window. + dim->set_size((i % 3) + 1); + dim->set_stride(1); + dim->set_padding_low(0); + dim->set_padding_high(0); + dim->set_window_dilation(1); + dim->set_base_dilation(1); + } + // Create add computation. + HloComputation* add_computation = nullptr; + { + HloComputation::Builder builder(TestName() + ".add"); + const Shape scalar_shape = ShapeUtil::MakeShape(F32, {}); + HloInstruction* p0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, scalar_shape, "p0")); + HloInstruction* p1 = builder.AddInstruction( + HloInstruction::CreateParameter(1, scalar_shape, "p1")); + builder.AddInstruction( + HloInstruction::CreateBinary(scalar_shape, HloOpcode::kAdd, p0, p1)); + add_computation = m->AddEmbeddedComputation(builder.Build()); + } + builder.AddInstruction(HloInstruction::CreateReduceWindow( + ShapeUtil::MakeShape(F32, {1, 1, 1, 4}), param, + builder.AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(0.0f))), + window, add_computation)); + m->AddEntryComputation(builder.Build()); + HloPassFix simplifier(default_options_); + EXPECT_THAT(m->entry_computation()->root_instruction(), + GmockMatch(m::ReduceWindow(m::Parameter(0), m::Constant()))); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); + EXPECT_THAT( + m->entry_computation()->root_instruction(), + GmockMatch(m::Reshape(m::Reduce(m::Parameter(0), m::Constant())))); } TEST_F(AlgebraicSimplifierTest, ZeroSizedReduceWindow) { + auto m = CreateNewVerifiedModule(); auto builder = HloComputation::Builder(TestName()); HloInstruction* param = builder.AddInstruction(HloInstruction::CreateParameter( @@ -1148,24 +1426,24 @@ TEST_F(AlgebraicSimplifierTest, ZeroSizedReduceWindow) { HloInstruction::CreateParameter(1, scalar_shape, "p1")); builder.AddInstruction( HloInstruction::CreateBinary(scalar_shape, HloOpcode::kAdd, p0, p1)); - add_computation = module().AddEmbeddedComputation(builder.Build()); + add_computation = m->AddEmbeddedComputation(builder.Build()); } builder.AddInstruction(HloInstruction::CreateReduceWindow( ShapeUtil::MakeShape(F32, {5, 2}), param, builder.AddInstruction( HloInstruction::CreateConstant(LiteralUtil::CreateR0(0.0f))), window, add_computation)); - module().AddEntryComputation(builder.Build()); - HloPassFix simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - EXPECT_THAT(module().entry_computation()->root_instruction(), - op::ReduceWindow(param, op::Constant())); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); - EXPECT_THAT(module().entry_computation()->root_instruction(), - op::Broadcast(op::Constant())); + m->AddEntryComputation(builder.Build()); + HloPassFix simplifier(default_options_); + EXPECT_THAT(m->entry_computation()->root_instruction(), + GmockMatch(m::ReduceWindow(m::Parameter(0), m::Constant()))); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); + EXPECT_THAT(m->entry_computation()->root_instruction(), + GmockMatch(m::Broadcast(m::Constant()))); } TEST_F(AlgebraicSimplifierTest, ZeroSizedPad) { + auto m = CreateNewVerifiedModule(); auto builder = HloComputation::Builder(TestName()); HloInstruction* param = builder.AddInstruction(HloInstruction::CreateParameter( @@ -1182,17 +1460,17 @@ TEST_F(AlgebraicSimplifierTest, ZeroSizedPad) { builder.AddInstruction( HloInstruction::CreateConstant(LiteralUtil::CreateR0(0.0f))), padding)); - module().AddEntryComputation(builder.Build()); - EXPECT_THAT(module().entry_computation()->root_instruction(), - op::Pad(param, op::Constant())); - HloPassFix simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); - EXPECT_THAT(module().entry_computation()->root_instruction(), - op::Broadcast(op::Constant())); + m->AddEntryComputation(builder.Build()); + EXPECT_THAT(m->entry_computation()->root_instruction(), + GmockMatch(m::Pad(m::Parameter(0), m::Constant()))); + HloPassFix simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); + EXPECT_THAT(m->entry_computation()->root_instruction(), + GmockMatch(m::Broadcast(m::Constant()))); } TEST_F(AlgebraicSimplifierTest, ReshapeBroadcast) { + auto m = CreateNewVerifiedModule(); Shape r0f32 = ShapeUtil::MakeShape(F32, {}); auto builder = HloComputation::Builder(TestName()); @@ -1206,39 +1484,40 @@ TEST_F(AlgebraicSimplifierTest, ReshapeBroadcast) { ShapeUtil::MakeShape(F32, {3, 2}), broadcast)); auto computation = builder.Build(); - module().AddEntryComputation(std::move(computation)); + m->AddEntryComputation(std::move(computation)); - EXPECT_THAT(module().entry_computation()->root_instruction(), - op::Reshape(op::Broadcast(op::Reshape(op)))); + EXPECT_THAT(m->entry_computation()->root_instruction(), + GmockMatch(m::Reshape(m::Broadcast(m::Reshape(m::Op().Is(op)))))); - HloPassFix simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + HloPassFix simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); - EXPECT_THAT(module().entry_computation()->root_instruction(), op); + EXPECT_THAT(m->entry_computation()->root_instruction(), op); } // Test that convert(A, $TYPE) is simplified to A if A is of type $TYPE. TEST_F(AlgebraicSimplifierTest, ConvertBetweenSameType) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); HloInstruction* input = builder.AddInstruction( HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0f))); builder.AddInstruction( HloInstruction::CreateConvert(ShapeUtil::MakeShape(F32, {}), input)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); - EXPECT_THAT(computation->root_instruction(), op::Convert(input)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Convert(m::Op().Is(input)))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), input); } // Test that copies are removed. TEST_F(AlgebraicSimplifierTest, RemoveCopy) { + auto m = CreateNewVerifiedModule(); Shape r0f32 = ShapeUtil::MakeShape(F32, {}); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction( @@ -1246,46 +1525,107 @@ TEST_F(AlgebraicSimplifierTest, RemoveCopy) { builder.AddInstruction( HloInstruction::CreateUnary(param0->shape(), HloOpcode::kCopy, param0)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); - EXPECT_THAT(computation->root_instruction(), op::Copy(param0)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Copy(m::Parameter(0)))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); 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}); - auto computation = module().AddEntryComputation(builder.Build()); - EXPECT_THAT(computation->root_instruction(), op::Copy(param)); + 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)))); +} - AlgebraicSimplifier simplifier1(/*is_layout_sensitive=*/true, - non_bitcasting_callback()); - ASSERT_FALSE(simplifier1.Run(&module()).ValueOrDie()); +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( + [](const Shape&, const Shape&) { return false; }); + options.set_is_layout_sensitive(true); + AlgebraicSimplifier simplifier1(options); + ASSERT_FALSE(simplifier1.Run(m.get()).ValueOrDie()); // Verify that the copy is not replaced. - EXPECT_THAT(computation->root_instruction(), op::Copy(param)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Copy(m::Parameter(0)))); - AlgebraicSimplifier simplifier2(/*is_layout_sensitive=*/true, - bitcasting_callback()); - ASSERT_TRUE(simplifier2.Run(&module()).ValueOrDie()); + AlgebraicSimplifierOptions options2; + options2.set_is_layout_sensitive(true); + AlgebraicSimplifier simplifier2(options2); + EXPECT_TRUE(simplifier2.Run(m.get()).ValueOrDie()); // Verify that the copy is replaced. - EXPECT_THAT(computation->root_instruction(), op::Bitcast(param)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Bitcast(m::Parameter(0)))); } // Test that unary concatenates are removed. TEST_F(AlgebraicSimplifierTest, RemoveUnaryConcatenate) { + auto m = CreateNewVerifiedModule(); Shape r1f32 = ShapeUtil::MakeShape(F32, {100}); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction( @@ -1293,19 +1633,20 @@ TEST_F(AlgebraicSimplifierTest, RemoveUnaryConcatenate) { builder.AddInstruction( HloInstruction::CreateConcatenate(param0->shape(), {param0}, 0)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); - EXPECT_THAT(computation->root_instruction(), op::Concatenate(param0)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Concatenate(m::Parameter(0)))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), param0); } // Test that empty operands of concatenates are removed. TEST_F(AlgebraicSimplifierTest, RemoveEmptyConcatenateOperands) { + auto m = CreateNewVerifiedModule(); const int kParamLength = 100; Shape r1f32 = ShapeUtil::MakeShape(F32, {kParamLength}); HloComputation::Builder builder(TestName()); @@ -1322,22 +1663,24 @@ TEST_F(AlgebraicSimplifierTest, RemoveEmptyConcatenateOperands) { builder.AddInstruction(HloInstruction::CreateConcatenate( result_shape, {empty_literal, param0, param0, empty_slice, param1}, 0)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); - EXPECT_THAT( - computation->root_instruction(), - op::Concatenate(empty_literal, param0, param0, empty_slice, param1)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Concatenate( + m::Op().Is(empty_literal), m::Parameter(0), m::Parameter(0), + m::Op().Is(empty_slice), m::Parameter(1)))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), - op::Concatenate(param0, param0, param1)); + GmockMatch(m::Concatenate(m::Parameter(0), m::Parameter(0), + m::Parameter(1)))); } // Test that reduce of concat is simplified. TEST_F(AlgebraicSimplifierTest, SimplifyReduceOfConcat) { + auto m = CreateNewVerifiedModule(); const int kParamLength = 100; Shape r3f32 = ShapeUtil::MakeShape(F32, {kParamLength, kParamLength, kParamLength}); @@ -1363,7 +1706,7 @@ TEST_F(AlgebraicSimplifierTest, SimplifyReduceOfConcat) { HloInstruction::CreateParameter(1, scalar_shape, "p1")); builder.AddInstruction( HloInstruction::CreateBinary(scalar_shape, HloOpcode::kAdd, p0, p1)); - add_computation = module().AddEmbeddedComputation(builder.Build()); + add_computation = m->AddEmbeddedComputation(builder.Build()); } Shape r4f32 = ShapeUtil::MakeShape(F32, {4, 5, 6, 7}); Shape reduce_shape = ShapeUtil::MakeShape(F32, {kParamLength}); @@ -1373,20 +1716,21 @@ TEST_F(AlgebraicSimplifierTest, SimplifyReduceOfConcat) { builder.AddInstruction(HloInstruction::CreateReduce( reduce_shape, Concatenate, zero, {1, 2}, add_computation)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); EXPECT_THAT( computation->root_instruction(), - op::Map(op::Map(op::Reduce(param0, zero), op::Reduce(param1, zero)), - op::Reduce(param2, zero))); + GmockMatch(m::Map(m::Map(m::Reduce(m::Parameter(0), m::Op().Is(zero)), + m::Reduce(m::Parameter(1), m::Op().Is(zero))), + m::Reduce(m::Parameter(2), m::Op().Is(zero))))); } // Test a concatenate with only empty operands is removed. TEST_F(AlgebraicSimplifierTest, OnlyEmptyConcatenateOperands) { + auto m = CreateNewVerifiedModule(); const int kParamLength = 100; Shape r1f32 = ShapeUtil::MakeShape(F32, {kParamLength}); HloComputation::Builder builder(TestName()); @@ -1401,20 +1745,21 @@ TEST_F(AlgebraicSimplifierTest, OnlyEmptyConcatenateOperands) { builder.AddInstruction(HloInstruction::CreateConcatenate( result_shape, {empty_literal, empty_slice}, 0)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), - op::Concatenate(empty_literal, empty_slice)); + GmockMatch(m::Concatenate(m::Op().Is(empty_literal), + m::Op().Is(empty_slice)))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); EXPECT_EQ(computation->root_instruction(), empty_literal); } // Test that concat with a scalar broadcast becomes a pad. TEST_F(AlgebraicSimplifierTest, ConcatenateOfBroadcastBecomesPad) { + auto m = CreateNewVerifiedModule(); Shape r1f32 = ShapeUtil::MakeShape(F32, {100}); Shape r0f32 = ShapeUtil::MakeShape(F32, {}); HloComputation::Builder builder(TestName()); @@ -1427,17 +1772,88 @@ TEST_F(AlgebraicSimplifierTest, ConcatenateOfBroadcastBecomesPad) { builder.AddInstruction(HloInstruction::CreateConcatenate( ShapeUtil::MakeShape(F32, {200}), {broadcast, param0}, 0)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); - EXPECT_THAT(computation->root_instruction(), op::Pad(param0, param1)); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Pad(m::Parameter(0), m::Parameter(1)))); +} + +TEST_F(AlgebraicSimplifierTest, SimplifyConcatenateOfSlices) { + auto m = CreateNewVerifiedModule(); + Shape r2f32 = ShapeUtil::MakeShape(F32, {100, 99}); + Shape concat_shape = ShapeUtil::MakeShape(F32, {50, 80}); + HloComputation::Builder builder(TestName()); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, r2f32, "param0")); + HloInstruction* param1 = builder.AddInstruction( + HloInstruction::CreateParameter(1, r2f32, "param1")); + + HloInstruction* slice0 = builder.AddInstruction(HloInstruction::CreateSlice( + ShapeUtil::MakeShape(F32, {50, 10}), param0, /*start_indices=*/{0, 0}, + /*limit_indices=*/{50, 10}, /*strides=*/{1, 1})); + + // Cannot merge 'slice0' and 'slice1' because of different start indices in + // dimension 0. + HloInstruction* slice1 = builder.AddInstruction(HloInstruction::CreateSlice( + ShapeUtil::MakeShape(F32, {50, 10}), param0, /*start_indices=*/{50, 10}, + /*limit_indices=*/{100, 20}, /*strides=*/{1, 1})); + + // Cannot merge 'slice1' and 'slice2' because of stride in dimension 2. + HloInstruction* slice2 = builder.AddInstruction(HloInstruction::CreateSlice( + ShapeUtil::MakeShape(F32, {50, 10}), param0, /*start_indices=*/{50, 20}, + /*limit_indices=*/{100, 40}, /*strides=*/{1, 2})); + + // Cannot merge 'slice2' and 'slice3' because of stride in dimension 2. + HloInstruction* slice3 = builder.AddInstruction(HloInstruction::CreateSlice( + ShapeUtil::MakeShape(F32, {50, 10}), param0, /*start_indices=*/{50, 40}, + /*limit_indices=*/{100, 50}, /*strides=*/{1, 1})); + + // Can merge 'slice3' and 'slice4'. + HloInstruction* slice4 = builder.AddInstruction(HloInstruction::CreateSlice( + ShapeUtil::MakeShape(F32, {50, 10}), param0, /*start_indices=*/{50, 50}, + /*limit_indices=*/{100, 60}, /*strides=*/{1, 1})); + + // Can merge 'slice4' and 'slice5'. + HloInstruction* slice5 = builder.AddInstruction(HloInstruction::CreateSlice( + ShapeUtil::MakeShape(F32, {50, 10}), param0, /*start_indices=*/{50, 60}, + /*limit_indices=*/{100, 70}, /*strides=*/{1, 1})); + + // Cannot merge 'slice5' and 'slice6' because of overlap. + HloInstruction* slice6 = builder.AddInstruction(HloInstruction::CreateSlice( + ShapeUtil::MakeShape(F32, {50, 10}), param0, /*start_indices=*/{50, 69}, + /*limit_indices=*/{100, 79}, /*strides=*/{1, 1})); + + // Cannot merge 'slice6' and 'slice7' because of slicing from a different + // parameter. + HloInstruction* slice7 = builder.AddInstruction(HloInstruction::CreateSlice( + ShapeUtil::MakeShape(F32, {50, 10}), param1, /*start_indices=*/{50, 79}, + /*limit_indices=*/{100, 89}, /*strides=*/{1, 1})); + + builder.AddInstruction(HloInstruction::CreateConcatenate( + concat_shape, + {slice0, slice1, slice2, slice3, slice4, slice5, slice6, slice7}, 1)); + auto computation = m->AddEntryComputation(builder.Build()); + + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); + auto s = m::Slice(m::Parameter(0)); + EXPECT_THAT( + computation->root_instruction(), + GmockMatch(m::Concatenate(s, s, s, s, s, m::Slice(m::Parameter(1))))); + // The operand 3 should be a merge of 'slice3', 'slice4' and 'slice5', so its + // shape should have dimensions {50, 30}. + EXPECT_TRUE( + ShapeUtil::Equal(computation->root_instruction()->operand(3)->shape(), + ShapeUtil::MakeShape(F32, {50, 30}))); + EXPECT_EQ(computation->root_instruction()->operand(3)->slice_starts(1), 40); } // Test that a simplification which changes layouts is not performed if layout // sensitive is true. TEST_F(AlgebraicSimplifierTest, CopyWithDifferentLayout) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction(HloInstruction::CreateParameter( @@ -1445,25 +1861,29 @@ TEST_F(AlgebraicSimplifierTest, CopyWithDifferentLayout) { HloInstruction* copy = builder.AddInstruction( HloInstruction::CreateUnary(param0->shape(), HloOpcode::kCopy, param0)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); // Set to different layouts. *param0->mutable_shape()->mutable_layout() = LayoutUtil::MakeLayout({0, 1}); *copy->mutable_shape()->mutable_layout() = LayoutUtil::MakeLayout({1, 0}); - EXPECT_THAT(computation->root_instruction(), op::Copy(param0)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Copy(m::Parameter(0)))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/true, - non_bitcasting_callback()); - EXPECT_FALSE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifierOptions options; + options.set_is_layout_sensitive(true); + AlgebraicSimplifier simplifier(options); + EXPECT_FALSE(simplifier.Run(m.get()).ValueOrDie()); // Copy has not been removed. - EXPECT_THAT(computation->root_instruction(), op::Copy(param0)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Copy(m::Parameter(0)))); } // Test that a simplification which preserves layouts is performed if layout // sensitive is true. TEST_F(AlgebraicSimplifierTest, CopyWithSameLayout) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction(HloInstruction::CreateParameter( @@ -1471,17 +1891,19 @@ TEST_F(AlgebraicSimplifierTest, CopyWithSameLayout) { HloInstruction* copy = builder.AddInstruction( HloInstruction::CreateUnary(param0->shape(), HloOpcode::kCopy, param0)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); // Set to same layouts. *param0->mutable_shape()->mutable_layout() = LayoutUtil::MakeLayout({0, 1}); *copy->mutable_shape()->mutable_layout() = LayoutUtil::MakeLayout({0, 1}); - EXPECT_THAT(computation->root_instruction(), op::Copy(param0)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Copy(m::Parameter(0)))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/true, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifierOptions options; + options.set_is_layout_sensitive(true); + AlgebraicSimplifier simplifier(options); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); // Copy has been removed. EXPECT_THAT(computation->root_instruction(), param0); @@ -1490,6 +1912,7 @@ TEST_F(AlgebraicSimplifierTest, CopyWithSameLayout) { // Test that a reshape which could be replaced with a bitcast is not if // add_bitcasts is false. TEST_F(AlgebraicSimplifierTest, NoBitcastAdded) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction(HloInstruction::CreateParameter( @@ -1502,20 +1925,25 @@ TEST_F(AlgebraicSimplifierTest, NoBitcastAdded) { *reshape->mutable_shape()->mutable_layout() = LayoutUtil::MakeLayout({0, 1, 2, 3, 4, 5}); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); - EXPECT_THAT(computation->root_instruction(), op::Reshape(param0)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Reshape(m::Parameter(0)))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/true, - non_bitcasting_callback()); - EXPECT_FALSE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifierOptions options( + [](const Shape&, const Shape&) { return false; }); + options.set_is_layout_sensitive(true); + AlgebraicSimplifier simplifier(options); + EXPECT_FALSE(simplifier.Run(m.get()).ValueOrDie()); // Reshape is not replaced with a bitcast. - EXPECT_THAT(computation->root_instruction(), op::Reshape(param0)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Reshape(m::Parameter(0)))); } // Test transforming reshapes and transposes of rng. TEST_F(AlgebraicSimplifierTest, ReshapeOfTransposeOfRngToRng) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); HloInstruction* zero = builder.AddInstruction( HloInstruction::CreateConstant(LiteralUtil::CreateR0(0.0f))); @@ -1532,21 +1960,21 @@ TEST_F(AlgebraicSimplifierTest, ReshapeOfTransposeOfRngToRng) { ShapeUtil::MakeShape(F32, {4}), transpose)) ->shape(); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - bitcasting_callback()); - EXPECT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(AlgebraicSimplifierOptions{}); + EXPECT_TRUE(simplifier.Run(m.get()).ValueOrDie()); - // Verify that that reshape(transpose(rng)) is replace by a single rng of the + // Verify that reshape(transpose(rng)) is replace by a single rng of the // same shape as the reshape. - EXPECT_THAT(computation->root_instruction(), op::Rng()); + EXPECT_THAT(computation->root_instruction(), GmockMatch(m::Rng())); EXPECT_TRUE(ShapeUtil::Equal(computation->root_instruction()->shape(), reshape_shape)); } // Test transforming reshapes to bitcasts under various conditions. TEST_F(AlgebraicSimplifierTest, ReshapeReplacedWithBitcast) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction(HloInstruction::CreateParameter( @@ -1578,25 +2006,29 @@ TEST_F(AlgebraicSimplifierTest, ReshapeReplacedWithBitcast) { builder.AddInstruction(HloInstruction::CreateTuple( {transformable_reshape, dimensions_wrong_reshape, layout_wrong_reshape})); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), - op::Tuple(transformable_reshape, dimensions_wrong_reshape, - layout_wrong_reshape)); + GmockMatch(m::Tuple(m::Op().Is(transformable_reshape), + m::Op().Is(dimensions_wrong_reshape), + m::Op().Is(layout_wrong_reshape)))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/true, - bitcasting_callback()); - simplifier.Run(&module()).ValueOrDie(); + AlgebraicSimplifierOptions options; + options.set_is_layout_sensitive(true); + AlgebraicSimplifier simplifier(options); + simplifier.Run(m.get()).ValueOrDie(); // Verify that only the first reshape is replaced. EXPECT_THAT( computation->root_instruction(), - op::Tuple(op::Bitcast(), dimensions_wrong_reshape, layout_wrong_reshape)); + GmockMatch(m::Tuple(m::Bitcast(), m::Op().Is(dimensions_wrong_reshape), + m::Op().Is(layout_wrong_reshape)))); } // Regression test for a bug where if we failed to sink a reshape, we'd set the // 'changed' bit in AlgebraicSimplifier to false. TEST_F(AlgebraicSimplifierTest, FailureToSinkReshapeDoesntAffectChangedBit) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); // This add (param0 + 0) can be simplified. @@ -1611,15 +2043,15 @@ TEST_F(AlgebraicSimplifierTest, FailureToSinkReshapeDoesntAffectChangedBit) { builder.AddInstruction( HloInstruction::CreateReshape(ShapeUtil::MakeShape(F32, {4}), add)); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - bitcasting_callback()); - module().AddEntryComputation(builder.Build()); - EXPECT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(AlgebraicSimplifierOptions{}); + m->AddEntryComputation(builder.Build()); + EXPECT_TRUE(simplifier.Run(m.get()).ValueOrDie()); } // Regression test for a bug where if we failed to sink a reshape, we'd set the // 'changed' bit in AlgebraicSimplifier to false. TEST_F(AlgebraicSimplifierTest, FailureToSinkBroadcastDoesntAffectChangedBit) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); // This add (param0 + 0) can be simplified. @@ -1635,13 +2067,13 @@ TEST_F(AlgebraicSimplifierTest, FailureToSinkBroadcastDoesntAffectChangedBit) { HloInstruction::CreateBroadcast(ShapeUtil::MakeShape(F32, {2, 2, 2}), add, /*broadcast_dimensions=*/{0, 1})); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - bitcasting_callback()); - module().AddEntryComputation(builder.Build()); - EXPECT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(AlgebraicSimplifierOptions{}); + m->AddEntryComputation(builder.Build()); + EXPECT_TRUE(simplifier.Run(m.get()).ValueOrDie()); } TEST_F(AlgebraicSimplifierTest, TransposeEqualsBitcast1) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); HloInstruction* param = builder.AddInstruction(HloInstruction::CreateParameter( @@ -1655,19 +2087,23 @@ TEST_F(AlgebraicSimplifierTest, TransposeEqualsBitcast1) { *transpose->mutable_shape()->mutable_layout() = LayoutUtil::MakeLayout({0, 1, 2, 3}); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); - EXPECT_THAT(computation->root_instruction(), op::Transpose(param)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Transpose(m::Parameter(0)))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/true, - bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifierOptions options; + options.set_is_layout_sensitive(true); + AlgebraicSimplifier simplifier(options); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); // Verify that the reshape is replaced. - EXPECT_THAT(computation->root_instruction(), op::Bitcast(param)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Bitcast(m::Parameter(0)))); } TEST_F(AlgebraicSimplifierTest, TransposeEqualsBitcast2) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); HloInstruction* param = builder.AddInstruction(HloInstruction::CreateParameter( @@ -1681,19 +2117,23 @@ TEST_F(AlgebraicSimplifierTest, TransposeEqualsBitcast2) { *transpose->mutable_shape()->mutable_layout() = LayoutUtil::MakeLayout({3, 1, 2, 0}); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); - EXPECT_THAT(computation->root_instruction(), op::Transpose(param)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Transpose(m::Parameter(0)))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/true, - bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifierOptions options; + options.set_is_layout_sensitive(true); + AlgebraicSimplifier simplifier(options); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); // Verify that the reshape is replaced. - EXPECT_THAT(computation->root_instruction(), op::Bitcast(param)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Bitcast(m::Parameter(0)))); } TEST_F(AlgebraicSimplifierTest, ReshapesMerged) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction(HloInstruction::CreateParameter( @@ -1706,19 +2146,20 @@ TEST_F(AlgebraicSimplifierTest, ReshapesMerged) { builder.AddInstruction(HloInstruction::CreateReshape( ShapeUtil::MakeShape(F32, {1, 2, 1, 1, 2, 1}), reshape1)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), - op::Reshape(op::Reshape(param0))); + GmockMatch(m::Reshape(m::Reshape(m::Parameter(0))))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); - EXPECT_THAT(computation->root_instruction(), op::Reshape(param0)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Reshape(m::Parameter(0)))); } TEST_F(AlgebraicSimplifierTest, CopiesMerged) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction(HloInstruction::CreateParameter( @@ -1733,18 +2174,22 @@ TEST_F(AlgebraicSimplifierTest, CopiesMerged) { ShapeUtil::MakeShapeWithLayout(F32, {2, 2, 2}, {0, 2, 1}), HloOpcode::kCopy, copy1)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); - EXPECT_THAT(computation->root_instruction(), op::Copy(op::Copy(param0))); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Copy(m::Copy(m::Parameter(0))))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/true, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifierOptions options; + options.set_is_layout_sensitive(true); + AlgebraicSimplifier simplifier(options); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); - EXPECT_THAT(computation->root_instruction(), op::Copy(param0)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Copy(m::Parameter(0)))); } TEST_F(AlgebraicSimplifierTest, TransposesMerged) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction(HloInstruction::CreateParameter( @@ -1757,21 +2202,43 @@ TEST_F(AlgebraicSimplifierTest, TransposesMerged) { builder.AddInstruction(HloInstruction::CreateTranspose( ShapeUtil::MakeShape(F32, {4, 3, 2}), transpose1, {1, 0, 2})); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); - EXPECT_THAT(computation->root_instruction(), op::Transpose(transpose1)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Transpose(m::Op().Is(transpose1)))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); - EXPECT_THAT(computation->root_instruction(), op::Transpose(param0)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Transpose(m::Parameter(0)))); EXPECT_EQ(std::vector({2, 1, 0}), computation->root_instruction()->dimensions()); } +TEST_F(AlgebraicSimplifierTest, TransposeIsReshape) { + const char* hlo_string = R"( + HloModule module + + ENTRY test { + param = f32[10] parameter(0) + reshaped = f32[1,1,10] reshape(f32[10] param) + transposed = f32[10,1,1] transpose(f32[1,1,10] reshaped), dimensions={2,1,0} + ROOT reshaped_again = f32[10] reshape(f32[10,1,1] transposed) + } + )"; + TF_ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(hlo_string)); + + HloPassFix simplifier(default_options_); + EXPECT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + auto root = module->entry_computation()->root_instruction(); + EXPECT_THAT(root, GmockMatch(m::Parameter())); +} + // Test merging reshape and broadcast. TEST_F(AlgebraicSimplifierTest, ReshapeAndBroadcastMerged) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); auto param0 = builder.AddInstruction(HloInstruction::CreateParameter( 0, ShapeUtil::MakeShape(F32, {5}), "param0")); @@ -1780,20 +2247,21 @@ TEST_F(AlgebraicSimplifierTest, ReshapeAndBroadcastMerged) { builder.AddInstruction(HloInstruction::CreateBroadcast( ShapeUtil::MakeShape(F32, {1, 2, 3, 5, 1}), reshape1, {0, 3, 2})); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), - op::Broadcast(op::Reshape(param0))); + GmockMatch(m::Broadcast(m::Reshape(m::Parameter(0))))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); - EXPECT_THAT(computation->root_instruction(), op::Broadcast(param0)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Broadcast(m::Parameter(0)))); } // Test merging broadcast and reshape. TEST_F(AlgebraicSimplifierTest, BroadcastAndReshapeMerged) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); auto param0 = builder.AddInstruction(HloInstruction::CreateParameter( 0, ShapeUtil::MakeShape(F32, {2, 3}), "param0")); @@ -1802,19 +2270,20 @@ TEST_F(AlgebraicSimplifierTest, BroadcastAndReshapeMerged) { builder.AddInstruction(HloInstruction::CreateReshape( ShapeUtil::MakeShape(F32, {2, 3, 7, 2, 1, 3, 2}), broadcast1)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), - op::Reshape(op::Broadcast(param0))); + GmockMatch(m::Reshape(m::Broadcast(m::Parameter(0))))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); - EXPECT_THAT(computation->root_instruction(), op::Broadcast(param0)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Broadcast(m::Parameter(0)))); } TEST_F(AlgebraicSimplifierTest, BroadcastAndReshape_1_3x1_3) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); auto param = builder.AddInstruction(HloInstruction::CreateParameter( 0, ShapeUtil::MakeShape(F32, {1}), "param")); @@ -1823,20 +2292,20 @@ TEST_F(AlgebraicSimplifierTest, BroadcastAndReshape_1_3x1_3) { builder.AddInstruction( HloInstruction::CreateReshape(ShapeUtil::MakeShape(F32, {3}), broadcast)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), - op::Reshape(op::Broadcast(param))); + GmockMatch(m::Reshape(m::Broadcast(m::Parameter(0))))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - EXPECT_FALSE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + EXPECT_FALSE(simplifier.Run(m.get()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), - op::Reshape(op::Broadcast(param))); + GmockMatch(m::Reshape(m::Broadcast(m::Parameter(0))))); } TEST_F(AlgebraicSimplifierTest, BroadcastAndReshape_4_3x2x4_6x1x1x4) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); auto param = builder.AddInstruction(HloInstruction::CreateParameter( 0, ShapeUtil::MakeShape(F32, {4}), "param")); @@ -1845,21 +2314,22 @@ TEST_F(AlgebraicSimplifierTest, BroadcastAndReshape_4_3x2x4_6x1x1x4) { builder.AddInstruction(HloInstruction::CreateReshape( ShapeUtil::MakeShape(F32, {6, 1, 1, 4}), broadcast)); - HloComputation* computation = module().AddEntryComputation(builder.Build()); + HloComputation* computation = m->AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), - op::Reshape(op::Broadcast(param))); + GmockMatch(m::Reshape(m::Broadcast(m::Parameter(0))))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); - EXPECT_THAT(computation->root_instruction(), op::Broadcast(param)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Broadcast(m::Parameter(0)))); EXPECT_THAT(computation->root_instruction()->dimensions(), ::testing::ElementsAre(3)); } TEST_F(AlgebraicSimplifierTest, BroadcastAndReshape_1_3x2x1_6x1x1x1) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); auto param = builder.AddInstruction(HloInstruction::CreateParameter( 0, ShapeUtil::MakeShape(F32, {1}), "param")); @@ -1868,16 +2338,16 @@ TEST_F(AlgebraicSimplifierTest, BroadcastAndReshape_1_3x2x1_6x1x1x1) { builder.AddInstruction(HloInstruction::CreateReshape( ShapeUtil::MakeShape(F32, {6, 1, 1, 1}), broadcast)); - HloComputation* computation = module().AddEntryComputation(builder.Build()); + HloComputation* computation = m->AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), - op::Reshape(op::Broadcast(param))); + GmockMatch(m::Reshape(m::Broadcast(m::Parameter(0))))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); - EXPECT_THAT(computation->root_instruction(), op::Broadcast(param)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Broadcast(m::Parameter(0)))); const std::vector broadcast_dims = computation->root_instruction()->dimensions(); EXPECT_EQ(1, broadcast_dims.size()); @@ -1885,6 +2355,7 @@ TEST_F(AlgebraicSimplifierTest, BroadcastAndReshape_1_3x2x1_6x1x1x1) { } TEST_F(AlgebraicSimplifierTest, BroadcastAndReshape_4_3x2x4x2_6x8) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); auto param = builder.AddInstruction(HloInstruction::CreateParameter( 0, ShapeUtil::MakeShape(F32, {4}), "param")); @@ -1893,115 +2364,119 @@ TEST_F(AlgebraicSimplifierTest, BroadcastAndReshape_4_3x2x4x2_6x8) { builder.AddInstruction(HloInstruction::CreateReshape( ShapeUtil::MakeShape(F32, {6, 8}), broadcast)); - HloComputation* computation = module().AddEntryComputation(builder.Build()); + HloComputation* computation = m->AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), - op::Reshape(op::Broadcast(param))); + GmockMatch(m::Reshape(m::Broadcast(m::Parameter(0))))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - EXPECT_FALSE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + EXPECT_FALSE(simplifier.Run(m.get()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), - op::Reshape(op::Broadcast(param))); + GmockMatch(m::Reshape(m::Broadcast(m::Parameter(0))))); } TEST_F(AlgebraicSimplifierTest, IotaAndReshapeMerged) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); auto iota = builder.AddInstruction(HloInstruction::CreateIota( ShapeUtil::MakeShape(F32, {1, 2, 3, 7, 12, 1}), 2)); Shape result_shape = ShapeUtil::MakeShape(F32, {2, 3, 7, 2, 1, 3, 2}); builder.AddInstruction(HloInstruction::CreateReshape(result_shape, iota)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); - EXPECT_THAT(computation->root_instruction(), op::Reshape(op::Iota())); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Reshape(m::Iota()))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); - EXPECT_THAT(computation->root_instruction(), op::Iota()); + EXPECT_THAT(computation->root_instruction(), GmockMatch(m::Iota())); EXPECT_TRUE( ShapeUtil::Equal(computation->root_instruction()->shape(), result_shape)); } TEST_F(AlgebraicSimplifierTest, IotaEffectiveScalar) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); auto iota = builder.AddInstruction( HloInstruction::CreateIota(ShapeUtil::MakeShape(F32, {1, 1}), 0)); auto result_shape = iota->shape(); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); - EXPECT_THAT(computation->root_instruction(), op::Iota()); + EXPECT_THAT(computation->root_instruction(), GmockMatch(m::Iota())); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); auto root = computation->root_instruction(); - EXPECT_THAT(root, op::Broadcast(op::Constant())); + EXPECT_THAT(root, GmockMatch(m::Broadcast(m::Constant()))); EXPECT_EQ(0.0f, root->operand(0)->literal().GetFirstElement()); EXPECT_TRUE( ShapeUtil::Equal(computation->root_instruction()->shape(), result_shape)); } TEST_F(AlgebraicSimplifierTest, IotaAndReshape_1_3x2_6) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); auto iota = builder.AddInstruction( HloInstruction::CreateIota(ShapeUtil::MakeShape(F32, {3, 2}), 1)); builder.AddInstruction( HloInstruction::CreateReshape(ShapeUtil::MakeShape(F32, {6}), iota)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); - EXPECT_THAT(computation->root_instruction(), op::Reshape(op::Iota())); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Reshape(m::Iota()))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - EXPECT_FALSE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + EXPECT_FALSE(simplifier.Run(m.get()).ValueOrDie()); - EXPECT_THAT(computation->root_instruction(), op::Reshape(op::Iota())); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Reshape(m::Iota()))); } TEST_F(AlgebraicSimplifierTest, IotaAndReshape_4_3x2x4_6x1x1x4) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); auto iota = builder.AddInstruction( HloInstruction::CreateIota(ShapeUtil::MakeShape(F32, {3, 2, 4}), 2)); builder.AddInstruction(HloInstruction::CreateReshape( ShapeUtil::MakeShape(F32, {6, 1, 1, 4}), iota)); - HloComputation* computation = module().AddEntryComputation(builder.Build()); + HloComputation* computation = m->AddEntryComputation(builder.Build()); - EXPECT_THAT(computation->root_instruction(), op::Reshape(op::Iota())); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Reshape(m::Iota()))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); - EXPECT_THAT(computation->root_instruction(), op::Iota()); + EXPECT_THAT(computation->root_instruction(), GmockMatch(m::Iota())); EXPECT_EQ(Cast(computation->root_instruction()) ->iota_dimension(), 3); } TEST_F(AlgebraicSimplifierTest, IotaAndReshape_1_3x2x2_6x1x1x2) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); auto iota = builder.AddInstruction( HloInstruction::CreateIota(ShapeUtil::MakeShape(F32, {3, 2, 2}), 2)); builder.AddInstruction(HloInstruction::CreateReshape( ShapeUtil::MakeShape(F32, {6, 1, 1, 2}), iota)); - HloComputation* computation = module().AddEntryComputation(builder.Build()); + HloComputation* computation = m->AddEntryComputation(builder.Build()); - EXPECT_THAT(computation->root_instruction(), op::Reshape(op::Iota())); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Reshape(m::Iota()))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); - EXPECT_THAT(computation->root_instruction(), op::Iota()); + EXPECT_THAT(computation->root_instruction(), GmockMatch(m::Iota())); const int64 iota_dim = Cast(computation->root_instruction()) ->iota_dimension(); @@ -2009,21 +2484,23 @@ TEST_F(AlgebraicSimplifierTest, IotaAndReshape_1_3x2x2_6x1x1x2) { } TEST_F(AlgebraicSimplifierTest, IotaAndReshape_4_3x2x4x2_6x8) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); auto iota = builder.AddInstruction( HloInstruction::CreateIota(ShapeUtil::MakeShape(F32, {3, 2, 4, 2}), 2)); builder.AddInstruction( HloInstruction::CreateReshape(ShapeUtil::MakeShape(F32, {6, 8}), iota)); - HloComputation* computation = module().AddEntryComputation(builder.Build()); + HloComputation* computation = m->AddEntryComputation(builder.Build()); - EXPECT_THAT(computation->root_instruction(), op::Reshape(op::Iota())); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Reshape(m::Iota()))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - EXPECT_FALSE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + EXPECT_FALSE(simplifier.Run(m.get()).ValueOrDie()); - EXPECT_THAT(computation->root_instruction(), op::Reshape(op::Iota())); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Reshape(m::Iota()))); } TEST_F(AlgebraicSimplifierTest, RemoveNoopPad) { @@ -2043,14 +2520,14 @@ TEST_F(AlgebraicSimplifierTest, RemoveNoopPad) { builder.AddInstruction(HloInstruction::CreatePad( ShapeUtil::MakeShape(F32, {2, 2}), param, zero, no_padding)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* computation = module->AddEntryComputation(builder.Build()); - EXPECT_THAT(computation->root_instruction(), op::Pad(param, zero)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Pad(m::Parameter(0), m::Op().Is(zero)))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), param); } @@ -2076,11 +2553,10 @@ TEST_F(AlgebraicSimplifierTest, NegativePadding) { HloInstruction* pad = builder.AddInstruction(HloInstruction::CreatePad( ShapeUtil::MakeShape(F32, {11, 5}), param, zero, padding)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* computation = module->AddEntryComputation(builder.Build()); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); + AlgebraicSimplifier simplifier(default_options_); auto has_negative_padding = [](const HloInstruction* pad) { for (auto& padding_dimension : pad->padding_config().dimensions()) { @@ -2092,16 +2568,54 @@ TEST_F(AlgebraicSimplifierTest, NegativePadding) { return false; }; - EXPECT_THAT(computation->root_instruction(), op::Pad(param, zero)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Pad(m::Parameter(0), m::Op().Is(zero)))); EXPECT_TRUE(has_negative_padding(pad)); - ASSERT_TRUE(simplifier.Run(module).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); - EXPECT_THAT(computation->root_instruction(), op::Slice(op::Pad(param, zero))); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Slice(m::Pad(m::Parameter(0), m::Op().Is(zero))))); EXPECT_FALSE( has_negative_padding(computation->root_instruction()->operand(0))); } +TEST_F(AlgebraicSimplifierTest, TrivialInteriorPadding) { + // Verify that a pad instruction with interior padding on one-sized + // dimensions, removes the interior padding. + HloComputation::Builder builder(TestName()); + HloInstruction* param = + builder.AddInstruction(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShape(F32, {2, 1}), "param")); + HloInstruction* zero = builder.AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(0.0f))); + PaddingConfig padding; + for (int i = 0; i < 2; ++i) { + auto dimension = padding.add_dimensions(); + dimension->set_edge_padding_low(3); + dimension->set_edge_padding_high(3); + dimension->set_interior_padding(i * 3); + } + HloInstruction* pad = builder.AddInstruction(HloInstruction::CreatePad( + ShapeUtil::MakeShape(F32, {8, 7}), param, zero, padding)); + + auto module = CreateNewVerifiedModule(); + HloComputation* computation = module->AddEntryComputation(builder.Build()); + + AlgebraicSimplifier simplifier(default_options_); + + ASSERT_THAT(computation->root_instruction(), + GmockMatch(m::Pad(m::Parameter(0), m::Op().Is(zero)))); + ASSERT_TRUE(HasInteriorPadding(pad->padding_config())); + + EXPECT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Pad(m::Parameter(0), m::Op().Is(zero)))); + EXPECT_FALSE( + HasInteriorPadding(computation->root_instruction()->padding_config())); +} + TEST_F(AlgebraicSimplifierTest, RemoveNoopReshape) { HloComputation::Builder builder(TestName()); HloInstruction* param = @@ -2110,14 +2624,14 @@ TEST_F(AlgebraicSimplifierTest, RemoveNoopReshape) { builder.AddInstruction( HloInstruction::CreateReshape(ShapeUtil::MakeShape(F32, {2, 3}), param)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* computation = module->AddEntryComputation(builder.Build()); - EXPECT_THAT(computation->root_instruction(), op::Reshape(param)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Reshape(m::Parameter(0)))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), param); } @@ -2133,14 +2647,14 @@ TEST_F(AlgebraicSimplifierTest, RemoveNoopSlice) { ShapeUtil::MakeShape(F32, {dim0, dim1}), param, /*start_indices=*/{0, 0}, /*limit_indices=*/{dim0, dim1}, /*strides=*/{1, 1})); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* computation = module->AddEntryComputation(builder.Build()); - EXPECT_THAT(computation->root_instruction(), op::Slice(param)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Slice(m::Parameter(0)))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), param); } @@ -2162,22 +2676,75 @@ TEST_F(AlgebraicSimplifierTest, SliceOfSliceToSlice) { ShapeUtil::MakeShape(F32, {dim0 - 5, dim1 - 9}), original_slice, /*start_indices=*/{2, 3}, /*limit_indices=*/{dim0 - 3, dim1 - 6}, /*strides=*/{1, 1})); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* computation = module->AddEntryComputation(builder.Build()); - EXPECT_THAT(computation->root_instruction(), op::Slice(op::Slice(param))); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Slice(m::Slice(m::Parameter(0))))); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); - EXPECT_THAT(computation->root_instruction(), op::Slice(param)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Slice(m::Parameter(0)))); EXPECT_EQ(computation->root_instruction()->slice_starts(0), 3); EXPECT_EQ(computation->root_instruction()->slice_starts(1), 5); EXPECT_EQ(computation->root_instruction()->slice_limits(0), dim0 - 2); EXPECT_EQ(computation->root_instruction()->slice_limits(1), dim1 - 4); } +TEST_F(AlgebraicSimplifierTest, SliceOfReshapeToReshapeOfSlice) { + HloComputation::Builder builder(TestName()); + const int64 dim0 = 11; + const int64 dim1 = 12; + const int64 dim2 = 13; + HloInstruction* param = + builder.AddInstruction(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShape(F32, {dim0 * dim1, dim2}), "param")); + HloInstruction* original_reshape = + builder.AddInstruction(HloInstruction::CreateReshape( + ShapeUtil::MakeShape(F32, {dim0, dim1, dim2}), param)); + + builder.AddInstruction(HloInstruction::CreateSlice( + ShapeUtil::MakeShape(F32, {dim0 - 2, dim1, dim2}), original_reshape, + /*start_indices=*/{0, 0, 0}, + /*limit_indices=*/{dim0 - 2, dim1, dim2}, /*strides=*/{1, 1, 1})); + auto module = CreateNewVerifiedModule(); + HloComputation* computation = module->AddEntryComputation(builder.Build()); + + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Slice(m::Reshape(m::Parameter(0))))); + + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Reshape(m::Slice(m::Parameter(0))))); +} + +TEST_F(AlgebraicSimplifierTest, SliceOfReshapeUnchanged) { + HloComputation::Builder builder(TestName()); + HloInstruction* param = + builder.AddInstruction(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShape(F32, {1, 144, 25, 1, 512}), "param")); + HloInstruction* original_reshape = + builder.AddInstruction(HloInstruction::CreateReshape( + ShapeUtil::MakeShape(F32, {3600, 512}), param)); + + builder.AddInstruction(HloInstruction::CreateSlice( + ShapeUtil::MakeShape(F32, {960, 512}), original_reshape, + /*start_indices=*/{0, 0}, + /*limit_indices=*/{960, 512}, /*strides=*/{1, 1})); + auto module = CreateNewVerifiedModule(); + HloComputation* computation = module->AddEntryComputation(builder.Build()); + + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Slice(m::Reshape(m::Parameter(0))))); + + AlgebraicSimplifier simplifier(default_options_); + ASSERT_FALSE(simplifier.Run(module.get()).ValueOrDie()); +} + TEST_F(AlgebraicSimplifierTest, RemoveNoopSort) { auto builder = HloComputation::Builder(TestName()); @@ -2185,14 +2752,86 @@ TEST_F(AlgebraicSimplifierTest, RemoveNoopSort) { auto keys = builder.AddInstruction( HloInstruction::CreateParameter(0, keys_shape, "keys")); builder.AddInstruction(HloInstruction::CreateSort(keys_shape, 0, keys)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* computation = module->AddEntryComputation(builder.Build()); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module).ValueOrDie()); + 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; + 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; + 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; + options.set_enable_permutation_sort_replacement(true); + AlgebraicSimplifier simplifier(options); + EXPECT_FALSE(simplifier.Run(module.get()).ValueOrDie()); +} + TEST_F(AlgebraicSimplifierTest, ReplaceEffectiveScalarKeyValueSortWithTuple) { auto builder = HloComputation::Builder(TestName()); @@ -2207,13 +2846,181 @@ TEST_F(AlgebraicSimplifierTest, ReplaceEffectiveScalarKeyValueSortWithTuple) { builder.AddInstruction(HloInstruction::CreateSort( ShapeUtil::MakeTupleShape({keys_shape, values_shape, values_shape}), 0, keys, {values0, values1})); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* computation = module->AddEntryComputation(builder.Build()); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), - op::Tuple(keys, values0, values1)); + GmockMatch(m::Tuple(m::Op().Is(keys), m::Op().Is(values0), + m::Op().Is(values1)))); +} + +// Test that A && True is simplified to A +TEST_F(AlgebraicSimplifierTest, AndTrue) { + auto m = CreateNewVerifiedModule(); + Shape r0pred = ShapeUtil::MakeShape(PRED, {}); + HloComputation::Builder builder(TestName()); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, r0pred, "param0")); + HloInstruction* const_true = builder.AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(true))); + builder.AddInstruction(HloInstruction::CreateBinary(r0pred, HloOpcode::kAnd, + param0, const_true)); + + auto computation = m->AddEntryComputation(builder.Build()); + HloInstruction* root = computation->root_instruction(); + EXPECT_EQ(root->opcode(), HloOpcode::kAnd); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); + root = computation->root_instruction(); + EXPECT_EQ(root, param0); +} + +// Test that True && A is simplified to A +TEST_F(AlgebraicSimplifierTest, AndTrue2) { + auto m = CreateNewVerifiedModule(); + Shape r0pred = ShapeUtil::MakeShape(PRED, {}); + HloComputation::Builder builder(TestName()); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, r0pred, "param0")); + HloInstruction* const_true = builder.AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(true))); + builder.AddInstruction(HloInstruction::CreateBinary(r0pred, HloOpcode::kAnd, + const_true, param0)); + + auto computation = m->AddEntryComputation(builder.Build()); + HloInstruction* root = computation->root_instruction(); + EXPECT_EQ(root->opcode(), HloOpcode::kAnd); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); + root = computation->root_instruction(); + EXPECT_EQ(root, param0); +} + +// Test that A && False is simplified to False +TEST_F(AlgebraicSimplifierTest, AndFalse) { + auto m = CreateNewVerifiedModule(); + Shape r0pred = ShapeUtil::MakeShape(PRED, {}); + HloComputation::Builder builder(TestName()); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, r0pred, "param0")); + HloInstruction* const_false = builder.AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(false))); + builder.AddInstruction(HloInstruction::CreateBinary(r0pred, HloOpcode::kAnd, + param0, const_false)); + + auto computation = m->AddEntryComputation(builder.Build()); + HloInstruction* root = computation->root_instruction(); + EXPECT_EQ(root->opcode(), HloOpcode::kAnd); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); + root = computation->root_instruction(); + EXPECT_EQ(root, const_false); +} + +// Test that False && A is simplified to False +TEST_F(AlgebraicSimplifierTest, AndFalse2) { + auto m = CreateNewVerifiedModule(); + Shape r0pred = ShapeUtil::MakeShape(PRED, {}); + HloComputation::Builder builder(TestName()); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, r0pred, "param0")); + HloInstruction* const_false = builder.AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(false))); + builder.AddInstruction(HloInstruction::CreateBinary(r0pred, HloOpcode::kAnd, + const_false, param0)); + + auto computation = m->AddEntryComputation(builder.Build()); + HloInstruction* root = computation->root_instruction(); + EXPECT_EQ(root->opcode(), HloOpcode::kAnd); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); + root = computation->root_instruction(); + EXPECT_EQ(root, const_false); +} + +// Test that A || True is simplified to True +TEST_F(AlgebraicSimplifierTest, OrTrue) { + auto m = CreateNewVerifiedModule(); + Shape r0pred = ShapeUtil::MakeShape(PRED, {}); + HloComputation::Builder builder(TestName()); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, r0pred, "param0")); + HloInstruction* const_true = builder.AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(true))); + builder.AddInstruction( + HloInstruction::CreateBinary(r0pred, HloOpcode::kOr, param0, const_true)); + + auto computation = m->AddEntryComputation(builder.Build()); + HloInstruction* root = computation->root_instruction(); + EXPECT_EQ(root->opcode(), HloOpcode::kOr); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); + root = computation->root_instruction(); + EXPECT_EQ(root, const_true); +} + +// Test that True || A is simplified to True +TEST_F(AlgebraicSimplifierTest, OrTrue2) { + auto m = CreateNewVerifiedModule(); + Shape r0pred = ShapeUtil::MakeShape(PRED, {}); + HloComputation::Builder builder(TestName()); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, r0pred, "param0")); + HloInstruction* const_true = builder.AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(true))); + builder.AddInstruction( + HloInstruction::CreateBinary(r0pred, HloOpcode::kOr, const_true, param0)); + + auto computation = m->AddEntryComputation(builder.Build()); + HloInstruction* root = computation->root_instruction(); + EXPECT_EQ(root->opcode(), HloOpcode::kOr); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); + root = computation->root_instruction(); + EXPECT_EQ(root, const_true); +} + +// Test that A || False is simplified to A +TEST_F(AlgebraicSimplifierTest, OrFalse) { + auto m = CreateNewVerifiedModule(); + Shape r0pred = ShapeUtil::MakeShape(PRED, {}); + HloComputation::Builder builder(TestName()); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, r0pred, "param0")); + HloInstruction* const_false = builder.AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(false))); + builder.AddInstruction(HloInstruction::CreateBinary(r0pred, HloOpcode::kOr, + param0, const_false)); + + auto computation = m->AddEntryComputation(builder.Build()); + HloInstruction* root = computation->root_instruction(); + EXPECT_EQ(root->opcode(), HloOpcode::kOr); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); + root = computation->root_instruction(); + EXPECT_EQ(root, param0); +} + +// Test that False || A is simplified to A +TEST_F(AlgebraicSimplifierTest, OrFalse2) { + auto m = CreateNewVerifiedModule(); + Shape r0pred = ShapeUtil::MakeShape(PRED, {}); + HloComputation::Builder builder(TestName()); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, r0pred, "param0")); + HloInstruction* const_false = builder.AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(false))); + builder.AddInstruction(HloInstruction::CreateBinary(r0pred, HloOpcode::kOr, + const_false, param0)); + + auto computation = m->AddEntryComputation(builder.Build()); + HloInstruction* root = computation->root_instruction(); + EXPECT_EQ(root->opcode(), HloOpcode::kOr); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); + root = computation->root_instruction(); + EXPECT_EQ(root, param0); } // Used for TEST_Ps that test merging (or not) of a kPad instruction into a @@ -2261,7 +3068,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. @@ -2332,23 +3139,23 @@ TEST_P(ConvInputPaddingTest, DoTest) { .ValueOrDie(); builder.AddInstruction(HloInstruction::CreateConvolve( ShapeInference::InferConvolveShape(lhs_pad->shape(), filter->shape(), - /*feature_group_count=*/1, window, - dnums) + /*feature_group_count=*/1, + /*batch_group_count=*/1, window, dnums) .ValueOrDie(), - lhs_pad, filter, /*feature_group_count=*/1, window, dnums, - DefaultPrecisionConfig(2))); - auto module = CreateNewModule(); + lhs_pad, filter, /*feature_group_count=*/1, /*batch_group_count=*/1, + window, dnums, DefaultPrecisionConfig(2))); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); + AlgebraicSimplifier simplifier(default_options_); if (testcase.expected_conv_window.empty()) { - ASSERT_FALSE(simplifier.Run(module).ValueOrDie()); + ASSERT_FALSE(simplifier.Run(module.get()).ValueOrDie()); } else { - ASSERT_TRUE(simplifier.Run(module).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); auto* conv = module->entry_computation()->root_instruction(); SCOPED_TRACE(module->ToString()); - ASSERT_THAT(conv, op::Convolution(op::Parameter(), op::Parameter())); + ASSERT_THAT(conv, + GmockMatch(m::Convolution(m::Parameter(), m::Parameter()))); EXPECT_EQ(window_util::ToString(conv->window()), absl::StrCat("size=3x3 ", testcase.expected_conv_window)); } @@ -2369,7 +3176,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; @@ -2449,24 +3256,24 @@ TEST_P(ConvFilterPaddingTest, DoIt) { builder.AddInstruction(HloInstruction::CreateConvolve( ShapeInference::InferConvolveShape(input->shape(), rhs_pad->shape(), - /*feature_group_count=*/1, window, - dnums) + /*feature_group_count=*/1, + /*batch_group_count=*/1, window, dnums) .ValueOrDie(), - input, rhs_pad, /*feature_group_count=*/1, window, dnums, - precision_config)); + input, rhs_pad, /*feature_group_count=*/1, /*batch_group_count=*/1, + window, dnums, precision_config)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); + AlgebraicSimplifier simplifier(default_options_); if (testcase.expected_conv_window.empty()) { - ASSERT_FALSE(simplifier.Run(module).ValueOrDie()); + ASSERT_FALSE(simplifier.Run(module.get()).ValueOrDie()); } else { - ASSERT_TRUE(simplifier.Run(module).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); auto* conv = module->entry_computation()->root_instruction(); SCOPED_TRACE(module->ToString()); - ASSERT_THAT(conv, op::Convolution(op::Parameter(), op::Parameter())); + ASSERT_THAT(conv, + GmockMatch(m::Convolution(m::Parameter(), m::Parameter()))); EXPECT_EQ(window_util::ToString(conv->window()), absl::StrFormat("size=%dx%d %s", conv->operand(1)->shape().dimensions(2), @@ -2601,14 +3408,16 @@ TEST_F(AlgebraicSimplifierTest, ConvertConvToMatmul) { b.AddInstruction(HloInstruction::CreateConvolve( out_shape, input, filter, - /*feature_group_count=*/1, window, dnums, DefaultPrecisionConfig(2))); + /*feature_group_count=*/1, /*batch_group_count=*/1, window, dnums, + DefaultPrecisionConfig(2))); // TODO(b/80488902): verify this module. - auto module = HloTestBase::CreateNewModule(); + auto module = CreateNewUnverifiedModule(); auto* computation = module->AddEntryComputation(b.Build()); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/true, - bitcasting_callback()); + AlgebraicSimplifierOptions simplifier_options; + simplifier_options.set_is_layout_sensitive(true); + AlgebraicSimplifier simplifier(simplifier_options); if (!simplifier.Run(module.get()).ValueOrDie()) { return "NO_CHANGE"; } @@ -2724,24 +3533,22 @@ TEST_F(AlgebraicSimplifierTest, ScalarBroadcastToSlice) { HloInstruction* slice = builder.AddInstruction(HloInstruction::CreateSlice( slice_shape, broadcast, {0, 1, 2, 3}, {2, 3, 5, 6}, {1, 1, 1, 1})); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root, slice); EXPECT_TRUE(ShapeUtil::Equal(root->shape(), slice_shape)); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); + AlgebraicSimplifier simplifier(default_options_); - ASSERT_TRUE(simplifier.Run(module).ValueOrDie()); + ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); // Running simplification again should not result in any further changes. - ASSERT_FALSE(simplifier.Run(module).ValueOrDie()); - - root = computation->root_instruction(); - EXPECT_THAT(root, op::Broadcast(scalar_param)); - EXPECT_TRUE(ShapeUtil::Equal(root->shape(), slice_shape)); + ASSERT_FALSE(simplifier.Run(module.get()).ValueOrDie()); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Broadcast(m::Op().Is(scalar_param)) + .WithShapeEqualTo(&slice_shape))); } // Test that reshape(transpose(broadcast(/*scalar value*/))) simplifies to a @@ -2763,26 +3570,24 @@ TEST_F(AlgebraicSimplifierTest, ScalarBroadcastToTransposeReshape) { HloInstruction* reshape = builder.AddInstruction( HloInstruction::CreateReshape(reshape_shape, transpose)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root, reshape); EXPECT_TRUE(ShapeUtil::Equal(root->shape(), reshape_shape)); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module).ValueOrDie()); - - root = computation->root_instruction(); - EXPECT_THAT(root, op::Broadcast(forty_two)); - EXPECT_TRUE(ShapeUtil::Equal(root->shape(), reshape_shape)); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Broadcast(m::Op().Is(forty_two)) + .WithShapeEqualTo(&reshape_shape))); } // Test that ReduceWindow(Pad(op, x), y) can simplify to ReduceWindow(op, x). TEST_F(AlgebraicSimplifierTest, FoldPadIntoReduceWindow) { // TODO(b/80488902): verify this module. - auto module = HloTestBase::CreateNewModule(); + auto module = CreateNewUnverifiedModule(); HloComputation::Builder builder(TestName()); // Create operand to the pad. @@ -2816,7 +3621,7 @@ TEST_F(AlgebraicSimplifierTest, FoldPadIntoReduceWindow) { // Create the reduce-window. Window window; - for (int64 i = 0; i < ShapeUtil::Rank(pad->shape()); ++i) { + for (int64 i = 0; i < pad->shape().rank(); ++i) { auto* dim = window.add_dimensions(); dim->set_size(1); dim->set_padding_low(10); @@ -2837,8 +3642,7 @@ TEST_F(AlgebraicSimplifierTest, FoldPadIntoReduceWindow) { auto computation = module->AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root, reduce_window); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); + AlgebraicSimplifier simplifier(default_options_); ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); // Running simplification again should not result in any further changes. @@ -2846,7 +3650,8 @@ TEST_F(AlgebraicSimplifierTest, FoldPadIntoReduceWindow) { // Verify the result root = computation->root_instruction(); - EXPECT_THAT(root, op::ReduceWindow(operand, op::Constant())); + EXPECT_THAT(root, + GmockMatch(m::ReduceWindow(m::Op().Is(operand), m::Constant()))); EXPECT_TRUE(ShapeUtil::Equal(root->shape(), reduce_window_shape)) << ShapeUtil::HumanString(root->shape()) << " vs " << ShapeUtil::HumanString(reduce_window_shape); @@ -2864,7 +3669,7 @@ TEST_F(AlgebraicSimplifierTest, FoldPadIntoReduceWindow) { // ReduceWindow(Convert(op), x). TEST_F(AlgebraicSimplifierTest, FoldConvertedPadIntoReduceWindow) { // TODO(b/80488902): verify this module. - auto module = HloTestBase::CreateNewModule(); + auto module = CreateNewUnverifiedModule(); HloComputation::Builder builder(TestName()); // Create operand to the pad. @@ -2902,7 +3707,7 @@ TEST_F(AlgebraicSimplifierTest, FoldConvertedPadIntoReduceWindow) { // Create the reduce-window. Window window; - for (int64 i = 0; i < ShapeUtil::Rank(pad->shape()); ++i) { + for (int64 i = 0; i < pad->shape().rank(); ++i) { auto* dim = window.add_dimensions(); dim->set_size(1); dim->set_padding_low(10); @@ -2923,8 +3728,7 @@ TEST_F(AlgebraicSimplifierTest, FoldConvertedPadIntoReduceWindow) { auto computation = module->AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root, reduce_window); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); + AlgebraicSimplifier simplifier(default_options_); ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); // Running simplification again should not result in any further changes. @@ -2932,7 +3736,8 @@ TEST_F(AlgebraicSimplifierTest, FoldConvertedPadIntoReduceWindow) { // Verify the result root = computation->root_instruction(); - EXPECT_THAT(root, op::ReduceWindow(op::Convert(parameter), op::Constant())); + EXPECT_THAT(root, GmockMatch(m::ReduceWindow(m::Convert(m::Parameter(0)), + m::Constant()))); EXPECT_TRUE(ShapeUtil::Equal(root->shape(), reduce_window_shape)) << ShapeUtil::HumanString(root->shape()) << " vs " << ShapeUtil::HumanString(reduce_window_shape); @@ -2954,12 +3759,11 @@ TEST_F(AlgebraicSimplifierTest, ReversalOfTrivialDimensionsToBitcast) { builder.AddInstruction( HloInstruction::CreateReverse(shape, a, /*dimensions=*/{2, 3})); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(module).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(a, root); @@ -2970,6 +3774,7 @@ TEST_F(AlgebraicSimplifierTest, IteratorInvalidation) { // Dots add computations to the parent module. Test that, when the HloModule's // computations are updated, then iterator invalidation doesn't occur // when running on subsequent computations. + auto m = CreateNewVerifiedModule(); Shape r1f32 = ShapeUtil::MakeShape(F32, {1}); HloComputation::Builder builder(TestName() + ".Dot"); HloInstruction* x = @@ -2991,15 +3796,15 @@ TEST_F(AlgebraicSimplifierTest, IteratorInvalidation) { call_builder.AddInstruction( HloInstruction::CreateCall(r1f32, {zero, one}, dot_computation.get())); - module().AddEmbeddedComputation(std::move(dot_computation)); - module().AddEntryComputation(call_builder.Build()); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + m->AddEmbeddedComputation(std::move(dot_computation)); + m->AddEntryComputation(call_builder.Build()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); } // Test that a constant with tuple shape becomes a tuple of constants. TEST_F(AlgebraicSimplifierTest, ConstantTupleBecomesTupleOfConstants) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); const float constant_scalar = 7.3f; std::initializer_list constant_vector = {1.1f, 2.0f, 3.3f}; @@ -3008,73 +3813,84 @@ TEST_F(AlgebraicSimplifierTest, ConstantTupleBecomesTupleOfConstants) { Literal value = LiteralUtil::MakeTuple({&elements[0], &elements[1]}); builder.AddInstruction(HloInstruction::CreateConstant(std::move(value))); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), - op::Tuple(op::Constant(), op::Constant())); + GmockMatch(m::Tuple(m::Constant(), m::Constant()))); } // A dynamic-slice is trivial if its start indices are all zeroes and the size // of its input equals the size of its output. In this case, the dynamic slice // is equal to its input. TEST_F(AlgebraicSimplifierTest, TrivialDynamicSlice) { + auto m = CreateNewVerifiedModule(); 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 = module().AddEntryComputation(builder.Build()); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); - EXPECT_THAT(computation->root_instruction(), op::Parameter()); + auto computation = m->AddEntryComputation(builder.Build()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); + EXPECT_THAT(computation->root_instruction(), GmockMatch(m::Parameter())); } // A dynamic-update-slice is trivial if its start indices are all zeroes and the // size of its "update" equals the size of its output. In this case, the // dynamic-update-slice is equal to its update. TEST_F(AlgebraicSimplifierTest, TrivialDynamicUpdateSlice) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); 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 = module().AddEntryComputation(builder.Build()); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + auto computation = m->AddEntryComputation(builder.Build()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), - op::DynamicSlice(op::Parameter(), op::Parameter())); + GmockMatch(m::DynamicSlice(m::Parameter(), m::Parameter(), + m::Parameter(), m::Parameter()))); } // Test that two consecutive broadcasts can be merged to one. TEST_F(AlgebraicSimplifierTest, MergeBroadcasts) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); Shape r2f32 = ShapeUtil::MakeShape(F32, {2, 2}); HloInstruction* input_array = builder.AddInstruction( @@ -3085,19 +3901,19 @@ TEST_F(AlgebraicSimplifierTest, MergeBroadcasts) { builder.AddInstruction( HloInstruction::CreateBroadcast(r3f32, inner_bcast, {0, 2})); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root->opcode(), HloOpcode::kBroadcast); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); root = computation->root_instruction(); - EXPECT_THAT(root, op::Broadcast(op::Constant())); + EXPECT_THAT(root, GmockMatch(m::Broadcast(m::Constant()))); EXPECT_THAT(root->dimensions(), ElementsAre(2)); } // Test that two consecutive broadcasts can be merged to one. TEST_F(AlgebraicSimplifierTest, MergeBroadcasts2) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); Shape r2f32 = ShapeUtil::MakeShape(F32, {2, 3}); Shape r3f32 = ShapeUtil::MakeShape(F32, {2, 5, 3}); @@ -3111,19 +3927,19 @@ TEST_F(AlgebraicSimplifierTest, MergeBroadcasts2) { builder.AddInstruction( HloInstruction::CreateBroadcast(r4f32, inner_bcast, {1, 2, 3})); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root->opcode(), HloOpcode::kBroadcast); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); root = computation->root_instruction(); - EXPECT_THAT(root, op::Broadcast(op::Parameter(0))); + EXPECT_THAT(root, GmockMatch(m::Broadcast(m::Parameter(0)))); EXPECT_THAT(root->dimensions(), ElementsAre(1, 3)); } // Test that a broadcast of an iota can be merged to one iota. TEST_F(AlgebraicSimplifierTest, MergeBroadcastAndIota) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); Shape r2f32 = ShapeUtil::MakeShape(F32, {2, 2}); HloInstruction* iota = @@ -3131,19 +3947,19 @@ TEST_F(AlgebraicSimplifierTest, MergeBroadcastAndIota) { Shape r3f32 = ShapeUtil::MakeShape(F32, {2, 2, 2}); builder.AddInstruction(HloInstruction::CreateBroadcast(r3f32, iota, {0, 2})); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root->opcode(), HloOpcode::kBroadcast); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); root = computation->root_instruction(); - EXPECT_THAT(root, op::Iota()); + EXPECT_THAT(root, GmockMatch(m::Iota())); EXPECT_EQ(Cast(root)->iota_dimension(), 2); } // Test that a broadcast of an iota can be merged to one iota. TEST_F(AlgebraicSimplifierTest, MergeBroadcastAndIota2) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); Shape r3f32 = ShapeUtil::MakeShape(F32, {2, 5, 3}); HloInstruction* iota = @@ -3152,17 +3968,184 @@ TEST_F(AlgebraicSimplifierTest, MergeBroadcastAndIota2) { builder.AddInstruction( HloInstruction::CreateBroadcast(r4f32, iota, {1, 2, 3})); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root->opcode(), HloOpcode::kBroadcast); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); root = computation->root_instruction(); - EXPECT_THAT(root, op::Iota()); + EXPECT_THAT(root, GmockMatch(m::Iota())); EXPECT_EQ(Cast(root)->iota_dimension(), 2); } +TEST_F(AlgebraicSimplifierTest, SliceOfPadLow) { + const char* hlo_string = R"( + HloModule module + + ENTRY test { + param = f32[3,4] parameter(0) + constant = f32[] constant(0.0) + pad = f32[8,10] pad(f32[3,4] param, f32[] constant), padding=3_2x1_5 + ROOT slice = f32[1,1] slice(f32[8,10] pad), slice={[2:3],[0:1]} + } + )"; + 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::Reshape(m::Constant()))); +} + +TEST_F(AlgebraicSimplifierTest, SliceOfPadHigh) { + const char* hlo_string = R"( + HloModule module + + ENTRY test { + param = f32[3,4] parameter(0) + constant = f32[] constant(0.0) + pad = f32[8,10] pad(f32[3,4] param, f32[] constant), padding=3_2x1_5 + ROOT slice = f32[1,1] slice(f32[8,10] pad), slice={[6:7],[9:10]} + } + )"; + 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::Reshape(m::Constant()))); +} + +TEST_F(AlgebraicSimplifierTest, SliceOfPadMidNonScalar) { + const char* hlo_string = R"( + HloModule module + + ENTRY test { + param = f32[3,4] parameter(0) + constant = f32[] constant(0.0) + pad = f32[8,10] pad(f32[3,4] param, f32[] constant), padding=3_2x1_5 + ROOT slice = f32[1,1] slice(f32[8,10] pad), slice={[5:6],[9:10]} + } + )"; + TF_ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(hlo_string)); + + AlgebraicSimplifierOptions options; + AlgebraicSimplifier simplifier(options); + EXPECT_FALSE(simplifier.Run(module.get()).ValueOrDie()); +} + +TEST_F(AlgebraicSimplifierTest, SliceOfPadMidScalar) { + const char* hlo_string = R"( + HloModule module + + ENTRY test { + param = f32[1,1] parameter(0) + constant = f32[] constant(0.0) + pad = f32[8,10] pad(f32[1,1] param, f32[] constant), padding=3_4x4_5 + ROOT slice = f32[1,1] slice(f32[8,10] pad), slice={[3:4],[4:5]} + } + )"; + 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::Parameter())); +} + +TEST_F(AlgebraicSimplifierTest, SliceOfConcatScalarInput) { + const char* hlo_string = R"( + HloModule module + + ENTRY test { + param.0 = f32[2] parameter(0) + param.1 = f32[1] parameter(1) + param.2 = f32[3] parameter(2) + concat = f32[6] concatenate(param.0, param.1, param.2), dimensions={0} + ROOT slice = f32[1] slice(concat), slice={[2:3]} + } + )"; + 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::Parameter(1))); +} + +TEST_F(AlgebraicSimplifierTest, SliceOfConcatNonScalarInput) { + const char* hlo_string = R"( + HloModule module + + ENTRY test { + param.0 = f32[2] parameter(0) + param.1 = f32[1] parameter(1) + param.2 = f32[3] parameter(2) + concat = f32[6] concatenate(param.0, param.1, param.2), dimensions={0} + ROOT slice = f32[1] slice(concat), slice={[4:5]} + } + )"; + 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::Slice(m::Parameter(2)))); + EXPECT_EQ(root->slice_starts(0), 1); + EXPECT_EQ(root->slice_limits(0), 2); +} + +TEST_F(AlgebraicSimplifierTest, NegateNegate) { + const char* hlo_string = R"( + HloModule module + + ENTRY test { + param.0 = f32[2] parameter(0) + neg.0 = f32[2] negate(param.0) + ROOT neg.1 = f32[2] negate(neg.0) + } + )"; + 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::Parameter(0))); +} + +TEST_F(AlgebraicSimplifierTest, NotNot) { + const char* hlo_string = R"( + HloModule module + + ENTRY test { + param.0 = pred[2] parameter(0) + not.0 = pred[2] not(param.0) + ROOT not.1 = pred[2] not(not.0) + } + )"; + 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::Parameter(0))); +} + struct PadReduceWindowEffectiveBroadcastCase { std::vector input_spatials; std::vector symmetric_pad_spatials; @@ -3192,6 +4175,7 @@ class PadReduceWindowEffectiveBroadcastTest PadReduceWindowEffectiveBroadcastCase> {}; TEST_P(PadReduceWindowEffectiveBroadcastTest, DoIt) { + auto m = CreateNewVerifiedModule(); const auto& param = GetParam(); // a and b are parallel bounds we can either turn into a B F S0 S1 or @@ -3240,7 +4224,7 @@ TEST_P(PadReduceWindowEffectiveBroadcastTest, DoIt) { HloInstruction::CreateParameter(1, scalar_shape, "p1")); builder.AddInstruction( HloInstruction::CreateBinary(scalar_shape, HloOpcode::kAdd, p0, p1)); - add_computation = module().AddEmbeddedComputation(builder.Build()); + add_computation = m->AddEmbeddedComputation(builder.Build()); } Window window = window_util::MakeWindow( @@ -3254,20 +4238,19 @@ TEST_P(PadReduceWindowEffectiveBroadcastTest, DoIt) { builder.AddInstruction(HloInstruction::CreateReduceWindow( output_shape, pad, zero, window, add_computation)); - auto computation = module().AddEntryComputation(builder.Build()); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - TF_ASSERT_OK_AND_ASSIGN(bool run_successful, simplifier.Run(&module())); + auto computation = m->AddEntryComputation(builder.Build()); + AlgebraicSimplifier simplifier(default_options_); + TF_ASSERT_OK_AND_ASSIGN(bool run_successful, simplifier.Run(m.get())); ASSERT_TRUE(run_successful); EXPECT_TRUE( ShapeUtil::Equal(computation->root_instruction()->shape(), output_shape)); if (param.should_become_broadcast) { - EXPECT_THAT(computation->root_instruction(), op::Broadcast(::testing::_)); + EXPECT_THAT(computation->root_instruction(), GmockMatch(m::Broadcast())); } else { EXPECT_THAT(computation->root_instruction(), - op::ReduceWindow(::testing::_, zero)); + GmockMatch(m::ReduceWindow(m::Op(), m::Op().Is(zero)))); } } @@ -3283,9 +4266,6 @@ PadReduceWindowEffectiveBroadcastCases() { {/*input_spatials=*/{2, 2}, /*symmetric_pad_amount=*/{6, 6}, /*reduce_window_spatials=*/{7, 7}, /*prepend_a=*/true, /*should_become_broadcast=*/false}, // - {/*input_spatials=*/{1, 1}, /*symmetric_pad_amount=*/{2, 2}, - /*reduce_window_spatials=*/{5, 5}, /*prepend_a=*/true, - /*should_become_broadcast=*/true}, // {/*input_spatials=*/{1, 1}, /*symmetric_pad_amount=*/{2, 2}, /*reduce_window_spatials=*/{1, 1}, /*prepend_a=*/true, /*should_become_broadcast=*/false}, // @@ -3296,16 +4276,68 @@ PadReduceWindowEffectiveBroadcastCases() { return *cases; } -INSTANTIATE_TEST_CASE_P( +INSTANTIATE_TEST_SUITE_P( PadReduceWindowEffectiveBroadcastInstantiation, PadReduceWindowEffectiveBroadcastTest, ::testing::ValuesIn(PadReduceWindowEffectiveBroadcastCases())); +class BatchDotStrengthReductionTest + : public AlgebraicSimplifierTest, + public ::testing::WithParamInterface< + ::testing::tuple> {}; +TEST_P(BatchDotStrengthReductionTest, BatchDotStrengthReduction) { + auto module = CreateNewVerifiedModule(); + int m, k, n; + PrimitiveType element_type; + 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}); + HloComputation::Builder builder(TestName()); + + auto lhs = builder.AddInstruction( + HloInstruction::CreateParameter(0, lhs_shape, "lhs")); + auto rhs = builder.AddInstruction( + HloInstruction::CreateParameter(1, rhs_shape, "rhs")); + DotDimensionNumbers dot_dnums; + dot_dnums.add_lhs_batch_dimensions(0); + dot_dnums.add_lhs_batch_dimensions(1); + dot_dnums.add_lhs_batch_dimensions(2); + 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); + 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 computation_should_be_modified = dot_should_be_transformed; + EXPECT_EQ(changed, computation_should_be_modified); + bool has_no_dot = true; + for (const auto& hlo : computation->instructions()) { + if (hlo->opcode() == HloOpcode::kDot) { + has_no_dot = false; + break; + } + } + EXPECT_EQ(has_no_dot, dot_should_be_transformed); +} + +INSTANTIATE_TEST_SUITE_P( + BatchDotStrengthReductionTestInstantiation, BatchDotStrengthReductionTest, + ::testing::Combine(::testing::Values(1, 2), ::testing::Values(1, 2), + ::testing::Values(1, 2), ::testing::Values(F32, BF16))); + class DotStrengthReductionTest : public AlgebraicSimplifierTest, public ::testing::WithParamInterface< ::testing::tuple> {}; TEST_P(DotStrengthReductionTest, DotStrengthReduction) { + auto module = CreateNewVerifiedModule(); int m, k, n; bool transpose_lhs, transpose_rhs; PrimitiveType element_type; @@ -3335,10 +4367,9 @@ TEST_P(DotStrengthReductionTest, DotStrengthReduction) { dot_dnums.add_rhs_contracting_dimensions(0); builder.AddInstruction(HloInstruction::CreateDot( dot_shape, lhs, rhs, dot_dnums, DefaultPrecisionConfig(2))); - auto computation = module().AddEntryComputation(builder.Build()); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - TF_ASSERT_OK_AND_ASSIGN(bool changed, simplifier.Run(&module())); + 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 computation_should_be_modified = dot_should_be_transformed || (transpose_lhs && transpose_rhs); @@ -3353,7 +4384,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(), @@ -3366,7 +4397,7 @@ struct DotOfConcatTestSpec { }; class DotOfConcatSimplificationTest - : public HloVerifiedTestBase, + : public AlgebraicSimplifierTest, public ::testing::WithParamInterface {}; // Test that we transform @@ -3374,6 +4405,7 @@ class DotOfConcatSimplificationTest // to // add(dot(const_0, A), dot(const_1, B), dot(const_2, C)) TEST_P(DotOfConcatSimplificationTest, ConstantLHS) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); DotOfConcatTestSpec spec = GetParam(); @@ -3412,20 +4444,20 @@ TEST_P(DotOfConcatSimplificationTest, ConstantLHS) { builder.AddInstruction(HloInstruction::CreateDot( dot_shape, lhs, rhs, dot_dnums, DefaultPrecisionConfig(2))); - auto computation = module().AddEntryComputation(builder.Build()); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - TF_ASSERT_OK_AND_ASSIGN(bool run_successful, simplifier.Run(&module())); + auto computation = m->AddEntryComputation(builder.Build()); + AlgebraicSimplifier simplifier(default_options_); + TF_ASSERT_OK_AND_ASSIGN(bool run_successful, simplifier.Run(m.get())); ASSERT_TRUE(run_successful); EXPECT_TRUE( ShapeUtil::Equal(computation->root_instruction()->shape(), dot_shape)); - auto match_dot_0 = op::Dot(op::Slice(op::Constant()), op::Parameter(0)); - auto match_dot_1 = op::Dot(op::Slice(op::Constant()), op::Parameter(1)); - auto match_dot_2 = op::Dot(op::Slice(op::Constant()), op::Parameter(2)); - EXPECT_THAT(computation->root_instruction(), - op::Add(op::Add(match_dot_0, match_dot_1), match_dot_2)); + auto match_dot_0 = m::Dot(m::Slice(m::Constant()), m::Parameter(0)); + auto match_dot_1 = m::Dot(m::Slice(m::Constant()), m::Parameter(1)); + auto match_dot_2 = m::Dot(m::Slice(m::Constant()), m::Parameter(2)); + EXPECT_THAT( + computation->root_instruction(), + GmockMatch(m::Add(m::Add(match_dot_0, match_dot_1), match_dot_2))); } // Test that we transform @@ -3433,6 +4465,7 @@ TEST_P(DotOfConcatSimplificationTest, ConstantLHS) { // to // add(dot(A, const_0), dot(B, const_1), dot(C, const_2)) TEST_P(DotOfConcatSimplificationTest, ConstantRHS) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); DotOfConcatTestSpec spec = GetParam(); @@ -3476,21 +4509,21 @@ TEST_P(DotOfConcatSimplificationTest, ConstantRHS) { builder.AddInstruction(HloInstruction::CreateDot( dot_shape, lhs, rhs, dot_dnums, DefaultPrecisionConfig(2))); - auto computation = module().AddEntryComputation(builder.Build()); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - TF_ASSERT_OK_AND_ASSIGN(bool run_successful, simplifier.Run(&module())); + auto computation = m->AddEntryComputation(builder.Build()); + AlgebraicSimplifier simplifier(default_options_); + TF_ASSERT_OK_AND_ASSIGN(bool run_successful, simplifier.Run(m.get())); ASSERT_TRUE(run_successful); EXPECT_TRUE( ShapeUtil::Equal(computation->root_instruction()->shape(), dot_shape)); - auto match_dot_0 = op::Dot(op::Parameter(0), op::Slice(op::Constant())); - auto match_dot_1 = op::Dot(op::Parameter(1), op::Slice(op::Constant())); - auto match_dot_2 = op::Dot(op::Parameter(2), op::Slice(op::Constant())); - auto match_dot_3 = op::Dot(op::Parameter(3), op::Slice(op::Constant())); - EXPECT_THAT(computation->root_instruction(), - op::Add(op::Add(op::Add(match_dot_0, match_dot_1), match_dot_2), - match_dot_3)); + auto match_dot_0 = m::Dot(m::Parameter(0), m::Slice(m::Constant())); + auto match_dot_1 = m::Dot(m::Parameter(1), m::Slice(m::Constant())); + auto match_dot_2 = m::Dot(m::Parameter(2), m::Slice(m::Constant())); + auto match_dot_3 = m::Dot(m::Parameter(3), m::Slice(m::Constant())); + EXPECT_THAT( + computation->root_instruction(), + GmockMatch(m::Add(m::Add(m::Add(match_dot_0, match_dot_1), match_dot_2), + match_dot_3))); } DotOfConcatTestSpec kDotOfConcatTestSpecs[] = { @@ -3504,6 +4537,7 @@ DotOfConcatTestSpec kDotOfConcatTestSpecs[] = { // Test that DynamicUpdateSlice update param with any dimension equal to zero // gets removed. TEST_F(AlgebraicSimplifierTest, DynamicUpdateSliceZeroUpdate) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); const Shape dslice_shape = ShapeUtil::MakeShape(F32, {10}); HloInstruction* const operand = builder.AddInstruction( @@ -3512,21 +4546,21 @@ 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 = - module().AddEntryComputation(builder.Build()); + m->AddEntryComputation(builder.Build()); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - ASSERT_TRUE(simplifier.Run(&module()).ValueOrDie()); + AlgebraicSimplifier simplifier(default_options_); + ASSERT_TRUE(simplifier.Run(m.get()).ValueOrDie()); 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; @@ -3539,7 +4573,7 @@ struct DotOfGatherTestSpec { }; class DotOfGatherSimplificationTest - : public HloVerifiedTestBase, + : public AlgebraicSimplifierTest, public ::testing::WithParamInterface {}; // input: dot(DS(ctA), ctB)) @@ -3548,6 +4582,7 @@ class DotOfGatherSimplificationTest // output: DS(dot(ctA, ctB)) // => output dimensions: DS ({M x N}, {s, 0}, {1, N}) => {1 x N}. TEST_P(DotOfGatherSimplificationTest, ConstantRHS) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); DotOfGatherTestSpec spec = GetParam(); @@ -3567,14 +4602,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::CreateR0(start_row))), builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR1({start_row, start_col}))); + 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; @@ -3594,10 +4632,9 @@ TEST_P(DotOfGatherSimplificationTest, ConstantRHS) { builder.AddInstruction(HloInstruction::CreateDot( dot_shape, ds, rhs, dot_dnums, DefaultPrecisionConfig(2))); - auto computation = module().AddEntryComputation(builder.Build()); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - TF_ASSERT_OK_AND_ASSIGN(bool run_successful, simplifier.Run(&module())); + auto computation = m->AddEntryComputation(builder.Build()); + AlgebraicSimplifier simplifier(default_options_); + TF_ASSERT_OK_AND_ASSIGN(bool run_successful, simplifier.Run(m.get())); ASSERT_TRUE(run_successful); EXPECT_TRUE( ShapeUtil::Equal(computation->root_instruction()->shape(), dot_shape)); @@ -3607,8 +4644,8 @@ TEST_P(DotOfGatherSimplificationTest, ConstantRHS) { HloOpcode::kDynamicSlice); } else { EXPECT_THAT(computation->root_instruction(), - op::DynamicSlice(op::Dot(op::Constant(), op::Constant()), - op::Concatenate())); + GmockMatch(m::DynamicSlice(m::Dot(m::Constant(), m::Constant()), + m::Constant(), m::Constant()))); } } @@ -3618,6 +4655,7 @@ TEST_P(DotOfGatherSimplificationTest, ConstantRHS) { // output: DS(dot(ctA, ctB)) // => output dimensions: DS ({M x N}, {0, s}, {M, 1}) => {M x 1}. TEST_P(DotOfGatherSimplificationTest, ConstantLHS) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); DotOfGatherTestSpec spec = GetParam(); @@ -3645,14 +4683,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); @@ -3664,10 +4705,9 @@ TEST_P(DotOfGatherSimplificationTest, ConstantLHS) { builder.AddInstruction(HloInstruction::CreateDot( dot_shape, lhs, ds, dot_dnums, DefaultPrecisionConfig(2))); - auto computation = module().AddEntryComputation(builder.Build()); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, - non_bitcasting_callback()); - TF_ASSERT_OK_AND_ASSIGN(bool run_successful, simplifier.Run(&module())); + auto computation = m->AddEntryComputation(builder.Build()); + AlgebraicSimplifier simplifier(default_options_); + TF_ASSERT_OK_AND_ASSIGN(bool run_successful, simplifier.Run(m.get())); ASSERT_TRUE(run_successful); EXPECT_TRUE( ShapeUtil::Equal(computation->root_instruction()->shape(), dot_shape)); @@ -3677,8 +4717,8 @@ TEST_P(DotOfGatherSimplificationTest, ConstantLHS) { HloOpcode::kDynamicSlice); } else { EXPECT_THAT(computation->root_instruction(), - op::DynamicSlice(op::Dot(op::Constant(), op::Constant()), - op::Concatenate())); + GmockMatch(m::DynamicSlice(m::Dot(m::Constant(), m::Constant()), + m::Constant(), m::Constant()))); } } @@ -3726,7 +4766,7 @@ std::vector DotOfGatherPositiveNegativeTests() { return all; } -INSTANTIATE_TEST_CASE_P( +INSTANTIATE_TEST_SUITE_P( DotOfGatherSimplificationTestInstantiation, DotOfGatherSimplificationTest, ::testing::ValuesIn(DotOfGatherPositiveNegativeTests())); diff --git a/tensorflow/compiler/xla/service/allocation_tracker.cc b/tensorflow/compiler/xla/service/allocation_tracker.cc index ef5e211646e7b0b66b8e6c09948be58063422943..6cb0e985e57016e5a22fba50c3e3ad6970f1b178 100644 --- a/tensorflow/compiler/xla/service/allocation_tracker.cc +++ b/tensorflow/compiler/xla/service/allocation_tracker.cc @@ -142,13 +142,13 @@ StatusOr> AllocationTracker::DeconstructTuple( // We only need to care about replica id 0 here, since the GlobalDataHandle is // the same for all buffers across replicas. const ShapedBuffer* shaped_buffer = replicated_buffers[0]; - if (!ShapeUtil::IsTuple(shaped_buffer->on_host_shape())) { + if (!shaped_buffer->on_host_shape().IsTuple()) { return InvalidArgument("global data handle %d is not a tuple", data.handle()); } // If the on-host representation is a tuple, then the on-device one should be // as well. - TF_RET_CHECK(ShapeUtil::IsTuple(shaped_buffer->on_device_shape())); + TF_RET_CHECK(shaped_buffer->on_device_shape().IsTuple()); if (ShapeUtil::IsNestedTuple(shaped_buffer->on_device_shape())) { return Unimplemented("Deconstructing nested tuples is not implemented."); diff --git a/tensorflow/compiler/xla/service/ar_crs_combiner.cc b/tensorflow/compiler/xla/service/ar_crs_combiner.cc new file mode 100644 index 0000000000000000000000000000000000000000..f8dff6a700cc9d5843053e3d451a7b005539ca26 --- /dev/null +++ b/tensorflow/compiler/xla/service/ar_crs_combiner.cc @@ -0,0 +1,342 @@ +/* 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/ar_crs_combiner.h" + +#include +#include +#include + +#include "tensorflow/compiler/xla/literal.h" +#include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/compiler/xla/service/call_graph.h" +#include "tensorflow/compiler/xla/service/hlo_computation.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/service/hlo_opcode.h" +#include "tensorflow/compiler/xla/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 { + +namespace { + +namespace m = match; + +// 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; + } + 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) { + return c->instruction_count() == 3 && + Match(c->root_instruction(), m::Add(m::Parameter(), m::Parameter())); + }; + + if (!instruction->IsCrossModuleAllReduce() || + !computation_is_addition(instruction->called_computations()[0]) || + instruction->user_count() != 1) { + return absl::nullopt; + } + auto next = instruction->users()[0]; + while (!next->IsCrossReplicaAllReduce()) { + if (can_ar_move_past_instruction(next)) { + next = next->users()[0]; + } else { + return absl::nullopt; + } + } + if (!Cast(next)->IsNoop() && + computation_is_addition(next->called_computations()[0])) { + return absl::optional(next); + } else { + return absl::nullopt; + } +} + +} // namespace + +absl::optional ArCrsCombiner::WhileFromBodyParameter( + HloInstruction* instruction) { + CHECK_EQ(HloOpcode::kParameter, instruction->opcode()); + HloComputation* computation = instruction->parent(); + auto caller_instructions = call_graph_->GetComputationCallers(computation); + if (caller_instructions.size() == 1) { + auto caller_instruction = caller_instructions[0]; + if (caller_instruction->opcode() == HloOpcode::kWhile) { + return caller_instruction; + } + } + return absl::nullopt; +} + +std::vector ArCrsCombiner::GetAllTuples( + HloInstruction* instruction) { + if (instruction->opcode() == HloOpcode::kTuple) { + return {instruction}; + } + if (instruction->opcode() == HloOpcode::kDomain) { + return GetAllTuples(instruction->operands()[0]); + } + if (instruction->opcode() == HloOpcode::kParameter) { + auto maybe_while = WhileFromBodyParameter(instruction); + if (!maybe_while) { + return {}; + } + auto while_instr = *maybe_while; + auto init_tuples = GetAllTuples(while_instr->while_init()); + auto body_tuples = + GetAllTuples(while_instr->while_body()->root_instruction()); + if (init_tuples.empty() || body_tuples.empty()) { + return {}; + } + init_tuples.insert(init_tuples.end(), body_tuples.begin(), + body_tuples.end()); + return init_tuples; + } + if (instruction->opcode() == HloOpcode::kGetTupleElement) { + std::vector result_tuples; + for (auto tuple : GetAllTuples(instruction->operands()[0])) { + auto tmp_tuples = + GetAllTuples(tuple->mutable_operand(instruction->tuple_index())); + if (tmp_tuples.empty()) { + return {}; + } + result_tuples.insert(result_tuples.end(), tmp_tuples.begin(), + tmp_tuples.end()); + } + return result_tuples; + } + return {}; +} + +bool ArCrsCombiner::TupleElementsComputeSameValue( + HloInstruction* tuple_shaped_instruction, int64 i1, int64 i2, + absl::flat_hash_map* visited_pairs) { + auto tuples = GetAllTuples(tuple_shaped_instruction); + if (tuples.empty()) { + return false; + } + for (auto tuple : tuples) { + CHECK_EQ(tuple->opcode(), HloOpcode::kTuple); + if (!InstructionsComputeSameValue(tuple->mutable_operand(i1), + tuple->mutable_operand(i2), + visited_pairs)) { + return false; + } + } + return true; +} + +/* static */ +bool ArCrsCombiner::TestInstructionsComputeSameValue(HloInstruction* i1, + HloInstruction* i2) { + ArCrsCombiner combiner(/*num_spatial_partitions=*/2); + auto module = i1->parent()->parent(); + CHECK_EQ(module, i2->parent()->parent()); + combiner.call_graph_ = CallGraph::Build(module); + absl::flat_hash_map visited_pairs; + return combiner.InstructionsComputeSameValue(i1, i2, &visited_pairs); +} + +bool ArCrsCombiner::InstructionsComputeSameValue( + HloInstruction* i1, HloInstruction* i2, + absl::flat_hash_map* visited_pairs) { + if (i1 == i2) { + return true; + } + auto uid1 = i1->unique_id(); + auto uid2 = i2->unique_id(); + auto min_uid = std::min(uid1, uid2); + auto max_uid = std::max(uid1, uid2); + auto it = visited_pairs->find(min_uid); + if (it != visited_pairs->end() && max_uid == it->second) { + return true; + } + auto opcode1 = i1->opcode(); + auto operands1 = i1->operands(); + 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]; + auto operand2 = i2->operands()[i]; + if (!InstructionsComputeSameValue(operand1, operand2, visited_pairs)) { + return false; + } + } + if (opcode1 == HloOpcode::kParameter) { + // In the general case, we don't try to prove equality of parameters. + // We only try in the context of get-tuple-element + // (see TupleElementsComputeSameValue). + return false; + } + if (opcode1 == HloOpcode::kGetTupleElement) { + return i1->tuple_index() == i2->tuple_index() || + TupleElementsComputeSameValue(operands1[0], i1->tuple_index(), + i2->tuple_index(), visited_pairs); + } + // Don't check that the operands are identical, because Identical can + // return false for instructions that compute the same value but are not + // identical, which we don't want. We have checked the arguments with + // InstructionsComputeSameValue earlier. + auto eq_instructions = [](const HloInstruction* i1, + const HloInstruction* i2) -> bool { return true; }; + return i1->Identical(*i2, eq_instructions, eq_computations, + /*layout_sensitive=*/false); +} + +void ArCrsCombiner::GroupAllReducesById(HloModule* module) { + for (HloComputation* computation : module->MakeNonfusionComputations()) { + for (HloInstruction* instruction : computation->instructions()) { + 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; + } + } + } + } +} + +void ArCrsCombiner::KeepProvablyEqualInstructionGroups() { + for (auto it : all_reduce_map_) { + auto all_reduce_id = it.first; + auto instruction_vec = it.second; + CHECK_EQ(instruction_vec.size(), num_spatial_partitions_); + auto instr_0 = instruction_vec[0]; + for (int i = 1; i < instruction_vec.size(); ++i) { + auto instr_i = instruction_vec[i]; + auto next_0 = instr_0->users()[0]; + auto next_i = instr_i->users()[0]; + absl::flat_hash_map visited_pairs; + 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]; + } + } + } +} + +StatusOr ArCrsCombiner::RewriteGraph() { + if (all_reduce_map_.empty()) { + return false; + } + for (auto it : all_reduce_map_) { + auto instruction_vec = it.second; + for (auto all_reduce : instruction_vec) { + auto parent_computation = all_reduce->parent(); + auto all_reduce_id = all_reduce->all_reduce_id(); + auto prev = all_reduce->mutable_operand(0); + auto next = all_reduce->users()[0]; + TF_CHECK_OK(all_reduce->ReplaceUseWith(next, prev)); + TF_CHECK_OK(parent_computation->RemoveInstruction(all_reduce)); + while (!next->IsCrossReplicaAllReduce()) { + switch (next->opcode()) { + case HloOpcode::kBitcast: + case HloOpcode::kTranspose: + case HloOpcode::kReshape: + case HloOpcode::kConvert: + case HloOpcode::kMultiply: + break; + case HloOpcode::kAdd: + case HloOpcode::kSubtract: { + auto other_operand = (next->operands()[0] == prev) + ? next->operands()[1] + : next->operands()[0]; + // To move the AR past the addition/subtraction, we need to divide + // other_operand by the number of spatial partitions. + auto shape = other_operand->shape(); + Literal lit(shape); + lit.PopulateWithValue(num_spatial_partitions_); + auto divisor = parent_computation->AddInstruction( + HloInstruction::CreateConstant(lit.Clone())); + auto division = + parent_computation->AddInstruction(HloInstruction::CreateBinary( + shape, HloOpcode::kDivide, other_operand, divisor)); + TF_CHECK_OK(other_operand->ReplaceUseWith(next, division)); + break; + } + default: + LOG(FATAL) << "Unexpected instruction: " << next->ToShortString(); + } + prev = next; + next = next->users()[0]; + } + // The AllReduce and the CRS are combined to an all-core AllReduce. + next->set_all_reduce_id(all_reduce_id); + } + } + return true; +} + +StatusOr ArCrsCombiner::Run(HloModule* module) { + call_graph_ = CallGraph::Build(module); + + GroupAllReducesById(module); + + KeepProvablyEqualInstructionGroups(); + + return RewriteGraph(); +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/ar_crs_combiner.h b/tensorflow/compiler/xla/service/ar_crs_combiner.h new file mode 100644 index 0000000000000000000000000000000000000000..e61ef5d4f9072979a6c356a9456c91e19405b01e --- /dev/null +++ b/tensorflow/compiler/xla/service/ar_crs_combiner.h @@ -0,0 +1,96 @@ +/* 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_AR_CRS_COMBINER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_AR_CRS_COMBINER_H_ + +#include "absl/container/flat_hash_map.h" +#include "absl/strings/string_view.h" +#include "tensorflow/compiler/xla/service/call_graph.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 { + +// When the HLO graph contains a cross-module AllReduce, followed by some simple +// linear operations, followed by a cross-replica AllReduce, we can combine the +// CMAR and the CRAR, to use an efficient AllReduce implementation that fully +// utilizes the interconnect bandwidth. +// Such sequences appear in spatially partitioned models. +// This pass must run right after spatial partitioning. +class ArCrsCombiner : public HloModulePass { + public: + ArCrsCombiner(int num_spatial_partitions) + : num_spatial_partitions_(num_spatial_partitions) {} + absl::string_view name() const override { return "ar-crs-combiner"; } + StatusOr Run(HloModule* module) override; + + // Helper method to allow testing of InstructionsComputeSameValue. + static bool TestInstructionsComputeSameValue(HloInstruction* i1, + HloInstruction* i2); + + private: + // If the passed instruction is a while parameter, and the while body is only + // called by a single while instruction, return the while instruction. + absl::optional WhileFromBodyParameter( + HloInstruction* instruction); + + // Returns a vector of tuple instructions. + // If all instructions that flow to "instruction" are tuples, return them. + // Otherwise, return an empty vector. + std::vector GetAllTuples(HloInstruction* instruction); + + // Checks whether two different elements in the same tuple compute the same + // value. + bool TupleElementsComputeSameValue( + HloInstruction* tuple_shaped_instruction, int64 i1, int64 i2, + absl::flat_hash_map* visited_pairs); + + // Returns whether the instructions i1 and i2 can be shown to evaluate to the + // same value. Handling WHILE requires recursion, which may cause us to visit + // the same instruction again. To avoid infinite loops, we pass a cache of + // visited instruction pairs. + bool InstructionsComputeSameValue( + HloInstruction* i1, HloInstruction* i2, + absl::flat_hash_map* visited_pairs); + + // Populates all_reduce_map_. + void GroupAllReducesById(HloModule* module); + + // Looks at each AllReduce group in all_reduce_map_, and keeps only the + // groups for which it's safe to move the AllReduce later in the HLO graph. + void KeepProvablyEqualInstructionGroups(); + + // Performs the graph rewrite that eliminates the early AllReduce and turns + // the later CRS into an AllReduce. + StatusOr RewriteGraph(); + + int num_spatial_partitions_; + + // 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_; +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_AR_CRS_COMBINER_H_ diff --git a/tensorflow/compiler/xla/service/ar_crs_combiner_test.cc b/tensorflow/compiler/xla/service/ar_crs_combiner_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..5152f0dc884a153f9b0ade06acd479832d87ff25 --- /dev/null +++ b/tensorflow/compiler/xla/service/ar_crs_combiner_test.cc @@ -0,0 +1,1175 @@ +/* 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/ar_crs_combiner.h" +#include "tensorflow/compiler/xla/service/hlo_matchers.h" +#include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" +#include "tensorflow/core/lib/core/status_test_util.h" + +namespace xla { +namespace { + +namespace op = xla::testing::opcode_matchers; + +class ArCrsCombinerTest : public HloTestBase {}; + +TEST_F(ArCrsCombinerTest, SameValueTestBasecase) { + const char* module_str = R"( +HloModule foobar + +ENTRY %entrycomp (p: f32[2,2]) -> (f32[2,2], f32[2,2]) { + %p = f32[2,2] parameter(0) + %constant.f32.1 = f32[2,2] constant({{1, 2}, {3, 4}}) + %constant.f32.2 = f32[2,2] constant({{1, 2}, {3, 4}}) + ROOT %tuple = (f32[2,2], f32[2,2]) tuple(%constant.f32.1, %constant.f32.2) +} +)"; + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(module_str)); + auto root_tuple = module->entry_computation()->root_instruction(); + auto i1 = root_tuple->operands()[0]; + auto i2 = root_tuple->operands()[1]; + EXPECT_FALSE(ArCrsCombiner::TestInstructionsComputeSameValue( + i1, module->entry_computation()->parameter_instruction(0))); + EXPECT_TRUE(ArCrsCombiner::TestInstructionsComputeSameValue(i1, i2)); +} + +TEST_F(ArCrsCombinerTest, SameValueTestBasecase2) { + const char* module_str = R"( +HloModule foobar + +ENTRY %entrycomp (x: f32[]) -> (f32[], f32[]) { + %x = f32[] parameter(0) + ROOT %tuple = (f32[], f32[]) tuple(%x, %x) +} +)"; + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(module_str)); + auto root_tuple = module->entry_computation()->root_instruction(); + auto i1 = root_tuple->operands()[0]; + auto i2 = root_tuple->operands()[1]; + EXPECT_TRUE(ArCrsCombiner::TestInstructionsComputeSameValue(i1, i2)); +} + +TEST_F(ArCrsCombinerTest, SameValueTestBasecase3) { + const char* module_str = R"( +HloModule foobar + +ENTRY %entrycomp (x: f32[], y: f32[]) -> (f32[], f32[]) { + %x = f32[] parameter(0) + %y = f32[] parameter(1) + ROOT %tuple = (f32[], f32[]) tuple(%x, %y) +} +)"; + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(module_str)); + auto root_tuple = module->entry_computation()->root_instruction(); + auto i1 = root_tuple->operands()[0]; + auto i2 = root_tuple->operands()[1]; + EXPECT_FALSE(ArCrsCombiner::TestInstructionsComputeSameValue(i1, i2)); +} + +TEST_F(ArCrsCombinerTest, SameValueTestNumOperands) { + const char* module_str = R"( +HloModule foobar + +ENTRY %entrycomp (p: f32[2,2]) -> ((f32[2,2]), (f32[2,2], f32[2,2])) { + %p = f32[2,2] parameter(0) + %constant.f32 = f32[2,2] constant({{1, 2}, {3, 4}}) + %tuple1 = (f32[2,2]) tuple(%constant.f32) + %tuple2 = (f32[2,2], f32[2,2]) tuple(%constant.f32, %constant.f32) + ROOT %tuple = ((f32[2,2]), (f32[2,2], f32[2,2])) tuple(%tuple1, %tuple2) +} +)"; + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(module_str)); + auto root_tuple = module->entry_computation()->root_instruction(); + auto i1 = root_tuple->operands()[0]; + auto i2 = root_tuple->operands()[1]; + EXPECT_FALSE(ArCrsCombiner::TestInstructionsComputeSameValue(i1, i2)); +} + +TEST_F(ArCrsCombinerTest, SameValueTestSliceIndicesMatch) { + const char* module_str = R"( +HloModule foobar + +ENTRY %entrycomp (p: f32[2]) -> (f32[1], f32[1]) { + %p = f32[2] parameter(0) + %slice.1 = f32[1] slice(f32[2] %p), slice={[0:1]} + %slice.2 = f32[1] slice(f32[2] %p), slice={[0:1]} + ROOT %tuple = (f32[1], f32[1]) tuple(%slice.1, %slice.2) +} +)"; + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(module_str)); + auto root_tuple = module->entry_computation()->root_instruction(); + auto i1 = root_tuple->operands()[0]; + auto i2 = root_tuple->operands()[1]; + EXPECT_TRUE(ArCrsCombiner::TestInstructionsComputeSameValue(i1, i2)); +} + +TEST_F(ArCrsCombinerTest, SameValueTestSliceIndicesDontMatch) { + const char* module_str = R"( +HloModule foobar + +ENTRY %entrycomp (p: f32[2]) -> (f32[1], f32[1]) { + %p = f32[2] parameter(0) + %slice.1 = f32[1] slice(f32[2] %p), slice={[0:1]} + %slice.2 = f32[1] slice(f32[2] %p), slice={[1:2]} + ROOT %tuple = (f32[1], f32[1]) tuple(%slice.1, %slice.2) +} +)"; + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(module_str)); + auto root_tuple = module->entry_computation()->root_instruction(); + auto i1 = root_tuple->operands()[0]; + auto i2 = root_tuple->operands()[1]; + EXPECT_FALSE(ArCrsCombiner::TestInstructionsComputeSameValue(i1, i2)); +} + +TEST_F(ArCrsCombinerTest, SameValueTestTupleElementSameIndex) { + const char* module_str = R"( +HloModule foobar + +ENTRY %entrycomp (p: f32[2,2]) -> (f32[2,2], f32[2,2]) { + %p = f32[2,2] parameter(0) + %constant.f32 = f32[2,2] constant({{1, 2}, {3, 4}}) + %tuple.1 = (f32[2,2], f32[2,2]) tuple(%constant.f32, %constant.f32) + %get-tuple-element.1 = f32[2,2] get-tuple-element(%tuple.1), index=0 + %get-tuple-element.2 = f32[2,2] get-tuple-element(%tuple.1), index=0 + ROOT %tuple = (f32[2,2], f32[2,2]) tuple(%get-tuple-element.1, %get-tuple-element.2) +} +)"; + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(module_str)); + auto root_tuple = module->entry_computation()->root_instruction(); + auto i1 = root_tuple->operands()[0]; + auto i2 = root_tuple->operands()[1]; + EXPECT_TRUE(ArCrsCombiner::TestInstructionsComputeSameValue(i1, i2)); +} + +TEST_F(ArCrsCombinerTest, SameValueTestTupleElementDifferentIndex1) { + const char* module_str = R"( +HloModule foobar + +ENTRY %entrycomp (p: f32[2,2]) -> (f32[2,2], f32[2,2]) { + %p = f32[2,2] parameter(0) + %constant.f32 = f32[2,2] constant({{1, 2}, {3, 4}}) + %tuple.1 = (f32[2,2], f32[2,2]) tuple(%constant.f32, %constant.f32) + %get-tuple-element.1 = f32[2,2] get-tuple-element(%tuple.1), index=0 + %get-tuple-element.2 = f32[2,2] get-tuple-element(%tuple.1), index=1 + ROOT %tuple = (f32[2,2], f32[2,2]) tuple(%get-tuple-element.1, %get-tuple-element.2) +} +)"; + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(module_str)); + auto root_tuple = module->entry_computation()->root_instruction(); + auto i1 = root_tuple->operands()[0]; + auto i2 = root_tuple->operands()[1]; + EXPECT_TRUE(ArCrsCombiner::TestInstructionsComputeSameValue(i1, i2)); +} + +TEST_F(ArCrsCombinerTest, SameValueTestTupleElementDifferentIndex2) { + const char* module_str = R"( +HloModule foobar + +ENTRY %entrycomp (p: f32[2,2]) -> (f32[2,2], f32[2,2]) { + %p = f32[2,2] parameter(0) + %constant.f32.1 = f32[2,2] constant({{1, 2}, {3, 4}}) + %constant.f32.2 = f32[2,2] constant({{2, 3}, {4, 5}}) + %tuple.1 = (f32[2,2], f32[2,2]) tuple(%constant.f32.1, %constant.f32.2) + %get-tuple-element.1 = f32[2,2] get-tuple-element(%tuple.1), index=0 + %get-tuple-element.2 = f32[2,2] get-tuple-element(%tuple.1), index=1 + ROOT %tuple = (f32[2,2], f32[2,2]) tuple(%get-tuple-element.1, %get-tuple-element.2) +} +)"; + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(module_str)); + auto root_tuple = module->entry_computation()->root_instruction(); + auto i1 = root_tuple->operands()[0]; + auto i2 = root_tuple->operands()[1]; + EXPECT_FALSE(ArCrsCombiner::TestInstructionsComputeSameValue(i1, i2)); +} + +TEST_F(ArCrsCombinerTest, SameValueTestWhile1) { + const char* module_str = R"( +HloModule foobar + +%condition (x: (f32[2,2], f32[2,2])) -> pred[] { + %x = (f32[2,2], f32[2,2]) parameter(0) + %constant.0 = s32[] constant(0) + %constant.1 = s32[] constant(1) + ROOT %greater-than = pred[] greater-than(s32[] %constant.1, s32[] %constant.0) +} + +%body (x: (f32[2,2], f32[2,2])) -> (f32[2,2], f32[2,2]) { + %x = (f32[2,2], f32[2,2]) parameter(0) + %constant.f32 = f32[2,2] constant({{1, 2}, {3, 4}}) + %get-tuple-element.1 = f32[2,2] get-tuple-element(%x), index=0 + %get-tuple-element.2 = f32[2,2] get-tuple-element(%x), index=1 + %add.1 = f32[2,2] add(%get-tuple-element.1, %constant.f32) + %add.2 = f32[2,2] add(%get-tuple-element.2, %constant.f32) + ROOT %tuple = (f32[2,2], f32[2,2]) tuple(%add.1, %add.2) +} + +ENTRY %WhileLoop () -> (f32[2,2], f32[2,2]) { + %constant.f32 = f32[2,2] constant({{3, 4}, {5, 6}}) + %init.tuple = (f32[2,2], f32[2,2]) tuple(%constant.f32, %constant.f32) + ROOT %while = (f32[2,2], f32[2,2]) while(%init.tuple), condition=%condition, body=%body +} +)"; + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(module_str)); + auto root_while = module->entry_computation()->root_instruction(); + auto body_tuple = root_while->while_body()->root_instruction(); + auto i1 = body_tuple->operands()[0]; + auto i2 = body_tuple->operands()[1]; + EXPECT_TRUE(ArCrsCombiner::TestInstructionsComputeSameValue(i1, i2)); +} + +TEST_F(ArCrsCombinerTest, SameValueTestWhile2) { + const char* module_str = R"( +HloModule foobar + +%condition (x: (f32[2,2], f32[2,2])) -> pred[] { + %x = (f32[2,2], f32[2,2]) parameter(0) + %constant.0 = s32[] constant(0) + %constant.1 = s32[] constant(1) + ROOT %greater-than = pred[] greater-than(s32[] %constant.1, s32[] %constant.0) +} + +%body (x: (f32[2,2], f32[2,2])) -> (f32[2,2], f32[2,2]) { + %x = (f32[2,2], f32[2,2]) parameter(0) + %constant.f32 = f32[2,2] constant({{1, 2}, {3, 4}}) + %get-tuple-element.1 = f32[2,2] get-tuple-element(%x), index=0 + %get-tuple-element.2 = f32[2,2] get-tuple-element(%x), index=1 + %add.1 = f32[2,2] add(%get-tuple-element.1, %constant.f32) + %add.2 = f32[2,2] add(%get-tuple-element.2, %constant.f32) + ROOT %tuple = (f32[2,2], f32[2,2]) tuple(%add.1, %add.2) +} + +ENTRY %WhileLoop () -> (f32[2,2], f32[2,2]) { + %constant.f32.1 = f32[2,2] constant({{3, 4}, {5, 6}}) + %constant.f32.2 = f32[2,2] constant({{3, 4}, {7, 8}}) + %init.tuple = (f32[2,2], f32[2,2]) tuple(%constant.f32.1, %constant.f32.2) + ROOT %while = (f32[2,2], f32[2,2]) while(%init.tuple), condition=%condition, body=%body +} +)"; + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(module_str)); + auto root_while = module->entry_computation()->root_instruction(); + auto body_tuple = root_while->while_body()->root_instruction(); + auto i1 = body_tuple->operands()[0]; + auto i2 = body_tuple->operands()[1]; + EXPECT_FALSE(ArCrsCombiner::TestInstructionsComputeSameValue(i1, i2)); +} + +TEST_F(ArCrsCombinerTest, SameValueTestWhile3) { + const char* module_str = R"( +HloModule foobar + +%condition (x: (f32[2,2], f32[2,2])) -> pred[] { + %x = (f32[2,2], f32[2,2]) parameter(0) + %constant.0 = s32[] constant(0) + %constant.1 = s32[] constant(1) + ROOT %greater-than = pred[] greater-than(s32[] %constant.1, s32[] %constant.0) +} + +%body (x: (f32[2,2], f32[2,2])) -> (f32[2,2], f32[2,2]) { + %x = (f32[2,2], f32[2,2]) parameter(0) + %constant.f32.1 = f32[2,2] constant({{1, 2}, {3, 4}}) + %constant.f32.2 = f32[2,2] constant({{3, 4}, {1, 2}}) + %get-tuple-element.1 = f32[2,2] get-tuple-element(%x), index=0 + %get-tuple-element.2 = f32[2,2] get-tuple-element(%x), index=1 + %add.1 = f32[2,2] add(%get-tuple-element.1, %constant.f32.1) + %add.2 = f32[2,2] add(%get-tuple-element.2, %constant.f32.2) + ROOT %tuple = (f32[2,2], f32[2,2]) tuple(%add.1, %add.2) +} + +ENTRY %WhileLoop () -> (f32[2,2], f32[2,2]) { + %constant.f32 = f32[2,2] constant({{3, 4}, {5, 6}}) + %init.tuple = (f32[2,2], f32[2,2]) tuple(%constant.f32, %constant.f32) + ROOT %while = (f32[2,2], f32[2,2]) while(%init.tuple), condition=%condition, body=%body +} +)"; + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(module_str)); + auto root_while = module->entry_computation()->root_instruction(); + auto body_tuple = root_while->while_body()->root_instruction(); + auto i1 = body_tuple->operands()[0]->operands()[0]; // %get-tuple-element.1 + auto i2 = body_tuple->operands()[1]->operands()[0]; // %get-tuple-element.2 + EXPECT_FALSE(ArCrsCombiner::TestInstructionsComputeSameValue(i1, i2)); +} + +void CompareReplicaGroups(const std::vector& groups_before, + const std::vector& groups_after) { + ASSERT_EQ(groups_before.size(), groups_after.size()); + for (int i = 0; i < groups_before.size(); ++i) { + // Somewhat verbose way to compare the replica_ids, because EqualsProto + // is not available in the open-source build. + auto group_before = groups_before[i]; + std::vector ids_before(group_before.replica_ids().begin(), + group_before.replica_ids().end()); + auto group_after = groups_after[i]; + std::vector ids_after(group_after.replica_ids().begin(), + group_after.replica_ids().end()); + EXPECT_EQ(ids_before, ids_after); + } +} + +TEST_F(ArCrsCombinerTest, RewriteArConvertCrs) { + 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},{1}}, + 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,1}}, + to_apply=%sum.f32, + sharding={maximal device=0} + + %all-reduce.ar.2 = bf16[] + all-reduce(%constant.bf16), + replica_groups={{0},{1}}, + 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,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::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(); + CompareReplicaGroups(replica_groups_before, replica_groups_after); +} + +TEST_F(ArCrsCombinerTest, RewriteArBitcastCrs) { + const char* module_str = R"( +HloModule foobar + +%sum.1 (a: f32[2,1], b: f32[2,1]) -> f32[2,1] { + %a = f32[2,1] parameter(0) + %b = f32[2,1] parameter(1) + ROOT %add = f32[2,1] add(%a, %b) +} + +%sum.2 (x: f32[2], y: f32[2]) -> f32[2] { + %x = f32[2] parameter(0) + %y = f32[2] parameter(1) + ROOT %add = f32[2] add(%x, %y) +} + +ENTRY %entrycomp (p: f32[2,1]) -> (f32[2], f32[2]) { + %p = f32[2,1] parameter(0) + + %all-reduce.ar.1 = f32[2,1] + all-reduce(%p), + replica_groups={{0},{1}}, + all_reduce_id=1, + to_apply=%sum.1, + sharding={maximal device=0} + %bitcast.1 = f32[2]{0} bitcast(f32[2,1]{1,0} %all-reduce.ar.1) + %all-reduce.1 = f32[2] + all-reduce(%bitcast.1), + replica_groups={{0,1}}, + to_apply=%sum.2, + sharding={maximal device=0} + + %all-reduce.ar.2 = f32[2,1] + all-reduce(%p), + replica_groups={{0},{1}}, + all_reduce_id=1, + to_apply=%sum.1, + sharding={maximal device=1} + %bitcast.2 = f32[2]{0} bitcast(f32[2,1]{1,0} %all-reduce.ar.2) + %all-reduce.2 = f32[2] + all-reduce(%bitcast.2), + replica_groups={{0,1}}, + to_apply=%sum.2, + 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::Bitcast(op::Parameter())), + op::AllReduce(op::Bitcast(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, RewriteArMultiplyCrs) { + 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} + %multiply.1 = f32[] + multiply(%all-reduce.ar.1, %constant.f32), + sharding={maximal device=0} + %all-reduce.1 = f32[] + all-reduce(%multiply.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} + %multiply.2 = f32[] + multiply(%all-reduce.ar.2, %constant.f32), + sharding={maximal device=1} + %all-reduce.2 = f32[] + all-reduce(%multiply.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::Multiply(op::Parameter(), op::Constant())), + op::AllReduce(op::Multiply(op::Parameter(), 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, RewriteArConvertAddCrs) { + 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: f32[]) -> (f32[], f32[]) { + %p = f32[] parameter(0) + %constant.bf16 = bf16[] constant(1) + %constant.f32 = f32[] constant(2) + + %all-reduce.ar.1 = bf16[] + all-reduce(%constant.bf16), + replica_groups={{0},{1}}, + all_reduce_id=1, + to_apply=%sum.bf16, + sharding={maximal device=0} + %convert.1 = f32[] + convert(%all-reduce.ar.1), + sharding={maximal device=0} + %add.1 = f32[] + add(%constant.f32, %convert.1), + sharding={maximal device=0} + %all-reduce.1 = f32[] + all-reduce(%add.1), + replica_groups={{0,1}}, + to_apply=%sum.f32, + sharding={maximal device=0} + + %all-reduce.ar.2 = bf16[] + all-reduce(%constant.bf16), + replica_groups={{0},{1}}, + all_reduce_id=1, + to_apply=%sum.bf16, + sharding={maximal device=1} + %convert.2 = f32[] + convert(%all-reduce.ar.2), + sharding={maximal device=1} + %add.2 = f32[] + add(%constant.f32, %convert.2), + sharding={maximal device=1} + %all-reduce.2 = f32[] + all-reduce(%add.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::Add(op::Divide(op::Constant(), op::Constant()), + op::Convert())), + op::AllReduce(op::Add(op::Divide(op::Constant(), op::Constant()), + op::Convert())))); + 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, OtherSummandNotTheSameDontRewrite) { + 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: f32[]) -> (f32[], f32[]) { + %p = f32[] parameter(0) + %constant.bf16 = bf16[] constant(1) + %constant.f32.1 = f32[] constant(2) + %constant.f32.2 = f32[] constant(3) + + %all-reduce.ar.1 = bf16[] + all-reduce(%constant.bf16), + replica_groups={{0},{1}}, + all_reduce_id=1, + to_apply=%sum.bf16, + sharding={maximal device=0} + %convert.1 = f32[] + convert(%all-reduce.ar.1), + sharding={maximal device=0} + %add.1 = f32[] + add(%constant.f32.1, %convert.1), + sharding={maximal device=0} + %all-reduce.1 = f32[] + all-reduce(%add.1), + replica_groups={{0,1}}, + to_apply=%sum.f32, + sharding={maximal device=0} + + %all-reduce.ar.2 = bf16[] + all-reduce(%constant.bf16), + replica_groups={{0},{1}}, + all_reduce_id=1, + to_apply=%sum.bf16, + sharding={maximal device=1} + %convert.2 = f32[] + convert(%all-reduce.ar.2), + sharding={maximal device=1} + %add.2 = f32[] + add(%constant.f32.2, %convert.2), + sharding={maximal device=1} + %all-reduce.2 = f32[] + all-reduce(%add.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)); + ArCrsCombiner combiner(2); + auto changed = combiner.Run(module.get()).ValueOrDie(); + 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 5c180cbdd492031e133b81149f0f4698619b7788..215e8ced4bb3f98a26ac4eb9912a7fd4d917852f 100644 --- a/tensorflow/compiler/xla/service/backend.cc +++ b/tensorflow/compiler/xla/service/backend.cc @@ -57,6 +57,16 @@ int BackendOptions::intra_op_parallelism_threads() const { return intra_op_parallelism_threads_; } +BackendOptions& BackendOptions::set_allowed_devices( + const absl::optional>& allowed_devices) { + allowed_devices_ = allowed_devices; + return *this; +} + +const absl::optional>& BackendOptions::allowed_devices() const { + return allowed_devices_; +} + // Define this in .cc file to avoid having to include eigen or forward declare // these types in the header. struct Backend::EigenThreadPoolWrapper { @@ -76,8 +86,9 @@ struct Backend::EigenThreadPoolWrapper { const BackendOptions& options) { se::Platform* platform = options.platform(); TF_ASSIGN_OR_RETURN(auto compiler, Compiler::GetForPlatform(platform)); - TF_ASSIGN_OR_RETURN(auto stream_executors, - PlatformUtil::GetStreamExecutors(platform)); + TF_ASSIGN_OR_RETURN( + auto stream_executors, + PlatformUtil::GetStreamExecutors(platform, options.allowed_devices())); TF_ASSIGN_OR_RETURN(auto transfer_manager, TransferManager::GetForPlatform(platform)); TF_ASSIGN_OR_RETURN(auto computation_placer, @@ -104,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 a2dafbe803f8bd5f23e4e9f3f6d3e6f744c9fab9..c35f033dc0180409ae3888c2050021da83f5c72a 100644 --- a/tensorflow/compiler/xla/service/backend.h +++ b/tensorflow/compiler/xla/service/backend.h @@ -18,9 +18,11 @@ limitations under the License. #include #include +#include #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" @@ -53,9 +55,16 @@ class BackendOptions { BackendOptions& set_intra_op_parallelism_threads(int num_threads); int intra_op_parallelism_threads() const; + // Sets the allowed_devices for selectively constructing stream executors + // on the platform. + BackendOptions& set_allowed_devices( + const absl::optional>& allowed_devices); + const absl::optional>& allowed_devices() const; + private: se::Platform* platform_ = nullptr; int intra_op_parallelism_threads_ = -1; + absl::optional> allowed_devices_; }; // Class which encapsulates an XLA backend. It includes everything necessary @@ -167,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/batch_dot_simplification_test.cc b/tensorflow/compiler/xla/service/batch_dot_simplification_test.cc index 38f1a5d3a645f98220ec445bb9bbdf2b9b842109..52ec1a794c5e9f4452a4bf2b648f453d8acfe976 100644 --- a/tensorflow/compiler/xla/service/batch_dot_simplification_test.cc +++ b/tensorflow/compiler/xla/service/batch_dot_simplification_test.cc @@ -17,14 +17,13 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_matchers.h" #include "tensorflow/compiler/xla/test.h" #include "tensorflow/compiler/xla/tests/hlo_test_base.h" -#include "tensorflow/compiler/xla/tests/hlo_verified_test_base.h" namespace xla { namespace { namespace op = xla::testing::opcode_matchers; -class BatchDotSimplificationTest : public HloVerifiedTestBase {}; +class BatchDotSimplificationTest : public HloTestBase {}; TEST_F(BatchDotSimplificationTest, ElideSingleDegenerateBatchDotDim_VectorVector) { @@ -38,11 +37,12 @@ main { } )"; - ParseAndVerifyModule(hlo_text); + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr m, + ParseAndReturnVerifiedModule(hlo_text)); BatchDotSimplification pass; - ASSERT_TRUE(pass.Run(&module()).ValueOrDie()); + ASSERT_TRUE(pass.Run(m.get()).ValueOrDie()); - HloInstruction* root = module().entry_computation()->root_instruction(); + HloInstruction* root = m->entry_computation()->root_instruction(); EXPECT_THAT(root, op::Reshape(op::Dot( op::Reshape(op::Parameter(0)), op::Reshape(op::Parameter(1)), @@ -61,11 +61,12 @@ main { } )"; - ParseAndVerifyModule(hlo_text); + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr m, + ParseAndReturnVerifiedModule(hlo_text)); BatchDotSimplification pass; - ASSERT_TRUE(pass.Run(&module()).ValueOrDie()); + ASSERT_TRUE(pass.Run(m.get()).ValueOrDie()); - HloInstruction* root = module().entry_computation()->root_instruction(); + HloInstruction* root = m->entry_computation()->root_instruction(); EXPECT_THAT(root, op::Reshape(op::Dot( op::Reshape(op::Parameter(0)), op::Reshape(op::Parameter(1)), @@ -84,11 +85,12 @@ main { } )"; - ParseAndVerifyModule(hlo_text); + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr m, + ParseAndReturnVerifiedModule(hlo_text)); BatchDotSimplification pass; - ASSERT_TRUE(pass.Run(&module()).ValueOrDie()); + ASSERT_TRUE(pass.Run(m.get()).ValueOrDie()); - HloInstruction* root = module().entry_computation()->root_instruction(); + HloInstruction* root = m->entry_computation()->root_instruction(); EXPECT_THAT(root, op::Reshape(op::Dot( op::Reshape(op::Parameter(0)), op::Reshape(op::Parameter(1)), @@ -107,11 +109,12 @@ main { } )"; - ParseAndVerifyModule(hlo_text); + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr m, + ParseAndReturnVerifiedModule(hlo_text)); BatchDotSimplification pass; - ASSERT_TRUE(pass.Run(&module()).ValueOrDie()); + ASSERT_TRUE(pass.Run(m.get()).ValueOrDie()); - HloInstruction* root = module().entry_computation()->root_instruction(); + HloInstruction* root = m->entry_computation()->root_instruction(); EXPECT_THAT(root, op::Reshape(op::Dot( op::Reshape(op::Parameter(0)), op::Reshape(op::Parameter(1)), @@ -130,11 +133,12 @@ main { } )"; - ParseAndVerifyModule(hlo_text); + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr m, + ParseAndReturnVerifiedModule(hlo_text)); BatchDotSimplification pass; - ASSERT_TRUE(pass.Run(&module()).ValueOrDie()); + ASSERT_TRUE(pass.Run(m.get()).ValueOrDie()); - HloInstruction* root = module().entry_computation()->root_instruction(); + HloInstruction* root = m->entry_computation()->root_instruction(); EXPECT_THAT(root, op::Reshape(op::Dot( op::Reshape(op::Parameter(0)), op::Reshape(op::Parameter(1)), @@ -153,11 +157,12 @@ main { } )"; - ParseAndVerifyModule(hlo_text); + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr m, + ParseAndReturnVerifiedModule(hlo_text)); BatchDotSimplification pass; - ASSERT_TRUE(pass.Run(&module()).ValueOrDie()); + ASSERT_TRUE(pass.Run(m.get()).ValueOrDie()); - HloInstruction* root = module().entry_computation()->root_instruction(); + HloInstruction* root = m->entry_computation()->root_instruction(); EXPECT_THAT(root, op::Reshape(op::Dot( op::Reshape(op::Parameter(0)), op::Reshape(op::Parameter(1)), diff --git a/tensorflow/compiler/xla/service/batchnorm_expander.cc b/tensorflow/compiler/xla/service/batchnorm_expander.cc index f70f6ddfec69c0113a1afe2073a2392098f49456..e5f5c3edb2ac0c217317fbf809463aa31af9af59 100644 --- a/tensorflow/compiler/xla/service/batchnorm_expander.cc +++ b/tensorflow/compiler/xla/service/batchnorm_expander.cc @@ -107,19 +107,37 @@ class BatchNormExpanderVisitor : public DfsHloVisitorWithDefault { } std::unique_ptr Mean( - int64 element_count, HloInstruction* operand, + HloInstruction* element_count, HloInstruction* operand, const std::function)>& add_instruction) { - HloInstruction* elem_count_recip = - add_instruction(HloInstruction::CreateBroadcast( - operand->shape(), - add_instruction(HloInstruction::CreateConvert( - ShapeUtil::MakeShape(operand->shape().element_type(), {}), - add_instruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR0(1.0 / element_count))))), - {})); - return HloInstruction::CreateBinary(operand->shape(), HloOpcode::kMultiply, - operand, elem_count_recip); + auto broadcast = add_instruction( + HloInstruction::CreateBroadcast(operand->shape(), element_count, {})); + return HloInstruction::CreateBinary(operand->shape(), HloOpcode::kDivide, + operand, broadcast); + } + + std::unique_ptr DynamicElementCountPerFeature( + HloInstruction* operand, int64 feature_index, + const std::function)>& + add_instruction) { + auto elements_per_feature_u32 = add_instruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(1))); + + for (int64 i = 0; i < operand->shape().rank(); ++i) { + if (i == feature_index) { + continue; + } + auto dynamic_dimension_size = + add_instruction(HloInstruction::CreateGetDimensionSize( + ShapeUtil::MakeShape(U32, {}), operand, i)); + elements_per_feature_u32 = add_instruction(HloInstruction::CreateBinary( + ShapeUtil::MakeShape(U32, {}), HloOpcode::kMultiply, + dynamic_dimension_size, elements_per_feature_u32)); + } + + return HloInstruction::CreateConvert( + ShapeUtil::MakeShape(operand->shape().element_type(), {}), + elements_per_feature_u32); } // Replaces the existing HLO instruction old_instruction, with @@ -195,9 +213,6 @@ Status BatchNormExpanderVisitor::HandleBatchNormTraining( const Shape operand_shape = operand->shape(); PrimitiveType ptype = operand_shape.element_type(); int64 feature_index = batch_norm->feature_index(); - const int64 feature_count = operand_shape.dimensions(feature_index); - const int64 size_in_elements = ShapeUtil::ElementsIn(operand_shape); - int64 elements_per_feature_int64 = size_in_elements / feature_count; HloInstruction* scale = batch_norm->mutable_operand(1); HloInstruction* offset = batch_norm->mutable_operand(2); @@ -214,12 +229,15 @@ Status BatchNormExpanderVisitor::HandleBatchNormTraining( add(HloInstruction::CreateConstant(std::move(epsilon_literal))), {})); std::vector dimensions_without_feature; - for (int64 i = 0; i < ShapeUtil::Rank(operand_shape); ++i) { + for (int64 i = 0; i < operand_shape.rank(); ++i) { if (i != feature_index) { dimensions_without_feature.push_back(i); } } + auto elements_per_feature = + add(DynamicElementCountPerFeature(operand, feature_index, add)); + auto scale_broadcasted = add( HloInstruction::CreateBroadcast(operand_shape, scale, {feature_index})); @@ -243,13 +261,13 @@ Status BatchNormExpanderVisitor::HandleBatchNormTraining( add_reduce_computation)); // E[X]. - auto mean = add(Mean(elements_per_feature_int64, sum, add)); + auto mean = add(Mean(elements_per_feature, sum, add)); auto mean_broadcasted = add( HloInstruction::CreateBroadcast(operand_shape, mean, {feature_index})); // E[X^2]. - auto square_mean = add(Mean(elements_per_feature_int64, squared_sum, add)); + auto square_mean = add(Mean(elements_per_feature, squared_sum, add)); // E^2[X]. auto mean_square = @@ -339,7 +357,7 @@ Status BatchNormExpanderVisitor::HandleBatchNormInference( std::vector dimensions_without_feature; - for (int64 i = 0; i < ShapeUtil::Rank(operand_shape); ++i) { + for (int64 i = 0; i < operand_shape.rank(); ++i) { if (i != feature_index) { dimensions_without_feature.push_back(i); } @@ -458,9 +476,8 @@ Status BatchNormExpanderVisitor::HandleBatchNormGrad( int64 feature_index = batch_norm->feature_index(); - const int64 size_in_elements = ShapeUtil::ElementsIn(activation_shape); - const int64 feature_count = activation_shape.dimensions(feature_index); - const int64 elements_per_feature_int64 = size_in_elements / feature_count; + auto elements_per_feature = + add(DynamicElementCountPerFeature(activation, feature_index, add)); auto zero_literal = LiteralUtil::CreateR0(0.0f); TF_ASSIGN_OR_RETURN(zero_literal, zero_literal.Convert(ptype)); @@ -477,7 +494,7 @@ Status BatchNormExpanderVisitor::HandleBatchNormGrad( std::vector dimensions_without_feature; - for (int64 i = 0; i < ShapeUtil::Rank(activation_shape); ++i) { + for (int64 i = 0; i < activation_shape.rank(); ++i) { if (i != feature_index) { dimensions_without_feature.push_back(i); } @@ -553,15 +570,9 @@ Status BatchNormExpanderVisitor::HandleBatchNormGrad( add_binary(activation_shape, HloOpcode::kMultiply, scale_broadcasted, rsqrt_var_add_epsilon_broadcasted); - scale_times_rsqrt_var_add_epsilon = add( - Mean(elements_per_feature_int64, scale_times_rsqrt_var_add_epsilon, add)); + scale_times_rsqrt_var_add_epsilon = + add(Mean(elements_per_feature, scale_times_rsqrt_var_add_epsilon, add)); - auto elements_per_feature_literal = - LiteralUtil::CreateR0(elements_per_feature_int64); - TF_ASSIGN_OR_RETURN(elements_per_feature_literal, - elements_per_feature_literal.Convert(ptype)); - auto elements_per_feature = add( - HloInstruction::CreateConstant(std::move(elements_per_feature_literal))); auto i1 = add_binary(activation_shape, HloOpcode::kMultiply, grad_output, add(HloInstruction::CreateBroadcast( activation_shape, elements_per_feature, {}))); diff --git a/tensorflow/compiler/xla/service/batchnorm_expander_test.cc b/tensorflow/compiler/xla/service/batchnorm_expander_test.cc index f7ac8f5482908af104554a1cf812370b9098cda7..8e8fbbd935b154e5a77d68e60d861601d740bf03 100644 --- a/tensorflow/compiler/xla/service/batchnorm_expander_test.cc +++ b/tensorflow/compiler/xla/service/batchnorm_expander_test.cc @@ -29,14 +29,28 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_pass_fix.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/test.h" -#include "tensorflow/compiler/xla/tests/hlo_verified_test_base.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" namespace xla { namespace { -using BatchNormExpanderTest = HloVerifiedTestBase; +class BatchNormExpanderTest : public HloTestBase { + protected: + // BatchNorm should have a dynamic sized dividor for mean operations. + int64 CountGetDimensionSize(const HloModule& module) { + int64 count = 0; + for (HloComputation* comp : module.computations()) { + for (HloInstruction* inst : comp->instructions()) { + if (inst->opcode() == HloOpcode::kGetDimensionSize) { + count++; + } + } + } + return count; + } +}; // Test that we expand BatchNormTraining. TEST_F(BatchNormExpanderTest, BatchNormTraining) { @@ -59,15 +73,16 @@ TEST_F(BatchNormExpanderTest, BatchNormTraining) { param0, param1, param2, /*epsilon=*/0.001, /*feature_index=*/3)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root->opcode(), HloOpcode::kBatchNormTraining); BatchNormExpander rewriter(/*rewrite_training_op=*/true, /*rewrite_inference_op=*/true, /*rewrite_grad_op=*/true); - ASSERT_TRUE(rewriter.Run(module).ValueOrDie()); + ASSERT_TRUE(rewriter.Run(module.get()).ValueOrDie()); root = computation->root_instruction(); + EXPECT_EQ(CountGetDimensionSize(*module), 3); // Make sure this operation is expanded. EXPECT_EQ(root->opcode(), HloOpcode::kTuple); } @@ -101,15 +116,16 @@ TEST_F(BatchNormExpanderTest, BatchNormGrad) { param1, param2, param3, param4, /*epsilon=*/0.001, /*feature_index=*/3)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root->opcode(), HloOpcode::kBatchNormGrad); BatchNormExpander rewriter(/*rewrite_training_op=*/true, /*rewrite_inference_op=*/true, /*rewrite_grad_op=*/true); - ASSERT_TRUE(rewriter.Run(module).ValueOrDie()); + ASSERT_TRUE(rewriter.Run(module.get()).ValueOrDie()); root = computation->root_instruction(); + EXPECT_EQ(CountGetDimensionSize(*module), 3); // Make sure this operation is expanded. EXPECT_EQ(root->opcode(), HloOpcode::kTuple); } @@ -126,13 +142,13 @@ ENTRY entry { epsilon=0.001, feature_index=1, sharding={maximal device=1} })"; - ParseAndVerifyModule(module_str); + TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(module_str)); BatchNormExpander rewriter(/*rewrite_training_op=*/true, /*rewrite_inference_op=*/true, /*rewrite_grad_op=*/true); - ASSERT_TRUE(rewriter.Run(&module()).ValueOrDie()); + ASSERT_TRUE(rewriter.Run(m.get()).ValueOrDie()); - for (auto* instruction : module().entry_computation()->instructions()) { + for (auto* instruction : m->entry_computation()->instructions()) { if (instruction->opcode() == HloOpcode::kParameter) { continue; } diff --git a/tensorflow/compiler/xla/service/bfloat16_conversion_folding.cc b/tensorflow/compiler/xla/service/bfloat16_conversion_folding.cc index d63287539dfde5bb4890ab8303ef2205133d8125..e62d72b323bd1d113e9d87bf8602bfb434c40d61 100644 --- a/tensorflow/compiler/xla/service/bfloat16_conversion_folding.cc +++ b/tensorflow/compiler/xla/service/bfloat16_conversion_folding.cc @@ -34,8 +34,8 @@ class BFloat16ConversionFoldingVisitor : public DfsHloVisitorWithDefault { Status DefaultAction(HloInstruction* hlo) override; - // Special handling for cross-replica-sum which can have a tuple output. - Status HandleCrossReplicaSum(HloInstruction* crs) override; + // Special handling for all-reduce which can have a tuple output. + Status HandleAllReduce(HloInstruction* crs) override; static bool Run(HloComputation* computation, const BFloat16Support* bfloat16_support) { @@ -151,15 +151,10 @@ Status BFloat16ConversionFoldingVisitor::TryFoldBF16Conversions( Status BFloat16ConversionFoldingVisitor::DefaultAction(HloInstruction* hlo) { // Do not fold BF16 conversions for instructions related to tuples, entry and - // exit of a computation, fusion, convert, and control flow. + // exit of a computation, fusion, convert, side-effecting instructions and + // control flow. if (hlo->opcode() == HloOpcode::kTuple || // hlo->opcode() == HloOpcode::kGetTupleElement || // - hlo->opcode() == HloOpcode::kInfeed || // - hlo->opcode() == HloOpcode::kOutfeed || // - hlo->opcode() == HloOpcode::kSend || // - hlo->opcode() == HloOpcode::kSendDone || // - hlo->opcode() == HloOpcode::kRecv || // - hlo->opcode() == HloOpcode::kRecvDone || // hlo->opcode() == HloOpcode::kConstant || // hlo->opcode() == HloOpcode::kParameter || // hlo->opcode() == HloOpcode::kFusion || // @@ -167,7 +162,8 @@ Status BFloat16ConversionFoldingVisitor::DefaultAction(HloInstruction* hlo) { hlo->opcode() == HloOpcode::kCall || // hlo->opcode() == HloOpcode::kCustomCall || // hlo->opcode() == HloOpcode::kWhile || // - hlo->opcode() == HloOpcode::kConditional) { + hlo->opcode() == HloOpcode::kConditional || // + hlo->HasSideEffectNoRecurse()) { return Status::OK(); } if (hlo == computation_->root_instruction() && @@ -180,8 +176,11 @@ Status BFloat16ConversionFoldingVisitor::DefaultAction(HloInstruction* hlo) { return TryFoldBF16Conversions(hlo); } -Status BFloat16ConversionFoldingVisitor::HandleCrossReplicaSum( - HloInstruction* crs) { +Status BFloat16ConversionFoldingVisitor::HandleAllReduce(HloInstruction* crs) { + if (crs->IsCrossModuleAllReduce()) { + // Cross-module all-reduce has side effect. + return Status::OK(); + } // First use DefaultAction() to handle the operands. It can't handle // tuple-shaped output. TF_RETURN_IF_ERROR(DefaultAction(crs)); @@ -191,7 +190,7 @@ Status BFloat16ConversionFoldingVisitor::HandleCrossReplicaSum( } // If the output is not a tuple, we don't need special handling. - if (!ShapeUtil::IsTuple(crs->shape())) { + if (!crs->shape().IsTuple()) { return Status::OK(); } diff --git a/tensorflow/compiler/xla/service/bfloat16_conversion_folding_test.cc b/tensorflow/compiler/xla/service/bfloat16_conversion_folding_test.cc index 5f93740887aa7e61458990992fe0573883ff056d..2232a2cbdfe0cf64dc4fb10d4598c0ad8b51ee5e 100644 --- a/tensorflow/compiler/xla/service/bfloat16_conversion_folding_test.cc +++ b/tensorflow/compiler/xla/service/bfloat16_conversion_folding_test.cc @@ -22,7 +22,7 @@ limitations under the License. #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_verified_test_base.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" #include "tensorflow/compiler/xla/xla_data.pb.h" namespace xla { @@ -38,7 +38,7 @@ class TestBFloat16Support : public BFloat16Support { hlo.opcode() == HloOpcode::kSubtract || hlo.opcode() == HloOpcode::kTuple || hlo.opcode() == HloOpcode::kGetTupleElement || - hlo.opcode() == HloOpcode::kCrossReplicaSum) { + hlo.opcode() == HloOpcode::kAllReduce) { return true; } return false; @@ -49,7 +49,7 @@ class TestBFloat16Support : public BFloat16Support { hlo.opcode() == HloOpcode::kSubtract || hlo.opcode() == HloOpcode::kTuple || hlo.opcode() == HloOpcode::kGetTupleElement || - hlo.opcode() == HloOpcode::kCrossReplicaSum) { + hlo.opcode() == HloOpcode::kAllReduce) { return true; } return false; @@ -58,18 +58,18 @@ class TestBFloat16Support : public BFloat16Support { bool SupportsMixedPrecisions(const HloInstruction& hlo) const override { if (hlo.opcode() == HloOpcode::kAdd || hlo.opcode() == HloOpcode::kTuple || hlo.opcode() == HloOpcode::kGetTupleElement || - hlo.opcode() == HloOpcode::kCrossReplicaSum) { + hlo.opcode() == HloOpcode::kAllReduce) { return true; } return false; } }; -class BFloat16ConversionFoldingTest : public HloVerifiedTestBase { +class BFloat16ConversionFoldingTest : public HloTestBase { protected: BFloat16ConversionFoldingTest() - : HloVerifiedTestBase(/*layout_sensitive=*/false, - /*allow_mixed_precision=*/true) {} + : HloTestBase(/*verifier_layout_sensitive=*/false, + /*allow_mixed_precision_in_hlo_verifier=*/true) {} bool FoldConversions(HloModule* module) { TestBFloat16Support bfloat16_support_; @@ -103,10 +103,10 @@ TEST_F(BFloat16ConversionFoldingTest, FoldIfSupported) { HloInstruction::CreateBinary(f32_shape, HloOpcode::kAdd, convert1, c)); builder.AddInstruction(HloInstruction::CreateConvert(bf16_shape, add1)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(FoldConversions(module)); + EXPECT_TRUE(FoldConversions(module.get())); EXPECT_EQ(computation->root_instruction(), add1); EXPECT_EQ(add0->shape().element_type(), BF16); @@ -138,10 +138,10 @@ TEST_F(BFloat16ConversionFoldingTest, DoNotFoldIfUnsupported) { HloInstruction* convert2 = builder.AddInstruction(HloInstruction::CreateConvert(bf16_shape, mul1)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_FALSE(FoldConversions(module)); + EXPECT_FALSE(FoldConversions(module.get())); EXPECT_EQ(computation->root_instruction(), convert2); EXPECT_EQ(mul0->shape().element_type(), F32); @@ -173,10 +173,10 @@ TEST_F(BFloat16ConversionFoldingTest, DoNotFoldUnsupportedMixedPrecision) { HloInstruction* convert2 = builder.AddInstruction(HloInstruction::CreateConvert(bf16_shape, sub1)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_FALSE(FoldConversions(module)); + EXPECT_FALSE(FoldConversions(module.get())); EXPECT_EQ(computation->root_instruction(), convert2); EXPECT_EQ(sub0->shape().element_type(), F32); @@ -203,20 +203,20 @@ TEST_F(BFloat16ConversionFoldingTest, DoNotFoldTuple) { HloInstruction* convert1 = builder.AddInstruction(HloInstruction::CreateConvert(bf16_shape, gte)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_FALSE(FoldConversions(module)); + EXPECT_FALSE(FoldConversions(module.get())); EXPECT_EQ(computation->root_instruction(), convert1); EXPECT_EQ(gte->shape().element_type(), F32); EXPECT_EQ(tuple->operand(1), convert0); } -TEST_F(BFloat16ConversionFoldingTest, FoldCrossReplicaSumTupleOutput) { +TEST_F(BFloat16ConversionFoldingTest, FoldAllReduceTupleOutput) { auto builder = HloComputation::Builder(TestName()); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation::Builder sum_builder("add"); auto x = sum_builder.AddInstruction(HloInstruction::CreateParameter( /*parameter_number=*/0, ShapeUtil::MakeShape(F32, {}), "x")); @@ -236,11 +236,10 @@ TEST_F(BFloat16ConversionFoldingTest, FoldCrossReplicaSumTupleOutput) { HloInstruction* b = builder.AddInstruction( HloInstruction::CreateParameter(1, f32_shape, "b")); - HloInstruction* crs = - builder.AddInstruction(HloInstruction::CreateCrossReplicaSum( - ShapeUtil::MakeTupleShape({f32_shape, f32_shape}), {convert_a, b}, - sum, /*replica_groups=*/{}, /*barrier=*/"", - /*all_reduce_id=*/absl::nullopt)); + HloInstruction* crs = builder.AddInstruction(HloInstruction::CreateAllReduce( + ShapeUtil::MakeTupleShape({f32_shape, f32_shape}), {convert_a, b}, sum, + /*replica_groups=*/{}, /*barrier=*/"", + /*all_reduce_id=*/absl::nullopt)); HloInstruction* gte_a = builder.AddInstruction( HloInstruction::CreateGetTupleElement(f32_shape, crs, 0)); HloInstruction* gte_b = builder.AddInstruction( @@ -252,7 +251,7 @@ TEST_F(BFloat16ConversionFoldingTest, FoldCrossReplicaSumTupleOutput) { auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(FoldConversions(module)); + EXPECT_TRUE(FoldConversions(module.get())); EXPECT_EQ(computation->root_instruction(), tuple); EXPECT_EQ(tuple->operand(0), gte_a); diff --git a/tensorflow/compiler/xla/service/bfloat16_normalization.cc b/tensorflow/compiler/xla/service/bfloat16_normalization.cc index 1251f0258f5d43a490ad654f519fee9076590453..d1b14d604f0559b6b18f7d1fba127669c241c8a3 100644 --- a/tensorflow/compiler/xla/service/bfloat16_normalization.cc +++ b/tensorflow/compiler/xla/service/bfloat16_normalization.cc @@ -346,11 +346,9 @@ Status BFloat16NormalizationVisitor::HandleInstruction(HloInstruction* hlo) { Status BFloat16NormalizationVisitor::DefaultAction(HloInstruction* hlo) { // Do not change instructions related to entry and exit of a computation, - // tuples, fusion, convert, and control flow. + // tuples, fusion, convert, side-effecting instructions, and control flow. if (hlo->opcode() == HloOpcode::kTuple || // hlo->opcode() == HloOpcode::kGetTupleElement || // - hlo->opcode() == HloOpcode::kInfeed || // - hlo->opcode() == HloOpcode::kOutfeed || // hlo->opcode() == HloOpcode::kConstant || // hlo->opcode() == HloOpcode::kParameter || // hlo->opcode() == HloOpcode::kFusion || // @@ -358,13 +356,14 @@ Status BFloat16NormalizationVisitor::DefaultAction(HloInstruction* hlo) { hlo->opcode() == HloOpcode::kCall || // hlo->opcode() == HloOpcode::kCustomCall || // hlo->opcode() == HloOpcode::kWhile || // - hlo->opcode() == HloOpcode::kConditional) { + hlo->opcode() == HloOpcode::kConditional || // + hlo->HasSideEffectNoRecurse()) { return Status::OK(); } // TODO(b/112040122): Correctly normalize variadic reduce. if ((hlo->opcode() == HloOpcode::kSort || - hlo->opcode() == HloOpcode::kCrossReplicaSum) && - ShapeUtil::IsTuple(hlo->shape())) { + hlo->opcode() == HloOpcode::kAllReduce) && + hlo->shape().IsTuple()) { return HandleMultipleOutputs(hlo); } return HandleInstruction(hlo); diff --git a/tensorflow/compiler/xla/service/bfloat16_normalization_test.cc b/tensorflow/compiler/xla/service/bfloat16_normalization_test.cc index cb075a5e38a5ea9db2ceb432b2b59f8db5e2e640..551ac4be73a7630d213a53ca3606aa7f890cd794 100644 --- a/tensorflow/compiler/xla/service/bfloat16_normalization_test.cc +++ b/tensorflow/compiler/xla/service/bfloat16_normalization_test.cc @@ -23,7 +23,7 @@ limitations under the License. #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_verified_test_base.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" #include "tensorflow/compiler/xla/xla_data.pb.h" namespace xla { @@ -68,11 +68,11 @@ class TestBFloat16Support : public BFloat16Support { } }; -class BFloat16NormalizationTest : public HloVerifiedTestBase { +class BFloat16NormalizationTest : public HloTestBase { protected: BFloat16NormalizationTest() - : HloVerifiedTestBase(/*layout_sensitive=*/false, - /*allow_mixed_precision=*/true) {} + : HloTestBase(/*verifier_layout_sensitive=*/false, + /*allow_mixed_precision_in_hlo_verifier=*/true) {} bool Normalize(HloModule* module) { TestBFloat16Support bfloat16_support_; @@ -106,10 +106,10 @@ TEST_F(BFloat16NormalizationTest, NoopIfSupported) { HloInstruction* add1 = builder.AddInstruction( HloInstruction::CreateBinary(f32_shape, HloOpcode::kAdd, add0, c)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_FALSE(Normalize(module)); + EXPECT_FALSE(Normalize(module.get())); EXPECT_EQ(computation->root_instruction(), add1); EXPECT_EQ(add0->shape().element_type(), BF16); @@ -134,10 +134,10 @@ TEST_F(BFloat16NormalizationTest, ResolveIfUnsupportedBF16) { HloInstruction* mul1 = builder.AddInstruction( HloInstruction::CreateBinary(bf16_shape, HloOpcode::kMultiply, mul0, c)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(Normalize(module)); + EXPECT_TRUE(Normalize(module.get())); EXPECT_EQ(computation->root_instruction()->opcode(), HloOpcode::kConvert); EXPECT_EQ(computation->root_instruction()->operand(0), mul1); @@ -164,10 +164,10 @@ TEST_F(BFloat16NormalizationTest, ResolveUnsupportedMixedPrecisionSubtraction) { HloInstruction* sub1 = builder.AddInstruction( HloInstruction::CreateBinary(bf16_shape, HloOpcode::kSubtract, sub0, c)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(Normalize(module)); + EXPECT_TRUE(Normalize(module.get())); EXPECT_EQ(computation->root_instruction()->opcode(), HloOpcode::kConvert); EXPECT_EQ(computation->root_instruction()->operand(0), sub1); @@ -191,7 +191,7 @@ TEST_F(BFloat16NormalizationTest, ResolveUnsupportedMixedPrecisionReduce) { HloInstruction::CreateBinary(bf16_scalar_shape, HloOpcode::kAdd, reduce_comp_param0, reduce_comp_param1)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto reduce_computation = module->AddEmbeddedComputation(reduce_comp_builder.Build()); @@ -205,7 +205,7 @@ TEST_F(BFloat16NormalizationTest, ResolveUnsupportedMixedPrecisionReduce) { auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(Normalize(module)); + EXPECT_TRUE(Normalize(module.get())); EXPECT_EQ(computation->root_instruction(), reduce); EXPECT_EQ(reduce->called_computations().size(), 1); @@ -232,8 +232,8 @@ TEST_F(BFloat16NormalizationTest, ResolveUnsupportedMixedPrecisionReduce) { EXPECT_EQ(reduce->operand(1)->shape().element_type(), F32); } -TEST_F(BFloat16NormalizationTest, ResolveMixedPrecisionTupleCrossReplicaSum) { - auto module = CreateNewModule(); +TEST_F(BFloat16NormalizationTest, ResolveMixedPrecisionTupleAllReduce) { + auto module = CreateNewVerifiedModule(); HloComputation::Builder sum_builder("sum"); auto x = sum_builder.AddInstruction(HloInstruction::CreateParameter( /*parameter_number=*/0, ShapeUtil::MakeShape(F32, {}), "x")); @@ -253,17 +253,16 @@ TEST_F(BFloat16NormalizationTest, ResolveMixedPrecisionTupleCrossReplicaSum) { HloInstruction* b = builder.AddInstruction( HloInstruction::CreateParameter(1, bf16_shape, "b")); - HloInstruction* crs = - builder.AddInstruction(HloInstruction::CreateCrossReplicaSum( - ShapeUtil::MakeTupleShape({f32_shape, bf16_shape}), {a, b}, reduction, - /*replica_groups=*/{}, /*barrier=*/"", - /*all_reduce_id=*/absl::nullopt)); + HloInstruction* crs = builder.AddInstruction(HloInstruction::CreateAllReduce( + ShapeUtil::MakeTupleShape({f32_shape, bf16_shape}), {a, b}, reduction, + /*replica_groups=*/{}, /*barrier=*/"", + /*all_reduce_id=*/absl::nullopt)); HloInstruction* gte = builder.AddInstruction( HloInstruction::CreateGetTupleElement(bf16_shape, crs, 1)); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(Normalize(module)); + EXPECT_TRUE(Normalize(module.get())); EXPECT_EQ(computation->root_instruction(), gte); EXPECT_EQ(gte->shape().element_type(), BF16); @@ -272,7 +271,7 @@ TEST_F(BFloat16NormalizationTest, ResolveMixedPrecisionTupleCrossReplicaSum) { } TEST_F(BFloat16NormalizationTest, ResolveMixedPrecisionTupleSort) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto builder = HloComputation::Builder(TestName()); Shape f32_shape = ShapeUtil::MakeShape(F32, {1024}); Shape bf16_shape = ShapeUtil::MakeShape(BF16, {1024}); @@ -290,7 +289,7 @@ TEST_F(BFloat16NormalizationTest, ResolveMixedPrecisionTupleSort) { auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(Normalize(module)); + EXPECT_TRUE(Normalize(module.get())); EXPECT_EQ(computation->root_instruction(), gte); EXPECT_EQ(gte->shape().element_type(), BF16); @@ -299,7 +298,7 @@ TEST_F(BFloat16NormalizationTest, ResolveMixedPrecisionTupleSort) { } TEST_F(BFloat16NormalizationTest, ResolveMixedPrecisionTupleSortRoot) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto builder = HloComputation::Builder(TestName()); Shape f32_shape = ShapeUtil::MakeShape(F32, {1024}); Shape bf16_shape = ShapeUtil::MakeShape(BF16, {1024}); @@ -314,7 +313,7 @@ TEST_F(BFloat16NormalizationTest, ResolveMixedPrecisionTupleSortRoot) { auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(Normalize(module)); + EXPECT_TRUE(Normalize(module.get())); EXPECT_EQ(sort->operand(0)->shape().element_type(), F32); EXPECT_EQ(ShapeUtil::GetSubshape(sort->shape(), {0}).element_type(), F32); @@ -342,10 +341,10 @@ TEST_F(BFloat16NormalizationTest, DoNotAddUnsupportedMixedPrecision) { HloInstruction* dot = builder.AddInstruction( HloInstruction::CreateDot(bf16_shape, a, b, dot_dnums, precision_config)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(Normalize(module)); + EXPECT_TRUE(Normalize(module.get())); EXPECT_EQ(computation->root_instruction()->opcode(), HloOpcode::kConvert); EXPECT_EQ(dot->shape().element_type(), F32); diff --git a/tensorflow/compiler/xla/service/bfloat16_propagation.cc b/tensorflow/compiler/xla/service/bfloat16_propagation.cc index 002be9c97098ef1f73446c458dae24bbc826a626..bab63f66d83b712d756078bef84926eed235f6b5 100644 --- a/tensorflow/compiler/xla/service/bfloat16_propagation.cc +++ b/tensorflow/compiler/xla/service/bfloat16_propagation.cc @@ -236,6 +236,10 @@ bool BFloat16Propagation::AllUsersConsumeBF16(const HloInstruction& hlo, // the end of the BFloat16Propagation pass. continue; } + if (use.instruction->HasSideEffectNoRecurse()) { + // Keep side-effecting instruction's operands unchanged. + return false; + } // Any visited user that can accept BF16 has already been updated if // necessary, e.g., the output has been changed to BF16 if it propagates // precision, or a called computation's parameters have been changed to @@ -272,8 +276,8 @@ bool BFloat16Propagation::AllUsersConsumeBF16(const HloInstruction& hlo, if (bfloat16_support_->EffectiveOperandPrecisionIsOutputPrecision( *use.instruction, use.operand_number)) { if (use.instruction->opcode() == HloOpcode::kTuple || - (use.instruction->opcode() == HloOpcode::kCrossReplicaSum && - ShapeUtil::IsTuple(use.instruction->shape()))) { + (use.instruction->opcode() == HloOpcode::kAllReduce && + use.instruction->shape().IsTuple())) { ShapeIndex use_output_index{use.operand_number}; for (int64 i : use.operand_index) { use_output_index.push_back(i); @@ -329,22 +333,6 @@ void BFloat16Propagation::DetermineInstructionPrecision(HloInstruction* hlo, return; } - // Do not change precision for instructions related to entry and exit of a - // computation, and control flow, because this pass might break the interfaces - // or assumptions for them. - if (hlo->opcode() == HloOpcode::kInfeed || // - hlo->opcode() == HloOpcode::kOutfeed || // - hlo->opcode() == HloOpcode::kSend || // - hlo->opcode() == HloOpcode::kSendDone || // - hlo->opcode() == HloOpcode::kRecv || // - hlo->opcode() == HloOpcode::kRecvDone || // - hlo->opcode() == HloOpcode::kCustomCall || // - hlo->opcode() == HloOpcode::kCall || // - hlo->opcode() == HloOpcode::kConditional || // - (hlo->opcode() == HloOpcode::kParameter && skip_parameters)) { - return; - } - // Prevent root instructions from having their output modified by recording // all F32 output values as needing to stay as F32. CHECK(hlo->parent() != nullptr); @@ -366,6 +354,17 @@ void BFloat16Propagation::DetermineInstructionPrecision(HloInstruction* hlo, return; } + // Do not change precision for instructions related to entry and exit of a + // computation, side-effecting instructions, and control flow, because this + // pass might break the interfaces or assumptions for them. + if (hlo->opcode() == HloOpcode::kCustomCall || // + hlo->opcode() == HloOpcode::kCall || // + hlo->opcode() == HloOpcode::kConditional || // + hlo->HasSideEffectNoRecurse() || // + (hlo->opcode() == HloOpcode::kParameter && skip_parameters)) { + return; + } + if (!ContainsKey(consider_using_bfloat16_, hlo)) { return; } diff --git a/tensorflow/compiler/xla/service/bfloat16_propagation_test.cc b/tensorflow/compiler/xla/service/bfloat16_propagation_test.cc index e032b5c624c0151fd63c870e0f21ec97656d625f..a9b5d9916e400b39039248098c22a715e44ccfd2 100644 --- a/tensorflow/compiler/xla/service/bfloat16_propagation_test.cc +++ b/tensorflow/compiler/xla/service/bfloat16_propagation_test.cc @@ -22,7 +22,7 @@ limitations under the License. #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_verified_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/xla_data.pb.h" @@ -55,11 +55,11 @@ class TestBFloat16Support : public BFloat16Support { } }; -class BFloat16PropagationTest : public HloVerifiedTestBase { +class BFloat16PropagationTest : public HloTestBase { protected: BFloat16PropagationTest() - : HloVerifiedTestBase(/*layout_sensitive=*/false, - /*allow_mixed_precision=*/true) {} + : HloTestBase(/*verifier_layout_sensitive=*/false, + /*allow_mixed_precision_in_hlo_verifier=*/true) {} // Runs the propagation pass on the given module, and returns whether the // module is changed after this pass. @@ -121,10 +121,10 @@ TEST_F(BFloat16PropagationTest, PropagateThroughSelectButNotAdd) { HloInstruction* root = builder.AddInstruction(HloInstruction::CreateBinary( ShapeUtil::MakeShape(F32, {4, 4}), HloOpcode::kAdd, dot, dot)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(PropagatePrecision(module)); + EXPECT_TRUE(PropagatePrecision(module.get())); EXPECT_EQ(computation->root_instruction(), root); EXPECT_TRUE(OutputsBF16(xpose)); @@ -136,6 +136,96 @@ TEST_F(BFloat16PropagationTest, PropagateThroughSelectButNotAdd) { EXPECT_FALSE(OutputsBF16(c)); } +TEST_F(BFloat16PropagationTest, PropagateThroughMaxPoolReduceWindow) { + auto module = CreateNewVerifiedModule(); + + auto sub_builder = HloComputation::Builder("max"); + HloInstruction* p0 = sub_builder.AddInstruction( + HloInstruction::CreateParameter(0, ShapeUtil::MakeShape(F32, {}), "a")); + HloInstruction* p1 = sub_builder.AddInstruction( + HloInstruction::CreateParameter(1, ShapeUtil::MakeShape(F32, {}), "b")); + sub_builder.AddInstruction(HloInstruction::CreateBinary( + ShapeUtil::MakeShape(F32, {}), HloOpcode::kMaximum, p0, p1)); + auto max_computation = module->AddEmbeddedComputation(sub_builder.Build()); + + auto builder = HloComputation::Builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {2, 4}); + + HloInstruction* a = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "a")); + HloInstruction* b = + builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "b")); + HloInstruction* c = + builder.AddInstruction(HloInstruction::CreateParameter(2, shape, "c")); + HloInstruction* add = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, a, b)); + Window window; + WindowDimension dim; + dim.set_size(2); + dim.set_stride(1); + dim.set_padding_high(1); + dim.set_window_dilation(1); + dim.set_base_dilation(1); + *window.add_dimensions() = dim; + *window.add_dimensions() = dim; + HloInstruction* rw = + builder.AddInstruction(HloInstruction::CreateReduceWindow( + shape, add, + builder.AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::Zero(F32))), + window, max_computation)); + HloInstruction* xpose = + builder.AddInstruction(HloInstruction::CreateTranspose( + ShapeUtil::MakeShape(F32, {4, 2}), c, {1, 0})); + HloInstruction* dot = builder.AddInstruction( + CreateDot(ShapeUtil::MakeShape(F32, {4, 4}), xpose, rw)); + HloInstruction* root = builder.AddInstruction(HloInstruction::CreateBinary( + ShapeUtil::MakeShape(F32, {4, 4}), HloOpcode::kAdd, dot, dot)); + + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_TRUE(PropagatePrecision(module.get())); + + EXPECT_EQ(computation->root_instruction(), root); + EXPECT_TRUE(OutputsBF16(add)); + EXPECT_TRUE(OutputsBF16(xpose)); + EXPECT_TRUE(OutputsBF16(rw)); +} + +// Tests that side-effecting all-reduce should not be changed. +TEST_F(BFloat16PropagationTest, DoNotChangeAllReduce) { + auto module = CreateNewVerifiedModule(); + + auto builder = HloComputation::Builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {4, 4}); + HloInstruction* a = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "a")); + HloInstruction* b = + builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "b")); + auto rb = HloComputation::Builder(TestName()); + rb.AddInstruction(HloInstruction::CreateBinary( + shape, HloOpcode::kAdd, + rb.AddInstruction(HloInstruction::CreateParameter(0, shape, "p0")), + rb.AddInstruction(HloInstruction::CreateParameter(1, shape, "p1")))); + auto reduction = module->AddEmbeddedComputation(rb.Build()); + HloInstruction* all_reduce = + builder.AddInstruction(HloInstruction::CreateAllReduce( + ShapeUtil::MakeTupleShape({shape, shape}), {a, b}, reduction, + /*replica_groups=*/{}, /*barrier=*/"", /*all_reduce_id=*/1)); + HloInstruction* gte0 = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(shape, all_reduce, 0)); + HloInstruction* gte1 = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(shape, all_reduce, 1)); + HloInstruction* dot = builder.AddInstruction(CreateDot(shape, gte0, gte1)); + HloInstruction* root = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, dot, dot)); + + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_FALSE(PropagatePrecision(module.get())); + EXPECT_EQ(computation->root_instruction(), root); +} + // Tests that if a constant is converted to BF16 then its literal must also be // converted. TEST_F(BFloat16PropagationTest, ConvertConstantLiteral) { @@ -152,10 +242,10 @@ TEST_F(BFloat16PropagationTest, ConvertConstantLiteral) { HloInstruction::CreateConstant(LiteralUtil::CreateFromArray(array_b))); HloInstruction* dot = builder.AddInstruction(CreateDot(shape, a, b)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(PropagatePrecision(module)); + EXPECT_TRUE(PropagatePrecision(module.get())); EXPECT_EQ(computation->root_instruction(), dot); EXPECT_TRUE(OutputsBF16(dot->operand(0))); @@ -208,10 +298,10 @@ TEST_F(BFloat16PropagationTest, PropagateThroughTuples) { HloInstruction* output_tuple = builder.AddInstruction(HloInstruction::CreateTuple({dot, add2})); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(PropagatePrecision(module)); + EXPECT_TRUE(PropagatePrecision(module.get())); EXPECT_EQ(computation->root_instruction(), output_tuple); EXPECT_TRUE(OutputsBF16(xpose)); @@ -247,10 +337,10 @@ TEST_F(BFloat16PropagationTest, SameValueReferencedTwice) { HloInstruction* dot = builder.AddInstruction( CreateDot(ShapeUtil::MakeShape(F32, {4, 4}), lhs, rhs)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(PropagatePrecision(module)); + EXPECT_TRUE(PropagatePrecision(module.get())); EXPECT_EQ(computation->root_instruction(), dot); EXPECT_TRUE(OutputsBF16(add1)); @@ -276,10 +366,10 @@ TEST_F(BFloat16PropagationTest, DoNotChangeComputationRoot) { HloInstruction* tuple = builder.AddInstruction(HloInstruction::CreateTuple({add, dot})); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_FALSE(PropagatePrecision(module)); + EXPECT_FALSE(PropagatePrecision(module.get())); EXPECT_EQ(computation->root_instruction(), tuple); EXPECT_FALSE(OutputsBF16(add)); @@ -287,7 +377,7 @@ TEST_F(BFloat16PropagationTest, DoNotChangeComputationRoot) { // Tests that BF16 is propagated properly through fused computations. TEST_F(BFloat16PropagationTest, PropagateThroughFusion) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto builder = HloComputation::Builder(TestName()); Shape shape = ShapeUtil::MakeShape(F32, {4, 4}); @@ -322,7 +412,7 @@ TEST_F(BFloat16PropagationTest, PropagateThroughFusion) { auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(PropagatePrecision(module)); + EXPECT_TRUE(PropagatePrecision(module.get())); EXPECT_EQ(computation->root_instruction(), fusion1); EXPECT_TRUE(OutputsBF16(add)); @@ -335,7 +425,7 @@ TEST_F(BFloat16PropagationTest, PropagateThroughFusion) { // Tests that changes to BF16 that cannot be propagated outside a fusion are // discarded. TEST_F(BFloat16PropagationTest, DiscardFusionInternalBF16Changes) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto builder = HloComputation::Builder(TestName()); Shape shape = ShapeUtil::MakeShape(F32, {4, 4}); @@ -359,7 +449,7 @@ TEST_F(BFloat16PropagationTest, DiscardFusionInternalBF16Changes) { auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_FALSE(PropagatePrecision(module)); + EXPECT_FALSE(PropagatePrecision(module.get())); EXPECT_EQ(computation->root_instruction(), fusion); } @@ -374,7 +464,7 @@ TEST_F(BFloat16PropagationTest, DiscardFusionInternalBF16Changes) { // (BF16, BF16) fusion_computation(F32 a, F32 b) // = tuple(BF16 convert(a), BF16 add(F32 a, F32 b)) TEST_F(BFloat16PropagationTest, ConvertTupleFusionElementIfUsedByAdd) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto builder = HloComputation::Builder(TestName()); Shape shape = ShapeUtil::MakeShape(F32, {4, 4}); @@ -405,7 +495,7 @@ TEST_F(BFloat16PropagationTest, ConvertTupleFusionElementIfUsedByAdd) { auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(PropagatePrecision(module)); + EXPECT_TRUE(PropagatePrecision(module.get())); EXPECT_EQ(computation->root_instruction(), dot); EXPECT_TRUE(OutputsBF16(gte0)); @@ -424,7 +514,7 @@ TEST_F(BFloat16PropagationTest, ConvertTupleFusionElementIfUsedByAdd) { // on_true and on_false must match, so that as long as one of them is F32, the // other must be F32 as well. TEST_F(BFloat16PropagationTest, SelectOverTuples) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto builder = HloComputation::Builder(TestName()); Shape shape = ShapeUtil::MakeShape(F32, {2, 4}); @@ -455,7 +545,7 @@ TEST_F(BFloat16PropagationTest, SelectOverTuples) { auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(PropagatePrecision(module)); + EXPECT_TRUE(PropagatePrecision(module.get())); EXPECT_EQ(computation->root_instruction(), dot); EXPECT_FALSE(OutputsBF16(add0)); @@ -468,7 +558,7 @@ TEST_F(BFloat16PropagationTest, SelectOverTuples) { // Tests that BF16 is propagated properly through a while computation with // non-tuple input/output. TEST_F(BFloat16PropagationTest, PropagateThroughSimpleWhile) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto builder = HloComputation::Builder(TestName()); Shape shape = ShapeUtil::MakeShape(F32, {4, 4}); @@ -511,7 +601,7 @@ TEST_F(BFloat16PropagationTest, PropagateThroughSimpleWhile) { auto dot = builder.AddInstruction(CreateDot(shape, while_hlo, while_hlo)); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(PropagatePrecision(module)); + EXPECT_TRUE(PropagatePrecision(module.get())); EXPECT_EQ(computation->root_instruction(), dot); EXPECT_TRUE( @@ -527,7 +617,7 @@ TEST_F(BFloat16PropagationTest, PropagateThroughSimpleWhile) { // made to the while body and thus the fusion node inside it. TEST_F(BFloat16PropagationTest, ConditionPreventsPropagationForFusionInsideWhile) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto builder = HloComputation::Builder(TestName()); Shape shape = ShapeUtil::MakeShape(F32, {4, 4}); @@ -576,7 +666,7 @@ TEST_F(BFloat16PropagationTest, auto dot = builder.AddInstruction(CreateDot(shape, while_hlo, while_hlo)); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_FALSE(PropagatePrecision(module)); + EXPECT_FALSE(PropagatePrecision(module.get())); EXPECT_EQ(computation->root_instruction(), dot); EXPECT_FALSE(OutputsBF16(add)); EXPECT_FALSE(OutputsBF16(body_fusion)); @@ -588,7 +678,7 @@ TEST_F(BFloat16PropagationTest, // Tests that BF16 is propagated properly through while computations with // tuple-shaped input/output. TEST_F(BFloat16PropagationTest, PropagateThroughTupleWhile) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto builder = HloComputation::Builder(TestName()); Shape shape = ShapeUtil::MakeShape(F32, {4, 4}); @@ -656,7 +746,7 @@ TEST_F(BFloat16PropagationTest, PropagateThroughTupleWhile) { auto dot = builder.AddInstruction(CreateDot(shape, lhs, rhs)); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(PropagatePrecision(module)); + EXPECT_TRUE(PropagatePrecision(module.get())); EXPECT_EQ(computation->root_instruction(), dot); EXPECT_TRUE(OutputsBF16(lhs)); @@ -675,7 +765,7 @@ TEST_F(BFloat16PropagationTest, PropagateThroughTupleWhile) { // Tests that BF16 is not propagated through multiple whiles that invoke the // same computation as long as one while prevents the propagation. TEST_F(BFloat16PropagationTest, DoNotPropagateWhilesCallingSameComputation) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto builder = HloComputation::Builder(TestName()); Shape shape = ShapeUtil::MakeShape(F32, {4, 4}); @@ -786,7 +876,7 @@ TEST_F(BFloat16PropagationTest, DoNotPropagateWhilesCallingSameComputation) { auto dot = builder.AddInstruction(CreateDot(shape, lhs, rhs)); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(PropagatePrecision(module)); + EXPECT_TRUE(PropagatePrecision(module.get())); EXPECT_FALSE(OutputsBF16(body_dot)); EXPECT_FALSE(OutputsBF16(body_rhs)); EXPECT_FALSE(OutputsBF16(body_lhs)); @@ -825,10 +915,10 @@ TEST_F(BFloat16PropagationTest, NoopConversionRemoved) { HloInstruction* add2 = builder.AddInstruction(HloInstruction::CreateBinary( bf16_shape, HloOpcode::kAdd, convert0, convert1)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(PropagatePrecision(module)); + EXPECT_TRUE(PropagatePrecision(module.get())); EXPECT_EQ(computation->root_instruction(), add2); EXPECT_EQ(add2->operand(0), add0); @@ -861,10 +951,10 @@ TEST_F(BFloat16PropagationTest, TupleDomain) { HloInstruction* root = builder.AddInstruction( HloInstruction::CreateBinary(shape, HloOpcode::kAdd, dot, dot)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(PropagatePrecision(module)); + EXPECT_TRUE(PropagatePrecision(module.get())); EXPECT_EQ(computation->root_instruction(), root); // test BF16 propagated through domain @@ -907,10 +997,10 @@ TEST_F(BFloat16PropagationTest, TupleDomainNoPropagation) { HloInstruction* root = builder.AddInstruction( HloInstruction::CreateBinary(shape, HloOpcode::kAdd, dot, dot)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(PropagatePrecision(module)); + EXPECT_TRUE(PropagatePrecision(module.get())); EXPECT_EQ(computation->root_instruction(), root); EXPECT_TRUE(OutputsBF16(a_trans)); diff --git a/tensorflow/compiler/xla/service/bfloat16_support.cc b/tensorflow/compiler/xla/service/bfloat16_support.cc index 5b48f10505e78c035608d4c575501e4623218987..2b9502f63a821f3675ddfb506f41bb2390cf4136 100644 --- a/tensorflow/compiler/xla/service/bfloat16_support.cc +++ b/tensorflow/compiler/xla/service/bfloat16_support.cc @@ -14,6 +14,7 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/compiler/xla/service/bfloat16_support.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" @@ -107,6 +108,21 @@ bool BFloat16Support::EffectiveOperandPrecisionIsOutputPrecision( case HloOpcode::kSelect: case HloOpcode::kTupleSelect: return operand_index == 1 || operand_index == 2; + case HloOpcode::kReduce: + case HloOpcode::kReduceWindow: { + HloComputation* reduce_comp = hlo.called_computations()[0]; + for (HloInstruction* inst : reduce_comp->instructions()) { + if (inst->opcode() == HloOpcode::kParameter) { + continue; + } + for (int64 i = 0; i < inst->operand_count(); ++i) { + if (!EffectiveOperandPrecisionIsOutputPrecision(*inst, i)) { + return false; + } + } + } + return true; + } default: break; } diff --git a/tensorflow/compiler/xla/service/buffer_assignment.cc b/tensorflow/compiler/xla/service/buffer_assignment.cc index d5d6a044a81303425495202d8a98c6735b0b8b89..e1b91b500191c7756f3d1a4b160a0dd1e09cfe7d 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. @@ -186,7 +185,7 @@ Status GatherComputationsByAllocationType( worklist.push_back(std::make_pair(subcomputation, false)); // Not thread local. break; - case HloOpcode::kCrossReplicaSum: + case HloOpcode::kAllReduce: case HloOpcode::kMap: case HloOpcode::kReduce: case HloOpcode::kReduceWindow: @@ -207,9 +206,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 +218,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 +232,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 +271,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 +308,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 +339,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( @@ -378,6 +371,20 @@ const BufferAllocation& BufferAssignment::GetAllocation( return allocations_[index]; } +const BufferAllocation* BufferAssignment::GetInstructionAllocation( + const HloInstruction* hlo, const ShapeIndex& shape_index) const { + const PointsToSet& points_to_set = points_to_analysis().GetPointsToSet(hlo); + const LogicalBuffer* buffer = points_to_set.element(shape_index)[0]; + + if (!HasAllocation(*buffer)) { + return nullptr; + } + + const BufferAllocation& instruction_allocation = + GetAssignedAllocation(*buffer); + return &instruction_allocation; +} + BufferAllocation* BufferAssignment::GetMutableAllocation( BufferAllocation::Index index) { return const_cast(&GetAllocation(index)); @@ -387,7 +394,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; } } @@ -445,8 +452,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 { @@ -473,10 +479,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 @@ -505,7 +510,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: " @@ -514,6 +519,9 @@ void BufferAssignment::AddAssignment(BufferAllocation* allocation, TF_CHECK_OK(points_to_analysis().VerifyBuffer(buffer)); allocation->AddAssignment(buffer, offset, size); + if (liveness().MaybeLiveOut(buffer)) { + allocation->set_maybe_live_out(true); + } allocation_index_for_buffer_[&buffer] = allocation->index(); } @@ -624,7 +632,7 @@ Status BufferAssignment::ComputeSummaryStats() { bool schedule_complete = true; for (const auto& computation : module_->computations()) { if (!computation->IsFusionComputation()) { - const std::vector* sequence = + const HloInstructionSequence* sequence = liveness_->hlo_ordering().SequentialOrder(*computation); if (sequence == nullptr) { schedule_complete = false; @@ -728,14 +736,90 @@ StatusOr> BufferAssigner::Run( LogicalBuffer::SizeFunction buffer_size, LogicalBuffer::AlignmentFunction color_alignment, bool allow_input_output_aliasing, bool allocate_buffers_for_constants, - BufferLiveness::Colorer colorer) { - BufferAssigner assigner(allow_input_output_aliasing, - allocate_buffers_for_constants, std::move(colorer)); + BufferLiveness::Colorer colorer, ReuseAllocationFunction reuse_checker) { + BufferAssigner assigner(allocate_buffers_for_constants, std::move(colorer), + std::move(reuse_checker)); return assigner.CreateAssignment(module, std::move(hlo_ordering), std::move(buffer_size), std::move(color_alignment)); } +namespace { + +// a and b are in different subcomputations. Check for the case +// where a is inside the while body, and b is outside, part of the same while's +// init-operand or while-result. +bool MayInterfereAcrossSubcomputations(BufferAssignment* assignment, + const LogicalBuffer& a_buffer, + const LogicalBuffer& b_buffer) { + const CallGraph& call_graph = + assignment->liveness().hlo_ordering().call_graph(); + const HloInstruction* a_ancestor; + const HloInstruction* b_ancestor; + std::tie(a_ancestor, b_ancestor) = + call_graph.NearestAncestorsInSameComputation(a_buffer.instruction(), + b_buffer.instruction()); + if (a_ancestor == nullptr) { + // No common ancestor. + return true; + } + if (a_ancestor->opcode() == HloOpcode::kWhile && + call_graph.InstructionIsNestedIn(a_buffer.instruction(), + a_ancestor->while_body())) { + const PointsToSet& init_set = + assignment->liveness().points_to_analysis().GetPointsToSet( + a_ancestor->operand(0)); + if (init_set.ContainsBuffer(b_buffer)) { + VLOG(4) << "Can't interfere: " << a_buffer << " and " << b_buffer + << " (part of while-operand)"; + return false; + } + const PointsToSet& while_set = + assignment->liveness().points_to_analysis().GetPointsToSet(a_ancestor); + if (while_set.ContainsBuffer(b_buffer)) { + VLOG(4) << "Can't interfere: " << a_buffer << " and " << b_buffer + << " (part of while)"; + return false; + } + } + return true; +} + +// Return true, if a and b can't possibly interfere (and therefore further +// checking for interference can be skipped). This function checks for special +// cases where copy insertion guarantees no interference, but the regular buffer +// liveness is too conservative: +// +// Operations inside a while-body can't interfere with operations outside the +// while op if their last use is at the while-loop itself as part of the +// while-init op, or the while-result. For ops that are live across a +// while-loop, copy insertion will already insert the necessary copies to avoid +// such interference. +// +// This allows sharing buffers in cases like this: +// init = {...} +// while (init): +// p = param(0) +// gte = get-tuple-element(p), index=i +// t1 = op1 (gte) +// t2 = op2 (t1) +// ROOT tuple = {..., t2, ...} +// +// where t1 and t2 can share the same buffer. +bool MaySkipInterferenceCheck(BufferAssignment* assignment, + const LogicalBuffer& a_buffer, + const LogicalBuffer& b_buffer) { + if (a_buffer.instruction()->parent() == b_buffer.instruction()->parent()) { + // Ops within the same computation are not handled here. Assume that they + // may interfere. + return false; + } + return !MayInterfereAcrossSubcomputations(assignment, a_buffer, b_buffer) || + !MayInterfereAcrossSubcomputations(assignment, b_buffer, a_buffer); +} + +} // namespace + bool BufferAssigner::MaybeAssignBuffer(BufferAllocation* allocation, const LogicalBuffer& buffer, BufferAssignment* assignment) { @@ -763,6 +847,12 @@ bool BufferAssigner::MaybeAssignBuffer(BufferAllocation* allocation, return false; } + if (reuse_checker_ != nullptr && + !reuse_checker_(*assignment, *allocation, buffer)) { + VLOG(4) << "Can't assign: reuse_checker_(allocation, buffer) == false"; + return false; + } + if (!allocation->is_reusable()) { VLOG(4) << "Can't assign: allocation is not reusable"; return false; @@ -770,6 +860,9 @@ bool BufferAssigner::MaybeAssignBuffer(BufferAllocation* allocation, for (const auto& buffer_offset_size : allocation->assigned_buffers()) { const LogicalBuffer& assigned_buffer = *buffer_offset_size.first; + if (MaySkipInterferenceCheck(assignment, buffer, assigned_buffer)) { + continue; + } if (assignment->liveness().MayInterfere(assigned_buffer, buffer)) { VLOG(4) << "Can't assign: assignee " << assigned_buffer << " may interfere with " << buffer; @@ -859,35 +952,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; @@ -919,10 +1012,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; } @@ -935,7 +1032,7 @@ Status BufferAssigner::AssignBuffersForComputation( continue; } - if (ShapeUtil::IsTuple(buffer->shape())) { + if (buffer->shape().IsTuple()) { BufferAllocation* allocation = assignment->NewAllocation(*buffer, buffer_size); allocation->set_is_tuple(true); @@ -955,7 +1052,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, @@ -986,7 +1083,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, @@ -1078,7 +1175,7 @@ Status BufferAssigner::AssignBuffersWithSequentialOrdering( const HloComputation* computation = pair.first; const flat_hash_set& buffers_to_assign = pair.second; - const std::vector* instruction_sequence = + const HloInstructionSequence* instruction_sequence = hlo_ordering.SequentialOrder(*computation); CHECK(instruction_sequence != nullptr) << computation->name(); schedule.set_sequence(computation, *instruction_sequence); @@ -1113,7 +1210,7 @@ Status BufferAssigner::AssignBuffersWithSequentialOrdering( const HloComputation* computation = pair.first; const flat_hash_set& buffers_to_assign = pair.second; - const std::vector* instruction_sequence = + const HloInstructionSequence* instruction_sequence = hlo_ordering.SequentialOrder(*computation); CHECK(instruction_sequence != nullptr) << computation->name(); auto color_map = SplitBuffersByColor(buffers_to_assign); @@ -1128,7 +1225,7 @@ Status BufferAssigner::AssignBuffersWithSequentialOrdering( TF_ASSIGN_OR_RETURN( const HeapSimulator::Result result, HeapSimulator::Run(get_heap_algorithm(alignment), *computation, - HloInstructionSequence(*instruction_sequence), + *instruction_sequence, assignment->points_to_analysis(), assignment->buffer_size_, options)); AssignBuffersFromHeapSimulator(result, assignment, @@ -1212,10 +1309,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; } @@ -1275,7 +1372,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], @@ -1324,41 +1421,50 @@ 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); + + // Do not merge if one of the sets includes live outs, entry parameters or + // constants. + // + // Buffer liveness does not report the correct live range for entry + // parameter and live out buffers so we have to special case them here. On + // backends that support constant buffer allocations, constant buffers are + // assigned globals in readonly storage so we can't merge colocated buffer + // sets containing constants with colocated buffer sets containing writing + // instructions or other constants. + // + // Moreover (on the CPU/GPU backends) the entry parameter buffers belong to + // the caller of the executable so we can't write to entry parameters + // either, and the argument for not merging constants also applies to entry + // parameters. + for (int64 i = 0; i < colocated_buffer_sets.size(); ++i) { + for (auto& buffer : colocated_buffer_sets[i]) { + if (buffer_liveness.MaybeLiveOut(*buffer) || + is_readonly_entry_parameter(*buffer) || + buffer->instruction()->opcode() == HloOpcode::kConstant) { + set_can_be_merged[i] = false; + break; + } + } + } + // Returns true if the two colocated buffer sets (specified by their indices // into the colocated_buffer_sets) can be merged into a single set. auto cannot_merge_buffer_sets = [&colocated_buffer_sets, &buffer_liveness, &buffer_size, - &is_entry_parameter](int64 i, int64 j) { - // Do not merge if one of the sets includes live outs, entry parameters or - // constants. - // - // Buffer liveness does not report the correct live range for entry - // parameter and live out buffers so we have to special case them here. On - // backends that support constant buffer allocations, constant buffers are - // assigned globals in readonly storage so we can't merge colocated buffer - // sets containing constants with colocated buffer sets containing writing - // instructions or other constants. - // - // Moreover (on the CPU/GPU backends) the entry parameter buffers belong to - // the caller of the executable so we can't write to entry parameters - // either, and the argument for not merging constants also applies to entry - // parameters. - for (int64 key : {i, j}) { - for (auto& buffer : colocated_buffer_sets[key]) { - if (buffer_liveness.MaybeLiveOut(*buffer) || - is_entry_parameter(*buffer) || - buffer->instruction()->opcode() == HloOpcode::kConstant) { - return true; - } - } + &set_can_be_merged](int64 i, int64 j) { + if (!set_can_be_merged[i] || !set_can_be_merged[j]) { + return true; } // Colocated sets satisfy the invariant that all buffers within a set have @@ -1428,16 +1534,19 @@ void BufferAssigner::BuildColocatedBufferSets( buffer_liveness.points_to_analysis(); // Set up colocated buffer set for input and output. + 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); }); @@ -1574,6 +1683,13 @@ void BufferAssigner::BuildColocatedBufferSets( return; } + int64 i = 0; + for (const auto& colocated_set : *colocated_buffer_sets) { + VLOG(4) << "Colocated set " << i++ << ":"; + for (const auto& buffer : colocated_set) { + VLOG(4) << " " << buffer->ToString(); + } + } // Try to find more coalescing opportunities among the colocated buffer sets. // // TODO(b/32491382): We should be able to remove this by using the @@ -1624,10 +1740,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); } @@ -1641,6 +1753,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 899cd36e1f98c9e7b8ba7e42c06ced5c3e8afcc8..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; @@ -362,6 +372,11 @@ class BufferAssignment { // with the given index. const BufferAllocation& GetAllocation(BufferAllocation::Index index) const; + // Returns the allocation with the given instruction and shape index. nullptr + // if no allocation exists. + const BufferAllocation* GetInstructionAllocation( + const HloInstruction* hlo, const ShapeIndex& shape_index) const; + // Builds and returns a vector containing the slices which might contain the // subvalue at the given index of given instruction. std::set GetAllSlices( @@ -520,6 +535,11 @@ class BufferAssignment { // A class which constructs a buffer assignment. class BufferAssigner { public: + // Returns false if a buffer cannot be assigned to given allocation. + using ReuseAllocationFunction = std::function; + // Build and return a BufferAssignment for the given module. The given // HloOrdering is used to determine buffer liveness. buffer_size and // color_alignment are functions which returns the size and alignment of a @@ -531,15 +551,16 @@ class BufferAssigner { LogicalBuffer::AlignmentFunction color_alignment, bool allow_input_output_aliasing = false, bool allocate_buffers_for_constants = false, - BufferLiveness::Colorer colorer = BufferLiveness::DefaultColorer()); + BufferLiveness::Colorer colorer = BufferLiveness::DefaultColorer(), + ReuseAllocationFunction reuse_checker = nullptr); private: - BufferAssigner(bool allow_input_output_aliasing, - bool allocate_buffers_for_constants, - BufferLiveness::Colorer colorer) - : allow_input_output_aliasing_(allow_input_output_aliasing), - allocate_buffers_for_constants_(allocate_buffers_for_constants), - colorer_(colorer) {} + BufferAssigner(bool allocate_buffers_for_constants, + BufferLiveness::Colorer colorer, + ReuseAllocationFunction reuse_checker) + : allocate_buffers_for_constants_(allocate_buffers_for_constants), + colorer_(colorer), + reuse_checker_(reuse_checker) {} virtual ~BufferAssigner() = default; // Create a buffer assignment. @@ -627,16 +648,15 @@ class BufferAssigner { LogicalBuffer::Color::Hasher> SplitBuffersByColor(const absl::flat_hash_set& buffers); - // If true, buffer assignments assumes that input parameter buffers and output - // buffers can be shared if their sizes match. - bool allow_input_output_aliasing_; - // If true, allocate buffers for constant instructions. bool allocate_buffers_for_constants_; // Functor used to assign colors to newly allocated logical buffers. BufferLiveness::Colorer colorer_; + // Functor to check if a buffer can reuse an allocation. + ReuseAllocationFunction reuse_checker_; + TF_DISALLOW_COPY_AND_ASSIGN(BufferAssigner); }; diff --git a/tensorflow/compiler/xla/service/buffer_assignment_test.cc b/tensorflow/compiler/xla/service/buffer_assignment_test.cc index 795beb9ff5ceb2998a85fbd03d8bb1d3b2febc12..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" @@ -38,7 +39,7 @@ limitations under the License. #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_verified_test_base.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" @@ -81,7 +82,7 @@ const std::vector GetInstructions(HloInstruction* root) { return main_list.GetInstructions(); } -class BufferAssignmentTest : public HloVerifiedTestBase { +class BufferAssignmentTest : public HloTestBase { protected: ~BufferAssignmentTest() override {} @@ -107,6 +108,24 @@ class BufferAssignmentTest : public HloVerifiedTestBase { .ConsumeValueOrDie(); } + std::unique_ptr RunBufferAssignmentNoBuffersReuseForAdd( + HloModule* module, int64 alignment = 1) { + auto reuse_checker = [](const BufferAssignment& assignment, + const BufferAllocation& alloc, + const LogicalBuffer& buffer) { + return (buffer.instruction()->opcode() != HloOpcode::kAdd); + }; + return BufferAssigner::Run( + module, absl::make_unique(module), + backend().compiler()->BufferSizeBytesFunction(), + [alignment](LogicalBuffer::Color) { return alignment; }, + /*allow_input_output_aliasing=*/false, + /*allocate_buffers_for_constants=*/false, + /*colorer=*/BufferLiveness::DefaultColorer(), + /*reuse_checker=*/reuse_checker) + .ConsumeValueOrDie(); + } + std::unique_ptr RunColoredBufferAssignment( HloModule* module, BufferLiveness::Colorer colorer, int64 alignment = 1) { return BufferAssigner::Run( @@ -119,8 +138,7 @@ class BufferAssignmentTest : public HloVerifiedTestBase { } std::unique_ptr RunBufferAssignmentWithInstructionSequence( - HloModule* module, - absl::Span instruction_sequence, + HloModule* module, absl::Span instruction_sequence, int64 alignment = 1) { HloSchedule schedule(module); schedule.set_sequence(module->entry_computation(), instruction_sequence); @@ -292,7 +310,7 @@ class BufferAssignmentTest : public HloVerifiedTestBase { 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( @@ -302,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; } } @@ -316,16 +334,16 @@ TEST_F(BufferAssignmentTest, ScalarConstant) { auto builder = HloComputation::Builder(TestName()); auto const0 = builder.AddInstruction( HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0))); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); { - auto buffers = RunBufferAssignment(module); + auto buffers = RunBufferAssignment(module.get()); EXPECT_TRUE(buffers->HasTopLevelAllocation(const0)); } { - auto buffers = RunBufferAssignmentNoBuffersForConstants(module); + auto buffers = RunBufferAssignmentNoBuffersForConstants(module.get()); EXPECT_FALSE(buffers->HasTopLevelAllocation(const0)); } } @@ -340,17 +358,17 @@ TEST_F(BufferAssignmentTest, BufferForConst) { LiteralUtil::CreateR1({4.1f, 4.2f, 4.3f, 4.4f}))); auto add = builder.AddInstruction( HloInstruction::CreateBinary(f32vec4_, HloOpcode::kAdd, const0, const1)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); { - auto buffers = RunBufferAssignment(module); + auto buffers = RunBufferAssignment(module.get()); EXPECT_TRUE(buffers->HasTopLevelAllocation(const0)); EXPECT_TRUE(buffers->HasTopLevelAllocation(const1)); GetAssignedOutputAllocation(*buffers, add); } { - auto buffers = RunBufferAssignmentNoBuffersForConstants(module); + auto buffers = RunBufferAssignmentNoBuffersForConstants(module.get()); EXPECT_FALSE(buffers->HasTopLevelAllocation(const0)); EXPECT_FALSE(buffers->HasTopLevelAllocation(const1)); GetAssignedOutputAllocation(*buffers, add); @@ -369,10 +387,10 @@ TEST_F(BufferAssignmentTest, HasAllocationAt) { HloInstruction::CreateUnary(f32vec100_, HloOpcode::kNegate, param0)); auto tuple = builder.AddInstruction( HloInstruction::CreateTuple({negate, param0, constant})); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); - auto buffers = RunBufferAssignment(module); + auto buffers = RunBufferAssignment(module.get()); // Make sure that HasAllocationAt() agrees with what HasTopLevelAllocation() // reports for the instruction directly. EXPECT_EQ(buffers->HasTopLevelAllocation(tuple), @@ -392,10 +410,10 @@ TEST_F(BufferAssignmentTest, BufferForOutputConst) { LiteralUtil::CreateR1({1.1f, 2.2f, 3.3f, 4.4f}))); auto copy = builder.AddInstruction( HloInstruction::CreateUnary(const0->shape(), HloOpcode::kCopy, const0)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); - auto buffers = RunBufferAssignment(module); + auto buffers = RunBufferAssignment(module.get()); // The copy node now has an output buffer. GetAssignedOutputAllocation(*buffers, copy); } @@ -421,10 +439,10 @@ TEST_F(BufferAssignmentTest, Basic) { HloInstruction::CreateBinary(f32vec100_, HloOpcode::kAdd, mul, param1)); auto sub = builder.AddInstruction(HloInstruction::CreateBinary( f32vec100_, HloOpcode::kSubtract, add, param1)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); - auto buffers = RunBufferAssignment(module); + auto buffers = RunBufferAssignment(module.get()); // Distinct input buffers were assigned for parameters. BufferAllocation paramscalar_buffer = @@ -447,6 +465,90 @@ 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. + // + // paramscalar ------- (mul) -- (add) -- (sub) + // / / / + // param0[100] -------/ / / + // / / + // param1[100] --------------/--------/ + auto builder = HloComputation::Builder(TestName()); + auto paramscalar = + builder.AddInstruction(HloInstruction::CreateParameter(0, r0f32_, "p")); + auto broadcast = builder.AddInstruction( + HloInstruction::CreateBroadcast(f32vec100_, paramscalar, {})); + auto param0 = builder.AddInstruction( + HloInstruction::CreateParameter(1, f32vec100_, "p1")); + auto param1 = builder.AddInstruction( + HloInstruction::CreateParameter(2, f32vec100_, "p2")); + auto mul = builder.AddInstruction(HloInstruction::CreateBinary( + f32vec100_, HloOpcode::kMultiply, broadcast, param0)); + auto add = builder.AddInstruction( + HloInstruction::CreateBinary(f32vec100_, HloOpcode::kAdd, mul, param1)); + auto sub = builder.AddInstruction(HloInstruction::CreateBinary( + f32vec100_, HloOpcode::kSubtract, add, param1)); + auto module = CreateNewVerifiedModule(); + module->AddEntryComputation(builder.Build()); + + auto buffers = RunBufferAssignmentNoBuffersReuseForAdd(module.get()); + + // Distinct input buffers were assigned for parameters. + BufferAllocation paramscalar_buffer = + GetAssignedInputAllocation(*buffers, paramscalar); + BufferAllocation param0_buffer = GetAssignedInputAllocation(*buffers, param0); + BufferAllocation param1_buffer = GetAssignedInputAllocation(*buffers, param1); + EXPECT_NE(paramscalar_buffer.index(), param0_buffer.index()); + EXPECT_NE(paramscalar_buffer.index(), param1_buffer.index()); + EXPECT_NE(param0_buffer.index(), param1_buffer.index()); + + // The mul node has a valid buffer assigned, doesn't share with input. + const BufferAllocation& mul_buffer = GetTopLevelAllocation(*buffers, mul); + EXPECT_NE(mul_buffer.index(), param0_buffer.index()); + + // The add node cannot reuse the mul node's buffer since we told buffer + // assignment so. + const BufferAllocation& add_buffer = GetTopLevelAllocation(*buffers, add); + EXPECT_NE(add_buffer.index(), mul_buffer.index()); + + // The sub node has a valid output buffer assigned. + GetAssignedOutputAllocation(*buffers, sub); +} + TEST_F(BufferAssignmentTest, BasicUniquelyColored) { // paramscalar ------- (mul) -- (add) -- (sub) // / / / @@ -470,7 +572,7 @@ TEST_F(BufferAssignmentTest, BasicUniquelyColored) { HloInstruction::CreateBinary(f32vec100_, HloOpcode::kAdd, mul, param1)); auto sub = builder.AddInstruction(HloInstruction::CreateBinary( f32vec100_, HloOpcode::kSubtract, add, param1)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); auto colorer = [](const BufferLiveness& buffer_liveness) { @@ -485,7 +587,7 @@ TEST_F(BufferAssignmentTest, BasicUniquelyColored) { return Status::OK(); }; - auto buffers = RunColoredBufferAssignment(module, colorer); + auto buffers = RunColoredBufferAssignment(module.get(), colorer); // Distinct input buffers were assigned for parameters. BufferAllocation paramscalar_buffer = @@ -531,7 +633,7 @@ TEST_F(BufferAssignmentTest, BasicPartiallyColored) { HloInstruction::CreateBinary(f32vec100_, HloOpcode::kAdd, mul, param1)); auto sub = builder.AddInstruction(HloInstruction::CreateBinary( f32vec100_, HloOpcode::kSubtract, add, param1)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); auto colorer = [](const BufferLiveness& buffer_liveness) { @@ -554,7 +656,7 @@ TEST_F(BufferAssignmentTest, BasicPartiallyColored) { return Status::OK(); }; - auto buffers = RunColoredBufferAssignment(module, colorer); + auto buffers = RunColoredBufferAssignment(module.get(), colorer); // Distinct input buffers were assigned for parameters. BufferAllocation paramscalar_buffer = @@ -603,10 +705,10 @@ TEST_F(BufferAssignmentTest, MultipleUsersForNode) { HloInstruction::CreateBinary(f32vec100_, HloOpcode::kAdd, mul, param1)); auto sub = builder.AddInstruction( HloInstruction::CreateBinary(f32vec100_, HloOpcode::kSubtract, add, mul)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); - auto buffers = RunBufferAssignment(module); + auto buffers = RunBufferAssignment(module.get()); // Input buffers were assigned for parameters. BufferAllocation paramscalar_buffer = @@ -638,7 +740,7 @@ TEST_F(BufferAssignmentTest, TrivialMap) { // param0[100x10] ---> (map x+1) // // Builds the map function. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto map_computation = module->AddEmbeddedComputation(BuildMapComputationPlus1("f32+1")); auto inner_last = map_computation->root_instruction(); @@ -657,7 +759,7 @@ TEST_F(BufferAssignmentTest, TrivialMap) { EXPECT_EQ(3, level1.size()) << "Invalid nested add+1 size"; // Assigns buffers and fetches sizes. - auto buffers = RunBufferAssignment(module); + auto buffers = RunBufferAssignment(module.get()); int64 size0 = ValidateBuffers(level0, *buffers); int64 size1 = ValidateBuffers(level1, *buffers); @@ -693,7 +795,7 @@ TEST_F(BufferAssignmentTest, CannotReuseInputBufferOfReduce) { // out-of-order reductions could overwrite an element before a use.) // // param0[100] --- (exp1) --- (exp2) --- (reduce x+y) --- (exp3) - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto reduce_computation = module->AddEmbeddedComputation(BuildReduceComputation("f32+f32")); @@ -716,7 +818,7 @@ TEST_F(BufferAssignmentTest, CannotReuseInputBufferOfReduce) { module->AddEntryComputation(builder.Build()); - auto buffers = RunBufferAssignment(module); + auto buffers = RunBufferAssignment(module.get()); const std::vector instrs = GetInstructions(exp3); ValidateBuffers(instrs, *buffers); @@ -744,7 +846,7 @@ TEST_F(BufferAssignmentTest, ExampleWhile) { // const4[f32[4]] --- tuple --- while[condition, body] // // Builds the nested condition and body. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto condition_computation = module->AddEmbeddedComputation(BuildWhileConditionComputation("if<4")); auto body_computation = @@ -772,7 +874,7 @@ TEST_F(BufferAssignmentTest, ExampleWhile) { EXPECT_EQ(8, levelb.size()) << "Invalid nested body size"; // Assigns buffers and fetches sizes. - auto buffers = RunBufferAssignment(module); + auto buffers = RunBufferAssignment(module.get()); int64 size0 = ValidateBuffers(level0, *buffers); int64 sizec = ValidateBuffers(levelc, *buffers); int64 sizeb = ValidateBuffers(levelb, *buffers); @@ -810,7 +912,7 @@ TEST_F(BufferAssignmentTest, ExampleWhile) { } TEST_F(BufferAssignmentTest, ExampleConditional) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto true_computation = module->AddEmbeddedComputation( BuildR0F32UnaryOpComputation(HloOpcode::kCeil, "Ceil")); auto false_computation = module->AddEmbeddedComputation( @@ -837,7 +939,7 @@ TEST_F(BufferAssignmentTest, ExampleConditional) { EXPECT_EQ(2, true_instrs.size()); EXPECT_EQ(2, false_instrs.size()); - auto buffers = RunBufferAssignment(module); + auto buffers = RunBufferAssignment(module.get()); ValidateBuffers(conditional_instrs, *buffers); ValidateBuffers(true_instrs, *buffers); ValidateBuffers(false_instrs, *buffers); @@ -873,9 +975,9 @@ TEST_F(BufferAssignmentTest, UnaryOpReuseChain) { auto neg = builder.AddInstruction( HloInstruction::CreateUnary(f32vec100_, HloOpcode::kNegate, exp2)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); - auto assignment = RunBufferAssignment(module); + auto assignment = RunBufferAssignment(module.get()); // tanh and exp2 can reuse exp1's buffer EXPECT_TRUE(assignment->HasTopLevelAllocation(exp1)); @@ -902,9 +1004,9 @@ TEST_F(BufferAssignmentTest, ReuseNonOperandBuffer) { auto broadcast = builder.AddInstruction( HloInstruction::CreateBroadcast(f32a100x10_, slice, {1})); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); - auto assignment = RunBufferAssignment(module); + auto assignment = RunBufferAssignment(module.get()); // negate and broadcast should share a buffer. EXPECT_TRUE(assignment->HasTopLevelAllocation(broadcast)); @@ -935,9 +1037,9 @@ TEST_F(BufferAssignmentTest, NoReuseLiveBuffer) { HloInstruction::CreateBroadcast(f32a100x10_, slice, {1})); builder.AddInstruction(HloInstruction::CreateTuple({negate, broadcast})); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); - auto assignment = RunBufferAssignment(module); + auto assignment = RunBufferAssignment(module.get()); // The instructions should not share buffers. EXPECT_NE(GetTopLevelAllocation(*assignment, broadcast), @@ -972,9 +1074,9 @@ TEST_F(BufferAssignmentTest, NoReuseAliasedBuffer) { HloInstruction::CreateBroadcast(f32a100x10_, slice, {1})); builder.AddInstruction(HloInstruction::CreateTuple({tuple, broadcast})); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); - auto assignment = RunBufferAssignment(module); + auto assignment = RunBufferAssignment(module.get()); // The instructions should not share buffers. EXPECT_NE(GetTopLevelAllocation(*assignment, broadcast), @@ -1007,9 +1109,9 @@ TEST_F(BufferAssignmentTest, DoNotReuseOversizedOutputBuffer) { auto broadcast = builder.AddInstruction(HloInstruction::CreateBroadcast( ShapeUtil::MakeShape(F32, {10, 4}), slice, {0})); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); - auto assignment = RunBufferAssignment(module); + auto assignment = RunBufferAssignment(module.get()); // The broadcast output buffer cannot be shared. EXPECT_NE(GetTopLevelAllocation(*assignment, broadcast), @@ -1039,9 +1141,9 @@ TEST_F(BufferAssignmentTest, ReuseOutputBufferIfExactlySized) { auto broadcast = builder.AddInstruction(HloInstruction::CreateBroadcast( ShapeUtil::MakeShape(F32, {10, 10}), slice, {0})); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); - auto assignment = RunBufferAssignment(module); + auto assignment = RunBufferAssignment(module.get()); // negate and broadcast should share a buffer. EXPECT_TRUE(assignment->HasTopLevelAllocation(broadcast)); @@ -1077,9 +1179,9 @@ TEST_F(BufferAssignmentTest, DoNotReuseOversizedOutputBufferInTuple) { ShapeUtil::MakeShape(F32, {10, 4}), slice, {0})); builder.AddInstruction(HloInstruction::CreateTuple({broadcast})); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); - auto assignment = RunBufferAssignment(module); + auto assignment = RunBufferAssignment(module.get()); // The broadcast output buffer cannot be shared. EXPECT_NE(GetTopLevelAllocation(*assignment, broadcast), @@ -1092,7 +1194,7 @@ TEST_F(BufferAssignmentTest, EmbeddedComputationBuffers) { // Verify that buffers for embedded computations are properly marked as // thread-local and that embedded parameters are not marked as // is_entry_computation_parameter. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto vec_shape = ShapeUtil::MakeShape(F32, {42}); auto scalar_shape = ShapeUtil::MakeShape(F32, {}); @@ -1123,7 +1225,7 @@ TEST_F(BufferAssignmentTest, EmbeddedComputationBuffers) { HloInstruction::CreateMap(vec_shape, {call}, map_computation)); module->AddEntryComputation(builder.Build()); - auto assignment = RunBufferAssignment(module); + auto assignment = RunBufferAssignment(module.get()); // Allocations for the map computation should be thread-local and not // live-out. @@ -1170,9 +1272,9 @@ TEST_F(BufferAssignmentTest, TupleParameterAsOutput) { ShapeUtil::MakeShape(S32, {42})}), "param0")); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); - auto assignment = RunBufferAssignment(module); + auto assignment = RunBufferAssignment(module.get()); // There should be four allocations: one for vector of pointers, and one for // each tuple element. @@ -1206,9 +1308,9 @@ TEST_F(BufferAssignmentTest, ElementOfNestedTupleParameterAsOutput) { builder.AddInstruction(HloInstruction::CreateGetTupleElement( ShapeUtil::GetSubshape(tuple_param->shape(), {1}), tuple_param, 1)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); - auto assignment = RunBufferAssignment(module); + auto assignment = RunBufferAssignment(module.get()); // Only some of the elements of the input param are liveout. EXPECT_FALSE( @@ -1250,9 +1352,9 @@ TEST_F(BufferAssignmentTest, TupleConstantAsOutput) { builder.AddInstruction(HloInstruction::CreateConstant( LiteralUtil::MakeTuple({&elements[0], &elements[1]}))); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); - auto assignment = RunBufferAssignment(module); + auto assignment = RunBufferAssignment(module.get()); EXPECT_EQ(3, assignment->Allocations().size()); } @@ -1264,9 +1366,9 @@ TEST_F(BufferAssignmentTest, TupleCustomCallAsOutput) { ShapeUtil::MakeTupleShape({ShapeUtil::MakeShape(PRED, {1, 2, 3, 4}), ShapeUtil::MakeShape(S32, {101})}), /*operands=*/{}, /*custom_call_target=*/"foo_function")); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); - auto assignment = RunBufferAssignment(module); + auto assignment = RunBufferAssignment(module.get()); EXPECT_EQ(3, assignment->Allocations().size()); EXPECT_TRUE( @@ -1279,7 +1381,7 @@ TEST_F(BufferAssignmentTest, TupleCustomCallAsOutput) { TEST_F(BufferAssignmentTest, TupleCallAsOutput) { // Test a computation which returns a tuple call value. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto elem_shape = f32vec4_; auto tuple_shape = ShapeUtil::MakeTupleShape({elem_shape}); @@ -1297,7 +1399,7 @@ TEST_F(BufferAssignmentTest, TupleCallAsOutput) { HloInstruction::CreateCall(tuple_shape, {param}, sub_computation)); module->AddEntryComputation(builder.Build()); - auto assignment = RunBufferAssignment(module); + auto assignment = RunBufferAssignment(module.get()); EXPECT_EQ(2, assignment->Allocations().size()); // Buffers for call are colocated with the sub-computation. @@ -1320,7 +1422,7 @@ TEST_F(BufferAssignmentTest, TupleChainedCallAsOutput) { // B: call(C, param) // C: call(D, param) // D: param - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto elem_shape = f32vec4_; auto tuple_shape = ShapeUtil::MakeTupleShape({elem_shape}); @@ -1359,7 +1461,7 @@ TEST_F(BufferAssignmentTest, TupleChainedCallAsOutput) { module->AddEntryComputation(std::move(a_computation)); module->AddEmbeddedComputation(std::move(b_computation)); - auto assignment = RunBufferAssignment(module); + auto assignment = RunBufferAssignment(module.get()); // Buffers for call are colocated with the sub-computations. EXPECT_EQ(GetAllocation(*assignment, a_call, /*index=*/{}), @@ -1393,9 +1495,9 @@ TEST_F(BufferAssignmentTest, BitcastAsOutput) { auto bitcast = builder.AddInstruction( HloInstruction::CreateUnary(param->shape(), HloOpcode::kBitcast, param)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); - auto assignment = RunBufferAssignment(module); + auto assignment = RunBufferAssignment(module.get()); // Bitcast should get the same allocation as the param. EXPECT_EQ(1, assignment->Allocations().size()); @@ -1420,9 +1522,9 @@ TEST_F(BufferAssignmentTest, AmbiguousBufferAsOutput) { HloInstruction::CreateTernary(tuple_shape, HloOpcode::kTupleSelect, pred_param, tuple_param0, tuple_param1)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); - auto assignment = RunBufferAssignment(module); + auto assignment = RunBufferAssignment(module.get()); // Select shallow copies one of its operands so it defines its own top-level // buffer and receives its own allocation. @@ -1458,9 +1560,9 @@ TEST_F(BufferAssignmentTest, TupleBufferNotReused) { auto copy = builder.AddInstruction(HloInstruction::CreateUnary( scalar_shape, HloOpcode::kCopy, tuple_element)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); - auto assignment = RunBufferAssignment(module); + auto assignment = RunBufferAssignment(module.get()); // There should be no buffer reuse. The copy should not reuse the tuple // buffer. @@ -1500,9 +1602,9 @@ TEST_F(BufferAssignmentTest, OneTempAllocation) { HloInstruction::CreateConcatenate(shape_5x4, {dot_ab, dot_bc}, 0)); // Run buffer assignment with alignment=1. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); - auto assignment = RunBufferAssignment(module, /*alignment=*/1); + auto assignment = RunBufferAssignment(module.get(), /*alignment=*/1); // There are 5 allocations: 3 parameters, 1 output, and 1 temp. EXPECT_EQ(5, assignment->Allocations().size()); @@ -1521,7 +1623,7 @@ TEST_F(BufferAssignmentTest, OneTempAllocation) { EXPECT_EQ(80, slice_bc.allocation()->size()); // Re-run buffer assignment with alignment=64. - assignment = RunBufferAssignment(module, /*alignment=*/64); + assignment = RunBufferAssignment(module.get(), /*alignment=*/64); EXPECT_EQ(5, assignment->Allocations().size()); slice_ab = assignment->GetUniqueTopLevelSlice(dot_ab).ConsumeValueOrDie(); slice_bc = assignment->GetUniqueTopLevelSlice(dot_bc).ConsumeValueOrDie(); @@ -1564,10 +1666,10 @@ TEST_F(BufferAssignmentTest, TrivialPeakBuffers) { HloInstruction::CreateBinary(f32vec100_, HloOpcode::kAdd, mul, param1)); builder.AddInstruction(HloInstruction::CreateBinary( f32vec100_, HloOpcode::kSubtract, add, param1)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); - auto buffers = RunBufferAssignment(module); + auto buffers = RunBufferAssignment(module.get()); const BufferAllocation& mul_buffer = GetTopLevelAllocation(*buffers, mul); const std::vector& peak_buffers = @@ -1605,11 +1707,11 @@ TEST_F(BufferAssignmentTest, PeakBuffers) { ShapeUtil::MakeShape(F32, {1}), concat, {0}, {1}, {1})); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); auto buffers = RunBufferAssignmentWithInstructionSequence( - module, {param, log, rev, neg, concat, root}); + module.get(), {param, log, rev, neg, concat, root}); // The temporary buffer should hold the 4 interior instructions. const BufferAllocation& buffer = GetTopLevelAllocation(*buffers, concat); @@ -1630,7 +1732,7 @@ TEST_F(BufferAssignmentTest, PeakBuffers) { } TEST_F(BufferAssignmentTest, PeakBuffersWhile) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); const Shape shape = ShapeUtil::MakeShape(F32, {123, 123}); HloComputation* condition; { @@ -1665,7 +1767,7 @@ TEST_F(BufferAssignmentTest, PeakBuffersWhile) { ShapeUtil::MakeShape(F32, {123, 123, 123}), bcast, {0})); module->AddEntryComputation(builder.Build()); - auto buffers = RunBufferAssignment(module); + auto buffers = RunBufferAssignment(module.get()); const BufferAllocation& buffer = GetTopLevelAllocation(*buffers, bcast); const std::vector& peak_buffers = buffer.PeakMemoryLogicalBuffers(); @@ -1715,13 +1817,13 @@ ENTRY main { } )"; - ParseAndVerifyModule(hlo_text); + TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(hlo_text)); HloInstruction* constant_1 = - module().entry_computation()->GetInstructionWithName("constant.1.1"); + m->entry_computation()->GetInstructionWithName("constant.1.1"); HloInstruction* constant_2 = - module().entry_computation()->GetInstructionWithName("constant.1.2"); + m->entry_computation()->GetInstructionWithName("constant.1.2"); - auto buffers = RunBufferAssignment(&module()); + auto buffers = RunBufferAssignment(m.get()); { const BufferAllocation& allocation_for_const_1 = @@ -1750,7 +1852,7 @@ ENTRY main { } } -class WhileBufferAssignmentTest : public HloVerifiedTestBase { +class WhileBufferAssignmentTest : public HloTestBase { protected: std::unique_ptr BuildWhileConditionComputation( const string& name) { @@ -1785,7 +1887,7 @@ class WhileBufferAssignmentTest : public HloVerifiedTestBase { std::unique_ptr RunBufferAssignment(HloModule* module, int64 alignment = 1) { HloSchedule schedule = - ScheduleModule(*module, ByteSizeOf).ConsumeValueOrDie(); + ScheduleModule(module, ByteSizeOf).ConsumeValueOrDie(); return BufferAssigner::Run( module, absl::make_unique(schedule), ByteSizeOf, @@ -1810,7 +1912,7 @@ static void RunCopyInsertion(HloModule* module) { } TEST_F(WhileBufferAssignmentTest, TwoForwardWhileLoops) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto builder = HloComputation::Builder("entry"); auto input0 = builder.AddInstruction( @@ -1849,8 +1951,8 @@ TEST_F(WhileBufferAssignmentTest, TwoForwardWhileLoops) { HloInstruction::CreateWhile(loop_state_shape_, cond1, body1, tuple1)); module->AddEntryComputation(builder.Build()); - RunCopyInsertion(module); - auto assignment = RunBufferAssignment(module); + RunCopyInsertion(module.get()); + auto assignment = RunBufferAssignment(module.get()); // Verify 'input0' and read-only use while0{0} alias. EXPECT_EQ(assignment->GetUniqueSlice(input0, {}).ConsumeValueOrDie(), @@ -1906,20 +2008,19 @@ ENTRY %test_module { ROOT %bcast = s32[1024,1024]{1,0} broadcast(s32[] %while.1), dimensions={} })"; - ParseAndVerifyModule(module_str); + TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(module_str)); // Run CopyInsertion and check if the graph constructed above doesn't need // any copies inserted for BufferAssignment to run. - int64 instruction_count = module().instruction_count(); + int64 instruction_count = m->instruction_count(); CopyInsertion copy_insertion; - ASSERT_IS_OK(copy_insertion.Run(&module()).status()); - ASSERT_EQ(instruction_count, module().instruction_count()); + ASSERT_IS_OK(copy_insertion.Run(m.get()).status()); + ASSERT_EQ(instruction_count, m->instruction_count()); // Get the instructions in the module. - const HloInstruction* bcast = - module().entry_computation()->root_instruction(); + const HloInstruction* bcast = m->entry_computation()->root_instruction(); const HloInstruction* param = - module().entry_computation()->parameter_instruction(0); + m->entry_computation()->parameter_instruction(0); ASSERT_EQ(bcast->opcode(), HloOpcode::kBroadcast); const HloInstruction* while1 = bcast->operand(0); ASSERT_EQ(while1->opcode(), HloOpcode::kWhile); @@ -1927,7 +2028,7 @@ ENTRY %test_module { ASSERT_EQ(while0->opcode(), HloOpcode::kWhile); // Run buffer assignment. - auto assignment = RunBufferAssignment(&module()); + auto assignment = RunBufferAssignment(m.get()); TF_ASSERT_OK_AND_ASSIGN(auto slice_param, assignment->GetUniqueSlice(param, {})); TF_ASSERT_OK_AND_ASSIGN(auto slice_while0, @@ -1974,20 +2075,19 @@ ENTRY %test_module { ROOT %bcast = s32[1024,1024]{1,0} broadcast(s32[] %while.1), dimensions={} })"; - ParseAndVerifyModule(module_str); + TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(module_str)); // Run CopyInsertion and check if the graph constructed above doesn't need // any copies inserted for BufferAssignment to run. - int64 instruction_count = module().instruction_count(); + int64 instruction_count = m->instruction_count(); CopyInsertion copy_insertion; - ASSERT_IS_OK(copy_insertion.Run(&module()).status()); - ASSERT_EQ(instruction_count, module().instruction_count()); + ASSERT_IS_OK(copy_insertion.Run(m.get()).status()); + ASSERT_EQ(instruction_count, m->instruction_count()); // Get the instructions in the module. - const HloInstruction* bcast = - module().entry_computation()->root_instruction(); + const HloInstruction* bcast = m->entry_computation()->root_instruction(); const HloInstruction* constant = - module().entry_computation()->GetInstructionWithName("constant.42"); + m->entry_computation()->GetInstructionWithName("constant.42"); ASSERT_EQ(bcast->opcode(), HloOpcode::kBroadcast); const HloInstruction* while1 = bcast->operand(0); ASSERT_EQ(while1->opcode(), HloOpcode::kWhile); @@ -1995,7 +2095,7 @@ ENTRY %test_module { ASSERT_EQ(while0->opcode(), HloOpcode::kWhile); // Run buffer assignment. - auto assignment = RunBufferAssignment(&module()); + auto assignment = RunBufferAssignment(m.get()); TF_ASSERT_OK_AND_ASSIGN(auto slice_constant, assignment->GetUniqueSlice(constant, {})); TF_ASSERT_OK_AND_ASSIGN(auto slice_while0, @@ -2053,7 +2153,7 @@ TEST_F(WhileBufferAssignmentTest, ColocatedBuffers) { }; // Build the entry computation as described in the comment above. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto builder = HloComputation::Builder("entry"); auto token = builder.AddInstruction(HloInstruction::CreateToken()); @@ -2088,7 +2188,7 @@ TEST_F(WhileBufferAssignmentTest, ColocatedBuffers) { // any copies inserted for BufferAssignment to run. int64 instruction_count = module->instruction_count(); CopyInsertion copy_insertion; - ASSERT_IS_OK(copy_insertion.Run(module).status()); + ASSERT_IS_OK(copy_insertion.Run(module.get()).status()); ASSERT_EQ(instruction_count, module->instruction_count()); // Create a sequential order among all the instructions in the entry @@ -2096,7 +2196,7 @@ TEST_F(WhileBufferAssignmentTest, ColocatedBuffers) { // nodes are traversed during BufferAssignment. TF_ASSERT_OK_AND_ASSIGN( HloSchedule schedule, - ScheduleModule(*module, [](const BufferValue& buffer) { + ScheduleModule(module.get(), [](const BufferValue& buffer) { return ShapeUtil::ByteSizeOf(buffer.shape(), /*pointer_size=*/sizeof(void*)); })); @@ -2107,12 +2207,12 @@ TEST_F(WhileBufferAssignmentTest, ColocatedBuffers) { TF_ASSERT_OK_AND_ASSIGN( auto assignment, - BufferAssigner::Run(module, - absl::make_unique(schedule), - backend().compiler()->BufferSizeBytesFunction(), - [](LogicalBuffer::Color) { return 1; }, - /*allow_input_output_aliasing=*/false, - /*allocate_buffers_for_constants=*/true)); + BufferAssigner::Run( + module.get(), absl::make_unique(schedule), + backend().compiler()->BufferSizeBytesFunction(), + [](LogicalBuffer::Color) { return 1; }, + /*allow_input_output_aliasing=*/false, + /*allocate_buffers_for_constants=*/true)); // The result tuple elements must be assigned with different buffers. TF_ASSERT_OK_AND_ASSIGN(auto slice0, assignment->GetUniqueSlice(tuple, {0})); @@ -2134,7 +2234,7 @@ TEST_F(WhileBufferAssignmentTest, ColocatedBuffers) { } TEST_F(WhileBufferAssignmentTest, OneForwardBackwardWhileLoopSet) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto builder = HloComputation::Builder("entry"); auto input0 = builder.AddInstruction( @@ -2166,8 +2266,8 @@ TEST_F(WhileBufferAssignmentTest, OneForwardBackwardWhileLoopSet) { HloInstruction::CreateWhile(loop_state_shape_, cond1, body1, while0)); module->AddEntryComputation(builder.Build()); - RunCopyInsertion(module); - auto assignment = RunBufferAssignment(module); + RunCopyInsertion(module.get()); + auto assignment = RunBufferAssignment(module.get()); // while0 and while1 buffers should be completely aligned. EXPECT_EQ(assignment->GetUniqueSlice(while0, {0}).ConsumeValueOrDie(), @@ -2179,7 +2279,7 @@ TEST_F(WhileBufferAssignmentTest, OneForwardBackwardWhileLoopSet) { } TEST_F(BufferAssignmentTest, TwoCalls) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); Shape r0f32 = ShapeUtil::MakeShape(xla::F32, {}); HloComputation* sub_computation; { @@ -2209,13 +2309,13 @@ TEST_F(BufferAssignmentTest, TwoCalls) { { FlattenCallGraph flatten; - TF_ASSERT_OK_AND_ASSIGN(bool result, flatten.Run(module)); + TF_ASSERT_OK_AND_ASSIGN(bool result, flatten.Run(module.get())); EXPECT_TRUE(result); - std::unique_ptr call_graph = CallGraph::Build(module); + std::unique_ptr call_graph = CallGraph::Build(module.get()); } - RunCopyInsertion(module); - auto assignment = RunBufferAssignment(module); + RunCopyInsertion(module.get()); + auto assignment = RunBufferAssignment(module.get()); EXPECT_TRUE(BuffersDistinct({call1}, {call2}, *assignment)); } @@ -2240,13 +2340,14 @@ ENTRY Main { )"; HloModuleConfig config; - config.set_debug_options(legacy_flags::GetDebugOptionsFromFlags()); - ParseAndVerifyModule(hlo_text, config); + config.set_debug_options(GetDebugOptionsFromFlags()); + TF_ASSERT_OK_AND_ASSIGN(auto m, + ParseAndReturnVerifiedModule(hlo_text, config)); - auto buffers = RunBufferAssignment(&module()); + auto buffers = RunBufferAssignment(m.get()); - HloComputation* main = module().entry_computation(); - HloComputation* callee = module().GetComputationWithName("Callee"); + HloComputation* main = m->entry_computation(); + HloComputation* callee = m->GetComputationWithName("Callee"); EXPECT_NE(callee, nullptr); HloInstruction* param0 = callee->parameter_instruction(0); @@ -2270,7 +2371,7 @@ ENTRY Main { } TEST_F(WhileBufferAssignmentTest, WhileLoopsInterferingResultRange) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto builder = HloComputation::Builder(TestName()); auto zero = builder.AddInstruction( @@ -2317,40 +2418,41 @@ TEST_F(WhileBufferAssignmentTest, WhileLoopsInterferingResultRange) { { FlattenCallGraph flatten; - TF_ASSERT_OK_AND_ASSIGN(bool result, flatten.Run(module)); + TF_ASSERT_OK_AND_ASSIGN(bool result, flatten.Run(module.get())); EXPECT_TRUE(result); } - RunCopyInsertion(module); + RunCopyInsertion(module.get()); HloSchedule schedule = - ScheduleModule(*module, ByteSizeOf).ConsumeValueOrDie(); + ScheduleModule(module.get(), ByteSizeOf).ConsumeValueOrDie(); // To trigger b/38494731, we want a specific Hlo schedule for the // root computation, so we overwrite that entry with a manually // crafted sequence. - schedule.set_sequence(module->entry_computation(), - {input1, weights1, one, output1, while1->operand(0), - while1, input0, weights0, zero, output0, - while0->operand(0), while0, gte0, gte1, root_add}); + schedule.set_sequence( + module->entry_computation(), + {input1, weights1, one, output1, while1->mutable_operand(0), while1, + input0, weights0, zero, output0, while0->mutable_operand(0), while0, + gte0, gte1, root_add}); // If this ASSERT fails, we constructed a bogus sequence above and this test // itself is buggy. TF_ASSERT_OK(schedule.Verify()); auto assignment = - BufferAssigner::Run(module, - absl::make_unique(schedule), - ByteSizeOf, [](LogicalBuffer::Color) { return 1; }, - /*allow_input_output_aliasing=*/false, - /*allocate_buffers_for_constants=*/true) + BufferAssigner::Run( + module.get(), absl::make_unique(schedule), + ByteSizeOf, [](LogicalBuffer::Color) { return 1; }, + /*allow_input_output_aliasing=*/false, + /*allocate_buffers_for_constants=*/true) .ConsumeValueOrDie(); EXPECT_TRUE(BuffersDistinct({while0}, {while1}, *assignment)); } TEST_F(WhileBufferAssignmentTest, WhilesDontShareEntryParamIfLiveOut) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto builder = HloComputation::Builder("entry"); auto input0 = builder.AddInstruction( @@ -2394,8 +2496,8 @@ TEST_F(WhileBufferAssignmentTest, WhilesDontShareEntryParamIfLiveOut) { HloInstruction::CreateGetTupleElement(data_shape_, while1, 2)); module->AddEntryComputation(builder.Build()); - RunCopyInsertion(module); - auto assignment = RunBufferAssignment(module); + RunCopyInsertion(module.get()); + auto assignment = RunBufferAssignment(module.get()); // Get BufferAllocation for root instruction. auto* root_alloc = assignment->GetUniqueTopLevelSlice(while1_out) .ConsumeValueOrDie() @@ -2406,5 +2508,58 @@ TEST_F(WhileBufferAssignmentTest, WhilesDontShareEntryParamIfLiveOut) { EXPECT_FALSE(root_alloc->is_entry_computation_parameter()); } +TEST_F(WhileBufferAssignmentTest, WhileWithDynamicUpdateSliceShare) { + const char* const hlo_string = R"( +HloModule test + +while_body { + state = (s32[], f32[1280,1,128]{2,1,0}) parameter(0) + constant.1 = f32[] constant(0) + broadcast.6 = f32[128,1,128]{2,1,0} broadcast(constant.1), dimensions={} + get-tuple-element.4 = f32[1280,1,128]{2,1,0} get-tuple-element(state), index=1 + 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[] 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) +} + +while_condition { + state = (s32[], f32[1280,1,128]{2,1,0}) parameter(0) + get-tuple-element = s32[] get-tuple-element(state), index=0 + get-tuple-element.1 = s32[] constant(3) + ROOT less-than.339.338 = pred[] less-than(get-tuple-element, get-tuple-element.1) +} + +ENTRY entry_computation { + constant.7 = s32[] constant(0) + copy.1 = s32[] copy(constant.7) + constant.6 = f32[] constant(0) + broadcast.6 = f32[1280,1,128]{2,1,0} broadcast(constant.6), dimensions={} + tuple.1 = (s32[], f32[1280,1,128]{2,1,0}) tuple(copy.1, broadcast.6) + while.0 = (s32[], f32[1280,1,128]{2,1,0}) while(tuple.1), condition=while_condition, body=while_body + ROOT get-tuple-element.2 = s32[] get-tuple-element(while.0), index=0 +} + +)"; + auto module_or_status = + HloRunner::CreateModuleFromString(hlo_string, GetDebugOptionsForTest()); + auto module = module_or_status.ConsumeValueOrDie(); + + RunCopyInsertion(module.get()); + auto assignment = RunBufferAssignment(module.get()); + // Get BufferAllocation for root instruction. + auto dus9 = FindInstruction(module.get(), "dynamic-update-slice.9"); + auto dus9_alloc_slice = + assignment->GetUniqueTopLevelSlice(dus9).ConsumeValueOrDie(); + auto dus5 = FindInstruction(module.get(), "dynamic-update-slice.5"); + auto dus5_alloc_slice = + assignment->GetUniqueTopLevelSlice(dus5).ConsumeValueOrDie(); + // Test that the two dynamic-update-slice ops share the same allocation slice. + EXPECT_EQ(dus9_alloc_slice.allocation(), dus5_alloc_slice.allocation()); + EXPECT_EQ(dus9_alloc_slice, dus5_alloc_slice); +} } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/service/buffer_liveness_test.cc b/tensorflow/compiler/xla/service/buffer_liveness_test.cc index 17e50905059ad2c92784d14132c1cb1f46f35ade..23b9af0281b0d5ee1ef6ca2315f0cc1042285609 100644 --- a/tensorflow/compiler/xla/service/buffer_liveness_test.cc +++ b/tensorflow/compiler/xla/service/buffer_liveness_test.cc @@ -52,8 +52,8 @@ class BufferLivenessTest : public HloTestBase { // interfere. Precondition: 'a' and 'b' are array-shaped. bool InstructionsMayInterfere(const BufferLiveness& liveness, HloInstruction* a, HloInstruction* b) { - EXPECT_FALSE(ShapeUtil::IsTuple(a->shape())); - EXPECT_FALSE(ShapeUtil::IsTuple(b->shape())); + EXPECT_FALSE(a->shape().IsTuple()); + EXPECT_FALSE(b->shape().IsTuple()); return liveness.MayInterfere( GetBuffer(liveness, /*instruction=*/a, /*index=*/{}), GetBuffer(liveness, /*instruction=*/b, /*index=*/{})); @@ -66,8 +66,8 @@ class BufferLivenessTest : public HloTestBase { HloInstruction* a, HloInstruction* b, const ShapeIndex& index) { // Check that top-level shapes are tuple and tuple element shapes are equal. - EXPECT_TRUE(ShapeUtil::IsTuple(a->shape())); - EXPECT_TRUE(ShapeUtil::IsTuple(b->shape())); + EXPECT_TRUE(a->shape().IsTuple()); + EXPECT_TRUE(b->shape().IsTuple()); EXPECT_TRUE( ShapeUtil::Compatible(ShapeUtil::GetSubshape(a->shape(), index), ShapeUtil::GetSubshape(b->shape(), index))); @@ -117,7 +117,7 @@ TEST_F(BufferLivenessTest, ElementwiseChain) { auto log = builder.AddInstruction( HloInstruction::CreateUnary(vec_, HloOpcode::kLog, exp)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); auto liveness = @@ -164,7 +164,7 @@ TEST_F(BufferLivenessTest, MultipleEntryParameters_Sequential) { auto add = builder.AddInstruction( HloInstruction::CreateBinary(vec_, HloOpcode::kAdd, negate, exp)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* entry = module->AddEntryComputation(builder.Build()); HloSchedule schedule(module.get()); @@ -213,7 +213,7 @@ TEST_F(BufferLivenessTest, NonElementwiseOperand) { auto reverse = builder.AddInstruction(HloInstruction::CreateReverse(vec_, negate, {0})); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); auto liveness = @@ -247,7 +247,7 @@ TEST_F(BufferLivenessTest, OverlappedBuffers) { auto add = builder.AddInstruction( HloInstruction::CreateBinary(vec_, HloOpcode::kAdd, negate, exp)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); auto liveness = @@ -289,7 +289,7 @@ TEST_F(BufferLivenessTest, OverlappedBuffersSequentialOrder) { auto add = builder.AddInstruction( HloInstruction::CreateBinary(vec_, HloOpcode::kAdd, negate, exp)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); HloSchedule schedule(module.get()); @@ -336,7 +336,7 @@ TEST_F(BufferLivenessTest, RootInstructionIsNotLastInSequentialOrder) { HloInstruction::CreateSend(recv_done, token, /*channel_id=*/1)); auto send_done = builder.AddInstruction(HloInstruction::CreateSendDone(send)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build(add)); HloSchedule schedule(module.get()); @@ -373,7 +373,7 @@ TEST_F(BufferLivenessTest, TupleLiveOut) { auto outer_tuple = builder.AddInstruction(HloInstruction::CreateTuple({inner_tuple, exp})); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); auto liveness = @@ -393,7 +393,7 @@ TEST_F(BufferLivenessTest, TupleLiveOut) { TEST_F(BufferLivenessTest, EmbeddedComputation) { // Test MaybeLiveOut and MayInterfere for embedded computation. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto embedded_builder = HloComputation::Builder(TestName() + "_embedded"); auto embedded_param = embedded_builder.AddInstruction( @@ -450,7 +450,7 @@ TEST_F(BufferLivenessTest, TupleConstantLiveOut) { builder.AddInstruction(HloInstruction::CreateGetTupleElement( inner_tuple0.shape(), tuple_constant, 0)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); auto liveness = @@ -514,7 +514,7 @@ TEST_F(BufferLivenessTest, IndependentTupleElements) { auto tuple_root = builder.AddInstruction(HloInstruction::CreateTuple({add0, add1})); - auto module = CreateNewModule(); + auto module = CreateNewUnverifiedModule(); module->AddEntryComputation(BuildDummyComputation()); module->AddEmbeddedComputation(builder.Build()); @@ -576,7 +576,7 @@ TEST_F(BufferLivenessTest, DependentTupleElements) { auto tuple_root = builder.AddInstruction(HloInstruction::CreateTuple({add0, add1})); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(BuildDummyComputation()); module->AddEmbeddedComputation(builder.Build()); @@ -611,8 +611,8 @@ TEST_F(BufferLivenessTest, DependentTupleElements) { class FusedDynamicUpdateSliceLivenessTest : public BufferLivenessTest { protected: // Builds and runs a computation (see test case computation graphs below). - std::unique_ptr BuildModule(const bool update_uses_tuple_element1, - const bool fuse_gte0) { + std::unique_ptr BuildModule( + const bool update_uses_tuple_element1, const bool fuse_gte0) { auto builder = HloComputation::Builder(TestName()); // Create param0 Tuple. Shape data_shape = ShapeUtil::MakeShape(F32, {8}); @@ -638,15 +638,15 @@ 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})); // Build module and get reference to entry computation. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); auto* computation = module->entry_computation(); // Create fusion instruction based on number of tuple element 1 users. @@ -794,15 +794,15 @@ 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})); // Build module and get reference to entry computation. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(BuildDummyComputation()); module->AddEmbeddedComputation(builder.Build()); // Run BufferLiveness on 'module'. diff --git a/tensorflow/compiler/xla/service/buffer_value.cc b/tensorflow/compiler/xla/service/buffer_value.cc index fdf822c666b15afbc7553ca89d4f92ab08201869..b1abba20689915b03304aacd7a5fcca5443c2c60 100644 --- a/tensorflow/compiler/xla/service/buffer_value.cc +++ b/tensorflow/compiler/xla/service/buffer_value.cc @@ -29,8 +29,8 @@ BufferValue::BufferValue(HloInstruction* instruction, const ShapeIndex& index, Id id) : id_(id) { const Shape& shape = ShapeUtil::GetSubshape(instruction->shape(), index); - is_array_ = ShapeUtil::IsArray(shape); - is_tuple_ = ShapeUtil::IsTuple(shape); + is_array_ = shape.IsArray(); + is_tuple_ = shape.IsTuple(); } BufferValue::~BufferValue() {} diff --git a/tensorflow/compiler/xla/service/call_graph.cc b/tensorflow/compiler/xla/service/call_graph.cc index bdd5069632e84fe6c67ca129f726432479ac1b35..94af788c54f6c722997311bec50da3ed93aa3cee 100644 --- a/tensorflow/compiler/xla/service/call_graph.cc +++ b/tensorflow/compiler/xla/service/call_graph.cc @@ -58,7 +58,7 @@ CallContext GetInstructionCallContext(HloOpcode opcode) { case HloOpcode::kConditional: case HloOpcode::kWhile: return CallContext::kSequential; - case HloOpcode::kCrossReplicaSum: + case HloOpcode::kAllReduce: case HloOpcode::kMap: case HloOpcode::kReduce: case HloOpcode::kReduceWindow: @@ -236,6 +236,41 @@ void CallGraph::SetCallContexts() { } } +void CallGraph::SetNodeDepths() { + std::queue worklist; + + // Initialize node depths to -1. + for (CallGraphNode& node : nodes_) { + node.set_depth(-1); + } + + // Initialize worklist with all roots of the call graph (computations without + // callers). + for (const HloComputation* computation : module_->computations()) { + CallGraphNode& node = GetNode(computation); + if (node.callers().empty()) { + node.set_depth(0); + worklist.push(&node); + } + } + + while (!worklist.empty()) { + CallGraphNode* node = worklist.front(); + worklist.pop(); + for (const HloComputation* callee : node->callees()) { + CallGraphNode& callee_node = GetNode(callee); + if (callee_node.depth() < node->depth() + 1) { + callee_node.set_depth(node->depth() + 1); + worklist.push(&callee_node); + } + } + } + + for (CallGraphNode& node : nodes_) { + CHECK_NE(node.depth(), -1); + } +} + /* static */ std::unique_ptr CallGraph::Build(const HloModule* module) { // Constructor for CallGraph is private so absl::make_unique can't be used. @@ -271,6 +306,8 @@ std::unique_ptr CallGraph::Build(const HloModule* module) { } call_graph->SetCallContexts(); + call_graph->SetNodeDepths(); + XLA_VLOG_LINES(1, call_graph->ToString()); return call_graph; @@ -325,6 +362,15 @@ bool CallGraph::IsFlattened() const { return true; } +std::vector CallGraph::GetComputationCallers( + HloComputation* c) { + std::vector callers; + for (auto callsite : GetNode(c).caller_callsites()) { + callers.push_back(callsite.instruction()); + } + return callers; +} + std::pair CallGraph::NearestAncestorsInSameComputation(HloInstruction* a, HloInstruction* b) const { @@ -343,15 +389,38 @@ CallGraph::NearestAncestorsInSameComputation(HloInstruction* a, // Iterate through the callee->caller chains and find the earliest common // element. - for (HloInstruction* a_ancestor = a; a_ancestor != nullptr; - a_ancestor = next_caller(a_ancestor)) { - for (HloInstruction* b_ancestor = b; b_ancestor != nullptr; - b_ancestor = next_caller(b_ancestor)) { - if (a_ancestor->parent() == b_ancestor->parent()) { - return {a_ancestor, b_ancestor}; + HloInstruction* a_ancestor = a; + HloInstruction* b_ancestor = b; + int a_depth = GetNode(a->parent()).depth(); + int b_depth = GetNode(b->parent()).depth(); + + // Advance a_ancestor (b_ancestor) up the call chain until the call depth of + // a_ancestor or b_ancestor are the same. Necessarily each call to next_caller + // reduces the depth by exactly one. + if (a_depth > b_depth) { + for (int i = 0; i < a_depth - b_depth; ++i) { + a_ancestor = next_caller(a_ancestor); + if (a_ancestor == nullptr) { + return {nullptr, nullptr}; + } + } + } else if (b_depth > a_depth) { + for (int i = 0; i < b_depth - a_depth; ++i) { + b_ancestor = next_caller(b_ancestor); + if (b_ancestor == nullptr) { + return {nullptr, nullptr}; } } } + + while ((a_ancestor != nullptr) && (b_ancestor != nullptr)) { + if (a_ancestor->parent() == b_ancestor->parent()) { + return {a_ancestor, b_ancestor}; + } + + a_ancestor = next_caller(a_ancestor); + b_ancestor = next_caller(b_ancestor); + } return {nullptr, nullptr}; } diff --git a/tensorflow/compiler/xla/service/call_graph.h b/tensorflow/compiler/xla/service/call_graph.h index cb56f4789d06ac33acdaadc8b619b9e37f683d58..c02ffda575278905f6549b362e5e7d94f5713b36 100644 --- a/tensorflow/compiler/xla/service/call_graph.h +++ b/tensorflow/compiler/xla/service/call_graph.h @@ -121,6 +121,11 @@ class CallGraphNode { // Returns the context in which this computation is called. CallContext context() const { return context_; } + // Returns the depth of this node in the call graph. The depth is defined as + // the length of the longest call chain from a computation with no callers + // (usually the entry computation node) to this node. + int depth() const { return depth_; } + string ToString() const; private: @@ -130,6 +135,9 @@ class CallGraphNode { // Sets the context in which this computation is called. void set_context(CallContext value) { context_ = value; } + // Sets the depth of this node in the graph. + void set_depth(int value) { depth_ = value; } + // Adds a callsite which calls this computation. Updates callers to include // the calling computation. void AddCallerCallSite(const CallSite& caller_callsite); @@ -164,6 +172,9 @@ class CallGraphNode { // The context in which this computation is called. CallContext context_ = CallContext::kNone; + + // The depth of this node in the call graph. + int depth_ = 0; }; // The call graph for an HLO module. The graph includes a node for each @@ -236,14 +247,25 @@ class CallGraph { // FlattenCallGraph. bool IsFlattened() const; + // Returns a vector of instructions calling the passed computation. + // (Often a vector of size 1.) + std::vector GetComputationCallers(HloComputation* c); + string ToString() const; 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(); + // Sets the call node depths for every node in the graph. + void SetNodeDepths(); + // Helper method for VisitNodes(). Traverses the call graph from 'node' in DFS // post order (callee before caller) calling visitor_func on each node. Adds // nodes to 'visited' as each node is visited. Skips nodes already in diff --git a/tensorflow/compiler/xla/service/call_graph_test.cc b/tensorflow/compiler/xla/service/call_graph_test.cc index 34f3f914d593bc603c4964663f9cafb70a136fd3..5de724f8924b78008ba4c56603b61bf93fbc5e7c 100644 --- a/tensorflow/compiler/xla/service/call_graph_test.cc +++ b/tensorflow/compiler/xla/service/call_graph_test.cc @@ -21,7 +21,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/tests/hlo_verified_test_base.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" #include "tensorflow/compiler/xla/util.h" #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/core/status_test_util.h" @@ -31,7 +31,7 @@ namespace { using ::testing::UnorderedElementsAre; -class CallGraphTest : public HloVerifiedTestBase { +class CallGraphTest : public HloTestBase { protected: // Build and return a trivial computation taking and returning a scalar. std::unique_ptr MakeScalarComputation( @@ -93,15 +93,16 @@ class CallGraphTest : public HloVerifiedTestBase { TEST_F(CallGraphTest, SingletonComputation) { // Test the call graph of a module with a single computation. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* computation = module->AddEntryComputation(MakeScalarComputation()); - std::unique_ptr call_graph = CallGraph::Build(module); + std::unique_ptr call_graph = CallGraph::Build(module.get()); EXPECT_EQ(1, call_graph->nodes().size()); EXPECT_TRUE(call_graph->IsFlattened()); const CallGraphNode& node = call_graph->GetNode(computation); EXPECT_EQ(computation, node.computation()); + EXPECT_EQ(node.depth(), 0); EXPECT_TRUE(node.callsites().empty()); EXPECT_TRUE(node.callees().empty()); EXPECT_TRUE(node.caller_callsites().empty()); @@ -112,21 +113,23 @@ TEST_F(CallGraphTest, SingletonComputation) { TEST_F(CallGraphTest, UnreachableComputation) { // Test the call graph of a module with an entry computation and an // unreachable computation. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* entry_computation = module->AddEntryComputation(MakeScalarComputation()); HloComputation* unreachable_computation = module->AddEmbeddedComputation(MakeScalarComputation()); - std::unique_ptr call_graph = CallGraph::Build(module); + std::unique_ptr call_graph = CallGraph::Build(module.get()); EXPECT_EQ(2, call_graph->nodes().size()); const CallGraphNode& entry_node = call_graph->GetNode(entry_computation); + EXPECT_EQ(entry_node.depth(), 0); EXPECT_EQ(entry_computation, entry_node.computation()); EXPECT_EQ(CallContext::kSequential, entry_node.context()); const CallGraphNode& unreachable_node = call_graph->GetNode(unreachable_computation); + EXPECT_EQ(unreachable_node.depth(), 0); EXPECT_EQ(unreachable_computation, unreachable_node.computation()); EXPECT_EQ(CallContext::kSequential, unreachable_node.context()); } @@ -134,17 +137,18 @@ TEST_F(CallGraphTest, UnreachableComputation) { TEST_F(CallGraphTest, ParallelComputation) { // Test a call graph of a module with an entry computation which calls another // computation in a parallel context via kMap. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* map_computation = module->AddEmbeddedComputation(MakeScalarComputation()); HloComputation* entry_computation = module->AddEntryComputation( MakeMappingComputation(map_computation, /*callsites=*/5)); - std::unique_ptr call_graph = CallGraph::Build(module); + std::unique_ptr call_graph = CallGraph::Build(module.get()); EXPECT_EQ(2, call_graph->nodes().size()); const CallGraphNode& entry_node = call_graph->GetNode(entry_computation); EXPECT_EQ(entry_computation, entry_node.computation()); + EXPECT_EQ(entry_node.depth(), 0); EXPECT_EQ(CallContext::kSequential, entry_node.context()); EXPECT_EQ(5, entry_node.callsites().size()); EXPECT_EQ(1, entry_node.callees().size()); @@ -153,6 +157,7 @@ TEST_F(CallGraphTest, ParallelComputation) { const CallGraphNode& map_node = call_graph->GetNode(map_computation); EXPECT_EQ(map_computation, map_node.computation()); + EXPECT_EQ(map_node.depth(), 1); EXPECT_EQ(CallContext::kParallel, map_node.context()); EXPECT_TRUE(map_node.callsites().empty()); EXPECT_TRUE(map_node.callees().empty()); @@ -163,13 +168,13 @@ TEST_F(CallGraphTest, ParallelComputation) { TEST_F(CallGraphTest, SequentialComputations) { // Test a call graph of a module with an entry computation which calls another // computation in a sequential context via kCall. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* called_computation = module->AddEmbeddedComputation(MakeScalarComputation()); HloComputation* entry_computation = module->AddEntryComputation( MakeCallingComputation(called_computation, /*callsites=*/3)); - std::unique_ptr call_graph = CallGraph::Build(module); + std::unique_ptr call_graph = CallGraph::Build(module.get()); EXPECT_EQ(2, call_graph->nodes().size()); // The called computation is only called from one other computation, but there @@ -196,7 +201,7 @@ TEST_F(CallGraphTest, SequentialComputations) { TEST_F(CallGraphTest, ContextBothComputations) { // Test a call graph of a module with an entry computation which calls another // computation in both a parallel and sequential context. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* subcomputation = module->AddEmbeddedComputation(MakeScalarComputation()); @@ -210,7 +215,7 @@ TEST_F(CallGraphTest, ContextBothComputations) { HloComputation* entry_computation = module->AddEntryComputation(builder.Build()); - std::unique_ptr call_graph = CallGraph::Build(module); + std::unique_ptr call_graph = CallGraph::Build(module.get()); EXPECT_EQ(2, call_graph->nodes().size()); EXPECT_FALSE(call_graph->IsFlattened()); @@ -234,12 +239,13 @@ TEST_F(CallGraphTest, ContextBothComputations) { EXPECT_EQ(entry_node.GetCallSite(map), &map_callsite); const CallGraphNode& sub_node = call_graph->GetNode(subcomputation); + EXPECT_EQ(sub_node.depth(), 1); EXPECT_EQ(CallContext::kBoth, sub_node.context()); } TEST_F(CallGraphTest, ComputationWithConditional) { // Test a call graph of a module with a conditional. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* true_computation = module->AddEmbeddedComputation(MakeScalarComputation(HloOpcode::kCeil)); HloComputation* false_computation = @@ -259,11 +265,12 @@ TEST_F(CallGraphTest, ComputationWithConditional) { HloComputation* entry_computation = module->AddEntryComputation(builder.Build()); - std::unique_ptr call_graph = CallGraph::Build(module); + std::unique_ptr call_graph = CallGraph::Build(module.get()); EXPECT_EQ(3, call_graph->nodes().size()); const CallGraphNode& entry_node = call_graph->GetNode(entry_computation); + EXPECT_EQ(entry_node.depth(), 0); EXPECT_EQ(entry_computation, entry_node.computation()); EXPECT_EQ(1, entry_node.callsites().size()); @@ -275,11 +282,13 @@ TEST_F(CallGraphTest, ComputationWithConditional) { EXPECT_EQ(entry_node.GetCallSite(conditional), &conditional_callsite); const CallGraphNode& true_node = call_graph->GetNode(true_computation); + EXPECT_EQ(true_node.depth(), 1); EXPECT_TRUE(true_node.callees().empty()); EXPECT_EQ(1, true_node.callers().size()); EXPECT_EQ(entry_computation, true_node.callers()[0]); const CallGraphNode& false_node = call_graph->GetNode(false_computation); + EXPECT_EQ(false_node.depth(), 1); EXPECT_TRUE(false_node.callees().empty()); EXPECT_EQ(1, false_node.callers().size()); EXPECT_EQ(entry_computation, false_node.callers()[0]); @@ -298,7 +307,7 @@ TEST_F(CallGraphTest, ComplexGraph) { // c // // Calls are made via kCall, kWhile, and kMap instructions. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* cond_computation = module->AddEmbeddedComputation(MakeConditionComputation()); HloComputation* c_computation = @@ -328,13 +337,25 @@ TEST_F(CallGraphTest, ComplexGraph) { entry_computation = module->AddEntryComputation(builder.Build()); } - std::unique_ptr call_graph = CallGraph::Build(module); + std::unique_ptr call_graph = CallGraph::Build(module.get()); EXPECT_EQ(5, call_graph->nodes().size()); EXPECT_FALSE(call_graph->IsFlattened()); + const CallGraphNode& entry_node = call_graph->GetNode(entry_computation); + const CallGraphNode& a_node = call_graph->GetNode(a_computation); + const CallGraphNode& b_node = call_graph->GetNode(b_computation); + const CallGraphNode& c_node = call_graph->GetNode(c_computation); + const CallGraphNode& cond_node = call_graph->GetNode(cond_computation); + + // Verify depths. + EXPECT_EQ(entry_node.depth(), 0); + EXPECT_EQ(a_node.depth(), 1); + EXPECT_EQ(b_node.depth(), 2); + EXPECT_EQ(c_node.depth(), 3); + EXPECT_EQ(cond_node.depth(), 2); + // Entry computation has one while instruction calling two computations // (cond_computation and a_computation). - const CallGraphNode& entry_node = call_graph->GetNode(entry_computation); ASSERT_EQ(1, entry_node.callsites().size()); const std::vector& called_computations = entry_node.callsites()[0].called_computations(); @@ -342,7 +363,6 @@ TEST_F(CallGraphTest, ComplexGraph) { UnorderedElementsAre(cond_computation, a_computation)); EXPECT_EQ(CallContext::kSequential, entry_node.context()); - const CallGraphNode& c_node = call_graph->GetNode(c_computation); EXPECT_TRUE(c_node.callsites().empty()); EXPECT_THAT(c_node.callers(), UnorderedElementsAre(a_computation, b_computation)); @@ -364,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); }; @@ -418,7 +438,7 @@ TEST_F(CallGraphTest, ComplexGraphNearestAncestors) { // c // // Calls are made via kCall, kWhile, and kMap instructions. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* cond_computation = module->AddEmbeddedComputation(MakeConditionComputation()); HloComputation* c_computation = @@ -452,7 +472,7 @@ TEST_F(CallGraphTest, ComplexGraphNearestAncestors) { entry_computation = module->AddEntryComputation(builder.Build()); } - std::unique_ptr call_graph = CallGraph::Build(module); + std::unique_ptr call_graph = CallGraph::Build(module.get()); EXPECT_EQ(5, call_graph->nodes().size()); // Verify NearestAncestorsInSameComputation for various instructions in the @@ -479,10 +499,10 @@ TEST_F(CallGraphTest, ComplexGraphNearestAncestors) { TEST_F(CallGraphTest, VisitSingletonComputation) { // Test the call graph visitor with a call graph with a single node. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* computation = module->AddEntryComputation(MakeScalarComputation()); - std::unique_ptr call_graph = CallGraph::Build(module); + std::unique_ptr call_graph = CallGraph::Build(module.get()); std::vector visited; TF_ASSERT_OK(call_graph->VisitNodes([&visited](const CallGraphNode& node) { @@ -494,12 +514,12 @@ TEST_F(CallGraphTest, VisitSingletonComputation) { TEST_F(CallGraphTest, VisitUnreachableComputation) { // Test the call graph visitor with a call graph with an unreachable node. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* entry_computation = module->AddEntryComputation(MakeScalarComputation()); HloComputation* unreachable_computation = module->AddEmbeddedComputation(MakeScalarComputation()); - std::unique_ptr call_graph = CallGraph::Build(module); + std::unique_ptr call_graph = CallGraph::Build(module.get()); // Test visitation of only reachable nodes. { @@ -531,9 +551,9 @@ TEST_F(CallGraphTest, VisitUnreachableComputation) { TEST_F(CallGraphTest, VisitWithError) { // Test that the call graph visitor properly propagates errors. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(MakeScalarComputation()); - std::unique_ptr call_graph = CallGraph::Build(module); + std::unique_ptr call_graph = CallGraph::Build(module.get()); Status status = call_graph->VisitNodes( [](const CallGraphNode&) { return InternalError("Visitation failed"); }); diff --git a/tensorflow/compiler/xla/service/call_inliner_test.cc b/tensorflow/compiler/xla/service/call_inliner_test.cc index e6b566543594a86eb5369ee9b7440f62618f6c5a..0b6e323f75c7a5dae127e20d2a4b92a83a72df3b 100644 --- a/tensorflow/compiler/xla/service/call_inliner_test.cc +++ b/tensorflow/compiler/xla/service/call_inliner_test.cc @@ -28,7 +28,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_pass_fix.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/test.h" -#include "tensorflow/compiler/xla/tests/hlo_verified_test_base.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" @@ -40,7 +40,7 @@ namespace { // Tests for call inlining that are most tractable at the HLO level (vs // ComputationBuilder API in call_test.cc). -using CallInlinerTest = HloVerifiedTestBase; +using CallInlinerTest = HloTestBase; TEST_F(CallInlinerTest, ControlDependenciesAreCarriedToCaller) { // "inner" computation just has a control dependency from the "zero" value to @@ -51,7 +51,7 @@ TEST_F(CallInlinerTest, ControlDependenciesAreCarriedToCaller) { HloInstruction* one = inner.AddInstruction( HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0f))); TF_ASSERT_OK(zero->AddControlDependencyTo(one)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* inner_computation = module->AddEmbeddedComputation(inner.Build()); @@ -64,7 +64,7 @@ TEST_F(CallInlinerTest, ControlDependenciesAreCarriedToCaller) { auto computation = module->AddEntryComputation(outer.Build()); CallInliner call_inliner; - TF_ASSERT_OK_AND_ASSIGN(bool mutated, call_inliner.Run(module)); + TF_ASSERT_OK_AND_ASSIGN(bool mutated, call_inliner.Run(module.get())); ASSERT_TRUE(mutated); EXPECT_THAT(computation->root_instruction(), op::Constant()); EXPECT_EQ(computation->root_instruction()->literal().GetFirstElement(), @@ -79,7 +79,7 @@ TEST_F(CallInlinerTest, ControlDependenciesAreCarriedToCaller) { // returns false should be identical to just returning false). TEST_F(CallInlinerTest, CallsWithinWhileBodiesAreInlined) { const Shape pred = ShapeUtil::MakeShape(PRED, {}); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); // Create a lambda that calls a function that returns the false predicate. // Note we also use this lambda twice by reference, just to make the test a @@ -107,7 +107,7 @@ TEST_F(CallInlinerTest, CallsWithinWhileBodiesAreInlined) { auto computation = module->AddEntryComputation(outer.Build()); CallInliner call_inliner; - TF_ASSERT_OK_AND_ASSIGN(bool mutated, call_inliner.Run(module)); + TF_ASSERT_OK_AND_ASSIGN(bool mutated, call_inliner.Run(module.get())); ASSERT_TRUE(mutated); EXPECT_THAT( computation->root_instruction()->while_condition()->root_instruction(), @@ -120,7 +120,7 @@ TEST_F(CallInlinerTest, CallsWithinWhileBodiesAreInlined) { // whole pass. TEST_F(CallInlinerTest, InlineWithoutRunningPass) { const Shape pred = ShapeUtil::MakeShape(PRED, {}); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation::Builder just_false(TestName() + ".false"); auto* true_constant = just_false.AddInstruction( @@ -144,7 +144,7 @@ TEST_F(CallInlinerTest, InlineWithoutRunningPass) { TEST_F(CallInlinerTest, CallToOutfeedComputationIsInlined) { const Shape f32 = ShapeUtil::MakeShape(F32, {}); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation::Builder outfeeder(TestName() + ".outfeeder"); auto value = outfeeder.AddInstruction( @@ -163,7 +163,7 @@ TEST_F(CallInlinerTest, CallToOutfeedComputationIsInlined) { module->AddEntryComputation(outer.Build()); CallInliner call_inliner; - TF_ASSERT_OK_AND_ASSIGN(bool mutated, call_inliner.Run(module)); + TF_ASSERT_OK_AND_ASSIGN(bool mutated, call_inliner.Run(module.get())); ASSERT_TRUE(mutated); } 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/compilation_cache.cc b/tensorflow/compiler/xla/service/compilation_cache.cc new file mode 100644 index 0000000000000000000000000000000000000000..2662fe46705f4936ce0d654df0943e7d30890ebe --- /dev/null +++ b/tensorflow/compiler/xla/service/compilation_cache.cc @@ -0,0 +1,70 @@ +/* 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/compiler/xla/service/compilation_cache.h" + +#include + +#include "tensorflow/compiler/xla/types.h" +#include "tensorflow/compiler/xla/util.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/lib/strings/strcat.h" +#include "tensorflow/core/platform/logging.h" + +namespace xla { + +namespace { + +int64 GetUniqueId() { + static tensorflow::mutex mu(tensorflow::LINKER_INITIALIZED); + static int64 counter = 0; + tensorflow::mutex_lock loc(mu); + const int64 id = counter++; + return id; +} + +} // namespace + +ExecutionHandle CompilationCache::Insert( + std::unique_ptr executable) { + tensorflow::mutex_lock lock(mutex_); + + CacheKey key = GetUniqueId(); + VLOG(2) << "inserting cache key: " << key; + CHECK_EQ(cache_.count(key), 0); + cache_.emplace(key, std::move(executable)); + + ExecutionHandle handle; + handle.set_handle(key); + return handle; +} + +StatusOr> CompilationCache::LookUp( + const ExecutionHandle& handle) const { + tensorflow::mutex_lock lock(mutex_); + + CacheKey key = handle.handle(); + VLOG(2) << "looking up cache key: " << key; + if (cache_.count(key) == 0) { + VLOG(2) << "cache key not found: " << key; + return InvalidArgumentStrCat("can not find executable with handle ", key); + } else { + auto& result = cache_.at(key); + VLOG(2) << "hit executable: " << result->module().name(); + return result; + } +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/compilation_cache.h b/tensorflow/compiler/xla/service/compilation_cache.h new file mode 100644 index 0000000000000000000000000000000000000000..5f94def509d4d4a8950272cb498af5056a698ce0 --- /dev/null +++ b/tensorflow/compiler/xla/service/compilation_cache.h @@ -0,0 +1,62 @@ +/* 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_SERVICE_COMPILATION_CACHE_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_COMPILATION_CACHE_H_ + +#include +#include +#include + +#include "absl/container/flat_hash_map.h" +#include "tensorflow/compiler/xla/service/executable.h" +#include "tensorflow/compiler/xla/service/hlo_module_config.h" +#include "tensorflow/compiler/xla/types.h" +#include "tensorflow/core/platform/macros.h" +#include "tensorflow/core/platform/mutex.h" +#include "tensorflow/core/platform/thread_annotations.h" + +namespace xla { + +// A cache which stores Executables indexed by computation handle and version. +// +// TODO(b/119042872): Provide mechanism for removing computations from the +// compilation cache. +class CompilationCache { + public: + CompilationCache() {} + + ExecutionHandle Insert(std::unique_ptr executable); + + // Lookup the Executable for the specified handle in the cache. Return a + // shared_ptr to the Executable if it exists in the cache. + StatusOr> LookUp( + const ExecutionHandle& handle) const; + + protected: + mutable tensorflow::mutex mutex_; + + using CacheKey = int64; + + absl::flat_hash_map> cache_ + GUARDED_BY(mutex_); + + private: + TF_DISALLOW_COPY_AND_ASSIGN(CompilationCache); +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_COMPILATION_CACHE_H_ diff --git a/tensorflow/compiler/xla/service/compile_only_service.cc b/tensorflow/compiler/xla/service/compile_only_service.cc index 6d67f970020d278cc7bf61b56350200d3e5cb926..1965925fa7f6d50b1d7af918bc3468d4b4d5d0a2 100644 --- a/tensorflow/compiler/xla/service/compile_only_service.cc +++ b/tensorflow/compiler/xla/service/compile_only_service.cc @@ -20,7 +20,7 @@ limitations under the License. #include #include "absl/strings/str_cat.h" -#include "tensorflow/compiler/xla/legacy_flags/debug_options_flags.h" +#include "tensorflow/compiler/xla/debug_options_flags.h" #include "tensorflow/compiler/xla/service/backend.h" #include "tensorflow/compiler/xla/service/computation_layout.h" #include "tensorflow/compiler/xla/service/platform_util.h" @@ -86,15 +86,15 @@ CompileOnlyService::CompileAheadOfTime( Executable::DumpToDirectory(per_host_path, filename, hlo_snapshot)); } - const auto& program_shape = instance.computation.host_program_shape(); ExecutionOptions execution_options; *execution_options.mutable_debug_options() = debug_options; *execution_options.mutable_shape_with_output_layout() = - *instance.result_layout; + instance.result_layout->ToProto(); TF_ASSIGN_OR_RETURN( std::unique_ptr module_config, - CreateModuleConfig(program_shape, instance.argument_layouts, - &execution_options)); + CreateModuleConfig( + ProgramShape(instance.computation.host_program_shape()), + instance.argument_layouts, &execution_options)); TF_ASSIGN_OR_RETURN( std::unique_ptr hlo_module, diff --git a/tensorflow/compiler/xla/service/compiler.cc b/tensorflow/compiler/xla/service/compiler.cc index 80c630c6201503d88a690f04a88f6fca6f3a438a..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. @@ -110,6 +117,6 @@ Compiler::GetPlatformCompilers() { } AotCompilationOptions::AotCompilationOptions() - : debug_options_(legacy_flags::GetDebugOptionsFromFlags()) {} + : debug_options_(GetDebugOptionsFromFlags()) {} } // namespace xla diff --git a/tensorflow/compiler/xla/service/compiler.h b/tensorflow/compiler/xla/service/compiler.h index 9ab179303b3e792c1f94c08626d7bc1afd2099f8..d4db95da8eb901af8a6675f2991def73ccfe8ee6 100644 --- a/tensorflow/compiler/xla/service/compiler.h +++ b/tensorflow/compiler/xla/service/compiler.h @@ -139,14 +139,15 @@ class Compiler { // Optimizes a HLO module group, a set of module which runs concurrently on // multiple devices potentially communicating data between the modules. virtual Status RunHloPassesOnModuleGroup( - HloModuleGroup* module_group, se::StreamExecutor* executor, + HloModuleGroup* module_group, + absl::Span executors, DeviceMemoryAllocator* device_allocator) = 0; // Compiles the HLO module for execution on a device given by the executor, // and returns an executable object or an error status. No HLO passes are // applied to module. Generally a module should be passed through RunHloPasses // prior to calling this method because some HLO passes are required for - // correctness. Takes ownership of the HLO module and is free to transform it. + // correctness. Takes ownership of the HLO module. // // The compiler may optionally specialize to the individual device // (not just type of device) indicated by the executor. 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/computation_placer.h b/tensorflow/compiler/xla/service/computation_placer.h index c899ffb9dc562426ef14c0d414469c04debeec70..844b42a38d7539cccd5c4e30071c0ea6693e3bba 100644 --- a/tensorflow/compiler/xla/service/computation_placer.h +++ b/tensorflow/compiler/xla/service/computation_placer.h @@ -105,8 +105,6 @@ class ComputationPlacer { // Map from platform kind to computation placer singleton. static std::map* GetPlatformComputationPlacers(); - se::Platform::Id platform_id_; - TF_DISALLOW_COPY_AND_ASSIGN(ComputationPlacer); }; diff --git a/tensorflow/compiler/xla/service/conditional_simplifier_test.cc b/tensorflow/compiler/xla/service/conditional_simplifier_test.cc index c43a31b167d47af3c92ed35fa52594fa5da1e4af..289eb6d90239a72ecc0f3312a7e0e8453f946858 100644 --- a/tensorflow/compiler/xla/service/conditional_simplifier_test.cc +++ b/tensorflow/compiler/xla/service/conditional_simplifier_test.cc @@ -25,7 +25,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_opcode.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/test.h" -#include "tensorflow/compiler/xla/tests/hlo_verified_test_base.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.h" @@ -37,7 +37,7 @@ namespace { namespace op = xla::testing::opcode_matchers; -class ConditionalSimplifierTest : public HloVerifiedTestBase { +class ConditionalSimplifierTest : public HloTestBase { public: // Makes a computation that contains a conditional with constant predicate. HloComputation* MakeConditional(HloModule* module); @@ -96,25 +96,28 @@ HloComputation* ConditionalSimplifierTest::MakeConditional(HloModule* module) { } TEST_F(ConditionalSimplifierTest, ConditionalGetsInlined) { - HloComputation* computation = MakeConditional(&module()); - ASSERT_TRUE(ConditionalSimplifier().Run(&module()).ValueOrDie()); + auto m = CreateNewVerifiedModule(); + HloComputation* computation = MakeConditional(m.get()); + ASSERT_TRUE(ConditionalSimplifier().Run(m.get()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::Add(op::Parameter(), op::Constant())); } TEST_F(ConditionalSimplifierTest, ConditionalWithControlDependency) { - HloComputation* computation = MakeConditional(&module()); + auto m = CreateNewVerifiedModule(); + HloComputation* computation = MakeConditional(m.get()); auto* true_op = computation->AddInstruction( HloInstruction::CreateConstant(LiteralUtil::CreateR0(true))); TF_ASSERT_OK( true_op->AddControlDependencyTo(computation->root_instruction())); - EXPECT_FALSE(ConditionalSimplifier().Run(&module()).ValueOrDie()); + EXPECT_FALSE(ConditionalSimplifier().Run(m.get()).ValueOrDie()); } TEST_F(ConditionalSimplifierTest, NotRemovedIfContainsSend) { - HloComputation* computation = MakeConditional(&module()); + auto m = CreateNewVerifiedModule(); + HloComputation* computation = MakeConditional(m.get()); auto* conditional = computation->root_instruction(); ASSERT_EQ(conditional->opcode(), HloOpcode::kConditional); @@ -125,11 +128,12 @@ TEST_F(ConditionalSimplifierTest, NotRemovedIfContainsSend) { HloInstruction::CreateConstant(LiteralUtil::CreateR0(true))), token, /*channel_id=*/0)); true_computation->AddInstruction(HloInstruction::CreateSendDone(send)); - EXPECT_FALSE(ConditionalSimplifier().Run(&module()).ValueOrDie()); + EXPECT_FALSE(ConditionalSimplifier().Run(m.get()).ValueOrDie()); } TEST_F(ConditionalSimplifierTest, NotRemovedIfContainsRecv) { - HloComputation* computation = MakeConditional(&module()); + auto m = CreateNewVerifiedModule(); + HloComputation* computation = MakeConditional(m.get()); auto* conditional = computation->root_instruction(); ASSERT_EQ(conditional->opcode(), HloOpcode::kConditional); @@ -138,18 +142,19 @@ TEST_F(ConditionalSimplifierTest, NotRemovedIfContainsRecv) { auto* recv = true_computation->AddInstruction(HloInstruction::CreateRecv( ShapeUtil::MakeShape(F32, {1}), token, /*channel_id=*/0)); true_computation->AddInstruction(HloInstruction::CreateRecvDone(recv)); - EXPECT_FALSE(ConditionalSimplifier().Run(&module()).ValueOrDie()); + EXPECT_FALSE(ConditionalSimplifier().Run(m.get()).ValueOrDie()); } TEST_F(ConditionalSimplifierTest, NotRemovedIfContainsNonRemovableInstruction) { - HloComputation* computation = MakeConditional(&module()); + auto m = CreateNewVerifiedModule(); + HloComputation* computation = MakeConditional(m.get()); auto* conditional = computation->root_instruction(); ASSERT_EQ(conditional->opcode(), HloOpcode::kConditional); auto* false_computation = conditional->false_computation(); auto token = false_computation->AddInstruction(HloInstruction::CreateToken()); false_computation->AddInstruction(HloInstruction::CreateInfeed( ShapeUtil::MakeShape(F32, {1}), token, "config")); - EXPECT_FALSE(ConditionalSimplifier().Run(&module()).ValueOrDie()); + EXPECT_FALSE(ConditionalSimplifier().Run(m.get()).ValueOrDie()); } } // namespace diff --git a/tensorflow/compiler/xla/service/convolution_feature_group_converter.cc b/tensorflow/compiler/xla/service/convolution_feature_group_converter.cc deleted file mode 100644 index 0ac4a65ec6ae55fabd2b48ea2982b94f9551c8d2..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/service/convolution_feature_group_converter.cc +++ /dev/null @@ -1,249 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/compiler/xla/service/convolution_feature_group_converter.h" - -#include -#include - -#include "absl/memory/memory.h" -#include "tensorflow/compiler/xla/literal.h" -#include "tensorflow/compiler/xla/literal_util.h" -#include "tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h" -#include "tensorflow/compiler/xla/service/hlo_computation.h" -#include "tensorflow/compiler/xla/service/hlo_instruction.h" -#include "tensorflow/compiler/xla/service/hlo_opcode.h" -#include "tensorflow/compiler/xla/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/core/status.h" -#include "tensorflow/core/platform/logging.h" - -namespace xla { - -namespace { - -// ConvolutionVisitor traverses the HLO computation and rewrites Convolution -// operations with feature_group_count > 1 into convolutions with -// feature_group_count = 1. -class ConvolutionVisitor : public DfsHloVisitorWithDefault { - public: - // Default visitor action is to do nothing and return OK. - Status DefaultAction(HloInstruction* /*hlo_instruction*/) override { - return Status::OK(); - } - - Status HandleConvolution(HloInstruction* convolution) override; - - // Runs the visitor on a computation. - static bool Run(HloComputation* computation); - - // Returns whether any convolution ops were rewritten. - const bool changed() const { return changed_; } - - ~ConvolutionVisitor() override = default; - - private: - explicit ConvolutionVisitor(HloComputation* computation) - : computation_(computation) {} - - // Current HloComputation instance the ConvolutionVisitor is traversing. - HloComputation* computation_; - - // Whether rewrite has occurred. - bool changed_ = false; -}; - -bool ConvolutionVisitor::Run(HloComputation* computation) { - ConvolutionVisitor visitor(computation); - TF_CHECK_OK(computation->Accept(&visitor)); - return visitor.changed_; -} - -Shape ExpandedFilterShape(const Shape& shape, int64 group_count, - int64 input_feature_dim) { - int64 num_dims = shape.dimensions_size(); - CHECK_GE(num_dims, 2); - Shape expanded_shape = shape; - expanded_shape.set_dimensions( - input_feature_dim, shape.dimensions(input_feature_dim) * group_count); - return expanded_shape; -} - -// Returns a vector with 'group_count' many groups, where the i-th group -// consists of 'group_size' times the value i. -std::vector GetMaskIds(int64 group_size, int64 group_count) { - std::vector values; - for (int i = 0; i < group_count; ++i) { - for (int j = 0; j < group_size; ++j) { - values.push_back(i); - } - } - return values; -} - -// Create a mask for grouped convolution that will make a normal convolution -// produce the same results as a grouped convolution. For a [2, 1, 6] -// filter this returns a [2, 3, 6] mask -// 1 1 0 0 0 0 -// 0 0 1 1 0 0 -// 0 0 0 0 1 1 -// -// 1 1 0 0 0 0 -// 0 0 1 1 0 0 -// 0 0 0 0 1 1 -// -// The first step is to create a rank 1 constant: -// 0 1 2 -// -// This is broadcasted to -// 0 0 0 0 0 0 -// 1 1 1 1 1 1 -// 2 2 2 2 2 2 -// -// 0 0 0 0 0 0 -// 1 1 1 1 1 1 -// 2 2 2 2 2 2 -// -// Then we create another rank 1 constant -// 0 0 1 1 2 2 -// -// This is broadcasted to -// 0 0 1 1 2 2 -// 0 0 1 1 2 2 -// 0 0 1 1 2 2 -// -// 0 0 1 1 2 2 -// 0 0 1 1 2 2 -// 0 0 1 1 2 2 -// -// Finally we use the Eq op of these two broadcasted constants and get the -// desired mask. -HloInstruction* GetExpandedFilterMask( - const Shape& filter_shape, int64 input_feature_dim, - int64 output_feature_dim, int64 group_count, - const std::function)>& - add_instruction) { - Shape expanded_filter_shape = - ExpandedFilterShape(filter_shape, group_count, input_feature_dim); - Shape mask_shape = ShapeUtil::MakeShape( - S32, AsInt64Slice(expanded_filter_shape.dimensions())); - int64 output_feature = filter_shape.dimensions(output_feature_dim); - int64 group_size = filter_shape.dimensions(input_feature_dim); - - // Create a 'input_feature' sized linspace and 'output_feature' sized linspace - // that will be broadcasted into perpendicular dimensions and compared. - const std::vector input_feature_filter_mask = - GetMaskIds(group_size, group_count); - const std::vector output_feature_filter_mask = - GetMaskIds(output_feature / group_count, group_count); - - auto mask1 = add_instruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR1(input_feature_filter_mask))); - auto broadcasted_mask1 = add_instruction( - HloInstruction::CreateBroadcast(mask_shape, mask1, {input_feature_dim})); - auto mask2 = add_instruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR1(output_feature_filter_mask))); - auto broadcasted_mask2 = add_instruction( - HloInstruction::CreateBroadcast(mask_shape, mask2, {output_feature_dim})); - - // Compare the broadcasted output feature linspace to the input feature - // linspace to create a diagonal predicate. - Shape predicate_shape = ShapeUtil::MakeShape( - PRED, AsInt64Slice(expanded_filter_shape.dimensions())); - return add_instruction(HloInstruction::CreateBinary( - predicate_shape, HloOpcode::kEq, broadcasted_mask1, broadcasted_mask2)); -} - -Status ConvolutionVisitor::HandleConvolution(HloInstruction* convolution) { - int64 group_count = convolution->feature_group_count(); - if (group_count == 1) { - return Status::OK(); - } - auto filter = convolution->mutable_operand(1); - changed_ = true; - auto add = [&](std::unique_ptr inst) { - return computation_->AddInstruction(std::move(inst)); - }; - - auto dim_numbers = convolution->convolution_dimension_numbers(); - int64 input_feature_dim = dim_numbers.kernel_input_feature_dimension(); - int64 group_size = filter->shape().dimensions(input_feature_dim); - int64 output_feature_dim = dim_numbers.kernel_output_feature_dimension(); - auto expanded_filter_shape = - ExpandedFilterShape(filter->shape(), group_count, input_feature_dim); - HloInstruction* filter_mask = GetExpandedFilterMask( - filter->shape(), input_feature_dim, output_feature_dim, group_count, add); - HloInstruction* expanded_filter; - // We want to repeat 'filter' in the 'input_feature_dim' dimension - // 'group_count' times. - if (group_size == 1) { - Shape reshaped_filter_shape = - ShapeUtil::DeleteDimension(input_feature_dim, filter->shape()); - auto reshaped_filter = - add(HloInstruction::CreateReshape(reshaped_filter_shape, filter)); - std::vector broadcast_dims; - for (int64 i = 0; i < filter->shape().dimensions_size(); ++i) { - if (i == input_feature_dim) { - continue; - } - broadcast_dims.push_back(i); - } - expanded_filter = add(HloInstruction::CreateBroadcast( - expanded_filter_shape, reshaped_filter, broadcast_dims)); - } else { - // We could possibly also use reshape, broadcast, reshape instead of concat - // here, but it would require more complex code, and for depthwise - // convolution we would never end up in this branch. - std::vector concat_operands(group_count, filter); - expanded_filter = add(HloInstruction::CreateConcatenate( - expanded_filter_shape, concat_operands, input_feature_dim)); - } - auto zero = add(HloInstruction::CreateConstant( - LiteralUtil::Zero(expanded_filter_shape.element_type()))); - auto zero_filter = - add(HloInstruction::CreateBroadcast(expanded_filter_shape, zero, {})); - auto new_filter = add( - HloInstruction::CreateTernary(expanded_filter_shape, HloOpcode::kSelect, - filter_mask, expanded_filter, zero_filter)); - auto new_convolution = HloInstruction::CreateConvolve( - convolution->shape(), convolution->mutable_operand(0), new_filter, - /*feature_group_count=*/1, convolution->window(), dim_numbers, - convolution->precision_config()); - TF_RETURN_IF_ERROR(computation_->ReplaceWithNewInstruction( - convolution, std::move(new_convolution))); - return Status::OK(); -} - -} // namespace - -StatusOr ConvolutionFeatureGroupConverter::Run(HloModule* module) { - XLA_VLOG_LINES(2, "ConvolutionFeatureGroupConverter::Run(), before:\n" + - module->ToString()); - bool changed = false; - for (auto* comp : module->MakeNonfusionComputations()) { - if (ConvolutionVisitor::Run(comp)) { - changed = true; - } - } - XLA_VLOG_LINES(2, "ConvolutionFeatureGroupConverter::Run(), after:\n" + - module->ToString()); - return changed; -} - -} // namespace xla diff --git a/tensorflow/compiler/xla/service/convolution_feature_group_converter.h b/tensorflow/compiler/xla/service/convolution_feature_group_converter.h deleted file mode 100644 index ce0138e56fbd51daaf5d3ac329ccbe31a9fdbde7..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/service/convolution_feature_group_converter.h +++ /dev/null @@ -1,43 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_CONVOLUTION_FEATURE_GROUP_CONVERTER_H_ -#define TENSORFLOW_COMPILER_XLA_SERVICE_CONVOLUTION_FEATURE_GROUP_CONVERTER_H_ - -#include "absl/strings/string_view.h" -#include "tensorflow/compiler/xla/service/hlo_module.h" -#include "tensorflow/compiler/xla/service/hlo_pass_interface.h" -#include "tensorflow/compiler/xla/status_macros.h" - -namespace xla { - -// A pass which rewrites convolutions with feature_group_count > 1 into -// convolutions with feature_group_count = 1. -class ConvolutionFeatureGroupConverter : public HloModulePass { - public: - ConvolutionFeatureGroupConverter() {} - - absl::string_view name() const override { - return "convolution-feature-group-converter"; - } - - // Run convolution rewriting on the given computation. Returns whether the - // computation was changed. - StatusOr Run(HloModule* module) override; -}; - -} // namespace xla - -#endif // TENSORFLOW_COMPILER_XLA_SERVICE_CONVOLUTION_FEATURE_GROUP_CONVERTER_H_ diff --git a/tensorflow/compiler/xla/service/convolution_feature_group_converter_test.cc b/tensorflow/compiler/xla/service/convolution_feature_group_converter_test.cc deleted file mode 100644 index 28373ebf636c7b6b3059dcf6cd931901ebc87fc2..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/service/convolution_feature_group_converter_test.cc +++ /dev/null @@ -1,100 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/compiler/xla/service/convolution_feature_group_converter.h" - -#include -#include - -#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_opcode.h" -#include "tensorflow/compiler/xla/service/hlo_parser.h" -#include "tensorflow/compiler/xla/test.h" -#include "tensorflow/compiler/xla/tests/hlo_test_base.h" -#include "tensorflow/compiler/xla/types.h" - -namespace xla { -namespace { - -using ConvolutionFeatureGroupConverterTest = HloTestBase; -namespace op = testing::opcode_matchers; - -TEST_F(ConvolutionFeatureGroupConverterTest, - ConvertFeatureGroupCountEqualToInputFeatureDim) { - string hlo_string = R"(HloModule Convolve1D1Window_0_module - -ENTRY %Convolve1D1Window_0.v3 (input: f32[1,2,2], filter: f32[1,1,2]) -> f32[1,2,2] { - %input = f32[1,2,2]{2,1,0} parameter(0) - %copy = f32[1,2,2]{2,0,1} copy(f32[1,2,2]{2,1,0} %input) - %filter = f32[1,1,2]{2,1,0} parameter(1) - ROOT %convolution = f32[1,2,2]{2,0,1} convolution(f32[1,2,2]{2,0,1} %copy, f32[1,1,2]{2,1,0} %filter), window={size=1}, dim_labels=b0f_0io->b0f, feature_group_count=2 -})"; - 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); - ConvolutionFeatureGroupConverter converter; - ASSERT_TRUE(converter.Run(module.get()).ValueOrDie()); - root = computation->root_instruction(); - // Make sure the convolution is converted to one with feature_group_count = 1. - EXPECT_EQ(root->opcode(), HloOpcode::kConvolution); - EXPECT_EQ(root->feature_group_count(), 1); - // Verify that the filter operand has been replaced. - EXPECT_THAT(root->operand(1), - op::Select(op::Eq(op::Broadcast(op::Constant()), - op::Broadcast(op::Constant())), - op::Broadcast(op::Reshape(op::Parameter())), - op::Broadcast(op::Constant()))); -} - -TEST_F(ConvolutionFeatureGroupConverterTest, - ConvertFeatureGroupCountDivisorOfInputFeatureDim) { - string hlo_string = R"(HloModule Convolve1D1Window_0_module - -ENTRY %Convolve1D1Window_0.v3 (input: f32[1,2,4], filter: f32[1,2,2]) -> f32[1,2,2] { - %input = f32[1,2,4]{2,1,0} parameter(0) - %copy = f32[1,2,4]{2,0,1} copy(f32[1,2,4]{2,1,0} %input) - %filter = f32[1,2,2]{2,1,0} parameter(1) - ROOT %convolution = f32[1,2,2]{2,0,1} convolution(f32[1,2,4]{2,0,1} %copy, f32[1,2,2]{2,1,0} %filter), window={size=1}, dim_labels=b0f_0io->b0f, feature_group_count=2 -})"; - 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); - ConvolutionFeatureGroupConverter converter; - ASSERT_TRUE(converter.Run(module.get()).ValueOrDie()); - root = computation->root_instruction(); - // Make sure the convolution is converted to one with feature_group_count = 1. - EXPECT_EQ(root->opcode(), HloOpcode::kConvolution); - EXPECT_EQ(root->feature_group_count(), 1); - // Verify that the filter operand has been replaced. - EXPECT_THAT(root->operand(1), - op::Select(op::Eq(op::Broadcast(op::Constant()), - op::Broadcast(op::Constant())), - // We expect to see Concatenate here instead of - // Broadcast, because feature_group_count < input - // feature dimension. - op::Concatenate(op::Parameter(), op::Parameter()), - op::Broadcast(op::Constant()))); -} - -} // namespace -} // namespace xla diff --git a/tensorflow/compiler/xla/service/convolution_group_converter.cc b/tensorflow/compiler/xla/service/convolution_group_converter.cc new file mode 100644 index 0000000000000000000000000000000000000000..b3ed27d9a849eced006eb3b01977ad2fe7ed7367 --- /dev/null +++ b/tensorflow/compiler/xla/service/convolution_group_converter.cc @@ -0,0 +1,621 @@ +/* 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/convolution_group_converter.h" + +#include +#include + +#include "absl/memory/memory.h" +#include "tensorflow/compiler/xla/literal.h" +#include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h" +#include "tensorflow/compiler/xla/service/hlo_computation.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/service/hlo_opcode.h" +#include "tensorflow/compiler/xla/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/core/status.h" +#include "tensorflow/core/platform/logging.h" + +namespace xla { + +namespace { + +// ConvolutionVisitor traverses the HLO computation and rewrites Convolution +// operations with feature_group_count > 1 into convolutions with +// feature_group_count = 1. +class ConvolutionVisitor : public DfsHloVisitorWithDefault { + public: + // Default visitor action is to do nothing and return OK. + Status DefaultAction(HloInstruction* /*hlo_instruction*/) override { + return Status::OK(); + } + + Status HandleConvolution(HloInstruction* convolution) override; + + Status HandleBatchGroupCount(HloInstruction* convolution); + + // Runs the visitor on a computation. + static bool Run(HloComputation* computation, + std::function is_cost_viable, + bool convert_batch_groups_only, + bool canonicalize_depthwise_filter); + + // Returns whether any convolution ops were rewritten. + const bool changed() const { return changed_; } + + ~ConvolutionVisitor() override = default; + + private: + explicit ConvolutionVisitor( + HloComputation* computation, + std::function is_cost_viable, + bool convert_batch_groups_only, + bool canonicalize_depthwise_filter = false) + : computation_(computation), + filter_expansion_(!canonicalize_depthwise_filter), + convert_batch_groups_only_(convert_batch_groups_only), + is_cost_viable_(is_cost_viable) {} + + // Current HloComputation instance the ConvolutionVisitor is traversing. + HloComputation* computation_; + + // Whether rewrite has occurred. + bool changed_ = false; + + // Whether filter expansion is required. + bool filter_expansion_; + + // Decides whether to convert batch groups or feature groups. + bool convert_batch_groups_only_; + + // std::function(int64, int64)> chunk_fetcher + std::function is_cost_viable_; +}; + +bool ConvolutionVisitor::Run( + HloComputation* computation, + std::function is_cost_viable, + bool convert_batch_groups_only, bool canonicalize_depthwise_filter) { + ConvolutionVisitor visitor(computation, is_cost_viable, + convert_batch_groups_only, + canonicalize_depthwise_filter); + TF_CHECK_OK(computation->Accept(&visitor)); + return visitor.changed_; +} + +Shape ExpandedFilterShape(const Shape& shape, int64 group_count, + int64 input_feature_dim) { + int64 num_dims = shape.dimensions_size(); + CHECK_GE(num_dims, 2); + Shape expanded_shape = shape; + expanded_shape.set_dimensions( + input_feature_dim, shape.dimensions(input_feature_dim) * group_count); + return expanded_shape; +} + +// Returns a vector with 'group_count' many groups, where the i-th group +// consists of 'group_size' times the value i. +std::vector GetMaskIds(int64 group_size, int64 group_count) { + std::vector values; + for (int i = 0; i < group_count; ++i) { + for (int j = 0; j < group_size; ++j) { + values.push_back(i); + } + } + return values; +} + +// Create a mask for grouped convolution that will make a normal convolution +// produce the same results as a grouped convolution. For a [2, 1, 6] +// filter this returns a [2, 3, 6] mask +// 1 1 0 0 0 0 +// 0 0 1 1 0 0 +// 0 0 0 0 1 1 +// +// 1 1 0 0 0 0 +// 0 0 1 1 0 0 +// 0 0 0 0 1 1 +// +// The first step is to create a rank 1 constant: +// 0 1 2 +// +// This is broadcasted to +// 0 0 0 0 0 0 +// 1 1 1 1 1 1 +// 2 2 2 2 2 2 +// +// 0 0 0 0 0 0 +// 1 1 1 1 1 1 +// 2 2 2 2 2 2 +// +// Then we create another rank 1 constant +// 0 0 1 1 2 2 +// +// This is broadcasted to +// 0 0 1 1 2 2 +// 0 0 1 1 2 2 +// 0 0 1 1 2 2 +// +// 0 0 1 1 2 2 +// 0 0 1 1 2 2 +// 0 0 1 1 2 2 +// +// Finally we use the Eq op of these two broadcasted constants and get the +// desired mask. +HloInstruction* GetExpandedFilterMask( + const Shape& filter_shape, int64 kernel_input_feature_dim, + int64 kernel_output_feature_dim, int64 group_count, + const std::function)>& + add_instruction) { + Shape expanded_filter_shape = + ExpandedFilterShape(filter_shape, group_count, kernel_input_feature_dim); + Shape mask_shape = ShapeUtil::MakeShape( + S32, AsInt64Slice(expanded_filter_shape.dimensions())); + int64 output_feature = filter_shape.dimensions(kernel_output_feature_dim); + int64 group_size = filter_shape.dimensions(kernel_input_feature_dim); + + // Create a 'input_feature' sized linspace and 'output_feature' sized linspace + // that will be broadcasted into perpendicular dimensions and compared. + const std::vector input_feature_filter_mask = + GetMaskIds(group_size, group_count); + const std::vector output_feature_filter_mask = + GetMaskIds(output_feature / group_count, group_count); + auto mask1 = add_instruction(HloInstruction::CreateConstant( + LiteralUtil::CreateR1(input_feature_filter_mask))); + auto broadcasted_mask1 = add_instruction(HloInstruction::CreateBroadcast( + mask_shape, mask1, {kernel_input_feature_dim})); + auto mask2 = add_instruction(HloInstruction::CreateConstant( + LiteralUtil::CreateR1(output_feature_filter_mask))); + auto broadcasted_mask2 = add_instruction(HloInstruction::CreateBroadcast( + mask_shape, mask2, {kernel_output_feature_dim})); + + // Compare the broadcasted output feature linspace to the input feature + // linspace to create a diagonal predicate. + Shape predicate_shape = ShapeUtil::MakeShape( + PRED, AsInt64Slice(expanded_filter_shape.dimensions())); + return add_instruction(HloInstruction::CreateBinary( + predicate_shape, HloOpcode::kEq, broadcasted_mask1, broadcasted_mask2)); +} + +// This function handles batch_group_counts which are relevant only for +// depthwise backprop filter convolutions. +Status ConvolutionVisitor::HandleBatchGroupCount(HloInstruction* convolution) { + auto dim_numbers = convolution->convolution_dimension_numbers(); + auto activation = convolution->mutable_operand(0); + auto filter = convolution->mutable_operand(1); + int64 batch_group_count = convolution->batch_group_count(); + + if (batch_group_count == 1) { + return Status::OK(); + } + + 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 { + // 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 + // out the expanded regions. Next we reduce the filter in the batch + // dimension to obtain the original filter size. + + HloInstruction* filter_mask = + GetExpandedFilterMask(convolution->shape(), output_batch_dimension, + output_feature_dimension, batch_group_count, add); + auto expanded_filter_shape = ExpandedFilterShape( + convolution->shape(), batch_group_count, output_batch_dimension); + + auto new_convolution = add(HloInstruction::CreateConvolve( + expanded_filter_shape, activation, filter, + /*feature_group_count=*/1, /*batch_group_count=*/1, + convolution->window(), dim_numbers, convolution->precision_config())); + + auto zero = add(HloInstruction::CreateConstant( + LiteralUtil::Zero(expanded_filter_shape.element_type()))); + auto zero_filter = + add(HloInstruction::CreateBroadcast(expanded_filter_shape, zero, {})); + + auto new_filter = add(HloInstruction::CreateTernary( + 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)); + auto zero_scalar = + add(HloInstruction::CreateConstant(std::move(zero_literal))); + + auto reduce_function = [&]() -> HloComputation* { + HloComputation::Builder b("add_computation"); + Shape shape = ShapeUtil::MakeShape(F32, {}); + auto lhs = + b.AddInstruction(HloInstruction::CreateParameter(0, shape, "lhs")); + auto rhs = + b.AddInstruction(HloInstruction::CreateParameter(1, shape, "rhs")); + auto scalar_op = b.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, lhs, rhs)); + return computation_->parent()->AddEmbeddedComputation(b.Build(scalar_op)); + }; + + auto reduce_window_shape = new_convolution->shape(); + reduce_window_shape.set_dimensions(output_batch_dimension, 1); + + // Ensure that data input to reduce window is of type F32. + if (primitive_util::BitWidth(new_filter->shape().element_type()) < + primitive_util::BitWidth(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)); + } + + // Create the reduce window. + Window window; + for (int64 i = 0; i < new_convolution->shape().dimensions_size(); ++i) { + auto* dim = window.add_dimensions(); + dim->set_padding_low(0); + dim->set_padding_high(0); + dim->set_window_dilation(1); + dim->set_base_dilation(1); + if (i == output_batch_dimension) { + dim->set_stride(batch_group_count); + dim->set_size(batch_group_count); + } else { + dim->set_stride(1); + dim->set_size(1); + } + } + auto reduce_window = add(HloInstruction::CreateReduceWindow( + reduce_window_shape, new_filter, zero_scalar, window, + reduce_function())); + + Shape convert_back_shape = reduce_window->shape(); + convert_back_shape.set_element_type(activation->shape().element_type()); + + // Convert reduced data back to the original data type. + auto reduce_window_converted = + HloInstruction::CreateConvert(convert_back_shape, reduce_window); + + TF_RETURN_IF_ERROR(computation_->ReplaceWithNewInstruction( + convolution, std::move(reduce_window_converted))); + } + + return Status::OK(); +} + +Status ConvolutionVisitor::HandleConvolution(HloInstruction* convolution) { + if (convert_batch_groups_only_) { + return HandleBatchGroupCount(convolution); + } + + auto add = [&](std::unique_ptr inst) { + return computation_->AddInstruction(std::move(inst)); + }; + + int64 group_count = convolution->feature_group_count(); + if (group_count == 1) { + return Status::OK(); + } + + changed_ = true; + auto dim_numbers = convolution->convolution_dimension_numbers(); + auto filter = convolution->mutable_operand(1); + int64 kernel_input_feature_dim = dim_numbers.kernel_input_feature_dimension(); + int64 group_size = filter->shape().dimensions(kernel_input_feature_dim); + int64 kernel_output_feature_dim = + dim_numbers.kernel_output_feature_dimension(); + auto expanded_filter_shape = ExpandedFilterShape(filter->shape(), group_count, + kernel_input_feature_dim); + HloInstruction* filter_mask = + GetExpandedFilterMask(filter->shape(), kernel_input_feature_dim, + kernel_output_feature_dim, group_count, add); + HloInstruction* expanded_filter; + + if (group_size == 1) { + bool depthwise_separable = + (group_count == filter->shape().dimensions(kernel_output_feature_dim)); + // If the code generator handles depthwise separable convolutions + // inherently, then no filter expansion is needed. + if (!filter_expansion_ && depthwise_separable) { + changed_ = false; + return Status::OK(); + } + // We want to repeat 'filter' in the 'input_feature_dim' dimension + // 'group_count' times. + Shape reshaped_filter_shape = + ShapeUtil::DeleteDimension(kernel_input_feature_dim, filter->shape()); + auto reshaped_filter = + add(HloInstruction::CreateReshape(reshaped_filter_shape, filter)); + std::vector broadcast_dims; + for (int64 i = 0; i < filter->shape().dimensions_size(); ++i) { + if (i == kernel_input_feature_dim) { + continue; + } + broadcast_dims.push_back(i); + } + expanded_filter = add(HloInstruction::CreateBroadcast( + expanded_filter_shape, reshaped_filter, broadcast_dims)); + + auto zero = add(HloInstruction::CreateConstant( + LiteralUtil::Zero(expanded_filter_shape.element_type()))); + auto zero_filter = + add(HloInstruction::CreateBroadcast(expanded_filter_shape, zero, {})); + auto new_filter = add(HloInstruction::CreateTernary( + expanded_filter_shape, HloOpcode::kSelect, filter_mask, expanded_filter, + zero_filter)); + + auto new_convolution = HloInstruction::CreateConvolve( + convolution->shape(), convolution->mutable_operand(0), new_filter, + /*feature_group_count=*/1, /*batch_group_count=*/1, + convolution->window(), dim_numbers, convolution->precision_config()); + TF_RETURN_IF_ERROR(computation_->ReplaceWithNewInstruction( + convolution, std::move(new_convolution))); + } else { + int64 activation_input_feature_dim = dim_numbers.input_feature_dimension(); + + int64 output_feature = + filter->shape().dimensions(kernel_output_feature_dim); + + // If group_count == output_feature, then we map those grouped convolutions + // onto depthwise convolution. This is done by adding an additional spatial + // dimension to the activations, kernel, and the output. + // E.g., we would turn + // [2, 12]{B, IF} conv [3, 4]{IF, OF} into + // [3, 2, 4]{S, B, IF} depth conv [3, 1, 4]{S, IF, OF}, where S is the + // additional spatial dimension. The generated convolution output will be + // [1, 2, 4]{S, B, OF} and then reshape the output back to [2, 4] {B, OF}. + + if (group_count == output_feature && !filter_expansion_) { + auto filter = convolution->mutable_operand(1); + auto activation = convolution->mutable_operand(0); + + // Add spatial dimension to the activation, and reshape. + Shape reshaped_activation_shape = activation->shape(); + ShapeUtil::AppendMajorDimension(group_size, &reshaped_activation_shape); + + int64 new_spatial_dim = reshaped_activation_shape.dimensions().size() - 1; + + reshaped_activation_shape.set_dimensions(activation_input_feature_dim, + group_count); + activation = add( + HloInstruction::CreateReshape(reshaped_activation_shape, activation)); + + // Add spatial 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)); + + Shape new_output_shape = convolution->shape(); + ShapeUtil::AppendMajorDimension(1, &new_output_shape); + + // Edit convolution dimension numbers. Note that kernel_input_feature_dim + // now becomes a spatial dimension, and the newly added dimension of size + // 1 is the new kernel_input_feature_dim. + dim_numbers.add_input_spatial_dimensions(new_spatial_dim); + dim_numbers.add_kernel_spatial_dimensions(kernel_input_feature_dim); + dim_numbers.set_kernel_input_feature_dimension(new_spatial_dim); + dim_numbers.add_output_spatial_dimensions(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(group_size); + + auto new_convolution = add(HloInstruction::CreateConvolve( + new_output_shape, activation, filter, 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, new_convolution->shape()); + auto reshaped_convolution = HloInstruction::CreateReshape( + reshaped_convolution_shape, new_convolution); + + TF_RETURN_IF_ERROR(computation_->ReplaceWithNewInstruction( + convolution, std::move(reshaped_convolution))); + + } else { + // The filter expansion mechanism adds zeroes in the kernel. + // For an OF = 12, IF = 6, and kernel IF = 2, the expanded filter mask + // would look like (IF on the Y-axis, OF on the X-axis) + // 1 1 1 1 0 0 0 0 0 0 0 0 + // 1 1 1 1 0 0 0 0 0 0 0 0 + // 0 0 0 0 1 1 1 1 0 0 0 0 + // 0 0 0 0 1 1 1 1 0 0 0 0 + // 0 0 0 0 0 0 0 0 1 1 1 1 + // 0 0 0 0 0 0 0 0 1 1 1 1 + // + // Instead of convolving the above with the input, we instead slice the + // kernel into three kernels, each containing islands of 1s from the + // filter above. We also slice the activations in the IF dimension with + // each slice of size = group_size. For each slice, we perform + // convolutions, and concatenate the generated outputs in the output OF + // dimension. + + std::vector sliced_convolutions; + auto activation = convolution->mutable_operand(0); + std::vector slice_strides(filter->shape().dimensions_size(), 1); + std::vector filter_slice_starts(filter->shape().dimensions_size(), + 0); + std::vector filter_slice_limits( + filter->shape().dimensions().begin(), + filter->shape().dimensions().end()); + std::vector activation_slice_starts( + activation->shape().dimensions_size(), 0); + std::vector activation_slice_limits( + activation->shape().dimensions().begin(), + activation->shape().dimensions().end()); + + int64 output_feature = + filter->shape().dimensions(kernel_output_feature_dim); + auto output_feature_dim = dim_numbers.output_feature_dimension(); + int64 filter_slice_width = output_feature / group_count; + + int64 activation_input_feature_dim = + dim_numbers.input_feature_dimension(); + + for (int64 i = 0; i < group_count; i++) { + filter_slice_starts[kernel_output_feature_dim] = i * filter_slice_width; + filter_slice_limits[kernel_output_feature_dim] = + (i + 1) * filter_slice_width; + auto filter_sliced_shape = filter->shape(); + filter_sliced_shape.set_dimensions(kernel_output_feature_dim, + filter_slice_width); + auto filter_slice = add(HloInstruction::CreateSlice( + filter_sliced_shape, filter, filter_slice_starts, + filter_slice_limits, slice_strides)); + + activation_slice_starts[activation_input_feature_dim] = i * group_size; + activation_slice_limits[activation_input_feature_dim] = + (i + 1) * group_size; + auto activation_sliced_shape = activation->shape(); + activation_sliced_shape.set_dimensions(activation_input_feature_dim, + group_size); + auto activation_slice = add(HloInstruction::CreateSlice( + activation_sliced_shape, activation, activation_slice_starts, + activation_slice_limits, slice_strides)); + + auto conv_slice_shape = convolution->shape(); + conv_slice_shape.set_dimensions(output_feature_dim, filter_slice_width); + + auto new_convolution = add(HloInstruction::CreateConvolve( + conv_slice_shape, activation_slice, filter_slice, + /*feature_group_count=*/1, /*batch_group_count=*/1, + convolution->window(), dim_numbers, + convolution->precision_config())); + + sliced_convolutions.push_back(new_convolution); + } + + auto new_conv = HloInstruction::CreateConcatenate( + convolution->shape(), sliced_convolutions, output_feature_dim); + TF_RETURN_IF_ERROR(computation_->ReplaceWithNewInstruction( + convolution, std::move(new_conv))); + } + } + + return Status::OK(); +} + +} // namespace + +StatusOr ConvolutionGroupConverter::Run(HloModule* module) { + XLA_VLOG_LINES( + 2, "ConvolutionGroupConverter::Run(), before:\n" + module->ToString()); + bool changed = false; + for (auto* comp : module->MakeNonfusionComputations()) { + if (ConvolutionVisitor::Run(comp, is_cost_viable_, + convert_batch_groups_only_, + filter_expansion_)) { + changed = true; + } + } + XLA_VLOG_LINES( + 2, "ConvolutionGroupConverter::Run(), after:\n" + module->ToString()); + return changed; +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/convolution_group_converter.h b/tensorflow/compiler/xla/service/convolution_group_converter.h new file mode 100644 index 0000000000000000000000000000000000000000..1caf1841119a965044502435fe0f5b38ca94f6a5 --- /dev/null +++ b/tensorflow/compiler/xla/service/convolution_group_converter.h @@ -0,0 +1,58 @@ +/* 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_CONVOLUTION_GROUP_CONVERTER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_CONVOLUTION_GROUP_CONVERTER_H_ + +#include "absl/strings/string_view.h" +#include "tensorflow/compiler/xla/service/hlo_module.h" +#include "tensorflow/compiler/xla/service/hlo_pass_interface.h" +#include "tensorflow/compiler/xla/status_macros.h" + +namespace xla { + +// A pass which rewrites convolutions with feature_group_count > 1 into +// convolutions with feature_group_count = 1. +class ConvolutionGroupConverter : public HloModulePass { + public: + ConvolutionGroupConverter(std::function is_cost_viable, + bool convert_batch_groups_only, + bool canonicalize_depthwise_filter = false) + : is_cost_viable_(is_cost_viable), + convert_batch_groups_only_(convert_batch_groups_only), + filter_expansion_(canonicalize_depthwise_filter) {} + + absl::string_view name() const override { + return "convolution-group-converter"; + } + + // Run convolution rewriting on the given computation. Returns whether the + // computation was changed. + StatusOr Run(HloModule* module) override; + + // Lambda containing cost model that decides whether to expand + // batch_group_count. + std::function is_cost_viable_; + + // Decides whether to convert batch groups or feature groups. + bool convert_batch_groups_only_; + + // Tells whether filter expansion is required. + bool filter_expansion_; +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_CONVOLUTION_GROUP_CONVERTER_H_ diff --git a/tensorflow/compiler/xla/service/convolution_group_converter_test.cc b/tensorflow/compiler/xla/service/convolution_group_converter_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..585b81a5db632901be863893bf723fcba19388ea --- /dev/null +++ b/tensorflow/compiler/xla/service/convolution_group_converter_test.cc @@ -0,0 +1,125 @@ +/* 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/convolution_group_converter.h" + +#include +#include + +#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_opcode.h" +#include "tensorflow/compiler/xla/service/hlo_parser.h" +#include "tensorflow/compiler/xla/test.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" +#include "tensorflow/compiler/xla/types.h" + +namespace xla { +namespace { + +using ConvolutionGroupConverterTest = HloTestBase; +namespace op = testing::opcode_matchers; + +TEST_F(ConvolutionGroupConverterTest, + ConvertFeatureGroupCountEqualToInputFeatureDim) { + string hlo_string = R"(HloModule Convolve1D1Window_0_module + +ENTRY %Convolve1D1Window_0.v3 (input: f32[1,2,2], filter: f32[1,1,2]) -> f32[1,2,2] { + %input = f32[1,2,2]{2,1,0} parameter(0) + %copy = f32[1,2,2]{2,0,1} copy(f32[1,2,2]{2,1,0} %input) + %filter = f32[1,1,2]{2,1,0} parameter(1) + ROOT %convolution = f32[1,2,2]{2,0,1} convolution(f32[1,2,2]{2,0,1} %copy, f32[1,1,2]{2,1,0} %filter), window={size=1}, dim_labels=b0f_0io->b0f, feature_group_count=2 +})"; + 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); + ConvolutionGroupConverter converter(nullptr, /*convert_batch_groups_only=*/ + false); + ASSERT_TRUE(converter.Run(module.get()).ValueOrDie()); + root = computation->root_instruction(); + // Make sure the convolution is converted to one with feature_group_count = 1. + EXPECT_EQ(root->opcode(), HloOpcode::kConvolution); + EXPECT_EQ(root->feature_group_count(), 1); + // Verify that the filter operand has been replaced. + EXPECT_THAT(root->operand(1), + op::Select(op::Eq(op::Broadcast(op::Constant()), + op::Broadcast(op::Constant())), + op::Broadcast(op::Reshape(op::Parameter())), + op::Broadcast(op::Constant()))); +} + +TEST_F(ConvolutionGroupConverterTest, + ConvertFeatureGroupCountDivisorOfInputFeatureDim) { + string hlo_string = R"(HloModule Convolve1D1Window_0_module + +ENTRY %Convolve1D1Window_0.v3 (input: f32[1,2,4], filter: f32[1,2,2]) -> f32[1,2,2] { + %input = f32[1,2,4]{2,1,0} parameter(0) + %copy = f32[1,2,4]{2,0,1} copy(f32[1,2,4]{2,1,0} %input) + %filter = f32[1,2,2]{2,1,0} parameter(1) + ROOT %convolution = f32[1,2,2]{2,0,1} convolution(f32[1,2,4]{2,0,1} %copy, f32[1,2,2]{2,1,0} %filter), window={size=1}, dim_labels=b0f_0io->b0f, feature_group_count=2 +})"; + 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); + ConvolutionGroupConverter converter(nullptr, /*convert_batch_groups_only=*/ + false); + ASSERT_TRUE(converter.Run(module.get()).ValueOrDie()); + root = computation->root_instruction(); + // Make sure the convolution is replaced with a concatenate. + EXPECT_EQ(root->opcode(), HloOpcode::kConcatenate); + // And the operands of the concatenate are convolutions, each with a feature + // group count = 1. + EXPECT_EQ(root->operand(0)->opcode(), HloOpcode::kConvolution); + EXPECT_EQ(root->operand(1)->opcode(), HloOpcode::kConvolution); + EXPECT_EQ(root->operand(0)->feature_group_count(), 1); + 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 true; }; + ConvolutionGroupConverter converter(cost_model, /*convert_batch_groups_only=*/ + true); + ASSERT_TRUE(converter.Run(module.get()).ValueOrDie()); + root = computation->root_instruction(); + // Make sure the convolution is converted to one with batch_group_count = 1. + EXPECT_EQ(root->operand(0)->opcode(), HloOpcode::kConvolution); + EXPECT_EQ(root->operand(0)->batch_group_count(), 1); + // Verify that the convolution is replaced by a reshape. + EXPECT_EQ(root->opcode(), HloOpcode::kReshape); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/service/copy_insertion.cc b/tensorflow/compiler/xla/service/copy_insertion.cc index 245db6be2a400a7447f1e87317018cbb1572c405..5e26a63cebfa9b2e50f4b13335c10c246999d4df 100644 --- a/tensorflow/compiler/xla/service/copy_insertion.cc +++ b/tensorflow/compiler/xla/service/copy_insertion.cc @@ -341,18 +341,20 @@ Status AddCopiesForAliasedInputOutputs(HloModule* module) { HloInstruction* root = entry->root_instruction(); ShapeTree output_indices_to_copy(root->shape()); - std::vector> copied_parameters; + std::vector>> copied_parameters( + entry->num_parameters()); bool has_alias = false; for (auto* param : entry->parameter_instructions()) { bool param_has_alias = false; 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; } }); @@ -361,6 +363,9 @@ Status AddCopiesForAliasedInputOutputs(HloModule* module) { continue; } + TF_RET_CHECK(param->parameter_number() < entry->num_parameters()); + TF_RET_CHECK(!copied_parameters[param->parameter_number()]); + has_alias = true; // Store a snapshot of users before DeepCopyInstruction, as // DeepCopyInstruction introduces new users of the instruction. @@ -374,7 +379,7 @@ Status AddCopiesForAliasedInputOutputs(HloModule* module) { TF_RETURN_IF_ERROR(param->ReplaceUseWith(user, copied)); } - copied_parameters.push_back(param_copy_tree); + copied_parameters[param->parameter_number()] = param_copy_tree; } if (!has_alias) { @@ -391,10 +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 { + [&](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); @@ -435,7 +444,6 @@ class CopyRemover { const HloOrdering& ordering, HloModule* module) : module_(module), alias_analysis_(alias_analysis), - ordering_(ordering), buffer_value_tracker_(*module, alias_analysis, ordering) {} // Try to elide the given copy. The copy is elided if the instruction is not @@ -516,7 +524,7 @@ class CopyRemover { // between copies added around aliased operations (kWhile) guarantees // this strict order. for (const HloValue* value_a : buffer.values()) { - if (ShapeUtil::IsToken(value_a->shape())) { + if (value_a->shape().IsToken()) { // Token values have no representation and cannot interfere. continue; } @@ -533,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); @@ -836,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); @@ -996,7 +1002,6 @@ class CopyRemover { HloModule* module_; const HloAliasAnalysis& alias_analysis_; - const HloOrdering& ordering_; // Object tracking the HLO values contained in each HLO buffer. BufferValueTracker buffer_value_tracker_; diff --git a/tensorflow/compiler/xla/service/copy_insertion.h b/tensorflow/compiler/xla/service/copy_insertion.h index c097089e30d59936a32f69c49123c398f1611ea3..8866b5050bf1e7419dda6496ea95d034178d25d8 100644 --- a/tensorflow/compiler/xla/service/copy_insertion.h +++ b/tensorflow/compiler/xla/service/copy_insertion.h @@ -94,10 +94,12 @@ class CopyInsertion : public HloModulePass { Status VerifyNoLiveRangeInterference(const HloOrdering& ordering, HloModule* module); - private: + protected: // Override which requires the caller to pass in a call graph. - Status AddSpecialCaseCopies(const CallGraph& call_graph, HloModule* module); + virtual Status AddSpecialCaseCopies(const CallGraph& call_graph, + HloModule* module); + private: Status AddCopiesToResolveInterference(HloModule* module); // Backend specific function that decides whether a fusion can share buffer diff --git a/tensorflow/compiler/xla/service/copy_insertion_test.cc b/tensorflow/compiler/xla/service/copy_insertion_test.cc index 4533ebb99bbba854a029fb8a9a1e31b023be720d..4391bdcba532661a0fde789e2c4ed324c40bcd32 100644 --- a/tensorflow/compiler/xla/service/copy_insertion_test.cc +++ b/tensorflow/compiler/xla/service/copy_insertion_test.cc @@ -17,7 +17,7 @@ limitations under the License. #include -#include "tensorflow/compiler/xla/legacy_flags/debug_options_flags.h" +#include "tensorflow/compiler/xla/debug_options_flags.h" #include "tensorflow/compiler/xla/literal.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" @@ -94,7 +94,7 @@ TEST_F(CopyInsertionTest, SingleParameter) { EXPECT_THAT(x->users(), UnorderedElementsAre(tuple)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); InsertCopies(module.get()); @@ -114,7 +114,7 @@ TEST_F(CopyInsertionTest, SingleConstant) { EXPECT_THAT(constant->users(), UnorderedElementsAre(tuple)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); InsertCopies(module.get()); @@ -127,7 +127,7 @@ TEST_F(CopyInsertionTest, SingleConstant) { TEST_F(CopyInsertionTest, ExistingCopiesNotRemoved) { // Verify that kCopy instructions which change layout and exist before // copy-insertion remain in the graph after copy-insertion. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto builder = HloComputation::Builder(TestName()); HloInstruction* constant = @@ -181,7 +181,7 @@ TEST_F(CopyInsertionTest, MultipleConstantsAndParameters) { builder.AddInstruction(HloInstruction::CreateTuple({constant2, x, add})); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); InsertCopies(module.get()); @@ -217,7 +217,7 @@ TEST_F(CopyInsertionTest, AmbiguousPointsToSet) { EXPECT_THAT(constant2->users(), UnorderedElementsAre(tuple1, tuple2)); EXPECT_THAT(constant3->users(), UnorderedElementsAre(tuple2)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); HloInstruction* old_root = module->entry_computation()->root_instruction(); @@ -238,7 +238,7 @@ TEST_F(CopyInsertionTest, BitcastParameter) { HloInstruction* bitcast = builder.AddInstruction(HloInstruction::CreateUnary( ShapeUtil::MakeShape(F32, {2, 2}), HloOpcode::kBitcast, x)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); EXPECT_THAT(x->users(), UnorderedElementsAre(bitcast)); @@ -261,7 +261,7 @@ TEST_F(CopyInsertionTest, BitcastConstant) { HloInstruction* bitcast = builder.AddInstruction(HloInstruction::CreateUnary( ShapeUtil::MakeShape(F32, {2, 2}), HloOpcode::kBitcast, constant)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); EXPECT_THAT(constant->users(), UnorderedElementsAre(bitcast)); @@ -283,7 +283,7 @@ TEST_F(CopyInsertionTest, BitcastTupleElementParameter) { ShapeUtil::MakeShape(F32, {2, 2}), HloOpcode::kBitcast, x)); builder.AddInstruction(HloInstruction::CreateTuple({bitcast})); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); EXPECT_THAT(x->users(), UnorderedElementsAre(bitcast)); @@ -310,7 +310,7 @@ TEST_F(CopyInsertionTest, NestedTupleParameter) { ShapeUtil::MakeShape(F32, {42})}), "param0")); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); EXPECT_EQ(HloOpcode::kParameter, @@ -351,7 +351,7 @@ TEST_F(CopyInsertionTest, ElementOfNestedTupleParameter) { auto gte = builder.AddInstruction(HloInstruction::CreateGetTupleElement( ShapeUtil::GetSubshape(param->shape(), {0}), param, 0)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); EXPECT_EQ(gte, module->entry_computation()->root_instruction()); @@ -388,7 +388,7 @@ TEST_F(CopyInsertionTest, AmbiguousTopLevelRoot) { builder.AddInstruction(HloInstruction::CreateGetTupleElement( ShapeUtil::GetSubshape(select->shape(), {0}), select, 0)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); EXPECT_EQ(gte, module->entry_computation()->root_instruction()); @@ -403,7 +403,7 @@ TEST_F(CopyInsertionTest, AmbiguousTopLevelRoot) { class WhileCopyInsertionTest : public CopyInsertionTest { protected: - WhileCopyInsertionTest() : module_(CreateNewModule()) {} + WhileCopyInsertionTest() : module_(CreateNewUnverifiedModule()) {} // Builds a While condition computation which reads the induction variable // from the tuple parameter, and returns a predicate indicating whether this @@ -1295,7 +1295,7 @@ TEST_F(WhileCopyInsertionTest, InitPointsToNonDistinctUsedByTwoWhileLoops) { TEST_F(CopyInsertionTest, SwizzlingWhile) { // Test a while instruction with a body which permutes its tuple parameter // elements. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); const Shape loop_state_shape = ShapeUtil::MakeTupleShape({scalar_shape_, scalar_shape_}); @@ -1362,7 +1362,7 @@ TEST_F(CopyInsertionTest, CrossingParameters) { // | / \ | // | / \| // (p1 , p0) - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); const Shape tuple_shape = ShapeUtil::MakeTupleShape({scalar_shape_, scalar_shape_}); @@ -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); @@ -1395,7 +1397,7 @@ TEST_F(CopyInsertionTest, ParametersAliasing) { // | | // | | // (p0 , p1) - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); const Shape tuple_shape = ShapeUtil::MakeTupleShape({scalar_shape_, scalar_shape_}); @@ -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); @@ -1428,7 +1432,7 @@ TEST_F(CopyInsertionTest, ParameterWithNoAliasing) { // | | // | | // (p0 , p1) - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); const Shape tuple_shape = ShapeUtil::MakeTupleShape({scalar_shape_, scalar_shape_}); @@ -1461,7 +1465,7 @@ TEST_F(CopyInsertionTest, ParameterWithPartialAliasing) { // | | // | | // (p0 , p1) - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); const Shape tuple_shape = ShapeUtil::MakeTupleShape({scalar_shape_, scalar_shape_}); @@ -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(), @@ -1496,7 +1501,7 @@ TEST_F(CopyInsertionTest, ParameterAndParallelOpsWithPartialAliasing) { // | | | // | | | // +-- (p0 , p1) - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); const Shape tuple_shape = ShapeUtil::MakeTupleShape({scalar_shape_, scalar_shape_}); @@ -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); @@ -1534,7 +1540,7 @@ TEST_F(CopyInsertionTest, ParameterAndOpsWithPartialAliasing) { // | Add----+ // | | | // +-- (p0 , p1) - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); const Shape tuple_shape = ShapeUtil::MakeTupleShape({scalar_shape_, scalar_shape_}); @@ -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); @@ -1569,7 +1576,7 @@ TEST_F(CopyInsertionTest, SwizzlingWhileWithOneOp) { // the operation (instruction) on the element makes the live range of the // respective input and output elements different than if the instruction were // not there (as in the SwizzlingWhile test above). - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); const Shape loop_state_shape = ShapeUtil::MakeTupleShape({scalar_shape_, scalar_shape_}); @@ -1632,7 +1639,7 @@ TEST_F(CopyInsertionTest, SwizzlingWhileSharedInput) { // the while body is a single constant (both loop state elements are the same // constant). This means no copies are necessary because both loop state // elements are the same so interchanging them is a no-op. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); const Shape loop_state_shape = ShapeUtil::MakeTupleShape({scalar_shape_, scalar_shape_}); @@ -1693,7 +1700,7 @@ TEST_F(CopyInsertionTest, SequentialWhiles) { const Shape loop_state_shape = ShapeUtil::MakeTupleShape( {element_shape, element_shape, element_shape, element_shape}); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto builder = HloComputation::Builder(TestName()); auto param_0 = builder.AddInstruction( HloInstruction::CreateParameter(0, element_shape, "param_0")); @@ -1783,7 +1790,7 @@ TEST_F(CopyInsertionTest, SequentialWhiles) { TEST_F(CopyInsertionTest, WhileBodyWithConstantRoot) { // Test a while body and condition which are each simply a constant (root of // computation is a constant). The body constant should be copied. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto builder = HloComputation::Builder(TestName()); auto param_0 = builder.AddInstruction( HloInstruction::CreateParameter(0, scalar_shape_, "param_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. @@ -1896,7 +1902,7 @@ void BM_SequentialWhiles(int num_iters, int num_whiles) { tensorflow::testing::StopTiming(); for (int i = 0; i < num_iters; ++i) { HloModuleConfig config; - config.set_debug_options(legacy_flags::GetDebugOptionsFromFlags()); + config.set_debug_options(GetDebugOptionsFromFlags()); HloModule module("BM_SequentialWhiles", config); auto builder = HloComputation::Builder("BM_SequentialWhiles"); @@ -1936,7 +1942,7 @@ void BM_ParallelWhiles(int num_iters, int num_whiles) { tensorflow::testing::StopTiming(); for (int i = 0; i < num_iters; ++i) { HloModuleConfig config; - config.set_debug_options(legacy_flags::GetDebugOptionsFromFlags()); + config.set_debug_options(GetDebugOptionsFromFlags()); HloModule module("BM_SequentialWhiles", config); auto builder = HloComputation::Builder("BM_ParallelWhiles"); @@ -2003,7 +2009,7 @@ std::unique_ptr MakeBenchmarkWhileBody( void BM_ManyElementTuple(int num_iters, const int num_tuple_inputs) { tensorflow::testing::StopTiming(); HloModuleConfig config; - config.set_debug_options(legacy_flags::GetDebugOptionsFromFlags()); + config.set_debug_options(GetDebugOptionsFromFlags()); CopyInsertion copy_insertion; const Shape element_shape = ShapeUtil::MakeShape(F32, {}); std::vector tuple_params(num_tuple_inputs); @@ -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 58abb330a6e31e9b7a8081cd7964cf89a5b64a09..d4535b204d7f3ad8d4e24beea5d0dd79e7a15ab0 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", @@ -51,6 +50,7 @@ cc_library( "//tensorflow/core:lib", "//tensorflow/core:stream_executor_no_cuda", "//tensorflow/stream_executor", + "@com_google_absl//absl/base", "@com_google_absl//absl/memory", "@com_google_absl//absl/types:span", ], @@ -94,7 +94,9 @@ 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", "//tensorflow/compiler/xla:literal", "//tensorflow/compiler/xla:protobuf_util", @@ -110,8 +112,9 @@ cc_library( "//tensorflow/compiler/xla/service:buffer_liveness", "//tensorflow/compiler/xla/service:call_inliner", "//tensorflow/compiler/xla/service:conditional_simplifier", - "//tensorflow/compiler/xla/service:convolution_feature_group_converter", + "//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", @@ -131,6 +134,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", @@ -239,6 +243,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", @@ -362,15 +367,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", @@ -378,6 +401,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", @@ -570,6 +594,7 @@ cc_library( ":runtime_matvec", "//tensorflow/compiler/xla:executable_run_options", "//tensorflow/core:framework_lite", + "//tensorflow/core/kernels:eigen_contraction_kernel", "//third_party/eigen3", ], ) @@ -628,6 +653,7 @@ cc_library( deps = [ ":runtime_matvec", "//tensorflow/core:framework_lite", + "//tensorflow/core/kernels:eigen_contraction_kernel", "//third_party/eigen3", ], ) @@ -823,7 +849,6 @@ tf_cc_test( "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", "//tensorflow/compiler/xla/tests:xla_internal_test_main", ], ) @@ -845,7 +870,6 @@ tf_cc_test( "//tensorflow/compiler/xla:test_helpers", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", "//tensorflow/compiler/xla/tests:xla_internal_test_main", ], ) @@ -886,7 +910,6 @@ tf_cc_test( "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/service:hlo_matchers", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", "//tensorflow/compiler/xla/tests:test_utils", "//tensorflow/core:lib", "//tensorflow/core:test", @@ -960,17 +983,16 @@ tf_cc_test( srcs = ["cpu_copy_insertion_test.cc"], deps = [ ":cpu_copy_insertion", + "//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/legacy_flags:debug_options_flags", "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/service:hlo_graph_dumper", "//tensorflow/compiler/xla/service:hlo_matchers", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:test", ], @@ -996,7 +1018,6 @@ tf_cc_test( "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:protos_all_cc", "//tensorflow/core:test", @@ -1008,7 +1029,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 73b03440cbb936017257b8a92f16dcc25d41e21c..796a7cf94d02b0ad42366387a9d3f8d589b8840a 100644 --- a/tensorflow/compiler/xla/service/cpu/compiler_functor.cc +++ b/tensorflow/compiler/xla/service/cpu/compiler_functor.cc @@ -61,19 +61,6 @@ Disabling these as a starting point. // TODO(b/64227304) Creating a custom pass pipeline will replace this. namespace { -class FilteredFunctionPassManager : public llvm::legacy::FunctionPassManager { - public: - FilteredFunctionPassManager(llvm::Module* m, bool disable_expensive_passes) - : llvm::legacy::FunctionPassManager(m), - disable_expensive_passes_(disable_expensive_passes) {} - void add(llvm::Pass* p) override { - llvm::legacy::FunctionPassManager::add(p); - } - - private: - bool disable_expensive_passes_; -}; - class FilteredPassManager : public llvm::legacy::PassManager { public: explicit FilteredPassManager(bool disable_expensive_passes) @@ -96,8 +83,7 @@ class FilteredPassManager : public llvm::legacy::PassManager { std::unique_ptr CompilerFunctor::operator()( llvm::Module& module) const { FilteredPassManager module_passes(disable_expensive_passes_); - FilteredFunctionPassManager function_passes(&module, - disable_expensive_passes_); + llvm::legacy::FunctionPassManager function_passes(&module); VLOG(2) << "IR before optimizations"; XLA_VLOG_LINES(2, llvm_ir::DumpModuleToString(module)); diff --git a/tensorflow/compiler/xla/service/cpu/conv_canonicalization.cc b/tensorflow/compiler/xla/service/cpu/conv_canonicalization.cc index 2d9978404cc9ec1e40fc61aaf794a8f1f06050bb..8e55267a67d330e7e721f9b5fb25451357a49a9d 100644 --- a/tensorflow/compiler/xla/service/cpu/conv_canonicalization.cc +++ b/tensorflow/compiler/xla/service/cpu/conv_canonicalization.cc @@ -132,7 +132,8 @@ StatusOr ConvCanonicalization::Run(HloModule* module) { HloInstruction* new_conv = module->entry_computation()->AddInstruction( HloInstruction::CreateConvolve( new_conv_shape, new_input, new_kernel, hlo->feature_group_count(), - hlo->window(), new_dnums, hlo->precision_config())); + hlo->batch_group_count(), hlo->window(), new_dnums, + hlo->precision_config())); // Reshape the output back to the shape of the original convolution. TF_RETURN_IF_ERROR(module->entry_computation()->ReplaceWithNewInstruction( diff --git a/tensorflow/compiler/xla/service/cpu/conv_canonicalization_test.cc b/tensorflow/compiler/xla/service/cpu/conv_canonicalization_test.cc index 2083f440fdd971db1b675d005664d25e6de53dbe..02085108a081358cd4f8aed6dc12557cbd8eea85 100644 --- a/tensorflow/compiler/xla/service/cpu/conv_canonicalization_test.cc +++ b/tensorflow/compiler/xla/service/cpu/conv_canonicalization_test.cc @@ -22,7 +22,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/test.h" -#include "tensorflow/compiler/xla/tests/hlo_verified_test_base.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" #include "tensorflow/compiler/xla/util.h" #include "tensorflow/compiler/xla/test_helpers.h" @@ -32,7 +32,7 @@ namespace cpu { using ::testing::ElementsAre; -class ConvCanonicalizationTest : public HloVerifiedTestBase { +class ConvCanonicalizationTest : public HloTestBase { public: ConvCanonicalizationTest() { for (int i = 0; i < 2; ++i) { @@ -84,10 +84,10 @@ TEST_F(ConvCanonicalizationTest, NonCanonicalToCanonical) { builder.AddInstruction(HloInstruction::CreateConvolve( ShapeUtil::MakeShape( F32, {kOutputFeatureCount, kBatchSize, output_size, output_size}), - input, kernel, /*feature_group_count=*/1, conv_window_, dnums, - DefaultPrecisionConfig(2))); + input, kernel, /*feature_group_count=*/1, /*batch_group_count=*/1, + conv_window_, dnums, DefaultPrecisionConfig(2))); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* entry_computation = module->AddEntryComputation(builder.Build()); @@ -96,7 +96,7 @@ TEST_F(ConvCanonicalizationTest, NonCanonicalToCanonical) { return cpu::TargetMachineFeatures::kEigenExpectedTensorAlignment; }); ConvCanonicalization conv_canonicalization(&target_machine_features); - EXPECT_TRUE(conv_canonicalization.Run(module).ValueOrDie()); + EXPECT_TRUE(conv_canonicalization.Run(module.get()).ValueOrDie()); const HloInstruction* output_reshape = entry_computation->root_instruction(); EXPECT_EQ(HloOpcode::kTranspose, output_reshape->opcode()); @@ -147,10 +147,10 @@ TEST_F(ConvCanonicalizationTest, CanonicalStaysTheSame) { builder.AddInstruction(HloInstruction::CreateConvolve( ShapeUtil::MakeShape( F32, {kBatchSize, output_size, output_size, kOutputFeatureCount}), - input, kernel, /*feature_group_count=*/1, conv_window_, dnums, - DefaultPrecisionConfig(2))); + input, kernel, /*feature_group_count=*/1, /*batch_group_count=*/1, + conv_window_, dnums, DefaultPrecisionConfig(2))); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); cpu::TargetMachineFeaturesWithFakeAlignmentLogic target_machine_features( @@ -158,7 +158,7 @@ TEST_F(ConvCanonicalizationTest, CanonicalStaysTheSame) { return cpu::TargetMachineFeatures::kEigenExpectedTensorAlignment; }); ConvCanonicalization conv_canonicalization(&target_machine_features); - EXPECT_FALSE(conv_canonicalization.Run(module).ValueOrDie()); + EXPECT_FALSE(conv_canonicalization.Run(module.get()).ValueOrDie()); } } // namespace cpu diff --git a/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc b/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc index da01c0caf2a6665f71cc087270b21fffdd6caa0d..eafda68510d93ee54f2aead60a84f3e97b3fe1f4 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc @@ -51,7 +51,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/buffer_liveness.h" #include "tensorflow/compiler/xla/service/call_inliner.h" #include "tensorflow/compiler/xla/service/conditional_simplifier.h" -#include "tensorflow/compiler/xla/service/convolution_feature_group_converter.h" +#include "tensorflow/compiler/xla/service/convolution_group_converter.h" #include "tensorflow/compiler/xla/service/cpu/buffer_info_util.h" #include "tensorflow/compiler/xla/service/cpu/compiler_functor.h" #include "tensorflow/compiler/xla/service/cpu/conv_canonicalization.h" @@ -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" @@ -76,6 +77,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_cse.h" #include "tensorflow/compiler/xla/service/hlo_dce.h" #include "tensorflow/compiler/xla/service/hlo_element_type_converter.h" +#include "tensorflow/compiler/xla/service/hlo_get_dimension_size_rewriter.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_memory_scheduler.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" @@ -91,6 +93,7 @@ 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/tuple_simplifier.h" #include "tensorflow/compiler/xla/service/while_loop_constant_sinking.h" @@ -243,6 +246,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( @@ -255,8 +259,17 @@ Status CpuCompiler::RunHloPassesThroughLayoutAssn( // pass. pipeline.AddPass(); 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; + }; + pipeline.AddPass( + cost_model, + /*convert_batch_groups_only=*/true); + pipeline.AddPass( + cost_model, + /*convert_batch_groups_only=*/false); pipeline.AddPass(target_machine_features); { auto& pass = @@ -268,10 +281,11 @@ Status CpuCompiler::RunHloPassesThroughLayoutAssn( /*rewrite_training_op=*/true, /*rewrite_inference_op=*/true, /*rewrite_grad_op=*/true); - pass.AddPass( - /*is_layout_sensitive=*/false, - [](const Shape&, const Shape&) { return false; }, - /*enable_dot_strength_reduction=*/false); + pipeline.AddPass(); + 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 @@ -291,7 +305,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{}; }, @@ -334,10 +349,10 @@ Status CpuCompiler::RunHloPassesAfterLayoutAssn( pass.AddInvariantChecker( /*layout_sensitive=*/true, /*allow_mixed_precision=*/false); - pass.AddPass>( - /*is_layout_sensitive=*/true, - [](const Shape&, const Shape&) { return true; }, - /*enable_dot_strength_reduction=*/false); + AlgebraicSimplifierOptions options; + options.set_is_layout_sensitive(true); + options.set_enable_dot_strength_reduction(false); + pass.AddPass>(options); pass.AddPass(); pass.AddPass(/*is_layout_sensitive=*/true); } @@ -494,7 +509,7 @@ Status CreateHloProfilingArtifacts( auto shape_size_bytes = [](const Shape& shape) { // On the cpu, opaques are pointers. - if (ShapeUtil::IsOpaque(shape)) { + if (shape.IsOpaque()) { return static_cast(sizeof(void*)); } return ShapeUtil::ByteSizeOf(shape, sizeof(void*)); @@ -502,8 +517,8 @@ Status CreateHloProfilingArtifacts( HloCostAnalysis cost_analysis(shape_size_bytes); TF_RETURN_IF_ERROR(entry_computation.Accept(&cost_analysis)); - *hlo_profile_printer_data = - CreateHloProfilePrinterData(**hlo_profile_index_map, cost_analysis); + *hlo_profile_printer_data = CreateHloProfilePrinterData( + **hlo_profile_index_map, cost_analysis, entry_computation.name()); *computation_to_profile_idx = (*hlo_profile_index_map)->computation_to_profile_idx(); @@ -587,9 +602,9 @@ StatusOr> CpuCompiler::RunBackend( // Select an order for emitting the HLO instructions for each // computation. Using this sequence enables tighter buffer liveness analysis // and reduced memory usage (as compared to using DependencyHloOrdering). - TF_ASSIGN_OR_RETURN( - HloSchedule schedule, - ScheduleModule(*module, BufferSizeBytesFunction(), DFSMemoryScheduler)); + TF_ASSIGN_OR_RETURN(HloSchedule schedule, + ScheduleModule(module.get(), BufferSizeBytesFunction(), + DFSMemoryScheduler)); // Run buffer allocation on the HLO graph. TF_ASSIGN_OR_RETURN( @@ -632,18 +647,17 @@ StatusOr> CpuCompiler::RunBackend( .EmitComputation( embedded_computation, embedded_computation->name(), /*is_top_level_computation=*/false, - &schedule.sequence(embedded_computation).instructions()) + schedule.sequence(embedded_computation).instructions()) .status()); } string function_name_prefix = entry_computation->name().empty() ? "__compute" : entry_computation->name(); - TF_ASSIGN_OR_RETURN( - llvm::Function * entry_function, - ir_emitter.EmitComputation( - entry_computation, function_name_prefix, - /*is_top_level_computation=*/true, - &schedule.sequence(entry_computation).instructions())); + TF_ASSIGN_OR_RETURN(llvm::Function * entry_function, + ir_emitter.EmitComputation( + entry_computation, function_name_prefix, + /*is_top_level_computation=*/true, + schedule.sequence(entry_computation).instructions())); string function_name = [&]() { llvm::SmallVector function_name_vector; @@ -779,7 +793,7 @@ CpuCompiler::CompileAheadOfTime(std::unique_ptr module_group, XLA_VLOG_LINES(2, module->ToString()); TF_ASSIGN_OR_RETURN(HloSchedule schedule, - ScheduleModule(*module, BufferSizeBytesFunction())); + ScheduleModule(module, BufferSizeBytesFunction())); // Run buffer analysis on the HLO graph. This analysis figures out which // temporary buffers are required to run the computation. @@ -832,7 +846,7 @@ CpuCompiler::CompileAheadOfTime(std::unique_ptr module_group, .EmitComputation( embedded_computation, embedded_computation->name(), /*is_top_level_computation=*/false, - &schedule.sequence(embedded_computation).instructions()) + schedule.sequence(embedded_computation).instructions()) .status()); } const string& entry_point_name = options.entry_point_name(); @@ -840,7 +854,7 @@ CpuCompiler::CompileAheadOfTime(std::unique_ptr module_group, ir_emitter.EmitComputation( computation, entry_point_name, /*is_top_level_computation=*/true, - &schedule.sequence(computation).instructions())); + schedule.sequence(computation).instructions())); CHECK(entry_function->getName() == llvm_ir::AsStringRef(entry_point_name)); diff --git a/tensorflow/compiler/xla/service/cpu/cpu_copy_insertion_test.cc b/tensorflow/compiler/xla/service/cpu/cpu_copy_insertion_test.cc index c9fb34be1cd582c71618c770c892058c233c571a..c085f85fb73e98e4c7ba15af8db8bb19c2499f5f 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_copy_insertion_test.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_copy_insertion_test.cc @@ -15,7 +15,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/cpu/cpu_copy_insertion.h" -#include "tensorflow/compiler/xla/legacy_flags/debug_options_flags.h" +#include "tensorflow/compiler/xla/debug_options_flags.h" #include "tensorflow/compiler/xla/literal.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" @@ -25,7 +25,7 @@ limitations under the License. #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_verified_test_base.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/platform/test_benchmark.h" @@ -52,7 +52,7 @@ int64 CountCopies(const HloModule& module) { return count; } -class CpuCopyInsertionTest : public HloVerifiedTestBase { +class CpuCopyInsertionTest : public HloTestBase { protected: void InsertCopies(HloModule* module) { CpuCopyInsertion copy_insertion; @@ -65,7 +65,7 @@ class CpuCopyInsertionTest : public HloVerifiedTestBase { TEST_F(CpuCopyInsertionTest, WhileBodyWithConstantRoot) { // Test a while body and condition which are each simply a constant (root of // computation is a constant). Each constant should be copied. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto builder = HloComputation::Builder(TestName()); auto param_0 = builder.AddInstruction( HloInstruction::CreateParameter(0, scalar_shape_, "param_0")); @@ -90,7 +90,7 @@ TEST_F(CpuCopyInsertionTest, WhileBodyWithConstantRoot) { module->AddEntryComputation(builder.Build()); - InsertCopies(module); + InsertCopies(module.get()); EXPECT_EQ(CountCopies(*module), 3); @@ -103,7 +103,7 @@ TEST_F(CpuCopyInsertionTest, TupleCall) { // Test a kCall instruction which calls a computation which produces a three // element tuple: one is a constant, one is a parameter, and one is produced // in the computation. The constant and parameter should be copied. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto builder = HloComputation::Builder(TestName()); auto param = builder.AddInstruction( HloInstruction::CreateParameter(0, scalar_shape_, "param_0")); @@ -127,7 +127,7 @@ TEST_F(CpuCopyInsertionTest, TupleCall) { module->AddEntryComputation(builder.Build()); - InsertCopies(module); + InsertCopies(module.get()); EXPECT_EQ(CountCopies(*subcomputation), 2); EXPECT_THAT(subcomputation->root_instruction(), 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 29abf38e439d919ff93629ed992cb3ff93a929bd..23d0af34233858515af21df5e92346742a5b5dc3 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_executable.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_executable.cc @@ -51,8 +51,7 @@ namespace cpu { CpuExecutable::CpuExecutable( std::unique_ptr jit, std::unique_ptr assignment, - std::unique_ptr hlo_module, - const string& entry_function_name, + std::unique_ptr hlo_module, const string& entry_function_name, std::unique_ptr hlo_profile_printer_data, std::unique_ptr hlo_profile_index_map) : Executable(std::move(hlo_module), std::move(hlo_profile_printer_data), @@ -214,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 @@ -233,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); @@ -327,7 +347,7 @@ StatusOr CpuExecutable::ExecuteAsyncOnStreamImpl( /*static*/ int64 CpuExecutable::ShapeSizeBytes(const Shape& shape) { // On the cpu, opaques are pointers. - if (ShapeUtil::IsOpaque(shape)) { + if (shape.IsOpaque()) { return sizeof(void*); } return ShapeUtil::ByteSizeOf(shape, sizeof(void*)); diff --git a/tensorflow/compiler/xla/service/cpu/cpu_executable.h b/tensorflow/compiler/xla/service/cpu/cpu_executable.h index 3c3c047bfe8ee0d1ad90ede2432a86264f47870b..3b91b15ba9b5603b50f78f489e9a3fdad354c083 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_executable.h +++ b/tensorflow/compiler/xla/service/cpu/cpu_executable.h @@ -49,7 +49,7 @@ class CpuExecutable : public Executable { public: CpuExecutable(std::unique_ptr jit, std::unique_ptr assignment, - std::unique_ptr hlo_module, + std::unique_ptr hlo_module, const string& entry_function_name, std::unique_ptr hlo_profile_printer_data, std::unique_ptr hlo_profile_index_map); diff --git a/tensorflow/compiler/xla/service/cpu/cpu_hlo_support_checker_test.cc b/tensorflow/compiler/xla/service/cpu/cpu_hlo_support_checker_test.cc index e6b6fcdf684eadb3702e490bbe24dbb7b3b52ec7..9cbfb88834bf51f4df54e97efe6cd7bf88b12334 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_hlo_support_checker_test.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_hlo_support_checker_test.cc @@ -16,7 +16,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/cpu/cpu_hlo_support_checker.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/test.h" -#include "tensorflow/compiler/xla/tests/hlo_verified_test_base.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" #include "tensorflow/core/lib/core/error_codes.pb.h" #include "tensorflow/core/lib/core/status_test_util.h" @@ -25,7 +25,7 @@ namespace { using ::testing::HasSubstr; -class CpuHloSupportCheckerTest : public HloVerifiedTestBase { +class CpuHloSupportCheckerTest : public HloTestBase { protected: CpuHloSupportChecker& checker() { return checker_; } @@ -42,10 +42,10 @@ TEST_F(CpuHloSupportCheckerTest, Add) { HloInstruction::CreateParameter(1, scalar_shape, "param1")); builder.AddInstruction(HloInstruction::CreateBinary( scalar_shape, HloOpcode::kAdd, param0, param1)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); - TF_ASSERT_OK(checker().Run(module).status()); + TF_ASSERT_OK(checker().Run(module.get()).status()); } TEST_F(CpuHloSupportCheckerTest, SparseUnimplemented) { @@ -60,7 +60,7 @@ TEST_F(CpuHloSupportCheckerTest, SparseUnimplemented) { // Since verifier is reporting sparse layouts as errors, we should // use a regular HloModule instead of VerifiedHloModule to avoid // verifier errors being triggered in the destructor. - auto module = HloTestBase::CreateNewModule(); + auto module = CreateNewUnverifiedModule(); module->AddEntryComputation(builder.Build()); Status status = checker().Run(module.get()).status(); diff --git a/tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion.cc b/tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion.cc index f9cd61bea3dc86cadff99d4a90eca44c16520823..6f79ad7c1468f27c74d84770ec6358fbcd1c1f09 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion.cc @@ -48,10 +48,15 @@ bool IsMatrixVectorDot(const HloInstruction* hlo) { (hlo_shape.dimensions(0) == 1 || hlo_shape.dimensions(1) == 1); } +bool HasExactlyOneUse(const HloInstruction& hlo_instr) { + return hlo_instr.user_count() == 1 && + absl::c_count(hlo_instr.users().front()->operands(), &hlo_instr) == 1; +} + bool CanBeOutputFused(const HloInstruction* producer, const HloInstruction* consumer) { return consumer->opcode() == HloOpcode::kAdd && IsMatrixVectorDot(producer) && - producer->user_count() == 1; + HasExactlyOneUse(*producer) == 1; } bool CanBeOutputFusedIntoSomeOperand(const HloInstruction* consumer) { 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 7d99b914d4f5e5d27722bcd098d2ae0c54a36a23..c4bde837e57e82584c2a007858ed8d55608acd3c 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion_test.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion_test.cc @@ -58,7 +58,7 @@ TEST_F(InstructionFusionTest, DotOperationFusion_Basic_0) { HloInstruction* dot = builder.AddInstruction( MakeDot(ShapeUtil::MakeShape(F32, {1024, 1}), exp0, arg1)); - auto module = CreateNewModule(); + auto module = CreateNewUnverifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(dot, computation->root_instruction()); EXPECT_TRUE(CpuInstructionFusion().Run(module.get()).ValueOrDie()); @@ -77,7 +77,7 @@ TEST_F(InstructionFusionTest, DotOperationFusion_Basic_1) { HloInstruction* dot = builder.AddInstruction( MakeDot(ShapeUtil::MakeShape(F32, {1, 1024}), arg0, exp1)); - auto module = CreateNewModule(); + auto module = CreateNewUnverifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(dot, computation->root_instruction()); EXPECT_TRUE(CpuInstructionFusion().Run(module.get()).ValueOrDie()); @@ -98,7 +98,7 @@ TEST_F(InstructionFusionTest, DotOperationNoFusion_Bitcast) { HloInstruction* dot = builder.AddInstruction( MakeDot(ShapeUtil::MakeShape(F32, {1024, 1}), bitcast0, arg1)); - auto module = CreateNewModule(); + auto module = CreateNewUnverifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(dot, computation->root_instruction()); EXPECT_FALSE(CpuInstructionFusion().Run(module.get()).ValueOrDie()); @@ -119,7 +119,7 @@ TEST_F(InstructionFusionTest, DotOperationFusion_Reshape) { HloInstruction* dot = builder.AddInstruction( MakeDot(ShapeUtil::MakeShape(F32, {1024, 1}), reshape0, arg1)); - auto module = CreateNewModule(); + auto module = CreateNewUnverifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(dot, computation->root_instruction()); EXPECT_TRUE(CpuInstructionFusion().Run(module.get()).ValueOrDie()); @@ -138,7 +138,7 @@ TEST_F(InstructionFusionTest, DotOperationFusion_TooLarge) { HloInstruction* dot = builder.AddInstruction( MakeDot(ShapeUtil::MakeShape(F32, {1, 32 * 1024}), arg0, exp1)); - auto module = CreateNewModule(); + auto module = CreateNewUnverifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(dot, computation->root_instruction()); EXPECT_FALSE(CpuInstructionFusion().Run(module.get()).ValueOrDie()); @@ -157,7 +157,7 @@ TEST_F(InstructionFusionTest, DotOperationFusion_ElementReuse) { HloInstruction* dot = builder.AddInstruction( MakeDot(ShapeUtil::MakeShape(F32, {2, 1024}), arg0, exp1)); - auto module = CreateNewModule(); + auto module = CreateNewUnverifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(dot, computation->root_instruction()); EXPECT_FALSE(CpuInstructionFusion().Run(module.get()).ValueOrDie()); @@ -321,7 +321,7 @@ TEST_F(OpcodeFusionTest, Exponential_Reshape_Negate) { builder.AddInstruction( HloInstruction::CreateUnary(result_shape, HloOpcode::kNegate, reshape2)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); RunFusionAndCheckOpcodesWereFused( @@ -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,23 +340,26 @@ 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)); - auto module = CreateNewModule(); + auto module = CreateNewUnverifiedModule(); module->AddEntryComputation(builder.Build()); 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) { @@ -370,7 +373,7 @@ TEST_F(OpcodeFusionTest, Broadcast_Negate) { builder.AddInstruction(HloInstruction::CreateUnary( result_shape, HloOpcode::kNegate, broadcast1)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); RunFusionAndCheckOpcodesWereFused( @@ -381,18 +384,18 @@ 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)); - auto module = CreateNewModule(); + auto module = CreateNewUnverifiedModule(); module->AddEntryComputation(builder.Build()); RunFusionAndCheckOpcodesWereFused( @@ -410,7 +413,7 @@ TEST_F(OpcodeFusionTest, Exponential_Negate) { builder.AddInstruction( HloInstruction::CreateUnary(param_shape, HloOpcode::kNegate, exp1)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); RunFusionAndCheckOpcodesWereFused( @@ -429,7 +432,7 @@ TEST_F(OpcodeFusionTest, Reshape_Negate) { builder.AddInstruction( HloInstruction::CreateUnary(result_shape, HloOpcode::kNegate, reshape1)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); RunFusionAndCheckOpcodesWereFused( @@ -447,7 +450,7 @@ TEST_F(OpcodeFusionTest, Reverse_Negate) { builder.AddInstruction( HloInstruction::CreateUnary(param_shape, HloOpcode::kNegate, reverse1)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); RunFusionAndCheckOpcodesWereFused( @@ -466,7 +469,7 @@ TEST_F(OpcodeFusionTest, Slice_Negate) { builder.AddInstruction(HloInstruction::CreateUnary( ShapeUtil::MakeShape(S32, {2}), HloOpcode::kNegate, slice1)); - auto module = CreateNewModule(); + auto module = CreateNewUnverifiedModule(); module->AddEntryComputation(builder.Build()); RunFusionAndCheckOpcodesWereFused( @@ -489,7 +492,7 @@ TEST_F(OpcodeFusionTest, Exponential_Transpose_Negate) { builder.AddInstruction(HloInstruction::CreateUnary( result_shape, HloOpcode::kNegate, transpose2)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); RunFusionAndCheckOpcodesWereFused( @@ -498,7 +501,7 @@ TEST_F(OpcodeFusionTest, Exponential_Transpose_Negate) { } TEST_F(OpcodeFusionTest, UnaryMapOfExp) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); Shape shape = ShapeUtil::MakeShape(F32, {3, 4}); @@ -517,7 +520,7 @@ TEST_F(OpcodeFusionTest, UnaryMapOfExp) { } TEST_F(OpcodeFusionTest, BinaryMapOfExps) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); Shape shape = ShapeUtil::MakeShape(F32, {3, 4}); @@ -542,85 +545,84 @@ TEST_F(OpcodeFusionTest, BinaryMapOfExps) { } TEST_F(OpcodeFusionTest, DynamicSliceWithDynamicUpdateSlice) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); 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) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); 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}); } @@ -641,7 +643,7 @@ TEST_F(OpcodeFusionTest, ReuseViaImplicitBroadcastUnary) { builder.AddInstruction( HloInstruction::CreateUnary(large_shape, HloOpcode::kExp, small_exp)); - std::unique_ptr module = CreateNewModule(); + std::unique_ptr module = CreateNewUnverifiedModule(); module->AddEntryComputation(builder.Build()); auto did_fusion = CpuInstructionFusion().Run(module.get()); @@ -670,7 +672,7 @@ TEST_F(OpcodeFusionTest, ReuseViaImplicitBroadcastBinary) { builder.AddInstruction(HloInstruction::CreateBinary( large_shape, HloOpcode::kAdd, small_exp, large_param)); - std::unique_ptr module = CreateNewModule(); + std::unique_ptr module = CreateNewUnverifiedModule(); module->AddEntryComputation(builder.Build()); auto did_fusion = CpuInstructionFusion().Run(module.get()); @@ -712,7 +714,7 @@ void CreateComputationForDotAddOutputFusionTest(const string& test_name, } TEST_F(OpcodeFusionTest, DotAddOutputFusion_1x50x19) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); CreateComputationForDotAddOutputFusionTest(TestName(), module.get(), /*m=*/1, /*k=*/50, /*n=*/19, /*add_extra_use_for_dot=*/false); @@ -725,7 +727,7 @@ TEST_F(OpcodeFusionTest, DotAddOutputFusion_1x50x19) { } TEST_F(OpcodeFusionTest, DotAddOutputFusion_19x50x1) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); CreateComputationForDotAddOutputFusionTest(TestName(), module.get(), /*m=*/19, /*k=*/50, /*n=*/1, /*add_extra_use_for_dot=*/false); @@ -738,7 +740,7 @@ TEST_F(OpcodeFusionTest, DotAddOutputFusion_19x50x1) { } TEST_F(OpcodeFusionTest, DotAddOutputFusion_19x50x19) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); CreateComputationForDotAddOutputFusionTest(TestName(), module.get(), /*m=*/19, /*k=*/50, /*n=*/19, /*add_extra_use_for_dot=*/false); @@ -751,7 +753,7 @@ TEST_F(OpcodeFusionTest, DotAddOutputFusion_19x50x19) { } TEST_F(OpcodeFusionTest, DotAddOutputFusion_19x50x1_multi_use) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); CreateComputationForDotAddOutputFusionTest(TestName(), module.get(), /*m=*/19, /*k=*/50, /*n=*/1, /*add_extra_use_for_dot=*/true); @@ -763,6 +765,28 @@ TEST_F(OpcodeFusionTest, DotAddOutputFusion_19x50x1_multi_use) { Not(op::Fusion())); } +TEST_F(InstructionFusionTest, + DotOperationFusion_DontOutputFuseDuplicateOperands) { + absl::string_view module_string = R"( +HloModule module + +ENTRY main { + a = f32[50,60]{1,0} parameter(0) + b = f32[60,1]{1,0} parameter(1) + c = f32[50,1]{1,0} dot(a, b), lhs_contracting_dims={1}, rhs_contracting_dims={0} + ROOT d = f32[50,1]{1,0} add(c, c) +} +)"; + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(module_string)); + TF_ASSERT_OK_AND_ASSIGN(bool fused_something, + CpuInstructionFusion().Run(module.get())); + EXPECT_FALSE(fused_something); + EXPECT_THAT(module->entry_computation()->root_instruction(), + Not(op::Fusion())); +} + struct GatherLoopFusionTestSpec { string test_name; string hlo_computation_text; @@ -908,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 c291bf2d1ba2eaff4192051840768c037bece86f..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) { @@ -160,7 +137,7 @@ Status CpuLayoutAssignment::AddBackendConstraints( continue; } // Skip operands with non-array shapes. - if (!ShapeUtil::IsArray(instruction->operand(operand_no)->shape())) { + if (!instruction->operand(operand_no)->shape().IsArray()) { continue; } Shape operand_shape( @@ -175,7 +152,7 @@ Status CpuLayoutAssignment::AddBackendConstraints( } // Skip instructions which don't produce array shapes (tuples, opaque, // etc.). - if (!ShapeUtil::IsArray(instruction->shape())) { + if (!instruction->shape().IsArray()) { continue; } } diff --git a/tensorflow/compiler/xla/service/cpu/cpu_layout_assignment_test.cc b/tensorflow/compiler/xla/service/cpu/cpu_layout_assignment_test.cc index 97659b88a7974d7caf91ab0d4741f3635e4dae4a..6c61b64758ede160e2d50e4429590a789ec253c3 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_layout_assignment_test.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_layout_assignment_test.cc @@ -73,7 +73,7 @@ TEST_F(CpuLayoutAssignmentTest, DotWithConstantRhsTensor) { auto result = builder.AddInstruction( CreateCanonicalDot(result_shape, dot_lhs, dot_rhs)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* computation = module->AddEntryComputation(builder.Build()); ComputationLayout computation_layout(computation->ComputeProgramShape()); @@ -114,7 +114,7 @@ TEST_F(CpuLayoutAssignmentTest, MultipleDotsWithSameConstantRhsTensor0) { builder.AddInstruction(HloInstruction::CreateBinary( result_shape, HloOpcode::kAdd, dot_a_result, dot_b_result)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* computation = module->AddEntryComputation(builder.Build()); ComputationLayout computation_layout(computation->ComputeProgramShape()); @@ -158,7 +158,7 @@ TEST_F(CpuLayoutAssignmentTest, MultipleDotsWithSameConstantRhsTensor1) { auto tuple_result = builder.AddInstruction( HloInstruction::CreateTuple({dot_a_result, dot_b_result})); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* computation = module->AddEntryComputation(builder.Build()); ComputationLayout computation_layout(computation->ComputeProgramShape()); @@ -192,7 +192,7 @@ TEST_F(CpuLayoutAssignmentTest, DotWithConstantLhsTensor) { auto dot_result = builder.AddInstruction( CreateCanonicalDot(result_shape, dot_lhs, dot_rhs)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* computation = module->AddEntryComputation(builder.Build()); ComputationLayout computation_layout(computation->ComputeProgramShape()); @@ -232,7 +232,7 @@ TEST_F(CpuLayoutAssignmentTest, DotWithConstantRhsTensorThroughGTE) { auto dot_result = builder.AddInstruction( CreateCanonicalDot(result_shape, dot_lhs, dot_rhs)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* computation = module->AddEntryComputation(builder.Build()); ComputationLayout computation_layout(computation->ComputeProgramShape()); @@ -353,7 +353,7 @@ static void AssertCorrectLayoutForDotOutputFusion( } TEST_F(CpuLayoutAssignmentTest, DotOutputFusion_1x50x19_dot_idx_0) { - std::unique_ptr module = CreateNewModule(); + std::unique_ptr module = CreateNewVerifiedModule(); TF_ASSERT_OK_AND_ASSIGN( DotOutputFusionLayoutAssignmentResult layout_assignment_result, RunDotOutputFusion(module.get(), TestName(), /*m=*/1, /*k=*/50, /*n=*/19, @@ -365,7 +365,7 @@ TEST_F(CpuLayoutAssignmentTest, DotOutputFusion_1x50x19_dot_idx_0) { } TEST_F(CpuLayoutAssignmentTest, DotOutputFusion_1x50x19_dot_idx_1) { - std::unique_ptr module = CreateNewModule(); + std::unique_ptr module = CreateNewVerifiedModule(); TF_ASSERT_OK_AND_ASSIGN( DotOutputFusionLayoutAssignmentResult layout_assignment_result, RunDotOutputFusion(module.get(), TestName(), /*m=*/1, /*k=*/50, /*n=*/19, @@ -377,7 +377,7 @@ TEST_F(CpuLayoutAssignmentTest, DotOutputFusion_1x50x19_dot_idx_1) { } TEST_F(CpuLayoutAssignmentTest, DotOutputFusion_19x50x1_dot_idx_0) { - std::unique_ptr module = CreateNewModule(); + std::unique_ptr module = CreateNewVerifiedModule(); TF_ASSERT_OK_AND_ASSIGN( DotOutputFusionLayoutAssignmentResult layout_assignment_result, RunDotOutputFusion(module.get(), TestName(), /*m=*/19, /*k=*/50, /*n=*/1, @@ -389,7 +389,7 @@ TEST_F(CpuLayoutAssignmentTest, DotOutputFusion_19x50x1_dot_idx_0) { } TEST_F(CpuLayoutAssignmentTest, DotOutputFusion_19x50x1_dot_idx_1) { - std::unique_ptr module = CreateNewModule(); + std::unique_ptr module = CreateNewVerifiedModule(); TF_ASSERT_OK_AND_ASSIGN( DotOutputFusionLayoutAssignmentResult layout_assignment_result, RunDotOutputFusion(module.get(), TestName(), /*m=*/19, /*k=*/50, /*n=*/1, @@ -401,7 +401,7 @@ TEST_F(CpuLayoutAssignmentTest, DotOutputFusion_19x50x1_dot_idx_1) { } TEST_F(CpuLayoutAssignmentTest, DotOutputFusion_19x50x19_dot_idx_0) { - std::unique_ptr module = CreateNewModule(); + std::unique_ptr module = CreateNewVerifiedModule(); TF_ASSERT_OK_AND_ASSIGN( DotOutputFusionLayoutAssignmentResult layout_assignment_result, RunDotOutputFusion(module.get(), TestName(), /*m=*/19, /*k=*/50, /*n=*/19, @@ -413,7 +413,7 @@ TEST_F(CpuLayoutAssignmentTest, DotOutputFusion_19x50x19_dot_idx_0) { } TEST_F(CpuLayoutAssignmentTest, DotOutputFusion_19x50x19_dot_idx_1) { - std::unique_ptr module = CreateNewModule(); + std::unique_ptr module = CreateNewVerifiedModule(); TF_ASSERT_OK_AND_ASSIGN( DotOutputFusionLayoutAssignmentResult layout_assignment_result, RunDotOutputFusion(module.get(), TestName(), /*m=*/19, /*k=*/50, /*n=*/19, diff --git a/tensorflow/compiler/xla/service/cpu/cpu_options.cc b/tensorflow/compiler/xla/service/cpu/cpu_options.cc index b8ace5702688096822573c7afae234cbcbe77b28..ff654c83d61e7cc09ac7839feccaf2bc9cb3c63c 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_options.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_options.cc @@ -22,10 +22,9 @@ limitations under the License. namespace { const char* const kXlaOptimizeForSizeCpuOption = "xla_cpu_optimize_for_size"; -const char* const kXlaDisableVectorizedReduce = "xla_disable_vectorized_reduce"; 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 @@ -58,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 1cc2844470376ceb61601f6d1361def84eac5b45..3361a5973f5e8c91802b26d68477347b196d3cac 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_transfer_manager.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_transfer_manager.cc @@ -19,6 +19,7 @@ limitations under the License. #include #include +#include "absl/base/casts.h" #include "absl/memory/memory.h" #include "tensorflow/compiler/xla/literal.h" #include "tensorflow/compiler/xla/literal_util.h" @@ -29,7 +30,6 @@ limitations under the License. #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/util.h" #include "tensorflow/compiler/xla/xla_data.pb.h" -#include "tensorflow/core/lib/core/casts.h" #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/gtl/cleanup.h" #include "tensorflow/core/platform/logging.h" @@ -97,7 +97,7 @@ Status CpuTransferManager::TransferLiteralToInfeed( VLOG(2) << "Transferring literal to infeed with shape: " << ShapeUtil::HumanString(shape); - if (!ShapeUtil::IsTuple(shape)) { + if (!shape.IsTuple()) { int64 size = GetByteSizeRequirement(shape); return TransferBufferToInfeed(executor, size, literal.untyped_data()); } @@ -178,12 +178,12 @@ CpuTransferManager::TransferBufferToInfeedInternal(se::StreamExecutor* executor, Status CpuTransferManager::TransferLiteralFromOutfeed( se::StreamExecutor* executor, const Shape& literal_shape, MutableBorrowingLiteral literal) { - if (!ShapeUtil::IsTuple(literal_shape)) { + if (!literal_shape.IsTuple()) { int64 size = GetByteSizeRequirement(literal_shape); // Note: OSS build didn't like implicit conversion from // literal_shape.dimensions() to the array slice on 2017-07-10. absl::Span dimensions( - tensorflow::bit_cast(literal_shape.dimensions().data()), + absl::bit_cast(literal_shape.dimensions().data()), literal_shape.dimensions().size()); TF_ASSIGN_OR_RETURN( Shape received_shape, 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 99fa707c959854e50c6d954fe92b87e93e267dc6..48510181bd01c87c9db764396b556fdf34e6c8c4 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,932 +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_.ForReturnVoid("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_.ForReturnVoid( - "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_.ForReturnVoid( - "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_.ForReturnVoid( - "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_.IfReturnVoid( - 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); + // 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 EmitInnerLoopTiled(MemoryTile* lhs_memory_tile, int64 rows, - std::vector* vector_accumulators); + // Emits a call to the CPU runtime to perform the matrix multiply. + Status EmitCallToRuntime(); - void EmitInnerLoopEpilogue(llvm::Value* current_tile_row, int64 rows, - std::vector* scalar_accumulators); + // Represents the dimensions of a matrix-matrix multiply operation. + struct MatMultDims { + // The number of rows in the LHS. + int64 m; - Config config_; - llvm::Value* lhs_; - llvm::Value* rhs_; - llvm::Value* addend_; - llvm::Value* result_; - llvm::IRBuilder<>* b_; - KernelSupportLibrary ksl_; - VectorSupportLibrary vsl_; -}; + // 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::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)); - } + // The number of columns on the RHS. + int64 n; - 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); - } -} + // True if the LHS matrix is column major. + bool lhs_column_major; -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; + // True if the LHS contraction dimension is not 1. + bool lhs_non_canonical; - ksl_.ForReturnVoid( - "dot.outer.tiled", - /*start=*/0, /*end=*/row_limit, /*step=*/tile_rows(), - [&](llvm::Value* row) { EmitOuterLoopBody(row, tile_rows()); }); + // True if the RHS matrix is column major. + bool rhs_column_major; - if (row_remainder != 0) { - EmitOuterLoopBody(b_->getInt64(row_limit), row_remainder); - } -} + // True if the RHS contraction dimension is not 0. + bool rhs_non_canonical; -void RowMajorMatrixVectorProductEmitter::EmitInnerLoopTiled( - MemoryTile* lhs_memory_tile, int64 rows, - std::vector* vector_accumulators) { - int64 column_limit = k() - (k() % tile_cols()); - - ksl_.ForReturnVoid("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_.ForReturnVoid( - "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_; + // 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_.ForReturnVoid("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_.ForReturnVoid( - "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_.ForReturnVoid( - "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_.ForReturnVoid( - "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, @@ -975,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), @@ -985,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(); @@ -1051,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( @@ -1063,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; - } +void DotOpEmitter::EmitTiledLlvmIrGemv() { + PrimitiveType primitive_type = dot_info_.result_shape.element_type(); - PrimitiveType primitive_type = dot_.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; @@ -1144,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); @@ -1178,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() { @@ -1241,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(); @@ -1256,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); @@ -1286,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( @@ -1391,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() { @@ -1406,16 +562,20 @@ Status DotOpEmitter::EmitScalarDot() { llvm::Value* rhs_value = rhs_array_.EmitReadArrayElement(/*index=*/element_index, b_); if (ShapeUtil::ElementIsComplex(lhs_array_.GetShape())) { -#define REAL(x) b_->CreateExtractValue(x, {0}) -#define IMAG(x) b_->CreateExtractValue(x, {1}) - llvm::Value* real = - b_->CreateFSub(b_->CreateFMul(REAL(lhs_value), REAL(rhs_value)), - b_->CreateFMul(IMAG(lhs_value), IMAG(rhs_value))); - llvm::Value* imag = - b_->CreateFAdd(b_->CreateFMul(REAL(lhs_value), IMAG(rhs_value)), - b_->CreateFMul(IMAG(lhs_value), REAL(rhs_value))); -#undef IMAG -#undef REAL + auto get_real = [&](llvm::Value* x) { + return b_->CreateExtractValue(x, {0}); + }; + + auto get_imag = [&](llvm::Value* x) { + return b_->CreateExtractValue(x, {1}); + }; + + llvm::Value* real = b_->CreateFSub( + b_->CreateFMul(get_real(lhs_value), get_real(rhs_value)), + b_->CreateFMul(get_imag(lhs_value), get_imag(rhs_value))); + llvm::Value* imag = b_->CreateFAdd( + b_->CreateFMul(get_real(lhs_value), get_imag(rhs_value)), + b_->CreateFMul(get_imag(lhs_value), get_real(rhs_value))); result = llvm::ConstantAggregateZero::get(lhs_array_.GetElementLlvmType()); result = b_->CreateInsertValue(result, real, {0}); result = b_->CreateInsertValue(result, imag, {1}); @@ -1435,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; @@ -1528,11 +688,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)), @@ -1546,77 +706,6 @@ DotOpEmitter::MatMultDims DotOpEmitter::GetMatMultDims() const { LayoutUtil::Minor(target_array_.GetShape().layout(), 0) == 0}; } -// Return whether the given shape is a matrix with no padding. -static bool IsRank2WithNoPadding(const Shape& shape) { - return ShapeUtil::Rank(shape) == 2 && !LayoutUtil::IsPadded(shape); -} - -// 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 (!(IsRank2WithNoPadding(lhs_shape) && IsRank2WithNoPadding(rhs_shape) && - IsRank2WithNoPadding(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 +744,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/ir_emission_utils.cc b/tensorflow/compiler/xla/service/cpu/ir_emission_utils.cc index 1a8bedfe6afb4f096ddd4703c312b84d521a7ba5..a8b139aec9e96b6bb580baf74789df7c998cebf8 100644 --- a/tensorflow/compiler/xla/service/cpu/ir_emission_utils.cc +++ b/tensorflow/compiler/xla/service/cpu/ir_emission_utils.cc @@ -26,7 +26,7 @@ namespace cpu { int64 GetMinimumAlignmentForArray( const Shape& shape, const TargetMachineFeatures& target_machine_features) { - CHECK(ShapeUtil::IsArray(shape)); + CHECK(shape.IsArray()); CHECK(!LayoutUtil::HasLayout(shape) || LayoutUtil::IsDense(shape.layout())); // We don't require a layout to be set on `shape`. This only works on CPU diff --git a/tensorflow/compiler/xla/service/cpu/ir_emitter.cc b/tensorflow/compiler/xla/service/cpu/ir_emitter.cc index b2abdb39a598871a7cc44760e464f48b9a200874..0226e8275cb0e1de39c4c2e9a06d4cfa1a4854d3 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" @@ -54,6 +52,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/cpu/simple_orc_jit.h" #include "tensorflow/compiler/xla/service/elemental_ir_emitter.h" #include "tensorflow/compiler/xla/service/hlo_casting_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" #include "tensorflow/compiler/xla/service/llvm_ir/buffer_assignment_util.h" @@ -69,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 { @@ -110,10 +111,9 @@ IrEmitter::IrEmitter( StatusOr IrEmitter::EmitComputation( HloComputation* computation, const string& function_name_prefix, bool is_top_level_computation, - const std::vector* instruction_order) { + absl::Span instruction_order) { string function_name = name_uniquer_.GetUniqueName(function_name_prefix); - VLOG(2) << "Emitting IR for CPU function [" << function_name_prefix - << "]; ordered? " << (instruction_order != nullptr); + VLOG(2) << "Emitting IR for CPU function [" << function_name_prefix << "]"; is_top_level_computation_ = is_top_level_computation; num_dynamic_loop_bounds_ = 0; if (!computation->root_instruction()->outer_dimension_partitions().empty()) { @@ -139,12 +139,8 @@ StatusOr IrEmitter::EmitComputation( // readcyclecounter if it is unavailable. bool use_rdtscp = arch_type_ == llvm::Triple::ArchType::x86 || arch_type_ == llvm::Triple::ArchType::x86_64; - profiling_state_ = ProfilingState(use_rdtscp, GetProfileCountersArgument()); - if (instruction_order == nullptr) { - TF_RETURN_IF_ERROR(computation->Accept(this)); - } else { - TF_RETURN_IF_ERROR(computation->AcceptOrdered(this, *instruction_order)); - } + profiling_state_ = ProfilingState(use_rdtscp); + TF_RETURN_IF_ERROR(computation->AcceptOrdered(this, instruction_order)); llvm::Function* ir_function = compute_function_->function(); InsertOrDie(&emitted_functions_, computation, ir_function); // Delete 'compute_function', finalizing 'ir_function' and restoring caller @@ -227,11 +223,11 @@ Status IrEmitter::HandleConstant(HloInstruction* constant) { } Status IrEmitter::HandleCopy(HloInstruction* copy) { - if (ShapeUtil::IsTuple(copy->shape())) { + if (copy->shape().IsTuple()) { // kCopy shallow copies a tuple so just memcpy the top-level buffer. TF_RETURN_IF_ERROR(EmitTargetAddressForOp(copy)); return EmitMemcpy(*(copy->operand(0)), *copy); - } else if (ShapeUtil::IsArray(copy->shape())) { + } else if (copy->shape().IsArray()) { // Use the elemental emitter for array shapes. return DefaultAction(copy); } @@ -243,10 +239,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 { @@ -320,7 +318,7 @@ Status IrEmitter::HandleTupleSelect(HloInstruction* tuple_select) { auto on_false = tuple_select->operand(2); TF_RET_CHECK(pred->shape().element_type() == PRED); TF_RET_CHECK(ShapeUtil::IsScalar(pred->shape())); - TF_RET_CHECK(ShapeUtil::IsTuple(tuple_select->shape())); + TF_RET_CHECK(tuple_select->shape().IsTuple()); TF_RETURN_IF_ERROR(EmitTargetAddressForOp(tuple_select)); llvm_ir::EmitTupleSelect(GetIrArrayFor(tuple_select), GetIrArrayFor(pred), GetEmittedValueFor(on_true), @@ -350,7 +348,7 @@ Status IrEmitter::HandleInfeed(HloInstruction* instruction) { llvm_ir::EmitTuple(GetIrArrayFor(infeed), {data_address, token_address}, &b_, module_); - if (ShapeUtil::IsTuple(data_shape)) { + if (data_shape.IsTuple()) { TF_RET_CHECK(!ShapeUtil::IsNestedTuple(data_shape)); // For a tuple, we first copy each of the internal elements to @@ -474,7 +472,7 @@ Status IrEmitter::HandleOutfeed(HloInstruction* outfeed) { const Shape& operand_shape = operand->shape(); llvm::Value* value = GetEmittedValueFor(operand); - if (!ShapeUtil::IsTuple(operand_shape)) { + if (!operand_shape.IsTuple()) { return EmitXfeedTransfer(XfeedKind::kOutfeed, operand_shape, value); } @@ -493,53 +491,64 @@ Status IrEmitter::HandleOutfeed(HloInstruction* outfeed) { return Status::OK(); } -Status IrEmitter::HandleSort(HloInstruction* sort) { +Status IrEmitter::HandleSort(HloInstruction* hlo) { + const HloSortInstruction* sort = Cast(hlo); TF_RETURN_IF_ERROR(EmitTargetAddressForOp(sort)); - auto keys = sort->operand(0); - auto values = sort->operand_count() > 1 ? sort->operand(1) : nullptr; - ShapeIndex keys_shape_index({}); - ShapeIndex values_shape_index({}); - if (values != nullptr) { - keys_shape_index = ShapeIndex({0}); - values_shape_index = ShapeIndex({1}); - } - auto keys_destination = GetAllocationSlice(*sort, keys_shape_index); - auto keys_destination_address = - EmitBufferPointer(keys_destination, keys->shape()); - auto values_destination = GetAllocationSlice(*sort, values_shape_index); - llvm::Value* values_destination_address = nullptr; - - // The sort is implemented in-place, therefore we first copy the operand - // buffer to the output buffer if they are not the same. - if (keys_destination != GetAllocationSlice(*keys)) { - int64 primitive_type_size = - ShapeUtil::ByteSizeOfPrimitiveType(keys->shape().element_type()); - auto source_buffer = GetEmittedValueFor(keys); - int64 keys_size = ByteSizeOf(keys->shape()); - MemCpy(keys_destination_address, /*DstAlign=*/primitive_type_size, - source_buffer, - /*SrcAlign=*/primitive_type_size, keys_size); - } - if (values != nullptr) { - values_destination_address = - EmitBufferPointer(values_destination, values->shape()); - if (values_destination != GetAllocationSlice(*values)) { + 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 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 = + sort->values_count() > 0 ? ShapeIndex({i}) : ShapeIndex({}); + const HloInstruction* operand = sort->operand(i); + // We assume that the layout of all involved operands and outputs is the + // same. + TF_RET_CHECK( + LayoutUtil::LayoutsInShapesEqual(keys_shape, operand->shape())); + TF_RET_CHECK(LayoutUtil::LayoutsInShapesEqual( + keys_shape, ShapeUtil::GetSubshape(sort->shape(), shape_index))); + + // The sort is implemented in-place, therefore we first copy the operand + // buffer to the output buffer if they are not the same. + auto destination_buffer = GetAllocationSlice(*sort, shape_index); + destination_addresses[i] = + EmitBufferPointer(destination_buffer, operand->shape()); + auto source_address = GetAllocationSlice(*operand); + if (destination_buffer != source_address) { int64 primitive_type_size = - ShapeUtil::ByteSizeOfPrimitiveType(values->shape().element_type()); - auto source_buffer = GetEmittedValueFor(values); - int64 values_size = ByteSizeOf(values->shape()); - MemCpy(values_destination_address, /*DstAlign=*/primitive_type_size, + ShapeUtil::ByteSizeOfPrimitiveType(operand->shape().element_type()); + auto source_buffer = GetEmittedValueFor(operand); + int64 size = ByteSizeOf(operand->shape()); + MemCpy(destination_addresses[i], /*DstAlign=*/primitive_type_size, source_buffer, - /*SrcAlign=*/primitive_type_size, values_size); + /*SrcAlign=*/primitive_type_size, size); } } // Normalize the shape and the dimension to sort. Shape normalized_keys_shape = - ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout( - keys->shape()); + ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout(keys_shape); int64 physical_dimension_to_sort = LayoutUtil::MakeLogicalToPhysical( - keys->shape().layout())[sort->dimensions(0)]; + keys_shape.layout())[sort->sort_dimension()]; int64 sort_dimension_elements = normalized_keys_shape.dimensions(physical_dimension_to_sort); @@ -548,94 +557,110 @@ Status IrEmitter::HandleSort(HloInstruction* sort) { higher_dimensions *= normalized_keys_shape.dimensions(i); } int64 lower_dimensions = 1; - for (int64 i = ShapeUtil::Rank(normalized_keys_shape) - 1; + for (int64 i = normalized_keys_shape.rank() - 1; i > physical_dimension_to_sort; --i) { 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)); + llvm::FunctionType* less_than_type = llvm::FunctionType::get( + b_.getInt1Ty(), {b_.getInt8PtrTy(), b_.getInt8PtrTy()}, + /*isVarArg=*/false); + auto less_than_function = llvm_ir::CreateFunction( + less_than_type, llvm::GlobalValue::InternalLinkage, + /*enable_fast_math=*/false, + /*optimize_for_size=*/true, absl::StrCat(IrName(sort), "_comparator"), + module_); + // Emit the code for the less_than function. + { + llvm::IRBuilder<>::InsertPointGuard guard(b_); + + auto* entry_bb = + llvm::BasicBlock::Create(b_.getContext(), "entry", less_than_function); + + b_.SetInsertPoint(entry_bb); + auto keys_ir_type = llvm_ir::PrimitiveTypeToIrType(keys_type, module_); + CHECK_EQ(less_than_function->arg_size(), 2); + llvm::Value* keys_lhs_ptr = less_than_function->arg_begin(); + keys_lhs_ptr = PointerCast(keys_lhs_ptr, keys_ir_type->getPointerTo()); + llvm::Value* keys_rhs_ptr = less_than_function->arg_begin() + 1; + keys_rhs_ptr = PointerCast(keys_rhs_ptr, keys_ir_type->getPointerTo()); + + // TODO(b/122298745): Replace the custom compare logic with a call to the + // computation specified for the Sort op. + llvm::Value* keys_lhs = Load(keys_ir_type, keys_lhs_ptr); + llvm::Value* keys_rhs = Load(keys_ir_type, keys_rhs_ptr); + bool is_signed_comparison = true; + if (primitive_util::IsFloatingPointType(keys_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( + keys_lhs->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); + }; + keys_lhs = b_.CreateBitCast(keys_lhs, comparison_type); + keys_rhs = b_.CreateBitCast(keys_rhs, comparison_type); + keys_lhs = maybe_flip(keys_lhs); + keys_rhs = maybe_flip(keys_rhs); + } else if (!primitive_util::IsSignedIntegralType(keys_type)) { + is_signed_comparison = false; + } + llvm::Value* result = + b_.CreateICmp(is_signed_comparison ? llvm::ICmpInst::ICMP_SLT + : llvm::ICmpInst::ICMP_ULT, + keys_lhs, keys_rhs); + llvm::ReturnInst::Create(b_.getContext(), + /*retVal=*/result, entry_bb); } llvm::FunctionType* key_value_sort_type = llvm::FunctionType::get( b_.getVoidTy(), - {keys_native_type, b_.getInt64Ty(), b_.getInt64Ty(), b_.getInt64Ty(), - b_.getInt8PtrTy(), b_.getInt32Ty()}, + {b_.getInt64Ty(), b_.getInt64Ty(), b_.getInt64Ty(), + b_.getInt8PtrTy()->getPointerTo(), b_.getInt32Ty(), + b_.getInt32Ty()->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::cast(module_->getOrInsertFunction( + runtime::kKeyValueSortSymbolName, key_value_sort_type)); key_value_sort_func->setCallingConv(llvm::CallingConv::C); key_value_sort_func->setDoesNotThrow(); - key_value_sort_func->setOnlyAccessesArgMemory(); + 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(keys_destination_address, keys_native_type), - b_.getInt64(higher_dimensions), b_.getInt64(sort_dimension_elements), - b_.getInt64(lower_dimensions), - values != nullptr - ? PointerCast(values_destination_address, b_.getInt8PtrTy()) - : llvm::Constant::getNullValue(b_.getInt8PtrTy()), - b_.getInt32(values != nullptr ? ShapeUtil::ByteSizeOfPrimitiveType( - values->shape().element_type()) - : 0)}); - - if (values != nullptr) { - llvm_ir::EmitTuple(GetIrArrayFor(sort), - {keys_destination_address, values_destination_address}, - &b_, module_); + {b_.getInt64(higher_dimensions), b_.getInt64(sort_dimension_elements), + b_.getInt64(lower_dimensions), values, + b_.getInt32(sort->operand_count()), sizes, less_than_function}); + + if (sort->values_count() > 0) { + llvm_ir::EmitTuple(GetIrArrayFor(sort), destination_addresses, &b_, + module_); } return Status::OK(); } @@ -772,8 +797,8 @@ Status IrEmitter::HandleSelectAndScatter(HloInstruction* select_and_scatter) { const auto init_value = select_and_scatter->operand(2); const Window& window = select_and_scatter->window(); PrimitiveType operand_element_type = operand->shape().element_type(); - const int64 rank = ShapeUtil::Rank(operand->shape()); - CHECK_EQ(rank, ShapeUtil::Rank(source->shape())); + const int64 rank = operand->shape().rank(); + CHECK_EQ(rank, source->shape().rank()); CHECK_EQ(rank, window.dimensions_size()); // TODO(b/31410564): Implement dilation for select-and-scatter. @@ -935,12 +960,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. @@ -963,10 +984,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( @@ -1111,7 +1132,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. @@ -1326,11 +1347,11 @@ Status IrEmitter::HandleFft(HloInstruction* fft) { return Status::OK(); } -Status IrEmitter::HandleCrossReplicaSum(HloInstruction* crs) { +Status IrEmitter::HandleAllReduce(HloInstruction* crs) { if (hlo_module_config_.replica_count() != 1) { // TODO(b/33011107): Support nontrivial cross replica sum on CPU. return Unimplemented( - "CrossReplicaSum with >1 replica is not implemented on CPU."); + "AllReduce with >1 replica is not implemented on CPU."); } // When there is a single replica, a cross replica sum is the identity @@ -1355,8 +1376,8 @@ Status IrEmitter::HandleCrossReplicaSum(HloInstruction* crs) { assignment_.GetUniqueSlice(crs, {i})); const Shape& operand_shape = crs->operand(i)->shape(); - CHECK(ShapeUtil::IsArray(operand_shape)) - << "Operands to cross-replica-sum must be arrays: " << crs->ToString(); + CHECK(operand_shape.IsArray()) + << "Operands to all-reduce must be arrays: " << crs->ToString(); operand_ptrs.push_back(EmitBufferPointer(out_slice, operand_shape)); // TODO(b/63762267): Be more aggressive about specifying alignment. @@ -1367,33 +1388,6 @@ Status IrEmitter::HandleCrossReplicaSum(HloInstruction* crs) { return Status::OK(); } -// Fills up the free variables in 'index_with_free_var' with values from -// 'filler_index'. The size of free variables must be the same as the -// size of 'filler_index'. -// -// This is often used after dimension reduction, where -// 'index_with_free_var' has one or more dimensions reduced, which serves as -// free variables (represented as nullptr). For example, if we have a 4 -// dimensional input and index for the dimension being reduced is -// 2 (third dimension), we will have an index like [i, j, NULL, k] -// after reduced dimension. -// -// Here we fill up that free variable by 'filler_index', which contains -// the value in the reduced dimension. -static llvm_ir::IrArray::Index FillReducedDimensionIndex( - llvm_ir::IrArray::Index index_with_free_var, - llvm_ir::IrArray::Index filler_index) { - llvm_ir::IrArray::Index::const_iterator it = filler_index.begin(); - - for (size_t i = 0; i < index_with_free_var.size(); ++i) { - if (index_with_free_var[i] == nullptr) { - index_with_free_var[i] = *it++; - } - } - CHECK(filler_index.end() == it); - return index_with_free_var; -} - Status IrEmitter::HandleParameter(HloInstruction* parameter) { VLOG(2) << "HandleParameter: " << parameter->ToString(); return EmitTargetAddressForOp(parameter); @@ -1419,7 +1413,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); @@ -1432,7 +1426,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; @@ -1524,7 +1518,8 @@ IrEmitter::ReductionGenerator IrEmitter::MatchReductionGenerator( case HloOpcode::kMaximum: return [root_is_floating_point, root_is_signed]( - llvm::IRBuilder<>* b, llvm::Value* lhs, llvm::Value* rhs) { + llvm::IRBuilder<>* b, llvm::Value* lhs, + llvm::Value* rhs) -> llvm::Value* { if (root_is_floating_point) { return llvm_ir::EmitCallToIntrinsic(llvm::Intrinsic::maxnum, {lhs, rhs}, {lhs->getType()}, b); @@ -1539,7 +1534,8 @@ IrEmitter::ReductionGenerator IrEmitter::MatchReductionGenerator( case HloOpcode::kMinimum: return [root_is_floating_point, root_is_signed]( - llvm::IRBuilder<>* b, llvm::Value* lhs, llvm::Value* rhs) { + llvm::IRBuilder<>* b, llvm::Value* lhs, + llvm::Value* rhs) -> llvm::Value* { if (root_is_floating_point) { return llvm_ir::EmitCallToIntrinsic(llvm::Intrinsic::minnum, {lhs, rhs}, {lhs->getType()}, b); @@ -1727,10 +1723,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()), @@ -1742,7 +1736,7 @@ StatusOr IrEmitter::EmitVectorizedReduce( return false; } - CHECK(!ShapeUtil::IsTuple(reduce->shape())); + CHECK(!reduce->shape().IsTuple()); TF_RETURN_IF_ERROR(EmitTargetAddressForOp(reduce)); // We know we're not reducing over the most minor dimension, which means we @@ -1909,7 +1903,7 @@ StatusOr IrEmitter::EmitTargetElementLoopBodyForReduce( Status IrEmitter::HandleReduce(HloInstruction* reduce) { // TODO(b/112040122): Support variadic reduce. - if (!ShapeUtil::IsArray(reduce->shape())) { + if (!reduce->shape().IsArray()) { return Unimplemented("Variadic reduce is not supported on CPU"); } auto arg = reduce->mutable_operand(0); @@ -2008,7 +2002,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 = @@ -2180,30 +2174,22 @@ Status IrEmitter::HandlePad(HloInstruction* pad) { return Status::OK(); } -// If `hlo` is a Transpose, returns its operand; otherwise returns `hlo` itself. -static const HloInstruction* StripTranspose(const HloInstruction& hlo) { - if (hlo.IsRank2Transpose()) { - return hlo.operand(0); - } - return &hlo; -} - Status IrEmitter::HandleFusion(HloInstruction* fusion) { auto* root = fusion->fused_expression_root(); if (llvm_ir::CanEmitFusedDynamicUpdateSliceInPlace(fusion, assignment_)) { VLOG(3) << "HandleFusion FusedDynamicUpdateSliceInPlace"; CpuElementalIrEmitter elemental_emitter(hlo_module_config_, this, module_); TF_RETURN_IF_ERROR(EmitTargetAddressForOp(fusion)); - // Delegate to common implementation of fused in-place dynamic-update-slice. - auto operands = GetIrArraysForOperandsOf(fusion); return llvm_ir::EmitFusedDynamicUpdateSliceInPlace( - fusion, operands, GetIrArrayFor(fusion), &elemental_emitter, &b_); + fusion, GetGeneratorForOperandIrArrays(fusion), GetIrArrayFor(fusion), + &elemental_emitter, &b_); } else if (fusion->fusion_kind() == HloInstruction::FusionKind::kLoop) { VLOG(3) << "HandleFusion kLoop"; CpuElementalIrEmitter elemental_emitter(hlo_module_config_, this, module_); auto operands = GetIrArraysForOperandsOf(fusion); - FusedIrEmitter fused_emitter(operands, &elemental_emitter); + FusedIrEmitter fused_emitter(GetGeneratorForOperandIrArrays(fusion), + &elemental_emitter); TF_RETURN_IF_ERROR(fusion->fused_expression_root()->Accept(&fused_emitter)); return EmitTargetElementLoop(fusion, fused_emitter.GetRootGenerator()); @@ -2231,10 +2217,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"); @@ -2292,6 +2278,21 @@ Status IrEmitter::HandleCustomCall(HloInstruction* custom_call) { /*isVarArg=*/false))); TF_RETURN_IF_ERROR(EmitTargetAddressForOp(custom_call)); + // Write the tuple table if the output is a tuple. + if (custom_call->shape().IsTuple()) { + std::vector base_ptrs; + for (int i = 0; i < ShapeUtil::TupleElementCount(custom_call->shape()); + ++i) { + const Shape& elem_shape = + ShapeUtil::GetTupleElementShape(custom_call->shape(), i); + TF_RET_CHECK(!elem_shape.IsTuple()) << "Nested tuples not implemented"; + TF_ASSIGN_OR_RETURN(const BufferAllocation::Slice slice, + assignment_.GetUniqueSlice(custom_call, {i})); + llvm::Value* addr = EmitBufferPointer(slice, elem_shape); + base_ptrs.push_back(addr); + } + llvm_ir::EmitTuple(GetIrArrayFor(custom_call), base_ptrs, &b_, module_); + } auto* output_address_arg = PointerCast(GetEmittedValueFor(custom_call), i8_ptr_type); @@ -2403,14 +2404,8 @@ StatusOr IrEmitter::EmitFastConcatenate( *failure_reason = "operand has mismatching layouts"; return false; } - if (LayoutUtil::IsPadded(op->shape())) { - *failure_reason = "operand has padded layout"; - return false; - } } - CHECK(!LayoutUtil::IsPadded(concatenate->shape())); - // We split the dimensions into three categories: the dimension over which we // are concatenating (concat_dim), the dimensions that are minor to it // (inner_dims) and the dimensions that are major to it (outer_dims). @@ -2418,8 +2413,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), @@ -2592,10 +2586,17 @@ Status IrEmitter::HandleConditional(HloInstruction* conditional) { return Status::OK(); } -Status IrEmitter::HandleAfterAll(HloInstruction* gen_token) { - TF_RET_CHECK(ByteSizeOf(gen_token->shape()) == 0); +Status IrEmitter::HandleAfterAll(HloInstruction* after_all) { + TF_RET_CHECK(ByteSizeOf(after_all->shape()) == 0); // No code to generate, but we need to emit an address for book-keeping. - TF_RETURN_IF_ERROR(EmitTargetAddressForOp(gen_token)); + TF_RETURN_IF_ERROR(EmitTargetAddressForOp(after_all)); + return Status::OK(); +} + +Status IrEmitter::HandleAddDependency(HloInstruction* add_dependency) { + // AddDedendency just forwards its zero-th operand. + emitted_value_[add_dependency] = + GetEmittedValueFor(add_dependency->operand(0)); return Status::OK(); } @@ -2812,7 +2813,7 @@ llvm::Value* IrEmitter::EmitThreadLocalBufferPointer( llvm_ir::EmitBufferIndexingGEP(params, param_number, &b_); llvm::LoadInst* param_address_untyped = Load(param_address_offset); - if (!ShapeUtil::IsOpaque(target_shape)) { + if (!target_shape.IsOpaque()) { AttachAlignmentMetadataForLoad(param_address_untyped, target_shape); AttachDereferenceableMetadataForLoad(param_address_untyped, target_shape); @@ -2871,7 +2872,9 @@ llvm::Value* IrEmitter::EmitBufferPointer(const BufferAllocation::Slice& slice, if (slice.allocation()->is_thread_local()) { return EmitThreadLocalBufferPointer(slice, target_shape); } else if (slice.allocation()->is_constant()) { - return FindOrDie(constant_buffer_to_global_, slice.allocation()->index()); + return BitCast( + FindOrDie(constant_buffer_to_global_, slice.allocation()->index()), + IrShapeType(target_shape)->getPointerTo()); } else { return EmitGlobalBufferPointer(slice, target_shape); } @@ -2964,8 +2967,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 586f27b104ed706a3b128903c6a90abbf3667e59..974dd7cd3f2254bfbc86fffae02c06c481af8902 100644 --- a/tensorflow/compiler/xla/service/cpu/ir_emitter.h +++ b/tensorflow/compiler/xla/service/cpu/ir_emitter.h @@ -59,6 +59,9 @@ namespace cpu { class IrEmitter : public DfsHloVisitorWithDefault, public IrBuilderMixin { public: + using GeneratorForOperandIrArrays = + std::function()>; + // Create a new LLVM IR emitter. // // hlo_module: the HLO module we are emitting IR for. @@ -98,7 +101,7 @@ class IrEmitter : public DfsHloVisitorWithDefault, StatusOr EmitComputation( HloComputation* computation, const string& function_name_prefix, bool is_top_level_computation, - const std::vector* instruction_order); + absl::Span instruction_order); llvm::IRBuilder<>* b() { return &b_; } @@ -131,7 +134,7 @@ class IrEmitter : public DfsHloVisitorWithDefault, Status HandleDot(HloInstruction* dot) override; Status HandleConvolution(HloInstruction* convolution) override; Status HandleFft(HloInstruction* fft) override; - Status HandleCrossReplicaSum(HloInstruction* crs) override; + Status HandleAllReduce(HloInstruction* crs) override; Status HandleInfeed(HloInstruction* infeed) override; Status HandleOutfeed(HloInstruction* outfeed) override; Status HandleSort(HloInstruction* sort) override; @@ -156,7 +159,8 @@ class IrEmitter : public DfsHloVisitorWithDefault, Status HandleConcatenate(HloInstruction* concatenate) override; Status HandleConditional(HloInstruction* conditional) override; Status HandleScatter(HloInstruction* scatter) override; - Status HandleAfterAll(HloInstruction* gen_token) override; + Status HandleAfterAll(HloInstruction* after_all) override; + Status HandleAddDependency(HloInstruction* add_dependency) override; Status HandleRng(HloInstruction* rng) override; Status FinishVisit(HloInstruction* root) override; @@ -208,6 +212,11 @@ class IrEmitter : public DfsHloVisitorWithDefault, std::vector GetIrArraysForOperandsOf( const HloInstruction* hlo); + GeneratorForOperandIrArrays GetGeneratorForOperandIrArrays( + HloInstruction* unnested_hlo) { + return [=]() { return GetIrArraysForOperandsOf(unnested_hlo); }; + } + // Augments IrArray with aliasing information. void AddAliasingInformationToIrArray(const HloInstruction& hlo, llvm_ir::IrArray* array) { @@ -241,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. @@ -439,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_; @@ -459,9 +460,8 @@ class IrEmitter : public DfsHloVisitorWithDefault, // profiling a computation. class ProfilingState { public: - ProfilingState() : use_rdtscp_(false), prof_counters_(nullptr) {} - ProfilingState(bool use_rdtscp, llvm::Value* prof_counters) - : use_rdtscp_(use_rdtscp), prof_counters_(prof_counters) {} + ProfilingState() : use_rdtscp_(false) {} + explicit ProfilingState(bool use_rdtscp) : use_rdtscp_(use_rdtscp) {} // Record the cycle counter before an HLO executes. void RecordCycleStart(llvm::IRBuilder<>* b, HloInstruction* hlo); @@ -486,9 +486,6 @@ class IrEmitter : public DfsHloVisitorWithDefault, // intrinsic? bool use_rdtscp_; - // The argument which corresponds to the profile counter buffer. - llvm::Value* prof_counters_; - // The first read cycle counter in the program. llvm::Value* first_read_cycle_start_ = nullptr; diff --git a/tensorflow/compiler/xla/service/cpu/llvm_ir_runtime.cc b/tensorflow/compiler/xla/service/cpu/llvm_ir_runtime.cc index cef5e57b0b12b7ae93af0d2508b2b9d6a592d390..f9722ffadac801521ddcbb568dd4435fd02e951b 100644 --- a/tensorflow/compiler/xla/service/cpu/llvm_ir_runtime.cc +++ b/tensorflow/compiler/xla/service/cpu/llvm_ir_runtime.cc @@ -22,7 +22,6 @@ limitations under the License. #include "llvm/Transforms/Utils/Cloning.h" #include "tensorflow/compiler/xla/service/cpu/vector_support_library.h" #include "tensorflow/compiler/xla/service/llvm_ir/math_ops.h" -#include "tensorflow/core/lib/core/casts.h" #include "tensorflow/core/platform/logging.h" namespace xla { diff --git a/tensorflow/compiler/xla/service/cpu/parallel_loop_emitter.cc b/tensorflow/compiler/xla/service/cpu/parallel_loop_emitter.cc index f8441c3e345504616485c6b34b4302acd5cc23a3..a6f4273a5a70aab0bc88383283d2a55b1ecb1681 100644 --- a/tensorflow/compiler/xla/service/cpu/parallel_loop_emitter.cc +++ b/tensorflow/compiler/xla/service/cpu/parallel_loop_emitter.cc @@ -34,7 +34,7 @@ ParallelLoopEmitter::EmitIndexAndSetExitBasicBlock(absl::string_view loop_name, llvm::Type* index_type) { CHECK_NE(index_type, nullptr); - CHECK(!ShapeUtil::IsTuple(shape_)); + CHECK(!shape_.IsTuple()); CHECK(!ShapeUtil::IsScalar(shape_)); llvm_ir::ForLoopNest loop_nest(loop_name, b_); diff --git a/tensorflow/compiler/xla/service/cpu/parallel_task_assignment.cc b/tensorflow/compiler/xla/service/cpu/parallel_task_assignment.cc index ede7f433ca6b2cc5629115f800348be9dfb2b93b..6121d1ca9a5c785cedd947200d3e7e320aa06bc2 100644 --- a/tensorflow/compiler/xla/service/cpu/parallel_task_assignment.cc +++ b/tensorflow/compiler/xla/service/cpu/parallel_task_assignment.cc @@ -146,11 +146,9 @@ 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) || - ShapeUtil::IsTuple(instruction->shape())) { + instruction->shape().IsTuple()) { return 1; } diff --git a/tensorflow/compiler/xla/service/cpu/parallel_task_assignment_test.cc b/tensorflow/compiler/xla/service/cpu/parallel_task_assignment_test.cc index fad76338a57cd9eb21d9469ca8552efa8ea0129b..35ae62b42dfa768c6abd0508097d6b235b2ebf54 100644 --- a/tensorflow/compiler/xla/service/cpu/parallel_task_assignment_test.cc +++ b/tensorflow/compiler/xla/service/cpu/parallel_task_assignment_test.cc @@ -17,13 +17,13 @@ limitations under the License. #include "tensorflow/compiler/xla/service/cpu/cpu_executable.h" #include "tensorflow/compiler/xla/service/cpu/target_machine_features_fake.h" #include "tensorflow/compiler/xla/test.h" -#include "tensorflow/compiler/xla/tests/hlo_verified_test_base.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" #include "tensorflow/core/lib/core/status_test_util.h" namespace xla { namespace { -class ParallelTaskAssignmentTest : public HloVerifiedTestBase { +class ParallelTaskAssignmentTest : public HloTestBase { protected: const HloCostAnalysis::ShapeSizeFunction shape_size_func_ = cpu::CpuExecutable::ShapeSizeBytes; @@ -35,7 +35,7 @@ class ParallelTaskAssignmentTest : public HloVerifiedTestBase { cpu::TargetMachineFeaturesWithFakeAlignmentLogic target_machine_features_; ParallelTaskAssignmentTest() - : HloVerifiedTestBase(), target_machine_features_([](int64 shape_size) { + : HloTestBase(), target_machine_features_([](int64 shape_size) { return cpu::TargetMachineFeatures::kEigenExpectedTensorAlignment; }) {} @@ -57,8 +57,9 @@ TEST_F(ParallelTaskAssignmentTest, DotOperationNotParallelized) { } )"; - ParseAndVerifyModule(hlo_string); - TF_ASSERT_OK_AND_ASSIGN(bool changed, RunParallelTaskAssigner(&module())); + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr m, + ParseAndReturnVerifiedModule(hlo_string)); + TF_ASSERT_OK_AND_ASSIGN(bool changed, RunParallelTaskAssigner(m.get())); EXPECT_FALSE(changed); } @@ -84,8 +85,9 @@ TEST_F(ParallelTaskAssignmentTest, } )"; - ParseAndVerifyModule(hlo_string); - TF_ASSERT_OK_AND_ASSIGN(bool changed, RunParallelTaskAssigner(&module())); + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr m, + ParseAndReturnVerifiedModule(hlo_string)); + TF_ASSERT_OK_AND_ASSIGN(bool changed, RunParallelTaskAssigner(m.get())); EXPECT_FALSE(changed); } @@ -100,8 +102,9 @@ TEST_F(ParallelTaskAssignmentTest, RngOperationNotParallelized) { } )"; - ParseAndVerifyModule(hlo_string); - TF_ASSERT_OK_AND_ASSIGN(bool changed, RunParallelTaskAssigner(&module())); + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr m, + ParseAndReturnVerifiedModule(hlo_string)); + TF_ASSERT_OK_AND_ASSIGN(bool changed, RunParallelTaskAssigner(m.get())); EXPECT_FALSE(changed); } @@ -109,15 +112,16 @@ TEST_F(ParallelTaskAssignmentTest, InfeedOutfeedOperationNotParallelized) { const string hlo_string = R"( HloModule TestTaskParallel_infeed_outfeed ENTRY InfeedOutfeed { - token = token[] after-all() - infeed0 = (u32[12345678,2]{1,0}, token[]) infeed(token) + token0 = token[] after-all() + infeed0 = (u32[12345678,2]{1,0}, token[]) infeed(token0) infeed0.data = u32[12345678,2]{1,0} get-tuple-element((u32[12345678,2]{1,0}, token[]) infeed0), index=0 - ROOT outfeed0 = token[] outfeed(infeed0.data, token) + ROOT outfeed0 = token[] outfeed(infeed0.data, token0) } )"; - ParseAndVerifyModule(hlo_string); - TF_ASSERT_OK_AND_ASSIGN(bool changed, RunParallelTaskAssigner(&module())); + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr m, + ParseAndReturnVerifiedModule(hlo_string)); + TF_ASSERT_OK_AND_ASSIGN(bool changed, RunParallelTaskAssigner(m.get())); EXPECT_FALSE(changed); } 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 e0e7deb98e579c090c8fae320a3ba8a3ce8dbe5f..a0667d0d9d1cde246f4b74626859955beeec08b0 100644 --- a/tensorflow/compiler/xla/service/cpu/runtime_key_value_sort.cc +++ b/tensorflow/compiler/xla/service/cpu/runtime_key_value_sort.cc @@ -15,11 +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 "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" #include "tensorflow/core/platform/dynamic_annotations.h" @@ -27,80 +26,20 @@ 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); -} - -// For floating point numbers, we want a total order comparator. -NaN and NaN -// should appear at the beginning and end of the ordering, and -0.0 should -// appear before 0.0. Also we want to have a stable sort, so if the keys are the -// same, we compare the index values. -template -bool LessThan(KeyType lhs, int64 lhs_index, KeyType rhs, int64 rhs_index) { - bool lhs_is_negative = std::signbit(lhs); - bool rhs_is_negative = std::signbit(rhs); - // 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(lhs); - bool rhs_nan = std::isnan(rhs); - // Exactly one number is nan? - if (lhs_nan != rhs_nan) { - if (lhs_nan) { - return lhs_is_negative; - } - return !rhs_is_negative; - } - if (lhs != rhs) { - return lhs < rhs; - } - return lhs_index < rhs_index; -} - -template <> -void KeyValueSort(std::pair* row_to_sort, int64 num_elements) { - std::sort(row_to_sort, row_to_sort + num_elements, - [](const std::pair& lhs, - const std::pair& rhs) -> bool { - return LessThan(lhs.first, lhs.second, rhs.first, rhs.second); - }); -} - -template <> -void KeyValueSort(std::pair* row_to_sort, int64 num_elements) { - std::sort(row_to_sort, row_to_sort + num_elements, - [](const std::pair& lhs, - const std::pair& rhs) -> bool { - return LessThan(lhs.first, lhs.second, rhs.first, rhs.second); - }); -} +} // namespace -template <> -void KeyValueSort(std::pair* row_to_sort, - int64 num_elements) { - std::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), lhs.second, - Eigen::half_impl::half_to_float(rhs.first), rhs.second); - }); -} +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, + bool (*less_than)(char*, char*)) { + // '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*)); + TF_ANNOTATE_MEMORY_IS_INITIALIZED(values_primitive_type_size_in_bytes, + values_count * sizeof(int32)); -template -void KeyValueSortImpl(KeyType* keys, int64 a, int64 b, int64 c, char* values, - int32 values_primitive_type_size_in_bytes) { // High-level idea of the iteration/sorting logic: // Conceptually we have a 3-dimensional shape [a, b, c]. b corresponds to the // dimension to sort, c is the product of the more minor dimensions (set to 1 @@ -114,8 +53,8 @@ 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::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) { @@ -128,109 +67,34 @@ 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 - // both 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; - } - if (values == nullptr) { - continue; - } - - // Reorder the values according to the order defined by the keys. - for (int64 i = 0; i < sort_dimension_elements; ++i) { - int64 memory_index = - (base_offset + row_to_sort[i].second * sort_dimension_offset) * - values_primitive_type_size_in_bytes; - - reordered_values[i] = std::string(values + memory_index, - values_primitive_type_size_in_bytes); - } - for (int64 i = 0; i < sort_dimension_elements; ++i) { - int64 memory_index = (base_offset + i * sort_dimension_offset) * - values_primitive_type_size_in_bytes; - memcpy(values + memory_index, reordered_values[i].c_str(), - values_primitive_type_size_in_bytes); + std::stable_sort( + indices.get(), indices.get() + sort_dimension_elements, + [&](int64 a, int64 b) { + 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]; + return less_than(values[0] + memory_index_lhs, + values[0] + memory_index_rhs); + }); + + // 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 + indices[i] * sort_dimension_offset) * + values_primitive_type_size_in_bytes[idx]; + + reordered_values[i] = + std::string(values[idx] + memory_index, + values_primitive_type_size_in_bytes[idx]); + } + for (int64 i = 0; i < sort_dimension_elements; ++i) { + int64 memory_index = (base_offset + i * sort_dimension_offset) * + values_primitive_type_size_in_bytes[idx]; + memcpy(values[idx] + memory_index, reordered_values[i].c_str(), + values_primitive_type_size_in_bytes[idx]); + } } } } -} // namespace - -TF_ATTRIBUTE_NO_SANITIZE_MEMORY void __xla_cpu_runtime_KeyValueSortPRED( - bool* keys, int64 a, int64 b, int64 c, char* values, - int32 values_primitive_type_size_in_bytes) { - KeyValueSortImpl(keys, a, b, c, values, 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_primitive_type_size_in_bytes) { - KeyValueSortImpl(keys, a, b, c, values, 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_primitive_type_size_in_bytes) { - KeyValueSortImpl(keys, a, b, c, values, 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_primitive_type_size_in_bytes) { - KeyValueSortImpl(keys, a, b, c, values, 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_primitive_type_size_in_bytes) { - KeyValueSortImpl(keys, a, b, c, values, 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_primitive_type_size_in_bytes) { - KeyValueSortImpl(keys, a, b, c, values, 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_primitive_type_size_in_bytes) { - KeyValueSortImpl(keys, a, b, c, values, 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_primitive_type_size_in_bytes) { - KeyValueSortImpl(keys, a, b, c, values, 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_primitive_type_size_in_bytes) { - KeyValueSortImpl(keys, a, b, c, values, 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_primitive_type_size_in_bytes) { - KeyValueSortImpl(keys, a, b, c, values, 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_primitive_type_size_in_bytes) { - KeyValueSortImpl(keys, a, b, c, values, 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_primitive_type_size_in_bytes) { - KeyValueSortImpl(keys, a, b, c, values, 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 28e35e82c18cbf078f8a1e7f5b818bf839d3d3df..5460af3485b94aaef1a5822a79e4fa325bcb67ea 100644 --- a/tensorflow/compiler/xla/service/cpu/runtime_key_value_sort.h +++ b/tensorflow/compiler/xla/service/cpu/runtime_key_value_sort.h @@ -21,68 +21,19 @@ 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. 'values' can be nullptr. -// If 'values' is not nullptr, the elements in 'values' 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 'values' 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, - char* values, 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_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_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_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_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_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_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_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_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_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_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_primitive_type_size_in_bytes); +// 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. +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, + bool (*less_than)(char*, char*)); } #endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_RUNTIME_KEY_VALUE_SORT_H_ diff --git a/tensorflow/compiler/xla/service/cpu/runtime_matmul.cc b/tensorflow/compiler/xla/service/cpu/runtime_matmul.cc index a71a85913cfef271bc2a226cb0cf2dd4204499a4..fe7e87a197b6cf571195537eaea2898659cd5e2e 100644 --- a/tensorflow/compiler/xla/service/cpu/runtime_matmul.cc +++ b/tensorflow/compiler/xla/service/cpu/runtime_matmul.cc @@ -23,12 +23,20 @@ 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; namespace { -template +bool Is16BytesAligned(void* ptr) { + return reinterpret_cast(ptr) % 16 == 0; +} + +template void MatMul(const void* run_options_ptr, T* out, T* lhs, T* rhs, int64 m, int64 n, int64 k, int32 transpose_lhs, int32 transpose_rhs) { const xla::ExecutableRunOptions* run_options = @@ -46,11 +54,11 @@ void MatMul(const void* run_options_ptr, T* out, T* lhs, T* rhs, int64 m, std::swap(rhs_rows, rhs_cols); } - const Eigen::TensorMap, Eigen::Aligned> A( - lhs, lhs_rows, lhs_cols); - const Eigen::TensorMap, Eigen::Aligned> B( - rhs, rhs_rows, rhs_cols); - Eigen::TensorMap, Eigen::Aligned> C(out, m, n); + const Eigen::TensorMap, Alignment> A(lhs, lhs_rows, + lhs_cols); + const Eigen::TensorMap, Alignment> B(rhs, rhs_rows, + rhs_cols); + Eigen::TensorMap, Alignment> C(out, m, n); typedef typename Eigen::Tensor::DimensionPair DimPair; int lhs_contract_dim = transpose_lhs ? 0 : 1; @@ -65,14 +73,24 @@ void MatMul(const void* run_options_ptr, T* out, T* lhs, T* rhs, int64 m, } template -void MatMulImpl(const void* run_options_ptr, T* out, T* lhs, T* rhs, int64 m, - int64 n, int64 k, int32 transpose_lhs, int32 transpose_rhs) { +void MatMulDispatch(const void* run_options_ptr, T* out, T* lhs, T* rhs, + int64 m, int64 n, int64 k, int32 transpose_lhs, + int32 transpose_rhs) { + bool all_buffers_16b_aligned = + Is16BytesAligned(out) && Is16BytesAligned(lhs) && Is16BytesAligned(rhs); + + if (!all_buffers_16b_aligned) { + MatMul(run_options_ptr, out, lhs, rhs, m, n, k, + transpose_lhs, transpose_rhs); + return; + } + if (m == 1 || n == 1) { // Despite being single threaded, this version of matrix * vector is faster. xla::EigenMatVec(out, lhs, rhs, m, n, k, transpose_lhs, transpose_rhs); } else { - MatMul(run_options_ptr, out, lhs, rhs, m, n, k, transpose_lhs, - transpose_rhs); + MatMul(run_options_ptr, out, lhs, rhs, m, n, k, + transpose_lhs, transpose_rhs); } } @@ -82,20 +100,20 @@ TF_ATTRIBUTE_NO_SANITIZE_MEMORY void __xla_cpu_runtime_EigenMatMulF16( const void* run_options_ptr, Eigen::half* out, Eigen::half* lhs, Eigen::half* rhs, int64 m, int64 n, int64 k, int32 transpose_lhs, int32 transpose_rhs) { - MatMulImpl(run_options_ptr, out, lhs, rhs, m, n, k, - transpose_lhs, transpose_rhs); + MatMulDispatch(run_options_ptr, out, lhs, rhs, m, n, k, + transpose_lhs, transpose_rhs); } TF_ATTRIBUTE_NO_SANITIZE_MEMORY void __xla_cpu_runtime_EigenMatMulF32( const void* run_options_ptr, float* out, float* lhs, float* rhs, int64 m, int64 n, int64 k, int32 transpose_lhs, int32 transpose_rhs) { - MatMulImpl(run_options_ptr, out, lhs, rhs, m, n, k, transpose_lhs, - transpose_rhs); + MatMulDispatch(run_options_ptr, out, lhs, rhs, m, n, k, transpose_lhs, + transpose_rhs); } TF_ATTRIBUTE_NO_SANITIZE_MEMORY void __xla_cpu_runtime_EigenMatMulF64( const void* run_options_ptr, double* out, double* lhs, double* rhs, int64 m, int64 n, int64 k, int32 transpose_lhs, int32 transpose_rhs) { - MatMulImpl(run_options_ptr, out, lhs, rhs, m, n, k, transpose_lhs, - transpose_rhs); + MatMulDispatch(run_options_ptr, out, lhs, rhs, m, n, k, transpose_lhs, + transpose_rhs); } 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 16692e7f2e6145b2649b67987eef47916e958be2..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,12 +20,20 @@ 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; namespace { -template +bool Is16BytesAligned(void* ptr) { + return reinterpret_cast(ptr) % 16 == 0; +} + +template void MatMul(const void* run_options_ptr, T* out, T* lhs, T* rhs, int64 m, int64 n, int64 k, int32 transpose_lhs, int32 transpose_rhs) { int64 lhs_rows = m; @@ -40,11 +48,11 @@ void MatMul(const void* run_options_ptr, T* out, T* lhs, T* rhs, int64 m, std::swap(rhs_rows, rhs_cols); } - const Eigen::TensorMap, Eigen::Aligned> A( - lhs, lhs_rows, lhs_cols); - const Eigen::TensorMap, Eigen::Aligned> B( - rhs, rhs_rows, rhs_cols); - Eigen::TensorMap, Eigen::Aligned> C(out, m, n); + const Eigen::TensorMap, Alignment> A(lhs, lhs_rows, + lhs_cols); + const Eigen::TensorMap, Alignment> B(rhs, rhs_rows, + rhs_cols); + Eigen::TensorMap, Alignment> C(out, m, n); typedef typename Eigen::Tensor::DimensionPair DimPair; int lhs_contract_dim = transpose_lhs ? 0 : 1; @@ -59,14 +67,22 @@ void MatMul(const void* run_options_ptr, T* out, T* lhs, T* rhs, int64 m, } template -void SingleThreadedMatMul(const void* run_options_ptr, T* out, T* lhs, T* rhs, - int64 m, int64 n, int64 k, int32 transpose_lhs, - int32 transpose_rhs) { +void SingleThreadedMatMulDispatch(const void* run_options_ptr, T* out, T* lhs, + T* rhs, int64 m, int64 n, int64 k, + int32 transpose_lhs, int32 transpose_rhs) { + bool all_buffers_16b_aligned = + Is16BytesAligned(out) && Is16BytesAligned(lhs) && Is16BytesAligned(rhs); + + if (!all_buffers_16b_aligned) { + MatMul(run_options_ptr, out, lhs, rhs, m, n, k, + transpose_lhs, transpose_rhs); + } + if (m == 1 || n == 1) { xla::EigenMatVec(out, lhs, rhs, m, n, k, transpose_lhs, transpose_rhs); } else { - MatMul(run_options_ptr, out, lhs, rhs, m, n, k, transpose_lhs, - transpose_rhs); + MatMul(run_options_ptr, out, lhs, rhs, m, n, k, + transpose_lhs, transpose_rhs); } } @@ -77,8 +93,8 @@ __xla_cpu_runtime_EigenSingleThreadedMatMulF16( const void* run_options_ptr, Eigen::half* out, Eigen::half* lhs, Eigen::half* rhs, int64 m, int64 n, int64 k, int32 transpose_lhs, int32 transpose_rhs) { - SingleThreadedMatMul(run_options_ptr, out, lhs, rhs, m, n, k, - transpose_lhs, transpose_rhs); + SingleThreadedMatMulDispatch(run_options_ptr, out, lhs, rhs, m, + n, k, transpose_lhs, transpose_rhs); } TF_ATTRIBUTE_NO_SANITIZE_MEMORY void @@ -87,8 +103,8 @@ __xla_cpu_runtime_EigenSingleThreadedMatMulF32(const void* run_options_ptr, float* rhs, int64 m, int64 n, int64 k, int32 transpose_lhs, int32 transpose_rhs) { - SingleThreadedMatMul(run_options_ptr, out, lhs, rhs, m, n, k, - transpose_lhs, transpose_rhs); + SingleThreadedMatMulDispatch(run_options_ptr, out, lhs, rhs, m, n, k, + transpose_lhs, transpose_rhs); } TF_ATTRIBUTE_NO_SANITIZE_MEMORY void @@ -97,6 +113,6 @@ __xla_cpu_runtime_EigenSingleThreadedMatMulF64(const void* run_options_ptr, double* rhs, int64 m, int64 n, int64 k, int32 transpose_lhs, int32 transpose_rhs) { - SingleThreadedMatMul(run_options_ptr, out, lhs, rhs, m, n, k, - transpose_lhs, transpose_rhs); + SingleThreadedMatMulDispatch(run_options_ptr, out, lhs, rhs, m, n, k, + transpose_lhs, transpose_rhs); } diff --git a/tensorflow/compiler/xla/service/cpu/shape_partition_test.cc b/tensorflow/compiler/xla/service/cpu/shape_partition_test.cc index 1a3d82de954318368d61e3feeb0345dc592dcd8b..7d8e51f909e3db699b745f94a6c625407bc4a6e3 100644 --- a/tensorflow/compiler/xla/service/cpu/shape_partition_test.cc +++ b/tensorflow/compiler/xla/service/cpu/shape_partition_test.cc @@ -19,14 +19,14 @@ limitations under the License. #include #include "tensorflow/compiler/xla/test_helpers.h" -#include "tensorflow/compiler/xla/tests/hlo_verified_test_base.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" #include "tensorflow/compiler/xla/util.h" namespace xla { namespace cpu { namespace { -class ShapePartitionAssignerTest : public HloVerifiedTestBase { +class ShapePartitionAssignerTest : public HloTestBase { protected: typedef std::vector Vec; @@ -91,7 +91,7 @@ TEST_F(ShapePartitionAssignerTest, Shape532WithLayout201) { expected_partitions); } -class ShapePartitionIteratorTest : public HloVerifiedTestBase { +class ShapePartitionIteratorTest : public HloTestBase { protected: typedef std::vector> Partition; }; @@ -145,7 +145,7 @@ TEST_F(ShapePartitionIteratorTest, Shape532WithLayout210) { } } -class RandomShapePartitionIteratorTest : public HloVerifiedTestBase { +class RandomShapePartitionIteratorTest : public HloTestBase { protected: typedef std::vector> Partition; RandomShapePartitionIteratorTest() diff --git a/tensorflow/compiler/xla/service/cpu/simple_orc_jit.cc b/tensorflow/compiler/xla/service/cpu/simple_orc_jit.cc index f77641eb7da71117092730c1fd5090c61c939813..9c2685674fbc133de1220caef81ac3b60a1c0f7c 100644 --- a/tensorflow/compiler/xla/service/cpu/simple_orc_jit.cc +++ b/tensorflow/compiler/xla/service/cpu/simple_orc_jit.cc @@ -116,20 +116,43 @@ 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(); } llvm::JITSymbol SimpleOrcJIT::ResolveRuntimeSymbol(const std::string& name) { - void* func_addr = CustomCallTargetRegistry::Global()->Lookup(name); + void* func_addr = nullptr; + if (name.size() > 1 && name.front() == data_layout_.getGlobalPrefix()) { + // On Mac OS X, 'name' may have a leading underscore prefix, even though the + // registered name may not. + std::string stripped_name(name.begin() + 1, name.end()); + func_addr = CustomCallTargetRegistry::Global()->Lookup(stripped_name); + } else { + func_addr = CustomCallTargetRegistry::Global()->Lookup(name); + } + if (func_addr == nullptr) { + LOG(ERROR) << "Unable to resolve runtime symbol: " << name; return nullptr; } llvm::JITEvaluatedSymbol symbol_info(reinterpret_cast(func_addr), @@ -137,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(); @@ -203,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)); @@ -286,6 +312,9 @@ bool RegisterKnownJITSymbols() { REGISTER_LIBM_SYMBOL(sin, double (*)(double)); #ifdef __APPLE__ REGISTER_LIBM_SYMBOL(__sincos, void (*)(double, double*, double*)); + registry->Register("__sincosf_stret", + reinterpret_cast(__sincosf_stret)); + registry->Register("__sincos_stret", reinterpret_cast(__sincos_stret)); #else REGISTER_LIBM_SYMBOL(sincos, void (*)(double, double*, double*)); #endif @@ -301,6 +330,13 @@ bool RegisterKnownJITSymbols() { registry->Register("memcpy", reinterpret_cast(memcpy)); registry->Register("memmove", reinterpret_cast(memmove)); registry->Register("memset", reinterpret_cast(memset)); + +#ifdef __APPLE__ + registry->Register("__bzero", reinterpret_cast(bzero)); + registry->Register("memset_pattern16", + reinterpret_cast(memset_pattern16)); +#endif + return true; } 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/BUILD b/tensorflow/compiler/xla/service/cpu/tests/BUILD index 4b129c95d46d8b5a119e5d23eef387daf7863cce..382dfd0d99df87bbadfe541ddaa32cd6da8e8068 100644 --- a/tensorflow/compiler/xla/service/cpu/tests/BUILD +++ b/tensorflow/compiler/xla/service/cpu/tests/BUILD @@ -48,7 +48,6 @@ tf_cc_test( "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/service/cpu:cpu_instruction_fusion", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", "//tensorflow/compiler/xla/tests:literal_test_util", "//tensorflow/core:test", "//tensorflow/core:test_main", 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 18ee25ba9158c28baaf01492c290638b9673f1ec..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" @@ -50,7 +49,7 @@ class CpuEigenDotOperationTest /*entry_point_name=*/"entry", /*relocation_model=*/CpuAotCompilationOptions::RelocationModel::Static}; - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewVerifiedModule(); hlo_module->AddEntryComputation(std::move(entry_computation)); CompileAheadOfTimeAndVerifyIr(std::move(hlo_module), options, @@ -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_external_constants_test.cc b/tensorflow/compiler/xla/service/cpu/tests/cpu_external_constants_test.cc index 00a7aa2ad2f6bac4877302296ccb76222557535c..e30f95311fce229f9c559d3bb40142151e8bf3e3 100644 --- a/tensorflow/compiler/xla/service/cpu/tests/cpu_external_constants_test.cc +++ b/tensorflow/compiler/xla/service/cpu/tests/cpu_external_constants_test.cc @@ -46,7 +46,7 @@ class CpuExternalConstantsTest : public CpuCodegenTest { builder.AddInstruction( HloInstruction::CreateBinary(shape, HloOpcode::kAdd, param, constant)); - std::unique_ptr module = CreateNewModule(); + std::unique_ptr module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); CompileAndVerifyIr(std::move(module), filecheck_pattern, diff --git a/tensorflow/compiler/xla/service/cpu/tests/cpu_fusion_test.cc b/tensorflow/compiler/xla/service/cpu/tests/cpu_fusion_test.cc index 1deb412064b02988a8d4a6d726969c948d354d47..04a81dfd35f459ff1fdb3181dc8fc65c62a37d4f 100644 --- a/tensorflow/compiler/xla/service/cpu/tests/cpu_fusion_test.cc +++ b/tensorflow/compiler/xla/service/cpu/tests/cpu_fusion_test.cc @@ -25,7 +25,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" #include "tensorflow/compiler/xla/shape_util.h" -#include "tensorflow/compiler/xla/tests/hlo_verified_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/xla_data.pb.h" #include "tensorflow/core/platform/test.h" @@ -34,7 +34,7 @@ namespace xla { namespace cpu { namespace { -class CpuFusionTest : public HloVerifiedTestBase { +class CpuFusionTest : public HloTestBase { protected: CpuFusionTest() {} @@ -57,11 +57,11 @@ TEST_F(CpuFusionTest, FuseTwoElementwiseOps) { builder.AddInstruction( HloInstruction::CreateUnary(vshape, HloOpcode::kNegate, add1)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); CpuInstructionFusion fusion; - EXPECT_TRUE(fusion.Run(module).ValueOrDie()); + EXPECT_TRUE(fusion.Run(module.get()).ValueOrDie()); // The computation root instruction was fused. Verify the fusion instruction // is now the root. @@ -104,11 +104,11 @@ TEST_F(CpuFusionTest, FuseElementwiseOpChain) { builder.AddInstruction( HloInstruction::CreateBinary(vshape, HloOpcode::kMultiply, two, floor)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); CpuInstructionFusion fusion; - EXPECT_TRUE(fusion.Run(module).ValueOrDie()); + EXPECT_TRUE(fusion.Run(module.get()).ValueOrDie()); // The computation root instruction was fused. Verify the fusion instruction // is now the root. @@ -131,7 +131,7 @@ TEST_F(CpuFusionTest, FuseElementwiseOpChain) { TEST_F(CpuFusionTest, ElementwiseOpChainWithNonfusibleInstruction) { // Test a chain of fusible ops with a non-fusible op (a reduce) thrown in the // middle. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto builder = HloComputation::Builder(TestName()); auto input_literal = LiteralUtil::CreateR1({-1.5, -2.5, -3.0}); Shape vshape = input_literal.shape(); @@ -183,7 +183,7 @@ TEST_F(CpuFusionTest, ElementwiseOpChainWithNonfusibleInstruction) { module->AddEntryComputation(builder.Build()); CpuInstructionFusion fusion; - EXPECT_TRUE(fusion.Run(module).ValueOrDie()); + EXPECT_TRUE(fusion.Run(module.get()).ValueOrDie()); // The computation root instruction was fused. Verify the fusion instruction // is now the root. @@ -250,12 +250,12 @@ TEST_F(CpuFusionTest, TestOperandOrderToAvoidDuplication) { builder.AddInstruction(HloInstruction::CreateTuple({add1, add2})); // Create computation and module. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); // Run fusion. CpuInstructionFusion fusion; - EXPECT_TRUE(fusion.Run(module).ValueOrDie()); + EXPECT_TRUE(fusion.Run(module.get()).ValueOrDie()); auto fusion1 = result->operand(0); auto fusion2 = result->operand(1); @@ -310,11 +310,11 @@ TEST_F(CpuFusionTest, DoNotDuplicateExpensiveOps) { auto tuple = builder.AddInstruction( HloInstruction::CreateTuple({negate1, negate2, exp2})); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); CpuInstructionFusion fusion; - EXPECT_TRUE(fusion.Run(module).ValueOrDie()); + EXPECT_TRUE(fusion.Run(module.get()).ValueOrDie()); // The only fusion instruction should be operand 0 of the tuple (formerly // negate1). diff --git a/tensorflow/compiler/xla/service/cpu/tests/cpu_infeed_test.cc b/tensorflow/compiler/xla/service/cpu/tests/cpu_infeed_test.cc index 5cc6d01c0f15d4209cbc1fb259a0078fb9957f6e..f0f897e9635600b22e0c389ba056899e4d6ab3d4 100644 --- a/tensorflow/compiler/xla/service/cpu/tests/cpu_infeed_test.cc +++ b/tensorflow/compiler/xla/service/cpu/tests/cpu_infeed_test.cc @@ -48,7 +48,7 @@ class InfeedTest : public ClientLibraryTestBase { ASSERT_IS_OK(client_->TransferToInfeed(literal)); XlaBuilder builder(TestName()); Infeed(&builder, literal.shape()); - if (ShapeUtil::IsTuple(literal.shape())) { + if (literal.shape().IsTuple()) { // TODO(b/30609564): Use ComputeAndCompareLiteral instead. ComputeAndCompareTuple(&builder, literal, {}); } else { 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 a434c04a980b9b3cd849792b97a0d9e965ba09f2..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 = ""; } @@ -91,7 +92,7 @@ TEST_P(CpuUnaryIntrinsicTest, DoIt) { /*entry_point_name=*/"entry", /*relocation_model=*/CpuAotCompilationOptions::RelocationModel::Static}; - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewVerifiedModule(); hlo_module->AddEntryComputation(std::move(computation)); string check_lines{spec.check_lines.data(), spec.check_lines.size()}; @@ -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_literal_caching_test.cc b/tensorflow/compiler/xla/service/cpu/tests/cpu_literal_caching_test.cc index 3b87683ffffefd2aa24dd234cc072425bef00a24..0584c0484f810a03ccccd522163f54535440ef8b 100644 --- a/tensorflow/compiler/xla/service/cpu/tests/cpu_literal_caching_test.cc +++ b/tensorflow/compiler/xla/service/cpu/tests/cpu_literal_caching_test.cc @@ -31,29 +31,27 @@ HloModule RepeatedConstants while_body { arg_body = f32[2,3,2] parameter(0) ROOT const = f32[2,3,2] constant( - f32[2,3,2] {{{1, 2}, {1001, 1002}, {2001, 2002}}, {{2, 1}, {2001, 3002}, {2001, 2002}}}) } while_cond { arg_cond = f32[2,3,2] parameter(0) - token = token[] after-all() - infeed = (pred[], token[]) infeed(token) + token0 = token[] after-all() + infeed = (pred[], token[]) infeed(token0) ROOT unknown = pred[] get-tuple-element((pred[], token[]) infeed), index=0 } ENTRY main { param = f32[2,3,2] parameter(0) const_a = f32[2,3,2] constant( - f32[2,3,2] {{{1, 2}, {1001, 1002}, {2001, 2002}}, {{2, 1}, {2001, 3002}, {2001, 2002}}}) const_b = f32[2,3,2] while(f32[2,3,2] const_a), condition=while_cond, body=while_body - token = token[] after-all() - out0 = token[] outfeed(f32[2,3,2] const_a, token[] token) - ROOT out1 = token[] outfeed(f32[2,3,2] const_b, token[] token) + token0 = token[] after-all() + out0 = token[] outfeed(f32[2,3,2] const_a, token[] token0) + ROOT out1 = token[] outfeed(f32[2,3,2] const_b, token[] token0) } )"; @@ -63,7 +61,7 @@ CHECK-NOT: private constant [48 x i8] )"; TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseHloString(hlo_text)); + ParseAndReturnVerifiedModule(hlo_text)); CpuAotCompilationOptions options{ /*triple=*/"x86_64-pc-linux", /*cpu_name=*/"", /*features=*/"", @@ -82,36 +80,36 @@ HloModule RepeatedConstants while_body { arg_body = (f32[2,1]{1,0}, f32[1]{0}) parameter(0) - ROOT const = (f32[2,1]{1,0}, f32[1]{0}) constant((f32[2,1], f32[1]) ( f32[2,1] { { 1 }, { 2 } }, {2} )) + ROOT const = (f32[2,1]{1,0}, f32[1]{0}) constant(({ { 1 }, { 2 } }, {2} )) } while_cond { arg_cond = (f32[2,1]{1,0}, f32[1]{0}) parameter(0) - token = token[] after-all() - infeed = (pred[], token[]) infeed(token) + token0 = token[] after-all() + infeed = (pred[], token[]) infeed(token0) ROOT unknown = pred[] get-tuple-element((pred[], token[]) infeed), index=0 } ENTRY main { param = f32[2,3,2] parameter(0) - const_a = (f32[2,1]{1,0}, f32[1]{0}) constant((f32[2,1], f32[1]) ( f32[2,1] { { 1 }, { 2 } }, {2} )) + const_a = (f32[2,1]{1,0}, f32[1]{0}) constant(( { { 1 }, { 2 } }, {2} )) const_b = (f32[2,1]{1,0}, f32[1]{0}) while((f32[2,1]{1,0}, f32[1]{0}) const_a), condition=while_cond, body=while_body - token = token[] after-all() - out0 = () outfeed((f32[2,1]{1,0}, f32[1]{0}) const_a, token[] token) - ROOT out1 = () outfeed((f32[2,1]{1,0}, f32[1]{0}) const_b, token[] token) + token0 = token[] after-all() + out0 = () outfeed((f32[2,1]{1,0}, f32[1]{0}) const_a, token[] token0) + ROOT out1 = () outfeed((f32[2,1]{1,0}, f32[1]{0}) const_b, token[] token0) } )"; string filecheck_pattern = R"( -CHECK: private constant [4 x i8] -CHECK: private constant [8 x i8] +CHECK-DAG: private constant [4 x i8] +CHECK-DAG: private constant [8 x i8] CHECK-NOT: private constant [4 x i8] CHECK-NOT: private constant [8 x i8] )"; TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseHloString(hlo_text)); + ParseAndReturnVerifiedModule(hlo_text)); CpuAotCompilationOptions options{ /*triple=*/"x86_64-pc-linux", /*cpu_name=*/"", /*features=*/"", 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 b35fd9dad877c319c3d0110c96a00aeefa78769e..a7702c2aeeaff8a46a2c4f2785ccb873ea2c08e5 100644 --- a/tensorflow/compiler/xla/service/cpu/tests/cpu_noalias_test.cc +++ b/tensorflow/compiler/xla/service/cpu/tests/cpu_noalias_test.cc @@ -56,7 +56,7 @@ TEST_F(CpuNoAliasTest, Concat) { std::unique_ptr computation = builder.Build(); - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewVerifiedModule(); hlo_module->AddEntryComputation(std::move(computation)); // Now that we have an HLO module, build an llvm_ir::AliasAnalysis for it. diff --git a/tensorflow/compiler/xla/service/cpu/tests/cpu_outfeed_test.cc b/tensorflow/compiler/xla/service/cpu/tests/cpu_outfeed_test.cc index e2c7af541eede5265f274c72f55305549f059839..aab7f0b393881642437f1891256bd138823a3b87 100644 --- a/tensorflow/compiler/xla/service/cpu/tests/cpu_outfeed_test.cc +++ b/tensorflow/compiler/xla/service/cpu/tests/cpu_outfeed_test.cc @@ -28,12 +28,11 @@ HloModule Outfeed ENTRY main { const_a = f32[2,3,2] constant( - f32[2,3,2] {{{1, 2}, {1001, 1002}, {2001, 2002}}, {{2, 1}, {2001, 3002}, {2001, 2002}}}) - token = token[] after-all() - outfeed = token[] outfeed(f32[2,3,2] const_a, token) + token0 = token[] after-all() + outfeed = token[] outfeed(f32[2,3,2] const_a, token0) ROOT root = () tuple() } )"; 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/xfeed_manager.h b/tensorflow/compiler/xla/service/cpu/xfeed_manager.h index 990ff94ba2338cb663b655ca3106bda83ab718a3..70008947f371d25e95d02839c30ba822fce7a292 100644 --- a/tensorflow/compiler/xla/service/cpu/xfeed_manager.h +++ b/tensorflow/compiler/xla/service/cpu/xfeed_manager.h @@ -23,6 +23,7 @@ limitations under the License. #include #include "absl/types/span.h" +#include "tensorflow/compiler/xla/shape.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" diff --git a/tensorflow/compiler/xla/service/defuser_test.cc b/tensorflow/compiler/xla/service/defuser_test.cc index e727ba49cb6321e499b5d50d5f45e7f7f6bb6fef..64fb50318394918b277fd717994f5366d762ac36 100644 --- a/tensorflow/compiler/xla/service/defuser_test.cc +++ b/tensorflow/compiler/xla/service/defuser_test.cc @@ -18,19 +18,19 @@ limitations under the License. #include "tensorflow/compiler/xla/literal.h" #include "tensorflow/compiler/xla/service/hlo_matchers.h" #include "tensorflow/compiler/xla/shape_util.h" -#include "tensorflow/compiler/xla/tests/hlo_verified_test_base.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" namespace op = xla::testing::opcode_matchers; namespace xla { namespace { -class DefuserTest : public HloVerifiedTestBase { +class DefuserTest : public HloTestBase { protected: // Returns the number of fusion instructions in the module. - int FusionCount() { + int FusionCount(const HloModule* m) { int count = 0; - for (HloComputation* computation : module().computations()) { + for (HloComputation* computation : m->computations()) { if (computation->IsFusionComputation()) { count++; } @@ -43,6 +43,7 @@ class DefuserTest : public HloVerifiedTestBase { }; TEST_F(DefuserTest, NoFusionInstruction) { + auto m = CreateNewVerifiedModule(); auto builder = HloComputation::Builder(TestName()); auto param0 = builder.AddInstruction(HloInstruction::CreateParameter(0, shape_, "p0")); @@ -51,13 +52,14 @@ TEST_F(DefuserTest, NoFusionInstruction) { builder.AddInstruction( HloInstruction::CreateBinary(shape_, HloOpcode::kAdd, param0, param1)); - module().AddEntryComputation(builder.Build()); - EXPECT_EQ(0, FusionCount()); + m->AddEntryComputation(builder.Build()); + EXPECT_EQ(0, FusionCount(m.get())); - EXPECT_FALSE(defuser_.Run(&module()).ValueOrDie()); + EXPECT_FALSE(defuser_.Run(m.get()).ValueOrDie()); } TEST_F(DefuserTest, TrivialFusionInstructionAsRoot) { + auto m = CreateNewVerifiedModule(); auto builder = HloComputation::Builder(TestName()); auto param0 = builder.AddInstruction(HloInstruction::CreateParameter(0, shape_, "p0")); @@ -66,21 +68,22 @@ TEST_F(DefuserTest, TrivialFusionInstructionAsRoot) { auto add = builder.AddInstruction( HloInstruction::CreateBinary(shape_, HloOpcode::kAdd, param0, param1)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); computation->CreateFusionInstruction({add}, HloInstruction::FusionKind::kLoop); EXPECT_THAT(computation->root_instruction(), op::Fusion()); - EXPECT_EQ(1, FusionCount()); - EXPECT_TRUE(defuser_.Run(&module()).ValueOrDie()); - EXPECT_EQ(0, FusionCount()); + EXPECT_EQ(1, FusionCount(m.get())); + EXPECT_TRUE(defuser_.Run(m.get()).ValueOrDie()); + EXPECT_EQ(0, FusionCount(m.get())); EXPECT_THAT(computation->root_instruction(), op::Add(op::Parameter(), op::Parameter())); } TEST_F(DefuserTest, TrivialFusionInstructionNotAsRoot) { + auto m = CreateNewVerifiedModule(); auto builder = HloComputation::Builder(TestName()); auto param0 = builder.AddInstruction(HloInstruction::CreateParameter(0, shape_, "p0")); @@ -91,21 +94,22 @@ TEST_F(DefuserTest, TrivialFusionInstructionNotAsRoot) { builder.AddInstruction( HloInstruction::CreateUnary(shape_, HloOpcode::kNegate, add)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); computation->CreateFusionInstruction({add}, HloInstruction::FusionKind::kLoop); EXPECT_THAT(computation->root_instruction(), op::Negate(op::Fusion())); - EXPECT_EQ(1, FusionCount()); - EXPECT_TRUE(defuser_.Run(&module()).ValueOrDie()); - EXPECT_EQ(0, FusionCount()); + EXPECT_EQ(1, FusionCount(m.get())); + EXPECT_TRUE(defuser_.Run(m.get()).ValueOrDie()); + EXPECT_EQ(0, FusionCount(m.get())); EXPECT_THAT(computation->root_instruction(), op::Negate(op::Add(op::Parameter(), op::Parameter()))); } TEST_F(DefuserTest, NonTrivialFusionInstruction) { + auto m = CreateNewVerifiedModule(); auto builder = HloComputation::Builder(TestName()); auto param0 = builder.AddInstruction(HloInstruction::CreateParameter(0, shape_, "p0")); @@ -128,22 +132,23 @@ TEST_F(DefuserTest, NonTrivialFusionInstruction) { auto add2 = builder.AddInstruction( HloInstruction::CreateBinary(shape_, HloOpcode::kAdd, constant, div)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); computation->CreateFusionInstruction( {add2, constant, div, mul, sub, negate, add}, HloInstruction::FusionKind::kLoop); EXPECT_THAT(computation->root_instruction(), op::Fusion()); - EXPECT_EQ(1, FusionCount()); - EXPECT_TRUE(defuser_.Run(&module()).ValueOrDie()); - EXPECT_EQ(0, FusionCount()); + EXPECT_EQ(1, FusionCount(m.get())); + EXPECT_TRUE(defuser_.Run(m.get()).ValueOrDie()); + EXPECT_EQ(0, FusionCount(m.get())); EXPECT_THAT(computation->root_instruction(), op::Add(op::Constant(), op::Divide())); } TEST_F(DefuserTest, MultipleFusionInstructions) { + auto m = CreateNewVerifiedModule(); auto builder = HloComputation::Builder(TestName()); auto param0 = builder.AddInstruction(HloInstruction::CreateParameter(0, shape_, "p0")); @@ -166,7 +171,7 @@ TEST_F(DefuserTest, MultipleFusionInstructions) { auto add2 = builder.AddInstruction( HloInstruction::CreateBinary(shape_, HloOpcode::kAdd, constant, div)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); computation->CreateFusionInstruction({add2, constant, div, mul}, HloInstruction::FusionKind::kLoop); computation->CreateFusionInstruction({sub, negate, add}, @@ -174,15 +179,16 @@ TEST_F(DefuserTest, MultipleFusionInstructions) { EXPECT_THAT(computation->root_instruction(), op::Fusion()); - EXPECT_EQ(2, FusionCount()); - EXPECT_TRUE(defuser_.Run(&module()).ValueOrDie()); - EXPECT_EQ(0, FusionCount()); + EXPECT_EQ(2, FusionCount(m.get())); + EXPECT_TRUE(defuser_.Run(m.get()).ValueOrDie()); + EXPECT_EQ(0, FusionCount(m.get())); EXPECT_THAT(computation->root_instruction(), op::Add(op::Constant(), op::Divide())); } TEST_F(DefuserTest, NestedFusionInstructions) { + auto m = CreateNewVerifiedModule(); auto builder = HloComputation::Builder(TestName()); auto param0 = builder.AddInstruction(HloInstruction::CreateParameter(0, shape_, "p0")); @@ -193,7 +199,7 @@ TEST_F(DefuserTest, NestedFusionInstructions) { auto negate = builder.AddInstruction( HloInstruction::CreateUnary(shape_, HloOpcode::kNegate, add)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); auto outer_fusion = computation->CreateFusionInstruction( {negate, add}, HloInstruction::FusionKind::kLoop); HloInstruction* fused_negate = outer_fusion->fused_expression_root(); @@ -203,9 +209,9 @@ TEST_F(DefuserTest, NestedFusionInstructions) { EXPECT_THAT(computation->root_instruction(), op::Fusion()); - EXPECT_EQ(2, FusionCount()); - EXPECT_TRUE(defuser_.Run(&module()).ValueOrDie()); - EXPECT_EQ(0, FusionCount()); + EXPECT_EQ(2, FusionCount(m.get())); + EXPECT_TRUE(defuser_.Run(m.get()).ValueOrDie()); + EXPECT_EQ(0, FusionCount(m.get())); EXPECT_THAT(computation->root_instruction(), op::Negate(op::Add())); } diff --git a/tensorflow/compiler/xla/service/despecializer.cc b/tensorflow/compiler/xla/service/despecializer.cc index b3549acfc291a54b2345b006310613c3a45a4b47..ed37099a5428075928ec98b134632867d58bbfe7 100644 --- a/tensorflow/compiler/xla/service/despecializer.cc +++ b/tensorflow/compiler/xla/service/despecializer.cc @@ -17,6 +17,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/bfloat16_normalization.h" #include "tensorflow/compiler/xla/service/defuser.h" +#include "tensorflow/compiler/xla/service/hlo_memory_scheduler.h" #include "tensorflow/compiler/xla/service/implicit_broadcast_remover.h" namespace xla { @@ -45,6 +46,7 @@ class ControlDepRemover : public HloModulePass { Despecializer::Despecializer() : pipeline_("despecializer") { // TODO(b/70588125): Also deal with window reversal in a fast way. + pipeline_.AddPass(); pipeline_.AddPass(); pipeline_.AddPass(); pipeline_.AddPass(); diff --git a/tensorflow/compiler/xla/service/device_memory_allocator.cc b/tensorflow/compiler/xla/service/device_memory_allocator.cc index edbcb25247421cdb50a845df1ec8b1851970efe3..e1e3b156fb34fd128864ed34c6d9d055294672bf 100644 --- a/tensorflow/compiler/xla/service/device_memory_allocator.cc +++ b/tensorflow/compiler/xla/service/device_memory_allocator.cc @@ -20,6 +20,7 @@ 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/core/lib/strings/numbers.h" namespace xla { @@ -39,6 +40,10 @@ StatusOr StreamExecutorMemoryAllocator::Allocate( "Failed to allocate request for %s (%uB) on device ordinal %d", tensorflow::strings::HumanReadableNumBytes(size), size, device_ordinal); } + VLOG(3) << absl::StreamFormat( + "Allocated %s (%uB) on device ordinal %d: %p", + tensorflow::strings::HumanReadableNumBytes(size), size, device_ordinal, + result.opaque()); return OwningDeviceMemory(result, device_ordinal, this); } @@ -47,6 +52,8 @@ Status StreamExecutorMemoryAllocator::Deallocate(int device_ordinal, if (!mem.is_null()) { TF_ASSIGN_OR_RETURN(se::StreamExecutor * stream_executor, GetStreamExecutor(device_ordinal)); + VLOG(3) << absl::StreamFormat("Freeing %p on device ordinal %d", + mem.opaque(), device_ordinal); stream_executor->Deallocate(&mem); } return Status::OK(); diff --git a/tensorflow/compiler/xla/service/dfs_hlo_visitor.h b/tensorflow/compiler/xla/service/dfs_hlo_visitor.h index 4159aa281fa2b66d310d7c135f123a5a3bb83270..2132468b9067ad4d5644d6cf3908a488a20ced05 100644 --- a/tensorflow/compiler/xla/service/dfs_hlo_visitor.h +++ b/tensorflow/compiler/xla/service/dfs_hlo_visitor.h @@ -105,9 +105,10 @@ class DfsHloVisitorBase { } virtual Status HandleConvolution(HloInstructionPtr hlo) = 0; virtual Status HandleFft(HloInstructionPtr fft) = 0; - virtual Status HandleCrossReplicaSum(HloInstructionPtr hlo) = 0; + virtual Status HandleAllReduce(HloInstructionPtr hlo) = 0; virtual Status HandleAllToAll(HloInstructionPtr hlo) = 0; virtual Status HandleCollectivePermute(HloInstructionPtr hlo) = 0; + virtual Status HandleGetDimensionSize(HloInstructionPtr hlo) = 0; virtual Status HandleCompare(HloInstructionPtr hlo) { return HandleElementwiseBinary(hlo); } @@ -250,6 +251,7 @@ class DfsHloVisitorBase { virtual Status HandleBatchNormGrad(HloInstructionPtr hlo) = 0; + virtual Status HandleAddDependency(HloInstructionPtr add_dependency) = 0; virtual Status HandleAfterAll(HloInstructionPtr token) = 0; // Invoked to inform the visitor that the traversal has completed, and that 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 4cd10ab06cd3b804406607212d3f3c316d6cff95..680dd256bb15bd3a9eaff7241174c1d2833002c6 100644 --- a/tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h +++ b/tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h @@ -91,7 +91,7 @@ class DfsHloVisitorWithDefaultBase Status HandleFft(HloInstructionPtr fft) override { return DefaultAction(fft); } - Status HandleCrossReplicaSum(HloInstructionPtr crs) override { + Status HandleAllReduce(HloInstructionPtr crs) override { return DefaultAction(crs); } Status HandleAllToAll(HloInstructionPtr hlo) override { @@ -203,6 +203,12 @@ class DfsHloVisitorWithDefaultBase Status HandleAfterAll(HloInstructionPtr token) override { return DefaultAction(token); } + Status HandleGetDimensionSize(HloInstructionPtr get_size) override { + return DefaultAction(get_size); + } + Status HandleAddDependency(HloInstructionPtr add_dependency) override { + return DefaultAction(add_dependency); + } // Invoked to inform the visitor that the traversal has completed, and that // the root was "root". 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..e8bc6d05716a2ef02e0280e86c7df4ac22fe78c4 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 new file mode 100644 index 0000000000000000000000000000000000000000..2b158d7a6ec510ce4cbc56bddc5cca71ac4f14f4 --- /dev/null +++ b/tensorflow/compiler/xla/service/dynamic_dimension_inference.cc @@ -0,0 +1,459 @@ +/* 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/dynamic_dimension_inference.h" +#include "tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.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" + +namespace xla { + +namespace { +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; +} +} // namespace + +class DynamicDimensionInferenceVisitor : public DfsHloVisitorWithDefault { + public: + explicit DynamicDimensionInferenceVisitor( + const DynamicParameterBinding& param_bindings, + DynamicDimensionInference* parent) + : param_bindings_(param_bindings), parent_(parent) {} + + Status DefaultAction(HloInstruction* hlo) override; + + static Status Run(HloComputation* computation, + const DynamicParameterBinding& param_bindings, + DynamicDimensionInference* parent) { + DynamicDimensionInferenceVisitor visitor(param_bindings, parent); + return computation->Accept(&visitor); + } + + Status HandleParameter(HloInstruction* hlo) override; + + Status HandleReduce(HloInstruction* hlo) override; + + Status HandleDot(HloInstruction* hlo) override; + + Status HandleTranspose(HloInstruction* hlo) override; + + Status HandleReshape(HloInstruction* hlo) override; + + Status HandlePad(HloInstruction* hlo) override; + + Status HandleBroadcast(HloInstruction* hlo) override; + + Status HandleGetDimensionSize(HloInstruction* hlo) override; + + Status HandleSelect(HloInstruction* hlo) override; + + Status HandleConvolution(HloInstruction* hlo) override; + + Status HandleReduceWindow(HloInstruction* hlo) override; + + Status HandleSelectAndScatter(HloInstruction* hlo) override; + + Status HandleGetTupleElement(HloInstruction* hlo) override; + + Status HandleElementwiseUnary(HloInstruction* hlo) override; + + Status HandleElementwiseBinary(HloInstruction* hlo) override; + + private: + using OperandDynamicDimensionFn = std::function; + + Status ForEachOperandDynamicDimension(HloInstruction* inst, + const OperandDynamicDimensionFn&); + + // Pass through a dynamic dimension from the input to the output with the same + // value and index in the shape. This is a helper function to handle trivial + // instructions like elementwise operations. + Status PassThroughDynamicDimension(HloInstruction*); + + // The dynamic parameter bindings of this computation. + const DynamicParameterBinding& param_bindings_; + + // A pointer to DynamicDimensionInference, used to update the dynamic mapping. + DynamicDimensionInference* parent_; +}; + +Status DynamicDimensionInferenceVisitor::DefaultAction(HloInstruction* hlo) { + return ForEachOperandDynamicDimension( + hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension, + int64 operand_index, HloInstruction* dynamic_size) { + return UnimplementedStrCat( + "Asked to propagate a dynamic dimension from hlo ", + operand->ToString(), "@", index.ToString(), "@", dimension, + " to hlo ", hlo->ToString(), ", which is not implemented."); + }); +} + +Status DynamicDimensionInferenceVisitor::HandleGetTupleElement( + HloInstruction* hlo) { + return ForEachOperandDynamicDimension( + hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension, + int64 operand_index, HloInstruction* dynamic_size) { + if (hlo->tuple_index() == index[0]) { + ShapeIndex new_index = + ShapeIndexView(index).ConsumeFront().ToShapeIndex(); + parent_->SetDynamicSize(hlo, new_index, dimension, dynamic_size); + } + return Status::OK(); + }); +} + +Status DynamicDimensionInferenceVisitor::HandleBroadcast(HloInstruction* hlo) { + return ForEachOperandDynamicDimension( + hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension, + int64 operand_index, HloInstruction* dynamic_size) { + int64 broadcast_dim = hlo->dimensions(dimension); + parent_->SetDynamicSize(hlo, index, broadcast_dim, dynamic_size); + return Status::OK(); + }); +} + +Status DynamicDimensionInferenceVisitor::HandlePad(HloInstruction* hlo) { + return ForEachOperandDynamicDimension( + hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension, + int64 operand_index, HloInstruction* dynamic_size) { + if (operand_index != 0) { + return Unimplemented( + "Dynamic dimension on padding value is not supported"); + } + const PaddingConfig_PaddingConfigDimension& padding_config = + hlo->padding_config().dimensions(dimension); + if (padding_config.interior_padding() == 0 && + padding_config.edge_padding_low() == 0 && + padding_config.edge_padding_high() == 0) { + parent_->SetDynamicSize(hlo, {}, dimension, dynamic_size); + return Status::OK(); + } else { + return Unimplemented( + "Dynamic dimension propagation on padding dimension is not " + "supported."); + } + }); +} + +Status DynamicDimensionInferenceVisitor::HandleReduce(HloInstruction* hlo) { + return ForEachOperandDynamicDimension( + hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension, + int64 operand_index, HloInstruction* dynamic_size) { + HloInstruction* reduce = hlo; + int64 operand_count = reduce->operand_count(); + CHECK_EQ(operand_count % 2, 0); + if (operand_index >= operand_count / 2) { + // Init values doesn't have dynamic size. + return Status::OK(); + } + if ((absl::c_count(reduce->dimensions(), dimension) != 0)) { + // Dimension is to be reduce, stop tracing. + return Status::OK(); + } + + // Find out the new dynamic dimension after reduce. + int64 dimensions_not_reduced_count = 0; + for (int i = 0; i < operand->shape().rank(); ++i) { + if (dimension == i) { + parent_->SetDynamicSize(reduce, {}, dimensions_not_reduced_count, + dynamic_size); + + return Status::OK(); + } + if (absl::c_count(reduce->dimensions(), i) == 0) { + dimensions_not_reduced_count++; + } + } + + return Status::OK(); + }); +} + +Status DynamicDimensionInferenceVisitor::HandleDot(HloInstruction* hlo) { + return ForEachOperandDynamicDimension( + hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension, + int64 operand_index, HloInstruction* dynamic_size) { + HloInstruction* dot = hlo; + const DotDimensionNumbers& dimension_numbers = + dot->dot_dimension_numbers(); + // A map from the operand dimensions to result dimension. + absl::flat_hash_map result_dim_mapping; + int64 current_result_dims = 0; + std::unordered_set batch_dims( + dimension_numbers.rhs_batch_dimensions().begin(), + dimension_numbers.rhs_batch_dimensions().end()); + + for (int64 i : dimension_numbers.rhs_batch_dimensions()) { + result_dim_mapping[i] = current_result_dims++; + } + + for (int64 i = 0; i < dot->operand(0)->shape().rank(); i++) { + if (!absl::c_linear_search( + dimension_numbers.lhs_contracting_dimensions(), i)) { + if (operand_index == 0) { + result_dim_mapping[i] = current_result_dims; + } + current_result_dims++; + } + } + + for (int64 i = 0; i < dot->operand(1)->shape().rank(); i++) { + if (!absl::c_linear_search( + dimension_numbers.rhs_contracting_dimensions(), i) && + !absl::c_linear_search(dimension_numbers.rhs_batch_dimensions(), + i)) { + if (operand_index == 1) { + result_dim_mapping[i] = current_result_dims; + } + current_result_dims++; + } + } + + // Check if the operand dim is in the result shape. If so, add another + // work item to trace that dimension. + auto iter = result_dim_mapping.find(dimension); + if (iter != result_dim_mapping.end()) { + parent_->SetDynamicSize(dot, {}, iter->second, dynamic_size); + } + + return Status::OK(); + }); +} + +Status DynamicDimensionInferenceVisitor::HandleTranspose(HloInstruction* hlo) { + return ForEachOperandDynamicDimension( + hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension, + int64 operand_index, HloInstruction* dynamic_size) { + parent_->SetDynamicSize(hlo, {}, hlo->dimensions()[dimension], + dynamic_size); + return Status::OK(); + }); +} + +Status DynamicDimensionInferenceVisitor::HandleConvolution( + HloInstruction* hlo) { + return ForEachOperandDynamicDimension( + hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension, + int64 operand_index, HloInstruction* dynamic_size) { + HloInstruction* conv = hlo; + const ConvolutionDimensionNumbers& dimension_numbers = + conv->convolution_dimension_numbers(); + + if (operand_index == 0) { + if (dimension == dimension_numbers.input_batch_dimension()) { + parent_->SetDynamicSize(conv, {}, + dimension_numbers.output_batch_dimension(), + dynamic_size); + return Status::OK(); + } + + if (dimension == dimension_numbers.input_feature_dimension()) { + return Status::OK(); + } + } else { + if (dimension == dimension_numbers.kernel_input_feature_dimension()) { + return Status::OK(); + } + } + + return Unimplemented("Dynamic Spatial Convolution is not supported: %s", + conv->ToString()); + }); +} + +Status DynamicDimensionInferenceVisitor::HandleGetDimensionSize( + HloInstruction*) { + // Dynamic dimension doesn't propagate through GetDimensionSize: + // + // Input: F32[x, y, z] + // | + // GetDimensionSize(1): U32[] + // + // The returned value is a scalar, which doesn't have any dynamic dimension in + // the shape (although the value contains the real size of the dynamic + // dimension of the input). + return Status::OK(); +} + +Status DynamicDimensionInferenceVisitor::PassThroughDynamicDimension( + HloInstruction* hlo) { + return ForEachOperandDynamicDimension( + hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension, + int64 operand_index, HloInstruction* dynamic_size) { + parent_->SetDynamicSize(hlo, index, dimension, dynamic_size); + return Status::OK(); + }); +} + +Status DynamicDimensionInferenceVisitor::HandleElementwiseUnary( + HloInstruction* hlo) { + return PassThroughDynamicDimension(hlo); +} + +Status DynamicDimensionInferenceVisitor::HandleSelect(HloInstruction* hlo) { + return PassThroughDynamicDimension(hlo); +} + +Status DynamicDimensionInferenceVisitor::HandleElementwiseBinary( + HloInstruction* hlo) { + return PassThroughDynamicDimension(hlo); +} + +Status DynamicDimensionInferenceVisitor::HandleReshape(HloInstruction* hlo) { + return ForEachOperandDynamicDimension( + hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension, + int64 operand_index, HloInstruction* dynamic_size) { + HloInstruction* reshape = hlo; + std::vector> unmodified_dims = + ShapeUtil::DimensionsUnmodifiedByReshape(operand->shape(), + reshape->shape()); + for (auto& unmodified : unmodified_dims) { + if (unmodified.first == dimension) { + parent_->SetDynamicSize(reshape, {}, unmodified.second, + dynamic_size); + return Status::OK(); + } + } + return Unimplemented( + "Dynamic Reshape on modified dimensions is yet not supported: %s", + reshape->ToString()); + }); +} + +Status DynamicDimensionInferenceVisitor::HandleReduceWindow( + HloInstruction* hlo) { + return ForEachOperandDynamicDimension( + hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension, + int64 operand_index, HloInstruction* dynamic_size) { + HloInstruction* reduce_window = hlo; + const WindowDimension& window_dimension = + reduce_window->window().dimensions(dimension); + + if (!IsTrivialWindowDimension(window_dimension)) { + return Unimplemented( + "Dynamic Spatial reduce window is not supported: %s", + reduce_window->ToString()); + } + + parent_->SetDynamicSize(reduce_window, {}, dimension, dynamic_size); + + return Status::OK(); + }); +} + +Status DynamicDimensionInferenceVisitor::HandleSelectAndScatter( + HloInstruction* hlo) { + return ForEachOperandDynamicDimension( + hlo, [&](HloInstruction* operand, ShapeIndex index, int64 dimension, + int64 operand_index, HloInstruction* dynamic_size) { + HloInstruction* select_and_scatter = hlo; + const WindowDimension& window_dimension = + select_and_scatter->window().dimensions(dimension); + + if (!IsTrivialWindowDimension(window_dimension)) { + return Unimplemented( + "Dynamic Spatial select and scatter is not supported: %s", + select_and_scatter->ToString()); + } + + parent_->SetDynamicSize(select_and_scatter, {}, dimension, + dynamic_size); + + return Status::OK(); + }); +} + +Status DynamicDimensionInferenceVisitor::HandleParameter(HloInstruction* hlo) { + return param_bindings_.ForEachBinding( + [&](const DynamicParameterBinding::DynamicParameter& dynamic_parameter, + const DynamicParameterBinding::DynamicDimension& dynamic_dimension) { + if (dynamic_dimension.parameter_num != hlo->parameter_number()) { + return Status::OK(); + } + HloComputation* computation = hlo->parent(); + HloInstruction* target_parameter = + computation->parameter_instruction(dynamic_dimension.parameter_num); + + HloInstruction* dynamic_size = + computation->parameter_instruction(dynamic_parameter.parameter_num); + for (int64 i : dynamic_parameter.parameter_index) { + dynamic_size = + computation->AddInstruction(HloInstruction::CreateGetTupleElement( + ShapeUtil::GetSubshape(dynamic_size->shape(), {i}), + dynamic_size, i)); + } + + parent_->SetDynamicSize(target_parameter, + dynamic_dimension.parameter_index, + dynamic_dimension.dimension, dynamic_size); + return Status::OK(); + }); +} + +Status DynamicDimensionInferenceVisitor::ForEachOperandDynamicDimension( + HloInstruction* inst, const OperandDynamicDimensionFn& fn) { + for (int64 operand_index = 0; operand_index < inst->operand_count(); + ++operand_index) { + auto iter = + parent_->per_hlo_dynamic_dimensions_.find(inst->operand(operand_index)); + if (iter != parent_->per_hlo_dynamic_dimensions_.end()) { + for (auto& dynamic_dimension : iter->second) { + HloInstruction* dynamic_size = parent_->GetDynamicSize( + dynamic_dimension.inst, dynamic_dimension.index, + dynamic_dimension.dim); + TF_RETURN_IF_ERROR(fn(dynamic_dimension.inst, dynamic_dimension.index, + dynamic_dimension.dim, operand_index, + dynamic_size)); + } + } + } + return Status::OK(); +} + +/* static */ +StatusOr DynamicDimensionInference::Run( + HloModule* module) { + VLOG(2) << "Param Config " << module->dynamic_parameter_binding().ToString(); + DynamicDimensionInference inference(module); + TF_RETURN_IF_ERROR(inference.AnalyzeDynamicDimensions()); + return inference; +} + +DynamicDimensionInference::DynamicDimensionInference(HloModule* module) + : module_(module) {} + +Status DynamicDimensionInference::AnalyzeDynamicDimensions() { + return DynamicDimensionInferenceVisitor::Run( + module_->entry_computation(), module_->dynamic_parameter_binding(), this); +} + +HloInstruction* DynamicDimensionInference::GetDynamicSize( + HloInstruction* inst, const ShapeIndex& index, int64 dim) const { + auto iter = dynamic_mapping_.find(DynamicDimension{inst, index, dim}); + if (iter != dynamic_mapping_.end()) { + return iter->second; + } + return nullptr; +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/dynamic_dimension_inference.h b/tensorflow/compiler/xla/service/dynamic_dimension_inference.h new file mode 100644 index 0000000000000000000000000000000000000000..164d15bf111a92e3da957f609b54ee0662ef18b1 --- /dev/null +++ b/tensorflow/compiler/xla/service/dynamic_dimension_inference.h @@ -0,0 +1,112 @@ +/* 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_DYNAMIC_DIMENSION_INFERENCE_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_DYNAMIC_DIMENSION_INFERENCE_H_ + +#include +#include +#include + +#include "absl/container/flat_hash_map.h" +#include "absl/types/span.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/service/hlo_module.h" +#include "tensorflow/compiler/xla/status.h" +#include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/compiler/xla/types.h" +#include "tensorflow/core/platform/macros.h" + +namespace xla { + +// DynamicDimensionInference analyzes each HLO instruction in a graph and +// inferences which dimensions are dynamic and which scalar instructions +// represent the runtime real size of those dynamic dimensions. +class DynamicDimensionInference { + public: + static StatusOr Run(HloModule* module); + + string ToString() const; + + // If the dimension `dim` of instruction `inst` at `index` has a dynamic size, + // returns a scalar HloInstruction that represents the runtime size of that + // dimension. Otherwise returns nullptr. + HloInstruction* GetDynamicSize(HloInstruction* inst, const ShapeIndex& index, + int64 dim) const; + + friend class DynamicDimensionInferenceVisitor; + + private: + explicit DynamicDimensionInference(HloModule* module); + + // DynamicDimension is used as a key in the dynamic key-value mapping. It + // unambiguously represents a dynamic dimension of a instruction at a given + // index. + struct DynamicDimension { + // HloInstruction that holds the dimension. + HloInstruction* inst; + // Subshape of the instruction that holds the dimension. + ShapeIndex index; + // The dimension number of the dynamic dimension at given index of a given + // instruction. + int64 dim; + + // Artifacts needed to make this struct able to be used as a `key` in absl + // maps. "friend" keywords are added so these functions can be found through + // ADL. + template + friend H AbslHashValue(H h, const DynamicDimension& m) { + return H::combine(std::move(h), m.inst, m.index, m.dim); + } + + friend bool operator==(const DynamicDimension& lhs, + const DynamicDimension& rhs) { + return lhs.inst == rhs.inst && lhs.index == rhs.index && + lhs.dim == rhs.dim; + } + }; + + // Update the dynamic mapping so that we know dimension `dim` of instruction + // `inst` at `index` has a dynamic size, and its runtime size is represented + // by a scalar instruction `size`. + void SetDynamicSize(HloInstruction* inst, const ShapeIndex& index, int64 dim, + HloInstruction* size) { + dynamic_mapping_.try_emplace(DynamicDimension{inst, index, dim}, size); + auto iter = per_hlo_dynamic_dimensions_.try_emplace(inst); + iter.first->second.emplace(DynamicDimension{inst, index, dim}); + } + + // AnalyzeDynamicDimensions starts the analysis of the dynamic dimensions in + // module_. + Status AnalyzeDynamicDimensions(); + + // HloModule being analyzed. + HloModule* module_; + + // dynamic_mapping_ holds the result of the analysis. It maps a dynamic + // dimension to a scalar HloInstruction that represents the real dynamic size + // of the dynamic dimension. + using DynamicMapping = absl::flat_hash_map; + DynamicMapping dynamic_mapping_; + + using PerHloDynamicDimensions = + absl::flat_hash_map>; + PerHloDynamicDimensions per_hlo_dynamic_dimensions_; +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_DYNAMIC_DIMENSION_INFERENCE_H_ diff --git a/tensorflow/compiler/xla/service/dynamic_dimension_inference_test.cc b/tensorflow/compiler/xla/service/dynamic_dimension_inference_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..b42e67b4bbcf731d89dd8af9e46b405235a92d8a --- /dev/null +++ b/tensorflow/compiler/xla/service/dynamic_dimension_inference_test.cc @@ -0,0 +1,551 @@ +/* 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/dynamic_dimension_inference.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 DynamicDimensionInferenceTest : public HloTestBase { + protected: + DynamicDimensionInferenceTest() : HloTestBase() { + module_ = CreateNewVerifiedModule(); + } + + Status RunInference() { + hlo_graph_dumper::MaybeDumpHloModule(*module_, "Before alias analysis"); + TF_ASSIGN_OR_RETURN(DynamicDimensionInference inference, + DynamicDimensionInference::Run(module_.get())); + + inference_ = absl::make_unique(inference); + return Status::OK(); + } + + HloComputation* GetAdd() { + 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()); + } + + 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, {}); +}; + +TEST_F(DynamicDimensionInferenceTest, ParamTest) { + auto builder = HloComputation::Builder(TestName()); + auto input_shape = ShapeUtil::MakeShape(F32, {1, 2, 2}); + + auto param = builder.AddInstruction( + HloInstruction::CreateParameter(0, input_shape, "param")); + auto param2 = builder.AddInstruction( + HloInstruction::CreateParameter(1, scalar_shape_, "param")); + + 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(RunInference()); + EXPECT_EQ(inference_->GetDynamicSize(param, {}, 1), param2); + EXPECT_EQ(inference_->GetDynamicSize(param, {}, 0), nullptr); + EXPECT_EQ(inference_->GetDynamicSize(param2, {}, 0), nullptr); +} + +TEST_F(DynamicDimensionInferenceTest, ParamTestTuple) { + auto builder = HloComputation::Builder(TestName()); + auto input_shape = ShapeUtil::MakeShape(F32, {1, 2, 2}); + + auto param = builder.AddInstruction(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeTupleShape({input_shape, scalar_shape_}), "param")); + + module_->AddEntryComputation(builder.Build()); + // Set up dynamic parameter binding. + TF_CHECK_OK(module_->dynamic_parameter_binding().Bind( + DynamicParameterBinding::DynamicParameter{0, {1}}, + DynamicParameterBinding::DynamicDimension{0, {0}, 1})); + + TF_ASSERT_OK(RunInference()); + EXPECT_THAT(inference_->GetDynamicSize(param, {0}, 1), + op::GetTupleElement(param, 1)); + + EXPECT_EQ(inference_->GetDynamicSize(param, {0}, 0), nullptr); +} + +TEST_F(DynamicDimensionInferenceTest, GetTupleElement) { + // When data flows through GTE, the dynamic dimension size keeps the + // same, and the index has its front popped. + auto builder = HloComputation::Builder(TestName()); + auto input_shape = ShapeUtil::MakeShape(F32, {1, 2, 2}); + + auto param = builder.AddInstruction(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeTupleShape({input_shape, scalar_shape_}), "param")); + + auto gte = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(input_shape, param, 0)); + + module_->AddEntryComputation(builder.Build()); + // Set up dynamic parameter binding. + TF_CHECK_OK(module_->dynamic_parameter_binding().Bind( + DynamicParameterBinding::DynamicParameter{0, {1}}, + DynamicParameterBinding::DynamicDimension{0, {0}, 1})); + + TF_ASSERT_OK(RunInference()); + EXPECT_THAT(inference_->GetDynamicSize(param, {0}, 1), + op::GetTupleElement(param, 1)); + + EXPECT_THAT(inference_->GetDynamicSize(gte, {}, 1), + op::GetTupleElement(param, 1)); + + EXPECT_EQ(inference_->GetDynamicSize(param, {0}, 0), nullptr); +} + +TEST_F(DynamicDimensionInferenceTest, ElementwiseTest) { + // When data flows through elementwise, the dynamic dimension size keeps the + // same. + auto builder = HloComputation::Builder(TestName()); + auto input_shape = ShapeUtil::MakeShape(F32, {1, 2, 2}); + + auto data_param = builder.AddInstruction( + HloInstruction::CreateParameter(0, input_shape, "data_param")); + auto size_param = builder.AddInstruction( + HloInstruction::CreateParameter(1, scalar_shape_, "size_param")); + + auto* negate = builder.AddInstruction( + HloInstruction::CreateUnary(input_shape, HloOpcode::kNegate, data_param)); + + 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(RunInference()); + EXPECT_EQ(inference_->GetDynamicSize(negate, {}, 1), size_param); +} + +TEST_F(DynamicDimensionInferenceTest, ReduceTestI) { + 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")); + auto size_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}, GetAdd())); + + 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(RunInference()); + EXPECT_EQ(inference_->GetDynamicSize(reduce, {}, 0), size_param); +} + +TEST_F(DynamicDimensionInferenceTest, ReduceTestII) { + // Same as ReduceTestI, but only reduce one dimension. + auto builder = HloComputation::Builder(TestName()); + auto input_shape = ShapeUtil::MakeShape(F32, {1, 2, 2}); + auto reduce_shape = ShapeUtil::MakeShape(F32, {1, 2}); + + auto data_param = builder.AddInstruction( + HloInstruction::CreateParameter(0, input_shape, "data_param")); + auto size_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, {1}, GetAdd())); + + module_->AddEntryComputation(builder.Build()); + + // Set up dynamic parameter binding. + TF_CHECK_OK(module_->dynamic_parameter_binding().Bind( + DynamicParameterBinding::DynamicParameter{1, {}}, + DynamicParameterBinding::DynamicDimension{0, {}, 2})); + + TF_ASSERT_OK(RunInference()); + EXPECT_EQ(inference_->GetDynamicSize(reduce, {}, 1), size_param); + EXPECT_EQ(inference_->GetDynamicSize(reduce, {}, 0), nullptr); +} + +TEST_F(DynamicDimensionInferenceTest, DotTest) { + 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 xz_shape = ShapeUtil::MakeShape(F32, {xdim, zdim}); + + 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")); + auto* size_param = builder.AddInstruction(HloInstruction::CreateParameter( + /*parameter_number=*/2, scalar_shape_, "size_param")); + + DotDimensionNumbers dot_dnums; + dot_dnums.add_lhs_contracting_dimensions(1); + dot_dnums.add_rhs_contracting_dimensions(0); + auto dot = builder.AddInstruction( + HloInstruction::CreateDot(xz_shape, a_param, b_param, dot_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_CHECK_OK(module_->dynamic_parameter_binding().Bind( + DynamicParameterBinding::DynamicParameter{2, {}}, + DynamicParameterBinding::DynamicDimension{1, {}, 0})); + + TF_ASSERT_OK(RunInference()); + EXPECT_EQ(inference_->GetDynamicSize(dot, {}, 0), size_param); + EXPECT_EQ(inference_->GetDynamicSize(dot, {}, 1), nullptr); +} + +TEST_F(DynamicDimensionInferenceTest, 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")); + auto* size_param = 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(RunInference()); + EXPECT_EQ(inference_->GetDynamicSize(conv, {}, 1), size_param); + EXPECT_EQ(inference_->GetDynamicSize(conv, {}, 0), nullptr); +} + +TEST_F(DynamicDimensionInferenceTest, TransposeTest) { + // Test the ability to trace unmodified dimensions + auto builder = HloComputation::Builder(TestName()); + auto input_shape = ShapeUtil::MakeShape(F32, {1, 2, 3}); + auto output_shape = ShapeUtil::MakeShape(F32, {3, 2, 1}); + + auto* a_param = builder.AddInstruction(HloInstruction::CreateParameter( + /*parameter_number=*/0, input_shape, "A")); + auto* size_param_1 = builder.AddInstruction(HloInstruction::CreateParameter( + /*parameter_number=*/1, scalar_shape_, "size_param")); + auto* size_param_2 = builder.AddInstruction(HloInstruction::CreateParameter( + /*parameter_number=*/2, scalar_shape_, "size_param")); + auto* size_param_3 = builder.AddInstruction(HloInstruction::CreateParameter( + /*parameter_number=*/3, scalar_shape_, "size_param")); + + auto* transpose = builder.AddInstruction( + HloInstruction::CreateTranspose(output_shape, a_param, {2, 1, 0})); + + 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{2, {}}, + DynamicParameterBinding::DynamicDimension{0, {}, 1})); + + TF_CHECK_OK(module_->dynamic_parameter_binding().Bind( + DynamicParameterBinding::DynamicParameter{3, {}}, + DynamicParameterBinding::DynamicDimension{0, {}, 2})); + + TF_ASSERT_OK(RunInference()); + EXPECT_EQ(inference_->GetDynamicSize(transpose, {}, 0), size_param_3); + EXPECT_EQ(inference_->GetDynamicSize(transpose, {}, 1), size_param_2); + EXPECT_EQ(inference_->GetDynamicSize(transpose, {}, 2), size_param_1); +} + +TEST_F(DynamicDimensionInferenceTest, ReshapeTest) { + // Test the ability to trace unmodified reshape dimensions. + auto builder = HloComputation::Builder(TestName()); + auto input_shape = ShapeUtil::MakeShape(F32, {2, 3, 4, 5, 6}); + auto output_shape = ShapeUtil::MakeShape(F32, {6, 4, 1, 5, 2, 3}); + + auto* a_param = builder.AddInstruction(HloInstruction::CreateParameter( + /*parameter_number=*/0, input_shape, "A")); + auto* size_param = builder.AddInstruction(HloInstruction::CreateParameter( + /*parameter_number=*/1, scalar_shape_, "size_param")); + + auto* reshape = builder.AddInstruction( + HloInstruction::CreateReshape(output_shape, a_param)); + + module_->AddEntryComputation(builder.Build()); + + TF_CHECK_OK(module_->dynamic_parameter_binding().Bind( + DynamicParameterBinding::DynamicParameter{1, {}}, + DynamicParameterBinding::DynamicDimension{0, {}, 2})); + + TF_CHECK_OK(module_->dynamic_parameter_binding().Bind( + DynamicParameterBinding::DynamicParameter{1, {}}, + DynamicParameterBinding::DynamicDimension{0, {}, 3})); + + TF_ASSERT_OK(RunInference()); + EXPECT_EQ(inference_->GetDynamicSize(reshape, {}, 0), nullptr); + EXPECT_EQ(inference_->GetDynamicSize(reshape, {}, 1), size_param); + EXPECT_EQ(inference_->GetDynamicSize(reshape, {}, 2), nullptr); + EXPECT_EQ(inference_->GetDynamicSize(reshape, {}, 3), size_param); + EXPECT_EQ(inference_->GetDynamicSize(reshape, {}, 4), nullptr); + EXPECT_EQ(inference_->GetDynamicSize(reshape, {}, 5), nullptr); +} + +TEST_F(DynamicDimensionInferenceTest, ReshapeTestUnimplemented) { + // Test the ability to trace unmodified reshape dimensions. + auto builder = HloComputation::Builder(TestName()); + auto input_shape = ShapeUtil::MakeShape(F32, {2, 3, 4, 5, 6}); + auto output_shape = ShapeUtil::MakeShape(F32, {6, 4, 1, 5, 2, 3}); + + auto* a_param = builder.AddInstruction(HloInstruction::CreateParameter( + /*parameter_number=*/0, input_shape, "A")); + + builder.AddInstruction(HloInstruction::CreateParameter( + /*parameter_number=*/1, scalar_shape_, "size_param")); + + builder.AddInstruction(HloInstruction::CreateReshape(output_shape, a_param)); + + module_->AddEntryComputation(builder.Build()); + + TF_CHECK_OK(module_->dynamic_parameter_binding().Bind( + DynamicParameterBinding::DynamicParameter{1, {}}, + DynamicParameterBinding::DynamicDimension{0, {}, 1})); + + Status status = RunInference(); + EXPECT_EQ(status.code(), tensorflow::error::UNIMPLEMENTED); +} + +TEST_F(DynamicDimensionInferenceTest, BroadcastTest) { + // Test the ability to trace broadcast dimension. + auto builder = HloComputation::Builder(TestName()); + auto input_shape = ShapeUtil::MakeShape(F32, {2}); + auto output_shape = ShapeUtil::MakeShape(F32, {3, 2, 4}); + + auto* a_param = builder.AddInstruction(HloInstruction::CreateParameter( + /*parameter_number=*/0, input_shape, "A")); + auto* size_param = builder.AddInstruction(HloInstruction::CreateParameter( + /*parameter_number=*/1, scalar_shape_, "size_param")); + + auto* broadcast = builder.AddInstruction( + HloInstruction::CreateBroadcast(output_shape, a_param, {1})); + + module_->AddEntryComputation(builder.Build()); + + TF_CHECK_OK(module_->dynamic_parameter_binding().Bind( + DynamicParameterBinding::DynamicParameter{1, {}}, + DynamicParameterBinding::DynamicDimension{0, {}, 0})); + + TF_ASSERT_OK(RunInference()); + EXPECT_EQ(inference_->GetDynamicSize(broadcast, {}, 0), nullptr); + EXPECT_EQ(inference_->GetDynamicSize(broadcast, {}, 1), size_param); + EXPECT_EQ(inference_->GetDynamicSize(broadcast, {}, 2), nullptr); +} + +TEST_F(DynamicDimensionInferenceTest, ReduceWindowBatchTest) { + // Test the ability to trace reduce window 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}); + + Window window; + // First dimension is unchanged. + WindowDimension* batch_dim = window.add_dimensions(); + batch_dim->set_size(1); + batch_dim->set_stride(1); + batch_dim->set_padding_low(0); + batch_dim->set_padding_high(0); + batch_dim->set_window_dilation(1); + batch_dim->set_base_dilation(1); + + // Second and third dimension are reduced. + for (int64 i = 0; i < 2; ++i) { + WindowDimension* dim = window.add_dimensions(); + dim->set_size(2); + dim->set_stride(2); + dim->set_padding_low(0); + dim->set_padding_high(0); + dim->set_window_dilation(1); + dim->set_base_dilation(1); + } + + auto* a_param = builder.AddInstruction(HloInstruction::CreateParameter( + /*parameter_number=*/0, input_shape, "A")); + auto* size_param = builder.AddInstruction(HloInstruction::CreateParameter( + /*parameter_number=*/1, scalar_shape_, "size_param")); + + auto init = builder.AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(0.0))); + + auto* reduce_window = + builder.AddInstruction(HloInstruction::CreateReduceWindow( + output_shape, a_param, init, window, GetAdd())); + + module_->AddEntryComputation(builder.Build()); + + TF_CHECK_OK(module_->dynamic_parameter_binding().Bind( + DynamicParameterBinding::DynamicParameter{1, {}}, + DynamicParameterBinding::DynamicDimension{0, {}, 0})); + + TF_ASSERT_OK(RunInference()); + EXPECT_EQ(inference_->GetDynamicSize(reduce_window, {}, 0), size_param); +} + +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 source_shape = ShapeUtil::MakeShape(F32, {2, 2, 2}); + + Window window; + // First dimension is unchanged. + WindowDimension* batch_dim = window.add_dimensions(); + batch_dim->set_size(1); + batch_dim->set_stride(1); + batch_dim->set_padding_low(0); + batch_dim->set_padding_high(0); + batch_dim->set_window_dilation(1); + batch_dim->set_base_dilation(1); + + // Second and third dimension are reduced. + for (int64 i = 0; i < 2; ++i) { + WindowDimension* dim = window.add_dimensions(); + dim->set_size(2); + dim->set_stride(2); + dim->set_padding_low(0); + dim->set_padding_high(0); + dim->set_window_dilation(1); + dim->set_base_dilation(1); + } + + auto* a_param = builder.AddInstruction(HloInstruction::CreateParameter( + /*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* 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(sns, {}, 0), size_param); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/service/dynamic_index_splitter.cc b/tensorflow/compiler/xla/service/dynamic_index_splitter.cc new file mode 100644 index 0000000000000000000000000000000000000000..e34adfd2d2bbb7214cfa2da28291b133538845e5 --- /dev/null +++ b/tensorflow/compiler/xla/service/dynamic_index_splitter.cc @@ -0,0 +1,99 @@ +/* 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/dynamic_index_splitter.h" + +#include + +#include "absl/container/flat_hash_map.h" +#include "absl/container/flat_hash_set.h" +#include "absl/container/inlined_vector.h" +#include "tensorflow/compiler/xla/service/hlo_casting_utils.h" +#include "tensorflow/compiler/xla/service/hlo_computation.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" +#include "tensorflow/compiler/xla/shape_util.h" + +namespace xla { + +StatusOr DynamicIndexSplitter::Run(HloModule* module) { + bool changed = false; + + std::vector computations = + module->MakeNonfusionComputations(); + for (HloComputation* computation : computations) { + for (HloInstruction* dynamic_op : computation->MakeInstructionPostOrder()) { + switch (dynamic_op->opcode()) { + case HloOpcode::kDynamicSlice: + case HloOpcode::kDynamicUpdateSlice: + break; + default: + continue; + } + auto parent = dynamic_op->parent(); + bool is_update = dynamic_op->opcode() == HloOpcode::kDynamicUpdateSlice; + int64 num_indices = dynamic_op->operand(0)->shape().rank(); + + if (num_indices == 0) { + // 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( + dynamic_op, dynamic_op->mutable_operand(1))); + } else { + TF_CHECK_OK(parent->ReplaceInstruction( + dynamic_op, dynamic_op->mutable_operand(0))); + } + 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) { + auto slice = parent->AddInstruction(HloInstruction::CreateSlice( + ShapeUtil::MakeShape(index_element_type, {1}), index_operand, {dim}, + {dim + 1}, {1})); + auto bitcast = parent->AddInstruction(HloInstruction::CreateReshape( + ShapeUtil::MakeShape(index_element_type, {}), slice)); + index_array.push_back(bitcast); + } + auto new_dynamic_op = + is_update + ? HloInstruction::CreateDynamicUpdateSlice( + dynamic_op->shape(), dynamic_op->mutable_operand(0), + dynamic_op->mutable_operand(1), absl::MakeSpan(index_array)) + : HloInstruction::CreateDynamicSlice( + dynamic_op->shape(), dynamic_op->mutable_operand(0), + absl::MakeSpan(index_array), + dynamic_op->dynamic_slice_sizes()); + TF_CHECK_OK(parent->ReplaceWithNewInstruction(dynamic_op, + std::move(new_dynamic_op))); + changed = true; + } + } + return changed; +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/dynamic_index_splitter.h b/tensorflow/compiler/xla/service/dynamic_index_splitter.h new file mode 100644 index 0000000000000000000000000000000000000000..3c12e3a4af287ad2272a08ba54cd99c2cad9d451 --- /dev/null +++ b/tensorflow/compiler/xla/service/dynamic_index_splitter.h @@ -0,0 +1,37 @@ +/* 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_DYNAMIC_INDEX_SPLITTER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_DYNAMIC_INDEX_SPLITTER_H_ + +#include "absl/strings/string_view.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 { + +// Convert R1 index operands to DynamicSlice and DynamicUpdateSlice ops into +// separate scalars. +class DynamicIndexSplitter : public HloModulePass { + public: + DynamicIndexSplitter() = default; + absl::string_view name() const override { return "dynamic-index-splitter"; } + StatusOr Run(HloModule* module) override; +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_DYNAMIC_INDEX_SPLITTER_H_ diff --git a/tensorflow/compiler/xla/service/dynamic_index_splitter_test.cc b/tensorflow/compiler/xla/service/dynamic_index_splitter_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..98029d1faff7d669730f6b66e38fcefece70f0eb --- /dev/null +++ b/tensorflow/compiler/xla/service/dynamic_index_splitter_test.cc @@ -0,0 +1,134 @@ +/* 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/dynamic_index_splitter.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_opcode.h" +#include "tensorflow/compiler/xla/test.h" +#include "tensorflow/compiler/xla/test_helpers.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" + +namespace xla { +namespace { + +namespace op = xla::testing::opcode_matchers; +class DynamicIndexSplitterTest : public HloTestBase {}; + +TEST_F(DynamicIndexSplitterTest, DynamicSlice) { + const char* const kDynamicSlice = R"( + HloModule DynamicSlice_module + + ENTRY entry (operand: s32[4,5,6], indices: s32[3]) -> s32[1,1,1] { + operand = s32[4,5,6] parameter(0) + indices = s32[3] parameter(1) + ROOT dynamic-slice = s32[1,1,1] dynamic-slice(operand, indices), dynamic_slice_sizes={1,1,1} + } + )"; + + HloModuleConfig config; + DebugOptions debug_options = config.debug_options(); + debug_options.set_xla_allow_scalar_index_dynamic_ops(true); + config.set_debug_options(debug_options); + + TF_ASSERT_OK_AND_ASSIGN(auto module, ParseHloString(kDynamicSlice, config)); + TF_ASSERT_OK_AND_ASSIGN(bool changed, + DynamicIndexSplitter().Run(module.get())); + EXPECT_TRUE(changed); + ASSERT_THAT(module->entry_computation()->root_instruction(), + op::DynamicSlice(op::Parameter(0), + op::Reshape(op::Slice(op::Parameter(1))), + op::Reshape(op::Slice(op::Parameter(1))), + op::Reshape(op::Slice(op::Parameter(1))))); + + for (int i = 0; i < 3; ++i) { + const HloInstruction* slice = module->entry_computation() + ->root_instruction() + ->operand(i + 1) + ->operand(0); + EXPECT_EQ(slice->slice_starts(0), i); + EXPECT_EQ(slice->slice_limits(0), i + 1); + } +} + +TEST_F(DynamicIndexSplitterTest, DynamicUpdateSlice) { + const char* const kDynamicUpdateSlice = R"( + HloModule DynamicUpdatedSlice_module + + ENTRY entry (operand: s32[4,5,6], indices: s32[3], update: s32[1,1,1]) -> s32[4,5,6] { + operand = s32[4,5,6] parameter(0) + indices = s32[3] parameter(1) + update = s32[1,1,1] parameter(2) + ROOT dynamic-update-slice = s32[4,5,6] dynamic-update-slice(operand, update, indices) + } + )"; + + HloModuleConfig config; + DebugOptions debug_options = config.debug_options(); + debug_options.set_xla_allow_scalar_index_dynamic_ops(true); + config.set_debug_options(debug_options); + + TF_ASSERT_OK_AND_ASSIGN(auto module, + ParseHloString(kDynamicUpdateSlice, config)); + TF_ASSERT_OK_AND_ASSIGN(bool changed, + DynamicIndexSplitter().Run(module.get())); + EXPECT_TRUE(changed); + ASSERT_THAT(module->entry_computation()->root_instruction(), + op::DynamicUpdateSlice(op::Parameter(0), op::Parameter(2), + op::Reshape(op::Slice(op::Parameter(1))), + op::Reshape(op::Slice(op::Parameter(1))), + op::Reshape(op::Slice(op::Parameter(1))))); + + for (int i = 0; i < 3; ++i) { + const HloInstruction* slice = module->entry_computation() + ->root_instruction() + ->operand(i + 2) + ->operand(0); + EXPECT_EQ(slice->slice_starts(0), i); + EXPECT_EQ(slice->slice_limits(0), i + 1); + } +} + +TEST_F(DynamicIndexSplitterTest, AlreadyScalar) { + const char* const kDynamicSlice = R"( + HloModule DynamicSlice_module + + ENTRY entry (operand: s32[4,5,6], index.0: s32[], index.1: s32[], index.2: s32[]) -> s32[1,1,1] { + operand = s32[4,5,6] parameter(0) + index.0 = s32[] parameter(1) + index.1 = s32[] parameter(2) + index.2 = s32[] parameter(3) + ROOT dynamic-slice = s32[1,1,1] dynamic-slice(operand, index.0, index.1, index.2), dynamic_slice_sizes={1,1,1} + } + )"; + + HloModuleConfig config; + DebugOptions debug_options = config.debug_options(); + debug_options.set_xla_allow_scalar_index_dynamic_ops(true); + config.set_debug_options(debug_options); + + TF_ASSERT_OK_AND_ASSIGN(auto module, ParseHloString(kDynamicSlice, config)); + TF_ASSERT_OK_AND_ASSIGN(bool changed, + DynamicIndexSplitter().Run(module.get())); + EXPECT_FALSE(changed); + EXPECT_THAT(module->entry_computation()->root_instruction(), + op::DynamicSlice(op::Parameter(0), op::Parameter(1), + op::Parameter(2), op::Parameter(3))); +} + +} // namespace +} // namespace xla 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 new file mode 100644 index 0000000000000000000000000000000000000000..e9c8aa03e2aa8f4866daf2a2f8d846e50fa68793 --- /dev/null +++ b/tensorflow/compiler/xla/service/dynamic_parameter_binding.cc @@ -0,0 +1,140 @@ +/* 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/dynamic_parameter_binding.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" + +namespace xla { + +Status DynamicParameterBinding::Bind( + const DynamicParameter& dynamic_parameter, + const DynamicDimension& dynamic_dimension) { + auto result = bindings_.emplace(dynamic_dimension, dynamic_parameter); + TF_RET_CHECK(result.second); + return Status::OK(); +} + +absl::optional +DynamicParameterBinding::GetBinding( + const DynamicDimension& dynamic_dimension) const { + auto param_iter = bindings_.find(dynamic_dimension); + if (param_iter == bindings_.end()) { + return absl::nullopt; + } + return param_iter->second; +} + +DynamicParameterBindingProto DynamicParameterBinding::ToProto() const { + DynamicParameterBindingProto result; + for (const auto& binding : bindings_) { + const DynamicDimension& dynamic_dimension = binding.first; + const DynamicParameter& dynamic_param = binding.second; + DynamicParameterBindingProto::Binding binding_proto; + binding_proto.set_dynamic_param_num(dynamic_param.parameter_num); + for (int64 i : dynamic_param.parameter_index) { + binding_proto.add_dynamic_param_index(i); + } + + binding_proto.set_target_param_num(dynamic_dimension.parameter_num); + + for (int64 i : dynamic_dimension.parameter_index) { + binding_proto.add_target_param_index(i); + } + + binding_proto.set_target_param_dim_num(dynamic_dimension.dimension); + result.add_entries()->Swap(&binding_proto); + } + return result; +} + +StatusOr DynamicParameterBinding::CreateFromProto( + const DynamicParameterBindingProto& proto) { + DynamicParameterBinding result; + for (const DynamicParameterBindingProto::Binding& binding : proto.entries()) { + int64 dynamic_param_num = binding.dynamic_param_num(); + ShapeIndex dynamic_param_index(binding.dynamic_param_index().begin(), + binding.dynamic_param_index().end()); + 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_dim_num(); + + TF_RETURN_IF_ERROR( + result.Bind(DynamicParameter{dynamic_param_num, dynamic_param_index}, + DynamicDimension{target_param_num, target_param_index, + target_dim_num})); + } + + return result; +} + +string DynamicParameterBinding::ToString() const { + std::vector pieces; + pieces.push_back("DynamicParameterBinding: "); + for (const auto& binding : bindings_) { + const DynamicDimension& dynamic_dimension = binding.first; + const DynamicParameter& dynamic_param = binding.second; + pieces.push_back(absl::StrFormat( + " -- Input param number %lld at %s has dim %lld as dynamic" + " dimension, which is represented by param number %lld at " + "%s", + dynamic_dimension.parameter_num, + dynamic_dimension.parameter_index.ToString(), + dynamic_dimension.dimension, dynamic_param.parameter_num, + dynamic_param.parameter_index.ToString())); + } + return absl::StrJoin(pieces, "\n"); +} + +Status DynamicParameterBinding::ForEachBinding(BindingFn fn) const { + for (const auto& binding : bindings_) { + TF_RETURN_IF_ERROR(fn(binding.second, binding.first)); + } + return Status::OK(); +} + +Status DynamicParameterBinding::Verify(const HloModule& module) const { + const HloComputation* entry = module.entry_computation(); + 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_dimension.parameter_num < entry->num_parameters()); + TF_RET_CHECK(ShapeUtil::IndexIsValid( + entry->parameter_instruction(dynamic_parameter.parameter_num)->shape(), + dynamic_parameter.parameter_index)); + TF_RET_CHECK(ShapeUtil::IndexIsValid( + entry->parameter_instruction(dynamic_dimension.parameter_num)->shape(), + dynamic_dimension.parameter_index)); + TF_RET_CHECK( + dynamic_dimension.dimension < + ShapeUtil::GetSubshape( + entry->parameter_instruction(dynamic_dimension.parameter_num) + ->shape(), + dynamic_dimension.parameter_index) + .rank()); + return Status::OK(); + }); +} + +std::ostream& operator<<(std::ostream& out, + const DynamicParameterBinding& binding) { + out << binding.ToString(); + return out; +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/dynamic_parameter_binding.h b/tensorflow/compiler/xla/service/dynamic_parameter_binding.h new file mode 100644 index 0000000000000000000000000000000000000000..57af2c43d3c65f7340e6a9f04e5abbf052ebceea --- /dev/null +++ b/tensorflow/compiler/xla/service/dynamic_parameter_binding.h @@ -0,0 +1,125 @@ +/* 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_DYNAMIC_PARAMETER_BINDING_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_DYNAMIC_PARAMETER_BINDING_H_ + +#include + +#include "absl/container/flat_hash_map.h" +#include "absl/types/optional.h" +#include "tensorflow/compiler/xla/service/hlo.pb.h" +#include "tensorflow/compiler/xla/shape_tree.h" +#include "tensorflow/compiler/xla/shape_util.h" + +namespace xla { + +class HloModule; +// We currently use an explicit API that takes an extra parameter to indicate +// the runtime size of a dynamic dimension. DynamicParameterBinding indicates +// the relationship between parameter: We can have a dynamic parameter that +// points to another target parameter to indicate that the target parameter is +// dynamic. +// +// +// TODO(b/119520625): Remove this API once we have more dynamic shape infra +// ready. +class DynamicParameterBinding { + public: + // DynamicParameter represents a special parameter that is used to represent + // the runtime size of a dimension of another parameter. A dynamic parameter + // has to be a scalar value. + struct DynamicParameter { + // The parameter number of dynamic parameter. + int64 parameter_num; + // The index of the parameter. + ShapeIndex parameter_index; + }; + + // DynamicDimension represents a dimension whose size is determined at + // runtime. A DynamicDimension's runtime size is determined by the binded + // DynamicParameter using `DynamicParameterBinding::Bind` method. + struct DynamicDimension { + // The parameter number of dynamic dimension. + int64 parameter_num; + // The subshape index of the parameter. + ShapeIndex parameter_index; + // The dimension number in the subshape. + int64 dimension; + + // "friend" keyword are added so these functions can be found by ADL. + template + friend H AbslHashValue(H h, const DynamicDimension& m) { + return H::combine(std::move(h), m.parameter_num, m.parameter_index, + m.dimension); + } + + friend bool operator==(const DynamicDimension& lhs, + const DynamicDimension& rhs) { + return lhs.parameter_num == rhs.parameter_num && + lhs.parameter_index == rhs.parameter_index && + lhs.dimension == rhs.dimension; + } + }; + + DynamicParameterBinding() = default; + + virtual ~DynamicParameterBinding() = default; + + // Adds binding which indicates that the dimension indicated by + // `dynamic_dimension` is dynamic, and its runtime size is represented by + // `dynamic_parameter`. + Status Bind(const DynamicParameter& dynamic_parameter, + const DynamicDimension& dynamic_dimension); + + // Returns the parameter and the index representing the runtime size of + // dimension `dim_num` of parameter `param_num` at `param_index`. + // + // Returns nullopt if the binding is not set. + absl::optional GetBinding( + const DynamicDimension& dynamic_dimension) const; + + using BindingFn = + std::function; + + // Iterate through each binding. + Status ForEachBinding(BindingFn fn) const; + + DynamicParameterBindingProto ToProto() const; + + static StatusOr CreateFromProto( + const DynamicParameterBindingProto& proto); + + string ToString() const; + + // Verifies that the given binding is valid for the given module. + // Specifically, the binding's parameter and parameter size should be valid. + Status Verify(const HloModule& module) const; + + private: + // Keeps track of mappings from DynamicDimension to DynamicParameter. The + // direction of is chosen so that we can easily query if a dimension is + // dynamic and which dynamic parameter represents the real size of that + // dimension. + absl::flat_hash_map bindings_; +}; + +std::ostream& operator<<(std::ostream& out, + const DynamicParameterBinding& binding); + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_DYNAMIC_PARAMETER_BINDING_H_ diff --git a/tensorflow/compiler/xla/service/dynamic_parameter_binding_test.cc b/tensorflow/compiler/xla/service/dynamic_parameter_binding_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..b5d57cda4f469a384dc0affdae9e5f93a70ac418 --- /dev/null +++ b/tensorflow/compiler/xla/service/dynamic_parameter_binding_test.cc @@ -0,0 +1,180 @@ +/* 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/compiler/xla/service/dynamic_parameter_binding.h" + +#include +#include + +#include "absl/algorithm/container.h" +#include "tensorflow/compiler/xla/service/hlo_computation.h" +#include "tensorflow/compiler/xla/service/hlo_dce.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/service/hlo_memory_scheduler.h" +#include "tensorflow/compiler/xla/service/hlo_opcode.h" +#include "tensorflow/compiler/xla/service/hlo_ordering.h" +#include "tensorflow/compiler/xla/service/hlo_parser.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" +#include "tensorflow/compiler/xla/types.h" +#include "tensorflow/core/lib/core/status_test_util.h" + +namespace xla { +namespace { +class DynamicParameterBindingTest : public HloTestBase { + protected: + // Serialize and then deserialize a binding. + void SerializeAndDeserialize(DynamicParameterBinding* binding) { + DynamicParameterBindingProto proto = binding->ToProto(); + 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 + // dimension. + const string module_str = R"( +HloModule TEST + +ENTRY main { + a = f32[] parameter(0) + b = f32[10] parameter(1) + ROOT root = (f32[], f32[10]) tuple(%a, %b) +} +)"; + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseHloString(module_str)); + + DynamicParameterBinding binding; + + TF_EXPECT_OK( + binding.Bind(DynamicParameterBinding::DynamicParameter{0, {}}, + DynamicParameterBinding::DynamicDimension{1, {}, 0})); + + 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) { + // 'gte2' is a dynamic shape; 'gte1' represents the real size of gte2's first + // dimension. + const string module_str = R"( +HloModule TEST + +ENTRY main { + param = (f32[], f32[10]) parameter(0) + gte1 = f32[] get-tuple-element(%param), index=0 + gte2 = f32[10] get-tuple-element(%param), index=1 + ROOT root = (f32[], f32[10]) tuple(%gte1, %gte2) +} +)"; + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseHloString(module_str)); + + DynamicParameterBinding binding; + + TF_EXPECT_OK( + binding.Bind(DynamicParameterBinding::DynamicParameter{0, {0}}, + DynamicParameterBinding::DynamicDimension{0, {1}, 0})); + + 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) { + // 'gte2' is a dynamic shape; 'gte1' represents the real size of gte2's both + // dimensions. + const string module_str = R"( +HloModule TEST + +ENTRY main { + param = (f32[], f32[10, 10]) parameter(0) + gte1 = f32[] get-tuple-element(%param), index=0 + gte2 = f32[10, 10] get-tuple-element(%param), index=1 + ROOT root = (f32[], f32[10, 10]) tuple(%gte1, %gte2) +} +)"; + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseHloString(module_str)); + + DynamicParameterBinding binding; + + TF_EXPECT_OK( + binding.Bind(DynamicParameterBinding::DynamicParameter{0, {0}}, + DynamicParameterBinding::DynamicDimension{0, {1}, 0})); + + TF_EXPECT_OK( + binding.Bind(DynamicParameterBinding::DynamicParameter{0, {0}}, + DynamicParameterBinding::DynamicDimension{0, {1}, 1})); + + 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 +} // namespace xla diff --git a/tensorflow/compiler/xla/service/elemental_ir_emitter.cc b/tensorflow/compiler/xla/service/elemental_ir_emitter.cc index 515267edd7caf42e04ebe638b99006db8967ea30..727e0bfa52d45b6f8c67d7d04613e4865f18a53c 100644 --- a/tensorflow/compiler/xla/service/elemental_ir_emitter.cc +++ b/tensorflow/compiler/xla/service/elemental_ir_emitter.cc @@ -22,6 +22,7 @@ limitations under the License. // IWYU pragma: no_include "llvm/IR/Intrinsics.gen.inc" #include "absl/algorithm/container.h" +#include "absl/container/flat_hash_map.h" #include "absl/strings/str_cat.h" #include "llvm/IR/BasicBlock.h" #include "llvm/IR/Instructions.h" @@ -811,11 +812,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); @@ -827,7 +831,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'", @@ -1326,9 +1336,9 @@ llvm_ir::IrArray::Index ElementalIrEmitter::ElementwiseSourceIndex( // If implicit broadcast is needed, the source dimensions that are broadcast // have index 0. - CHECK_EQ(ShapeUtil::Rank(operand_shape), ShapeUtil::Rank(hlo.shape())); + CHECK_EQ(operand_shape.rank(), hlo.shape().rank()); llvm_ir::IrArray::Index source_index(target_index.GetType()); - for (int64 i = 0; i < ShapeUtil::Rank(hlo.shape()); ++i) { + for (int64 i = 0; i < hlo.shape().rank(); ++i) { if (hlo.shape().dimensions(i) == operand_shape.dimensions(i)) { source_index.push_back(target_index[i]); } else { @@ -1671,26 +1681,66 @@ StatusOr ElementalIrEmitter::EmitElementalConcatenate( b_->SetInsertPoint(init_block); + // Assign a unique id for each *different* operand, and count how often each + // operand is used. If all operands are different, the usage count will be 1 + // for each operand. + absl::flat_hash_map to_unique_operand_id; + std::vector operand_usage_count; + for (const auto* operand : hlo->operands()) { + if (to_unique_operand_id.contains(operand)) { + ++operand_usage_count[to_unique_operand_id[operand]]; + } else { + int64 unique_operand_id = to_unique_operand_id.size(); + to_unique_operand_id[operand] = unique_operand_id; + operand_usage_count.push_back(1); + } + } + + // To avoid that we emit the same operand more than once, we create one basic + // block for each *different* operand with a PHI node for the different source + // index inputs. + std::vector emit_operand_blocks( + to_unique_operand_id.size(), nullptr); + std::vector source_index_phis(to_unique_operand_id.size(), + nullptr); + for (const auto* operand : hlo->operands()) { + int64 operand_id = to_unique_operand_id[operand]; + if (emit_operand_blocks[operand_id] != nullptr) { + continue; + } + + emit_operand_blocks[operand_id] = llvm_ir::CreateBasicBlock( + exit_block, StrCat("concat_index_from_operand_id", operand_id), b_); + auto saved_insert_point = b_->GetInsertPoint(); + llvm_ir::SetToFirstInsertPoint(emit_operand_blocks[operand_id], b_); + source_index_phis[operand_id] = + PHI(source_index.GetType(), operand_usage_count[operand_id]); + auto operand_index = source_index; + operand_index[concat_dim] = source_index_phis[operand_id]; + + // Create the terminator of the block before calling operand generators, + // because they require non-degenerate basic blocks. + b_->SetInsertPoint(llvm::BranchInst::Create( + exit_block, /*InsertAtEnd=*/emit_operand_blocks[operand_id])); + TF_ASSIGN_OR_RETURN(llvm::Value * value, + operand_to_generator.at(operand)(operand_index)); + output->addIncoming(value, b_->GetInsertBlock()); + b_->SetInsertPoint(init_block, saved_insert_point); + } + for (int64 operand_idx = 0; operand_idx < hlo->operand_count(); ++operand_idx) { const HloInstruction* operand = hlo->operand(operand_idx); - auto true_block = llvm_ir::CreateBasicBlock( - exit_block, StrCat("concat_index_from_operand", operand_idx), b_); auto false_block = llvm_ir::CreateBasicBlock( exit_block, StrCat("concat_index_not_from_operand", operand_idx), b_); auto concat_dim_size = llvm::ConstantInt::get(source_index[concat_dim]->getType(), operand->shape().dimensions(concat_dim)); - CondBr(ICmpULT(source_index[concat_dim], concat_dim_size), true_block, - false_block); - - // Create the terminator of the true block before calling operand - // generators, because they require non-degenerate basic blocks. - b_->SetInsertPoint( - llvm::BranchInst::Create(exit_block, /*InsertAtEnd=*/true_block)); - TF_ASSIGN_OR_RETURN(llvm::Value * value, - operand_to_generator.at(operand)(source_index)); - output->addIncoming(value, b_->GetInsertBlock()); + int64 operand_id = to_unique_operand_id[operand]; + source_index_phis[operand_id]->addIncoming(source_index[concat_dim], + b_->GetInsertBlock()); + CondBr(ICmpULT(source_index[concat_dim], concat_dim_size), + emit_operand_blocks[operand_id], false_block); // Subtract the size of the concat dimension of the current operand // from the source index. @@ -1709,7 +1759,7 @@ StatusOr ElementalIrEmitter::EmitElementalDynamicSlice( const llvm_ir::IrArray::Index& index) { // Emit IR to read dynamic start indices from hlo->operand(1). const HloInstruction* input_hlo = hlo->operand(0); - const int64 rank = ShapeUtil::Rank(input_hlo->shape()); + const int64 rank = input_hlo->shape().rank(); // Use the same index type for all tensor accesses in the same kernel. llvm::Type* index_type = index.GetType(); llvm_ir::IrArray::Index slice_start_index(index_type, rank); @@ -1717,9 +1767,18 @@ StatusOr ElementalIrEmitter::EmitElementalDynamicSlice( auto index_typed_const = [&](uint64 c) -> llvm::Constant* { return llvm::ConstantInt::get(index_type, c); }; - llvm_ir::IrArray::Index dim_index(1, index_typed_const(i)); - TF_ASSIGN_OR_RETURN(llvm::Value * start_index_value, - operand_to_generator.at(hlo->operand(1))(dim_index)); + // 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)); + } // 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) @@ -1815,8 +1874,6 @@ StatusOr ElementalIrEmitter::EmitElementalGather( // Clamp the gather index so that the gather region fits in the operand. // gather_dim_component_extended_inbound = // clamp(gather_dim_component_extended, 0, largest_valid_start_index); - - // TODO(b/111078873): This is implementation defined behavior. bool is_signed = ShapeUtil::ElementIsSigned(indices_shape); auto gather_dim_component_extended_inbound = EmitIntegralMin( index.GetConstantWithIndexType(largest_valid_start_index), @@ -1854,7 +1911,7 @@ StatusOr ElementalIrEmitter::EmitElementalDynamicUpdateSlice( const HloInstruction* update_hlo = hlo->operand(1); const HloInstruction* start_hlo = hlo->operand(2); // Calculate slice start/end indices. - const int64 rank = ShapeUtil::Rank(input_hlo->shape()); + const int64 rank = input_hlo->shape().rank(); llvm_ir::IrArray::Index slice_start_index(index.GetType(), rank); llvm_ir::IrArray::Index slice_limit_index(index.GetType(), rank); // Slice intersection gathers (ANDs) conditions on all ranks for which @@ -1866,9 +1923,19 @@ StatusOr ElementalIrEmitter::EmitElementalDynamicUpdateSlice( auto index_typed_const = [&](uint64 c) -> llvm::Constant* { return llvm::ConstantInt::get(index_type, c); }; - llvm_ir::IrArray::Index dim_index(1, index_typed_const(i)); - TF_ASSIGN_OR_RETURN(llvm::Value * start_index_value, - operand_to_generator.at(start_hlo)(dim_index)); + + 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)); + } // 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) @@ -2186,7 +2253,7 @@ llvm_ir::ElementGenerator ElementalIrEmitter::MakeElementGenerator( auto* iota = Cast(hlo); PrimitiveType element_type = iota->shape().element_type(); IrArray::Index elem_index = - ShapeUtil::Rank(iota->shape()) > 1 + iota->shape().rank() > 1 ? target_index.SourceIndexOfBroadcast( iota->shape(), ShapeUtil::MakeShapeWithDescendingLayout( @@ -2206,13 +2273,15 @@ llvm_ir::ElementGenerator ElementalIrEmitter::MakeElementGenerator( : iota->shape(); PrimitiveType component_element_type = component_shape.element_type(); llvm::Value* iota_result; - if (ShapeUtil::ElementIsIntegral(component_shape)) { + if (primitive_util::IsIntegralType(component_element_type) || + component_element_type == PRED) { iota_result = b_->CreateIntCast( elem_index_linear, llvm_ir::PrimitiveTypeToIrType(component_element_type, module_), /*isSigned=*/false); } else { - TF_RET_CHECK(ShapeUtil::ElementIsFloating(component_shape)) + TF_RET_CHECK( + primitive_util::IsFloatingPointType(component_element_type)) << component_element_type; llvm::Type* float_ir_type; if (component_element_type == BF16) { diff --git a/tensorflow/compiler/xla/service/executable.cc b/tensorflow/compiler/xla/service/executable.cc index 47c56e2f7fbd9f53be6a2b189c5c36cf4fdcdccb..10b8c01ff1383658fcfb2271c177ba54347f985a 100644 --- a/tensorflow/compiler/xla/service/executable.cc +++ b/tensorflow/compiler/xla/service/executable.cc @@ -17,7 +17,7 @@ limitations under the License. #include "absl/memory/memory.h" #include "absl/strings/str_format.h" -#include "tensorflow/compiler/xla/legacy_flags/debug_options_flags.h" +#include "tensorflow/compiler/xla/debug_options_flags.h" #include "tensorflow/compiler/xla/service/hlo_graph_dumper.h" #include "tensorflow/compiler/xla/status.h" #include "tensorflow/compiler/xla/status_macros.h" diff --git a/tensorflow/compiler/xla/service/executable.h b/tensorflow/compiler/xla/service/executable.h index 3a6780f2a67f230cae626ea00cfbf93b4e60d968..b34bca55a48b113c325dbf28c03f7a0f5b71f658 100644 --- a/tensorflow/compiler/xla/service/executable.h +++ b/tensorflow/compiler/xla/service/executable.h @@ -22,7 +22,7 @@ limitations under the License. #include "absl/types/span.h" #include "absl/types/variant.h" -#include "tensorflow/compiler/xla/legacy_flags/debug_options_flags.h" +#include "tensorflow/compiler/xla/debug_options_flags.h" #include "tensorflow/compiler/xla/service/computation_layout.h" #include "tensorflow/compiler/xla/service/device_memory_allocator.h" #include "tensorflow/compiler/xla/service/hlo.pb.h" @@ -61,7 +61,7 @@ struct ExecutionOutput { class Executable { public: explicit Executable( - std::unique_ptr hlo_module, + std::unique_ptr hlo_module, std::unique_ptr hlo_profile_printer_data, std::unique_ptr hlo_profile_index_map) : hlo_module_(std::move(hlo_module)), @@ -162,7 +162,7 @@ class Executable { return hlo_profile_printer_data_ != nullptr; } - const HloModule& module() const { return *hlo_module_; } + HloModule& module() const { return *hlo_module_; } const bool has_module() const { return hlo_module_ != nullptr; } @@ -199,7 +199,7 @@ class Executable { // HloModule this was compiled from. BufferAssignment keeps pointers to // HloInstructions owned by the HloModule so we need to keep the HloModule // around. - const std::unique_ptr hlo_module_; + const std::unique_ptr hlo_module_; // HloSnapshot this was compiled from. Null if not dumping executions. std::unique_ptr hlo_snapshot_; diff --git a/tensorflow/compiler/xla/service/flatten_call_graph_test.cc b/tensorflow/compiler/xla/service/flatten_call_graph_test.cc index 5fbd73a5363b4cdbcaafedbe6f4e7bd6bb2a92d8..8eeb930b48165a2e3c622581e05cb5f7063fa1fa 100644 --- a/tensorflow/compiler/xla/service/flatten_call_graph_test.cc +++ b/tensorflow/compiler/xla/service/flatten_call_graph_test.cc @@ -22,7 +22,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/tests/hlo_verified_test_base.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" #include "tensorflow/compiler/xla/util.h" #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/core/status_test_util.h" @@ -30,7 +30,7 @@ limitations under the License. namespace xla { namespace { -class FlattenCallGraphTest : public HloVerifiedTestBase { +class FlattenCallGraphTest : public HloTestBase { protected: // Build and return a trivial computation taking and returning a scalar. std::unique_ptr MakeScalarComputation() { @@ -108,7 +108,7 @@ TEST_F(FlattenCallGraphTest, ComplexGraph) { // c // // Calls are made via kCall, kWhile, and kMap instructions. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* cond_computation = module->AddEmbeddedComputation(MakeConditionComputation()); HloComputation* c_computation = @@ -139,9 +139,9 @@ TEST_F(FlattenCallGraphTest, ComplexGraph) { } { - TF_ASSERT_OK_AND_ASSIGN(bool result, RunFlattenCallGraph(module)); + TF_ASSERT_OK_AND_ASSIGN(bool result, RunFlattenCallGraph(module.get())); EXPECT_TRUE(result); - std::unique_ptr flat_call_graph = CallGraph::Build(module); + std::unique_ptr flat_call_graph = CallGraph::Build(module.get()); const CallGraphNode& c_node = flat_call_graph->GetNode(c_computation); EXPECT_EQ(1, c_node.caller_callsites().size()); } @@ -149,7 +149,7 @@ TEST_F(FlattenCallGraphTest, ComplexGraph) { // Test corner case of a computation used as a body and a loop condition. TEST_F(FlattenCallGraphTest, SharedWhileConditionAndBody) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* cond_computation; { HloComputation::Builder builder(TestName() + ".cond"); @@ -176,15 +176,15 @@ TEST_F(FlattenCallGraphTest, SharedWhileConditionAndBody) { } { - std::unique_ptr call_graph = CallGraph::Build(module); + std::unique_ptr call_graph = CallGraph::Build(module.get()); const CallGraphNode& cond_node = call_graph->GetNode(cond_computation); EXPECT_EQ(2, cond_node.caller_callsites().size()); } { - TF_ASSERT_OK_AND_ASSIGN(bool result, RunFlattenCallGraph(module)); + TF_ASSERT_OK_AND_ASSIGN(bool result, RunFlattenCallGraph(module.get())); EXPECT_TRUE(result); - std::unique_ptr call_graph = CallGraph::Build(module); + std::unique_ptr call_graph = CallGraph::Build(module.get()); const CallGraphNode& cond_node = call_graph->GetNode(cond_computation); EXPECT_EQ(1, cond_node.caller_callsites().size()); } @@ -201,7 +201,7 @@ TEST_F(FlattenCallGraphTest, SharedWhileConditionAndBody) { // C // TEST_F(FlattenCallGraphTest, FlattenCalls) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* c_computation = module->AddEmbeddedComputation(MakeScalarComputation()); @@ -211,9 +211,9 @@ TEST_F(FlattenCallGraphTest, FlattenCalls) { module->AddEntryComputation( MakeCallingComputation(b_computation, /*callsites=*/2, ".Entry")); - TF_ASSERT_OK_AND_ASSIGN(bool result, RunFlattenCallGraph(module)); + TF_ASSERT_OK_AND_ASSIGN(bool result, RunFlattenCallGraph(module.get())); EXPECT_TRUE(result); - std::unique_ptr call_graph = CallGraph::Build(module); + std::unique_ptr call_graph = CallGraph::Build(module.get()); EXPECT_EQ(7, module->computation_count()); const CallGraphNode& c_node = call_graph->GetNode(c_computation); @@ -224,7 +224,7 @@ TEST_F(FlattenCallGraphTest, FlattenCalls) { } TEST_F(FlattenCallGraphTest, FlattenCallsInConditional) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* sub_computation = module->AddEmbeddedComputation(MakeScalarComputation()); @@ -243,9 +243,9 @@ TEST_F(FlattenCallGraphTest, FlattenCallsInConditional) { module->AddEntryComputation(builder.Build()); EXPECT_EQ(2, module->computation_count()); - TF_ASSERT_OK_AND_ASSIGN(bool result, RunFlattenCallGraph(module)); + TF_ASSERT_OK_AND_ASSIGN(bool result, RunFlattenCallGraph(module.get())); EXPECT_TRUE(result); - std::unique_ptr call_graph = CallGraph::Build(module); + std::unique_ptr call_graph = CallGraph::Build(module.get()); // The true and false computations must now be different. EXPECT_EQ(3, module->computation_count()); diff --git a/tensorflow/compiler/xla/service/gather_expander.cc b/tensorflow/compiler/xla/service/gather_expander.cc index cb86c9857936f21d9d2ac6bc22c725b89cca6482..590942cddcdd138981ee829f090ae17b0d038e1a 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) { @@ -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( @@ -346,7 +343,8 @@ StatusOr GatherExpander::ExpandGather( [&](HloInstruction* indvar, const std::vector& loop_state) { return GatherLoopBody(*gather_instr, indvar, loop_state); - }); + }, + gather_instr->metadata()); TF_ASSIGN_OR_RETURN(std::vector gather_loop_result, gather_loop_result_or_error); diff --git a/tensorflow/compiler/xla/service/gather_expander_test.cc b/tensorflow/compiler/xla/service/gather_expander_test.cc index 141dd4d6f10272ce749edc4e91153c365ed322e6..e1ea5c39d58b6d23b076740626ca0ad63dc341ee 100644 --- a/tensorflow/compiler/xla/service/gather_expander_test.cc +++ b/tensorflow/compiler/xla/service/gather_expander_test.cc @@ -89,7 +89,7 @@ ENTRY main { // an implementation detail from WhileUtil::MakeCountedLoop). const Shape& while_shape = while_instr->shape(); - ASSERT_TRUE(ShapeUtil::IsTuple(while_shape)); + ASSERT_TRUE(while_shape.IsTuple()); ASSERT_EQ(ShapeUtil::TupleElementCount(while_shape), 4); EXPECT_TRUE(ShapeUtil::SameDimensions( @@ -104,5 +104,44 @@ ENTRY main { ShapeUtil::MakeShape(S32, {2, 3}), ShapeUtil::GetTupleElementShape(while_shape, 3))); } + +TEST(GatherExpanderTest, CheckOpMetadata) { + const string hlo_text = R"( +HloModule TensorFlowGatherV2 + +ENTRY main { + operand = s32[3,3] parameter(0) + indices = s32[2] parameter(1) + ROOT gather = s32[3,2] gather(operand, indices), + offset_dims={0}, + collapsed_slice_dims={1}, + start_index_map={1}, + index_vector_dim=1, + slice_sizes={3, 1} +} +)"; + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseHloString(hlo_text)); + OpMetadata metadata; + metadata.set_op_name("Gather"); + module->entry_computation()->root_instruction()->set_metadata(metadata); + TF_ASSERT_OK_AND_ASSIGN(bool changed, GatherExpander{}.Run(module.get())); + ASSERT_TRUE(changed); + + HloInstruction* while_instr = nullptr; + for (auto* instr : module->entry_computation()->instructions()) { + if (instr->opcode() == HloOpcode::kWhile) { + ASSERT_EQ(while_instr, nullptr) + << "Expected exactly one while instruction in the entry computation " + "after gather expansion"; + while_instr = instr; + } + } + + ASSERT_NE(while_instr, nullptr) + << "Expected exactly one while instruction in the entry computation " + "after gather expansion"; + EXPECT_EQ(while_instr->metadata().op_name(), "Gather"); +} } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/service/generic_transfer_manager.cc b/tensorflow/compiler/xla/service/generic_transfer_manager.cc index bec02e14f951c6d905b7329be5c02896984279d0..7d450f4b53cdea209f2ef10ba785be6ec3b8bf8d 100644 --- a/tensorflow/compiler/xla/service/generic_transfer_manager.cc +++ b/tensorflow/compiler/xla/service/generic_transfer_manager.cc @@ -83,7 +83,7 @@ Status GenericTransferManager::TransferLiteralFromDeviceInternal( TF_RETURN_IF_ERROR(ShapeUtil::ForEachSubshapeWithStatus( device_buffer.on_host_shape(), [&](const Shape& subshape, const ShapeIndex& index) -> Status { - if (ShapeUtil::IsArray(subshape)) { + if (subshape.IsArray()) { TF_RETURN_IF_ERROR(executor->SynchronousMemcpyD2H( /*source=*/device_buffer.buffer(index), /*size=*/GetByteSizeRequirement(subshape), @@ -120,7 +120,7 @@ Status GenericTransferManager::TransferLiteralToDeviceAsync( device_buffer.on_host_shape(), [&](const Shape& device_subshape, const ShapeIndex& index) -> Status { se::DeviceMemoryBase device_memory = device_buffer.buffer(index); - if (ShapeUtil::IsArray(device_subshape)) { + if (device_subshape.IsArray()) { TF_RET_CHECK(GetByteSizeRequirement(device_subshape) == device_memory.size()); // Element is array-shaped: transfer array data to device buffer. diff --git a/tensorflow/compiler/xla/service/gpu/BUILD b/tensorflow/compiler/xla/service/gpu/BUILD index 449fd919d64612ba10932c7cc0865f1fca96424a..dc17aa4426236f54e5f03c28634278d45f462158 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", ], ) @@ -111,7 +110,6 @@ tf_cc_test( "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", "//tensorflow/compiler/xla/tests:test_utils", "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:lib", @@ -136,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", @@ -264,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", ], @@ -363,6 +365,7 @@ cc_library( "//tensorflow/core/platform/default/build_config:stream_executor_cuda", # build_cleaner: keep "//tensorflow/stream_executor", "@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", @@ -463,7 +466,7 @@ tf_cc_test( "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/service:hlo_matchers", "//tensorflow/compiler/xla/service:shape_inference", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", + "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/compiler/xla/tests:xla_internal_test_main", # fixdeps: keep "//tensorflow/core:test", ], @@ -627,7 +630,7 @@ tf_cc_test( "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla/service:hlo_matchers", "//tensorflow/compiler/xla/service:hlo_parser", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", + "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/compiler/xla/tests:xla_internal_test_main", # build_cleaner: keep ], ) @@ -683,6 +686,7 @@ cc_library( ":partition_assignment", ":stream_assignment", ":stream_executor_util", + ":variadic_op_splitter", "//tensorflow/compiler/xla:protobuf_util", "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:statusor", @@ -694,6 +698,9 @@ cc_library( "//tensorflow/compiler/xla/service:buffer_liveness", "//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", @@ -701,6 +708,7 @@ cc_library( "//tensorflow/compiler/xla/service:hlo_cse", "//tensorflow/compiler/xla/service:hlo_dce", "//tensorflow/compiler/xla/service:hlo_element_type_converter", + "//tensorflow/compiler/xla/service:hlo_get_dimension_size_rewriter", "//tensorflow/compiler/xla/service:hlo_pass", "//tensorflow/compiler/xla/service:hlo_pass_pipeline", "//tensorflow/compiler/xla/service:hlo_proto", @@ -710,6 +718,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", @@ -723,6 +732,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", @@ -848,7 +858,6 @@ tf_cc_test( "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", "//tensorflow/compiler/xla/tests:test_utils", "//tensorflow/compiler/xla/tests:xla_internal_test_main", "@com_google_absl//absl/memory", @@ -908,7 +917,6 @@ tf_cc_test( "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/compiler/xla/tests:hlo_verified_test_base", "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:protos_all_cc", "//tensorflow/core:test", @@ -999,3 +1007,38 @@ tf_cc_test( "@com_google_absl//absl/strings", ], ) + +cc_library( + name = "variadic_op_splitter", + srcs = ["variadic_op_splitter.cc"], + hdrs = ["variadic_op_splitter.h"], + deps = [ + "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:util", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla/service:hlo", + "//tensorflow/compiler/xla/service:hlo_pass", + "//tensorflow/core:lib", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/types:span", + ], +) + +tf_cc_test( + name = "variadic_op_splitter_test", + srcs = ["variadic_op_splitter_test.cc"], + deps = [ + ":ir_emission_utils", + ":variadic_op_splitter", + "//tensorflow/compiler/xla:literal_util", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:status_macros", + "//tensorflow/compiler/xla:util", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla/service:hlo", + "//tensorflow/compiler/xla/service:hlo_matchers", + "//tensorflow/compiler/xla/service:hlo_parser", + "//tensorflow/compiler/xla/service:pattern_matcher", + "//tensorflow/compiler/xla/tests:hlo_test_base", + ], +) 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..309b0aca64954e64509d731dce28ce9d8da4ee43 100644 --- a/tensorflow/compiler/xla/service/gpu/cudnn_conv_algorithm_picker.cc +++ b/tensorflow/compiler/xla/service/gpu/cudnn_conv_algorithm_picker.cc @@ -146,7 +146,8 @@ tensorflow::mutex_lock LockGpu(const se::StreamExecutor* stream_exec) { // 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) { +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 = @@ -249,12 +250,13 @@ CudnnConvAlgorithmPicker::PickBestAlgorithm(HloCustomCallInstruction* instr) { 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)); + // Use assignment instead of brace-list to make GCC 4.9 happy. + RunConvOptions options; + options.profile_result = &profile_result; + options.algo_override = alg; bool launch_ok = RunCudnnConv(instr, absl::MakeSpan(operand_buffers), result_buffer, - &scratch_allocator, &stream, &profile_result) + &scratch_allocator, &stream, options) .ok(); if (launch_ok && profile_result.is_valid()) { 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..4991db0948589e479a202f4082d96df275f6e088 100644 --- a/tensorflow/compiler/xla/service/gpu/cudnn_conv_algorithm_picker.h +++ b/tensorflow/compiler/xla/service/gpu/cudnn_conv_algorithm_picker.h @@ -56,7 +56,8 @@ class CudnnConvAlgorithmPicker : public HloModulePass { 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_pad_for_tensor_cores.cc b/tensorflow/compiler/xla/service/gpu/cudnn_conv_pad_for_tensor_cores.cc index 5aa4f839f4be5f1060480fea98775f8ffada0bdd..958e0b9c6e7b7885f87b90d61ee5b3bbf6ab2702 100644 --- a/tensorflow/compiler/xla/service/gpu/cudnn_conv_pad_for_tensor_cores.cc +++ b/tensorflow/compiler/xla/service/gpu/cudnn_conv_pad_for_tensor_cores.cc @@ -50,10 +50,10 @@ static HloInstruction* PadInstruction(HloInstruction* instr, auto* zero = comp->AddInstruction( HloInstruction::CreateConstant(LiteralUtil::Zero(shape.element_type()))); - PaddingConfig pad_config = MakeNoPaddingConfig(ShapeUtil::Rank(shape)); + PaddingConfig pad_config = MakeNoPaddingConfig(shape.rank()); bool added_padding = false; - for (int64 dim = 0; dim < ShapeUtil::Rank(shape); ++dim) { + for (int64 dim = 0; dim < shape.rank(); ++dim) { if (shape.dimensions(dim) == new_shape.dimensions(dim)) { continue; } diff --git a/tensorflow/compiler/xla/service/gpu/cudnn_conv_pad_for_tensor_cores_test.cc b/tensorflow/compiler/xla/service/gpu/cudnn_conv_pad_for_tensor_cores_test.cc index fa3afa6a5d318c399dc38e8934199b5a1393669e..af9303a5b761b99705945f1c02303156e3f874de 100644 --- a/tensorflow/compiler/xla/service/gpu/cudnn_conv_pad_for_tensor_cores_test.cc +++ b/tensorflow/compiler/xla/service/gpu/cudnn_conv_pad_for_tensor_cores_test.cc @@ -19,7 +19,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_matchers.h" #include "tensorflow/compiler/xla/service/hlo_parser.h" #include "tensorflow/compiler/xla/status_macros.h" -#include "tensorflow/compiler/xla/tests/hlo_verified_test_base.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" #include "tensorflow/compiler/xla/util.h" namespace xla { @@ -29,7 +29,7 @@ namespace { namespace op = xla::testing::opcode_matchers; using ::testing::_; -class CudnnConvPadForTensorCoresTest : public HloVerifiedTestBase {}; +class CudnnConvPadForTensorCoresTest : public HloTestBase {}; TEST_F(CudnnConvPadForTensorCoresTest, PadF16ForwardConvInputChannels) { auto module = ParseAndReturnVerifiedModule(R"( diff --git a/tensorflow/compiler/xla/service/gpu/cudnn_conv_padding_legalization.cc b/tensorflow/compiler/xla/service/gpu/cudnn_conv_padding_legalization.cc index d7829045cc127deaa4c2c9b705dca5285d704af2..17d0f7aa7bf6031148aae79f74f7878d6fca9574 100644 --- a/tensorflow/compiler/xla/service/gpu/cudnn_conv_padding_legalization.cc +++ b/tensorflow/compiler/xla/service/gpu/cudnn_conv_padding_legalization.cc @@ -43,13 +43,14 @@ bool IsForwardConvolutionCanonical(const HloInstruction& conv) { // dilation), returns kPad and/or kSlice instructions that explicitly apply the // padding; otherwise returns the original input operand. When there is both // positive padding (including dilation) and negative padding, we insert both -// kPad and kSlice. +// kPad and kSlice. Modifies 'conv_window' accordingly if any padding was moved +// into a kPad or kSlice op. HloInstruction* MaybePaddedAndSlicedInput( - const Window& conv_window, const ConvolutionDimensionNumbers& conv_dnums, + Window* conv_window, const ConvolutionDimensionNumbers& conv_dnums, HloInstruction* input) { HloComputation* computation = input->parent(); - if (!window_util::HasSymmetricPadding(conv_window) || - window_util::HasBaseDilation(conv_window)) { + if (!window_util::HasSymmetricPadding(*conv_window) || + window_util::HasBaseDilation(*conv_window)) { // If padding is uneven or has dilation, we insert a kPad instruction that // applies positive padding and dilation. // @@ -62,12 +63,21 @@ HloInstruction* MaybePaddedAndSlicedInput( MakeNoPaddingConfig(input->shape().dimensions_size()); for (size_t i = 0; i < conv_dnums.input_spatial_dimensions().size(); ++i) { int64 dim = conv_dnums.input_spatial_dimensions(i); - padding_config.mutable_dimensions(dim)->set_edge_padding_low( - std::max(0LL, conv_window.dimensions(i).padding_low())); - padding_config.mutable_dimensions(dim)->set_edge_padding_high( - std::max(0LL, conv_window.dimensions(i).padding_high())); - padding_config.mutable_dimensions(dim)->set_interior_padding( - conv_window.dimensions(i).base_dilation() - 1); + if (conv_window->dimensions(i).padding_low() > 0) { + padding_config.mutable_dimensions(dim)->set_edge_padding_low( + conv_window->dimensions(i).padding_low()); + conv_window->mutable_dimensions(i)->set_padding_low(0); + } + if (conv_window->dimensions(i).padding_high() > 0) { + padding_config.mutable_dimensions(dim)->set_edge_padding_high( + conv_window->dimensions(i).padding_high()); + conv_window->mutable_dimensions(i)->set_padding_high(0); + } + if (conv_window->dimensions(i).base_dilation() != 1) { + padding_config.mutable_dimensions(dim)->set_interior_padding( + conv_window->dimensions(i).base_dilation() - 1); + conv_window->mutable_dimensions(i)->set_base_dilation(1); + } } PrimitiveType element_type = input->shape().element_type(); HloInstruction* padding = computation->AddInstruction( @@ -75,7 +85,7 @@ HloInstruction* MaybePaddedAndSlicedInput( input = MakePadHlo(input, padding, padding_config).ValueOrDie(); } - if (window_util::HasNegativePadding(conv_window)) { + if (window_util::HasNegativePadding(*conv_window)) { // If the window has negative padding, insert a kSlice that explicitly // applies negative padding. // @@ -89,10 +99,14 @@ HloInstruction* MaybePaddedAndSlicedInput( int64 dim = conv_dnums.input_spatial_dimensions(i); // If dimension "dim" has negative padding, increase the start index or // decrement the limit index by the amount of negative padding. - start_indices[dim] += - std::max(0LL, -conv_window.dimensions(i).padding_low()); - limit_indices[dim] -= - std::max(0LL, -conv_window.dimensions(i).padding_high()); + if (conv_window->dimensions(i).padding_low() < 0) { + start_indices[dim] += -conv_window->dimensions(i).padding_low(); + conv_window->mutable_dimensions(i)->set_padding_low(0); + } + if (conv_window->dimensions(i).padding_high() < 0) { + limit_indices[dim] -= -conv_window->dimensions(i).padding_high(); + conv_window->mutable_dimensions(i)->set_padding_high(0); + } } input = @@ -140,25 +154,22 @@ bool CudnnConvPaddingLegalization::CanonicalizeForwardConvolution( // Insert slices and/or pads between the convolution and its input and/or // kernel operand. + Window new_conv_window = conv->window(); HloInstruction* new_input = MaybePaddedAndSlicedInput( - conv->window(), conv->convolution_dimension_numbers(), + &new_conv_window, conv->convolution_dimension_numbers(), conv->mutable_operand(0)); HloInstruction* new_kernel = - MaybePaddedKernel(conv->window(), conv->convolution_dimension_numbers(), + MaybePaddedKernel(new_conv_window, conv->convolution_dimension_numbers(), conv->mutable_operand(1)); - // Remove the padding from convolution's window field. These paddings are - // made explicit with the inserted pads. - Window new_conv_window = conv->window(); + // Remove the window dilation from convolution's window field. These paddings + // are made explicit with the pads inserted by MaybePaddedKernel(). for (size_t i = 0; i < new_conv_window.dimensions_size(); ++i) { WindowDimension* dim = new_conv_window.mutable_dimensions(i); // The size of the kernel may have changed so update the Window to match. dim->set_size(new_kernel->shape().dimensions( conv->convolution_dimension_numbers().kernel_spatial_dimensions(i))); - dim->set_padding_low(0); - dim->set_padding_high(0); - dim->set_base_dilation(1); dim->set_window_dilation(1); } @@ -208,7 +219,7 @@ bool CudnnConvPaddingLegalization::CanonicalizeBackwardFilterConvolution( Window new_backward_conv_window = backward_conv->window(); // input_padding_config is the config of the kPad to be inserted. PaddingConfig input_padding_config = - MakeNoPaddingConfig(ShapeUtil::Rank(input->shape())); + MakeNoPaddingConfig(input->shape().rank()); ConvolutionDimensionNumbers backward_conv_dnums = backward_conv->convolution_dimension_numbers(); for (size_t i = 0; i < backward_conv->window().dimensions_size(); ++i) { diff --git a/tensorflow/compiler/xla/service/gpu/cudnn_conv_rewriter.cc b/tensorflow/compiler/xla/service/gpu/cudnn_conv_rewriter.cc index 01de110aa9361f5813231767ad01e4aac03cfe0a..e81850db69edced29ea31bb2a526b0503bf8a453 100644 --- a/tensorflow/compiler/xla/service/gpu/cudnn_conv_rewriter.cc +++ b/tensorflow/compiler/xla/service/gpu/cudnn_conv_rewriter.cc @@ -77,7 +77,11 @@ bool CanImplementAsCudnnForwardConv(HloInstruction* conv) { return false; } - if (window_util::HasWindowReversal(conv->window())) { + // CuDNN can perform either cross correlation (no reversal), + // or convolution (all dimensions reversed). + if (dnums.input_spatial_dimensions_size() == 2 + ? !window_util::AllOrNoneReversed(conv->window()) + : window_util::HasWindowReversal(conv->window())) { return false; } return true; @@ -167,6 +171,8 @@ std::tuple MatchBackwardFilter( // The window's low padding is the same as the low padding of the // activations. dim->set_padding_low(conv->window().dimensions(i).padding_low()); + dim->set_base_dilation(1); + dim->set_window_dilation(1); int64 input_size = conv->operand(0)->shape().dimensions(input_spatial_dims[i]); @@ -252,7 +258,7 @@ MatchBackwardInput(HloInstruction* conv) { const auto no_match_result = std::make_tuple(false, Window(), ConvolutionDimensionNumbers(), nullptr); - // TODO(b/31709653): Theoretically cuDNN supports grouped convolutions also + // TODO(b/119479517): Theoretically cuDNN supports grouped convolutions also // for the backward input convolution, but at least for now with version 7.1.4 // it is slower. This needs to be re-evaluated for future cuDNN versions. // Note that we already have the necessary code down below, the only thing to @@ -335,6 +341,7 @@ MatchBackwardInput(HloInstruction* conv) { // = the base dilation factor of the forward convolution auto dim = new_window.mutable_dimensions(i); dim->set_stride(old_window.dimensions(i).base_dilation()); + dim->set_base_dilation(1); // The low padding = kernel_size - 1 - low padding on the gradients // Make sure the low padding is not negative. diff --git a/tensorflow/compiler/xla/service/gpu/cudnn_conv_rewriter_test.cc b/tensorflow/compiler/xla/service/gpu/cudnn_conv_rewriter_test.cc index a6980850af370645389f1e10922097f6a16cdee9..dbcdc2b075bc72f3194af8e555faabb1511376e0 100644 --- a/tensorflow/compiler/xla/service/gpu/cudnn_conv_rewriter_test.cc +++ b/tensorflow/compiler/xla/service/gpu/cudnn_conv_rewriter_test.cc @@ -24,7 +24,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/shape_inference.h" #include "tensorflow/compiler/xla/test.h" #include "tensorflow/compiler/xla/test_helpers.h" -#include "tensorflow/compiler/xla/tests/hlo_verified_test_base.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" #include "tensorflow/core/platform/test.h" namespace xla { @@ -34,11 +34,11 @@ namespace { namespace op = xla::testing::opcode_matchers; using ::testing::_; -class CudnnConvRewriterTest : public HloVerifiedTestBase { +class CudnnConvRewriterTest : public HloTestBase { public: CudnnConvRewriterTest() - : HloVerifiedTestBase(/*layout_sensitive=*/true, - /*allow_mixed_precision=*/false) { + : HloTestBase(/*layout_sensitive=*/true, + /*allow_mixed_precision=*/false) { for (int i = 0; i < 2; ++i) { WindowDimension* window_dim = default_conv_window_.add_dimensions(); window_dim->set_size(1); @@ -109,19 +109,21 @@ TEST_F(CudnnConvRewriterTest, BackwardFilterConvolve) { auto* conv = builder.AddInstruction(HloInstruction::CreateConvolve( ShapeInference::InferConvolveShape( activations->shape(), gradients->shape(), /*feature_group_count=*/1, - conv_window, tf_default_dnums_for_backward_filter_) + /*batch_group_count=*/1, conv_window, + tf_default_dnums_for_backward_filter_) .ConsumeValueOrDie(), - activations, gradients, /*feature_group_count=*/1, conv_window, + activations, gradients, /*feature_group_count=*/1, + /*batch_group_count=*/1, conv_window, tf_default_dnums_for_backward_filter_, DefaultPrecisionConfig(2))); OpMetadata metadata; metadata.set_op_name("foo"); conv->set_metadata(metadata); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* entry_computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(RunPass(module)); + EXPECT_TRUE(RunPass(module.get())); ASSERT_THAT(entry_computation->root_instruction(), op::GetTupleElement( op::CustomCall(kCudnnConvBackwardFilterCallTarget), 0)); @@ -147,15 +149,17 @@ TEST_F(CudnnConvRewriterTest, builder.AddInstruction(HloInstruction::CreateConvolve( ShapeInference::InferConvolveShape( activations->shape(), gradients->shape(), /*feature_group_count=*/1, - conv_window, tf_default_dnums_for_backward_filter_) + /*batch_group_count=*/1, conv_window, + tf_default_dnums_for_backward_filter_) .ConsumeValueOrDie(), - activations, gradients, /*feature_group_count=*/1, conv_window, + activations, gradients, /*feature_group_count=*/1, + /*batch_group_count=*/1, conv_window, tf_default_dnums_for_backward_filter_, DefaultPrecisionConfig(2))); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* entry_computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(RunPass(module)); + EXPECT_TRUE(RunPass(module.get())); EXPECT_THAT(entry_computation->root_instruction(), op::GetTupleElement( op::CustomCall(kCudnnConvBackwardFilterCallTarget), 0)); @@ -179,13 +183,13 @@ TEST_F(CudnnConvRewriterTest, BackwardFilterConvolveWithPaddedActivations) { } builder.AddInstruction(HloInstruction::CreateConvolve( ShapeUtil::MakeShape(F32, {32, 3, 3, 32}), activations, gradients, - /*feature_group_count=*/1, conv_window, + /*feature_group_count=*/1, /*batch_group_count=*/1, conv_window, tf_default_dnums_for_backward_filter_, DefaultPrecisionConfig(2))); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* entry_computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(RunPass(module)); + EXPECT_TRUE(RunPass(module.get())); EXPECT_THAT(entry_computation->root_instruction(), op::GetTupleElement( op::CustomCall(kCudnnConvBackwardFilterCallTarget), 0)); @@ -209,13 +213,13 @@ TEST_F(CudnnConvRewriterTest, BackwardFilterConvolveWithPaddedGradients) { } builder.AddInstruction(HloInstruction::CreateConvolve( ShapeUtil::MakeShape(F32, {320, 3, 3, 192}), activations, gradients, - /*feature_group_count=*/1, conv_window, + /*feature_group_count=*/1, /*batch_group_count=*/1, conv_window, tf_default_dnums_for_backward_filter_, DefaultPrecisionConfig(2))); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* entry_computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(RunPass(module)); + EXPECT_TRUE(RunPass(module.get())); EXPECT_THAT(entry_computation->root_instruction(), op::GetTupleElement( op::CustomCall(kCudnnConvBackwardFilterCallTarget), 0)); @@ -238,13 +242,13 @@ TEST_F(CudnnConvRewriterTest, BackwardFilterConvolveWithUnevenPadding) { } builder.AddInstruction(HloInstruction::CreateConvolve( ShapeUtil::MakeShape(F32, {32, 2, 2, 32}), activations, gradients, - /*feature_group_count=*/1, conv_window, + /*feature_group_count=*/1, /*batch_group_count=*/1, conv_window, tf_default_dnums_for_backward_filter_, DefaultPrecisionConfig(2))); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* entry_computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(RunPass(module)); + EXPECT_TRUE(RunPass(module.get())); EXPECT_THAT(entry_computation->root_instruction(), op::GetTupleElement( op::CustomCall(kCudnnConvBackwardFilterCallTarget), 0)); @@ -283,19 +287,21 @@ TEST_F(CudnnConvRewriterTest, BackwardInputConvolveEvenPadding) { HloInstruction* conv = builder.AddInstruction(HloInstruction::CreateConvolve( ShapeUtil::MakeShape(F32, {4, 3, 16, 16}), /*lhs=*/output, - /*rhs=*/reverse_kernel, /*feature_group_count=*/1, conv_window, - conv_dnums, DefaultPrecisionConfig(2))); + /*rhs=*/reverse_kernel, /*feature_group_count=*/1, + /*batch_group_count=*/1, conv_window, conv_dnums, + DefaultPrecisionConfig(2))); // Verify the convolution's shape is consistent with ShapeInference. CHECK(ShapeUtil::Compatible( conv->shape(), ShapeInference::InferConvolveShape( output->shape(), reverse_kernel->shape(), - /*feature_group_count=*/1, conv_window, conv_dnums) + /*feature_group_count=*/1, /*batch_group_count=*/1, + conv_window, conv_dnums) .ValueOrDie())); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* entry_computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(RunPass(module)); + EXPECT_TRUE(RunPass(module.get())); ASSERT_THAT(entry_computation->root_instruction(), op::GetTupleElement( @@ -309,6 +315,7 @@ TEST_F(CudnnConvRewriterTest, BackwardInputConvolveEvenPadding) { EXPECT_EQ(3, window_dim.padding_low()); EXPECT_EQ(3, window_dim.padding_high()); EXPECT_EQ(1, window_dim.stride()); + EXPECT_EQ(1, window_dim.base_dilation()); } } @@ -331,16 +338,18 @@ TEST_F(CudnnConvRewriterTest, BackwardInputConvolve1x1Filter) { builder.AddInstruction(HloInstruction::CreateConvolve( ShapeInference::InferConvolveShape(output->shape(), kernel->shape(), - /*feature_group_count=*/1, conv_window, + /*feature_group_count=*/1, + /*batch_group_count=*/1, conv_window, tf_default_dnums_for_backward_input_) .ConsumeValueOrDie(), - /*lhs=*/output, /*rhs=*/kernel, /*feature_group_count=*/1, conv_window, + /*lhs=*/output, /*rhs=*/kernel, /*feature_group_count=*/1, + /*batch_group_count=*/1, conv_window, tf_default_dnums_for_backward_input_, DefaultPrecisionConfig(2))); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* entry_computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(RunPass(module)); + EXPECT_TRUE(RunPass(module.get())); EXPECT_THAT(entry_computation->root_instruction(), op::GetTupleElement( op::CustomCall(kCudnnConvBackwardInputCallTarget), 0)); @@ -364,16 +373,17 @@ TEST_F(CudnnConvRewriterTest, builder.AddInstruction(HloInstruction::CreateConvolve( ShapeInference::InferConvolveShape( output->shape(), kernel->shape(), /*feature_group_count=*/1, - default_conv_window_, tf_default_dnums_for_backward_input_) + /*batch_group_count=*/1, default_conv_window_, + tf_default_dnums_for_backward_input_) .ConsumeValueOrDie(), /*lhs=*/output, /*rhs=*/kernel, /*feature_group_count=*/1, - default_conv_window_, tf_default_dnums_for_backward_input_, - DefaultPrecisionConfig(2))); + /*batch_group_count=*/1, default_conv_window_, + tf_default_dnums_for_backward_input_, DefaultPrecisionConfig(2))); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* entry_computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(RunPass(module)); + EXPECT_TRUE(RunPass(module.get())); EXPECT_THAT( entry_computation->root_instruction(), op::GetTupleElement(op::CustomCall(kCudnnConvForwardCallTarget), 0)); @@ -414,20 +424,20 @@ TEST_F(CudnnConvRewriterTest, BackwardInputConvolveUnevenPaddingOnGradients) { } HloInstruction* conv = builder.AddInstruction(HloInstruction::CreateConvolve( ShapeUtil::MakeShape(F32, {20, 10, 10, 192}), output, reverse_kernel, - /*feature_group_count=*/1, conv_window, + /*feature_group_count=*/1, /*batch_group_count=*/1, conv_window, tf_default_dnums_for_backward_input_, DefaultPrecisionConfig(2))); // Verify the convolution's shape is consistent with ShapeInference. CHECK(ShapeUtil::Compatible( - conv->shape(), - ShapeInference::InferConvolveShape( - output->shape(), reverse_kernel->shape(), /*feature_group_count=*/1, - conv_window, tf_default_dnums_for_backward_input_) - .ValueOrDie())); + conv->shape(), ShapeInference::InferConvolveShape( + output->shape(), reverse_kernel->shape(), + /*feature_group_count=*/1, /*batch_group_count=*/1, + conv_window, tf_default_dnums_for_backward_input_) + .ValueOrDie())); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* entry_computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(RunPass(module)); + EXPECT_TRUE(RunPass(module.get())); ASSERT_THAT(entry_computation->root_instruction(), op::GetTupleElement( op::CustomCall(kCudnnConvBackwardInputCallTarget), 0)); @@ -438,6 +448,7 @@ TEST_F(CudnnConvRewriterTest, BackwardInputConvolveUnevenPaddingOnGradients) { EXPECT_EQ(0, window_dim.padding_low()); EXPECT_EQ(0, window_dim.padding_high()); EXPECT_EQ(2, window_dim.stride()); + EXPECT_EQ(1, window_dim.base_dilation()); } } @@ -463,26 +474,26 @@ TEST_F(CudnnConvRewriterTest, BackwardInputConvolveLowPaddingTooLarge) { } HloInstruction* conv = builder.AddInstruction(HloInstruction::CreateConvolve( ShapeUtil::MakeShape(F32, {20, 10, 10, 192}), output, reverse_kernel, - /*feature_group_count=*/1, conv_window, + /*feature_group_count=*/1, /*batch_group_count=*/1, conv_window, tf_default_dnums_for_backward_input_, DefaultPrecisionConfig(2))); // Verify the convolution's shape is consistent with ShapeInference. CHECK(ShapeUtil::Compatible( - conv->shape(), - ShapeInference::InferConvolveShape( - output->shape(), reverse_kernel->shape(), /*feature_group_count=*/1, - conv_window, tf_default_dnums_for_backward_input_) - .ValueOrDie())); + conv->shape(), ShapeInference::InferConvolveShape( + output->shape(), reverse_kernel->shape(), + /*feature_group_count=*/1, /*batch_group_count=*/1, + conv_window, tf_default_dnums_for_backward_input_) + .ValueOrDie())); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* entry_computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(RunPass(module)); + EXPECT_TRUE(RunPass(module.get())); EXPECT_THAT( entry_computation->root_instruction(), op::GetTupleElement(op::CustomCall(kCudnnConvForwardCallTarget), 0)); } -// Extracted from //learning/brain/google/xla/benchmarks/resnet.py +// Extracted from Resnet-50. // // For simplicity, we focus on the column dimension and ignore other dimensions. // We use [?] to represent the shape instead of the content. @@ -517,20 +528,20 @@ TEST_F(CudnnConvRewriterTest, BackwardInputConvolveUnevenPaddingOnActivations) { forward_conv_col_dim->set_base_dilation(2); HloInstruction* conv = builder.AddInstruction(HloInstruction::CreateConvolve( ShapeUtil::MakeShape(F32, {1, 1, 14, 1}), output, reverse_kernel, - /*feature_group_count=*/1, conv_window, + /*feature_group_count=*/1, /*batch_group_count=*/1, conv_window, tf_default_dnums_for_backward_input_, DefaultPrecisionConfig(2))); // Verify the convolution's shape is consistent with ShapeInference. CHECK(ShapeUtil::Compatible( - conv->shape(), - ShapeInference::InferConvolveShape( - output->shape(), reverse_kernel->shape(), /*feature_group_count=*/1, - conv_window, tf_default_dnums_for_backward_input_) - .ValueOrDie())); + conv->shape(), ShapeInference::InferConvolveShape( + output->shape(), reverse_kernel->shape(), + /*feature_group_count=*/1, /*batch_group_count=*/1, + conv_window, tf_default_dnums_for_backward_input_) + .ValueOrDie())); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); const HloComputation* entry_computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(RunPass(module)); + EXPECT_TRUE(RunPass(module.get())); ASSERT_THAT(entry_computation->root_instruction(), op::GetTupleElement( op::CustomCall(kCudnnConvBackwardInputCallTarget), 0)); @@ -572,20 +583,20 @@ TEST_F(CudnnConvRewriterTest, forward_conv_col_dim->set_padding_high(2); HloInstruction* conv = builder.AddInstruction(HloInstruction::CreateConvolve( ShapeUtil::MakeShape(F32, {1, 1, 4, 1}), output, reverse_kernel, - /*feature_group_count=*/1, conv_window, + /*feature_group_count=*/1, /*batch_group_count=*/1, conv_window, tf_default_dnums_for_backward_input_, DefaultPrecisionConfig(2))); // Verify the convolution's shape is consistent with ShapeInference. CHECK(ShapeUtil::Compatible( - conv->shape(), - ShapeInference::InferConvolveShape( - output->shape(), reverse_kernel->shape(), /*feature_group_count=*/1, - conv_window, tf_default_dnums_for_backward_input_) - .ValueOrDie())); + conv->shape(), ShapeInference::InferConvolveShape( + output->shape(), reverse_kernel->shape(), + /*feature_group_count=*/1, /*batch_group_count=*/1, + conv_window, tf_default_dnums_for_backward_input_) + .ValueOrDie())); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* entry_computation = module->AddEntryComputation(builder.Build()); - EXPECT_TRUE(RunPass(module)); + EXPECT_TRUE(RunPass(module.get())); EXPECT_THAT( entry_computation->root_instruction(), op::GetTupleElement(op::CustomCall(kCudnnConvForwardCallTarget), 0)); @@ -597,8 +608,9 @@ TEST_F(CudnnConvRewriterTest, BackwardInputConvolveConstantFilter) { Array4D constant_arr(4, 4, 2, 2); constant_arr.FillIota(0); string constant_str = - LiteralUtil::CreateR4FromArray4D(constant_arr).ToString(); - ParseAndVerifyModule(absl::StrFormat(R"( + LiteralUtil::CreateR4FromArray4D(constant_arr).ToStringWithoutShape(); + + const string module_str = absl::StrFormat(R"( HloModule test ENTRY entry_computation { @@ -608,10 +620,12 @@ TEST_F(CudnnConvRewriterTest, BackwardInputConvolveConstantFilter) { window={size=4x4 pad=2_2x2_2 lhs_dilate=2x2}, dim_labels=bf01_01oi->bf01, feature_group_count=1 })", - constant_str)); - EXPECT_TRUE(RunPass(&module())); + constant_str); + TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(module_str)); + + EXPECT_TRUE(RunPass(m.get())); EXPECT_THAT( - module().entry_computation()->root_instruction(), + m->entry_computation()->root_instruction(), op::GetTupleElement(op::CustomCall(kCudnnConvBackwardInputCallTarget, _, op::Reverse(op::Constant())), 0)); diff --git a/tensorflow/compiler/xla/service/gpu/cudnn_conv_runner.cc b/tensorflow/compiler/xla/service/gpu/cudnn_conv_runner.cc index 0b4fdf71623e1597168c6873a0d2b60176e518ce..b628f27f4b2ba8ccf17fd531d8a0c25cb99d9396 100644 --- a/tensorflow/compiler/xla/service/gpu/cudnn_conv_runner.cc +++ b/tensorflow/compiler/xla/service/gpu/cudnn_conv_runner.cc @@ -126,9 +126,9 @@ Status RunCudnnConvImpl(CudnnConvParams params, int64 feature_group_count = params.feature_group_count; AlgorithmConfig algorithm = params.algorithm; - VLOG(3) << "Convolution Algorithm: " << algorithm.algorithm().algo_id(); + VLOG(3) << "Convolution Algorithm: " << algorithm.algorithm()->algo_id(); VLOG(3) << "tensor_ops_enabled: " - << algorithm.algorithm().tensor_ops_enabled(); + << algorithm.algorithm()->tensor_ops_enabled(); VLOG(3) << "Convolution kind: " << CudnnConvKindToString(kind); VLOG(3) << "input shape: " << ShapeUtil::HumanStringWithLayout(input_shape); VLOG(3) << "filter shape: " << ShapeUtil::HumanStringWithLayout(filter_shape); @@ -138,6 +138,7 @@ Status RunCudnnConvImpl(CudnnConvParams params, const int num_dimensions = window.dimensions_size(); CHECK_LE(num_dimensions, 3); + CHECK_GE(num_dimensions, 1); // cuDNN does not support 1D convolutions. We therefore express 1D // convolutions as 2D convolutions where the first spatial dimension is 1. // This matches the behavior of TF (see definition of conv1d in @@ -148,11 +149,22 @@ Status RunCudnnConvImpl(CudnnConvParams params, output_shape.element_type()) << ShapeUtil::HumanString(output_shape); + // If one dimension is reversed, we need to have all dimensions reversed (so + // we're doing convolution not cross correlation). + const bool dims_reversed = window.dimensions()[0].window_reversal(); + CHECK_EQ(num_dimensions, dnums.input_spatial_dimensions_size()); CHECK_EQ(num_dimensions, dnums.kernel_spatial_dimensions_size()); CHECK_EQ(num_dimensions, dnums.output_spatial_dimensions_size()); for (const WindowDimension& dim : window.dimensions()) { + CHECK_EQ(dims_reversed, dim.window_reversal()); CHECK_EQ(dim.padding_low(), dim.padding_high()); + CHECK_EQ(dim.base_dilation(), 1) + << "cudnn does not support base dilation; it " + "must be made explicit with a kPad"; + CHECK_EQ(dim.window_dilation(), 1) + << "XLA does not support window dilation (although cudnn does); it " + "must be made explicit with a kPad"; } // cuDNN's convolution APIs support the BDYX layout for activations/output and @@ -192,6 +204,7 @@ Status RunCudnnConvImpl(CudnnConvParams params, ConvolutionDescriptor convolution_descriptor(effective_num_dimensions); convolution_descriptor.set_group_count(feature_group_count); + convolution_descriptor.set_convolution_not_crosscorr(dims_reversed); for (int dim = 0; dim < num_dimensions; ++dim) { convolution_descriptor .set_zero_padding( @@ -296,8 +309,8 @@ Status RunCudnnConvImpl(CudnnConvParams params, if (!stream->ok()) { return InternalError( "Unable to launch convolution with type %s and algorithm (%d, %d)", - CudnnConvKindToString(kind), algorithm.algorithm().algo_id(), - algorithm.algorithm_no_scratch().algo_id()); + CudnnConvKindToString(kind), algorithm.algorithm()->algo_id(), + algorithm.algorithm_no_scratch()->algo_id()); } return Status::OK(); } @@ -357,14 +370,12 @@ StatusOr GetCudnnConvParams( params.output_shape = &conv_result_shape; params.fusion.emplace(); auto& fusion = *params.fusion; - if (backend_config.activation_mode() < - static_cast(se::dnn::ActivationMode::kNumActivationModes)) { - fusion.mode = static_cast( - backend_config.activation_mode()); - } else { + if (!se::dnn::ActivationMode_IsValid(backend_config.activation_mode())) { return InternalError("Bad activation mode: %s", backend_config.ShortDebugString()); } + fusion.mode = static_cast( + backend_config.activation_mode()); fusion.side_input_scale = backend_config.side_input_scale(); params.input_buf = operand_buffers[0]; params.filter_buf = operand_buffers[1]; @@ -384,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 6dcdaf1cfe06e446deed847aaf29088a7ed10e13..2ab754a471070d5f90a3eaebd0600ff180d2fe5d 100644 --- a/tensorflow/compiler/xla/service/gpu/elemental_ir_emitter.cc +++ b/tensorflow/compiler/xla/service/gpu/elemental_ir_emitter.cc @@ -161,6 +161,16 @@ StatusOr GpuElementalIrEmitter::EmitFloatBinaryOp( PrimitiveType lhs_input_type = op->operand(0)->shape().element_type(); PrimitiveType rhs_input_type = op->operand(1)->shape().element_type(); PrimitiveType output_type = op->shape().element_type(); + HloOpcode opcode = op->opcode(); + + if (hlo_module_config_.debug_options().xla_gpu_enable_fast_min_max() && + (opcode == HloOpcode::kMaximum || opcode == HloOpcode::kMinimum)) { + return llvm_ir::EmitCallToIntrinsic( + opcode == HloOpcode::kMaximum ? llvm::Intrinsic::maxnum + : llvm::Intrinsic::minnum, + {lhs_value, rhs_value}, {lhs_value->getType()}, b_); + } + switch (op->opcode()) { case HloOpcode::kRemainder: { return EmitLibdeviceMathCall("__nv_fmod", {lhs_value, rhs_value}, diff --git a/tensorflow/compiler/xla/service/gpu/fusion_merger.cc b/tensorflow/compiler/xla/service/gpu/fusion_merger.cc index 30c1f9088968305ad0207164ecb07ba13cc89ee6..91930eccdff94bb2fc85636f3a4b2d661c618d87 100644 --- a/tensorflow/compiler/xla/service/gpu/fusion_merger.cc +++ b/tensorflow/compiler/xla/service/gpu/fusion_merger.cc @@ -35,7 +35,7 @@ namespace { // Traverses users of tuple shape, adding leaf instructions to 'instructions'. void MaybeResolveTupleElements(HloInstruction* instruction, std::vector* instructions) { - if (ShapeUtil::IsTuple(instruction->shape())) { + if (instruction->shape().IsTuple()) { for (auto tuple_user : instruction->users()) { MaybeResolveTupleElements(tuple_user, instructions); } @@ -229,7 +229,7 @@ Status FusionInstructionMerger::HandleFusion(HloInstruction* fusion) { if (!absl::c_all_of(fusion->users(), [&](const HloInstruction* user) { return user->opcode() == HloOpcode::kFusion && (user->fusion_kind() == HloInstruction::FusionKind::kLoop || - (user->fusion_kind() == HloInstruction::FusionKind::kInput && + (IsReduceInputFusion(*user) && LayoutsAreReduceInputFusionFriendly(*fusion, *user))); })) { VLOG(3) << "Not merging " << fusion->name() diff --git a/tensorflow/compiler/xla/service/gpu/gemm_thunk.cc b/tensorflow/compiler/xla/service/gpu/gemm_thunk.cc index 9c4a4903667ea1a6c99ce9e912c9d0497b8e389f..86c9bc6a345047fb5329af0be45c8981cc427f50 100644 --- a/tensorflow/compiler/xla/service/gpu/gemm_thunk.cc +++ b/tensorflow/compiler/xla/service/gpu/gemm_thunk.cc @@ -18,6 +18,7 @@ limitations under the License. #include #include "absl/strings/str_cat.h" +#include "tensorflow/compiler/xla/service/gpu/ir_emission_utils.h" #include "tensorflow/compiler/xla/util.h" #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/stream_executor_no_cuda.h" @@ -51,7 +52,8 @@ struct MatrixDescriptor { // rhs_matrix, and stores the result to output_matrix. template bool DoGemm(MatrixDescriptor lhs_matrix, MatrixDescriptor rhs_matrix, - MatrixDescriptor output_matrix, double alpha, se::Stream* stream) { + MatrixDescriptor output_matrix, double alpha, double beta, + se::Stream* stream) { DCHECK(!output_matrix.transpose); const int64 batch_size = lhs_matrix.batch_size; @@ -73,7 +75,7 @@ bool DoGemm(MatrixDescriptor lhs_matrix, MatrixDescriptor rhs_matrix, lhs_transpose, rhs_transpose, output_matrix.num_rows, output_matrix.num_cols, /*size of reduce dim=*/k, /*alpha=*/alpha, lhs_data, /*leading dim of LHS=*/lhs_matrix.num_rows, rhs_data, - /*leading dim of RHS=*/rhs_matrix.num_rows, /*beta=*/0.0, + /*leading dim of RHS=*/rhs_matrix.num_rows, /*beta=*/beta, &output_data, /*leading dim of output=*/output_matrix.num_rows) .ok(); } @@ -88,7 +90,7 @@ bool DoGemm(MatrixDescriptor lhs_matrix, MatrixDescriptor rhs_matrix, /*alpha=*/alpha, lhs_data, /*leading dim of LHS=*/lhs_matrix.num_rows, lhs_stride, rhs_data, /*leading dim of RHS=*/rhs_matrix.num_rows, rhs_stride, - /*beta=*/0.0, &output_data, + /*beta=*/beta, &output_data, /*leading dim of output=*/output_matrix.num_rows, output_stride, batch_size) .ok(); @@ -112,6 +114,7 @@ template bool DoGemmWithAlgorithm(MatrixDescriptor lhs_matrix, MatrixDescriptor rhs_matrix, MatrixDescriptor output_matrix, double alpha, + double beta, se::blas::ComputationType computation_type, se::blas::AlgorithmType algorithm, se::Stream* stream, se::blas::ProfileResult* output_profile_result) { @@ -138,7 +141,7 @@ bool DoGemmWithAlgorithm(MatrixDescriptor lhs_matrix, /*alpha=*/static_cast(alpha), lhs_data, /*leading dim of LHS=*/lhs_matrix.num_rows, rhs_data, /*leading dim of RHS=*/rhs_matrix.num_rows, - /*beta=*/static_cast(0.0f), &output_data, + /*beta=*/static_cast(beta), &output_data, /*leading dim of output=*/output_matrix.num_rows, computation_type, algorithm, output_profile_result) .ok(); @@ -153,7 +156,7 @@ bool DoGemmWithAlgorithm(MatrixDescriptor lhs_matrix, template StatusOr DoGemmAutotune( MatrixDescriptor lhs_matrix, MatrixDescriptor rhs_matrix, - MatrixDescriptor output_matrix, double alpha, + MatrixDescriptor output_matrix, double alpha, double beta, se::blas::ComputationType computation_type, se::Stream* stream) { std::vector algorithms; CHECK(stream->parent()->GetBlasGemmAlgorithms(&algorithms)); @@ -166,7 +169,7 @@ StatusOr DoGemmAutotune( // non-null ProfileResult, DoGemmWithAlgorithm should always return true, // and the actual success-ness is returned in ProfileResult::is_valid. CHECK(DoGemmWithAlgorithm(lhs_matrix, rhs_matrix, output_matrix, - alpha, computation_type, algorithm, + alpha, beta, computation_type, algorithm, stream, &profile_result)); if (profile_result.is_valid()) { @@ -203,6 +206,8 @@ auto GetGemmFn(PrimitiveType type) -> decltype(&DoGemm) { return &DoGemm; case C64: return &DoGemm>; + case C128: + return &DoGemm>; default: LOG(FATAL) << "Unsupported type."; } @@ -218,6 +223,8 @@ auto GetGemmWithAlgorithmFn(PrimitiveType type) return &DoGemmWithAlgorithm; case C64: return &DoGemmWithAlgorithm>; + case C128: + return &DoGemmWithAlgorithm>; default: LOG(FATAL) << "Unsupported type."; } @@ -232,6 +239,8 @@ auto GetGemmAutotuneFn(PrimitiveType type) -> decltype(&DoGemmAutotune) { return &DoGemmAutotune; case C64: return &DoGemmAutotune>; + case C128: + return &DoGemmAutotune>; default: LOG(FATAL) << "Unsupported type."; } @@ -252,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."; } @@ -263,8 +274,9 @@ DotDimensionNumbers GetDimensionNumbers(const HloInstruction& hlo_instruction) { } CHECK_EQ(hlo_instruction.opcode(), HloOpcode::kFusion); CHECK_EQ(hlo_instruction.fusion_kind(), HloInstruction::FusionKind::kOutput); - CHECK_EQ(hlo_instruction.fused_expression_root()->opcode(), - HloOpcode::kMultiply); + CHECK(hlo_instruction.fused_expression_root()->opcode() == HloOpcode::kAdd || + hlo_instruction.fused_expression_root()->opcode() == + HloOpcode::kMultiply); // Try to find the dot inside the output fusion node. const HloInstruction* dot = hlo_instruction.fused_expression_root()->operand(0); @@ -282,8 +294,9 @@ GemmThunk::GemmThunk(const BufferAllocation::Slice& lhs_buffer, const BufferAllocation::Slice& rhs_buffer, const BufferAllocation::Slice& output_buffer, const Shape& lhs_shape, const Shape& rhs_shape, - const Shape& output_shape, double alpha, - const HloInstruction* hlo_instruction) + const Shape& output_shape, double alpha, double beta, + const HloInstruction* hlo_instruction, + bool implements_whole_instruction) : Thunk(Kind::kGemm, hlo_instruction), lhs_buffer_(lhs_buffer), rhs_buffer_(rhs_buffer), @@ -291,7 +304,9 @@ GemmThunk::GemmThunk(const BufferAllocation::Slice& lhs_buffer, lhs_shape_(lhs_shape), rhs_shape_(rhs_shape), output_shape_(output_shape), - alpha_(alpha) {} + alpha_(alpha), + beta_(beta), + implements_whole_instruction_(implements_whole_instruction) {} Status GemmThunk::ExecuteOnStream(const BufferAllocations& buffer_allocations, se::Stream* stream, @@ -308,8 +323,7 @@ Status GemmThunk::ExecuteOnStream(const BufferAllocations& buffer_allocations, DotDimensionNumbers dim_nums = GetDimensionNumbers(*hlo_instruction()); CHECK_EQ(dim_nums.lhs_batch_dimensions_size(), dim_nums.rhs_batch_dimensions_size()); - CHECK_EQ(dim_nums.lhs_batch_dimensions_size() + 2, - ShapeUtil::Rank(output_shape_)); + CHECK_EQ(dim_nums.lhs_batch_dimensions_size() + 2, output_shape_.rank()); int64 row_dim = dim_nums.lhs_batch_dimensions_size(); int64 col_dim = dim_nums.lhs_batch_dimensions_size() + 1; @@ -386,7 +400,7 @@ Status GemmThunk::ExecuteOnStream(const BufferAllocations& buffer_allocations, // TODO(b/112111608): Implement auto tune for batched gemm. if (batch_size != 1) { return GetGemmFn(element_type)(lhs_matrix, rhs_matrix, output_matrix, - alpha_, stream); + alpha_, beta_, stream); } auto thunk_name = [&] { @@ -398,9 +412,27 @@ Status GemmThunk::ExecuteOnStream(const BufferAllocations& buffer_allocations, auto autotune_it = autotune_results_.find(device_name); if (autotune_it == autotune_results_.end()) { VLOG(3) << "Starting autotune of GemmThunk " << thunk_name(); - StatusOr best_algorithm = - GetGemmAutotuneFn(element_type)(lhs_matrix, rhs_matrix, output_matrix, - alpha_, computation_type, stream); + + // If the output buffer already contains a bias then autotune into a + // scratch buffer. This avoids overwriting the bias buffer. The scratch + // buffer may contain arbitrary garbage values. + se::DeviceMemoryBase scratch_data = output_data; + std::unique_ptr> scratch_mem; + if (beta_ != 0.0) { + auto temp_status = stream->AllocateTemporaryArray( + ShapeUtil::ByteSizeOf(output_shape_)); + if (!temp_status.ok()) { + return false; + } + scratch_mem = std::move(temp_status).ValueOrDie(); + scratch_data = scratch_mem->device_memory(); + } + const MatrixDescriptor scratch_descriptor( + scratch_data, false, output_num_cols, output_num_rows, batch_size); + + StatusOr best_algorithm = GetGemmAutotuneFn( + element_type)(lhs_matrix, rhs_matrix, scratch_descriptor, alpha_, + beta_, computation_type, stream); autotune_it = autotune_results_.insert({device_name, best_algorithm}).first; @@ -421,18 +453,19 @@ Status GemmThunk::ExecuteOnStream(const BufferAllocations& buffer_allocations, VLOG(2) << "Using algorithm " << algorithm << " chosen by autotuning on GemmThunk " << thunk_name(); return GetGemmWithAlgorithmFn(element_type)( - lhs_matrix, rhs_matrix, output_matrix, alpha_, computation_type, - algorithm, stream, + lhs_matrix, rhs_matrix, output_matrix, alpha_, beta_, + computation_type, algorithm, stream, /*output_profile_result=*/nullptr); } // Autotune will fail when CUDA 8 and GPU sm_50 or older are used. // Use the older Gemm API in this case. return GetGemmFn(element_type)(lhs_matrix, rhs_matrix, output_matrix, - alpha_, stream); + alpha_, beta_, stream); }; - auto op_profiler = profiler->MakeScopedInstructionProfiler(hlo_instruction()); + auto op_profiler = profiler->MakeScopedInstructionProfiler( + implements_whole_instruction_ ? hlo_instruction() : nullptr); bool launch_ok; if (LayoutUtil::Minor(output_shape_.layout(), row_dim) == 0) { launch_ok = launch(lhs_descriptor, rhs_descriptor, diff --git a/tensorflow/compiler/xla/service/gpu/gemm_thunk.h b/tensorflow/compiler/xla/service/gpu/gemm_thunk.h index 12c81f9bfc6bfdac63edf9c826b835057107fa41..cc2d12a39c045fc081292dcf53053f6613d3d9ef 100644 --- a/tensorflow/compiler/xla/service/gpu/gemm_thunk.h +++ b/tensorflow/compiler/xla/service/gpu/gemm_thunk.h @@ -41,8 +41,9 @@ class GemmThunk : public Thunk { const BufferAllocation::Slice& rhs_buffer, const BufferAllocation::Slice& output_buffer, const Shape& lhs_shape, const Shape& rhs_shape, - const Shape& output_shape, double alpha, - const HloInstruction* hlo_instruction); + const Shape& output_shape, double alpha, double beta, + const HloInstruction* hlo_instruction, + bool implements_whole_instruction); GemmThunk(const GemmThunk&) = delete; GemmThunk& operator=(const GemmThunk&) = delete; @@ -70,6 +71,9 @@ class GemmThunk : public Thunk { const Shape output_shape_; const double alpha_; + const double beta_; + + const bool implements_whole_instruction_; // Maps device names (StreamExecutor::DeviceDescription::name()) to autotune // results. The map's value is the best algorithm we've found for this thunk diff --git a/tensorflow/compiler/xla/service/gpu/gpu_executable.cc b/tensorflow/compiler/xla/service/gpu/gpu_executable.cc index 57426327822d95a42f407ed7488f35acfd3623d2..434060ad89dac7ad65c790c8c0a7f3d6ad62a25a 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_executable.cc +++ b/tensorflow/compiler/xla/service/gpu/gpu_executable.cc @@ -51,7 +51,7 @@ GpuExecutable::GpuExecutable( const string& ptx, const std::vector& cubin, std::pair compute_capability, std::unique_ptr thunk_schedule, - std::unique_ptr hlo_module, + std::unique_ptr hlo_module, std::unique_ptr assignment, std::unique_ptr hlo_profile_printer_data, std::unique_ptr hlo_profile_index_map) @@ -218,7 +218,7 @@ GpuExecutable::ResolveConstantGlobals(se::StreamExecutor* executor) { const Literal& literal = llvm_ir::LiteralForConstantAllocation(allocation); - CHECK(ShapeUtil::IsArray(literal.shape())); + CHECK(literal.shape().IsArray()); if (!ShouldEmitLiteralInLlvmIr(literal)) { VLOG(3) << "H2D memcpy for constant with shape " << ShapeUtil::HumanString(literal.shape()); @@ -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_executable.h b/tensorflow/compiler/xla/service/gpu/gpu_executable.h index 0e276282e40fba0ae4881a51dad0c7c9e8d1c081..2b3c77f5b82aa94f44d8de56caf0f4d31c05e0cb 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_executable.h +++ b/tensorflow/compiler/xla/service/gpu/gpu_executable.h @@ -54,7 +54,7 @@ class GpuExecutable : public Executable { GpuExecutable(const string& ptx, const std::vector& cubin, std::pair compute_capability, std::unique_ptr thunk_schedule, - std::unique_ptr hlo_module, + std::unique_ptr hlo_module, std::unique_ptr assignment, std::unique_ptr hlo_profile_printer_data, std::unique_ptr hlo_profile_index_map); diff --git a/tensorflow/compiler/xla/service/gpu/gpu_fusible.cc b/tensorflow/compiler/xla/service/gpu/gpu_fusible.cc index 2d31fd5570c468b0c42fa308535fd335f3588a79..842ba2fdcd31a451cec1be543e102e0a46077f38 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_fusible.cc +++ b/tensorflow/compiler/xla/service/gpu/gpu_fusible.cc @@ -42,20 +42,18 @@ bool LayoutsAreReduceInputFusionFriendly(const HloInstruction& producer, int64 max_rank = -1; const Layout* max_rank_layout; for (HloInstruction* param : params) { - if (ShapeUtil::IsArray(param->shape()) && - ShapeUtil::Rank(param->shape()) > max_rank) { - max_rank = ShapeUtil::Rank(param->shape()); + if (param->shape().IsArray() && param->shape().rank() > max_rank) { + max_rank = param->shape().rank(); max_rank_layout = ¶m->shape().layout(); } } return absl::c_all_of(params, [&](HloInstruction* param) { - return (!ShapeUtil::IsArray(param->shape())) || - (ShapeUtil::Rank(param->shape()) < max_rank) || + return (!param->shape().IsArray()) || (param->shape().rank() < max_rank) || (LayoutUtil::Equal(param->shape().layout(), *max_rank_layout)); }); } -bool IsInputFusibleReduction(const HloInstruction& instr) { +bool IsReduceInputFusion(const HloInstruction& instr) { if (instr.IsMultiOutputFusion()) { for (const HloInstruction* operand : instr.fused_expression_root()->operands()) { @@ -67,17 +65,70 @@ bool IsInputFusibleReduction(const HloInstruction& instr) { return true; } } - return false; - } else if (instr.opcode() == HloOpcode::kFusion) { - if (IsReductionToVector(*instr.fused_expression_root())) { - CHECK(instr.fusion_kind() == HloInstruction::FusionKind::kInput) - << " Fusion rooted at reduction-to-vector op must be of kind kInput: " - << instr.ToString(); - return true; + } else if (instr.opcode() == HloOpcode::kFusion && + IsReductionToVector(*instr.fused_expression_root())) { + CHECK(instr.fusion_kind() == HloInstruction::FusionKind::kInput) + << " Fusion rooted at reduction-to-vector op must be of kind kInput: " + << instr.ToString(); + return true; + } + return false; +} + +bool IsInputFusibleReduction(const HloInstruction& instr) { + return IsReduceInputFusion(instr) || IsReductionToVector(instr); +} + +bool ShapesCompatibleForMultiOutputFusion(const HloInstruction& instr1, + const HloInstruction& instr2) { + // Returns the instructions that determines the emitter used for lowering, + // sometimes referred to as "the real hero". + auto get_real_hero = + [&](const HloInstruction* instr) -> const HloInstruction* { + if (instr->opcode() == HloOpcode::kFusion) { + auto fused_expression_root = instr->fused_expression_root(); + if (instr->IsMultiOutputFusion()) { + // If possible, we want to pick a reduction-to-vector operand of the + // fusion root, because it has the most constraints. + for (const auto* inst : fused_expression_root->operands()) { + if (IsReductionToVector(*inst)) { + return inst; + } + } + return fused_expression_root->operands()[0]; + } + return fused_expression_root; } + return instr; + }; + + // Multi-output fusion kernels share a common parallel loop. The loop + // dimenstions are determined by instruction shapes. + auto get_loop_shape = [&](const HloInstruction* element_instr) { + // Special-case reduction-to-vector ops: The loop dimensions are determined + // by the shape of the first operand. + if (IsReductionToVector(*element_instr)) { + return element_instr->operand(0)->shape(); + } + return element_instr->shape(); + }; + + // All shapes of the root tuple of multi-output fusions should agree, i.e. all + // root ops should have equal output shapes. An exception are + // reduction-to-vector ops. Here the input shapes of the reduction (first + // operand shape) and the reduction dimensions need to match. + auto* instr_1 = get_real_hero(&instr1); + auto* instr_2 = get_real_hero(&instr2); + // TODO(tjoerg): Relax the shape constraint. The datatype does not matter. + if (IsReductionToVector(*instr_1) && IsReductionToVector(*instr_2) && + (!ShapeUtil::Equal(instr_1->shape(), instr_2->shape()) || + instr_1->dimensions() != instr_2->dimensions())) { return false; } - return IsReductionToVector(instr); + // The elementwise output shapes must be the same (including layout). + // TODO(tjoerg): Further relax the constraint. The datatype does not matter. + return ShapeUtil::EqualIgnoringFpPrecision(get_loop_shape(instr_1), + get_loop_shape(instr_2)); } } // namespace gpu diff --git a/tensorflow/compiler/xla/service/gpu/gpu_fusible.h b/tensorflow/compiler/xla/service/gpu/gpu_fusible.h index f7c24a0d5bbfcc61389ea19ae7f769671e4e974d..e9d7ba1c4cfa865532a0d06c2ed883a2fea4e2cd 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_fusible.h +++ b/tensorflow/compiler/xla/service/gpu/gpu_fusible.h @@ -33,16 +33,29 @@ namespace gpu { bool LayoutsAreReduceInputFusionFriendly(const HloInstruction& producer, const HloInstruction& reduce); -// Whether `instr` is fusible as root of a reduce input fusions, i.e. `instr` -// is either an unfused reduction-to-vector op, an input fusion rooted at a -// reduction-to-vector op, or a multi-output input fusion with at least one -// reduction-to-vector op root. // Note that reduction ops are lowered in different ways. Reduce input fusions // are lowered by IrEmitterUnnested::EmitReductionToVector and must be rooted at // reduction-to-vector ops. Other reduction ops are lowered by // GpuElementalIrEmitter and fused like elementwise ops. + +// Whether `instr` is an input fusion rooted at a reduction-to-vector op or a +// multi-output input fusion with at least one reduction-to-vector op root. +bool IsReduceInputFusion(const HloInstruction& instr); + +// Whether `instr` is fusible as root of a reduce input fusions, i.e. `instr` +// is either an unfused reduction-to-vector op or a reduce input fusion. bool IsInputFusibleReduction(const HloInstruction& instr); +// Whether instruction shapes are compatible for multi-output fusion, i.e. +// whether the emitters support lowering the resulting fusion. +// This function works for both, sibling and producer-conumser multi-output +// fusion. +// So far, multi-output fusion is supported for loop fusions and reduce +// input fusions only. It is up to the caller to ensure the instructions +// themselves are fusible! +bool ShapesCompatibleForMultiOutputFusion(const HloInstruction& instr1, + const HloInstruction& instr2); + } // namespace gpu } // namespace xla diff --git a/tensorflow/compiler/xla/service/gpu/gpu_fusible_test.cc b/tensorflow/compiler/xla/service/gpu/gpu_fusible_test.cc index d91b7bc61fda5a07c163a07ec0e1644d2ad9db49..15d4ee206ce8debcb8a5dbc6ec65d29ba257d302 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_fusible_test.cc +++ b/tensorflow/compiler/xla/service/gpu/gpu_fusible_test.cc @@ -178,7 +178,7 @@ TEST_F(GpuFusibleTest, EXPECT_TRUE(LayoutsAreReduceInputFusionFriendly(*loop_fusion, *reduce)); } -TEST_F(GpuFusibleTest, IsInputFusibleReduction_ReductionToVector) { +TEST_F(GpuFusibleTest, IsReduceInputFusion_ReductionToVector) { auto module = ParseHloString(absl::StrCat(kModulePrefix, R"( ENTRY entry { c0 = f32[] parameter(0) @@ -191,10 +191,11 @@ TEST_F(GpuFusibleTest, IsInputFusibleReduction_ReductionToVector) { const HloInstruction* reduce = module->entry_computation()->root_instruction(); ASSERT_EQ(reduce->opcode(), HloOpcode::kReduce); + EXPECT_FALSE(IsReduceInputFusion(*reduce)); EXPECT_TRUE(IsInputFusibleReduction(*reduce)); } -TEST_F(GpuFusibleTest, IsInputFusibleReduction_ElementalReduction) { +TEST_F(GpuFusibleTest, IsReduceInputFusion_ElementalReduction) { auto module = ParseHloString(absl::StrCat(kModulePrefix, R"( ENTRY entry { c0 = f32[] parameter(0) @@ -207,10 +208,11 @@ TEST_F(GpuFusibleTest, IsInputFusibleReduction_ElementalReduction) { const HloInstruction* reduce = module->entry_computation()->root_instruction(); ASSERT_EQ(reduce->opcode(), HloOpcode::kReduce); + EXPECT_FALSE(IsReduceInputFusion(*reduce)); EXPECT_FALSE(IsInputFusibleReduction(*reduce)); } -TEST_F(GpuFusibleTest, IsInputFusibleReduction_SingleOutputInputReduceFusion) { +TEST_F(GpuFusibleTest, IsReduceInputFusion_SingleOutputInputReduceFusion) { auto module = ParseHloString(absl::StrCat(kModulePrefix, R"( fused_reduction { c0 = f32[] parameter(0) @@ -225,10 +227,11 @@ TEST_F(GpuFusibleTest, IsInputFusibleReduction_SingleOutputInputReduceFusion) { const HloInstruction* reduce = module->entry_computation()->root_instruction(); ASSERT_EQ(reduce->opcode(), HloOpcode::kFusion); + EXPECT_TRUE(IsReduceInputFusion(*reduce)); EXPECT_TRUE(IsInputFusibleReduction(*reduce)); } -TEST_F(GpuFusibleTest, IsInputFusibleReduction_SingleOutputLoopReduceFusion) { +TEST_F(GpuFusibleTest, IsReduceInputFusion_SingleOutputLoopReduceFusion) { auto module = ParseHloString(absl::StrCat(kModulePrefix, R"( fused_reduction { c0 = f32[] parameter(0) @@ -243,10 +246,11 @@ TEST_F(GpuFusibleTest, IsInputFusibleReduction_SingleOutputLoopReduceFusion) { const HloInstruction* reduce = module->entry_computation()->root_instruction(); ASSERT_EQ(reduce->opcode(), HloOpcode::kFusion); + EXPECT_FALSE(IsReduceInputFusion(*reduce)); EXPECT_FALSE(IsInputFusibleReduction(*reduce)); } -TEST_F(GpuFusibleTest, IsInputFusibleReduction_MultiOutputInputReduceFusion) { +TEST_F(GpuFusibleTest, IsReduceInputFusion_MultiOutputInputReduceFusion) { auto module = ParseHloString(absl::StrCat(kModulePrefix, R"( fused_reduction { c0 = f32[] parameter(0) @@ -263,11 +267,12 @@ TEST_F(GpuFusibleTest, IsInputFusibleReduction_MultiOutputInputReduceFusion) { const HloInstruction* reduce = module->entry_computation()->root_instruction(); ASSERT_EQ(reduce->opcode(), HloOpcode::kFusion); + EXPECT_TRUE(IsReduceInputFusion(*reduce)); EXPECT_TRUE(IsInputFusibleReduction(*reduce)); } TEST_F(GpuFusibleTest, - IsInputFusibleReduction_MultiOutputInputReduceFusionWithExtraOutputs) { + IsReduceInputFusion_MultiOutputInputReduceFusionWithExtraOutputs) { auto module = ParseHloString(absl::StrCat(kModulePrefix, R"( fused_reduction { c0 = f32[] parameter(0) @@ -284,10 +289,11 @@ TEST_F(GpuFusibleTest, const HloInstruction* reduce = module->entry_computation()->root_instruction(); ASSERT_EQ(reduce->opcode(), HloOpcode::kFusion); + EXPECT_TRUE(IsReduceInputFusion(*reduce)); EXPECT_TRUE(IsInputFusibleReduction(*reduce)); } -TEST_F(GpuFusibleTest, IsInputFusibleReduction_MultiOutputLoopReduceFusion) { +TEST_F(GpuFusibleTest, IsReduceInputFusion_MultiOutputLoopReduceFusion) { auto module = ParseHloString(absl::StrCat(kModulePrefix, R"( fused_reduction { c0 = f32[] parameter(0) @@ -304,11 +310,12 @@ TEST_F(GpuFusibleTest, IsInputFusibleReduction_MultiOutputLoopReduceFusion) { const HloInstruction* reduce = module->entry_computation()->root_instruction(); ASSERT_EQ(reduce->opcode(), HloOpcode::kFusion); + EXPECT_FALSE(IsReduceInputFusion(*reduce)); EXPECT_FALSE(IsInputFusibleReduction(*reduce)); } TEST_F(GpuFusibleTest, - IsInputFusibleReduction_MultiOutputLoopFusionReduceAndElementwiseOp) { + IsReduceInputFusion_MultiOutputLoopFusionReduceAndElementwiseOp) { auto module = ParseHloString(absl::StrCat(kModulePrefix, R"( fused_reduction { c0 = f32[] parameter(0) @@ -325,8 +332,304 @@ TEST_F(GpuFusibleTest, const HloInstruction* reduce = module->entry_computation()->root_instruction(); ASSERT_EQ(reduce->opcode(), HloOpcode::kFusion); + EXPECT_FALSE(IsReduceInputFusion(*reduce)); EXPECT_FALSE(IsInputFusibleReduction(*reduce)); } +TEST_F(GpuFusibleTest, ShapesCompatibleForMultiOutputFusion_LoopFusions) { + auto module = ParseHloString(absl::StrCat(kModulePrefix, R"( + fused_computation_1 { + p0.1 = f32[6400]{0} parameter(0) + ROOT mul = f32[6400]{0} multiply(p0.1, p0.1) + } + + fused_computation_2 { + p0.2 = f32[6400]{0} parameter(0) + const.2 = f32[] constant(1) + ROOT div = f32[6400]{0} divide(p0.2, const.2) + } + + ENTRY entry { + p0 = f32[6400]{0} parameter(0) + fusion.1 = f32[6400]{0} fusion(p0), kind=kLoop, calls=fused_computation_1 + fusion.2 = f32[6400]{0} fusion(p0), kind=kLoop, calls=fused_computation_2 + ROOT root = (f32[6400]{0}, f32[6400]{0}) tuple(fusion.1, fusion.2) + })")) + .ValueOrDie(); + const HloInstruction* fusion_1 = + module->entry_computation()->root_instruction()->operand(0); + const HloInstruction* fusion_2 = + module->entry_computation()->root_instruction()->operand(1); + EXPECT_TRUE(ShapesCompatibleForMultiOutputFusion(*fusion_1, *fusion_2)); +} + +TEST_F(GpuFusibleTest, ShapesCompatibleForMultiOutputFusion_IgnoreFpPrecision) { + auto module = ParseHloString(absl::StrCat(kModulePrefix, R"( + fused_computation_1 { + p0.1 = f32[6400]{0} parameter(0) + ROOT mul = f32[6400]{0} multiply(p0.1, p0.1) + } + + fused_computation_2 { + p0.2 = f32[6400]{0} parameter(0) + ROOT convert = f16[6400]{0} convert(p0.2) + } + + ENTRY entry { + p0 = f32[6400]{0} parameter(0) + fusion.1 = f32[6400]{0} fusion(p0), kind=kLoop, calls=fused_computation_1 + fusion.2 = f32[6400]{0} fusion(p0), kind=kLoop, calls=fused_computation_2 + ROOT root = (f32[6400]{0}, f32[6400]{0}) tuple(fusion.1, fusion.2) + })")) + .ValueOrDie(); + const HloInstruction* fusion_1 = + module->entry_computation()->root_instruction()->operand(0); + const HloInstruction* fusion_2 = + module->entry_computation()->root_instruction()->operand(1); + EXPECT_TRUE(ShapesCompatibleForMultiOutputFusion(*fusion_1, *fusion_2)); +} + +TEST_F(GpuFusibleTest, ShapesCompatibleForMultiOutputFusion_Reduce) { + auto module = ParseHloString(absl::StrCat(kModulePrefix, R"( + fused_computation_1 { + p0.1 = f32[6400]{0} parameter(0) + ROOT mul = f32[6400]{0} multiply(p0.1, p0.1) + } + + ENTRY entry { + p0 = f32[6400]{0} parameter(0) + fusion.1 = f32[6400]{0} fusion(p0), kind=kLoop, calls=fused_computation_1 + const.2 = f32[] constant(0) + reduce = f32[] reduce(p0, const.2), dimensions={0}, to_apply=scalar_add + ROOT root = (f32[6400]{0}, f32[]) tuple(fusion.1, reduce) + })")) + .ValueOrDie(); + const HloInstruction* fusion = + module->entry_computation()->root_instruction()->operand(0); + const HloInstruction* reduce = + module->entry_computation()->root_instruction()->operand(1); + EXPECT_TRUE(ShapesCompatibleForMultiOutputFusion(*fusion, *reduce)); +} + +TEST_F(GpuFusibleTest, ShapesCompatibleForMultiOutputFusion_Elementwise) { + auto module = ParseHloString(absl::StrCat(kModulePrefix, R"( + fused_computation_1 { + p0.1 = f32[6400]{0} parameter(0) + ROOT mul = f32[6400]{0} multiply(p0.1, p0.1) + } + + ENTRY entry { + p0 = f32[6400]{0} parameter(0) + fusion.1 = f32[6400]{0} fusion(p0), kind=kLoop, calls=fused_computation_1 + const.2 = f32[] constant(1) + div = f32[6400]{0} divide(p0, const.2) + ROOT root = (f32[6400]{0}, f32[6400]{0}) tuple(fusion.1, div) + })")) + .ValueOrDie(); + const HloInstruction* fusion = + module->entry_computation()->root_instruction()->operand(0); + const HloInstruction* div = + module->entry_computation()->root_instruction()->operand(1); + EXPECT_TRUE(ShapesCompatibleForMultiOutputFusion(*fusion, *div)); +} + +TEST_F(GpuFusibleTest, + ShapesCompatibleForMultiOutputFusion_MultiOutputLoopFusion) { + auto module = ParseHloString(absl::StrCat(kModulePrefix, R"( + fused_computation_1 { + p0.1 = f32[8,1,5,16,1,1]{5,4,3,2,1,0} parameter(0) + mul = f32[8,1,5,16,1,1]{5,4,3,2,1,0} multiply(p0.1, p0.1) + exp = f32[8,1,5,16,1,1]{5,4,3,2,1,0} exponential(p0.1) + ROOT tuple = (f32[8,1,5,16,1,1]{5,4,3,2,1,0}, f32[8,1,5,16,1,1]{5,4,3,2,1,0}) tuple(mul, exp) + } + + fused_computation_2 { + p0.2 = f32[8,1,5,16,1,1]{5,4,3,2,1,0} parameter(0) + const.2 = f32[] constant(0) + ROOT add = f32[8,1,5,16,1,1]{5,4,3,2,1,0} add(p0.2, const.2) + } + + ENTRY entry { + p0 = f32[8,1,5,16,1,1]{5,4,3,2,1,0} parameter(0) + fusion.1 = (f32[8,1,5,16,1,1]{5,4,3,2,1,0}, f32[8,1,5,16,1,1]{5,4,3,2,1,0}) fusion(p0), kind=kLoop, calls=fused_computation_1 + fusion.2 = f32[8,1,5,16,1,1]{5,4,3,2,1,0} fusion(p0), kind=kLoop, calls=fused_computation_2 + gte0 = f32[8,1,5,16,1,1]{5,4,3,2,1,0} get-tuple-element(fusion.1), index=0 + gte1 = f32[8,1,5,16,1,1]{5,4,3,2,1,0} get-tuple-element(fusion.1), index=1 + ROOT root = (f32[8,1,5,16,1,1]{5,4,3,2,1,0}, f32[8,1,5,16,1,1]{5,4,3,2,1,0}, f32[8,1,5,16,1,1]{5,4,3,2,1,0}) tuple(gte0, gte1, fusion.2) + })")) + .ValueOrDie(); + const HloInstruction* fusion_1 = + module->entry_computation()->root_instruction()->operand(0)->operand(0); + const HloInstruction* fusion_2 = + module->entry_computation()->root_instruction()->operand(1)->operand(0); + EXPECT_TRUE(ShapesCompatibleForMultiOutputFusion(*fusion_1, *fusion_2)); +} + +TEST_F(GpuFusibleTest, ShapesCompatibleForMultiOutputFusion_UnfusedOps) { + auto module = ParseHloString(absl::StrCat(kModulePrefix, R"( + ENTRY reduce { + p0 = f32[2,2,2]{2,1,0} parameter(0) + c0 = f32[] constant(0) + exp = f32[2,2,2]{2,1,0} exponential(p0) + reduce = f32[2,2]{1,0} reduce(exp, c0), dimensions={2}, to_apply=scalar_add + ROOT root = (f32[2,2]{1,0}, f32[2,2,2]{2,1,0}) tuple(reduce, exp) + })")) + .ValueOrDie(); + const HloInstruction* reduce = + module->entry_computation()->root_instruction()->operand(0); + const HloInstruction* exp = + module->entry_computation()->root_instruction()->operand(1); + EXPECT_TRUE(ShapesCompatibleForMultiOutputFusion(*reduce, *exp)); +} + +TEST_F(GpuFusibleTest, ShapesCompatibleForMultiOutputFusion_DifferentLayouts) { + auto module = ParseHloString(absl::StrCat(kModulePrefix, R"( + ENTRY reduce { + p0 = f32[2,2,2]{2,1,0} parameter(0) + p1 = f32[2,2,2]{0,1,2} parameter(1) + c0 = f32[] constant(0) + exp = f32[2,2,2]{2,1,0} exponential(p0) + reduce = f32[2,2]{0,1} reduce(p1, c0), dimensions={2}, to_apply=scalar_add + ROOT root = (f32[2,2]{0,1}, f32[2,2,2]{2,1,0}) tuple(reduce, exp) + })")) + .ValueOrDie(); + const HloInstruction* reduce = + module->entry_computation()->root_instruction()->operand(0); + const HloInstruction* exp = + module->entry_computation()->root_instruction()->operand(1); + EXPECT_FALSE(ShapesCompatibleForMultiOutputFusion(*reduce, *exp)); +} + +TEST_F(GpuFusibleTest, + ShapesCompatibleForMultiOutputFusion_MultiOutputReduceFusion) { + auto module = ParseHloString(absl::StrCat(kModulePrefix, R"( + fused_select { + p1.1 = f32[2,2,2]{2,1,0} parameter(1) + c0 = f32[] constant(0) + broadcast = f32[2,2,2]{2,1,0} broadcast(f32[] c0), dimensions={} + greater-than = pred[2,2,2]{2,1,0} greater-than(f32[2,2,2]{2,1,0} p1.1, f32[2,2,2]{2,1,0} broadcast) + p0.1 = f32[2,2,2]{2,1,0} parameter(0) + ROOT select = f32[2,2,2]{2,1,0} select(pred[2,2,2]{2,1,0} greater-than, f32[2,2,2]{2,1,0} p0.1, f32[2,2,2]{2,1,0} broadcast) + } + + fused_reduce { + p0.2 = f32[2,2,2]{2,1,0} parameter(0) + c1 = f32[] constant(0) + r1 = f32[2,2]{1,0} reduce(p0.2, c1), dimensions={2}, to_apply=scalar_add + mul = f32[2,2,2]{2,1,0} multiply(p0.2, p0.2) + r2 = f32[2,2]{1,0} reduce(mul, c1), dimensions={2}, to_apply=scalar_add + ROOT tuple = (f32[2,2]{1,0}, f32[2,2]{1,0}) tuple(r1, r2) + } + + ENTRY reduce { + p0 = f32[2,2,2]{2,1,0} parameter(0) + p1 = f32[2,2,2]{2,1,0} parameter(1) + select = f32[2,2,2]{2,1,0} fusion(p0, p1), kind=kLoop, calls=fused_select + fusion = (f32[2,2]{1,0}, f32[2,2]{1,0}) fusion(select), kind=kInput, calls=fused_reduce + gte0 = f32[2,2]{1,0} get-tuple-element(fusion), index=0 + gte1 = f32[2,2]{1,0} get-tuple-element(fusion), index=1 + ROOT root = (f32[2,2]{1,0}, f32[2,2]{1,0}, f32[2,2,2]{2,1,0}) tuple(gte1, gte1, select) + })")) + .ValueOrDie(); + const HloInstruction* fusion_1 = + module->entry_computation()->root_instruction()->operand(0)->operand(0); + const HloInstruction* fusion_2 = + module->entry_computation()->root_instruction()->operand(1)->operand(0); + EXPECT_TRUE(ShapesCompatibleForMultiOutputFusion(*fusion_1, *fusion_2)); +} + +TEST_F(GpuFusibleTest, ShapesCompatibleForMultiOutputFusion_ReduceFusions) { + auto module = ParseHloString(absl::StrCat(kModulePrefix, R"( + fused_reduce_1 { + p0.1 = f32[2,2,2]{2,1,0} parameter(0) + c0 = f32[] constant(0) + ROOT reduce = f32[2,2]{1,0} reduce(f32[2,2,2]{2,1,0} p0.1, f32[] c0), dimensions={0}, to_apply=scalar_add + } + + fused_reduce_2 { + p0.2 = f32[2,2,2]{2,1,0} parameter(0) + mul = f32[2,2,2]{2,1,0} multiply(f32[2,2,2]{2,1,0} p0.2, f32[2,2,2]{2,1,0} p0.2) + c1 = f32[] constant(0) + ROOT reduce = f32[2,2]{1,0} reduce(f32[2,2,2]{2,1,0} mul, f32[] c1), dimensions={0}, to_apply=scalar_add + } + + ENTRY reduce { + p0 = f32[2,2,2]{2,1,0} parameter(0) + p1 = f32[2,2,2]{2,1,0} parameter(1) + reduce_1 = f32[2,2]{1,0} fusion(p0), kind=kLoop, calls=fused_reduce_1 + reduce_2 = f32[2,2]{1,0} fusion(p1), kind=kLoop, calls=fused_reduce_2 + ROOT root = (f32[2,2]{1,0}, f32[2,2,2]{2,1,0}) tuple(reduce_1, reduce_2) + })")) + .ValueOrDie(); + const HloInstruction* fusion_1 = + module->entry_computation()->root_instruction()->operand(0); + const HloInstruction* fusion_2 = + module->entry_computation()->root_instruction()->operand(1); + EXPECT_TRUE(ShapesCompatibleForMultiOutputFusion(*fusion_1, *fusion_2)); +} + +TEST_F(GpuFusibleTest, + ShapesCompatibleForMultiOutputFusion_DifferentReduceDimensions) { + auto module = ParseHloString(absl::StrCat(kModulePrefix, R"( + fused_reduce_1 { + p0.1 = f32[2,2,2]{2,1,0} parameter(0) + c0 = f32[] constant(0) + ROOT reduce = f32[2,2]{1,0} reduce(f32[2,2,2]{2,1,0} p0.1, f32[] c0), dimensions={0}, to_apply=scalar_add + } + + fused_reduce_2 { + p0.2 = f32[2,2,2]{2,1,0} parameter(0) + mul = f32[2,2,2]{2,1,0} multiply(f32[2,2,2]{2,1,0} p0.2, f32[2,2,2]{2,1,0} p0.2) + c1 = f32[] constant(0) + ROOT reduce = f32[2,2]{1,0} reduce(f32[2,2,2]{2,1,0} mul, f32[] c1), dimensions={2}, to_apply=scalar_add + } + + ENTRY reduce { + p0 = f32[2,2,2]{2,1,0} parameter(0) + p1 = f32[2,2,2]{2,1,0} parameter(1) + reduce_1 = f32[2,2]{1,0} fusion(p0), kind=kLoop, calls=fused_reduce_1 + reduce_2 = f32[2,2]{1,0} fusion(p1), kind=kLoop, calls=fused_reduce_2 + ROOT root = (f32[2,2]{1,0}, f32[2,2,2]{2,1,0}) tuple(reduce_1, reduce_2) + })")) + .ValueOrDie(); + const HloInstruction* fusion_1 = + module->entry_computation()->root_instruction()->operand(0); + const HloInstruction* fusion_2 = + module->entry_computation()->root_instruction()->operand(1); + EXPECT_FALSE(ShapesCompatibleForMultiOutputFusion(*fusion_1, *fusion_2)); +} + +TEST_F(GpuFusibleTest, + ShapesCompatibleForMultiOutputFusion_NoReductionToVector) { + auto module = ParseHloString(absl::StrCat(kModulePrefix, R"( + fused_element_wise { + p0.1 = f32[2,2,2]{2,1,0} parameter(0) + p1.1 = f32[2,2,2]{2,1,0} parameter(1) + ROOT add = f32[2,2,2]{2,1,0} add(p0.1, p1.1) + } + + fused_reduce { + p0.2 = f32[2,2,2]{2,1,0} parameter(0) + mul = f32[2,2,2]{2,1,0} multiply(f32[2,2,2]{2,1,0} p0.2, f32[2,2,2]{2,1,0} p0.2) + c1 = f32[] constant(0) + // Note that reduce is not a reduction-to-vector. + ROOT reduce = f32[2,2]{1,0} reduce(f32[2,2,2]{2,1,0} mul, f32[] c1), dimensions={1}, to_apply=scalar_add + } + + ENTRY reduce { + p0 = f32[2,2,2]{2,1,0} parameter(0) + p1 = f32[2,2,2]{2,1,0} parameter(1) + element_wise = f32[2,2,2]{2,1,0} fusion(p0, p1), kind=kLoop, calls=fused_element_wise + fusion = (f32[2,2]{1,0}, f32[2,2]{1,0}) fusion(element_wise), kind=kLoop, calls=fused_reduce + ROOT root = (f32[2,2]{1,0}, f32[2,2,2]{2,1,0}) tuple(fusion, element_wise) + })")) + .ValueOrDie(); + const HloInstruction* fusion_1 = + module->entry_computation()->root_instruction()->operand(0); + const HloInstruction* fusion_2 = + module->entry_computation()->root_instruction()->operand(1); + EXPECT_FALSE(ShapesCompatibleForMultiOutputFusion(*fusion_1, *fusion_2)); +} + } // namespace gpu } // namespace xla diff --git a/tensorflow/compiler/xla/service/gpu/gpu_hlo_schedule.cc b/tensorflow/compiler/xla/service/gpu/gpu_hlo_schedule.cc index 02a0d028c118aba23996f9b97d05443bb4a00c88..1126943624a3771433ecac591545d335c1890115 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_hlo_schedule.cc +++ b/tensorflow/compiler/xla/service/gpu/gpu_hlo_schedule.cc @@ -37,12 +37,12 @@ class GpuHloOrdering : public PredecessorHloOrdering { public: GpuHloOrdering(const HloModule* module, const StreamAssignment& stream_assignment, - const std::vector& thunk_launch_order); + const std::vector& thunk_launch_order); ~GpuHloOrdering() override = default; // Only the entry computation can possibly be sequentially ordered, and only // if we've assigned all instructions to a single stream. - const std::vector* SequentialOrder( + const HloInstructionSequence* SequentialOrder( const HloComputation& computation) const override { return &computation == module_->entry_computation() ? entry_sequence_.get() : nullptr; @@ -51,17 +51,17 @@ class GpuHloOrdering : public PredecessorHloOrdering { string ToString() const override { return ToStringHelper("GpuHloOrdering"); } private: - std::unique_ptr> entry_sequence_; + std::unique_ptr entry_sequence_; }; GpuHloOrdering::GpuHloOrdering( const HloModule* module, const StreamAssignment& stream_assignment, - const std::vector& thunk_launch_order) + const std::vector& thunk_launch_order) : PredecessorHloOrdering(module) { // The entry computation has a total order when there's only one stream. if (stream_assignment.StreamCount() == 1) { - entry_sequence_ = absl::make_unique>( - thunk_launch_order); + entry_sequence_ = + absl::make_unique(thunk_launch_order); } // The ordering of instructions for the entry computation is determined by the @@ -124,7 +124,8 @@ GpuHloOrdering::GpuHloOrdering( for (auto* computation : module->computations()) { if (computation != module->entry_computation() && !computation->IsFusionComputation()) { - predecessors_.emplace(computation, computation->ComputeReachability()); + predecessors_.emplace(computation, + HloReachabilityMap::Build(computation)); } } } @@ -149,7 +150,7 @@ GpuHloOrdering::GpuHloOrdering( // However, if the total order is A,B,D,C,E, then C and E can run // concurrently. void BFSLaunchOrder(const HloComputation* computation, - std::vector* launch_order) { + std::vector* launch_order) { // This topological sort uses two data structures: // 1. `incoming_edge_count` which keeps track of the number of incoming // edges to each HLO; @@ -157,9 +158,9 @@ void BFSLaunchOrder(const HloComputation* computation, // // The sorting algorithm repeatedly pops the top from the queue and deletes // that HLO from the graph, making more HLOs incoming-edge free. - std::deque queue; + std::deque queue; std::unordered_map incoming_edge_count; - for (const auto& hlo : computation->instructions()) { + for (auto* hlo : computation->instructions()) { if (hlo->operand_count() == 0) { queue.push_back(hlo); } else { @@ -171,10 +172,10 @@ void BFSLaunchOrder(const HloComputation* computation, } while (!queue.empty()) { - const HloInstruction* x = queue.front(); + HloInstruction* x = queue.front(); queue.pop_front(); launch_order->push_back(x); - for (const HloInstruction* y : x->users()) { + for (HloInstruction* y : x->users()) { --incoming_edge_count[y]; if (incoming_edge_count[y] == 0) { queue.push_back(y); @@ -194,14 +195,14 @@ StatusOr> GpuHloSchedule::Build( std::unique_ptr schedule(new GpuHloSchedule); // Initialize thunk_launch_order_, the total order of thunk launches. - const HloComputation* entry_computation = module.entry_computation(); + HloComputation* entry_computation = module.entry_computation(); if (stream_assignment.StreamCount() == 1) { // All kernels are launched on a single stream, so there's no loss of // concurrency by optimizing for minimal memory usage. TF_ASSIGN_OR_RETURN( HloInstructionSequence sequence, ScheduleComputation( - *entry_computation, [pointer_size](const BufferValue& buffer) { + entry_computation, [pointer_size](const BufferValue& buffer) { return ShapeUtil::ByteSizeOf(buffer.shape(), pointer_size); })); schedule->thunk_launch_order_ = sequence.instructions(); diff --git a/tensorflow/compiler/xla/service/gpu/gpu_hlo_schedule.h b/tensorflow/compiler/xla/service/gpu/gpu_hlo_schedule.h index 07a7fc67aa555845c3de57e574ab582403ec0490..7f224ffe4f03f8f05b0f1907628d99d9df387770 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_hlo_schedule.h +++ b/tensorflow/compiler/xla/service/gpu/gpu_hlo_schedule.h @@ -46,7 +46,7 @@ class GpuHloSchedule { // Returns the total order of thunk launches, represented in terms of HLO // instructions. - const std::vector& ThunkLaunchOrder() const { + const std::vector& ThunkLaunchOrder() const { return thunk_launch_order_; } @@ -60,7 +60,7 @@ class GpuHloSchedule { private: GpuHloSchedule(); - std::vector thunk_launch_order_; + std::vector thunk_launch_order_; std::unique_ptr hlo_ordering_; }; diff --git a/tensorflow/compiler/xla/service/gpu/gpu_hlo_schedule_test.cc b/tensorflow/compiler/xla/service/gpu/gpu_hlo_schedule_test.cc index b857fa775a76ec999b505a2a64332cc0c54cf00b..91db7151f22fd75b20244878bee86d65acd1d304 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_hlo_schedule_test.cc +++ b/tensorflow/compiler/xla/service/gpu/gpu_hlo_schedule_test.cc @@ -24,16 +24,16 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" #include "tensorflow/compiler/xla/test_helpers.h" -#include "tensorflow/compiler/xla/tests/hlo_verified_test_base.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" #include "tensorflow/compiler/xla/tests/test_utils.h" #include "tensorflow/compiler/xla/types.h" namespace xla { namespace gpu { -class GpuHloScheduleTest : public HloVerifiedTestBase { +class GpuHloScheduleTest : public HloTestBase { protected: - using HloVec = std::vector; + using HloVec = std::vector; // Pre-canned shapes. Shape f32_2x2_ = ShapeUtil::MakeShape(F32, {2, 2}); @@ -44,7 +44,7 @@ class GpuHloScheduleTest : public HloVerifiedTestBase { .ConsumeValueOrDie(); } - std::unique_ptr CreateNewModule() { + std::unique_ptr CreateNewVerifiedModule() { HloModuleConfig config; auto debug_options = GetDebugOptionsForTest(); debug_options.set_xla_gpu_disable_multi_streaming(false); @@ -79,7 +79,7 @@ TEST_F(GpuHloScheduleTest, SequentialMatMul) { HloInstruction* dot2 = builder.AddInstruction(CreateCanonicalDot(f32_2x2_, dot1, z)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build(dot2)); std::unique_ptr streams = AssignStreams(*module); @@ -139,7 +139,7 @@ TEST_F(GpuHloScheduleTest, SequentialAdd) { HloInstruction* add3 = builder.AddInstruction( HloInstruction::CreateBinary(f32_2x2_, HloOpcode::kAdd, add1, add2)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build(add3)); std::unique_ptr streams = AssignStreams(*module); @@ -209,7 +209,7 @@ TEST_F(GpuHloScheduleTest, ConcurrentMatMul) { HloInstruction* add = builder.AddInstruction(CreateCanonicalDot(f32_2x2_, dot1, dot2)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build(add)); std::unique_ptr streams = AssignStreams(*module); @@ -288,7 +288,7 @@ TEST_F(GpuHloScheduleTest, LatticeMatMul) { HloInstruction* d40 = builder.AddInstruction(CreateCanonicalDot(f32_2x2_, d30, d31)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build(d40)); std::unique_ptr streams = AssignStreams(*module); diff --git a/tensorflow/compiler/xla/service/gpu/gpu_hlo_support_checker_test.cc b/tensorflow/compiler/xla/service/gpu/gpu_hlo_support_checker_test.cc index 7d01eeb02567d710e9de089c7f29ffcc5f959f9a..b511155f85fb24adc1828cbef7f3fb60778ef7ab 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_hlo_support_checker_test.cc +++ b/tensorflow/compiler/xla/service/gpu/gpu_hlo_support_checker_test.cc @@ -16,7 +16,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/gpu/gpu_hlo_support_checker.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/test.h" -#include "tensorflow/compiler/xla/tests/hlo_verified_test_base.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" #include "tensorflow/core/lib/core/error_codes.pb.h" #include "tensorflow/core/lib/core/status_test_util.h" @@ -25,7 +25,7 @@ namespace { using ::testing::HasSubstr; -class GpuHloSupportCheckerTest : public HloVerifiedTestBase { +class GpuHloSupportCheckerTest : public HloTestBase { protected: GpuHloSupportChecker& checker() { return checker_; } @@ -42,10 +42,10 @@ TEST_F(GpuHloSupportCheckerTest, Add) { HloInstruction::CreateParameter(1, scalar_shape, "param1")); builder.AddInstruction(HloInstruction::CreateBinary( scalar_shape, HloOpcode::kAdd, param0, param1)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); - TF_ASSERT_OK(checker().Run(module).status()); + TF_ASSERT_OK(checker().Run(module.get()).status()); } TEST_F(GpuHloSupportCheckerTest, SparseUnimplemented) { @@ -60,7 +60,7 @@ TEST_F(GpuHloSupportCheckerTest, SparseUnimplemented) { // Since verifier is reporting sparse layouts as errors, we should // use a regular HloModule instead of VerifiedHloModule to avoid // verifier errors being triggered in the destructor. - auto module = HloTestBase::CreateNewModule(); + auto module = CreateNewUnverifiedModule(); module->AddEntryComputation(builder.Build()); Status status = checker().Run(module.get()).status(); diff --git a/tensorflow/compiler/xla/service/gpu/gpu_layout_assignment.cc b/tensorflow/compiler/xla/service/gpu/gpu_layout_assignment.cc index 1c0a23fa3eb38961d420aff05e412c3b4d8524e7..58bdd4209a2315cdb7d29e920faded4d1a6a5876 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_layout_assignment.cc +++ b/tensorflow/compiler/xla/service/gpu/gpu_layout_assignment.cc @@ -65,8 +65,8 @@ HeuristicLayoutAssignment(const HloInstruction* instr, VLOG(2) << "Using heuristic to figure out layouts for " << instr->ToString(); - // Empirically we've found with Volta and cudnn 7 that backward-input convs - // with stride are significantly faster with NCHW layouts. + // Empirically we've found with Volta and cudnn <= 7.3 that backward-input + // convs with stride are significantly faster with NCHW layouts. // // We could have used a mixed layout combination, e.g. (NHWC, NCHW, NCHW), // which on paper gives good performance. However, there are two observations: @@ -75,11 +75,17 @@ HeuristicLayoutAssignment(const HloInstruction* instr, // * we've also observed that for mixed layouts, cuDNN transposes data back // and forth from a different layout combination. If we end up with // transposes anyway, we prefer to have them in XLA, as they can be fused. - // TODO(timshen): Figure out the exact condition. This may be achieved by - // auto-tuning layouts offline. - if (instr->custom_call_target() == kCudnnConvBackwardInputCallTarget && - window_util::HasStride(instr->window())) { - return kAllNCHW; + if (auto* dnn = stream_executor->AsDnn()) { + auto version_status = dnn->GetVersion(); + if (version_status.ok()) { + auto version = version_status.ConsumeValueOrDie(); + if (std::make_tuple(version.major_version(), version.minor_version()) <= + std::make_tuple(7, 3) && + instr->custom_call_target() == kCudnnConvBackwardInputCallTarget && + window_util::HasStride(instr->window())) { + return kAllNCHW; + } + } } // For other Volta f16 convolutions, use NHWC. @@ -190,9 +196,9 @@ Status GpuLayoutAssignment::AddBackendConstraints( CHECK_EQ(dim_nums.lhs_batch_dimensions_size(), dim_nums.rhs_batch_dimensions_size()); CHECK_EQ(dim_nums.lhs_batch_dimensions_size() + 2, - ShapeUtil::Rank(instruction->shape())); + instruction->shape().rank()); for (int64 batch_dim : dim_nums.lhs_batch_dimensions()) { - CHECK_LT(batch_dim, ShapeUtil::Rank(instruction->shape()) - 2); + CHECK_LT(batch_dim, instruction->shape().rank() - 2); } // Set both inputs and the output to default layout. @@ -209,18 +215,18 @@ Status GpuLayoutAssignment::AddBackendConstraints( TF_RETURN_IF_ERROR( constraints->SetInstructionLayout(output_shape, instruction)); } else if (instruction->opcode() == HloOpcode::kSort && - ShapeUtil::Rank(instruction->operand(0)->shape()) > 1) { + instruction->operand(0)->shape().rank() > 1) { // Make sure that all the operands and the output(s) have the same layout. Shape keys_shape = instruction->operand(0)->shape(); Layout keys_layout = - LayoutUtil::GetDefaultLayoutForRank(ShapeUtil::Rank(keys_shape)); + LayoutUtil::GetDefaultLayoutForRank(keys_shape.rank()); for (int64 i = 0; i < instruction->operand_count(); ++i) { Shape shape = instruction->operand(i)->shape(); *shape.mutable_layout() = keys_layout; TF_RETURN_IF_ERROR( constraints->SetOperandLayout(shape, instruction, i)); const LogicalBuffer* output_buffer; - if (ShapeUtil::IsArray(instruction->shape())) { + if (instruction->shape().IsArray()) { TF_ASSIGN_OR_RETURN( output_buffer, constraints->points_to_analysis().GetBufferDefinedAt(instruction, 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 4822b820f4e229336e2b26cfbd0097c8c31a50c8..29756d27260b0f41b2dd4b649ea9b1610ff90268 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_layout_assignment_test.cc +++ b/tensorflow/compiler/xla/service/gpu/gpu_layout_assignment_test.cc @@ -61,7 +61,7 @@ TEST_F(LayoutAssignmentTest, Elementwise) { HloInstruction::CreateParameter(1, ashape, "y")); auto add = builder.AddInstruction( HloInstruction::CreateBinary(ashape, HloOpcode::kAdd, x, y)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* computation = module->AddEntryComputation(builder.Build(add)); @@ -148,7 +148,7 @@ TEST_F(LayoutAssignmentTest, BatchNormInference) { {operand, scale, offset, mean, variance, epsilon, feature_index}, kCudnnBatchNormForwardInferenceCallTarget)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* computation = module->AddEntryComputation(builder.Build(batchnorm)); @@ -217,7 +217,7 @@ TEST_F(LayoutAssignmentTest, BatchNormTraining) { batchnorm_shape, {operand, scale, offset, epsilon, feature_index}, kCudnnBatchNormForwardTrainingCallTarget)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* computation = module->AddEntryComputation(builder.Build(batchnorm)); @@ -298,7 +298,7 @@ TEST_F(LayoutAssignmentTest, BatchNormGrad) { feature_index}, kCudnnBatchNormBackwardCallTarget)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* computation = module->AddEntryComputation(builder.Build(batchnorm)); @@ -369,7 +369,7 @@ TEST_F(LayoutAssignmentTest, SortLayout) { const char* hlo_text = R"( HloModule SortLayout ENTRY sort { - keys = f32[3,2]{0,1} constant(f32[3,2]{0,1}{{0,1},{0,1},{0,1}}) + 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), diff --git a/tensorflow/compiler/xla/service/gpu/gpu_transfer_manager.cc b/tensorflow/compiler/xla/service/gpu/gpu_transfer_manager.cc index f3c274429242d5c989146d14ea523b5910408cff..8c6a6914792a96ab517fa5f20ff2215e4785490e 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_transfer_manager.cc +++ b/tensorflow/compiler/xla/service/gpu/gpu_transfer_manager.cc @@ -59,7 +59,7 @@ Status GpuTransferManager::TransferLiteralToInfeed( TF_RETURN_IF_ERROR(ShapeUtil::ForEachSubshapeWithStatus( shape, [&](const Shape& literal_subshape, const ShapeIndex& index) { - if (ShapeUtil::IsArray(literal_subshape)) { + if (literal_subshape.IsArray()) { int64 tuple_element_size = GetByteSizeRequirement(literal_subshape); TF_ASSIGN_OR_RETURN( *buffer_tree.mutable_element(index), @@ -126,13 +126,12 @@ static void ShapeTreeToLiteral( ShapeTree>* shape_tree, ShapeIndex* index) { const Shape& shape = ShapeUtil::GetSubshape(shape_tree->shape(), *index); - if (ShapeUtil::IsArray(shape)) { + if (shape.IsArray()) { (*shape_tree->mutable_element(*index))->WaitUntilAvailable(); return; } - CHECK(ShapeUtil::IsTuple(shape)) - << ShapeUtil::HumanStringWithLayout(shape); + CHECK(shape.IsTuple()) << ShapeUtil::HumanStringWithLayout(shape); const int64 tuple_element_count = ShapeUtil::TupleElementCount(shape); index->push_back(0); for (int64 i = 0; i < tuple_element_count; ++i) { @@ -158,7 +157,7 @@ Status GpuTransferManager::TransferLiteralFromOutfeed( std::unique_ptr* buffer) { const Shape& shape = ShapeUtil::GetSubshape(literal_shape, index); // Do not transfer tuple index buffers. - if (ShapeUtil::IsTuple(shape)) { + if (shape.IsTuple()) { return; } *buffer = absl::make_unique( 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 51627402b45f594dab3480129ba182d54d01b811..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); @@ -280,7 +281,7 @@ string HloToIrBindings::ToString() const { StrAppend(&s, " ", instr->ToString()); const ShapeTree& shape_tree = it->second; - if (!ShapeUtil::IsTuple(instr->shape())) { + if (!instr->shape().IsTuple()) { const llvm::Value* val = shape_tree.begin()->second; StrAppend(&s, " -> ", llvm_ir::DumpToString(*val), "\n"); continue; 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/infeed_thunk.cc b/tensorflow/compiler/xla/service/gpu/infeed_thunk.cc index 8c3a026740851767855beae59d6a3c92f7a0d6bd..676380c3b10f9a20c641eea0d9a948a26becaddc 100644 --- a/tensorflow/compiler/xla/service/gpu/infeed_thunk.cc +++ b/tensorflow/compiler/xla/service/gpu/infeed_thunk.cc @@ -36,6 +36,21 @@ Status InfeedThunk::ExecuteOnStream(const BufferAllocations& buffer_allocations, ShapeTree infeed_buffers = GetOrCreateInfeedManager()->BlockingGetNextDestination(); + // infeed_slices_'s shape should be a tuple of shape (buffers, token). + const auto& infeed_shape = infeed_slices_.shape(); + TF_RET_CHECK(infeed_shape.IsTuple()) + << ShapeUtil::HumanStringWithLayout(infeed_shape); + TF_RET_CHECK(infeed_shape.tuple_shapes().size() == 2) + << ShapeUtil::HumanStringWithLayout(infeed_shape); + TF_RET_CHECK(infeed_shape.tuple_shapes(1).IsToken()) + << ShapeUtil::HumanStringWithLayout(infeed_shape); + TF_RET_CHECK( + ShapeUtil::Equal(infeed_buffers.shape(), infeed_shape.tuple_shapes(0))) + << "Expected infeed of shape " + << ShapeUtil::HumanStringWithLayout(infeed_shape.tuple_shapes(0)) + << " but was " + << ShapeUtil::HumanStringWithLayout(infeed_buffers.shape()); + { // The infeed buffer has an extra outer tuple with a token. Adjust the index // accordingly. @@ -45,7 +60,7 @@ Status InfeedThunk::ExecuteOnStream(const BufferAllocations& buffer_allocations, const Shape& shape = ShapeUtil::GetSubshape(infeed_buffers.shape(), ShapeIndexView(index, 1)); // For the leaf buffers of the tuple copy the elements directly. - if (ShapeUtil::IsArray(shape)) { + if (shape.IsArray()) { const BufferAllocation::Slice& tuple_element_buffer = infeed_slices_.element(index); se::DeviceMemoryBase tuple_element_address = diff --git a/tensorflow/compiler/xla/service/gpu/instruction_fusion.cc b/tensorflow/compiler/xla/service/gpu/instruction_fusion.cc index 1d66787d8927ad818cbc66d19429c1816fc51748..f07141029cbf8b034b74548f6fca8f1628589f0c 100644 --- a/tensorflow/compiler/xla/service/gpu/instruction_fusion.cc +++ b/tensorflow/compiler/xla/service/gpu/instruction_fusion.cc @@ -47,6 +47,7 @@ bool IsFusible(const HloInstruction& hlo) { hlo.opcode() == HloOpcode::kReduce || hlo.opcode() == HloOpcode::kReduceWindow || hlo.opcode() == HloOpcode::kReshape || + hlo.opcode() == HloOpcode::kReverse || hlo.opcode() == HloOpcode::kScatter || hlo.opcode() == HloOpcode::kSlice || hlo.opcode() == HloOpcode::kTranspose; @@ -79,7 +80,7 @@ bool IsIEEEFloatingPointScalarConstant(const HloInstruction* constant) { // This function limits the maximum number of operands to a fusion. // // There's a cap on how many parameters we can pass to a CUDA kernel, but -// exactly what that limit is is hazy, as it depends on (among other things) how +// exactly what that limit is hazy, as it depends on (among other things) how // much GPU constant memory is in use for other purposes. // // Moreover, we don't even know at the point that we're running fusion how many @@ -179,6 +180,11 @@ bool GpuInstructionFusion::ShouldFuse(HloInstruction* consumer, IsIEEEFloatingPointScalarConstant(alpha->operand(0))) { return true; } + } else if (consumer->operand_count() == 2 && + consumer->opcode() == HloOpcode::kAdd && + consumer->operand(other_operand_index) != producer) { + // Fuse a bias add into the output of the dot. + return true; } } @@ -252,12 +258,17 @@ bool GpuInstructionFusion::ShouldFuse(HloInstruction* consumer, return false; } - // Fuse scalar constants into loop fusion nodes, this reduces the number of + // Fuse scalar constants into loop fusion nodes. This reduces the number of // parameters and makes matching scalar broadcasts easier. - if (ShapeUtil::IsEffectiveScalar(producer->shape()) && - consumer->opcode() == HloOpcode::kFusion && - producer->opcode() == HloOpcode::kConstant) { - return true; + // + // Don't fuse other constants: Unfused constants in GPU land can be + // represented as an external constant (i.e. not emitted in LLVM IR / PTX), + // but fused constants are handled by shrared CPU/GPU code and always emitted + // in the IR/PTX. The external constant representation makes for faster + // compiles and significantly smaller assembly code. + if (producer->opcode() == HloOpcode::kConstant) { + return ShapeUtil::IsEffectiveScalar(producer->shape()) && + consumer->opcode() == HloOpcode::kFusion; } if (!IsFusible(*producer) || !IsFusible(*consumer) || @@ -271,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 fd9b7cee80bdad9a8ed625872ae68bede10200b3..a05ab86cf77a134a1fc387d93cb482aa1ff5345b 100644 --- a/tensorflow/compiler/xla/service/gpu/instruction_fusion_test.cc +++ b/tensorflow/compiler/xla/service/gpu/instruction_fusion_test.cc @@ -41,7 +41,7 @@ TEST_F(InstructionFusionTest, builder.AddInstruction(HloInstruction::CreateBroadcast( ShapeUtil::MakeShape(S32, {1}), exp1, {0})); - auto module = CreateNewModule(); + auto module = CreateNewUnverifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(broadcast2, computation->root_instruction()); EXPECT_FALSE(GpuInstructionFusion(/*may_duplicate=*/true) @@ -61,7 +61,7 @@ TEST_F(InstructionFusionTest, builder.AddInstruction(HloInstruction::CreateBroadcast( ShapeUtil::MakeShape(S32, {1}), negate1, {0})); - auto module = CreateNewModule(); + auto module = CreateNewUnverifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(broadcast2, computation->root_instruction()); EXPECT_TRUE(GpuInstructionFusion(/*may_duplicate=*/true) @@ -80,7 +80,7 @@ TEST_F(InstructionFusionTest, HloInstruction* reshape2 = builder.AddInstruction( HloInstruction::CreateReshape(ShapeUtil::MakeShape(S32, {}), exp1)); - auto module = CreateNewModule(); + auto module = CreateNewUnverifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(reshape2, computation->root_instruction()); EXPECT_TRUE(GpuInstructionFusion(/*may_duplicate=*/true) @@ -99,7 +99,7 @@ TEST_F(InstructionFusionTest, HloInstruction* transpose2 = builder.AddInstruction( HloInstruction::CreateTranspose(ShapeUtil::MakeShape(S32, {}), exp1, {})); - auto module = CreateNewModule(); + auto module = CreateNewUnverifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(transpose2, computation->root_instruction()); EXPECT_TRUE(GpuInstructionFusion(/*may_duplicate=*/true) @@ -117,7 +117,7 @@ TEST_F(InstructionFusionTest, PotentialBitcastReshapeOfDotUnfused) { auto reshape2 = builder.AddInstruction(HloInstruction::CreateReshape( ShapeUtil::MakeShape(S32, {1, 1, 1}), dot1)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(reshape2, computation->root_instruction()); EXPECT_FALSE(GpuInstructionFusion(/*may_duplicate=*/true) @@ -134,7 +134,7 @@ TEST_F(InstructionFusionTest, PotentialBitcastTransposeOfDotUnfused) { auto transpose2 = builder.AddInstruction(HloInstruction::CreateTranspose( ShapeUtil::MakeShape(S32, {1, 1}), dot1, {0, 1})); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(transpose2, computation->root_instruction()); EXPECT_FALSE(GpuInstructionFusion(/*may_duplicate=*/true) @@ -331,6 +331,56 @@ TEST_F(InstructionFusionTest, DotOutputFusion) { op::Broadcast(op::Constant()))); } +TEST_F(InstructionFusionTest, DotOutputFusionBiasAdd) { + auto module = ParseHloString(R"( + HloModule test_module + ENTRY OutputFusion { + alpha = f32[] constant(3) + broadcast = f32[4,4]{1,0} broadcast(alpha), dimensions={} + p0 = f32[4,3]{1,0} parameter(0) + p1 = f32[4,3]{1,0} parameter(1) + p2 = f32[4,4]{1,0} parameter(2) + transpose = f32[3,4]{1,0} transpose(p1), dimensions={1, 0} + dot = f32[4,4]{1,0} dot(p0, transpose), lhs_contracting_dims={1}, rhs_contracting_dims={0} + ROOT add = f32[4,4] add(dot, p2) + })") + .ValueOrDie(); + + EXPECT_TRUE(GpuInstructionFusion(/*may_duplicate=*/true) + .Run(module.get()) + .ValueOrDie()); + + HloInstruction* root = module->entry_computation()->root_instruction(); + EXPECT_THAT(root, op::Fusion()); + EXPECT_EQ(root->fusion_kind(), HloInstruction::FusionKind::kOutput); + EXPECT_THAT(root->fused_expression_root(), + op::Add(op::Dot(op::Parameter(), op::Transpose(op::Parameter())), + op::Parameter())); +} + +TEST_F(InstructionFusionTest, + DotOperationFusion_DontOutputFuseDuplicateOperands) { + absl::string_view module_string = R"( +HloModule module + +ENTRY main { + a = f32[50,60]{1,0} parameter(0) + b = f32[60,1]{1,0} parameter(1) + c = f32[50,1]{1,0} dot(a, b), lhs_contracting_dims={1}, rhs_contracting_dims={0} + ROOT d = f32[50,1]{1,0} add(c, c) +} +)"; + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(module_string)); + TF_ASSERT_OK_AND_ASSIGN( + bool fused_something, + GpuInstructionFusion(/*may_duplicate=*/false).Run(module.get())); + EXPECT_FALSE(fused_something); + EXPECT_THAT(module->entry_computation()->root_instruction(), + Not(op::Fusion())); +} + // Compute sum(1/p0), where p0 has type f32, twice. Check that the division is // duplicated and fused into both reduces. TEST_F(InstructionFusionTest, FloatingPointDivIsCheap) { @@ -456,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. - 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. + // 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(); + .ValueOrDie()); } TEST_F(InstructionFusionTest, FuseScalarConstant) { @@ -696,7 +555,7 @@ TEST_F(InstructionFusionTest, AvoidsLargeFusion) { sum = b.AddInstruction( HloInstruction::CreateBinary(shape, HloOpcode::kAdd, sum, param)); } - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(b.Build()); EXPECT_TRUE(GpuInstructionFusion(/*may_duplicate=*/true) .Run(module.get()) @@ -748,5 +607,56 @@ TEST_F(InstructionFusionTest, FuseIntoScatter) { op::Scatter(op::Add(), op::Add(), op::Add())); } +TEST_F(InstructionFusionTest, NonscalarConstantsNotFused) { + auto module = ParseHloString(R"( + HloModule test_module + + add { + lhs = f32[] parameter(0) + rhs = f32[] parameter(1) + ROOT add = f32[] add(lhs, rhs) + } + + ENTRY BroadcastIntoReduce { + constant = f32[16] constant({0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15}) + broadcast = f32[16,16,16,16]{3,2,1,0} broadcast(constant), dimensions={0} + constant.1 = f32[] constant(0) + ROOT reduce = f32[] reduce(broadcast, constant.1), dimensions={0,1,2,3}, + to_apply=add + })") + .ValueOrDie(); + + EXPECT_TRUE(GpuInstructionFusion(/*may_duplicate=*/true) + .Run(module.get()) + .ValueOrDie()); + // The f32[16] constant should not be fused into the reduce, but the f32[] + // constant should be. + auto* root = module->entry_computation()->root_instruction(); + ASSERT_THAT(root, op::Fusion()); + EXPECT_THAT(root->fused_instructions_computation()->root_instruction(), + op::Reduce(op::Broadcast(op::Parameter()), op::Constant())); +} + +TEST_F(InstructionFusionTest, FuseReverse) { + auto module = ParseHloString(R"( + HloModule test_module + + ENTRY Reverse { + p0 = f32[50,96,1024]{2,1,0} parameter(0) + add = f32[50,96,1024]{2,1,0} add(p0, p0) + ROOT reverse = f32[50,96,1024] reverse(add), dimensions={0} + })") + .ValueOrDie(); + + EXPECT_TRUE(GpuInstructionFusion(/*may_duplicate=*/true) + .Run(module.get()) + .ValueOrDie()); + + HloInstruction* root = module->entry_computation()->root_instruction(); + EXPECT_THAT(root, op::Fusion()); + EXPECT_THAT(root->fused_expression_root(), + op::Reverse(op::Add(op::Parameter(), op::Parameter()))); +} + } // namespace gpu } // namespace xla diff --git a/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc b/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc index ec3d8f9405840bb7be97ba5cd5725a4ac68a15a8..82bdd677d96d3d0826bb4127b32d074eb632b1a3 100644 --- a/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc +++ b/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc @@ -38,10 +38,9 @@ namespace gpu { namespace { -// Return whether the given shape is a matrix with no padding. -bool IsRank2WithNoPadding(const Shape& shape, int64 batch_dimensions_size) { - return ShapeUtil::Rank(shape) == batch_dimensions_size + 2 && - !LayoutUtil::IsPadded(shape); +// Return whether the given shape is rank 2 excluding the batch dimensions. +bool IsRank2(const Shape& shape, int64 batch_dimensions_size) { + return shape.rank() == batch_dimensions_size + 2; } // In a gemm operation where output = lhs * rhs, check whether the given shapes @@ -55,11 +54,11 @@ 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); - return type_is_allowed && - IsRank2WithNoPadding(lhs_shape, batch_dimensions_size) && - IsRank2WithNoPadding(rhs_shape, batch_dimensions_size) && - IsRank2WithNoPadding(output_shape, batch_dimensions_size) && + 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) && !ShapeUtil::IsZeroElementArray(lhs_shape) && !ShapeUtil::IsZeroElementArray(rhs_shape); } @@ -93,7 +92,8 @@ bool ImplementedAsGemm(const HloInstruction& hlo) { if (hlo.opcode() == HloOpcode::kFusion && hlo.fusion_kind() == HloInstruction::FusionKind::kOutput && - hlo.fused_expression_root()->opcode() == HloOpcode::kMultiply) { + (hlo.fused_expression_root()->opcode() == HloOpcode::kMultiply || + hlo.fused_expression_root()->opcode() == HloOpcode::kAdd)) { // Try to find the dot inside the output fusion node. const HloInstruction* dot = hlo.fused_expression_root()->operand(0); if (dot->opcode() != HloOpcode::kDot) { @@ -155,20 +155,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 @@ -269,5 +266,17 @@ string CudnnConvKindToString(CudnnConvKind kind) { } } +llvm::Value* IsBlock0Thread0(llvm::IRBuilder<>* b) { + return b->CreateAnd( + b->CreateICmpEQ( + b->getInt32(0), + llvm_ir::EmitCallToIntrinsic( + llvm::Intrinsic::nvvm_read_ptx_sreg_tid_x, {}, {}, b)), + b->CreateICmpEQ( + b->getInt32(0), + llvm_ir::EmitCallToIntrinsic( + llvm::Intrinsic::nvvm_read_ptx_sreg_ctaid_x, {}, {}, b))); +} + } // namespace gpu } // namespace xla diff --git a/tensorflow/compiler/xla/service/gpu/ir_emission_utils.h b/tensorflow/compiler/xla/service/gpu/ir_emission_utils.h index f373d4a8393a047aba599b0fae954e98a740161e..ebf4d926b7a280e10b09a2532caba7ad6ab3ceb2 100644 --- a/tensorflow/compiler/xla/service/gpu/ir_emission_utils.h +++ b/tensorflow/compiler/xla/service/gpu/ir_emission_utils.h @@ -155,6 +155,10 @@ llvm::Value* EmitPrintf(absl::string_view fmt, llvm::Value* EmitFullWarpShuffleDown(llvm::Value* value, llvm::Value* offset, llvm::IRBuilder<>* builder); +// Emits code that determines whether the current thread is thread 0 within +// block 0 of the kernel. +llvm::Value* IsBlock0Thread0(llvm::IRBuilder<>* b); + } // namespace gpu } // namespace xla diff --git a/tensorflow/compiler/xla/service/gpu/ir_emitter.cc b/tensorflow/compiler/xla/service/gpu/ir_emitter.cc index a3821e077ecf6b1dce1e2c8785fe3a59516db2be..0007a9a8a3369d8ac010640127e1561615a6d813 100644 --- a/tensorflow/compiler/xla/service/gpu/ir_emitter.cc +++ b/tensorflow/compiler/xla/service/gpu/ir_emitter.cc @@ -63,9 +63,6 @@ IrEmitter::IrEmitter(const HloModuleConfig& hlo_module_config, &ir_emitter_context->buffer_assignment(), &b_, module_, is_nested), hlo_module_config_(hlo_module_config) { - b_.setFastMathFlags(llvm_ir::GetFastMathFlags( - /*fast_math_enabled=*/hlo_module_config.debug_options() - .xla_gpu_enable_fast_math())); } Status IrEmitter::DefaultAction(HloInstruction* hlo) { @@ -97,6 +94,18 @@ Status IrEmitter::HandleBitcast(HloInstruction* bitcast) { return Status::OK(); } +Status IrEmitter::HandleAddDependency(HloInstruction* add_dependency) { + VLOG(2) << "HandleAddDependency: " << add_dependency->ToString(); + const HloInstruction* operand = add_dependency->operand(0); + // Add_Dependency is a no-op, but we still want to bind it to an llvm::Value + // sometimes, e.g., when it's operand is a constant or a bitcast of a + // constant. + if (bindings_.BoundToIrValue(*operand)) { + bindings_.BindHloToIrValue(*add_dependency, GetBasePointer(*operand)); + } + return Status::OK(); +} + Status IrEmitter::HandleGetTupleElement(HloInstruction* get_tuple_element) { auto operand = get_tuple_element->operand(0); CHECK(bindings_.BoundToIrValue(*operand)); @@ -421,7 +430,7 @@ Status IrEmitter::HandleTupleSelect(HloInstruction* tuple_select) { auto on_false = tuple_select->operand(2); TF_RET_CHECK(pred->shape().element_type() == PRED); TF_RET_CHECK(ShapeUtil::IsScalar(pred->shape())); - TF_RET_CHECK(ShapeUtil::IsTuple(tuple_select->shape())); + TF_RET_CHECK(tuple_select->shape().IsTuple()); llvm_ir::EmitTupleSelect(GetIrArray(*tuple_select, *tuple_select), GetIrArray(*pred, *tuple_select), GetBasePointer(*on_true), GetBasePointer(*on_false), @@ -628,9 +637,9 @@ Status IrEmitter::HandleFft(HloInstruction* fft) { return Unimplemented("Hit a case for fft that is not implemented on GPU."); } -Status IrEmitter::HandleCrossReplicaSum(HloInstruction* crs) { +Status IrEmitter::HandleAllReduce(HloInstruction* crs) { // TODO(b/33011107): Support cross replica sum on GPU. - return Unimplemented("CrossReplicaSum is not implemented on GPU."); + return Unimplemented("AllReduce is not implemented on GPU."); } Status IrEmitter::HandleParameter(HloInstruction* parameter) { @@ -639,7 +648,7 @@ Status IrEmitter::HandleParameter(HloInstruction* parameter) { Status IrEmitter::HandleReduce(HloInstruction* reduce) { // TODO(b/112040122): Support variadic reduce. - if (!ShapeUtil::IsArray(reduce->shape())) { + if (!reduce->shape().IsArray()) { return Unimplemented("Variadic reduce is not supported on GPU"); } auto arg = reduce->operand(0); @@ -697,15 +706,11 @@ Status IrEmitter::HandleReduce(HloInstruction* reduce) { Status IrEmitter::HandleFusion(HloInstruction* fusion) { // kFusion for library calls should be handled by // IrEmitterUnnested::HandleFusion. - CHECK(HloInstruction::FusionKind::kLoop == fusion->fusion_kind()); - - std::vector parameter_arrays; - for (HloInstruction* operand : fusion->operands()) { - parameter_arrays.push_back(GetIrArray(*operand, *fusion)); - } + CHECK_EQ(HloInstruction::FusionKind::kLoop, fusion->fusion_kind()); GpuElementalIrEmitter elemental_emitter(hlo_module_config_, module_, &b_, GetNestedComputer()); - FusedIrEmitter fused_emitter(parameter_arrays, &elemental_emitter); + FusedIrEmitter fused_emitter(GetGeneratorForOperandIrArrays(fusion), + &elemental_emitter); TF_RETURN_IF_ERROR(fusion->fused_expression_root()->Accept(&fused_emitter)); return EmitTargetElementLoop(*fusion, fused_emitter.GetRootGenerator()); @@ -778,7 +783,7 @@ StatusOr IrEmitter::ComputeNestedElement( std::vector IrEmitter::ConstructIrArrayForOutputs( const HloInstruction& hlo) { std::vector output_arrays; - if (ShapeUtil::IsTuple(hlo.shape())) { + if (hlo.shape().IsTuple()) { int64 num_outputs = ShapeUtil::TupleElementCount(hlo.shape()); output_arrays.reserve(num_outputs); for (int64 i = 0; i < num_outputs; ++i) { diff --git a/tensorflow/compiler/xla/service/gpu/ir_emitter.h b/tensorflow/compiler/xla/service/gpu/ir_emitter.h index 880520148005838cc25a5be9e26c8bc9028a70ce..f380aee9d3c06a29b503c81c7bd3846dbccf6ce5 100644 --- a/tensorflow/compiler/xla/service/gpu/ir_emitter.h +++ b/tensorflow/compiler/xla/service/gpu/ir_emitter.h @@ -68,6 +68,9 @@ namespace gpu { class IrEmitter : public DfsHloVisitorWithDefault, public IrBuilderMixin { public: + using GeneratorForOperandIrArrays = + std::function()>; + IrEmitter(const IrEmitter&) = delete; IrEmitter& operator=(const IrEmitter&) = delete; @@ -78,7 +81,7 @@ class IrEmitter : public DfsHloVisitorWithDefault, Status HandleDot(HloInstruction* dot) override; Status HandleConvolution(HloInstruction* convolution) override; Status HandleFft(HloInstruction* fft) override; - Status HandleCrossReplicaSum(HloInstruction* crs) override; + Status HandleAllReduce(HloInstruction* crs) override; Status HandleInfeed(HloInstruction* infeed) override; Status HandleOutfeed(HloInstruction* outfeed) override; Status HandleSend(HloInstruction* send) override; @@ -97,6 +100,7 @@ class IrEmitter : public DfsHloVisitorWithDefault, Status HandleBatchNormInference(HloInstruction* batch_norm) override; Status HandleBatchNormTraining(HloInstruction* batch_norm) override; Status HandleBatchNormGrad(HloInstruction* batch_norm) override; + Status HandleAddDependency(HloInstruction* add_dependency) override; Status FinishVisit(HloInstruction* root) override { return Status::OK(); } @@ -179,6 +183,20 @@ class IrEmitter : public DfsHloVisitorWithDefault, // Hlo configuration data used during code generation. const HloModuleConfig& hlo_module_config_; + protected: + GeneratorForOperandIrArrays GetGeneratorForOperandIrArrays( + HloInstruction* fusion) { + return [=]() { + std::vector ir_arrays; + ir_arrays.reserve(fusion->operand_count()); + absl::c_transform(fusion->operands(), std::back_inserter(ir_arrays), + [&](const HloInstruction* operand) { + return GetIrArray(*operand, *fusion); + }); + return ir_arrays; + }; + } + private: // A helper method for EmitAtomicOperationForNestedComputation. Certain // computations, such as floating-point addition and integer maximization, can diff --git a/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc b/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc index eb8aaaea4f91f552c2f21f104b83924fd604ebfa..294a454931b5cfa368bf094c428a1e942f4556b8 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 @@ -22,7 +23,6 @@ limitations under the License. #include "tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.h" #include "absl/algorithm/container.h" -#include "absl/container/inlined_vector.h" #include "absl/memory/memory.h" #include "absl/strings/str_cat.h" #include "absl/types/optional.h" @@ -65,11 +65,11 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_casting_utils.h" #include "tensorflow/compiler/xla/service/hlo_computation.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" #include "tensorflow/compiler/xla/service/llvm_ir/buffer_assignment_util.h" #include "tensorflow/compiler/xla/service/llvm_ir/dynamic_update_slice_util.h" #include "tensorflow/compiler/xla/service/llvm_ir/fused_ir_emitter.h" -#include "tensorflow/compiler/xla/service/llvm_ir/kernel_support_library.h" #include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h" #include "tensorflow/compiler/xla/service/llvm_ir/sort_util.h" #include "tensorflow/compiler/xla/service/llvm_ir/tuple_ops.h" @@ -88,6 +88,11 @@ limitations under the License. namespace xla { namespace gpu { +using llvm_ir::KernelMappingScheme; +using EmitElementFunction = + std::function; + namespace { using absl::InlinedVector; @@ -291,13 +296,12 @@ llvm::Type* GetIndexTypeForKernel(const HloInstruction* hlo, int64 launch_size, auto shape_in_range = [&](const Shape& s) { bool in_range = true; - ShapeUtil::ForEachSubshape( - s, [&](const Shape& sub_shape, const ShapeIndex& /*index*/) { - if (ShapeUtil::IsArray(sub_shape) && - !IsInt32(ShapeUtil::ElementsIn(sub_shape))) { - in_range = false; - } - }); + ShapeUtil::ForEachSubshape(s, [&](const Shape& sub_shape, + const ShapeIndex& /*index*/) { + if (sub_shape.IsArray() && !IsInt32(ShapeUtil::ElementsIn(sub_shape))) { + in_range = false; + } + }); return in_range; }; @@ -337,34 +341,26 @@ llvm::Type* GetIndexTypeForKernel(const HloInstruction* hlo, int64 launch_size, } // namespace Status IrEmitterUnnested::DefaultAction(HloInstruction* hlo) { - int unroll_factor = 1; - // Unfused elementwise operations are usually memory bound, unroll them. - if (hlo->IsElementwise()) { - unroll_factor = ComputeMaxUnrollFactor(hlo); - } - - thunk_sequence_->emplace_back(BuildKernelThunk( - hlo, /*implements_whole_instruction=*/true, unroll_factor)); return IrEmitter::DefaultAction(hlo); } Status IrEmitterUnnested::HandleDot(HloInstruction* dot) { if (ImplementedAsGemm(*dot)) { - thunk_sequence_->emplace_back(BuildGemmThunk(dot)); + AddThunkToThunkSequence(BuildGemmThunk(dot)); return Status::OK(); } - thunk_sequence_->emplace_back( + AddThunkToThunkSequence( BuildKernelThunk(dot, /*implements_whole_instruction=*/true)); return IrEmitter::HandleDot(dot); } Status IrEmitterUnnested::HandleConditional(HloInstruction* conditional) { - thunk_sequence_->emplace_back(BuildConditionalThunk(conditional)); + AddThunkToThunkSequence(BuildConditionalThunk(conditional)); return Status::OK(); } Status IrEmitterUnnested::HandleConvolution(HloInstruction* convolution) { - thunk_sequence_->emplace_back( + AddThunkToThunkSequence( BuildKernelThunk(convolution, /*implements_whole_instruction=*/true)); return IrEmitter::HandleConvolution(convolution); } @@ -386,7 +382,7 @@ Status IrEmitterUnnested::HandleCustomCall(HloInstruction* custom_call) { CHECK(feature_index->IsConstant()); int64 feature_index_value = feature_index->literal().Get({}); - thunk_sequence_->emplace_back( + AddThunkToThunkSequence( absl::make_unique( /*operand=*/GetAllocationSlice(*custom_call->operand(0)), /*scale=*/GetAllocationSlice(*custom_call->operand(1)), @@ -416,7 +412,7 @@ Status IrEmitterUnnested::HandleCustomCall(HloInstruction* custom_call) { auto output_data = assn.GetUniqueSlice(custom_call, {0}).ValueOrDie(); auto output_mean = assn.GetUniqueSlice(custom_call, {1}).ValueOrDie(); auto output_inv_stddev = assn.GetUniqueSlice(custom_call, {2}).ValueOrDie(); - thunk_sequence_->emplace_back( + AddThunkToThunkSequence( absl::make_unique( /*operand=*/GetAllocationSlice(*custom_call->operand(0)), /*scale=*/GetAllocationSlice(*custom_call->operand(1)), @@ -447,20 +443,19 @@ Status IrEmitterUnnested::HandleCustomCall(HloInstruction* custom_call) { auto output_grad_scale = assn.GetUniqueSlice(custom_call, {1}).ValueOrDie(); auto output_grad_offset = assn.GetUniqueSlice(custom_call, {2}).ValueOrDie(); - thunk_sequence_->emplace_back( - absl::make_unique( - /*operand=*/GetAllocationSlice(*custom_call->operand(0)), - /*scale=*/GetAllocationSlice(*custom_call->operand(1)), - /*mean=*/GetAllocationSlice(*custom_call->operand(2)), - /*inv_stddev=*/GetAllocationSlice(*custom_call->operand(3)), - /*grad_output=*/GetAllocationSlice(*custom_call->operand(4)), - /*epsilon=*/epsilon_value, - /*feature_index=*/feature_index_value, - /*output_grad_data=*/output_grad_data, - /*output_grad_scale=*/output_grad_scale, - /*output_grad_offset=*/output_grad_offset, - /*output_tuple=*/GetAllocationSlice(*custom_call), - /*hlo=*/custom_call)); + AddThunkToThunkSequence(absl::make_unique( + /*operand=*/GetAllocationSlice(*custom_call->operand(0)), + /*scale=*/GetAllocationSlice(*custom_call->operand(1)), + /*mean=*/GetAllocationSlice(*custom_call->operand(2)), + /*inv_stddev=*/GetAllocationSlice(*custom_call->operand(3)), + /*grad_output=*/GetAllocationSlice(*custom_call->operand(4)), + /*epsilon=*/epsilon_value, + /*feature_index=*/feature_index_value, + /*output_grad_data=*/output_grad_data, + /*output_grad_scale=*/output_grad_scale, + /*output_grad_offset=*/output_grad_offset, + /*output_tuple=*/GetAllocationSlice(*custom_call), + /*hlo=*/custom_call)); return Status::OK(); } @@ -475,7 +470,7 @@ Status IrEmitterUnnested::HandleCustomCall(HloInstruction* custom_call) { auto conv_result_slice = assn.GetUniqueSlice(custom_call, {0}).ValueOrDie(); auto scratch_slice = assn.GetUniqueSlice(custom_call, {1}).ValueOrDie(); - thunk_sequence_->emplace_back(absl::make_unique( + AddThunkToThunkSequence(absl::make_unique( Cast(custom_call), std::move(operand_slices), conv_result_slice, scratch_slice, tuple_result_slice)); return Status::OK(); @@ -488,7 +483,7 @@ Status IrEmitterUnnested::HandleFft(HloInstruction* fft) { TF_RET_CHECK( LayoutUtil::IsMonotonicWithDim0Major(fft->operand(0)->shape().layout())); TF_RET_CHECK(LayoutUtil::IsMonotonicWithDim0Major(fft->shape().layout())); - thunk_sequence_->emplace_back(BuildFftThunk(fft)); + AddThunkToThunkSequence(BuildFftThunk(fft)); return Status::OK(); } @@ -506,15 +501,12 @@ Status IrEmitterUnnested::HandleFusion(HloInstruction* fusion) { thunks.push_back(BuildKernelThunk( fusion, /*implements_whole_instruction=*/false, unroll_factor)); - std::vector operand_parameter_arrays; - for (HloInstruction* operand : fusion->operands()) { - operand_parameter_arrays.push_back(GetIrArray(*operand, *fusion)); - } GpuElementalIrEmitter operand_elemental_emitter( hlo_module_config_, ir_emitter_context_->llvm_module(), &b_, GetNestedComputer()); - FusedIrEmitter operand_fused_emitter(operand_parameter_arrays, - &operand_elemental_emitter); + FusedIrEmitter operand_fused_emitter( + GetGeneratorForOperandIrArrays(fusion), + &operand_elemental_emitter); TF_RETURN_IF_ERROR( root->mutable_operand(0)->Accept(&operand_fused_emitter)); @@ -530,15 +522,12 @@ Status IrEmitterUnnested::HandleFusion(HloInstruction* fusion) { BuildKernelThunk(fusion, /*implements_whole_instruction=*/false)); // Spin up a new fused emitter for the scatter kernel and emit it. - std::vector scatter_parameter_arrays; - for (HloInstruction* operand : fusion->operands()) { - scatter_parameter_arrays.push_back(GetIrArray(*operand, *fusion)); - } GpuElementalIrEmitter scatter_elemental_emitter( hlo_module_config_, ir_emitter_context_->llvm_module(), &b_, GetNestedComputer()); - FusedIrEmitter scatter_fused_emitter(scatter_parameter_arrays, - &scatter_elemental_emitter); + FusedIrEmitter scatter_fused_emitter( + GetGeneratorForOperandIrArrays(fusion), + &scatter_elemental_emitter); TF_RETURN_IF_ERROR(root->Accept(&scatter_fused_emitter)); TF_RETURN_IF_ERROR(EmitScatter( thunks.back().get(), root, @@ -547,7 +536,7 @@ Status IrEmitterUnnested::HandleFusion(HloInstruction* fusion) { /*updates_gen=*/ scatter_fused_emitter.GetGenerator(root->operand(2)))); } - thunk_sequence_->emplace_back( + AddThunkToThunkSequence( absl::make_unique(std::move(thunks), fusion)); return Status::OK(); } @@ -556,96 +545,11 @@ Status IrEmitterUnnested::HandleFusion(HloInstruction* fusion) { // HandleFusion specializes reduction from a multi-dimensional array to // 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 && - ShapeUtil::IsTuple(root->shape())) { + if (root->opcode() == HloOpcode::kReduce && root->shape().IsTuple()) { // TODO(b/112040122): Support variadic reduce. return Unimplemented("Variadic reduce is not supported on GPU"); } - VLOG(3) << "Emitting fused reduction to vector: " << fusion->ToString(); - std::vector> thunks; - absl::Span output_instructions = - root->opcode() == HloOpcode::kTuple - ? root->operands() - : absl::Span(&root, 1); - - // For multi-output fusion emit an initializer for each tuple element. - // Otherwise it's sufficient to just initialize the single output. - HloInstruction* first_reduce = nullptr; - for (int i = 0, e = output_instructions.size(); i != e; ++i) { - if (output_instructions[i]->opcode() == HloOpcode::kReduce) { - TF_ASSIGN_OR_RETURN( - std::unique_ptr initializer_thunk, - BuildInitializerThunk(fusion, output_instructions[i] == root - ? ShapeIndex() - : ShapeIndex({i}))); - thunks.push_back(std::move(initializer_thunk)); - first_reduce = - first_reduce == nullptr ? output_instructions[i] : first_reduce; - } - } - CHECK(first_reduce != nullptr); - thunks.push_back( - BuildKernelThunk(fusion, /*implements_whole_instruction=*/false)); - thunk_sequence_->emplace_back( - absl::make_unique(std::move(thunks), fusion)); - std::vector parameter_arrays; - for (HloInstruction* operand : fusion->operands()) { - parameter_arrays.push_back(GetIrArray(*operand, *fusion)); - } - GpuElementalIrEmitter elemental_emitter( - hlo_module_config_, ir_emitter_context_->llvm_module(), &b_, - GetNestedComputer()); - FusedIrEmitter fused_emitter(parameter_arrays, &elemental_emitter); - TF_RETURN_IF_ERROR(root->Accept(&fused_emitter)); - - // For multi-output fusion CHECK the constraints and feed all the - // reduces into a single loop code generator. Single-output reduce - // fusion is a special case of that. - InlinedVector input_gens; - InlinedVector init_value_gens; - std::vector> - extra_output_gens; - InlinedVector reducers; - InlinedVector reduce_output_shapes; - for (int i = 0, e = output_instructions.size(); i != e; ++i) { - const HloInstruction* inst = output_instructions[i]; - ShapeIndex output_shape_index; - if (root->opcode() == HloOpcode::kTuple) { - output_shape_index = {i}; - } - if (inst->opcode() == HloOpcode::kReduce) { - CHECK(IsReductionToVector(*inst)) - << "Only reductions to vector are supported"; - // Shapes, layouts and dimensions must be the same for all reduces - // inside of this fusion. - CHECK(ShapeUtil::Equal(first_reduce->shape(), inst->shape())); - CHECK(ShapeUtil::Equal(first_reduce->operand(0)->shape(), - inst->operand(0)->shape())); - CHECK(ShapeUtil::Equal(first_reduce->operand(1)->shape(), - inst->operand(1)->shape())); - CHECK(first_reduce->dimensions() == inst->dimensions()); - input_gens.push_back(fused_emitter.GetGenerator(inst->operand(0))); - init_value_gens.push_back( - fused_emitter.GetGenerator(inst->operand(1))); - reducers.push_back(inst->to_apply()); - reduce_output_shapes.push_back(std::move(output_shape_index)); - } else { - // For extra outputs we can relax shape equality to allow different - // types (with the same number of elements). Layouts still have to - // match. - CHECK(ShapeUtil::CompatibleIgnoringElementType( - first_reduce->operand(0)->shape(), inst->shape())); - CHECK(LayoutUtil::Equal(first_reduce->operand(0)->shape().layout(), - inst->shape().layout())); - extra_output_gens.emplace_back(fused_emitter.GetGenerator(inst), - std::move(output_shape_index)); - } - } - const Shape& input_shape = first_reduce->operand(0)->shape(); - return EmitReductionToVector(first_reduce, input_shape, input_gens, - init_value_gens, - first_reduce->dimensions(), reducers, - reduce_output_shapes, extra_output_gens); + return EmitReductionToVector(fusion); } default: LOG(FATAL) << "Bad opcode for input fusion: " @@ -659,12 +563,8 @@ Status IrEmitterUnnested::HandleFusion(HloInstruction* fusion) { // touching the un-updated elements. // Set up kernel thunk and fused ir emitter. - thunk_sequence_->emplace_back( - BuildKernelThunk(fusion, /*implements_whole_instruction=*/true)); - std::vector operand_arrays; - for (HloInstruction* operand : fusion->operands()) { - operand_arrays.push_back(GetIrArray(*operand, *fusion)); - } + std::unique_ptr fusion_thunk = + BuildKernelThunk(fusion, /*implements_whole_instruction=*/true); GpuElementalIrEmitter elemental_emitter(hlo_module_config_, ir_emitter_context_->llvm_module(), &b_, GetNestedComputer()); @@ -678,18 +578,17 @@ Status IrEmitterUnnested::HandleFusion(HloInstruction* fusion) { LaunchDimensions launch_dimensions = CalculateLaunchDimensions( update_shape, ir_emitter_context_->device_description()); - CHECK(Thunk::Kind::kKernel == LastThunk()->kind()); - UpdateLaunchDimensions(launch_dimensions, - static_cast(LastThunk()), + UpdateLaunchDimensions(launch_dimensions, fusion_thunk.get(), ir_emitter_context_->llvm_module()); + AddThunkToThunkSequence(std::move(fusion_thunk)); return llvm_ir::EmitParallelFusedDynamicUpdateSliceInPlace( - fusion, operand_arrays, output_array, &elemental_emitter, - launch_dimensions, &b_); + fusion, GetGeneratorForOperandIrArrays(fusion), output_array, + &elemental_emitter, launch_dimensions, &b_); } if (ImplementedAsGemm(*fusion)) { - thunk_sequence_->emplace_back(BuildGemmThunk(fusion)); + AddThunkToThunkSequence(BuildGemmThunk(fusion)); return Status::OK(); } @@ -699,10 +598,6 @@ Status IrEmitterUnnested::HandleFusion(HloInstruction* fusion) { return Status::OK(); } - int unroll_factor = ComputeMaxUnrollFactor(fusion); - - thunk_sequence_->emplace_back(BuildKernelThunk( - fusion, /*implements_whole_instruction=*/true, unroll_factor)); return IrEmitter::HandleFusion(fusion); } @@ -713,7 +608,7 @@ Status IrEmitterUnnested::HandleCopy(HloInstruction* copy) { if (LayoutUtil::Equal(copy->operand(0)->shape().layout(), copy->shape().layout()) && buffer_assignment.GetUniqueTopLevelSlice(copy->operand(0)).ok()) { - thunk_sequence_->emplace_back(BuildDeviceToDeviceCopyThunk(copy)); + AddThunkToThunkSequence(BuildDeviceToDeviceCopyThunk(copy)); return Status::OK(); } if (CheckAndEmitHloWithTile021(copy)) { @@ -724,13 +619,12 @@ Status IrEmitterUnnested::HandleCopy(HloInstruction* copy) { } Status IrEmitterUnnested::EmitExtraOutputsForReduce( - const HloInstruction* reduce, const IrArray::Index& index, + const HloInstruction* unnested_hlo, const IrArray::Index& index, absl::Span> extra_output_gens) { for (int i = 0; i != extra_output_gens.size(); ++i) { - const HloInstruction* output = reduce->parent()->FusionInstruction(); llvm::Value* extra_output_address = - GetIrArray(*output, *output, extra_output_gens[i].second) + GetIrArray(*unnested_hlo, *unnested_hlo, extra_output_gens[i].second) .EmitArrayElementAddress(index, &b_, "extra_output_element_address"); TF_ASSIGN_OR_RETURN(llvm::Value* const extra_output_ir_value, @@ -740,990 +634,15 @@ Status IrEmitterUnnested::EmitExtraOutputsForReduce( return Status::OK(); } -Status IrEmitterUnnested::EmitReductionToScalar( - HloInstruction* reduce, const Shape& input_shape, - absl::Span input_gens, - absl::Span init_value_gens, - absl::Span reducers, - absl::Span reduce_output_shapes, - absl::Span> - extra_output_gens) { - // Number of elements processed by a single thread. - constexpr int64 kTileSize = 16; - int64 num_elems = ShapeUtil::ElementsIn(input_shape); - - // Round up the number of tiles to a multiple of the warp size. This is - // necessary for correctness. We launch one thread per tile, and if the - // number of threads isn't a multiple of the number of the warp size, our - // shuffles will read from inactive threads, producing undefined values. - int64 num_tiles = - RoundUpToNearest(CeilOfRatio(num_elems, kTileSize), kWarpSize); - - Shape tiled_input_shape = ShapeUtil::MakeShapeWithLayout( - reduce->shape().element_type(), {num_tiles}, {0}); - LaunchDimensions launch_dimensions = CalculateLaunchDimensions( - tiled_input_shape, ir_emitter_context_->device_description()); - - llvm::Type* index_ty = - GetIndexTypeForKernel(reduce, launch_dimensions.launch_bound(), &b_); - - auto index_typed_constant = [&](uint64 c) -> llvm::Constant* { - return llvm::ConstantInt::get(index_ty, c); - }; - - // Check whether every thread will process a full tile's worth of elements - // without reading outside the bounds of the input. If this is true, we can - // skip some bounds checks in the final algorithm. - bool all_threads_in_bounds = num_tiles * kTileSize == num_elems; - - // __global__ void full_reduce_kernel() { - // x_in_tiles = threadIdx.x + blockIdx.x * blockDim.x; - // x = x_in_tiles * kTileSize; - // - // partial_result = init_value; - // if (all_threads_in_bounds || x + kTileSize <= num_elems) { - // for (i = 0; i < kTileSize; ++i) { - // partial_result = Reducer(partial_result, input[x + i]); - // } - // } else { - // for (i = 0; i < kTileSize; ++i) { - // if (x + i < num_elems) { - // partial_result = Reducer(partial_result, input[x + i]); - // } - // } - // } - // for (i = warpSize / 2; i > 0; i /= 2) { - // partial_result = Reducer(partial_result, - // __shfl_down(partial_result, i)); - // } - // if (lane_id == 0) { - // AtomicReducer(&output[y], partial_result); - // } - // } - // - // // Choose num_blocks and threads_per_block such that: - // // - // // num_blocks * threads_per_block = - // // RoundUpToNextMultipleOf(Ceil(num_elems / kTileSize), warpSize), - // // - // // and threads_per_block is a multiple of warpSize. - // reduce_kernel // - auto loop_body_emitter = [=](const IrArray::Index& tile_index) -> Status { - const int num_reduces = reducers.size(); - llvm::Type* element_ir_type = - llvm_ir::PrimitiveTypeToIrType(input_shape.element_type(), module_); - std::vector partial_reduction_result_addresses; - for (int i = 0; i != num_reduces; ++i) { - llvm::Value* partial_reduction_result_address = - Alloca(element_ir_type, /*ArraySize=*/nullptr, - "partial_reduction_result." + llvm::Twine(i)); - TF_ASSIGN_OR_RETURN(llvm::Value* const init_ir_value, - init_value_gens[i](IrArray::Index(index_ty))); - Store(init_ir_value, partial_reduction_result_address); - partial_reduction_result_addresses.push_back( - partial_reduction_result_address); - } - - llvm::Value* x_in_tiles = tile_index[0]; - x_in_tiles = ZExtOrTrunc(x_in_tiles, index_ty); - - // Emit an inner for-loop that reduces the elements in the tile. - auto emit_tile_element_loop = [=](bool tile_in_bounds) -> Status { - std::unique_ptr tile_element_loop = - llvm_ir::ForLoop::EmitForLoop( - "element_id_in_tile", index_typed_constant(0), - index_typed_constant(kTileSize), index_typed_constant(1), &b_); - - // Emit the body of the partial reduction loop. - llvm_ir::SetToFirstInsertPoint(tile_element_loop->GetBodyBasicBlock(), - &b_); - llvm::Value* x = - NSWAdd(NSWMul(x_in_tiles, index_typed_constant(kTileSize)), - tile_element_loop->GetIndVarValue()); - // Unless we know the tile is entirely in bounds, we have to emit a - // x-in-bounds check before reading from the input. - if (!tile_in_bounds) { - llvm_ir::LlvmIfData if_data = llvm_ir::EmitIfThenElse( - ICmpULT(x, index_typed_constant(num_elems)), "x_in_bounds", &b_); - - // Emit code that reads the input element and accumulates it to - // the partial reduction result. - llvm_ir::SetToFirstInsertPoint(if_data.true_block, &b_); - } - - IrArray::Index input_index( - /*linear=*/x, input_shape, &b_); - llvm::Value* input_address = Alloca(element_ir_type); - for (int i = 0; i != num_reduces; ++i) { - TF_ASSIGN_OR_RETURN(llvm::Value* const input_ir_value, - input_gens[i](input_index)); - Store(input_ir_value, input_address); - TF_RETURN_IF_ERROR(EmitCallToNestedComputation( - *reducers[i], - {partial_reduction_result_addresses[i], input_address}, - partial_reduction_result_addresses[i])); - } - return EmitExtraOutputsForReduce(reduce, input_index, extra_output_gens); - }; - - // x_end = kTileSize + x_in_tiles * kTileSize, i.e., the location that's - // immediately beyond the tile. - llvm::Value* x_end = - NSWAdd(index_typed_constant(kTileSize), - NSWMul(x_in_tiles, index_typed_constant(kTileSize))); - // The tile is entirely in bound if all_threads_in_bounds or - // x_end <= num_elems. - llvm::Value* tile_in_bounds = - Or(ICmpULE(x_end, index_typed_constant(num_elems)), - b_.getInt1(all_threads_in_bounds)); - llvm_ir::LlvmIfData if_tile_in_bounds_data = - llvm_ir::EmitIfThenElse(tile_in_bounds, "tile_in_bounds", &b_); - llvm_ir::SetToFirstInsertPoint(if_tile_in_bounds_data.true_block, &b_); - TF_RETURN_IF_ERROR(emit_tile_element_loop(/*tile_in_bounds=*/true)); - llvm_ir::SetToFirstInsertPoint(if_tile_in_bounds_data.false_block, &b_); - TF_RETURN_IF_ERROR(emit_tile_element_loop(/*tile_in_bounds=*/false)); - - // After the if-then-else statement on tile_in_bounds, emit calls to - // shfl_down that accumulate the partial reduction results of all threads - // from the warp. - llvm_ir::SetToFirstInsertPoint(if_tile_in_bounds_data.after_block, &b_); - int bit_width = llvm_ir::GetSizeInBits(element_ir_type); - // bitcast cannot be applied to aggregate types (even packed ones), so we - // instead bitcast addresses of load/store to intN* of the same bit-width. - llvm::Type* shuffle_ir_type = element_ir_type->isStructTy() - ? b_.getIntNTy(bit_width) - : element_ir_type; - for (int shuffle_distance = kWarpSize / 2; shuffle_distance >= 1; - shuffle_distance /= 2) { - llvm::Value* result_from_other_lane = - Alloca(element_ir_type, nullptr, "result_from_other_lane"); - for (int i = 0; i != num_reduces; ++i) { - llvm::Value* partial_reduction_result = - Load(BitCast(partial_reduction_result_addresses[i], - shuffle_ir_type->getPointerTo()), - "partial_reduction_result"); - CHECK_EQ(launch_dimensions.threads_per_block() % kWarpSize, 0) - << "Requires block size a multiple of the warp size, otherwise we " - "will read undefined elements."; - Store(EmitFullWarpShuffleDown(partial_reduction_result, - b_.getInt32(shuffle_distance), &b_), - BitCast(result_from_other_lane, shuffle_ir_type->getPointerTo())); - TF_RETURN_IF_ERROR(EmitCallToNestedComputation( - *reducers[i], - {partial_reduction_result_addresses[i], result_from_other_lane}, - partial_reduction_result_addresses[i])); - } - } - - const HloInstruction* output = - reduce->IsFused() ? reduce->parent()->FusionInstruction() : reduce; - - // Emit an atomic operation that accumulates the partial reduction result of - // lane 0 (which holds the partially accumulated result for its warp) to the - // output element. - llvm::Value* lane_id = - URem(x_in_tiles, index_typed_constant(kWarpSize), "lane_id"); - llvm_ir::LlvmIfData if_lane_id_is_zero_data = llvm_ir::EmitIfThenElse( - ICmpEQ(lane_id, index_typed_constant(0)), "lane_id_is_zero", &b_); - llvm_ir::SetToFirstInsertPoint(if_lane_id_is_zero_data.true_block, &b_); - - for (int i = 0; i != num_reduces; ++i) { - llvm::Value* output_address = - GetIrArray(*output, *output, reduce_output_shapes[i]) - .EmitArrayElementAddress( - IrArray::Index( - /*linear=*/b_.getInt64(0), - ShapeUtil::GetSubshape(output->shape(), - reduce_output_shapes[i]), - &b_), - &b_, "output_element_address"); - TF_RETURN_IF_ERROR(EmitAtomicOperationForNestedComputation( - *reducers[i], output_address, partial_reduction_result_addresses[i])); - } - return Status::OK(); - }; - - // Emit a parallel loop that iterates through all input tiles, one per thread. - CHECK(LastThunk()->kind() == Thunk::Kind::kSequential); - UpdateLaunchDimensions( - launch_dimensions, - static_cast(LastThunk())->thunks().back().get(), - ir_emitter_context_->llvm_module()); - return ParallelLoopEmitter(loop_body_emitter, tiled_input_shape, - launch_dimensions, &b_) - .EmitLoop(IrName(reduce), index_ty); -} - -Status IrEmitterUnnested::EmitColumnReduction( - int64 height, int64 width, HloInstruction* reduce, const Shape& input_shape, - absl::Span input_gens, - absl::Span init_value_gens, - absl::Span reducers, - absl::Span reduce_output_shapes, - absl::Span> - extra_output_gens) { - // Divide the input matrix into tiles of size KxL. For example, when the - // input matrix is 4x4, K=2, and L=1 the tiled matrix looks like - // - // 0123 - // 0123 - // 4567 - // 4567 // Numbers indicate tile IDs. - // - // Each tile is first partially reduced to a scalar by a thread, and then the - // scalar is accumulated to the output vector using atomic operations. - // - // We choose 128 as the tile size based on empirical evidence. It's big enough - // to reduce the amount of atomic adds in the end, maximizing the memory - // bandwidth. A tile width of 2 allows for high memory bandwidth utilization - // on 16b input data. - constexpr int64 kTileHeight = 128; - constexpr int64 kTileWidth = 2; - - // If the height is not a multiple of kTileHeight, we pad the bottom of the - // input matrix. - const int64 height_in_tiles = CeilOfRatio(height, kTileHeight); - // If width is not a multiple of kTileWidth the rightmost thread will process - // fewer input elements. - const int64 width_in_tiles = CeilOfRatio(width, kTileWidth); - Shape tiled_input_shape = - ShapeUtil::MakeShapeWithLayout(reduce->shape().element_type(), - {height_in_tiles, width_in_tiles}, {1, 0}); - LaunchDimensions launch_dimensions = CalculateLaunchDimensions( - tiled_input_shape, ir_emitter_context_->device_description()); - - // TODO(b/110211620): Convert to use i32 index_type when it is possible. - llvm::Type* index_ty = b_.getInt64Ty(); - - auto index_typed_constant = [&](uint64 c) -> llvm::Constant* { - return llvm::ConstantInt::get(index_ty, c); - }; - - // for (linear_index = threadIdx.x + blockIdx.x * blockDim.x; - // linear_index < height_in_tiles * width_in_tiles; - // linear_index += blockDim.x * gridDim.x) { - // y_in_tiles = linear_index / width_in_tiles; - // x_in_tiles = linear_index % width_in_tiles; - // - // partial_results[kTileWidth] = init_values; - // tile_in_y_bounds = height % kTileHeight == 0 || - // y_in_tiles * kTileHeight + kTileHeight <= height; - // tile_in_x_bounds = width % kTileWidth == 0 || - // x_in_tiles * kTileWidth + kTileWidth <= width; - // // The implementation handles y and x bound checks separately. - // if (tile_in_y_bounds && tile_in_x_bounds) { - // for (y_offset : range(kTileHeight)) { - // y = y_in_tiles * kTileHeight + y_offset; - // for (x_offset : range(kTileWidth)) { - // x = x_in_tiles * kTileWidth + x_offset; - // partial_result = Reducer(partial_result[x_offset], input[y][x]); - // } - // } - // } else { - // for (y_offset : range(kTileHeight)) { - // y = y_in_tiles * kTileHeight + y_offset; - // for (y_offset : range(kTileHeight)) { - // x = x_in_tiles * kTileWidth + x_offset; - // if (y < height && x < width) { - // partial_result = Reducer(partial_result, input[y][x]); - // } - // } - // } - // } - // for (x_offset : range(kTileWidth)) { - // AtomicReducer(&output[x + x_offset], partial_result[x_offset]); - // } - // } - auto loop_body_emitter = [=](const IrArray::Index& tile_index) -> Status { - const int num_reduces = reducers.size(); - // Emit the loop body that reduces one tile. - llvm::Type* element_ir_type = - llvm_ir::PrimitiveTypeToIrType(input_shape.element_type(), module_); - std::vector partial_reduction_result_addresses; - for (int i = 0; i != num_reduces; ++i) { - for (int x_offset = 0; x_offset < kTileWidth; ++x_offset) { - llvm::Value* partial_reduction_result_address = - Alloca(element_ir_type, /*ArraySize=*/nullptr, - "partial_reduction_result." + - llvm::Twine(i * kTileWidth + x_offset)); - TF_ASSIGN_OR_RETURN(llvm::Value* const init_ir_value, - init_value_gens[i](IrArray::Index(index_ty))); - Store(init_ir_value, partial_reduction_result_address); - partial_reduction_result_addresses.push_back( - partial_reduction_result_address); - } - } - - // Emit an inner for-loop that partially reduces the elements in the given - // tile. - llvm::Value* y_in_tiles = tile_index[0]; - llvm::Value* x_in_tiles = tile_index[1]; - - y_in_tiles = ZExtOrTrunc(y_in_tiles, index_ty); - x_in_tiles = ZExtOrTrunc(x_in_tiles, index_ty); - - auto emit_tile_element_loop = [=](bool tile_in_y_bounds, - bool tile_in_x_bounds) -> Status { - std::unique_ptr tile_element_loop = - llvm_ir::ForLoop::EmitForLoop( - "element_id_in_tile", index_typed_constant(0), - index_typed_constant(kTileHeight), index_typed_constant(1), &b_); - - // Emit the body of the partial reduction loop. - llvm_ir::SetToFirstInsertPoint(tile_element_loop->GetBodyBasicBlock(), - &b_); - llvm::Value* y = - NSWAdd(NSWMul(y_in_tiles, index_typed_constant(kTileHeight)), - tile_element_loop->GetIndVarValue()); - - // Unless we know that y is in bounds, we have to emit a check before - // reading from the input. - if (!tile_in_y_bounds) { - llvm_ir::LlvmIfData if_data = llvm_ir::EmitIfThenElse( - ICmpULT(y, index_typed_constant(height)), "y_in_bounds", &b_); - - // Emit code that reads the input element and accumulates it to - // the partial reduction result. - llvm_ir::SetToFirstInsertPoint(if_data.true_block, &b_); - } - for (int x_offset = 0; x_offset < kTileWidth; ++x_offset) { - llvm::Value* x = - NSWAdd(NSWMul(x_in_tiles, index_typed_constant(kTileWidth)), - index_typed_constant(x_offset)); - // Unless we know that x is in bounds, we have to emit a check before - // reading from the input. - if (!tile_in_x_bounds) { - llvm_ir::LlvmIfData if_data = llvm_ir::EmitIfThenElse( - ICmpULT(x, index_typed_constant(width)), "x_in_bounds", &b_); - llvm_ir::SetToFirstInsertPoint(if_data.true_block, &b_); - } - llvm::Value* input_address = Alloca(element_ir_type); - // {y,x} is an index to input_matrix_shape [height,width]. We need to - // convert that to an index to input_shape (the shape of the operand of - // "reduce"). This conversion is composed of a transposition from - // input_shape to normalized_input_shape and a reshape from - // normalized_input_shape to input_matrix_shape. - const Shape normalized_input_shape = - ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout( - input_shape); - auto input_shape_min2maj = LayoutUtil::MinorToMajor(input_shape); - const std::vector transpose_dimension_mapping( - input_shape_min2maj.rbegin(), input_shape_min2maj.rend()); - - const Shape input_matrix_shape = - ShapeUtil::MakeShapeWithDescendingLayout(input_shape.element_type(), - {height, width}); - const IrArray::Index input_matrix_index({y, x}, input_matrix_shape, - &b_); - const IrArray::Index input_index = - input_matrix_index - .SourceIndexOfReshape(input_matrix_shape, - normalized_input_shape, &b_) - .SourceIndexOfTranspose(normalized_input_shape, input_shape, - transpose_dimension_mapping, &b_); - for (int i = 0; i != num_reduces; ++i) { - TF_ASSIGN_OR_RETURN(llvm::Value* const input_ir_value, - input_gens[i](input_index)); - Store(input_ir_value, input_address); - TF_RETURN_IF_ERROR(EmitCallToNestedComputation( - *reducers[i], - {partial_reduction_result_addresses[i * kTileWidth + x_offset], - input_address}, - partial_reduction_result_addresses[i * kTileWidth + x_offset])); - TF_RETURN_IF_ERROR(EmitExtraOutputsForReduce(reduce, input_index, - extra_output_gens)); - } - } - return Status::OK(); - }; - - // y_end = kTileHeight + y_in_tiles * kTileHeight, i.e., the y location - // that's immediately beyond the tile. - llvm::Value* y_end = - NSWAdd(index_typed_constant(kTileHeight), - NSWMul(y_in_tiles, index_typed_constant(kTileHeight))); - // x_end = kTileWidth + x_in_tiles * kTileWidth, i.e., the x location - // that's immediately beyond the tile. - llvm::Value* x_end = - NSWAdd(index_typed_constant(kTileWidth), - NSWMul(x_in_tiles, index_typed_constant(kTileWidth))); - llvm::Value* tile_in_y_bounds = - Or(ICmpULE(y_end, index_typed_constant(height)), - b_.getInt1(height % kTileHeight == 0)); - llvm::Value* tile_in_x_bounds = - Or(ICmpULE(x_end, index_typed_constant(width)), - b_.getInt1(width % kTileWidth == 0)); - // The tile is in y bounds if "height" is a multiple of kTileHeight or - // y_end <= height. - llvm_ir::LlvmIfData if_tile_in_y_bounds_data = - llvm_ir::EmitIfThenElse(tile_in_y_bounds, "tile_in_y_bounds", &b_); - llvm_ir::SetToFirstInsertPoint(if_tile_in_y_bounds_data.true_block, &b_); - // The tile is in x bounds if "width" is a multiple of kTileWidth or - // x_end <= width. - llvm_ir::LlvmIfData if_tile_in_x_bounds_data = - llvm_ir::EmitIfThenElse(tile_in_x_bounds, "tile_in_x_bounds", &b_); - llvm_ir::SetToFirstInsertPoint(if_tile_in_x_bounds_data.true_block, &b_); - TF_RETURN_IF_ERROR(emit_tile_element_loop(/*tile_in_y_bounds=*/true, - /*tile_in_x_bounds=*/true)); - llvm_ir::SetToFirstInsertPoint(if_tile_in_x_bounds_data.false_block, &b_); - TF_RETURN_IF_ERROR(emit_tile_element_loop(/*tile_in_y_bounds=*/true, - /*tile_in_x_bounds=*/false)); - llvm_ir::SetToFirstInsertPoint(if_tile_in_y_bounds_data.false_block, &b_); - if_tile_in_x_bounds_data = - llvm_ir::EmitIfThenElse(tile_in_x_bounds, "tile_in_x_bounds", &b_); - llvm_ir::SetToFirstInsertPoint(if_tile_in_x_bounds_data.true_block, &b_); - TF_RETURN_IF_ERROR(emit_tile_element_loop(/*tile_in_y_bounds=*/false, - /*tile_in_x_bounds=*/true)); - llvm_ir::SetToFirstInsertPoint(if_tile_in_x_bounds_data.false_block, &b_); - TF_RETURN_IF_ERROR(emit_tile_element_loop(/*tile_in_y_bounds=*/false, - /*tile_in_x_bounds=*/false)); - - // After the nested if-then-else statement on tile_in_y_bounds and - // tile_in_x_bounds, emit atomic operations to accumulate the partial - // reduction result to the output element. - llvm_ir::SetToFirstInsertPoint(if_tile_in_y_bounds_data.after_block, &b_); - const HloInstruction* output = - reduce->IsFused() ? reduce->parent()->FusionInstruction() : reduce; - for (int i = 0; i != num_reduces; ++i) { - for (int x_offset = 0; x_offset < kTileWidth; ++x_offset) { - llvm::Value* x = - NSWAdd(NSWMul(x_in_tiles, index_typed_constant(kTileWidth)), - index_typed_constant(x_offset)); - llvm::Value* output_address = - GetIrArray(*output, *output, reduce_output_shapes[i]) - .EmitArrayElementAddress( - IrArray::Index( - x, - ShapeUtil::GetSubshape(output->shape(), - reduce_output_shapes[i]), - &b_), - &b_, "output_element_address"); - TF_RETURN_IF_ERROR(EmitAtomicOperationForNestedComputation( - *reducers[i], output_address, - partial_reduction_result_addresses[i * kTileWidth + x_offset])); - } - } - return Status::OK(); - }; - - // Emit a parallel loop that iterate through all input tiles. - CHECK(LastThunk()->kind() == Thunk::Kind::kSequential); - UpdateLaunchDimensions( - launch_dimensions, - static_cast(LastThunk())->thunks().back().get(), - ir_emitter_context_->llvm_module()); - return ParallelLoopEmitter(loop_body_emitter, tiled_input_shape, - launch_dimensions, &b_) - .EmitLoop(IrName(reduce), index_ty); -} - -static std::pair ComputeTilingSchemeForReduction( - int64 depth, int64 width, int64 kWarpSize) { - constexpr int64 kTargetNumElementsPerThread = 64; - int64 x_tile_size = kTargetNumElementsPerThread; - int64 z_tile_size = 1; - - // Only tile along the x dimension with tile size kTargetNumElementsPerThread - // if doing so doesn't require a slow version of loop with bound check on each - // dimension. A more sophisticated heuristics is to enable tile along the - // x dimension with tile size kTargetNumElementsPerThread when either width is - // a factor of (kWarpSize * kTargetNumElementsPerThread) or width is big - // enough so that only a small fraction of the threads execute the slow - // version of loop with bound check. - if (width % (kWarpSize * kTargetNumElementsPerThread) != 0) { - x_tile_size = 8; - z_tile_size = 8; - while (depth % z_tile_size != 0) { - z_tile_size -= 1; - } - } - - return std::pair(x_tile_size, z_tile_size); -} - -Status IrEmitterUnnested::EmitRowReduction( - int64 depth, int64 height, int64 width, HloInstruction* reduce, - const Shape& input_shape, - absl::Span input_gens, - absl::Span init_value_gens, - absl::Span reducers, - absl::Span reduce_output_shapes, - absl::Span> - extra_output_gens) { - // A naive algorithm is: - // 1. Divide the x dimension of the input tensor into tiles of size 1x1xX. - // 2. Partially reduces each tile to a scalar using one thread. - // 3. Accumulates that scalar to the output vector using atomic operations. - // - // for (linear_index = threadIdx.x + blockIdx.x * blockDim.x; - // linear_index < depth * height * width_in_tiles; - // linear_index += blockDim.x * gridDim.x) { - // int x_in_tiles = linear_index % width_in_tiles; - // int y = linear_index / width_in_tiles % height; - // int z = linear_index / (height * width_in_tiles); - // float partial_result = 0; - // for (element_id_in_tile : range(x_tile_size)) { - // int x = x_in_tiles * x_tile_size + element_id_in_tile; - // if (x < width) - // partial_result = reducer(partial_result, input[z][y][x]); - // } - // AtomicReducer(&output[y], partial_result); - // } - // - // Four optimizations are performed. - // - // 1. To coalesce global memory accesses, dilate the tile with a factor of 32 - // (i.e. the warp size). For example, suppose the width is 8x32=256. Instead - // of making each tile consecutive, we let make tile 0 column - // [0,32,64,...,224], tile 1 column [1,33,65,...,225], and so on. This ensures - // that threads in a warp access consecutive memory in one iteration (i.e. - // coalesced). In the above example, the warp that contains thread 0-31 - // accesses column 0-31 in the first iteration, and 32-63 in the second - // iteration, and so on. - // - // 2. Partially accumulate partial reduced results computed by threads in the - // same warp using shfl_down. Using shfl_down is faster than directly using - // atomic operations because shfl_down transfers the data between threads - // using shared memory and threads in the same warp run in lock step (thus no - // extra synchronization needed). See - // https://devblogs.nvidia.com/parallelforall/faster-parallel-reductions-kepler/ - // for details. The downside is, to produce correct results when using - // shfl_down, we need to guarantee threads in the same warp work on input - // elements with the same y, so the number of tiles in each row must be a - // multiple of 32. - // - // 3. Specialize the case that the entire tile is in bounds. When that is - // true, we don't need to emit "if(x 0; shuffle_distance /= 2) - // partial_result = Reducer( - // partial_result, - // __shfl_down_sync(CUDA_WARP_ALL, partial_result, shuffle_distance)); - // if (lane_id == 0) - // AtomicReducer(&output[y], partial_result); - // } - // - - int64 x_tile_size; - int64 z_tile_size; - std::tie(x_tile_size, z_tile_size) = - ComputeTilingSchemeForReduction(depth, width, kWarpSize); - - // Round the width in tiles up to the nearest multiple of kWarpSize, so that - // the use of shfl_down is valid. - const int64 width_in_tiles = - RoundUpToNearest(CeilOfRatio(width, x_tile_size), kWarpSize); - Shape tiled_input_shape = ShapeUtil::MakeShapeWithLayout( - reduce->shape().element_type(), - {depth / z_tile_size, height, width_in_tiles}, {2, 1, 0}); - LaunchDimensions launch_dimensions = CalculateLaunchDimensions( - tiled_input_shape, ir_emitter_context_->device_description()); - llvm::Type* index_ty = - GetIndexTypeForKernel(reduce, launch_dimensions.launch_bound(), &b_); - - auto index_typed_constant = [&](uint64 c) -> llvm::Constant* { - return llvm::ConstantInt::get(index_ty, c); - }; - - auto loop_body_emitter = [=](const IrArray::Index& tile_index) { - const int num_reduces = reducers.size(); - llvm::Type* element_ir_type = llvm_ir::PrimitiveTypeToIrType( - input_shape.element_type(), ir_emitter_context_->llvm_module()); - std::vector partial_reduction_result_addresses; - for (int i = 0; i != num_reduces; ++i) { - llvm::Value* partial_reduction_result_address = - Alloca(element_ir_type, /*ArraySize=*/nullptr, - "partial_reduction_result." + llvm::Twine(i)); - TF_ASSIGN_OR_RETURN(llvm::Value* const init_ir_value, - init_value_gens[i](IrArray::Index(index_ty))); - Store(init_ir_value, partial_reduction_result_address); - partial_reduction_result_addresses.push_back( - partial_reduction_result_address); - } - - llvm::Value* z_tile = tile_index[0]; - llvm::Value* y = tile_index[1]; - llvm::Value* x_tile = tile_index[2]; - - x_tile = ZExtOrTrunc(x_tile, index_ty); - - llvm::Value* warp_id = - UDiv(x_tile, index_typed_constant(kWarpSize), "warp_id"); - llvm::Value* lane_id = - URem(x_tile, index_typed_constant(kWarpSize), "lane_id"); - - // The x-location of the last element in this z-x-tile. - // last_x = lane_id + warpSize * (x_tile_size - 1 + warp_id * x_tile_size); - llvm::Value* last_x = NSWAdd( - lane_id, - NSWMul(index_typed_constant(kWarpSize), - NSWAdd(index_typed_constant(x_tile_size - 1), - NSWMul(warp_id, index_typed_constant(x_tile_size))))); - - KernelSupportLibrary ksl( - &b_, - /*unroll_mode=*/xla::llvm_ir::UnrollMode::kFullyUnroll, - /*prevent_vectorization=*/false); - - // Emit a for-loop that partially reduces the elements in the given - // z-x-tile. - auto emit_z_x_tile_element_loop = [&](bool x_tile_in_bounds, - int64 x_tile_loop_bound) -> Status { - auto emit_z_tile_element_loop = [&](llvm::Value* z_indvar) -> Status { - llvm::Value* z = - NSWAdd(z_indvar, NSWMul(index_typed_constant(z_tile_size), z_tile)); - TF_RETURN_IF_ERROR(ksl.For( - "x_tile", - /*start=*/index_typed_constant(0), - /*end=*/index_typed_constant(x_tile_loop_bound), - /*step=*/1, [&](llvm::Value* x_indvar) -> Status { - // x = lane_id + - // warpSize * (element_id_in_x_tile + warp_id * x_tile_size); - llvm::Value* x = NSWAdd( - lane_id, - NSWMul(index_typed_constant(kWarpSize), - NSWAdd(x_indvar, - NSWMul(warp_id, llvm::ConstantInt::get( - index_ty, x_tile_size))))); - - // Unless we know the x-tile is entirely in bounds, we have to - // emit a x-in-bounds check before reading from the input. - if (!x_tile_in_bounds) { - llvm_ir::LlvmIfData if_x_in_bounds_data = - llvm_ir::EmitIfThenElse( - ICmpULT(x, index_typed_constant(width)), "x_in_bounds", - &b_); - // Points b_ to the then-block. - llvm_ir::SetToFirstInsertPoint(if_x_in_bounds_data.true_block, - &b_); - } - - // Emit code that reads the input element and accumulates it - // to the partial reduction result. - llvm::Value* input_address = Alloca(element_ir_type); - { - // {z,y,x} is an index to input_3d_tensor_shape - // [depth,height,width]. We need to convert that to an index - // to input_shape (the shape of the operand of "reduce"). - // This conversion is composed of a transposition from - // input_shape to normalized_input_shape and a reshape from - // normalized_input_shape to input_3d_tensor_shape. - const Shape normalized_input_shape = ShapeUtil:: - MakeShapeWithDescendingLayoutAndSamePhysicalLayout( - input_shape); - auto input_shape_min2maj = - LayoutUtil::MinorToMajor(input_shape); - const std::vector transpose_dimension_mapping( - input_shape_min2maj.rbegin(), input_shape_min2maj.rend()); - const Shape input_3d_tensor_shape = - ShapeUtil::MakeShapeWithDescendingLayout( - input_shape.element_type(), {depth, height, width}); - const IrArray::Index input_3d_tensor_index( - {z, y, x}, input_3d_tensor_shape, &b_); - const IrArray::Index input_index = - input_3d_tensor_index - .SourceIndexOfReshape(input_3d_tensor_shape, - normalized_input_shape, &b_) - .SourceIndexOfTranspose( - normalized_input_shape, input_shape, - transpose_dimension_mapping, &b_); - - for (int i = 0; i != num_reduces; ++i) { - TF_ASSIGN_OR_RETURN(llvm::Value* const input_ir_value, - input_gens[i](input_index)); - Store(input_ir_value, input_address); - TF_RETURN_IF_ERROR(EmitCallToNestedComputation( - *reducers[i], - {partial_reduction_result_addresses[i], input_address}, - partial_reduction_result_addresses[i])); - } - return EmitExtraOutputsForReduce(reduce, input_index, - extra_output_gens); - } - })); - return Status::OK(); - }; - - return ksl.For("z_tile", - /*start=*/index_typed_constant(0), - /*end=*/index_typed_constant(z_tile_size), - /*step=*/1, emit_z_tile_element_loop); - }; - - llvm::Value* tile_in_bounds = - Or(b_.getInt1(width % (x_tile_size * kWarpSize) == 0), - ICmpULT(last_x, index_typed_constant(width))); - - TF_RETURN_IF_ERROR( - ksl.If(tile_in_bounds, - /*true_block_generator=*/ - [&]() -> Status { - return emit_z_x_tile_element_loop(/*x_tile_in_bounds=*/true, - x_tile_size); - }, - /*false_block_generator=*/ - [&]() -> Status { - return emit_z_x_tile_element_loop( - /*x_tile_in_bounds=*/false, - CeilOfRatio(width % (x_tile_size * kWarpSize), kWarpSize)); - })); - - // After accumulating the elements of the z_x_tile, emit calls to - // shfl_down that accumulate the partial reduction results of all - // threads in a warp. - int bit_width = llvm_ir::GetSizeInBits(element_ir_type); - // bitcast cannot be applied to aggregate types (even packed ones), so we - // instead bitcast addresses of load/store to intN* of the same bit-width. - llvm::Type* shuffle_ir_type = element_ir_type->isStructTy() - ? b_.getIntNTy(bit_width) - : element_ir_type; - for (int shuffle_distance = 16; shuffle_distance >= 1; - shuffle_distance /= 2) { - llvm::Value* result_from_other_lane = - Alloca(element_ir_type, nullptr, "result_from_other_lane"); - for (int i = 0; i != num_reduces; ++i) { - llvm::Value* partial_reduction_result = - Load(BitCast(partial_reduction_result_addresses[i], - shuffle_ir_type->getPointerTo()), - "partial_reduction_result"); - CHECK_EQ(launch_dimensions.threads_per_block() % kWarpSize, 0) - << "Requires block size a multiple of the warp size, otherwise we " - "will read undefined elements."; - Store(EmitFullWarpShuffleDown(partial_reduction_result, - b_.getInt32(shuffle_distance), &b_), - BitCast(result_from_other_lane, shuffle_ir_type->getPointerTo())); - TF_RETURN_IF_ERROR(EmitCallToNestedComputation( - *reducers[i], - {partial_reduction_result_addresses[i], result_from_other_lane}, - partial_reduction_result_addresses[i])); - } - } - - const HloInstruction* output = - reduce->IsFused() ? reduce->parent()->FusionInstruction() : reduce; - - // Emit an atomic operation that accumulates the partial reduction result of - // lane 0 (which holds the partially accumulated result for its warp) to the - // output element. - llvm_ir::LlvmIfData if_lane_id_is_zero_data = llvm_ir::EmitIfThenElse( - ICmpEQ(lane_id, index_typed_constant(0)), "lane_id_is_zero", &b_); - llvm_ir::SetToFirstInsertPoint(if_lane_id_is_zero_data.true_block, &b_); - for (int i = 0; i != num_reduces; ++i) { - llvm::Value* output_address = - GetIrArray(*output, *output, reduce_output_shapes[i]) - .EmitArrayElementAddress( - IrArray::Index(y, - ShapeUtil::GetSubshape( - output->shape(), reduce_output_shapes[i]), - &b_), - &b_, "output_element_address"); - // We don't need to emit atomic operations if there is only one tile of - // results. 'depth' is the z dimension, 'width' is the x dimension. - if (z_tile_size >= depth && x_tile_size >= width) { - TF_RETURN_IF_ERROR(EmitCallToNestedComputation( - *reducers[i], - {output_address, partial_reduction_result_addresses[i]}, - output_address)); - } else { - TF_RETURN_IF_ERROR(EmitAtomicOperationForNestedComputation( - *reducers[i], output_address, - partial_reduction_result_addresses[i])); - } - } - return Status::OK(); - }; - - // Emit a parallel loop that iterates through every input tiles. - CHECK(LastThunk()->kind() == Thunk::Kind::kSequential); - UpdateLaunchDimensions( - launch_dimensions, - static_cast(LastThunk())->thunks().back().get(), - ir_emitter_context_->llvm_module()); - return ParallelLoopEmitter(loop_body_emitter, tiled_input_shape, - launch_dimensions, &b_) - .EmitLoop(IrName(reduce), index_ty); -} - -// Figures out whether `reduce` is a row or column reduction, and which -// dimensions to reduce, and calls either `EmitRowReduction` or -// `EmitColumnReduction` as appropriate. -// Prerequisite: all the dimensions to keep are contiguous in the input layout -// and, if `reduce` is fused, the fused subgraph is pure -// elementwise. -Status IrEmitterUnnested::EmitReductionToVector( - HloInstruction* reduce, const Shape& input_shape, - absl::Span input_gens, - absl::Span init_value_gens, - absl::Span dimensions_to_reduce, - absl::Span reducers, - absl::Span reduce_output_shapes, - absl::Span> - extra_output_gens) { - // This emission requires "reduce" to have an input layout. It is either set - // by LayoutAssignment (for a top-level kReduce) or by InstructionFusion (for - // a fused kReduce). - CHECK(input_shape.has_layout()) << "LayoutAssignment or InstructionFusion " - "doesn't set the input layout of " - << reduce->ToString(); - - // Specialize multi-dimensional-array-to-vector reduction. - std::vector input_dims_to_keep; - for (int64 input_dim = 0; input_dim < ShapeUtil::Rank(input_shape); - ++input_dim) { - if (std::find(dimensions_to_reduce.begin(), dimensions_to_reduce.end(), - input_dim) == dimensions_to_reduce.end()) { - input_dims_to_keep.push_back(input_dim); - } - } - - // Sort the dimensions to keep from minor to major, to facilitate checking - // whether another dimension is major or minor of them. - std::sort(input_dims_to_keep.begin(), input_dims_to_keep.end(), - [&input_shape](int64 dim_a, int64 dim_b) { - return PositionInContainer(LayoutUtil::MinorToMajor(input_shape), - dim_a) < - PositionInContainer(LayoutUtil::MinorToMajor(input_shape), - dim_b); - }); - // Now, if output rank is at least 1, `input_dims_to_keep.front()` is - // minormost and `input_dims_to_keep.back()` is majormost. - - // If the dimensions to keep are minormost, emit a column reduction. As all - // the dimensions to keep are contiguous, by prerequisite of - // `EmitReductionToVector`, we only need to check whether the minormost - // dimension of the input is to keep. - if (input_dims_to_keep.empty()) { - return EmitReductionToScalar(reduce, input_shape, input_gens, - init_value_gens, reducers, - reduce_output_shapes, extra_output_gens); - } else if (input_dims_to_keep.front() == - LayoutUtil::Minor(input_shape.layout(), 0)) { - // Column reduction. Treat the result of "input" as a matrix whose width - // is the most minor dimension and height the product of other dimensions, - // and treat "reduce" as a column reduction of the input matrix. - const int64 width = ShapeUtil::ElementsIn(reduce->shape()); - // "width" can be zero, so don't do - // height = ShapeUtil::ElementsIn(input_shape) / width; - int64 height = 1; - for (int64 input_dim = 0; input_dim < ShapeUtil::Rank(input_shape); - ++input_dim) { - if (!std::count(input_dims_to_keep.begin(), input_dims_to_keep.end(), - input_dim)) { - height *= input_shape.dimensions(input_dim); - } - } - return EmitColumnReduction(height, width, reduce, input_shape, input_gens, - init_value_gens, reducers, reduce_output_shapes, - extra_output_gens); - } else { - // Reduce the row dimension of a matrix or reduce dimension 0 and 2 in a - // 3D tensor. The size of dimension 1 (the height) is the size of the - // dimension to keep, the size of dimension 0 (the depth) is the product - // of dimensions that are more major than the dimension to keep, and the - // size of dimension 2 (the width) is the product of more minor - // dimensions. - int64 depth = 1; - int64 width = 1; - for (int64 input_dim = 0; input_dim < ShapeUtil::Rank(input_shape); - ++input_dim) { - if (PositionInContainer(LayoutUtil::MinorToMajor(input_shape), - input_dim) > - PositionInContainer(LayoutUtil::MinorToMajor(input_shape), - input_dims_to_keep.back())) { - depth *= input_shape.dimensions(input_dim); - } else if (PositionInContainer(LayoutUtil::MinorToMajor(input_shape), - input_dim) < - PositionInContainer(LayoutUtil::MinorToMajor(input_shape), - input_dims_to_keep.front())) { - width *= input_shape.dimensions(input_dim); - } - } - const int64 height = ShapeUtil::ElementsIn(reduce->shape()); - return EmitRowReduction(depth, height, width, reduce, input_shape, - input_gens, init_value_gens, reducers, - reduce_output_shapes, extra_output_gens); - } -} - Status IrEmitterUnnested::HandleReduce(HloInstruction* reduce) { // TODO(b/112040122): Support multi-output reduce. - if (!ShapeUtil::IsArray(reduce->shape())) { + if (!reduce->shape().IsArray()) { return Unimplemented("Multi-output reduce is not supported on GPU"); } - auto input = reduce->operand(0); - auto init_value = reduce->operand(1); - absl::Span dimensions_to_reduce(reduce->dimensions()); - HloComputation* reducer = reduce->to_apply(); - // HandleReduce specializes reduction from a multi-dimensional array to a 1D - // array. The specialized version requires an initializer thunk that - // initializes the output array to the initial value of the reduce. if (IsReductionToVector(*reduce)) { - TF_ASSIGN_OR_RETURN(std::unique_ptr initializer_thunk, - BuildInitializerThunk(reduce)); - std::vector> thunks; - thunks.push_back(std::move(initializer_thunk)); - thunks.push_back( - BuildKernelThunk(reduce, /*implements_whole_instruction=*/false)); - thunk_sequence_->emplace_back( - absl::make_unique(std::move(thunks), reduce)); - - return EmitReductionToVector( - reduce, input->shape(), {[&](const IrArray::Index& index) { - return GetIrArray(*input, *reduce).EmitReadArrayElement(index, &b_); - }}, - {[&](const IrArray::Index& index) { - return GetIrArray(*init_value, *reduce) - .EmitReadArrayElement(index, &b_); - }}, - dimensions_to_reduce, {reducer}, {{}}, {}); - } - - thunk_sequence_->emplace_back( - BuildKernelThunk(reduce, /*implements_whole_instruction=*/true)); + return EmitReductionToVector(reduce); + } + return IrEmitter::HandleReduce(reduce); } @@ -1759,11 +678,11 @@ Status IrEmitterUnnested::HandleTuple(HloInstruction* tuple) { for (const HloInstruction* tuple_element : tuple->operands()) { tuple_element_buffers.push_back(GetAllocationSlice(*tuple_element)); } - thunk_sequence_->emplace_back(absl::make_unique( + AddThunkToThunkSequence(absl::make_unique( tuple_element_buffers, GetAllocationSlice(*tuple), tuple)); return Status::OK(); } - thunk_sequence_->emplace_back( + AddThunkToThunkSequence( BuildKernelThunk(tuple, /*implements_whole_instruction=*/true)); return IrEmitter::HandleTuple(tuple); } @@ -1781,8 +700,8 @@ Status IrEmitterUnnested::HandleSelectAndScatter( const auto* source = select_and_scatter->operand(1); const Window& window = select_and_scatter->window(); PrimitiveType operand_element_type = operand->shape().element_type(); - const int64 rank = ShapeUtil::Rank(operand->shape()); - CHECK_EQ(rank, ShapeUtil::Rank(source->shape())); + const int64 rank = operand->shape().rank(); + CHECK_EQ(rank, source->shape().rank()); CHECK_EQ(rank, window.dimensions_size()); TF_ASSIGN_OR_RETURN(std::unique_ptr initializer_thunk, @@ -1791,8 +710,8 @@ Status IrEmitterUnnested::HandleSelectAndScatter( thunks.push_back(std::move(initializer_thunk)); thunks.push_back(BuildKernelThunk(select_and_scatter, /*implements_whole_instruction=*/false)); - thunk_sequence_->emplace_back(absl::make_unique( - std::move(thunks), select_and_scatter)); + std::unique_ptr select_and_scatter_thunk = + absl::make_unique(std::move(thunks), select_and_scatter); // TODO(b/31410564): Implement dilation rate for select-and-scatter. if (window_util::HasDilation(window)) { @@ -1846,7 +765,7 @@ Status IrEmitterUnnested::HandleSelectAndScatter( // Create the inner loop to iterate over the window. llvm_ir::ForLoopNest window_loops(IrName(select_and_scatter, "inner"), &b_, index_type); - std::vector window_size; + DimensionVector window_size; for (const auto& dim : window.dimensions()) { window_size.push_back(dim.size()); CHECK_GT(dim.size(), 0); @@ -1958,8 +877,9 @@ Status IrEmitterUnnested::HandleSelectAndScatter( // consisting of two thunks, an initializer KernelThunk that initializes // the output and another KernelThunk that accumulates the scattered // elements. - static_cast(LastThunk())->thunks().back().get(), + select_and_scatter_thunk->thunks().back().get(), ir_emitter_context_->llvm_module()); + AddThunkToThunkSequence(std::move(select_and_scatter_thunk)); return ParallelLoopEmitter(loop_body_emitter, source->shape(), launch_dimensions, &b_) .EmitLoop(IrName(select_and_scatter), index_type); @@ -1973,10 +893,10 @@ Status IrEmitterUnnested::HandleWhile(HloInstruction* xla_while) { // Build ForThunk for conformant while loops, otherwise build WhileThunk. // TODO(b/112163966): Move trip count computation earlier in the pipeline. if (auto loop_trip_count = ComputeWhileLoopTripCount(xla_while)) { - thunk_sequence_->emplace_back(BuildForThunk(xla_while, *loop_trip_count)); + AddThunkToThunkSequence(BuildForThunk(xla_while, *loop_trip_count)); VLOG(3) << "Built ForThunk for while: " << xla_while->name(); } else { - thunk_sequence_->emplace_back(BuildWhileThunk(xla_while)); + AddThunkToThunkSequence(BuildWhileThunk(xla_while)); VLOG(3) << "Built WhileThunk for while: " << xla_while->name(); } return Status::OK(); @@ -1987,36 +907,32 @@ Status IrEmitterUnnested::HandleRng(HloInstruction* rng) { // // Unroll the kernel so that the duplicated computation that calculates the // 128 bit sample can be optimized away by LLVM. - thunk_sequence_->emplace_back( - BuildKernelThunk(rng, /*implements_whole_instruction=*/false, - ComputeMaxUnrollFactor(rng))); + std::unique_ptr rng_thunk = BuildKernelThunk( + rng, /*implements_whole_instruction=*/false, ComputeMaxUnrollFactor(rng)); ElementalIrEmitter::HloToElementGeneratorMap operand_to_generator; for (const HloInstruction* operand : rng->operands()) { operand_to_generator[operand] = [=](const llvm_ir::IrArray::Index& index) { return GetIrArray(*operand, *rng).EmitReadArrayElement(index, &b_); }; } - TF_RETURN_IF_ERROR(EmitTargetElementLoop( - *rng, GpuElementalIrEmitter(hlo_module_config_, module_, &b_, - GetNestedComputer()) - .MakeElementGenerator(rng, operand_to_generator))); - std::unique_ptr rng_thunk = std::move(thunk_sequence_->back()); - thunk_sequence_->pop_back(); + TF_RETURN_IF_ERROR(EmitTargetElementLoopInThunk( + *rng, + GpuElementalIrEmitter(hlo_module_config_, module_, &b_, + GetNestedComputer()) + .MakeElementGenerator(rng, operand_to_generator), + rng_thunk.get())); // Emit a kernel to increment the global state for Philox RNG algorithm. - thunk_sequence_->emplace_back( - BuildKernelThunk(rng, /*implements_whole_instruction=*/false)); - llvm_ir::IncrementVariableForPhiloxRngState(1, module_, &b_); std::unique_ptr increment_seed_thunk = - std::move(thunk_sequence_->back()); - thunk_sequence_->pop_back(); + BuildKernelThunk(rng, /*implements_whole_instruction=*/false); + llvm_ir::IncrementVariableForPhiloxRngState(1, module_, &b_); // Build the SequentialThunk for the RNG hlo. std::vector> thunks; thunks.reserve(2); thunks.push_back(std::move(rng_thunk)); thunks.push_back(std::move(increment_seed_thunk)); - thunk_sequence_->emplace_back( + AddThunkToThunkSequence( absl::make_unique(std::move(thunks), rng)); return Status::OK(); @@ -2058,11 +974,12 @@ Status IrEmitterUnnested::HandleScatter(HloInstruction* scatter) { // Elide the sequential thunk if there's no copy. if (thunks.size() == 1) { - thunk_sequence_->push_back(std::move(thunks[0])); + AddThunkToThunkSequence(std::move(thunks[0])); } else { - thunk_sequence_->emplace_back( + AddThunkToThunkSequence( absl::make_unique(std::move(thunks), scatter)); } + return Status::OK(); } @@ -2100,7 +1017,7 @@ Status IrEmitterUnnested::EmitScatter( int64 raw_window_multidim_idx = 0; std::vector input_window_multidim; std::vector input_window_bounds; - for (int64 i = 0, e = ShapeUtil::Rank(operand->shape()); i != e; ++i) { + for (int64 i = 0, e = operand->shape().rank(); i != e; ++i) { if (absl::c_binary_search(dim_numbers.inserted_window_dims(), i)) { input_window_bounds.push_back(1); // Trivial dimension. input_window_multidim.push_back(index.GetConstantWithIndexType(0)); @@ -2112,12 +1029,11 @@ Status IrEmitterUnnested::EmitScatter( ++raw_window_multidim_idx; } } - DCHECK_EQ(input_window_multidim.size(), ShapeUtil::Rank(operand->shape())); + DCHECK_EQ(input_window_multidim.size(), operand->shape().rank()); // Insert a 1 dimension at the end if index_vector_dim requests one. Shape scatter_indices_shape = scatter_indices->shape(); - if (dim_numbers.index_vector_dim() == - ShapeUtil::Rank(scatter_indices_shape)) { + if (dim_numbers.index_vector_dim() == scatter_indices_shape.rank()) { scatter_indices_shape.add_dimensions(1); scatter_indices_shape.mutable_layout()->add_minor_to_major( dim_numbers.index_vector_dim()); @@ -2195,15 +1111,24 @@ Status IrEmitterUnnested::EmitScatter( } Status IrEmitterUnnested::HandleSelect(HloInstruction* select) { - thunk_sequence_->push_back( - BuildKernelThunk(select, /*implements_whole_instruction=*/true)); return IrEmitter::HandleSelect(select); } 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 @@ -2228,10 +1153,10 @@ Status IrEmitterUnnested::HandleSort(HloInstruction* sort) { } } - int64 dimension_to_sort = sort->dimensions(0); - int64 dimension_to_sort_bound = keys_shape.dimensions(dimension_to_sort); + uint64 dimension_to_sort_bound = keys_shape.dimensions(dimension_to_sort); int64 num_stages = tensorflow::Log2Ceiling(dimension_to_sort_bound); - auto index_type = b_.getInt64Ty(); + CHECK_GE(1ULL << num_stages, dimension_to_sort_bound); + CHECK_LT(1ULL << (num_stages - 1), dimension_to_sort_bound); // Naive C++ code for the outer loops: // @@ -2245,59 +1170,137 @@ Status IrEmitterUnnested::HandleSort(HloInstruction* sort) { // } // } // - // This follows the algorithm described on Wikipedia: - // https://en.wikipedia.org/wiki/Bitonic_sorter + // This follows the alternative representation of the algorithm described on + // Wikipedia: https://en.wikipedia.org/wiki/Bitonic_sorter + // + // Each mask specifies how to derive from one position in the array the + // position with which it should be compared (we calculate the xor of the + // position with the mask). + // As an optimization, we can move the 'mask' loop to inside the + // sorting/comparison loop if the comparisons happen within a small block of + // the array. To make this work, we collect all consecutive masks that are + // smaller than our chosen power of 2 tile size, and pass them to SortInPlace. + // Each thread then processes one tile of data. + + const uint64 kTileSize = std::min(2048ULL, 1ULL << num_stages); + + // If we cannot combine several xor masks together, we don't use tiling, so we + // calculate the standard launch dimensions for the shape. However we only + // need to iterate through ~half of the dimension to sort (rounded up to the + // next highest power of 2), because each iteration compares one pair of + // elements. + Shape standard_iteration_shape = keys_shape; + uint64 standard_num_iterations_in_sort_dim = 1ULL << (num_stages - 1); + standard_iteration_shape.set_dimensions(dimension_to_sort, + standard_num_iterations_in_sort_dim); + LaunchDimensions standard_launch_dimensions = CalculateLaunchDimensions( + standard_iteration_shape, ir_emitter_context_->device_description()); + + // Calculate the launch dimensions for the case where we use tiling. We split + // the dimension that should be sorted into tiles of size 'kTileSize'. This + // means we first need to round 'dimension_to_sort_bound' up to be a multiple + // of the tile size. + int64 rounded_bound = RoundUpToNearest(dimension_to_sort_bound, kTileSize); + Shape iteration_shape = keys_shape; + + // We iterate through the element pairs that should be compared. + uint64 num_iterations_in_sort_dim = rounded_bound / 2; + iteration_shape.set_dimensions(dimension_to_sort, num_iterations_in_sort_dim); + uint64 num_iterations = ShapeUtil::ElementsIn(iteration_shape); + + // For correctness reasons we need exactly 'kTileSize' / 2 many threads per + // block. Each thread is responsible for copying exactly two adjacent elements + // into shared memory, and then does a comparison of two possibly different + // elements taken from shared memory. + const uint64 kThreadsPerBlock = kTileSize / 2; + + // Check whether we should use any tiling. We might not be able to use it if + // we have not enough threads, or not enough shared memory. Also it does not + // give a speedup if the tile size is < 128. + int64 total_shared_memory_needed = 0; + for (int64 i = 0; i < sort->operand_count(); ++i) { + total_shared_memory_needed += + kTileSize * ShapeUtil::ByteSizeOfPrimitiveType( + sort->operand(i)->shape().element_type()); + } + bool no_tiling = + kTileSize < 128 || + kThreadsPerBlock > + ir_emitter_context_->device_description().threads_per_block_limit() || + total_shared_memory_needed > + ir_emitter_context_->device_description().shared_memory_per_block(); + uint64 num_blocks = CeilOfRatio(num_iterations, kThreadsPerBlock); + LaunchDimensions tiled_launch_dimensions(num_blocks, kThreadsPerBlock); + + auto emit_kernel = [&](absl::Span xor_masks) { + thunks.push_back( + BuildKernelThunk(sort, /*implements_whole_instruction=*/false)); + LaunchDimensions launch_dimensions = xor_masks.size() > 1 + ? tiled_launch_dimensions + : 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); + 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)); + } + } + return llvm_ir::EmitSortInPlace( + dimension_to_sort, keys_array, values_arrays, + iota_values_parameter_index, IrName(sort), xor_masks, &b_, + launch_dimensions, + xor_masks.size() > 1 ? num_iterations_in_sort_dim + : standard_num_iterations_in_sort_dim, + kTileSize); + }; + std::vector xor_masks; for (int64 stage = 0; stage < num_stages; ++stage) { for (int64 mask = stage; mask >= 0; --mask) { - thunks.push_back( - BuildKernelThunk(sort, /*implements_whole_instruction=*/false)); - LaunchDimensions launch_dimensions = CalculateLaunchDimensions( - keys_shape, ir_emitter_context_->device_description()); - UpdateLaunchDimensions(launch_dimensions, thunks.back().get(), - ir_emitter_context_->llvm_module()); - - llvm::Value* xor_mask; + int64 xor_mask; if (mask == stage) { - xor_mask = llvm::ConstantInt::get(index_type, (1LL << (stage + 1)) - 1); + xor_mask = (1LL << (stage + 1)) - 1; } else { - xor_mask = llvm::ConstantInt::get(index_type, 1LL << mask); + xor_mask = 1LL << mask; } - - IrArray keys_array; - std::vector values_arrays; - values_arrays.reserve(sort->operand_count() - 1); - 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)); + if (xor_mask >= kTileSize || no_tiling) { + if (!xor_masks.empty()) { + TF_RETURN_IF_ERROR(emit_kernel(xor_masks)); + xor_masks.clear(); } + TF_RETURN_IF_ERROR(emit_kernel({xor_mask})); + } else { + xor_masks.push_back(xor_mask); } - TF_RETURN_IF_ERROR(llvm_ir::EmitSortInPlace( - dimension_to_sort, keys_array, values_arrays, IrName(sort), xor_mask, - &b_, &launch_dimensions)); } } + if (!xor_masks.empty()) { + TF_RETURN_IF_ERROR(emit_kernel(xor_masks)); + } - thunk_sequence_->emplace_back( + AddThunkToThunkSequence( absl::make_unique(std::move(thunks), sort)); return Status::OK(); } Status IrEmitterUnnested::HandleTupleSelect(HloInstruction* tuple_select) { - thunk_sequence_->push_back( + AddThunkToThunkSequence( BuildKernelThunk(tuple_select, /*implements_whole_instruction=*/true)); return IrEmitter::HandleTupleSelect(tuple_select); } -Status IrEmitterUnnested::HandleCrossReplicaSum(HloInstruction* crs) { +Status IrEmitterUnnested::HandleAllReduce(HloInstruction* crs) { if (hlo_module_config_.replica_count() != 1) { // TODO(b/33011107): Support nontrivial cross replica sum on GPU. return Unimplemented( - "CrossReplicaSum with >1 replica is not implemented on GPU."); + "AllReduce with >1 replica is not implemented on GPU."); } // CRS with one operand and one replica is simply the identity function. @@ -2308,9 +1311,9 @@ Status IrEmitterUnnested::HandleCrossReplicaSum(HloInstruction* crs) { // HloModuleConfig::num_replicas changes between when the module is compiled // and when it's run. if (crs->operand_count() == 1) { - CHECK(ShapeUtil::IsArray(crs->operand(0)->shape())) - << "Operands to cross-replica-sum must be arrays: " << crs->ToString(); - thunk_sequence_->push_back(absl::make_unique( + CHECK(crs->operand(0)->shape().IsArray()) + << "Operands to all-reduce must be arrays: " << crs->ToString(); + AddThunkToThunkSequence(absl::make_unique( /*source_address=*/GetAllocationSlice(*crs->operand(0)), /*destination_buffer=*/GetAllocationSlice(*crs), /*mem_size=*/ShapeUtil::ByteSizeOf(crs->shape()), crs)); @@ -2334,22 +1337,22 @@ Status IrEmitterUnnested::HandleCrossReplicaSum(HloInstruction* crs) { // Output a tuple of the buffers above. thunks.push_back(absl::make_unique( tuple_element_buffers, GetAllocationSlice(*crs), nullptr)); - thunk_sequence_->push_back( + AddThunkToThunkSequence( absl::make_unique(std::move(thunks), crs)); return Status::OK(); } -Status IrEmitterUnnested::HandleAfterAll(HloInstruction* gen_token) { +Status IrEmitterUnnested::HandleAfterAll(HloInstruction* after_all) { return Status::OK(); } Status IrEmitterUnnested::HandleInfeed(HloInstruction* infeed) { - thunk_sequence_->emplace_back(BuildInfeedThunk(infeed)); + AddThunkToThunkSequence(BuildInfeedThunk(infeed)); return Status::OK(); } Status IrEmitterUnnested::HandleOutfeed(HloInstruction* outfeed) { - thunk_sequence_->emplace_back(BuildOutfeedThunk(outfeed)); + AddThunkToThunkSequence(BuildOutfeedThunk(outfeed)); return Status::OK(); } @@ -2507,10 +1510,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); @@ -2658,28 +1661,43 @@ std::unique_ptr IrEmitterUnnested::BuildGemmThunk( rhs->shape(), // The shape of RHS. inst->shape(), // The shape of the output. 1.0, // alpha. - inst); + 0.0, // beta. + inst, /*implements_whole_instruction=*/true); } if (inst->opcode() == HloOpcode::kFusion) { CHECK_EQ(inst->fusion_kind(), HloInstruction::FusionKind::kOutput); - const HloInstruction* mul = inst->fused_expression_root(); - const HloInstruction* dot = mul->operand(0); - const HloInstruction* alpha = mul->operand(1); - if (dot->opcode() != HloOpcode::kDot) { - std::swap(dot, alpha); - } - if (alpha->opcode() == HloOpcode::kBroadcast) { - alpha = alpha->operand(0); - } - if (alpha->opcode() == HloOpcode::kParameter) { - alpha = inst->operand(alpha->parameter_number()); - } - // TODO(b/74185543): Remove the following if block once we support fusion - // with a non-constant as well. Then we will just always use the constant - // on the device. - if (alpha->opcode() == HloOpcode::kCopy) { - alpha = alpha->operand(0); + const HloInstruction* output_fused_op = inst->fused_expression_root(); + + double alpha_value = 1.0; + const HloInstruction* bias = nullptr; + const HloInstruction* dot = output_fused_op->operand(0); + if (output_fused_op->opcode() == HloOpcode::kMultiply) { + const HloInstruction* alpha = output_fused_op->operand(1); + if (dot->opcode() != HloOpcode::kDot) { + std::swap(dot, alpha); + } + if (alpha->opcode() == HloOpcode::kBroadcast) { + alpha = alpha->operand(0); + } + if (alpha->opcode() == HloOpcode::kParameter) { + alpha = inst->operand(alpha->parameter_number()); + } + // TODO(b/74185543): Remove the following if block once we support fusion + // with a non-constant as well. Then we will just always use the constant + // on the device. + if (alpha->opcode() == HloOpcode::kCopy) { + alpha = alpha->operand(0); + } + alpha_value = GetScalarConstantAsDouble(alpha->literal()); + } else { + // Fused bias add. + CHECK_EQ(output_fused_op->opcode(), HloOpcode::kAdd); + bias = output_fused_op->operand(1); + if (dot->opcode() != HloOpcode::kDot) { + std::swap(dot, bias); + } + bias = inst->operand(bias->parameter_number()); } DCHECK(dot->opcode() == HloOpcode::kDot); @@ -2692,15 +1710,38 @@ std::unique_ptr IrEmitterUnnested::BuildGemmThunk( const HloInstruction* rhs = inst->operand(rhs_parameter->parameter_number()); + // The bias is passed inside the output buffer. If those buffers are shared + // we can just use it, otherwise copy the bias values into the output buffer + // first. + if (bias != nullptr && + GetAllocationSlice(*bias) != GetAllocationSlice(*inst)) { + std::vector> thunks; + thunks.push_back(absl::make_unique( + /*source_buffer=*/GetAllocationSlice(*bias), + /*destination_buffer=*/GetAllocationSlice(*inst), + /*mem_size=*/ShapeUtil::ByteSizeOf(inst->shape()), nullptr)); + thunks.push_back(absl::make_unique( + GetAllocationSlice(*lhs), // The buffer assigned to LHS. + GetAllocationSlice(*rhs), // The buffer assigned to RHS. + GetAllocationSlice(*inst), // The output buffer. + lhs->shape(), // The shape of LHS. + rhs->shape(), // The shape of RHS. + inst->shape(), // The shape of the output. + alpha_value, // alpha. + 1.0, // beta. + inst, /*implements_whole_instruction=*/false)); + return absl::make_unique(std::move(thunks), inst); + } return absl::make_unique( - GetAllocationSlice(*lhs), // The buffer assigned to LHS. - GetAllocationSlice(*rhs), // The buffer assigned to RHS. - GetAllocationSlice(*inst), // The output buffer. - lhs->shape(), // The shape of LHS. - rhs->shape(), // The shape of RHS. - inst->shape(), // The shape of the output. - GetScalarConstantAsDouble(alpha->literal()), // alpha. - inst); + GetAllocationSlice(*lhs), // The buffer assigned to LHS. + GetAllocationSlice(*rhs), // The buffer assigned to RHS. + GetAllocationSlice(*inst), // The output buffer. + lhs->shape(), // The shape of LHS. + rhs->shape(), // The shape of RHS. + inst->shape(), // The shape of the output. + alpha_value, // alpha. + bias != nullptr ? 1.0 : 0.0, // beta. + inst, /*implements_whole_instruction=*/true); } LOG(FATAL) << "Cannot build a GemmThunk for " << inst->ToString(); @@ -2809,15 +1850,12 @@ StatusOr> IrEmitterUnnested::BuildInitializerThunk( if (fused) { // If init_value was fused into this reduce we have to generate it first. - std::vector parameter_arrays; - for (HloInstruction* operand : hlo->operands()) { - parameter_arrays.push_back(GetIrArray(*operand, *hlo)); - } GpuElementalIrEmitter elemental_emitter(hlo_module_config_, ir_emitter_context_->llvm_module(), &b_, GetNestedComputer()); - FusedIrEmitter fused_emitter(parameter_arrays, &elemental_emitter); + FusedIrEmitter fused_emitter(GetGeneratorForOperandIrArrays(hlo), + &elemental_emitter); TF_RETURN_IF_ERROR(init_value_operand->Accept(&fused_emitter)); TF_RETURN_IF_ERROR( ParallelLoopEmitter(fused_emitter.GetGenerator(init_value_operand), @@ -3022,8 +2060,16 @@ Status IrEmitterUnnested::EmitTargetElementLoopInThunk( GetIndexTypeForKernel(&hlo, launch_dimensions.launch_bound(), &b_)); } - // For multioutput fusion, we need to emit each operand and the root. + // Emit the tuple pointers in one thread. We could do this at any point in + // the kernel, but we do it at the beginning in the hopes of reducing register + // pressure, since we touch threadIdx.x and blockIdx.x at the beginning of the + // kernel *anyway*. std::vector output_arrays = ConstructIrArrayForOutputs(hlo); + KernelSupportLibrary{&b_}.If("emit_mof_tuple", IsBlock0Thread0(&b_), [&] { + llvm_ir::EmitTuple(GetIrArray(hlo, hlo), output_arrays, &b_, module_); + }); + + // For multioutput fusion, we need to emit each operand and the root. TF_RETURN_IF_ERROR( ParallelLoopEmitter(element_generator, output_arrays, launch_dimensions, &b_, unroll_factor) @@ -3032,17 +2078,49 @@ Status IrEmitterUnnested::EmitTargetElementLoopInThunk( &hlo, launch_dimensions.launch_bound(), &b_))); b_.SetInsertPoint(b_.GetInsertBlock()->getTerminator()); - llvm_ir::EmitTuple(GetIrArray(hlo, hlo), output_arrays, &b_, module_); - 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) { - CHECK_EQ(Thunk::Kind::kKernel, LastThunk()->kind()); - return EmitTargetElementLoopInThunk(hlo, element_generator, - static_cast(LastThunk())); + int unroll_factor = 1; + // Unfused elementwise operations are usually memory bound, unroll them. + if (hlo.IsElementwise() || + (hlo.opcode() == HloOpcode::kFusion && !MayPreventVectorization(hlo))) { + unroll_factor = ComputeMaxUnrollFactor(&hlo); + } + + std::unique_ptr kernel_thunk = BuildKernelThunk( + &hlo, /*implements_whole_instruction=*/true, unroll_factor); + Status emit_status = + EmitTargetElementLoopInThunk(hlo, element_generator, kernel_thunk.get()); + thunk_sequence_->emplace_back(std::move(kernel_thunk)); + + return emit_status; } std::vector IrEmitterUnnested::ConstructIrArrayForInputs( @@ -3055,31 +2133,6 @@ std::vector IrEmitterUnnested::ConstructIrArrayForInputs( return param_arrays; } -int IrEmitterUnnested::ConstructOutputReducedShapeAndCastOutputIrArrayToShape( - const HloInstruction& hlo, const std::vector& output_arrays, - absl::Span reduced_output_dims, - std::vector* output_reduced_shapes, - std::vector* output_in_reduced_shape_arrays) { - int64 num_outputs = 1; - if (hlo.IsMultiOutputFusion()) { - num_outputs = ShapeUtil::TupleElementCount(hlo.shape()); - output_in_reduced_shape_arrays->reserve(num_outputs); - output_reduced_shapes->reserve(num_outputs); - for (int64 i = 0; i < num_outputs; ++i) { - output_reduced_shapes->push_back(ShapeUtil::MakeShapeWithDescendingLayout( - ShapeUtil::GetSubshape(hlo.shape(), {i}).element_type(), - reduced_output_dims)); - output_in_reduced_shape_arrays->push_back( - output_arrays[i].CastToShape((*output_reduced_shapes)[i], &b_)); - } - } else { - output_reduced_shapes->push_back(ShapeUtil::MakeShapeWithDescendingLayout( - hlo.shape().element_type(), reduced_output_dims)); - output_in_reduced_shape_arrays->push_back( - output_arrays[0].CastToShape((*output_reduced_shapes)[0], &b_)); - } - return num_outputs; -} int IrEmitterUnnested::ConstructInputReducedShapeAndCastInputIrArrayToShape( const HloInstruction& hlo, const std::vector& param_arrays, @@ -3108,338 +2161,1082 @@ int IrEmitterUnnested::ConstructInputReducedShapeAndCastInputIrArrayToShape( namespace { -// Reads thread_idx.x and converts it to a (y,x) coordinate, assuming that the -// thread lives within a square tile of size tile_size (so thread blocks are of -// size tile_size * tile_size). -std::tuple CalculateYXCoordinateWithinTile( - llvm::IRBuilder<>* builder, llvm::Value* tile_size, - int64 threads_per_tile) { - // Calculate the starting element coordinate within a tile for the current - // thread, (y, x) from thread_id. - llvm::Value* thread_id = llvm_ir::EmitCallToIntrinsic( - llvm::Intrinsic::nvvm_read_ptx_sreg_tid_x, {}, {}, builder); - llvm_ir::AddRangeMetadata(0, threads_per_tile, - llvm::cast(thread_id)); - thread_id = builder->CreateIntCast(thread_id, tile_size->getType(), - /*isSigned=*/true, "thread.id.x"); - auto x = builder->CreateURem(thread_id, tile_size); - auto y = builder->CreateUDiv(thread_id, tile_size); - return std::make_tuple(y, x); -} - -// Reads block_idx.x, casts it to type index_ty, and adds the assumption that -// it's in the range [0, num_blocks]. -llvm::Value* GetBlockIdx(llvm::IRBuilder<>* builder, llvm::Type* index_ty, - int64 num_blocks) { - llvm::Value* block_id = llvm_ir::EmitCallToIntrinsic( - llvm::Intrinsic::nvvm_read_ptx_sreg_ctaid_x, {}, {}, builder); - llvm_ir::AddRangeMetadata(0, num_blocks, - llvm::cast(block_id)); - return builder->CreateIntCast(block_id, index_ty, /*isSigned=*/true, - "block.id.x"); -} - -// Emits code to process up to (tile_size/num_rows) elements in a tile, given -// `emit_elem_function` is the function to emit code to process one element, `y` -// and `x` are the coordinates for the first element to process, and `index` is -// the index for the origin of the tile. Emits bounds check to ensure that each -// processed element is within the boundary defined by `tile_width` and -// `tile_height`. -void EmitTiledElementalCodeWithBoundsCheck( - int64 tile_size, int64 num_rows, const IrArray::Index& index, - const string& loop_name, KernelSupportLibrary* ksl, - llvm::IRBuilder<>* builder, llvm::Value* y, llvm::Value* x, - llvm::Value* tile_width, llvm::Value* tile_height, - 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 = 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 / 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 * 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 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(); + + 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), + [&] { + // tile_height_bound = + // ceil(tile_height / num_threads_y) * num_threads_y + llvm::Value* ceiling_of_ratio = builder->CreateUDiv( + builder->CreateAdd(tile_height, llvm::ConstantInt::get( + index_ty, num_threads_y - 1)), + llvm::ConstantInt::get(index_ty, num_threads_y)); + llvm::Value* tile_height_bound = builder->CreateMul( + ceiling_of_ratio, + llvm::ConstantInt::get(index_ty, num_threads_y)); + ksl->For( + loop_name, /*start=*/llvm::ConstantInt::get(index_ty, 0), + /*end=*/tile_height_bound, + /*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_x.AddOffsetToDim( + y_indvar, KernelMappingScheme::DimY, builder), + y_loc, x_loc, j); + }); + }); + }); + } +} + +// Emits code to process up to +// (tile_size_x/num_threads_x * tile_size_y/num_threads_y) elements in a tile, +// given `emit_elem_function` is the function to emit code to process one +// element, `y` and `x` are the intra-tile coordinates for the first element +// to process, and `index` is the index for the origin of the tile. Information +// about tile_size_x/y and num_threads_x/y are stored in `mapping_scheme`. Emits +// bounds check to ensure that each processed element is within the boundary +// defined by `tile_width` and `tile_height`. +void EmitTiledElementalCodeWithBoundsCheck( + 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, + 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(); - // Emits a constant value with index type. + + ksl->If( + loop_name + "_full_tile", + builder->CreateAnd( + builder->CreateICmpEQ(llvm::ConstantInt::get(index_ty, tile_size_x), + tile_width), + builder->CreateICmpEQ(llvm::ConstantInt::get(index_ty, tile_size_y), + tile_height)), + [&] { + EmitFullElementalTile(mapping_scheme, tile_origin_index, loop_name, ksl, + builder, y, x, index_ty, emit_elem_function); + }, + [&] { + EmitPartialElementalTile(mapping_scheme, tile_origin_index, loop_name, + ksl, builder, y, x, tile_height, tile_width, + index_ty, emit_elem_function); + }); +} +} // namespace + +// Emits code to process a tensor element in a tile for the given kCopy HLO that +// performs a 0-2-1 transpose. +// +// index: The index for the first output element in the normalized tensor. The +// normalized tensor is the resulting tensor after collapsing contiguous +// dimensions that play the same role in the transpose. +// 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. +void IrEmitterUnnested::EmitTileElementForCopy( + HloInstruction* hlo, const llvm_ir::IrArray::Index& index, + const KernelCodegenInfo* kernel_info, llvm::Value* y_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. + llvm::Instruction* load_from_shmem_buffer = + Load(GEP(tiled_param_info->GetBufferForParameter(0), + {b_.getInt64(0), x_loc, y_loc}), + "output_element"); + llvm_ir::IrArray output_array = GetIrArray(*hlo, *hlo); + Shape output_reduced_shape = ShapeUtil::MakeShapeWithDescendingLayout( + hlo->shape().element_type(), + kernel_info->GetKernelMappingScheme()->GetDimensionsInElements()); + // When the output_reduced_shape is a 0-2-1 transpose of the input shape, + // the 0-2-1 transpose is achieved through EmitWriteArrayElement. + output_array.CastToShape(output_reduced_shape, &b_) + .EmitWriteArrayElement(index, load_from_shmem_buffer, &b_); +} + +// 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. +// +// index: The index for the first output element in the normalized tensor, that +// is the resulting tensor after collapsing contiguous dimensions that play +// the same role in the transpose. +// kernel_info: Other information to support the kernel code generation. +// y_loc: The y coordinate within a tile. +// x_loc: The x coordinate within a tile. +void IrEmitterUnnested::EmitTileElementForFusion( + HloInstruction* hlo, const llvm_ir::IrArray::Index& index, + const KernelCodegenInfo* kernel_info, llvm::Value* y_loc, + llvm::Value* x_loc, int64 /*x_iter_num*/) { + llvm_ir::TiledParameterInfo* tiled_param_info = + kernel_info->GetTiledParameterInfo(); + std::vector output_arrays = ConstructIrArrayForOutputs(*hlo); + GpuElementalIrEmitter elem_emitter(hlo_module_config_, module_, &b_, + GetNestedComputer()); + FusedIrEmitter fused_emitter(GetGeneratorForOperandIrArrays(hlo), + &elem_emitter); + tiled_param_info->set_y(y_loc); + tiled_param_info->set_x(x_loc); + fused_emitter.SetTiledParameterInfo(tiled_param_info); + TF_CHECK_OK(hlo->fused_expression_root()->Accept(&fused_emitter)); + IrArray::Index untiled_index = + kernel_info->GetKernelMappingScheme()->GetUnnormalizedIndex( + index, output_arrays[0].GetShape()); + const llvm_ir::ElementGenerator& output_generator = + fused_emitter.GetRootGenerator(); + llvm::Value* output_value = output_generator(untiled_index).ValueOrDie(); + if (hlo->IsMultiOutputFusion()) { + DCHECK(output_value->getType()->isStructTy()); + DCHECK_EQ(output_value->getType()->getStructNumElements(), + output_arrays.size()); + for (int64 i = 0; i < output_arrays.size(); ++i) { + output_arrays[i].EmitWriteArrayElement( + untiled_index, ExtractValue(output_value, i), &b_); + } + } else { + output_arrays[0].EmitWriteArrayElement(untiled_index, output_value, &b_); + } +} + +// Information to support the code generation for a tiled reduction kernel. +using AddressVector = InlinedVector; +class ReductionCodegenInfo : public IrEmitterUnnested::KernelCodegenInfo { + public: + explicit ReductionCodegenInfo(llvm_ir::KernelMappingScheme* mapping_scheme, + bool is_row_reduction) + : KernelCodegenInfo(mapping_scheme), + current_output_linear_index_address_(nullptr), + current_output_inbound_address_(nullptr), + is_row_reduction_(is_row_reduction) {} + + void SetCurrentOutputLinearIndexAddress(llvm::AllocaInst* a) { + current_output_linear_index_address_ = a; + } + // Returns the address of the memory that stores the linear index of the + // current output. Since we are processing reduction to contiguous physical + // dimensions, this linear index is the linear index of the 1D output array. + llvm::AllocaInst* GetCurrentOutputLinearIndexAddress() const { + return current_output_linear_index_address_; + } + + void SetCurrentOutputInboundAddress(llvm::AllocaInst* a) { + current_output_inbound_address_ = a; + } + + llvm::AllocaInst* GetCurrentOutputInboundAddress() const { + return current_output_inbound_address_; + } + + AddressVector* GetMutablePartialResultAddresses() { + return &partial_result_addresses_; + } + absl::Span GetPartialResultAddresses() const { + return partial_result_addresses_; + } + + AddressVector* GetMutableReductionInputAddresses() { + return &reduction_input_addresses_; + } + absl::Span GetReductionInputAddresses() const { + return reduction_input_addresses_; + } + + InlinedVector* GetMutableReducers() { return &reducers_; } + const InlinedVector& GetReducers() const { + return reducers_; + } + int GetNumberOfReduces() const { return reducers_.size(); } + + InlinedVector* GetMutableReductionOutputShapeIndices() { + return &reduction_output_shape_indices_; + } + absl::Span GetReductionOutputShapeIndices() const { + return reduction_output_shape_indices_; + } + + bool IsRowReduction() const { return is_row_reduction_; } + + // Return the dimension that is being reduced between DimX and DimY. + int GetReducedDimensionEnum() const { + return IsRowReduction() ? llvm_ir::KernelMappingScheme::DimX + : llvm_ir::KernelMappingScheme::DimY; + } + + // Return the dimension that is being ketp between DimX and DimY. + int GetKeptDimensionEnum() const { + return IsRowReduction() ? llvm_ir::KernelMappingScheme::DimY + : 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_; + InlinedVector reducers_; + InlinedVector reduction_output_shape_indices_; + llvm::AllocaInst* current_output_linear_index_address_; + llvm::AllocaInst* current_output_inbound_address_; + bool is_row_reduction_; +}; + +namespace { +// Returns a group of instructions that generate the output for the kernel +// containing the given HLO instruction. The result may be an unnested kReduce +// HLO, a nested kReduce HLO of a kInput fusion, or the operands of the tuple +// for a multiple output fusion. +absl::Span GetOutputInstructions( + HloInstruction* const* reduce_or_tuple_pointer) { + HloOpcode opcode = (*reduce_or_tuple_pointer)->opcode(); + CHECK(opcode == HloOpcode::kReduce || opcode == HloOpcode::kTuple); + return opcode == HloOpcode::kTuple + ? (*reduce_or_tuple_pointer)->operands() + : absl::Span(reduce_or_tuple_pointer, 1); +} + +const HloInstruction* GetFirstReduceInstruction( + absl::Span instructions) { + auto first_reduce_iter = + absl::c_find_if(instructions, [](const HloInstruction* inst) { + return inst->opcode() == HloOpcode::kReduce; + }); + CHECK_NE(first_reduce_iter, instructions.end()); + return *first_reduce_iter; +} + +}; // namespace + +void IrEmitterUnnested::EmitPrologueForOneReduction( + HloInstruction* unnested_hlo, HloInstruction* reduce_inst, int reduce_idx, + KernelCodegenInfo* kernel_info, GpuElementalIrEmitter* elemental_emitter, + ShapeIndex output_shape_index) { + ReductionCodegenInfo* reduction_info = + static_cast(kernel_info); + + InlinedVector* reducers = + reduction_info->GetMutableReducers(); + CHECK(IsReductionToVector(*reduce_inst)); + reducers->push_back(reduce_inst->to_apply()); + + InlinedVector* reduction_output_shape_indices = + reduction_info->GetMutableReductionOutputShapeIndices(); + reduction_output_shape_indices->push_back(std::move(output_shape_index)); + + AddressVector* reduction_input_addresses = + reduction_info->GetMutableReductionInputAddresses(); + llvm::Type* element_type = llvm_ir::PrimitiveTypeToIrType( + reduce_inst->shape().element_type(), ir_emitter_context_->llvm_module()); + 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=*/b_.getInt32(num_partial_results), + "partial_reduction_result." + llvm::Twine(reduce_idx)); + partial_result_addresses->push_back(partial_result_address); + + // Initialize the partial result with the initial value of the reduction. + llvm::Value* init_ir_value; + if (unnested_hlo->opcode() == HloOpcode::kFusion) { + HloInstruction* init_value_operand = reduce_inst->mutable_operand(1); + FusedIrEmitter fused_emitter(GetGeneratorForOperandIrArrays(unnested_hlo), + elemental_emitter); + + TF_CHECK_OK(init_value_operand->Accept(&fused_emitter)); + init_ir_value = + fused_emitter + .GetGenerator(init_value_operand)(IrArray::Index(b_.getInt32Ty())) + .ValueOrDie(); + } else { + const HloInstruction* init_value = unnested_hlo->operand(1); + init_ir_value = + GetIrArray(*init_value, *unnested_hlo) + .EmitReadArrayElement(IrArray::Index(b_.getInt32Ty()), &b_); + } + + for (int i = 0; i < num_partial_results; ++i) { + Store(init_ir_value, InBoundsGEP(partial_result_address, {b_.getInt32(i)})); + } +} + +void IrEmitterUnnested::EmitPrologueForReduction( + HloInstruction* unnested_hlo, KernelCodegenInfo* kernel_info) { + VLOG(10) << "Emit prologue for reduction " << unnested_hlo->ToString(); + // Find the unnested kReduce or the tuple that contains a list of kReduce. + HloInstruction* reduce_or_tuple = unnested_hlo->opcode() == HloOpcode::kFusion + ? unnested_hlo->fused_expression_root() + : unnested_hlo; + absl::Span output_instructions = + GetOutputInstructions(&reduce_or_tuple); + ReductionCodegenInfo* reduction_info = + static_cast(kernel_info); + GpuElementalIrEmitter elemental_emitter(hlo_module_config_, + ir_emitter_context_->llvm_module(), + &b_, GetNestedComputer()); + const HloInstruction* first_reduce = nullptr; + for (int i = 0, e = output_instructions.size(); i != e; ++i) { + if (output_instructions[i]->opcode() != HloOpcode::kReduce) { + continue; + } + HloInstruction* reduce_inst = output_instructions[i]; + if (first_reduce == nullptr) { + first_reduce = reduce_inst; + } else { + CHECK(first_reduce->dimensions() == reduce_inst->dimensions()); + } + ShapeIndex output_shape_index; + if (reduce_or_tuple->opcode() == HloOpcode::kTuple) { + output_shape_index = {i}; + } + + EmitPrologueForOneReduction(unnested_hlo, reduce_inst, i, kernel_info, + &elemental_emitter, + std::move(output_shape_index)); + } + + 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(), + /*ArraySize=*/b_.getInt32(num_partial_results), + "current_output_linear_index_address")); + + if (!reduction_info->IsRowReduction()) { + llvm::Type* bool_ty = b_.getInt1Ty(); + llvm::AllocaInst* output_inbound_addr = Alloca(bool_ty); + Store(llvm::ConstantInt::get(bool_ty, 0), output_inbound_addr); + reduction_info->SetCurrentOutputInboundAddress(output_inbound_addr); + } +} + +void IrEmitterUnnested::EmitFullWarpShuffleDownLoopForAllReduces( + absl::Span reducers, + absl::Span partial_result_addresses) { + for (int distance = 16; distance >= 1; distance /= 2) { + for (int i = 0; i != reducers.size(); ++i) { + llvm::Type* element_type = + partial_result_addresses[i]->getType()->getElementType(); + int bit_width = llvm_ir::GetSizeInBits(element_type); + llvm::Value* result_from_other_lane = Alloca( + element_type, nullptr, "result_from_other_lane" + llvm::Twine(i)); + // Bitcast cannot be applied to aggregate types (even packed ones), so + // we bitcast addresses of load/store to intN* of the same bit-width. + llvm::Type* shuffled_value_type = + element_type->isStructTy() ? b_.getIntNTy(bit_width) : element_type; + auto convert_pointer_for_shuffle = [&](llvm::Value* ptr) { + return BitCast(ptr, shuffled_value_type->getPointerTo()); + }; + llvm::Value* partial_result = + Load(convert_pointer_for_shuffle(partial_result_addresses[i]), + "partial_reduction_result"); + Store(EmitFullWarpShuffleDown(partial_result, b_.getInt32(distance), &b_), + convert_pointer_for_shuffle(result_from_other_lane)); + TF_CHECK_OK(EmitCallToNestedComputation( + *reducers[i], {partial_result_addresses[i], result_from_other_lane}, + partial_result_addresses[i])); + } + } +} + +void IrEmitterUnnested::EmitEpilogueForReduction( + HloInstruction* unnested_hlo, KernelCodegenInfo* kernel_info) { + ReductionCodegenInfo* reduction_info = + static_cast(kernel_info); + int num_reduces = reduction_info->GetNumberOfReduces(); + absl::Span partial_result_addresses = + reduction_info->GetPartialResultAddresses(); + const InlinedVector& reducers = + reduction_info->GetReducers(); + absl::Span reduction_output_shape_indices = + reduction_info->GetReductionOutputShapeIndices(); + + if (reduction_info->IsRowReduction()) { + EmitFullWarpShuffleDownLoopForAllReduces(reducers, + partial_result_addresses); + llvm::Value* lane_id = reduction_info->GetLaneId(); + llvm_ir::LlvmIfData if_lane_id_is_zero_data = llvm_ir::EmitIfThenElse( + ICmpEQ(lane_id, llvm::ConstantInt::get(lane_id->getType(), 0)), + "lane_id_is_zero", &b_); + llvm_ir::SetToFirstInsertPoint(if_lane_id_is_zero_data.true_block, &b_); + } else { + llvm::Value* output_inbound_addr = + reduction_info->GetCurrentOutputInboundAddress(); + llvm::Value* output_inbound = Load(output_inbound_addr); + llvm_ir::LlvmIfData if_output_inbound_data = llvm_ir::EmitIfThenElse( + ICmpEQ(output_inbound, + llvm::ConstantInt::get(output_inbound->getType(), 1)), + "output_inbound", &b_); + 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) { + 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)}))); + } + } + } +} + +void IrEmitterUnnested::EmitTileElementForReduction( + HloInstruction* unnested_hlo, const llvm_ir::IrArray::Index& index, + const KernelCodegenInfo* kernel_info, llvm::Value* y_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() + : unnested_hlo; + llvm_ir::TiledParameterInfo* tiled_param_info = + kernel_info->GetTiledParameterInfo(); + tiled_param_info->set_y(y_loc); + tiled_param_info->set_x(x_loc); + + // 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()], + InBoundsGEP(reduction_info->GetCurrentOutputLinearIndexAddress(), + {b_.getInt32(partial_result_index)})); + if (!reduction_info->IsRowReduction()) { + llvm::Type* bool_ty = b_.getInt1Ty(); + llvm::AllocaInst* output_inbound_addr = + reduction_info->GetCurrentOutputInboundAddress(); + Store(llvm::ConstantInt::get(bool_ty, 1), output_inbound_addr); + } + + InlinedVector input_gens; + std::vector> + extra_output_gens; + GpuElementalIrEmitter elem_emitter(hlo_module_config_, module_, &b_, + GetNestedComputer()); + FusedIrEmitter fused_emitter(GetGeneratorForOperandIrArrays(unnested_hlo), + &elem_emitter); + absl::Span output_instructions = + GetOutputInstructions(&reduce_or_tuple); + // Construct the ElementGenerator for each reduction and extra output in the + // the group of output instructions. + if (unnested_hlo->opcode() == HloOpcode::kFusion) { + fused_emitter.SetTiledParameterInfo(tiled_param_info); + TF_CHECK_OK(unnested_hlo->fused_expression_root()->Accept(&fused_emitter)); + + for (int i = 0, e = output_instructions.size(); i != e; ++i) { + const HloInstruction* inst = output_instructions[i]; + ShapeIndex output_shape_index; + if (reduce_or_tuple->opcode() == HloOpcode::kTuple) { + output_shape_index = {i}; + } + if (inst->opcode() == HloOpcode::kReduce) { + input_gens.push_back(fused_emitter.GetGenerator(inst->operand(0))); + } else { + extra_output_gens.emplace_back(fused_emitter.GetGenerator(inst), + std::move(output_shape_index)); + } + } + } else { + input_gens.push_back([&](const IrArray::Index& index) { + return GetIrArray(*unnested_hlo->operand(0), *unnested_hlo) + .EmitReadArrayElement(index, &b_); + }); + } + + IrArray::Index input_index = + 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 = + reduction_info->GetReductionInputAddresses(); + const InlinedVector& reducers = + reduction_info->GetReducers(); + + // Emit code to generate the input and perform the reduction computation for + // each reduction instruction. + 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_result_address, reduction_input_addresses[i]}, + partial_result_address)); + } + + // Emit code to generate the output for the non-reduction instructions in the + // fusion, if any. + TF_CHECK_OK( + EmitExtraOutputsForReduce(unnested_hlo, input_index, extra_output_gens)); +} + +// Emits a kernel for the hlo instruction using the given tiling scheme. +void IrEmitterUnnested::EmitBlock(const TileGenerator& emit_one_tile, + KernelCodegenInfo* kernel_info, + KernelSupportLibrary* ksl, + llvm::Type* index_ty) { + KernelMappingScheme* mapping_scheme = kernel_info->GetKernelMappingScheme(); + absl::Span dims_in_tile = mapping_scheme->GetDimensionsInTiles(); + absl::Span dims_in_block = + mapping_scheme->GetDimensionsInBlocks(); + absl::Span block_sizes = mapping_scheme->GetBlockSizes(); auto index_typed_constant = [&](uint64 c) -> llvm::Constant* { return llvm::ConstantInt::get(index_ty, c); }; - // Adds `addend` to the given `dim` of `index`. - auto offset_dim = [&](IrArray::Index index, llvm::Value* addend, int64 dim) { - index[dim] = builder->CreateAdd(index[dim], addend); - return index; - }; - auto emit_full_tile = [&] { - for (int64 i = 0; i < tile_size; i += num_rows) { - auto source_idx = offset_dim(index, index_typed_constant(i), /*dim=*/1); - auto y_loc = builder->CreateAdd(index_typed_constant(i), y); - emit_elem_function(source_idx, y_loc); + // Emit all the tiles for a given dimension in a tile block. + auto emit_tiles_for_block_dim = + [&](const string& loop_name, const IrArray::Index& starting_tile, + int dim_id, + const std::function + emit_next_block_dim) { + if (block_sizes[dim_id] == 1) { + emit_next_block_dim(starting_tile); + } else { + llvm::Value* starting_tile_index_for_dim = starting_tile[dim_id]; + llvm::Value* block_size_for_dim = + index_typed_constant(block_sizes[dim_id]); + llvm::Value* block_id_for_dim = + b_.CreateUDiv(starting_tile_index_for_dim, block_size_for_dim); + llvm::Value* last_block_for_dim = + index_typed_constant(dims_in_block[dim_id] - 1); + llvm::Value* last_block_size_for_dim = index_typed_constant( + dims_in_tile[dim_id] - + (dims_in_block[dim_id] - 1) * block_sizes[dim_id]); + 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); + }); + } + }; + + absl::Span reduced_dims = + mapping_scheme->GetDimensionsInElements(); + const bool block_contains_multi_tiles = + mapping_scheme->GetNumberOfTilesInOneBlock() > 1; + + // Emit the tile with a given tile_index, by calculating the tight bounds for + // each dimension of the tile and then calling emit_one_tile. + auto emit_one_tile_for_tile_index = [&](const IrArray::Index& tile_index) { + std::vector output_tile_bounds(3); + for (int i = KernelMappingScheme::DimY; i < KernelMappingScheme::DimTot; + ++i) { + int64 tile_size_for_dim = mapping_scheme->GetTileSizeForDimension(i); + // Only last row or column may not have full size. + llvm::Value* is_last_row = + ICmpEQ(tile_index[i], index_typed_constant(dims_in_tile[i] - 1)); + int64 partial_row_size = + reduced_dims[i] - (dims_in_tile[i] - 1) * tile_size_for_dim; + output_tile_bounds[i] = + Select(is_last_row, index_typed_constant(partial_row_size), + index_typed_constant(tile_size_for_dim), "tile_bound"); } - }; - auto emit_last_row = [&] { - ksl->IfReturnVoid("x_in_tile", builder->CreateICmpULT(x, tile_width), [&] { - // tile_height_upper_bound = - // ceil(tile_height / num_rows) * num_rows - auto tile_height_upper_bound = builder->CreateMul( - builder->CreateUDiv( - builder->CreateAdd(tile_height, - index_typed_constant(num_rows - 1)), - index_typed_constant(num_rows)), - index_typed_constant(num_rows)); - ksl->ForReturnVoid( - loop_name, /*start=*/index_typed_constant(0), - /*end=*/tile_height_upper_bound, - /*step=*/index_typed_constant(num_rows), [&](llvm::Value* y_indvar) { - auto y_loc = builder->CreateAdd(y_indvar, y); - ksl->IfReturnVoid( - "y_in_tile", builder->CreateICmpULT(y_loc, tile_height), [&] { - emit_elem_function(offset_dim(index, y_indvar, /*dim=*/1), - y_loc); - }); - }); - }); + IrArray::Index tile_origin = + mapping_scheme->GetElementIndexForTileOrigin(tile_index); + emit_one_tile(tile_origin, output_tile_bounds, block_contains_multi_tiles); }; - ksl->IfReturnVoid( - "full_tile", - builder->CreateAnd( - builder->CreateICmpEQ(index_typed_constant(tile_size), tile_width), - builder->CreateICmpEQ(index_typed_constant(tile_size), tile_height)), - emit_full_tile, emit_last_row); + + const IrArray::Index starting_block = + mapping_scheme->EmitBlockIndex(index_ty); + const IrArray::Index starting_tile_for_dim_z = + mapping_scheme->GetTileIndexForBlockOrigin(starting_block); + + // Emit the three dimensional block of tiles. + emit_tiles_for_block_dim( + "block_dim_z", starting_tile_for_dim_z, KernelMappingScheme::DimZ, + [&](const IrArray::Index& starting_tile_for_dim_y) { + emit_tiles_for_block_dim( + "block_dim_y", starting_tile_for_dim_y, KernelMappingScheme::DimY, + [&](const IrArray::Index& starting_tile_for_dim_x) { + emit_tiles_for_block_dim("block_dim_x", starting_tile_for_dim_x, + KernelMappingScheme::DimX, + emit_one_tile_for_tile_index); + }); + }); } -} // namespace -// 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 -// which have a shape that is a 0-2-1 transpose of the output tensors. -// -// 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 of -// three components 0-1-2 in the order major to minor. The x- and y- dimensions -// of the tensors are tiled in square tiles of edge length `kTileSize`. Each -// thread block of `kTileSize` x `kNumRows` threads transposes one tile: each -// thread copies kTileSize/kNumRows elements from the input to a shared memory -// tile, then the otherwise "regular hlo kernel" reads from the shared memory -// instead of the original input. -// -// This is similar to the following CUDA algorithm in TensorFlow: -// https://goo.gl/MStRV6. +// Emits a kernel for the hlo instruction using the given kernel mapping scheme. // -// `kTileSize` should usually be same as warp size. We currently choose 32 for -// `kTileSize` and 4 for `kNumRows`. The CUDA algorithm uses 8 for `kNumRows`. -// -// TODO(b/33320379): Here each block transposes 1 tile. It may be more efficient -// to launch fewer blocks so each transposes many tiles. -LaunchDimensions IrEmitterUnnested::EmitHlo021Tile( - HloInstruction* hlo, absl::Span reduced_output_dims, - absl::Span tiled_param_ids) { - // Parameters for the tiling algorithm. - constexpr int64 kTileSize = 32; - constexpr int64 kNumRows = 4; - constexpr int64 kThreadsPerTile = kTileSize * kNumRows; - - // Construct IrArrays for the inputs and outputs. - std::vector output_arrays = ConstructIrArrayForOutputs(*hlo); - int64 num_outputs = output_arrays.size(); - std::vector param_arrays = ConstructIrArrayForInputs(*hlo); +// 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 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. +// kernel_info: Represent other information to support the code generation +// of the tiled kernel for the hlo. +LaunchDimensions IrEmitterUnnested::EmitKernel( + HloInstruction* unnested_hlo, absl::Span tiled_param_ids, + const KernelCodeGenerator& kernel_generator, + KernelCodegenInfo* kernel_info) { + KernelMappingScheme* mapping_scheme = kernel_info->GetKernelMappingScheme(); + + std::vector param_arrays = ConstructIrArrayForInputs(*unnested_hlo); int64 num_params = param_arrays.size(); - // Allocate shared memory buffers to store the tiled inputs. std::vector param_shmem_buffers(num_params, nullptr); for (int64 id : tiled_param_ids) { - const HloInstruction* param = hlo->operand(id); - // Add 1 to the minor dimension to reduce shared memory bank conflicts. - llvm::Type* tile_type = llvm::ArrayType::get( - llvm::ArrayType::get(llvm_ir::PrimitiveTypeToIrType( - param->shape().element_type(), module_), - kTileSize + 1), - kTileSize); - const int kNVPTXSharedMemoryAddrSpace = 3; - auto* tile_base_ptr = new llvm::GlobalVariable( - *b_.GetInsertBlock()->getParent()->getParent(), tile_type, - /*isConstant=*/false, llvm::GlobalValue::PrivateLinkage, - llvm::UndefValue::get(tile_type), - llvm_ir::AsStringRef(IrName(hlo, StrCat("tile", id))), nullptr, - llvm::GlobalValue::NotThreadLocal, kNVPTXSharedMemoryAddrSpace); - param_shmem_buffers[id] = tile_base_ptr; + const HloInstruction* param = unnested_hlo->operand(id); + param_shmem_buffers[id] = + mapping_scheme->GetSharedMemoryBufferForElementType( + llvm_ir::PrimitiveTypeToIrType(param->shape().element_type(), + module_), + IrName(unnested_hlo, StrCat("tile", id))); VLOG(3) << "Added shmem buffer for parameter " << id << ": " - << llvm_ir::DumpToString(*tile_base_ptr); + << llvm_ir::DumpToString(*param_shmem_buffers[id]); } - // The 0-2-1 shape of the tiling scheme is the reduced shape of the HLO result - // for the purpose of tiling. Calculate the logical output dimensions in the - // tile from the reduced output dimensions. - std::vector output_dims_in_tiles = std::vector( - reduced_output_dims.begin(), reduced_output_dims.end()); - CHECK_EQ(output_dims_in_tiles.size(), 3); - for (int i = 1; i < 3; ++i) { - output_dims_in_tiles[i] = - CeilOfRatio(output_dims_in_tiles[i], kTileSize); - } - const int64 num_tiles = - absl::c_accumulate(output_dims_in_tiles, 1, std::multiplies()); - LaunchDimensions launch_dimensions(num_tiles, kThreadsPerTile); + const ReductionCodegenInfo* reduction_info = + dynamic_cast(kernel_info); + bool is_column_reduction = + (reduction_info && !reduction_info->IsRowReduction()); + + LaunchDimensions launch_dimensions = + LaunchDimensions(mapping_scheme->GetNumberOfBlocks(), + mapping_scheme->GetThreadsPerBlock()); + // TODO(b/110211620): Enable int32 index type for column reduction. llvm::Type* index_ty = - GetIndexTypeForKernel(hlo, launch_dimensions.launch_bound(), &b_); + is_column_reduction + ? b_.getInt64Ty() + : GetIndexTypeForKernel(unnested_hlo, + launch_dimensions.launch_bound(), &b_); + auto index_typed_constant = [&](uint64 c) -> llvm::Constant* { return llvm::ConstantInt::get(index_ty, c); }; - // Cast each output IrArray to its corresponding reduced shape and keep the - // reduced shape live during IR emission. - std::vector output_in_reduced_shape_arrays; - std::vector output_reduced_shapes; - CHECK_EQ(ConstructOutputReducedShapeAndCastOutputIrArrayToShape( - *hlo, output_arrays, reduced_output_dims, &output_reduced_shapes, - &output_in_reduced_shape_arrays), - num_outputs); + // For multioutput fusion, one thread needs to output a tuple with pointers to + // all the individual outputs. We could do this at any point in the kernel, + // but we do it at the beginning in the hopes of reducing register pressure, + // 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_); + }); + } // For each tiled parameter, cast its input IrArray to the corresponding // reduced shape and keep the reduced shape live during IR emission. std::vector param_in_reduced_shape_arrays; std::vector param_reduced_shapes; - CHECK_EQ(ConstructInputReducedShapeAndCastInputIrArrayToShape( - *hlo, param_arrays, param_shmem_buffers, reduced_output_dims, - ¶m_reduced_shapes, ¶m_in_reduced_shape_arrays), - num_params); + absl::Span reduced_dims = + mapping_scheme->GetDimensionsInElements(); + int num_shapes = ConstructInputReducedShapeAndCastInputIrArrayToShape( + *unnested_hlo, param_arrays, param_shmem_buffers, reduced_dims, + ¶m_reduced_shapes, ¶m_in_reduced_shape_arrays); + DCHECK_EQ(num_shapes, num_params); // Calculate the starting element coordinate within a tile for the current // thread, (y, x) from thread_id. llvm::Value* x; llvm::Value* y; - std::tie(y, x) = CalculateYXCoordinateWithinTile( - &b_, index_typed_constant(kTileSize), kThreadsPerTile); - - // Calculate the index for the current output tile from block_id. - const IrArray::Index output_tile_index( - GetBlockIdx(&b_, index_ty, num_tiles), - ShapeUtil::MakeShapeWithDescendingLayout(PRED /*arbitrary*/, - output_dims_in_tiles), - &b_); - - // Output tile origin is the index for the first element of the current output - // tile. - const IrArray::Index output_tile_origin = [&] { - IrArray::Index index = output_tile_index; - for (int i = 1; i < 3; ++i) { - index[i] = Mul(output_tile_index[i], index_typed_constant(kTileSize), - "tile_origin." + std::to_string(i)); - } - return index; - }(); - - // Calculate the input tile origin from the output tile origin. - const IrArray::Index input_tile_origin( - Permute({0, 2, 1}, output_tile_origin.multidim())); + std::tie(y, x) = mapping_scheme->EmitThreadYXCoordinate(index_ty); - // Calculate the current output tile bounds in each of the logical dimensions. - std::vector output_tile_bounds(3); - for (int i = 1; i < 3; ++i) { - // Only last row or column may not have full size. - output_tile_bounds[i] = - Select(ICmpEQ(output_tile_index[i], - index_typed_constant(output_dims_in_tiles[i] - 1)), - index_typed_constant(reduced_output_dims[i] - - (output_dims_in_tiles[i] - 1) * kTileSize), - index_typed_constant(kTileSize), "kTileSize"); - } + kernel_info->SetLaneId( + mapping_scheme->GetNumberOfThreadsForDimensionX() == kWarpSize ? x + : nullptr); + kernel_info->SetIndexType(index_ty); KernelSupportLibrary ksl(&b_, llvm_ir::UnrollMode::kDefaultUnroll); - // Curry a few parameters to EmitTiledElementalCodeWithBoundsCheck. auto emit_tiled_elemental_code_with_bounds_check = [&](const IrArray::Index& index, const string& loop_name, - llvm::Value* tile_width, llvm::Value* tile_height, - const std::function& - emit_elem_function) { - EmitTiledElementalCodeWithBoundsCheck( - kTileSize, kNumRows, index, loop_name, &ksl, &b_, y, x, tile_width, - tile_height, emit_elem_function); + llvm::Value* tile_height, llvm::Value* tile_width, + const EmitElementFunction& emit_elem_function) { + EmitTiledElementalCodeWithBoundsCheck(mapping_scheme, index, loop_name, + &ksl, &b_, y, x, tile_height, + tile_width, emit_elem_function); }; - // Adds `addend` to the given `dim` of `index`. - auto offset_dim = [&](IrArray::Index index, llvm::Value* addend, int64 dim) { - index[dim] = Add(index[dim], addend); - return index; + auto emit_one_tile = [&](const IrArray::Index& output_tile_origin, + absl::Span output_tile_bounds, + bool block_contains_multi_tiles) { + // Calculate the input tile origin from the output tile origin. + const IrArray::Index input_tile_origin( + Permute({0, 2, 1}, output_tile_origin.multidim())); + + // 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. + if (!tiled_param_ids.empty()) { + // Copy input parameter values to shared memory buffers: + // tile[y, x] = input[index] + // 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_tile_origin, "input", output_tile_bounds[2], + output_tile_bounds[1], + [&](const IrArray::Index& index, llvm::Value* y_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]; + llvm::Value* shmem_buffer = param_shmem_buffers[id]; + // TODO(jlebar): Add AA metadata to this store. Tile buffers are + // global variables, so LLVM can't infer much about it. + Store(input_in_logical_shape.EmitReadArrayElement( + index, &b_, "input_element"), + GEP(shmem_buffer, {index_typed_constant(0), y_loc, x_loc})); + } + }); + + // Wait for all threads to reach this point using `__syncthreads` in CUDA. + llvm_ir::EmitCallToIntrinsic(llvm::Intrinsic::nvvm_barrier0, {}, {}, &b_); + } + + llvm_ir::TiledParameterInfo tiled_param_info(param_shmem_buffers, y, x); + kernel_info->SetTiledParamInfo(&tiled_param_info); + + // 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_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 + // used, we need to wait for all threads to finish using the shared memory + // buffer for the current tile before we move on to process the next tile + // and overwrite the shared memory buffers. + if (block_contains_multi_tiles && !tiled_param_ids.empty()) { + llvm_ir::EmitCallToIntrinsic(llvm::Intrinsic::nvvm_barrier0, {}, {}, &b_); + } }; - const IrArray::Index input_index = - offset_dim(offset_dim(input_tile_origin, x, /*dim=*/2), y, /*dim=*/1); - - // Copy input parameter values to shared memory buffers: - // tile[y, x] = input[index] - emit_tiled_elemental_code_with_bounds_check( - input_index, "input", output_tile_bounds[1], output_tile_bounds[2], - [&](const IrArray::Index& index, llvm::Value* y_loc) { - for (int64 id : tiled_param_ids) { - IrArray& input_in_logical_shape = param_in_reduced_shape_arrays[id]; - llvm::Value* shmem_buffer = param_shmem_buffers[id]; - // TODO(jlebar): Add AA metadata to this store. Tile buffers are - // global variables, so LLVM can't infer much about it. - Store(input_in_logical_shape.EmitReadArrayElement(index, &b_, - "input_element"), - GEP(shmem_buffer, {index_typed_constant(0), y_loc, x})); - } - }); - // 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. - llvm_ir::EmitCallToIntrinsic(llvm::Intrinsic::nvvm_barrier0, {}, {}, &b_); + const BlockPrologueGenerator& block_prologue_generator = + kernel_generator.GetBlockPrologueGenerator(); + if (block_prologue_generator) { + block_prologue_generator(unnested_hlo, kernel_info); + } + + EmitBlock(std::move(emit_one_tile), kernel_info, &ksl, index_ty); - llvm_ir::TiledParameterInfo tiled_param_info(param_shmem_buffers, y, x); + const BlockEpilogueGenerator& block_epilogue_generator = + kernel_generator.GetBlockEpilogueGenerator(); + if (block_epilogue_generator) { + block_epilogue_generator(unnested_hlo, kernel_info); + } - const IrArray::Index output_index = - offset_dim(offset_dim(output_tile_origin, x, /*dim=*/2), y, /*dim=*/1); + return launch_dimensions; +} - // Write to output[index] by emitting code like normal, except that values for - // the tiled parameters are read from the shmem buffers. +// 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. 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 +// of three components 0-1-2 in the order major to minor. The x- and y- +// dimensions of the tensors are tiled in square tiles with an edge length +// `kTileSize`. Each thread block of `kTileSize` x `kNumRows` threads +// transposes one tile: each thread copies kTileSize/kNumRows elements from +// the input to a shared memory tile, then the otherwise "regular HLO kernel" +// reads from the shared memory instead of the original input. +// +// This is similar to the following CUDA algorithm in TensorFlow: +// https://goo.gl/MStRV6. +// +// `kTileSize` should usually be same as warp size. We currently choose 32 for +// `kTileSize` and 4 for `kNumRows`. The CUDA algorithm uses 8 for `kNumRows`. +// +// TODO(b/33320379): Here each block transposes 1 tile. It may be more +// efficient to launch fewer blocks so each transposes many tiles. +LaunchDimensions IrEmitterUnnested::EmitHlo021Tile( + HloInstruction* hlo, absl::Span reduced_output_dims, + absl::Span tiled_param_ids) { + constexpr int kNumRows = 4; + KernelMappingScheme mapping_scheme( + reduced_output_dims, /*tile_size_y=*/kWarpSize, + /*tile_size_x=*/kWarpSize, /*req_block_sizes=*/{1, 1, 1}, + /*num_threads_y=*/kNumRows, + /*num_threads_x=*/kWarpSize, &b_); + TileElementGenerator element_generator; if (hlo->opcode() == HloOpcode::kCopy) { - emit_tiled_elemental_code_with_bounds_check( - output_index, "output", output_tile_bounds[2], output_tile_bounds[1], - [&](const IrArray::Index& index, llvm::Value* y_loc) { - // TODO(jlebar): Add AA metadata to this load. - llvm::Instruction* load_from_shmem_buffer = - Load(GEP(param_shmem_buffers[0], {b_.getInt64(0), x, y_loc}), - "output_element"); - output_in_reduced_shape_arrays[0].EmitWriteArrayElement( - index, load_from_shmem_buffer, &b_); - }); + 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) { + EmitTileElementForCopy(hlo, index, kernel_info, y_loc, x_loc, x_iter_num); + }; } else { - CHECK_EQ(hlo->opcode(), HloOpcode::kFusion); - emit_tiled_elemental_code_with_bounds_check( - output_index, "output", output_tile_bounds[2], output_tile_bounds[1], - [&](const IrArray::Index& index, llvm::Value* y_loc) { - GpuElementalIrEmitter elem_emitter(hlo_module_config_, module_, &b_, - GetNestedComputer()); - FusedIrEmitter fused_emitter(param_arrays, &elem_emitter); - tiled_param_info.set_y(y_loc); - fused_emitter.SetTiledParameterInfo(&tiled_param_info); - TF_CHECK_OK(hlo->fused_expression_root()->Accept(&fused_emitter)); - IrArray::Index untiled_index = llvm_ir::GetUnreducedOutputIndex( - index, output_reduced_shapes[0], output_arrays[0].GetShape(), - &b_); - const llvm_ir::ElementGenerator& output_generator = - fused_emitter.GetRootGenerator(); - llvm::Value* output_value = - output_generator(untiled_index).ValueOrDie(); - if (hlo->IsMultiOutputFusion()) { - CHECK(output_value->getType()->isStructTy()); - CHECK_EQ(output_value->getType()->getStructNumElements(), - output_in_reduced_shape_arrays.size()); - for (int64 i = 0; i < output_in_reduced_shape_arrays.size(); ++i) { - output_in_reduced_shape_arrays[i].EmitWriteArrayElement( - index, ExtractValue(output_value, i), &b_); - } - } else { - output_in_reduced_shape_arrays[0].EmitWriteArrayElement( - index, output_value, &b_); - } - }); + 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, 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)); + return EmitKernel(hlo, tiled_param_ids, kernel_generator, &kernel_info); +} - // For multioutput fusion, emit a tuple with all the individual outputs. - if (hlo->IsMultiOutputFusion()) { - llvm_ir::EmitTuple(GetIrArray(*hlo, *hlo), output_arrays, &b_, module_); +namespace { +// 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. +// +// 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); + }); } - return launch_dimensions; + 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; + } } +// 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) { HloOpcode opcode = hlo->opcode(); CHECK(opcode == HloOpcode::kFusion || opcode == HloOpcode::kCopy); @@ -3451,8 +3248,8 @@ bool IrEmitterUnnested::CheckAndEmitHloWithTile021(HloInstruction* hlo) { ? ShapeUtil::GetSubshape(hlo->shape(), {0}) : hlo->shape(); - // If the output_shape is reduced to 021 shape, find all the parameters of the - // hlo that are in the corresponding 012 shape. + // If the output_shape is reduced to 021 shape, find all the parameters of + // the HLO that are in the corresponding 012 shape. std::vector params_012; optional> reduced_dims_021; for (int64 operand_idx = 0; operand_idx < hlo->operand_count(); @@ -3484,10 +3281,17 @@ bool IrEmitterUnnested::CheckAndEmitHloWithTile021(HloInstruction* 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 - // elements are of size 4 bytes), and CUDA has an architectural limit of 48kb - // shared memory per SM. (This is increased to 96kb in Volta, but we don't - // use this, in part because it eats into our L1 cache space.) + // elements are of size 4 bytes), and CUDA has an architectural limit of + // 48kb shared memory per SM. (This is increased to 96kb in Volta, but we + // don't use this, in part because it eats into our L1 cache space.) // // For correctness we need to ensure that we don't make more than 48kb worth // of shmem tiles per block. And for performance, we'd probably like to use @@ -3495,9 +3299,9 @@ bool IrEmitterUnnested::CheckAndEmitHloWithTile021(HloInstruction* hlo) { // gpu core. // // We say without benchmarks that we want at least 3 threads/block, - // corresponding to 3 shmem tiles if the elements are 32 bits wide. We choose - // which params get the shmem transpose treatment arbitrarily; it's not clear - // if there's a Right Choice. + // corresponding to 3 shmem tiles if the elements are 32 bits wide. We + // choose which params get the shmem transpose treatment arbitrarily; it's + // not clear if there's a Right Choice. // // This is only sound if tiled transposes are the only place where we use // shared memory in fusions. If in the future other fusible ops use shared @@ -3519,16 +3323,361 @@ bool IrEmitterUnnested::CheckAndEmitHloWithTile021(HloInstruction* hlo) { } VLOG(3) << "EmitHlo021Tile Emitting hlo tile 0-2-1" << hlo->ToString(); - thunk_sequence_->emplace_back( - BuildKernelThunk(hlo, /*implements_whole_instruction=*/true)); + std::unique_ptr kernel_thunk = + BuildKernelThunk(hlo, /*implements_whole_instruction=*/true); const LaunchDimensions launch_dimensions = EmitHlo021Tile(hlo, *reduced_dims_021, params_012); - UpdateLaunchDimensions(launch_dimensions, LastThunk(), + UpdateLaunchDimensions(launch_dimensions, kernel_thunk.get(), ir_emitter_context_->llvm_module()); + AddThunkToThunkSequence(std::move(kernel_thunk)); return true; } +namespace { +// Checks that the outputs of a fusion with reduction are consistent. +Status AreFusedReductionOutputsConsistent( + absl::Span output_instructions, + const HloInstruction* first_reduce) { + for (const HloInstruction* inst : output_instructions) { + if (inst->opcode() == HloOpcode::kReduce) { + // Shapes, layouts and dimensions must be the same for all reduces + // inside of this fusion. + TF_RET_CHECK(ShapeUtil::Equal(first_reduce->shape(), inst->shape())); + TF_RET_CHECK(ShapeUtil::Equal(first_reduce->operand(0)->shape(), + inst->operand(0)->shape())); + TF_RET_CHECK(ShapeUtil::Equal(first_reduce->operand(1)->shape(), + inst->operand(1)->shape())); + TF_RET_CHECK(first_reduce->dimensions() == inst->dimensions()); + } else { + // For extra outputs we can relax shape equality to allow different + // types (with the same number of elements). Layouts still have to + // match. + TF_RET_CHECK(ShapeUtil::CompatibleIgnoringElementType( + first_reduce->operand(0)->shape(), inst->shape())); + TF_RET_CHECK(LayoutUtil::Equal(first_reduce->operand(0)->shape().layout(), + inst->shape().layout())); + } + } + return Status::OK(); +} + +// Finds the dimensions to keep for the reduction, sorts and returns the +// dimensions from minor to major. +DimensionVector GetDimensionsToKeepMinorToMajor( + const Shape& input_shape, absl::Span dims_to_reduce) { + DimensionVector input_dims(input_shape.rank(), 0); + absl::c_iota(input_dims, 0); + DimensionVector input_dims_to_keep; + for (int input_dim : input_dims) { + auto it = absl::c_find_if(dims_to_reduce, [&](int64 dim_to_reduce) { + return dim_to_reduce == input_dim; + }); + if (it == dims_to_reduce.end()) { + input_dims_to_keep.push_back(input_dim); + } + } + + // Sort the dimensions to keep from minor to major. + absl::c_sort(input_dims_to_keep, [&input_shape](int64 dim_a, int64 dim_b) { + return PositionInContainer(LayoutUtil::MinorToMajor(input_shape), dim_a) < + PositionInContainer(LayoutUtil::MinorToMajor(input_shape), dim_b); + }); + + VLOG(10) << "dims to keep minor to major" + << absl::StrJoin(input_dims_to_keep, ","); + return input_dims_to_keep; +} + +// Given the input shape and dimensions to reduce for the reduction to vector, +// returns : +// num_kept: the number of elements in the contiguous dimensions to keep. +// num_reduced_major: the number of elements in the dimensions to reduce that +// are more major than the dimensions to keep. +// num_reduced_minor: the number of elements in the dimensions to reduce that +// are more minor than the dimensions to kept. +std::tuple GetReductionToVectorDimensions( + const Shape& input_shape, absl::Span dims_to_reduce) { + DimensionVector input_dims_to_keep_minor_to_major = + GetDimensionsToKeepMinorToMajor(input_shape, dims_to_reduce); + CHECK(LayoutUtil::AreDimensionsConsecutive( + input_shape.layout(), input_dims_to_keep_minor_to_major)); + int num_reduced_major = 1, num_kept = 1, num_reduced_minor = 1; + if (input_dims_to_keep_minor_to_major.empty()) { + return std::make_tuple(num_reduced_major, num_kept, num_reduced_minor); + } + DimensionVector input_dims(input_shape.rank(), 0); + absl::c_iota(input_dims, 0); + absl::Span minor_to_major = + LayoutUtil::MinorToMajor(input_shape); + for (int input_dim : input_dims) { + int64 curr_dim_size = input_shape.dimensions(input_dim); + if (PositionInContainer(minor_to_major, input_dim) > + PositionInContainer(minor_to_major, + input_dims_to_keep_minor_to_major.back())) { + num_reduced_major *= curr_dim_size; + } else if (PositionInContainer(minor_to_major, input_dim) < + PositionInContainer(minor_to_major, + input_dims_to_keep_minor_to_major.front())) { + num_reduced_minor *= curr_dim_size; + } else { + num_kept *= curr_dim_size; + } + } + + 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* unnested_hlo, const HloInstruction* first_reduce) { + int64 depth = 1; + int64 height = 1; + int64 width = 1; + bool is_row_reduction = true; + int64 tile_size_x = 1; + int64 tile_size_y = 1; + int64 block_size_z = 1; + int64 num_threads_x = 1; + int64 num_threads_y = 1; + const Shape& input_shape = first_reduce->operand(0)->shape(); + int64 num_input_elems = ShapeUtil::ElementsIn(input_shape); + int64 num_output_elems = ShapeUtil::ElementsIn(first_reduce->shape()); + int64 num_reduced_major, num_kept, num_reduced_minor; + 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. + width = num_input_elems; + tile_size_x = kWarpSize * 16; + num_threads_x = kWarpSize; + } else if (num_reduced_minor == 1) { + // Column reduction reduces inputs with dimension [height, width], where + // width is the minor dimension, to dimension [width]. + height = num_reduced_major; + width = num_kept; + 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 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()); + 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 { + // Row reduction reduces inputs with dimension [depth, height, width], + // where width is the most minor dimension, to dimension [height] . + depth = num_reduced_major; + height = num_kept; + width = num_reduced_minor; + num_threads_x = kWarpSize; + if (width % (kWarpSize * 64) == 0) { + tile_size_x = kWarpSize * 64; + } else { + tile_size_x = kWarpSize * 8; + block_size_z = 8; + while (depth % block_size_z != 0) { + block_size_z -= 1; + } + } + } + DCHECK_EQ(depth * height * width, num_input_elems); + VLOG(10) << "is_row_reduction " << is_row_reduction << depth << " " << height + << " " << width; + + DimensionVector dims_in_elem{depth, height, width}; + DimensionVector req_block_sizes{block_size_z, 1, 1}; + 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); +} + +Status IrEmitterUnnested::EmitReductionToVector(HloInstruction* unnested_hlo) { + VLOG(10) << "Emitting reduction to vector " << unnested_hlo->ToString(); + + HloInstruction* reduce_or_tuple = unnested_hlo->opcode() == HloOpcode::kFusion + ? unnested_hlo->fused_expression_root() + : unnested_hlo; + absl::Span output_instructions = + GetOutputInstructions(&reduce_or_tuple); + const HloInstruction* first_reduce = + GetFirstReduceInstruction(output_instructions); + + if (output_instructions.size() > 1) { + TF_RETURN_IF_ERROR( + AreFusedReductionOutputsConsistent(output_instructions, first_reduce)); + } + + // Build an initializer thunk to initialize each reduction output. + std::vector> thunks; + for (int i = 0, e = output_instructions.size(); i != e; ++i) { + if (output_instructions[i]->opcode() != HloOpcode::kReduce) { + continue; + } + TF_ASSIGN_OR_RETURN( + std::unique_ptr initializer_thunk, + BuildInitializerThunk(unnested_hlo, + (output_instructions[i] == reduce_or_tuple) + ? ShapeIndex() + : ShapeIndex({i}))); + thunks.push_back(std::move(initializer_thunk)); + } + + // Build a kernel thunk to compute all the outputs. + std::unique_ptr kernel_thunk = + BuildKernelThunk(unnested_hlo, /*implements_whole_instruction=*/false); + + const Shape& input_shape = first_reduce->operand(0)->shape(); + // The layout of a reduction input is either set by LayoutAssignment for + // unnested kReduce or by InstructionFusion for fused kReduce. + CHECK(input_shape.has_layout()) << "LayoutAssignment or InstructionFusion " + "doesn't set the input layout of " + << first_reduce->ToString(); + + bool is_row_reduction; + llvm_ir::KernelMappingScheme mapping_scheme; + std::tie(mapping_scheme, is_row_reduction) = + 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, int64 x_iter_num) { + EmitTileElementForReduction(hlo, index, kernel_info, y_loc, x_loc, + x_iter_num); + }, + /*block_prologue_generator=*/ + [&](HloInstruction* hlo, KernelCodegenInfo* kernel_info) { + EmitPrologueForReduction(hlo, kernel_info); + }, + /*block_epilogue_generator*/ + [&](HloInstruction* hlo, KernelCodegenInfo* kernel_info) { + EmitEpilogueForReduction(hlo, kernel_info); + }); + + LaunchDimensions launch_dimensions = + EmitKernel(unnested_hlo, {}, kernel_generator, &reduction_info); + UpdateLaunchDimensions(launch_dimensions, kernel_thunk.get(), + ir_emitter_context_->llvm_module()); + + thunks.push_back(std::move(kernel_thunk)); + std::unique_ptr sequential_thunk = + absl::make_unique(std::move(thunks), unnested_hlo); + AddThunkToThunkSequence(std::move(sequential_thunk)); + + return Status::OK(); +} + Status IrEmitterUnnested::EmitConstantGlobals() { for (const BufferAllocation& allocation : ir_emitter_context_->buffer_assignment().Allocations()) { @@ -3550,10 +3699,10 @@ Status IrEmitterUnnested::EmitConstantGlobals() { } // These globals will be looked up by name by GpuExecutable so we need to - // give them an external linkage. Not all of their uses are visible in the - // LLVM IR (e.g. TupleThunk) so we can't give then a linkage that merely - // preserves their names (like available_externally), we also need to ensure - // that they stick around even if they're "unused". + // give them an external linkage. Not all of their uses are visible in + // the LLVM IR (e.g. TupleThunk) so we can't give then a linkage that + // merely preserves their names (like available_externally), we also need + // to ensure that they stick around even if they're "unused". // // We may have to be more more clever here in the future if we notice that // we're keeping around too many globals because of their linkage. diff --git a/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.h b/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.h index 93f11c069a4cebdf3c79cba17c824eded4f4b1db..21b842bb2cd63ac454f85556df20ae5877cecbe1 100644 --- a/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.h +++ b/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.h @@ -16,8 +16,11 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_GPU_IR_EMITTER_UNNESTED_H_ #define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_IR_EMITTER_UNNESTED_H_ +#include "absl/container/inlined_vector.h" #include "tensorflow/compiler/xla/service/gpu/ir_emitter.h" +#include "tensorflow/compiler/xla/service/gpu/sequential_thunk.h" #include "tensorflow/compiler/xla/service/gpu/thunk.h" +#include "tensorflow/compiler/xla/service/llvm_ir/kernel_support_library.h" #include "tensorflow/compiler/xla/service/llvm_ir/kernel_tiling.h" namespace xla { @@ -46,6 +49,100 @@ namespace gpu { // class IrEmitterUnnested : public IrEmitter { public: + // Parameter block_contains_multi_tiles indicates whether a tile block + // consists of multiple tiles or not. If the tile block contains only one + // tile, there is no need to use atomic operation to accumulate a local result + // to a global result to implement reduction. + using TileGenerator = + std::function output_tile_bounds, + bool block_contains_multi_tiles)>; + // KernelCodegenInfo records the common information to support the code + // generation for a kernel to process tensor elements by blocks. A block of + // tensor elements may contain one or multiple tiles. The code generators that + // generate code for tile elements or block prologue/epilogue refer to this + // class in their prototypes. If the implementations of such code generators + // require other information that are specific to the HLO instructions, the + // implementations need to define and use derived classes of this class. + class KernelCodegenInfo { + public: + explicit KernelCodegenInfo(llvm_ir::KernelMappingScheme* mapping_scheme) + : mapping_scheme_(mapping_scheme), + tiled_param_info_(nullptr), + lane_id_(nullptr), + index_ty_(nullptr) {} + virtual ~KernelCodegenInfo() {} + + void SetLaneId(llvm::Value* v) { lane_id_ = v; } + void SetIndexType(llvm::Type* t) { index_ty_ = t; } + void SetTiledParamInfo(llvm_ir::TiledParameterInfo* tiled_param_info) { + tiled_param_info_ = tiled_param_info; + } + + llvm::Value* GetLaneId() const { return lane_id_; } + llvm_ir::KernelMappingScheme* GetKernelMappingScheme() const { + return mapping_scheme_; + } + llvm_ir::TiledParameterInfo* GetTiledParameterInfo() const { + return tiled_param_info_; + } + llvm::Type* GetIndexType() const { return index_ty_; } + + protected: + llvm_ir::KernelMappingScheme* mapping_scheme_; + llvm_ir::TiledParameterInfo* tiled_param_info_; + llvm::Value* lane_id_; + llvm::Type* index_ty_; + }; + + // A function object to prepare for the code generation for a tile block. + using BlockPrologueGenerator = + std::function; + // A function object to finalize the code generation for a tile block. + using BlockEpilogueGenerator = + std::function; + // A function object to generate code to process one element in a tile. + // + // hlo: the instruction for which the code is generated for. + // index: the index for the first output element of the current thread. + // 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; + + // KernelCodeGenerator records the code generator objects that generate code + // for tile elements or tile block prologue/epilogue. + class KernelCodeGenerator { + public: + explicit KernelCodeGenerator( + TileElementGenerator tile_element_generator, + BlockPrologueGenerator block_prologue_generator = {}, + BlockEpilogueGenerator block_epilogue_generator = {}) + : tile_element_generator_(std::move(tile_element_generator)), + block_prologue_generator_(std::move(block_prologue_generator)), + block_epilogue_generator_(std::move(block_epilogue_generator)) {} + + const TileElementGenerator& GetTileElementGenerator() const { + return tile_element_generator_; + } + const BlockPrologueGenerator& GetBlockPrologueGenerator() const { + return block_prologue_generator_; + } + const BlockEpilogueGenerator& GetBlockEpilogueGenerator() const { + return block_epilogue_generator_; + } + + private: + TileElementGenerator tile_element_generator_; + BlockPrologueGenerator block_prologue_generator_; + BlockEpilogueGenerator block_epilogue_generator_; + }; + IrEmitterUnnested(const HloModuleConfig& hlo_module_config, const HloComputation* hlo_computation, IrEmitterContext* ir_emitter_context); @@ -80,8 +177,8 @@ class IrEmitterUnnested : public IrEmitter { Status HandleSelect(HloInstruction* select) override; Status HandleSort(HloInstruction* sort) override; Status HandleTupleSelect(HloInstruction* tuple_select) override; - Status HandleCrossReplicaSum(HloInstruction* crs) override; - Status HandleAfterAll(HloInstruction* gen_token) override; + Status HandleAllReduce(HloInstruction* crs) override; + Status HandleAfterAll(HloInstruction* after_all) override; Status EmitTargetElementLoop( const HloInstruction& hlo, @@ -97,10 +194,10 @@ class IrEmitterUnnested : public IrEmitter { Status EmitConstantGlobals(); private: - // Builds the appropriate thunk for the instruction hlo and returns the owning - // pointer to it. The caller needs to make sure `inst` outlives the lifetime - // of the returned Thunk object. - std::unique_ptr BuildThunk(const HloInstruction* hlo); + // Add a owning Thunk object to the thunk sequence. + void AddThunkToThunkSequence(std::unique_ptr thunk) { + thunk_sequence_->emplace_back(std::move(thunk)); + } // Builds the prototype of the IR kernel for `inst` and adds it to the module. // This kernel takes as arguments pointers to the given buffer allocations. @@ -110,80 +207,23 @@ class IrEmitterUnnested : public IrEmitter { // Helper for writing extra outputs from inside a reduce kernel. Status EmitExtraOutputsForReduce( - const HloInstruction* reduce, const llvm_ir::IrArray::Index& index, - absl::Span> - extra_output_gens); - - // EmitColumnReduction and EmitRowReduction emit code for column and row - // reduction of a matrix and/or 3D tensor. Row and column reduction have - // different memory access pattern, so for performance their implementations - // are significantly different. - // - // Emits code that reduces a matrix of shape [height x width] to a vector of - // [width]. Other parameters have the same meaning as those of - // `EmitReductionToVector`. Note that input shape might not be - // [height x width], but can be bitcast to [height x width] with "height" - // being the major dimension. - Status EmitColumnReduction( - int64 height, int64 width, HloInstruction* reduce, - const Shape& input_shape, - absl::Span input_gens, - absl::Span init_value_gens, - absl::Span reducers, - absl::Span reduce_output_shapes, - absl::Span> - extra_output_gens); - - // Emits code that reduces a 3D tensor of shape [depth x height x width] to a - // vector of shape [height]. Other parameters have the same meaning as those - // of `EmitReductionToVector`. Note that input shape might not be - // [depth x height x width], but can be bitcast to [depth x height x width] - // with "depth" being the most major dimension. - Status EmitRowReduction( - int64 depth, int64 height, int64 width, HloInstruction* reduce, - const Shape& input_shape, - absl::Span input_gens, - absl::Span init_value_gens, - absl::Span reducers, - absl::Span reduce_output_shapes, + const HloInstruction* unnested_hlo, const llvm_ir::IrArray::Index& index, absl::Span> extra_output_gens); - // Emits code that reduces a tensor of arbitrary rank to a scalar. - Status EmitReductionToScalar( - HloInstruction* reduce, const Shape& input_shape, - absl::Span input_gens, - absl::Span init_value_gens, - absl::Span reducers, - absl::Span reduce_output_shapes, - absl::Span> - extra_output_gens); - - // Figures out whether `reduce` is a row or column reduction, and which - // dimensions to reduce, and calls either `EmitRowReduction` or - // `EmitColumnReduction` as appropriate. `input_shape` is the shape of the - // input array, which is the operand of the Reduce instruction if unfused or - // of the Fusion instruction if fused. `input_gen` and `init_value_gen` - // generate elements of the input and the initial value. Other parameters mean - // the same as for `HandleReduce`. - // - // Multiple reduces can be emitted in the same loop, assuming they have the - // same input and output shapes, and the same reduce dimensions. - // - // extra_output_gens can contain extra generators for intermediate outputs. - // These must have the same shape as the reduce input as they are computed - // when the reduce inputs are being read. + // Generates code for reduction to contiguous dimensions. // - // Prerequisite: `IsReductionToVector(*reduce)` - Status EmitReductionToVector( - HloInstruction* reduce, const Shape& input_shape, - absl::Span input_gens, - absl::Span init_value_gens, - absl::Span dimensions_to_reduce, - absl::Span reducers, - absl::Span reduce_output_shapes, - absl::Span> - extra_output_gens); + // Prerequisite: `IsReductionToVector(*unnested_hlo)` + Status EmitReductionToVector(HloInstruction* unnested_hlo); + + // Computes the KernelMappingScheme for the reduce HLO and indicates whether + // 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* 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 @@ -202,22 +242,57 @@ class IrEmitterUnnested : public IrEmitter { LaunchDimensions EmitHlo021Tile(HloInstruction* hlo, absl::Span reduced_output_dims, absl::Span tiled_param_ids); + // Emits a kernel for an unnested HLO instruction. + LaunchDimensions EmitKernel(HloInstruction* unnested_hlo, + absl::Span param_ids, + const KernelCodeGenerator& kernel_generator, + KernelCodegenInfo* kernel_info); + void EmitBlock(const TileGenerator& emit_one_tile, + 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, + 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, + 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, + 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); + void EmitPrologueForOneReduction(HloInstruction* unnested_hlo, + HloInstruction* reduce_inst, int reduce_idx, + KernelCodegenInfo* kernel_info, + GpuElementalIrEmitter* elemental_emitter, + ShapeIndex output_shape_index); + // Wraps up the code generation for a tile block of a reduction kernel. + void EmitEpilogueForReduction(HloInstruction* unnested_hlo, + KernelCodegenInfo* kernel_info); + // For each reducer, emits the shuffle-down loop to accumulate the partial + // result to the global result. + void EmitFullWarpShuffleDownLoopForAllReduces( + absl::Span reducers, + absl::Span partial_result_addresses); // Generates the IrArray for each input of an hlo and returns a vector that // constains such IrArrays. std::vector ConstructIrArrayForInputs( const HloInstruction& hlo); - // For each output of the `hlo` instruction, constructs the reduced shape for - // the output with the given `reduced_output_dims` and cast the original - // output IrArray element in `output_arrays` to the reduced shape. Returns - // the number of outputs. - int ConstructOutputReducedShapeAndCastOutputIrArrayToShape( - const HloInstruction& hlo, - const std::vector& output_arrays, - absl::Span reduced_output_dims, - std::vector* output_reduced_shapes, - std::vector* output_in_reduced_shape_arrays); // For each input of the `hlo` instruction, checks its value in // `param_buffers` to find out whether the input has a reduced shape. If the // input has a reduced shape, constructs the reduced shape for the input and 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 8751e3a9c2a4c8da46d3ecd8437629450d4a2ba2..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}, @@ -125,8 +123,9 @@ static string GetSmName(std::pair compute_capability) { {{6, 2}, 62}, {{7, 0}, 70}, {{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; @@ -177,13 +176,6 @@ std::unique_ptr GetTargetMachine( } TargetOptions target_options = InitTargetOptionsFromCodeGenFlags(); - llvm_ir::SetTargetOptions( - /*fast_math_enabled=*/hlo_module_config.debug_options() - .xla_gpu_enable_fast_math(), - &target_options); - - // Enable FMA synthesis. - target_options.AllowFPOpFusion = FPOpFusion::Fast; // Set the verbose assembly options. target_options.MCOptions.AsmVerbose = false; @@ -206,8 +198,7 @@ std::unique_ptr GetTargetMachine( } return absl::WrapUnique(target->createTargetMachine( triple.str(), llvm_ir::AsStringRef(cpu_name), "+ptx60", target_options, - Optional(RelocModel), Optional(CMModel), - codegen_opt_level)); + getRelocModel(), getCodeModel(), codegen_opt_level)); } // Adds the standard LLVM optimization passes, based on the speed optimization @@ -401,8 +392,16 @@ StatusOr CompileModuleToPtx(llvm::Module* module, int32 opt_level = hlo_module_config.debug_options().xla_backend_optimization_level(); - CHECK_GE(opt_level, 2) - << "The XLA GPU backend doesn't support unoptimized code generation"; + if (opt_level < 2) { + LOG(ERROR) << std::string(80, '*'); + LOG(ERROR) << "The XLA GPU backend doesn't support unoptimized code " + "generation but "; + LOG(ERROR) << "--xla_backend_optimization_level is set to " << opt_level + << "!"; + LOG(ERROR) << "(Supported configuration is " + "--xla_backend_optimization_level >= 2.)"; + LOG(ERROR) << std::string(80, '*'); + } AddOptimizationPasses(opt_level, /*size_level=*/0, target_machine.get(), &module_passes, @@ -453,18 +452,21 @@ void GPUBackendInit(const HloModuleConfig& hlo_module_config) { // * 3-6 gives similar results as 2; // * >6 start hurting the performance of at least dot product kernels. // - // TODO(jingyue): The current threshold only considers the numbr of IR + // TODO(jingyue): The current threshold only considers the number of IR // instructions which do not accurately reflect the true cost. We need a // better cost model. FeedLLVMWithFlags({"-bonus-inst-threshold=2"}); - // TODO(b/22073864): Increase limit when scan memory dependency. - // This helps to reduce more redundant load instructions. + // Increase limit when scanning memory dependencies. This helps to reduce + // more redundant load instructions. // // The specific value is currently large enough for s3d in shoc benchmark, // which contains a lot of load instructions and many arithmetic instructions // between those loads. FeedLLVMWithFlags({"-memdep-block-scan-limit=500"}); + // Use div.approx -- it matters for some float-division heavy benchmarks. + FeedLLVMWithFlags({"-nvptx-prec-divf32=0"}); + llvm_ir::InitializeLLVMCommandLineOptions(hlo_module_config); // Initialize the NVPTX target; it's the only target we link with, so call its diff --git a/tensorflow/compiler/xla/service/gpu/multi_output_fusion.cc b/tensorflow/compiler/xla/service/gpu/multi_output_fusion.cc index 835924024b7b7de79624a369a69b07d72ac751ab..02e1207f377b8c28bf2566bee8cf3bcbc66794fb 100644 --- a/tensorflow/compiler/xla/service/gpu/multi_output_fusion.cc +++ b/tensorflow/compiler/xla/service/gpu/multi_output_fusion.cc @@ -41,50 +41,7 @@ GpuMultiOutputFusion::GpuMultiOutputFusion() : MultiOutputFusion(INT64_MAX) {} bool GpuMultiOutputFusion::ShapesCompatibleForFusion(HloInstruction* instr1, HloInstruction* instr2) { - auto get_element_instr = - [&](const HloInstruction* instr) -> const HloInstruction* { - const HloInstruction* element_instr = instr; - if (instr->opcode() == HloOpcode::kFusion) { - auto fused_expression_root = instr->fused_expression_root(); - if (instr->IsMultiOutputFusion()) { - // If possible, we want to pick a reduce operand of the fusion root, - // because it has the most constraints. - for (const auto* inst : fused_expression_root->operands()) { - if (IsReductionToVector(*inst)) { - return inst; - } - } - return fused_expression_root->operands()[0]; - } else { - element_instr = fused_expression_root; - } - } - return element_instr; - }; - - auto get_element_shape = [&](const HloInstruction* element_instr) { - // Special handling of kReduce instructions -- the fusion - // applies to the first operand. - if (IsReductionToVector(*element_instr)) { - return element_instr->operand(0)->shape(); - } - return element_instr->shape(); - }; - - // The shapes in all tuple operands should agree, unless it is a reduce. - // In that case, the operand of the reduce needs to have the same shape - // as the other tuple operands, but also we need to compare the output - // shapes of the reduces. - auto* element_instr_1 = get_element_instr(instr1); - auto* element_instr_2 = get_element_instr(instr2); - if (element_instr_1->opcode() == HloOpcode::kReduce && - element_instr_2->opcode() == HloOpcode::kReduce && - !ShapeUtil::Equal(element_instr_1->shape(), element_instr_2->shape())) { - return false; - } - // The elementwise output shapes must be the same (including layout). - return ShapeUtil::EqualIgnoringFpPrecision( - get_element_shape(element_instr_1), get_element_shape(element_instr_2)); + return ShapesCompatibleForMultiOutputFusion(*instr1, *instr2); } bool GpuMultiOutputFusion::IsFusible(HloInstruction* instr) { @@ -110,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()); @@ -140,6 +97,18 @@ bool GpuMultiOutputFusion::LegalToFuse(HloInstruction* instr1, return false; } + // The emitter only supports in-place DUS for fusions with a single DUS at the + // root. Don't sibling fuse DUS for now. + // TODO(b/119178699): Multi-output fusing DUS can improve performance if we + // share the input and output buffers and add support to the emitter. + if (instr1->fused_expression_root()->opcode() == + HloOpcode::kDynamicUpdateSlice || + (instr2->opcode() == HloOpcode::kFusion && + instr2->fused_expression_root()->opcode() == + HloOpcode::kDynamicUpdateSlice)) { + return false; + } + // Do this check last, as it may be expensive. return !GpuInstructionFusion::FusionWouldBeTooLarge(instr1, instr2); } @@ -180,6 +149,12 @@ bool GpuMultiOutputFusion::DoProducerConsumerMultiOutputFusion() { VLOG(3) << producer->name() << " is not fusible."; continue; } + // Never multi-output fuse constants. To the extent that we want to fuse + // constants, that should be handled by the regular fusion pass. + if (producer->opcode() == HloOpcode::kConstant) { + VLOG(3) << producer->name() << " is a constant."; + continue; + } const bool is_loop_fusion = producer->opcode() == HloOpcode::kFusion && producer->fusion_kind() == HloInstruction::FusionKind::kLoop; @@ -187,7 +162,7 @@ bool GpuMultiOutputFusion::DoProducerConsumerMultiOutputFusion() { VLOG(3) << producer->name() << " is not a loop fusion."; continue; } - if (!ShapesCompatibleForFusion(producer, consumer)) { + if (!ShapesCompatibleForMultiOutputFusion(*producer, *consumer)) { VLOG(3) << producer->name() << " has an incompatible shape."; continue; } 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 8a6e5327e082791ff857a89e840c6a4f045f0edb..40b87b16a195564c9b98497f79a70f1db0539d87 100644 --- a/tensorflow/compiler/xla/service/gpu/multi_output_fusion_test.cc +++ b/tensorflow/compiler/xla/service/gpu/multi_output_fusion_test.cc @@ -505,7 +505,7 @@ TEST_F(MultiOutputFusionTest, p1.1 = f16[2,2,2]{2,1,0} parameter(1) c0 = f16[] constant(0) broadcast = f16[2,2,2]{2,1,0} broadcast(f16[] c0), dimensions={} - greater-than = pred[2,2,2]{2,1,0} greater-than(f32[2,2,2]{2,1,0} p1.1, f32[2,2,2]{2,1,0} broadcast) + greater-than = pred[2,2,2]{2,1,0} greater-than(f16[2,2,2]{2,1,0} p1.1, f16[2,2,2]{2,1,0} broadcast) p0.1 = f16[2,2,2]{2,1,0} parameter(0) ROOT select = f16[2,2,2]{2,1,0} select(pred[2,2,2]{2,1,0} greater-than, f16[2,2,2]{2,1,0} p0.1, f16[2,2,2]{2,1,0} broadcast) } @@ -580,7 +580,7 @@ TEST_F(MultiOutputFusionTest, AvoidsLargeFusion) { // ... // where each of the (pi * pj)'s is represented as a fusion node so that // multi-output fusion will pay attention to it. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation::Builder b(TestName()); Shape shape = ShapeUtil::MakeShape(F32, {10, 100}); @@ -621,5 +621,38 @@ TEST_F(MultiOutputFusionTest, AvoidsLargeFusion) { } } +TEST_F(MultiOutputFusionTest, MultiOutputFusionDUS) { + auto module = ParseHloString(R"(HloModule dus_mof + fusion.1 { + p.0 = f16[50,96,1024]{2,1,0} parameter(0) + p.1 = s32[1]{0} parameter(1) + p.2 = f16[1,96,1024]{2,1,0} parameter(2) + c.0 = s32[] constant(0) + 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 { + p.0 = f16[50,96,1024]{2,1,0} parameter(0) + 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, p.1, c.0, c.0) + } + + ENTRY entry { + p.00 = f16[50,96,1024]{2,1,0} parameter(0) + p.01 = f16[50,96,1024]{2,1,0} parameter(1) + p.1 = s32[1]{0} parameter(2) + p.2 = f16[1,96,1024]{2,1,0} parameter(3) + + f1 = f16[50,96,1024] fusion(p.00, p.1, p.2), kind=kLoop, calls=fusion.1 + f2 = f16[50,96,1024] fusion(p.01, p.1, p.2), kind=kLoop, calls=fusion.2 + ROOT tuple = (f16[50,96,1024],f16[50,96,1024]) tuple(f1, f2) + })") + .ValueOrDie(); + ASSERT_FALSE(GpuMultiOutputFusion().Run(module.get()).ValueOrDie()); +} + } // namespace gpu } // namespace xla diff --git a/tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc b/tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc index 791d414c915e6f23d84a38ae99dcfa9a59ab6353..48f718b514cc9809d4100627f85af7aa05445d36 100644 --- a/tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc +++ b/tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc @@ -36,6 +36,9 @@ limitations under the License. #include "tensorflow/compiler/xla/service/buffer_liveness.h" #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" @@ -60,12 +63,14 @@ limitations under the License. #include "tensorflow/compiler/xla/service/gpu/stream_assignment.h" #include "tensorflow/compiler/xla/service/gpu/stream_executor_util.h" #include "tensorflow/compiler/xla/service/gpu/thunk_schedule.h" +#include "tensorflow/compiler/xla/service/gpu/variadic_op_splitter.h" #include "tensorflow/compiler/xla/service/hlo.pb.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/service/hlo_constant_folding.h" #include "tensorflow/compiler/xla/service/hlo_cse.h" #include "tensorflow/compiler/xla/service/hlo_dce.h" #include "tensorflow/compiler/xla/service/hlo_element_type_converter.h" +#include "tensorflow/compiler/xla/service/hlo_get_dimension_size_rewriter.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_pass_fix.h" #include "tensorflow/compiler/xla/service/hlo_pass_pipeline.h" @@ -75,6 +80,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" @@ -106,27 +112,34 @@ namespace { namespace tracing = tensorflow::tracing; -// Returns the directory containing nvvm libdevice files. config_cuda_data_dir -// should be equal to config().debug_options().xla_gpu_cuda_data_dir() of the -// HloModule being compiled. -string GetLibdeviceDir(const string& config_cuda_data_dir) { - std::vector potential_libdevice_dirs; - if (!config_cuda_data_dir.empty()) { - potential_libdevice_dirs.push_back(config_cuda_data_dir); - } - potential_libdevice_dirs.push_back(tensorflow::LibdeviceRoot()); - - // Tries all potential libdevice directories in the order they are inserted. - // Returns the first directory that exists in the file system. - for (const string& potential_libdevice_dir : potential_libdevice_dirs) { - if (tensorflow::Env::Default()->IsDirectory(potential_libdevice_dir).ok()) { - VLOG(2) << "Found libdevice dir " << potential_libdevice_dir; - return potential_libdevice_dir; - } - VLOG(2) << "Unable to find potential libdevice dir " - << potential_libdevice_dir; +// Returns a vector of potential locations of the CUDA root directory. +std::vector GetCudaRootCandidates( + const HloModuleConfig& hlo_module_config) { + std::vector potential_cuda_roots = tensorflow::CandidateCudaRoots(); + + // CUDA location explicitly specified by user via --xla_gpu_cuda_data_dir has + // highest priority. + string xla_gpu_cuda_data_dir = + hlo_module_config.debug_options().xla_gpu_cuda_data_dir(); + if (!xla_gpu_cuda_data_dir.empty()) { + potential_cuda_roots.insert(potential_cuda_roots.begin(), + xla_gpu_cuda_data_dir); } + return potential_cuda_roots; +} +// Returns the directory containing nvvm libdevice files. +string GetLibdeviceDir(const HloModuleConfig& hlo_module_config) { + for (const string& cuda_root : GetCudaRootCandidates(hlo_module_config)) { + string libdevice_dir = + tensorflow::io::JoinPath(cuda_root, "nvvm", "libdevice"); + VLOG(2) << "Looking for libdevice at " << libdevice_dir; + if (tensorflow::Env::Default()->IsDirectory(libdevice_dir).ok()) { + VLOG(2) << "Found libdevice dir " << libdevice_dir; + return libdevice_dir; + } + } + LOG(WARNING) << "Unable to find libdevice dir. Using '.'"; // Last resort: maybe in the current folder. return "."; } @@ -142,6 +155,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(), @@ -149,6 +163,14 @@ Status OptimizeHloModule(HloModule* hlo_module, se::StreamExecutor* stream_exec, // TODO(b/64094172): make Call work on GPU instead of inlining. pipeline.AddPass(); + auto cost_model = [](HloInstruction* conv) { + // We need a cost model for GPUs. Currently, do nothing. + return false; + }; + pipeline.AddPass(false); + pipeline.AddPass( + cost_model, + /*convert_batch_groups_only=*/true); // Convert BF16 operations to F32 operations so that the GPU backend can // support BF16 operations without directly implementing a BF16 lowering for // most ops. @@ -171,13 +193,16 @@ Status OptimizeHloModule(HloModule* hlo_module, se::StreamExecutor* stream_exec, /*rewrite_inference_op=*/true, /*rewrite_grad_op=*/true); + pipeline.AddPass(); + // BatchNormExpander can create zero-sized ops, so zero-sized HLO // elimination has to come after that pass. pipeline.AddPass(); - pass.AddPass( - /*is_layout_sensitive=*/false, - [](const Shape&, const Shape&) { return false; }); + AlgebraicSimplifierOptions options; + options.set_enable_permutation_sort_replacement(true); + pass.AddPass(options); + pass.AddPass(); pass.AddPass(); pass.AddPass(); pass.AddPass(); @@ -237,6 +262,8 @@ Status OptimizeHloModule(HloModule* hlo_module, se::StreamExecutor* stream_exec, { HloPassPipeline pipeline("post-layout_assignment"); + /* TODO(b/117531509): Use LayoutAssignment::InstructionCanChangeLayout after + * fixing the ticket. */ pipeline.AddInvariantChecker( /*layout_sensitive=*/true, /*allow_mixed_precision=*/false, @@ -244,11 +271,10 @@ 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. - pipeline.AddPass>( - /*is_layout_sensitive=*/true, - /*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. // @@ -286,6 +312,11 @@ Status OptimizeHloModule(HloModule* hlo_module, se::StreamExecutor* stream_exec, { HloPassFix fusion("fusion"); + // We try to split variadic ops with many parameters into several such ops + // to avoid exceeding the parameter space. + fusion.AddPass(); + /* TODO(b/117531509): Use LayoutAssignment::InstructionCanChangeLayout after + * fixing the ticket. */ fusion.AddInvariantChecker( /*layout_sensitive=*/true, /*allow_mixed_precision=*/false, @@ -300,6 +331,8 @@ Status OptimizeHloModule(HloModule* hlo_module, se::StreamExecutor* stream_exec, TF_RETURN_IF_ERROR(fusion.Run(hlo_module).status()); HloPassPipeline reduce_pipeline("reduce-precision"); + /* TODO(b/117531509): Use LayoutAssignment::InstructionCanChangeLayout after + * fixing the ticket. */ reduce_pipeline.AddInvariantChecker( /*is_layout_sensitive=*/true, /*allow_mixed_precision=*/false, LayoutAssignment::InstructionCanChangeLayout); @@ -328,6 +361,8 @@ Status PrepareHloModuleForIrEmitting(HloModule* hlo_module) { // (b/27180329). Therefore, in that case, we set the output to be a copy of // the parameter. HloPassPipeline pipeline("GPU-ir-emit-prepare"); + /* TODO(b/117531509): Use LayoutAssignment::InstructionCanChangeLayout after + * fixing the ticket. */ pipeline.AddInvariantChecker( /*layout_sensitive=*/true, /*allow_mixed_precision=*/false, @@ -459,14 +494,21 @@ void WarnIfBadDriverJITVersion() { // Compiles the given PTX string using ptxas and returns the resulting machine // code (i.e. a cubin) as a byte array. -StatusOr> CompilePtx(const string& ptx, int cc_major, - int cc_minor) { +StatusOr> CompilePtx( + const string& ptx, int cc_major, int cc_minor, + const HloModuleConfig& hlo_module_config) { tracing::ScopedActivity activity("Compile PTX", /*is_expensive=*/true); - const string ptxas_path = - tensorflow::io::JoinPath(tensorflow::CudaRoot(), "bin", "ptxas"); - VLOG(2) << "Using ptxas at " << ptxas_path; auto env = tensorflow::Env::Default(); + string ptxas_path; + for (const string& cuda_root : GetCudaRootCandidates(hlo_module_config)) { + ptxas_path = tensorflow::io::JoinPath(cuda_root, "bin", "ptxas"); + VLOG(2) << "Looking for ptxas at " << ptxas_path; + if (env->FileExists(ptxas_path).ok()) { + break; + } + } TF_RETURN_IF_ERROR(env->FileExists(ptxas_path)); + VLOG(2) << "Using ptxas at " << ptxas_path; WarnIfBadPtxasVersion(ptxas_path); @@ -499,6 +541,9 @@ StatusOr> CompilePtx(const string& ptx, int cc_major, if (VLOG_IS_ON(2)) { ptxas_args.push_back("-v"); } + if (hlo_module_config.debug_options().xla_gpu_disable_ptxas_optimizations()) { + ptxas_args.push_back("-O0"); + } ptxas_info_dumper.SetProgram(ptxas_path, ptxas_args); ptxas_info_dumper.SetChannelAction(tensorflow::CHAN_STDERR, tensorflow::ACTION_PIPE); @@ -532,14 +577,17 @@ StatusOr> NVPTXCompiler::RunHloPasses( std::unique_ptr module, se::StreamExecutor* stream_exec, DeviceMemoryAllocator* device_allocator) { // We dump the post-optimization HLO in RunBackend so no need to dump it here. - VLOG(2) << "*** HLO Before Optimization"; - XLA_VLOG_LINES(2, module->ToString()); + VLOG(3) << "*** HLO Before Optimization"; + XLA_VLOG_LINES(3, module->ToString()); XLA_SCOPED_LOGGING_TIMER("NVPTXCompiler::RunHloPasses"); tracing::ScopedActivity activity("HLO Transforms", module->name(), /*is_expensive=*/true); TF_RETURN_IF_ERROR( OptimizeHloModule(module.get(), stream_exec, device_allocator, this)); + + TF_RETURN_IF_ERROR(PrepareHloModuleForIrEmitting(module.get())); + return std::move(module); } @@ -550,8 +598,6 @@ StatusOr> NVPTXCompiler::RunBackend( TF_RET_CHECK(stream_exec != nullptr); - TF_RETURN_IF_ERROR(PrepareHloModuleForIrEmitting(module.get())); - llvm::LLVMContext llvm_context; std::string buffer; llvm::raw_string_ostream error(buffer); @@ -591,8 +637,8 @@ StatusOr> NVPTXCompiler::RunBackend( // include headers, so no need for us to print them ourselves. XLA_VLOG_LINES(1, buffer_assignment->GetStats().ToString()); XLA_VLOG_LINES(2, buffer_assignment->ToString()); - VLOG(2) << "*** HLO After Optimization"; - XLA_VLOG_LINES(2, module->ToString()); + VLOG(3) << "*** HLO After Optimization"; + XLA_VLOG_LINES(3, module->ToString()); const string xla_dump_optimized_hlo_proto_to = module->config().debug_options().xla_dump_optimized_hlo_proto_to(); if (!xla_dump_optimized_hlo_proto_to.empty()) { @@ -622,10 +668,10 @@ StatusOr> NVPTXCompiler::RunBackend( string ir_module_string_before_opt; const bool embed_ir_in_executable = module->config().debug_options().xla_embed_ir_in_executable(); - if (VLOG_IS_ON(2) || embed_ir_in_executable) { + if (VLOG_IS_ON(3) || embed_ir_in_executable) { ir_module_string_before_opt = llvm_ir::DumpModuleToString(llvm_module); - VLOG(2) << "LLVM module before optimizations:"; - XLA_VLOG_LINES(2, ir_module_string_before_opt); + VLOG(3) << "LLVM module before optimizations:"; + XLA_VLOG_LINES(3, ir_module_string_before_opt); } const string& ir_dump_directory = @@ -660,15 +706,13 @@ StatusOr> NVPTXCompiler::RunBackend( // Find the directory containing libdevice. To avoid searching for it every // time, we have a one-element cache, keyed on the module's config's // cuda_data_dir. - const auto& config_cuda_data_dir = - module->config().debug_options().xla_gpu_cuda_data_dir(); - if (cached_libdevice_dir_.empty() || - cached_cuda_data_dir_ != config_cuda_data_dir) { - cached_cuda_data_dir_ = config_cuda_data_dir; - cached_libdevice_dir_ = GetLibdeviceDir(config_cuda_data_dir); + if (cached_libdevice_dir_.empty()) { + cached_libdevice_dir_ = GetLibdeviceDir(module->config()); } libdevice_dir = cached_libdevice_dir_; } + VLOG(2) << "Libdevice dir = " << libdevice_dir << "\n"; + int cc_major, cc_minor; if (!stream_exec->GetDeviceDescription().cuda_compute_capability(&cc_major, &cc_minor)) { @@ -695,10 +739,10 @@ StatusOr> NVPTXCompiler::RunBackend( if (user_post_optimization_hook_) { TF_CHECK_OK(user_post_optimization_hook_(llvm_module)); } - VLOG(2) << "LLVM module after optimizations:"; - XLA_VLOG_LINES(2, llvm_ir::DumpModuleToString(llvm_module)); - VLOG(2) << "PTX:"; - XLA_VLOG_LINES(2, ptx); + VLOG(3) << "LLVM module after optimizations:"; + XLA_VLOG_LINES(3, llvm_ir::DumpModuleToString(llvm_module)); + VLOG(3) << "PTX:"; + XLA_VLOG_LINES(3, ptx); // Write PTX to IR dump directory, if IR dumping was requested. if (!ir_dump_directory.empty()) { @@ -717,13 +761,13 @@ StatusOr> NVPTXCompiler::RunBackend( } const std::vector cubin = - CompilePtxOrGetCachedResult(ptx, cc_major, cc_minor); + CompilePtxOrGetCachedResult(ptx, cc_major, cc_minor, module->config()); auto thunk_schedule = absl::make_unique( ir_emitter.ConsumeThunkSequence(), std::move(stream_assignment), hlo_schedule->ThunkLaunchOrder()); - VLOG(2) << "Printing the thunk schedule..."; - XLA_VLOG_LINES(2, thunk_schedule->ToString()); + VLOG(3) << "Printing the thunk schedule..."; + XLA_VLOG_LINES(3, thunk_schedule->ToString()); std::unique_ptr profile_index_map; std::unique_ptr profile_printer; @@ -734,8 +778,8 @@ StatusOr> NVPTXCompiler::RunBackend( 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); + profile_printer = CreateHloProfilePrinterData( + *profile_index_map, cost_analysis, entry_computation->name()); } auto* gpu_executable = new GpuExecutable( @@ -749,9 +793,9 @@ StatusOr> NVPTXCompiler::RunBackend( return std::unique_ptr(gpu_executable); } -std::vector NVPTXCompiler::CompilePtxOrGetCachedResult(const string& ptx, - int cc_major, - int cc_minor) { +std::vector NVPTXCompiler::CompilePtxOrGetCachedResult( + const string& ptx, int cc_major, int cc_minor, + const HloModuleConfig& hlo_module_config) { XLA_SCOPED_LOGGING_TIMER("NVPTXCompiler::CompilePtxOrGetCachedResult"); tracing::ScopedActivity activity("PTX->CUBIN", /*is_expensive=*/true); bool inserted; @@ -780,7 +824,7 @@ std::vector NVPTXCompiler::CompilePtxOrGetCachedResult(const string& ptx, CHECK(!cache_value->compilation_done); if (!ptx.empty()) { StatusOr> maybe_cubin = - CompilePtx(*cache_ptx, cc_major, cc_minor); + CompilePtx(*cache_ptx, cc_major, cc_minor, hlo_module_config); if (maybe_cubin.ok()) { cache_value->cubin_data = std::move(maybe_cubin).ValueOrDie(); VLOG(2) << "Compiled PTX size:" << ptx.size() @@ -793,7 +837,7 @@ std::vector NVPTXCompiler::CompilePtxOrGetCachedResult(const string& ptx, // binaries are not available. We don't want to spam logs with // identical warnings in this case. - // TODO(zhengxq): we should implement a LOG_FIRST_N and LOG_EVERY_N + // TODO(jlebar): we should implement a LOG_FIRST_N and LOG_EVERY_N // for more general usage. static std::atomic warning_done(false); log_warning = !warning_done.exchange(true); diff --git a/tensorflow/compiler/xla/service/gpu/nvptx_compiler.h b/tensorflow/compiler/xla/service/gpu/nvptx_compiler.h index f79ae2990ae7d6e6985b15727a72358289121aa9..b2077f42fd097330703fde063d80a20704fa48e2 100644 --- a/tensorflow/compiler/xla/service/gpu/nvptx_compiler.h +++ b/tensorflow/compiler/xla/service/gpu/nvptx_compiler.h @@ -97,8 +97,9 @@ class NVPTXCompiler : public LLVMCompiler { // Tries to compile the given ptx string to cubin. Returns a vector with the // compiled cubin. If compilation was unsuccessful, returns an empty vector. - std::vector CompilePtxOrGetCachedResult(const string& ptx, - int cc_major, int cc_minor); + std::vector CompilePtxOrGetCachedResult( + const string& ptx, int cc_major, int cc_minor, + const HloModuleConfig& hlo_module_config); // The compilation_cache_ map is a cache from {ptx string, cc_major, cc_minor} // -> cubin so we don't recompile the same ptx twice. This is important for diff --git a/tensorflow/compiler/xla/service/gpu/partition_assignment.cc b/tensorflow/compiler/xla/service/gpu/partition_assignment.cc index 375f68a15957936151aee068582a714b62694af2..bfed4f5230dfe37bca48560ce83a2dd82c8950a4 100644 --- a/tensorflow/compiler/xla/service/gpu/partition_assignment.cc +++ b/tensorflow/compiler/xla/service/gpu/partition_assignment.cc @@ -39,6 +39,25 @@ std::ostream& operator<<(std::ostream& out, return out; } +int64 ThreadsPerBlockLimit(const se::DeviceDescription& device_desc) { + int64 threads_per_block = device_desc.threads_per_block_limit(); + if (threads_per_block == 0) { + static std::atomic log_count{0}; + if (log_count.fetch_add(1) < 8) { + LOG(WARNING) << "Attempting to calculate launch dimensions for GPU " + "without full information about its capabilities. " + "StreamExecutor's PopulateDeviceDescription should be " + "updated for this device."; + } + threads_per_block = device_desc.threads_per_warp(); + if (threads_per_block == 0) { + // Fall back to *something* if we can't even get num threads per warp. + threads_per_block = 32; + } + } + return threads_per_block; +} + // Calculates the launch dimensions used to invoke `hlo`. LaunchDimensions CalculateLaunchDimensions( const Shape& shape, const se::DeviceDescription& device_desc, @@ -62,21 +81,7 @@ LaunchDimensions CalculateLaunchDimensions( // // * = - int64 threads_per_block = device_desc.threads_per_block_limit(); - if (threads_per_block == 0) { - static std::atomic log_count{0}; - if (log_count.fetch_add(1) < 8) { - LOG(WARNING) << "Attempting to calculate launch dimensions for GPU " - "without full information about its capabilities. " - "StreamExecutor's PopulateDeviceDescription should be " - "updated for this device."; - } - threads_per_block = device_desc.threads_per_warp(); - if (threads_per_block == 0) { - // Fall back to *something* if we can't even get num threads per warp. - threads_per_block = 32; - } - } + int64 threads_per_block = ThreadsPerBlockLimit(device_desc); if (num_elements < threads_per_block) { threads_per_block = num_elements; diff --git a/tensorflow/compiler/xla/service/gpu/partition_assignment.h b/tensorflow/compiler/xla/service/gpu/partition_assignment.h index 02471129e004b4876ce20a62cade34060c65b478..eb41dcccb938ccc088c2371def96ca73276771ab 100644 --- a/tensorflow/compiler/xla/service/gpu/partition_assignment.h +++ b/tensorflow/compiler/xla/service/gpu/partition_assignment.h @@ -57,6 +57,9 @@ class LaunchDimensions { std::ostream& operator<<(std::ostream& out, const LaunchDimensions& launch_dims); +// Returns the maximum number of threads per block allowed by the device. +int64 ThreadsPerBlockLimit(const se::DeviceDescription& device_desc); + LaunchDimensions CalculateLaunchDimensions( const Shape& shape, const se::DeviceDescription& device_desc, int unroll_factor = 1); diff --git a/tensorflow/compiler/xla/service/gpu/stream_assignment.cc b/tensorflow/compiler/xla/service/gpu/stream_assignment.cc index 5b6cf2c04d05378a363232e33a6df6432cd6848e..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; } } @@ -122,7 +123,7 @@ std::unique_ptr AssignStreams(const HloModule& module) { auto stream_assignment = absl::make_unique(); const HloComputation& computation = *module.entry_computation(); std::unique_ptr reachability = - computation.ComputeReachability(); + HloReachabilityMap::Build(&computation); std::vector seen_gemms; // The execution of different RNG Hlo instructions in the same module updates // a common global variable. To avoid a race condition, we simply assign all diff --git a/tensorflow/compiler/xla/service/gpu/stream_assignment_test.cc b/tensorflow/compiler/xla/service/gpu/stream_assignment_test.cc index c4f43cc9a614283acb376b5f98e4976615b590ad..31a5d7a8c04e9863830e2026fc73cd7ded8c322e 100644 --- a/tensorflow/compiler/xla/service/gpu/stream_assignment_test.cc +++ b/tensorflow/compiler/xla/service/gpu/stream_assignment_test.cc @@ -21,16 +21,16 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" #include "tensorflow/compiler/xla/test_helpers.h" -#include "tensorflow/compiler/xla/tests/hlo_verified_test_base.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" #include "tensorflow/compiler/xla/tests/test_utils.h" #include "tensorflow/compiler/xla/types.h" namespace xla { namespace gpu { -class StreamAssignmentTest : public HloVerifiedTestBase { +class StreamAssignmentTest : public HloTestBase { protected: - std::unique_ptr CreateNewModule() { + std::unique_ptr CreateNewVerifiedModule() { HloModuleConfig config; auto debug_options = GetDebugOptionsForTest(); debug_options.set_xla_gpu_disable_multi_streaming(false); @@ -55,7 +55,7 @@ TEST_F(StreamAssignmentTest, SequentialMatMul) { HloInstruction* dot2 = builder.AddInstruction(CreateCanonicalDot(f32_2x2_, dot1, z)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build(dot2)); std::unique_ptr assignment = AssignStreams(*module); @@ -76,7 +76,7 @@ TEST_F(StreamAssignmentTest, ConcurrentMatMul) { HloInstruction* add = builder.AddInstruction( HloInstruction::CreateBinary(f32_2x2_, HloOpcode::kAdd, dot1, dot2)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build(add)); std::unique_ptr assignment = AssignStreams(*module); @@ -120,7 +120,7 @@ TEST_F(StreamAssignmentTest, LatticeMatMul) { HloInstruction* d40 = builder.AddInstruction(CreateCanonicalDot(f32_2x2_, d30, d31)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build(d40)); std::unique_ptr assignment = AssignStreams(*module); diff --git a/tensorflow/compiler/xla/service/gpu/stream_executor_util.h b/tensorflow/compiler/xla/service/gpu/stream_executor_util.h index 1fc46bafa10e7ba6c896f081d5c836bd400886c9..92e4d6dbbc1bd564657f8a5de09d23d5ae81a93e 100644 --- a/tensorflow/compiler/xla/service/gpu/stream_executor_util.h +++ b/tensorflow/compiler/xla/service/gpu/stream_executor_util.h @@ -16,6 +16,7 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_GPU_STREAM_EXECUTOR_UTIL_H_ #define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_STREAM_EXECUTOR_UTIL_H_ +#include "tensorflow/compiler/xla/layout.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" diff --git a/tensorflow/compiler/xla/service/gpu/tests/BUILD b/tensorflow/compiler/xla/service/gpu/tests/BUILD index ed46f08d5970d479db33a7b9ad416a1480535764..d798b31643782eb25bba08227e29903ec0e7a597 100644 --- a/tensorflow/compiler/xla/service/gpu/tests/BUILD +++ b/tensorflow/compiler/xla/service/gpu/tests/BUILD @@ -37,7 +37,7 @@ cc_library( hdrs = ["gpu_codegen_test.h"], tags = tf_cuda_tests_tags(), deps = [ - "//tensorflow/compiler/xla/legacy_flags:debug_options_flags", + "//tensorflow/compiler/xla:debug_options_flags", "//tensorflow/compiler/xla/service:gpu_plugin", "//tensorflow/compiler/xla/service/gpu:gpu_executable", "//tensorflow/compiler/xla/tests:filecheck", diff --git a/tensorflow/compiler/xla/service/gpu/tests/gpu_codegen_test.cc b/tensorflow/compiler/xla/service/gpu/tests/gpu_codegen_test.cc index 79e77d4c4d649020cf52ac25c220c3f90e8469b9..9e3ff8750b88d08bcbc1aae3faead5aecfa19848 100644 --- a/tensorflow/compiler/xla/service/gpu/tests/gpu_codegen_test.cc +++ b/tensorflow/compiler/xla/service/gpu/tests/gpu_codegen_test.cc @@ -15,7 +15,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/gpu/tests/gpu_codegen_test.h" #include "absl/memory/memory.h" -#include "tensorflow/compiler/xla/legacy_flags/debug_options_flags.h" +#include "tensorflow/compiler/xla/debug_options_flags.h" #include "tensorflow/compiler/xla/service/gpu/gpu_executable.h" #include "tensorflow/compiler/xla/tests/filecheck.h" #include "tensorflow/core/platform/logging.h" @@ -23,9 +23,10 @@ limitations under the License. namespace xla { namespace gpu { -std::unique_ptr GpuCodegenTest::CreateNewModuleWithFTZ(bool ftz) { +std::unique_ptr GpuCodegenTest::CreateNewUnverifiedModuleWithFTZ( + bool ftz) { HloModuleConfig config; - auto debug_options = legacy_flags::GetDebugOptionsFromFlags(); + auto debug_options = GetDebugOptionsFromFlags(); debug_options.set_xla_gpu_ftz(ftz); debug_options.set_xla_gpu_max_kernel_unroll_factor(1); // TODO(b/38354253): Change tests to use Parameters instead of Constants. diff --git a/tensorflow/compiler/xla/service/gpu/tests/gpu_codegen_test.h b/tensorflow/compiler/xla/service/gpu/tests/gpu_codegen_test.h index e4a3573babb7ed746504c1466f85b582aa4d044f..d917320e36363c4fa7e4c0055e8f3345cbc610a2 100644 --- a/tensorflow/compiler/xla/service/gpu/tests/gpu_codegen_test.h +++ b/tensorflow/compiler/xla/service/gpu/tests/gpu_codegen_test.h @@ -26,9 +26,9 @@ namespace gpu { // Tests that verify IR or PTX emitted by the GPU backend is as expected. class GpuCodegenTest : public LlvmIrGenTestBase { protected: - // Like HloTestBase::CreateNewModule(), with a flag for configuring the ftz - // option. - std::unique_ptr CreateNewModuleWithFTZ(bool ftz); + // Like HloTestBase::CreateNewVerifiedModule(), with a flag for configuring + // the ftz option. + std::unique_ptr CreateNewUnverifiedModuleWithFTZ(bool ftz); // Compiles the given HLO module to PTX and verifies the PTX matches the given // FileCheck pattern. (See http://llvm.org/docs/CommandGuide/FileCheck.html). 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 780539c164277f14c2bd964024f7c3ca179f4ada..a1ed8499040359fe7265a7317b0577a990a2234c 100644 --- a/tensorflow/compiler/xla/service/gpu/tests/gpu_copy_test.cc +++ b/tensorflow/compiler/xla/service/gpu/tests/gpu_copy_test.cc @@ -46,7 +46,7 @@ TEST_F(GpuCopyTest, UseMemcpy) { std::unique_ptr computation = builder.Build(); - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewVerifiedModule(); hlo_module->AddEntryComputation(std::move(computation)); // There should not be any kernel prefixed "copy". diff --git a/tensorflow/compiler/xla/service/gpu/tests/gpu_ftz_test.cc b/tensorflow/compiler/xla/service/gpu/tests/gpu_ftz_test.cc index 177b94934c7f519172508b5cc6e088f908401193..5e524faab18947f5793dc2ae34e9329a446d4235 100644 --- a/tensorflow/compiler/xla/service/gpu/tests/gpu_ftz_test.cc +++ b/tensorflow/compiler/xla/service/gpu/tests/gpu_ftz_test.cc @@ -39,7 +39,7 @@ class GpuFtzTest : public GpuCodegenTest { /* parameter_number=*/1, param_shape, "y")); builder.AddInstruction(HloInstruction::CreateBinary(param_shape, op, x, y)); - auto hlo_module = CreateNewModuleWithFTZ(ftz_); + auto hlo_module = CreateNewUnverifiedModuleWithFTZ(ftz_); hlo_module->AddEntryComputation(builder.Build()); return hlo_module; } @@ -54,7 +54,7 @@ class GpuFtzTest : public GpuCodegenTest { /* parameter_number=*/0, param_shape, "x")); builder.AddInstruction(HloInstruction::CreateUnary(param_shape, op, x)); - auto hlo_module = CreateNewModuleWithFTZ(ftz_); + auto hlo_module = CreateNewUnverifiedModuleWithFTZ(ftz_); hlo_module->AddEntryComputation(builder.Build()); return hlo_module; } @@ -75,16 +75,16 @@ class GpuFtzDisabledTest : public GpuFtzTest { // Check that we emit mul.ftz.f32 when in ftz mode, and plain mul.f32 otherwise. TEST_F(GpuFtzEnabledTest, MultiplyFtz) { CompileAndVerifyPtx(CreateBinaryOpModule(HloOpcode::kMultiply), R"( - CHECK-NOT: mul.f32 - CHECK: mul.ftz.f32 - CHECK-NOT: mul.f32 + CHECK-NOT: mul.rn.f32 + CHECK: mul.rn.ftz.f32 + CHECK-NOT: mul.rn.f32 )"); } TEST_F(GpuFtzDisabledTest, MultiplyFtz) { CompileAndVerifyPtx(CreateBinaryOpModule(HloOpcode::kMultiply), R"( - CHECK-NOT: mul.ftz.f32 - CHECK: mul.f32 - CHECK-NOT: mul.ftz.f32 + CHECK-NOT: mul.rn.ftz.f32 + CHECK: mul.rn.f32 + CHECK-NOT: mul.rn.ftz.f32 )"); } diff --git a/tensorflow/compiler/xla/service/gpu/tests/gpu_index_test.cc b/tensorflow/compiler/xla/service/gpu/tests/gpu_index_test.cc index a06576df7b874745236a8d9075355a01ec42e777..6814be779e0b02c38e3bc7008f036b845d88cb6f 100644 --- a/tensorflow/compiler/xla/service/gpu/tests/gpu_index_test.cc +++ b/tensorflow/compiler/xla/service/gpu/tests/gpu_index_test.cc @@ -51,7 +51,7 @@ TEST_F(GpuIndexTest, CompatibleUseLinearIndex) { builder.AddInstruction(HloInstruction::CreateBinary( ShapeUtil::MakeShape(PRED, {5, 7, 2}), HloOpcode::kGe, param_x, param_y)); - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewVerifiedModule(); hlo_module->AddEntryComputation(builder.Build()); // Check the optimized IR as the unoptimized IR contains dead udiv and urem. 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 15d1e269cc22b88f5269175084f20600f165011c..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,12 +187,230 @@ 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, TransposedInputWithUserReverseNotTiled) { + const char *const kHloString = R"( + HloModule FusionTransposeWithReverseNotTiled + fused_computation.1 { + arg0 = f32[128,64]{1,0} parameter(0) + copy0 = f32[128,64]{0,1} copy(arg0) + ROOT reverse0 = f32[128,64]{0,1} reverse(copy0), dimensions={0} + } + + ENTRY reverse_break_assumption { + param0 = f32[128,64]{1,0} parameter(0) + ROOT fusion0 = f32[128,64]{0,1} fusion(param0), kind=kLoop, + calls=fused_computation.1 + })"; + + // 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); +} + +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/gpu_ldg_test.cc b/tensorflow/compiler/xla/service/gpu/tests/gpu_ldg_test.cc index 6a9ecd9dae7c9ddde0b56d8615e4a39fb3df0af9..3019215c015a4e0aa094a62424d650ced0de2a0e 100644 --- a/tensorflow/compiler/xla/service/gpu/tests/gpu_ldg_test.cc +++ b/tensorflow/compiler/xla/service/gpu/tests/gpu_ldg_test.cc @@ -48,7 +48,7 @@ TEST_F(GpuLdgTest, LdgForParamRead) { HloInstruction::CreateBinary(shape, HloOpcode::kAdd, param, param)); std::unique_ptr computation = builder.Build(); - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewVerifiedModule(); hlo_module->AddEntryComputation(std::move(computation)); CompileAndVerifyPtx(std::move(hlo_module), R"( @@ -73,7 +73,7 @@ TEST_F(GpuLdgTest, LdgForNonParamRead) { builder.AddInstruction(HloInstruction::CreateTuple({add, square})); std::unique_ptr computation = builder.Build(); - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewVerifiedModule(); hlo_module->AddEntryComputation(std::move(computation)); CompileAndVerifyPtx(std::move(hlo_module), R"( @@ -95,7 +95,7 @@ TEST_F(GpuLdgTest, LdgForNonParamRead) { // reduce in the foreseeable future. But if that turns out to be wrong, I give // you, future reader, permission to delete this test. TEST_F(GpuLdgTest, NoLdgWhenSharingBuffer) { - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); HloComputation* reduce_computation; diff --git a/tensorflow/compiler/xla/service/gpu/tests/gpu_noalias_test.cc b/tensorflow/compiler/xla/service/gpu/tests/gpu_noalias_test.cc index 15198865bda98f9718342d5a444a20305f923b48..ca0a78034d7dc83d17ad72202914d95f37ac122b 100644 --- a/tensorflow/compiler/xla/service/gpu/tests/gpu_noalias_test.cc +++ b/tensorflow/compiler/xla/service/gpu/tests/gpu_noalias_test.cc @@ -47,7 +47,7 @@ TEST_F(GpuNoAliasTest, Concat) { std::unique_ptr computation = builder.Build(); - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewVerifiedModule(); hlo_module->AddEntryComputation(std::move(computation)); CompileAndVerifyIr(std::move(hlo_module), diff --git a/tensorflow/compiler/xla/service/gpu/tests/gpu_unrolling_test.cc b/tensorflow/compiler/xla/service/gpu/tests/gpu_unrolling_test.cc index 0f2d5568cafc9db0f5f067437fdd5e2e775ad2c8..4636f1d9d20b8c213ffadec427b3820a89c68a7f 100644 --- a/tensorflow/compiler/xla/service/gpu/tests/gpu_unrolling_test.cc +++ b/tensorflow/compiler/xla/service/gpu/tests/gpu_unrolling_test.cc @@ -85,7 +85,7 @@ TEST_F(GpuUnrollingTest, UnrollFourTimes) { TEST_F(GpuUnrollingTest, UnrollDefaultTimes) { // The default unrolling factor is 4. HloModuleConfig config; - config.set_debug_options(legacy_flags::GetDebugOptionsFromFlags()); + config.set_debug_options(GetDebugOptionsFromFlags()); auto hlo_module = ParseHloString(kAddModule, config).ValueOrDie(); CompileAndVerifyIr(std::move(hlo_module), diff --git a/tensorflow/compiler/xla/service/gpu/tests/infeed_test.cc b/tensorflow/compiler/xla/service/gpu/tests/infeed_test.cc index f8120a5fa00ce38644cd85c54d5ef65701be1eda..f91a22d482bc8bc046977870a7a4d18ca1acde68 100644 --- a/tensorflow/compiler/xla/service/gpu/tests/infeed_test.cc +++ b/tensorflow/compiler/xla/service/gpu/tests/infeed_test.cc @@ -43,7 +43,7 @@ class InfeedTest : public ClientLibraryTestBase { ASSERT_IS_OK(client_->TransferToInfeed(literal)); XlaBuilder builder(TestName()); Infeed(&builder, literal.shape()); - if (ShapeUtil::IsTuple(literal.shape())) { + if (literal.shape().IsTuple()) { // TODO(b/30609564): Use ComputeAndCompareLiteral instead. ComputeAndCompareTuple(&builder, literal, {}); } else { diff --git a/tensorflow/compiler/xla/service/gpu/thunk_schedule.cc b/tensorflow/compiler/xla/service/gpu/thunk_schedule.cc index 141f3219387940a08ef22cbcc0be0971a14c2cd6..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. @@ -45,17 +47,17 @@ void ThunkSchedule::AddDependenciesOnTransitiveOperands( ThunkSchedule::ThunkSchedule( std::unique_ptr thunks, std::unique_ptr stream_assignment, - const std::vector& hlo_total_order) + 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 (const HloInstruction* hlo : hlo_total_order) { - if (hlo_to_thunk.count(hlo)) { - thunk_total_order_.push_back(FindOrDie(hlo_to_thunk, hlo)); + for (HloInstruction* hlo : hlo_total_order) { + 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 d3352994f845a535233612a17e19107511ce0622..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" @@ -46,7 +48,7 @@ class ThunkSchedule { public: ThunkSchedule(std::unique_ptr thunks, std::unique_ptr stream_assignment, - const std::vector& hlo_total_order); + const std::vector& hlo_total_order); // Returns the total order of executing all the thunks. const std::vector& TotalOrder() const { return thunk_total_order_; } @@ -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/variadic_op_splitter.cc b/tensorflow/compiler/xla/service/gpu/variadic_op_splitter.cc new file mode 100644 index 0000000000000000000000000000000000000000..c552c2925497f1c4808d74a615d35cdbeeba1858 --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/variadic_op_splitter.cc @@ -0,0 +1,106 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "tensorflow/compiler/xla/service/gpu/variadic_op_splitter.h" + +#include + +#include "absl/types/span.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" +#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 { +namespace gpu { + +namespace { +// The parameter space on the GPU device is limited. We pick an arbitrary low +// constant here to try to prevent exceeding this parameter space. For a proper +// fix, we would have to take into account which parameters share a buffer, and +// how big these buffers are. +constexpr int32 kMaxParameters = 128; + +StatusOr SplitConcatenate(HloInstruction* concat, HloComputation* comp) { + auto operands = concat->operands(); + std::vector operands_to_split(operands.begin(), + operands.end()); + while (operands_to_split.size() > 1) { + std::vector new_operands; + absl::Span operands_span(operands_to_split); + for (int64 offset = 0; offset < operands_to_split.size(); + offset += kMaxParameters) { + // Check if there is a remainder of operands that does not completely fill + // one "batch" of exactly 'kMaxParameters' operands. If there are only + // less than 'kMaxParameters' operands left, then we still put them into a + // concat together. Otherwise, we spare them for another round so that + // they can be put together into a concat with some of the newly created + // concats. + if (offset > 0 && offset + kMaxParameters > operands_to_split.size()) { + new_operands.insert(new_operands.end(), + operands_to_split.begin() + offset, + operands_to_split.end()); + } else { + Shape new_shape = concat->shape(); + int64 concat_dimension_size = 0; + for (int64 i = 0; + i < kMaxParameters && offset + i < operands_to_split.size(); ++i) { + concat_dimension_size += + operands_to_split[i + offset]->shape().dimensions( + concat->concatenate_dimension()); + } + new_shape.set_dimensions(concat->concatenate_dimension(), + concat_dimension_size); + auto new_concat = comp->AddInstruction(concat->CloneWithNewOperands( + new_shape, operands_span.subspan(offset, kMaxParameters))); + new_operands.push_back(new_concat); + } + } + operands_to_split = new_operands; + } + TF_RETURN_IF_ERROR(comp->ReplaceInstruction(concat, operands_to_split[0])); + return true; +} + +std::vector GetRelevantVariadicOps(HloComputation* comp) { + std::vector ops; + for (HloInstruction* instr : comp->instructions()) { + if (instr->opcode() == HloOpcode::kConcatenate && + instr->operand_count() > kMaxParameters) { + ops.push_back(instr); + } + } + return ops; +} + +} // namespace + +StatusOr VariadicOpSplitter::Run(HloModule* module) { + bool changed = false; + for (HloComputation* comp : module->MakeNonfusionComputations()) { + for (HloInstruction* op : GetRelevantVariadicOps(comp)) { + // TODO(b/112613927): Handle also other ops than concatenate. + TF_ASSIGN_OR_RETURN(bool result, SplitConcatenate(op, comp)); + changed |= result; + } + } + return changed; +} + +} // namespace gpu +} // namespace xla diff --git a/tensorflow/compiler/xla/service/gpu/variadic_op_splitter.h b/tensorflow/compiler/xla/service/gpu/variadic_op_splitter.h new file mode 100644 index 0000000000000000000000000000000000000000..7673ad0d48a04508987025dac84b60e396e3d7dc --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/variadic_op_splitter.h @@ -0,0 +1,38 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_GPU_VARIADIC_OP_SPLITTER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_VARIADIC_OP_SPLITTER_H_ + +#include "absl/strings/string_view.h" +#include "tensorflow/compiler/xla/service/hlo_pass_interface.h" +#include "tensorflow/compiler/xla/statusor.h" + +namespace xla { +namespace gpu { + +// Splits variadic ops with many operands into pieces such that we don't exceed +// the parameter space on the GPU. Currently only concatenate ops are split up. +class VariadicOpSplitter : public HloModulePass { + public: + absl::string_view name() const override { return "variadic-op-splitter"; } + + StatusOr Run(HloModule* module) override; +}; + +} // namespace gpu +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_VARIADIC_OP_SPLITTER_H_ diff --git a/tensorflow/compiler/xla/service/gpu/variadic_op_splitter_test.cc b/tensorflow/compiler/xla/service/gpu/variadic_op_splitter_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..3d00ac4dc7b57664a317157c093d7ffaa01b4fd6 --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/variadic_op_splitter_test.cc @@ -0,0 +1,82 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "tensorflow/compiler/xla/service/gpu/variadic_op_splitter.h" + +#include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/compiler/xla/service/gpu/ir_emission_utils.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" +#include "tensorflow/compiler/xla/service/hlo_parser.h" +#include "tensorflow/compiler/xla/service/pattern_matcher.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" +#include "tensorflow/compiler/xla/util.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" + +namespace xla { +namespace gpu { +namespace { +using match::Concatenate; + +class VariadicOpSplitterTest : public HloTestBase {}; + +TEST_F(VariadicOpSplitterTest, DontSplit) { + auto module = ParseAndReturnVerifiedModule(R"( + HloModule TestModule + + ENTRY TestComputation { + p0 = f16[30,41] parameter(0) + p1 = f16[30,41] parameter(1) + ROOT result = f16[60, 41] concatenate(p0, p1), dimensions={0} + })") + .ValueOrDie(); + EXPECT_FALSE(VariadicOpSplitter().Run(module.get()).ValueOrDie()); +} + +TEST_F(VariadicOpSplitterTest, SplitInto2) { + auto builder = HloComputation::Builder(TestName()); + auto operand = builder.AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR1({42}))); + std::vector concat_operands(255, operand); + builder.AddInstruction(HloInstruction::CreateConcatenate( + ShapeUtil::MakeShape(S32, {255}), concat_operands, 0)); + auto module = CreateNewVerifiedModule(); + auto entry_computation = module->AddEntryComputation(builder.Build()); + EXPECT_TRUE(VariadicOpSplitter().Run(module.get()).ValueOrDie()); + EXPECT_TRUE(Match(entry_computation->root_instruction(), + Concatenate().WithNumOperands(128).WithOperand( + 0, Concatenate().WithNumOperands(128)))); +} + +TEST_F(VariadicOpSplitterTest, SplitInto3) { + auto builder = HloComputation::Builder(TestName()); + auto operand = builder.AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR1({42}))); + std::vector concat_operands(256, operand); + builder.AddInstruction(HloInstruction::CreateConcatenate( + ShapeUtil::MakeShape(S32, {256}), concat_operands, 0)); + auto module = CreateNewVerifiedModule(); + auto entry_computation = module->AddEntryComputation(builder.Build()); + EXPECT_TRUE(VariadicOpSplitter().Run(module.get()).ValueOrDie()); + EXPECT_TRUE(Match(entry_computation->root_instruction(), + Concatenate(Concatenate().WithNumOperands(128), + Concatenate().WithNumOperands(128)))); +} + +} // namespace +} // namespace gpu +} // namespace xla diff --git a/tensorflow/compiler/xla/service/gpu/while_transformer_test.cc b/tensorflow/compiler/xla/service/gpu/while_transformer_test.cc index 926b59a1b854bd3d7d2699124e10b70147e52e2a..2dce7749bbd8da2673ae607eee3d731d9917e8fe 100644 --- a/tensorflow/compiler/xla/service/gpu/while_transformer_test.cc +++ b/tensorflow/compiler/xla/service/gpu/while_transformer_test.cc @@ -29,7 +29,7 @@ namespace { class WhileTransformerTest : public HloTestBase { protected: WhileTransformerTest() - : module_(CreateNewModule()), + : module_(CreateNewVerifiedModule()), induction_variable_shape_(ShapeUtil::MakeShape(S32, {})), data_shape_(ShapeUtil::MakeShape(F32, {8})), condition_result_shape_(ShapeUtil::MakeShape(PRED, {})) {} 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/heap_simulator_test.cc b/tensorflow/compiler/xla/service/heap_simulator_test.cc index e30e7667f3015bc7bfe67c65147a5016332780f7..dc40b9446ad1bffcb757543e52fc9ab20de6d52e 100644 --- a/tensorflow/compiler/xla/service/heap_simulator_test.cc +++ b/tensorflow/compiler/xla/service/heap_simulator_test.cc @@ -30,16 +30,16 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_value.h" #include "tensorflow/compiler/xla/service/tuple_points_to_analysis.h" #include "tensorflow/compiler/xla/status_macros.h" -#include "tensorflow/compiler/xla/tests/hlo_verified_test_base.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" #include "tensorflow/core/lib/core/status_test_util.h" namespace xla { namespace { -class MinimumMemoryForSequenceTest : public HloVerifiedTestBase {}; +class MinimumMemoryForSequenceTest : public HloTestBase {}; TEST_F(MinimumMemoryForSequenceTest, MultiComputation) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); const Shape scalar_shape = ShapeUtil::MakeShape(xla::F32, {}); const Shape tuple_shape = ShapeUtil::MakeTupleShape({scalar_shape, scalar_shape}); @@ -86,7 +86,7 @@ TEST_F(MinimumMemoryForSequenceTest, MultiComputation) { return ShapeUtil::ByteSizeOf(buffer.shape(), /*pointer_size=*/8); }; - HloSchedule schedule(module); + HloSchedule schedule(module.get()); schedule.set_sequence(cond_computation, {cond_param, cond_iter, cond_data, cond_lt}); schedule.set_sequence(body_computation, {body_param}); @@ -258,7 +258,7 @@ class HeapSimulatorTracker { // Constructor for testing a single entry computation. HeapSimulatorTracker( const string& name, std::unique_ptr computation, - const std::vector& instruction_sequence) { + const std::vector& instruction_sequence) { HloModuleConfig config; module_ = absl::make_unique(name, config); module_->AddEntryComputation(std::move(computation)); @@ -286,7 +286,7 @@ class HeapSimulatorTracker { // Similar to the single entry computation constructor above, but runs the // simulation over the entire module. void RunWholeModule( - const std::vector& full_module_sequence) { + const std::vector& full_module_sequence) { points_to_analysis_ = TuplePointsToAnalysis::Run(module_.get()).ConsumeValueOrDie(); @@ -294,7 +294,7 @@ class HeapSimulatorTracker { HloSchedule schedule(module_.get()); absl::flat_hash_map reverse_position; for (int i = 0; i < full_module_sequence.size(); ++i) { - const HloInstruction* instruction = full_module_sequence[i]; + HloInstruction* instruction = full_module_sequence[i]; schedule.GetOrCreateSequence(instruction->parent()) .push_back(instruction); reverse_position[instruction] = full_module_sequence.size() - i; @@ -351,7 +351,7 @@ class HeapSimulatorTracker { HeapSimulator::Result result_; }; -class HeapSimulatorTest : public HloVerifiedTestBase { +class HeapSimulatorTest : public HloTestBase { protected: HeapSimulatorTest() {} ~HeapSimulatorTest() override {} diff --git a/tensorflow/compiler/xla/service/hlo.proto b/tensorflow/compiler/xla/service/hlo.proto index dbab62f847e8ca5e0b46dfd4162a0f4222640252..263b42a29dbb0dbc0fb6eca7968674ff242f45ed 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: 58 +// Next ID: 59 message HloInstructionProto { reserved 10; reserved "parameter_name"; @@ -51,7 +51,7 @@ message HloInstructionProto { string name = 1; string opcode = 2; - xla.Shape shape = 3; + xla.ShapeProto shape = 3; xla.OpMetadata metadata = 7; @@ -82,6 +82,8 @@ message HloInstructionProto { // it will use a default value of 1. int64 feature_group_count = 50; + int64 batch_group_count = 58; + // Describes the [begin, end) index range and stride for slices. message SliceDimensions { int64 start = 1; @@ -132,7 +134,7 @@ message HloInstructionProto { string custom_call_opaque = 53; // Shape of outfeed request. - xla.Shape outfeed_shape = 29; + xla.ShapeProto outfeed_shape = 29; // Describes the dimension numbers used for a dot operation xla.DotDimensionNumbers dot_dimension_numbers = 30; @@ -166,7 +168,7 @@ message HloInstructionProto { // Cross replica op fields. repeated ReplicaGroup replica_groups = 49; int64 all_reduce_id = 45; - string cross_replica_sum_barrier = 46; + string all_reduce_barrier = 46; // Whether this Send/Recv instruction transfers data to/from the host. Only // present for Send and Recv instructions and their SendDone and RecvDone @@ -190,7 +192,7 @@ message HloInstructionProto { // 'operand_shapes_with_layout' must contain a shape with layout for each // operand. bool constrain_layout = 56; - repeated Shape operand_shapes_with_layout = 57; + repeated xla.ShapeProto operand_shapes_with_layout = 57; } // Serialization of HloComputation. @@ -205,7 +207,8 @@ message HloComputationProto { repeated HloInstructionProto instructions = 2; // The program shape (with layout) of this computation. - xla.ProgramShape program_shape = 4; + + xla.ProgramShapeProto program_shape = 4; // The id of this computation. int64 id = 5; @@ -226,6 +229,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 @@ -246,11 +261,48 @@ 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; } +message DynamicParameterBindingProto { + // A list of bindings which indicates that the `target_dim_num` in + // the subshape `target_param_index` of parameter `target_param_num` + // is a dynamic dimension and its real dynamic size is represented + // by `dynamic_param_index` in parameter `dynamic_param_num`. + // + // As an example, imagine we have a program: + // + // ENTRY main { + // a = f32[] parameter(0) + // b = f32[10] parameter(1) + // ROOT root = (f32[], f32[10]) tuple(%a, %b) + // } + // + // Let's say 'b' (param index 1) is a dynamic shape whose input has + // an upperbound of 10 and real size is determined at runtime.'a' + // represents the real size of b's first dimension. + // + // In this case, the fields are set in the following way: + // dynamic_param_num = 1 + // dynamic_param_index = {} + // target_param_num = 0 + // target_param_index = {} + // target_param_dim = 0 + message Binding { + int64 dynamic_param_num = 1; + repeated int64 dynamic_param_index = 2; + int64 target_param_num = 3; + repeated int64 target_param_index = 4; + int64 target_param_dim_num = 5; + } + + repeated Binding entries = 1; +} + // Serialization of HloModule. message HloModuleProto { string name = 1; @@ -262,7 +314,7 @@ message HloModuleProto { repeated HloComputationProto computations = 3; // The host program shape (with layout) of the entry computation. - xla.ProgramShape host_program_shape = 4; + xla.ProgramShapeProto host_program_shape = 4; // The id of this module. int64 id = 5; @@ -272,6 +324,8 @@ message HloModuleProto { // Describes alias information between inputs and outputs. HloInputOutputAliasProto input_output_alias = 8; + + DynamicParameterBindingProto dynamic_parameter_binding = 9; } // Serialization of LogicalBuffer. diff --git a/tensorflow/compiler/xla/service/hlo_alias_analysis.cc b/tensorflow/compiler/xla/service/hlo_alias_analysis.cc index cf8e6594cbe5ffd28ca75dd5006e8817f1e8581c..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) { @@ -457,7 +455,7 @@ string HloAliasAnalysis::ToString() const { for (const HloComputation* computation : module_->computations()) { for (const HloInstruction* instruction : computation->instructions()) { StrAppend(&out, " ", instruction->name(), ":\n"); - if (ShapeUtil::IsTuple(instruction->shape())) { + if (instruction->shape().IsTuple()) { ShapeUtil::ForEachSubshape( instruction->shape(), [&out, &instruction, this](const Shape&, const ShapeIndex& index) { @@ -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] = @@ -533,11 +531,11 @@ bool HloAliasAnalysis::HasLiveRangeInterference( const HloOrdering& ordering) const { for (const HloBuffer& buffer : buffers()) { CHECK(!buffer.values().empty()); - if (ShapeUtil::IsToken(buffer.values().front()->shape())) { + if (buffer.values().front()->shape().IsToken()) { // Tokens have no on-device representation and cannot interfere. for (const HloValue* value : buffer.values()) { // If one of the values is a token, all values must be a token. - DCHECK(ShapeUtil::IsToken(value->shape())); + DCHECK(value->shape().IsToken()); } continue; } @@ -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 5c8d97b2d15e15d15cb8014a7d25b37437ce8aec..b6dbf07959c541bceaa8eda5a0101503970ee832 100644 --- a/tensorflow/compiler/xla/service/hlo_alias_analysis_test.cc +++ b/tensorflow/compiler/xla/service/hlo_alias_analysis_test.cc @@ -28,7 +28,7 @@ limitations under the License. #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_verified_test_base.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/logging.h" @@ -39,17 +39,17 @@ namespace { using ::testing::UnorderedElementsAre; -class HloAliasAnalysisTest : public HloVerifiedTestBase { +class HloAliasAnalysisTest : public HloTestBase { protected: - HloAliasAnalysisTest() : HloVerifiedTestBase() { - module_ = CreateNewModule(); + HloAliasAnalysisTest() : HloTestBase() { + module_ = CreateNewVerifiedModule(); } // Run alias analysis on the member module. For convenience returns a // reference to the generated analysis stored in analysis_. HloAliasAnalysis& RunAnalysis() { hlo_graph_dumper::MaybeDumpHloModule(*module_, "Before alias analysis"); - analysis_ = HloAliasAnalysis::Run(module_, + analysis_ = HloAliasAnalysis::Run(module_.get(), /*fusion_can_share_buffer=*/nullptr) .ConsumeValueOrDie(); return *analysis_; @@ -93,7 +93,7 @@ class HloAliasAnalysisTest : public HloVerifiedTestBase { // never occurs, but HLO graphs with interference can be explicitly // constructed. bool AnyValuesInSameBufferInterfere() { - DependencyHloOrdering ordering(module_); + DependencyHloOrdering ordering(module_.get()); for (const HloBuffer& buffer : analysis_->buffers()) { for (const HloValue* value_a : buffer.values()) { for (const HloValue* value_b : buffer.values()) { @@ -110,7 +110,7 @@ class HloAliasAnalysisTest : public HloVerifiedTestBase { return false; } - HloModule* module_; + std::unique_ptr module_; std::unique_ptr analysis_; const Shape scalar_shape_ = ShapeUtil::MakeShape(F32, {}); @@ -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(); @@ -638,7 +646,7 @@ TEST_F(HloAliasAnalysisTest, SequentialWhiles) { module_->AddEntryComputation(builder.Build()); FlattenCallGraph flattener; - TF_ASSERT_OK(flattener.Run(module_).status()); + TF_ASSERT_OK(flattener.Run(module_.get()).status()); const HloAliasAnalysis& analysis = RunAnalysis(); @@ -1012,7 +1020,7 @@ TEST_F(HloAliasAnalysisTest, BitcastInterference) { const HloAliasAnalysis& analysis = RunAnalysis(); - DependencyHloOrdering ordering(module_); + DependencyHloOrdering ordering(module_.get()); EXPECT_FALSE(analysis.HasLiveRangeInterference(ordering)); } @@ -1054,13 +1062,13 @@ TEST_F(HloAliasAnalysisTest, WhileInterference) { { // Dependency ordering should interfere because the negate and while are // unordered. - DependencyHloOrdering ordering(module_); + DependencyHloOrdering ordering(module_.get()); EXPECT_TRUE(analysis.HasLiveRangeInterference(ordering)); } // For a sequential order, if there is interference iff the negate is after // the while. - HloSchedule schedule(module_); + HloSchedule schedule(module_.get()); schedule.set_sequence(body, {body_param, body_root}); schedule.set_sequence(condition, {cond_param, cond_root}); { 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 b0f7cd91ad1db0a59c09cfbfc1885813dc57e01e..98164752b96d552eac4975800e7498c573882a7b 100644 --- a/tensorflow/compiler/xla/service/hlo_computation.cc +++ b/tensorflow/compiler/xla/service/hlo_computation.cc @@ -15,8 +15,8 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_computation.h" -#include #include +#include #include #include #include @@ -33,6 +33,7 @@ limitations under the License. #include "tensorflow/compiler/xla/layout_util.h" #include "tensorflow/compiler/xla/map_util.h" #include "tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" #include "tensorflow/compiler/xla/shape_util.h" @@ -206,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; @@ -321,7 +322,7 @@ void HloComputation::ComputeInstructionPostOrder( // Add the operands to the stack in reverse order so the first operand is // processed first. This will produce a more natural ordering and a nicer - // result for thigns like HLO stringification. + // result for things like HLO stringification. const auto& operands = current->operands(); for (int64 i = operands.size() - 1; i >= 0; --i) { dfs_stack.emplace_back(operands[i]); @@ -331,7 +332,7 @@ void HloComputation::ComputeInstructionPostOrder( dfs_stack.emplace_back(op); } - // Add inputs for send->recv_done dependencies and cross-replica-sum + // Add inputs for send->recv_done dependencies and all-reduce // dependencies. switch (current->opcode()) { case HloOpcode::kRecvDone: { @@ -343,7 +344,7 @@ void HloComputation::ComputeInstructionPostOrder( } break; } - case HloOpcode::kCrossReplicaSum: { + case HloOpcode::kAllReduce: { auto all_reduce_id = current->all_reduce_id(); if (all_reduce_id) { auto it = channel_dependency_map.find(all_reduce_id.value()); @@ -371,7 +372,7 @@ HloComputation::ComputeChannelDependencies() const { instruction.get()); break; } - case HloOpcode::kCrossReplicaSum: { + case HloOpcode::kAllReduce: { auto all_reduce_id = instruction->all_reduce_id(); if (all_reduce_id) { auto& dependencies = channel_dependency_map[all_reduce_id.value()]; @@ -395,6 +396,7 @@ std::vector HloComputation::MakeInstructionPostOrder() const { post_order.reserve(instruction_count()); std::vector trace_instructions; absl::flat_hash_map visited; + visited.reserve(instruction_count()); for (auto& instruction : instructions_) { if (instruction->opcode() == HloOpcode::kTrace) { // Trace instructions aren't handled by the DFS visitor. Add trace @@ -498,7 +500,7 @@ HloComputationProto HloComputation::ToProto() const { proto.add_instructions()->Swap(&instruction_proto); } proto.set_root_id(root_instruction()->unique_id()); - *proto.mutable_program_shape() = ComputeProgramShape(); + *proto.mutable_program_shape() = ComputeProgramShape().ToProto(); return proto; } @@ -529,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); @@ -598,7 +599,7 @@ StatusOr HloComputation::DeepCopyHelper( const std::function< HloInstruction*(HloInstruction* leaf, const ShapeIndex& leaf_index, HloComputation* computation)>& copy_leaf) { - if (ShapeUtil::IsTuple(instruction->shape())) { + if (instruction->shape().IsTuple()) { std::vector elements; for (int64 i = 0; i < ShapeUtil::TupleElementCount(instruction->shape()); i++) { @@ -615,14 +616,14 @@ StatusOr HloComputation::DeepCopyHelper( } return AddInstruction(HloInstruction::CreateTuple(elements)); } - if (ShapeUtil::IsToken(instruction->shape())) { + if (instruction->shape().IsToken()) { // Tokens have no on-device representation and cannot be copied. Pass // through transparently. return instruction; } // Array shape. - TF_RET_CHECK(ShapeUtil::IsArray(instruction->shape())); + TF_RET_CHECK(instruction->shape().IsArray()); return copy_leaf(instruction, *index, this); } @@ -692,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( @@ -739,72 +754,6 @@ Status HloComputation::ReplaceInstruction(HloInstruction* old_instruction, return RemoveInstructionAndUnusedOperands(old_instruction); } -std::unique_ptr HloComputation::ComputeReachability() - const { - const auto& all = MakeInstructionPostOrder(); - auto result = absl::make_unique(all); - auto channel_dependency_map = ComputeChannelDependencies(); - - std::vector inputs; - for (const HloInstruction* hlo : all) { - inputs.assign(hlo->operands().begin(), hlo->operands().end()); - inputs.insert(inputs.end(), hlo->control_predecessors().begin(), - hlo->control_predecessors().end()); - - switch (hlo->opcode()) { - case HloOpcode::kRecvDone: { - auto it = channel_dependency_map.find(hlo->channel_id()); - if (it != channel_dependency_map.end()) { - absl::c_copy(it->second, std::back_inserter(inputs)); - } - break; - } - case HloOpcode::kCrossReplicaSum: { - auto all_reduce_id = hlo->all_reduce_id(); - if (all_reduce_id) { - auto it = channel_dependency_map.find(all_reduce_id.value()); - if (it != channel_dependency_map.end()) { - absl::c_copy(it->second, std::back_inserter(inputs)); - } - } - break; - } - default: - break; - } - - result->FastSetReachabilityToUnion(inputs, hlo); - } - return result; -} - -void HloComputation::UpdateReachabilityThroughInstruction( - const HloInstruction* instruction, HloReachabilityMap* reachability_map) { - std::queue worklist; - worklist.push(instruction); - - std::vector inputs; - - while (!worklist.empty()) { - const HloInstruction* item = worklist.front(); - worklist.pop(); - - inputs.assign(item->operands().begin(), item->operands().end()); - inputs.insert(inputs.end(), item->control_predecessors().begin(), - item->control_predecessors().end()); - - if (reachability_map->SetReachabilityToUnion(inputs, item)) { - // Add immediate successors to worklist. - for (const HloInstruction* user : item->users()) { - worklist.push(user); - } - for (const HloInstruction* succ : item->control_successors()) { - worklist.push(succ); - } - } - } -} - std::vector HloComputation::CollectUnreachableRoots() const { std::vector unreachable_roots; for (auto* instruction : instructions()) { @@ -860,20 +809,19 @@ Status HloComputation::AcceptWithOperandOrder( template Status HloComputation::AcceptOrdered( DfsHloVisitorBase* visitor, - const std::vector& order) const { + 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 = @@ -890,9 +838,9 @@ Status HloComputation::AcceptOrdered( // Explicit instantiations. template Status HloComputation::AcceptOrdered( - DfsHloVisitor*, const std::vector&) const; + DfsHloVisitor*, absl::Span) const; template Status HloComputation::AcceptOrdered( - ConstDfsHloVisitor*, const std::vector&) const; + ConstDfsHloVisitor*, absl::Span) const; Status HloComputation::Accept( const std::function& visitor_func) { @@ -911,14 +859,50 @@ std::unique_ptr HloComputation::Clone( return CloneWithReplacements( /*replacements=*/std::unordered_map>(), - /*extras=*/{}, context, suffix); + /*extra_parameters=*/{}, context, suffix); +} + +std::unique_ptr HloComputation::CloneWithReplacementPairs( + std::pair> r1, + HloCloneContext* context, const string& suffix) { + std::unordered_map> + replacements; + replacements.emplace(std::move(r1)); + 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> + replacements; + replacements.emplace(std::move(r1)); + replacements.emplace(std::move(r2)); + return CloneWithReplacements(std::move(replacements), /*extra_parameters=*/{}, + context, suffix); +} + +std::unique_ptr HloComputation::CloneWithReplacementPairs( + std::pair> r1, + std::pair> r2, + std::pair> r3, + HloCloneContext* context, const string& suffix) { + std::unordered_map> + replacements; + replacements.emplace(std::move(r1)); + replacements.emplace(std::move(r2)); + replacements.emplace(std::move(r3)); + return CloneWithReplacements(std::move(replacements), /*extra_parameters=*/{}, + context, suffix); } std::unique_ptr HloComputation::CloneWithReplacements( std::unordered_map> replacements, - absl::Span extras, HloCloneContext* context, - const string& suffix) { + absl::Span extra_parameters, + HloCloneContext* context, const string& suffix) { std::unique_ptr context_ptr; if (context == nullptr) { context_ptr = absl::make_unique(parent(), suffix); @@ -939,18 +923,56 @@ std::unique_ptr HloComputation::CloneWithReplacements( }; VLOG(1) << "Cloning " << name() << " --> " << suffix << "\n"; + + // We want to do a postorder walk over [replace(i) for i in instructions_]. + // We can't reuse MakeInstructionPostOrder() for this, because that will + // generate a postorder of plain instructions_, and our replacements may + // change the postorder! + // + // The postorder we want here is simpler than what MakeInstructionPostOrder() + // does -- we only care about operand dependencies -- so let's just do it + // ourselves. std::vector postorder; - for (HloInstruction* instr : extras) { - postorder.push_back(instr); - } - for (HloInstruction* instr : MakeInstructionPostOrder()) { - if (HloInstruction* replacement = replace(instr)) { - postorder.push_back(replacement); + absl::flat_hash_map visited; + for (const auto& instr : instructions_) { + std::vector dfs_stack; + HloInstruction* new_instr = replace(instr.get()); + if (!new_instr) { + continue; + } + dfs_stack.push_back(new_instr); + + while (!dfs_stack.empty()) { + auto* cur = dfs_stack.back(); + auto it = visited.find(cur); + if (it != visited.end()) { + dfs_stack.pop_back(); + if (it->second == kVisited) { + continue; + } + CHECK_EQ(it->second, kVisiting); + postorder.push_back(cur); + it->second = kVisited; + continue; + } + + visited.insert({cur, kVisiting}); + for (HloInstruction* operand : cur->operands()) { + HloInstruction* new_operand = replace(operand); + if (new_operand) { + dfs_stack.emplace_back(new_operand); + } + } } } std::vector> instructions; - std::unique_ptr new_instr; + // 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()) { @@ -960,9 +982,8 @@ std::unique_ptr HloComputation::CloneWithReplacements( << operand->ToString() << ", used by " << instr->ToString(); new_operands.push_back(context->GetInstruction(replaced_operand)); } - new_instr = - instr->CloneWithNewOperands(instr->shape(), new_operands, context); - instructions.push_back(std::move(new_instr)); + instructions.push_back( + instr->CloneWithNewOperands(instr->shape(), new_operands, context)); } Builder builder(name() + "." + suffix); for (auto& instr : instructions) { diff --git a/tensorflow/compiler/xla/service/hlo_computation.h b/tensorflow/compiler/xla/service/hlo_computation.h index dec96d11a93cf56d3c40a6bb7882ffb7336aeeb0..e6a1eb89cfdb474f79c184ea0eb77dba8ccd5f03 100644 --- a/tensorflow/compiler/xla/service/hlo_computation.h +++ b/tensorflow/compiler/xla/service/hlo_computation.h @@ -35,7 +35,6 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo.pb.h" #include "tensorflow/compiler/xla/service/hlo_clone_context.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" -#include "tensorflow/compiler/xla/service/hlo_reachability.h" #include "tensorflow/compiler/xla/service/name_uniquer.h" #include "tensorflow/compiler/xla/shape_tree.h" #include "tensorflow/compiler/xla/statusor.h" @@ -215,19 +214,6 @@ class HloComputation { // this order, definitions of values always appear before their uses. std::vector MakeInstructionPostOrder() const; - // Computes and returns the reachability between HLO instructions in the - // computation. The returned HloReachabilityMap is constructed such that - // HloReachabilityMap::IsReachable(a, b) returns true iff there exists a - // directed path (from producer to consumer) from 'a' to 'b'. Both data - // dependencies (operands) and control dependencies are considered for - // reachability. Trivially an instruction is reachable from itself. - std::unique_ptr ComputeReachability() const; - - // Updates the given reachability map after the immediate predecessor set - // (operands and control predecessors) of 'instruction' has changed. - void UpdateReachabilityThroughInstruction( - const HloInstruction* instruction, HloReachabilityMap* reachability_map); - int64 instruction_count() const { return instruction_iterators_.size(); } // Creates and returns a list of the embedded computations called by this @@ -315,7 +301,7 @@ class HloComputation { // be a topological sort of all instructions in the computation. template Status AcceptOrdered(DfsHloVisitorBase* visitor, - const std::vector& order) const; + absl::Span order) const; // Same as Accept() above, but the visitor is given as a function. Status Accept(const std::function& visitor_func); @@ -333,14 +319,42 @@ class HloComputation { // the map's value to replace that instruction in the cloned computation. // // If replacements maps a key to nullptr, we remove that instruction from the - // new computation. - // If additional instructions are used by instructions in replacement map, - // they must be passed in post-order in the extras span. + // new computation. If an element of `replacements` references an instruction + // 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> replacements, - absl::Span extras, HloCloneContext* context = nullptr, - const string& suffix = "clone"); + absl::Span extra_parameters = {}, + HloCloneContext* context = nullptr, const string& suffix = "clone"); + + // Convenience overloads for CloneWithReplacements. You want to do + // + // CloneWithReplacements({{a, std::move(b)}, {c, std::move(d)}}) // ERROR + // + // but that doesn't work because std::initializer_list is not movable. These + // overloads let you do + // + // CloneWithReplacementPairs({a, std::move(b)}, {c, std::move(d)}); // OK + // + std::unique_ptr CloneWithReplacementPairs( + std::pair> r1, + HloCloneContext* context = nullptr, const string& suffix = "clone"); + std::unique_ptr CloneWithReplacementPairs( + std::pair> r1, + std::pair> r2, + HloCloneContext* context = nullptr, const string& suffix = "clone"); + std::unique_ptr CloneWithReplacementPairs( + std::pair> r1, + std::pair> r2, + std::pair> r3, + HloCloneContext* context = nullptr, const string& suffix = "clone"); // Returns true if the given instruction can be removed from the computation. // Parameter instructions cannot be removed without violating invariants of @@ -355,6 +369,14 @@ class HloComputation { // channel complete). bool IsRemovable(const HloInstruction* instruction); + // Returns a map from channel-id to directed dependencies of the channel + // instructions. For send&recv pairs it means the send instruction and for + // all-reduce the union of the dependencies for all participating + // instructions. + using ChannelDependencyMap = + absl::flat_hash_map>; + ChannelDependencyMap ComputeChannelDependencies() const; + // Returns true if this computation has a side effect. A computation has a // side effect if it contains one or more instructions with a side effect. bool HasSideEffect() const; @@ -410,14 +432,6 @@ class HloComputation { // Internal helper to collect unreachable roots. std::vector CollectUnreachableRoots() const; - // Returns a map from channel-id to directed dependencies of the channel - // instructions. For send&recv pairs it means the send instruction and for - // cross-replica-sum the union of the dependencies for all participating - // instructions. - using ChannelDependencyMap = - absl::flat_hash_map>; - ChannelDependencyMap ComputeChannelDependencies() const; - enum VisitState { kVisiting, kVisited }; void ComputeInstructionPostOrder( const HloComputation::ChannelDependencyMap& channel_dependency_map, diff --git a/tensorflow/compiler/xla/service/hlo_computation_test.cc b/tensorflow/compiler/xla/service/hlo_computation_test.cc index 2aaaef1d36d58bcce18db4aa37ff05ea352e484b..35ee532704550a8f91a9778356f8e98e495b5f19 100644 --- a/tensorflow/compiler/xla/service/hlo_computation_test.cc +++ b/tensorflow/compiler/xla/service/hlo_computation_test.cc @@ -15,24 +15,28 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_computation.h" +#include #include +#include +#include +#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" -#include "tensorflow/compiler/xla/service/hlo_matchers.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" +#include "tensorflow/compiler/xla/service/pattern_matcher.h" +#include "tensorflow/compiler/xla/service/pattern_matcher_gmock.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" -namespace op = xla::testing::opcode_matchers; - namespace xla { namespace { +namespace m = match; using ::testing::ElementsAre; using ::testing::UnorderedElementsAre; @@ -65,7 +69,7 @@ class HloComputationTest : public HloTestBase { }; TEST_F(HloComputationTest, GetEmbeddedComputationsEmpty) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto negate_computation = module->AddEntryComputation(CreateNegateComputation()); EXPECT_TRUE(negate_computation->MakeEmbeddedComputationsList().empty()); @@ -73,7 +77,7 @@ TEST_F(HloComputationTest, GetEmbeddedComputationsEmpty) { TEST_F(HloComputationTest, GetEmbeddedComputationsOneComputation) { // Create computation which calls one other computation. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto negate_computation = module->AddEmbeddedComputation(CreateNegateComputation()); auto map_computation = @@ -85,7 +89,7 @@ TEST_F(HloComputationTest, GetEmbeddedComputationsOneComputation) { TEST_F(HloComputationTest, GetEmbeddedComputationsDiamond) { // Create computations with a diamond-shaped callgraph. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto negate_computation = module->AddEmbeddedComputation(CreateNegateComputation()); auto map1_computation = @@ -119,7 +123,7 @@ TEST_F(HloComputationTest, PostOrderSingleton) { auto builder = HloComputation::Builder(TestName()); auto constant = builder.AddInstruction( HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0f))); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_THAT(computation->MakeInstructionPostOrder(), ElementsAre(constant)); } @@ -134,7 +138,7 @@ TEST_F(HloComputationTest, PostOrderSimple) { HloInstruction::CreateUnary(r0f32_, HloOpcode::kNegate, constant)); auto negate2 = builder.AddInstruction( HloInstruction::CreateUnary(r0f32_, HloOpcode::kNegate, negate1)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_THAT(computation->MakeInstructionPostOrder(), ElementsAre(constant, negate1, negate2)); @@ -151,7 +155,7 @@ TEST_F(HloComputationTest, PostOrderTrace) { builder.AddInstruction(HloInstruction::CreateTrace("foobar", negate1)); auto negate2 = builder.AddInstruction( HloInstruction::CreateUnary(r0f32_, HloOpcode::kNegate, negate1)); - auto module = CreateNewModule(); + auto module = CreateNewUnverifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); // Trace instructions should be at the end of the sort. EXPECT_THAT(computation->MakeInstructionPostOrder(), @@ -170,7 +174,7 @@ TEST_F(HloComputationTest, PostOrderDisconnectedInstructions) { HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0f))); auto constant4 = builder.AddInstruction( HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0f))); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_THAT(computation->MakeInstructionPostOrder(), UnorderedElementsAre(constant1, constant2, constant3, constant4)); @@ -192,7 +196,7 @@ TEST_F(HloComputationTest, PostOrderWithMultipleRoots) { r0f32_, HloOpcode::kAdd, constant2, constant3)); auto add3 = builder.AddInstruction(HloInstruction::CreateBinary( r0f32_, HloOpcode::kAdd, constant1, constant3)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); auto post_order = computation->MakeInstructionPostOrder(); EXPECT_EQ(6, post_order.size()); @@ -217,7 +221,7 @@ TEST_F(HloComputationTest, VisitWithMultipleRoots) { constant2, constant3)); builder.AddInstruction(HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, constant1, constant3)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); // Visitor which keeps track of which instructions have been visited. class TestVisitor : public DfsHloVisitorWithDefault { @@ -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; }; @@ -257,11 +261,11 @@ TEST_F(HloComputationTest, DeepCopyArray) { auto builder = HloComputation::Builder(TestName()); auto constant = builder.AddInstruction(HloInstruction::CreateConstant( LiteralUtil::CreateR1({1.0, 2.0, 3.0}))); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); auto copy = computation->DeepCopyInstruction(constant).ValueOrDie(); - EXPECT_THAT(copy, op::Copy(constant)); + EXPECT_THAT(copy, GmockMatch(m::Copy(m::Op().Is(constant)))); } TEST_F(HloComputationTest, DeepCopyTuple) { @@ -274,12 +278,13 @@ TEST_F(HloComputationTest, DeepCopyTuple) { auto tuple = builder.AddInstruction( HloInstruction::CreateTuple({constant1, constant2})); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); auto tuple_copy = computation->DeepCopyInstruction(tuple).ValueOrDie(); - EXPECT_THAT(tuple_copy, op::Tuple(op::Copy(op::GetTupleElement(tuple)), - op::Copy(op::GetTupleElement(tuple)))); + EXPECT_THAT(tuple_copy, GmockMatch(m::Tuple( + m::Copy(m::GetTupleElement(m::Op().Is(tuple))), + m::Copy(m::GetTupleElement(m::Op().Is(tuple)))))); EXPECT_EQ(0, tuple_copy->operand(0)->operand(0)->tuple_index()); EXPECT_EQ(1, tuple_copy->operand(1)->operand(0)->tuple_index()); } @@ -297,7 +302,7 @@ TEST_F(HloComputationTest, DeepCopyArrayAtIndices) { ShapeTree indices_to_copy(constant->shape(), /*init_value=*/true); EXPECT_THAT(computation->DeepCopyInstruction(constant, &indices_to_copy) .ValueOrDie(), - op::Copy(constant)); + GmockMatch(m::Copy(m::Op().Is(constant)))); } { @@ -330,10 +335,11 @@ TEST_F(HloComputationTest, DeepCopyTupleAtIndices) { computation->DeepCopyInstruction(tuple, &indices_to_copy, &copies_added) .ValueOrDie(); - EXPECT_THAT(deep_copy, op::Tuple(op::Copy(op::GetTupleElement(tuple)), - op::Copy(op::GetTupleElement(tuple)))); - EXPECT_THAT(deep_copy, op::Tuple(copies_added.element({0}), - copies_added.element({1}))); + EXPECT_THAT(deep_copy, GmockMatch(m::Tuple( + m::Copy(m::GetTupleElement(m::Op().Is(tuple))) + .Is(copies_added.element({0})), + m::Copy(m::GetTupleElement(m::Op().Is(tuple))) + .Is(copies_added.element({1}))))); } { @@ -346,8 +352,9 @@ TEST_F(HloComputationTest, DeepCopyTupleAtIndices) { computation->DeepCopyInstruction(tuple, &indices_to_copy, &copies_added) .ValueOrDie(); - EXPECT_THAT(deep_copy, op::Tuple(op::GetTupleElement(tuple), - op::GetTupleElement(tuple))); + EXPECT_THAT(deep_copy, + GmockMatch(m::Tuple(m::GetTupleElement(m::Op().Is(tuple)), + m::GetTupleElement(m::Op().Is(tuple))))); EXPECT_TRUE(copies_added.element({}) == nullptr); EXPECT_TRUE(copies_added.element({0}) == nullptr); EXPECT_TRUE(copies_added.element({1}) == nullptr); @@ -363,8 +370,9 @@ TEST_F(HloComputationTest, DeepCopyTupleAtIndices) { computation->DeepCopyInstruction(tuple, &indices_to_copy, &copies_added) .ValueOrDie(); - EXPECT_THAT(deep_copy, op::Tuple(op::Copy(op::GetTupleElement(tuple)), - op::GetTupleElement(tuple))); + EXPECT_THAT(deep_copy, GmockMatch(m::Tuple( + m::Copy(m::GetTupleElement(m::Op().Is(tuple))), + m::GetTupleElement(m::Op().Is(tuple))))); EXPECT_TRUE(copies_added.element({}) == nullptr); EXPECT_TRUE(copies_added.element({0}) != nullptr); EXPECT_TRUE(copies_added.element({1}) == nullptr); @@ -376,12 +384,12 @@ TEST_F(HloComputationTest, DeepCopyToken) { // copied. auto builder = HloComputation::Builder(TestName()); auto token = builder.AddInstruction(HloInstruction::CreateToken()); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); auto copy = computation->DeepCopyInstruction(token).ValueOrDie(); // No copy should be added. - EXPECT_THAT(copy, op::AfterAll()); + EXPECT_THAT(copy, GmockMatch(m::AfterAll())); } TEST_F(HloComputationTest, DeepCopyTokenTuple) { @@ -393,14 +401,15 @@ TEST_F(HloComputationTest, DeepCopyTokenTuple) { HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0))); auto tuple = builder.AddInstruction(HloInstruction::CreateTuple({token, constant})); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); auto copy = computation->DeepCopyInstruction(tuple).ValueOrDie(); // Only the array (second tuple element) should be copied. The token is passed // through transparently. - EXPECT_THAT(copy, op::Tuple(op::GetTupleElement(tuple), - op::Copy(op::GetTupleElement(tuple)))); + EXPECT_THAT(copy, GmockMatch(m::Tuple( + m::GetTupleElement(m::Op().Is(tuple)), + m::Copy(m::GetTupleElement(m::Op().Is(tuple)))))); } TEST_F(HloComputationTest, CycleDetection) { @@ -412,7 +421,7 @@ TEST_F(HloComputationTest, CycleDetection) { HloInstruction::CreateUnary(r0f32_, HloOpcode::kNegate, constant)); auto add = builder.AddInstruction( HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, negate, negate)); - auto module = CreateNewModule(); + auto module = CreateNewUnverifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); // Add a control dependency to create a cycle. ASSERT_IS_OK(add->AddControlDependencyTo(negate)); @@ -440,16 +449,18 @@ TEST_F(HloComputationTest, RemoveInstructionWithDuplicateOperand) { r0f32_, HloOpcode::kAdd, dead_negate, dead_negate)); auto negate = builder.AddInstruction( HloInstruction::CreateUnary(r0f32_, HloOpcode::kNegate, constant)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(4, computation->instruction_count()); - EXPECT_THAT(computation->root_instruction(), op::Negate(constant)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Negate(m::Op().Is(constant)))); EXPECT_EQ(negate, computation->root_instruction()); ASSERT_IS_OK(computation->RemoveInstructionAndUnusedOperands(dead_add)); EXPECT_EQ(2, computation->instruction_count()); - EXPECT_THAT(computation->root_instruction(), op::Negate(constant)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Negate(m::Op().Is(constant)))); EXPECT_EQ(negate, computation->root_instruction()); } @@ -466,7 +477,7 @@ TEST_F(HloComputationTest, CloneWithControlDependency) { HloInstruction::CreateParameter(0, r0f32_, "param0")); auto negate = builder.AddInstruction( HloInstruction::CreateUnary(r0f32_, HloOpcode::kNegate, param)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build(/*root_instruction=*/add)); @@ -484,105 +495,39 @@ TEST_F(HloComputationTest, CloneWithControlDependency) { EXPECT_THAT(successors, ::testing::ElementsAre(cloned_add)); } -TEST_F(HloComputationTest, Reachability) { - // Test reachability of a non-trivial computation: - // - // const1 const2 - // | | - // | +-------+ - // | | | - // add .. negate - // | . | - // | .... exp - // | | - // +---+ +-+---+ - // | | | - // multiply copy - // - // There is a control dependency from 'add' to 'exp'. +TEST_F(HloComputationTest, CloneWithReplacements) { auto builder = HloComputation::Builder(TestName()); - auto constant1 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0f))); - auto constant2 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(2.0f))); - auto add = builder.AddInstruction(HloInstruction::CreateBinary( - r0f32_, HloOpcode::kAdd, constant1, constant2)); - auto negate = builder.AddInstruction( - HloInstruction::CreateUnary(r0f32_, HloOpcode::kNegate, constant2)); - auto exp = builder.AddInstruction( - HloInstruction::CreateUnary(r0f32_, HloOpcode::kExp, negate)); - auto mul = builder.AddInstruction( - HloInstruction::CreateBinary(r0f32_, HloOpcode::kMultiply, add, exp)); - auto copy = builder.AddInstruction( - HloInstruction::CreateUnary(r0f32_, HloOpcode::kCopy, exp)); - - auto module = CreateNewModule(); + 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=*/mul)); - - TF_CHECK_OK(add->AddControlDependencyTo(exp)); - auto reachability = computation->ComputeReachability(); - - EXPECT_TRUE(reachability->IsReachable(constant1, constant1)); - EXPECT_FALSE(reachability->IsReachable(constant1, constant2)); - EXPECT_TRUE(reachability->IsReachable(constant1, add)); - EXPECT_FALSE(reachability->IsReachable(constant1, negate)); - EXPECT_TRUE(reachability->IsReachable(constant1, exp)); - EXPECT_TRUE(reachability->IsReachable(constant1, mul)); - EXPECT_TRUE(reachability->IsReachable(constant1, copy)); - - EXPECT_FALSE(reachability->IsReachable(constant2, constant1)); - EXPECT_TRUE(reachability->IsReachable(constant2, constant2)); - EXPECT_TRUE(reachability->IsReachable(constant2, add)); - EXPECT_TRUE(reachability->IsReachable(constant2, negate)); - EXPECT_TRUE(reachability->IsReachable(constant2, exp)); - EXPECT_TRUE(reachability->IsReachable(constant2, mul)); - EXPECT_TRUE(reachability->IsReachable(constant2, copy)); - - EXPECT_FALSE(reachability->IsReachable(exp, constant1)); - EXPECT_FALSE(reachability->IsReachable(exp, constant2)); - EXPECT_FALSE(reachability->IsReachable(exp, add)); - EXPECT_FALSE(reachability->IsReachable(exp, negate)); - EXPECT_TRUE(reachability->IsReachable(exp, exp)); - EXPECT_TRUE(reachability->IsReachable(exp, mul)); - EXPECT_TRUE(reachability->IsReachable(exp, copy)); - - EXPECT_FALSE(reachability->IsReachable(mul, constant1)); - EXPECT_FALSE(reachability->IsReachable(mul, constant2)); - EXPECT_FALSE(reachability->IsReachable(mul, add)); - EXPECT_FALSE(reachability->IsReachable(mul, negate)); - EXPECT_FALSE(reachability->IsReachable(mul, exp)); - EXPECT_TRUE(reachability->IsReachable(mul, mul)); - EXPECT_FALSE(reachability->IsReachable(mul, copy)); - - EXPECT_TRUE(reachability->IsConnected(constant1, copy)); - EXPECT_TRUE(reachability->IsConnected(copy, constant1)); - EXPECT_FALSE(reachability->IsConnected(negate, add)); - EXPECT_FALSE(reachability->IsConnected(add, negate)); - - // Remove the control dependency then update and verify the reachability map - ASSERT_IS_OK(add->RemoveControlDependencyTo(exp)); - computation->UpdateReachabilityThroughInstruction(exp, reachability.get()); - - EXPECT_TRUE(reachability->IsReachable(constant1, constant1)); - EXPECT_FALSE(reachability->IsReachable(constant1, constant2)); - EXPECT_TRUE(reachability->IsReachable(constant1, add)); - EXPECT_FALSE(reachability->IsReachable(constant1, negate)); - EXPECT_FALSE(reachability->IsReachable(constant1, exp)); - EXPECT_TRUE(reachability->IsReachable(constant1, mul)); - EXPECT_FALSE(reachability->IsReachable(constant1, copy)); - - // Change a use within the graph then update and verify the reachability map - ASSERT_IS_OK(constant2->ReplaceUseWith(negate, constant1)); - computation->UpdateReachabilityThroughInstruction(negate, reachability.get()); - - EXPECT_FALSE(reachability->IsReachable(constant2, constant1)); - EXPECT_TRUE(reachability->IsReachable(constant2, constant2)); - EXPECT_TRUE(reachability->IsReachable(constant2, add)); - EXPECT_FALSE(reachability->IsReachable(constant2, negate)); - EXPECT_FALSE(reachability->IsReachable(constant2, exp)); - EXPECT_TRUE(reachability->IsReachable(constant2, mul)); - EXPECT_FALSE(reachability->IsReachable(constant2, copy)); + module->AddEntryComputation(builder.Build(/*root_instruction=*/lt)); + std::unordered_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) { @@ -606,7 +551,7 @@ TEST_F(HloComputationTest, Stringification) { 2, PrecisionConfig::DEFAULT); builder.AddInstruction( HloInstruction::CreateDot(sout, x, reshape, dot_dnums, precision_config)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto* computation = module->AddEntryComputation(builder.Build()); auto options = HloPrintOptions().set_print_metadata(false); @@ -641,7 +586,7 @@ TEST_F(HloComputationTest, StringificationIndent) { 2, PrecisionConfig::DEFAULT); builder.AddInstruction( HloInstruction::CreateDot(sout, x, reshape, dot_dnums, precision_config)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto* computation = module->AddEntryComputation(builder.Build()); auto options = @@ -677,7 +622,7 @@ TEST_F(HloComputationTest, StringificationCanonical) { 2, PrecisionConfig::DEFAULT); builder.AddInstruction( HloInstruction::CreateDot(sout, x, reshape, dot_dnums, precision_config)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto* computation = module->AddEntryComputation(builder.Build()); auto options = HloPrintOptions().set_print_metadata(false); @@ -700,26 +645,27 @@ TEST_F(HloComputationTest, StringificationCanonical) { EXPECT_EQ(computation->ToString(options), expected_computation2); } -TEST_F(HloComputationTest, ChannelReachability) { - const Shape shape = ShapeUtil::MakeShape(F32, {5, 7}); - HloComputation::Builder builder("ChannelReachability"); - auto param = builder.AddInstruction( - HloInstruction::CreateParameter(0, shape, "param")); - auto token0 = builder.AddInstruction(HloInstruction::CreateToken()); - auto send = - builder.AddInstruction(HloInstruction::CreateSend(param, token0, 1)); - auto send_done = builder.AddInstruction(HloInstruction::CreateSendDone(send)); - auto token1 = builder.AddInstruction(HloInstruction::CreateToken()); - auto recv = - builder.AddInstruction(HloInstruction::CreateRecv(shape, token1, 1)); - auto recv_done = builder.AddInstruction(HloInstruction::CreateRecvDone(recv)); - - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build(recv_done)); - auto reachability = computation->ComputeReachability(); - EXPECT_TRUE(reachability->IsReachable(param, recv_done)); - EXPECT_FALSE(reachability->IsReachable(send, recv)); - EXPECT_FALSE(reachability->IsReachable(send_done, recv)); +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 diff --git a/tensorflow/compiler/xla/service/hlo_constant_folding.cc b/tensorflow/compiler/xla/service/hlo_constant_folding.cc index 4f898ce61c3f36e83e4b13130a404dbb4a2c36c6..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,22 +80,24 @@ StatusOr HloConstantFolding::Run(HloModule* module) { computation->root_instruction() != instruction) { continue; } - // Skip Constant, Parameter, and AfterAll operation. - // TODO(b/64407269): Enable Tuple once the timeout issue is resolved. - // 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. @@ -76,12 +106,23 @@ 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 (ShapeUtil::IsArray(instruction->shape())) { + if (instruction->shape().IsArray()) { int64 elements_in_removed_operands = 0; for (HloInstruction* operand : instruction->operands()) { - if (operand->user_count() == 1 && - ShapeUtil::IsArray(operand->shape())) { + if (operand->user_count() == 1 && operand->shape().IsArray()) { elements_in_removed_operands += ShapeUtil::ElementsIn(operand->shape()); } diff --git a/tensorflow/compiler/xla/service/hlo_constant_folding_test.cc b/tensorflow/compiler/xla/service/hlo_constant_folding_test.cc index e45f905f7152c37a9ab2b41d407310671310c2a3..4bdc980c9ac4fb79cde0242f407aea7057474b27 100644 --- a/tensorflow/compiler/xla/service/hlo_constant_folding_test.cc +++ b/tensorflow/compiler/xla/service/hlo_constant_folding_test.cc @@ -22,22 +22,23 @@ limitations under the License. #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_opcode.h" #include "tensorflow/compiler/xla/service/hlo_parser.h" #include "tensorflow/compiler/xla/service/hlo_pass_fix.h" +#include "tensorflow/compiler/xla/service/pattern_matcher.h" +#include "tensorflow/compiler/xla/service/pattern_matcher_gmock.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/test.h" -#include "tensorflow/compiler/xla/tests/hlo_verified_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/types.h" -namespace op = xla::testing::opcode_matchers; - namespace xla { namespace { -using HloConstantFoldingTest = HloVerifiedTestBase; +namespace m = xla::match; + +using HloConstantFoldingTest = HloTestBase; TEST_F(HloConstantFoldingTest, ConvertF32ToS64) { HloComputation::Builder builder(TestName()); @@ -46,16 +47,17 @@ TEST_F(HloConstantFoldingTest, ConvertF32ToS64) { builder.AddInstruction( HloInstruction::CreateConvert(ShapeUtil::MakeShape(S64, {}), input)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_THAT(computation->root_instruction(), op::Convert(input)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Convert().WithOperand(0, m::Op().Is(input)))); HloConstantFolding const_folder; - TF_ASSERT_OK_AND_ASSIGN(bool result, const_folder.Run(module)); + TF_ASSERT_OK_AND_ASSIGN(bool result, const_folder.Run(module.get())); EXPECT_TRUE(result); - EXPECT_THAT(computation->root_instruction(), op::Constant()); + EXPECT_THAT(computation->root_instruction(), GmockMatch(m::Constant())); EXPECT_EQ(computation->root_instruction()->literal().GetFirstElement(), 42); } @@ -67,16 +69,17 @@ TEST_F(HloConstantFoldingTest, ConvertS64ToF32) { builder.AddInstruction( HloInstruction::CreateConvert(ShapeUtil::MakeShape(F32, {}), input)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_THAT(computation->root_instruction(), op::Convert(input)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Convert().WithOperand(0, m::Op().Is(input)))); HloConstantFolding const_folder; - TF_ASSERT_OK_AND_ASSIGN(bool result, const_folder.Run(module)); + TF_ASSERT_OK_AND_ASSIGN(bool result, const_folder.Run(module.get())); EXPECT_TRUE(result); - EXPECT_THAT(computation->root_instruction(), op::Constant()); + EXPECT_THAT(computation->root_instruction(), GmockMatch(m::Constant())); EXPECT_EQ(computation->root_instruction()->literal().GetFirstElement(), 42.0f); } @@ -88,16 +91,17 @@ TEST_F(HloConstantFoldingTest, ConvertF32ArrayToS64Array) { builder.AddInstruction( HloInstruction::CreateConvert(ShapeUtil::MakeShape(S64, {2}), input)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_THAT(computation->root_instruction(), op::Convert(input)); + EXPECT_THAT(computation->root_instruction(), + GmockMatch(m::Convert().WithOperand(0, m::Op().Is(input)))); HloConstantFolding const_folder; - TF_ASSERT_OK_AND_ASSIGN(bool result, const_folder.Run(module)); + TF_ASSERT_OK_AND_ASSIGN(bool result, const_folder.Run(module.get())); EXPECT_TRUE(result); - EXPECT_THAT(computation->root_instruction(), op::Constant()); + EXPECT_THAT(computation->root_instruction(), GmockMatch(m::Constant())); EXPECT_EQ(computation->root_instruction()->literal().Get({0}), 42); EXPECT_EQ(computation->root_instruction()->literal().Get({1}), 19); } @@ -130,15 +134,15 @@ TEST_F(HloConstantFoldingTest, Concatenate) { Shape shape = ShapeUtil::MakeShape(F32, dimensions); builder.AddInstruction(HloInstruction::CreateConcatenate( shape, operands, test_config.concat_dimension)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); HloConstantFolding const_folder; - TF_ASSERT_OK_AND_ASSIGN(bool result, const_folder.Run(module)); + TF_ASSERT_OK_AND_ASSIGN(bool result, const_folder.Run(module.get())); EXPECT_TRUE(result); HloInstruction* root = computation->root_instruction(); - EXPECT_THAT(root, op::Constant()); + EXPECT_THAT(root, GmockMatch(m::Constant())); EXPECT_TRUE(ShapeUtil::Equal(root->shape(), shape)); } } @@ -157,15 +161,15 @@ TEST_F(HloConstantFoldingTest, Slice) { Shape shape = ShapeUtil::MakeShape(F32, {6, 6, 3, 4, 4}); builder.AddInstruction(HloInstruction::CreateSlice( shape, literal_instruction, slice_start, slice_limits, slice_strides)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); HloConstantFolding const_folder; - TF_ASSERT_OK_AND_ASSIGN(bool result, const_folder.Run(module)); + TF_ASSERT_OK_AND_ASSIGN(bool result, const_folder.Run(module.get())); EXPECT_TRUE(result); HloInstruction* root = computation->root_instruction(); - EXPECT_THAT(root, op::Constant()); + EXPECT_THAT(root, GmockMatch(m::Constant())); EXPECT_TRUE(ShapeUtil::Equal(root->shape(), shape)); } @@ -182,15 +186,15 @@ TEST_F(HloConstantFoldingTest, TransposeConstantFold) { const int64 permutation[] = {1, 2, 0, 4, 3}; builder.AddInstruction( HloInstruction::CreateTranspose(shape, literal_instruction, permutation)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); HloConstantFolding const_folder; - TF_ASSERT_OK_AND_ASSIGN(bool result, const_folder.Run(module)); + TF_ASSERT_OK_AND_ASSIGN(bool result, const_folder.Run(module.get())); EXPECT_TRUE(result); HloInstruction* root = computation->root_instruction(); - EXPECT_THAT(root, op::Constant()); + EXPECT_THAT(root, GmockMatch(m::Constant())); EXPECT_TRUE(ShapeUtil::Compatible(root->shape(), shape)); using NativeT = typename primitive_util::PrimitiveTypeToNative::type; @@ -219,34 +223,36 @@ const char* const kConstantFoldReduce = R"( })"; TEST_F(HloConstantFoldingTest, ConstantFoldReduce) { - ParseAndVerifyModule(kConstantFoldReduce); + TF_ASSERT_OK_AND_ASSIGN(auto m, + ParseAndReturnVerifiedModule(kConstantFoldReduce)); HloConstantFolding const_folder; - TF_ASSERT_OK_AND_ASSIGN(bool result, const_folder.Run(&module())); + TF_ASSERT_OK_AND_ASSIGN(bool result, const_folder.Run(m.get())); EXPECT_TRUE(result); - EXPECT_EQ(6, module() - .entry_computation() + EXPECT_EQ(6, m->entry_computation() ->root_instruction() ->literal() .GetFirstElement()); } TEST_F(HloConstantFoldingTest, ConstantFoldReduceNoLayout) { - ParseAndVerifyModule(kConstantFoldReduce); - HloInstruction* add = module().computations().begin()->root_instruction(); + TF_ASSERT_OK_AND_ASSIGN(auto m, + ParseAndReturnVerifiedModule(kConstantFoldReduce)); + HloInstruction* add = m->computations().begin()->root_instruction(); LayoutUtil::ClearLayout(add->mutable_shape()); HloConstantFolding const_folder; - TF_ASSERT_OK_AND_ASSIGN(bool result, const_folder.Run(&module())); + TF_ASSERT_OK_AND_ASSIGN(bool result, const_folder.Run(m.get())); EXPECT_FALSE(result); - EXPECT_THAT(module().entry_computation()->root_instruction(), op::Reduce()); + EXPECT_THAT(m->entry_computation()->root_instruction(), + GmockMatch(m::Reduce())); } const char* const kConstantFoldLargePad = R"( HloModule ConstantFoldLargePad ENTRY r { - a = f32[1,1,1] constant(f32[1,1,1]{{{7}}}) + a = f32[1,1,1] constant({{{7}}}) b = f32[] constant(42) ROOT pad = f32[2048,2048,128] pad(a, b), padding=1024_1023x1024_1023x64_63 })"; @@ -259,7 +265,53 @@ TEST_F(HloConstantFoldingTest, DoesNotFoldLargePad) { EXPECT_FALSE(result); EXPECT_THAT(module->entry_computation()->root_instruction(), - op::Pad(op::Constant(), op::Constant())); + 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 diff --git a/tensorflow/compiler/xla/service/hlo_cost_analysis.cc b/tensorflow/compiler/xla/service/hlo_cost_analysis.cc index 108aeea097d7170d236b988c414b517a1a284640..76fd402b2c25c8dbed7902a458cd3af44f89cbd1 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_widht FMA operations. + current_properties_[kFlopsKey] = + kFmaFlops * ShapeUtil::ElementsIn(dot_shape) * reduction_width; return Status::OK(); } @@ -269,7 +262,7 @@ Status HloCostAnalysis::HandleOutfeed(const HloInstruction*) { Status HloCostAnalysis::HandleMap(const HloInstruction* map) { // Compute properties of the mapped function. TF_ASSIGN_OR_RETURN(const Properties sub_properties, - ProcessSubcomputation(map->to_apply())); + ProcessNestedSubcomputation(map->to_apply())); // Compute the cost of all elements for this Map operation. const int64 element_count = ShapeUtil::ElementsIn(map->shape()); @@ -285,14 +278,14 @@ Status HloCostAnalysis::HandleReduce(const HloInstruction* reduce) { HloComputation* function = reduce->to_apply(); // Compute the cost of the user function. TF_ASSIGN_OR_RETURN(const Properties sub_properties, - ProcessSubcomputation(function)); + ProcessNestedSubcomputation(function)); // Compute the cost of all elements for this Reduce operation. // This counts the number of times the reduction function is applied, so it // does not need to be multiplied by the number of input tensors - that's // already "priced in" by the sub-computation doing more work. auto arg = reduce->operand(0); - auto output_shape = ShapeUtil::IsArray(reduce->shape()) + auto output_shape = reduce->shape().IsArray() ? reduce->shape() : reduce->shape().tuple_shapes(0); int64 reduction_count = @@ -311,7 +304,7 @@ Status HloCostAnalysis::HandleReduceWindow( auto function = reduce_window->to_apply(); // Compute the properties of the reduction function. TF_ASSIGN_OR_RETURN(const Properties sub_properties, - ProcessSubcomputation(function)); + ProcessNestedSubcomputation(function)); // Compute the cost of all elements for this ReduceWindow operation. For each // output element there are window_size - 1 reductions to perform. @@ -336,9 +329,9 @@ Status HloCostAnalysis::HandleSelectAndScatter( // Compute the properties of the select and scatter function. // Compute the properties of the reduction function. TF_ASSIGN_OR_RETURN(const Properties select_properties, - ProcessSubcomputation(instruction->select())); + ProcessNestedSubcomputation(instruction->select())); TF_ASSIGN_OR_RETURN(const Properties scatter_properties, - ProcessSubcomputation(instruction->scatter())); + ProcessNestedSubcomputation(instruction->scatter())); // Compute the cost of all elements for this operation. For each scatter // source element there are window_size - 1 select computations to perform and @@ -419,6 +412,21 @@ Status HloCostAnalysis::HandleTranspose(const HloInstruction*) { } Status HloCostAnalysis::HandleAfterAll(const HloInstruction*) { + // This instruction is used to enforce ordering at compile time. No code is + // emitted. + current_should_compute_bottleneck_time_ = false; + current_properties_[kBytesAccessedKey] = 0; + current_properties_[kOptimalSecondsKey] = 0; + return Status::OK(); +} + +Status HloCostAnalysis::HandleAddDependency( + const HloInstruction* add_dependency) { + // This instruction is used to enforce ordering at compile time. No code is + // emitted. + current_should_compute_bottleneck_time_ = false; + current_properties_[kBytesAccessedKey] = 0; + current_properties_[kOptimalSecondsKey] = 0; return Status::OK(); } @@ -524,7 +532,7 @@ Status HloCostAnalysis::HandleConvolution(const HloInstruction* convolution) { Status HloCostAnalysis::HandleFft(const HloInstruction* fft) { auto real_shape = - ShapeUtil::IsTuple(fft->operand(0)->shape()) + fft->operand(0)->shape().IsTuple() ? ShapeUtil::GetTupleElementShape(fft->operand(0)->shape(), 0) : fft->operand(0)->shape(); constexpr int kFmaPerComplexMul = 4; @@ -537,7 +545,7 @@ Status HloCostAnalysis::HandleFft(const HloInstruction* fft) { return Status::OK(); } -Status HloCostAnalysis::HandleCrossReplicaSum(const HloInstruction* crs) { +Status HloCostAnalysis::HandleAllReduce(const HloInstruction* crs) { // We assume 2 replicas, so that each output element is the sum of two input // elements. // @@ -546,7 +554,7 @@ Status HloCostAnalysis::HandleCrossReplicaSum(const HloInstruction* crs) { double flops = 0.0; ShapeUtil::ForEachSubshape(crs->shape(), [&](const Shape& subshape, const ShapeIndex&) { - if (ShapeUtil::IsArray(subshape)) { + if (subshape.IsArray()) { flops += ShapeUtil::ElementsIn(subshape); } }); @@ -574,7 +582,7 @@ Status HloCostAnalysis::HandleRng(const HloInstruction* random) { Status HloCostAnalysis::HandleFusion(const HloInstruction* fusion) { TF_ASSIGN_OR_RETURN( current_properties_, - ProcessSubcomputation(fusion->fused_instructions_computation())); + ProcessNestedSubcomputation(fusion->fused_instructions_computation())); // Fusion nodes that produce a tuple also produce the entries in the tuple. // Ignore the memory accessed inside fused ops, since fusion is supposed to @@ -595,7 +603,7 @@ Status HloCostAnalysis::HandleFusion(const HloInstruction* fusion) { Status HloCostAnalysis::HandleCall(const HloInstruction* call) { TF_ASSIGN_OR_RETURN(current_properties_, - ProcessSubcomputation(call->to_apply())); + ProcessUnnestedSubcomputation(call->to_apply())); current_should_compute_bottleneck_time_ = false; return Status::OK(); } @@ -624,13 +632,12 @@ Status HloCostAnalysis::HandleWhile(const HloInstruction* xla_while) { // Since the number of iterations of the while node will not always be // something that we can statically analyze, we cannot precisely compute the // cost of a while node. For now compute the cost of a single iteration. - // - // TODO(b/26346211): Improve the cost analysis for while nodes. TF_ASSIGN_OR_RETURN(const Properties body_properties, - ProcessSubcomputation(xla_while->while_body())); + ProcessUnnestedSubcomputation(xla_while->while_body())); - TF_ASSIGN_OR_RETURN(const Properties condition_properties, - ProcessSubcomputation(xla_while->while_condition())); + TF_ASSIGN_OR_RETURN( + const Properties condition_properties, + ProcessUnnestedSubcomputation(xla_while->while_condition())); current_properties_.clear(); for (const auto& property : body_properties) { @@ -647,10 +654,12 @@ Status HloCostAnalysis::HandleWhile(const HloInstruction* xla_while) { Status HloCostAnalysis::HandleConditional(const HloInstruction* conditional) { // Compute the cost of the true and false computations and take the maximum // from those for each property. - TF_ASSIGN_OR_RETURN(const Properties true_computation_properties, - ProcessSubcomputation(conditional->true_computation())); - TF_ASSIGN_OR_RETURN(const Properties false_computation_properties, - ProcessSubcomputation(conditional->false_computation())); + TF_ASSIGN_OR_RETURN( + const Properties true_computation_properties, + ProcessUnnestedSubcomputation(conditional->true_computation())); + TF_ASSIGN_OR_RETURN( + const Properties false_computation_properties, + ProcessUnnestedSubcomputation(conditional->false_computation())); current_properties_ = true_computation_properties; for (const auto& property : false_computation_properties) { if (!tensorflow::gtl::InsertIfNotPresent(¤t_properties_, property)) { @@ -680,7 +689,7 @@ Status HloCostAnalysis::HandleScatter(const HloInstruction* scatter) { const int64 element_count = ShapeUtil::ElementsIn(scatter->operand(2)->shape()); TF_ASSIGN_OR_RETURN(const Properties sub_properties, - ProcessSubcomputation(scatter->to_apply())); + ProcessNestedSubcomputation(scatter->to_apply())); for (const auto& property : sub_properties) { if (property.first != kBytesAccessedKey) { current_properties_[property.first] = property.second * element_count; @@ -689,6 +698,11 @@ Status HloCostAnalysis::HandleScatter(const HloInstruction* scatter) { return Status::OK(); } +Status HloCostAnalysis::HandleGetDimensionSize( + const HloInstruction* /*get_size*/) { + return Status::OK(); +} + Status HloCostAnalysis::FinishVisit(const HloInstruction*) { return Status::OK(); } @@ -725,10 +739,19 @@ float HloCostAnalysis::optimal_seconds(const HloInstruction& hlo) const { return GetPropertyForHlo(hlo, kOptimalSecondsKey, hlo_properties_); } -StatusOr HloCostAnalysis::ProcessSubcomputation( - HloComputation* computation) { +StatusOr +HloCostAnalysis::ProcessNestedSubcomputation(HloComputation* computation) { + HloCostAnalysis visitor(shape_size_, per_second_rates_); + TF_RETURN_IF_ERROR(computation->Accept(&visitor)); + return visitor.properties(); +} + +StatusOr +HloCostAnalysis::ProcessUnnestedSubcomputation(HloComputation* computation) { HloCostAnalysis visitor(shape_size_, per_second_rates_); TF_RETURN_IF_ERROR(computation->Accept(&visitor)); + hlo_properties_.insert(visitor.hlo_properties_.begin(), + visitor.hlo_properties_.end()); return visitor.properties(); } diff --git a/tensorflow/compiler/xla/service/hlo_cost_analysis.h b/tensorflow/compiler/xla/service/hlo_cost_analysis.h index 46b4bbeef222e6de581360fc01b293e812f1dedd..b52305626dd67336eb31098d086ad357f12d96c7 100644 --- a/tensorflow/compiler/xla/service/hlo_cost_analysis.h +++ b/tensorflow/compiler/xla/service/hlo_cost_analysis.h @@ -71,7 +71,7 @@ class HloCostAnalysis : public ConstDfsHloVisitor { Status HandleDot(const HloInstruction* dot) override; Status HandleConvolution(const HloInstruction* convolution) override; Status HandleFft(const HloInstruction* fft) override; - Status HandleCrossReplicaSum(const HloInstruction* crs) override; + Status HandleAllReduce(const HloInstruction* crs) override; Status HandleAllToAll(const HloInstruction* hlo) override; Status HandleCollectivePermute(const HloInstruction* hlo) override; Status HandleInfeed(const HloInstruction* infeed) override; @@ -101,12 +101,14 @@ class HloCostAnalysis : public ConstDfsHloVisitor { Status HandleBroadcast(const HloInstruction* broadcast) override; Status HandlePad(const HloInstruction* pad) override; Status HandleReshape(const HloInstruction* reshape) override; + Status HandleAddDependency(const HloInstruction* add_dependency) override; Status HandleAfterAll(const HloInstruction* token) override; Status HandleTranspose(const HloInstruction* transpose) override; Status HandleWhile(const HloInstruction* xla_while) override; Status HandleConditional(const HloInstruction* conditional) override; Status HandleGather(const HloInstruction* gather) override; Status HandleScatter(const HloInstruction* scatter) override; + Status HandleGetDimensionSize(const HloInstruction* get_size) override; Status FinishVisit(const HloInstruction* root) override; Status Preprocess(const HloInstruction* hlo) override; @@ -153,7 +155,24 @@ class HloCostAnalysis : public ConstDfsHloVisitor { // Returns the properties computed from visiting the computation rooted at the // given hlo. - StatusOr ProcessSubcomputation(HloComputation* computation); + // + // The difference between ProcessNestedSubcomputation and + // ProcessUnnestedSubcomputation is that we expect to get profile results for + // an unnested subcomputation's individual instructions, while we expect that + // a nested subcomputation is completely subsumed by its parent. + // + // For example, subcomputations inside kFusion and kMap are considered nested, + // while subcomputations inside kWhile and kConditional are considered + // unnested. + // + // Another way of thinking of this is, kFusion is implemented on the GPU + // backend using just one GPU kernel, while kWhile's body is implemented as a + // sequence of kernels, one for each HLO therein. Backends don't necessarily + // need to follow this same implementation strategy, but we assume they do for + // the purposes of this platform-generic cost analysis. + StatusOr ProcessNestedSubcomputation(HloComputation* computation); + StatusOr ProcessUnnestedSubcomputation( + HloComputation* computation); // Utility function to handle all element-wise operations. Status HandleElementwiseOp(const HloInstruction* hlo_instruction); diff --git a/tensorflow/compiler/xla/service/hlo_cost_analysis_test.cc b/tensorflow/compiler/xla/service/hlo_cost_analysis_test.cc index 9acee892d5993be3498d51ed66d7fa4647d7de88..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"); @@ -387,7 +469,7 @@ TEST_F(FusionCostAnalysis, LoopFusion) { HloInstruction::CreateBinary(r2f32, HloOpcode::kSubtract, mul, clamp)); auto tuple = HloInstruction::CreateTuple({sub, sub, mul, c1}); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto* computation = module->AddEntryComputation(builder.Build()); auto* fusion = computation->CreateFusionInstruction( {sub, mul, exp, clamp, add}, HloInstruction::FusionKind::kLoop); @@ -429,7 +511,7 @@ TEST_F(FusionCostAnalysis, NoLayout) { auto add = builder.AddInstruction(HloInstruction::CreateBinary( shape_with_layout, HloOpcode::kAdd, c1, broadcast)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto* computation = module->AddEntryComputation(builder.Build()); auto* fusion = computation->CreateFusionInstruction( {add, broadcast}, HloInstruction::FusionKind::kLoop); @@ -472,7 +554,7 @@ TEST_F(DomainCostAnalysis, DomainCost) { auto domain = builder.AddInstruction( HloInstruction::CreateDomain(tuple->shape(), tuple, nullptr, nullptr)); - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewVerifiedModule(); hlo_module->AddEntryComputation(builder.Build()); EXPECT_EQ(hlo_module->entry_computation()->root_instruction(), domain); @@ -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 b2005d3c210d4ae7e3702cb9624c3ad98056984c..d56f673455f9129b72e9d85eaf8cbf03cfee4302 100644 --- a/tensorflow/compiler/xla/service/hlo_creation_utils.cc +++ b/tensorflow/compiler/xla/service/hlo_creation_utils.cc @@ -19,6 +19,7 @@ limitations under the License. #include "absl/strings/str_cat.h" #include "tensorflow/compiler/xla/literal.h" #include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/service/shape_inference.h" #include "tensorflow/compiler/xla/util.h" @@ -69,11 +70,11 @@ StatusOr MakeConvolveHlo( CHECK_EQ(computation, rhs->parent()); TF_ASSIGN_OR_RETURN(Shape convolve_shape, ShapeInference::InferConvolveShape( - lhs->shape(), rhs->shape(), feature_group_count, + lhs->shape(), rhs->shape(), feature_group_count, 1, window, dimension_numbers)); return computation->AddInstruction(HloInstruction::CreateConvolve( - convolve_shape, lhs, rhs, feature_group_count, window, dimension_numbers, - precision_config)); + convolve_shape, lhs, rhs, feature_group_count, 1, window, + dimension_numbers, precision_config)); } StatusOr MakeTransposeHlo(HloInstruction* operand, @@ -105,12 +106,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 +134,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); @@ -189,8 +218,7 @@ StatusOr MakeMapHlo(absl::Span operands, for (const HloInstruction* operand : operands) { CHECK_EQ(computation, operand->parent()); operand_shapes.push_back(&operand->shape()); - max_operand_rank = - std::max(max_operand_rank, ShapeUtil::Rank(operand->shape())); + max_operand_rank = std::max(max_operand_rank, operand->shape().rank()); } std::vector map_dims(max_operand_rank); std::iota(map_dims.begin(), map_dims.end(), 0); @@ -207,7 +235,7 @@ StatusOr MakeReduceHlo(HloInstruction* operand, HloOpcode binary_opcode, HloModule* module) { DCHECK_NE(nullptr, module); - std::vector all_dims(ShapeUtil::Rank(operand->shape())); + std::vector all_dims(operand->shape().rank()); std::iota(all_dims.begin(), all_dims.end(), 0); auto scalar_shape = ShapeUtil::MakeShape(operand->shape().element_type(), {}); @@ -366,9 +394,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..1c3174e9c89c16cb11589e7c0235bdf13eae6b85 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`. @@ -198,9 +198,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 e07a196d1154dc0ea45ccd2f15b0b9b56f7c41f8..6025e6a77941369f75ebaa98bdf0979669b3a03c 100644 --- a/tensorflow/compiler/xla/service/hlo_creation_utils_test.cc +++ b/tensorflow/compiler/xla/service/hlo_creation_utils_test.cc @@ -19,22 +19,22 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/test.h" -#include "tensorflow/compiler/xla/tests/hlo_verified_test_base.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" #include "tensorflow/core/platform/test.h" namespace xla { namespace { -class HloCreationUtilsTest : public HloVerifiedTestBase { +class HloCreationUtilsTest : public HloTestBase { protected: - HloModule* CreateModuleWithProgramShape( + std::unique_ptr CreateModuleWithProgramShape( PrimitiveType primitive_type, absl::Span input_shape_dims, absl::Span output_shape_dims, HloInstruction** param, HloComputation** entry_computation) { Shape input_shape = ShapeUtil::MakeShape(primitive_type, input_shape_dims); Shape output_shape = ShapeUtil::MakeShape(primitive_type, output_shape_dims); - auto module = CreateNewModule("test"); + auto module = CreateNewVerifiedModule("test"); *entry_computation = module->AddEntryComputation( CreateComputationWithSignature({&input_shape}, output_shape, "entry") .ValueOrDie()); @@ -47,19 +47,18 @@ TEST_F(HloCreationUtilsTest, CollapseFirst1Dim) { HloInstruction* param; HloComputation* entry_computation; - HloModule* module = CreateModuleWithProgramShape(S32, - /*input_shape_dims=*/{2}, - /*output_shape_dims=*/{2}, - ¶m, &entry_computation); + auto module = CreateModuleWithProgramShape(S32, /*input_shape_dims=*/{2}, + /*output_shape_dims=*/{2}, ¶m, + &entry_computation); TF_ASSERT_OK_AND_ASSIGN(HloInstruction * first_1_dims_collapsed, CollapseFirstNDims(param, 1)); entry_computation->set_root_instruction(first_1_dims_collapsed); HloEvaluator evaluator; - TF_ASSERT_OK_AND_ASSIGN(Literal result_literal, - evaluator.Evaluate( - *module, {LiteralUtil::CreateR1({3, 4})})); + TF_ASSERT_OK_AND_ASSIGN( + Literal result_literal, + evaluator.Evaluate(*module, {LiteralUtil::CreateR1({3, 4})})); CHECK_EQ(result_literal, LiteralUtil::CreateR1({3, 4})); } @@ -67,9 +66,8 @@ TEST_F(HloCreationUtilsTest, CollapseFirst2Dims) { HloInstruction* param; HloComputation* entry_computation; - HloModule* module = CreateModuleWithProgramShape( - S32, - /*input_shape_dims=*/{2, 3, 2}, /*output_shape_dims=*/{6, 2}, ¶m, + auto module = CreateModuleWithProgramShape( + S32, /*input_shape_dims=*/{2, 3, 2}, /*output_shape_dims=*/{6, 2}, ¶m, &entry_computation); TF_ASSERT_OK_AND_ASSIGN(HloInstruction * first_2_dims_collapsed, @@ -79,10 +77,9 @@ TEST_F(HloCreationUtilsTest, CollapseFirst2Dims) { HloEvaluator evaluator; TF_ASSERT_OK_AND_ASSIGN( Literal result_literal, - evaluator.Evaluate( - *module, - {LiteralUtil::CreateR3( - {{{1, 2}, {3, 4}, {5, 6}}, {{-1, -2}, {-3, -4}, {-5, -6}}})})); + evaluator.Evaluate(*module, {LiteralUtil::CreateR3( + {{{1, 2}, {3, 4}, {5, 6}}, + {{-1, -2}, {-3, -4}, {-5, -6}}})})); CHECK_EQ(result_literal, LiteralUtil::CreateR2( {{1, 2}, {3, 4}, {5, 6}, {-1, -2}, {-3, -4}, {-5, -6}})); @@ -92,10 +89,9 @@ TEST_F(HloCreationUtilsTest, Prepend1DegenerateDim) { HloInstruction* param; HloComputation* entry_computation; - HloModule* module = CreateModuleWithProgramShape(S32, - /*input_shape_dims=*/{2}, - /*output_shape_dims=*/{1, 2}, - ¶m, &entry_computation); + auto module = CreateModuleWithProgramShape(S32, /*input_shape_dims=*/{2}, + /*output_shape_dims=*/{1, 2}, + ¶m, &entry_computation); TF_ASSERT_OK_AND_ASSIGN(HloInstruction * with_1_degenerate_dim_prepended, PrependDegenerateDims(param, 1)); @@ -104,8 +100,7 @@ TEST_F(HloCreationUtilsTest, Prepend1DegenerateDim) { HloEvaluator evaluator; TF_ASSERT_OK_AND_ASSIGN( Literal result_literal, - evaluator.Evaluate(*module, - {LiteralUtil::CreateR1({9, 10})})); + evaluator.Evaluate(*module, {LiteralUtil::CreateR1({9, 10})})); CHECK_EQ(result_literal, LiteralUtil::CreateR2({{9, 10}})); } @@ -113,10 +108,9 @@ TEST_F(HloCreationUtilsTest, Prepend2DegenerateDims) { HloInstruction* param; HloComputation* entry_computation; - HloModule* module = CreateModuleWithProgramShape( - S32, - /*input_shape_dims=*/{2}, /*output_shape_dims=*/{1, 1, 2}, ¶m, - &entry_computation); + auto module = CreateModuleWithProgramShape(S32, /*input_shape_dims=*/{2}, + /*output_shape_dims=*/{1, 1, 2}, + ¶m, &entry_computation); TF_ASSERT_OK_AND_ASSIGN(HloInstruction * with_2_degenerate_dims_prepended, PrependDegenerateDims(param, 2)); @@ -125,8 +119,7 @@ TEST_F(HloCreationUtilsTest, Prepend2DegenerateDims) { HloEvaluator evaluator; TF_ASSERT_OK_AND_ASSIGN( Literal result_literal, - evaluator.Evaluate(*module, - {LiteralUtil::CreateR1({9, 10})})); + evaluator.Evaluate(*module, {LiteralUtil::CreateR1({9, 10})})); CHECK_EQ(result_literal, LiteralUtil::CreateR3({{{9, 10}}})); } @@ -134,10 +127,9 @@ TEST_F(HloCreationUtilsTest, Prepend2DegenerateDimsToScalar) { HloInstruction* param; HloComputation* entry_computation; - HloModule* module = CreateModuleWithProgramShape(S32, - /*input_shape_dims=*/{}, - /*output_shape_dims=*/{1, 1}, - ¶m, &entry_computation); + auto module = CreateModuleWithProgramShape(S32, /*input_shape_dims=*/{}, + /*output_shape_dims=*/{1, 1}, + ¶m, &entry_computation); TF_ASSERT_OK_AND_ASSIGN(HloInstruction * with_2_degenerate_dims_prepended, PrependDegenerateDims(param, 2)); @@ -146,7 +138,7 @@ TEST_F(HloCreationUtilsTest, Prepend2DegenerateDimsToScalar) { HloEvaluator evaluator; TF_ASSERT_OK_AND_ASSIGN( Literal result_literal, - evaluator.Evaluate(*module, {LiteralUtil::CreateR0(9)})); + evaluator.Evaluate(*module, {LiteralUtil::CreateR0(9)})); CHECK_EQ(result_literal, LiteralUtil::CreateR2({{9}})); } @@ -154,10 +146,9 @@ TEST_F(HloCreationUtilsTest, ExpandFirstDimInto3Dims) { HloInstruction* param; HloComputation* entry_computation; - HloModule* module = CreateModuleWithProgramShape( - S32, - /*input_shape_dims=*/{6}, /*output_shape_dims=*/{3, 1, 2}, ¶m, - &entry_computation); + auto module = CreateModuleWithProgramShape(S32, /*input_shape_dims=*/{6}, + /*output_shape_dims=*/{3, 1, 2}, + ¶m, &entry_computation); TF_ASSERT_OK_AND_ASSIGN(HloInstruction * first_dim_expanded, ExpandFirstDimIntoNDims(param, {3, 1, 2})); @@ -166,8 +157,8 @@ TEST_F(HloCreationUtilsTest, ExpandFirstDimInto3Dims) { HloEvaluator evaluator; TF_ASSERT_OK_AND_ASSIGN( Literal result_literal, - evaluator.Evaluate( - *module, {LiteralUtil::CreateR1({1, 2, 3, 4, 5, 6})})); + evaluator.Evaluate(*module, + {LiteralUtil::CreateR1({1, 2, 3, 4, 5, 6})})); CHECK_EQ(result_literal, LiteralUtil::CreateR3({{{1, 2}}, {{3, 4}}, {{5, 6}}})); } @@ -176,10 +167,9 @@ TEST_F(HloCreationUtilsTest, PadVectorWithZeros) { HloInstruction* param; HloComputation* entry_computation; - HloModule* module = CreateModuleWithProgramShape(S32, - /*input_shape_dims=*/{2}, - /*output_shape_dims=*/{6}, - ¶m, &entry_computation); + auto module = CreateModuleWithProgramShape(S32, /*input_shape_dims=*/{2}, + /*output_shape_dims=*/{6}, ¶m, + &entry_computation); TF_ASSERT_OK_AND_ASSIGN( HloInstruction * zero_padded_param, @@ -187,9 +177,9 @@ TEST_F(HloCreationUtilsTest, PadVectorWithZeros) { entry_computation->set_root_instruction(zero_padded_param); HloEvaluator evaluator; - TF_ASSERT_OK_AND_ASSIGN(Literal result_literal, - evaluator.Evaluate( - *module, {LiteralUtil::CreateR1({3, 4})})); + TF_ASSERT_OK_AND_ASSIGN( + Literal result_literal, + evaluator.Evaluate(*module, {LiteralUtil::CreateR1({3, 4})})); CHECK_EQ(result_literal, LiteralUtil::CreateR1({0, 0, 0, 3, 4, 0})); } @@ -197,20 +187,18 @@ TEST_F(HloCreationUtilsTest, BroadcastZeros_S32) { HloInstruction* param; HloComputation* entry_computation; - HloModule* module = CreateModuleWithProgramShape(S32, - /*input_shape_dims=*/{}, - /*output_shape_dims=*/{2, 2}, - ¶m, &entry_computation); + auto module = CreateModuleWithProgramShape(S32, /*input_shape_dims=*/{}, + /*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; TF_ASSERT_OK_AND_ASSIGN( Literal result_literal, - evaluator.Evaluate(*module, {LiteralUtil::CreateR0(0)})); + evaluator.Evaluate(*module, {LiteralUtil::CreateR0(0)})); CHECK_EQ(result_literal, LiteralUtil::CreateR2({{0, 0}, {0, 0}})); } @@ -218,20 +206,18 @@ TEST_F(HloCreationUtilsTest, BroadcastZeros_F32) { HloInstruction* param; HloComputation* entry_computation; - HloModule* module = CreateModuleWithProgramShape(F32, - /*input_shape_dims=*/{}, - /*output_shape_dims=*/{2, 2}, - ¶m, &entry_computation); + auto module = CreateModuleWithProgramShape(F32, /*input_shape_dims=*/{}, + /*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; - TF_ASSERT_OK_AND_ASSIGN(Literal result_literal, - evaluator.Evaluate( - *module, {LiteralUtil::CreateR0(0.0f)})); + TF_ASSERT_OK_AND_ASSIGN( + Literal result_literal, + evaluator.Evaluate(*module, {LiteralUtil::CreateR0(0.0f)})); CHECK_EQ(result_literal, LiteralUtil::CreateR2({{0.0f, 0.0f}, {0.0f, 0.0f}})); } diff --git a/tensorflow/compiler/xla/service/hlo_cse_test.cc b/tensorflow/compiler/xla/service/hlo_cse_test.cc index 9b18b0284f63c25934c1b7118dc8973caa62cadc..1eb0260468c4560985027947e89c62cc21139e7e 100644 --- a/tensorflow/compiler/xla/service/hlo_cse_test.cc +++ b/tensorflow/compiler/xla/service/hlo_cse_test.cc @@ -29,7 +29,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" #include "tensorflow/compiler/xla/shape_util.h" -#include "tensorflow/compiler/xla/tests/hlo_verified_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_utils.h" #include "tensorflow/compiler/xla/util.h" @@ -44,7 +44,7 @@ namespace op = xla::testing::opcode_matchers; namespace xla { namespace { -class HloCseTest : public HloVerifiedTestBase { +class HloCseTest : public HloTestBase { protected: HloCseTest() {} }; @@ -59,13 +59,13 @@ TEST_F(HloCseTest, CombineTwoConstants) { builder.AddInstruction(HloInstruction::CreateBinary( constant1->shape(), HloOpcode::kAdd, constant1, constant2)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(3, computation->instruction_count()); HloCSE cse(/*is_layout_sensitive=*/false); - EXPECT_TRUE(cse.Run(module).ValueOrDie()); + EXPECT_TRUE(cse.Run(module.get()).ValueOrDie()); EXPECT_EQ(2, computation->instruction_count()); HloInstruction* constant = *computation->instructions().begin(); @@ -89,14 +89,14 @@ TEST_F(HloCseTest, CombineTwoConstantsDifferentLayoutsAndInsensitive) { auto add = builder.AddInstruction(HloInstruction::CreateBinary( constant1->shape(), HloOpcode::kAdd, constant1, constant2)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(3, computation->instruction_count()); EXPECT_THAT(add, op::Add(constant1, constant2)); HloCSE cse(/*is_layout_sensitive=*/false); - EXPECT_TRUE(cse.Run(module).ValueOrDie()); + EXPECT_TRUE(cse.Run(module.get()).ValueOrDie()); EXPECT_EQ(2, computation->instruction_count()); auto first_operand = add->operand(0); @@ -121,14 +121,14 @@ TEST_F(HloCseTest, CombineTwoConstantsDifferentLayoutsAndSensitive) { auto add = builder.AddInstruction(HloInstruction::CreateBinary( constant1->shape(), HloOpcode::kAdd, constant1, constant2)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(3, computation->instruction_count()); EXPECT_THAT(add, op::Add(constant1, constant2)); HloCSE cse(/*is_layout_sensitive=*/true); - EXPECT_FALSE(cse.Run(module).ValueOrDie()); + EXPECT_FALSE(cse.Run(module.get()).ValueOrDie()); EXPECT_EQ(3, computation->instruction_count()); EXPECT_THAT(add, op::Add(constant1, constant2)); @@ -171,13 +171,13 @@ TEST_F(HloCseTest, ConstantsSameValueDifferentType) { shape_r0, HloOpcode::kAdd, root, constants[i])); } - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(20, computation->instruction_count()); HloCSE cse(/*is_layout_sensitive=*/false); - EXPECT_TRUE(cse.Run(module).ValueOrDie()); + EXPECT_TRUE(cse.Run(module.get()).ValueOrDie()); // CSE will remove both the second float(42.0f) and the corresponding // convert/cast. @@ -201,7 +201,7 @@ TEST_F(HloCseTest, NonscalarConstants) { auto tuple = builder.AddInstruction(HloInstruction::CreateTuple( {common_constant1, common_constant2, uncommon_constant})); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(4, computation->instruction_count()); @@ -209,7 +209,7 @@ TEST_F(HloCseTest, NonscalarConstants) { op::Tuple(common_constant1, common_constant2, uncommon_constant)); HloCSE cse(/*is_layout_sensitive=*/false); - EXPECT_TRUE(cse.Run(module).ValueOrDie()); + EXPECT_TRUE(cse.Run(module.get()).ValueOrDie()); EXPECT_EQ(3, computation->instruction_count()); auto first_operand = tuple->operand(0); @@ -233,14 +233,14 @@ TEST_F(HloCseTest, IdenticalInstructions) { auto tuple = builder.AddInstruction(HloInstruction::CreateTuple({exp1, exp2, exp3})); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(5, computation->instruction_count()); EXPECT_THAT(tuple, op::Tuple(exp1, exp2, exp3)); HloCSE cse(/*is_layout_sensitive=*/true); - EXPECT_TRUE(cse.Run(module).ValueOrDie()); + EXPECT_TRUE(cse.Run(module.get()).ValueOrDie()); EXPECT_EQ(3, computation->instruction_count()); auto first_operand = tuple->operand(0); @@ -250,7 +250,7 @@ TEST_F(HloCseTest, IdenticalInstructions) { // Test two identical while loops with same inputs TEST_F(HloCseTest, WhileLoopsIdenticalConditionsAndBodiesSameInput) { - ParseAndVerifyModule(R"( + const char* const hlo_string = R"( HloModule WhileLoopsIdenticalConditionsAndBodiesSameInput %body (param: (f32[], f32[])) -> (f32[], f32[]) { @@ -277,21 +277,21 @@ index=1 %add = f32[] add(f32[] %get-tuple-element, f32[] %get-tuple-element.1) f32[]) while((f32[], f32[]) %tuple.1), condition=%condition, body=%body ROOT %while.1 = (f32[], f32[]) while((f32[], f32[]) %tuple.1), condition=%condition.1, body=%body - } - )"); + })"; - auto computation = module().entry_computation(); + TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(hlo_string)); + auto computation = m->entry_computation(); EXPECT_EQ(5, computation->instruction_count()); HloCSE cse(true); - EXPECT_TRUE(cse.Run(&module()).ValueOrDie()); + EXPECT_TRUE(cse.Run(m.get()).ValueOrDie()); EXPECT_EQ(4, computation->instruction_count()); } // Test two while loops with same conditions, same inputs, but different // bodies TEST_F(HloCseTest, WhileLoopsIdenticalConditionsSameInputAndDifferentBodies) { - ParseAndVerifyModule(R"( + const char* const hlo_string = R"( HloModule WhileLoopsIdenticalConditionsSameInputAndDifferentBodies %body (param: (f32[], f32[])) -> (f32[], f32[]) { @@ -327,20 +327,20 @@ index=1 %sub = f32[] subtract(f32[] %get-tuple-element.2, f32[] %while = (f32[], f32[]) while((f32[], f32[]) %tuple.1), condition=%condition, body=%body ROOT %while.1 = (f32[], f32[]) while((f32[], f32[]) %tuple.1), condition=%condition.1, body=%body2 - } - )"); + })"; - auto computation = module().entry_computation(); + TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(hlo_string)); + auto computation = m->entry_computation(); EXPECT_EQ(5, computation->instruction_count()); HloCSE cse(true); - EXPECT_FALSE(cse.Run(&module()).ValueOrDie()); + EXPECT_FALSE(cse.Run(m.get()).ValueOrDie()); EXPECT_EQ(5, computation->instruction_count()); } // Test two identical while loops with different inputs TEST_F(HloCseTest, WhileLoopsIdenticalConditionsAndBodiesDifferentInput) { - ParseAndVerifyModule(R"( + const char* const hlo_string = R"( HloModule WhileLoopsIdenticalConditionsAndBodiesDifferentInput %body (param: (f32[], f32[])) -> (f32[], f32[]) { @@ -369,22 +369,21 @@ condition=%condition, body=%body %constant.4 = f32[] constant(1) %constant.5 = f32[] constant(2) %tuple.2 = (f32[], f32[]) tuple(f32[] %constant.4, f32[] %constant.5) ROOT %while.1 = (f32[], f32[]) while((f32[], f32[]) %tuple.2), condition=%condition.1, body=%body - } - - )"); + })"; - auto computation = module().entry_computation(); + TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(hlo_string)); + auto computation = m->entry_computation(); EXPECT_EQ(8, computation->instruction_count()); HloCSE cse(true); - EXPECT_FALSE(cse.Run(&module()).ValueOrDie()); + EXPECT_FALSE(cse.Run(m.get()).ValueOrDie()); EXPECT_EQ(8, computation->instruction_count()); } // Test two while loops with identical bodies and same inputs, but different // conditions TEST_F(HloCseTest, WhileLoopsIdenticalBodiesAndInputDifferntConditions) { - ParseAndVerifyModule(R"( + const char* const hlo_string = R"( HloModule WhileLoopsIdenticalBodiesAndInputDifferntConditions %body (param: (f32[], f32[])) -> (f32[], f32[]) { @@ -411,13 +410,14 @@ f32[]) { %constant.2 = f32[] constant(1) %constant.3 = f32[] constant(2) %while = (f32[], f32[]) while((f32[], f32[]) %tuple.1), condition=%condition, body=%body ROOT %while.1 = (f32[], f32[]) while((f32[], f32[]) %tuple.1), condition=%condition.1, body=%body - })"); + })"; - auto computation = module().entry_computation(); + TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(hlo_string)); + auto computation = m->entry_computation(); EXPECT_EQ(5, computation->instruction_count()); HloCSE cse(true); - EXPECT_FALSE(cse.Run(&module()).ValueOrDie()); + EXPECT_FALSE(cse.Run(m.get()).ValueOrDie()); EXPECT_EQ(5, computation->instruction_count()); } @@ -439,14 +439,14 @@ TEST_F(HloCseTest, IdenticalInstructionsDifferentLayoutsSensitive) { auto tuple = builder.AddInstruction(HloInstruction::CreateTuple({exp1, exp2})); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(4, computation->instruction_count()); EXPECT_THAT(tuple, op::Tuple(exp1, exp2)); HloCSE cse(/*is_layout_sensitive=*/true); - EXPECT_FALSE(cse.Run(module).ValueOrDie()); + EXPECT_FALSE(cse.Run(module.get()).ValueOrDie()); EXPECT_EQ(4, computation->instruction_count()); EXPECT_THAT(tuple, op::Tuple(exp1, exp2)); @@ -470,14 +470,14 @@ TEST_F(HloCseTest, IdenticalInstructionsDifferentLayoutsInsensitive) { auto tuple = builder.AddInstruction(HloInstruction::CreateTuple({exp1, exp2})); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(4, computation->instruction_count()); EXPECT_THAT(tuple, op::Tuple(exp1, exp2)); HloCSE cse(/*is_layout_sensitive=*/false); - EXPECT_TRUE(cse.Run(module).ValueOrDie()); + EXPECT_TRUE(cse.Run(module.get()).ValueOrDie()); EXPECT_EQ(3, computation->instruction_count()); auto first_operand = tuple->operand(0); @@ -488,7 +488,7 @@ TEST_F(HloCseTest, IdenticalInstructionsDifferentLayoutsInsensitive) { TEST_F(HloCseTest, FusionInternalCSE) { // Test that we can CSE expressions that live within a fusion node // computation. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto builder = HloComputation::Builder(TestName()); const Shape shape_r0 = ShapeUtil::MakeShape(F32, {}); @@ -512,7 +512,7 @@ TEST_F(HloCseTest, FusionInternalCSE) { EXPECT_EQ(5, fused_computation->instruction_count()); HloCSE cse(/*is_layout_sensitive=*/false); - EXPECT_TRUE(cse.Run(module).ValueOrDie()); + EXPECT_TRUE(cse.Run(module.get()).ValueOrDie()); EXPECT_EQ(4, fused_computation->instruction_count()); auto root = fused_computation->root_instruction(); @@ -554,14 +554,14 @@ TEST_F(HloCseTest, IdenticalExpressions) { auto tuple = builder.AddInstruction(HloInstruction::CreateTuple({add1, add2})); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(8, computation->instruction_count()); EXPECT_THAT(tuple, op::Tuple(op::Add(negate1, exp1), op::Add(negate2, exp2))); HloCSE cse(/*is_layout_sensitive=*/false); - EXPECT_TRUE(cse.Run(module).ValueOrDie()); + EXPECT_TRUE(cse.Run(module.get()).ValueOrDie()); EXPECT_EQ(5, computation->instruction_count()); auto operand = tuple->operand(0); @@ -586,7 +586,7 @@ TEST_F(HloCseTest, DoNotCombineRng) { builder.AddInstruction(HloInstruction::CreateBinary( constant1->shape(), HloOpcode::kAdd, rng1, rng2)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); @@ -595,7 +595,7 @@ TEST_F(HloCseTest, DoNotCombineRng) { uint32 count_before = computation->instruction_count(); HloCSE cse(/*is_layout_sensitive=*/false); - EXPECT_FALSE(cse.Run(module).ValueOrDie()); + EXPECT_FALSE(cse.Run(module.get()).ValueOrDie()); uint32 count_after = computation->instruction_count(); EXPECT_EQ(count_before, count_after); @@ -607,7 +607,7 @@ TEST_F(HloCseTest, DoNotCombineCallsToImpureFunctions) { // Test that two calls to an impure function are not commoned. RNG // is the source of the impurity. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); // rng_function is an impure function because it does RNG. HloComputation* rng_function = nullptr; @@ -649,7 +649,7 @@ TEST_F(HloCseTest, DoNotCombineCallsToImpureFunctions) { VLOG(3) << "before: " << module->ToString(); HloCSE cse(/*is_layout_sensitive=*/false); - EXPECT_FALSE(cse.Run(module).ValueOrDie()); + EXPECT_FALSE(cse.Run(module.get()).ValueOrDie()); VLOG(3) << "after: " << module->ToString(); @@ -659,7 +659,7 @@ TEST_F(HloCseTest, DoNotCombineCallsToImpureFunctions) { } TEST_F(HloCseTest, CompareComputations) { - ParseAndVerifyModule(R"( + const char* const hlo_string = R"( HloModule m add_computation { @@ -680,11 +680,12 @@ TEST_F(HloCseTest, CompareComputations) { r1 = f32[] reduce(p, c), dimensions={0}, to_apply=add_computation r2 = f32[] reduce(p, c), dimensions={0}, to_apply=add_computation2 ROOT f2 = (f32[],f32[]) tuple(r1, r2) - })"); + })"; + TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(hlo_string)); HloCSE cse(/*is_layout_sensitive=*/false); - EXPECT_TRUE(cse.Run(&module()).ValueOrDie()); - HloInstruction* root = module().entry_computation()->root_instruction(); + EXPECT_TRUE(cse.Run(m.get()).ValueOrDie()); + HloInstruction* root = m->entry_computation()->root_instruction(); EXPECT_EQ(root->operand(0), root->operand(1)); } @@ -697,19 +698,19 @@ TEST_F(HloCseTest, ConstantsSameValueInDifferentDomains) { builder.AddInstruction( HloInstruction::CreateConstant(LiteralUtil::CreateR0(42))); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(2, computation->instruction_count()); HloCSE cse(/*is_layout_sensitive=*/false); - EXPECT_FALSE(cse.Run(module).ValueOrDie()); + EXPECT_FALSE(cse.Run(module.get()).ValueOrDie()); EXPECT_EQ(2, computation->instruction_count()); } TEST_F(HloCseTest, Domain) { - ParseAndVerifyModule(R"( + const char* const hlo_string = R"( HloModule module ENTRY %entry { %param = f32[] parameter(0), sharding={maximal device=0} @@ -730,11 +731,12 @@ ENTRY %entry { domain={kind="sharding", entry={maximal device=2}, exit={maximal device=0}} %add = f32[] add(%domain.3, %domain.4) ROOT %sub = f32[] subtract(%add, %domain.5) -})"); +})"; + TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(hlo_string)); HloCSE cse(/*is_layout_sensitive=*/false); - EXPECT_TRUE(cse.Run(&module()).ValueOrDie()); - const HloInstruction* sub = module().entry_computation()->root_instruction(); + EXPECT_TRUE(cse.Run(m.get()).ValueOrDie()); + const HloInstruction* sub = m->entry_computation()->root_instruction(); const HloInstruction* add = sub->operand(0); EXPECT_EQ(add->operand(0), add->operand(1)); EXPECT_NE(add->operand(0), sub->operand(1)); diff --git a/tensorflow/compiler/xla/service/hlo_dataflow_analysis.cc b/tensorflow/compiler/xla/service/hlo_dataflow_analysis.cc index 5dcf6bc985ff18fa6fc1ab5a5692914b4597d065..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); } } @@ -190,7 +190,7 @@ string HloDataflowAnalysis::ToString() const { for (const HloComputation* computation : module_.computations()) { for (const HloInstruction* instruction : computation->instructions()) { StrAppend(&out, " ", instruction->name(), ":\n"); - if (ShapeUtil::IsTuple(instruction->shape())) { + if (instruction->shape().IsTuple()) { GetInstructionValueSet(instruction) .ForEachElement([this, &instruction, &out]( const ShapeIndex& index, @@ -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); } @@ -466,6 +465,21 @@ bool HloDataflowAnalysis::UpdateDomainValueSet(HloInstruction* domain) { return changed; } +bool HloDataflowAnalysis::UpdateAddDependencyValueSet( + HloInstruction* add_dependency) { + // AddDependency just forwards the value of its zero-th operand. + CHECK_EQ(add_dependency->opcode(), HloOpcode::kAddDependency); + const InstructionValueSet& operand_set = + GetInstructionValueSet(add_dependency->operand(0)); + InstructionValueSet& add_dependency_set = + GetInstructionValueSet(add_dependency); + if (operand_set != add_dependency_set) { + add_dependency_set = operand_set; + return true; + } + return false; +} + bool HloDataflowAnalysis::UpdateGetTupleElementValueSet(HloInstruction* gte) { CHECK_EQ(gte->opcode(), HloOpcode::kGetTupleElement); bool changed = false; @@ -622,6 +636,8 @@ bool HloDataflowAnalysis::UpdateInstructionValueSet( HloInstruction* instruction) { // Recompute from operands. switch (instruction->opcode()) { + case HloOpcode::kAddDependency: + return UpdateAddDependencyValueSet(instruction); case HloOpcode::kBitcast: return UpdateBitcastValueSet(instruction); case HloOpcode::kDomain: @@ -795,6 +811,7 @@ Status HloDataflowAnalysis::InitializeInstructionValueSets() { define_all_values(); } break; + case HloOpcode::kAddDependency: case HloOpcode::kWhile: case HloOpcode::kCall: case HloOpcode::kConditional: @@ -903,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()); @@ -919,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(); } @@ -936,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(); } @@ -1023,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; } @@ -1082,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.h b/tensorflow/compiler/xla/service/hlo_dataflow_analysis.h index abac398c04fc4c418d8814a0097db4434bc1cd9c..ece17fc4c3ea0261474df5d53c088dd05016e1e4 100644 --- a/tensorflow/compiler/xla/service/hlo_dataflow_analysis.h +++ b/tensorflow/compiler/xla/service/hlo_dataflow_analysis.h @@ -193,6 +193,7 @@ class HloDataflowAnalysis { bool UpdateSendValueSet(HloInstruction* send); bool UpdateTupleValueSet(HloInstruction* tuple); bool UpdateWhileValueSet(HloInstruction* xla_while); + bool UpdateAddDependencyValueSet(HloInstruction* add_dependency); // Propagate the dataflow through the module. void Propagate(); diff --git a/tensorflow/compiler/xla/service/hlo_dataflow_analysis_test.cc b/tensorflow/compiler/xla/service/hlo_dataflow_analysis_test.cc index 909853106d57d181e85e3e4134b4039be2b176f5..4a7c4963b7b399e625da907b3810c42df7ee2bd3 100644 --- a/tensorflow/compiler/xla/service/hlo_dataflow_analysis_test.cc +++ b/tensorflow/compiler/xla/service/hlo_dataflow_analysis_test.cc @@ -43,7 +43,7 @@ using ::testing::UnorderedElementsAre; class HloDataflowAnalysisTest : public HloTestBase, public ::testing::WithParamInterface { protected: - HloDataflowAnalysisTest() : module_(CreateNewModule()) {} + HloDataflowAnalysisTest() : module_(CreateNewVerifiedModule()) {} // Run dataflow analysis on the member module. For convenience returns a // reference to the generated analysis stored in analysis_. @@ -73,8 +73,8 @@ class HloDataflowAnalysisTest : public HloTestBase, bool InstructionsMayInterfere(const HloOrdering& ordering, const HloInstruction* a, const HloInstruction* b) { - EXPECT_FALSE(ShapeUtil::IsTuple(a->shape())); - EXPECT_FALSE(ShapeUtil::IsTuple(b->shape())); + EXPECT_FALSE(a->shape().IsTuple()); + EXPECT_FALSE(b->shape().IsTuple()); return ordering.MayInterfere(analysis_->GetValueDefinedAt(a), analysis_->GetValueDefinedAt(b), *analysis_); } @@ -1877,14 +1877,38 @@ TEST_P(HloDataflowAnalysisTest, NestedConditionals) { } } -INSTANTIATE_TEST_CASE_P(HloDataflowAnalysisInstantiation, - HloDataflowAnalysisTest, - ::testing::Values(false, true)); +TEST_P(HloDataflowAnalysisTest, AddDependency) { + string module_string = R"( +HloModule AddDependency +ENTRY %AddDependency (p: f32[3]) -> f32[3] { + %p = f32[3] parameter(0) + %token0 = token[] after-all() + ROOT %add_dep = f32[3] add-dependency(f32[3] %p, token[] %token0) +} +)"; + TF_ASSERT_OK_AND_ASSIGN( + std::unique_ptr module, + ParseHloString(module_string, GetModuleConfigForTest())); + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr analysis, + HloDataflowAnalysis::Run(*module)); + const HloInstruction* root = module->entry_computation()->root_instruction(); + EXPECT_EQ(root->opcode(), HloOpcode::kAddDependency); + + // The after-all and parameter should define a value. Add-dependency should + // not. + EXPECT_EQ(analysis->values().size(), 2); + EXPECT_FALSE(analysis->ValueIsDefinedAt(root)); +} + +INSTANTIATE_TEST_SUITE_P(HloDataflowAnalysisInstantiation, + HloDataflowAnalysisTest, + ::testing::Values(false, true)); class HloDataflowAnalysisTestBase : public HloTestBase { protected: void BuildModule(std::unique_ptr computation) { - module_ = CreateNewModule(); + module_ = CreateNewUnverifiedModule(); computation_ = module_->AddEntryComputation(std::move(computation)); } @@ -1946,12 +1970,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})); @@ -1988,12 +2013,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})); @@ -2126,17 +2152,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( @@ -2195,12 +2221,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})); @@ -2235,12 +2262,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)); @@ -2266,10 +2294,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()); @@ -2280,7 +2311,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) { @@ -2476,7 +2509,7 @@ TEST_F(CanShareOperandBufferWithUserTest, WhileCanShare) { return builder.Build(); }; - module_ = CreateNewModule(); + module_ = CreateNewUnverifiedModule(); HloComputation* cond_computation = module_->AddEmbeddedComputation(make_cond()); HloComputation* body_computation = @@ -2511,7 +2544,7 @@ TEST_F(CanShareOperandBufferWithUserTest, CallToComputationWithFusionRoot) { auto add = sub_builder.AddInstruction( HloInstruction::CreateBinary(shape, HloOpcode::kAdd, sub_param, ones)); - module_ = CreateNewModule(); + module_ = CreateNewUnverifiedModule(); auto sub_computation = module_->AddEmbeddedComputation(sub_builder.Build()); sub_computation->CreateFusionInstruction({add, ones}, HloInstruction::FusionKind::kLoop); 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 3b5cde2996c4195ef458662cd21de85a832d8d55..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); } }; @@ -59,7 +57,7 @@ TEST_F(HloDceTest, NoDeadCode) { builder.AddInstruction(HloInstruction::CreateBinary( constant1->shape(), HloOpcode::kAdd, constant1, constant2)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(3, computation->instruction_count()); @@ -80,7 +78,7 @@ TEST_F(HloDceTest, InstructionsWithSideEffect) { HloInstruction::CreateSend(constant, token, /*channel_id=*/0)); builder.AddInstruction(HloInstruction::CreateTuple({})); - auto module = CreateNewModule(); + auto module = CreateNewUnverifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(4, computation->instruction_count()); @@ -110,7 +108,7 @@ TEST_F(HloDceTest, DeadParameters) { builder.AddInstruction(HloInstruction::CreateUnary( live_param->shape(), HloOpcode::kNegate, live_param)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(5, computation->instruction_count()); @@ -150,7 +148,7 @@ TEST_F(HloDceTest, ControlDependencies) { builder.AddInstruction(HloInstruction::CreateBinary( constant1->shape(), HloOpcode::kAdd, constant1, constant2)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); // Add a control dependency between two instructions. @@ -175,7 +173,7 @@ TEST_F(HloDceTest, ControlDependencies) { // Tests that a dead call instruction is removed. TEST_F(HloDceTest, DeadInstructionWithCalledComputation) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); Shape shape = ShapeUtil::MakeShape(F32, {}); // Called computation for the call instruction. @@ -215,7 +213,7 @@ TEST_F(HloDceTest, DeadInstructionWithCalledComputation) { // Tests that a while instruction with an infeed (effectul instruction) in its // body is not removed, even its user count is 0. TEST_F(HloDceTest, CalledComputationWithSideEffect) { - auto module = CreateNewModule(); + auto module = CreateNewUnverifiedModule(); Shape shape = ShapeUtil::MakeShape(F32, {}); // Condition computation of a while instruction. @@ -270,7 +268,7 @@ TEST_F(HloDceTest, CalledComputationWithSideEffect) { // Tests that a nested call instruction with a side effect is not removed. TEST_F(HloDceTest, CalledComputationWithNestedSideEffect) { - auto module = CreateNewModule(); + auto module = CreateNewUnverifiedModule(); Shape shape = ShapeUtil::MakeShape(F32, {}); // Nested called computation with a side effect. @@ -323,7 +321,7 @@ TEST_F(HloDceTest, CalledComputationWithNestedSideEffect) { } TEST_F(HloDceTest, RemoveDeadSubcomputation) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); HloComputation::Builder subcomp_builder("reduction_subcomp"); @@ -364,7 +362,7 @@ TEST_F(HloDceTest, RemoveDeadSubcomputation) { } TEST_F(HloDceTest, KeepUsedSubcomputation) { - auto module = CreateNewModule(); + auto module = CreateNewUnverifiedModule(); HloComputation::Builder builder(TestName()); HloComputation::Builder subcomp_builder("reduction_subcomp"); diff --git a/tensorflow/compiler/xla/service/hlo_domain_isolator.cc b/tensorflow/compiler/xla/service/hlo_domain_isolator.cc index 72185698c9bdcbf2bebed7ee82bc4ed082ce6a14..19b5734825df833fd34d634e4c1630dd75e96c4c 100644 --- a/tensorflow/compiler/xla/service/hlo_domain_isolator.cc +++ b/tensorflow/compiler/xla/service/hlo_domain_isolator.cc @@ -23,23 +23,14 @@ limitations under the License. namespace xla { -class HloDomainIsolator::RunContext { - public: - RunContext(HloModule* module, HloDomainIsolator* isolator) - : module_(module), isolator_(isolator) {} +namespace { - StatusOr Run(); - - private: - HloModule* module_; - HloDomainIsolator* isolator_; -}; - -StatusOr HloDomainIsolator::RunContext::Run() { - hlo_graph_dumper::MaybeDumpHloModule(*module_, "Before Domain Isolator"); +StatusOr RunInternal(HloModule* module, + HloDomainIsolator::DomainCreator* creator) { + hlo_graph_dumper::MaybeDumpHloModule(*module, "Before Domain Isolator"); int64 added_domains = 0; - for (HloComputation* computation : module_->computations()) { + for (HloComputation* computation : module->computations()) { // Walk in post order and place all the required kDomain instructions. for (HloInstruction* instruction : computation->MakeInstructionPostOrder()) { @@ -55,8 +46,7 @@ StatusOr HloDomainIsolator::RunContext::Run() { root = root->mutable_operand(0); } // Check whether a kDomain is necessary between instruction and operand. - HloInstruction* domain = - isolator_->creator_(instruction, root, operand); + HloInstruction* domain = (*creator)(instruction, root, operand); if (domain != nullptr) { VLOG(4) << "New domain: " << domain->ToString(); TF_RETURN_IF_ERROR(operand->ReplaceUseWith(instruction, domain)); @@ -67,17 +57,19 @@ StatusOr HloDomainIsolator::RunContext::Run() { } VLOG(3) << "Added " << added_domains << " kDomain instructions"; if (added_domains > 0) { - hlo_graph_dumper::MaybeDumpHloModule(*module_, "After Domain Isolator"); + hlo_graph_dumper::MaybeDumpHloModule(*module, "After Domain Isolator"); } return added_domains > 0; } -HloDomainIsolator::HloDomainIsolator(DomainCreator creator) - : creator_(std::move(creator)) {} +} // namespace + +HloDomainIsolator::HloDomainIsolator(DomainCreatorFactory creator_factory) + : creator_factory_(std::move(creator_factory)) {} StatusOr HloDomainIsolator::Run(HloModule* module) { - RunContext run_context(module, this); - return run_context.Run(); + DomainCreator creator = creator_factory_(); + return RunInternal(module, &creator); } } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_domain_isolator.h b/tensorflow/compiler/xla/service/hlo_domain_isolator.h index c0bf1b9e16b52d81365db277abeb06defeb12d44..2274c3a96c2bdd1f4dbd454782699ccb0404529d 100644 --- a/tensorflow/compiler/xla/service/hlo_domain_isolator.h +++ b/tensorflow/compiler/xla/service/hlo_domain_isolator.h @@ -40,17 +40,15 @@ class HloDomainIsolator : public HloModulePass { // Returns nullptr in case no domain separation is necessary. using DomainCreator = std::function; - - explicit HloDomainIsolator(DomainCreator creator); + using DomainCreatorFactory = std::function; + explicit HloDomainIsolator(DomainCreatorFactory creator_factory_); absl::string_view name() const override { return "domain_isolator"; } StatusOr Run(HloModule* module) override; private: - class RunContext; - - DomainCreator creator_; + DomainCreatorFactory creator_factory_; }; } // namespace xla 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_domain_test.cc b/tensorflow/compiler/xla/service/hlo_domain_test.cc index 43e74d2f6f07bd685ad8683401138a4f06cd2ad2..fd4fb0246d8d42ab7329c05dc23e386303cdce3c 100644 --- a/tensorflow/compiler/xla/service/hlo_domain_test.cc +++ b/tensorflow/compiler/xla/service/hlo_domain_test.cc @@ -14,7 +14,7 @@ limitations under the License. ==============================================================================*/ #include "absl/memory/memory.h" -#include "tensorflow/compiler/xla/legacy_flags/debug_options_flags.h" +#include "tensorflow/compiler/xla/debug_options_flags.h" #include "tensorflow/compiler/xla/service/hlo_domain_isolator.h" #include "tensorflow/compiler/xla/service/hlo_domain_metadata.h" #include "tensorflow/compiler/xla/service/hlo_domain_remover.h" @@ -22,13 +22,12 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_sharding_metadata.h" #include "tensorflow/compiler/xla/test.h" #include "tensorflow/compiler/xla/tests/hlo_test_base.h" -#include "tensorflow/compiler/xla/tests/hlo_verified_test_base.h" #include "tensorflow/core/lib/core/status_test_util.h" namespace xla { namespace { -class HloDomainTest : public HloVerifiedTestBase { +class HloDomainTest : public HloTestBase { protected: bool FindUserViaDomainPath(HloInstruction* instruction, HloInstruction* operand) const { @@ -64,13 +63,6 @@ class HloDomainTest : public HloVerifiedTestBase { } return false; } - - StatusOr ParseModule(absl::string_view hlo_string) { - HloModuleConfig config; - config.set_debug_options(legacy_flags::GetDebugOptionsFromFlags()); - ParseAndVerifyModule(hlo_string, config); - return &module(); - } }; // Dummy DomainMetadata implementation which create kDomain boundaries around @@ -106,20 +98,22 @@ class OpNameMetadata : public DomainMetadata { }; // Creator function for OpNameMetadata domains. -HloInstruction* OpNameDomainCreator(HloInstruction* instruction, - HloInstruction* root, - HloInstruction* operand) { - if (instruction->metadata().op_name() == root->metadata().op_name()) { - return nullptr; +class OpNameDomainCreator { + public: + HloInstruction* operator()(HloInstruction* instruction, HloInstruction* root, + HloInstruction* operand) { + if (instruction->metadata().op_name() == root->metadata().op_name()) { + return nullptr; + } + std::unique_ptr operand_side_metadata = + absl::make_unique(root->metadata().op_name()); + std::unique_ptr user_side_metadata = + absl::make_unique(instruction->metadata().op_name()); + return operand->parent()->AddInstruction(HloInstruction::CreateDomain( + operand->shape(), operand, std::move(operand_side_metadata), + std::move(user_side_metadata))); } - std::unique_ptr operand_side_metadata = - absl::make_unique(root->metadata().op_name()); - std::unique_ptr user_side_metadata = - absl::make_unique(instruction->metadata().op_name()); - return operand->parent()->AddInstruction(HloInstruction::CreateDomain( - operand->shape(), operand, std::move(operand_side_metadata), - std::move(user_side_metadata))); -} +}; Status OpNameDomainNormalizer(const DomainMetadata::Domain& domain, const DomainMetadata* metadata) { @@ -142,31 +136,32 @@ ENTRY entry { } )"; - TF_ASSERT_OK_AND_ASSIGN(HloModule * module, ParseModule(hlo_string)); + TF_ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(hlo_string)); LOG(INFO) << "Original module:\n" << module->ToString(); - HloDomainIsolator isolator(ShardingDomainCreator{}); - TF_ASSERT_OK_AND_ASSIGN(bool isolator_changed, isolator.Run(module)); + HloDomainIsolator isolator([]() { return ShardingDomainCreator{}; }); + TF_ASSERT_OK_AND_ASSIGN(bool isolator_changed, isolator.Run(module.get())); EXPECT_TRUE(isolator_changed); - EXPECT_TRUE(HasDomainEdge(module, "c", "a")); - EXPECT_TRUE(HasDomainEdge(module, "c", "b")); - EXPECT_TRUE(HasDomainEdge(module, "d", "a")); - EXPECT_TRUE(HasDomainEdge(module, "d", "b")); - EXPECT_FALSE(HasDomainEdge(module, "e", "c")); - EXPECT_FALSE(HasDomainEdge(module, "e", "d")); + EXPECT_TRUE(HasDomainEdge(module.get(), "c", "a")); + EXPECT_TRUE(HasDomainEdge(module.get(), "c", "b")); + EXPECT_TRUE(HasDomainEdge(module.get(), "d", "a")); + EXPECT_TRUE(HasDomainEdge(module.get(), "d", "b")); + EXPECT_FALSE(HasDomainEdge(module.get(), "e", "c")); + EXPECT_FALSE(HasDomainEdge(module.get(), "e", "d")); HloDomainRemover remover(ShardingMetadata::KindName(), ShardingMetadata::NormalizeShardingDomain); - TF_ASSERT_OK_AND_ASSIGN(bool remover_changed, remover.Run(module)); + TF_ASSERT_OK_AND_ASSIGN(bool remover_changed, remover.Run(module.get())); EXPECT_TRUE(remover_changed); - EXPECT_FALSE(HasDomainEdge(module, "c", "a")); - EXPECT_FALSE(HasDomainEdge(module, "c", "b")); - EXPECT_FALSE(HasDomainEdge(module, "d", "a")); - EXPECT_FALSE(HasDomainEdge(module, "d", "b")); - EXPECT_FALSE(HasDomainEdge(module, "e", "c")); - EXPECT_FALSE(HasDomainEdge(module, "e", "d")); + EXPECT_FALSE(HasDomainEdge(module.get(), "c", "a")); + EXPECT_FALSE(HasDomainEdge(module.get(), "c", "b")); + EXPECT_FALSE(HasDomainEdge(module.get(), "d", "a")); + EXPECT_FALSE(HasDomainEdge(module.get(), "d", "b")); + EXPECT_FALSE(HasDomainEdge(module.get(), "e", "c")); + EXPECT_FALSE(HasDomainEdge(module.get(), "e", "d")); } TEST_F(HloDomainTest, CheckNoDomainAddedIfNoSharding) { @@ -184,11 +179,12 @@ ENTRY entry { } )"; - TF_ASSERT_OK_AND_ASSIGN(HloModule * module, ParseModule(hlo_string)); + TF_ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(hlo_string)); LOG(INFO) << "Original module:\n" << module->ToString(); - HloDomainIsolator isolator(ShardingDomainCreator{}); - TF_ASSERT_OK_AND_ASSIGN(bool isolator_changed, isolator.Run(module)); + HloDomainIsolator isolator([]() { return ShardingDomainCreator{}; }); + TF_ASSERT_OK_AND_ASSIGN(bool isolator_changed, isolator.Run(module.get())); EXPECT_TRUE(!isolator_changed); } @@ -199,10 +195,10 @@ HloModule Module ENTRY entry { p0 = (f32[4]) parameter(0) a = f32[4] get-tuple-element(p0), index=0 - token = token[] after-all() - b = (f32[4], u32[], token[]) send(a, token), channel_id=1, sharding={maximal device=0} + token0 = token[] after-all() + b = (f32[4], u32[], token[]) send(a, token0), channel_id=1, sharding={maximal device=0} c = token[] send-done(b), channel_id=1, sharding={maximal device=0} - d = (f32[4], u32[], token[]) recv(token), channel_id=2, sharding={maximal device=0} + d = (f32[4], u32[], token[]) recv(token0), channel_id=2, sharding={maximal device=0} e = (f32[4], token[]) recv-done(d), channel_id=2, sharding={maximal device=0} e_element = f32[4] get-tuple-element(e), index=0, sharding={maximal device=0} f = f32[4] add(a, e_element) @@ -211,26 +207,27 @@ ENTRY entry { } )"; - TF_ASSERT_OK_AND_ASSIGN(HloModule * module, ParseModule(hlo_string)); + TF_ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(hlo_string)); LOG(INFO) << "Original module:\n" << module->ToString(); - HloDomainIsolator isolator(ShardingDomainCreator{}); - TF_ASSERT_OK_AND_ASSIGN(bool isolator_changed, isolator.Run(module)); + HloDomainIsolator isolator([]() { return ShardingDomainCreator{}; }); + TF_ASSERT_OK_AND_ASSIGN(bool isolator_changed, isolator.Run(module.get())); EXPECT_TRUE(isolator_changed); - EXPECT_TRUE(HasDomainEdge(module, "b", "a")); - EXPECT_TRUE(HasDomainEdge(module, "f", "e_element")); - EXPECT_FALSE(HasDomainEdge(module, "a", "p0")); - EXPECT_FALSE(HasDomainEdge(module, "c", "b")); - EXPECT_FALSE(HasDomainEdge(module, "e", "d")); + EXPECT_TRUE(HasDomainEdge(module.get(), "b", "a")); + EXPECT_TRUE(HasDomainEdge(module.get(), "f", "e_element")); + EXPECT_FALSE(HasDomainEdge(module.get(), "a", "p0")); + EXPECT_FALSE(HasDomainEdge(module.get(), "c", "b")); + EXPECT_FALSE(HasDomainEdge(module.get(), "e", "d")); HloDomainRemover remover(ShardingMetadata::KindName(), ShardingMetadata::NormalizeShardingDomain); - TF_ASSERT_OK_AND_ASSIGN(bool remover_changed, remover.Run(module)); + TF_ASSERT_OK_AND_ASSIGN(bool remover_changed, remover.Run(module.get())); EXPECT_TRUE(remover_changed); - EXPECT_FALSE(HasDomainEdge(module, "b", "a")); - EXPECT_FALSE(HasDomainEdge(module, "f", "e_element")); + EXPECT_FALSE(HasDomainEdge(module.get(), "b", "a")); + EXPECT_FALSE(HasDomainEdge(module.get(), "f", "e_element")); } TEST_F(HloDomainTest, CheckNoDomainAddedOnPureIOComputation) { @@ -238,21 +235,22 @@ TEST_F(HloDomainTest, CheckNoDomainAddedOnPureIOComputation) { HloModule Module ENTRY entry { - token = token[] after-all(), sharding={maximal device=-1} - a = (f32[4], u32[], token[]) recv(token), channel_id=1, sharding={maximal device=-1} + token0 = token[] after-all(), sharding={maximal device=-1} + a = (f32[4], u32[], token[]) recv(token0), channel_id=1, sharding={maximal device=-1} b = (f32[4], token[]) recv-done(a), channel_id=1, sharding={maximal device=-1} b_element = f32[4] get-tuple-element(b), index=0, sharding={maximal device=-1} c = f32[4] add(b_element, b_element), sharding={maximal device=-1} - d = (f32[4], u32[], token[]) send(c, token), channel_id=2, sharding={maximal device=-1} + d = (f32[4], u32[], token[]) send(c, token0), channel_id=2, sharding={maximal device=-1} ROOT e = token[] send-done(d), channel_id=2, sharding={maximal device=-1} } )"; - TF_ASSERT_OK_AND_ASSIGN(HloModule * module, ParseModule(hlo_string)); + TF_ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(hlo_string)); LOG(INFO) << "Original module:\n" << module->ToString(); - HloDomainIsolator isolator(ShardingDomainCreator{}); - TF_ASSERT_OK_AND_ASSIGN(bool isolator_changed, isolator.Run(module)); + HloDomainIsolator isolator([]() { return ShardingDomainCreator{}; }); + TF_ASSERT_OK_AND_ASSIGN(bool isolator_changed, isolator.Run(module.get())); EXPECT_FALSE(isolator_changed); } @@ -261,25 +259,26 @@ TEST_F(HloDomainTest, CheckNormalizationOnPureIOComputation) { HloModule Module ENTRY entry { - token = token[] after-all(), sharding={maximal device=0} - a = (f32[4], u32[], token[]) recv(token), channel_id=1, sharding={maximal device=0} + token0 = token[] after-all(), sharding={maximal device=0} + a = (f32[4], u32[], token[]) recv(token0), channel_id=1, sharding={maximal device=0} b = (f32[4], token[]) recv-done(a), channel_id=1, sharding={maximal device=0} b_element = f32[4] get-tuple-element(b), index=0, sharding={maximal device=0} c = f32[4] add(b_element, b_element) - d = (f32[4], u32[], token[]) send(c, token), channel_id=2, sharding={maximal device=0} + d = (f32[4], u32[], token[]) send(c, token0), channel_id=2, sharding={maximal device=0} ROOT e = token[] send-done(d), channel_id=2, sharding={maximal device=0} } )"; - TF_ASSERT_OK_AND_ASSIGN(HloModule * module, ParseModule(hlo_string)); + TF_ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(hlo_string)); LOG(INFO) << "Original module:\n" << module->ToString(); HloDomainRemover remover(ShardingMetadata::KindName(), ShardingMetadata::NormalizeShardingDomain); - TF_ASSERT_OK_AND_ASSIGN(bool remover_changed, remover.Run(module)); + TF_ASSERT_OK_AND_ASSIGN(bool remover_changed, remover.Run(module.get())); EXPECT_FALSE(remover_changed); - HloInstruction* add = FindInstruction(module, "c"); + HloInstruction* add = FindInstruction(module.get(), "c"); ASSERT_NE(add, nullptr); auto device = add->sharding_unique_device(); EXPECT_TRUE(device.has_value()); @@ -302,41 +301,42 @@ ENTRY entry { } )"; - TF_ASSERT_OK_AND_ASSIGN(HloModule * module, ParseModule(hlo_string)); + TF_ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(hlo_string)); LOG(INFO) << "Original module:\n" << module->ToString(); - HloDomainIsolator sharding_isolator(ShardingDomainCreator{}); + HloDomainIsolator sharding_isolator([]() { return ShardingDomainCreator{}; }); TF_ASSERT_OK_AND_ASSIGN(bool sharding_isolator_changed, - sharding_isolator.Run(module)); + sharding_isolator.Run(module.get())); EXPECT_TRUE(sharding_isolator_changed); - HloDomainIsolator opname_isolator(OpNameDomainCreator); + HloDomainIsolator opname_isolator([]() { return OpNameDomainCreator{}; }); TF_ASSERT_OK_AND_ASSIGN(bool opname_isolator_changed, - opname_isolator.Run(module)); + opname_isolator.Run(module.get())); EXPECT_TRUE(opname_isolator_changed); - EXPECT_TRUE(HasDomainEdge(module, "c", "a")); - EXPECT_TRUE(HasDomainEdge(module, "c", "b")); - EXPECT_TRUE(HasDomainEdge(module, "d", "a")); - EXPECT_TRUE(HasDomainEdge(module, "d", "c")); - EXPECT_FALSE(HasDomainEdge(module, "e", "d")); + EXPECT_TRUE(HasDomainEdge(module.get(), "c", "a")); + EXPECT_TRUE(HasDomainEdge(module.get(), "c", "b")); + EXPECT_TRUE(HasDomainEdge(module.get(), "d", "a")); + EXPECT_TRUE(HasDomainEdge(module.get(), "d", "c")); + EXPECT_FALSE(HasDomainEdge(module.get(), "e", "d")); HloDomainRemover sharding_remover(ShardingMetadata::KindName(), ShardingMetadata::NormalizeShardingDomain); TF_ASSERT_OK_AND_ASSIGN(bool sharding_remover_changed, - sharding_remover.Run(module)); + sharding_remover.Run(module.get())); EXPECT_TRUE(sharding_remover_changed); HloDomainRemover opname_remover(OpNameMetadata::KindName(), OpNameDomainNormalizer); TF_ASSERT_OK_AND_ASSIGN(bool opname_remover_changed, - opname_remover.Run(module)); + opname_remover.Run(module.get())); EXPECT_TRUE(opname_remover_changed); - EXPECT_FALSE(HasDomainEdge(module, "c", "a")); - EXPECT_FALSE(HasDomainEdge(module, "c", "b")); - EXPECT_FALSE(HasDomainEdge(module, "d", "a")); - EXPECT_FALSE(HasDomainEdge(module, "d", "c")); + EXPECT_FALSE(HasDomainEdge(module.get(), "c", "a")); + EXPECT_FALSE(HasDomainEdge(module.get(), "c", "b")); + EXPECT_FALSE(HasDomainEdge(module.get(), "d", "a")); + EXPECT_FALSE(HasDomainEdge(module.get(), "d", "c")); } TEST_F(HloDomainTest, CheckNormalizationOnInfeedTuple) { @@ -344,8 +344,8 @@ TEST_F(HloDomainTest, CheckNormalizationOnInfeedTuple) { HloModule Module ENTRY entry { - token = token[] after-all() - infeed = ((f32[4], f32[4]), token[]) infeed(token), + token0 = token[] after-all() + infeed = ((f32[4], f32[4]), token[]) infeed(token0), sharding={{maximal device=1}, {maximal device=0}, {maximal device=0}} infeed.data = (f32[4], f32[4]) get-tuple-element(infeed), index=0, sharding={{maximal device=1}, {maximal device=0}} @@ -357,16 +357,17 @@ ENTRY entry { } )"; - TF_ASSERT_OK_AND_ASSIGN(HloModule * module, ParseModule(hlo_string)); + TF_ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(hlo_string)); LOG(INFO) << "Original module:\n" << module->ToString(); - HloDomainIsolator isolator(ShardingDomainCreator{}); - TF_ASSERT_OK_AND_ASSIGN(bool isolator_changed, isolator.Run(module)); + HloDomainIsolator isolator([]() { return ShardingDomainCreator{}; }); + TF_ASSERT_OK_AND_ASSIGN(bool isolator_changed, isolator.Run(module.get())); EXPECT_TRUE(isolator_changed); - EXPECT_TRUE(HasDomainEdge(module, "infeed.data", "infeed")); - EXPECT_FALSE(HasDomainEdge(module, "copy0", "gte0")); - EXPECT_FALSE(HasDomainEdge(module, "copy1", "gte1")); + EXPECT_TRUE(HasDomainEdge(module.get(), "infeed.data", "infeed")); + EXPECT_FALSE(HasDomainEdge(module.get(), "copy0", "gte0")); + EXPECT_FALSE(HasDomainEdge(module.get(), "copy1", "gte1")); // Inject unassigned tuple/gte within the infeed domain, to simulate the // HLO passes adding unexpected instructions. @@ -382,7 +383,7 @@ ENTRY entry { // \ / // TUPLE // | - HloInstruction* infeed_data = FindInstruction(module, "infeed.data"); + HloInstruction* infeed_data = FindInstruction(module.get(), "infeed.data"); ASSERT_NE(infeed_data, nullptr); auto infeed_data_users = infeed_data->users(); @@ -408,7 +409,7 @@ ENTRY entry { HloDomainRemover remover(ShardingMetadata::KindName(), ShardingMetadata::NormalizeShardingDomain); - TF_ASSERT_OK_AND_ASSIGN(bool remover_changed, remover.Run(module)); + TF_ASSERT_OK_AND_ASSIGN(bool remover_changed, remover.Run(module.get())); EXPECT_TRUE(remover_changed); struct Assignment { @@ -444,25 +445,26 @@ ENTRY entry { sharding={maximal device=1} })"; - TF_ASSERT_OK_AND_ASSIGN(HloModule * module, ParseModule(hlo_string)); + TF_ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(hlo_string)); - HloDomainIsolator isolator(ShardingDomainCreator{}); - TF_ASSERT_OK_AND_ASSIGN(bool isolator_changed, isolator.Run(module)); + HloDomainIsolator isolator([]() { return ShardingDomainCreator{}; }); + TF_ASSERT_OK_AND_ASSIGN(bool isolator_changed, isolator.Run(module.get())); EXPECT_TRUE(isolator_changed); - EXPECT_TRUE(HasDomainEdge(module, "tuple", "param")); - EXPECT_FALSE(HasDomainEdge(module, "gte", "tuple")); + EXPECT_TRUE(HasDomainEdge(module.get(), "tuple", "param")); + EXPECT_FALSE(HasDomainEdge(module.get(), "gte", "tuple")); // Remove %tuple and %gte (tuple simplification) - HloInstruction* gte = FindInstruction(module, "gte"); - HloInstruction* tuple = FindInstruction(module, "tuple"); + HloInstruction* gte = FindInstruction(module.get(), "gte"); + HloInstruction* tuple = FindInstruction(module.get(), "tuple"); module->entry_computation()->set_root_instruction(tuple->mutable_operand(0)); TF_EXPECT_OK(module->entry_computation()->RemoveInstruction(gte)); TF_EXPECT_OK(module->entry_computation()->RemoveInstruction(tuple)); HloDomainRemover remover(ShardingMetadata::KindName(), ShardingMetadata::NormalizeShardingDomain); - TF_ASSERT_OK_AND_ASSIGN(bool remover_changed, remover.Run(module)); + TF_ASSERT_OK_AND_ASSIGN(bool remover_changed, remover.Run(module.get())); EXPECT_TRUE(remover_changed); const HloInstruction* root = module->entry_computation()->root_instruction(); @@ -484,11 +486,11 @@ TEST_F(HloDomainTest, DumpParseNullSharding) { builder.AddInstruction( HloInstruction::CreateBinary(shape, HloOpcode::kAdd, domain, domain)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); auto hlo_string = module->ToString(); - ASSERT_TRUE(ParseModule(hlo_string).status().ok()); + ASSERT_TRUE(ParseAndReturnVerifiedModule(hlo_string).status().ok()); } // Tuple inputs are domain instructions. @@ -505,20 +507,21 @@ ENTRY entry { } )"; - TF_ASSERT_OK_AND_ASSIGN(HloModule * module, ParseModule(hlo_string)); + TF_ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(hlo_string)); - HloDomainIsolator isolator(ShardingDomainCreator{}); - TF_ASSERT_OK_AND_ASSIGN(bool isolator_changed, isolator.Run(module)); + HloDomainIsolator isolator([]() { return ShardingDomainCreator{}; }); + TF_ASSERT_OK_AND_ASSIGN(bool isolator_changed, isolator.Run(module.get())); EXPECT_TRUE(isolator_changed); // Clear sharding of tpl instruction, in order to test domain sharding // application. - auto tpl = FindInstruction(module, "tpl"); + auto tpl = FindInstruction(module.get(), "tpl"); tpl->clear_sharding(); HloDomainRemover remover(ShardingMetadata::KindName(), ShardingMetadata::NormalizeShardingDomain); - TF_ASSERT_OK_AND_ASSIGN(bool remover_changed, remover.Run(module)); + TF_ASSERT_OK_AND_ASSIGN(bool remover_changed, remover.Run(module.get())); EXPECT_TRUE(remover_changed); EXPECT_EQ(HloSharding::Tuple(tpl->shape(), {HloSharding::AssignDevice(1), @@ -553,36 +556,37 @@ ENTRY %entry (p0: (f32[4], f32[4])) -> (f32[4], f32[4], f32[4]) { ROOT %g = (f32[4], f32[4], f32[4]) tuple(%domain.2, %domain.3, %domain.4) })"; - TF_ASSERT_OK_AND_ASSIGN(HloModule * module, ParseModule(hlo_string)); + TF_ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(hlo_string)); LOG(INFO) << "Original module:\n" << module->ToString(); - HloDomainIsolator opname_isolator(OpNameDomainCreator); + HloDomainIsolator opname_isolator([]() { return OpNameDomainCreator{}; }); TF_ASSERT_OK_AND_ASSIGN(bool opname_isolator_changed, - opname_isolator.Run(module)); + opname_isolator.Run(module.get())); EXPECT_TRUE(opname_isolator_changed); - EXPECT_TRUE(HasDomainEdge(module, "c", "a")); - EXPECT_TRUE(HasDomainEdge(module, "c", "b")); - EXPECT_TRUE(HasDomainEdge(module, "d", "a")); - EXPECT_TRUE(HasDomainEdge(module, "d", "c")); - EXPECT_FALSE(HasDomainEdge(module, "e", "d")); + EXPECT_TRUE(HasDomainEdge(module.get(), "c", "a")); + EXPECT_TRUE(HasDomainEdge(module.get(), "c", "b")); + EXPECT_TRUE(HasDomainEdge(module.get(), "d", "a")); + EXPECT_TRUE(HasDomainEdge(module.get(), "d", "c")); + EXPECT_FALSE(HasDomainEdge(module.get(), "e", "d")); HloDomainRemover sharding_remover(ShardingMetadata::KindName(), ShardingMetadata::NormalizeShardingDomain); TF_ASSERT_OK_AND_ASSIGN(bool sharding_remover_changed, - sharding_remover.Run(module)); + sharding_remover.Run(module.get())); EXPECT_TRUE(sharding_remover_changed); HloDomainRemover opname_remover(OpNameMetadata::KindName(), OpNameDomainNormalizer); TF_ASSERT_OK_AND_ASSIGN(bool opname_remover_changed, - opname_remover.Run(module)); + opname_remover.Run(module.get())); EXPECT_TRUE(opname_remover_changed); - EXPECT_FALSE(HasDomainEdge(module, "c", "a")); - EXPECT_FALSE(HasDomainEdge(module, "c", "b")); - EXPECT_FALSE(HasDomainEdge(module, "d", "a")); - EXPECT_FALSE(HasDomainEdge(module, "d", "c")); + EXPECT_FALSE(HasDomainEdge(module.get(), "c", "a")); + EXPECT_FALSE(HasDomainEdge(module.get(), "c", "b")); + EXPECT_FALSE(HasDomainEdge(module.get(), "d", "a")); + EXPECT_FALSE(HasDomainEdge(module.get(), "d", "c")); } // Emulate instructions inserted at top and bottom within nested tuple domain. @@ -601,15 +605,16 @@ ENTRY entry { } )"; - TF_ASSERT_OK_AND_ASSIGN(HloModule * module, ParseModule(hlo_string)); + TF_ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(hlo_string)); - HloDomainIsolator isolator(ShardingDomainCreator{}); - TF_ASSERT_OK_AND_ASSIGN(bool isolator_changed, isolator.Run(module)); + HloDomainIsolator isolator([]() { return ShardingDomainCreator{}; }); + TF_ASSERT_OK_AND_ASSIGN(bool isolator_changed, isolator.Run(module.get())); EXPECT_TRUE(isolator_changed); // Clear sharding of tuple.0 instruction, in order to test domain sharding // application. - auto tuple0 = FindInstruction(module, "tuple.0"); + auto tuple0 = FindInstruction(module.get(), "tuple.0"); tuple0->clear_sharding(); // Insert the following instructons above and below tuple.0, to emulate other @@ -653,7 +658,7 @@ ENTRY entry { HloDomainRemover remover(ShardingMetadata::KindName(), ShardingMetadata::NormalizeShardingDomain); - TF_ASSERT_OK_AND_ASSIGN(bool remover_changed, remover.Run(module)); + TF_ASSERT_OK_AND_ASSIGN(bool remover_changed, remover.Run(module.get())); EXPECT_TRUE(remover_changed); EXPECT_TRUE(tuple0->has_sharding()); diff --git a/tensorflow/compiler/xla/service/hlo_element_type_converter.cc b/tensorflow/compiler/xla/service/hlo_element_type_converter.cc index 72006e17e7e7ec09b62e88d05b695ec9f4c49647..9b0f2b2a0f4dd5d1d1191e9ab0637cc3034b50da 100644 --- a/tensorflow/compiler/xla/service/hlo_element_type_converter.cc +++ b/tensorflow/compiler/xla/service/hlo_element_type_converter.cc @@ -68,7 +68,7 @@ Shape GetConvertedTupleShape(const Shape& shape, PrimitiveType from_type, std::vector new_tuple_subshapes; for (int64 i = 0; i < ShapeUtil::TupleElementCount(shape); ++i) { Shape subshape = ShapeUtil::GetTupleElementShape(shape, i); - CHECK(!ShapeUtil::IsTuple(subshape)); + CHECK(!subshape.IsTuple()); if (subshape.element_type() == from_type) { subshape = ShapeUtil::ChangeElementType(subshape, to_type); } @@ -92,7 +92,7 @@ HloInstruction* ConvertTupleElements(HloInstruction* hlo, HloInstruction* element = computation->AddInstruction( HloInstruction::CreateGetTupleElement(ele_shape, hlo, i)); const Shape& to_ele_shape = ShapeUtil::GetTupleElementShape(to_shape, i); - CHECK(!ShapeUtil::IsTuple(ele_shape)); + CHECK(!ele_shape.IsTuple()); if (ele_shape.element_type() != to_ele_shape.element_type()) { element = computation->AddInstruction( HloInstruction::CreateConvert(to_ele_shape, element)); @@ -141,10 +141,9 @@ StatusOr HloElementTypeConverter::Run(HloModule* module) { // These are ops with embedded computations where it suffices to convert // the embedded computations instead of converting the ops themselves. if (opcode == HloOpcode::kWhile || opcode == HloOpcode::kCall || - opcode == HloOpcode::kCrossReplicaSum || - opcode == HloOpcode::kFusion || opcode == HloOpcode::kMap || - opcode == HloOpcode::kReduce || opcode == HloOpcode::kReduceWindow || - opcode == HloOpcode::kScatter || + opcode == HloOpcode::kAllReduce || opcode == HloOpcode::kFusion || + opcode == HloOpcode::kMap || opcode == HloOpcode::kReduce || + opcode == HloOpcode::kReduceWindow || opcode == HloOpcode::kScatter || opcode == HloOpcode::kSelectAndScatter || opcode == HloOpcode::kConditional) { continue; @@ -191,7 +190,7 @@ StatusOr HloElementTypeConverter::Run(HloModule* module) { TF_RETURN_IF_ERROR(new_hlo->CopyAllControlDepsFrom(hlo)); new_hlo = ToElementType(new_hlo, eliminate_type_); - } else if (ShapeUtil::IsTuple(hlo->shape())) { + } else if (hlo->shape().IsTuple()) { Shape old_shape = hlo->shape(); Shape new_shape = GetConvertedTupleShape(hlo->shape(), eliminate_type_, replace_with_type_); 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 c170e36c73ad2bef830e528de3ec72d38683d888..5b633784e2f306290ca6c096f67c657be1f188c8 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); @@ -57,13 +49,13 @@ TEST_F(HloElementTypeConverterTest, InfeedsOutfeedsNotConverted) { const string& hlo_string = R"( HloModule InfeedOutfeed ENTRY RoundTrip16MiBR1.v2 { - token = token[] after-all() - infeed = (bf16[4]{0}, token[]) infeed(token) + token0 = token[] after-all() + infeed = (bf16[4]{0}, token[]) infeed(token0) ROOT infeed.data = bf16[4]{0} get-tuple-element(infeed), index=0 - outfeed = token[] outfeed(infeed.data, token) + 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); @@ -96,13 +87,13 @@ TEST_F(HloElementTypeConverterTest, BatchNormGradBF16Converted) { const string& hlo_string = R"( HloModule BatchNormGrad ENTRY BatchNormGrad.v6 { - constant.4 = bf16[2,2,2,1]{3,2,1,0} constant(bf16[2,2,2,1] { { /*i0=0*/ + constant.4 = bf16[2,2,2,1]{3,2,1,0} constant({ { /*i0=0*/ { /*i1=0*/ {0}, {0} }, { /*i1=1*/ {0}, {0} } }, { /*i0=1*/ { /*i1=0*/ {0}, {0} }, { /*i1=1*/ {0}, {0} } } }) constant.5 = bf16[2]{0} constant({1, 1}) constant.6 = bf16[2]{0} constant({0, 0}) constant.7 = bf16[2]{0} constant({1, 1}) - constant.8 = bf16[2,2,2,1]{3,2,1,0} constant(bf16[2,2,2,1] { { /*i0=0*/ + constant.8 = bf16[2,2,2,1]{3,2,1,0} constant({ { /*i0=0*/ { /*i1=0*/ {1}, {2} }, { /*i1=1*/ {3}, {4} } }, { /*i0=1*/ { /*i1=0*/ {5}, {6} }, { /*i1=1*/ {7}, {8} } } }) ROOT batch-norm-grad = (bf16[2,2,2,1]{3,2,1,0}, bf16[2]{0}, bf16[2]{0}) @@ -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())); diff --git a/tensorflow/compiler/xla/service/hlo_evaluator.cc b/tensorflow/compiler/xla/service/hlo_evaluator.cc index c2998883851481b3cda5a3423baa3454018117b2..a40ded76cf051c255ba7c501c9275aa23df52a90 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" @@ -33,17 +33,18 @@ limitations under the License. #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/map_util.h" #include "tensorflow/compiler/xla/primitive_util.h" +#include "tensorflow/compiler/xla/service/cpu/runtime_single_threaded_matmul.h" #include "tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" #include "tensorflow/compiler/xla/service/hlo_query.h" #include "tensorflow/compiler/xla/service/shape_inference.h" #include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/util.h" #include "tensorflow/compiler/xla/window_util.h" #include "tensorflow/core/lib/core/bitmap.h" -#include "tensorflow/core/lib/core/casts.h" #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/platform/logging.h" @@ -135,8 +136,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] = @@ -144,22 +181,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] = @@ -172,6 +201,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 @@ -197,65 +228,30 @@ HloEvaluator::HloEvaluator(int64 max_loop_iterations) }); } -template -StatusOr HloEvaluator::Evaluate( - const HloModule& module, absl::Span arg_literals) { - XLA_VLOG_LINES(2, "HloEvaluator::Evaluate module:\n" + module.ToString()); - - evaluated_.clear(); - arg_literals_.clear(); - for (const auto& literal_ptr : arg_literals) { - arg_literals_.push_back(&*literal_ptr); - } - - TF_RETURN_IF_ERROR(module.entry_computation()->Accept(this)); - - return GetEvaluatedLiteralFor(module.entry_computation()->root_instruction()) - .Clone(); -} - -template <> -StatusOr HloEvaluator::Evaluate( - const HloModule& module, absl::Span arg_literals) { - std::vector arg_literal_ptrs; - for (const auto& literal_ptr : arg_literals) { - arg_literal_ptrs.push_back(&literal_ptr); - } - return Evaluate(module, arg_literal_ptrs); -} - -template StatusOr HloEvaluator::Evaluate( const HloComputation& computation, - absl::Span arg_literals) { + absl::Span arg_literals) { CHECK(computation.parent() != nullptr); XLA_VLOG_LINES( 2, "HloEvaluator::Evaluate computation:\n" + computation.ToString()); - evaluated_.clear(); - arg_literals_.clear(); - for (const auto& literal_ptr : arg_literals) { - arg_literals_.push_back(&*literal_ptr); + if (arg_literals.size() != computation.num_parameters()) { + return InvalidArgument( + "Expected %d argument%s, but got %d.", computation.num_parameters(), + computation.num_parameters() == 1 ? "" : "s", arg_literals.size()); } - - TF_RETURN_IF_ERROR(computation.Accept(this)); - return GetEvaluatedLiteralFor(computation.root_instruction()).Clone(); -} - -template <> -StatusOr HloEvaluator::Evaluate( - const HloComputation& computation, absl::Span arg_literals) { - std::vector arg_literal_ptrs; - for (const auto& literal_ptr : arg_literals) { - arg_literal_ptrs.push_back(&literal_ptr); + for (int64 i = 0; i < arg_literals.size(); ++i) { + const auto& computation_shape = + computation.parameter_instruction(i)->shape(); + const auto& arg_shape = arg_literals[i]->shape(); + if (!ShapeUtil::Equal(computation_shape, arg_shape)) { + return InvalidArgument( + "Shape mismatch at parameter %d. Computation expected %s, but arg " + "was %s.", + i, ShapeUtil::HumanStringWithLayout(computation_shape), + ShapeUtil::HumanString(arg_shape)); + } } - return Evaluate(computation, arg_literal_ptrs); -} - -template -StatusOr HloEvaluator::Evaluate( - HloInstruction* instruction, absl::Span arg_literals) { - TF_RET_CHECK(hlo_query::AllOperandsAreParametersOrConstants(*instruction)); evaluated_.clear(); arg_literals_.clear(); @@ -263,33 +259,20 @@ StatusOr HloEvaluator::Evaluate( arg_literals_.push_back(&*literal_ptr); } - // Evaluate operands of Parameter type against the input literals which - // caches the evaluated literal results. - for (const auto operand : instruction->operands()) { - if (operand->opcode() == HloOpcode::kParameter) { - const Literal* input_literal = arg_literals_[operand->parameter_number()]; - VLOG(2) << "Parameter operand evaluated to: " - << input_literal->ToString(); - TF_RET_CHECK(ShapeUtil::Equal(operand->shape(), input_literal->shape())); - - evaluated_[operand] = input_literal->Clone(); - } + // 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(Preprocess(instruction)); - TF_RETURN_IF_ERROR(instruction->Visit(this)); - TF_RETURN_IF_ERROR(Postprocess(instruction)); - return GetEvaluatedLiteralFor(instruction).Clone(); -} - -template <> -StatusOr HloEvaluator::Evaluate( - HloInstruction* instruction, absl::Span arg_literals) { - std::vector arg_literal_ptrs; - for (const auto& literal : arg_literals) { - arg_literal_ptrs.push_back(&literal); - } - return Evaluate(instruction, arg_literal_ptrs); + TF_RETURN_IF_ERROR(computation.Accept(this)); + return GetEvaluatedLiteralFor(computation.root_instruction()).Clone(); } StatusOr HloEvaluator::Evaluate(HloInstruction* instruction) { @@ -397,16 +380,55 @@ StatusOr HloEvaluator::EvaluateDotOp( return Evaluate(cloned_instruction.get()); } +Status HloEvaluator::HandleBitcast(HloInstruction* bitcast) { + const Literal& operand_literal = GetEvaluatedLiteralFor(bitcast->operand(0)); + Literal result(bitcast->shape()); + TF_RET_CHECK(operand_literal.size_bytes() == result.size_bytes()); + memcpy(result.untyped_data(), operand_literal.untyped_data(), + operand_literal.size_bytes()); + evaluated_[bitcast] = std::move(result); + 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(); } @@ -431,8 +453,8 @@ Status HloEvaluator::HandleConcatenate(HloInstruction* concatenate) { // The result concatenate dimension is going to be the sum of all // concatenate dimensions of the operands taking part of the operation. const Shape& reference_shape = operands[0]->shape(); - CHECK(ShapeUtil::IsArray(reference_shape)); - const int64 rank = ShapeUtil::Rank(reference_shape); + CHECK(reference_shape.IsArray()); + const int64 rank = reference_shape.rank(); const int64 concat_dim = concatenate->dimensions()[0]; CHECK_GE(concat_dim, 0); CHECK_LT(concat_dim, rank); @@ -442,7 +464,7 @@ Status HloEvaluator::HandleConcatenate(HloInstruction* concatenate) { for (int64 i = 1; i < operands.size(); ++i) { const Shape& operand_shape = operands[i]->shape(); - CHECK(ShapeUtil::IsArray(operand_shape)); + CHECK(operand_shape.IsArray()); // Accumulate the concat dimension from all tensors taking part to the // operation. concat_dimensions[concat_dim] += @@ -519,6 +541,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; }, @@ -549,11 +578,61 @@ 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(); +} + +Status HloEvaluator::HandleComplex(HloInstruction* complex) { + const Literal& real = GetEvaluatedLiteralFor(complex->operand(0)); + const Literal& imag = GetEvaluatedLiteralFor(complex->operand(1)); + TF_RET_CHECK(ShapeUtil::Compatible(real.shape(), imag.shape())); + + Literal result(complex->shape()); + 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(); } @@ -590,8 +669,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], @@ -607,8 +689,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], @@ -619,8 +704,11 @@ Status HloEvaluator::HandleCompare(HloInstruction* compare) { evaluated_[compare], Compare(compare->shape(), opcode, lhs_literal, rhs_literal)); } break; - case F16: - return Unimplemented("unhandled primitive type: F16."); + case F16: { + TF_ASSIGN_OR_RETURN( + evaluated_[compare], + Compare(compare->shape(), opcode, lhs_literal, rhs_literal)); + } break; case BF16: { TF_ASSIGN_OR_RETURN(evaluated_[compare], Compare(compare->shape(), opcode, @@ -641,6 +729,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()); @@ -1022,11 +1115,9 @@ Status HloEvaluator::HandleGather(HloInstruction* gather) { Status HloEvaluator::HandleBroadcast(HloInstruction* broadcast) { const Literal& operand = GetEvaluatedLiteralFor(broadcast->operand(0)); - TF_RET_CHECK(broadcast->dimensions().size() == - ShapeUtil::Rank(operand.shape())) + TF_RET_CHECK(broadcast->dimensions().size() == operand.shape().rank()) << "broadcast dimensions is of size: " << broadcast->dimensions().size() - << " and rank of operand_to_broadcast is: " - << ShapeUtil::Rank(operand.shape()); + << " and rank of operand_to_broadcast is: " << operand.shape().rank(); // Checks that operand's dimensions are the same as the broadcast's // dimensions along the dimensions to be broadcasted. for (int64 i = 0; i < broadcast->dimensions().size(); ++i) { @@ -1047,8 +1138,15 @@ Status HloEvaluator::HandleBroadcast(HloInstruction* broadcast) { return Status::OK(); } -Status HloEvaluator::HandleAfterAll(HloInstruction* token) { - evaluated_[token] = LiteralUtil::CreateToken(); +Status HloEvaluator::HandleAfterAll(HloInstruction* after_all) { + evaluated_[after_all] = LiteralUtil::CreateToken(); + return Status::OK(); +} + +Status HloEvaluator::HandleAddDependency(HloInstruction* add_dependency) { + // AddDedendency just forwards its zero-th operand. + evaluated_[add_dependency] = + GetEvaluatedLiteralFor(add_dependency->operand(0)).Clone(); return Status::OK(); } @@ -1092,9 +1190,10 @@ Status HloEvaluator::HandleCall(HloInstruction* call) { } HloEvaluator embedded_evaluator; - Literal result = - embedded_evaluator.Evaluate(*computation, arg_literals) - .ConsumeValueOrDie(); + embedded_evaluator.set_dynamic_dimension_inference( + dynamic_dimension_inference_); + Literal result = embedded_evaluator.Evaluate(*computation, arg_literals) + .ConsumeValueOrDie(); evaluated_[call] = std::move(result); return Status::OK(); @@ -1110,7 +1209,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)); @@ -1124,9 +1225,10 @@ 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) + embedded_evaluator.Evaluate(*readded_computation, arg_literals) .ConsumeValueOrDie(); evaluated_[fusion] = std::move(result); @@ -1144,16 +1246,16 @@ 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 = embedded_evaluator - .Evaluate(*true_computation, - {&true_computation_arg}) - .ConsumeValueOrDie(); + result = + embedded_evaluator.Evaluate(*true_computation, {&true_computation_arg}) + .ConsumeValueOrDie(); } else { result = embedded_evaluator - .Evaluate(*false_computation, - {&false_computation_arg}) + .Evaluate(*false_computation, {&false_computation_arg}) .ConsumeValueOrDie(); } @@ -1200,18 +1302,21 @@ 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).", while_hlo->name(), max_loop_iterations_); } TF_ASSIGN_OR_RETURN(auto cond_val, - cond_evaluator.Evaluate(*cond_comp, {&lcv})); + cond_evaluator.Evaluate(*cond_comp, {&lcv})); keep_going = cond_val.GetFirstElement(); if (keep_going) { - TF_ASSIGN_OR_RETURN(auto body_val, loop_body_evaluator.Evaluate( - *body_comp, {&lcv})); + TF_ASSIGN_OR_RETURN(auto body_val, + loop_body_evaluator.Evaluate(*body_comp, {&lcv})); VLOG(3) << "Loop iteration result: " << body_val.ToString(); lcv = std::move(body_val); cond_evaluator.ResetVisitStates(); @@ -1222,173 +1327,172 @@ 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 = ShapeUtil::Rank(keys_literal.shape()); - 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}); - } - - Literal keys_result_literal(keys_literal.shape()); - Literal values_result_literal(values_literal.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); - increment[sort_dim] = sort_dim_elements; - // Iterate through each dimension except 'sort_dim'. - TF_RETURN_IF_ERROR(ShapeUtil::ForEachIndexWithStatus( - keys_literal.shape(), zero_base, - AsInt64Slice(keys_literal.shape().dimensions()), increment, - [&](absl::Span indices) -> StatusOr { - // Extract a slice from the keys and values literals that correspond 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 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::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); - } - 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)); - 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); +StatusOr ExtractFromIndexPositions(const Literal& from, + absl::Span indices) { + PrimitiveType type = from.shape().element_type(); + switch (type) { + case PRED: { + // 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); + } + case F32: { + std::vector values; + for (int64 index : indices) { + values.push_back(from.Get({index})); + } + return LiteralUtil::CreateR1(values); + } + case U32: { + std::vector values; + for (int64 index : indices) { + values.push_back(from.Get({index})); + } + return LiteralUtil::CreateR1(values); + } + case S32: { + std::vector values; + for (int64 index : indices) { + values.push_back(from.Get({index})); + } + return LiteralUtil::CreateR1(values); + } + case BF16: { + std::vector values; + for (int64 index : indices) { + values.push_back(from.Get({index})); + } + return LiteralUtil::CreateR1(values); + } default: - return InvalidArgument("Unsupported type for Sort"); + return InvalidArgument("Unsupported type for Sort: %s", + PrimitiveType_Name(type)); } } } // namespace Status HloEvaluator::HandleSort(HloInstruction* sort) { - if (!ShapeUtil::IsTuple(sort->shape())) { + if (!sort->shape().IsTuple()) { return DefaultAction(sort); } 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(); + TF_RET_CHECK(sort->operand_count() >= 2) << "Expected key-value 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"; + } + + 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(); } - 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), + Shape key_shape = sort->operand(0)->shape(); + auto rank = key_shape.rank(); + PrimitiveType keys_type = key_shape.element_type(); + if (keys_type != F32 && keys_type != U32 && keys_type != S32 && + keys_type != BF16) { + return InvalidArgument("Unsupported type for Sort: %s", + PrimitiveType_Name(keys_type)); + } + 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 = key_shape.dimensions(sort_dim); + increment[sort_dim] = sort_dim_elements; + // Iterate through each dimension except 'sort_dim'. + TF_RETURN_IF_ERROR(ShapeUtil::ForEachIndexWithStatus( + key_shape, zero_base, AsInt64Slice(key_shape.dimensions()), increment, + [&](absl::Span indices) -> StatusOr { + // 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()); + absl::c_for_each(limit_indices, [](int64& index) { ++index; }); + limit_indices[sort_dim] = sort_dim_elements; + 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::vector indices_to_sort(sort_dim_elements); + std::iota(indices_to_sort.begin(), indices_to_sort.end(), 0); + std::stable_sort( + indices_to_sort.begin(), indices_to_sort.end(), + [keys_type, &literals_to_sort](int64 a, int64 b) { + switch (keys_type) { + case F32: { + auto key_lhs = literals_to_sort[0].Get({a}); + auto key_rhs = literals_to_sort[0].Get({b}); + return SafeLess(key_lhs, key_rhs); + } + case U32: { + auto key_lhs = literals_to_sort[0].Get({a}); + auto key_rhs = literals_to_sort[0].Get({b}); + return SafeLess(key_lhs, key_rhs); + } + case S32: { + auto key_lhs = literals_to_sort[0].Get({a}); + auto key_rhs = literals_to_sort[0].Get({b}); + return SafeLess(key_lhs, key_rhs); + } + case BF16: { + auto key_lhs = literals_to_sort[0].Get({a}); + auto key_rhs = literals_to_sort[0].Get({b}); + return SafeLess(key_lhs, key_rhs); + } + default: + // We should never reach here, because we checked earlier + // that 'key_type' is one of the cases above. + LOG(FATAL) << "Invalid key type in Sort: %s", + PrimitiveType_Name(keys_type); + return false; + } + }); + std::vector slice_dimensions(rank, 1); + slice_dimensions[sort_dim] = sort_dim_elements; + std::vector start_indices(rank, 0); + 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; + })); + + 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); + + 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) { - if (!ShapeUtil::IsTuple(reduce->shape())) { + if (!reduce->shape().IsTuple()) { return DefaultAction(reduce); } else { auto first_element_type = reduce->shape().tuple_shapes(0).element_type(); @@ -1420,16 +1524,46 @@ Status HloEvaluator::Postprocess(HloInstruction* hlo) { return Status::OK(); } -// Explicit instantiation of templatized Evaluate* methods. -// -template StatusOr HloEvaluator::Evaluate( - const HloModule& module, absl::Span arg_literals); +namespace { +template +std::unique_ptr> MatmulArray2DImpl( + const Array2D& lhs, const Array2D& rhs, + const std::function& impl_fn) { + CHECK_EQ(lhs.width(), rhs.height()); + int m = lhs.height(); + int n = rhs.width(); + int k = lhs.width(); + auto result = absl::make_unique>(m, n); + // Because Eigen is a header-oriented library, make sure that the Eigen code + // is the same as the code used by the CPU backend (otherwise the linker will + // randomly pick *some* definition). + impl_fn( + /*run_options_ptr=*/nullptr, result->data(), rhs.data(), lhs.data(), n, m, + k, + /*transpose_lhs=*/0, + /*transpose_rhs=*/0); + return result; +} +} // namespace -template StatusOr HloEvaluator::Evaluate( - const HloComputation& computation, - absl::Span arg_literals); +std::unique_ptr> HloEvaluator::MatmulArray2D( + const Array2D& lhs, const Array2D& rhs) { + return MatmulArray2DImpl( + lhs, rhs, __xla_cpu_runtime_EigenSingleThreadedMatMulF16); +} -template StatusOr HloEvaluator::Evaluate( - HloInstruction* instruction, absl::Span arg_literals); +std::unique_ptr> HloEvaluator::MatmulArray2D( + const Array2D& lhs, const Array2D& rhs) { + return MatmulArray2DImpl( + lhs, rhs, __xla_cpu_runtime_EigenSingleThreadedMatMulF32); +} + +std::unique_ptr> HloEvaluator::MatmulArray2D( + const Array2D& lhs, const Array2D& rhs) { + return MatmulArray2DImpl( + lhs, rhs, __xla_cpu_runtime_EigenSingleThreadedMatMulF64); +} } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_evaluator.h b/tensorflow/compiler/xla/service/hlo_evaluator.h index 07f8d0aad4af0b07303b4e485b3630cc75bcb519..ccb8af4fb07fedb054693b78e8bab49527d38700 100644 --- a/tensorflow/compiler/xla/service/hlo_evaluator.h +++ b/tensorflow/compiler/xla/service/hlo_evaluator.h @@ -21,7 +21,9 @@ limitations under the License. #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/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" @@ -42,16 +44,24 @@ class HloEvaluator : public DfsHloVisitorWithDefault { // specified. explicit HloEvaluator(int64 max_loop_iterations = -1); - // Evaluates an HLO module and an array of pointers to literals. - // Returns the evaluated result as a literal if successful. + // Evaluates an HLO module and an array of pointers to literals. Returns the + // evaluated result as a literal if successful. + // // Precondition: The indices of arg_literals correspond to the parameter // numbers of the HLO parameters in the computation. See comment below for an // example. - // `LiteralPtr` accepts either Literal or const Literal* - // type. - template + // + // (Dummy template arg is to reduce the overloading priority of one overload + // so that Evaluate(module, {}) resolves unambiguously.) + StatusOr Evaluate(const HloModule& module, + absl::Span arg_literals) { + return Evaluate(*module.entry_computation(), arg_literals); + } + template StatusOr Evaluate(const HloModule& module, - absl::Span arg_literals); + absl::Span arg_literals) { + return Evaluate(*module.entry_computation(), arg_literals); + } // Evaluates an HLO computation and an array of pointers to literals. // Returns the evaluated result as a literal if successful. @@ -69,29 +79,24 @@ class HloEvaluator : public DfsHloVisitorWithDefault { // where Parameter0 has parameter_number 0 and Parameter1 has parameter_number // 1 in this computation. The input literals array will then have its first // literal map to Parameter0 and the second map to Parameter1. - // `LiteralPtr` accepts either Literal or const Literal* - // type. - template + // + // (Dummy template arg is to reduce the overloading priority of one overload + // so that Evaluate(module, {}) resolves unambiguously.) StatusOr Evaluate(const HloComputation& computation, - absl::Span arg_literals); - - // Evaluates a single HLO instruction and an array of pointers to literals. - // Return the evaluated result as literal if successful. - // Precondition: - // 1. argument literals correspond to the input instruction's parameters in - // their post-ordering. - // 2. the instruction's operands must be of either Parameter or Constant type. - // `LiteralPtr` accepts either Literal or const Literal* - // type. - template - StatusOr Evaluate(HloInstruction* instruction, - absl::Span arg_literals); - - // Evaluates a single HLO instruction with constant operands. - // Returns the evaluated result as literal if successful. - // Precondition: - // 1. all operands of the input instruction are constants. - // 2. the instruction is not a Parameter operation. + absl::Span arg_literals); + template + StatusOr Evaluate(const HloComputation& computation, + absl::Span arg_literals) { + std::vector arg_literal_ptrs; + for (const auto& l : arg_literals) { + arg_literal_ptrs.push_back(&l); + } + return Evaluate(computation, arg_literal_ptrs); + } + + // Gets the value of running a single HLO instruction. + // + // All of the operands to this instruction must be constants. StatusOr Evaluate(HloInstruction* instruction); // Same as Evaluate, except returning false on error and accepts an output @@ -119,6 +124,22 @@ 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; } + + // Returns the result of a matrix multiply `lhs x rhs`. + static std::unique_ptr> MatmulArray2D( + const Array2D& lhs, const Array2D& rhs); + static std::unique_ptr> MatmulArray2D( + const Array2D& lhs, const Array2D& rhs); + static std::unique_ptr> MatmulArray2D( + const Array2D& lhs, const Array2D& rhs); + protected: // Make HloEvaluatorTypedVisitor a friend because it is logically part of this // class. @@ -144,6 +165,10 @@ class HloEvaluator : public DfsHloVisitorWithDefault { // Operations that are type-agnostic or always return a specific type, such as // HandleIsFinite where boolean is always returned. // + Status HandleBitcast(HloInstruction* bitcast) override; + + Status HandleGetDimensionSize(HloInstruction* get_dimension_size) override; + Status HandleParameter(HloInstruction* parameter) override; Status HandleConstant(HloInstruction* constant) override; @@ -180,7 +205,9 @@ class HloEvaluator : public DfsHloVisitorWithDefault { Status HandleBroadcast(HloInstruction* broadcast) override; - Status HandleAfterAll(HloInstruction* token) override; + Status HandleAfterAll(HloInstruction* after_all) override; + + Status HandleAddDependency(HloInstruction* add_dependency) override; Status HandleSort(HloInstruction* sort) override; @@ -188,16 +215,49 @@ class HloEvaluator : public DfsHloVisitorWithDefault { Status HandleImag(HloInstruction* imag) override; + Status HandleComplex(HloInstruction* complex) override; + Status HandleReduce(HloInstruction* reduce) 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(); @@ -205,14 +265,23 @@ 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. + bool use_fast_path_ = false; + private: template static StatusOr ElementWiseUnaryOpImpl( @@ -221,16 +290,7 @@ class HloEvaluator : public DfsHloVisitorWithDefault { const Literal& operand_literal) { const auto shape = instruction->shape(); const auto* operand = instruction->operand(0); - - // TODO(b/35950897, b/27796129): add DCHECK back once implicit broadcast is - // removed. - if (!ShapeUtil::SameDimensions(shape, operand->shape())) { - return Unimplemented( - "Implicit broadcasting is currently unsupported in HLO evaluator " - "Shape Mismatch: %s vs %s", - ShapeUtil::HumanString(shape), - ShapeUtil::HumanString(operand->shape())); - } + TF_RET_CHECK(ShapeUtil::SameDimensions(shape, operand->shape())); Literal result(shape); TF_RETURN_IF_ERROR( @@ -252,9 +312,20 @@ class HloEvaluator : public DfsHloVisitorWithDefault { // Max loop iterations to execute with no maximum if negative. int64 max_loop_iterations_; + // Module-level seed handle. + uint64 seed_; + // 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_; + TF_DISALLOW_COPY_AND_ASSIGN(HloEvaluator); }; +std::unique_ptr> MatmulArray2D(const Array2D& lhs, + const Array2D& rhs); } // namespace xla #endif // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_EVALUATOR_H_ diff --git a/tensorflow/compiler/xla/service/hlo_evaluator_test.cc b/tensorflow/compiler/xla/service/hlo_evaluator_test.cc index 608a42bb60702aa075daca39535ca1672dcc5467..590f76f4727a30dfe9c8a1641d7a9f4c16c0ad35 100644 --- a/tensorflow/compiler/xla/service/hlo_evaluator_test.cc +++ b/tensorflow/compiler/xla/service/hlo_evaluator_test.cc @@ -22,6 +22,7 @@ limitations under the License. #include #include "absl/memory/memory.h" +#include "absl/strings/str_format.h" #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/compiler/xla/literal.h" #include "tensorflow/compiler/xla/reference_util.h" @@ -33,8 +34,9 @@ limitations under the License. #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/test.h" -#include "tensorflow/compiler/xla/tests/hlo_verified_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_utils.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/util.h" #include "tensorflow/compiler/xla/xla_data.pb.h" @@ -49,46 +51,40 @@ namespace { static std::array use_bf16_params{true, false}; -class HloEvaluatorTest : public ::testing::WithParamInterface, - public HloVerifiedTestBase { - protected: - HloEvaluatorTest() : HloVerifiedTestBase(), use_bfloat16_(GetParam()) { - evaluator_ = absl::make_unique(); - } +// Test fixture for the HloEvaluator. +// +// In bf16 mode, all f32 shapes are converted to bf16 before running. +class HloEvaluatorTest : public HloTestBase { + public: + HloEvaluatorTest() : use_bfloat16_(false) {} Literal Evaluate(absl::Span arg_literals = {}) { if (use_bfloat16_) { - // In BF16 mode, we convert all F32 type to BF16 and evaluate the module. - auto type_converter = HloElementTypeConverter(F32, BF16); - type_converter.Run(&module()).ValueOrDie(); + HloElementTypeConverter(F32, BF16).Run(m_.get()).ValueOrDie(); } - return evaluator_->Evaluate(*module().entry_computation(), arg_literals) + return evaluator_.Evaluate(*m_->entry_computation(), arg_literals) .ConsumeValueOrDie(); } - // Evaluate function that takes in a local module instead of using module_ - // that is in HloVerifiedTestBase. Once module_ in HloVerifiedTestBase is + // Evaluate function that takes in a local module instead of using m_ + // that is in HloTestBase. Once m_ in HloTestBase is // removed, this should be the default Evaluate function. Literal EvaluateWithModule( HloModule* module, absl::Span arg_literals = {}) { if (use_bfloat16_) { - // In BF16 mode, we convert all F32 type to BF16 and evaluate the module. - auto type_converter = HloElementTypeConverter(F32, BF16); - type_converter.Run(module).ValueOrDie(); + HloElementTypeConverter(F32, BF16).Run(m_.get()).ValueOrDie(); } - return evaluator_->Evaluate(*module->entry_computation(), arg_literals) + return evaluator_.Evaluate(*module->entry_computation(), arg_literals) .ConsumeValueOrDie(); } - std::unique_ptr evaluator_; - void TestUnaryOp(HloOpcode opcode, Literal expected, Literal input, float aabs = 0) { HloComputation::Builder b(TestName()); auto c1 = b.AddInstruction(HloInstruction::CreateConstant(std::move(input))); b.AddInstruction(HloInstruction::CreateUnary(expected.shape(), opcode, c1)); - module().AddEntryComputation(b.Build()); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate(); @@ -108,22 +104,34 @@ class HloEvaluatorTest : public ::testing::WithParamInterface, auto c2 = b.AddInstruction(HloInstruction::CreateConstant(std::move(rhs))); b.AddInstruction( HloInstruction::CreateBinary(expected.shape(), opcode, c1, c2)); - module().AddEntryComputation(b.Build()); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate(); EXPECT_TRUE(LiteralTestUtil::Equal(expected, result)); } - bool use_bfloat16_; + protected: + explicit HloEvaluatorTest(bool use_bfloat16) : use_bfloat16_(use_bfloat16) {} + HloEvaluator evaluator_; + + const bool use_bfloat16_; + std::unique_ptr m_ = CreateNewVerifiedModule(); +}; + +// Lets you write TEST_Ps that run twice, once with and once without bf16. +class HloEvaluatorBf16Test : public ::testing::WithParamInterface, + public HloEvaluatorTest { + protected: + HloEvaluatorBf16Test() : HloEvaluatorTest(/*use_bfloat16=*/GetParam()) {} }; -#define XLA_TYPED_TEST_P(test_case_name, test_name, test_type1) \ - TEST_P(test_case_name, test_name) +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. -TEST_P(HloEvaluatorTest, DoesClamp) { +TEST_P(HloEvaluatorBf16Test, DoesClamp) { auto low = LiteralUtil::CreateR2({{0.f, 2.f}, {2.f, 4.f}}); auto value = LiteralUtil::CreateR2({{0.f, 5.f}, {0.f, 4.f}}); auto high = LiteralUtil::CreateR2({{2.f, 4.f}, {4.f, 4.f}}); @@ -135,7 +143,7 @@ TEST_P(HloEvaluatorTest, DoesClamp) { auto c3 = b.AddInstruction(HloInstruction::CreateConstant(std::move(high))); b.AddInstruction( HloInstruction::CreateTernary(shape, HloOpcode::kClamp, c1, c2, c3)); - module().AddEntryComputation(b.Build()); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate(); @@ -144,7 +152,7 @@ TEST_P(HloEvaluatorTest, DoesClamp) { EXPECT_TRUE(LiteralTestUtil::Equal(expected, result)); } -TEST_P(HloEvaluatorTest, DISABLED_DoesClampSpecialBroadcast) { +TEST_P(HloEvaluatorBf16Test, DISABLED_DoesClampSpecialBroadcast) { auto low = LiteralUtil::CreateR0(0.f); auto value = LiteralUtil::CreateR2({{-1.f, 0.f}, {1.f, 2.f}}); auto high = LiteralUtil::CreateR0(1.f); @@ -156,7 +164,7 @@ TEST_P(HloEvaluatorTest, DISABLED_DoesClampSpecialBroadcast) { auto c3 = b.AddInstruction(HloInstruction::CreateConstant(std::move(high))); b.AddInstruction( HloInstruction::CreateTernary(shape, HloOpcode::kClamp, c1, c2, c3)); - module().AddEntryComputation(b.Build()); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate(); @@ -167,7 +175,7 @@ TEST_P(HloEvaluatorTest, DISABLED_DoesClampSpecialBroadcast) { // Verifies that HloEvaluator evaluates a HLO instruction that performs select // with 3 operands. -TEST_P(HloEvaluatorTest, DoesSelect) { +TEST_P(HloEvaluatorBf16Test, DoesSelect) { auto pred = LiteralUtil::CreateR2({{true, false}, {false, true}}); auto on_true = LiteralUtil::CreateR2({{2.f, 4.f}, {4.f, 4.f}}); auto on_false = LiteralUtil::CreateR2({{0.f, 5.f}, {0.f, 4.f}}); @@ -181,7 +189,7 @@ TEST_P(HloEvaluatorTest, DoesSelect) { b.AddInstruction(HloInstruction::CreateConstant(std::move(on_false))); b.AddInstruction( HloInstruction::CreateTernary(shape, HloOpcode::kSelect, c1, c2, c3)); - module().AddEntryComputation(b.Build()); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate({}); @@ -192,7 +200,7 @@ TEST_P(HloEvaluatorTest, DoesSelect) { // Verifies that HloEvaluator evaluates a HLO instruction that performs // element-wise addition with 2 operands. -TEST_P(HloEvaluatorTest, DoesAdd) { +TEST_F(HloEvaluatorTest, DoesAdd) { auto lhs = LiteralUtil::CreateR2({{1, 0}, {-100, 4}}); auto rhs = LiteralUtil::CreateR2({{2, 4}, {4, 4}}); auto expected = LiteralUtil::CreateR2({{3, 4}, {-96, 8}}); @@ -201,7 +209,7 @@ TEST_P(HloEvaluatorTest, DoesAdd) { } // Verifies that HloEvaluator evaluates a HLO instruction that performs // element-wise and with 2 operands. -TEST_P(HloEvaluatorTest, DoesAnd) { +TEST_P(HloEvaluatorBf16Test, DoesAnd) { auto lhs = LiteralUtil::CreateR2({{1, 0}, {-100, 4}}); auto rhs = LiteralUtil::CreateR2({{2, 4}, {4, 4}}); auto expected = LiteralUtil::CreateR2({{0, 0}, {4, 4}}); @@ -210,7 +218,7 @@ TEST_P(HloEvaluatorTest, DoesAnd) { } // Verifies that HloEvaluator evaluates a HLO instruction that performs // element-wise or with 2 operands. -TEST_P(HloEvaluatorTest, DoesOr) { +TEST_F(HloEvaluatorTest, DoesOr) { auto lhs = LiteralUtil::CreateR2({{1, 0}, {-100, 4}}); auto rhs = LiteralUtil::CreateR2({{2, 4}, {4, 4}}); auto expected = LiteralUtil::CreateR2({{3, 4}, {-100, 4}}); @@ -219,7 +227,7 @@ TEST_P(HloEvaluatorTest, DoesOr) { } // Verifies that HloEvaluator evaluates a HLO instruction that performs // element-wise or with 2 operands. -TEST_P(HloEvaluatorTest, DoesXor) { +TEST_F(HloEvaluatorTest, DoesXor) { auto lhs = LiteralUtil::CreateR2({{1, 0}, {-100, 4}}); auto rhs = LiteralUtil::CreateR2({{2, 4}, {4, 4}}); auto expected = LiteralUtil::CreateR2({{3, 4}, {-104, 0}}); @@ -228,7 +236,7 @@ TEST_P(HloEvaluatorTest, DoesXor) { } // Verifies that HloEvaluator evaluates a HLO instruction that performs // element-wise multiply with 2 operands. -TEST_P(HloEvaluatorTest, DoesMultiply) { +TEST_F(HloEvaluatorTest, DoesMultiply) { auto lhs = LiteralUtil::CreateR2({{-1, 0}, {-100, 4}}); auto rhs = LiteralUtil::CreateR2( {{std::numeric_limits::min(), 4}, {4, 4}}); @@ -239,14 +247,14 @@ TEST_P(HloEvaluatorTest, DoesMultiply) { } // Verifies that HloEvaluator evaluates a HLO instruction that performs // element-wise divide with 2 operands. -TEST_P(HloEvaluatorTest, DoesDivideInt64) { +TEST_F(HloEvaluatorTest, DoesDivideInt64) { auto lhs = LiteralUtil::CreateR2({{1, 0}, {-100, 4}}); auto rhs = LiteralUtil::CreateR2({{2, 4}, {4, 4}}); auto expected = LiteralUtil::CreateR2({{0, 0}, {-25, 1}}); TestBinaryOp(HloOpcode::kDivide, std::move(expected), std::move(lhs), std::move(rhs)); } -TEST_P(HloEvaluatorTest, DoesDivideDouble) { +TEST_P(HloEvaluatorBf16Test, DoesDivideDouble) { auto lhs = LiteralUtil::CreateR2({{1.0, 0.0}, {-100.0, 4.0}}); auto rhs = LiteralUtil::CreateR2({{2.2, 4.0}, {4.0, 4.0}}); auto expected = @@ -257,41 +265,41 @@ TEST_P(HloEvaluatorTest, DoesDivideDouble) { // Verifies that HloEvaluator evaluates a HLO instruction that performs // element-wise abs op with 1 operand. -TEST_P(HloEvaluatorTest, DoesAbsR2) { +TEST_F(HloEvaluatorTest, DoesAbsR2) { auto operand = LiteralUtil::CreateR2({{1, -20}, {-100, 4}}); auto expected = LiteralUtil::CreateR2({{1, 20}, {100, 4}}); TestUnaryOp(HloOpcode::kAbs, std::move(expected), std::move(operand)); } -TEST_P(HloEvaluatorTest, DoesAbsR0) { +TEST_P(HloEvaluatorBf16Test, DoesAbsR0) { auto operand = LiteralUtil::CreateR0(-1.0f); auto expected = LiteralUtil::CreateR0(1.0f); TestUnaryOp(HloOpcode::kAbs, std::move(expected), std::move(operand)); } -TEST_P(HloEvaluatorTest, DoesAbsR1WithZeroSize) { +TEST_P(HloEvaluatorBf16Test, DoesAbsR1WithZeroSize) { auto operand = LiteralUtil::CreateR1({}); auto expected = LiteralUtil::CreateR1({}); TestUnaryOp(HloOpcode::kAbs, std::move(expected), std::move(operand)); } -TEST_P(HloEvaluatorTest, DoesNegateR2) { +TEST_F(HloEvaluatorTest, DoesNegateR2) { auto operand = LiteralUtil::CreateR2( {{0, std::numeric_limits::min()}, {-1, 4}}); auto expected = LiteralUtil::CreateR2( {{0, std::numeric_limits::min()}, {1, -4}}); TestUnaryOp(HloOpcode::kNegate, std::move(expected), std::move(operand)); } -TEST_P(HloEvaluatorTest, DoesCosR2) { +TEST_P(HloEvaluatorBf16Test, DoesCosR2) { auto operand = LiteralUtil::CreateR2({{0, M_PI}, {-M_PI, 2 * M_PI}}); auto expected = LiteralUtil::CreateR2({{1, -1}, {-1, 1}}); TestUnaryOp(HloOpcode::kCos, std::move(expected), std::move(operand), use_bfloat16_ ? 0.031250 : 9.5367431640625E-7); } -TEST_P(HloEvaluatorTest, DoesSinR2) { +TEST_P(HloEvaluatorBf16Test, DoesSinR2) { auto operand = LiteralUtil::CreateR2({{0, M_PI}, {-M_PI, 2 * M_PI}}); auto expected = LiteralUtil::CreateR2({{0, 0}, {0, 0}}); TestUnaryOp(HloOpcode::kSin, std::move(expected), std::move(operand), use_bfloat16_ ? 0.031250 : 9.5367431640625E-7); } -TEST_P(HloEvaluatorTest, DoesNotR2) { +TEST_F(HloEvaluatorTest, DoesNotR2) { auto operand = LiteralUtil::CreateR2({{0, std::numeric_limits::min()}, {-1, std::numeric_limits::max()}}); @@ -300,9 +308,22 @@ TEST_P(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_P(HloEvaluatorTest, DoesTraverseInstructions) { +TEST_F(HloEvaluatorTest, DoesTraverseInstructions) { auto lhs = LiteralUtil::CreateR2({{1, 0}, {-100, 4}}); auto rhs = LiteralUtil::CreateR2({{2, 4}, {4, 4}}); auto rhs2 = LiteralUtil::CreateR2({{1, -20}, {-100, 4}}); @@ -322,7 +343,7 @@ TEST_P(HloEvaluatorTest, DoesTraverseInstructions) { b.AddInstruction(HloInstruction::CreateParameter(2, shape, "rhs2")); b.AddInstruction(HloInstruction::CreateBinary(shape, HloOpcode::kAdd, lhs_instruction, param_rhs2)); - module().AddEntryComputation(b.Build()); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate(args); @@ -332,7 +353,7 @@ TEST_P(HloEvaluatorTest, DoesTraverseInstructions) { } // Verifies Reshape operation is correctly evaluated. -TEST_P(HloEvaluatorTest, DoesReshape) { +TEST_F(HloEvaluatorTest, DoesReshape) { HloComputation::Builder b(TestName()); const int64 dimensions[] = {11, 8, 7, 5, 9}; TF_ASSERT_OK_AND_ASSIGN(auto literal, @@ -346,7 +367,7 @@ TEST_P(HloEvaluatorTest, DoesReshape) { const int64 permutation[] = {1, 2, 0, 4, 3}; b.AddInstruction( HloInstruction::CreateTranspose(shape, literal_instruction, permutation)); - module().AddEntryComputation(b.Build()); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate({}); @@ -358,7 +379,7 @@ TEST_P(HloEvaluatorTest, DoesReshape) { } // Verifies Broadcast operation is correctly evaluated. -TEST_P(HloEvaluatorTest, DoesBroadcast) { +TEST_F(HloEvaluatorTest, DoesBroadcast) { HloComputation::Builder b(TestName()); auto input_literal = LiteralUtil::CreateR2({{1, 2}, {3, 4}, {5, 6}}); auto output_literal = LiteralUtil::CreateR3( @@ -367,14 +388,14 @@ TEST_P(HloEvaluatorTest, DoesBroadcast) { HloInstruction::CreateConstant(std::move(input_literal))); b.AddInstruction(HloInstruction::CreateBroadcast( output_literal.shape(), literal_instruction, {1, 2})); - module().AddEntryComputation(b.Build()); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate({}); EXPECT_TRUE(LiteralTestUtil::Equal(result, output_literal)); } -TEST_P(HloEvaluatorTest, DoesBroadcastScalar) { +TEST_F(HloEvaluatorTest, DoesBroadcastScalar) { HloComputation::Builder b(TestName()); auto input_literal = LiteralUtil::CreateR0(111); auto output_literal = LiteralUtil::CreateR2( @@ -386,14 +407,14 @@ TEST_P(HloEvaluatorTest, DoesBroadcastScalar) { b.AddInstruction(HloInstruction::CreateBroadcast( output_literal.shape(), literal_instruction, /*broadcast_dimensions=*/{})); - module().AddEntryComputation(b.Build()); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate({}); EXPECT_TRUE(LiteralTestUtil::Equal(result, output_literal)); } -TEST_P(HloEvaluatorTest, DoesConcatenateSimple) { +TEST_F(HloEvaluatorTest, DoesConcatenateSimple) { HloComputation::Builder b(TestName()); HloInstruction* operand1 = b.AddInstruction(HloInstruction::CreateConstant( @@ -406,7 +427,7 @@ TEST_P(HloEvaluatorTest, DoesConcatenateSimple) { Shape shape = ShapeUtil::MakeShape(S64, {4, 2}); b.AddInstruction(HloInstruction::CreateConcatenate(shape, operands, 0)); - module().AddEntryComputation(b.Build()); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate(); @@ -415,7 +436,7 @@ TEST_P(HloEvaluatorTest, DoesConcatenateSimple) { EXPECT_TRUE(LiteralTestUtil::Equal(expected, result)); } -TEST_P(HloEvaluatorTest, ConcatenateHandlesShapeWithZeroElement) { +TEST_F(HloEvaluatorTest, ConcatenateHandlesShapeWithZeroElement) { HloComputation::Builder b(TestName()); HloInstruction* operand1 = b.AddInstruction( @@ -428,7 +449,7 @@ TEST_P(HloEvaluatorTest, ConcatenateHandlesShapeWithZeroElement) { Shape shape = ShapeUtil::MakeShape(S64, {2}); b.AddInstruction(HloInstruction::CreateConcatenate(shape, operands, 0)); - module().AddEntryComputation(b.Build()); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate(); @@ -436,7 +457,7 @@ TEST_P(HloEvaluatorTest, ConcatenateHandlesShapeWithZeroElement) { EXPECT_TRUE(LiteralTestUtil::Equal(expected, result)); } -TEST_P(HloEvaluatorTest, ConvertWithSameLayout) { +TEST_P(HloEvaluatorBf16Test, ConvertWithSameLayout) { HloComputation::Builder b(TestName()); auto input_literal = LiteralUtil::CreateR2({{1, 2}, {3, 4}, {5, 6}}); @@ -448,14 +469,14 @@ TEST_P(HloEvaluatorTest, ConvertWithSameLayout) { HloInstruction* constant = b.AddInstruction( HloInstruction::CreateConstant(std::move(input_literal))); b.AddInstruction(HloInstruction::CreateConvert(expected.shape(), constant)); - module().AddEntryComputation(b.Build()); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate(); EXPECT_TRUE(LiteralTestUtil::Equal(result, expected)); } -TEST_P(HloEvaluatorTest, ConvertWithDifferentLayout) { +TEST_P(HloEvaluatorBf16Test, ConvertWithDifferentLayout) { HloComputation::Builder b(TestName()); auto input_literal = LiteralUtil::CreateR2WithLayout( @@ -468,7 +489,7 @@ TEST_P(HloEvaluatorTest, ConvertWithDifferentLayout) { HloInstruction* constant = b.AddInstruction( HloInstruction::CreateConstant(std::move(input_literal))); b.AddInstruction(HloInstruction::CreateConvert(expected.shape(), constant)); - module().AddEntryComputation(b.Build()); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate(); @@ -488,7 +509,7 @@ PaddingConfig CreatePaddingConfig( return padding_config; } -TEST_P(HloEvaluatorTest, Pad2DIntegerArrayWithZeroDimension) { +TEST_F(HloEvaluatorTest, Pad2DIntegerArrayWithZeroDimension) { auto operand = LiteralUtil::CreateR2({{}, {}}); HloComputation::Builder b(TestName()); auto operand_instruction = @@ -503,7 +524,7 @@ TEST_P(HloEvaluatorTest, Pad2DIntegerArrayWithZeroDimension) { Shape shape = ShapeUtil::MakeShape(S32, {5, 2}); b.AddInstruction(HloInstruction::CreatePad( shape, operand_instruction, padding_value_instruction, padding_config)); - module().AddEntryComputation(b.Build()); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate(); @@ -513,7 +534,7 @@ TEST_P(HloEvaluatorTest, Pad2DIntegerArrayWithZeroDimension) { EXPECT_TRUE(LiteralTestUtil::Equal(expected, result)); } -TEST_P(HloEvaluatorTest, Pad4DFloatArrayWithInteriorPadding) { +TEST_P(HloEvaluatorBf16Test, Pad4DFloatArrayWithInteriorPadding) { HloComputation::Builder b(TestName()); Array4D input_array(3, 2, 1, 1, {1, 2, 3, 4, 5, 6}); @@ -530,7 +551,7 @@ TEST_P(HloEvaluatorTest, Pad4DFloatArrayWithInteriorPadding) { CreatePaddingConfig({{{1, 0, 2}}, {{0, 2, 1}}, {{0, 0, 0}}, {{0, 0, 0}}}); b.AddInstruction(HloInstruction::CreatePad( shape, input_instruction, pad_instruction, r4_padding_on_dim0_dim1)); - module().AddEntryComputation(b.Build()); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate(); @@ -548,7 +569,7 @@ TEST_P(HloEvaluatorTest, Pad4DFloatArrayWithInteriorPadding) { EXPECT_TRUE(LiteralTestUtil::Equal(expected, result)); } -TEST_P(HloEvaluatorTest, NegativePadding2D) { +TEST_P(HloEvaluatorBf16Test, NegativePadding2D) { HloComputation::Builder b(TestName()); // input_array: @@ -574,7 +595,7 @@ TEST_P(HloEvaluatorTest, NegativePadding2D) { pad_value_instruction, r2_padding_on_dim0_dim1)); - module().AddEntryComputation(b.Build()); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate(); @@ -590,7 +611,7 @@ TEST_P(HloEvaluatorTest, NegativePadding2D) { EXPECT_TRUE(LiteralTestUtil::Near(expected, result, ErrorSpec(0.031250))); } -TEST_P(HloEvaluatorTest, NegativeAndInteriorPadding2D) { +TEST_P(HloEvaluatorBf16Test, NegativeAndInteriorPadding2D) { HloComputation::Builder b(TestName()); // f32[4,3] { @@ -619,7 +640,7 @@ TEST_P(HloEvaluatorTest, NegativeAndInteriorPadding2D) { pad_value_instruction, r2_padding_on_dim0_dim1)); - module().AddEntryComputation(b.Build()); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate(); @@ -629,7 +650,7 @@ TEST_P(HloEvaluatorTest, NegativeAndInteriorPadding2D) { EXPECT_TRUE(LiteralTestUtil::Equal(expected, result)); } -TEST_P(HloEvaluatorTest, DotRank2AndRank1) { +TEST_P(HloEvaluatorBf16Test, DotRank2AndRank1) { HloComputation::Builder b(TestName()); // lhs: @@ -658,7 +679,7 @@ TEST_P(HloEvaluatorTest, DotRank2AndRank1) { b.AddInstruction(HloInstruction::CreateDot(shape, lhs_instruction, rhs_instruction, dot_dnums, DefaultPrecisionConfig(2))); - module().AddEntryComputation(b.Build()); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate(); @@ -675,7 +696,7 @@ TEST_P(HloEvaluatorTest, DotRank2AndRank1) { EXPECT_TRUE(LiteralTestUtil::Equal(expected, result)); } -TEST_P(HloEvaluatorTest, DotRank1AndRank2) { +TEST_P(HloEvaluatorBf16Test, DotRank1AndRank2) { HloComputation::Builder b(TestName()); // lhs: @@ -704,7 +725,7 @@ TEST_P(HloEvaluatorTest, DotRank1AndRank2) { b.AddInstruction(HloInstruction::CreateDot(shape, lhs_instruction, rhs_instruction, dot_dnums, DefaultPrecisionConfig(2))); - module().AddEntryComputation(b.Build()); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate(); @@ -713,7 +734,7 @@ TEST_P(HloEvaluatorTest, DotRank1AndRank2) { EXPECT_TRUE(LiteralTestUtil::Equal(expected, result)); } -TEST_P(HloEvaluatorTest, DotRank2AndRank2) { +TEST_P(HloEvaluatorBf16Test, DotRank2AndRank2) { HloComputation::Builder b(TestName()); // lhs: @@ -748,7 +769,7 @@ TEST_P(HloEvaluatorTest, DotRank2AndRank2) { b.AddInstruction(HloInstruction::CreateDot(shape, lhs_instruction, rhs_instruction, dot_dnums, DefaultPrecisionConfig(2))); - module().AddEntryComputation(b.Build()); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate(); @@ -763,7 +784,51 @@ TEST_P(HloEvaluatorTest, DotRank2AndRank2) { EXPECT_TRUE(LiteralTestUtil::Equal(expected, result)); } -TEST_P(HloEvaluatorTest, SimpleConv1D) { +TEST_P(HloEvaluatorBf16Test, DotRank4AndRank4) { + HloComputation::Builder b(TestName()); + + auto lhs_array = absl::make_unique>(2, 2, 3, 1); + lhs_array->FillIota(1.0f); + auto lhs_literal = LiteralUtil::CreateR4FromArray4D(*lhs_array); + HloInstruction* lhs_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(lhs_literal))); + + auto rhs_array = absl::make_unique>(2, 2, 3, 1); + rhs_array->FillIota(2.0f); + auto rhs_literal = LiteralUtil::CreateR4FromArray4D(*rhs_array); + HloInstruction* rhs_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(rhs_literal))); + + Shape shape = ShapeUtil::MakeShape(F32, {2, 1, 1}); + DotDimensionNumbers dot_dnums; + + dot_dnums.add_lhs_batch_dimensions(0); + dot_dnums.add_rhs_batch_dimensions(0); + dot_dnums.add_lhs_contracting_dimensions(1); + dot_dnums.add_lhs_contracting_dimensions(2); + dot_dnums.add_rhs_contracting_dimensions(1); + dot_dnums.add_rhs_contracting_dimensions(2); + b.AddInstruction(HloInstruction::CreateDot(shape, lhs_instruction, + rhs_instruction, dot_dnums, + DefaultPrecisionConfig(2))); + m_->AddEntryComputation(b.Build()); + + Literal result = Evaluate(); + float expected_1 = 0; + for (float i = 1.0f; i < 7.0f; ++i) { + expected_1 += i * i + i; + } + float expected_2 = 0; + for (float i = 7.0f; i < 13.0f; ++i) { + expected_2 += i * i + i; + } + auto expected_array = Array3D({{{expected_1}}, {{expected_2}}}); + auto expected = LiteralUtil::CreateR3FromArray3D(expected_array); + + EXPECT_TRUE(LiteralTestUtil::Equal(expected, result)); +} + +TEST_P(HloEvaluatorBf16Test, SimpleConv1D) { HloComputation::Builder b(TestName()); Array3D lhs_array = {{{1, 2, 3}}}; @@ -801,8 +866,8 @@ TEST_P(HloEvaluatorTest, SimpleConv1D) { Shape shape = ShapeUtil::MakeShape(F32, {1, 1, 3}); b.AddInstruction(HloInstruction::CreateConvolve( shape, lhs_instruction, rhs_instruction, /*feature_group_count=*/1, - window, dnums, DefaultPrecisionConfig(2))); - module().AddEntryComputation(b.Build()); + /*batch_group_count=*/1, window, dnums, DefaultPrecisionConfig(2))); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate(); @@ -812,7 +877,7 @@ TEST_P(HloEvaluatorTest, SimpleConv1D) { EXPECT_TRUE(LiteralTestUtil::Equal(expected, result)); } -TEST_P(HloEvaluatorTest, Simple4x4Conv2DWith2x2Kernel) { +TEST_P(HloEvaluatorBf16Test, Simple4x4Conv2DWith2x2Kernel) { HloComputation::Builder b(TestName()); Array4D lhs_array(1, 1, 4, 4); @@ -856,8 +921,8 @@ TEST_P(HloEvaluatorTest, Simple4x4Conv2DWith2x2Kernel) { Shape shape = ShapeUtil::MakeShape(F32, {1, 1, 4, 4}); b.AddInstruction(HloInstruction::CreateConvolve( shape, lhs_instruction, rhs_instruction, /*feature_group_count=*/1, - window, dnums, DefaultPrecisionConfig(2))); - module().AddEntryComputation(b.Build()); + /*batch_group_count=*/1, window, dnums, DefaultPrecisionConfig(2))); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate(); @@ -875,7 +940,7 @@ TEST_P(HloEvaluatorTest, Simple4x4Conv2DWith2x2Kernel) { EXPECT_TRUE(LiteralTestUtil::Equal(expected, result)); } -TEST_P(HloEvaluatorTest, Conv2DGeneralDimensionsReversed) { +TEST_P(HloEvaluatorBf16Test, Conv2DGeneralDimensionsReversed) { HloComputation::Builder b(TestName()); // clang-format off @@ -940,8 +1005,8 @@ TEST_P(HloEvaluatorTest, Conv2DGeneralDimensionsReversed) { Shape shape = ShapeUtil::MakeShape(F32, {1, 1, 1, 2}); b.AddInstruction(HloInstruction::CreateConvolve( shape, lhs_instruction, rhs_instruction, /*feature_group_count=*/1, - window, dnums, DefaultPrecisionConfig(2))); - module().AddEntryComputation(b.Build()); + /*batch_group_count=*/1, window, dnums, DefaultPrecisionConfig(2))); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate(); @@ -956,7 +1021,7 @@ TEST_P(HloEvaluatorTest, Conv2DGeneralDimensionsReversed) { EXPECT_TRUE(LiteralTestUtil::Equal(expected, result)); } -TEST_P(HloEvaluatorTest, Conv2DGeneralDimensions) { +TEST_P(HloEvaluatorBf16Test, Conv2DGeneralDimensions) { HloComputation::Builder b(TestName()); // clang-format off @@ -1018,8 +1083,8 @@ TEST_P(HloEvaluatorTest, Conv2DGeneralDimensions) { Shape shape = ShapeUtil::MakeShape(F32, {1, 1, 1, 2}); b.AddInstruction(HloInstruction::CreateConvolve( shape, lhs_instruction, rhs_instruction, /*feature_group_count=*/1, - window, dnums, DefaultPrecisionConfig(2))); - module().AddEntryComputation(b.Build()); + /*batch_group_count=*/1, window, dnums, DefaultPrecisionConfig(2))); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate(); @@ -1034,7 +1099,7 @@ TEST_P(HloEvaluatorTest, Conv2DGeneralDimensions) { EXPECT_TRUE(LiteralTestUtil::Equal(expected, result)); } -TEST_P(HloEvaluatorTest, DilatedBaseConv2DWithHighPadding) { +TEST_P(HloEvaluatorBf16Test, DilatedBaseConv2DWithHighPadding) { HloComputation::Builder b(TestName()); Array4D lhs_array(1, 1, 4, 4); @@ -1078,8 +1143,8 @@ TEST_P(HloEvaluatorTest, DilatedBaseConv2DWithHighPadding) { Shape shape = ShapeUtil::MakeShape(F32, {1, 1, 7, 7}); b.AddInstruction(HloInstruction::CreateConvolve( shape, lhs_instruction, rhs_instruction, /*feature_group_count=*/1, - window, dnums, DefaultPrecisionConfig(2))); - module().AddEntryComputation(b.Build()); + /*batch_group_count=*/1, window, dnums, DefaultPrecisionConfig(2))); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate(); @@ -1098,7 +1163,7 @@ TEST_P(HloEvaluatorTest, DilatedBaseConv2DWithHighPadding) { EXPECT_TRUE(LiteralTestUtil::Equal(expected, result)); } -TEST_P(HloEvaluatorTest, DilatedBaseConv2DWithLowAndHighPadding) { +TEST_P(HloEvaluatorBf16Test, DilatedBaseConv2DWithLowAndHighPadding) { HloComputation::Builder b(TestName()); Array4D lhs_array(1, 1, 4, 4); @@ -1142,8 +1207,8 @@ TEST_P(HloEvaluatorTest, DilatedBaseConv2DWithLowAndHighPadding) { Shape shape = ShapeUtil::MakeShape(F32, {1, 1, 8, 8}); b.AddInstruction(HloInstruction::CreateConvolve( shape, lhs_instruction, rhs_instruction, /*feature_group_count=*/1, - window, dnums, DefaultPrecisionConfig(2))); - module().AddEntryComputation(b.Build()); + /*batch_group_count=*/1, window, dnums, DefaultPrecisionConfig(2))); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate(); @@ -1163,7 +1228,7 @@ TEST_P(HloEvaluatorTest, DilatedBaseConv2DWithLowAndHighPadding) { EXPECT_TRUE(LiteralTestUtil::Equal(expected, result)); } -TEST_P(HloEvaluatorTest, +TEST_P(HloEvaluatorBf16Test, DilatedWindowAndBaseConv2DWithDifferentLowAndHighPaddingAndStrides) { HloComputation::Builder b(TestName()); @@ -1214,8 +1279,8 @@ TEST_P(HloEvaluatorTest, Shape shape = ShapeUtil::MakeShape(F32, {1, 1, 9, 3}); b.AddInstruction(HloInstruction::CreateConvolve( shape, lhs_instruction, rhs_instruction, /*feature_group_count=*/1, - window, dnums, DefaultPrecisionConfig(2))); - module().AddEntryComputation(b.Build()); + /*batch_group_count=*/1, window, dnums, DefaultPrecisionConfig(2))); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate(); @@ -1236,7 +1301,7 @@ TEST_P(HloEvaluatorTest, EXPECT_TRUE(LiteralTestUtil::Equal(expected, result)); } -TEST_P(HloEvaluatorTest, Conv2DGroupedConvolution) { +TEST_P(HloEvaluatorBf16Test, Conv2DGroupedConvolution) { HloComputation::Builder b(TestName()); std::vector input_dims = {1, 2, 2, 4}; std::vector filter_dims = {2, 2, 2, 8}; @@ -1285,8 +1350,9 @@ TEST_P(HloEvaluatorTest, Conv2DGroupedConvolution) { Shape shape = ShapeUtil::MakeShape(F32, {1, 1, 1, 8}); b.AddInstruction(HloInstruction::CreateConvolve( shape, lhs_instruction, rhs_instruction, - /*feature_group_count=*/2, window, dnums, DefaultPrecisionConfig(2))); - module().AddEntryComputation(b.Build()); + /*feature_group_count=*/2, /*batch_group_count=*/1, window, dnums, + DefaultPrecisionConfig(2))); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate(); @@ -1297,11 +1363,12 @@ TEST_P(HloEvaluatorTest, Conv2DGroupedConvolution) { EXPECT_TRUE(LiteralTestUtil::Equal(expected, result)); } -class HloEvaluatorPreciseReduceTest : public HloVerifiedTestBase {}; +class HloEvaluatorPreciseReduceTest : public HloTestBase {}; // Tests that Reduce doesn't lose precision when adding many numbers (because // it accumulates its result in a double). TEST_F(HloEvaluatorPreciseReduceTest, AddReductionPrecisionTest) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder b(TestName()); constexpr int kNumElements = 1 << 25; // float += 1 saturates at 1<<24 @@ -1319,12 +1386,12 @@ TEST_F(HloEvaluatorPreciseReduceTest, AddReductionPrecisionTest) { HloInstruction::CreateParameter(1, scalar_shape, "rhs")); add_computation.AddInstruction(HloInstruction::CreateBinary( scalar_shape, HloOpcode::kAdd, param_lhs, param_rhs)); - auto add_func = module().AddEmbeddedComputation(add_computation.Build()); + auto add_func = m->AddEmbeddedComputation(add_computation.Build()); HloInstruction* reduce_instruction = b.AddInstruction( HloInstruction::CreateReduce(scalar_shape, arg_instruction, init_value, /*dimensions_to_reduce=*/{0}, add_func)); - module().AddEntryComputation(b.Build()); + m->AddEntryComputation(b.Build()); HloEvaluator hlo_eval; Literal result = hlo_eval.Evaluate(reduce_instruction).ConsumeValueOrDie(); @@ -1337,7 +1404,7 @@ void BM_ReducePrecisely(int num_iters) { tensorflow::testing::StopTiming(); HloComputation::Builder b("BM_ReducePrecisely"); HloModuleConfig config; - config.set_debug_options(legacy_flags::GetDebugOptionsFromFlags()); + config.set_debug_options(GetDebugOptionsFromFlags()); HloModule module("BM_ReducePrecisely", config); constexpr int kNumElements = 1 << 25; // float += 1 saturates at 1<<24 @@ -1370,7 +1437,7 @@ void BM_ReducePrecisely(int num_iters) { BENCHMARK(BM_ReducePrecisely); -TEST_P(HloEvaluatorTest, ReduceAdd) { +TEST_P(HloEvaluatorBf16Test, ReduceAdd) { HloComputation::Builder b(TestName()); // arg: @@ -1396,14 +1463,14 @@ TEST_P(HloEvaluatorTest, ReduceAdd) { HloInstruction::CreateParameter(1, scalar_shape, "rhs")); add_computation.AddInstruction(HloInstruction::CreateBinary( scalar_shape, HloOpcode::kAdd, param_lhs, param_rhs)); - auto add_func = module().AddEmbeddedComputation(add_computation.Build()); + auto add_func = m_->AddEmbeddedComputation(add_computation.Build()); Shape shape = ShapeUtil::MakeShape(F32, {2}); b.AddInstruction( HloInstruction::CreateReduce(shape, arg_instruction, init_value, /*dimensions_to_reduce=*/{1}, add_func)); - module().AddEntryComputation(b.Build()); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate(); @@ -1412,7 +1479,7 @@ TEST_P(HloEvaluatorTest, ReduceAdd) { EXPECT_TRUE(LiteralTestUtil::Equal(expected, result)); } -TEST_P(HloEvaluatorTest, ReduceWindowMax) { +TEST_P(HloEvaluatorBf16Test, ReduceWindowMax) { HloComputation::Builder b(TestName()); // arg: @@ -1438,7 +1505,7 @@ TEST_P(HloEvaluatorTest, ReduceWindowMax) { HloInstruction::CreateParameter(1, scalar_shape, "rhs")); max_computation.AddInstruction(HloInstruction::CreateBinary( scalar_shape, HloOpcode::kMaximum, param_lhs, param_rhs)); - auto max_func = module().AddEmbeddedComputation(max_computation.Build()); + auto max_func = m_->AddEmbeddedComputation(max_computation.Build()); Window window; WindowDimension dim; @@ -1455,7 +1522,7 @@ TEST_P(HloEvaluatorTest, ReduceWindowMax) { b.AddInstruction(HloInstruction::CreateReduceWindow( shape, arg_instruction, init_value, window, max_func)); - module().AddEntryComputation(b.Build()); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate(); @@ -1463,7 +1530,7 @@ TEST_P(HloEvaluatorTest, ReduceWindowMax) { EXPECT_TRUE(LiteralTestUtil::Equal(expected, result)); } -TEST_P(HloEvaluatorTest, ReduceWindowMaxWindowDilation) { +TEST_P(HloEvaluatorBf16Test, ReduceWindowMaxWindowDilation) { HloComputation::Builder b(TestName()); // arg: @@ -1490,7 +1557,7 @@ TEST_P(HloEvaluatorTest, ReduceWindowMaxWindowDilation) { HloInstruction::CreateParameter(1, scalar_shape, "rhs")); max_computation.AddInstruction(HloInstruction::CreateBinary( scalar_shape, HloOpcode::kMaximum, param_lhs, param_rhs)); - auto max_func = module().AddEmbeddedComputation(max_computation.Build()); + auto max_func = m_->AddEmbeddedComputation(max_computation.Build()); Window window; WindowDimension dim; @@ -1507,7 +1574,7 @@ TEST_P(HloEvaluatorTest, ReduceWindowMaxWindowDilation) { b.AddInstruction(HloInstruction::CreateReduceWindow( shape, arg_instruction, init_value, window, max_func)); - module().AddEntryComputation(b.Build()); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate(); @@ -1515,7 +1582,7 @@ TEST_P(HloEvaluatorTest, ReduceWindowMaxWindowDilation) { EXPECT_TRUE(LiteralTestUtil::Equal(expected, result)); } -TEST_P(HloEvaluatorTest, ReduceWindowAdd) { +TEST_P(HloEvaluatorBf16Test, ReduceWindowAdd) { HloComputation::Builder b(TestName()); // arg: @@ -1541,7 +1608,7 @@ TEST_P(HloEvaluatorTest, ReduceWindowAdd) { HloInstruction::CreateParameter(1, scalar_shape, "rhs")); add_computation.AddInstruction(HloInstruction::CreateBinary( scalar_shape, HloOpcode::kAdd, param_lhs, param_rhs)); - auto add_func = module().AddEmbeddedComputation(add_computation.Build()); + auto add_func = m_->AddEmbeddedComputation(add_computation.Build()); Window window; WindowDimension dim; @@ -1564,7 +1631,7 @@ TEST_P(HloEvaluatorTest, ReduceWindowAdd) { b.AddInstruction(HloInstruction::CreateReduceWindow( shape, arg_instruction, init_value, window, add_func)); - module().AddEntryComputation(b.Build()); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate(); @@ -1572,7 +1639,7 @@ TEST_P(HloEvaluatorTest, ReduceWindowAdd) { EXPECT_TRUE(LiteralTestUtil::Equal(expected, result)); } -TEST_P(HloEvaluatorTest, ReduceWindowAdd6D) { +TEST_P(HloEvaluatorBf16Test, ReduceWindowAdd6D) { HloComputation::Builder b(TestName()); // arg: f32[4,4,4,4,4,4] full of ones. Using small dims to limit run-time. @@ -1594,7 +1661,7 @@ TEST_P(HloEvaluatorTest, ReduceWindowAdd6D) { HloInstruction::CreateParameter(1, scalar_shape, "rhs")); add_computation.AddInstruction(HloInstruction::CreateBinary( scalar_shape, HloOpcode::kAdd, param_lhs, param_rhs)); - auto add_func = module().AddEmbeddedComputation(add_computation.Build()); + auto add_func = m_->AddEmbeddedComputation(add_computation.Build()); Window window; @@ -1625,7 +1692,7 @@ TEST_P(HloEvaluatorTest, ReduceWindowAdd6D) { b.AddInstruction(HloInstruction::CreateReduceWindow( shape, arg_instruction, init_value, window, add_func)); - module().AddEntryComputation(b.Build()); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate(); @@ -1635,7 +1702,7 @@ TEST_P(HloEvaluatorTest, ReduceWindowAdd6D) { EXPECT_TRUE(LiteralTestUtil::Equal(result_literal, result)); } -TEST_P(HloEvaluatorTest, StridedSlice) { +TEST_P(HloEvaluatorBf16Test, StridedSlice) { HloComputation::Builder b(TestName()); // arg: @@ -1657,7 +1724,7 @@ TEST_P(HloEvaluatorTest, StridedSlice) { /*start_indices=*/{0, 2}, /*limit_indices=*/{3, 5}, /*strides=*/{2, 3})); - module().AddEntryComputation(b.Build()); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate(); @@ -1669,7 +1736,7 @@ TEST_P(HloEvaluatorTest, StridedSlice) { EXPECT_TRUE(LiteralTestUtil::Equal(expected, result)); } -TEST_P(HloEvaluatorTest, DynamicSlice) { +TEST_P(HloEvaluatorBf16Test, DynamicSlice) { HloComputation::Builder b(TestName()); // arg: @@ -1685,13 +1752,15 @@ TEST_P(HloEvaluatorTest, 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})); - module().AddEntryComputation(b.Build()); + b.AddInstruction( + HloInstruction::CreateDynamicSlice(shape, operand, {zero, one}, {2, 3})); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate(); @@ -1705,7 +1774,7 @@ TEST_P(HloEvaluatorTest, DynamicSlice) { // Verifies that the HloEvaluator's implementation goes along with existing // backends' behavior, although this is not required by the spec. -TEST_P(HloEvaluatorTest, DynamicSliceModSlice) { +TEST_P(HloEvaluatorBf16Test, DynamicSliceModSlice) { HloComputation::Builder b(TestName()); // arg: @@ -1721,13 +1790,15 @@ TEST_P(HloEvaluatorTest, 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})); - module().AddEntryComputation(b.Build()); + b.AddInstruction( + HloInstruction::CreateDynamicSlice(shape, operand, {two, one}, {2, 3})); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate(); @@ -1739,7 +1810,7 @@ TEST_P(HloEvaluatorTest, DynamicSliceModSlice) { EXPECT_TRUE(LiteralTestUtil::Equal(expected, result)); } -TEST_P(HloEvaluatorTest, DynamicSliceUpdate) { +TEST_P(HloEvaluatorBf16Test, DynamicSliceUpdate) { HloComputation::Builder b(TestName()); // arg: @@ -1755,16 +1826,18 @@ TEST_P(HloEvaluatorTest, 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)); - module().AddEntryComputation(b.Build()); + shape, operand, update, {zero, one})); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate(); @@ -1776,7 +1849,7 @@ TEST_P(HloEvaluatorTest, DynamicSliceUpdate) { EXPECT_TRUE(LiteralTestUtil::Equal(expected, result)); } -TEST_P(HloEvaluatorTest, SetAndGetTuples) { +TEST_P(HloEvaluatorBf16Test, SetAndGetTuples) { HloComputation::Builder b(TestName()); // arg: @@ -1800,7 +1873,7 @@ TEST_P(HloEvaluatorTest, SetAndGetTuples) { Shape shape = ShapeUtil::MakeShape(F64, {2, 3}); b.AddInstruction(HloInstruction::CreateGetTupleElement(shape, tuple, 1)); - module().AddEntryComputation(b.Build()); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate(); @@ -1812,7 +1885,7 @@ TEST_P(HloEvaluatorTest, SetAndGetTuples) { EXPECT_TRUE(LiteralTestUtil::Equal(expected, result)); } -TEST_P(HloEvaluatorTest, SetAndGetNestedTuples) { +TEST_P(HloEvaluatorBf16Test, SetAndGetNestedTuples) { HloComputation::Builder b(TestName()); // arg: @@ -1839,7 +1912,7 @@ TEST_P(HloEvaluatorTest, SetAndGetNestedTuples) { b.AddInstruction( HloInstruction::CreateGetTupleElement(tuple2->shape(), outer_tuple, 1)); - module().AddEntryComputation(b.Build()); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate(); @@ -1851,7 +1924,7 @@ TEST_P(HloEvaluatorTest, SetAndGetNestedTuples) { EXPECT_TRUE(LiteralTestUtil::Equal(expected, result)); } -TEST_P(HloEvaluatorTest, Reverse) { +TEST_P(HloEvaluatorBf16Test, Reverse) { HloComputation::Builder b(TestName()); // Input shape is float[4x3x2x1]. @@ -1877,7 +1950,7 @@ TEST_P(HloEvaluatorTest, Reverse) { const Shape shape = ShapeUtil::MakeShape(F32, {4, 3, 2, 1}); b.AddInstruction(HloInstruction::CreateReverse(shape, operand, {0, 1})); - module().AddEntryComputation(b.Build()); + m_->AddEntryComputation(b.Build()); Literal result = Evaluate(); @@ -1904,7 +1977,7 @@ TEST_P(HloEvaluatorTest, Reverse) { EXPECT_TRUE(LiteralTestUtil::Equal(expected, result)); } -TEST_P(HloEvaluatorTest, EvaluateWithSubstitutions) { +TEST_P(HloEvaluatorBf16Test, EvaluateWithSubstitutions) { HloComputation::Builder b(TestName()); Shape shape = ShapeUtil::MakeShape(F32, {4}); @@ -1928,7 +2001,7 @@ TEST_P(HloEvaluatorTest, EvaluateWithSubstitutions) { // Check that EvaluateWithSubstitutions works if one of the operands to the op // we're evaluating is a constant. -TEST_P(HloEvaluatorTest, EvaluateWithSubstitutionsWithConstantOperand) { +TEST_P(HloEvaluatorBf16Test, EvaluateWithSubstitutionsWithConstantOperand) { HloComputation::Builder b(TestName()); Shape shape = ShapeUtil::MakeShape(F32, {4}); @@ -1951,7 +2024,7 @@ TEST_P(HloEvaluatorTest, EvaluateWithSubstitutionsWithConstantOperand) { LiteralUtil::CreateR1({11, 22, 33, 44}), result.ValueOrDie())); } -TEST_P(HloEvaluatorTest, EvaluateGather_TensorFlowGatherV1) { +TEST_F(HloEvaluatorTest, EvaluateGather_TensorFlowGatherV1) { const char* hlo_text = R"( HloModule TensorFlowGatherV1 @@ -1966,7 +2039,7 @@ ENTRY main { slice_sizes={1, 3} } )"; - ParseAndVerifyModule(hlo_text); + TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_text)); Literal operand = LiteralUtil::CreateR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}); Literal start_indices = LiteralUtil::CreateR1({0, 2}); @@ -1975,7 +2048,7 @@ ENTRY main { Evaluate({&operand, &start_indices}))); } -TEST_P(HloEvaluatorTest, EvaluateGather_TensorFlowGatherV2) { +TEST_F(HloEvaluatorTest, EvaluateGather_TensorFlowGatherV2) { const char* hlo_text = R"( HloModule TensorFlowGatherV2 @@ -1990,7 +2063,7 @@ ENTRY main { slice_sizes={3, 1} } )"; - ParseAndVerifyModule(hlo_text); + TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_text)); Literal operand = LiteralUtil::CreateR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}); Literal start_indices = LiteralUtil::CreateR1({0, 2}); @@ -1999,7 +2072,7 @@ ENTRY main { Evaluate({&operand, &start_indices}))); } -TEST_P(HloEvaluatorTest, EvaluateGather_TensorFlowGatherMultipleBatchDims) { +TEST_F(HloEvaluatorTest, EvaluateGather_TensorFlowGatherMultipleBatchDims) { const char* hlo_text = R"( HloModule TensorFlowGatherMultipleBatchDims @@ -2014,7 +2087,7 @@ ENTRY main { slice_sizes={3, 1} } )"; - ParseAndVerifyModule(hlo_text); + TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_text)); Literal operand = LiteralUtil::CreateR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}); Literal start_indices = LiteralUtil::CreateR2({{0, 2}, {2, 1}}); @@ -2024,7 +2097,7 @@ ENTRY main { Evaluate({&operand, &start_indices}))); } -TEST_P(HloEvaluatorTest, EvaluateGather_TensorFlowGatherNd) { +TEST_F(HloEvaluatorTest, EvaluateGather_TensorFlowGatherNd) { const char* hlo_text = R"( HloModule TensorFlowGatherNd @@ -2039,7 +2112,7 @@ ENTRY main { slice_sizes={1,1,2} } )"; - ParseAndVerifyModule(hlo_text); + TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_text)); Literal operand = LiteralUtil::CreateR3({{{-1, 1}, {-2, 2}, {-3, 3}}, // {{-4, 4}, {-5, 5}, {-6, 6}}, // @@ -2050,7 +2123,7 @@ ENTRY main { Evaluate({&operand, &start_indices}))); } -TEST_P(HloEvaluatorTest, +TEST_F(HloEvaluatorTest, EvaluateGather_TensorFlowGatherNdNonDefaultIndexVectorDim) { const char* hlo_text = R"( HloModule TensorFlowGatherNd @@ -2066,7 +2139,7 @@ ENTRY main { slice_sizes={1,1,2} } )"; - ParseAndVerifyModule(hlo_text); + TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_text)); Literal operand = LiteralUtil::CreateR3({{{-1, 1}, {-2, 2}, {-3, 3}}, // {{-4, 4}, {-5, 5}, {-6, 6}}, // @@ -2077,7 +2150,7 @@ ENTRY main { Evaluate({&operand, &start_indices}))); } -TEST_P(HloEvaluatorTest, EvaluateGather_DynamicSlice) { +TEST_F(HloEvaluatorTest, EvaluateGather_DynamicSlice) { const char* hlo_text = R"( HloModule DynamicSlice @@ -2092,7 +2165,7 @@ ENTRY main { slice_sizes={1,1} } )"; - ParseAndVerifyModule(hlo_text); + TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_text)); Literal operand = LiteralUtil::CreateR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}); Literal start_indices = LiteralUtil::CreateR1({1, 1}); @@ -2100,7 +2173,7 @@ ENTRY main { Evaluate({&operand, &start_indices}))); } -TEST_P(HloEvaluatorTest, EvaluateGather_BatchDynamicSlice) { +TEST_F(HloEvaluatorTest, EvaluateGather_BatchDynamicSlice) { const char* hlo_text = R"( HloModule BatchDynamicSlice @@ -2115,7 +2188,7 @@ ENTRY main { slice_sizes={1,1} } )"; - ParseAndVerifyModule(hlo_text); + TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_text)); Literal operand = LiteralUtil::CreateR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}); Literal start_indices = LiteralUtil::CreateR2({{2, 1}, {1, 1}}); @@ -2124,7 +2197,7 @@ ENTRY main { Evaluate({&operand, &start_indices}))); } -TEST_P(HloEvaluatorTest, EvaluateGather_ZeroDimBounds) { +TEST_F(HloEvaluatorTest, EvaluateGather_ZeroDimBounds) { const char* hlo_text = R"( HloModule TensorFlowGatherV1 @@ -2139,14 +2212,14 @@ ENTRY main { slice_sizes={1, 0} } )"; - ParseAndVerifyModule(hlo_text); + TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_text)); Literal operand = LiteralUtil::CreateR2({{}, {}, {}}); Literal start_indices = LiteralUtil::CreateR1({0, 2}); EXPECT_TRUE(LiteralTestUtil::Equal(LiteralUtil::CreateR2({{}, {}}), Evaluate({&operand, &start_indices}))); } -TEST_P(HloEvaluatorTest, EvaluateGather_NoOutputWindowDims) { +TEST_F(HloEvaluatorTest, EvaluateGather_NoOutputWindowDims) { const string hlo_text = R"( HloModule GatherXd @@ -2161,7 +2234,7 @@ ENTRY main { slice_sizes={1} } )"; - ParseAndVerifyModule(hlo_text); + TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_text)); Literal operand = LiteralUtil::CreateR1({0, 1, 2}); Literal start_indices = @@ -2171,7 +2244,7 @@ ENTRY main { Evaluate({&operand, &start_indices}))); } -TEST_P(HloEvaluatorTest, EvaluateScatter_TensorFlowScatterV1_Update) { +TEST_F(HloEvaluatorTest, EvaluateScatter_TensorFlowScatterV1_Update) { const char* hlo_text = R"( HloModule TensorFlowScatterV1 @@ -2192,7 +2265,7 @@ ENTRY main { index_vector_dim=1 } )"; - ParseAndVerifyModule(hlo_text); + TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_text)); Literal operand = LiteralUtil::CreateR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}); Literal scatter_indices = LiteralUtil::CreateR1({0, 2}); @@ -2202,7 +2275,7 @@ ENTRY main { Evaluate({&operand, &scatter_indices, &updates}))); } -TEST_P(HloEvaluatorTest, EvaluateScatter_TensorFlowScatterV2_Update) { +TEST_F(HloEvaluatorTest, EvaluateScatter_TensorFlowScatterV2_Update) { const char* hlo_text = R"( HloModule TensorFlowScatterV2 @@ -2223,7 +2296,7 @@ ENTRY main { index_vector_dim=1 } )"; - ParseAndVerifyModule(hlo_text); + TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_text)); Literal operand = LiteralUtil::CreateR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}); Literal scatter_indices = LiteralUtil::CreateR1({0, 2}); @@ -2234,7 +2307,7 @@ ENTRY main { Evaluate({&operand, &scatter_indices, &updates}))); } -TEST_P(HloEvaluatorTest, EvaluateScatter_TensorFlowScatter_Add) { +TEST_F(HloEvaluatorTest, EvaluateScatter_TensorFlowScatter_Add) { const char* hlo_text = R"( HloModule TensorFlowScatter @@ -2256,7 +2329,7 @@ ENTRY main { index_vector_dim=1 } )"; - ParseAndVerifyModule(hlo_text); + TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_text)); Literal operand = LiteralUtil::CreateR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}); Literal scatter_indices = LiteralUtil::CreateR1({0, 2}); @@ -2266,7 +2339,7 @@ ENTRY main { Evaluate({&operand, &scatter_indices, &updates}))); } -TEST_P(HloEvaluatorTest, EvaluateScatter_TensorFlowScatter_Mul) { +TEST_F(HloEvaluatorTest, EvaluateScatter_TensorFlowScatter_Mul) { const char* hlo_text = R"( HloModule TensorFlowScatter @@ -2288,7 +2361,7 @@ ENTRY main { index_vector_dim=1 } )"; - ParseAndVerifyModule(hlo_text); + TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_text)); Literal operand = LiteralUtil::CreateR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}); Literal scatter_indices = LiteralUtil::CreateR1({0, 2}); @@ -2298,7 +2371,7 @@ ENTRY main { Evaluate({&operand, &scatter_indices, &updates}))); } -TEST_P(HloEvaluatorTest, EvaluateScatter_TensorFlowScatter_F32) { +TEST_P(HloEvaluatorBf16Test, EvaluateScatter_TensorFlowScatter_F32) { const char* hlo_text = R"( HloModule TensorFlowScatter @@ -2320,7 +2393,7 @@ ENTRY main { index_vector_dim=1 } )"; - ParseAndVerifyModule(hlo_text); + TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_text)); Literal operand = LiteralUtil::CreateR2( {{1.1, 2.2, 3.3}, {4.4, 5.5, 6.6}, {7.7, 8.8, 9.9}}); Literal scatter_indices = LiteralUtil::CreateR1({2, 1}); @@ -2332,7 +2405,7 @@ ENTRY main { Evaluate({&operand, &scatter_indices, &updates}), ErrorSpec{0.1, 0.01})); } -TEST_P(HloEvaluatorTest, EvaluateScatter_TensorFlowScatter_RepeatedIndices) { +TEST_F(HloEvaluatorTest, EvaluateScatter_TensorFlowScatter_RepeatedIndices) { const char* hlo_text = R"( HloModule TensorFlowScatter @@ -2354,7 +2427,7 @@ ENTRY main { index_vector_dim=1 } )"; - ParseAndVerifyModule(hlo_text); + TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_text)); Literal operand = LiteralUtil::CreateR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}); Literal scatter_indices = LiteralUtil::CreateR1({1, 1}); @@ -2364,7 +2437,7 @@ ENTRY main { Evaluate({&operand, &scatter_indices, &updates}))); } -TEST_P(HloEvaluatorTest, EvaluateScatter_TensorFlowScatter_MultipleBatchDims) { +TEST_F(HloEvaluatorTest, EvaluateScatter_TensorFlowScatter_MultipleBatchDims) { const char* hlo_text = R"( HloModule TensorFlowScatterMultipleBatchDims @@ -2386,7 +2459,7 @@ ENTRY main { index_vector_dim=2 } )"; - ParseAndVerifyModule(hlo_text); + TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_text)); Literal operand = LiteralUtil::CreateR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}); Literal scatter_indices = LiteralUtil::CreateR2({{0, 2}, {2, 1}}); @@ -2397,7 +2470,7 @@ ENTRY main { Evaluate({&operand, &scatter_indices, &updates}))); } -TEST_P(HloEvaluatorTest, EvaluateScatter_TensorFlowScatterNd) { +TEST_F(HloEvaluatorTest, EvaluateScatter_TensorFlowScatterNd) { const char* hlo_text = R"( HloModule TensorFlowScatterNd @@ -2418,7 +2491,7 @@ ENTRY main { index_vector_dim=1 } )"; - ParseAndVerifyModule(hlo_text); + TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_text)); Literal operand = LiteralUtil::CreateR3({{{-1, 1}, {-2, 2}, {-3, 3}}, // {{-4, 4}, {-5, 5}, {-6, 6}}, // @@ -2433,7 +2506,7 @@ ENTRY main { expected, Evaluate({&operand, &scatter_indices, &updates}))); } -TEST_P(HloEvaluatorTest, +TEST_F(HloEvaluatorTest, EvaluateScatter_TensorFlowScatterNd_NonDefaultIndexVectorDim) { const char* hlo_text = R"( HloModule TensorFlowScatterNdNonDefaultIndexVectorDim @@ -2455,7 +2528,7 @@ ENTRY main { index_vector_dim=0 } )"; - ParseAndVerifyModule(hlo_text); + TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_text)); Literal operand = LiteralUtil::CreateR3({{{-1, 1}, {-2, 2}, {-3, 3}}, // {{-4, 4}, {-5, 5}, {-6, 6}}, // @@ -2470,7 +2543,7 @@ ENTRY main { expected, Evaluate({&operand, &scatter_indices, &updates}))); } -TEST_P(HloEvaluatorTest, EvaluateScatter_DynamicUpdateSlice) { +TEST_F(HloEvaluatorTest, EvaluateScatter_DynamicUpdateSlice) { const char* hlo_text = R"( HloModule DynamicUpdateSlice @@ -2491,7 +2564,7 @@ ENTRY main { index_vector_dim=0 } )"; - ParseAndVerifyModule(hlo_text); + TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_text)); Literal operand = LiteralUtil::CreateR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}); Literal scatter_indices = LiteralUtil::CreateR1({1, 1}); @@ -2502,7 +2575,7 @@ ENTRY main { expected, Evaluate({&operand, &scatter_indices, &updates}))); } -TEST_P(HloEvaluatorTest, EvaluateScatter_BatchDynamicUpdateSlice) { +TEST_F(HloEvaluatorTest, EvaluateScatter_BatchDynamicUpdateSlice) { const char* hlo_text = R"( HloModule BatchDynamicUpdateSlice @@ -2523,7 +2596,7 @@ ENTRY main { index_vector_dim=0 } )"; - ParseAndVerifyModule(hlo_text); + TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_text)); Literal operand = LiteralUtil::CreateR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}); Literal scatter_indices = LiteralUtil::CreateR2({{2, 1}, {1, 1}}); @@ -2534,7 +2607,7 @@ ENTRY main { expected, Evaluate({&operand, &scatter_indices, &updates}))); } -TEST_P(HloEvaluatorTest, EvaluateScatter_ZeroDimBounds) { +TEST_F(HloEvaluatorTest, EvaluateScatter_ZeroDimBounds) { const char* hlo_text = R"( HloModule TensorFlowScatter_ZeroDimBounds @@ -2555,7 +2628,7 @@ ENTRY main { index_vector_dim=1 } )"; - ParseAndVerifyModule(hlo_text); + TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_text)); Literal operand = LiteralUtil::CreateR2({{}, {}, {}}); Literal scatter_indices = LiteralUtil::CreateR1({0, 2}); Literal updates = LiteralUtil::CreateR2({{}, {}}); @@ -2563,7 +2636,7 @@ ENTRY main { operand, Evaluate({&operand, &scatter_indices, &updates}))); } -TEST_P(HloEvaluatorTest, EvaluateScatter_NoUpdateWindowDims) { +TEST_F(HloEvaluatorTest, EvaluateScatter_NoUpdateWindowDims) { const string hlo_text = R"( HloModule Scatter_NoUpdateWindowDims @@ -2585,7 +2658,7 @@ ENTRY main { index_vector_dim=2 } )"; - ParseAndVerifyModule(hlo_text); + TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_text)); Literal operand = LiteralUtil::CreateR1({0, 1, 2}); Literal scatter_indices = @@ -2596,7 +2669,7 @@ ENTRY main { expected, Evaluate({&operand, &scatter_indices, &updates}))); } -TEST_P(HloEvaluatorTest, EvaluateScatter_NegativeIndices) { +TEST_F(HloEvaluatorTest, EvaluateScatter_NegativeIndices) { const char* hlo_text = R"( HloModule TensorFlowScatter_NegativeIndices @@ -2631,7 +2704,7 @@ ENTRY main { {&operand, &scatter_indices, &updates}))); } -TEST_P(HloEvaluatorTest, EvaluateScatter_OobIndices) { +TEST_F(HloEvaluatorTest, EvaluateScatter_OobIndices) { const string hlo_text = R"( HloModule BatchDynamicUpdateSlice @@ -2667,7 +2740,7 @@ ENTRY main { {&operand, &scatter_indices, &updates}))); } -TEST_P(HloEvaluatorTest, EvaluateScatter_OobUpdateWindow) { +TEST_F(HloEvaluatorTest, EvaluateScatter_OobUpdateWindow) { const char* hlo_text = R"( HloModule TensorFlowScatterNd_OobUpdateWindow @@ -2706,7 +2779,7 @@ ENTRY main { // Verifies that HloEvaluator evaluates a HLO instruction that performs // element-wise comparison with 2 bfloat16 operands. -TEST_P(HloEvaluatorTest, DoesCompareBF16) { +TEST_F(HloEvaluatorTest, DoesCompareBF16) { // lhs >= rhs auto lhs = LiteralUtil::CreateR2( {{bfloat16(0.25), bfloat16(0.35), bfloat16(0.125)}, @@ -2720,7 +2793,7 @@ TEST_P(HloEvaluatorTest, DoesCompareBF16) { std::move(rhs)); } -TEST_P(HloEvaluatorTest, Bf16Reduction) { +TEST_P(HloEvaluatorBf16Test, Bf16Reduction) { const string hlo_text = R"( HloModule Bf16Reduction @@ -2736,7 +2809,7 @@ ENTRY main { ROOT %reduce = bf16[] reduce(arg0, init), dimensions={0}, to_apply=add_bf16 } )"; - ParseAndVerifyModule(hlo_text); + TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_text)); Literal arg = LiteralUtil::CreateR1( {bfloat16(1.0f), bfloat16(3.0f), bfloat16(-2.0f), bfloat16(42.0f)}); @@ -2744,7 +2817,7 @@ ENTRY main { EXPECT_TRUE(LiteralTestUtil::Equal(expected, Evaluate({&arg}))); } -TEST_P(HloEvaluatorTest, SliceWithDifferentLayout) { +TEST_P(HloEvaluatorBf16Test, SliceWithDifferentLayout) { // Regression test for b/114735354. const string hlo_text = R"( HloModule SliceWithDifferentLayout @@ -2754,7 +2827,7 @@ ENTRY main { ROOT %slice = f32[2,2,2]{1,0,2} slice(f32[2,2,2]{0,1,2} %arg), slice={[0:2], [0:2], [0:2]} } )"; - ParseAndVerifyModule(hlo_text); + TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_text)); Literal arg = LiteralUtil::CreateR3WithLayout( {{{1.0f, 2.0f}, {3.0f, 4.0f}}, {{5.0f, 6.0f}, {7.0f, 8.0f}}}, @@ -2763,8 +2836,210 @@ ENTRY main { EXPECT_TRUE(LiteralTestUtil::Equal(arg, actual)); } -INSTANTIATE_TEST_CASE_P(HloEvaluatorTest_Instantiation, HloEvaluatorTest, - ::testing::ValuesIn(use_bf16_params)); +TEST_P(HloEvaluatorBf16Test, Bitcast) { + // Regression test for b/114735354. + constexpr absl::string_view hlo_text_base = R"( +HloModule Bitcast + +ENTRY main { + param = %s[32,121]{1,0} parameter(0) + ROOT bitcast = %s[121,32,1]{0,1,2} bitcast(%s[32,121]{1,0} param) +} +)"; + string hlo_text; + if (use_bfloat16_) { + hlo_text = absl::StrFormat(hlo_text_base, "bf16", "bf16", "bf16"); + } else { + hlo_text = absl::StrFormat(hlo_text_base, "f32", "f32", "f32"); + } + TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_text)); + auto args = MakeFakeArguments(m_.get()).ConsumeValueOrDie(); + Literal actual = Evaluate({&args[0]}); + if (use_bfloat16_) { + EXPECT_TRUE( + absl::c_equal(args[0].data(), actual.data())); + } else { + EXPECT_TRUE(absl::c_equal(args[0].data(), actual.data())); + } +} + +// Check that s32 under/overflow doesn't trigger a ubsan failure. +TEST_F(HloEvaluatorTest, Int32Overflow) { + constexpr absl::string_view hlo_text = R"( +HloModule Test + +ENTRY main { + c1 = s32[] constant(1073741824) // 2^30 + sum = s32[] add(c1, c1) // 2^31, i.e. INT_MIN + + c2 = s32[] constant(-2147483648) // -2^31 + sub = s32[] subtract(c2, c1) // -2^31 - 2^30, underflows + + mul = s32[] multiply(c1, c1) + ROOT tuple = (s32[], s32[], s32[]) tuple(sum, sub, mul) +} +)"; + TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_text)); + std::vector actual = Evaluate({}).DecomposeTuple(); + ASSERT_EQ(actual.size(), 3); + + uint32 pow30 = uint32{1} << 30; + uint32 pow31 = uint32{1} << 31; + EXPECT_EQ(actual[0].GetFirstElement(), static_cast(pow31)); + EXPECT_EQ(actual[1].GetFirstElement(), + static_cast(-(pow31 + pow30))); + EXPECT_EQ(actual[2].GetFirstElement(), + 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"( +HloModule Test + +ENTRY main { + p0 = s32[1] parameter(0) + ROOT sum = s32[1] add(p0, p0) +} +)"; + TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_text)); + Literal input_wrong_shape = LiteralUtil::CreateR1({0, 1}); + + EXPECT_EQ(HloEvaluator() + .Evaluate(*m_, {&input_wrong_shape}) + .status() + .error_message(), + "Shape mismatch at parameter 0. Computation expected s32[1]{0}, " + "but arg was s32[2]."); + EXPECT_EQ(HloEvaluator() + .Evaluate(*m_->entry_computation(), {&input_wrong_shape}) + .status() + .error_message(), + "Shape mismatch at parameter 0. Computation expected s32[1]{0}, " + "but arg was s32[2]."); +} + +// Check that we get a useful error if we pass too many or too few inputs. +TEST_F(HloEvaluatorTest, EvaluateWithWrongNumberOfInputs) { + constexpr absl::string_view hlo_text = R"( +HloModule Test + +ENTRY main { + p0 = s32[1] parameter(0) + ROOT sum = s32[1] add(p0, p0) +} +)"; + TF_ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_text)); + Literal input = LiteralUtil::CreateR1({0}); + + EXPECT_EQ( + HloEvaluator().Evaluate(*m_, {&input, &input}).status().error_message(), + "Expected 1 argument, but got 2."); + EXPECT_EQ(HloEvaluator() + .Evaluate(*m_->entry_computation(), {&input, &input}) + .status() + .error_message(), + "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())); +} } // 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 84fbbd3e0c3ddb704b8db601897f3b199dc99626..648c7d0e676cd85ea255557bd969d92659aeeca7 100644 --- a/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor.h +++ b/tensorflow/compiler/xla/service/hlo_evaluator_typed_visitor.h @@ -19,15 +19,17 @@ limitations under the License. #include #include "absl/algorithm/container.h" +#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" #include "tensorflow/compiler/xla/service/hlo_casting_utils.h" #include "tensorflow/compiler/xla/service/hlo_evaluator.h" #include "tensorflow/compiler/xla/service/hlo_instructions.h" #include "tensorflow/compiler/xla/service/shape_inference.h" -#include "tensorflow/core/lib/core/casts.h" namespace xla { @@ -38,9 +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; +using is_complex_t = + absl::disjunction, 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 @@ -82,6 +83,26 @@ bool SafeLess(const NativeT& a, const NativeT& b) { return SafeLess(static_cast(a), static_cast(b)); } +// ToArithmeticSafeType(T t): +// - converts `t` to the bitwise-equivalent `unsigned T` if T is a signed +// integer, and +// - otherwise returns `t` unchanged. +// +// It's UB in C++ to under/overflow a signed integer, so we wrap all arithmetic +// in this type to force 2's complement behavior. +template ::value && + std::is_signed::value>::type* = nullptr> +typename std::make_unsigned::type ToArithmeticSafeType(T t) { + return static_cast::type>(t); +} +template ::value || + !std::is_signed::value>::type* = nullptr> +T ToArithmeticSafeType(T t) { + return std::move(t); +} + // Templated DfsHloVisitor for use by HloEvaluator. // // Typically ReturnT here indicates the resulting literal type of each evaluated @@ -105,6 +126,12 @@ bool SafeLess(const NativeT& a, const NativeT& b) { template class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { private: + Status UnsupportedTypeError(HloInstruction* instruction) { + return InvalidArgument( + "Unsupported type for %s: %s", HloOpcodeString(instruction->opcode()), + PrimitiveType_Name(instruction->shape().element_type())); + } + // Get the value in the given literal static_cast as a double. template < typename NativeT, @@ -161,9 +188,6 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { HloOpcodeString(hlo_instruction->opcode())); } - // TODO(b/35950897): many of the stl functions used in the handlers are not - // overloaded for every XLA primitive type. - template ::value>::type* = nullptr> @@ -188,7 +212,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)); @@ -207,6 +231,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); } @@ -227,7 +253,7 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { typename NativeT, typename std::enable_if::value>::type* = nullptr> Status HandleRound(HloInstruction* round) { - return InvalidArgument("Unsupported type for Round"); + return UnsupportedTypeError(round); } Status HandleRound(HloInstruction* round) override { @@ -249,7 +275,7 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { typename NativeT, typename std::enable_if::value>::type* = nullptr> Status HandleCeil(HloInstruction* ceil) { - return InvalidArgument("Unsupported type for Ceil"); + return UnsupportedTypeError(ceil); } Status HandleCeil(HloInstruction* ceil) override { @@ -300,8 +326,8 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { template < typename NativeT, typename std::enable_if::value>::type* = nullptr> - Status HandleExpm1(HloInstruction* floor) { - return InvalidArgument("Unsupported type for Expm1"); + Status HandleExpm1(HloInstruction* expm1) { + return UnsupportedTypeError(expm1); } Status HandleExpm1(HloInstruction* floor) override { @@ -324,7 +350,7 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { typename NativeT, typename std::enable_if::value>::type* = nullptr> Status HandleFloor(HloInstruction* floor) { - return InvalidArgument("Unsupported type for Floor"); + return UnsupportedTypeError(floor); } Status HandleFloor(HloInstruction* floor) override { @@ -354,12 +380,12 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { template < typename NativeT, typename std::enable_if::value>::type* = nullptr> - Status HandleLog1p(HloInstruction* floor) { - return InvalidArgument("Unsupported type for Log1p"); + Status HandleLog1p(HloInstruction* log1p) { + return UnsupportedTypeError(log1p); } - Status HandleLog1p(HloInstruction* floor) override { - return HandleLog1p(floor); + Status HandleLog1p(HloInstruction* log1p) override { + return HandleLog1p(log1p); } template ::value>::type* = nullptr> Status HandleNot(HloInstruction* not_) { - return InvalidArgument("Unsupported type for Not"); + return UnsupportedTypeError(not_); } Status HandleNot(HloInstruction* not_) override { @@ -479,7 +505,7 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { template ::value>::type* = nullptr> Status HandleAtan2(HloInstruction* atan2) { - return InvalidArgument("Unsupported type for Atan2"); + return UnsupportedTypeError(atan2); } Status HandleAtan2(HloInstruction* atan2) override { @@ -494,47 +520,25 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { return Status::OK(); } - template ::value && - !std::is_floating_point::value>::type* = nullptr> - Status HandleMultiply(HloInstruction* multiply) { - using type = typename std::make_unsigned::type; - TF_ASSIGN_OR_RETURN( - parent_->evaluated_[multiply], - ElementWiseBinaryOp(multiply, - [](ElementwiseT lhs_elem, ElementwiseT rhs_elem) { - return NativeT(type(lhs_elem) * type(rhs_elem)); - })); - return Status::OK(); - } - - template < - typename NativeT, - typename std::enable_if::value || - std::is_floating_point::value || - is_complex_t::value>::type* = nullptr> - Status HandleMultiply(HloInstruction* multiply) { + Status HandleMultiply(HloInstruction* multiply) override { TF_ASSIGN_OR_RETURN( parent_->evaluated_[multiply], - ElementWiseBinaryOp(multiply, - [](ElementwiseT lhs_elem, ElementwiseT rhs_elem) { - return lhs_elem * rhs_elem; - })); + ElementWiseBinaryOp( + multiply, [](ElementwiseT lhs_elem, ElementwiseT rhs_elem) { + return ElementwiseT(ToArithmeticSafeType(lhs_elem) * + ToArithmeticSafeType(rhs_elem)); + })); return Status::OK(); } - Status HandleMultiply(HloInstruction* multiply) override { - return HandleMultiply(multiply); - } - Status HandleSubtract(HloInstruction* subtract) override { TF_ASSIGN_OR_RETURN( parent_->evaluated_[subtract], - ElementWiseBinaryOp(subtract, - [](ElementwiseT lhs_elem, ElementwiseT rhs_elem) { - return lhs_elem - rhs_elem; - })); + ElementWiseBinaryOp( + subtract, [](ElementwiseT lhs_elem, ElementwiseT rhs_elem) { + return ElementwiseT(ToArithmeticSafeType(lhs_elem) - + ToArithmeticSafeType(rhs_elem)); + })); return Status::OK(); } @@ -542,7 +546,8 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { TF_ASSIGN_OR_RETURN(parent_->evaluated_[add], ElementWiseBinaryOp(add, [](ElementwiseT lhs_elem, ElementwiseT rhs_elem) { - return lhs_elem + rhs_elem; + return ElementwiseT(ToArithmeticSafeType(lhs_elem) + + ToArithmeticSafeType(rhs_elem)); })); return Status::OK(); } @@ -596,7 +601,7 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { return Status::OK(); } - Status HandleDivide(HloInstruction* divide) { + Status HandleDivide(HloInstruction* divide) override { return HandleDivide(divide); } @@ -627,7 +632,7 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { typename NativeT, typename std::enable_if::value>::type* = nullptr> Status HandleMaximum(HloInstruction* maximum) { - return InvalidArgument("Unsupported type for Maximum"); + return UnsupportedTypeError(maximum); } Status HandleMaximum(HloInstruction* maximum) override { @@ -662,7 +667,7 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { typename NativeT, typename std::enable_if::value>::type* = nullptr> Status HandleMinimum(HloInstruction* minimum) { - return InvalidArgument("Unsupported type for Minimum"); + return UnsupportedTypeError(minimum); } Status HandleMinimum(HloInstruction* minimum) override { @@ -670,11 +675,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(); } @@ -727,7 +735,7 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { typename NativeT, typename std::enable_if::value>::type* = nullptr> Status HandleRemainder(HloInstruction* remainder) { - return InvalidArgument("Unsupported type for Remainder"); + return UnsupportedTypeError(remainder); } Status HandleRemainder(HloInstruction* remainder) override { @@ -749,14 +757,14 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { template ::value>::type* = nullptr> Status HandleAnd(HloInstruction* and_) { - return InvalidArgument("Unsupported type for And"); + return UnsupportedTypeError(and_); } template < typename NativeT, typename std::enable_if::value>::type* = nullptr> Status HandleAnd(HloInstruction* and_) { - return InvalidArgument("Unsupported type for And"); + return UnsupportedTypeError(and_); } Status HandleAnd(HloInstruction* and_) override { @@ -778,7 +786,7 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { template ::value>::type* = nullptr> Status HandleOr(HloInstruction* or_) { - return InvalidArgument("Unsupported type for Or"); + return UnsupportedTypeError(or_); } template < @@ -807,14 +815,14 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { template ::value>::type* = nullptr> Status HandleXor(HloInstruction* xor_) { - return InvalidArgument("Unsupported type for Xor"); + return UnsupportedTypeError(xor_); } template < typename NativeT, typename std::enable_if::value>::type* = nullptr> Status HandleXor(HloInstruction* xor_) { - return InvalidArgument("Unsupported type for Xor"); + return UnsupportedTypeError(xor_); } Status HandleXor(HloInstruction* xor_) override { @@ -839,8 +847,8 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { typename std::enable_if::value || std::is_same::value>::type* = nullptr> - Status HandleShiftLeft(HloInstruction*) { - return InvalidArgument("Unsupported type for ShiftLeft"); + Status HandleShiftLeft(HloInstruction* shift) { + return UnsupportedTypeError(shift); } Status HandleShiftLeft(HloInstruction* shl) override { @@ -869,8 +877,8 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { typename std::enable_if::value || std::is_same::value>::type* = nullptr> - Status HandleShiftRightArithmetic(HloInstruction*) { - return InvalidArgument("Unsupported type for ShiftRightArithmetic"); + Status HandleShiftRightArithmetic(HloInstruction* shift) { + return UnsupportedTypeError(shift); } Status HandleShiftRightArithmetic(HloInstruction* shra) override { @@ -900,8 +908,8 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { typename std::enable_if::value || std::is_same::value>::type* = nullptr> - Status HandleShiftRightLogical(HloInstruction*) { - return InvalidArgument("Unsupported type for ShiftRightLogical"); + Status HandleShiftRightLogical(HloInstruction* shift) { + return UnsupportedTypeError(shift); } Status HandleShiftRightLogical(HloInstruction* shrl) override { @@ -914,7 +922,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], @@ -926,8 +938,8 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { template < typename NativeT, typename std::enable_if::value>::type* = nullptr> - Status HandleClamp(HloInstruction*) { - return InvalidArgument("Unsupported type for Clamp"); + Status HandleClamp(HloInstruction* clamp) { + return UnsupportedTypeError(clamp); } Status HandleClamp(HloInstruction* clamp) override { @@ -936,7 +948,7 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { Status HandleSelect(HloInstruction* select) override { CHECK(!ShapeUtil::IsScalar(select->operand(0)->shape())); - CHECK(ShapeUtil::IsArray(select->shape())); + CHECK(select->shape().IsArray()); std::function select_op = [](bool pred, ReturnT on_true, ReturnT on_false) { if (pred) { @@ -989,8 +1001,8 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { TF_CHECK_OK(ShapeUtil::ValidateShape(lhs_shape)); TF_CHECK_OK(ShapeUtil::ValidateShape(rhs_shape)); - CHECK(ShapeUtil::IsArray(lhs_shape)); - CHECK(ShapeUtil::IsArray(rhs_shape)); + CHECK(lhs_shape.IsArray()); + CHECK(rhs_shape.IsArray()); CHECK(ShapeUtil::SameElementType(lhs_shape, rhs_shape)); CHECK(ShapeUtil::SameElementType(lhs_shape, result_shape)); @@ -1001,16 +1013,16 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { CHECK_GE(num_spatial_dims, 0); CHECK_EQ(window.dimensions_size(), num_spatial_dims); - const auto lhs_rank = ShapeUtil::Rank(lhs_shape); - const auto rhs_rank = ShapeUtil::Rank(rhs_shape); + const auto lhs_rank = lhs_shape.rank(); + const auto rhs_rank = rhs_shape.rank(); CHECK_EQ(num_spatial_dims + 2, lhs_rank); CHECK_EQ(num_spatial_dims + 2, rhs_rank); - TF_ASSIGN_OR_RETURN( - auto inferred_return_shape, - ShapeInference::InferConvolveShape( - lhs_shape, rhs_shape, conv->feature_group_count(), window, dnums)); + TF_ASSIGN_OR_RETURN(auto inferred_return_shape, + ShapeInference::InferConvolveShape( + lhs_shape, rhs_shape, conv->feature_group_count(), + conv->batch_group_count(), window, dnums)); CHECK(ShapeUtil::Compatible(result_shape, inferred_return_shape)) << "return shape set to: " << ShapeUtil::HumanString(result_shape) << " but is inferred to be: " @@ -1033,12 +1045,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(); + const int64 feature_group_count = conv->feature_group_count(); + const int64 batch_group_count = conv->batch_group_count(); 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(); @@ -1051,6 +1064,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; @@ -1066,11 +1085,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; @@ -1123,11 +1146,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]; @@ -1151,23 +1187,31 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { } Status HandleDot(HloInstruction* dot) override { - auto lhs = dot->operand(0); - auto rhs = dot->operand(1); - CHECK(ShapeUtil::IsArray(dot->shape())); - CHECK(ShapeUtil::IsArray(lhs->shape())); - CHECK(ShapeUtil::IsArray(rhs->shape())); + if (dot->dot_dimension_numbers().rhs_contracting_dimensions_size() == 1 && + parent_->use_fast_path_) { + return HandleDot(dot); + } + return HandleDotSlowPath(dot); + } + + template ::value>::type* = nullptr> + Status HandleDot(HloInstruction* dot) { + const HloInstruction* lhs = dot->operand(0); + const HloInstruction* rhs = dot->operand(1); + CHECK(dot->shape().IsArray()); + CHECK(lhs->shape().IsArray()); + CHECK(rhs->shape().IsArray()); const auto& dnums = dot->dot_dimension_numbers(); - const auto lhs_rank = ShapeUtil::Rank(lhs->shape()); - const auto rhs_rank = ShapeUtil::Rank(rhs->shape()); + const int64 lhs_rank = lhs->shape().rank(); + const int64 rhs_rank = rhs->shape().rank(); CHECK(ShapeUtil::SameElementType(lhs->shape(), rhs->shape())); CHECK(ShapeUtil::SameElementType(lhs->shape(), dot->shape())); // There must be 1 and only 1 Contracting dimension for lhs and rhs. - CHECK_EQ(dnums.lhs_contracting_dimensions_size(), 1); - CHECK_EQ(dnums.rhs_contracting_dimensions_size(), 1); const int64 lhs_contracting_dimension = dnums.lhs_contracting_dimensions(0); const int64 rhs_contracting_dimension = dnums.rhs_contracting_dimensions(0); // Contracted dimension sizes must be the same. @@ -1177,8 +1221,56 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { << lhs->shape().dimensions(lhs_contracting_dimension) << " rhs contracted dimension: " << rhs->shape().dimensions(rhs_contracting_dimension); - const int64 contracted_dimension_size = - lhs->shape().dimensions(lhs_contracting_dimension); + + // The fast path is for a simple rank 2 dot with default layout operands. + if (lhs_rank == 2 && rhs_rank == 2 && lhs_contracting_dimension == 1 && + rhs_contracting_dimension == 0 && + LayoutUtil::Equal(lhs->shape().layout(), + LayoutUtil::GetDefaultLayoutForR2()) && + LayoutUtil::Equal(rhs->shape().layout(), + LayoutUtil::GetDefaultLayoutForR2()) && + LayoutUtil::Equal(dot->shape().layout(), + LayoutUtil::GetDefaultLayoutForR2())) { + const Literal& lhs_literal = parent_->GetEvaluatedLiteralFor(lhs); + const Literal& rhs_literal = parent_->GetEvaluatedLiteralFor(rhs); + const int64 contracted_dimension_size = + lhs->shape().dimensions(lhs_contracting_dimension); + Array2D lhs_array(lhs->shape().dimensions(0), + contracted_dimension_size); + lhs_array.SetValues(lhs_literal.data()); + Array2D rhs_array(contracted_dimension_size, + rhs->shape().dimensions(1)); + rhs_array.SetValues(rhs_literal.data()); + std::unique_ptr> result_array = + HloEvaluator::MatmulArray2D(lhs_array, rhs_array); + Literal result(dot->shape()); + result.PopulateR2FromArray2D(*result_array); + parent_->evaluated_[dot] = std::move(result); + return Status::OK(); + } + return HandleDotSlowPath(dot); + } + + template ::value>::type* = nullptr> + Status HandleDot(HloInstruction* dot) { + return HandleDotSlowPath(dot); + } + + Status HandleDotSlowPath(HloInstruction* dot) { + auto lhs = dot->operand(0); + auto rhs = dot->operand(1); + CHECK(dot->shape().IsArray()); + CHECK(lhs->shape().IsArray()); + CHECK(rhs->shape().IsArray()); + + const auto& dnums = dot->dot_dimension_numbers(); + + const auto lhs_rank = lhs->shape().rank(); + const auto rhs_rank = rhs->shape().rank(); + + CHECK(ShapeUtil::SameElementType(lhs->shape(), rhs->shape())); + CHECK(ShapeUtil::SameElementType(lhs->shape(), dot->shape())); const Literal& lhs_literal = parent_->GetEvaluatedLiteralFor(lhs); const Literal& rhs_literal = parent_->GetEvaluatedLiteralFor(rhs); @@ -1193,7 +1285,9 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { // in lhs_index or rhs_index where the i'th result index should go. absl::InlinedVector, kInlineRank> result_index_locations; - result_index_locations.reserve(lhs_rank + rhs_rank - 2); + result_index_locations.reserve( + (lhs_rank - dnums.lhs_contracting_dimensions_size()) + + (rhs_rank - dnums.rhs_contracting_dimensions_size())); // The first components in the output shape are the LHS and RHS batch // dimensions: @@ -1205,18 +1299,32 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { // Then we have the LHS and RHS non-contracting dimensions, if any: for (int64 i = 0; i < lhs_rank; i++) { - if (i != lhs_contracting_dimension && + if (!absl::c_linear_search(dnums.lhs_contracting_dimensions(), i) && !absl::c_linear_search(dnums.lhs_batch_dimensions(), i)) { result_index_locations.push_back({&lhs_index[i], nullptr}); } } for (int64 i = 0; i < rhs_rank; i++) { - if (i != rhs_contracting_dimension && + if (!absl::c_linear_search(dnums.rhs_contracting_dimensions(), i) && !absl::c_linear_search(dnums.rhs_batch_dimensions(), i)) { result_index_locations.push_back({&rhs_index[i], nullptr}); } } + absl::InlinedVector accumulate_index_sizes; + accumulate_index_sizes.reserve(dnums.lhs_contracting_dimensions_size()); + absl::InlinedVector, kInlineRank> + accumulate_index_locations; + accumulate_index_locations.reserve(dnums.lhs_contracting_dimensions_size()); + for (int64 i = 0; i < dnums.lhs_contracting_dimensions_size(); ++i) { + const int64 lhs_dnum = dnums.lhs_contracting_dimensions(i); + const int64 rhs_dnum = dnums.rhs_contracting_dimensions(i); + accumulate_index_locations.push_back( + {&lhs_index[lhs_dnum], &rhs_index[rhs_dnum]}); + const int64 dim_size = lhs->shape().dimensions(lhs_dnum); + accumulate_index_sizes.push_back(dim_size); + } + const int64 total_contraction_size = Product(accumulate_index_sizes); Literal result(dot->shape()); TF_RETURN_IF_ERROR( result.Populate([&](absl::Span result_index) { @@ -1230,13 +1338,30 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { } // Accumulates resulting product along the contracted dimension. - for (int64 i = 0; i < contracted_dimension_size; ++i) { - lhs_index[lhs_contracting_dimension] = i; - rhs_index[rhs_contracting_dimension] = i; + absl::InlinedVector accumulate_index( + accumulate_index_sizes.size(), 0); + for (int64 k = 0; k < total_contraction_size; k++) { + for (int64 i = 0; i < accumulate_index_sizes.size(); ++i) { + *(accumulate_index_locations[i].first) = accumulate_index[i]; + *(accumulate_index_locations[i].second) = accumulate_index[i]; + } result_val += static_cast(lhs_literal.Get(lhs_index)) * static_cast(rhs_literal.Get(rhs_index)); + + // 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; + } + } } return static_cast(result_val); @@ -1247,10 +1372,10 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { } Status HandlePad(HloInstruction* pad) override { - CHECK(ShapeUtil::IsArray(pad->operand(0)->shape())); + CHECK(pad->operand(0)->shape().IsArray()); // Padding value must be scalar. CHECK(ShapeUtil::IsScalar(pad->operand(1)->shape())); - CHECK_EQ(ShapeUtil::Rank(pad->operand(0)->shape()), + CHECK_EQ(pad->operand(0)->shape().rank(), pad->padding_config().dimensions_size()); TF_ASSIGN_OR_RETURN(auto inferred_return_shape, @@ -1273,9 +1398,8 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { const Literal& evaluated_operand = parent_->GetEvaluatedLiteralFor(pad->operand(0)); - std::vector input_index(ShapeUtil::Rank(evaluated_operand.shape()), - 0); - std::vector target_index(ShapeUtil::Rank(result.shape()), 0); + std::vector input_index(evaluated_operand.shape().rank(), 0); + std::vector target_index(result.shape().rank(), 0); // Loop through each element of the operand, assign them to the // corresponding index of the resulting padded literal. @@ -1318,10 +1442,12 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { auto operand = dynamic_slice->operand(0); auto start_indices = dynamic_slice->operand(1); auto result_shape = dynamic_slice->shape(); - TF_ASSIGN_OR_RETURN(auto inferred_return_shape, - ShapeInference::InferDynamicSliceShape( - operand->shape(), start_indices->shape(), - dynamic_slice->dynamic_slice_sizes())); + TF_ASSIGN_OR_RETURN( + auto inferred_return_shape, + ShapeInference::InferDynamicSliceShape( + operand->shape(), + Cast(dynamic_slice)->index_shapes(), + 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: " @@ -1330,33 +1456,39 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { primitive_util::IsIntegralType(start_indices->shape().element_type())); const Literal& operand_literal = parent_->GetEvaluatedLiteralFor(operand); - const Literal& start_indices_literal = - parent_->GetEvaluatedLiteralFor(start_indices); switch (start_indices->shape().element_type()) { case S32: { TF_ASSIGN_OR_RETURN( parent_->evaluated_[dynamic_slice], - DynamicSlice(operand_literal, start_indices_literal, - result_shape)); + DynamicSlice( + operand_literal, + absl::MakeConstSpan(dynamic_slice->operands()).subspan(1), + result_shape)); } break; case S64: { TF_ASSIGN_OR_RETURN( parent_->evaluated_[dynamic_slice], - DynamicSlice(operand_literal, start_indices_literal, - result_shape)); + DynamicSlice( + operand_literal, + absl::MakeConstSpan(dynamic_slice->operands()).subspan(1), + result_shape)); } break; case U32: { TF_ASSIGN_OR_RETURN( parent_->evaluated_[dynamic_slice], - DynamicSlice(operand_literal, start_indices_literal, - result_shape)); + DynamicSlice( + operand_literal, + absl::MakeConstSpan(dynamic_slice->operands()).subspan(1), + result_shape)); } break; case U64: { TF_ASSIGN_OR_RETURN( parent_->evaluated_[dynamic_slice], - DynamicSlice(operand_literal, start_indices_literal, - result_shape)); + DynamicSlice( + operand_literal, + absl::MakeConstSpan(dynamic_slice->operands()).subspan(1), + result_shape)); } break; default: LOG(FATAL) << "HandleDynamicSlice: unhandled primitive type for " @@ -1376,7 +1508,9 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { TF_ASSIGN_OR_RETURN( auto inferred_return_shape, ShapeInference::InferDynamicUpdateSliceShape( - operand->shape(), update->shape(), start_indices->shape())); + operand->shape(), update->shape(), + Cast(dynamic_update_slice) + ->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: " @@ -1387,33 +1521,39 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { const Literal& operand_literal = parent_->GetEvaluatedLiteralFor(operand); const Literal& update_literal = parent_->GetEvaluatedLiteralFor(update); - const Literal& start_indices_literal = - parent_->GetEvaluatedLiteralFor(start_indices); switch (start_indices->shape().element_type()) { case S32: { TF_ASSIGN_OR_RETURN( parent_->evaluated_[dynamic_update_slice], - DynamicUpdateSlice(operand_literal, update_literal, - start_indices_literal)); + DynamicUpdateSlice( + operand_literal, update_literal, + absl::MakeConstSpan(dynamic_update_slice->operands()) + .subspan(2))); } break; case S64: { TF_ASSIGN_OR_RETURN( parent_->evaluated_[dynamic_update_slice], - DynamicUpdateSlice(operand_literal, update_literal, - start_indices_literal)); + DynamicUpdateSlice( + operand_literal, update_literal, + absl::MakeConstSpan(dynamic_update_slice->operands()) + .subspan(2))); } break; case U32: { TF_ASSIGN_OR_RETURN( parent_->evaluated_[dynamic_update_slice], - DynamicUpdateSlice(operand_literal, update_literal, - start_indices_literal)); + DynamicUpdateSlice( + operand_literal, update_literal, + absl::MakeConstSpan(dynamic_update_slice->operands()) + .subspan(2))); } break; case U64: { TF_ASSIGN_OR_RETURN( parent_->evaluated_[dynamic_update_slice], - DynamicUpdateSlice(operand_literal, update_literal, - start_indices_literal)); + DynamicUpdateSlice( + operand_literal, update_literal, + absl::MakeConstSpan(dynamic_update_slice->operands()) + .subspan(2))); } break; default: LOG(FATAL) << "HandleDynamicUpdateSlice: unhandled primitive type for " @@ -1450,7 +1590,7 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { } Literal computed_result = - embedded_evaluator.Evaluate(*computation, arg_literals) + embedded_evaluator.Evaluate(*computation, arg_literals) .ConsumeValueOrDie(); // Clear visit states so that the we can use the evaluate again on // the same computation. @@ -1508,6 +1648,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: " @@ -1530,7 +1674,7 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { const Literal& keys_literal = parent_->GetEvaluatedLiteralFor(keys); int64 sort_dim = sort->dimensions(0); int64 sort_dim_elements = keys->shape().dimensions(sort_dim); - int64 rank = ShapeUtil::Rank(keys->shape()); + int64 rank = keys->shape().rank(); if (rank == 0) { // Nothing to sort. parent_->evaluated_[sort] = keys_literal.Clone(); @@ -1547,8 +1691,7 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { // 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; }); + absl::c_for_each(limit_indices, [](int64& index) { ++index; }); limit_indices[sort_dim] = sort_dim_elements; TF_ASSIGN_OR_RETURN(auto row_to_sort, keys_literal.Slice(indices, limit_indices) @@ -1556,10 +1699,10 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { const auto& row_data = row_to_sort.data(); std::vector result_data(row_data.begin(), row_data.end()); - std::sort(result_data.begin(), result_data.end(), - [](const NativeT& a, const NativeT& b) { - return SafeLess(a, b); - }); + 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)); @@ -1581,7 +1724,7 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { std::is_same::value>::type* = nullptr> Status HandleSort(HloInstruction* sort) { - return InvalidArgument("Unsupported type for Sort"); + return UnsupportedTypeError(sort); } Status HandleSort(HloInstruction* sort) override { @@ -1591,7 +1734,7 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { Status HandleReduce(HloInstruction* hlo) override { HloReduceInstruction* reduce = Cast(hlo); int64 num_args = reduce->inputs().size(); - bool has_tuple_output = ShapeUtil::IsTuple(reduce->shape()); + bool has_tuple_output = reduce->shape().IsTuple(); absl::Span dimensions(reduce->dimensions()); HloComputation* function = reduce->to_apply(); @@ -1622,7 +1765,7 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { // All args and results have the same dimensions, so pick an arbitrary one. const Shape& arg_shape = arg_literals[0]->shape(); - const Shape& result_shape = ShapeUtil::IsTuple(reduce->shape()) + const Shape& result_shape = reduce->shape().IsTuple() ? reduce->shape().tuple_shapes(0) : reduce->shape(); const auto arg_dimensions = AsInt64Slice(arg_shape.dimensions()); @@ -1711,7 +1854,7 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { [](Literal& literal) { return &literal; }); TF_ASSIGN_OR_RETURN(Literal computed_result, - embedded_evaluator.Evaluate( + embedded_evaluator.Evaluate( *function, embedded_operands_ptrs)); // Clear visit states so that we can use the evaluator again on // the same computation. @@ -1789,7 +1932,7 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { const Literal& operand_literal = parent_->GetEvaluatedLiteralFor(operand); const Literal& source_literal = parent_->GetEvaluatedLiteralFor(source); - int64 rank = ShapeUtil::Rank(operand_literal.shape()); + int64 rank = operand_literal.shape().rank(); HloEvaluator embedded_evaluator(parent_->max_loop_iterations_); DimensionVector source_index(rank, 0); @@ -1827,8 +1970,8 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { selected_val_literal.Set({}, *selected_val); Literal computed_result = embedded_evaluator - .Evaluate( - *select, {&selected_val_literal, &curr_val_literal}) + .Evaluate(*select, + {&selected_val_literal, &curr_val_literal}) .ConsumeValueOrDie(); bool selected = !computed_result.Get({}); if (selected) { @@ -1849,9 +1992,8 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { scattered_literal.Set({}, scattered); Literal computed_result = embedded_evaluator - .Evaluate( - *scatter, - {&source_literal_scatter, &scattered_literal}) + .Evaluate(*scatter, + {&source_literal_scatter, &scattered_literal}) .ConsumeValueOrDie(); result.Set(operand_index, computed_result.Get({})); // Clear visit states so that the we can use the evaluator again @@ -1901,7 +2043,7 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { operand->shape().element_type(), window_dimension_sizes); DimensionVector window_index(window.dimensions_size()); - DimensionVector operand_index(ShapeUtil::Rank(operand_literal.shape())); + DimensionVector operand_index(operand_literal.shape().rank()); HloEvaluator embedded_evaluator(parent_->max_loop_iterations_); Literal result(reduce_window->shape()); @@ -1925,8 +2067,8 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { LiteralUtil::CreateR0(result_val); Literal computed_result = embedded_evaluator - .Evaluate( - *function, {&result_val_literal, &curr_val_literal}) + .Evaluate(*function, + {&result_val_literal, &curr_val_literal}) .ConsumeValueOrDie(); // Clear visit states so that the we can use the evaluate again @@ -2288,9 +2430,8 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { LiteralUtil::CreateR0(updates.Get(update_index)); Literal updated_result = embedded_evaluator - .Evaluate( - *scatter->to_apply(), - {&result_value_literal, &update_value_literal}) + .Evaluate(*scatter->to_apply(), + {&result_value_literal, &update_value_literal}) .ConsumeValueOrDie(); // Clear visit states so that the we can use the evaluate again on the // same computation. @@ -2332,7 +2473,7 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { << " but is inferred to be: " << ShapeUtil::HumanString(inferred_return_shape); - const int64 rank = ShapeUtil::Rank(operand->shape()); + const int64 rank = operand->shape().rank(); const Literal& operand_literal = parent_->GetEvaluatedLiteralFor(operand); auto func = [&](absl::Span out_index) { DimensionVector operand_index(rank); @@ -2360,7 +2501,7 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { std::is_same::value || std::is_same::value)>::type* = nullptr> Status HandleClz(HloInstruction* clz) { - return InvalidArgument("Unsupported type for Clz"); + return UnsupportedTypeError(clz); } template ::value || is_complex_t::value>::type* = nullptr> Status HandleSin(HloInstruction* sin) { - return InvalidArgument("Unsupported type for Sin"); + return UnsupportedTypeError(sin); } Status HandleSin(HloInstruction* sin) override { @@ -2428,7 +2569,7 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { typename std::enable_if::value || is_complex_t::value>::type* = nullptr> Status HandleCos(HloInstruction* cos) { - return InvalidArgument("Unsupported type for Cos"); + return UnsupportedTypeError(cos); } Status HandleCos(HloInstruction* cos) override { @@ -2442,7 +2583,7 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { parent_->evaluated_[reduce_precision], ElementWiseUnaryOp(reduce_precision, [reduce_precision]( ElementwiseT elem) { - uint32_t value_as_int = tensorflow::bit_cast(elem); + uint32_t value_as_int = absl::bit_cast(elem); const uint32_t mantissa_bits = reduce_precision->mantissa_bits(); const uint32_t exponent_bits = reduce_precision->exponent_bits(); @@ -2515,7 +2656,7 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { value_as_int = x_underflows ? x_signed_zero : value_as_int; } - float reduced_result = tensorflow::bit_cast(value_as_int); + float reduced_result = absl::bit_cast(value_as_int); if (std::isnan(elem)) { reduced_result = mantissa_bits > 0 ? elem @@ -2529,7 +2670,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 < @@ -2537,47 +2678,161 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { typename std::enable_if::value || is_complex_t::value>::type* = nullptr> Status HandleReducePrecision(HloInstruction* reduce_precision) { - return InvalidArgument("Unsupported type for reduce precision"); + return UnsupportedTypeError(reduce_precision); } Status HandleReducePrecision(HloInstruction* reduce_precision) override { return HandleReducePrecision(reduce_precision); } - template ::value || - std::is_same::value || - std::is_same::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); - std::vector data(iota->shape().dimensions(iota->iota_dimension())); - std::iota(data.begin(), data.end(), 0); + const int64 iota_size = iota->shape().dimensions(iota->iota_dimension()); + // Avoid using std::vector since std::vector does not convert to + // absl::Span. + absl::InlinedVector data(iota_size); + // We don't use std::iota for two reasons: + // + // (1) std:iota does not support bfloat16 and float16. + // + // (2) std::iota saturates for floating point types when the value is not + // representable, but the definition of HLO iota is the value as a + // 64-bit integer cast to the native type. + for (int64 i = 0; i < iota_size; ++i) { + // static_cast is required for Eigen::half (F16). + data[i] = static_cast(i); + } auto result = LiteralUtil::CreateR1(data); - if (ShapeUtil::Rank(iota->shape()) > 1) { + if (iota->shape().rank() > 1) { TF_ASSIGN_OR_RETURN( parent_->evaluated_[iota], result.Broadcast(iota->shape(), {iota->iota_dimension()})); } else { - TF_RET_CHECK(ShapeUtil::Rank(iota->shape()) == 1); + TF_RET_CHECK(iota->shape().rank() == 1); parent_->evaluated_[iota] = std::move(result); } return Status::OK(); } - template ::value || - std::is_same::value || - std::is_same::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 InvalidArgument("Unsupported type for iota"); + return UnsupportedTypeError(iota); } Status HandleIota(HloInstruction* iota) override { 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. @@ -2589,7 +2844,7 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { // // This lets you calculate LI given the multidimensional indices in any order. static DimensionVector MakeDimMultipliers(const Shape& shape) { - DimensionVector v(ShapeUtil::Rank(shape)); + DimensionVector v(shape.rank()); int64 scale = 1; for (auto dim : LayoutUtil::MinorToMajor(shape)) { v[dim] = scale; @@ -2606,7 +2861,7 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { const Shape& window_shape, const Window& window, const Shape& base_shape, const absl::Span& window_count_index, const std::function&)>& f) { - const int64 rank = ShapeUtil::Rank(base_shape); + const int64 rank = base_shape.rank(); DimensionVector window_index(rank); std::fill(window_index.begin(), window_index.end(), 0); do { @@ -2637,12 +2892,27 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { } template - StatusOr DynamicSlice(const Literal& operand_literal, - const Literal& start_indices_literal, - const Shape& result_shape) { - auto start_indices_typed = start_indices_literal.data(); - std::vector start(start_indices_typed.begin(), - start_indices_typed.end()); + StatusOr DynamicSlice( + const Literal& operand_literal, + 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()); + } + } // Clamp the start indices so the slice is in-bounds w.r.t the operand. for (int64 i = 0; i < start.size(); ++i) { @@ -2668,14 +2938,28 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { } template - StatusOr DynamicUpdateSlice(const Literal& operand_literal, - const Literal& update_literal, - const Literal& start_indices_literal) { + StatusOr DynamicUpdateSlice( + const Literal& operand_literal, const Literal& update_literal, + absl::Span start_indices) { auto result = operand_literal.Clone(); - auto start_indices_typed = start_indices_literal.data(); - const auto rank = ShapeUtil::Rank(result.shape()); - std::vector start(start_indices_typed.begin(), - start_indices_typed.end()); + 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()); + } + } // Clamp the update start indices so the slice is in-bounds w.r.t the // operand. for (int64 i = 0; i < rank; ++i) { @@ -2722,17 +3006,8 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { const auto shape = instruction->shape(); const auto* lhs = instruction->operand(0); const auto* rhs = instruction->operand(1); - - // TODO(b/35950897, b/27796129): add DCHECK back once implicit broadcast - // is removed. - if (!(ShapeUtil::SameDimensions(shape, rhs->shape()) && - ShapeUtil::SameDimensions(lhs->shape(), rhs->shape()))) { - return Unimplemented( - "Implicit broadcasting is currently unsupported in HLO evaluator " - "Shape Mismatch: %s vs %s vs %s: ", - ShapeUtil::HumanString(shape), ShapeUtil::HumanString(lhs->shape()), - ShapeUtil::HumanString(rhs->shape())); - } + TF_RET_CHECK(ShapeUtil::SameDimensions(shape, rhs->shape())); + TF_RET_CHECK(ShapeUtil::SameDimensions(lhs->shape(), rhs->shape())); const Literal& lhs_literal = parent_->GetEvaluatedLiteralFor(lhs); const Literal& rhs_literal = parent_->GetEvaluatedLiteralFor(rhs); @@ -2756,19 +3031,9 @@ class HloEvaluatorTypedVisitor : public DfsHloVisitorWithDefault { const auto* lhs = instruction->operand(0); const auto* rhs = instruction->operand(1); const auto* ehs = instruction->operand(2); - - // TODO(b/35950897, b/27796129): add DCHECK back once implicit - // broadcast is removed. - if (!(ShapeUtil::SameDimensions(shape, lhs->shape()) && - ShapeUtil::SameDimensions(lhs->shape(), rhs->shape()) && - ShapeUtil::SameDimensions(rhs->shape(), ehs->shape()))) { - return Unimplemented( - "Implicit broadcasting is currently unsupported in HLO evaluator " - "Shape Mismatch: %s vs %s vs %s vs %s: ", - ShapeUtil::HumanString(shape), ShapeUtil::HumanString(lhs->shape()), - ShapeUtil::HumanString(rhs->shape()), - ShapeUtil::HumanString(ehs->shape())); - } + TF_RET_CHECK(ShapeUtil::SameDimensions(shape, lhs->shape())); + TF_RET_CHECK(ShapeUtil::SameDimensions(lhs->shape(), rhs->shape())); + TF_RET_CHECK(ShapeUtil::SameDimensions(rhs->shape(), ehs->shape())); const Literal& lhs_literal = parent_->GetEvaluatedLiteralFor(lhs); const Literal& rhs_literal = parent_->GetEvaluatedLiteralFor(rhs); @@ -2811,6 +3076,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_execution_profile.cc b/tensorflow/compiler/xla/service/hlo_execution_profile.cc index ce4cad42355ec5881f2ae14f4dd52a0588d51cf7..2df8eb962ae54eb5b9492fdeb274eec309a8288f 100644 --- a/tensorflow/compiler/xla/service/hlo_execution_profile.cc +++ b/tensorflow/compiler/xla/service/hlo_execution_profile.cc @@ -28,7 +28,8 @@ limitations under the License. #include "tensorflow/compiler/xla/util.h" namespace xla { -HloProfileIndexMap::HloProfileIndexMap(const HloModule& module) { +HloProfileIndexMap::HloProfileIndexMap(const HloModule& module, + absl::Span extra_metrics) { size_t current_profile_index = 0; for (xla::HloComputation* computation : module.MakeComputationPostOrder()) { InsertOrDie(&computation_to_profile_idx_, computation, @@ -40,11 +41,15 @@ HloProfileIndexMap::HloProfileIndexMap(const HloModule& module) { current_profile_index++); } } + for (const string& key : extra_metrics) { + InsertOrDie(&extra_metric_to_profile_idx_, key, current_profile_index++); + } } std::unique_ptr CreateHloProfilePrinterData( const HloProfileIndexMap& hlo_profile_index_map, - const HloCostAnalysis& cost_analysis) { + const HloCostAnalysis& cost_analysis, + const string& entry_computation_name) { using HloComputationInfo = HloProfilePrinterData::HloComputationInfo; using HloInstructionInfo = HloProfilePrinterData::HloInstructionInfo; @@ -105,6 +110,14 @@ std::unique_ptr CreateHloProfilePrinterData( } } + // Add extra metrics if any. + for (const auto& pair : hlo_profile_index_map.extra_metric_to_profile_idx()) { + profile_printer_data->mutable_extra_metrics()->insert( + {pair.first, pair.second}); + } + + profile_printer_data->set_entry_computation(entry_computation_name); + return profile_printer_data; } diff --git a/tensorflow/compiler/xla/service/hlo_execution_profile.h b/tensorflow/compiler/xla/service/hlo_execution_profile.h index be989846ef5cd2645da88ac9bbfea9534dd47821..da30e15908328f9aa7fe282656a6d44ab7348195 100644 --- a/tensorflow/compiler/xla/service/hlo_execution_profile.h +++ b/tensorflow/compiler/xla/service/hlo_execution_profile.h @@ -17,6 +17,7 @@ limitations under the License. #define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_EXECUTION_PROFILE_H_ #include +#include #include "tensorflow/compiler/xla/map_util.h" #include "tensorflow/compiler/xla/service/hlo_cost_analysis.h" @@ -34,7 +35,10 @@ class HloInstruction; class HloProfileIndexMap { public: // Scans `module` to populate this instance of HloProfileIndexMap. - explicit HloProfileIndexMap(const HloModule& module); + explicit HloProfileIndexMap(const HloModule& module) + : HloProfileIndexMap(module, {}) {} + explicit HloProfileIndexMap(const HloModule& module, + absl::Span extra_metrics); HloProfileIndexMap(const HloProfileIndexMap&) = default; HloProfileIndexMap(HloProfileIndexMap&&) = default; @@ -50,6 +54,10 @@ class HloProfileIndexMap { return FindOrDie(computation_to_profile_idx(), &computation); } + size_t GetProfileIndexFor(const string& key) const { + return xla::FindOrDie(extra_metric_to_profile_idx(), key); + } + size_t instruction_count() const { return instruction_to_profile_idx().size(); } @@ -58,8 +66,12 @@ class HloProfileIndexMap { return computation_to_profile_idx().size(); } + size_t extra_metrics_count() const { + return extra_metric_to_profile_idx().size(); + } + size_t total_count() const { - return instruction_count() + computation_count(); + return instruction_count() + computation_count() + extra_metrics_count(); } const std::unordered_map& @@ -72,15 +84,20 @@ class HloProfileIndexMap { return computation_to_profile_idx_; } + const std::unordered_map& extra_metric_to_profile_idx() const { + return extra_metric_to_profile_idx_; + } + private: std::unordered_map instruction_to_profile_idx_; std::unordered_map computation_to_profile_idx_; + std::unordered_map extra_metric_to_profile_idx_; }; // Create an instance of `HloProfilePrinterData`. std::unique_ptr CreateHloProfilePrinterData( const HloProfileIndexMap& hlo_profile_index_map, - const HloCostAnalysis& cost_analysis); + const HloCostAnalysis& cost_analysis, const string& entry_computation_name); // Describes how much time each HLO operation took. // @@ -113,6 +130,12 @@ class HloExecutionProfile { total_cycles_executed; } + // Record extra metric. + void set_extra_metrics(const string& metric, uint64 value) { + profile_counters_[hlo_profile_index_map_.GetProfileIndexFor(metric)] = + value; + } + // Returns a version of the execution profile suitable for performance // debugging; e.g. emits cycle counts, execution time at the nominal device // frequency, and the effective throughput given the provided cost_analysis diff --git a/tensorflow/compiler/xla/service/hlo_execution_profile_test.cc b/tensorflow/compiler/xla/service/hlo_execution_profile_test.cc index 460ae2b5eca78659f86df1227e6a0a4e57508611..df06cf8c53ec8407f8b44c9126ed4fb5409f8ef3 100644 --- a/tensorflow/compiler/xla/service/hlo_execution_profile_test.cc +++ b/tensorflow/compiler/xla/service/hlo_execution_profile_test.cc @@ -45,7 +45,7 @@ TEST_F(HloExecutionProfileTest, Basic) { auto shape_size_function = [&](const Shape& shape) { const int64 pointer_size = 8; - if (ShapeUtil::IsOpaque(shape)) { + if (shape.IsOpaque()) { return pointer_size; } return ShapeUtil::ByteSizeOf(shape, pointer_size); @@ -54,7 +54,8 @@ TEST_F(HloExecutionProfileTest, Basic) { HloCostAnalysis cost_analysis(shape_size_function); HloProfileIndexMap profile_index_map(*hlo_module); std::unique_ptr profile_printer = - CreateHloProfilePrinterData(profile_index_map, cost_analysis); + CreateHloProfilePrinterData(profile_index_map, cost_analysis, + hlo_module->entry_computation()->name()); HloExecutionProfile execution_profile(profile_printer.get(), &profile_index_map); diff --git a/tensorflow/compiler/xla/service/hlo_get_dimension_size_rewriter.cc b/tensorflow/compiler/xla/service/hlo_get_dimension_size_rewriter.cc new file mode 100644 index 0000000000000000000000000000000000000000..862b2029718bbd802b69d789b66683a4edfa2367 --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_get_dimension_size_rewriter.cc @@ -0,0 +1,75 @@ +/* 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_get_dimension_size_rewriter.h" + +#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" + +namespace xla { + +namespace { + +StatusOr ReplaceGetSize( + HloInstruction* instr, + const DynamicDimensionInference* dynamic_dimension_inference) { + if (instr->opcode() != HloOpcode::kGetDimensionSize) { + return false; + } + HloComputation* computation = instr->parent(); + + TF_ASSIGN_OR_RETURN(auto legal_shape, + ShapeInference::InferGetDimensionSizeShape( + instr->operand(0)->shape(), instr->dimension())); + TF_RET_CHECK(ShapeUtil::Equal(instr->shape(), legal_shape)); + TF_RET_CHECK(ShapeUtil::HasPrimitiveType(instr->shape(), U32)); + 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; +} + +} // namespace + +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, &inference)); + changed = changed || replaced; + } + } + return changed; +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_get_dimension_size_rewriter.h b/tensorflow/compiler/xla/service/hlo_get_dimension_size_rewriter.h new file mode 100644 index 0000000000000000000000000000000000000000..9aa79fe66b665c48ec871c4188e44ba2056de3ad --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_get_dimension_size_rewriter.h @@ -0,0 +1,38 @@ +/* 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_HLO_GET_DIMENSION_SIZE_REWRITER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_GET_DIMENSION_SIZE_REWRITER_H_ + +#include "tensorflow/compiler/xla/service/hlo_module.h" +#include "tensorflow/compiler/xla/service/hlo_pass_interface.h" + +namespace xla { + +// 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 { + return "hlo-get-dimension-size-rewriter"; + } + + StatusOr Run(HloModule* module) override; +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_GET_DIMENSION_SIZE_REWRITER_H_ diff --git a/tensorflow/compiler/xla/service/hlo_get_dimension_size_rewriter_test.cc b/tensorflow/compiler/xla/service/hlo_get_dimension_size_rewriter_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..a86aebdd5b64240e6e07d8e8050c0c8681cce765 --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_get_dimension_size_rewriter_test.cc @@ -0,0 +1,83 @@ +/* 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_get_dimension_size_rewriter.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_parser.h" +#include "tensorflow/compiler/xla/shape_util.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_utils.h" +#include "tensorflow/compiler/xla/types.h" +#include "tensorflow/core/lib/core/status_test_util.h" +#include "tensorflow/core/platform/types.h" + +namespace xla { +namespace { + +namespace op = xla::testing::opcode_matchers; + +class HloGetDimensionSizeRewriterTest : public HloTestBase { + protected: + HloGetDimensionSizeRewriterTest() {} +}; + +TEST_F(HloGetDimensionSizeRewriterTest, Ok) { + auto module = ParseHloString(R"( +HloModule _ +ENTRY gds { + p = s32[3,4] parameter(0) + size0 = u32[] get-dimension-size(p), dimensions={0} + size1 = u32[] get-dimension-size(p), dimensions={1} + ROOT mul = u32[] multiply(size0, size1) +})") + .ValueOrDie(); + HloGetDimensionSizeRewriter pass; + EXPECT_TRUE(pass.Run(module.get()).ValueOrDie()); + EXPECT_THAT(module->entry_computation()->root_instruction(), + op::Multiply(op::Constant(), op::Constant())); +} + +TEST_F(HloGetDimensionSizeRewriterTest, IllegalType) { + auto module = ParseHloString(R"( +HloModule _ +ENTRY gds { + p = s32[3]{0} parameter(0) + ROOT gds = s64[] get-dimension-size(p), dimensions={0} +})") + .ValueOrDie(); + HloGetDimensionSizeRewriter pass; + EXPECT_FALSE(pass.Run(module.get()).ok()); +} + +TEST_F(HloGetDimensionSizeRewriterTest, IllegalDimension) { + auto module = ParseHloString(R"( +HloModule _ +ENTRY gds { + p = f32[2,5] parameter(0) + ROOT gds = u32[] get-dimension-size(p), dimensions={2} +})") + .ValueOrDie(); + HloGetDimensionSizeRewriter pass; + EXPECT_FALSE(pass.Run(module.get()).ok()); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_graph_dumper.cc b/tensorflow/compiler/xla/service/hlo_graph_dumper.cc index 13a74fd8a115c5dc9a9518b226dfee4445cc7180..4c7f5e9e7dfb12a8cb699bdf397eab21983342a1 100644 --- a/tensorflow/compiler/xla/service/hlo_graph_dumper.cc +++ b/tensorflow/compiler/xla/service/hlo_graph_dumper.cc @@ -21,11 +21,12 @@ limitations under the License. #include #include #include +#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" @@ -38,6 +39,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_instructions.h" #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/service/hlo_tfgraph_builder.h" +#include "tensorflow/compiler/xla/service/pattern_matcher.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/window_util.h" @@ -111,11 +113,6 @@ class NodeFilter { result == kSomeUsersOmitted; } - bool ShowFusionSubcomputation(const HloInstruction* instr) const { - CHECK_EQ(instr->opcode(), HloOpcode::kFusion); - return Show(instr) && !SomeOrAllOperandsOmitted(instr); - } - private: std::function filter_; }; @@ -240,34 +237,28 @@ string HtmlLikeStringSanitize(absl::string_view s) { // it to a short string lets us tell the user what the subcomputation is without // drawing it as a graph. optional MatchTrivialComputation(const HloComputation* computation) { + namespace m = match; + if (computation->instruction_count() != 3) { return nullopt; } - HloInstruction* root = computation->root_instruction(); - if (root->operand_count() != 2) { - return nullopt; - } - - // Check that both of the operands to the root are parameters. - const HloInstruction* operand0 = root->operand(0); - const HloInstruction* operand1 = root->operand(1); - if (operand0->opcode() != HloOpcode::kParameter || - operand1->opcode() != HloOpcode::kParameter) { + const HloInstruction *param0, *param1; + if (!Match(root, m::Op() + .WithNumOperands(2) + .WithShape(m::Shape().IsEffectiveScalar()) + .WithBinaryOperandsAnyOrder( + m::Parameter(¶m0, 0) + .WithShape(m::Shape().IsEffectiveScalar()), + m::Parameter(¶m1, 1) + .WithShape(m::Shape().IsEffectiveScalar())))) { return nullopt; } - // Check that the two operands of root are param0 and param1. All of the - // opcodes we recognize are commutative, so we're OK with either order. - auto n0 = operand0->parameter_number(); - auto n1 = operand1->parameter_number(); - if (!(n0 == 0 && n1 == 1) && !(n1 == 0 && n0 == 1)) { - return nullopt; - } - - // If the params are reversed, check that the operation being performed is - // commutative. - if (n0 == 1) { + // If the params are reversed (i.e. operand0 is param1 and operand1 is + // param0), check that the operation being performed is commutative. + if (root->operand(0) == param1) { + CHECK_EQ(root->operand(1), param0); switch (root->opcode()) { case HloOpcode::kLe: case HloOpcode::kGe: @@ -279,13 +270,6 @@ optional MatchTrivialComputation(const HloComputation* computation) { } } - // Check that the root and params are all effective scalars. - if (!ShapeUtil::IsEffectiveScalar(root->shape()) || - !ShapeUtil::IsEffectiveScalar(operand0->shape()) || - !ShapeUtil::IsEffectiveScalar(operand1->shape())) { - return nullopt; - } - // If we recognize the root's opcode, we've successfully pattern-matched! switch (root->opcode()) { case HloOpcode::kAdd: @@ -396,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. @@ -413,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 @@ -423,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; }; @@ -577,8 +561,8 @@ bool HloDotDumper::ShouldShowSubcomputation(const HloComputation* subcomp) { } // Show the subcomputation if we're showing any of its members. - return std::any_of( - computation_->instructions().begin(), computation_->instructions().end(), + return absl::c_any_of( + subcomp->instructions(), [&](const HloInstruction* instr) { return filter_.Show(instr); }); } @@ -749,17 +733,16 @@ bool HloDotDumper::ShouldMergeIntoUsers(const HloInstruction* instr) const { return true; } const int kMinUsersToOmit = 3; - return instr->opcode() == HloOpcode::kParameter && - ShapeUtil::IsTuple(instr->shape()) && !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; - }); + return instr->opcode() == HloOpcode::kParameter && instr->shape().IsTuple() && + !instr->IsFused() && + 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) { @@ -832,7 +815,7 @@ string HloDotDumper::GetInstructionNodeInlinedOperands( // Print the literal value of constants with <= K elements. optional elem_count; - if (ShapeUtil::IsArray(shape)) { + if (shape.IsArray()) { elem_count = 1; for (int64 dim : shape.dimensions()) { *elem_count *= dim; @@ -916,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; } @@ -987,6 +969,7 @@ ColorScheme HloDotDumper::GetInstructionColor(const HloInstruction* instr) { case HloOpcode::kGetTupleElement: case HloOpcode::kTrace: case HloOpcode::kAfterAll: + case HloOpcode::kAddDependency: case HloOpcode::kTuple: return kWhite; case HloOpcode::kBroadcast: @@ -1043,8 +1026,9 @@ ColorScheme HloDotDumper::GetInstructionColor(const HloInstruction* instr) { case HloOpcode::kDomain: case HloOpcode::kFusion: case HloOpcode::kMap: + case HloOpcode::kGetDimensionSize: return kGray; - case HloOpcode::kCrossReplicaSum: + case HloOpcode::kAllReduce: case HloOpcode::kAllToAll: case HloOpcode::kCollectivePermute: case HloOpcode::kInfeed: @@ -1266,12 +1250,12 @@ const HloInstruction* HloDotDumper::GetNodeForEdge( class GraphRendererRegistry { public: - void AddRenderer(GraphRendererInterface* graph_renderer) { + void SetRenderer(std::shared_ptr graph_renderer) { tensorflow::mutex_lock lock(mu_); graph_renderer_ = graph_renderer; } - GraphRendererInterface* GetDefaultRenderer() { + std::shared_ptr GetDefaultRenderer() { tensorflow::mutex_lock lock(mu_); return graph_renderer_; } @@ -1283,23 +1267,24 @@ class GraphRendererRegistry { private: tensorflow::mutex mu_; - GraphRendererInterface* graph_renderer_ = nullptr; + std::shared_ptr graph_renderer_ GUARDED_BY(mu_); }; } // namespace -Registrar::Registrar(GraphRendererInterface* dumper) { - GraphRendererRegistry::Default()->AddRenderer(dumper); +Registrar::Registrar(std::shared_ptr dumper) { + GraphRendererRegistry::Default()->SetRenderer(dumper); } namespace { // Gets a NodeFilter that includes roughly all instructions whose distance from // root is <= radius. -NodeFilter MakeNodeFilter(const HloInstruction* root, int64 radius) { +NodeFilter MakeNodeRadiusAroundFilter(const HloInstruction* root, + int64 radius) { // 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()) { @@ -1320,7 +1305,7 @@ NodeFilter MakeNodeFilter(const HloInstruction* root, int64 radius) { // 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}); } } @@ -1348,7 +1333,7 @@ NodeFilter MakeNodeFilter(const HloInstruction* root, int64 radius) { continue; } for (const HloInstruction* user : instr->users()) { - if (!nodes.count(user)) { + if (!nodes.contains(user)) { worklist.push_back({user, depth + 1}); } } @@ -1357,7 +1342,7 @@ NodeFilter MakeNodeFilter(const HloInstruction* root, int64 radius) { 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(); }; @@ -1368,12 +1353,11 @@ NodeFilter MakeNodeFilter(const HloInstruction* root, int64 radius) { 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; } @@ -1381,8 +1365,7 @@ NodeFilter MakeNodeFilter(const HloInstruction* root, int64 radius) { // 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; } } @@ -1403,6 +1386,56 @@ NodeFilter MakeNodeFilter(const HloInstruction* root, int64 radius) { }); } +// Gets a node filter that includes nodes on all paths from `from` to `to`. If +// the all-paths set contains more than max_nodes elements, includes the nodes +// on the shortest paths and sets hit_limit to true. +NodeFilter MakeNodeFromToFilter(const HloInstruction* from, + const HloInstruction* to, int64 max_nodes, + bool* hit_limit) { + *hit_limit = false; + + // Elements in the queue are paths through the graph. + std::deque> queue; + queue.push_front({from}); + + // Compute the set of nodes we want to show using a slightly-modified + // Djikstra's algorithm. The only real difference is, rather than stopping + // when we find a (shortest) path, we continue until we've found max_nodes + // nodes on some path. + std::unordered_set visited; + std::unordered_set to_display = {from, to}; + while (!queue.empty() && to_display.size() < max_nodes) { + std::vector path = std::move(queue.front()); + queue.pop_front(); + if (!visited.insert(path.back()).second) { + continue; + } + + for (const auto* user : path.back()->users()) { + if (user == to) { + auto it = path.begin(); + for (; it != path.end() && to_display.size() < max_nodes; ++it) { + to_display.insert(*it); + } + if (it != path.end()) { + *hit_limit = true; + } + } else if (!visited.count(user)) { + auto new_path = path; + new_path.push_back(user); + queue.push_back(std::move(new_path)); + } + } + } + + return NodeFilter([=](const HloInstruction* instr) { + if (instr == from || instr == to) { + return kHighlightNode; + } + return to_display.count(instr) ? kNormalNode : kHideNode; + }); +} + string SaveGraph(const string& graph, GraphRendererInterface::GraphKind graph_kind, const string& dest_path) { @@ -1437,14 +1470,15 @@ string ExportGraph(const string& graph, GraphRendererInterface::GraphKind graph_kind, const DebugOptions& debug_options) { string path = debug_options.xla_hlo_graph_path(); - if (!path.empty()) { + if (!path.empty() && !debug_options.xla_hlo_dump_as_html()) { return SaveGraph(graph, graph_kind, path); } else { auto graph_renderer = GraphRendererRegistry::Default()->GetDefaultRenderer(); CHECK(graph_renderer != nullptr) << "No registered renderer for the HLO graph. " - "Use --xla_hlo_graph_path=PATH to export to local file system"; + "Use --xla_hlo_graph_path=PATH --xla_hlo_dump_as_html=false to " + "export to local file system"; return graph_renderer->RenderGraph(graph, graph_kind, debug_options); } } @@ -1482,7 +1516,7 @@ string DumpNeighborhoodAround(const HloInstruction& node, int radius, auto debug_options = node.GetModule()->config().debug_options(); string label = StrCat("Neighborhood of ", radius, " nodes around ", node.name()); - NodeFilter filter = MakeNodeFilter(&node, radius); + NodeFilter filter = MakeNodeRadiusAroundFilter(&node, radius); string graph = HloDotDumper(node.parent(), label, debug_options, show_backend_config, /*profile=*/nullptr, filter) @@ -1490,6 +1524,29 @@ string DumpNeighborhoodAround(const HloInstruction& node, int radius, return ExportGraph(graph, GraphRendererInterface::DOT_GRAPH, debug_options); } +string DumpAllPathsFromTo(const HloInstruction& from, const HloInstruction& to, + int64 max_nodes, bool show_backend_config) { + CHECK_EQ(from.parent(), to.parent()) << "Nodes must be in same computation!"; + auto debug_options = from.GetModule()->config().debug_options(); + + bool hit_limit = false; + NodeFilter filter = MakeNodeFromToFilter(&from, &to, max_nodes, &hit_limit); + string label; + if (!hit_limit) { + label = StrCat("All paths from ", from.name(), " to ", to.name()); + } else { + label = StrCat(max_nodes, " nodes on the shortest paths from ", from.name(), + " to ", to.name(), + "

***SHOWING ONLY A SUBSET OF ALL PATHS BETWEEN " + "NODES***

"); + } + string graph = + HloDotDumper(from.parent(), label, debug_options, show_backend_config, + /*profile=*/nullptr, filter) + .Dump(); + return ExportGraph(graph, GraphRendererInterface::DOT_GRAPH, debug_options); +} + void DumpText(const HloModule& module, const string& label, const string& directory_path, bool do_prefix) { Env* env = Env::Default(); @@ -1529,5 +1586,145 @@ string MaybeDumpHloModule(const HloModule& module, const string& label, return graph_url; } +string WrapDotInHTML(const string& dot) { + static const char html_prefix[] = R"html( + + + + + + + + + + + +
+ + + +)html"; + + return html_prefix + dot + html_suffix; +} + +string RenderDotAsHTMLFile(const string& dot, + const DebugOptions& debug_options) { + string html = WrapDotInHTML(dot); + + auto env = tensorflow::Env::Default(); + std::vector dirs; + string output_dir = debug_options.xla_hlo_graph_path(); + if (output_dir.empty()) { + env->GetLocalTempDirectories(&dirs); + } else { + dirs.push_back(output_dir); + } + // Try each directory, as they might be full, have inappropriate + // permissions or have different problems at times. + string output; + for (const string& dir : dirs) { + string filename = tensorflow::io::JoinPath(dir, "graph-"); + if (env->CreateUniqueFileName(&filename, ".html")) { + output = filename; + break; + } + } + if (output.empty()) { + LOG(FATAL) << "Failed to create unique output file name."; + } + TF_CHECK_OK(tensorflow::WriteStringToFile(env, output, html)); + return "file://" + output; +} + } // namespace hlo_graph_dumper } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_graph_dumper.h b/tensorflow/compiler/xla/service/hlo_graph_dumper.h index 0b11f34abb7f0d937a24d11f4dc5d2d6a0aae6e7..8e51454ef1cf992386cc7325e32705c08bf7712f 100644 --- a/tensorflow/compiler/xla/service/hlo_graph_dumper.h +++ b/tensorflow/compiler/xla/service/hlo_graph_dumper.h @@ -66,6 +66,12 @@ string DumpGraph(const HloComputation& computation, const string& label, string DumpNeighborhoodAround(const HloInstruction& node, int radius, bool show_backend_config = false); +// 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 +// paths. +string DumpAllPathsFromTo(const HloInstruction& from, const HloInstruction& to, + int64 max_nodes, bool show_backend_config = false); + // Dumps the HloModule::ToString() as a file into the provided directory path // suffixed with the provided label. // @@ -75,6 +81,12 @@ string DumpNeighborhoodAround(const HloInstruction& node, int radius, void DumpText(const HloModule& module, const string& label, const string& directory_path, bool do_prefix = true); +// Renders DOT graph as inline SVG and saves it in an HTML file in a temprary +// directory or directory specified via --xla_hlo_graph_path. Returns the file +// URI pointing to the file. +string RenderDotAsHTMLFile(const string& dot, + const DebugOptions& debug_options); + // Graph renderers may be added using a registration mechanism, e.g.: // XLA_REGISTER_GRAPH_RENDERER(AGraphRendererClass, 100) // The renderer with the highest numeric priority value is used. @@ -87,13 +99,13 @@ void DumpText(const HloModule& module, const string& label, // Class that registers a graph renderer. class Registrar { public: - Registrar(GraphRendererInterface* dumper); + Registrar(std::shared_ptr dumper); }; -#define XLA_INTERNAL_REGISTER_GRAPH_RENDERER(factory, ctr, ...) \ - static ::xla::hlo_graph_dumper::Registrar \ - XLA_INTERNAL_REGISTER_GRAPH_RENDERER_NAME(ctr)(new factory, \ - ##__VA_ARGS__) +#define XLA_INTERNAL_REGISTER_GRAPH_RENDERER(factory, ctr, ...) \ + static ::xla::hlo_graph_dumper::Registrar \ + XLA_INTERNAL_REGISTER_GRAPH_RENDERER_NAME(ctr)( \ + std::make_shared(), ##__VA_ARGS__) // __COUNTER__ must go through another macro to be properly expanded #define XLA_INTERNAL_REGISTER_GRAPH_RENDERER_NAME(ctr) ___##ctr##__object_ diff --git a/tensorflow/compiler/xla/service/hlo_graph_html_renderer.cc b/tensorflow/compiler/xla/service/hlo_graph_html_renderer.cc new file mode 100644 index 0000000000000000000000000000000000000000..84c4cf18df69816c611f4eb159ba247320ebc20e --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_graph_html_renderer.cc @@ -0,0 +1,43 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +// Implementation of an DOT graph renderer that uses Javascript to render DOT to +// SVG in a browser. + +#include "tensorflow/compiler/xla/service/hlo_graph_dumper.h" +#include "tensorflow/compiler/xla/service/hlo_module.h" + +namespace xla { +namespace hlo_graph_dumper { +namespace { + +class GraphHtmlRenderer : public GraphRendererInterface { + public: + string RenderGraph(const string& graph, GraphKind graph_kind, + const DebugOptions& debug_options) override { + switch (graph_kind) { + case DOT_GRAPH: + return RenderDotAsHTMLFile(graph, debug_options); + default: + LOG(FATAL) << "Only DOT graphs can be rendered"; + } + } +}; + +XLA_REGISTER_GRAPH_RENDERER(GraphHtmlRenderer); + +} // namespace +} // namespace hlo_graph_dumper +} // namespace xla 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 450c620576922e657b90fa4a8888e8c0cdee24c1..b01c00121b3363630b83a1e49d0027a66f3a9e1a 100644 --- a/tensorflow/compiler/xla/service/hlo_input_output_alias_config.cc +++ b/tensorflow/compiler/xla/service/hlo_input_output_alias_config.cc @@ -17,37 +17,61 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_module.h" namespace xla { + +bool HloInputOutputAliasConfig::OutputHasAlias( + const ShapeIndex& output_index) const { + return alias_.element(output_index).has_value(); +} + Status HloInputOutputAliasConfig::SetUpAlias(const ShapeIndex& output_index, int64 param_number, - const ShapeIndex& param_index) { + 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); + 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(); return Status::OK(); } 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); @@ -63,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; } @@ -78,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)); @@ -125,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); } }); } @@ -136,16 +162,17 @@ 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(); }); } -Status HloInputOutputAliasConfig::Verify(const HloModule& module) const { +Status HloInputOutputAliasConfig::Verify( + const HloModule& module, + std::function size_func) const { std::vector> param_has_seen; const HloComputation* entry = module.entry_computation(); for (int64 i = 0; i < entry->num_parameters(); ++i) { @@ -153,24 +180,43 @@ Status HloInputOutputAliasConfig::Verify(const HloModule& module) const { 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)); - // 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; - + const Shape& param_subshape = + ShapeUtil::GetSubshape(param_shape, alias.parameter_index); + const Shape& output_subshape = + ShapeUtil::GetSubshape(output_shape, output_index); + TF_RET_CHECK(LayoutUtil::IsDenseArray(param_subshape)); + TF_RET_CHECK(LayoutUtil::IsDenseArray(output_subshape)); + + if (size_func(param_subshape) != size_func(output_subshape)) { + return InternalError( + "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", + 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 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 0fae75842ba28da5dcb59e5952cd60c1d1c5ea68..b0b71dece81b561f492767db8c1ccbe3fde442d4 100644 --- a/tensorflow/compiler/xla/service/hlo_input_output_alias_config.h +++ b/tensorflow/compiler/xla/service/hlo_input_output_alias_config.h @@ -18,6 +18,7 @@ limitations under the License. #include +#include "absl/container/flat_hash_set.h" #include "absl/types/optional.h" #include "tensorflow/compiler/xla/service/hlo.pb.h" #include "tensorflow/compiler/xla/shape_tree.h" @@ -31,6 +32,28 @@ 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) {} @@ -40,12 +63,22 @@ class HloInputOutputAliasConfig { // 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; // (De)Serializes an HloInputOutoutAliasConfig to/from an // HloInputOutoutAliasProto. @@ -63,24 +96,23 @@ 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 // the aliased buffers should match. - Status Verify(const HloModule& module) const; + Status Verify(const HloModule& module, + std::function size_func_) const; Status ForEachAliasWithStatus(AliasFnWithStatus fn) const; @@ -89,9 +121,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_; + // 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 3b61ff04e6d7eeaa5876775fa18a85af82164b3d..a46a107723de30176241aae01b268a8c10d991d3 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 @@ -45,11 +45,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 +61,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 +86,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 +118,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,13 +157,45 @@ 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)); + ASSERT_IS_NOT_OK(config.Verify(*module, [](const Shape& shape) { + return ShapeUtil::ByteSizeOf(shape); + })); +} + +TEST_F(HloInputOutputAliasConfigTest, SizesMustMatch) { + const string module_str = R"( +HloModule TEST + +ENTRY main { + a = f32[] parameter(0) + b = f32[4096] parameter(1) + ROOT root = (f32[], f32[4096]) tuple(%a, %b) +} +)"; + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseHloString(module_str)); + + HloInputOutputAliasConfig config( + module->entry_computation()->root_instruction()->shape()); + + 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); + })); } TEST_F(HloInputOutputAliasConfigTest, OutputDoNotAliasTwice) { @@ -174,11 +214,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 f6ed86b41650fd331201814559386ff644092c23..3c92554ad4ec48686d64c74a00f732a3bfee87bc 100644 --- a/tensorflow/compiler/xla/service/hlo_instruction.cc +++ b/tensorflow/compiler/xla/service/hlo_instruction.cc @@ -31,6 +31,7 @@ limitations under the License. #include "absl/strings/numbers.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_join.h" +#include "absl/types/span.h" #include "tensorflow/compiler/xla/layout_util.h" #include "tensorflow/compiler/xla/literal.h" #include "tensorflow/compiler/xla/protobuf_util.h" @@ -82,18 +83,18 @@ 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)"; - TF_RETURN_IF_ERROR(ShapeUtil::ValidateShapeWithOptionalLayout(proto.shape())); + Shape shape(proto.shape()); + TF_RETURN_IF_ERROR(ShapeUtil::ValidateShapeWithOptionalLayout(shape)); switch (opcode) { // Ops migrated to subclasses. @@ -101,23 +102,23 @@ StatusOr> HloInstruction::CreateFromProto( TF_RET_CHECK(proto.operand_ids_size() == 3) << "BatchNormTraining instruction should have 3 operands but sees " << proto.operand_ids_size(); - instruction = CreateBatchNormTraining( - proto.shape(), operands(0), operands(1), operands(2), proto.epsilon(), - proto.feature_index()); + instruction = + CreateBatchNormTraining(shape, operands(0), operands(1), operands(2), + proto.epsilon(), proto.feature_index()); break; case HloOpcode::kBatchNormInference: TF_RET_CHECK(proto.operand_ids_size() == 5) << "BatchNormInference instruction should have 5 operands but sees " << proto.operand_ids_size(); instruction = CreateBatchNormInference( - proto.shape(), operands(0), operands(1), operands(2), operands(3), + shape, operands(0), operands(1), operands(2), operands(3), operands(4), proto.epsilon(), proto.feature_index()); break; case HloOpcode::kBatchNormGrad: TF_RET_CHECK(proto.operand_ids_size() == 5) << "BatchNormGrad instruction should have 5 operands but sees " << proto.operand_ids_size(); - instruction = CreateBatchNormGrad(proto.shape(), operands(0), operands(1), + instruction = CreateBatchNormGrad(shape, operands(0), operands(1), operands(2), operands(3), operands(4), proto.epsilon(), proto.feature_index()); break; @@ -127,7 +128,7 @@ StatusOr> HloInstruction::CreateFromProto( << proto.operand_ids_size(); std::vector fft_length(proto.fft_length().begin(), proto.fft_length().end()); - instruction = CreateFft(proto.shape(), operands(0), proto.fft_type(), + instruction = CreateFft(shape, operands(0), proto.fft_type(), absl::Span(fft_length)); break; } @@ -148,7 +149,7 @@ StatusOr> HloInstruction::CreateFromProto( TF_RET_CHECK(proto.operand_ids_size() == 1) << "Recv instruction should have 1 operand but sees " << proto.operand_ids_size(); - instruction = CreateRecv(proto.shape().tuple_shapes(0), operands(0), + instruction = CreateRecv(shape.tuple_shapes(0), operands(0), proto.channel_id(), proto.is_host_transfer()); break; case HloOpcode::kRecvDone: @@ -161,7 +162,7 @@ StatusOr> HloInstruction::CreateFromProto( TF_RET_CHECK(proto.operand_ids_size() == 1) << "Reverse instruction should have 1 operand but sees " << proto.operand_ids_size(); - instruction = CreateReverse(proto.shape(), operands(0), + instruction = CreateReverse(shape, operands(0), std::vector(proto.dimensions().begin(), proto.dimensions().end())); break; @@ -170,7 +171,7 @@ StatusOr> HloInstruction::CreateFromProto( << "Concatenate instruction should have 1 dimension but sees " << proto.dimensions_size(); instruction = - CreateConcatenate(proto.shape(), all_operands(), proto.dimensions(0)); + CreateConcatenate(shape, all_operands(), proto.dimensions(0)); break; case HloOpcode::kReduce: TF_RET_CHECK(proto.operand_ids_size() % 2 == 0) @@ -188,7 +189,7 @@ StatusOr> HloInstruction::CreateFromProto( absl::MakeSpan(reduce_operands) .subspan(reduce_operands.size() / 2, reduce_operands.size()); instruction = - CreateReduce(proto.shape(), inputs, init_values, + CreateReduce(shape, inputs, init_values, std::vector(proto.dimensions().begin(), proto.dimensions().end()), computations(0)); @@ -203,7 +204,7 @@ StatusOr> HloInstruction::CreateFromProto( auto sort_operands = all_operands(); HloInstruction* keys = sort_operands[0]; instruction = CreateSort( - proto.shape(), proto.dimensions(0), keys, + shape, proto.dimensions(0), keys, absl::Span(sort_operands).subspan(1)); break; } @@ -212,7 +213,7 @@ StatusOr> HloInstruction::CreateFromProto( << "Transpose instruction should have 1 operand but sees " << proto.operand_ids_size(); instruction = - CreateTranspose(proto.shape(), operands(0), + CreateTranspose(shape, operands(0), std::vector(proto.dimensions().begin(), proto.dimensions().end())); break; @@ -221,7 +222,7 @@ StatusOr> HloInstruction::CreateFromProto( << "Broadcast instruction should have 1 operand but sees " << proto.operand_ids_size(); instruction = - CreateBroadcast(proto.shape(), operands(0), + CreateBroadcast(shape, operands(0), std::vector(proto.dimensions().begin(), proto.dimensions().end())); break; @@ -229,7 +230,7 @@ StatusOr> HloInstruction::CreateFromProto( TF_RET_CHECK(proto.called_computation_ids_size() == 1) << "Map instruction should have 1 called computation but sees " << proto.called_computation_ids_size(); - instruction = CreateMap(proto.shape(), all_operands(), computations(0)); + instruction = CreateMap(shape, all_operands(), computations(0)); break; case HloOpcode::kSlice: { TF_RET_CHECK(proto.operand_ids_size() == 1) @@ -242,8 +243,8 @@ StatusOr> HloInstruction::CreateFromProto( slice_limits.push_back(slice_dimensions.limit()); slice_strides.push_back(slice_dimensions.stride()); } - instruction = CreateSlice(proto.shape(), operands(0), slice_starts, - slice_limits, slice_strides); + instruction = CreateSlice(shape, operands(0), slice_starts, slice_limits, + slice_strides); break; } case HloOpcode::kConstant: { @@ -253,7 +254,7 @@ StatusOr> HloInstruction::CreateFromProto( Literal::CreateFromProto(proto.literal())); instruction = CreateConstant(std::move(literal)); } else { - instruction = absl::make_unique(proto.shape()); + instruction = absl::make_unique(shape); } break; } @@ -284,71 +285,74 @@ StatusOr> HloInstruction::CreateFromProto( tensorflow::gtl::FindPtrOrNull(computation_map, fusion_id); TF_RET_CHECK(fused_computation != nullptr) << "No fusion computation with id " << fusion_id; - instruction = CreateFusion(proto.shape(), fusion_kind, all_operands(), - fused_computation); + instruction = + CreateFusion(shape, fusion_kind, all_operands(), fused_computation); break; } case HloOpcode::kRng: - instruction = - CreateRng(proto.shape(), proto.distribution(), all_operands()); + instruction = CreateRng(shape, proto.distribution(), all_operands()); break; case HloOpcode::kParameter: - instruction = CreateParameter(proto.parameter_number(), proto.shape(), - proto.name()); + instruction = + CreateParameter(proto.parameter_number(), shape, proto.name()); break; case HloOpcode::kGetTupleElement: TF_RET_CHECK(proto.operand_ids_size() == 1) << "GetTupleElement instruction should have 1 operand but sees " << proto.operand_ids_size(); - instruction = CreateGetTupleElement(proto.shape(), operands(0), - proto.tuple_index()); + instruction = + CreateGetTupleElement(shape, operands(0), proto.tuple_index()); break; case HloOpcode::kReducePrecision: TF_RET_CHECK(proto.operand_ids_size() == 1) << "ReducePrecision instruction should have 1 operand but sees " << proto.operand_ids_size(); - instruction = - CreateReducePrecision(proto.shape(), operands(0), - proto.exponent_bits(), proto.mantissa_bits()); + instruction = CreateReducePrecision( + shape, operands(0), proto.exponent_bits(), proto.mantissa_bits()); break; case HloOpcode::kInfeed: { - const Shape& data_shape = - ShapeUtil::GetTupleElementShape(proto.shape(), 0); + TF_RET_CHECK(shape.IsTuple() && + (ShapeUtil::TupleElementCount(shape) == 2)) + << "Infeed should have a tuple shape with 2 operands, but has: " + << shape; + const Shape& data_shape = ShapeUtil::GetTupleElementShape(shape, 0); TF_RET_CHECK(proto.operand_ids_size() == 1) << "Infeed instruction should have 1 operand but sees " << proto.operand_ids_size(); instruction = CreateInfeed(data_shape, operands(0), proto.infeed_config()); } break; - case HloOpcode::kOutfeed: + case HloOpcode::kOutfeed: { TF_RET_CHECK(proto.operand_ids_size() == 2) << "Outfeed instruction should have 2 operands but sees " << proto.operand_ids_size(); + Shape outfeed_shape(proto.outfeed_shape()); TF_RETURN_IF_ERROR( - ShapeUtil::ValidateShapeWithOptionalLayout(proto.outfeed_shape())); - instruction = CreateOutfeed(proto.outfeed_shape(), operands(0), - operands(1), proto.outfeed_config()); + ShapeUtil::ValidateShapeWithOptionalLayout(outfeed_shape)); + instruction = CreateOutfeed(outfeed_shape, operands(0), operands(1), + proto.outfeed_config()); break; - case HloOpcode::kCrossReplicaSum: { + } + case HloOpcode::kAllReduce: { TF_RET_CHECK(proto.called_computation_ids_size() == 1) - << "CrossReplicaSum should have 1 called computation but sees " + << "AllReduce should have 1 called computation but sees " << proto.called_computation_ids_size(); absl::optional all_reduce_id; if (proto.all_reduce_id() > 0) { all_reduce_id = proto.all_reduce_id(); } - instruction = CreateCrossReplicaSum( - proto.shape(), all_operands(), computations(0), + instruction = CreateAllReduce( + shape, all_operands(), computations(0), /*replica_groups=*/ std::vector(proto.replica_groups().begin(), proto.replica_groups().end()), - /*barrier=*/proto.cross_replica_sum_barrier(), + /*barrier=*/proto.all_reduce_barrier(), /*all_reduce_id=*/all_reduce_id); break; } case HloOpcode::kAllToAll: { instruction = CreateAllToAll( - proto.shape(), all_operands(), + shape, all_operands(), /*replica_groups=*/ std::vector(proto.replica_groups().begin(), proto.replica_groups().end())); @@ -364,8 +368,8 @@ StatusOr> HloInstruction::CreateFromProto( source_target_pairs[i].first = proto.source_target_pairs(i).source(); source_target_pairs[i].second = proto.source_target_pairs(i).target(); } - instruction = CreateCollectivePermute(proto.shape(), operands(0), - source_target_pairs); + instruction = + CreateCollectivePermute(shape, operands(0), source_target_pairs); break; } case HloOpcode::kConvolution: { @@ -378,8 +382,9 @@ StatusOr> HloInstruction::CreateFromProto( precision_config.mutable_operand_precision()->Resize( proto.operand_ids_size(), PrecisionConfig::DEFAULT); instruction = CreateConvolve( - proto.shape(), operands(0), operands(1), - std::max(proto.feature_group_count(), 1), proto.window(), + shape, operands(0), operands(1), + std::max(proto.feature_group_count(), 1), + std::max(proto.batch_group_count(), 1), proto.window(), proto.convolution_dimension_numbers(), precision_config); break; } @@ -390,7 +395,7 @@ StatusOr> HloInstruction::CreateFromProto( TF_RET_CHECK(proto.called_computation_ids_size() == 1) << "ReduceWindow should have 1 called computation but sees " << proto.called_computation_ids_size(); - instruction = CreateReduceWindow(proto.shape(), operands(0), operands(1), + instruction = CreateReduceWindow(shape, operands(0), operands(1), proto.window(), computations(0)); break; case HloOpcode::kSelectAndScatter: @@ -400,9 +405,9 @@ StatusOr> HloInstruction::CreateFromProto( TF_RET_CHECK(proto.called_computation_ids_size() == 2) << "SelectAndScatter should have 2 called computations but sees " << proto.called_computation_ids_size(); - instruction = CreateSelectAndScatter( - proto.shape(), operands(0), computations(0), proto.window(), - operands(1), operands(2), computations(1)); + instruction = CreateSelectAndScatter(shape, operands(0), computations(0), + proto.window(), operands(1), + operands(2), computations(1)); break; case HloOpcode::kCustomCall: if (proto.constrain_layout()) { @@ -410,16 +415,17 @@ StatusOr> HloInstruction::CreateFromProto( // vector of pointers essentially) so create a vector of shapes to pass // in. std::vector operand_shapes; - for (const Shape& shape : proto.operand_shapes_with_layout()) { - operand_shapes.push_back(shape); + for (const ShapeProto& shape_proto : + proto.operand_shapes_with_layout()) { + operand_shapes.emplace_back(shape_proto); } - instruction = CreateCustomCall( - proto.shape(), all_operands(), proto.custom_call_target(), - operand_shapes, proto.custom_call_opaque()); + instruction = + CreateCustomCall(shape, all_operands(), proto.custom_call_target(), + operand_shapes, proto.custom_call_opaque()); } else { - instruction = CreateCustomCall(proto.shape(), all_operands(), - proto.custom_call_target(), - proto.custom_call_opaque()); + instruction = + CreateCustomCall(shape, all_operands(), proto.custom_call_target(), + proto.custom_call_opaque()); } if (proto.has_window()) { static_cast(instruction.get()) @@ -433,23 +439,56 @@ StatusOr> HloInstruction::CreateFromProto( static_cast(instruction.get()) ->set_feature_group_count( std::max(static_cast(proto.feature_group_count()), 1LL)); + static_cast(instruction.get()) + ->set_batch_group_count( + std::max(static_cast(proto.batch_group_count()), 1LL)); break; case HloOpcode::kPad: TF_RET_CHECK(proto.operand_ids_size() == 2) << "Pad instruction should have 2 operands but sees " << proto.operand_ids_size(); TF_RET_CHECK(proto.has_padding_config()); - instruction = CreatePad(proto.shape(), operands(0), operands(1), - proto.padding_config()); + instruction = + CreatePad(shape, operands(0), operands(1), proto.padding_config()); break; case HloOpcode::kDynamicSlice: { - TF_RET_CHECK(proto.operand_ids_size() == 2) - << "DynamicSlice instruction should have 2 operands but sees " - << proto.operand_ids_size(); std::vector slice_sizes(proto.dynamic_slice_sizes_size()); absl::c_copy(proto.dynamic_slice_sizes(), slice_sizes.begin()); - instruction = CreateDynamicSlice(proto.shape(), operands(0), operands(1), - slice_sizes); + TF_RET_CHECK(proto.operand_ids_size() >= 1) + << "DynamicSlice instruction should have at least 1 operands but " + "sees " + << proto.operand_ids_size(); + // 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); + break; + } + case HloOpcode::kDynamicUpdateSlice: { + TF_RET_CHECK(proto.operand_ids_size() >= 2) + << "DynamicUpdateSlice instruction should have at least 2 operands " + "but sees " + << proto.operand_ids_size(); + // 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)); + break; } case HloOpcode::kGather: { @@ -465,7 +504,7 @@ StatusOr> HloInstruction::CreateFromProto( for (int64 bound : proto.gather_slice_sizes()) { gather_slice_sizes.push_back(bound); } - instruction = CreateGather(proto.shape(), operands(0), operands(1), + instruction = CreateGather(shape, operands(0), operands(1), *gather_dimension_numbers, gather_slice_sizes); break; } @@ -481,16 +520,15 @@ StatusOr> HloInstruction::CreateFromProto( auto scatter_dimension_numbers = absl::make_unique( proto.scatter_dimension_numbers()); - instruction = - CreateScatter(proto.shape(), operands(0), operands(1), operands(2), - computations(0), *scatter_dimension_numbers); + instruction = CreateScatter(shape, operands(0), operands(1), operands(2), + computations(0), *scatter_dimension_numbers); break; } case HloOpcode::kIota: TF_RET_CHECK(proto.dimensions_size() == 1) << "Iota instruction should have 1 dimension but sees " << proto.dimensions_size(); - instruction = CreateIota(proto.shape(), proto.dimensions(0)); + instruction = CreateIota(shape, proto.dimensions(0)); break; case HloOpcode::kDot: { TF_RET_CHECK(proto.has_dot_dimension_numbers()) @@ -502,8 +540,8 @@ StatusOr> HloInstruction::CreateFromProto( precision_config.mutable_operand_precision()->Resize( proto.operand_ids_size(), PrecisionConfig::DEFAULT); instruction = absl::make_unique( - proto.shape(), operands(0), operands(1), - proto.dot_dimension_numbers(), precision_config); + shape, operands(0), operands(1), proto.dot_dimension_numbers(), + precision_config); break; } case HloOpcode::kDomain: { @@ -525,13 +563,19 @@ StatusOr> HloInstruction::CreateFromProto( exit_hlo_sharding = std::make_shared(sharding); } instruction = absl::make_unique( - proto.shape(), operands(0), + shape, operands(0), absl::make_unique(entry_hlo_sharding), absl::make_unique(exit_hlo_sharding)); break; } + case HloOpcode::kGetDimensionSize: + TF_RET_CHECK(proto.operand_ids_size() == 1); + TF_RET_CHECK(proto.dimensions_size() == 1); + instruction = + CreateGetDimensionSize(shape, operands(0), proto.dimensions(0)); + break; default: { - instruction = absl::WrapUnique(new HloInstruction(opcode, proto.shape())); + instruction = absl::WrapUnique(new HloInstruction(opcode, shape)); for (const int64 operand_id : proto.operand_ids()) { instruction->AppendOperand(instruction_map.at(operand_id)); } @@ -559,6 +603,11 @@ StatusOr> HloInstruction::CreateFromProto( instruction->SetAndSanitizeName(proto.name()); instruction->metadata_ = proto.metadata(); instruction->backend_config_ = proto.backend_config(); + + TF_RET_CHECK(proto.id() >= 0) + << "Instruction with negative id: " << proto.id(); + TF_RET_CHECK(proto.id() <= INT_MAX) + << "Instruction with id > INT_MAX: " << proto.id(); instruction->unique_id_ = proto.id(); if (proto.has_sharding()) { @@ -609,7 +658,7 @@ HloInstruction::CreateGetTupleElement(const Shape& shape, absl::Span operands) { if (opcode == HloOpcode::kCopy) { // It is impossible to copy an opaque shape, we don't know how big it is. - CHECK(!ShapeUtil::IsOpaque(shape)); + CHECK(!shape.IsOpaque()); } auto instruction = absl::WrapUnique(new HloInstruction(opcode, shape)); for (auto operand : operands) { @@ -719,12 +768,12 @@ HloInstruction::CreateGetTupleElement(const Shape& shape, /* static */ std::unique_ptr HloInstruction::CreateConvolve( const Shape& shape, HloInstruction* lhs, HloInstruction* rhs, - int64 feature_group_count, const Window& window, + int64 feature_group_count, int64 batch_group_count, const Window& window, const ConvolutionDimensionNumbers& dimension_numbers, const PrecisionConfig& precision_config) { return absl::make_unique( - shape, lhs, rhs, feature_group_count, window, dimension_numbers, - precision_config); + shape, lhs, rhs, feature_group_count, batch_group_count, window, + dimension_numbers, precision_config); } /* static */ std::unique_ptr HloInstruction::CreateFft( @@ -751,8 +800,7 @@ HloInstruction::CreateReducePrecision(const Shape& shape, shape, operand, exponent_bits, mantissa_bits); } -/* static */ std::unique_ptr -HloInstruction::CreateCrossReplicaSum( +/* static */ std::unique_ptr HloInstruction::CreateAllReduce( const Shape& shape, absl::Span operands, HloComputation* reduce_computation, const std::vector& replica_groups, absl::string_view barrier, @@ -845,6 +893,16 @@ HloInstruction::CreateCollectivePermute( new HloInstruction(HloOpcode::kAfterAll, ShapeUtil::MakeTokenShape())); } +/* static */ std::unique_ptr +HloInstruction::CreateAddDependency(HloInstruction* data_operand, + HloInstruction* token_operand) { + auto instruction = absl::WrapUnique( + new HloInstruction(HloOpcode::kAddDependency, data_operand->shape())); + instruction->AppendOperand(data_operand); + instruction->AppendOperand(token_operand); + return instruction; +} + /* static */ std::unique_ptr HloInstruction::CreateWhile( const Shape& shape, HloComputation* condition, HloComputation* body, HloInstruction* init) { @@ -883,23 +941,19 @@ HloInstruction::CreateCollectivePermute( } /* static */ std::unique_ptr HloInstruction::CreateDynamicSlice( - const Shape& shape, HloInstruction* operand, HloInstruction* start_indices, + const Shape& shape, HloInstruction* operand, + absl::Span start_indices, absl::Span slice_sizes) { return absl::make_unique( shape, operand, start_indices, slice_sizes); } /* static */ std::unique_ptr -HloInstruction::CreateDynamicUpdateSlice(const Shape& shape, - HloInstruction* operand, - HloInstruction* update, - HloInstruction* start_indices) { - auto instruction = absl::WrapUnique( - new HloInstruction(HloOpcode::kDynamicUpdateSlice, shape)); - instruction->AppendOperand(operand); - instruction->AppendOperand(update); - instruction->AppendOperand(start_indices); - return instruction; +HloInstruction::CreateDynamicUpdateSlice( + const Shape& shape, HloInstruction* operand, HloInstruction* update, + absl::Span start_indices) { + return absl::make_unique( + shape, operand, update, start_indices); } /* static */ std::unique_ptr HloInstruction::CreateConcatenate( @@ -1001,13 +1055,21 @@ HloInstruction::CreateSelectAndScatter( broadcast_dimensions); } +/* static */ std::unique_ptr +HloInstruction::CreateGetDimensionSize(const Shape& shape, + HloInstruction* operand, + int64 dimension) { + return absl::make_unique(shape, operand, + dimension); +} + /* static */ std::unique_ptr HloInstruction::CreateBroadcastSequence( const Shape& output_shape, HloInstruction* operand, const std::function)>& adder) { CHECK(ShapeUtil::IsScalar(operand->shape()) || - ShapeUtil::Rank(operand->shape()) == ShapeUtil::Rank(output_shape)); + operand->shape().rank() == output_shape.rank()); Shape broadcast_shape = ShapeUtil::ChangeElementType( output_shape, operand->shape().element_type()); // Do explicit broadcast for scalar. @@ -1023,7 +1085,7 @@ HloInstruction::CreateBroadcastSequence( // Do explicit broadcast for degenerate broadcast. std::vector broadcast_dimensions; std::vector reshaped_dimensions; - for (int i = 0; i < ShapeUtil::Rank(operand->shape()); i++) { + for (int i = 0; i < operand->shape().rank(); i++) { if (operand->shape().dimensions(i) == output_shape.dimensions(i)) { broadcast_dimensions.push_back(i); reshaped_dimensions.push_back(operand->shape().dimensions(i)); @@ -1100,7 +1162,7 @@ HloInstruction::CreateBroadcastSequence( void HloInstruction::set_single_sharding(const HloSharding& sharding) { CHECK(!sharding.IsTuple()) << sharding; - if (ShapeUtil::IsTuple(shape())) { + if (shape().IsTuple()) { set_sharding(HloSharding::Tuple(sharding.GetAsShapeTree(shape()))); } else { set_sharding(sharding); @@ -1109,7 +1171,11 @@ void HloInstruction::set_single_sharding(const HloSharding& sharding) { void HloInstruction::SetupDerivedInstruction( HloInstruction* derived_instruction) const { - if (sharding_ != nullptr) { + if (sharding_ != nullptr && ShapeUtil::CompatibleIgnoringElementType( + shape_, derived_instruction->shape())) { + // Only copy sharding if the shape of the two instruction is compatible + // because copying it between differently shaped instructions can produce + // invalid shardings. derived_instruction->set_sharding(*sharding_); } else { derived_instruction->clear_sharding(); @@ -1128,7 +1194,7 @@ bool HloInstruction::HasSideEffectNoRecurse() const { case HloOpcode::kOutfeed: case HloOpcode::kTrace: return true; - case HloOpcode::kCrossReplicaSum: + case HloOpcode::kAllReduce: return all_reduce_id().has_value(); default: return false; @@ -1251,7 +1317,7 @@ std::unique_ptr HloInstruction::CloneWithNewOperands( case HloOpcode::kParameter: case HloOpcode::kGetTupleElement: case HloOpcode::kReducePrecision: - case HloOpcode::kCrossReplicaSum: + case HloOpcode::kAllReduce: case HloOpcode::kAllToAll: case HloOpcode::kCollectivePermute: case HloOpcode::kInfeed: @@ -1268,6 +1334,7 @@ std::unique_ptr HloInstruction::CloneWithNewOperands( case HloOpcode::kIota: case HloOpcode::kDot: case HloOpcode::kDomain: + case HloOpcode::kGetDimensionSize: clone = CloneWithNewOperandsImpl(shape, new_operands, context); break; // Unary ops. @@ -1345,9 +1412,8 @@ std::unique_ptr HloInstruction::CloneWithNewOperands( clone = CreateReshape(shape, new_operands[0]); break; case HloOpcode::kDynamicUpdateSlice: - CHECK_EQ(new_operands.size(), 3); clone = CreateDynamicUpdateSlice(shape, new_operands[0], new_operands[1], - new_operands[2]); + new_operands.subspan(2)); break; case HloOpcode::kTuple: clone = CreateTuple(new_operands); @@ -1371,6 +1437,10 @@ std::unique_ptr HloInstruction::CloneWithNewOperands( clone = CreateAfterAll(new_operands); } break; + case HloOpcode::kAddDependency: + CHECK_EQ(new_operands.size(), 2); + clone = CreateAddDependency(new_operands[0], new_operands[1]); + break; } // SetupDerivedInstruction will setup the precision_config_ field. SetupDerivedInstruction(clone.get()); @@ -1505,12 +1575,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(); @@ -1657,6 +1725,7 @@ bool HloInstruction::IdenticalSlowPath( // This opcode has complex or special behavior so just return false. case HloOpcode::kAfterAll: + case HloOpcode::kAddDependency: return false; // Remaining instructions with special values. @@ -1702,7 +1771,7 @@ bool HloInstruction::IdenticalSlowPath( case HloOpcode::kReducePrecision: case HloOpcode::kInfeed: case HloOpcode::kOutfeed: - case HloOpcode::kCrossReplicaSum: + case HloOpcode::kAllReduce: case HloOpcode::kAllToAll: case HloOpcode::kCollectivePermute: case HloOpcode::kConvolution: @@ -1715,19 +1784,50 @@ bool HloInstruction::IdenticalSlowPath( case HloOpcode::kScatter: case HloOpcode::kDot: case HloOpcode::kDomain: + case HloOpcode::kGetDimensionSize: LOG(FATAL) << "Base class impl called for opcode with subclass: " << opcode(); } return false; } +static uint64 HashOperand(const HloInstruction* hlo) { + return ShapeUtil::Hash(hlo->shape()); +} + +uint64 HloInstruction::Hash( + const std::function& hash_operand) const { + using tensorflow::Hash64Combine; + + uint64 hash_value = Hash64Combine(0, static_cast(opcode())); + hash_value = Hash64Combine(hash_value, ShapeUtil::Hash(shape())); + + if (!IsCrossModuleAllReduce()) { + if (!operands().empty()) { + for (size_t i = 0; i < operands().size(); ++i) { + hash_value = Hash64Combine(hash_value, hash_operand(operand(i))); + } + } + } + + hash_value = Hash64Combine(hash_value, InnerHash()); + return hash_value; +} + +uint64 HloInstruction::Hash() const { + // Use HashOperand as an argument to prevent non-termination. + return Hash(HashOperand); +} + +uint64 HloInstruction::InnerHash() const { return 13; } + void HloInstruction::RemoveUser(HloInstruction* user) { auto set_it = user_set_.find(user); CHECK(set_it != user_set_.end()); 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); } @@ -1745,8 +1845,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); @@ -1759,6 +1858,16 @@ Status HloInstruction::ReplaceUseWith(HloInstruction* user, Status HloInstruction::ReplaceOperandWith(int64 operand_num, HloInstruction* new_operand) { + auto old_operand = operand(operand_num); + TF_RET_CHECK(ShapeUtil::CompatibleIgnoringFpPrecision(old_operand->shape(), + new_operand->shape())) + << old_operand->shape() << " is not compatible with " + << new_operand->shape(); + return ReplaceOperandWithDifferentShape(operand_num, new_operand); +} + +Status HloInstruction::ReplaceOperandWithDifferentShape( + int64 operand_num, HloInstruction* new_operand) { TF_RET_CHECK(operand_num >= 0); TF_RET_CHECK(operand_num < operand_count()); HloInstruction* old_operand = mutable_operand(operand_num); @@ -1766,17 +1875,12 @@ Status HloInstruction::ReplaceOperandWith(int64 operand_num, return Status::OK(); } - TF_RET_CHECK(ShapeUtil::CompatibleIgnoringFpPrecision(old_operand->shape(), - new_operand->shape())) - << old_operand->shape() << " is not compatible with " - << new_operand->shape(); operands_[operand_num] = new_operand; 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); @@ -1784,6 +1888,14 @@ Status HloInstruction::ReplaceOperandWith(int64 operand_num, } Status HloInstruction::ReplaceAllUsesWith(HloInstruction* new_producer) { + TF_RET_CHECK( + ShapeUtil::CompatibleIgnoringFpPrecision(shape(), new_producer->shape())) + << shape() << " is not compatible with " << new_producer->shape(); + return ReplaceAllUsesWithDifferentShape(new_producer); +} + +Status HloInstruction::ReplaceAllUsesWithDifferentShape( + HloInstruction* new_producer) { bool new_producer_is_user = false; for (HloInstruction* user : users()) { if (user == new_producer) { @@ -1808,7 +1920,8 @@ Status HloInstruction::ReplaceAllUsesWith(HloInstruction* new_producer) { AddUser(new_producer); } if (parent_ && parent_->root_instruction() == this) { - parent_->set_root_instruction(new_producer); + parent_->set_root_instruction(new_producer, + /*accept_different_shape=*/true); } return Status::OK(); @@ -1820,7 +1933,7 @@ HloComputation* HloInstruction::to_apply() const { case HloOpcode::kMap: case HloOpcode::kReduceWindow: case HloOpcode::kReduce: - case HloOpcode::kCrossReplicaSum: + case HloOpcode::kAllReduce: case HloOpcode::kScatter: CHECK_EQ(called_computations_.size(), 1); return called_computations_[0]; @@ -1839,7 +1952,7 @@ void HloInstruction::set_to_apply(HloComputation* computation) { case HloOpcode::kMap: case HloOpcode::kReduceWindow: case HloOpcode::kReduce: - case HloOpcode::kCrossReplicaSum: + case HloOpcode::kAllReduce: case HloOpcode::kScatter: CHECK_EQ(called_computations_.size(), 1); called_computations_[0] = computation; @@ -1876,6 +1989,11 @@ void HloInstruction::set_while_body(HloComputation* computation) { called_computations_[kBodyComputationIndex] = computation; } +HloInstruction* HloInstruction::while_init() const { + CHECK_EQ(HloOpcode::kWhile, opcode_); + return operands_[0]; +} + HloComputation* HloInstruction::true_computation() const { CHECK_EQ(HloOpcode::kConditional, opcode_); return called_computations_[kTrueComputationIndex]; @@ -1992,7 +2110,11 @@ bool HloInstruction::IsElementwiseImpl( } bool HloInstruction::IsCrossModuleAllReduce() const { - return opcode() == HloOpcode::kCrossReplicaSum && all_reduce_id(); + return opcode() == HloOpcode::kAllReduce && all_reduce_id(); +} + +bool HloInstruction::IsCrossReplicaAllReduce() const { + return opcode() == HloOpcode::kAllReduce && !all_reduce_id(); } string HloInstruction::ToStringWithCanonicalNameMap( @@ -2103,7 +2225,7 @@ std::vector HloInstruction::ExtraAttributesToString( } else if (opcode() == HloOpcode::kCall || opcode() == HloOpcode::kMap || opcode() == HloOpcode::kReduceWindow || opcode() == HloOpcode::kReduce || - opcode() == HloOpcode::kCrossReplicaSum || + opcode() == HloOpcode::kAllReduce || opcode() == HloOpcode::kScatter) { extra.push_back( StrCat("to_apply=", PrintName(to_apply()->name(), options))); @@ -2139,7 +2261,7 @@ std::vector HloInstruction::ExtraAttributesToString( case HloOpcode::kMap: case HloOpcode::kReduceWindow: case HloOpcode::kReduce: - case HloOpcode::kCrossReplicaSum: + case HloOpcode::kAllReduce: case HloOpcode::kScatter: extra.push_back( StrCat("to_apply=\n", to_apply()->ToString(new_options))); @@ -2190,7 +2312,7 @@ HloInstructionProto HloInstruction::ToProto() const { proto.set_id(unique_id_); proto.set_name(name_); proto.set_opcode(HloOpcodeString(opcode_)); - *proto.mutable_shape() = shape_; + *proto.mutable_shape() = shape_.ToProto(); for (const HloInstruction* operand : operands_) { proto.add_operand_ids(operand->unique_id()); } @@ -2336,8 +2458,8 @@ Status HloInstruction::Visit(DfsHloVisitorBase* visitor) { return visitor->HandleConvolution(this); case HloOpcode::kFft: return visitor->HandleFft(this); - case HloOpcode::kCrossReplicaSum: - return visitor->HandleCrossReplicaSum(this); + case HloOpcode::kAllReduce: + return visitor->HandleAllReduce(this); case HloOpcode::kAllToAll: return visitor->HandleAllToAll(this); case HloOpcode::kCollectivePermute: @@ -2438,8 +2560,12 @@ Status HloInstruction::Visit(DfsHloVisitorBase* visitor) { return visitor->HandleDomain(this); case HloOpcode::kAfterAll: return visitor->HandleAfterAll(this); + case HloOpcode::kAddDependency: + return visitor->HandleAddDependency(this); case HloOpcode::kIota: return visitor->HandleIota(this); + case HloOpcode::kGetDimensionSize: + return visitor->HandleGetDimensionSize(this); // These opcodes are not handled here. case HloOpcode::kTrace: @@ -2597,36 +2723,6 @@ Status HloInstruction::AcceptWithOperandOrder( return Status::OK(); } -namespace { - -// Returns true if the given order is a topological sort of the instructions -// it contains. -bool OrderIsTopologicalSort(const std::vector& order) { - // Create a map from instruction to its position in 'order'. - std::unordered_map order_position; - for (int i = 0; i < order.size(); i++) { - if (!order_position.insert({order[i], i}).second) { - // Instruction order[i] is duplicated in the order. - return false; - } - } - // Verify that the operand of each instruction in the order is also in the - // order *and* the operand's position is earlier (defs are before uses for - // all ops). - for (auto* instruction : order) { - for (auto* operand : instruction->operands()) { - if (!ContainsKey(order_position, operand) || - order_position.at(operand) >= order_position.at(instruction)) { - return false; - } - } - } - - return true; -} - -} // namespace - Status HloInstruction::Accept( const std::function& visitor_func) { FunctionVisitor visitor(visitor_func); @@ -2639,49 +2735,7 @@ Status HloInstruction::Accept( return this->Accept(&visitor); } -Status HloInstruction::AcceptOrdered( - DfsHloVisitor* visitor, const std::vector& order) { - VLOG(2) << "HloInstruction::AcceptOrdered(%" << name() << ")"; - TF_RET_CHECK(OrderIsTopologicalSort(order)); - - // Compute the predecessors of this instruction. - std::unordered_set predecessors; - TF_RETURN_IF_ERROR(this->Accept([&predecessors](HloInstruction* instruction) { - predecessors.insert(instruction); - return Status::OK(); - })); - - for (auto* const_instruction : order) { - if (!ContainsKey(predecessors, const_instruction)) { - // Instruction is not a predecessors of 'this'. - continue; - } - - // The visitor can mark instructions as visited to skip particular - // instructions. - if (visitor->DidVisit(*const_instruction)) { - VLOG(3) << "Not visiting HLO %" << const_instruction->name() - << " as it was already visited."; - continue; - } - - // TODO(b/78350259): Eliminate const laundering. - HloInstruction* instruction = - const_cast(const_instruction); - - TF_RETURN_IF_ERROR(visitor->Preprocess(instruction)); - VLOG(2) << "Visiting HLO %" << instruction->name(); - TF_RETURN_IF_ERROR(instruction->Visit(visitor)); - visitor->SetVisited(*instruction); - TF_RETURN_IF_ERROR(visitor->Postprocess(instruction)); - } - - return visitor->FinishVisit(this); -} - -const Shape& HloInstruction::shape() const { - return shape_; -} +const Shape& HloInstruction::shape() const { return shape_; } std::vector HloInstruction::OperandIndices( const HloInstruction* operand) const { @@ -2810,7 +2864,7 @@ HloInstruction::UseKind HloInstruction::OperandElementUse(int64 i) const { } return UseKind::kReuse; case HloOpcode::kDynamicUpdateSlice: - // Dynamic-update-slice reuses only operand 2 (start_indices). + // Dynamic-update-slice reuses only start_indices. if (i == 0 || i == 1) { return UseKind::kUse; } @@ -2863,10 +2917,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) { @@ -3038,6 +3092,16 @@ const PrecisionConfig& HloInstruction::precision_config() const { LOG(FATAL) << "Unimplemented method."; } +PrecisionConfig* HloInstruction::mutable_precision_config() { + if (auto* convolution = DynCast(this)) { + return convolution->mutable_precision_config(); + } + if (auto* dot = DynCast(this)) { + return dot->mutable_precision_config(); + } + LOG(FATAL) << "Unimplemented method."; +} + HloModule* HloInstruction::GetModule() const { if (parent_) { return parent_->parent(); @@ -3080,6 +3144,10 @@ int64 HloInstruction::concatenate_dimension() const { return Cast(this)->concatenate_dimension(); } +int64 HloInstruction::dimension() const { + return Cast(this)->dimension(); +} + bool HloInstruction::IsRank2Transpose() const { auto transpose = DynCast(this); return transpose != nullptr && transpose->IsRank2Transpose(); @@ -3246,19 +3314,23 @@ HloInstruction::source_target_pairs() const { return Cast(this)->source_target_pairs(); } -string HloInstruction::cross_replica_sum_barrier() const { - return Cast(this)->cross_replica_sum_barrier(); +string HloInstruction::all_reduce_barrier() const { + return Cast(this)->all_reduce_barrier(); } -void HloInstruction::set_cross_replica_sum_barrier(const string& barrier) { - return Cast(this)->set_cross_replica_sum_barrier( - barrier); +void HloInstruction::set_all_reduce_barrier(const string& barrier) { + return Cast(this)->set_all_reduce_barrier(barrier); } absl::optional HloInstruction::all_reduce_id() const { return Cast(this)->all_reduce_id(); } +void HloInstruction::set_all_reduce_id( + const absl::optional& all_reduce_id) { + return Cast(this)->set_all_reduce_id(all_reduce_id); +} + const ConvolutionDimensionNumbers& HloInstruction::convolution_dimension_numbers() const { if (auto convolution = DynCast(this)) { @@ -3293,6 +3365,18 @@ void HloInstruction::set_feature_group_count(int64 feature_group_count) { feature_group_count); } +int64 HloInstruction::batch_group_count() const { + if (auto convolution = DynCast(this)) { + return convolution->batch_group_count(); + } + return Cast(this)->batch_group_count(); +} + +void HloInstruction::set_batch_group_count(int64 batch_group_count) { + Cast(this)->set_batch_group_count( + batch_group_count); +} + HloComputation* HloInstruction::select() const { return Cast(this)->select(); } diff --git a/tensorflow/compiler/xla/service/hlo_instruction.h b/tensorflow/compiler/xla/service/hlo_instruction.h index 15a4da8dbe0053aad314989a6718ebd61532ab8b..2c29b6c243bffccc346af12277dd4fc061250cbe 100644 --- a/tensorflow/compiler/xla/service/hlo_instruction.h +++ b/tensorflow/compiler/xla/service/hlo_instruction.h @@ -426,7 +426,7 @@ class HloInstruction { // and window describes how the filter is applied to lhs. static std::unique_ptr CreateConvolve( const Shape& shape, HloInstruction* lhs, HloInstruction* rhs, - int64 feature_group_count, const Window& window, + int64 feature_group_count, int64 batch_group_count, const Window& window, const ConvolutionDimensionNumbers& dimension_numbers, const PrecisionConfig& precision_config); @@ -462,9 +462,7 @@ class HloInstruction { // `all_reduce_id`: for Allreduce nodes from different modules, if they have // the same all_reduce_id, they will be 'Allreduce'd. If empty, Allreduce will // not be applied cross modules. - // - // TODO(b/117564385): Rename this to AllReduce. - static std::unique_ptr CreateCrossReplicaSum( + static std::unique_ptr CreateAllReduce( const Shape& shape, absl::Span operands, HloComputation* reduce_computation, const std::vector& replica_groups, @@ -560,13 +558,14 @@ class HloInstruction { // 'slice_sizes'. static std::unique_ptr CreateDynamicSlice( const Shape& shape, HloInstruction* operand, - HloInstruction* start_indices, absl::Span slice_sizes); + absl::Span start_indices, + absl::Span slice_sizes); // 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); + absl::Span start_indices); // Creates a concatenate instruction, where the operands are concatenated on // the provided dimension. @@ -767,6 +766,12 @@ class HloInstruction { // when we plumb a primordial token from the entry computation. static std::unique_ptr CreateToken(); + static std::unique_ptr CreateGetDimensionSize( + const Shape& shape, HloInstruction* operand, int64 dimension); + + static std::unique_ptr CreateAddDependency( + HloInstruction* data_operand, HloInstruction* token_operand); + // Returns the opcode for this instruction. HloOpcode opcode() const { return opcode_; } @@ -880,11 +885,15 @@ class HloInstruction { return false; } - // Use an explicit loop rather than ContainerEquals, because copying around - // std::functions may be too expensive in some cases. - for (size_t i = 0; i < operands().size(); ++i) { - if (!eq_operands(operand(i), other.operand(i))) { - return false; + // Two AllReduces are Identical if they have the same all_reduce_id. + // Their operands don't have to be Identical. + if (!IsCrossModuleAllReduce()) { + // Use an explicit loop rather than ContainerEquals, because copying + // around std::functions may be too expensive in some cases. + for (size_t i = 0; i < operands().size(); ++i) { + if (!eq_operands(operand(i), other.operand(i))) { + return false; + } } } @@ -895,6 +904,20 @@ class HloInstruction { return IdenticalSlowPath(other, eq_computations); } + // Generates a hash value of an HLO instruction. Hash considers + // information on opcode, shape, operands, and typically a root instruction. + // This function returns the same hash value for equivalent HLO instructions, + // with respect to HloInstruction::Identical() method. + // + // Uses hash_operand function to compute hash values of its operands. + // At the very top level, hash_operand should be non-recursive to prevent + // non-termination. + uint64 Hash( + const std::function& hash_operand) const; + + // Calls the above method with non-recursive hash_operand function. + uint64 Hash() const; + // Returns whether the instruction has a constant operand. bool HasConstantOperand() const; @@ -906,11 +929,16 @@ class HloInstruction { // operands of it which could be created due to this replacement. Status ReplaceUseWith(HloInstruction* user, HloInstruction* new_producer); - // Replaces the specified operand with new_operand. + // Replaces the specified operand with new_operand. The old and new operands + // must have compatible shapes ignoring floating-point precision. // // This function does NOT remove duplicated operands even if this instruction // is a fusion, so that the existing operand numbers do not change. - Status ReplaceOperandWith(int64 operand_no, HloInstruction* new_operand); + Status ReplaceOperandWith(int64 operand_num, HloInstruction* new_operand); + + // Same as ReplaceOperandWith(), but new_operand can have a different shape. + Status ReplaceOperandWithDifferentShape(int64 operand_num, + HloInstruction* new_operand); // Replaces all uses of this instruction with the new producer. If // new_producer is a user of this instruction then new_producer remains a use @@ -919,10 +947,16 @@ class HloInstruction { // If this instruction is the root of its computation, sets the computation's // root to new_producer. // + // The new producer must have a compatible shape ignoring floating-point + // precision. + // // If a user is a fusion instruction, this function will remove any duplicated // operands of it which could be created due to this replacement. Status ReplaceAllUsesWith(HloInstruction* new_producer); + // Same as ReplaceAllUsesWith, but new_producer can have a different shape. + Status ReplaceAllUsesWithDifferentShape(HloInstruction* new_producer); + // Performs a postorder DFS visit using this node as the root. If // call_finish_visit is true, then DfsHloVisitor::FinishVisit is called when // complete. If ignore_control_predecessors is true, instructions only @@ -954,16 +988,6 @@ class HloInstruction { Status Accept( const std::function& visitor_func) const; - // Visits all instructions rooted at this instruction using the given visitor - // in the given order. 'order' must contain at least the set of instructions - // rooted at this node (ie, those accessible from a DFS traversal from this - // instruction). Instructions contained in 'order' which are not in the set of - // instructions rooted at this node are ignored. 'order' must also be a valid - // topological sort of these instructions (defs appear before uses) though - // need not be a DFS post-order. - Status AcceptOrdered(DfsHloVisitor* visitor, - const std::vector& order); - // Visit this instruction and only this instruction with the given visitor. template Status Visit(DfsHloVisitorBase* visitor); @@ -1004,6 +1028,8 @@ class HloInstruction { void set_while_condition(HloComputation* while_condition); void set_while_body(HloComputation* while_body); + HloInstruction* while_init() const; + // Gets/sets the true and false HloComputation for Conditional. The setters // should only be called by HloModule or HloComputation methods. // @@ -1166,9 +1192,12 @@ class HloInstruction { // Returns true if this instruction is elementwise on all its operands. bool IsElementwise() const; - // Returns true if this is an cross module all-reduce instrucion. + // Returns true if this is a cross module all-reduce instruction. bool IsCrossModuleAllReduce() const; + // Returns true if this is a cross-replica all-reduce instruction. + bool IsCrossReplicaAllReduce() const; + // Returns true if this elementwise instruction implicitly broadcasts operand // `operand_idx`. // @@ -1264,6 +1293,7 @@ class HloInstruction { // superior. // Precondition: opcode must be kConvolution or kDot. const PrecisionConfig& precision_config() const; + PrecisionConfig* mutable_precision_config(); // Sets the debug metadata for this instruction. void set_metadata(const OpMetadata& metadata) { metadata_ = metadata; } @@ -1324,6 +1354,9 @@ class HloInstruction { // Delegates to HloConcatenateInstruction::concatenate_dimension. int64 concatenate_dimension() const; + // Delegates to HloGetDimensionSizeInstruction::dimension. + int64 dimension() const; + // Returns whether this instruction does a rank-2 transposition. bool IsRank2Transpose() const; @@ -1436,12 +1469,13 @@ class HloInstruction { // Delegates to HloCollectivePermuteInstruction::source_target_pairs. const std::vector>& source_target_pairs() const; - // Delegates to HloAllReduceInstruction::cross_replica_sum_barrier. - string cross_replica_sum_barrier() const; - void set_cross_replica_sum_barrier(const string& barrier); + // Delegates to HloAllReduceInstruction::all_reduce_barrier. + string all_reduce_barrier() const; + void set_all_reduce_barrier(const string& barrier); // Delegates to HloAllReduceInstruction::all_reduce_id. absl::optional all_reduce_id() const; + void set_all_reduce_id(const absl::optional& all_reduce_id); // Returns data on the window in a windowed operation such as // convolution. @@ -1471,6 +1505,11 @@ class HloInstruction { void set_feature_group_count(int64 feature_group_count); + // The number of batch groups. Must be a divisor of the input batch dimension + int64 batch_group_count() const; + + void set_batch_group_count(int64 batch_group_count); + // Delegates to HloSelectAndScatterInstruction::select. HloComputation* select() const; @@ -1606,6 +1645,10 @@ class HloInstruction { const std::function& eq_computations) const; + // Generates a hash value specific to a particular type of an instruction. + // This function typically considers the inner root instruction. + virtual uint64 InnerHash() const; + // Creates an n-ary elementwise operation. static std::unique_ptr CreateNary( const Shape& shape, HloOpcode opcode, diff --git a/tensorflow/compiler/xla/service/hlo_instruction_test.cc b/tensorflow/compiler/xla/service/hlo_instruction_test.cc index d93351fe0435b5f29035dc4ea0621a8c576bfd5a..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" @@ -29,7 +30,7 @@ limitations under the License. #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_verified_test_base.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" #include "tensorflow/compiler/xla/util.h" #include "tensorflow/compiler/xla/window_util.h" @@ -39,7 +40,7 @@ namespace { using ::testing::ElementsAre; using ::testing::UnorderedElementsAre; -class HloInstructionTest : public HloVerifiedTestBase { +class HloInstructionTest : public HloTestBase { protected: Shape r0f32_ = ShapeUtil::MakeShape(F32, {}); }; @@ -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) { @@ -151,7 +152,7 @@ TEST_F(HloInstructionTest, UserWithTwoOperands) { builder.AddInstruction(HloInstruction::CreateParameter(1, r0f32_, "bar")); auto add = builder.AddInstruction( HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, foo, bar)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); EXPECT_THAT(add->operands(), UnorderedElementsAre(foo, bar)); @@ -188,7 +189,7 @@ TEST_F(HloInstructionTest, MultipleUsers) { HloInstruction::CreateUnary(r0f32_, HloOpcode::kExp, foo)); auto add = builder.AddInstruction( HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, foo, bar)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); EXPECT_EQ(3, foo->user_count()); @@ -221,7 +222,7 @@ TEST_F(HloInstructionTest, RepeatedUser) { builder.AddInstruction(HloInstruction::CreateParameter(0, r0f32_, "foo")); auto add = builder.AddInstruction( HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, foo, foo)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); EXPECT_EQ(1, foo->user_count()); @@ -256,7 +257,7 @@ TEST_F(HloInstructionTest, MultipleUsersAndOperands) { HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, c0, param1)); auto addtotal = builder.AddInstruction( HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, addleft, addright)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); OpAndUserCollectingVisitor visitor; @@ -305,7 +306,7 @@ TEST_F(HloInstructionTest, MultipleUsersAndOperandsWithUnaryOps) { HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, addleft, addright)); auto neg2 = builder.AddInstruction( HloInstruction::CreateUnary(r0f32_, HloOpcode::kNegate, addtotal)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); OpAndUserCollectingVisitor visitor; @@ -327,7 +328,7 @@ TEST_F(HloInstructionTest, TrivialMap) { // Shape r0f32 = ShapeUtil::MakeShape(F32, {}); Shape f32a100x10 = ShapeUtil::MakeShape(F32, {100, 10}); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); // Builds an x+1.0 computation to use in a Map. auto embedded_builder = HloComputation::Builder("f32+1"); @@ -375,7 +376,7 @@ TEST_F(HloInstructionTest, TrivialReduce) { HloInstruction::CreateParameter(1, r0f32, "y")); embedded_builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kAdd, paramx, paramy)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto add_f32 = module->AddEmbeddedComputation(embedded_builder.Build()); // Builds a parameter and an initial value and feeds them to the reduce. @@ -416,7 +417,7 @@ TEST_F(HloInstructionTest, ReplaceUseInBinaryOps) { HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, foo, foo)); builder.AddInstruction(HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, add_foobar, add_foofoo)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); EXPECT_EQ(2, foo->user_count()); @@ -451,7 +452,7 @@ TEST_F(HloInstructionTest, ReplaceUseInVariadicOp) { builder.AddInstruction(HloInstruction::CreateTuple({foo, bar, baz, foo})); auto add_foobar = builder.AddInstruction( HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, foo, bar)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); EXPECT_EQ(2, foo->user_count()); @@ -479,7 +480,7 @@ TEST_F(HloInstructionTest, ReplaceUseInUnaryOp) { HloInstruction::CreateUnary(r0f32_, HloOpcode::kExp, foo)); auto log = builder.AddInstruction( HloInstruction::CreateUnary(r0f32_, HloOpcode::kLog, foo)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); EXPECT_EQ(2, foo->user_count()); @@ -516,7 +517,7 @@ TEST_F(HloInstructionTest, ReplaceAllUsesWithInBinaryOps) { HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, foo, foo)); builder.AddInstruction(HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, add_foobar, add_foofoo)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); EXPECT_EQ(2, foo->user_count()); @@ -546,7 +547,7 @@ TEST_F(HloInstructionTest, ReplaceAllUsesInMultipleOps) { auto exp = builder.AddInstruction( HloInstruction::CreateUnary(r0f32_, HloOpcode::kExp, foo)); auto tuple = builder.AddInstruction(HloInstruction::CreateTuple({foo, bar})); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); EXPECT_EQ(3, foo->user_count()); @@ -611,7 +612,7 @@ TEST_F(HloInstructionTest, PostProcessAllVisitedNodes) { HloInstruction::CreateUnary(r0f32_, HloOpcode::kLog, foo)); auto add = builder.AddInstruction( HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, exp, log)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); NodeCollectorAndPostProcessor visitor; @@ -629,7 +630,7 @@ TEST_F(HloInstructionTest, SingletonFusionOp) { HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.1f))); auto exp = builder.AddInstruction( HloInstruction::CreateUnary(r0f32_, HloOpcode::kExp, constant)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto* computation = module->AddEntryComputation(builder.Build()); auto* fusion = computation->CreateFusionInstruction( {exp}, HloInstruction::FusionKind::kLoop); @@ -647,7 +648,7 @@ TEST_F(HloInstructionTest, BinaryFusionOp) { HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.1f))); auto add = builder.AddInstruction(HloInstruction::CreateBinary( r0f32_, HloOpcode::kAdd, constant1, constant2)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto* computation = module->AddEntryComputation(builder.Build()); auto* fusion = computation->CreateFusionInstruction( {add}, HloInstruction::FusionKind::kLoop); @@ -669,7 +670,7 @@ TEST_F(HloInstructionTest, ChainFusionOp) { auto exp3 = builder.AddInstruction( HloInstruction::CreateUnary(r0f32_, HloOpcode::kExp, exp2)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto* computation = module->AddEntryComputation(builder.Build()); auto* fusion = computation->CreateFusionInstruction( {exp3, exp2, exp1}, HloInstruction::FusionKind::kLoop); @@ -692,7 +693,7 @@ TEST_F(HloInstructionTest, PreserveMetadataInFusionAndClone) { exp1->set_metadata(metadata); exp2->set_metadata(metadata); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto* computation = module->AddEntryComputation(builder.Build()); auto* fusion = computation->CreateFusionInstruction( {exp2, exp1}, HloInstruction::FusionKind::kLoop); @@ -749,7 +750,7 @@ TEST_F(HloInstructionTest, PreserveTupleShapeThroughClone) { TEST_F(HloInstructionTest, FusionOpWithCalledComputations) { // Create a fusion instruction containing a single unary operation. const Shape scalar_shape = ShapeUtil::MakeShape(F32, {}); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto make_map_computation = [&]() { auto builder = HloComputation::Builder("FusionMap"); @@ -817,7 +818,7 @@ TEST_F(HloInstructionTest, ComplexFusionOp) { auto tuple = builder.AddInstruction(HloInstruction::CreateTuple({sub, sub, mul, c1})); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto* computation = module->AddEntryComputation(builder.Build()); auto* fusion = computation->CreateFusionInstruction( {tuple, sub, mul, exp, clamp, add}, HloInstruction::FusionKind::kLoop); @@ -977,13 +978,13 @@ TEST_F(HloInstructionTest, FunctionVisitor) { HloInstruction::CreateUnary(f32, HloOpcode::kExp, param)); auto add = builder.AddInstruction( HloInstruction::CreateBinary(f32, HloOpcode::kAdd, negate, exp)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); 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(); @@ -1006,7 +1007,7 @@ TEST_F(HloInstructionTest, FullyElementwise) { builder.AddInstruction(HloInstruction::CreateParameter(1, r1f32, "y")); auto add = builder.AddInstruction( HloInstruction::CreateBinary(r1f32, HloOpcode::kAdd, x, y)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); EXPECT_TRUE(add->IsElementwise()); @@ -1016,7 +1017,7 @@ TEST_F(HloInstructionTest, FullyElementwise) { } TEST_F(HloInstructionTest, MapIsElementwise) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); const Shape r2f32 = ShapeUtil::MakeShapeWithLayout(F32, {10, 10}, {1, 0}); HloComputation::Builder builder(TestName()); HloComputation::Builder map_builder("id"); @@ -1067,7 +1068,7 @@ TEST_F(HloInstructionTest, PartiallyElementwise) { HloInstruction* max = builder.AddInstruction( HloInstruction::CreateBinary(r2f32, HloOpcode::kMaximum, div, broadcast)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto* computation = module->AddEntryComputation(builder.Build()); HloInstruction* fusion = computation->CreateFusionInstruction( {max, broadcast, div, mul}, HloInstruction::FusionKind::kLoop); @@ -1108,7 +1109,7 @@ TEST_F(HloInstructionTest, PartiallyElementwiseWithReuse) { HloInstruction* sub = builder.AddInstruction(HloInstruction::CreateBinary( r1f32, HloOpcode::kSubtract, min, broadcast)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto* computation = module->AddEntryComputation(builder.Build()); HloInstruction* fusion = computation->CreateFusionInstruction( {sub, broadcast, min}, HloInstruction::FusionKind::kLoop); @@ -1151,7 +1152,7 @@ TEST_F(HloInstructionTest, CloneOfFusionPreservesShape) { HloInstruction* dot = builder.AddInstruction(HloInstruction::CreateDot( sout, x, reshape, dot_dnums, DefaultPrecisionConfig(2))); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto* computation = module->AddEntryComputation(builder.Build()); HloInstruction* fusion = computation->CreateFusionInstruction( {dot, reshape}, HloInstruction::FusionKind::kLoop); @@ -1192,7 +1193,7 @@ TEST_F(HloInstructionTest, NoRedundantFusionOperandsAfterReplacingUse) { HloInstruction* dot = builder.AddInstruction(HloInstruction::CreateDot( s, x, reshape, dot_dnums, DefaultPrecisionConfig(2))); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto* computation = module->AddEntryComputation(builder.Build()); HloInstruction* fusion = computation->CreateFusionInstruction( {dot, reshape}, HloInstruction::FusionKind::kLoop); @@ -1204,7 +1205,7 @@ TEST_F(HloInstructionTest, NoRedundantFusionOperandsAfterReplacingUse) { } TEST_F(HloInstructionTest, FusionEquality) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); // Create two fusion instructions containing a single unary operation. @@ -1226,7 +1227,7 @@ TEST_F(HloInstructionTest, FusionEquality) { } TEST_F(HloInstructionTest, NestedFusionEquality) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); // Build a nested fusion computation. @@ -1330,7 +1331,7 @@ TEST_F(HloInstructionTest, Stringification) { "%dot = f32[5,20]{1,0} dot(f32[5,10]{1,0} %x, f32[10,20]{1,0} " "%transpose), lhs_contracting_dims={1}, rhs_contracting_dims={0}"); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto* computation = module->AddEntryComputation(builder.Build()); HloInstruction* loop = builder.AddInstruction( @@ -1373,7 +1374,7 @@ TEST_F(HloInstructionTest, StringifyGather_0) { /*index_vector_dim=*/4), /*slice_sizes=*/{30, 29, 28, 27, 26})); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); EXPECT_EQ(gather_instruction->ToString(), @@ -1408,7 +1409,7 @@ TEST_F(HloInstructionTest, StringifyGather_1) { /*index_vector_dim=*/2), /*slice_sizes=*/{30, 29, 28, 27, 26})); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); EXPECT_EQ(gather_instruction->ToString(), @@ -1443,7 +1444,7 @@ TEST_F(HloInstructionTest, StringifyScatter) { update_builder.AddInstruction( HloInstruction::CreateParameter(1, ShapeUtil::MakeShape(F32, {}), "p2")); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto* update_computation = module->AddEmbeddedComputation(update_builder.Build()); @@ -1495,7 +1496,7 @@ TEST_F(HloInstructionTest, CanonnicalStringificationFusion) { "f32[5,20]{1,0} dot(f32[5,10]{1,0}, f32[10,20]{1,0}), " "lhs_contracting_dims={1}, rhs_contracting_dims={0}"); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto* computation = module->AddEntryComputation(builder.Build()); HloInstruction* fusion = computation->CreateFusionInstruction( {dot, reshape}, HloInstruction::FusionKind::kLoop); @@ -1531,7 +1532,7 @@ TEST_F(HloInstructionTest, CanonnicalStringificationWhile) { HloInstruction* dot = builder.AddInstruction(HloInstruction::CreateDot( sout, x, reshape, dot_dnums, DefaultPrecisionConfig(2))); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto* computation = module->AddEntryComputation(builder.Build()); computation->CreateFusionInstruction({dot, reshape}, HloInstruction::FusionKind::kLoop); @@ -1587,7 +1588,7 @@ TEST_F(HloInstructionTest, CanonnicalStringificationConditional) { HloInstruction* dot = builder.AddInstruction(HloInstruction::CreateDot( sout, x, reshape, dot_dnums, DefaultPrecisionConfig(2))); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto* computation = module->AddEntryComputation(builder.Build()); computation->CreateFusionInstruction({dot, reshape}, HloInstruction::FusionKind::kLoop); diff --git a/tensorflow/compiler/xla/service/hlo_instructions.cc b/tensorflow/compiler/xla/service/hlo_instructions.cc index 88495e80000c4f87a778c4fad747f6bdf09b7a14..b01f01ef012b4c366035dc16b44508d71ad07d79 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) { @@ -363,29 +361,43 @@ HloAllReduceInstruction::HloAllReduceInstruction( HloComputation* reduce_computation, const std::vector& replica_groups, absl::string_view barrier, const absl::optional& all_reduce_id) - : HloCollectiveInstruction(HloOpcode::kCrossReplicaSum, shape, operands, + : HloCollectiveInstruction(HloOpcode::kAllReduce, shape, operands, replica_groups), - cross_replica_sum_barrier_(barrier), + all_reduce_barrier_(barrier), all_reduce_id_(all_reduce_id) { AppendComputation(reduce_computation); } +void HloAllReduceInstruction::set_all_reduce_id( + const absl::optional& all_reduce_id) { + all_reduce_id_ = all_reduce_id; +} + HloInstructionProto HloAllReduceInstruction::ToProto() const { HloInstructionProto proto = HloCollectiveInstruction::ToProto(); // Proto3 is so sad. if (all_reduce_id_) { proto.set_all_reduce_id(*all_reduce_id_); } - proto.set_cross_replica_sum_barrier(cross_replica_sum_barrier_); + proto.set_all_reduce_barrier(all_reduce_barrier_); 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 = HloCollectiveInstruction::ExtraAttributesToStringImpl(options); - if (!cross_replica_sum_barrier().empty()) { - result.push_back(StrCat("barrier=\"", cross_replica_sum_barrier(), "\"")); + if (!all_reduce_barrier().empty()) { + result.push_back(StrCat("barrier=\"", all_reduce_barrier(), "\"")); } if (all_reduce_id_) { result.push_back(StrCat("all_reduce_id=", *all_reduce_id_)); @@ -400,8 +412,7 @@ bool HloAllReduceInstruction::IdenticalSlowPath( const auto& casted_other = static_cast(other); return HloCollectiveInstruction::IdenticalSlowPath(other, eq_computations) && eq_computations(to_apply(), casted_other.to_apply()) && - cross_replica_sum_barrier() == - casted_other.cross_replica_sum_barrier() && + all_reduce_barrier() == casted_other.all_reduce_barrier() && all_reduce_id() == casted_other.all_reduce_id(); } @@ -410,8 +421,8 @@ HloAllReduceInstruction::CloneWithNewOperandsImpl( const Shape& shape, absl::Span new_operands, HloCloneContext* /*context*/) const { return absl::make_unique( - shape, new_operands, to_apply(), replica_groups(), - cross_replica_sum_barrier(), all_reduce_id()); + shape, new_operands, to_apply(), replica_groups(), all_reduce_barrier(), + all_reduce_id()); } HloAllToAllInstruction::HloAllToAllInstruction( @@ -730,7 +741,7 @@ HloMapInstruction::HloMapInstruction(const Shape& shape, AppendComputation(map_computation); // TODO(b/65689298) Remove code below once Map is generalized to accept // arbitrary map dimensions. - dimensions_.resize(ShapeUtil::Rank(shape)); + dimensions_.resize(shape.rank()); std::iota(dimensions_.begin(), dimensions_.end(), 0); } @@ -810,8 +821,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( @@ -862,7 +872,7 @@ void HloConstantInstruction::RelayoutConstant(const Layout& new_layout, const ShapeIndex& shape_index) { Shape* mutable_array_subshape = ShapeUtil::GetMutableSubshape(mutable_shape(), shape_index); - CHECK(ShapeUtil::IsArray(*mutable_array_subshape)); + CHECK(mutable_array_subshape->IsArray()); // Normally array_subshape will always have a layout, but this invariant is // temporarily broken in LayoutAssignment::AssignLayouts. @@ -896,11 +906,11 @@ string HloConstantInstruction::OperandsToStringWithCanonicalNameMap( string operands; // For constants, show the actual value in place of an empty operand list. if (literal_.has_value() && - ((ShapeUtil::IsArray(shape()) && ShapeUtil::ElementsIn(shape()) <= 10) || + ((shape().IsArray() && ShapeUtil::ElementsIn(shape()) <= 10) || options.print_large_constants())) { // Literal::ToString emits multidimensional arrays over multiple // lines. Compact this into one line by stripping out white space. - string tmp = literal().ToString(); + string tmp = literal().ToStringWithoutShape(); std::replace(tmp.begin(), tmp.end(), '\n', ' '); std::vector v = absl::StrSplit(tmp, ' '); bool first = true; @@ -1047,8 +1057,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 = @@ -1215,8 +1224,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 @@ -1320,7 +1329,7 @@ HloInstruction* HloFusionInstruction::CloneAndFuseInternal( if (newly_created_tuple_instr) { HloInstruction* new_instr = parent()->AddInstruction( HloInstruction::CreateGetTupleElement(fused_root->shape(), this, 0)); - TF_CHECK_OK(ReplaceAllUsesWith(new_instr)); + TF_CHECK_OK(ReplaceAllUsesWithDifferentShape(new_instr)); } int64 index = tuple_elements.size(); if (instruction_to_fuse->opcode() == HloOpcode::kTuple) { @@ -1367,6 +1376,16 @@ bool HloFusionInstruction::IdenticalSlowPath( other.fused_instructions_computation()); } +static uint64 HashOperandRecursive(const HloInstruction* hlo) { + return hlo->Hash(HashOperandRecursive); +} + +uint64 HloFusionInstruction::InnerHash() const { + // Use HashOperandRecursive to recursively compute hash on inner operands. + return fused_instructions_computation()->root_instruction()->Hash( + HashOperandRecursive); +} + std::unique_ptr HloFusionInstruction::CloneWithNewOperandsImpl( const Shape& shape, absl::Span new_operands, HloCloneContext* context) const { @@ -1610,7 +1629,7 @@ HloOutfeedInstruction::HloOutfeedInstruction(const Shape& outfeed_shape, HloInstructionProto HloOutfeedInstruction::ToProto() const { HloInstructionProto proto = HloInstruction::ToProto(); proto.set_outfeed_config(outfeed_config()); - *proto.mutable_outfeed_shape() = outfeed_shape(); + *proto.mutable_outfeed_shape() = outfeed_shape().ToProto(); return proto; } @@ -1640,11 +1659,12 @@ std::unique_ptr HloOutfeedInstruction::CloneWithNewOperandsImpl( HloConvolutionInstruction::HloConvolutionInstruction( const Shape& shape, HloInstruction* lhs, HloInstruction* rhs, - int64 feature_group_count, const Window& window, + int64 feature_group_count, int64 batch_group_count, const Window& window, const ConvolutionDimensionNumbers& dimension_numbers, const PrecisionConfig& precision_config) : HloInstruction(HloOpcode::kConvolution, shape), feature_group_count_(feature_group_count), + batch_group_count_(batch_group_count), window_(window), convolution_dimension_numbers_(dimension_numbers), precision_config_(precision_config) { @@ -1691,6 +1711,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); @@ -1722,8 +1746,9 @@ HloConvolutionInstruction::CloneWithNewOperandsImpl( HloCloneContext* context) const { CHECK_EQ(new_operands.size(), 2); return absl::make_unique( - shape, new_operands[0], new_operands[1], feature_group_count_, window(), - convolution_dimension_numbers_, precision_config_); + shape, new_operands[0], new_operands[1], feature_group_count_, + batch_group_count_, window(), convolution_dimension_numbers_, + precision_config_); } HloReduceWindowInstruction::HloReduceWindowInstruction( @@ -1862,7 +1887,7 @@ HloInstructionProto HloCustomCallInstruction::ToProto() const { if (layout_constrained()) { proto.set_constrain_layout(true); for (const Shape& shape : operand_shapes_with_layout_) { - *proto.add_operand_shapes_with_layout() = shape; + *proto.add_operand_shapes_with_layout() = shape.ToProto(); } } return proto; @@ -1985,12 +2010,44 @@ std::unique_ptr HloPadInstruction::CloneWithNewOperandsImpl( HloDynamicSliceInstruction::HloDynamicSliceInstruction( const Shape& shape, HloInstruction* operand, HloInstruction* start_indices, absl::Span slice_sizes) - : HloInstruction(HloOpcode::kDynamicSlice, shape), + : HloDynamicIndexInstruction(HloOpcode::kDynamicSlice, shape), dynamic_slice_sizes_(slice_sizes.begin(), slice_sizes.end()) { AppendOperand(operand); AppendOperand(start_indices); } +HloDynamicSliceInstruction::HloDynamicSliceInstruction( + const Shape& shape, HloInstruction* operand, + absl::Span start_indices, + absl::Span slice_sizes) + : HloDynamicIndexInstruction(HloOpcode::kDynamicSlice, shape), + dynamic_slice_sizes_(slice_sizes.begin(), slice_sizes.end()) { + AppendOperand(operand); + for (HloInstruction* index : start_indices) { + AppendOperand(index); + } +} + +HloDynamicUpdateSliceInstruction::HloDynamicUpdateSliceInstruction( + const Shape& shape, HloInstruction* operand, HloInstruction* update, + HloInstruction* start_indices) + : HloDynamicIndexInstruction(HloOpcode::kDynamicUpdateSlice, shape) { + AppendOperand(operand); + AppendOperand(update); + AppendOperand(start_indices); +} + +HloDynamicUpdateSliceInstruction::HloDynamicUpdateSliceInstruction( + const Shape& shape, HloInstruction* operand, HloInstruction* update, + absl::Span start_indices) + : HloDynamicIndexInstruction(HloOpcode::kDynamicUpdateSlice, shape) { + AppendOperand(operand); + AppendOperand(update); + for (HloInstruction* index : start_indices) { + AppendOperand(index); + } +} + HloInstructionProto HloDynamicSliceInstruction::ToProto() const { HloInstructionProto proto = HloInstruction::ToProto(); for (int64 slice_size : dynamic_slice_sizes_) { @@ -2016,9 +2073,14 @@ std::unique_ptr HloDynamicSliceInstruction::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], dynamic_slice_sizes_); + if (new_operands.size() == 2 && new_operands[1]->shape().rank() == 1) { + // TODO(b/118437727): Old form, remove this path. + return absl::make_unique( + shape, new_operands[0], new_operands[1], dynamic_slice_sizes_); + } else { + return absl::make_unique( + shape, new_operands[0], new_operands.subspan(1), dynamic_slice_sizes_); + } } HloGatherInstruction::HloGatherInstruction( @@ -2349,4 +2411,43 @@ HloInstructionProto HloDomainInstruction::ToProto() const { return proto; } + +HloGetDimensionSizeInstruction::HloGetDimensionSizeInstruction( + const Shape& shape, HloInstruction* operand, int64 dimension) + : HloInstruction(HloOpcode::kGetDimensionSize, shape), + dimension_(dimension) { + AppendOperand(operand); +} + +HloInstructionProto HloGetDimensionSizeInstruction::ToProto() const { + HloInstructionProto proto = HloInstruction::ToProto(); + proto.add_dimensions(dimension()); + return proto; +} + +std::vector HloGetDimensionSizeInstruction::ExtraAttributesToStringImpl( + const HloPrintOptions& /*options*/) const { + return {StrCat("dimensions={", dimension(), "}")}; +} + +bool HloGetDimensionSizeInstruction::IdenticalSlowPath( + const HloInstruction& other, + const std::function& + /*eq_computations*/) const { + const auto& casted_other = + static_cast(other); + return dimension() == casted_other.dimension(); +} + +std::unique_ptr +HloGetDimensionSizeInstruction::CloneWithNewOperandsImpl( + const Shape& shape, absl::Span new_operands, + HloCloneContext* /*context*/) const { + if (new_operands.size() != 1) { + LOG(FATAL) << "expects 1 operand"; + } + return absl::make_unique( + shape, new_operands[0], dimension()); +} + } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_instructions.h b/tensorflow/compiler/xla/service/hlo_instructions.h index 5f06dc093248e1d4d36ec845ced1e68c2b9d0752..1b4a94753cda8aba8d50836b9d51b7c3fd5807f6 100644 --- a/tensorflow/compiler/xla/service/hlo_instructions.h +++ b/tensorflow/compiler/xla/service/hlo_instructions.h @@ -242,20 +242,21 @@ class HloAllReduceInstruction : public HloCollectiveInstruction { const std::vector& replica_groups, absl::string_view barrier, const absl::optional& all_reduce_id); - // Returns the barrier config used for the CrossReplicaSum implementation of + // Returns the barrier config used for the AllReduce implementation of // each backend. - string cross_replica_sum_barrier() const { - return cross_replica_sum_barrier_; - } - void set_cross_replica_sum_barrier(string barrier) { - cross_replica_sum_barrier_ = barrier; - } + string all_reduce_barrier() const { return all_reduce_barrier_; } + void set_all_reduce_barrier(string barrier) { all_reduce_barrier_ = barrier; } absl::optional all_reduce_id() const { return all_reduce_id_; } + void set_all_reduce_id(const absl::optional& all_reduce_id); // 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; @@ -269,8 +270,8 @@ class HloAllReduceInstruction : public HloCollectiveInstruction { const Shape& shape, absl::Span new_operands, HloCloneContext* context) const override; - // The string representation of the barrier config used for CrossReplicaSum. - string cross_replica_sum_barrier_; + // The string representation of the barrier config used for AllReduce. + string all_reduce_barrier_; // For Allreduce nodes from different modules, if they have the same // all_reduce_id, they will be 'Allreduce'd. If empty, Allreduce will not be @@ -742,6 +743,8 @@ class HloFusionInstruction : public HloInstruction { const HloInstruction& other, const std::function& eq_computations) const override; + uint64 InnerHash() const override; + // Implementation for non-common logic of CloneWithNewOperands. std::unique_ptr CloneWithNewOperandsImpl( const Shape& shape, absl::Span new_operands, @@ -930,7 +933,7 @@ class HloConvolutionInstruction : public HloInstruction { public: explicit HloConvolutionInstruction( const Shape& shape, HloInstruction* lhs, HloInstruction* rhs, - int64 feature_group_count, const Window& window, + int64 feature_group_count, int64 batch_group_count, const Window& window, const ConvolutionDimensionNumbers& dimension_numbers, const PrecisionConfig& precision_config); const Window& window() const override { return window_; } @@ -946,6 +949,10 @@ class HloConvolutionInstruction : public HloInstruction { // dimension and output feature dimension. int64 feature_group_count() const { return feature_group_count_; } + // The number of feature groups. Must be a divisor of the input batch + // dimension. + int64 batch_group_count() const { return batch_group_count_; } + // Returns the information used to tell the implementation information about // what sort of precision is requested. The meaning of the field is backend // specific. At the moment, it is only supported for kConvolution and kDot. @@ -954,6 +961,7 @@ class HloConvolutionInstruction : public HloInstruction { // information but it is presumed that the alternate lowering is strictly // superior. const PrecisionConfig& precision_config() const { return precision_config_; } + PrecisionConfig* mutable_precision_config() { return &precision_config_; } string ToCategory() const override; // Returns a serialized representation of this instruction. @@ -973,6 +981,9 @@ class HloConvolutionInstruction : public HloInstruction { // The number of feature groups. Must be a divisor of the input feature // dimension and output feature dimension. int64 feature_group_count_; + // The number of feature groups. Must be a divisor of the input batch + // dimension. + int64 batch_group_count_; // Describes the window used for a convolution. Window window_; // Describes the dimension numbers used for a convolution. @@ -1095,7 +1106,11 @@ class HloCustomCallInstruction : public HloInstruction { void set_feature_group_count(int64 feature_group_count) { feature_group_count_ = feature_group_count; } + void set_batch_group_count(int64 batch_group_count) { + batch_group_count_ = batch_group_count; + } int64 feature_group_count() const { return feature_group_count_; } + int64 batch_group_count() const { return batch_group_count_; } // Returns a serialized representation of this instruction. HloInstructionProto ToProto() const override; @@ -1130,6 +1145,7 @@ class HloCustomCallInstruction : public HloInstruction { std::unique_ptr convolution_dimension_numbers_; // The number of feature groups. This is used for grouped convolutions. int64 feature_group_count_; + int64 batch_group_count_; // Whether the result and operand layouts are constrained. bool layout_constrained_; // For layout-constrained custom calls, this vector holds the shape with @@ -1144,6 +1160,9 @@ class HloPadInstruction : public HloInstruction { const PaddingConfig& padding_config); // Returns the padding configuration for a pad node. const PaddingConfig& padding_config() const { return padding_config_; } + // Returns the padding value. + const HloInstruction* padding_value() const { return operand(1); } + HloInstruction* mutable_padding_value() { return mutable_operand(1); } // Returns a serialized representation of this instruction. HloInstructionProto ToProto() const override; @@ -1164,12 +1183,38 @@ class HloPadInstruction : public HloInstruction { PaddingConfig padding_config_; }; -class HloDynamicSliceInstruction : public HloInstruction { +class HloDynamicIndexInstruction : public HloInstruction { + public: + explicit HloDynamicIndexInstruction(HloOpcode opcode, const Shape& shape) + : HloInstruction(opcode, shape) {} + virtual int64 first_index_operand_number() const = 0; + + // Returns a subspan of operands which represent the start indices. + absl::Span index_operands() const { + return absl::MakeSpan(operands()).subspan(first_index_operand_number()); + } + + // Returns the shapes of the index operands. + std::vector index_shapes() const { + std::vector shapes; + auto indices = index_operands(); + for (const HloInstruction* index : indices) { + shapes.push_back(index->shape()); + } + return shapes; + } +}; + +class HloDynamicSliceInstruction : public HloDynamicIndexInstruction { public: explicit HloDynamicSliceInstruction(const Shape& shape, HloInstruction* operand, HloInstruction* start_indices, absl::Span slice_sizes); + explicit HloDynamicSliceInstruction( + const Shape& shape, HloInstruction* operand, + absl::Span start_indices, + absl::Span slice_sizes); // Old methods kept for smooth subclassing transition END. // Returns the size of the slice in the given dimension for a dynamic // slice node. @@ -1182,6 +1227,8 @@ class HloDynamicSliceInstruction : public HloInstruction { // Returns a serialized representation of this instruction. HloInstructionProto ToProto() const override; + int64 first_index_operand_number() const override { return 1; } + private: std::vector ExtraAttributesToStringImpl( const HloPrintOptions& options) const override; @@ -1199,6 +1246,19 @@ class HloDynamicSliceInstruction : public HloInstruction { std::vector dynamic_slice_sizes_; }; +class HloDynamicUpdateSliceInstruction : public HloDynamicIndexInstruction { + public: + explicit HloDynamicUpdateSliceInstruction(const Shape& shape, + HloInstruction* operand, + HloInstruction* update, + HloInstruction* start_indices); + explicit HloDynamicUpdateSliceInstruction( + const Shape& shape, HloInstruction* operand, HloInstruction* update, + absl::Span start_indices); + + int64 first_index_operand_number() const override { return 2; } +}; + class HloGatherInstruction : public HloInstruction { public: explicit HloGatherInstruction( @@ -1322,6 +1382,7 @@ class HloDotInstruction : public HloInstruction { // information but it is presumed that the alternate lowering is strictly // superior. const PrecisionConfig& precision_config() const { return precision_config_; } + PrecisionConfig* mutable_precision_config() { return &precision_config_; } // Returns a serialized representation of this instruction. HloInstructionProto ToProto() const override; @@ -1382,6 +1443,33 @@ class HloDomainInstruction : public HloInstruction { std::unique_ptr operand_side_metadata_; std::unique_ptr user_side_metadata_; }; + +class HloGetDimensionSizeInstruction : public HloInstruction { + public: + explicit HloGetDimensionSizeInstruction(const Shape& shape, + HloInstruction* operand, + int64 dimension); + + // Returns the dimension sizes or numbers associated with this instruction. + int64 dimension() const { return dimension_; } + // 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; + + int64 dimension_; +}; + } // namespace xla #endif // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_INSTRUCTIONS_H_ diff --git a/tensorflow/compiler/xla/service/hlo_lexer.cc b/tensorflow/compiler/xla/service/hlo_lexer.cc index 971a9a20636c80820306d512af9e7ff4a14b79b5..798760885dcd55e0a1cbdf403fa160347d67fc3a 100644 --- a/tensorflow/compiler/xla/service/hlo_lexer.cc +++ b/tensorflow/compiler/xla/service/hlo_lexer.cc @@ -17,8 +17,10 @@ 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" #include "absl/types/optional.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/statusor.h" @@ -36,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 @@ -82,15 +84,29 @@ tensorflow::RegexpStringPiece HloLexer::RegexpStringPieceFromPointers( return tensorflow::RegexpStringPiece(begin, end - begin); } +TokKind HloLexer::LookAhead() { + if (GetKind() == TokKind::kEof || GetKind() == TokKind::kError) { + return GetKind(); + } + + const char* old_current_ptr = current_ptr_; + TokenState old_token_state = token_state_; + Lex(); + TokKind kind = GetKind(); + token_state_ = old_token_state; + current_ptr_ = old_current_ptr; + return kind; +} + TokKind HloLexer::LexToken() { while (true) { - token_start_ = current_ptr_; + token_state_.token_start = current_ptr_; int current_char = GetNextChar(); 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(); } @@ -125,6 +141,12 @@ TokKind HloLexer::LexToken() { return LexNumberOrPattern(); case '=': return TokKind::kEqual; + case '<': + if (current_char == '<' && PeekCurrentChar() == '=') { + current_ptr_++; + return TokKind::kLeq; + } + return TokKind::kError; case ',': return TokKind::kComma; case '%': @@ -163,6 +185,9 @@ TokKind HloLexer::LexToken() { current_ptr_ = comment_start; return TokKind::kError; } + if (current == kError) { + return TokKind::kError; + } } // Return no token for the comment. Keep lexing. continue; @@ -177,6 +202,9 @@ TokKind HloLexer::LexToken() { if (current == kEOF || current == '\n' || current == '\r') { break; } + if (current == kError) { + return TokKind::kError; + } current_ptr_++; } continue; @@ -200,43 +228,37 @@ TokKind HloLexer::LexToken() { // dim_labels_pattern ::= [0-9bf]{2,}_[0-9io]{2,}->[0-9bf]{2,} // identifiers ::= other cases that match [a-zA-Z_][a-zA-Z0-9_.-]* TokKind HloLexer::LexIdentifier() { - { - auto consumable = RegexpStringPieceFromPointers(token_start_, buf_.end()); - // 'consumable' will be advanced iff its prefix matches the pattern. - static LazyRE2 shape_pattern = { - R"(^(\w*\d*)\[([\d,\s]*)\](?:(dense|sparse)?{([\d,\s]+)})?)"}; - if (RE2::Consume(&consumable, *shape_pattern)) { - auto status_or_shape = ShapeUtil::ParseShapeString( - StringPieceFromPointers(token_start_, consumable.begin())); - if (status_or_shape.ok()) { - // This is a shape string. - shape_val_ = status_or_shape.ValueOrDie(); - current_ptr_ = consumable.begin(); - return TokKind::kShape; - } - } - } - while (IsIdentifierChar(PeekCurrentChar())) { current_ptr_++; } // If followed by ':', it's a name. if (PeekCurrentChar() == ':') { - str_val_.assign(token_start_, current_ptr_); + token_state_.str_val.assign(token_state_.token_start, current_ptr_); current_ptr_++; // skip ':' return TokKind::kName; } // If followed by '=', it's a attribute name. if (PeekCurrentChar() == '=') { - str_val_.assign(token_start_, current_ptr_); + token_state_.str_val.assign(token_state_.token_start, current_ptr_); current_ptr_++; // skip '=' return TokKind::kAttributeName; } absl::string_view identifier = - StringPieceFromPointers(token_start_, current_ptr_); + StringPieceFromPointers(token_state_.token_start, current_ptr_); + + // Primitive type strings are reserved words. The exception is 'tuple' whose + // type is represented using nested parentheses without the string 'tuple'. + if (primitive_util::IsPrimitiveTypeName(identifier)) { + PrimitiveType primitive_type = + primitive_util::StringToPrimitiveType(identifier).ValueOrDie(); + if (primitive_type != TUPLE) { + token_state_.primitive_type_val = primitive_type; + return TokKind::kPrimitiveType; + } + } // See if this is a keyword. #define KEYWORD(STR) \ @@ -255,21 +277,23 @@ TokKind HloLexer::LexIdentifier() { KEYWORD(ROOT); KEYWORD(maximal); KEYWORD(replicated); + KEYWORD(sparse); #undef KEYWORD { - auto consumable = RegexpStringPieceFromPointers(token_start_, buf_.end()); + auto consumable = + RegexpStringPieceFromPointers(token_state_.token_start, buf_.end()); static LazyRE2 dim_labels_pattern = { R"([0-9bf]{2,}_[0-9io]{2,}->[0-9bf]{2,})"}; if (RE2::Consume(&consumable, *dim_labels_pattern)) { current_ptr_ = consumable.begin(); - str_val_.assign(token_start_, current_ptr_); + token_state_.str_val.assign(token_state_.token_start, current_ptr_); return TokKind::kDimLabels; } } - str_val_ = string(identifier); + token_state_.str_val = string(identifier); return TokKind::kIdent; } @@ -277,13 +301,13 @@ 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())) { current_ptr_++; } - str_val_.assign(name_start, current_ptr_); + token_state_.str_val.assign(name_start, current_ptr_); return TokKind::kName; } return TokKind::kError; @@ -301,12 +325,14 @@ TokKind HloLexer::LexPercent() { // int ::= [-]?[0-9]+ // negative inf ::= '-inf' TokKind HloLexer::LexNumberOrPattern() { - auto consumable = RegexpStringPieceFromPointers(token_start_, buf_.end()); + auto consumable = + RegexpStringPieceFromPointers(token_state_.token_start, buf_.end()); static LazyRE2 float_pattern = { R"([-]?((\d+|\d+[.]\d*|\d*[.]\d+)([eE][+-]?\d+))|[-]?(\d+[.]\d*|\d*[.]\d+))"}; if (RE2::Consume(&consumable, *float_pattern)) { current_ptr_ = consumable.begin(); - CHECK(absl::SimpleAtod(string(token_start_, current_ptr_), &decimal_val_)); + CHECK(absl::SimpleAtod(string(token_state_.token_start, current_ptr_), + &token_state_.decimal_val)); return TokKind::kDecimal; } @@ -318,27 +344,28 @@ TokKind HloLexer::LexNumberOrPattern() { if (RE2::Consume(&consumable, *dim_labels_pattern)) { current_ptr_ = consumable.begin(); - str_val_.assign(token_start_, current_ptr_); + token_state_.str_val.assign(token_state_.token_start, current_ptr_); return TokKind::kDimLabels; } if (RE2::Consume(&consumable, *dxd_pattern)) { current_ptr_ = consumable.begin(); - str_val_.assign(token_start_, current_ptr_); + token_state_.str_val.assign(token_state_.token_start, current_ptr_); return TokKind::kDxD; } if (RE2::Consume(&consumable, *pad_pattern)) { current_ptr_ = consumable.begin(); - str_val_.assign(token_start_, current_ptr_); + token_state_.str_val.assign(token_state_.token_start, current_ptr_); return TokKind::kPad; } static LazyRE2 int_pattern = {R"([-]?\d+)"}; if (RE2::Consume(&consumable, *int_pattern)) { current_ptr_ = consumable.begin(); - auto slice = StringPieceFromPointers(token_start_, current_ptr_); - if (absl::SimpleAtoi(slice, &int64_val_)) { + auto slice = + StringPieceFromPointers(token_state_.token_start, current_ptr_); + if (absl::SimpleAtoi(slice, &token_state_.int64_val)) { return TokKind::kInt; } LOG(ERROR) << "Failed to parse int literal: " << slice; @@ -397,16 +424,17 @@ absl::string_view HloLexer::GetLine(LocTy loc) const { } // Lexes quoted string with escaping characters. If matched, the quoted string -// will be unescaped and stored to str_val_. +// will be unescaped and stored to token_state_.str_val. TokKind HloLexer::LexString() { - auto consumable = RegexpStringPieceFromPointers(token_start_, buf_.end()); + auto consumable = + RegexpStringPieceFromPointers(token_state_.token_start, buf_.end()); static LazyRE2 escaping_pattern = {R"("([^"\\]|\\.)*")"}; if (RE2::Consume(&consumable, *escaping_pattern)) { current_ptr_ = consumable.begin(); absl::string_view raw = - StringPieceFromPointers(token_start_ + 1, current_ptr_ - 1); + StringPieceFromPointers(token_state_.token_start + 1, current_ptr_ - 1); string error; - if (!absl::CUnescape(raw, &str_val_, &error)) { + if (!absl::CUnescape(raw, &token_state_.str_val, &error)) { LOG(ERROR) << "Failed unescaping string: " << raw << ". error: " << error; return TokKind::kError; } @@ -441,6 +469,8 @@ string TokKindToString(TokKind kind) { return "kRparen"; case TokKind::kArrow: return "kArrow"; + case TokKind::kLeq: + return "kLeq"; case TokKind::kw_HloModule: return "kw_HloModule"; case TokKind::kw_ENTRY: @@ -461,6 +491,10 @@ string TokKindToString(TokKind kind) { return "kw_inf"; case TokKind::kNegInf: return "kNegInf"; + case TokKind::kw_sparse: + return "kw_sparse"; + case TokKind::kPrimitiveType: + return "kPrimitiveType"; case TokKind::kName: return "kName"; case TokKind::kAttributeName: @@ -475,8 +509,6 @@ string TokKindToString(TokKind kind) { return "kIdent"; case TokKind::kString: return "kString"; - case TokKind::kShape: - return "kShape"; case TokKind::kInt: return "kInt"; case TokKind::kDecimal: diff --git a/tensorflow/compiler/xla/service/hlo_lexer.h b/tensorflow/compiler/xla/service/hlo_lexer.h index 3e2f8bcd52f9043f161197756a2060b28dded1d9..94fac3cd8e9da7f273e7e521e21510f5188702e6 100644 --- a/tensorflow/compiler/xla/service/hlo_lexer.h +++ b/tensorflow/compiler/xla/service/hlo_lexer.h @@ -19,7 +19,7 @@ limitations under the License. #include #include "absl/strings/string_view.h" -#include "tensorflow/compiler/xla/service/hlo_token.h" +#include "tensorflow/compiler/xla/shape.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/platform/logging.h" @@ -28,6 +28,58 @@ limitations under the License. namespace xla { +// Defines different kinds of tokens used by the HLO lexer. +// +// You shouldn't need to use this directly unless you're using HloLexer +// directly, and you probably don't need to do that. Use hlo_parser instead. +enum class TokKind { + // Markers + kEof, + kError, + + // Tokens with no info. + kEqual, // = + kComma, // , + kColon, // : + kLsquare, + kRsquare, // [ ] + kLbrace, + kRbrace, // { } + kLparen, + kRparen, // ( ) + + kArrow, // -> + kLeq, // <= + + // Keywords + kw_HloModule, + kw_ENTRY, + kw_ROOT, + kw_true, + kw_false, + kw_maximal, + kw_replicated, + kw_nan, + kw_inf, + kw_sparse, + + kNegInf, // -inf + + // Typed tokens. + kPrimitiveType, // F32, PRED, etc. + kName, // %foo + kAttributeName, // dimensions= + kDimLabels, // [0-9bf]{2,}_[0-9io]{2,}->[0-9bf]{2,} + kDxD, // [0-9]+(x[0-9]+)+ + kPad, // [0-9]+_[0-9]+(_[0-9]+)?(x[0-9]+_[0-9]+(_[0-9]+)?)* + kIdent, // other identifiers + kString, // "abcd\"\n" + kInt, // 42 + kDecimal, // 4.2 +}; + +string TokKindToString(TokKind kind); + // Lexer for the HloModule::ToString() format text. // // This class is meant to be used by hlo_parser.cc. You shouldn't need to use @@ -38,9 +90,9 @@ class HloLexer { current_ptr_ = buf_.begin(); } - TokKind Lex() { return current_kind_ = LexToken(); } + TokKind Lex() { return token_state_.current_kind = LexToken(); } - TokKind GetKind() const { return current_kind_; } + TokKind GetKind() const { return token_state_.current_kind; } string GetStrVal() const { switch (GetKind()) { case TokKind::kName: @@ -50,28 +102,28 @@ class HloLexer { case TokKind::kPad: case TokKind::kString: case TokKind::kIdent: - return str_val_; + return token_state_.str_val; default: LOG(FATAL) << "This token does not have string value"; } } - Shape GetShapeVal() const { - CHECK(GetKind() == TokKind::kShape); - return shape_val_; - } tensorflow::int64 GetInt64Val() const { CHECK(GetKind() == TokKind::kInt); - return int64_val_; + return token_state_.int64_val; } double GetDecimalVal() const { CHECK(GetKind() == TokKind::kDecimal); - return decimal_val_; + return token_state_.decimal_val; + } + PrimitiveType GetPrimitiveTypeVal() const { + CHECK(GetKind() == TokKind::kPrimitiveType); + return token_state_.primitive_type_val; } typedef const char* LocTy; // Returns the location of the current token. - LocTy GetLoc() const { return token_start_; } + LocTy GetLoc() const { return token_state_.token_start; } // Returns the line and column of a location in the buffer. std::pair GetLineAndColumn(LocTy location) const; @@ -79,6 +131,9 @@ class HloLexer { // Returns the whole line given the location. absl::string_view GetLine(LocTy loc) const; + // Looks ahead one token and returns it. Lexer state is unchanged. + TokKind LookAhead(); + private: // Returns the current character. If it's neither the end of input buffer nor // an invalid character, moves the pointer forward. @@ -111,12 +166,15 @@ class HloLexer { const char* current_ptr_; // Information about the current token. - const char* token_start_ = nullptr; - TokKind current_kind_; - string str_val_; - Shape shape_val_; - tensorflow::int64 int64_val_; - double decimal_val_; + struct TokenState { + const char* token_start = nullptr; + TokKind current_kind; + string str_val; + tensorflow::int64 int64_val; + double decimal_val; + PrimitiveType primitive_type_val; + }; + TokenState token_state_; struct LineNoCacheTy { const char* last_query; 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_liveness_analysis_test.cc b/tensorflow/compiler/xla/service/hlo_liveness_analysis_test.cc index e0ae1173c6114f0bc6ef18b2cfff9d54ccfe2faf..436cccb1fb9ecf6f4efad772c700c611b28ce628 100644 --- a/tensorflow/compiler/xla/service/hlo_liveness_analysis_test.cc +++ b/tensorflow/compiler/xla/service/hlo_liveness_analysis_test.cc @@ -403,9 +403,9 @@ TEST_F(HloLivenessAnalysisTest, WhileWithOutfeed) { HloModule OutfeedLoop WhileBody { body_param = (s32[]) parameter(0) - token = token[] after-all() + token0 = token[] after-all() constant.2 = s32[] constant(2) - outfeed_tuple = (s32[]) outfeed(constant.2, token) + outfeed_tuple = (s32[]) outfeed(constant.2, token0) get-tuple-element.1 = s32[] get-tuple-element(body_param), index=0 constant.1 = s32[] constant(1) add = s32[] add(get-tuple-element.1, constant.1) @@ -436,9 +436,9 @@ TEST_F(HloLivenessAnalysisTest, NestedWhileWithOutfeed) { HloModule OutfeedLoop InnerWhileBody { body_param = (s32[]) parameter(0) - token = token[] after-all() + token0 = token[] after-all() constant.2 = s32[] constant(2) - outfeed_tuple = (s32[]) outfeed(constant.2, token) + outfeed_tuple = (s32[]) outfeed(constant.2, token0) get-tuple-element.1 = s32[] get-tuple-element(body_param), index=0 constant.1 = s32[] constant(1) add = s32[] add(get-tuple-element.1, constant.1) diff --git a/tensorflow/compiler/xla/service/hlo_matchers.cc b/tensorflow/compiler/xla/service/hlo_matchers.cc index 5269cad94d35be3dd1c009588bbe422ff1533364..d28e79d41ad5d58a8881cfb80d488684af26564f 100644 --- a/tensorflow/compiler/xla/service/hlo_matchers.cc +++ b/tensorflow/compiler/xla/service/hlo_matchers.cc @@ -237,8 +237,4 @@ void PrintTo(const HloInstruction* inst, ::std::ostream* os) { *os << (inst ? inst->ToString() : "nullptr"); } -void PrintTo(HloInstruction* inst, ::std::ostream* os) { - PrintTo(const_cast(inst), os); -} - } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_matchers.h b/tensorflow/compiler/xla/service/hlo_matchers.h index 1717770301e3666b0a1c23d20b7f2e3bac5f62e4..67488a6a9a0c9cba7f576f9036c3a0cbe1900fff 100644 --- a/tensorflow/compiler/xla/service/hlo_matchers.h +++ b/tensorflow/compiler/xla/service/hlo_matchers.h @@ -165,6 +165,7 @@ namespace opcode_matchers { } HLO_MATCHER(Abs); HLO_MATCHER(Add); +HLO_MATCHER(AllToAll); HLO_MATCHER(Bitcast); HLO_MATCHER(Broadcast); HLO_MATCHER(BatchNormGrad); @@ -177,7 +178,8 @@ HLO_MATCHER(Constant); HLO_MATCHER(Convert); HLO_MATCHER(Convolution); HLO_MATCHER(Copy); -HLO_MATCHER(CrossReplicaSum); +HLO_MATCHER(AllReduce); +HLO_MATCHER(CollectivePermute); HLO_MATCHER(Divide); HLO_MATCHER(Domain); HLO_MATCHER(DynamicSlice); @@ -310,8 +312,8 @@ inline ::testing::Matcher Shape( } inline ::testing::Matcher Shape( absl::string_view shape) { - return ::testing::MakeMatcher(new ::xla::testing::HloShapeMatcher( - ShapeUtil::ParseShapeString(shape).ValueOrDie())); + return ::testing::MakeMatcher( + new ::xla::testing::HloShapeMatcher(ParseShape(shape).ValueOrDie())); } inline ::testing::Matcher ShapeWithLayout( const class Shape& shape) { @@ -321,7 +323,7 @@ inline ::testing::Matcher ShapeWithLayout( inline ::testing::Matcher ShapeWithLayout( absl::string_view shape) { return ::testing::MakeMatcher(new ::xla::testing::HloShapeAndLayoutMatcher( - ShapeUtil::ParseShapeString(shape).ValueOrDie())); + ParseShape(shape).ValueOrDie())); } // Verifies the value of the HloSharing against the provided sharding object. @@ -383,7 +385,6 @@ std::vector Pointers(const Container& container) { // Tell GMock to print HloInstruction* by value, so error messages are nice. // Has to be in the same namespace as 'HloInstruction'. void PrintTo(const HloInstruction* inst, ::std::ostream* os); -void PrintTo(HloInstruction* inst, ::std::ostream* os); } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_memory_scheduler.cc b/tensorflow/compiler/xla/service/hlo_memory_scheduler.cc index 5cee865b7ad34eded1743d9d5455bb40febf6182..d2740bcce26f04c5d7c8b64cfdaea53e3c697855 100644 --- a/tensorflow/compiler/xla/service/hlo_memory_scheduler.cc +++ b/tensorflow/compiler/xla/service/hlo_memory_scheduler.cc @@ -73,7 +73,7 @@ class ListScheduler { // Construct and return a memory-minimizing sequence of HLO instructions // containing the given HLO computation. static StatusOr Run( - const HloComputation& computation, + HloComputation* computation, const TuplePointsToAnalysis& points_to_analysis, const LogicalBuffer::SizeFunction& size_function, const absl::flat_hash_map& @@ -98,7 +98,7 @@ class ListScheduler { // comparison operators. using Priority = std::pair; - ListScheduler(const HloComputation& computation, + ListScheduler(HloComputation* computation, const TuplePointsToAnalysis& points_to_analysis, const LogicalBuffer::SizeFunction& size_function, const absl::flat_hash_map& @@ -111,7 +111,7 @@ class ListScheduler { // instruction. An HLO instruction "uses" a LogicalBuffer if the // LogicalBuffer is in an operand of the instruction as indicated by // points-to analysis. - for (auto* instruction : computation.instructions()) { + for (auto* instruction : computation->instructions()) { absl::flat_hash_set instr_uses; for (auto* operand : instruction->operands()) { points_to_analysis.GetPointsToSet(operand).ForEachElement( @@ -126,13 +126,13 @@ class ListScheduler { // Create map containing the number of unscheduled uses (hlo instructions) // of each logical buffer. - for (auto* instruction : computation.instructions()) { + for (auto* instruction : computation->instructions()) { for (auto* buffer : points_to_analysis.GetBuffersDefinedByInstruction(instruction)) { unscheduled_use_count_[buffer] = 0; } } - for (auto* instruction : computation.instructions()) { + for (auto* instruction : computation->instructions()) { for (const LogicalBuffer* buffer : buffer_uses_.at(instruction)) { ++unscheduled_use_count_[buffer]; } @@ -141,7 +141,7 @@ class ListScheduler { // Buffers live out of the computation have an implicit use at the end of // the computation. for (const LogicalBuffer* live_out_buffer : - points_to_analysis.GetPointsToSet(computation.root_instruction()) + points_to_analysis.GetPointsToSet(computation->root_instruction()) .CreateFlattenedSet()) { ++unscheduled_use_count_[live_out_buffer]; } @@ -157,7 +157,7 @@ class ListScheduler { // HloInstruction, plus some cached metadata, saved for the purposes of making // BytesFreedIfScheduled fast. struct ReadyListEntry { - const HloInstruction* instruction; + HloInstruction* instruction; // The total size of all buffers defined by this instruction. int64 bytes_defined; @@ -171,7 +171,7 @@ class ListScheduler { }; // Creates a ReadyListEntry for the given instruction. - ReadyListEntry MakeReadyListEntry(const HloInstruction* instruction) { + ReadyListEntry MakeReadyListEntry(HloInstruction* instruction) { ReadyListEntry entry; entry.instruction = instruction; @@ -250,13 +250,13 @@ class ListScheduler { // Populate the ready list with instructions which have no operands or // control predecessors. absl::flat_hash_map unscheduled_pred_count; - for (auto* instruction : computation_.instructions()) { + for (auto* instruction : computation_->instructions()) { // TODO(b/34466113): Replace this and above with successors() or // predecessors() when these methods are added to HloInstruction. - for (const HloInstruction* user : instruction->users()) { + for (HloInstruction* user : instruction->users()) { unscheduled_pred_count[user]++; } - for (const HloInstruction* succ : instruction->control_successors()) { + for (HloInstruction* succ : instruction->control_successors()) { unscheduled_pred_count[succ]++; } } @@ -275,7 +275,7 @@ class ListScheduler { ready_instructions[inst] = it; }; - for (auto* instruction : computation_.instructions()) { + for (auto* instruction : computation_->instructions()) { if (instruction->operands().empty() && instruction->control_predecessors().empty()) { add_to_ready_queue(instruction); @@ -287,7 +287,7 @@ class ListScheduler { // schedule. auto best_it = ready_queue.end(); --best_it; - const HloInstruction* best = best_it->second.instruction; + HloInstruction* best = best_it->second.instruction; VLOG(2) << "Schedule instruction: " << best->ToShortString() << " Bytes freed: " << best_it->first.first; ready_queue.erase(best_it); @@ -348,13 +348,13 @@ class ListScheduler { } } } - CHECK_EQ(schedule.size(), computation_.instruction_count()); - CHECK_EQ(scheduled_instructions_.size(), computation_.instruction_count()); + CHECK_EQ(schedule.size(), computation_->instruction_count()); + CHECK_EQ(scheduled_instructions_.size(), computation_->instruction_count()); return schedule; } - const HloComputation& computation_; + HloComputation* computation_; const TuplePointsToAnalysis& points_to_analysis_; const LogicalBuffer::SizeFunction& size_function_; // Computations are analyzed in post-order. When scheduling an instruction @@ -386,13 +386,13 @@ int64 SumLogicalBufferSizes( } StatusOr ScheduleComputationHelper( - const HloComputation& computation, + HloComputation* computation, const TuplePointsToAnalysis& points_to_analysis, const LogicalBuffer::SizeFunction& size_function, const MemorySchedulerAlgorithm& algorithm, const absl::flat_hash_map& memory_by_computation) { - VLOG(2) << "Computation: " << computation.name(); + VLOG(2) << "Computation: " << computation->name(); if (algorithm) { return algorithm(computation, points_to_analysis, size_function, memory_by_computation); @@ -404,17 +404,17 @@ StatusOr ScheduleComputationHelper( } // namespace StatusOr DFSMemoryScheduler( - const HloComputation& computation, + HloComputation* computation, const TuplePointsToAnalysis& points_to_analysis, const LogicalBuffer::SizeFunction& size_function, const absl::flat_hash_map& memory_by_computation) { // These variables are a hack to prevent overflows. int64 cumulative_total_size = 0; - int64 total_hlos = computation.parent()->instruction_count(); + int64 total_hlos = computation->parent()->instruction_count(); absl::flat_hash_map extra_users; absl::flat_hash_map total_sizes; - for (const HloInstruction* hlo : computation.MakeInstructionPostOrder()) { + for (const HloInstruction* hlo : computation->MakeInstructionPostOrder()) { if (ListScheduler::IgnoreInstruction(*hlo)) { extra_users[hlo] = 0; total_sizes[hlo] = 0; @@ -448,8 +448,8 @@ StatusOr DFSMemoryScheduler( total_sizes[hlo] = std::min(total_sizes[hlo], cumulative_total_size); extra_users[hlo] = std::min(extra_users[hlo], total_hlos); } - CHECK_EQ(extra_users.size(), computation.instruction_count()); - CHECK_EQ(total_sizes.size(), computation.instruction_count()); + CHECK_EQ(extra_users.size(), computation->instruction_count()); + CHECK_EQ(total_sizes.size(), computation->instruction_count()); // Construct a total order based on DFS post-order, visiting operands in // decreasing cumulative extra user order, and next by cumulative size, with a @@ -459,7 +459,7 @@ StatusOr DFSMemoryScheduler( sequence.push_back(hlo); return Status::OK(); }); - TF_RETURN_IF_ERROR(computation.AcceptWithOperandOrder( + TF_RETURN_IF_ERROR(computation->AcceptWithOperandOrder( &visitor, [&extra_users, &total_sizes](const HloInstruction* a, const HloInstruction* b) { if (extra_users[a] != extra_users[b]) { @@ -470,12 +470,12 @@ StatusOr DFSMemoryScheduler( } return a->name() < b->name(); })); - CHECK_EQ(sequence.size(), computation.instruction_count()); + CHECK_EQ(sequence.size(), computation->instruction_count()); return sequence; } // namespace xla StatusOr ListMemoryScheduler( - const HloComputation& computation, + HloComputation* computation, const TuplePointsToAnalysis& points_to_analysis, const LogicalBuffer::SizeFunction& size_function, const absl::flat_hash_map& @@ -485,16 +485,16 @@ StatusOr ListMemoryScheduler( } StatusOr PostOrderMemoryScheduler( - const HloComputation& computation, + HloComputation* computation, const TuplePointsToAnalysis& points_to_analysis, const LogicalBuffer::SizeFunction& size_function, const absl::flat_hash_map& memory_by_computation) { - return HloInstructionSequence(computation.MakeInstructionPostOrder()); + return HloInstructionSequence(computation->MakeInstructionPostOrder()); } StatusOr DefaultMemoryScheduler( - const HloComputation& computation, + HloComputation* computation, const TuplePointsToAnalysis& points_to_analysis, const LogicalBuffer::SizeFunction& size_function, const absl::flat_hash_map& @@ -513,7 +513,7 @@ StatusOr DefaultMemoryScheduler( memory_by_computation)); TF_ASSIGN_OR_RETURN(const int64 list_memory, HeapSimulator::MinimumMemoryForComputation( - computation, list_sequence, points_to_analysis, + *computation, list_sequence, points_to_analysis, size_function, &memory_by_computation)); VLOG(2) << "Min-memory list sequence: " << HumanReadableNumBytes(list_memory); @@ -522,7 +522,7 @@ StatusOr DefaultMemoryScheduler( size_function, memory_by_computation)); TF_ASSIGN_OR_RETURN(const int64 dfs_memory, HeapSimulator::MinimumMemoryForComputation( - computation, dfs_sequence, points_to_analysis, + *computation, dfs_sequence, points_to_analysis, size_function, &memory_by_computation)); VLOG(2) << "Min-memory dfs sequence: " << HumanReadableNumBytes(dfs_memory); @@ -532,7 +532,7 @@ StatusOr DefaultMemoryScheduler( memory_by_computation)); TF_ASSIGN_OR_RETURN(const int64 post_order_memory, HeapSimulator::MinimumMemoryForComputation( - computation, post_order_sequence, points_to_analysis, + *computation, post_order_sequence, points_to_analysis, size_function, &memory_by_computation)); VLOG(2) << "Min-memory post order sequence: " << HumanReadableNumBytes(post_order_memory); @@ -555,17 +555,17 @@ StatusOr DefaultMemoryScheduler( } StatusOr ScheduleModule( - const HloModule& module, const LogicalBuffer::SizeFunction& size_function, + HloModule* module, const LogicalBuffer::SizeFunction& size_function, const MemorySchedulerAlgorithm& algorithm) { - HloSchedule schedule(&module); + HloSchedule schedule(module); TF_ASSIGN_OR_RETURN(std::unique_ptr points_to_analysis, - TuplePointsToAnalysis::Run(&module)); + TuplePointsToAnalysis::Run(module)); absl::flat_hash_map memory_by_computation; - for (const auto* computation : module.MakeComputationPostOrder()) { + for (auto* computation : module->MakeComputationPostOrder()) { if (!computation->IsFusionComputation()) { TF_ASSIGN_OR_RETURN(HloInstructionSequence computation_sequence, ScheduleComputationHelper( - *computation, *points_to_analysis, size_function, + computation, *points_to_analysis, size_function, algorithm, memory_by_computation)); memory_by_computation[computation] = HeapSimulator::MinimumMemoryForComputation( @@ -583,11 +583,11 @@ StatusOr ScheduleModule( } StatusOr ScheduleComputation( - const HloComputation& computation, + HloComputation* computation, const LogicalBuffer::SizeFunction& size_function) { - CHECK(!computation.IsFusionComputation()); + CHECK(!computation->IsFusionComputation()); TF_ASSIGN_OR_RETURN(std::unique_ptr points_to_analysis, - TuplePointsToAnalysis::Run(computation.parent())); + TuplePointsToAnalysis::Run(computation->parent())); absl::flat_hash_map empty_map; return ScheduleComputationHelper(computation, *points_to_analysis, size_function, nullptr, empty_map); @@ -600,7 +600,24 @@ HloMemoryScheduler::HloMemoryScheduler( StatusOr HloMemoryScheduler::Run(HloModule* module) { TF_ASSIGN_OR_RETURN(HloSchedule schedule, - ScheduleModule(*module, size_function_, algorithm_)); + ScheduleModule(module, size_function_, algorithm_)); + TF_RETURN_IF_ERROR(module->set_schedule(std::move(schedule))); + return true; +} + +StatusOr HloTrivialScheduler::Run(HloModule* module) { + HloSchedule schedule(module); + for (HloComputation* computation : module->MakeComputationPostOrder()) { + if (!computation->IsFusionComputation()) { + HloInstructionSequence& computation_sequence = + schedule.GetOrCreateSequence(computation); + TF_RETURN_IF_ERROR(computation->Accept( + [&computation_sequence](HloInstruction* instruction) { + computation_sequence.push_back(instruction); + return Status::OK(); + })); + } + } TF_RETURN_IF_ERROR(module->set_schedule(std::move(schedule))); return true; } diff --git a/tensorflow/compiler/xla/service/hlo_memory_scheduler.h b/tensorflow/compiler/xla/service/hlo_memory_scheduler.h index a4c1d3db8170a1725043def576f913e09b352e5d..7227bfb27c74758d2b79e404afc9eb97a1ca894d 100644 --- a/tensorflow/compiler/xla/service/hlo_memory_scheduler.h +++ b/tensorflow/compiler/xla/service/hlo_memory_scheduler.h @@ -36,14 +36,14 @@ namespace xla { // that describes buffer aliasing, together with a target-specific size function // that maps a tensor's logical size to its padded size. typedef std::function( - const HloComputation&, const TuplePointsToAnalysis&, + HloComputation*, const TuplePointsToAnalysis&, const LogicalBuffer::SizeFunction&, const absl::flat_hash_map&)> MemorySchedulerAlgorithm; // List scheduler StatusOr ListMemoryScheduler( - const HloComputation& computation, + HloComputation* computation, const TuplePointsToAnalysis& points_to_analysis, const LogicalBuffer::SizeFunction& size_function, const absl::flat_hash_map& @@ -51,7 +51,7 @@ StatusOr ListMemoryScheduler( // DFS-order scheduler StatusOr DFSMemoryScheduler( - const HloComputation& computation, + HloComputation* computation, const TuplePointsToAnalysis& points_to_analysis, const LogicalBuffer::SizeFunction& size_function, const absl::flat_hash_map& @@ -59,7 +59,7 @@ StatusOr DFSMemoryScheduler( // Naive Post Order scheduler StatusOr PostOrderMemoryScheduler( - const HloComputation& computation, + HloComputation* computation, const TuplePointsToAnalysis& points_to_analysis, const LogicalBuffer::SizeFunction& size_function, const absl::flat_hash_map& @@ -69,7 +69,7 @@ StatusOr PostOrderMemoryScheduler( // and the DFS scheduler, and chooses whichever returns a lower min-memory, // not accounting for fragmentation. StatusOr DefaultMemoryScheduler( - const HloComputation& computation, + HloComputation* computation, const TuplePointsToAnalysis& points_to_analysis, const LogicalBuffer::SizeFunction& size_function, const absl::flat_hash_map& @@ -79,13 +79,13 @@ StatusOr DefaultMemoryScheduler( // the computation. size_function is the function returning the number of bytes // required for a LogicalBuffer. StatusOr ScheduleModule( - const HloModule& module, const LogicalBuffer::SizeFunction& size_function, + HloModule* module, const LogicalBuffer::SizeFunction& size_function, const MemorySchedulerAlgorithm& algorithm = {}); // Computes the schedule for a single computation. // Currently only used by the GPU backend. StatusOr ScheduleComputation( - const HloComputation& computation, + HloComputation* computation, const LogicalBuffer::SizeFunction& size_function); // A pass which schedules the HLO instructions in a module. The HloModule's @@ -108,6 +108,15 @@ class HloMemoryScheduler : public HloModulePass { MemorySchedulerAlgorithm algorithm_; }; +// A pass which produces a naive, but correct schedule. The schedule is produced +// using a DFS traversal of the graph with no attempt to minimize memory use. +class HloTrivialScheduler : public HloModulePass { + public: + absl::string_view name() const override { return "hlo-trivial-scheduler"; } + + StatusOr Run(HloModule* module) override; +}; + // A trivial pass which clears the schedule currently set on the // HloModule. After this pass runs HloModudle::has_schedule will return false. class HloDescheduler : public HloModulePass { diff --git a/tensorflow/compiler/xla/service/hlo_memory_scheduler_test.cc b/tensorflow/compiler/xla/service/hlo_memory_scheduler_test.cc index 214119fba881c4411a262cd4227b5cc49cef0d14..bc0d7e2bc00eab014f2660c95a51b966642eaee9 100644 --- a/tensorflow/compiler/xla/service/hlo_memory_scheduler_test.cc +++ b/tensorflow/compiler/xla/service/hlo_memory_scheduler_test.cc @@ -65,7 +65,7 @@ TEST_F(HloSchedulingTest, LastUseScheduledFirst) { auto sub = builder.AddInstruction( HloInstruction::CreateBinary(vec, HloOpcode::kSubtract, add, negate)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); HloMemoryScheduler scheduler([](const BufferValue& buffer) { @@ -78,7 +78,7 @@ TEST_F(HloSchedulingTest, LastUseScheduledFirst) { TF_ASSERT_OK(module->schedule().Verify()); // Verify that all instructions are in the sequence. - const std::vector& sequence = + const std::vector& sequence = module->schedule().sequence(module->entry_computation()).instructions(); EXPECT_EQ(module->entry_computation()->instruction_count(), sequence.size()); @@ -124,9 +124,9 @@ ENTRY root { }; TF_ASSERT_OK_AND_ASSIGN( HloSchedule schedule, - ScheduleModule(*module, size_fn, ListMemoryScheduler)); + ScheduleModule(module.get(), size_fn, ListMemoryScheduler)); // Verify that all instructions are in the sequence. - const std::vector& sequence = + const std::vector& sequence = schedule.sequence(module->entry_computation()).instructions(); EXPECT_EQ(module->entry_computation()->instruction_count(), sequence.size()); @@ -172,15 +172,16 @@ TEST_F(HloSchedulingTest, TuplesAreAccountedCorrectly) { builder.AddInstruction(HloInstruction::CreateBinary(r1f32, HloOpcode::kAdd, tuple_elm, abs_abs2)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); TF_ASSERT_OK_AND_ASSIGN(HloSchedule schedule, - ScheduleModule(*module, - [](const BufferValue& buffer) { - return ShapeUtil::ByteSizeOf( - buffer.shape(), TUPLE_SIZE); - }, - ListMemoryScheduler)); + ScheduleModule( + module.get(), + [](const BufferValue& buffer) { + return ShapeUtil::ByteSizeOf(buffer.shape(), + TUPLE_SIZE); + }, + ListMemoryScheduler)); // Verify that all instructions are in the sequence. EXPECT_EQ(module->entry_computation()->instruction_count(), @@ -218,19 +219,19 @@ TEST_F(HloSchedulingTest, MultiOutputFusionAccountedCorrectly) { builder.AddInstruction( HloInstruction::CreateBinary(r1f32, HloOpcode::kAdd, tuple_elm, exp)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto* computation = module->AddEntryComputation(builder.Build()); auto fusion = computation->CreateFusionInstruction( {tuple, mul, add}, HloInstruction::FusionKind::kLoop); TF_ASSERT_OK_AND_ASSIGN(HloSchedule schedule, - ScheduleModule(*module, - [](const BufferValue& buffer) { - return ShapeUtil::ByteSizeOf( - buffer.shape(), 2); - }, - ListMemoryScheduler)); + ScheduleModule( + module.get(), + [](const BufferValue& buffer) { + return ShapeUtil::ByteSizeOf(buffer.shape(), 2); + }, + ListMemoryScheduler)); // Verify that all instructions are in the sequence. EXPECT_EQ(module->entry_computation()->instruction_count(), @@ -242,7 +243,7 @@ TEST_F(HloSchedulingTest, MultiOutputFusionAccountedCorrectly) { } TEST_F(HloSchedulingTest, HeapSimulatorAccountsForSubcomputations) { - auto module = CreateNewModule(); + auto module = CreateNewUnverifiedModule(); const Shape r1f32 = ShapeUtil::MakeShape(F32, {4}); // param != 0 @@ -252,7 +253,7 @@ TEST_F(HloSchedulingTest, HeapSimulatorAccountsForSubcomputations) { HloInstruction::CreateParameter(0, r1f32, "cond_param")); HloInstruction* zero_vector = cond_builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR2({{0, 0, 0, 0}}))); + LiteralUtil::CreateR1({0, 0, 0, 0}))); cond_builder.AddInstruction(HloInstruction::CreateBinary( ShapeUtil::MakeShape(PRED, {}), HloOpcode::kNe, cond_param, zero_vector)); auto cond_computation = module->AddEmbeddedComputation(cond_builder.Build()); @@ -284,7 +285,7 @@ TEST_F(HloSchedulingTest, HeapSimulatorAccountsForSubcomputations) { }; TF_ASSERT_OK_AND_ASSIGN( HloSchedule schedule, - ScheduleModule(*module, size_fn, ListMemoryScheduler)); + ScheduleModule(module.get(), size_fn, ListMemoryScheduler)); // Verify that all instructions are in the sequence. auto entry_computation = module->entry_computation(); EXPECT_EQ(module->entry_computation()->instruction_count(), @@ -309,5 +310,40 @@ TEST_F(HloSchedulingTest, HeapSimulatorAccountsForSubcomputations) { .ValueOrDie()); } +TEST_F(HloSchedulingTest, TrivialScheduler) { + const char* const hlo_string = R"( +HloModule ModuleWithWhile + +body { + param.b = (s32[], s32[]) parameter(0) + gte.0 = s32[] get-tuple-element(param.b), index=0 + gte.1 = s32[] get-tuple-element(param.b), index=1 + add = s32[] add(gte.0, gte.1) + ROOT tuple = (s32[], s32[]) tuple(gte.0, add) +} + +cond { + param.c = (s32[], s32[]) parameter(0) + ROOT constant = pred[] constant(true) +} + +ENTRY main { + init = (s32[], s32[]) parameter(0) + ROOT while = (s32[], s32[]) while(init), condition=cond, body=body +} +)"; + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseHloString(hlo_string)); + EXPECT_FALSE(module->has_schedule()); + TF_ASSERT_OK(HloTrivialScheduler().Run(module.get()).status()); + ASSERT_TRUE(module->has_schedule()); + TF_ASSERT_OK(module->schedule().Verify()); + + // Verify that a clone of the module also has a schedule. + std::unique_ptr clone = module->Clone(); + ASSERT_TRUE(clone->has_schedule()); + TF_ASSERT_OK(clone->schedule().Verify()); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_module.cc b/tensorflow/compiler/xla/service/hlo_module.cc index 6845c27a91845ef971dc2d82266200bfccb25533..258f918f47a313b4b89fb260457b1b119dc16177 100644 --- a/tensorflow/compiler/xla/service/hlo_module.cc +++ b/tensorflow/compiler/xla/service/hlo_module.cc @@ -41,18 +41,6 @@ HloModule::HloModule(const string& name, const HloModuleConfig& config) config_(config), unique_id_(next_unique_module_id_++) {} -StatusOr HloModule::LaunderConstInstructionFromModule( - const HloInstruction* hlo) { - if (hlo == nullptr) { - return nullptr; - } - - TF_RET_CHECK(hlo->GetModule() == this); - - // TODO(b/78350259): Eliminate const laundering. - return const_cast(hlo); -} - Status HloModule::set_schedule(HloSchedule schedule) { TF_RET_CHECK(schedule.module() == this); TF_RETURN_IF_ERROR(schedule.Verify()); @@ -119,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(); @@ -252,8 +239,10 @@ HloModuleProto HloModule::ToProto() const { *proto.mutable_schedule() = schedule().ToProto().ValueOrDie(); } *proto.mutable_host_program_shape() = - entry_computation_layout().ComputeProgramShape(); + entry_computation_layout().ComputeProgramShape().ToProto(); *proto.mutable_input_output_alias() = input_output_alias_config().ToProto(); + *proto.mutable_dynamic_parameter_binding() = + dynamic_parameter_binding().ToProto(); return proto; } @@ -267,7 +256,7 @@ StatusOr> HloModule::CreateFromProto( // the entry parameters and root. TF_RET_CHECK(proto.has_host_program_shape()) << "No program shape found in the proto"; - const auto& expected_program_shape = proto.host_program_shape(); + ProgramShape expected_program_shape(proto.host_program_shape()); TF_RET_CHECK(expected_program_shape.parameters_size() == module_config.entry_computation_layout().parameter_count()); for (int i = 0; i < expected_program_shape.parameters_size(); ++i) { @@ -314,11 +303,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) { @@ -330,12 +318,17 @@ StatusOr> HloModule::CreateFromProto( } TF_RET_CHECK(module->entry_computation_ != nullptr); - TF_ASSIGN_OR_RETURN(module->input_output_alias_config_, - HloInputOutputAliasConfig::CreateFromProto( - result_shape, proto.input_output_alias())); + TF_ASSIGN_OR_RETURN( + module->input_output_alias_config_, + HloInputOutputAliasConfig::CreateFromProto( + entry->ComputeProgramShape().result(), proto.input_output_alias())); // Because we didn't uniquify the names or the ids, double-check that the // instruction and computation names and ids are unique from the proto. + TF_ASSIGN_OR_RETURN(module->dynamic_parameter_binding_, + DynamicParameterBinding::CreateFromProto( + proto.dynamic_parameter_binding())); + absl::flat_hash_set computation_names; absl::flat_hash_set instruction_names; absl::flat_hash_set computation_ids; @@ -374,9 +367,9 @@ StatusOr HloModule::CreateModuleConfigFromProto( const HloModuleProto& module, const DebugOptions& debug_options) { TF_RET_CHECK(module.has_host_program_shape()) << "No program shape found in the proto"; - const auto& program_shape = module.host_program_shape(); + ProgramShape program_shape(module.host_program_shape()); - HloModuleConfig module_config(program_shape); + HloModuleConfig module_config(ProgramShape{program_shape}); module_config.set_debug_options(debug_options); // The module config is constructed with default layouts regardless of what is @@ -397,15 +390,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 @@ -416,9 +406,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; @@ -507,7 +497,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 : @@ -520,19 +510,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()); } @@ -570,11 +560,28 @@ std::unique_ptr HloModule::Clone(const string& suffix) const { std::unique_ptr HloModule::Clone(const HloModuleConfig& config, const string& suffix) const { VLOG(1) << "Cloning module :" << name_ << " --> " << suffix << "\n"; - auto module = absl::make_unique(name_ + "-" + suffix, config); + auto module = absl::make_unique( + absl::StrCat(name_, suffix.empty() ? "" : "-", suffix), config); HloCloneContext context(module.get(), suffix); auto cloned_computation = entry_computation_->Clone(suffix, &context); module->AddEntryComputation(std::move(cloned_computation)); + + if (has_schedule() && schedule().Verify().ok()) { + HloSchedule clone_schedule(module.get()); + for (HloComputation* computation : computations()) { + if (schedule().is_computation_scheduled(computation)) { + HloInstructionSequence& clone_sequence = + clone_schedule.GetOrCreateSequence( + context.GetComputation(computation)); + for (const HloInstruction* instruction : + schedule().sequence(computation).instructions()) { + clone_sequence.push_back(context.GetInstruction(instruction)); + } + } + } + TF_CHECK_OK(module->set_schedule(std::move(clone_schedule))); + } return module; } diff --git a/tensorflow/compiler/xla/service/hlo_module.h b/tensorflow/compiler/xla/service/hlo_module.h index 5dc795fabec5d8d794635ef6965c4d065b0b75a6..f1310e4b270898a21dbb4f86123edde4ba8993d0 100644 --- a/tensorflow/compiler/xla/service/hlo_module.h +++ b/tensorflow/compiler/xla/service/hlo_module.h @@ -28,6 +28,7 @@ limitations under the License. #include "absl/types/optional.h" #include "absl/types/span.h" #include "tensorflow/compiler/xla/iterator_util.h" +#include "tensorflow/compiler/xla/service/dynamic_parameter_binding.h" #include "tensorflow/compiler/xla/service/hlo.pb.h" #include "tensorflow/compiler/xla/service/hlo_clone_context.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" @@ -103,11 +104,7 @@ class HloModule { HloCloneContext* context = nullptr); // Return a pointer to the entry computation of the module. - const HloComputation* entry_computation() const { - CHECK_NE(nullptr, entry_computation_); - return entry_computation_; - } - HloComputation* entry_computation() { + HloComputation* entry_computation() const { CHECK_NE(nullptr, entry_computation_); return entry_computation_; } @@ -135,6 +132,14 @@ class HloModule { return config_.entry_computation_layout(); } + // Generates a hash value of an HLO module. Hash considers + // information on opcode, shape, operands, and typically a root instruction. + // This function returns the same hash value for equivalent HLO modules, + // with respect to HloInstruction::Identical() method. + uint64 Hash() const { + return entry_computation()->root_instruction()->Hash(); + } + // Gets the computations in this module. // // Returns a view of HloComputation*s, so you can iterate over this in the @@ -232,29 +237,20 @@ class HloModule { return input_output_alias_config_; } + // DynamicParameterBinding holds the list of bindings that indicates which + // parameter dimensions are dynamic and which parameters represent their + // runtime value. + DynamicParameterBinding& dynamic_parameter_binding() { + return dynamic_parameter_binding_; + } + const DynamicParameterBinding& dynamic_parameter_binding() const { + return dynamic_parameter_binding_; + } + // Returns an id that is unique to this module across all modules created over // the lifetime of this process. int unique_id() const { return unique_id_; } - // Returns a non-const version of the passed-in const HloInstruction*. This is - // safe on the argument that if you have a non-const module, then you can - // access all instructions in the module as non-const. - // - // Returns an error if the passed-in instruction is not from this module, - // except that it is allowed to pass in a null pointer. - // - // TODO(b/78350259): Eliminate const laundering. The argument above is not - // reliable since at any time someone could add or discover a way for a - // non-const module to transitively contain a const HloInstruction. The - // reliable way to do this would be to create a const laundering map from a - // module, mapping each encountered HloInstruction to its non-const version - // and then look up each instruction in need of laundering in that map, but - // this is much more expensive and complicated. This returns a Status instead - // of doing a CHECK-failure in part to make it strongly apparent that this is - // something that can fail. - StatusOr LaunderConstInstructionFromModule( - const HloInstruction* hlo); - // Sets the schedule of the module to the given schedule. Status set_schedule(HloSchedule schedule); @@ -304,6 +300,9 @@ class HloModule { // alias_config indicates the alias information of input/output buffers that // are expected from the module. HloInputOutputAliasConfig input_output_alias_config_; + + // Bindings for dynamic parameter mapping. + DynamicParameterBinding dynamic_parameter_binding_; }; } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_module_dce.cc b/tensorflow/compiler/xla/service/hlo_module_dce.cc index 31d26cc51e8217234526bbfeb83510aadf2c27b5..6b72ba128664d27c51aa8dcfa61fe959a0160c73 100644 --- a/tensorflow/compiler/xla/service/hlo_module_dce.cc +++ b/tensorflow/compiler/xla/service/hlo_module_dce.cc @@ -49,7 +49,7 @@ StatusOr RunWhileDCE(HloModule* module, HloLivenessAnalysis* liveness) { auto* while_body_param = while_body_comp->parameter_instruction(0); auto* while_body_root = while_body_comp->root_instruction(); - if (!ShapeUtil::IsTuple(xla_while->shape()) || + if (!xla_while->shape().IsTuple() || while_body_root->opcode() != HloOpcode::kTuple) { // Only run DCE on tuple-shaped while loops where body root is Tuple, // with no I/O instructions. diff --git a/tensorflow/compiler/xla/service/hlo_module_dce_test.cc b/tensorflow/compiler/xla/service/hlo_module_dce_test.cc index bf66cc6bc37a5e11c9ecfc07a62ba0ea5ca11a03..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 @@ -373,9 +371,9 @@ TEST_F(HloModuleDceTest, WhileWithOutfeed) { HloModule OutfeedLoop WhileBody { body_param = (s32[]) parameter(0) - token = token[] after-all() + token0 = token[] after-all() constant.2 = s32[] constant(2) - outfeed_tuple = (s32[]) outfeed(constant.2, token) + outfeed_tuple = (s32[]) outfeed(constant.2, token0) get-tuple-element.1 = s32[] get-tuple-element(body_param), index=0 constant.1 = s32[] constant(1) add = s32[] add(get-tuple-element.1, constant.1) diff --git a/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc b/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc index b4aac4c8076cb69647d42c6243bc969d06d0709e..47734bc55cc00d605f4e318400be88639450343c 100644 --- a/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc +++ b/tensorflow/compiler/xla/service/hlo_module_group_metadata.cc @@ -79,36 +79,36 @@ Status HloModuleGroupMetadata::Build() { return Status::OK(); } - std::vector peers; - if (IsChannelInstruction(hlo)) { - peers.push_back(PeerComputation(hlo)); - } else if (hlo->IsCrossModuleAllReduce()) { - for (HloInstruction* instr : GetAllReduceGroup(*hlo->all_reduce_id())) { - if (instr == hlo) { - continue; + if (IsChannelInstruction(hlo) || hlo->IsCrossModuleAllReduce()) { + std::vector peers; + if (IsChannelInstruction(hlo)) { + peers.push_back(PeerComputation(hlo)); + } else if (hlo->IsCrossModuleAllReduce()) { + for (HloInstruction* instr : GetAllReduceGroup(*hlo->all_reduce_id())) { + if (instr == hlo) { + continue; + } + peers.push_back(instr->parent()); } - peers.push_back(instr->parent()); } - } - - // Add the parent computation of this channel (or all-reduce) instruction - // and its peer computation(s) (both must be while computations) as - // companions. - for (HloComputation* peer_computation : peers) { - const TrackedInstruction* peer_tracked = - GetTrackedInstruction(peer_computation); - TF_RET_CHECK(peer_tracked != nullptr) - << "Peer instruction is not a possible companion"; - TF_RET_CHECK(*tracked == *peer_tracked) - << "Peer instruction does not match the computation kind"; - TF_RETURN_IF_ERROR( - AddCompanion(tracked->instruction(), peer_tracked->instruction())); - tracked_instructions_comms_[tracked->instruction()].push_back(hlo); - } - // Add the parents of companion instructions (they must be all of the same - // kind of instructions, opcode wise) as companions. - if (IsCompanionInstruction(hlo)) { + // Add the parent computation of this channel (or all-reduce) instruction + // and its peer computation(s) (both must be while computations) as + // companions. + for (HloComputation* peer_computation : peers) { + const TrackedInstruction* peer_tracked = + GetTrackedInstruction(peer_computation); + TF_RET_CHECK(peer_tracked != nullptr) + << "Peer instruction is not a possible companion"; + TF_RET_CHECK(*tracked == *peer_tracked) + << "Peer instruction does not match the computation kind"; + TF_RETURN_IF_ERROR( + AddCompanion(tracked->instruction(), peer_tracked->instruction())); + tracked_instructions_comms_[tracked->instruction()].push_back(hlo); + } + } else if (IsCompanionInstruction(hlo)) { + // Add the parents of companion instructions (they must be all of the same + // kind of instructions, opcode wise) as companions. for (HloInstruction* companion : Companions(hlo)) { const TrackedInstruction* companion_tracked = GetTrackedInstruction(companion->parent()); @@ -118,6 +118,7 @@ Status HloModuleGroupMetadata::Build() { companion_tracked->instruction())); } } + return Status::OK(); }; @@ -198,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( @@ -509,7 +510,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 928df0f5a7444ad877961a5de970c752e1d024da..3ed95c10504141139d83eb8679a0b8144b15ad0d 100644 --- a/tensorflow/compiler/xla/service/hlo_module_group_metadata.h +++ b/tensorflow/compiler/xla/service/hlo_module_group_metadata.h @@ -38,7 +38,7 @@ namespace xla { // Class for bookkeeping the information on the given modules, in particular on // the interaction between computations. // -// Companion instructions are one of the information collected as we build the +// Companion instructions are one piece of information collected as we build the // metadata. For example, for each While instruction, companion instructions // refer to a set of While instructions in other computations that communicate // with each other. @@ -51,6 +51,13 @@ namespace xla { // } While_4() { Recv(0) } // } // +// Each instruction can belong to at most one companion set: While_0 and While_5 +// are in the same set even though they don't communicate with each other, +// because they both communicate with While_2. +// +// A send and the matching recv must both have the same level of nesting of +// companion instructions. +// // Companion instructions are used to detect cycles in the graph and also for // global scheduling. class HloModuleGroupMetadata { @@ -171,7 +178,7 @@ class HloModuleGroupMetadata { // 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)); } @@ -215,11 +222,8 @@ class HloModuleGroupMetadata { // * Each channel has all 4 instructions (Send, Recv, SendDone, RecvDone). // * The shape of channel instructions match. // * The nest level of channel instructions match. - // * Channel instructions are used in allowed computations; i.e., in the + // * Channel instructions are used in allowed computations, i.e., in the // entry computation of the module or condition/body of While computations. - // - // TODO(b/62064342): Currently, HloModuleGroupScheduler checks if there is a - // cycle in the graph, but it would be good to verify here. Status VerifyChannelInstructions(); // Adds metadata that the given two instructions are companions. @@ -231,8 +235,8 @@ class HloModuleGroupMetadata { Status CheckCommunicatingInstruction(HloInstruction* instruction) const; // Performs a consistency check on the companion sets built for the input - // modules. Check that a companion set does not include instructions from the - // same module/device. + // modules. Checks that each instruction in a companion set is in a different + // module/device. Status VerifyCompanionSets() const; // Retrieves a pointer to the stored TrackedInstruction associated with a diff --git a/tensorflow/compiler/xla/service/hlo_module_group_util.cc b/tensorflow/compiler/xla/service/hlo_module_group_util.cc index fddeb5f0a27a43ff9ca8b2b5d314bcfe91aaf0e6..91417bd2d9a6ca8a5192a37302e6a91e49a94d77 100644 --- a/tensorflow/compiler/xla/service/hlo_module_group_util.cc +++ b/tensorflow/compiler/xla/service/hlo_module_group_util.cc @@ -198,6 +198,8 @@ std::vector HloModuleGroupUtil::RootInstructions( for (HloComputation* computation : computations) { for (HloInstruction* instruction : computation->instructions()) { if (GlobalSuccessors(instruction).empty()) { + // An instruction that has no successors, e.g., an unused instruction, + // is in roots, even though it's not the ROOT of its computation. roots.push_back(instruction); } } diff --git a/tensorflow/compiler/xla/service/hlo_module_group_util.h b/tensorflow/compiler/xla/service/hlo_module_group_util.h index f21b44bcd98d77b831de5d8a6afa4f9ddd91d15d..862666b48c9aa423ba4eeea3052c17fcc1064fd2 100644 --- a/tensorflow/compiler/xla/service/hlo_module_group_util.h +++ b/tensorflow/compiler/xla/service/hlo_module_group_util.h @@ -49,7 +49,7 @@ class HloModuleGroupUtil { // Returns all unique successors of the instruction. This includes: // * successors in the same computation: users and control successors // * Send is a successor of Recv - // * RecvDone is a predecessor of Send + // * RecvDone is a successor of Send // * successors of companions (if the instruction is a companion while) // * successors' companions (for any successor that is a companion while) std::vector GlobalSuccessors(HloInstruction* instruction); diff --git a/tensorflow/compiler/xla/service/hlo_module_test.cc b/tensorflow/compiler/xla/service/hlo_module_test.cc index 39f38b417ab0e8b54864176d8d1e0ad1a422eca6..620cb7e01ad1a060915f5b73474f6950ab18122a 100644 --- a/tensorflow/compiler/xla/service/hlo_module_test.cc +++ b/tensorflow/compiler/xla/service/hlo_module_test.cc @@ -63,7 +63,7 @@ class HloModuleTest : public HloTestBase { TEST_F(HloModuleTest, OneComputationPostOrder) { // Create a module with a single computation. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(CreateConstantComputation()); EXPECT_THAT(module->MakeComputationPostOrder(), @@ -72,7 +72,7 @@ TEST_F(HloModuleTest, OneComputationPostOrder) { TEST_F(HloModuleTest, TwoComputationsPostOrder) { // Create a module with two unconnected computations. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation1 = module->AddEntryComputation(CreateConstantComputation()); auto computation2 = module->AddEmbeddedComputation(CreateConstantComputation()); @@ -88,7 +88,7 @@ TEST_F(HloModuleTest, TwoComputationsPostOrder) { TEST_F(HloModuleTest, CloneTest) { // Create and copy a module with a diamond call graph of computations. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation1 = module->AddEmbeddedComputation(CreateConstantComputation()); auto computation2 = @@ -111,7 +111,7 @@ TEST_F(HloModuleTest, CloneTest) { } TEST_F(HloModuleTest, CloneHasFusion) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); // Create the fused computation. HloComputation* fused_computation; @@ -154,7 +154,7 @@ TEST_F(HloModuleTest, CloneHasFusion) { TEST_F(HloModuleTest, DiamondComputationsPostOrder) { // Create a module with a diamond call graph of computations. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation1 = module->AddEmbeddedComputation(CreateConstantComputation()); auto computation2 = @@ -174,7 +174,7 @@ TEST_F(HloModuleTest, DiamondComputationsPostOrder) { TEST_F(HloModuleTest, LargeConstantToString) { // Create a module with a single computation. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto builder = HloComputation::Builder("Constant"); std::vector values(16, 42.0); builder.AddInstruction( @@ -194,8 +194,8 @@ TEST_F(HloModuleTest, LargeConstantToString) { } TEST_F(HloModuleTest, UniqueModuleId) { - auto module_a = CreateNewModule(); - auto module_b = CreateNewModule(); + auto module_a = CreateNewVerifiedModule(); + auto module_b = CreateNewVerifiedModule(); EXPECT_NE(module_a->unique_id(), module_b->unique_id()); } diff --git a/tensorflow/compiler/xla/service/hlo_opcode.h b/tensorflow/compiler/xla/service/hlo_opcode.h index e6bfb8025d4bfeba1d334d1f946e33841a2da092..94122ac38ff2a3f7053b19e55f9a400c80ae2134 100644 --- a/tensorflow/compiler/xla/service/hlo_opcode.h +++ b/tensorflow/compiler/xla/service/hlo_opcode.h @@ -47,6 +47,9 @@ namespace xla { #define HLO_OPCODE_LIST(V) \ V(kAbs, "abs") \ V(kAdd, "add") \ + V(kAddDependency, "add-dependency") \ + V(kAfterAll, "after-all", kHloOpcodeIsVariadic) \ + V(kAllReduce, "all-reduce") \ V(kAllToAll, "all-to-all") \ V(kAtan2, "atan2") \ V(kBatchNormGrad, "batch-norm-grad") \ @@ -68,7 +71,6 @@ namespace xla { V(kConvolution, "convolution") \ V(kCopy, "copy") \ V(kCos, "cosine") \ - V(kCrossReplicaSum, "cross-replica-sum") \ V(kCustomCall, "custom-call") \ V(kDivide, "divide") \ V(kDomain, "domain") \ @@ -83,7 +85,7 @@ namespace xla { V(kFusion, "fusion", kHloOpcodeIsVariadic) \ V(kGather, "gather") \ V(kGe, "greater-than-or-equal-to", kHloOpcodeIsComparison) \ - V(kAfterAll, "after-all", kHloOpcodeIsVariadic) \ + V(kGetDimensionSize, "get-dimension-size") \ V(kGetTupleElement, "get-tuple-element") \ V(kGt, "greater-than", kHloOpcodeIsComparison) \ V(kImag, "imag") \ diff --git a/tensorflow/compiler/xla/service/hlo_ordering.cc b/tensorflow/compiler/xla/service/hlo_ordering.cc index 23d41d91d6969ddf9062507e926ae39c1e1315d4..0cec61c257bb84e467290fb52ec9063a32ed558d 100644 --- a/tensorflow/compiler/xla/service/hlo_ordering.cc +++ b/tensorflow/compiler/xla/service/hlo_ordering.cc @@ -334,7 +334,7 @@ DependencyHloOrdering::DependencyHloOrdering(const HloModule* module) // ordering based on dependencies. ExecutesBefore will return true iff there // exists a path in the HLO computation graph from 'a' to 'b'. for (auto* computation : module->MakeNonfusionComputations()) { - predecessors_.emplace(computation, computation->ComputeReachability()); + predecessors_.emplace(computation, HloReachabilityMap::Build(computation)); } } @@ -356,8 +356,7 @@ void SequentialHloOrdering::Initialize() { // Create a map from instruction to its order position. TF_DCHECK_OK(schedule_.Verify()); for (const auto& computation_sequence : schedule_.sequences()) { - const std::vector& order = - computation_sequence.second.instructions(); + const auto& order = computation_sequence.second.instructions(); for (int i = 0; i < order.size(); ++i) { InsertOrDie(&order_position_, order[i], i); } @@ -368,17 +367,16 @@ 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); } -const std::vector* -SequentialHloOrdering::SequentialOrder( +const HloInstructionSequence* SequentialHloOrdering::SequentialOrder( const HloComputation& computation) const { return schedule_.is_computation_scheduled(&computation) - ? &schedule_.sequence(&computation).instructions() + ? &schedule_.sequence(&computation) : nullptr; } diff --git a/tensorflow/compiler/xla/service/hlo_ordering.h b/tensorflow/compiler/xla/service/hlo_ordering.h index 66313492eb2dd10ac9a6000639ddb8991b367c0f..a07214c22c0989a438f12219e136a7e76ee0dcce 100644 --- a/tensorflow/compiler/xla/service/hlo_ordering.h +++ b/tensorflow/compiler/xla/service/hlo_ordering.h @@ -26,6 +26,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_dataflow_analysis.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_module.h" +#include "tensorflow/compiler/xla/service/hlo_reachability.h" #include "tensorflow/compiler/xla/service/hlo_schedule.h" #include "tensorflow/compiler/xla/service/hlo_value.h" #include "tensorflow/compiler/xla/types.h" @@ -64,7 +65,7 @@ class HloOrdering { // Returns the sequential instruction order for the given computation, or // nullptr if the computation does not have a sequential ordering. - virtual const std::vector* SequentialOrder( + virtual const HloInstructionSequence* SequentialOrder( const HloComputation& computation) const = 0; // Return the call graph of the module used to compute ordering. @@ -96,7 +97,7 @@ class PredecessorHloOrdering : public HloOrdering { // Returns nullptr indicating the computation does not have a sequential // ordering. - const std::vector* SequentialOrder( + const HloInstructionSequence* SequentialOrder( const HloComputation& computation) const override { return nullptr; } @@ -185,7 +186,7 @@ class SequentialHloOrdering : public HloOrdering { ~SequentialHloOrdering() override = default; // Returns the sequential instruction order for the given computation. - const std::vector* SequentialOrder( + const HloInstructionSequence* SequentialOrder( const HloComputation& computation) const override; string ToString() const override; diff --git a/tensorflow/compiler/xla/service/hlo_ordering_test.cc b/tensorflow/compiler/xla/service/hlo_ordering_test.cc index b045adc9640ac0ca8cf4a127fea2fbfcbb1aaf3f..3ca77e60cd5275c22eb0e338cd5437fc44b49958 100644 --- a/tensorflow/compiler/xla/service/hlo_ordering_test.cc +++ b/tensorflow/compiler/xla/service/hlo_ordering_test.cc @@ -53,7 +53,7 @@ TEST_F(HloOrderingTest, InstructionsInDifferentComputations) { // %c = Constant(42.0f) // // This results in a diamond-shaped callgraph. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); const Shape scalar_shape = ShapeUtil::MakeShape(xla::F32, {}); auto builder_c = HloComputation::Builder("C"); @@ -126,7 +126,7 @@ TEST_F(HloOrderingTest, InstructionsInWhileComputations) { // %constant = Constant(1.0) // return While(%constant, body, condition) // - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); const Shape scalar_shape = ShapeUtil::MakeShape(xla::F32, {}); auto body_builder = HloComputation::Builder("body"); @@ -176,7 +176,7 @@ TEST_F(HloOrderingTest, InstructionsInWhileComputations) { TEST_F(HloOrderingTest, ParametersDefinedBeforeOthers) { // Entry parameter should always be defined before other instruction. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); const Shape scalar_shape = ShapeUtil::MakeShape(xla::F32, {}); auto builder = HloComputation::Builder(TestName()); auto constant = builder.AddInstruction( @@ -209,7 +209,7 @@ TEST_F(HloOrderingTest, ValuesInWhileComputations) { // %while = While(%constant, body, condition) // %add = Add(%constant, %while) // - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); const Shape scalar_shape = ShapeUtil::MakeShape(xla::F32, {}); auto body_builder = HloComputation::Builder("body"); @@ -407,7 +407,7 @@ TEST_F(HloOrderingTest, // %dead = Constant(123.0) // // %root should interfere with %dead. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); const Shape scalar_shape = ShapeUtil::MakeShape(xla::F32, {}); auto builder = HloComputation::Builder(TestName()); @@ -455,7 +455,7 @@ TEST_F(HloOrderingTest, // ROOT %call = call({%c}), subcomputation // // %root should interfere with %dead. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); const Shape scalar_shape = ShapeUtil::MakeShape(xla::F32, {}); auto subbuilder = HloComputation::Builder(TestName() + ".sub"); diff --git a/tensorflow/compiler/xla/service/hlo_parser.cc b/tensorflow/compiler/xla/service/hlo_parser.cc index 81f091238e5725f64b953f70b82d52cc90aef8ea..638396308c2a9c1f20e47f78b594d54f07c0c4e5 100644 --- a/tensorflow/compiler/xla/service/hlo_parser.cc +++ b/tensorflow/compiler/xla/service/hlo_parser.cc @@ -23,6 +23,7 @@ limitations under the License. #include "absl/strings/str_split.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_opcode.h" @@ -47,11 +48,11 @@ 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(const HloModule* module) { +HloSchedule ScheduleFromInstructionOrder(HloModule* module) { HloSchedule schedule(module); - for (const HloComputation* computation : module->computations()) { + for (HloComputation* computation : module->computations()) { if (!computation->IsFusionComputation()) { - for (const HloInstruction* instruction : computation->instructions()) { + for (HloInstruction* instruction : computation->instructions()) { schedule.GetOrCreateSequence(computation).push_back(instruction); } } @@ -74,6 +75,7 @@ class HloParser { string GetError() const { return StrJoin(error_, "\n"); } // Stand alone parsing utils for various aggregate data types. + StatusOr ParseShapeOnly(); StatusOr ParseShardingOnly(); StatusOr ParseWindowOnly(); StatusOr ParseConvolutionDimensionNumbersOnly(); @@ -108,7 +110,7 @@ class HloParser { bool ParseInstructionList(HloComputation** computation, const string& computation_name); bool ParseInstruction(HloComputation::Builder* builder, string* root_name); - bool ParseInstruciontRhs(HloComputation::Builder* builder, const string& name, + bool ParseInstructionRhs(HloComputation::Builder* builder, const string& name, LocTy name_loc); bool ParseControlPredecessors(HloInstruction* instruction); bool ParseLiteral(Literal* literal, const Shape& shape); @@ -255,7 +257,10 @@ class HloParser { bool ParseName(string* result); bool ParseAttributeName(string* result); bool ParseString(string* result); + bool ParseDimensionSizes(std::vector* dimension_sizes, + std::vector* dynamic_dimensions); bool ParseShape(Shape* result); + bool ParseLayout(Layout* layout); bool ParseOpcode(HloOpcode* result); bool ParseFftType(FftType* result); bool ParseFusionKind(HloInstruction::FusionKind* result); @@ -279,9 +284,6 @@ class HloParser { // If the current token is 'kind', eats it (i.e. lexes the next token) and // returns true. bool EatIfPresent(TokKind kind); - // Parses a shape, and returns true if the result is compatible with the given - // shape. - bool EatShapeAndCheckCompatible(const Shape& shape); // Adds the instruction to the pool. Returns false and emits an error if the // instruction already exists. @@ -418,6 +420,18 @@ std::pair* HloParser::FindInstruction( } return create_missing_instruction_(name, *shape); } + + if (instr != nullptr && shape.has_value() && + !ShapeUtil::Compatible(instr->first->shape(), shape.value())) { + Error( + lexer_.GetLoc(), + StrCat("The declared operand shape ", + ShapeUtil::HumanStringWithLayout(shape.value()), + " is not compatible with the shape of the operand instruction ", + ShapeUtil::HumanStringWithLayout(instr->first->shape()), ".")); + return nullptr; + } + return instr; } @@ -596,10 +610,10 @@ bool HloParser::ParseInstruction(HloComputation::Builder* builder, *root_name = name; } - return ParseInstruciontRhs(builder, name, name_loc); + return ParseInstructionRhs(builder, name, name_loc); } -bool HloParser::ParseInstruciontRhs(HloComputation::Builder* builder, +bool HloParser::ParseInstructionRhs(HloComputation::Builder* builder, const string& name, LocTy name_loc) { Shape shape; HloOpcode opcode; @@ -754,7 +768,7 @@ bool HloParser::ParseInstruciontRhs(HloComputation::Builder* builder, HloInstruction::CreateBitcastConvert(shape, operands[0])); break; } - case HloOpcode::kCrossReplicaSum: { + case HloOpcode::kAllReduce: { optional>> tmp_groups; optional to_apply; optional> replica_group_ids; @@ -774,10 +788,9 @@ bool HloParser::ParseInstruciontRhs(HloComputation::Builder* builder, if (tmp_groups) { replica_groups = CreateReplicaGroups(*tmp_groups); } - instruction = - builder->AddInstruction(HloInstruction::CreateCrossReplicaSum( - shape, operands, *to_apply, replica_groups, - barrier ? *barrier : "", all_reduce_id)); + instruction = builder->AddInstruction(HloInstruction::CreateAllReduce( + shape, operands, *to_apply, replica_groups, barrier ? *barrier : "", + all_reduce_id)); break; } case HloOpcode::kAllToAll: { @@ -838,6 +851,15 @@ bool HloParser::ParseInstruciontRhs(HloComputation::Builder* builder, } break; } + case HloOpcode::kAddDependency: { + if (!ParseOperands(&operands, /*expected_size=*/2) || + !ParseAttributes(attrs)) { + return false; + } + instruction = builder->AddInstruction( + HloInstruction::CreateAddDependency(operands[0], operands[1])); + break; + } case HloOpcode::kSort: { optional> dimensions; attrs["dimensions"] = {/*required=*/true, AttrTy::kBracedInt64List, @@ -985,11 +1007,14 @@ bool HloParser::ParseInstruciontRhs(HloComputation::Builder* builder, optional window; optional dnums; optional feature_group_count; + optional batch_group_count; attrs["window"] = {/*required=*/false, AttrTy::kWindow, &window}; attrs["dim_labels"] = {/*required=*/true, 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}; optional> operand_precision; attrs["operand_precision"] = {/*required=*/false, AttrTy::kPrecisionList, &operand_precision}; @@ -1003,6 +1028,9 @@ bool HloParser::ParseInstruciontRhs(HloComputation::Builder* builder, if (!feature_group_count) { feature_group_count = 1; } + if (!batch_group_count) { + batch_group_count = 1; + } PrecisionConfig precision_config; if (operand_precision) { *precision_config.mutable_operand_precision() = { @@ -1013,7 +1041,8 @@ bool HloParser::ParseInstruciontRhs(HloComputation::Builder* builder, } instruction = builder->AddInstruction(HloInstruction::CreateConvolve( shape, /*lhs=*/operands[0], /*rhs=*/operands[1], - feature_group_count.value(), *window, *dnums, precision_config)); + feature_group_count.value(), batch_group_count.value(), *window, + *dnums, precision_config)); break; } case HloOpcode::kFft: { @@ -1089,8 +1118,8 @@ bool HloParser::ParseInstruciontRhs(HloComputation::Builder* builder, absl::Span(operands).subspan( 0, operands.size() / 2), /*init_values=*/ - absl::Span(operands).subspan( - operands.size() / 2, operands.size()), + absl::Span(operands).subspan(operands.size() / + 2), *dimensions_to_reduce, *reduce_computation)); break; } @@ -1142,24 +1171,39 @@ bool HloParser::ParseInstruciontRhs(HloComputation::Builder* builder, optional> dynamic_slice_sizes; attrs["dynamic_slice_sizes"] = { /*required=*/true, AttrTy::kBracedInt64List, &dynamic_slice_sizes}; - if (!ParseOperands(&operands, /*expected_size=*/2) || - !ParseAttributes(attrs)) { + LocTy loc = lexer_.GetLoc(); + if (!ParseOperands(&operands) || !ParseAttributes(attrs)) { return false; } + if (operands.empty()) { + return Error(loc, "Expected at least one operand."); + } + if (!(operands.size() == 2 && operands[1]->shape().rank() == 1) && + operands.size() != 1 + operands[0]->shape().rank()) { + return Error(loc, "Wrong number of operands."); + } instruction = builder->AddInstruction(HloInstruction::CreateDynamicSlice( - shape, /*operand=*/operands[0], /*start_indices=*/operands[1], + shape, /*operand=*/operands[0], + /*start_indices=*/absl::MakeSpan(operands).subspan(1), *dynamic_slice_sizes)); break; } case HloOpcode::kDynamicUpdateSlice: { - if (!ParseOperands(&operands, /*expected_size=*/3) || - !ParseAttributes(attrs)) { + LocTy loc = lexer_.GetLoc(); + if (!ParseOperands(&operands) || !ParseAttributes(attrs)) { return false; } + if (operands.size() < 2) { + return Error(loc, "Expected at least two operands."); + } + if (!(operands.size() == 3 && operands[2]->shape().rank() == 1) && + operands.size() != 2 + operands[0]->shape().rank()) { + return Error(loc, "Wrong number of operands."); + } instruction = builder->AddInstruction(HloInstruction::CreateDynamicUpdateSlice( shape, /*operand=*/operands[0], /*update=*/operands[1], - /*start_indices=*/operands[2])); + /*start_indices=*/absl::MakeSpan(operands).subspan(2))); break; } case HloOpcode::kTranspose: { @@ -1259,7 +1303,7 @@ bool HloParser::ParseInstruciontRhs(HloComputation::Builder* builder, // the infeed instruction. ShapeUtil::GetTupleElementShape will check fail // if the shape is not a non-empty tuple, so add guard so an error message // can be emitted instead of a check fail - if (!ShapeUtil::IsTuple(shape) && !ShapeUtil::IsEmptyTuple(shape)) { + if (!shape.IsTuple() && !ShapeUtil::IsEmptyTuple(shape)) { return Error(lexer_.GetLoc(), "infeed must have a non-empty tuple shape"); } @@ -1535,6 +1579,18 @@ bool HloParser::ParseInstruciontRhs(HloComputation::Builder* builder, case HloOpcode::kTrace: return TokenError(StrCat("parsing not yet implemented for op: ", HloOpcodeString(opcode))); + case HloOpcode::kGetDimensionSize: + optional> dimensions; + attrs["dimensions"] = {/*required=*/true, AttrTy::kBracedInt64List, + &dimensions}; + if (!ParseOperands(&operands, /*expected_size=*/1) || + !ParseAttributes(attrs)) { + return false; + } + instruction = + builder->AddInstruction(HloInstruction::CreateGetDimensionSize( + shape, operands[0], (*dimensions)[0])); + break; } instruction->SetAndSanitizeName(name); @@ -1664,11 +1720,6 @@ bool HloParser::ParseSingleSharding(OpSharding* sharding, } break; } - case TokKind::kShape: - // TODO(b/112302613): Left here for backward compatibility to ignore the - // removed tile shape data. - lexer_.Lex(); - break; case TokKind::kRbrace: break; default: @@ -1794,6 +1845,10 @@ bool HloParser::SetValueInLiteral(tensorflow::int64 value, case U64: return SetValueInLiteralHelper(value, linear_index, literal); + case PRED: + // Bool type literals with rank >= 1 are printed in 0s and 1s. + return SetValueInLiteralHelper(static_cast(value), + linear_index, literal); default: LOG(FATAL) << "unknown integral primitive type " << PrimitiveType_Name(shape.element_type()); @@ -1888,25 +1943,12 @@ bool HloParser::SetValueInLiteralHelper(ParsedElemT value, return true; } -bool HloParser::EatShapeAndCheckCompatible(const Shape& shape) { - Shape new_shape; - if (!ParseShape(&new_shape)) { - return TokenError(StrCat("expects shape ", ShapeUtil::HumanString(shape))); - } - if (!ShapeUtil::Compatible(shape, new_shape)) { - return TokenError(StrCat( - "expects shape ", ShapeUtil::HumanString(shape), - ", but sees a different shape: ", ShapeUtil::HumanString(new_shape))); - } - return true; -} - // literal // ::= tuple // ::= non_tuple bool HloParser::ParseLiteral(Literal* literal, const Shape& shape) { - return ShapeUtil::IsTuple(shape) ? ParseTupleLiteral(literal, shape) - : ParseNonTupleLiteral(literal, shape); + return shape.IsTuple() ? ParseTupleLiteral(literal, shape) + : ParseNonTupleLiteral(literal, shape); } // tuple @@ -1915,10 +1957,6 @@ bool HloParser::ParseLiteral(Literal* literal, const Shape& shape) { // ::= /*empty*/ // ::= literal (',' literal)* bool HloParser::ParseTupleLiteral(Literal* literal, const Shape& shape) { - if (!EatShapeAndCheckCompatible(shape)) { - return TokenError(StrCat("expects tuple constant in shape ", - ShapeUtil::HumanString(shape))); - } if (!ParseToken(TokKind::kLparen, "expects '(' in front of tuple elements")) { return false; } @@ -1953,16 +1991,12 @@ bool HloParser::ParseNonTupleLiteral(Literal* literal, const Shape& shape) { return ParseSparseLiteral(literal, shape); } - CHECK(LayoutUtil::IsDenseArray(shape)); + CHECK(LayoutUtil::IsDenseArray(shape)) << shape.ToString(true); return ParseDenseLiteral(literal, shape); } bool HloParser::ParseDenseLiteral(Literal* literal, const Shape& shape) { - const tensorflow::int64 rank = ShapeUtil::Rank(shape); - if (rank > 1 && !EatShapeAndCheckCompatible(shape)) { - return false; - } - + const tensorflow::int64 rank = shape.rank(); // Create a literal with the given shape in default layout. *literal = LiteralUtil::CreateFromDimensions( shape.element_type(), AsInt64Slice(shape.dimensions())); @@ -2048,14 +2082,13 @@ bool HloParser::ParseDenseLiteral(Literal* literal, const Shape& shape) { } if (lexer_.GetKind() == TokKind::kw_true || lexer_.GetKind() == TokKind::kw_false) { - // TODO(congliu): bool type literals with rank >= 1 are actually - // printed in a compact form instead of "true" or "false". Fix that. if (!SetValueInLiteral(lexer_.GetKind() == TokKind::kw_true, linear_index++, literal)) { return false; } lexer_.Lex(); - } else if (primitive_util::IsIntegralType(shape.element_type())) { + } else if (primitive_util::IsIntegralType(shape.element_type()) || + shape.element_type() == PRED) { LocTy loc = lexer_.GetLoc(); tensorflow::int64 value; if (!ParseInt64(&value)) { @@ -2090,10 +2123,6 @@ bool HloParser::ParseDenseLiteral(Literal* literal, const Shape& shape) { } bool HloParser::ParseSparseLiteral(Literal* literal, const Shape& shape) { - if (!EatShapeAndCheckCompatible(shape)) { - return false; - } - switch (shape.element_type()) { case PRED: return ParseSparseLiteralHelper(literal, shape); @@ -2132,7 +2161,7 @@ template bool HloParser::ParseSparseLiteralHelper(Literal* literal, const Shape& shape) { std::vector index; - tensorflow::int64 rank = ShapeUtil::Rank(shape); + tensorflow::int64 rank = shape.rank(); *literal = Literal(shape); @@ -2693,7 +2722,7 @@ bool HloParser::ParseConvolutionDimensionNumbers( // The str is expected to have 3 items, lhs, rhs, out, and it must look like // lhs_rhs->out, that is, the first separator is "_" and the second is "->". - std::vector split1 = absl::StrSplit(str, "_"); + std::vector split1 = absl::StrSplit(str, '_'); if (split1.size() != 2) { LOG(FATAL) << "expects 3 items: lhs, rhs, and output dims, but sees " << str; @@ -2717,7 +2746,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(); }; @@ -2958,6 +2987,50 @@ bool HloParser::ParseParamList() { return ParseToken(TokKind::kRparen, "expects ')' at the end of param list"); } +// dimension_sizes ::= '[' dimension_list ']' +// dimension_list +// ::= /*empty*/ +// ::= <=? int64 (',' param)* +// param ::= name shape +bool HloParser::ParseDimensionSizes(std::vector* dimension_sizes, + std::vector* dynamic_dimensions) { + auto parse_and_add_item = [&]() { + tensorflow::int64 i; + bool is_dynamic = false; + if (lexer_.GetKind() == TokKind::kLeq) { + is_dynamic = true; + lexer_.Lex(); + } + if (!ParseInt64(&i)) { + return false; + } + dimension_sizes->push_back(i); + dynamic_dimensions->push_back(is_dynamic); + return true; + }; + return ParseList(TokKind::kLsquare, TokKind::kRsquare, TokKind::kComma, + parse_and_add_item); +} + +// layout ::= '{' int64_list '}' +bool HloParser::ParseLayout(Layout* layout) { + std::vector minor_to_major; + auto parse_and_add_item = [&]() { + tensorflow::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)) { + return false; + } + *layout = LayoutUtil::MakeLayout(minor_to_major); + return true; +} + // shape ::= shape_val_ // shape ::= '(' tuple_elements ')' // tuple_elements @@ -2981,19 +3054,67 @@ bool HloParser::ParseShape(Shape* result) { return ParseToken(TokKind::kRparen, "expects ')' at the end of tuple."); } - if (lexer_.GetKind() != TokKind::kShape) { - return TokenError(absl::StrCat("expected shape, saw ", + if (lexer_.GetKind() != TokKind::kPrimitiveType) { + return TokenError(absl::StrCat("expected primitive type, saw ", TokKindToString(lexer_.GetKind()))); } - *result = lexer_.GetShapeVal(); + PrimitiveType primitive_type = lexer_.GetPrimitiveTypeVal(); lexer_.Lex(); + + // Each element contains a dimension size and a bool indicating whether this + // is a dynamic dimension. + std::vector dimension_sizes; + std::vector dynamic_dimensions; + if (!ParseDimensionSizes(&dimension_sizes, &dynamic_dimensions)) { + return false; + } + result->set_element_type(primitive_type); + for (int i = 0; i < dimension_sizes.size(); ++i) { + result->add_dimensions(dimension_sizes[i]); + result->set_dynamic_dimension(i, dynamic_dimensions[i]); + } + LayoutUtil::SetToDefaultLayout(result); + + if (lexer_.GetKind() == TokKind::kw_sparse) { + lexer_.Lex(); + const string message = + "expects a brace-bracketed integer for sparse layout"; + tensorflow::int64 max_sparse_elements; + if (!ParseToken(TokKind::kLbrace, message) || + !ParseInt64(&max_sparse_elements) || + !ParseToken(TokKind::kRbrace, message)) { + return false; + } + *result->mutable_layout() = + LayoutUtil::MakeSparseLayout(max_sparse_elements); + return true; + } + + // We need to lookahead to see if a following open brace is the start of a + // layout. The specific problematic case is: + // + // ENTRY %foo (x: f32[42]) -> f32[123] { + // ... + // } + // + // 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 + if (lexer_.GetKind() == TokKind::kLbrace && + lexer_.LookAhead() == TokKind::kInt) { + Layout layout; + if (!ParseLayout(&layout)) { + return false; + } + *result->mutable_layout() = layout; + } return true; } bool HloParser::CanBeShape() { - // A non-tuple shape starts with a kShape token; a tuple shape starts with - // '('. - return lexer_.GetKind() == TokKind::kShape || + // A non-tuple shape starts with a kPrimitiveType token; a tuple shape starts + // with '('. + return lexer_.GetKind() == TokKind::kPrimitiveType || lexer_.GetKind() == TokKind::kLparen; } @@ -3296,6 +3417,18 @@ bool HloParser::AddComputation(const string& name, HloComputation* computation, return true; } +StatusOr HloParser::ParseShapeOnly() { + lexer_.Lex(); + Shape shape; + if (!ParseShape(&shape)) { + return InvalidArgument("Syntax error:\n%s", GetError()); + } + if (lexer_.GetKind() != TokKind::kEof) { + return InvalidArgument("Syntax error:\nExtra content after shape"); + } + return shape; +} + StatusOr HloParser::ParseShardingOnly() { lexer_.Lex(); OpSharding op_sharding; @@ -3374,7 +3507,7 @@ bool HloParser::ParseSingleInstruction(HloModule* module) { // e.g. // // f32[10] fusion(...), calls={...} - if (!ParseInstruciontRhs(&builder, module->name(), lexer_.GetLoc())) { + if (!ParseInstructionRhs(&builder, module->name(), lexer_.GetLoc())) { return false; } } else { @@ -3439,4 +3572,9 @@ StatusOr ParsePaddingConfig(absl::string_view str) { return parser.ParsePaddingConfigOnly(); } +StatusOr ParseShape(absl::string_view str) { + HloParser parser(str); + return parser.ParseShapeOnly(); +} + } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_parser.h b/tensorflow/compiler/xla/service/hlo_parser.h index 81eeb9f13bf7f06123c0b35e9f3352c197866a7a..450a54c54c156c2ae27475d145a8e83dc841b431 100644 --- a/tensorflow/compiler/xla/service/hlo_parser.h +++ b/tensorflow/compiler/xla/service/hlo_parser.h @@ -44,7 +44,9 @@ Status ParseHloString(absl::string_view str, HloModule* module); // creates a HloModule with default config. StatusOr> ParseHloString(absl::string_view str); -// Parses the result of HloSharding::ToString(), e.g. "{replicated}". +// ParseHloString sharding from str. str is supposed to contain the body of the +// sharding, i.e. just the rhs of the "sharding={...}" attribute string, +// e.g., "{replicated}". StatusOr ParseSharding(absl::string_view str); // Parses the result of window_util::ToString(const Window&). @@ -55,13 +57,12 @@ StatusOr ParseWindow(absl::string_view str); StatusOr ParseConvolutionDimensionNumbers( absl::string_view str); -// ParseHloString sharding from str. str is supposed to contain the body of the -// sharding, i.e. just the rhs of the "sharding={...}" attribute string. -StatusOr ParseSharding(absl::string_view str); - // Parses the result of PaddingConfigToString(), e.g. "0_0x1_1". StatusOr ParsePaddingConfig(absl::string_view str); +// Parses and returns a Shape::ToString-format string. +StatusOr ParseShape(absl::string_view str); + } // namespace xla #endif // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_PARSER_H_ diff --git a/tensorflow/compiler/xla/service/hlo_parser_test.cc b/tensorflow/compiler/xla/service/hlo_parser_test.cc index 19f84d8bd28371518e44e38614b8a81fa920985f..6ba16cc82ac1da2a30610d9dfb56cacc100ae05f 100644 --- a/tensorflow/compiler/xla/service/hlo_parser_test.cc +++ b/tensorflow/compiler/xla/service/hlo_parser_test.cc @@ -21,7 +21,8 @@ limitations under the License. #include "absl/strings/string_view.h" #include "tensorflow/compiler/xla/service/hlo_casting_utils.h" #include "tensorflow/compiler/xla/service/hlo_instructions.h" -#include "tensorflow/compiler/xla/service/hlo_matchers.h" +#include "tensorflow/compiler/xla/service/pattern_matcher.h" +#include "tensorflow/compiler/xla/service/pattern_matcher_gmock.h" #include "tensorflow/compiler/xla/window_util.h" #include "tensorflow/core/lib/core/status_test_util.h" #include "tensorflow/core/platform/test.h" @@ -29,7 +30,7 @@ limitations under the License. namespace xla { namespace { -namespace op = ::xla::testing::opcode_matchers; +namespace m = ::xla::match; using absl::string_view; struct TestData { @@ -75,6 +76,18 @@ ENTRY %constant_pred () -> pred[] { )" }, +// pred array constant +{ +"ConstantPredArray", +R"(HloModule module + +ENTRY %constant_pred_array () -> pred[2,3] { + ROOT %constant = pred[2,3]{1,0} constant({ { 0, 1, 0 }, { 1, 0, 1 } }) +} + +)" +}, + // s32 constant { "ConstantS32", @@ -115,7 +128,7 @@ ENTRY %ConstantF32Empty.v4 () -> f32[0] { R"(HloModule ConstantF32R4Empty_module ENTRY %ConstantF32R4Empty.v4 () -> f32[2,0,4,3] { - ROOT %constant = f32[2,0,4,3]{3,2,1,0} constant(f32[2,0,4,3] { { /*i0=0*/ }, { /*i0=1*/ } }) + ROOT %constant = f32[2,0,4,3]{3,2,1,0} constant({ { /*i0=0*/ }, { /*i0=1*/ } }) } )" @@ -126,7 +139,7 @@ ENTRY %ConstantF32R4Empty.v4 () -> f32[2,0,4,3] { R"(HloModule Small_3x2x1x1_module ENTRY %Small_3x2x1x1.v1 () -> f32[3,2,1,1] { - ROOT %constant = f32[3,2,1,1]{3,2,1,0} constant(f32[3,2,1,1] { { /*i0=0*/ { /*i1=0*/ {-1} }, { /*i1=1*/ {4.1} } }, { /*i0=1*/ { /*i1=0*/ {2} }, { /*i1=1*/ {4.1} } }, { /*i0=2*/ { /*i1=0*/ {5} }, { /*i1=1*/ {4.4} } } }) + ROOT %constant = f32[3,2,1,1]{3,2,1,0} constant({ { /*i0=0*/ { /*i1=0*/ {-1} }, { /*i1=1*/ {4.1} } }, { /*i0=1*/ { /*i1=0*/ {2} }, { /*i1=1*/ {4.1} } }, { /*i0=2*/ { /*i1=0*/ {5} }, { /*i1=1*/ {4.4} } } }) } )" @@ -183,7 +196,7 @@ ENTRY %add_constants () -> f32[] { R"(HloModule TupleConstant_module ENTRY %TupleConstant.v1 () -> (f32[2,1], f32[2]) { - ROOT %constant = (f32[2,1]{1,0}, f32[2]{0}) constant((f32[2,1], f32[2]) ( f32[2,1] { { 1 }, { 2 } }, {2, 42} )) + ROOT %constant = (f32[2,1]{1,0}, f32[2]{0}) constant(( { {1}, {2} }, {2, 42} )) } )" @@ -282,11 +295,11 @@ ENTRY %WhileWithScalarS32Result.v2 () -> s32[] { R"(HloModule TwoSendRecvBothWayRecvFist_module ENTRY %TwoSendRecvBothWayRecvFist.v3 () -> (f32[], token[]) { - %token = token[] after-all() - %recv = (f32[], u32[], token[]) recv(token[] %token), channel_id=15, sharding={maximal device=1} + %token0 = token[] after-all() + %recv = (f32[], u32[], token[]) recv(token[] %token0), channel_id=15, sharding={maximal device=1} ROOT %recv-done = (f32[], token[]) recv-done((f32[], u32[], token[]) %recv), channel_id=15, sharding={maximal device=1} %constant = f32[] constant(2.1), sharding={maximal device=0} - %send = (f32[], u32[], token[]) send(f32[] %constant, token[] %token), channel_id=16, sharding={maximal device=0}, control-predecessors={%recv} + %send = (f32[], u32[], token[]) send(f32[] %constant, token[] %token0), channel_id=16, sharding={maximal device=0}, control-predecessors={%recv} %send-done = token[] send-done((f32[], u32[], token[]) %send), channel_id=16, sharding={maximal device=0} } @@ -297,11 +310,11 @@ ENTRY %TwoSendRecvBothWayRecvFist.v3 () -> (f32[], token[]) { R"(HloModule HostTransferSendRecv_module ENTRY %TwoSendRecvBothWayRecvFist.v3 () -> (f32[], token[]) { - %token = token[] after-all() - %recv = (f32[], u32[], token[]) recv(token[] %token), channel_id=15, is_host_transfer=true + %token0 = token[] after-all() + %recv = (f32[], u32[], token[]) recv(token[] %token0), channel_id=15, is_host_transfer=true ROOT %recv-done = (f32[], token[]) recv-done((f32[], u32[], token[]) %recv), channel_id=15, is_host_transfer=true %constant = f32[] constant(2.1), sharding={maximal device=0} - %send = (f32[], u32[], token[]) send(f32[] %constant, token[] %token), channel_id=16, is_host_transfer=true + %send = (f32[], u32[], token[]) send(f32[] %constant, token[] %token0), channel_id=16, is_host_transfer=true %send-done = token[] send-done((f32[], u32[], token[]) %send), channel_id=16, is_host_transfer=true } @@ -314,7 +327,7 @@ R"(HloModule GetTupleElement_module ENTRY %GetTupleElement.v4 () -> s32[2,3] { %constant = f32[3]{0} constant({1, 2, 3}) - %constant.1 = s32[2,3]{1,0} constant(s32[2,3] { { 1, 2, 3 }, { 4, 5, 6 } }) + %constant.1 = s32[2,3]{1,0} constant({ { 1, 2, 3 }, { 4, 5, 6 } }) %tuple = (f32[3]{0}, s32[2,3]{1,0}) tuple(f32[3]{0} %constant, s32[2,3]{1,0} %constant.1) ROOT %get-tuple-element = s32[2,3]{1,0} get-tuple-element((f32[3]{0}, s32[2,3]{1,0}) %tuple), index=1, sharding={maximal device=0} } @@ -421,7 +434,7 @@ ENTRY %ConvolveBackward (input: f32[128,7,7,512], filter: f32[3,3,512,512]) -> f R"(HloModule Reverse4DFloatArrayOnDim01_module ENTRY %Reverse4DFloatArrayOnDim01.v2 () -> f32[4,3,2,1] { - %constant = f32[4,3,2,1]{0,1,2,3} constant(f32[4,3,2,1] { { /*i0=0*/ { /*i1=0*/ {1}, {2} }, { /*i1=1*/ {3}, {4} }, { /*i1=2*/ {5}, {6} } }, { /*i0=1*/ { /*i1=0*/ {7}, {8} }, { /*i1=1*/ {9}, {10} }, { /*i1=2*/ {11}, {12} } }, { /*i0=2*/ { /*i1=0*/ {13}, {14} }, { /*i1=1*/ {15}, {16} }, { /*i1=2*/ {17}, {18} } }, { /*i0=3*/ { /*i1=0*/ {19}, {20} }, { /*i1=1*/ {21}, {22} }, { /*i1=2*/ {23}, {24} } } }) + %constant = f32[4,3,2,1]{0,1,2,3} constant({ { /*i0=0*/ { /*i1=0*/ {1}, {2} }, { /*i1=1*/ {3}, {4} }, { /*i1=2*/ {5}, {6} } }, { /*i0=1*/ { /*i1=0*/ {7}, {8} }, { /*i1=1*/ {9}, {10} }, { /*i1=2*/ {11}, {12} } }, { /*i0=2*/ { /*i1=0*/ {13}, {14} }, { /*i1=1*/ {15}, {16} }, { /*i1=2*/ {17}, {18} } }, { /*i0=3*/ { /*i1=0*/ {19}, {20} }, { /*i1=1*/ {21}, {22} }, { /*i1=2*/ {23}, {24} } } }) ROOT %reverse = f32[4,3,2,1]{0,1,2,3} reverse(f32[4,3,2,1]{0,1,2,3} %constant), dimensions={0,1} } @@ -433,8 +446,8 @@ ENTRY %Reverse4DFloatArrayOnDim01.v2 () -> f32[4,3,2,1] { R"(HloModule Concat2x3With2x5_module ENTRY %Concat2x3With2x5.v3 () -> f32[2,8] { - %constant = f32[2,3]{1,0} constant(f32[2,3] { { 0, 1, 2 }, { 1000, 1001, 1002 } }) - %constant.1 = f32[2,5]{1,0} constant(f32[2,5] { { 64, 65, 66, 67, 68 }, { 1064, 1065, 1066, 1067, 1068 } }) + %constant = f32[2,3]{1,0} constant({ { 0, 1, 2 }, { 1000, 1001, 1002 } }) + %constant.1 = f32[2,5]{1,0} constant({ { 64, 65, 66, 67, 68 }, { 1064, 1065, 1066, 1067, 1068 } }) ROOT %concatenate = f32[2,8]{1,0} concatenate(f32[2,3]{1,0} %constant, f32[2,5]{1,0} %constant.1), dimensions={1} } @@ -458,8 +471,8 @@ R"(HloModule R4F32OverlapSmall_module } ENTRY %R4F32OverlapSmall.v4 () -> f32[4,5,1,1] { - %constant = f32[4,5,1,1]{3,2,1,0} constant(f32[4,5,1,1] { { /*i0=0*/ { /*i1=0*/ {7} }, { /*i1=1*/ {2} }, { /*i1=2*/ {5} }, { /*i1=3*/ {3} }, { /*i1=4*/ {8} } }, { /*i0=1*/ { /*i1=0*/ {3} }, { /*i1=1*/ {8} }, { /*i1=2*/ {9} }, { /*i1=3*/ {3} }, { /*i1=4*/ {4} } }, { /*i0=2*/ { /*i1=0*/ {1} }, { /*i1=1*/ {5} }, { /*i1=2*/ {7} }, { /*i1=3*/ {5} }, { /*i1=4*/ {6} } }, { /*i0=3*/ { /*i1=0*/ {0} }, { /*i1=1*/ {6} }, { /*i1=2*/ {2} }, { /*i1=3*/ {10} }, { /*i1=4*/ {2} } } }) - %constant.1 = f32[2,2,1,1]{3,2,1,0} constant(f32[2,2,1,1] { { /*i0=0*/ { /*i1=0*/ {2} }, { /*i1=1*/ {6} } }, { /*i0=1*/ { /*i1=0*/ {3} }, { /*i1=1*/ {1} } } }) + %constant = f32[4,5,1,1]{3,2,1,0} constant({ { /*i0=0*/ { /*i1=0*/ {7} }, { /*i1=1*/ {2} }, { /*i1=2*/ {5} }, { /*i1=3*/ {3} }, { /*i1=4*/ {8} } }, { /*i0=1*/ { /*i1=0*/ {3} }, { /*i1=1*/ {8} }, { /*i1=2*/ {9} }, { /*i1=3*/ {3} }, { /*i1=4*/ {4} } }, { /*i0=2*/ { /*i1=0*/ {1} }, { /*i1=1*/ {5} }, { /*i1=2*/ {7} }, { /*i1=3*/ {5} }, { /*i1=4*/ {6} } }, { /*i0=3*/ { /*i1=0*/ {0} }, { /*i1=1*/ {6} }, { /*i1=2*/ {2} }, { /*i1=3*/ {10} }, { /*i1=4*/ {2} } } }) + %constant.1 = f32[2,2,1,1]{3,2,1,0} constant({ { /*i0=0*/ { /*i1=0*/ {2} }, { /*i1=1*/ {6} } }, { /*i0=1*/ { /*i1=0*/ {3} }, { /*i1=1*/ {1} } } }) %constant.2 = f32[] constant(0) ROOT %select-and-scatter = f32[4,5,1,1]{3,2,1,0} select-and-scatter(f32[4,5,1,1]{3,2,1,0} %constant, f32[2,2,1,1]{3,2,1,0} %constant.1, f32[] %constant.2), window={size=2x3x1x1 stride=2x2x1x1}, select=%ge_F32.v3, scatter=%add_F32.v3 } @@ -510,7 +523,7 @@ ENTRY %slice.v2 (p0: f32[3,3,4,4]) -> f32[3,3,2,4] { R"(HloModule Slice3x3x3_To_1x3x3_F32_module ENTRY %Slice3x3x3_To_1x3x3_F32.v2 () -> f32[1,3,3] { - %constant = f32[3,3,3]{2,1,0} constant(f32[3,3,3] { { { 0, 1, 2 }, { 3, 4, 5 }, { 6, 7, 8 } }, { { 9, 10, 11 }, { 12, 13, 14 }, { 15, 16, 17 } }, { { 18, 19, 20 }, { 21, 22, 23 }, { 24, 25, 26 } } }) + %constant = f32[3,3,3]{2,1,0} constant({ { { 0, 1, 2 }, { 3, 4, 5 }, { 6, 7, 8 } }, { { 9, 10, 11 }, { 12, 13, 14 }, { 15, 16, 17 } }, { { 18, 19, 20 }, { 21, 22, 23 }, { 24, 25, 26 } } }) ROOT %slice = f32[1,3,3]{2,1,0} slice(f32[3,3,3]{2,1,0} %constant), slice={[0:1], [0:3], [0:3]} } @@ -534,10 +547,21 @@ ENTRY %SliceR0.v2 () -> s32[] { R"(HloModule Transpose_module ENTRY %Transpose.v2 () -> s32[1,2,3] { - %constant = s32[1,2,3]{2,1,0} constant(s32[1,2,3] { { { 1, 2, 3 }, { 4, 5, 6 } } }) + %constant = s32[1,2,3]{2,1,0} constant({ { { 1, 2, 3 }, { 4, 5, 6 } } }) 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 @@ -553,12 +577,26 @@ ENTRY %DynamicSlice.v5 (original_parameter: s32[2,2,258], start_index: s32[1]) - ROOT %dynamic-slice = s32[2,2,258]{2,1,0} dynamic-slice(s32[2,2,258]{2,1,0} %original_parameter, s32[3]{0} %concatenate), dynamic_slice_sizes={2,2,258} } +)" +}, +// Dynamic slice with scalar indices +{ +"DynamicSliceScalarIndices", +R"(HloModule DynamicSlice_module + +ENTRY %DynamicSlice.v5 (original_parameter: s32[2,2,258], start_index: s32[]) -> s32[2,2,258] { + %original_parameter = s32[2,2,258]{2,1,0} parameter(0) + %constant = s32[] constant(0) + %start_index = s32[] parameter(1) + ROOT %dynamic-slice = s32[2,2,258]{2,1,0} dynamic-slice(s32[2,2,258]{2,1,0} %original_parameter, s32[] %constant, s32[] %constant, s32[] %start_index), dynamic_slice_sizes={2,2,258} +} + )" }, // Dynamic update slice { "DynamicUpdateSlice", -R"(HloModule DynamicUpdateSlice_module +R"(HloModule DynamicSlice_module ENTRY %DynamicUpdateSlice.v4 (input: s32[1,1,25,1], update: s32[1,1,2,1], start_indices: s32[4]) -> s32[1,1,25,1] { %input = s32[1,1,25,1]{3,2,1,0} parameter(0) @@ -567,6 +605,23 @@ ENTRY %DynamicUpdateSlice.v4 (input: s32[1,1,25,1], update: s32[1,1,2,1], start_ ROOT %dynamic-update-slice = s32[1,1,25,1]{3,2,1,0} dynamic-update-slice(s32[1,1,25,1]{3,2,1,0} %input, s32[1,1,2,1]{3,2,1,0} %update, s32[4]{0} %start_indices) } +)" +}, +// Dynamic update slice with scalar indices +{ +"DynamicUpdateSliceScalarIndex", +R"(HloModule DynamicUpdateSlice_module + +ENTRY %DynamicUpdateSlice.v4 (input: s32[1,1,25,1], update: s32[1,1,2,1], start_index.0: s32[], start_index.1: s32[], start_index.2: s32[], start_index.3: s32[]) -> s32[1,1,25,1] { + %input = s32[1,1,25,1]{3,2,1,0} parameter(0) + %update = s32[1,1,2,1]{3,2,1,0} parameter(1) + %start_index.0 = s32[] parameter(2) + %start_index.1 = s32[] parameter(3) + %start_index.2 = s32[] parameter(4) + %start_index.3 = s32[] parameter(5) + ROOT %dynamic-update-slice = s32[1,1,25,1]{3,2,1,0} dynamic-update-slice(s32[1,1,25,1]{3,2,1,0} %input, s32[1,1,2,1]{3,2,1,0} %update, s32[] %start_index.0, s32[] %start_index.1, s32[] %start_index.2, s32[] %start_index.3) +} + )" }, // batch norm training @@ -575,7 +630,7 @@ ENTRY %DynamicUpdateSlice.v4 (input: s32[1,1,25,1], update: s32[1,1,2,1], start_ R"(HloModule BasicTraining_module ENTRY %BasicTraining.v4 () -> (f32[2,2,1,2], f32[2], f32[2]) { - %constant = f32[2,2,1,2]{3,2,1,0} constant(f32[2,2,1,2] { { /*i0=0*/ { /*i1=0*/ {1, 2} }, { /*i1=1*/ {3, 4} } }, { /*i0=1*/ { /*i1=0*/ {5, 6} }, { /*i1=1*/ {7, 8} } } }) + %constant = f32[2,2,1,2]{3,2,1,0} constant({ { /*i0=0*/ { /*i1=0*/ { 1, 2 } }, { /*i1=1*/ { 3, 4 } } }, { /*i0=1*/ { /*i1=0*/ { 5, 6 } }, { /*i1=1*/ { 7, 8 } } } }) %constant.1 = f32[2]{0} constant({2, 3}) %constant.2 = f32[2]{0} constant({1, 2}) ROOT %batch-norm-training = (f32[2,2,1,2]{3,2,1,0}, f32[2]{0}, f32[2]{0}) batch-norm-training(f32[2,2,1,2]{3,2,1,0} %constant, f32[2]{0} %constant.1, f32[2]{0} %constant.2), epsilon=0.001, feature_index=3 @@ -715,7 +770,7 @@ R"(HloModule fusion_module } ENTRY %fusion.v3 () -> f32[3,2,1,1] { - %constant = f32[3,2,1,1]{3,2,1,0} constant(f32[3,2,1,1] { { /*i0=0*/ { /*i1=0*/ {-1} }, { /*i1=1*/ {4.1} } }, { /*i0=1*/ { /*i1=0*/ {2} }, { /*i1=1*/ {4.1} } }, { /*i0=2*/ { /*i1=0*/ {5} }, { /*i1=1*/ {4.4} } } }) + %constant = f32[3,2,1,1]{3,2,1,0} constant({ { /*i0=0*/ { /*i1=0*/ {-1} }, { /*i1=1*/ {4.1} } }, { /*i0=1*/ { /*i1=0*/ {2} }, { /*i1=1*/ {4.1} } }, { /*i0=2*/ { /*i1=0*/ {5} }, { /*i1=1*/ {4.4} } } }) %constant.1 = f32[2]{0} constant({3.14, 4.25}) ROOT %fusion = f32[3,2,1,1]{3,2,1,0} fusion(f32[3,2,1,1]{3,2,1,0} %constant, f32[2]{0} %constant.1), kind=kLoop, calls=%fused_computation } @@ -727,7 +782,7 @@ 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(f32[2,3,4]{[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, 3]: 2, [2, 3, 4]: 3}) } )" @@ -737,7 +792,7 @@ ENTRY %sparse () -> f32[2,3,4] { R"(HloModule sparse_f32_empty ENTRY %sparse_f32_empty () -> f32[2,3,4] { - ROOT %foo = f32[2,3,4]sparse{10} constant(f32[2,3,4]{}) + ROOT %foo = f32[2,3,4]sparse{10} constant({}) } )" @@ -747,7 +802,7 @@ ENTRY %sparse_f32_empty () -> f32[2,3,4] { R"(HloModule sparse_f32_r1 ENTRY %sparse_f32_r1 () -> f32[9] { - ROOT %foo = f32[9]sparse{10} constant(f32[9]{1: 2, 3: 4, 5: 6}) + ROOT %foo = f32[9]sparse{10} constant({1: 2, 3: 4, 5: 6}) } )" @@ -918,11 +973,11 @@ ENTRY reduce_entry { R"(HloModule outfeed_module ENTRY InfeedToOutfeed { - token = token[] after-all() - infeed = ((u32[3]{0}, pred[]), token[]) infeed(token) + token0 = token[] after-all() + infeed = ((u32[3]{0}, pred[]), token[]) infeed(token0) infeed.data = (u32[3]{0}, pred[]) get-tuple-element(infeed), index=0 - outfeed = token[] outfeed(infeed.data, token) - ROOT infeed.1 = ((u32[3]{0}, pred[]), token[]) infeed(token) + outfeed = token[] outfeed(infeed.data, token0) + ROOT infeed.1 = ((u32[3]{0}, pred[]), token[]) infeed(token0) infeed.1.data = (u32[3]{0}, pred[]) get-tuple-element(infeed.1), index=0 infeed.1.token = token[] get-tuple-element(infeed.1), index=1 outfeed.1 = token[] outfeed(infeed.1.data, infeed.1.token) @@ -1104,9 +1159,9 @@ ENTRY Gather { )" }, -// cross-replica-sum +// all-reduce { -"CrossReplicaSum", +"AllReduce", R"(HloModule CRS add { @@ -1117,14 +1172,14 @@ add { ENTRY CRS { input = f32[8]{0} parameter(0) - ROOT crs = f32[8]{0} cross-replica-sum(input), replica_groups={}, to_apply=add + ROOT crs = f32[8]{0} all-reduce(input), replica_groups={}, to_apply=add } )" }, -// cross-replica-sum with subgroups +// all-reduce with subgroups { -"CrossReplicaSumWithSubgroups", +"AllReduceWithSubgroups", R"(HloModule CRS_Subgroups add { @@ -1133,9 +1188,28 @@ add { ROOT add = f32[] add(lhs, rhs) } -ENTRY CrossReplicaSumWithSubgroups { +ENTRY AllReduceWithSubgroups { input = f32[128,32]{0,1} parameter(0) - ROOT cross-replica-sum = f32[128,32]{0,1} cross-replica-sum(input), replica_groups={{0,1},{2,3}}, barrier="abc", to_apply=add + ROOT all-reduce = f32[128,32]{0,1} all-reduce(input), replica_groups={{0,1},{2,3}}, barrier="abc", to_apply=add +} + +)" +}, +// all-reduce with all-reduce-id +{ +"AllReduceAllReduce", +R"(HloModule CRS + +add { + lhs = f32[] parameter(0) + rhs = f32[] parameter(1) + ROOT add = f32[] add(lhs, rhs) +} + +ENTRY CRS { + input = f32[8]{0} parameter(0) + crs.1 = f32[8]{0} all-reduce(input), replica_groups={{0}}, all_reduce_id=1, to_apply=add + ROOT crs.0 = f32[8]{0} all-reduce(input), replica_groups={{0}}, all_reduce_id=1, to_apply=add } )" @@ -1210,7 +1284,38 @@ ENTRY Sort { } )" + }, +// AfterAll with multiple operands +{ +"AfterAllWithMultipleOperands", +R"(HloModule AfterAllWithMultipleOperands + +ENTRY AfterAllWithMultipleOperands { + p0 = f32[] parameter(0) + token0 = token[] after-all() + token1 = token[] after-all() + ROOT after-all = token[] after-all(p0, token0, token1) +} + +)" +}, +// AddDependency +// A dependency chain is created from 'neg' to 'exp' using tokens. +{ +"AddDependency", +R"(HloModule AddDependency + +ENTRY AddDependency { + p = f32[] parameter(0) + neg = f32[] negate(p) + token0 = token[] after-all(neg) + p_after_token = f32[] add-dependency(p, token0) + exp = f32[] exponential(p_after_token) + ROOT sum = f32[] add(neg, exp) } + +)" +}, }); // clang-format on } @@ -1266,20 +1371,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: @@ -1356,7 +1461,7 @@ TEST_F(HloParserTest, MoreConstants) { ENTRY %SelectScalarS32True.v4 () -> s32[] { %constant.2 = pred[] constant(true) - %constant.1 = s32[] constant(-42), sharding={s32[5,6] devices=[2,2]1,2,3,4} + %constant.1 = s32[] constant(-42), sharding={devices=[2,2]1,2,3,4} %constant = s32[] constant(42) %select = s32[] select(pred[] %constant.2, s32[] %constant.1, s32[] %constant) } @@ -1399,7 +1504,7 @@ TEST_F(HloParserTest, LiteralDimensionsMismatch_2) { const string original = R"(HloModule some_2x3_module ENTRY %some_2x3 () -> f32[2,3] { - ROOT %constant = f32[2,3]{1,0} constant(f32[2,3] {1, 2, 3, 4, 5, 6}) + ROOT %constant = f32[2,3]{1,0} constant({1, 2, 3, 4, 5, 6}) } )"; @@ -1413,7 +1518,7 @@ TEST_F(HloParserTest, LiteralDimensionsMismatch_3) { const string original = R"(HloModule some_2x3x2_module ENTRY %some_2x3x2 () -> f32[2,3,2] { - ROOT %constant = f32[2,3,2]{2,1,0} constant(f32[2,3,2] {{{1, 2}, {3, 4}, {5, 6}, {7, 8}, {9, 10}, {11, 12}}}) + ROOT %constant = f32[2,3,2]{2,1,0} constant({{{1, 2}, {3, 4}, {5, 6}, {7, 8}, {9, 10}, {11, 12}}}) } )"; @@ -1531,11 +1636,11 @@ TEST_F(HloParserTest, UnexpectedAttribute) { const string original = R"(HloModule unexpected_attr_module ENTRY %TwoSendRecvBothWayRecvFist.v3 () -> f32[] { - %token = token[] after-all() - %recv = (f32[], u32[], token[]) recv(token[] %token), channel_id=15 + %token0 = token[] after-all() + %recv = (f32[], u32[], token[]) recv(token[] %token0), channel_id=15 %recv-done = (f32[], token[]) recv-done((f32[], u32[], token[]) %recv), channel_id=15 ROOT %constant = f32[] constant(2.1) - %send = (f32[], u32[], token[]) send(f32[] %constant, token[] %token), channel_id=16, calls=%recv + %send = (f32[], u32[], token[]) send(f32[] %constant, token[] %token0), channel_id=16, calls=%recv %send-done = token[] send-done((f32[], u32[], token[]) %send), channel_id=16 } @@ -1548,11 +1653,11 @@ TEST_F(HloParserTest, MissingAttribute) { const string original = R"(HloModule missing_attr_module ENTRY %TwoSendRecvBothWayRecvFist.v3 () -> f32[] { - %token = token[] after-all() - %recv = (f32[], u32[], token[]) recv(token[] %token), channel_id=15 + %token0 = token[] after-all() + %recv = (f32[], u32[], token[]) recv(token[] %token0), channel_id=15 %recv-done = (f32[], token[]) recv-done((f32[], u32[], token[]) %recv), channel_id=15 ROOT %constant = f32[] constant(-2.1) - %send = (f32[], u32[], token[]) send(f32[] %constant, token[] %token) + %send = (f32[], u32[], token[]) send(f32[] %constant, token[] %token0) %send-done = token[] send-done((f32[], u32[], token[]) %send), channel_id=16 } @@ -1565,11 +1670,11 @@ TEST_F(HloParserTest, PredecessorUndefined) { const string original = R"(HloModule pre_not_found_module ENTRY %TwoSendRecvBothWayRecvFist.v3 () -> f32[] { - %token = token[] after-all() - %recv = (f32[], u32[], token[]) recv(token[] %token), channel_id=15 + %token0 = token[] after-all() + %recv = (f32[], u32[], token[]) recv(token[] %token0), channel_id=15 %recv-done = (f32[], token[]) recv-done((f32[], u32[], token[]) %recv), channel_id=15 ROOT %constant = f32[] constant(2.1) - %send = (f32[], u32[], token[]) send(f32[] %constant, token[] %token), channel_id=16, control-predecessors={%done} + %send = (f32[], u32[], token[]) send(f32[] %constant, token[] %token0), channel_id=16, control-predecessors={%done} %send-done = token[] send-done((f32[], u32[], token[]) %send), channel_id=16 } @@ -1831,7 +1936,8 @@ ENTRY ReduceR3ToR2 { )"; TF_ASSERT_OK_AND_ASSIGN(auto module, ParseHloString(original)); ASSERT_NE(module->entry_computation(), nullptr); - EXPECT_THAT(module->entry_computation()->root_instruction(), op::Reduce()); + EXPECT_THAT(module->entry_computation()->root_instruction(), + GmockMatch(m::Reduce())); } TEST_F(HloParserTest, ParseSharding) { @@ -1876,8 +1982,8 @@ TEST_F(HloParserTest, ParsePaddingConfigInteriorPaddingImplicitZeroDim) { TEST_F(HloParserTest, NontupleInfeed) { const string original = R"(HloModule nontuple_infeed: ENTRY nontuple_infeed { - token = token[] after-all() - ROOT infeed = pred[] infeed(token) + token0 = token[] after-all() + ROOT infeed = pred[] infeed(token0) })"; ExpectHasSubstr(ParseHloString(original).status().error_message(), "infeed must have a non-empty tuple shape"); @@ -1891,7 +1997,7 @@ TEST(HloParserSingleOpTest, SingleOp) { const HloComputation* computation = module->entry_computation(); ASSERT_NE(computation, nullptr); EXPECT_THAT(computation->root_instruction(), - op::Multiply(op::Parameter(0), op::Parameter(1))); + GmockMatch(m::Multiply(m::Parameter(0), m::Parameter(1)))); } TEST(HloParserSingleOpTest, SingleOpNoShapeProducesError) { @@ -1919,7 +2025,7 @@ TEST(HloParserSingleOpTest, SingleOpNoNames) { const HloComputation* computation = module->entry_computation(); ASSERT_NE(computation, nullptr); EXPECT_THAT(computation->root_instruction(), - op::Multiply(op::Parameter(0), op::Parameter(1))); + GmockMatch(m::Multiply(m::Parameter(0), m::Parameter(1)))); } TEST(HloParserSingleOpTest, CanonicalOp) { @@ -1928,7 +2034,7 @@ TEST(HloParserSingleOpTest, CanonicalOp) { const HloComputation* computation = module->entry_computation(); ASSERT_NE(computation, nullptr); EXPECT_THAT(computation->root_instruction(), - op::Multiply(op::Parameter(0), op::Parameter(1))); + GmockMatch(m::Multiply(m::Parameter(0), m::Parameter(1)))); EXPECT_EQ( computation->root_instruction()->ToString(HloPrintOptions::Canonical()), text); @@ -1982,7 +2088,11 @@ TEST(HloParserSingleOpTest, SingleOpWithNested) { const HloComputation* computation = module->entry_computation(); ASSERT_NE(computation, nullptr); EXPECT_THAT(computation->root_instruction(), - op::Fusion(op::Parameter(0), op::Parameter(1))); + GmockMatch(m::Op() + .WithOpcode(HloOpcode::kFusion) + .WithNumOperands(2) + .WithOperand(0, m::Parameter(0)) + .WithOperand(1, m::Parameter(1)))); } TEST(HloParserSingleOpTest, SingleOpWithNested_DoesNotExist) { @@ -2026,7 +2136,7 @@ TEST(HloParserSingleOpTest, ConvolutionTrivialFeatureGroupCount) { const HloComputation* computation = module->entry_computation(); ASSERT_NE(computation, nullptr); EXPECT_THAT(computation->root_instruction(), - op::Convolution(op::Parameter(0), op::Parameter(1))); + GmockMatch(m::Convolution(m::Parameter(0), m::Parameter(1)))); auto* convolution = Cast(computation->root_instruction()); EXPECT_EQ(convolution->feature_group_count(), 1); @@ -2090,8 +2200,10 @@ ENTRY %axpy.v5 (alpha: f32[], x: f32[2,4], y: f32[2,4]) -> f32[2,4] { module->schedule().is_computation_scheduled(module->entry_computation())); EXPECT_THAT( module->schedule().sequence(module->entry_computation()).instructions(), - ::testing::ElementsAre(op::Parameter(), op::Broadcast(), op::Parameter(), - op::Multiply(), op::Parameter(), op::Add())); + ::testing::ElementsAre( + GmockMatch(m::Parameter()), GmockMatch(m::Broadcast()), + GmockMatch(m::Parameter()), GmockMatch(m::Multiply()), + GmockMatch(m::Parameter()), GmockMatch(m::Add()))); } TEST_F(HloParserTest, IsScheduledIsTrueDifferentOrder) { @@ -2117,8 +2229,10 @@ ENTRY %axpy.v5 (alpha: f32[], x: f32[2,4], y: f32[2,4]) -> f32[2,4] { module->schedule().is_computation_scheduled(module->entry_computation())); EXPECT_THAT( module->schedule().sequence(module->entry_computation()).instructions(), - ::testing::ElementsAre(op::Parameter(), op::Parameter(), op::Parameter(), - op::Broadcast(), op::Multiply(), op::Add())); + ::testing::ElementsAre( + GmockMatch(m::Parameter()), GmockMatch(m::Parameter()), + GmockMatch(m::Parameter()), GmockMatch(m::Broadcast()), + GmockMatch(m::Multiply()), GmockMatch(m::Add()))); } TEST_F(HloParserTest, CustomCallWrongNumberofOperandConstraints) { @@ -2161,7 +2275,121 @@ ENTRY entry { ParseHloString(text)); } -// custom call incompatible shape. +TEST_F(HloParserTest, ShapeMismatchInOperand) { + const string text = R"( +HloModule foobar + +ENTRY %entrycomp (p: f32[2,2]) -> f32[2,2] { + %p = f32[2,2] parameter(0) + %constant.1 = f32[2,2] constant({{1, 2}, {3, 4}}) + ROOT %add.1 = f32[2,2] add(f32[2,2] %p, f32[2,5] %constant.1) +} +)"; + + ExpectHasSubstr(ParseHloString(text).status().error_message(), + "The declared operand shape f32[2,5]{1,0} is not compatible" + " with the shape of the operand instruction f32[2,2]{1,0}."); +} + +TEST_F(HloParserTest, ParseShapeStringR2F32) { + string shape_string = "f32[123,456]"; + TF_ASSERT_OK_AND_ASSIGN(Shape actual, ParseShape(shape_string)); + Shape expected = ShapeUtil::MakeShape(F32, {123, 456}); + ASSERT_TRUE(ShapeUtil::Equal(expected, actual)) + << "expected: " << ShapeUtil::HumanString(expected) + << "actual: " << ShapeUtil::HumanString(actual); +} + +TEST_F(HloParserTest, ParseShapeStringTupleOfArrays) { + string shape_string = "(f32[1572864],s8[5120,1024])"; + TF_ASSERT_OK_AND_ASSIGN(Shape actual, ParseShape(shape_string)); + Shape expected = + ShapeUtil::MakeTupleShape({ShapeUtil::MakeShape(F32, {1572864}), + ShapeUtil::MakeShape(S8, {5120, 1024})}); + ASSERT_TRUE(ShapeUtil::Equal(expected, actual)) + << "expected: " << ShapeUtil::HumanString(expected) + << "actual: " << ShapeUtil::HumanString(actual); +} + +TEST_F(HloParserTest, ParseShapeStringNestedTuple) { + string shape_string = "(f32[1],(f32[2], token[]), opaque[], f32[3])"; + TF_ASSERT_OK_AND_ASSIGN(Shape actual, ParseShape(shape_string)); + Shape expected = ShapeUtil::MakeTupleShape({ + ShapeUtil::MakeShape(F32, {1}), + ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeShape(F32, {2}), ShapeUtil::MakeTokenShape()}), + ShapeUtil::MakeOpaqueShape(), + ShapeUtil::MakeShape(F32, {3}), + }); + ASSERT_TRUE(ShapeUtil::Equal(expected, actual)) + << "expected: " << ShapeUtil::HumanString(expected) + << "actual: " << ShapeUtil::HumanString(actual); +} + +TEST_F(HloParserTest, ParseShapeStringWithLayout) { + string shape_string = "f32[123,456]{0,1}"; + TF_ASSERT_OK_AND_ASSIGN(Shape actual, ParseShape(shape_string)); + Shape expected = ShapeUtil::MakeShapeWithLayout(F32, {123, 456}, {0, 1}); + ASSERT_TRUE(ShapeUtil::Equal(expected, actual)) + << "expected: " << ShapeUtil::HumanString(expected) + << "actual: " << ShapeUtil::HumanString(actual); +} + +TEST_F(HloParserTest, ParseShapeStringWithSparseLayout) { + string shape_string = "f32[123,456]sparse{10}"; + TF_ASSERT_OK_AND_ASSIGN(Shape actual, ParseShape(shape_string)); + Shape expected = ShapeUtil::MakeShapeWithSparseLayout(F32, {123, 456}, 10); + ASSERT_TRUE(ShapeUtil::Equal(expected, actual)) + << "expected: " << ShapeUtil::HumanString(expected) + << "actual: " << ShapeUtil::HumanString(actual); +} + +TEST_F(HloParserTest, ParseOpaqueType) { + TF_ASSERT_OK_AND_ASSIGN(Shape actual, ParseShape("opaque[]")); + Shape expected = ShapeUtil::MakeOpaqueShape(); + ASSERT_TRUE(ShapeUtil::Equal(expected, actual)) + << "expected: " << ShapeUtil::HumanString(expected) + << "actual: " << ShapeUtil::HumanString(actual); +} + +TEST_F(HloParserTest, ParseTokenType) { + TF_ASSERT_OK_AND_ASSIGN(Shape actual, ParseShape("token[]")); + Shape expected = ShapeUtil::MakeTokenShape(); + ASSERT_TRUE(ShapeUtil::Equal(expected, actual)) + << "expected: " << ShapeUtil::HumanString(expected) + << "actual: " << ShapeUtil::HumanString(actual); +} + +TEST_F(HloParserTest, ParseInvalidShapeString) { + string shape_strings[] = { + "f32[123,456]foobar{0,1}", "f32[123,456]sparse{0,1}", "f32[123,456]{foo}", + "f32[123,456]dense{foo}", "f32[123,456]sparse{foo}", + }; + for (const string& shape_string : shape_strings) { + StatusOr result = ParseShape(shape_string); + ASSERT_FALSE(result.ok()) << "shape: " << shape_string; + } +} + +TEST_F(HloParserTest, ParseDynamicArray) { + string shape_string = "f32[123,<=456]"; + TF_ASSERT_OK_AND_ASSIGN(Shape actual, ParseShape(shape_string)); + Shape expected = ShapeUtil::MakeShape(F32, {123, 456}, {false, true}); + ASSERT_TRUE(ShapeUtil::Equal(expected, actual)) + << "expected: " << ShapeUtil::HumanString(expected) + << "actual: " << ShapeUtil::HumanString(actual); +} + +TEST_F(HloParserTest, ParseDynamicTuple) { + string shape_string = "(f32[42], u32[<=123,<=456])"; + TF_ASSERT_OK_AND_ASSIGN(Shape actual, ParseShape(shape_string)); + Shape expected = ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeShape(F32, {42}), + ShapeUtil::MakeShape(U32, {123, 456}, {true, true})}); + ASSERT_TRUE(ShapeUtil::Equal(expected, actual)) + << "expected: " << ShapeUtil::HumanString(expected) + << "actual: " << ShapeUtil::HumanString(actual); +} } // 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 5e004ce78ac1fd6da18ab2a54d23ef27e9586cf6..ae8c08cf1d16ad6738962f3be7c1b5512110b1d1 100644 --- a/tensorflow/compiler/xla/service/hlo_pass_pipeline.cc +++ b/tensorflow/compiler/xla/service/hlo_pass_pipeline.cc @@ -77,6 +77,11 @@ std::vector HloPassPipeline::GetEnabledPasses( auto repeated_field = debug_options.xla_disable_hlo_passes(); absl::flat_hash_set disabled_pass_names(repeated_field.begin(), repeated_field.end()); + if (debug_options.xla_disable_all_hlo_passes()) { + VLOG(1) << "*All* passes disabled by --xla_disable_all_hlo_passes."; + return {}; + } + if (!disabled_pass_names.empty()) { VLOG(1) << "Passes disabled by --xla_disable_hlo_passes: " << absl::StrJoin(disabled_pass_names, ", "); @@ -84,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()); } } @@ -113,9 +118,10 @@ void HloPassPipeline::MaybeDumpHlo(const HloModule& module, } const string message = - StrCat("after ", after_pass_name, ", before ", before_pass_name); + absl::StrCat("after ", after_pass_name, ", before ", before_pass_name); hlo_graph_dumper::MaybeDumpHloModule(module, message); VLOG(3) << "HLO " << message << ":"; + VLOG(3) << module.entry_computation_layout().ToString(); XLA_VLOG_LINES(3, module.ToString()); } diff --git a/tensorflow/compiler/xla/service/hlo_pass_pipeline.h b/tensorflow/compiler/xla/service/hlo_pass_pipeline.h index 09e7033ea4ed88849d2f3665d04f74f3f388b3f5..60d72b9d296d71f7bc2f1637bcbec1675513e5df 100644 --- a/tensorflow/compiler/xla/service/hlo_pass_pipeline.h +++ b/tensorflow/compiler/xla/service/hlo_pass_pipeline.h @@ -105,8 +105,6 @@ class HloPassPipeline : public HloPassInterface { std::vector> passes_; std::vector> invariant_checkers_; bool run_called_ = false; - - TF_DISALLOW_COPY_AND_ASSIGN(HloPassPipeline); }; } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_pass_pipeline_test.cc b/tensorflow/compiler/xla/service/hlo_pass_pipeline_test.cc index ee8cb12b231718e09f6ac0d05d7a6887f4c4d746..20384b9da6be4bab447b474f0e2240bcb277a620 100644 --- a/tensorflow/compiler/xla/service/hlo_pass_pipeline_test.cc +++ b/tensorflow/compiler/xla/service/hlo_pass_pipeline_test.cc @@ -19,14 +19,14 @@ 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_parser.h" -#include "tensorflow/compiler/xla/tests/hlo_verified_test_base.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" #include "tensorflow/compiler/xla/util.h" #include "tensorflow/core/lib/core/status_test_util.h" namespace xla { namespace { -class HloPassPipelineTest : public HloVerifiedTestBase { +class HloPassPipelineTest : public HloTestBase { protected: StatusOr ParseModuleGroup( absl::Span hlo_strings) { diff --git a/tensorflow/compiler/xla/service/hlo_profile_printer.cc b/tensorflow/compiler/xla/service/hlo_profile_printer.cc index dcc22793015147aaf3229875078b2989e4ef7559..9cc202aa9f5fe5a20a9da05251ea811137ccaadb 100644 --- a/tensorflow/compiler/xla/service/hlo_profile_printer.cc +++ b/tensorflow/compiler/xla/service/hlo_profile_printer.cc @@ -15,6 +15,8 @@ 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" namespace xla { @@ -25,14 +27,18 @@ string PrintHloProfile(const HloProfilePrinterData& hlo_profile_printer_data, string result; + for (const auto& item : hlo_profile_printer_data.extra_metrics()) { + absl::StrAppend(&result, "Extra metric ", item.first, ": ", + counters[item.second], "\n"); + } + 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; @@ -41,8 +47,9 @@ string PrintHloProfile(const HloProfilePrinterData& hlo_profile_printer_data, // Once we start using this in AOT for real, we will probably need a more // minimal version of HumanReadableProfileBuilder. HumanReadableProfileBuilder builder( - computation_info.name(), counters[computation_info.profile_index()], - clock_rate_ghz); + computation_info.name(), + hlo_profile_printer_data.entry_computation() == computation_info.name(), + counters[computation_info.profile_index()], clock_rate_ghz); for (const auto& instruction_info : instruction_infos) { builder.AddOp( diff --git a/tensorflow/compiler/xla/service/hlo_profile_printer_data.proto b/tensorflow/compiler/xla/service/hlo_profile_printer_data.proto index 9f22b733fe1d676b177039a9d7a3064b8638d7bc..ee66c86ffcb4fb74a24033e05f588a2f4d27dfe4 100644 --- a/tensorflow/compiler/xla/service/hlo_profile_printer_data.proto +++ b/tensorflow/compiler/xla/service/hlo_profile_printer_data.proto @@ -57,4 +57,10 @@ message HloProfilePrinterData { // The size of the profile counters array we will pretty-print. int64 profile_counters_size = 2; + + // Maps extra metric name to the index into the profile counters array. + map extra_metrics = 3; + + // Name of the entry computation. + string entry_computation = 4; } diff --git a/tensorflow/compiler/xla/service/hlo_proto_util.cc b/tensorflow/compiler/xla/service/hlo_proto_util.cc index cf33668f5bfa64a7843efc76e9f6768d18533240..3a9ee57e5551ae5b608f02d9f8bd0428ff16db13 100644 --- a/tensorflow/compiler/xla/service/hlo_proto_util.cc +++ b/tensorflow/compiler/xla/service/hlo_proto_util.cc @@ -39,6 +39,7 @@ HloProto MakeHloProto(const HloModule& module) { StatusOr> CreateModuleFromProto( const HloModuleProto& proto, const HloModuleConfig& module_config) { + VLOG(4) << proto.ShortDebugString(); TF_ASSIGN_OR_RETURN(std::unique_ptr module, HloModule::CreateFromProto(proto, module_config)); TF_RETURN_IF_ERROR( @@ -48,7 +49,7 @@ StatusOr> CreateModuleFromProto( return std::move(module); } -StatusOr> EntryComputationParameterShapes( +StatusOr> EntryComputationParameterShapes( const HloProto& hlo_proto) { if (!hlo_proto.has_hlo_module()) { return NotFound("HloProto missing HloModuleProto."); @@ -57,15 +58,16 @@ StatusOr> EntryComputationParameterShapes( return NotFound("HloProto missing program shape."); } - std::vector parameter_shapes; + std::vector parameter_shapes; const auto& program_shape = hlo_proto.hlo_module().host_program_shape(); - for (const Shape& shape : program_shape.parameters()) { + for (const ShapeProto& shape : program_shape.parameters()) { parameter_shapes.push_back(&shape); } return parameter_shapes; } -StatusOr EntryComputationOutputShape(const HloProto& hlo_proto) { +StatusOr EntryComputationOutputShape( + const HloProto& hlo_proto) { if (!hlo_proto.has_hlo_module()) { return NotFound("HloProto missing HloModuleProto."); } diff --git a/tensorflow/compiler/xla/service/hlo_proto_util.h b/tensorflow/compiler/xla/service/hlo_proto_util.h index 1db82dd6fcaa5d7fe7d65894c1021105f0b26266..31ea2aaffd9cdb76d21edbd0d4a03aa5f865f4f0 100644 --- a/tensorflow/compiler/xla/service/hlo_proto_util.h +++ b/tensorflow/compiler/xla/service/hlo_proto_util.h @@ -43,12 +43,13 @@ StatusOr> CreateModuleFromProto( // Returns the shapes of the parameters of the entry computation. Shape pointers // refer to shapes inside of the given HloProto. -StatusOr> EntryComputationParameterShapes( +StatusOr> EntryComputationParameterShapes( const HloProto& hlo_proto); // Returns the shape of the output of the entry computation. The shape pointer // refers to the output shape inside of the given HloProto. -StatusOr EntryComputationOutputShape(const HloProto& hlo_proto); +StatusOr EntryComputationOutputShape( + const HloProto& hlo_proto); } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_query.cc b/tensorflow/compiler/xla/service/hlo_query.cc index 2d5197be9e6f69f698729e06b7506a5bc6260bcd..f968a4a94453f678f5c17e0b8d1df4aea70c93ea 100644 --- a/tensorflow/compiler/xla/service/hlo_query.cc +++ b/tensorflow/compiler/xla/service/hlo_query.cc @@ -104,5 +104,20 @@ bool IsScalarConstant(const HloInstruction* instruction) { return instruction->IsConstant() && ShapeUtil::IsScalar(instruction->shape()); } +bool ContainsInstrWithOpcode(const HloComputation* comp, + const absl::flat_hash_set& opcodes) { + for (const auto* instr : comp->instructions()) { + if (opcodes.count(instr->opcode())) { + return true; + } + for (const HloComputation* subcomp : instr->called_computations()) { + if (ContainsInstrWithOpcode(subcomp, opcodes)) { + return true; + } + } + } + return false; +} + } // namespace hlo_query } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_query.h b/tensorflow/compiler/xla/service/hlo_query.h index c0826a6aee1f693484207a86ec258c6604d92318..215051f8834fc94eb9e32b508f34b13626ac9349 100644 --- a/tensorflow/compiler/xla/service/hlo_query.h +++ b/tensorflow/compiler/xla/service/hlo_query.h @@ -16,6 +16,8 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_HLO_QUERY_H_ #define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_QUERY_H_ +#include "absl/container/flat_hash_set.h" +#include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" namespace xla { @@ -41,6 +43,12 @@ bool AllOperandsAreConstants(const HloInstruction& instruction); // Returns whether the instruction is a scalar constant. bool IsScalarConstant(const HloInstruction* instruction); +// Determines whether the given computation contains an instruction with one of +// the given opcodes. Checks both comp's instructions and the instructions of +// any computations nested within it. +bool ContainsInstrWithOpcode(const HloComputation* comp, + const absl::flat_hash_set& opcodes); + // Returns an operand of an instruction with the given opcode. If there are // multiple matching operands, then the first matching operand is returned. If // there are no matching operands then nullptr is returned. diff --git a/tensorflow/compiler/xla/service/hlo_reachability.cc b/tensorflow/compiler/xla/service/hlo_reachability.cc index 961930f0a888e90f86e4354fa1373a303af8ec2f..0fced7f15bdaf1dbe349e3b0fc6ada68393c6512 100644 --- a/tensorflow/compiler/xla/service/hlo_reachability.cc +++ b/tensorflow/compiler/xla/service/hlo_reachability.cc @@ -13,6 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include + #include "tensorflow/compiler/xla/service/hlo_reachability.h" namespace xla { @@ -22,7 +24,7 @@ HloReachabilityMap::HloReachabilityMap( : size_(instructions.size()) { bit_vectors_.reserve(size_); for (const HloInstruction* hlo : instructions) { - indices_[hlo] = bit_vectors_.size(); + indices_[GetKey(hlo)] = bit_vectors_.size(); bit_vectors_.emplace_back(size_); } CHECK_EQ(size_, indices_.size()); // instructions should be unique @@ -47,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)); @@ -71,4 +73,70 @@ bool HloReachabilityMap::IsConnected(const HloInstruction* a, return IsReachable(a, b) || IsReachable(b, a); } +std::unique_ptr HloReachabilityMap::Build( + const HloComputation* computation) { + const auto& all = computation->MakeInstructionPostOrder(); + auto result = absl::make_unique(all); + auto channel_dependency_map = computation->ComputeChannelDependencies(); + + std::vector inputs; + for (const HloInstruction* hlo : all) { + inputs.assign(hlo->operands().begin(), hlo->operands().end()); + inputs.insert(inputs.end(), hlo->control_predecessors().begin(), + hlo->control_predecessors().end()); + + switch (hlo->opcode()) { + case HloOpcode::kRecvDone: { + auto it = channel_dependency_map.find(hlo->channel_id()); + if (it != channel_dependency_map.end()) { + absl::c_copy(it->second, std::back_inserter(inputs)); + } + break; + } + case HloOpcode::kAllReduce: { + auto all_reduce_id = hlo->all_reduce_id(); + if (all_reduce_id) { + auto it = channel_dependency_map.find(all_reduce_id.value()); + if (it != channel_dependency_map.end()) { + absl::c_copy(it->second, std::back_inserter(inputs)); + } + } + break; + } + default: + break; + } + + result->FastSetReachabilityToUnion(inputs, hlo); + } + return result; +} + +void HloReachabilityMap::UpdateReachabilityThroughInstruction( + const HloInstruction* instruction) { + std::queue worklist; + worklist.push(instruction); + + std::vector inputs; + + while (!worklist.empty()) { + const HloInstruction* item = worklist.front(); + worklist.pop(); + + inputs.assign(item->operands().begin(), item->operands().end()); + inputs.insert(inputs.end(), item->control_predecessors().begin(), + item->control_predecessors().end()); + + if (SetReachabilityToUnion(inputs, item)) { + // Add immediate successors to worklist. + for (const HloInstruction* user : item->users()) { + worklist.push(user); + } + for (const HloInstruction* succ : item->control_successors()) { + worklist.push(succ); + } + } + } +} + } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_reachability.h b/tensorflow/compiler/xla/service/hlo_reachability.h index 5a5f01f8fd647c74217c80ce4a7633b8957e335f..7823b06a41b3052f6f50f7ffa358de5b23ba679f 100644 --- a/tensorflow/compiler/xla/service/hlo_reachability.h +++ b/tensorflow/compiler/xla/service/hlo_reachability.h @@ -16,27 +16,30 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_HLO_REACHABILITY_H_ #define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_REACHABILITY_H_ +#include #include #include +#include "absl/base/casts.h" #include "absl/container/flat_hash_map.h" #include "absl/types/span.h" #include "tensorflow/compiler/xla/map_util.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" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/platform/types.h" namespace xla { -class HloInstruction; - // A class for representing reachability between HloInstructions. // -// !!! THIS CLASS DOES NOT COMPUTE REACHABILITY !!! It has an adjacency matrix -// and it is up to the user of the class to set the adjacency matrix such that -// it represents reachability, i.e. such that it is transitive. That the graph -// be transitive is thus not an invariant of this class, but it is required for -// the name of the class and its methods to make sense. +// It has an adjacency matrix and it is up to the user of the class to set the +// adjacency matrix such that it represents reachability, i.e. such that it is +// transitive. That the graph be transitive is thus not an invariant of this +// class, but it is required for the name of the class and its methods to make +// sense. class HloReachabilityMap { public: // Sets up a graph with no edges and where the nodes correspond to the given @@ -44,6 +47,15 @@ class HloReachabilityMap { explicit HloReachabilityMap( absl::Span instructions); + // Computes and returns the reachability between HLO instructions in the + // computation. The returned HloReachabilityMap is constructed such that + // HloReachabilityMap::IsReachable(a, b) returns true iff there exists a + // directed path (from producer to consumer) from 'a' to 'b'. Both data + // dependencies (operands) and control dependencies are considered for + // reachability. Trivially an instruction is reachable from itself. + static std::unique_ptr Build( + const HloComputation* computation); + // Set the reachability set of 'instruction' to the union of the reachability // sets of 'inputs'. Upon return, IsReachable(x, instruction) where // 'x' is not 'instruction' will return true iff IsReachable(x, input) is true @@ -70,6 +82,10 @@ class HloReachabilityMap { // adjacency matrix. void SetReachable(const HloInstruction* a, const HloInstruction* b); + // Updates the given reachability map after the immediate predecessor set + // (operands and control predecessors) of 'instruction' has changed. + void UpdateReachabilityThroughInstruction(const HloInstruction* instruction); + // Returns true if "b" is reachable from "a" // // Note that this function only correctly answers queries about reachability @@ -82,6 +98,11 @@ class HloReachabilityMap { // if the set of edges that have been provided to this class are transitive. bool IsConnected(const HloInstruction* a, const HloInstruction* b) const; + // Checks if an instruction is in the Reachability map. + bool IsPresent(const HloInstruction* a) const { + return indices_.contains(GetKey(a)); + } + private: // A bit-vector implementation specialized for this use case which provides a // fast bitwise OR operation not available in tensorflow::gtl::BitMap. @@ -143,18 +164,24 @@ class HloReachabilityMap { absl::Span inputs, const HloInstruction* instruction, BitVector* bit_vector); + uint64 GetKey(const HloInstruction* instruction) const { + uint64 unique_id = absl::bit_cast(instruction->unique_id()); + uint64 module_id = + absl::bit_cast(instruction->parent()->parent()->unique_id()); + return (module_id << 32) | unique_id; + } // Return the index of the given instruction. The value is used to index into // the vector of BitVectors and the BitVectors themselves. int GetIndex(const HloInstruction* instruction) const { - return FindOrDie(indices_, instruction); + return FindOrDie(indices_, GetKey(instruction)); } // The number of instructions in the reachability map. const size_t size_; - // Dense assignment from HloInstruction* to number. These numbers index - // into the bit_vectors_ vector and into the bits within a BitVector. - absl::flat_hash_map indices_; + // Dense assignment from HloInstruction::unique_id to number. These numbers + // index into the bit_vectors_ vector and into the bits within a BitVector. + absl::flat_hash_map indices_; // Bitvectors holding the reachability to each instruction. The bit vector for // instruction X includes ones for each instruction which X is reachable from. diff --git a/tensorflow/compiler/xla/service/hlo_reachability_test.cc b/tensorflow/compiler/xla/service/hlo_reachability_test.cc index d9848cee0bfa904a90aea4626c3ee62c2cbb45b6..595176709806d54fc7c7c5ea301654717096b2d6 100644 --- a/tensorflow/compiler/xla/service/hlo_reachability_test.cc +++ b/tensorflow/compiler/xla/service/hlo_reachability_test.cc @@ -20,13 +20,13 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/test.h" #include "tensorflow/compiler/xla/test_helpers.h" -#include "tensorflow/compiler/xla/tests/hlo_verified_test_base.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" namespace xla { namespace { -class HloReachabilityTest : public HloVerifiedTestBase {}; +class HloReachabilityTest : public HloTestBase {}; TEST_F(HloReachabilityTest, Reachability) { // Construct and test a reachability graph of the following form: @@ -48,7 +48,8 @@ TEST_F(HloReachabilityTest, Reachability) { HloInstruction::CreateConstant(LiteralUtil::CreateR0(0.0f))); auto e = builder.AddInstruction( HloInstruction::CreateConstant(LiteralUtil::CreateR0(0.0f))); - builder.Build(); + auto module = CreateNewVerifiedModule(); + module->AddEntryComputation(builder.Build()); HloReachabilityMap reachability({a, b, c, d, e}); reachability.SetReachable(a, a); @@ -81,6 +82,130 @@ TEST_F(HloReachabilityTest, Reachability) { EXPECT_FALSE(reachability.SetReachabilityToUnion({b, c}, d)); } +TEST_F(HloReachabilityTest, NonTrivialReachability) { + // Test reachability of a non-trivial computation: + // + // const1 const2 + // | | + // | +-------+ + // | | | + // add .. negate + // | . | + // | .... exp + // | | + // +---+ +-+---+ + // | | | + // multiply copy + // + // There is a control dependency from 'add' to 'exp'. + Shape r0f32 = ShapeUtil::MakeShape(F32, {}); + auto builder = HloComputation::Builder(TestName()); + auto constant1 = builder.AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0f))); + auto constant2 = builder.AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(2.0f))); + auto add = builder.AddInstruction(HloInstruction::CreateBinary( + r0f32, HloOpcode::kAdd, constant1, constant2)); + auto negate = builder.AddInstruction( + HloInstruction::CreateUnary(r0f32, HloOpcode::kNegate, constant2)); + auto exp = builder.AddInstruction( + HloInstruction::CreateUnary(r0f32, HloOpcode::kExp, negate)); + auto mul = builder.AddInstruction( + HloInstruction::CreateBinary(r0f32, HloOpcode::kMultiply, add, exp)); + auto copy = builder.AddInstruction( + HloInstruction::CreateUnary(r0f32, HloOpcode::kCopy, exp)); + + auto module = CreateNewVerifiedModule(); + auto computation = + module->AddEntryComputation(builder.Build(/*root_instruction=*/mul)); + + TF_CHECK_OK(add->AddControlDependencyTo(exp)); + auto reachability = HloReachabilityMap::Build(computation); + + EXPECT_TRUE(reachability->IsReachable(constant1, constant1)); + EXPECT_FALSE(reachability->IsReachable(constant1, constant2)); + EXPECT_TRUE(reachability->IsReachable(constant1, add)); + EXPECT_FALSE(reachability->IsReachable(constant1, negate)); + EXPECT_TRUE(reachability->IsReachable(constant1, exp)); + EXPECT_TRUE(reachability->IsReachable(constant1, mul)); + EXPECT_TRUE(reachability->IsReachable(constant1, copy)); + + EXPECT_FALSE(reachability->IsReachable(constant2, constant1)); + EXPECT_TRUE(reachability->IsReachable(constant2, constant2)); + EXPECT_TRUE(reachability->IsReachable(constant2, add)); + EXPECT_TRUE(reachability->IsReachable(constant2, negate)); + EXPECT_TRUE(reachability->IsReachable(constant2, exp)); + EXPECT_TRUE(reachability->IsReachable(constant2, mul)); + EXPECT_TRUE(reachability->IsReachable(constant2, copy)); + + EXPECT_FALSE(reachability->IsReachable(exp, constant1)); + EXPECT_FALSE(reachability->IsReachable(exp, constant2)); + EXPECT_FALSE(reachability->IsReachable(exp, add)); + EXPECT_FALSE(reachability->IsReachable(exp, negate)); + EXPECT_TRUE(reachability->IsReachable(exp, exp)); + EXPECT_TRUE(reachability->IsReachable(exp, mul)); + EXPECT_TRUE(reachability->IsReachable(exp, copy)); + + EXPECT_FALSE(reachability->IsReachable(mul, constant1)); + EXPECT_FALSE(reachability->IsReachable(mul, constant2)); + EXPECT_FALSE(reachability->IsReachable(mul, add)); + EXPECT_FALSE(reachability->IsReachable(mul, negate)); + EXPECT_FALSE(reachability->IsReachable(mul, exp)); + EXPECT_TRUE(reachability->IsReachable(mul, mul)); + EXPECT_FALSE(reachability->IsReachable(mul, copy)); + + EXPECT_TRUE(reachability->IsConnected(constant1, copy)); + EXPECT_TRUE(reachability->IsConnected(copy, constant1)); + EXPECT_FALSE(reachability->IsConnected(negate, add)); + EXPECT_FALSE(reachability->IsConnected(add, negate)); + + // Remove the control dependency then update and verify the reachability map + ASSERT_IS_OK(add->RemoveControlDependencyTo(exp)); + reachability->UpdateReachabilityThroughInstruction(exp); + + EXPECT_TRUE(reachability->IsReachable(constant1, constant1)); + EXPECT_FALSE(reachability->IsReachable(constant1, constant2)); + EXPECT_TRUE(reachability->IsReachable(constant1, add)); + EXPECT_FALSE(reachability->IsReachable(constant1, negate)); + EXPECT_FALSE(reachability->IsReachable(constant1, exp)); + EXPECT_TRUE(reachability->IsReachable(constant1, mul)); + EXPECT_FALSE(reachability->IsReachable(constant1, copy)); + + // Change a use within the graph then update and verify the reachability map + ASSERT_IS_OK(constant2->ReplaceUseWith(negate, constant1)); + reachability->UpdateReachabilityThroughInstruction(negate); + + EXPECT_FALSE(reachability->IsReachable(constant2, constant1)); + EXPECT_TRUE(reachability->IsReachable(constant2, constant2)); + EXPECT_TRUE(reachability->IsReachable(constant2, add)); + EXPECT_FALSE(reachability->IsReachable(constant2, negate)); + EXPECT_FALSE(reachability->IsReachable(constant2, exp)); + EXPECT_TRUE(reachability->IsReachable(constant2, mul)); + EXPECT_FALSE(reachability->IsReachable(constant2, copy)); +} + +TEST_F(HloReachabilityTest, ChannelReachability) { + const Shape shape = ShapeUtil::MakeShape(F32, {5, 7}); + HloComputation::Builder builder("ChannelReachability"); + auto param = builder.AddInstruction( + HloInstruction::CreateParameter(0, shape, "param")); + auto token0 = builder.AddInstruction(HloInstruction::CreateToken()); + auto send = + builder.AddInstruction(HloInstruction::CreateSend(param, token0, 1)); + auto send_done = builder.AddInstruction(HloInstruction::CreateSendDone(send)); + auto token1 = builder.AddInstruction(HloInstruction::CreateToken()); + auto recv = + builder.AddInstruction(HloInstruction::CreateRecv(shape, token1, 1)); + auto recv_done = builder.AddInstruction(HloInstruction::CreateRecvDone(recv)); + + auto module = CreateNewVerifiedModule(); + auto computation = module->AddEntryComputation(builder.Build(recv_done)); + auto reachability = HloReachabilityMap::Build(computation); + EXPECT_TRUE(reachability->IsReachable(param, recv_done)); + EXPECT_FALSE(reachability->IsReachable(send, recv)); + EXPECT_FALSE(reachability->IsReachable(send_done, recv)); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_rematerialization.cc b/tensorflow/compiler/xla/service/hlo_rematerialization.cc index 49e46ecd00ee4370f3e93746348373b79febed3d..a175e4643de2ac6ce07ac00da914d7ab7acca541 100644 --- a/tensorflow/compiler/xla/service/hlo_rematerialization.cc +++ b/tensorflow/compiler/xla/service/hlo_rematerialization.cc @@ -57,13 +57,22 @@ 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()) { case HloOpcode::kCall: case HloOpcode::kConstant: case HloOpcode::kConditional: - case HloOpcode::kCrossReplicaSum: + case HloOpcode::kAllReduce: case HloOpcode::kCustomCall: case HloOpcode::kParameter: case HloOpcode::kWhile: @@ -130,10 +139,10 @@ using ItemList = absl::InlinedVector; // before arbitrary elements. class InstructionList { public: - explicit InstructionList(const std::vector& order) { + explicit InstructionList(const HloInstructionSequence& order) { int64 position = 0; Item* last = nullptr; - for (const HloInstruction* inst : order) { + for (HloInstruction* inst : order.instructions()) { // Add a new item to the linked list. Item* item = new Item; item->next = nullptr; @@ -151,7 +160,7 @@ class InstructionList { // to be monotonically increasing through the list, and so is still useful // for quickly(-ish) determining the order of arbitrary instructions in // the list. - item->instruction = const_cast(inst); + item->instruction = inst; item->position = position; position++; @@ -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; } @@ -927,7 +933,7 @@ Item* PickRematerializationCandidate( StatusOr HloRematerialization::ComputePeakMemory( const HloComputation* computation, - const std::vector& order) const { + const HloInstructionSequence& order) const { InstructionList instruction_list(order); MemoryUsageTracker tracker(computation, size_function_, *points_to_analysis_, instruction_list); @@ -971,8 +977,7 @@ StatusOr HloRematerialization::RematerializeComputation( << HumanReadableNumBytes(computation_peak_memory_.at(computation)); CHECK(!ContainsKey(rematerialized_computations_, computation)); - InstructionList instruction_list( - schedule->sequence(computation).instructions()); + InstructionList instruction_list(schedule->sequence(computation)); MemoryUsageTracker memory_tracker(computation, size_function_, *points_to_analysis_, instruction_list); bool changed = false; @@ -1095,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); @@ -1165,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 " @@ -1184,7 +1189,7 @@ StatusOr HloRematerialization::RematerializeComputation( sequence.clear(); for (auto* item = instruction_list.first(); item != nullptr; item = instruction_list.next(item)) { - const HloInstruction* instruction = item->instruction; + HloInstruction* instruction = item->instruction; sequence.push_back(instruction); } rematerialized_computations_.insert(computation); @@ -1235,10 +1240,8 @@ StatusOr HloRematerialization::Run(HloModule* module) { if (node.context() == CallContext::kSequential) { TF_ASSIGN_OR_RETURN( computation_peak_memory_[node.computation()], - ComputePeakMemory(node.computation(), - module->schedule() - .sequence(node.computation()) - .instructions())); + ComputePeakMemory(node.computation(), module->schedule().sequence( + node.computation()))); } return Status::OK(); }, diff --git a/tensorflow/compiler/xla/service/hlo_rematerialization.h b/tensorflow/compiler/xla/service/hlo_rematerialization.h index 70d83c04f07ca7fd0139f586869e8fe688f958f4..a07d348041b72bba45c6fd1f726f2a0065d01e53 100644 --- a/tensorflow/compiler/xla/service/hlo_rematerialization.h +++ b/tensorflow/compiler/xla/service/hlo_rematerialization.h @@ -87,9 +87,8 @@ class HloRematerialization : public HloModulePass { // peak memory is the maximum total size of all live HLO instruction values at // any program point. 'order' is the order in which the HLO instructions will // be emitted which is used to determine lifespans of HLO values. - StatusOr ComputePeakMemory( - const HloComputation* computation, - const std::vector& order) const; + StatusOr ComputePeakMemory(const HloComputation* computation, + const HloInstructionSequence& order) const; // Returns the peak memory usage of the called computations for the given // instruction. Zero is returned if the instruction calls no computations. diff --git a/tensorflow/compiler/xla/service/hlo_rematerialization_test.cc b/tensorflow/compiler/xla/service/hlo_rematerialization_test.cc index f7e82fb1f88e856305f6f481a451d4cd64ba4acf..102a360ad8116d8781baf9cb7627a920f4a687c4 100644 --- a/tensorflow/compiler/xla/service/hlo_rematerialization_test.cc +++ b/tensorflow/compiler/xla/service/hlo_rematerialization_test.cc @@ -24,7 +24,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_opcode.h" #include "tensorflow/compiler/xla/service/hlo_ordering.h" #include "tensorflow/compiler/xla/shape_util.h" -#include "tensorflow/compiler/xla/tests/hlo_verified_test_base.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" @@ -36,7 +36,7 @@ namespace op = xla::testing::opcode_matchers; using ::testing::_; -class HloRematerializationTest : public HloVerifiedTestBase { +class HloRematerializationTest : public HloTestBase { protected: // Creates and returns a computation which can benefit from // rematerialization. The computation looks like: @@ -162,7 +162,7 @@ class HloRematerializationTest : public HloVerifiedTestBase { // Test rematerialization of a single computation produced by // MakeRematerializableComputation. TEST_F(HloRematerializationTest, SingleComputation) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* computation = module->AddEntryComputation(MakeRematerializableComputation()); @@ -177,7 +177,7 @@ TEST_F(HloRematerializationTest, SingleComputation) { // with rematerialization so pick a memory limit between these values (14KB). TF_ASSERT_OK_AND_ASSIGN(bool changed, RunHloRematerialization( - /*memory_limit_bytes=*/14 * 1024, module)); + /*memory_limit_bytes=*/14 * 1024, module.get())); EXPECT_TRUE(changed); // Root should not have changed. @@ -203,7 +203,7 @@ TEST_F(HloRematerializationTest, SingleComputation) { // MakeRematerializableComputation but with a sufficiently high memory limit // such that no instructions are rematerialized. TEST_F(HloRematerializationTest, SingleComputationNoRematerialization) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* computation = module->AddEntryComputation(MakeRematerializableComputation()); @@ -211,7 +211,7 @@ TEST_F(HloRematerializationTest, SingleComputationNoRematerialization) { TF_ASSERT_OK_AND_ASSIGN(bool changed, RunHloRematerialization( - /*memory_limit_bytes=*/20 * 1024, module)); + /*memory_limit_bytes=*/20 * 1024, module.get())); // No instructions should have been materialized. EXPECT_FALSE(changed); @@ -225,7 +225,7 @@ TEST_F(HloRematerializationTest, SingleComputationNoRematerialization) { // computation should be the one chosen because rematerialization in the while // will presumably be more expensive. TEST_F(HloRematerializationTest, RematerializeAroundWhile) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto cond_builder = HloComputation::Builder(TestName() + ".cond"); cond_builder.AddInstruction( @@ -249,7 +249,7 @@ TEST_F(HloRematerializationTest, RematerializeAroundWhile) { // bit lower (17KB) to force rematerialization of the entry computation. TF_ASSERT_OK_AND_ASSIGN(bool changed, RunHloRematerialization( - /*memory_limit_bytes=*/17 * 1024, module)); + /*memory_limit_bytes=*/17 * 1024, module.get())); EXPECT_TRUE(changed); // Only the entry computation should have a rematerialized instruction added. @@ -261,7 +261,7 @@ TEST_F(HloRematerializationTest, RematerializeAroundWhile) { // while. Both the entry computation and while body computation should have // computations rematerialized. TEST_F(HloRematerializationTest, RematerializeEntryAndWhileBody) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto cond_builder = HloComputation::Builder(TestName() + ".cond"); cond_builder.AddInstruction( @@ -282,7 +282,7 @@ TEST_F(HloRematerializationTest, RematerializeEntryAndWhileBody) { TF_ASSERT_OK_AND_ASSIGN(bool changed, RunHloRematerialization( - /*memory_limit_bytes=*/15 * 1024, module)); + /*memory_limit_bytes=*/15 * 1024, module.get())); EXPECT_TRUE(changed); // Both computations should have rematerialized instructions added. @@ -293,7 +293,7 @@ TEST_F(HloRematerializationTest, RematerializeEntryAndWhileBody) { // Test rematerialization of a doubly nested computation. All computations // should have an instruction rematerialized. TEST_F(HloRematerializationTest, RematerializeNestedComputations) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto cond_builder = HloComputation::Builder(TestName() + ".cond"); cond_builder.AddInstruction( @@ -321,7 +321,7 @@ TEST_F(HloRematerializationTest, RematerializeNestedComputations) { // ~12K so pick something slightly larger. TF_ASSERT_OK_AND_ASSIGN(bool changed, RunHloRematerialization( - /*memory_limit_bytes=*/13 * 1024, module)); + /*memory_limit_bytes=*/13 * 1024, module.get())); EXPECT_TRUE(changed); // All computations should have rematerialized instructions added. @@ -346,7 +346,7 @@ TEST_F(HloRematerializationTest, RngNotRematerialized) { // // F32[1024] add_2 = add(rng, add(tanh, add_1)) // LIVE: add_2 + add_1 + // // rng + tanh + exp - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto builder = HloComputation::Builder(TestName()); auto param = builder.AddInstruction( @@ -390,7 +390,7 @@ TEST_F(HloRematerializationTest, RngNotRematerialized) { TF_ASSERT_OK_AND_ASSIGN( bool changed, RunHloRematerialization( - /*memory_limit_bytes=*/4 * ByteSizeOf(vec1024_shape_), module)); + /*memory_limit_bytes=*/4 * ByteSizeOf(vec1024_shape_), module.get())); EXPECT_TRUE(changed); // The rng should not have been rematerialized. EXPECT_EQ(count_rngs(entry_computation), 1); @@ -420,7 +420,7 @@ TEST_F(HloRematerializationTest, InstructionRematerializedMultipleTimes) { // The value %bcast is live across each call of Subcomputation (which requires // 8KB) though the value is not used in the calls. Rematerializing %bcast // across these calls reduces peak memory use from ~20KB down to ~16KB. - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* subcomputation = nullptr; { @@ -482,7 +482,7 @@ TEST_F(HloRematerializationTest, InstructionRematerializedMultipleTimes) { // rematerialization). TF_ASSERT_OK_AND_ASSIGN(bool changed, RunHloRematerialization( - /*memory_limit_bytes=*/22 * 1024, module)); + /*memory_limit_bytes=*/22 * 1024, module.get())); EXPECT_TRUE(changed); // The broadcast should have been rematerialized 3 times. @@ -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 {}; @@ -533,7 +579,7 @@ TEST_P(IndirectUseTest, IndirectUseNotRematerialized) { // (ie %bcast is used indirectly by %negate), otherwise the %negate operand // aliases %add_2. const bool indirectly_used = GetParam(); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloComputation* subcomputation = nullptr; { @@ -576,7 +622,7 @@ TEST_P(IndirectUseTest, IndirectUseNotRematerialized) { // rematerialization). TF_ASSERT_OK_AND_ASSIGN(bool changed, RunHloRematerialization( - /*memory_limit_bytes=*/22 * 1024, module)); + /*memory_limit_bytes=*/22 * 1024, module.get())); // Rematerialization should only occur if the rematerializable instruction has // no indirect uses. if (indirectly_used) { @@ -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 3f0ca342b4c84216ddd5ee553848360d8bd1ff0b..d7d66ae1c4592723ca991d5ee971fa72cc1af90a 100644 --- a/tensorflow/compiler/xla/service/hlo_runner.cc +++ b/tensorflow/compiler/xla/service/hlo_runner.cc @@ -205,6 +205,40 @@ StatusOr HloRunner::ExecuteWithDeviceBuffers( /*profile=*/profile); } +StatusOr HloRunner::ExecuteWithDeviceBuffers( + std::unique_ptr executable, + const absl::Span arguments, + ExecutionProfile* profile) { + // Get service run options. + se::Stream stream(backend().default_stream_executor()); + stream.Init(); + ServiceExecutableRunOptions service_run_options = + GetServiceRunOptionsForDevice(backend().default_device_ordinal(), &stream, + nullptr); + + TF_ASSIGN_OR_RETURN( + ScopedShapedBuffer retval, + executable->ExecuteOnStreamWrapper(&service_run_options, + /*profile=*/profile, arguments)); + TF_RETURN_IF_ERROR(stream.BlockHostUntilDone()); + return std::move(retval); +} + +StatusOr HloRunner::ExecuteWithDeviceBuffers( + std::unique_ptr executable, + const absl::Span arguments, + ExecutionProfile* profile) { + std::vector argument_pointers; + argument_pointers.reserve(arguments.size()); + for (const auto& argument : arguments) { + argument_pointers.push_back(&argument); + } + return ExecuteWithDeviceBuffers( + /*executable=*/std::move(executable), + /*arguments=*/argument_pointers, + /*profile=*/profile); +} + StatusOr> HloRunner::ExecuteReplicated( std::unique_ptr module, const ReplicatedExecuteOptions& options) { @@ -349,9 +383,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 2e934bf66ae43ea412f242030b874dddb6d3722d..bb792cf8c9825ff67ca33bbcf2c3c32b1a0ecb85 100644 --- a/tensorflow/compiler/xla/service/hlo_runner.h +++ b/tensorflow/compiler/xla/service/hlo_runner.h @@ -136,6 +136,21 @@ class HloRunner { const absl::Span arguments, bool run_hlo_passes = true, ExecutionProfile* profile = nullptr); + StatusOr ExecuteWithDeviceBuffers( + std::unique_ptr executable, + const absl::Span arguments, + ExecutionProfile* profile = nullptr); + + StatusOr ExecuteWithDeviceBuffers( + std::unique_ptr executable, + const absl::Span arguments, + ExecutionProfile* profile = nullptr); + + // Creates an executable object given an HLO module. If run_hlo_passes is + // true, the HLO passes will be run as part of compilation. + StatusOr> CreateExecutable( + std::unique_ptr module, bool run_hlo_passes); + // Executes a given HLO module into a set of replicas, and returns a map // with the replica number as key, and the corresponding returned literal as // value. @@ -152,11 +167,6 @@ class HloRunner { const Backend& backend() const; private: - // Creates an executable object given an HLO module. If run_hlo_passes is - // true, the HLO passes will be run before. - StatusOr> CreateExecutable( - std::unique_ptr module, bool run_hlo_passes); - // Creates a ServiceExecutableRunOptions object to configure a run on device, // using the provided stream object. If device_assignment is not nullptr, it // will be used to configure the replication parameters. Replicated executions diff --git a/tensorflow/compiler/xla/service/hlo_schedule.cc b/tensorflow/compiler/xla/service/hlo_schedule.cc index 0778ff52174ef89c476950f2c830268a63888382..e75373501cffac6a736be89e9f6139b6ff2cdbc1 100644 --- a/tensorflow/compiler/xla/service/hlo_schedule.cc +++ b/tensorflow/compiler/xla/service/hlo_schedule.cc @@ -46,8 +46,8 @@ namespace xla { << "No computation exists in HLO module with id " << computation_id; const HloComputation* computation = comp_it->second; - absl::flat_hash_map id_to_instruction; - for (const HloInstruction* instruction : computation->instructions()) { + absl::flat_hash_map id_to_instruction; + for (HloInstruction* instruction : computation->instructions()) { id_to_instruction[instruction->unique_id()] = instruction; } @@ -81,9 +81,8 @@ StatusOr HloSchedule::ToProto() const { return std::move(proto); } -void HloSchedule::set_sequence( - const HloComputation* computation, - absl::Span sequence) { +void HloSchedule::set_sequence(const HloComputation* computation, + absl::Span sequence) { set_sequence(computation, HloInstructionSequence(sequence)); } @@ -114,8 +113,8 @@ Status HloSchedule::UpdateComputationSchedule( const HloComputation* computation) { // Map from unique ID to HloInstruction pointer for instructions in the // computation. - absl::flat_hash_map id_to_instruction; - for (const HloInstruction* instruction : computation->instructions()) { + absl::flat_hash_map id_to_instruction; + for (HloInstruction* instruction : computation->instructions()) { InsertOrDie(&id_to_instruction, instruction->unique_id(), instruction); } @@ -128,7 +127,7 @@ Status HloSchedule::UpdateComputationSchedule( // Map from HloInstruction X to newly added instructions (instruction is in // computation, but not in schedule) which use X. If an instruction is not in // the map, then it has no users which are newly added instructions. - absl::flat_hash_map> + absl::flat_hash_map> new_instruction_uses; // For each newly added instruction, this is the count of the instruction's @@ -138,10 +137,10 @@ Status HloSchedule::UpdateComputationSchedule( // Create a worklist of newly added instructions which are ready to be added // to the schedule. Initialize worklist with those that have zero operands. - std::queue worklist; + std::queue worklist; - for (const HloInstruction* instruction : computation->instructions()) { - if (ids_in_schedule.count(instruction->unique_id()) == 0) { + for (HloInstruction* instruction : computation->instructions()) { + 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); @@ -161,17 +160,17 @@ Status HloSchedule::UpdateComputationSchedule( // Lambda which schedules all instructions on the worklist. auto schedule_worklist = [&]() { while (!worklist.empty()) { - const HloInstruction* instruction = worklist.front(); + HloInstruction* instruction = worklist.front(); worklist.pop(); new_sequence.push_back(instruction); - std::vector* new_users = + std::vector* new_users = tensorflow::gtl::FindOrNull(new_instruction_uses, instruction); if (new_users != nullptr) { // This just-scheduled instruction has users which are newly added to // the module. Update the number of unscheduled operands and push the // newly added instruction to the worklist if it is ready to // schedule. - for (const HloInstruction* new_user : *new_users) { + for (HloInstruction* new_user : *new_users) { unscheduled_operand_count.at(new_user)--; CHECK_GE(unscheduled_operand_count.at(new_user), 0); if (unscheduled_operand_count.at(new_user) == 0) { @@ -205,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()) { @@ -216,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; @@ -245,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."; } @@ -264,9 +263,12 @@ Status HloSchedule::Verify() const { } TF_RET_CHECK(instruction_position.size() == - computation->instruction_count()); + computation->instruction_count()) + << "Schedule for computation " << computation->name() << " has " + << 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 0a714101ee587aa847fa674bbde5586287c51f33..a5f54ae2c33259d080631061dff9ae40b41495dc 100644 --- a/tensorflow/compiler/xla/service/hlo_schedule.h +++ b/tensorflow/compiler/xla/service/hlo_schedule.h @@ -35,14 +35,14 @@ class HloInstructionSequence { public: HloInstructionSequence() = default; explicit HloInstructionSequence( - absl::Span instructions) { - for (const HloInstruction* instruction : instructions) { + absl::Span instructions) { + for (HloInstruction* instruction : instructions) { push_back(instruction); } } // Adds the instruction to the end of the sequence. - void push_back(const HloInstruction* instruction) { + void push_back(HloInstruction* instruction) { instruction_sequence_.push_back(instruction); id_sequence_.push_back(instruction->unique_id()); } @@ -56,7 +56,7 @@ class HloInstructionSequence { int64 size() const { return instruction_sequence_.size(); } // Returns the sequence of HLO instructions. - const std::vector& instructions() const { + const std::vector& instructions() const { return instruction_sequence_; } @@ -65,7 +65,7 @@ class HloInstructionSequence { private: // The sequence as HloInstructions. - std::vector instruction_sequence_; + std::vector instruction_sequence_; // The sequence of HLO instructions, represented by their unique IDs. The // sequence is stored as both HloInstructions and unique IDs because the @@ -98,7 +98,7 @@ class HloSchedule { // Sets the sequence for the given computation to the given sequence. void set_sequence(const HloComputation* computation, - absl::Span sequence); + absl::Span sequence); void set_sequence(const HloComputation* computation, HloInstructionSequence sequence); @@ -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_schedule_test.cc b/tensorflow/compiler/xla/service/hlo_schedule_test.cc index 1424569ac1f62e4b965876141f1eb40be4f15bea..0e56e6f760e35ddcb45c6f58771d78405a09acfe 100644 --- a/tensorflow/compiler/xla/service/hlo_schedule_test.cc +++ b/tensorflow/compiler/xla/service/hlo_schedule_test.cc @@ -56,10 +56,10 @@ ENTRY main { ParseHloString(module_str)); TF_ASSERT_OK_AND_ASSIGN( HloSchedule schedule, - ScheduleModule(*module, [](const BufferValue& buffer) { + ScheduleModule(module.get(), [](const BufferValue& buffer) { return ShapeUtil::ByteSizeOf(buffer.shape()); })); - const std::vector& entry_schedule = + const auto& entry_schedule = schedule.sequence(module->entry_computation()).instructions(); EXPECT_EQ(entry_schedule.size(), 6); @@ -90,7 +90,7 @@ ENTRY main { ParseHloString(module_str)); TF_ASSERT_OK_AND_ASSIGN( HloSchedule schedule, - ScheduleModule(*module, [](const BufferValue& buffer) { + ScheduleModule(module.get(), [](const BufferValue& buffer) { return ShapeUtil::ByteSizeOf(buffer.shape()); })); @@ -139,7 +139,7 @@ ENTRY main { ParseHloString(module_str)); TF_ASSERT_OK_AND_ASSIGN( HloSchedule schedule, - ScheduleModule(*module, [](const BufferValue& buffer) { + ScheduleModule(module.get(), [](const BufferValue& buffer) { return ShapeUtil::ByteSizeOf(buffer.shape()); })); @@ -183,7 +183,7 @@ ENTRY main { ParseHloString(module_str)); TF_ASSERT_OK_AND_ASSIGN( HloSchedule schedule, - ScheduleModule(*module, [](const BufferValue& buffer) { + ScheduleModule(module.get(), [](const BufferValue& buffer) { return ShapeUtil::ByteSizeOf(buffer.shape()); })); @@ -244,7 +244,7 @@ ENTRY %WhileLoop () -> s32[] { ParseHloString(module_str)); TF_ASSERT_OK_AND_ASSIGN( HloSchedule schedule, - ScheduleModule(*module, [](const BufferValue& buffer) { + ScheduleModule(module.get(), [](const BufferValue& buffer) { return ShapeUtil::ByteSizeOf(buffer.shape(), /*pointer_size=*/sizeof(void*)); })); @@ -313,7 +313,7 @@ ENTRY %WhileLoop () -> s32[] { ParseHloString(module_str)); TF_ASSERT_OK_AND_ASSIGN( HloSchedule schedule, - ScheduleModule(*module, [](const BufferValue& buffer) { + ScheduleModule(module.get(), [](const BufferValue& buffer) { return ShapeUtil::ByteSizeOf(buffer.shape(), /*pointer_size=*/sizeof(void*)); })); diff --git a/tensorflow/compiler/xla/service/hlo_sharding.cc b/tensorflow/compiler/xla/service/hlo_sharding.cc index 70a860c356ca2fb1c4c973ea3d96c50fabc2c7c2..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" @@ -30,7 +31,7 @@ HloSharding HloSharding::AssignDevice(int64 device_id) { } HloSharding HloSharding::Tile1D(const Shape& input_shape, int64 num_tiles) { - CHECK_EQ(1, ShapeUtil::Rank(input_shape)); + CHECK_EQ(1, input_shape.rank()); CHECK_GT(num_tiles, 1); std::vector dimensions(1, num_tiles); Array assignment(dimensions); @@ -57,7 +58,7 @@ HloSharding HloSharding::Tuple(const ShapeTree& sub_shardings) { HloSharding HloSharding::Tuple(const Shape& tuple_shape, absl::Span shardings) { - CHECK(ShapeUtil::IsTuple(tuple_shape)) << ShapeUtil::HumanString(tuple_shape); + CHECK(tuple_shape.IsTuple()) << ShapeUtil::HumanString(tuple_shape); for (auto& sharding : shardings) { CHECK(!sharding.IsTuple()) << sharding.ToString(); } @@ -70,7 +71,7 @@ HloSharding HloSharding::Tuple(const Shape& tuple_shape, HloSharding HloSharding::SingleTuple(const Shape& tuple_shape, const HloSharding& sharding) { - CHECK(ShapeUtil::IsTuple(tuple_shape)) << ShapeUtil::HumanString(tuple_shape); + CHECK(tuple_shape.IsTuple()) << ShapeUtil::HumanString(tuple_shape); CHECK(!sharding.IsTuple()) << sharding.ToString(); int64 leaf_count = RequiredLeaves(tuple_shape); std::vector flattened_list; @@ -80,7 +81,7 @@ HloSharding HloSharding::SingleTuple(const Shape& tuple_shape, HloSharding HloSharding::Single(const Shape& shape, const HloSharding& sharding) { - return ShapeUtil::IsTuple(shape) ? SingleTuple(shape, sharding) : sharding; + return shape.IsTuple() ? SingleTuple(shape, sharding) : sharding; } string HloSharding::ToString() const { @@ -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 { @@ -269,7 +269,7 @@ int64 HloSharding::GetUniqueDevice() const { } Status HloSharding::ValidateTuple(const Shape& shape, int64 num_devices) const { - if (!ShapeUtil::IsTuple(shape)) { + if (!shape.IsTuple()) { return tensorflow::errors::InvalidArgument( StrCat("Sharding is tuple-shaped but validation shape is not.")); } @@ -305,7 +305,7 @@ Status HloSharding::Validate(const Shape& shape, int64 num_devices) const { Status HloSharding::ValidateNonTuple(const Shape& shape, int64 num_devices) const { - if (ShapeUtil::IsTuple(shape)) { + if (shape.IsTuple()) { return tensorflow::errors::InvalidArgument( StrCat("Validation shape is a tuple but sharding is not.")); } @@ -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")); } @@ -340,7 +340,7 @@ Status HloSharding::ValidateNonTuple(const Shape& shape, } // The tile assignment tensor must have the same rank as the input. - if (ShapeUtil::Rank(shape) != tile_assignment_.num_dimensions()) { + if (shape.rank() != tile_assignment_.num_dimensions()) { return tensorflow::errors::InvalidArgument( "Number of tile assignment dimensions is different to the input rank. " "sharding=", @@ -437,8 +437,8 @@ Shape HloSharding::TileShape(const Shape& shape) const { } Shape result_shape = shape; for (int64 i = 0; i < shape.dimensions_size(); ++i) { - (*result_shape.mutable_dimensions())[i] = - CeilOfRatio(shape.dimensions(i), tile_assignment_.dim(i)); + result_shape.set_dimensions( + i, CeilOfRatio(shape.dimensions(i), tile_assignment_.dim(i))); } return result_shape; } @@ -455,7 +455,7 @@ HloSharding HloSharding::GetSubSharding(const Shape& shape, } sub_shape = &ShapeUtil::GetSubshape(*sub_shape, {idx}); } - if (ShapeUtil::IsTuple(*sub_shape)) { + if (sub_shape->IsTuple()) { auto begin_it = tuple_elements_.begin() + sharding_index; std::vector sub_shardings( begin_it, begin_it + ShapeUtil::GetLeafCount(*sub_shape)); 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 88329c899794a6e0f5102d181d6161fe17f89932..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(); @@ -234,7 +234,7 @@ StatusOr ApplyShardingFromUsers(HloInstruction* instruction, if (instruction->users().empty()) { // No sharding from users, use domain_sharding, after checking // compatibility. - TF_RET_CHECK(ShapeUtil::IsTuple(instruction->shape()) && + TF_RET_CHECK(instruction->shape().IsTuple() && ShapeUtil::GetLeafCount(instruction->shape()) == domain_sharding.tuple_elements().size()); instruction->set_sharding(domain_sharding); @@ -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(const_cast(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. @@ -266,7 +266,7 @@ StatusOr ApplyShardingFromUsers(HloInstruction* instruction, AssignmentKind sub_assigned = AssignmentKind::kUnassigned; TF_ASSIGN_OR_RETURN(ShapeTree user_sharding_tree, GetShardingTreeFromUser(*instruction, *user)); - if (ShapeUtil::IsTuple(instruction->shape())) { + if (instruction->shape().IsTuple()) { // For tuple-shaped instructions collect individual tuple subshardings // from the uses, and then combine them into the tuple sharding. // If the user is a GTE its sharding concerns only the subtree of @@ -298,7 +298,7 @@ StatusOr ApplyShardingFromUsers(HloInstruction* instruction, } if (assigned == AssignmentKind::kAssigned) { - if (ShapeUtil::IsTuple(instruction->shape())) { + if (instruction->shape().IsTuple()) { instruction->set_sharding(HloSharding::Tuple(sharding_tree)); } else { TF_RET_CHECK(sharding_tree.leaf_count() == 1); @@ -361,7 +361,7 @@ Status ApplyDomainSharding(const DomainMetadata::Domain& domain, // kUnassignedDevice. Indeed in case of doubt it is better to leave the // entire tuple unassigned, and let the device placer decide for it. if (instruction->sharding().UsesDevice(kUnassignedDevice)) { - TF_RET_CHECK(ShapeUtil::IsTuple(instruction->shape())) + TF_RET_CHECK(instruction->shape().IsTuple()) << "Only tuples can have kUnassignedDevice sub shardings"; instruction->clear_sharding(); } diff --git a/tensorflow/compiler/xla/service/hlo_subcomputation_unification_test.cc b/tensorflow/compiler/xla/service/hlo_subcomputation_unification_test.cc index 45c684d66752862eec301b8943d350804f070309..c1073911ea9dc3811c195e27bcbae9b00929ad17 100644 --- a/tensorflow/compiler/xla/service/hlo_subcomputation_unification_test.cc +++ b/tensorflow/compiler/xla/service/hlo_subcomputation_unification_test.cc @@ -66,7 +66,7 @@ class HloSubcomputationUnificationTest : public HloTestBase { }; TEST_F(HloSubcomputationUnificationTest, UnifyIdentities) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto builder = HloComputation::Builder(TestName()); auto callee1 = @@ -103,7 +103,7 @@ TEST_F(HloSubcomputationUnificationTest, UnifyIdentities) { } TEST_F(HloSubcomputationUnificationTest, UnifyAdditions) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto builder = HloComputation::Builder(TestName()); auto callee1 = @@ -143,7 +143,7 @@ TEST_F(HloSubcomputationUnificationTest, UnifyAdditions) { // Do not unify subcomputations with different parameter shapes. TEST_F(HloSubcomputationUnificationTest, DifferentParameterShapes) { - auto module = CreateNewModule(); + auto module = CreateNewUnverifiedModule(); auto builder = HloComputation::Builder(TestName()); auto callee1 = @@ -184,7 +184,7 @@ TEST_F(HloSubcomputationUnificationTest, DifferentParameterShapes) { // Regression test for b/31466798. Checks that entry_computation is still valid // after unification. TEST_F(HloSubcomputationUnificationTest, TwoIdenticalComputations) { - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); for (int i = 0; i < 2; ++i) { HloComputation::Builder builder("pow"); auto x = diff --git a/tensorflow/compiler/xla/service/hlo_tfgraph_builder.cc b/tensorflow/compiler/xla/service/hlo_tfgraph_builder.cc index 487653344976a10e18ba667085525ba1ecbb8612..c1f69db74eafb7743e85f499f2f4828ed0375501 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, '_'); } @@ -159,7 +158,7 @@ void HloTfGraphBuilder::SetNodeAttrs(const HloInstruction* instruction, // Set the layout. if (LayoutUtil::HasLayout(instruction->shape())) { string layout_string; - if (ShapeUtil::IsTuple(instruction->shape())) { + if (instruction->shape().IsTuple()) { // 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()); diff --git a/tensorflow/compiler/xla/service/hlo_tfgraph_builder_test.cc b/tensorflow/compiler/xla/service/hlo_tfgraph_builder_test.cc index 6fd734a2b9e6c8c9fca76a944ca3df4c3b8a212f..1e2b31a1f2bb4865faafc3d14e2b194e3aa171a1 100644 --- a/tensorflow/compiler/xla/service/hlo_tfgraph_builder_test.cc +++ b/tensorflow/compiler/xla/service/hlo_tfgraph_builder_test.cc @@ -14,7 +14,7 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/compiler/xla/service/hlo_tfgraph_builder.h" -#include "tensorflow/compiler/xla/tests/hlo_verified_test_base.h" +#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" @@ -24,7 +24,7 @@ namespace { using ::tensorflow::GraphDef; -class HloTfGraphBuilderTest : public HloVerifiedTestBase { +class HloTfGraphBuilderTest : public HloTestBase { protected: HloTfGraphBuilderTest() {} HloTfGraphBuilder generator_; diff --git a/tensorflow/compiler/xla/service/hlo_token.h b/tensorflow/compiler/xla/service/hlo_token.h deleted file mode 100644 index 4458c251dee4af365e39027dd4289925c8890efd..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/service/hlo_token.h +++ /dev/null @@ -1,78 +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_SERVICE_HLO_TOKEN_H_ -#define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_TOKEN_H_ - -#include - -#include "tensorflow/compiler/xla/types.h" -#include "tensorflow/core/platform/types.h" - -namespace xla { - -// Defines different kinds of tokens in a hlo module string. -// -// You shouldn't need to use this directly unless you're using HloLexer -// directly, and you probably don't need to do that. Use hlo_parser instead. -enum class TokKind { - // Markers - kEof, - kError, - - // Tokens with no info. - kEqual, // = - kComma, // , - kColon, // : - kLsquare, - kRsquare, // [ ] - kLbrace, - kRbrace, // { } - kLparen, - kRparen, // ( ) - - kArrow, // -> - - // Keywords - kw_HloModule, - kw_ENTRY, - kw_ROOT, - kw_true, - kw_false, - kw_maximal, - kw_replicated, - kw_nan, - kw_inf, - - kNegInf, // -inf - - // Typed tokens. - kName, // %foo - kAttributeName, // dimensions= - kDimLabels, // [0-9bf]{2,}_[0-9io]{2,}->[0-9bf]{2,} - kDxD, // [0-9]+(x[0-9]+)+ - kPad, // [0-9]+_[0-9]+(_[0-9]+)?(x[0-9]+_[0-9]+(_[0-9]+)?)* - kIdent, // other identifiers - kString, // "abcd\"\n" - kShape, // f32[2,3]{1,0} - kInt, // 42 - kDecimal, // 4.2 -}; - -string TokKindToString(TokKind kind); - -} // namespace xla - -#endif // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_TOKEN_H_ diff --git a/tensorflow/compiler/xla/service/hlo_value.cc b/tensorflow/compiler/xla/service/hlo_value.cc index 59594ab2f0f70a206c73e998dbfa69c2c5c7ba43..218b33b2ac2b86edc30b2f014ba206c71da37682 100644 --- a/tensorflow/compiler/xla/service/hlo_value.cc +++ b/tensorflow/compiler/xla/service/hlo_value.cc @@ -46,7 +46,7 @@ const Shape& HloPosition::shape() const { string HloPosition::ToString() const { string index_str = - ShapeUtil::IsTuple(instruction->shape()) ? (" " + index.ToString()) : ""; + instruction->shape().IsTuple() ? (" " + index.ToString()) : ""; return StrCat(instruction->name(), index_str); } @@ -56,10 +56,9 @@ std::ostream& operator<<(std::ostream& out, const HloPosition& position) { } string HloUse::ToString() const { - string index_str = - ShapeUtil::IsTuple(instruction->operand(operand_number)->shape()) - ? (" " + operand_index.ToString()) - : ""; + string index_str = instruction->operand(operand_number)->shape().IsTuple() + ? (" " + operand_index.ToString()) + : ""; return StrCat(instruction->name(), ", operand ", operand_number, index_str); } @@ -88,7 +87,7 @@ bool HloValue::operator!=(const HloValue& other) const { } string HloValue::ToShortString() const { - string index_str = ShapeUtil::IsTuple(defining_instruction()->shape()) + string index_str = defining_instruction()->shape().IsTuple() ? defining_index().ToString() : ""; return StrCat(id(), " ", is_phi_ ? "PHI " : "", @@ -210,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_value.h b/tensorflow/compiler/xla/service/hlo_value.h index b6670d409b92e8be42f5cdb40fba8d662ae83958..1f01b0bb365450a933da9cc443db5223c06903f0 100644 --- a/tensorflow/compiler/xla/service/hlo_value.h +++ b/tensorflow/compiler/xla/service/hlo_value.h @@ -166,9 +166,6 @@ class HloValue : public BufferValue { // Whether this value is live out of the HLO module. bool live_out_of_module_ = false; - - // Whether this value is live out of its computation. - bool live_out_of_computation_ = false; }; std::ostream& operator<<(std::ostream& out, const HloValue& hlo_value); diff --git a/tensorflow/compiler/xla/service/hlo_verifier.cc b/tensorflow/compiler/xla/service/hlo_verifier.cc index e955b905ad84923d385624ee7d592ce6c4d305d5..144c01eac1c06bb067c9f29f29b536c459ea273e 100644 --- a/tensorflow/compiler/xla/service/hlo_verifier.cc +++ b/tensorflow/compiler/xla/service/hlo_verifier.cc @@ -18,6 +18,7 @@ limitations under the License. #include "absl/container/flat_hash_map.h" #include "absl/strings/str_join.h" #include "tensorflow/compiler/xla/service/hlo_casting_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" #include "tensorflow/compiler/xla/service/hlo_verifier.h" @@ -43,7 +44,7 @@ bool IsCallerInstruction(HloInstruction* hlo) { case HloOpcode::kCall: case HloOpcode::kConditional: case HloOpcode::kWhile: - case HloOpcode::kCrossReplicaSum: + case HloOpcode::kAllReduce: case HloOpcode::kMap: case HloOpcode::kReduce: case HloOpcode::kReduceWindow: @@ -65,7 +66,9 @@ Status ShapeVerifier::Preprocess(HloInstruction* hlo) { return VerifyNotSparse(hlo->shape()); } -static Status CheckOperandCount(const HloInstruction* hlo, int expected) { +namespace { + +Status CheckOperandCount(const HloInstruction* hlo, int expected) { if (hlo->operand_count() != expected) { return InternalError("Expected %d operands for %s instruction: %s", expected, HloOpcodeString(hlo->opcode()), @@ -74,6 +77,19 @@ static Status CheckOperandCount(const HloInstruction* hlo, int expected) { return Status::OK(); } +Status CheckParameterCount(const HloInstruction* calling_instruction, + const HloComputation* computation, int expected) { + if (computation->num_parameters() != expected) { + return InternalError( + "Expected computation %s called from %s to have %d parameters, has %d", + computation->name(), calling_instruction->name(), expected, + computation->num_parameters()); + } + return Status::OK(); +} + +} // namespace + Status ShapeVerifier::HandleElementwiseUnary(HloInstruction* hlo) { return CheckUnaryShape(hlo); } @@ -137,8 +153,8 @@ Status ShapeVerifier::HandleConvolution(HloInstruction* convolution) { const Shape expected, ShapeInference::InferConvolveShape( convolution->operand(0)->shape(), convolution->operand(1)->shape(), - convolution->feature_group_count(), convolution->window(), - convolution->convolution_dimension_numbers())); + convolution->feature_group_count(), convolution->batch_group_count(), + convolution->window(), convolution->convolution_dimension_numbers())); return CheckShape(convolution, expected); } @@ -151,13 +167,12 @@ Status ShapeVerifier::HandleFft(HloInstruction* fft) { return CheckShape(fft, expected); } -Status ShapeVerifier::HandleCrossReplicaSum(HloInstruction* crs) { +Status ShapeVerifier::HandleAllReduce(HloInstruction* crs) { std::vector operand_shapes; for (const HloInstruction* operand : crs->operands()) { operand_shapes.push_back(&operand->shape()); } - return CheckShape(crs, - ShapeInference::InferCrossReplicaSumShape(operand_shapes)); + return CheckShape(crs, ShapeInference::InferAllReduceShape(operand_shapes)); } Status ShapeVerifier::HandleAllToAll(HloInstruction* hlo) { @@ -334,7 +349,10 @@ Status ShapeVerifier::HandleConstant(HloInstruction* constant) { Status ShapeVerifier::HandleIota(HloInstruction* instruction) { TF_RETURN_IF_ERROR(CheckOperandCount(instruction, 0)); auto* iota = Cast(instruction); - const int64 rank = ShapeUtil::Rank(iota->shape()); + 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."); } @@ -372,6 +390,14 @@ Status ShapeVerifier::HandleReduce(HloInstruction* reduce) { Status ShapeVerifier::HandleBitcast(HloInstruction* bitcast) { TF_RETURN_IF_ERROR(CheckOperandCount(bitcast, 1)); + // Bitcasts are not allowed to change the element type. + if (bitcast->operand(0)->shape().element_type() != + bitcast->shape().element_type()) { + 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(); } @@ -382,13 +408,12 @@ Status ShapeVerifier::HandleBroadcast(HloInstruction* broadcast) { const Shape& operand_shape = broadcast->operand(0)->shape(); // Check for mixed precision. TF_RET_CHECK(SameElementType(broadcast->shape(), operand_shape)); - TF_RET_CHECK(ShapeUtil::Rank(operand_shape) == - broadcast->dimensions().size()); - for (int64 operand_dimension = 0; - operand_dimension < ShapeUtil::Rank(operand_shape); + TF_RET_CHECK(operand_shape.rank() == broadcast->dimensions().size()); + for (int64 operand_dimension = 0; operand_dimension < operand_shape.rank(); ++operand_dimension) { int64 output_dimension = broadcast->dimensions()[operand_dimension]; - TF_RET_CHECK((output_dimension < ShapeUtil::Rank(broadcast->shape())) && + TF_RET_CHECK((output_dimension < broadcast->shape().rank()) && + output_dimension >= 0 && (broadcast->shape().dimensions(output_dimension) == operand_shape.dimensions(operand_dimension))) << broadcast->ToString() << " operand shape " << operand_shape; @@ -440,6 +465,8 @@ Status ShapeVerifier::HandleFusion(HloInstruction* fusion) { } Status ShapeVerifier::HandleCall(HloInstruction* call) { + TF_RETURN_IF_ERROR( + CheckParameterCount(call, call->to_apply(), call->operand_count())); for (int64 i = 0; i < call->to_apply()->num_parameters(); ++i) { TF_RETURN_IF_ERROR(CheckOperandAndParameter(call, i, call->to_apply(), i)); } @@ -462,7 +489,9 @@ Status ShapeVerifier::HandleCustomCall(HloInstruction* instruction) { const Shape& operand_shape_with_layout = custom_call->operand_shapes_with_layout()[i]; TF_RET_CHECK(ShapeUtil::Compatible(custom_call->operand(i)->shape(), - operand_shape_with_layout)); + operand_shape_with_layout)) + << custom_call->operand(i)->shape().ToString() << " operand " + << operand_shape_with_layout.ToString(); TF_RET_CHECK(LayoutUtil::HasLayout(operand_shape_with_layout)); } } @@ -478,21 +507,23 @@ Status ShapeVerifier::HandleSlice(HloInstruction* slice) { } Status ShapeVerifier::HandleDynamicSlice(HloInstruction* dynamic_slice) { - TF_RETURN_IF_ERROR(CheckOperandCount(dynamic_slice, 2)); - return CheckShape(dynamic_slice, ShapeInference::InferDynamicSliceShape( - dynamic_slice->operand(0)->shape(), - dynamic_slice->operand(1)->shape(), - dynamic_slice->dynamic_slice_sizes())); + return CheckShape( + dynamic_slice, + ShapeInference::InferDynamicSliceShape( + dynamic_slice->operand(0)->shape(), + Cast(dynamic_slice)->index_shapes(), + dynamic_slice->dynamic_slice_sizes())); } Status ShapeVerifier::HandleDynamicUpdateSlice( HloInstruction* dynamic_update_slice) { - 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(), - dynamic_update_slice->operand(2)->shape())); + 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())); } Status ShapeVerifier::HandleTuple(HloInstruction* tuple) { @@ -504,8 +535,7 @@ Status ShapeVerifier::HandleMap(HloInstruction* map) { int64 max_operand_rank = 0; for (const HloInstruction* operand : map->operands()) { operand_shapes.push_back(&operand->shape()); - max_operand_rank = - std::max(max_operand_rank, ShapeUtil::Rank(operand->shape())); + max_operand_rank = std::max(max_operand_rank, operand->shape().rank()); } // TODO(b/65689298) Remove code below once Map is generalized to accept // arbitrary map dimensions. @@ -539,6 +569,10 @@ Status ShapeVerifier::HandleSelectAndScatter(HloInstruction* instruction) { Status ShapeVerifier::HandleWhile(HloInstruction* xla_while) { TF_RETURN_IF_ERROR(CheckOperandCount(xla_while, 1)); + TF_RETURN_IF_ERROR( + CheckParameterCount(xla_while, xla_while->while_body(), 1)); + TF_RETURN_IF_ERROR( + CheckParameterCount(xla_while, xla_while->while_condition(), 1)); TF_RETURN_IF_ERROR( CheckOperandAndParameter(xla_while, 0, xla_while->while_body(), 0)); TF_RETURN_IF_ERROR( @@ -559,6 +593,10 @@ Status ShapeVerifier::HandleWhile(HloInstruction* xla_while) { Status ShapeVerifier::HandleConditional(HloInstruction* conditional) { TF_RETURN_IF_ERROR(CheckOperandCount(conditional, 3)); + TF_RETURN_IF_ERROR( + CheckParameterCount(conditional, conditional->true_computation(), 1)); + TF_RETURN_IF_ERROR( + CheckParameterCount(conditional, conditional->false_computation(), 1)); TF_RETURN_IF_ERROR(CheckOperandAndParameter( conditional, 1, conditional->true_computation(), 0)); TF_RETURN_IF_ERROR(CheckOperandAndParameter( @@ -656,7 +694,7 @@ Status CheckMixedPrecisionOperands(const HloInstruction* instruction) { case HloOpcode::kCall: case HloOpcode::kConditional: case HloOpcode::kConstant: - case HloOpcode::kCrossReplicaSum: + case HloOpcode::kAllReduce: case HloOpcode::kCustomCall: case HloOpcode::kDomain: case HloOpcode::kFusion: @@ -667,7 +705,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: @@ -726,7 +763,19 @@ Status ShapeVerifier::HandleAfterAll(HloInstruction* token) { for (const HloInstruction* operand : token->operands()) { operand_shapes.push_back(&operand->shape()); } - return CheckShape(token, ShapeInference::InferAfterAllShape(operand_shapes)); + return CheckShape(token, ShapeUtil::MakeTokenShape()); +} + +Status ShapeVerifier::HandleAddDependency(HloInstruction* add_dependency) { + TF_RETURN_IF_ERROR(CheckOperandCount(add_dependency, 2)); + TF_RETURN_IF_ERROR(CheckIsTokenOperand(add_dependency, 1)); + return CheckShape(add_dependency, add_dependency->operand(0)->shape()); +} + +Status ShapeVerifier::HandleGetDimensionSize(HloInstruction* get_size) { + return CheckShape(get_size, + ShapeInference::InferGetDimensionSizeShape( + get_size->operand(0)->shape(), get_size->dimension())); } Status ShapeVerifier::CheckShape(const HloInstruction* instruction, @@ -828,6 +877,50 @@ Status ShapeVerifier::CheckVariadicShape(const HloInstruction* instruction) { instruction->opcode(), instruction->operands())); } +Status ShapeVerifier::VerifyEntryComputationLayout(const HloModule& module) { + const HloComputation* computation = module.entry_computation(); + const auto& layout = module.entry_computation_layout(); + const ShapeLayout& result_layout = layout.result_layout(); + + TF_RETURN_IF_ERROR( + ShapeUtil::ValidateShapeWithOptionalLayout(result_layout.shape())); + + TF_RETURN_IF_ERROR(VerifyNotSparse(result_layout.shape())); + + if (!ShapeUtil::Compatible(computation->root_instruction()->shape(), + result_layout.shape())) { + return InternalError( + "Shape of the root instruction of entry computation (%s) should be " + "compatible to one specified in module's entry computation layout (%s)", + ShapeUtil::HumanString(computation->root_instruction()->shape()), + ShapeUtil::HumanString(result_layout.shape())); + } + + if (computation->num_parameters() != layout.parameter_count()) { + return InternalError( + "Number of parameters in entry computation layout (%d) must be same " + "as number of parameters of entry computation computation (%d)", + layout.parameter_count(), computation->num_parameters()); + } + + for (int i = 0; i < computation->num_parameters(); ++i) { + const HloInstruction* parameter = computation->parameter_instruction(i); + TF_RETURN_IF_ERROR( + ShapeUtil::ValidateShapeWithOptionalLayout(layout.parameter_shape(i))); + TF_RETURN_IF_ERROR(VerifyNotSparse(layout.parameter_shape(i))); + if (!ShapeUtil::Compatible(parameter->shape(), layout.parameter_shape(i))) { + return InternalError( + "Shape of the entry computation parameter %d is %s should be " + "compatible to the one specified in module's entry computation " + "layout %s", + i, ShapeUtil::HumanString(parameter->shape()), + ShapeUtil::HumanString(layout.parameter_shape(i))); + } + } + + return Status::OK(); +} + string ComputationsToString(absl::Span computations) { return absl::StrJoin(computations, ",", [](string* s, const HloComputation* computation) { @@ -899,7 +992,7 @@ bool ShapeContainsToken(const Shape& shape) { bool contains_token = false; ShapeUtil::ForEachSubshape( shape, [&contains_token](const Shape& subshape, const ShapeIndex&) { - if (ShapeUtil::IsToken(subshape)) { + if (subshape.IsToken()) { contains_token = true; } }); @@ -923,52 +1016,6 @@ Status VerifyEntryAndExitShapes(const HloModule& module) { return Status::OK(); } -// Verifies that entry computation layout matches characteristics of -// entry computation. -Status CheckEntryComputationLayout(const HloModule& module) { - const HloComputation* computation = module.entry_computation(); - const auto& layout = module.entry_computation_layout(); - const ShapeLayout& result_layout = layout.result_layout(); - - TF_RETURN_IF_ERROR( - ShapeUtil::ValidateShapeWithOptionalLayout(result_layout.shape())); - - TF_RETURN_IF_ERROR(VerifyNotSparse(result_layout.shape())); - - if (!ShapeUtil::Compatible(computation->root_instruction()->shape(), - result_layout.shape())) { - return InternalError( - "Shape of the root instruction of entry computation (%s) should be " - "compatible to one specified in module's entry computation layout (%s)", - ShapeUtil::HumanString(computation->root_instruction()->shape()), - ShapeUtil::HumanString(result_layout.shape())); - } - - if (computation->num_parameters() != layout.parameter_count()) { - return InternalError( - "Number of parameters in entry computation layout (%d) must be same " - "as number of parameters of entry computation computation (%d)", - layout.parameter_count(), computation->num_parameters()); - } - - for (int i = 0; i < computation->num_parameters(); ++i) { - const HloInstruction* parameter = computation->parameter_instruction(i); - TF_RETURN_IF_ERROR( - ShapeUtil::ValidateShapeWithOptionalLayout(layout.parameter_shape(i))); - TF_RETURN_IF_ERROR(VerifyNotSparse(layout.parameter_shape(i))); - if (!ShapeUtil::Compatible(parameter->shape(), layout.parameter_shape(i))) { - return InternalError( - "Shape of the entry computation parameter %d is %s should be " - "compatible to the one specified in module's entry computation " - "layout %s", - i, ShapeUtil::HumanString(parameter->shape()), - ShapeUtil::HumanString(layout.parameter_shape(i))); - } - } - - return Status::OK(); -} - // Checks if the given two instructions share the same channel id. Status CheckSameChannel(const HloInstruction* instr1, const HloInstruction* instr2) { @@ -1233,11 +1280,11 @@ class InstructionVerifier : public DfsHloVisitorWithDefault { // op. See https://groups.google.com/forum/#!topic/xla-dev/9LqijHmTt_I // or ComputationLowerer::Visit() TF_RET_CHECK(broadcast->dimensions().size() == - ShapeUtil::Rank(broadcast->operand(0)->shape())) + broadcast->operand(0)->shape().rank()) << "Broadcast HLO (" << broadcast->ToShortString() << ") has invalid number of dimensions: " << broadcast->dimensions().size() - << " != " << ShapeUtil::Rank(broadcast->operand(0)->shape()); + << " != " << broadcast->operand(0)->shape().rank(); return Status::OK(); } @@ -1287,7 +1334,7 @@ class InstructionVerifier : public DfsHloVisitorWithDefault { } Status HandleGetTupleElement(HloInstruction* gte) override { - TF_RET_CHECK(ShapeUtil::IsTuple(gte->operand(0)->shape())); + TF_RET_CHECK(gte->operand(0)->shape().IsTuple()); return Status::OK(); } @@ -1307,6 +1354,15 @@ class InstructionVerifier : public DfsHloVisitorWithDefault { return Status::OK(); } + Status HandleAllReduce(HloInstruction* crs) override { + if (crs->all_reduce_id().has_value()) { + TF_RET_CHECK(crs->all_reduce_id().value() > 0) + << "All reduce id must be greater than 0 for " + << crs->ToShortString(); + } + return Status::OK(); + } + Status Preprocess(HloInstruction* instruction) override { auto previous = instructions_by_name_.find(instruction->name()); TF_RET_CHECK(previous == instructions_by_name_.end()) @@ -1329,13 +1385,12 @@ class InstructionVerifier : public DfsHloVisitorWithDefault { for (HloInstruction* operand : instruction->operands()) { const Shape& operand_shape = operand->shape(); if (LayoutUtil::IsDenseArray(operand_shape) && - ShapeUtil::Rank(operand_shape) == ShapeUtil::Rank(result_shape)) { + operand_shape.rank() == result_shape.rank()) { const Layout& operand_layout = operand_shape.layout(); TF_RET_CHECK(LayoutUtil::Equal(result_layout, operand_layout)) << "Instruction shouldn't change layouts " - << instruction->ToString() << " From " - << ShapeUtil::HumanString(result_shape) << " To " - << ShapeUtil::HumanString(operand_shape); + << instruction->ToString() << " From " << result_shape << " To " + << operand_shape; } } } @@ -1363,8 +1418,9 @@ StatusOr HloVerifier::Run(HloModule* module) { TF_RETURN_IF_ERROR(VerifyHloStructure(module)); TF_RETURN_IF_ERROR(VerifySendsAndRecvs(*module)); + std::unique_ptr shape_verifier = + target_metadata_->GetVerifier(); for (auto* computation : module->computations()) { - std::unique_ptr shape_verifier = shape_verifier_factory_(); TF_RETURN_IF_ERROR(computation->Accept(shape_verifier.get())); InstructionVerifier instruction_verifier( @@ -1372,7 +1428,7 @@ StatusOr HloVerifier::Run(HloModule* module) { TF_RETURN_IF_ERROR(computation->Accept(&instruction_verifier)); } - TF_RETURN_IF_ERROR(CheckEntryComputationLayout(*module)); + TF_RETURN_IF_ERROR(shape_verifier->VerifyEntryComputationLayout(*module)); TF_RETURN_IF_ERROR(VerifyEntryAndExitShapes(*module)); // If the module has a schedule, it must be valid. @@ -1380,7 +1436,12 @@ StatusOr HloVerifier::Run(HloModule* module) { TF_RETURN_IF_ERROR(module->schedule().Verify()); } - TF_RETURN_IF_ERROR(module->input_output_alias_config().Verify(*module)); + TF_RETURN_IF_ERROR(module->input_output_alias_config().Verify( + *module, [this](const Shape& shape) { + return target_metadata_->ShapeSize(shape); + })); + + TF_RETURN_IF_ERROR(module->dynamic_parameter_binding().Verify(*module)); return false; } diff --git a/tensorflow/compiler/xla/service/hlo_verifier.h b/tensorflow/compiler/xla/service/hlo_verifier.h index 5e65d7687755c2a664d20495dc4f2c3b2499f487..479905b317d5639ff2cebc4d1044e21b527693f6 100644 --- a/tensorflow/compiler/xla/service/hlo_verifier.h +++ b/tensorflow/compiler/xla/service/hlo_verifier.h @@ -16,6 +16,7 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_HLO_VERIFIER_H_ #define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_VERIFIER_H_ +#include #include "tensorflow/compiler/xla/service/hlo_pass_interface.h" #include "absl/memory/memory.h" @@ -28,10 +29,14 @@ namespace xla { // TODO(b/26024837): Check output shape for all instruction types. class ShapeVerifier : public DfsHloVisitor { public: - explicit ShapeVerifier(bool layout_sensitive, bool allow_mixed_precision) + ShapeVerifier(bool layout_sensitive, bool allow_mixed_precision) : layout_sensitive_(layout_sensitive), allow_mixed_precision_(allow_mixed_precision) {} + // Verifies that entry computation layout matches parameters and root shape of + // the module's entry computation. + virtual Status VerifyEntryComputationLayout(const HloModule& module); + Status Preprocess(HloInstruction* hlo) override; Status HandleElementwiseUnary(HloInstruction* hlo) override; @@ -47,7 +52,7 @@ class ShapeVerifier : public DfsHloVisitor { Status HandleDot(HloInstruction* dot) override; Status HandleConvolution(HloInstruction* convolution) override; Status HandleFft(HloInstruction* fft) override; - Status HandleCrossReplicaSum(HloInstruction* crs) override; + Status HandleAllReduce(HloInstruction* crs) override; Status HandleAllToAll(HloInstruction* hlo) override; Status HandleCollectivePermute(HloInstruction* hlo) override; Status HandleReducePrecision(HloInstruction* reduce_precision) override; @@ -89,6 +94,8 @@ class ShapeVerifier : public DfsHloVisitor { Status HandleGather(HloInstruction* gather) override; Status HandleScatter(HloInstruction* scatter) override; Status HandleAfterAll(HloInstruction* token) override; + Status HandleGetDimensionSize(HloInstruction* get_size) override; + Status HandleAddDependency(HloInstruction* add_dependency) override; Status FinishVisit(HloInstruction*) override { return Status::OK(); } @@ -158,27 +165,75 @@ class ShapeVerifier : public DfsHloVisitor { bool allow_mixed_precision_; }; +// 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. + int64 ShapeSize(const Shape& shape) const { + return shape_size_function_(shape); + } + + virtual std::unique_ptr GetVerifier() const = 0; + + TargetVerifierMetadata() {} + virtual ~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, + std::function shape_size_function) + : TargetVerifierMetadata(shape_size_function), + layout_sensitive_(layout_sensitive), + allow_mixed_precision_(allow_mixed_precision) {} + + // 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 + // the verifier. + std::unique_ptr GetVerifier() const override { + return absl::make_unique(layout_sensitive_, + allow_mixed_precision_); + } + + private: + bool layout_sensitive_; + bool allow_mixed_precision_; +}; + // HLO pass that verifies invariants of HLO instructions for each computation in // the module. class HloVerifier : public HloModulePass { public: - using ShapeVerifierFactory = std::function()>; - - explicit HloVerifier(bool layout_sensitive, bool allow_mixed_precision, - std::function - instruction_can_change_layout_func = {}) - : shape_verifier_factory_([layout_sensitive, allow_mixed_precision] { - return absl::make_unique(layout_sensitive, - allow_mixed_precision); - }), + 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, shape_size_func)), instruction_can_change_layout_func_( std::move(instruction_can_change_layout_func)) { CHECK(instruction_can_change_layout_func_ == nullptr || layout_sensitive); } - // Uses custom shape verification. - explicit HloVerifier(ShapeVerifierFactory shape_verifier_factory) - : shape_verifier_factory_(std::move(shape_verifier_factory)) {} + // Uses custom target metadata + explicit HloVerifier(std::unique_ptr target_metadata) + : target_metadata_(std::move(target_metadata)) {} ~HloVerifier() override = default; absl::string_view name() const override { return "verifier"; } @@ -187,11 +242,7 @@ class HloVerifier : public HloModulePass { StatusOr Run(HloModule* module) override; private: - // Creates a ShapeVerifier that checks that shapes match inferred - // expectations. This is a factory function because ShapeVerifier, - // being a DfsHloVisitor, is stateful. We want a clean object - // for each run of the verifier. - ShapeVerifierFactory shape_verifier_factory_; + std::unique_ptr target_metadata_; // Determines whether an instruction can change layouts. std::function diff --git a/tensorflow/compiler/xla/service/hlo_verifier_test.cc b/tensorflow/compiler/xla/service/hlo_verifier_test.cc index afe01e5487c3225815e01343d86e9fe894c2cde8..4f69bd155b8713041ba539098808125956e86259 100644 --- a/tensorflow/compiler/xla/service/hlo_verifier_test.cc +++ b/tensorflow/compiler/xla/service/hlo_verifier_test.cc @@ -20,6 +20,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_config.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" #include "tensorflow/compiler/xla/service/hlo_parser.h" #include "tensorflow/compiler/xla/service/layout_assignment.h" @@ -27,6 +28,7 @@ limitations under the License. #include "tensorflow/compiler/xla/test.h" #include "tensorflow/compiler/xla/tests/hlo_test_base.h" #include "tensorflow/compiler/xla/types.h" +#include "tensorflow/compiler/xla/xla.pb.h" #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/core/status_test_util.h" @@ -35,7 +37,11 @@ namespace { using ::testing::HasSubstr; -// This class cannot be converted to use HloVerifiedTestBase. It explicitly +std::unique_ptr CreateUnverifiedModule() { + return absl::make_unique("module", HloModuleConfig()); +} + +// This class cannot be converted to use HloTestBase. It explicitly // uses HloTestBase to create and test malformed HLOs. class HloVerifierTest : public HloTestBase { public: @@ -66,7 +72,7 @@ TEST_F(HloVerifierTest, NullInstructionParent) { HloInstruction::CreateParameter(0, scalar_shape, "param")); HloInstruction* negate = builder.AddInstruction( HloInstruction::CreateUnary(scalar_shape, HloOpcode::kNegate, param)); - auto module = CreateNewModule(); + auto module = CreateUnverifiedModule(); module->AddEntryComputation(builder.Build()); TF_ASSERT_OK(verifier().Run(module.get()).status()); @@ -85,7 +91,7 @@ TEST_F(HloVerifierTest, NullComputationParent) { HloInstruction::CreateParameter(0, scalar_shape, "param")); builder.AddInstruction( HloInstruction::CreateUnary(scalar_shape, HloOpcode::kNegate, param)); - auto module = CreateNewModule(); + auto module = CreateUnverifiedModule(); HloComputation* computation = module->AddEntryComputation(builder.Build()); TF_ASSERT_OK(verifier().Run(module.get()).status()); @@ -104,7 +110,7 @@ TEST_F(HloVerifierTest, DifferentOperandParents) { HloInstruction::CreateParameter(0, scalar_shape, "param")); HloInstruction* negate = builder.AddInstruction( HloInstruction::CreateUnary(scalar_shape, HloOpcode::kNegate, param)); - auto module = CreateNewModule(); + auto module = CreateUnverifiedModule(); module->AddEntryComputation(builder.Build()); HloComputation::Builder emb_builder(TestName()); @@ -138,7 +144,7 @@ TEST_F(HloVerifierTest, ResetsShapeVerifierState) { builder.AddInstruction( HloInstruction::CreateBinary(s2, HloOpcode::kMultiply, add, add)); - auto module = CreateNewModule(); + auto module = CreateUnverifiedModule(); module->AddEntryComputation(builder.Build()); // Run the verifier twice. It should fail both times, because it shouldn't @@ -303,7 +309,7 @@ TEST_F(HloVerifierTest, NegativeInteriorPaddingNotAllowed) { HloInstruction::CreateConstant(LiteralUtil::Zero(F32))), padding_config)); - auto module = CreateNewModule(); + auto module = CreateUnverifiedModule(); module->AddEntryComputation(builder.Build()); auto status = verifier().Run(module.get()).status(); @@ -327,7 +333,7 @@ TEST_F(HloVerifierTest, PadNegativeInteriorDilationNotAllowed) { HloInstruction::CreateConstant(LiteralUtil::Zero(F32).Clone())), padding_config)); - auto module = CreateNewModule(); + auto module = CreateUnverifiedModule(); module->AddEntryComputation(builder.Build()); EXPECT_THAT(verifier().Run(module.get()).status().error_message(), @@ -382,6 +388,55 @@ TEST_F(HloVerifierTest, AddWithLayoutChange) { ASSERT_TRUE(status.ok()); } +TEST_F(HloVerifierTest, ScalarIndexDynamicSlice) { + const char* const kScalarIndexDynamicSlice = R"( + HloModule DynamicSlice_module + + ENTRY %DynamicSlice.v5 (original_parameter: s32[2,2,258], start_index: s32[]) -> s32[2,2,258] { + %original_parameter = s32[2,2,258] parameter(0) + %constant = s32[] constant(0) + %start_index = s32[] parameter(1) + ROOT %dynamic-slice = s32[2,2,258] dynamic-slice(s32[2,2,258] %original_parameter, s32[] %constant, s32[] %constant, s32[] %start_index), dynamic_slice_sizes={2,2,258} + } + )"; + + HloModuleConfig config; + DebugOptions debug_options = config.debug_options(); + debug_options.set_xla_allow_scalar_index_dynamic_ops(true); + config.set_debug_options(debug_options); + + TF_ASSERT_OK_AND_ASSIGN(auto module, + ParseHloString(kScalarIndexDynamicSlice, config)); + auto status = verifier().Run(module.get()).status(); + ASSERT_TRUE(status.ok()); +} + +TEST_F(HloVerifierTest, ScalarIndexDynamicUpdateSlice) { + const char* const kScalarIndexDynamicSlice = R"( + HloModule DynamicUpdateSlice_module + + ENTRY %DynamicUpdateSlice.v4 (input: s32[1,1,25,1], update: s32[1,1,2,1], start_index.0: s32[], start_index.1: s32[], start_index.2: s32[], start_index.3: s32[]) -> s32[1,1,25,1] { + %input = s32[1,1,25,1]{3,2,1,0} parameter(0) + %update = s32[1,1,2,1]{3,2,1,0} parameter(1) + %start_index.0 = s32[] parameter(2) + %start_index.1 = s32[] parameter(3) + %start_index.2 = s32[] parameter(4) + %start_index.3 = s32[] parameter(5) + ROOT %dynamic-update-slice = s32[1,1,25,1]{3,2,1,0} dynamic-update-slice(s32[1,1,25,1]{3,2,1,0} %input, s32[1,1,2,1]{3,2,1,0} %update, s32[] %start_index.0, s32[] %start_index.1, s32[] %start_index.2, s32[] %start_index.3) + } + )"; + + HloModuleConfig config; + DebugOptions debug_options = config.debug_options(); + debug_options.set_xla_allow_scalar_index_dynamic_ops(true); + config.set_debug_options(debug_options); + + TF_ASSERT_OK_AND_ASSIGN(auto module, + ParseHloString(kScalarIndexDynamicSlice, config)); + auto status = verifier().Run(module.get()).status(); + ASSERT_TRUE(status.ok()); +} + TEST_F(HloVerifierTestLayoutSensitive, AddWithLayoutChangeNotAllowed) { TF_ASSERT_OK_AND_ASSIGN(auto module, ParseHloString(kAddWithLayoutChangeHlo)); auto status = verifier().Run(module.get()).status(); @@ -395,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} } )"; @@ -425,5 +481,76 @@ 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")); +} + } // 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 e76b93107c923b41666f6b0a388dda143a8cb50a..88fc62bd1e2a7830b3f61738a8642308ef4225a7 100644 --- a/tensorflow/compiler/xla/service/human_readable_profile_builder.cc +++ b/tensorflow/compiler/xla/service/human_readable_profile_builder.cc @@ -90,20 +90,29 @@ string HumanReadableProfileBuilder::ToString() const { op.optimal_seconds < 0 ? "" : StrFormat("(%12.1f optimal)", op.optimal_seconds * 1e6), - op.flop_count <= 0 ? "" : HumanReadableNumFlops(op.flop_count, nsecs), - op.transcendental_count <= 0 - ? "" - : HumanReadableNumTranscendentalOps(op.transcendental_count, nsecs), + op.flop_count > 0 && nsecs > 0 + ? HumanReadableNumFlops(op.flop_count, nsecs) + : "", + op.transcendental_count > 0 && nsecs > 0 + ? HumanReadableNumTranscendentalOps(op.transcendental_count, nsecs) + : "", bytes_per_sec, bytes_per_cycle, op.name); }; - float optimal_seconds_sum = 0.0; + double optimal_seconds_sum = 0; int64 total_flops = 0.; int64 total_transcendentals = 0.; int64 total_bytes = 0; for (const auto& op : op_infos_) { if (op.optimal_seconds > 0) { - optimal_seconds_sum += op.optimal_seconds; + // An op can run faster than the estimated optimum. For example, we might + // estimate a fusion's speed by looking at the size of its operands and + // result, but perhaps the fusion doesn't read the entirety of all of its + // inputs. For the purposes of summing the instructions' optimal speeds, + // we treat the "optimum" as the smallest of either the estimated optimum + // and the actual speed. + optimal_seconds_sum += + std::min(double{op.optimal_seconds}, CyclesToSeconds(op.cycles)); } total_flops += std::max(op.flop_count, int64{0}); total_transcendentals += std::max(op.transcendental_count, int64{0}); @@ -112,15 +121,16 @@ string HumanReadableProfileBuilder::ToString() const { VLOG(1) << "Total floating point ops: " << total_flops; - print_op({"[total]", "[total]", /*category=*/"", total_cycles_, total_flops, - total_transcendentals, total_bytes, optimal_seconds_sum}, + print_op({is_entry_computation_ ? "[total] [entry]" : "[total]", "[total]", + /*category=*/"", total_cycles_, total_flops, total_transcendentals, + total_bytes, static_cast(optimal_seconds_sum)}, /*is_total=*/true); // 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); } @@ -154,8 +164,10 @@ string HumanReadableProfileBuilder::ToString() const { entry.text = op.name; entry.short_text = op.short_name; entry.category_text = op.category; - entry.metric = - CyclesToMicroseconds(op.cycles) - op.optimal_seconds * 1e6; + // Ignore ops that run faster than the estimated optimal here, as we do + // when calculating optimal_seconds_sum. + entry.metric = std::max( + 0., CyclesToMicroseconds(op.cycles) - op.optimal_seconds * 1e6); total_discrepancy_in_microseconds += entry.metric; table.AddEntry(std::move(entry)); } diff --git a/tensorflow/compiler/xla/service/human_readable_profile_builder.h b/tensorflow/compiler/xla/service/human_readable_profile_builder.h index 925111fa1f1e48650b0089f402d92e431043eabe..d4e5cbbe27418ddf3c81ebe00bc8aa979d3c2d5e 100644 --- a/tensorflow/compiler/xla/service/human_readable_profile_builder.h +++ b/tensorflow/compiler/xla/service/human_readable_profile_builder.h @@ -30,9 +30,11 @@ namespace xla { class HumanReadableProfileBuilder { public: explicit HumanReadableProfileBuilder(absl::string_view computation_name, + bool is_entry_computation, int64 total_cycles, double clock_rate_ghz) : computation_name_(computation_name), + is_entry_computation_(is_entry_computation), total_cycles_(total_cycles), clock_rate_ghz_(clock_rate_ghz) { CHECK_GE(clock_rate_ghz, 1e-9); @@ -75,6 +77,7 @@ class HumanReadableProfileBuilder { } string computation_name_; + bool is_entry_computation_; int64 total_cycles_; double clock_rate_ghz_; std::vector op_infos_; diff --git a/tensorflow/compiler/xla/service/implicit_broadcast_remover_test.cc b/tensorflow/compiler/xla/service/implicit_broadcast_remover_test.cc index f85d31d5225b8012b68f851b2bfec219d736ba0d..cf6cf897fe11eda01ba6b22119bba34ac2bef8fe 100644 --- a/tensorflow/compiler/xla/service/implicit_broadcast_remover_test.cc +++ b/tensorflow/compiler/xla/service/implicit_broadcast_remover_test.cc @@ -18,19 +18,20 @@ limitations under the License. #include "tensorflow/compiler/xla/literal.h" #include "tensorflow/compiler/xla/service/hlo_matchers.h" #include "tensorflow/compiler/xla/shape_util.h" -#include "tensorflow/compiler/xla/tests/hlo_verified_test_base.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" namespace op = xla::testing::opcode_matchers; namespace xla { namespace { -class ImplicitBroadcastRemoverTest : public HloVerifiedTestBase { +class ImplicitBroadcastRemoverTest : public HloTestBase { protected: ImplicitBroadcastRemover remover_; }; TEST_F(ImplicitBroadcastRemoverTest, NoImplicitBroadcast) { + auto m = CreateNewVerifiedModule(); auto builder = HloComputation::Builder(TestName()); const Shape shape = ShapeUtil::MakeShape(F32, {2, 4}); @@ -41,15 +42,16 @@ TEST_F(ImplicitBroadcastRemoverTest, NoImplicitBroadcast) { builder.AddInstruction( HloInstruction::CreateBinary(shape, HloOpcode::kAdd, param0, param1)); - HloComputation* computation = module().AddEntryComputation(builder.Build()); + HloComputation* computation = m->AddEntryComputation(builder.Build()); - EXPECT_FALSE(remover_.Run(&module()).ValueOrDie()); + EXPECT_FALSE(remover_.Run(m.get()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::Add(op::Parameter(), op::Parameter())); } TEST_F(ImplicitBroadcastRemoverTest, ScalarBroadcast) { + auto m = CreateNewVerifiedModule(); auto builder = HloComputation::Builder(TestName()); const Shape shape = ShapeUtil::MakeShape(F32, {2, 4}); @@ -60,13 +62,13 @@ TEST_F(ImplicitBroadcastRemoverTest, ScalarBroadcast) { builder.AddInstruction( HloInstruction::CreateBinary(shape, HloOpcode::kPower, param0, param1)); - HloComputation* computation = module().AddEntryComputation(builder.Build()); + HloComputation* computation = m->AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_FALSE(ShapeUtil::Compatible(root->shape(), root->operand(0)->shape())); EXPECT_TRUE(ShapeUtil::Compatible(root->shape(), root->operand(1)->shape())); - EXPECT_TRUE(remover_.Run(&module()).ValueOrDie()); + EXPECT_TRUE(remover_.Run(m.get()).ValueOrDie()); root = computation->root_instruction(); EXPECT_THAT(root, op::Power(op::Broadcast(op::Parameter()), op::Parameter())); @@ -76,6 +78,7 @@ TEST_F(ImplicitBroadcastRemoverTest, ScalarBroadcast) { } TEST_F(ImplicitBroadcastRemoverTest, DegenerateDimensionBroadcast) { + auto m = CreateNewVerifiedModule(); auto builder = HloComputation::Builder(TestName()); const Shape shape = ShapeUtil::MakeShape(F32, {2, 4, 6}); @@ -86,9 +89,9 @@ TEST_F(ImplicitBroadcastRemoverTest, DegenerateDimensionBroadcast) { builder.AddInstruction(HloInstruction::CreateBinary( shape, HloOpcode::kSubtract, param0, param1)); - HloComputation* computation = module().AddEntryComputation(builder.Build()); + HloComputation* computation = m->AddEntryComputation(builder.Build()); - EXPECT_TRUE(remover_.Run(&module()).ValueOrDie()); + EXPECT_TRUE(remover_.Run(m.get()).ValueOrDie()); HloInstruction* root = computation->root_instruction(); EXPECT_THAT(root, op::Subtract(op::Parameter(), @@ -98,6 +101,7 @@ TEST_F(ImplicitBroadcastRemoverTest, DegenerateDimensionBroadcast) { } TEST_F(ImplicitBroadcastRemoverTest, ScalarBroadcastToDegenerateDimensions) { + auto m = CreateNewVerifiedModule(); auto builder = HloComputation::Builder(TestName()); const Shape shape = ShapeUtil::MakeShape(F32, {1, 4, 1}); @@ -108,9 +112,9 @@ TEST_F(ImplicitBroadcastRemoverTest, ScalarBroadcastToDegenerateDimensions) { builder.AddInstruction(HloInstruction::CreateBinary( shape, HloOpcode::kSubtract, param0, param1)); - HloComputation* computation = module().AddEntryComputation(builder.Build()); + HloComputation* computation = m->AddEntryComputation(builder.Build()); - EXPECT_TRUE(remover_.Run(&module()).ValueOrDie()); + EXPECT_TRUE(remover_.Run(m.get()).ValueOrDie()); HloInstruction* root = computation->root_instruction(); EXPECT_THAT(root, @@ -120,6 +124,7 @@ TEST_F(ImplicitBroadcastRemoverTest, ScalarBroadcastToDegenerateDimensions) { } TEST_F(ImplicitBroadcastRemoverTest, TernaryDegenerateDimensionBroadcast) { + auto m = CreateNewVerifiedModule(); auto builder = HloComputation::Builder(TestName()); const Shape shape = ShapeUtil::MakeShape(F32, {2, 4, 6, 8}); @@ -132,9 +137,9 @@ TEST_F(ImplicitBroadcastRemoverTest, TernaryDegenerateDimensionBroadcast) { builder.AddInstruction(HloInstruction::CreateTernary(shape, HloOpcode::kClamp, param0, param1, param2)); - HloComputation* computation = module().AddEntryComputation(builder.Build()); + HloComputation* computation = m->AddEntryComputation(builder.Build()); - EXPECT_TRUE(remover_.Run(&module()).ValueOrDie()); + EXPECT_TRUE(remover_.Run(m.get()).ValueOrDie()); HloInstruction* root = computation->root_instruction(); EXPECT_THAT(root, op::Clamp(op::Broadcast(op::Reshape(op::Parameter())), @@ -147,6 +152,7 @@ TEST_F(ImplicitBroadcastRemoverTest, TernaryDegenerateDimensionBroadcast) { TEST_F(ImplicitBroadcastRemoverTest, TernaryScalarAndDegenerateDimensionBroadcast) { + auto m = CreateNewVerifiedModule(); auto builder = HloComputation::Builder(TestName()); const Shape shape = ShapeUtil::MakeShape(F32, {2, 4, 6}); @@ -159,9 +165,9 @@ TEST_F(ImplicitBroadcastRemoverTest, builder.AddInstruction(HloInstruction::CreateTernary(shape, HloOpcode::kClamp, param0, param1, param2)); - HloComputation* computation = module().AddEntryComputation(builder.Build()); + HloComputation* computation = m->AddEntryComputation(builder.Build()); - EXPECT_TRUE(remover_.Run(&module()).ValueOrDie()); + EXPECT_TRUE(remover_.Run(m.get()).ValueOrDie()); HloInstruction* root = computation->root_instruction(); EXPECT_THAT(root, op::Clamp(op::Broadcast(op::Parameter()), diff --git a/tensorflow/compiler/xla/service/indexed_array_analysis.cc b/tensorflow/compiler/xla/service/indexed_array_analysis.cc index 1ebb3319779c00fd4afe90606bf336e16349429d..76bf48870d55e82497ba5f63e9e2e2a322cb330e 100644 --- a/tensorflow/compiler/xla/service/indexed_array_analysis.cc +++ b/tensorflow/compiler/xla/service/indexed_array_analysis.cc @@ -103,7 +103,7 @@ Status IndexedArrayAnalysis::TraverseAndPopulateCache( do { const HloInstruction* instr = stack.back(); - if (cache_.count(instr)) { + if (cache_.contains(instr)) { stack.pop_back(); continue; } @@ -111,9 +111,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; } @@ -1002,7 +1002,7 @@ bool CanFoldDotIntoIndexedArray( absl::Span contracting_dims, absl::Span batch_dims) { absl::optional non_contracting_non_batch_dim = - GetOnlyNonContractingNonBatchDim(ShapeUtil::Rank(indexed_array->shape()), + GetOnlyNonContractingNonBatchDim(indexed_array->shape().rank(), contracting_dims, batch_dims); if (!non_contracting_non_batch_dim.has_value()) { VLOG(3) << tag << ": multiple or no non-contracting non-batch dimensions"; @@ -1015,7 +1015,7 @@ bool CanFoldDotIntoIndexedArray( return false; } - int64 indexed_array_rank = ShapeUtil::Rank(indexed_array->shape()); + int64 indexed_array_rank = indexed_array->shape().rank(); if (indexed_array->source_dim() < (indexed_array_rank - 2)) { // This restriction can be lifted by inserting reshape nodes. VLOG(3) << tag @@ -1043,7 +1043,7 @@ IndexedArrayAnalysis::ComputeArrayForDotWithIndexedLhs( return nullptr; } - int64 lhs_rank = ShapeUtil::Rank(lhs->shape()); + int64 lhs_rank = lhs->shape().rank(); DotDimensionNumbers new_dim_numbers = dim_numbers; new_dim_numbers.set_lhs_contracting_dimensions( 0, lhs->source_dim() == (lhs_rank - 1) ? (lhs_rank - 2) : (lhs_rank - 1)); @@ -1078,7 +1078,7 @@ IndexedArrayAnalysis::ComputeArrayForDotWithIndexedRhs( return nullptr; } - int64 rhs_rank = ShapeUtil::Rank(rhs->shape()); + int64 rhs_rank = rhs->shape().rank(); DotDimensionNumbers new_dim_numbers = dim_numbers; new_dim_numbers.set_rhs_contracting_dimensions( diff --git a/tensorflow/compiler/xla/service/indexed_array_analysis_test.cc b/tensorflow/compiler/xla/service/indexed_array_analysis_test.cc index 2d03aebc1aca4c55cca588072233b7a18e70a306..62107b5a88d4e37552fa5a6384700a9291a9c655 100644 --- a/tensorflow/compiler/xla/service/indexed_array_analysis_test.cc +++ b/tensorflow/compiler/xla/service/indexed_array_analysis_test.cc @@ -13,15 +13,14 @@ 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 "tensorflow/compiler/xla/tests/hlo_verified_test_base.h" +#include "absl/strings/ascii.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" #include "tensorflow/compiler/xla/tests/test_utils.h" namespace xla { namespace { -class IndexedArrayAnalysisTest : public HloVerifiedTestBase { +class IndexedArrayAnalysisTest : public HloTestBase { protected: void AssertArrayForRootExpressionIs(const string& hlo_text, const string& root_expression) { @@ -43,7 +42,7 @@ class IndexedArrayAnalysisTest : public HloVerifiedTestBase { 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(' '); @@ -61,12 +60,12 @@ class IndexedArrayAnalysisTest : public HloVerifiedTestBase { const string& root_expression, bool print_constants) { IndexedArrayAnalysis indexed_tensor_analysis; - ParseAndVerifyModule(hlo_text); + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr m, + ParseAndReturnVerifiedModule(hlo_text)); - TF_ASSERT_OK_AND_ASSIGN( - IndexedArrayAnalysis::Array* const array_result, - indexed_tensor_analysis.GetArrayFor( - module().entry_computation()->root_instruction())); + TF_ASSERT_OK_AND_ASSIGN(IndexedArrayAnalysis::Array* const array_result, + indexed_tensor_analysis.GetArrayFor( + m->entry_computation()->root_instruction())); string string_result = CanonicalizeWhitespace( indexed_tensor_analysis.ToString(array_result, print_constants)); LOG(INFO) << string_result; @@ -99,7 +98,7 @@ TEST_F(IndexedArrayAnalysisTest, SimpleOneToOneConstantGather) { HloModule SimpleGather ENTRY main { - operand = s32[3,3] constant(s32[3,3]{{1,2,3},{1,2,3},{1,2,3}}) + operand = s32[3,3] constant({{1,2,3},{1,2,3},{1,2,3}}) indices = s32[5] parameter(0) ROOT gather = s32[5,3] gather(operand, indices), offset_dims={1}, @@ -119,7 +118,7 @@ TEST_F(IndexedArrayAnalysisTest, GatherIsNotScalarIndexed0) { HloModule SimpleGather ENTRY main { - operand = s32[3,3] constant(s32[3,3]{{1,2,3},{1,2,3},{1,2,3}}) + operand = s32[3,3] constant({{1,2,3},{1,2,3},{1,2,3}}) indices = s32[5,2] parameter(0) ROOT gather = s32[5] gather(operand, indices), offset_dims={}, @@ -195,7 +194,7 @@ TEST_F(IndexedArrayAnalysisTest, GatherOfGather_OneToOne) { HloModule SimpleGather ENTRY main { - operand = s32[3,3] constant(s32[3,3]{{1,2,3},{1,2,3},{1,2,3}}) + operand = s32[3,3] constant({{1,2,3},{1,2,3},{1,2,3}}) indices_a = s32[5] parameter(0) indices_b = s32[2] parameter(1) gather_a = s32[5,3] gather(operand, indices_a), @@ -309,7 +308,7 @@ TEST_F(IndexedArrayAnalysisTest, ReshapeOfGather0) { HloModule ReshapeOfGather ENTRY main { - operand = s32[3,4] constant(s32[3,4]{{1,2,3,4},{1,2,3,4},{1,2,3,4}}) + operand = s32[3,4] constant({{1,2,3,4},{1,2,3,4},{1,2,3,4}}) indices = s32[5] parameter(0) gather = s32[5,4] gather(operand, indices), offset_dims={1}, @@ -330,7 +329,7 @@ TEST_F(IndexedArrayAnalysisTest, ReshapeOfGather1) { HloModule ReshapeOfGather ENTRY main { - operand = s32[3,4] constant(s32[3,4]{{1,2,3,4},{1,2,3,4},{1,2,3,4}}) + operand = s32[3,4] constant({{1,2,3,4},{1,2,3,4},{1,2,3,4}}) indices = s32[5,7] parameter(0) gather = s32[5,4,7] gather(operand, indices), offset_dims={1}, @@ -352,7 +351,7 @@ TEST_F(IndexedArrayAnalysisTest, ReshapeOfGather2) { HloModule ReshapeOfGather ENTRY main { - operand = s32[3,2,6] constant(s32[3,2,6]{ + operand = s32[3,2,6] constant({ {{1,2,3,4,5,6},{1,2,3,4,5,6}}, {{1,2,3,4,5,6},{1,2,3,4,5,6}}, {{1,2,3,4,5,6},{1,2,3,4,5,6}}}) @@ -377,7 +376,7 @@ TEST_F(IndexedArrayAnalysisTest, ReshapeOfGather3) { HloModule ReshapeOfGather ENTRY main { - operand = s32[2,6] constant(s32[2,6]{ + operand = s32[2,6] constant({ {1,2,3,4,5,6},{1,2,3,4,5,6}}) indices = s32[1] parameter(0) gather = s32[1,6] gather(operand, indices), @@ -405,7 +404,7 @@ TEST_F(IndexedArrayAnalysisTest, ReshapeOfGather4) { HloModule ReshapeOfGather ENTRY main { - operand = s32[2,3]{1,0} constant(s32[2,3] { { 1, 2, 3 }, { 1, 2, 3 } }) + operand = s32[2,3]{1,0} constant({ { 1, 2, 3 }, { 1, 2, 3 } }) i.0 = s64[1,3]{1,0} parameter(0) g.0 = s32[1,3,3]{2,1,0} gather(operand, i.0), offset_dims={2}, @@ -438,7 +437,7 @@ TEST_F(IndexedArrayAnalysisTest, ReshapeOfGather5) { HloModule ReshapeOfGather ENTRY main { - operand = s32[1,6] constant(s32[1,6]{{1,2,3,4,5,6}}) + operand = s32[1,6] constant({{1,2,3,4,5,6}}) indices = s32[1] parameter(0) gather = s32[1,6] gather(operand, indices), offset_dims={1}, @@ -465,7 +464,7 @@ TEST_F(IndexedArrayAnalysisTest, ReshapeOfGather6) { HloModule ReshapeOfGather ENTRY main { - operand = s32[1,2,6] constant(s32[1,2,6]{{ + operand = s32[1,2,6] constant({{ {1,2,3,4,5,6},{1,2,3,4,5,6}}}) indices = s32[1] parameter(0) gather = s32[1,1,6] gather(operand, indices), @@ -481,8 +480,8 @@ ENTRY main { const char* expected_root_expression = R"( (scalar-indexed-const (constant s32[2,1,1,1,6] s32[2,1,1,1,6] { - { /*i0=0*/ { /*i1=0*/ { /*i2=0*/ {1, 2, 3, 4, 5, 6} } } }, - { /*i0=1*/ { /*i1=0*/ { /*i2=0*/ {1, 2, 3, 4, 5, 6} } } } }) + { /*i0=0*/ { /*i1=0*/ { /*i2=0*/ { 1, 2, 3, 4, 5, 6 } } } }, + { /*i0=1*/ { /*i1=0*/ { /*i2=0*/ { 1, 2, 3, 4, 5, 6 } } } } }) (reshape %indices to s32[]) 0->[]) )"; @@ -496,7 +495,7 @@ TEST_F(IndexedArrayAnalysisTest, ReshapeOfGather7) { HloModule ReshapeOfGather ENTRY main { - operand = s32[2,6] constant(s32[2,6]{ + operand = s32[2,6] constant({ {1,2,3,4,5,6},{1,2,3,4,5,6}}) indices = s32[1,5] parameter(0) gather = s32[1,5,6] gather(operand, indices), @@ -512,8 +511,8 @@ ENTRY main { const char* expected_root_expression = R"( (scalar-indexed-const (constant s32[2,1,1,6] s32[2,1,1,6] { - { /*i0=0*/ { /*i1=0*/ {1, 2, 3, 4, 5, 6} } }, - { /*i0=1*/ { /*i1=0*/ {1, 2, 3, 4, 5, 6} } } }) + { /*i0=0*/ { /*i1=0*/ { 1, 2, 3, 4, 5, 6 } } }, + { /*i0=1*/ { /*i1=0*/ { 1, 2, 3, 4, 5, 6 } } } }) (reshape %indices to s32[5]) 0->[2]) )"; @@ -527,7 +526,7 @@ TEST_F(IndexedArrayAnalysisTest, ReshapeOfGatherNoFold0) { HloModule ReshapeOfGather ENTRY main { - operand = s32[3,4] constant(s32[3,4]{{1,2,3,4},{1,2,3,4},{1,2,3,4}}) + operand = s32[3,4] constant({{1,2,3,4},{1,2,3,4},{1,2,3,4}}) indices = s32[5,6] parameter(0) gather = s32[5,4,6] gather(operand, indices), offset_dims={1}, @@ -556,7 +555,7 @@ TEST_F(IndexedArrayAnalysisTest, ReshapeOfGatherNoFold1) { HloModule ReshapeOfGather ENTRY main { - operand = s32[3,5,2] constant(s32[3,5,2]{ + operand = s32[3,5,2] constant({ {{1,2},{3,4},{5,6},{7,8},{9,10}}, {{1,2},{3,4},{5,6},{7,8},{9,10}}, {{1,2},{3,4},{5,6},{7,8},{9,10}}}) @@ -588,7 +587,7 @@ TEST_F(IndexedArrayAnalysisTest, ReshapeOfGatherNoFold2) { HloModule ReshapeOfGather ENTRY main { - operand = s32[3,4,1] constant(s32[3,4,1]{ + operand = s32[3,4,1] constant({ {{1},{2},{3},{4}}, {{1},{2},{3},{4}}, {{1},{2},{3},{4}}}) @@ -620,7 +619,7 @@ TEST_F(IndexedArrayAnalysisTest, UnaryOpOfGather) { HloModule UnaryOpOfGather ENTRY main { - operand = f32[3,4] constant(f32[3,4]{{1,2,3,4},{1,3,2,4},{4,3,2,1}}) + operand = f32[3,4] constant({{1,2,3,4},{1,3,2,4},{4,3,2,1}}) indices = s32[5] parameter(0) gather = f32[5,4] gather(operand, indices), offset_dims={1}, @@ -645,7 +644,7 @@ TEST_F(IndexedArrayAnalysisTest, AddBroadcastedScalarWithGather) { HloModule AddBroadcastedScalarWithGather ENTRY main { - gather_operand = s32[3,4] constant(s32[3,4]{{1,2,3,4},{1,3,2,4},{4,3,2,1}}) + gather_operand = s32[3,4] constant({{1,2,3,4},{1,3,2,4},{4,3,2,1}}) constant = s32[] constant(5) constant_broadcasted = s32[5,4] broadcast(constant), dimensions={} indices = s32[5] parameter(0) @@ -673,7 +672,7 @@ TEST_F(IndexedArrayAnalysisTest, HloModule SubtractBroadcastedScalarWithGather ENTRY main { - gather_operand = s32[3,4] constant(s32[3,4]{{1,2,3,4},{1,3,2,4},{4,3,2,1}}) + gather_operand = s32[3,4] constant({{1,2,3,4},{1,3,2,4},{4,3,2,1}}) constant = s32[] constant(5) constant_broadcasted = s32[5,4] broadcast(constant), dimensions={} indices = s32[5] parameter(0) @@ -701,7 +700,7 @@ TEST_F(IndexedArrayAnalysisTest, HloModule SubtractBroadcastedScalarWithGather ENTRY main { - gather_operand = s32[3,4] constant(s32[3,4]{{1,2,3,4},{1,3,2,4},{4,3,2,1}}) + gather_operand = s32[3,4] constant({{1,2,3,4},{1,3,2,4},{4,3,2,1}}) constant = s32[] constant(5) constant_broadcasted = s32[5,4] broadcast(constant), dimensions={} indices = s32[5] parameter(0) @@ -728,7 +727,7 @@ TEST_F(IndexedArrayAnalysisTest, AddBroadcastedVectorWithGather) { HloModule AddBroadcastedVectorWithGather ENTRY main { - gather_operand = s32[3,4] constant(s32[3,4]{{1,2,3,4},{1,3,2,4},{4,3,2,1}}) + gather_operand = s32[3,4] constant({{1,2,3,4},{1,3,2,4},{4,3,2,1}}) constant_vect = s32[4] constant({10,11,12,13}) constant_broadcasted = s32[5,4] broadcast(constant_vect), dimensions={1} indices = s32[5] parameter(0) @@ -755,7 +754,7 @@ TEST_F(IndexedArrayAnalysisTest, AddBroadcastedVectorWithGather_Negative) { HloModule AddBroadcastedVectorWithGather ENTRY main { - gather_operand = s32[3,4] constant(s32[3,4]{{1,2,3,4},{1,3,2,4},{4,3,2,1}}) + gather_operand = s32[3,4] constant({{1,2,3,4},{1,3,2,4},{4,3,2,1}}) constant_vect = s32[5] constant({10,11,12,13,14}) constant_broadcasted = s32[5,4] broadcast(constant_vect), dimensions={0} indices = s32[5] parameter(0) @@ -804,8 +803,8 @@ TEST_F(IndexedArrayAnalysisTest, DotOpBasic_0) { HloModule DotOp ENTRY main { - gather_operand = s32[3,4] constant(s32[3,4]{{1,2,3,4},{5,6,7,8},{9,10,11,12}}) - dot_rhs_constant = s32[4,3] constant(s32[4,3]{{1,2,3},{4,5,6},{7,8,9},{10,11,12}}) + gather_operand = s32[3,4] constant({{1,2,3,4},{5,6,7,8},{9,10,11,12}}) + dot_rhs_constant = s32[4,3] constant({{1,2,3},{4,5,6},{7,8,9},{10,11,12}}) indices = s32[5] parameter(0) dot_lhs = s32[5,4] gather(gather_operand, indices), offset_dims={1}, @@ -831,8 +830,8 @@ TEST_F(IndexedArrayAnalysisTest, DotOpBasic_1) { HloModule DotOp ENTRY main { - gather_operand = s32[3,4] constant(s32[3,4]{{1,2,3,4},{5,6,7,8},{9,10,11,12}}) - dot_rhs_constant = s32[3,3] constant(s32[3,3]{{1,2,3},{4,5,6},{7,8,9}}) + gather_operand = s32[3,4] constant({{1,2,3,4},{5,6,7,8},{9,10,11,12}}) + dot_rhs_constant = s32[3,3] constant({{1,2,3},{4,5,6},{7,8,9}}) indices = s32[5] parameter(0) dot_lhs = s32[3,5] gather(gather_operand, indices), offset_dims={0}, @@ -859,8 +858,8 @@ TEST_F(IndexedArrayAnalysisTest, DotOpBasic_2) { HloModule DotOp ENTRY main { - gather_operand = s32[3,4] constant(s32[3,4]{{1,2,3,4},{5,6,7,8},{9,10,11,12}}) - dot_lhs_constant = s32[4,3] constant(s32[4,3]{{1,2,3},{4,5,6},{7,8,9},{10,11,12}}) + gather_operand = s32[3,4] constant({{1,2,3,4},{5,6,7,8},{9,10,11,12}}) + dot_lhs_constant = s32[4,3] constant({{1,2,3},{4,5,6},{7,8,9},{10,11,12}}) indices = s32[5] parameter(0) dot_rhs = s32[3,5] gather(gather_operand, indices), offset_dims={0}, @@ -888,8 +887,8 @@ TEST_F(IndexedArrayAnalysisTest, DotOpBasic_3) { HloModule DotOp ENTRY main { - gather_operand = s32[4,3] constant(s32[4,3]{{1,2,3},{4,5,6},{7,8,9},{10,11,12}}) - dot_lhs_constant = s32[4,3] constant(s32[4,3]{{1,2,3},{4,5,6},{7,8,9},{10,11,12}}) + gather_operand = s32[4,3] constant({{1,2,3},{4,5,6},{7,8,9},{10,11,12}}) + dot_lhs_constant = s32[4,3] constant({{1,2,3},{4,5,6},{7,8,9},{10,11,12}}) indices = s32[5] parameter(0) dot_rhs = s32[5,3] gather(gather_operand, indices), offset_dims={1}, @@ -917,8 +916,8 @@ TEST_F(IndexedArrayAnalysisTest, DotOpWithBatch) { HloModule DotOp ENTRY main { - gather_operand = s32[2,3,2] constant(s32[2,3,2]{{{1,2},{3,4},{5,6}},{{7,8},{9,10},{11,12}}}) - dot_lhs_constant = s32[2,2,3] constant(s32[2,2,3]{{{1,2,3},{4,5,6}},{{7,8,9},{10,11,12}}}) + gather_operand = s32[2,3,2] constant({{{1,2},{3,4},{5,6}},{{7,8},{9,10},{11,12}}}) + dot_lhs_constant = s32[2,2,3] constant({{{1,2,3},{4,5,6}},{{7,8,9},{10,11,12}}}) indices = s32[4] parameter(0) dot_rhs = s32[2,3,4] gather(gather_operand, indices), offset_dims={0,1}, @@ -948,8 +947,8 @@ TEST_F(IndexedArrayAnalysisTest, DotOpNegative) { HloModule DotOp ENTRY main { - gather_operand = s32[3,4] constant(s32[3,4]{{1,2,3,4},{5,6,7,8},{9,10,11,12}}) - dot_rhs_constant = s32[2,3] constant(s32[2,3]{{1,2,3},{4,5,6}}) + gather_operand = s32[3,4] constant({{1,2,3,4},{5,6,7,8},{9,10,11,12}}) + dot_rhs_constant = s32[2,3] constant({{1,2,3},{4,5,6}}) indices = s32[2] parameter(0) dot_lhs = s32[3,2] gather(gather_operand, indices), offset_dims={0}, diff --git a/tensorflow/compiler/xla/service/instruction_fusion.cc b/tensorflow/compiler/xla/service/instruction_fusion.cc index 69a4c160ee5c4539272c3085338dc6de1b9347ff..b97060535d998e174639dceca5cde517cef01e30 100644 --- a/tensorflow/compiler/xla/service/instruction_fusion.cc +++ b/tensorflow/compiler/xla/service/instruction_fusion.cc @@ -23,10 +23,13 @@ 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/memory/memory.h" #include "tensorflow/compiler/xla/map_util.h" #include "tensorflow/compiler/xla/service/fusion_queue.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" +#include "tensorflow/compiler/xla/service/hlo_reachability.h" #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/platform/logging.h" @@ -101,7 +104,6 @@ bool IsAlwaysDuplicable(const HloInstruction& instruction) { case HloOpcode::kShiftRightLogical: case HloOpcode::kSlice: case HloOpcode::kSubtract: - case HloOpcode::kAfterAll: case HloOpcode::kTranspose: case HloOpcode::kTuple: case HloOpcode::kTupleSelect: @@ -114,7 +116,10 @@ bool IsAlwaysDuplicable(const HloInstruction& instruction) { case HloOpcode::kSin: return ShapeUtil::ElementIsComplex(instruction.shape()); - // Expensive instructions. + // Expensive instructions or unusual instructions for which fusion is + // nonsensical. + case HloOpcode::kAddDependency: + case HloOpcode::kAfterAll: case HloOpcode::kAtan2: case HloOpcode::kBatchNormGrad: case HloOpcode::kBatchNormInference: @@ -122,7 +127,7 @@ bool IsAlwaysDuplicable(const HloInstruction& instruction) { case HloOpcode::kCall: case HloOpcode::kConditional: case HloOpcode::kConvolution: - case HloOpcode::kCrossReplicaSum: + case HloOpcode::kAllReduce: case HloOpcode::kAllToAll: case HloOpcode::kCollectivePermute: case HloOpcode::kCustomCall: @@ -153,6 +158,7 @@ bool IsAlwaysDuplicable(const HloInstruction& instruction) { case HloOpcode::kTanh: case HloOpcode::kTrace: case HloOpcode::kWhile: + case HloOpcode::kGetDimensionSize: return true; } @@ -168,23 +174,22 @@ bool InstructionFusion::EffectivelyAtMostUnary(HloInstruction* hlo) { ShapeUtil::ForEachSubshape( hlo->shape(), [&output_rank](const Shape& subshape, const ShapeIndex& shape_index) { - if (ShapeUtil::IsArray(subshape)) { + if (subshape.IsArray()) { 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( @@ -268,7 +273,7 @@ InstructionFusion::ComputeGloballyUnfusible( ShapeUtil::ForEachSubshape( shape, [&size](const Shape& subshape, const ShapeIndex& shape_index) { - if (ShapeUtil::IsArray(subshape)) { + if (subshape.IsArray()) { size += ShapeUtil::ElementsIn(subshape); } }); @@ -403,9 +408,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)) > @@ -437,8 +441,7 @@ class ReversePostOrderFusionQueue : public FusionQueue { } // namespace std::unique_ptr InstructionFusion::GetFusionQueue( - HloComputation* computation, - const std::function& skip_producer) { + HloComputation* computation) { return absl::make_unique(computation); } @@ -451,14 +454,16 @@ StatusOr InstructionFusion::Run(HloModule* module) { for (auto* computation : module->MakeNonfusionComputations()) { CHECK(!computation->IsFusionComputation()); computation_ = computation; - reachability_ = computation_->ComputeReachability(); - - HloInstructionSet do_not_duplicate = - ComputeGloballyUnfusible(computation_->MakeInstructionPostOrder()); - auto fusion_queue = - GetFusionQueue(computation_, [&](HloInstruction* producer) { - return do_not_duplicate.count(producer) > 0; - }); + reachability_ = HloReachabilityMap::Build(computation_); + + HloInstructionSet do_not_duplicate; + // If we allow duplications, we need to compute which instructions we do not + // want to duplicate based on a global analysis of the graph. + if (may_duplicate_) { + do_not_duplicate = + ComputeGloballyUnfusible(computation_->MakeInstructionPostOrder()); + } + auto fusion_queue = GetFusionQueue(computation_); // Instruction fusion effectively fuses edges in the computation graph // (producer instruction -> consumer instruction) so we iterate over all @@ -489,9 +494,8 @@ StatusOr InstructionFusion::Run(HloModule* module) { HloInstruction* fusion_instruction; // Try "regular" fusion if the operand may be duplicated. Otherwise, // perform multi-output fusion, unless this creates a cycle. - // TODO(tjoerg): Consider making multi-output fusion the default. - if (ShouldFuse(instruction, i) && - do_not_duplicate.count(operand) == 0) { + if (do_not_duplicate.count(operand) == 0 && + ShouldFuse(instruction, i)) { fusion_queue->PreFusion(operand, instruction); fusion_instruction = Fuse(operand, instruction); } else if (ShouldFuseIntoMultiOutput(instruction, i) && @@ -565,15 +569,42 @@ HloInstruction* InstructionFusion::FuseIntoMultiOutput( bool InstructionFusion::MultiOutputFusionCreatesCycle( HloInstruction* producer, HloInstruction* consumer) { - return absl::c_any_of( - consumer->operands(), [&](const HloInstruction* consumer_operand) { - // The fusion algorithm traverses the HLO graph in reverse post order. - // Thus `cosumers` is visited before its operands (including - // `producer`). Therefore, consumer operands cannot have been fused yet. - // It is thus safe to use the pre-computed reachability map. - return consumer_operand != producer && - reachability_->IsReachable(producer, consumer_operand); - }); + absl::flat_hash_set operands; + for (const HloInstruction* operand : consumer->operands()) { + if (operand == producer) { + continue; + } + + // If the reachability map already contains the producer and the operand of + // the consumer, and the producer can reach the operand, then we know for + // sure MultiOutputFusion would create a cycle. If not, we need to do a DFS + // traversal of the computation to verify that this multioutput fusion would + // not create a cycle. + if (reachability_->IsPresent(producer) && + reachability_->IsPresent(operand) && + reachability_->IsReachable(producer, operand)) { + return true; + } + operands.insert(operand->unique_id()); + } + + // Do a DFS on the producer to see if any of the other consumer operands are + // reachable in the current state of the graph. + std::vector worklist = producer->users(); + absl::flat_hash_set visits; + while (!worklist.empty()) { + const HloInstruction* user = worklist.back(); + worklist.pop_back(); + if (operands.count(user->unique_id()) != 0) { + return true; + } + if (visits.count(user->unique_id()) == 0) { + visits.insert(user->unique_id()); + worklist.insert(worklist.end(), user->users().begin(), + user->users().end()); + } + } + return false; } bool InstructionFusion::ShouldFuse(HloInstruction* consumer, diff --git a/tensorflow/compiler/xla/service/instruction_fusion.h b/tensorflow/compiler/xla/service/instruction_fusion.h index f14c6675208c72112aea0179c238b58709d625b5..198bd7fce5f392e5e895b959523d4fe9cf208ba2 100644 --- a/tensorflow/compiler/xla/service/instruction_fusion.h +++ b/tensorflow/compiler/xla/service/instruction_fusion.h @@ -22,6 +22,7 @@ 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_pass_interface.h" +#include "tensorflow/compiler/xla/service/hlo_reachability.h" #include "tensorflow/core/platform/macros.h" namespace xla { @@ -54,8 +55,7 @@ class InstructionFusion : public HloModulePass { // fused. The default implementation processes consumers in reverse post // order. virtual std::unique_ptr GetFusionQueue( - HloComputation* computation, - const std::function& skip_producer); + HloComputation* computation); // Returns whether the given producer instruction should be fused into the // given consumer instruction. producer is necessarily an operand of consumer. @@ -111,6 +111,10 @@ class InstructionFusion : public HloModulePass { return is_expensive_(instruction); } + // Whether multi-output fusion would introduce a cycle into the HLO graph. + bool MultiOutputFusionCreatesCycle(HloInstruction* producer, + HloInstruction* consumer); + // Current HloComputation instance the loop fuser is traversing. HloComputation* computation_; HloModule* module_; @@ -145,10 +149,6 @@ class InstructionFusion : public HloModulePass { // duplicated. std::function is_expensive_; - // Whether multi-output fusion would introduce a cycle into the HLO graph. - bool MultiOutputFusionCreatesCycle(HloInstruction* producer, - HloInstruction* consumer); - // Returns whether we may duplicate an instruction if we want to fuse it. bool may_duplicate_; diff --git a/tensorflow/compiler/xla/service/instruction_fusion_test.cc b/tensorflow/compiler/xla/service/instruction_fusion_test.cc index da1ad90959dc0ab1a840b3390281ce9d4999651e..611cfd404d7622f561f0acc86fc9b05e16eea22e 100644 --- a/tensorflow/compiler/xla/service/instruction_fusion_test.cc +++ b/tensorflow/compiler/xla/service/instruction_fusion_test.cc @@ -117,7 +117,7 @@ TEST_F(InstructionFusionTest, PotentialBitcastReshapeOfParameterUnfused) { auto reshape1 = builder.AddInstruction( HloInstruction::CreateReshape(ShapeUtil::MakeShape(S32, {1, 1}), param0)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(reshape1, computation->root_instruction()); EXPECT_FALSE( @@ -133,7 +133,7 @@ TEST_F(InstructionFusionTest, PotentialBitcastSimpleReshapeOfParameterUnfused) { auto reshape1 = builder.AddInstruction( HloInstruction::CreateReshape(ShapeUtil::MakeShape(S32, {1, 1}), param0)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(reshape1, computation->root_instruction()); EXPECT_FALSE( @@ -149,7 +149,7 @@ TEST_F(InstructionFusionTest, PotentialBitcastTransposeOfParameterUnfused) { auto transpose1 = builder.AddInstruction(HloInstruction::CreateTranspose( ShapeUtil::MakeShape(S32, {}), param0, {})); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(transpose1, computation->root_instruction()); EXPECT_FALSE( @@ -172,7 +172,7 @@ TEST_F(InstructionFusionTest, AvoidDuplicationIfNotAllFusible) { HloInstruction* unary = builder.AddInstruction( HloInstruction::CreateUnary(shape, HloOpcode::kAbs, binary1)); - auto module = CreateNewModule(); + auto module = CreateNewUnverifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(unary, computation->root_instruction()); EXPECT_FALSE( @@ -259,8 +259,8 @@ TEST_F(InstructionFusionTest, AvoidDuplicationIfNotAllFusibleRecursively) { add = f32[4,3]{1,0} add(p0, p0) abs1 = f32[4,3]{1,0} abs(add) log = f32[4,3]{1,0} log(abs1) - token = token[] after-all() - send = f32[4,3]{1,0} send(log, token), channel_id=0 + token0 = token[] after-all() + send = f32[4,3]{1,0} send(log, token0), channel_id=0 abs2 = f32[4,3]{1,0} abs(log) ROOT root = f32[4,3]{1,0} subtract(abs2, add) })") @@ -290,8 +290,8 @@ TEST_F(InstructionFusionTest, AvoidDuplicationIfNotAllFusibleRecursively) { p0 = f32[4,3]{1,0} parameter(0) add1 = f32[4,3]{1,0} add(p0, p0) log = f32[4,3]{1,0} log(p0) - token = token[] after-all() - send = f32[4,3]{1,0} send(log, token), channel_id=0 + token0 = token[] after-all() + send = f32[4,3]{1,0} send(log, token0), channel_id=0 add2 = f32[4,3]{1,0} add(log, add1) ROOT root = f32[4,3]{1,0} subtract(add1, add2) })") @@ -324,8 +324,8 @@ TEST_F(InstructionFusionTest, AvoidDuplicationIfNotAllFusibleRecursively) { add1 = f32[4,3]{1,0} add(p0, p0) add2 = f32[4,3]{1,0} add(add1, add1) log = f32[4,3]{1,0} log(add2) - token = token[] after-all() - send = f32[4,3]{1,0} send(log, token), channel_id=0 + token0 = token[] after-all() + send = f32[4,3]{1,0} send(log, token0), channel_id=0 sub1 = f32[4,3]{1,0} subtract(log, add2) sub2 = f32[4,3]{1,0} subtract(add2, add1) ROOT root = (f32[4,3]{1,0}, f32[4,3]{1,0}) tuple(sub1, sub2) @@ -361,7 +361,7 @@ TEST_F(InstructionFusionTest, AllowUnaryDuplication) { HloInstruction* unary2 = builder.AddInstruction( HloInstruction::CreateUnary(shape, HloOpcode::kAbs, unary1)); - auto module = CreateNewModule(); + auto module = CreateNewUnverifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(unary2, computation->root_instruction()); EXPECT_TRUE( @@ -385,7 +385,7 @@ TEST_F(InstructionFusionTest, AllowEffectiveUnaryDuplication) { HloInstruction* unary = builder.AddInstruction( HloInstruction::CreateUnary(shape, HloOpcode::kAbs, binary1)); - auto module = CreateNewModule(); + auto module = CreateNewUnverifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(unary, computation->root_instruction()); EXPECT_TRUE( @@ -394,6 +394,56 @@ TEST_F(InstructionFusionTest, AllowEffectiveUnaryDuplication) { .ValueOrDie()); } +TEST_F(InstructionFusionTest, FuseDiamondGraphsNoDuplication) { + auto module = ParseHloString(R"( + HloModule test_module + ENTRY Test { + p0 = f32[100] parameter(0) + p1 = f32[100] parameter(1) + add = f32[100] add(p0, p1) + slice1 = f32[99] slice(add), slice={[0:99:1]} + slice2 = f32[99] slice(add), slice={[1:100:1]} + ROOT add2 = f32[99] add(slice1, slice2) + })") + .ValueOrDie(); + EXPECT_TRUE( + InstructionFusion(InstructionFusion::IsExpensive, /*may_duplicate=*/false) + .Run(module.get()) + .ValueOrDie()) + << module->ToString(); + + HloInstruction* root = module->entry_computation()->root_instruction(); + // 'add' would originally need to be duplicated if fused. However after its + // two users 'slice1' and 'slice2' are fused into 'add2', 'add' has only one + // user and can now be also fused. + EXPECT_THAT(root, op::Fusion(op::Parameter(), op::Parameter())); +} + +TEST_F(InstructionFusionTest, FuseDiamondGraphsAllowDuplication) { + auto module = ParseHloString(R"( + HloModule test_module + ENTRY Test { + p0 = f32[100] parameter(0) + p1 = f32[100] parameter(1) + add = f32[100] add(p0, p1) + slice1 = f32[99] slice(add), slice={[0:99:1]} + slice2 = f32[99] slice(add), slice={[1:100:1]} + ROOT add2 = f32[99] add(slice1, slice2) + })") + .ValueOrDie(); + EXPECT_TRUE( + InstructionFusion(InstructionFusion::IsExpensive, /*may_duplicate=*/true) + .Run(module.get()) + .ValueOrDie()) + << module->ToString(); + + HloInstruction* root = module->entry_computation()->root_instruction(); + // 'add' would originally need to be duplicated if fused. However after its + // two users 'slice1' and 'slice2' are fused into 'add2', 'add' has only one + // user and can now be also fused. + EXPECT_THAT(root, op::Fusion(op::Parameter(), op::Parameter())); +} + TEST_F(InstructionFusionTest, WideningConvertsAreAlwaysDuplicableIntoConsumers) { auto module = ParseHloString(R"( diff --git a/tensorflow/compiler/xla/service/interpreter/BUILD b/tensorflow/compiler/xla/service/interpreter/BUILD index 1484e14df10d94841c5a2e849761779f5800392d..545662543cca40e42b0f0302e14152e5283f9e4f 100644 --- a/tensorflow/compiler/xla/service/interpreter/BUILD +++ b/tensorflow/compiler/xla/service/interpreter/BUILD @@ -1,12 +1,12 @@ -licenses(["restricted"]) - -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,6 +48,7 @@ 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:while_loop_simplifier", "//tensorflow/core:lib", @@ -115,6 +117,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 a1fe97cffa4de1993396d2443166321bc795b553..4818b2dae0a9951346600a9b2906488c3ef7e06e 100644 --- a/tensorflow/compiler/xla/service/interpreter/compiler.cc +++ b/tensorflow/compiler/xla/service/interpreter/compiler.cc @@ -21,6 +21,7 @@ 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/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,6 +32,7 @@ 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/while_loop_simplifier.h" #include "tensorflow/compiler/xla/status_macros.h" @@ -43,9 +45,15 @@ namespace interpreter { Status InterpreterCompiler::RunHloOptimization(HloModule* hlo_module) { HloPassPipeline pipeline("Interpreter"); + 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(); } @@ -58,7 +66,8 @@ StatusOr> InterpreterCompiler::RunHloPasses( } Status InterpreterCompiler::RunHloPassesOnModuleGroup( - HloModuleGroup* module_group, se::StreamExecutor* executor, + HloModuleGroup* module_group, + absl::Span executors, DeviceMemoryAllocator* device_allocator) { return Unimplemented("Module group compilation not supported on Interpreter"); } @@ -75,9 +84,12 @@ StatusOr> InterpreterCompiler::RunBackend( // 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()); std::unique_ptr executable = - absl::make_unique( - std::move(hlo_module), absl::make_unique()); + absl::make_unique(std::move(hlo_module), + std::move(evaluator)); return std::move(executable); } diff --git a/tensorflow/compiler/xla/service/interpreter/compiler.h b/tensorflow/compiler/xla/service/interpreter/compiler.h index d8cb32c0beb279ae6484b1b8f5f99085c2d67c67..591272951a01a3e2aa3b615673dceced8e94f674 100644 --- a/tensorflow/compiler/xla/service/interpreter/compiler.h +++ b/tensorflow/compiler/xla/service/interpreter/compiler.h @@ -47,7 +47,8 @@ class InterpreterCompiler : public Compiler { std::unique_ptr hlo_module, se::StreamExecutor* stream_exec, DeviceMemoryAllocator* device_allocator) override; Status RunHloPassesOnModuleGroup( - HloModuleGroup* module_group, se::StreamExecutor* executor, + HloModuleGroup* module_group, + absl::Span executors, DeviceMemoryAllocator* device_allocator) override; StatusOr> RunBackend( diff --git a/tensorflow/compiler/xla/service/interpreter/executable.cc b/tensorflow/compiler/xla/service/interpreter/executable.cc index a06d6113e84630df14ff68280c248cccb9afaf06..7a6ebdef708bcc3a92fbd8618db0c42c35e6ce8b 100644 --- a/tensorflow/compiler/xla/service/interpreter/executable.cc +++ b/tensorflow/compiler/xla/service/interpreter/executable.cc @@ -37,7 +37,7 @@ namespace xla { namespace interpreter { InterpreterExecutable::InterpreterExecutable( - std::unique_ptr hlo_module, + std::unique_ptr hlo_module, std::unique_ptr evaluator) : Executable(std::move(hlo_module), /*hlo_profile_printer=*/nullptr, /*hlo_profile_index_map=*/nullptr), @@ -68,6 +68,18 @@ StatusOr InterpreterExecutable::ExecuteOnStream( "Mismatch between argument count and graph parameter count."); } + // Check that the args have the right shape. + for (int64 i = 0; i < computation->num_parameters(); ++i) { + const auto& expected_shape = computation->parameter_instruction(i)->shape(); + const auto& actual_shape = arguments[i]->on_device_shape(); + if (!ShapeUtil::Equal(expected_shape, actual_shape)) { + return InvalidArgument( + "Shape mismatch on parameter %d. Expected %s, but was %s.", i, + ShapeUtil::HumanString(expected_shape), + ShapeUtil::HumanString(actual_shape)); + } + } + TF_ASSIGN_OR_RETURN(TransferManager * transfer_manager, TransferManager::GetForPlatform(platform)); @@ -85,8 +97,9 @@ StatusOr InterpreterExecutable::ExecuteOnStream( Literal result_literal; { tensorflow::mutex_lock lock(evaluator_lock_); - TF_ASSIGN_OR_RETURN(result_literal, evaluator_->Evaluate( - *computation, arg_literals)); + evaluator_->ResetVisitStates(); + TF_ASSIGN_OR_RETURN(result_literal, + evaluator_->Evaluate(*computation, arg_literals)); } // Transform the result literal back into a ShapedBuffer. @@ -116,7 +129,7 @@ StatusOr InterpreterExecutable::ExecuteAsyncOnStream( } /*static*/ int64 InterpreterExecutable::ShapeSizeBytes(const Shape& shape) { - if (ShapeUtil::IsOpaque(shape)) { + if (shape.IsOpaque()) { return sizeof(void*); } return ShapeUtil::ByteSizeOf(shape, sizeof(void*)); diff --git a/tensorflow/compiler/xla/service/interpreter/executable.h b/tensorflow/compiler/xla/service/interpreter/executable.h index 3b1ebce0c75457d65e6834c809fe488a9c4a159a..bda13d376360306c81230e41b01cefc6caff230d 100644 --- a/tensorflow/compiler/xla/service/interpreter/executable.h +++ b/tensorflow/compiler/xla/service/interpreter/executable.h @@ -42,7 +42,7 @@ namespace interpreter { // buffer allocation. Refer to interpreter/README.md for more. class InterpreterExecutable : public Executable { public: - InterpreterExecutable(std::unique_ptr hlo_module, + InterpreterExecutable(std::unique_ptr hlo_module, std::unique_ptr evaluator); ~InterpreterExecutable() override; diff --git a/tensorflow/compiler/xla/service/interpreter/executor.cc b/tensorflow/compiler/xla/service/interpreter/executor.cc index 4fb67bd0b72fc591c1ffa76ebb0513bf14ed3737..e3e5fa71543baa309b3a68888b1b9bdfd43cfbd5 100644 --- a/tensorflow/compiler/xla/service/interpreter/executor.cc +++ b/tensorflow/compiler/xla/service/interpreter/executor.cc @@ -78,9 +78,14 @@ port::Status XlaInterpreterExecutor::SynchronousMemcpy( return port::Status::OK(); } -bool XlaInterpreterExecutor::HostCallback(Stream *stream, - std::function callback) { - AsExecutorStream(stream)->EnqueueTask(callback); +bool XlaInterpreterExecutor::HostCallback( + Stream *stream, std::function callback) { + AsExecutorStream(stream)->EnqueueTask([callback]() { + port::Status s = callback(); + if (!s.ok()) { + LOG(WARNING) << "Host callback failed: " << s; + } + }); return true; } diff --git a/tensorflow/compiler/xla/service/interpreter/executor.h b/tensorflow/compiler/xla/service/interpreter/executor.h index fbb99457847dca69a1901006d5d8ff713882f918..400c30515464ed5b00251fba303fef303a26b97b 100644 --- a/tensorflow/compiler/xla/service/interpreter/executor.h +++ b/tensorflow/compiler/xla/service/interpreter/executor.h @@ -125,7 +125,8 @@ class XlaInterpreterExecutor : public internal::StreamExecutorInterface { return port::Status{port::error::UNIMPLEMENTED, ""}; } - bool HostCallback(Stream *stream, std::function callback) override; + bool HostCallback(Stream *stream, + std::function callback) override; port::Status AllocateEvent(Event *event) override { return port::Status{port::error::UNIMPLEMENTED, ""}; diff --git a/tensorflow/compiler/xla/service/interpreter/platform.cc b/tensorflow/compiler/xla/service/interpreter/platform.cc index c9b40d3c6195f80a19272a0d98890049d02315b9..b0fc1af8b89d7327a00f77f471e90d143a92de7c 100644 --- a/tensorflow/compiler/xla/service/interpreter/platform.cc +++ b/tensorflow/compiler/xla/service/interpreter/platform.cc @@ -110,3 +110,5 @@ REGISTER_MODULE_INITIALIZER( // open-source project, so this will be a no-op there. REGISTER_MODULE_INITIALIZER_SEQUENCE(interpreter_platform, multi_platform_manager); +REGISTER_MODULE_INITIALIZER_SEQUENCE(multi_platform_manager_listener, + interpreter_platform); diff --git a/tensorflow/compiler/xla/service/layout_assignment.cc b/tensorflow/compiler/xla/service/layout_assignment.cc index 232d1dc0879cd6931158e642e01fe68e43e6c655..5d9e3392fd86c587a0bd998a282c52d145cc710e 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, @@ -256,7 +253,7 @@ Status LayoutConstraints::SetArrayOperandLayout( const Layout& layout, const HloInstruction* instruction, int64 operand_no, bool mandatory, bool dfs) { const HloInstruction* operand = instruction->operand(operand_no); - TF_RET_CHECK(ShapeUtil::IsArray(operand->shape())); + TF_RET_CHECK(operand->shape().IsArray()); Shape shape(operand->shape()); *shape.mutable_layout() = layout; TF_RETURN_IF_ERROR(LayoutUtil::ValidateLayoutInShape(shape)); @@ -314,7 +311,7 @@ Status LayoutConstraints::SetInstructionLayout( CHECK_EQ(1, buffers.size()); CHECK_EQ(buffers[0]->instruction(), instruction); - if (ShapeUtil::IsArray(subshape)) { + if (subshape.IsArray()) { return SetBufferLayout(subshape.layout(), *buffers[0], mandatory); } else { return Status::OK(); @@ -406,7 +403,7 @@ Status LayoutAssignment::BuildHostChannelConstraints( instruction->opcode() == HloOpcode::kRecv) { const Shape& data_shape = ShapeUtil::GetTupleElementShape(send_recv_instr->shape(), 0); - TF_RET_CHECK(ShapeUtil::IsArray(data_shape)); + TF_RET_CHECK(data_shape.IsArray()); TF_RET_CHECK(LayoutUtil::HasLayout(data_shape)); const Layout* prev_layout = host_channel_constraints_.ConstrainChannel( send_recv_instr->channel_id(), data_shape.layout()); @@ -449,7 +446,6 @@ Status LayoutAssignment::AddMandatoryConstraints( // instruction. // TODO(b/31425034): Change infeeds to be more like parameters, with // shapes in the ComputationLayout. - DCHECK(!LayoutUtil::IsPadded(instruction->shape())); TF_RETURN_IF_ERROR( constraints->SetInstructionLayout(instruction->shape(), instruction)); } else if (instruction->opcode() == HloOpcode::kOutfeed) { @@ -490,7 +486,7 @@ Status LayoutAssignment::AddMandatoryConstraints( if (instruction->opcode() == HloOpcode::kSend) { // TODO(b/68493863): Change to use SetOperandLayout(). const Shape send_buffer_shape = instruction->operand(0)->shape(); - TF_RET_CHECK(ShapeUtil::IsArray(send_buffer_shape)); + TF_RET_CHECK(send_buffer_shape.IsArray()); Shape new_buffer_shape = get_channel_constraints(instruction) ->LayoutShapeForChannel(send_buffer_shape, @@ -500,7 +496,7 @@ Status LayoutAssignment::AddMandatoryConstraints( } else { const Shape recv_buffer_shape = ShapeUtil::GetTupleElementShape(instruction->shape(), 0); - TF_RET_CHECK(ShapeUtil::IsArray(recv_buffer_shape)); + TF_RET_CHECK(recv_buffer_shape.IsArray()); TF_ASSIGN_OR_RETURN( const LogicalBuffer* buffer, constraints->points_to_analysis().GetBufferDefinedAt(instruction, @@ -521,7 +517,7 @@ Status LayoutAssignment::AddMandatoryConstraints( } // TODO(b/68493863): Change to use SetOperandLayout(). const Shape& buffer_shape = instruction->operand(0)->shape(); - TF_RET_CHECK(ShapeUtil::IsArray(buffer_shape)); + TF_RET_CHECK(buffer_shape.IsArray()); Shape new_buffer_shape = get_channel_constraints(instruction) ->LayoutShapeForChannel(buffer_shape, all_reduce_id); @@ -781,7 +777,7 @@ StatusOr LayoutAssignment::CreateCopyWithNewLayout( << ShapeUtil::HumanString(instruction->shape()) << " instruction: " << instruction->ToString(); - if (ShapeUtil::IsTuple(instruction->shape())) { + if (instruction->shape().IsTuple()) { // Copy tuple elements which have differing layouts. std::vector element_copies; for (int64 i = 0; i < ShapeUtil::TupleElementCount(instruction->shape()); @@ -812,7 +808,7 @@ StatusOr LayoutAssignment::CreateCopyWithNewLayout( TF_RETURN_IF_ERROR(LayoutUtil::CopyLayoutBetweenShapes( shape_with_layout, tuple_copy->mutable_shape())); return tuple_copy; - } else if (ShapeUtil::IsArray(instruction->shape())) { + } else if (instruction->shape().IsArray()) { HloInstruction* copy = instruction->parent()->AddInstruction(HloInstruction::CreateUnary( instruction->shape(), HloOpcode::kCopy, instruction)); @@ -989,13 +985,10 @@ std::unique_ptr LayoutAssignment::ChooseOperandLayoutFromOutputLayout( const Layout& output_layout, const HloInstruction* instruction, int64 operand_no) { const HloInstruction* operand = instruction->operand(operand_no); - - CHECK(ShapeUtil::IsArray(instruction->shape())); - CHECK(ShapeUtil::IsArray(operand->shape())); - + CHECK(instruction->shape().IsArray()); + CHECK(operand->shape().IsArray()); if (!ShapeUtil::IsScalar(operand->shape()) && - ShapeUtil::Rank(operand->shape()) == - ShapeUtil::Rank(instruction->shape()) && + operand->shape().rank() == instruction->shape().rank() && !instruction_can_change_layout_func_(instruction)) { // Propagate the result layout to the operand layout if the instruction // requires the same layout out for the result and the operand. @@ -1015,7 +1008,7 @@ std::unique_ptr LayoutAssignment::ChooseOperandLayoutFromOutputLayout( // operations. For similar reasons, if the operand and output have the same // rank, try to match the operand's layout to the output. if (ShapeUtil::TrueRank(operand->shape()) == 1 && - ShapeUtil::Rank(instruction->shape()) == 1) { + instruction->shape().rank() == 1) { // Don't assign a layout in case of R1 -> effective R1 reshape. return nullptr; } @@ -1029,7 +1022,7 @@ std::unique_ptr LayoutAssignment::ChooseOperandLayoutFromOutputLayout( if (ShapeUtil::ReshapeIsBitcast(operand_shape, output_shape_with_layout)) { return absl::make_unique(operand_shape.layout()); } - if (ShapeUtil::Rank(operand_shape) == ShapeUtil::Rank(output_shape)) { + if (operand_shape.rank() == output_shape.rank()) { *operand_shape.mutable_layout() = output_layout; if (ShapeUtil::ReshapeIsBitcast(operand_shape, output_shape_with_layout)) { @@ -1048,7 +1041,7 @@ std::unique_ptr LayoutAssignment::ChooseOperandLayoutFromOutputLayout( if (instruction->opcode() == HloOpcode::kTranspose) { // Pick the operand layout that makes the transpose a bitcast. - int64 rank = ShapeUtil::Rank(instruction->shape()); + int64 rank = instruction->shape().rank(); std::vector new_minor_to_major(rank); for (int64 i = 0; i < rank; ++i) { int64 output_dim = LayoutUtil::Minor(output_layout, i); @@ -1069,11 +1062,10 @@ std::unique_ptr LayoutAssignment::ChooseOutputLayoutFromOperandLayout( int64 operand_no) { const HloInstruction* operand = user->operand(operand_no); - CHECK(ShapeUtil::IsArray(user->shape()) && - ShapeUtil::IsArray(operand->shape())); + CHECK(user->shape().IsArray() && operand->shape().IsArray()); if (!ShapeUtil::IsScalar(operand->shape()) && - ShapeUtil::Rank(operand->shape()) == ShapeUtil::Rank(user->shape()) && + operand->shape().rank() == user->shape().rank() && !instruction_can_change_layout_func_(user)) { // Assign users the same layout as the operand. return absl::make_unique(operand_layout); @@ -1086,7 +1078,7 @@ std::unique_ptr LayoutAssignment::ChooseOutputLayoutFromOperandLayout( // reshape is a bitcast when using the same layout. This may avoid copy // operations. For similar reasons, if the operand and output have the same // rank, try to match the outputs's layout to the operand. - if (ShapeUtil::Rank(operand->shape()) == 1 && + if (operand->shape().rank() == 1 && ShapeUtil::TrueRank(user->shape()) == 1) { // Don't assign a layout in case of R1 -> effective R1 reshape. return nullptr; @@ -1101,7 +1093,7 @@ std::unique_ptr LayoutAssignment::ChooseOutputLayoutFromOperandLayout( if (ShapeUtil::ReshapeIsBitcast(output_shape, operand_shape_with_layout)) { return absl::make_unique(output_shape.layout()); } - if (ShapeUtil::Rank(operand->shape()) == ShapeUtil::Rank(output_shape)) { + if (operand->shape().rank() == output_shape.rank()) { *output_shape.mutable_layout() = operand_layout; if (ShapeUtil::ReshapeIsBitcast(output_shape, operand_shape_with_layout)) { @@ -1120,7 +1112,7 @@ std::unique_ptr LayoutAssignment::ChooseOutputLayoutFromOperandLayout( if (user->opcode() == HloOpcode::kTranspose) { // Pick the user layout that makes the transpose a bitcast. - int64 rank = ShapeUtil::Rank(user->shape()); + int64 rank = user->shape().rank(); std::vector new_minor_to_major(rank); auto inverse_dimensions = InversePermutation(user->dimensions()); for (int64 i = 0; i < rank; ++i) { @@ -1196,7 +1188,7 @@ std::vector> GetArrayUsesOfBuffer( CHECK(buffer.IsArray()); std::vector> uses; for (const auto& buffer_alias : points_to_analysis.GetBufferAliases(buffer)) { - if (!ShapeUtil::IsArray(buffer_alias.instruction()->shape())) { + if (!buffer_alias.instruction()->shape().IsArray()) { continue; } // This alias must be the top-level (index == {}) of the instruction's @@ -1230,7 +1222,7 @@ Status LayoutAssignment::PropagateUseConstraintToDefs( if (ShapeUtil::IsLeafIndex(shape_layout.shape(), index)) { for (const LogicalBuffer* buffer : buffers) { if (constraints->BufferLayout(*buffer) == nullptr && - ShapeUtil::IsArray(buffer->shape())) { + buffer->shape().IsArray()) { TF_RETURN_IF_ERROR(constraints->SetBufferLayout( ShapeUtil::GetSubshape(shape_layout.shape(), index).layout(), *buffer, /*mandatory=*/true)); @@ -1241,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) { @@ -1251,13 +1260,20 @@ Status LayoutAssignment::PropagateOperandConstraint( operand_constraint.operand(), constraints)); // For array-shaped operands and user instructions try to pick a minimum cost - // layout. For example, if the operand of a elementwise instruction is - // constained to a certain layout we want the output of the instruction to + // layout. For example, if the operand of an elementwise instruction is + // constrained to a certain layout we want the output of the instruction to // have the same layout. + // + // If the user is not array-shaped, we still want to propagate the layout + // to siblings if the instruction can't change layout. This is to represent + // the information that non-layout-changing instructions should have the same + // layout for the operands with the same ranks. const HloInstruction* operand = operand_constraint.operand(); const HloInstruction* user = operand_constraint.instruction(); - if (!ShapeUtil::IsArray(operand->shape()) || - !ShapeUtil::IsArray(user->shape())) { + if (!operand->shape().IsArray()) { + return Status::OK(); + } + if (instruction_can_change_layout_func_(user) && !user->shape().IsArray()) { return Status::OK(); } @@ -1267,52 +1283,181 @@ Status LayoutAssignment::PropagateOperandConstraint( operand_constraint.operand_no())) { return Status::OK(); } - TF_ASSIGN_OR_RETURN( - const LogicalBuffer* buffer, - constraints->points_to_analysis().GetBufferDefinedAt(user, /*index=*/{})); - if (constraints->BufferLayout(*buffer) == nullptr) { - std::unique_ptr layout = ChooseOutputLayoutFromOperandLayout( - operand_constraint.shape_layout().layout(), user, - operand_constraint.operand_no()); - if (layout != nullptr) { - TF_RETURN_IF_ERROR( - constraints->SetBufferLayout(*layout, *buffer, /*mandatory=*/false)); + int64 operand_rank = operand->shape().rank(); + if (operand_rank <= 1) { + return Status::OK(); + } + + // Propagate layouts between operands of the same instruction. This is a + // constraint on non-layout-changing instructions. + if (!instruction_can_change_layout_func_(user)) { + // Make sure all siblings have the same layout as the operand. + for (int64 operand_no = 0; operand_no < user->operand_count(); + ++operand_no) { + if (user->operand(operand_no) == operand) { + continue; + } + const HloInstruction* sibling = user->operand(operand_no); + const int64 sibling_rank = sibling->shape().rank(); + if (sibling_rank <= 1) { + continue; + } + if (operand_rank != sibling_rank) { + continue; + } + const OperandLayoutConstraint* constraint = + constraints->GetOperandLayoutConstraint(user, operand_no); + if (constraint != nullptr) { + // Due to the DFS of the propagation we can end up here when operand_no + // has a layout set that hasn't been propagated yet (is still on the + // stack of layouts to propagate). + // We can continue here and leave the operands with different layouts, + // as we will either: + // - overwrite the current operand when the DFS gets back to propagating + // operand(operand_no) to its siblings + // - overwrite operand(operand_no)'s layout with a mandatory layout if + // we continue to propagate our layout to the result, and then + // backwards into all operands (if the result is an array of rank > 1) + continue; + } + TF_RETURN_IF_ERROR(constraints->SetArrayOperandLayout( + operand_constraint.shape_layout().layout(), user, operand_no, + /*mandatory=*/false)); } + TF_RETURN_IF_ERROR(ShapeUtil::ForEachSubshapeWithStatus( + user->shape(), + [&](const Shape& subshape, const ShapeIndex& shape_index) { + if (subshape.IsTuple()) { + return Status::OK(); + } + if (subshape.rank() <= 1) { + return Status::OK(); + } + + // Assign the right layout to input fusion of higher rank reduce + // operations. + if (subshape.rank() != operand->shape().rank()) { + return Status::OK(); + } + // TODO(b/67641796): Are there cases except fusion that use this code + // path? + TF_ASSIGN_OR_RETURN( + const LogicalBuffer* buffer, + constraints->points_to_analysis().GetBufferDefinedAt( + user, shape_index)); + // Make sure the output has the same layout as the operand. + const BufferLayoutConstraint* constraint = + constraints->GetBufferLayoutConstraint(*buffer); + // If we already have a constraint for the buffer it was assigned but + // hasn't propagated yet. This can happen with diamond-shaped graphs + // where one path is first evaluated in depth-first order (we're here) + // and the other path is propagated later. We don't set the layout + // here as it will always be overwritten later. + if (constraint == nullptr) { + TF_RETURN_IF_ERROR(constraints->SetBufferLayout( + operand_constraint.shape_layout().layout(), *buffer, + /*mandatory=*/false)); + } + return Status::OK(); + })); + return Status::OK(); } + TF_RETURN_IF_ERROR(ShapeUtil::ForEachSubshapeWithStatus( + user->shape(), [&](const Shape& subshape, const ShapeIndex& shape_index) { + if (subshape.IsTuple()) { + return Status::OK(); + } + if (subshape.rank() <= 1) { + return Status::OK(); + } + TF_ASSIGN_OR_RETURN( + const LogicalBuffer* buffer, + constraints->points_to_analysis().GetBufferDefinedAt(user, + shape_index)); + if (constraints->BufferLayout(*buffer) == nullptr || + !constraints->GetBufferLayoutConstraint(*buffer)->mandatory()) { + std::unique_ptr layout = ChooseOutputLayoutFromOperandLayout( + operand_constraint.shape_layout().layout(), user, + operand_constraint.operand_no()); + if (layout != nullptr) { + TF_RETURN_IF_ERROR(constraints->SetBufferLayout( + *layout, *buffer, + /*mandatory=*/user->opcode() == HloOpcode::kReduce, + /*dfs=*/InstructionShouldPropagateDepthFirst(*user))); + } + } + return Status::OK(); + })); return Status::OK(); } -Status LayoutAssignment::PropagateBufferConstraint( +Status LayoutAssignment::PropagateBufferConstraintToOperands( const BufferLayoutConstraint& buffer_constraint, LayoutConstraints* constraints) { - // Only propagate array layouts. + VLOG(5) << "PropagateBufferConstraintToOperands: " + << buffer_constraint.ToString(); const LogicalBuffer& buffer = buffer_constraint.buffer(); - if (!buffer.IsArray()) { + + const HloInstruction* instruction = buffer.instruction(); + if (IsAtMostRank1(instruction->shape())) { return Status::OK(); } - // If this buffer is the result of an array-shaped op (as opposed to an array - // element in a tuple) try to propagate the layout to its operands. - if (buffer.IsTopLevel()) { - const HloInstruction* instruction = buffer.instruction(); - // Propagate the def-constraint on an instruction to the use-constraints on - // its operands (use-def propagation). - for (int64 operand_no = 0; operand_no < instruction->operand_count(); - ++operand_no) { - if (constraints->OperandLayout(instruction, operand_no) == nullptr && - ShapeUtil::IsArray(instruction->operand(operand_no)->shape())) { + for (int64 operand_no = 0; operand_no < instruction->operand_count(); + ++operand_no) { + const HloInstruction* operand = instruction->operand(operand_no); + if (IsAtMostRank1(operand->shape())) { + continue; + } + if (!instruction_can_change_layout_func_(instruction)) { + // Copy the layout to the operand. + if (buffer.IsArray() && operand->shape().IsArray() && + operand->shape().rank() == + LayoutUtil::MinorToMajor(buffer_constraint.layout()).size()) { + TF_RETURN_IF_ERROR(constraints->SetArrayOperandLayout( + buffer_constraint.layout(), instruction, operand_no, + /*mandatory=*/true)); + } + } else { + if (!buffer.IsTopLevel() || + !instruction->operand(operand_no)->shape().IsArray()) { + continue; // Don't touch buffers that are internal to a tuple. + } + VLOG(6) << "Propagating constraint to operand " << operand_no << " of " + << instruction->ToShortString(); + // Assign a layout if there is no constraint already. + const OperandLayoutConstraint* constraint = + constraints->GetOperandLayoutConstraint(instruction, operand_no); + if (constraint == nullptr || !constraint->mandatory()) { std::unique_ptr operand_layout = ChooseOperandLayoutFromOutputLayout(buffer_constraint.layout(), instruction, operand_no); if (operand_layout != nullptr) { TF_RETURN_IF_ERROR(constraints->SetArrayOperandLayout( - *operand_layout, instruction, operand_no, /*mandatory=*/true)); + *operand_layout, instruction, operand_no, /*mandatory=*/false, + /*dfs=*/InstructionShouldPropagateDepthFirst(*instruction))); } + } else { + VLOG(6) << "Operand already has a constraint " + << constraint->ToString(); } } } - return PropagateBufferConstraintToUses(buffer_constraint, constraints); + return Status::OK(); +} + +Status LayoutAssignment::PropagateBufferConstraint( + const BufferLayoutConstraint& buffer_constraint, + LayoutConstraints* constraints) { + // Only propagate array layouts. + const LogicalBuffer& buffer = buffer_constraint.buffer(); + if (!buffer.IsArray()) { + return Status::OK(); + } + TF_RETURN_IF_ERROR( + PropagateBufferConstraintToUses(buffer_constraint, constraints)); + return PropagateBufferConstraintToOperands(buffer_constraint, constraints); } Status LayoutAssignment::PropagateBufferConstraintToUses( @@ -1340,12 +1485,12 @@ Status LayoutAssignment::PropagateBufferConstraintToUses( } Status LayoutAssignment::PropagateResultConstraint( - const ResultLayoutConstraint& result_constraint, + const ResultLayoutConstraint& layout_constraint, LayoutConstraints* constraints) { // Propagate the use constraint of the root instruction up to the logical // buffers which make up the result. return PropagateUseConstraintToDefs( - result_constraint.shape_layout(), + layout_constraint.shape_layout(), constraints->computation()->root_instruction(), constraints); } @@ -1361,7 +1506,7 @@ StatusOr InferArrayLayout( // This function should only be called for array shapes which don't yet have // layouts. const Shape& subshape = ShapeUtil::GetSubshape(instruction->shape(), index); - TF_RET_CHECK(ShapeUtil::IsArray(subshape)); + TF_RET_CHECK(subshape.IsArray()); TF_RET_CHECK(!subshape.has_layout()); // The instruction should not define the buffer at this index. @@ -1440,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()); } } @@ -1479,7 +1625,7 @@ Status LayoutAssignment::AssignLayouts(const LayoutConstraints& constraints, for (const LogicalBuffer* buffer : constraints.points_to_analysis().GetBuffersDefinedByInstruction( instruction)) { - if (!ShapeUtil::IsArray(buffer->shape())) { + if (!buffer->shape().IsArray()) { continue; } @@ -1503,7 +1649,7 @@ Status LayoutAssignment::AssignLayouts(const LayoutConstraints& constraints, TF_RETURN_IF_ERROR(ShapeUtil::ForEachMutableSubshapeWithStatus( instruction->mutable_shape(), [instruction, &constraints](Shape* subshape, const ShapeIndex& index) { - if (subshape->has_layout() || !ShapeUtil::IsArray(*subshape)) { + if (subshape->has_layout() || !subshape->IsArray()) { return Status::OK(); } // Set Layout of subshape to match layout of LogicalBuffer which @@ -1864,6 +2010,7 @@ bool LayoutAssignment::InstructionCanChangeLayout( switch (instruction->opcode()) { case HloOpcode::kAbs: case HloOpcode::kAdd: + case HloOpcode::kAddDependency: case HloOpcode::kAnd: case HloOpcode::kAtan2: case HloOpcode::kBitcastConvert: @@ -1875,7 +2022,7 @@ bool LayoutAssignment::InstructionCanChangeLayout( case HloOpcode::kConditional: case HloOpcode::kConvert: case HloOpcode::kCos: - case HloOpcode::kCrossReplicaSum: + case HloOpcode::kAllReduce: case HloOpcode::kAllToAll: case HloOpcode::kCollectivePermute: case HloOpcode::kDivide: @@ -1956,10 +2103,21 @@ bool LayoutAssignment::InstructionCanChangeLayout( case HloOpcode::kTrace: case HloOpcode::kTranspose: case HloOpcode::kTuple: + case HloOpcode::kGetDimensionSize: return true; } } +/* static */ +bool LayoutAssignment::IsAtMostRank1(const Shape& shape) { + if (shape.IsArray()) { + return shape.rank() <= 1; + } + return absl::c_all_of(shape.tuple_shapes(), [](const Shape& subshape) { + return IsAtMostRank1(subshape); + }); +} + Status LayoutAssignment::Init() { computation_layouts_.clear(); *entry_computation_layout_ = saved_entry_computation_layout_; @@ -1975,7 +2133,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 cb56f4cd19ded036ef521a579eb7d6ea7f3b6268..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 @@ -315,6 +315,10 @@ class LayoutAssignment : public HloModulePass { // rank as the output to have the same layout as the output. static bool InstructionCanChangeLayout(const HloInstruction* instruction); + // In case of an array shape returns true iff it is at most rank 1. In case of + // a tuple shape returns true iff all leaf shapes are at most rank 1. + static bool IsAtMostRank1(const Shape& shape); + protected: // These methods, invoked by PropagateConstraints, propagate a layout // constraint to its neighbors (i.e. operands and users) in order to minimize @@ -362,7 +366,7 @@ class LayoutAssignment : public HloModulePass { // `user` that minimizes its cost on that operand. Returns null if it can't // decide the best layout. // Precondition: `user` and the operand are array-shaped. - std::unique_ptr ChooseOutputLayoutFromOperandLayout( + virtual std::unique_ptr ChooseOutputLayoutFromOperandLayout( const Layout& operand_layout, const HloInstruction* user, int64 operand_no); @@ -408,6 +412,10 @@ class LayoutAssignment : public HloModulePass { // required for correctness. Status PropagateConstraints(LayoutConstraints* constraints); + Status PropagateBufferConstraintToOperands( + const BufferLayoutConstraint& buffer_constraint, + LayoutConstraints* constraints); + // Check that all layouts in the module have been set and satisfy all // necessary conditions. Status CheckLayouts(HloModule* module); diff --git a/tensorflow/compiler/xla/service/layout_assignment_test.cc b/tensorflow/compiler/xla/service/layout_assignment_test.cc index a831751fa96f8cef233e16fe02378ac036efc8ab..c8cf3c47d380012fdb0206c0d20d67e6a13017ae 100644 --- a/tensorflow/compiler/xla/service/layout_assignment_test.cc +++ b/tensorflow/compiler/xla/service/layout_assignment_test.cc @@ -27,42 +27,41 @@ limitations under the License. #include "tensorflow/compiler/xla/service/computation_layout.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_parser.h" +#include "tensorflow/compiler/xla/service/pattern_matcher.h" +#include "tensorflow/compiler/xla/service/pattern_matcher_gmock.h" #include "tensorflow/compiler/xla/shape_layout.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_verified_test_base.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" #include "tensorflow/compiler/xla/tests/test_utils.h" #include "tensorflow/compiler/xla/util.h" #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/lib/core/status_test_util.h" -namespace op = xla::testing::opcode_matchers; - namespace xla { namespace { +namespace m = xla::match; using ::testing::ElementsAre; -class LayoutAssignmentTest : public HloVerifiedTestBase { +class LayoutAssignmentTest : public HloTestBase { protected: - void AssignLayouts(HloModule* module, - ComputationLayout* entry_computation_layout, + void AssignLayouts(HloModule* m, ComputationLayout* entry_computation_layout, ChannelLayoutConstraints* channel_constraints = nullptr) { LayoutAssignment layout_assignment( entry_computation_layout, LayoutAssignment::InstructionCanChangeLayout, /*channel_constraints=*/channel_constraints); - EXPECT_IS_OK(layout_assignment.Run(module).status()); + EXPECT_IS_OK(layout_assignment.Run(m).status()); } - std::vector LayoutOf(HloModule* module, absl::string_view name) { + std::vector LayoutOf(HloModule* m, absl::string_view name) { auto minor_to_major = - FindInstruction(module, name)->shape().layout().minor_to_major(); + FindInstruction(m, name)->shape().layout().minor_to_major(); return std::vector(minor_to_major.begin(), minor_to_major.end()); } @@ -91,7 +90,7 @@ class LayoutAssignmentTest : public HloVerifiedTestBase { TEST_F(LayoutAssignmentTest, ComputationLayout) { // Verify the layouts of the root and parameter instructions of a computation // match the ComputationLayout for two different layouts. - std::vector> minor_to_majors = {{0, 1}, {1, 0}}; + std::vector> minor_to_majors = {{0, 1}, {1, 0}}; for (auto& minor_to_major : minor_to_majors) { auto builder = HloComputation::Builder(TestName()); Shape ashape = ShapeUtil::MakeShape(F32, {42, 12}); @@ -101,8 +100,8 @@ TEST_F(LayoutAssignmentTest, ComputationLayout) { HloInstruction::CreateParameter(1, ashape, "param1")); auto add = builder.AddInstruction( HloInstruction::CreateBinary(ashape, HloOpcode::kAdd, param0, param1)); - auto module = CreateNewModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); + auto m = CreateNewVerifiedModule(); + HloComputation* computation = m->AddEntryComputation(builder.Build()); Layout layout = LayoutUtil::MakeLayout(minor_to_major); Shape shape(ashape); @@ -113,7 +112,7 @@ TEST_F(LayoutAssignmentTest, ComputationLayout) { *computation_layout.mutable_parameter_layout(0) = shape_layout; *computation_layout.mutable_parameter_layout(1) = shape_layout; *computation_layout.mutable_result_layout() = shape_layout; - AssignLayouts(module, &computation_layout); + AssignLayouts(m.get(), &computation_layout); EXPECT_TRUE(LayoutUtil::Equal(layout, param0->shape().layout())); EXPECT_TRUE(LayoutUtil::Equal(layout, param1->shape().layout())); EXPECT_TRUE(LayoutUtil::Equal(layout, add->shape().layout())); @@ -131,8 +130,8 @@ TEST_F(LayoutAssignmentTest, ComputationLayoutMixedLayout) { HloInstruction::CreateParameter(1, ashape, "param1")); builder.AddInstruction( HloInstruction::CreateBinary(ashape, HloOpcode::kAdd, param0, param1)); - auto module = CreateNewModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); + auto m = CreateNewVerifiedModule(); + HloComputation* computation = m->AddEntryComputation(builder.Build()); Layout col_major_layout = LayoutUtil::MakeLayout({1, 0}); Shape col_major_shape(ashape); @@ -149,7 +148,7 @@ TEST_F(LayoutAssignmentTest, ComputationLayoutMixedLayout) { *computation_layout.mutable_parameter_layout(1) = row_major; *computation_layout.mutable_result_layout() = col_major; - AssignLayouts(module, &computation_layout); + AssignLayouts(m.get(), &computation_layout); EXPECT_TRUE(LayoutUtil::Equal(col_major_layout, param0->shape().layout())); EXPECT_TRUE(LayoutUtil::Equal(row_major_layout, param1->shape().layout())); EXPECT_TRUE(LayoutUtil::Equal( @@ -160,7 +159,7 @@ TEST_F(LayoutAssignmentTest, FusionInstruction) { // Verify that the layout of the fused parameters in a fusion instruction // match that of the fusion operands. Other fused instructions should have no // layout. - std::vector> minor_to_majors = {{0, 1}, {1, 0}}; + std::vector> minor_to_majors = {{0, 1}, {1, 0}}; for (auto& minor_to_major : minor_to_majors) { auto builder = HloComputation::Builder(TestName()); auto constant_literal1 = LiteralUtil::CreateR2WithLayout( @@ -180,8 +179,8 @@ TEST_F(LayoutAssignmentTest, FusionInstruction) { auto negate2 = builder.AddInstruction( HloInstruction::CreateUnary(ashape, HloOpcode::kNegate, negate1)); - auto module = CreateNewModule(); - HloComputation* computation = module->AddEntryComputation(builder.Build()); + auto m = CreateNewVerifiedModule(); + HloComputation* computation = m->AddEntryComputation(builder.Build()); auto fusion = computation->CreateFusionInstruction( {negate2, negate1, add}, HloInstruction::FusionKind::kLoop); @@ -194,7 +193,7 @@ TEST_F(LayoutAssignmentTest, FusionInstruction) { ComputationLayout computation_layout(computation->ComputeProgramShape()); *computation_layout.mutable_result_layout() = shape_layout; - AssignLayouts(module, &computation_layout); + AssignLayouts(m.get(), &computation_layout); EXPECT_TRUE(LayoutUtil::Equal( layout, fusion->fused_parameter(0)->shape().layout())); @@ -229,13 +228,13 @@ TEST_F(LayoutAssignmentTest, TupleLayout) { auto negate = builder.AddInstruction(HloInstruction::CreateUnary( constant0->shape(), HloOpcode::kNegate, get_element0)); - auto module = CreateNewModule(); - module->AddEntryComputation(builder.Build()); + auto m = CreateNewVerifiedModule(); + m->AddEntryComputation(builder.Build()); ComputationLayout computation_layout( - module->entry_computation()->ComputeProgramShape()); + m->entry_computation()->ComputeProgramShape()); - AssignLayouts(module, &computation_layout); + AssignLayouts(m.get(), &computation_layout); EXPECT_TRUE( LayoutUtil::LayoutsInShapesEqual(constant0->shape(), constant1->shape())); @@ -267,17 +266,17 @@ TEST_F(LayoutAssignmentTest, TupleSelect) { auto select = builder.AddInstruction(HloInstruction::CreateTernary( tuple0->shape(), HloOpcode::kTupleSelect, pred, tuple0, tuple1)); - auto module = CreateNewModule(); - module->AddEntryComputation(builder.Build()); + auto m = CreateNewVerifiedModule(); + m->AddEntryComputation(builder.Build()); ComputationLayout computation_layout( - module->entry_computation()->ComputeProgramShape()); + m->entry_computation()->ComputeProgramShape()); Shape result_shape = ShapeUtil::MakeTupleShape({constant0->shape(), constant1->shape()}); TF_CHECK_OK(computation_layout.mutable_result_layout()->CopyLayoutFromShape( result_shape)); - AssignLayouts(module, &computation_layout); + AssignLayouts(m.get(), &computation_layout); EXPECT_TRUE(LayoutUtil::LayoutsInShapesEqual(result_shape, select->shape())); } @@ -302,11 +301,11 @@ TEST_F(LayoutAssignmentTest, ConflictingLayoutTuple) { auto nested_tuple = builder.AddInstruction( HloInstruction::CreateTuple({inner_tuple, inner_tuple})); - auto module = CreateNewModule(); - module->AddEntryComputation(builder.Build()); + auto m = CreateNewVerifiedModule(); + m->AddEntryComputation(builder.Build()); ComputationLayout computation_layout( - module->entry_computation()->ComputeProgramShape()); + m->entry_computation()->ComputeProgramShape()); Shape result_shape = nested_tuple->shape(); *ShapeUtil::GetMutableSubshape(&result_shape, /*index=*/{0, 0}) = ShapeUtil::MakeShapeWithLayout(F32, {2, 2}, {1, 0}); @@ -316,7 +315,7 @@ TEST_F(LayoutAssignmentTest, ConflictingLayoutTuple) { result_shape)); LayoutAssignment layout_assignment(&computation_layout); - AssignLayouts(module, &computation_layout); + AssignLayouts(m.get(), &computation_layout); // Layout assignment should have deep copied the result of the computation to // address the layout conflict. This results in several Tuple() and @@ -329,12 +328,11 @@ TEST_F(LayoutAssignmentTest, ConflictingLayoutTuple) { // %tuple.1 = Tuple(%copy) layout=({0,1}) // %tuple.2 = Tuple(%tuple.0, %tuple.1) layout=(({1,0}), ({0,1})) // - EXPECT_TRUE( - AlgebraicSimplifier(/*is_layout_sensitive=*/true, - [](const Shape&, const Shape&) { return false; }) - .Run(module) - .ValueOrDie()); - HloInstruction* root = module->entry_computation()->root_instruction(); + AlgebraicSimplifierOptions options( + [](const Shape&, const Shape&) { return false; }); + options.set_is_layout_sensitive(true); + EXPECT_TRUE(AlgebraicSimplifier(options).Run(m.get()).ValueOrDie()); + HloInstruction* root = m->entry_computation()->root_instruction(); // Verify layout of the root and the root's operands. EXPECT_TRUE(ShapeUtil::Equal(result_shape, root->shape())); EXPECT_TRUE(ShapeUtil::Equal(ShapeUtil::GetSubshape(result_shape, {0}), @@ -344,7 +342,8 @@ TEST_F(LayoutAssignmentTest, ConflictingLayoutTuple) { // Verify the structure of the HLO graph. EXPECT_THAT(root, - op::Tuple(op::Tuple(constant), op::Tuple(op::Copy(constant)))); + GmockMatch(m::Tuple(m::Tuple(m::Op().Is(constant)), + m::Tuple(m::Copy(m::Op().Is(constant)))))); } TEST_F(LayoutAssignmentTest, ElementwiseAndReshape) { @@ -361,9 +360,8 @@ TEST_F(LayoutAssignmentTest, ElementwiseAndReshape) { auto tanh = builder.AddInstruction( HloInstruction::CreateUnary(bshape, HloOpcode::kTanh, reshape)); - auto module = CreateNewModule(); - HloComputation* computation = - module->AddEntryComputation(builder.Build(tanh)); + auto m = CreateNewVerifiedModule(); + HloComputation* computation = m->AddEntryComputation(builder.Build(tanh)); Shape ashape_with_layout(ashape); Shape bshape_with_layout(bshape); @@ -374,7 +372,7 @@ TEST_F(LayoutAssignmentTest, ElementwiseAndReshape) { *computation_layout.mutable_parameter_layout(0) = ShapeLayout(ashape_with_layout); *computation_layout.mutable_result_layout() = ShapeLayout(bshape_with_layout); - AssignLayouts(module, &computation_layout); + AssignLayouts(m.get(), &computation_layout); auto log_minor_to_major = AsInt64Slice(log->shape().layout().minor_to_major()); @@ -403,8 +401,8 @@ TEST_F(LayoutAssignmentTest, ElementwiseAndTranspose) { HloInstruction::CreateTranspose(bshape, log, {1, 0})); auto tanh = builder.AddInstruction( HloInstruction::CreateUnary(bshape, HloOpcode::kTanh, transpose)); - auto module = CreateNewModule(); - auto computation = module->AddEntryComputation(builder.Build(tanh)); + auto m = CreateNewVerifiedModule(); + auto computation = m->AddEntryComputation(builder.Build(tanh)); Shape ashape_with_layout(ashape); Shape bshape_with_layout(bshape); @@ -415,7 +413,7 @@ TEST_F(LayoutAssignmentTest, ElementwiseAndTranspose) { *computation_layout.mutable_parameter_layout(0) = ShapeLayout(ashape_with_layout); *computation_layout.mutable_result_layout() = ShapeLayout(bshape_with_layout); - AssignLayouts(module, &computation_layout); + AssignLayouts(m.get(), &computation_layout); EXPECT_TRUE( LayoutUtil::Equal(ashape_with_layout.layout(), log->shape().layout())); @@ -439,9 +437,9 @@ TEST_F(LayoutAssignmentTest, BroadcastAndTranspose) { HloInstruction::CreateBroadcast(bshape, param, {1, 2})); auto transpose = builder.AddInstruction( HloInstruction::CreateTranspose(cshape, broadcast, {2, 1, 0})); - auto module = CreateNewModule(); + auto m = CreateNewVerifiedModule(); HloComputation* computation = - module->AddEntryComputation(builder.Build(transpose)); + m->AddEntryComputation(builder.Build(transpose)); Shape input_shape_with_layout(ashape); Shape output_shape_with_layout(cshape); @@ -454,7 +452,7 @@ TEST_F(LayoutAssignmentTest, BroadcastAndTranspose) { ShapeLayout(input_shape_with_layout); *computation_layout.mutable_result_layout() = ShapeLayout(output_shape_with_layout); - AssignLayouts(module, &computation_layout); + AssignLayouts(m.get(), &computation_layout); EXPECT_THAT(broadcast->shape().layout().minor_to_major(), ElementsAre(0, 1, 2)); @@ -488,9 +486,8 @@ TEST_F(LayoutAssignmentTest, ReshapeOperandHasMultipleUsers) { HloInstruction::CreateBroadcast(f32_234, tanh, {1, 2})); auto tuple = builder.AddInstruction( HloInstruction::CreateTuple({transpose, broadcast2})); - auto module = CreateNewModule(); - HloComputation* computation = - module->AddEntryComputation(builder.Build(tuple)); + auto m = CreateNewVerifiedModule(); + HloComputation* computation = m->AddEntryComputation(builder.Build(tuple)); ComputationLayout computation_layout(computation->ComputeProgramShape()); Shape param_shape_with_layout(f32_4); @@ -507,7 +504,7 @@ TEST_F(LayoutAssignmentTest, ReshapeOperandHasMultipleUsers) { *computation_layout.mutable_result_layout() = ShapeLayout(ShapeUtil::MakeTupleShape( {transpose_shape_with_layout, broadcast2_shape_with_layout})); - AssignLayouts(module, &computation_layout); + AssignLayouts(m.get(), &computation_layout); EXPECT_THAT(broadcast->shape().layout().minor_to_major(), ElementsAre(0, 1)); EXPECT_THAT(transpose->shape().layout().minor_to_major(), ElementsAre(1, 0)); @@ -531,8 +528,7 @@ class OperandsMustBeTheSameLayoutAssignment : public LayoutAssignment { for (int64 operand_no = 0; operand_no < instruction->operand_count(); ++operand_no) { const HloInstruction* operand = instruction->operand(operand_no); - if (ShapeUtil::Rank(instruction->shape()) != - ShapeUtil::Rank(operand->shape())) { + if (instruction->shape().rank() != operand->shape().rank()) { continue; } TF_RETURN_IF_ERROR(constraints->SetArrayOperandLayout( @@ -558,9 +554,8 @@ TEST_F(LayoutAssignmentTest, MakeOperandsTheSame) { HloInstruction::CreateConcatenate(bshape, {param0, param1}, 1)); auto reshape = builder.AddInstruction( HloInstruction::CreateReshape(cshape, concatenate)); - auto module = CreateNewModule(); - HloComputation* computation = - module->AddEntryComputation(builder.Build(reshape)); + auto m = CreateNewVerifiedModule(); + HloComputation* computation = m->AddEntryComputation(builder.Build(reshape)); Shape param0_shape_with_layout(ashape); Shape param1_shape_with_layout(ashape); @@ -573,7 +568,7 @@ TEST_F(LayoutAssignmentTest, MakeOperandsTheSame) { *computation_layout.mutable_parameter_layout(1) = ShapeLayout(param1_shape_with_layout); OperandsMustBeTheSameLayoutAssignment layout_assignment(&computation_layout); - EXPECT_IS_OK(layout_assignment.Run(module).status()); + EXPECT_IS_OK(layout_assignment.Run(m.get()).status()); EXPECT_EQ(HloOpcode::kCopy, concatenate->operand(0)->opcode()); EXPECT_THAT(concatenate->operand(0)->shape().layout().minor_to_major(), @@ -593,11 +588,11 @@ TEST_F(LayoutAssignmentTest, TransposeToBitcastFromOperand) { HloInstruction::CreateParameter(0, input_shape_with_layout, "param")); auto transpose = builder.AddInstruction(HloInstruction::CreateTranspose( ShapeUtil::MakeShape(F32, {6, 7, 3, 5}), param, {2, 3, 0, 1})); - auto module = CreateNewModule(); + auto m = CreateNewVerifiedModule(); HloComputation* computation = - module->AddEntryComputation(builder.Build(transpose)); + m->AddEntryComputation(builder.Build(transpose)); ComputationLayout computation_layout(computation->ComputeProgramShape()); - AssignLayouts(module, &computation_layout); + AssignLayouts(m.get(), &computation_layout); EXPECT_TRUE(ShapeUtil::TransposeIsBitcast(transpose->operand(0)->shape(), transpose->shape(), {2, 3, 0, 1})); } @@ -611,11 +606,11 @@ TEST_F(LayoutAssignmentTest, TransposeToBitcastToUser) { HloInstruction::CreateBroadcast(input_shape, constant, {})); auto transpose = builder.AddInstruction(HloInstruction::CreateTranspose( ShapeUtil::MakeShape(F32, {6, 7, 3, 5}), broadcast, {2, 3, 0, 1})); - auto module = CreateNewModule(); + auto m = CreateNewVerifiedModule(); HloComputation* computation = - module->AddEntryComputation(builder.Build(transpose)); + m->AddEntryComputation(builder.Build(transpose)); ComputationLayout computation_layout(computation->ComputeProgramShape()); - AssignLayouts(module, &computation_layout); + AssignLayouts(m.get(), &computation_layout); EXPECT_TRUE(ShapeUtil::TransposeIsBitcast(transpose->operand(0)->shape(), transpose->shape(), {2, 3, 0, 1})); } @@ -681,12 +676,12 @@ TEST_F(LayoutAssignmentTest, TransposeWithinFusionDoesNotCrash) { } )"; - ParseAndVerifyModule(module_str); - + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr m, + ParseAndReturnVerifiedModule(module_str)); std::unique_ptr compiled_module = backend() .compiler() - ->RunHloPasses(module().Clone(), backend().default_stream_executor(), + ->RunHloPasses(m->Clone(), backend().default_stream_executor(), /*device_allocator=*/nullptr) .ConsumeValueOrDie(); @@ -721,9 +716,10 @@ TEST_F(LayoutAssignmentTest, GTEInheritsLayoutFromOperand) { } )"; - ParseAndVerifyModule(module_str); + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr m, + ParseAndReturnVerifiedModule(module_str)); ComputationLayout computation_layout( - module().entry_computation()->ComputeProgramShape()); + m->entry_computation()->ComputeProgramShape()); Shape param_shape = ShapeUtil::MakeTupleShape( {ShapeUtil::MakeShapeWithLayout(F32, {2, 2, 2}, {0, 1, 2}), ShapeUtil::MakeTupleShape({ @@ -735,19 +731,19 @@ TEST_F(LayoutAssignmentTest, GTEInheritsLayoutFromOperand) { param_shape)); computation_layout.mutable_result_layout()->ResetLayout( LayoutUtil::MakeLayout({2, 1, 0})); - AssignLayouts(&module(), &computation_layout); + AssignLayouts(m.get(), &computation_layout); - EXPECT_THAT(LayoutOf(&module(), "gte0"), ElementsAre(0, 1, 2)); - EXPECT_THAT(LayoutOf(&module(), "gte1a"), ElementsAre(1, 2, 0)); - EXPECT_THAT(LayoutOf(&module(), "gte1b"), ElementsAre(2, 0, 1)); - EXPECT_THAT(LayoutOf(&module(), "fresult"), ElementsAre(2, 1, 0)); - EXPECT_THAT(FindInstruction(&module(), "gte1") + EXPECT_THAT(LayoutOf(m.get(), "gte0"), ElementsAre(0, 1, 2)); + EXPECT_THAT(LayoutOf(m.get(), "gte1a"), ElementsAre(1, 2, 0)); + EXPECT_THAT(LayoutOf(m.get(), "gte1b"), ElementsAre(2, 0, 1)); + EXPECT_THAT(LayoutOf(m.get(), "fresult"), ElementsAre(2, 1, 0)); + EXPECT_THAT(FindInstruction(m.get(), "gte1") ->shape() .tuple_shapes(0) .layout() .minor_to_major(), ElementsAre(1, 2, 0)); - EXPECT_THAT(FindInstruction(&module(), "gte1") + EXPECT_THAT(FindInstruction(m.get(), "gte1") ->shape() .tuple_shapes(1) .layout() @@ -757,7 +753,7 @@ TEST_F(LayoutAssignmentTest, GTEInheritsLayoutFromOperand) { TEST_F(LayoutAssignmentTest, ConditionalAsymmetricLayout) { auto builder = HloComputation::Builder(TestName()); - auto module = CreateNewModule(); + auto m = CreateNewVerifiedModule(); Shape shape = ShapeUtil::MakeShape(F32, {128, 8}); Shape tshape = ShapeUtil::MakeTupleShape({shape, shape}); Shape result_tshape = ShapeUtil::MakeTupleShape({shape}); @@ -784,7 +780,7 @@ TEST_F(LayoutAssignmentTest, ConditionalAsymmetricLayout) { true_builder.AddInstruction(HloInstruction::CreateTuple({add})); } HloComputation* true_computation = - module->AddEmbeddedComputation(true_builder.Build()); + m->AddEmbeddedComputation(true_builder.Build()); auto false_builder = HloComputation::Builder(TestName() + "_FalseBranch"); { @@ -800,14 +796,14 @@ TEST_F(LayoutAssignmentTest, ConditionalAsymmetricLayout) { false_builder.AddInstruction(HloInstruction::CreateTuple({infeed_data})); } HloComputation* false_computation = - module->AddEmbeddedComputation(false_builder.Build()); + m->AddEmbeddedComputation(false_builder.Build()); builder.AddInstruction(HloInstruction::CreateConditional( result_tshape, pred, tuple, true_computation, tuple, false_computation)); - HloComputation* computation = module->AddEntryComputation(builder.Build()); + HloComputation* computation = m->AddEntryComputation(builder.Build()); ComputationLayout computation_layout(computation->ComputeProgramShape()); - AssignLayouts(module, &computation_layout); + AssignLayouts(m.get(), &computation_layout); const HloInstruction* true_root = true_computation->root_instruction(); const HloInstruction* false_root = false_computation->root_instruction(); @@ -828,13 +824,13 @@ TEST_F(LayoutAssignmentTest, InternalErrorOnBitcast) { {{1.0, 2.0}, {3.0, 4.0}}, LayoutUtil::MakeLayout({0, 1})))); builder.AddInstruction(HloInstruction::CreateUnary( constant0->shape(), HloOpcode::kBitcast, constant0)); - auto module = CreateNewModule(); - module->AddEntryComputation(builder.Build()); + auto m = CreateNewVerifiedModule(); + m->AddEntryComputation(builder.Build()); ComputationLayout computation_layout( - module->entry_computation()->ComputeProgramShape()); + m->entry_computation()->ComputeProgramShape()); LayoutAssignment layout_assignment(&computation_layout); - Status error_status = layout_assignment.Run(module).status(); + Status error_status = layout_assignment.Run(m.get()).status(); EXPECT_FALSE(error_status.ok()); EXPECT_THAT( error_status.error_message(), @@ -850,20 +846,21 @@ TEST_F(LayoutAssignmentTest, ChannelLayoutMismatch) { ENTRY entry_computation { param = (f32[2,2]) parameter(0) gte = f32[2,2] get-tuple-element(param), index=0 - token = token[] after-all() - recv = (f32[2,2], u32[], token[]) recv(token), channel_id=1, sharding={maximal device=1} + token0 = token[] after-all() + recv = (f32[2,2], u32[], token[]) recv(token0), channel_id=1, sharding={maximal device=1} recv-done = (f32[2,2], token[]) recv-done(recv), channel_id=1, sharding={maximal device=1} ROOT root = f32[2,2] get-tuple-element(recv-done), index=0 - send = (f32[2,2], u32[], token[]) send(gte, token), channel_id=1, + send = (f32[2,2], u32[], token[]) send(gte, token0), channel_id=1, sharding={maximal device=0} send-done = token[] send-done(send), channel_id=1, sharding={maximal device=0} } )"; - ParseAndVerifyModule(module_str); + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr m, + ParseAndReturnVerifiedModule(module_str)); ComputationLayout computation_layout( - module().entry_computation()->ComputeProgramShape()); + m->entry_computation()->ComputeProgramShape()); Shape param_shape = ShapeUtil::MakeTupleShape( {ShapeUtil::MakeShapeWithLayout(F32, {2, 2}, {0, 1})}); TF_ASSERT_OK( @@ -873,12 +870,12 @@ TEST_F(LayoutAssignmentTest, ChannelLayoutMismatch) { LayoutUtil::MakeLayout({1, 0})); ChannelLayoutConstraints channel_constraints; - AssignLayouts(&module(), &computation_layout, &channel_constraints); + AssignLayouts(m.get(), &computation_layout, &channel_constraints); - EXPECT_THAT(LayoutOf(&module(), "gte"), ElementsAre(0, 1)); - EXPECT_THAT(LayoutOf(&module(), "root"), ElementsAre(1, 0)); + EXPECT_THAT(LayoutOf(m.get(), "gte"), ElementsAre(0, 1)); + EXPECT_THAT(LayoutOf(m.get(), "root"), ElementsAre(1, 0)); EXPECT_TRUE(ShapeUtil::Equal( - ShapeUtil::GetSubshape(FindInstruction(&module(), "send")->shape(), {0}), + ShapeUtil::GetSubshape(FindInstruction(m.get(), "send")->shape(), {0}), ShapeUtil::MakeShapeWithLayout(F32, {2, 2}, {1, 0}))); } @@ -896,18 +893,18 @@ TEST_F(LayoutAssignmentTest, AllReduceLayoutMissmatch) { ENTRY entry_computation { param = (f32[2,2]) parameter(0) gte = f32[2,2] get-tuple-element(param), index=0 - ar.0 = f32[2,2] cross-replica-sum(gte), - all_reduce_id=0, replica_groups={{0}}, to_apply=add, + ar.0 = f32[2,2] all-reduce(gte), + all_reduce_id=1, replica_groups={{0}}, to_apply=add, sharding={maximal device=0} - const = f32[2,2] constant(f32[2,2]{{0,1},{2,3}}) - ROOT ar.1 = f32[2,2] cross-replica-sum(const), - all_reduce_id=0, replica_groups={{0}}, to_apply=add, + const = f32[2,2] constant({{0,1},{2,3}}) + ROOT ar.1 = f32[2,2] all-reduce(const), + all_reduce_id=1, replica_groups={{0}}, to_apply=add, sharding={maximal device=1} })"; - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr m, ParseAndReturnVerifiedModule(module_str)); ComputationLayout computation_layout( - module->entry_computation()->ComputeProgramShape()); + m->entry_computation()->ComputeProgramShape()); Shape param_shape = ShapeUtil::MakeTupleShape( {ShapeUtil::MakeShapeWithLayout(F32, {2, 2}, {0, 1})}); TF_ASSERT_OK( @@ -917,12 +914,12 @@ TEST_F(LayoutAssignmentTest, AllReduceLayoutMissmatch) { LayoutUtil::MakeLayout({1, 0})); ChannelLayoutConstraints channel_constraints; - AssignLayouts(module.get(), &computation_layout, &channel_constraints); + AssignLayouts(m.get(), &computation_layout, &channel_constraints); - EXPECT_THAT(LayoutOf(module.get(), "gte"), ElementsAre(0, 1)); - EXPECT_THAT(LayoutOf(module.get(), "ar.0"), ElementsAre(0, 1)); - EXPECT_THAT(LayoutOf(module.get(), "ar.1"), ElementsAre(0, 1)); - const HloInstruction* root = module->entry_computation()->root_instruction(); + EXPECT_THAT(LayoutOf(m.get(), "gte"), ElementsAre(0, 1)); + EXPECT_THAT(LayoutOf(m.get(), "ar.0"), ElementsAre(0, 1)); + EXPECT_THAT(LayoutOf(m.get(), "ar.1"), ElementsAre(0, 1)); + const HloInstruction* root = m->entry_computation()->root_instruction(); EXPECT_THAT(root->shape().layout().minor_to_major(), ElementsAre(1, 0)); } @@ -938,19 +935,22 @@ TEST_F(LayoutAssignmentTest, CopySliceOperandToAvoidImplicitLayoutChange) { } )"; - ParseAndVerifyModule(module_str); + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr m, + ParseAndReturnVerifiedModule(module_str)); auto compiled_module = backend() .compiler() - ->RunHloPasses(module().Clone(), backend().default_stream_executor(), + ->RunHloPasses(m->Clone(), backend().default_stream_executor(), /*device_allocator=*/nullptr) .ConsumeValueOrDie(); HloInstruction* root = compiled_module->entry_computation()->root_instruction(); Shape shape_copy = ShapeUtil::MakeShapeWithLayout(F32, {4, 5}, {1, 0}); - EXPECT_THAT(root, op::Add(op::Parameter(), - op::Slice(AllOf(op::Copy(op::Parameter(1)), - op::ShapeWithLayout(shape_copy))))); + EXPECT_THAT( + root, + GmockMatch(m::Add( + m::Parameter(), + m::Slice(m::Copy(m::Parameter(1)).WithShapeEqualTo(&shape_copy))))); } TEST_F(LayoutAssignmentTest, CopyDSliceOperandToAvoidImplicitLayoutChange) { @@ -960,27 +960,30 @@ 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) } )"; - ParseAndVerifyModule(module_str); + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr m, + ParseAndReturnVerifiedModule(module_str)); auto compiled_module = backend() .compiler() - ->RunHloPasses(module().Clone(), backend().default_stream_executor(), + ->RunHloPasses(m->Clone(), backend().default_stream_executor(), /*device_allocator=*/nullptr) .ConsumeValueOrDie(); HloInstruction* root = compiled_module->entry_computation()->root_instruction(); Shape shape_copy = ShapeUtil::MakeShapeWithLayout(F32, {4, 5}, {1, 0}); EXPECT_THAT(root, - op::Add(op::Parameter(), - op::DynamicSlice(AllOf(op::Copy(op::Parameter(1)), - op::ShapeWithLayout(shape_copy)), - op::Parameter(2)))); + GmockMatch(m::Add( + m::Parameter(), + m::DynamicSlice( + m::Copy(m::Parameter(1)).WithShapeEqualTo(&shape_copy), + m::Parameter(2), m::Parameter(3))))); } TEST_F(LayoutAssignmentTest, CopyConcatOperandToAvoidImplicitLayoutChange) { @@ -997,21 +1000,23 @@ TEST_F(LayoutAssignmentTest, CopyConcatOperandToAvoidImplicitLayoutChange) { } )"; - ParseAndVerifyModule(module_str); + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr m, + ParseAndReturnVerifiedModule(module_str)); auto compiled_module = backend() .compiler() - ->RunHloPasses(module().Clone(), backend().default_stream_executor(), + ->RunHloPasses(m->Clone(), backend().default_stream_executor(), /*device_allocator=*/nullptr) .ConsumeValueOrDie(); HloInstruction* root = compiled_module->entry_computation()->root_instruction(); Shape shape_copy = ShapeUtil::MakeShapeWithLayout(F32, {3, 5}, {1, 0}); - EXPECT_THAT(root, - op::Add(op::Parameter(), - op::Concatenate(AllOf(op::Copy(op::Parameter(1)), - op::ShapeWithLayout(shape_copy)), - op::Parameter(2)))); + EXPECT_THAT( + root, + GmockMatch(m::Add( + m::Parameter(), + m::Concatenate(m::Copy(m::Parameter(1)).WithShapeEqualTo(&shape_copy), + m::Parameter(2))))); } TEST_F(LayoutAssignmentTest, @@ -1028,16 +1033,18 @@ TEST_F(LayoutAssignmentTest, } )"; - ParseAndVerifyModule(module_str); + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr m, + ParseAndReturnVerifiedModule(module_str)); auto compiled_module = backend() .compiler() - ->RunHloPasses(module().Clone(), backend().default_stream_executor(), + ->RunHloPasses(m->Clone(), backend().default_stream_executor(), /*device_allocator=*/nullptr) .ConsumeValueOrDie(); HloInstruction* root = compiled_module->entry_computation()->root_instruction(); - EXPECT_THAT(root, op::Convolution(op::Parameter(0), op::Parameter(1))); + EXPECT_THAT(root, + GmockMatch(m::Convolution(m::Parameter(0), m::Parameter(1)))); } TEST_F(LayoutAssignmentTest, PropagatingLayoutFromResultToOperand) { @@ -1050,18 +1057,20 @@ TEST_F(LayoutAssignmentTest, PropagatingLayoutFromResultToOperand) { } )"; - ParseAndVerifyModule(module_str); + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr m, + ParseAndReturnVerifiedModule(module_str)); auto compiled_module = backend() .compiler() - ->RunHloPasses(module().Clone(), backend().default_stream_executor(), + ->RunHloPasses(m->Clone(), backend().default_stream_executor(), /*device_allocator=*/nullptr) .ConsumeValueOrDie(); HloInstruction* root = compiled_module->entry_computation()->root_instruction(); Shape shape_copy = ShapeUtil::MakeShapeWithLayout(F32, {4, 5}, {0, 1}); - EXPECT_THAT(root, op::Slice(AllOf(op::Copy(op::Parameter(0)), - op::ShapeWithLayout(shape_copy)))); + EXPECT_THAT(root, + GmockMatch(m::Slice( + m::Copy(m::Parameter(0)).WithShapeEqualTo(&shape_copy)))); } TEST_F(LayoutAssignmentTest, TupleCopyOnLayoutMismatch) { @@ -1107,20 +1116,21 @@ TEST_F(LayoutAssignmentTest, TupleCopyOnLayoutMismatch) { } )"; - ParseAndVerifyModule(module_str); + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr m, + ParseAndReturnVerifiedModule(module_str)); ComputationLayout computation_layout( - module().entry_computation()->ComputeProgramShape()); + m->entry_computation()->ComputeProgramShape()); // Sanity check to verify that there's a layout mismatch. - EXPECT_THAT(LayoutOf(&module(), "ibuf"), ElementsAre(0, 1)); - EXPECT_THAT(LayoutOf(&module(), "next_buf"), ElementsAre(1, 0)); + EXPECT_THAT(LayoutOf(m.get(), "ibuf"), ElementsAre(0, 1)); + EXPECT_THAT(LayoutOf(m.get(), "next_buf"), ElementsAre(1, 0)); - AssignLayouts(&module(), &computation_layout); + AssignLayouts(m.get(), &computation_layout); // Make sure that layout assignment did not magically eliminate the mismatch, // in which case the test didn't prove anything. - EXPECT_THAT(LayoutOf(&module(), "ibuf"), ElementsAre(0, 1)); - EXPECT_THAT(LayoutOf(&module(), "next_buf"), ElementsAre(1, 0)); + EXPECT_THAT(LayoutOf(m.get(), "ibuf"), ElementsAre(0, 1)); + EXPECT_THAT(LayoutOf(m.get(), "next_buf"), ElementsAre(1, 0)); } TEST_F(LayoutAssignmentTest, CustomCallNotLayoutConstrained) { @@ -1136,33 +1146,33 @@ ENTRY %CustomCallWithNotLayoutConstrained (p: f32[42,2,3]) -> f32[1,2,3,4] { // and result layout should match that of the computation. { TF_ASSERT_OK_AND_ASSIGN( - std::unique_ptr module, + std::unique_ptr m, ParseAndReturnVerifiedModule(module_str, GetModuleConfigForTest())); - ComputationLayout computation_layout = module->entry_computation_layout(); + ComputationLayout computation_layout = m->entry_computation_layout(); *computation_layout.mutable_parameter_layout(0) = ShapeLayout(ShapeUtil::MakeShapeWithLayout(F32, {42, 2, 3}, {0, 2, 1})); *computation_layout.mutable_result_layout() = ShapeLayout( ShapeUtil::MakeShapeWithLayout(F32, {1, 2, 3, 4}, {3, 2, 0, 1})); - AssignLayouts(module.get(), &computation_layout); + AssignLayouts(m.get(), &computation_layout); - HloInstruction* root = module->entry_computation()->root_instruction(); - ASSERT_THAT(root, op::CustomCall(op::Parameter())); + HloInstruction* root = m->entry_computation()->root_instruction(); + ASSERT_THAT(root, GmockMatch(m::CustomCall(m::Parameter()))); ExpectLayoutIs(root->shape(), {3, 2, 0, 1}); ExpectLayoutIs(root->operand(0)->shape(), {0, 2, 1}); } { TF_ASSERT_OK_AND_ASSIGN( - std::unique_ptr module, + std::unique_ptr m, ParseAndReturnVerifiedModule(module_str, GetModuleConfigForTest())); - ComputationLayout computation_layout = module->entry_computation_layout(); + ComputationLayout computation_layout = m->entry_computation_layout(); *computation_layout.mutable_parameter_layout(0) = ShapeLayout(ShapeUtil::MakeShapeWithLayout(F32, {42, 2, 3}, {0, 1, 2})); *computation_layout.mutable_result_layout() = ShapeLayout( ShapeUtil::MakeShapeWithLayout(F32, {1, 2, 3, 4}, {0, 2, 3, 1})); - AssignLayouts(module.get(), &computation_layout); + AssignLayouts(m.get(), &computation_layout); - HloInstruction* root = module->entry_computation()->root_instruction(); - ASSERT_THAT(root, op::CustomCall(op::Parameter())); + HloInstruction* root = m->entry_computation()->root_instruction(); + ASSERT_THAT(root, GmockMatch(m::CustomCall(m::Parameter()))); ExpectLayoutIs(root->shape(), {0, 2, 3, 1}); ExpectLayoutIs(root->operand(0)->shape(), {0, 1, 2}); } @@ -1179,24 +1189,24 @@ ENTRY %CustomCallWithLayoutConstraints (p0: f32[4,4], p1: f32[2,3]) -> f32[1,2,3 } )"; TF_ASSERT_OK_AND_ASSIGN( - std::unique_ptr module, + std::unique_ptr m, ParseAndReturnVerifiedModule(module_str, GetModuleConfigForTest())); - ComputationLayout computation_layout = module->entry_computation_layout(); + ComputationLayout computation_layout = m->entry_computation_layout(); *computation_layout.mutable_parameter_layout(0) = ShapeLayout(ShapeUtil::MakeShapeWithLayout(F32, {4, 4}, {1, 0})); *computation_layout.mutable_parameter_layout(1) = ShapeLayout(ShapeUtil::MakeShapeWithLayout(F32, {2, 3}, {1, 0})); *computation_layout.mutable_result_layout() = ShapeLayout( ShapeUtil::MakeShapeWithLayout(F32, {1, 2, 3, 4}, {2, 1, 0, 3})); - AssignLayouts(module.get(), &computation_layout); + AssignLayouts(m.get(), &computation_layout); // The custom call should be partially encapsulated in kCopy instructions // because of the layout mismatches. - ASSERT_THAT(module->entry_computation()->root_instruction(), - op::Copy(op::CustomCall(op::Copy(), op::Parameter()))); + ASSERT_THAT(m->entry_computation()->root_instruction(), + GmockMatch(m::Copy(m::CustomCall(m::Copy(), m::Parameter())))); const HloInstruction* custom_call = - module->entry_computation()->root_instruction()->operand(0); + m->entry_computation()->root_instruction()->operand(0); ExpectLayoutIs(custom_call->shape(), {3, 2, 0, 1}); ExpectLayoutIs(custom_call->operand(0)->shape(), {0, 1}); ExpectLayoutIs(custom_call->operand(1)->shape(), {1, 0}); @@ -1211,18 +1221,18 @@ ENTRY %CustomCallLayoutConstrainedZeroOperands () -> f32[1,2,3,4] { } )"; TF_ASSERT_OK_AND_ASSIGN( - std::unique_ptr module, + std::unique_ptr m, ParseAndReturnVerifiedModule(module_str, GetModuleConfigForTest())); - ComputationLayout computation_layout = module->entry_computation_layout(); + ComputationLayout computation_layout = m->entry_computation_layout(); *computation_layout.mutable_result_layout() = ShapeLayout( ShapeUtil::MakeShapeWithLayout(F32, {1, 2, 3, 4}, {2, 1, 0, 3})); - AssignLayouts(module.get(), &computation_layout); + AssignLayouts(m.get(), &computation_layout); - ASSERT_THAT(module->entry_computation()->root_instruction(), - op::Copy(op::CustomCall())); + ASSERT_THAT(m->entry_computation()->root_instruction(), + GmockMatch(m::Copy(m::CustomCall()))); const HloInstruction* custom_call = - module->entry_computation()->root_instruction()->operand(0); + m->entry_computation()->root_instruction()->operand(0); ExpectLayoutIs(custom_call->shape(), {3, 2, 0, 1}); } @@ -1238,25 +1248,25 @@ ENTRY %CustomCallLayoutConstrainedTupleOperand (p0: f32[4,4], p1: f32[2,3]) -> f } )"; TF_ASSERT_OK_AND_ASSIGN( - std::unique_ptr module, + std::unique_ptr m, ParseAndReturnVerifiedModule(module_str, GetModuleConfigForTest())); - ComputationLayout computation_layout = module->entry_computation_layout(); + ComputationLayout computation_layout = m->entry_computation_layout(); *computation_layout.mutable_parameter_layout(0) = ShapeLayout(ShapeUtil::MakeShapeWithLayout(F32, {4, 4}, {1, 0})); *computation_layout.mutable_parameter_layout(1) = ShapeLayout(ShapeUtil::MakeShapeWithLayout(F32, {2, 3}, {1, 0})); *computation_layout.mutable_result_layout() = ShapeLayout( ShapeUtil::MakeShapeWithLayout(F32, {1, 2, 3, 4}, {2, 1, 0, 3})); - AssignLayouts(module.get(), &computation_layout); + AssignLayouts(m.get(), &computation_layout); - HloInstruction* root = module->entry_computation()->root_instruction(); + HloInstruction* root = m->entry_computation()->root_instruction(); ExpectLayoutIs(root->shape(), {2, 1, 0, 3}); - ASSERT_THAT(module->entry_computation()->root_instruction(), - op::Copy(op::CustomCall(op::Tuple()))); + ASSERT_THAT(m->entry_computation()->root_instruction(), + GmockMatch(m::Copy(m::CustomCall(m::Tuple())))); const HloInstruction* custom_call = - module->entry_computation()->root_instruction()->operand(0); + m->entry_computation()->root_instruction()->operand(0); ExpectLayoutIs(custom_call->shape(), {3, 2, 0, 1}); ExpectTupleLayoutIs(custom_call->operand(0)->shape(), {{1, 0}, {0, 1}}); } @@ -1273,23 +1283,75 @@ ENTRY %CustomCallLayoutConstrainedTupleResult (p0: f32[4,4]) -> (f32[4,4]{1,0}, // Try with a couple different layouts. In each case the custom calls operand // and result layout should match that of the computation. TF_ASSERT_OK_AND_ASSIGN( - std::unique_ptr module, + std::unique_ptr m, ParseAndReturnVerifiedModule(module_str, GetModuleConfigForTest())); - ComputationLayout computation_layout = module->entry_computation_layout(); + ComputationLayout computation_layout = m->entry_computation_layout(); *computation_layout.mutable_parameter_layout(0) = ShapeLayout(ShapeUtil::MakeShapeWithLayout(F32, {4, 4}, {1, 0})); *computation_layout.mutable_result_layout() = ShapeLayout(ShapeUtil::MakeTupleShape( {ShapeUtil::MakeShapeWithLayout(F32, {4, 4}, {1, 0}), ShapeUtil::MakeShapeWithLayout(F32, {2, 3}, {1, 0})})); - AssignLayouts(module.get(), &computation_layout); + AssignLayouts(m.get(), &computation_layout); - ExpectTupleLayoutIs(module->result_shape(), {{1, 0}, {1, 0}}); + ExpectTupleLayoutIs(m->result_shape(), {{1, 0}, {1, 0}}); - const HloInstruction* custom_call = - FindInstruction(module.get(), "custom-call"); + const HloInstruction* custom_call = FindInstruction(m.get(), "custom-call"); ExpectTupleLayoutIs(custom_call->shape(), {{1, 0}, {0, 1}}); } +Status AssignLayoutsToComputation( + HloModule* m, ChannelLayoutConstraints* channel_constraints = nullptr) { + if (!m->entry_computation_layout().result_layout().LayoutIsSet()) { + m->mutable_entry_computation_layout() + ->mutable_result_layout() + ->SetToDefaultLayout(); + } + LayoutAssignment layout_assignment( + m->mutable_entry_computation_layout(), + LayoutAssignment::InstructionCanChangeLayout, channel_constraints); + return layout_assignment.Run(m).status(); +} + +TEST_F(LayoutAssignmentTest, OverwriteDiamondShapedConstraintsX) { + // Check that we handle a diamond-shaped graph correctly. + // transpose + // / \ + // add | + // \ / + // tuple + + auto b = HloComputation::Builder(TestName()); + Shape ashape = ShapeUtil::MakeShape(F32, {12, 8}); + Shape bshape = ShapeUtil::MakeShape(F32, {8, 12}); + auto param0 = + b.AddInstruction(HloInstruction::CreateParameter(0, bshape, "input")); + auto param1 = + b.AddInstruction(HloInstruction::CreateParameter(1, ashape, "input")); + auto transpose = + b.AddInstruction(HloInstruction::CreateTranspose(ashape, param0, {1, 0})); + auto add = b.AddInstruction( + HloInstruction::CreateBinary(ashape, HloOpcode::kAdd, transpose, param1)); + b.AddInstruction(HloInstruction::CreateTuple({add, transpose})); + auto m = CreateNewVerifiedModule(); + m->AddEntryComputation(b.Build()); + Shape ashape_major = ShapeUtil::MakeShapeWithLayout(F32, {12, 8}, {1, 0}); + Shape ashape_minor = ShapeUtil::MakeShapeWithLayout(F32, {12, 8}, {0, 1}); + *m->mutable_entry_computation_layout()->mutable_result_layout() = + ShapeLayout(ShapeUtil::MakeTupleShape({ashape_major, ashape_minor})); + const Layout r2_dim0major = LayoutUtil::MakeLayout({1, 0}); + ForceParameterLayout(m.get(), 0, r2_dim0major); + ForceParameterLayout(m.get(), 1, r2_dim0major); + TF_ASSERT_OK(AssignLayoutsToComputation(m.get())); + + EXPECT_THAT(add->shape().layout().minor_to_major(), ElementsAre(1, 0)); + EXPECT_THAT(add->operand(0)->shape().layout().minor_to_major(), + ElementsAre(1, 0)); + EXPECT_THAT(add->operand(1)->shape().layout().minor_to_major(), + ElementsAre(1, 0)); + + EXPECT_THAT(transpose->shape().layout().minor_to_major(), ElementsAre(0, 1)); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/service/llvm_compiler.cc b/tensorflow/compiler/xla/service/llvm_compiler.cc index d287aa4ec7bbcd11f51ea07cd2a1572e59f0d6c6..382b575120277ffb0e63e693757591681a78479e 100644 --- a/tensorflow/compiler/xla/service/llvm_compiler.cc +++ b/tensorflow/compiler/xla/service/llvm_compiler.cc @@ -22,7 +22,8 @@ limitations under the License. namespace xla { Status LLVMCompiler::RunHloPassesOnModuleGroup( - HloModuleGroup* module_group, se::StreamExecutor* executor, + HloModuleGroup* module_group, + absl::Span executors, DeviceMemoryAllocator* device_allocator) { return Unimplemented( "Model partitioning not implemented for the CPU/GPU compilers!"); diff --git a/tensorflow/compiler/xla/service/llvm_compiler.h b/tensorflow/compiler/xla/service/llvm_compiler.h index 86abd5da0189feb0eadfde3d6dbab446eb2be900..182d8edbe30da292f28aeab53be646ce6651839f 100644 --- a/tensorflow/compiler/xla/service/llvm_compiler.h +++ b/tensorflow/compiler/xla/service/llvm_compiler.h @@ -70,7 +70,8 @@ class LLVMCompiler : public Compiler { using Compiler::RunHloPasses; Status RunHloPassesOnModuleGroup( - HloModuleGroup* module_group, se::StreamExecutor* executor, + HloModuleGroup* module_group, + absl::Span executors, DeviceMemoryAllocator* device_allocator) override; StatusOr>> RunBackendOnModuleGroup( diff --git a/tensorflow/compiler/xla/service/llvm_ir/BUILD b/tensorflow/compiler/xla/service/llvm_ir/BUILD index 5f7ad81d82978d0a752b33d12b72e16f0c1c6826..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", ], @@ -72,6 +71,7 @@ cc_library( "//tensorflow/compiler/xla/service:hlo_module_config", "//tensorflow/compiler/xla/service:name_uniquer", "//tensorflow/core:lib", + "@com_google_absl//absl/base", "@com_google_absl//absl/strings", "@com_google_absl//absl/types:span", "@llvm//:core", @@ -168,6 +168,8 @@ 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", ], @@ -196,14 +198,17 @@ cc_library( hdrs = ["sort_util.h"], deps = [ ":ir_array", + ":kernel_support_library", ":llvm_loop", ":llvm_util", ":loop_emitter", "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla/service/gpu:parallel_loop_emitter", "//tensorflow/compiler/xla/service/gpu:partition_assignment", "//tensorflow/core:lib", "@com_google_absl//absl/strings", + "@com_google_absl//absl/types:span", "@llvm//:core", "@llvm//:support", ], 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/dynamic_update_slice_util.cc b/tensorflow/compiler/xla/service/llvm_ir/dynamic_update_slice_util.cc index cc2e862f2eb9a49099c5f90efe1b29fb77c8f106..c66eaec8fb0e4c03f6967fec0cf0ae9661cdf470 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 @@ -36,19 +36,20 @@ bool CanUpdateDynamicSliceInPlace(HloInstruction* dynamic_update_slice, // EmitFusedDynamicUpdateSliceInPlace. // // Emits a sequential loop if launch_dimensions is null. +using IndexGenerator = std::function(int64)>; + static Status EmitDynamicUpdateSliceInPlaceImpl( - const Shape& update_shape, const ElementGenerator& start_indices_generator, + const Shape& update_shape, const IndexGenerator& start_indices_generator, bool is_signed, ElementGenerator update_array_generator, const IrArray& output_array, const gpu::LaunchDimensions* launch_dimensions, absl::string_view name, llvm::IRBuilder<>* b) { const Shape& output_shape = output_array.GetShape(); // Read start indices from start_indices_generator. - const int64 rank = ShapeUtil::Rank(output_shape); + const int64 rank = output_shape.rank(); IrArray::Index start_index(b->getInt64Ty(), rank); for (int64 i = 0; i < rank; ++i) { - IrArray::Index dim_index({b->getInt64(i)}); - TF_ASSIGN_OR_RETURN(start_index[i], start_indices_generator(dim_index)); + TF_ASSIGN_OR_RETURN(start_index[i], start_indices_generator(i)); llvm::Value* output_dim_size = llvm::ConstantInt::get( start_index[i]->getType(), output_shape.dimensions(i)); llvm::Value* update_dim_size = llvm::ConstantInt::get( @@ -112,9 +113,20 @@ Status EmitDynamicUpdateSliceInPlace(absl::Span operand_arrays, Shape output_shape = output_array.GetShape(); Shape update_shape = update_array.GetShape(); - ElementGenerator start_indices_generator = [&](const IrArray::Index& index) { - return start_indices_array.EmitReadArrayElement(index, b); - }; + 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); + }; + } + ElementGenerator update_array_generator = [&](const IrArray::Index& index) { return update_array.EmitReadArrayElement(index, b); }; @@ -130,7 +142,8 @@ Status EmitDynamicUpdateSliceInPlace(absl::Span operand_arrays, // // Emits a sequential loop if launch_dimensions is null. static Status EmitFusedDynamicUpdateSliceInPlaceImpl( - HloInstruction* fusion, absl::Span fusion_operand_arrays, + HloInstruction* fusion, + GeneratorForOperandIrArrays operand_arrays_generator, const IrArray& fusion_output_array, ElementalIrEmitter* elemental_emitter, const gpu::LaunchDimensions* launch_dimensions, llvm::IRBuilder<>* b) { CHECK_EQ(fusion->opcode(), HloOpcode::kFusion); @@ -160,11 +173,25 @@ static Status EmitFusedDynamicUpdateSliceInPlaceImpl( LayoutUtil::CopyLayoutBetweenShapes(fusion->shape(), &update_shape)); // Create element generators for update and start_indices. - FusedIrEmitter fused_emitter(fusion_operand_arrays, elemental_emitter); + FusedIrEmitter fused_emitter(std::move(operand_arrays_generator), + elemental_emitter); TF_RETURN_IF_ERROR(dynamic_update_slice->Accept(&fused_emitter)); ElementGenerator update_array_generator = fused_emitter.GetGenerator(update); - ElementGenerator start_indices_generator = - fused_emitter.GetGenerator(start_indices); + + // 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())); + }; + } bool is_signed = ShapeUtil::ElementIsSigned(start_indices->shape()); return EmitDynamicUpdateSliceInPlaceImpl( @@ -173,21 +200,24 @@ static Status EmitFusedDynamicUpdateSliceInPlaceImpl( } Status EmitFusedDynamicUpdateSliceInPlace( - HloInstruction* fusion, absl::Span fusion_operand_arrays, + HloInstruction* fusion, + GeneratorForOperandIrArrays operand_arrays_generator, const IrArray& fusion_output_array, ElementalIrEmitter* elemental_emitter, llvm::IRBuilder<>* b) { return EmitFusedDynamicUpdateSliceInPlaceImpl( - fusion, fusion_operand_arrays, fusion_output_array, elemental_emitter, + fusion, std::move(operand_arrays_generator), fusion_output_array, + elemental_emitter, /*launch_dimensions=*/nullptr, b); } Status EmitParallelFusedDynamicUpdateSliceInPlace( - HloInstruction* fusion, absl::Span fusion_operand_arrays, + HloInstruction* fusion, + GeneratorForOperandIrArrays operand_arrays_generator, const IrArray& fusion_output_array, ElementalIrEmitter* elemental_emitter, const gpu::LaunchDimensions& launch_dimensions, llvm::IRBuilder<>* b) { return EmitFusedDynamicUpdateSliceInPlaceImpl( - fusion, fusion_operand_arrays, fusion_output_array, elemental_emitter, - &launch_dimensions, b); + fusion, std::move(operand_arrays_generator), fusion_output_array, + elemental_emitter, &launch_dimensions, b); } } // namespace llvm_ir diff --git a/tensorflow/compiler/xla/service/llvm_ir/dynamic_update_slice_util.h b/tensorflow/compiler/xla/service/llvm_ir/dynamic_update_slice_util.h index fb3e4eb97cae06f2a0c87dd7118b8332048df56e..7fe803d1f8da5251c99f0a8fd97f99e9ca031175 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/dynamic_update_slice_util.h +++ b/tensorflow/compiler/xla/service/llvm_ir/dynamic_update_slice_util.h @@ -27,6 +27,9 @@ limitations under the License. namespace xla { namespace llvm_ir { +using GeneratorForOperandIrArrays = + std::function()>; + // Checks if we can emit code for the given DynamicUpdateSlice node that updates // its input in place. Returns true if the dynamic-update-slice's // array-to-be-updated and output share the same BufferAllocation::Slice. @@ -73,14 +76,16 @@ Status EmitDynamicUpdateSliceInPlace(absl::Span operand_arrays, // (sequential) code for a fusion node that does the dynamic-update-slice in // place. Status EmitFusedDynamicUpdateSliceInPlace( - HloInstruction* fusion, absl::Span fusion_operand_arrays, + HloInstruction* fusion, + GeneratorForOperandIrArrays operand_arrays_generator, const IrArray& fusion_output_array, ElementalIrEmitter* elemental_emitter, llvm::IRBuilder<>* b); // Same as EmitFusedDynamicUpdateSliceInPlace, except emits a parallel loop with // the given launch dimensions. Status EmitParallelFusedDynamicUpdateSliceInPlace( - HloInstruction* fusion, absl::Span fusion_operand_arrays, + HloInstruction* fusion, + GeneratorForOperandIrArrays operand_arrays_generator, const IrArray& fusion_output_array, ElementalIrEmitter* elemental_emitter, const gpu::LaunchDimensions& launch_dimensions, llvm::IRBuilder<>* b); 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 b606c993a2d58a6d177af10de7b214de130c2279..e440f05e2b2f0d4a2a4c7b326b4881183de4d235 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/fused_ir_emitter.cc +++ b/tensorflow/compiler/xla/service/llvm_ir/fused_ir_emitter.cc @@ -33,9 +33,9 @@ namespace xla { using llvm_ir::IrArray; Status FusedIrEmitter::DefaultAction(HloInstruction* hlo) { - generators_[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; @@ -63,25 +63,26 @@ Status FusedIrEmitter::DefaultAction(HloInstruction* hlo) { << llvm_ir::AsString(b_->GetInsertBlock()->getName()) << ")."; } - TF_ASSIGN_OR_RETURN( - generated_value_cache_[hlo][index.multidim()], - elemental_emitter_->MakeElementGenerator(hlo, generators_)(index)); + TF_ASSIGN_OR_RETURN(generated_value_cache_[hlo][index.multidim()], + elemental_emitter_->MakeElementGenerator( + hlo, indexed_generators_)(index)); return generated_value_cache_[hlo][index.multidim()]; }; return Status::OK(); } Status FusedIrEmitter::HandleConstant(HloInstruction* constant) { - const Literal& literal = constant->literal(); - llvm::Constant* initializer = - llvm_ir::ConvertLiteralToIrConstant(literal, module_); - llvm::GlobalVariable* global = new llvm::GlobalVariable( - *b_->GetInsertBlock()->getModule(), initializer->getType(), - /*isConstant=*/true, llvm::GlobalValue::ExternalLinkage, initializer, - /*Name=*/""); - llvm::Constant* shape_constant = llvm::ConstantExpr::getBitCast( - global, llvm_ir::ShapeToIrType(literal.shape(), module_)->getPointerTo()); - generators_[constant] = [=](const IrArray::Index& index) { + indexed_generators_[constant] = [=](const IrArray::Index& index) { + const Literal& literal = constant->literal(); + llvm::Constant* initializer = + llvm_ir::ConvertLiteralToIrConstant(literal, module_); + llvm::GlobalVariable* global = new llvm::GlobalVariable( + *b_->GetInsertBlock()->getModule(), initializer->getType(), + /*isConstant=*/true, llvm::GlobalValue::ExternalLinkage, initializer, + /*Name=*/""); + llvm::Constant* shape_constant = llvm::ConstantExpr::getBitCast( + global, + llvm_ir::ShapeToIrType(literal.shape(), module_)->getPointerTo()); return IrArray(shape_constant, constant->shape()) .EmitReadArrayElement(index, b_); }; @@ -91,34 +92,47 @@ Status FusedIrEmitter::HandleConstant(HloInstruction* constant) { Status FusedIrEmitter::HandleGetTupleElement( HloInstruction* get_tuple_element) { - // Lookup ir value for 'operand'. - auto operand = get_tuple_element->operand(0); - auto it = gte_values_.find(operand); - if (it == gte_values_.end()) { - return Unimplemented( - "GetTupleElement fusion currently only supports" - " parameter operands, but found operand: %s", - operand->name()); - } - // Emit code to lookup tuple element pointer, and store it in 'gte_values_'. - llvm::Value* tuple_element_ptr = llvm_ir::EmitGetTupleElement( - get_tuple_element->shape(), get_tuple_element->tuple_index(), - /*alignment=*/1, it->second, b_, module_); - gte_values_.insert(std::make_pair(get_tuple_element, tuple_element_ptr)); - // Emit code to read base tuple element array (if non-tuple shaped). - if (!ShapeUtil::IsTuple(get_tuple_element->shape())) { - generators_[get_tuple_element] = + auto emit_tuple_element_ptr = [=]() -> StatusOr { + const HloInstruction* tuple_operand = get_tuple_element->operand(0); + llvm::Value* tuple_ptr; + if (tuple_operand->opcode() == HloOpcode::kGetTupleElement) { + TF_ASSIGN_OR_RETURN(tuple_ptr, non_indexed_generators_[tuple_operand]()); + } else { + if (tuple_operand->opcode() != HloOpcode::kParameter) { + return Unimplemented( + "GetTupleElement fusion currently only supports parameter or " + "nested" + "GetTupleElement as tuple operand, found an exception: %s", + tuple_operand->name()); + } + tuple_ptr = + GetBasePointerForFusedParameter(tuple_operand->parameter_number()); + } + + // Lookup tuple element pointer. + return llvm_ir::EmitGetTupleElement( + get_tuple_element->shape(), get_tuple_element->tuple_index(), + /*alignment=*/1, tuple_ptr, b_, module_); + }; + + if (!get_tuple_element->shape().IsTuple()) { + indexed_generators_[get_tuple_element] = [=](const IrArray::Index& index) -> StatusOr { // TODO(b/34080002) Add aliasing information to tuple element IrArray. + TF_ASSIGN_OR_RETURN(llvm::Value * tuple_element_ptr, + emit_tuple_element_ptr()); return IrArray(tuple_element_ptr, get_tuple_element->shape()) .EmitReadArrayElement(index, b_); }; + } else { + non_indexed_generators_[get_tuple_element] = emit_tuple_element_ptr; } return Status::OK(); } Status FusedIrEmitter::HandleParameter(HloInstruction* parameter) { - generators_[parameter] = [=](const IrArray::Index& index) -> llvm::Value* { + indexed_generators_[parameter] = + [=](const IrArray::Index& index) -> llvm::Value* { if (tiled_parameter_info_) { if (llvm::Value* param_tile_buffer = tiled_parameter_info_->GetBufferForParameter( @@ -135,14 +149,9 @@ Status FusedIrEmitter::HandleParameter(HloInstruction* parameter) { "tiled_buffer"); } } - return parameter_arrays_[parameter->parameter_number()] + return GetIrArrayForFusedParameter(parameter->parameter_number()) .EmitReadArrayElement(index, b_); }; - // Store ir value for fusion operand associated with fusion parameter to be - // accessed by subsequent fused GetTupleElement instructions. - gte_values_.insert(std::make_pair( - parameter, - parameter_arrays_[parameter->parameter_number()].GetBasePointer())); return Status::OK(); } @@ -153,12 +162,13 @@ Status FusedIrEmitter::HandleTuple(HloInstruction* tuple) { operand_elemental_ir_types.push_back(llvm_ir::PrimitiveTypeToIrType( operand->shape().element_type(), module_)); } - generators_[tuple] = + indexed_generators_[tuple] = [=](const IrArray::Index& index) -> StatusOr { llvm::Value* ret = llvm::UndefValue::get( llvm::StructType::get(b_->getContext(), operand_elemental_ir_types)); for (size_t i = 0; i < ShapeUtil::TupleElementCount(tuple->shape()); ++i) { - TF_ASSIGN_OR_RETURN(llvm::Value * val_i, generators_[operands[i]](index)); + TF_ASSIGN_OR_RETURN(llvm::Value * val_i, + indexed_generators_[operands[i]](index)); ret = b_->CreateInsertValue(ret, val_i, i); } return ret; @@ -171,15 +181,15 @@ Status FusedIrEmitter::FinishVisit(HloInstruction* root) { return Status::OK(); } -FusedIrEmitter::Generator FusedIrEmitter::GetRootGenerator() const { +FusedIrEmitter::IndexedGenerator FusedIrEmitter::GetRootGenerator() const { CHECK_NE(nullptr, fused_root_) << "GetRootGenerator should be called after Accept."; - return generators_.at(fused_root_); + return indexed_generators_.at(fused_root_); } -FusedIrEmitter::Generator FusedIrEmitter::GetGenerator( +FusedIrEmitter::IndexedGenerator FusedIrEmitter::GetGenerator( const HloInstruction* instruction) const { - return generators_.at(instruction); + return indexed_generators_.at(instruction); } } // namespace xla 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 44d21fa750a532633f46614002d59c90fc0b5d40..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,8 @@ 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" #include "llvm/IR/Value.h" @@ -52,11 +54,15 @@ namespace xla { // same length. class FusedIrEmitter : public DfsHloVisitorWithDefault { public: - using Generator = llvm_ir::ElementGenerator; + using IndexedGenerator = llvm_ir::ElementGenerator; + using NonIndexedGenerator = std::function()>; + using GeneratorForOperandIrArrays = + std::function()>; - FusedIrEmitter(absl::Span parameter_arrays, + FusedIrEmitter(GeneratorForOperandIrArrays operand_arrays_generator, ElementalIrEmitter* elemental_emitter) - : parameter_arrays_(parameter_arrays), + : operand_arrays_(), + operand_arrays_generator_(std::move(operand_arrays_generator)), tiled_parameter_info_(nullptr), elemental_emitter_(elemental_emitter), b_(elemental_emitter->b()), @@ -76,25 +82,34 @@ class FusedIrEmitter : public DfsHloVisitorWithDefault { Status FinishVisit(HloInstruction* root) override; // Returns the generator function for the root of the fused computation. - Generator GetRootGenerator() const; + IndexedGenerator GetRootGenerator() const; // Returns the generator function for the given instruction. - Generator GetGenerator(const HloInstruction* instruction) const; - - // Returns the ir value for instruction 'hlo'. - llvm::Value* GetIrValueForGTE(const HloInstruction* hlo) const { - auto it = gte_values_.find(hlo); - CHECK(it != gte_values_.end()); - return it->second; - } + IndexedGenerator GetGenerator(const HloInstruction* instruction) const; void SetTiledParameterInfo(const llvm_ir::TiledParameterInfo* info) { tiled_parameter_info_ = info; } + protected: + // Returns the IrArrays for the fusion instruction operands. + llvm_ir::IrArray& GetIrArrayForFusedParameter(int64 parameter_number) { + if (!operand_arrays_.has_value()) { + operand_arrays_ = operand_arrays_generator_(); + } + return operand_arrays_.value()[parameter_number]; + } + + llvm::Value* GetBasePointerForFusedParameter(int64 parameter_number) { + return GetIrArrayForFusedParameter(parameter_number).GetBasePointer(); + } + private: - // Arrays of parameters of fusion instruction - absl::Span parameter_arrays_; + // IrArrays for the fusion instruction operands, whose base addresses are the + // base address of the corresponding parameters in the fused computation. + absl::optional> operand_arrays_; + GeneratorForOperandIrArrays operand_arrays_generator_; + const llvm_ir::TiledParameterInfo* tiled_parameter_info_; ElementalIrEmitter* elemental_emitter_; @@ -106,19 +121,24 @@ class FusedIrEmitter : public DfsHloVisitorWithDefault { llvm::IRBuilder<>* b_; llvm::Module* module_; - // Map from instruction pointers to functions to generate elements of their - // outputs - std::unordered_map generators_; + // Map from instructions to functions that generate code for the output + // elements. If an instruction is a GetTupleElement instruction, the + // instruction produces non-tuple result. + std::unordered_map + indexed_generators_; + + // Map from tuple-result-producing GetTupleELement instructions to functions + // that generate the base pointers for the output elements. This is used to + // support the translation of nested GetTupleElement instructions. + std::unordered_map + non_indexed_generators_; // 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_; - - // Stores ir values required to emit fused (and possibly nested) - // GetTupleElement instructions. - std::unordered_map gte_values_; }; } // namespace xla diff --git a/tensorflow/compiler/xla/service/llvm_ir/ir_array.cc b/tensorflow/compiler/xla/service/llvm_ir/ir_array.cc index 67f7423121177e2ca1e3384341dad2644c8f5e34..8ee07ae8331e986f9d271be5e39065f0d87853b1 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/ir_array.cc +++ b/tensorflow/compiler/xla/service/llvm_ir/ir_array.cc @@ -61,7 +61,7 @@ void IrArray::Index::Delinearize(std::vector* multidim, IrArray::Index::Index(llvm::Value* linear, const Shape& shape, llvm::IRBuilder<>* b) - : multidim_(ShapeUtil::Rank(shape)), + : multidim_(shape.rank()), linear_(linear), layout_(shape.layout()), dims_(shape.dimensions().begin(), shape.dimensions().end()) { @@ -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; @@ -117,10 +117,10 @@ IrArray::IrArray(llvm::Value* base_ptr, const Shape& shape) ++depth; } - if (!ShapeUtil::IsArray(*shape_) || ShapeUtil::IsScalar(*shape_)) { + if (!shape_->IsArray() || ShapeUtil::IsScalar(*shape_)) { DCHECK(depth == 1 || depth == 0) << depth; } else { - DCHECK_EQ(depth, ShapeUtil::Rank(*shape_)) << shape.ShortDebugString(); + DCHECK_EQ(depth, shape_->rank()) << shape.ShortDebugString(); } } @@ -137,12 +137,12 @@ IrArray::Index IrArray::Index::SourceIndexOfReshape( const Shape& output_shape, const Shape& input_shape, llvm::IRBuilder<>* builder) const { const auto& target_index = *this; - CHECK_EQ(target_index.size(), ShapeUtil::Rank(output_shape)); + CHECK_EQ(target_index.size(), output_shape.rank()); std::vector> common_factors = CommonFactors(AsInt64Slice(input_shape.dimensions()), AsInt64Slice(output_shape.dimensions())); std::vector source_multidim_index( - ShapeUtil::Rank(input_shape), llvm::UndefValue::get(index_type_)); + input_shape.rank(), llvm::UndefValue::get(index_type_)); // We compute the source indices in each common factor from only the target // indices in the same common factor. for (ssize_t k = common_factors.size() - 2; k >= 0; --k) { @@ -257,7 +257,7 @@ IrArray::Index IrArray::Index::SourceIndexOfBroadcast( const Shape& shape, const Shape& operand_shape, absl::Span dimension_mapping, llvm::IRBuilder<>* builder) const { - int64 rank = ShapeUtil::Rank(operand_shape); + int64 rank = operand_shape.rank(); std::vector source_index(rank); for (int64 i = 0; i < rank; ++i) { source_index[i] = multidim_[dimension_mapping[i]]; @@ -271,7 +271,7 @@ IrArray::Index IrArray::Index::SourceIndexOfBroadcast( // The other dimensions can be masked out with a div and a mod operation. std::vector logical_to_physical = LayoutUtil::MakeLogicalToPhysical(shape.layout()); - int64 output_rank = ShapeUtil::Rank(shape); + int64 output_rank = shape.rank(); // The minimum physical dimension that is broadcasted. int64 min_broadcasted_dimension = output_rank; // The maximum physical dimension that is broadcasted. @@ -348,7 +348,7 @@ llvm::Value* IrArray::EmitArrayElementAddress(const IrArray::Index& index, // over higher-rank arrays. return base_ptr_; } - CHECK_EQ(index.size(), ShapeUtil::Rank(*shape_)); + CHECK_EQ(index.size(), shape_->rank()); if (index.LinearValidOnShape(*shape_)) { llvm::Module* module = b->GetInsertBlock()->getParent()->getParent(); diff --git a/tensorflow/compiler/xla/service/llvm_ir/ir_array.h b/tensorflow/compiler/xla/service/llvm_ir/ir_array.h index f4b05f29c38529b3cce81b4c8ee6fae5c00cafcc..b706ebd311cbb706e7e4698b93319e37e664d10a 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/ir_array.h +++ b/tensorflow/compiler/xla/service/llvm_ir/ir_array.h @@ -25,6 +25,7 @@ limitations under the License. #include "llvm/IR/IRBuilder.h" #include "llvm/IR/Value.h" #include "tensorflow/compiler/xla/map_util.h" +#include "tensorflow/compiler/xla/shape.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/platform/logging.h" @@ -108,6 +109,14 @@ class IrArray { Index(absl::Span multidim, llvm::Value* linear, const Shape& shape); + // Returns an index that adds `addend` to the given `dim` of the object. + Index AddOffsetToDim(llvm::Value* addend, int64 dim, + llvm::IRBuilder<>* b) const { + IrArray::Index index = *this; + index[dim] = b->CreateAdd(index[dim], addend); + return index; + } + const std::vector& multidim() const { return multidim_; } llvm::Value* linear() const { return linear_; } @@ -121,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; @@ -180,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() { @@ -211,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; @@ -227,7 +243,6 @@ class IrArray { llvm::Type* GetElementLlvmType() const { return element_type_; } const Shape& GetShape() const { - CHECK(shape_ != nullptr); return *shape_; } @@ -322,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..cf5083e8c13b9485035923895cec1ad05049c644 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)...); diff --git a/tensorflow/compiler/xla/service/llvm_ir/kernel_support_library.cc b/tensorflow/compiler/xla/service/llvm_ir/kernel_support_library.cc index bd0139f85b6a5c5dc23dad962263038451921e65..5eeb29c478a371dae83251771f2dc4844672d3e9 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/kernel_support_library.cc +++ b/tensorflow/compiler/xla/service/llvm_ir/kernel_support_library.cc @@ -18,28 +18,29 @@ limitations under the License. #include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h" namespace xla { -Status KernelSupportLibrary::For( +Status KernelSupportLibrary::ForWithStatus( absl::string_view name, llvm::Value* start, llvm::Value* end, llvm::Value* step, const std::function& for_body_generator) { - return If(b_->CreateICmpSLT(start, end), [&]() -> Status { + return IfWithStatus(b_->CreateICmpSLT(start, end), [&]() -> Status { TF_RETURN_IF_ERROR(for_body_generator(start, /*is_first_iteration=*/true)); - return For(name, b_->CreateAdd(start, step), end, step, - [&](llvm::Value* iv) { return for_body_generator(iv, false); }); + return ForWithStatus( + name, b_->CreateAdd(start, step), end, step, + [&](llvm::Value* iv) { return for_body_generator(iv, false); }); }); } -Status KernelSupportLibrary::For( +Status KernelSupportLibrary::ForWithStatus( absl::string_view name, llvm::Value* start, llvm::Value* end, llvm::Value* step, bool peel_first_iteration, const std::function& for_body_generator) { if (peel_first_iteration) { - return For(name, start, end, step, true, - [&](llvm::Value* indvar, bool is_first_iteration) -> Status { - return for_body_generator(indvar, - b_->getInt1(is_first_iteration)); - }); + return ForWithStatus( + name, start, end, step, true, + [&](llvm::Value* indvar, bool is_first_iteration) -> Status { + return for_body_generator(indvar, b_->getInt1(is_first_iteration)); + }); } else { std::unique_ptr loop = llvm_ir::ForLoop::EmitForLoop( name, start, end, step, b_, @@ -55,7 +56,7 @@ Status KernelSupportLibrary::For( } } -Status KernelSupportLibrary::If( +Status KernelSupportLibrary::IfWithStatus( absl::string_view name, llvm::Value* condition, const std::function& true_block_generator, const std::function& false_block_generator) { diff --git a/tensorflow/compiler/xla/service/llvm_ir/kernel_support_library.h b/tensorflow/compiler/xla/service/llvm_ir/kernel_support_library.h index 43fec311f150d6054f6ad24f99db332f90ff94a3..612b839cfa15711061e1ae53358a72d5220e1801 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/kernel_support_library.h +++ b/tensorflow/compiler/xla/service/llvm_ir/kernel_support_library.h @@ -48,41 +48,42 @@ class KernelSupportLibrary { // for (i64 i = `start` + `step`; i s< `end`; i += `step`) // `for_body_generator(/*ind_var=*/,i, /*is_first_iteration=*/false)`; // } - Status For( + Status ForWithStatus( absl::string_view name, llvm::Value* start, llvm::Value* end, llvm::Value* step, const std::function& for_body_generator); - void ForReturnVoid( + void For( absl::string_view name, llvm::Value* start, llvm::Value* end, llvm::Value* step, const std::function& for_body_generator) { CHECK_EQ(Status::OK(), - For(name, start, end, step, + ForWithStatus( + name, start, end, step, [&](llvm::Value* ind_var, bool is_first_iteration) -> Status { for_body_generator(ind_var, is_first_iteration); return Status::OK(); })); } - Status For(absl::string_view name, int64 start, int64 end, int64 step, - const std::function& - for_body_generator) { - return For(name, /*start=*/b_->getInt64(start), - /*end=*/b_->getInt64(end), - /*step=*/b_->getInt64(step), for_body_generator); + Status ForWithStatus( + absl::string_view name, int64 start, int64 end, int64 step, + const std::function& for_body_generator) { + return ForWithStatus(name, /*start=*/b_->getInt64(start), + /*end=*/b_->getInt64(end), + /*step=*/b_->getInt64(step), for_body_generator); } - void ForReturnVoid( + void For( absl::string_view name, int64 start, int64 end, int64 step, const std::function& for_body_generator) { - ForReturnVoid(name, /*start=*/b_->getInt64(start), - /*end=*/b_->getInt64(end), - /*step=*/b_->getInt64(step), for_body_generator); + For(name, /*start=*/b_->getInt64(start), + /*end=*/b_->getInt64(end), + /*step=*/b_->getInt64(step), for_body_generator); } // Generates the following control flow structure if `peel_first_iteration` is @@ -99,19 +100,19 @@ class KernelSupportLibrary { // for (i64 i = `start`; i s< `end`; i += `step`) // `for_body_generator(/*ind_var=*/,i, // /*is_first_iteration=*/,(i != `start`))`; - Status For(absl::string_view name, llvm::Value* start, llvm::Value* end, - llvm::Value* step, bool peel_first_iteration, - const std::function& - for_body_generator); + Status ForWithStatus( + absl::string_view name, llvm::Value* start, llvm::Value* end, + llvm::Value* step, bool peel_first_iteration, + const std::function& + for_body_generator); - void ForReturnVoid(absl::string_view name, llvm::Value* start, - llvm::Value* end, llvm::Value* step, - bool peel_first_iteration, - const std::function& - for_body_generator) { - TF_CHECK_OK(For( + void For(absl::string_view name, llvm::Value* start, llvm::Value* end, + llvm::Value* step, bool peel_first_iteration, + const std::function& + for_body_generator) { + TF_CHECK_OK(ForWithStatus( name, start, end, step, peel_first_iteration, [&](llvm::Value* ind_var, llvm::Value* is_first_iteration) -> Status { for_body_generator(ind_var, is_first_iteration); @@ -119,80 +120,81 @@ class KernelSupportLibrary { })); } - Status For(absl::string_view name, llvm::Value* start, llvm::Value* end, - int64 step, bool peel_first_iteration, - const std::function& - for_body_generator) { - return For(name, /*start=*/start, /*end=*/end, - /*step=*/llvm::ConstantInt::get(start->getType(), step), - peel_first_iteration, for_body_generator); + Status ForWithStatus( + absl::string_view name, llvm::Value* start, llvm::Value* end, int64 step, + bool peel_first_iteration, + const std::function& + for_body_generator) { + return ForWithStatus( + name, /*start=*/start, /*end=*/end, + /*step=*/llvm::ConstantInt::get(start->getType(), step), + peel_first_iteration, for_body_generator); } - void ForReturnVoid(absl::string_view name, llvm::Value* start, - llvm::Value* end, int64 step, bool peel_first_iteration, - const std::function& - for_body_generator) { - ForReturnVoid(name, /*start=*/start, /*end=*/end, - /*step=*/llvm::ConstantInt::get(start->getType(), step), - peel_first_iteration, for_body_generator); + void For(absl::string_view name, llvm::Value* start, llvm::Value* end, + int64 step, bool peel_first_iteration, + const std::function& + for_body_generator) { + For(name, /*start=*/start, /*end=*/end, + /*step=*/llvm::ConstantInt::get(start->getType(), step), + peel_first_iteration, for_body_generator); } - Status For( + Status ForWithStatus( absl::string_view name, llvm::Value* start, llvm::Value* end, llvm::Value* step, const std::function& for_body_generator) { - return For(name, start, end, step, - /*peel_first_iteration=*/false, - [&](llvm::Value* indvar, llvm::Value*) -> Status { - return for_body_generator(indvar); - }); + return ForWithStatus(name, start, end, step, + /*peel_first_iteration=*/false, + [&](llvm::Value* indvar, llvm::Value*) -> Status { + return for_body_generator(indvar); + }); } - void ForReturnVoid( + void For( absl::string_view name, llvm::Value* start, llvm::Value* end, llvm::Value* step, const std::function& for_body_generator) { - ForReturnVoid(name, start, end, step, - /*peel_first_iteration=*/false, - [&](llvm::Value* indvar, llvm::Value*) { - return for_body_generator(indvar); - }); + For(name, start, end, step, + /*peel_first_iteration=*/false, [&](llvm::Value* indvar, llvm::Value*) { + return for_body_generator(indvar); + }); } - Status For( + Status ForWithStatus( absl::string_view name, llvm::Value* start, llvm::Value* end, int64 step, const std::function& for_body_generator) { - return For(name, start, end, llvm::ConstantInt::get(start->getType(), step), - /*peel_first_iteration=*/false, - [&](llvm::Value* indvar, llvm::Value*) -> Status { - return for_body_generator(indvar); - }); + return ForWithStatus(name, start, end, + llvm::ConstantInt::get(start->getType(), step), + /*peel_first_iteration=*/false, + [&](llvm::Value* indvar, llvm::Value*) -> Status { + return for_body_generator(indvar); + }); } - void ForReturnVoid( + void For( absl::string_view name, llvm::Value* start, llvm::Value* end, int64 step, const std::function& for_body_generator) { - ForReturnVoid(name, start, end, - llvm::ConstantInt::get(start->getType(), step), - for_body_generator); + For(name, start, end, llvm::ConstantInt::get(start->getType(), step), + for_body_generator); } - Status For( + Status ForWithStatus( absl::string_view name, int64 start, int64 end, int64 step, const std::function& for_body_generator) { - return For(name, /*start=*/b_->getInt64(start), - /*end=*/b_->getInt64(end), - /*step=*/b_->getInt64(step), for_body_generator); + return ForWithStatus(name, /*start=*/b_->getInt64(start), + /*end=*/b_->getInt64(end), + /*step=*/b_->getInt64(step), for_body_generator); } - void ForReturnVoid( + void For( absl::string_view name, int64 start, int64 end, int64 step, const std::function& for_body_generator) { - ForReturnVoid(name, /*start=*/b_->getInt64(start), - /*end=*/b_->getInt64(end), - /*step=*/b_->getInt64(step), for_body_generator); + For(name, /*start=*/b_->getInt64(start), + /*end=*/b_->getInt64(end), + /*step=*/b_->getInt64(step), for_body_generator); } // Generates the following control flow structure: @@ -201,38 +203,43 @@ class KernelSupportLibrary { // `true_block_generator()`; // else // `false_block_generator()`; - Status If(absl::string_view name, llvm::Value* condition, - const std::function& true_block_generator, - const std::function& false_block_generator = - []() -> Status { return Status::OK(); }); + Status IfWithStatus( + absl::string_view name, llvm::Value* condition, + const std::function& true_block_generator, + const std::function& false_block_generator = []() -> Status { + return Status::OK(); + }); - Status If(llvm::Value* condition, - const std::function& true_block_generator, - const std::function& false_block_generator = - []() -> Status { return Status::OK(); }) { - return If("", condition, true_block_generator, false_block_generator); + Status IfWithStatus( + llvm::Value* condition, + const std::function& true_block_generator, + const std::function& false_block_generator = []() -> Status { + return Status::OK(); + }) { + return IfWithStatus("", condition, true_block_generator, + false_block_generator); } - void IfReturnVoid(llvm::Value* condition, - const std::function& true_block_generator, - const std::function& false_block_generator = []() { - }) { - IfReturnVoid("", condition, true_block_generator, false_block_generator); + void If( + llvm::Value* condition, const std::function& true_block_generator, + const std::function& false_block_generator = []() {}) { + If("", condition, true_block_generator, false_block_generator); } - void IfReturnVoid(absl::string_view name, llvm::Value* condition, - const std::function& true_block_generator, - const std::function& false_block_generator = []() { - }) { - TF_CHECK_OK(If(name, condition, - [&]() { - true_block_generator(); - return Status::OK(); - }, - [&]() { - false_block_generator(); - return Status::OK(); - })); + void If( + absl::string_view name, llvm::Value* condition, + const std::function& true_block_generator, + const std::function& false_block_generator = []() {}) { + TF_CHECK_OK(IfWithStatus( + name, condition, + [&]() { + true_block_generator(); + return Status::OK(); + }, + [&]() { + false_block_generator(); + return Status::OK(); + })); } using ArgumentVector = absl::Span; diff --git a/tensorflow/compiler/xla/service/llvm_ir/kernel_tiling.cc b/tensorflow/compiler/xla/service/llvm_ir/kernel_tiling.cc index e5fbdbd51b8a9aa14decadedd1eeb3bdbf831738..cd8dd72cd775d5e0b52f96a2326367da0775e7eb 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/kernel_tiling.cc +++ b/tensorflow/compiler/xla/service/llvm_ir/kernel_tiling.cc @@ -52,6 +52,29 @@ Shape MergeDimensions(absl::Span segs, const Shape& shape) { return ShapeUtil::MakeShapeWithDescendingLayout(shape.element_type(), dimensions); } + +// Given an index for a shape, return the equivalent new index if the shape is +// reshaped to another shape. +IrArray::Index GetReshapedIndex(const IrArray::Index& index, const Shape& shape, + const Shape& reshaped_shape, + llvm::IRBuilder<>* b) { + auto bounds = shape.dimensions(); + auto minor_to_major = shape.layout().minor_to_major(); + llvm::Value* linear_index = index.GetConstantWithIndexType(0); + int64 multiplier = 1; + for (int i = 0; i < index.size(); ++i) { + int64 dim = minor_to_major[i]; + llvm::Value* addend = b->CreateMul( + index[dim], index.GetConstantWithIndexType(multiplier), "linearizing", + /*HasNUW=*/true, /*HasNSW=*/true); + linear_index = b->CreateAdd(linear_index, addend, "", + /*HasNUW=*/true, /*HasNSW=*/true); + multiplier *= bounds[dim]; + } + + return IrArray::Index(linear_index, reshaped_shape, b); +} + } // namespace absl::optional > FindTranspose021(const Shape& a, @@ -60,28 +83,30 @@ absl::optional > FindTranspose021(const Shape& a, return absl::nullopt; } - std::vector perm(a.dimensions().size()); - { - auto layout_a_orig = LayoutUtil::MinorToMajor(a); - std::vector layout_a(layout_a_orig.rbegin(), layout_a_orig.rend()); - auto layout_b_orig = LayoutUtil::MinorToMajor(b); - std::vector layout_b(layout_b_orig.rbegin(), layout_b_orig.rend()); - for (size_t i = 0; i < perm.size(); ++i) { - perm[i] = PositionInContainer(layout_b, layout_a[i]); - } + std::vector permutation(a.dimensions().size()); + absl::Span minor_to_major_a = LayoutUtil::MinorToMajor(a); + std::vector major_to_minor_a(minor_to_major_a.rbegin(), + minor_to_major_a.rend()); + absl::Span minor_to_major_b = LayoutUtil::MinorToMajor(b); + std::vector major_to_minor_b(minor_to_major_b.rbegin(), + minor_to_major_b.rend()); + for (size_t i = 0; i < permutation.size(); ++i) { + permutation[i] = PositionInContainer(major_to_minor_b, major_to_minor_a[i]); } - auto segs = ConsecutiveSegments(perm); - if ((3 == segs.size() && 0 == perm[0]) || 2 == segs.size()) { - Shape norm_a = + + std::vector segments = ConsecutiveSegments(permutation); + if ((3 == segments.size() && 0 == permutation[0]) || 2 == segments.size()) { + Shape descending_layout_shape = ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout(a); - Shape reduced_a = MergeDimensions(segs, norm_a); - auto reduced_a_dims = reduced_a.dimensions(); + Shape normalized_shape = MergeDimensions(segments, descending_layout_shape); + absl::Span normalized_dims = + AsInt64Slice(normalized_shape.dimensions()); std::vector dims_021; - if (2 == segs.size()) { + if (2 == segments.size()) { // The logical component-0 is of size one. - dims_021 = {1, reduced_a_dims[1], reduced_a_dims[0]}; + dims_021 = {1, normalized_dims[1], normalized_dims[0]}; } else { - dims_021 = {reduced_a_dims[0], reduced_a_dims[2], reduced_a_dims[1]}; + dims_021 = {normalized_dims[0], normalized_dims[2], normalized_dims[1]}; } return dims_021; @@ -90,27 +115,120 @@ absl::optional > FindTranspose021(const Shape& a, return absl::nullopt; } -IrArray::Index GetUnreducedOutputIndex( - const IrArray::Index& reduced_output_index, - const Shape& reduced_output_shape, const Shape& unreduced_output_shape, - llvm::IRBuilder<>* b) { - auto bounds = reduced_output_shape.dimensions(); - auto minor_to_major = reduced_output_shape.layout().minor_to_major(); - llvm::Value* linear_index = reduced_output_index.GetConstantWithIndexType(0); - int64 multiplier = 1; - for (int i = 0; i < reduced_output_index.size(); ++i) { - int64 dim = minor_to_major[i]; - llvm::Value* addend = - b->CreateMul(reduced_output_index[dim], - reduced_output_index.GetConstantWithIndexType(multiplier), - "linearizing", - /*HasNUW=*/true, /*HasNSW=*/true); - linear_index = b->CreateAdd(linear_index, addend, "", - /*HasNUW=*/true, /*HasNSW=*/true); - multiplier *= bounds[dim]; +KernelMappingScheme::KernelMappingScheme( + absl::Span dims_in_elems, int64 tile_size_y, int64 tile_size_x, + absl::Span req_block_sizes, int64 num_threads_y, + int64 num_threads_x, llvm::IRBuilder<>* b) + : b_(b), + 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), + dilated_x_(true) { + DCHECK_EQ(dims_in_elems_.size(), 3); + DCHECK_EQ(req_block_sizes.size(), 3); + + DCHECK_EQ(tile_size_y % num_threads_y_, 0); + DCHECK_EQ(tile_size_x % num_threads_x_, 0); + + dims_in_tiles_ = ElementWiseCeilOfRatio(dims_in_elems_, tile_sizes_); + block_sizes_.reserve(req_block_sizes.size()); + absl::c_transform(req_block_sizes, dims_in_tiles_, + std::back_inserter(block_sizes_), + [](const int64 requested_size, const int64 max_size) { + return std::min(requested_size, max_size); + }); + dims_in_blocks_ = ElementWiseCeilOfRatio(dims_in_tiles_, block_sizes_); + + VLOG(10) << "dims_in_elems_ = [" << absl::StrJoin(dims_in_elems_, ",") << "]"; + VLOG(10) << "dims_in_tiles_ = [" << absl::StrJoin(dims_in_tiles_, ",") << "]"; + VLOG(10) << "dims_in_blocks_ = [" << absl::StrJoin(dims_in_blocks_, ",") + << "]"; +} + +IrArray::Index KernelMappingScheme::GetUnnormalizedIndex( + const IrArray::Index& normalized_shape_index, + const Shape& unnormalized_shape) { + DCHECK_EQ(normalized_shape_index.size(), dims_in_elems_.size()); + Shape output_shape = ShapeUtil::MakeShapeWithDescendingLayout( + unnormalized_shape.element_type(), GetDimensionsInElements()); + return GetReshapedIndex(normalized_shape_index, output_shape, + unnormalized_shape, b_); +} + +IrArray::Index KernelMappingScheme::EmitBlockIndex(llvm::Type* index_ty) { + llvm::Value* block_id = llvm_ir::EmitCallToIntrinsic( + llvm::Intrinsic::nvvm_read_ptx_sreg_ctaid_x, {}, {}, b_); + llvm_ir::AddRangeMetadata(0, GetNumberOfBlocks(), + llvm::cast(block_id)); + llvm::Value* linear_block_id = + b_->CreateIntCast(block_id, index_ty, /*isSigned=*/true, "block.id.x"); + return IrArray::Index(linear_block_id, + ShapeUtil::MakeShapeWithDescendingLayout( + PRED /*arbitrary*/, dims_in_blocks_), + b_); +} + +IrArray::Index KernelMappingScheme::GetTileIndexForBlockOrigin( + const IrArray::Index& block_index) { + DCHECK_EQ(block_index.size(), block_sizes_.size()); + std::vector multidim; + multidim.reserve(block_sizes_.size()); + for (int i = 0; i < block_sizes_.size(); ++i) { + multidim.push_back(b_->CreateMul( + block_index[i], + llvm::ConstantInt::get(block_index[i]->getType(), block_sizes_[i]), + "block_origin." + std::to_string(i))); + } + return IrArray::Index(multidim, block_index[0]->getType()); +} + +IrArray::Index KernelMappingScheme::GetElementIndexForTileOrigin( + const IrArray::Index& tile_index) { + IrArray::Index elem_index = tile_index; + for (int i = DimY; i < DimTot; ++i) { + elem_index[i] = + b_->CreateMul(tile_index[i], + llvm::ConstantInt::get(tile_index[i]->getType(), + GetTileSizeForDimension(i)), + "tile_origin." + std::to_string(i)); } + return elem_index; +} + +llvm::GlobalVariable* KernelMappingScheme::GetSharedMemoryBufferForElementType( + llvm::Type* elem_ty, absl::string_view buffer_name) { + // If shared memory tranpose is needed, we use square tiles. + CHECK_EQ(GetTileSizeForDimensionX(), GetTileSizeForDimensionY()); + + // For Nvidia GPUs, the warp size is 32 threads and the shared memory bank is + // organized into 32-way. We usually use the warp size or a multiplier or a + // the warp size as the size for tiling. This may cause all elements in the + // same column of a tile use the same memory bank and therefore shared memory + // bank conflicts. Adding 1 to the minor dimension of the shared memory buffer + // can reduce such shared memory bank conflicts. + llvm::Type* buffer_type = llvm::ArrayType::get( + llvm::ArrayType::get(elem_ty, GetTileSizeForDimension(DimX) + 1), + GetTileSizeForDimension(DimY)); + return llvm_ir::AllocateSharedMemoryTile(b_->GetInsertBlock()->getModule(), + buffer_type, buffer_name); +} - return IrArray::Index(linear_index, unreduced_output_shape, b); +std::tuple +KernelMappingScheme::EmitThreadYXCoordinate(llvm::Type* index_ty) { + // Calculate (y, x) coordinate of the thread in the 2D view of thread block + // defined by (num_thread_y, num_thread_x) from thread_id. + llvm::CallInst* thread_id_raw = llvm_ir::EmitCallToIntrinsic( + llvm::Intrinsic::nvvm_read_ptx_sreg_tid_x, {}, {}, b_); + llvm_ir::AddRangeMetadata(0, GetThreadsPerBlock(), thread_id_raw); + llvm::Value* thread_id_int = + b_->CreateIntCast(thread_id_raw, index_ty, + /*isSigned=*/true, "thread.id.x"); + llvm::Value* num_thread_x = + llvm::ConstantInt::get(index_ty, GetNumberOfThreadsForDimensionX()); + llvm::Value* x = b_->CreateURem(thread_id_int, num_thread_x, "thread.x"); + llvm::Value* y = b_->CreateUDiv(thread_id_int, num_thread_x, "thread.y"); + return std::make_tuple(y, x); } } // namespace llvm_ir diff --git a/tensorflow/compiler/xla/service/llvm_ir/kernel_tiling.h b/tensorflow/compiler/xla/service/llvm_ir/kernel_tiling.h index 5ea05b3188a1c0881e4c0c41625d530aff1b1205..f802cc27d519e621262f328903697373aa8c284c 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/kernel_tiling.h +++ b/tensorflow/compiler/xla/service/llvm_ir/kernel_tiling.h @@ -28,23 +28,185 @@ namespace llvm_ir { // If a shape can be viewed as three logical components 0-1-2 in the order of // major to minor, a 0-2-1-transpose changes the order of such logical // components to 0-2-1. We call the shape being transposed the input shape and -// the transposed shape the output shape. The logical view of the input and -// output shapes for the transpose are called the 0-1-2 shape or reduced input -// shape and the 0-2-1 shape or the reduced output shape respectively. The -// original input and output shapes are called the unreduced input and output -// shapes. - +// the transposed shape the output shape. The logical view of the input/output +// shapes for the transpose are called the 0-1-2/0-2-1 shapes or the normalized +// shapes. The original input/output shapes are called unnormalized shapes. +// // If `b` is a 0-2-1 transpose of `a` in 0-1-2, return the dimensions for the -// reduced shape of `b` or the 0-2-1 shape. +// normalized shape of `b` or the 0-2-1 shape. absl::optional > FindTranspose021(const Shape& a, const Shape& b); -// Return the unreduced output index corresponding to the given reduced output -// index. -IrArray::Index GetUnreducedOutputIndex( - const IrArray::Index& reduced_output_index, - const Shape& reduced_output_shape, const Shape& unreduced_output_shape, - llvm::IRBuilder<>* b); +// A tile is a spatial subdivision of a tensor. We group tensor elements into +// tiles so that we can launch kernels to process the tensor elements in blocks +// of tiles. +// +// A kernel mapping scheme describes a method to partition the tensors accessed +// by an unnested HLO instruction into tiles and blocks of tiles, and the +// associated information to use hardware threads to process the tensor elements +// in blocks of tiles. +// +// Currently, there are two main use cases for a tiling scheme. First, we +// implement kernels with 0-2-1 memory transpose using shared memory to improve +// memory access pattern. Second, we implement reduction to contiguous +// dimensions in layout, with or without memory tranpsose, to achieve better +// memory access pattern as well as to reduce the need numbers of executed +// expensive instructions, such as thread synchronization related instructions +// and atomic operations. For both use cases, we can apply a normalization to +// the original tensors, to collapse contiguous dimensions for the same purpose +// and produce normlized three dimensional tensors. For this reason, the tiling +// scheme class only needs to handle normalized three dimensional tensors and +// two dimensional tiles. +// +// The current implementation of the class is somewhat NVIDIA GPU oriented. This +// situation can be improved when there is a need though. The idea of 0-2-1 +// transpose using shared memory can be found in the following CUDA algorithm in +// TensorFlow: https://goo.gl/MStRV6. +// +// We use a thread block to process a tile because we want to use the HW thread +// block synchronization primitives to synchronize the processing of all the +// elements in the same tile. A thread block can be viewed as a two dimensional +// array of threads, described by the number of threads for the Y and X +// dimensions. A thread block (num_threads_y, num_threads_x) processes a tile of +// (tile_size_y, tile_size_x) as follows: each thread in the thread block +// processes one element in the tile so that all the threads in the thread block +// together process a subdivision of the tile that has the same dimension as the +// thread block array. Then the thread block moves on to process the next +// subdivision of the tile until the whole tile is processed. Therefore, each +// thread in the thread block processes +// tile_size_x/num_threads_x * tile_size_y/num_threads_y elements in a tile. +// +// There are situations where we want a thread block to process multiple +// tiles. We can't group those tiles into a bigger tiles because we limit a tile +// to a two dimensional spatial subdivision of a tensor. For example, when we +// use tiling to implement reduction with tranpose, we want the partial sum +// produced by each thread to accumulate values for more elements before using +// shlf_down and atomic_add instructions for further reduction, to amortize the +// cost of such expensive instructions. The concept of tile block is introduced +// for this purpose. A tile block is a three dimensional array of tiles, of +// which some dimensions may be degenerated to only one tile. +class KernelMappingScheme { + public: + enum { DimZ = 0, DimY, DimX, DimTot }; + + public: + KernelMappingScheme() {} + // dims_in_elems: the normalized tensor dimensions. + // req_block_sizes: the requested block size in number of tiles for each + // dimension. The actual block size is set to min(req_block_size, + // dims_in_number_of_blocks). + KernelMappingScheme(absl::Span dims_in_elems, int64 tile_size_y, + int64 tile_size_x, + absl::Span req_block_sizes, + int64 num_threads_y, int64 num_threads_x, + llvm::IRBuilder<>* b); + + absl::Span GetDimensionsInElements() const { + return dims_in_elems_; + } + absl::Span GetDimensionsInTiles() const { + return dims_in_tiles_; + } + absl::Span GetDimensionsInBlocks() const { + return dims_in_blocks_; + } + + int64 GetNumberOfTilesInTotal() const { + return absl::c_accumulate(dims_in_tiles_, 1LL, std::multiplies()); + } + 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()); + } + + int64 GetTileSizeForDimension(int d) const { + DCHECK(d >= DimZ && d <= DimX); + return tile_sizes_[d]; + } + int64 GetTileSizeForDimensionX() const { + return GetTileSizeForDimension(DimX); + } + int64 GetTileSizeForDimensionY() const { + return GetTileSizeForDimension(DimY); + } + + absl::Span GetBlockSizes() const { return block_sizes_; } + int64 GetTileBlockSizeForDimension(int d) const { + DCHECK(d >= DimZ && d <= DimX); + return dims_in_blocks_[d]; + } + + int64 GetNumberOfThreadsForDimensionX() const { return num_threads_x_; } + int64 GetNumberOfThreadsForDimensionY() const { return num_threads_y_; } + + int64 GetThreadsPerBlock() const { + return GetNumberOfThreadsForDimensionX() * + 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. + IrArray::Index GetTileIndexForBlockOrigin(const IrArray::Index& block_index); + // Returns the index for the first element in the tile with the given tile + // index. + IrArray::Index GetElementIndexForTileOrigin(const IrArray::Index& tile_index); + + std::tuple EmitThreadYXCoordinate( + llvm::Type* index_ty); + + IrArray::Index GetUnnormalizedIndex( + const IrArray::Index& normalized_shape_index, + const Shape& unnormalized_shape); + + llvm::GlobalVariable* GetSharedMemoryBufferForElementType( + llvm::Type* elem_ty, absl::string_view buffer_name); + + private: + llvm::IRBuilder<>* b_; + // The number of elements in each dimension. + std::vector dims_in_elems_; + + // The number of elements for each dimension of a tile. + std::vector tile_sizes_; + // The number of tiles in each dimension. It is computed from dims_in_elem_ + // and tile_sizes_. + std::vector dims_in_tiles_; + + // The number of tiles for each dimension of a tile block. + std::vector block_sizes_; + // The number of blocks in each dimension of a tile block. It is computed from + // dims_in_tile_ and block_sizes_. + std::vector dims_in_blocks_; + + // Number of threads used to process elements in the X direction of a tile. + 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 // for 021 transpose. diff --git a/tensorflow/compiler/xla/service/llvm_ir/llvm_loop.cc b/tensorflow/compiler/xla/service/llvm_ir/llvm_loop.cc index 219a9f221fbd116cdfbaf17985e21d82aefd079d..fe320bbe727111fbc986cc1fbc217feed74d30f1 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/llvm_loop.cc +++ b/tensorflow/compiler/xla/service/llvm_ir/llvm_loop.cc @@ -235,7 +235,7 @@ std::unique_ptr ForLoopNest::AddLoop(int64 start_index, IrArray::Index ForLoopNest::AddLoopsForShape(const Shape& shape, absl::string_view suffix) { - std::vector dimensions(ShapeUtil::Rank(shape)); + std::vector dimensions(shape.rank()); std::iota(dimensions.begin(), dimensions.end(), 0); return AddLoopsForShapeOnDimensions(shape, dimensions, suffix); } diff --git a/tensorflow/compiler/xla/service/llvm_ir/llvm_util.cc b/tensorflow/compiler/xla/service/llvm_ir/llvm_util.cc index 1a53c026be340ca3bec3a49b11666d6124728130..807296329c07b8e4ac630486a1e1f59e4fdfa009 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/llvm_util.cc +++ b/tensorflow/compiler/xla/service/llvm_ir/llvm_util.cc @@ -19,10 +19,12 @@ limitations under the License. #include #include +#include "absl/base/casts.h" #include "absl/strings/match.h" #include "absl/strings/str_cat.h" #include "llvm/IR/DerivedTypes.h" #include "llvm/IR/GlobalValue.h" +#include "llvm/IR/GlobalVariable.h" #include "llvm/IR/MDBuilder.h" #include "llvm/IR/Operator.h" #include "llvm/Target/TargetOptions.h" @@ -33,7 +35,6 @@ limitations under the License. #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/util.h" -#include "tensorflow/core/lib/core/casts.h" #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/io/path.h" #include "tensorflow/core/platform/byte_order.h" @@ -83,10 +84,9 @@ string DumpModuleToString(const llvm::Module& module) { return AsString(buffer_string); } -llvm::Value* EmitCallToIntrinsic(llvm::Intrinsic::ID intrinsic_id, - absl::Span operands, - absl::Span overloaded_types, - llvm::IRBuilder<>* b) { +llvm::CallInst* EmitCallToIntrinsic( + llvm::Intrinsic::ID intrinsic_id, absl::Span operands, + absl::Span overloaded_types, llvm::IRBuilder<>* b) { llvm::Module* module = ModuleFromIRBuilder(b); llvm::Function* intrinsic = llvm::Intrinsic::getDeclaration( module, intrinsic_id, AsArrayRef(overloaded_types)); @@ -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: @@ -219,10 +228,10 @@ int GetSizeInBits(llvm::Type* type) { llvm::Type* ShapeToIrType(const Shape& shape, llvm::Module* module) { llvm::Type* result_type = PrimitiveTypeToIrType(shape.element_type(), module); - if (ShapeUtil::IsTuple(shape)) { + if (shape.IsTuple()) { // A tuple buffer is an array of pointers. result_type = llvm::ArrayType::get(result_type, shape.tuple_shapes_size()); - } else if (ShapeUtil::IsArray(shape)) { + } else if (shape.IsArray()) { for (int64 dimension : LayoutUtil::MinorToMajor(shape)) { result_type = llvm::ArrayType::get(result_type, shape.dimensions(dimension)); @@ -244,10 +253,11 @@ StatusOr EncodeSelfDescribingShapeConstant(const Shape& shape, StatusOr DecodeSelfDescribingShapeConstant(const void* shape_ptr, int32 size_bytes) { - Shape shape; - TF_RET_CHECK(shape.ParseFromArray(shape_ptr, size_bytes)); + ShapeProto shape_proto; + TF_RET_CHECK(shape_proto.ParseFromArray(shape_ptr, size_bytes)); + Shape shape(shape_proto); TF_RETURN_IF_ERROR(ShapeUtil::ValidateShape(shape)); - return shape; + return std::move(shape); } llvm::Constant* ConvertLiteralToIrConstant(const Literal& literal, @@ -260,6 +270,17 @@ llvm::Constant* ConvertLiteralToIrConstant(const Literal& literal, /*AddNull=*/false); } +llvm::GlobalVariable* AllocateSharedMemoryTile(llvm::Module* module, + llvm::Type* tile_type, + absl::string_view name) { + const int kNVPTXSharedMemoryAddrSpace = 3; + return new llvm::GlobalVariable( + *module, tile_type, + /*isConstant=*/false, llvm::GlobalValue::PrivateLinkage, + llvm::UndefValue::get(tile_type), AsStringRef(name), nullptr, + llvm::GlobalValue::NotThreadLocal, kNVPTXSharedMemoryAddrSpace); +} + llvm::AllocaInst* EmitAllocaAtFunctionEntry(llvm::Type* type, absl::string_view name, llvm::IRBuilder<>* b, @@ -362,11 +383,10 @@ static void LogS64(const char* tag, int64 value) { void EmitLogging(const char* tag, llvm::Value* value, llvm::IRBuilder<>* b) { llvm::FunctionType* log_function_type = llvm::FunctionType::get( b->getVoidTy(), {b->getInt64Ty(), b->getInt64Ty()}, /*isVarArg=*/false); - b->CreateCall( - log_function_type, - b->CreateIntToPtr(b->getInt64(tensorflow::bit_cast(&LogS64)), - log_function_type->getPointerTo()), - {b->getInt64(tensorflow::bit_cast(tag)), value}); + b->CreateCall(log_function_type, + b->CreateIntToPtr(b->getInt64(absl::bit_cast(&LogS64)), + log_function_type->getPointerTo()), + {b->getInt64(absl::bit_cast(tag)), value}); } void SetAlignmentMetadataForLoad(llvm::LoadInst* load, uint64_t alignment) { @@ -610,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/llvm_util.h b/tensorflow/compiler/xla/service/llvm_ir/llvm_util.h index f59baff263fe7184c6b0821c9dbd9eee205586a6..c604c7c870adf734a29017e6accbd159317a9548 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/llvm_util.h +++ b/tensorflow/compiler/xla/service/llvm_ir/llvm_util.h @@ -24,6 +24,7 @@ limitations under the License. #include "absl/types/span.h" #include "llvm/ADT/StringRef.h" #include "llvm/IR/BasicBlock.h" +#include "llvm/IR/GlobalVariable.h" #include "llvm/IR/IRBuilder.h" #include "llvm/IR/Instructions.h" #include "llvm/IR/Module.h" @@ -101,10 +102,9 @@ string SanitizeFunctionName(string function_name); // intrinsics (for example, "minnum") must include a type in overloaded_types // for each overloaded type. Typically, overloaded intrinsics have only a single // overloaded type. -llvm::Value* EmitCallToIntrinsic(llvm::Intrinsic::ID intrinsic_id, - absl::Span operands, - absl::Span overloaded_types, - llvm::IRBuilder<>* b); +llvm::CallInst* EmitCallToIntrinsic( + llvm::Intrinsic::ID intrinsic_id, absl::Span operands, + absl::Span overloaded_types, llvm::IRBuilder<>* b); // Emit float max. Emit maxnum intrinsic is fast math is disabled, or // fcmp+select otherwise @@ -155,6 +155,11 @@ StatusOr DecodeSelfDescribingShapeConstant(const void* shape_ptr, llvm::Constant* ConvertLiteralToIrConstant(const Literal& literal, llvm::Module* module); +// Allocates a tile of shared memory. +llvm::GlobalVariable* AllocateSharedMemoryTile(llvm::Module* module, + llvm::Type* tile_type, + absl::string_view name); + // Inserts an allocate of the requested type at the entry point of the // function that the builder is currently building. The insert point // of the builder is set to the same place after calling this function diff --git a/tensorflow/compiler/xla/service/llvm_ir/sort_util.cc b/tensorflow/compiler/xla/service/llvm_ir/sort_util.cc index 05ba4a40da413f0e774214e55ef69d023afc48e2..89b6a36f96beedbcb7322e6164ac59221650d3d8 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/sort_util.cc +++ b/tensorflow/compiler/xla/service/llvm_ir/sort_util.cc @@ -18,7 +18,9 @@ limitations under the License. #include // IWYU pragma: no_include "llvm/IR/Intrinsics.gen.inc" +#include "absl/strings/str_cat.h" #include "absl/strings/string_view.h" +#include "absl/types/span.h" #include "llvm/ADT/APInt.h" #include "llvm/IR/BasicBlock.h" #include "llvm/IR/Constants.h" @@ -28,10 +30,12 @@ limitations under the License. #include "tensorflow/compiler/xla/service/gpu/parallel_loop_emitter.h" #include "tensorflow/compiler/xla/service/gpu/partition_assignment.h" #include "tensorflow/compiler/xla/service/llvm_ir/ir_array.h" +#include "tensorflow/compiler/xla/service/llvm_ir/kernel_support_library.h" #include "tensorflow/compiler/xla/service/llvm_ir/llvm_loop.h" #include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h" #include "tensorflow/compiler/xla/service/llvm_ir/loop_emitter.h" #include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/util.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/platform/types.h" @@ -39,147 +43,364 @@ namespace xla { namespace llvm_ir { namespace { -// Adds the inner comparison loop where we compare elements pointed to by -// 'keys_index' and 'compare_keys_index'. -void EmitCompareLoop(int64 dimension_to_sort, const IrArray::Index& keys_index, - const IrArray::Index& compare_keys_index, - const IrArray& keys_array, - const std::vector& values_arrays, - llvm::IRBuilder<>* b) { - // if (is_smaller_index && - // compare_keys[dimension_to_sort] < dimension_to_sort_bound) - llvm::Value* is_smaller_index = b->CreateICmpSLT( - keys_index[dimension_to_sort], compare_keys_index[dimension_to_sort]); - int64 dimension_to_sort_bound = - keys_array.GetShape().dimensions(dimension_to_sort); - auto if_data = EmitIfThenElse( - b->CreateAnd(is_smaller_index, - b->CreateICmpSLT(compare_keys_index[dimension_to_sort], - keys_index.GetConstantWithIndexType( - dimension_to_sort_bound))), - "smaller_comparison_index", b, /*emit_else=*/false); - SetToFirstInsertPoint(if_data.true_block, b); - auto key1 = keys_array.EmitReadArrayElement(keys_index, b); - auto key2 = keys_array.EmitReadArrayElement(compare_keys_index, b); - auto compare_key1 = key1; - auto compare_key2 = key2; - auto key_type = keys_array.GetShape().element_type(); - 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; + +// 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, + int64 xor_mask, llvm::Type* index_type, + std::function read_element, + std::function + write_element, + llvm::IRBuilder<>* b, bool needs_bounds_checks = true) { + auto index_typed_constant = [&](int64 value) { + return llvm::ConstantInt::get(index_type, value); + }; + // The 'xor_mask' determines which elements are compared against each other. + // Index 'current_keys_index' will be compared with 'current_keys_index' xor + // 'xor_mask'. This means that we will always compare a block of consecutive + // elements against elements from the adjacent block of the same size. When + // 'xor_mask' is a power of 2, it immediately identifies the size of such a + // block. We can also have 'xor_mask' being 2^k - 1 (for some value of k). In + // that case, we essentially flip the last 'k' - 1 bits when computing the + // position of the element to compare to, so the block size is 2^(k - 1). + int64 block_size = xor_mask; + // Check if it is a value 2^k - 1. + if (xor_mask > 1 && (xor_mask & (xor_mask + 1)) == 0) { + block_size = (xor_mask + 1) / 2; } - auto comparison = - b->CreateICmp(is_signed_comparison ? llvm::ICmpInst::ICMP_SLT - : llvm::ICmpInst::ICMP_ULT, - compare_key2, compare_key1); - // If key2 < key1 - auto if_smaller_data = - EmitIfThenElse(comparison, "is_smaller_than", b, /*emit_else=*/false); - SetToFirstInsertPoint(if_smaller_data.true_block, b); - // Swap key1 with key2. - keys_array.EmitWriteArrayElement(keys_index, key2, b); - keys_array.EmitWriteArrayElement(compare_keys_index, key1, b); - for (const auto& values_array : values_arrays) { - // Also swap the values. - auto value1 = values_array.EmitReadArrayElement(keys_index, b); - auto value2 = values_array.EmitReadArrayElement(compare_keys_index, b); - values_array.EmitWriteArrayElement(keys_index, value2, b); - values_array.EmitWriteArrayElement(compare_keys_index, value1, b); + auto current_keys_index = element_pair_index; + if (block_size == 1) { + // If the block size is 1, we take every second element and compare it to + // the next one. + current_keys_index = + b->CreateMul(current_keys_index, index_typed_constant(2)); + } else if (block_size * 2 < iteration_bound) { + // current_keys_index iterates through the 'left' elements of the element + // pairs to be compared. We first need to compute the comparison block to + // which the element belongs. The block id of that block is index / + // block_size. + auto block_id = + b->CreateUDiv(current_keys_index, index_typed_constant(block_size)); + // The index of the 'left' element within its block is simply the remainder + // when dividing by 'block_size'. + auto index_within_block = + b->CreateURem(current_keys_index, index_typed_constant(block_size)); + // The first element of the 'left' block of elements that is compared + // against elements from the adjacent 'right' block of elements is + // 'block_id' * (2 * 'block_size'). + auto first_element_in_block = + b->CreateMul(block_id, index_typed_constant(2 * block_size)); + current_keys_index = + b->CreateAdd(first_element_in_block, index_within_block); } + auto compare_keys_index = + b->CreateXor(current_keys_index, index_typed_constant(xor_mask)); + // current_keys_index < compare_keys_index + llvm::Value* is_smaller_index = + b->CreateICmpSLT(current_keys_index, compare_keys_index); + // compare_keys_index < iteration_bound + llvm::Value* index_is_inbounds = b->CreateICmpSLT( + compare_keys_index, index_typed_constant(iteration_bound)); + llvm::Value* do_comparison = + needs_bounds_checks ? b->CreateAnd(is_smaller_index, index_is_inbounds) + : b->getInt1(true); + + // 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)); + } + 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); + } + }); + }); +} + +void 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) { + KernelSupportLibrary ksl(b); + llvm::Value* thread_id = llvm_ir::EmitCallToIntrinsic( + llvm::Intrinsic::nvvm_read_ptx_sreg_tid_x, {}, {}, b); + llvm_ir::AddRangeMetadata(0, tile_size / 2, + llvm::cast(thread_id)); + thread_id = b->CreateIntCast(thread_id, tiled_keys_index.GetType(), + /*isSigned=*/true, "thread.id.x"); + + auto copy_loop_body = + [&](std::function + read_or_write) { + auto value_one = tiled_keys_index.GetConstantWithIndexType(1); + auto current_keys_index = + b->CreateShl(tiled_keys_index[dimension_to_sort], value_one); + // We want to copy two adjacent elements. We first check whether the + // first index position is within bounds. + ksl.If( + "smaller_keys_index", + b->CreateICmpSLT(current_keys_index, + tiled_keys_index.GetConstantWithIndexType( + dimension_to_sort_bound)), + [&]() { + auto cache_index = b->CreateShl(thread_id, value_one); + read_or_write(cache_index, current_keys_index); + // Increment to go 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", + b->CreateICmpSLT(current_keys_index, + tiled_keys_index.GetConstantWithIndexType( + dimension_to_sort_bound)), + [&]() { + cache_index = b->CreateAdd(cache_index, value_one); + read_or_write(cache_index, current_keys_index); + }); + }); + }; + + // Copy operand tiles from the operand buffers to shared memory. + IrArray::Index keys_index = tiled_keys_index; + for (int64 i = 0; i < params.size(); ++i) { + copy_loop_body([&](llvm::Value* cache_index, llvm::Value* index) { + keys_index[dimension_to_sort] = index; + auto value = params[i].EmitReadArrayElement(keys_index, b); + b->CreateStore(value, + b->CreateGEP(param_shmem_buffers[i], + {tiled_keys_index.GetConstantWithIndexType(0), + cache_index})); + }); + } + // Wait until all reads have happened. + 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( + b->CreateGEP(param_shmem_buffers[operand], + {tiled_keys_index.GetConstantWithIndexType(0), index})); + }; + auto write_element = [&](int64 operand, llvm::Value* index, + llvm::Value* value) { + b->CreateStore( + value, + b->CreateGEP(param_shmem_buffers[operand], + {tiled_keys_index.GetConstantWithIndexType(0), index})); + }; + for (int64 xor_mask : xor_masks) { + // The index of the element pair to be compared within the tile stored in + // shared memory. We order the element pairs by the element with the smaller + // index. + auto element_pair_index = thread_id; + // If 'dimension_to_sort_bound' is evenly divisible by 'tile_size', we don't + // need any bounds checks. + 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( + "is_last_tile", + b->CreateICmpUGE( + b->CreateMul(tiled_keys_index[dimension_to_sort], + tiled_keys_index.GetConstantWithIndexType(2)), + 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); + }, + [&]() { + 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); + }); + } 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); + } + // Wait until all comparisons have happened. + llvm_ir::EmitCallToIntrinsic(llvm::Intrinsic::nvvm_barrier0, {}, {}, b); + } + + // Copy the operand tiles back from shared memory to the operand buffers. + for (int64 i = 0; i < params.size(); ++i) { + copy_loop_body([&](llvm::Value* cache_index, llvm::Value* index) { + keys_index[dimension_to_sort] = index; + auto value = b->CreateLoad(b->CreateGEP( + param_shmem_buffers[i], + {tiled_keys_index.GetConstantWithIndexType(0), cache_index})); + params[i].EmitWriteArrayElement(keys_index, value, b); + }); + } + // We should normally synchronize here to make sure all writes have happened. + // However the very next thing each thread does is reading 2 elements from the + // operand buffer and writing it into the same location in shared memory from + // which it previously copied it to the operand buffer, and we synchronize + // after this has happened. We can be sure that a thread always writes to the + // 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; } } // namespace Status EmitSortInPlace(int64 dimension_to_sort, const IrArray& keys_array, const std::vector& values_arrays, - absl::string_view name, llvm::Value* xor_mask, - llvm::IRBuilder<>* b, - const gpu::LaunchDimensions* launch_dimensions) { - const Shape& keys_shape = keys_array.GetShape(); + 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) { + // 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 + // in tiles of 64 elements each, so we use another loop that happens within + // one thread to process this tile worth of data (thereby combining several + // comparison stages of the bitonic sort algorithm because they all happen + // within those 64 elements and are therefore independent of the other + // comparisons). - // Create loop nests which loop through the operand dimensions. The sort - // dimension is handled in the innermost loop which performs the sorting. - ForLoopNest loop_nest(name, b); - IrArray::Index keys_index = - loop_nest.EmitOperandArrayLoopNest(keys_array, dimension_to_sort, "keys"); - if (loop_nest.GetInnerLoopBodyBasicBlock() != nullptr) { - SetToFirstInsertPoint(loop_nest.GetInnerLoopBodyBasicBlock(), b); + const Shape& keys_shape = keys_array.GetShape(); + int64 rank = keys_shape.rank(); + int64 dimension_to_sort_bound = keys_shape.dimensions(dimension_to_sort); + std::vector dimensions_in_iteration_order(rank); + std::vector iteration_order_to_logical_order(rank); + int64 dim = 0; + for (int64 dimension : LayoutUtil::MinorToMajor(keys_shape)) { + if (dimension != dimension_to_sort) { + dimensions_in_iteration_order[dim] = keys_shape.dimensions(dimension); + iteration_order_to_logical_order[dim++] = dimension; + } } + dimensions_in_iteration_order[dim] = num_iterations_in_sort_dim; + iteration_order_to_logical_order[dim] = dimension_to_sort; - // 'compare_keys_index' is the index of the element that 'keys_index' should - // be compared to. - IrArray::Index compare_keys_index(keys_index.GetType()); - for (size_t dimension = 0; dimension < keys_index.size(); ++dimension) { - if (dimension != dimension_to_sort) { - compare_keys_index.push_back(keys_index[dimension]); - } else { - compare_keys_index.push_back(nullptr); + 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); + 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); + param_shmem_buffers[i] = llvm_ir::AllocateSharedMemoryTile( + module, tile_type, absl::StrCat(name, "_tile_param_", i)); } } - // Naive C++ code for the inner compare loop: - // - // for (int64 i = 0; i < dimension_to_sort_bound; ++i) { - // int64 j = i ^ xor_mask; - // if (i < j && j < dimension_to_sort_bound) { - // int64 min_key = std::min(keys[i], keys[j]); - // keys[j] = std::max(keys[i], keys[j]); - // keys[i] = min_key; - // } - // } - // - // This follows the algorithm described on Wikipedia: - // https://en.wikipedia.org/wiki/Bitonic_sorter - - int64 dimension_to_sort_bound = - keys_array.GetShape().dimensions(dimension_to_sort); - Shape compare_shape = ShapeUtil::MakeShape(keys_shape.element_type(), - {dimension_to_sort_bound}); auto compare_loop_body_emitter = - [&](const IrArray::Index& compare_index) -> Status { - keys_index[dimension_to_sort] = compare_index[0]; - compare_keys_index[dimension_to_sort] = - b->CreateXor(compare_index[0], xor_mask); - EmitCompareLoop(dimension_to_sort, keys_index, compare_keys_index, - keys_array, values_arrays, b); + [&](const IrArray::Index& tiles_index) -> Status { + // Naive C++ code for the inner compare loop: + // + // for (int64 i = 0; i < dimension_to_sort_bound; ++i) { + // int64 j = i ^ xor_mask; + // /* emitted in EmitCompareLoopBody() */ + // if (i < j && j < dimension_to_sort_bound) { + // int64 min_key = std::min(keys[i], keys[j]); + // keys[j] = std::max(keys[i], keys[j]); + // keys[i] = min_key; + // } + // } + // + // This follows the algorithm described on Wikipedia: + // https://en.wikipedia.org/wiki/Bitonic_sorter + IrArray::Index keys_index(tiles_index.GetType(), rank); + for (int64 i = 0; i < rank; ++i) { + 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); + } else { + auto read_element = [&](int64 operand, llvm::Value* index) { + keys_index[dimension_to_sort] = index; + return params[operand].EmitReadArrayElement(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); + }; + 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); + } return Status::OK(); }; - if (launch_dimensions != nullptr) { - TF_RETURN_IF_ERROR(gpu::ParallelLoopEmitter(compare_loop_body_emitter, - compare_shape, - *launch_dimensions, b) - .EmitLoop(name)); - } else { - TF_RETURN_IF_ERROR(LoopEmitter(compare_loop_body_emitter, compare_shape, b) - .EmitLoop(name)); - } - - // Set the IR builder insert point to the exit basic block of the outer most - // loop. This ensures later instructions are inserted after this loop nest. - b->SetInsertPoint(loop_nest.GetOuterLoopExitBasicBlock()); - - return Status::OK(); + return gpu::ParallelLoopEmitter(compare_loop_body_emitter, iteration_shape, + launch_dimensions, b) + .EmitLoop(name); } } // namespace llvm_ir diff --git a/tensorflow/compiler/xla/service/llvm_ir/sort_util.h b/tensorflow/compiler/xla/service/llvm_ir/sort_util.h index 2f3bcda2307bcbb35a03b9e71dbbe44e366b3820..685f9383acba416f51681270e4037d56abb4b6ea 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/sort_util.h +++ b/tensorflow/compiler/xla/service/llvm_ir/sort_util.h @@ -19,6 +19,7 @@ limitations under the License. #include #include "absl/strings/string_view.h" +#include "absl/types/span.h" #include "llvm/IR/Value.h" #include "tensorflow/compiler/xla/service/gpu/partition_assignment.h" #include "tensorflow/compiler/xla/service/llvm_ir/ir_array.h" @@ -29,13 +30,17 @@ namespace xla { namespace llvm_ir { // 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. If 'launch_dimensions' is nullptr, -// the inner compare loop will not be parallelized. +// 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, - absl::string_view name, llvm::Value* xor_mask, - llvm::IRBuilder<>* b, - const gpu::LaunchDimensions* launch_dimensions); + 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); } // namespace llvm_ir } // namespace xla diff --git a/tensorflow/compiler/xla/service/llvm_ir/tuple_ops.cc b/tensorflow/compiler/xla/service/llvm_ir/tuple_ops.cc index a60643bc754f896d096b3ca4e1216e77d7e384c6..d8d2700e1934fd202d44a1dc60e71a99913d4537 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/tuple_ops.cc +++ b/tensorflow/compiler/xla/service/llvm_ir/tuple_ops.cc @@ -93,7 +93,7 @@ llvm::Value* EmitGetTupleElement(const Shape& target_shape, int64 index, llvm::LoadInst* src_buffer = b->CreateLoad(element_ptr); // Mark the loaded pointer as dereferenceable if we know its shape. - if (!ShapeUtil::IsOpaque(target_shape)) { + if (!target_shape.IsOpaque()) { SetDereferenceableMetadataForLoad( src_buffer, ByteSizeOf(target_shape, src_buffer->getModule()->getDataLayout())); diff --git a/tensorflow/compiler/xla/service/local_service.cc b/tensorflow/compiler/xla/service/local_service.cc index cca37556173bb95ef062b59ab0a4bf9ca7c496fe..3470fe5b2c34bf832207ed546fad176319446f31 100644 --- a/tensorflow/compiler/xla/service/local_service.cc +++ b/tensorflow/compiler/xla/service/local_service.cc @@ -52,8 +52,10 @@ namespace xla { } BackendOptions backend_options; - backend_options.set_platform(platform).set_intra_op_parallelism_threads( - options.intra_op_parallelism_threads()); + backend_options.set_platform(platform) + .set_intra_op_parallelism_threads(options.intra_op_parallelism_threads()) + .set_allowed_devices(options.allowed_devices()); + TF_ASSIGN_OR_RETURN(std::unique_ptr backend, Backend::CreateBackend(backend_options)); @@ -96,44 +98,19 @@ ExecutionOptions CreateExecutionOptions( const ExecutableBuildOptions& build_options, const ProgramShape* program_shape) { ExecutionOptions execution_options = CreateDefaultExecutionOptions(); - if (build_options.hlo_profile().has_value()) { - execution_options.mutable_debug_options()->set_xla_hlo_profile( - *build_options.hlo_profile()); - } - if (build_options.generate_hlo_graph().has_value()) { - execution_options.mutable_debug_options()->set_xla_generate_hlo_graph( - build_options.generate_hlo_graph().value()); - } - if (build_options.dump_optimized_hlo_proto_to().has_value()) { - execution_options.mutable_debug_options() - ->set_xla_dump_optimized_hlo_proto_to( - build_options.dump_optimized_hlo_proto_to().value()); - } - if (build_options.dump_unoptimized_hlo_proto_to().has_value()) { - execution_options.mutable_debug_options() - ->set_xla_dump_unoptimized_hlo_proto_to( - build_options.dump_unoptimized_hlo_proto_to().value()); - } - if (build_options.dump_per_pass_hlo_proto_to().has_value()) { - execution_options.mutable_debug_options() - ->set_xla_dump_per_pass_hlo_proto_to( - build_options.dump_per_pass_hlo_proto_to().value()); + if (build_options.has_debug_options()) { + *execution_options.mutable_debug_options() = build_options.debug_options(); } if (build_options.result_layout() != nullptr) { *execution_options.mutable_shape_with_output_layout() = - *build_options.result_layout(); + build_options.result_layout()->ToProto(); } else { + Shape result_shape(program_shape->result()); + LayoutUtil::SetToDefaultLayout(&result_shape); *execution_options.mutable_shape_with_output_layout() = - program_shape->result(); - LayoutUtil::SetToDefaultLayout( - execution_options.mutable_shape_with_output_layout()); + result_shape.ToProto(); } - - for (const std::string& disabled_pass : build_options.disabled_hlo_passes()) { - execution_options.mutable_debug_options()->add_xla_disable_hlo_passes( - disabled_pass); - } - + execution_options.set_num_replicas(build_options.num_replicas()); return execution_options; } @@ -145,7 +122,7 @@ StatusOr> LocalService::CompileExecutable( const ExecutableBuildOptions& build_options) { const HloModuleProto& proto = computation.proto(); TF_RET_CHECK(proto.has_host_program_shape()); - const ProgramShape& program_shape = proto.host_program_shape(); + ProgramShape program_shape(proto.host_program_shape()); // Validate incoming layouts. if (argument_layouts.size() != program_shape.parameters_size()) { @@ -220,4 +197,10 @@ StatusOr LocalService::GlobalDataToShapedBuffer( return buffers[replica_number]; } +StatusOr LocalService::RegisterReplicatedBuffers( + std::vector replicated_buffers, const string& tag) { + return allocation_tracker_.RegisterReplicatedBuffers( + std::move(replicated_buffers), tag); +} + } // namespace xla diff --git a/tensorflow/compiler/xla/service/local_service.h b/tensorflow/compiler/xla/service/local_service.h index 3b4f0b50832d6d2b64528ffb63eb5c7375396aec..f56ba32b04b9bf3aba75654bdb98887ad22e6791 100644 --- a/tensorflow/compiler/xla/service/local_service.h +++ b/tensorflow/compiler/xla/service/local_service.h @@ -63,6 +63,11 @@ class LocalService : public Service { StatusOr GlobalDataToShapedBuffer( const GlobalDataHandle& data, int replica_number); + // Registers a vector of shaped buffers of device memory, one per replica, and + // returns a corresponding handle that can be used for talking to XLA clients. + StatusOr RegisterReplicatedBuffers( + std::vector replicated_buffers, const string& tag); + private: explicit LocalService(const ServiceOptions& options, std::unique_ptr backend); diff --git a/tensorflow/compiler/xla/service/logical_buffer_analysis.cc b/tensorflow/compiler/xla/service/logical_buffer_analysis.cc index ec52a24d782a44fda961feab3230886072e755c7..972a5b9ced0d84387ef8308efe2a7aff7317d047 100644 --- a/tensorflow/compiler/xla/service/logical_buffer_analysis.cc +++ b/tensorflow/compiler/xla/service/logical_buffer_analysis.cc @@ -113,6 +113,13 @@ Status LogicalBufferAnalysis::HandleGetTupleElement(HloInstruction*) { return Status::OK(); } +Status LogicalBufferAnalysis::HandleAddDependency( + HloInstruction* add_dependency) { + // AddDependency just forwards the value of its zero-th operand and does not + // create buffers. + return Status::OK(); +} + Status LogicalBufferAnalysis::HandleCopy(HloInstruction* copy) { // The top-level buffer (index={}) for kCopy is newly created, but all other // buffers (in the case of a tuple shape) come from the operand diff --git a/tensorflow/compiler/xla/service/logical_buffer_analysis.h b/tensorflow/compiler/xla/service/logical_buffer_analysis.h index 81f524d84a8091e1fff13dc7c55b401143a02753..7ffca943d0f7805ad4420343fcdbf860415c4c40 100644 --- a/tensorflow/compiler/xla/service/logical_buffer_analysis.h +++ b/tensorflow/compiler/xla/service/logical_buffer_analysis.h @@ -64,6 +64,7 @@ class LogicalBufferAnalysis : public DfsHloVisitorWithDefault { Status HandleRecvDone(HloInstruction* recv_done) override; Status HandleSend(HloInstruction* send) override; Status HandleTupleSelect(HloInstruction* tuple_select) override; + Status HandleAddDependency(HloInstruction* add_dependency) override; // A map from the buffer ID to the logical buffer std::vector> logical_buffers_; diff --git a/tensorflow/compiler/xla/service/map_inliner_test.cc b/tensorflow/compiler/xla/service/map_inliner_test.cc index 84059dd0f71ee8fc0a25703cbab2268d7dc149a8..fd18bfdc3e7f4b5f94237c554c3e6ca8bd065a35 100644 --- a/tensorflow/compiler/xla/service/map_inliner_test.cc +++ b/tensorflow/compiler/xla/service/map_inliner_test.cc @@ -26,7 +26,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_opcode.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/test.h" -#include "tensorflow/compiler/xla/tests/hlo_verified_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/xla_data.pb.h" @@ -35,7 +35,7 @@ namespace op = xla::testing::opcode_matchers; namespace xla { namespace { -using MapInlinerTest = HloVerifiedTestBase; +using MapInlinerTest = HloTestBase; // Test that `map` with `max` is transformed to `max` TEST_F(MapInlinerTest, MapMax) { @@ -59,12 +59,12 @@ TEST_F(MapInlinerTest, MapMax) { HloInstruction::CreateMap(lhs->shape(), {lhs, rhs}, max_f32.get())); auto computation = builder.Build(); - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewVerifiedModule(); hlo_module->AddEmbeddedComputation(std::move(max_f32)); hlo_module->AddEntryComputation(std::move(computation)); MapInliner inliner; - EXPECT_TRUE(inliner.Run(hlo_module).ValueOrDie()); + EXPECT_TRUE(inliner.Run(hlo_module.get()).ValueOrDie()); EXPECT_THAT(hlo_module->entry_computation()->root_instruction(), op::Maximum(lhs, rhs)); @@ -93,12 +93,12 @@ TEST_F(MapInlinerTest, MapConstant) { HloInstruction::CreateMap(lhs->shape(), {lhs}, const2_f32.get())); auto computation = builder.Build(); - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewVerifiedModule(); hlo_module->AddEmbeddedComputation(std::move(const2_f32)); hlo_module->AddEntryComputation(std::move(computation)); HloInstruction* root = hlo_module->entry_computation()->root_instruction(); MapInliner inliner; - EXPECT_TRUE(inliner.Run(hlo_module).ValueOrDie()); + EXPECT_TRUE(inliner.Run(hlo_module.get()).ValueOrDie()); root = hlo_module->entry_computation()->root_instruction(); EXPECT_THAT(root, op::Broadcast(op::Constant())); @@ -131,12 +131,12 @@ TEST_F(MapInlinerTest, MapSubtractOppositeOrder) { HloInstruction::CreateMap(lhs->shape(), {lhs, rhs}, max_f32.get())); auto computation = builder.Build(); - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewVerifiedModule(); hlo_module->AddEmbeddedComputation(std::move(max_f32)); hlo_module->AddEntryComputation(std::move(computation)); MapInliner inliner; - EXPECT_TRUE(inliner.Run(hlo_module).ValueOrDie()); + EXPECT_TRUE(inliner.Run(hlo_module.get()).ValueOrDie()); EXPECT_THAT(hlo_module->entry_computation()->root_instruction(), op::Subtract(rhs, lhs)); diff --git a/tensorflow/compiler/xla/service/multi_output_fusion.cc b/tensorflow/compiler/xla/service/multi_output_fusion.cc index 2ca527bc4cb8f66a085c1e6a7cbb8ddaedbfc07e..53d52d9a3d918fa6dee093668923fcfff963d084 100644 --- a/tensorflow/compiler/xla/service/multi_output_fusion.cc +++ b/tensorflow/compiler/xla/service/multi_output_fusion.cc @@ -18,6 +18,7 @@ limitations under the License. #include "absl/container/flat_hash_set.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" +#include "tensorflow/compiler/xla/service/hlo_reachability.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/core/platform/types.h" @@ -197,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); @@ -257,7 +258,7 @@ bool MultiOutputFusion::LegalToFuse(HloInstruction* instr1, } void MultiOutputFusion::RecomputeReachability() { - reachability_ = computation_->ComputeReachability(); + reachability_ = HloReachabilityMap::Build(computation_); } void MultiOutputFusion::UpdateReachability( @@ -317,9 +318,9 @@ bool MultiOutputFusion::Perform() { << instr2->fused_instructions_computation()->ToString( HloPrintOptions().set_indent_amount(1)); } + Update(instr1, instr2); HloInstruction* ret = Fuse(instr1, instr2); set_is_fused(ret == instr1 ? instr2 : instr1); - Update(instr1, instr2); changed = true; VLOG(2) << "After fusion, \t this: " << ret->name() << "\n" << ret->fused_instructions_computation()->ToString( diff --git a/tensorflow/compiler/xla/service/multi_output_fusion.h b/tensorflow/compiler/xla/service/multi_output_fusion.h index 9508ab2ed1d38ec40983d8892ec8875b848fb21b..1c7583ece720f9e4d4b71a6279b976fed40e10cb 100644 --- a/tensorflow/compiler/xla/service/multi_output_fusion.h +++ b/tensorflow/compiler/xla/service/multi_output_fusion.h @@ -23,6 +23,7 @@ limitations under the License. #include "absl/strings/string_view.h" #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/service/hlo_pass_interface.h" +#include "tensorflow/compiler/xla/service/hlo_reachability.h" #include "tensorflow/compiler/xla/statusor.h" namespace xla { diff --git a/tensorflow/compiler/xla/service/name_uniquer.cc b/tensorflow/compiler/xla/service/name_uniquer.cc index ac2f79674feceff436c0e9c65338967f498e4473..e55b83d17e90bc2ca0053a0421cf80ef6edd5bca 100644 --- a/tensorflow/compiler/xla/service/name_uniquer.cc +++ b/tensorflow/compiler/xla/service/name_uniquer.cc @@ -15,8 +15,10 @@ 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" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/types.h" @@ -27,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; } @@ -42,9 +44,10 @@ NameUniquer::NameUniquer(const string& separator) { if (name.empty()) { return ""; } + 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++) { @@ -52,6 +55,13 @@ NameUniquer::NameUniquer(const string& separator) { result[i] = '_'; } } + + // HLO primitive type names (with the exception of 'tuple') are keywords in + // the HLO text representation and cannot be names, so append an underscore if + // the name is a primitive type. + if (primitive_util::IsPrimitiveTypeName(result) && result != "tuple") { + result += "_"; + } return result; } diff --git a/tensorflow/compiler/xla/service/name_uniquer_test.cc b/tensorflow/compiler/xla/service/name_uniquer_test.cc index 3e2592c6ac626143f1421e545a31d9be91e376bc..d0d04147e0c29c66cba447550c0a9c703f35573a 100644 --- a/tensorflow/compiler/xla/service/name_uniquer_test.cc +++ b/tensorflow/compiler/xla/service/name_uniquer_test.cc @@ -104,5 +104,21 @@ TEST_F(NameUniquerTest, KeepNamesInRandomOrder) { EXPECT_EQ("foo.3", uniquer.GetUniqueName("foo.3")); } +TEST_F(NameUniquerTest, AvoidKeywords) { + NameUniquer uniquer("."); + + EXPECT_EQ("f32_", uniquer.GetUniqueName("f32")); + EXPECT_EQ("s64_", uniquer.GetUniqueName("s64")); + EXPECT_EQ("pred_", uniquer.GetUniqueName("pred")); + + // Though a primitive type, "tuple" is not a keyword. + EXPECT_EQ("tuple", uniquer.GetUniqueName("tuple")); + + // Keywords are not capitalized. + EXPECT_EQ("F32", uniquer.GetUniqueName("F32")); + EXPECT_EQ("S32", uniquer.GetUniqueName("S32")); + EXPECT_EQ("Pred", uniquer.GetUniqueName("Pred")); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/service/pattern_matcher.h b/tensorflow/compiler/xla/service/pattern_matcher.h index 380cde0e6a858c7800445be94bb08dc22f3e776a..9e3d1060210790f60243195a1c1dff13f1fc7fc5 100644 --- a/tensorflow/compiler/xla/service/pattern_matcher.h +++ b/tensorflow/compiler/xla/service/pattern_matcher.h @@ -16,6 +16,7 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_PATTERN_MATCHER_H_ #define TENSORFLOW_COMPILER_XLA_SERVICE_PATTERN_MATCHER_H_ +#include "absl/strings/str_replace.h" #include "absl/strings/string_view.h" #include "absl/utility/utility.h" #include "tensorflow/compiler/xla/layout_util.h" @@ -44,32 +45,48 @@ namespace xla { // // This pattern will match Add instructions whose first operand is a constant. // -// Each pattern type has the following modifiers: +// Each pattern type has the following modifiers, which are described where +// nontrivial. // // Op(): -// - WithName: match operations with the given name -// - WithOpcode: match operations with the given opcode -// - WithShape: match operations whose shape matches the given pattern -// - WithOperand: match operations whose operand matches the given pattern +// - Is: is the given HloInstruction* (i.e. pointer equality) +// - WithName +// - WithOpcode +// - WithoutOpcode: anything other than the given opcode +// - WithShape: instr's shape matches the given pattern +// - WithShapeEqualTo: instr's shape is equal to the given Shape +// - WithShapeCompatibleTo: instr's shape is compatible with the given Shape +// - WithNumOperands +// - WithOperand: operand at the given index matches the given pattern +// - IsConstant +// - IsNonConstant +// - IsConstantScalar/IsEffectiveConstantScalar: Optionally accepts a value, +// e.g. IsConstantScalar() or IsConstantScalar(42). +// - WithFusionKind +// - WithTupleIndex: get-tuple-element operations with the given tuple index +// - WithOneUse: Instruction is used as an operand exactly once. +// - WithOneUser: Instruction is used by exactly one other instruction, but +// is possibly used more than once as an operand (e.g. multiply(x,x)). // // Shape(): -// - EqualTo: matches shapes that are equal to the argument -// - CompatibleTo: matches shapes that are compatible to the argument -// - IsScalar/IsArray/IsTuple: matches scalar/array/tuple shapes -// - IsDenseArray/IsSparseArray: matches arrays with dense/sparse format -// - WithLayout: match shapes whose layout matches the given pattern -// - WithLayoutEqualTo: matches shapes whose layouts equal the argument -// - WithSubshape: matches tuple shapes whose subshape matches the given -// pattern -// - WithSubshapeEqualTo: matches shapes with a subshape equal the argument -// - WithElementType: matches array/scalar shapes with the given element -// type -// - WithRank: matches array/scalar types with the given rank +// - EqualTo +// - CompatibleTo +// - IsScalar/IsEffectiveScalar/IsArray/IsTuple +// - IsDenseArray/IsSparseArray +// - WithLayout: layout shape's layout matches the given pattern (e.g. +// Layout().WithDenseFormat()) +// - WithLayoutEqualTo: shape's layout equals the argument (i.e. another +// Layout, but not the result of Layout().foo()) +// - WithSubshape: shape is a tuple whose subshape matches the given pattern +// (e.g. Shape().IsScalar()). +// - WithSubshapeEqualTo: shape is a tuple with a subshape equal to the arg +// (i.e. another Shape, but not the result of Shape().foo()) +// - WithElementType: shape is an array/scalar with the given elem type +// - WithRank: shape is an array/scalar with the given rank // // Layout(): -// - EqualTo: matches layouts that are equal to the argument -// - WithDenseFormat/WithSparseFormat: matches layouts with dense/sparse -// format +// - EqualTo +// - WithDenseFormat/WithSparseFormat // // Op(), Shape(), and Layout() may be passed an argument of type // HloInstruction**, Shape**, or Layout**, respectively, or const versions of @@ -82,53 +99,55 @@ namespace xla { // CHECK(Match(foo, // match::Op().WithOperand(0, match::Op(&matched_operand)))); // -// Helpers are provided for common nullary, unary, binary, and ternary -// instructions. These helpers can be called with no arguments, in which case -// they will match any instruction matching the opcode. They may also be called -// with matches for the operands and with an optional capture. (The capture must -// be the first argument.) Some examples of these helpers and their equivalents -// are provided below. -// +// Helpers are provided for most HLO instructions. These helpers can be called +// with no arguments, in which case they will match any instruction matching the +// opcode. They may also be called with matches for the operands and with an +// optional capture. (The capture must be the first argument.) Some examples of +// these helpers and their equivalents are provided below. + // Example nullary instruction: -// Param() == Op().WithOpcode(HloOpcode::kParam) -// Param(&a) == Op(&a).WithOpcode(HloOpcode::kParam) +// Parameter() == Op().WithOpcode(HloOpcode::kParameter) +// Parameter(&a) == Op(&a).WithOpcode(HloOpcode::kParameter) // // Example unary instruction: -// Abs() == Op().WithOpcode(HloOpcode::kAbs) -// Abs(Op(&a)) == Op().WithOpcode(HloOpcode::kAbs) -// .WithOperand(0, Op(&a))) -// Abs(&a, Op(&b)) == Op(&a).WithOpcode(HloOpcode::kAbs) -// .WithOperand(0, Op(&b)) +// Abs() == Op().WithOpcode(HloOpcode::kAbs) +// Abs(Op(&a)) == Op().WithOpcode(HloOpcode::kAbs) +// .WithOperand(0, Op(&a))) +// Abs(&a, Op(&b)) == Op(&a).WithOpcode(HloOpcode::kAbs) +// .WithOperand(0, Op(&b)) +// +// Commutative binary instructions have a special form that accepts either order +// of args, e.g.: +// +// AddAnyOrder(Parameter(1), Abs()) == +// Op().WithOpcode(HloOpcode::kAdd) +// .WithBinaryOperandsAnyOrder(Op().WithParameterNum(1), Abs()); // -// Example binary instruction: -// Add() == Op().WithOpcode(HloOpcode::kAdd) -// Add(Op(&a), Op(&b)) == Op().WithOpcode(HloOpcode::kAdd) -// .WithOperand(0, Op(&a)) -// .WithOperand(1, Op(&b)) -// Add(&a, Op(&b), Op(&c)) == Op(&a).WithOpcode(HloOpcode::kAdd) -// .WithOperand(0, Op(&b)) -// .WithOperand(1, Op(&c)) +// MultiplyAnyOrder(&a, Parameter(), Abs()) // Captures the mul in `a`. // -// Example ternary instruction: -// Clamp() == Op().WithOpcode(HloOpcode::kClamp) -// Clamp(Op(&a), Op(&b), Op(&c)) == Op().WithOpcode(HloOpcode::kClamp) -// .WithOperand(0, Op(&a)) -// .WithOperand(1, Op(&b)) -// .WithOperand(2, Op(&c)) -// Clamp(&a, Op(&b), Op(&c), Op(&d)) == Op(&a).WithOpcode(HloOpcode::kClamp) -// .WithOperand(0, Op(&b)) -// .WithOperand(1, Op(&c)) -// .WithOperand(2, Op(&d)) +// The following additional helpers are provided. In all cases, `&a` is +// optional. // +// ConstantScalar(&a) == Op(&a).IsConstantScalar(); +// ConstantScalar(&a, v) == Op(&a).IsConstantScalar(v); +// ConstantEffectiveScalar(&a) == Op(&a).IsConstantEffectiveScalar(); +// ConstantEffectiveScalar(&a, v) == Op(&a).IsConstantEffectiveScalar(&a, v) +// NonConstant(&a) == Op(&a).IsNonConstant() +// GetTupleElement(&a, b, index) == Op(&a).WithTupleIndex(index) +// .WithOperand(0, b); +// Parameter(&a, n) == Op(&a).WithParameterNum(n); struct MatchOption { // If true, actually capture matched item into the user pointer. bool capture; + + // An explanation for why we failed to match is streamed here, if not-null. + std::ostream* explain_os; }; template bool Match(Value* value, const Pattern& pattern, - MatchOption option = {/*.capture=*/true}) { + MatchOption option = {/*.capture=*/true, /*.explain_os=*/nullptr}) { if (option.capture) { auto new_option = option; new_option.capture = false; @@ -143,6 +162,77 @@ namespace match { namespace detail { +// Macro for streaming to option.explain_os if it's not null. +// +// EXPLAIN << "value of foo(): " << foo() +// +#pragma push_macro("EXPLAIN") +#define EXPLAIN \ + if (option.explain_os) *option.explain_os + +// kIndentInc is the additional number of spaces that we indent by when we +// increase the indent "by one". +enum { + kIndentInc = 2, +}; + +// Writes a newline and then `indent` spaces. +// +// We follow an unintuitive convention in this file's pretty-printers: Indents +// are performed by the caller, not the callee. For example, if you want to +// print +// +// foo: +// - bar +// +// you'd do: +// +// Foo::DescribeTo(std::ostream* os, int64 indent) { +// *os << "foo:"; +// Indent(os, indent) // Create a newline at the *current* indent level. +// *os << " - "; +// bar.DescribeTo(os, indent + 3); // + 3 because strlen(" * ") == 3. +// } +// +// Bar::DescribeTo(std::ostream* os, int64 indent) { *os << "bar"; } +// +// Notice that Bar::DescribeTo() does not call Indent; the indenting is +// performed by Foo. This convention allows the caller to decide whether a +// matcher is preceded by a newline, which is important e.g. for the AllOf +// matcher. +// +// (Incidentally, indenting in Match's explanations is handled differently. +// Indents are a common case in DescribeTo [we're printing a whole tree], but +// they're a special case in Match [we're printing only a path through the tree +// that encounters a failing node]. Indents in Match only appear when we +// encounter a failing disjunction, so we just handle them as a special case +// there.) +inline void Indent(std::ostream* os, int64 indent) { + *os << "\n"; + for (int64 i = 0; i < indent; ++i) { + *os << " "; + } +} + +// SFINAE template that determines whether T declares a static member +// kIsTrivialMatcher. +// +// Trivial matchers get special treatment. For example, when printing +// a conjunction of matchers, we don't print "and" after a trivial matcher. This +// yields e.g. +// "a shape compatible with f32[1,2]" +// rather than +// "a shape AND compatible with f32[1,2]" +template +struct IsTrivialMatcher { + static constexpr bool value = false; +}; +template +struct IsTrivialMatcher::type> { + static constexpr bool value = true; +}; + template class AllOfPattern { public: @@ -162,10 +252,19 @@ class AllOfPattern { return matched; } + void DescribeTo(std::ostream* os, int64 indent = 0) const { + DescribeToImpl(os, std::integral_constant(), indent); + } + + // Accessor for patterns_. Please don't use this outside of this file. + const std::tuple& patterns() const { return patterns_; } + private: template bool MatchImpl(ItemType* item, MatchOption option, std::integral_constant) const { + // We don't need to do any EXPLAINing here; it's all correctly handled by + // our sub-matchers (if any fail). return std::get(patterns_).Match(item, option) && MatchImpl(item, option, std::integral_constant()); } @@ -176,6 +275,73 @@ class AllOfPattern { return true; } + // Pretty-printing a conjunction has some special cases to make it easy to + // read in the simple (common) case. + // + // If sizeof...(Patterns) == 1, prints as e.g. + // + // a shape + // + // If sizeof...(Patterns) == 2 and patterns_[0] is a trivial matcher (e.g. "a + // shape") prints as + // + // a shape compatible with f32[1,2] + // + // If sizeof...(Patterns) > 2 and patterns_[0] is a trivial matcher, prints as + // + // a shape: + // * compatible with f32[1,2] AND + // * that represents a scalar + // + // Otherwise prints as: + // + // all of: + // * foo AND + // * bar + // + template + void DescribeToImpl(std::ostream* os, std::integral_constant, + int64 indent) const { + constexpr bool first_is_trivial = + IsTrivialMatcher(patterns_))>::type>::value; + constexpr bool is_last = index == sizeof...(Patterns) - 1; + const auto& submatcher = std::get(patterns_); + + auto print_bulleted_item = [&] { + *os << " * "; + submatcher.DescribeTo(os, indent + 3); + if (!is_last) { + *os << " AND"; + Indent(os, indent); + } + }; + + if (index == 0) { + if (first_is_trivial || is_last) { + submatcher.DescribeTo(os, indent + kIndentInc); + if (sizeof...(Patterns) > 2) { + *os << ":"; + Indent(os, indent); + } + } else { + *os << "all of:"; + Indent(os, indent); + print_bulleted_item(); + } + } else if (first_is_trivial && index == 1 && sizeof...(Patterns) == 2) { + *os << " "; + submatcher.DescribeTo(os, indent); + } else { + print_bulleted_item(); + } + DescribeToImpl(os, std::integral_constant(), indent); + } + + void DescribeToImpl(std::ostream* os, + std::integral_constant, + int64 indent) const {} + std::tuple patterns_; }; @@ -183,10 +349,6 @@ class AllOfPattern { // Returns a pattern that represents the conjunction of all input patterns. All // patterns need to match in order to have the AllOf pattern match. -// -// TODO(timshen): Currently AllOf is still nested, e.g. AllOf, B> is -// not AllOf. We might want to flatten the AllOf type structure if the -// C++ compile error message gets annoying. template detail::AllOfPattern::type, Patterns...> AllOf( const Patterns&... patterns) { @@ -194,6 +356,25 @@ detail::AllOfPattern::type, Patterns...> AllOf( Patterns...>(patterns...); } +// AllOf, X, Y, ...> => AllOf. +// +// This transformation is necessary for good pretty-printing. +template +detail::AllOfPattern::type, InnerPs..., + OuterPs...> +AllOf(const detail::AllOfPattern& inner_p, + const OuterPs&... outer_ps) { + // Invoke constructor of AllOfPattern. + auto make_all_of = [](const InnerPs&... inner_ps, + const OuterPs&... outer_ps) { + return detail::AllOfPattern::type, + InnerPs..., OuterPs...>(inner_ps..., + outer_ps...); + }; + return absl::apply(make_all_of, std::tuple_cat(inner_p.patterns(), + std::make_tuple(outer_ps...))); +} + namespace detail { template @@ -204,8 +385,18 @@ class LayoutPattern; class LayoutPatternBaseImpl { public: bool Match(const ::xla::Layout* layout, MatchOption option) const { - return layout != nullptr; + if (layout == nullptr) { + EXPLAIN << "Layout is null"; + return false; + } + return true; + } + + void DescribeTo(std::ostream* os, int64 indent = 0) const { + *os << "a layout"; } + + static constexpr bool kIsTrivialMatcher = true; }; // A LayoutPattern implementation that matches only if the layout equals a @@ -216,7 +407,17 @@ class LayoutPatternEqualImpl { : layout_(layout) {} bool Match(const ::xla::Layout* layout, MatchOption option) const { - return LayoutUtil::Equal(*layout_, *layout); + if (!LayoutUtil::Equal(*layout_, *layout)) { + EXPLAIN << "Layout " << LayoutUtil::HumanString(*layout) + << " is not equal to expected " + << LayoutUtil::HumanString(*layout_); + return false; + } + return true; + } + + void DescribeTo(std::ostream* os, int64 indent = 0) const { + *os << "equal to " << LayoutUtil::HumanString(*layout_); } private: @@ -230,7 +431,16 @@ class LayoutPatternFormatImpl { explicit constexpr LayoutPatternFormatImpl(Format format) : format_(format) {} bool Match(const ::xla::Layout* layout, MatchOption option) const { - return layout->format() == format_; + if (layout->format() != format_) { + EXPLAIN << "Layout has format " << Format_Name(layout->format()) + << " but expected " << Format_Name(format_); + return false; + } + return true; + } + + void DescribeTo(std::ostream* os, int64 indent = 0) const { + *os << "with format " << Format_Name(format_); } private: @@ -242,11 +452,13 @@ template class LayoutPattern { private: template - LayoutPattern> - AppendImpl(NewImpl new_impl) const { - return LayoutPattern>( - AllOf(impl_, std::move(new_impl)), matched_layout_); + auto AppendImpl(NewImpl new_impl) const + -> LayoutPattern(std::declval(), + std::move(new_impl)))> { + auto new_allof = AllOf(impl_, std::move(new_impl)); + return LayoutPattern(std::move(new_allof), + matched_layout_); } public: @@ -276,6 +488,10 @@ class LayoutPattern { return false; } + void DescribeTo(std::ostream* os, int64 indent = 0) const { + impl_.DescribeTo(os, indent); + } + // Modifies the pattern to match only if the layout equals the given proto. // The layout must outlive the returned pattern. constexpr auto EqualTo(const ::xla::Layout* layout) const @@ -306,19 +522,48 @@ class AnyOfPattern { explicit AnyOfPattern(const Patterns&... patterns) : patterns_(patterns...) {} bool Match(const Item* item, MatchOption option) const { - return MatchImpl(item, option, std::integral_constant()); + return MatchImpl(item, option); } bool Match(Item* item, MatchOption option) const { - return MatchImpl(item, option, std::integral_constant()); + return MatchImpl(item, option); + } + + void DescribeTo(std::ostream* os, int64 indent = 0) const { + *os << "any of:"; + Indent(os, indent); + DescribeToImpl(os, std::integral_constant(), indent); } private: + template + bool MatchImpl(ItemType* item, MatchOption option) const { + // If we're generating an explanation, buffer it until we know we failed. + absl::optional explanation; + MatchOption new_option = option; + if (option.explain_os) { + new_option.explain_os = &explanation.emplace(); + } + bool rv = MatchRecursiveImpl(item, new_option, + std::integral_constant()); + if (!rv && option.explain_os) { + EXPLAIN << "None of the following matchers succeeded:"; + EXPLAIN << explanation->str(); + } + return rv; + } + template - bool MatchImpl(ItemType* item, MatchOption option, - std::integral_constant) const { + bool MatchRecursiveImpl(ItemType* item, MatchOption option, + std::integral_constant) const { auto new_option = option; new_option.capture = false; + + absl::optional explanation; + if (option.explain_os) { + new_option.explain_os = &explanation.emplace(); + } + // Try to match the sub-pattern without capturing behavior. if (std::get(patterns_).Match(item, new_option)) { // Capture the branch. @@ -337,20 +582,46 @@ class AnyOfPattern { // AnyOf will be a runtime number indicate which sub-pattern is matched. // Then we run another pass to do captures only with the help of the // trace. - bool ret = std::get(patterns_).Match(item, option); - DCHECK(ret); + bool matched = std::get(patterns_).Match(item, option); + DCHECK(matched); } return true; } - return MatchImpl(item, option, std::integral_constant()); + if (option.explain_os) { + EXPLAIN << "\nMatcher #" << index + 1; + EXPLAIN << "\n - "; + std::get(patterns_).DescribeTo(option.explain_os, /*indent=*/3); + EXPLAIN << "\nfailed with"; + EXPLAIN << "\n - "; + EXPLAIN << absl::StrReplaceAll(explanation->str(), {{"\n", "\n "}}); + } + return MatchRecursiveImpl(item, option, + std::integral_constant()); } template - bool MatchImpl(ItemType* item, MatchOption option, - std::integral_constant) const { + bool MatchRecursiveImpl( + ItemType* item, MatchOption option, + std::integral_constant) const { return false; } + template + void DescribeToImpl(std::ostream* os, std::integral_constant, + int64 indent) const { + *os << " - "; + std::get(patterns_).DescribeTo(os, indent + 3); + if (index != sizeof...(Patterns) - 1) { + *os << " OR"; + Indent(os, indent); + } + DescribeToImpl(os, std::integral_constant(), indent); + } + + void DescribeToImpl(std::ostream* os, + std::integral_constant, + int64 indent) const {} + std::tuple patterns_; }; @@ -395,8 +666,17 @@ class ShapePattern; class ShapePatternBaseImpl { public: bool Match(const ::xla::Shape* shape, MatchOption option) const { + if (shape == nullptr) { + EXPLAIN << "Shape is null"; + } return shape != nullptr; } + + void DescribeTo(std::ostream* os, int64 indent = 0) const { + *os << "a shape"; + } + + static constexpr bool kIsTrivialMatcher = true; }; // A ShapePattern implementation that matches only if the shape equals a Shape @@ -407,7 +687,16 @@ class ShapePatternEqualImpl { : shape_(shape) {} bool Match(const ::xla::Shape* shape, MatchOption option) const { - return ShapeUtil::Equal(*shape_, *shape); + if (!ShapeUtil::Equal(*shape_, *shape)) { + EXPLAIN << "Shape not equal to " + << ShapeUtil::HumanStringWithLayout(*shape_); + return false; + } + return true; + } + + void DescribeTo(std::ostream* os, int64 indent = 0) const { + *os << "equal to " << ShapeUtil::HumanStringWithLayout(*shape_); } private: @@ -422,7 +711,16 @@ class ShapePatternCompatibleImpl { : shape_(shape) {} bool Match(const ::xla::Shape* shape, MatchOption option) const { - return ShapeUtil::Compatible(*shape_, *shape); + if (!ShapeUtil::Compatible(*shape_, *shape)) { + EXPLAIN << "Shape not compatible with " + << ShapeUtil::HumanString(*shape_); + return false; + } + return true; + } + + void DescribeTo(std::ostream* os, int64 indent = 0) const { + *os << "compatible with " << ShapeUtil::HumanString(*shape_); } private: @@ -437,7 +735,16 @@ class ShapePatternElementTypeImpl { : element_type_(element_type) {} bool Match(const ::xla::Shape* shape, MatchOption option) const { - return shape->element_type() == element_type_; + if (shape->element_type() != element_type_) { + EXPLAIN << "Shape does not have element type " + << PrimitiveType_Name(element_type_); + return false; + } + return true; + } + + void DescribeTo(std::ostream* os, int64 indent = 0) const { + *os << "with element type " << PrimitiveType_Name(element_type_); } private: @@ -450,7 +757,15 @@ class ShapePatternIsScalarImpl { explicit constexpr ShapePatternIsScalarImpl() {} bool Match(const ::xla::Shape* shape, MatchOption option) const { - return ShapeUtil::IsScalar(*shape); + if (!ShapeUtil::IsScalar(*shape)) { + EXPLAIN << "Shape is not a scalar"; + return false; + } + return true; + } + + void DescribeTo(std::ostream* os, int64 indent = 0) const { + *os << "that represents a scalar"; } }; @@ -460,7 +775,15 @@ class ShapePatternIsArrayImpl { explicit constexpr ShapePatternIsArrayImpl() {} bool Match(const ::xla::Shape* shape, MatchOption option) const { - return ShapeUtil::IsArray(*shape); + if (!shape->IsArray()) { + EXPLAIN << "Shape is not an array"; + return false; + } + return true; + } + + void DescribeTo(std::ostream* os, int64 indent = 0) const { + *os << "that represents an array"; } }; @@ -470,7 +793,34 @@ class ShapePatternIsTupleImpl { explicit constexpr ShapePatternIsTupleImpl() {} bool Match(const ::xla::Shape* shape, MatchOption option) const { - return ShapeUtil::IsTuple(*shape); + if (!shape->IsTuple()) { + EXPLAIN << "Shape is not a tuple"; + return false; + } + return true; + } + + void DescribeTo(std::ostream* os, int64 indent = 0) const { + *os << "that represents a tuple"; + } +}; + +// A ShapePattern implementation that matches only if the shape is an effective +// scalar. +class ShapePatternEffectiveScalarImpl { + public: + explicit constexpr ShapePatternEffectiveScalarImpl() {} + + bool Match(const ::xla::Shape* shape, MatchOption option) const { + if (!ShapeUtil::IsEffectiveScalar(*shape)) { + EXPLAIN << "Shape is not an effective scalar"; + return false; + } + return true; + } + + void DescribeTo(std::ostream* os, int64 indent = 0) const { + *os << "that is an effective scalar"; } }; @@ -481,7 +831,23 @@ class ShapePatternRankImpl { explicit constexpr ShapePatternRankImpl(int64 rank) : rank_(rank) {} bool Match(const ::xla::Shape* shape, MatchOption option) const { - return ShapeUtil::Rank(*shape) == rank_; + if (shape->rank() != rank_) { + if (rank_ == 0) { + EXPLAIN << "Shape is not a scalar"; + } else { + EXPLAIN << "Shape does not have rank " << rank_; + } + return false; + } + return true; + } + + void DescribeTo(std::ostream* os, int64 indent = 0) const { + if (rank_ == 0) { + *os << "that is a scalar"; + } else { + *os << "that has " << rank_ << " dimension" << (rank_ != 1 ? "s" : ""); + } } private: @@ -503,8 +869,21 @@ class ShapePatternLayoutImpl { } bool Match(Shape* shape, MatchOption option) const { - return LayoutUtil::HasLayout(*shape) && - layout_.Match(shape->mutable_layout(), option); + if (!LayoutUtil::HasLayout(*shape)) { + EXPLAIN << "Shape does not have a layout"; + return false; + } + if (!layout_.Match(shape->mutable_layout(), option)) { + EXPLAIN << "\nin layout"; + return false; + } + return true; + } + + void DescribeTo(std::ostream* os, int64 indent = 0) const { + *os << "with"; + Indent(os, indent + kIndentInc); + layout_.DescribeTo(os, indent + kIndentInc); } private: @@ -522,17 +901,40 @@ class ShapePatternSubshapeImpl { : index_(index), subshape_(subshape) {} bool Match(const ::xla::Shape* shape, MatchOption option) const { - return ShapeUtil::IndexIsValid(*shape, index_) && - subshape_.Match(&ShapeUtil::GetSubshape(*shape, index_), option); + return MatchImpl(shape, option); } bool Match(::xla::Shape* shape, MatchOption option) const { - return ShapeUtil::IndexIsValid(*shape, index_) && - subshape_.Match(ShapeUtil::GetMutableSubshape(shape, index_), - option); + return MatchImpl(shape, option); + } + + void DescribeTo(std::ostream* os, int64 indent = 0) const { + *os << "with subshape at index " << index_.ToString() << " which is"; + Indent(os, indent + kIndentInc); + subshape_.DescribeTo(os, indent + kIndentInc); } private: + Shape* GetSubshape(Shape* shape) const { + return ShapeUtil::GetMutableSubshape(shape, index_); + } + const Shape* GetSubshape(const Shape* shape) const { + return &ShapeUtil::GetSubshape(*shape, index_); + } + + template + bool MatchImpl(ShapeType* shape, MatchOption option) const { + if (!ShapeUtil::IndexIsValid(*shape, index_)) { + EXPLAIN << "No subshape at " << index_.ToString(); + return false; + } + if (!subshape_.Match(GetSubshape(shape), option)) { + EXPLAIN << "\nin subshape at " << index_.ToString(); + return false; + } + return true; + } + ShapeIndexView index_; ShapePattern subshape_; }; @@ -542,10 +944,12 @@ template class ShapePattern { private: template - ShapePattern> AppendImpl( - NewImpl new_impl) const { - return ShapePattern>( - AllOf(impl_, std::move(new_impl)), matched_shape_); + auto AppendImpl(NewImpl new_impl) const + -> ShapePattern(std::declval(), + std::move(new_impl)))> { + auto new_all_of = AllOf(impl_, std::move(new_impl)); + return ShapePattern(std::move(new_all_of), + matched_shape_); } public: @@ -560,6 +964,11 @@ class ShapePattern { } return true; } + if (shape) { + EXPLAIN << "\nin " + << (shape->has_layout() ? ShapeUtil::HumanStringWithLayout(*shape) + : ShapeUtil::HumanString(*shape)); + } return false; } @@ -571,9 +980,16 @@ class ShapePattern { } return true; } + EXPLAIN << "\nin " + << (shape->has_layout() ? ShapeUtil::HumanStringWithLayout(*shape) + : ShapeUtil::HumanString(*shape)); return false; } + void DescribeTo(std::ostream* os, int64 indent = 0) const { + return impl_.DescribeTo(os, indent); + } + // Modifies the pattern to match only if the shape equals the given proto. // The layout must outlive the returned pattern. constexpr auto EqualTo(const ::xla::Shape* shape) const @@ -612,6 +1028,11 @@ class ShapePattern { return AppendImpl(ShapePatternIsTupleImpl()); } + constexpr auto IsEffectiveScalar() const + -> decltype(this->AppendImpl(ShapePatternEffectiveScalarImpl())) { + return AppendImpl(ShapePatternEffectiveScalarImpl()); + } + // Modifies the pattern to match only if the shape has the given rank. constexpr auto WithRank(int64 rank) const -> decltype(this->AppendImpl(ShapePatternRankImpl(rank))) { @@ -706,6 +1127,22 @@ Shape(::xla::Shape** matched_shape) { namespace detail { +// Overloads to get a const or non-const operand out of an instruction. +inline HloInstruction* HloOperand(HloInstruction* instr, int64 idx) { + return instr->mutable_operand(idx); +} +inline const HloInstruction* HloOperand(const HloInstruction* instr, + int64 idx) { + return instr->operand(idx); +} + +// Pretty-printer for HloInstruction. Sort of like ToShortString, but with +// fewer %s and more shapes. +inline string InstToString(const HloInstruction* inst) { + return inst->ToString( + HloPrintOptions().set_print_metadata(false).set_print_percent(false)); +} + template class HloInstructionPattern; @@ -714,8 +1151,18 @@ class HloInstructionPattern; class HloInstructionPatternBaseImpl { public: bool Match(const ::xla::HloInstruction* inst, MatchOption option) const { - return inst != nullptr; + if (inst == nullptr) { + EXPLAIN << "HloInstruction* is null"; + return false; + } + return true; + } + + void DescribeTo(std::ostream* os, int64 indent = 0) const { + *os << "an HloInstruction"; } + + static constexpr bool kIsTrivialMatcher = true; }; // An HloInstructionPattern implementation that matches only if the instruction @@ -726,13 +1173,44 @@ class HloInstructionPatternNameImpl { : name_(name) {} bool Match(const ::xla::HloInstruction* inst, MatchOption option) const { - return inst->name() == name_; + if (inst->name() != name_) { + EXPLAIN << "HloInstruction not named \"" << name_ << "\""; + return false; + } + return true; + } + + void DescribeTo(std::ostream* os, int64 indent = 0) const { + *os << "named \"" << name_ << "\""; } private: absl::string_view name_; }; +// An HloInstructionPattern implementation that matches only if the instruction +// equals a particular pointer. +class HloInstructionIsImpl { + public: + explicit HloInstructionIsImpl(const HloInstruction* inst) : inst_(inst) {} + + bool Match(const ::xla::HloInstruction* inst, MatchOption option) const { + if (inst != inst_) { + EXPLAIN << "HloInstruction " << inst << " is not " << inst_ << " (" + << InstToString(inst_) << ")"; + return false; + } + return true; + } + + void DescribeTo(std::ostream* os, int64 indent = 0) const { + *os << "which is " << inst_ << " (" << InstToString(inst_) << ")"; + } + + private: + const HloInstruction* inst_; +}; + // An HloInstructionPattern implementation that matches only if the instruction // has a given opcode. class HloInstructionPatternOpcodeImpl { @@ -742,7 +1220,25 @@ class HloInstructionPatternOpcodeImpl { : opcode_(opcode), invert_(invert) {} bool Match(const ::xla::HloInstruction* inst, MatchOption option) const { - return (invert_ ^ (inst->opcode() == opcode_)); + if (invert_ && inst->opcode() == opcode_) { + EXPLAIN << "HloInstruction has opcode " << HloOpcodeString(opcode_) + << ", expected anything else"; + return false; + } + if (!invert_ && inst->opcode() != opcode_) { + EXPLAIN << "HloInstruction doesn't have opcode " + << HloOpcodeString(opcode_); + return false; + } + return true; + } + + void DescribeTo(std::ostream* os, int64 indent = 0) const { + if (!invert_) { + *os << "with opcode " << HloOpcodeString(opcode_); + } else { + *os << "with any opcode other than " << HloOpcodeString(opcode_); + } } private: @@ -750,6 +1246,30 @@ class HloInstructionPatternOpcodeImpl { bool invert_; }; +// An HloInstructionPattern implementation that matches only if the instruction +// has the given number of operands. +class HloInstructionPatternNumOperandsImpl { + public: + explicit constexpr HloInstructionPatternNumOperandsImpl(int64 num_operands) + : num_operands_(num_operands) {} + + bool Match(const ::xla::HloInstruction* inst, MatchOption option) const { + if (inst->operand_count() != num_operands_) { + EXPLAIN << "HloInstruction doesn't have " << num_operands_ << " operands"; + return false; + } + return true; + } + + void DescribeTo(std::ostream* os, int64 indent = 0) const { + *os << "with " << num_operands_ << " operand" + << (num_operands_ != 1 ? "s" : ""); + } + + private: + int64 num_operands_; +}; + // An HloInstructionPattern implementation that matches only if the instruction // has a shape that matches a given pattern. template @@ -760,11 +1280,25 @@ class HloInstructionPatternShapeImpl { : shape_(shape) {} bool Match(const ::xla::HloInstruction* inst, MatchOption option) const { - return shape_.Match(&inst->shape(), option); + if (!shape_.Match(&inst->shape(), option)) { + EXPLAIN << "\nin output shape"; + return false; + } + return true; } bool Match(::xla::HloInstruction* inst, MatchOption option) const { - return shape_.Match(inst->mutable_shape(), option); + if (!shape_.Match(inst->mutable_shape(), option)) { + EXPLAIN << "\nin output shape"; + return false; + } + return true; + } + + void DescribeTo(std::ostream* os, int64 indent = 0) const { + *os << "outputting"; + Indent(os, indent + kIndentInc); + shape_.DescribeTo(os, indent + kIndentInc); } private: @@ -782,20 +1316,197 @@ class HloInstructionPatternOperandImpl { : operand_index_(operand_index), operand_(operand) {} bool Match(const ::xla::HloInstruction* inst, MatchOption option) const { - return operand_index_ < inst->operand_count() && - operand_.Match(inst->operand(operand_index_), option); + return MatchImpl(inst, option); } bool Match(::xla::HloInstruction* inst, MatchOption option) const { - return operand_index_ < inst->operand_count() && - operand_.Match(inst->mutable_operand(operand_index_), option); + return MatchImpl(inst, option); + } + + void DescribeTo(std::ostream* os, int64 indent = 0) const { + *os << "with operand " << operand_index_ << " which is:"; + Indent(os, indent + kIndentInc); + operand_.DescribeTo(os, indent + kIndentInc); } private: + template + bool MatchImpl(HloInstructionType* inst, MatchOption option) const { + if (operand_index_ >= inst->operand_count()) { + EXPLAIN << "desired operand index " << operand_index_ + << " is out of bounds"; + return false; + } + if (!operand_.Match(HloOperand(inst, operand_index_), option)) { + EXPLAIN << "\nin operand " << operand_index_; + return false; + } + return true; + } + int64 operand_index_; HloInstructionPattern operand_; }; +// Matches a binary instruction whose operands come in any order. +template +class HloInstructionPatternBinaryOperandsAnyOrderImpl { + public: + explicit constexpr HloInstructionPatternBinaryOperandsAnyOrderImpl( + const HloInstructionPattern& op1, + const HloInstructionPattern& op2) + : op1_(op1), op2_(op2) {} + + bool Match(HloInstruction* inst, MatchOption option) const { + return MatchImpl(inst, option); + } + + bool Match(const HloInstruction* inst, MatchOption option) const { + return MatchImpl(inst, option); + } + + void DescribeTo(std::ostream* os, int64 indent = 0) const { + *os << "with two operands in either order:"; + Indent(os, indent); + *os << " - "; + op1_.DescribeTo(os, indent + 3); + Indent(os, indent); + *os << " - "; + op2_.DescribeTo(os, indent + 3); + } + + private: + HloInstruction* operand(HloInstruction* inst, int64 idx) const { + return inst->mutable_operand(idx); + } + const HloInstruction* operand(const HloInstruction* inst, int64 idx) const { + return inst->operand(idx); + } + + template + bool MatchImpl(HloInstructionType* inst, MatchOption option) const { + // We could implement this using AnyOf and AllOf matchers, but the templates + // get pretty difficult to debug, since any compile error herein becomes + // not-an-error via SFINAE. Also this way lets us give better messages on + // failure. + if (inst->operand_count() != 2) { + EXPLAIN << "HloInstruction did not have two operands"; + return false; + } + + // If we're not generating explanations, this is pretty simple. + if (!option.explain_os) { + auto try_match = [&](int64 idx1, int64 idx2) { + MatchOption new_option = option; + new_option.capture = false; + if (op1_.Match(operand(inst, idx1), new_option) && + op2_.Match(operand(inst, idx2), new_option)) { + if (option.capture) { + bool matched = op1_.Match(operand(inst, idx1), option) && + op2_.Match(operand(inst, idx2), option); + DCHECK(matched); + } + return true; + } + return false; + }; + return try_match(0, 1) || try_match(1, 0); + } + + // If we are generating explanations, we have some work to do in order to + // generate a helpful error. + // + // First, try all four operand/matcher combinations, recording the + // failure explanations separately from option.explain_os. matches[i][j] + // tells us if matcher_i matches operand j. + bool matches[/*matcher*/ 2][/*operand*/ 2]; + std::stringstream explanations[/*matcher*/ 2][/*operand*/ 2]; + for (int i = 0; i < 2; ++i) { + for (int j = 0; j < 2; ++j) { + MatchOption new_option = option; + new_option.capture = false; + new_option.explain_os = &explanations[i][j]; + matches[i][j] = i == 0 ? op1_.Match(operand(inst, j), new_option) + : op2_.Match(operand(inst, j), new_option); + } + } + + // Check if the match succeeded. + for (int i = 0; i < 2; ++i) { + if (matches[0][i] && matches[1][(i + 1) % 2]) { + // Rerun the matches with capture enabled if necessary. + if (option.capture) { + auto* operand1 = operand(inst, i); + auto* operand2 = operand(inst, (i + 1) % 2); + bool matched = + op1_.Match(operand1, option) && op2_.Match(operand2, option); + DCHECK(matched); + } + return true; + } + } + + auto describe_matcher = [&](int matcher_idx) { + EXPLAIN << "\n - "; + if (matcher_idx == 0) { + op1_.DescribeTo(option.explain_os, /*indent=*/3); + } else { + CHECK_EQ(matcher_idx, 1); + op2_.DescribeTo(option.explain_os, /*indent=*/3); + } + for (int i = 0; i < 2; ++i) { + if (matches[matcher_idx][/*operand*/ i]) { + continue; + } + EXPLAIN << "\ndoes not match " << (i == 0 ? "LHS" : "RHS") << ":\n"; + EXPLAIN << " - "; + EXPLAIN << absl::StrReplaceAll( + explanations[matcher_idx][/*operand*/ i].str(), {{"\n", "\n "}}); + } + }; + + // If we failed to match, one of the following is true: + // 1. op1 (op2) matches neither LHS nor RHS, or + // 2. op1 and op2 both match LHS (RHS), but neither matches RHS (LHS). + // We print different explanations depending on which case we're in. + + // Case 1. + bool wrote_explanation = false; + for (int i = 0; !wrote_explanation && i < 2; ++i) { + if (!matches[i][0] && !matches[i][1]) { + EXPLAIN << "HloInstruction's operands (ignoring order) did not match " + << (i == 0 ? "first" : "second") << " matcher. Specifically,"; + describe_matcher(i); + wrote_explanation = true; + } + } + + // Case 2. + for (int i = 0; !wrote_explanation && i < 2; ++i) { + if (matches[/*matcher*/ 0][/*operand*/ i] && + matches[/*matcher*/ 1][/*operand*/ i]) { + CHECK(!matches[0][(i + 1) % 2]); + CHECK(!matches[1][(i + 1) % 2]); + CHECK(!wrote_explanation); + EXPLAIN << "HloInstruction's " << (i == 1 ? "LHS" : "RHS") + << " operand did not match either of the two matchers. " + "Specifically,"; + describe_matcher(0); + EXPLAIN << "\nand"; + describe_matcher(1); + wrote_explanation = true; + } + } + + CHECK(wrote_explanation); + return false; + } + + HloInstructionPattern op1_; + HloInstructionPattern op2_; +}; + // An HloInstructionPattern implementation that matches only if the instruction // is a fusion node with a particular kind. class HloInstructionPatternFusionKindImpl { @@ -805,14 +1516,32 @@ class HloInstructionPatternFusionKindImpl { : kind_(kind) {} bool Match(const ::xla::HloInstruction* inst, MatchOption option) const { - return inst->opcode() == HloOpcode::kFusion && inst->fusion_kind() == kind_; + return MatchImpl(inst, option); } bool Match(::xla::HloInstruction* inst, MatchOption option) const { - return inst->opcode() == HloOpcode::kFusion && inst->fusion_kind() == kind_; + return MatchImpl(inst, option); + } + + void DescribeTo(std::ostream* os, int64 indent = 0) const { + *os << "with fusion kind " << ToString(kind_); } private: + template + bool MatchImpl(HloInstructionType* inst, MatchOption option) const { + if (inst->opcode() != HloOpcode::kFusion) { + EXPLAIN << "HloInstruction does not have fusion kind " << ToString(kind_) + << "; it's not a fusion"; + return false; + } + if (inst->fusion_kind() != kind_) { + EXPLAIN << "HloInstruction does not have fusion kind " << ToString(kind_); + return false; + } + return true; + } + ::xla::HloInstruction::FusionKind kind_; }; @@ -824,47 +1553,212 @@ class HloInstructionPatternTupleIndexImpl { : tuple_index_(tuple_index) {} bool Match(const ::xla::HloInstruction* inst, MatchOption option) const { - return inst->opcode() == HloOpcode::kGetTupleElement && - inst->tuple_index() == tuple_index_; + return MatchImpl(inst, option); } bool Match(::xla::HloInstruction* inst, MatchOption option) const { - return inst->opcode() == HloOpcode::kGetTupleElement && - inst->tuple_index() == tuple_index_; + return MatchImpl(inst, option); + } + + void DescribeTo(std::ostream* os, int64 indent = 0) const { + *os << "which is a GTE with index " << tuple_index_; } private: + template + bool MatchImpl(HloInstructionType* inst, MatchOption option) const { + if (inst->opcode() != HloOpcode::kGetTupleElement) { + EXPLAIN << "HloInstruction is not a GTE with index " << tuple_index_ + << "; it's not a GTE at all"; + return false; + } + if (inst->tuple_index() != tuple_index_) { + EXPLAIN << "HloInstruction is not a GTE with index " << tuple_index_; + return false; + } + return true; + } + int64 tuple_index_; }; -template -class HloPredicatePatternImpl { +class HloInstructionPatternParameterNumImpl { public: - explicit HloPredicatePatternImpl(Predicate pred) : pred_(std::move(pred)) {} + explicit constexpr HloInstructionPatternParameterNumImpl(int64 parameter_num) + : parameter_num_(parameter_num) {} - bool Match(const ItemType* item, MatchOption option) const { - return pred_(item); + bool Match(const ::xla::HloInstruction* inst, MatchOption option) const { + return MatchImpl(inst, option); } - bool Match(ItemType* item, MatchOption option) const { return pred_(item); } + bool Match(::xla::HloInstruction* inst, MatchOption option) const { + return MatchImpl(inst, option); + } + + void DescribeTo(std::ostream* os, int64 indent = 0) const { + *os << "which is parameter " << parameter_num_; + } private: - Predicate pred_; + template + bool MatchImpl(HloInstructionType* inst, MatchOption option) const { + if (inst->opcode() != HloOpcode::kParameter || + inst->parameter_number() != parameter_num_) { + EXPLAIN << "HloInstruction is not parameter " << parameter_num_; + return false; + } + return true; + } + + int64 parameter_num_; }; -struct PatternFriend; +// Superclass that contains common code used by Op::WithOneUse() and +// Op::WithOneUser(). +class HloInstructionPatternOneUseOrUserImpl { + protected: + bool MatchOneUser(const HloInstruction* inst, MatchOption option) const { + if (inst->user_count() != 1) { + EXPLAIN << "HloInstruction has " << inst->user_count() + << " users, but expected exactly one."; + if (inst->user_count() > 1) { + EXPLAIN << "\nAll users:"; + for (const HloInstruction* user : inst->users()) { + EXPLAIN << "\n - " << InstToString(user); + } + } + return false; + } + return true; + } +}; + +class HloInstructionPatternOneUseImpl + : public HloInstructionPatternOneUseOrUserImpl { + public: + bool Match(const HloInstruction* inst, MatchOption option) const { + if (!MatchOneUser(inst, option)) { + return false; + } + + int64 use_count = absl::c_count_if( + inst->users()[0]->operands(), + [&](const HloInstruction* operand) { return operand == inst; }); + if (use_count != 1) { + EXPLAIN << "HloInstruction is used " << use_count + << " times by its user, but is expected to be used just once: " + << InstToString(inst->users()[0]); + return false; + } + return true; + } + + void DescribeTo(std::ostream* os, int64 indent = 0) const { + *os << "which has exactly one use"; + } +}; + +class HloInstructionPatternOneUserImpl + : public HloInstructionPatternOneUseOrUserImpl { + public: + bool Match(const HloInstruction* inst, MatchOption option) const { + return MatchOneUser(inst, option); + } + + void DescribeTo(std::ostream* os, int64 indent = 0) const { + *os << "which has exactly one user (but possibly is used multiple times by " + "that instruction)"; + } +}; + +// Matches a constant scalar or effective scalar, optionally with a given value. +template +class HloConstantScalarImpl { + public: + explicit constexpr HloConstantScalarImpl(bool match_effective_scalar) + : val_(absl::nullopt), match_effective_scalar_(match_effective_scalar) {} + + constexpr HloConstantScalarImpl(ScalarTy val, bool match_effective_scalar) + : val_(val), match_effective_scalar_(match_effective_scalar) {} + + bool Match(const ::xla::HloInstruction* inst, MatchOption option) const { + return MatchImpl(inst, option); + } + + bool Match(::xla::HloInstruction* inst, MatchOption option) const { + return MatchImpl(inst, option); + } + + void DescribeTo(std::ostream* os, int64 indent = 0) const { + *os << "which is a constant " + << (match_effective_scalar_ ? "effective " : "") << "scalar"; + if (val_.has_value()) { + *os << " with value " << *val_; + } + } + + private: + template + bool MatchImpl(InstTy* inst, MatchOption option) const { + const auto* const_inst = DynCast(inst); + if (!const_inst) { + EXPLAIN << "HloInstruction is not a constant"; + return false; + } + if (match_effective_scalar_ && + !ShapeUtil::IsEffectiveScalar(inst->shape())) { + EXPLAIN << "HloInstruction is not an effective scalar"; + return false; + } + if (!match_effective_scalar_ && !ShapeUtil::IsScalar(inst->shape())) { + EXPLAIN << "HloInstruction is not a scalar"; + return false; + } + if (!val_.has_value()) { + return true; + } + + // Check that literal == static_cast(val) and + // val == static_cast(literal). This is sufficient to ensure that + // the two constant scalars are actually "equal". + auto val_literal = LiteralUtil::CreateR0(*val_); + auto literal_r0_or = const_inst->literal().Reshape({}); + auto val_as_literal_ty_or = + val_literal.Convert(const_inst->shape().element_type()); + if (!literal_r0_or.ok() || !val_as_literal_ty_or.ok()) { + EXPLAIN << "could not construct relevant Literals (how did this happen?)"; + return false; + } + auto literal_r0 = std::move(literal_r0_or).ValueOrDie(); + auto val_as_literal_ty = std::move(val_as_literal_ty_or).ValueOrDie(); + auto literal_r0_as_val_ty_or = + literal_r0.Convert(val_literal.shape().element_type()); + bool rv = literal_r0_as_val_ty_or.ok() && // + literal_r0_as_val_ty_or.ValueOrDie() == val_literal && + literal_r0 == val_as_literal_ty; + if (!rv) { + EXPLAIN << "HloInstruction's constant value " + << literal_r0.ToStringWithoutShape() + << " did not match expected value " << *val_; + } + return rv; + } + + absl::optional val_; + bool match_effective_scalar_; +}; // A pattern that matches HloInstructions. template class HloInstructionPattern { private: template - HloInstructionPattern> - AppendImpl(NewImpl new_impl) const { - return HloInstructionPattern< - HloInstructionType, AllOfPattern<::xla::HloInstruction, Impl, NewImpl>>( - AllOf(impl_, std::move(new_impl)), matched_inst_); + auto AppendImpl(NewImpl new_impl) const -> HloInstructionPattern< + HloInstructionType, decltype(AllOf( + std::declval(), std::move(new_impl)))> { + auto new_allof = AllOf(impl_, std::move(new_impl)); + return HloInstructionPattern( + std::move(new_allof), matched_inst_); } public: @@ -880,6 +1774,9 @@ class HloInstructionPattern { } return true; } + if (inst != nullptr) { + EXPLAIN << "\nin " << InstToString(inst); + } return false; } @@ -891,6 +1788,7 @@ class HloInstructionPattern { } return true; } + EXPLAIN << "\nin " << InstToString(inst); return false; } @@ -907,6 +1805,11 @@ class HloInstructionPattern { return AppendImpl(HloInstructionPatternOpcodeImpl(opcode, false)); } + auto WithNumOperands(int64 num_operands) const -> decltype( + this->AppendImpl(HloInstructionPatternNumOperandsImpl(num_operands))) { + return AppendImpl(HloInstructionPatternNumOperandsImpl(num_operands)); + } + // Modifies the pattern to match only if the instruction does not have the // given opcode. auto WithoutOpcode(HloOpcode opcode) const @@ -915,12 +1818,47 @@ class HloInstructionPattern { return AppendImpl(HloInstructionPatternOpcodeImpl(opcode, true)); } + constexpr auto Is(const HloInstruction* instr) const + -> decltype(this->AppendImpl(HloInstructionIsImpl(instr))) { + return AppendImpl(HloInstructionIsImpl(instr)); + } + // Modifies the pattern to match only if the instruction is a constant. constexpr auto IsConstant() const -> decltype(this->WithOpcode(HloOpcode::kConstant)) { return WithOpcode(HloOpcode::kConstant); } + constexpr auto IsConstantScalar() const -> decltype(this->AppendImpl( + HloConstantScalarImpl(/*match_effective_scalar=*/false))) { + return AppendImpl( + HloConstantScalarImpl(/*match_effective_scalar=*/false)); + } + + // This does not check that T has the same type as the instruction, so e.g. + // IsConstantScalar(1.0) may match a constant of shape int32[]. + template + constexpr auto IsConstantScalar(const ScalarTy& val) const + -> decltype(this->AppendImpl(HloConstantScalarImpl( + val, /*match_effective_scalar=*/false))) { + return AppendImpl( + HloConstantScalarImpl(val, /*match_effective_scalar=*/false)); + } + + constexpr auto IsConstantEffectiveScalar() const -> decltype(this->AppendImpl( + HloConstantScalarImpl(/*match_effective_scalar=*/true))) { + return AppendImpl( + HloConstantScalarImpl(/*match_effective_scalar=*/true)); + } + + template + constexpr auto IsConstantEffectiveScalar(const ScalarTy& val) const + -> decltype(this->AppendImpl(HloConstantScalarImpl( + val, /*match_effective_scalar=*/true))) { + return AppendImpl( + HloConstantScalarImpl(val, /*match_effective_scalar=*/true)); + } + // Modifies the pattern to match only if the instruction is not a constant. constexpr auto IsNonConstant() const -> decltype(this->WithoutOpcode(HloOpcode::kConstant)) { @@ -937,6 +1875,22 @@ class HloInstructionPattern { HloInstructionPatternShapeImpl(shape)); } + // Make this a templated function to work around gcc 4.9.4 template infinite + // recursion bug. + template + constexpr auto WithShapeEqualTo(const ::xla::Shape* shape) const + -> decltype(this->WithShape(Shape().EqualTo(shape))) { + return WithShape(Shape().EqualTo(shape)); + } + + // Make this a templated function to work around gcc 4.9.4 template infinite + // recursion bug. + template + constexpr auto WithShapeCompatibleTo(const ::xla::Shape* shape) const + -> decltype(this->WithShape(Shape().CompatibleTo(shape))) { + return WithShape(Shape().CompatibleTo(shape)); + } + // Modifies the pattern to match only if the instruction has an operand that // matches the given pattern. template @@ -951,6 +1905,20 @@ class HloInstructionPattern { operand_index, operand)); } + template + constexpr auto WithBinaryOperandsAnyOrder( + const HloInstructionPattern& op1, + const HloInstructionPattern& op2) const + -> decltype(this->AppendImpl( + HloInstructionPatternBinaryOperandsAnyOrderImpl< + OperandType1, OperandImpl1, OperandType2, OperandImpl2>(op1, + op2))) { + return AppendImpl( + HloInstructionPatternBinaryOperandsAnyOrderImpl< + OperandType1, OperandImpl1, OperandType2, OperandImpl2>(op1, op2)); + } + // Modifies the pattern to match only if the instruction is a fusion node with // the given kind. constexpr auto WithFusionKind(HloInstruction::FusionKind kind) const @@ -965,17 +1933,34 @@ class HloInstructionPattern { return AppendImpl(HloInstructionPatternTupleIndexImpl(tuple_index)); } - private: - template - constexpr auto WithPredicate(Predicate pred) const -> decltype( - this->AppendImpl(HloPredicatePatternImpl( - std::move(pred)))) { - return AppendImpl( - HloPredicatePatternImpl(std::move(pred))); + // Modifies the pattern to match only if the instruction is a parameter + // with the given parameter number. + constexpr auto WithParameterNum(int64 parameter_num) const -> decltype( + this->AppendImpl(HloInstructionPatternParameterNumImpl(parameter_num))) { + return AppendImpl(HloInstructionPatternParameterNumImpl(parameter_num)); } - friend struct PatternFriend; + // Modifies the pattern to match if the instruction is used exactly once. + // Does not match if the instruction is used twice by the same user (e.g. + // multiply(x,x)). + constexpr auto WithOneUse() const + -> decltype(this->AppendImpl(HloInstructionPatternOneUseImpl())) { + return AppendImpl(HloInstructionPatternOneUseImpl()); + } + // Modifies the pattern to match if the instruction is used by exactly one + // other instruction. Will match if the instruction is used twice, so long as + // it's by the same user (e.g. multiply(x,x)). + constexpr auto WithOneUser() const + -> decltype(this->AppendImpl(HloInstructionPatternOneUserImpl())) { + return AppendImpl(HloInstructionPatternOneUserImpl()); + } + + void DescribeTo(std::ostream* os, int64 indent = 0) const { + impl_.DescribeTo(os, indent); + } + + private: Impl impl_; HloInstructionType** matched_inst_; }; @@ -1016,6 +2001,7 @@ Op(::xla::HloInstruction** matched_inst) { XLA_NULLOP_PATTERN(Constant) XLA_NULLOP_PATTERN(Parameter) XLA_NULLOP_PATTERN(Iota) +XLA_NULLOP_PATTERN(Rng) #undef XLA_NULLOP_PATTERN // Helpers for unary instructions. @@ -1047,8 +2033,10 @@ XLA_UNOP_PATTERN(RoundNearestAfz) XLA_UNOP_PATTERN(Bitcast) XLA_UNOP_PATTERN(Broadcast) XLA_UNOP_PATTERN(Ceil) +XLA_UNOP_PATTERN(Convert) XLA_UNOP_PATTERN(Copy) XLA_UNOP_PATTERN(Cos) +XLA_UNOP_PATTERN(AllReduce) XLA_UNOP_PATTERN(Exp) XLA_UNOP_PATTERN(Fft) XLA_UNOP_PATTERN(Floor) @@ -1062,14 +2050,13 @@ XLA_UNOP_PATTERN(Negate) XLA_UNOP_PATTERN(Real) XLA_UNOP_PATTERN(Recv) XLA_UNOP_PATTERN(RecvDone) -XLA_UNOP_PATTERN(Reduce) XLA_UNOP_PATTERN(ReducePrecision) XLA_UNOP_PATTERN(Reshape) XLA_UNOP_PATTERN(Reverse) XLA_UNOP_PATTERN(SendDone) XLA_UNOP_PATTERN(Sign) XLA_UNOP_PATTERN(Sin) -XLA_UNOP_PATTERN(Sort) +XLA_UNOP_PATTERN(Slice) XLA_UNOP_PATTERN(Tanh) XLA_UNOP_PATTERN(Transpose) #undef XLA_UNOP_PATTERN @@ -1106,24 +2093,30 @@ XLA_UNOP_PATTERN(Transpose) #define XLA_COMMUTATIVE_BINOP_PATTERN(NAME) \ XLA_BINOP_PATTERN(NAME) \ \ - template \ - inline auto NAME##AnyOrder(Lhs&& lhs, Rhs&& rhs) \ - ->decltype(AnyOf(NAME(lhs, rhs), NAME(rhs, lhs))) { \ - return AnyOf(NAME(lhs, rhs), NAME(rhs, lhs)); \ - } \ - \ template \ inline auto NAME##AnyOrder(HloInstructionType** matched_inst, Lhs&& lhs, \ Rhs&& rhs) \ - ->decltype(AnyOf(NAME(matched_inst, lhs, rhs), \ - NAME(matched_inst, rhs, lhs))) { \ - return AnyOf(NAME(matched_inst, lhs, rhs), \ - NAME(matched_inst, rhs, lhs)); \ + ->decltype(Op(matched_inst) \ + .WithOpcode(HloOpcode::k##NAME) \ + .WithBinaryOperandsAnyOrder(std::forward(lhs), \ + std::forward(rhs))) { \ + return Op(matched_inst) \ + .WithOpcode(HloOpcode::k##NAME) \ + .WithBinaryOperandsAnyOrder(std::forward(lhs), \ + std::forward(rhs)); \ + } \ + template \ + inline auto NAME##AnyOrder(Lhs&& lhs, Rhs&& rhs) \ + ->decltype(NAME##AnyOrder( \ + nullptr, std::forward(lhs), std::forward(rhs))) { \ + return NAME##AnyOrder( \ + nullptr, std::forward(lhs), std::forward(rhs)); \ } XLA_COMMUTATIVE_BINOP_PATTERN(Add) XLA_BINOP_PATTERN(Atan2) XLA_BINOP_PATTERN(Divide) XLA_BINOP_PATTERN(Complex) +XLA_BINOP_PATTERN(Convolution) XLA_BINOP_PATTERN(Dot) XLA_COMMUTATIVE_BINOP_PATTERN(Eq) XLA_BINOP_PATTERN(Gather) @@ -1136,7 +2129,9 @@ XLA_COMMUTATIVE_BINOP_PATTERN(Minimum) XLA_COMMUTATIVE_BINOP_PATTERN(Multiply) XLA_COMMUTATIVE_BINOP_PATTERN(Ne) XLA_BINOP_PATTERN(Outfeed) +XLA_BINOP_PATTERN(Pad) XLA_BINOP_PATTERN(Power) +XLA_BINOP_PATTERN(ReduceWindow) XLA_BINOP_PATTERN(Remainder) XLA_BINOP_PATTERN(Send) XLA_BINOP_PATTERN(Subtract) @@ -1183,33 +2178,68 @@ XLA_BINOP_PATTERN(ShiftRightLogical) .WithOperand(2, std::forward(arg2)); \ } XLA_TERNOP_PATTERN(Clamp); +XLA_TERNOP_PATTERN(Scatter); XLA_TERNOP_PATTERN(Select); #undef XLA_TERNOP_PATTERN namespace detail { -struct PatternFriend { - template - static auto ConstantScalar(T constant) -> decltype( - Constant() - .WithShape(match::Shape().IsScalar()) - .WithPredicate( - std::declval>())) { - std::function pred = - [constant](const HloInstruction* instr) { - const auto& literal = Cast(instr)->literal(); - auto status_or_const = LiteralUtil::CreateR0(constant).Convert( - literal.shape().element_type()); - return status_or_const.ok() && - literal == status_or_const.ConsumeValueOrDie(); - }; - - return Constant() - .WithShape(match::Shape().IsScalar()) - .WithPredicate(std::move(pred)); - } -}; +template +inline auto WithOperands(Matcher&& m, int64 operand_num, FirstArg&& first_arg) + -> decltype(m.WithOperand(operand_num, std::forward(first_arg))) { + return m.WithOperand(operand_num, std::forward(first_arg)); +} + +template +inline auto WithOperands(Matcher&& m, int64 operand_num, FirstArg&& first_arg, + Args&&... args) + -> decltype(WithOperands(m.WithOperand(operand_num, + std::forward(first_arg)), + operand_num + 1, std::forward(args)...)) { + return WithOperands( + m.WithOperand(operand_num, std::forward(first_arg)), + operand_num + 1, std::forward(args)...); +} } // namespace detail +#define XLA_VARIADIC_OP_PATTERN(NAME) \ + inline auto NAME()->decltype(Op().WithOpcode(HloOpcode::k##NAME)) { \ + return Op().WithOpcode(HloOpcode::k##NAME); \ + } \ + \ + template \ + inline auto NAME(Args&&... args) \ + ->decltype(detail::WithOperands(Op().WithOpcode(HloOpcode::k##NAME) \ + .WithNumOperands(sizeof...(Args)), \ + 0, std::forward(args)...)) { \ + return detail::WithOperands( \ + Op().WithOpcode(HloOpcode::k##NAME).WithNumOperands(sizeof...(Args)), \ + /*operand_num=*/0, std::forward(args)...); \ + } \ + \ + template \ + inline auto NAME(HloInstructionType** matched_inst, Args&&... args) \ + ->decltype(detail::WithOperands(Op(matched_inst) \ + .WithOpcode(HloOpcode::k##NAME) \ + .WithNumOperands(sizeof...(Args)), \ + 0, std::forward(args)...)) { \ + return detail::WithOperands(Op(matched_inst) \ + .WithOpcode(HloOpcode::k##NAME) \ + .WithNumOperands(sizeof...(Args)), \ + /*operand_num=*/0, \ + std::forward(args)...); \ + } + +// We could implement all ops as "variadic" ops, but it would make the +// already-bad compile errors even worse. +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); +XLA_VARIADIC_OP_PATTERN(Tuple); + // Helpers for matching non-constant instructions. inline auto NonConstant() -> decltype(Op().IsNonConstant()) { return Op().IsNonConstant(); @@ -1247,14 +2277,71 @@ inline auto GetTupleElement(HloInstructionType** matched_inst, Arg&& arg, .WithTupleIndex(tuple_index); } -template -inline auto ConstantScalar(T constant) - -> decltype(detail::PatternFriend::ConstantScalar(constant)) { - return detail::PatternFriend::ConstantScalar(constant); +// Add overloads for Parameter which take an int64 specifying the parameter +// number. +inline auto Parameter(int64 parameter_num) -> decltype( + Op().WithOpcode(HloOpcode::kParameter).WithParameterNum(parameter_num)) { + return Op().WithOpcode(HloOpcode::kParameter).WithParameterNum(parameter_num); +} +template +inline auto Parameter(HloInstructionType** matched_inst, int64 parameter_num) + -> decltype(Op(matched_inst) + .WithOpcode(HloOpcode::kParameter) + .WithParameterNum(parameter_num)) { + return Op(matched_inst) + .WithOpcode(HloOpcode::kParameter) + .WithParameterNum(parameter_num); +} + +inline auto ConstantScalar() -> decltype(Op().IsConstantScalar()) { + return Op().IsConstantScalar(); +} + +template +inline auto ConstantScalar(HloInstructionType** matched_inst) + -> decltype(Op(matched_inst).IsConstantScalar()) { + return Op(matched_inst).IsConstantScalar(); +} + +template +inline auto ConstantScalar(ScalarTy val) + -> decltype(Op().IsConstantScalar(val)) { + return Op().IsConstantScalar(val); +} + +template +inline auto ConstantScalar(HloInstructionType** matched_inst, ScalarTy val) + -> decltype(Op(matched_inst).IsConstantScalar(val)) { + return Op(matched_inst).IsConstantScalar(val); +} + +inline auto ConstantEffectiveScalar() -> decltype(Op().IsConstantScalar()) { + return Op().IsConstantEffectiveScalar(); +} + +template +inline auto ConstantEffectiveScalar(HloInstructionType** matched_inst) + -> decltype(Op(matched_inst).IsConstantScalar()) { + return Op(matched_inst).IsConstantEffectiveScalar(); +} + +template +inline auto ConstantEffectiveScalar(ScalarTy val) + -> decltype(Op().IsConstantEffectiveScalar(val)) { + return Op().IsConstantEffectiveScalar(val); +} + +template +inline auto ConstantEffectiveScalar(HloInstructionType** matched_inst, + ScalarTy val) + -> decltype(Op(matched_inst).IsConstantEffectiveScalar(val)) { + return Op(matched_inst).IsConstantEffectiveScalar(val); } } // namespace match } // namespace xla +#undef EXPLAIN +#pragma pop_macro("EXPLAIN") #endif // TENSORFLOW_COMPILER_XLA_SERVICE_PATTERN_MATCHER_H_ diff --git a/tensorflow/compiler/xla/service/pattern_matcher_gmock.h b/tensorflow/compiler/xla/service/pattern_matcher_gmock.h new file mode 100644 index 0000000000000000000000000000000000000000..8fe2d10a11b5b2d26ee222c63e0db2d55e361d12 --- /dev/null +++ b/tensorflow/compiler/xla/service/pattern_matcher_gmock.h @@ -0,0 +1,92 @@ +/* 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_PATTERN_MATCHER_GMOCK_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_PATTERN_MATCHER_GMOCK_H_ + +#include +#include "tensorflow/compiler/xla/service/pattern_matcher.h" +#include "tensorflow/compiler/xla/test.h" +#include "tensorflow/core/platform/test.h" + +namespace xla { + +namespace pattern_matcher_gmock_detail { +template +class GmockMatcher { + public: + explicit GmockMatcher(Pattern p) : pattern_(std::move(p)) {} + + // In service of better error messages, list out the overloads explicitly + // rather than just using a template. gMock's polymorphism plus + // pattern_matcher yields some pretty gnarly stuff. + bool MatchAndExplain(const Layout& l, + ::testing::MatchResultListener* listener) const { + return MatchAndExplainImpl(&l, listener); + } + bool MatchAndExplain(const Layout* l, + ::testing::MatchResultListener* listener) const { + return MatchAndExplainImpl(l, listener); + } + + bool MatchAndExplain(const Shape& s, + ::testing::MatchResultListener* listener) const { + return MatchAndExplainImpl(&s, listener); + } + bool MatchAndExplain(const Shape* s, + ::testing::MatchResultListener* listener) const { + return MatchAndExplainImpl(s, listener); + } + + bool MatchAndExplain(const HloInstruction& instr, + ::testing::MatchResultListener* listener) const { + return MatchAndExplainImpl(&instr, listener); + } + bool MatchAndExplain(const HloInstruction* instr, + ::testing::MatchResultListener* listener) const { + return MatchAndExplainImpl(instr, listener); + } + + void DescribeTo(std::ostream* os) const { pattern_.DescribeTo(os); } + + void DescribeNegationTo(std::ostream* os) const { + *os << "is NOT: "; + DescribeTo(os); + } + + private: + template + bool MatchAndExplainImpl(const T* t, + ::testing::MatchResultListener* listener) const { + MatchOption options{/*.capture=*/true, /*.explain_os=*/listener->stream()}; + return Match(t, pattern_, options); + } + + Pattern pattern_; +}; +} // namespace pattern_matcher_gmock_detail + +template +::testing::PolymorphicMatcher< + pattern_matcher_gmock_detail::GmockMatcher> +GmockMatch(Pattern&& p) { + return ::testing::MakePolymorphicMatcher( + pattern_matcher_gmock_detail::GmockMatcher( + std::forward(p))); +} + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_PATTERN_MATCHER_GMOCK_H_ diff --git a/tensorflow/compiler/xla/service/pattern_matcher_gmock_test.cc b/tensorflow/compiler/xla/service/pattern_matcher_gmock_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..f51a18b13894d75300c46835fabd82a4ce0699af --- /dev/null +++ b/tensorflow/compiler/xla/service/pattern_matcher_gmock_test.cc @@ -0,0 +1,75 @@ +/* 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/pattern_matcher_gmock.h" +#include "tensorflow/compiler/xla/service/pattern_matcher.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/test.h" +#include "tensorflow/core/platform/test.h" + +namespace xla { +namespace { + +namespace m = ::xla::match; +using ::testing::Not; + +template +string Describe(const ::testing::Matcher& m) { + std::stringstream ss; + m.DescribeTo(&ss); + return ss.str(); +} + +template +string Explain( + const MatchedTy& val, + const ::testing::Matcher::type>& m) { + ::testing::StringMatchResultListener listener; + EXPECT_THAT(val, ::testing::Not(m)); // For the error message. + EXPECT_FALSE(m.MatchAndExplain(val, &listener)); + return listener.str(); +} + +// This file tests the GmockMatch function. The actual explanation and +// description returned by matchers is tested in pattern_matchers_test. +TEST(PatternMatcherGmock, MatchShape) { + Shape s = ShapeUtil::MakeShape(F32, {10, 100}); + // You can pass const Shape& or a const Shape*. + EXPECT_THAT(s, GmockMatch(m::Shape())); + EXPECT_THAT(&s, Not(GmockMatch(m::Shape().WithElementType(F16)))); + EXPECT_THAT(Describe(GmockMatch(m::Shape().IsArray())), + "a shape that represents an array"); +} + +TEST(PatternMatcherGmock, MatchLayout) { + Layout l = LayoutUtil::MakeLayout({0, 1}); + EXPECT_THAT(l, GmockMatch(m::Layout())); + EXPECT_THAT(&l, Not(GmockMatch(m::Layout().WithSparseFormat()))); + EXPECT_THAT(Describe(GmockMatch(m::Layout().WithSparseFormat())), + "a layout with format SPARSE"); +} + +TEST(PatternMatchGmock, MatchInstruction) { + auto instr = + HloInstruction::CreateParameter(0, ShapeUtil::MakeShape(F32, {42}), "p"); + EXPECT_THAT(instr.get(), GmockMatch(m::Parameter())); + EXPECT_THAT(*instr, GmockMatch(m::Parameter(0))); + EXPECT_THAT(*instr, Not(GmockMatch(m::Parameter(1)))); + EXPECT_THAT(Describe(GmockMatch(m::Parameter())), + "an HloInstruction with opcode parameter"); +} + +} // anonymous namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/service/pattern_matcher_test.cc b/tensorflow/compiler/xla/service/pattern_matcher_test.cc index 3ab7b7fd7168d7ddd1470fdb03a04ba7b171fddb..5c3c009a68bffbda8642fceedfb724879fbf1530 100644 --- a/tensorflow/compiler/xla/service/pattern_matcher_test.cc +++ b/tensorflow/compiler/xla/service/pattern_matcher_test.cc @@ -14,14 +14,18 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/compiler/xla/service/pattern_matcher.h" +#include "absl/strings/str_cat.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" #include "tensorflow/compiler/xla/service/hlo_parser.h" +#include "tensorflow/compiler/xla/test.h" #include "tensorflow/core/platform/test.h" namespace xla { namespace { +namespace m = match; + TEST(PatternMatcherTest, AddOp) { constexpr char kModuleStr[] = R"(HloModule two_plus_two_module ENTRY %two_plus_two_computation () -> f32[] { @@ -229,23 +233,74 @@ TEST(PatternMatcherTest, AnyOf) { } TEST(PatternMatcherTest, ConstantScalar) { - constexpr char kModuleStr[] = R"( - HloModule test_module ENTRY test { ROOT constant = f16[] constant(42) })"; - TF_ASSERT_OK_AND_ASSIGN(auto hlo_module, ParseHloString(kModuleStr)); - auto* root = hlo_module->entry_computation()->root_instruction(); - - EXPECT_TRUE(Match(root, match::ConstantScalar(42))); - EXPECT_FALSE(Match(root, match::ConstantScalar(41))); - EXPECT_FALSE(Match(root, match::ConstantScalar(0))); -} + using match::ConstantEffectiveScalar; + using match::ConstantScalar; + using match::Op; + using match::Tuple; -TEST(PatternMatcherTest, NoMatchConstantScalar) { constexpr char kModuleStr[] = R"( - HloModule test_module ENTRY test { ROOT v = f16[] parameter(0) })"; + HloModule test_module + ENTRY test { + a = s32[] constant(1) + b = s32[1,1] constant({{2}}) + c = s32[1,2] constant({{2,2}}) + d = f32[] constant(1) + e = f32[] constant(1.25) + ROOT tuple = (s32[], s32[1,1], s32[1,2], f32[], f32[]) tuple(a,b,c,d,e) + })"; TF_ASSERT_OK_AND_ASSIGN(auto hlo_module, ParseHloString(kModuleStr)); auto* root = hlo_module->entry_computation()->root_instruction(); - EXPECT_FALSE(Match(root, match::ConstantScalar(42))); + const HloInstruction* a = root->operand(0); + const HloInstruction* b = root->operand(1); + const HloInstruction* c = root->operand(2); + const HloInstruction* d = root->operand(3); + const HloInstruction* e = root->operand(4); + EXPECT_TRUE(Match(a, ConstantScalar())); + EXPECT_TRUE(Match(a, ConstantScalar(1))); + EXPECT_TRUE(Match(a, ConstantEffectiveScalar())); + EXPECT_TRUE(Match(a, ConstantEffectiveScalar(1))); + EXPECT_FALSE(Match(a, ConstantScalar(2))); + EXPECT_FALSE(Match(a, ConstantScalar(2.01))); + EXPECT_FALSE(Match(a, ConstantEffectiveScalar(2))); + EXPECT_FALSE(Match(a, ConstantEffectiveScalar(1.01))); + + EXPECT_FALSE(Match(b, ConstantScalar())); + EXPECT_FALSE(Match(b, ConstantScalar(2))); + EXPECT_TRUE(Match(b, ConstantEffectiveScalar())); + EXPECT_TRUE(Match(b, ConstantEffectiveScalar(2))); + + EXPECT_FALSE(Match(c, ConstantScalar())); + EXPECT_FALSE(Match(c, ConstantScalar(2))); + EXPECT_FALSE(Match(c, ConstantEffectiveScalar())); + EXPECT_FALSE(Match(c, ConstantEffectiveScalar(2))); + + EXPECT_TRUE(Match(d, ConstantScalar(1))); + EXPECT_TRUE(Match(d, ConstantEffectiveScalar(1))); + EXPECT_TRUE(Match(d, ConstantScalar(1.0))); + EXPECT_TRUE(Match(d, ConstantEffectiveScalar(1.0))); + + EXPECT_TRUE(Match(e, ConstantScalar(1.25f))); + EXPECT_TRUE(Match(e, ConstantScalar(1.25))); + EXPECT_TRUE(Match(e, ConstantEffectiveScalar(1.25))); + EXPECT_FALSE(Match(e, ConstantScalar(1))); + EXPECT_FALSE(Match(e, ConstantEffectiveScalar(1))); + + const HloInstruction* instr = nullptr; + EXPECT_TRUE(Match(a, ConstantScalar(&instr))); + EXPECT_EQ(instr, a); + + instr = nullptr; + EXPECT_TRUE(Match(a, ConstantScalar(&instr, 1))); + EXPECT_EQ(instr, a); + + instr = nullptr; + EXPECT_TRUE(Match(a, ConstantEffectiveScalar(&instr))); + EXPECT_EQ(instr, a); + + instr = nullptr; + EXPECT_TRUE(Match(a, ConstantEffectiveScalar(&instr, 1))); + EXPECT_EQ(instr, a); } TEST(PatternMatcherTest, MultiplyAnyOrder) { @@ -267,6 +322,15 @@ TEST(PatternMatcherTest, MultiplyAnyOrder) { root, MultiplyAnyOrder(&instr, ConstantScalar(42), ConstantScalar(52)))); EXPECT_TRUE(Match( root, MultiplyAnyOrder(&instr, ConstantScalar(52), ConstantScalar(42)))); + + // Check that MultiplyAnyOrder exposes the same API as Op(), so we can call + // e.g. IsNonConstant() on it. + EXPECT_TRUE(Match( + root, MultiplyAnyOrder(&instr, ConstantScalar(42), ConstantScalar(52)) + .IsNonConstant())); + EXPECT_TRUE( + Match(root, MultiplyAnyOrder(ConstantScalar(42), ConstantScalar(52)) + .IsNonConstant())); } TEST(PatternMatcherTest, AnyOfShortCircuit) { @@ -315,14 +379,22 @@ TEST(PatternMatcherTest, AllOf) { TF_ASSERT_OK_AND_ASSIGN(auto hlo_module, ParseHloString(kModuleStr)); auto* root = hlo_module->entry_computation()->root_instruction(); + auto f16_scalar = ShapeUtil::MakeShape(F16, {}); + auto f16_pattern = Constant().WithShapeEqualTo(&f16_scalar); + auto f16_compatible_pattern = Constant().WithShapeCompatibleTo(&f16_scalar); auto scalar_pattern = Constant().WithShape(match::Shape().IsScalar()); - auto f16_pattern = Constant().WithShape(match::Shape().WithElementType(F16)); ASSERT_TRUE(Match(root, scalar_pattern)); ASSERT_TRUE(Match(root, f16_pattern)); - EXPECT_TRUE(Match(root, AllOf(scalar_pattern, f16_pattern))); - EXPECT_TRUE(Match(root, AllOf(f16_pattern, scalar_pattern))); + ASSERT_TRUE(Match(root, f16_compatible_pattern)); + EXPECT_TRUE(Match(root, AllOf(scalar_pattern, f16_pattern, + f16_compatible_pattern))); + EXPECT_TRUE( + Match(root, AllOf(f16_pattern, f16_compatible_pattern, + scalar_pattern))); EXPECT_FALSE( Match(root, AllOf(Broadcast(Op()), f16_pattern))); + EXPECT_FALSE(Match( + root, AllOf(Broadcast(Op()), f16_compatible_pattern))); EXPECT_FALSE( Match(root, AllOf(Broadcast(Op()), scalar_pattern))); } @@ -394,5 +466,470 @@ TEST(PatternMatcherTest, TestCaptureMatchedSubPatternForAnyOf) { EXPECT_EQ(nullptr, addend2); } +TEST(PatternMatcherTest, TestConcat) { + using match::Concatenate; + using match::ConstantScalar; + using match::Op; + using match::Reshape; + + constexpr char kModuleStr[] = R"( + HloModule test_module + ENTRY test { + c1 = u32[] constant(1) + c2 = u32[] constant(2) + c3 = u32[] constant(3) + c4 = u32[] constant(4) + r1 = u32[1] reshape(c1) + r2 = u32[1] reshape(c2) + r3 = u32[1] reshape(c3) + r4 = u32[1] reshape(c4) + ROOT concat = u32[4] concatenate(r1, r2, r3, r4), dimensions={0} + })"; + TF_ASSERT_OK_AND_ASSIGN(auto hlo_module, ParseHloString(kModuleStr)); + auto* root = hlo_module->entry_computation()->root_instruction(); + ASSERT_TRUE(Match( + root, + Concatenate(Reshape(ConstantScalar(1)), Reshape(ConstantScalar(2)), + Reshape(ConstantScalar(3)), Reshape(ConstantScalar(4))))); + ASSERT_FALSE(Match( + root, + Concatenate(Reshape(ConstantScalar(2)), Reshape(ConstantScalar(1)), + Reshape(ConstantScalar(3)), Reshape(ConstantScalar(4))))); + ASSERT_FALSE(Match( + root, Concatenate(Reshape(ConstantScalar(1)), Reshape(ConstantScalar(2)), + Reshape(ConstantScalar(3))))); + ASSERT_FALSE(Match( + root, Concatenate(Reshape(ConstantScalar(2)), Reshape(ConstantScalar(3)), + Reshape(ConstantScalar(4))))); +} + +template +string Description(const Pattern& pattern) { + std::stringstream ss; + pattern.DescribeTo(&ss); + return ss.str(); +} + +template +string Explanation(Elem* elem, const Pattern& pattern) { + std::stringstream ss; + MatchOption options{/*.capture=*/true, /*.explain_os=*/&ss}; + Match(elem, pattern, options); + return ss.str(); +} +template +string Explanation(const std::unique_ptr& elem, const Pattern& pattern) { + return Explanation(elem.get(), pattern); +} +template +string Explanation(const Elem& elem, const Pattern& pattern) { + return Explanation(&elem, pattern); +} + +// Helper macro for checking a pattern's description and the explanation printed +// when attempting to match (and presumably failing) on a given object. +// +// We use a macro rather than a function because we want good line numbers in +// errors. We use this rather than writing a helper that returns a pair of +// (description, explanation) and doing something like +// +// EXPECT_THAT(DescAndExplanation(...), ::testing::Pair(..., ...)); +// +// because EXPECT_EQ prints a unified diff if multiline string comparison fails, +// while EXPECT_THAT does not. This unified diff makes the errors much easier +// to read. +#define EXPECT_DESC_AND_EXPLANATION(elem, pattern, expected_desc, \ + expected_explanation) \ + do { \ + EXPECT_EQ(Description(pattern), (expected_desc)); \ + EXPECT_EQ(Explanation((elem), (pattern)), expected_explanation); \ + } while (0) + +TEST(PatternMatcherTest, LayoutDescribeToAndExplain) { + auto layout = LayoutUtil::MakeLayout({1, 2}); + auto layout2 = LayoutUtil::MakeLayout({2, 2}); + + EXPECT_DESC_AND_EXPLANATION(static_cast(nullptr), m::Layout(), + "a layout", "Layout is null"); + EXPECT_DESC_AND_EXPLANATION(layout2, m::Layout().EqualTo(&layout), + "a layout equal to {1,2}", + "Layout {2,2} is not equal to expected {1,2}"); + EXPECT_DESC_AND_EXPLANATION(layout2, m::Layout().WithSparseFormat(), + "a layout with format SPARSE", + "Layout has format DENSE but expected SPARSE"); + EXPECT_DESC_AND_EXPLANATION(layout, + m::Layout().EqualTo(&layout).WithSparseFormat(), + "a layout:\n" + " * equal to {1,2} AND\n" + " * with format SPARSE", + "Layout has format DENSE but expected SPARSE"); +} + +TEST(PatternMatcherTest, ShapeDescribeToAndExplain) { + auto shape = ShapeUtil::MakeShapeWithLayout(F32, {1, 2}, {0, 1}); + auto layout = shape.layout(); + + EXPECT_DESC_AND_EXPLANATION(static_cast(nullptr), m::Shape(), + "a shape", "Shape is null"); + EXPECT_DESC_AND_EXPLANATION( + ShapeUtil::MakeShapeWithLayout(F32, {1, 2}, {1, 0}), + m::Shape().EqualTo(&shape), "a shape equal to f32[1,2]{0,1}", + "Shape not equal to f32[1,2]{0,1}\n" + "in f32[1,2]{1,0}"); + EXPECT_DESC_AND_EXPLANATION(ShapeUtil::MakeShape(F32, {2, 2}), + m::Shape().CompatibleTo(&shape), + "a shape compatible with f32[1,2]", + "Shape not compatible with f32[1,2]\n" + "in f32[2,2]{1,0}"); + EXPECT_DESC_AND_EXPLANATION(shape, m::Shape().WithElementType(F16), + "a shape with element type F16", + "Shape does not have element type F16\n" + "in f32[1,2]{0,1}"); + EXPECT_DESC_AND_EXPLANATION(shape, m::Shape().IsScalar(), + "a shape that represents a scalar", + "Shape is not a scalar\n" + "in f32[1,2]{0,1}"); + EXPECT_DESC_AND_EXPLANATION(ShapeUtil::MakeNil(), m::Shape().IsArray(), + "a shape that represents an array", + "Shape is not an array\n" + "in ()"); + EXPECT_DESC_AND_EXPLANATION(shape, m::Shape().IsTuple(), + "a shape that represents a tuple", + "Shape is not a tuple\n" + "in f32[1,2]{0,1}"); + EXPECT_DESC_AND_EXPLANATION(shape, m::Shape().IsEffectiveScalar(), + "a shape that is an effective scalar", + "Shape is not an effective scalar\n" + "in f32[1,2]{0,1}"); + EXPECT_DESC_AND_EXPLANATION(shape, m::Shape().WithRank(42), + "a shape that has 42 dimensions", + "Shape does not have rank 42\n" + "in f32[1,2]{0,1}"); + EXPECT_DESC_AND_EXPLANATION(shape, m::Shape().WithRank(0), + "a shape that is a scalar", + "Shape is not a scalar\n" + "in f32[1,2]{0,1}"); + EXPECT_DESC_AND_EXPLANATION(shape, m::Shape().WithRank(1).IsArray(), + "a shape:\n" + " * that has 1 dimension AND\n" + " * that represents an array", + "Shape does not have rank 1\n" + "in f32[1,2]{0,1}"); + EXPECT_DESC_AND_EXPLANATION(ShapeUtil::MakeNil(), + m::Shape().IsArray().WithRank(1), + "a shape:\n" + " * that represents an array AND\n" + " * that has 1 dimension", + "Shape is not an array\n" + "in ()"); + EXPECT_DESC_AND_EXPLANATION( + ShapeUtil::MakeShapeWithLayout(F32, {1, 2}, {1, 0}), + m::Shape().WithLayoutEqualTo(&layout), + "a shape with\n a layout equal to {0,1}", + "Layout {1,0} is not equal to expected {0,1}\n" + "in f32[1,2]{1,0}"); + EXPECT_DESC_AND_EXPLANATION( + shape, m::Shape().WithLayout(m::Layout().WithSparseFormat()), + "a shape with\n a layout with format SPARSE", + "Layout has format DENSE but expected SPARSE\n" + "in f32[1,2]{0,1}"); + EXPECT_DESC_AND_EXPLANATION(shape, + m::Shape().WithSubshapeEqualTo({10}, &shape), + "a shape with subshape at index {10} which is\n" + " a shape equal to f32[1,2]{0,1}", + "No subshape at {10}\n" + "in f32[1,2]{0,1}"); + EXPECT_DESC_AND_EXPLANATION( + ShapeUtil::MakeTupleShape({ShapeUtil::MakeShape(F32, {2, 2})}), + m::Shape().WithSubshapeEqualTo({0}, &shape), + "a shape with subshape at index {0} which is\n" + " a shape equal to f32[1,2]{0,1}", + "Shape not equal to f32[1,2]{0,1}\n" + "in f32[2,2]{1,0}\n" + "in subshape at {0}\n" + "in (f32[2,2])"); + EXPECT_DESC_AND_EXPLANATION(shape, + m::Shape().WithSubshapeCompatibleTo({10}, &shape), + "a shape with subshape at index {10} which is\n" + " a shape compatible with f32[1,2]", + "No subshape at {10}\n" + "in f32[1,2]{0,1}"); + EXPECT_DESC_AND_EXPLANATION( + ShapeUtil::MakeTupleShape({ShapeUtil::MakeShape(F32, {2, 2})}), + m::Shape().WithSubshapeCompatibleTo({0}, &shape), + "a shape with subshape at index {0} which is\n" + " a shape compatible with f32[1,2]", + "Shape not compatible with f32[1,2]\n" + "in f32[2,2]{1,0}\n" + "in subshape at {0}\n" + "in (f32[2,2])"); + EXPECT_DESC_AND_EXPLANATION( + ShapeUtil::MakeTupleShape({ShapeUtil::MakeTupleShape({shape})}), + m::Shape().WithSubshape({0, 0}, m::Shape().IsScalar()), + "a shape with subshape at index {0,0} which is\n" + " a shape that represents a scalar", + "Shape is not a scalar\n" + "in f32[1,2]{0,1}\n" + "in subshape at {0,0}\n" + "in ((f32[1,2]))"); +} + +std::unique_ptr SetName(absl::string_view name, + std::unique_ptr instr) { + instr->SetAndSanitizeName(string(name)); + return instr; +} + +TEST(PatternMatcherTest, HloInstructionDescribeToAndExplain) { + std::unique_ptr iota = + SetName("i", HloInstruction::CreateIota(ShapeUtil::MakeShape(S32, {42}), + /*iota_dimension=*/0)); + std::unique_ptr constant = + SetName("c", HloInstruction::CreateConstant(LiteralUtil::CreateR0(0))); + + EXPECT_DESC_AND_EXPLANATION(static_cast(nullptr), + m::Op(), "an HloInstruction", + "HloInstruction* is null"); + EXPECT_DESC_AND_EXPLANATION(iota, m::Op().WithName("foo"), + "an HloInstruction named \"foo\"", + "HloInstruction not named \"foo\"\n" + "in i = s32[42]{0} iota(), iota_dimension=0"); + EXPECT_DESC_AND_EXPLANATION(iota, m::Op().WithOpcode(HloOpcode::kAdd), + "an HloInstruction with opcode add", + "HloInstruction doesn't have opcode add\n" + "in i = s32[42]{0} iota(), iota_dimension=0"); + EXPECT_DESC_AND_EXPLANATION( + constant, m::Op().IsNonConstant(), + "an HloInstruction with any opcode other than constant", + "HloInstruction has opcode constant, expected anything else\n" + "in c = s32[] constant(0)"); + EXPECT_DESC_AND_EXPLANATION(iota, m::Op().WithNumOperands(42), + "an HloInstruction with 42 operands", + "HloInstruction doesn't have 42 operands\n" + "in i = s32[42]{0} iota(), iota_dimension=0"); + EXPECT_DESC_AND_EXPLANATION(iota, m::Op().WithShape(m::Shape().IsTuple()), + "an HloInstruction outputting\n" + " a shape that represents a tuple", + "Shape is not a tuple\n" + "in s32[42]{0}\n" + "in output shape\n" + "in i = s32[42]{0} iota(), iota_dimension=0"); + EXPECT_DESC_AND_EXPLANATION( + iota, m::Op().WithOperand(2, m::Op().WithOpcode(HloOpcode::kAdd)), + "an HloInstruction with operand 2 which is:\n" + " an HloInstruction with opcode add", + "desired operand index 2 is out of bounds\n" + "in i = s32[42]{0} iota(), iota_dimension=0"); + + EXPECT_DESC_AND_EXPLANATION( + SetName("a", HloInstruction::CreateBinary(ShapeUtil::MakeShape(S32, {}), + HloOpcode::kAdd, constant.get(), + constant.get())), + m::Op().WithOperand(1, m::Op().IsNonConstant()), + "an HloInstruction with operand 1 which is:\n" + " an HloInstruction with any opcode other than constant", + "HloInstruction has opcode constant, expected anything else\n" + "in c = s32[] constant(0)\n" + "in operand 1\n" + "in a = s32[] add(s32[] c, s32[] c)"); + EXPECT_DESC_AND_EXPLANATION( + iota, m::Op().WithFusionKind(HloInstruction::FusionKind::kLoop), + "an HloInstruction with fusion kind kLoop", + "HloInstruction does not have fusion kind kLoop; it's not a fusion\n" + "in i = s32[42]{0} iota(), iota_dimension=0"); + EXPECT_DESC_AND_EXPLANATION( + iota, m::Op().WithTupleIndex(42), + "an HloInstruction which is a GTE with index 42", + "HloInstruction is not a GTE with index 42; it's not a GTE at all\n" + "in i = s32[42]{0} iota(), iota_dimension=0"); + EXPECT_DESC_AND_EXPLANATION(iota, m::Op().IsConstantScalar(), + "an HloInstruction which is a constant scalar", + "HloInstruction is not a constant\n" + "in i = s32[42]{0} iota(), iota_dimension=0"); + EXPECT_DESC_AND_EXPLANATION( + SetName("c", HloInstruction::CreateConstant( + LiteralUtil::CreateR1({1, 2}))), + m::Op().IsConstantEffectiveScalar(), + "an HloInstruction which is a constant effective scalar", + "HloInstruction is not an effective scalar\n" + "in c = s32[2]{0} constant({1, 2})"); + EXPECT_DESC_AND_EXPLANATION( + SetName("c", HloInstruction::CreateConstant(LiteralUtil::CreateR0(10))), + m::Op().IsConstantScalar(42), + "an HloInstruction which is a constant scalar with value 42", + "HloInstruction's constant value 10 did not match expected value 42\n" + "in c = s32[] constant(10)"); + EXPECT_DESC_AND_EXPLANATION( + SetName("c", HloInstruction::CreateConstant(LiteralUtil::CreateR0(2.25))), + m::Op().IsConstantEffectiveScalar(1.25), + "an HloInstruction which is a constant effective scalar with value 1.25", + "HloInstruction's constant value 2.25 did not match expected value 1.25\n" + "in c = f64[] constant(2.25)"); + EXPECT_DESC_AND_EXPLANATION( + constant, m::Op().Is(iota.get()), + absl::StrCat("an HloInstruction which is 0x", absl::Hex(iota.get()), + " (i = s32[42]{0} iota(), iota_dimension=0)"), + absl::StrCat("HloInstruction 0x", absl::Hex(constant.get()), " is not 0x", + absl::Hex(iota.get()), + " (i = s32[42]{0} iota(), iota_dimension=0)\n" + "in c = s32[] constant(0)")); +} + +TEST(PatternMatcherTest, HloInstructionMatcherAnyOrderDescribeTo) { + auto scalar_s32 = ShapeUtil::MakeShape(S32, {}); + EXPECT_DESC_AND_EXPLANATION( + SetName("a", HloInstruction::CreateBinary( + scalar_s32, HloOpcode::kAdd, + SetName("b", HloInstruction::CreateConstant( + LiteralUtil::CreateR0(0))) + .get(), + SetName("c", HloInstruction::CreateConstant( + LiteralUtil::CreateR0(0))) + .get())), + m::AddAnyOrder(m::Op().WithName("b"), m::Op().WithName("bar")), + "an HloInstruction:\n" + " * with opcode add AND\n" + " * with two operands in either order:\n" + " - an HloInstruction named \"b\"\n" + " - an HloInstruction named \"bar\"", + "HloInstruction's operands (ignoring order) did not match second " + "matcher. Specifically,\n" + " - an HloInstruction named \"bar\"\n" + "does not match LHS:\n" + " - HloInstruction not named \"bar\"\n" + " in b = s32[] constant(0)\n" + "does not match RHS:\n" + " - HloInstruction not named \"bar\"\n" + " in c = s32[] constant(0)\n" + "in a = s32[] add(s32[] b, s32[] c)"); + + EXPECT_DESC_AND_EXPLANATION( + SetName("a", + HloInstruction::CreateBinary( + scalar_s32, HloOpcode::kAdd, + HloInstruction::CreateParameter(0, scalar_s32, "p").get(), + SetName("c", HloInstruction::CreateConstant( + LiteralUtil::CreateR0(0))) + .get())), + m::AddAnyOrder(m::Op().IsConstantScalar(), m::Op().IsConstant()), + "an HloInstruction:\n" + " * with opcode add AND\n" + " * with two operands in either order:\n" + " - an HloInstruction which is a constant scalar\n" + " - an HloInstruction with opcode constant", + "HloInstruction's LHS operand did not match either of the two matchers. " + "Specifically,\n" + " - an HloInstruction which is a constant scalar\n" + "does not match LHS:\n" + " - HloInstruction is not a constant\n" + " in p = s32[] parameter(0)\n" + "and\n" + " - an HloInstruction with opcode constant\n" + "does not match LHS:\n" + " - HloInstruction doesn't have opcode constant\n" + " in p = s32[] parameter(0)\n" + "in a = s32[] add(s32[] p, s32[] c)"); +} + +TEST(PatternMatcherTest, AnyOfMatcherDescribeToAndExplain) { + EXPECT_DESC_AND_EXPLANATION( + SetName("c", HloInstruction::CreateConstant(LiteralUtil::CreateR0(0))), + m::AnyOf(m::Op().WithName("foo"), + m::Op().WithName("bar")), + "any of:\n" + " - an HloInstruction named \"foo\" OR\n" + " - an HloInstruction named \"bar\"", + "None of the following matchers succeeded:\n" + "Matcher #1\n" + " - an HloInstruction named \"foo\"\n" + "failed with\n" + " - HloInstruction not named \"foo\"\n" + " in c = s32[] constant(0)\n" + "Matcher #2\n" + " - an HloInstruction named \"bar\"\n" + "failed with\n" + " - HloInstruction not named \"bar\"\n" + " in c = s32[] constant(0)"); +} + +TEST(PatternMatcherTest, Parameter) { + auto param = + HloInstruction::CreateParameter(1, ShapeUtil::MakeShape(F32, {}), "p1"); + auto non_param = + SetName("c", HloInstruction::CreateConstant(LiteralUtil::CreateR0(0))); + EXPECT_FALSE(Match(param.get(), m::Parameter(0))); + EXPECT_TRUE(Match(param.get(), m::Parameter())); + EXPECT_TRUE(Match(param.get(), m::Parameter(1))); + EXPECT_FALSE(Match(non_param.get(), m::Parameter())); + EXPECT_FALSE(Match(non_param.get(), m::Parameter(1))); + + EXPECT_DESC_AND_EXPLANATION(non_param, m::Parameter(1), + "an HloInstruction:\n" + " * with opcode parameter AND\n" + " * which is parameter 1", + "HloInstruction doesn't have opcode parameter\n" + "in c = s32[] constant(0)"); + EXPECT_EQ(Explanation(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShape(F32, {}), "p0"), + m::Parameter(1)), + "HloInstruction is not parameter 1\n" + "in p0 = f32[] parameter(0)"); +} + +TEST(PatternMatcherTest, OneUseAndOneUser) { + auto param = + HloInstruction::CreateParameter(0, ShapeUtil::MakeShape(F32, {}), "p0"); + + EXPECT_FALSE(Match(param.get(), m::Op().WithOneUse())); + EXPECT_DESC_AND_EXPLANATION( + param, m::Op().WithOneUse(), + "an HloInstruction which has exactly one use", + "HloInstruction has 0 users, but expected exactly one.\n" + "in p0 = f32[] parameter(0)"); + + EXPECT_FALSE(Match(param.get(), m::Op().WithOneUser())); + EXPECT_DESC_AND_EXPLANATION( + param, m::Op().WithOneUser(), + "an HloInstruction which has exactly one user (but possibly is used " + "multiple times by that instruction)", + "HloInstruction has 0 users, but expected exactly one.\n" + "in p0 = f32[] parameter(0)"); + + { + auto reshape = + SetName("r", HloInstruction::CreateReshape( + ShapeUtil::MakeShape(F32, {1}), param.get())); + EXPECT_TRUE(Match(param.get(), m::Op().WithOneUse())); + EXPECT_TRUE(Match(param.get(), m::Op().WithOneUser())); + + auto reshape1 = + SetName("r1", HloInstruction::CreateReshape( + ShapeUtil::MakeShape(F32, {1}), param.get())); + EXPECT_FALSE(Match(param.get(), m::Op().WithOneUse())); + EXPECT_FALSE(Match(param.get(), m::Op().WithOneUser())); + + const char* kMultipleUserExplanation = + "HloInstruction has 2 users, but expected exactly one.\n" + "All users:\n" + " - r = f32[1]{0} reshape(f32[] p0)\n" + " - r1 = f32[1]{0} reshape(f32[] p0)\n" + "in p0 = f32[] parameter(0)"; + EXPECT_EQ(Explanation(param.get(), m::Op().WithOneUse()), + kMultipleUserExplanation); + EXPECT_EQ(Explanation(param.get(), m::Op().WithOneUser()), + kMultipleUserExplanation); + } + + auto add = SetName("add", HloInstruction::CreateBinary( + ShapeUtil::MakeShape(F32, {}), HloOpcode::kAdd, + param.get(), param.get())); + EXPECT_TRUE(Match(param.get(), m::Op().WithOneUser())); + EXPECT_FALSE(Match(param.get(), m::Op().WithOneUse())); + EXPECT_EQ(Explanation(param.get(), m::Op().WithOneUse()), + "HloInstruction is used 2 times by its user, but is expected to be " + "used just once: add = f32[] add(f32[] p0, f32[] p0)\n" + "in p0 = f32[] parameter(0)"); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/service/platform_util.cc b/tensorflow/compiler/xla/service/platform_util.cc index c522e7ae23b734090f85d241bf365fccc37f0adb..886a0545624927fa77528141f61d8ecb6bec180a 100644 --- a/tensorflow/compiler/xla/service/platform_util.cc +++ b/tensorflow/compiler/xla/service/platform_util.cc @@ -21,7 +21,7 @@ limitations under the License. #include "absl/strings/ascii.h" #include "absl/strings/str_join.h" -#include "tensorflow/compiler/xla/legacy_flags/debug_options_flags.h" +#include "tensorflow/compiler/xla/debug_options_flags.h" #include "tensorflow/compiler/xla/service/compiler.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/statusor.h" @@ -59,22 +59,20 @@ string CanonicalPlatformName(const string& name) { /* static */ StatusOr> PlatformUtil::GetSupportedPlatforms() { - se::MultiPlatformManager::PlatformMap platform_map; - se::port::Status platforms_status = se::MultiPlatformManager::WithPlatforms( - [&platform_map](se::MultiPlatformManager::PlatformMap* map) { - platform_map = *map; - return se::port::Status::OK(); - }); - if (platform_map.empty()) { + std::vector all_platforms = + se::MultiPlatformManager::AllPlatforms(); + if (all_platforms.empty()) { LOG(WARNING) << "no executor platforms available: platform map is empty"; } // Gather all platforms which have an XLA compiler. std::vector platforms; - for (auto& platform_pair : platform_map) { - auto* platform = platform_pair.second; + 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 " @@ -210,7 +208,9 @@ static bool IsDeviceSupported(se::StreamExecutor* executor) { } /* static */ StatusOr> -PlatformUtil::GetStreamExecutors(se::Platform* platform) { +PlatformUtil::GetStreamExecutors( + se::Platform* platform, + const absl::optional>& allowed_devices) { int device_count = platform->VisibleDeviceCount(); if (device_count <= 0) { return NotFound("no %s devices found", platform->Name()); @@ -222,8 +222,8 @@ PlatformUtil::GetStreamExecutors(se::Platform* platform) { // fix the number of devices to one. However we do let the user override // this behavior to help run tests on the host that run models in parallel // across multiple devices. - device_count = legacy_flags::GetDebugOptionsFromFlags() - .xla_force_host_platform_device_count(); + device_count = + GetDebugOptionsFromFlags().xla_force_host_platform_device_count(); } std::vector stream_executors(device_count, nullptr); VLOG(1) << "Initializing devices"; @@ -231,6 +231,17 @@ PlatformUtil::GetStreamExecutors(se::Platform* platform) { tensorflow::thread::ThreadPool thread_pool( tensorflow::Env::Default(), "device_initialization", device_count); for (int i = 0; i < device_count; ++i) { + // Once a stream executor is instantiated it will cause allocations on + // the device, for example for GPUs cuda context, cudnn handles etc. will + // be constructed. By constructing stream executors only on the + // allowed_devices, we don't make any allocations on other devices. + // This helps in multi-process executions on the same host like horovod or + // shared hosts. + if (allowed_devices && allowed_devices->count(i) == 0) { + VLOG(1) << "Not initializing StreamExecutor for device " << i + << " since it is not in the visible device list"; + continue; + } thread_pool.Schedule([platform, i, &stream_executors]() { VLOG(1) << "Started device init " << i; se::StreamExecutorConfig config; @@ -252,8 +263,8 @@ PlatformUtil::GetStreamExecutors(se::Platform* platform) { // 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/platform_util.h b/tensorflow/compiler/xla/service/platform_util.h index 571451ba43a81d19b70e4954e45d3447f15dcedc..592b20282f334e12e0d7a7f683c9a6ab59d21fea 100644 --- a/tensorflow/compiler/xla/service/platform_util.h +++ b/tensorflow/compiler/xla/service/platform_util.h @@ -16,6 +16,7 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_PLATFORM_UTIL_H_ #define TENSORFLOW_COMPILER_XLA_SERVICE_PLATFORM_UTIL_H_ +#include #include #include @@ -60,10 +61,14 @@ class PlatformUtil { // Returns a vector of StreamExecutors for the given platform. The vector is // indexed by device ordinal (device numbering used by StreamExecutor). If an // element is nullptr, then the device is present by not supported by XLA. + // If populated, only the devices in allowed_devices will have + // their StreamExecutors initialized, otherwise all StreamExecutors will be + // initialized and returned. // // If the platform has no visible devices, a not-found error is returned. static StatusOr> GetStreamExecutors( - se::Platform* platform); + se::Platform* platform, + const absl::optional>& allowed_devices = absl::nullopt); private: TF_DISALLOW_COPY_AND_ASSIGN(PlatformUtil); diff --git a/tensorflow/compiler/xla/service/reduce_precision_insertion.cc b/tensorflow/compiler/xla/service/reduce_precision_insertion.cc index 688cceff0cd10df62a4093f00ad3331ca77652e0..b70cb7057477a338bfb36ebab76237b30d018e41 100644 --- a/tensorflow/compiler/xla/service/reduce_precision_insertion.cc +++ b/tensorflow/compiler/xla/service/reduce_precision_insertion.cc @@ -111,7 +111,7 @@ StatusOr ReducePrecisionInsertion::insert_on_inputs( VLOG(2) << "Adding to operand " << i << ": " << operand; if (!is_valid_shape(operand->shape())) { - VLOG(2) << "Skipped: value is not an F32 vector"; + VLOG(2) << "Skipped: value is not of type F32"; continue; } @@ -168,7 +168,7 @@ StatusOr ReducePrecisionInsertion::insert_on_outputs( << instruction->ToString(); if (!is_valid_shape(instruction->shape())) { - VLOG(2) << "Skipped: value is not an F32 nonscalar array"; + VLOG(2) << "Skipped: value is not of type F32"; continue; } diff --git a/tensorflow/compiler/xla/service/reduce_precision_insertion.h b/tensorflow/compiler/xla/service/reduce_precision_insertion.h index 0b4e82e8d606cf2cacfab42d07c2201939d5e10b..76c6a87f176ec9c6f8e49c25278c6dad703e3c7c 100644 --- a/tensorflow/compiler/xla/service/reduce_precision_insertion.h +++ b/tensorflow/compiler/xla/service/reduce_precision_insertion.h @@ -118,13 +118,7 @@ class ReducePrecisionInsertion : public HloModulePass { // equivalent behavior can be obtained by adding ReducePrecision // instructions after the instructions that pull the F32 arrays out of // the tuples. - // - // TODO(b/64093391): Remove the IsScalar check once this won't cause - // failures on the GPU backend if the ReducePrecision instruction ends up - // inserted between a scalar constant and the init_value argument of a - // Reduce operation. - return shape.element_type() == PrimitiveType::F32 && - !ShapeUtil::IsScalar(shape); + return shape.element_type() == PrimitiveType::F32; } // Is this instruction one such that following or preceding it with a new diff --git a/tensorflow/compiler/xla/service/reduce_precision_insertion_test.cc b/tensorflow/compiler/xla/service/reduce_precision_insertion_test.cc index 69e4b534bd8e3aeab8b729f3e594a10b4368f15f..efeec96571455d8a9e4b7837dd7286392c12f1a3 100644 --- a/tensorflow/compiler/xla/service/reduce_precision_insertion_test.cc +++ b/tensorflow/compiler/xla/service/reduce_precision_insertion_test.cc @@ -54,7 +54,34 @@ TEST_F(ReducePrecisionInsertionTest, BeforeUnaryInstruction) { HloInstruction* b = builder.AddInstruction( HloInstruction::CreateUnary(shape, HloOpcode::kCos, a)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + // Confirm expected state before adding ops. + EXPECT_EQ(computation->root_instruction(), b); + EXPECT_EQ(b->operand(0), a); + + EXPECT_TRUE(InsertOps(module.get(), HloReducePrecisionOptions::OP_INPUTS, + [](const HloInstruction* instruction) { + return instruction->opcode() == HloOpcode::kCos; + })); + + // Confirm expected graph after adding ops. + EXPECT_EQ(computation->root_instruction(), b); + EXPECT_THAT(b->operand(0), op::ReducePrecision(a)); +} + +TEST_F(ReducePrecisionInsertionTest, BeforeUnaryScalarInstruction) { + auto builder = HloComputation::Builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {}); + + // Create a simple graph with a parameter feeding a unary cosine function. + HloInstruction* a = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "a")); + HloInstruction* b = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kCos, a)); + + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); // Confirm expected state before adding ops. @@ -84,7 +111,7 @@ TEST_F(ReducePrecisionInsertionTest, BeforeBinaryInstruction) { HloInstruction* c = builder.AddInstruction( HloInstruction::CreateBinary(shape, HloOpcode::kAdd, a, b)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); // Confirm expected state before adding ops. @@ -113,7 +140,7 @@ TEST_F(ReducePrecisionInsertionTest, BeforeZeroInputInstruction) { HloInstruction* b = builder.AddInstruction( HloInstruction::CreateUnary(shape, HloOpcode::kCos, a)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); // Confirm expected state before adding ops. @@ -146,7 +173,7 @@ TEST_F(ReducePrecisionInsertionTest, AvoidAddingDuplicateInstructions) { HloInstruction* d = builder.AddInstruction( HloInstruction::CreateBinary(shape, HloOpcode::kAdd, b, c)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); // Confirm expected state before adding ops. @@ -178,7 +205,7 @@ TEST_F(ReducePrecisionInsertionTest, AfterRootInstruction) { HloInstruction* b = builder.AddInstruction( HloInstruction::CreateUnary(shape, HloOpcode::kCos, a)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); // Confirm expected state before adding ops. @@ -215,7 +242,7 @@ TEST_F(ReducePrecisionInsertionTest, AfterNonRootInstruction) { HloInstruction* c = builder.AddInstruction( HloInstruction::CreateBinary(shape, HloOpcode::kAdd, a_cos, b_cos)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); // Confirm expected graph before adding ops. @@ -242,7 +269,7 @@ TEST_F(ReducePrecisionInsertionTest, OutputIsNotFloat) { HloInstruction* y = builder.AddInstruction( HloInstruction::CreateUnary(shape, HloOpcode::kCos, x)); - auto module = CreateNewModule(); + auto module = CreateNewUnverifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); // Confirm expected graph before adding ops. @@ -268,7 +295,7 @@ TEST_F(ReducePrecisionInsertionTest, ShouldReduceOutputPrecisionIsFalse) { HloInstruction* y = builder.AddInstruction( HloInstruction::CreateUnary(shape, HloOpcode::kCos, x)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); // Confirm expected graph before adding ops. @@ -294,7 +321,7 @@ TEST_F(ReducePrecisionInsertionTest, InsertionIsNotRecursive) { HloInstruction* b = builder.AddInstruction( HloInstruction::CreateReducePrecision(shape, a, 8, 23)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); // Confirm expected state before adding ops. @@ -321,7 +348,7 @@ TEST_F(ReducePrecisionInsertionTest, SkipRedundantReducePrecisionAfter) { HloInstruction* y = builder.AddInstruction( HloInstruction::CreateReducePrecision(shape, x, 5, 10)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); // Confirm expected graph before adding ops. @@ -349,7 +376,7 @@ TEST_F(ReducePrecisionInsertionTest, AddNonRedundantReducePrecision) { HloInstruction* y = builder.AddInstruction( HloInstruction::CreateReducePrecision(shape, x, 8, 23)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); // Confirm expected graph before adding ops. @@ -375,7 +402,7 @@ TEST_F(ReducePrecisionInsertionTest, IgnoreOpsInsideFusionNode) { builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "x")); HloInstruction* y = builder.AddInstruction( HloInstruction::CreateUnary(shape, HloOpcode::kCos, x)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); // Manually fuse the kCos operation into a fusion operation. @@ -411,7 +438,7 @@ TEST_F(ReducePrecisionInsertionTest, OpGetsInsertedInHeadOfFusionNode) { builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "x")); HloInstruction* y = builder.AddInstruction( HloInstruction::CreateUnary(shape, HloOpcode::kCos, x)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); // Manually fuse the kCos operation into a fusion operation. @@ -458,7 +485,7 @@ TEST_F(ReducePrecisionInsertionTest, OpGetsInsertedInTailOfFusionNode) { builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "x")); HloInstruction* y = builder.AddInstruction( HloInstruction::CreateUnary(shape, HloOpcode::kCos, x)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); // Manually fuse the kCos operation into a fusion operation. diff --git a/tensorflow/compiler/xla/service/reshape_mover.cc b/tensorflow/compiler/xla/service/reshape_mover.cc index 4df746fca9f8320eed72911726f33bb01f06fed5..a62118df157edf67114ff41befbdce3da129fe93 100644 --- a/tensorflow/compiler/xla/service/reshape_mover.cc +++ b/tensorflow/compiler/xla/service/reshape_mover.cc @@ -226,7 +226,10 @@ StatusOr PerformSinkReshapeOrTranspose( // changes, so all the fused instructions have the same dimensions. for (const auto& fused_instruction : instruction->fused_instructions()) { Shape* shape = fused_instruction->mutable_shape(); - *shape->mutable_dimensions() = new_operand_shape.dimensions(); + shape->clear_dimensions(); + for (int64 i : new_operand_shape.dimensions()) { + shape->add_dimensions(i); + } *shape->mutable_layout() = new_operand_shape.layout(); } } diff --git a/tensorflow/compiler/xla/service/reshape_mover_test.cc b/tensorflow/compiler/xla/service/reshape_mover_test.cc index fcf269eee925c2ddb7511d70e71bd815e4b8c24a..341659b15c4c7355d39739ee171a4a749d87e929 100644 --- a/tensorflow/compiler/xla/service/reshape_mover_test.cc +++ b/tensorflow/compiler/xla/service/reshape_mover_test.cc @@ -25,7 +25,7 @@ limitations under the License. #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_verified_test_base.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" @@ -34,9 +34,10 @@ namespace { namespace op = xla::testing::opcode_matchers; -class ReshapeMoverTest : public HloVerifiedTestBase {}; +class ReshapeMoverTest : public HloTestBase {}; TEST_F(ReshapeMoverTest, ReshapesWithDifferentInputShapesNotMoved) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); auto root_shape = ShapeUtil::MakeShape(F32, {8, 7}); auto param0 = builder.AddInstruction(HloInstruction::CreateParameter( @@ -50,12 +51,12 @@ TEST_F(ReshapeMoverTest, ReshapesWithDifferentInputShapesNotMoved) { builder.AddInstruction(HloInstruction::CreateBinary( root_shape, HloOpcode::kAdd, reshape0, reshape1)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Add(op::Reshape(param0), op::Reshape(param1))); - EXPECT_FALSE(ReshapeMover().Run(&module()).ValueOrDie()); + EXPECT_FALSE(ReshapeMover().Run(m.get()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::Add(op::Reshape(param0), op::Reshape(param1))); @@ -74,6 +75,7 @@ TEST_F(ReshapeMoverTest, ReshapesWithDifferentInputShapesNotMoved) { // Verifies that the reshape is not moved, since rng0 is trivially reshapable // and therefore there is no nontrivial reshapes to move. TEST_F(ReshapeMoverTest, 1ConstantAnd1ReshapesOnRngNotMoved) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); auto root_shape = ShapeUtil::MakeShape(F32, {8, 7}); auto rng0 = builder.AddInstruction(HloInstruction::CreateRng( @@ -92,18 +94,19 @@ TEST_F(ReshapeMoverTest, 1ConstantAnd1ReshapesOnRngNotMoved) { builder.AddInstruction(HloInstruction::CreateBinary( root_shape, HloOpcode::kAdd, reshape0, const1)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Add(op::Reshape(rng0), const1)); - EXPECT_FALSE(ReshapeMover().Run(&module()).ValueOrDie()); + EXPECT_FALSE(ReshapeMover().Run(m.get()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::Add(op::Reshape(rng0), const1)); } TEST_F(ReshapeMoverTest, ScalarReshapesNotMoved) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); auto root_shape = ShapeUtil::MakeShape(F32, {}); auto param0 = builder.AddInstruction(HloInstruction::CreateParameter( @@ -117,12 +120,12 @@ TEST_F(ReshapeMoverTest, ScalarReshapesNotMoved) { builder.AddInstruction(HloInstruction::CreateBinary( root_shape, HloOpcode::kAdd, reshape0, reshape1)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Add(op::Reshape(param0), op::Reshape(param1))); - EXPECT_FALSE(ReshapeMover().Run(&module()).ValueOrDie()); + EXPECT_FALSE(ReshapeMover().Run(m.get()).ValueOrDie()); EXPECT_THAT( computation->root_instruction(), @@ -130,6 +133,7 @@ TEST_F(ReshapeMoverTest, ScalarReshapesNotMoved) { } TEST_F(ReshapeMoverTest, EquivalentReshapesMoved) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); auto root_shape = ShapeUtil::MakeShape(F32, {8, 7}); auto param0 = builder.AddInstruction(HloInstruction::CreateParameter( @@ -143,11 +147,11 @@ TEST_F(ReshapeMoverTest, EquivalentReshapesMoved) { builder.AddInstruction(HloInstruction::CreateBinary( root_shape, HloOpcode::kAdd, reshape0, reshape1)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Add(op::Reshape(param0), op::Reshape(param1))); - EXPECT_TRUE(ReshapeMover().Run(&module()).ValueOrDie()); + EXPECT_TRUE(ReshapeMover().Run(m.get()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::Reshape(op::Add(param0, param1))); @@ -177,6 +181,7 @@ TEST_F(ReshapeMoverTest, EquivalentReshapesMoved) { // | // reshape4 TEST_F(ReshapeMoverTest, 1ConstantAnd2ReshapesMoved) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); auto root_shape = ShapeUtil::MakeShape(F32, {2, 3}); auto const0 = builder.AddInstruction( @@ -196,12 +201,12 @@ TEST_F(ReshapeMoverTest, 1ConstantAnd2ReshapesMoved) { builder.AddInstruction(HloInstruction::CreateTernary( root_shape, HloOpcode::kSelect, const0, reshape1, reshape2)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Select(const0, reshape1, reshape2)); - EXPECT_TRUE(ReshapeMover().Run(&module()).ValueOrDie()); + EXPECT_TRUE(ReshapeMover().Run(m.get()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::Reshape(op::Select(op::Reshape(const0), param1, param2))); @@ -221,6 +226,7 @@ TEST_F(ReshapeMoverTest, 1ConstantAnd2ReshapesMoved) { // Verifies that the reshape0 does not sink below add, because param1 is not // trivially reshapable nor is a Reshape/Transpose. TEST_F(ReshapeMoverTest, 1ParameterAnd1ReshapeNotMoved) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); auto root_shape = ShapeUtil::MakeShape(F32, {8, 7}); auto param0 = builder.AddInstruction(HloInstruction::CreateParameter( @@ -232,11 +238,11 @@ TEST_F(ReshapeMoverTest, 1ParameterAnd1ReshapeNotMoved) { builder.AddInstruction(HloInstruction::CreateBinary( root_shape, HloOpcode::kAdd, reshape0, param1)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Add(op::Reshape(param0), param1)); - EXPECT_FALSE(ReshapeMover().Run(&module()).ValueOrDie()); + EXPECT_FALSE(ReshapeMover().Run(m.get()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::Add(op::Reshape(param0), param1)); @@ -257,6 +263,7 @@ TEST_F(ReshapeMoverTest, 1ParameterAnd1ReshapeNotMoved) { // Verifies that we don't unnecessarily sink reshapes, which are in fact // trivial reshapes. TEST_F(ReshapeMoverTest, 2TrivialConstantReshapeNotMoved) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); auto root_shape = ShapeUtil::MakeShape(F32, {3, 2}); auto const0 = builder.AddInstruction(HloInstruction::CreateConstant( @@ -275,12 +282,12 @@ TEST_F(ReshapeMoverTest, 2TrivialConstantReshapeNotMoved) { builder.AddInstruction(HloInstruction::CreateTernary( root_shape, HloOpcode::kSelect, pred, reshape0, reshape1)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Select(pred, op::Reshape(const0), op::Reshape(const1))); - EXPECT_FALSE(ReshapeMover().Run(&module()).ValueOrDie()); + EXPECT_FALSE(ReshapeMover().Run(m.get()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::Select(pred, op::Reshape(const0), op::Reshape(const1))); @@ -309,6 +316,7 @@ TEST_F(ReshapeMoverTest, 2TrivialConstantReshapeNotMoved) { // // (note that reshape1 here is trivial). TEST_F(ReshapeMoverTest, 1NonTrivialReshapeMoved) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); auto root_shape = ShapeUtil::MakeShape(F32, {2, 3}); auto param0 = builder.AddInstruction(HloInstruction::CreateParameter( @@ -320,12 +328,12 @@ TEST_F(ReshapeMoverTest, 1NonTrivialReshapeMoved) { builder.AddInstruction(HloInstruction::CreateBinary( root_shape, HloOpcode::kAdd, reshape0, const1)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Add(op::Reshape(param0), const1)); - EXPECT_TRUE(ReshapeMover().Run(&module()).ValueOrDie()); + EXPECT_TRUE(ReshapeMover().Run(m.get()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::Reshape(op::Add(param0, op::Reshape(const1)))); @@ -348,6 +356,7 @@ TEST_F(ReshapeMoverTest, 1NonTrivialReshapeMoved) { // For now we treat it as non-trivial, so we verify that we don't sink the // reshapes in this case. TEST_F(ReshapeMoverTest, 1NonTrivialReshapeWith1ReshapedConstNotMoved) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); auto root_shape = ShapeUtil::MakeShape(F32, {1, 1, 3}); auto param0 = builder.AddInstruction(HloInstruction::CreateParameter( @@ -362,12 +371,12 @@ TEST_F(ReshapeMoverTest, 1NonTrivialReshapeWith1ReshapedConstNotMoved) { builder.AddInstruction(HloInstruction::CreateBinary( root_shape, HloOpcode::kAdd, reshape0, reshape1)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Add(op::Reshape(param0), op::Reshape(const1))); - EXPECT_FALSE(ReshapeMover().Run(&module()).ValueOrDie()); + EXPECT_FALSE(ReshapeMover().Run(m.get()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::Add(op::Reshape(param0), op::Reshape(const1))); @@ -376,6 +385,7 @@ TEST_F(ReshapeMoverTest, 1NonTrivialReshapeWith1ReshapedConstNotMoved) { } TEST_F(ReshapeMoverTest, EquivalentReshapesMovedAcrossFusion) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); auto root_shape = ShapeUtil::MakeShape(F32, {8, 7}); auto param0 = builder.AddInstruction(HloInstruction::CreateParameter( @@ -389,14 +399,14 @@ TEST_F(ReshapeMoverTest, EquivalentReshapesMovedAcrossFusion) { auto add = builder.AddInstruction(HloInstruction::CreateBinary( root_shape, HloOpcode::kAdd, reshape0, reshape1)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); computation->CreateFusionInstruction({add}, HloInstruction::FusionKind::kLoop); EXPECT_THAT(computation->root_instruction(), op::Fusion(op::Reshape(param0), op::Reshape(param1))); - EXPECT_TRUE(ReshapeMover().Run(&module()).ValueOrDie()); + EXPECT_TRUE(ReshapeMover().Run(m.get()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::Reshape(op::Fusion(param0, param1))); @@ -405,6 +415,7 @@ TEST_F(ReshapeMoverTest, EquivalentReshapesMovedAcrossFusion) { } TEST_F(ReshapeMoverTest, EquivalentReshapesMovedAcrossSelect) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); auto root_shape = ShapeUtil::MakeShape(F32, {8, 7}); auto pred_shape = ShapeUtil::MakeShape(PRED, {8, 7}); @@ -423,13 +434,13 @@ TEST_F(ReshapeMoverTest, EquivalentReshapesMovedAcrossSelect) { builder.AddInstruction(HloInstruction::CreateTernary( root_shape, HloOpcode::kSelect, reshape_pred, reshape0, reshape1)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); EXPECT_THAT( computation->root_instruction(), op::Select(op::Reshape(pred), op::Reshape(param0), op::Reshape(param1))); - EXPECT_TRUE(ReshapeMover().Run(&module()).ValueOrDie()); + EXPECT_TRUE(ReshapeMover().Run(m.get()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::Reshape(op::Select(pred, param0, param1))); @@ -438,6 +449,7 @@ TEST_F(ReshapeMoverTest, EquivalentReshapesMovedAcrossSelect) { } TEST_F(ReshapeMoverTest, ScalarReshapeNotMovedAcrossSelect) { + auto m = CreateNewVerifiedModule(); HloComputation::Builder builder(TestName()); auto root_shape = ShapeUtil::MakeShape(F32, {}); auto pred_shape = ShapeUtil::MakeShape(PRED, {}); @@ -452,11 +464,11 @@ TEST_F(ReshapeMoverTest, ScalarReshapeNotMovedAcrossSelect) { auto select = builder.AddInstruction(HloInstruction::CreateTernary( root_shape, HloOpcode::kSelect, reshape_pred, param0, param1)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Select(op::Reshape(pred), param0, param1)); - EXPECT_FALSE(ReshapeMover().Run(&module()).ValueOrDie()); + EXPECT_FALSE(ReshapeMover().Run(m.get()).ValueOrDie()); EXPECT_THAT(computation->root_instruction(), op::Select(op::Reshape(pred), param0, param1)); @@ -477,6 +489,7 @@ TEST_F(ReshapeMoverTest, ScalarReshapeNotMovedAcrossSelect) { // // We expect reshape{0,1} AND reshape{2,3} to be lifted. TEST_F(ReshapeMoverTest, MultiplePasses) { + auto m = CreateNewVerifiedModule(); auto shape1 = ShapeUtil::MakeShape(F32, {1, 8, 1, 7}); auto shape2 = ShapeUtil::MakeShape(F32, {8, 7, 1}); auto shape3 = ShapeUtil::MakeShape(F32, {8, 7}); @@ -500,14 +513,14 @@ TEST_F(ReshapeMoverTest, MultiplePasses) { builder.AddInstruction(HloInstruction::CreateBinary(shape3, HloOpcode::kAdd, reshape2, reshape3)); - auto computation = module().AddEntryComputation(builder.Build()); + auto computation = m->AddEntryComputation(builder.Build()); EXPECT_THAT( computation->root_instruction(), op::Add(op::Reshape(param2), op::Reshape(op::Add(op::Reshape(param0), op::Reshape(param1))))); - EXPECT_TRUE(ReshapeMover().Run(&module()).ValueOrDie()); + EXPECT_TRUE(ReshapeMover().Run(m.get()).ValueOrDie()); EXPECT_THAT( computation->root_instruction(), @@ -526,11 +539,11 @@ TEST_F(ReshapeMoverTest, SinkTransposeAcrossBroadcastScalar) { } )"; - ParseAndVerifyModule(hlo_string); - TF_ASSERT_OK_AND_ASSIGN(bool changed, ReshapeMover().Run(&module())); + TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(hlo_string)); + TF_ASSERT_OK_AND_ASSIGN(bool changed, ReshapeMover().Run(m.get())); EXPECT_TRUE(changed); - EXPECT_THAT(module().entry_computation()->root_instruction(), + EXPECT_THAT(m->entry_computation()->root_instruction(), op::Transpose(op::Multiply())); } @@ -555,8 +568,8 @@ TEST_F(ReshapeMoverTest, ReshapeWithUsersOutsideCandidatesNotSink) { } )"; - ParseAndVerifyModule(hlo_string); - TF_ASSERT_OK_AND_ASSIGN(bool changed, ReshapeMover().Run(&module())); + TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(hlo_string)); + TF_ASSERT_OK_AND_ASSIGN(bool changed, ReshapeMover().Run(m.get())); EXPECT_FALSE(changed); } @@ -580,10 +593,10 @@ TEST_F(ReshapeMoverTest, ReshapeNoUsersOutsideCandidatesSink1) { } )"; - ParseAndVerifyModule(hlo_string); - TF_ASSERT_OK_AND_ASSIGN(bool changed, ReshapeMover().Run(&module())); + TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(hlo_string)); + TF_ASSERT_OK_AND_ASSIGN(bool changed, ReshapeMover().Run(m.get())); EXPECT_TRUE(changed); - EXPECT_THAT(module().entry_computation()->root_instruction(), + EXPECT_THAT(m->entry_computation()->root_instruction(), op::Tuple(op::Reshape(), op::Reshape(), op::Reshape())); } @@ -597,10 +610,10 @@ TEST_F(ReshapeMoverTest, ReshapeNoUsersOutsideCandidatesSink2) { } )"; - ParseAndVerifyModule(hlo_string); - TF_ASSERT_OK_AND_ASSIGN(bool changed, ReshapeMover().Run(&module())); + TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(hlo_string)); + TF_ASSERT_OK_AND_ASSIGN(bool changed, ReshapeMover().Run(m.get())); EXPECT_TRUE(changed); - EXPECT_THAT(module().entry_computation()->root_instruction(), + EXPECT_THAT(m->entry_computation()->root_instruction(), op::Reshape(op::Add())); } diff --git a/tensorflow/compiler/xla/service/scatter_expander.cc b/tensorflow/compiler/xla/service/scatter_expander.cc index de7aee262e61195b37099fc661a95508d0539e18..acad871c4d427b174ffce3a462a0a3918a1e0c33 100644 --- a/tensorflow/compiler/xla/service/scatter_expander.cc +++ b/tensorflow/compiler/xla/service/scatter_expander.cc @@ -26,7 +26,6 @@ limitations under the License. namespace xla { - // Transposes the given scatter_indices such that the index_vector_dim becomes // the most-minor dimension. static StatusOr TransposeIndexVectorDimToLast( @@ -60,6 +59,13 @@ static StatusOr CanonicalizeScatterIndices( TF_ASSIGN_OR_RETURN( HloInstruction * transposed_scatter_indices, TransposeIndexVectorDimToLast(scatter_indices, index_vector_dim)); + 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()); + TF_ASSIGN_OR_RETURN(scatter_indices, + MakeReshapeHlo(new_shape, scatter_indices)); + } bool indices_are_scalar = index_vector_dim == scatter_indices->shape().dimensions_size(); @@ -88,7 +94,7 @@ static StatusOr CanonicalizeScatterIndices( static StatusOr PermuteScatterAndWindowDims( HloInstruction* updates, absl::Span update_window_dims) { std::vector permutation; - const int64 updates_rank = ShapeUtil::Rank(updates->shape()); + const int64 updates_rank = updates->shape().rank(); permutation.reserve(updates_rank); for (int64 i = 0; i < updates_rank; ++i) { @@ -165,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)); @@ -214,15 +219,11 @@ static StatusOr> ScatterLoopBody( HloInstruction* updates = loop_state[2]; bool has_scalar_indices = scatter_indices->shape().dimensions_size() == 1; - CHECK_EQ(has_scalar_indices, - dim_numbers.index_vector_dim() == - scatter->operand(1)->shape().dimensions_size()); // 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. @@ -392,7 +393,8 @@ StatusOr ScatterExpander::ExpandScatter( [&](HloInstruction* induction_var, const std::vector& loop_state) { return ScatterLoopBody(scatter, induction_var, loop_state); - }); + }, + scatter->metadata()); TF_ASSIGN_OR_RETURN(std::vector scatter_loop_result, scatter_loop_result_status); return scatter_loop_result.front(); diff --git a/tensorflow/compiler/xla/service/scatter_expander.h b/tensorflow/compiler/xla/service/scatter_expander.h index 559a85dccfef27816e7dbf746fd71c44bbf46f60..533af060bc9f943e5bc2882db626e25c77484029 100644 --- a/tensorflow/compiler/xla/service/scatter_expander.h +++ b/tensorflow/compiler/xla/service/scatter_expander.h @@ -25,7 +25,7 @@ class ScatterExpander : public HloModulePass { absl::string_view name() const override { return "scatter_expander"; } StatusOr Run(HloModule* module) override; - private: + protected: StatusOr ExpandScatter(HloInstruction* scatter); }; diff --git a/tensorflow/compiler/xla/service/service.cc b/tensorflow/compiler/xla/service/service.cc index 75465359f8f37e56369c0976ba7434e3c3f202cc..83434528a21b16cad7c831e7d9cc42d436634540 100644 --- a/tensorflow/compiler/xla/service/service.cc +++ b/tensorflow/compiler/xla/service/service.cc @@ -23,12 +23,13 @@ limitations under the License. #include "absl/memory/memory.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_format.h" +#include "tensorflow/compiler/xla/debug_options_flags.h" #include "tensorflow/compiler/xla/execution_options_util.h" #include "tensorflow/compiler/xla/layout_util.h" -#include "tensorflow/compiler/xla/legacy_flags/debug_options_flags.h" #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" @@ -41,6 +42,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/source_map_util.h" #include "tensorflow/compiler/xla/service/stream_pool.h" #include "tensorflow/compiler/xla/service/transfer_manager.h" +#include "tensorflow/compiler/xla/shape.h" #include "tensorflow/compiler/xla/shape_layout.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/status_macros.h" @@ -112,6 +114,16 @@ int ServiceOptions::intra_op_parallelism_threads() const { return intra_op_parallelism_threads_; } +ServiceOptions& ServiceOptions::set_allowed_devices( + const absl::optional>& allowed_devices) { + allowed_devices_ = allowed_devices; + return *this; +} + +const absl::optional>& ServiceOptions::allowed_devices() const { + return allowed_devices_; +} + /* static */ StatusOr> Service::NewService( se::Platform* platform) { ServiceOptions default_options; @@ -128,6 +140,7 @@ int ServiceOptions::intra_op_parallelism_threads() const { } BackendOptions backend_options; backend_options.set_platform(platform); + backend_options.set_allowed_devices(options.allowed_devices()); TF_ASSIGN_OR_RETURN(execute_backend, Backend::CreateBackend(backend_options)); std::unique_ptr service( @@ -149,17 +162,13 @@ Service::Service(const ServiceOptions& options, LOG(INFO) << StrFormat( "XLA service %p executing computations on platform %s. Devices:", this, execute_backend_->platform()->Name()); + auto stream_executors = execute_backend_->stream_executors(); for (int i = 0; i < execute_backend_->device_count(); ++i) { - if (execute_backend_->device_ordinal_supported(i)) { - se::StreamExecutor* executor = - execute_backend_->stream_executor(i).ValueOrDie(); - const auto& description = executor->GetDeviceDescription(); - LOG(INFO) << StrFormat(" StreamExecutor device (%d): %s, %s", i, - description.name(), - description.platform_version()); - } else { - LOG(INFO) << StrFormat(" StreamExecutor device (%d) not supported", i); - } + se::StreamExecutor* executor = stream_executors.at(i); + const auto& description = executor->GetDeviceDescription(); + LOG(INFO) << StrFormat(" StreamExecutor device (%d): %s, %s", i, + description.name(), + description.platform_version()); } } else { VLOG(1) << "XLA compile-only service constructed"; @@ -175,7 +184,14 @@ Status Service::CreateChannelHandle(const CreateChannelHandleRequest* arg, Status Service::Unregister(const UnregisterRequest* arg, UnregisterResponse* result) { - return allocation_tracker_.Unregister(arg->data()); + Status status; + for (auto& data : arg->data()) { + Status unregister_status = allocation_tracker_.Unregister(data); + if (!unregister_status.ok() && status.ok()) { + status = unregister_status; + } + } + return status; } // Deconstructs a previously-allocated global handle. @@ -268,8 +284,8 @@ StatusOr> Service::CreateModuleConfig( } if (execution_options != nullptr && execution_options->has_shape_with_output_layout()) { - const auto& shape_with_output_layout = - execution_options->shape_with_output_layout(); + const Shape shape_with_output_layout( + execution_options->shape_with_output_layout()); TF_RETURN_IF_ERROR( ValidateResultShape(shape_with_output_layout, program_shape.result())); TF_RETURN_IF_ERROR( @@ -280,12 +296,17 @@ 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_debug_options(legacy_flags::GetDebugOptionsFromFlags()); + config->set_replica_count(options_.number_of_replicas()); + config->set_debug_options(GetDebugOptionsFromFlags()); } if (execute_backend_ != nullptr && @@ -508,13 +529,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, @@ -522,10 +543,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; @@ -537,9 +559,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) { @@ -651,9 +671,9 @@ Status Service::ExecuteGraphParallel(const ExecuteGraphParallelRequest* arg, // replica 0. TF_ASSIGN_OR_RETURN( std::unique_ptr module_config, - CreateModuleConfig(request.computation().host_program_shape(), - replicated_arguments.front(), - request.execution_options())); + CreateModuleConfig( + ProgramShape{request.computation().host_program_shape()}, + replicated_arguments.front(), request.execution_options())); VLOG(3) << "ExecuteGraphParallel created HloModuleConfig computation layout: " << module_config->entry_computation_layout().ToString(); @@ -696,14 +716,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; @@ -738,9 +777,9 @@ Status Service::GetDeviceHandles(const GetDeviceHandlesRequest* arg, } if (available_device_count < arg->device_count() * replica_count) { return ResourceExhausted( - "Requested device count (%d) exceeds the number of available devices " - "on the target (%d)", - arg->device_count(), available_device_count); + "Requested logical device count (%d) with replica count (%d) exceeds " + "the number of available physical devices on the target (%d)", + arg->device_count(), replica_count, available_device_count); } for (int64 i = 0; i < arg->device_count(); ++i) { @@ -753,38 +792,6 @@ Status Service::GetDeviceHandles(const GetDeviceHandlesRequest* arg, return Status::OK(); } -Status Service::ExecuteOneToN(const ExecuteGraphRequest* arg, - ExecuteResponse* result) { - ExecuteGraphParallelRequest parallel_arg; - *parallel_arg.add_requests() = *arg; - ExecuteParallelResponse parallel_result; - TF_RETURN_IF_ERROR(ExecuteGraphParallel(¶llel_arg, ¶llel_result)); - return PickParallelResponse(parallel_result, result); -} - -Status Service::PickParallelResponse( - const ExecuteParallelResponse& parallel_result, ExecuteResponse* result) { - // The "result device" selection is a bit hacky, but better than assuming it - // is device 0. We have b/76035356 for restructuring the client API to clean - // up the current asymmetries and support more functionalities. - for (int64 i = 0; i < parallel_result.responses_size(); ++i) { - TF_ASSIGN_OR_RETURN(const ShapedBuffer* buffer, - allocation_tracker_.ResolveForReplica( - parallel_result.responses(i).output(), 0)); - const Shape& shape = buffer->on_host_shape(); - if (!ShapeUtil::IsEmptyTuple(shape)) { - *result = parallel_result.responses(i); - VLOG(3) << "Fetching result from device " << i << ": " - << ShapeUtil::HumanString(shape); - return Status::OK(); - } - } - TF_RET_CHECK(parallel_result.responses_size() > 0); - *result = parallel_result.responses(0); - VLOG(1) << "Defaulting to device 0 result"; - return Status::OK(); -} - StatusOr> Service::BuildExecutable( const HloModuleProto& module_proto, std::unique_ptr module_config, Backend* backend, @@ -829,10 +836,8 @@ StatusOr> Service::BuildExecutable( return std::move(executable); } -Status Service::ExecuteGraph(const ExecuteGraphRequest* arg, - ExecuteResponse* result) { - VLOG(1) << "running execute-graph request"; - +Status Service::Compile(const CompileRequest* arg, CompileResponse* result) { + VLOG(1) << "running compile request"; if (!arg->has_computation()) { return InvalidArgument("computations may not be empty"); } @@ -840,22 +845,24 @@ Status Service::ExecuteGraph(const ExecuteGraphRequest* arg, return InvalidArgument("programe shape may not be empty"); } - // If we received multiple device handles, we must partition the module. if (arg->execution_options().device_handles_size() > 1) { - return ExecuteOneToN(arg, result); + return InvalidArgument( + "The compile request does not support multiple device handles."); } - TF_ASSIGN_OR_RETURN(auto replicas, Replicas(*execute_backend_, - SingleComputationDeviceHandle())); - TF_ASSIGN_OR_RETURN( - std::vector> replicated_arguments, - ResolveAndValidateArguments(arg->arguments(), replicas)); - + std::vector argument_shapes; + argument_shapes.reserve(arg->input_shape_with_layout_size()); + std::vector argument_shape_ptrs; + for (const ShapeProto& shape_proto : arg->input_shape_with_layout()) { + argument_shapes.push_back(Shape(shape_proto)); + argument_shape_ptrs.push_back(&argument_shapes.back()); + } TF_ASSIGN_OR_RETURN( std::unique_ptr module_config, - CreateModuleConfig(arg->computation().host_program_shape(), - replicated_arguments.front(), - arg->execution_options())); + CreateModuleConfig(ProgramShape{arg->computation().host_program_shape()}, + argument_shape_ptrs, &arg->execution_options())); + VLOG(3) << "Compile created HloModuleConfig computation layout: " + << module_config->entry_computation_layout().ToString(); TF_ASSIGN_OR_RETURN( std::unique_ptr executable, @@ -864,6 +871,48 @@ Status Service::ExecuteGraph(const ExecuteGraphRequest* arg, execute_backend_->default_stream_executor(), /*device_allocator=*/nullptr)); + *result->mutable_handle() = compilation_cache_.Insert(std::move(executable)); + + VLOG(1) << "successfully completed 'compile' request"; + return Status::OK(); +} + +Status Service::Execute(const ExecuteRequest* arg, ExecuteResponse* result) { + VLOG(1) << "running execute request"; + if (!arg->has_handle()) { + return InvalidArgument("execution handle should not be empty"); + } + TF_ASSIGN_OR_RETURN(auto executable, + compilation_cache_.LookUp(arg->handle())); + + TF_ASSIGN_OR_RETURN(auto replicas, Replicas(*execute_backend_, + SingleComputationDeviceHandle())); + TF_ASSIGN_OR_RETURN( + std::vector> replicated_arguments, + ResolveAndValidateArguments(arg->arguments(), replicas)); + + // Check that the replicated_arguments has the same shape and layout as the + // module config used when creating the exectuable. + const int64 num_module_args = + executable->module_config().entry_computation_layout().parameter_count(); + if (num_module_args != arg->arguments_size()) { + return InvalidArgument( + "The executable expects %lld arguments, but sees %lld.", + num_module_args, arg->arguments_size()); + } + for (int64 i = 0; i < num_module_args; i++) { + const Shape& shape_module = + executable->module_config().entry_computation_layout().parameter_shape( + i); + const Shape& shape_arg = replicated_arguments.front()[i]->on_host_shape(); + if (!ShapeUtil::Equal(shape_module, shape_arg)) { + return InvalidArgumentStrCat( + "The executable exepcts the ", i, "th argument in shape ", + ShapeUtil::HumanStringWithLayout(shape_module), " but sees ", + ShapeUtil::HumanStringWithLayout(shape_arg)); + } + } + TF_ASSIGN_OR_RETURN(auto stream, execute_backend_->BorrowStream( execute_backend_->default_stream_executor())); @@ -877,9 +926,11 @@ Status Service::ExecuteGraph(const ExecuteGraphRequest* arg, TF_ASSIGN_OR_RETURN( *result->mutable_output(), - ExecuteAndRegisterResult( - executable.get(), replicated_arguments, execute_backend_.get(), - "result of " + arg->computation().name(), result->mutable_profile())); + ExecuteAndRegisterResult(executable.get(), replicated_arguments, + execute_backend_.get(), + SingleComputationDeviceHandle(), + "result of " + executable->module().name(), + result->mutable_profile())); if (executable->dumping_snapshot()) { TF_ASSIGN_OR_RETURN( @@ -891,7 +942,7 @@ Status Service::ExecuteGraph(const ExecuteGraphRequest* arg, TF_RETURN_IF_ERROR(executable->DumpHloSnapshot()); } - VLOG(1) << "successfully completed 'execute-graph' request"; + VLOG(1) << "successfully completed 'execute' request"; return Status::OK(); } @@ -915,14 +966,14 @@ Status Service::TransferToClient(const TransferToClientRequest* arg, TF_ASSIGN_OR_RETURN(const ShapedBuffer* shaped_buffer, allocation_tracker_.ResolveForReplica(arg->data(), 0)); - const Shape* return_shape; + Shape return_shape; if (arg->has_shape_with_layout()) { - if (!LayoutUtil::HasLayout(arg->shape_with_layout())) { + return_shape = Shape(arg->shape_with_layout()); + if (!LayoutUtil::HasLayout(return_shape)) { return InvalidArgument("shape_with_layout must have layout if present."); } - return_shape = &arg->shape_with_layout(); } else { - return_shape = &shaped_buffer->on_host_shape(); + return_shape = Shape(shaped_buffer->on_host_shape()); } TF_ASSIGN_OR_RETURN(auto stream, execute_backend_->BorrowStream( @@ -933,30 +984,15 @@ Status Service::TransferToClient(const TransferToClientRequest* arg, execute_backend_->transfer_manager()->TransferLiteralFromDevice( stream.get(), *shaped_buffer)); - if (LayoutUtil::LayoutsInShapesEqual(*return_shape, result_literal.shape())) { + if (LayoutUtil::LayoutsInShapesEqual(return_shape, result_literal.shape())) { *result->mutable_literal() = result_literal.ToProto(); } else { *result->mutable_literal() = - result_literal.Relayout(*return_shape).ToProto(); + result_literal.Relayout(return_shape).ToProto(); } return Status::OK(); } -namespace { - -// Creates a clone of the given shaped buffer with the given device ordinal. The -// shape and DeviceMemoryBase values of the clone are identical to the original. -std::unique_ptr CloneShapedBufferOnDevice( - const ShapedBuffer& shaped_buffer, int device_ordinal) { - auto clone = absl::make_unique( - shaped_buffer.on_host_shape(), shaped_buffer.on_device_shape(), - shaped_buffer.platform(), device_ordinal); - clone->buffers() = shaped_buffer.buffers(); - return clone; -} - -} // namespace - Status Service::TransferToServer(const TransferToServerRequest* arg, TransferToServerResponse* result) { TF_ASSIGN_OR_RETURN(Literal literal, @@ -1045,11 +1081,11 @@ Status Service::TransferFromOutfeed(const TransferFromOutfeedRequest* arg, executor = replicas[arg->replica_id()]; } - auto literal = Literal::CreateFromShape(arg->shape_with_layout()); + auto literal = Literal::CreateFromShape(Shape(arg->shape_with_layout())); TF_RETURN_IF_ERROR( execute_backend_->transfer_manager()->TransferLiteralFromOutfeed( - executor, arg->shape_with_layout(), literal)); + executor, Shape(arg->shape_with_layout()), literal)); *result->mutable_literal() = literal.ToProto(); return Status::OK(); } @@ -1072,11 +1108,13 @@ Status Service::ComputeConstantGraph(const ComputeConstantGraphRequest* arg, "constant computation may not depend on any parameters."); } - ProgramShape program_shape = arg->computation().host_program_shape(); + ProgramShape program_shape(arg->computation().host_program_shape()); TF_DCHECK_OK(ShapeUtil::ValidateShape(program_shape.result())); + absl::optional output_layout; if (arg->has_output_layout()) { + output_layout = Layout::CreateFromProto(arg->output_layout()); TF_RETURN_IF_ERROR(LayoutUtil::ValidateLayoutForShape( - arg->output_layout(), program_shape.result())); + *output_layout, program_shape.result())); } HloModuleConfig config(program_shape); @@ -1084,16 +1122,19 @@ 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; - TF_ASSIGN_OR_RETURN(auto result_literal, evaluator.Evaluate( - *module, /*arg_literals=*/{})); + 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 // relayout here. // // TODO(b/77824332): Make HloEvaluator take care of the re-layout. - if (arg->has_output_layout()) { - result_literal = result_literal.Relayout(arg->output_layout()); + if (output_layout.has_value()) { + result_literal = result_literal.Relayout(*output_layout); } *result->mutable_literal() = result_literal.ToProto(); @@ -1103,7 +1144,7 @@ Status Service::ComputeConstantGraph(const ComputeConstantGraphRequest* arg, Status Service::GetShape(const GetShapeRequest* arg, GetShapeResponse* result) { TF_ASSIGN_OR_RETURN(const ShapedBuffer* buffer, allocation_tracker_.ResolveForReplica(arg->data(), 0)); - *result->mutable_shape() = buffer->on_host_shape(); + *result->mutable_shape() = buffer->on_host_shape().ToProto(); return Status::OK(); } @@ -1116,7 +1157,7 @@ Status Service::GetComputationGraphStats( return InvalidArgument("Program shape may not be empty."); } - HloModuleConfig config(arg->computation().host_program_shape()); + HloModuleConfig config(ProgramShape{arg->computation().host_program_shape()}); config.set_debug_options(arg->debug_options()); TF_ASSIGN_OR_RETURN(std::unique_ptr module, CreateModuleFromProto(arg->computation(), config)); diff --git a/tensorflow/compiler/xla/service/service.h b/tensorflow/compiler/xla/service/service.h index 8cf1a7b9f01fbb3572c6849c8b18e14174ced89f..fd907d07daef9e8337aeed198ef4fd23d069df21 100644 --- a/tensorflow/compiler/xla/service/service.h +++ b/tensorflow/compiler/xla/service/service.h @@ -18,15 +18,17 @@ limitations under the License. #include #include +#include #include #include #include "absl/types/span.h" +#include "tensorflow/compiler/xla/debug_options_flags.h" #include "tensorflow/compiler/xla/executable_run_options.h" -#include "tensorflow/compiler/xla/legacy_flags/debug_options_flags.h" #include "tensorflow/compiler/xla/service/allocation_tracker.h" #include "tensorflow/compiler/xla/service/backend.h" #include "tensorflow/compiler/xla/service/channel_tracker.h" +#include "tensorflow/compiler/xla/service/compilation_cache.h" #include "tensorflow/compiler/xla/service/device_memory_allocator.h" #include "tensorflow/compiler/xla/service/executable.h" #include "tensorflow/compiler/xla/service/execution_tracker.h" @@ -51,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; @@ -60,10 +62,17 @@ class ServiceOptions { ServiceOptions& set_intra_op_parallelism_threads(int num_threads); int intra_op_parallelism_threads() const; + // Sets the allowed_devices set for selectively constructing stream executors + // on the platform. + ServiceOptions& set_allowed_devices( + const absl::optional>& allowed_devices); + const absl::optional>& allowed_devices() const; + private: se::Platform* platform_ = nullptr; int number_of_replicas_ = 1; int intra_op_parallelism_threads_ = -1; + absl::optional> allowed_devices_; }; // The XLA service object, which is the same across all platforms. It maintains @@ -90,11 +99,14 @@ class Service : public ServiceInterface { Status DeconstructTuple(const DeconstructTupleRequest* arg, DeconstructTupleResponse* result) override; - // Executes a computation with the provided global data passed as - // immutable arguments. The request contains the whole computation graph. - // Returns global data output and execution timing. - Status ExecuteGraph(const ExecuteGraphRequest* arg, - ExecuteResponse* result) override; + // Compiles a computation into an executable. The request contains the whole + // computation graph. Returns the handle to the executable. + Status Compile(const CompileRequest* arg, CompileResponse* result) override; + + // Executes an executable with the provided global data passes as immutable + // arguments. The request contains the handle to the executable. Returns + // global data output and execution timing. + Status Execute(const ExecuteRequest* arg, ExecuteResponse* result) override; // Executes one or more computations in parallel with the provided global data // passed as immutable arguments. Returns global data output for each @@ -179,10 +191,6 @@ class Service : public ServiceInterface { absl::Span arguments, const ExecutionOptions& execution_options); - // Picks a parallel response and fills the result. - Status PickParallelResponse(const ExecuteParallelResponse& parallel_result, - ExecuteResponse* result); - // Prepare the executors for executing parallel. StatusOr> GetExecutors( const ExecutionOptions& execution_options, int64 requests_size, @@ -242,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 @@ -254,11 +263,6 @@ class Service : public ServiceInterface { Backend* backend, absl::Span device_handles, absl::Span result_tags, ExecutionProfile* profile); - // Executes a single computation which has more than one target device. - // The N devices are expected to all return an empty tuple, but one, which - // will be the result of this computation. - Status ExecuteOneToN(const ExecuteGraphRequest* arg, ExecuteResponse* result); - // Convenience function which checks whether the given client_shape // (presumably passed by the client to set the result layout) is valid for the // given computation result shape. @@ -281,6 +285,9 @@ class Service : public ServiceInterface { ServiceOptions options_; + // Cache containing previously built Executables. + CompilationCache compilation_cache_; + // Tracks channels created via the API. ChannelTracker channel_tracker_; 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 2e02b256cb84f6504bfcf73832c24b7003c7fb60..946577d55d43f04fe2dbabb3dd11c3468f2c7edf 100644 --- a/tensorflow/compiler/xla/service/shape_inference.cc +++ b/tensorflow/compiler/xla/service/shape_inference.cc @@ -15,8 +15,8 @@ limitations under the License. #include "tensorflow/compiler/xla/service/shape_inference.h" -#include #include +#include #include #include #include @@ -50,7 +50,7 @@ bool AllUnique(absl::Span slice) { } Status ExpectArray(const Shape& shape, absl::string_view op_type) { - if (!ShapeUtil::IsArray(shape)) { + if (!shape.IsArray()) { return InvalidArgument("Expected array argument for %s, but got %s.", string(op_type), ShapeUtil::HumanString(shape)); } @@ -70,7 +70,7 @@ Status VerifyReducerShape(const ProgramShape& reducer_shape, const Shape& accumulator_shape = reducer_shape.result(); std::vector accumulator_subshapes; - if (ShapeUtil::IsArray(accumulator_shape)) { + if (accumulator_shape.IsArray()) { if (inputs != 1) { return InvalidArgument( "Reduction function must produce a tuple with %d elements, but " @@ -78,7 +78,7 @@ Status VerifyReducerShape(const ProgramShape& reducer_shape, inputs); } accumulator_subshapes.push_back(&accumulator_shape); - } else if (ShapeUtil::IsTuple(accumulator_shape)) { + } else if (accumulator_shape.IsTuple()) { if (ShapeUtil::TupleElementCount(accumulator_shape) != inputs) { return InvalidArgument( "Reduction function must produce a tuple with %d elements, but has " @@ -96,7 +96,7 @@ Status VerifyReducerShape(const ProgramShape& reducer_shape, } for (const Shape* element_shape : accumulator_subshapes) { - if (ShapeUtil::Rank(*element_shape) != 0) { + if (element_shape->rank() != 0) { return InvalidArgument( "Reduction function must return a scalar or tuple of scalars but " "returns shape: %s", @@ -156,17 +156,26 @@ 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, bool allow_negative_padding) { - if (window.dimensions_size() != ShapeUtil::Rank(base_shape)) { + if (window.dimensions_size() != base_shape.rank()) { return InvalidArgument( "Window has dimension %d but base shape has dimension %d.", - window.dimensions_size(), ShapeUtil::Rank(base_shape)); + window.dimensions_size(), base_shape.rank()); } 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 @@ -338,7 +355,7 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, if (arg_shapes.empty()) { return InvalidArgument("Concatenate expects at least one argument."); } - if (dimension < 0 || dimension >= ShapeUtil::Rank(*arg_shapes[0])) { + if (dimension < 0 || dimension >= arg_shapes[0]->rank()) { return InvalidArgument("Concatenate dimension out of bounds: %d.", dimension); } @@ -351,12 +368,12 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, element_type = arg_shape->element_type(); continue; } - if (ShapeUtil::Rank(*arg_shape) != ShapeUtil::Rank(*shape)) { + if (arg_shape->rank() != shape->rank()) { return InvalidArgument( "Cannot concatenate arrays with different ranks: %d (%s) vs %d " "(%s).", - ShapeUtil::Rank(*arg_shape), ShapeUtil::HumanString(*arg_shape), - ShapeUtil::Rank(*shape), ShapeUtil::HumanString(*shape)); + arg_shape->rank(), ShapeUtil::HumanString(*arg_shape), shape->rank(), + ShapeUtil::HumanString(*shape)); } if (!ShapeUtil::SameElementTypeIgnoringFpPrecision(*arg_shape, *shape)) { return InvalidArgument( @@ -364,8 +381,8 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, PrimitiveType_Name(arg_shape->element_type()), PrimitiveType_Name(shape->element_type())); } - for (int64 dimension_number = 0; - dimension_number < ShapeUtil::Rank(*arg_shape); ++dimension_number) { + for (int64 dimension_number = 0; dimension_number < arg_shape->rank(); + ++dimension_number) { if (arg_shape->dimensions(dimension_number) != shape->dimensions(dimension_number)) { if (dimension_number == dimension) { @@ -391,17 +408,6 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, return ShapeUtil::MakeShape(element_type, new_dimensions); } -/* static */ StatusOr ShapeInference::InferAfterAllShape( - absl::Span arg_shapes) { - for (const Shape* arg_shape : arg_shapes) { - if (arg_shape->element_type() != TOKEN) { - return InvalidArgument( - "Operands of token instructions must be TOKEN types."); - } - } - return ShapeUtil::MakeTokenShape(); -} - /* static */ StatusOr ShapeInference::InferConvertShape( const Shape& operand_shape, PrimitiveType new_element_type) { auto old_element_type = operand_shape.element_type(); @@ -412,7 +418,7 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, ShapeUtil::HumanString(operand_shape), PrimitiveType_Name(new_element_type)); } - if (!ShapeUtil::IsArray(operand_shape) || + if (!operand_shape.IsArray() || !primitive_util::IsArrayType(new_element_type)) { // Note: we may want to support tuple conversions via this operation in the // future, by recursing into the tuple elements to check all sub-conversions @@ -435,7 +441,7 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, ShapeUtil::HumanString(operand_shape), PrimitiveType_Name(new_element_type)); } - if (!ShapeUtil::IsArray(operand_shape) || + if (!operand_shape.IsArray() || !primitive_util::IsArrayType(new_element_type)) { // Note: we may want to support tuple conversions via this operation in the // future, by recursing into the tuple elements to check all sub-conversions @@ -483,7 +489,7 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, /* static */ StatusOr ShapeInference::InferPadShape( const Shape& operand_shape, const Shape& padding_value_shape, const PaddingConfig& padding_config) { - if (!ShapeUtil::IsArray(operand_shape)) { + if (!operand_shape.IsArray()) { return InvalidArgument( "Pad operation does not support tuple-shape operands."); } @@ -491,7 +497,7 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, return InvalidArgument( "Pad operation does not support non-scalar padding values."); } - if (ShapeUtil::Rank(operand_shape) != padding_config.dimensions_size()) { + if (operand_shape.rank() != padding_config.dimensions_size()) { return InvalidArgument( "The rank of the operand and the padding configuration do not match: " "%s vs %s.", @@ -511,35 +517,40 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, padding_config.ShortDebugString()); } - std::vector dimensions(ShapeUtil::Rank(operand_shape)); + 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(); + 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: // // Contracting Dimensions: -// *) Exactly one contracting dimension on both lhs and rhs. +// *) Same number of contracting dimensions on both lhs and rhs. // *) Contracting dimension size must be the same on both lhs and rhs. -// *) Contracting dimension numbers do not need to be the same (i.e. transposes -// are passed on to emitter implementations). // // Batch Dimensions: // *) Same number of batch dimensions on both lhs and rhs. -// *) Same batch dimension numbers (and sizes) on both lhs and rhs. -// *) Batch dimension numbers must be ordered before contracting and -// non-contracting/non-batch dimension numbers. -// -// Non-Contracting-Non-Batch Dimensions: -// *) Can be 0 (matrix-vector) or 1 (matrix-matrix). +// *) Same batch dimension sizes on both lhs and rhs. // namespace { @@ -552,9 +563,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 = @@ -566,9 +576,9 @@ Status ValidateDotDimensionNumbers( absl::Span rhs_batch_dimensions = AsInt64Slice(dimension_numbers.rhs_batch_dimensions()); - if (!dims_in_range(ShapeUtil::Rank(lhs), lhs_contracting_dimensions, + if (!dims_in_range(lhs.rank(), lhs_contracting_dimensions, lhs_batch_dimensions) || - !dims_in_range(ShapeUtil::Rank(rhs), rhs_contracting_dimensions, + !dims_in_range(rhs.rank(), rhs_contracting_dimensions, rhs_batch_dimensions)) { return InvalidArgument("A dimension number is out of range in Dot: %s.", dimension_numbers.DebugString()); @@ -581,9 +591,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) || @@ -592,36 +601,6 @@ Status ValidateDotDimensionNumbers( dimension_numbers.DebugString()); } - // Check that the count of non-contracting-non-batch dimensions is in {0, 1}. - const int64 lhs_non_contracting_non_batch_dims = - ShapeUtil::Rank(lhs) - - dimension_numbers.lhs_contracting_dimensions_size() - - dimension_numbers.lhs_batch_dimensions_size(); - const int64 rhs_non_contracting_non_batch_dims = - ShapeUtil::Rank(rhs) - - dimension_numbers.rhs_contracting_dimensions_size() - - dimension_numbers.rhs_batch_dimensions_size(); - if (lhs_non_contracting_non_batch_dims < 0 || - lhs_non_contracting_non_batch_dims > 1 || - rhs_non_contracting_non_batch_dims < 0 || - rhs_non_contracting_non_batch_dims > 1) { - return InvalidArgument( - "Batch and contracting dimension number mismatch with rank."); - } - - // Check that batch dimension numbers are ordered before all others, and - // that they are monotonically increasing. - std::vector batch_dim_numbers(lhs_batch_dimensions.size()); - std::iota(batch_dim_numbers.begin(), batch_dim_numbers.end(), 0); - if (!std::equal(batch_dim_numbers.begin(), batch_dim_numbers.end(), - lhs_batch_dimensions.begin()) || - !std::equal(batch_dim_numbers.begin(), batch_dim_numbers.end(), - rhs_batch_dimensions.begin())) { - return InvalidArgument( - "Batch dimension numbers must precede non-batch dimensions and be" - "monotonically increasing."); - } - return Status::OK(); } @@ -648,28 +627,33 @@ Status ValidateDotDimensionNumbers( return fail("Element types do not match."); } - if ((ShapeUtil::Rank(lhs) < 1) || (ShapeUtil::Rank(rhs) < 1)) { + if ((lhs.rank() < 1) || (rhs.rank() < 1)) { return fail("Dot only supports rank 1 or above."); } // Validate basic properties of dot dimension numbers. TF_RETURN_IF_ERROR(ValidateDotDimensionNumbers(lhs, rhs, dimension_numbers)); - // Check that there is only one contracting dimension for both lhs and rhs. + // Check that number of contracting dimensions match. if (dimension_numbers.lhs_contracting_dimensions_size() != - dimension_numbers.rhs_contracting_dimensions_size() || - dimension_numbers.lhs_contracting_dimensions_size() != 1) { - return fail("Must specify one contracting dimension for both lhs and rhs."); + dimension_numbers.rhs_contracting_dimensions_size()) { + return fail( + "Must specify the same number of contracting dimensions for lhs and " + "rhs."); } - // Check that contracting dimension sizes match. - const int64 lhs_contracting_dimension = - dimension_numbers.lhs_contracting_dimensions(0); - const int64 rhs_contracting_dimension = - dimension_numbers.rhs_contracting_dimensions(0); - if (lhs.dimensions(lhs_contracting_dimension) != - rhs.dimensions(rhs_contracting_dimension)) { - return fail("Contracting dimension sizes do not match."); + for (int64 i = 0; i < dimension_numbers.lhs_contracting_dimensions_size(); + ++i) { + const int64 lhs_contracting_dimension = + dimension_numbers.lhs_contracting_dimensions(i); + const int64 rhs_contracting_dimension = + dimension_numbers.rhs_contracting_dimensions(i); + if (lhs.dimensions(lhs_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."); + } } // Check that number of batch dimensions match. @@ -680,11 +664,12 @@ Status ValidateDotDimensionNumbers( // Check that batch dimension numbers and sizes match. for (int64 i = 0; i < dimension_numbers.lhs_batch_dimensions_size(); ++i) { - if (dimension_numbers.lhs_batch_dimensions(i) != - dimension_numbers.rhs_batch_dimensions(i) || - lhs.dimensions(dimension_numbers.lhs_batch_dimensions(i)) != - rhs.dimensions(dimension_numbers.rhs_batch_dimensions(i))) { - return fail("Batch dimension numbers and sizes must match for lhs/rhs."); + if (lhs.dimensions(dimension_numbers.lhs_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."); } } @@ -694,21 +679,29 @@ Status ValidateDotDimensionNumbers( // Generate the result dimensions in order, rhs dimensions followed by lhs // dimensions except the contracted and batch dimensions. std::vector dimensions; - std::unordered_set rhs_batch_dims( - dimension_numbers.rhs_batch_dimensions().begin(), - dimension_numbers.rhs_batch_dimensions().end()); - for (int64 i = 0; i < ShapeUtil::Rank(lhs); i++) { - if (i != lhs_contracting_dimension) { + 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 < ShapeUtil::Rank(rhs); i++) { - if (i != rhs_contracting_dimension && rhs_batch_dims.count(i) == 0) { + for (int64 i = 0; i < rhs.rank(); i++) { + if (!absl::c_linear_search(dimension_numbers.rhs_contracting_dimensions(), + 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); @@ -719,20 +712,24 @@ Status ValidateDotDimensionNumbers( ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, const Shape& lhs, const Shape& rhs) { - TF_RET_CHECK(ShapeUtil::Rank(lhs) == ShapeUtil::Rank(rhs)); + TF_RET_CHECK(lhs.rank() == rhs.rank()); // The shapes have to be compatible. That is, if some dimension d has a // different size in the two shapes, one of them has to be 1 (a "degenerate" // 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(ShapeUtil::Rank(lhs)); - for (int64 i = 0; i < ShapeUtil::Rank(lhs); ++i) { + 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.", @@ -741,7 +738,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, } } return ShapeUtil::MakeShape(ShapeUtil::HigherPrecisionElementType(lhs, rhs), - output_dimensions); + output_dimensions, output_dimensions_is_dynamic); } /* static */ StatusOr ShapeInference::InferInDimBroadcastShape( @@ -754,13 +751,13 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, return InvalidArgument("Automatic shape inference not supported: %s and %s", ShapeUtil::HumanString(smaller_shape), ShapeUtil::HumanString(larger_shape)); - } else if (broadcast_dimensions.size() != ShapeUtil::Rank(smaller_shape)) { + } else if (broadcast_dimensions.size() != smaller_shape.rank()) { return InvalidArgument( "Size of broadcast_dimensions has to match lower-rank operand's " "rank; " " lower-rank operand's rank is %d, size of broadcast_dimensions is " "%u.", - ShapeUtil::Rank(smaller_shape), broadcast_dimensions.size()); + smaller_shape.rank(), broadcast_dimensions.size()); } // broadcast_dimensions is a sequence of dimensions; its length is equal to @@ -820,6 +817,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). @@ -831,6 +831,12 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, ShapeUtil::HumanString(smaller_shape), ShapeUtil::HumanString(larger_shape)); } + if (small_is_dynamic != large_is_dynamic) { + 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) { @@ -840,6 +846,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; @@ -858,8 +865,8 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, ShapeUtil::HumanString(rhs)); } - if (ShapeUtil::Rank(lhs) == ShapeUtil::Rank(rhs)) { - std::vector identity_dims(ShapeUtil::Rank(lhs)); + if (lhs.rank() == rhs.rank()) { + std::vector identity_dims(lhs.rank()); std::iota(identity_dims.begin(), identity_dims.end(), 0); if (!broadcast_dimensions.empty() && broadcast_dimensions != identity_dims) { @@ -876,15 +883,13 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, lhs, ShapeUtil::HigherPrecisionElementType(lhs, rhs)); } - if (ShapeUtil::Rank(lhs) == ShapeUtil::Rank(rhs)) { + if (lhs.rank() == rhs.rank()) { return InferDegenerateDimensionBroadcastShape(operation, lhs, rhs); } else { // Ranks do not match, so perform InDim broadcasting using // broadcast_dimensions. Scalar broadcasting is a special case of this. - const Shape& larger_shape = - ShapeUtil::Rank(lhs) > ShapeUtil::Rank(rhs) ? lhs : rhs; - const Shape& smaller_shape = - ShapeUtil::Rank(lhs) > ShapeUtil::Rank(rhs) ? rhs : lhs; + const Shape& larger_shape = lhs.rank() > rhs.rank() ? lhs : rhs; + const Shape& smaller_shape = lhs.rank() > rhs.rank() ? rhs : lhs; // After InDim broadcasting, perform degenerate dimensions broadcasting. TF_ASSIGN_OR_RETURN(Shape indim_broadcast_shape, @@ -953,6 +958,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."); } @@ -1029,7 +1036,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, switch (opcode) { case HloOpcode::kTuple: { Shape result = ShapeUtil::MakeTupleShape({}); - result.mutable_tuple_shapes()->Reserve(operand_shapes.size()); + result.mutable_tuple_shapes()->reserve(operand_shapes.size()); for (const Shape* shape : operand_shapes) { ShapeUtil::AppendShapeToTuple(*shape, &result); } @@ -1173,12 +1180,12 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, TF_RET_CHECK(ShapeUtil::ValidateShapeWithOptionalLayout(scale_shape) == Status::OK()); - if (feature_index >= ShapeUtil::Rank(operand_shape)) { + if (feature_index >= operand_shape.rank()) { return InvalidArgument( "Expected feature_index of batch-norm-training to be " "smaller than the rank of operand_shape; " "got feature_index %d, and rank %d.", - feature_index, ShapeUtil::Rank(operand_shape)); + feature_index, operand_shape.rank()); } if (feature_index < 0) { @@ -1188,25 +1195,25 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, feature_index); } - if (ShapeUtil::Rank(operand_shape) < 1) { + if (operand_shape.rank() < 1) { return InvalidArgument( "Expected the rank of operand to " "batch-norm-training to be at least 1; got %d.", - ShapeUtil::Rank(operand_shape)); + operand_shape.rank()); } - if (ShapeUtil::Rank(offset_shape) != 1) { + if (offset_shape.rank() != 1) { return InvalidArgument( "Offset input of batch-norm-training must have" " rank 1, but has rank %d.", - ShapeUtil::Rank(offset_shape)); + offset_shape.rank()); } - if (ShapeUtil::Rank(scale_shape) != 1) { + if (scale_shape.rank() != 1) { return InvalidArgument( "Scale input of batch-norm-training must have" " rank 1, but has rank %d.", - ShapeUtil::Rank(scale_shape)); + scale_shape.rank()); } if (!ShapeUtil::ElementIsFloating(operand_shape)) { @@ -1283,12 +1290,12 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, TF_RET_CHECK(ShapeUtil::ValidateShapeWithOptionalLayout(variance_shape) == Status::OK()); - if (feature_index >= ShapeUtil::Rank(operand_shape)) { + if (feature_index >= operand_shape.rank()) { return InvalidArgument( "Expected feature_index of batch-norm-inference to be " "smaller than the rank of operand_shape; " "got feature_index %d, and rank %d.", - feature_index, ShapeUtil::Rank(operand_shape)); + feature_index, operand_shape.rank()); } if (feature_index < 0) { @@ -1298,25 +1305,25 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, feature_index); } - if (ShapeUtil::Rank(operand_shape) < 1) { + if (operand_shape.rank() < 1) { return InvalidArgument( "Expected the rank of operand to " "batch-norm-inference to be at least 1; got %d.", - ShapeUtil::Rank(operand_shape)); + operand_shape.rank()); } - if (ShapeUtil::Rank(offset_shape) != 1) { + if (offset_shape.rank() != 1) { return InvalidArgument( "Offset input of batch-norm-inference must have" " rank 1, but has rank %d.", - ShapeUtil::Rank(offset_shape)); + offset_shape.rank()); } - if (ShapeUtil::Rank(scale_shape) != 1) { + if (scale_shape.rank() != 1) { return InvalidArgument( "Scale input of batch-norm-inference must have" " rank 1, but has rank %d.", - ShapeUtil::Rank(scale_shape)); + scale_shape.rank()); } if (!ShapeUtil::ElementIsFloating(operand_shape)) { @@ -1428,41 +1435,41 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, TF_RETURN_IF_ERROR( ShapeUtil::ValidateShapeWithOptionalLayout(output_grad_shape)); - if (feature_index >= ShapeUtil::Rank(operand_shape)) { + if (feature_index >= operand_shape.rank()) { return InvalidArgument( "Expected feature_index of batch-norm-grad to be " "smaller than the rank of operand_shape; " "got feature_index %d, and rank %d.", - feature_index, ShapeUtil::Rank(operand_shape)); + feature_index, operand_shape.rank()); } - if (ShapeUtil::Rank(operand_shape) != ShapeUtil::Rank(output_grad_shape)) { + if (operand_shape.rank() != output_grad_shape.rank()) { return InvalidArgument( "Expected operand_shape of batch-norm-grad to have the same rank as" " output_grad_shape; got rank(oprand_shape) %d, and" " rank(output_grad_shape) %d.", - ShapeUtil::Rank(operand_shape), ShapeUtil::Rank(output_grad_shape)); + operand_shape.rank(), output_grad_shape.rank()); } - if (ShapeUtil::Rank(mean_shape) != 1) { + if (mean_shape.rank() != 1) { return InvalidArgument( "Mean input of batch-norm-grad must have" " rank 1, but has rank %d.", - ShapeUtil::Rank(mean_shape)); + mean_shape.rank()); } - if (ShapeUtil::Rank(scale_shape) != 1) { + if (scale_shape.rank() != 1) { return InvalidArgument( "Scale input of batch-norm-grad must have" " rank 1, but has rank %d.", - ShapeUtil::Rank(scale_shape)); + scale_shape.rank()); } - if (ShapeUtil::Rank(var_shape) != 1) { + if (var_shape.rank() != 1) { return InvalidArgument( "Var input of batch-norm-grad must have" " rank 1, but has rank %d.", - ShapeUtil::Rank(var_shape)); + var_shape.rank()); } if (!ShapeUtil::ElementIsFloating(operand_shape)) { @@ -1549,7 +1556,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, } // Verify operand_shape and output_grad_shape have same bounds. - for (int64 i = 0; i < ShapeUtil::Rank(operand_shape); ++i) { + for (int64 i = 0; i < operand_shape.rank(); ++i) { if (ShapeUtil::GetDimension(operand_shape, i) != ShapeUtil::GetDimension(output_grad_shape, i)) { return InvalidArgument( @@ -1567,10 +1574,30 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, /* static */ StatusOr ShapeInference::InferConvolveShape( const Shape& lhs, const Shape& rhs, int64 feature_group_count, - const Window& window, const ConvolutionDimensionNumbers& dnums) { + int64 batch_group_count, const Window& window, + const ConvolutionDimensionNumbers& dnums) { TF_RETURN_IF_ERROR(ExpectArray(lhs, "lhs of convolution")); TF_RETURN_IF_ERROR(ExpectArray(rhs, "rhs of convolution")); + if (feature_group_count <= 0) { + return InvalidArgument( + "feature_group_count must be a positive number, got %d", + feature_group_count); + } + + if (batch_group_count <= 0) { + return InvalidArgument( + "batch_group_count must be a positive number, got %d", + 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.", @@ -1601,12 +1628,12 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, } const int num_dims = num_spatial_dims + 2; - if (ShapeUtil::Rank(lhs) != num_dims) { + if (lhs.rank() != num_dims) { return InvalidArgument( "The LHS argument to a convolution should have rank %d; lhs: %s.", num_dims, ShapeUtil::HumanString(lhs)); } - if (ShapeUtil::Rank(rhs) != num_dims) { + if (rhs.rank() != num_dims) { return InvalidArgument( "The RHS argument to a convolution should have rank %d; rhs: %s.", num_dims, ShapeUtil::HumanString(rhs)); @@ -1621,29 +1648,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()); @@ -1684,14 +1711,26 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, const int64 kernel_output_features = rhs.dimensions(dnums.kernel_output_feature_dimension()); - if (input_features != kernel_input_features * feature_group_count) { + 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( - "Expected LHS feature dimension (value %d) to match RHS " - "input feature dimension * feature_group_count (value %d * %d = %d); " + "Expected LHS feature dimension (value %d) to be a multiple of " + "feature_group_count (value %d), and LHS feature dimension / " + "feature_group_count = RHS feature dimension (value %d); " "got (%s, %s)\n" "Dimension numbers: {%s}.", - input_features, kernel_input_features, feature_group_count, - kernel_input_features * feature_group_count, + input_features, feature_group_count, kernel_input_features, ShapeUtil::HumanString(lhs), ShapeUtil::HumanString(rhs), dnums.DebugString()); } @@ -1705,6 +1744,17 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, ShapeUtil::HumanString(lhs), ShapeUtil::HumanString(rhs), dnums.DebugString()); } + + if (input_batch % batch_group_count > 0) { + return InvalidArgument( + "Expected input batch dimension (value %d) to be divisible by " + "batch_group_count (value %d); " + "got (%s, %s)\n" + "Dimension numbers: {%s}.", + input_batch, batch_group_count, ShapeUtil::HumanString(lhs), + ShapeUtil::HumanString(rhs), dnums.DebugString()); + } + std::vector window_dims(num_spatial_dims); for (int i = 0; i < num_spatial_dims; ++i) { window_dims[i] = window.dimensions(i).size(); @@ -1727,14 +1777,39 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, /*allow_negative_padding=*/true)); std::vector dimensions(num_dims); - dimensions[dnums.output_batch_dimension()] = input_batch; + dimensions[dnums.output_batch_dimension()] = input_batch / batch_group_count; dimensions[dnums.output_feature_dimension()] = kernel_output_features; for (int i = 0; i < num_spatial_dims; ++i) { 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( @@ -1755,7 +1830,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())); } @@ -1819,7 +1894,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, #undef RET_CHECK_RANK } -/* static */ StatusOr ShapeInference::InferCrossReplicaSumShape( +/* static */ StatusOr ShapeInference::InferAllReduceShape( absl::Span operand_shapes) { for (const Shape* operand_shape : operand_shapes) { TF_RETURN_IF_ERROR( @@ -1839,12 +1914,12 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, const Shape& shape, int64 split_dimension, int64 concat_dimension, int64 split_count) { TF_RET_CHECK(split_count > 0); - if (split_dimension >= ShapeUtil::Rank(shape) || split_dimension < 0) { + if (split_dimension >= shape.rank() || split_dimension < 0) { return InvalidArgument( "AllToAll split_dimension %d is out-of-bounds in shape %s.", split_dimension, ShapeUtil::HumanString(shape)); } - if (concat_dimension >= ShapeUtil::Rank(shape) || concat_dimension < 0) { + if (concat_dimension >= shape.rank() || concat_dimension < 0) { return InvalidArgument( "AllToAll concat_dimension %d is out-of-bounds in shape %s.", concat_dimension, ShapeUtil::HumanString(shape)); @@ -1882,7 +1957,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, /* static */ StatusOr ShapeInference::InferCollectivePermuteShape( const Shape& shape) { - TF_RET_CHECK(ShapeUtil::IsArray(shape)); + TF_RET_CHECK(shape.IsArray()); return shape; } @@ -1906,7 +1981,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])); @@ -1918,7 +1993,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, // doesn't matter which one we choose. const Shape& arg = *reduced_args[0]; for (int64 dimension : dimensions_to_reduce) { - if (dimension >= ShapeUtil::Rank(arg) || dimension < 0) { + if (dimension >= arg.rank() || dimension < 0) { return InvalidArgument("Reducing out-of-bounds dimension %d in shape %s.", dimension, ShapeUtil::HumanString(arg)); } @@ -1935,20 +2010,22 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, std::set dimensions_to_reduce_set(dimensions_to_reduce.begin(), dimensions_to_reduce.end()); std::vector new_dimensions; - for (int i = 0; i < ShapeUtil::Rank(arg); ++i) { + 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); } @@ -2022,9 +2099,29 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, ShapeUtil::HumanString(source_shape), ShapeUtil::HumanString(window_result_shape)); } + return operand_shape; } +/* static */ StatusOr ShapeInference::InferGetDimensionSizeShape( + const Shape& shape, int64 dimension) { + if (dimension < 0 || dimension >= shape.rank()) { + return InvalidArgument("GetDimensionSize dimension out of bounds: %d.", + dimension); + } + + // TODO(b/119580730): Remove this restriction when very large dimension size + // is needed. + if (shape.dimensions(dimension) > std::numeric_limits::max()) { + return InvalidArgument( + "GetDimensionSize's input shape is %s, the %dth dimension exceeds the " + "UINT_MAX limit.", + ShapeUtil::HumanString(shape), dimension); + } + + return ShapeUtil::MakeShape(U32, {}); +} + /* static */ StatusOr ShapeInference::InferSliceShape( const Shape& arg, absl::Span starts, absl::Span limits, absl::Span strides) { @@ -2050,10 +2147,10 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, starts.size(), strides.size())); } - if (starts.size() != ShapeUtil::Rank(arg)) { + if (starts.size() != arg.rank()) { return InvalidArgument( "Slice index count does not match argument rank: %u vs %d.", - starts.size(), ShapeUtil::Rank(arg)); + starts.size(), arg.rank()); } std::vector sizes; @@ -2088,41 +2185,87 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, } /* static */ StatusOr ShapeInference::InferDynamicSliceShape( - const Shape& operand_shape, const Shape& start_indices_shape, - absl::Span slice_sizes) { + const Shape& operand_shape, absl::Span start_index_shapes, + absl::Span slice_sizes, bool allow_scalar_indices) { TF_RETURN_IF_ERROR(ExpectArray(operand_shape, "operand of dynamic slice")); - TF_RETURN_IF_ERROR( - ExpectArray(start_indices_shape, "start indices of dynamic slice")); + auto number_of_indices = start_index_shapes.size(); + // TODO(b/118437727): Remove this path. + if (!allow_scalar_indices || + (number_of_indices >= 1 && start_index_shapes[0].rank() == 1)) { + if (number_of_indices != 1) { + return InvalidArgument( + "Dynamic slice should have exactly 1 index operand, has %d.", + number_of_indices); + } - 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, ", ")); + const Shape& start_indices_shape = start_index_shapes[0]; + 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, ", ")); - if (ShapeUtil::Rank(start_indices_shape) != 1) { - return InvalidArgument( - "Dynamic slice start indices of rank %d must be rank1.", - ShapeUtil::Rank(start_indices_shape)); - } + TF_RETURN_IF_ERROR( + ExpectArray(start_indices_shape, "start indices of dynamic slice")); - if (!ShapeUtil::ElementIsIntegral(start_indices_shape)) { - return InvalidArgument( - "Dynamic slice start indices must be of integral type."); - } + if (start_indices_shape.rank() != 1) { + return InvalidArgument( + "Dynamic slice start indices of rank %d must be rank1.", + start_indices_shape.rank()); + } - const int64 start_num_dims = start_indices_shape.dimensions(0); - if (ShapeUtil::Rank(operand_shape) != start_num_dims) { - return InvalidArgument( - "Dynamic slice start number of dimensions %d (%s) must match rank " - "%d of slice input (%s).", - start_num_dims, ShapeUtil::HumanString(start_indices_shape), - ShapeUtil::Rank(operand_shape), ShapeUtil::HumanString(operand_shape)); + if (!ShapeUtil::ElementIsIntegral(start_indices_shape)) { + return InvalidArgument( + "Dynamic slice start indices must be of integral type."); + } + + const int64 start_num_dims = start_indices_shape.dimensions(0); + if (operand_shape.rank() != start_num_dims) { + return InvalidArgument( + "Dynamic slice start number of dimensions %d (%s) must match rank " + "%d of slice input (%s).", + start_num_dims, ShapeUtil::HumanString(start_indices_shape), + operand_shape.rank(), ShapeUtil::HumanString(operand_shape)); + } + } else { + VLOG(2) << StrFormat("slicing shape %s a with slice_sizes={%s}", + ShapeUtil::HumanString(operand_shape), + StrJoin(slice_sizes, ", ")); + + if (operand_shape.rank() != number_of_indices) { + return InvalidArgument( + "Dynamic slice start number of dimensions %d must match rank " + "%d of slice input (%s).", + number_of_indices, operand_shape.rank(), + ShapeUtil::HumanString(operand_shape)); + } + + if (number_of_indices > 0) { + const Shape& first_index_shape = start_index_shapes[0]; + if (!ShapeUtil::IsScalar(first_index_shape)) { + return InvalidArgument("Dynamic slice indices must be scalar, not %s.", + ShapeUtil::HumanString(first_index_shape)); + } + if (!ShapeUtil::ElementIsIntegral(first_index_shape)) { + return InvalidArgument( + "Dynamic slice start indices must be of integral type."); + } + for (const Shape& index_shape : start_index_shapes) { + 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.", + ShapeUtil::HumanString(first_index_shape), + ShapeUtil::HumanString(index_shape)); + } + } + } } - if (slice_sizes.size() != ShapeUtil::Rank(operand_shape)) { + if (slice_sizes.size() != operand_shape.rank()) { return InvalidArgument( "Dynamic slice index count does not match argument rank: %u vs %d.", - slice_sizes.size(), ShapeUtil::Rank(operand_shape)); + slice_sizes.size(), operand_shape.rank()); } for (int64 dim = 0; dim < slice_sizes.size(); ++dim) { @@ -2145,46 +2288,92 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, /* static */ StatusOr ShapeInference::InferDynamicUpdateSliceShape( const Shape& operand_shape, const Shape& update_shape, - const Shape& start_indices_shape) { + absl::Span start_index_shapes, bool allow_scalar_indices) { TF_RETURN_IF_ERROR( ExpectArray(operand_shape, "operand of dynamic update slice")); TF_RETURN_IF_ERROR( ExpectArray(update_shape, "update of dynamic update slice")); - TF_RETURN_IF_ERROR(ExpectArray(start_indices_shape, - "start indices of dynamic update slice")); - VLOG(2) << StrFormat( - "updating slice of shape %s at dynamic start_indices %s with update " - "shape %s", - ShapeUtil::HumanString(operand_shape), - ShapeUtil::HumanString(start_indices_shape), - ShapeUtil::HumanString(update_shape)); + auto number_of_indices = start_index_shapes.size(); + // TODO(b/118437727): Remove this path. + if (!allow_scalar_indices || + (number_of_indices >= 1 && start_index_shapes[0].rank() == 1)) { + if (number_of_indices != 1) { + return InvalidArgument( + "Dynamic update slice should have exactly 1 index operand, has %d.", + number_of_indices); + } + const Shape& start_indices_shape = start_index_shapes[0]; + TF_RETURN_IF_ERROR(ExpectArray(start_indices_shape, + "start indices of dynamic update slice")); - if (ShapeUtil::Rank(start_indices_shape) != 1) { - return InvalidArgument( - "Dynamic update slice start indices of rank %d must be rank1.", - ShapeUtil::Rank(start_indices_shape)); - } + VLOG(2) << StrFormat( + "updating slice of shape %s at dynamic start_indices %s with update " + "shape %s", + ShapeUtil::HumanString(operand_shape), + ShapeUtil::HumanString(start_indices_shape), + ShapeUtil::HumanString(update_shape)); - if (!ShapeUtil::ElementIsIntegral(start_indices_shape)) { - return InvalidArgument( - "Dynamic update slice start indices must be of integral type."); - } + if (start_indices_shape.rank() != 1) { + return InvalidArgument( + "Dynamic update slice start indices of rank %d must be rank1.", + start_indices_shape.rank()); + } - const int64 start_num_dims = start_indices_shape.dimensions(0); - if (ShapeUtil::Rank(operand_shape) != start_num_dims) { - return InvalidArgument( - "Dynamic update slice start number of dimensions %d (%s) must match " - "rank %d of slice input (%s).", - start_num_dims, ShapeUtil::HumanString(start_indices_shape), - ShapeUtil::Rank(operand_shape), ShapeUtil::HumanString(operand_shape)); + if (!ShapeUtil::ElementIsIntegral(start_indices_shape)) { + return InvalidArgument( + "Dynamic update slice start indices must be of integral type."); + } + + const int64 start_num_dims = start_indices_shape.dimensions(0); + if (operand_shape.rank() != start_num_dims) { + return InvalidArgument( + "Dynamic update slice start number of dimensions %d (%s) must match " + "rank %d of slice input (%s).", + start_num_dims, ShapeUtil::HumanString(start_indices_shape), + operand_shape.rank(), ShapeUtil::HumanString(operand_shape)); + } + } else { + VLOG(2) << StrFormat("updating slice of shape %s with update shape %s", + ShapeUtil::HumanString(operand_shape), + ShapeUtil::HumanString(update_shape)); + + 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).", + number_of_indices, operand_shape.rank(), + ShapeUtil::HumanString(operand_shape)); + } + + if (number_of_indices > 0) { + const Shape& first_index_shape = start_index_shapes[0]; + if (!ShapeUtil::IsScalar(first_index_shape)) { + return InvalidArgument( + "Dynamic update slice indices must be scalar, not %s.", + ShapeUtil::HumanString(first_index_shape)); + } + if (!ShapeUtil::ElementIsIntegral(first_index_shape)) { + return InvalidArgument( + "Dynamic update slice start indices must be of integral type."); + } + for (const Shape& index_shape : start_index_shapes) { + 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.", + ShapeUtil::HumanString(first_index_shape), + ShapeUtil::HumanString(index_shape)); + } + } + } } - if (ShapeUtil::Rank(update_shape) != ShapeUtil::Rank(operand_shape)) { + if (update_shape.rank() != operand_shape.rank()) { return InvalidArgument( "Dynamic update slice update rank does not match argument rank: " "%d vs %d.", - ShapeUtil::Rank(update_shape), ShapeUtil::Rank(operand_shape)); + update_shape.rank(), operand_shape.rank()); } if (!ShapeUtil::SameElementTypeIgnoringFpPrecision(operand_shape, @@ -2196,7 +2385,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, PrimitiveType_Name(update_shape.element_type())); } - for (int64 dim = 0; dim < ShapeUtil::Rank(operand_shape); ++dim) { + for (int64 dim = 0; dim < operand_shape.rank(); ++dim) { const int64 input_dim_size = operand_shape.dimensions(dim); const int64 update_dim_size = update_shape.dimensions(dim); if (update_dim_size < 0) { @@ -2222,7 +2411,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, return InvalidArgument("a dimension number is duplicated in reverse"); } for (int64 dimension : dimensions) { - if (dimension >= ShapeUtil::Rank(operand_shape) || dimension < 0) { + if (dimension >= operand_shape.rank() || dimension < 0) { return InvalidArgument( "One of the reverse dimensions (%d) is out-of-bounds in shape %s.", dimension, ShapeUtil::HumanString(operand_shape)); @@ -2233,7 +2422,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, /* static */ StatusOr ShapeInference::InferGetTupleElementShape( const Shape& arg, int64 index) { - if (!ShapeUtil::IsTuple(arg)) { + if (!arg.IsTuple()) { return InvalidArgument( "Cannot infer shape: attempting to index into non-tuple: %s.", ShapeUtil::HumanString(arg)); @@ -2269,7 +2458,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()); } @@ -2289,7 +2478,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)); } @@ -2359,6 +2548,58 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, return ShapeUtil::MakeShape(operand.element_type(), dimensions); } +/* static */ StatusOr ShapeInference::InferBroadcastShape( + const Shape& operand_shape, const Shape& output_shape, + absl::Span broadcast_dimensions) { + TF_RETURN_IF_ERROR(ExpectArray(operand_shape, "operand of broadcast")); + TF_RETURN_IF_ERROR(ExpectArray(output_shape, "operand of broadcast")); + const int64 operand_rank = operand_shape.rank(); + const int64 output_rank = output_shape.rank(); + if (operand_rank > output_rank) { + return InvalidArgument( + "InDim style broadcast must be to an equal or higher ranked shape; " + "operand rank: %lld; output rank: %lld", + operand_rank, output_rank); + } + if (operand_rank != broadcast_dimensions.size()) { + return InvalidArgument( + "Size of broadcast_dimensions has to match operand's rank; operand " + "rank: %lld, size of broadcast_dimensions %u.", + operand_rank, broadcast_dimensions.size()); + } + for (int64 i = 0; i < operand_rank; i++) { + if (broadcast_dimensions[i] < 0 || broadcast_dimensions[i] >= output_rank) { + return InvalidArgument("Broadcast dimension %lld is out of bound", + broadcast_dimensions[i]); + } + if (operand_shape.dimensions(i) != + output_shape.dimensions(broadcast_dimensions[i]) && + 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 " + "%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]) { + return InvalidArgument( + "Broadcast dimensions order is wrong: %d comes after %d.", + broadcast_dimensions[i], broadcast_dimensions.at(i - 1)); + } + } + + return output_shape; +} + /* static */ StatusOr ShapeInference::InferReshapeShape( const Shape& operand, absl::Span dimensions, absl::Span new_sizes) { @@ -2378,9 +2619,9 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, ShapeUtil::HumanString(inferred_shape)); } - std::vector indices(ShapeUtil::Rank(operand)); + std::vector indices(operand.rank()); std::iota(indices.begin(), indices.end(), 0); - if (dimensions.size() != ShapeUtil::Rank(operand) || + if (dimensions.size() != operand.rank() || !std::is_permutation(dimensions.begin(), dimensions.end(), indices.begin())) { return InvalidArgument( @@ -2389,6 +2630,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; } @@ -2396,9 +2645,9 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, const Shape& operand, absl::Span dimensions) { TF_RETURN_IF_ERROR(ExpectArray(operand, "transpose")); - std::vector indices(ShapeUtil::Rank(operand)); + std::vector indices(operand.rank()); std::iota(indices.begin(), indices.end(), 0); - if (dimensions.size() != ShapeUtil::Rank(operand) || + if (dimensions.size() != operand.rank() || !std::is_permutation(dimensions.begin(), dimensions.end(), indices.begin())) { return InvalidArgument( @@ -2469,12 +2718,11 @@ ShapeInference::InferDegenerateDimensionBroadcastShape(HloOpcode operation, // dimensions as on_true and on_false. return 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)); } + return InvalidArgument( + "Select operation with non-scalar predicate with dimensionality " + "different from the other operands: %s.", + ShapeUtil::HumanString(pred)); } /* static */ StatusOr ShapeInference::InferTupleSelectShape( @@ -2750,7 +2998,7 @@ Status ValidateScatterDimensionNumbers( "update_window_dims in scatter op must not repeat; got: %s.", StrJoin(dim_numbers.update_window_dims(), ", ")); } - const int64 updates_rank = ShapeUtil::Rank(updates_shape); + const int64 updates_rank = updates_shape.rank(); for (int64 window_dim : dim_numbers.update_window_dims()) { if (window_dim < 0 || window_dim >= updates_rank) { return InvalidArgument( @@ -2781,6 +3029,15 @@ Status ValidateScatterDimensionNumbers( } } + // Validate window size. + auto window_size = dim_numbers.update_window_dims_size() + + dim_numbers.inserted_window_dims_size(); + if (window_size != operand_shape.rank()) { + return InvalidArgument( + "Scatter op has window of size %d; doesn't match operand of rank %d.", + window_size, operand_shape.rank()); + } + // Validate scatter_dims_to_operand_dims in ScatterDimensionNumbers. if (dim_numbers.scatter_dims_to_operand_dims_size() != scatter_indices_shape[dim_numbers.index_vector_dim()]) { @@ -2863,10 +3120,9 @@ Status ValidateScatterDimensionNumbers( int64 expected_updates_rank = expanded_scatter_indices_shape.size() - 1 + scatter_dim_numbers.update_window_dims_size(); - if (ShapeUtil::Rank(updates_shape) != expected_updates_rank) { + if (updates_shape.rank() != expected_updates_rank) { return InvalidArgument("Updates tensor must be of rank %d; got %d.", - expected_updates_rank, - ShapeUtil::Rank(updates_shape)); + expected_updates_rank, updates_shape.rank()); } TF_RETURN_IF_ERROR(ValidateScatterDimensionNumbers( @@ -2897,7 +3153,7 @@ Status ValidateScatterDimensionNumbers( } int64 scatter_dims_seen = 0; - for (int64 i = 0; i < ShapeUtil::Rank(updates_shape); ++i) { + for (int64 i = 0; i < updates_shape.rank(); ++i) { bool is_update_window_dim = absl::c_binary_search(scatter_dim_numbers.update_window_dims(), i); if (is_update_window_dim) { diff --git a/tensorflow/compiler/xla/service/shape_inference.h b/tensorflow/compiler/xla/service/shape_inference.h index 96a0ee165d46753da4fef119e7072f66637bf2c4..7d39ef38e05abf0a81683c1fb0f3999908b27d23 100644 --- a/tensorflow/compiler/xla/service/shape_inference.h +++ b/tensorflow/compiler/xla/service/shape_inference.h @@ -109,7 +109,7 @@ class ShapeInference { // filter (rhs) to lhs in the way specified by the fields on window. static StatusOr InferConvolveShape( const Shape& lhs, const Shape& rhs, int64 feature_group_count, - const Window& window, + int64 batch_group_count, const Window& window, const ConvolutionDimensionNumbers& dimension_numbers); // Infers the shape produced by the given FFT type on the given operand. @@ -118,7 +118,7 @@ class ShapeInference { // Infers the shape produced by a cross replica sum with the given operand // shapes. - static StatusOr InferCrossReplicaSumShape( + static StatusOr InferAllReduceShape( absl::Span operand_shapes); // Infers final shape of an Alltoall operation that is created by the xla @@ -176,14 +176,15 @@ class ShapeInference { // Infers the shape produced by a dynamic slice operation of size specified // in 'slice_sizes', with dynamic start indices shape 'start_indices_shape'. static StatusOr InferDynamicSliceShape( - const Shape& operand_shape, const Shape& start_indices_shape, - absl::Span slice_sizes); + const Shape& operand_shape, absl::Span start_index_shapes, + 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, - const Shape& start_indices_shape); + absl::Span start_index_shapes, + 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 @@ -210,6 +211,12 @@ class ShapeInference { static StatusOr InferBroadcastShape( const Shape& operand, absl::Span broadcast_sizes); + // Checks whether the given parameters can form a broadcast. Returns the same + // output_shape if it's legal. + static StatusOr InferBroadcastShape( + const Shape& operand_shape, const Shape& output_shape, + absl::Span broadcast_dimensions); + // Infers the shape produced by a reshape operation from the element type of // its operand and the new dimension sizes specified. static StatusOr InferReshapeShape(const Shape& operand, @@ -226,13 +233,6 @@ class ShapeInference { static StatusOr InferConcatOpShape( absl::Span arg_shapes, int64 dimension); - // Infers the shape produced by a kAfterAll. Trivially this shape is always a - // TOKEN shape. However, ShapeInference serves two purposes: inferring shapes - // and checking operand shapes. This method verifies that the operand shapes - // are all TOKENs. - static StatusOr InferAfterAllShape( - absl::Span arg_shapes); - // Helper that validates the given operand shape can be converted to the // target output_shape via a convert instruction -- the requirement is that // the shape is identical except for the element type. @@ -285,6 +285,9 @@ class ShapeInference { const Shape& updates_shape, const ProgramShape& to_apply_shape, const ScatterDimensionNumbers& scatter_dim_numbers); + static StatusOr InferGetDimensionSizeShape(const Shape& shape, + int64 dimension); + private: // Helper that infers the shape produced by performing an element-wise binary // operation with the given LHS and RHS shapes. diff --git a/tensorflow/compiler/xla/service/shape_inference_test.cc b/tensorflow/compiler/xla/service/shape_inference_test.cc index 7b65e8c1c9d2bc730c6c8550e9265b69fdde71cf..26120a06b823c9fddf378991cec434a880fb888d 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, {}); @@ -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) { @@ -420,7 +424,8 @@ TEST_F(ShapeInferenceTest, Convolve) { dim1->set_window_dilation(1); dim1->set_base_dilation(1); auto inferred_status = ShapeInference::InferConvolveShape( - lhs_shape, rhs_shape, /*feature_group_count=*/1, window, dnums); + lhs_shape, rhs_shape, /*feature_group_count=*/1, /*batch_group_count=*/1, + window, dnums); ASSERT_IS_OK(inferred_status.status()); Shape inferred_shape = inferred_status.ValueOrDie(); ASSERT_TRUE(ShapeUtil::Equal(ShapeUtil::MakeShape(F32, {10, 12, 2, 3}), @@ -465,7 +470,8 @@ TEST_F(ShapeInferenceTest, ConvolveWithWindowDilation) { dim1->set_window_dilation(2); dim1->set_base_dilation(1); auto inferred_status = ShapeInference::InferConvolveShape( - lhs_shape, rhs_shape, /*feature_group_count=*/1, window, dnums); + lhs_shape, rhs_shape, /*feature_group_count=*/1, /*batch_group_count=*/1, + window, dnums); ASSERT_IS_OK(inferred_status.status()); Shape inferred_shape = inferred_status.ValueOrDie(); ASSERT_TRUE(ShapeUtil::Equal(ShapeUtil::MakeShape(F32, {10, 12, 31, 5}), @@ -510,7 +516,8 @@ TEST_F(ShapeInferenceTest, ConvolveWithBaseDilation) { dim1->set_window_dilation(1); dim1->set_base_dilation(2); auto inferred_status = ShapeInference::InferConvolveShape( - lhs_shape, rhs_shape, /*feature_group_count=*/1, window, dnums); + lhs_shape, rhs_shape, /*feature_group_count=*/1, /*batch_group_count=*/1, + window, dnums); ASSERT_IS_OK(inferred_status.status()); Shape inferred_shape = inferred_status.ValueOrDie(); ASSERT_TRUE(ShapeUtil::Equal(ShapeUtil::MakeShape(F32, {10, 12, 4, 9}), @@ -548,7 +555,8 @@ TEST_F(ShapeInferenceTest, ConvolveDimensionNumbersOverlapError) { dim1->set_padding_low(1); dim1->set_padding_high(1); auto inferred_status = ShapeInference::InferConvolveShape( - lhs_shape, rhs_shape, /*feature_group_count=*/1, window, dnums); + lhs_shape, rhs_shape, /*feature_group_count=*/1, /*batch_group_count=*/1, + window, dnums); ASSERT_FALSE(inferred_status.ok()); ASSERT_THAT(inferred_status.status().error_message(), HasSubstr("each dimension exactly once")); @@ -1002,9 +1010,9 @@ TEST_F(ShapeInferenceTest, DotWithRankHigherThanTwo) { dot_dnums.add_rhs_contracting_dimensions(0); auto inferred_status = ShapeInference::InferDotOpShape( ShapeUtil::MakeShape(F32, {32, 32, 32}), matrix_32_64_, dot_dnums); - ASSERT_FALSE(inferred_status.ok()); - ASSERT_THAT(inferred_status.status().error_message(), - HasSubstr("Batch and contracting dimension number mismatch")); + EXPECT_TRUE(inferred_status.ok()); + EXPECT_TRUE(ShapeUtil::Equal(inferred_status.ValueOrDie(), + ShapeUtil::MakeShape(F32, {32, 32, 64}))); } // vector vector -> scalar @@ -1096,7 +1104,6 @@ TEST_F(ShapeInferenceTest, DotGeneral) { TEST_F(ShapeInferenceTest, DotWithTwoContractingDimsFails) { Shape lhs_shape = ShapeUtil::MakeShape(F32, {2, 11, 3, 2}); Shape rhs_shape = ShapeUtil::MakeShape(F32, {2, 3, 14}); - Shape output_shape = ShapeUtil::MakeShape(F32, {2, 11, 14}); DotDimensionNumbers dot_dnums; dot_dnums.add_lhs_contracting_dimensions(2); @@ -1110,8 +1117,28 @@ TEST_F(ShapeInferenceTest, DotWithTwoContractingDimsFails) { ShapeInference::InferDotOpShape(lhs_shape, rhs_shape, dot_dnums); ASSERT_FALSE(inferred_status.ok()); ASSERT_THAT(inferred_status.status().error_message(), - HasSubstr("Must specify one contracting dimension for both " - "lhs and rhs")); + HasSubstr("Must specify the same number of contracting " + "dimensions for lhs and rhs.")); +} + +TEST_F(ShapeInferenceTest, DotWithTwoContractingDimsPasses) { + Shape lhs_shape = ShapeUtil::MakeShape(F32, {2, 11, 3, 2}); + Shape rhs_shape = ShapeUtil::MakeShape(F32, {2, 3, 2, 14}); + Shape output_shape = ShapeUtil::MakeShape(F32, {2, 11, 14}); + + DotDimensionNumbers dot_dnums; + dot_dnums.add_lhs_contracting_dimensions(2); + dot_dnums.add_lhs_contracting_dimensions(3); + dot_dnums.add_lhs_batch_dimensions(0); + + dot_dnums.add_rhs_contracting_dimensions(1); + dot_dnums.add_rhs_contracting_dimensions(2); + dot_dnums.add_rhs_batch_dimensions(0); + + auto inferred_status = + ShapeInference::InferDotOpShape(lhs_shape, rhs_shape, dot_dnums); + EXPECT_TRUE(inferred_status.ok()); + EXPECT_TRUE(ShapeUtil::Equal(inferred_status.ValueOrDie(), output_shape)); } // BatchMatMul with different batch dimension sizes fails. @@ -1130,11 +1157,11 @@ TEST_F(ShapeInferenceTest, DotWithMisatchedBatchDimSizesFails) { ShapeInference::InferDotOpShape(lhs_shape, rhs_shape, dot_dnums); ASSERT_FALSE(inferred_status.ok()); ASSERT_THAT(inferred_status.status().error_message(), - HasSubstr("Batch dimension numbers and sizes must match")); + HasSubstr("Batch dimension sizes must match")); } -// BatchMatMul with different batch dimension numbers fails. -TEST_F(ShapeInferenceTest, DotWithMisatchedBatchDimNumbersFails) { +// BatchMatMul with different batch dimension numbers passes +TEST_F(ShapeInferenceTest, DotWithMisatchedBatchDimNumbersPasses) { Shape lhs_shape = ShapeUtil::MakeShape(F32, {2, 11, 3}); Shape rhs_shape = ShapeUtil::MakeShape(F32, {3, 2, 14}); @@ -1147,9 +1174,9 @@ TEST_F(ShapeInferenceTest, DotWithMisatchedBatchDimNumbersFails) { auto inferred_status = ShapeInference::InferDotOpShape(lhs_shape, rhs_shape, dot_dnums); - ASSERT_FALSE(inferred_status.ok()); - ASSERT_THAT(inferred_status.status().error_message(), - HasSubstr("Batch dimension numbers must precede non-batch")); + ASSERT_TRUE(inferred_status.ok()); + ASSERT_TRUE(ShapeUtil::Equal(inferred_status.ValueOrDie(), + ShapeUtil::MakeShape(F32, {2, 11, 14}))); } // BatchMatMul with out-of-range dimension numbers fails. @@ -2673,5 +2700,23 @@ TEST_F(ScatterGatherShapeInferenceTest, << statusor.status(); } +TEST_F(ScatterGatherShapeInferenceTest, + InvalidScatterDimNumbers_InsufficientWindowDims) { + StatusOr statusor = ShapeInference::InferScatterShape( + f32_5d_tensor_50_49_48_47_46_, s64_scalar_, + ShapeUtil::MakeShape(F32, {30, 29, 28, 27}), to_apply_, + HloScatterInstruction::MakeScatterDimNumbers( + /*update_window_dims=*/{0, 1, 2, 3}, + /*inserted_window_dims=*/{}, + /*scatter_dims_to_operand_dims=*/{0}, + /*index_vector_dim=*/0)); + ASSERT_FALSE(statusor.ok()); + EXPECT_THAT( + statusor.status().error_message(), + HasSubstr( + "Scatter op has window of size 4; doesn't match operand of rank 5.")) + << statusor.status(); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/service/shaped_buffer.cc b/tensorflow/compiler/xla/service/shaped_buffer.cc index 56952e3adae59656605a12fd499162504a2a3379..d90dde3b13d3aa9e1de10dd9e1d11a8e6da170de 100644 --- a/tensorflow/compiler/xla/service/shaped_buffer.cc +++ b/tensorflow/compiler/xla/service/shaped_buffer.cc @@ -85,7 +85,7 @@ string ShapedBuffer::ToString() const { on_device_shape(), [this, &s](const Shape& subshape, const ShapeIndex& index) { string shape_str; - if (ShapeUtil::IsTuple(subshape)) { + if (subshape.IsTuple()) { shape_str = "tuple"; } else { shape_str = ShapeUtil::HumanStringWithLayout(subshape); @@ -157,4 +157,23 @@ void ScopedShapedBuffer::Deallocate() { } } +ScopedShapedBuffer ScopedShapedBuffer::TakeSubTree(ShapeIndexView index) { + const xla::Shape& sub_on_host_shape = + xla::ShapeUtil::GetSubshape(on_host_shape(), {index}); + const xla::Shape& sub_on_device_shape = + xla::ShapeUtil::GetSubshape(on_device_shape(), {index}); + + ScopedShapedBuffer output(sub_on_host_shape, sub_on_device_shape, + memory_allocator(), device_ordinal()); + auto src_it = buffers().find(index); + auto dst_it = output.buffers().begin(); + while (dst_it != output.buffers().end()) { + dst_it->second = src_it->second; + src_it->second = tensorflow::se::DeviceMemoryBase(nullptr, 0); + ++src_it; + ++dst_it; + } + return output; +} + } // namespace xla diff --git a/tensorflow/compiler/xla/service/shaped_buffer.h b/tensorflow/compiler/xla/service/shaped_buffer.h index e1d26da4a20c0105be304b1a34c81515fcdc6b7f..f5210c9cfa6b29853bcd0f5bfd581ee3e116a509 100644 --- a/tensorflow/compiler/xla/service/shaped_buffer.h +++ b/tensorflow/compiler/xla/service/shaped_buffer.h @@ -176,6 +176,11 @@ class ScopedShapedBuffer : public ShapedBuffer { // It's the caller's job to ensure that the memory contained therein is freed. TF_MUST_USE_RESULT ShapedBuffer release(); + // Extracts the sub-tree rooted at 'index' and returns a ScopedShapedBuffer + // that holds ownership of the subtree. Sets the buffers corresponding to the + // subtree to null in 'this'. + ScopedShapedBuffer TakeSubTree(ShapeIndexView index); + protected: void Deallocate(); diff --git a/tensorflow/compiler/xla/service/shaped_buffer_test.cc b/tensorflow/compiler/xla/service/shaped_buffer_test.cc index d69e6362e91e4696dab3c46d99a981c67b593a1c..ca64bd3c8dd2baa686db2b85c937a034b37ab22b 100644 --- a/tensorflow/compiler/xla/service/shaped_buffer_test.cc +++ b/tensorflow/compiler/xla/service/shaped_buffer_test.cc @@ -20,6 +20,8 @@ limitations under the License. #include "tensorflow/compiler/xla/service/platform_util.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/test.h" +#include "tensorflow/core/platform/stream_executor_no_cuda.h" +#include "tensorflow/core/platform/test_benchmark.h" #include "tensorflow/core/util/ptr_util.h" namespace xla { @@ -107,5 +109,79 @@ TEST(ScopedShapedBufferTest, TestMoveAssignmentOperator) { // TestAllocator's destructor checks that all memory was freed. } +TEST(ScopedShapedBufferTest, TestTakeSubTree) { + TestAllocator allocator; + + Shape s = ShapeUtil::MakeShape(F32, {1}); + s = xla::ShapeUtil::MakeTupleShape(std::vector(2, s)); + s = xla::ShapeUtil::MakeTupleShape(std::vector(3, s)); + + ScopedShapedBuffer sb(s, s, &allocator, /*device_ordinal=*/0); + sb.buffers().ForEachMutableElement( + [&](const xla::ShapeIndex& index, se::DeviceMemoryBase* buffer) { + TF_ASSERT_OK_AND_ASSIGN( + OwningDeviceMemory m, + allocator.Allocate(/*device_ordinal=*/0, /*size=*/77)); + *buffer = m.Forget(); + }); + ShapeTree buffers = sb.buffers(); + + // Takes a subtree out of 'sb', and verifies the buffers are as expected. + xla::ShapeIndex subtree_index = {1}; + ScopedShapedBuffer output = sb.TakeSubTree(subtree_index); + + output.buffers().ForEachElement([&](const xla::ShapeIndex& sub_index, + const se::DeviceMemoryBase& buffer) { + xla::ShapeIndex orig_index = subtree_index; + for (int i : sub_index) { + orig_index.push_back(i); + } + EXPECT_TRUE(buffers.find(orig_index)->second.IsSameAs(buffer)); + }); + sb.buffers().ForEachElement( + [&](const xla::ShapeIndex& index, const se::DeviceMemoryBase& buffer) { + if (ShapeIndexView(index).StartsWith(subtree_index)) { + EXPECT_TRUE(buffer.is_null()); + } else { + EXPECT_TRUE(buffers.find(index)->second.IsSameAs(buffer)); + } + }); +} + +// Test TakeSubTree with different depths (depth of ShapeTree) and fan-outs +// (cardinality of each non-leaf node's children). +void BM_TakeSubTree(int iters, int depth, int fan_out) { + tensorflow::testing::StopTiming(); + TestAllocator allocator; + xla::Shape shape = xla::ShapeUtil::MakeShape(xla::F32, {32, 64, 128}); + for (int i = 0; i < depth; ++i) { + std::vector shapes(fan_out, shape); + shape = xla::ShapeUtil::MakeTupleShape(shapes); + } + xla::ScopedShapedBuffer shaped_buffer(shape, shape, /*allocator=*/&allocator, + /*device_ordinal=*/0); + tensorflow::testing::StartTiming(); + for (int i = 0; i < iters; ++i) { + // Extract a buffer from approximately the middle of the first level of the + // tree. + (void)shaped_buffer.TakeSubTree(/*index=*/{fan_out / 2}).release(); + } + tensorflow::testing::StopTiming(); +} + +BENCHMARK(BM_TakeSubTree) + ->ArgPair(1, 4) + ->ArgPair(1, 8) + ->ArgPair(1, 32) + ->ArgPair(1, 64) + ->ArgPair(1, 128) + ->ArgPair(1, 256) + ->ArgPair(1, 512) + ->ArgPair(2, 4) + ->ArgPair(2, 8) + ->ArgPair(2, 32) + ->ArgPair(2, 64) + ->ArgPair(2, 128); + } // anonymous namespace } // namespace xla diff --git a/tensorflow/compiler/xla/service/sort_simplifier.cc b/tensorflow/compiler/xla/service/sort_simplifier.cc new file mode 100644 index 0000000000000000000000000000000000000000..4a00e8d7b227f14d462ca53f695189f3f48754ee --- /dev/null +++ b/tensorflow/compiler/xla/service/sort_simplifier.cc @@ -0,0 +1,126 @@ +/* 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_computation.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/statusor.h" + +#include "absl/container/flat_hash_map.h" +#include "absl/container/flat_hash_set.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; + } + + // 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; + } + 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..cd05fcf830d32e8bac4f8b260d3dd143ab98ad7b --- /dev/null +++ b/tensorflow/compiler/xla/service/sort_simplifier_test.cc @@ -0,0 +1,102 @@ +/* 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 + + 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)); + + 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 + + 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)); + + 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 + + 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)); + + SortSimplifier simplifier; + EXPECT_FALSE(simplifier.Run(module.get()).ValueOrDie()); +} +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/service/transfer_manager.cc b/tensorflow/compiler/xla/service/transfer_manager.cc index a21e586efadb85d18e88e44999283b28f7f65eac..15ef623cc7b2dbc31e9cba5c4783c39b8805a5aa 100644 --- a/tensorflow/compiler/xla/service/transfer_manager.cc +++ b/tensorflow/compiler/xla/service/transfer_manager.cc @@ -142,7 +142,7 @@ Status TransferManager::TransferArrayToDeviceAsync( se::Stream* stream, const LiteralSlice& literal, const se::DeviceMemoryBase& dest) { const Shape on_device_shape = HostShapeToDeviceShape(literal.shape()); - TF_RET_CHECK(ShapeUtil::IsArray(on_device_shape)) + TF_RET_CHECK(on_device_shape.IsArray()) << "On-device representation of " << ShapeUtil::HumanString(literal.shape()) << " is not an array: " << ShapeUtil::HumanString(on_device_shape); @@ -227,7 +227,7 @@ Status TransferManager::WriteTupleIndexTablesAsync( return ShapeUtil::ForEachSubshapeWithStatus( device_buffer.on_device_shape(), [&](const Shape& device_subshape, const ShapeIndex& index) -> Status { - if (ShapeUtil::IsTuple(device_subshape)) { + if (device_subshape.IsTuple()) { se::DeviceMemoryBase device_memory = device_buffer.buffer(index); TF_RET_CHECK(GetByteSizeRequirement(device_subshape) == device_memory.size()); @@ -248,6 +248,22 @@ Status TransferManager::WriteTupleIndexTablesAsync( }); } +Status TransferManager::WriteRootTupleIndexTable( + se::Stream* stream, const ShapedBuffer& device_buffer) { + TF_RET_CHECK(device_buffer.on_device_shape().IsTuple()); + se::DeviceMemoryBase device_memory = device_buffer.buffer({}); + TF_RET_CHECK(GetByteSizeRequirement(device_buffer.on_device_shape()) == + device_memory.size()); + + std::vector elements; + for (int64 i = 0; + i < ShapeUtil::TupleElementCount(device_buffer.on_device_shape()); ++i) { + elements.push_back(device_buffer.buffer({i})); + } + return WriteSingleTupleIndexTable( + stream, elements, device_buffer.on_device_shape(), &device_memory); +} + Status TransferManager::TransferBufferFromDevice( se::Stream* stream, const se::DeviceMemoryBase& source, int64 size, void* destination) { diff --git a/tensorflow/compiler/xla/service/transfer_manager.h b/tensorflow/compiler/xla/service/transfer_manager.h index f952e64af2b675b9c0f8a30e9a2bc3c855e34efa..43a50487c636da75224547286a31625db3f91330 100644 --- a/tensorflow/compiler/xla/service/transfer_manager.h +++ b/tensorflow/compiler/xla/service/transfer_manager.h @@ -95,7 +95,13 @@ class TransferManager { // but need not have the same layout. // // This operation is performed asynchronously on the given stream. It returns - // once the transfer is enqueued. + // once the transfer is enqueued, and may return before the transfer has + // completed. + // + // The caller may free the data structures 'literal' and 'device_buffer' + // immediately after this function returns, however their constituent buffers + // on both host and device must remain valid until the enqueued transfer has + // completed on 'stream'. virtual Status TransferLiteralToDeviceAsync( se::Stream* stream, const LiteralSlice& literal, const ShapedBuffer& device_buffer) = 0; @@ -140,6 +146,12 @@ class TransferManager { Status WriteTupleIndexTablesAsync(se::Stream* stream, const ShapedBuffer& device_buffer); + // Writes a tuple index buffer for the root of 'device_buffer', which must + // be a tuple. Unlike WriteTupleIndexTables, only writes the root buffer, + // rather than writing all subbuffers. This method is always asynchronous. + Status WriteRootTupleIndexTable(se::Stream* stream, + const ShapedBuffer& device_buffer); + // Determines the byte size requirement for the given shape on the underlying // architecture. This will be used to allocate an appropriately sized memory // region for a host-to-device transfer. diff --git a/tensorflow/compiler/xla/service/transpose_folding.cc b/tensorflow/compiler/xla/service/transpose_folding.cc index 7c1f4b5cc67dd2a84271b4f2b8015fdb2ff6e846..a95ca2bf2a8fcd700eb9234cafbfce9b62f2370c 100644 --- a/tensorflow/compiler/xla/service/transpose_folding.cc +++ b/tensorflow/compiler/xla/service/transpose_folding.cc @@ -45,7 +45,7 @@ TransposeFolding::OperandIndices CanFoldOperandsIntoDot( auto& operand = *dot.operand(i); if (operand.IsRank2Transpose()) { operand_set.push_back(i); - } else if (ShapeUtil::Rank(operand.shape()) != 2) { + } else if (operand.shape().rank() != 2) { return {}; } } @@ -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); @@ -178,7 +176,8 @@ bool FoldTransposeIntoConvolution(InstructionOperandsPair pair) { auto new_conv = HloInstruction::CreateConvolve( convolution.shape(), new_lhs, new_rhs, convolution.feature_group_count(), - convolution.window(), new_dnums, convolution.precision_config()); + convolution.batch_group_count(), convolution.window(), new_dnums, + convolution.precision_config()); TF_CHECK_OK(convolution.parent()->ReplaceWithNewInstruction( &convolution, std::move(new_conv))); diff --git a/tensorflow/compiler/xla/service/transpose_folding_test.cc b/tensorflow/compiler/xla/service/transpose_folding_test.cc index 79b5c09abb355cd067a4891af558c8c44d80d88e..f8a5fa0215007310d6bec35d20fc643afc824dda 100644 --- a/tensorflow/compiler/xla/service/transpose_folding_test.cc +++ b/tensorflow/compiler/xla/service/transpose_folding_test.cc @@ -139,9 +139,9 @@ TEST_F(TransposeFoldingTest, FoldDotTransposeConstant) { HloModule FoldDotTransposeConstant ENTRY entry_computation { - constant = f32[2,1]{1,0} constant(f32[2,1] { { 1 }, { 2 } }) + constant = f32[2,1]{1,0} constant({ { 1 }, { 2 } }) transpose = f32[1,2]{1,0} transpose(constant), dimensions={1,0} - constant.1 = f32[3,2]{1,0} constant(f32[3,2] { { 1, 2 }, { 3, 4 }, { 5, 6 } }) + constant.1 = f32[3,2]{1,0} constant({ { 1, 2 }, { 3, 4 }, { 5, 6 } }) transpose.1 = f32[2,3]{1,0} transpose(constant.1), dimensions={1,0} ROOT dot = f32[1,3]{1,0} dot(transpose, transpose.1), lhs_contracting_dims={1}, rhs_contracting_dims={0} } @@ -172,7 +172,7 @@ TEST_F(TransposeFoldingTest, FuseDotWithConstantOperands) { HloInstruction* mul = builder.AddInstruction(HloInstruction::CreateBinary( add->shape(), HloOpcode::kMultiply, add, sub)); - auto module = CreateNewModule("fuse_with_constant_operands"); + auto module = CreateNewVerifiedModule("fuse_with_constant_operands"); HloComputation* entry_computation = module->AddEntryComputation(builder.Build(mul)); HloInstruction* call = module->OutlineExpressionFromComputation( @@ -240,14 +240,15 @@ TEST_F(TransposeFoldingTest, FoldConvDimSwapTransposeRhs) { transpose_y->shape().dimensions(dnums.kernel_spatial_dimensions(i))); } StatusOr conv_shape = ShapeInference::InferConvolveShape( - x->shape(), transpose_y->shape(), /*feature_group_count=*/1, window, - dnums); + x->shape(), transpose_y->shape(), /*feature_group_count=*/1, + /*batch_group_count=*/1, window, dnums); EXPECT_IS_OK(conv_shape); HloInstruction* conv = builder.AddInstruction(HloInstruction::CreateConvolve( conv_shape.ValueOrDie(), x, transpose_y, - /*feature_group_count=*/1, window, dnums, DefaultPrecisionConfig(2))); + /*feature_group_count=*/1, /*batch_group_count=*/1, window, dnums, + DefaultPrecisionConfig(2))); - auto module = CreateNewModule("test_module"); + auto module = CreateNewVerifiedModule("test_module"); HloComputation* entry_computation = module->AddEntryComputation(builder.Build(conv)); FoldTranspose(module.get()); @@ -295,14 +296,15 @@ TEST_F(TransposeFoldingTest, FoldConvComplexTransposeRhs) { transpose_y->shape().dimensions(dnums.kernel_spatial_dimensions(i))); } StatusOr conv_shape = ShapeInference::InferConvolveShape( - x->shape(), transpose_y->shape(), /*feature_group_count=*/1, window, - dnums); + x->shape(), transpose_y->shape(), /*feature_group_count=*/1, + /*batch_group_count=*/1, window, dnums); EXPECT_IS_OK(conv_shape); HloInstruction* conv = builder.AddInstruction(HloInstruction::CreateConvolve( conv_shape.ValueOrDie(), x, transpose_y, - /*feature_group_count=*/1, window, dnums, DefaultPrecisionConfig(2))); + /*feature_group_count=*/1, /*batch_group_count=*/1, window, dnums, + DefaultPrecisionConfig(2))); - auto module = CreateNewModule("test_module"); + auto module = CreateNewVerifiedModule("test_module"); HloComputation* entry_computation = module->AddEntryComputation(builder.Build(conv)); FoldTranspose(module.get()); @@ -355,14 +357,15 @@ TEST_F(TransposeFoldingTest, FoldConvTransposeLhs) { dim->set_size(y->shape().dimensions(dnums.kernel_spatial_dimensions(i))); } StatusOr conv_shape = ShapeInference::InferConvolveShape( - transpose_x->shape(), y->shape(), /*feature_group_count=*/1, window, - dnums); + transpose_x->shape(), y->shape(), /*feature_group_count=*/1, + /*batch_group_count=*/1, window, dnums); EXPECT_IS_OK(conv_shape); HloInstruction* conv = builder.AddInstruction(HloInstruction::CreateConvolve( conv_shape.ValueOrDie(), transpose_x, y, - /*feature_group_count=*/1, window, dnums, DefaultPrecisionConfig(2))); + /*feature_group_count=*/1, /*batch_group_count=*/1, window, dnums, + DefaultPrecisionConfig(2))); - auto module = CreateNewModule("test_module"); + auto module = CreateNewVerifiedModule("test_module"); HloComputation* entry_computation = module->AddEntryComputation(builder.Build(conv)); FoldTranspose(module.get()); @@ -421,14 +424,15 @@ TEST_F(TransposeFoldingTest, FoldConvComplexTransposeLhs) { dim->set_size(y->shape().dimensions(dnums.kernel_spatial_dimensions(i))); } StatusOr conv_shape = ShapeInference::InferConvolveShape( - transpose_x->shape(), y->shape(), /*feature_group_count=*/1, window, - dnums); + transpose_x->shape(), y->shape(), /*feature_group_count=*/1, + /*batch_group_count=*/1, window, dnums); EXPECT_IS_OK(conv_shape); HloInstruction* conv = builder.AddInstruction(HloInstruction::CreateConvolve( conv_shape.ValueOrDie(), transpose_x, y, - /*feature_group_count=*/1, window, dnums, DefaultPrecisionConfig(2))); + /*feature_group_count=*/1, /*batch_group_count=*/1, window, dnums, + DefaultPrecisionConfig(2))); - auto module = CreateNewModule("test_module"); + auto module = CreateNewVerifiedModule("test_module"); HloComputation* entry_computation = module->AddEntryComputation(builder.Build(conv)); FoldTranspose(module.get()); diff --git a/tensorflow/compiler/xla/service/tuple_points_to_analysis.cc b/tensorflow/compiler/xla/service/tuple_points_to_analysis.cc index 96f3055c98e0611dfe25517cb490014a6d1f7c76..5e505aaf02f157d0cba9dff42b1a9b89a6691504 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, @@ -210,7 +207,7 @@ Status TuplePointsToAnalysis::DefaultAction(HloInstruction* hlo_instruction) { &logical_buffer_analysis_->GetBuffer(hlo_instruction, index)); }); - if (ShapeUtil::IsTuple(hlo_instruction->shape())) { + if (hlo_instruction->shape().IsTuple()) { // If the hlo instruction is a tuple-shaped, then trivially the instruction // itself is the source of the tuple. points_to_set.add_tuple_source({}, hlo_instruction); @@ -280,6 +277,13 @@ Status TuplePointsToAnalysis::HandleDomain(HloInstruction* domain) { return Status::OK(); } +Status TuplePointsToAnalysis::HandleAddDependency( + HloInstruction* add_dependency) { + // AddDependency just forwards the value of its zero-th operand. + CreateCopiedPointsToSet(add_dependency, add_dependency->operand(0)); + return Status::OK(); +} + Status TuplePointsToAnalysis::HandleRecvDone(HloInstruction* recv_done) { // RecvDone aliases its input (Recv) tuple element {0} to element {0} of its // output. The other indices ({} and {1}) define their own buffers. @@ -597,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()); @@ -665,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()) { @@ -736,11 +738,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; } diff --git a/tensorflow/compiler/xla/service/tuple_points_to_analysis.h b/tensorflow/compiler/xla/service/tuple_points_to_analysis.h index bcfcb388f95b0bedb35a8c399e804034816867b3..0a1d5649d6d69fea12263e6986ce76af62615ec7 100644 --- a/tensorflow/compiler/xla/service/tuple_points_to_analysis.h +++ b/tensorflow/compiler/xla/service/tuple_points_to_analysis.h @@ -252,6 +252,7 @@ class TuplePointsToAnalysis : public DfsHloVisitorWithDefault { Status HandleRecvDone(HloInstruction* recv_done) override; Status HandleSend(HloInstruction* send) override; Status HandleTupleSelect(HloInstruction* tuple_select) override; + Status HandleAddDependency(HloInstruction* add_dependency) override; string ToString() const; 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 d9ebebf74ed846aa05326a4df72019ef3e71ad88..fd5759e44230db8223822d6ae0f511027f73d8f9 100644 --- a/tensorflow/compiler/xla/service/tuple_points_to_analysis_test.cc +++ b/tensorflow/compiler/xla/service/tuple_points_to_analysis_test.cc @@ -48,7 +48,7 @@ class TuplePointsToAnalysisTest : public HloTestBase { } void BuildModule(std::unique_ptr computation) { - module_ = CreateNewModule(); + module_ = CreateNewUnverifiedModule(); module_->AddEntryComputation(std::move(computation)); } @@ -264,6 +264,22 @@ TEST_F(TuplePointsToAnalysisTest, GetTupleElement) { UnorderedElementsAre(inner_tuple)); } +TEST_F(TuplePointsToAnalysisTest, AddDependency) { + auto builder = HloComputation::Builder(TestName()); + auto constant = builder.AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0))); + auto token = builder.AddInstruction(HloInstruction::CreateToken()); + auto add_dependency = builder.AddInstruction( + HloInstruction::CreateAddDependency(constant, token)); + BuildModuleAndRunAnalysis(builder.Build()); + + auto& points_to_set = points_to_analysis_->GetPointsToSet(add_dependency); + EXPECT_EQ(1, points_to_set.size()); + EXPECT_FALSE(points_to_set.IsAmbiguous()); + EXPECT_TRUE(points_to_set.IsDistinct()); + ExpectHasTopLevelBuffers(points_to_set.CreateFlattenedSet(), {constant}); +} + TEST_F(TuplePointsToAnalysisTest, DuplicatedElement) { // Create a tuple which contains duplicate elements. auto builder = HloComputation::Builder(TestName()); @@ -607,7 +623,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}); @@ -641,7 +657,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()); @@ -705,9 +721,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; }); @@ -718,7 +733,7 @@ class FusionPointsToAnalysisTest : public TuplePointsToAnalysisTest { // Returns all users of 'fusion_paran' at 'tuple_index'. std::vector GetFusionParameterUsersAt( HloInstruction* fusion_param, int64 tuple_index) { - CHECK(ShapeUtil::IsTuple(fusion_param->shape())); + CHECK(fusion_param->shape().IsTuple()); std::vector users_at_tuple_index; for (auto user : fusion_param->users()) { CHECK_EQ(HloOpcode::kGetTupleElement, user->opcode()); @@ -809,7 +824,7 @@ TEST_F(FusionPointsToAnalysisTest, FusionParam0TwoUsers) { class PointsToAnalysisTestBase : public HloTestBase { protected: void BuildModule(std::unique_ptr computation) { - module_ = CreateNewModule(); + module_ = CreateNewUnverifiedModule(); computation_ = module_->AddEntryComputation(std::move(computation)); } @@ -867,12 +882,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})); @@ -961,12 +976,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})); @@ -988,7 +1003,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( @@ -996,7 +1011,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()); @@ -1176,7 +1191,7 @@ TEST_F(CanShareOperandBufferWithUserTest, WhileCanShare) { return builder.Build(); }; - module_ = CreateNewModule(); + module_ = CreateNewUnverifiedModule(); HloComputation* cond_computation = module_->AddEmbeddedComputation(make_cond()); HloComputation* body_computation = @@ -1211,7 +1226,7 @@ TEST_F(CanShareOperandBufferWithUserTest, CallToComputationWithFusionRoot) { auto add = sub_builder.AddInstruction( HloInstruction::CreateBinary(shape, HloOpcode::kAdd, sub_param, ones)); - module_ = CreateNewModule(); + module_ = CreateNewUnverifiedModule(); auto sub_computation = module_->AddEmbeddedComputation(sub_builder.Build()); sub_computation->CreateFusionInstruction({add, ones}, HloInstruction::FusionKind::kLoop); diff --git a/tensorflow/compiler/xla/service/tuple_simplifier_test.cc b/tensorflow/compiler/xla/service/tuple_simplifier_test.cc index 516754e2110ee50a597818c4a8bcfbfbb76c5cec..65b0f8c804475d8f22fff9798e79c9881a51f1f1 100644 --- a/tensorflow/compiler/xla/service/tuple_simplifier_test.cc +++ b/tensorflow/compiler/xla/service/tuple_simplifier_test.cc @@ -25,7 +25,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_opcode.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/test.h" -#include "tensorflow/compiler/xla/tests/hlo_verified_test_base.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/core/lib/core/status_test_util.h" @@ -34,7 +34,7 @@ namespace op = xla::testing::opcode_matchers; namespace xla { namespace { -class TupleSimplifierTest : public HloVerifiedTestBase { +class TupleSimplifierTest : public HloTestBase { protected: void Run(HloModule* module, bool change_expected) { TupleSimplifier simplifier; @@ -65,10 +65,10 @@ TEST_F(TupleSimplifierTest, TupleOfParameters) { HloInstruction* param2 = builder.AddInstruction( HloInstruction::CreateParameter(2, scalar_shape_, "param2")); builder.AddInstruction(HloInstruction::CreateTuple({param0, param1, param2})); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); - Run(module, /*change_expected=*/false); + Run(module.get(), /*change_expected=*/false); } TEST_F(TupleSimplifierTest, GteOfTupleOfParameter) { @@ -78,10 +78,10 @@ TEST_F(TupleSimplifierTest, GteOfTupleOfParameter) { HloInstruction::CreateParameter(0, tuple_shape_, "param")); builder.AddInstruction( HloInstruction::CreateGetTupleElement(scalar_shape_, param, 1)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); module->AddEntryComputation(builder.Build()); - Run(module, /*change_expected=*/false); + Run(module.get(), /*change_expected=*/false); } TEST_F(TupleSimplifierTest, GteOfTuple) { @@ -98,12 +98,12 @@ TEST_F(TupleSimplifierTest, GteOfTuple) { HloInstruction* gte = builder.AddInstruction( HloInstruction::CreateGetTupleElement(scalar_shape_, tuple, 1)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), gte); - Run(module, /*change_expected=*/true); + Run(module.get(), /*change_expected=*/true); EXPECT_THAT(computation->root_instruction(), param1); } @@ -125,13 +125,13 @@ TEST_F(TupleSimplifierTest, GteOfTupleChain) { builder.AddInstruction( HloInstruction::CreateUnary(scalar_shape_, HloOpcode::kNegate, element)); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), op::Negate(op::GetTupleElement(op::Tuple()))); - Run(module, /*change_expected=*/true); + Run(module.get(), /*change_expected=*/true); EXPECT_THAT(computation->root_instruction(), op::Negate(op::Parameter())); } @@ -157,12 +157,12 @@ TEST_F(TupleSimplifierTest, NestedGteOfTuples) { ShapeUtil::GetTupleElementShape(element->shape(), 0), element, 0)); } - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), element); - Run(module, /*change_expected=*/true); + Run(module.get(), /*change_expected=*/true); EXPECT_THAT(computation->root_instruction(), param); } @@ -182,12 +182,12 @@ TEST_F(TupleSimplifierTest, TupleOfGteInstructions) { HloInstruction* tuple = builder.AddInstruction(HloInstruction::CreateTuple({gte0, gte1, gte2})); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), tuple); - Run(module, /*change_expected=*/true); + Run(module.get(), /*change_expected=*/true); EXPECT_THAT(computation->root_instruction(), tuple_param); } @@ -207,19 +207,19 @@ TEST_F(TupleSimplifierTest, IncompatibleTuples) { HloInstruction* tuple = builder.AddInstruction(HloInstruction::CreateTuple({gte0, gte1})); - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_THAT(computation->root_instruction(), tuple); - Run(module, /*change_expected=*/false); + Run(module.get(), /*change_expected=*/false); EXPECT_THAT(computation->root_instruction(), tuple); } TEST_F(TupleSimplifierTest, CanExcludeEntryComputation) { // Verify that the root computation can be excluded - auto module = CreateNewModule(); + auto module = CreateNewVerifiedModule(); HloInstruction* p0; HloInstruction* p1; @@ -281,7 +281,7 @@ TEST_F(TupleSimplifierTest, CanExcludeEntryComputation) { entry = module->AddEntryComputation(builder.Build()); } - Run(module, /*change_expected=*/true, /*exclude_entry=*/true); + Run(module.get(), /*change_expected=*/true, /*exclude_entry=*/true); EXPECT_THAT(c0->root_instruction(), p0); EXPECT_THAT(c1->root_instruction(), p1); diff --git a/tensorflow/compiler/xla/service/tuple_util.cc b/tensorflow/compiler/xla/service/tuple_util.cc index cfb0c787d09557fd1aec3517eb9698cfec323369..90ea79ec263a038556ccbd2cd345b337c5a5dcf3 100644 --- a/tensorflow/compiler/xla/service/tuple_util.cc +++ b/tensorflow/compiler/xla/service/tuple_util.cc @@ -21,7 +21,7 @@ namespace xla { /*static*/ HloInstruction* TupleUtil::ExtractPrefix(HloInstruction* input_tuple, int64 elements) { - CHECK(ShapeUtil::IsTuple(input_tuple->shape())); + CHECK(input_tuple->shape().IsTuple()); HloComputation* computation = input_tuple->parent(); const Shape& input_shape = input_tuple->shape(); @@ -41,7 +41,7 @@ namespace xla { /*static*/ HloInstruction* TupleUtil::AppendSuffix( HloInstruction* input_tuple, absl::Span trailing_values) { - CHECK(ShapeUtil::IsTuple(input_tuple->shape())); + CHECK(input_tuple->shape().IsTuple()); HloComputation* computation = input_tuple->parent(); const Shape& input_shape = input_tuple->shape(); diff --git a/tensorflow/compiler/xla/service/while_loop_analysis.cc b/tensorflow/compiler/xla/service/while_loop_analysis.cc index 541b117e0299c94de330604ec5c16e20f07c425f..c93a9ba3176002a34fe84a29e62075de4d19168f 100644 --- a/tensorflow/compiler/xla/service/while_loop_analysis.cc +++ b/tensorflow/compiler/xla/service/while_loop_analysis.cc @@ -15,6 +15,9 @@ limitations under the License. #include "tensorflow/compiler/xla/service/while_loop_analysis.h" #include "tensorflow/compiler/xla/service/hlo_evaluator.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/service/hlo_module.h" +#include "tensorflow/compiler/xla/service/hlo_opcode.h" namespace xla { @@ -229,4 +232,96 @@ optional ComputeWhileLoopTripCount(HloInstruction* while_op, return nullopt; } +// If the only user of this instruction is a get-tuple-element, return that +// get-tuple-element, otherwise return null. If this runs before CSE/DCE, we may +// get a false negative if there are several copies of the same GTE, or there +// are unused GTEs, but we can live with this. +static HloInstruction* GetOnlyGTE(HloInstruction* inst) { + if (inst->user_count() != 1) { + return nullptr; + } + + HloInstruction* user = inst->users().back(); + if (user->opcode() != HloOpcode::kGetTupleElement) { + return nullptr; + } + return user; +} + +optional ComputeWhileLoopTripCountUpperBound(HloInstruction* while_op) { + // If we know the exact trip count, it's also the upper bound. + auto exact_trip_count = ComputeWhileLoopTripCount(while_op); + if (exact_trip_count) { + VLOG(2) << "Loop has exact trip count."; + return exact_trip_count; + } + + // There is one more case we know how to handle. If the loop condition only + // looks at one element of the tuple, and the loop body sets this element to a + // constant, there are two options: + // 1) Evaluating the condition on this constant returns true. In this case, + // the loop either executes 0 times, or is an infinite loop, depending on the + // init value. + // 2) Evaluating the condition on this constant returns false. In this case, + // the loop executes 0 or 1 times, depending on the init value. This means + // that, regardless of the init value, the upper bound on the trip count is 1. + + // Check whether the condition depends on a single parameter, and find out + // which. + auto* while_cond = while_op->while_condition(); + auto* while_cond_param = while_cond->parameter_instruction(0); + auto* cond_gte = GetOnlyGTE(while_cond_param); + if (!cond_gte) { + VLOG(2) << "Induction variable not found in loop condition: " + << while_cond->root_instruction()->ToString(); + return nullopt; + } + + // Now check whether this gets set to a constant by the while body. + auto* while_body = while_op->while_body(); + auto* while_body_root = while_body->root_instruction(); + if (while_body_root->opcode() != HloOpcode::kTuple) { + VLOG(3) << "While body's root is not a tuple instruction: " + << while_body_root->ToString(); + return nullopt; + } + + int64 indvar_index = cond_gte->tuple_index(); + auto* while_body_indvar = while_body_root->operand(indvar_index); + if (while_body_indvar->opcode() != HloOpcode::kConstant) { + VLOG(3) << "While body does not set the IV to a constant: " + << while_body_indvar->ToString(); + return nullopt; + } + + // We have a constant. Evaluate the condition on this constant. + HloEvaluator evaluator(/*max_loop_iterations=*/0); + Literal fake_input = Literal::CreateFromShape(while_cond_param->shape()); + TF_CHECK_OK(fake_input.CopyFrom(while_body_indvar->literal(), + /*dest_shape_index=*/{indvar_index}, + /*src_shape_index=*/{})); + StatusOr eval_result = + evaluator.Evaluate(*while_cond, {std::move(fake_input)}); + + if (!eval_result.ok()) { + VLOG(2) << "Couldn't evaluate while loop condition."; + return nullopt; + } + + Literal cond_result_pred = std::move(eval_result.ValueOrDie()); + CHECK(ShapeUtil::Equal(cond_result_pred.shape(), + ShapeUtil::MakeShape(PRED, {}))); + + // Per the explanation above, if the evaluated condition returns false, the + // loop executes at most once. + bool cond_returns_true = cond_result_pred.GetFirstElement(); + if (!cond_returns_true) { + VLOG(2) << "Upper bound on the trip count is 1"; + return 1; + } + + VLOG(2) << "Loop has no known upper bound on the trip count."; + return nullopt; +} + } // namespace xla diff --git a/tensorflow/compiler/xla/service/while_loop_analysis.h b/tensorflow/compiler/xla/service/while_loop_analysis.h index bf497f4892b95c927379411468a66d8961465413..ac69a727bd6b403672a676400993fb7d8afc0a55 100644 --- a/tensorflow/compiler/xla/service/while_loop_analysis.h +++ b/tensorflow/compiler/xla/service/while_loop_analysis.h @@ -28,6 +28,10 @@ namespace xla { absl::optional ComputeWhileLoopTripCount(HloInstruction *while_op, int64 max_value_returned = 128); +// Returns an upper bound on the trip count of the loop if it's statically +// known, nullopt otherwise. +absl::optional ComputeWhileLoopTripCountUpperBound( + HloInstruction *while_op); } // namespace xla #endif // TENSORFLOW_COMPILER_XLA_SERVICE_WHILE_LOOP_ANALYSIS_H_ diff --git a/tensorflow/compiler/xla/service/while_loop_analysis_test.cc b/tensorflow/compiler/xla/service/while_loop_analysis_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..1da0fbeac89a93eaaef893e5f25dd3b87cc1d5d5 --- /dev/null +++ b/tensorflow/compiler/xla/service/while_loop_analysis_test.cc @@ -0,0 +1,124 @@ +/* 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/compiler/xla/service/while_loop_analysis.h" + +#include "tensorflow/compiler/xla/service/hlo_parser.h" +#include "tensorflow/compiler/xla/test.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" +#include "tensorflow/core/lib/core/status_test_util.h" + +namespace xla { +namespace { + +class WhileLoopAnalysisTest : public HloTestBase {}; + +TEST_F(WhileLoopAnalysisTest, SingleIterationUpperBound) { + const char* const kHloModule = R"( + HloModule ModuleWithWhile + + body { + p_body = (f32[2], s32[]) parameter(0) + val = f32[2] get-tuple-element(p_body), index=0 + const = s32[] constant(-1) + ROOT root = (f32[2], s32[]) tuple(val, const) + } + + condition { + p_cond = (f32[2], s32[]) parameter(0) + gte = s32[] get-tuple-element(p_cond), index=1 + const = s32[] constant(42) + ROOT result = pred[] equal-to(gte, const) + } + + ENTRY entry { + param.0 = f32[2] parameter(0) + param.1 = s32[] parameter(1) + while_init = (f32[2], s32[]) tuple(param.0, param.1) + ROOT while = (f32[2], s32[]) while(while_init), condition=condition, body=body + })"; + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(kHloModule)); + + HloInstruction* while_op = module->entry_computation()->root_instruction(); + EXPECT_EQ(*ComputeWhileLoopTripCountUpperBound(while_op), 1); +} + +TEST_F(WhileLoopAnalysisTest, NoUpperBound) { + const char* const kHloModule = R"( + HloModule ModuleWithWhile + + body { + p_body = (f32[2], s32[]) parameter(0) + val = f32[2] get-tuple-element(p_body), index=0 + const = s32[] constant(42) + ROOT root = (f32[2], s32[]) tuple(val, const) + } + + condition { + p_cond = (f32[2], s32[]) parameter(0) + gte = s32[] get-tuple-element(p_cond), index=1 + const = s32[] constant(42) + ROOT result = pred[] equal-to(gte, const) + } + + ENTRY entry { + param.0 = f32[2] parameter(0) + param.1 = s32[] parameter(1) + while_init = (f32[2], s32[]) tuple(param.0, param.1) + ROOT while = (f32[2], s32[]) while(while_init), condition=condition, body=body + })"; + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(kHloModule)); + + HloInstruction* while_op = module->entry_computation()->root_instruction(); + EXPECT_EQ(ComputeWhileLoopTripCountUpperBound(while_op), absl::nullopt); +} + +TEST_F(WhileLoopAnalysisTest, ExactBound) { + const char* const kHloModule = R"( + HloModule ModuleWithWhile + + body { + p_body = (f32[2], s32[]) parameter(0) + val = f32[2] get-tuple-element(p_body), index=0 + index = s32[] get-tuple-element(p_body), index=1 + one = s32[] constant(1) + inc = s32[] add(index, one) + ROOT root = (f32[2], s32[]) tuple(val, inc) + } + + condition { + p_cond = (f32[2], s32[]) parameter(0) + gte = s32[] get-tuple-element(p_cond), index=1 + const = s32[] constant(42) + ROOT result = pred[] less-than(gte, const) + } + + ENTRY entry { + param.0 = f32[2] parameter(0) + param.1 = s32[] constant(0) + while_init = (f32[2], s32[]) tuple(param.0, param.1) + ROOT while = (f32[2], s32[]) while(while_init), condition=condition, body=body + })"; + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(kHloModule)); + + HloInstruction* while_op = module->entry_computation()->root_instruction(); + EXPECT_EQ(*ComputeWhileLoopTripCountUpperBound(while_op), 42); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/service/while_loop_constant_sinking.cc b/tensorflow/compiler/xla/service/while_loop_constant_sinking.cc index 067cfcc17d65860a249de4d9e31703df12091d3a..8b381dec07397c1427e98bc30511ac21dc577610 100644 --- a/tensorflow/compiler/xla/service/while_loop_constant_sinking.cc +++ b/tensorflow/compiler/xla/service/while_loop_constant_sinking.cc @@ -46,8 +46,9 @@ static Status ReplaceUsesWhileKeepingLoopInvariance( return Status::OK(); } -StatusOr WhileLoopConstantSinking::TrySinkingConstantsIntoWhileBody( +StatusOr WhileLoopConstantSinking::TrySinkingConstantsIntoWhileLoop( HloInstruction* while_instr) { + HloComputation* while_cond = while_instr->while_condition(); HloComputation* while_body = while_instr->while_body(); const HloInstruction& init_value = *while_instr->operand(0); @@ -57,24 +58,48 @@ StatusOr WhileLoopConstantSinking::TrySinkingConstantsIntoWhileBody( bool changed = false; - for (HloInstruction* invariant_gte : - WhileUtil::GetInvariantGTEsForWhileBody(*while_body)) { - int64 index = invariant_gte->tuple_index(); + absl::flat_hash_map> + conditional_gte_index_to_insts = + WhileUtil::GetGTEsMapForWhileConditional(*while_cond); + std::vector invariant_body_gtes = + WhileUtil::GetInvariantGTEsForWhileBody(*while_body); + + for (HloInstruction* invariant_body_gte : invariant_body_gtes) { + int64 index = invariant_body_gte->tuple_index(); const HloInstruction& invariant_value = *init_value.operand(index); - // Should have at least one user that's not while_body_root. - if (invariant_gte->user_count() <= 1) { + // Original value should be a constant. + if (invariant_value.opcode() != HloOpcode::kConstant) { continue; } - if (invariant_value.opcode() == HloOpcode::kConstant) { - auto* constant_instr = + // Sink into the while_body. + // Should have at least one user that's not while_body_root. + if (invariant_body_gte->user_count() > 1) { + HloInstruction* constant_instr = while_body->AddInstruction(invariant_value.Clone(/*suffix=*/".sunk")); TF_RETURN_IF_ERROR(ReplaceUsesWhileKeepingLoopInvariance( - invariant_gte, constant_instr, while_body->root_instruction(), + invariant_body_gte, constant_instr, while_body->root_instruction(), index)); changed = true; } + + // Check if there is a corresponding GTE in while_conditional. + auto it = conditional_gte_index_to_insts.find(index); + if (it == conditional_gte_index_to_insts.end()) { + continue; + } + + for (HloInstruction* invariant_cond_gte : it->second) { + // Should have at least one user. + if (invariant_cond_gte->user_count() > 0) { + HloInstruction* constant_instr = while_cond->AddInstruction( + invariant_value.Clone(/*suffix=*/".sunk")); + TF_RETURN_IF_ERROR( + invariant_cond_gte->ReplaceAllUsesWith(constant_instr)); + changed = true; + } + } } return changed; @@ -115,10 +140,8 @@ StatusOr WhileLoopConstantSinking::Run(HloModule* module) { } for (HloInstruction* while_instr : while_instrs) { - // We only sink into while loop bodies, but this can be extended to - // transform conditions as well. TF_ASSIGN_OR_RETURN(bool result, - TrySinkingConstantsIntoWhileBody(while_instr)); + TrySinkingConstantsIntoWhileLoop(while_instr)); changed |= result; } diff --git a/tensorflow/compiler/xla/service/while_loop_constant_sinking.h b/tensorflow/compiler/xla/service/while_loop_constant_sinking.h index 577bad6c7062d2ee40271e407e8eed7655fa13bf..a866bc1264b4013bb7530b5e02b546e6f78d676b 100644 --- a/tensorflow/compiler/xla/service/while_loop_constant_sinking.h +++ b/tensorflow/compiler/xla/service/while_loop_constant_sinking.h @@ -23,8 +23,8 @@ limitations under the License. namespace xla { // Sinks while loop invariant values that happen to be constants into the while -// loop body. This is probably not a win in isolation but may unlock further -// optimizations like constant folding. +// loop body and conditional. This is probably not a win in isolation but may +// unlock further optimizations like constant folding. // // state = (..., const, ...) // while (pred(state)) { @@ -46,22 +46,19 @@ namespace xla { // tuple trivially loop invariant. WhileLoopSimplifier will later get rid of // `v`. // -// We only sink into while loop bodies, but this can be extended to transform -// conditions as well. -// // TODO(b/79121449): We should also sink broadcasts of constants. class WhileLoopConstantSinking : public HloModulePass { public: ~WhileLoopConstantSinking() override = default; absl::string_view name() const override { - return "while-loop-invariant-code-motion"; + return "while-loop-constant-sinking"; } StatusOr Run(HloModule* module) override; private: - StatusOr TrySinkingConstantsIntoWhileBody(HloInstruction* while_instr); + StatusOr TrySinkingConstantsIntoWhileLoop(HloInstruction* while_instr); }; } // namespace xla diff --git a/tensorflow/compiler/xla/service/while_loop_constant_sinking_test.cc b/tensorflow/compiler/xla/service/while_loop_constant_sinking_test.cc index 0e7667de832c54f647d071e3c9563091d0f994aa..3bcf5c38309a86e9e3cab3268f3f065005f7a923 100644 --- a/tensorflow/compiler/xla/service/while_loop_constant_sinking_test.cc +++ b/tensorflow/compiler/xla/service/while_loop_constant_sinking_test.cc @@ -114,7 +114,7 @@ HloModule ModuleWithWhile body { p_b = (f32[2],(f32[2],f32[2])) parameter(0) - p_b.0 = f32[2] get-tuple-element((f32[2],f32[2],f32[2]) p_b), index=0 + p_b.0 = f32[2] get-tuple-element((f32[2],(f32[2],f32[2])) p_b), index=0 p_b.1 = (f32[2],f32[2]) get-tuple-element((f32[2],(f32[2],f32[2])) p_b), index=1 p_b.1.1 = f32[2] get-tuple-element(p_b.1), index=0 @@ -129,7 +129,7 @@ condition { ENTRY entry { const_0 = f32[2] constant({1, 2}) - const_1 = (f32[2], f32[2]) constant((f32[2], f32[2]) ({2, 1},{3,1})) + const_1 = (f32[2], f32[2]) constant(({2, 1},{3,1})) while_init = (f32[2],(f32[2],f32[2])) tuple(const_0, const_1) ROOT while = (f32[2],(f32[2],f32[2])) while(while_init), condition=condition, body=body } @@ -206,8 +206,8 @@ body { p_body.0 = f32[2] get-tuple-element((f32[2],f32[2]) p_body), index=0 p_body.1 = f32[2] get-tuple-element((f32[2],f32[2]) p_body), index=1 - token = token[] after-all() - outfeed = token[] outfeed(p_body.0, token) + token0 = token[] after-all() + outfeed = token[] outfeed(p_body.0, token0) ROOT root = (f32[2],f32[2],f32[2]) tuple(p_body.0, p_body.1, p_body.1) } @@ -242,5 +242,178 @@ ENTRY entry { } } } + +TEST_F(WhileLoopConstantSinkingTest, ConditionalSinkConstant) { + const char* const hlo_string = R"( +HloModule ModuleWithWhile + +body { + p_body = (f32[],f32[]) parameter(0) + p_body.0 = f32[] get-tuple-element((f32[],f32[]) p_body), index=0 + const = f32[] constant(1) + add = f32[] add(p_body.0, const) + p_body.1 = f32[] get-tuple-element((f32[],f32[]) p_body), index=1 + ROOT root = (f32[],f32[]) tuple(add, p_body.1) +} + +condition { + p_cond = (f32[],f32[]) parameter(0) + p_cond.0 = f32[] get-tuple-element((f32[],f32[]) p_cond), index=0 + p_cond.1 = f32[] get-tuple-element((f32[],f32[]) p_cond), index=1 + ROOT result = pred[] less-than(p_cond.0, p_cond.1) +} + +ENTRY entry { + const_0 = f32[] constant(0) + const_1 = f32[] constant(10) + while_init = (f32[],f32[]) tuple(const_0, const_1) + ROOT while = (f32[],f32[]) while(while_init), condition=condition, body=body +} +)"; + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseHloString(hlo_string)); + + TF_ASSERT_OK_AND_ASSIGN(bool changed, + WhileLoopConstantSinking{}.Run(module.get())); + ASSERT_TRUE(changed); + + auto* while_condition = module->GetComputationWithName("condition"); + EXPECT_THAT(while_condition->root_instruction(), op::Lt(_, op::Constant())); +} + +TEST_F(WhileLoopConstantSinkingTest, ConditionalTupleShapedConstants) { + const char* const hlo_string = R"( +HloModule ModuleWithWhile + +body { + p_b = (f32[],(f32[],f32[])) parameter(0) + p_b.0 = f32[] get-tuple-element((f32[],(f32[],f32[])) p_b), index=0 + p_b.1 = (f32[],f32[]) get-tuple-element((f32[],(f32[],f32[])) p_b), index=1 + p_b.1.0 = f32[] get-tuple-element((f32[],f32[]) p_b.1), index=0 + add = f32[] add(p_b.0, p_b.1.0) + ROOT root = (f32[],(f32[],f32[])) tuple(add, p_b.1) +} + +condition { + p_c = (f32[],(f32[],f32[])) parameter(0) + p_c.0 = f32[] get-tuple-element((f32[],(f32[],f32[])) p_c), index=0 + p_c.1 = (f32[],f32[]) get-tuple-element((f32[],(f32[],f32[])) p_c), index=1 + p_c.1.1 = f32[] get-tuple-element((f32[],f32[]) p_c.1), index=1 + ROOT result = pred[] less-than(p_c.0, p_c.1.1) +} + +ENTRY entry { + const_0 = f32[] constant(0) + const_1 = (f32[], f32[]) constant((1, 10)) + while_init = (f32[],(f32[],f32[])) tuple(const_0, const_1) + ROOT while = (f32[],(f32[],f32[])) while(while_init), condition=condition, body=body +} +)"; + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseHloString(hlo_string)); + + TF_ASSERT_OK_AND_ASSIGN(bool changed, + WhileLoopConstantSinking{}.Run(module.get())); + ASSERT_TRUE(changed); + + auto* while_condition = module->GetComputationWithName("condition"); + EXPECT_THAT(while_condition->root_instruction(), + op::Lt(_, op::GetTupleElement(op::Constant()))); +} + +TEST_F(WhileLoopConstantSinkingTest, ConditionalDontCreateDeadConstant) { + const char* const hlo_string = R"( +HloModule ModuleWithWhile + +body { + p_body = (f32[],f32[],f32[]) parameter(0) + p_body.0 = f32[] get-tuple-element((f32[],f32[],f32[]) p_body), index=0 + const = f32[] constant(1) + add = f32[] add(p_body.0, const) + p_body.1 = f32[] get-tuple-element((f32[],f32[],f32[]) p_body), index=1 + p_body.2 = f32[] get-tuple-element((f32[],f32[],f32[]) p_body), index=2 + ROOT root = (f32[],f32[],f32[]) tuple(add, p_body.1, p_body.2) +} + +condition { + p_cond = (f32[],f32[],f32[]) parameter(0) + p_cond.0 = f32[] get-tuple-element((f32[],f32[],f32[]) p_cond), index=0 + p_cond.1 = f32[] get-tuple-element((f32[],f32[],f32[]) p_cond), index=1 + p_cond.2 = f32[] get-tuple-element((f32[],f32[],f32[]) p_cond), index=2 + ROOT result = pred[] less-than(p_cond.0, p_cond.1) +} + +ENTRY entry { + const_0 = f32[] constant(0) + const_1 = f32[] constant(10) + const_2 = f32[] constant(12) + while_init = (f32[],f32[],f32[]) tuple(const_0, const_1, const_2) + ROOT while = (f32[],f32[],f32[]) while(while_init), condition=condition, body=body +} +)"; + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseHloString(hlo_string)); + + TF_ASSERT_OK_AND_ASSIGN(bool changed, + WhileLoopConstantSinking{}.Run(module.get())); + ASSERT_TRUE(changed); + + auto* while_condition = module->GetComputationWithName("condition"); + EXPECT_THAT(while_condition->root_instruction(), op::Lt(_, op::Constant())); + for (const HloInstruction* inst : while_condition->instructions()) { + if (inst->opcode() == HloOpcode::kConstant) { + EXPECT_GT(inst->user_count(), 0); + } + } +} + +TEST_F(WhileLoopConstantSinkingTest, ConditionalMultipleSameIndexGTEs) { + const char* const hlo_string = R"( +HloModule ModuleWithWhile + +body { + p_body = (f32[],f32[],f32[]) parameter(0) + p_body.0 = f32[] get-tuple-element((f32[],f32[],f32[]) p_body), index=0 + const = f32[] constant(1) + add.0 = f32[] add(p_body.0, const) + p_body.1 = f32[] get-tuple-element((f32[],f32[],f32[]) p_body), index=1 + add.1 = f32[] add(p_body.1, const) + p_body.2 = f32[] get-tuple-element((f32[],f32[],f32[]) p_body), index=2 + ROOT root = (f32[],f32[],f32[]) tuple(add.0, add.1, p_body.2) +} + +condition { + p_cond = (f32[],f32[],f32[]) parameter(0) + p_cond.0 = f32[] get-tuple-element((f32[],f32[],f32[]) p_cond), index=0 + p_cond.2 = f32[] get-tuple-element((f32[],f32[],f32[]) p_cond), index=2 + lt.0 = pred[] less-than(p_cond.0, p_cond.2) + p_cond.1 = f32[] get-tuple-element((f32[],f32[],f32[]) p_cond), index=1 + p_cond.2.c = f32[] get-tuple-element((f32[],f32[],f32[]) p_cond), index=2 + lt.1 = pred[] less-than(p_cond.1, p_cond.2.c) + ROOT result = pred[] and(lt.0, lt.1) +} + +ENTRY entry { + const_0 = f32[] constant(0) + const_1 = f32[] constant(0) + const_2 = f32[] constant(12) + while_init = (f32[],f32[],f32[]) tuple(const_0, const_1, const_2) + ROOT while = (f32[],f32[],f32[]) while(while_init), condition=condition, body=body +} +)"; + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseHloString(hlo_string)); + TF_ASSERT_OK_AND_ASSIGN(bool changed, + WhileLoopConstantSinking{}.Run(module.get())); + ASSERT_TRUE(changed); + + auto* while_condition = module->GetComputationWithName("condition"); + EXPECT_THAT(while_condition->root_instruction(), + op::And(op::Lt(_, op::Constant()), op::Lt(_, op::Constant()))); +} } // namespace } // namespace xla 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 9795b2830b6d9add82b89ac76b5438ddc3d2bfe8..69cc8feb3f31ad782b9d3437d81d0ab8ce10aadb 100644 --- a/tensorflow/compiler/xla/service/while_loop_invariant_code_motion.cc +++ b/tensorflow/compiler/xla/service/while_loop_invariant_code_motion.cc @@ -19,7 +19,9 @@ limitations under the License. #include "absl/container/flat_hash_set.h" #include "absl/container/inlined_vector.h" #include "tensorflow/compiler/xla/service/tuple_util.h" +#include "tensorflow/compiler/xla/service/while_loop_analysis.h" #include "tensorflow/compiler/xla/service/while_util.h" +#include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/util.h" namespace xla { @@ -87,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; } @@ -125,7 +127,7 @@ WhileLoopInvariantCodeMotion::TryHoistingInvariantInstructionsFromWhileBody( HloInstruction* while_instr) { auto print_no_metadata = HloPrintOptions{}.set_print_metadata(false); - if (!ShapeUtil::IsTuple(while_instr->shape())) { + if (!while_instr->shape().IsTuple()) { // This restriction leaves one interesting pattern on the table: // // while_body(f32[1024, 1024] %param) { @@ -143,6 +145,12 @@ WhileLoopInvariantCodeMotion::TryHoistingInvariantInstructionsFromWhileBody( string while_instr_name = while_instr->ToString(print_no_metadata); VLOG(2) << "Trying to hoist from " << while_instr_name; + auto maybe_upper_bound = ComputeWhileLoopTripCountUpperBound(while_instr); + if (maybe_upper_bound && *maybe_upper_bound <= 1) { + VLOG(2) << "Loop has a trip count of at most 1, skipping."; + return false; + } + HloComputation* while_body = while_instr->while_body(); // Maps instructions in the while body to instructions hoisted outside the @@ -160,7 +168,7 @@ WhileLoopInvariantCodeMotion::TryHoistingInvariantInstructionsFromWhileBody( // is no benefit to hoisting them unless something that uses it is also // hoisted. for (auto* instr : WhileUtil::GetInvariantGTEsForWhileBody(*while_body)) { - if (ShapeUtil::IsArray(instr->shape())) { + if (instr->shape().IsArray()) { // TODO(b/79147885): We should try to generalize this to tuples for // uniformity's sake, if nothing else. InsertOrDie(&unhoisted_invariant_instructions, instr); @@ -180,6 +188,13 @@ WhileLoopInvariantCodeMotion::TryHoistingInvariantInstructionsFromWhileBody( return false; } + // LICM in the presence of domain instructions is complex, bail. + for (auto* instruction : while_body->MakeInstructionPostOrder()) { + if (instruction->opcode() == HloOpcode::kDomain) { + return false; + } + } + // instructions_to_replace[i] is hoisted into a loop invariant instruction // replacement_instructions[i]. std::vector instructions_to_replace; @@ -193,9 +208,40 @@ WhileLoopInvariantCodeMotion::TryHoistingInvariantInstructionsFromWhileBody( continue; } + if (!hoist_size_inflating_ops_) { + // Check that hoisting the instruction doesn't cause a significant memory + // blow-up. LICM extends the live-range of the output of the hoisted + // instruction to be the entire while loop, which may be problematic on + // platforms where memory is limited. This can be especially harmful if + // the instruction has a significantly larger output than its input, e.g. + // kIota, kBroadcast or kConstant. + int64 input_size = 0, output_size = 0; + + for (auto* operand : instruction->operands()) { + ShapeUtil::ForEachSubshape( + operand->shape(), + [&input_size](const Shape& subshape, const ShapeIndex& /*index*/) { + if (subshape.IsArray()) { + input_size += ShapeUtil::ByteSizeOfElements(subshape); + } + }); + } + ShapeUtil::ForEachSubshape( + instruction->shape(), + [&output_size](const Shape& subshape, const ShapeIndex& /*index*/) { + if (subshape.IsArray()) { + output_size += ShapeUtil::ByteSizeOfElements(subshape); + } + }); + + if (output_size > input_size) { + continue; + } + } + 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.h b/tensorflow/compiler/xla/service/while_loop_invariant_code_motion.h index 3031899f71e0fd77f20448d9d7489798af01615c..bd6232dc0a988775a0490abbf6125daad8476295 100644 --- a/tensorflow/compiler/xla/service/while_loop_invariant_code_motion.h +++ b/tensorflow/compiler/xla/service/while_loop_invariant_code_motion.h @@ -34,8 +34,14 @@ class WhileLoopInvariantCodeMotion : public HloModulePass { // Setting `hoist_constants` to false can be help if LICM is run in the mid // level HLO pipeline because hoisting constants out of while loop bodies can // break optimizations like constant folding. - explicit WhileLoopInvariantCodeMotion(bool hoist_constants = false) - : hoist_constants_(hoist_constants) {} + // Setting `hoist_size_inflating_ops` to false will forbid hoisting + // instructions where the size of the output(s) is larger than the size of the + // input(s). This is useful on platforms on which it's important to prevent + // blow-ups in memory size. + explicit WhileLoopInvariantCodeMotion(bool hoist_constants = false, + bool hoist_size_inflating_ops = true) + : hoist_constants_(hoist_constants), + hoist_size_inflating_ops_(hoist_size_inflating_ops) {} ~WhileLoopInvariantCodeMotion() override = default; absl::string_view name() const override { @@ -49,6 +55,7 @@ class WhileLoopInvariantCodeMotion : public HloModulePass { HloInstruction* while_instr); bool hoist_constants_; + bool hoist_size_inflating_ops_; }; } // namespace xla 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 32e69c335b713c438bd7fcb2053709b0624f58ed..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 @@ -18,7 +18,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_matchers.h" #include "tensorflow/compiler/xla/service/hlo_parser.h" #include "tensorflow/compiler/xla/test.h" -#include "tensorflow/compiler/xla/tests/hlo_verified_test_base.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" #include "tensorflow/core/lib/core/status_test_util.h" namespace xla { @@ -26,7 +26,7 @@ namespace { namespace op = xla::testing::opcode_matchers; -class WhileLoopInvariantCodeMotionTest : public HloVerifiedTestBase { +class WhileLoopInvariantCodeMotionTest : public HloTestBase { public: // Makes a computation which has one parameter, of the given shape, and always // returns PRED[]{true}. This is useful as a dummy loop condition. @@ -58,6 +58,7 @@ HloComputation* WhileLoopInvariantCodeMotionTest::MakeAlwaysTrueComputation( } TEST_F(WhileLoopInvariantCodeMotionTest, HoistOneInvariantOperation) { + auto m = CreateNewVerifiedModule(); auto scalar_s32 = ShapeUtil::MakeShape(S32, {}); Shape while_shape = ShapeUtil::MakeTupleShape({scalar_s32, scalar_s32, scalar_s32}); @@ -76,19 +77,18 @@ TEST_F(WhileLoopInvariantCodeMotionTest, HoistOneInvariantOperation) { builder.AddInstruction( HloInstruction::CreateTuple({gte_0, gte_1, add_result})); - return module().AddEmbeddedComputation(builder.Build()); + return m->AddEmbeddedComputation(builder.Build()); }(); HloComputation::Builder builder(TestName()); auto* init_value = builder.AddInstruction( HloInstruction::CreateParameter(0, while_shape, "init_value")); builder.AddInstruction(HloInstruction::CreateWhile( - while_shape, MakeAlwaysTrueComputation(while_shape, &module()), - while_body, init_value)); - HloComputation* entry_computation = - module().AddEntryComputation(builder.Build()); + while_shape, MakeAlwaysTrueComputation(while_shape, m.get()), while_body, + init_value)); + HloComputation* entry_computation = m->AddEntryComputation(builder.Build()); TF_ASSERT_OK_AND_ASSIGN(bool simplified_loop, - WhileLoopInvariantCodeMotion{}.Run(&module())); + WhileLoopInvariantCodeMotion{}.Run(m.get())); EXPECT_TRUE(simplified_loop); HloInstruction* transformed_while; @@ -100,6 +100,7 @@ TEST_F(WhileLoopInvariantCodeMotionTest, HoistOneInvariantOperation) { } TEST_F(WhileLoopInvariantCodeMotionTest, HoistInvariantOperationTree) { + auto m = CreateNewVerifiedModule(); auto scalar_s32 = ShapeUtil::MakeShape(S32, {}); Shape while_shape = ShapeUtil::MakeTupleShape({scalar_s32, scalar_s32, scalar_s32}); @@ -135,19 +136,18 @@ TEST_F(WhileLoopInvariantCodeMotionTest, HoistInvariantOperationTree) { builder.AddInstruction( HloInstruction::CreateTuple({gte_0, gte_1, divide_result})); - return module().AddEmbeddedComputation(builder.Build()); + return m->AddEmbeddedComputation(builder.Build()); }(); HloComputation::Builder builder(TestName()); auto* init_value = builder.AddInstruction( HloInstruction::CreateParameter(0, while_shape, "init_value")); builder.AddInstruction(HloInstruction::CreateWhile( - while_shape, MakeAlwaysTrueComputation(while_shape, &module()), - while_body, init_value)); - HloComputation* entry_computation = - module().AddEntryComputation(builder.Build()); + while_shape, MakeAlwaysTrueComputation(while_shape, m.get()), while_body, + init_value)); + HloComputation* entry_computation = m->AddEntryComputation(builder.Build()); TF_ASSERT_OK_AND_ASSIGN(bool simplified_loop, - WhileLoopInvariantCodeMotion{}.Run(&module())); + WhileLoopInvariantCodeMotion{}.Run(m.get())); EXPECT_TRUE(simplified_loop); HloInstruction* transformed_while; @@ -173,6 +173,7 @@ TEST_F(WhileLoopInvariantCodeMotionTest, HoistInvariantOperationTree) { TEST_F(WhileLoopInvariantCodeMotionTest, DontHoistTriviallyLoopVaryingComputation) { // Basic negative test: the add expression is not loop invariant. + auto m = CreateNewVerifiedModule(); auto scalar_s32 = ShapeUtil::MakeShape(S32, {}); Shape while_shape = ShapeUtil::MakeTupleShape({scalar_s32, scalar_s32}); @@ -189,20 +190,20 @@ TEST_F(WhileLoopInvariantCodeMotionTest, scalar_s32, HloOpcode::kAdd, gte_0, gte_1)); builder.AddInstruction(HloInstruction::CreateTuple({gte_0, add_result})); - return module().AddEmbeddedComputation(builder.Build()); + return m->AddEmbeddedComputation(builder.Build()); }(); HloComputation::Builder builder(TestName()); auto* init_value = builder.AddInstruction( HloInstruction::CreateParameter(0, while_shape, "init_value")); auto* while_inst = builder.AddInstruction(HloInstruction::CreateWhile( - while_shape, MakeAlwaysTrueComputation(while_shape, &module()), - while_body, init_value)); + while_shape, MakeAlwaysTrueComputation(while_shape, m.get()), while_body, + init_value)); - module().AddEntryComputation(builder.Build()); + m->AddEntryComputation(builder.Build()); TF_ASSERT_OK_AND_ASSIGN(bool simplified_loop, - WhileLoopInvariantCodeMotion{}.Run(&module())); + WhileLoopInvariantCodeMotion{}.Run(m.get())); EXPECT_FALSE(simplified_loop); EXPECT_THAT(while_inst->while_body()->instructions(), Contains(op::Add())); @@ -210,6 +211,7 @@ TEST_F(WhileLoopInvariantCodeMotionTest, TEST_F(WhileLoopInvariantCodeMotionTest, DontHoistLoopVaryingComputationWithAlternatingTuples) { + auto m = CreateNewVerifiedModule(); auto scalar_s32 = ShapeUtil::MakeShape(S32, {}); Shape while_shape = ShapeUtil::MakeTupleShape({scalar_s32, scalar_s32, scalar_s32}); @@ -228,25 +230,26 @@ TEST_F(WhileLoopInvariantCodeMotionTest, builder.AddInstruction( HloInstruction::CreateTuple({gte_1, gte_0, add_result})); - return module().AddEmbeddedComputation(builder.Build()); + return m->AddEmbeddedComputation(builder.Build()); }(); HloComputation::Builder builder(TestName()); auto* init_value = builder.AddInstruction( HloInstruction::CreateParameter(0, while_shape, "init_value")); auto* while_inst = builder.AddInstruction(HloInstruction::CreateWhile( - while_shape, MakeAlwaysTrueComputation(while_shape, &module()), - while_body, init_value)); + while_shape, MakeAlwaysTrueComputation(while_shape, m.get()), while_body, + init_value)); - module().AddEntryComputation(builder.Build()); + m->AddEntryComputation(builder.Build()); TF_ASSERT_OK_AND_ASSIGN(bool simplified_loop, - WhileLoopInvariantCodeMotion{}.Run(&module())); + WhileLoopInvariantCodeMotion{}.Run(m.get())); EXPECT_FALSE(simplified_loop); EXPECT_THAT(while_inst->while_body()->instructions(), Contains(op::Add())); } TEST_F(WhileLoopInvariantCodeMotionTest, DontHoistInstructionWithSideEffects) { + auto m = CreateNewVerifiedModule(); auto scalar_s32 = ShapeUtil::MakeShape(S32, {}); auto token_shape = ShapeUtil::MakeTokenShape(); Shape while_shape = @@ -267,7 +270,7 @@ TEST_F(WhileLoopInvariantCodeMotionTest, DontHoistInstructionWithSideEffects) { builder.AddInstruction( HloInstruction::CreateTuple({gte_0, gte_1, out_token})); - return module().AddEmbeddedComputation(builder.Build()); + return m->AddEmbeddedComputation(builder.Build()); }(); HloComputation::Builder builder(TestName()); @@ -277,14 +280,14 @@ TEST_F(WhileLoopInvariantCodeMotionTest, DontHoistInstructionWithSideEffects) { auto* init_value = builder.AddInstruction( HloInstruction::CreateTuple({scalar_param, scalar_param, token})); auto* while_inst = builder.AddInstruction(HloInstruction::CreateWhile( - while_shape, MakeAlwaysTrueComputation(while_shape, &module()), - while_body, init_value)); + while_shape, MakeAlwaysTrueComputation(while_shape, m.get()), while_body, + init_value)); builder.AddInstruction( HloInstruction::CreateGetTupleElement(scalar_s32, while_inst, 0)); - module().AddEntryComputation(builder.Build()); + m->AddEntryComputation(builder.Build()); TF_ASSERT_OK_AND_ASSIGN(bool simplified_loop, - WhileLoopInvariantCodeMotion{}.Run(&module())); + WhileLoopInvariantCodeMotion{}.Run(m.get())); ASSERT_FALSE(simplified_loop); EXPECT_THAT(while_inst->while_body()->instructions(), @@ -294,8 +297,9 @@ TEST_F(WhileLoopInvariantCodeMotionTest, DontHoistInstructionWithSideEffects) { TEST_F(WhileLoopInvariantCodeMotionTest, DontHoistBitcastAlone) { // The bitcast's user, an outfeed, can't be hoisted, so don't hoist the // 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}); @@ -310,14 +314,16 @@ 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})); - return module().AddEmbeddedComputation(builder.Build()); + return m->AddEmbeddedComputation(builder.Build()); }(); HloComputation::Builder builder(TestName()); @@ -327,15 +333,15 @@ TEST_F(WhileLoopInvariantCodeMotionTest, DontHoistBitcastAlone) { auto* init_value = builder.AddInstruction( HloInstruction::CreateTuple({scalar_param, scalar_param, token})); auto* while_inst = builder.AddInstruction(HloInstruction::CreateWhile( - while_shape, MakeAlwaysTrueComputation(while_shape, &module()), - while_body, init_value)); + while_shape, MakeAlwaysTrueComputation(while_shape, m.get()), while_body, + init_value)); builder.AddInstruction( HloInstruction::CreateGetTupleElement(scalar_s32, while_inst, 0)); - module().AddEntryComputation(builder.Build()); + m->AddEntryComputation(builder.Build()); TF_ASSERT_OK_AND_ASSIGN(bool simplified_loop, - WhileLoopInvariantCodeMotion{}.Run(&module())); + WhileLoopInvariantCodeMotion{}.Run(m.get())); EXPECT_FALSE(simplified_loop); EXPECT_THAT(while_inst->while_body()->instructions(), @@ -346,10 +352,11 @@ TEST_F(WhileLoopInvariantCodeMotionTest, DontHoistBitcastAlone) { 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"); @@ -358,30 +365,30 @@ 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})); - return module().AddEmbeddedComputation(builder.Build()); + return m->AddEmbeddedComputation(builder.Build()); }(); HloComputation::Builder builder(TestName()); auto* init_value = builder.AddInstruction( HloInstruction::CreateParameter(0, while_shape, "init_value")); builder.AddInstruction(HloInstruction::CreateWhile( - while_shape, MakeAlwaysTrueComputation(while_shape, &module()), - while_body, init_value)); + while_shape, MakeAlwaysTrueComputation(while_shape, m.get()), while_body, + init_value)); - HloComputation* entry_computation = - module().AddEntryComputation(builder.Build()); + HloComputation* entry_computation = m->AddEntryComputation(builder.Build()); TF_ASSERT_OK_AND_ASSIGN(bool simplified_loop, - WhileLoopInvariantCodeMotion{}.Run(&module())); + WhileLoopInvariantCodeMotion{}.Run(m.get())); EXPECT_TRUE(simplified_loop); HloInstruction* transformed_while; @@ -396,6 +403,7 @@ TEST_F(WhileLoopInvariantCodeMotionTest, HoistBitcastIfNeeded) { } TEST_F(WhileLoopInvariantCodeMotionTest, DontHoistControlDependencies) { + auto m = CreateNewVerifiedModule(); auto scalar_s32 = ShapeUtil::MakeShape(S32, {}); Shape while_shape = ShapeUtil::MakeTupleShape({scalar_s32, scalar_s32, scalar_s32}); @@ -416,22 +424,23 @@ TEST_F(WhileLoopInvariantCodeMotionTest, DontHoistControlDependencies) { builder.AddInstruction( HloInstruction::CreateTuple({gte_0, gte_1, add_result})); - while_body = module().AddEmbeddedComputation(builder.Build()); + while_body = m->AddEmbeddedComputation(builder.Build()); } HloComputation::Builder builder(TestName()); auto* init_value = builder.AddInstruction( HloInstruction::CreateParameter(0, while_shape, "init_value")); builder.AddInstruction(HloInstruction::CreateWhile( - while_shape, MakeAlwaysTrueComputation(while_shape, &module()), - while_body, init_value)); - module().AddEntryComputation(builder.Build()); + while_shape, MakeAlwaysTrueComputation(while_shape, m.get()), while_body, + init_value)); + m->AddEntryComputation(builder.Build()); TF_ASSERT_OK_AND_ASSIGN(bool simplified_loop, - WhileLoopInvariantCodeMotion{}.Run(&module())); + WhileLoopInvariantCodeMotion{}.Run(m.get())); EXPECT_FALSE(simplified_loop); } TEST_F(WhileLoopInvariantCodeMotionTest, BodyHasNonTupleRoot) { + auto m = CreateNewVerifiedModule(); auto scalar_s32 = ShapeUtil::MakeShape(S32, {}); Shape while_shape = ShapeUtil::MakeTupleShape({scalar_s32, scalar_s32}); @@ -439,7 +448,7 @@ TEST_F(WhileLoopInvariantCodeMotionTest, BodyHasNonTupleRoot) { HloComputation::Builder builder(TestName() + ".passthrough"); HloInstruction* param = builder.AddInstruction( HloInstruction::CreateParameter(0, while_shape, "param")); - HloComputation* result = module().AddEmbeddedComputation(builder.Build()); + HloComputation* result = m->AddEmbeddedComputation(builder.Build()); result->AddInstruction( HloInstruction::CreateGetTupleElement(scalar_s32, param, 1)); @@ -450,11 +459,11 @@ TEST_F(WhileLoopInvariantCodeMotionTest, BodyHasNonTupleRoot) { auto* init_value = builder.AddInstruction( HloInstruction::CreateParameter(0, while_shape, "init_value")); builder.AddInstruction(HloInstruction::CreateWhile( - while_shape, MakeAlwaysTrueComputation(while_shape, &module()), - while_body, init_value)); - module().AddEntryComputation(builder.Build()); + while_shape, MakeAlwaysTrueComputation(while_shape, m.get()), while_body, + init_value)); + m->AddEntryComputation(builder.Build()); TF_ASSERT_OK_AND_ASSIGN(bool simplified_loop, - WhileLoopInvariantCodeMotion{}.Run(&module())); + WhileLoopInvariantCodeMotion{}.Run(m.get())); EXPECT_FALSE(simplified_loop); } @@ -482,14 +491,14 @@ ENTRY entry { )"; TEST_F(WhileLoopInvariantCodeMotionTest, HoistsConstantWhenAsked) { - ParseAndVerifyModule(kConstantHoistingTestCase); + auto m = ParseAndReturnVerifiedModule(kConstantHoistingTestCase).ValueOrDie(); TF_ASSERT_OK_AND_ASSIGN( bool simplified_loop, - WhileLoopInvariantCodeMotion{/*hoist_constants=*/true}.Run(&module())); + WhileLoopInvariantCodeMotion{/*hoist_constants=*/true}.Run(m.get())); EXPECT_TRUE(simplified_loop); - HloComputation* while_body = module().GetComputationWithName("wide.body"); + HloComputation* while_body = m->GetComputationWithName("wide.body"); ASSERT_NE(while_body, nullptr); // We expect the while body to be the equivalent of: @@ -523,10 +532,98 @@ TEST_F(WhileLoopInvariantCodeMotionTest, HoistsConstantWhenAsked) { } TEST_F(WhileLoopInvariantCodeMotionTest, DoesNotHoistConstantByDefault) { - ParseAndVerifyModule(kConstantHoistingTestCase); + auto m = ParseAndReturnVerifiedModule(kConstantHoistingTestCase).ValueOrDie(); TF_ASSERT_OK_AND_ASSIGN(bool simplified_loop, - WhileLoopInvariantCodeMotion{}.Run(&module())); + WhileLoopInvariantCodeMotion{}.Run(m.get())); + EXPECT_FALSE(simplified_loop); +} + +TEST_F(WhileLoopInvariantCodeMotionTest, DoNotHoistOutOfSingleIteration) { + const char* const kHloModule = R"( + HloModule ModuleWithWhile + + body { + p_body = (f32[2], f32[2], f32[2], s32[]) parameter(0) + val.0 = f32[2] get-tuple-element(p_body), index=0 + val.1 = f32[2] get-tuple-element(p_body), index=1 + add = f32[2] add(val.0, val.1) + const = s32[] constant(-1) + ROOT root = (f32[2], f32[2], f32[2], s32[]) tuple(val.0, val.1, add, const) + } + + condition { + p_cond = (f32[2], f32[2], f32[2], s32[]) parameter(0) + gte = s32[] get-tuple-element(p_cond), index=3 + const = s32[] constant(42) + ROOT result = pred[] equal-to(gte, const) + } + + ENTRY entry { + param.0 = f32[2] parameter(0) + param.1 = s32[] parameter(1) + while_init = (f32[2], f32[2], f32[2], s32[]) tuple(param.0, param.0, param.0, param.1) + ROOT while = (f32[2], f32[2], f32[2], s32[]) while(while_init), condition=condition, body=body + })"; + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseAndReturnVerifiedModule(kHloModule)); + + TF_ASSERT_OK_AND_ASSIGN(bool simplified_loop, + WhileLoopInvariantCodeMotion{}.Run(module.get())); + EXPECT_FALSE(simplified_loop); +} + +const char* const kInflatingTestCase = R"( +HloModule ModuleWithWhile + +mul { + lhs = f32[] parameter(0) + rhs = f32[] parameter(1) + ROOT mul = f32[] multiply(lhs, rhs) +} + +body { + p_body = (f32[]) parameter(0) + iota = f32[1024, 1024] iota(), iota_dimension=0 + add = f32[1024, 1024] add(iota, iota) + constant = f32[] constant(1.0) + reduce = f32[] reduce(f32[1024, 1024] add, f32[] constant), dimensions={0,1}, to_apply=mul + ROOT root = (f32[]) tuple(reduce) +} + +condition { + p_cond = (f32[]) parameter(0) + ROOT result = pred[] constant(true) +} + +ENTRY entry { + param = f32[] parameter(0) + while_init = (f32[]) tuple(param) + ROOT while = (f32[]) while(while_init), condition=condition, body=body +} +)"; + +TEST_F(WhileLoopInvariantCodeMotionTest, HoistsInflatingByDefault) { + auto m = ParseAndReturnVerifiedModule(kInflatingTestCase).ValueOrDie(); + + TF_ASSERT_OK_AND_ASSIGN( + bool simplified_loop, + WhileLoopInvariantCodeMotion(/*hoist_constants=*/true).Run(m.get())); + EXPECT_TRUE(simplified_loop); + + HloComputation* while_body = m->GetComputationWithName("wide.body"); + ASSERT_NE(while_body, nullptr); + EXPECT_THAT(while_body->instructions(), Not(Contains(op::Iota()))); +} + +TEST_F(WhileLoopInvariantCodeMotionTest, NoHoistInflating) { + auto m = ParseAndReturnVerifiedModule(kInflatingTestCase).ValueOrDie(); + + TF_ASSERT_OK_AND_ASSIGN( + bool simplified_loop, + WhileLoopInvariantCodeMotion(/*hoist_constants=*/true, + /*hoist_size_inflating_ops=*/false) + .Run(m.get())); EXPECT_FALSE(simplified_loop); } diff --git a/tensorflow/compiler/xla/service/while_loop_simplifier.cc b/tensorflow/compiler/xla/service/while_loop_simplifier.cc index 630d71e5ca25e9d282ce6283284a32d6f725a193..09d54095718029541a7a25aa62f9a2e9a177960d 100644 --- a/tensorflow/compiler/xla/service/while_loop_simplifier.cc +++ b/tensorflow/compiler/xla/service/while_loop_simplifier.cc @@ -19,41 +19,19 @@ limitations under the License. #include "absl/strings/str_cat.h" #include "absl/strings/str_join.h" #include "absl/types/optional.h" +#include "tensorflow/compiler/xla/primitive_util.h" #include "tensorflow/compiler/xla/service/call_inliner.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/service/hlo_instructions.h" +#include "tensorflow/compiler/xla/service/hlo_query.h" +#include "tensorflow/compiler/xla/service/pattern_matcher.h" #include "tensorflow/compiler/xla/service/while_loop_analysis.h" namespace xla { +namespace m = match; using absl::optional; - -// Determines whether the given instruction is a send/recv node, or has a -// subcomputation which contains a send/recv node. -static bool IsOrContainsSendOrRecv(const HloInstruction* instr); - -// Determines whether the given computation contains a send or recv node. -static bool ContainsSendOrRecv(const HloComputation* comp) { - for (const auto* instr : comp->instructions()) { - if (IsOrContainsSendOrRecv(instr)) { - return true; - } - } - return false; -} - -static bool IsOrContainsSendOrRecv(const HloInstruction* instr) { - if (instr->opcode() == HloOpcode::kSend || - instr->opcode() == HloOpcode::kSendDone || - instr->opcode() == HloOpcode::kRecv || - instr->opcode() == HloOpcode::kRecvDone) { - return true; - } - for (const auto& subcomp : instr->called_computations()) { - if (ContainsSendOrRecv(subcomp)) { - return true; - } - } - return false; -} +using hlo_query::ContainsInstrWithOpcode; // Tries to remove elements in a while loop's tuple that aren't used within the // loop. @@ -80,7 +58,7 @@ static StatusOr TryRemoveDeadWhileParams(HloInstruction* while_op) { HloComputation* while_body = while_op->while_body(); HloInstruction* while_body_root = while_body->root_instruction(); - if (!ShapeUtil::IsTuple(while_init->shape())) { + if (!while_init->shape().IsTuple()) { VLOG(2) << "While op's carried value isn't tuple shaped."; return false; } @@ -131,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; } @@ -149,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; } @@ -180,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) { @@ -253,7 +230,7 @@ static StatusOr TryRemoveDeadWhileParams(HloInstruction* while_op) { // Create the new while condition, body, and init value. std::unique_ptr new_while_cond = while_cond->CloneWithReplacements( - make_while_computation_replacements(while_cond), /*extras=*/{}); + make_while_computation_replacements(while_cond)); std::unordered_map> while_body_replacements = make_while_computation_replacements(while_body); @@ -266,8 +243,7 @@ static StatusOr TryRemoveDeadWhileParams(HloInstruction* while_op) { while_body_replacements.emplace( while_body_root, HloInstruction::CreateTuple(new_while_body_root_elems)); std::unique_ptr new_while_body = - while_body->CloneWithReplacements(std::move(while_body_replacements), - /*extras=*/{}); + while_body->CloneWithReplacements(std::move(while_body_replacements)); // Add a new while_init instruction that repackages the old while_init // instruction's elements. We rely on the AlgebraicSimplifier and DCE to @@ -329,6 +305,147 @@ static StatusOr TryRemoveDeadWhileParams(HloInstruction* while_op) { return true; } +// Removes each loop parameter (i.e. member of the while loop tuple) that is a +// constant and is the same in the while loop body and the while loop init. +static StatusOr TryRemoveConstantParams(HloInstruction* while_op) { + HloModule* module = while_op->GetModule(); + HloComputation* computation = while_op->parent(); + auto* while_init = while_op->mutable_operand(0); + auto* while_body = while_op->while_body(); + auto* while_cond = while_op->while_condition(); + auto* while_body_root = while_body->root_instruction(); + if (while_init->opcode() != HloOpcode::kTuple || + while_body_root->opcode() != HloOpcode::kTuple) { + return false; + } + + TF_RET_CHECK(while_cond->num_parameters() == 1); + TF_RET_CHECK(while_body->num_parameters() == 1); + TF_RET_CHECK( + ShapeUtil::Compatible(while_init->shape(), while_body_root->shape())); + + absl::flat_hash_set constant_tuple_indices; + const auto& while_shape = while_init->shape(); + for (int64 i = 0; i < while_shape.tuple_shapes_size(); ++i) { + auto* init_elem = while_init->operand(i); + auto* body_elem = while_body_root->operand(i); + if (init_elem->opcode() == HloOpcode::kConstant && + body_elem->opcode() == HloOpcode::kConstant && + init_elem->literal() == body_elem->literal()) { + constant_tuple_indices.insert(i); + } + } + + if (constant_tuple_indices.empty()) { + return false; + } + + // OK, we found some constant elements of the while parameter! Eliminate + // them. + std::vector new_while_shape_elems; + for (int64 i = 0; i < while_shape.tuple_shapes_size(); ++i) { + if (!constant_tuple_indices.count(i)) { + new_while_shape_elems.push_back(while_shape.tuple_shapes(i)); + } + } + Shape new_while_shape = ShapeUtil::MakeTupleShape(new_while_shape_elems); + + // `new_instrs` holds instructions created outside of a computation for + // cloning. Elements added here just need to live until the end of the + // relevant CloneWithReplacement call. + std::vector> new_instrs; + auto add_new_instr = [&](std::unique_ptr instr) { + new_instrs.push_back(std::move(instr)); + return new_instrs.back().get(); + }; + + // Returns a new tuple without the elements of constant_tuple_indices. + auto remove_constant_elems = [&](HloInstruction* instr) { + CHECK(ShapeUtil::Compatible(instr->shape(), while_shape)); + + std::vector tuple_elems; + for (int64 i = 0; i < while_shape.tuple_shapes_size(); ++i) { + if (!constant_tuple_indices.count(i)) { + tuple_elems.push_back( + add_new_instr(HloInstruction::CreateGetTupleElement( + while_shape.tuple_shapes(i), instr, i))); + } + } + return HloInstruction::CreateTuple(tuple_elems); + }; + + auto add_constant_elems = [&](HloInstruction* instr) { + CHECK(ShapeUtil::Compatible(instr->shape(), new_while_shape)); + + std::vector tuple_elems; + int64 j = 0; + for (int64 i = 0; i < while_shape.tuple_shapes_size(); ++i) { + if (constant_tuple_indices.count(i)) { + tuple_elems.push_back(while_init->mutable_operand(i)); + } else { + tuple_elems.push_back( + add_new_instr(HloInstruction::CreateGetTupleElement( + while_shape.tuple_shapes(i), instr, j))); + ++j; + } + } + return HloInstruction::CreateTuple(tuple_elems); + }; + + // Special case: constant_tuple_indices covers the whole while parameter, so + // the new while shape is the empty tuple. In this case, the value of the + // while loop is simply equal to the value of `init`. + // + // It's unfortunate to special-case this, but it's simpler than the + // alternative. The problem is that if our while parameter has no + // non-constant elems, the tuple returned by `add_constant_elems` won't depend + // on instr (the loop body/cond parameter), and therefore + // CloneWithReplacementPairs will *leave the parameter out entirely*, creating + // invalid HLO. + if (ShapeUtil::IsEmptyTuple(new_while_shape)) { + TF_RETURN_IF_ERROR(computation->ReplaceInstruction(while_op, while_init)); + return true; + } + + std::unique_ptr new_while_cond = + while_cond->CloneWithReplacementPairs({ + while_cond->parameter_instruction(0), + add_constant_elems(add_new_instr(HloInstruction::CreateParameter( + 0, new_while_shape, + while_cond->parameter_instruction(0)->name()))), + }); + + std::unique_ptr new_while_body = + while_body->CloneWithReplacementPairs( + { + while_body->parameter_instruction(0), + add_constant_elems(add_new_instr(HloInstruction::CreateParameter( + 0, new_while_shape, + while_cond->parameter_instruction(0)->name()))), + }, + { + while_body->root_instruction(), + remove_constant_elems( + add_new_instr(while_body->root_instruction()->Clone())), + }); + + // Create the final while loop, and add any new instructions created to + // `computation`. + new_instrs.clear(); + TF_RETURN_IF_ERROR(computation->ReplaceWithNewInstruction( + while_op, + add_constant_elems( + computation->AddInstruction(HloInstruction::CreateWhile( + new_while_shape, + module->AddEmbeddedComputation(std::move(new_while_cond)), + module->AddEmbeddedComputation(std::move(new_while_body)), + add_new_instr(remove_constant_elems(while_init))))))); + for (auto& instr : new_instrs) { + computation->AddInstruction(std::move(instr)); + } + return true; +} + // Tries to remove a while loop from the graph. // // - Loops with trip count of 0 can be replaced by the loop's "init" value. @@ -408,16 +525,14 @@ static StatusOr TryPropagateConstant(HloInstruction* while_op) { // performance by forcing us to copy constants. absl::flat_hash_map index_to_constant; for (int i = 0; i < root_operands.size(); i++) { - HloInstruction* instr = root_operands[i]; - if (instr->opcode() == HloOpcode::kGetTupleElement && - instr->tuple_index() == i && instr->operand(0) == while_body_param && - ShapeUtil::IsScalar(instr->shape())) { - auto tuple_element = while_init->operand(i); - if (tuple_element->IsConstant()) { - VLOG(3) << "Found loop invariant tuple element " << i << " " - << tuple_element->ToString(); - index_to_constant[i] = tuple_element; - } + const HloInstruction* init_tuple_elem = nullptr; + if (Match(root_operands[i], + m::GetTupleElement(m::Op().Is(while_body_param), i) + .WithShape(m::Shape().IsScalar())) && + Match(while_init->operand(i), m::Constant(&init_tuple_elem))) { + VLOG(3) << "Found loop invariant tuple element " << i << " " + << init_tuple_elem->ToString(); + index_to_constant[i] = init_tuple_elem; } } @@ -458,6 +573,408 @@ static StatusOr TryPropagateConstant(HloInstruction* while_op) { return changed_cond || changed_body; } +// Converts a flat list of instructions into a tuple of the desired shape. For +// example, given a tuple shape ((x, x), x) and instructions {A, B, C}, returns +// a tuple of value ((A, B), C). +// +// desired_shape must be a tuple. (This precondition allows us to return a +// unique_ptr rather than a raw ptr.) +static std::unique_ptr UnflattenTupleInstr( + absl::Span instrs, const Shape& desired_shape, + std::vector>* new_instrs) { + CHECK(desired_shape.IsTuple()) << ShapeUtil::HumanString(desired_shape); + + // For each child shape in `desired_shape`, slice out the correct number of + // `instrs` and call UnflattenTupleInstr recursively. At each step we remove + // elements from `instrs` so that it only contains instructions we have not + // yet processed. + std::vector elems; + for (int64 i = 0; i < desired_shape.tuple_shapes_size(); ++i) { + const Shape& subshape = desired_shape.tuple_shapes(i); + if (!subshape.IsTuple()) { + elems.push_back(instrs[0]); + instrs.remove_prefix(1); + continue; + } + + // Count the number of leaf nodes underneath desired_shape[i]. + int64 num_leaves = 0; + ShapeUtil::ForEachSubshape( + subshape, [&](const Shape& s, const ShapeIndex& /*index*/) { + if (!s.IsTuple()) { + ++num_leaves; + } + }); + + std::unique_ptr subinstr = + UnflattenTupleInstr(instrs.subspan(0, num_leaves), + desired_shape.tuple_shapes(i), new_instrs); + elems.push_back(subinstr.get()); + new_instrs->push_back(std::move(subinstr)); + instrs.remove_prefix(num_leaves); + } + return HloInstruction::CreateTuple(elems); +} + +// Builds a vector whose elements are the values in the flattened tuple for +// `instr`. For example, if `instr` is a tuple of form ((A, B), C), returns the +// vector {A, B, C} (or kGetTupleElement ops which point to A, B, and C). +static std::vector GetFlatTupleElems( + HloInstruction* instr, + std::vector>* new_instrs) { + const auto& shape = instr->shape(); + if (!shape.IsTuple()) { + return {instr}; + } + std::vector elems; + for (int64 i = 0; i < shape.tuple_shapes_size(); ++i) { + const Shape& subshape = shape.tuple_shapes(i); + new_instrs->push_back( + HloInstruction::CreateGetTupleElement(subshape, instr, i)); + auto* gte = new_instrs->back().get(); + auto flattened_subshape = GetFlatTupleElems(gte, new_instrs); + elems.insert(elems.end(), flattened_subshape.begin(), + flattened_subshape.end()); + } + return elems; +} + +static StatusOr TryFlattenNestedTuples(HloInstruction* while_op) { + HloModule* module = while_op->GetModule(); + HloComputation* computation = while_op->parent(); + auto* while_init = while_op->mutable_operand(0); + auto* while_body = while_op->while_body(); + auto* while_cond = while_op->while_condition(); + auto* while_body_root = while_body->root_instruction(); + if (while_init->opcode() != HloOpcode::kTuple || + while_body_root->opcode() != HloOpcode::kTuple) { + return false; + } + + TF_RET_CHECK(while_cond->num_parameters() == 1); + TF_RET_CHECK(while_body->num_parameters() == 1); + TF_RET_CHECK( + ShapeUtil::Compatible(while_init->shape(), while_body_root->shape())); + Shape while_shape = while_init->shape(); + if (!ShapeUtil::IsNestedTuple(while_shape)) { + return false; + } + + std::vector flattened_shape_elems; + ShapeUtil::ForEachSubshape(while_shape, + [&](const Shape& s, const ShapeIndex& /*index*/) { + if (!s.IsTuple()) { + flattened_shape_elems.push_back(s); + } + }); + Shape flattened_shape = ShapeUtil::MakeTupleShape(flattened_shape_elems); + + // `new_instrs` holds instructions created outside of a computation for + // cloning. Elements added here just need to live until the end of the + // relevant CloneWithReplacement call. + std::vector> new_instrs; + auto add_new_instr = [&](std::unique_ptr instr) { + new_instrs.push_back(std::move(instr)); + return new_instrs.back().get(); + }; + + auto nested = [&](HloInstruction* instr) { + std::vector gtes; + const Shape& flat_shape = instr->shape(); + for (int64 i = 0; i < flat_shape.tuple_shapes_size(); ++i) { + gtes.push_back(add_new_instr(HloInstruction::CreateGetTupleElement( + flat_shape.tuple_shapes(i), instr, i))); + } + auto nested_instr = + UnflattenTupleInstr(absl::MakeSpan(gtes), while_shape, &new_instrs); + CHECK(ShapeUtil::Compatible(nested_instr->shape(), while_shape)) + << ShapeUtil::HumanString(nested_instr->shape()) << " vs " + << ShapeUtil::HumanString(while_shape); + return nested_instr; + }; + + auto flattened = [&](HloInstruction* instr) { + return HloInstruction::CreateTuple(GetFlatTupleElems(instr, &new_instrs)); + }; + + // Create a new while-condition computation, where parameter 0 has flat shape + // but all uses of it go through the nested shape. + std::unique_ptr new_while_cond = + while_cond->CloneWithReplacementPairs({ + while_cond->parameter_instruction(0), + nested(add_new_instr(HloInstruction::CreateParameter( + 0, flattened_shape, + while_cond->parameter_instruction(0)->name()))), + }); + + // Create a new while-body computation, where parameter 0 has a flat shape and + // all uses of it go through the nested shape, and where the root has a flat + // shape constructed from the old nested root. + std::unique_ptr new_while_body = + while_body->CloneWithReplacementPairs( + { + while_body->parameter_instruction(0), + nested(add_new_instr(HloInstruction::CreateParameter( + 0, flattened_shape, + while_body->parameter_instruction(0)->name()))), + }, + { + while_body->root_instruction(), + flattened(add_new_instr(while_body->root_instruction()->Clone())), + }); + + // Create the final while loop, and add any new instructions created to + // `computation`. + new_instrs.clear(); + TF_RETURN_IF_ERROR(computation->ReplaceWithNewInstruction( + while_op, nested(computation->AddInstruction(HloInstruction::CreateWhile( + flattened_shape, + module->AddEmbeddedComputation(std::move(new_while_cond)), + module->AddEmbeddedComputation(std::move(new_while_body)), + computation->AddInstruction(flattened(while_init))))))); + for (auto& instr : new_instrs) { + computation->AddInstruction(std::move(instr)); + } + return true; +} + +// Tries to merge loop induction variables of a given type. +// +// In this pass we're only concerned with elements of the loop's tuple that +// are effective-scalars of type `elem_ty`. Some terminology: +// +// - The trip counter is the first element of the loop's tuple that starts at +// 0 and does x++ on each iteration. +// +// - An induction variable is an element of the loop's tuple that is not the +// trip counter and does `x += ` on each iteration of the loop. +// Negative constants are OK. +// +// This pass adds a trip counter if one isn't already present, then replaces +// each induction variable with +// +// + * . +// +// This reduces the number of scalar operations in the loop, which is important +// e.g. on GPUs, where each scalar operation is nontrivially expensive because +// it's a separate kernel launch. +// +// Returns the new loop if a change was made, or null if no change was made. +// Note that the new loop is not a valid replacement for the old loop; it may +// need to be wrapped in a tuple that changes its shape. We return the loop +// itself so that you can call TryMergeInductionVariables in a loop, once for +// each integral type elem_ty. +static StatusOr TryMergeInductionVariables( + HloInstruction* while_op, PrimitiveType elem_ty) { + CHECK(primitive_util::IsIntegralType(elem_ty)) << PrimitiveType_Name(elem_ty); + HloModule* module = while_op->GetModule(); + HloComputation* computation = while_op->parent(); + auto* while_init = while_op->mutable_operand(0); + auto* while_body = while_op->while_body(); + auto* while_cond = while_op->while_condition(); + auto* while_body_root = while_body->root_instruction(); + if (while_init->opcode() != HloOpcode::kTuple || + while_body_root->opcode() != HloOpcode::kTuple) { + return nullptr; + } + + TF_RET_CHECK(while_cond->num_parameters() == 1); + TF_RET_CHECK(while_body->num_parameters() == 1); + TF_RET_CHECK( + ShapeUtil::Compatible(while_init->shape(), while_body_root->shape())); + Shape while_shape = while_init->shape(); + + // The tuple index of the trip counter, if one is present. + absl::optional trip_counter; + // Maps the tuple index of each induction variable to its constant increment. + absl::flat_hash_map induction_vars; + for (int64 i = 0; i < while_body_root->operand_count(); ++i) { + HloInstruction* constant; + if (!Match(while_body_root->mutable_operand(i), + m::AddAnyOrder(m::GetTupleElement(m::Parameter(), i), + m::ConstantScalar(&constant)) + .WithShape(m::Shape().WithElementType(elem_ty)))) { + continue; + } + if (!trip_counter && constant->literal().IsAll(1) && + while_init->operand(i)->IsConstant() && + while_init->operand(i)->literal().IsAll(0)) { + VLOG(10) << "Found existing trip counter at index " << i; + trip_counter = i; + } else { + VLOG(10) << "Found induction variable at index " << i; + induction_vars.emplace(i, Cast(constant)); + } + } + + // There's only something to simplify if we can either: + // + // - combine one or more induction vars with an existing trip counter, or + // - replace two or more induction variables with a new trip counter. + // + // Put another way, there's only something to simplify if the number of + // induction vars plus the number of existing trip counters (0 or 1) is >= 2. + if (induction_vars.size() + (trip_counter.has_value() ? 1 : 0) < 2) { + return nullptr; + } + + // OK, we're going to do the transformation! Set up some helpers. + + // `new_instrs` holds instructions created outside of a computation for + // cloning. Elements added here just need to live until the end of the + // relevant CloneWithReplacement call. + std::vector> new_instrs; + auto add_new_instr = [&](std::unique_ptr instr) { + new_instrs.push_back(std::move(instr)); + return new_instrs.back().get(); + }; + + auto add_binary_op = [&](const Shape& shape, HloOpcode opcode, + HloInstruction* lhs, HloInstruction* rhs) { + // Reshape lhs/rhs to the output shape if necessary. This deals with the + // fact that induction variables need only be effective scalars, not true + // scalars. + if (!ShapeUtil::Compatible(shape, lhs->shape())) { + lhs = add_new_instr(HloInstruction::CreateReshape(shape, lhs)); + } + if (!ShapeUtil::Compatible(shape, rhs->shape())) { + rhs = add_new_instr(HloInstruction::CreateReshape(shape, rhs)); + } + return add_new_instr(HloInstruction::CreateBinary(shape, opcode, lhs, rhs)); + }; + + auto add_gte = [&](HloInstruction* src, int64 idx) { + return add_new_instr(HloInstruction::CreateGetTupleElement( + src->shape().tuple_shapes(idx), src, idx)); + }; + + // Our new while loop will have the same shape as the old while loop, except + // we'll add a trip counter to the end if it wasn't originally present. + Shape new_while_shape = while_shape; + bool added_trip_counter = false; + if (!trip_counter) { + VLOG(10) << "Adding new trip counter to end of loop's tuple."; + trip_counter = new_while_shape.tuple_shapes_size(); + *new_while_shape.add_tuple_shapes() = + ShapeUtil::MakeShape(elem_ty, /*dimensions=*/{}); + added_trip_counter = true; + } + + // Converts `instr` into a tuple of the "old" form -- that is, to a tuple with + // shape `while_body->shape()` and where the induction variables are "reified" + // (i.e. they have value + * ). + auto convert_to_old_form = [&](HloInstruction* instr) { + CHECK(ShapeUtil::Compatible(instr->shape(), new_while_shape)); + std::vector tuple_elems; + for (int64 i = 0; i < while_shape.tuple_shapes_size(); ++i) { + const auto& elem_shape = while_shape.tuple_shapes(i); + if (!induction_vars.count(i)) { + tuple_elems.push_back(add_gte(instr, i)); + continue; + } + tuple_elems.push_back(add_binary_op( + elem_shape, HloOpcode::kAdd, add_gte(instr, i), + add_binary_op(elem_shape, HloOpcode::kMultiply, + add_gte(instr, *trip_counter), + add_new_instr(induction_vars.at(i)->Clone())))); + } + return HloInstruction::CreateTuple(tuple_elems); + }; + + // Converts `root` into a tuple of the "new" form -- that is, to a tuple with + // shape `new_while_shape` and where the induction variables (but not trip + // counters) are replaced with their unchanging values. + auto convert_to_new_form = [&](HloInstruction* old_root, + HloParameterInstruction* loop_body_param) { + CHECK(ShapeUtil::Compatible(old_root->shape(), while_shape)); + std::vector tuple_elems; + + // In the new form, induction variables come from `init`, everything else + // (including the trip counter if it's not one we created ourselves) comes + // from the `root` tuple unmodified. + for (int64 i = 0; i < while_shape.tuple_shapes_size(); ++i) { + tuple_elems.push_back( + add_gte((induction_vars.count(i) ? loop_body_param : old_root), i)); + } + // If we created a trip counter ourselves, add 1 to it in the next + // iteration. + if (added_trip_counter) { + tuple_elems.push_back(add_binary_op( + new_while_shape.tuple_shapes(*trip_counter), HloOpcode::kAdd, + add_gte(loop_body_param, *trip_counter), + add_new_instr( + HloInstruction::CreateConstant(LiteralUtil::One(elem_ty))))); + } + + return HloInstruction::CreateTuple(tuple_elems); + }; + + // Creates a new init tuple, which is the same as the old init tuple except if + // we added a trip counter, it's set to 0. + auto get_new_while_init = [&](HloInstruction* init) { + CHECK(ShapeUtil::Compatible(init->shape(), while_shape)); + if (!added_trip_counter) { + return init; + } + std::vector tuple_elems; + for (int64 i = 0; i < while_shape.tuple_shapes_size(); ++i) { + tuple_elems.push_back(add_gte(init, i)); + } + tuple_elems.push_back(add_new_instr( + HloInstruction::CreateConstant(LiteralUtil::Zero(elem_ty)))); + return add_new_instr(HloInstruction::CreateTuple(tuple_elems)); + }; + + std::unique_ptr new_while_cond = + while_cond->CloneWithReplacementPairs({ + while_cond->parameter_instruction(0), + convert_to_old_form(add_new_instr(HloInstruction::CreateParameter( + 0, new_while_shape, + while_cond->parameter_instruction(0)->name()))), + }); + + // Creating the new while body proceeds in two steps. First we convert the + // users of the parameter to the old form. Then as a second + // CloneWithReplacement operation we convert the root to the new form. We + // have to do this in two steps because the new root needs to use the new + // param0, and during the first clone operation, only the *old-form* param0 is + // accessible. + // + // We have to add temp_new_while_body to the module because cloning a + // computation touches the module (to get its NameUniquer). + HloComputation* temp_new_while_body = + module->AddEmbeddedComputation(while_body->CloneWithReplacementPairs({ + while_body->parameter_instruction(0), + convert_to_old_form(add_new_instr(HloInstruction::CreateParameter( + 0, new_while_shape, + while_body->parameter_instruction(0)->name()))), + })); + std::unique_ptr new_while_body = + temp_new_while_body->CloneWithReplacementPairs({ + temp_new_while_body->root_instruction(), + convert_to_new_form( + add_new_instr(temp_new_while_body->root_instruction()->Clone()), + Cast( + temp_new_while_body->parameter_instruction(0))), + }); + TF_RETURN_IF_ERROR(module->RemoveEmbeddedComputation(temp_new_while_body)); + + // Create the final while loop, and add any new instructions created to + // `computation`. + new_instrs.clear(); + auto* new_while = computation->AddInstruction(HloInstruction::CreateWhile( + new_while_shape, + module->AddEmbeddedComputation(std::move(new_while_cond)), + module->AddEmbeddedComputation(std::move(new_while_body)), + get_new_while_init(while_init))); + TF_RETURN_IF_ERROR(computation->ReplaceWithNewInstruction( + while_op, convert_to_old_form(new_while))); + for (auto& instr : new_instrs) { + computation->AddInstruction(std::move(instr)); + } + return new_while; +} + StatusOr WhileLoopSimplifier::Run(HloModule* module) { XLA_VLOG_LINES(3, "WhileLoopSimplifier::Run(), before:\n" + module->ToString()); @@ -478,32 +995,77 @@ StatusOr WhileLoopSimplifier::Run(HloModule* module) { for (HloInstruction* while_op : while_ops) { // We can't remove while loops that contain send/recv nodes, because we rely // on the particular loop structure around the node matching on the send and - // recv sides. Removing dead while params requires us to remove the loop + // recv sides. Other while simplifications require us to remove the loop // and replace it with a new one, so we can't do that either. - if (ContainsSendOrRecv(while_op->while_body()) || - ContainsSendOrRecv(while_op->while_condition())) { + if (ContainsInstrWithOpcode(while_op->while_body(), + {HloOpcode::kSend, HloOpcode::kSendDone, + HloOpcode::kRecv, HloOpcode::kRecvDone}) || + ContainsInstrWithOpcode(while_op->while_condition(), + {HloOpcode::kSend, HloOpcode::kSendDone, + HloOpcode::kRecv, HloOpcode::kRecvDone})) { VLOG(2) << "Not attempting to simplify while loop because it contains a " "send/recv node: " << while_op->ToShortString(); continue; } - StatusOr result = TryPropagateConstant(while_op); - TF_RETURN_IF_ERROR(result.status()); - changed |= result.ValueOrDie(); + TF_ASSIGN_OR_RETURN(bool result, TryPropagateConstant(while_op)); + changed |= result; + + TF_ASSIGN_OR_RETURN(result, TryRemoveWhileLoop(while_op)); + changed |= result; + if (result) { + // Don't continue simplifying after successfully removing the while loop + // -- that would result in use-after-free nastiness. + continue; + } + + // TODO(b/119281462): Cowardly refuse to perform any of the following + // optimizations in the presence of kDomain instructions. It seems that + // modifying a while loop's tuple doesn't work when kDomain is present. + if (ContainsInstrWithOpcode(while_op->while_body(), {HloOpcode::kDomain}) || + ContainsInstrWithOpcode(while_op->while_condition(), + {HloOpcode::kDomain})) { + continue; + } + + // Each of the optimizations below modifies the while loop itself if it's + // successful, meaning that `while_op` is no longer valid after one of these + // transformations returns true. - result = TryRemoveWhileLoop(while_op); - TF_RETURN_IF_ERROR(result.status()); - if (result.ValueOrDie()) { - changed = true; - // Don't try to remove dead while params after successfully removing the - // while loop -- that would result in use-after-free nastiness. + TF_ASSIGN_OR_RETURN(result, TryFlattenNestedTuples(while_op)); + changed |= result; + if (result) { continue; } - result = TryRemoveDeadWhileParams(while_op); - TF_RETURN_IF_ERROR(result.status()); - changed |= result.ValueOrDie(); + TF_ASSIGN_OR_RETURN(result, TryRemoveDeadWhileParams(while_op)); + changed |= result; + if (result) { + continue; + } + + TF_ASSIGN_OR_RETURN(result, TryRemoveConstantParams(while_op)); + changed |= result; + if (result) { + continue; + } + + bool merged_induction_vars = false; + // Notably missing from this list are S16 and U16. These don't currently + // work because S/U16 literals are not implemented. + for (auto elem_ty : {S8, U8, S32, U32, S64, U64}) { + TF_ASSIGN_OR_RETURN(auto* new_while_op, + TryMergeInductionVariables(while_op, elem_ty)); + if (new_while_op) { + while_op = new_while_op; + changed = true; + merged_induction_vars = true; + } + } + if (merged_induction_vars) { + continue; + } } XLA_VLOG_LINES(3, diff --git a/tensorflow/compiler/xla/service/while_loop_simplifier.h b/tensorflow/compiler/xla/service/while_loop_simplifier.h index 0bc5a0107bbcfb3b29a01d593fb79b89a863e49b..a378f179c63c788cd205ddbb784dee0e6b2106d7 100644 --- a/tensorflow/compiler/xla/service/while_loop_simplifier.h +++ b/tensorflow/compiler/xla/service/while_loop_simplifier.h @@ -25,11 +25,22 @@ namespace xla { // HLO pass that makes the following transformations on while loops: // // - A while loop with static trip count of 0 is deleted. +// // - A while loop with static trip count of 1 is replaced by its body (sans // loop). +// // - Elements of a while loop's tuple that the loop doesn't use are removed // from the tuple. // +// - If the while loop's parameter is a nested tuple, it's flattened to a +// single-level tuple. This is good because it usually reduces the number of +// kTuple instructions, but also because it unlocks additional optimizations +// (e.g. removing unused loop parameters). +// +// Flattening nested while loop tuples adds a whole mess of likely unnecessary +// kGetTupleElement and kTuple operations to the graph. We expect that tuple +// simplifier will be run afterwards. +// class WhileLoopSimplifier : public HloModulePass { public: ~WhileLoopSimplifier() override {} diff --git a/tensorflow/compiler/xla/service/while_loop_simplifier_test.cc b/tensorflow/compiler/xla/service/while_loop_simplifier_test.cc index 1c892ba179ec67ccc9dbfe93d925551d6977ba15..ecca76b1e86d833c73fbb9bad6a341660a7d2669 100644 --- a/tensorflow/compiler/xla/service/while_loop_simplifier_test.cc +++ b/tensorflow/compiler/xla/service/while_loop_simplifier_test.cc @@ -17,28 +17,46 @@ limitations under the License. #include "absl/strings/str_cat.h" #include "absl/strings/str_replace.h" +#include "tensorflow/compiler/xla/service/algebraic_simplifier.h" +#include "tensorflow/compiler/xla/service/hlo_cse.h" +#include "tensorflow/compiler/xla/service/hlo_dce.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_matchers.h" +#include "tensorflow/compiler/xla/service/hlo_parser.h" +#include "tensorflow/compiler/xla/service/tuple_simplifier.h" #include "tensorflow/compiler/xla/test.h" -#include "tensorflow/compiler/xla/tests/hlo_verified_test_base.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" #include "tensorflow/core/lib/core/status_test_util.h" namespace xla { namespace { +using ::testing::_; namespace op = xla::testing::opcode_matchers; -class WhileLoopSimplifierTest : public HloVerifiedTestBase { +// Returns the first kWhile instruction within m's entry computation. +HloInstruction* FindFirstWhile(HloModule* m) { + const auto& instrs = m->entry_computation()->instructions(); + return *absl::c_find_if(instrs, [](const HloInstruction* instr) { + return instr->opcode() == HloOpcode::kWhile; + }); +} + +class WhileLoopSimplifierTest : public HloTestBase { protected: // Makes an HloModule that contains a loop with `num_iters` iteration. - void MakeModuleWithSimpleLoop(int num_iters); + TF_MUST_USE_RESULT std::unique_ptr + MakeModuleWithSimpleLoop(int num_iters); // Similar to MakeModuleWithSimpleLoop except that the loop bound is passed to // the loop-condition through an element of a tuple which is the // loop-condition parameter. - void MakeModuleWithSimpleLoopTupleElementLoopBound(int num_iters); + TF_MUST_USE_RESULT std::unique_ptr + MakeModuleWithSimpleLoopTupleElementLoopBound(int num_iters); }; -void WhileLoopSimplifierTest::MakeModuleWithSimpleLoop(int num_iters) { +std::unique_ptr +WhileLoopSimplifierTest::MakeModuleWithSimpleLoop(int num_iters) { string hlo_string_template = R"( HloModule SimpleLoop SimpleLoop.body { @@ -67,10 +85,11 @@ void WhileLoopSimplifierTest::MakeModuleWithSimpleLoop(int num_iters) { string hlo_string = absl::StrReplaceAll( hlo_string_template, {{"{{LOOP_BOUND}}", absl::StrCat(42 + num_iters)}}); - ParseAndVerifyModule(hlo_string); + return ParseAndReturnVerifiedModule(hlo_string).ValueOrDie(); } -void WhileLoopSimplifierTest::MakeModuleWithSimpleLoopTupleElementLoopBound( +std::unique_ptr +WhileLoopSimplifierTest::MakeModuleWithSimpleLoopTupleElementLoopBound( int num_iters) { string hlo_string_template = R"( HloModule SimpleLoopWithIndirectLoopBound @@ -104,60 +123,55 @@ void WhileLoopSimplifierTest::MakeModuleWithSimpleLoopTupleElementLoopBound( string hlo_string = absl::StrReplaceAll( hlo_string_template, {{"{{LOOP_BOUND}}", absl::StrCat(42 + num_iters)}}); - ParseAndVerifyModule(hlo_string); + return ParseAndReturnVerifiedModule(hlo_string).ValueOrDie(); } TEST_F(WhileLoopSimplifierTest, LoopWithZeroIterationSimiplified) { - MakeModuleWithSimpleLoop(/*num_iters=*/0); - HloModule* the_module = &module(); - ASSERT_TRUE(WhileLoopSimplifier().Run(the_module).ValueOrDie()); - EXPECT_THAT(the_module->entry_computation()->root_instruction(), + auto m = MakeModuleWithSimpleLoop(/*num_iters=*/0); + ASSERT_TRUE(WhileLoopSimplifier().Run(m.get()).ValueOrDie()); + EXPECT_THAT(m->entry_computation()->root_instruction(), op::Tuple(op::Constant(), op::Constant())); } TEST_F(WhileLoopSimplifierTest, LoopWithZeroIterationTupleElementLoopBoundSimplified) { - MakeModuleWithSimpleLoopTupleElementLoopBound(/*num_iters=*/0); - HloModule* the_module = &module(); - ASSERT_TRUE(WhileLoopSimplifier().Run(the_module).ValueOrDie()); - EXPECT_THAT(the_module->entry_computation()->root_instruction(), + auto m = MakeModuleWithSimpleLoopTupleElementLoopBound(/*num_iters=*/0); + ASSERT_TRUE(WhileLoopSimplifier().Run(m.get()).ValueOrDie()); + EXPECT_THAT(m->entry_computation()->root_instruction(), op::Tuple(op::Constant(), op::Constant(), op::Constant())); } TEST_F(WhileLoopSimplifierTest, LoopWithOneIterationSimplified) { - MakeModuleWithSimpleLoop(/*num_iters=*/1); - HloModule* the_module = &module(); - ASSERT_TRUE(WhileLoopSimplifier().Run(the_module).ValueOrDie()); - EXPECT_THAT(the_module->entry_computation()->root_instruction(), + auto m = MakeModuleWithSimpleLoop(/*num_iters=*/1); + ASSERT_TRUE(WhileLoopSimplifier().Run(m.get()).ValueOrDie()); + EXPECT_THAT(m->entry_computation()->root_instruction(), op::Tuple(op::Add(), op::Multiply())); } TEST_F(WhileLoopSimplifierTest, LoopWithOneIterationTupleELementLoopBoundSimplified) { - MakeModuleWithSimpleLoopTupleElementLoopBound(/*num_iters=*/1); - HloModule* the_module = &module(); - ASSERT_TRUE(WhileLoopSimplifier().Run(the_module).ValueOrDie()); - EXPECT_THAT(the_module->entry_computation()->root_instruction(), + auto m = MakeModuleWithSimpleLoopTupleElementLoopBound(/*num_iters=*/1); + ASSERT_TRUE(WhileLoopSimplifier().Run(m.get()).ValueOrDie()); + EXPECT_THAT(m->entry_computation()->root_instruction(), op::Tuple(op::Add(), op::Multiply(), op::Constant())); } TEST_F(WhileLoopSimplifierTest, LoopWithTwoIterationsNotSimplified) { - MakeModuleWithSimpleLoop(/*num_iters=*/2); - EXPECT_FALSE(WhileLoopSimplifier().Run(&module()).ValueOrDie()); + auto m = MakeModuleWithSimpleLoop(/*num_iters=*/2); + EXPECT_FALSE(WhileLoopSimplifier().Run(m.get()).ValueOrDie()); } TEST_F(WhileLoopSimplifierTest, LoopWithControlDependencySimplifiedDependencyPreserved) { - MakeModuleWithSimpleLoop(/*num_iters=*/1); - HloModule* the_module = &module(); - HloComputation* computation = the_module->entry_computation(); + auto m = MakeModuleWithSimpleLoop(/*num_iters=*/1); + HloComputation* computation = m->entry_computation(); auto* while_op = computation->root_instruction(); ASSERT_EQ(while_op->opcode(), HloOpcode::kWhile); auto* true_op = while_op->while_body()->AddInstruction( HloInstruction::CreateConstant(LiteralUtil::CreateR0(true))); TF_ASSERT_OK(true_op->AddControlDependencyTo( while_op->while_body()->root_instruction())); - ASSERT_TRUE(WhileLoopSimplifier().Run(the_module).ValueOrDie()); + ASSERT_TRUE(WhileLoopSimplifier().Run(m.get()).ValueOrDie()); EXPECT_THAT(computation->root_instruction()->control_predecessors(), ElementsAre(op::Constant())) << computation->ToString(); @@ -166,9 +180,8 @@ TEST_F(WhileLoopSimplifierTest, // Loops that contain send/recv nodes can't be simplified; the loop structure // around send/recv nodes must be preserved. TEST_F(WhileLoopSimplifierTest, LoopWithSendNotSimplified) { - MakeModuleWithSimpleLoop(/*num_iters=*/1); - HloModule* the_module = &module(); - HloComputation* computation = the_module->entry_computation(); + auto m = MakeModuleWithSimpleLoop(/*num_iters=*/1); + HloComputation* computation = m->entry_computation(); auto* while_op = computation->root_instruction(); ASSERT_EQ(while_op->opcode(), HloOpcode::kWhile); auto* while_body = while_op->while_body(); @@ -179,13 +192,12 @@ TEST_F(WhileLoopSimplifierTest, LoopWithSendNotSimplified) { token, /*channel_id=*/0)); while_body->AddInstruction(HloInstruction::CreateSendDone(send)); - EXPECT_FALSE(WhileLoopSimplifier().Run(the_module).ValueOrDie()); + EXPECT_FALSE(WhileLoopSimplifier().Run(m.get()).ValueOrDie()); } TEST_F(WhileLoopSimplifierTest, LoopWithRecvNotSimplified) { - MakeModuleWithSimpleLoop(/*num_iters=*/1); - HloModule* the_module = &module(); - HloComputation* computation = the_module->entry_computation(); + auto m = MakeModuleWithSimpleLoop(/*num_iters=*/1); + HloComputation* computation = m->entry_computation(); auto* while_op = computation->root_instruction(); ASSERT_EQ(while_op->opcode(), HloOpcode::kWhile); auto* while_body = while_op->while_body(); @@ -194,7 +206,7 @@ TEST_F(WhileLoopSimplifierTest, LoopWithRecvNotSimplified) { HloInstruction::CreateRecv(ShapeUtil::MakeShape(F32, {1}), token, /*channel_id=*/0)); while_body->AddInstruction(HloInstruction::CreateRecvDone(recv)); - EXPECT_FALSE(WhileLoopSimplifier().Run(the_module).ValueOrDie()); + EXPECT_FALSE(WhileLoopSimplifier().Run(m.get()).ValueOrDie()); } // The limitation on not being able to simplify loops that contain infeeds (and @@ -202,16 +214,15 @@ TEST_F(WhileLoopSimplifierTest, LoopWithRecvNotSimplified) { // fact that our infrastructure sees simplifying such a loop as tantamount to // removing the non-removable instruction. TEST_F(WhileLoopSimplifierTest, LoopWithInfeedNotSimplified) { - MakeModuleWithSimpleLoop(/*num_iters=*/1); - HloModule* the_module = &module(); - HloComputation* computation = the_module->entry_computation(); + auto m = MakeModuleWithSimpleLoop(/*num_iters=*/1); + HloComputation* computation = m->entry_computation(); auto* while_op = computation->root_instruction(); ASSERT_EQ(while_op->opcode(), HloOpcode::kWhile); auto* while_body = while_op->while_body(); auto token = while_body->AddInstruction(HloInstruction::CreateToken()); while_body->AddInstruction(HloInstruction::CreateInfeed( ShapeUtil::MakeShape(F32, {1}), token, "config")); - EXPECT_FALSE(WhileLoopSimplifier().Run(the_module).ValueOrDie()); + EXPECT_FALSE(WhileLoopSimplifier().Run(m.get()).ValueOrDie()); } // A non-tuple shaped loop shouldn't be simplified or crash the compiler. @@ -236,8 +247,8 @@ TEST_F(WhileLoopSimplifierTest, NonTupleShapedLoopNotSimplified) { } )"; - ParseAndVerifyModule(hlo_string); - EXPECT_FALSE(WhileLoopSimplifier().Run(&module()).ValueOrDie()); + auto m = ParseAndReturnVerifiedModule(hlo_string).ValueOrDie(); + EXPECT_FALSE(WhileLoopSimplifier().Run(m.get()).ValueOrDie()); } // A while loop that does nothing else besides swapping tuple elements @@ -268,8 +279,8 @@ TEST_F(WhileLoopSimplifierTest, LoopSwappingTupleElementsNotSimplified) { } )"; - ParseAndVerifyModule(hlo_string); - EXPECT_FALSE(WhileLoopSimplifier().Run(&module()).ValueOrDie()); + auto m = ParseAndReturnVerifiedModule(hlo_string).ValueOrDie(); + EXPECT_FALSE(WhileLoopSimplifier().Run(m.get()).ValueOrDie()); } // Construct a loop where we assign a constant to tuple element 0 in each @@ -297,8 +308,8 @@ TEST_F(WhileLoopSimplifierTest, } )"; - ParseAndVerifyModule(hlo_string); - EXPECT_FALSE(WhileLoopSimplifier().Run(&module()).ValueOrDie()); + auto m = ParseAndReturnVerifiedModule(hlo_string).ValueOrDie(); + EXPECT_FALSE(WhileLoopSimplifier().Run(m.get()).ValueOrDie()); } // Nothing to simplify in a while loop whose tuple has 0 elements. @@ -320,8 +331,8 @@ TEST_F(WhileLoopSimplifierTest, LoopWithEmptyTupleNotSimplified) { } )"; - ParseAndVerifyModule(hlo_string); - EXPECT_FALSE(WhileLoopSimplifier().Run(&module()).ValueOrDie()); + auto m = ParseAndReturnVerifiedModule(hlo_string).ValueOrDie(); + EXPECT_FALSE(WhileLoopSimplifier().Run(m.get()).ValueOrDie()); } // While loop where one tuple element is used twice in the body, and thus can't @@ -348,8 +359,8 @@ TEST_F(WhileLoopSimplifierTest, LoopWithElemUsedTwiceNotSimplified) { } )"; - ParseAndVerifyModule(hlo_string); - EXPECT_FALSE(WhileLoopSimplifier().Run(&module()).ValueOrDie()); + auto m = ParseAndReturnVerifiedModule(hlo_string).ValueOrDie(); + EXPECT_FALSE(WhileLoopSimplifier().Run(m.get()).ValueOrDie()); } // This while loop has three tuple elements. Element 0 is unused and should be @@ -390,20 +401,18 @@ TEST_F(WhileLoopSimplifierTest, RemoveUnusedLoopOperands) { } )"; - ParseAndVerifyModule(hlo_string); - HloModule* the_module = &module(); - EXPECT_TRUE(WhileLoopSimplifier().Run(the_module).ValueOrDie()); + auto m = ParseAndReturnVerifiedModule(hlo_string).ValueOrDie(); + EXPECT_TRUE(WhileLoopSimplifier().Run(m.get()).ValueOrDie()); // 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(the_module->entry_computation()->instructions().begin(), - the_module->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( @@ -440,8 +449,8 @@ TEST_F(WhileLoopSimplifierTest, LoopWithNonTupleBodyShapeNotSimplified) { } )"; - ParseAndVerifyModule(hlo_string); - EXPECT_FALSE(WhileLoopSimplifier().Run(&module()).ValueOrDie()); + auto m = ParseAndReturnVerifiedModule(hlo_string).ValueOrDie(); + EXPECT_FALSE(WhileLoopSimplifier().Run(m.get()).ValueOrDie()); } TEST_F(WhileLoopSimplifierTest, @@ -473,8 +482,8 @@ TEST_F(WhileLoopSimplifierTest, } )"; - ParseAndVerifyModule(hlo_string); - EXPECT_FALSE(WhileLoopSimplifier().Run(&module()).ValueOrDie()); + auto m = ParseAndReturnVerifiedModule(hlo_string).ValueOrDie(); + EXPECT_FALSE(WhileLoopSimplifier().Run(m.get()).ValueOrDie()); } TEST_F(WhileLoopSimplifierTest, LoopWithArrayConstantNotSimplified) { @@ -505,8 +514,230 @@ TEST_F(WhileLoopSimplifierTest, LoopWithArrayConstantNotSimplified) { } )"; - ParseAndVerifyModule(hlo_string); - EXPECT_FALSE(WhileLoopSimplifier().Run(&module()).ValueOrDie()); + auto m = ParseAndReturnVerifiedModule(hlo_string).ValueOrDie(); + EXPECT_FALSE(WhileLoopSimplifier().Run(m.get()).ValueOrDie()); +} + +TEST_F(WhileLoopSimplifierTest, FlattenNestedTuple) { + const string hlo_string = R"( + HloModule Test + Body { + param = ((s32[1]), (s32[2], s32[3], (s32[4]))) parameter(0) + ta = (s32[1]) get-tuple-element(param), index=0 + a = s32[1] get-tuple-element(ta), index=0 + a.1 = s32[1] add(a, a) + tbcd = (s32[2], s32[3], (s32[4])) get-tuple-element(param), index=1 + ROOT tuple = ((s32[1]), (s32[2], s32[3], (s32[4]))) tuple(ta, tbcd) + } + Cond { + param = ((s32[1]), (s32[2], s32[3], (s32[4]))) parameter(0) + ROOT cond = pred[] constant(true) + } + ENTRY Loop { + a = s32[1] constant({0}) + b = s32[2] constant({0,1}) + c = s32[3] constant({0,1,2}) + d = s32[4] constant({0,1,2,3}) + ta = (s32[1]) tuple(a) + td = (s32[4]) tuple(d) + tbcd = (s32[2], s32[3], (s32[4])) tuple(b, c, td) + init = ((s32[1]), (s32[2], s32[3], (s32[4]))) tuple(ta, tbcd) + ROOT while = ((s32[1]), (s32[2], s32[3], (s32[4]))) while(init), + condition=Cond, body=Body + })"; + + auto m = ParseAndReturnVerifiedModule(hlo_string).ValueOrDie(); + EXPECT_TRUE(WhileLoopSimplifier().Run(m.get()).ValueOrDie()); + // DCE away the old loop so there's just one while loop in the module, making + // it easy to find. + EXPECT_TRUE(HloDCE().Run(m.get()).ok()); + + HloInstruction* new_while = FindFirstWhile(m.get()); + Shape flat_tuple = + ParseShape("(s32[1], s32[2], s32[3], s32[4])").ValueOrDie(); + SCOPED_TRACE(m->ToString()); + EXPECT_TRUE(ShapeUtil::Equal(new_while->shape(), flat_tuple)); + EXPECT_TRUE(ShapeUtil::Equal( + new_while->while_body()->root_instruction()->shape(), flat_tuple)); + EXPECT_TRUE(ShapeUtil::Equal( + new_while->while_body()->parameter_instruction(0)->shape(), flat_tuple)); + EXPECT_TRUE(ShapeUtil::Equal( + new_while->while_condition()->parameter_instruction(0)->shape(), + flat_tuple)); + EXPECT_TRUE(ShapeUtil::Equal( + m->entry_computation()->root_instruction()->shape(), + ParseShape("((s32[1]), (s32[2], s32[3], (s32[4])))").ValueOrDie())); +} + +// Edge-case: All elements of the loop carry are constants which can be removed, +// leaving us with a nullary loop. This is a special case, we just replace the +// loop with its init. +TEST_F(WhileLoopSimplifierTest, OnlyConstantsInLoopCarry) { + const string hlo_string = R"( + HloModule Test + Body { + param = (s32[1]) parameter(0) + a = s32[1] constant({0}) + ROOT tuple = (s32[1]) tuple(a) + } + Cond { + param = (s32[1]) parameter(0) + ROOT cond = pred[] constant(true) + } + ENTRY Loop { + a = s32[1] constant({0}) + init = (s32[1]) tuple(a) + ROOT while = (s32[1]) while(init), condition=Cond, body=Body + })"; + + auto m = ParseAndReturnVerifiedModule(hlo_string).ValueOrDie(); + EXPECT_TRUE(WhileLoopSimplifier().Run(m.get()).ValueOrDie()); + EXPECT_TRUE(HloDCE().Run(m.get()).ok()); + EXPECT_TRUE(TupleSimplifier().Run(m.get()).ok()); + EXPECT_THAT(m->entry_computation()->root_instruction(), + op::Tuple(op::Constant())); +} + +TEST_F(WhileLoopSimplifierTest, RemoveConstantFromLoopCarry) { + const string hlo_string = R"( + HloModule Test + Body { + param = (s32[1], s32[2], s32[3]) parameter(0) + a = s32[1] get-tuple-element(param), index=0 + a.1 = s32[1] add(a, a) + b = s32[2] constant({1,1}) + c = s32[3] constant({10,10,10}) + ROOT tuple = (s32[1], s32[2], s32[3]) tuple(a.1, b, c) + } + Cond { + param = (s32[1], s32[2], s32[3]) parameter(0) + /* Use each tuple element. The verifier will then ensure that if any of + * these get modified, they're replaced with values of the correct shape. */ + a = s32[1] get-tuple-element(param), index=0 + b = s32[2] get-tuple-element(param), index=1 + c = s32[3] get-tuple-element(param), index=2 + ROOT cond = pred[] constant(true) + } + ENTRY Loop { + /* Only `b` should be simplified away. `a` is not a constant within the + * loop, and `c`'s value changes depending on whether we run 0 or 1 + * iterations of the loop. */ + a = s32[1] constant({0}) + b = s32[2] constant({1,1}) + c = s32[3] constant({2,2,2}) + init = (s32[1], s32[2], s32[3]) tuple(a,b,c) + ROOT while = (s32[1], s32[2], s32[3]) while(init), + condition=Cond, body=Body + })"; + + auto m = ParseAndReturnVerifiedModule(hlo_string).ValueOrDie(); + EXPECT_TRUE(WhileLoopSimplifier().Run(m.get()).ValueOrDie()); + // DCE away the old loop so there's just one while loop in the module, making + // it easy to find. + EXPECT_TRUE(HloDCE().Run(m.get()).ok()); + // Run the tuple simplifier to make the resulting HLO a bit easier to check. + EXPECT_TRUE(TupleSimplifier().Run(m.get()).ok()); + + HloInstruction* new_while = FindFirstWhile(m.get()); + Shape new_while_shape = ParseShape("(s32[1], s32[3])").ValueOrDie(); + EXPECT_TRUE(ShapeUtil::Equal(new_while->shape(), new_while_shape)); + EXPECT_TRUE(ShapeUtil::Equal( + new_while->while_body()->root_instruction()->shape(), new_while_shape)); + EXPECT_TRUE(ShapeUtil::Equal( + new_while->while_body()->parameter_instruction(0)->shape(), + new_while_shape)); + EXPECT_TRUE(ShapeUtil::Equal( + new_while->while_condition()->parameter_instruction(0)->shape(), + new_while_shape)); + EXPECT_TRUE( + ShapeUtil::Equal(m->entry_computation()->root_instruction()->shape(), + ParseShape("(s32[1], s32[2], s32[3])").ValueOrDie())); + EXPECT_THAT(m->entry_computation()->root_instruction(), + op::Tuple(_, op::Constant(), _)); +} + +const char* const kSimpleMergeInductionVariablesModule = R"( + HloModule Test + Body { + param = (TYPE[], TYPE[], TYPE[]) parameter(0) + + a = TYPE[] get-tuple-element(param), index=0 + one = TYPE[] constant(1) + a1 = TYPE[] add(a, one) + + b = TYPE[] get-tuple-element(param), index=1 + negone = TYPE[] constant(-1) + b1 = TYPE[] add(b, negone) + + c = TYPE[] add(a, b) + + ROOT tuple = (TYPE[], TYPE[], TYPE[]) tuple(a1,b1,c) + } + Cond { + param = (TYPE[], TYPE[], TYPE[]) parameter(0) + a = TYPE[] get-tuple-element(param), index=0 + b = TYPE[] get-tuple-element(param), index=1 + sum = TYPE[] power(a, b) + ten = TYPE[] constant(10) + ROOT cond = pred[] less-than(sum, ten) + } + ENTRY Loop { + a = TYPE[] constant(10) + b = TYPE[] constant(100) + c = TYPE[] constant(0) + init = (TYPE[], TYPE[], TYPE[]) tuple(a,b,c) + while = (TYPE[], TYPE[], TYPE[]) while(init), condition=Cond, body=Body + + a1 = TYPE[] get-tuple-element(while), index=0 + b1 = TYPE[] get-tuple-element(while), index=1 + ROOT sum = TYPE[] add(a1, b1) + })"; + +TEST_F(WhileLoopSimplifierTest, MergeInductionVariables_Simple) { + string hlo_string = absl::StrReplaceAll(kSimpleMergeInductionVariablesModule, + {{"TYPE", "s32"}}); + + auto m = ParseAndReturnVerifiedModule(hlo_string).ValueOrDie(); + EXPECT_TRUE(WhileLoopSimplifier().Run(m.get()).ValueOrDie()); + // DCE away the old loop so there's just one while loop in the module, making + // it easy to find, and run the tuple simplifier to make the resulting HLO + // easier to check. + EXPECT_TRUE(HloDCE().Run(m.get()).ok()); + EXPECT_TRUE(TupleSimplifier().Run(m.get()).ok()); + + HloInstruction* new_while = FindFirstWhile(m.get()); + // We should have added a new loop counter for s32[] to the end of the tuple. + SCOPED_TRACE(m->ToString()); + Shape new_while_shape = + ParseShape("(s32[], s32[], s32[], s32[])").ValueOrDie(); + EXPECT_TRUE(ShapeUtil::Equal(new_while->shape(), new_while_shape)); + EXPECT_TRUE(ShapeUtil::Equal( + new_while->while_body()->root_instruction()->shape(), new_while_shape)); + EXPECT_TRUE(ShapeUtil::Equal( + new_while->while_body()->parameter_instruction(0)->shape(), + new_while_shape)); + EXPECT_TRUE(ShapeUtil::Equal( + new_while->while_condition()->parameter_instruction(0)->shape(), + new_while_shape)); + + EXPECT_THAT(new_while->while_body()->root_instruction(), + op::Tuple(op::GetTupleElement(op::Parameter(), 0), + op::GetTupleElement(op::Parameter(), 1), op::Add(), + op::Add(op::GetTupleElement(op::Parameter(), 3), + op::Constant()))); + EXPECT_THAT(new_while->while_condition()->root_instruction(), + op::Lt(op::Power(op::Add(), op::Add()), op::Constant())); +} + +// We shouldn't merge S16 induction variables; we can't create constants of this +// type because S16 literals are not implemented. +TEST_F(WhileLoopSimplifierTest, MergeInductionVariables_SkipS16) { + string hlo_string = absl::StrReplaceAll(kSimpleMergeInductionVariablesModule, + {{"TYPE", "s16"}}); + EXPECT_FALSE( + WhileLoopSimplifier() + .Run(ParseAndReturnVerifiedModule(hlo_string).ValueOrDie().get()) + .ValueOrDie()); } } // namespace diff --git a/tensorflow/compiler/xla/service/while_util.cc b/tensorflow/compiler/xla/service/while_util.cc index f90ac91f9d07aded8cafccf82dae894c9a149bd1..d77386497a14b3e52be2ea7f655fa330f60e4a97 100644 --- a/tensorflow/compiler/xla/service/while_util.cc +++ b/tensorflow/compiler/xla/service/while_util.cc @@ -15,6 +15,8 @@ limitations under the License. #include "tensorflow/compiler/xla/service/while_util.h" #include "absl/algorithm/container.h" +#include "absl/container/flat_hash_map.h" +#include "absl/container/inlined_vector.h" #include "absl/strings/str_cat.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" @@ -95,7 +97,7 @@ WidenWhileBody(HloComputation* narrow_body, const Shape& wide_shape) { WhileUtil::MakeInstructionsLiveIn( HloInstruction* while_instr, absl::Span instructions) { - CHECK(ShapeUtil::IsTuple(while_instr->shape())); + CHECK(while_instr->shape().IsTuple()); int64 elements_in_old_while_shape = while_instr->shape().tuple_shapes_size(); Shape new_while_shape = while_instr->shape(); @@ -225,7 +227,8 @@ static Shape MakeLoopStateShape(const WhileUtil::LoopStateTy& init_values) { /*static*/ StatusOr WhileUtil::MakeCountedLoop( HloComputation* computation, int32 trip_count, const WhileUtil::LoopStateTy& init_values, - const WhileUtil::LoopBodyGeneratorTy& loop_body_generator) { + const WhileUtil::LoopBodyGeneratorTy& loop_body_generator, + const OpMetadata& metadata) { CHECK_GE(trip_count, 0); Shape loop_state_shape = MakeLoopStateShape(init_values); @@ -242,6 +245,7 @@ static Shape MakeLoopStateShape(const WhileUtil::LoopStateTy& init_values) { computation->AddInstruction(HloInstruction::CreateWhile( loop_state_shape, module->AddEmbeddedComputation(std::move(cond)), module->AddEmbeddedComputation(std::move(body)), init_tuple)); + while_instr->set_metadata(metadata); std::vector result; for (int64 i = 0, e = init_values.size(); i < e; i++) { @@ -268,4 +272,17 @@ static Shape MakeLoopStateShape(const WhileUtil::LoopStateTy& init_values) { return result; } +/*static*/ absl::flat_hash_map> +WhileUtil::GetGTEsMapForWhileConditional( + const HloComputation& while_conditional) { + absl::flat_hash_map> result; + for (HloInstruction* user : + while_conditional.parameter_instruction(0)->users()) { + if (user->opcode() == HloOpcode::kGetTupleElement) { + result[user->tuple_index()].push_back(user); + } + } + return result; +} + } // namespace xla diff --git a/tensorflow/compiler/xla/service/while_util.h b/tensorflow/compiler/xla/service/while_util.h index b1c4486887ae0ddbe2ba4e79f45a265689111017..cba41ccd8b184ba3d867bc170724aee71e777788 100644 --- a/tensorflow/compiler/xla/service/while_util.h +++ b/tensorflow/compiler/xla/service/while_util.h @@ -16,6 +16,8 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_WHILE_UTIL_H_ #define TENSORFLOW_COMPILER_XLA_SERVICE_WHILE_UTIL_H_ +#include "absl/container/flat_hash_map.h" +#include "absl/container/inlined_vector.h" #include "tensorflow/compiler/xla/service/call_inliner.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" @@ -77,13 +79,21 @@ class WhileUtil { static StatusOr MakeCountedLoop( HloComputation* computation, int32 trip_count, const LoopStateTy& init_values, - const LoopBodyGeneratorTy& loop_body_generator); + const LoopBodyGeneratorTy& loop_body_generator, + const OpMetadata& metadata); // Returns the GetTupleElement instructions in `while_body` that access // elements in the parameter tuple that don't change across iterations. // Assumes `while_body` is the body computation of the while loop in question. static std::vector GetInvariantGTEsForWhileBody( const HloComputation& while_body); + + // Returns a map of index to GetTupleElement instructions in + // `while_conditional` that access elements in the parameter tuple. Assumes + // `while_conditional` is the conditional computation of the while loop in + // question. + static absl::flat_hash_map> + GetGTEsMapForWhileConditional(const HloComputation& while_conditional); }; } // namespace xla diff --git a/tensorflow/compiler/xla/service/while_util_test.cc b/tensorflow/compiler/xla/service/while_util_test.cc index 5e6941933330fde29bc9c779aae4bb3c36914660..d92b9870f373564ae8fd904c8bf9f0d1afbff9c4 100644 --- a/tensorflow/compiler/xla/service/while_util_test.cc +++ b/tensorflow/compiler/xla/service/while_util_test.cc @@ -180,8 +180,8 @@ body { cond { param.c = (s32[], s32[]) parameter(0) - token = token[] after-all() - infeed = (pred[], token[]) infeed(token) + token0 = token[] after-all() + infeed = (pred[], token[]) infeed(token0) ROOT condition = pred[] get-tuple-element(infeed), index=0 } diff --git a/tensorflow/compiler/xla/service/zero_sized_hlo_elimination.cc b/tensorflow/compiler/xla/service/zero_sized_hlo_elimination.cc index 83d696fe0915086c3c98b6d7cbdaeaeb4d9d0bdb..661b7aa7d99ca549da6a509812760a1665d60919 100644 --- a/tensorflow/compiler/xla/service/zero_sized_hlo_elimination.cc +++ b/tensorflow/compiler/xla/service/zero_sized_hlo_elimination.cc @@ -31,16 +31,21 @@ StatusOr ZeroSizedHloElimination::Run(HloModule* module) { bool changed = false; for (HloComputation* comp : module->MakeNonfusionComputations()) { for (HloInstruction* instruction : comp->MakeInstructionPostOrder()) { - if (instruction->HasSideEffect() || - !ShapeUtil::IsArray(instruction->shape()) || + if (instruction->HasSideEffect() || !instruction->shape().IsArray() || instruction->opcode() == HloOpcode::kConstant) { continue; } 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 b9ef18892d7aa859f6b0b505db4c004e4f5c5066..572a79609e7a912277af0fd2ba43f9a1e14a6f52 100644 --- a/tensorflow/compiler/xla/service/zero_sized_hlo_elimination_test.cc +++ b/tensorflow/compiler/xla/service/zero_sized_hlo_elimination_test.cc @@ -45,7 +45,8 @@ class ZeroSizedHloEliminationTest : public HloTestBase { 0, ShapeUtil::MakeShape(F32, {3, 0}), "zero sized param"))) {} StatusOr RunZeroSizedElimination() { - auto module = CreateNewModule("zero_sized_elimination_test_module"); + auto module = + CreateNewUnverifiedModule("zero_sized_elimination_test_module"); module->AddEntryComputation(builder_.Build()); return ZeroSizedHloElimination{}.Run(module.get()); } @@ -81,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/service_interface.h b/tensorflow/compiler/xla/service_interface.h index 14c35e7b84f07bebac33a9753ac26a8ee1418f1e..33edbd1b20d01bf132f2a152625d5f49a45f26f9 100644 --- a/tensorflow/compiler/xla/service_interface.h +++ b/tensorflow/compiler/xla/service_interface.h @@ -47,8 +47,11 @@ class ServiceInterface { virtual Status ResetDevice(const ResetDeviceRequest* arg, ResetDeviceResponse* result) = 0; - virtual Status ExecuteGraph(const ExecuteGraphRequest* arg, - ExecuteResponse* result) = 0; + virtual Status Compile(const CompileRequest* arg, + CompileResponse* result) = 0; + + virtual Status Execute(const ExecuteRequest* arg, + ExecuteResponse* result) = 0; virtual Status ExecuteGraphParallel(const ExecuteGraphParallelRequest* arg, ExecuteParallelResponse* result) = 0; diff --git a/tensorflow/compiler/xla/shape.cc b/tensorflow/compiler/xla/shape.cc new file mode 100644 index 0000000000000000000000000000000000000000..a36d3547a0987422c2658b0f3046f7b1f83369c6 --- /dev/null +++ b/tensorflow/compiler/xla/shape.cc @@ -0,0 +1,231 @@ +/* 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/shape.h" + +#include "absl/strings/str_cat.h" +#include "absl/strings/str_join.h" +#include "tensorflow/compiler/xla/shape_util.h" + +namespace xla { + +Shape::Shape(const ShapeProto& shape_proto) { + set_element_type(shape_proto.element_type()); + dimensions_.reserve(shape_proto.dimensions_size()); + for (const int64 dimension : shape_proto.dimensions()) { + add_dimensions(dimension); + } + // 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()) { + LOG(ERROR) << "Malformed shape proto: number of is_dynamic_dimension " + "fields does not match number of dimension fields"; + } + 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()); + for (const ShapeProto& element_shape : shape_proto.tuple_shapes()) { + *add_tuple_shapes() = Shape(element_shape); + } + if (shape_proto.has_layout()) { + *mutable_layout() = Layout::CreateFromProto(shape_proto.layout()); + } +} + +ShapeProto Shape::ToProto() const { + ShapeProto proto; + proto.set_element_type(element_type_); + proto.mutable_dimensions()->Reserve(dimensions_size()); + for (const int64 dimension : dimensions()) { + proto.add_dimensions(dimension); + } + for (const bool dynamic : dynamic_dimensions_) { + proto.add_is_dynamic_dimension(dynamic); + } + proto.mutable_tuple_shapes()->Reserve(tuple_shapes_size()); + for (const Shape& shape : tuple_shapes()) { + *proto.add_tuple_shapes() = shape.ToProto(); + } + if (has_layout()) { + *proto.mutable_layout() = layout().ToProto(); + } + return proto; +} + +string Shape::ToString(bool print_layout) const { + if (print_layout) { + return ShapeUtil::HumanStringWithLayout(*this); + } else { + return ShapeUtil::HumanString(*this); + } +} + +bool Shape::is_static() const { + if (IsTuple()) { + for (const Shape& subshape : tuple_shapes_) { + if (!subshape.is_static()) { + return false; + } + } + } + return !absl::c_any_of(dynamic_dimensions_, [](bool b) { return b; }); +} + +void Shape::DeleteDimension(int64 dim_to_delete) { + CHECK(IsArray()); + CHECK_GE(dim_to_delete, 0); + CHECK_LT(dim_to_delete, dimensions_.size()); + dimensions_.erase(dimensions_.begin() + dim_to_delete); + dynamic_dimensions_.erase(dynamic_dimensions_.begin() + dim_to_delete); + if (LayoutUtil::HasLayout(*this)) { + layout_.set_format(DENSE); + for (int64 i = 0; i < layout_.minor_to_major().size();) { + if (layout_.minor_to_major(i) == dim_to_delete) { + layout_.mutable_minor_to_major()->erase( + layout_.mutable_minor_to_major()->begin() + i); + continue; + } + if (layout_.minor_to_major(i) > dim_to_delete) { + (*layout_.mutable_minor_to_major())[i] -= 1; + } + ++i; + } + } +} + +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 (!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; + } + + 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; +} + +ProgramShape::ProgramShape(const ProgramShapeProto& program_shape_proto) { + for (const ShapeProto& shape_proto : program_shape_proto.parameters()) { + *add_parameters() = Shape(shape_proto); + } + *mutable_result() = Shape(program_shape_proto.result()); + for (const string& name : program_shape_proto.parameter_names()) { + add_parameter_names(name); + } +} + +ProgramShapeProto ProgramShape::ToProto() const { + ProgramShapeProto proto; + for (const Shape& shape : parameters()) { + *proto.add_parameters() = shape.ToProto(); + } + *proto.mutable_result() = result().ToProto(); + for (const string& name : parameter_names()) { + proto.add_parameter_names(name); + } + return proto; +} + +string ProgramShape::ToString() const { + std::vector parameter_strings(parameters_size()); + for (int i = 0; i < parameters_size(); ++i) { + parameter_strings[i] = absl::StrCat( + i < parameter_names_size() ? parameter_names(i) : "(unknown)", ": ", + ShapeUtil::HumanString(parameters(i))); + } + return absl::StrCat("(", absl::StrJoin(parameter_strings, ", "), ") -> ", + ShapeUtil::HumanString(result())); +} + +std::ostream& operator<<(std::ostream& out, const ProgramShape& program_shape) { + out << program_shape.ToString() << "\n"; + return out; +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/shape.h b/tensorflow/compiler/xla/shape.h new file mode 100644 index 0000000000000000000000000000000000000000..e6b4e872f69e16ea407dc18cadfc83618080084f --- /dev/null +++ b/tensorflow/compiler/xla/shape.h @@ -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. +==============================================================================*/ + +#ifndef TENSORFLOW_COMPILER_XLA_SHAPE_H_ +#define TENSORFLOW_COMPILER_XLA_SHAPE_H_ + +#include +#include + +#include "absl/types/optional.h" +#include "tensorflow/compiler/xla/layout.h" +#include "tensorflow/compiler/xla/primitive_util.h" +#include "tensorflow/compiler/xla/types.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/platform/types.h" + +namespace xla { + +// A shape describes the number of dimensions in a array, the bounds of each +// dimension, and the primitive component type. For tuples, shape describes the +// structure (number of elements and nesting). +class Shape { + public: + Shape() = default; + + // Construct a shape from a ShapeProto. + explicit Shape(const ShapeProto& shape_proto); + + // Returns a ShapeProto representation of the Shape. + ShapeProto ToProto() const; + + // Returns a human-readable string that represents the given shape, with or + // without layout. e.g. "F32[42,12] {0, 1}" or "F32[64]". + string ToString(bool print_layout = false) const; + + // Returns the rank (number of dimensions) of the given shape. Shape must be + // an array. + int64 rank() const { + CHECK(IsArray()) << "Non-arrays do not have a rank, shape: " << ToString(); + return dimensions_.size(); + } + + // Returns whether the shape is of the specified type (array, tuple, etc). + bool IsArray() const { return primitive_util::IsArrayType(element_type()); } + bool IsTuple() const { return element_type() == TUPLE; } + bool IsToken() const { return element_type() == TOKEN; } + bool IsOpaque() const { return element_type() == OPAQUE; } + + // Returns true if no array dimension in the shape is dynamically sized. Tuple + // shapes are traversed recursively. + bool is_static() const; + + // Returns true if the given dimension is dynamically-sized. + bool is_dynamic_dimension(int dimension) const { + return dynamic_dimensions_.at(dimension); + } + + // Sets whether or not the given dimension is dynamically-sized. + void set_dynamic_dimension(int dimension, bool is_dynamic) { + 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 + // adjusted to match the modified shape. + void DeleteDimension(int64 dim_to_delete); + + // The following methods mirror the protobuf generated code interface for the + // message ShapeProto. This enabled easy migration of this data structure + // from a proto to a proper C++ class. + // TODO(b/29771030): Replace or augment these methods with a more ergonomic + // interface. + + // Methods for accessing the primitive type. + PrimitiveType element_type() const { return element_type_; } + void set_element_type(PrimitiveType value) { element_type_ = value; } + + // Methods for accessing the dimensions array. + int dimensions_size() const { return dimensions_.size(); } + int64 dimensions(int index) const { return dimensions_.at(index); } + void set_dimensions(int index, int64 value) { dimensions_.at(index) = value; } + void add_dimensions(int64 value) { + dimensions_.push_back(value); + dynamic_dimensions_.push_back(false); + } + void clear_dimensions() { + dimensions_.clear(); + dynamic_dimensions_.clear(); + } + const std::vector& dimensions() const { return dimensions_; } + absl::Span mutable_dimensions() { return absl::MakeSpan(dimensions_); } + + // Methods for accessing the tuple subshapes. This field only non-empty for + // tuple shapes. + int tuple_shapes_size() const { return tuple_shapes_.size(); } + const Shape& tuple_shapes(int index) const { return tuple_shapes_.at(index); } + Shape* mutable_tuple_shapes(int index) { return &tuple_shapes_.at(index); } + Shape* add_tuple_shapes() { + tuple_shapes_.push_back(Shape()); + return &tuple_shapes_.back(); + } + void clear_tuple_shapes() { tuple_shapes_.clear(); } + const std::vector& tuple_shapes() const { return tuple_shapes_; } + std::vector* mutable_tuple_shapes() { return &tuple_shapes_; } + + // Methods for accessing the layout field. + bool has_layout() const { return layout_.format() != INVALID_FORMAT; } + const Layout& layout() const { return layout_; } + Layout* mutable_layout() { return &layout_; } + void clear_layout() { layout_.Clear(); } + + void Swap(Shape* other) { + using std::swap; + swap(*this, *other); + } + + void Clear() { + element_type_ = PRIMITIVE_TYPE_INVALID; + dimensions_.clear(); + tuple_shapes_.clear(); + clear_layout(); + } + + string SerializeAsString() const { return ToProto().SerializeAsString(); } + 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; + }; + + private: + // The element type of this shape (tuple, array, etc). + PrimitiveType element_type_ = PRIMITIVE_TYPE_INVALID; + + // The array bounds of the dimensions. This is nonempty only for array + // shapes. For a dynamically-sized dimension, the respective value in this + // vector is an inclusive upper limit of the array bound. + std::vector dimensions_; + + // This vector is the same size as 'dimensions_' and indicates whether the + // respective dimension is dynamically sized. + std::vector dynamic_dimensions_; + + // The tuple element subshapes. This is nonempty only for tuple shapes. + std::vector tuple_shapes_; + + // The layout of the shape. Only relevant for arrays. + Layout layout_; +}; + +// Shape of the parameters and output of an XLA computation. This is analogous +// to a traditional function signature. +class ProgramShape { + public: + ProgramShape() = default; + + // Creates a ProgramShape from a ProgramShapeProto protobuf. + explicit ProgramShape(const ProgramShapeProto& program_shape_proto); + + // Returns a proto representation of the object. + ProgramShapeProto ToProto() const; + + string ToString() const; + + // The following methods mirror the protobuf generated code interface for the + // message ProgramShapeProto. This enabled easy migration of this data + // structure from a proto to a proper C++ class. + // TODO(b/29771030): Replace or augment these methods with a more ergonomic + // interface. + + // Methods for accessing and manipulating the Shape of the parameters. + int parameters_size() const { return parameters_.size(); } + const Shape& parameters(int index) const { return parameters_.at(index); } + Shape* mutable_parameters(int index) { return ¶meters_.at(index); } + Shape* add_parameters() { + parameters_.emplace_back(); + return ¶meters_.back(); + } + void clear_parameters() { parameters_.clear(); } + const std::vector& parameters() const { return parameters_; } + std::vector* mutable_parameters() { return ¶meters_; } + + // Methods for accessing and manipulating the Shape of the result. + const Shape& result() const { return result_; } + Shape* mutable_result() { return &result_; } + + // Methods for accessing and manipulating the names of the parameters. + int parameter_names_size() const { return parameter_names_.size(); } + const string& parameter_names(int index) const { + return parameter_names_.at(index); + } + void set_parameter_names(int index, const string& value) { + parameter_names_.at(index) = value; + } + string* mutable_parameter_names(int index) { + return ¶meter_names_.at(index); + } + void add_parameter_names(const string& value) { + parameter_names_.push_back(value); + } + string* add_parameter_names() { + parameter_names_.push_back(""); + return ¶meter_names_.back(); + } + void clear_parameter_names() { parameter_names_.clear(); } + const std::vector& parameter_names() const { + return parameter_names_; + } + std::vector* mutable_parameter_names() { return ¶meter_names_; } + + string ShortDebugString() const { return ToProto().ShortDebugString(); } + string DebugString() const { return ToProto().DebugString(); } + + private: + // The shapes of the parameters of the computation represented by this object. + std::vector parameters_; + + // The names of the parameters of the computation represented by this object. + std::vector parameter_names_; + + // The shape of the result of the computation represented by this object. + Shape result_; +}; + +std::ostream& operator<<(std::ostream& out, const Shape& shape); +std::ostream& operator<<(std::ostream& out, const ProgramShape& program_shape); + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SHAPE_H_ diff --git a/tensorflow/compiler/xla/shape_layout.cc b/tensorflow/compiler/xla/shape_layout.cc index d44db89d571891ecef554cd45c050017833982bb..a000886d60d06a4a598910c901accb6dfd0a8f1a 100644 --- a/tensorflow/compiler/xla/shape_layout.cc +++ b/tensorflow/compiler/xla/shape_layout.cc @@ -52,7 +52,7 @@ bool ShapeLayout::MatchesLayoutInShape(const Shape& shape) const { const Layout& ShapeLayout::layout() const { CHECK(LayoutIsSet()); - CHECK(!ShapeUtil::IsTuple(shape_)); + CHECK(!shape_.IsTuple()); return shape_.layout(); } @@ -61,15 +61,15 @@ void ShapeLayout::Clear() { LayoutUtil::ClearLayout(&shape_); } bool ShapeLayout::LayoutIsSet() const { return LayoutUtil::HasLayout(shape_); } void ShapeLayout::ResetLayout(const Layout& layout) { - CHECK(!ShapeUtil::IsTuple(shape_)); - CHECK(!ShapeUtil::IsOpaque(shape_)); + CHECK(!shape_.IsTuple()); + CHECK(!shape_.IsOpaque()); *shape_.mutable_layout() = layout; TF_CHECK_OK(ShapeUtil::ValidateShape(shape_)); } void ShapeLayout::ResetLayout(const Layout& layout, ShapeIndexView shape_index) { - CHECK(ShapeUtil::IsTuple(shape_)); + CHECK(shape_.IsTuple()); *ShapeUtil::GetMutableSubshape(&shape_, shape_index)->mutable_layout() = layout; TF_CHECK_OK(ShapeUtil::ValidateShape(shape_)); diff --git a/tensorflow/compiler/xla/shape_test.cc b/tensorflow/compiler/xla/shape_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..55ce5fe884e98e474253be9ef694f1b8137b4b01 --- /dev/null +++ b/tensorflow/compiler/xla/shape_test.cc @@ -0,0 +1,192 @@ +/* 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/compiler/xla/shape.h" + +#include +#include "absl/strings/str_cat.h" +#include "absl/strings/str_join.h" +#include "tensorflow/compiler/xla/layout_util.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/status_macros.h" +#include "tensorflow/compiler/xla/test.h" +#include "tensorflow/compiler/xla/test_helpers.h" +#include "tensorflow/compiler/xla/types.h" +#include "tensorflow/compiler/xla/util.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" + +namespace xla { +namespace { + +class ShapeTest : public ::testing::Test { + protected: + const Shape opaque_ = ShapeUtil::MakeOpaqueShape(); + const Shape token_ = ShapeUtil::MakeTokenShape(); + const Shape scalar_ = ShapeUtil::MakeShape(F32, {}); + const Shape matrix_ = ShapeUtil::MakeShape(U32, {1, 2}); + const Shape matrix2_ = ShapeUtil::MakeShapeWithLayout(S32, {3, 4}, {0, 1}); + const Shape tuple_ = + ShapeUtil::MakeTupleShape({opaque_, scalar_, matrix_, matrix2_}); + const Shape nested_tuple_ = + ShapeUtil::MakeTupleShape({tuple_, matrix_, token_}); + const Shape dyanmic_matrix_ = + ShapeUtil::MakeShape(S32, {5, 2}, {true, false}); +}; + +TEST_F(ShapeTest, ShapeToFromProto) { + for (const Shape& shape : {opaque_, token_, scalar_, matrix_, matrix2_, + tuple_, nested_tuple_, dyanmic_matrix_}) { + Shape shape_copy(shape.ToProto()); + EXPECT_TRUE(ShapeUtil::Equal(shape, shape_copy)) + << shape << " != " << shape_copy; + } +} + +TEST_F(ShapeTest, ShapeToString) { + EXPECT_EQ("opaque[]", opaque_.ToString()); + EXPECT_EQ("token[]", token_.ToString()); + EXPECT_EQ("f32[]", scalar_.ToString()); + EXPECT_EQ("u32[1,2]", matrix_.ToString()); + EXPECT_EQ("s32[3,4]", matrix2_.ToString()); + EXPECT_EQ("(opaque[], f32[], u32[1,2], s32[3,4])", tuple_.ToString()); + EXPECT_EQ("((opaque[], f32[], u32[1,2], s32[3,4]), u32[1,2], token[])", + nested_tuple_.ToString()); + + EXPECT_EQ("opaque[]", opaque_.ToString(/*print_layout=*/true)); + EXPECT_EQ("f32[]", scalar_.ToString(/*print_layout=*/true)); + EXPECT_EQ("u32[1,2]{1,0}", matrix_.ToString(/*print_layout=*/true)); + EXPECT_EQ("s32[3,4]{0,1}", matrix2_.ToString(/*print_layout=*/true)); + EXPECT_EQ("(opaque[], f32[], u32[1,2]{1,0}, s32[3,4]{0,1})", + tuple_.ToString(/*print_layout=*/true)); + EXPECT_EQ( + "((opaque[], f32[], u32[1,2]{1,0}, s32[3,4]{0,1}), u32[1,2]{1,0}, " + "token[])", + nested_tuple_.ToString(/*print_layout=*/true)); +} + +TEST_F(ShapeTest, DynamicShapeToString) { + Shape array_shape = + ShapeUtil::MakeShape(F32, {23, 44, 55}, {true, false, true}); + EXPECT_EQ("f32[<=23,44,<=55]", array_shape.ToString()); + + array_shape.set_dynamic_dimension(2, false); + EXPECT_EQ("f32[<=23,44,55]", array_shape.ToString()); +} + +TEST_F(ShapeTest, IsStatic) { + EXPECT_TRUE(opaque_.is_static()); + EXPECT_TRUE(token_.is_static()); + EXPECT_TRUE(matrix_.is_static()); + EXPECT_TRUE(tuple_.is_static()); + EXPECT_TRUE(nested_tuple_.is_static()); + + Shape dynamic_matrix = matrix_; + EXPECT_TRUE(dynamic_matrix.is_static()); + dynamic_matrix.set_dynamic_dimension(1, true); + EXPECT_FALSE(dynamic_matrix.is_static()); + + Shape dynamic_tuple = tuple_; + EXPECT_TRUE(dynamic_tuple.is_static()); + ShapeUtil::GetMutableSubshape(&dynamic_tuple, {2}) + ->set_dynamic_dimension(1, true); + EXPECT_FALSE(dynamic_tuple.is_static()); +} + +TEST_F(ShapeTest, IsDynamicDimension) { + Shape dynamic_matrix = matrix_; + dynamic_matrix.set_dynamic_dimension(1, true); + EXPECT_FALSE(dynamic_matrix.is_dynamic_dimension(0)); + EXPECT_TRUE(dynamic_matrix.is_dynamic_dimension(1)); + + Shape dynamic_tuple = tuple_; + EXPECT_TRUE(dynamic_tuple.is_static()); + ShapeUtil::GetMutableSubshape(&dynamic_tuple, {2}) + ->set_dynamic_dimension(1, true); + EXPECT_FALSE(dynamic_tuple.is_static()); +} + +TEST_F(ShapeTest, ProgramShapeToFromProto) { + ProgramShape program_shape; + *program_shape.add_parameters() = ShapeUtil::MakeShape(F32, {1, 2, 3}); + *program_shape.add_parameters() = ShapeUtil::MakeTokenShape(); + *program_shape.add_parameters() = ShapeUtil::MakeShape(S64, {}); + *program_shape.add_parameters() = ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeShape(S32, {}), + ShapeUtil::MakeTupleShape({ShapeUtil::MakeTokenShape()}), + ShapeUtil::MakeShape(F32, {42, 42})}); + + *program_shape.mutable_result() = ShapeUtil::MakeShape(F32, {7}); + + program_shape.add_parameter_names("foo"); + program_shape.add_parameter_names("bar"); + program_shape.add_parameter_names("baz"); + program_shape.add_parameter_names("qux qux"); + + // Create a copy of the program shape by round-tripping through a proto. + ProgramShape program_shape_copy(program_shape.ToProto()); + ASSERT_EQ(program_shape.parameters_size(), + program_shape_copy.parameters_size()); + for (int i = 0; i < program_shape.parameters_size(); ++i) { + EXPECT_TRUE(ShapeUtil::Equal(program_shape.parameters(i), + program_shape_copy.parameters(i))); + } + + EXPECT_TRUE( + ShapeUtil::Equal(program_shape.result(), program_shape_copy.result())); + + ASSERT_EQ(program_shape.parameter_names_size(), + program_shape_copy.parameter_names_size()); + for (int i = 0; i < program_shape.parameter_names_size(); ++i) { + EXPECT_EQ(program_shape.parameter_names(i), + program_shape_copy.parameter_names(i)); + } +} + +TEST_F(ShapeTest, ProgramShapeToString) { + ProgramShape prog = ShapeUtil::MakeProgramShape( + {opaque_, scalar_, matrix_, matrix2_, tuple_, nested_tuple_}, + nested_tuple_); + EXPECT_EQ( + "((unknown): opaque[], " + "(unknown): f32[], " + "(unknown): u32[1,2], " + "(unknown): s32[3,4], " + "(unknown): (opaque[], f32[], u32[1,2], s32[3,4]), " + "(unknown): ((opaque[], f32[], u32[1,2], s32[3,4]), u32[1,2], token[])) " + "-> " + "((opaque[], f32[], u32[1,2], s32[3,4]), u32[1,2], token[])", + prog.ToString()); + + prog.add_parameter_names("arg0"); + prog.add_parameter_names("scalar"); + prog.add_parameter_names("matrix"); + prog.add_parameter_names("matrix2"); + prog.add_parameter_names("tuple"); + prog.add_parameter_names("nested_tuple"); + EXPECT_EQ( + "(arg0: opaque[], " + "scalar: f32[], " + "matrix: u32[1,2], " + "matrix2: s32[3,4], " + "tuple: (opaque[], f32[], u32[1,2], s32[3,4]), " + "nested_tuple: ((opaque[], f32[], u32[1,2], s32[3,4]), u32[1,2], " + "token[])) " + "-> " + "((opaque[], f32[], u32[1,2], s32[3,4]), u32[1,2], token[])", + prog.ToString()); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/shape_tree.h b/tensorflow/compiler/xla/shape_tree.h index df610102b4c7fa08c0b7030124939009130f89f4..089120179e2a77518eb5b18c11a35670b03e9b77 100644 --- a/tensorflow/compiler/xla/shape_tree.h +++ b/tensorflow/compiler/xla/shape_tree.h @@ -395,7 +395,7 @@ class ShapeTreeIterator template int64 ShapeTree::CountSubshapes(const Shape& shape) { int64 current_count = 1; - if (ShapeUtil::IsTuple(shape)) { + if (shape.IsTuple()) { int64 count = ShapeUtil::TupleElementCount(shape); for (int i = 0; i < count; ++i) { current_count += CountSubshapes(shape.tuple_shapes(i)); @@ -407,7 +407,7 @@ int64 ShapeTree::CountSubshapes(const Shape& shape) { template void ShapeTree::InitChildren(const Shape& shape, const T& init_value, Node* node, Index* index) { - if (ShapeUtil::IsTuple(shape)) { + if (shape.IsTuple()) { const int64 size = ShapeUtil::TupleElementCount(shape); #ifndef NDEBUG index->children_count = size; @@ -443,7 +443,7 @@ void ShapeTree::InitChildren(const Shape& shape, const T& init_value, template void ShapeTree::InitChildren(const Shape& shape, Node* node, Index* index) { - if (ShapeUtil::IsTuple(shape)) { + if (shape.IsTuple()) { const int64 size = ShapeUtil::TupleElementCount(shape); #ifndef NDEBUG index->children_count = size; @@ -667,12 +667,11 @@ void ShapeTree::CopySubtreeFrom(const ShapeTree& other, template bool ShapeTree::operator==(const ShapeTree& other) const { bool equal = true; - ForEachElement( - [this, &other, &equal](const ShapeIndex& index, const T& data) { - if (data != other.element(index)) { - equal = false; - } - }); + ForEachElement([&other, &equal](const ShapeIndex& index, const T& data) { + if (data != other.element(index)) { + equal = false; + } + }); return equal; } diff --git a/tensorflow/compiler/xla/shape_tree_test.cc b/tensorflow/compiler/xla/shape_tree_test.cc index c8ff55e7845785d9292516b823fb591cc28cbfad..2b6c484bc4f205be0180403eeac2dd391029b110 100644 --- a/tensorflow/compiler/xla/shape_tree_test.cc +++ b/tensorflow/compiler/xla/shape_tree_test.cc @@ -52,10 +52,10 @@ class ShapeTreeTest : public ::testing::Test { TEST_F(ShapeTreeTest, DefaultConstructor) { ShapeTree int_tree; - EXPECT_TRUE(ShapeUtil::IsNil(int_tree.shape())); + EXPECT_TRUE(ShapeUtil::IsEmptyTuple(int_tree.shape())); ShapeTree bool_tree; - EXPECT_TRUE(ShapeUtil::IsNil(bool_tree.shape())); + EXPECT_TRUE(ShapeUtil::IsEmptyTuple(bool_tree.shape())); } void ShapeTreeTest::TestShapeConstructor(const Shape& shape, diff --git a/tensorflow/compiler/xla/shape_util.cc b/tensorflow/compiler/xla/shape_util.cc index 43b10a4af4d6f6c29a02d10e3e2d461064da30ea..1ada4bc0362f86bc770d4adfcd4d4b0ff7379c77 100644 --- a/tensorflow/compiler/xla/shape_util.cc +++ b/tensorflow/compiler/xla/shape_util.cc @@ -74,68 +74,17 @@ std::ostream& operator<<(std::ostream& out, const ShapeIndexView& shape_index) { return out; } -namespace { - -// Returns whether the given primitive type corresponds to an array shape. -bool IsArrayPrimitiveType(PrimitiveType primitive_type) { - return primitive_type != PRIMITIVE_TYPE_INVALID && primitive_type != TUPLE && - primitive_type != OPAQUE && primitive_type != TOKEN; -} - -// 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 (ShapeUtil::IsTuple(lhs)) { - 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 (!ShapeUtil::IsArray(lhs)) { - // 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; - } - if (!absl::c_equal(lhs.layout().padded_dimensions(), - rhs.layout().padded_dimensions())) { - VLOG(3) - << "CompareShapes: lhs padded_dimensions != rhs padded_dimensions"; - return false; - } - if (lhs.layout().padding_value() != rhs.layout().padding_value()) { - VLOG(3) << "CompareShapes: lhs padding value != rhs padding_value"; - return false; - } - } - } +bool ShapeIndexView::StartsWith(ShapeIndexView prefix) const { + return size() >= prefix.size() && + indices_.subspan(0, prefix.size()) == prefix.indices_; +} - if (!ShapeUtil::SameDimensions(lhs, rhs)) { - VLOG(3) << "CompareShapes: lhs dimensions != rhs dimensions"; - return false; - } - return true; +/* static */ bool ShapeUtil::IsArrayPrimitiveType( + PrimitiveType primitive_type) { + return primitive_util::IsArrayType(primitive_type); } +namespace { // Constructs and returns the new shape with the given minor_to_major order in // its Layout. StatusOr MakeShapeWithLayoutInternal( @@ -149,11 +98,12 @@ StatusOr MakeShapeWithLayoutInternal( return InvalidArgument("Unsupported element type: %s", PrimitiveType_Name(element_type)); } - Shape shape = ShapeUtil::MakeShape(element_type, dimensions); + TF_ASSIGN_OR_RETURN(Shape shape, + ShapeUtil::MakeValidatedShape(element_type, dimensions)); auto min2maj = shape.mutable_layout()->mutable_minor_to_major(); - min2maj->Clear(); + min2maj->clear(); for (int64 value : minor_to_major) { - min2maj->Add(value); + min2maj->push_back(value); } if (!shape.has_layout()) { return InvalidArgument("Shape has no layout."); @@ -161,12 +111,11 @@ StatusOr MakeShapeWithLayoutInternal( 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(); @@ -177,8 +126,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(); @@ -187,12 +135,6 @@ StatusOr MakeShapeWithLayoutInternal( return equal; } -/* static */ int64 ShapeUtil::Rank(const Shape& shape) { - CHECK(ShapeUtil::IsArray(shape)) - << "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()) { @@ -207,7 +149,7 @@ StatusOr MakeShapeWithLayoutInternal( /* static */ ProgramShape ShapeUtil::MakeProgramShape( std::initializer_list parameters, Shape result) { ProgramShape program_shape; - for (const auto& shape : parameters) { + for (const Shape& shape : parameters) { *program_shape.add_parameters() = shape; } *program_shape.mutable_result() = std::move(result); @@ -219,14 +161,32 @@ StatusOr MakeShapeWithLayoutInternal( return MakeValidatedShape(element_type, dimensions).ValueOrDie(); } +/* static */ Shape ShapeUtil::MakeShape( + PrimitiveType element_type, absl::Span dimensions, + const std::vector& dynamic_dimensions) { + return MakeValidatedShape(element_type, dimensions, dynamic_dimensions) + .ValueOrDie(); +} + /* static */ StatusOr ShapeUtil::MakeValidatedShape( PrimitiveType element_type, absl::Span dimensions) { - CHECK(IsArrayPrimitiveType(element_type)); + CHECK(IsArrayPrimitiveType(element_type)) << element_type; Shape result; TF_RETURN_IF_ERROR(PopulateShape(element_type, dimensions, &result)); return result; } +/* static */ StatusOr ShapeUtil::MakeValidatedShape( + PrimitiveType element_type, absl::Span dimensions, + const std::vector& dynamic_dimensions) { + TF_ASSIGN_OR_RETURN(Shape shape, + MakeValidatedShape(element_type, dimensions)); + for (int i = 0; i < dynamic_dimensions.size(); ++i) { + shape.set_dynamic_dimension(i, dynamic_dimensions[i]); + } + return shape; +} + /* static */ Shape ShapeUtil::MakeShapeWithLayout( PrimitiveType element_type, absl::Span dimensions, absl::Span minor_to_major) { @@ -276,7 +236,7 @@ ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout( /* static */ Shape ShapeUtil::MakeTupleShape(absl::Span shapes) { Shape result; result.set_element_type(TUPLE); - result.mutable_tuple_shapes()->Reserve(shapes.size()); + result.mutable_tuple_shapes()->reserve(shapes.size()); for (const auto& shape : shapes) { AppendShapeToTuple(shape, &result); } @@ -306,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)); } @@ -321,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; @@ -345,6 +305,7 @@ ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout( case U32: case U64: case C64: + case C128: case TUPLE: case OPAQUE: case TOKEN: @@ -363,31 +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; -} - -/* static */ bool ShapeUtil::IsNil(const Shape& shape) { - return IsEmptyTuple(shape); + 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); @@ -403,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)); @@ -420,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]; } @@ -438,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; @@ -463,7 +411,7 @@ ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout( } /* static */ bool ShapeUtil::IsZeroElementArray(const Shape& shape) { - return ShapeUtil::IsArray(shape) && ElementsIn(shape) == 0; + return shape.IsArray() && ElementsIn(shape) == 0; } /* static */ bool ShapeUtil::IsScalarWithElementType( @@ -471,56 +419,8 @@ ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout( return IsScalar(shape) && shape.element_type() == element_type; } -namespace { - -// Class to memoize the computation of -// absl::AsciiStrToLower(PrimitiveType_Name(p)) -// for all PrimitiveType values "p" -class PrimitiveTypeNameGenerator { - public: - PrimitiveTypeNameGenerator() { - for (int i = 0; i < PrimitiveType_ARRAYSIZE; i++) { - if (PrimitiveType_IsValid(i)) { - lowercase_name_[i] = absl::AsciiStrToLower( - PrimitiveType_Name(static_cast(i))); - } - } - } - const string& LowercaseName(PrimitiveType t) { - return lowercase_name_[static_cast(t)]; - } - - private: - string lowercase_name_[PrimitiveType_ARRAYSIZE]; -}; - -const string& LowercasePrimitiveTypeName(PrimitiveType s) { - static PrimitiveTypeNameGenerator* gen = new PrimitiveTypeNameGenerator(); - return gen->LowercaseName(s); -} - -StatusOr StringToPrimitiveType(const string& name) { - static std::unordered_map* name_to_type = [] { - static auto* map = new std::unordered_map; - for (int i = 0; i < PrimitiveType_ARRAYSIZE; i++) { - if (PrimitiveType_IsValid(i)) { - auto value = static_cast(i); - (*map)[LowercasePrimitiveTypeName(value)] = value; - } - } - return map; - }(); - auto found = name_to_type->find(name); - if (found == name_to_type->end()) { - return InvalidArgument("Invalid element type string: \"%s\".", name); - } - return found->second; -} - -} // namespace - /* 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()) { @@ -530,12 +430,21 @@ StatusOr StringToPrimitiveType(const string& name) { text += ")"; return text; } - return StrCat(LowercasePrimitiveTypeName(shape.element_type()), "[", - absl::StrJoin(shape.dimensions(), ","), "]"); + std::vector dim_elements; + for (int i = 0; i < shape.dimensions_size(); ++i) { + if (shape.is_dynamic_dimension(i)) { + dim_elements.push_back(StrCat("<=", shape.dimensions(i))); + } else { + dim_elements.push_back(StrCat(shape.dimensions(i))); + } + } + return StrCat( + primitive_util::LowercasePrimitiveTypeName(shape.element_type()), "[", + absl::StrJoin(dim_elements, ","), "]"); } /* 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()) { @@ -545,12 +454,14 @@ StatusOr StringToPrimitiveType(const string& name) { text += ")"; return text; } - string result = StrCat(LowercasePrimitiveTypeName(shape.element_type()), "["); + 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())); } @@ -571,154 +482,25 @@ StatusOr StringToPrimitiveType(const string& name) { HumanString(program_shape.result())); } -namespace { -// Parses shapes with simple recursive descent structure -- consumes from the -// front of s and passes that view recursively as required. -StatusOr ParseShapeStringInternal(absl::string_view* s) { - *s = StripLeadingAsciiWhitespace(*s); - - if (absl::ConsumePrefix(s, "(")) { // Tuple. - std::vector shapes; - bool must_end = false; - while (true) { - if (absl::ConsumePrefix(s, ")")) { - break; - } else if (must_end) { - return InvalidArgument("Expected end of tuple; got: \"%s\"", *s); - } - shapes.emplace_back(); - TF_ASSIGN_OR_RETURN(shapes.back(), ParseShapeStringInternal(s)); - *s = StripLeadingAsciiWhitespace(*s); - must_end = !absl::ConsumePrefix(s, ","); - } - return ShapeUtil::MakeTupleShape(shapes); - } - - string element_type_string; - string dimensions_string; - string format_string; - string layout_string; - // absl::string_view is not compatible with internal RE2 StringPiece, so - // we convert in to the RE2-consumable type and then consume the corresponding - // amount from our string_view type. - static LazyRE2 shape_pattern = { - "^(\\w*\\d*)\\[([\\d,\\s]*)\\](?:\\s*(dense|sparse)?\\s*{([\\d,\\s]+)})" - "?"}; - tensorflow::RegexpStringPiece s_consumable(s->data(), s->size()); - if (RE2::Consume(&s_consumable, *shape_pattern, &element_type_string, - &dimensions_string, &format_string, &layout_string)) { - size_t consumed = s->size() - s_consumable.size(); - s->remove_prefix(consumed); - auto string_to_int64 = [&s](absl::string_view input) -> StatusOr { - int64 element; - if (!absl::SimpleAtoi(input, &element)) { - return InvalidArgument( - "Invalid s64 value in parsed shape string: \"%s\" in \"%s\"", input, - *s); - } - return element; - }; - - auto comma_list_to_int64s = - [string_to_int64](const string& input) -> StatusOr> { - std::vector results; - for (const auto& piece : absl::StrSplit(input, ',', absl::SkipEmpty())) { - TF_ASSIGN_OR_RETURN(int64 element, string_to_int64(piece)); - results.push_back(element); - } - return results; - }; - - // Extract the dimensions. - TF_ASSIGN_OR_RETURN(std::vector dimensions, - comma_list_to_int64s(dimensions_string)); - - // Extract the primitive element type. - TF_ASSIGN_OR_RETURN(const PrimitiveType primitive_type, - StringToPrimitiveType(element_type_string)); - if (primitive_type == PRIMITIVE_TYPE_INVALID || primitive_type == TUPLE) { - return InvalidArgument("Invalid element type string: \"%s\".", - element_type_string); - } - - Shape result; - if (primitive_type == OPAQUE) { - result = ShapeUtil::MakeOpaqueShape(); - } else if (primitive_type == TOKEN) { - result = ShapeUtil::MakeTokenShape(); - } else if (format_string.empty() && layout_string.empty()) { - // Create a shape without a layout set. - result = ShapeUtil::MakeShape(primitive_type, dimensions); - } else if (format_string == "sparse") { - TF_ASSIGN_OR_RETURN(int64 max_elements, string_to_int64(layout_string)); - result = ShapeUtil::MakeShapeWithSparseLayout(primitive_type, dimensions, - max_elements); - } else if (format_string.empty() || format_string == "dense") { - // Extract the layout minor-to-major and set it. - TF_ASSIGN_OR_RETURN(std::vector min2maj, - comma_list_to_int64s(layout_string)); - TF_ASSIGN_OR_RETURN(result, MakeShapeWithLayoutInternal( - primitive_type, dimensions, min2maj)); - } else { - // This should not be reached. - LOG(FATAL) << "Unhandled condition when parsing shape; format: \"" - << format_string << "\", layout: \"" << layout_string << "\""; - } - TF_RETURN_IF_ERROR(ShapeUtil::ValidateShape(result)); - return std::move(result); - } - - return InvalidArgument("Invalid shape string to parse: \"%s\"", *s); -} -} // namespace - -/* static */ StatusOr ShapeUtil::ParseShapeString(absl::string_view s) { - TF_ASSIGN_OR_RETURN(Shape shape, ParseShapeStringInternal(&s)); - if (!s.empty()) { - return InvalidArgument("Invalid shape string to parse: \"%s\"", s); - } - return shape; -} - /* static */ bool ShapeUtil::SameDimensions(const Shape& lhs, const Shape& rhs) { - CHECK(ShapeUtil::IsArray(lhs)); - CHECK(ShapeUtil::IsArray(rhs)); + CHECK(lhs.IsArray()); + CHECK(rhs.IsArray()); return absl::c_equal(lhs.dimensions(), rhs.dimensions()); } /* 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, @@ -729,7 +511,7 @@ StatusOr ParseShapeStringInternal(absl::string_view* s) { /* 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; @@ -766,6 +548,8 @@ StatusOr ParseShapeStringInternal(absl::string_view* s) { return sizeof(double); case C64: return sizeof(complex64); + case C128: + return sizeof(complex128); case TOKEN: // Tokens require no space. return 0; @@ -783,7 +567,7 @@ StatusOr ParseShapeStringInternal(absl::string_view* s) { 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); @@ -809,24 +593,14 @@ StatusOr ParseShapeStringInternal(absl::string_view* s) { /* static */ int64 ShapeUtil::ByteSizeOfElements(const Shape& shape) { TF_DCHECK_OK(ValidateShape(shape)); - CHECK(ShapeUtil::IsArray(shape)); + CHECK(shape.IsArray()); int64 allocated_element_count; if (LayoutUtil::IsSparseArray(shape)) { allocated_element_count = LayoutUtil::MaxSparseElements(shape.layout()); } else { CHECK(LayoutUtil::IsDenseArray(shape)) << shape.ShortDebugString(); - absl::Span padded_dimensions = - LayoutUtil::PaddedDimensions(shape); - if (!padded_dimensions.empty()) { - CHECK_EQ(Rank(shape), padded_dimensions.size()); - allocated_element_count = 1; - for (int64 dimension_size : padded_dimensions) { - allocated_element_count *= dimension_size; - } - } else { - allocated_element_count = ElementsIn(shape); - } + allocated_element_count = ElementsIn(shape); } return allocated_element_count * ByteSizeOfPrimitiveType(shape.element_type()); @@ -835,8 +609,8 @@ StatusOr ParseShapeStringInternal(absl::string_view* s) { /* static */ int64 ShapeUtil::ByteSizeOfSparseIndices(const Shape& shape) { TF_DCHECK_OK(ValidateShape(shape)); CHECK(LayoutUtil::IsSparseArray(shape)); - return LayoutUtil::MaxSparseElements(shape.layout()) * - ShapeUtil::Rank(shape) * sizeof(int64); + return LayoutUtil::MaxSparseElements(shape.layout()) * shape.rank() * + sizeof(int64); } /* static */ Status ShapeUtil::ValidateShapeWithOptionalLayoutInternal( @@ -867,22 +641,22 @@ StatusOr ParseShapeStringInternal(absl::string_view* s) { if (shape.dimensions_size() != 0) { return InvalidArgument( "shape has %s element type, but has dimensions field: %s", - LowercasePrimitiveTypeName(shape.element_type()), + primitive_util::LowercasePrimitiveTypeName(shape.element_type()), shape.ShortDebugString()); } if (shape.has_layout()) { return InvalidArgument( "shape has %s element type, but has layout field: %s", - LowercasePrimitiveTypeName(shape.element_type()), + primitive_util::LowercasePrimitiveTypeName(shape.element_type()), shape.ShortDebugString()); } 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( @@ -898,12 +672,17 @@ StatusOr ParseShapeStringInternal(absl::string_view* s) { /* static */ Status ShapeUtil::ValidateShapeSize(const Shape& shape) { VLOG(3) << "Validating shape size: " << ShapeUtil::HumanString(shape); - if (!IsArray(shape)) { + if (!shape.IsArray()) { return Status::OK(); } - int64 shape_size = [&shape]() { - if (LayoutUtil::IsSparseArray(shape)) { + // We can only reason about some aspects of array's shape if it has a valid + // layout, these aspects will be ignored otherwise. + bool shape_has_valid_layout = LayoutUtil::HasLayout(shape) && + LayoutUtil::ValidateLayoutInShape(shape).ok(); + + int64 shape_size = [&]() { + if (shape_has_valid_layout && LayoutUtil::IsSparseArray(shape)) { int64 max_sparse_elements = LayoutUtil::MaxSparseElements(shape.layout()); if (max_sparse_elements < 0) { return max_sparse_elements; @@ -914,7 +693,7 @@ StatusOr ParseShapeStringInternal(absl::string_view* s) { return sparse_elements_size; } int64 sparse_indices_size = - MultiplyWithoutOverflow(max_sparse_elements, ShapeUtil::Rank(shape)); + MultiplyWithoutOverflow(max_sparse_elements, shape.rank()); if (sparse_indices_size < 0) { return sparse_indices_size; } @@ -939,11 +718,8 @@ StatusOr ParseShapeStringInternal(absl::string_view* s) { return dense_shape_size; } - bool is_padded = - LayoutUtil::IsDenseArray(shape) && LayoutUtil::IsPadded(shape); absl::Span shape_max_dimensions = - is_padded ? LayoutUtil::PaddedDimensions(shape) - : AsInt64Slice(shape.dimensions()); + AsInt64Slice(shape.dimensions()); for (int64 dim : shape_max_dimensions) { dense_shape_size = MultiplyWithoutOverflow(dense_shape_size, dim); if (dense_shape_size < 0) { @@ -989,7 +765,7 @@ StatusOr ParseShapeStringInternal(absl::string_view* s) { 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); @@ -1001,7 +777,7 @@ StatusOr ParseShapeStringInternal(absl::string_view* s) { 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); } @@ -1012,7 +788,7 @@ StatusOr ParseShapeStringInternal(absl::string_view* s) { 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", @@ -1027,7 +803,7 @@ StatusOr ParseShapeStringInternal(absl::string_view* s) { 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; @@ -1035,11 +811,11 @@ StatusOr ParseShapeStringInternal(absl::string_view* s) { /* 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; @@ -1061,10 +837,15 @@ bool ShapeUtil::IsLeafIndex(const Shape& shape, const ShapeIndex& index) { } /* static */ bool ShapeUtil::HasDegenerateDimensions(const Shape& shape) { - CHECK(ShapeUtil::IsArray(shape)); + CHECK(shape.IsArray()); return absl::c_linear_search(shape.dimensions(), 1); } +/* static */ Shape ShapeUtil::DropDegenerateDimensions(const Shape& shape) { + return FilterDimensions( + [&](int64 dim) -> bool { return shape.dimensions()[dim] != 1; }, shape); +} + namespace { // Helper for ForEachSubshape which visits the subshapes of the given shape in @@ -1073,7 +854,7 @@ Status ForEachSubshapeHelper(const Shape& shape, const ShapeUtil::StatusVisitorFunction& func, ShapeIndex* index) { TF_RETURN_IF_ERROR(func(shape, *index)); - if (ShapeUtil::IsTuple(shape)) { + if (shape.IsTuple()) { for (int64 i = 0; i < ShapeUtil::TupleElementCount(shape); ++i) { index->push_back(i); TF_RETURN_IF_ERROR(ForEachSubshapeHelper( @@ -1090,7 +871,7 @@ Status ForEachMutableSubshapeHelper( Shape* shape, const ShapeUtil::MutatingStatusVisitorFunction& func, ShapeIndex* index) { TF_RETURN_IF_ERROR(func(shape, *index)); - if (ShapeUtil::IsTuple(*shape)) { + if (shape->IsTuple()) { for (int64 i = 0; i < ShapeUtil::TupleElementCount(*shape); ++i) { index->push_back(i); TF_RETURN_IF_ERROR(ForEachMutableSubshapeHelper( @@ -1148,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. @@ -1166,7 +951,7 @@ Status ForEachMutableSubshapeHelper( // Let the argument `permutation` be P. This is a permutation over `shape`'s // dimensions, so our return value will be a shape with dims P.I = P. Our // goal is to construct a layout permutation L* that we can apply to P such - // that that the physical dimension ordering of the returned shape is the same + // that the physical dimension ordering of the returned shape is the same // as that of the original shape, namely L'. // // Our returned shape has dims P and layout L*, so its in-memory layout is @@ -1185,13 +970,6 @@ Status ForEachMutableSubshapeHelper( permutation, AsInt64Slice(shape.layout().minor_to_major()))) { new_layout->add_minor_to_major(index); } - if (shape.layout().padded_dimensions_size() > 0) { - new_layout->clear_padded_dimensions(); - for (auto dim : - Permute(permutation, shape.layout().padded_dimensions())) { - new_layout->add_padded_dimensions(dim); - } - } // The permutation accepted by TransposeIsBitcast is the inverse of the // permutation here. CHECK(TransposeIsBitcast(shape, new_shape, InversePermutation(permutation))) @@ -1205,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()); @@ -1253,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; } @@ -1265,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()), @@ -1294,11 +1072,6 @@ ShapeUtil::DimensionsUnmodifiedByReshape(const Shape& input_shape, return false; } - // Padding is not handled. - if (LayoutUtil::IsPadded(input_shape) && LayoutUtil::IsPadded(output_shape)) { - return false; - } - // Check the reshape permutes the positions of each dimension in the // minor-to-major order. positions[i]=k means dimension `i` is k-th minor. // input_positions = apply(dimension_mapping, output_positions) @@ -1321,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)); @@ -1330,11 +1103,6 @@ ShapeUtil::DimensionsUnmodifiedByReshape(const Shape& input_shape, return false; } - // Padding is not handled. - if (LayoutUtil::IsPadded(input_shape) || LayoutUtil::IsPadded(output_shape)) { - return false; - } - CHECK_EQ(ElementsIn(input_shape), ElementsIn(output_shape)); if (ElementsIn(input_shape) == 0) { return true; @@ -1455,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, @@ -1486,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 @@ -1627,29 +1395,14 @@ ShapeUtil::DimensionsUnmodifiedByReshape(const Shape& input_shape, /* static */ Shape ShapeUtil::DeleteDimension(int64 dim_to_delete, Shape shape) { - CHECK(IsArray(shape)); - shape.mutable_dimensions()->erase(shape.dimensions().begin() + dim_to_delete); - if (LayoutUtil::HasLayout(shape)) { - Layout* layout = shape.mutable_layout(); - layout->set_format(DENSE); - for (size_t i = 0; i < layout->minor_to_major().size();) { - if (layout->minor_to_major(i) == dim_to_delete) { - layout->mutable_minor_to_major()->erase( - layout->minor_to_major().begin() + i); - continue; - } - if (layout->minor_to_major(i) > dim_to_delete) { - (*layout->mutable_minor_to_major())[i] -= 1; - } - ++i; - } - } + 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)) { @@ -1662,11 +1415,6 @@ ShapeUtil::DimensionsUnmodifiedByReshape(const Shape& input_shape, return shape; } -std::ostream& operator<<(std::ostream& out, const Shape& shape) { - out << ShapeUtil::HumanStringWithLayout(shape); - return out; -} - /*static*/ size_t ShapeUtil::Hash(const Shape& shape) { using tensorflow::hash; using tensorflow::Hash64Combine; @@ -1674,8 +1422,11 @@ std::ostream& operator<<(std::ostream& out, const Shape& shape) { size_t hash_value = hash()(shape.element_type()); if (shape.tuple_shapes().empty()) { - for (int64 dim : shape.dimensions()) { - hash_value = Hash64Combine(hash_value, hash()(dim)); + for (int i = 0; i < shape.dimensions_size(); ++i) { + hash_value = + Hash64Combine(hash_value, hash()(shape.dimensions(i))); + hash_value = Hash64Combine(hash_value, + hash()(shape.is_dynamic_dimension(i))); } hash_value = Hash64Combine(hash_value, LayoutUtil::Hash(shape.layout())); diff --git a/tensorflow/compiler/xla/shape_util.h b/tensorflow/compiler/xla/shape_util.h index 191ab04759f2d0ae87d988cba0d303f1ab696432..fb6da7460e2475732d6f02758e5519fbdb7c0f8d 100644 --- a/tensorflow/compiler/xla/shape_util.h +++ b/tensorflow/compiler/xla/shape_util.h @@ -28,6 +28,7 @@ limitations under the License. #include "absl/types/span.h" #include "tensorflow/compiler/xla/layout_util.h" #include "tensorflow/compiler/xla/primitive_util.h" +#include "tensorflow/compiler/xla/shape.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/types.h" @@ -37,6 +38,7 @@ limitations under the License. #include "tensorflow/core/platform/cpu_info.h" #include "tensorflow/core/platform/env.h" #include "tensorflow/core/platform/macros.h" +#include "tensorflow/core/platform/mutex.h" #include "tensorflow/core/platform/types.h" namespace xla { @@ -100,6 +102,11 @@ class ShapeIndex { string ToString() const; + template + friend H AbslHashValue(H h, const ShapeIndex& index) { + return H::combine(std::move(h), index.indices_); + } + private: container_type indices_; }; @@ -147,6 +154,9 @@ class ShapeIndexView { string ToString() const; + // Returns true if this shape index starts with 'prefix'. + bool StartsWith(ShapeIndexView prefix) const; + private: absl::Span indices_; }; @@ -175,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. @@ -197,7 +207,7 @@ class ShapeUtil { // Returns the number of bytes used to store the primitive_type. // - // Precondition: ShapeUtil::IsArray(shape) + // Precondition: shape.IsArray() static int64 ByteSizeOfPrimitiveType(PrimitiveType primitive_type); // Returns the number of bytes required to store the tuple member pointers for @@ -231,10 +241,6 @@ class ShapeUtil { // (param_name: f32[42x12], ...) -> f32[24x42] static string HumanString(const ProgramShape& program_shape); - // Parses a ShapeUtil::HumanString-format shape string back into a shape - // object. - static StatusOr ParseShapeString(absl::string_view s); - // Returns whether the LHS and RHS shapes have the same dimensions; note: does // not check element type. // Precondition: IsArray(lhs) && IsArray(rhs) @@ -256,7 +262,7 @@ class ShapeUtil { } // Returns the higher-precision element type if a and b are both floating - // point types; otherwise, checks that that they have the same element type + // point types; otherwise, checks that they have the same element type // and returns it. static PrimitiveType HigherPrecisionElementType(const Shape& a, const Shape& b) { @@ -284,16 +290,12 @@ class ShapeUtil { // being F32. Tuple elements are compared recursively for compatibility. static bool CompatibleIgnoringFpPrecision(const Shape& lhs, const Shape& rhs); - // Returns whether the lhs and rhs shapes are identical protobufs. + // Returns whether the lhs and rhs shapes are identical. static bool Equal(const Shape& lhs, const Shape& rhs); // 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) - 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., @@ -307,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. @@ -365,11 +367,24 @@ class ShapeUtil { static Shape MakeShape(PrimitiveType element_type, absl::Span dimensions); + // Constructs a new shape with the given element type and sequence of + // potentially dynamic dimensions. The argument 'dynamic_dimensions' indicates + // with a true value that the respective dimension is dynamic. If the + // dimension is dynamic then the respective value in 'dimension' is an upper + // bound on the dimension size. 'dimensions' and 'dynamic_dimensions' must be + // the same size. + static Shape MakeShape(PrimitiveType element_type, + absl::Span dimensions, + const std::vector& dynamic_dimensions); + // Constructs a new shape with the given element type and sequence of // dimensions. Method checks if the element type is valid and the shape's // size fits in std::numeric_limits::max(). static StatusOr MakeValidatedShape(PrimitiveType element_type, absl::Span dimensions); + static StatusOr MakeValidatedShape( + PrimitiveType element_type, absl::Span dimensions, + const std::vector& dynamic_dimensions); // Creates a Shape with element type corresponding to T and the given // dimensions @@ -437,26 +452,8 @@ class ShapeUtil { // that floating point numbers are signed. static bool ElementIsSigned(const Shape& shape); - // Returns whether the shape is a tuple. - 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). - 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. - static bool IsToken(const Shape& shape) { - return shape.element_type() == TOKEN; - } - - // Returns whether the shape is an array. Note that scalars are considered - // arrays. - static bool IsArray(const Shape& shape); + // Returns whether the given primitive type corresponds to an array shape. + static bool IsArrayPrimitiveType(PrimitiveType primitive_type); // Returns whether the shape is a tuple with at least one element which is // also a tuple. @@ -465,9 +462,6 @@ class ShapeUtil { // Returns true if shape is an empty tuple. static bool IsEmptyTuple(const Shape& shape); - // Returns true if shape is the nil shape (an empty tuple). - static bool IsNil(const Shape& shape); - // Returns the number of elements in the given tuple shape. // Precondition: IsTuple(shape) static int64 TupleElementCount(const Shape& shape); @@ -487,12 +481,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); @@ -541,6 +529,9 @@ class ShapeUtil { // (dimensions with bound 1). static bool HasDegenerateDimensions(const Shape& shape); + // Drops any degenerate dimensions (i.e. dimensions of size 1) + static Shape DropDegenerateDimensions(const Shape& shape); + // Permutes the dimensions by the given permutation, so // return_value.dimensions[permutation[i]] = argument.dimensions[i]. // @@ -684,11 +675,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 @@ -737,7 +726,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(); @@ -751,10 +740,18 @@ class ShapeUtil { pool.emplace(tensorflow::Env::Default(), "foreach", kNumThreads); } + tensorflow::mutex mu; + Status status; // Guarded by mu + while (n < rank) { if (pool != absl::nullopt) { - pool->Schedule( - [indexes, &visitor_function] { visitor_function(indexes); }); + pool->Schedule([indexes, &visitor_function, &mu, &status] { + StatusOr result = visitor_function(indexes); + if (!result.ok()) { + tensorflow::mutex_lock lock(mu); + status = status.ok() ? result.status() : status; + } + }); } else { TF_ASSIGN_OR_RETURN(bool should_continue, visitor_function(indexes)); if (!should_continue) { @@ -772,14 +769,14 @@ class ShapeUtil { } } - return Status::OK(); + // Waits for the scheduled work to complete. + pool.reset(); + return status; } TF_DISALLOW_COPY_AND_ASSIGN(ShapeUtil); }; -std::ostream& operator<<(std::ostream& out, const Shape& shape); - } // namespace xla #endif // TENSORFLOW_COMPILER_XLA_SHAPE_UTIL_H_ diff --git a/tensorflow/compiler/xla/shape_util_test.cc b/tensorflow/compiler/xla/shape_util_test.cc index c622ecdca1fd66604d1a6ceaf705f2e70edaee55..126ae58293d12182e9b6e30f779f681829729526 100644 --- a/tensorflow/compiler/xla/shape_util_test.cc +++ b/tensorflow/compiler/xla/shape_util_test.cc @@ -82,102 +82,6 @@ TEST(ShapeUtilTest, Rank4DimensionIndexing) { ASSERT_EQ(3, shape.dimensions(0)); } -TEST(ShapeUtilTest, ParseShapeStringR2F32) { - string shape_string = "f32[123,456]"; - TF_ASSERT_OK_AND_ASSIGN(Shape actual, - ShapeUtil::ParseShapeString(shape_string)); - Shape expected = ShapeUtil::MakeShape(F32, {123, 456}); - ASSERT_TRUE(ShapeUtil::Equal(expected, actual)) - << "expected: " << ShapeUtil::HumanString(expected) - << "actual: " << ShapeUtil::HumanString(actual); -} - -TEST(ShapeUtilTest, ParseShapeStringTupleOfArrays) { - string shape_string = "(f32[1572864],s8[5120,1024])"; - TF_ASSERT_OK_AND_ASSIGN(Shape actual, - ShapeUtil::ParseShapeString(shape_string)); - Shape expected = - ShapeUtil::MakeTupleShape({ShapeUtil::MakeShape(F32, {1572864}), - ShapeUtil::MakeShape(S8, {5120, 1024})}); - ASSERT_TRUE(ShapeUtil::Equal(expected, actual)) - << "expected: " << ShapeUtil::HumanString(expected) - << "actual: " << ShapeUtil::HumanString(actual); -} - -TEST(ShapeUtilTest, ParseShapeStringNestedTuple) { - string shape_string = "(f32[1],(f32[2], token[]), opaque[], f32[3])"; - TF_ASSERT_OK_AND_ASSIGN(Shape actual, - ShapeUtil::ParseShapeString(shape_string)); - Shape expected = ShapeUtil::MakeTupleShape({ - ShapeUtil::MakeShape(F32, {1}), - ShapeUtil::MakeTupleShape( - {ShapeUtil::MakeShape(F32, {2}), ShapeUtil::MakeTokenShape()}), - ShapeUtil::MakeOpaqueShape(), - ShapeUtil::MakeShape(F32, {3}), - }); - ASSERT_TRUE(ShapeUtil::Equal(expected, actual)) - << "expected: " << ShapeUtil::HumanString(expected) - << "actual: " << ShapeUtil::HumanString(actual); -} - -TEST(ShapeUtilTest, ParseShapeStringWithLayout) { - string shape_string = "f32[123,456]{0,1}"; - TF_ASSERT_OK_AND_ASSIGN(Shape actual, - ShapeUtil::ParseShapeString(shape_string)); - Shape expected = ShapeUtil::MakeShapeWithLayout(F32, {123, 456}, {0, 1}); - ASSERT_TRUE(ShapeUtil::Equal(expected, actual)) - << "expected: " << ShapeUtil::HumanString(expected) - << "actual: " << ShapeUtil::HumanString(actual); -} - -TEST(ShapeUtilTest, ParseShapeStringWithExplicitDenseLayout) { - string shape_string = "f32[123,456]dense{0,1}"; - TF_ASSERT_OK_AND_ASSIGN(Shape actual, - ShapeUtil::ParseShapeString(shape_string)); - Shape expected = ShapeUtil::MakeShapeWithLayout(F32, {123, 456}, {0, 1}); - ASSERT_TRUE(ShapeUtil::Equal(expected, actual)) - << "expected: " << ShapeUtil::HumanString(expected) - << "actual: " << ShapeUtil::HumanString(actual); -} - -TEST(ShapeUtilTest, ParseShapeStringWithSparseLayout) { - string shape_string = "f32[123,456]sparse{10}"; - TF_ASSERT_OK_AND_ASSIGN(Shape actual, - ShapeUtil::ParseShapeString(shape_string)); - Shape expected = ShapeUtil::MakeShapeWithSparseLayout(F32, {123, 456}, 10); - ASSERT_TRUE(ShapeUtil::Equal(expected, actual)) - << "expected: " << ShapeUtil::HumanString(expected) - << "actual: " << ShapeUtil::HumanString(actual); -} - -TEST(ShapeUtilTest, ParseOpaqueType) { - TF_ASSERT_OK_AND_ASSIGN(Shape actual, - ShapeUtil::ParseShapeString("opaque[]")); - Shape expected = ShapeUtil::MakeOpaqueShape(); - ASSERT_TRUE(ShapeUtil::Equal(expected, actual)) - << "expected: " << ShapeUtil::HumanString(expected) - << "actual: " << ShapeUtil::HumanString(actual); -} - -TEST(ShapeUtilTest, ParseTokenType) { - TF_ASSERT_OK_AND_ASSIGN(Shape actual, ShapeUtil::ParseShapeString("token[]")); - Shape expected = ShapeUtil::MakeTokenShape(); - ASSERT_TRUE(ShapeUtil::Equal(expected, actual)) - << "expected: " << ShapeUtil::HumanString(expected) - << "actual: " << ShapeUtil::HumanString(actual); -} - -TEST(ShapeUtilTest, ParseInvalidShapeString) { - string shape_strings[] = { - "f32[123,456]foobar{0,1}", "f32[123,456]sparse{0,1}", "f32[123,456]{foo}", - "f32[123,456]dense{foo}", "f32[123,456]sparse{foo}", - }; - for (const string& shape_string : shape_strings) { - StatusOr result = ShapeUtil::ParseShapeString(shape_string); - ASSERT_FALSE(result.ok()) << "shape: " << shape_string; - } -} - TEST(ShapeUtilTest, CompatibleIdenticalShapes) { Shape shape1 = ShapeUtil::MakeShape(F32, {3, 2}); Shape shape2 = ShapeUtil::MakeShape(F32, {3, 2}); @@ -272,6 +176,28 @@ TEST(ShapeUtilTest, UnequalIgnoringFpPrecision) { ShapeUtil::MakeShapeWithLayout(PRED, {4, 3}, {0, 1}))); } +TEST(ShapeUtilTest, EqualDynamicShapes) { + EXPECT_TRUE( + ShapeUtil::Equal(ShapeUtil::MakeShape(F32, {4, 3}, {true, false}), + ShapeUtil::MakeShape(F32, {4, 3}, {true, false}))); + EXPECT_FALSE( + ShapeUtil::Equal(ShapeUtil::MakeShape(F32, {4, 3}, {true, false}), + ShapeUtil::MakeShape(F32, {4, 3}, {false, false}))); +} + +TEST(ShapeUtilTest, CompatibleDynamicShapes) { + Shape shape_a = ShapeUtil::MakeShape(F32, {4, 3}, {true, false}); + *shape_a.mutable_layout() = Layout({1, 0}); + Shape shape_b = ShapeUtil::MakeShape(F32, {4, 3}, {true, false}); + *shape_b.mutable_layout() = Layout({0, 1}); + Shape shape_c = ShapeUtil::MakeShape(F32, {4, 3}, {false, true}); + *shape_c.mutable_layout() = Layout({0, 1}); + + EXPECT_TRUE(ShapeUtil::Compatible(shape_a, shape_a)); + EXPECT_TRUE(ShapeUtil::Compatible(shape_a, shape_b)); + EXPECT_FALSE(ShapeUtil::Compatible(shape_a, shape_c)); +} + TEST(ShapeUtilTest, CompatibleTuples) { Shape tuple1 = ShapeUtil::MakeTupleShape( {ShapeUtil::MakeShape(F32, {3, 2}), ShapeUtil::MakeShape(PRED, {4, 5})}); @@ -345,26 +271,6 @@ TEST(ShapeUtilTest, OpaqueVsArray) { EXPECT_FALSE(ShapeUtil::CompatibleIgnoringElementType(shape2, shape1)); } -TEST(ShapeUtilTest, CompareShapesWithPaddedDimensionsMismatch) { - Shape shape1 = ShapeUtil::MakeShape(F32, {20, 30}); - shape1.mutable_layout()->add_padded_dimensions(10); - - Shape shape2 = ShapeUtil::MakeShape(F32, {20, 30}); - shape2.mutable_layout()->add_padded_dimensions(11); - - EXPECT_FALSE(ShapeUtil::Equal(shape1, shape2)); -} - -TEST(ShapeUtilTest, CompareShapesWithPaddingValueMismatch) { - Shape shape1 = ShapeUtil::MakeShape(F32, {20, 30}); - shape1.mutable_layout()->set_padding_value(ZERO_PAD); - - Shape shape2 = ShapeUtil::MakeShape(F32, {20, 30}); - shape2.mutable_layout()->set_padding_value(LOWEST_PAD); - - EXPECT_FALSE(ShapeUtil::Equal(shape1, shape2)); -} - TEST(ShapeUtilTest, ScalarDefaultLayoutEqualsScalarEmptyMin2Maj) { Shape scalar_default_layout = ShapeUtil::MakeShape(F32, {}); ASSERT_TRUE(scalar_default_layout.has_layout()) @@ -395,23 +301,13 @@ TEST(ShapeUtilTest, ByteSizeOfWithoutPadding) { EXPECT_EQ(0, ShapeUtil::ByteSizeOf(ShapeUtil::MakeTokenShape())); } -TEST(ShapeUtilTest, ByteSizeOfWithPadding) { - EXPECT_EQ(4, ShapeUtil::ByteSizeOfPrimitiveType(F32)); - Shape shape = ShapeUtil::MakeShape(F32, {10, 20}); - EXPECT_EQ(800, ShapeUtil::ByteSizeOf(shape)); - - shape.mutable_layout()->add_padded_dimensions(15); - shape.mutable_layout()->add_padded_dimensions(21); - EXPECT_EQ(15 * 21 * 4, ShapeUtil::ByteSizeOf(shape)); -} - TEST(ShapeUtilTest, NilShape) { - EXPECT_TRUE(ShapeUtil::IsNil(ShapeUtil::MakeNil())); - EXPECT_FALSE(ShapeUtil::IsNil(ShapeUtil::MakeShape(F32, {1, 2, 3}))); - EXPECT_FALSE(ShapeUtil::IsNil(ShapeUtil::MakeShape(F32, {0, 1}))); - EXPECT_FALSE(ShapeUtil::IsNil( + EXPECT_TRUE(ShapeUtil::IsEmptyTuple(ShapeUtil::MakeNil())); + EXPECT_FALSE(ShapeUtil::IsEmptyTuple(ShapeUtil::MakeShape(F32, {1, 2, 3}))); + EXPECT_FALSE(ShapeUtil::IsEmptyTuple(ShapeUtil::MakeShape(F32, {0, 1}))); + EXPECT_FALSE(ShapeUtil::IsEmptyTuple( ShapeUtil::MakeTupleShape({ShapeUtil::MakeShape(S32, {})}))); - EXPECT_FALSE(ShapeUtil::IsNil( + EXPECT_FALSE(ShapeUtil::IsEmptyTuple( ShapeUtil::MakeTupleShape({ShapeUtil::MakeShape(F32, {0})}))); } @@ -576,68 +472,6 @@ TEST(ShapeUtilTest, IsLeafIndex) { EXPECT_TRUE(ShapeUtil::IsLeafIndex(nested_tuple_shape, {1, 1})); } -TEST(ShapeUtilTest, HumanString) { - Shape opaque = ShapeUtil::MakeOpaqueShape(); - Shape token = ShapeUtil::MakeTokenShape(); - Shape scalar = ShapeUtil::MakeShape(F32, {}); - Shape matrix = ShapeUtil::MakeShape(U32, {1, 2}); - Shape matrix2 = ShapeUtil::MakeShapeWithLayout(S32, {3, 4}, {0, 1}); - Shape tuple = ShapeUtil::MakeTupleShape({opaque, scalar, matrix, matrix2}); - Shape nested_tuple = ShapeUtil::MakeTupleShape({tuple, matrix, token}); - - EXPECT_EQ("opaque[]", ShapeUtil::HumanString(opaque)); - EXPECT_EQ("token[]", ShapeUtil::HumanString(token)); - EXPECT_EQ("f32[]", ShapeUtil::HumanString(scalar)); - EXPECT_EQ("u32[1,2]", ShapeUtil::HumanString(matrix)); - EXPECT_EQ("s32[3,4]", ShapeUtil::HumanString(matrix2)); - EXPECT_EQ("(opaque[], f32[], u32[1,2], s32[3,4])", - ShapeUtil::HumanString(tuple)); - EXPECT_EQ("((opaque[], f32[], u32[1,2], s32[3,4]), u32[1,2], token[])", - ShapeUtil::HumanString(nested_tuple)); - - EXPECT_EQ("opaque[]", ShapeUtil::HumanStringWithLayout(opaque)); - EXPECT_EQ("f32[]", ShapeUtil::HumanStringWithLayout(scalar)); - EXPECT_EQ("u32[1,2]{1,0}", ShapeUtil::HumanStringWithLayout(matrix)); - EXPECT_EQ("s32[3,4]{0,1}", ShapeUtil::HumanStringWithLayout(matrix2)); - EXPECT_EQ("(opaque[], f32[], u32[1,2]{1,0}, s32[3,4]{0,1})", - ShapeUtil::HumanStringWithLayout(tuple)); - EXPECT_EQ( - "((opaque[], f32[], u32[1,2]{1,0}, s32[3,4]{0,1}), u32[1,2]{1,0}, " - "token[])", - ShapeUtil::HumanStringWithLayout(nested_tuple)); - - ProgramShape prog = ShapeUtil::MakeProgramShape( - {opaque, scalar, matrix, matrix2, tuple, nested_tuple}, nested_tuple); - EXPECT_EQ( - "((unknown): opaque[], " - "(unknown): f32[], " - "(unknown): u32[1,2], " - "(unknown): s32[3,4], " - "(unknown): (opaque[], f32[], u32[1,2], s32[3,4]), " - "(unknown): ((opaque[], f32[], u32[1,2], s32[3,4]), u32[1,2], token[])) " - "-> " - "((opaque[], f32[], u32[1,2], s32[3,4]), u32[1,2], token[])", - ShapeUtil::HumanString(prog)); - - prog.add_parameter_names("arg0"); - prog.add_parameter_names("scalar"); - prog.add_parameter_names("matrix"); - prog.add_parameter_names("matrix2"); - prog.add_parameter_names("tuple"); - prog.add_parameter_names("nested_tuple"); - EXPECT_EQ( - "(arg0: opaque[], " - "scalar: f32[], " - "matrix: u32[1,2], " - "matrix2: s32[3,4], " - "tuple: (opaque[], f32[], u32[1,2], s32[3,4]), " - "nested_tuple: ((opaque[], f32[], u32[1,2], s32[3,4]), u32[1,2], " - "token[])) " - "-> " - "((opaque[], f32[], u32[1,2], s32[3,4]), u32[1,2], token[])", - ShapeUtil::HumanString(prog)); -} - TEST(ShapeUtilTest, ForEachSubshapeArray) { const Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); int calls = 0; @@ -704,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; @@ -880,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.cc b/tensorflow/compiler/xla/sparse_index_array.cc index a40bb7875e7ea53a8959a9a67ec09ec260ba9c37..82091bdee65c709bb6020f40acc15f13d8599c1d 100644 --- a/tensorflow/compiler/xla/sparse_index_array.cc +++ b/tensorflow/compiler/xla/sparse_index_array.cc @@ -79,7 +79,7 @@ void SparseIndexArray::Resize(int64 num_indices) { } bool SparseIndexArray::Validate(const Shape& shape) const { - if (rank_ == 0 || rank_ != ShapeUtil::Rank(shape)) { + if (rank_ == 0 || rank_ != shape.rank()) { return false; } int64 num_indices = index_count(); 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/tests/BUILD b/tensorflow/compiler/xla/tests/BUILD index 3ad6960b4ea7f1cab3dbeb77c4d2ecf33ade534f..5dd9f3b408efd3c12fc8d40a38b901c03615028a 100644 --- a/tensorflow/compiler/xla/tests/BUILD +++ b/tensorflow/compiler/xla/tests/BUILD @@ -1,6 +1,13 @@ # Description: # Base testing infrastructure for XLA. +load("//tensorflow/compiler/xla/tests:build_defs.bzl", "generate_backend_suites", "generate_backend_test_macros", "xla_test", "xla_test_library") +load( + "//tensorflow/core:platform/default/build_config_root.bzl", + "tf_cuda_tests_tags", +) +load("//tensorflow:tensorflow.bzl", "tf_cc_binary", "tf_cc_test") + licenses(["notice"]) # Apache 2.0 package( @@ -23,17 +30,6 @@ filegroup( ]), ) -load("//tensorflow/compiler/xla/tests:build_defs.bzl", "xla_test") -load("//tensorflow/compiler/xla/tests:build_defs.bzl", "xla_test_library") -load("//tensorflow/compiler/xla/tests:build_defs.bzl", "generate_backend_suites") -load("//tensorflow/compiler/xla/tests:build_defs.bzl", "generate_backend_test_macros") -load("//tensorflow:tensorflow.bzl", "tf_cc_binary") -load("//tensorflow:tensorflow.bzl", "tf_cc_test") -load( - "//tensorflow/core:platform/default/build_config_root.bzl", - "tf_cuda_tests_tags", -) - # Generate test_suites for all backends, named "${backend}_tests". generate_backend_suites() @@ -44,7 +40,7 @@ cc_library( testonly = True, srcs = ["xla_internal_test_main.cc"], deps = [ - "//tensorflow/compiler/xla/legacy_flags:debug_options_flags", + "//tensorflow/compiler/xla:debug_options_flags", "//tensorflow/core:lib", "//tensorflow/core:test", "@com_google_absl//absl/strings", @@ -75,10 +71,12 @@ 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", "//tensorflow/core:lib", + "@com_google_absl//absl/base", "@com_google_absl//absl/memory", "@com_google_absl//absl/types:span", ], @@ -117,12 +115,12 @@ cc_library( deps = [ ":literal_test_util", ":test_utils", + "//tensorflow/compiler/xla:debug_options_flags", "//tensorflow/compiler/xla:shape_layout", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/legacy_flags:debug_options_flags", "//tensorflow/compiler/xla/service:backend", "//tensorflow/compiler/xla/service:computation_layout", "//tensorflow/compiler/xla/service:hlo", @@ -135,50 +133,13 @@ cc_library( "//tensorflow/core:stream_executor_no_cuda", "//tensorflow/core:test", "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/base:core_headers", "@com_google_absl//absl/memory", "@com_google_absl//absl/types:optional", "@com_google_absl//absl/types:span", ], ) -cc_library( - name = "hlo_verified_test_base", - testonly = True, - srcs = ["hlo_verified_test_base.cc"], - hdrs = ["hlo_verified_test_base.h"], - deps = [ - ":hlo_test_base", - "//tensorflow/compiler/xla:shape_util", - "//tensorflow/compiler/xla:status_macros", - "//tensorflow/compiler/xla/service:hlo", - "//tensorflow/compiler/xla/service:hlo_parser", - "//tensorflow/compiler/xla/service:hlo_verifier", - "//tensorflow/core:lib", - "@com_google_absl//absl/base:core_headers", - "@com_google_absl//absl/memory", - ], -) - -tf_cc_test( - name = "hlo_verified_test_base_test", - srcs = ["hlo_verified_test_base_test.cc"], - deps = [ - ":hlo_test_base", - ":hlo_verified_test_base", - ":test_macros_cpu", - ":test_utils", - "//tensorflow/compiler/xla:shape_util", - "//tensorflow/compiler/xla/client:xla_builder", - "//tensorflow/compiler/xla/client:xla_computation", - "//tensorflow/compiler/xla/service:hlo", - "//tensorflow/compiler/xla/service:hlo_parser", - "//tensorflow/compiler/xla/service:hlo_verifier", - "//tensorflow/compiler/xla/tests:xla_internal_test_main", - "//tensorflow/core:lib", - "//tensorflow/core:test", - ], -) - tf_cc_binary( name = "local_client_aot_test_helper", srcs = ["local_client_aot_test_helper.cc"], @@ -316,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", @@ -335,12 +293,75 @@ xla_test( ], ) +xla_test( + name = "conv_depthwise_test", + timeout = "long", + srcs = ["conv_depthwise_test.cc"], + shard_count = 50, + 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 = "conv_depthwise_backprop_filter_test", + timeout = "long", + srcs = ["conv_depthwise_backprop_filter_test.cc"], + shard_count = 1, + 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", + srcs = ["grouped_convolution_test.cc"], + blacklisted_backends = [ + # disabled because of a break b/119590850. + "gpu", + # disabled because it times out. + "cpu", + ], + shard_count = 50, + 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 = "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", @@ -361,9 +382,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", @@ -381,9 +399,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", @@ -407,6 +422,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", @@ -430,9 +449,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", @@ -447,7 +463,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", @@ -500,9 +515,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", @@ -518,9 +530,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", @@ -538,7 +547,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", @@ -556,7 +564,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", @@ -617,9 +624,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", @@ -642,7 +646,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", @@ -660,6 +663,7 @@ xla_test( "//tensorflow/compiler/xla/tests:literal_test_util", "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:lib", + "@com_google_absl//absl/base", "@com_google_absl//absl/types:span", ], ) @@ -684,13 +688,13 @@ xla_test( "//tensorflow/compiler/xla/client:xla_builder", "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:lib", + "@com_google_absl//absl/base", ], ) 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", @@ -707,6 +711,7 @@ xla_test( "//tensorflow/compiler/xla/tests:literal_test_util", "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:lib", + "@com_google_absl//absl/base", "@com_google_absl//absl/strings", ], ) @@ -716,7 +721,6 @@ xla_test( srcs = ["dot_operation_test.cc"], shard_count = 20, tags = [ - "enable_for_xla_interpreter", "optonly", ], deps = [ @@ -726,7 +730,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", @@ -783,7 +789,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", @@ -797,9 +805,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", @@ -819,9 +824,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", @@ -832,7 +834,9 @@ 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", "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:lib", @@ -865,7 +869,8 @@ xla_test( name = "convolution_test", timeout = "long", srcs = ["convolution_test.cc"], - shard_count = 25, + shard_count = 40, + tags = ["optonly"], deps = CONVOLUTION_TEST_DEPS + [ "@com_google_absl//absl/memory", "@com_google_absl//absl/strings", @@ -940,6 +945,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", @@ -1031,9 +1041,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", @@ -1054,9 +1061,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", @@ -1074,9 +1078,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", @@ -1102,9 +1103,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", @@ -1128,9 +1126,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", @@ -1151,7 +1146,6 @@ xla_test( srcs = ["reduce_test.cc"], shard_count = 40, tags = [ - "enable_for_xla_interpreter", "optonly", ], deps = [ @@ -1174,6 +1168,7 @@ xla_test( "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:lib", "//tensorflow/core:test", + "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/strings", "@com_google_absl//absl/strings:str_format", "@com_google_absl//absl/types:span", @@ -1217,7 +1212,6 @@ xla_test( srcs = [], shard_count = 20, tags = [ - "enable_for_xla_interpreter", "optonly", ], xla_test_library_deps = [":reduce_window_test_library"], @@ -1229,7 +1223,6 @@ xla_test( timeout = "long", srcs = ["select_and_scatter_test.cc"], tags = [ - "enable_for_xla_interpreter", "optonly", ], deps = [ @@ -1255,9 +1248,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", @@ -1278,9 +1268,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", @@ -1294,10 +1281,8 @@ 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", "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/compiler/xla/tests:xla_internal_test_main", @@ -1310,9 +1295,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", @@ -1332,6 +1314,7 @@ xla_test( xla_test( name = "custom_call_test", srcs = ["custom_call_test.cc"], + backends = ["cpu"], deps = [ "//tensorflow/compiler/xla:literal", "//tensorflow/compiler/xla:literal_util", @@ -1354,9 +1337,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", @@ -1374,9 +1354,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", @@ -1396,9 +1373,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", @@ -1420,9 +1394,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", @@ -1437,9 +1408,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", @@ -1454,9 +1422,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", @@ -1503,9 +1468,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", @@ -1531,9 +1493,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", @@ -1552,9 +1511,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", @@ -1578,9 +1534,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", @@ -1601,9 +1554,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", @@ -1616,12 +1566,17 @@ xla_test( "//tensorflow/core:stream_executor_no_cuda", "//tensorflow/core:test", "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/base", ], ) xla_test( - name = "cross_replica_sum_test", - srcs = ["cross_replica_sum_test.cc"], + 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", @@ -1646,9 +1601,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", @@ -1688,9 +1640,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", @@ -1752,6 +1701,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", @@ -1766,6 +1719,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", @@ -1779,9 +1736,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", @@ -1804,9 +1758,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", @@ -1860,6 +1811,7 @@ xla_test( "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:lib", "//tensorflow/core:test", + "@com_google_absl//absl/base", "@com_google_absl//absl/types:span", ], ) @@ -1867,9 +1819,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", @@ -1896,6 +1845,7 @@ xla_test( xla_test( name = "multioutput_fusion_test", srcs = ["multioutput_fusion_test.cc"], + backends = ["gpu"], deps = [ "//tensorflow/compiler/xla:literal", "//tensorflow/compiler/xla:shape_util", @@ -1986,6 +1936,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", @@ -2152,6 +2106,7 @@ xla_test( "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:lib", "//tensorflow/core:test", + "@com_google_absl//absl/base", "@com_google_absl//absl/container:flat_hash_set", ], ) @@ -2161,7 +2116,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", ], @@ -2189,3 +2143,18 @@ 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", + ], +) diff --git a/tensorflow/compiler/xla/tests/all_reduce_test.cc b/tensorflow/compiler/xla/tests/all_reduce_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..7e695f829e39831e2c8558cb07d0689e560bbafa --- /dev/null +++ b/tensorflow/compiler/xla/tests/all_reduce_test.cc @@ -0,0 +1,102 @@ +/* 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/literal.h" +#include "tensorflow/compiler/xla/service/hlo_parser.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/test.h" +#include "tensorflow/compiler/xla/test_helpers.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" +#include "tensorflow/compiler/xla/tests/test_macros.h" + +namespace xla { +namespace { + +class TrivialCrossReplicaSumTest : public HloTestBase {}; + +// Currently the CPU and GPU backends only support CrossReplicaSum with one +// replica. But we can at least check this. + +XLA_TEST_F(TrivialCrossReplicaSumTest, OneOperand) { + const char* module_str = R"( + HloModule test + + add { + x = f32[] parameter(0) + y = f32[] parameter(1) + add = f32[] add(x, y) + } + + ENTRY test_computation { + p = f32[3] parameter(0) + ROOT crs = f32[3] all-reduce(p), to_apply=add + })"; + auto module = + ParseHloString(module_str, GetModuleConfigForTest()).ValueOrDie(); + auto literal = LiteralUtil::CreateR1({1, 2, 3}); + EXPECT_EQ(literal, ExecuteAndTransfer(std::move(module), {&literal})); +} + +XLA_TEST_F(TrivialCrossReplicaSumTest, MultipleOperands) { + const char* module_str = R"( + HloModule test + + add { + x = f32[] parameter(0) + y = f32[] parameter(1) + add = f32[] add(x, y) + } + + ENTRY test_computation { + p0 = f32[3] parameter(0) + p1 = f32[2] parameter(1) + ROOT crs = (f32[3], f32[2]) all-reduce(p0, p1), to_apply=add + })"; + auto module = + ParseHloString(module_str, GetModuleConfigForTest()).ValueOrDie(); + auto literal0 = LiteralUtil::CreateR1({1, 2, 3}); + auto literal1 = LiteralUtil::CreateR1({10, 20}); + EXPECT_EQ(LiteralUtil::MakeTuple({&literal0, &literal1}), + ExecuteAndTransfer(std::move(module), {&literal0, &literal1})); +} + +// On the GPU backend, constants get special handling. Someone might pass a +// constant to CRS to e.g. count the number of replicas -- we need to make sure +// it works. +XLA_TEST_F(TrivialCrossReplicaSumTest, ConstantOperand) { + const char* module_str = R"( + HloModule test + + add { + x = f32[] parameter(0) + y = f32[] parameter(1) + add = f32[] add(x, y) + } + + ENTRY test_computation { + p0 = f32[3] parameter(0) + p1 = f32[2] constant({10, 20}) + ROOT crs = (f32[3], f32[2]) all-reduce(p0, p1), to_apply=add + })"; + auto module = + ParseHloString(module_str, GetModuleConfigForTest()).ValueOrDie(); + auto literal0 = LiteralUtil::CreateR1({1, 2, 3}); + auto literal1 = LiteralUtil::CreateR1({10, 20}); + EXPECT_EQ(LiteralUtil::MakeTuple({&literal0, &literal1}), + ExecuteAndTransfer(std::move(module), {&literal0})); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/tests/array_elementwise_ops_test.cc b/tensorflow/compiler/xla/tests/array_elementwise_ops_test.cc index c257566fb218d4769aec0c793efb9256b023b7ea..7379fbcc22745f46f2a29732c4bda46f352d07e7 100644 --- a/tensorflow/compiler/xla/tests/array_elementwise_ops_test.cc +++ b/tensorflow/compiler/xla/tests/array_elementwise_ops_test.cc @@ -19,6 +19,7 @@ limitations under the License. #include #include +#include "absl/base/casts.h" #include "absl/types/span.h" #include "tensorflow/compiler/xla/array2d.h" #include "tensorflow/compiler/xla/array3d.h" @@ -35,7 +36,6 @@ limitations under the License. #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/lib/core/casts.h" #include "tensorflow/core/platform/types.h" namespace xla { @@ -139,7 +139,7 @@ XLA_TEST_F(ArrayElementwiseOpTest, IsFiniteZeroElementF32s) { } // A non-canonical quiet NaN value. -static const float kNonCanonicalNaN = tensorflow::bit_cast(0x7FD01234); +static const float kNonCanonicalNaN = absl::bit_cast(0x7FD01234); XLA_TEST_F(ArrayElementwiseOpTest, IsFiniteScalarF32) { XlaBuilder builder(TestName()); @@ -329,13 +329,13 @@ TEST_P(ArrayElementwiseOpTestParamCount, AddManyValues) { Literal b_literal = LiteralUtil::CreateR1({b_values}); std::unique_ptr b_data = client_->TransferToServer(b_literal).ConsumeValueOrDie(); - auto b_constant = Parameter(&builder, 1, a_literal.shape(), "b_param"); - auto b_param = ConstantR1(&builder, b_values); + auto b_param = Parameter(&builder, 1, a_literal.shape(), "b_param"); + auto b_constant = ConstantR1(&builder, b_values); - auto sum1 = Add(a_constant, b_constant); - auto sum2 = Add(a_constant, b_param); - auto sum3 = Add(a_param, b_constant); - auto sum4 = Add(a_param, b_param); + auto sum1 = Add(a_constant, b_param); + auto sum2 = Add(a_constant, b_constant); + auto sum3 = Add(a_param, b_param); + auto sum4 = Add(a_param, b_constant); auto sum = Add(sum1, sum2); sum = Add(sum, sum3); @@ -350,6 +350,44 @@ TEST_P(ArrayElementwiseOpTestParamCount, AddManyValues) { error_spec_); } +// TODO(b/119692968): This test runs OOM on the GPU and CPU backend. +XLA_TEST_F(ArrayElementwiseOpTest, + DISABLED_ON_GPU(DISABLED_ON_CPU(DeeplyNestedAddWithSlices))) { + XlaBuilder builder(TestName()); + std::vector values(30, 0.0); + auto a_literal = LiteralUtil::CreateR1(values); + auto a = Parameter(&builder, 0, a_literal.shape(), "x"); + auto b_literal = LiteralUtil::CreateR1(values); + auto b = Parameter(&builder, 1, b_literal.shape(), "x"); + + // Construct a sequence of diamond-shaped gadgets like this: + // + // add + // / \ + // slice slice + // \ / + // add + // + // Each 'left' slice removes the last element, each 'right' slice removes the + // first element. In this way, we index into the add with different + // multi-dimensional index arrays, which defeats the caching we use to avoid + // exponential compile time. + std::function generate_recursive = + [&](int64 slice_size) -> XlaOp { + if (slice_size == values.size()) { + return Add(a, b); + } + XlaOp param = generate_recursive(slice_size + 1); + auto slice1 = Slice(param, {0}, {slice_size}, {1}); + auto slice2 = Slice(param, {1}, {slice_size + 1}, {1}); + return Add(slice1, slice2); + }; + generate_recursive(1); + auto a_data = client_->TransferToServer(a_literal).ConsumeValueOrDie(); + auto b_data = client_->TransferToServer(b_literal).ConsumeValueOrDie(); + ComputeAndCompareR1(&builder, {0.0}, {a_data.get(), b_data.get()}); +} + XLA_TEST_F(ArrayElementwiseOpTest, SubTwoConstantF32s) { XlaBuilder builder(TestName()); auto a = ConstantR1(&builder, {-2.5f, 3.14f, 2.25f, -10.0f, 6.0f}); @@ -1405,6 +1443,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, {}); @@ -2009,6 +2068,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); @@ -2478,8 +2550,8 @@ XLA_TEST_F(ArrayElementwiseOpTest, Compare1DTo2DS32Ne) { Ne(v, m, /*broadcast_dimensions=*/{1}); const string expected = R"(pred[2,2] { - { 00 }, - { 01 } + { 0, 0 }, + { 0, 1 } })"; EXPECT_EQ(expected, ExecuteToString(&builder, {})); } @@ -2492,8 +2564,8 @@ XLA_TEST_F(ArrayElementwiseOpTest, Compare1DTo2DS32Ge) { Ge(v, m, /*broadcast_dimensions=*/{1}); const string expected = R"(pred[2,4] { - { 1100 }, - { 0001 } + { 1, 1, 0, 0 }, + { 0, 0, 0, 1 } })"; EXPECT_EQ(expected, ExecuteToString(&builder, {})); } @@ -2506,8 +2578,8 @@ XLA_TEST_F(ArrayElementwiseOpTest, Compare1DTo2DS32Gt) { Gt(v, m, /*broadcast_dimensions=*/{1}); const string expected = R"(pred[2,4] { - { 0100 }, - { 0000 } + { 0, 1, 0, 0 }, + { 0, 0, 0, 0 } })"; EXPECT_EQ(expected, ExecuteToString(&builder, {})); } @@ -2520,8 +2592,8 @@ XLA_TEST_F(ArrayElementwiseOpTest, Compare1DTo2DS32Le) { Le(v, m, /*broadcast_dimensions=*/{1}); const string expected = R"(pred[2,4] { - { 1011 }, - { 1111 } + { 1, 0, 1, 1 }, + { 1, 1, 1, 1 } })"; EXPECT_EQ(expected, ExecuteToString(&builder, {})); } @@ -2534,8 +2606,8 @@ XLA_TEST_F(ArrayElementwiseOpTest, Compare1DTo2DS32Lt) { Lt(v, m, /*broadcast_dimensions=*/{1}); const string expected = R"(pred[2,4] { - { 0011 }, - { 1110 } + { 0, 0, 1, 1 }, + { 1, 1, 1, 0 } })"; EXPECT_EQ(expected, ExecuteToString(&builder, {})); } @@ -2744,12 +2816,16 @@ XLA_TEST_F(ArrayElementwiseOpTest, CompareGtR3F32sWithDegenerateDim2) { Array3D expected_3d( {{{0, 1}, {0, 0}, {0, 0}}, {{0, 1}, {1, 0}, {0, 1}}}); const string expected = R"(pred[2,3,2] { -{ { 01 }, - { 00 }, - { 00 } }, -{ { 01 }, - { 10 }, - { 01 } } +{ + { 0, 1 }, + { 0, 0 }, + { 0, 0 } +}, +{ + { 0, 1 }, + { 1, 0 }, + { 0, 1 } +} })"; EXPECT_EQ(expected, ExecuteToString(&builder, {})); } 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/broadcast_simple_test.cc b/tensorflow/compiler/xla/tests/broadcast_simple_test.cc index dde19fb65d65064c9452a6ac49c70e20cf113336..702fb32adfc8a0ded26845c92245776a79777c34 100644 --- a/tensorflow/compiler/xla/tests/broadcast_simple_test.cc +++ b/tensorflow/compiler/xla/tests/broadcast_simple_test.cc @@ -161,8 +161,7 @@ XLA_TEST_F(BroadcastSimpleTest, 1DTo2D) { XLA_TEST_F(BroadcastSimpleTest, 1DTo2D_WithDimsUsual) { XlaBuilder b(TestName()); - BroadcastInDim(ConstantR1(&b, {1, 2}), - ShapeUtil::MakeShape(F32, {2, 2}), {1}); + BroadcastInDim(ConstantR1(&b, {1, 2}), {2, 2}, {1}); Array2D expected(2, 2); expected(0, 0) = 1; @@ -175,8 +174,7 @@ XLA_TEST_F(BroadcastSimpleTest, 1DTo2D_WithDimsUsual) { XLA_TEST_F(BroadcastSimpleTest, 1DTo2D_WithDimsTranspose) { XlaBuilder b(TestName()); - BroadcastInDim(ConstantR1(&b, {1, 2}), - ShapeUtil::MakeShape(F32, {2, 2}), {0}); + BroadcastInDim(ConstantR1(&b, {1, 2}), {2, 2}, {0}); Array2D expected(2, 2); expected(0, 0) = 1; @@ -189,8 +187,8 @@ XLA_TEST_F(BroadcastSimpleTest, 1DTo2D_WithDimsTranspose) { XLA_TEST_F(BroadcastSimpleTest, 2DTo3D_WithDims) { XlaBuilder b(TestName()); - BroadcastInDim(ConstantR2(&b, {{1.0, 5.0}, {2.0, 6.0}}), - ShapeUtil::MakeShape(F32, {2, 2, 2}), {0, 1}); + BroadcastInDim(ConstantR2(&b, {{1.0, 5.0}, {2.0, 6.0}}), {2, 2, 2}, + {0, 1}); Array3D expected(2, 2, 2); expected(0, 0, 0) = 1.0; @@ -207,8 +205,8 @@ XLA_TEST_F(BroadcastSimpleTest, 2DTo3D_WithDims) { XLA_TEST_F(BroadcastSimpleTest, 2DTo3D_WithDimsNotPossibleWithBroadCast) { XlaBuilder b(TestName()); - BroadcastInDim(ConstantR2(&b, {{1.0, 5.0}, {2.0, 6.0}}), - ShapeUtil::MakeShape(F32, {2, 2, 2}), {0, 2}); + BroadcastInDim(ConstantR2(&b, {{1.0, 5.0}, {2.0, 6.0}}), {2, 2, 2}, + {0, 2}); Array3D expected(2, 2, 2); expected(0, 0, 0) = 1.0; @@ -225,8 +223,7 @@ XLA_TEST_F(BroadcastSimpleTest, 2DTo3D_WithDimsNotPossibleWithBroadCast) { XLA_TEST_F(BroadcastSimpleTest, 1DTo2D_WithDimsNotPossibleWithBroadCast) { XlaBuilder b(TestName()); - BroadcastInDim(ConstantR1(&b, {1, 2}), - ShapeUtil::MakeShape(F32, {3, 2}), {1}); + BroadcastInDim(ConstantR1(&b, {1, 2}), {3, 2}, {1}); Array2D expected(3, 2); expected(0, 0) = 1; diff --git a/tensorflow/compiler/xla/tests/broadcast_test.cc b/tensorflow/compiler/xla/tests/broadcast_test.cc index 9966e4606ef7f104487182e0240e64e4c9e4d834..9930bfc95c297093584d427397cac042c296050f 100644 --- a/tensorflow/compiler/xla/tests/broadcast_test.cc +++ b/tensorflow/compiler/xla/tests/broadcast_test.cc @@ -42,7 +42,7 @@ XLA_TEST_F(BroadcastTest, BroadcastScalarToScalar) { ShapeUtil::MakeShape(F32, {}), input, {})); // Create HLO module, compile, and execute. - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); hlo_module->AddEntryComputation(builder.Build()); auto result = ExecuteAndTransfer(std::move(hlo_module), {}); @@ -58,7 +58,7 @@ XLA_TEST_F(BroadcastTest, BroadcastScalarTo2D) { ShapeUtil::MakeShape(F32, {2, 2}), input, {})); // Create HLO module, compile, and execute. - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); hlo_module->AddEntryComputation(builder.Build()); auto result = ExecuteAndTransfer(std::move(hlo_module), {}); @@ -81,7 +81,7 @@ XLA_TEST_F(BroadcastTest, BroadcastVectorTo2D) { builder.AddInstruction(HloInstruction::CreateTuple({element1, element2})); // Create HLO module, compile, and execute. - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); hlo_module->AddEntryComputation(builder.Build()); auto result = ExecuteAndTransfer(std::move(hlo_module), {}); @@ -102,7 +102,7 @@ XLA_TEST_F(BroadcastTest, Broadcast2DTo2D) { ShapeUtil::MakeShape(F32, {2, 2}), input, {0, 1})); // Create HLO module, compile, and execute. - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); hlo_module->AddEntryComputation(builder.Build()); auto result = ExecuteAndTransfer(std::move(hlo_module), {}); @@ -121,7 +121,7 @@ XLA_TEST_F(BroadcastTest, Broadcast2DTo2DTranspose) { ShapeUtil::MakeShape(F32, {2, 2}), input, {1, 0})); // Create HLO module, compile, and execute. - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); hlo_module->AddEntryComputation(builder.Build()); auto result = ExecuteAndTransfer(std::move(hlo_module), {}); @@ -138,7 +138,7 @@ XLA_TEST_F(BroadcastTest, Broadcast2DTo3D) { ShapeUtil::MakeShape(F32, {2, 3, 2}), input, {0, 2})); // Create HLO module, compile, and execute. - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); hlo_module->AddEntryComputation(builder.Build()); auto result = ExecuteAndTransfer(std::move(hlo_module), {}); @@ -158,7 +158,7 @@ TEST_F(BroadcastTest, Broadcast_R1_2_To_R4_2x2x3x3) { ShapeUtil::MakeShape(F32, {2, 2, 3, 3}), input, {1})); // Create HLO module, compile, and execute. - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); hlo_module->AddEntryComputation(builder.Build()); auto result = ExecuteAndTransfer(std::move(hlo_module), {}); @@ -183,7 +183,7 @@ TEST_F(BroadcastTest, Broadcast_R1_1025_To_R4_3x3x3x1025) { ShapeUtil::MakeShape(F32, {3, 3, 3, r1_size}), input, {3})); // Create HLO module, compile, and execute. - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); hlo_module->AddEntryComputation(builder.Build()); auto result = ExecuteAndTransfer(std::move(hlo_module), {}); @@ -214,7 +214,7 @@ XLA_TEST_F(BroadcastTest, Broadcast_R1_64_To_R4_32x64x7x7) { ShapeUtil::MakeShape(F32, {32, 64, 7, 7}), input, {1})); // Create HLO module, compile, and execute. - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); hlo_module->AddEntryComputation(builder.Build()); auto result = ExecuteAndTransfer(std::move(hlo_module), {}); @@ -230,7 +230,7 @@ TEST_F(BroadcastTest, Broadcast_R0_to_R4_64x64x3x3) { ShapeUtil::MakeShape(F32, {64, 64, 3, 3}), input, {})); // Create HLO module, compile, and execute. - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); hlo_module->AddEntryComputation(builder.Build()); LOG(INFO) << hlo_module->ToString(); auto result = ExecuteAndTransfer(std::move(hlo_module), {}); @@ -253,7 +253,7 @@ TEST_F(BroadcastTest, Broadcast_R2_2x2_To_R4_3x3x2x2) { ShapeUtil::MakeShape(F32, {3, 3, 2, 2}), input, {2, 3})); // Create HLO module, compile, and execute. - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); hlo_module->AddEntryComputation(builder.Build()); auto result = ExecuteAndTransfer(std::move(hlo_module), {}); @@ -287,7 +287,7 @@ TEST_F(BroadcastTest, Broadcast_R3_2x3x4_to_R4_2x3x4x5) { ShapeUtil::MakeShape(F32, {2, 3, 4, 5}), input, {0, 1, 2})); // Create HLO module, compile, and execute. - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); hlo_module->AddEntryComputation(builder.Build()); auto result = ExecuteAndTransfer(std::move(hlo_module), {}); diff --git a/tensorflow/compiler/xla/tests/client_library_test_base.cc b/tensorflow/compiler/xla/tests/client_library_test_base.cc index fbdf0fcb6543f09dedefef55cfe0f8a5d9067d5a..edb95c973b70e30702ed8490c15a48d4d5604170 100644 --- a/tensorflow/compiler/xla/tests/client_library_test_base.cc +++ b/tensorflow/compiler/xla/tests/client_library_test_base.cc @@ -15,6 +15,7 @@ limitations under the License. #include "tensorflow/compiler/xla/tests/client_library_test_base.h" +#include #include #include "absl/memory/memory.h" @@ -74,6 +75,9 @@ ClientLibraryTestBase::ClientLibraryTestBase( // default. execution_options_.mutable_debug_options()->add_xla_disable_hlo_passes( "constant_folding"); + + execution_options_.mutable_debug_options() + ->set_xla_hlo_evaluator_use_fast_path(true); } ClientLibraryTestBase::ClientLibraryTestBase(se::Platform* platform) @@ -88,6 +92,9 @@ ClientLibraryTestBase::ClientLibraryTestBase(se::Platform* platform) execution_options_.mutable_debug_options()->add_xla_disable_hlo_passes( "constant_folding"); + + execution_options_.mutable_debug_options() + ->set_xla_hlo_evaluator_use_fast_path(true); } string ClientLibraryTestBase::TestName() const { @@ -107,7 +114,7 @@ StatusOr ClientLibraryTestBase::ExecuteAndTransfer( ExecutionOptions execution_options = execution_options_; if (shape_with_output_layout != nullptr) { *execution_options.mutable_shape_with_output_layout() = - *shape_with_output_layout; + shape_with_output_layout->ToProto(); } return client_->ExecuteAndTransfer(computation, arguments, &execution_options); @@ -127,7 +134,7 @@ StatusOr ClientLibraryTestBase::ExecuteAndTransferReference( ExecutionOptions execution_options = execution_options_; if (shape_with_output_layout != nullptr) { *execution_options.mutable_shape_with_output_layout() = - *shape_with_output_layout; + shape_with_output_layout->ToProto(); } execution_options.clear_device_handles(); return ref_client_->ExecuteAndTransfer(computation, arguments, @@ -184,7 +191,7 @@ Status ClientLibraryTestBase::ComputeAndCompareLiteralWithAllOutputLayouts( verify_output(actual, ""); // Try with all output layouts. - std::vector minor_to_major(ShapeUtil::Rank(expected.shape())); + std::vector minor_to_major(expected.shape().rank()); std::iota(minor_to_major.begin(), minor_to_major.end(), 0); do { auto layout = ShapeUtil::MakeShapeWithLayout( @@ -217,7 +224,7 @@ Status ClientLibraryTestBase::ComputeAndCompareLiteralWithAllInputLayouts( TF_ASSIGN_OR_RETURN(auto literal, client_->Transfer(*arguments[index], nullptr)); // Skip tuples because they don't have a rank. - if (ShapeUtil::IsTuple(literal.shape())) { + if (literal.shape().IsTuple()) { layout_strings.push_back( ShapeUtil::HumanStringWithLayout(literal.shape())); arguments_with_layout.push_back(arguments[index]); @@ -227,7 +234,7 @@ Status ClientLibraryTestBase::ComputeAndCompareLiteralWithAllInputLayouts( return Status::OK(); } - std::vector minor_to_major(ShapeUtil::Rank(literal.shape())); + std::vector minor_to_major(literal.shape().rank()); std::iota(minor_to_major.begin(), minor_to_major.end(), 0); do { auto literal_relayout = @@ -262,6 +269,29 @@ Status ClientLibraryTestBase::ComputeAndCompareLiteralWithAllInputLayouts( return choose(0); } +StatusOr ClientLibraryTestBase::ComputeAndTransfer( + XlaBuilder* builder, absl::Span arguments_passed_in, + const Shape* shape_with_layout) { + std::vector arguments(arguments_passed_in.begin(), + arguments_passed_in.end()); + + // Transfer and use elements of arguments_, if the AddParam() API was used. + std::vector> owning_arguments; + if (!arguments_.empty()) { + CHECK(arguments.empty()); + for (const auto& argument : arguments_) { + TF_ASSIGN_OR_RETURN( + std::unique_ptr owned_argument, + client_->TransferToServer(MaybeConvertLiteralToBfloat16(argument))); + owning_arguments.push_back(std::move(owned_argument)); + arguments.push_back(owning_arguments.back().get()); + } + } + + TF_ASSIGN_OR_RETURN(auto computation, builder->Build()); + return ExecuteAndTransfer(computation, arguments, shape_with_layout); +} + Status ClientLibraryTestBase::ComputeAndCompareLiteralWithStatus( XlaBuilder* builder, const Literal& expected, absl::Span arguments_passed_in, @@ -274,9 +304,10 @@ Status ClientLibraryTestBase::ComputeAndCompareLiteralWithStatus( if (!arguments_.empty()) { CHECK(arguments.empty()); for (const auto& argument : arguments_) { - owning_arguments.push_back( - client_->TransferToServer(MaybeConvertLiteralToBfloat16(argument)) - .ValueOrDie()); + TF_ASSIGN_OR_RETURN( + std::unique_ptr owned_argument, + client_->TransferToServer(MaybeConvertLiteralToBfloat16(argument))); + owning_arguments.push_back(std::move(owned_argument)); arguments.push_back(owning_arguments.back().get()); } } @@ -334,9 +365,10 @@ Status ClientLibraryTestBase::ComputeAndCompareLiteralWithStatus( if (!arguments_.empty()) { CHECK(arguments.empty()); for (const auto& argument : arguments_) { - owning_arguments.push_back( - client_->TransferToServer(MaybeConvertLiteralToBfloat16(argument)) - .ValueOrDie()); + TF_ASSIGN_OR_RETURN( + std::unique_ptr owned_argument, + client_->TransferToServer(MaybeConvertLiteralToBfloat16(argument))); + owning_arguments.push_back(std::move(owned_argument)); arguments.push_back(owning_arguments.back().get()); } } diff --git a/tensorflow/compiler/xla/tests/client_library_test_base.h b/tensorflow/compiler/xla/tests/client_library_test_base.h index 9d32f4f5174a57a53a9d3e6477b46fa4de852f7f..3f65ed7fce4ff4b5c3781ac2581935bfacc69ce1 100644 --- a/tensorflow/compiler/xla/tests/client_library_test_base.h +++ b/tensorflow/compiler/xla/tests/client_library_test_base.h @@ -76,7 +76,7 @@ class ClientLibraryTestBase : public ::testing::Test { void SetFastMathDisabled(bool disabled) { auto* opts = execution_options_.mutable_debug_options(); opts->set_xla_cpu_enable_fast_math(!disabled); - opts->set_xla_gpu_enable_fast_math(!disabled); + opts->set_xla_gpu_enable_fast_min_max(!disabled); } void SetSeed(uint64 seed) { execution_options_.set_seed(seed); } @@ -188,6 +188,13 @@ class ClientLibraryTestBase : public ::testing::Test { ErrorSpec error, const Shape* shape_with_layout = nullptr); + // Build and run the computation and return the result as a literal. + // shape_with_layout indicates the result layout to request when calling + // Execute. + StatusOr ComputeAndTransfer( + XlaBuilder* builder, absl::Span arguments, + const Shape* shape_with_layout = nullptr); + // ComputeAndCompare variant which returns an error status. Status ComputeAndCompareLiteralWithStatus( XlaBuilder* builder, const Literal& expected, @@ -424,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, @@ -448,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, @@ -473,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); @@ -499,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); @@ -525,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 6f2ca84bb646e88af221ab80b727911ff7d990eb..247328b730f3af936d933f824da491b593b27c90 100644 --- a/tensorflow/compiler/xla/tests/client_test.cc +++ b/tensorflow/compiler/xla/tests/client_test.cc @@ -50,7 +50,8 @@ XLA_TEST_F(ClientTest, ExecuteWithLayout) { ExecutionOptions execution_options = execution_options_; *execution_options.mutable_shape_with_output_layout() = ShapeUtil::MakeShapeWithLayout(S32, /*dimensions=*/{2, 2}, - execute_layout); + execute_layout) + .ToProto(); TF_ASSERT_OK_AND_ASSIGN( std::unique_ptr data, client_->Execute(computation, {}, &execution_options)); @@ -84,7 +85,8 @@ XLA_TEST_F(ClientTest, ExecuteWithTupleLayout) { {ShapeUtil::MakeShapeWithLayout(S32, /*dimensions=*/{2, 2}, /*minor_to_major=*/{0, 1}), ShapeUtil::MakeShapeWithLayout(S32, /*dimensions=*/{2, 2}, - /*minor_to_major=*/{1, 0})}); + /*minor_to_major=*/{1, 0})}) + .ToProto(); TF_ASSERT_OK_AND_ASSIGN( auto result, @@ -94,7 +96,7 @@ XLA_TEST_F(ClientTest, ExecuteWithTupleLayout) { LiteralTestUtil::ExpectR2Equal({{10, 20}, {30, 40}}, LiteralSlice(result, {1})); - EXPECT_TRUE(ShapeUtil::IsTuple(result.shape())); + EXPECT_TRUE(result.shape().IsTuple()); EXPECT_EQ(2, ShapeUtil::TupleElementCount(result.shape())); EXPECT_TRUE(ShapeUtil::Equal( @@ -107,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/concat_test.cc b/tensorflow/compiler/xla/tests/concat_test.cc index 9811a015e91d866d6f4de6ebb6dac536ed6c7e06..4f5b525a34252db9e967a55af0d1bf39a2dd830e 100644 --- a/tensorflow/compiler/xla/tests/concat_test.cc +++ b/tensorflow/compiler/xla/tests/concat_test.cc @@ -492,6 +492,32 @@ XLA_TEST_F(ConcatTest, ConcatR3WeirdDims) { ComputeAndCompareR3(&builder, expected, {p0.get(), p1.get()}); } +XLA_TEST_F(ConcatTest, ConcatDeeplyNested) { + XlaBuilder builder(TestName()); + auto a_literal = LiteralUtil::CreateR1({256.0}); + auto a = Parameter(&builder, 0, a_literal.shape(), "x"); + auto b = ConcatInDim(&builder, {a, a}, 0); + auto c = ConcatInDim(&builder, {b, b}, 0); + auto d = ConcatInDim(&builder, {c, c}, 0); + auto e = ConcatInDim(&builder, {d, d}, 0); + auto f = ConcatInDim(&builder, {e, e}, 0); + auto g = ConcatInDim(&builder, {f, f}, 0); + auto h = ConcatInDim(&builder, {g, g}, 0); + auto i = ConcatInDim(&builder, {h, h}, 0); + auto j = ConcatInDim(&builder, {i, i}, 0); + auto k = ConcatInDim(&builder, {j, j}, 0); + auto l = ConcatInDim(&builder, {k, k}, 0); + auto m = ConcatInDim(&builder, {l, l}, 0); + auto n = ConcatInDim(&builder, {m, m}, 0); + auto o = ConcatInDim(&builder, {n, n}, 0); + auto p = ConcatInDim(&builder, {o, o}, 0); + auto q = ConcatInDim(&builder, {p, p}, 0); + ConcatInDim(&builder, {q, q}, 0); + std::vector expected(131072, 256.0); + auto a_data = client_->TransferToServer(a_literal).ConsumeValueOrDie(); + ComputeAndCompareR1(&builder, expected, {a_data.get()}); +} + // Describes a binary rank-2 concatenation test. struct R2BinarySpec { int64 lhs_dim0; diff --git a/tensorflow/compiler/xla/tests/constants_test.cc b/tensorflow/compiler/xla/tests/constants_test.cc index 72ff1e74a47c8584cb5336c86a1c978c4637a902..6530007871ced1d0bbffe2b44ccc8cf9bddd79e1 100644 --- a/tensorflow/compiler/xla/tests/constants_test.cc +++ b/tensorflow/compiler/xla/tests/constants_test.cc @@ -21,11 +21,14 @@ 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" #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/core/lib/core/status_test_util.h" #include "tensorflow/core/platform/test.h" #include "tensorflow/core/platform/types.h" @@ -178,5 +181,54 @@ 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. +XLA_TEST_F(ConstantsHloTest, DISABLED_ON_GPU(BitcastOfConstant)) { + const char* testcase = R"( + HloModule module, is_scheduled=true + + func { + lhs = s32[] parameter(0) + rhs = s32[] parameter(1) + ROOT mul = s32[] add(lhs, rhs) + } + + ENTRY test { + constant.0 = s32[1]{0} constant({0}) + parameter.0 = s32[] parameter(0) + constant-as-scalar = s32[] bitcast(constant.0) + ROOT result = s32[] call(parameter.0, constant-as-scalar), to_apply=func + } + )"; + auto module = ParseAndReturnVerifiedModule(testcase).ValueOrDie(); + auto param = LiteralUtil::CreateR0(1); + auto result = ExecuteNoHloPasses(std::move(module), {¶m}); + EXPECT_TRUE(LiteralTestUtil::Equal(param, result)); +} + } // namespace } // namespace xla 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..90c197140359d0021d08931b73f221d659e71144 --- /dev/null +++ b/tensorflow/compiler/xla/tests/conv_depthwise_backprop_filter_test.cc @@ -0,0 +1,150 @@ +/* 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 DepthwiseConvolution2DSpec { + int64 output_batch, window; + std::vector activation_dims; + std::vector activation_layout; + std::vector kernel_dims; + std::vector kernel_layout; + std::vector output_dims; + std::vector output_layout; +}; + +class DepthwiseConvolution2DTest + : public HloTestBase, + public ::testing::WithParamInterface< + ::testing::tuple> {}; + +static std::vector GetConv2DTestCases() { + std::vector config_set; + std::vector> config_options = { + {16, 5, 5, 2}, {64, 4, 4, 16}, {2, 5, 5, 256}}; + + for (auto option : config_options) { + int64 feature = option[3]; + int64 activation_size = option[1]; + int64 kernel_size = option[2]; + int64 batch = option[0]; + + std::vector kernel_layout = {3, 2, 1, 0}; + DepthwiseConvolution2DSpec config; + config.output_batch = feature; + config.window = kernel_size; + + config.activation_dims = {batch, activation_size, activation_size, feature}; + config.activation_layout = {0, 3, 2, 1}; + + config.kernel_dims = {batch, kernel_size, kernel_size, feature}; + config.kernel_layout = {0, 2, 3, 1}; + + config.output_dims = {3, 3, feature, 1}; + + // Try this layout for all kernel shapes. + config.output_layout = {3, 2, 0, 1}; + config_set.push_back(config); + } + + return config_set; +} + +string DepthwiseConvolution2DTestDataToString( + 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"), + "_activation_layout_", absl::StrJoin(spec.activation_layout, "_"), + "_kernel_dims_", absl::StrJoin(spec.kernel_dims, "x"), "_kernel_layout_", + absl::StrJoin(spec.kernel_layout, "_"), "_output_dims_", + absl::StrJoin(spec.output_dims, "x"), "_output_layout_", + absl::StrJoin(spec.output_layout, "_"), data_type); + + // Test names are not allowed to contain the '-' character. + absl::c_replace(str, '-', 'n'); + return str; +} + +string BuildHloTextDepthwiseConvolution2D( + const DepthwiseConvolution2DSpec& spec, bool use_bfloat16) { + const string data_type = GetFloatDataType(use_bfloat16); + return absl::StrFormat( + R"( + HloModule TensorFlowDepthwiseConv + + 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_1x1_1}, dim_labels=f01b_i01o->01fb, + batch_group_count=%d + } + )", + data_type, absl::StrJoin(spec.activation_dims, ","), + absl::StrJoin(spec.activation_layout, ","), data_type, + absl::StrJoin(spec.kernel_dims, ","), + absl::StrJoin(spec.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_layout, ","), data_type, + absl::StrJoin(spec.kernel_dims, ","), + absl::StrJoin(spec.kernel_layout, ","), spec.window, spec.window, + spec.output_batch); +} + +XLA_TEST_P(DepthwiseConvolution2DTest, DoIt) { + const DepthwiseConvolution2DSpec& spec = ::testing::get<0>(GetParam()); + bool use_bfloat16 = ::testing::get<1>(GetParam()); + const string hlo_text = + BuildHloTextDepthwiseConvolution2D(spec, use_bfloat16); + + EXPECT_TRUE(RunAndCompare(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( + DepthwiseConvolution2DTestWithRandomIndices, DepthwiseConvolution2DTest, + ::testing::Combine(::testing::ValuesIn(GetConv2DTestCases()), + ::testing::Bool()), + DepthwiseConvolution2DTestDataToString); + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/tests/conv_depthwise_test.cc b/tensorflow/compiler/xla/tests/conv_depthwise_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..627a17a0ca114085240dbaf28211bb3511cf0cab --- /dev/null +++ b/tensorflow/compiler/xla/tests/conv_depthwise_test.cc @@ -0,0 +1,234 @@ +/* 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 DepthwiseConvolution2DSpec { + int64 output_feature, window, stride, pad, lhs_dilate; + std::vector activation_dims; + std::vector activation_layout; + std::vector kernel_dims; + std::vector kernel_layout; + std::vector output_dims; + std::vector output_layout; +}; + +class DepthwiseConvolution2DTest + : public HloTestBase, + public ::testing::WithParamInterface< + ::testing::tuple> {}; + +static std::vector GetConv2DTestCases() { + std::vector config_set; + std::vector> config_options = { + {128, 6, 3, 64}, {256, 5, 3, 256}, {256, 5, 2, 144}, {144, 5, 3, 64}, + {144, 5, 2, 256}, {8, 48, 17, 8}, {128, 20, 6, 64}, {64, 14, 12, 172}, + {16, 9, 4, 16}, {128, 1, 2, 144}, {256, 1, 2, 64}}; + + for (auto option : config_options) { + int64 feature = option[0]; + int64 activation_size = option[1]; + int64 kernel_size = option[2]; + int64 batch = option[3]; + + std::vector kernel_layout = {3, 2, 1, 0}; + DepthwiseConvolution2DSpec config; + config.output_feature = feature; + config.window = kernel_size; + + config.activation_dims = {batch, activation_size, activation_size, feature}; + config.activation_layout = {3, 0, 2, 1}; + + config.kernel_dims = {kernel_size, kernel_size, 1, feature}; + config.kernel_layout = {3, 2, 1, 0}; + + if (activation_size == 1 && kernel_size == 2) { + // Test for outer dim. + config.output_dims = {batch, activation_size + kernel_size - 1, + activation_size + kernel_size, feature}; + } else if (feature == 256) { + // Restrict dilation-based tests only to one feature configuration. + config.stride = activation_size - 1; + config.pad = 0; + config.lhs_dilate = feature / 32; + config.output_dims = {batch, feature / 32, + activation_size - kernel_size + 1, feature}; + } else { + config.stride = config.pad = config.lhs_dilate = -1; + config.output_dims = {batch, activation_size - kernel_size + 1, + activation_size - kernel_size + 1, feature}; + } + + // Try this layout for all kernel shapes. + config.output_layout = {3, 0, 2, 1}; + config_set.push_back(config); + + // Try other layouts only for certain kernel shapes. + if (kernel_size % 2 == 0) { + config.activation_layout = {0, 3, 2, 1}; + config_set.push_back(config); + + config.output_layout = {0, 3, 2, 1}; + config_set.push_back(config); + + config.activation_layout = {3, 0, 2, 1}; + config_set.push_back(config); + } + } + + return config_set; +} + +string DepthwiseConvolution2DTestDataToString( + 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"), + "_activation_layout_", absl::StrJoin(spec.activation_layout, "_"), + "_kernel_dims_", absl::StrJoin(spec.kernel_dims, "x"), "_kernel_layout_", + absl::StrJoin(spec.kernel_layout, "_"), "_output_dims_", + absl::StrJoin(spec.output_dims, "x"), "_output_layout_", + absl::StrJoin(spec.output_layout, "_"), data_type); + // -1 indicates non-existence. + if (spec.stride != -1) { + absl::StrAppend(&str, "_lhs_dilation_", spec.lhs_dilate, "x1"); + } + + // Test names are not allowed to contain the '-' character. + absl::c_replace(str, '-', 'n'); + return str; +} + +string BuildHloTextDepthwiseConvolution2D( + const DepthwiseConvolution2DSpec& spec, bool use_bfloat16) { + const string data_type = GetFloatDataType(use_bfloat16); + if (spec.activation_dims[1] == 1 && spec.kernel_dims[1] == 2) { + return absl::StrFormat( + R"( + HloModule TensorFlowDepthwiseConv + + 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_1x%d_%d rhs_dilate=1x%d}, dim_labels=b01f_01io->b01f, + feature_group_count=%d + } + )", + data_type, absl::StrJoin(spec.activation_dims, ","), + absl::StrJoin(spec.activation_layout, ","), data_type, + absl::StrJoin(spec.kernel_dims, ","), + absl::StrJoin(spec.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_layout, ","), data_type, + absl::StrJoin(spec.kernel_dims, ","), + absl::StrJoin(spec.kernel_layout, ","), spec.window, spec.window, + spec.window, spec.window, spec.window, spec.output_feature); + + } else if (spec.stride == -1) { + return absl::StrFormat( + R"( + HloModule TensorFlowDepthwiseConv + + 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}, dim_labels=b01f_01io->b01f, + feature_group_count=%d + } + )", + data_type, absl::StrJoin(spec.activation_dims, ","), + absl::StrJoin(spec.activation_layout, ","), data_type, + absl::StrJoin(spec.kernel_dims, ","), + absl::StrJoin(spec.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_layout, ","), data_type, + absl::StrJoin(spec.kernel_dims, ","), + absl::StrJoin(spec.kernel_layout, ","), spec.window, spec.window, + spec.output_feature); + } else { + return absl::StrFormat( + R"( + HloModule TensorFlowDepthwiseConv + + 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 stride=%dx1 pad=%d_%dx0_0 lhs_dilate=%dx1}, + dim_labels=b01f_01io->b01f, feature_group_count=%d + } + )", + data_type, absl::StrJoin(spec.activation_dims, ","), + absl::StrJoin(spec.activation_layout, ","), data_type, + absl::StrJoin(spec.kernel_dims, ","), + absl::StrJoin(spec.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_layout, ","), data_type, + absl::StrJoin(spec.kernel_dims, ","), + absl::StrJoin(spec.kernel_layout, ","), spec.window, spec.window, + spec.stride, 0, 0, spec.lhs_dilate, spec.output_feature); + } +} + +XLA_TEST_P(DepthwiseConvolution2DTest, DoIt) { + const DepthwiseConvolution2DSpec& spec = ::testing::get<0>(GetParam()); + bool use_bfloat16 = ::testing::get<1>(GetParam()); + const string hlo_text = + BuildHloTextDepthwiseConvolution2D(spec, use_bfloat16); + + EXPECT_TRUE(RunAndCompare(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( + DepthwiseConvolution2DTestWithRandomIndices, DepthwiseConvolution2DTest, + ::testing::Combine(::testing::ValuesIn(GetConv2DTestCases()), + ::testing::Bool()), + DepthwiseConvolution2DTestDataToString); + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/tests/convert_test.cc b/tensorflow/compiler/xla/tests/convert_test.cc index 5f063e67847487f1d18bf4ee80b1634ebdf4183a..20bf3c317986c30c12dca7dca14dbf80c70b42f6 100644 --- a/tensorflow/compiler/xla/tests/convert_test.cc +++ b/tensorflow/compiler/xla/tests/convert_test.cc @@ -20,6 +20,7 @@ limitations under the License. #include #include "absl/algorithm/container.h" +#include "absl/base/casts.h" #include "tensorflow/compiler/xla/client/local_client.h" #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/compiler/xla/shape_util.h" @@ -27,7 +28,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/xla_data.pb.h" -#include "tensorflow/core/lib/core/casts.h" #include "tensorflow/core/lib/math/math_util.h" #include "tensorflow/core/platform/stream_executor_no_cuda.h" #include "tensorflow/core/platform/test.h" @@ -429,11 +429,9 @@ TEST_F(ConvertTest, ConvertReshape) { std::vector GetInterestingF16ConversionTestCases() { float infinity = std::numeric_limits::infinity(); - float half_min_positive_normal = - tensorflow::bit_cast(0x38800000); - float half_max_subnormal = tensorflow::bit_cast(0x387fc000); - float half_min_positive_subnormal = - tensorflow::bit_cast(0x33800000); + float half_min_positive_normal = absl::bit_cast(0x38800000); + float half_max_subnormal = absl::bit_cast(0x387fc000); + float half_min_positive_subnormal = absl::bit_cast(0x33800000); float half_max = 65504.0f; std::vector test_cases( diff --git a/tensorflow/compiler/xla/tests/convolution_test.cc b/tensorflow/compiler/xla/tests/convolution_test.cc index 3aebf784664dac14ba2ea45c5a229b7b2e4fc39d..9db9f2563b636c4f929585eb13a9c7f809833eda 100644 --- a/tensorflow/compiler/xla/tests/convolution_test.cc +++ b/tensorflow/compiler/xla/tests/convolution_test.cc @@ -98,7 +98,7 @@ class ForwardPassConvolution_3x3x256_256_OutputZ_Iota : public ConvolutionTest { precision.add_operand_precision(PrecisionConfig::HIGHEST); precision.add_operand_precision(PrecisionConfig::DEFAULT); Conv(lhs, rhs, {1, 1}, Padding::kValid, /*feature_group_count=*/1, - &precision); + /*batch_group_count=*/1, &precision); ComputeAndCompare(&builder, {}, error_spec_); } @@ -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 @@ -597,7 +597,692 @@ TYPED_TEST(Convolve2D_1x3x3x5_3x3x1x15_Depthwise_Valid, Types) { } template -class Convolve2D_1x2x2x6_2x2x1x12_Grouped_Valid : public ConvolutionTest { +class Convolve2D_1x4x4x5_3x3x1x5_Depthwise_Valid : public ConvolutionTest { + public: + void RunTest() { + XlaBuilder builder(TestName()); + std::vector input_dims = {1, 4, 4, 5}; + std::vector filter_dims = {3, 3, 1, 5}; + Shape input_shape = ShapeUtil::MakeShapeWithType(input_dims); + Shape filter_shape = ShapeUtil::MakeShapeWithType(filter_dims); + { + auto input = Parameter(&builder, 0, input_shape, "input"); + auto filter = Parameter(&builder, 1, filter_shape, "filter"); + + // Tensorflow dimension numbers for 2D convolution. + 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=*/5); + } + + std::vector input_elems(ShapeUtil::ElementsIn(input_shape)); + iota_int_init_value(input_elems, 1); + auto input_r1 = LiteralUtil::CreateR1(input_elems); + auto input_r4 = input_r1.Reshape(input_dims).ConsumeValueOrDie(); + + std::vector filter_elems(ShapeUtil::ElementsIn(filter_shape)); + iota_int_init_value(filter_elems, 1); + auto filter_r1 = LiteralUtil::CreateR1(filter_elems); + auto filter_r4 = filter_r1.Reshape(filter_dims).ConsumeValueOrDie(); + + auto expected_r1 = LiteralUtil::CreateR1( + {static_cast(6864), static_cast(7296), static_cast(7746), + static_cast(8214), static_cast(8700), static_cast(7809), + static_cast(8286), static_cast(8781), static_cast(9294), + static_cast(9825), static_cast(10644), static_cast(11256), + static_cast(11886), static_cast(12534), static_cast(13200), + static_cast(11589), static_cast(12246), static_cast(12921), + static_cast(13614), static_cast(14325)}); + auto expected_r4 = expected_r1.Reshape({1, 2, 2, 5}).ConsumeValueOrDie(); + + auto input_literal = + client_->TransferToServer(input_r4).ConsumeValueOrDie(); + auto filter_literal = + client_->TransferToServer(filter_r4).ConsumeValueOrDie(); + + ComputeAndCompareLiteral(&builder, expected_r4, + {input_literal.get(), filter_literal.get()}, + error_spec_); + + auto filter_r = filter_r1.Reshape(filter_dims); + } +}; + +TYPED_TEST_CASE(Convolve2D_1x4x4x5_3x3x1x5_Depthwise_Valid, TestTypes); +TYPED_TEST(Convolve2D_1x4x4x5_3x3x1x5_Depthwise_Valid, Types) { + this->RunTest(); +} + +template +class Convolve2D_1x4x4x512_3x3x1x512_Depthwise_Valid : public ConvolutionTest { + public: + void RunTest() { + XlaBuilder builder(TestName()); + std::vector input_dims = {1, 4, 4, 512}; + std::vector filter_dims = {3, 3, 1, 512}; + Shape input_shape = ShapeUtil::MakeShapeWithType(input_dims); + Shape filter_shape = ShapeUtil::MakeShapeWithType(filter_dims); + { + auto input = Parameter(&builder, 0, input_shape, "input"); + auto filter = Parameter(&builder, 1, filter_shape, "filter"); + + // Tensorflow dimension numbers for 2D convolution. + 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=*/512); + } + + std::vector input_elems(ShapeUtil::ElementsIn(input_shape), + static_cast(1)); + auto input_r1 = LiteralUtil::CreateR1(input_elems); + auto input_r4 = input_r1.Reshape(input_dims).ConsumeValueOrDie(); + + std::vector filter_elems(ShapeUtil::ElementsIn(filter_shape), + static_cast(2)); + auto filter_r1 = LiteralUtil::CreateR1(filter_elems); + auto filter_r4 = filter_r1.Reshape(filter_dims).ConsumeValueOrDie(); + + std::vector output_elems(2048, static_cast(18)); + + auto expected_r1 = LiteralUtil::CreateR1(output_elems); + auto expected_r4 = expected_r1.Reshape({1, 2, 2, 512}).ConsumeValueOrDie(); + + auto input_literal = + client_->TransferToServer(input_r4).ConsumeValueOrDie(); + auto filter_literal = + client_->TransferToServer(filter_r4).ConsumeValueOrDie(); + + ComputeAndCompareLiteral(&builder, expected_r4, + {input_literal.get(), filter_literal.get()}, + error_spec_); + } +}; + +TYPED_TEST_CASE(Convolve2D_1x4x4x512_3x3x1x512_Depthwise_Valid, TestTypes); +TYPED_TEST(Convolve2D_1x4x4x512_3x3x1x512_Depthwise_Valid, Types) { + this->RunTest(); +} + +template +class Convolve2D_1x4x4x512_3x3x1x512_Depthwise_Valid_Output_Batch_In_Lanes + : public ConvolutionTest { + public: + void RunTest() { + XlaBuilder builder(TestName()); + std::vector input_dims = {1, 4, 4, 512}; + std::vector filter_dims = {3, 3, 1, 512}; + Shape input_shape = ShapeUtil::MakeShapeWithType(input_dims); + Shape filter_shape = ShapeUtil::MakeShapeWithType(filter_dims); + { + auto input = Parameter(&builder, 0, input_shape, "input"); + auto filter = Parameter(&builder, 1, filter_shape, "filter"); + + // Tensorflow dimension numbers for 2D convolution. + 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=*/512); + } + + std::vector input_elems(ShapeUtil::ElementsIn(input_shape), + static_cast(1)); + auto input_r1 = LiteralUtil::CreateR1(input_elems); + auto input_r4 = input_r1.Reshape(input_dims).ConsumeValueOrDie(); + + std::vector filter_elems(ShapeUtil::ElementsIn(filter_shape), + static_cast(2)); + auto filter_r1 = LiteralUtil::CreateR1(filter_elems); + auto filter_r4 = filter_r1.Reshape(filter_dims).ConsumeValueOrDie(); + + std::vector output_elems(2048, static_cast(18)); + + auto expected_r1 = LiteralUtil::CreateR1(output_elems); + auto expected_r4 = expected_r1.Reshape({1, 2, 2, 512}).ConsumeValueOrDie(); + auto expected_r4_relaid = + expected_r4.Relayout(LayoutUtil::MakeLayout({0, 3, 2, 1})); + + auto input_literal = + client_->TransferToServer(input_r4).ConsumeValueOrDie(); + auto filter_literal = + client_->TransferToServer(filter_r4).ConsumeValueOrDie(); + + ComputeAndCompareLiteral(&builder, expected_r4_relaid, + {input_literal.get(), filter_literal.get()}, + error_spec_, &expected_r4_relaid.shape()); + } +}; + +TYPED_TEST_CASE( + Convolve2D_1x4x4x512_3x3x1x512_Depthwise_Valid_Output_Batch_In_Lanes, + TestTypes); +TYPED_TEST(Convolve2D_1x4x4x512_3x3x1x512_Depthwise_Valid_Output_Batch_In_Lanes, + Types) { + this->RunTest(); +} + +template +class Convolve2D_256x4x4x512_3x3x1x512_Depthwise_Input_Batch_in_Lanes + : public ConvolutionTest { + public: + void RunTest() { + XlaBuilder builder(TestName()); + std::vector input_dims = {256, 4, 4, 512}; + std::vector filter_dims = {3, 3, 1, 512}; + Shape input_shape = ShapeUtil::MakeShapeWithType(input_dims); + Shape filter_shape = ShapeUtil::MakeShapeWithType(filter_dims); + { + auto input = Parameter(&builder, 0, input_shape, "input"); + auto filter = Parameter(&builder, 1, filter_shape, "filter"); + + // Tensorflow dimension numbers for 2D convolution. + 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=*/512); + } + + std::vector input_elems(ShapeUtil::ElementsIn(input_shape), + static_cast(1)); + auto input_r1 = LiteralUtil::CreateR1(input_elems); + auto input_r4 = input_r1.Reshape(input_dims).ConsumeValueOrDie(); + auto input_r4_relaid = + input_r4.Relayout(LayoutUtil::MakeLayout({0, 3, 2, 1})); + + std::vector filter_elems(ShapeUtil::ElementsIn(filter_shape), + static_cast(2)); + auto filter_r1 = LiteralUtil::CreateR1(filter_elems); + auto filter_r4 = filter_r1.Reshape(filter_dims).ConsumeValueOrDie(); + + std::vector output_elems(2048 * 256, static_cast(18)); + + auto expected_r1 = LiteralUtil::CreateR1(output_elems); + auto expected_r4 = + expected_r1.Reshape({256, 2, 2, 512}).ConsumeValueOrDie(); + + auto input_literal = + client_->TransferToServer(input_r4_relaid).ConsumeValueOrDie(); + auto filter_literal = + client_->TransferToServer(filter_r4).ConsumeValueOrDie(); + + ComputeAndCompareLiteral(&builder, expected_r4, + {input_literal.get(), filter_literal.get()}, + error_spec_); + } +}; + +TYPED_TEST_CASE(Convolve2D_256x4x4x512_3x3x1x512_Depthwise_Input_Batch_in_Lanes, + TestTypes); +TYPED_TEST(Convolve2D_256x4x4x512_3x3x1x512_Depthwise_Input_Batch_in_Lanes, + Types) { + this->RunTest(); +} + +template +class Convolve2D_256x4x4x512_3x3x1x512_Depthwise_Both_Batch_in_Lanes + : public ConvolutionTest { + public: + void RunTest() { + XlaBuilder builder(TestName()); + std::vector input_dims = {256, 4, 4, 512}; + std::vector filter_dims = {3, 3, 1, 512}; + Shape input_shape = ShapeUtil::MakeShapeWithType(input_dims); + Shape filter_shape = ShapeUtil::MakeShapeWithType(filter_dims); + { + auto input = Parameter(&builder, 0, input_shape, "input"); + auto filter = Parameter(&builder, 1, filter_shape, "filter"); + + // Tensorflow dimension numbers for 2D convolution. + 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=*/512); + } + + std::vector input_elems(ShapeUtil::ElementsIn(input_shape), + static_cast(1)); + auto input_r1 = LiteralUtil::CreateR1(input_elems); + auto input_r4 = input_r1.Reshape(input_dims).ConsumeValueOrDie(); + auto input_r4_relaid = + input_r4.Relayout(LayoutUtil::MakeLayout({0, 3, 2, 1})); + + std::vector filter_elems(ShapeUtil::ElementsIn(filter_shape), + static_cast(2)); + auto filter_r1 = LiteralUtil::CreateR1(filter_elems); + auto filter_r4 = filter_r1.Reshape(filter_dims).ConsumeValueOrDie(); + + std::vector output_elems(2048 * 256, static_cast(18)); + + auto expected_r1 = LiteralUtil::CreateR1(output_elems); + auto expected_r4 = + expected_r1.Reshape({256, 2, 2, 512}).ConsumeValueOrDie(); + auto expected_r4_relaid = + expected_r4.Relayout(LayoutUtil::MakeLayout({0, 3, 2, 1})); + + auto input_literal = + client_->TransferToServer(input_r4_relaid).ConsumeValueOrDie(); + auto filter_literal = + client_->TransferToServer(filter_r4).ConsumeValueOrDie(); + + ComputeAndCompareLiteral(&builder, expected_r4_relaid, + {input_literal.get(), filter_literal.get()}, + error_spec_, &expected_r4_relaid.shape()); + } +}; + +TYPED_TEST_CASE(Convolve2D_256x4x4x512_3x3x1x512_Depthwise_Both_Batch_in_Lanes, + TestTypes); +TYPED_TEST(Convolve2D_256x4x4x512_3x3x1x512_Depthwise_Both_Batch_in_Lanes, + Types) { + this->RunTest(); +} + +template +class Convolve2D_1x4x4x5_3x3x1x5_Depthwise_Valid_Output_Batch_In_Lanes + : public ConvolutionTest { + public: + void RunTest() { + XlaBuilder builder(TestName()); + std::vector input_dims = {1, 4, 4, 5}; + std::vector filter_dims = {3, 3, 1, 5}; + Shape input_shape = ShapeUtil::MakeShapeWithType(input_dims); + Shape filter_shape = ShapeUtil::MakeShapeWithType(filter_dims); + { + auto input = Parameter(&builder, 0, input_shape, "input"); + auto filter = Parameter(&builder, 1, filter_shape, "filter"); + + // Tensorflow dimension numbers for 2D convolution. + 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=*/5); + } + + std::vector input_elems(ShapeUtil::ElementsIn(input_shape)); + iota_int_init_value(input_elems, 1); + auto input_r1 = LiteralUtil::CreateR1(input_elems); + auto input_r4 = input_r1.Reshape(input_dims).ConsumeValueOrDie(); + auto input_r4_relaid = + input_r4.Relayout(LayoutUtil::MakeLayout({0, 3, 2, 1})); + + std::vector filter_elems(ShapeUtil::ElementsIn(filter_shape)); + iota_int_init_value(filter_elems, 1); + auto filter_r1 = LiteralUtil::CreateR1(filter_elems); + auto filter_r4 = filter_r1.Reshape(filter_dims).ConsumeValueOrDie(); + + auto expected_r1 = LiteralUtil::CreateR1( + {static_cast(6864), static_cast(7296), static_cast(7746), + static_cast(8214), static_cast(8700), static_cast(7809), + static_cast(8286), static_cast(8781), static_cast(9294), + static_cast(9825), static_cast(10644), static_cast(11256), + static_cast(11886), static_cast(12534), static_cast(13200), + static_cast(11589), static_cast(12246), static_cast(12921), + static_cast(13614), static_cast(14325)}); + auto expected_r4 = expected_r1.Reshape({1, 2, 2, 5}).ConsumeValueOrDie(); + auto expected_r4_relaid = + expected_r4.Relayout(LayoutUtil::MakeLayout({0, 3, 2, 1})); + + auto input_literal = + client_->TransferToServer(input_r4_relaid).ConsumeValueOrDie(); + auto filter_literal = + client_->TransferToServer(filter_r4).ConsumeValueOrDie(); + + ComputeAndCompareLiteral(&builder, expected_r4_relaid, + {input_literal.get(), filter_literal.get()}, + error_spec_, &expected_r4_relaid.shape()); + } +}; + +TYPED_TEST_CASE( + Convolve2D_1x4x4x5_3x3x1x5_Depthwise_Valid_Output_Batch_In_Lanes, + TestTypes); +TYPED_TEST(Convolve2D_1x4x4x5_3x3x1x5_Depthwise_Valid_Output_Batch_In_Lanes, + Types) { + this->RunTest(); +} + +template +class Convolve2D_1x4x4x160_3x3x1x160_Depthwise_Valid : public ConvolutionTest { + public: + void RunTest() { + XlaBuilder builder(TestName()); + std::vector input_dims = {1, 4, 4, 160}; + std::vector filter_dims = {3, 3, 1, 160}; + Shape input_shape = ShapeUtil::MakeShapeWithType(input_dims); + Shape filter_shape = ShapeUtil::MakeShapeWithType(filter_dims); + { + auto input = Parameter(&builder, 0, input_shape, "input"); + auto filter = Parameter(&builder, 1, filter_shape, "filter"); + + // Tensorflow dimension numbers for 2D convolution. + 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=*/160); + } + + std::vector input_elems(ShapeUtil::ElementsIn(input_shape), + static_cast(1)); + auto input_r1 = LiteralUtil::CreateR1(input_elems); + auto input_r4 = input_r1.Reshape(input_dims).ConsumeValueOrDie(); + + std::vector filter_elems(ShapeUtil::ElementsIn(filter_shape), + static_cast(2)); + auto filter_r1 = LiteralUtil::CreateR1(filter_elems); + auto filter_r4 = filter_r1.Reshape(filter_dims).ConsumeValueOrDie(); + + std::vector output_elems(640, static_cast(18)); + + auto expected_r1 = LiteralUtil::CreateR1(output_elems); + auto expected_r4 = expected_r1.Reshape({1, 2, 2, 160}).ConsumeValueOrDie(); + + auto input_literal = + client_->TransferToServer(input_r4).ConsumeValueOrDie(); + auto filter_literal = + client_->TransferToServer(filter_r4).ConsumeValueOrDie(); + + ComputeAndCompareLiteral(&builder, expected_r4, + {input_literal.get(), filter_literal.get()}, + error_spec_); + } +}; + +TYPED_TEST_CASE(Convolve2D_1x4x4x160_3x3x1x160_Depthwise_Valid, TestTypes); +TYPED_TEST(Convolve2D_1x4x4x160_3x3x1x160_Depthwise_Valid, Types) { + this->RunTest(); +} + +template +class Convolve2D_1x4x4x160_3x3x1x160_Depthwise_Input_Batch_In_Lanes + : public ConvolutionTest { + public: + void RunTest() { + XlaBuilder builder(TestName()); + std::vector input_dims = {1, 4, 4, 160}; + std::vector filter_dims = {3, 3, 1, 160}; + Shape input_shape = ShapeUtil::MakeShapeWithType(input_dims); + Shape filter_shape = ShapeUtil::MakeShapeWithType(filter_dims); + { + auto input = Parameter(&builder, 0, input_shape, "input"); + auto filter = Parameter(&builder, 1, filter_shape, "filter"); + + // Tensorflow dimension numbers for 2D convolution. + 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=*/160); + } + + std::vector input_elems(ShapeUtil::ElementsIn(input_shape), + static_cast(1)); + auto input_r1 = LiteralUtil::CreateR1(input_elems); + auto input_r4 = input_r1.Reshape(input_dims).ConsumeValueOrDie(); + auto input_r4_relaid = + input_r4.Relayout(LayoutUtil::MakeLayout({0, 3, 2, 1})); + + std::vector filter_elems(ShapeUtil::ElementsIn(filter_shape), + static_cast(2)); + auto filter_r1 = LiteralUtil::CreateR1(filter_elems); + auto filter_r4 = filter_r1.Reshape(filter_dims).ConsumeValueOrDie(); + + std::vector output_elems(640, static_cast(18)); + + auto expected_r1 = LiteralUtil::CreateR1(output_elems); + auto expected_r4 = expected_r1.Reshape({1, 2, 2, 160}).ConsumeValueOrDie(); + auto expected_r4_relaid = + expected_r4.Relayout(LayoutUtil::MakeLayout({3, 0, 2, 1})); + + auto input_literal = + client_->TransferToServer(input_r4_relaid).ConsumeValueOrDie(); + auto filter_literal = + client_->TransferToServer(filter_r4).ConsumeValueOrDie(); + + ComputeAndCompareLiteral(&builder, expected_r4_relaid, + {input_literal.get(), filter_literal.get()}, + error_spec_, &expected_r4_relaid.shape()); + } +}; + +TYPED_TEST_CASE(Convolve2D_1x4x4x160_3x3x1x160_Depthwise_Input_Batch_In_Lanes, + TestTypes); +TYPED_TEST(Convolve2D_1x4x4x160_3x3x1x160_Depthwise_Input_Batch_In_Lanes, + Types) { + this->RunTest(); +} + +template +class Convolve2D_1x4x4x160_3x3x1x160_Dephtwise_Both_Batch_In_Lanes + : public ConvolutionTest { + public: + void RunTest() { + XlaBuilder builder(TestName()); + std::vector input_dims = {1, 4, 4, 160}; + std::vector filter_dims = {3, 3, 1, 160}; + Shape input_shape = ShapeUtil::MakeShapeWithType(input_dims); + Shape filter_shape = ShapeUtil::MakeShapeWithType(filter_dims); + { + auto input = Parameter(&builder, 0, input_shape, "input"); + auto filter = Parameter(&builder, 1, filter_shape, "filter"); + + // Tensorflow dimension numbers for 2D convolution. + 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=*/160); + } + + std::vector input_elems(ShapeUtil::ElementsIn(input_shape), + static_cast(1)); + auto input_r1 = LiteralUtil::CreateR1(input_elems); + auto input_r4 = input_r1.Reshape(input_dims).ConsumeValueOrDie(); + auto input_r4_relaid = + input_r4.Relayout(LayoutUtil::MakeLayout({0, 3, 2, 1})); + + std::vector filter_elems(ShapeUtil::ElementsIn(filter_shape), + static_cast(2)); + auto filter_r1 = LiteralUtil::CreateR1(filter_elems); + auto filter_r4 = filter_r1.Reshape(filter_dims).ConsumeValueOrDie(); + + std::vector output_elems(640, static_cast(18)); + + auto expected_r1 = LiteralUtil::CreateR1(output_elems); + auto expected_r4 = expected_r1.Reshape({1, 2, 2, 160}).ConsumeValueOrDie(); + auto expected_r4_relaid = + expected_r4.Relayout(LayoutUtil::MakeLayout({0, 3, 2, 1})); + + auto input_literal = + client_->TransferToServer(input_r4_relaid).ConsumeValueOrDie(); + auto filter_literal = + client_->TransferToServer(filter_r4).ConsumeValueOrDie(); + + ComputeAndCompareLiteral(&builder, expected_r4_relaid, + {input_literal.get(), filter_literal.get()}, + error_spec_, &expected_r4_relaid.shape()); + } +}; + +TYPED_TEST_CASE(Convolve2D_1x4x4x160_3x3x1x160_Dephtwise_Both_Batch_In_Lanes, + TestTypes); +TYPED_TEST(Convolve2D_1x4x4x160_3x3x1x160_Dephtwise_Both_Batch_In_Lanes, + Types) { + this->RunTest(); +} + +template +class Convolve2D_1x4x4x1024_3x3x1x1024_Depthwise_Valid + : public ConvolutionTest { + public: + void RunTest() { + XlaBuilder builder(TestName()); + std::vector input_dims = {1, 4, 4, 1024}; + std::vector filter_dims = {3, 3, 1, 1024}; + Shape input_shape = ShapeUtil::MakeShapeWithType(input_dims); + Shape filter_shape = ShapeUtil::MakeShapeWithType(filter_dims); + { + auto input = Parameter(&builder, 0, input_shape, "input"); + auto filter = Parameter(&builder, 1, filter_shape, "filter"); + + // Tensorflow dimension numbers for 2D convolution. + 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=*/1024); + } + + std::vector input_elems(ShapeUtil::ElementsIn(input_shape), + static_cast(1)); + auto input_r1 = LiteralUtil::CreateR1(input_elems); + auto input_r4 = input_r1.Reshape(input_dims).ConsumeValueOrDie(); + + std::vector filter_elems(ShapeUtil::ElementsIn(filter_shape), + static_cast(2)); + auto filter_r1 = LiteralUtil::CreateR1(filter_elems); + auto filter_r4 = filter_r1.Reshape(filter_dims).ConsumeValueOrDie(); + + std::vector output_elems(4096, static_cast(18)); + + auto expected_r1 = LiteralUtil::CreateR1(output_elems); + auto expected_r4 = expected_r1.Reshape({1, 2, 2, 1024}).ConsumeValueOrDie(); + + auto input_literal = + client_->TransferToServer(input_r4).ConsumeValueOrDie(); + auto filter_literal = + client_->TransferToServer(filter_r4).ConsumeValueOrDie(); + + ComputeAndCompareLiteral(&builder, expected_r4, + {input_literal.get(), filter_literal.get()}, + error_spec_); + } +}; + +TYPED_TEST_CASE(Convolve2D_1x4x4x1024_3x3x1x1024_Depthwise_Valid, TestTypes); +TYPED_TEST(Convolve2D_1x4x4x1024_3x3x1x1024_Depthwise_Valid, Types) { + this->RunTest(); +} + +template +class Convolve2D_1x2x2x6_2x2x2x12_Grouped_Valid : public ConvolutionTest { public: void RunTest() { XlaBuilder builder(TestName()); @@ -625,7 +1310,200 @@ class Convolve2D_1x2x2x6_2x2x1x12_Grouped_Valid : public ConvolutionTest { dnums.set_kernel_output_feature_dimension(3); ConvWithGeneralDimensions(input, filter, {1, 1}, Padding::kValid, dnums, - /*feature_group_count=*/3); + /*feature_group_count=*/3); + } + + std::vector input_elems(ShapeUtil::ElementsIn(input_shape)); + iota_int_init_value(input_elems, 1); + auto input_r1 = LiteralUtil::CreateR1(input_elems); + auto input_r4 = input_r1.Reshape(input_dims).ConsumeValueOrDie(); + + std::vector filter_elems(ShapeUtil::ElementsIn(filter_shape)); + iota_int_init_value(filter_elems, 1); + auto filter_r1 = LiteralUtil::CreateR1(filter_elems); + auto filter_r4 = filter_r1.Reshape(filter_dims).ConsumeValueOrDie(); + + auto expected_r1 = LiteralUtil::CreateR1( + {static_cast(5076), static_cast(5160), static_cast(5244), + static_cast(5328), static_cast(6164), static_cast(6264), + static_cast(6364), static_cast(6464), static_cast(7380), + static_cast(7496), static_cast(7612), static_cast(7728)}); + auto expected_r4 = expected_r1.Reshape({1, 1, 1, 12}).ConsumeValueOrDie(); + + auto input_literal = + client_->TransferToServer(input_r4).ConsumeValueOrDie(); + auto filter_literal = + client_->TransferToServer(filter_r4).ConsumeValueOrDie(); + + ComputeAndCompareLiteral(&builder, expected_r4, + {input_literal.get(), filter_literal.get()}, + error_spec_); + } +}; + +TYPED_TEST_CASE(Convolve2D_1x2x2x6_2x2x2x12_Grouped_Valid, TestTypes); +TYPED_TEST(Convolve2D_1x2x2x6_2x2x2x12_Grouped_Valid, Types) { + this->RunTest(); +} + +template +class Convolve2D_1x2x2x1024_2x2x128x512_Grouped_Valid : public ConvolutionTest { + public: + void RunTest() { + XlaBuilder builder(TestName()); + std::vector input_dims = {1, 2, 2, 1024}; + std::vector filter_dims = {2, 2, 128, 512}; + Shape input_shape = ShapeUtil::MakeShapeWithType(input_dims); + Shape filter_shape = ShapeUtil::MakeShapeWithType(filter_dims); + { + auto input = Parameter(&builder, 0, input_shape, "input"); + auto filter = Parameter(&builder, 1, filter_shape, "filter"); + + // Tensorflow dimension numbers for 2D convolution. + 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=*/8); + } + + std::vector input_elems(ShapeUtil::ElementsIn(input_shape), + static_cast(1)); + + auto input_r1 = LiteralUtil::CreateR1(input_elems); + auto input_r4 = input_r1.Reshape(input_dims).ConsumeValueOrDie(); + + std::vector filter_elems(ShapeUtil::ElementsIn(filter_shape), + static_cast(2)); + + auto filter_r1 = LiteralUtil::CreateR1(filter_elems); + auto filter_r4 = filter_r1.Reshape(filter_dims).ConsumeValueOrDie(); + + std::vector output_elems(512, static_cast(1024)); + auto expected_r1 = LiteralUtil::CreateR1(output_elems); + auto expected_r4 = expected_r1.Reshape({1, 1, 1, 512}).ConsumeValueOrDie(); + + auto input_literal = + client_->TransferToServer(input_r4).ConsumeValueOrDie(); + auto filter_literal = + client_->TransferToServer(filter_r4).ConsumeValueOrDie(); + + ComputeAndCompareLiteral(&builder, expected_r4, + {input_literal.get(), filter_literal.get()}, + error_spec_); + } +}; + +TYPED_TEST_CASE(Convolve2D_1x2x2x1024_2x2x128x512_Grouped_Valid, TestTypes); +TYPED_TEST(Convolve2D_1x2x2x1024_2x2x128x512_Grouped_Valid, Types) { + this->RunTest(); +} + +template +class Convolve2D_1x2x2x1024_2x2x128x8_Grouped_Valid : public ConvolutionTest { + public: + void RunTest() { + XlaBuilder builder(TestName()); + std::vector input_dims = {1, 2, 2, 1024}; + std::vector filter_dims = {2, 2, 128, 8}; + Shape input_shape = ShapeUtil::MakeShapeWithType(input_dims); + Shape filter_shape = ShapeUtil::MakeShapeWithType(filter_dims); + { + auto input = Parameter(&builder, 0, input_shape, "input"); + auto filter = Parameter(&builder, 1, filter_shape, "filter"); + + // Tensorflow dimension numbers for 2D convolution. + 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=*/8); + } + + std::vector input_elems(ShapeUtil::ElementsIn(input_shape), + static_cast(1)); + + auto input_r1 = LiteralUtil::CreateR1(input_elems); + auto input_r4 = input_r1.Reshape(input_dims).ConsumeValueOrDie(); + + std::vector filter_elems(ShapeUtil::ElementsIn(filter_shape), + static_cast(2)); + + auto filter_r1 = LiteralUtil::CreateR1(filter_elems); + auto filter_r4 = filter_r1.Reshape(filter_dims).ConsumeValueOrDie(); + + std::vector output_elems(8, static_cast(1024)); + auto expected_r1 = LiteralUtil::CreateR1(output_elems); + auto expected_r4 = expected_r1.Reshape({1, 1, 1, 8}).ConsumeValueOrDie(); + + auto input_literal = + client_->TransferToServer(input_r4).ConsumeValueOrDie(); + auto filter_literal = + client_->TransferToServer(filter_r4).ConsumeValueOrDie(); + + ComputeAndCompareLiteral(&builder, expected_r4, + {input_literal.get(), filter_literal.get()}, + error_spec_); + } +}; + +TYPED_TEST_CASE(Convolve2D_1x2x2x1024_2x2x128x8_Grouped_Valid, TestTypes); +TYPED_TEST(Convolve2D_1x2x2x1024_2x2x128x8_Grouped_Valid, Types) { + this->RunTest(); +} + +template +class Convolve2D_1x2x2x12_2x2x3x4_Grouped_Valid : public ConvolutionTest { + public: + void RunTest() { + XlaBuilder builder(TestName()); + std::vector input_dims = {1, 2, 2, 12}; + std::vector filter_dims = {2, 2, 3, 4}; + Shape input_shape = ShapeUtil::MakeShapeWithType(input_dims); + Shape filter_shape = ShapeUtil::MakeShapeWithType(filter_dims); + { + auto input = Parameter(&builder, 0, input_shape, "input"); + auto filter = Parameter(&builder, 1, filter_shape, "filter"); + + // Tensorflow dimension numbers for 2D convolution. + 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=*/4); } std::vector input_elems(ShapeUtil::ElementsIn(input_shape)); @@ -638,12 +1516,140 @@ class Convolve2D_1x2x2x6_2x2x1x12_Grouped_Valid : public ConvolutionTest { auto filter_r1 = LiteralUtil::CreateR1(filter_elems); auto filter_r4 = filter_r1.Reshape(filter_dims).ConsumeValueOrDie(); + auto expected_r1 = + LiteralUtil::CreateR1({static_cast(7712), static_cast(8816), + static_cast(9992), static_cast(11240)}); + auto expected_r4 = expected_r1.Reshape({1, 1, 1, 4}).ConsumeValueOrDie(); + + auto input_literal = + client_->TransferToServer(input_r4).ConsumeValueOrDie(); + auto filter_literal = + client_->TransferToServer(filter_r4).ConsumeValueOrDie(); + + ComputeAndCompareLiteral(&builder, expected_r4, + {input_literal.get(), filter_literal.get()}, + error_spec_); + } +}; + +TYPED_TEST_CASE(Convolve2D_1x2x2x12_2x2x3x4_Grouped_Valid, TestTypes); +TYPED_TEST(Convolve2D_1x2x2x12_2x2x3x4_Grouped_Valid, Types) { + this->RunTest(); +} + +template +class Convolve2D_1x2x2x12_2x2x3x4_Grouped_Valid_Filter_OF_In_Sublanes + : public ConvolutionTest { + public: + void RunTest() { + XlaBuilder builder(TestName()); + std::vector input_dims = {1, 2, 2, 12}; + std::vector filter_dims = {2, 2, 4, 3}; + Shape input_shape = ShapeUtil::MakeShapeWithType(input_dims); + Shape filter_shape = ShapeUtil::MakeShapeWithType(filter_dims); + { + auto input = Parameter(&builder, 0, input_shape, "input"); + auto filter = Parameter(&builder, 1, filter_shape, "filter"); + + // Tensorflow dimension numbers for 2D convolution. + 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(3); + dnums.set_kernel_output_feature_dimension(2); + + ConvWithGeneralDimensions(input, filter, {1, 1}, Padding::kValid, dnums, + /*feature_group_count=*/4); + } + + std::vector input_elems(ShapeUtil::ElementsIn(input_shape)); + iota_int_init_value(input_elems, 1); + auto input_r1 = LiteralUtil::CreateR1(input_elems); + auto input_r4 = input_r1.Reshape(input_dims).ConsumeValueOrDie(); + + std::vector filter_elems(ShapeUtil::ElementsIn(filter_shape)); + iota_int_init_value(filter_elems, 1); + auto filter_r1 = LiteralUtil::CreateR1(filter_elems); + auto filter_r4 = filter_r1.Reshape(filter_dims).ConsumeValueOrDie(); + auto filter_r4_relaid = + filter_r4.Relayout(LayoutUtil::MakeLayout({3, 2, 1, 0})); auto expected_r1 = LiteralUtil::CreateR1( - {static_cast(5076), static_cast(5160), static_cast(5244), - static_cast(5328), static_cast(6164), static_cast(6264), - static_cast(6364), static_cast(6464), static_cast(7380), - static_cast(7496), static_cast(7612), static_cast(7728)}); - auto expected_r4 = expected_r1.Reshape({1, 1, 1, 12}).ConsumeValueOrDie(); + {static_cast(6968), static_cast(8516), static_cast(10280), + static_cast(12260)}); + auto expected_r4 = expected_r1.Reshape({1, 1, 1, 4}).ConsumeValueOrDie(); + + auto input_literal = + client_->TransferToServer(input_r4).ConsumeValueOrDie(); + auto filter_literal = + client_->TransferToServer(filter_r4_relaid).ConsumeValueOrDie(); + + ComputeAndCompareLiteral(&builder, expected_r4, + {input_literal.get(), filter_literal.get()}, + error_spec_); + } +}; + +TYPED_TEST_CASE(Convolve2D_1x2x2x12_2x2x3x4_Grouped_Valid_Filter_OF_In_Sublanes, + TestTypes); +TYPED_TEST(Convolve2D_1x2x2x12_2x2x3x4_Grouped_Valid_Filter_OF_In_Sublanes, + Types) { + this->RunTest(); +} + +template +class Convolve2D_1x1x1x12_1x1x3x4_Grouped_Valid : public ConvolutionTest { + public: + void RunTest() { + XlaBuilder builder(TestName()); + std::vector input_dims = {1, 1, 1, 12}; + std::vector filter_dims = {1, 1, 3, 4}; + Shape input_shape = ShapeUtil::MakeShapeWithType(input_dims); + Shape filter_shape = ShapeUtil::MakeShapeWithType(filter_dims); + { + auto input = Parameter(&builder, 0, input_shape, "input"); + auto filter = Parameter(&builder, 1, filter_shape, "filter"); + + // Tensorflow dimension numbers for 2D convolution. + 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=*/4); + } + + std::vector input_elems(ShapeUtil::ElementsIn(input_shape)); + iota_int_init_value(input_elems, 1); + auto input_r1 = LiteralUtil::CreateR1(input_elems); + auto input_r4 = input_r1.Reshape(input_dims).ConsumeValueOrDie(); + + std::vector filter_elems(ShapeUtil::ElementsIn(filter_shape)); + iota_int_init_value(filter_elems, 1); + auto filter_r1 = LiteralUtil::CreateR1(filter_elems); + auto filter_r4 = filter_r1.Reshape(filter_dims).ConsumeValueOrDie(); + + auto expected_r1 = + LiteralUtil::CreateR1({static_cast(38), static_cast(98), + static_cast(176), static_cast(272)}); + auto expected_r4 = expected_r1.Reshape({1, 1, 1, 4}).ConsumeValueOrDie(); auto input_literal = client_->TransferToServer(input_r4).ConsumeValueOrDie(); @@ -656,8 +1662,8 @@ class Convolve2D_1x2x2x6_2x2x1x12_Grouped_Valid : public ConvolutionTest { } }; -TYPED_TEST_CASE(Convolve2D_1x2x2x6_2x2x1x12_Grouped_Valid, TestTypes); -TYPED_TEST(Convolve2D_1x2x2x6_2x2x1x12_Grouped_Valid, Types) { +TYPED_TEST_CASE(Convolve2D_1x1x1x12_1x1x3x4_Grouped_Valid, TestTypes); +TYPED_TEST(Convolve2D_1x1x1x12_1x1x3x4_Grouped_Valid, Types) { this->RunTest(); } @@ -951,6 +1957,18 @@ ENTRY Test { EXPECT_TRUE(RunAndCompare(kHlo, ErrorSpec{0.001})); } +XLA_TEST_F(ConvolutionHloTest, DISABLED_ON_CPU(ConvolveF32ForwardReversed)) { + constexpr char kHlo[] = R"( +HloModule TestModule + +ENTRY Test { + %arg0 = f32[3,56,56,16] parameter(0) + %arg1 = f32[3,3,3,32] parameter(1) + ROOT %conv = f32[54,54,16,32] convolution(%arg0, %arg1), window={size=3x3 rhs_reversal=1x1}, dim_labels=f01b_i01o->01bf +})"; + EXPECT_TRUE(RunAndCompare(kHlo, ErrorSpec{0.001})); +} + XLA_TEST_F(ConvolutionHloTest, DISABLED_ON_CPU(ConvolveF64BackwardFilter)) { constexpr char kHlo[] = R"( HloModule TestModule diff --git a/tensorflow/compiler/xla/tests/copy_test.cc b/tensorflow/compiler/xla/tests/copy_test.cc index 1407e68d9a336b6bb1c960711015430f872aa912..df005a67097bb8aaf070c57d1c51acd1909fee12 100644 --- a/tensorflow/compiler/xla/tests/copy_test.cc +++ b/tensorflow/compiler/xla/tests/copy_test.cc @@ -45,7 +45,7 @@ class CopyOpTest : public HloTestBase { builder.AddInstruction(HloInstruction::CreateUnary( constant->shape(), HloOpcode::kCopy, constant)); auto computation = builder.Build(); - auto module = CreateNewModule(); + auto module = CreateNewUnverifiedModule(); module->AddEntryComputation(std::move(computation)); Literal result = ExecuteAndTransfer(std::move(module), {}); @@ -98,7 +98,7 @@ XLA_TEST_F(CopyOpTest, CopyParameterScalar) { auto computation = builder.Build(); - auto module = CreateNewModule(); + auto module = CreateNewUnverifiedModule(); module->AddEntryComputation(std::move(computation)); Literal result = ExecuteAndTransfer(std::move(module), {&literal}); @@ -119,7 +119,7 @@ XLA_TEST_F(CopyOpTest, CopyConstantR2Twice) { auto computation = builder.Build(); - auto module = CreateNewModule(); + auto module = CreateNewUnverifiedModule(); module->AddEntryComputation(std::move(computation)); Literal result = ExecuteAndTransfer(std::move(module), {}); LiteralTestUtil::ExpectR2Near({{1.0, 2.0}, {3.0, 4.0}}, result, @@ -133,7 +133,9 @@ XLA_TEST_F(CopyOpTest, CopyConstantR2DifferentLayouts) { // Reverse the minor-to-major order of the literal. Layout* literal_layout = literal.mutable_shape_do_not_use()->mutable_layout(); ASSERT_EQ(2, literal_layout->minor_to_major_size()); - literal_layout->mutable_minor_to_major()->SwapElements(0, 1); + // Swap the first and second elements. + *literal_layout->mutable_minor_to_major() = { + literal_layout->minor_to_major(1), literal_layout->minor_to_major(0)}; HloInstruction* constant = builder.AddInstruction( HloInstruction::CreateConstant(std::move(literal))); @@ -143,7 +145,7 @@ XLA_TEST_F(CopyOpTest, CopyConstantR2DifferentLayouts) { std::unique_ptr computation = builder.Build(); - auto module = CreateNewModule(); + auto module = CreateNewUnverifiedModule(); module->AddEntryComputation(std::move(computation)); Literal result = ExecuteAndTransfer(std::move(module), {}); @@ -175,7 +177,7 @@ void CopyOpTest::TestCopyConstantLayout021(size_t n1, size_t n2, size_t n3) { std::unique_ptr computation = builder.Build(); - auto module = CreateNewModule(); + auto module = CreateNewUnverifiedModule(); module->AddEntryComputation(std::move(computation)); ForceResultLayout(module.get(), LayoutUtil::MakeLayout({1, 2, 0})); Literal result = ExecuteAndTransfer(std::move(module), {}); @@ -209,7 +211,7 @@ void CopyOpTest::TestCopyConstantLayoutR4(size_t n1, size_t n2, size_t n3, std::unique_ptr computation = builder.Build(); - auto module = CreateNewModule(); + auto module = CreateNewUnverifiedModule(); module->AddEntryComputation(std::move(computation)); ForceResultLayout(module.get(), LayoutUtil::MakeLayout(permutation)); Literal result = ExecuteAndTransfer(std::move(module), {}); diff --git a/tensorflow/compiler/xla/tests/cross_replica_sum_test.cc b/tensorflow/compiler/xla/tests/cross_replica_sum_test.cc deleted file mode 100644 index 410732c07b7b6d3ece33ab11f4778241dc53ca50..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/tests/cross_replica_sum_test.cc +++ /dev/null @@ -1,102 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/compiler/xla/literal.h" -#include "tensorflow/compiler/xla/service/hlo_parser.h" -#include "tensorflow/compiler/xla/shape_util.h" -#include "tensorflow/compiler/xla/test.h" -#include "tensorflow/compiler/xla/test_helpers.h" -#include "tensorflow/compiler/xla/tests/hlo_test_base.h" -#include "tensorflow/compiler/xla/tests/test_macros.h" - -namespace xla { -namespace { - -class TrivialCrossReplicaSumTest : public HloTestBase {}; - -// Currently the CPU and GPU backends only support CrossReplicaSum with one -// replica. But we can at least check this. - -XLA_TEST_F(TrivialCrossReplicaSumTest, OneOperand) { - const char* module_str = R"( - HloModule test - - add { - x = f32[] parameter(0) - y = f32[] parameter(1) - add = f32[] add(x, y) - } - - ENTRY test_computation { - p = f32[3] parameter(0) - ROOT crs = f32[3] cross-replica-sum(p), to_apply=add - })"; - auto module = - ParseHloString(module_str, GetModuleConfigForTest()).ValueOrDie(); - auto literal = LiteralUtil::CreateR1({1, 2, 3}); - EXPECT_EQ(literal, ExecuteAndTransfer(std::move(module), {&literal})); -} - -XLA_TEST_F(TrivialCrossReplicaSumTest, MultipleOperands) { - const char* module_str = R"( - HloModule test - - add { - x = f32[] parameter(0) - y = f32[] parameter(1) - add = f32[] add(x, y) - } - - ENTRY test_computation { - p0 = f32[3] parameter(0) - p1 = f32[2] parameter(1) - ROOT crs = (f32[3], f32[2]) cross-replica-sum(p0, p1), to_apply=add - })"; - auto module = - ParseHloString(module_str, GetModuleConfigForTest()).ValueOrDie(); - auto literal0 = LiteralUtil::CreateR1({1, 2, 3}); - auto literal1 = LiteralUtil::CreateR1({10, 20}); - EXPECT_EQ(LiteralUtil::MakeTuple({&literal0, &literal1}), - ExecuteAndTransfer(std::move(module), {&literal0, &literal1})); -} - -// On the GPU backend, constants get special handling. Someone might pass a -// constant to CRS to e.g. count the number of replicas -- we need to make sure -// it works. -XLA_TEST_F(TrivialCrossReplicaSumTest, ConstantOperand) { - const char* module_str = R"( - HloModule test - - add { - x = f32[] parameter(0) - y = f32[] parameter(1) - add = f32[] add(x, y) - } - - ENTRY test_computation { - p0 = f32[3] parameter(0) - p1 = f32[2] constant({10, 20}) - ROOT crs = (f32[3], f32[2]) cross-replica-sum(p0, p1), to_apply=add - })"; - auto module = - ParseHloString(module_str, GetModuleConfigForTest()).ValueOrDie(); - auto literal0 = LiteralUtil::CreateR1({1, 2, 3}); - auto literal1 = LiteralUtil::CreateR1({10, 20}); - EXPECT_EQ(LiteralUtil::MakeTuple({&literal0, &literal1}), - ExecuteAndTransfer(std::move(module), {&literal0})); -} - -} // namespace -} // namespace xla diff --git a/tensorflow/compiler/xla/tests/custom_call_test.cc b/tensorflow/compiler/xla/tests/custom_call_test.cc index 001490c6a8c568656437465054ee4db40d0d8dee..cad43d1b5547d74701760fa623e50466fc15c263 100644 --- a/tensorflow/compiler/xla/tests/custom_call_test.cc +++ b/tensorflow/compiler/xla/tests/custom_call_test.cc @@ -54,11 +54,20 @@ void Add1ToValues(float* out, float** in) { out[2] = array[2] + 1; out[3] = array[3] + 1; } + +void F32TupleSwap(float** out, float** in) { + TF_ANNOTATE_MEMORY_IS_INITIALIZED(in[0], sizeof(float)); + TF_ANNOTATE_MEMORY_IS_INITIALIZED(in[1], sizeof(float)); + *out[0] = *in[1]; + *out[1] = *in[0]; +} + } // namespace REGISTER_CUSTOM_CALL_TARGET(R0F32Add2); REGISTER_CUSTOM_CALL_TARGET(R2F32ReduceSum); REGISTER_CUSTOM_CALL_TARGET(Add1ToValues); +REGISTER_CUSTOM_CALL_TARGET(F32TupleSwap); namespace xla { namespace { @@ -69,8 +78,8 @@ class CustomCallTest : public HloTestBase { Shape r2f32_ = ShapeUtil::MakeShape(F32, {2, 2}); }; -XLA_TEST_F(CustomCallTest, DISABLED_ON_GPU(CustomCallR0F32Add2)) { - auto module = CreateNewModule(); +XLA_TEST_F(CustomCallTest, CustomCallR0F32Add2) { + auto module = CreateNewUnverifiedModule(); auto builder = HloComputation::Builder(TestName()); auto constant = builder.AddInstruction( @@ -84,8 +93,8 @@ XLA_TEST_F(CustomCallTest, DISABLED_ON_GPU(CustomCallR0F32Add2)) { LiteralTestUtil::ExpectR0Near(44.0f, result, error_spec_); } -XLA_TEST_F(CustomCallTest, DISABLED_ON_GPU(CustomCallR2F32Reduce)) { - auto module = CreateNewModule(); +XLA_TEST_F(CustomCallTest, CustomCallR2F32Reduce) { + auto module = CreateNewUnverifiedModule(); auto builder = HloComputation::Builder(TestName()); Array2D array(2, 2); @@ -105,8 +114,8 @@ XLA_TEST_F(CustomCallTest, DISABLED_ON_GPU(CustomCallR2F32Reduce)) { LiteralTestUtil::ExpectR0Near(10.0f, result, error_spec_); } -XLA_TEST_F(CustomCallTest, DISABLED_ON_GPU(UsedInOtherComputations)) { - auto module = CreateNewModule(); +XLA_TEST_F(CustomCallTest, UsedInOtherComputations) { + auto module = CreateNewUnverifiedModule(); auto b = HloComputation::Builder(TestName()); auto input = b.AddInstruction( @@ -129,8 +138,8 @@ XLA_TEST_F(CustomCallTest, DISABLED_ON_GPU(UsedInOtherComputations)) { Array3D{{{2, 3}, {4, 5}}, {{3, 4}, {5, 6}}}, result); } -XLA_TEST_F(CustomCallTest, DISABLED_ON_GPU(InputAndOutputLayoutDiffer)) { - auto module = CreateNewModule(); +XLA_TEST_F(CustomCallTest, InputAndOutputLayoutDiffer) { + auto module = CreateNewUnverifiedModule(); auto b = HloComputation::Builder(TestName()); auto input = @@ -151,11 +160,11 @@ XLA_TEST_F(CustomCallTest, DISABLED_ON_GPU(InputAndOutputLayoutDiffer)) { LiteralTestUtil::ExpectR2Equal({{2.f, 4.f}, {3.f, 5.f}}, result); } -XLA_TEST_F(CustomCallTest, DISABLED_ON_GPU(LayoutConstrained)) { +XLA_TEST_F(CustomCallTest, LayoutConstrained) { // The argument and result of the computation are set to different layouts, // but the custom call is layout constrained to a fixed operand and result // layout, so the correct result should be produced. - auto module = CreateNewModule(); + auto module = CreateNewUnverifiedModule(); auto b = HloComputation::Builder(TestName()); auto input = @@ -176,6 +185,26 @@ XLA_TEST_F(CustomCallTest, DISABLED_ON_GPU(LayoutConstrained)) { LiteralTestUtil::ExpectR2Equal({{2.f, 3.f}, {4.f, 5.f}}, result); } +XLA_TEST_F(CustomCallTest, TupleOutput) { + const char* kModuleStr = R"( + HloModule m + test { + p0 = f32[] parameter(0) + p1 = f32[] parameter(1) + ROOT %custom-call = (f32[], f32[]) custom-call(f32[] %p0, f32[] %p1), custom_call_target="F32TupleSwap", operand_layout_constraints={f32[], f32[]} + } + )"; + TF_ASSERT_OK_AND_ASSIGN(auto module, + ParseAndReturnVerifiedModule(kModuleStr)); + + Literal arg0 = LiteralUtil::CreateR0(7.f); + Literal arg1 = LiteralUtil::CreateR0(42.f); + + Literal expected = LiteralUtil::MakeTuple({&arg1, &arg0}); + Literal result = ExecuteAndTransfer(std::move(module), {&arg0, &arg1}); + EXPECT_EQ(result, expected); +} + class CustomCallClientAPITest : public ClientLibraryTestBase {}; // When using the client API, CustomCall targets can't begin with '$' -- these diff --git a/tensorflow/compiler/xla/tests/dot_operation_test.cc b/tensorflow/compiler/xla/tests/dot_operation_test.cc index 6c0847a875798870b4362a99ac2ab65d99f9f3e6..f740f4815810727890583405b2244fceaec0bd3f 100644 --- a/tensorflow/compiler/xla/tests/dot_operation_test.cc +++ b/tensorflow/compiler/xla/tests/dot_operation_test.cc @@ -19,18 +19,19 @@ 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" #include "tensorflow/core/platform/test.h" #include "tensorflow/core/platform/types.h" -#include "tensorflow/core/util/command_line_flags.h" namespace xla { namespace { @@ -637,6 +638,76 @@ XLA_TYPED_TEST(DotOperationTest_F16F32F64CF64, GeneralMatMul) { {x_data.get(), y_data.get()}, this->error_spec_); } +#ifndef XLA_TEST_BACKEND_CPU +// TODO(b/74459949): failed on CPU on 2018-10-29. +XLA_TYPED_TEST(DotOperationTest_F16F32F64CF64, GeneralMatMulR3LhsR2Rhs) { + using T = TypeParam; + + XlaBuilder builder(this->TestName()); + auto x = + Parameter(&builder, 0, ShapeUtil::MakeShapeWithType({2, 2, 2}), "x"); + auto y = Parameter(&builder, 1, ShapeUtil::MakeShapeWithType({2, 2}), "y"); + + DotDimensionNumbers dnums; + dnums.add_lhs_contracting_dimensions(1); + dnums.add_rhs_contracting_dimensions(1); + dnums.add_lhs_batch_dimensions(0); + dnums.add_rhs_batch_dimensions(0); + + DotGeneral(x, y, dnums); + + auto x_data = + this->client_ + ->TransferToServer(LiteralUtil::CreateR3FromArray3D( + {{{1.0f, 2.0f}, {3.0f, 4.0f}}, {{5.0f, 6.0f}, {7.0f, 8.0f}}})) + .ConsumeValueOrDie(); + + auto y_data = this->client_ + ->TransferToServer(LiteralUtil::CreateR2FromArray2D( + {{1.0f, 0.0f}, {0.0f, 1.0f}})) + .ConsumeValueOrDie(); + + this->template ComputeAndCompareR2( + &builder, + /*expected=*/{{1.0f, 2.0f}, {7.0f, 8.0f}}, {x_data.get(), y_data.get()}, + this->error_spec_); +} + +// TODO(b/74459949): failed on CPU on 2018-10-29. +XLA_TYPED_TEST(DotOperationTest_F16F32F64CF64, GeneralMatMulR2LhsR3Rhs) { + using T = TypeParam; + + XlaBuilder builder(this->TestName()); + auto x = Parameter(&builder, 0, ShapeUtil::MakeShapeWithType({2, 2}), "x"); + auto y = + Parameter(&builder, 1, ShapeUtil::MakeShapeWithType({2, 2, 2}), "y"); + + DotDimensionNumbers dnums; + dnums.add_lhs_contracting_dimensions(1); + dnums.add_rhs_contracting_dimensions(1); + dnums.add_lhs_batch_dimensions(0); + dnums.add_rhs_batch_dimensions(0); + + DotGeneral(x, y, dnums); + + auto x_data = this->client_ + ->TransferToServer(LiteralUtil::CreateR2FromArray2D( + {{1.0f, 0.0f}, {0.0f, 1.0f}})) + .ConsumeValueOrDie(); + + auto y_data = + this->client_ + ->TransferToServer(LiteralUtil::CreateR3FromArray3D( + {{{1.0f, 2.0f}, {3.0f, 4.0f}}, {{5.0f, 6.0f}, {7.0f, 8.0f}}})) + .ConsumeValueOrDie(); + + this->template ComputeAndCompareR2( + &builder, + /*expected=*/{{1.0f, 2.0f}, {7.0f, 8.0f}}, {x_data.get(), y_data.get()}, + this->error_spec_); +} +#endif // XLA_TEST_BACKEND_CPU + XLA_TYPED_TEST(DotOperationTest_F16F32F64CF64, GeneralMatMulMultipleBatch) { using T = TypeParam; @@ -849,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); @@ -876,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); @@ -905,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); @@ -932,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); @@ -964,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); @@ -996,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); @@ -1020,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); @@ -1044,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); @@ -1078,5 +1157,105 @@ 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"}, + }; + 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 51b50d456e496c9c01c38fb8539bb3737de16937..c84973e17b234c24c84f02a369ce0185f5772cca 100644 --- a/tensorflow/compiler/xla/tests/exhaustive_f32_elementwise_op_test.cc +++ b/tensorflow/compiler/xla/tests/exhaustive_f32_elementwise_op_test.cc @@ -13,11 +13,11 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include "absl/base/casts.h" #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/compiler/xla/tests/client_library_test_base.h" #include "tensorflow/compiler/xla/tests/literal_test_util.h" #include "tensorflow/compiler/xla/tests/test_macros.h" -#include "tensorflow/core/lib/core/casts.h" namespace xla { namespace { @@ -47,7 +47,7 @@ class ExhaustiveF32ElementwiseOpTest // 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}, tensorflow::bit_cast(i)); + input_literal.Set({i - begin}, absl::bit_cast(i)); } } diff --git a/tensorflow/compiler/xla/tests/filecheck.cc b/tensorflow/compiler/xla/tests/filecheck.cc index dcb469087e0064d17ce3b04fdeaf0b6136069a55..1b0bebe2d03a9a153cd0c80329ed0c49c91333a3 100644 --- a/tensorflow/compiler/xla/tests/filecheck.cc +++ b/tensorflow/compiler/xla/tests/filecheck.cc @@ -48,7 +48,7 @@ StatusOr RunFileCheck(const string& input, const string& pattern) { tensorflow::SubProcess file_check_process; file_check_process.SetProgram(file_check_path, - {file_check_path, pattern_path}); + {file_check_path, "-v", pattern_path}); file_check_process.SetChannelAction(tensorflow::CHAN_STDIN, tensorflow::ACTION_PIPE); file_check_process.SetChannelAction(tensorflow::CHAN_STDERR, @@ -71,9 +71,7 @@ StatusOr RunFileCheck(const string& input, const string& pattern) { LOG(WARNING) << "NOTE: FileCheck binary does not exist!"; } - LOG(WARNING) << "FileCheck error: " << standard_error; - LOG(WARNING) << "FileCheck input was:"; - XLA_LOG_LINES(tensorflow::WARNING, input); + LOG(WARNING) << "FileCheck error:\n" << standard_error; LOG(WARNING) << "FileCheck pattern was:"; XLA_LOG_LINES(tensorflow::WARNING, pattern); } else if (!standard_error.empty()) { diff --git a/tensorflow/compiler/xla/tests/fusion_test.cc b/tensorflow/compiler/xla/tests/fusion_test.cc index 4d4b676a538947c8dd92a7e34db72e45766cae2c..2178c9b3f3d39ac034c59585c6836d2bc59162c1 100644 --- a/tensorflow/compiler/xla/tests/fusion_test.cc +++ b/tensorflow/compiler/xla/tests/fusion_test.cc @@ -81,7 +81,7 @@ class FusionTest : public HloTestBase { } auto builder = HloComputation::Builder(TestName()); - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); auto prim_type = primitive_util::NativeToPrimitiveType(); @@ -183,7 +183,7 @@ XLA_TEST_F(FusionTest, Test) { // (-{{1.0, 1.0, 1.0}, {0.0, 0.0, 0.0}}), // {{0.5, 0.5, 0.5}, {0.5, 0.5, 0.5}})) = {{0.5}, {2.72}} auto builder = HloComputation::Builder(TestName()); - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); auto const0 = builder.AddInstruction(HloInstruction::CreateConstant( LiteralUtil::CreateR2({{1.0}, {2.0}, {3.0}}))); auto const1 = builder.AddInstruction(HloInstruction::CreateConstant( @@ -231,7 +231,7 @@ XLA_TEST_F(FusionTest, Parameter) { // Build a computation and fuse part of it so the fusion instruction has an // operand parameter. auto builder = HloComputation::Builder(TestName()); - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); auto const0 = builder.AddInstruction(HloInstruction::CreateConstant( LiteralUtil::CreateR2({{1.0, 2.0, 3.0}}))); auto copy1 = builder.AddInstruction(HloInstruction::CreateUnary( @@ -266,7 +266,7 @@ XLA_TEST_F(FusionTest, RandomizedParallelPartition) { ShapeUtil::MakeShapeWithLayout(F32, {rand_dim0_size, dim1_size}, {1, 0}); // Build simple fusion computation: y = x^2 (elementwise). auto builder = HloComputation::Builder(TestName()); - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); auto two = builder.AddInstruction( HloInstruction::CreateConstant(LiteralUtil::CreateR0(2.0))); @@ -290,7 +290,7 @@ XLA_TEST_F(FusionTest, RandomizedParallelPartition) { XLA_TEST_F(FusionTest, BroadcastIntoBinaryOp) { auto builder = HloComputation::Builder(TestName()); - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); auto const_vector = builder.AddInstruction(HloInstruction::CreateConstant( LiteralUtil::CreateR1({1.0, 2.0, 3.0}))); auto const_array = builder.AddInstruction(HloInstruction::CreateConstant( @@ -314,7 +314,7 @@ XLA_TEST_F(FusionTest, BroadcastIntoBinaryOp) { XLA_TEST_F(FusionTest, ReshapeToScalar) { auto builder = HloComputation::Builder(TestName()); - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); auto single_element_array = builder.AddInstruction( HloInstruction::CreateConstant(LiteralUtil::CreateR2({{5}}))); auto reshape = builder.AddInstruction(HloInstruction::CreateReshape( @@ -329,7 +329,7 @@ XLA_TEST_F(FusionTest, ReshapeToScalar) { XLA_TEST_F(FusionTest, Reshape_3by2_1by2by3) { auto builder = HloComputation::Builder(TestName()); - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); auto const0 = builder.AddInstruction(HloInstruction::CreateConstant( LiteralUtil::CreateR2({{1, 2}, {3, 4}, {5, 6}}))); auto reshape1 = builder.AddInstruction(HloInstruction::CreateReshape( @@ -344,7 +344,7 @@ XLA_TEST_F(FusionTest, Reshape_3by2_1by2by3) { XLA_TEST_F(FusionTest, Reshape_1by2by3_3by2) { auto builder = HloComputation::Builder(TestName()); - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); auto const0 = builder.AddInstruction(HloInstruction::CreateConstant( LiteralUtil::CreateR3({{{1, 2, 3}, {4, 5, 6}}}))); auto reshape1 = builder.AddInstruction( @@ -359,7 +359,7 @@ XLA_TEST_F(FusionTest, Reshape_1by2by3_3by2) { XLA_TEST_F(FusionTest, Reshape_1by1by1_) { auto builder = HloComputation::Builder(TestName()); - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); auto const0 = builder.AddInstruction( HloInstruction::CreateConstant(LiteralUtil::CreateR3({{{7}}}))); auto reshape1 = builder.AddInstruction( @@ -374,7 +374,7 @@ XLA_TEST_F(FusionTest, Reshape_1by1by1_) { XLA_TEST_F(FusionTest, Reshape__1by1by1) { auto builder = HloComputation::Builder(TestName()); - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); auto const0 = builder.AddInstruction( HloInstruction::CreateConstant(LiteralUtil::CreateR0(7))); auto reshape1 = builder.AddInstruction(HloInstruction::CreateReshape( @@ -389,7 +389,7 @@ XLA_TEST_F(FusionTest, Reshape__1by1by1) { XLA_TEST_F(FusionTest, Reshape__) { auto builder = HloComputation::Builder(TestName()); - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); auto const0 = builder.AddInstruction( HloInstruction::CreateConstant(LiteralUtil::CreateR0(7))); auto reshape1 = builder.AddInstruction( @@ -404,7 +404,7 @@ XLA_TEST_F(FusionTest, Reshape__) { XLA_TEST_F(FusionTest, Reshape_3by3_3by3) { auto builder = HloComputation::Builder(TestName()); - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); auto const0 = builder.AddInstruction(HloInstruction::CreateConstant( LiteralUtil::CreateR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}))); auto reshape1 = builder.AddInstruction( @@ -419,7 +419,7 @@ XLA_TEST_F(FusionTest, Reshape_3by3_3by3) { XLA_TEST_F(FusionTest, Transpose_2by3) { auto builder = HloComputation::Builder(TestName()); - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); auto const0 = builder.AddInstruction(HloInstruction::CreateConstant( LiteralUtil::CreateR2({{1, 2, 3}, {4, 5, 6}}))); auto reshape1 = builder.AddInstruction(HloInstruction::CreateTranspose( @@ -434,7 +434,7 @@ XLA_TEST_F(FusionTest, Transpose_2by3) { XLA_TEST_F(FusionTest, Transpose_3by3) { auto builder = HloComputation::Builder(TestName()); - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); auto const0 = builder.AddInstruction(HloInstruction::CreateConstant( LiteralUtil::CreateR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}))); auto reshape1 = builder.AddInstruction(HloInstruction::CreateTranspose( @@ -449,7 +449,7 @@ XLA_TEST_F(FusionTest, Transpose_3by3) { XLA_TEST_F(FusionTest, Reverse) { auto builder = HloComputation::Builder(TestName()); - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); auto const0 = builder.AddInstruction( HloInstruction::CreateConstant(LiteralUtil::CreateR1({1, 2, 3}))); auto reverse1 = builder.AddInstruction(HloInstruction::CreateReverse( @@ -465,7 +465,7 @@ XLA_TEST_F(FusionTest, Reverse) { XLA_TEST_F(FusionTest, ReverseNegate) { auto builder = HloComputation::Builder(TestName()); - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); auto const0 = builder.AddInstruction( HloInstruction::CreateConstant(LiteralUtil::CreateR1({1, 2, 3}))); auto reverse1 = builder.AddInstruction(HloInstruction::CreateReverse( @@ -483,7 +483,7 @@ XLA_TEST_F(FusionTest, ReverseNegate) { XLA_TEST_F(FusionTest, BroadcastNegate) { auto builder = HloComputation::Builder(TestName()); - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); auto const0 = builder.AddInstruction( HloInstruction::CreateConstant(LiteralUtil::CreateR0(1))); auto broadcast1 = builder.AddInstruction(HloInstruction::CreateBroadcast( @@ -501,7 +501,7 @@ XLA_TEST_F(FusionTest, BroadcastNegate) { XLA_TEST_F(FusionTest, SliceNegate) { auto builder = HloComputation::Builder(TestName()); - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); auto const0 = builder.AddInstruction(HloInstruction::CreateConstant( LiteralUtil::CreateR1({1, 2, 3, 4}))); auto slice1 = builder.AddInstruction(HloInstruction::CreateSlice( @@ -519,14 +519,14 @@ XLA_TEST_F(FusionTest, SliceNegate) { XLA_TEST_F(FusionTest, DynamicSliceNegate) { auto builder = HloComputation::Builder(TestName()); - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); 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()) @@ -541,7 +541,7 @@ XLA_TEST_F(FusionTest, DynamicSliceNegate) { XLA_TEST_F(FusionTest, ReshapeNegate) { auto builder = HloComputation::Builder(TestName()); - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); auto const0 = builder.AddInstruction(HloInstruction::CreateConstant( LiteralUtil::CreateR1({1, 2, 3, 4}))); auto reshape1 = builder.AddInstruction( @@ -559,7 +559,7 @@ XLA_TEST_F(FusionTest, ReshapeNegate) { XLA_TEST_F(FusionTest, TransposeNegate) { auto builder = HloComputation::Builder(TestName()); - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); auto const0 = builder.AddInstruction(HloInstruction::CreateConstant( LiteralUtil::CreateR2({{1, 2}, {3, 4}}))); auto transpose1 = builder.AddInstruction(HloInstruction::CreateTranspose( @@ -587,7 +587,7 @@ std::unique_ptr MakeReduceTestComputation() { } XLA_TEST_F(FusionTest, DISABLED_ON_CPU(Reduce)) { - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); auto builder = HloComputation::Builder(TestName()); auto const0 = builder.AddInstruction(HloInstruction::CreateConstant( @@ -607,7 +607,7 @@ XLA_TEST_F(FusionTest, DISABLED_ON_CPU(Reduce)) { } XLA_TEST_F(FusionTest, DISABLED_ON_CPU(ReduceImplicitBroadcast)) { - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); auto builder = HloComputation::Builder(TestName()); auto const0 = builder.AddInstruction(HloInstruction::CreateConstant( @@ -630,7 +630,7 @@ XLA_TEST_F(FusionTest, DISABLED_ON_CPU(ReduceImplicitBroadcast)) { XLA_TEST_F(FusionTest, DISABLED_ON_CPU(ReduceWindow)) { auto builder = HloComputation::Builder(TestName()); - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); auto const0 = builder.AddInstruction(HloInstruction::CreateConstant( LiteralUtil::CreateR2({{2, 3, 5}, {7, 11, 13}, {17, 19, 23}}))); auto const1 = builder.AddInstruction( @@ -682,7 +682,7 @@ XLA_TEST_F(FusionTest, DISABLED_ON_CPU(ReduceWindow)) { // into a fusion, it should remain shared, rather than being duplicated // within the fusion. XLA_TEST_F(FusionTest, SharedConstant) { - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); auto builder = HloComputation::Builder(TestName()); auto const0 = builder.AddInstruction( 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/grouped_convolution_test.cc b/tensorflow/compiler/xla/tests/grouped_convolution_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..8f7049910e70c4e591636a47c1b6ba72cf2c234f --- /dev/null +++ b/tensorflow/compiler/xla/tests/grouped_convolution_test.cc @@ -0,0 +1,245 @@ +/* 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 GroupedConvolution2DSpec { + int64 input_feature, output_feature, window, stride, pad, lhs_dilate; + int64 group_size, group_count; + std::vector activation_dims; + std::vector activation_layout; + std::vector kernel_dims; + std::vector kernel_layout; + std::vector output_dims; + std::vector output_layout; +}; + +class GroupedConvolution2DTest + : public HloTestBase, + public ::testing::WithParamInterface< + ::testing::tuple> {}; + +static std::vector GetConv2DTestCases() { + std::vector config_set; + // Add to this set if you want a new test configuration. + // Rule : the penultimate number must be divisible by the last number. + std::vector> config_options = {{8, 2, 2, 1, 1024, 128}, + {512, 3, 3, 144, 1024, 16}, + {256, 3, 3, 129, 512, 64}, + {64, 1, 2, 127, 32, 8}, + {256, 3, 3, 256, 1024, 4}}; + + for (auto option : config_options) { + int64 output_feature = option[0]; + int64 activation_size = option[1]; + int64 kernel_size = option[2]; + int64 batch = option[3]; + int64 input_feature = option[4]; + int64 group_size = option[5]; + + std::vector kernel_layout = {3, 2, 1, 0}; + GroupedConvolution2DSpec config; + config.group_size = group_size; + config.group_count = input_feature / group_size; + config.output_feature = output_feature; + config.window = kernel_size; + + config.activation_dims = {batch, activation_size, activation_size, + input_feature}; + config.activation_layout = {3, 0, 2, 1}; + + config.kernel_dims = {kernel_size, kernel_size, group_size, output_feature}; + config.kernel_layout = {3, 2, 1, 0}; + + if (activation_size == 1 && kernel_size == 2) { + // Test for outer dim. + config.output_dims = {batch, activation_size + kernel_size - 1, + activation_size + kernel_size, output_feature}; + } else if (output_feature == 256) { + // Restrict dilation-based tests only to one feature configuration. + config.stride = activation_size - 1; + config.pad = 0; + config.lhs_dilate = output_feature / 32; + config.output_dims = {batch, output_feature / 32, + activation_size - kernel_size + 1, output_feature}; + } else { + config.stride = config.pad = config.lhs_dilate = -1; + config.output_dims = {batch, activation_size - kernel_size + 1, + activation_size - kernel_size + 1, output_feature}; + } + + // Try this layout for all kernel shapes. + config.output_layout = {3, 0, 2, 1}; + config_set.push_back(config); + + // Try other layouts only for certain kernel shapes. + if (kernel_size % 2 == 0) { + config.activation_layout = {0, 3, 2, 1}; + config_set.push_back(config); + + config.output_layout = {0, 3, 2, 1}; + config_set.push_back(config); + + config.activation_layout = {3, 0, 2, 1}; + config_set.push_back(config); + } + } + + return config_set; +} + +string GroupedConvolution2DTestDataToString( + 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"), + "_activation_layout_", absl::StrJoin(spec.activation_layout, "_"), + "_kernel_dims_", absl::StrJoin(spec.kernel_dims, "x"), "_kernel_layout_", + absl::StrJoin(spec.kernel_layout, "_"), "_output_dims_", + absl::StrJoin(spec.output_dims, "x"), "_output_layout_", + absl::StrJoin(spec.output_layout, "_"), data_type); + // -1 indicates non-existence. + if (spec.stride != -1) { + absl::StrAppend(&str, "_lhs_dilation_", spec.lhs_dilate, "x1"); + } + + // Test names are not allowed to contain the '-' character. + absl::c_replace(str, '-', 'n'); + return str; +} + +string BuildHloTextGroupedConvolution2D(const GroupedConvolution2DSpec& spec, + bool use_bfloat16) { + const string data_type = GetFloatDataType(use_bfloat16); + if (spec.activation_dims[1] == 1 && spec.kernel_dims[1] == 2) { + // Check for outer dim. + return absl::StrFormat( + R"( + HloModule TensorFlowDepthwiseConv + + 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_1x%d_%d rhs_dilate=1x%d}, dim_labels=b01f_01io->b01f, + feature_group_count=%d + } + )", + data_type, absl::StrJoin(spec.activation_dims, ","), + absl::StrJoin(spec.activation_layout, ","), data_type, + absl::StrJoin(spec.kernel_dims, ","), + absl::StrJoin(spec.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_layout, ","), data_type, + absl::StrJoin(spec.kernel_dims, ","), + absl::StrJoin(spec.kernel_layout, ","), spec.window, spec.window, + spec.window, spec.window, spec.window, spec.group_count); + + } else if (spec.stride == -1) { + // Check for basic, non-dilated cases. + return absl::StrFormat( + R"( + HloModule TensorFlowDepthwiseConv + + 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}, dim_labels=b01f_01io->b01f, + feature_group_count=%d + } + )", + data_type, absl::StrJoin(spec.activation_dims, ","), + absl::StrJoin(spec.activation_layout, ","), data_type, + absl::StrJoin(spec.kernel_dims, ","), + absl::StrJoin(spec.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_layout, ","), data_type, + absl::StrJoin(spec.kernel_dims, ","), + absl::StrJoin(spec.kernel_layout, ","), spec.window, spec.window, + spec.group_count); + } else { + // Check for base dilations. + return absl::StrFormat( + R"( + HloModule TensorFlowDepthwiseConv + + 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 stride=%dx1 pad=%d_%dx0_0 lhs_dilate=%dx1}, + dim_labels=b01f_01io->b01f, feature_group_count=%d + } + )", + data_type, absl::StrJoin(spec.activation_dims, ","), + absl::StrJoin(spec.activation_layout, ","), data_type, + absl::StrJoin(spec.kernel_dims, ","), + absl::StrJoin(spec.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_layout, ","), data_type, + absl::StrJoin(spec.kernel_dims, ","), + absl::StrJoin(spec.kernel_layout, ","), spec.window, spec.window, + spec.stride, 0, 0, spec.lhs_dilate, spec.group_count); + } +} + +XLA_TEST_P(GroupedConvolution2DTest, DoIt) { + const GroupedConvolution2DSpec& spec = ::testing::get<0>(GetParam()); + bool use_bfloat16 = ::testing::get<1>(GetParam()); + const string hlo_text = BuildHloTextGroupedConvolution2D(spec, use_bfloat16); + + EXPECT_TRUE(RunAndCompare(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( + GroupedConvolution2DTestWithRandomIndices, GroupedConvolution2DTest, + ::testing::Combine(::testing::ValuesIn(GetConv2DTestCases()), + ::testing::Bool()), + GroupedConvolution2DTestDataToString); + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/tests/hlo_test_base.cc b/tensorflow/compiler/xla/tests/hlo_test_base.cc index 7ab2ecda58666acd7e9b8587d200a902b75822f3..66f72ba8d20b8ef1f436da4425b2bb6518ee9a94 100644 --- a/tensorflow/compiler/xla/tests/hlo_test_base.cc +++ b/tensorflow/compiler/xla/tests/hlo_test_base.cc @@ -23,8 +23,8 @@ limitations under the License. #include "absl/algorithm/container.h" #include "absl/memory/memory.h" #include "absl/types/span.h" +#include "tensorflow/compiler/xla/debug_options_flags.h" #include "tensorflow/compiler/xla/layout_util.h" -#include "tensorflow/compiler/xla/legacy_flags/debug_options_flags.h" #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/service/hlo_parser.h" #include "tensorflow/compiler/xla/service/platform_util.h" @@ -85,6 +85,25 @@ ProgramShape GetProgramShapeWithLayout(const HloModule& module) { } // namespace +Status VerifiedHloModule::Verify() { + if (computation_count() == 0) { + // The computation was never built. Nothing to verify. + return Status::OK(); + } + return verifier_.Run(this).status(); +} + +void VerifiedHloModule::VerifyOrAddFailure(const string& message) { + Status status = Verify(); + if (!status.ok()) { + ADD_FAILURE() << "HloVerifier failed on module " << name() + << (message.empty() ? "" : absl::StrCat(" (", message, ")")) + << ": " << status; + LOG(ERROR) << "Contents of bad module:"; + XLA_LOG_LINES(tensorflow::ERROR, ToString()); + } +} + HloTestBase::HloTestBase(bool verifier_layout_sensitive, bool allow_mixed_precision_in_hlo_verifier, std::function @@ -100,17 +119,42 @@ HloTestBase::HloTestBase(se::Platform* test_platform, bool allow_mixed_precision_in_hlo_verifier, std::function instruction_can_change_layout_func) - : test_runner_(test_platform), reference_runner_(reference_platform) { + : test_runner_(test_platform), + reference_runner_(reference_platform), + verifier_layout_sensitive_(verifier_layout_sensitive), + allow_mixed_precision_in_hlo_verifier_( + allow_mixed_precision_in_hlo_verifier) { hlo_verifier_ = absl::make_unique( /*layout_sensitive=*/verifier_layout_sensitive, /*allow_mixed_precision=*/allow_mixed_precision_in_hlo_verifier, instruction_can_change_layout_func); } -std::unique_ptr HloTestBase::CreateNewModule(const string& name) { +std::unique_ptr HloTestBase::CreateNewUnverifiedModule( + const string& name) { return absl::make_unique(name, GetModuleConfigForTest()); } +std::unique_ptr HloTestBase::CreateNewVerifiedModule( + const string& name) { + return absl::make_unique( + name, GetModuleConfigForTest(), verifier_layout_sensitive_, + allow_mixed_precision_in_hlo_verifier_, + backend().compiler()->ShapeSizeBytesFunction()); +} + +StatusOr> +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_, + backend().compiler()->ShapeSizeBytesFunction()); + TF_RETURN_IF_ERROR(ParseHloString(hlo_text, module.get())); + TF_RETURN_IF_ERROR(module->Verify()); + return std::move(module); +} + /* static */ StatusOr HloTestBase::RunHloPass(HloPassInterface* hlo_pass, HloModule* module) { @@ -135,10 +179,11 @@ PrecisionConfig HloTestBase::DefaultPrecisionConfig(int operands) { } DebugOptions HloTestBase::GetDebugOptionsForTest() { - auto debug_options = legacy_flags::GetDebugOptionsFromFlags(); + auto debug_options = GetDebugOptionsFromFlags(); // TODO(b/38354253): Change tests to use Parameters instead of Constants. debug_options.add_xla_disable_hlo_passes("constant_folding"); debug_options.set_xla_gpu_max_kernel_unroll_factor(1); + debug_options.set_xla_hlo_evaluator_use_fast_path(true); return debug_options; } diff --git a/tensorflow/compiler/xla/tests/hlo_test_base.h b/tensorflow/compiler/xla/tests/hlo_test_base.h index 217428befa474448cf2dcbae2eb6cb5b0e61d44c..69a4f96288c7285010e9adbdc33f1b394f58d8d2 100644 --- a/tensorflow/compiler/xla/tests/hlo_test_base.h +++ b/tensorflow/compiler/xla/tests/hlo_test_base.h @@ -20,6 +20,7 @@ limitations under the License. #include #include +#include "absl/base/macros.h" #include "absl/types/optional.h" #include "absl/types/span.h" #include "tensorflow/compiler/xla/service/backend.h" @@ -38,6 +39,33 @@ limitations under the License. namespace xla { +// An HLO module derived class which verifies itself on destruction. This class +// is intended to be used in unit tests. Any verification errors are raised via +// ADD_FAILURE. +class VerifiedHloModule : public HloModule { + public: + VerifiedHloModule(const string& name, const HloModuleConfig& config, + bool verifier_layout_sensitive, + 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, + /*instruction_can_change_layout_func=*/{}, shape_size_function) {} + + ~VerifiedHloModule() override { VerifyOrAddFailure("in destructor"); } + + // Verifies the module using HloVerifier and returns the status. + Status Verify(); + + // Verifies the module and flags any error with ADD_FAILURE. 'message' is + // included in the failure message. + void VerifyOrAddFailure(const string& message); + + private: + HloVerifier verifier_; +}; + // A base class for tests which build and/or run HLO code. The class includes // support for running an HLO module on two platforms and compare the results. // This is a lower level of abstraction than using the client interface and @@ -72,7 +100,22 @@ class HloTestBase : public ::testing::Test { // options from command-line flags. If you want a fresh HloModule object and // then add HloComputations to it, it's recommended to use this method in your // tests. - std::unique_ptr CreateNewModule(const string& name = TestName()); + // + // This returns a vanilla HloModule that doesn't run the HLO verifier on + // destruction. + ABSL_DEPRECATED("Use CreateNewVerifiedModule instead.") + std::unique_ptr CreateNewUnverifiedModule( + const string& name = TestName()); + + // Like CreateNewUnverifiedModule, except the HloModule returned here runs the + // HLO verifier on destruction. + std::unique_ptr CreateNewVerifiedModule( + const string& name = TestName()); + + // Parses the given string and returns module as a VerifiedHloModule. + StatusOr> ParseAndReturnVerifiedModule( + absl::string_view hlo_text, + const HloModuleConfig& config = HloModuleConfig()); // Runs the hlo_pass with the provided module and returns the result. This // function also verifies that the module remains unchanged when hlo_pass @@ -247,6 +290,8 @@ class HloTestBase : public ::testing::Test { HloRunner test_runner_; HloRunner reference_runner_; + bool verifier_layout_sensitive_; + bool allow_mixed_precision_in_hlo_verifier_; std::unique_ptr hlo_verifier_; ErrorSpec error_spec_{0.0001}; diff --git a/tensorflow/compiler/xla/tests/hlo_verified_test_base.cc b/tensorflow/compiler/xla/tests/hlo_verified_test_base.cc deleted file mode 100644 index 8bd0a729b77f3ec14204952cb0062103c823883e..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/tests/hlo_verified_test_base.cc +++ /dev/null @@ -1,88 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/compiler/xla/tests/hlo_verified_test_base.h" - -#include "absl/memory/memory.h" -#include "tensorflow/compiler/xla/service/hlo_parser.h" -#include "tensorflow/compiler/xla/service/hlo_verifier.h" -#include "tensorflow/compiler/xla/shape_util.h" -#include "tensorflow/compiler/xla/status_macros.h" -#include "tensorflow/core/platform/logging.h" - -namespace xla { - -Status VerifiedHloModule::Verify() { - if (computation_count() == 0) { - // The computation was never built. Nothing to verify. - return Status::OK(); - } - return verifier_.Run(this).status(); -} - -void VerifiedHloModule::VerifyOrAddFailure(const string& message) { - Status status = Verify(); - if (!status.ok()) { - ADD_FAILURE() << "HloVerifier failed on module " << name() - << (message.empty() ? "" : absl::StrCat(" (", message, ")")) - << ": " << status; - } -} - -HloVerifiedTestBase::HloVerifiedTestBase(bool layout_sensitive, - bool allow_mixed_precision) - : HloTestBase( - /*verifier_layout_sensitive=*/layout_sensitive, - /*allow_mixed_precision_in_hlo_verifier=*/allow_mixed_precision), - verifier_layout_sensitive_(layout_sensitive), - allow_mixed_precision_in_hlo_verifier_(allow_mixed_precision) {} - -HloModule& HloVerifiedTestBase::module() { - if (!module_) { - module_ = CreateNewVerifiedModule(TestName()); - } - return *module_; -} - -HloModule* HloVerifiedTestBase::CreateNewModule(const string& name) { - modules_.emplace_back(CreateNewVerifiedModule(name)); - return modules_.back().get(); -} - -void HloVerifiedTestBase::ParseAndVerifyModule(absl::string_view hlo_text, - const HloModuleConfig& config) { - CHECK(!module_) << "Called ParseModule when test already has a module."; - module_ = CreateNewVerifiedModule(TestName()); - TF_CHECK_OK(ParseHloString(hlo_text, module_.get())); - module_->VerifyOrAddFailure("after parsing"); -} - -StatusOr> -HloVerifiedTestBase::ParseAndReturnVerifiedModule( - absl::string_view hlo_text, const HloModuleConfig& config) { - auto module = CreateNewVerifiedModule(TestName()); - TF_RETURN_IF_ERROR(ParseHloString(hlo_text, module.get())); - TF_RETURN_IF_ERROR(module->Verify()); - return std::move(module); -} - -std::unique_ptr HloVerifiedTestBase::CreateNewVerifiedModule( - const string& name) { - return absl::make_unique( - name, GetModuleConfigForTest(), verifier_layout_sensitive_, - allow_mixed_precision_in_hlo_verifier_); -} - -} // namespace xla diff --git a/tensorflow/compiler/xla/tests/hlo_verified_test_base.h b/tensorflow/compiler/xla/tests/hlo_verified_test_base.h deleted file mode 100644 index 388a99bb36408665edbc20ade6c6a733d64db88d..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/tests/hlo_verified_test_base.h +++ /dev/null @@ -1,105 +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_TESTS_HLO_VERIFIED_TEST_BASE_H_ -#define TENSORFLOW_COMPILER_XLA_TESTS_HLO_VERIFIED_TEST_BASE_H_ - -#include -#include -#include - -#include "absl/base/macros.h" -#include "tensorflow/compiler/xla/service/hlo_module.h" -#include "tensorflow/compiler/xla/tests/hlo_test_base.h" - -namespace xla { - -// An HLO module derived class which verifies itself on destruction. This class -// is intended to be used in unit tests. Any verification errors are raised via -// ADD_FAILURE. -class VerifiedHloModule : public HloModule { - public: - VerifiedHloModule(const string& name, const HloModuleConfig& config, - bool verifier_layout_sensitive, - bool allow_mixed_precision_in_hlo_verifier) - : HloModule(name, config), - verifier_(verifier_layout_sensitive, - allow_mixed_precision_in_hlo_verifier) {} - - ~VerifiedHloModule() override { VerifyOrAddFailure("in destructor"); } - - // Verifies the module using HloVerifier and returns the status. - Status Verify(); - - // Verifies the module and flags any error with ADD_FAILURE. 'message' is - // included in the failure message. - void VerifyOrAddFailure(const string& message); - - private: - HloVerifier verifier_; -}; - -// A base class for HLO tests that stores a default VerifiedHloModule. -class HloVerifiedTestBase : public HloTestBase { - protected: - HloVerifiedTestBase(bool layout_sensitive = false, - bool allow_mixed_precision = false); - - // Constructs a default shape verifier. - std::unique_ptr MakeShapeVerifier(); - - // Returns the default HloModule, lazily creating it if necessary via - // HloTestBase::CreateNewModule(). - ABSL_DEPRECATED("Use CreateNewVerifiedModule() instead.") - HloModule& module(); - - ABSL_DEPRECATED("Use ParseAndReturnVerifiedModule() instead.") - void ParseAndVerifyModule(absl::string_view hlo_text, - const HloModuleConfig& config = HloModuleConfig()); - - // Parses the given string and returns module as a VerifiedHloModule. - StatusOr> ParseAndReturnVerifiedModule( - absl::string_view hlo_text, - const HloModuleConfig& config = HloModuleConfig()); - - // Creates a new module for a test, and stores it in modules_ so it can be - // verified. Intentionally hides HloTestBase::CreateNewModule, to prevent - // creation of unverified modules. - ABSL_DEPRECATED("Use CreateNewVerifiedModule() instead.") - HloModule* CreateNewModule(const string& name = TestName()); - - // Creates and returns a verified HLO module with the given name. - std::unique_ptr CreateNewVerifiedModule( - const string& name = TestName()); - - private: - // It is confusing to store modules created by module() and CreateNewModule() - // in different fields, but it allows us to migrate tests to - // HloVerifiedTestBase more easily, so it's a win because we can verify more - // modules. See b/80488902. - // - // Lazily populated. Access via module(). - std::unique_ptr module_; - - // Populated by calls to CreateNewModule. - std::vector> modules_; - - bool verifier_layout_sensitive_; - bool allow_mixed_precision_in_hlo_verifier_; -}; - -} // namespace xla - -#endif // TENSORFLOW_COMPILER_XLA_TESTS_HLO_VERIFIED_TEST_BASE_H_ diff --git a/tensorflow/compiler/xla/tests/hlo_verified_test_base_test.cc b/tensorflow/compiler/xla/tests/hlo_verified_test_base_test.cc deleted file mode 100644 index 5c0263e811f94c90a69a460525ffa0c65127ebb5..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/tests/hlo_verified_test_base_test.cc +++ /dev/null @@ -1,158 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/compiler/xla/tests/hlo_verified_test_base.h" - -#include "tensorflow/compiler/xla/service/hlo_opcode.h" -#include "tensorflow/compiler/xla/service/hlo_verifier.h" -#include "tensorflow/compiler/xla/shape_util.h" -#include "tensorflow/compiler/xla/tests/test_macros.h" -#include "tensorflow/core/lib/core/status_test_util.h" -#include "tensorflow/core/platform/test.h" -#include "tensorflow/core/platform/types.h" - -namespace xla { -namespace { - -// This class includes unit tests which are expected to fail because invalid HLO -// modules are intentionally built. Unfortunately, Tensorflow doesn't appear to -// include the necessary gunit parts to test this test machinery (needs the -// macro EXPECT_NONFATAL_FAILURE). The disabled tests can be run with the -// disabled tests enabled and failures can be manually compared against -// expectations. -class HloVerifiedTestBaseTest : public HloVerifiedTestBase {}; - -XLA_TEST_F(HloVerifiedTestBaseTest, NoModule) { - // Test shouldn't fail if no module is created at all. -} - -XLA_TEST_F(HloVerifiedTestBaseTest, GoodLazilyCreatedModule) { - // Use module() to lazily create an empty module, build it up, and verify no - // failures. - HloModule& hlo_module = module(); - auto builder = HloComputation::Builder(TestName()); - auto input = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0))); - builder.AddInstruction( - HloInstruction::CreateUnary(input->shape(), HloOpcode::kNegate, input)); - hlo_module.AddEntryComputation(builder.Build()); -} - -// This test is expected to fail. See test class comment. -XLA_TEST_F(HloVerifiedTestBaseTest, DISABLED_BadLazilyCreatedModule) { - // Use module() to lazily create an empty module and build up an invalid - // module. - HloModule& hlo_module = module(); - auto builder = HloComputation::Builder(TestName()); - auto input = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0))); - builder.AddInstruction( - HloInstruction::CreateUnary(input->shape(), HloOpcode::kNegate, input)); - hlo_module.AddEntryComputation(builder.Build()); - - *hlo_module.entry_computation()->root_instruction()->mutable_shape() = - ShapeUtil::MakeShape(PRED, {1, 2, 3}); -} - -XLA_TEST_F(HloVerifiedTestBaseTest, GoodCreateNewModule) { - // Call CreateNewModule and build up a valid module. - HloModule* module = CreateNewModule(); - auto builder = HloComputation::Builder(TestName()); - auto input = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0))); - builder.AddInstruction( - HloInstruction::CreateUnary(input->shape(), HloOpcode::kNegate, input)); - module->AddEntryComputation(builder.Build()); -} - -// This test is expected to fail. See test class comment. -XLA_TEST_F(HloVerifiedTestBaseTest, DISABLED_BadCreateNewModule) { - // Call CreateNewModule and build up a invalid module. - HloModule* module = CreateNewModule(); - auto builder = HloComputation::Builder(TestName()); - auto input = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0))); - builder.AddInstruction( - HloInstruction::CreateUnary(input->shape(), HloOpcode::kNegate, input)); - module->AddEntryComputation(builder.Build()); - - *module->entry_computation()->root_instruction()->mutable_shape() = - ShapeUtil::MakeShape(PRED, {1, 2, 3}); -} - -XLA_TEST_F(HloVerifiedTestBaseTest, ParseAndVerifyModuleGood) { - const char* const hlo_string = R"( -HloModule ParseAndVerifyModuleGood - -ENTRY entry { - x = f32[] parameter(0) - y = f32[] parameter(1) - ROOT add = f32[] add(x,y) -} -)"; - - ParseAndVerifyModule(hlo_string); - EXPECT_EQ(module().entry_computation()->instruction_count(), 3); -} - -XLA_TEST_F(HloVerifiedTestBaseTest, ParseAndReturnVerifiedModuleGood) { - const char* const hlo_string = R"( -HloModule ParseAndReturnVerifiedModuleGood - -ENTRY entry { - x = f32[] parameter(0) - y = f32[] parameter(1) - ROOT add = f32[] add(x,y) -} -)"; - - TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(hlo_string)); - EXPECT_EQ(module->entry_computation()->instruction_count(), 3); -} - -XLA_TEST_F(HloVerifiedTestBaseTest, ParseAndReturnVerifiedModuleInvalidText) { - const char* const hlo_string = R"( -HloModule ParseAndReturnVerifiedModuleGood - -ENTRY entry { - x = f32[] parameter(0) - y = f32[] parameter(1) - ROOT add = f32[] add(x,y) -} - -RANDOM GARBAGE -)"; - - ASSERT_IS_NOT_OK(ParseAndReturnVerifiedModule(hlo_string).status()); -} - -// This test is expected to fail. See test class comment. -XLA_TEST_F(HloVerifiedTestBaseTest, DISABLED_ParseAndReturnVerifiedModuleBad) { - const char* const hlo_string = R"( -HloModule ParseAndReturnVerifiedModuleBad - -ENTRY entry { - x = f32[] parameter(0) - y = f32[] parameter(1) - ROOT add = f32[1234] add(x,y) -} -)"; - - ASSERT_IS_NOT_OK(ParseAndReturnVerifiedModule(hlo_string).status()); -} - -} // namespace -} // namespace xla diff --git a/tensorflow/compiler/xla/tests/iota_test.cc b/tensorflow/compiler/xla/tests/iota_test.cc index 310f3495922250d68aa463fcbb24ef0b04603d09..37b2c635eebe57590e1ba73c62f015ccf399b548 100644 --- a/tensorflow/compiler/xla/tests/iota_test.cc +++ b/tensorflow/compiler/xla/tests/iota_test.cc @@ -80,7 +80,7 @@ TEST_P(IotaR2Test, DoIt) { } INSTANTIATE_TEST_CASE_P(IotaR2TestInstantiation, IotaR2Test, - ::testing::Combine(::testing::Values(F32, S32), + ::testing::Combine(::testing::Values(F32, S32, BF16), ::testing::Range(/*start=*/10, /*end=*/1001, /*step=*/10), @@ -113,5 +113,26 @@ INSTANTIATE_TEST_CASE_P(IotaR3TestInstantiation, IotaR3Test, /*step=*/10), ::testing::Values(0, 1, 2))); +class IotaR3PredTest : public ClientLibraryTestBase, + public ::testing::WithParamInterface {}; + +TEST_P(IotaR3PredTest, DoIt) { + const auto element_type = PRED; + const int64 num_elements = 2; + const int64 iota_dim = GetParam(); + XlaBuilder builder(TestName() + "_" + PrimitiveType_Name(element_type)); + std::vector dimensions = {42, 19}; + dimensions.insert(dimensions.begin() + iota_dim, num_elements); + Iota(&builder, ShapeUtil::MakeShape(element_type, dimensions), iota_dim); + if (primitive_util::IsFloatingPointType(element_type)) { + ComputeAndCompare(&builder, {}, ErrorSpec{0.0001}); + } else { + ComputeAndCompare(&builder, {}); + } +} + +INSTANTIATE_TEST_CASE_P(IotaR3PredTestInstantiation, IotaR3PredTest, + ::testing::Values(0, 1, 2)); + } // namespace } // namespace xla 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/literal_test_util_test.cc b/tensorflow/compiler/xla/tests/literal_test_util_test.cc index b6f9b8156b51144e4f74d285b1e4111d098f13c2..ea9b3037cf482e41238413179888f125822d161c 100644 --- a/tensorflow/compiler/xla/tests/literal_test_util_test.cc +++ b/tensorflow/compiler/xla/tests/literal_test_util_test.cc @@ -89,11 +89,11 @@ TEST(LiteralTestUtilTest, ExpectNearFailurePlacesResultsInTemporaryDirectory) { Literal literal = Literal::CreateFromProto(literal_proto).ConsumeValueOrDie(); if (result.find("expected") != string::npos) { - EXPECT_EQ("2", literal.ToString()); + EXPECT_EQ("f32[] 2", literal.ToString()); } else if (result.find("actual") != string::npos) { - EXPECT_EQ("4", literal.ToString()); + EXPECT_EQ("f32[] 4", literal.ToString()); } else if (result.find("mismatches") != string::npos) { - EXPECT_EQ("true", literal.ToString()); + EXPECT_EQ("pred[] true", literal.ToString()); } else { FAIL() << "unknown file in temporary directory: " << result; } @@ -105,9 +105,9 @@ TEST(LiteralTestUtilTest, NotEqualHasValuesInMessage) { auto actual = LiteralUtil::CreateR1({4, 5, 6}); ::testing::AssertionResult result = LiteralTestUtil::Equal(expected, actual); EXPECT_THAT(result.message(), - ::testing::HasSubstr("Expected literal:\n{1, 2, 3}")); + ::testing::HasSubstr("Expected literal:\ns32[3] {1, 2, 3}")); EXPECT_THAT(result.message(), - ::testing::HasSubstr("Actual literal:\n{4, 5, 6}")); + ::testing::HasSubstr("Actual literal:\ns32[3] {4, 5, 6}")); } TEST(LiteralTestUtilTest, NearComparatorR1) { diff --git a/tensorflow/compiler/xla/tests/llvm_compiler_test.cc b/tensorflow/compiler/xla/tests/llvm_compiler_test.cc index c622b295094e53e63d0ed692d428bc97724c787c..a78ccacec114858740bf1b9c04e9b688bca5818d 100644 --- a/tensorflow/compiler/xla/tests/llvm_compiler_test.cc +++ b/tensorflow/compiler/xla/tests/llvm_compiler_test.cc @@ -68,7 +68,7 @@ class LLVMCompilerTest : public ::testing::Test { builder.AddInstruction( HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0))); - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); hlo_module->AddEntryComputation(builder.Build()); compiler->SetPreOptimizationHook(pre_opt_hook); @@ -90,7 +90,7 @@ class LLVMCompilerTest : public ::testing::Test { builder.AddInstruction( HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0))); - std::unique_ptr hlo_module = CreateNewModule(); + std::unique_ptr hlo_module = CreateNewUnverifiedModule(); hlo_module->AddEntryComputation(builder.Build()); auto module_group = absl::make_unique("test_module_group"); @@ -124,9 +124,9 @@ class LLVMCompilerTest : public ::testing::Test { return ::testing::UnitTest::GetInstance()->current_test_info()->name(); } - static std::unique_ptr CreateNewModule() { + static std::unique_ptr CreateNewUnverifiedModule() { HloModuleConfig config; - config.set_debug_options(legacy_flags::GetDebugOptionsFromFlags()); + config.set_debug_options(GetDebugOptionsFromFlags()); return absl::make_unique(TestName(), config); } }; diff --git a/tensorflow/compiler/xla/tests/local_client_execute_test.cc b/tensorflow/compiler/xla/tests/local_client_execute_test.cc index a99b43f4690b3063f76e2cda1e58c9b4ba9a1df4..96527886b718bc1ea4ce8cc2d7dbeb2e3ef1d1eb 100644 --- a/tensorflow/compiler/xla/tests/local_client_execute_test.cc +++ b/tensorflow/compiler/xla/tests/local_client_execute_test.cc @@ -205,7 +205,7 @@ XLA_TEST_F(LocalClientExecuteTest, TupleResult) { ScopedShapedBuffer result = ExecuteLocallyOrDie(computation, {&x_array, &y_array}); - EXPECT_TRUE(ShapeUtil::IsTuple(result.on_host_shape())); + EXPECT_TRUE(result.on_host_shape().IsTuple()); EXPECT_EQ(3, ShapeUtil::TupleElementCount(result.on_host_shape())); Literal result_literal = ShapedBufferToLiteral(result); @@ -233,7 +233,7 @@ XLA_TEST_F(LocalClientExecuteTest, NestedTupleResult) { ScopedShapedBuffer result = ExecuteLocallyOrDie(computation, {&x_array, &y_array}); - EXPECT_TRUE(ShapeUtil::IsTuple(result.on_host_shape())); + EXPECT_TRUE(result.on_host_shape().IsTuple()); EXPECT_EQ(2, ShapeUtil::TupleElementCount(result.on_host_shape())); Literal result_literal = ShapedBufferToLiteral(result); @@ -311,7 +311,7 @@ XLA_TEST_F(LocalClientExecuteTest, TupleArguments) { ScopedShapedBuffer result = ExecuteLocallyOrDie(computation, {&x_buffer, &y_buffer}); - EXPECT_TRUE(ShapeUtil::IsTuple(result.on_host_shape())); + EXPECT_TRUE(result.on_host_shape().IsTuple()); EXPECT_EQ(2, ShapeUtil::TupleElementCount(result.on_host_shape())); Literal result_literal = ShapedBufferToLiteral(result); @@ -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 56aaeb0e6878737e6c689e8065d8f1e1871b3472..1fd9cb055c0bebc0f31496eb82f53a7b7a6cbfba 100644 --- a/tensorflow/compiler/xla/tests/multioutput_fusion_test.cc +++ b/tensorflow/compiler/xla/tests/multioutput_fusion_test.cc @@ -62,7 +62,7 @@ class MultiOutputFusionTest : public HloTestBase { void RunTest2D(bool manual_fusion, int64 size) { auto builder = HloComputation::Builder(TestName()); - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); const Shape elem_shape0 = ShapeUtil::MakeShapeWithLayout(F32, {}, {}); const Shape elem_shape2 = @@ -122,7 +122,7 @@ class MultiOutputFusionTest : public HloTestBase { void RunTest1D(bool manual_fusion, int size) { auto builder = HloComputation::Builder(TestName()); - auto hlo_module = CreateNewModule(); + auto hlo_module = CreateNewUnverifiedModule(); const Shape elem_shape_F32 = ShapeUtil::MakeShapeWithDescendingLayout(F32, {size}); @@ -192,7 +192,7 @@ XLA_TEST_F(MultiOutputFusionTest, DiffentTypesFusion) { RunTest1D(true, 8); } XLA_TEST_F(MultiOutputFusionTest, FusionNodeIsRoot) { const char* testcase = R"( - HloModule m + HloModule m, is_scheduled=true fused_computation { x.param_0 = (((s32[]), f32[]), (f32[], s32[])) parameter(0) @@ -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)), @@ -224,7 +222,7 @@ XLA_TEST_F(MultiOutputFusionTest, FusionNodeIsRoot) { XLA_TEST_F(MultiOutputFusionTest, MultiOutputLoopFusion) { const char* testcase = R"( - HloModule m + HloModule m, is_scheduled=true fused_computation { p = f32[4] parameter(0) @@ -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); @@ -251,7 +247,7 @@ XLA_TEST_F(MultiOutputFusionTest, MultiOutputLoopFusion) { XLA_TEST_F(MultiOutputFusionTest, MultiOutputLoopFeedingMap) { const char* testcase = R"( - HloModule m + HloModule m, is_scheduled=true fused_computation { p = f32[] parameter(0) @@ -273,16 +269,14 @@ 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); } const char* const kScalarOps = R"( - HloModule m + HloModule m, is_scheduled=true Add { lhsadd = f32[] parameter(0) @@ -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/pred_test.cc b/tensorflow/compiler/xla/tests/pred_test.cc index 58539e6b061b0cec1cc660b52e78894e5deeea56..774eb8d2a85914c52597144e70838ee117ee1134 100644 --- a/tensorflow/compiler/xla/tests/pred_test.cc +++ b/tensorflow/compiler/xla/tests/pred_test.cc @@ -87,8 +87,8 @@ TEST_F(PredTest, ConstantR2Pred) { XlaBuilder builder(TestName()); ConstantR2(&builder, {{false, true, true}, {true, false, false}}); const string expected = R"(pred[2,3] { - { 011 }, - { 100 } + { 0, 1, 1 }, + { 1, 0, 0 } })"; EXPECT_EQ(expected, ExecuteToString(&builder, {})); } 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 193e66969259f3a8dc18f959c5e72baee11dce24..f80d29b9de440b11c36e8c9bc65d4a93353a6267 100644 --- a/tensorflow/compiler/xla/tests/reduce_precision_test.cc +++ b/tensorflow/compiler/xla/tests/reduce_precision_test.cc @@ -19,6 +19,7 @@ limitations under the License. #include #include +#include "absl/base/casts.h" #include "absl/strings/str_cat.h" #include "tensorflow/compiler/xla/array2d.h" #include "tensorflow/compiler/xla/client/global_data.h" @@ -34,7 +35,6 @@ limitations under the License. #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/lib/core/casts.h" #include "tensorflow/core/platform/types.h" namespace xla { @@ -216,14 +216,13 @@ XLA_TEST_P(ReducePrecisionAccuracyTest, ReducePrecisionF32) { const uint32_t sign_bit = 1u << 31; for (const auto& test_value : test_values) { // Add positive values. - input_values.push_back(tensorflow::bit_cast(test_value[0])); - expected_values.push_back(tensorflow::bit_cast(test_value[index])); + input_values.push_back(absl::bit_cast(test_value[0])); + expected_values.push_back(absl::bit_cast(test_value[index])); // Add negative values. We do this in the bitwise representation so as to // avoid problems with NaN handling. - input_values.push_back( - tensorflow::bit_cast(test_value[0] ^ sign_bit)); + input_values.push_back(absl::bit_cast(test_value[0] ^ sign_bit)); expected_values.push_back( - tensorflow::bit_cast(test_value[index] ^ sign_bit)); + absl::bit_cast(test_value[index] ^ sign_bit)); } // This is required for proper handling of NaN values. diff --git a/tensorflow/compiler/xla/tests/reduce_test.cc b/tensorflow/compiler/xla/tests/reduce_test.cc index 83997cdac21c437d460dabdbdfdb31100b1359af..18c99490a387923aaf68e06041cd11ed3b954aa5 100644 --- a/tensorflow/compiler/xla/tests/reduce_test.cc +++ b/tensorflow/compiler/xla/tests/reduce_test.cc @@ -32,6 +32,7 @@ limitations under the License. #include #include +#include "absl/algorithm/container.h" #include "absl/strings/str_format.h" #include "absl/strings/str_join.h" #include "absl/types/span.h" @@ -980,5 +981,25 @@ XLA_TEST_F(ReduceTest, OrReduceU64) { ComputeAndCompareR1(&builder, expected, {}); } +XLA_TEST_F(ReduceTest, R0ReduceInDisguise) { + XlaBuilder builder(TestName()); + XlaComputation add_f32 = CreateScalarAddComputation(F32, &builder); + constexpr int element_count = 127; + const Shape input_shape = ShapeUtil::MakeShape(F32, {element_count, 1}); + auto input = Parameter(&builder, 0, input_shape, "input"); + auto zero = ConstantR0(&builder, 0.0); + Reduce(input, zero, add_f32, /*dimensions_to_reduce=*/{0}); + + Array2D input_data(element_count, 1); + input_data.FillRandom(3.0f); + Literal input_literal = LiteralUtil::CreateR2FromArray2D(input_data); + std::unique_ptr input_global_data = + client_->TransferToServer(input_literal).ConsumeValueOrDie(); + + float expected = absl::c_accumulate(input_data, 0.0f); + ComputeAndCompareR1(&builder, {expected}, {input_global_data.get()}, + ErrorSpec(0.001)); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/tests/reduce_window_test.cc b/tensorflow/compiler/xla/tests/reduce_window_test.cc index 22fe4a2670e2e0e1fedc45036a1ceec19f44e42e..16c67d94c76bcf8984a2b3e4cb092026a6924aeb 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"); diff --git a/tensorflow/compiler/xla/tests/replay_test.cc b/tensorflow/compiler/xla/tests/replay_test.cc index 5cf87e565bf493167f5173588e7afa3b96282488..34c7dc7c46427b2d18ea21fc286ee03175f70800 100644 --- a/tensorflow/compiler/xla/tests/replay_test.cc +++ b/tensorflow/compiler/xla/tests/replay_test.cc @@ -55,7 +55,8 @@ TEST_F(ReplayTest, TwoPlusTwoReplay) { client_->GetComputationShape(computation).ConsumeValueOrDie(); std::unique_ptr replayed_shape = client_->GetComputationShape(replayed).ConsumeValueOrDie(); - ASSERT_TRUE(protobuf_util::ProtobufEquals(*original_shape, *replayed_shape)); + ASSERT_TRUE(protobuf_util::ProtobufEquals(original_shape->ToProto(), + replayed_shape->ToProto())); // Run it. Literal literal = @@ -87,7 +88,8 @@ XLA_TEST_F(ReplayTest, XPlusYReplayWithParameters) { client_->GetComputationShape(computation).ConsumeValueOrDie(); std::unique_ptr replayed_shape = client_->GetComputationShape(replayed).ConsumeValueOrDie(); - ASSERT_TRUE(protobuf_util::ProtobufEquals(*original_shape, *replayed_shape)); + ASSERT_TRUE(protobuf_util::ProtobufEquals(original_shape->ToProto(), + replayed_shape->ToProto())); // Run it. std::unique_ptr x_data = @@ -133,7 +135,8 @@ TEST_F(ReplayTest, MapPlusTwoOverR1) { client_->GetComputationShape(computation).ConsumeValueOrDie(); std::unique_ptr replayed_shape = client_->GetComputationShape(replayed).ConsumeValueOrDie(); - ASSERT_TRUE(protobuf_util::ProtobufEquals(*original_shape, *replayed_shape)); + ASSERT_TRUE(protobuf_util::ProtobufEquals(original_shape->ToProto(), + replayed_shape->ToProto())); // Run it. Literal literal = diff --git a/tensorflow/compiler/xla/tests/reshape_test.cc b/tensorflow/compiler/xla/tests/reshape_test.cc index dedc95b5ae8315185a35f786af42aad53bd7ad96..298136002e9ef47188e0bae95af3f596596e6062 100644 --- a/tensorflow/compiler/xla/tests/reshape_test.cc +++ b/tensorflow/compiler/xla/tests/reshape_test.cc @@ -618,7 +618,8 @@ XLA_TEST_P(ReshapeTest, R4Dim0MinorLayoutToR2Dim0MajorLayout) { ExecutionOptions execution_options = execution_options_; *execution_options.mutable_shape_with_output_layout() = ShapeUtil::MakeShapeWithLayout(use_bfloat16() ? BF16 : F32, {2, 8}, - {1, 0}); + {1, 0}) + .ToProto(); Literal actual = client_ ->ExecuteAndTransfer(computation, {input.get()}, &execution_options) @@ -767,7 +768,8 @@ XLA_TEST_P(ReshapeTest, NoopReshape) { ExecutionOptions execution_options = execution_options_; *execution_options.mutable_shape_with_output_layout() = ShapeUtil::MakeShapeWithLayout(use_bfloat16() ? BF16 : F32, {7, 2, 3, 5}, - {2, 3, 0, 1}); + {2, 3, 0, 1}) + .ToProto(); Literal output_literal = client_ ->ExecuteAndTransfer(computation, {input_data.get()}, diff --git a/tensorflow/compiler/xla/tests/round_trip_packed_literal_test.cc b/tensorflow/compiler/xla/tests/round_trip_packed_literal_test.cc index 091a5d2cacce6ac5bf986776e5ec96612d08cc75..606a099ecbc4e5677034c6d57e7ba5c398c69ab9 100644 --- a/tensorflow/compiler/xla/tests/round_trip_packed_literal_test.cc +++ b/tensorflow/compiler/xla/tests/round_trip_packed_literal_test.cc @@ -15,6 +15,7 @@ limitations under the License. #include +#include "absl/base/casts.h" #include "absl/types/span.h" #include "tensorflow/compiler/xla/client/global_data.h" #include "tensorflow/compiler/xla/client/local_client.h" @@ -27,7 +28,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/xla_data.pb.h" -#include "tensorflow/core/lib/core/casts.h" #include "tensorflow/core/platform/env.h" #include "tensorflow/core/platform/test.h" #include "tensorflow/core/platform/types.h" @@ -47,7 +47,7 @@ class RoundTripPackedLiteralTest : public ClientLibraryTestBase { TEST_F(RoundTripPackedLiteralTest, RoundTripsR1F32Length2) { string data(sizeof(float) * 2, 0); - absl::Span floats(tensorflow::bit_cast(data.data()), 2); + absl::Span floats(absl::bit_cast(data.data()), 2); floats[0] = 42.0; floats[1] = 24.0; @@ -69,7 +69,7 @@ TEST_F(RoundTripPackedLiteralTest, RoundTripsR1F32Length2) { TEST_F(RoundTripPackedLiteralTest, RoundTripsR2F32Size2x2Dim0Minor) { string data(sizeof(float) * 4, 0); - absl::Span floats(tensorflow::bit_cast(data.data()), 4); + absl::Span floats(absl::bit_cast(data.data()), 4); // With x as the minor dimension, these will become: floats[0] = 42.0; // y=0,x=0 floats[1] = 24.0; // y=0,x=1 @@ -102,7 +102,7 @@ TEST_F(RoundTripPackedLiteralTest, RoundTripsR2F32Size2x2Dim0Minor) { TEST_F(RoundTripPackedLiteralTest, RoundTripsR2F32Size2x2Dim1Minor) { string data(sizeof(float) * 4, 0); - absl::Span floats(tensorflow::bit_cast(data.data()), 4); + absl::Span floats(absl::bit_cast(data.data()), 4); // With y as the minor dimension, these will become: floats[0] = 42.0; // y=0,x=0 floats[1] = 24.0; // y=1,x=0 diff --git a/tensorflow/compiler/xla/tests/scatter_test.cc b/tensorflow/compiler/xla/tests/scatter_test.cc index 7e1f4aa0eb4801876d9bdbac6a4d7f1d09f81ba8..32de0fdf78f9c442e17c55e1b951e39122dac5ef 100644 --- a/tensorflow/compiler/xla/tests/scatter_test.cc +++ b/tensorflow/compiler/xla/tests/scatter_test.cc @@ -129,6 +129,42 @@ ENTRY main { RunTest(hlo_text, &operand, &scatter_indices, &updates); } +XLA_TEST_F(ScatterTest, TensorFlowScatterV2_InversePermutation) { + const char* hlo_text = R"( +HloModule TensorFlowScatterV2 + +update_s32 (lhs: s32[], rhs: s32[]) -> s32[] { + lhs = s32[] parameter(0) + ROOT rhs = s32[] parameter(1) +} + +ENTRY main { + permutation = s32[3,4] parameter(0) + reshape = s32[3,4,1] reshape(permutation) + operand = s32[3,4] iota(), iota_dimension=1 + updates = s32[3,4,1,1] iota(), iota_dimension=1 + iota = s32[3,4,1] iota(), iota_dimension=0 + indices = s32[3,4,2] concatenate(iota, reshape), dimensions={2} + ROOT scatter = s32[3,4] scatter(operand, indices, updates), + to_apply=update_s32, + update_window_dims={2,3}, + inserted_window_dims={}, + scatter_dims_to_operand_dims={0,1}, + index_vector_dim=2 +} +)"; + Literal permutation = + LiteralUtil::CreateR2({{1, 3, 2, 0}, {3, 0, 2, 1}, {2, 3, 1, 0}}); + HloModuleConfig config; + config.set_debug_options(GetDebugOptionsForTest()); + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr module, + ParseHloString(hlo_text, config)); + auto actual = ExecuteAndTransfer(std::move(module), {&permutation}); + Literal expected = + LiteralUtil::CreateR2({{3, 0, 2, 1}, {1, 3, 2, 0}, {3, 2, 0, 1}}); + EXPECT_TRUE(LiteralTestUtil::Equal(expected, actual)); +} + XLA_TEST_F(ScatterTest, SimpleR4) { const char* hlo_text = R"( HloModule SimpleR4 diff --git a/tensorflow/compiler/xla/tests/slice_test.cc b/tensorflow/compiler/xla/tests/slice_test.cc index 2cc33ab0963afe8ba2d8e9a6972dcf0622e27c48..3fb69419e735bfd9c5054673e0687f5139a410cb 100644 --- a/tensorflow/compiler/xla/tests/slice_test.cc +++ b/tensorflow/compiler/xla/tests/slice_test.cc @@ -166,6 +166,26 @@ TEST_F(SliceTest, SliceR4ThreeDimsMiddleMinor) { ComputeAndCompareR4(&builder, *expected, {}, ErrorSpec(0.000001)); } +TEST_F(SliceTest, SliceOfReshape) { + Array2D values(2 * 3 * 24, 7); + values.FillIota(1); + XlaBuilder builder(TestName()); + auto original = ConstantR2FromArray2D(&builder, values); + auto reshape = Reshape(original, {24, 3, 2, 7}); + Slice(reshape, {0, 0, 0, 0}, {11, 3, 2, 7}, {1, 1, 1, 1}); + ComputeAndCompare(&builder, {}); +} + +TEST_F(SliceTest, SliceOfCollapsingReshape) { + Array4D values(2, 3, 5, 7); + values.FillIota(1); + XlaBuilder builder(TestName()); + auto original = ConstantR4FromArray4D(&builder, values); + auto reshape = Reshape(original, {2 * 3 * 5, 7}); + Slice(reshape, {0, 0}, {4, 7}, {1, 1}); + ComputeAndCompare(&builder, {}); +} + XLA_TEST_F(SliceTest, StridedSliceR4WithOutputLayout) { Array4D values(2, 4, 6, 8); values.FillRandom(3.14f); @@ -253,7 +273,6 @@ XLA_TEST_P(SliceR1LargeTest, DoIt_S64) { Run(GetParam()); } XLA_TEST_P(SliceR1Test, DoIt_PRED) { Run(GetParam()); } - // Tests for R1 slice ops. // The format for each testcase is {input size, start, limit, stride}. // clang-format off diff --git a/tensorflow/compiler/xla/tests/test_macros.h b/tensorflow/compiler/xla/tests/test_macros.h index 7ca99a91635e85cd0888e59ecde31e47fec21844..80a6868485c9162d1cb0de24f0adf3f1c1d2503a 100644 --- a/tensorflow/compiler/xla/tests/test_macros.h +++ b/tensorflow/compiler/xla/tests/test_macros.h @@ -79,30 +79,28 @@ string PrependDisabledIfIndicated(const string& test_case_name, // heuristic to decide whether the test case should be disabled, and we // determine whether the test case should be disabled by resolving the (test // case name, test name) in a manifest file. -#define XLA_GTEST_TEST_(test_case_name, test_name, parent_class, parent_id) \ - class GTEST_TEST_CLASS_NAME_(test_case_name, test_name) \ - : public parent_class { \ - public: \ - GTEST_TEST_CLASS_NAME_(test_case_name, test_name)() {} \ - \ - private: \ - virtual void TestBody(); \ - static ::testing::TestInfo* const test_info_ GTEST_ATTRIBUTE_UNUSED_; \ - GTEST_DISALLOW_COPY_AND_ASSIGN_(GTEST_TEST_CLASS_NAME_(test_case_name, \ - test_name)); \ - }; \ - \ - ::testing::TestInfo* const GTEST_TEST_CLASS_NAME_(test_case_name, \ - test_name)::test_info_ = \ - ::testing::internal::MakeAndRegisterTestInfo( \ - #test_case_name, \ - ::xla::PrependDisabledIfIndicated(#test_case_name, #test_name) \ - .c_str(), \ - nullptr, nullptr, \ - ::testing::internal::CodeLocation(__FILE__, __LINE__), (parent_id), \ - parent_class::SetUpTestCase, parent_class::TearDownTestCase, \ - new ::testing::internal::TestFactoryImpl); \ +#define XLA_GTEST_TEST_(test_case_name, test_name, parent_class) \ + class GTEST_TEST_CLASS_NAME_(test_case_name, test_name) \ + : public parent_class { \ + public: \ + GTEST_TEST_CLASS_NAME_(test_case_name, test_name)() {} \ + \ + private: \ + virtual void TestBody(); \ + static ::testing::TestInfo* const test_info_ GTEST_ATTRIBUTE_UNUSED_; \ + GTEST_DISALLOW_COPY_AND_ASSIGN_(GTEST_TEST_CLASS_NAME_(test_case_name, \ + test_name)); \ + }; \ + \ + ::testing::TestInfo* const GTEST_TEST_CLASS_NAME_(test_case_name, \ + test_name)::test_info_ = \ + ::testing::RegisterTest( \ + #test_case_name, \ + ::xla::PrependDisabledIfIndicated(#test_case_name, #test_name) \ + .c_str(), \ + nullptr, nullptr, __FILE__, __LINE__, []() -> parent_class* { \ + return new GTEST_TEST_CLASS_NAME_(test_case_name, test_name)(); \ + }); \ void GTEST_TEST_CLASS_NAME_(test_case_name, test_name)::TestBody() // This is identical to the TEST_F macro from "gtest", but it potentially @@ -111,9 +109,8 @@ string PrependDisabledIfIndicated(const string& test_case_name, // Per usual, you can see what tests are available via --gunit_list_tests and // choose to run tests that have been disabled via the manifest via // --gunit_also_run_disabled_tests. -#define XLA_TEST_F(test_fixture, test_name) \ - XLA_GTEST_TEST_(test_fixture, test_name, test_fixture, \ - ::testing::internal::GetTypeId()) +#define XLA_TEST_F(test_fixture, test_name) \ + XLA_GTEST_TEST_(test_fixture, test_name, test_fixture) // Likewise, this is identical to the TEST_P macro from "gtest", but // potentially disables the test based on the DISABLED_MANIFEST file. diff --git a/tensorflow/compiler/xla/tests/test_utils.cc b/tensorflow/compiler/xla/tests/test_utils.cc index 2f18036ff4c5b0bfa28723fb181c33fa6995eb80..95c89b0ba6f29c453abab88e29bca13ee006455a 100644 --- a/tensorflow/compiler/xla/tests/test_utils.cc +++ b/tensorflow/compiler/xla/tests/test_utils.cc @@ -15,9 +15,11 @@ limitations under the License. #include +#include "absl/base/casts.h" #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" @@ -28,65 +30,113 @@ namespace xla { namespace { template -void PopulateWithRandomFloatingPointDataImpl(Literal* literal, - std::minstd_rand0* engine, - bool no_duplicates) { +void PopulateWithRandomFloatingPointData(Literal* literal, + std::minstd_rand0* engine) { + std::uniform_real_distribution generator(-0.1f, 0.2f); + for (FloatT& value : literal->data()) { + value = static_cast(generator(*engine)); + } +} + +template +void PopulateWithIntNext(Literal* literal); + +template <> +void PopulateWithIntNext(Literal* literal) { + // Duplicates may be generated if we don't have enough bits. + uint16 next_value = 0; + for (half& value : literal->data()) { + // Zero-out the MSB of the exponent to avoid Infs and NaNs, and put it into + // the sign bit. We could be less wasteful, but this is best-effort anyway. + uint16 exponent_msb = next_value & 0x4000; + value.x = (next_value & 0xBFFF) | (exponent_msb << 1); + next_value++; + } +} + +template <> +void PopulateWithIntNext(Literal* literal) { + // Duplicates may be generated if we don't have enough bits. + // Start at 0x80 rather than 0 to avoid denormals. + uint16 next_value = 0x80; + for (bfloat16& value : literal->data()) { + // Zero-out the MSB of the exponent to avoid Infs and NaNs, and put it into + // the sign bit. We could be less wasteful, but this is best-effort anyway. + uint16 exponent_msb = next_value & 0x4000; + value.value = (next_value & 0xBFFF) | (exponent_msb << 1); + next_value++; + } +} + +template +void PopulateWithNextAfter(Literal* literal) { + // Duplicates may be generated if the number of elements in the literal + // exceeds the number of positive values supported by the type. + float next_value = std::numeric_limits::min(); + for (float& value : literal->data()) { + value = next_value; + next_value = std::nextafter(next_value, std::numeric_limits::max()); + } +} + +template ::value || + std::is_same::value, + int>::type = 0> +void PopulateWithNoDuplicateData(Literal* literal, std::minstd_rand0* engine) { + PopulateWithIntNext(literal); + std::shuffle(literal->data().begin(), literal->data().end(), + *engine); +} + +template ::value && + !std::is_same::value, + int>::type = 0> +void PopulateWithNoDuplicateData(Literal* literal, std::minstd_rand0* engine) { + PopulateWithNextAfter(literal); + std::shuffle(literal->data().begin(), literal->data().end(), + *engine); +} + +template +void PopulateWithFloatingPointData(Literal* literal, std::minstd_rand0* engine, + bool no_duplicates) { CHECK(engine != nullptr); CHECK_EQ(literal->shape().element_type(), primitive_util::NativeToPrimitiveType()); if (no_duplicates) { - // Duplicates may be generated if the number of elements in the literal - // exceeds the number of positive values supported by the type. - FloatT next_value = std::numeric_limits::min(); - for (FloatT& value : literal->data()) { - value = next_value; - next_value = - std::nextafter(next_value, std::numeric_limits::max()); - } - std::shuffle(literal->data().begin(), literal->data().end(), - *engine); + PopulateWithNoDuplicateData(literal, engine); } else { - std::uniform_real_distribution generator(-0.1f, 0.2f); - for (FloatT& value : literal->data()) { - value = static_cast(generator(*engine)); - } + PopulateWithRandomFloatingPointData(literal, engine); } } -template -void PopulateWithRandomFloatingPointData(Literal* literal, +template <> +void PopulateWithFloatingPointData(Literal* literal, std::minstd_rand0* engine, bool no_duplicates) { CHECK(engine != nullptr); - PopulateWithRandomFloatingPointDataImpl(literal, engine, - no_duplicates); -} - -template <> -void PopulateWithRandomFloatingPointData(Literal* literal, - std::minstd_rand0* engine, - bool no_duplicates) { - // no_duplicates is ignored for half types. Unique values can only be - // generated for arrays with fewer than ~2**16 elements and no_duplicates is - // best-effort anyway. - CHECK(engine != nullptr); - std::uniform_real_distribution generator(-0.1f, 0.2f); - for (half& value : literal->data()) { - value = static_cast(generator(*engine)); + CHECK_EQ(literal->shape().element_type(), + primitive_util::NativeToPrimitiveType()); + if (no_duplicates) { + PopulateWithNoDuplicateData(literal, engine); + } else { + PopulateWithRandomFloatingPointData(literal, engine); } } template <> -void PopulateWithRandomFloatingPointData(Literal* literal, - std::minstd_rand0* engine, - bool no_duplicates) { - // no_duplicates is ignored for bfloat types. Unique values can only be - // generated for arrays with fewer than ~2**16 elements and no_duplicates is - // best-effort anyway. +void PopulateWithFloatingPointData(Literal* literal, + std::minstd_rand0* engine, + bool no_duplicates) { CHECK(engine != nullptr); - std::uniform_real_distribution generator(-0.1f, 0.2f); - for (bfloat16& value : literal->data()) { - value = static_cast(generator(*engine)); + CHECK_EQ(literal->shape().element_type(), + primitive_util::NativeToPrimitiveType()); + if (no_duplicates) { + PopulateWithNoDuplicateData(literal, engine); + } else { + PopulateWithRandomFloatingPointData(literal, engine); } } @@ -119,7 +169,7 @@ void PopulateWithRandomIntegralData(Literal* literal, std::minstd_rand0* engine, StatusOr MakeFakeLiteralInternal(const Shape& shape, std::minstd_rand0* engine, bool no_duplicates) { - if (ShapeUtil::IsTuple(shape)) { + if (shape.IsTuple()) { std::vector elements; for (const Shape& element_shape : shape.tuple_shapes()) { TF_ASSIGN_OR_RETURN( @@ -135,20 +185,16 @@ StatusOr MakeFakeLiteralInternal(const Shape& shape, Literal literal(shape); switch (shape.element_type()) { case BF16: - PopulateWithRandomFloatingPointData(&literal, engine, - no_duplicates); + PopulateWithFloatingPointData(&literal, engine, no_duplicates); break; case F16: - PopulateWithRandomFloatingPointData(&literal, engine, - no_duplicates); + PopulateWithFloatingPointData(&literal, engine, no_duplicates); break; case F32: - PopulateWithRandomFloatingPointData(&literal, engine, - no_duplicates); + PopulateWithFloatingPointData(&literal, engine, no_duplicates); break; case F64: - PopulateWithRandomFloatingPointData(&literal, engine, - no_duplicates); + PopulateWithFloatingPointData(&literal, engine, no_duplicates); break; case S8: PopulateWithRandomIntegralData(&literal, engine, no_duplicates); @@ -229,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 @@ -255,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 = @@ -291,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; @@ -303,19 +342,16 @@ StatusOr CreateLiteralForConstrainedUses( const Shape& slice_shape = use->opcode() == HloOpcode::kDynamicSlice ? use->shape() : use->operand(1)->shape(); - const int64 rank = ShapeUtil::Rank(indexed_shape); - 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; @@ -343,13 +379,14 @@ 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()) { - 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) { case ConstantType::kZero: @@ -414,8 +451,8 @@ Status VerifyHloModule(HloModule* const module, bool layout_sensitive, std::unique_ptr CreateCanonicalDot(const Shape& shape, HloInstruction* lhs, HloInstruction* rhs) { - CHECK_EQ(ShapeUtil::Rank(lhs->shape()), 2); - CHECK_EQ(ShapeUtil::Rank(rhs->shape()), 2); + CHECK_EQ(lhs->shape().rank(), 2); + CHECK_EQ(rhs->shape().rank(), 2); PrecisionConfig precision_config; precision_config.mutable_operand_precision()->Resize( 2, PrecisionConfig::DEFAULT); diff --git a/tensorflow/compiler/xla/tests/test_utils_test.cc b/tensorflow/compiler/xla/tests/test_utils_test.cc index bc433eac8fcb02087d8e4eb10f638c85dc141b22..591d6c19228a313f530cdae18f4be37e7b517601 100644 --- a/tensorflow/compiler/xla/tests/test_utils_test.cc +++ b/tensorflow/compiler/xla/tests/test_utils_test.cc @@ -15,13 +15,13 @@ limitations under the License. #include "tensorflow/compiler/xla/tests/test_utils.h" +#include "absl/base/casts.h" #include "absl/container/flat_hash_set.h" #include "tensorflow/compiler/xla/client/xla_builder.h" #include "tensorflow/compiler/xla/service/hlo_parser.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/tests/local_client_test_base.h" #include "tensorflow/compiler/xla/tests/test_macros.h" -#include "tensorflow/core/lib/core/casts.h" #include "tensorflow/core/lib/core/status_test_util.h" namespace xla { @@ -61,11 +61,11 @@ XLA_TEST_F(TestUtilsTest, Token) { R"(HloModule outfeed_module ENTRY InfeedToOutfeed { - token = token[] parameter(0) - infeed = ((u32[3]{0}, pred[]), token[]) infeed(token) + token0 = token[] parameter(0) + infeed = ((u32[3]{0}, pred[]), token[]) infeed(token0) infeed.data = (u32[3]{0}, pred[]) get-tuple-element(infeed), index=0 - outfeed = token[] outfeed(infeed.data, token) - ROOT infeed.1 = ((u32[3]{0}, pred[]), token[]) infeed(token) + outfeed = token[] outfeed(infeed.data, token0) + ROOT infeed.1 = ((u32[3]{0}, pred[]), token[]) infeed(token0) infeed.1.data = (u32[3]{0}, pred[]) get-tuple-element(infeed.1), index=0 infeed.1.token = token[] get-tuple-element(infeed.1), index=1 outfeed.1 = token[] outfeed(infeed.1.data, infeed.1.token) @@ -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) { @@ -148,7 +150,7 @@ ENTRY %sort.148.1589 (parameter.0: f32[1048576], parameter.1: s32[1048576]) -> ( absl::flat_hash_set key_set; for (const float& value : key_arg.data()) { - EXPECT_TRUE(key_set.insert(tensorflow::bit_cast(value)).second); + EXPECT_TRUE(key_set.insert(absl::bit_cast(value)).second); } } @@ -171,9 +173,60 @@ ENTRY %sort.148.1589 (parameter.0: s32[1048576], parameter.1: s32[1048576]) -> ( absl::flat_hash_set key_set; for (const int32& value : key_arg.data()) { - EXPECT_TRUE(key_set.insert(tensorflow::bit_cast(value)).second); + EXPECT_TRUE(key_set.insert(absl::bit_cast(value)).second); + } +} + +XLA_TEST_F(TestUtilsTest, NoDuplicatesBfloat16) { + // Inputs which are sort keys in key/value sorts should have no duplicates. + auto module = ParseHloString(R"( +HloModule sort, is_scheduled=true + +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} +} +)") + .ValueOrDie(); + TF_ASSERT_OK_AND_ASSIGN(std::vector args, + MakeFakeArguments(module.get())); + ASSERT_EQ(args.size(), 2); + const Literal& key_arg = args[0]; + + absl::flat_hash_set key_set; + for (const bfloat16& value : key_arg.data()) { + EXPECT_TRUE(key_set.insert(absl::bit_cast(value)).second); } } +XLA_TEST_F(TestUtilsTest, MakeFakeArgumentsR0InputToDynamicSlice) { + auto module = ParseHloString(R"( +HloModule Test + +ENTRY %module (parameter.0: s32[], parameter.1: f32[20,20]) -> f32[] { + %parameter.1 = f32[20,20]{1,0} parameter(1) + %constant.1 = s32[1]{0} constant({0}) + %parameter.0 = s32[] parameter(0) + %bitcast.3 = s32[1]{0} bitcast(s32[] %parameter.0) + %concatenate.1 = s32[2]{0} concatenate(s32[1]{0} %constant.1, s32[1]{0} %bitcast.3), dimensions={0} + %dynamic-slice.2 = f32[20,1]{1,0} dynamic-slice(f32[20,20]{1,0} %parameter.1, s32[2]{0} %concatenate.1), dynamic_slice_sizes={20,1} + %bitcast.4 = f32[20]{0} bitcast(f32[20,1]{1,0} %dynamic-slice.2) + %dynamic-slice.3 = f32[1]{0} dynamic-slice(f32[20]{0} %bitcast.4, s32[1]{0} %bitcast.3), dynamic_slice_sizes={1} + ROOT %bitcast.5 = f32[] bitcast(f32[1]{0} %dynamic-slice.3) +} +)") + .ValueOrDie(); + + TF_ASSERT_OK_AND_ASSIGN(std::vector args, + MakeFakeArguments(module.get())); + ASSERT_EQ(args.size(), 2); + EXPECT_TRUE(ShapeUtil::Equal(args[0].shape(), ShapeUtil::MakeShape(S32, {}))) + << ShapeUtil::HumanString(args[0].shape()); + EXPECT_TRUE( + ShapeUtil::Equal(args[1].shape(), ShapeUtil::MakeShape(F32, {20, 20}))) + << ShapeUtil::HumanString(args[1].shape()); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/tests/token_hlo_test.cc b/tensorflow/compiler/xla/tests/token_hlo_test.cc index b34fd0f2e873214c509533f29553af914ddc984d..b77cf38ed8e29973985406015c0a3936916ad5e6 100644 --- a/tensorflow/compiler/xla/tests/token_hlo_test.cc +++ b/tensorflow/compiler/xla/tests/token_hlo_test.cc @@ -16,6 +16,7 @@ limitations under the License. #include #include "absl/strings/str_cat.h" +#include "tensorflow/compiler/xla/service/hlo_parser.h" #include "tensorflow/compiler/xla/service/hlo_verifier.h" #include "tensorflow/compiler/xla/tests/hlo_test_base.h" #include "tensorflow/compiler/xla/tests/test_macros.h" @@ -28,7 +29,7 @@ namespace { class TokenHloTest : public HloTestBase {}; XLA_TEST_F(TokenHloTest, SingleTokenInstruction) { - std::unique_ptr module = CreateNewModule(); + std::unique_ptr module = CreateNewUnverifiedModule(); auto builder = HloComputation::Builder(TestName()); builder.AddInstruction(HloInstruction::CreateToken()); @@ -38,8 +39,22 @@ XLA_TEST_F(TokenHloTest, SingleTokenInstruction) { EXPECT_TRUE(LiteralTestUtil::Equal(result, LiteralUtil::CreateToken())); } +XLA_TEST_F(TokenHloTest, TokenInTuple) { + std::unique_ptr module = CreateNewUnverifiedModule(); + auto builder = HloComputation::Builder(TestName()); + auto token = builder.AddInstruction(HloInstruction::CreateToken()); + builder.AddInstruction(HloInstruction::CreateTuple({token})); + + module->AddEntryComputation(builder.Build()); + + TF_ASSERT_OK_AND_ASSIGN(Literal result, Execute(std::move(module), {})); + Literal token_literal = LiteralUtil::CreateToken(); + EXPECT_TRUE( + LiteralTestUtil::Equal(result, LiteralUtil::MakeTuple({&token_literal}))); +} + XLA_TEST_F(TokenHloTest, TokenTree) { - std::unique_ptr module = CreateNewModule(); + std::unique_ptr module = CreateNewUnverifiedModule(); auto builder = HloComputation::Builder(TestName()); auto token0 = builder.AddInstruction(HloInstruction::CreateToken()); auto token1 = builder.AddInstruction(HloInstruction::CreateToken()); @@ -54,7 +69,7 @@ XLA_TEST_F(TokenHloTest, TokenTree) { } XLA_TEST_F(TokenHloTest, InvalidTokenShapedEntryParameter) { - std::unique_ptr module = CreateNewModule(); + std::unique_ptr module = CreateNewUnverifiedModule(); auto builder = HloComputation::Builder(TestName()); builder.AddInstruction( HloInstruction::CreateParameter(0, ShapeUtil::MakeShape(F32, {}), "p0")); @@ -75,7 +90,7 @@ XLA_TEST_F(TokenHloTest, InvalidTokenShapedEntryParameter) { } XLA_TEST_F(TokenHloTest, InvalidTupleTokenShapedEntryParameter) { - std::unique_ptr module = CreateNewModule(); + std::unique_ptr module = CreateNewUnverifiedModule(); auto builder = HloComputation::Builder(TestName()); builder.AddInstruction(HloInstruction::CreateParameter( 0, @@ -94,26 +109,6 @@ XLA_TEST_F(TokenHloTest, InvalidTupleTokenShapedEntryParameter) { ::testing::HasSubstr("Entry parameter 0 is or contains a token shape")); } -XLA_TEST_F(TokenHloTest, InvalidOperandToTokenInstruction) { - std::unique_ptr module = CreateNewModule(); - auto builder = HloComputation::Builder(TestName()); - auto param = builder.AddInstruction( - HloInstruction::CreateParameter(0, ShapeUtil::MakeShape(F32, {}), "p0")); - builder.AddInstruction(HloInstruction::CreateAfterAll({param})); - builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(123))); - module->AddEntryComputation(builder.Build()); - - Status status = - HloVerifier(/*layout_sensitive=*/false, /*allow_mixed_precision=*/false) - .Run(module.get()) - .status(); - ASSERT_IS_NOT_OK(status); - EXPECT_THAT(status.error_message(), - ::testing::HasSubstr( - "Operands of token instructions must be TOKEN types")); -} - XLA_TEST_F(TokenHloTest, TokenInWhileLoop) { // Thread a token around a while loop. Token is created and consumed by a // AfterAll instruction in the while body. @@ -206,5 +201,95 @@ ENTRY %TokenInConditional (param.3: pred[]) -> s32[] { } } +XLA_TEST_F(TokenHloTest, AddDependency) { + string module_string = R"( +HloModule AddDependency, is_scheduled=true + +// Computes (p0 + 42) * (-p1) +// where there is a dependency from the add to the negation using a token +// with after-all and add-dependency instructions. +ENTRY %AddDependency (p0: f32[], p1: f32[]) -> f32[] { + %p0 = f32[] parameter(0) + %p1 = f32[] parameter(1) + + %forty_two = f32[] constant(42.0) + %add = f32[] add(f32[] %p0, f32[] %forty_two) + %token0 = token[] after-all(f32[] %add) + %p1_after_token = f32[] add-dependency(f32[] %p1, token[] %token0) + %neg = f32[] negate(f32[] %p1_after_token) + ROOT %product = f32[] multiply(f32[] %add, f32[] %neg) +} +)"; + TF_ASSERT_OK_AND_ASSIGN( + std::unique_ptr module, + ParseHloString(module_string, GetModuleConfigForTest())); + auto p0 = LiteralUtil::CreateR0(10.0); + auto p1 = LiteralUtil::CreateR0(3.0); + auto expected = LiteralUtil::CreateR0(-156.0); + EXPECT_EQ(expected, ExecuteNoHloPasses(std::move(module), {&p0, &p1})); +} + +XLA_TEST_F(TokenHloTest, AddDependencyOfConstant) { + string module_string = R"( +HloModule AddDependencyOfConstant, is_scheduled=true + +ENTRY %AddDependency (p0: f32[]) -> f32[] { + %p0 = f32[] parameter(0) + %forty_two = f32[] constant(42.0) + %token0 = token[] after-all(f32[] %p0) + %forty_two_after_token = f32[] add-dependency(f32[] %forty_two, token[] %token0) + ROOT %product = f32[] multiply(f32[] %p0, f32[] %forty_two_after_token) +} +)"; + TF_ASSERT_OK_AND_ASSIGN( + std::unique_ptr module, + ParseHloString(module_string, GetModuleConfigForTest())); + auto p0 = LiteralUtil::CreateR0(10.0); + auto expected = LiteralUtil::CreateR0(420.0); + EXPECT_EQ(expected, ExecuteNoHloPasses(std::move(module), {&p0})); +} + +XLA_TEST_F(TokenHloTest, AddDependencyAsRoot) { + string module_string = R"( +HloModule AddDependencyAsRoot, is_scheduled=true +ENTRY %AddDependency (p: f32[3]) -> f32[3] { + %p = f32[3] parameter(0) + %neg = f32[3] negate(f32[3] %p) + %token0 = token[] after-all() + ROOT %add_dep = f32[3] add-dependency(f32[3] %neg, token[] %token0) +} +)"; + TF_ASSERT_OK_AND_ASSIGN( + std::unique_ptr module, + ParseHloString(module_string, GetModuleConfigForTest())); + auto input = LiteralUtil::CreateR1({1.0, 3.0, 7.0}); + auto expected = LiteralUtil::CreateR1({-1.0, -3.0, -7.0}); + EXPECT_EQ(expected, ExecuteNoHloPasses(std::move(module), {&input})); +} + +XLA_TEST_F(TokenHloTest, TupleShapedAddDependency) { + string module_string = R"( +HloModule TupleShapedAddDependency, is_scheduled=true +ENTRY %TupleShapedAddDependency (p0: f32[3], p1: f32[3]) -> f32[3] { + %p0 = f32[3] parameter(0) + %p1 = f32[3] parameter(1) + %forty_two = f32[] constant(42.0) + %token0 = token[] after-all() + %tuple = (f32[3], token[], f32[3], f32[]) tuple(f32[3] %p0, token[] %token0, f32[3] %p1, f32[] %forty_two) + %add_dep = (f32[3], token[], f32[3], f32[]) add-dependency((f32[3], token[], f32[3], f32[]) %tuple, token[] %token0) + %elem0 = f32[3] get-tuple-element((f32[3], token[], f32[3], f32[]) %add_dep), index=0 + %elem2 = f32[3] get-tuple-element((f32[3], token[], f32[3], f32[]) %add_dep), index=2 + ROOT %diff = f32[3] subtract(f32[3] %elem0, f32[3] %elem2) +} +)"; + TF_ASSERT_OK_AND_ASSIGN( + std::unique_ptr module, + ParseHloString(module_string, GetModuleConfigForTest())); + auto p0 = LiteralUtil::CreateR1({3.0, 3.0, 47.0}); + auto p1 = LiteralUtil::CreateR1({1.0, -2.0, 2.0}); + auto expected = LiteralUtil::CreateR1({2.0, 5.0, 45.0}); + EXPECT_EQ(expected, ExecuteNoHloPasses(std::move(module), {&p0, &p1})); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/tests/tuple_test.cc b/tensorflow/compiler/xla/tests/tuple_test.cc index 619d2a388b5646c31f0a61f709a2ab3067e39c03..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,10 +513,9 @@ 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 + HloModule m, is_scheduled=true ENTRY test { name.1 = (f32[3]{0}) parameter(0) @@ -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}); @@ -555,13 +553,11 @@ XLA_TEST_F(TupleHloTest, s = (f32[2],f32[2]) tuple-select(cond, tup0, tup1) gte = f32[2] get-tuple-element(s), index=0 tuple = (f32[2]) tuple(gte) - token = token[] after-all() - ROOT outfeed = token[] outfeed(tuple, token) + token0 = token[] after-all() + 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/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 a6e70eb6ca25ffac24a8ebaf0420238e109e4fad..c7337e8caae8f2ee25f4b25dc22439e08d2ecc25 100644 --- a/tensorflow/compiler/xla/tests/xla_hlo_profile_test.cc +++ b/tensorflow/compiler/xla/tests/xla_hlo_profile_test.cc @@ -91,16 +91,16 @@ Status ParseOneProfileOutputLine( string match_usecs = "([0-9.]+) usec"; string match_flops = "([^ ]*)"; string match_trops = "([^ ]*)"; - string match_bytes_per_sec = "([0-9.TGMKi]+)B/s"; - string match_bytes_per_cycle = "([0-9.TGMKi]+)B/cycle"; + string match_bytes_per_sec = "([0-9.TGMKi]*)(?:B/s)?"; + string match_bytes_per_cycle = "([0-9.TGMKi]*)(?:B/cycle)?"; // The underlined part is what we're trying to match with match_opcode: // // %dot33 = f32[256,256]{1,0} dot(...) // ^^^ - string match_opcode = - expect_hlo ? "%[^=]+= [^ ]+ ([^(]+)\\(.*" : "(\\[total\\])"; + string match_opcode = expect_hlo ? "%[^=]+= [^ ]+ ([^(]+)\\(.*" + : "(\\[total\\])( \\[entry\\])?"; string regexp_pattern = absl::StrCat( " +", match_cycles, separator, match_usecs, separator, match_flops, separator, match_trops, separator, match_bytes_per_sec, separator, @@ -125,6 +125,10 @@ Status ParseOneProfileOutputLine( return Status::OK(); } +bool IsExtraMetricProfileOutputLine(const string& line) { + return RE2::FullMatch(line, "Extra metric \\S+: \\d+"); +} + // Returns void so that we can ASSERT. void ExecuteAndFetchProfile(string* profile_output, LocalClient* client, const XlaComputation& computation, @@ -153,10 +157,12 @@ void ExecuteAndFetchProfile(string* profile_output, LocalClient* client, TF_ASSERT_OK(transfer_manager->TransferLiteralToDevice( stream_ptr.get(), Literal::CreateFromShape(rhs_arg_shape), rhs_arg)); + ExecutableBuildOptions build_options; + build_options.mutable_debug_options()->set_xla_hlo_profile(true); TF_ASSERT_OK_AND_ASSIGN( std::unique_ptr local_executable, client->Compile(computation, {&lhs_arg_shape, &rhs_arg_shape}, - ExecutableBuildOptions().set_hlo_profile(true))); + build_options)); Executable* executable = local_executable->executable(); HloExecutionProfile hlo_execution_profile( @@ -168,9 +174,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, @@ -204,20 +209,35 @@ XLA_TEST_F(HloProfileTest, ProfileSingleComputation) { string profile_output; ExecuteAndFetchProfile(&profile_output, client, computation, lhs_shape, rhs_shape); - + VLOG(4) << "Profile Output:\n" << profile_output; std::vector profile_output_lines = absl::StrSplit(profile_output, '\n'); absl::flat_hash_map parsed_profile_lines; - TF_ASSERT_OK(ParseOneProfileOutputLine( - profile_output_lines[1], /*expect_hlo=*/false, &parsed_profile_lines)); + int line_no = 0; + + // Skip extra metrics. + while (IsExtraMetricProfileOutputLine(profile_output_lines[line_no])) { + line_no++; + } + + line_no++; // Skip 'Execution profile for ....' - TF_ASSERT_OK(ParseOneProfileOutputLine( - profile_output_lines[2], /*expect_hlo=*/true, &parsed_profile_lines)); + ASSERT_LT(line_no, profile_output_lines.size()); + TF_ASSERT_OK(ParseOneProfileOutputLine(profile_output_lines[line_no++], + /*expect_hlo=*/false, + &parsed_profile_lines)); - TF_ASSERT_OK(ParseOneProfileOutputLine( - profile_output_lines[3], /*expect_hlo=*/true, &parsed_profile_lines)); + ASSERT_LT(line_no, profile_output_lines.size()); + TF_ASSERT_OK(ParseOneProfileOutputLine(profile_output_lines[line_no++], + /*expect_hlo=*/true, + &parsed_profile_lines)); + + ASSERT_LT(line_no, profile_output_lines.size()); + TF_ASSERT_OK(ParseOneProfileOutputLine(profile_output_lines[line_no++], + /*expect_hlo=*/true, + &parsed_profile_lines)); TF_ASSERT_OK_AND_ASSIGN(ParsedProfileOutputLine total_profile, MaybeFind(parsed_profile_lines, "[total]")); @@ -291,6 +311,7 @@ XLA_TEST_F(HloProfileTest, ProfileWhileComputation) { string profile_output; ExecuteAndFetchProfile(&profile_output, client, computation, matrix_shape, matrix_shape); + SCOPED_TRACE(profile_output); std::vector profile_output_lines = absl::StrSplit(profile_output, '\n'); @@ -302,14 +323,13 @@ XLA_TEST_F(HloProfileTest, ProfileWhileComputation) { ASSERT_NE(while_body_profile_start, profile_output_lines.cend()); - auto while_body_profile_end = std::find_if( - while_body_profile_start, profile_output_lines.end(), - [](absl::string_view s) { - return absl::StartsWith(s, "********** microseconds report **********"); - }); + auto while_body_profile_end = + std::find_if(while_body_profile_start, profile_output_lines.end(), + [](absl::string_view s) { + return absl::StartsWith(s, "********** microseconds "); + }); - // We emit a blank line before the "********** microseconds report **********" - // line. + // We emit a blank line before the "microseconds report" line. while_body_profile_end--; ASSERT_NE(while_body_profile_end, profile_output_lines.end()); @@ -364,7 +384,7 @@ static std::pair AddXlaHloProfileFlag(int argc, char** argv) { GTEST_API_ int main(int argc, char** argv) { std::vector flag_list; - xla::legacy_flags::AppendDebugOptionsFlags(&flag_list); + xla::AppendDebugOptionsFlags(&flag_list); std::tie(argc, argv) = AddXlaHloProfileFlag(argc, argv); auto usage = tensorflow::Flags::Usage(argv[0], flag_list); diff --git a/tensorflow/compiler/xla/tests/xla_internal_test_main.cc b/tensorflow/compiler/xla/tests/xla_internal_test_main.cc index 15603619b62d8f45cdce97ac7d83924a78f88cf3..dca0aa52a533130372759156a3238f1a3b10ca42 100644 --- a/tensorflow/compiler/xla/tests/xla_internal_test_main.cc +++ b/tensorflow/compiler/xla/tests/xla_internal_test_main.cc @@ -15,14 +15,14 @@ limitations under the License. #include "absl/strings/match.h" #include "absl/strings/string_view.h" -#include "tensorflow/compiler/xla/legacy_flags/debug_options_flags.h" +#include "tensorflow/compiler/xla/debug_options_flags.h" #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/test.h" #include "tensorflow/core/platform/test_benchmark.h" GTEST_API_ int main(int argc, char** argv) { std::vector flag_list; - xla::legacy_flags::AppendDebugOptionsFlags(&flag_list); + xla::AppendDebugOptionsFlags(&flag_list); auto usage = tensorflow::Flags::Usage(argv[0], flag_list); if (!tensorflow::Flags::Parse(&argc, argv, flag_list)) { LOG(ERROR) << "\n" << usage; @@ -49,7 +49,7 @@ GTEST_API_ int main(int argc, char** argv) { // different API than Tensorflow's. testing::InitGoogleTest(&argc, argv); #if defined(PLATFORM_GOOGLE) - base::SetFlag(&FLAGS_benchmarks, pattern); + absl::SetFlag(&FLAGS_benchmarks, pattern); RunSpecifiedBenchmarks(); #else tensorflow::testing::Benchmark::Run(pattern); diff --git a/tensorflow/compiler/xla/text_literal_reader.cc b/tensorflow/compiler/xla/text_literal_reader.cc index cdde88c1359416d423685f330e9cbdf77948034f..c78ec522aa5f13556c6d4602267544694887f250 100644 --- a/tensorflow/compiler/xla/text_literal_reader.cc +++ b/tensorflow/compiler/xla/text_literal_reader.cc @@ -27,6 +27,7 @@ limitations under the License. #include "absl/strings/string_view.h" #include "absl/strings/strip.h" #include "tensorflow/compiler/xla/literal.h" +#include "tensorflow/compiler/xla/service/hlo_parser.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/types.h" @@ -66,7 +67,7 @@ StatusOr TextLiteralReader::ReadAllLines() { } absl::StripAsciiWhitespace(&shape_string); - TF_ASSIGN_OR_RETURN(Shape shape, ShapeUtil::ParseShapeString(shape_string)); + TF_ASSIGN_OR_RETURN(Shape shape, ParseShape(shape_string)); if (shape.element_type() != F32) { return Unimplemented( "unsupported element type for text literal reading: %s", diff --git a/tensorflow/compiler/xla/tools/BUILD b/tensorflow/compiler/xla/tools/BUILD index 3a086c66bbb37965b1ad7c83a93f0054ae723e87..52fee4770ab940741723514d742e998b25765f24 100644 --- a/tensorflow/compiler/xla/tools/BUILD +++ b/tensorflow/compiler/xla/tools/BUILD @@ -14,7 +14,7 @@ filegroup( visibility = ["//tensorflow/compiler/xla:internal"], ) -load("//tensorflow:tensorflow.bzl", "tf_cc_binary") +load("//tensorflow:tensorflow.bzl", "tf_cc_binary", "tf_cc_test") tf_cc_binary( name = "hex_floats_to_packed_literal", @@ -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: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/legacy_flags:debug_options_flags", - "//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"], @@ -78,6 +51,7 @@ cc_library( name = "replay_computation_library", srcs = ["replay_computation.cc"], deps = [ + "//tensorflow/compiler/xla:debug_options_flags", "//tensorflow/compiler/xla:execution_options_util", "//tensorflow/compiler/xla:literal", "//tensorflow/compiler/xla:shape_util", @@ -91,10 +65,10 @@ cc_library( "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/client:xla_computation", "//tensorflow/compiler/xla/client/lib:testing", - "//tensorflow/compiler/xla/legacy_flags:debug_options_flags", "//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", @@ -207,13 +181,13 @@ tf_cc_binary( name = "dumped_computation_to_tf_graphdef", srcs = ["dumped_computation_to_tf_graphdef.cc"], deps = [ + "//tensorflow/compiler/xla:debug_options_flags", "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla:types", "//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/legacy_flags:debug_options_flags", "//tensorflow/compiler/xla/service", "//tensorflow/compiler/xla/service:hlo_graph_dumper", "//tensorflow/compiler/xla/service:hlo_proto", @@ -234,3 +208,56 @@ tf_cc_binary( "//tensorflow/core:lib", ], ) + +tf_cc_test( + name = "hlo_extractor_test", + srcs = ["hlo_extractor_test.cc"], + deps = [ + ":hlo_extractor", + "//tensorflow/compiler/xla/service:hlo_matchers", + "//tensorflow/compiler/xla/tests:hlo_test_base", + "//tensorflow/core:test", + ], +) + +cc_library( + name = "hlo_extractor", + srcs = ["hlo_extractor.cc"], + hdrs = ["hlo_extractor.h"], + deps = [ + "//tensorflow/compiler/xla:status", + "//tensorflow/compiler/xla/service:hlo", + "//tensorflow/compiler/xla/service:hlo_verifier", + "@com_google_absl//absl/container:flat_hash_map", + "@com_google_absl//absl/container:flat_hash_set", + "@com_google_absl//absl/memory", + ], +) + +tf_cc_binary( + name = "interactive_graphviz", + srcs = ["interactive_graphviz.cc"], + deps = [ + ":hlo_extractor", + "//tensorflow/compiler/xla/client:client_library", + "//tensorflow/compiler/xla/client:local_client", + "//tensorflow/compiler/xla/service:compiler", + "//tensorflow/compiler/xla/service:cpu_plugin", + "//tensorflow/compiler/xla/service:gpu_plugin", + "//tensorflow/compiler/xla/service:hlo_graph_dumper", + "//tensorflow/compiler/xla/service:hlo_proto", + "//tensorflow/compiler/xla/service:hlo_runner", + "//tensorflow/compiler/xla/service:local_service", + "//tensorflow/compiler/xla/service:platform_util", + "//tensorflow/core:framework_internal", + "//tensorflow/core:lib", + "@com_google_absl//absl/algorithm:container", + "@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 c866a13de7543fc948311f94708bc6b904717b62..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/legacy_flags/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 = legacy_flags::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::legacy_flags::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_tf_graphdef.cc b/tensorflow/compiler/xla/tools/dumped_computation_to_tf_graphdef.cc index 07ef5ff656bb48519a700a1d7d6c60b655a40ed6..f8bb9a6b1e217fc4e6e15c8a3302be61ed339c82 100644 --- a/tensorflow/compiler/xla/tools/dumped_computation_to_tf_graphdef.cc +++ b/tensorflow/compiler/xla/tools/dumped_computation_to_tf_graphdef.cc @@ -31,7 +31,7 @@ limitations under the License. #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/legacy_flags/debug_options_flags.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" @@ -53,7 +53,7 @@ void RealMain(absl::Span args) { tensorflow::ReadBinaryProto(tensorflow::Env::Default(), arg, &module)); XlaComputation computation = client->LoadSnapshot(module).ConsumeValueOrDie(); - DebugOptions debug_options = legacy_flags::GetDebugOptionsFromFlags(); + DebugOptions debug_options = GetDebugOptionsFromFlags(); debug_options.set_xla_generate_hlo_graph(".*"); debug_options.set_xla_hlo_dump_as_graphdef(true); ComputationStats stats = @@ -68,7 +68,7 @@ void RealMain(absl::Span args) { int main(int argc, char** argv) { std::vector flag_list; - xla::legacy_flags::AppendDebugOptionsFlags(&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) { diff --git a/tensorflow/compiler/xla/tools/hlo_extractor.cc b/tensorflow/compiler/xla/tools/hlo_extractor.cc new file mode 100644 index 0000000000000000000000000000000000000000..f3ce5f99b0c2a8e9ae5446f4bedc34b678c95b96 --- /dev/null +++ b/tensorflow/compiler/xla/tools/hlo_extractor.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/compiler/xla/tools/hlo_extractor.h" + +#include +#include + +#include "absl/container/flat_hash_map.h" +#include "absl/container/flat_hash_set.h" +#include "absl/memory/memory.h" +#include "tensorflow/compiler/xla/service/hlo_clone_context.h" +#include "tensorflow/compiler/xla/service/hlo_computation.h" +#include "tensorflow/compiler/xla/service/hlo_verifier.h" +#include "tensorflow/compiler/xla/status.h" + +namespace xla { +namespace { + +// Visitor that build a new HLO module with an entry computation and a root that +// is provided to the visit function. Only HLOs that are reachable from the new +// root instruction are included in the new module. +// +// The constructor allows specifying a set of boundary HLOs to prune the HLO +// graph. HLOs at the boundary are replaced with parameters. Can be nullptr +// which means no boundary, i.e. no HLOs are replaced with parameters. +class ExtractionVisitor : public ConstDfsHloVisitorWithDefault { + public: + explicit ExtractionVisitor( + const HloModule& old_module, + absl::flat_hash_set* boundary) + : old_module_(old_module), + module_(absl::make_unique("extracted", config_)), + clone_context_(module_.get()), + builder_("entry_computation"), + boundary_(boundary) {} + + Status HandleParameter(const HloInstruction* parameter) override { + // Entry parameters need renumbering. + auto new_parameter = HloInstruction::CreateParameter( + parameter_number_++, parameter->shape(), parameter->name()); + clone_context_.MapInstruction(parameter, new_parameter.get()); + builder_.AddInstruction(std::move(new_parameter)); + return Status::OK(); + } + + Status DefaultAction(const HloInstruction* hlo) override { + // Replace instructions at the boundary with parameters, but leave constants + // untouched. + if (boundary_ != nullptr && boundary_->count(hlo) > 0) { + auto new_parameter = HloInstruction::CreateParameter( + parameter_number_, hlo->shape(), hlo->name()); + parameter_number_++; + clone_context_.MapInstruction(hlo, new_parameter.get()); + builder_.AddInstruction(std::move(new_parameter)); + return Status::OK(); + } + std::vector new_operands; + for (auto operand : hlo->operands()) { + new_operands.push_back(clone_context_.GetInstruction(operand)); + } + auto instruction = + hlo->CloneWithNewOperands(hlo->shape(), new_operands, &clone_context_); + builder_.AddInstruction(std::move(instruction)); + return Status::OK(); + } + + Status FinishVisit(const HloInstruction* /*root*/) override { + module_->AddEntryComputation(builder_.Build()); + // Rename HLOs so that their name matches the original. By default, + // HLOs get new unique names when adding a new entry computation to + // a module. + for (auto computation : old_module_.MakeComputationPostOrder()) { + for (auto old_instruction : computation->MakeInstructionPostOrder()) { + if (auto new_instruction = + clone_context_.FindInstruction(old_instruction)) { + new_instruction->SetAndSanitizeName(old_instruction->name()); + } + } + } + return Status::OK(); + } + + HloModule* module() { return module_.get(); } + + std::unique_ptr ConsumeModule() { return std::move(module_); } + + private: + const HloModule& old_module_; + HloModuleConfig config_; + std::unique_ptr module_; + HloCloneContext clone_context_; + HloComputation::Builder builder_; + absl::flat_hash_set* boundary_; + int64 parameter_number_ = 0; +}; + +void ComputeBoundary(const HloInstruction* root, int64 limit, + absl::flat_hash_set* boundary) { + std::deque worklist; + absl::flat_hash_map visited; + worklist.push_back(root); + visited.emplace(root, 0); + while (!worklist.empty()) { + auto hlo = worklist.front(); + worklist.pop_front(); + int64 hops = visited[hlo]; + if (hops > limit) { + boundary->insert(hlo); + continue; + } + for (const HloInstruction* operand : hlo->operands()) { + if (visited.count(operand)) { + continue; + } + worklist.push_back(operand); + visited.emplace(operand, hops + 1); + } + } +} + +} // namespace + +std::unique_ptr ExtractModule(HloInstruction* instruction, + int64 height) { + absl::flat_hash_set boundary; + if (height != -1) { + ComputeBoundary(instruction, height, &boundary); + } + ExtractionVisitor visitor(*instruction->GetModule(), &boundary); + CHECK(instruction->Accept(&visitor).ok()); + + // The first pass may leave unused parameter instructions. Do another + // extraction pass to remove unused parameters. This is done because + // HloComputation does not allow removing parameters after the computation has + // been built. + ExtractionVisitor cleanup_visitor(*visitor.module(), /*boundary=*/nullptr); + TF_CHECK_OK(visitor.module()->entry_computation()->root_instruction()->Accept( + &cleanup_visitor)); + + HloVerifier verifier(/*layout_sensitive=*/false, + /*allow_mixed_precision=*/true); + TF_CHECK_OK(verifier.Run(cleanup_visitor.module()).status()); + return cleanup_visitor.ConsumeModule(); +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/tools/hlo_extractor.h b/tensorflow/compiler/xla/tools/hlo_extractor.h new file mode 100644 index 0000000000000000000000000000000000000000..bc13dc7e438fe0e64312746150af02df805e746a --- /dev/null +++ b/tensorflow/compiler/xla/tools/hlo_extractor.h @@ -0,0 +1,36 @@ +/* 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_TOOLS_HLO_EXTRACTOR_H_ +#define TENSORFLOW_COMPILER_XLA_TOOLS_HLO_EXTRACTOR_H_ + +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/service/hlo_module.h" + +namespace xla { + +// Creates a new HLO module rooted with an entry computation rooted at the given +// instruction. +// +// By default (height == -1), the new computation includes all transitive +// operands of `root`. If you specify a different height, the new computation +// will include all instructions <= `height` hops away from `root`. +// Instructions at the boundary are replaced by parameters. +std::unique_ptr ExtractModule(HloInstruction* instruction, + int64 height = -1); + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_TOOLS_HLO_EXTRACTOR_H_ diff --git a/tensorflow/compiler/xla/tools/hlo_extractor_test.cc b/tensorflow/compiler/xla/tools/hlo_extractor_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..4beb099b330cadf4540944979f38681bae07103c --- /dev/null +++ b/tensorflow/compiler/xla/tools/hlo_extractor_test.cc @@ -0,0 +1,139 @@ +/* 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/tools/hlo_extractor.h" + +#include "tensorflow/compiler/xla/service/hlo_matchers.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" + +namespace xla { +namespace { + +namespace op = testing::opcode_matchers; + +using HloExtractorTest = HloTestBase; + +TEST_F(HloExtractorTest, ExtractTopLevel) { + const string& hlo_string = R"( +HloModule test + +ENTRY %entry { + param.0 = f32[4]{0} parameter(0) + negate = f32[4]{0} negate(f32[4]{0} param.0) + ROOT exp = f32[4]{0} exponential(f32[4]{0} negate) +} +)"; + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr hlo_module, + ParseAndReturnVerifiedModule(hlo_string)); + + { + auto extracted_module = + ExtractModule(FindInstruction(hlo_module.get(), "exp")); + EXPECT_THAT(extracted_module->entry_computation()->root_instruction(), + op::Exp(op::Negate(op::Parameter(0)))); + } + + { + auto extracted_module = + ExtractModule(FindInstruction(hlo_module.get(), "exp"), /*height=*/0); + EXPECT_THAT(extracted_module->entry_computation()->root_instruction(), + op::Exp(op::Parameter(0))); + } + + { + auto extracted_module = ExtractModule( + FindInstruction(hlo_module.get(), "negate"), /*height=*/0); + EXPECT_THAT(extracted_module->entry_computation()->root_instruction(), + op::Negate(op::Parameter(0))); + } +} + +TEST_F(HloExtractorTest, ExtractDag) { + const string& hlo_string = R"( +HloModule test + +ENTRY %entry { + param.0 = f32[4]{0} parameter(0) + tanh = f32[4]{0} tanh(f32[4]{0} param.0) + negate = f32[4]{0} negate(f32[4]{0} tanh) + exp = f32[4]{0} exponential(f32[4]{0} negate) + ROOT add = f32[4]{0} add(f32[4]{0} negate, f32[4]{0} exp) +} +)"; + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr hlo_module, + ParseAndReturnVerifiedModule(hlo_string)); + + { + auto extracted_module = + ExtractModule(FindInstruction(hlo_module.get(), "exp")); + EXPECT_THAT(extracted_module->entry_computation()->root_instruction(), + op::Exp(op::Negate(op::Tanh(op::Parameter(0))))); + } + + { + auto extracted_module = + ExtractModule(FindInstruction(hlo_module.get(), "add"), /*height=*/0); + EXPECT_THAT(extracted_module->entry_computation()->root_instruction(), + op::Add(op::Parameter(0), op::Parameter(1))); + } + { + auto extracted_module = + ExtractModule(FindInstruction(hlo_module.get(), "add"), /*height=*/1); + EXPECT_THAT(extracted_module->entry_computation()->root_instruction(), + op::Add(op::Negate(op::Parameter(0)), + op::Exp(op::Negate(op::Parameter(0))))); + } + { + auto extracted_module = + ExtractModule(FindInstruction(hlo_module.get(), "add"), /*height=*/2); + EXPECT_THAT(extracted_module->entry_computation()->root_instruction(), + op::Add(op::Negate(op::Tanh(op::Parameter(0))), + op::Exp(op::Negate(op::Tanh(op::Parameter(0)))))); + } +} + +TEST_F(HloExtractorTest, ExtractWithConstant) { + const string& hlo_string = R"( +HloModule test + +ENTRY %entry { + p = f32[4]{0} parameter(0) + tanh = f32[4]{0} tanh(p) + c = f32[4]{0} constant({1, 2, 3, 4}) + ROOT add = f32[4]{0} add(tanh, c) +} +)"; + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr hlo_module, + ParseAndReturnVerifiedModule(hlo_string)); + + { + auto extracted_module = + ExtractModule(FindInstruction(hlo_module.get(), "add"), /*height=*/0); + EXPECT_THAT(extracted_module->entry_computation()->root_instruction(), + op::Add(op::Parameter(0), op::Parameter(1))); + } + { + auto extracted_module = + ExtractModule(FindInstruction(hlo_module.get(), "add"), /*height=*/1); + EXPECT_THAT(extracted_module->entry_computation()->root_instruction(), + op::Add(op::Tanh(op::Parameter(0)), op::Constant())); + } +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/tools/interactive_graphviz.cc b/tensorflow/compiler/xla/tools/interactive_graphviz.cc new file mode 100644 index 0000000000000000000000000000000000000000..ac865707f8697e0b94173a2a33e7be52a9564867 --- /dev/null +++ b/tensorflow/compiler/xla/tools/interactive_graphviz.cc @@ -0,0 +1,652 @@ +/* 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 tool for interactively exploring graphviz dumps of HLO graphs. +// +// Input can be a binary HloSnapshot proto, a binary HloProto proto, or a +// textual HLO string. +// +// Generated visualization is opened in a new default browser window using +// /usr/bin/sensible-browser. + +#include +#include + +#include "absl/algorithm/container.h" +#include "absl/strings/match.h" +#include "absl/strings/numbers.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/str_join.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" +#include "tensorflow/compiler/xla/service/hlo.pb.h" +#include "tensorflow/compiler/xla/service/hlo_runner.h" +#include "tensorflow/compiler/xla/service/local_service.h" +#include "tensorflow/compiler/xla/service/platform_util.h" +#include "tensorflow/compiler/xla/tools/hlo_extractor.h" +#include "tensorflow/core/platform/init_main.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/subprocess.h" +#include "tensorflow/core/util/command_line_flags.h" +#if defined(PLATFORM_GOOGLE) +#include "util/readline/readline.h" +#endif + +namespace xla { +namespace tools { +namespace { + +bool ReadLine(const char *prompt, string *line) { +#if defined(PLATFORM_GOOGLE) + return util::ReadLine(prompt, line); +#else + std::cout << prompt; + std::getline(std::cin, *line); + return std::cin.good(); +#endif +} + +// Command-line opts to this tool. See main() for descriptions of these +// fields. +struct Options { + string hlo_snapshot; + string hlo_proto; + string hlo_text; + string platform; + string browser; +}; + +const char* const kUsage = R"( +This tool lets you load an XLA dump and then interactively explore its graphical +representation. + +Most models are too large to visualize in their entirety using graphviz, but +it's still useful to be able to look at the nodes "near" a particular node of +interest. + +If you pass --platform, this tool will compile the HloModule for the given +platform. This means that if you acquired your proto from a binary running at a +particular CL, the HLO graph it ran isn't necessarily the same as the one shown +here, unless this program was built at the same CL (and our compiler is +deterministic :). + +Be patient when starting this program if you give it a large input; it has to +compile the whole thing. + +Usage: + + interactive_graphviz -- \ + --{hlo_snapshot,hlo_proto,hlo_text}=path/to/binary_proto + --platform={CUDA,CPU,...} +)"; + +// Unless an explicit width is specified, we will render a neighborhood of +// kDefaultWidth nodes around the requested instruction. +constexpr int64 kDefaultWidth = 2; + +// When printing all paths between two nodes, we print out only this many nodes +// by default, truncating the graph if there are more nodes than this in the +// all-paths set. +constexpr int64 kDefaultMaxNumNodesInAllPaths = 100; + +using absl::EqualsIgnoreCase; + +// A global control for whether backend configuration display is enabled. +bool show_backend_config = true; + +HloInstruction* FindInstruction(const HloModule& module, string node_name) { + if (absl::StartsWith(node_name, "%")) { + node_name.erase(node_name.begin()); + } + for (const auto& computation : module.computations()) { + auto instrs = computation->instructions(); + auto it = absl::c_find_if(instrs, [&](const HloInstruction* instr) { + // Try with and without "%" at the beginning of the node name. + return EqualsIgnoreCase(instr->name(), node_name) || + EqualsIgnoreCase(instr->name(), absl::StrCat("%", node_name)); + }); + if (it != instrs.end()) { + return *it; + } + } + return nullptr; +} + +HloComputation* FindComputation(const HloModule& module, + const string& comp_name) { + for (auto* computation : module.computations()) { + if (EqualsIgnoreCase(computation->name(), comp_name)) { + return computation; + } + } + return nullptr; +} + +// 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 )" + << kDefaultWidth << R"(. + allpaths [] + Renders a subset of all paths from one instruction to the other. Either + order of nodes is accepted. Shows the nodes in the all-paths set on the + shortest paths; default is )" + << kDefaultMaxNumNodesInAllPaths << R"(. + + Renders all nodes in . + backend_config [on|off] + Controls whether backend operation configuration information is printed. + list [name|op_name|op_type] + Lists all instructions whose name, metadata op_name, or metadata op_type + contains as a substring. + list computations + Lists all computations in the module. + info + info + Prints information about or . + extract + Creates a new HLO module with as entry computation root. If + is specified, the new computation contains nodes up to + nodes above the root. + help + Prints this usage information.)" + << std::endl; +} + +// Turn metadata-printing on or off. +void DoBackendConfigCommand(const std::vector& tokens) { + if (tokens.size() == 2 && tokens[1] == "on") { + show_backend_config = true; + } else if (tokens.size() == 2 && tokens[1] == "off") { + show_backend_config = false; + } else if (tokens.size() != 1) { + std::cerr << "(Illegal backend_config value. Use either 'on' or 'off'.)" + << std::endl; + } + std::cout << "Backend configuration display " + << (show_backend_config ? "ON" : "OFF") << std::endl; +} + +// List all computations in the module. +void DoListComputationsCommand(const HloModule& module, + const std::vector& tokens) { + if (tokens.size() > 2) { + std::cout << R"(Illegal syntax; "list computations" takes no arguments.)"; + return; + } + if (module.entry_computation() != nullptr) { + std::cout << "Entry computation:" << std::endl; + std::cout << " " << module.entry_computation()->name() << std::endl + << std::endl; + } + std::cout << "Subcomputations:" << std::endl; + std::vector names; + for (const auto& computation : module.computations()) { + if (computation == module.entry_computation()) { + continue; + } + std::cout << " " << computation->name() << std::endl; + } +} + +// List all instructions matching a pattern. +void DoListCommand(const HloModule& module, const std::vector& tokens) { + string pattern = ""; + string type = "name"; + if (tokens.size() == 2) { + pattern = tokens[1]; + } else if (tokens.size() == 3) { + type = tokens[1]; + pattern = tokens[2]; + } else { + std::cout << "Illegal list query syntax. Use " + << R"("list [name|op_name|op_type] pattern".)" << std::endl; + return; + } + + std::cout << "Query results:" << std::endl; + for (const auto& computation : module.computations()) { + for (const auto& instr : computation->instructions()) { + if ((type == "name" && instr->name().find(pattern) != string::npos) || + (type == "op_name" && + instr->metadata().op_name().find(pattern) != string::npos) || + (type == "op_type" && + instr->metadata().op_type().find(pattern) != string::npos)) { + std::cout << " " << instr->name(); + std::cout << ", op_name '" << instr->metadata().op_name() << "'"; + std::cout << ", op_type '" << instr->metadata().op_type() << "'"; + std::cout << std::endl; + } + } + } +} + +// Print info about an instruction or computation. +void DoInfoCommand(const HloModule& module, const std::vector& tokens) { + if (tokens.size() != 2) { + std::cerr << "Illegal info query syntax. Use " + << R"("info name".)"; + return; + } + string node_name = tokens[1]; + + const HloInstruction* instr = FindInstruction(module, node_name); + const HloComputation* comp = FindComputation(module, node_name); + if (!instr && !comp) { + std::cerr << "Couldn't find HloInstruction or HloComputation named " + << node_name << std::endl; + return; + } + + if (comp != nullptr) { + std::cout << "HloComputation " << comp->name() << std::endl; + if (comp->IsFusionComputation()) { + std::cout << " Fusion instruction: " << comp->FusionInstruction()->name() + << std::endl; + } + std::cout << " Parameters:" << std::endl; + for (const auto& param : comp->parameter_instructions()) { + std::cout << " " << param->name() << " (" + << ShapeUtil::HumanStringWithLayout(param->shape()) << ")" + << std::endl; + } + HloInstruction* root = comp->root_instruction(); + std::cout << " Root instruction: " << root->name() << " (" + << ShapeUtil::HumanStringWithLayout(root->shape()) << ")" + << std::endl; + + auto embedded_computations = comp->MakeEmbeddedComputationsList(); + std::cout << " " << embedded_computations.size() << " embedded computation" + << (embedded_computations.size() != 1 ? "s" : "") + << (!embedded_computations.empty() ? ":" : ".") << std::endl; + for (const HloComputation* c : embedded_computations) { + std::cout << " " << c->name() << std::endl; + } + + // Find which computations reference comp as an embedded computation. + std::vector users; + for (const HloComputation* c : module.computations()) { + if (absl::c_linear_search(c->MakeEmbeddedComputationsList(), comp)) { + users.push_back(c); + } + } + std::cout << " Used by " << users.size() << " computation" + << (users.size() != 1 ? "s" : "") << (!users.empty() ? ":" : "."); + for (const HloComputation* c : users) { + std::cout << " " << c->name() << std::endl; + } + } else { + std::cout << "HloInstruction " << instr->name() << std::endl; + std::cout << " Parent computation: " << instr->parent()->name() + << std::endl; + std::cout << " Opcode: " << HloOpcodeString(instr->opcode()) << std::endl; + std::cout << " Shape: " << ShapeUtil::HumanStringWithLayout(instr->shape()) + << std::endl; + std::cout << " Metadata:" << std::endl; + if (!instr->metadata().op_name().empty()) { + std::cout << " Name: " << instr->metadata().op_name() << std::endl; + } + if (!instr->metadata().op_type().empty()) { + std::cout << " Type: " << instr->metadata().op_type() << std::endl; + } + if (!instr->raw_backend_config_string().empty()) { + std::cout << " Backend configuration: " + << instr->raw_backend_config_string() << std::endl; + } + if (instr->opcode() == HloOpcode::kFusion) { + std::cout << " Fusion kind: " << xla::ToString(instr->fusion_kind()) + << std::endl; + std::cout << " Fusion computation: " + << instr->fused_instructions_computation()->name() << std::endl; + std::cout << " Fused computation root: " + << instr->fused_expression_root()->name() << std::endl; + } + std::cout << " Operands:" << std::endl; + for (HloInstruction* operand : instr->operands()) { + std::cout << " " << operand->name() << " (" + << ShapeUtil::HumanStringWithLayout(operand->shape()) << ")" + << std::endl; + } + std::cout << " Users:" << std::endl; + for (HloInstruction* user : instr->users()) { + std::cout << " " << user->name() << std::endl; + } + if (instr->parent()->root_instruction() == instr) { + std::cout << " Root instruction of " << instr->parent()->name() + << std::endl; + } + } +} + +void DoExtractCommand(const HloModule& module, + absl::Span tokens) { + if (tokens.size() > 3) { + std::cerr << R"(Illegal input. Enter e.g. "extract %fusion.1 2")" + << std::endl; + return; + } + + // Find the node with the given name. + string node_name = tokens[1]; + HloInstruction* instr = FindInstruction(module, node_name); + if (!instr) { + std::cerr << "Couldn't find HloInstruction named " << node_name << "." + << std::endl; + return; + } + + int64 height = -1; + if (tokens.size() == 3) { + if (!absl::SimpleAtoi(tokens[2], &height)) { + std::cerr << "Can't parse '" << tokens[2] << "' as an integer." + << std::endl; + return; + } + } + + auto extracted_module = ExtractModule(instr, height); + std::cout << extracted_module->ToString( + HloPrintOptions::ShortParsable().set_print_backend_config( + show_backend_config)) + << std::endl; +} + +// Checks if there is a use-def path from `from` to `to`. +bool ExistsPathFromTo(const HloInstruction* from, const HloInstruction* to) { + std::unordered_set visited; + std::vector to_visit = {from}; + while (!to_visit.empty()) { + auto* n = to_visit.back(); + if (n == to) { + return true; + } + to_visit.pop_back(); + visited.insert(n); + for (auto* user : n->users()) { + if (!visited.count(user)) { + to_visit.push_back(user); + } + } + } + return false; +} + +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 (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; + p.SetProgram(browser_bin, {browser_bin, handle}); + p.Start(); + } else if (handle.empty()) { + std::cerr << "Unable to render graph, perhaps due to graphviz server " + "timeout. Run with --logtostderr to see." + << std::endl; + } else { + std::cerr << "\nExpected a URL, but got strange graph result (dumped " + "above). If this isn't what you expected, maybe file a bug?" + << std::endl; + } +} + +void DoAllPathsCommand(const Options& opts, const HloModule& module, + const std::vector& tokens) { + if (tokens.size() > 4) { + std::cerr << R"(Illegal input. Enter e.g. "allpaths %add.4 %subtract.2" or +"allpaths add.4 subtract.2 42.)" + << std::endl; + return; + } + + int64 max_nodes = kDefaultMaxNumNodesInAllPaths; + if (tokens.size() == 4 && !absl::SimpleAtoi(tokens[3], &max_nodes)) { + std::cerr << "Can't parse '" << tokens[3] << "' as an integer." + << std::endl; + return; + } + + const HloInstruction* n1 = FindInstruction(module, tokens[1]); + if (!n1) { + std::cerr << "Couldn't find HloInstruction named " << tokens[1]; + return; + } + const HloInstruction* n2 = FindInstruction(module, tokens[2]); + if (!n2) { + std::cerr << "Couldn't find HloInstruction named " << tokens[2]; + return; + } + + // Is there a path from n1 to n2, or vice versa? + const HloInstruction* from; + const HloInstruction* to; + if (ExistsPathFromTo(n1, n2)) { + from = n1; + to = n2; + } else if (ExistsPathFromTo(n2, n1)) { + from = n2; + to = n1; + } else { + std::cerr << "No path from/to " << tokens[1] << " to/from " << tokens[2]; + return; + } + DisplayGraphHandle(opts, hlo_graph_dumper::DumpAllPathsFromTo( + *from, *to, max_nodes, /*show_backend_config=*/show_backend_config)); +} + +// 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. + const HloInstruction* instr = FindInstruction(module, node_name); + const HloComputation* comp = FindComputation(module, node_name); + if (!instr && !comp) { + std::cerr << "Couldn't find HloInstruction or HloComputation named " + << node_name << "." << std::endl; + return; + } + + uint64 graph_width = kDefaultWidth; + 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; + } + } + + // Generate the graph and print the resulting string, which should be a + // graphviz url. + if (comp) { + DisplayGraphHandle(opts, hlo_graph_dumper::DumpGraph( + *comp, "", comp->parent()->config().debug_options(), nullptr, + /*show_backend_config=*/show_backend_config)); + } else { + DisplayGraphHandle(opts, hlo_graph_dumper::DumpNeighborhoodAround( + *instr, graph_width, /*show_backend_config=*/show_backend_config)); + } +} + +// Run the main event loop, reading user commands and processing them. +void InteractiveDumpGraphs(const Options& opts, const HloModule& module) { + // This is an interactive tool, but some may use `extract` in non-tty + // environment anyway. Give them a clean hlo dump. + if (isatty(fileno(stdin))) { + std::cout << "\n\nLoaded module " << module.name() << "." << std::endl; + DoHelpCommand(); + } + for (string line; ReadLine("\ncommand: ", &line);) { + if (line.empty()) { + std::cout << R"(Enter e.g. "fusion.1 3" or "add.8".)" << std::endl + << R"(Enter "help" for help; ^D, "quit", or "exit" to exit.)" + << std::endl; + continue; + } + std::vector tokens = absl::StrSplit(line, ' '); + if (tokens[0] == "quit" || tokens[0] == "exit") { + break; + } else if (tokens[0] == "help") { + DoHelpCommand(); + } else if (tokens[0] == "backend_config") { + DoBackendConfigCommand(tokens); + } else if (tokens[0] == "list") { + if (tokens.size() > 1 && tokens[1] == "computations") { + DoListComputationsCommand(module, tokens); + } else { + DoListCommand(module, tokens); + } + } else if (tokens[0] == "info") { + DoInfoCommand(module, tokens); + } else if (tokens[0] == "extract") { + DoExtractCommand(module, tokens); + } else if (tokens[0] == "allpaths") { + DoAllPathsCommand(opts, module, tokens); + } else { + DoPlotCommand(opts, module, tokens); + } + } +} + +void CheckFlags(const Options &opts) { + std::vector nonempty_proto_flags; + if (!opts.hlo_proto.empty()) { + nonempty_proto_flags.push_back("--hlo_proto"); + } + if (!opts.hlo_snapshot.empty()) { + nonempty_proto_flags.push_back("--hlo_snapshot"); + } + if (!opts.hlo_text.empty()) { + nonempty_proto_flags.push_back("--hlo_text"); + } + switch (nonempty_proto_flags.size()) { + case 1: + // We're good to go. + break; + case 0: + LOG(FATAL) << "Need one of the following options: " + << absl::StrJoin(nonempty_proto_flags, ", "); + default: + LOG(FATAL) << "Can only specify one of " + << absl::StrJoin(nonempty_proto_flags, ", "); + } +} + +void RealMain(const Options& opts) { + if (!isatty(fileno(stdin))) { + LOG(ERROR) << "\n\n*****************************************\n" + << "This is an interactive tool, but stdin is not a tty.\n" + << "*****************************************\n\n"; + } + + CheckFlags(opts); + + std::unique_ptr module; + if (!opts.hlo_snapshot.empty()) { + HloSnapshot snapshot; + TF_CHECK_OK(tensorflow::ReadBinaryProto(tensorflow::Env::Default(), + opts.hlo_snapshot, &snapshot)) + << "Can't open, read, or parse HloSnapshot proto at " + << opts.hlo_snapshot; + auto config = + HloModule::CreateModuleConfigFromProto(snapshot.hlo().hlo_module(), + xla::GetDebugOptionsFromFlags()) + .ValueOrDie(); + module = HloModule::CreateFromProto(snapshot.hlo().hlo_module(), config) + .ValueOrDie(); + } else if (!opts.hlo_proto.empty()) { + module = HloRunner::ReadModuleFromBinaryProtoFile( + opts.hlo_proto, xla::GetDebugOptionsFromFlags()) + .ValueOrDie(); + } else if (!opts.hlo_text.empty()) { + module = HloRunner::ReadModuleFromHloTextFile( + opts.hlo_text, xla::GetDebugOptionsFromFlags()) + .ValueOrDie(); + } + + // If a platform was specified, compile the module for that platform. + if (!opts.platform.empty()) { + se::Platform* platform = + PlatformUtil::GetPlatform(opts.platform).ValueOrDie(); + LOG(INFO) << "Compiling module for " << platform->Name(); + + se::StreamExecutor* executor = + platform->ExecutorForDevice(/*ordinal=*/0).ValueOrDie(); + auto compiler = Compiler::GetForPlatform(platform).ValueOrDie(); + module = compiler + ->RunHloPasses(std::move(module), executor, + /*device_allocator=*/nullptr) + .ValueOrDie(); + auto executable = compiler + ->RunBackend(std::move(module), executor, + /*device_allocator=*/nullptr) + .ValueOrDie(); + InteractiveDumpGraphs(opts, executable->module()); + } else { + InteractiveDumpGraphs(opts, *module); + } +} + +} // namespace +} // namespace tools +} // namespace xla + +int main(int argc, char** argv) { + xla::tools::Options opts; + opts.browser = "/usr/bin/sensible-browser"; + bool need_help = false; + const std::vector flag_list = { + tensorflow::Flag("hlo_snapshot", &opts.hlo_snapshot, + "HloSnapshot proto to interactively dump to graphviz"), + tensorflow::Flag("hlo_proto", &opts.hlo_proto, + "XLA hlo proto to interactively dump to graphviz"), + tensorflow::Flag("hlo_text", &opts.hlo_text, + "XLA hlo proto to interactively dump to graphviz"), + tensorflow::Flag("platform", &opts.platform, + "Platform to compile for: CPU, CUDA, etc"), + tensorflow::Flag("browser", &opts.browser, + "Path to web browser used to display produced graphs."), + tensorflow::Flag("help", &need_help, + "Prints this help message"), + }; + xla::string usage = tensorflow::Flags::Usage(argv[0], flag_list); + bool parse_ok = tensorflow::Flags::Parse(&argc, argv, flag_list); + tensorflow::port::InitMain(argv[0], &argc, &argv); + if (argc != 1 || !parse_ok || need_help) { + LOG(QFATAL) << usage; + } + xla::tools::RealMain(opts); + return 0; +} 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 f910e980535c073562473978662f73f4ee4bee79..c01a47b510c0e4252e350960b995643b39b70d4a 100644 --- a/tensorflow/compiler/xla/tools/replay_computation.cc +++ b/tensorflow/compiler/xla/tools/replay_computation.cc @@ -47,10 +47,11 @@ limitations under the License. #include "tensorflow/compiler/xla/client/lib/testing.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/execution_options_util.h" -#include "tensorflow/compiler/xla/legacy_flags/debug_options_flags.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,7 +74,17 @@ 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; @@ -82,16 +93,97 @@ struct Options { std::unique_ptr CompileExecutable(const HloSnapshot& module, LocalClient* client) { XlaComputation computation(module.hlo().hlo_module()); - std::vector argument_layouts; - for (const auto& param : + std::vector argument_layouts; + argument_layouts.reserve( + computation.proto().host_program_shape().parameters_size()); + std::vector argument_layout_ptrs; + for (const ShapeProto& param : computation.proto().host_program_shape().parameters()) { - argument_layouts.push_back(¶m); + argument_layouts.push_back(Shape(param)); + argument_layout_ptrs.push_back(&argument_layouts.back()); } return client - ->Compile(computation, argument_layouts, ExecutableBuildOptions()) + ->Compile(computation, argument_layout_ptrs, ExecutableBuildOptions()) .ValueOrDie(); } +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) // parameter if use_fake_data, Otherwise use recorded data if available. // @@ -114,7 +206,12 @@ StatusOr ReplayComputation(const HloSnapshot& module, std::vector> global_data_arguments; std::vector argument_ptrs; if (opts.use_fake_data) { - global_data_arguments = MakeFakeArgumentsOrDie(computation, client); + // Run fake computations with debug options ignoring XLA_FLAGS. Users very + // likely want XLA_FLAGS only to apply to the "real" computation being run, + // not to the fake computations we use for generating arguments. + auto debug_opts = DefaultDebugOptionsIgnoringFlags(); + global_data_arguments = + MakeFakeArgumentsOrDie(computation, client, &debug_opts); for (const auto& data : global_data_arguments) { argument_ptrs.push_back( client->GlobalDataToShapedBuffer(data->handle(), /*device_ordinal=*/0) @@ -133,55 +230,37 @@ StatusOr ReplayComputation(const HloSnapshot& module, } } - bool provide_infeed = false; - Shape infeed_shape; - if (!opts.fake_infeed_shape.empty()) { - StatusOr shape_status = - ShapeUtil::ParseShapeString(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; - infeed_shape = instruction.shape(); - LOG(INFO) << "Generating fake infeed shape for inferred shape: " - << ShapeUtil::HumanString(infeed_shape); - } - } - } + 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)); + }); } - // 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(); - } - 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. @@ -191,16 +270,16 @@ StatusOr ReplayComputation(const HloSnapshot& module, // Run the computation num_runs times, and return the result from the last // execution. - const bool xla_hlo_profile = - legacy_flags::GetDebugOptionsFromFlags().xla_hlo_profile(); + const bool xla_hlo_profile = GetDebugOptionsFromFlags().xla_hlo_profile(); StreamExecutorMemoryAllocator allocator( client->platform(), {client->platform()->ExecutorForDevice(0).ValueOrDie()}); - absl::optional result; + absl::optional final_result; for (int i = 0; i < opts.num_runs; ++i) { // If xla_hlo_profile is enabled, print a noisy message before the last run, // making it easier to separate this profile from the others in the logspam. - if (xla_hlo_profile && i == opts.num_runs - 1) { + bool is_final_result = i == opts.num_runs - 1; + if (xla_hlo_profile && is_final_result) { LOG(INFO) << "\n\n***** Final run below ******"; } ExecutionProfile profile; @@ -208,14 +287,22 @@ StatusOr ReplayComputation(const HloSnapshot& module, run_options.set_execution_profile(&profile); run_options.set_allocator(&allocator); - TF_ASSIGN_OR_RETURN(result, executable->Run(argument_ptrs, run_options)); + TF_ASSIGN_OR_RETURN(ScopedShapedBuffer result, + executable->Run(argument_ptrs, run_options)); LOG(INFO) << "Done executing in " << static_cast(profile.compute_time_ns()) / 1e9 << "s: " << module.hlo().hlo_module().name(); + + // Save the result if this is for the final iteration. Otherwise discard + // the result before rerunning the computation, so as to free up the + // relevant memory. + if (is_final_result) { + final_result = std::move(result); + } } TF_ASSIGN_OR_RETURN(Literal result_literal, - client->ShapedBufferToLiteral(*result)); + client->ShapedBufferToLiteral(*final_result)); return result_literal; } @@ -288,8 +375,10 @@ int RealMain(absl::Span args, const Options& opts) { for (int64 i = 0; i < executables.size(); ++i) { LocalExecutable* executable = executables[i].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()); @@ -307,9 +396,10 @@ int RealMain(absl::Span args, const Options& opts) { if (snapshot.has_result()) { Literal literal = Literal::CreateFromProto(snapshot.result()).ConsumeValueOrDie(); - fprintf(stdout, "was %s:%s\n", - ShapeUtil::HumanString(snapshot.result().shape()).c_str(), - literal.ToString().c_str()); + fprintf( + stdout, "was %s:%s\n", + ShapeUtil::HumanString(Shape(snapshot.result().shape())).c_str(), + literal.ToString().c_str()); } } } @@ -333,9 +423,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/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..34b73b5206fa20d6dff7567afd78fd89897c8c33 100644 --- a/tensorflow/compiler/xla/util.cc +++ b/tensorflow/compiler/xla/util.cc @@ -86,7 +86,7 @@ bool IsPermutation(absl::Span permutation, int64 rank) { CHECK_LT(index, rank); output[index] = 0; } - return std::find(output.begin(), output.end(), -1) == output.end(); + return !absl::c_linear_search(output, -1); } std::vector InversePermutation( diff --git a/tensorflow/compiler/xla/util.h b/tensorflow/compiler/xla/util.h index 8ce741647414a1fa75e6d706ec1e719ace7b7cc8..f2fd17dc99455a921bf875aad2a3661b4d456823 100644 --- a/tensorflow/compiler/xla/util.h +++ b/tensorflow/compiler/xla/util.h @@ -152,6 +152,13 @@ static inline absl::Span AsInt64Slice( slice.size()); } +// TODO(b/29771030): This nop overload was added to simplify the migration of +// Shape from a proto to a C++ class. Remove after class has been migrated. +static inline absl::Span AsInt64Slice( + absl::Span slice) { + return slice; +} + // As above, but for uint64 types. static inline absl::Span AsUInt64Slice( const tensorflow::protobuf::RepeatedField& v) { @@ -317,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 @@ -387,6 +393,19 @@ T CeilOfRatio(T dividend, T divisor) { return tensorflow::MathUtil::CeilOfRatio(dividend, divisor); } +template +std::vector ElementWiseCeilOfRatio(absl::Span dividends, + absl::Span divisors) { + std::vector ceil_of_ratios; + CHECK_EQ(dividends.size(), divisors.size()); + ceil_of_ratios.reserve(dividends.size()); + absl::c_transform(dividends, divisors, std::back_inserter(ceil_of_ratios), + [](const T dividend, const T divisor) { + return CeilOfRatio(dividend, divisor); + }); + return ceil_of_ratios; +} + // Rounds the value up to a multiple of the divisor by first calling CeilOfRatio // then multiplying by the divisor. For example: RoundUpToNearest(13, 8) => 16 template diff --git a/tensorflow/compiler/xla/window_util.cc b/tensorflow/compiler/xla/window_util.cc index 8ea8dbab2574ca1e24271e7c1c7762d4a6b6a8de..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) { @@ -185,6 +184,16 @@ bool HasWindowReversal(const Window& window) { return false; } +bool AllOrNoneReversed(const Window& window) { + if (window.dimensions().empty()) { + return true; + } + bool reversed = window.dimensions()[0].window_reversal(); + return absl::c_all_of(window.dimensions(), [&](const WindowDimension& dim) { + return dim.window_reversal() == reversed; + }); +} + bool HasDilation(const Window& window) { return HasBaseDilation(window) || HasWindowDilation(window); } diff --git a/tensorflow/compiler/xla/window_util.h b/tensorflow/compiler/xla/window_util.h index 1fb9e855fc16f334eb0e83dfd27b307b2149628f..099d7ecdd5c732ffc8c6ff6370288a2fc4144fa2 100644 --- a/tensorflow/compiler/xla/window_util.h +++ b/tensorflow/compiler/xla/window_util.h @@ -56,6 +56,7 @@ bool HasWindowDilation(const Window& window); bool HasDilation(const Window& window); bool HasWindowReversal(const Window& window); +bool AllOrNoneReversed(const Window& window); // Returns true if the given logical dimension is inactive in the sense that it // has window bound 1, no striding and no padding. diff --git a/tensorflow/compiler/xla/xla.proto b/tensorflow/compiler/xla/xla.proto index 60d25a6407476cddba77aadd1df2e3939f5e40ac..92834dbb02cdcd6383ceec3ffd079834b163ee6a 100644 --- a/tensorflow/compiler/xla/xla.proto +++ b/tensorflow/compiler/xla/xla.proto @@ -100,6 +100,14 @@ message DebugOptions { // names as specified by the HloPassInterface::name() method. repeated string xla_disable_hlo_passes = 30; + // Disables all HLO passes. Notes that some passes are necessary for + // correctness and the invariants that must be satisfied by "fully optimized" + // HLO are different for different devices and may change over time. The only + // "guarantee", such as it is, is that if you compile XLA and dump the + // optimized HLO for some graph, you should be able to run it again on the + // same device with the same build of XLA. + bool xla_disable_all_hlo_passes = 104; + // Numerical optimization level for the XLA compiler backend; the specific // interpretation of this value is left to the backends. int32 xla_backend_optimization_level = 31; @@ -193,7 +201,11 @@ message DebugOptions { // - Assuming that operations never produce or consume NaN or +/- Inf. // - Assuming that +0 and -0 are indistinguishable. bool xla_cpu_enable_fast_math = 99; - bool xla_gpu_enable_fast_math = 100; + + // When true we lower the Minimum and Maximum hlos in the GPU backend such + // that Min(NotNaN, NaN) = Min(NaN, NotNaN) = NotNaN. In other words, if flag + // this is true we don't propagate NaNs through Min and Max. + bool xla_gpu_enable_fast_min_max = 100; // Crashes the program when any kind of verification fails, instead of just // logging the failures. One example is cross checking of convolution results @@ -209,6 +221,21 @@ message DebugOptions { // the host that run models in parallel across multiple devices. int32 xla_force_host_platform_device_count = 102; + // If set to true XLA:GPU invokes `ptxas` with -O0 (default is -O3). + bool xla_gpu_disable_ptxas_optimizations = 103; + + // Dump HLO graphs as an HTML (DOT -> SVG inlined in HTML) + bool xla_hlo_dump_as_html = 105; + + // Enable fast math with eigen in the HLO evaluator. + bool xla_hlo_evaluator_use_fast_path = 106; + + // Temporary option to allow support for both the R1 and the scalar index + // versions of DynamicSlice and DynamicUpdateSlice. Only used for testing. + bool xla_allow_scalar_index_dynamic_ops = 107; + + // Next id: 108 + // Extra options to pass to the compilation backend (e.g. LLVM); specific // interpretation of these values is left to the backend. map xla_backend_extra_options = 500; @@ -224,7 +251,7 @@ message ExecutionOptions { // may be faster when using this layout. // // We use a Shape here to accommodate computations that return a tuple. - Shape shape_with_output_layout = 2; + ShapeProto shape_with_output_layout = 2; // Used to seed random-number generators used in this computation. If this is // 0, we generate a seed ourselves. @@ -238,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 { @@ -253,7 +284,7 @@ message TransferToClientRequest { // This optional field directs the service to return the literal in this // layout. A shape is used to hold the layout to accommodate tuples. - Shape shape_with_layout = 2; + ShapeProto shape_with_layout = 2; } message TransferToClientResponse { @@ -281,7 +312,7 @@ message TransferToInfeedResponse { message TransferFromOutfeedRequest { // This optional field directs the service to return the literal in this // layout. A shape is used to hold the layout to accommodate tuples. - Shape shape_with_layout = 1; + ShapeProto shape_with_layout = 1; int64 replica_id = 2; DeviceHandle device_handle = 3; @@ -316,12 +347,40 @@ message CreateChannelHandleResponse { } message UnregisterRequest { - GlobalDataHandle data = 1; + repeated GlobalDataHandle data = 1; } message UnregisterResponse { } +message CompileRequest { + // The graph to be compiled. + HloModuleProto computation = 1; + + // Options that affect how XLA compiles code to service this request. + ExecutionOptions execution_options = 2; + + // The layouts of the input arguments. If not set, the default layout will be + // used. Although the real arguments are not needed in compilation, the + // layouts of the arguments can affect the compilation. + repeated ShapeProto input_shape_with_layout = 3; +} + +message CompileResponse { + // The handle to the executable. + ExecutionHandle handle = 1; +} + +message ExecuteRequest { + ExecutionHandle handle = 1; + + // The shape and layout of the arguments must be the same as the those of the + // executable's parameters. + repeated GlobalDataHandle arguments = 2; +} + +// TODO(b/118493728): Remove this and ExecuteGraphParallelRequest and replace +// the uses with calls to Compile and Execute. message ExecuteGraphRequest { HloModuleProto computation = 1; repeated GlobalDataHandle arguments = 2; @@ -354,7 +413,7 @@ message WaitForExecutionResponse { message ComputeConstantGraphRequest { HloModuleProto computation = 1; - Layout output_layout = 2; + LayoutProto output_layout = 2; } message ComputeConstantResponse { @@ -378,7 +437,7 @@ message LoadDataRequest { string columnio_field = 2; // Individual element shape, excluding rows. - Shape element_shape = 3; + ShapeProto element_shape = 3; // Warning: ColumnIO does not support random-access, so use offset with // caution in performance-critical scenarios. @@ -394,7 +453,7 @@ message LoadDataRequest { message LoadDataResponse { GlobalDataHandle data = 1; - Shape data_shape = 2; + ShapeProto data_shape = 2; int64 available_rows = 3; int64 rows_loaded = 4; int64 nanoseconds = 5; @@ -405,7 +464,7 @@ message GetShapeRequest { } message GetShapeResponse { - Shape shape = 1; + ShapeProto shape = 1; } message UnpackRequest { diff --git a/tensorflow/compiler/xla/xla_data.proto b/tensorflow/compiler/xla/xla_data.proto index 73b3589dbf12341ddb3f3e819a550467a7b4d166..a64e2f5df5cacca05e83f31c941c57abd5ccf4de 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,29 +76,7 @@ enum PrimitiveType { // primitive type will have empty dimensions and tuple_shapes fields. TOKEN = 17; - // Next = 18 -} - -// Describes the value held inside padding elements. -enum PaddingValue { - INVALID_PAD = 0; - - // Zero padding must be 0-values that correspond to the shape's element type. - ZERO_PAD = 1; - - // One padding must be 1-values that correspond to the shape's element type. - ONE_PAD = 2; - - // "Lowest" padding must be the lowest values in the shape's element type, - // used as padding for operations like max-accumulation. - LOWEST_PAD = 3; - - // "Highest" padding must be the largest values in the shape's element type, - // used as padding for operations like min-accumulation. - HIGHEST_PAD = 4; - - // Unknown padding could be anything; e.g. floating NaNs! - UNKNOWN_PAD = 5; + // Next = 19 } // Describes the padding configuration for Pad operation. The padding amount on @@ -122,18 +101,29 @@ message PaddingConfig { // A format specifies the method used by a layout to store an array in memory. enum Format { + // TODO(b/120869032): Rename this to FORMAT_NONE or something else which + // better corresponds to its meaning. INVALID_FORMAT = 0; - // The default layout, with exactly one storage location per element (ignoring - // padding). + // The default layout, with exactly one storage location per element. DENSE = 1; // A sparsely encoded layout, providing only the index/value pairs of non-zero // elements. SPARSE = 2; } +// Describes a tile used in tiling-based layout. Refer to +// g3doc/third_party/tensorflow/compiler/xla/g3doc/layout_with_tiling.md for +// details about tiling-based layout. +message TileProto { + // Number of elements in each dimension of the tile. It's ordered from the + // most major dimension of the tile to the most minor dimension of the tile. + // The dimensions correspond to a suffix of the dimensions of the shape being + // tiled. + repeated int64 dimensions = 1; +} + // A layout describes how the array is placed in (1D) memory space. This -// includes the minor-to-major ordering of dimensions within a shape, as well as -// any padding present in those dimensions. +// includes the minor-to-major ordering of dimensions within a shape. // // Clients must specify the layouts of input Literals to the // computation. Layouts specified in interior operations which take Shapes (for @@ -142,7 +132,7 @@ enum Format { // See the XLA documentation for more information on shapes and layouts. // // LINT.IfChange -message Layout { +message LayoutProto { // The method used to store the data in memory. The format determines which of // the other fields are used by the layout. Format format = 4; @@ -151,22 +141,31 @@ message Layout { // (slowest varying index). This field is required. repeated int64 minor_to_major = 1; - // The width to which the layout of each dimension is padded up to. If - // present, the size of the padded_dimensions must equal the rank of the - // shape. The padding appears at the end of a dimension, not at the - // beginning. This kind of padding, unlike padding in e.g. convolution, is not - // part of the shape. This field must be unset unless the format is DENSE. - repeated int64 padded_dimensions = 2; + reserved 2; + reserved "padded_dimensions"; - // Describes the values in the padding specified by padded_dimensions. This - // field must be unset unless the format is DENSE. - PaddingValue padding_value = 3; + reserved 3; + reserved "padding_value"; // The maximum number of elements that can be stored for SPARSE formats. This // can be used to determine the maximum size in bytes of arrays stored in // memory. This field must be unset unless the format is SPARSE. int64 max_sparse_elements = 5; + // A sequence of tiles, starting from the tile that's applied first to the + // Shape. + // + // TODO(b/119839262): implement tiling in each backend or add Unimplemented + // error. + repeated TileProto tiles = 6; + + // Bit size of each element. If the size is bigger than what the element + // type requires, the value is stored in the least significant + // bits and the additional most significant bits are filled with 0's. + // + // TODO(b/119839262): implement in each backend or add Unimplemented error. + int64 element_size_in_bits = 7; + // Important: if any field is added, be sure to modify ShapeUtil::Equal() and // LayoutUtil::Hash appropriately to account for the new field. } @@ -183,25 +182,34 @@ message Layout { // See the XLA documentation for more information on shapes and layouts. // // LINT.IfChange -message Shape { +message ShapeProto { reserved 1; reserved "rank"; // The element type for this shape. PrimitiveType element_type = 2; - // The size (number of elements) for each dimension. - // In XLA, dimensions are numbered from 0 to N-1 for an - // N-dimensional array. The first element of 'dimensions' is the size of - // dimension 0, the second element is the size of dimension 1, and so forth. - // Empty list indicates a scalar. + // The size (number of elements) for each dimension, or an upper bound on the + // size if the dimension is dynamic. In XLA, dimensions are numbered from 0 + // to N-1 for an N-dimensional array. The first element of 'dimensions' is the + // size of dimension 0, the second element is the size of dimension 1, and so + // forth. Empty list indicates a scalar. + // + // If the respective element in 'is_dimension_dynamic' is true then the value + // in this field represents an upper bound on the size of the dimension. repeated int64 dimensions = 3; // For tuples only, the shapes of constitutent shapes in the tuple sequence. - repeated Shape tuple_shapes = 4; + repeated ShapeProto tuple_shapes = 4; // The layout used to back this shape. - Layout layout = 5; + LayoutProto layout = 5; + + // For arrays, this indicates whether or not each dimension is + // dynamically-sized. The number of elements in this repeated field should be + // zero (indicating that no dimensions are dynamic) or equal to the number of + // elements in the 'dimensions' field. + repeated bool is_dynamic_dimension = 6; // Important: if any field is added, be sure to modify ShapeUtil::Equal(), // ShapeUtil::Compatible() and ShapeUtil::Hash() appropriately to account for @@ -212,9 +220,9 @@ message Shape { // Shape of the parameters and output of a computation (like a traditional // function signature). -message ProgramShape { - repeated Shape parameters = 1; - Shape result = 2; +message ProgramShapeProto { + repeated ShapeProto parameters = 1; + ShapeProto result = 2; repeated string parameter_names = 3; } @@ -349,7 +357,7 @@ message DeviceAssignmentProto { // Transfers to/from the client are encoded in literal form, and the structure // of the repeated fields is implied by the shape. message LiteralProto { - Shape shape = 1; + ShapeProto shape = 1; repeated bool preds = 2; bytes s8s = 15; bytes u8s = 3; @@ -360,12 +368,15 @@ 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 and BF16s are encoded in little endian byte order + // The F16s, BF16s, U16s and S16s are encoded in little endian byte order bytes f16s = 11; bytes bf16s = 13; + bytes u16s = 16; + bytes s16s = 17; repeated int64 sparse_indices = 14; - // Next = 16 + // Next = 19 } message WindowDimension { @@ -548,7 +559,7 @@ message OpSharding { } Type type = 1; // The shape of the sharded tile. - Shape tile_shape = 2; + ShapeProto tile_shape = 2; // The shape of the tile assignment tensor - this must be the same rank as // tile_shape and the product of its dimensions must equal // tile_assignment_devices.size(). diff --git a/tensorflow/compiler/xrt/BUILD b/tensorflow/compiler/xrt/BUILD index 2ff97914f862e0ec30fc54602ec5fee2a0a5ebca..2dae746d034a1bf52e84de74dfb0c6e23aaed4d1 100644 --- a/tensorflow/compiler/xrt/BUILD +++ b/tensorflow/compiler/xrt/BUILD @@ -22,6 +22,7 @@ xla_proto_library( deps = [ "//tensorflow/compiler/tf2xla:host_compute_metadata_proto", "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla:xla_proto", "//tensorflow/compiler/xla/service:hlo_proto", ], ) @@ -32,20 +33,25 @@ cc_library( "xrt_compilation_cache.cc", "xrt_device.cc", "xrt_state.cc", + "xrt_util.cc", ], hdrs = [ "xrt_compilation_cache.h", "xrt_device.h", "xrt_state.h", + "xrt_util.h", ], deps = [ "//tensorflow/compiler/jit:xla_device", "//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:types", "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla:xla_proto", "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/service:backend", "//tensorflow/compiler/xla/service:device_memory_allocator", diff --git a/tensorflow/compiler/xrt/kernels/BUILD b/tensorflow/compiler/xrt/kernels/BUILD index 9e3d2454d16730c1d1f93cb384db88544380f77e..dc02fd272fd8700c7f8fa64adf7ab57c88bab706 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: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/legacy_flags:debug_options_flags", - "//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,6 +50,7 @@ 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_proto", "//tensorflow/compiler/xrt:xrt_utils", "//tensorflow/core:core_cpu_internal", @@ -62,7 +58,7 @@ cc_library( "//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 dc62cf7a6b24e373374b458d2e4722e79500fb93..2ee1a6cd1aebcdbd65892b33e5044489070ab5c4 100644 --- a/tensorflow/compiler/xrt/kernels/xrt_compile_ops.cc +++ b/tensorflow/compiler/xrt/kernels/xrt_compile_ops.cc @@ -33,6 +33,7 @@ limitations under the License. #include "tensorflow/compiler/xrt/xrt.pb.h" #include "tensorflow/compiler/xrt/xrt_compilation_cache.h" #include "tensorflow/compiler/xrt/xrt_device.h" +#include "tensorflow/compiler/xrt/xrt_util.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/resource_mgr.h" #include "tensorflow/core/framework/tensor.h" @@ -108,19 +109,26 @@ Status XRTCompileOp::Compile(OpKernelContext* ctx, TF_ASSIGN_OR_RETURN(xla::XlaComputation computation, client->LoadSnapshot(computation_proto.hlo_snapshot())); - std::vector argument_layouts( + std::vector argument_layouts( + config.program_shape().parameters_size()); + std::vector argument_layout_ptrs( config.program_shape().parameters_size()); for (int i = 0; i < config.program_shape().parameters_size(); ++i) { - argument_layouts[i] = &config.program_shape().parameters(i); + argument_layouts[i] = xla::Shape(config.program_shape().parameters(i)); + argument_layout_ptrs[i] = &argument_layouts[i]; } xla::ExecutableBuildOptions build_options; build_options.set_device_ordinal(client->default_device_ordinal()); - build_options.set_result_layout(config.program_shape().result()); + build_options.set_result_layout(xla::Shape(config.program_shape().result())); build_options.set_device_allocator(device_ref.backend()->memory_allocator()); + if (config.has_debug_options()) { + *build_options.mutable_debug_options() = + BuildXlaDebugOptions(config.debug_options()); + } VLOG(1) << "Building executable"; auto compile_result = - client->Compile(computation, argument_layouts, build_options); + client->Compile(computation, argument_layout_ptrs, build_options); if (!compile_result.ok()) { return compile_result.status(); } @@ -174,11 +182,12 @@ void XRTCompileOp::Compute(OpKernelContext* ctx) { ctx->set_output(0, handle_output); xla::LocalExecutable* executable = entry->get().get_executable(); - xla::ProgramShape program_shape = executable->executable() - ->module() - .config() - .entry_computation_layout() - .ComputeProgramShape(); + xla::ProgramShapeProto program_shape = executable->executable() + ->module() + .config() + .entry_computation_layout() + .ComputeProgramShape() + .ToProto(); Tensor program_shape_output(DT_STRING, TensorShape({1})); program_shape_output.vec()(0) = program_shape.SerializeAsString(); ctx->set_output(1, program_shape_output); @@ -206,11 +215,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)); @@ -221,9 +225,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 8c6191ddc06ea7d85f5fd21a7d4058c669ffdeb2..116c193cab65410a5a7c3058f98cc2be2cbe9e67 100644 --- a/tensorflow/compiler/xrt/kernels/xrt_execute_op.cc +++ b/tensorflow/compiler/xrt/kernels/xrt_execute_op.cc @@ -19,6 +19,7 @@ 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" @@ -229,13 +230,53 @@ Status XRTExecuteOp::DoWork(OpKernelContext* context) { shaped_buffer, device_ref.backend(), device_ref.device_ordinal(), &output_tuple)); - Tensor* output_tensor; - TF_RETURN_IF_ERROR( - context->allocate_output(0, TensorShape({}), &output_tensor)); - int64 key; - TF_RETURN_IF_ERROR(output_tuple->Intern(rm, &key)); - output_tensor->scalar()() = key; - + // 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 = + xla::ShapeUtil::TupleElementCount(output_tuple->on_device_shape()); + Tensor* output_tensor; + TF_RETURN_IF_ERROR(context->allocate_output( + 0, TensorShape({tuple_element_count}), &output_tensor)); + + for (int64 i = 0; i < tuple_element_count; ++i) { + xla::ShapeIndex shape_index; + shape_index.push_back(i); + + XRTTupleAllocation* suballocation; + TF_RETURN_IF_ERROR(XRTTupleAllocation::MakeSubBuffer( + output_tuple, shape_index, &suballocation, + /*alias_parent_allocation=*/false)); + int64 key; + TF_RETURN_IF_ERROR(suballocation->Intern(rm, &key)); + output_tensor->vec()(i) = key; + } + output_tuple->Unref(); + } else { + Tensor* output_tensor; + TF_RETURN_IF_ERROR( + context->allocate_output(0, TensorShape({}), &output_tensor)); + int64 key; + TF_RETURN_IF_ERROR(output_tuple->Intern(rm, &key)); + output_tensor->scalar()() = key; + } return Status::OK(); } diff --git a/tensorflow/compiler/xrt/kernels/xrt_state_ops.cc b/tensorflow/compiler/xrt/kernels/xrt_state_ops.cc index ffea592491d43788b876a51866dc8a6611e8c734..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") @@ -87,6 +98,19 @@ REGISTER_KERNEL_BUILDER(Name("XRTReadLiteral") .HostMemory("literal"), XRTReadLiteralOp); +REGISTER_KERNEL_BUILDER(Name("XRTWriteLiteral") + .Device(DEVICE_XLA_GPU) + .HostMemory("handle") + .HostMemory("literal") + .HostMemory("output_handle"), + XRTWriteLiteralOp); +REGISTER_KERNEL_BUILDER(Name("XRTWriteLiteral") + .Device(DEVICE_XLA_CPU) + .HostMemory("handle") + .HostMemory("literal") + .HostMemory("output_handle"), + XRTWriteLiteralOp); + REGISTER_KERNEL_BUILDER(Name("XRTReadLiteralAndRelease") .Device(DEVICE_XLA_GPU) .HostMemory("handle") @@ -107,4 +131,9 @@ REGISTER_KERNEL_BUILDER(Name("XRTReleaseAllocationHandle") .HostMemory("handle"), XRTReleaseAllocationOp); +REGISTER_KERNEL_BUILDER(Name("XRTReleaseAllAllocations").Device(DEVICE_XLA_GPU), + XRTReleaseAllAllocationsOp); +REGISTER_KERNEL_BUILDER(Name("XRTReleaseAllAllocations").Device(DEVICE_XLA_CPU), + XRTReleaseAllAllocationsOp); + } // namespace tensorflow diff --git a/tensorflow/compiler/xrt/kernels/xrt_state_ops.h b/tensorflow/compiler/xrt/kernels/xrt_state_ops.h index 54b06558adcd8ef1f8f1bee52d210d558801afea..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" @@ -183,9 +188,7 @@ class XRTAllocateOp : public OpKernel { // 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, allocation_proto.device_ordinal(), &device_ref)); + OP_REQUIRES_OK(ctx, DeviceAccessor::InitScopedRef(ctx, &device_ref)); XRTTupleAllocation* allocation; OP_REQUIRES_OK(ctx, XRTTupleAllocation::CreateAndTransfer( @@ -202,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 @@ -393,6 +499,56 @@ class XRTReadLiteralOp : public OpKernel { } }; +// Op that writes a new literal value into device-resident memory. +template +class XRTWriteLiteralOp : public OpKernel { + public: + explicit XRTWriteLiteralOp(OpKernelConstruction* ctx) : OpKernel(ctx) {} + ~XRTWriteLiteralOp() override = default; + XRTWriteLiteralOp(const XRTWriteLiteralOp&) = delete; + XRTWriteLiteralOp& operator=(const XRTWriteLiteralOp&) = delete; + + void Compute(OpKernelContext* ctx) override { + VLOG(1) << "XRTWriteLiteralOp::Compute"; + + const Tensor& handle_tensor = ctx->input(0); + OP_REQUIRES( + ctx, TensorShapeUtils::IsScalar(handle_tensor.shape()), + errors::Internal("computation input should be an int64 scalar")); + int64 allocation_handle = handle_tensor.scalar()(); + + const Tensor& literal_info = ctx->input(1); + OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(literal_info.shape()), + errors::Internal("literal input should be a string scalar")); + xla::LiteralProto literal_proto; + OP_REQUIRES(ctx, + literal_proto.ParseFromString(literal_info.scalar()()), + errors::InvalidArgument( + "Unable to parse allocation input to LiteralProto")); + xla::Literal literal; + OP_REQUIRES_OK(ctx, XRTStateHelpers::MakeLiteral(literal_proto, &literal)); + + ResourceMgr* rm; + OP_REQUIRES_OK(ctx, DeviceAccessor::GetResourceManager(ctx, &rm)); + + XRTTupleAllocation* allocation; + OP_REQUIRES_OK( + ctx, XRTTupleAllocation::Lookup(rm, allocation_handle, &allocation)); + core::ScopedUnref allocation_unref(allocation); + // We are guaranteed that the underlying device object won't be deleted out + // from under us, while the ScopedRef is live. + typename DeviceAccessor::ScopedRef device_ref; + OP_REQUIRES_OK(ctx, DeviceAccessor::InitScopedRef( + ctx, allocation->device_ordinal(), &device_ref)); + OP_REQUIRES_OK(ctx, + allocation->WriteLiteral(device_ref.backend(), literal)); + + Tensor output(DT_INT64, TensorShape({})); + output.scalar()() = allocation_handle; + ctx->set_output(0, output); + } +}; + // Op that discards a handle to device memory. template class XRTReleaseAllocationOp : public OpKernel { @@ -405,17 +561,37 @@ 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)); + 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; + } + } +}; + +// Op that discards a handle to device memory. +template +class XRTReleaseAllAllocationsOp : public OpKernel { + public: + explicit XRTReleaseAllAllocationsOp(OpKernelConstruction* ctx) + : OpKernel(ctx) {} + ~XRTReleaseAllAllocationsOp() override = default; + XRTReleaseAllAllocationsOp(const XRTReleaseAllAllocationsOp&) = delete; + XRTReleaseAllAllocationsOp& operator=(const XRTReleaseAllAllocationsOp&) = + delete; + + void Compute(OpKernelContext* ctx) override { + VLOG(1) << "XRTReleaseAllAllocationsOp::Compute"; - VLOG(2) << "Released allocation handle " << key; + ResourceMgr* rm; + OP_REQUIRES_OK(ctx, DeviceAccessor::GetResourceManager(ctx, &rm)); + OP_REQUIRES_OK(ctx, XRTTupleAllocation::ReleaseAllAllocations(rm)); } }; 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 07d025ce343f229097b557d33ad41bf9612b0696..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") @@ -95,6 +124,20 @@ Copies an allocated tuple from device memory and returns it as a literal. 'literal' is a serialized xla::LiteralProto proto. )"); +REGISTER_OP("XRTWriteLiteral") + .Input("handle: int64") + .Input("literal: string") + .Output("output_handle: int64") + .SetShapeFn(tensorflow::shape_inference::ScalarShape) + .Doc( + R"( +Copies the input literal into the device memory pointed to by handle. +Returns the handle itself. + +'handle' is the id returned from the Op that produced the on-device allocation. +'literal' is a serialized xla::LiteralProto proto to be written to device memory. +)"); + REGISTER_OP("XRTReadLiteralAndRelease") .Input("handle: int64") .Output("literal: string") @@ -113,10 +156,18 @@ 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") + .SetShapeFn(tensorflow::shape_inference::NoOutputs) + .Doc( + R"( +Discards all the XRT allocations. All the client held handles will be invalid. )"); } // namespace tensorflow 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 25464b5554d21f4b936f3f4a442fd174a8b56a8b..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; } @@ -175,6 +203,18 @@ xla::XlaComputation AddAndTuple() { return builder.Build().ValueOrDie(); } +xla::XlaComputation AddAndSubTuple() { + xla::XlaBuilder builder("AddAndSubTuple"); + auto p0 = xla::Parameter(&builder, 0, xla::ShapeUtil::MakeShape(xla::F32, {}), + "P0"); + auto p1 = xla::Parameter(&builder, 1, xla::ShapeUtil::MakeShape(xla::F32, {}), + "P1"); + auto sum = xla::Add(p0, p1); + auto sub = xla::Sub(p0, p1); + xla::Tuple(&builder, {sum, sub}); + return builder.Build().ValueOrDie(); +} + void StoreComputationSnapshot(const xla::XlaComputation& computation, xla::HloSnapshot* dst) { auto snapshot = computation.Snapshot().ValueOrDie(); @@ -203,9 +243,295 @@ 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() = + xla::LiteralUtil::CreateR2({{4, 5}, {6, 7}}).ToProto(); + + Scope root = Scope::NewRootScope().WithDevice(DeviceFromFlag()); + auto value = + ops::Const(root.WithDevice("/device:CPU:0"), alloc.SerializeAsString()); + auto handle = ops::XRTAllocate(root, value); + auto read_back = ops::XRTReadLiteral(root, handle); + TF_ASSERT_OK(root.status()); + + tensorflow::ClientSession session(root); + std::vector outputs; + TF_EXPECT_OK(session.Run({read_back, handle}, &outputs)); + EXPECT_EQ(outputs.size(), 2); + + int64 allocation_handle = outputs[1].scalar()(); + xla::LiteralProto response; + EXPECT_TRUE(response.ParseFromString(outputs[0].scalar()())); + EXPECT_TRUE(CompareLiteralProtos(alloc.value(), response)); + outputs.clear(); + + xla::LiteralProto new_literal = + xla::LiteralUtil::CreateR2({{9, 2}, {4, 1}}).ToProto(); + auto new_value = ops::Const(root.WithDevice("/device:CPU:0"), + new_literal.SerializeAsString()); + auto write_op = + ops::XRTWriteLiteral(root, Input(allocation_handle), new_value); + TF_ASSERT_OK(root.status()); + TF_EXPECT_OK(session.Run({write_op}, &outputs)); + EXPECT_EQ(outputs.size(), 1); + EXPECT_EQ(allocation_handle, outputs[0].scalar()()); + outputs.clear(); + + auto read_after_write = ops::XRTReadLiteral(root, Input(allocation_handle)); + TF_EXPECT_OK(session.Run({read_after_write}, &outputs)); + EXPECT_EQ(outputs.size(), 1); + + xla::LiteralProto new_response; + EXPECT_TRUE(new_response.ParseFromString(outputs[0].scalar()())); + EXPECT_TRUE(CompareLiteralProtos(new_literal, new_response)); + + 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)); +} + +TEST(RawApiTest, AllocAndClearAll) { + xrt::XLAAllocation alloc; + *alloc.mutable_value() = + xla::LiteralUtil::CreateR2({{4, 5}, {6, 7}}).ToProto(); + + Scope root = Scope::NewRootScope().WithDevice(DeviceFromFlag()); + auto value = + ops::Const(root.WithDevice("/device:CPU:0"), alloc.SerializeAsString()); + auto handle = ops::XRTAllocate(root, value); + TF_ASSERT_OK(root.status()); + + tensorflow::ClientSession session(root); + std::vector outputs; + TF_EXPECT_OK(session.Run({handle}, &outputs)); + EXPECT_EQ(outputs.size(), 1); + + int64 allocation_handle = outputs[0].scalar()(); + + auto clear_all = ops::XRTReleaseAllAllocations(root); + + outputs.clear(); + TF_EXPECT_OK(session.Run(tensorflow::ClientSession::FeedType(), {}, + {clear_all}, &outputs)); + EXPECT_EQ(outputs.size(), 0); + + auto read_after_clear = ops::XRTReadLiteral(root, Input(allocation_handle)); + EXPECT_EQ(session.Run({read_after_clear}, &outputs).code(), + tensorflow::error::Code::NOT_FOUND); +} + TEST(RawApiTest, ReadAndWriteState) { xrt::XLAAllocation alloc; - alloc.set_device_ordinal(0); *alloc.mutable_value() = TwoElementTuple(); Scope root = Scope::NewRootScope().WithDevice(DeviceFromFlag()); @@ -230,7 +556,6 @@ TEST(RawApiTest, ReadAndWriteState) { TEST(RawApiTest, ReadAndWriteStateAutoFree) { xrt::XLAAllocation alloc; - alloc.set_device_ordinal(0); *alloc.mutable_value() = TwoElementTuple(); Scope root = Scope::NewRootScope().WithDevice(DeviceFromFlag()); @@ -251,7 +576,6 @@ TEST(RawApiTest, ReadAndWriteStateAutoFree) { TEST(RawApiTest, SubBuffer) { xrt::XLAAllocation alloc; - alloc.set_device_ordinal(0); *alloc.mutable_value() = NestedTuple(); Scope root = Scope::NewRootScope().WithDevice(DeviceFromFlag()); @@ -292,10 +616,8 @@ TEST(RawApiTest, SubBuffer) { TEST(RawApiTest, MakeTuple) { xrt::XLAAllocation alloc_0; - alloc_0.set_device_ordinal(0); *alloc_0.mutable_value() = TwoElementTuple(); xrt::XLAAllocation alloc_1; - alloc_1.set_device_ordinal(0); *alloc_1.mutable_value() = ScalarLiteral(); // The trivial tuple that just forwards its input and releases it. @@ -366,18 +688,19 @@ TEST(RawApiTest, MakeTuple) { TEST(RawApiTest, CompileAndExecute) { xrt::XLAAllocation p0; - p0.set_device_ordinal(0); *p0.mutable_value() = FloatVector({1.0f, 2.0f}); xrt::XLAAllocation p1; - p1.set_device_ordinal(0); *p1.mutable_value() = FloatVector({8.0f, 5.0f}); xrt::XLAComputation c; auto config = c.mutable_config(); auto shapes = config->mutable_program_shape(); - *shapes->add_parameters() = xla::ShapeUtil::MakeShape(xla::F32, {2}); - *shapes->add_parameters() = xla::ShapeUtil::MakeShape(xla::F32, {2}); - *shapes->mutable_result() = xla::ShapeUtil::MakeShape(xla::F32, {2}); + *shapes->add_parameters() = + xla::ShapeUtil::MakeShape(xla::F32, {2}).ToProto(); + *shapes->add_parameters() = + xla::ShapeUtil::MakeShape(xla::F32, {2}).ToProto(); + *shapes->mutable_result() = + xla::ShapeUtil::MakeShape(xla::F32, {2}).ToProto(); StoreComputationSnapshot(AddAndScale(), c.mutable_hlo_snapshot()); xrt::XRTExecutionConfig e; @@ -411,25 +734,26 @@ TEST(RawApiTest, CompileAndExecute) { auto expected = xla::LiteralUtil::CreateR1({27.0f, 21.0f}); EXPECT_TRUE(CompareLiteralToLiteralProto(expected, response)); - xla::ProgramShape program_shape; + xla::ProgramShapeProto program_shape; EXPECT_TRUE(program_shape.ParseFromString(outputs[1].vec()(0))); EXPECT_EQ(program_shape.parameters_size(), 2); } TEST(RawApiTest, CompileAndExecuteWithArgumentVector) { xrt::XLAAllocation p0; - p0.set_device_ordinal(0); *p0.mutable_value() = FloatVector({1.0f, 2.0f}); xrt::XLAAllocation p1; - p1.set_device_ordinal(0); *p1.mutable_value() = FloatVector({8.0f, 5.0f}); xrt::XLAComputation c; auto config = c.mutable_config(); auto shapes = config->mutable_program_shape(); - *shapes->add_parameters() = xla::ShapeUtil::MakeShape(xla::F32, {2}); - *shapes->add_parameters() = xla::ShapeUtil::MakeShape(xla::F32, {2}); - *shapes->mutable_result() = xla::ShapeUtil::MakeShape(xla::F32, {2}); + *shapes->add_parameters() = + xla::ShapeUtil::MakeShape(xla::F32, {2}).ToProto(); + *shapes->add_parameters() = + xla::ShapeUtil::MakeShape(xla::F32, {2}).ToProto(); + *shapes->mutable_result() = + xla::ShapeUtil::MakeShape(xla::F32, {2}).ToProto(); StoreComputationSnapshot(AddAndScale(), c.mutable_hlo_snapshot()); xrt::XRTExecutionConfig e; @@ -465,7 +789,7 @@ TEST(RawApiTest, CompileAndExecuteWithArgumentVector) { auto expected = xla::LiteralUtil::CreateR1({27.0f, 21.0f}); EXPECT_TRUE(CompareLiteralToLiteralProto(expected, response)); - xla::ProgramShape program_shape; + xla::ProgramShapeProto program_shape; EXPECT_TRUE(program_shape.ParseFromString(outputs[1].vec()(0))); EXPECT_EQ(program_shape.parameters_size(), 2); } @@ -494,8 +818,8 @@ TEST(RawApiTest, CompileWithXlaReturnShapes) { xrt::XLAComputation c; auto config = c.mutable_config(); auto shapes = config->mutable_program_shape(); - *shapes->add_parameters() = param_shape; - *shapes->mutable_result() = result_shape; + *shapes->add_parameters() = param_shape.ToProto(); + *shapes->mutable_result() = result_shape.ToProto(); StoreComputationSnapshot(xla_computation, c.mutable_hlo_snapshot()); Scope root = Scope::NewRootScope().WithDevice(DeviceFromFlag()); @@ -510,8 +834,9 @@ TEST(RawApiTest, CompileWithXlaReturnShapes) { TF_EXPECT_OK(session.Run(tensorflow::ClientSession::FeedType(), {c_handle.program_shape}, {release}, &outputs)); - xla::ProgramShape program_shape; - EXPECT_TRUE(program_shape.ParseFromString(outputs[0].vec()(0))); + xla::ProgramShapeProto program_shape_proto; + EXPECT_TRUE(program_shape_proto.ParseFromString(outputs[0].vec()(0))); + xla::ProgramShape program_shape(program_shape_proto); EXPECT_EQ(program_shape.parameters_size(), 1); VLOG(2) << "Param: " @@ -520,7 +845,7 @@ TEST(RawApiTest, CompileWithXlaReturnShapes) { << xla::ShapeUtil::HumanStringWithLayout(program_shape.result()); xla::ProgramShape xla_program_shape = - XlaCompiledProgramShape(xla_computation, *shapes); + XlaCompiledProgramShape(xla_computation, xla::ProgramShape(*shapes)); EXPECT_TRUE(xla::LayoutUtil::Equal( xla::ShapeUtil::GetSubshape(program_shape.parameters(0), {0}).layout(), xla::ShapeUtil::GetSubshape(xla_program_shape.parameters(0), {0}) @@ -537,21 +862,19 @@ TEST(RawApiTest, DotGeneralWithLayoutTest) { auto layout = xla::LayoutUtil::MakeLayout({0, 1}); xrt::XLAAllocation p0; - p0.set_device_ordinal(0); *p0.mutable_value() = FloatMatrix({{1.0f, 2.0f}, {3.0f, 4.0f}}, layout); xrt::XLAAllocation p1; - p1.set_device_ordinal(0); *p1.mutable_value() = FloatMatrix({{8.0f}, {5.0f}}, layout); xrt::XLAComputation c; auto config = c.mutable_config(); auto shapes = config->mutable_program_shape(); *shapes->add_parameters() = - xla::ShapeUtil::MakeShapeWithLayout(xla::F32, {2, 2}, {0, 1}); + xla::ShapeUtil::MakeShapeWithLayout(xla::F32, {2, 2}, {0, 1}).ToProto(); *shapes->add_parameters() = - xla::ShapeUtil::MakeShapeWithLayout(xla::F32, {2, 1}, {0, 1}); + xla::ShapeUtil::MakeShapeWithLayout(xla::F32, {2, 1}, {0, 1}).ToProto(); *shapes->mutable_result() = - xla::ShapeUtil::MakeShapeWithLayout(xla::F32, {2, 1}, {0, 1}); + xla::ShapeUtil::MakeShapeWithLayout(xla::F32, {2, 1}, {0, 1}).ToProto(); StoreComputationSnapshot(Dot(), c.mutable_hlo_snapshot()); xrt::XRTExecutionConfig e; @@ -592,7 +915,7 @@ TEST(RawApiTest, CompileAndExecuteZeroArg) { xrt::XLAComputation c; auto config = c.mutable_config(); auto shapes = config->mutable_program_shape(); - *shapes->mutable_result() = xla::ShapeUtil::MakeShape(xla::F32, {}); + *shapes->mutable_result() = xla::ShapeUtil::MakeShape(xla::F32, {}).ToProto(); xrt::XRTExecutionConfig e; e.set_release_input_handles(true); @@ -623,19 +946,20 @@ TEST(RawApiTest, CompileAndExecuteZeroArg) { TEST(RawApiTest, CompileAndExecuteReturnTuple) { xrt::XLAAllocation p0; - p0.set_device_ordinal(0); *p0.mutable_value() = FloatVector({1.0f, 2.0f}); xrt::XLAAllocation p1; - p1.set_device_ordinal(0); *p1.mutable_value() = FloatVector({8.0f, 5.0f}); xrt::XLAComputation c; auto config = c.mutable_config(); auto shapes = config->mutable_program_shape(); - *shapes->add_parameters() = xla::ShapeUtil::MakeShape(xla::F32, {2}); - *shapes->add_parameters() = xla::ShapeUtil::MakeShape(xla::F32, {2}); - *shapes->mutable_result() = xla::ShapeUtil::MakeTupleShape( - {xla::ShapeUtil::MakeShape(xla::F32, {2})}); + *shapes->add_parameters() = + xla::ShapeUtil::MakeShape(xla::F32, {2}).ToProto(); + *shapes->add_parameters() = + xla::ShapeUtil::MakeShape(xla::F32, {2}).ToProto(); + *shapes->mutable_result() = + xla::ShapeUtil::MakeTupleShape({xla::ShapeUtil::MakeShape(xla::F32, {2})}) + .ToProto(); StoreComputationSnapshot(AddAndTuple(), c.mutable_hlo_snapshot()); xrt::XRTExecutionConfig e; @@ -671,14 +995,79 @@ TEST(RawApiTest, CompileAndExecuteReturnTuple) { EXPECT_TRUE(CompareLiteralToLiteralProto(expected, response)); } +TEST(RawApiTest, CompileAndExecuteReturnExplodedTuple) { + xrt::XLAAllocation p0; + *p0.mutable_value() = xla::LiteralUtil::CreateR0(12.0f).ToProto(); + + xrt::XLAAllocation p1; + *p1.mutable_value() = xla::LiteralUtil::CreateR0(3.0f).ToProto(); + + xrt::XLAComputation c; + auto config = c.mutable_config(); + auto shapes = config->mutable_program_shape(); + *shapes->add_parameters() = xla::ShapeUtil::MakeShape(xla::F32, {}).ToProto(); + *shapes->add_parameters() = xla::ShapeUtil::MakeShape(xla::F32, {}).ToProto(); + *shapes->mutable_result() = + xla::ShapeUtil::MakeTupleShape({xla::ShapeUtil::MakeShape(xla::F32, {}), + xla::ShapeUtil::MakeShape(xla::F32, {})}) + .ToProto(); + StoreComputationSnapshot(AddAndSubTuple(), c.mutable_hlo_snapshot()); + + 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 = + ops::Const(root.WithDevice("/device:CPU:0"), e.SerializeAsString()); + auto computation = + ops::Const(root.WithDevice("/device:CPU:0"), c.SerializeAsString()); + auto c_handle = ops::XRTCompile(root, computation); + auto p0_value = + ops::Const(root.WithDevice("/device:CPU:0"), p0.SerializeAsString()); + auto p0_handle = ops::XRTAllocate(root, p0_value); + auto p1_value = + ops::Const(root.WithDevice("/device:CPU:0"), p1.SerializeAsString()); + auto p1_handle = ops::XRTAllocate(root, p1_value); + auto result = ops::XRTExecute(root, c_handle.handle, e_config, + {Output(p0_handle), Output(p1_handle)}); + TF_ASSERT_OK(root.status()); + + ClientSession session(root); + std::vector outputs; + TF_EXPECT_OK(session.Run({result}, &outputs)); + EXPECT_EQ(outputs.size(), 1); + + auto handles_vec = outputs.front().vec(); + EXPECT_EQ(handles_vec.size(), 2); + + const float kResults[2] = {15.0f, 9.0f}; + for (int64 i = 0; i < handles_vec.size(); ++i) { + auto read_back = ops::XRTReadLiteralAndRelease(root, Input(handles_vec(i))); + std::vector voutputs; + TF_EXPECT_OK(session.Run({read_back}, &voutputs)); + EXPECT_EQ(voutputs.size(), 1); + + xla::LiteralProto response; + EXPECT_TRUE(response.ParseFromString(voutputs[0].scalar()())); + + auto expected = xla::LiteralUtil::CreateR0(kResults[i]); + EXPECT_TRUE(CompareLiteralToLiteralProto(expected, response)); + } +} + TEST(RawApiTest, LeakCompilationReference) { xrt::XLAComputation c; auto config = c.mutable_config(); auto shapes = config->mutable_program_shape(); - *shapes->add_parameters() = xla::ShapeUtil::MakeShape(xla::F32, {2}); - *shapes->add_parameters() = xla::ShapeUtil::MakeShape(xla::F32, {2}); - *shapes->mutable_result() = xla::ShapeUtil::MakeTupleShape( - {xla::ShapeUtil::MakeShape(xla::F32, {2})}); + *shapes->add_parameters() = + xla::ShapeUtil::MakeShape(xla::F32, {2}).ToProto(); + *shapes->add_parameters() = + xla::ShapeUtil::MakeShape(xla::F32, {2}).ToProto(); + *shapes->mutable_result() = + xla::ShapeUtil::MakeTupleShape({xla::ShapeUtil::MakeShape(xla::F32, {2})}) + .ToProto(); StoreComputationSnapshot(AddAndTuple(), c.mutable_hlo_snapshot()); Scope root = Scope::NewRootScope().WithDevice(DeviceFromFlag()); @@ -692,25 +1081,125 @@ 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.set_device_ordinal(0); *p0.mutable_value() = xla::LiteralUtil::CreateR0(11031965).ToProto(); xrt::XLAAllocation p1; - p1.set_device_ordinal(0); *p1.mutable_value() = xla::LiteralUtil::CreateR0(4091934).ToProto(); xrt::XLAComputation c; auto config = c.mutable_config(); auto shapes = config->mutable_program_shape(); - *shapes->add_parameters() = xla::ShapeUtil::MakeShape(xla::S64, {}); - *shapes->add_parameters() = xla::ShapeUtil::MakeShape(xla::S64, {}); - *shapes->mutable_result() = xla::ShapeUtil::MakeShape(xla::S64, {}); + *shapes->add_parameters() = xla::ShapeUtil::MakeShape(xla::S64, {}).ToProto(); + *shapes->add_parameters() = xla::ShapeUtil::MakeShape(xla::S64, {}).ToProto(); + *shapes->mutable_result() = xla::ShapeUtil::MakeShape(xla::S64, {}).ToProto(); StoreComputationSnapshot(AddS64(), c.mutable_hlo_snapshot()); 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 = @@ -739,11 +1228,11 @@ TEST(RawApiTest, CompileAndExecuteWithS64Argument) { auto expected = xla::LiteralUtil::CreateR0(15123899); EXPECT_TRUE(CompareLiteralToLiteralProto(expected, response)); - xla::ProgramShape program_shape; + xla::ProgramShapeProto program_shape; EXPECT_TRUE(program_shape.ParseFromString(outputs[1].vec()(0))); EXPECT_EQ(program_shape.parameters_size(), 2); - EXPECT_TRUE( - xla::ShapeUtil::HasPrimitiveType(program_shape.result(), xla::S64)); + EXPECT_TRUE(xla::ShapeUtil::HasPrimitiveType( + xla::Shape(program_shape.result()), xla::S64)); } } // namespace diff --git a/tensorflow/compiler/xrt/xrt.proto b/tensorflow/compiler/xrt/xrt.proto index 5678f0905ff5b8956e0811026e7450acba8815e9..84adee7392825d408dd88dd74dc0c1bc7b06d7c4 100644 --- a/tensorflow/compiler/xrt/xrt.proto +++ b/tensorflow/compiler/xrt/xrt.proto @@ -3,9 +3,28 @@ syntax = "proto3"; package xrt; import "tensorflow/compiler/tf2xla/host_compute_metadata.proto"; +import "tensorflow/compiler/xla/xla.proto"; import "tensorflow/compiler/xla/xla_data.proto"; import "tensorflow/compiler/xla/service/hlo.proto"; +message DeviceAssignment { + message ComputationDevice { + message DeviceMeshCoordinates { + // The mesh coordinates for the device. Usually (X, Y, Core), in the order + // in which they are returned in the TopologyProto. + // X = value(0) + // Y = value(1) + // Core = value(2) + repeated int32 value = 1; + } + // As many replicas as there are in the replicated computation. + repeated DeviceMeshCoordinates replica_devices = 1; + } + // As many ComputationDevice as many there are computations (number + // of cores per replica). + repeated ComputationDevice computation_devices = 1; +} + // Options for an XLA compilation. message XLAComputationConfig { // The number of replicas the computation will be run on. If this is @@ -18,11 +37,18 @@ message XLAComputationConfig { tensorflow.tf2xla.HostComputeMetadata host_compute_metadata = 3; // The arg/result shapes for the whole computation. - xla.ProgramShape program_shape = 4; + xla.ProgramShapeProto program_shape = 4; // The arg/result shapes for each core of a model-parallel // computation. per_core_args_and_result_shapes is optional for a // single-core computation. - repeated xla.ProgramShape per_core_program_shape = 5; + repeated xla.ProgramShapeProto per_core_program_shape = 5; + // Describes how replicated computation instances should be assigned to + // devices. There are num_cores_per_replica computations, and each one will be + // sent and executed to the set of replica device numbers described in the + // DeviceAssignment proto. + DeviceAssignment device_assignment = 6; + // The debugging options to be passed to the XLA compilation process. + xla.DebugOptions debug_options = 7; } // Options and XLA computation for a compilation. @@ -33,7 +59,7 @@ message XLAComputation { // Literal to allocate space for, and transfer to, device memory. message XLAAllocation { - int32 device_ordinal = 1; + reserved 1; xla.LiteralProto value = 2; } @@ -75,4 +101,8 @@ message XRTExecutionConfig { bool release_input_handles = 5; // If true, release the handle to the computation after running. bool release_compilation_handle = 6; + // If set to true, and the result shape is a tuple, then instead of returning + // a single tuple allocation the execution will return a vector of + // allocations, one for each of the first-level elements of the result tuple. + bool return_exploded_tuple = 7; } 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_device.cc b/tensorflow/compiler/xrt/xrt_device.cc index ea40e6c895c4f6af13b74735685f2c342181ada9..34cb64742a20985b29d8e153bbaf5ee184fd385d 100644 --- a/tensorflow/compiler/xrt/xrt_device.cc +++ b/tensorflow/compiler/xrt/xrt_device.cc @@ -43,4 +43,12 @@ namespace tensorflow { return Status::OK(); } +/*static*/ Status XRTGenericDeviceAccessor::InitScopedRef( + OpKernelContext* ctx, ScopedRef* scoped_ref) { + const XlaDevice::Metadata* metadata; + TF_RETURN_IF_ERROR(XlaDevice::GetMetadata(ctx, &metadata)); + scoped_ref->Acquire(metadata->client()); + return Status::OK(); +} + } // namespace tensorflow diff --git a/tensorflow/compiler/xrt/xrt_device.h b/tensorflow/compiler/xrt/xrt_device.h index 1e3fddd2a72a3657d1e115375133c244772ea9f3..fb010651d9bf76c540517b9596e472c881241d8a 100644 --- a/tensorflow/compiler/xrt/xrt_device.h +++ b/tensorflow/compiler/xrt/xrt_device.h @@ -59,6 +59,8 @@ class XRTGenericDeviceAccessor { static Status InitScopedRef(OpKernelContext* ctx, int device_ordinal, ScopedRef* scoped_ref); + + static Status InitScopedRef(OpKernelContext* ctx, ScopedRef* scoped_ref); }; } // namespace tensorflow diff --git a/tensorflow/compiler/xrt/xrt_state.cc b/tensorflow/compiler/xrt/xrt_state.cc index 3a99820d7aa9e9546cc95385fd98c05f28988e9e..1e2a9584f88b73d7c92a929e93af60376a59170b 100644 --- a/tensorflow/compiler/xrt/xrt_state.cc +++ b/tensorflow/compiler/xrt/xrt_state.cc @@ -19,6 +19,7 @@ limitations under the License. #include "tensorflow/compiler/xrt/xrt_state.h" #include +#include #include #include #include @@ -34,6 +35,7 @@ limitations under the License. #include "tensorflow/core/framework/resource_mgr.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/lib/random/random.h" +#include "tensorflow/core/platform/mutex.h" #include "tensorflow/core/platform/types.h" #include "tensorflow/stream_executor/stream_executor.h" @@ -41,6 +43,34 @@ namespace tensorflow { namespace { +class BufferAllocStats { + public: + struct Stats { + int64 count = 0; + int64 size = 0; + }; + + Stats ReportAlloc(int64 device, int64 msize) { + mutex_lock lock(lock_); + Stats* device_stats = &stats_[device]; + device_stats->count += 1; + device_stats->size += msize; + return *device_stats; + } + + Stats ReportFree(int64 device, int64 msize) { + mutex_lock lock(lock_); + Stats* device_stats = &stats_[device]; + device_stats->count -= 1; + device_stats->size -= msize; + return *device_stats; + } + + private: + mutable mutex lock_; + std::map stats_; +}; + const char* kTupleContainer = "tuples"; int64 get_uid() { @@ -48,6 +78,11 @@ int64 get_uid() { return static_cast(unsigned_rand); } +BufferAllocStats* GetAllocStats() { + static BufferAllocStats* stats = new BufferAllocStats(); + return stats; +} + Status AllocateScopedShapedBuffer( xla::Backend* backend, int device_ordinal, const xla::Shape& shape, std::unique_ptr* buffer) { @@ -98,11 +133,22 @@ 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) {} + allocator_(allocator) { + if (VLOG_IS_ON(2)) { + auto stats = + GetAllocStats()->ReportAlloc(device_ordinal_, allocation_.size()); + LOG(INFO) << "XRT Allocation Stats: device=" << device_ordinal_ + << " count=" << stats.count << " size=" << stats.size; + } +} XRTBufferAllocation::~XRTBufferAllocation() { + if (VLOG_IS_ON(2)) { + GetAllocStats()->ReportFree(device_ordinal_, allocation_.size()); + } // Deallocate explicitly allows allocation_ to be null. Status s = allocator_->Deallocate(device_ordinal_, allocation_); // Nothing to do but check fail here if memory datastructures are corrupted. @@ -136,7 +182,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(); @@ -178,11 +224,36 @@ 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(); } +Status XRTTupleAllocation::WriteLiteral(xla::Backend* backend, + const xla::Literal& literal) { + if (!xla::ShapeUtil::Equal(literal.shape(), on_host_shape())) { + return errors::InvalidArgument( + "New literal shape not matching the existing one: literal=", + xla::ShapeUtil::HumanStringWithLayout(literal.shape()), + " device=", xla::ShapeUtil::HumanStringWithLayout(on_host_shape())); + } + auto transfer_manager = backend->transfer_manager(); + TF_ASSIGN_OR_RETURN(auto stream, backend->BorrowStream(device_ordinal())); + return transfer_manager->TransferLiteralToDevice(stream.get(), literal, + ToShapedBuffer()); +} + void XRTTupleAllocation::DiscardAllocation( const xla::ShapeIndex& buffer_index) { buffers_.element(buffer_index)->DiscardAllocation(); @@ -213,6 +284,11 @@ const se::DeviceMemoryBase& XRTTupleAllocation::root_allocation() { return rm->Delete(kTupleContainer, key_string); } +/* static */ Status XRTTupleAllocation::ReleaseAllAllocations(ResourceMgr* rm) { + VLOG(1) << "Releasing all XRT held device memory"; + return rm->Cleanup(kTupleContainer); +} + // Helper typedef to make ShapeTree ForEach helper lambda signatures more // readable. They need a type of const T& where in this case T is the // following pointer. @@ -441,11 +517,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 73b5584e38f781343fe6793af7ad28232fbfc184..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); @@ -129,6 +137,10 @@ class XRTTupleAllocation : public ResourceBase { // Deletes the reference in the rm to an allocation interned under key. static Status DeleteFromResourceManager(ResourceMgr* rm, int64 key); + // Releases all the device memory allocated by XRT within the resource + // manager. + static Status ReleaseAllAllocations(ResourceMgr* rm); + // Adds the allocation to a ResourceMgr and returns the key that will be used // to retrieve it. Transfers a reference on *this to rm. Status Intern(ResourceMgr* rm, int64* key); @@ -137,6 +149,9 @@ class XRTTupleAllocation : public ResourceBase { Status ToLiteral(xla::Backend* backend, int device_ordinal, xla::Literal* literal); + // Write a new literal value to the allocation. + Status WriteLiteral(xla::Backend* backend, const xla::Literal& literal); + // True if none of the buffers in the allocation are aliased by any other live // handle. bool IsExclusiveOwner(); @@ -161,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/compiler/xrt/xrt_util.cc b/tensorflow/compiler/xrt/xrt_util.cc new file mode 100644 index 0000000000000000000000000000000000000000..3ef8bedc7324696cd255c72a851f0f2410e03848 --- /dev/null +++ b/tensorflow/compiler/xrt/xrt_util.cc @@ -0,0 +1,76 @@ +/* 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/xrt/xrt_util.h" + +#include +#include + +#include "tensorflow/compiler/xla/debug_options_flags.h" +#include "tensorflow/compiler/xla/types.h" +#include "tensorflow/core/platform/logging.h" + +namespace tensorflow { +namespace { + +bool DebugOptionsPassThroughEnabled() { + const char* env = getenv("TF_XLA_DEBUG_OPTIONS_PASSTHROUGH"); + bool enabled = + env != nullptr && (strcmp(env, "1") == 0 || strcmp(env, "true") == 0); + if (enabled) { + LOG(WARNING) << "Passing through XLA debug options!"; + } else { + LOG(WARNING) << "TF_XLA_DEBUG_OPTIONS_PASSTHROUGH not set, not all options " + "will be retained"; + } + return enabled; +} + +string SafeDebugPath(const string& path) { + if (path.empty() || path.compare(0, 5, "gs://") == 0 || + path.compare(0, 11, "bigstore://") == 0) { + return path; + } + LOG(WARNING) << "Invalid config path (will be dropped): " << path; + return string(); +} + +} // namespace + +xla::DebugOptions BuildXlaDebugOptions(const xla::DebugOptions& ref_options) { + static const bool options_passthrough = DebugOptionsPassThroughEnabled(); + if (options_passthrough) { + return ref_options; + } + xla::DebugOptions options = xla::GetDebugOptionsFromFlags(); + options.set_xla_generate_hlo_text_to( + SafeDebugPath(ref_options.xla_generate_hlo_text_to())); + options.set_xla_dump_optimized_hlo_proto_to( + SafeDebugPath(ref_options.xla_dump_optimized_hlo_proto_to())); + options.set_xla_dump_computations_to( + SafeDebugPath(ref_options.xla_dump_computations_to())); + options.set_xla_dump_executions_to( + SafeDebugPath(ref_options.xla_dump_executions_to())); + for (auto& pass : ref_options.xla_disable_hlo_passes()) { + options.add_xla_disable_hlo_passes(pass); + } + options.set_xla_dump_unoptimized_hlo_proto_to( + SafeDebugPath(ref_options.xla_dump_unoptimized_hlo_proto_to())); + options.set_xla_dump_per_pass_hlo_proto_to( + SafeDebugPath(ref_options.xla_dump_per_pass_hlo_proto_to())); + return options; +} + +} // namespace tensorflow diff --git a/tensorflow/compiler/xrt/xrt_util.h b/tensorflow/compiler/xrt/xrt_util.h new file mode 100644 index 0000000000000000000000000000000000000000..d9c05a7f3406313f99ae214d67b34e8e7de8be3e --- /dev/null +++ b/tensorflow/compiler/xrt/xrt_util.h @@ -0,0 +1,34 @@ +/* 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. +==============================================================================*/ + +// Utility functions in support of the XRT API. + +#ifndef TENSORFLOW_COMPILER_XRT_XRT_UTIL_H_ +#define TENSORFLOW_COMPILER_XRT_XRT_UTIL_H_ + +#include "tensorflow/compiler/xla/xla.pb.h" + +namespace tensorflow { + +// Filters the debug options provided as argument according to the value of the +// TF_XLA_DEBUG_OPTIONS_PASSTHROUGH environment variable. If such variable is +// set to "1" or "true", the debug options will be returned as is. Otherwise +// only a subset of them will be set in the returned ones, and all the paths +// contained in it, will be limited to gs:// and bigstore:// ones. +xla::DebugOptions BuildXlaDebugOptions(const xla::DebugOptions& ref_options); + +} // namespace tensorflow + +#endif // TENSORFLOW_COMPILER_XRT_XRT_UTIL_H_ diff --git a/tensorflow/contrib/BUILD b/tensorflow/contrib/BUILD index 78ad19a4ab112be08569c857c8ed4e16ceed6d80..a4c3d9623adfe3133af0c6ea055586b9544e659b 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", @@ -72,7 +71,6 @@ py_library( "//tensorflow/contrib/metrics:metrics_py", "//tensorflow/contrib/mixed_precision:mixed_precision", "//tensorflow/contrib/model_pruning", - "//tensorflow/contrib/nccl:nccl_py", "//tensorflow/contrib/nearest_neighbor:nearest_neighbor_py", "//tensorflow/contrib/nn:nn_py", "//tensorflow/contrib/opt:opt_py", @@ -179,9 +177,7 @@ cc_library( "//tensorflow/contrib/tensor_forest:stats_ops_kernels", "//tensorflow/contrib/tensor_forest:tensor_forest_kernels", "//tensorflow/contrib/text:all_kernels", - ] + if_mpi(["//tensorflow/contrib/mpi_collectives:mpi_collectives_py"]) + if_cuda([ - "//tensorflow/contrib/nccl:nccl_kernels", - ]) + select({ + ] + if_mpi(["//tensorflow/contrib/mpi_collectives:mpi_collectives_py"]) + select({ "//tensorflow:android": [], "//tensorflow:ios": [], "//tensorflow:linux_s390x": [], @@ -200,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", ]), ) @@ -215,7 +211,6 @@ cc_library( "//tensorflow/contrib/hadoop:dataset_ops_op_lib", "//tensorflow/contrib/input_pipeline:input_pipeline_ops_op_lib", "//tensorflow/contrib/layers:sparse_feature_cross_op_op_lib", - "//tensorflow/contrib/nccl:nccl_ops_op_lib", "//tensorflow/contrib/nearest_neighbor:nearest_neighbor_ops_op_lib", "//tensorflow/contrib/rnn:all_ops", "//tensorflow/contrib/seq2seq:beam_search_ops_op_lib", diff --git a/tensorflow/contrib/__init__.py b/tensorflow/contrib/__init__.py index f52a1a7babceeae93cdd2e5a93dad413a1d30191..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 @@ -62,7 +63,6 @@ from tensorflow.contrib import memory_stats from tensorflow.contrib import metrics from tensorflow.contrib import mixed_precision from tensorflow.contrib import model_pruning -from tensorflow.contrib import nccl from tensorflow.contrib import nn from tensorflow.contrib import opt from tensorflow.contrib import periodic_resample @@ -92,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 @@ -104,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/all_reduce/BUILD b/tensorflow/contrib/all_reduce/BUILD index 881808a98bfd688c2efaa8beb5b8f11a2527fee8..f6c6560c1c354ed8a36b98b1f564835eb9958e55 100644 --- a/tensorflow/contrib/all_reduce/BUILD +++ b/tensorflow/contrib/all_reduce/BUILD @@ -9,8 +9,6 @@ licenses(["notice"]) # Apache 2.0 exports_files(["LICENSE"]) -load("//tensorflow:tensorflow.bzl", "tf_py_test") - py_library( name = "all_reduce_py", srcs = ["__init__.py"], @@ -29,29 +27,6 @@ py_library( srcs_version = "PY2AND3", visibility = ["//visibility:public"], deps = [ - "//tensorflow/contrib/nccl:nccl_py", - "//tensorflow/python:array_ops", - "//tensorflow/python:framework_ops", - "//tensorflow/python:math_ops", - ], -) - -tf_py_test( - name = "all_reduce_test", - srcs = ["python/all_reduce_test.py"], - additional_deps = [ - ":all_reduce", - "//third_party/py/numpy", - "//tensorflow/core:protos_all_py", - "//tensorflow/python:array_ops", - "//tensorflow/python:client", - "//tensorflow/python:framework_ops", - "//tensorflow/python:framework_test_lib", - "//tensorflow/python:math_ops", - "//tensorflow/python:constant_op", - "//tensorflow/python:client_testlib", - "//tensorflow/python:platform", - "//tensorflow/python:platform_test", - "//tensorflow/python:state_ops", + "//tensorflow/python/distribute:all_reduce", ], ) diff --git a/tensorflow/contrib/all_reduce/python/all_reduce.py b/tensorflow/contrib/all_reduce/python/all_reduce.py index 3b539734a236804026826a8117d9c668c0dd089a..238cdaf8a79812df3f043d9d070bbcfd443f6e1e 100644 --- a/tensorflow/contrib/all_reduce/python/all_reduce.py +++ b/tensorflow/contrib/all_reduce/python/all_reduce.py @@ -18,842 +18,5 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import collections -import math - -from tensorflow.contrib import nccl -from tensorflow.python.framework import device as device_lib -from tensorflow.python.framework import ops -from tensorflow.python.ops import array_ops -from tensorflow.python.ops import math_ops - - -def _flatten_tensors(tensors): - """Check tensors for isomorphism and flatten. - - Args: - tensors: list of T `tf.Tensor` which must all have the same shape. - - Returns: - tensors: a list of T `tf.Tensor` which are flattened (1D) views of tensors - shape: the original shape of each element of input tensors - - Raises: - ValueError: tensors are empty or non-isomorphic or have unknown shape. - """ - if not tensors: - raise ValueError("tensors cannot be empty") - shape = tensors[0].shape - for tensor in tensors: - shape = shape.merge_with(tensor.shape) - if not shape.is_fully_defined(): - raise ValueError("Tensors must have statically known shape.") - if len(shape) != 1: - reshaped = [] - for t in tensors: - with ops.colocate_with(t): - reshaped.append(array_ops.reshape(t, [-1])) - tensors = reshaped - return tensors, shape - - -def _reshape_tensors(tensors, shape): - """Reshape tensors flattened by _flatten_tensors. - - Args: - tensors: list of T `tf.Tensor` of identical length 1D tensors. - shape: list of integers describing the desired shape. Product of - the elements must equal the length of each tensor. - - Returns: - list of T `tf.Tensor` which are the reshaped inputs. - """ - reshaped = [] - for t in tensors: - with ops.colocate_with(t): - reshaped.append(array_ops.reshape(t, shape)) - return reshaped - - -def _padded_split(tensor, pieces): - """Like split for 1D tensors but pads-out case where len % pieces != 0. - - Args: - tensor: T `tf.Tensor` that must be 1D. - pieces: a positive integer specifying the number of pieces into which - tensor should be split. - - Returns: - list of T `tf.Tensor` of length pieces, which hold the values of - thin input tensor, in order. The final tensor may - be zero-padded on the end to make its size equal to those of all - of the other tensors. - - Raises: - ValueError: The input tensor is not 1D. - """ - shape = tensor.shape - if 1 != len(shape): - raise ValueError("input tensor must be 1D") - tensor_len = shape[0].value - with ops.colocate_with(tensor): - if tensor_len % pieces != 0: - # pad to an even length - chunk_size = 1 + tensor_len // pieces - if pieces > tensor_len: - # This is an edge case that should not come up in practice, - # i.e. a different reduction algorithm would be better, - # but we'll make it work just for completeness. - pad_len = pieces - tensor_len - extended_whole = array_ops.concat( - [tensor, array_ops.zeros([pad_len], dtype=tensor.dtype)], 0) - parts = array_ops.split(extended_whole, pieces) - return parts, pad_len - elif (pieces - 1) * chunk_size >= tensor_len: - # Another edge case of limited real interest. - pad_len = (pieces * chunk_size) % tensor_len - extended_whole = array_ops.concat( - [tensor, array_ops.zeros([pad_len], dtype=tensor.dtype)], 0) - parts = array_ops.split(extended_whole, pieces) - return parts, pad_len - else: - last_chunk_size = tensor_len - (pieces - 1) * chunk_size - pad_len = chunk_size - last_chunk_size - piece_lens = [chunk_size for _ in range(pieces - 1)] + [last_chunk_size] - parts = array_ops.split(tensor, piece_lens) - parts[-1] = array_ops.concat( - [parts[-1], array_ops.zeros([pad_len], dtype=tensor.dtype)], 0) - return parts, pad_len - else: - return array_ops.split(tensor, pieces), 0 - - -def _strip_padding(tensors, pad_len): - """Strip the suffix padding added by _padded_split. - - Args: - tensors: list of T `tf.Tensor` of identical length 1D tensors. - pad_len: number of elements to be stripped from the end of each tensor. - - Returns: - list of T `tf.Tensor` which are the stripped inputs. - - Raises: - ValueError: tensors must be a non-empty list of 1D tensors, and - each must be longer than pad_len. - """ - if not tensors: - raise ValueError("tensors cannot be empty") - shape = tensors[0].shape - if len(shape) > 1: - raise ValueError("tensors must be 1D") - prefix_len = int(shape[0] - pad_len) - if prefix_len < 0: - raise ValueError("pad_len longer than tensor") - stripped = [] - for t in tensors: - with ops.colocate_with(t): - stripped.append(array_ops.slice(t, [0], [prefix_len])) - return stripped - - -def _ragged_split(tensor, pieces): - """Like split for 1D tensors but allows case where len % pieces != 0. - - Args: - tensor: T `tf.Tensor` that must be 1D. - pieces: a positive integer specifying the number of pieces into which - tensor should be split. - - Returns: - list of T `tf.Tensor` of length pieces, which hold the values of - the input tensor, in order. The final tensor may be shorter - than the others, which will all be of equal length. - - Raises: - ValueError: input tensor must be 1D. - """ - shape = tensor.shape - if 1 != len(shape): - raise ValueError("input tensor must be 1D") - tensor_len = shape[0].value - chunk_size = tensor_len // pieces - with ops.colocate_with(tensor): - if tensor_len != (pieces * chunk_size): - # last piece will be short - assert pieces > 1 - last_chunk_size = tensor_len - ((pieces - 1) * chunk_size) - assert last_chunk_size > 0 - piece_lens = [chunk_size for _ in range(pieces - 1)] + [last_chunk_size] - return array_ops.split(tensor, piece_lens) - else: - return array_ops.split(tensor, pieces) - - -def _ring_permutations(num_workers, num_subchunks, gpu_perm): - """"Generate an array of device index arrays, one for each subchunk. - - In the basic ring reduction algorithm there are size(T)/num_devices - data chunks and each device process one chunk per tick, i.e. sending - one chunk and receiving one chunk. The idea of subchunking is that - each device processes num_subchunks smaller data regions per tick, - and the ring rank permutation is different for each subchunk index - so that a device is potentially sending to and receiving from - num_subchunks different other devices at each tick. Where multiple - independent data channels exist between devices, this strategy - supplies a method of using them in parallel. - - Args: - num_workers: number of worker tasks - num_subchunks: number of subchunks into which to divide each per-GPU chunk. - gpu_perm: an array of integers in [0, num_gpus-1] giving the default - ring order of GPUs at each worker. Other permutations will be generated - by rotating this array and splicing together per-worker instances. - - Raises: - ValueError: the number of subchunks may not exceed the number of GPUs. - - Returns: - pred_by_s_d: list of lists that maps (by index) from (subchunk, dev) to - preceding device in the permutation for that subchunk. The - device index of GPU i at worker j is i + (j * num_gpus). - rank_by_s_d: list of lists that maps (by index) from (subchunk, dev) to - local rank of device d in the permutation for that subchunk. - """ - num_gpus = len(gpu_perm) - devices = num_workers * num_gpus - if devices == 0: - return [], [] - if num_subchunks > num_gpus: - raise ValueError( - "num_subchunks %d must be <= num_gpus %d" % (num_subchunks, num_gpus)) - rotation_interval = max(1, int(num_gpus / num_subchunks)) - perms_by_s = [] - for s in range(0, num_subchunks): - full_order = [] - offset = s * rotation_interval - for w in range(0, num_workers): - default_order = [(w * num_gpus) + i for i in gpu_perm] - dev_order = default_order[offset:] + default_order[:offset] - full_order += dev_order - perms_by_s.append(full_order) - pred_by_s_d = [[-1 for d in range(0, devices)] - for s in range(0, num_subchunks)] - rank_by_s_d = [[-1 for d in range(0, devices)] - for s in range(0, num_subchunks)] - for s in range(0, num_subchunks): - for d in range(0, devices): - for t in range(0, devices): - if d == perms_by_s[s][t]: - rank_by_s_d[s][d] = t - pred_by_s_d[s][d] = perms_by_s[s][(t + devices - 1) % devices] - break - return (pred_by_s_d, rank_by_s_d) - - -def build_ring_all_reduce(input_tensors, num_workers, num_subchunks, - gpu_perm, red_op, un_op=None): - """Construct a subgraph performing a ring-style all-reduce of input_tensors. - - Args: - input_tensors: a list of T `tf.Tensor` objects, which must all - have the same shape and type. - num_workers: number of worker tasks spanned by input_tensors. - num_subchunks: number of subchunks each device should process in one tick. - gpu_perm: a list of ints giving a ring-wise rank ordering of GPUs at - each worker. All workers must have the same number of - GPUs with the same rank ordering. If NVLINK is available, this should - be a ring order supported by NVLINK edges. - red_op: a binary operator for elementwise reduction. - un_op: an optional unary operator to apply to fully reduced values. - - Raises: - ValueError: empty input_tensors or they don't all have same - size. - - Returns: - a list of T `tf.Tensor` identical sum-reductions of input_tensors. - """ - if len(input_tensors) < 2: - raise ValueError("input_tensors must be length 2 or longer") - input_tensors, shape = _flatten_tensors(input_tensors) - devices = [t.device for t in input_tensors] - (pred_by_s_d, rank_by_s_d) = _ring_permutations( - num_workers, num_subchunks, gpu_perm) - chunks_by_dev, pad_len = _build_ring_gather( - input_tensors, devices, - num_subchunks, pred_by_s_d, rank_by_s_d, red_op) - if un_op: - chunks_by_dev = _apply_unary_to_chunks(un_op, chunks_by_dev) - output_tensors = _build_ring_scatter(pred_by_s_d, rank_by_s_d, - chunks_by_dev) - if pad_len > 0: - output_tensors = _strip_padding(output_tensors, pad_len) - if len(shape) != 1: - output_tensors = _reshape_tensors(output_tensors, shape) - return output_tensors - - -def _build_ring_gather(input_tensors, devices, num_subchunks, - pred_by_s_d, rank_by_s_d, red_op): - """Construct a subgraph for the first (reduction) pass of ring all-reduce. - - Args: - input_tensors: a list of T `tf.Tensor` 1D input tensors of same - shape and type. - devices: array of device name strings - num_subchunks: number of subchunks each device should process in one tick. - pred_by_s_d: as produced by _ring_permutations - rank_by_s_d: as produced by _ring_permutations - red_op: a binary operator for elementwise reduction - - Raises: - ValueError: tensors must all be one dimensional. - - Returns: - list of list of T `tf.Tensor` of (partially) reduced values where - exactly num_subchunks chunks at each device are fully reduced. - """ - num_devices = len(input_tensors) - if num_devices == 0: - return [] - if num_devices == 1: - return input_tensors - shape = input_tensors[0].shape - if 1 != len(shape): - raise ValueError("input tensors must be 1D") - num_chunks = num_devices * num_subchunks - num_ticks = num_devices - 1 - # Initialize chunks_by_dev with splits of the input tensors. - chunks_by_dev = [] - split_pad_len = 0 - for d in range(0, num_devices): - with ops.device(devices[d]): - splits, split_pad_len = _padded_split(input_tensors[d], num_chunks) - chunks_by_dev.append(splits) - # Reduction phase - for tick in range(0, num_ticks): - # One new partial reduction for every chunk - new_partial_reductions = [None for _ in range(0, num_chunks)] - # Compute reductions with respect to last tick's values - for d in range(0, num_devices): - with ops.device(devices[d]): - for s in range(0, num_subchunks): - rank = rank_by_s_d[s][d] - seg_index = (rank + num_devices - (2 + tick)) % num_devices - pred_dev = pred_by_s_d[s][d] - chunk_index = (seg_index * num_subchunks) + s - new_partial_reductions[chunk_index] = red_op( - chunks_by_dev[pred_dev][chunk_index], - chunks_by_dev[d][chunk_index]) - # Update chunks_by_dev with the new values at the end of the tick. - for d in range(0, num_devices): - for s in range(0, num_subchunks): - rank = rank_by_s_d[s][d] - seg_index = (rank + num_devices - (2 + tick)) % num_devices - chunk_index = (seg_index * num_subchunks) + s - chunks_by_dev[d][chunk_index] = new_partial_reductions[chunk_index] - return chunks_by_dev, split_pad_len - - -def _apply_unary_to_chunks(f, chunks_by_dev): - """Apply a unary op to each tensor in chunks_by_dev, on same device. - - Args: - f: a unary function over T `tf.Tensor`. - chunks_by_dev: list of lists of T `tf.Tensor`. - - Returns: - new list of lists of T `tf.Tensor` with the same structure as - chunks_by_dev containing the derived tensors. - """ - output = [] - for x in chunks_by_dev: - with ops.colocate_with(x[0]): - output.append([f(t) for t in x]) - return output - - -def _build_ring_scatter(pred_by_s_d, rank_by_s_d, - chunks_by_dev): - """Construct subgraph for second (scatter) pass of ring all-reduce. - - Args: - pred_by_s_d: as produced by _ring_permutations - rank_by_s_d: as produced by _ring_permutations - chunks_by_dev: list of list of T `tf.Tensor` indexed by ints - (device, chunk) - - Raises: - ValueError: chunks_by_dev is not well-formed - - Returns: - list of T `tf.Tensor` which are the fully reduced tensors, one - at each device corresponding to the outer dimension of chunks_by_dev. - """ - num_devices = len(chunks_by_dev) - num_chunks = len(chunks_by_dev[0]) - if 0 != num_chunks % num_devices: - raise ValueError( - "Expect number of chunks per device to be divisible by num_devices") - num_subchunks = int(num_chunks / num_devices) - num_ticks = num_devices - 1 - for tick in range(0, num_ticks): - passed_values = [None for _ in range(0, num_chunks)] - for d in range(0, num_devices): - with ops.colocate_with(chunks_by_dev[d][0]): - for s in range(0, num_subchunks): - rank = rank_by_s_d[s][d] - seg_index = (rank + num_devices - (1 + tick)) % num_devices - pred_dev = pred_by_s_d[s][d] - chunk_index = (seg_index * num_subchunks) + s - passed_values[chunk_index] = array_ops.identity( - chunks_by_dev[pred_dev][chunk_index]) - for d in range(0, num_devices): - for s in range(0, num_subchunks): - rank = rank_by_s_d[s][d] - seg_index = (rank + num_devices - (1 + tick)) % num_devices - chunk_index = (seg_index * num_subchunks) + s - chunks_by_dev[d][chunk_index] = passed_values[chunk_index] - # Join chunks at each device. - output = [] - for x in chunks_by_dev: - with ops.colocate_with(x[0]): - output.append(array_ops.concat(x, 0)) - return output - - -def build_recursive_hd_all_reduce(input_tensors, red_op, un_op=None): - """Construct a subgraph for recursive halving-doubling all-reduce. - - The recursive halving-doubling algorithm is described in - http://www.mcs.anl.gov/~thakur/papers/ijhpca-coll.pdf - - The concept is to arrange the participating n devices in - a linear sequence where devices exchange data pairwise - with one other device in each round. During the gather - phase there are lg(n) rounds where devices exchange - increasingly smaller sub-tensors with another device - at increasingly greater distances, until at the top - each device has 1/n of the fully reduced values. During the - scatter phase each device exchanges its fully reduced - sub-tensor (which doubles in length at each round) - with one other device at increasingly smaller distances - until each device has all of the fully reduced values. - - Note: this preliminary version requires that len(input_tensors) be a - power of 2. TODO(tucker): relax this restriction. Also, the - number of elements in each tensor must be divisible by 2^h where h - is the number of hops in each phase. This will also be relaxed in - the future with edge-case specific logic. - - Args: - input_tensors: list of T `tf.Tensor` to be elementwise reduced. - red_op: a binary elementwise reduction Op. - un_op: an optional unary elementwise Op to apply to reduced values. - - Returns: - list of T `tf.Tensor` which are the fully reduced tensors, one - at each device of input_tensors. - - Raises: - ValueError: num_devices not a power of 2, or tensor len not divisible - by 2 the proper number of times. - """ - devices = [t.device for t in input_tensors] - input_tensors, shape = _flatten_tensors(input_tensors) - reduced_shards = _build_recursive_hd_gather(input_tensors, devices, red_op) - if un_op: - reduced_shards = [un_op(t) for t in reduced_shards] - output_tensors = _build_recursive_hd_scatter(reduced_shards, devices) - if len(shape) != 1: - output_tensors = _reshape_tensors(output_tensors, shape) - return output_tensors - - -def _build_recursive_hd_gather(input_tensors, devices, red_op): - """Construct the gather phase of recursive halving-doubling all-reduce. - - Args: - input_tensors: list of T `tf.Tensor` to be elementwise reduced. - devices: a list of strings naming the devices hosting input_tensors, - which will also be used to host the (partial) reduction values. - red_op: a binary elementwise reduction Op. - - Returns: - list of T `tf.Tensor` which are the fully reduced tensor shards. - - Raises: - ValueError: num_devices not a power of 2, or tensor len not divisible - by 2 the proper number of times. - """ - num_devices = len(devices) - num_hops = int(math.log(num_devices, 2)) - if num_devices != (2 ** num_hops): - raise ValueError("num_devices must be a power of 2") - chunks = input_tensors - for h in range(0, num_hops): - span = 2 ** h - group_size = span * 2 - new_chunks = [[] for _ in devices] - for d in range(0, num_devices): - if (d % group_size) >= (group_size / 2): - # skip right half of a pair - continue - left_dev = devices[d] - right_dev = devices[d + span] - left_split = array_ops.split(chunks[d], 2) - right_split = array_ops.split(chunks[d+span], 2) - with ops.device(left_dev): - new_chunks[d] = red_op(left_split[0], right_split[0]) - with ops.device(right_dev): - new_chunks[d + span] = red_op(left_split[1], right_split[1]) - chunks = new_chunks - return chunks - - -def _build_recursive_hd_scatter(input_tensors, devices): - """Construct the scatter phase of recursive halving-doublng all-reduce. - - Args: - input_tensors: list of T `tf.Tensor` that are fully-reduced shards. - devices: a list of strings naming the devices on which the reconstituted - full tensors should be placed. - - Returns: - list of T `tf.Tensor` which are the fully reduced tensors. - """ - num_devices = len(devices) - num_hops = int(math.log(num_devices, 2)) - assert num_devices == (2 ** num_hops), "num_devices must be a power of 2" - chunks = input_tensors - for h in reversed(range(0, num_hops)): - span = 2 ** h - group_size = span * 2 - new_chunks = [[] for _ in devices] - for d in range(0, num_devices): - if (d % group_size) >= (group_size / 2): - # skip right half of a pair - continue - left_idx = d - right_idx = d + span - left_dev = devices[left_idx] - right_dev = devices[right_idx] - with ops.device(left_dev): - new_chunks[left_idx] = array_ops.concat([chunks[left_idx], - chunks[right_idx]], 0) - with ops.device(right_dev): - new_chunks[right_idx] = array_ops.concat([chunks[left_idx], - chunks[right_idx]], 0) - chunks = new_chunks - return chunks - - -def build_shuffle_all_reduce(input_tensors, gather_devices, red_op, un_op=None): - """Construct a subgraph for shuffle all-reduce. - - Shuffle reduce is essentially the algorithm implemented when using - parameter servers. Suppose tensor length is n, there are d devices - and g gather shards. Each device sends a n/g length sub-tensor to - each gather shard. The gather shards perform a reduction across d - fragments, then broadcast the result back to each device. The - devices then join the g fully reduced fragments they receive from - the shards. The gather shards could perform d-1 pairwise - reductions, or one d-way reduction. The first is better where - reduction Op time is low compared to transmission time, the second - better in the other case. - - Args: - input_tensors: list of T @(tf.Tensor} values to be reduced. - gather_devices: list of names of devices on which reduction shards - should be placed. - red_op: an n-array elementwise reduction Op - un_op: optional elementwise unary Op to be applied to fully-reduced values. - - Returns: - list of T `tf.Tensor` which are the fully reduced tensors. - """ - input_tensors, shape = _flatten_tensors(input_tensors) - dst_devices = [t.device for t in input_tensors] - reduced_shards = _build_shuffle_gather(input_tensors, gather_devices, - red_op, un_op) - output_tensors = _build_shuffle_scatter(reduced_shards, dst_devices) - if len(shape) != 1: - output_tensors = _reshape_tensors(output_tensors, shape) - return output_tensors - - -def _build_shuffle_gather(input_tensors, gather_devices, red_op, un_op=None): - """Construct the gather (concentrate and reduce) phase of shuffle all-reduce. - - Args: - input_tensors: list of T @(tf.Tensor} values to be reduced. - gather_devices: list of names of devices on which reduction shards - should be placed. - red_op: the binary reduction Op - un_op: optional elementwise unary Op to be applied to fully-reduced values. - - Returns: - list of T `tf.Tensor` which are the fully reduced shards. - - Raises: - ValueError: inputs not well-formed. - """ - num_source_devices = len(input_tensors) - num_gather_devices = len(gather_devices) - shape = input_tensors[0].shape - if len(shape) != 1: - raise ValueError("input_tensors must be 1D") - shards_by_source = [] - for d in range(0, num_source_devices): - with ops.colocate_with(input_tensors[d]): - shards_by_source.append( - _ragged_split(input_tensors[d], num_gather_devices)) - reduced_shards = [] - for d in range(0, num_gather_devices): - with ops.device(gather_devices[d]): - values = [s[d] for s in shards_by_source] - red_shard = red_op(values) - if un_op: - red_shard = un_op(red_shard) - reduced_shards.append(red_shard) - return reduced_shards - - -def _build_shuffle_scatter(reduced_shards, dst_devices): - """Build the scatter phase of shuffle all-reduce. - - Args: - reduced_shards: list of T @(tf.Tensor} fully reduced shards - dst_devices: list of names of devices at which the fully-reduced value - should be reconstituted. - - Returns: - list of T `tf.Tensor` scattered tensors. - """ - num_devices = len(dst_devices) - out_tensors = [] - for d in range(0, num_devices): - with ops.device(dst_devices[d]): - out_tensors.append(array_ops.concat(reduced_shards, 0)) - return out_tensors - - -def _split_by_task(devices, values): - """Partition devices and values by common task. - - Args: - devices: list of device name strings - values: list of T `tf.tensor` of same length as devices. - - Returns: - (per_task_devices, per_task_values) where both values are - lists of lists with isomorphic structure: the outer list is - indexed by task, and the inner list has length of the number - of values belonging to that task. per_task_devices contains - the specific devices to which the values are local, and - per_task_values contains the corresponding values. - - Raises: - ValueError: devices must be same length as values. - """ - num_devices = len(devices) - if num_devices != len(values): - raise ValueError("len(devices) must equal len(values)") - per_task_devices = collections.OrderedDict() - per_task_values = collections.OrderedDict() - for d in range(num_devices): - d_spec = device_lib.DeviceSpec.from_string(devices[d]) - if not hasattr(d_spec, "task") or d_spec.task is None: - assert False, "failed to parse device %s" % devices[d] - index = (d_spec.job or "localhost", d_spec.replica or 0, d_spec.task) - if index not in per_task_devices: - per_task_devices[index] = [] - per_task_values[index] = [] - per_task_devices[index].append(devices[d]) - per_task_values[index].append(values[d]) - - return (list(per_task_devices.values()), list(per_task_values.values())) - - -def build_nccl_all_reduce(input_tensors, red_op, un_op=None): - """Build a subgraph that does one full all-reduce, using NCCL. - - Args: - input_tensors: list of T `tf.Tensor` of same-shape and type values to - be reduced. - red_op: binary elementwise reduction operator. Must be one of - {tf.add} - un_op: optional unary elementwise Op to apply to fully-reduce values. - - Returns: - list of T `tf.Tensor` of reduced values. - - Raises: - ValueError: red_op not supported. - """ - if red_op == math_ops.add: - output_tensors = nccl.all_sum(input_tensors) - else: - raise ValueError("red_op not supported by NCCL all-reduce: ", red_op) - if un_op: - un_op_wrapped = [] - for t in output_tensors: - with ops.colocate_with(t): - un_op_wrapped.append(un_op(t)) - output_tensors = un_op_wrapped - return output_tensors - - -def _build_nccl_hybrid(input_tensors, red_op, upper_level_f): - """Construct a subgraph for NCCL hybrid all-reduce. - - Args: - input_tensors: list of T `tf.Tensor` of same-shape and type values to - be reduced. - red_op: binary elementwise reduction operator. - upper_level_f: function for reducing one value per worker, across - workers. - - Returns: - list of T `tf.Tensor` of reduced values. - - Raises: - ValueError: inputs not well-formed. - """ - input_tensors, shape = _flatten_tensors(input_tensors) - devices = [t.device for t in input_tensors] - per_worker_devices, per_worker_values = _split_by_task(devices, input_tensors) - num_workers = len(per_worker_devices) - up_values = [None for w in range(0, num_workers)] - up_devices = up_values[:] - down_values = up_values[:] - # First stage: reduce within each worker using NCCL - for w in range(0, num_workers): - worker_values = build_nccl_all_reduce(per_worker_values[w], red_op) - # NOTE: these reductions will not run to completion unless - # every output value is used. Since we only need one, we - # need to put control dependencies on the rest. - with ops.control_dependencies(worker_values): - with ops.device(worker_values[0].device): - up_values[w] = array_ops.identity(worker_values[0]) - up_devices[w] = per_worker_devices[w][0] - # Second stage: Apply upper_level_f to reduce across first device at - # each worker - level_2_output = upper_level_f(up_values) - # Third stage: propagate within each worker using NCCL Broadcast - for w in range(0, num_workers): - dst_tensors = [] - with ops.device(per_worker_devices[w][0]): - broadcast_src = nccl.broadcast(array_ops.identity(level_2_output[w])) - for d in per_worker_devices[w]: - with ops.device(d): - dst_tensors.append(array_ops.identity(broadcast_src)) - down_values[w] = dst_tensors - output_tensors = [v for sublist in down_values for v in sublist] - if len(shape) != 1: - output_tensors = _reshape_tensors(output_tensors, shape) - return output_tensors - - -def _reduce_non_singleton(input_tensors, red_f, un_op): - """If input_tensors has more than one element apply red_f, else apply un_op.""" - if len(input_tensors) > 1: - return red_f(input_tensors) - else: - if not un_op: - return input_tensors - output_tensors = [] - for t in input_tensors: - with ops.colocate_with(t): - output_tensors.append(un_op(t)) - return output_tensors - - -def build_nccl_then_ring(input_tensors, subdiv, red_op, un_op=None): - """Construct hybrid of NCCL within workers, Ring across workers.""" - def upper_builder(y): - return build_ring_all_reduce(y, len(y), subdiv, [0], red_op, un_op) - def upper_level_f(x): - return _reduce_non_singleton(x, upper_builder, un_op) - return _build_nccl_hybrid(input_tensors, red_op, upper_level_f) - - -def build_nccl_then_recursive_hd(input_tensors, red_op, un_op=None): - """Construct hybrid of NCCL within workers, Recursive-HD across workers.""" - upper_level_f = lambda x: build_recursive_hd_all_reduce(x, red_op, un_op) - return _build_nccl_hybrid(input_tensors, red_op, upper_level_f) - - -def build_nccl_then_shuffle(input_tensors, gather_devices, nccl_red_op, - shuffle_red_op, un_op=None): - """Construct hybrid of NCCL within workers, Shuffle across workers.""" - upper_level_f = lambda x: build_shuffle_all_reduce(x, gather_devices, - shuffle_red_op, un_op) - return _build_nccl_hybrid(input_tensors, nccl_red_op, upper_level_f) - - -def _build_shuffle_hybrid(input_tensors, gather_devices, red_op, upper_level_f): - """Construct a subgraph for Shuffle hybrid all-reduce. - - Args: - input_tensors: list of T `tf.Tensor` of same-shape and type values to - be reduced. - gather_devices: list of device names on which to host gather shards. - red_op: binary elementwise reduction operator. - upper_level_f: function for reducing one value per worker, across - workers. - - Returns: - list of T `tf.Tensor` of reduced values. - - Raises: - ValueError: inputs not well-formed. - """ - input_tensors, shape = _flatten_tensors(input_tensors) - # First stage, reduce across each worker using gather_devices. - devices = [t.device for t in input_tensors] - per_worker_devices, per_worker_values = _split_by_task(devices, input_tensors) - num_workers = len(per_worker_devices) - up_values = [] - if len(gather_devices) != num_workers: - raise ValueError("For shuffle hybrid, gather_devices must contain one " - "device per worker. ") - for w in range(0, num_workers): - reduced_shards = _build_shuffle_gather( - per_worker_values[w], [gather_devices[w]], red_op) - up_values.append(reduced_shards[0]) - # Second stage, apply upper_level_f. - level_2_output = upper_level_f(up_values) - # Third stage, apply shuffle scatter at each worker. - output_tensors = [] - for w in range(0, num_workers): - output_tensors += _build_shuffle_scatter( - [level_2_output[w]], per_worker_devices[w]) - if len(shape) != 1: - output_tensors = _reshape_tensors(output_tensors, shape) - return output_tensors - - -def build_shuffle_then_ring(input_tensors, gather_devices, subdiv, - red_n_op, red_op, un_op=None): - """Construct hybrid of Shuffle within workers, Ring across workers.""" - def upper_builder(tensors): - return build_ring_all_reduce(tensors, len(tensors), subdiv, [0], - red_op, un_op) - def upper_level_f(tensors): - return _reduce_non_singleton(tensors, upper_builder, un_op) - return _build_shuffle_hybrid( - input_tensors, gather_devices, red_n_op, upper_level_f) - - -def build_shuffle_then_shuffle(input_tensors, first_gather_devices, - second_gather_devices, red_op, un_op=None): - """Construct hybrid of Shuffle within workers, Shuffle across workers.""" - def upper_builder(tensors): - return build_shuffle_all_reduce(tensors, second_gather_devices, - red_op, un_op) - def upper_level_f(tensors): - return _reduce_non_singleton(tensors, upper_builder, un_op) - return _build_shuffle_hybrid( - input_tensors, first_gather_devices, red_op, upper_level_f) +# pylint: disable=unused-import,wildcard-import +from tensorflow.python.distribute.all_reduce import * diff --git a/tensorflow/contrib/android/README.md b/tensorflow/contrib/android/README.md index db37bcf73d144eb81c32a461a276d10be7e2d193..27f8ac21323e6eb21a80dfab4d2239738c2fcf1e 100644 --- a/tensorflow/contrib/android/README.md +++ b/tensorflow/contrib/android/README.md @@ -52,6 +52,7 @@ Then, to build the native TF library: bazel build -c opt //tensorflow/contrib/android:libtensorflow_inference.so \ --crosstool_top=//external:android/crosstool \ --host_crosstool_top=@bazel_tools//tools/cpp:toolchain \ + --cxxopt=-std=c++11 \ --cpu=armeabi-v7a ``` diff --git a/tensorflow/contrib/android/cmake/build.gradle b/tensorflow/contrib/android/cmake/build.gradle index 17a57b99fd6c9efc09bda0ce1249b1f51bd5af5c..ddec08894f34f96b080610f1d27a6a436f7ffa91 100644 --- a/tensorflow/contrib/android/cmake/build.gradle +++ b/tensorflow/contrib/android/cmake/build.gradle @@ -22,8 +22,8 @@ android { } externalNativeBuild { cmake { - arguments '-DANDROID_TOOLCHAIN=gcc', - '-DANDROID_STL=gnustl_static' + arguments '-DANDROID_TOOLCHAIN=clang', + '-DANDROID_STL=c++_static' } } } @@ -70,7 +70,7 @@ if (ndkDir == null || ndkDir == "") { ndkDir = System.getenv('ANDROID_NDK_HOME') } -if(! Os.isFamily(Os.FAMILY_WINDOWS)) { +if (!Os.isFamily(Os.FAMILY_WINDOWS)) { // This script is for non-Windows OS. For Windows OS, MANUALLY build // (or copy the built) libs/headers to the // ${TENSORFLOW_ROOT_DIR}/tensorflow/contrib/makefile/gen diff --git a/tensorflow/contrib/autograph/examples/benchmarks/BUILD b/tensorflow/contrib/autograph/examples/benchmarks/BUILD new file mode 100644 index 0000000000000000000000000000000000000000..6d2d70c99b4cc804f2c8bf57afdc8c11f1f73516 --- /dev/null +++ b/tensorflow/contrib/autograph/examples/benchmarks/BUILD @@ -0,0 +1,36 @@ +licenses(["notice"]) # Apache 2.0 + +load("//tensorflow:tensorflow.bzl", "py_test") +load("//tensorflow/tools/test:performance.bzl", "tf_py_logged_benchmark") + +py_library( + name = "benchmark_base", + srcs = [ + "benchmark_base.py", + ], + deps = [ + "//tensorflow:tensorflow_py", + ], +) + +py_test( + name = "cartpole_benchmark", + size = "enormous", + srcs = ["cartpole_benchmark.py"], + tags = [ + "local", + "manual", + "no_oss", + "notap", + "nozapfhahn", + ], + deps = [ + ":benchmark_base", + # Note: required gym dependency may need to be added here. + ], +) + +tf_py_logged_benchmark( + name = "cartpole_logged_benchmark", + target = "//tensorflow/contrib/autograph/examples/benchmarks:cartpole_benchmark", +) diff --git a/tensorflow/contrib/autograph/examples/benchmarks/benchmark_base.py b/tensorflow/contrib/autograph/examples/benchmarks/benchmark_base.py new file mode 100644 index 0000000000000000000000000000000000000000..93c694849c4dc3faca71e7f9d8614649a7784f99 --- /dev/null +++ b/tensorflow/contrib/autograph/examples/benchmarks/benchmark_base.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. +# ============================================================================== +"""Common benchmarking code. + +See https://www.tensorflow.org/community/benchmarks for usage. +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import time + +import numpy as np + +import tensorflow as tf + + +class ReportingBenchmark(tf.test.Benchmark): + """Base class for a benchmark that reports general performance metrics. + + Subclasses only need to call one of the _profile methods, and optionally + report_results. + """ + + def time_execution(self, name, target, iters, warm_up_iters=5): + for _ in range(warm_up_iters): + target() + + all_times = [] + for _ in range(iters): + iter_time = time.time() + target() + all_times.append(time.time() - iter_time) + + avg_time = np.average(all_times) + + extras = dict() + extras['all_times'] = all_times + + if isinstance(name, tuple): + extras['name'] = name + name = '_'.join(str(piece) for piece in name) + + self.report_benchmark( + iters=iters, wall_time=avg_time, name=name, extras=extras) + + +if __name__ == '__main__': + tf.test.main() diff --git a/tensorflow/contrib/autograph/examples/benchmarks/cartpole_benchmark.py b/tensorflow/contrib/autograph/examples/benchmarks/cartpole_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..4f553be58e94f11e45f0697558348fbbd26bfb91 --- /dev/null +++ b/tensorflow/contrib/autograph/examples/benchmarks/cartpole_benchmark.py @@ -0,0 +1,492 @@ +# 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. +# ============================================================================== +"""A basic RL cartpole benchmark. + +The RL model uses the OpenAI Gym environment to train a simple network using +the policy gradients method. The training scales the gradients for each step +by the episode's cumulative discounted reward and averages these gradients over +a fixed number of games before applying the optimization step. + +For benchmarking purposes, we replace the OpenAI Gym environment to a fake +that returns random actions and rewards and never ends the episode. This way +the benchmarks compare the same amount of computation at each step. +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import gym +import numpy as np +import tensorflow as tf + +from tensorflow.contrib import eager +from tensorflow.contrib.autograph.examples.benchmarks import benchmark_base +from tensorflow.python import autograph as ag +from tensorflow.python.eager import context + +# +# AutoGraph implementation +# + + +@ag.convert() +def graph_append_discounted_rewards(destination, rewards, discount_rate): + """Discounts episode rewards and appends them to destination.""" + ag.set_element_type(rewards, tf.float32) + + cdr = 0.0 + reverse_discounted = [] + ag.set_element_type(reverse_discounted, tf.float32) + + for i in range(len(rewards) - 1, -1, -1): + cdr = cdr * discount_rate + rewards[i] + cdr.set_shape(()) + reverse_discounted.append(cdr) + + retval = destination + # Note: AutoGraph doesn't yet support reversed() so we use a loop instead. + for i in range(len(reverse_discounted) - 1, -1, -1): + retval.append(reverse_discounted[i]) + + return retval + + +class GraphPolicyNetwork(tf.keras.Model): + """Policy network for the cart-pole reinforcement learning problem. + + The forward path of the network takes an observation from the cart-pole + environment (length-4 vector) and outputs an action. + """ + + def __init__(self, hidden_size): + super(GraphPolicyNetwork, self).__init__() + self._hidden_layer = tf.keras.layers.Dense( + hidden_size, activation=tf.nn.elu) + self._output_layer = tf.keras.layers.Dense(1) + + def call(self, inputs): + """Calculates logits and action. + + Args: + inputs: Observations from a step in the cart-pole environment, of shape + `(batch_size, input_size)` + + Returns: + logits: the logits output by the output layer. This can be viewed as the + likelihood vales of choosing the left (0) action. Shape: + `(batch_size, 1)`. + actions: randomly selected actions ({0, 1}) based on the logits. Shape: + `(batch_size, 1)`. + """ + hidden = self._hidden_layer(inputs) + logits = self._output_layer(hidden) + + left_prob = tf.nn.sigmoid(logits) + action_probs = tf.concat([left_prob, 1.0 - left_prob], 1) + + actions = tf.multinomial(tf.log(action_probs), 1) + return logits, actions + + # TODO(mdan): Move this method out of the class. + @ag.convert() + def train(self, cart_pole_env, optimizer, discount_rate, num_games, + max_steps_per_game): + var_list = tf.trainable_variables() + grad_list = [ + tf.TensorArray(tf.float32, 0, dynamic_size=True) for _ in var_list + ] + + step_counts = [] + discounted_rewards = [] + ag.set_element_type(discounted_rewards, tf.float32) + ag.set_element_type(step_counts, tf.int32) + + # Note: we use a shared object, cart_pole_env here. Because calls to the + # object's method are made through py_func, TensorFlow cannot detect its + # data dependencies. Hence we must manually synchronize access to it + # and ensure the control dependencies are set in such a way that + # calls to reset(), take_one_step, etc. are made in the correct order. + sync_counter = tf.constant(0) + + for _ in tf.range(num_games): + with tf.control_dependencies([sync_counter]): + obs = cart_pole_env.reset() + with tf.control_dependencies([obs]): + sync_counter += 1 + + game_rewards = [] + ag.set_element_type(game_rewards, tf.float32) + + for step in tf.range(max_steps_per_game): + logits, actions = self(obs) # pylint:disable=not-callable + logits = tf.reshape(logits, ()) + actions = tf.reshape(actions, ()) + + labels = 1.0 - tf.cast(actions, tf.float32) + loss = tf.nn.sigmoid_cross_entropy_with_logits( + labels=labels, logits=logits) + grads = tf.gradients(loss, var_list) + + for i in range(len(grads)): + grad_list[i].append(grads[i]) + + with tf.control_dependencies([sync_counter]): + obs, reward, done = cart_pole_env.step(actions) + with tf.control_dependencies([obs]): + sync_counter += 1 + obs = tf.reshape(obs, (1, 4)) + + game_rewards.append(reward) + if reward < 0.1 or done: + step_counts.append(step + 1) + break + + discounted_rewards = graph_append_discounted_rewards( + discounted_rewards, game_rewards, discount_rate) + + discounted_rewards = ag.stack(discounted_rewards) + discounted_rewards.set_shape((None,)) + mean, variance = tf.nn.moments(discounted_rewards, [0]) + normalized_rewards = (discounted_rewards - mean) / tf.sqrt(variance) + + for i in range(len(grad_list)): + g = ag.stack(grad_list[i]) + + # This block just adjusts the shapes to match for multiplication. + r = normalized_rewards + if r.shape.ndims < g.shape.ndims: + r = tf.expand_dims(r, -1) + if r.shape.ndims < g.shape.ndims: + r = tf.expand_dims(r, -1) + + grad_list[i] = tf.reduce_mean(g * r, axis=0) + + optimizer.apply_gradients( + zip(grad_list, var_list), global_step=tf.train.get_global_step()) + + return ag.stack(step_counts) + + +@ag.convert() +def graph_train_model(policy_network, cart_pole_env, optimizer, iterations): + """Trains the policy network for a given number of iterations.""" + i = tf.constant(0) + mean_steps_per_iteration = [] + ag.set_element_type(mean_steps_per_iteration, tf.int32) + + while i < iterations: + steps_per_game = policy_network.train( + cart_pole_env, + optimizer, + discount_rate=0.95, + num_games=20, + max_steps_per_game=200) + mean_steps_per_iteration.append(tf.reduce_mean(steps_per_game)) + i += 1 + + return ag.stack(mean_steps_per_iteration) + + +class GraphGymCartpoleEnv(object): + """An env backed by OpenAI Gym's CartPole environment. + + Used to confirm a functional model only. + """ + + def __init__(self): + cart_pole_env = gym.make('CartPole-v1') + cart_pole_env.seed(0) + cart_pole_env.reset() + self.env = cart_pole_env + + def reset(self): + obs = ag.utils.wrap_py_func(self.env.reset, tf.float64, ()) + obs = tf.reshape(obs, (1, 4)) + obs = tf.cast(obs, tf.float32) + return obs + + def step(self, actions): + + def take_one_step(actions): + obs, reward, done, _ = self.env.step(actions) + obs = obs.astype(np.float32) + reward = np.float32(reward) + return obs, reward, done + + return ag.utils.wrap_py_func(take_one_step, + (tf.float32, tf.float32, tf.bool), (actions,)) + + +class GraphRandomCartpoleEnv(object): + """An environment that returns random actions and never finishes. + + Used during benchmarking, it will cause training to run a constant number of + steps. + """ + + def reset(self): + return tf.random.normal((1, 4)) + + def step(self, actions): + with tf.control_dependencies([actions]): + random_obs = tf.random.normal((1, 4)) + fixed_reward = tf.constant(0.001) + done = tf.constant(False) + return random_obs, fixed_reward, done + + +# +# Eager implementation +# + + +def eager_append_discounted_rewards(discounted_rewards, rewards, discount_rate): + cdr = 0.0 + reverse_discounted = [] + + for i in range(len(rewards) - 1, -1, -1): + cdr = cdr * discount_rate + rewards[i] + reverse_discounted.append(cdr) + + discounted_rewards.extend(reversed(reverse_discounted)) + return discounted_rewards + + +class EagerPolicyNetwork(tf.keras.Model): + """Policy network for the cart-pole reinforcement learning problem. + + The forward path of the network takes an observation from the cart-pole + environment (length-4 vector) and outputs an action. + """ + + def __init__(self, hidden_size): + super(EagerPolicyNetwork, self).__init__() + self._hidden_layer = tf.keras.layers.Dense( + hidden_size, activation=tf.nn.elu) + self._output_layer = tf.keras.layers.Dense(1) + + def call(self, inputs): + """Calculates logits and action. + + Args: + inputs: Observations from a step in the cart-pole environment, of shape + `(batch_size, input_size)` + + Returns: + logits: the logits output by the output layer. This can be viewed as the + likelihood vales of choosing the left (0) action. Shape: + `(batch_size, 1)`. + actions: randomly selected actions ({0, 1}) based on the logits. Shape: + `(batch_size, 1)`. + """ + hidden = self._hidden_layer(inputs) + logits = self._output_layer(hidden) + + left_prob = tf.nn.sigmoid(logits) + action_probs = tf.concat([left_prob, 1.0 - left_prob], 1) + + self._grad_fn = eager.implicit_gradients( + self._get_cross_entropy_and_save_actions) + + actions = tf.multinomial(tf.log(action_probs), 1) + return logits, actions + + def _get_cross_entropy_and_save_actions(self, inputs): + logits, actions = self(inputs) # pylint:disable=not-callable + self._current_actions = actions + labels = 1.0 - tf.cast(actions, tf.float32) + return tf.nn.sigmoid_cross_entropy_with_logits(labels=labels, logits=logits) + + def train(self, cart_pole_env, optimizer, discount_rate, num_games, + max_steps_per_game): + grad_list = None + + step_counts = [] + discounted_rewards = [] + + for _ in range(num_games): + obs = cart_pole_env.reset() + + game_rewards = [] + + for step in range(max_steps_per_game): + grads_and_vars = self._grad_fn(tf.constant([obs], dtype=tf.float32)) + grads, var_list = zip(*grads_and_vars) + actions = self._current_actions.numpy()[0][0] + + if grad_list is None: + grad_list = [[g] for g in grads] + else: + for i in range(len(grads)): + grad_list[i].append(grads[i]) + + obs, reward, done = cart_pole_env.step(actions) + + game_rewards.append(reward) + if reward < 0.1 or done: + step_counts.append(step + 1) + break + + discounted_rewards = eager_append_discounted_rewards( + discounted_rewards, game_rewards, discount_rate) + + discounted_rewards = tf.stack(discounted_rewards) + mean, variance = tf.nn.moments(discounted_rewards, [0]) + normalized_rewards = (discounted_rewards - mean) / tf.sqrt(variance) + + for i in range(len(grad_list)): + g = tf.stack(grad_list[i]) + + r = normalized_rewards + while r.shape.ndims < g.shape.ndims: + r = tf.expand_dims(r, -1) + + grad_list[i] = tf.reduce_mean(g * r, axis=0) + + optimizer.apply_gradients( + zip(grad_list, var_list), global_step=tf.train.get_global_step()) + + return tf.stack(step_counts) + + +def eager_train_model(policy_network, cart_pole_env, optimizer, iterations): + """Trains the policy network for a given number of iterations.""" + mean_steps_per_iteration = [] + + for _ in range(iterations): + steps_per_game = policy_network.train( + cart_pole_env, + optimizer, + discount_rate=0.95, + num_games=20, + max_steps_per_game=200) + mean_steps_per_iteration.append(tf.reduce_mean(steps_per_game)) + + return mean_steps_per_iteration + + +class EagerGymCartpoleEnv(object): + """An env backed by OpenAI Gym's CartPole environment. + + Used to confirm a functional model only. + """ + + def __init__(self): + cart_pole_env = gym.make('CartPole-v1') + cart_pole_env.seed(0) + cart_pole_env.reset() + self.env = cart_pole_env + + def reset(self): + return self.env.reset() + + def step(self, actions): + obs, reward, done, _ = self.env.step(actions) + return obs, reward, done + + +class EagerRandomCartpoleEnv(object): + """An environment that returns random actions and never finishes. + + Used during benchmarking, it will cause training to run a constant number of + steps. + """ + + def reset(self): + return np.random.normal(size=(4,)) + + def step(self, actions): + with tf.control_dependencies([actions]): + random_obs = np.random.normal(size=(4,)) + fixed_reward = 0.001 + done = False + return random_obs, fixed_reward, done + + +def graph_demo_training(): + """Not used in the benchmark. Used to confirm a functional model.""" + with tf.Graph().as_default(): + tf.set_random_seed(0) + + network = GraphPolicyNetwork(hidden_size=5) + network.build((1, 4)) + env = GraphGymCartpoleEnv() + opt = tf.train.AdamOptimizer(0.05) + + train_ops = graph_train_model(network, env, opt, iterations=5) + + with tf.Session() as sess: + sess.run(tf.global_variables_initializer()) + sess.run(tf.local_variables_initializer()) + steps_per_iteration = sess.run(train_ops) + for i, steps in enumerate(steps_per_iteration): + print('Step {} iterations: {}'.format(i, steps)) + + +def eager_demo_training(): + with context.eager_mode(): + network = EagerPolicyNetwork(hidden_size=5) + network.build((1, 4)) + env = EagerGymCartpoleEnv() + opt = tf.train.AdamOptimizer(0.05) + + steps_per_iteration = eager_train_model(network, env, opt, iterations=5) + for i, steps in enumerate(steps_per_iteration): + print('Step {} iterations: {}'.format(i, steps)) + + +class RLCartPoleBenchmark(benchmark_base.ReportingBenchmark): + """Actual benchmark. + + Trains the RL agent a fixed number of times, on random environments that + result in constant number of steps. + """ + + def benchmark_cartpole(self): + + def train_session(sess, ops): + return lambda: sess.run(ops) + + def train_eager(network, env, opt): + return lambda: eager_train_model(network, env, opt, iterations=10) + + for model_size in (10, 100, 1000): + with tf.Graph().as_default(): + network = GraphPolicyNetwork(hidden_size=model_size) + network.build((1, 4)) + env = GraphRandomCartpoleEnv() + opt = tf.train.AdamOptimizer(0.05) + train_ops = graph_train_model(network, env, opt, iterations=10) + + with tf.Session() as sess: + sess.run(tf.global_variables_initializer()) + sess.run(tf.local_variables_initializer()) + + self.time_execution(('cartpole', 'autograph', model_size), + train_session(sess, train_ops), 20) + + with context.eager_mode(): + network = EagerPolicyNetwork(hidden_size=model_size) + network.build((1, 4)) + env = EagerRandomCartpoleEnv() + opt = tf.train.AdamOptimizer(0.05) + + self.time_execution(('cartpole', 'eager', model_size), + train_eager(network, env, opt), 20) + + +if __name__ == '__main__': + tf.test.main() 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/batching/python/ops/batch_ops.py b/tensorflow/contrib/batching/python/ops/batch_ops.py index 55faad983f2bcf2f3fa633669bd371608e2e925b..3e4d0dc1cec76b068c1c846eb476eec615e4f613 100644 --- a/tensorflow/contrib/batching/python/ops/batch_ops.py +++ b/tensorflow/contrib/batching/python/ops/batch_ops.py @@ -18,8 +18,9 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -from tensorflow.python.framework import function +from tensorflow.python.eager import function from tensorflow.python.framework import ops +from tensorflow.python.framework import tensor_spec from tensorflow.python.ops import gen_batch_ops # go/tf-wildcard-import # pylint: disable=wildcard-import @@ -101,12 +102,15 @@ def batch_function(num_batch_threads, def decorator(fn): # pylint: disable=missing-docstring def decorated(*args): # pylint: disable=missing-docstring - types = [arg.dtype for arg in args] - @function.Defun(*types) + @function.defun() def computation(*computation_args): return fn(*computation_args) + computation = computation.get_concrete_function( + *[tensor_spec.TensorSpec(dtype=x.dtype, shape=x.shape, name=str(i)) + for i, x in enumerate(args)]) + with ops.name_scope("batch") as name: for a in args: if not isinstance(a, ops.Tensor): @@ -123,7 +127,7 @@ def batch_function(num_batch_threads, f=computation, in_tensors=list(args), captured_tensors=computation.captured_inputs, - Tout=[o.type for o in computation.definition.signature.output_arg]) + Tout=[o.dtype for o in computation.outputs]) return decorated diff --git a/tensorflow/contrib/batching/python/ops/batch_ops_test.py b/tensorflow/contrib/batching/python/ops/batch_ops_test.py index 01ee8703a93836d607ee9b765c51c79fe3bb974f..9109b9c1c91cefa4c52bad49de23336a6e05e1ef 100644 --- a/tensorflow/contrib/batching/python/ops/batch_ops_test.py +++ b/tensorflow/contrib/batching/python/ops/batch_ops_test.py @@ -219,6 +219,7 @@ class BatchOpsTest(test.TestCase): @batch_ops.batch_function(1, 10, 100000) def computation(in_t): + self.assertTrue(in_t.shape is not None) return in_t + 1 inp = array_ops.placeholder(dtype=dtypes.int32, shape=[1]) diff --git a/tensorflow/contrib/bayesflow/python/kernel_tests/monte_carlo_test.py b/tensorflow/contrib/bayesflow/python/kernel_tests/monte_carlo_test.py index 13215ffabf3a956d3f83697f867457b2fa72e7c9..8b6ed9f041b89a0da02a505ec261bca82b094f74 100644 --- a/tensorflow/contrib/bayesflow/python/kernel_tests/monte_carlo_test.py +++ b/tensorflow/contrib/bayesflow/python/kernel_tests/monte_carlo_test.py @@ -81,7 +81,7 @@ class ExpectationImportanceSampleTest(test.TestCase): # Compute E_p[X_1 * X_2 > 0], with X_i the ith component of X ~ p(x). # Should equal 1/2 because p is a spherical Gaussian centered at (0, 0). def indicator(x): - x1_times_x2 = math_ops.reduce_prod(x, reduction_indices=[-1]) + x1_times_x2 = math_ops.reduce_prod(x, axis=[-1]) return 0.5 * (math_ops.sign(x1_times_x2) + 1.0) prob = mc.expectation_importance_sampler( diff --git a/tensorflow/contrib/bayesflow/python/ops/monte_carlo_impl.py b/tensorflow/contrib/bayesflow/python/ops/monte_carlo_impl.py index 18d40fc1dff8e7c9aefffbe3ceba770598a42096..e83a54851195708eb7e6412b7400236f4bc06e6b 100644 --- a/tensorflow/contrib/bayesflow/python/ops/monte_carlo_impl.py +++ b/tensorflow/contrib/bayesflow/python/ops/monte_carlo_impl.py @@ -353,12 +353,12 @@ def expectation(f, samples, log_prob=None, use_reparametrization=True, def _sample_mean(values): """Mean over sample indices. In this module this is always [0].""" - return math_ops.reduce_mean(values, reduction_indices=[0]) + return math_ops.reduce_mean(values, axis=[0]) def _sample_max(values): """Max over sample indices. In this module this is always [0].""" - return math_ops.reduce_max(values, reduction_indices=[0]) + return math_ops.reduce_max(values, axis=[0]) def _get_samples(dist, z, n, seed): diff --git a/tensorflow/contrib/bigtable/README.md b/tensorflow/contrib/bigtable/README.md index 2c44abed5e1955cc666273e97e6b2378766f13d2..79052bee35c7895cb4048b10c1f73acb036d1587 100644 --- a/tensorflow/contrib/bigtable/README.md +++ b/tensorflow/contrib/bigtable/README.md @@ -51,25 +51,18 @@ BIGTABLE_TABLE_NAME = '' PREFIX = 'train-' def main(): + tf.enable_eager_execution() + client = tf.contrib.cloud.BigtableClient(GCP_PROJECT_ID, BIGTABLE_INSTANCE_ID) table = client.table(BIGTABLE_TABLE_NAME) dataset = table.keys_by_prefix_dataset(PREFIX) - iterator = dataset.make_initializable_iterator() - get_next_op = iterator.get_next() - with tf.Session() as sess: - print('Initializing the iterator.') - sess.run(iterator.initializer) - print('Retrieving rows:') - row_index = 0 - while True: - try: - row_key = sess.run(get_next_op) - print('Row key %d: %s' % (row_index, row_key)) - row_index += 1 - except tf.errors.OutOfRangeError: - print('Finished reading data!') - break + print('Retrieving rows:') + row_index = 0 + for row_key in dataset: + print('Row key %d: %s' % (row_index, row_key)) + row_index += 1 + print('Finished reading data!') if __name__ == '__main__': main() diff --git a/tensorflow/contrib/bigtable/kernels/bigtable_lib.cc b/tensorflow/contrib/bigtable/kernels/bigtable_lib.cc index 67bf14c17646cff81af707405b66c9fba2ded0bd..98f906408c230a4382ffafe412ee9990d4384930 100644 --- a/tensorflow/contrib/bigtable/kernels/bigtable_lib.cc +++ b/tensorflow/contrib/bigtable/kernels/bigtable_lib.cc @@ -29,8 +29,7 @@ Status GrpcStatusToTfStatus(const ::grpc::Status& status) { } return Status(static_cast<::tensorflow::error::Code>(status.error_code()), strings::StrCat("Error reading from Cloud Bigtable: ", - status.error_message(), - " (Details: ", status.error_details(), ")")); + status.error_message())); } string RegexFromStringSet(const std::vector& strs) { 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/bigtable_sample_key_pairs_dataset_op.cc b/tensorflow/contrib/bigtable/kernels/bigtable_sample_key_pairs_dataset_op.cc index 01608dc6bc07890c3a59577ef31c90c2694e6a87..f0c3ef4e2ecbf5f4d21e421be4fb527d4769200c 100644 --- a/tensorflow/contrib/bigtable/kernels/bigtable_sample_key_pairs_dataset_op.cc +++ b/tensorflow/contrib/bigtable/kernels/bigtable_sample_key_pairs_dataset_op.cc @@ -167,7 +167,7 @@ class BigtableSampleKeyPairsDatasetOp : public DatasetOpKernel { std::vector* out_tensors, bool* end_of_sequence) override { mutex_lock l(mu_); - if (index_ > keys_.size() - 2) { + if (index_ + 2 > keys_.size()) { *end_of_sequence = true; return Status::OK(); } 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 f083ce6f44b3c2a83d9b5d3235056eb94c4be4a8..e6fda9e61757f1441b3691c2a3d57c6f1a5a0d42 100644 --- a/tensorflow/contrib/bigtable/kernels/test_kernels/bigtable_test_client.cc +++ b/tensorflow/contrib/bigtable/kernels/test_kernels/bigtable_test_client.cc @@ -15,7 +15,7 @@ limitations under the License. #include "tensorflow/contrib/bigtable/kernels/test_kernels/bigtable_test_client.h" -#include "google/bigtable/v2/data.pb.h" +#include "external/com_github_googleapis_googleapis/google/bigtable/v2/data.pb.h" #include "google/protobuf/wrappers.pb.h" #include "re2/re2.h" #include "tensorflow/core/lib/strings/stringprintf.h" @@ -366,6 +366,61 @@ BigtableTestClient::MutateRows( return MakeUnique(request.entries_size()); } +std::unique_ptr> +BigtableTestClient::AsyncMutateRow( + grpc::ClientContext* context, + google::bigtable::v2::MutateRowRequest const& request, + grpc::CompletionQueue* cq) { + LOG(WARNING) << "Call to InMemoryDataClient::" << __func__ + << "(); this will likely cause a crash!"; + return nullptr; +} + +std::unique_ptr<::grpc::ClientAsyncReaderInterface< + ::google::bigtable::v2::SampleRowKeysResponse>> +BigtableTestClient::AsyncSampleRowKeys( + ::grpc::ClientContext* context, + const ::google::bigtable::v2::SampleRowKeysRequest& request, + ::grpc::CompletionQueue* cq, void* tag) { + LOG(WARNING) << "Call to InMemoryDataClient::" << __func__ + << "(); this will likely cause a crash!"; + return nullptr; +} + +std::unique_ptr<::grpc::ClientAsyncReaderInterface< + ::google::bigtable::v2::MutateRowsResponse>> +BigtableTestClient::AsyncMutateRows( + ::grpc::ClientContext* context, + const ::google::bigtable::v2::MutateRowsRequest& request, + ::grpc::CompletionQueue* cq, void* tag) { + LOG(WARNING) << "Call to InMemoryDataClient::" << __func__ + << "(); this will likely cause a crash!"; + return nullptr; +} + +std::unique_ptr> +BigtableTestClient::AsyncCheckAndMutateRow( + grpc::ClientContext* context, + const google::bigtable::v2::CheckAndMutateRowRequest& request, + grpc::CompletionQueue* cq) { + LOG(WARNING) << "Call to InMemoryDataClient::" << __func__ + << "(); this will likely cause a crash!"; + 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 dac2b16a216d26f02684c7401ed2ddaa4b7baddb..8e1326f2ce841368ea81fc7194a0588e5d6cd637 100644 --- a/tensorflow/contrib/bigtable/kernels/test_kernels/bigtable_test_client.h +++ b/tensorflow/contrib/bigtable/kernels/test_kernels/bigtable_test_client.h @@ -61,6 +61,38 @@ class BigtableTestClient : public ::google::cloud::bigtable::DataClient { MutateRows(grpc::ClientContext* context, google::bigtable::v2::MutateRowsRequest const& request) override; + std::unique_ptr> + AsyncMutateRow(grpc::ClientContext* context, + google::bigtable::v2::MutateRowRequest const& request, + grpc::CompletionQueue* cq) override; + + std::unique_ptr<::grpc::ClientAsyncReaderInterface< + ::google::bigtable::v2::SampleRowKeysResponse>> + AsyncSampleRowKeys( + ::grpc::ClientContext* context, + const ::google::bigtable::v2::SampleRowKeysRequest& request, + ::grpc::CompletionQueue* cq, void* tag) override; + + std::unique_ptr<::grpc::ClientAsyncReaderInterface< + ::google::bigtable::v2::MutateRowsResponse>> + AsyncMutateRows(::grpc::ClientContext* context, + const ::google::bigtable::v2::MutateRowsRequest& request, + ::grpc::CompletionQueue* cq, void* tag) override; + + std::unique_ptr> + AsyncCheckAndMutateRow( + grpc::ClientContext* context, + 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/python/kernel_tests/bigtable_ops_test.py b/tensorflow/contrib/bigtable/python/kernel_tests/bigtable_ops_test.py index 316da9ebe152ef52c7e7f846cf8c3eb1555ee8a6..197f5578eb010bee5a3aad7c05446393193f99e2 100644 --- a/tensorflow/contrib/bigtable/python/kernel_tests/bigtable_ops_test.py +++ b/tensorflow/contrib/bigtable/python/kernel_tests/bigtable_ops_test.py @@ -57,7 +57,7 @@ class BigtableOpsTest(test.TestCase): sess.run(write_op) def runReadKeyTest(self, read_ds): - itr = read_ds.make_initializable_iterator() + itr = dataset_ops.make_initializable_iterator(read_ds) n = itr.get_next() expected = list(self.COMMON_ROW_KEYS) expected.reverse() @@ -78,7 +78,7 @@ class BigtableOpsTest(test.TestCase): self.runReadKeyTest(self._table.keys_by_range_dataset("r1", "r4")) def runScanTest(self, read_ds): - itr = read_ds.make_initializable_iterator() + itr = dataset_ops.make_initializable_iterator(read_ds) n = itr.get_next() expected_keys = list(self.COMMON_ROW_KEYS) expected_keys.reverse() @@ -120,7 +120,7 @@ class BigtableOpsTest(test.TestCase): def testLookup(self): ds = self._table.keys_by_prefix_dataset("r") ds = ds.apply(self._table.lookup_columns(cf1="c1")) - itr = ds.make_initializable_iterator() + itr = dataset_ops.make_initializable_iterator(ds) n = itr.get_next() expected_keys = list(self.COMMON_ROW_KEYS) expected_values = list(self.COMMON_VALUES) @@ -141,7 +141,7 @@ class BigtableOpsTest(test.TestCase): def testSampleKeys(self): ds = self._table.sample_keys() - itr = ds.make_initializable_iterator() + itr = dataset_ops.make_initializable_iterator(ds) n = itr.get_next() expected_key = self.COMMON_ROW_KEYS[0] with self.cached_session() as sess: @@ -161,7 +161,7 @@ class BigtableOpsTest(test.TestCase): sess.run(n) def runSampleKeyPairsTest(self, ds, expected_key_pairs): - itr = ds.make_initializable_iterator() + itr = dataset_ops.make_initializable_iterator(ds) n = itr.get_next() with self.cached_session() as sess: self._writeCommonValues(sess) @@ -218,7 +218,7 @@ class BigtableOpsTest(test.TestCase): def testSampleKeyPairsPrefixAndStartKey(self): ds = bigtable_api._BigtableSampleKeyPairsDataset( self._table, prefix="r", start="r1", end="") - itr = ds.make_initializable_iterator() + itr = dataset_ops.make_initializable_iterator(ds) with self.cached_session() as sess: with self.assertRaises(errors.InvalidArgumentError): sess.run(itr.initializer) @@ -226,14 +226,14 @@ class BigtableOpsTest(test.TestCase): def testSampleKeyPairsPrefixAndEndKey(self): ds = bigtable_api._BigtableSampleKeyPairsDataset( self._table, prefix="r", start="", end="r3") - itr = ds.make_initializable_iterator() + itr = dataset_ops.make_initializable_iterator(ds) with self.cached_session() as sess: with self.assertRaises(errors.InvalidArgumentError): sess.run(itr.initializer) def testParallelScanPrefix(self): ds = self._table.parallel_scan_prefix(prefix="r", cf1="c1") - itr = ds.make_initializable_iterator() + itr = dataset_ops.make_initializable_iterator(ds) n = itr.get_next() with self.cached_session() as sess: self._writeCommonValues(sess) @@ -251,7 +251,7 @@ class BigtableOpsTest(test.TestCase): def testParallelScanRange(self): ds = self._table.parallel_scan_range(start="r1", end="r4", cf1="c1") - itr = ds.make_initializable_iterator() + itr = dataset_ops.make_initializable_iterator(ds) n = itr.get_next() with self.cached_session() as sess: self._writeCommonValues(sess) diff --git a/tensorflow/contrib/bigtable/python/ops/bigtable_api.py b/tensorflow/contrib/bigtable/python/ops/bigtable_api.py index 7c87b0daeb09950cc44c51f49c16534d413f0376..fa64055dfd65a134afdf46cebccb7f7d96106502 100644 --- a/tensorflow/contrib/bigtable/python/ops/bigtable_api.py +++ b/tensorflow/contrib/bigtable/python/ops/bigtable_api.py @@ -35,8 +35,8 @@ from tensorflow.contrib.util import loader from tensorflow.python.data.experimental.ops import interleave_ops from tensorflow.python.data.ops import dataset_ops from tensorflow.python.data.util import nest +from tensorflow.python.data.util import structure from tensorflow.python.framework import dtypes -from tensorflow.python.framework import ops from tensorflow.python.framework import tensor_shape from tensorflow.python.platform import resource_loader @@ -111,8 +111,7 @@ class BigtableClient(object): class BigtableTable(object): - """BigtableTable is the entrypoint for reading and writing data in Cloud - Bigtable. + """Entry point for reading and writing data in Cloud Bigtable. This BigtableTable class is the Python representation of the Cloud Bigtable table within TensorFlow. Methods on this class allow data to be read from and @@ -222,7 +221,7 @@ class BigtableTable(object): A `tf.data.Dataset`. containing `tf.string` Tensors corresponding to all of the row keys matching that prefix. """ - return _BigtablePrefixKeyDataset(self, prefix) + return dataset_ops.DatasetV1Adapter(_BigtablePrefixKeyDataset(self, prefix)) def sample_keys(self): """Retrieves a sampling of row keys from the Bigtable table. @@ -234,7 +233,7 @@ class BigtableTable(object): Returns: A `tf.data.Dataset` returning string row keys. """ - return _BigtableSampleKeysDataset(self) + return dataset_ops.DatasetV1Adapter(_BigtableSampleKeysDataset(self)) def scan_prefix(self, prefix, probability=None, columns=None, **kwargs): """Retrieves row (including values) from the Bigtable service. @@ -279,7 +278,8 @@ class BigtableTable(object): """ probability = _normalize_probability(probability) normalized = _normalize_columns(columns, kwargs) - return _BigtableScanDataset(self, prefix, "", "", normalized, probability) + return dataset_ops.DatasetV1Adapter( + _BigtableScanDataset(self, prefix, "", "", normalized, probability)) def scan_range(self, start, end, probability=None, columns=None, **kwargs): """Retrieves rows (including values) from the Bigtable service. @@ -324,7 +324,8 @@ class BigtableTable(object): """ probability = _normalize_probability(probability) normalized = _normalize_columns(columns, kwargs) - return _BigtableScanDataset(self, "", start, end, normalized, probability) + return dataset_ops.DatasetV1Adapter( + _BigtableScanDataset(self, "", start, end, normalized, probability)) def parallel_scan_prefix(self, prefix, @@ -380,7 +381,8 @@ class BigtableTable(object): """ probability = _normalize_probability(probability) normalized = _normalize_columns(columns, kwargs) - ds = _BigtableSampleKeyPairsDataset(self, prefix, "", "") + ds = dataset_ops.DatasetV1Adapter( + _BigtableSampleKeyPairsDataset(self, prefix, "", "")) return self._make_parallel_scan_dataset(ds, num_parallel_scans, probability, normalized) @@ -442,7 +444,8 @@ class BigtableTable(object): """ probability = _normalize_probability(probability) normalized = _normalize_columns(columns, kwargs) - ds = _BigtableSampleKeyPairsDataset(self, "", start, end) + ds = dataset_ops.DatasetV1Adapter( + _BigtableSampleKeyPairsDataset(self, "", start, end)) return self._make_parallel_scan_dataset(ds, num_parallel_scans, probability, normalized) @@ -486,7 +489,7 @@ class BigtableTable(object): "len(dataset.output_types))") return gen_bigtable_ops.dataset_to_bigtable( self._resource, - dataset._as_variant_tensor(), # pylint: disable=protected-access + dataset._variant_tensor, # pylint: disable=protected-access column_families, columns, timestamp) @@ -579,26 +582,19 @@ class _BigtableKeyDataset(dataset_ops.DatasetSource): """_BigtableKeyDataset is an abstract class representing the keys of a table. """ - def __init__(self, table): + def __init__(self, table, variant_tensor): """Constructs a _BigtableKeyDataset. Args: table: a Bigtable class. + variant_tensor: DT_VARIANT representation of the dataset. """ - super(_BigtableKeyDataset, self).__init__() + super(_BigtableKeyDataset, self).__init__(variant_tensor) self._table = table @property - def output_classes(self): - return ops.Tensor - - @property - def output_shapes(self): - return tensor_shape.TensorShape([]) - - @property - def output_types(self): - return dtypes.string + def _element_structure(self): + return structure.TensorStructure(dtypes.string, []) class _BigtablePrefixKeyDataset(_BigtableKeyDataset): @@ -606,13 +602,11 @@ class _BigtablePrefixKeyDataset(_BigtableKeyDataset): """ def __init__(self, table, prefix): - super(_BigtablePrefixKeyDataset, self).__init__(table) self._prefix = prefix - - def _as_variant_tensor(self): - return gen_bigtable_ops.bigtable_prefix_key_dataset( - table=self._table._resource, # pylint: disable=protected-access + variant_tensor = gen_bigtable_ops.bigtable_prefix_key_dataset( + table=table._resource, # pylint: disable=protected-access prefix=self._prefix) + super(_BigtablePrefixKeyDataset, self).__init__(table, variant_tensor) class _BigtableRangeKeyDataset(_BigtableKeyDataset): @@ -620,15 +614,13 @@ class _BigtableRangeKeyDataset(_BigtableKeyDataset): """ def __init__(self, table, start, end): - super(_BigtableRangeKeyDataset, self).__init__(table) self._start = start self._end = end - - def _as_variant_tensor(self): - return gen_bigtable_ops.bigtable_range_key_dataset( - table=self._table._resource, # pylint: disable=protected-access + variant_tensor = gen_bigtable_ops.bigtable_range_key_dataset( + table=table._resource, # pylint: disable=protected-access start_key=self._start, end_key=self._end) + super(_BigtableRangeKeyDataset, self).__init__(table, variant_tensor) class _BigtableSampleKeysDataset(_BigtableKeyDataset): @@ -638,11 +630,9 @@ class _BigtableSampleKeysDataset(_BigtableKeyDataset): # TODO(saeta): Expose the data size offsets into the keys. def __init__(self, table): - super(_BigtableSampleKeysDataset, self).__init__(table) - - def _as_variant_tensor(self): - return gen_bigtable_ops.bigtable_sample_keys_dataset( - table=self._table._resource) # pylint: disable=protected-access + variant_tensor = gen_bigtable_ops.bigtable_sample_keys_dataset( + table=table._resource) # pylint: disable=protected-access + super(_BigtableSampleKeysDataset, self).__init__(table, variant_tensor) class _BigtableLookupDataset(dataset_ops.DatasetSource): @@ -656,26 +646,17 @@ class _BigtableLookupDataset(dataset_ops.DatasetSource): self._normalized = normalized self._column_families = [i[0] for i in normalized] self._columns = [i[1] for i in normalized] - - @property - def output_classes(self): - return tuple([ops.Tensor] * self._num_outputs) - - @property - def output_shapes(self): - return tuple([tensor_shape.TensorShape([])] * self._num_outputs) - - @property - def output_types(self): - return tuple([dtypes.string] * self._num_outputs) - - def _as_variant_tensor(self): - # pylint: disable=protected-access - return gen_bigtable_ops.bigtable_lookup_dataset( - keys_dataset=self._dataset._as_variant_tensor(), - table=self._table._resource, + variant_tensor = gen_bigtable_ops.bigtable_lookup_dataset( + keys_dataset=self._dataset._variant_tensor, # pylint: disable=protected-access + table=self._table._resource, # pylint: disable=protected-access column_families=self._column_families, columns=self._columns) + super(_BigtableLookupDataset, self).__init__(variant_tensor) + + @property + def _element_structure(self): + return structure.NestedStructure(tuple( + [structure.TensorStructure(dtypes.string, [])] * self._num_outputs)) class _BigtableScanDataset(dataset_ops.DatasetSource): @@ -691,21 +672,7 @@ class _BigtableScanDataset(dataset_ops.DatasetSource): self._columns = [i[1] for i in normalized] self._probability = probability self._num_outputs = len(normalized) + 1 # 1 for row key - - @property - def output_classes(self): - return tuple([ops.Tensor] * self._num_outputs) - - @property - def output_shapes(self): - return tuple([tensor_shape.TensorShape([])] * self._num_outputs) - - @property - def output_types(self): - return tuple([dtypes.string] * self._num_outputs) - - def _as_variant_tensor(self): - return gen_bigtable_ops.bigtable_scan_dataset( + variant_tensor = gen_bigtable_ops.bigtable_scan_dataset( table=self._table._resource, # pylint: disable=protected-access prefix=self._prefix, start_key=self._start, @@ -713,6 +680,13 @@ class _BigtableScanDataset(dataset_ops.DatasetSource): column_families=self._column_families, columns=self._columns, probability=self._probability) + super(_BigtableScanDataset, self).__init__(variant_tensor) + + @property + def _element_structure(self): + return structure.NestedStructure( + tuple( + [structure.TensorStructure(dtypes.string, [])] * self._num_outputs)) class _BigtableSampleKeyPairsDataset(dataset_ops.DatasetSource): @@ -724,23 +698,15 @@ class _BigtableSampleKeyPairsDataset(dataset_ops.DatasetSource): self._prefix = prefix self._start = start self._end = end - - @property - def output_classes(self): - return (ops.Tensor, ops.Tensor) - - @property - def output_shapes(self): - return (tensor_shape.TensorShape([]), tensor_shape.TensorShape([])) - - @property - def output_types(self): - return (dtypes.string, dtypes.string) - - def _as_variant_tensor(self): - # pylint: disable=protected-access - return gen_bigtable_ops.bigtable_sample_key_pairs_dataset( - table=self._table._resource, + variant_tensor = gen_bigtable_ops.bigtable_sample_key_pairs_dataset( + table=self._table._resource, # pylint: disable=protected-access prefix=self._prefix, start_key=self._start, end_key=self._end) + super(_BigtableSampleKeyPairsDataset, self).__init__(variant_tensor) + + @property + def _element_structure(self): + return structure.NestedStructure( + (structure.TensorStructure(dtypes.string, []), + structure.TensorStructure(dtypes.string, []))) diff --git a/tensorflow/contrib/boosted_trees/estimator_batch/BUILD b/tensorflow/contrib/boosted_trees/estimator_batch/BUILD index 14b6fc4ac26f74f54628ae37ad6437c7d3e8caba..d3b23d949ee2c7674c3918d39e8b71d76eefcfec 100644 --- a/tensorflow/contrib/boosted_trees/estimator_batch/BUILD +++ b/tensorflow/contrib/boosted_trees/estimator_batch/BUILD @@ -132,6 +132,7 @@ py_library( srcs = ["estimator.py"], srcs_version = "PY2AND3", deps = [ + ":custom_loss_head", ":estimator_utils", ":model", "//tensorflow/contrib/boosted_trees:losses", diff --git a/tensorflow/contrib/boosted_trees/estimator_batch/custom_export_strategy.py b/tensorflow/contrib/boosted_trees/estimator_batch/custom_export_strategy.py index a3df272e6924792128fc38fd153b9527b58b486e..b314b4d74df882a421d9a2ecce2629a63d5c5248 100644 --- a/tensorflow/contrib/boosted_trees/estimator_batch/custom_export_strategy.py +++ b/tensorflow/contrib/boosted_trees/estimator_batch/custom_export_strategy.py @@ -41,7 +41,8 @@ def make_custom_export_strategy(name, convert_fn, feature_columns, export_input_fn, - use_core_columns=False): + use_core_columns=False, + feature_engineering_fn=None): """Makes custom exporter of GTFlow tree format. Args: @@ -52,6 +53,7 @@ def make_custom_export_strategy(name, export_input_fn: A function that takes no arguments and returns an `InputFnOps`. use_core_columns: A boolean, whether core feature columns were used. + feature_engineering_fn: Feature eng function to be called on the input. Returns: An `ExportStrategy`. @@ -59,9 +61,12 @@ def make_custom_export_strategy(name, base_strategy = saved_model_export_utils.make_export_strategy( serving_input_fn=export_input_fn, strip_default_attrs=True) input_fn = export_input_fn() + features = input_fn.features + if feature_engineering_fn is not None: + features, _ = feature_engineering_fn(features, labels=None) (sorted_feature_names, dense_floats, sparse_float_indices, _, _, sparse_int_indices, _, _) = gbdt_batch.extract_features( - input_fn.features, feature_columns, use_core_columns) + features, feature_columns, use_core_columns) def export_fn(estimator, export_dir, checkpoint_path=None, eval_result=None): """A wrapper to export to SavedModel, and convert it to other formats.""" diff --git a/tensorflow/contrib/boosted_trees/estimator_batch/dnn_tree_combined_estimator.py b/tensorflow/contrib/boosted_trees/estimator_batch/dnn_tree_combined_estimator.py index ca73e4af2fbd0a383d02fa7111f59161701661df..358404cd946bbc56d2f7228be8fe4223749c850b 100644 --- a/tensorflow/contrib/boosted_trees/estimator_batch/dnn_tree_combined_estimator.py +++ b/tensorflow/contrib/boosted_trees/estimator_batch/dnn_tree_combined_estimator.py @@ -36,7 +36,7 @@ from tensorflow.contrib.learn.python.learn.estimators import estimator from tensorflow.contrib.learn.python.learn.estimators import head as head_lib from tensorflow.python.estimator import estimator as core_estimator from tensorflow.contrib.learn.python.learn.estimators import model_fn -from tensorflow.python.feature_column import feature_column as feature_column_lib +from tensorflow.python.feature_column import feature_column_lib from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import control_flow_ops diff --git a/tensorflow/contrib/boosted_trees/estimator_batch/estimator.py b/tensorflow/contrib/boosted_trees/estimator_batch/estimator.py index 4c7a538b385ec19f520bff79bab20a121221c60f..a178820841c4c8bcb7f5742babdb6d0f4825de31 100644 --- a/tensorflow/contrib/boosted_trees/estimator_batch/estimator.py +++ b/tensorflow/contrib/boosted_trees/estimator_batch/estimator.py @@ -18,6 +18,8 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import functools + from tensorflow.contrib.boosted_trees.estimator_batch import model from tensorflow.contrib.boosted_trees.python.utils import losses from tensorflow.contrib.learn.python.learn.estimators import estimator @@ -26,7 +28,8 @@ from tensorflow.python.estimator.canned import head as core_head_lib from tensorflow.python.estimator import estimator as core_estimator from tensorflow.python.ops import math_ops from tensorflow.python.ops.losses import losses as core_losses - +from tensorflow.contrib.boosted_trees.estimator_batch import custom_loss_head +from tensorflow.python.ops import array_ops # ================== Old estimator interface=================================== # The estimators below were designed for old feature columns and old estimator @@ -414,30 +417,167 @@ class GradientBoostedDecisionTreeRanker(estimator.Estimator): config=config, feature_engineering_fn=feature_engineering_fn) +# When using this estimator, make sure to regularize the hessian (at least l2, +# min_node_weight)! +# TODO(nponomareva): extend to take multiple quantiles in one go. +class GradientBoostedDecisionTreeQuantileRegressor(estimator.Estimator): + """An estimator that does quantile regression and returns quantile estimates. + """ + + def __init__(self, + learner_config, + examples_per_layer, + quantiles, + label_dimension=1, + num_trees=None, + feature_columns=None, + weight_column_name=None, + model_dir=None, + config=None, + feature_engineering_fn=None, + logits_modifier_function=None, + center_bias=True, + use_core_libs=False, + output_leaf_index=False, + override_global_step_value=None, + num_quantiles=100): + """Initializes a GradientBoostedDecisionTreeQuantileRegressor instance. + + Args: + learner_config: A config for the learner. + examples_per_layer: Number of examples to accumulate before growing a + layer. It can also be a function that computes the number of examples + based on the depth of the layer that's being built. + quantiles: a list of quantiles for the loss, each between 0 and 1. + label_dimension: Dimension of regression label. This is the size + of the last dimension of the labels `Tensor` (typically, this has shape + `[batch_size, label_dimension]`). When label_dimension>1, it is + recommended to use multiclass strategy diagonal hessian or full hessian. + num_trees: An int, number of trees to build. + feature_columns: A list of feature columns. + weight_column_name: Name of the column for weights, or None if not + weighted. + model_dir: Directory for model exports, etc. + config: `RunConfig` object to configure the runtime settings. + feature_engineering_fn: Feature engineering function. Takes features and + labels which are the output of `input_fn` and returns features and + labels which will be fed into the model. + logits_modifier_function: A modifier function for the logits. + center_bias: Whether a separate tree should be created for first fitting + the bias. + use_core_libs: Whether feature columns and loss are from the core (as + opposed to contrib) version of tensorflow. + output_leaf_index: whether to output leaf indices along with predictions + during inference. The leaf node indexes are available in predictions + dict by the key 'leaf_index'. For example, + result_dict = classifier.predict(...) + for example_prediction_result in result_dict: + # access leaf index list by example_prediction_result["leaf_index"] + # which contains one leaf index per tree + override_global_step_value: If after the training is done, global step + value must be reset to this value. This should be used to reset global + step to a number > number of steps used to train the current ensemble. + For example, the usual way is to train a number of trees and set a very + large number of training steps. When the training is done (number of + trees were trained), this parameter can be used to set the global step + to a large value, making it look like that number of training steps ran. + If None, no override of global step will happen. + num_quantiles: Number of quantiles to build for numeric feature values. + """ + + if len(quantiles) > 1: + raise ValueError('For now, just one quantile per estimator is supported') + + def _quantile_regression_head(quantile): + # Use quantile regression. + head = custom_loss_head.CustomLossHead( + loss_fn=functools.partial( + losses.per_example_quantile_regression_loss, quantile=quantile), + link_fn=array_ops.identity, + logit_dimension=label_dimension) + return head + + learner_config.num_classes = max(2, label_dimension) + + super(GradientBoostedDecisionTreeQuantileRegressor, self).__init__( + model_fn=model.model_builder, + params={ + 'head': _quantile_regression_head(quantiles[0]), + 'feature_columns': feature_columns, + 'learner_config': learner_config, + 'num_trees': num_trees, + 'weight_column_name': weight_column_name, + 'examples_per_layer': examples_per_layer, + 'logits_modifier_function': logits_modifier_function, + 'center_bias': center_bias, + 'use_core_libs': use_core_libs, + 'output_leaf_index': False, + 'override_global_step_value': override_global_step_value, + 'num_quantiles': num_quantiles, + }, + model_dir=model_dir, + config=config, + feature_engineering_fn=feature_engineering_fn) + # ================== New Estimator interface=================================== # The estimators below use new core Estimator interface and must be used with # new feature columns and heads. + # For multiclass classification, use the following head since it uses loss # that is twice differentiable. -def core_multiclass_head(n_classes): +def core_multiclass_head( + n_classes, + weight_column=None, + loss_reduction=core_losses.Reduction.SUM_OVER_NONZERO_WEIGHTS): """Core head for multiclass problems.""" def loss_fn(labels, logits): result = losses.per_example_maxent_loss( - labels=labels, logits=logits, weights=None, num_classes=n_classes) + labels=labels, + logits=logits, + weights=weight_column, + num_classes=n_classes) return result[0] # pylint:disable=protected-access head_fn = core_head_lib._multi_class_head_with_softmax_cross_entropy_loss( n_classes=n_classes, loss_fn=loss_fn, - loss_reduction=core_losses.Reduction.SUM_OVER_NONZERO_WEIGHTS) + loss_reduction=loss_reduction, + weight_column=weight_column) # pylint:enable=protected-access return head_fn +# For quantile regression, use this head with Core..Estimator, or use +# Core..QuantileRegressor directly, +def core_quantile_regression_head( + quantiles, + label_dimension=1, + weight_column=None, + loss_reduction=core_losses.Reduction.SUM_OVER_NONZERO_WEIGHTS): + """Core head for quantile regression problems.""" + + def loss_fn(labels, logits): + result = losses.per_example_quantile_regression_loss( + labels=labels, + predictions=logits, + weights=weight_column, + quantile=quantiles) + return result[0] + + # pylint:disable=protected-access + head_fn = core_head_lib._regression_head( + label_dimension=label_dimension, + loss_fn=loss_fn, + loss_reduction=loss_reduction, + weight_column=weight_column) + # pylint:enable=protected-access + return head_fn + + class CoreGradientBoostedDecisionTreeEstimator(core_estimator.Estimator): """An estimator using gradient boosted decision trees. @@ -601,3 +741,104 @@ class CoreGradientBoostedDecisionTreeRanker(core_estimator.Estimator): super(CoreGradientBoostedDecisionTreeRanker, self).__init__( model_fn=_model_fn, model_dir=model_dir, config=config) + + +# When using this estimator, make sure to regularize the hessian (at least l2, +# min_node_weight)! +# TODO(nponomareva): extend to take multiple quantiles in one go. +class CoreGradientBoostedDecisionTreeQuantileRegressor( + core_estimator.Estimator): + """An estimator that does quantile regression and returns quantile estimates. + """ + + def __init__(self, + learner_config, + examples_per_layer, + quantiles, + label_dimension=1, + num_trees=None, + feature_columns=None, + weight_column_name=None, + model_dir=None, + config=None, + label_keys=None, + feature_engineering_fn=None, + logits_modifier_function=None, + center_bias=True, + output_leaf_index=False, + num_quantiles=100): + """Initializes a core version of GradientBoostedDecisionTreeEstimator. + + Args: + learner_config: A config for the learner. + examples_per_layer: Number of examples to accumulate before growing a + layer. It can also be a function that computes the number of examples + based on the depth of the layer that's being built. + quantiles: a list of quantiles for the loss, each between 0 and 1. + label_dimension: Dimension of regression label. This is the size + of the last dimension of the labels `Tensor` (typically, this has shape + `[batch_size, label_dimension]`). When label_dimension>1, it is + recommended to use multiclass strategy diagonal hessian or full hessian. + num_trees: An int, number of trees to build. + feature_columns: A list of feature columns. + weight_column_name: Name of the column for weights, or None if not + weighted. + model_dir: Directory for model exports, etc. + config: `RunConfig` object to configure the runtime settings. + label_keys: Optional list of strings with size `[n_classes]` defining the + label vocabulary. Only supported for `n_classes` > 2. + feature_engineering_fn: Feature engineering function. Takes features and + labels which are the output of `input_fn` and returns features and + labels which will be fed into the model. + logits_modifier_function: A modifier function for the logits. + center_bias: Whether a separate tree should be created for first fitting + the bias. + output_leaf_index: whether to output leaf indices along with predictions + during inference. The leaf node indexes are available in predictions + dict by the key 'leaf_index'. For example, + result_dict = classifier.predict(...) + for example_prediction_result in result_dict: + # access leaf index list by example_prediction_result["leaf_index"] + # which contains one leaf index per tree + num_quantiles: Number of quantiles to build for numeric feature values. + """ + if len(quantiles) > 1: + raise ValueError('For now, just one quantile per estimator is supported') + + def _model_fn(features, labels, mode, config): + return model.model_builder( + features=features, + labels=labels, + mode=mode, + config=config, + params={ + 'head': + core_quantile_regression_head( + quantiles[0], label_dimension=label_dimension), + 'feature_columns': + feature_columns, + 'learner_config': + learner_config, + 'num_trees': + num_trees, + 'weight_column_name': + weight_column_name, + 'examples_per_layer': + examples_per_layer, + 'center_bias': + center_bias, + 'logits_modifier_function': + logits_modifier_function, + 'use_core_libs': + True, + 'output_leaf_index': + output_leaf_index, + 'override_global_step_value': + None, + 'num_quantiles': + num_quantiles, + }, + output_type=model.ModelBuilderOutputType.ESTIMATOR_SPEC) + + super(CoreGradientBoostedDecisionTreeQuantileRegressor, self).__init__( + model_fn=_model_fn, model_dir=model_dir, config=config) diff --git a/tensorflow/contrib/boosted_trees/estimator_batch/estimator_test.py b/tensorflow/contrib/boosted_trees/estimator_batch/estimator_test.py index c155128c0e4ccf928349ee6453baff4384222096..47d910d42a27db4b857eeb12209dfbb429dd1be2 100644 --- a/tensorflow/contrib/boosted_trees/estimator_batch/estimator_test.py +++ b/tensorflow/contrib/boosted_trees/estimator_batch/estimator_test.py @@ -25,6 +25,7 @@ from tensorflow.contrib.boosted_trees.proto import learner_pb2 from tensorflow.contrib.layers.python.layers import feature_column as contrib_feature_column from tensorflow.contrib.learn.python.learn.estimators import run_config from tensorflow.python.estimator.canned import head as head_lib +from tensorflow.python.estimator.inputs import numpy_io from tensorflow.python.feature_column import feature_column_lib as core_feature_column from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes @@ -47,8 +48,8 @@ def _multiclass_train_input_fn(): features = { "x": constant_op.constant([[2.], [1.], [1.], [5.], [3.5], [4.6], [3.5]]) } - label = constant_op.constant( - [[1], [0], [0], [2], [2], [0], [1]], dtype=dtypes.int32) + label = constant_op.constant([[1], [0], [0], [2], [2], [0], [1]], + dtype=dtypes.int32) return features, label @@ -77,6 +78,59 @@ def _infer_ranking_train_input_fn(): return features, None +_QUANTILE_REGRESSION_SIZE = 1000 + + +def _quantile_regression_input_fns(two_dimension=False): + # The data generation is taken from + # http://scikit-learn.org/stable/auto_examples/ensemble/plot_gradient_boosting_quantile.html + np.random.seed(1) + + def f(x): + """The function to predict.""" + return x * np.sin(x) + + def g(x): + """The function to predict.""" + return x * np.cos(x) + + # Training data. + x = np.atleast_2d(np.random.uniform(0, 10.0, + size=_QUANTILE_REGRESSION_SIZE)).T + x = x.astype(np.float32) + + # Labels. + if not two_dimension: + y = f(x).ravel() + else: + y = np.column_stack((f(x).ravel(), g(x).ravel())) + + # Add random noise. + dy = 1.5 + 1.0 * np.random.random(y.shape) + noise = np.random.normal(0, dy) + y += noise + y_original = y.astype(np.float32) + if not two_dimension: + y = y.reshape(_QUANTILE_REGRESSION_SIZE, 1) + + train_input_fn = numpy_io.numpy_input_fn( + x=x, + y=y, + batch_size=_QUANTILE_REGRESSION_SIZE, + num_epochs=None, + shuffle=True) + + # Test on the training data to make sure the predictions are calibrated. + test_input_fn = numpy_io.numpy_input_fn( + x=x, + y=y, + batch_size=_QUANTILE_REGRESSION_SIZE, + num_epochs=1, + shuffle=False) + + return train_input_fn, test_input_fn, y_original + + class BoostedTreeEstimatorTest(test_util.TensorFlowTestCase): def setUp(self): @@ -341,6 +395,130 @@ class BoostedTreeEstimatorTest(test_util.TensorFlowTestCase): for prediction_dict in result_iter: self.assertTrue("classes" in prediction_dict) + # One dimensional quantile regression. + def testQuantileRegression(self): + learner_config = learner_pb2.LearnerConfig() + learner_config.num_classes = 2 + learner_config.constraints.max_tree_depth = 3 + learner_config.growing_mode = learner_pb2.LearnerConfig.WHOLE_TREE + learner_config.constraints.min_node_weight = 1 / _QUANTILE_REGRESSION_SIZE + learner_config.regularization.l2 = 1.0 / _QUANTILE_REGRESSION_SIZE + learner_config.regularization.l1 = 1.0 / _QUANTILE_REGRESSION_SIZE + learner_config.regularization.tree_complexity = ( + 1.0 / _QUANTILE_REGRESSION_SIZE) + + train_input_fn, test_input_fn, y = _quantile_regression_input_fns() + + # 95% percentile. + model_upper = estimator.GradientBoostedDecisionTreeQuantileRegressor( + quantiles=[0.95], + learner_config=learner_config, + num_trees=100, + examples_per_layer=_QUANTILE_REGRESSION_SIZE, + center_bias=False) + + model_upper.fit(input_fn=train_input_fn, steps=1000) + result_iter = model_upper.predict(input_fn=test_input_fn) + upper = [] + for prediction_dict in result_iter: + upper.append(prediction_dict["scores"]) + + frac_below_upper = round(1. * np.count_nonzero(upper > y) / len(y), 3) + # +/- 3% + self.assertTrue(frac_below_upper >= 0.92) + self.assertTrue(frac_below_upper <= 0.98) + + train_input_fn, test_input_fn, _ = _quantile_regression_input_fns() + model_lower = estimator.GradientBoostedDecisionTreeQuantileRegressor( + quantiles=[0.05], + learner_config=learner_config, + num_trees=100, + examples_per_layer=_QUANTILE_REGRESSION_SIZE, + center_bias=False) + + model_lower.fit(input_fn=train_input_fn, steps=1000) + result_iter = model_lower.predict(input_fn=test_input_fn) + lower = [] + for prediction_dict in result_iter: + lower.append(prediction_dict["scores"]) + + frac_above_lower = round(1. * np.count_nonzero(lower < y) / len(y), 3) + # +/- 3% + self.assertTrue(frac_above_lower >= 0.92) + self.assertTrue(frac_above_lower <= 0.98) + + # Multi-dimensional quantile regression. + def testQuantileRegressionMultiDimLabel(self): + learner_config = learner_pb2.LearnerConfig() + learner_config.num_classes = 2 + learner_config.constraints.max_tree_depth = 3 + learner_config.growing_mode = learner_pb2.LearnerConfig.WHOLE_TREE + learner_config.constraints.min_node_weight = 1 / _QUANTILE_REGRESSION_SIZE + learner_config.regularization.l2 = 1.0 / _QUANTILE_REGRESSION_SIZE + learner_config.regularization.l1 = 1.0 / _QUANTILE_REGRESSION_SIZE + learner_config.regularization.tree_complexity = ( + 1.0 / _QUANTILE_REGRESSION_SIZE) + + train_input_fn, test_input_fn, y = _quantile_regression_input_fns( + two_dimension=True) + + # 95% percentile. + model_upper = estimator.GradientBoostedDecisionTreeQuantileRegressor( + quantiles=[0.95], + learner_config=learner_config, + label_dimension=2, + num_trees=100, + examples_per_layer=_QUANTILE_REGRESSION_SIZE, + center_bias=False) + + model_upper.fit(input_fn=train_input_fn, steps=1000) + result_iter = model_upper.predict(input_fn=test_input_fn) + upper = [] + for prediction_dict in result_iter: + upper.append(prediction_dict["scores"]) + + count_below_upper = np.count_nonzero(upper > y, axis=0) + count_both_below_upper = np.count_nonzero(np.prod(upper > y, axis=1)) + frac_below_upper_0 = round(1. * count_below_upper[0] / len(y), 3) + frac_below_upper_1 = round(1. * count_below_upper[1] / len(y), 3) + frac_both_below_upper = round(1. * count_both_below_upper / len(y), 3) + # +/- 3% + self.assertTrue(frac_below_upper_0 >= 0.92) + 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.91) + self.assertTrue(frac_both_below_upper <= 0.99) + + train_input_fn, test_input_fn, _ = _quantile_regression_input_fns( + two_dimension=True) + model_lower = estimator.GradientBoostedDecisionTreeQuantileRegressor( + quantiles=[0.05], + learner_config=learner_config, + label_dimension=2, + num_trees=100, + examples_per_layer=_QUANTILE_REGRESSION_SIZE, + center_bias=False) + + model_lower.fit(input_fn=train_input_fn, steps=1000) + result_iter = model_lower.predict(input_fn=test_input_fn) + lower = [] + for prediction_dict in result_iter: + lower.append(prediction_dict["scores"]) + + count_above_lower = np.count_nonzero(lower < y, axis=0) + count_both_aboce_lower = np.count_nonzero(np.prod(lower < y, axis=1)) + frac_above_lower_0 = round(1. * count_above_lower[0] / len(y), 3) + frac_above_lower_1 = round(1. * count_above_lower[1] / len(y), 3) + frac_both_above_lower = round(1. * count_both_aboce_lower / len(y), 3) + # +/- 3% + self.assertTrue(frac_above_lower_0 >= 0.92) + 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.91) + self.assertTrue(frac_both_above_lower <= 0.99) + class CoreGradientBoostedDecisionTreeEstimators(test_util.TensorFlowTestCase): @@ -489,8 +667,8 @@ class CoreGradientBoostedDecisionTreeEstimators(test_util.TensorFlowTestCase): feature_columns = [ core_feature_column.weighted_categorical_column( - categorical_column=core_feature_column. - categorical_column_with_vocabulary_list( + categorical_column=core_feature_column + .categorical_column_with_vocabulary_list( key="word", vocabulary_list=["the", "cat", "dog"]), weight_feature_key="weight") ] @@ -509,8 +687,8 @@ class CoreGradientBoostedDecisionTreeEstimators(test_util.TensorFlowTestCase): # Weights for the words are 5 - cat, 6- dog and 1 -the. features_dict["word"] = sparse_tensor.SparseTensor( indices=[[0, 0], [0, 1], [1, 0], [3, 0]], - values=constant_op.constant( - ["the", "cat", "dog", "the"], dtype=dtypes.string), + values=constant_op.constant(["the", "cat", "dog", "the"], + dtype=dtypes.string), dense_shape=[4, 3]) features_dict["weight"] = sparse_tensor.SparseTensor( indices=[[0, 0], [0, 1], [1, 0], [3, 0]], @@ -534,6 +712,132 @@ class CoreGradientBoostedDecisionTreeEstimators(test_util.TensorFlowTestCase): est.evaluate(input_fn=input_fn, steps=1) est.predict(input_fn=input_fn) + # One dimensional quantile regression. + def testQuantileRegression(self): + learner_config = learner_pb2.LearnerConfig() + learner_config.num_classes = 2 + learner_config.constraints.max_tree_depth = 3 + learner_config.growing_mode = learner_pb2.LearnerConfig.WHOLE_TREE + learner_config.constraints.min_node_weight = 1 / _QUANTILE_REGRESSION_SIZE + learner_config.regularization.l2 = 1.0 / _QUANTILE_REGRESSION_SIZE + learner_config.regularization.l1 = 1.0 / _QUANTILE_REGRESSION_SIZE + learner_config.regularization.tree_complexity = ( + 1.0 / _QUANTILE_REGRESSION_SIZE) + + train_input_fn, test_input_fn, y = _quantile_regression_input_fns() + y = y.reshape(_QUANTILE_REGRESSION_SIZE, 1) + + # 95% percentile. + model_upper = estimator.CoreGradientBoostedDecisionTreeQuantileRegressor( + quantiles=[0.95], + learner_config=learner_config, + num_trees=100, + examples_per_layer=_QUANTILE_REGRESSION_SIZE, + center_bias=False) + + model_upper.train(input_fn=train_input_fn, steps=1000) + result_iter = model_upper.predict(input_fn=test_input_fn) + upper = [] + for prediction_dict in result_iter: + upper.append(prediction_dict["predictions"]) + + frac_below_upper = round(1. * np.count_nonzero(upper > y) / len(y), 3) + # +/- 3% + self.assertTrue(frac_below_upper >= 0.92) + self.assertTrue(frac_below_upper <= 0.98) + + train_input_fn, test_input_fn, _ = _quantile_regression_input_fns() + model_lower = estimator.CoreGradientBoostedDecisionTreeQuantileRegressor( + quantiles=[0.05], + learner_config=learner_config, + num_trees=100, + examples_per_layer=_QUANTILE_REGRESSION_SIZE, + center_bias=False) + + model_lower.train(input_fn=train_input_fn, steps=1000) + result_iter = model_lower.predict(input_fn=test_input_fn) + lower = [] + for prediction_dict in result_iter: + lower.append(prediction_dict["predictions"]) + + frac_above_lower = round(1. * np.count_nonzero(lower < y) / len(y), 3) + # +/- 3% + self.assertTrue(frac_above_lower >= 0.92) + self.assertTrue(frac_above_lower <= 0.98) + + # Multi-dimensional quantile regression. + def testQuantileRegressionMultiDimLabel(self): + learner_config = learner_pb2.LearnerConfig() + learner_config.num_classes = 2 + learner_config.constraints.max_tree_depth = 3 + learner_config.growing_mode = learner_pb2.LearnerConfig.WHOLE_TREE + learner_config.constraints.min_node_weight = 1 / _QUANTILE_REGRESSION_SIZE + learner_config.regularization.l2 = 1.0 / _QUANTILE_REGRESSION_SIZE + learner_config.regularization.l1 = 1.0 / _QUANTILE_REGRESSION_SIZE + learner_config.regularization.tree_complexity = ( + 1.0 / _QUANTILE_REGRESSION_SIZE) + + train_input_fn, test_input_fn, y = _quantile_regression_input_fns( + two_dimension=True) + y = y.reshape(_QUANTILE_REGRESSION_SIZE, 2) + + # 95% percentile. + model_upper = estimator.CoreGradientBoostedDecisionTreeQuantileRegressor( + quantiles=[0.95], + learner_config=learner_config, + num_trees=100, + label_dimension=2, + examples_per_layer=_QUANTILE_REGRESSION_SIZE, + center_bias=False) + + model_upper.train(input_fn=train_input_fn, steps=1000) + result_iter = model_upper.predict(input_fn=test_input_fn) + upper = [] + for prediction_dict in result_iter: + upper.append(prediction_dict["predictions"]) + + count_below_upper = np.count_nonzero(upper > y, axis=0) + count_both_below_upper = np.count_nonzero(np.prod(upper > y, axis=1)) + frac_below_upper_0 = round(1. * count_below_upper[0] / len(y), 3) + frac_below_upper_1 = round(1. * count_below_upper[1] / len(y), 3) + frac_both_below_upper = round(1. * count_both_below_upper / len(y), 3) + # +/- 3% + self.assertTrue(frac_below_upper_0 >= 0.92) + 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.91) + self.assertTrue(frac_both_below_upper <= 0.99) + + train_input_fn, test_input_fn, _ = _quantile_regression_input_fns( + two_dimension=True) + model_lower = estimator.CoreGradientBoostedDecisionTreeQuantileRegressor( + quantiles=[0.05], + learner_config=learner_config, + num_trees=100, + label_dimension=2, + examples_per_layer=_QUANTILE_REGRESSION_SIZE, + center_bias=False) + + model_lower.train(input_fn=train_input_fn, steps=1000) + result_iter = model_lower.predict(input_fn=test_input_fn) + lower = [] + for prediction_dict in result_iter: + lower.append(prediction_dict["predictions"]) + + count_above_lower = np.count_nonzero(lower < y, axis=0) + count_both_aboce_lower = np.count_nonzero(np.prod(lower < y, axis=1)) + frac_above_lower_0 = round(1. * count_above_lower[0] / len(y), 3) + frac_above_lower_1 = round(1. * count_above_lower[1] / len(y), 3) + frac_both_above_lower = round(1. * count_both_aboce_lower / len(y), 3) + # +/- 3% + self.assertTrue(frac_above_lower_0 >= 0.92) + 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.91) + self.assertTrue(frac_both_above_lower <= 0.99) + if __name__ == "__main__": googletest.main() diff --git a/tensorflow/contrib/boosted_trees/examples/boston.py b/tensorflow/contrib/boosted_trees/examples/boston.py index 54c4ff059e3408d2cb8fc689a9ae877f57485f58..09b240a7006a8ef53eb95108b3adbfae728cf8fc 100644 --- a/tensorflow/contrib/boosted_trees/examples/boston.py +++ b/tensorflow/contrib/boosted_trees/examples/boston.py @@ -90,13 +90,13 @@ def _make_experiment_fn(output_dir): (x_train, y_train), (x_test, y_test) = tf.keras.datasets.boston_housing.load_data() - train_input_fn = tf.estimator.inputs.numpy_input_fn( + train_input_fn = tf.compat.v1.estimator.inputs.numpy_input_fn( x={"x": x_train}, y=y_train, batch_size=FLAGS.batch_size, num_epochs=None, shuffle=True) - eval_input_fn = tf.estimator.inputs.numpy_input_fn( + eval_input_fn = tf.compat.v1.estimator.inputs.numpy_input_fn( x={"x": x_test}, y=y_test, num_epochs=1, shuffle=False) feature_columns = [ diff --git a/tensorflow/contrib/boosted_trees/examples/boston_combined.py b/tensorflow/contrib/boosted_trees/examples/boston_combined.py index e04b56afbfd266dc13a5b0d78d171ea273415ee3..d640af354f55423b7c9706900359f5e64c459f39 100644 --- a/tensorflow/contrib/boosted_trees/examples/boston_combined.py +++ b/tensorflow/contrib/boosted_trees/examples/boston_combined.py @@ -80,13 +80,13 @@ def _make_experiment_fn(output_dir): (x_train, y_train), (x_test, y_test) = tf.keras.datasets.boston_housing.load_data() - train_input_fn = tf.estimator.inputs.numpy_input_fn( + train_input_fn = tf.compat.v1.estimator.inputs.numpy_input_fn( x={"x": x_train}, y=y_train, batch_size=FLAGS.batch_size, num_epochs=None, shuffle=True) - eval_input_fn = tf.estimator.inputs.numpy_input_fn( + eval_input_fn = tf.compat.v1.estimator.inputs.numpy_input_fn( x={"x": x_test}, y=y_test, num_epochs=1, shuffle=False) feature_columns = [ diff --git a/tensorflow/contrib/boosted_trees/kernels/split_handler_ops.cc b/tensorflow/contrib/boosted_trees/kernels/split_handler_ops.cc index 8edb5d6c640611bbb90d7731b2fea4354e125563..6d78e27e8f69ea289b686af8402bd91967f997f4 100644 --- a/tensorflow/contrib/boosted_trees/kernels/split_handler_ops.cc +++ b/tensorflow/contrib/boosted_trees/kernels/split_handler_ops.cc @@ -834,8 +834,13 @@ class BuildCategoricalEqualitySplitsOp : public OpKernel { root_gradient_stats *= normalizer_ratio; NodeStats root_stats = state->ComputeNodeStats(root_gradient_stats); int32 best_feature_idx = 0; + bool best_feature_updated = false; NodeStats best_right_node_stats(0); NodeStats best_left_node_stats(0); + CHECK(end_index - start_index >= 2) + << "Partition should have a non bias feature. Start index " + << start_index << " and end index " << end_index; + for (int64 feature_idx = start_index + 1; feature_idx < end_index; ++feature_idx) { GradientStats left_gradient_stats(*gradients_t, *hessians_t, @@ -845,11 +850,13 @@ class BuildCategoricalEqualitySplitsOp : public OpKernel { root_gradient_stats - left_gradient_stats; NodeStats left_stats = state->ComputeNodeStats(left_gradient_stats); NodeStats right_stats = state->ComputeNodeStats(right_gradient_stats); - if (left_stats.gain + right_stats.gain > best_gain) { + if (!best_feature_updated || + left_stats.gain + right_stats.gain > best_gain) { best_gain = left_stats.gain + right_stats.gain; best_left_node_stats = left_stats; best_right_node_stats = right_stats; best_feature_idx = feature_idx; + best_feature_updated = true; } } SplitInfo split_info; @@ -864,7 +871,7 @@ class BuildCategoricalEqualitySplitsOp : public OpKernel { << feature_ids(best_feature_idx, 0) << ", " << feature_ids(best_feature_idx, 1) << "\nPartition IDS: " << partition_ids(start_index) << " " - << partition_ids(best_feature_idx); + << partition_ids(best_feature_idx) << " and best gain " << best_gain; equality_split->set_feature_id(feature_ids(best_feature_idx, 0)); auto* left_child = split_info.mutable_left_child(); auto* right_child = split_info.mutable_right_child(); 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/kernels/training_ops.cc b/tensorflow/contrib/boosted_trees/kernels/training_ops.cc index ab2853352a70073648f47e9835f8a66852ff584f..a30cfa663f4a4954f83224a7fd6448b369ad93b4 100644 --- a/tensorflow/contrib/boosted_trees/kernels/training_ops.cc +++ b/tensorflow/contrib/boosted_trees/kernels/training_ops.cc @@ -382,8 +382,7 @@ class GrowTreeEnsembleOp : public OpKernel { break; } case LearnerConfig::OBLIVIOUS_DECISION_TREE: { - FindBestSplitsPerPartitionOblivious(context, gains_list, splits_list, - &best_splits); + FindBestSplitOblivious(context, gains_list, splits_list, &best_splits); break; } } @@ -475,10 +474,10 @@ class GrowTreeEnsembleOp : public OpKernel { } } - void FindBestSplitsPerPartitionOblivious( - OpKernelContext* const context, const OpInputList& gains_list, - const OpInputList& splits_list, - std::map* best_splits) { + void FindBestSplitOblivious(OpKernelContext* const context, + const OpInputList& gains_list, + const OpInputList& splits_list, + std::map* best_splits) { // Find best split per partition going through every feature candidate. for (int64 handler_id = 0; handler_id < num_handlers_; ++handler_id) { const auto& gains = gains_list[handler_id].vec(); @@ -654,6 +653,12 @@ class GrowTreeEnsembleOp : public OpKernel { return dest; } + if (dest->leaf_case() == boosted_trees::trees::Leaf::LEAF_NOT_SET) { + // No merging is required. Just copy the source weights; + *dest = source; + return dest; + } + // Handle leaf merging based on type. switch (source.leaf_case()) { case boosted_trees::trees::Leaf::kVector: { diff --git a/tensorflow/contrib/boosted_trees/lib/BUILD b/tensorflow/contrib/boosted_trees/lib/BUILD index 3028c2281705bd7e34b212332160d25386559d4e..fd832de982a4a7a2bd39e450ad495e60c284ace7 100644 --- a/tensorflow/contrib/boosted_trees/lib/BUILD +++ b/tensorflow/contrib/boosted_trees/lib/BUILD @@ -67,6 +67,7 @@ tf_cc_test( "//tensorflow/core:tensor_testutil", "//tensorflow/core:test", "//tensorflow/core:test_main", + "@com_google_absl//absl/algorithm:container", ], ) diff --git a/tensorflow/contrib/boosted_trees/lib/learner/batch/base_split_handler.py b/tensorflow/contrib/boosted_trees/lib/learner/batch/base_split_handler.py index 5d4819b0f1cb598cfbe146f569aecd7883186339..efa2ab1dad8df9815c983afaa2e43982a49c5787 100644 --- a/tensorflow/contrib/boosted_trees/lib/learner/batch/base_split_handler.py +++ b/tensorflow/contrib/boosted_trees/lib/learner/batch/base_split_handler.py @@ -19,15 +19,17 @@ from __future__ import division from __future__ import print_function import abc + +import six + from tensorflow.contrib.boosted_trees.python.ops import batch_ops_utils from tensorflow.python.ops import control_flow_ops +@six.add_metaclass(abc.ABCMeta) class BaseSplitHandler(object): """Abstract Base class defining split handlers interface.""" - __metaclass__ = abc.ABCMeta - def __init__(self, l1_regularization, l2_regularization, diff --git a/tensorflow/contrib/boosted_trees/lib/learner/batch/categorical_split_handler.py b/tensorflow/contrib/boosted_trees/lib/learner/batch/categorical_split_handler.py index 4da25298cb82093ac501997cc21c48265df06860..d26af58419752170bbc58bba757ac43349fc2cff 100644 --- a/tensorflow/contrib/boosted_trees/lib/learner/batch/categorical_split_handler.py +++ b/tensorflow/contrib/boosted_trees/lib/learner/batch/categorical_split_handler.py @@ -119,7 +119,7 @@ class EqualitySplitHandler(base_split_handler.BaseSplitHandler): def not_active_inputs(): return (constant_op.constant([], dtype=dtypes.int32), - constant_op.constant([], dtype=dtypes.int64, shape=[1, 2]), + constant_op.constant_v1([], dtype=dtypes.int64, shape=[1, 2]), empty_gradients, empty_hessians) def active_inputs(): diff --git a/tensorflow/contrib/boosted_trees/lib/learner/batch/categorical_split_handler_test.py b/tensorflow/contrib/boosted_trees/lib/learner/batch/categorical_split_handler_test.py index a2f708081a4b484d649b5d09b172c2c60db69aeb..386dc19fc7b9529993a9625fb1298f6eb9a70d87 100644 --- a/tensorflow/contrib/boosted_trees/lib/learner/batch/categorical_split_handler_test.py +++ b/tensorflow/contrib/boosted_trees/lib/learner/batch/categorical_split_handler_test.py @@ -36,9 +36,9 @@ def get_empty_tensors(gradient_shape, hessian_shape): empty_hess_shape = [1] + hessian_shape.as_list() empty_grad_shape = [1] + gradient_shape.as_list() - empty_gradients = constant_op.constant( + empty_gradients = constant_op.constant_v1( [], dtype=dtypes.float32, shape=empty_grad_shape) - empty_hessians = constant_op.constant( + empty_hessians = constant_op.constant_v1( [], dtype=dtypes.float32, shape=empty_hess_shape) return empty_gradients, empty_hessians @@ -486,8 +486,8 @@ class EqualitySplitHandlerTest(test_util.TensorFlowTestCase): gradients = array_ops.constant([0.2, -0.5, 1.2, 4.0]) hessians = array_ops.constant([0.12, 0.07, 0.2, 0.13]) partition_ids = [0, 0, 0, 1] - indices = array_ops.constant([], dtype=dtypes.int64, shape=[0, 2]) - values = array_ops.constant([], dtype=dtypes.int64) + indices = constant_op.constant_v1([], dtype=dtypes.int64, shape=[0, 2]) + values = constant_op.constant_v1([], dtype=dtypes.int64) gradient_shape = tensor_shape.scalar() hessian_shape = tensor_shape.scalar() diff --git a/tensorflow/contrib/boosted_trees/lib/learner/batch/ordinal_split_handler.py b/tensorflow/contrib/boosted_trees/lib/learner/batch/ordinal_split_handler.py index f45010ec26ed25127ca78b97f4d6fd7ebd6467ae..0476bed2cd3f3ea5b47b10c51a819f17d6e37c74 100644 --- a/tensorflow/contrib/boosted_trees/lib/learner/batch/ordinal_split_handler.py +++ b/tensorflow/contrib/boosted_trees/lib/learner/batch/ordinal_split_handler.py @@ -142,7 +142,7 @@ class InequalitySplitHandler(base_split_handler.BaseSplitHandler): name="StatsAccumulator/{}".format(self._name)) # Allocate both stats accumulator and quantile accumulator on the same # device so that we can build splits with fewer RPCs. - with ops.colocate_with(self._stats_accumulator.resource()): + with ops.colocate_with(self._stats_accumulator.resource_handle): self._quantile_accumulator = quantile_ops.QuantileAccumulator( init_stamp_token, epsilon=epsilon, @@ -268,8 +268,8 @@ class DenseSplitHandler(InequalitySplitHandler): handler = make_dense_split_tensor are_splits_ready, partition_ids, gains, split_infos = ( - handler(self._quantile_accumulator.resource(), - self._stats_accumulator.resource(), stamp_token, + handler(self._quantile_accumulator.resource_handle, + self._stats_accumulator.resource_handle, stamp_token, next_stamp_token, self._multiclass_strategy, class_id, self._feature_column_group_id, self._l1_regularization, self._l2_regularization, self._tree_complexity_regularization, @@ -447,8 +447,8 @@ class SparseSplitHandler(InequalitySplitHandler): handler = make_sparse_split_tensor are_splits_ready, partition_ids, gains, split_infos = ( - handler(self._quantile_accumulator.resource(), - self._stats_accumulator.resource(), stamp_token, + handler(self._quantile_accumulator.resource_handle, + self._stats_accumulator.resource_handle, stamp_token, next_stamp_token, self._multiclass_strategy, class_id, self._feature_column_group_id, self._l1_regularization, self._l2_regularization, self._tree_complexity_regularization, @@ -605,7 +605,7 @@ def dense_make_stats_update(is_active, are_buckets_ready, float_column, quantile_buckets, example_partition_ids, gradients, hessians, weights, empty_gradients, empty_hessians): """Updates the state for dense split handler.""" - empty_float = constant_op.constant([], dtype=dtypes.float32) + empty_float = constant_op.constant_v1([], dtype=dtypes.float32) quantile_values, quantile_weights = control_flow_ops.cond( is_active[1], # For the next layer, this handler is inactive. @@ -621,8 +621,8 @@ def dense_make_stats_update(is_active, are_buckets_ready, float_column, return (example_partition_ids, quantized_feature, gradients, hessians) def not_ready_inputs_fn(): - return (constant_op.constant([], dtype=dtypes.int32), - constant_op.constant([[]], dtype=dtypes.int64, shape=[1, 2]), + return (constant_op.constant_v1([], dtype=dtypes.int32), + constant_op.constant_v1([[]], dtype=dtypes.int64, shape=[1, 2]), empty_gradients, empty_hessians) example_partition_ids, feature_ids, gradients, hessians = ( @@ -708,11 +708,11 @@ def sparse_make_stats_update( def quantiles_not_ready(): """The subgraph for when the quantiles are not ready.""" - return (constant_op.constant([], dtype=dtypes.int32), - constant_op.constant([], dtype=dtypes.int64, shape=[1, 2]), + return (constant_op.constant_v1([], dtype=dtypes.int32), + constant_op.constant_v1([], dtype=dtypes.int64, shape=[1, 2]), empty_gradients, empty_hessians) - empty_float = constant_op.constant([], dtype=dtypes.float32) + empty_float = constant_op.constant_v1([], dtype=dtypes.float32) handler_not_active = (constant_op.constant( [], dtype=dtypes.int64, shape=[0, 2]), empty_float, constant_op.constant([0, 1], dtype=dtypes.int64), diff --git a/tensorflow/contrib/boosted_trees/lib/learner/batch/ordinal_split_handler_test.py b/tensorflow/contrib/boosted_trees/lib/learner/batch/ordinal_split_handler_test.py index 74b0ea6989c65e83e7a466107d624712a0e72d1b..4a1b528646e7d2139d7eabb0264b8d280f8da133 100644 --- a/tensorflow/contrib/boosted_trees/lib/learner/batch/ordinal_split_handler_test.py +++ b/tensorflow/contrib/boosted_trees/lib/learner/batch/ordinal_split_handler_test.py @@ -39,9 +39,9 @@ def get_empty_tensors(gradient_shape, hessian_shape): empty_hess_shape = [1] + hessian_shape.as_list() empty_grad_shape = [1] + gradient_shape.as_list() - empty_gradients = constant_op.constant( + empty_gradients = constant_op.constant_v1( [], dtype=dtypes.float32, shape=empty_grad_shape) - empty_hessians = constant_op.constant( + empty_hessians = constant_op.constant_v1( [], dtype=dtypes.float32, shape=empty_hess_shape) return empty_gradients, empty_hessians @@ -1476,9 +1476,9 @@ class SparseSplitHandlerTest(test_util.TensorFlowTestCase): def testEmpty(self): with self.cached_session() as sess: - indices = array_ops.constant([], dtype=dtypes.int64, shape=[0, 2]) + indices = constant_op.constant_v1([], dtype=dtypes.int64, shape=[0, 2]) # No values in this feature column in this mini-batch. - values = array_ops.constant([], dtype=dtypes.float32) + values = constant_op.constant_v1([], dtype=dtypes.float32) sparse_column = sparse_tensor.SparseTensor(indices, values, [4, 1]) gradient_shape = tensor_shape.scalar() @@ -1549,8 +1549,9 @@ class SparseSplitHandlerTest(test_util.TensorFlowTestCase): sparse_column = array_ops.sparse_placeholder(dtypes.float32) # We have two batches - at first, a sparse feature is empty. - empty_indices = array_ops.constant([], dtype=dtypes.int64, shape=[0, 2]) - empty_values = array_ops.constant([], dtype=dtypes.float32) + empty_indices = constant_op.constant_v1([], dtype=dtypes.int64, + shape=[0, 2]) + empty_values = constant_op.constant_v1([], dtype=dtypes.float32) empty_sparse_column = sparse_tensor.SparseTensor(empty_indices, empty_values, [4, 2]) empty_sparse_column = empty_sparse_column.eval(session=sess) diff --git a/tensorflow/contrib/boosted_trees/lib/trees/decision_tree.cc b/tensorflow/contrib/boosted_trees/lib/trees/decision_tree.cc index 64921faf81c0ea8ae7fb1bbec71396ef3408e6ca..de30a7bde792e727ceab7798458566d4527f5867 100644 --- a/tensorflow/contrib/boosted_trees/lib/trees/decision_tree.cc +++ b/tensorflow/contrib/boosted_trees/lib/trees/decision_tree.cc @@ -81,9 +81,10 @@ int DecisionTree::Traverse(const DecisionTreeConfig& config, const auto& split = current_node.categorical_id_binary_split(); const auto& features = example.sparse_int_features[split.feature_column()]; - node_id = features.find(split.feature_id()) != features.end() - ? split.left_id() - : split.right_id(); + node_id = (std::find(features.begin(), features.end(), + split.feature_id()) == features.end()) + ? split.right_id() + : split.left_id(); break; } case TreeNode::kCategoricalIdSetMembershipBinarySplit: { @@ -117,7 +118,8 @@ int DecisionTree::Traverse(const DecisionTreeConfig& config, oblivious_leaf_idx <<= 1; const auto& features = example.sparse_int_features[split.feature_column()]; - if (features.find(split.feature_id()) == features.end()) { + if (std::find(features.begin(), features.end(), split.feature_id()) == + features.end()) { oblivious_leaf_idx++; } node_id++; diff --git a/tensorflow/contrib/boosted_trees/lib/utils/example.h b/tensorflow/contrib/boosted_trees/lib/utils/example.h index 1371ff337f78dd1c38f2bd0ba86911642f3aeb3e..445ffaaa714c4a69710f9a21d5f2775b8b0f6e22 100644 --- a/tensorflow/contrib/boosted_trees/lib/utils/example.h +++ b/tensorflow/contrib/boosted_trees/lib/utils/example.h @@ -20,6 +20,7 @@ #include #include #include "tensorflow/contrib/boosted_trees/lib/utils/optional_value.h" +#include "tensorflow/core/lib/gtl/inlined_vector.h" namespace tensorflow { namespace boosted_trees { @@ -124,7 +125,9 @@ struct Example { // Sparse integer features indexed by feature column. // Note that all integer features are assumed to be categorical, i.e. will // never be compared by order. Also these features can be multivalent. - std::vector> sparse_int_features; + // By default we allocate a InlinedVector of length 1 though since that is + // the most common case. + std::vector> sparse_int_features; }; } // namespace utils diff --git a/tensorflow/contrib/boosted_trees/lib/utils/examples_iterable.h b/tensorflow/contrib/boosted_trees/lib/utils/examples_iterable.h index 1b654e1c44e545fb97216ad950f3cd2d3240ffd0..3c5e0fbbb40a916e6a3c4197007fb2b562682aae 100644 --- a/tensorflow/contrib/boosted_trees/lib/utils/examples_iterable.h +++ b/tensorflow/contrib/boosted_trees/lib/utils/examples_iterable.h @@ -148,7 +148,7 @@ class ExamplesIterable { row_range.start); for (int64 row_idx = row_range.start; row_idx < row_range.end; ++row_idx) { - sparse_int_features[sparse_int_idx].insert( + sparse_int_features[sparse_int_idx].push_back( iter_->sparse_int_column_values_[sparse_int_idx](row_idx)); } } diff --git a/tensorflow/contrib/boosted_trees/lib/utils/examples_iterable_test.cc b/tensorflow/contrib/boosted_trees/lib/utils/examples_iterable_test.cc index 30c37435fe16ef29a9e29202850501098e9ac7f8..2f4f2495eaf799a35fb78e183e545f6a1e2d7790 100644 --- a/tensorflow/contrib/boosted_trees/lib/utils/examples_iterable_test.cc +++ b/tensorflow/contrib/boosted_trees/lib/utils/examples_iterable_test.cc @@ -13,6 +13,7 @@ // limitations under the License. // ============================================================================= #include "tensorflow/contrib/boosted_trees/lib/utils/examples_iterable.h" +#include "absl/algorithm/container.h" #include "tensorflow/core/framework/tensor_testutil.h" #include "tensorflow/core/lib/core/status_test_util.h" #include "tensorflow/core/platform/test.h" @@ -90,8 +91,8 @@ TEST_F(ExamplesIterableTest, Iterate) { EXPECT_EQ(1.0f, example.sparse_float_features[1][1].get_value()); EXPECT_EQ(2, example.sparse_int_features[0].size()); - EXPECT_EQ(1, example.sparse_int_features[0].count(1)); - EXPECT_EQ(1, example.sparse_int_features[0].count(8)); + EXPECT_EQ(1, absl::c_count(example.sparse_int_features[0], 1)); + EXPECT_EQ(1, absl::c_count(example.sparse_int_features[0], 8)); EXPECT_EQ(0, example.sparse_int_features[1].size()); } break; case 1: { @@ -105,9 +106,9 @@ TEST_F(ExamplesIterableTest, Iterate) { EXPECT_FALSE(example.sparse_float_features[1][1].has_value()); EXPECT_EQ(1, example.sparse_int_features[0].size()); - EXPECT_EQ(1, example.sparse_int_features[0].count(0)); + EXPECT_EQ(1, absl::c_count(example.sparse_int_features[0], 0)); EXPECT_EQ(1, example.sparse_int_features[1].size()); - EXPECT_EQ(1, example.sparse_int_features[1].count(7)); + EXPECT_EQ(1, absl::c_count(example.sparse_int_features[1], 7)); } break; case 2: { EXPECT_EQ(2, example.example_idx); @@ -122,7 +123,7 @@ TEST_F(ExamplesIterableTest, Iterate) { EXPECT_EQ(0, example.sparse_int_features[0].size()); EXPECT_EQ(1, example.sparse_int_features[1].size()); - EXPECT_EQ(1, example.sparse_int_features[1].count(13)); + EXPECT_EQ(1, absl::c_count(example.sparse_int_features[1], 13)); } break; case 3: { EXPECT_EQ(3, example.example_idx); @@ -136,10 +137,10 @@ TEST_F(ExamplesIterableTest, Iterate) { EXPECT_FALSE(example.sparse_float_features[1][1].has_value()); EXPECT_EQ(2, example.sparse_int_features[0].size()); - EXPECT_EQ(1, example.sparse_int_features[0].count(2)); - EXPECT_EQ(1, example.sparse_int_features[0].count(0)); + EXPECT_EQ(1, absl::c_count(example.sparse_int_features[0], 2)); + EXPECT_EQ(1, absl::c_count(example.sparse_int_features[0], 0)); EXPECT_EQ(1, example.sparse_int_features[1].size()); - EXPECT_EQ(1, example.sparse_int_features[1].count(4)); + EXPECT_EQ(1, absl::c_count(example.sparse_int_features[1], 4)); } break; case 4: { EXPECT_EQ(4, example.example_idx); @@ -154,7 +155,7 @@ TEST_F(ExamplesIterableTest, Iterate) { EXPECT_EQ(0, example.sparse_int_features[0].size()); EXPECT_EQ(1, example.sparse_int_features[1].size()); - EXPECT_EQ(1, example.sparse_int_features[1].count(0)); + EXPECT_EQ(1, absl::c_count(example.sparse_int_features[1], 0)); } break; case 5: { EXPECT_EQ(5, example.example_idx); @@ -191,7 +192,7 @@ TEST_F(ExamplesIterableTest, Iterate) { EXPECT_FALSE(example.sparse_float_features[1][1].has_value()); EXPECT_EQ(1, example.sparse_int_features[0].size()); - EXPECT_EQ(1, example.sparse_int_features[0].count(5)); + EXPECT_EQ(1, absl::c_count(example.sparse_int_features[0], 5)); } break; default: { LOG(QFATAL) << "Invalid example index."; } break; } 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/kernel_tests/stats_accumulator_ops_test.py b/tensorflow/contrib/boosted_trees/python/kernel_tests/stats_accumulator_ops_test.py index 05ce0884ccfff53484fdc0c26e596e7fb6fcdfd6..356ae337685d580319da16a20bbab27ccaa73255 100644 --- a/tensorflow/contrib/boosted_trees/python/kernel_tests/stats_accumulator_ops_test.py +++ b/tensorflow/contrib/boosted_trees/python/kernel_tests/stats_accumulator_ops_test.py @@ -34,7 +34,7 @@ class StatsAccumulatorScalarTest(test_util.TensorFlowTestCase): stamp_token=0, gradient_shape=tensor_shape.scalar(), hessian_shape=tensor_shape.scalar()) - with ops.control_dependencies([accumulator._create_op]): + with ops.control_dependencies([accumulator.initializer]): op1 = accumulator.add( stamp_token=0, partition_ids=[1, 2], @@ -62,7 +62,7 @@ class StatsAccumulatorScalarTest(test_util.TensorFlowTestCase): stamp_token=0, gradient_shape=tensor_shape.scalar(), hessian_shape=tensor_shape.scalar()) - with ops.control_dependencies([accumulator._create_op]): + with ops.control_dependencies([accumulator.initializer]): op1 = accumulator.add( stamp_token=0, partition_ids=[1, 2, 1], @@ -91,7 +91,7 @@ class StatsAccumulatorScalarTest(test_util.TensorFlowTestCase): stamp_token=0, gradient_shape=tensor_shape.scalar(), hessian_shape=tensor_shape.scalar()) - with ops.control_dependencies([accumulator._create_op]): + with ops.control_dependencies([accumulator.initializer]): op1 = accumulator.add( stamp_token=0, partition_ids=[1, 2], @@ -123,7 +123,7 @@ class StatsAccumulatorScalarTest(test_util.TensorFlowTestCase): stamp_token=0, gradient_shape=tensor_shape.scalar(), hessian_shape=tensor_shape.scalar()) - with ops.control_dependencies([accumulator._create_op]): + with ops.control_dependencies([accumulator.initializer]): op1 = accumulator.add( stamp_token=0, partition_ids=[1, 2], @@ -133,7 +133,7 @@ class StatsAccumulatorScalarTest(test_util.TensorFlowTestCase): with ops.control_dependencies([op1]): (stamp_token, num_updates, partition_1, feature_1, grads_1, - hessians_1) = accumulator.serialize() + hessians_1) = accumulator.saveable.serialize() # Make sure that the accumulator hasn't changed during serialization. with ops.control_dependencies([stamp_token]): num_updates_2, partition_2, feature_2, grads_2, hessians_2 = ( @@ -164,7 +164,7 @@ class StatsAccumulatorScalarTest(test_util.TensorFlowTestCase): stamp_token=0, gradient_shape=tensor_shape.scalar(), hessian_shape=tensor_shape.scalar()) - with ops.control_dependencies([accumulator._create_op]): + with ops.control_dependencies([accumulator.initializer]): # These will be deleted due to deserialize call. op1 = accumulator.add( stamp_token=0, @@ -175,7 +175,7 @@ class StatsAccumulatorScalarTest(test_util.TensorFlowTestCase): with ops.control_dependencies([op1]): deserialize = ( - accumulator.deserialize( + accumulator.saveable.deserialize( stamp_token=2, num_updates=3, partition_ids=[3, 4], @@ -223,7 +223,7 @@ class StatsAccumulatorTensorTest(test_util.TensorFlowTestCase): stamp_token=0, gradient_shape=tensor_shape.TensorShape([2]), hessian_shape=tensor_shape.TensorShape([2, 2])) - with ops.control_dependencies([accumulator._create_op]): + with ops.control_dependencies([accumulator.initializer]): op1 = accumulator.add( stamp_token=0, partition_ids=[1, 2], @@ -261,7 +261,7 @@ class StatsAccumulatorTensorTest(test_util.TensorFlowTestCase): stamp_token=0, gradient_shape=tensor_shape.TensorShape([2]), hessian_shape=tensor_shape.TensorShape([2, 2])) - with ops.control_dependencies([accumulator._create_op]): + with ops.control_dependencies([accumulator.initializer]): op1 = accumulator.add( stamp_token=0, partition_ids=[1, 2], @@ -299,7 +299,7 @@ class StatsAccumulatorTensorTest(test_util.TensorFlowTestCase): stamp_token=0, gradient_shape=tensor_shape.TensorShape([2]), hessian_shape=tensor_shape.TensorShape([2, 2])) - with ops.control_dependencies([accumulator._create_op]): + with ops.control_dependencies([accumulator.initializer]): op1 = accumulator.add( stamp_token=0, partition_ids=[1, 2], @@ -336,7 +336,7 @@ class StatsAccumulatorTensorTest(test_util.TensorFlowTestCase): stamp_token=0, gradient_shape=tensor_shape.TensorShape([2]), hessian_shape=tensor_shape.TensorShape([2, 2])) - with ops.control_dependencies([accumulator._create_op]): + with ops.control_dependencies([accumulator.initializer]): op1 = accumulator.add( stamp_token=0, partition_ids=[1, 2], @@ -349,7 +349,7 @@ class StatsAccumulatorTensorTest(test_util.TensorFlowTestCase): with ops.control_dependencies([op1]): (stamp_token, num_updates_1, partition_1, feature_1, grads_1, - hessians_1) = accumulator.serialize() + hessians_1) = accumulator.saveable.serialize() # Make sure that the accumulator hasn't changed during serialization. with ops.control_dependencies([stamp_token]): num_updates_2, partition_2, feature_2, grads_2, hessians_2 = ( @@ -386,7 +386,7 @@ class StatsAccumulatorTensorTest(test_util.TensorFlowTestCase): stamp_token=0, gradient_shape=tensor_shape.TensorShape([2]), hessian_shape=tensor_shape.TensorShape([2, 2])) - with ops.control_dependencies([accumulator._create_op]): + with ops.control_dependencies([accumulator.initializer]): # These will be deleted due to deserialize call. op1 = accumulator.add( stamp_token=0, @@ -399,7 +399,7 @@ class StatsAccumulatorTensorTest(test_util.TensorFlowTestCase): 0.08]]]) with ops.control_dependencies([op1]): - deserialize = accumulator.deserialize( + deserialize = accumulator.saveable.deserialize( stamp_token=2, num_updates=3, partition_ids=[3, 4], diff --git a/tensorflow/contrib/boosted_trees/python/ops/batch_ops_utils.py b/tensorflow/contrib/boosted_trees/python/ops/batch_ops_utils.py index 843420968ac6a6716fdf6b4967146e131139f67c..4dc764f95713ab788c282c2f3e7fb278a24f4822 100644 --- a/tensorflow/contrib/boosted_trees/python/ops/batch_ops_utils.py +++ b/tensorflow/contrib/boosted_trees/python/ops/batch_ops_utils.py @@ -20,6 +20,8 @@ from __future__ import print_function import abc import collections +import six + from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops @@ -27,11 +29,10 @@ from tensorflow.python.framework import tensor_shape from tensorflow.python.ops import array_ops +@six.add_metaclass(abc.ABCMeta) class ScheduledOp(object): """Represents a scheduled remote operation.""" - __metaclass__ = abc.ABCMeta - @abc.abstractmethod def batching_key(self): """Returns the key for batching operations.""" diff --git a/tensorflow/contrib/boosted_trees/python/ops/model_ops.py b/tensorflow/contrib/boosted_trees/python/ops/model_ops.py index 25b2c9e2fd72bd018717e8a87fce726f26bad968..c3685b54e201f73039f6623443c67ba2b217a51e 100644 --- a/tensorflow/contrib/boosted_trees/python/ops/model_ops.py +++ b/tensorflow/contrib/boosted_trees/python/ops/model_ops.py @@ -17,6 +17,8 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import functools + # pylint: disable=unused-import from tensorflow.contrib.boosted_trees.python.ops import boosted_trees_ops_loader # pylint: enable=unused-import @@ -31,6 +33,7 @@ from tensorflow.contrib.boosted_trees.python.ops.gen_model_ops import tree_ensem from tensorflow.python.framework import ops from tensorflow.python.ops import resources from tensorflow.python.training import saver +from tensorflow.python.training.checkpointable import tracking ops.NotDifferentiable("TreeEnsembleVariable") ops.NotDifferentiable("TreeEnsembleSerialize") @@ -59,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 @@ -82,6 +85,44 @@ class TreeEnsembleVariableSavable(saver.BaseSaverBuilder.SaveableObject): tree_ensemble_config=restored_tensors[1]) +class TreeEnsembleVariable(tracking.TrackableResource): + """A Tree ensemble model.""" + + def __init__(self, stamp_token, tree_ensemble_config, name, container=None): + self._stamp_token = stamp_token + self._tree_ensemble_config = tree_ensemble_config + self._name = name + self._container = container + self._init_op = None + super(TreeEnsembleVariable, self).__init__() + + def create_resource(self): + return gen_model_ops.decision_tree_ensemble_resource_handle_op( + self._container, shared_name=self._name, name=self._name) + + def initialize(self): + return gen_model_ops.create_tree_ensemble_variable( + self.resource_handle, self._stamp_token, self._tree_ensemble_config) + + @property + def initializer(self): + if self._init_op is None: + self._init_op = self.initialize() + return self._init_op + + def is_initialized(self): + return gen_model_ops.tree_ensemble_is_initialized_op(self.resource_handle) + + def _gather_saveables_for_checkpoint(self): + return { + self.resource_handle.op.name + "/tree_ensemble_variable": + functools.partial( + TreeEnsembleVariableSavable, + tree_ensemble_handle=self.resource_handle, + create_op=self.initializer) + } + + def tree_ensemble_variable(stamp_token, tree_ensemble_config, name, @@ -90,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 `""`. @@ -99,12 +140,11 @@ def tree_ensemble_variable(stamp_token, A `Tensor` of type mutable `string`. The handle to the tree ensemble. """ with ops.name_scope(name, "TreeEnsembleVariable") as name: - resource_handle = gen_model_ops.decision_tree_ensemble_resource_handle_op( - container, shared_name=name, name=name) - create_op = gen_model_ops.create_tree_ensemble_variable( - resource_handle, stamp_token, tree_ensemble_config) - is_initialized_op = gen_model_ops.tree_ensemble_is_initialized_op( - resource_handle) + tree_ensemble_var = TreeEnsembleVariable(stamp_token, tree_ensemble_config, + name, container) + resource_handle = tree_ensemble_var.resource_handle + create_op = tree_ensemble_var.initializer + is_initialized_op = tree_ensemble_var.is_initialized() # Adds the variable to the savable list. saveable = TreeEnsembleVariableSavable(resource_handle, create_op, resource_handle.name) diff --git a/tensorflow/contrib/boosted_trees/python/ops/quantile_ops.py b/tensorflow/contrib/boosted_trees/python/ops/quantile_ops.py index 19b6b3296db394b07f57a25dbde187eb9195af38..0c319cc9bd1f720eb404a9da05227c5807ec874f 100644 --- a/tensorflow/contrib/boosted_trees/python/ops/quantile_ops.py +++ b/tensorflow/contrib/boosted_trees/python/ops/quantile_ops.py @@ -33,59 +33,20 @@ from tensorflow.python.framework import ops from tensorflow.python.framework import sparse_tensor from tensorflow.python.ops import resources from tensorflow.python.training import saver +from tensorflow.python.training.checkpointable import tracking # Pattern to remove all non alpha numeric from a string. _PATTERN = re.compile(r"[\W_]+") -class QuantileAccumulator(saver.BaseSaverBuilder.SaveableObject): - """A resource that allows distributed quantile computation.""" - - def __init__(self, - init_stamp_token, - epsilon, - num_quantiles, - max_elements=None, - name=None, - container=None, - generate_quantiles=False): - """Creates a QuantileAccumulator object. - - Args: - init_stamp_token: The initial value for the stamp token. - epsilon: Error bound on the quantile computation. - num_quantiles: Number of quantiles to produce from the final summary. - max_elements: Maximum number of elements added to the accumulator. - name: the name to save the accumulator under. - container: An optional `string`. Defaults to `""` - generate_quantiles: Generate quantiles instead of approximate boundaries. - If true, exactly `num_quantiles` will be produced in the final summary. - """ - self._epsilon = epsilon - self._generate_quantiles = generate_quantiles +class QuantileAccumulatorSaveable(saver.BaseSaverBuilder.SaveableObject): + """SaveableObject implementation for QuantileAccumulator.""" - name = _PATTERN.sub("", name) - with ops.name_scope(name, "QuantileAccumulator") as name: - self._quantile_accumulator_handle = ( - gen_quantile_ops.quantile_stream_resource_handle_op( - container=container, shared_name=name, name=name)) - self._create_op = gen_quantile_ops.create_quantile_accumulator( - self._quantile_accumulator_handle, - init_stamp_token, - epsilon=epsilon, - max_elements=max_elements, - num_quantiles=num_quantiles, - generate_quantiles=generate_quantiles) - is_initialized_op = gen_quantile_ops.quantile_accumulator_is_initialized( - self._quantile_accumulator_handle) - resources.register_resource(self._quantile_accumulator_handle, - self._create_op, is_initialized_op) - self._make_savable(name) - - def _make_savable(self, name): + def __init__(self, resource_handle, create_op, name): + self._resource_handle = resource_handle + self._create_op = create_op stamp_token, state, are_buckets_ready, buckets = ( - gen_quantile_ops.quantile_accumulator_serialize( - self._quantile_accumulator_handle)) + gen_quantile_ops.quantile_accumulator_serialize(resource_handle)) # slice_spec is useful for saving a slice from a variable. # It's not meaningful in quantile accumulator. slice_spec = "" @@ -96,9 +57,8 @@ class QuantileAccumulator(saver.BaseSaverBuilder.SaveableObject): specs += [make_save_spec(state, "_state")] specs += [make_save_spec(are_buckets_ready, "_are_buckets_ready")] specs += [make_save_spec(buckets, "buckets")] - super(QuantileAccumulator, - self).__init__(self._quantile_accumulator_handle, specs, name) - ops.add_to_collection(ops.GraphKeys.SAVEABLE_OBJECTS, self) + super(QuantileAccumulatorSaveable, self).__init__(self._resource_handle, + specs, name) def restore(self, restored_tensors, unused_restored_shapes): """Restores the associated quantile accumulator from 'restored_tensors'. @@ -119,24 +79,94 @@ class QuantileAccumulator(saver.BaseSaverBuilder.SaveableObject): buckets = restored_tensors[3] with ops.control_dependencies([self._create_op]): return gen_quantile_ops.quantile_accumulator_deserialize( - self._quantile_accumulator_handle, + self._resource_handle, stamp_token=stamp_token, stream_state=state, are_buckets_ready=are_buckets_ready, buckets=buckets) + +class QuantileAccumulator(tracking.TrackableResource): + """A resource that allows distributed quantile computation.""" + + def __init__(self, + init_stamp_token, + epsilon, + num_quantiles, + max_elements=None, + name=None, + container=None, + generate_quantiles=False): + """Creates a QuantileAccumulator object. + + Args: + init_stamp_token: The initial value for the stamp token. + epsilon: Error bound on the quantile computation. + num_quantiles: Number of quantiles to produce from the final summary. + max_elements: Maximum number of elements added to the accumulator. + name: the name to save the accumulator under. + container: An optional `string`. Defaults to `""` + generate_quantiles: Generate quantiles instead of approximate boundaries. + If true, exactly `num_quantiles` will be produced in the final summary. + """ + self._init_stamp_token = init_stamp_token + self._epsilon = epsilon + self._num_quantiles = num_quantiles + self._max_elements = max_elements + self._container = container + self._generate_quantiles = generate_quantiles + super(QuantileAccumulator, self).__init__() + + name = _PATTERN.sub("", name) + with ops.name_scope(name, "QuantileAccumulator") as name: + self._name = name + self._resource_handle = self.create_resource() + self._init_op = self.initialize() + is_initialized_op = self.is_initialized() + resources.register_resource(self.resource_handle, self._init_op, + is_initialized_op) + self._saveable = QuantileAccumulatorSaveable(self.resource_handle, + self._init_op, name) + ops.add_to_collection(ops.GraphKeys.SAVEABLE_OBJECTS, self._saveable) + + def create_resource(self): + return gen_quantile_ops.quantile_stream_resource_handle_op( + container=self._container, shared_name=self._name, name=self._name) + + def initialize(self): + return gen_quantile_ops.create_quantile_accumulator( + self.resource_handle, + self._init_stamp_token, + epsilon=self._epsilon, + max_elements=self._max_elements, + num_quantiles=self._num_quantiles, + generate_quantiles=self._generate_quantiles) + + @property + def initializer(self): + if self._init_op is None: + self._init_op = self.initialize() + return self._init_op + + def is_initialized(self): + return gen_quantile_ops.quantile_accumulator_is_initialized( + self.resource_handle) + + def _gather_saveables_for_checkpoint(self): + return {"quantile_accumulator", self.saveable} + def get_buckets(self, stamp_token): """Returns quantile buckets created during previous flush.""" are_buckets_ready, buckets = ( gen_quantile_ops.quantile_accumulator_get_buckets( - quantile_accumulator_handles=[self._quantile_accumulator_handle], + quantile_accumulator_handles=[self.resource_handle], stamp_token=stamp_token)) return are_buckets_ready[0], buckets[0] def schedule_get_buckets(self): """Returns a scheduled read of buckets created during previous flush.""" return batch_ops_utils.ScheduledStampedResourceOp( - resource_handle=self._quantile_accumulator_handle, + resource_handle=self.resource_handle, op=gen_quantile_ops.quantile_accumulator_get_buckets) def _make_summary(self, column, example_weights): @@ -161,14 +191,14 @@ class QuantileAccumulator(saver.BaseSaverBuilder.SaveableObject): """Adds quantile summary to its stream in resource.""" summary = self._make_summary(column, example_weights) return gen_quantile_ops.quantile_accumulator_add_summaries( - quantile_accumulator_handles=[self._quantile_accumulator_handle], + quantile_accumulator_handles=[self.resource_handle], stamp_token=stamp_token, summaries=[summary]) def add_prebuilt_summary(self, stamp_token, summary): """Adds quantile summary to its stream in resource.""" return gen_quantile_ops.quantile_accumulator_add_summaries( - quantile_accumulator_handles=[self._quantile_accumulator_handle], + quantile_accumulator_handles=[self.resource_handle], stamp_token=stamp_token, summaries=[summary]) @@ -177,7 +207,7 @@ class QuantileAccumulator(saver.BaseSaverBuilder.SaveableObject): summary = self._make_summary(column, example_weights) return batch_ops_utils.ScheduledStampedResourceOp( op=gen_quantile_ops.quantile_accumulator_add_summaries, - resource_handle=self._quantile_accumulator_handle, + resource_handle=self.resource_handle, summaries=summary) def flush(self, stamp_token, next_stamp_token): @@ -190,17 +220,14 @@ class QuantileAccumulator(saver.BaseSaverBuilder.SaveableObject): The flush operation. """ return gen_quantile_ops.quantile_accumulator_flush( - quantile_accumulator_handle=self._quantile_accumulator_handle, + quantile_accumulator_handle=self.resource_handle, stamp_token=stamp_token, next_stamp_token=next_stamp_token) def flush_summary(self, stamp_token, next_stamp_token): """Finalizes quantile summary stream and resets it for next iteration.""" result = gen_quantile_ops.quantile_accumulator_flush_summary( - quantile_accumulator_handle=self._quantile_accumulator_handle, + quantile_accumulator_handle=self.resource_handle, stamp_token=stamp_token, next_stamp_token=next_stamp_token) return result - - def resource(self): - return self._quantile_accumulator_handle diff --git a/tensorflow/contrib/boosted_trees/python/ops/stats_accumulator_ops.py b/tensorflow/contrib/boosted_trees/python/ops/stats_accumulator_ops.py index 2e94e353f325f06eed2d290d3a7a461861820c39..ad1191d41236e71008bff8c8a7fbd42c16e3f9c5 100644 --- a/tensorflow/contrib/boosted_trees/python/ops/stats_accumulator_ops.py +++ b/tensorflow/contrib/boosted_trees/python/ops/stats_accumulator_ops.py @@ -26,12 +26,83 @@ from tensorflow.python.framework import ops from tensorflow.python.framework import tensor_shape from tensorflow.python.ops import resources from tensorflow.python.training import saver +from tensorflow.python.training.checkpointable import tracking # Pattern to remove all non alpha numeric from a string. _PATTERN = re.compile(r"[\W_]+") -class StatsAccumulator(saver.BaseSaverBuilder.SaveableObject): +class StatsAccumulatorSaveable(saver.BaseSaverBuilder.SaveableObject): + """SaveableObject implementation for StatsAccumulator.""" + + def __init__(self, resource_handle, create_op, is_scalar, name): + self._create_op = create_op + self._resource_handle = resource_handle + self._is_scalar = is_scalar + slice_spec = "" + saver_name = self._resource_handle.name + (stamp_token, num_updates, partition_ids, feature_ids, gradients, + hessians) = self.serialize() + specs = [ + saver.BaseSaverBuilder.SaveSpec(stamp_token, slice_spec, + saver_name + "_stamp"), + saver.BaseSaverBuilder.SaveSpec(num_updates, slice_spec, + saver_name + "_num_updates"), + saver.BaseSaverBuilder.SaveSpec(partition_ids, slice_spec, + saver_name + "_partition_ids"), + saver.BaseSaverBuilder.SaveSpec(feature_ids, slice_spec, + saver_name + "_feature_ids"), + saver.BaseSaverBuilder.SaveSpec(gradients, slice_spec, + saver_name + "_gradients"), + saver.BaseSaverBuilder.SaveSpec(hessians, slice_spec, + saver_name + "hessians"), + ] + super(StatsAccumulatorSaveable, self).__init__(self._resource_handle, specs, + name) + + def serialize(self): + """Serializes the stats accumulator state.""" + if self._is_scalar: + return gen_stats_accumulator_ops.stats_accumulator_scalar_serialize( + self._resource_handle) + else: + return gen_stats_accumulator_ops.stats_accumulator_tensor_serialize( + self._resource_handle) + + def deserialize(self, stamp_token, num_updates, partition_ids, feature_ids, + gradients, hessians): + """Resets the stats accumulator with the serialized state.""" + if self._is_scalar: + return gen_stats_accumulator_ops.stats_accumulator_scalar_deserialize( + self._resource_handle, stamp_token, num_updates, partition_ids, + feature_ids, gradients, hessians) + else: + return gen_stats_accumulator_ops.stats_accumulator_tensor_deserialize( + self._resource_handle, stamp_token, num_updates, partition_ids, + feature_ids, gradients, hessians) + + def restore(self, restored_tensors, unused_restored_shapes): + """Restores the associated tree ensemble from 'restored_tensors'. + + Args: + restored_tensors: the tensors that were loaded from a checkpoint. + unused_restored_shapes: the shapes this object should conform to after + restore. Not meaningful for trees. + + Returns: + The operation that restores the state of the tree ensemble variable. + """ + with ops.control_dependencies([self._create_op]): + return self.deserialize( + stamp_token=restored_tensors[0], + num_updates=restored_tensors[1], + partition_ids=restored_tensors[2], + feature_ids=restored_tensors[3], + gradients=restored_tensors[4], + hessians=restored_tensors[5]) + + +class StatsAccumulator(tracking.TrackableResource): """A resource that allows to accumulate gradients and hessians. For consistency guarantees, we use read and write stamp tokens. @@ -58,58 +129,69 @@ class StatsAccumulator(saver.BaseSaverBuilder.SaveableObject): Returns: A `Tensor` of type mutable `string`. The handle to the stats accumulator. """ + self._stamp_token = stamp_token + self._gradient_shape = gradient_shape + self._hessian_shape = hessian_shape + self._container = container + + if (gradient_shape == tensor_shape.scalar() and + hessian_shape == tensor_shape.scalar()): + self._is_scalar = True + else: + self._is_scalar = False + if name is not None: name = _PATTERN.sub("", name) with ops.name_scope(name, "StatsAccumulator") as name: - # Both values are scalars. - if (gradient_shape == tensor_shape.scalar() and - hessian_shape == tensor_shape.scalar()): - self._is_scalar = True - self._resource_handle = (gen_stats_accumulator_ops. - stats_accumulator_scalar_resource_handle_op( - container, name, name=name)) - - create_op = gen_stats_accumulator_ops.create_stats_accumulator_scalar( - self._resource_handle, stamp_token) - is_initialized_op = ( - gen_stats_accumulator_ops.stats_accumulator_scalar_is_initialized( - self._resource_handle)) - else: - self._is_scalar = False - self._resource_handle = (gen_stats_accumulator_ops. - stats_accumulator_tensor_resource_handle_op( - container, name, name=name)) - create_op = gen_stats_accumulator_ops.create_stats_accumulator_tensor( - self._resource_handle, stamp_token, gradient_shape.as_list(), - hessian_shape.as_list()) - is_initialized_op = ( - gen_stats_accumulator_ops.stats_accumulator_tensor_is_initialized( - self._resource_handle)) + self._name = name + self._resource_handle = self.create_resource() + self._init_op = self.initialize() + is_initialized_op = self.is_initialized() + resources.register_resource(self.resource_handle, self.initializer, + is_initialized_op) + self._saveable = StatsAccumulatorSaveable( + self.resource_handle, self.initializer, self._is_scalar, name) + ops.add_to_collection(ops.GraphKeys.SAVEABLE_OBJECTS, self._saveable) - self._create_op = create_op - slice_spec = "" - saver_name = self._resource_handle.name - (stamp_token, num_updates, partition_ids, feature_ids, gradients, - hessians) = self.serialize() - specs = [ - saver.BaseSaverBuilder.SaveSpec(stamp_token, slice_spec, - saver_name + "_stamp"), - saver.BaseSaverBuilder.SaveSpec(num_updates, slice_spec, - saver_name + "_num_updates"), - saver.BaseSaverBuilder.SaveSpec(partition_ids, slice_spec, - saver_name + "_partition_ids"), - saver.BaseSaverBuilder.SaveSpec(feature_ids, slice_spec, - saver_name + "_feature_ids"), - saver.BaseSaverBuilder.SaveSpec(gradients, slice_spec, - saver_name + "_gradients"), - saver.BaseSaverBuilder.SaveSpec(hessians, slice_spec, - saver_name + "hessians"), - ] + def create_resource(self): + if self._is_scalar: + return ( + gen_stats_accumulator_ops.stats_accumulator_scalar_resource_handle_op( + self._container, self._name, name=self._name)) + else: + return ( + gen_stats_accumulator_ops.stats_accumulator_tensor_resource_handle_op( + self._container, self._name, name=self._name)) - super(StatsAccumulator, self).__init__(self._resource_handle, specs, name) - resources.register_resource(self._resource_handle, create_op, - is_initialized_op) - ops.add_to_collection(ops.GraphKeys.SAVEABLE_OBJECTS, self) + def initialize(self): + if self._is_scalar: + return gen_stats_accumulator_ops.create_stats_accumulator_scalar( + self.resource_handle, self._stamp_token) + else: + return gen_stats_accumulator_ops.create_stats_accumulator_tensor( + self.resource_handle, self._stamp_token, + self._gradient_shape.as_list(), self._hessian_shape.as_list()) + + @property + def initializer(self): + if self._init_op is None: + self._init_op = self.initialize() + return self._init_op + + def is_initialized(self): + if self._is_scalar: + return gen_stats_accumulator_ops.stats_accumulator_scalar_is_initialized( + self.resource_handle) + else: + return gen_stats_accumulator_ops.stats_accumulator_tensor_is_initialized( + self.resource_handle) + + @property + def saveable(self): + return self._saveable + + def _gather_saveables_for_checkpoint(self): + return {"stats_accumulator", self.saveable} def add(self, stamp_token, partition_ids, feature_ids, gradients, hessians): """Updates the stats accumulator.""" @@ -117,11 +199,11 @@ class StatsAccumulator(saver.BaseSaverBuilder.SaveableObject): partition_ids, feature_ids, gradients, hessians)) if self._is_scalar: return gen_stats_accumulator_ops.stats_accumulator_scalar_add( - [self._resource_handle], stamp_token, [partition_ids], [feature_ids], + [self.resource_handle], stamp_token, [partition_ids], [feature_ids], [gradients], [hessians]) else: return gen_stats_accumulator_ops.stats_accumulator_tensor_add( - [self._resource_handle], stamp_token, [partition_ids], [feature_ids], + [self.resource_handle], stamp_token, [partition_ids], [feature_ids], [gradients], [hessians]) def schedule_add(self, partition_ids, feature_ids, gradients, hessians): @@ -131,7 +213,7 @@ class StatsAccumulator(saver.BaseSaverBuilder.SaveableObject): if self._is_scalar: return batch_ops_utils.ScheduledStampedResourceOp( op=gen_stats_accumulator_ops.stats_accumulator_scalar_add, - resource_handle=self._resource_handle, + resource_handle=self.resource_handle, partition_ids=partition_ids, feature_ids=feature_ids, gradients=gradients, @@ -139,7 +221,7 @@ class StatsAccumulator(saver.BaseSaverBuilder.SaveableObject): else: return batch_ops_utils.ScheduledStampedResourceOp( op=gen_stats_accumulator_ops.stats_accumulator_tensor_add, - resource_handle=self._resource_handle, + resource_handle=self.resource_handle, partition_ids=partition_ids, feature_ids=feature_ids, gradients=gradients, @@ -153,55 +235,11 @@ class StatsAccumulator(saver.BaseSaverBuilder.SaveableObject): return gen_stats_accumulator_ops.stats_accumulator_tensor_make_summary( partition_ids, feature_ids, gradients, hessians) - def deserialize(self, stamp_token, num_updates, partition_ids, feature_ids, - gradients, hessians): - """Resets the stats accumulator with the serialized state.""" - if self._is_scalar: - return gen_stats_accumulator_ops.stats_accumulator_scalar_deserialize( - self._resource_handle, stamp_token, num_updates, partition_ids, - feature_ids, gradients, hessians) - else: - return gen_stats_accumulator_ops.stats_accumulator_tensor_deserialize( - self._resource_handle, stamp_token, num_updates, partition_ids, - feature_ids, gradients, hessians) - def flush(self, stamp_token, next_stamp_token): """Flushes the stats accumulator.""" if self._is_scalar: return gen_stats_accumulator_ops.stats_accumulator_scalar_flush( - self._resource_handle, stamp_token, next_stamp_token) + self.resource_handle, stamp_token, next_stamp_token) else: return gen_stats_accumulator_ops.stats_accumulator_tensor_flush( - self._resource_handle, stamp_token, next_stamp_token) - - def serialize(self): - """Serializes the stats accumulator state.""" - if self._is_scalar: - return gen_stats_accumulator_ops.stats_accumulator_scalar_serialize( - self._resource_handle) - else: - return gen_stats_accumulator_ops.stats_accumulator_tensor_serialize( - self._resource_handle) - - def restore(self, restored_tensors, unused_restored_shapes): - """Restores the associated tree ensemble from 'restored_tensors'. - - Args: - restored_tensors: the tensors that were loaded from a checkpoint. - unused_restored_shapes: the shapes this object should conform to after - restore. Not meaningful for trees. - - Returns: - The operation that restores the state of the tree ensemble variable. - """ - with ops.control_dependencies([self._create_op]): - return self.deserialize( - stamp_token=restored_tensors[0], - num_updates=restored_tensors[1], - partition_ids=restored_tensors[2], - feature_ids=restored_tensors[3], - gradients=restored_tensors[4], - hessians=restored_tensors[5]) - - def resource(self): - return self._resource_handle + self.resource_handle, stamp_token, next_stamp_token) 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 bd5d5bb695684c7dcb5cc5c0038386074edd4dc4..e78ec476ab3b43e5eb56a2502008bb8020ae97e0 100644 --- a/tensorflow/contrib/boosted_trees/python/training/functions/gbdt_batch.py +++ b/tensorflow/contrib/boosted_trees/python/training/functions/gbdt_batch.py @@ -386,10 +386,21 @@ class GradientBoostedDecisionTreeModel(object): learner_pb2.LearnerConfig.GROWING_MODE_UNSPECIFIED): learner_config.growing_mode = learner_pb2.LearnerConfig.LAYER_BY_LAYER + if (learner_config.weak_learner_type == learner_pb2.LearnerConfig + .OBLIVIOUS_DECISION_TREE and learner_config.pruning_mode == learner_pb2 + .LearnerConfig.PRUNING_MODE_UNSPECIFIED): + learner_config.pruning_mode = learner_pb2.LearnerConfig.PRE_PRUNE + if (learner_config.pruning_mode == learner_pb2.LearnerConfig.PRUNING_MODE_UNSPECIFIED): learner_config.pruning_mode = learner_pb2.LearnerConfig.POST_PRUNE + if (learner_config.weak_learner_type == learner_pb2.LearnerConfig + .OBLIVIOUS_DECISION_TREE and + learner_config.pruning_mode == learner_pb2.LearnerConfig.POST_PRUNE): + raise ValueError( + "Post pruning is not implmented for oblivious decision trees.") + if learner_config.constraints.max_tree_depth == 0: # Use 6 as the default maximum depth. learner_config.constraints.max_tree_depth = 6 @@ -418,6 +429,11 @@ class GradientBoostedDecisionTreeModel(object): sparse_float_shapes, sparse_int_indices, sparse_int_values, sparse_int_shapes) = extract_features( features, self._feature_columns, use_core_columns) + if (learner_config.weak_learner_type == learner_pb2.LearnerConfig + .OBLIVIOUS_DECISION_TREE and sparse_float_indices): + raise ValueError("Oblivious trees don't handle sparse float features yet." + ) + logging.info("Active Feature Columns: " + str(fc_names)) logging.info("Learner config: " + str(learner_config)) self._fc_names = fc_names @@ -550,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]): @@ -598,13 +615,19 @@ class GradientBoostedDecisionTreeModel(object): predictions_dict[NUM_TREES_ATTEMPTED] % self._logits_dimension) return constant_op.constant(-1, dtype=dtypes.int32) - def update_stats(self, loss, predictions_dict): + def update_stats(self, loss, predictions_dict, gradients=None, hessians=None): """Update the accumulators with stats from this batch. Args: loss: A scalar tensor representing average loss of examples. predictions_dict: Dictionary of Rank 2 `Tensor` representing information about predictions per example. + gradients: A tensor with the gradients with the respect to logits from + predictions_dict. If not provided, tensorflow will do + autodifferentiation. + hessians: A tensor with the hessians with the respect to logits from + predictions_dict. If not provided, tensorflow will do + autodifferentiation. Returns: Three values: @@ -626,13 +649,14 @@ class GradientBoostedDecisionTreeModel(object): predictions = predictions_dict[PREDICTIONS] partition_ids = predictions_dict[PARTITION_IDS] ensemble_stamp = predictions_dict[ENSEMBLE_STAMP] - gradients = gradients_impl.gradients( - loss, - predictions, - name="Gradients", - colocate_gradients_with_ops=False, - gate_gradients=0, - aggregation_method=None)[0] + if gradients is None: + gradients = gradients_impl.gradients( + loss, + predictions, + name="Gradients", + colocate_gradients_with_ops=False, + gate_gradients=0, + aggregation_method=None)[0] strategy = self._learner_config.multi_class_strategy class_id = self._get_class_id(predictions_dict) @@ -641,17 +665,20 @@ class GradientBoostedDecisionTreeModel(object): # We build one vs rest trees. if self._logits_dimension == 1: # We have only 1 score, gradients is of shape [batch, 1]. - hessians = gradients_impl.gradients( - gradients, - predictions, - name="Hessian", - colocate_gradients_with_ops=False, - gate_gradients=0, - aggregation_method=None)[0] + if hessians is None: + hessians = gradients_impl.gradients( + gradients, + predictions, + name="Hessian", + colocate_gradients_with_ops=False, + gate_gradients=0, + aggregation_method=None)[0] squeezed_gradients = array_ops.squeeze(gradients, axis=[1]) squeezed_hessians = array_ops.squeeze(hessians, axis=[1]) else: + if hessians is not None: + raise ValueError("Providing hessians is not yet supported here.") hessian_list = self._diagonal_hessian(gradients, predictions) # Assemble hessian list into a tensor. hessians = array_ops.stack(hessian_list, axis=1) @@ -662,6 +689,8 @@ class GradientBoostedDecisionTreeModel(object): squeezed_hessians = array_ops.squeeze( _get_column_by_index(hessians, class_id)) else: + if hessians is not None: + raise ValueError("Providing hessians is not yet supported here.") # Other multiclass strategies. if strategy == learner_pb2.LearnerConfig.FULL_HESSIAN: hessian_list = self._full_hessian(gradients, predictions) @@ -819,9 +848,9 @@ class GradientBoostedDecisionTreeModel(object): stats_update_ops.append( control_flow_ops.cond( continue_centering, - self._make_update_bias_stats_fn( - ensemble_stamp, predictions, gradients, - bias_stats_accumulator), control_flow_ops.no_op)) + self._make_update_bias_stats_fn(ensemble_stamp, predictions, + gradients, bias_stats_accumulator, + hessians), control_flow_ops.no_op)) # Update handler stats. handler_reads = collections.OrderedDict() @@ -881,9 +910,9 @@ class GradientBoostedDecisionTreeModel(object): empty_hess_shape = [1] + self._hessian_shape.as_list() empty_grad_shape = [1] + self._gradient_shape.as_list() - empty_gradients = constant_op.constant( + empty_gradients = constant_op.constant_v1( [], dtype=dtypes.float32, shape=empty_grad_shape) - empty_hessians = constant_op.constant( + empty_hessians = constant_op.constant_v1( [], dtype=dtypes.float32, shape=empty_hess_shape) active_handlers = array_ops.unstack(active_handlers, axis=0) @@ -976,7 +1005,7 @@ class GradientBoostedDecisionTreeModel(object): # Get accumulated steps and examples for the current layer. _, _, _, _, acc_examples, acc_steps = ( - steps_accumulator.serialize()) + steps_accumulator.saveable.serialize()) acc_examples = math_ops.cast(acc_examples[0], dtypes.int64) acc_steps = math_ops.cast(acc_steps[0], dtypes.int64) ensemble_update_ops.append( @@ -1146,7 +1175,8 @@ class GradientBoostedDecisionTreeModel(object): def get_max_tree_depth(self): return self._max_tree_depth - def train(self, loss, predictions_dict, labels): + def train(self, loss, predictions_dict, labels, gradients=None, + hessians=None): """Updates the accumalator stats and grows the ensemble. Args: @@ -1155,6 +1185,12 @@ class GradientBoostedDecisionTreeModel(object): about predictions per example. labels: Rank 2 `Tensor` representing labels per example. Has no effect on the training and is only kept for backward compatibility. + gradients: A tensor with the gradients with the respect to logits from + predictions_dict. If not provided, tensorflow will do + autodifferentiation. + hessians: A tensor with the hessians with the respect to logits from + predictions_dict. If not provided, tensorflow will do + autodifferentiation. Returns: An op that adds a new tree to the ensemble. @@ -1163,7 +1199,8 @@ class GradientBoostedDecisionTreeModel(object): ValueError: if inputs are not valid. """ del labels # unused; kept for backward compatibility. - update_op, _, training_state = self.update_stats(loss, predictions_dict) + update_op, _, training_state = self.update_stats(loss, predictions_dict, + gradients, hessians) with ops.control_dependencies(update_op): return self.increment_step_counter_and_maybe_update_ensemble( predictions_dict, training_state) @@ -1241,13 +1278,12 @@ class GradientBoostedDecisionTreeModel(object): def _get_replica_device_setter(self, worker_device): """Creates a replica device setter.""" ps_tasks = self._num_ps_replicas - ps_ops = [ - "Variable", - "VariableV2", + ps_ops = list(device_setter.STANDARD_PS_OPS) + ps_ops.extend([ "DecisionTreeEnsembleResourceHandleOp", "StatsAccumulatorScalarResourceHandleOp", "StatsAccumulatorTensorResourceHandleOp", - ] + ]) ps_strategy = _OpRoundRobinStrategy(ps_ops, ps_tasks) return device_setter.replica_device_setter( worker_device=worker_device, @@ -1256,21 +1292,28 @@ class GradientBoostedDecisionTreeModel(object): ps_ops=ps_ops, ps_strategy=ps_strategy) - def _make_update_bias_stats_fn(self, ensemble_stamp, predictions, gradients, - bias_stats_accumulator): + def _make_update_bias_stats_fn(self, + ensemble_stamp, + predictions, + gradients, + bias_stats_accumulator, + hessians=None): """A method to create the function which updates the bias stats.""" def _update_bias_stats(): """A method to update the bias stats.""" # Get reduced gradients and hessians. grads_sum = math_ops.reduce_sum(gradients, 0) - hess = gradients_impl.gradients( - grads_sum, - predictions, - name="Hessians", - colocate_gradients_with_ops=False, - gate_gradients=0, - aggregation_method=None)[0] + if hessians is not None: + hess = hessians + else: + hess = gradients_impl.gradients( + grads_sum, + predictions, + name="Hessians", + colocate_gradients_with_ops=False, + gate_gradients=0, + aggregation_method=None)[0] hess_sum = math_ops.reduce_sum(hess, 0) # Accumulate gradients and hessians. 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 6d20a2e7f482953481fb1effe4c6e2e5a300786f..92068e88a76cb8bfdd394c1093347a8fb8a63449 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 @@ -1257,6 +1257,96 @@ class GbdtTest(test_util.TensorFlowTestCase): self.assertArrayNear(expected_leaf_2, output.trees[0].nodes[2].leaf.vector.value, 1e-3) + def testTrainFnMulticlassDiagonalHessianOblivious(self): + """Tests the GBDT train for multiclass diagonal hessian.""" + with self.cached_session(): + ensemble_handle = model_ops.tree_ensemble_variable( + stamp_token=0, tree_ensemble_config="", name="tree_ensemble") + + learner_config = learner_pb2.LearnerConfig() + learner_config.learning_rate_tuner.fixed.learning_rate = 1 + # Use full hessian multiclass strategy. + learner_config.multi_class_strategy = ( + learner_pb2.LearnerConfig.DIAGONAL_HESSIAN) + learner_config.num_classes = 5 + learner_config.regularization.l1 = 0 + # To make matrix inversible. + learner_config.regularization.l2 = 1e-5 + learner_config.weak_learner_type = ( + learner_pb2.LearnerConfig.OBLIVIOUS_DECISION_TREE) + learner_config.pruning_mode = learner_pb2.LearnerConfig.PRE_PRUNE + learner_config.constraints.max_tree_depth = 5 + learner_config.constraints.min_node_weight = 0 + batch_size = 3 + features = {} + features["sparse_int"] = sparse_tensor.SparseTensor( + array_ops.constant([[0, 0], [1, 0]], dtypes.int64), + array_ops.constant([1, 2], dtypes.int64), + array_ops.constant([3, 1], dtypes.int64)) + + gbdt_model = gbdt_batch.GradientBoostedDecisionTreeModel( + is_chief=True, + num_ps_replicas=0, + center_bias=False, + ensemble_handle=ensemble_handle, + examples_per_layer=1, + learner_config=learner_config, + logits_dimension=5, + features=features) + + labels = array_ops.constant([[2], [2], [3]], dtype=dtypes.float32) + weights = array_ops.ones([batch_size, 1], dtypes.float32) + + predictions_dict = gbdt_model.predict(learn.ModeKeys.TRAIN) + predictions = predictions_dict["predictions"] + + # Create train op. + train_op = gbdt_model.train( + loss=math_ops.reduce_mean( + losses.per_example_maxent_loss( + labels, + weights, + predictions, + num_classes=learner_config.num_classes)[0]), + predictions_dict=predictions_dict, + labels=labels) + variables.global_variables_initializer().run() + resources.initialize_resources(resources.shared_resources()).run() + + stamp_token, serialized = model_ops.tree_ensemble_serialize( + ensemble_handle) + + # Grow 2 layers. + train_op.run() + train_op.run() + + output = tree_config_pb2.DecisionTreeEnsembleConfig() + output.ParseFromString(serialized.eval()) + + stamp_token, serialized = model_ops.tree_ensemble_serialize( + ensemble_handle) + output.ParseFromString(serialized.eval()) + self.assertEqual(len(output.trees), 1) + # We got 6 nodes: one parent and 4 leafs. + self.assertEqual(len(output.trees[0].nodes), 6) + self.assertAllClose(output.tree_weights, [1]) + self.assertEqual(stamp_token.eval(), 2) + + print(output.trees[0]) + # Leafs should have a dense vector of size 5. + expected_leaf_1 = [-1.2497, -1.24976, 4.999, -1.24976, -1.2497] + expected_leaf_2 = [-2.2362, -2.2362, 6.0028, -2.2362, -2.2362] + expected_leaf_3 = [-2.2694, -2.2694, 4.0064, -0.0084, -2.2694] + expected_leaf_4 = [-2.2694, -2.2694, -0.0084, 4.0064, -2.2694] + self.assertArrayNear(expected_leaf_1, + output.trees[0].nodes[2].leaf.vector.value, 1e-3) + self.assertArrayNear(expected_leaf_2, + output.trees[0].nodes[3].leaf.vector.value, 1e-3) + self.assertArrayNear(expected_leaf_3, + output.trees[0].nodes[4].leaf.vector.value, 1e-3) + self.assertArrayNear(expected_leaf_4, + output.trees[0].nodes[5].leaf.vector.value, 1e-3) + def testTrainFnMulticlassTreePerClass(self): """Tests the GBDT train for multiclass tree per class strategy.""" with self.cached_session() as sess: diff --git a/tensorflow/contrib/boosted_trees/python/utils/losses.py b/tensorflow/contrib/boosted_trees/python/utils/losses.py index b5ebaf1999519f65110e8164fa20bace5ecc3ef6..220e981618b7c0bfb1e4e98c087d83b451b9b3cf 100644 --- a/tensorflow/contrib/boosted_trees/python/utils/losses.py +++ b/tensorflow/contrib/boosted_trees/python/utils/losses.py @@ -48,6 +48,47 @@ def per_example_logistic_loss(labels, weights, predictions): labels=labels, logits=predictions) return unweighted_loss * weights, control_flow_ops.no_op() +# MUST USE WITH HESSIAN REGULARIZATION, +# This loss can have zero hessian, so it must be used with l2 or min_node_weight +# regularization. +# An example config is +# learner_config.constraints.min_node_weight = 1 / num_examples_per_layer +# learner_config.regularization.l2 = 1.0 / num_examples_per_layer +# TODO(nponomareva): make it multidimensional so we can estimate several +# quantiles at once. +def per_example_quantile_regression_loss(labels, weights, predictions, + quantile): + """Smoothed loss for quantile regression. + + The standard quantile regression loss is quantile*(y-y') when y>y' and + (quantile-1)*(y-y') otherwise, y' is a prediction, y is a label. The impl + below is this loss but squared in the region where the loss value < 1. + + Args: + labels: Rank 2 (N, D) tensor of per-example labels. + weights: Rank 2 (N, 1) tensor of per-example weights. + predictions: Rank 2 (N, D) tensor of per-example predictions. + quantile: The quantile to use. + + Returns: + loss: A Rank 2 (N, 1) tensor of per-example quantile loss. + update_op: An update operation to update the loss's internal state. + """ + labels = math_ops.to_float(labels) + error = labels - predictions + square_loss_right = array_ops.where(error * quantile < 1.0, + math_ops.square(quantile * error), + quantile * error) + square_loss_left = array_ops.where(error * (quantile - 1) < 1, + math_ops.square((quantile - 1) * error), + (quantile - 1) * error) + + unweighted_loss = array_ops.where(error > 0, square_loss_right, + square_loss_left) + if weights is None: + return unweighted_loss, control_flow_ops.no_op() + else: + return unweighted_loss * weights, control_flow_ops.no_op() # This is classical form of Maximum entropy loss, that is twice differentiable # (sparse_softmax_cross_entropy which is what we go for is not twice @@ -78,8 +119,7 @@ def per_example_maxent_loss(labels, weights, logits, num_classes, eps=1e-15): labels = array_ops.expand_dims(labels, 1) # Labels are indices of classes, convert them to one hot encodings. target_one_hot = array_ops.one_hot(indices=labels, depth=num_classes) - labels = math_ops.reduce_sum( - input_tensor=target_one_hot, reduction_indices=[1]) + labels = math_ops.reduce_sum(input_tensor=target_one_hot, axis=[1]) labels = math_ops.to_float(labels) # Calculate softmax probabilities for each class. 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/BUILD b/tensorflow/contrib/checkpoint/python/BUILD index ada41687261ab63286933d01da4e286173042e0c..4e529322c7c76797938468b405cd175609dc0a73 100644 --- a/tensorflow/contrib/checkpoint/python/BUILD +++ b/tensorflow/contrib/checkpoint/python/BUILD @@ -2,7 +2,7 @@ licenses(["notice"]) # Apache 2.0 package(default_visibility = ["//tensorflow:internal"]) -load("//tensorflow:tensorflow.bzl", "py_test") +load("//tensorflow:tensorflow.bzl", "tf_py_test") py_library( name = "checkpoint", @@ -27,17 +27,17 @@ py_library( ], ) -py_test( +tf_py_test( name = "containers_test", srcs = ["containers_test.py"], - deps = [ + additional_deps = [ ":containers", + "@six_archive//:six", "//tensorflow/python:client_testlib", "//tensorflow/python:framework_test_lib", "//tensorflow/python:resource_variable_ops", "//tensorflow/python/training/checkpointable:base", "//tensorflow/python/training/checkpointable:util", - "@six_archive//:six", ], ) @@ -53,18 +53,18 @@ py_library( ], ) -py_test( +tf_py_test( name = "python_state_test", srcs = ["python_state_test.py"], - deps = [ + additional_deps = [ ":python_state", + "//third_party/py/numpy", "//tensorflow/python:framework_ops", "//tensorflow/python:framework_test_lib", "//tensorflow/python:session", "//tensorflow/python:variables", "//tensorflow/python/eager:test", "//tensorflow/python/training/checkpointable:util", - "//third_party/py/numpy", ], ) @@ -80,10 +80,10 @@ py_library( ], ) -py_test( +tf_py_test( name = "split_dependency_test", srcs = ["split_dependency_test.py"], - deps = [ + additional_deps = [ ":split_dependency", "//tensorflow/python:array_ops", "//tensorflow/python:framework_test_lib", @@ -106,10 +106,10 @@ py_library( ], ) -py_test( +tf_py_test( name = "visualize_test", srcs = ["visualize_test.py"], - deps = [ + additional_deps = [ ":visualize", "//tensorflow/python:constant_op", "//tensorflow/python:resource_variable_ops", diff --git a/tensorflow/contrib/checkpoint/python/containers.py b/tensorflow/contrib/checkpoint/python/containers.py index 242c1e8ba45e0b2f6f9a1a51695b824546382666..97936d9e9dfd5d6e62fdf8312707a276b63e1267 100644 --- a/tensorflow/contrib/checkpoint/python/containers.py +++ b/tensorflow/contrib/checkpoint/python/containers.py @@ -46,6 +46,10 @@ class UniqueNameTracker(data_structures.CheckpointableDataStructure): self._maybe_initialize_checkpointable() self._name_counts = {} + @property + def _values(self): + return [dep.ref for dep in self._checkpoint_dependencies] + def track(self, checkpointable, base_name): """Add a dependency on `checkpointable`. @@ -59,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/cloud/kernels/BUILD b/tensorflow/contrib/cloud/kernels/BUILD index 1311063ec023bdaa2588d6f1c826bf900f7dea09..20f8c2b2453a58fdbe5a3587fa6687debd9c06d3 100644 --- a/tensorflow/contrib/cloud/kernels/BUILD +++ b/tensorflow/contrib/cloud/kernels/BUILD @@ -27,7 +27,6 @@ tf_kernel_library( deps = [ ":bigquery_table_accessor", ":bigquery_table_partition_proto_cc", - "//tensorflow/contrib/cloud:bigquery_reader_ops_op_lib", "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:reader_base", @@ -79,7 +78,6 @@ tf_kernel_library( srcs = ["gcs_config_ops.cc"], visibility = ["//tensorflow:internal"], deps = [ - "//tensorflow/contrib/cloud:gcs_config_ops_op_lib", "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core/platform/cloud:curl_http_request", diff --git a/tensorflow/contrib/cluster_resolver/BUILD b/tensorflow/contrib/cluster_resolver/BUILD index 7b59f9774e239919f088c616de9df6ba50802aae..f944b7f88438ff257a44581170ead16640540e69 100644 --- a/tensorflow/contrib/cluster_resolver/BUILD +++ b/tensorflow/contrib/cluster_resolver/BUILD @@ -21,118 +21,25 @@ py_library( py_library( name = "cluster_resolver_py", - srcs = [ + srcs = glob([ "__init__.py", - "python/training/__init__.py", - ], + "python/training/*.py", + ]), srcs_version = "PY2AND3", visibility = ["//visibility:public"], - deps = [ - ":base_cluster_resolver_py", - ":gce_cluster_resolver_py", - ":slurm_cluster_resolver_py", - ":tpu_cluster_resolver_py", - "//tensorflow/python:util", - ], -) - -py_library( - name = "base_cluster_resolver_py", - srcs = ["python/training/cluster_resolver.py"], - srcs_version = "PY2AND3", - deps = [ - "//tensorflow/python:training", - ], -) - -py_library( - name = "gce_cluster_resolver_py", - srcs = ["python/training/gce_cluster_resolver.py"], - srcs_version = "PY2AND3", - deps = [ - ":base_cluster_resolver_py", - "//tensorflow/python:training", - ], -) - -py_library( - name = "tpu_cluster_resolver_py", - srcs = ["python/training/tpu_cluster_resolver.py"], - srcs_version = "PY2AND3", - deps = [ - ":base_cluster_resolver_py", - "//tensorflow/python:training", - ], -) - -py_library( - name = "slurm_cluster_resolver_py", - srcs = ["python/training/slurm_cluster_resolver.py"], - srcs_version = "PY2AND3", - deps = [ - ":base_cluster_resolver_py", - "//tensorflow/python:training", - ], -) - -tf_py_test( - name = "base_cluster_resolver_py_test", - srcs = ["python/training/cluster_resolver_test.py"], - additional_deps = [ - ":cluster_resolver_py", - "//tensorflow/python:client_testlib", - "//tensorflow/python:framework_for_generated_wrappers", - "//tensorflow/python:framework_test_lib", - "//tensorflow/python:platform_test", - "//tensorflow/python:training", - ], - main = "python/training/cluster_resolver_test.py", -) - -tf_py_test( - name = "gce_cluster_resolver_py_test", - size = "small", - srcs = ["python/training/gce_cluster_resolver_test.py"], - additional_deps = [ - ":cluster_resolver_py", - ":gce_cluster_resolver_py", - "//tensorflow/python:client_testlib", - "//tensorflow/python:framework_for_generated_wrappers", - "//tensorflow/python:framework_test_lib", - "//tensorflow/python:platform_test", - "//tensorflow/python:training", - ], - main = "python/training/gce_cluster_resolver_test.py", -) - -tf_py_test( - name = "tpu_cluster_resolver_py_test", - size = "small", - srcs = ["python/training/tpu_cluster_resolver_test.py"], - additional_deps = [ - ":tpu_cluster_resolver_py", - "//tensorflow/python:client_testlib", - "//tensorflow/python:framework_for_generated_wrappers", - "//tensorflow/python:framework_test_lib", - "//tensorflow/python:platform_test", - "//tensorflow/python:training", - ], - grpc_enabled = True, - main = "python/training/tpu_cluster_resolver_test.py", + deps = ["//tensorflow/python/distribute/cluster_resolver:cluster_resolver_lib"], ) tf_py_test( - name = "slurm_cluster_resolver_py_test", - size = "small", - srcs = ["python/training/slurm_cluster_resolver_test.py"], + name = "cluster_resolver_initialization_test", + srcs = ["cluster_resolver_initialization_test.py"], additional_deps = [ ":cluster_resolver_py", - ":slurm_cluster_resolver_py", "//tensorflow/python:client_testlib", "//tensorflow/python:framework_for_generated_wrappers", "//tensorflow/python:framework_test_lib", "//tensorflow/python:platform_test", "//tensorflow/python:training", ], - main = "python/training/slurm_cluster_resolver_test.py", + main = "cluster_resolver_initialization_test.py", ) diff --git a/tensorflow/contrib/cluster_resolver/__init__.py b/tensorflow/contrib/cluster_resolver/__init__.py index fd1263fe81ae826d5edfa8752460fb78fe52b32a..390b3e7550b3d991269bb84707c3500f2fa33290 100644 --- a/tensorflow/contrib/cluster_resolver/__init__.py +++ b/tensorflow/contrib/cluster_resolver/__init__.py @@ -20,12 +20,14 @@ from __future__ import division from __future__ import print_function # pylint: disable=wildcard-import,unused-import -from tensorflow.contrib.cluster_resolver.python.training.cluster_resolver import ClusterResolver -from tensorflow.contrib.cluster_resolver.python.training.cluster_resolver import SimpleClusterResolver -from tensorflow.contrib.cluster_resolver.python.training.cluster_resolver import UnionClusterResolver -from tensorflow.contrib.cluster_resolver.python.training.gce_cluster_resolver import GceClusterResolver -from tensorflow.contrib.cluster_resolver.python.training.slurm_cluster_resolver import SlurmClusterResolver -from tensorflow.contrib.cluster_resolver.python.training.tpu_cluster_resolver import TPUClusterResolver +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.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 +from tensorflow.python.distribute.cluster_resolver.tpu_cluster_resolver import TPUClusterResolver # pylint: enable=wildcard-import,unused-import from tensorflow.python.util.all_util import remove_undocumented @@ -35,6 +37,8 @@ _allowed_symbols = [ 'SimpleClusterResolver', 'UnionClusterResolver', 'GceClusterResolver', + 'KubernetesClusterResolver', + 'TFConfigClusterResolver', 'TPUClusterResolver', 'SlurmClusterResolver', ] diff --git a/tensorflow/contrib/cluster_resolver/cluster_resolver_initialization_test.py b/tensorflow/contrib/cluster_resolver/cluster_resolver_initialization_test.py new file mode 100644 index 0000000000000000000000000000000000000000..01ff1478c694cf0901aeed48b6e0f873d8abe65e --- /dev/null +++ b/tensorflow/contrib/cluster_resolver/cluster_resolver_initialization_test.py @@ -0,0 +1,53 @@ +# 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 to ensure ClusterResolvers are usable via the old contrib path.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.cluster_resolver import SimpleClusterResolver +from tensorflow.contrib.cluster_resolver.python.training import cluster_resolver +from tensorflow.contrib.cluster_resolver.python.training import UnionClusterResolver +from tensorflow.python.platform import test +from tensorflow.python.training import server_lib + + +class ClusterResolverInitializationTest(test.TestCase): + + def testCreateSimpleClusterResolverFromLib(self): + base_cluster_spec = server_lib.ClusterSpec({ + "ps": ["ps0:2222", "ps1:2222"], + "worker": ["worker0:2222", "worker1:2222", "worker2:2222"] + }) + cluster_resolver.SimpleClusterResolver(base_cluster_spec) + + def testCreateSimpleClusterResolver(self): + base_cluster_spec = server_lib.ClusterSpec({ + "ps": ["ps0:2222", "ps1:2222"], + "worker": ["worker0:2222", "worker1:2222", "worker2:2222"] + }) + SimpleClusterResolver(base_cluster_spec) + + def testCreateUnionClusterResolver(self): + base_cluster_spec = server_lib.ClusterSpec({ + "ps": ["ps0:2222", "ps1:2222"], + "worker": ["worker0:2222", "worker1:2222", "worker2:2222"] + }) + simple_cr = SimpleClusterResolver(base_cluster_spec) + UnionClusterResolver(simple_cr) + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/cluster_resolver/python/training/__init__.py b/tensorflow/contrib/cluster_resolver/python/training/__init__.py index f4dbadcbb5db923406119d825428975e33df26bd..10d93549ebbd4f7e900796d0516b0af1744224af 100644 --- a/tensorflow/contrib/cluster_resolver/python/training/__init__.py +++ b/tensorflow/contrib/cluster_resolver/python/training/__init__.py @@ -18,9 +18,36 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -from tensorflow.contrib.cluster_resolver.python.training.cluster_resolver import ClusterResolver -from tensorflow.contrib.cluster_resolver.python.training.cluster_resolver import SimpleClusterResolver -from tensorflow.contrib.cluster_resolver.python.training.cluster_resolver import UnionClusterResolver -from tensorflow.contrib.cluster_resolver.python.training.gce_cluster_resolver import GceClusterResolver -from tensorflow.contrib.cluster_resolver.python.training.slurm_cluster_resolver import SlurmClusterResolver -from tensorflow.contrib.cluster_resolver.python.training.tpu_cluster_resolver import TPUClusterResolver +# This file (and all files in this directory in general) is a backwards +# compatibility shim that exists to re-export ClusterResolvers such that +# existing OSS code will not be broken. + +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.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 +from tensorflow.python.distribute.cluster_resolver.tpu_cluster_resolver import TPUClusterResolver + +from tensorflow.python.util.all_util import remove_undocumented + +_allowed_symbols = [ + 'cluster_resolver', + 'gce_cluster_resolver', + 'kubernetes_cluster_resolver', + 'slurm_cluster_resolver', + 'tfconfig_cluster_resolver', + 'tpu_cluster_resolver', + 'ClusterResolver', + 'SimpleClusterResolver', + 'UnionClusterResolver', + 'GceClusterResolver', + 'KubernetesClusterResolver', + 'TFConfigClusterResolver', + 'TPUClusterResolver', + 'SlurmClusterResolver', +] + +remove_undocumented(__name__, _allowed_symbols) diff --git a/tensorflow/contrib/cluster_resolver/python/training/cluster_resolver.py b/tensorflow/contrib/cluster_resolver/python/training/cluster_resolver.py index 1c480b25134b1e54200e0ddb780bd7bb0f122341..99840fb5166dd739b3bee06a926e06b534011d1f 100644 --- a/tensorflow/contrib/cluster_resolver/python/training/cluster_resolver.py +++ b/tensorflow/contrib/cluster_resolver/python/training/cluster_resolver.py @@ -1,4 +1,4 @@ -# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# 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. @@ -12,181 +12,29 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Cluster Resolvers are used for dynamic cluster IP/hostname resolution.""" +"""Stub file for ClusterResolver to maintain backwards compatibility.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function -import abc +# This file (and all files in this directory in general) is a backwards +# compatibility shim that exists to re-export ClusterResolvers such that +# existing OSS code will not be broken. -from tensorflow.python.training.server_lib import ClusterSpec +# pylint: disable=unused-import +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 +# pylint: enable=unused-import +from tensorflow.python.util.all_util import remove_undocumented -class ClusterResolver(object): - """Abstract class for all implementations of ClusterResolvers. +_allowed_symbols = [ + 'ClusterResolver', + 'SimpleClusterResolver', + 'UnionClusterResolver', +] - This defines the skeleton for all implementations of ClusterResolvers. - ClusterResolvers are a way for TensorFlow to communicate with various cluster - management systems (e.g. GCE, AWS, etc...). +remove_undocumented(__name__, _allowed_symbols) - By letting TensorFlow communicate with these systems, we will be able to - automatically discover and resolve IP addresses for various TensorFlow - workers. This will eventually allow us to automatically recover from - underlying machine failures and scale TensorFlow worker clusters up and down. - """ - - @abc.abstractmethod - def cluster_spec(self): - """Retrieve the current state of the cluster and returns a ClusterSpec. - - Returns: - A ClusterSpec representing the state of the cluster at the moment this - function is called. - - Implementors of this function must take care in ensuring that the - ClusterSpec returned is up-to-date at the time of calling this function. - This usually means retrieving the information from the underlying cluster - management system every time this function is invoked and reconstructing - a cluster_spec, rather than attempting to cache anything. - """ - raise NotImplementedError( - 'cluster_spec is not implemented for {}.'.format(self)) - - @abc.abstractmethod - def master(self): - """...""" - raise NotImplementedError('master is not implemented for {}.'.format(self)) - - -class SimpleClusterResolver(ClusterResolver): - """Simple implementation of ClusterResolver that accepts a ClusterSpec.""" - - def __init__(self, cluster_spec, master=''): - """Creates a SimpleClusterResolver from a ClusterSpec.""" - super(SimpleClusterResolver, self).__init__() - - if not isinstance(cluster_spec, ClusterSpec): - raise TypeError('cluster_spec must be a ClusterSpec.') - self._cluster_spec = cluster_spec - - if not isinstance(master, str): - raise TypeError('master must be a string.') - self._master = master - - def cluster_spec(self): - """Returns the ClusterSpec passed into the constructor.""" - return self._cluster_spec - - def master(self): - """Returns the master address to use when creating a session.""" - return self._master - - -class UnionClusterResolver(ClusterResolver): - """Performs a union on underlying ClusterResolvers. - - This class performs a union given two or more existing ClusterResolvers. It - merges the underlying ClusterResolvers, and returns one unified ClusterSpec - when cluster_spec is called. The details of the merge function is - documented in the cluster_spec function. - """ - - def __init__(self, *args): - """Initializes a UnionClusterResolver with other ClusterResolvers. - - Args: - *args: `ClusterResolver` objects to be unionized. - - Raises: - TypeError: If any argument is not a subclass of `ClusterResolvers`. - ValueError: If there are no arguments passed. - """ - super(UnionClusterResolver, self).__init__() - - if not args: - raise ValueError('At least one ClusterResolver is required.') - - for cluster_resolver in args: - if not isinstance(cluster_resolver, ClusterResolver): - raise TypeError('All arguments must be a sub-class of ' - '`ClusterResolver.`') - self._cluster_resolvers = args - - def cluster_spec(self): - """Returns a union of all the ClusterSpecs from the ClusterResolvers. - - Returns: - A ClusterSpec containing host information merged from all the underlying - ClusterResolvers. - - Raises: - KeyError: If there are conflicting keys detected when merging two or - more dictionaries, this exception is raised. - - Note: If there are multiple ClusterResolvers exposing ClusterSpecs with the - same job name, we will merge the list/dict of workers. - - If *all* underlying ClusterSpecs expose the set of workers as lists, we will - concatenate the lists of workers, starting with the list of workers from - the first ClusterResolver passed into the constructor. - - If *any* of the ClusterSpecs expose the set of workers as a dict, we will - treat all the sets of workers as dicts (even if they are returned as lists) - and will only merge them into a dict if there is no conflicting keys. If - there is a conflicting key, we will raise a `KeyError`. - """ - - merged_cluster = {} - - # We figure out whether it is all lists for a particular job, or whether - # there are dicts inside. - for cluster_resolver in self._cluster_resolvers: - cluster_spec = cluster_resolver.cluster_spec() - cluster_dict = cluster_spec.as_dict() - - for job_name, tasks in cluster_dict.items(): - if job_name in merged_cluster: - # If we see a dict, then we write a dict out regardless. - if isinstance(tasks, dict): - merged_cluster[job_name] = {} - else: - # We take whichever type is present. - if isinstance(tasks, list): - merged_cluster[job_name] = [] - else: - merged_cluster[job_name] = {} - - # We then do the merge as appropriate in merged_cluster[job]. - for cluster_resolver in self._cluster_resolvers: - cluster_spec = cluster_resolver.cluster_spec() - cluster_dict = cluster_spec.as_dict() - - for job_name, tasks in cluster_dict.items(): - if isinstance(merged_cluster[job_name], list): - # We all have lists, we can just concatenate and be done. - merged_cluster[job_name].extend(tasks) - else: - if isinstance(tasks, list): - # We convert to a dictionary if the type is a list. - task_dict = dict(zip(range(0, len(tasks)), tasks)) - else: - # We can simply make a copy (for update) and be done. - task_dict = tasks.copy() - - # We detect if there are duplicates, and raise an error if so. - task_keys = set(task_dict) - merged_keys = set(merged_cluster[job_name].keys()) - intersected_keys = task_keys.intersection(merged_keys) - if intersected_keys: - raise KeyError('Duplicate keys detected when merging two ' - 'ClusterSpecs: %s' % repr(intersected_keys)) - - # We do the merge after all the processing. - merged_cluster[job_name].update(task_dict) - - return ClusterSpec(merged_cluster) - - def master(self): - """master returns the master address from the first cluster resolver.""" - return self._cluster_resolvers[0].master() diff --git a/tensorflow/contrib/cluster_resolver/python/training/cluster_resolver_test.py b/tensorflow/contrib/cluster_resolver/python/training/cluster_resolver_test.py deleted file mode 100644 index d9c97d53eb3663f6ab2f7b40395592dc7638b896..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/cluster_resolver/python/training/cluster_resolver_test.py +++ /dev/null @@ -1,240 +0,0 @@ -# Copyright 2017 The TensorFlow Authors. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== -"""Tests for Cluster Resolvers.""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -from tensorflow.contrib.cluster_resolver.python.training.cluster_resolver import SimpleClusterResolver -from tensorflow.contrib.cluster_resolver.python.training.cluster_resolver import UnionClusterResolver -from tensorflow.python.platform import test -from tensorflow.python.training import server_lib - - -class UnionClusterResolverTest(test.TestCase): - # TODO(frankchn): Transform to parameterized test after it is included in the - # TF open source codebase. - - def _verifyClusterSpecEquality(self, cluster_spec, expected_proto): - self.assertProtoEquals(expected_proto, cluster_spec.as_cluster_def()) - self.assertProtoEquals( - expected_proto, server_lib.ClusterSpec(cluster_spec).as_cluster_def()) - self.assertProtoEquals( - expected_proto, - server_lib.ClusterSpec(cluster_spec.as_cluster_def()).as_cluster_def()) - self.assertProtoEquals( - expected_proto, - server_lib.ClusterSpec(cluster_spec.as_dict()).as_cluster_def()) - - def testSingleClusterResolver(self): - base_cluster_spec = server_lib.ClusterSpec({ - "ps": ["ps0:2222", "ps1:2222"], - "worker": ["worker0:2222", "worker1:2222", "worker2:2222"] - }) - simple_resolver = SimpleClusterResolver(base_cluster_spec) - union_resolver = UnionClusterResolver(simple_resolver) - - expected_proto = """ - job { name: 'ps' tasks { key: 0 value: 'ps0:2222' } - tasks { key: 1 value: 'ps1:2222' } } - job { name: 'worker' tasks { key: 0 value: 'worker0:2222' } - tasks { key: 1 value: 'worker1:2222' } - tasks { key: 2 value: 'worker2:2222' } } - """ - actual_cluster_spec = union_resolver.cluster_spec() - self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto) - - def testTwoNonOverlappingJobMergedClusterResolver(self): - cluster_spec_1 = server_lib.ClusterSpec({ - "ps": [ - "ps0:2222", - "ps1:2222" - ] - }) - cluster_spec_2 = server_lib.ClusterSpec({ - "worker": [ - "worker0:2222", - "worker1:2222", - "worker2:2222" - ] - }) - cluster_resolver_1 = SimpleClusterResolver(cluster_spec_1) - cluster_resolver_2 = SimpleClusterResolver(cluster_spec_2) - - union_cluster = UnionClusterResolver(cluster_resolver_1, cluster_resolver_2) - cluster_spec = union_cluster.cluster_spec() - - expected_proto = """ - job { name: 'ps' tasks { key: 0 value: 'ps0:2222' } - tasks { key: 1 value: 'ps1:2222' } } - job { name: 'worker' tasks { key: 0 value: 'worker0:2222' } - tasks { key: 1 value: 'worker1:2222' } - tasks { key: 2 value: 'worker2:2222' } } - """ - self._verifyClusterSpecEquality(cluster_spec, expected_proto) - - def testOverlappingJobMergedClusterResolver(self): - cluster_spec_1 = server_lib.ClusterSpec({ - "worker": [ - "worker4:2222", - "worker5:2222" - ] - }) - cluster_spec_2 = server_lib.ClusterSpec({ - "worker": [ - "worker0:2222", - "worker1:2222", - "worker2:2222" - ] - }) - cluster_resolver_1 = SimpleClusterResolver(cluster_spec_1) - cluster_resolver_2 = SimpleClusterResolver(cluster_spec_2) - - union_cluster = UnionClusterResolver(cluster_resolver_1, cluster_resolver_2) - cluster_spec = union_cluster.cluster_spec() - - expected_proto = """ - job { name: 'worker' tasks { key: 0 value: 'worker4:2222' } - tasks { key: 1 value: 'worker5:2222' } - tasks { key: 2 value: 'worker0:2222' } - tasks { key: 3 value: 'worker1:2222' } - tasks { key: 4 value: 'worker2:2222' } } - """ - self._verifyClusterSpecEquality(cluster_spec, expected_proto) - - def testOverlappingSparseJobMergedClusterResolverThrowError(self): - cluster_spec_1 = server_lib.ClusterSpec({ - "worker": { - 7: "worker4:2222", - 9: "worker5:2222" - } - }) - cluster_spec_2 = server_lib.ClusterSpec({ - "worker": { - 3: "worker0:2222", - 6: "worker1:2222", - 7: "worker2:2222" - } - }) - cluster_resolver_1 = SimpleClusterResolver(cluster_spec_1) - cluster_resolver_2 = SimpleClusterResolver(cluster_spec_2) - - union_cluster = UnionClusterResolver(cluster_resolver_1, cluster_resolver_2) - self.assertRaises(KeyError, union_cluster.cluster_spec) - - def testOverlappingDictAndListThrowError(self): - cluster_spec_1 = server_lib.ClusterSpec({ - "worker": [ - "worker4:2222", - "worker5:2222" - ] - }) - cluster_spec_2 = server_lib.ClusterSpec({ - "worker": { - 1: "worker0:2222", - 2: "worker1:2222", - 3: "worker2:2222" - } - }) - cluster_resolver_1 = SimpleClusterResolver(cluster_spec_1) - cluster_resolver_2 = SimpleClusterResolver(cluster_spec_2) - - union_cluster = UnionClusterResolver(cluster_resolver_1, cluster_resolver_2) - self.assertRaises(KeyError, union_cluster.cluster_spec) - - def testOverlappingJobNonOverlappingKey(self): - cluster_spec_1 = server_lib.ClusterSpec({ - "worker": { - 5: "worker4:2222", - 9: "worker5:2222" - } - }) - cluster_spec_2 = server_lib.ClusterSpec({ - "worker": { - 3: "worker0:2222", - 6: "worker1:2222", - 7: "worker2:2222" - } - }) - cluster_resolver_1 = SimpleClusterResolver(cluster_spec_1) - cluster_resolver_2 = SimpleClusterResolver(cluster_spec_2) - - union_cluster = UnionClusterResolver(cluster_resolver_1, cluster_resolver_2) - cluster_spec = union_cluster.cluster_spec() - - expected_proto = """ - job { name: 'worker' tasks { key: 3 value: 'worker0:2222' } - tasks { key: 5 value: 'worker4:2222' } - tasks { key: 6 value: 'worker1:2222' } - tasks { key: 7 value: 'worker2:2222' } - tasks { key: 9 value: 'worker5:2222' }} - """ - self._verifyClusterSpecEquality(cluster_spec, expected_proto) - - def testMixedModeNonOverlappingKey(self): - cluster_spec_1 = server_lib.ClusterSpec({ - "worker": [ - "worker4:2222", - "worker5:2222" - ] - }) - cluster_spec_2 = server_lib.ClusterSpec({ - "worker": { - 3: "worker0:2222", - 6: "worker1:2222", - 7: "worker2:2222" - } - }) - cluster_resolver_1 = SimpleClusterResolver(cluster_spec_1) - cluster_resolver_2 = SimpleClusterResolver(cluster_spec_2) - - union_cluster = UnionClusterResolver(cluster_resolver_1, cluster_resolver_2) - cluster_spec = union_cluster.cluster_spec() - - expected_proto = """ - job { name: 'worker' tasks { key: 0 value: 'worker4:2222' } - tasks { key: 1 value: 'worker5:2222' } - tasks { key: 3 value: 'worker0:2222' } - tasks { key: 6 value: 'worker1:2222' } - tasks { key: 7 value: 'worker2:2222' }} - """ - self._verifyClusterSpecEquality(cluster_spec, expected_proto) - - def testRetainSparseJobWithNoMerging(self): - base_cluster_spec = server_lib.ClusterSpec({ - "worker": { - 1: "worker0:2222", - 3: "worker1:2222", - 5: "worker2:2222" - } - }) - - base_cluster_resolver = SimpleClusterResolver(base_cluster_spec) - union_cluster = UnionClusterResolver(base_cluster_resolver) - cluster_spec = union_cluster.cluster_spec() - - expected_proto = """ - job { name: 'worker' tasks { key: 1 value: 'worker0:2222' } - tasks { key: 3 value: 'worker1:2222' } - tasks { key: 5 value: 'worker2:2222' } } - """ - self._verifyClusterSpecEquality(cluster_spec, expected_proto) - - -# TODO(saeta): Include tests for master resolution - -if __name__ == "__main__": - test.main() 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 3f5824128948453634bc5e5a7d6fdeedae60f5bd..55e61155c683c928efab9bb018868faec3e3df8c 100644 --- a/tensorflow/contrib/cluster_resolver/python/training/gce_cluster_resolver.py +++ b/tensorflow/contrib/cluster_resolver/python/training/gce_cluster_resolver.py @@ -1,4 +1,4 @@ -# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# 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. @@ -12,128 +12,24 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Implementation of Cluster Resolvers for GCE Instance Groups.""" +"""Stub file for GceClusterResolver to maintain backwards compatibility.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function +# This file (and all files in this directory in general) is a backwards +# compatibility shim that exists to re-export ClusterResolvers such that +# existing OSS code will not be broken. -from tensorflow.contrib.cluster_resolver.python.training.cluster_resolver import ClusterResolver -from tensorflow.python.training.server_lib import ClusterSpec +# pylint: disable=unused-import +from tensorflow.python.distribute.cluster_resolver.gce_cluster_resolver import GceClusterResolver +# pylint: enable=unused-import -_GOOGLE_API_CLIENT_INSTALLED = True -try: - from googleapiclient import discovery # pylint: disable=g-import-not-at-top - from oauth2client.client import GoogleCredentials # pylint: disable=g-import-not-at-top -except ImportError: - _GOOGLE_API_CLIENT_INSTALLED = False +from tensorflow.python.util.all_util import remove_undocumented +_allowed_symbols = [ + 'GceClusterResolver', +] -class GceClusterResolver(ClusterResolver): - """Cluster Resolver for Google Compute Engine. - - This is an implementation of cluster resolvers for the Google Compute Engine - instance group platform. By specifying a project, zone, and instance group, - this will retrieve the IP address of all the instances within the instance - group and return a Cluster Resolver object suitable for use for distributed - TensorFlow. - """ - - def __init__(self, - project, - zone, - instance_group, - port, - job_name='worker', - credentials='default', - service=None): - """Creates a new GceClusterResolver object. - - This takes in a few parameters and creates a GceClusterResolver project. It - will then use these parameters to query the GCE API for the IP addresses of - each instance in the instance group. - - Args: - project: Name of the GCE project - zone: Zone of the GCE instance group - instance_group: Name of the GCE instance group - port: Port of the listening TensorFlow server (default: 8470) - job_name: Name of the TensorFlow job this set of instances belongs to - credentials: GCE Credentials. If nothing is specified, this defaults to - GoogleCredentials.get_application_default() - service: The GCE API object returned by the googleapiclient.discovery - function. (Default: discovery.build('compute', 'v1')). If you specify a - custom service object, then the credentials parameter will be ignored. - - Raises: - ImportError: If the googleapiclient is not installed. - """ - self._project = project - self._zone = zone - self._instance_group = instance_group - self._job_name = job_name - self._port = port - self._credentials = credentials - - if credentials == 'default': - if _GOOGLE_API_CLIENT_INSTALLED: - self._credentials = GoogleCredentials.get_application_default() - - if service is None: - if not _GOOGLE_API_CLIENT_INSTALLED: - raise ImportError('googleapiclient must be installed before using the ' - 'GCE cluster resolver') - self._service = discovery.build( - 'compute', 'v1', - credentials=self._credentials) - else: - self._service = service - - def cluster_spec(self): - """Returns a ClusterSpec object based on the latest instance group info. - - This returns a ClusterSpec object for use based on information from the - specified instance group. We will retrieve the information from the GCE APIs - every time this method is called. - - Returns: - A ClusterSpec containing host information retrieved from GCE. - """ - request_body = {'instanceState': 'RUNNING'} - request = self._service.instanceGroups().listInstances( - project=self._project, - zone=self._zone, - instanceGroups=self._instance_group, - body=request_body, - orderBy='name') - - worker_list = [] - - while request is not None: - response = request.execute() - - items = response['items'] - for instance in items: - instance_name = instance['instance'].split('/')[-1] - - instance_request = self._service.instances().get( - project=self._project, - zone=self._zone, - instance=instance_name) - - if instance_request is not None: - instance_details = instance_request.execute() - ip_address = instance_details['networkInterfaces'][0]['networkIP'] - instance_url = '%s:%s' % (ip_address, self._port) - worker_list.append(instance_url) - - request = self._service.instanceGroups().listInstances_next( - previous_request=request, - previous_response=response) - - worker_list.sort() - return ClusterSpec({self._job_name: worker_list}) - - def master(self): - return '' +remove_undocumented(__name__, _allowed_symbols) diff --git a/tensorflow/contrib/cluster_resolver/python/training/kubernetes_cluster_resolver.py b/tensorflow/contrib/cluster_resolver/python/training/kubernetes_cluster_resolver.py new file mode 100644 index 0000000000000000000000000000000000000000..a8eaf33629a6299d5da5f8a930e0cad7d07044e8 --- /dev/null +++ b/tensorflow/contrib/cluster_resolver/python/training/kubernetes_cluster_resolver.py @@ -0,0 +1,36 @@ +# 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. +# ============================================================================== +"""Stub file for KubernetesClusterResolver for backwards compatibility.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +# This file (and all files in this directory in general) is a backwards +# compatibility shim that exists to re-export ClusterResolvers such that +# existing OSS code will not be broken. + +# pylint: disable=unused-import +from tensorflow.python.distribute.cluster_resolver.kubernetes_cluster_resolver import KubernetesClusterResolver +# pylint: enable=unused-import + +from tensorflow.python.util.all_util import remove_undocumented + +_allowed_symbols = [ + 'KubernetesClusterResolver', +] + +remove_undocumented(__name__, _allowed_symbols) + diff --git a/tensorflow/contrib/cluster_resolver/python/training/slurm_cluster_resolver.py b/tensorflow/contrib/cluster_resolver/python/training/slurm_cluster_resolver.py index 227b31596d22ccbb1b5ff37478ec054beec3e2d5..fcd2a846eeb1be7ad4b5a98b067a125afbbebc7d 100644 --- a/tensorflow/contrib/cluster_resolver/python/training/slurm_cluster_resolver.py +++ b/tensorflow/contrib/cluster_resolver/python/training/slurm_cluster_resolver.py @@ -12,182 +12,24 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Implementation of Cluster Resolvers for Slurm workload manager.""" +"""Stub file for SlurmClusterResolver to maintain backwards compatibility.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function -import os -import subprocess +# This file (and all files in this directory in general) is a backwards +# compatibility shim that exists to re-export ClusterResolvers such that +# existing OSS code will not be broken. -from tensorflow.contrib.cluster_resolver.python.training.cluster_resolver import ClusterResolver -from tensorflow.python.training.server_lib import ClusterSpec +# pylint: disable=unused-import +from tensorflow.python.distribute.cluster_resolver.slurm_cluster_resolver import SlurmClusterResolver +# pylint: enable=unused-import +from tensorflow.python.util.all_util import remove_undocumented -class SlurmClusterResolver(ClusterResolver): - """Cluster Resolver for system with Slurm workload manager. +_allowed_symbols = [ + 'SlurmClusterResolver', +] - This is an implementation of cluster resolvers for Slurm clusters. This allows - the specification of jobs and task counts, number of tasks per node, number of - GPUs on each node and number of GPUs for each task, It retrieves system - attributes by Slurm environment variables, resolves allocated computing node - names, construct a cluster and return a Cluster Resolver object which an be - use for distributed TensorFlow. - """ - - def _resolve_hostnames(self): - """Resolve host names of nodes allocated in current jobs. - - Returns: - A list of node names as strings. - """ - hostlist = (subprocess.check_output(['scontrol', 'show', 'hostname']). - decode('utf-8').strip().split('\n')) - return hostlist - - def __init__(self, - jobs, - port_base=8888, - gpus_per_node=1, - gpus_per_task=1, - tasks_per_node=None, - auto_set_gpu=True): - """Creates a new SlurmClusterResolver object. - - This takes in parameters and creates a SlurmClusterResolver object. It uses - those parameters to check which nodes will processes reside and resolves - their hostnames. With the number of the GPUs on each node and number of GPUs - for each task it offsets the port number for each processes and allocate - GPUs to tasks by setting environment variables. The resolver currently - supports homogeneous tasks and default Slurm process allocation. - - Args: - jobs: Dictionary with job names as key and number of tasks in the job as - value - port_base: The first port number to start with for processes on a node. - gpus_per_node: Number of GPUs available on each node. - gpus_per_task: Number of GPUs to be used for each task. - tasks_per_node: Number of tasks to run on each node, if not set defaults - to Slurm's output environment variable SLURM_NTASKS_PER_NODE. - auto_set_gpu: Set the visible CUDA devices automatically while resolving - the cluster by setting CUDA_VISIBLE_DEVICES environment variable. - Defaults to True. - - Returns: - A ClusterResolver object which can be used with distributed TensorFlow. - - Raises: - RuntimeError: If requested more GPUs per node then available or requested - more tasks then assigned tasks. - """ - - # check if launched by mpirun - if 'OMPI_COMM_WORLD_RANK' in os.environ: - self._rank = int(os.environ['OMPI_COMM_WORLD_RANK']) - num_tasks = int(os.environ['OMPI_COMM_WORLD_SIZE']) - else: - self._rank = int(os.environ['SLURM_PROCID']) - num_tasks = int(os.environ['SLURM_NTASKS']) - - self._jobs = jobs - self._port_base = port_base - - # user specification overrides SLURM specification - if tasks_per_node is not None: - self._tasks_per_node = tasks_per_node - elif tasks_per_node is None and 'SLURM_NTASKS_PER_NODE' in os.environ: - self._tasks_per_node = int(os.environ['SLURM_NTASKS_PER_NODE']) - else: - raise RuntimeError('Neither `tasks_per_node` or ' - 'SLURM_NTASKS_PER_NODE is set.') - - self._gpus_per_node = gpus_per_node - self._gpus_per_task = gpus_per_task - - self._auto_set_gpu = auto_set_gpu - self._job_name = None - self._task_index = None - - self._gpu_allocation = [] - self._cluster_allocation = {} - - if self._tasks_per_node * self._gpus_per_task > self._gpus_per_node: - raise RuntimeError('Requested more GPUs per node then available.') - - if sum(self._jobs.values()) != num_tasks: - raise RuntimeError('Requested more tasks then assigned tasks.') - - def cluster_spec(self): - """Returns a ClusterSpec object based on the latest instance group info. - - This returns a ClusterSpec object for use based on information from the - specified initialization parameters and Slurm environment variables. The - cluster specification is resolved each time this function is called. The - resolver extract hostnames of nodes by scontrol and pack tasks in that - order until a node a has number of tasks that is equal to specification. - GPUs on nodes are allocated to tasks by specification through setting - CUDA_VISIBLE_DEVICES environment variable. - - Returns: - A ClusterSpec containing host information retrieved from Slurm's - environment variables. - """ - hostlist = self._resolve_hostnames() - - task_list = [] - self._gpu_allocation = [] - self._cluster_allocation = {} - - for host in hostlist: - for port_offset, gpu_offset in zip( - range(self._tasks_per_node), - range(0, self._gpus_per_node, self._gpus_per_task)): - - host_addr = '%s:%d' % (host, self._port_base + port_offset) - task_list.append(host_addr) - gpu_id_list = [] - - for gpu_id in range(gpu_offset, gpu_offset + self._gpus_per_task): - gpu_id_list.append(str(gpu_id)) - - self._gpu_allocation.append(','.join(gpu_id_list)) - - cluster_rank_offset_start = 0 - cluster_rank_offset_end = 0 - - for job_name, num_tasks in self._jobs.items(): - cluster_rank_offset_end = cluster_rank_offset_start + num_tasks - - self._cluster_allocation[job_name] = \ - task_list[cluster_rank_offset_start:cluster_rank_offset_end] - - if self._rank >= cluster_rank_offset_start and \ - self._rank < cluster_rank_offset_end: - - self._job_name = job_name - self._task_index = self._rank - cluster_rank_offset_start - - cluster_rank_offset_start = cluster_rank_offset_end - - if self._auto_set_gpu is True: - os.environ['CUDA_VISIBLE_DEVICES'] = self._gpu_allocation[self._rank] - - return ClusterSpec(self._cluster_allocation) - - def get_task_info(self): - """Returns job name and task_index for the process which calls this. - - This returns the job name and task index for the process which calls this - function according to its rank and cluster specification. The job name and - task index are set after a cluster is constructed by cluster_spec otherwise - defaults to None. - - Returns: - A string specifying job name the process belongs to and an integner - specifying the task index the process belongs to in that job. - """ - return self._job_name, self._task_index - - def master(self): - return self._cluster_allocation[str(self._job_name)][self._task_index] +remove_undocumented(__name__, _allowed_symbols) diff --git a/tensorflow/contrib/cluster_resolver/python/training/slurm_cluster_resolver_test.py b/tensorflow/contrib/cluster_resolver/python/training/slurm_cluster_resolver_test.py deleted file mode 100644 index d108fd1b49dcb9d7e852639a48456a3994a7f55d..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/cluster_resolver/python/training/slurm_cluster_resolver_test.py +++ /dev/null @@ -1,224 +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. -# ============================================================================== -"""Tests for SlurmClusterResolver.""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import os - -from tensorflow.contrib.cluster_resolver.python.training.slurm_cluster_resolver import SlurmClusterResolver -from tensorflow.python.platform import test -from tensorflow.python.training import server_lib - -mock = test.mock - - -class SlurmClusterResolverTest(test.TestCase): - - def mock_resolve_hostnames_output(self): - return ['t02n13', 't02n41', 't02n43', 't02n44'] - - def _verifyClusterSpecEquality(self, cluster_spec, expected_proto): - self.assertProtoEquals(expected_proto, cluster_spec.as_cluster_def()) - self.assertProtoEquals( - expected_proto, - server_lib.ClusterSpec(cluster_spec).as_cluster_def()) - self.assertProtoEquals( - expected_proto, - server_lib.ClusterSpec(cluster_spec.as_cluster_def()).as_cluster_def()) - self.assertProtoEquals( - expected_proto, - server_lib.ClusterSpec(cluster_spec.as_dict()).as_cluster_def()) - - @mock.patch.dict(os.environ, {'SLURM_PROCID': '0', 'SLURM_NTASKS': '3'}) - @mock.patch.object(SlurmClusterResolver, '_resolve_hostnames', - mock_resolve_hostnames_output) - def testSimpleSuccessfulRetrieval(self): - slurm_cluster_resolver = SlurmClusterResolver( - jobs={ - 'ps': 1, - 'worker': 2 - }, - port_base=8888, - tasks_per_node=1, - gpus_per_node=1, - gpus_per_task=1, - auto_set_gpu=False) - - actual_cluster_spec = slurm_cluster_resolver.cluster_spec() - expected_proto = """ - job { - name: "ps" - tasks { - value: "t02n13:8888" - } - } - job { - name: "worker" - tasks { - value: "t02n41:8888" - } - tasks { - key: 1 - value: "t02n43:8888" - } - } - """ - self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto) - - @mock.patch.dict(os.environ, { - 'SLURM_PROCID': '0', - 'SLURM_NTASKS': '3', - 'SLURM_NTASKS_PER_NODE': '1' - }) - @mock.patch.object(SlurmClusterResolver, '_resolve_hostnames', - mock_resolve_hostnames_output) - def testTaskPerNodeNotSetRetrieval(self): - slurm_cluster_resolver = SlurmClusterResolver( - jobs={ - 'ps': 1, - 'worker': 2 - }, - port_base=8888, - gpus_per_node=1, - gpus_per_task=1, - auto_set_gpu=False) - - actual_cluster_spec = slurm_cluster_resolver.cluster_spec() - expected_proto = """ - job { - name: "ps" - tasks { - value: "t02n13:8888" - } - } - job { - name: "worker" - tasks { - value: "t02n41:8888" - } - tasks { - key: 1 - value: "t02n43:8888" - } - } - """ - self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto) - - @mock.patch.dict( - os.environ, { - 'SLURM_PROCID': '1', - 'SLURM_NTASKS': '5', - 'SLURM_NTASKS_PER_NODE': '2', - 'CUDA_VISIBLE_DEVICES': '' - }) - @mock.patch.object(SlurmClusterResolver, '_resolve_hostnames', - mock_resolve_hostnames_output) - def testMultiTaskPerNodeRetrieval(self): - slurm_cluster_resolver = SlurmClusterResolver( - jobs={ - 'ps': 1, - 'worker': 4 - }, - port_base=8888, - gpus_per_node=2, - gpus_per_task=1, - auto_set_gpu=True) - - actual_cluster_spec = slurm_cluster_resolver.cluster_spec() - expected_proto = """ - job { - name: "ps" - tasks { - value: "t02n13:8888" - } - } - job { - name: "worker" - tasks { - value: "t02n13:8889" - } - tasks { - key: 1 - value: "t02n41:8888" - } - tasks { - key: 2 - value: "t02n41:8889" - } - tasks { - key: 3 - value: "t02n43:8888" - } - } - """ - self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto) - assert os.environ['CUDA_VISIBLE_DEVICES'] == '1' - - @mock.patch.dict( - os.environ, { - 'SLURM_PROCID': '1', - 'SLURM_NTASKS': '5', - 'SLURM_NTASKS_PER_NODE': '2', - 'CUDA_VISIBLE_DEVICES': '' - }) - @mock.patch.object(SlurmClusterResolver, '_resolve_hostnames', - mock_resolve_hostnames_output) - def testMultipleGpusPerTaskRetrieval(self): - slurm_cluster_resolver = SlurmClusterResolver( - jobs={ - 'ps': 1, - 'worker': 4 - }, - port_base=8888, - gpus_per_node=4, - gpus_per_task=2, - auto_set_gpu=True) - - actual_cluster_spec = slurm_cluster_resolver.cluster_spec() - expected_proto = """ - job { - name: "ps" - tasks { - value: "t02n13:8888" - } - } - job { - name: "worker" - tasks { - value: "t02n13:8889" - } - tasks { - key: 1 - value: "t02n41:8888" - } - tasks { - key: 2 - value: "t02n41:8889" - } - tasks { - key: 3 - value: "t02n43:8888" - } - } - """ - self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto) - assert os.environ['CUDA_VISIBLE_DEVICES'] == '2,3' - - -if __name__ == '__main__': - test.main() diff --git a/tensorflow/contrib/cluster_resolver/python/training/tfconfig_cluster_resolver.py b/tensorflow/contrib/cluster_resolver/python/training/tfconfig_cluster_resolver.py new file mode 100644 index 0000000000000000000000000000000000000000..9db7f47dcb49c499719b9002b1d2d6c4837a7bd2 --- /dev/null +++ b/tensorflow/contrib/cluster_resolver/python/training/tfconfig_cluster_resolver.py @@ -0,0 +1,36 @@ +# 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. +# ============================================================================== +"""Stub file for TFConfigClusterResolver to maintain backwards compatibility.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +# This file (and all files in this directory in general) is a backwards +# compatibility shim that exists to re-export ClusterResolvers such that +# existing OSS code will not be broken. + +# pylint: disable=unused-import +from tensorflow.python.distribute.cluster_resolver.tfconfig_cluster_resolver import TFConfigClusterResolver +# pylint: enable=unused-import + +from tensorflow.python.util.all_util import remove_undocumented + +_allowed_symbols = [ + 'TFConfigClusterResolver', +] + +remove_undocumented(__name__, _allowed_symbols) + diff --git a/tensorflow/contrib/cluster_resolver/python/training/tpu_cluster_resolver.py b/tensorflow/contrib/cluster_resolver/python/training/tpu_cluster_resolver.py index f4a8e16c99f464b813a98e981579bd0ff53bd464..3a1eaccd06e574babbe9a3232dacd1d66f3a4648 100644 --- a/tensorflow/contrib/cluster_resolver/python/training/tpu_cluster_resolver.py +++ b/tensorflow/contrib/cluster_resolver/python/training/tpu_cluster_resolver.py @@ -1,4 +1,4 @@ -# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# 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. @@ -12,311 +12,24 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Implementation of Cluster Resolvers for Cloud TPUs.""" +"""Stub file for TPUClusterResolver to maintain backwards compatibility.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function -import os +# This file (and all files in this directory in general) is a backwards +# compatibility shim that exists to re-export ClusterResolvers such that +# existing OSS code will not be broken. -from six.moves.urllib.request import Request -from six.moves.urllib.request import urlopen +# pylint: disable=unused-import +from tensorflow.python.distribute.cluster_resolver.tpu_cluster_resolver import TPUClusterResolver +# pylint: enable=unused-import -from tensorflow.contrib.cluster_resolver.python.training.cluster_resolver import ClusterResolver -from tensorflow.python.training import server_lib -from tensorflow.python.util import compat +from tensorflow.python.util.all_util import remove_undocumented -_GOOGLE_API_CLIENT_INSTALLED = True -try: - from googleapiclient import discovery # pylint: disable=g-import-not-at-top - from oauth2client.client import GoogleCredentials # pylint: disable=g-import-not-at-top -except ImportError: - _GOOGLE_API_CLIENT_INSTALLED = False +_allowed_symbols = [ + 'TPUClusterResolver', +] - -_GKE_ENV_VARIABLE = 'KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS' -_ENDPOINTS_SEPARATOR = ',' -_DEFAULT_ENV_VARIABLE = 'TPU_NAME' -_DISCOVERY_SERVICE_URL_ENV_VARIABLE = 'TPU_API_DISCOVERY_URL' - - -class TPUClusterResolver(ClusterResolver): - """Cluster Resolver for Google Cloud TPUs. - - This is an implementation of cluster resolvers for the Google Cloud TPU - service. As Cloud TPUs are in alpha, you will need to specify a API definition - file for this to consume, in addition to a list of Cloud TPUs in your Google - Cloud Platform project. - """ - - def _requestComputeMetadata(self, path): - req = Request('http://metadata/computeMetadata/v1/%s' % path, - headers={'Metadata-Flavor': 'Google'}) - resp = urlopen(req) - return compat.as_bytes(resp.read()) - - def _shouldResolve(self): - if (self._tpu == compat.as_bytes('') or - self._tpu == compat.as_bytes('local') or - self._tpu.startswith(compat.as_bytes('/bns')) or - self._tpu.startswith(compat.as_bytes('localhost:')) or - self._tpu.startswith(compat.as_bytes('grpc://'))): - return False - return True - - @staticmethod - def _inGke(): - """When running in GKE, the environment variable will be set.""" - return _GKE_ENV_VARIABLE in os.environ - - @staticmethod - def _gkeEndpoints(): - return os.environ[_GKE_ENV_VARIABLE] - - @staticmethod - def _envVarFallback(): - if _DEFAULT_ENV_VARIABLE in os.environ: - return os.environ[_DEFAULT_ENV_VARIABLE] - return None - - @staticmethod - def _discoveryUrl(): - return os.environ.get(_DISCOVERY_SERVICE_URL_ENV_VARIABLE) - - def __init__(self, - tpu=None, - zone=None, - project=None, - job_name='worker', - coordinator_name=None, - coordinator_address=None, - credentials='default', - service=None, - discovery_url=None): - """Creates a new TPUClusterResolver object. - - The ClusterResolver will then use the parameters to query the Cloud TPU APIs - for the IP addresses and ports of each Cloud TPU listed. - - Args: - tpu: Either a string, or a list of strings corresponding to the TPUs to - use. If the single string is the empty string, the string 'local', or a - string that begins with 'grpc://' or '/bns', then it is assumed to not - correspond with a Cloud TPU and will instead be passed as the session - master and no ClusterSpec propagation will be done. - zone: Zone where the TPUs are located. If omitted or empty, we will assume - that the zone of the TPU is the same as the zone of the GCE VM, which we - will try to discover from the GCE metadata service. - project: Name of the GCP project containing Cloud TPUs. If omitted or - empty, we will try to discover the project name of the GCE VM from the - GCE metadata service. - job_name: Name of the TensorFlow job the TPUs belong to. - coordinator_name: The name to use for the coordinator. Set to None if the - coordinator should not be included in the computed ClusterSpec. - coordinator_address: The address of the coordinator (typically an ip:port - pair). If set to None, a TF server will be started. If coordinator_name - is None, a TF server will not be started even if coordinator_address is - None. - credentials: GCE Credentials. If None, then we use default credentials - from the oauth2client - service: The GCE API object returned by the googleapiclient.discovery - function. If you specify a custom service object, then the credentials - parameter will be ignored. - discovery_url: A URL template that points to the location of - the discovery service. It should have two parameters {api} and - {apiVersion} that when filled in produce an absolute URL to the - discovery document for that service. The environment variable - 'TPU_API_DISCOVERY_URL' will override this. - - Raises: - ImportError: If the googleapiclient is not installed. - ValueError: If no TPUs are specified. - """ - if isinstance(tpu, list): - if not tpu: - raise ValueError('At least one TPU must be specified.') - if len(tpu) != 1: - raise NotImplementedError( - 'Using multiple TPUs in a single session is not yet implemented') - tpu = tpu[0] - - in_gke = self._inGke() - # When using GKE with Cloud TPUs, the env variable will be set. - if tpu is None: - if in_gke: - tpu = self._gkeEndpoints() - else: - tpu = self._envVarFallback() - - if tpu is None: - raise ValueError('Please provide a TPU Name to connect to.') - - self._tpu = compat.as_bytes(tpu) # self._tpu is always bytes - self._job_name = job_name - self._credentials = credentials - - should_resolve = self._shouldResolve() - - if not project and should_resolve: - project = compat.as_str( - self._requestComputeMetadata('project/project-id')) - - if not zone and should_resolve: - zone_path = compat.as_str(self._requestComputeMetadata('instance/zone')) - zone = zone_path.split('/')[-1] - - self._project = project - self._zone = zone - - if credentials == 'default' and should_resolve: - if _GOOGLE_API_CLIENT_INSTALLED: - self._credentials = GoogleCredentials.get_application_default() - - if service is None and should_resolve: - if not _GOOGLE_API_CLIENT_INSTALLED: - raise ImportError('googleapiclient and oauth2client must be installed ' - 'before using the TPU cluster resolver. Execute: ' - '`pip install --upgrade google-api-python-client` ' - 'and `pip install --upgrade oauth2client` to ' - 'install with pip.') - - final_discovery_url = self._discoveryUrl() or discovery_url - if final_discovery_url: - self._service = discovery.build( - 'tpu', 'v1alpha1', - credentials=self._credentials, - discoveryServiceUrl=final_discovery_url) - else: - self._service = discovery.build( - 'tpu', 'v1alpha1', - credentials=self._credentials) - else: - self._service = service - - self._coordinator_name = coordinator_name - if coordinator_name and not coordinator_address and (should_resolve or - in_gke): - self._start_local_server() - else: - self._coordinator_address = coordinator_address - - def master(self): - """Get the Master string to be used for the session. - - In the normal case, this returns the grpc path (grpc://1.2.3.4:8470) of - first instance in the ClusterSpec returned by the cluster_spec function. - - If a non-TPU name is used when constructing a TPUClusterResolver, that will - be returned instead (e.g. If the tpus argument's value when constructing - this TPUClusterResolver was 'grpc://10.240.1.2:8470', - 'grpc://10.240.1.2:8470' will be returned). - - Returns: - string, the connection string to use when creating a session. - - Raises: - ValueError: If none of the TPUs specified exists. - """ - if not self._shouldResolve(): - return self._tpu.split(compat.as_bytes(_ENDPOINTS_SEPARATOR))[0] - - job_tasks = self.cluster_spec().job_tasks(self._job_name) - if not job_tasks: - raise ValueError('No TPUs exists with the specified names exist.') - - return 'grpc://' + job_tasks[0] - - def get_master(self): - return self.master() - - def get_job_name(self): - if self._shouldResolve(): - return self._job_name - - def cluster_spec(self): - """Returns a ClusterSpec object based on the latest TPU information. - - We retrieve the information from the GCE APIs every time this method is - called. - - Returns: - A ClusterSpec containing host information returned from Cloud TPUs. - - Raises: - RuntimeError: If the provided TPU is not healthy. - """ - ############################################################################ - # There are 5 potential cases this code must handle: - # 1. [Normal case.] We should resolve the TPU name to a set of tasks, and - # a. Create a ClusterSpec that includes the coordinator job - # b. Create a ClusterSpec without the coordinator job. - # 2. [GKE / No API Access.] We should not resolve the TPU name to a set of - # tasks and - # a. Create a ClusterSpec with the coordinator - # b. Create a ClusterSpec without the coordinator - # 3. [Other (legacy non-gRPC).] We should return an empty ClusterSpec. - ############################################################################ - - if self._shouldResolve(): - # Case 1. - full_name = 'projects/%s/locations/%s/nodes/%s' % ( - self._project, self._zone, compat.as_text(self._tpu)) - request = self._service.projects().locations().nodes().get(name=full_name) - response = request.execute() - - if 'state' in response and response['state'] != 'READY': - raise RuntimeError('TPU "%s" is not yet ready; state: "%s"' % - (compat.as_text(self._tpu), response['state'])) - - if 'health' in response and response['health'] != 'HEALTHY': - raise RuntimeError('TPU "%s" is unhealthy: "%s"' % - (compat.as_text(self._tpu), response['health'])) - - if 'networkEndpoints' in response: - worker_list = [ - '%s:%s' % (endpoint['ipAddress'], endpoint['port']) - for endpoint in response['networkEndpoints'] - ] - else: - # Fall back to the deprecated response format - instance_url = '%s:%s' % (response['ipAddress'], response['port']) - worker_list = [instance_url] - - cluster_spec = {self._job_name: worker_list} - else: - if not self._tpu.startswith(compat.as_bytes('grpc://')): - # Case 3. - return None - # Case 2. - cluster_spec = { - self._job_name: [ - x[len(compat.as_bytes('grpc://')):] - for x in self._tpu.split(compat.as_bytes(_ENDPOINTS_SEPARATOR)) - ] - } - - if self._coordinator_address: - # {1, 2}.a - cluster_spec[self._coordinator_name] = [self._coordinator_address] - - return server_lib.ClusterSpec(cluster_spec) - - def _start_local_server(self): - address = self._requestComputeMetadata('instance/network-interfaces/0/ip') - self._server = server_lib.Server( - { - 'local': ['0.0.0.0:0'] - }, protocol='grpc', config=None, start=True) - # self._server.target is of the form: grpc://ipaddress:port - target = compat.as_bytes(self._server.target) - splits = target.split(compat.as_bytes(':')) - assert len(splits) == 3, self._server.target - assert splits[0] == compat.as_bytes('grpc'), self._server.target - self._coordinator_port = compat.as_text(splits[2]) - self._coordinator_address = '%s:%s' % ( - address, compat.as_text(self._coordinator_port)) - - def __deepcopy__(self, memo): - # TODO(b/73668574): Remove this once RunConfig avoids performing deepcopy. - return self +remove_undocumented(__name__, _allowed_symbols) diff --git a/tensorflow/contrib/cluster_resolver/python/training/tpu_cluster_resolver_test.py b/tensorflow/contrib/cluster_resolver/python/training/tpu_cluster_resolver_test.py deleted file mode 100644 index ad4f6432630be44a7de6e778f55f1fb7fd66f307..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/cluster_resolver/python/training/tpu_cluster_resolver_test.py +++ /dev/null @@ -1,468 +0,0 @@ -# Copyright 2017 The TensorFlow Authors. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== -"""Tests for TPUClusterResolver.""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import os - -from tensorflow.contrib.cluster_resolver.python.training.tpu_cluster_resolver import TPUClusterResolver -from tensorflow.python.platform import test -from tensorflow.python.training import server_lib -from tensorflow.python.util import compat - -mock = test.mock - - -class MockRequestClass(object): - - def __init__(self, name, tpu_map): - self._name = name - self._tpu_map = tpu_map - - def execute(self): - if self._name in self._tpu_map: - return self._tpu_map[self._name] - else: - raise KeyError('Resource %s was not found' % self._name) - - -class MockNodeClass(object): - - def __init__(self, tpu_map): - self._tpu_map = tpu_map - - def get(self, name): - return MockRequestClass(name, self._tpu_map) - - -def mock_request_compute_metadata(cls, *args, **kwargs): - del cls, kwargs # Unused. - if args[0] == 'project/project-id': - return 'test-project' - elif args[0] == 'instance/zone': - return 'projects/test-project/locations/us-central1-c' - elif args[0] == 'instance/network-interfaces/0/ip': - return '10.128.1.2' - return '' - - -class TPUClusterResolverTest(test.TestCase): - - def _verifyClusterSpecEquality(self, cluster_spec, expected_proto): - """Verifies that the ClusterSpec generates the correct proto. - - We are testing this four different ways to ensure that the ClusterSpec - returned by the TPUClusterResolver behaves identically to a normal - ClusterSpec when passed into the generic ClusterSpec libraries. - - Args: - cluster_spec: ClusterSpec returned by the TPUClusterResolver - expected_proto: Expected protobuf - """ - self.assertProtoEquals(expected_proto, cluster_spec.as_cluster_def()) - self.assertProtoEquals( - expected_proto, - server_lib.ClusterSpec(cluster_spec).as_cluster_def()) - self.assertProtoEquals(expected_proto, - server_lib.ClusterSpec( - cluster_spec.as_cluster_def()).as_cluster_def()) - self.assertProtoEquals(expected_proto, - server_lib.ClusterSpec( - cluster_spec.as_dict()).as_cluster_def()) - - def mock_service_client(self, tpu_map=None): - - if tpu_map is None: - tpu_map = {} - - mock_locations = mock.MagicMock() - mock_locations.nodes.return_value = MockNodeClass(tpu_map) - - mock_project = mock.MagicMock() - mock_project.locations.return_value = mock_locations - - mock_client = mock.MagicMock() - mock_client.projects.return_value = mock_project - - return mock_client - - @mock.patch.object(TPUClusterResolver, '_requestComputeMetadata', - mock_request_compute_metadata) - def testRetrieveProjectAndZoneFromMetadata(self): - tpu_map = { - 'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': { - 'ipAddress': '10.1.2.3', - 'port': '8470', - 'health': 'HEALTHY' - } - } - - tpu_cluster_resolver = TPUClusterResolver( - project=None, - zone=None, - tpu=['test-tpu-1'], - credentials=None, - service=self.mock_service_client(tpu_map=tpu_map), - coordinator_name='coordinator') - - actual_cluster_spec = tpu_cluster_resolver.cluster_spec() - expected_proto = """ - job { - name: 'coordinator' - tasks { key: 0 value: '10.128.1.2:%s' } - } - job { - name: 'worker' - tasks { key: 0 value: '10.1.2.3:8470' } - } - """ % tpu_cluster_resolver._coordinator_port - self._verifyClusterSpecEquality(actual_cluster_spec, str(expected_proto)) - - @mock.patch.object(TPUClusterResolver, '_requestComputeMetadata', - mock_request_compute_metadata) - def testRetrieveProjectAndZoneFromMetadataNoCoordinator(self): - tpu_map = { - 'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': { - 'ipAddress': '10.1.2.3', - 'port': '8470', - 'health': 'HEALTHY' - } - } - - tpu_cluster_resolver = TPUClusterResolver( - project=None, - zone=None, - tpu=['test-tpu-1'], - coordinator_name=None, - credentials=None, - service=self.mock_service_client(tpu_map=tpu_map)) - - actual_cluster_spec = tpu_cluster_resolver.cluster_spec() - expected_proto = """ - job { name: 'worker' tasks { key: 0 value: '10.1.2.3:8470' } } - """ - self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto) - - @mock.patch.object(TPUClusterResolver, '_requestComputeMetadata', - mock_request_compute_metadata) - def testUnhealthyCloudTpu(self): - tpu_map = { - 'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': { - 'ipAddress': '10.1.2.3', - 'port': '8470', - 'health': 'UNHEALTHY' - } - } - - tpu_cluster_resolver = TPUClusterResolver( - project=None, - zone=None, - tpu='test-tpu-1', - coordinator_name=None, - credentials=None, - service=self.mock_service_client(tpu_map=tpu_map)) - - with self.assertRaises(RuntimeError): - tpu_cluster_resolver.cluster_spec() - - @mock.patch.object(TPUClusterResolver, '_requestComputeMetadata', - mock_request_compute_metadata) - def testNotReadyCloudTpu(self): - tpu_map = { - 'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': { - 'ipAddress': '10.1.2.3', - 'port': '8470', - 'state': 'CREATING' - } - } - - tpu_cluster_resolver = TPUClusterResolver( - project=None, - zone=None, - tpu='test-tpu-1', - coordinator_name=None, - credentials=None, - service=self.mock_service_client(tpu_map=tpu_map)) - - with self.assertRaises(RuntimeError): - tpu_cluster_resolver.cluster_spec() - - def testSimpleSuccessfulRetrieval(self): - tpu_map = { - 'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': { - 'ipAddress': '10.1.2.3', - 'port': '8470', - 'health': 'HEALTHY' - } - } - - tpu_cluster_resolver = TPUClusterResolver( - project='test-project', - zone='us-central1-c', - tpu=['test-tpu-1'], - coordinator_name='coordinator', - coordinator_address='10.128.1.5:10203', - credentials=None, - service=self.mock_service_client(tpu_map=tpu_map)) - - actual_cluster_spec = tpu_cluster_resolver.cluster_spec() - expected_proto = """ - job { name: 'coordinator' tasks { key: 0 value: '10.128.1.5:10203' } } - job { name: 'worker' tasks { key: 0 value: '10.1.2.3:8470' } } - """ - self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto) - - def testNewNetworkEndpointFormat(self): - tpu_map = { - 'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': { - 'health': 'HEALTHY', - 'networkEndpoints': [{ - 'ipAddress': '10.2.3.4', - 'port': 8470, - }] - } - } - - tpu_cluster_resolver = TPUClusterResolver( - project='test-project', - zone='us-central1-c', - tpu='test-tpu-1', - coordinator_name='coordinator', - coordinator_address='10.128.1.5:10203', - credentials=None, - service=self.mock_service_client(tpu_map=tpu_map)) - - actual_cluster_spec = tpu_cluster_resolver.cluster_spec() - expected_proto = """ - job { name: 'coordinator' tasks { key: 0 value: '10.128.1.5:10203' } } - job { name: 'worker' tasks { key: 0 value: '10.2.3.4:8470' } } - """ - self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto) - self.assertEqual('grpc://10.2.3.4:8470', tpu_cluster_resolver.master()) - - @mock.patch.object(TPUClusterResolver, '_requestComputeMetadata', - mock_request_compute_metadata) - def testPodResolution(self): - tpu_map = { - 'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': { - 'health': - 'HEALTHY', - 'networkEndpoints': [ - { - 'ipAddress': '10.2.3.4', - 'port': 8470, - }, - { - 'ipAddress': '10.2.3.5', - 'port': 8470, - }, - { - 'ipAddress': '10.2.3.6', - 'port': 8470, - }, - { - 'ipAddress': '10.2.3.7', - 'port': 8470, - }, - ] - } - } - - tpu_cluster_resolver = TPUClusterResolver( - tpu='test-tpu-1', - credentials=None, - service=self.mock_service_client(tpu_map=tpu_map), - coordinator_name='coordinator') - - actual_cluster_spec = tpu_cluster_resolver.cluster_spec() - expected_proto = """ - job { - name: 'coordinator', - tasks { key: 0 value: '10.128.1.2:%s'} - } - job { - name: 'worker' - tasks { key: 0 value: '10.2.3.4:8470' } - tasks { key: 1 value: '10.2.3.5:8470' } - tasks { key: 2 value: '10.2.3.6:8470' } - tasks { key: 3 value: '10.2.3.7:8470' } - } - """ % tpu_cluster_resolver._coordinator_port - self._verifyClusterSpecEquality(actual_cluster_spec, str(expected_proto)) - - def testPodResolutionNoCoordinator(self): - tpu_map = { - 'projects/test-project/locations/us-central1-c/nodes/test-tpu-1': { - 'health': - 'HEALTHY', - 'networkEndpoints': [ - { - 'ipAddress': '10.2.3.4', - 'port': 8470, - }, - { - 'ipAddress': '10.2.3.5', - 'port': 8470, - }, - { - 'ipAddress': '10.2.3.6', - 'port': 8470, - }, - { - 'ipAddress': '10.2.3.7', - 'port': 8470, - }, - ] - } - } - - tpu_cluster_resolver = TPUClusterResolver( - project='test-project', - zone='us-central1-c', - tpu='test-tpu-1', - coordinator_name=None, - credentials=None, - service=self.mock_service_client(tpu_map=tpu_map)) - - actual_cluster_spec = tpu_cluster_resolver.cluster_spec() - expected_proto = """ - job { - name: 'worker' - tasks { key: 0 value: '10.2.3.4:8470' } - tasks { key: 1 value: '10.2.3.5:8470' } - tasks { key: 2 value: '10.2.3.6:8470' } - tasks { key: 3 value: '10.2.3.7:8470' } - } - """ - self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto) - - def testGetMasterNoEntries(self): - tpu_map = {} - - with self.assertRaises(ValueError): - TPUClusterResolver( - project='test-project', - zone='us-central1-c', - tpu=[], - coordinator_name=None, - credentials=None, - service=self.mock_service_client(tpu_map=tpu_map)) - - # TODO(saeta): Convert to parameterized test when included in OSS TF. - def verifyShouldResolve(self, tpu, should_resolve): - tpu_cluster_resolver = TPUClusterResolver( - project='test-project', - zone='us-central1-c', - tpu=tpu, - coordinator_name=None, - credentials=None, - service=self.mock_service_client(tpu_map={})) - self.assertEqual(should_resolve, tpu_cluster_resolver._shouldResolve(), - "TPU: '%s'" % tpu) - - def testShouldResolveNoName(self): - self.verifyShouldResolve('', False) - - def testShouldResolveLocal(self): - self.verifyShouldResolve('local', False) - - def testShouldResolveGrpc(self): - self.verifyShouldResolve('grpc://10.1.2.3:8470', False) - - def testShouldResolveBns(self): - self.verifyShouldResolve('/bns/foo/bar', False) - - def testShouldResolveName(self): - self.verifyShouldResolve('mytpu', True) - - def testShouldResolveList(self): - self.verifyShouldResolve(['myothertpu'], True) - - def testShouldResolveGrpcPrefix(self): - self.verifyShouldResolve('grpctpu', True) - - def testNoCallComputeMetadata(self): - tpu_cluster_resolver = TPUClusterResolver(tpu='/bns/foo/bar') - self.assertEqual( - compat.as_bytes('/bns/foo/bar'), tpu_cluster_resolver.master()) - self.assertEqual(None, tpu_cluster_resolver.cluster_spec()) - - def testGkeEnvironmentForDonut(self): - os.environ['KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS'] = 'grpc://10.120.27.5:8470' - - self.assertIn('KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS', os.environ) - self.assertTrue(TPUClusterResolver._inGke()) - self.assertEqual( - compat.as_bytes('grpc://10.120.27.5:8470'), - compat.as_bytes(TPUClusterResolver._gkeEndpoints())) - - tpu_cluster_resolver = TPUClusterResolver() - self.assertEqual( - compat.as_bytes('grpc://10.120.27.5:8470'), - compat.as_bytes(tpu_cluster_resolver.master())) - actual_cluster_spec = tpu_cluster_resolver.cluster_spec() - expected_proto = """ - job { - name: 'worker' - tasks { key: 0 value: '10.120.27.5:8470' } - } - """ - self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto) - - del os.environ['KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS'] - - def testGkeEnvironmentForPod(self): - os.environ['KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS'] = ('grpc://10.120.27.5:8470,' - 'grpc://10.120.27.6:8470,' - 'grpc://10.120.27.7:8470,' - 'grpc://10.120.27.8:8470') - - self.assertIn('KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS', os.environ) - self.assertTrue(TPUClusterResolver._inGke()) - self.assertEqual( - compat.as_bytes('grpc://10.120.27.5:8470,' - 'grpc://10.120.27.6:8470,' - 'grpc://10.120.27.7:8470,' - 'grpc://10.120.27.8:8470'), - compat.as_bytes(TPUClusterResolver._gkeEndpoints())) - - tpu_cluster_resolver = TPUClusterResolver() - self.assertEqual( - compat.as_bytes('grpc://10.120.27.5:8470'), - compat.as_bytes(tpu_cluster_resolver.master())) - actual_cluster_spec = tpu_cluster_resolver.cluster_spec() - expected_proto = """ - job { - name: 'worker' - tasks { key: 0 value: '10.120.27.5:8470' } - tasks { key: 1 value: '10.120.27.6:8470' } - tasks { key: 2 value: '10.120.27.7:8470' } - tasks { key: 3 value: '10.120.27.8:8470' } - } - """ - self._verifyClusterSpecEquality(actual_cluster_spec, expected_proto) - - del os.environ['KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS'] - - def testDiscoveryUrl(self): - os.environ['TPU_API_DISCOVERY_URL'] = 'https://{api}.internal/{apiVersion}' - self.assertEqual('https://{api}.internal/{apiVersion}', - TPUClusterResolver._discoveryUrl()) - -if __name__ == '__main__': - test.main() diff --git a/tensorflow/contrib/cmake/CMakeLists.txt b/tensorflow/contrib/cmake/CMakeLists.txt index fbdca497fcc3126d2086d289ebdb113370072d22..2ad9ae42a16f690d38b8e2652e853012ec1dd267 100644 --- a/tensorflow/contrib/cmake/CMakeLists.txt +++ b/tensorflow/contrib/cmake/CMakeLists.txt @@ -3,16 +3,16 @@ cmake_minimum_required(VERSION 3.5) if(WIN32) if(${CMAKE_VERSION} VERSION_LESS "3.8") - message(WARNING "Your current cmake version is ${CMAKE_VERSION} which does not support setting the toolset architecture to x64. This may cause \"compiler out of heap space\" errors when building. Consider upgrading your cmake to > 3.8 and using the flag -Thost=x64 when running cmake.") + message(WARNING "Your current cmake version is ${CMAKE_VERSION} which does not support setting the toolset architecture to x64. This may cause \"compiler out of heap space\" errors when building. Consider upgrading your cmake to > 3.8 and using the flag -Thost=x64 when running cmake. Ignore this if you are on CMake GUI.") else() if(NOT CMAKE_VS_PLATFORM_TOOLSET_HOST_ARCHITECTURE OR NOT "${CMAKE_VS_PLATFORM_TOOLSET_HOST_ARCHITECTURE}" STREQUAL "x64") - message(WARNING "Your current cmake generator is set to use 32 bit toolset architecture. This may cause \"compiler out of heap space\" errors when building. Consider using the flag -Thost=x64 when running cmake.") + message(WARNING "Your current cmake generator is set to use 32 bit toolset architecture. This may cause \"compiler out of heap space\" errors when building. Consider using the flag -Thost=x64 when running cmake. Ignore this if you are on CMake GUI.") endif() endif() endif() # Project -project(tensorflow C CXX) +project(tensorflow VERSION 1.12.0 LANGUAGES C CXX) # Set C++14 as standard for the whole project set(CMAKE_CXX_STANDARD 14) @@ -52,15 +52,19 @@ option(tensorflow_OPTIMIZE_FOR_NATIVE_ARCH "Enable compiler optimizations for th option(tensorflow_ENABLE_SNAPPY_SUPPORT "Enable SNAPPY compression support" ON) option(tensorflow_DISABLE_EIGEN_FORCEINLINE "Disable forceinline, to speed up build on windows." OFF) +if (WIN32) +SET(tensorflow_WIN_CPU_SIMD_OPTIONS "/arch:AVX" CACHE STRING "Enables CPU SIMD instructions") +SET_PROPERTY(CACHE tensorflow_WIN_CPU_SIMD_OPTIONS PROPERTY STRINGS /arch:AVX) +endif() + # SIMD, MKL and MKLDNN options option(tensorflow_WIN_CPU_SIMD_OPTIONS "Enables CPU SIMD instructions" OFF) option(tensorflow_ENABLE_MKL_SUPPORT "Enable Intel MKL support" OFF) option(tensorflow_ENABLE_MKLDNN_SUPPORT "Enable Intel MKLDNN support, requires MKL enabled" OFF) + # GPU, CUDA and cuDNN options option(tensorflow_ENABLE_GPU "Enable GPU support" OFF) -set(tensorflow_CUDA_VERSION "9.0" CACHE STRING "CUDA version to build against") -set(tensorflow_CUDNN_VERSION "7" CACHE STRING "cuDNN version to build against") if(HAIKU) option(tensorflow_ENABLE_POSITION_INDEPENDENT_CODE "Enable PIE support" OFF) @@ -72,25 +76,30 @@ endif() if (NOT WIN32) # Threads: defines CMAKE_THREAD_LIBS_INIT and adds -pthread compile option # for targets that link ${CMAKE_THREAD_LIBS_INIT}. - find_package (Threads) + find_package (Threads REQUIRED) # Options for linking CUDA/CUDNN libraries - option(tensorflow_PATH_STATIC_LIB "Additional library search path for libcudnn_static.a, libnccl_static.a, libculibos.a" /usr/local/cuda/lib64/) + option(tensorflow_PATH_CUDA_LIB "Additional library search path for cudnn, nccl, culibos" /usr/local/cuda/lib64/) option(tensorflow_CUDNN_INCLUDE "cudnn.h header install path" /usr/include/) if (NOT tensorflow_CUDNN_INCLUDE) # option's default value is OFF. Fill it with real default values set(tensorflow_CUDNN_INCLUDE /usr/include) endif (NOT tensorflow_CUDNN_INCLUDE) - option(tensorflow_PATH_CUDNN_STATIC_LIB "Override PATH_STATIC_LIB for libcudnn_static.a" ${tensorflow_PATH_STATIC_LIB}) - if (NOT tensorflow_PATH_CUDNN_STATIC_LIB) + option(tensorflow_NCCL_INCLUDE "nccl.h header install path" /usr/include/) + if (NOT tensorflow_NCCL_INCLUDE) # option's default value is OFF. Fill it with real default values - set (tensorflow_PATH_CUDNN_STATIC_LIB ${tensorflow_PATH_STATIC_LIB}) - endif (NOT tensorflow_PATH_CUDNN_STATIC_LIB) - option(tensorflow_PATH_NCCL_STATIC_LIB "Override PATH_STATIC_LIB for libnccl_static.a" ${tensorflow_PATH_STATIC_LIB}) - if (NOT tensorflow_PATH_NCCL_STATIC_LIB) + set(tensorflow_NCCL_INCLUDE /usr/include) + endif (NOT tensorflow_NCCL_INCLUDE) + option(tensorflow_PATH_CUDNN_LIB "Override PATH_CUDA_LIB for cudnn" ${tensorflow_PATH_CUDA_LIB}) + if (NOT tensorflow_PATH_CUDNN_LIB) # option's default value is OFF. Fill it with real default values - set (tensorflow_PATH_NCCL_STATIC_LIB ${tensorflow_PATH_STATIC_LIB}) - endif (NOT tensorflow_PATH_NCCL_STATIC_LIB) + set (tensorflow_PATH_CUDNN_LIB ${tensorflow_PATH_CUDA_LIB}) + endif (NOT tensorflow_PATH_CUDNN_LIB) + option(tensorflow_PATH_NCCL_LIB "Override PATH_CUDA_LIB for nccl" ${tensorflow_PATH_CUDA_LIB}) + if (NOT tensorflow_PATH_NCCL_LIB) + # option's default value is OFF. Fill it with real default values + set (tensorflow_PATH_NCCL_LIB ${tensorflow_PATH_CUDA_LIB}) + endif (NOT tensorflow_PATH_NCCL_LIB) option(tensorflow_CUDA_LIBRARY_PATH "Designate the default CUDA library paths" /usr/local/cuda/lib64) if (NOT tensorflow_CUDA_LIBRARY_PATH) # option's default value is OFF. Fill it with real default values @@ -195,6 +204,7 @@ if(WIN32) set(CMAKE_SUPPRESS_REGENERATION ON) endif() + if ("${CMAKE_CXX_COMPILER_ID}" STREQUAL "GNU") set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fno-exceptions -std=c++11") endif() @@ -210,14 +220,17 @@ endif() include(CheckCXXCompilerFlag) # OpenMP Support -CHECK_CXX_COMPILER_FLAG("-fopenmp" GCC_OPENMP_SUPPORT) -if (GCC_OPENMP_SUPPORT) - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fopenmp") -endif() -CHECK_CXX_COMPILER_FLAG("/openmp" MSVC_OPENMP_SUPPORT) -if (MSVC_OPENMP_SUPPORT) - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /openmp") -endif() +if (WIN32) + CHECK_CXX_COMPILER_FLAG("/openmp" MSVC_OPENMP_SUPPORT) + if (MSVC_OPENMP_SUPPORT) + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /openmp") + endif() +else (WIN32) + CHECK_CXX_COMPILER_FLAG("-fopenmp" GCC_OPENMP_SUPPORT) + if (GCC_OPENMP_SUPPORT) + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fopenmp") + endif() +endif (WIN32) # MSVC SIMD instructions if (tensorflow_WIN_CPU_SIMD_OPTIONS) @@ -280,6 +293,14 @@ else (systemlib_ZLIB) ${zlib_STATIC_LIBRARIES}) endif (systemlib_ZLIB) +if (systemlib_ABSEIL_CPP) + set(tensorflow_EXTERNAL_LIBRARIES ${tensorflow_EXTERNAL_LIBRARIES} + ${abseil_cpp_LIBRARIES}) +else (systemlib_ABSEIL_CPP) + set(tensorflow_EXTERNAL_LIBRARIES ${tensorflow_EXTERNAL_LIBRARIES} + ${abseil_cpp_STATIC_LIBRARIES}) +endif (systemlib_ABSEIL_CPP) + set(tensorflow_EXTERNAL_DEPENDENCIES zlib_copy_headers_to_destination gif_copy_headers_to_destination @@ -377,32 +398,23 @@ if (tensorflow_ENABLE_GPU) list(APPEND CMAKE_LIBRARY_PATH "${tensorflow_CUDA_LIBRARY_PATH}/stubs") endif (NOT WIN32) - # later command will make use of the value in tensorflow_CUDA_VERSION - find_package(CUDA ${tensorflow_CUDA_VERSION} REQUIRED EXACT) - - # Test compatibility of compiler on CUDA - try_compile(CUDA_TEST_COMPILE_C - ${CMAKE_CURRENT_BINARY_DIR}/tests/cuda - ${CMAKE_CURRENT_SOURCE_DIR}/tests/cuda/compatibility_test.c - CMAKE_FLAGS -DINCLUDE_DIRECTORIES=${CUDA_INCLUDE_DIRS}) - try_compile(CUDA_TEST_COMPILE_CXX - ${CMAKE_CURRENT_BINARY_DIR}/tests/cuda - ${CMAKE_CURRENT_SOURCE_DIR}/tests/cuda/compatibility_test.cc - CMAKE_FLAGS -DINCLUDE_DIRECTORIES=${CUDA_INCLUDE_DIRS}) - if(NOT (CUDA_TEST_COMPILE_C AND CUDA_TEST_COMPILE_CXX)) - message(FATAL_ERROR "Selected compiler (or version) is not supported for CUDA") + # minimum 9.0 in cuda version + find_package(CUDA 9.0 REQUIRED) + if(NOT CUDA_FOUND) + message(FATAL_ERROR "CUDA not found.") endif() - # by default we assume compute cabability 3.5 and 5.2. If you change this change it in - # CUDA_NVCC_FLAGS and cuda_config.h below - set(CUDA_NVCC_FLAGS ${CUDA_NVCC_FLAGS};-gencode arch=compute_37,code=\"sm_37,compute_37\") - set(CUDA_NVCC_FLAGS ${CUDA_NVCC_FLAGS};-gencode arch=compute_52,code=\"sm_52,compute_52\") - set(CUDA_NVCC_FLAGS ${CUDA_NVCC_FLAGS};-gencode arch=compute_60,code=\"sm_60,compute_60\") - set(CUDA_NVCC_FLAGS ${CUDA_NVCC_FLAGS};-gencode arch=compute_61,code=\"sm_61,compute_61\") - set(CUDA_NVCC_FLAGS ${CUDA_NVCC_FLAGS};-gencode arch=compute_70,code=\"sm_70,compute_70\") + # use cmake internal CUDA_ARCH_NAME switch + # e.g. CUDA_ARCH_NAME="Auto" will autodetect + # CUDA_ARCH_NAME="All" will use all arches + cuda_select_nvcc_arch_flags(NVCC_ARCH_FLAGS ${CUDA_ARCH_NAME}) + list(APPEND CUDA_NVCC_FLAGS ${NVCC_ARCH_FLAGS}) + message(STATUS "Using CUDA arch flags: ${NVCC_ARCH_FLAGS_readable}") + set(CUDA_NVCC_FLAGS ${CUDA_NVCC_FLAGS};--include-path ${PROJECT_BINARY_DIR}/$\{build_configuration\};--expt-relaxed-constexpr) set(CUDA_NVCC_FLAGS ${CUDA_NVCC_FLAGS};-ftz=true) # Flush denormals to zero set(CUDA_INCLUDE ${CUDA_TOOLKIT_TARGET_DIR} ${CUDA_TOOLKIT_TARGET_DIR}/extras/CUPTI/include) + include_directories(${CUDA_INCLUDE}) if (WIN32) add_definitions(-DGOOGLE_CUDA=1 -DTF_EXTRA_CUDA_CAPABILITIES=3.7,5.2,6.0,6.1,7.0) @@ -423,43 +435,94 @@ if (tensorflow_ENABLE_GPU) else (WIN32) set(CUDNN_INCLUDE "${tensorflow_CUDNN_INCLUDE}") - find_library(nccl_STATIC_LIBRARY NAMES libnccl_static.a PATHS ${tensorflow_PATH_NCCL_STATIC_LIB} ${CUDA_TOOLKIT_ROOT_DIR}) - if (NOT nccl_STATIC_LIBRARY) + if (tensorflow_BUILD_SHARED_LIB) + find_library(nccl_LIBRARY NAMES libnccl.so PATHS ${tensorflow_PATH_NCCL_LIB} ${CUDA_TOOLKIT_ROOT_DIR}) + else (tensorflow_BUILD_SHARED_LIB) + find_library(nccl_LIBRARY NAMES libnccl_static.a PATHS ${tensorflow_PATH_NCCL_LIB} ${CUDA_TOOLKIT_ROOT_DIR}) + endif (tensorflow_BUILD_SHARED_LIB) + if (NOT nccl_LIBRARY) message(FATAL_ERROR "NCCL is required for GPU-build") - else (NOT nccl_STATIC_LIBRARY) - message("nccl-static: ${nccl_STATIC_LIBRARY}") + else (NOT nccl_LIBRARY) + message("nccl: ${nccl_LIBRARY}") # something like /usr/lib64/libnccl_static.a - endif (NOT nccl_STATIC_LIBRARY) - - find_library(cudnn_STATIC_LIBRARY NAMES libcudnn_static.a PATHS ${tensorflow_PATH_CUDNN_STATIC_LIB} ${CUDA_TOOLKIT_ROOT_DIR}) - if (NOT cudnn_STATIC_LIBRARY) + endif (NOT nccl_LIBRARY) + + if (tensorflow_BUILD_SHARED_LIB) + find_library(cudnn_LIBRARY NAMES libcudnn.so PATHS ${tensorflow_PATH_CUDNN_LIB} ${CUDA_TOOLKIT_ROOT_DIR}) + else (tensorflow_BUILD_SHARED_LIB) + find_library(cudnn_LIBRARY NAMES libcudnn_static.a PATHS ${tensorflow_PATH_CUDNN_LIB} ${CUDA_TOOLKIT_ROOT_DIR}) + endif (tensorflow_BUILD_SHARED_LIB) + if (NOT cudnn_LIBRARY) message(FATAL_ERROR "CUDNN is required for GPU-build") - else (NOT cudnn_STATIC_LIBRARY) - message("cudnn-static: ${cudnn_STATIC_LIBRARY}") - endif (NOT cudnn_STATIC_LIBRARY) - - find_library(culibos_STATIC_LIBRARY NAMES libculibos.a PATHS ${tensorflow_PATH_STATIC_LIB} ${CUDA_TOOLKIT_ROOT_DIR}) - if (NOT culibos_STATIC_LIBRARY) + else (NOT cudnn_LIBRARY) + file(READ ${CUDNN_INCLUDE}/cudnn.h CUDNN_VERSION_FILE_CONTENTS) + # fetch cudnn version + string(REGEX MATCH "define CUDNN_MAJOR * +([0-9]+)" + CUDNN_VERSION_MAJOR "${CUDNN_VERSION_FILE_CONTENTS}") + string(REGEX REPLACE "define CUDNN_MAJOR * +([0-9]+)" "\\1" + CUDNN_VERSION_MAJOR "${CUDNN_VERSION_MAJOR}") + string(REGEX MATCH "define CUDNN_MINOR * +([0-9]+)" + CUDNN_VERSION_MINOR "${CUDNN_VERSION_FILE_CONTENTS}") + string(REGEX REPLACE "define CUDNN_MINOR * +([0-9]+)" "\\1" + CUDNN_VERSION_MINOR "${CUDNN_VERSION_MINOR}") + string(REGEX MATCH "define CUDNN_PATCHLEVEL * +([0-9]+)" + CUDNN_VERSION_PATCH "${CUDNN_VERSION_FILE_CONTENTS}") + string(REGEX REPLACE "define CUDNN_PATCHLEVEL * +([0-9]+)" "\\1" + CUDNN_VERSION_PATCH "${CUDNN_VERSION_PATCH}") + if(NOT CUDNN_VERSION_MAJOR) + set(CUDNN_VERSION "???") + else() + set(CUDNN_VERSION "${CUDNN_VERSION_MAJOR}.${CUDNN_VERSION_MINOR}.${CUDNN_VERSION_PATCH}") + endif() + message(STATUS "cudnn library: ${cudnn_LIBRARY} (found version: \"${CUDNN_VERSION}\")") + endif (NOT cudnn_LIBRARY) + + if (tensorflow_BUILD_SHARED_LIB) + # shared first (if exists) else static one + find_library(culibos_LIBRARY NAMES libculibos.so libculibos.a PATHS ${tensorflow_PATH_CUDA_LIB} ${CUDA_TOOLKIT_ROOT_DIR}) + else (tensorflow_BUILD_SHARED_LIB) + # only static version + find_library(culibos_LIBRARY NAMES libculibos.a PATHS ${tensorflow_PATH_CUDA_LIB} ${CUDA_TOOLKIT_ROOT_DIR}) + endif (tensorflow_BUILD_SHARED_LIB) + if (NOT culibos_LIBRARY) message(FATAL_ERROR "CULIBOS is required for GPU-build") - else (NOT culibos_STATIC_LIBRARY) - message("culibos-static: ${culibos_STATIC_LIBRARY}") - endif (NOT culibos_STATIC_LIBRARY) + else (NOT culibos_LIBRARY) + message("culibos: ${culibos_LIBRARY}") + endif (NOT culibos_LIBRARY) set(CUDA_LIBRARIES ${CUDA_LIBRARIES} ${CUDA_CUDA_LIBRARY} ${CUDA_CUBLAS_LIBRARIES} ${CUDA_CUFFT_LIBRARIES} - ${CUDA_curand_LIBRARY} ${CUDA_cupti_LIBRARY} ${CUDA_cusolver_LIBRARY} ${cudnn_STATIC_LIBRARY} ${culibos_STATIC_LIBRARY} ${nccl_STATIC_LIBRARY}) + ${CUDA_curand_LIBRARY} ${CUDA_cupti_LIBRARY} ${CUDA_cusolver_LIBRARY} ${cudnn_LIBRARY} ${culibos_LIBRARY} ${nccl_LIBRARY}) endif (WIN32) include_directories(${CUDNN_INCLUDE}) # Remove "." from CUDA version variable. - string(REPLACE "." "" short_CUDA_VER ${tensorflow_CUDA_VERSION}) + string(REPLACE "." "" short_CUDA_VER ${CUDA_VERSION}) + + # List of enumerated CUDA caps + string(REPLACE " " ";" NVCC_ARCH_LIST "${NVCC_ARCH_FLAGS_readable}") + set(list ${NVCC_ARCH_LIST}) + + # Construct capability string + foreach(NVCC_ARCH ${NVCC_ARCH_LIST}) + if (NVCC_ARCH MATCHES "sm_") + string(REGEX REPLACE "^.sm*" "" NVCC_ARCH ${NVCC_ARCH}) + math(EXPR NVCC_ARCH_MAJOR "${NVCC_ARCH} / 10") + math(EXPR NVCC_ARCH_MINOR "(${NVCC_ARCH} - (${NVCC_ARCH_MAJOR}*10))") + if (TF_CUDA_CAP) + set(TF_CUDA_CAP "${TF_CUDA_CAP},CudaVersion(\"${NVCC_ARCH_MAJOR}.${NVCC_ARCH_MINOR}\")") + else (TF_CUDA_CAP) + set(TF_CUDA_CAP "CudaVersion(\"${NVCC_ARCH_MAJOR}.${NVCC_ARCH_MINOR}\")") + endif (TF_CUDA_CAP) + endif() + endforeach() # create cuda_config.h FILE(WRITE ${tensorflow_source_dir}/third_party/gpus/cuda/cuda_config.h "#ifndef CUDA_CUDA_CONFIG_H_\n" "#define CUDA_CUDA_CONFIG_H_\n" - "#define TF_CUDA_CAPABILITIES CudaVersion(\"3.7\"),CudaVersion(\"5.2\"),CudaVersion(\"6.0\"),CudaVersion(\"6.1\"),CudaVersion(\"7.0\")\n" + "#define TF_CUDA_CAPABILITIES ${TF_CUDA_CAP}\n" "#define TF_CUDA_VERSION \"64_${short_CUDA_VER}\"\n" - "#define TF_CUDNN_VERSION \"64_${tensorflow_CUDNN_VERSION}\"\n" + "#define TF_CUDNN_VERSION \"64_${CUDNN_VERSION}\"\n" "#define TF_CUDA_TOOLKIT_PATH \"${CUDA_TOOLKIT_ROOT_DIR}\"\n" "#endif // CUDA_CUDA_CONFIG_H_\n" ) @@ -494,24 +557,30 @@ if (tensorflow_ENABLE_GPU) set(tensorflow_BUILD_INFO_FLAGS --build_config cuda --key_value msvcp_dll_name=msvcp140.dll cudart_dll_name=cudart64_${short_CUDA_VER}.dll - cuda_version_number=${tensorflow_CUDA_VERSION} + cuda_version_number=${CUDA_VERSION} nvcuda_dll_name=nvcuda.dll cudnn_dll_name=cudnn64_${tensorflow_CUDNN_VERSION}.dll cudnn_version_number=${tensorflow_CUDNN_VERSION}) else(WIN32) set(tensorflow_BUILD_INFO_FLAGS --build_config cuda --key_value - cuda_version_number=${tensorflow_CUDA_VERSION} - cudnn_version_number=${tensorflow_CUDNN_VERSION}) + cuda_version_number=${CUDA_VERSION} + cudnn_version_number=${tensorflow_CUDNN_VERSION}) endif(WIN32) else(tensorflow_ENABLE_GPU) - set(tensorflow_BUILD_INFO_FLAGS --build_config cpu --key_value - msvcp_dll_name=msvcp140.dll) + if(WIN32) + set(tensorflow_BUILD_INFO_FLAGS --build_config cpu --key_value + msvcp_dll_name=msvcp140.dll) + else() + set(tensorflow_BUILD_INFO_FLAGS --build_config cpu) + endif() endif(tensorflow_ENABLE_GPU) -# Find python executable -include(FindPythonInterp) -if(NOT ${PYTHONINTERP_FOUND}) - message(FATAL_ERROR "CMake was unable to find a python interpreter.") +if(tensorflow_BUILD_PYTHON_BINDINGS) + # Find python executable + include(FindPythonInterp) + if(NOT ${PYTHONINTERP_FOUND}) + message(FATAL_ERROR "CMake was unable to find a python interpreter.") + endif() endif() # Let's get to work! @@ -532,6 +601,7 @@ include(tf_cc_ops.cmake) include(tf_c.cmake) include(tf_grappler.cmake) include(tf_core_profiler.cmake) +include(tf_core_eager_runtime.cmake) if(tensorflow_BUILD_CC_EXAMPLE) include(tf_tutorials.cmake) include(tf_label_image_example.cmake) @@ -545,4 +615,4 @@ if(tensorflow_BUILD_SHARED_LIB) endif() if(tensorflow_BUILD_CC_TESTS OR tensorflow_BUILD_PYTHON_TESTS) include(tf_tests.cmake) -endif() +endif() \ No newline at end of file diff --git a/tensorflow/contrib/cmake/README.md b/tensorflow/contrib/cmake/README.md index 84c679162c3ed8ffc9babcd3af583b26fb62c2d6..60ee1b4b3fd7d0b6afaefcc05effd3bbae00cf2c 100644 --- a/tensorflow/contrib/cmake/README.md +++ b/tensorflow/contrib/cmake/README.md @@ -5,10 +5,10 @@ CMAKE build is deprecated for TensorFlow. Please use `bazel` to build TF for all platforms. For details, see the [TensorFlow install guide](https://www.tensorflow.org/install/). -This directory contains CMake files for building TensorFlow on Microsoft -Windows. [CMake](https://cmake.org) is a cross-platform tool that can -generate build scripts for multiple build systems, including Microsoft -Visual Studio. +This directory contains CMake files for building TensorFlow on Microsoft Windows +and Linux. [CMake](https://cmake.org) is a cross-platform tool that can generate +build scripts for multiple build systems, including Microsoft Visual Studio and +GCC. "The method has not been tested on Mac OS X. **N.B.** We provide Linux build instructions primarily for the purpose of testing the build. We recommend using the standard Bazel-based build on @@ -17,12 +17,17 @@ Linux. Current Status -------------- -CMake can be used to build TensorFlow on Windows. See the [getting started documentation](https://www.tensorflow.org/install/source_windows) -for instructions on how to install a pre-built TensorFlow package on Windows. +CMake can be used to build TensorFlow on all platforms. See the +[getting started documentation](https://www.tensorflow.org/install/install_windows) +for instructions on how to install a pre-built TensorFlow package on Windows and +Linux. The procedure in MacOS is similar to the Linux build. ### Current known limitations -* It is not possible to load a custom Op library. -* GCS file system is not supported. + +* It is not possible to load a custom Op library. +* GCS file system is not supported. +* Debug build is not available since Python for Windows is no longer + distributed with a debug library. ## Building with CMake @@ -32,70 +37,88 @@ bindings. ### Prerequisites -* CMake version 3.5 or later. +* CMake version 3.5 or later. + +* [Git](https://git-scm.com) + +* [SWIG](http://www.swig.org/download.html) + +* [Perl](https://www.perl.org/get.html) (optional, for SSL support build) + +* [Go](https://golang.org/) (optional, for SSL support build) + +* [NASM](http://www.nasm.us/)/[YASM](http://yasm.tortall.net/) (optional, for + SSL support build) + +* Additional pre-requisites for Microsoft Windows: + + - Visual Studio 2015 (latest version of MSVC 2017 is not supported by CUDA + yet, try it on your own risk) -* [Git](https://git-scm.com) + - Python 3.5 -* [SWIG](http://www.swig.org/download.html) +* Additional prerequisites for Linux: -* Additional prerequisites for Microsoft Windows: - - Visual Studio 2015 - - Python 3.5 + - Python 2.7 or later + - [Docker](https://www.docker.com/) (for automated testing) -* Additional prerequisites for Linux: - - Python 2.7 or later - - [Docker](https://www.docker.com/) (for automated testing) +* Python dependencies: -* Python dependencies: - - wheel - - NumPy 1.11.0 or later + - wheel + - NumPy 1.11.0 or later ### Known-good configurations -* Microsoft Windows 10 - - Microsoft Visual Studio Enterprise 2015 with Visual C++ 2015 - - [Anaconda 4.1.1 (Python 3.5 64-bit)](https://www.anaconda.com/download/) - - [Git for Windows version 2.9.2.windows.1](https://git-scm.com/download/win) - - [swigwin-3.0.10](http://www.swig.org/download.html) - - [NVidia CUDA Toolkit 8.0](https://developer.nvidia.com/cuda-downloads) - - [NVidia CUDNN 5.1](https://developer.nvidia.com/cudnn) - - [CMake 3.6](https://cmake.org/files/v3.6/cmake-3.6.3-win64-x64.msi) +* Microsoft Windows 10 -* Ubuntu 14.04 - - Makefile generator - - Docker 1.9.1 (for automated testing) + - Microsoft Visual Studio Enterprise/ Community 2015 with Visual C++ 2015 + - [Anaconda 4.1.1 (Python 3.5 64-bit)](https://www.anaconda.com/download/) + - [Git for Windows version 2.9.2.windows.1](https://git-scm.com/download/win) + - [swigwin-3.0.10](http://www.swig.org/download.html) + - [NVidia CUDA Toolkit 9.0](https://developer.nvidia.com/cuda-downloads) + - [NVidia CUDNN 7](https://developer.nvidia.com/cudnn) + - [CMake 3.6](https://cmake.org/files/v3.6/cmake-3.6.3-win64-x64.msi) + +* Ubuntu 14.04 + + - Makefile generator + - Docker 1.9.1 (for automated testing) ### Current known limitations - - The Python package supports **Python 3.5 only**, because that is the only - version for which standard Python binaries exist and those binaries are - compatible with the TensorFlow runtime. (On Windows, the standard Python + +- The Python package supports **Python 3.5/3.6 only**, because these are the + only versions for which standard Python binaries exist and those binaries + are compatible with the TensorFlow runtime. (On Windows, the standard Python binaries for versions earlier than 3.5 were compiled with older compilers that do not have all of the features (e.g. C++11 support) needed to compile - TensorFlow. We welcome patches for making TensorFlow work with Python 2.7 - on Windows, but have not yet committed to supporting that configuration.) - - - The following Python APIs are not currently implemented: - * Loading custom op libraries via `tf.load_op_library()`. In order to use your - custom op, please put the source code under the tensorflow/core/user_ops - directory, and a shape function is required (not optional) for each op. - * Path manipulation functions (such as `tf.gfile.ListDirectory()`) are not - functional. - - - The `tf.contrib` libraries are not currently included in the PIP package. - - - The following operations are not currently implemented: - * `DepthwiseConv2dNative` - * `Digamma` - * `Erf` - * `Erfc` - * `Igamma` - * `Igammac` - * `ImmutableConst` - * `Lgamma` - * `Polygamma` - * `Zeta` - - - Google Cloud Storage support is not currently implemented. The GCS library + TensorFlow. We welcome patches for making TensorFlow work with Python 2.7 on + Windows, but have not yet committed to supporting that configuration.) + +- The following Python APIs are not currently implemented: + + * Loading custom op libraries via `tf.load_op_library()`. In order to use + your custom op, please put the source code under the + tensorflow/core/user_ops directory, and a shape function is required + (not optional) for each op. + * Path manipulation functions (such as `tf.gfile.ListDirectory()`) are not + functional. + +- The `tf.contrib` libraries are not currently included in the PIP package. + +- The following operations are not currently implemented: + + * `DepthwiseConv2dNative` + * `Digamma` + * `Erf` + * `Erfc` + * `Igamma` + * `Igammac` + * `ImmutableConst` + * `Lgamma` + * `Polygamma` + * `Zeta` + +- Google Cloud Storage support is not currently implemented. The GCS library currently depends on `libcurl` and `boringssl`, and the Windows version could use standard Windows APIs for making HTTP requests and cryptography (for OAuth). Contributions are welcome for this feature. @@ -104,9 +127,211 @@ We are actively working on improving CMake and Windows support, and addressing these limitations. We would appreciate pull requests that implement missing ops or APIs. +# CMake GUI build (all platforms) + +Install from CMake GUI would be a convenient way to generate C++ build projects. +The software supports Windows, MacOS and Linux, while the posix platform +provides an extra ccmake binary to run command line GUI. Both working principal +of cmake, ccmake and cmake-gui are the same, the only difference is by providing +suitable interface for project configuration and dependency setting. + +1. Pre-buid checklist: The following binary/libraries should be setted in + system path, otherwise you need to set manualy via cmake. + * Compiler (GCC for Linux, MSVC for Windows) + * Make sure compiler directory has been set to system path + * CUDA 9.0 (GPU build) + * CUDNN (GPU build) + * NCCL (GPU build on Linux) + * SWIG (python binding) + * Perl (required if you need ssl support, optional) + * 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 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 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 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 + affect tensorflow function, turn it to `off` if you want a slim build. + (optional) + * `tensorflow_BUILD_PYTHON_BINDING` is default to be `on`. Set to `off` if + you don't need python interaface. If SWIG is not in system path, you + need set it manually. (optional) + * `tensorflow_BUILD_SHARED_LIB` is default to be `off`. Set to `on` if you + want the c++ interface. (optional) + * `tensorflow_ENABLE_GPU` is default to be `off`. Set to `on` if you want + GPU support. It will search CUDA and CUDNN dependecies if you have set + them to system path, otherwise CMake would prompt error and request you + to set it manually. (optional) + * `tensorflow_ENABLE_GRPC_SUPPORT` is default to be `on`. For Linux build, + this option must always be `on`. This need to be `on` for a gpu build. + Reminded that Perl, Go and NASM/YASM are required for this option if you + want to build grpc with offical SSL support. + * `tensorflow_ENABLE_POSITION_INDEPENDENT_CODE` should always be `on` + * `tensorflow_ENABLE_SNAPPY_SUPPORT` should always be `on` + * `tensorflow_OPTIMIZE_FOR_NATIVE_ARCH` should always be `on` + * `CMAKE_INSTALL_PREFIX` is the location where the final package will be + installed. You may change it to your own preferred path (optional) + +7. After changing the configuration in step 5, press `Configure` again + +8. If not error is found, press `Generate` + +#### Windows + +1. Open `tensorflow.sln` in the build folder (Windows). Change build type from + `Debug` to `Release`. Choose `Build`->`Build Solution`. This may take more + than hours of compilation. If everything is alright, the output window would + show no error. + + ##### Python + + In solution explorer, right click on `tf_python_build_pip_package` -> + `build`. It will generate the wheel file in + `/tf_python/dist`. Install with following command: + + `pip install --upgrade tensorflow-.whl` + + ***The wheel name varies depends on you config. Change to your own wheel + filename.*** + + Reminded that some pip installation requires administrator right command + prompt. + + ##### C++ + + You can directly use the build folder tree for C++ interface with cmake. If + you want to do installation for api releasing, right click on `Install` -> + `build`. The headers and library will be installed in the directory specify + by `CMAKE_INSTALL_PREFIX` during configuration. + +1. For smaller RAM computer, it is noticed that out of heap space error + appears. Change to command prompt build is an alternative to do step 1. + + Open `VS2015 x64 Native Tools Command Prompt`. You can open it by press + `Start`, then type the binary name. Use `VS2017 x64 Native Tools Command + Prompt` if you are using MSVC 2017. + + ##### Python + + Directly build python wheel package by following command: + + `MSBuild /p:Configuration=Release + ` + + Remember to change `` to the + actual path of the file, it can be found at the root of build directory + + Install the wheel file generated as instructed by step 1. + + ##### C++ interface + + Build from VS native toolchain with following command: `MSBuild + /p:Configuration=Release ` + + Headers are discretely located in the build folders. Tensorflow library can + be found at `/Release`, namely `tensorflow.dll` and + `tensorflow.lib`. + + * Build to install for api release (optional): `MSBuild + /p:Configuration=Release ` + + Remember to change `` and + `` to the actual path of the file, it can be found + at the root of build directory. + +#### Linux/MacOS (command line GNU build) + +1. Open the terminal, change working directory to the one specified in step 3. + +2. Type the following command: + + `make -sj all` + + ##### Python + + **Important Note** CMake generated python wheel for Linux/MacOs is currently + under development. Please use bazel build. + + Follow code is an expected Linux/MacOS python package build after + development work is completed. + + ``` + make -sj tf_python_build_pip_package + cd tf_python + pip install --upgrade tensorflow-.whl + ``` + + ##### C++ interface + + `make -sj install` + + Where `` is the threads used for the compilation, change + 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 + `tensorflow.dylib` (MacOS). + +#### Start a Tensorflow C++ project with CMake + +Here we assume that you have basic knowledge on gathering dependency with +`CMakeLists.txt`. Here we introduce how the C++ api works with +[official hello world tutorial](https://www.tensorflow.org/api_guides/cc/guide). + +1. Create a new working directory and create a new text file named + `CMakeLists.txt` and the c++ file `main.cxx` +2. Fill in the `main.cxx` with the code provided in + [official c++ api basic](https://www.tensorflow.org/api_guides/cc/guide). +3. Fill in the `CMakeLists.txt` with following code: ``` cmake + cmake_minimum_required (VERSION 2.6) project (tf_hello) + + # Tensorflow + + find_package(Tensorflow REQUIRED) + include_directories(${TENSORFLOW_INCLUDE_DIRS}) + + # compiler setting required by tensorflow, to be tested on all compilers + + # currently only tested on MSVC and GCC + + if (${CMAKE_CXX_COMPILER_ID} STREQUAL MSVC) add_definitions(-DCOMPILER_MSVC) + elseif (${CMAKE_CXX_COMPILER_ID} STREQUAL GNU) if + (${CMAKE_CXX_COMPILER_VERSION} VERSION_LESS "3") + add_definitions(-DCOMPILER_GCC3) else() add_definitions(-D__GNUC__) endif() + else() message(ERROR " compiler ${CMAKE_CXX_COMPILER_ID} not supported by + this CMakeList.txt, under development") endif() + + add_executable(tf_hello main.cxx) target_link_libraries(tf_hello + ${TENSORFLOW_LIBRARIES}) ``` + +4. Configure the folder with cmake-gui, an error should be prompted out, + requesting you to locate the folder containing `TensorflowConfig.cmake`. + This file can be found at `` or `` (for + those have build install in previous steps). + +5. Configure again, generate the project. + +6. Compile the project with `Release` config (Windows). For Linux users, just + compile the project. + +7. Copy the `tensorflow.dll`(Windows)/`tensorflow.so`(Linux) from build + directory to the build folder containing `tf_hello` binary. + +8. Run `tf_hello` binary -Step-by-step Windows build -========================== +# Step-by-step Windows build (command prompt) 1. Install the prerequisites detailed above, and set up your environment. diff --git a/tensorflow/contrib/cmake/TensorflowConfig.cmake.in b/tensorflow/contrib/cmake/TensorflowConfig.cmake.in new file mode 100644 index 0000000000000000000000000000000000000000..cc04db6e952f53b8bb5416dde60b8173e60bf60e --- /dev/null +++ b/tensorflow/contrib/cmake/TensorflowConfig.cmake.in @@ -0,0 +1,16 @@ +# - Config file for the Tensorflow package +# It defines the following variables +# TENSORFLOW_INCLUDE_DIRS - include directories for FooBar +# TENSORFLOW_LIBRARIES - libraries to link against + +# Compute paths +get_filename_component(TENSORFLOW_CMAKE_DIR "${CMAKE_CURRENT_LIST_FILE}" PATH) +set(TENSORFLOW_INCLUDE_DIRS "@CONF_INCLUDE_DIRS@") + +# Our library dependencies (contains definitions for IMPORTED targets) +if(NOT TENSORFLOW_BINARY_DIR) + include("${TENSORFLOW_CMAKE_DIR}/TensorflowTargets.cmake") +endif() + +# These are IMPORTED targets created by TensorflowTargets.cmake +set(TENSORFLOW_LIBRARIES tensorflow) \ No newline at end of file diff --git a/tensorflow/contrib/cmake/TensorflowConfigVersion.cmake.in b/tensorflow/contrib/cmake/TensorflowConfigVersion.cmake.in new file mode 100644 index 0000000000000000000000000000000000000000..2a9609ddb9c4ca864651818bdfae0f8fe290de31 --- /dev/null +++ b/tensorflow/contrib/cmake/TensorflowConfigVersion.cmake.in @@ -0,0 +1,11 @@ +set(PACKAGE_VERSION "@TENSORFLOW_VERSION@") + +# Check whether the requested PACKAGE_FIND_VERSION is compatible +if("${PACKAGE_VERSION}" VERSION_LESS "${PACKAGE_FIND_VERSION}") + set(PACKAGE_VERSION_COMPATIBLE FALSE) +else() + set(PACKAGE_VERSION_COMPATIBLE TRUE) + if ("${PACKAGE_VERSION}" VERSION_EQUAL "${PACKAGE_FIND_VERSION}") + set(PACKAGE_VERSION_EXACT TRUE) + endif() +endif() \ No newline at end of file diff --git a/tensorflow/contrib/cmake/external/abseil_cpp.cmake b/tensorflow/contrib/cmake/external/abseil_cpp.cmake index c6c5021f60b38ed05a19f3e439c9810251841f76..6c6a5df7f76723800740a81ccdcb137a0ec33846 100644 --- a/tensorflow/contrib/cmake/external/abseil_cpp.cmake +++ b/tensorflow/contrib/cmake/external/abseil_cpp.cmake @@ -20,6 +20,7 @@ if (systemlib_ABSEIL_CPP) absl_dynamic_annotations absl_malloc_internal absl_throw_delegate + absl_int128 absl_strings str_format_internal absl_bad_optional_access) @@ -38,20 +39,21 @@ else (systemlib_ABSEIL_CPP) include (ExternalProject) set(abseil_cpp_INCLUDE_DIR ${CMAKE_BINARY_DIR}/abseil_cpp/src/abseil_cpp_build) - set(abseil_cpp_URL https://github.com/abseil/abseil-cpp/archive/e01d95528ea2137a4a27a88d1f57c6cb260aafed.tar.gz) - set(abseil_cpp_HASH SHA256=84043ed402d2a2a6ba4cdddb7e85118b1158fd81fe4ac3a14adc343d054c1e2e) + 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.*") set(abseil_cpp_STATIC_LIBRARIES ${abseil_cpp_BUILD}/absl/base/Release/absl_base.lib - ${abseil_cpp_BUILD}/absl/base/Release/absl_spinlock_wait.lib ${abseil_cpp_BUILD}/absl/base/Release/absl_dynamic_annotations.lib - ${abseil_cpp_BUILD}/absl/base/Release/absl_malloc_internal.lib - ${abseil_cpp_BUILD}/absl/base/Release/absl_throw_delegate.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 @@ -60,8 +62,10 @@ else (systemlib_ABSEIL_CPP) ${abseil_cpp_BUILD}/absl/base/absl_dynamic_annotations.lib ${abseil_cpp_BUILD}/absl/base/absl_malloc_internal.lib ${abseil_cpp_BUILD}/absl/base/absl_throw_delegate.lib + ${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() @@ -71,15 +75,16 @@ else (systemlib_ABSEIL_CPP) ${abseil_cpp_BUILD}/absl/base/libabsl_dynamic_annotations.a ${abseil_cpp_BUILD}/absl/base/libabsl_malloc_internal.a ${abseil_cpp_BUILD}/absl/base/libabsl_throw_delegate.a + ${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_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} @@ -93,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_build) -endif (systemlib_ABSEIL_CPP) +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 b1e64aa55c80ad59cfdc0f4767c0282b4f73367f..e570c09ecb5e64130ed6f3375a51d74850cc3989 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 d184fa229d75d336aedea0041bd59cb93e7e267f) +set(GRPC_TAG 69b6c047bc767b4d80e7af4d00ccb7c45b683dae) if(WIN32) # We use unsecure gRPC because boringssl does not build on windows @@ -26,9 +26,9 @@ if(WIN32) set(grpc_SSL_PROVIDER NONE) if(${CMAKE_GENERATOR} MATCHES "Visual Studio.*") set(grpc_STATIC_LIBRARIES - ${CMAKE_CURRENT_BINARY_DIR}/grpc/src/grpc/Release/grpc++_unsecure.lib - ${CMAKE_CURRENT_BINARY_DIR}/grpc/src/grpc/Release/grpc_unsecure.lib - ${CMAKE_CURRENT_BINARY_DIR}/grpc/src/grpc/Release/gpr.lib) + ${CMAKE_CURRENT_BINARY_DIR}/grpc/src/grpc/$(Configuration)/grpc++_unsecure.lib + ${CMAKE_CURRENT_BINARY_DIR}/grpc/src/grpc/$(Configuration)/grpc_unsecure.lib + ${CMAKE_CURRENT_BINARY_DIR}/grpc/src/grpc/$(Configuration)/gpr.lib) else() set(grpc_STATIC_LIBRARIES ${CMAKE_CURRENT_BINARY_DIR}/grpc/src/grpc/grpc++_unsecure.lib @@ -43,8 +43,9 @@ else() ${CMAKE_CURRENT_BINARY_DIR}/grpc/src/grpc/libgrpc++.a ${CMAKE_CURRENT_BINARY_DIR}/grpc/src/grpc/libgrpc.a ${CMAKE_CURRENT_BINARY_DIR}/grpc/src/grpc/libaddress_sorting.a + ${CMAKE_CURRENT_BINARY_DIR}/grpc/src/grpc/libgpr.a ${CMAKE_CURRENT_BINARY_DIR}/grpc/src/grpc/third_party/cares/cares/lib/libcares.a - ${CMAKE_CURRENT_BINARY_DIR}/grpc/src/grpc/libgpr.a) + ${CMAKE_CURRENT_BINARY_DIR}/grpc/src/grpc/third_party/zlib/libz.a) endif() add_definitions(-DGRPC_ARES=0) @@ -66,7 +67,7 @@ ExternalProject_Add(grpc -DPROTOBUF_INCLUDE_DIRS:STRING=${PROTOBUF_INCLUDE_DIRS} -DPROTOBUF_LIBRARIES:STRING=${protobuf_STATIC_LIBRARIES} -DZLIB_ROOT:STRING=${ZLIB_INSTALL} - -DgRPC_SSL_PROVIDER:STRING=${grpc_SSL_PROVIDER} + -DgRPC_SSL_PROVIDER:STRING=${grpc_SSL_PROVIDER} ) # grpc/src/core/ext/census/tracing.c depends on the existence of openssl/rand.h. 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/external/png.cmake b/tensorflow/contrib/cmake/external/png.cmake index 1a147e9c8e5a9fee17a81e37c9babe3c9ec0290b..32e6d78e508e25f76bd263e9d52b6574ca315f6c 100644 --- a/tensorflow/contrib/cmake/external/png.cmake +++ b/tensorflow/contrib/cmake/external/png.cmake @@ -59,6 +59,7 @@ ExternalProject_Add(png -DCMAKE_VERBOSE_MAKEFILE:BOOL=OFF -DCMAKE_INSTALL_PREFIX:STRING=${png_INSTALL} -DZLIB_ROOT:STRING=${ZLIB_INSTALL} + -DPNG_TESTS:BOOL=OFF ) ## put png includes in the directory where they are expected diff --git a/tensorflow/contrib/cmake/external/protobuf.cmake b/tensorflow/contrib/cmake/external/protobuf.cmake index 56a57a2340ddc7f923c611c222a0399e279ad58a..773c37b309b1dff4ed28d24cd7d6140a63ec5bc6 100644 --- a/tensorflow/contrib/cmake/external/protobuf.cmake +++ b/tensorflow/contrib/cmake/external/protobuf.cmake @@ -16,7 +16,18 @@ include (ExternalProject) set(PROTOBUF_INCLUDE_DIRS ${CMAKE_CURRENT_BINARY_DIR}/protobuf/src/protobuf/src) set(PROTOBUF_URL https://github.com/google/protobuf.git) -set(PROTOBUF_TAG v3.6.1) + +# enable choose protobuf versions +SET(PROTOBUF_VERSION "3.6.1" CACHE STRING "Protobuf version") +SET_PROPERTY(CACHE PROTOBUF_VERSION PROPERTY STRINGS "3.4.0" "3.5.0" "3.6.1") + +if(${PROTOBUF_VERSION} STREQUAL "3.5.1") + set(PROTOBUF_TAG v3.6.1) +elseif(${PROTOBUF_VERSION} STREQUAL "3.5.0") + set(PROTOBUF_TAG 2761122b810fe8861004ae785cc3ab39f384d342) +elseif(${PROTOBUF_VERSION} STREQUAL "3.4.0") + set(PROTOBUF_TAG b04e5cba356212e4e8c66c61bbe0c3a20537c5b9) +endif() if(WIN32) if(${CMAKE_GENERATOR} MATCHES "Visual Studio.*") diff --git a/tensorflow/contrib/cmake/modules/FindAbseilCpp.cmake b/tensorflow/contrib/cmake/modules/FindAbseilCpp.cmake index d4f8bb1bec9ae8eff58dfe78168d8e71319c85e1..944ae3997a9489c13f65f93d9a7e61c21dd975c1 100644 --- a/tensorflow/contrib/cmake/modules/FindAbseilCpp.cmake +++ b/tensorflow/contrib/cmake/modules/FindAbseilCpp.cmake @@ -24,10 +24,10 @@ if(EXISTS "${ABSEIL_CPP_INCLUDE_DIR}" AND NOT "${ABSEIL_CPP_INCLUDE_DIR}" STREQU # search all libraries if no COMPONENTS was requested set(AbseilCpp_FIND_COMPONENTS "absl_algorithm;absl_any;absl_bad_any_cast" - "absl_bad_optional_access;absl_base absl_container;absl_debugging" + "absl_bad_optional_access;absl_base;absl_container;absl_debugging" "absl_dynamic_annotations;absl_examine_stack;absl_failure_signal_handler" - "absl_int128;absl_leak_check;absl_malloc_internal;absl_memory;absl_meta" - "absl_numeric;absl_optional;absl_span;absl_spinlock_wait;absl_stack_consumption" + "absl_int128;absl_leak_check;absl_internal_malloc_internal;absl_memory;absl_meta" + "absl_numeric;absl_optional;absl_span;absl_internal_spinlock_wait;absl_stack_consumption" "absl_stacktrace;absl_str_format;absl_strings;absl_symbolize;absl_synchronization" "absl_throw_delegate;absl_time;absl_utility;str_format_extension_internal" "str_format_internal;test_instance_tracker_lib") diff --git a/tensorflow/contrib/cmake/python_modules.txt b/tensorflow/contrib/cmake/python_modules.txt index 6e72670142d560a364350bb4769f1153f884b0f6..21ae9a08a6bb8f71e5935ddde2d7bb3ed0cd8bbc 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 @@ -57,6 +60,7 @@ tensorflow/python/ops tensorflow/python/ops/distributions tensorflow/python/ops/linalg tensorflow/python/ops/losses +tensorflow/python/ops/signal tensorflow/python/platform tensorflow/python/profiler tensorflow/python/profiler/internal @@ -279,10 +283,10 @@ tensorflow/contrib/linear_optimizer/kernels/g3doc tensorflow/contrib/linear_optimizer/python tensorflow/contrib/linear_optimizer/python/ops # TODO(drpngx): Fix failing imports -# tensorflow/contrib/lite -# tensorflow/contrib/lite/python -# tensorflow/contrib/lite/toco -# tensorflow/contrib/lite/toco/python +# tensorflow/lite +# tensorflow/lite/python +# tensorflow/lite/toco +# tensorflow/lite/toco/python tensorflow/contrib/lookup tensorflow/contrib/losses tensorflow/contrib/losses/python @@ -308,11 +312,6 @@ tensorflow/contrib/model_pruning/examples tensorflow/contrib/model_pruning/examples/cifar10 tensorflow/contrib/model_pruning/python tensorflow/contrib/model_pruning/python/layers -tensorflow/contrib/nccl -tensorflow/contrib/nccl/kernels -tensorflow/contrib/nccl/ops -tensorflow/contrib/nccl/python -tensorflow/contrib/nccl/python/ops tensorflow/contrib/nearest_neighbor tensorflow/contrib/nearest_neighbor/kernels tensorflow/contrib/nearest_neighbor/ops @@ -382,8 +381,6 @@ tensorflow/contrib/seq2seq/python/ops tensorflow/contrib/session_bundle tensorflow/contrib/session_bundle/example tensorflow/contrib/signal -tensorflow/contrib/signal/python -tensorflow/contrib/signal/python/ops tensorflow/contrib/slim tensorflow/contrib/slim/python tensorflow/contrib/slim/python/slim diff --git a/tensorflow/contrib/cmake/python_protos.txt b/tensorflow/contrib/cmake/python_protos.txt index 42afbd9105ef3789430606d909979ca308e2eaa8..013180c89083748b240ad061b342300e886d3568 100644 --- a/tensorflow/contrib/cmake/python_protos.txt +++ b/tensorflow/contrib/cmake/python_protos.txt @@ -6,7 +6,7 @@ tensorflow/contrib/boosted_trees/proto tensorflow/contrib/cloud/kernels tensorflow/contrib/decision_trees/proto tensorflow/contrib/gdr -tensorflow/contrib/lite/toco +tensorflow/lite/toco tensorflow/contrib/mpi tensorflow/contrib/mpi_collectives tensorflow/contrib/session_bundle diff --git a/tensorflow/contrib/cmake/tf_c.cmake b/tensorflow/contrib/cmake/tf_c.cmake index 7a30eb94f54b18a2a517615a315e23e09e1170d0..a04142bd249ed5e16beba11057d0efc1e191e31b 100644 --- a/tensorflow/contrib/cmake/tf_c.cmake +++ b/tensorflow/contrib/cmake/tf_c.cmake @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== + ######################################################## # tf_c_framework library ######################################################## diff --git a/tensorflow/contrib/cmake/tf_cc_ops.cmake b/tensorflow/contrib/cmake/tf_cc_ops.cmake index 6c90cf398c69c8c1b22ea75e0c407f258e2535f9..6514ae50a4a35b35ba100af6997079294c22f9b8 100644 --- a/tensorflow/contrib/cmake/tf_cc_ops.cmake +++ b/tensorflow/contrib/cmake/tf_cc_ops.cmake @@ -149,11 +149,7 @@ add_library(tf_cc OBJECT ${tf_cc_srcs}) add_dependencies(tf_cc tf_cc_framework tf_cc_ops) if (WIN32) - if(${CMAKE_GENERATOR} MATCHES "Visual Studio.*") - set (pywrap_tensorflow_lib "${CMAKE_CURRENT_BINARY_DIR}/${CMAKE_BUILD_TYPE}/pywrap_tensorflow_internal.lib") - else() - set (pywrap_tensorflow_lib "${CMAKE_CURRENT_BINARY_DIR}/pywrap_tensorflow_internal.lib") - endif() + set (pywrap_tensorflow_lib "${CMAKE_CURRENT_BINARY_DIR}/$(Configuration)/pywrap_tensorflow_internal.lib") else (WIN32) set (pywrap_tensorflow_lib "${CMAKE_CURRENT_BINARY_DIR}/libpywrap_tensorflow_internal${CMAKE_SHARED_LIBRARY_SUFFIX}") endif (WIN32) diff --git a/tensorflow/contrib/cmake/tf_core_cpu.cmake b/tensorflow/contrib/cmake/tf_core_cpu.cmake index a54cbff33b66d63d7229fa2f50b8a4ca962111ed..d8884d464fb5974d77506561a9ed36110a3804c0 100644 --- a/tensorflow/contrib/cmake/tf_core_cpu.cmake +++ b/tensorflow/contrib/cmake/tf_core_cpu.cmake @@ -39,6 +39,8 @@ file(GLOB_RECURSE tf_core_cpu_exclude_srcs "${tensorflow_source_dir}/tensorflow/core/*test*.h" "${tensorflow_source_dir}/tensorflow/core/*test*.cc" "${tensorflow_source_dir}/tensorflow/core/*main.cc" + "${tensorflow_source_dir}/tensorflow/core/common_runtime/eager/*.cc" + "${tensorflow_source_dir}/tensorflow/core/common_runtime/eager/*.h" "${tensorflow_source_dir}/tensorflow/core/common_runtime/gpu/*.cc" "${tensorflow_source_dir}/tensorflow/core/common_runtime/gpu_device_factory.cc" "${tensorflow_source_dir}/tensorflow/core/common_runtime/direct_session.cc" diff --git a/tensorflow/contrib/cmake/tf_core_eager_runtime.cmake b/tensorflow/contrib/cmake/tf_core_eager_runtime.cmake new file mode 100644 index 0000000000000000000000000000000000000000..78e4c0d3035cdaefa1d0950f4270d60152c805af --- /dev/null +++ b/tensorflow/contrib/cmake/tf_core_eager_runtime.cmake @@ -0,0 +1,57 @@ +# 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. +# ============================================================================== +######################################################## +# tf_core_eager_runtime library +######################################################## +file(GLOB_RECURSE tf_core_eager_runtime_srcs + "${tensorflow_source_dir}/tensorflow/core/common_runtime/eager/*.cc" + "${tensorflow_source_dir}/tensorflow/core/common_runtime/eager/*.h" +) + +file(GLOB_RECURSE tf_core_eager_runtime_exclude_srcs + "${tensorflow_source_dir}/tensorflow/core/common_runtime/eager/*test*.h" + "${tensorflow_source_dir}/tensorflow/core/common_runtime/eager/*test*.cc" +) + +list(REMOVE_ITEM tf_core_eager_runtime_srcs ${tf_core_eager_runtime_exclude_srcs}) + +add_library(tf_core_eager_runtime OBJECT ${tf_core_eager_runtime_srcs}) +add_dependencies( + tf_core_eager_runtime + tf_c + tf_core_lib) + + +file(GLOB_RECURSE tf_c_eager_srcs + "${tensorflow_source_dir}/tensorflow/c/eager/*.cc" + "${tensorflow_source_dir}/tensorflow/c/eager/*.h" +) + +file(GLOB_RECURSE tf_c_eager_exlclude_srcs + "${tensorflow_source_dir}/tensorflow/c/eager/*test*.h" + "${tensorflow_source_dir}/tensorflow/c/eager/*test*.cc" +) + +list(REMOVE_ITEM tf_c_eager_srcs ${tf_c_eager_exlclude_srcs}) + +add_library(tf_c_eager OBJECT ${tf_c_eager_srcs}) +add_dependencies( + tf_c_eager + tf_core_eager_runtime + tf_c + tf_cc_framework + tf_cc_while_loop + tf_core_lib + tf_protos_cc) \ No newline at end of file diff --git a/tensorflow/contrib/cmake/tf_core_framework.cmake b/tensorflow/contrib/cmake/tf_core_framework.cmake index 7e806685b8448cbd629985cdc00ed1193857abe6..d8d1cc3aa2ca4fff3c950654b7cbd7085c76010c 100644 --- a/tensorflow/contrib/cmake/tf_core_framework.cmake +++ b/tensorflow/contrib/cmake/tf_core_framework.cmake @@ -140,6 +140,7 @@ set(tf_proto_text_srcs "tensorflow/core/example/example.proto" "tensorflow/core/example/feature.proto" "tensorflow/core/framework/allocation_description.proto" + "tensorflow/core/framework/api_def.proto" "tensorflow/core/framework/attr_value.proto" "tensorflow/core/framework/cost_graph.proto" "tensorflow/core/framework/device_attributes.proto" @@ -150,6 +151,7 @@ set(tf_proto_text_srcs "tensorflow/core/framework/log_memory.proto" "tensorflow/core/framework/node_def.proto" "tensorflow/core/framework/op_def.proto" + "tensorflow/core/framework/reader_base.proto" "tensorflow/core/framework/remote_fused_graph_execute_info.proto" "tensorflow/core/framework/resource_handle.proto" "tensorflow/core/framework/step_stats.proto" @@ -159,6 +161,7 @@ set(tf_proto_text_srcs "tensorflow/core/framework/tensor_shape.proto" "tensorflow/core/framework/tensor_slice.proto" "tensorflow/core/framework/types.proto" + "tensorflow/core/framework/variable.proto" "tensorflow/core/framework/versions.proto" "tensorflow/core/lib/core/error_codes.proto" "tensorflow/core/protobuf/cluster.proto" @@ -204,10 +207,10 @@ file(GLOB tf_core_platform_srcs "${tensorflow_source_dir}/tensorflow/core/framework/resource_handle.h" "${tensorflow_source_dir}/tensorflow/core/framework/resource_handle.cc") if (NOT tensorflow_ENABLE_GPU) - file(GLOB tf_core_platform_gpu_srcs + file(GLOB tf_core_platform_gpu_srcs_exclude "${tensorflow_source_dir}/tensorflow/core/platform/cuda_libdevice_path.*" "${tensorflow_source_dir}/tensorflow/core/platform/default/cuda_libdevice_path.*") - list(REMOVE_ITEM tf_core_platform_srcs ${tf_core_platform_gpu_srcs}) + list(REMOVE_ITEM tf_core_platform_srcs ${tf_core_platform_gpu_srcs_exclude}) else() file(GLOB tf_core_platform_srcs_exclude "${tensorflow_source_dir}/tensorflow/core/platform/default/device_tracer.cc") @@ -298,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" ) @@ -313,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_core_kernels.cmake b/tensorflow/contrib/cmake/tf_core_kernels.cmake index 7b892ba248bc43cd885f295288c677ac97efaa06..d66e39ac07c7b7c9423fa7e878a9cefd94b867bd 100644 --- a/tensorflow/contrib/cmake/tf_core_kernels.cmake +++ b/tensorflow/contrib/cmake/tf_core_kernels.cmake @@ -68,14 +68,6 @@ if(tensorflow_BUILD_CONTRIB_KERNELS) "${tensorflow_source_dir}/tensorflow/contrib/coder/kernels/range_coder_ops.cc" "${tensorflow_source_dir}/tensorflow/contrib/coder/kernels/range_coder_ops_util.cc" "${tensorflow_source_dir}/tensorflow/contrib/coder/ops/coder_ops.cc" - "${tensorflow_source_dir}/tensorflow/contrib/data/kernels/assert_next_dataset_op.cc" - "${tensorflow_source_dir}/tensorflow/contrib/data/kernels/csv_dataset_op.cc" - "${tensorflow_source_dir}/tensorflow/contrib/data/kernels/directed_interleave_dataset_op.cc" - "${tensorflow_source_dir}/tensorflow/contrib/data/kernels/ignore_errors_dataset_op.cc" - "${tensorflow_source_dir}/tensorflow/contrib/data/kernels/prefetching_kernels.cc" - "${tensorflow_source_dir}/tensorflow/contrib/data/kernels/threadpool_dataset_op.cc" - "${tensorflow_source_dir}/tensorflow/contrib/data/kernels/unique_dataset_op.cc" - "${tensorflow_source_dir}/tensorflow/contrib/data/ops/dataset_ops.cc" "${tensorflow_source_dir}/tensorflow/contrib/factorization/kernels/clustering_ops.cc" "${tensorflow_source_dir}/tensorflow/contrib/factorization/kernels/masked_matmul_ops.cc" "${tensorflow_source_dir}/tensorflow/contrib/factorization/kernels/wals_solver_ops.cc" @@ -97,9 +89,6 @@ if(tensorflow_BUILD_CONTRIB_KERNELS) "${tensorflow_source_dir}/tensorflow/contrib/layers/ops/sparse_feature_cross_op.cc" "${tensorflow_source_dir}/tensorflow/contrib/libsvm/kernels/decode_libsvm_op.cc" "${tensorflow_source_dir}/tensorflow/contrib/libsvm/ops/libsvm_ops.cc" - "${tensorflow_source_dir}/tensorflow/contrib/nccl/kernels/nccl_manager.cc" - "${tensorflow_source_dir}/tensorflow/contrib/nccl/kernels/nccl_ops.cc" - "${tensorflow_source_dir}/tensorflow/contrib/nccl/ops/nccl_ops.cc" "${tensorflow_source_dir}/tensorflow/contrib/nearest_neighbor/kernels/hyperplane_lsh_probes.cc" "${tensorflow_source_dir}/tensorflow/contrib/nearest_neighbor/ops/nearest_neighbor_ops.cc" "${tensorflow_source_dir}/tensorflow/contrib/resampler/kernels/resampler_ops.cc" diff --git a/tensorflow/contrib/cmake/tf_core_ops.cmake b/tensorflow/contrib/cmake/tf_core_ops.cmake index bc753333dba4f67eee0114c4022743dd59a05982..310eed4ecbfdd30a3b3bdd4728c030fe70930797 100644 --- a/tensorflow/contrib/cmake/tf_core_ops.cmake +++ b/tensorflow/contrib/cmake/tf_core_ops.cmake @@ -13,13 +13,14 @@ # limitations under the License. # ============================================================================== set(tf_op_lib_names - "audio_ops" "array_ops" + "audio_ops" "batch_ops" "bitwise_ops" "boosted_trees_ops" "candidate_sampling_ops" "checkpoint_ops" + "collective_ops" "control_flow_ops" "ctc_ops" "cudnn_rnn_ops" @@ -27,13 +28,14 @@ set(tf_op_lib_names "dataset_ops" "decode_proto_ops" "encode_proto_ops" + "function_ops" "functional_ops" "image_ops" "io_ops" "linalg_ops" "list_ops" - "lookup_ops" "logging_ops" + "lookup_ops" "manip_ops" "math_ops" "nn_ops" @@ -43,10 +45,11 @@ set(tf_op_lib_names "remote_fused_graph_ops" "resource_variable_ops" "rpc_ops" + "scoped_allocator_ops" "script_ops" "sdca_ops" - "set_ops" "sendrecv_ops" + "set_ops" "sparse_ops" "spectral_ops" "state_ops" @@ -54,6 +57,7 @@ set(tf_op_lib_names "string_ops" "summary_ops" "training_ops" + "word2vec_ops" ) foreach(tf_op_lib_name ${tf_op_lib_names}) @@ -89,7 +93,6 @@ GENERATE_CONTRIB_OP_LIBRARY(boosted_trees_prediction "${tensorflow_source_dir}/t GENERATE_CONTRIB_OP_LIBRARY(boosted_trees_quantiles "${tensorflow_source_dir}/tensorflow/contrib/boosted_trees/ops/quantile_ops.cc") GENERATE_CONTRIB_OP_LIBRARY(boosted_trees_stats_accumulator "${tensorflow_source_dir}/tensorflow/contrib/boosted_trees/ops/stats_accumulator_ops.cc") GENERATE_CONTRIB_OP_LIBRARY(coder "${tensorflow_source_dir}/tensorflow/contrib/coder/ops/coder_ops.cc") -GENERATE_CONTRIB_OP_LIBRARY(data_dataset "${tensorflow_source_dir}/tensorflow/contrib/data/ops/dataset_ops.cc") GENERATE_CONTRIB_OP_LIBRARY(factorization_clustering "${tensorflow_source_dir}/tensorflow/contrib/factorization/ops/clustering_ops.cc") GENERATE_CONTRIB_OP_LIBRARY(factorization_factorization "${tensorflow_source_dir}/tensorflow/contrib/factorization/ops/factorization_ops.cc") GENERATE_CONTRIB_OP_LIBRARY(framework_variable "${tensorflow_source_dir}/tensorflow/contrib/framework/ops/variable_ops.cc") @@ -99,7 +102,6 @@ GENERATE_CONTRIB_OP_LIBRARY(image_distort_image "${tensorflow_source_dir}/tensor GENERATE_CONTRIB_OP_LIBRARY(image_sirds "${tensorflow_source_dir}/tensorflow/contrib/image/ops/single_image_random_dot_stereograms_ops.cc") GENERATE_CONTRIB_OP_LIBRARY(layers_sparse_feature_cross "${tensorflow_source_dir}/tensorflow/contrib/layers/ops/sparse_feature_cross_op.cc") GENERATE_CONTRIB_OP_LIBRARY(memory_stats "${tensorflow_source_dir}/tensorflow/contrib/memory_stats/ops/memory_stats_ops.cc") -GENERATE_CONTRIB_OP_LIBRARY(nccl "${tensorflow_source_dir}/tensorflow/contrib/nccl/ops/nccl_ops.cc") GENERATE_CONTRIB_OP_LIBRARY(periodic_resample "${tensorflow_source_dir}/tensorflow/contrib/periodic_resample/ops/array_ops.cc") GENERATE_CONTRIB_OP_LIBRARY(nearest_neighbor "${tensorflow_source_dir}/tensorflow/contrib/nearest_neighbor/ops/nearest_neighbor_ops.cc") GENERATE_CONTRIB_OP_LIBRARY(resampler "${tensorflow_source_dir}/tensorflow/contrib/resampler/ops/resampler_ops.cc") diff --git a/tensorflow/contrib/cmake/tf_python.cmake b/tensorflow/contrib/cmake/tf_python.cmake index 6d86daf5f174a3238ab92e5bba6085c904766766..1fe8795ddf00232eba5a60a130e0845a6f6a8e17 100755 --- a/tensorflow/contrib/cmake/tf_python.cmake +++ b/tensorflow/contrib/cmake/tf_python.cmake @@ -222,17 +222,17 @@ endforeach(python_module) add_custom_command(TARGET tf_python_touchup_modules PRE_BUILD COMMAND ${CMAKE_COMMAND} -E make_directory - "${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/contrib/lite") + "${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/lite") add_custom_command(TARGET tf_python_touchup_modules PRE_BUILD COMMAND ${CMAKE_COMMAND} -E make_directory - "${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/contrib/lite/python") + "${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/lite/python") add_custom_command(TARGET tf_python_touchup_modules PRE_BUILD COMMAND ${CMAKE_COMMAND} -E touch - "${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/contrib/lite/python/__init__.py") + "${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/lite/python/__init__.py") add_custom_command( TARGET tf_python_copy_scripts_to_destination PRE_BUILD COMMAND ${CMAKE_COMMAND} -E touch - ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/contrib/lite/python/lite.py) + ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/lite/python/lite.py) # Generate the tensorflow.python.platform.build_info module. set(BUILD_INFO_PY "${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/python/platform/build_info.py") @@ -313,15 +313,14 @@ function(GENERATE_PYTHON_OP_LIB tf_python_op_lib_name) ${GENERATE_PYTHON_OP_LIB_DESTINATION} PARENT_SCOPE) endfunction() -GENERATE_PYTHON_OP_LIB("audio_ops") GENERATE_PYTHON_OP_LIB("array_ops") +GENERATE_PYTHON_OP_LIB("audio_ops") GENERATE_PYTHON_OP_LIB("batch_ops") GENERATE_PYTHON_OP_LIB("bitwise_ops") GENERATE_PYTHON_OP_LIB("boosted_trees_ops") -GENERATE_PYTHON_OP_LIB("math_ops") -GENERATE_PYTHON_OP_LIB("functional_ops") GENERATE_PYTHON_OP_LIB("candidate_sampling_ops") GENERATE_PYTHON_OP_LIB("checkpoint_ops") +GENERATE_PYTHON_OP_LIB("collective_ops") GENERATE_PYTHON_OP_LIB("control_flow_ops" ADDITIONAL_LIBRARIES $) GENERATE_PYTHON_OP_LIB("ctc_ops") @@ -332,14 +331,18 @@ GENERATE_PYTHON_OP_LIB("decode_proto_ops" DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/contrib/proto/python/ops/gen_decode_proto_op.py) GENERATE_PYTHON_OP_LIB("encode_proto_ops" DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/contrib/proto/python/ops/gen_encode_proto_op.py) +GENERATE_PYTHON_OP_LIB("function_ops") +GENERATE_PYTHON_OP_LIB("functional_ops") GENERATE_PYTHON_OP_LIB("image_ops") GENERATE_PYTHON_OP_LIB("io_ops") GENERATE_PYTHON_OP_LIB("linalg_ops") GENERATE_PYTHON_OP_LIB("list_ops") GENERATE_PYTHON_OP_LIB("logging_ops") GENERATE_PYTHON_OP_LIB("lookup_ops") -GENERATE_PYTHON_OP_LIB("nn_ops") GENERATE_PYTHON_OP_LIB("manip_ops") +GENERATE_PYTHON_OP_LIB("math_ops") +GENERATE_PYTHON_OP_LIB("nn_ops") +GENERATE_PYTHON_OP_LIB("no_op") GENERATE_PYTHON_OP_LIB("parsing_ops") GENERATE_PYTHON_OP_LIB("random_ops") GENERATE_PYTHON_OP_LIB("remote_fused_graph_ops" @@ -347,17 +350,21 @@ GENERATE_PYTHON_OP_LIB("remote_fused_graph_ops" GENERATE_PYTHON_OP_LIB("resource_variable_ops") GENERATE_PYTHON_OP_LIB("rpc_ops" DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/contrib/rpc/python/ops/gen_rpc_op.py) +GENERATE_PYTHON_OP_LIB("scoped_allocator_ops") GENERATE_PYTHON_OP_LIB("script_ops") GENERATE_PYTHON_OP_LIB("sdca_ops") +GENERATE_PYTHON_OP_LIB("sendrecv_ops") GENERATE_PYTHON_OP_LIB("set_ops") -GENERATE_PYTHON_OP_LIB("state_ops") GENERATE_PYTHON_OP_LIB("sparse_ops") GENERATE_PYTHON_OP_LIB("spectral_ops") +GENERATE_PYTHON_OP_LIB("state_ops") +GENERATE_PYTHON_OP_LIB("stateless_random_ops") GENERATE_PYTHON_OP_LIB("string_ops") GENERATE_PYTHON_OP_LIB("summary_ops") GENERATE_PYTHON_OP_LIB("user_ops") GENERATE_PYTHON_OP_LIB("training_ops" DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/python/training/gen_training_ops.py) +GENERATE_PYTHON_OP_LIB("word2vec_ops") GENERATE_PYTHON_OP_LIB("contrib_boosted_trees_model_ops" DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/contrib/boosted_trees/python/ops/gen_model_ops.py) @@ -373,8 +380,6 @@ GENERATE_PYTHON_OP_LIB("contrib_boosted_trees_stats_accumulator_ops" DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/contrib/boosted_trees/python/ops/gen_stats_accumulator_ops.py) GENERATE_PYTHON_OP_LIB("contrib_coder_ops" DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/contrib/coder/python/ops/gen_coder_ops.py) -GENERATE_PYTHON_OP_LIB("contrib_data_dataset_ops" - DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/contrib/data/python/ops/gen_dataset_ops.py) GENERATE_PYTHON_OP_LIB("contrib_factorization_clustering_ops" DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/contrib/factorization/python/ops/gen_clustering_ops.py) GENERATE_PYTHON_OP_LIB("contrib_factorization_factorization_ops" @@ -393,11 +398,8 @@ GENERATE_PYTHON_OP_LIB("contrib_layers_sparse_feature_cross_ops" DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/contrib/layers/ops/gen_sparse_feature_cross_op.py) GENERATE_PYTHON_OP_LIB("contrib_memory_stats_ops" DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/contrib/memory_stats/ops/gen_memory_stats_ops.py) -GENERATE_PYTHON_OP_LIB("contrib_nccl_ops" - DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/contrib/nccl/ops/gen_nccl_ops.py) GENERATE_PYTHON_OP_LIB("contrib_periodic_resample_ops" DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/contrib/periodic_resample/python/ops/gen_periodic_resample_op.py) - GENERATE_PYTHON_OP_LIB("contrib_nearest_neighbor_ops" DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/contrib/nearest_neighbor/ops/gen_nearest_neighbor_ops.py) GENERATE_PYTHON_OP_LIB("contrib_resampler_ops" @@ -422,8 +424,6 @@ GENERATE_PYTHON_OP_LIB("contrib_bigquery_reader_ops" DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/contrib/cloud/python/ops/gen_bigquery_reader_ops.py) GENERATE_PYTHON_OP_LIB("contrib_gcs_config_ops" DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/contrib/cloud/python/ops/gen_gcs_config_ops.py) -GENERATE_PYTHON_OP_LIB("stateless_random_ops" - DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/contrib/stateless/gen_stateless_random_ops.py) GENERATE_PYTHON_OP_LIB("debug_ops" DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/tf_python/tensorflow/python/debug/ops/gen_debug_ops.py) @@ -526,11 +526,13 @@ if(WIN32) add_library(pywrap_tensorflow_internal_static STATIC ${pywrap_tensorflow_internal_src} $ + $ $ $ $ $ $ + $ $ $ $ @@ -583,11 +585,13 @@ endif(WIN32) add_library(pywrap_tensorflow_internal SHARED ${pywrap_tensorflow_internal_src} $ + $ $ $ $ $ $ + $ $ $ $ @@ -617,13 +621,28 @@ target_include_directories(pywrap_tensorflow_internal PUBLIC ${NUMPY_INCLUDE_DIR} ) -target_link_libraries(pywrap_tensorflow_internal PRIVATE +if(CMAKE_COMPILER_IS_GNUCC AND CMAKE_CXX_COMPILER_VERSION VERSION_GREATER 5.0) + # There is a bug in GCC 5 resulting in undefined reference to a __cpu_model function when + # linking to the tensorflow library. Adding the following libraries fixes it. + # See issue on github: https://github.com/tensorflow/tensorflow/issues/9593 + target_link_libraries(pywrap_tensorflow_internal PRIVATE ${tf_core_gpu_kernels_lib} ${tensorflow_EXTERNAL_LIBRARIES} tf_protos_cc tf_python_protos_cc ${PYTHON_LIBRARIES} + gcc_s + gcc ) +else() + target_link_libraries(pywrap_tensorflow_internal PRIVATE + ${tf_core_gpu_kernels_lib} + ${tensorflow_EXTERNAL_LIBRARIES} + tf_protos_cc + tf_python_protos_cc + ${PYTHON_LIBRARIES} +) +endif() if(WIN32) @@ -783,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} @@ -808,10 +828,10 @@ add_dependencies(tf_python_api tf_python_ops) ######################################################## # Parse tensorflow/python/tools/api/generator/BUILD to get list of generated files. -FILE(READ ${tensorflow_source_dir}/tensorflow/python/tools/api/generator/api_gen.bzl api_generator_BUILD_text) -STRING(REGEX MATCH "# BEGIN GENERATED ESTIMATOR FILES.*# END GENERATED ESTIMATOR FILES" api_init_files_text ${api_generator_BUILD_text}) -string(REPLACE "# BEGIN GENERATED ESTIMATOR FILES" "" api_init_files_text ${api_init_files_text}) -string(REPLACE "# END GENERATED ESTIMATOR FILES" "" api_init_files_text ${api_init_files_text}) +FILE(READ ${tensorflow_source_dir}/tensorflow/python/tools/api/generator/api_init_files.bzl api_generator_BUILD_text) +STRING(REGEX MATCH "# BEGIN GENERATED FILES.*# END GENERATED FILES" api_init_files_text ${api_generator_BUILD_text}) +string(REPLACE "# BEGIN GENERATED FILES" "" api_init_files_text ${api_init_files_text}) +string(REPLACE "# END GENERATED FILES" "" api_init_files_text ${api_init_files_text}) string(REPLACE "," ";" api_init_files_list ${api_init_files_text}) set(api_init_files "") diff --git a/tensorflow/contrib/cmake/tf_shared_lib.cmake b/tensorflow/contrib/cmake/tf_shared_lib.cmake index fdf522f1fd90ffc64acbe82381ef57a389645d61..62005dd113bfb80fbdf23afb6d4aa5f90a1e32de 100644 --- a/tensorflow/contrib/cmake/tf_shared_lib.cmake +++ b/tensorflow/contrib/cmake/tf_shared_lib.cmake @@ -23,6 +23,8 @@ if(WIN32) # we need. # add_library(tensorflow_static STATIC + $ + $ $ $ $ @@ -65,6 +67,8 @@ endif(WIN32) # tensorflow is a shared library containing all of the # TensorFlow runtime and the standard ops and kernels. add_library(tensorflow SHARED + $ + $ $ $ $ @@ -96,6 +100,27 @@ if(CMAKE_COMPILER_IS_GNUCC AND CMAKE_CXX_COMPILER_VERSION VERSION_GREATER 5.0) target_link_libraries(tensorflow PRIVATE gcc_s gcc) endif() +# Offer the user the choice of overriding the installation directories +set(INSTALL_LIB_DIR lib CACHE PATH "Installation directory for libraries") +set(INSTALL_BIN_DIR bin CACHE PATH "Installation directory for executables") +set(INSTALL_INCLUDE_DIR include CACHE PATH + "Installation directory for header files") +if(WIN32 AND NOT CYGWIN) + set(DEF_INSTALL_CMAKE_DIR cmake) +else() + set(DEF_INSTALL_CMAKE_DIR lib/cmake) +endif() +set(INSTALL_CMAKE_DIR ${DEF_INSTALL_CMAKE_DIR} CACHE PATH + "Installation directory for CMake files") + +# Make relative paths absolute (needed later on) +foreach(p LIB BIN INCLUDE CMAKE) + set(var INSTALL_${p}_DIR) + if(NOT IS_ABSOLUTE "${${var}}") + set(${var} "${CMAKE_INSTALL_PREFIX}/${${var}}") + endif() +endforeach() + if(WIN32) add_dependencies(tensorflow tensorflow_static) endif(WIN32) @@ -103,14 +128,57 @@ endif(WIN32) target_include_directories(tensorflow PUBLIC $) -install(TARGETS tensorflow EXPORT tensorflow_export - RUNTIME DESTINATION bin - LIBRARY DESTINATION lib - ARCHIVE DESTINATION lib) +# Add all targets to build-tree export set +export(TARGETS tensorflow + FILE ${PROJECT_BINARY_DIR}/TensorflowTargets.cmake) + +# Export the package for use from the build-tree +export(PACKAGE Tensorflow) + +# Create the TensorflowConfig.cmake and TensorflowConfigVersion files +file(RELATIVE_PATH REL_INCLUDE_DIR "${INSTALL_CMAKE_DIR}" + "${INSTALL_INCLUDE_DIR}") +# for the build tree +set(CONF_INCLUDE_DIRS "${tensorflow_source_dir}" + "${PROJECT_BINARY_DIR}" + "${CMAKE_CURRENT_BINARY_DIR}/protobuf/src/protobuf/src" + "${CMAKE_CURRENT_BINARY_DIR}/nsync/install/include" # Please if there is a better directory + "${CMAKE_CURRENT_BINARY_DIR}/eigen/src/eigen/Eigen/" + "${CMAKE_CURRENT_BINARY_DIR}/external/eigen_archive/" + "${tensorflow_source_dir}/third_party/eigen3/" + "${CMAKE_CURRENT_BINARY_DIR}/eigen/src/eigen/unsupported/Eigen/") +configure_file(TensorflowConfig.cmake.in + "${PROJECT_BINARY_DIR}/TensorflowConfig.cmake" @ONLY) +# for the install tree, yet to be complete +set(CONF_INCLUDE_DIRS "\${TENSORFLOW_CMAKE_DIR}/${REL_INCLUDE_DIR}") +configure_file(TensorflowConfig.cmake.in + "${PROJECT_BINARY_DIR}/${CMAKE_FILES_DIRECTORY}/TensorflowConfig.cmake" @ONLY) +# for both +configure_file(TensorflowConfigVersion.cmake.in + "${PROJECT_BINARY_DIR}/TensorflowConfigVersion.cmake" @ONLY) + +# install(TARGETS tensorflow EXPORT tensorflow_export +# RUNTIME DESTINATION ${INSTALL_BIN_DIR} +# LIBRARY DESTINATION ${INSTALL_LIB_DIR} +# ARCHIVE DESTINATION ${INSTALL_LIB_DIR}) + +# install(EXPORT tensorflow_export +# FILE TensorflowConfig.cmake +# DESTINATION ${INSTALL_CMAKE_DIR}) -install(EXPORT tensorflow_export - FILE TensorflowConfig.cmake - DESTINATION lib/cmake) +install(FILES + "${PROJECT_BINARY_DIR}/${CMAKE_FILES_DIRECTORY}/TensorflowConfig.cmake" + "${PROJECT_BINARY_DIR}/TensorflowConfigVersion.cmake" + DESTINATION "${INSTALL_CMAKE_DIR}" COMPONENT dev) + +# install the export set for use with the install-tree +install(EXPORT TensorflowTargets + DESTINATION ${INSTALL_CMAKE_DIR}) + +install(TARGETS tensorflow EXPORT TensorflowTargets + RUNTIME DESTINATION ${INSTALL_BIN_DIR} + LIBRARY DESTINATION ${INSTALL_LIB_DIR} + ARCHIVE DESTINATION ${INSTALL_LIB_DIR}) # install necessary headers # tensorflow headers @@ -145,6 +213,10 @@ install(DIRECTORY ${tensorflow_source_dir}/third_party/eigen3/ # unsupported Eigen directory install(DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/eigen/src/eigen/unsupported/Eigen/ DESTINATION include/unsupported/Eigen) +# absl directory +install(DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/abseil_cpp/src/abseil_cpp/absl/ + DESTINATION include/absl + FILES_MATCHING PATTERN "*.h") # mkl if (tensorflow_ENABLE_MKL_SUPPORT) install(DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/mkl/src/mkl/include/ diff --git a/tensorflow/contrib/compiler/BUILD b/tensorflow/contrib/compiler/BUILD index 9987dc9c95efdc7faff606a93c33bad7d0c9bf10..e32097ceddfec95b8677fc762d641d09078e5343 100644 --- a/tensorflow/contrib/compiler/BUILD +++ b/tensorflow/contrib/compiler/BUILD @@ -7,6 +7,7 @@ package_group( includes = ["//tensorflow/compiler/jit:friends"], packages = [ "//tensorflow/...", + "//tensorflow_models/...", "//third_party/py/tensor2tensor/...", ], ) @@ -57,6 +58,7 @@ py_library( srcs_version = "PY2AND3", deps = [ "//tensorflow/compiler/jit:xla_ops_py", + "//tensorflow/compiler/jit/ops:xla_ops_grad", "//tensorflow/python:array_ops", "//tensorflow/python:control_flow_ops", "//tensorflow/python:framework_ops", @@ -68,22 +70,30 @@ py_library( ], ) -tf_py_test( +cuda_py_test( name = "xla_test", srcs = ["xla_test.py"], additional_deps = [ ":xla", - "@six_archive//:six", + "@absl_py//absl/testing:parameterized", + "//tensorflow/compiler/tests:xla_test", + "//tensorflow/contrib/tpu:tpu_estimator", + "//tensorflow/contrib/tpu:tpu_lib", + "//tensorflow/python:client_testlib", "//tensorflow/python:constant_op", "//tensorflow/python:control_flow_ops", "//tensorflow/python:control_flow_util", "//tensorflow/python:math_ops", "//tensorflow/python:platform", - "//tensorflow/contrib/tpu:tpu_lib", "//tensorflow/python:state_ops", "//tensorflow/python:summary", "//tensorflow/python:training", "//tensorflow/python:variable_scope", + "//tensorflow/python/data/ops:dataset_ops", + ], + tags = [ + "no_mac", + "no_windows", ], - tags = ["no_pip"], + xla_enabled = True, ) diff --git a/tensorflow/contrib/compiler/xla.py b/tensorflow/contrib/compiler/xla.py index 146f4b51a5e5b8decd5f50945df072d6b6915d9a..0f1be500f499ebba7e1907de663f8bbfa889bb17 100644 --- a/tensorflow/contrib/compiler/xla.py +++ b/tensorflow/contrib/compiler/xla.py @@ -23,6 +23,7 @@ import contextlib from six.moves import xrange # pylint: disable=redefined-builtin from tensorflow.compiler.jit.ops import xla_ops +from tensorflow.compiler.jit.ops import xla_ops_grad # pylint: disable=unused-import from tensorflow.core.framework import attr_value_pb2 from tensorflow.python.estimator import model_fn as model_fn_lib from tensorflow.python.framework import ops @@ -33,6 +34,7 @@ from tensorflow.python.ops import variable_scope from tensorflow.python.platform import tf_logging as logging from tensorflow.python.util import compat from tensorflow.python.util import function_utils +from tensorflow.python.util import nest from tensorflow.python.util import tf_decorator from tensorflow.python.util import tf_inspect @@ -75,10 +77,22 @@ def compile(computation, inputs=None): # pylint: disable=redefined-builtin All `Operation`s returned from `computation` will be executed when evaluating any of the returned output tensors. - inputs: A list of input tensors or `None` (equivalent to an empty list). + 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-dimention list of scalar tensors rather than a single Rank-N + tensors. If you need different behavior, convert part of inputs to tensors + with `tf.convert_to_tensor`. 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) @@ -179,14 +193,11 @@ 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. - 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() + external_control_inputs = [ + array_ops.identity(x.outputs[0]).op + for x in external_control_inputs + if x.outputs + ] # pylint: disable=protected-access op._add_control_inputs(external_control_inputs) # pylint: enable=protected-access @@ -247,13 +258,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: @@ -262,17 +281,10 @@ def _compile_internal(computation, inputs=None): if not isinstance(inputs, collections.Sequence): raise TypeError('inputs must be a list') + # Flatten inputs. + flat_inputs = nest.flatten(inputs) # Converts inputs to Tensors. - inputs = [ops.convert_to_tensor(x) for x in inputs] - input_arity = len(inputs) - - arg_error = check_function_argument_count( - computation, input_arity, infeed_queue=None) - if arg_error is not None: - raise TypeError( - 'Supplied computation cannot be called with the specified inputs. You ' - 'specified %d inputs: %s, but the computation needs %s' % - (input_arity, str([i.name for i in inputs]), arg_error)) + flat_inputs = [ops.convert_to_tensor(x) for x in flat_inputs] cluster_name = ops.get_default_graph().unique_name('cluster') pivot = control_flow_ops.no_op(name=cluster_name + '/pivot') @@ -282,11 +294,15 @@ def _compile_internal(computation, inputs=None): # Add identity ops so even unused inputs are 'consumed' by the # computation. - computation_inputs = [ + flat_inputs = [ array_ops.identity(x, name='input_{}'.format(i)) - for i, x in enumerate(inputs) + for i, x in enumerate(flat_inputs) ] + # Re-pack flat_inputs in same structure as 'inputs'. + computation_inputs = nest.pack_sequence_as( + structure=inputs, flat_sequence=flat_inputs) + # Only resource variables work inside an XLA computation, so turn on # resource variables for the computation. vscope = variable_scope.get_variable_scope() @@ -299,66 +315,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/compiler/xla_test.py b/tensorflow/contrib/compiler/xla_test.py index 8d13dc7316a693657f1b6e102830808d35372fe9..c4384dcde75035dc55e67bd503e348fe19b97025 100644 --- a/tensorflow/contrib/compiler/xla_test.py +++ b/tensorflow/contrib/compiler/xla_test.py @@ -18,19 +18,34 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import re +from absl.testing import parameterized + from tensorflow.contrib.compiler import xla +from tensorflow.contrib.tpu.python.tpu import tpu_estimator from tensorflow.contrib.tpu.python.tpu import tpu_feed +from tensorflow.contrib.training.python.training import hparam from tensorflow.python import summary +from tensorflow.python.data.ops import dataset_ops +from tensorflow.python.estimator import model_fn as model_fn_lib 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 control_flow_ops from tensorflow.python.ops import control_flow_util 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 summary_ops from tensorflow.python.ops import variable_scope from tensorflow.python.platform import test +from tensorflow.python.training import training + + +_TRAIN = model_fn_lib.ModeKeys.TRAIN +_EVAL = model_fn_lib.ModeKeys.EVAL +_EXPECTED_LOSS = 1 +_EXPECTED_FEATURE = 2 +_EXPECTED_LABEL = 3 class XLACompileContextTest(test.TestCase): @@ -49,7 +64,7 @@ class XLACompileContextTest(test.TestCase): histogram_summary = summary.histogram('histogram_summary', dummy_tensor) image_summary = summary.image('image_summary', dummy_tensor) scalar_summary = summary.scalar('scalar_summary', dummy_tensor) - tensor_summary = summary_ops.tensor_summary('tensor_summary', dummy_tensor) + tensor_summary = summary.tensor_summary('tensor_summary', dummy_tensor) summary.merge( [ audio_summary, histogram_summary, image_summary, scalar_summary, @@ -253,5 +268,329 @@ class CheckFunctionArgumentCountTest(test.TestCase): xla.check_function_argument_count(func, 0, queue)) +def _test_train_model_fn(features, labels, mode, params): + """A dummy model_fn for testing purpose.""" + del features, labels, params + loss = constant_op.constant(_EXPECTED_LOSS) + return model_fn_lib.EstimatorSpec( + mode=mode, loss=loss, train_op=array_ops.identity(loss)) + + +@xla.estimator_model_fn +def decorated_model_fn(features, labels, mode, params): + return _test_train_model_fn(features, labels, mode, params) + + +def make_dummy_features_labels(): + # XLA CPU/GPU backend doesn't support guaranteed constant, thus use dataset + # container to work around. + features_dataset = dataset_ops.Dataset.from_tensors( + constant_op.constant(_EXPECTED_FEATURE)).repeat(10) + features_op = features_dataset.make_one_shot_iterator().get_next() + labels_dataset = dataset_ops.Dataset.from_tensors( + constant_op.constant(_EXPECTED_LABEL)).repeat(10) + labels_op = labels_dataset.make_one_shot_iterator().get_next() + return features_op, labels_op + + +class XlaDecoratorTest(test.TestCase, parameterized.TestCase): + + @parameterized.named_parameters( + ('test_use_as_decorator', decorated_model_fn, None), + ('test_use_as_function', xla.estimator_model_fn(_test_train_model_fn), + None), + ('test_use_tpu_false_hparams', decorated_model_fn, + hparam.HParams(use_tpu=False)), + ('test_use_tpu_false_dict_params', decorated_model_fn, { + 'use_tpu': False + }), + ) + def test_compile(self, model_fn, params): + """Calls model_fn and verifies it is compiled.""" + with test.mock.patch.object(xla, 'compile') as mock_xla_compile: + loss = constant_op.constant(_EXPECTED_LOSS) + mock_xla_compile.return_value = [loss] + + features, labels = make_dummy_features_labels() + estimator_spec = model_fn( + features=features, labels=labels, mode=_TRAIN, params=params or {}) + + self.assertEqual(mock_xla_compile.call_count, 1) + self.assertEqual(estimator_spec.mode, _TRAIN) + + with self.test_session() as sess: + self.assertEqual(sess.run(estimator_spec.loss), sess.run(loss)) + self.assertEqual(sess.run(estimator_spec.train_op), sess.run(loss)) + + @parameterized.named_parameters( + ('test_use_tpu_true_hparams', decorated_model_fn, + hparam.HParams(use_tpu=True)), + ('test_use_tpu_true_dict_params', decorated_model_fn, { + 'use_tpu': True + }), + ) + def test_not_compile(self, model_fn, params): + """Calls model_fn and verifies it is NOT compiled.""" + with test.mock.patch.object(xla, 'compile') as mock_xla_compile: + loss = constant_op.constant(_EXPECTED_LOSS) + mock_xla_compile.return_value = [loss] + + features, labels = make_dummy_features_labels() + estimator_spec = model_fn( + features=features, labels=labels, mode=_TRAIN, params=params or {}) + + mock_xla_compile.assert_not_called() + self.assertEqual(estimator_spec.mode, _TRAIN) + + with self.test_session() as sess: + self.assertEqual(sess.run(estimator_spec.loss), sess.run(loss)) + self.assertEqual(sess.run(estimator_spec.train_op), sess.run(loss)) + + def test_model_with_summary(self): + """Tests that summary ops are disabled.""" + + @xla.estimator_model_fn + def model_fn_with_summary(features, labels, mode, params): + del features, labels, params + loss = constant_op.constant(_EXPECTED_LOSS) + summary.scalar('loss_scalar_summary', loss) + summary.histogram('loss_histogram_summary', loss) + summary.image('loss_image_summary', loss) + return model_fn_lib.EstimatorSpec( + mode=mode, loss=loss, train_op=array_ops.identity(loss)) + + features, labels = make_dummy_features_labels() + estimator_spec = model_fn_with_summary( + features=features, labels=labels, mode=_TRAIN, params={}) + + with self.test_session() as sess: + self.assertEqual(sess.run(estimator_spec.loss), _EXPECTED_LOSS) + + +def _test_eval_metric_fn(eval_tensor_1, eval_tensor_2): + return { + 'metric_1': (eval_tensor_1, eval_tensor_1), + 'metric_2': (eval_tensor_2, eval_tensor_2), + } + + +class XlaDecoratorEvaluationTest(test.TestCase): + + def _verify_evaluation_result(self, eval_model_fn): + features, labels = make_dummy_features_labels() + estimator_spec = eval_model_fn( + features=features, labels=labels, mode=_EVAL, params={}) + + with self.test_session() as sess: + self.assertEqual(sess.run(estimator_spec.loss), _EXPECTED_LOSS) + self.assertEqual( + sess.run(estimator_spec.eval_metric_ops['metric_1'][0]), + _EXPECTED_FEATURE + _EXPECTED_LABEL) + self.assertEqual( + sess.run(estimator_spec.eval_metric_ops['metric_1'][1]), + _EXPECTED_FEATURE + _EXPECTED_LABEL) + self.assertEqual( + sess.run(estimator_spec.eval_metric_ops['metric_2'][0]), + _EXPECTED_FEATURE - _EXPECTED_LABEL) + self.assertEqual( + sess.run(estimator_spec.eval_metric_ops['metric_2'][1]), + _EXPECTED_FEATURE - _EXPECTED_LABEL) + + def test_eval_base_estimator_spec_eval_metric_ops_disallowed(self): + + @xla.estimator_model_fn + def eval_model_fn_return_estimator_spec(features, labels, mode, params): + del features, labels, params + loss = constant_op.constant(_EXPECTED_LOSS) + return model_fn_lib.EstimatorSpec( + mode=mode, + loss=loss, + eval_metric_ops={ + 'metric': (array_ops.identity(loss), control_flow_ops.no_op()) + }) + + with self.assertRaisesRegexp( + ValueError, 'EstimatorSpec.eval_metric_ops is not supported with XLA ' + 'compilation. Please use TPUEstimatorSpec.eval_metrics instead.'): + self._verify_evaluation_result(eval_model_fn_return_estimator_spec) + + def test_eval_base_estimator_spec_no_eval_metric_ops(self): + + @xla.estimator_model_fn + def eval_model_fn_no_eval_metric_ops(features, labels, mode, params): + del features, labels, params + return model_fn_lib.EstimatorSpec( + mode=mode, loss=constant_op.constant(_EXPECTED_LOSS)) + + features, labels = make_dummy_features_labels() + estimator_spec = eval_model_fn_no_eval_metric_ops( + features=features, labels=labels, mode=_EVAL, params={}) + with self.test_session() as sess: + self.assertEqual(sess.run(estimator_spec.loss), _EXPECTED_LOSS) + + def test_eval_no_eval_metrics(self): + + @xla.estimator_model_fn + def eval_model_fn_no_eval_metrics(features, labels, mode, params): + del features, labels, params + return tpu_estimator.TPUEstimatorSpec( + mode=mode, loss=constant_op.constant(_EXPECTED_LOSS)) + + features, labels = make_dummy_features_labels() + estimator_spec = eval_model_fn_no_eval_metrics( + features=features, labels=labels, mode=_EVAL, params={}) + + self.assertEqual(estimator_spec.eval_metric_ops, {}) + with self.test_session() as sess: + self.assertEqual(sess.run(estimator_spec.loss), _EXPECTED_LOSS) + + def test_eval_fn_missing_input_tensor(self): + + @xla.estimator_model_fn + def eval_model_fn(features, labels, mode, params): + del params + dummy_eval_metric_fn_tensors_dict = { + 'eval_tensor_1': features + labels, + } + return tpu_estimator.TPUEstimatorSpec( + mode=mode, + loss=constant_op.constant(_EXPECTED_LOSS), + eval_metrics=(_test_eval_metric_fn, + dummy_eval_metric_fn_tensors_dict)) + + with self.assertRaisesRegexp( + ValueError, + re.escape("Arguments ['eval_tensor_2'] are needed by metric_fn (first " + 'element of TPUEstimatorSpec.eval_metrics) but they are not ' + 'provided by evaluation tensors (second element of ' + 'TPUEstimatorSpec.eval_metrics).')): + self._verify_evaluation_result(eval_model_fn) + + def test_eval_fn_extraneous_input_tensor(self): + + @xla.estimator_model_fn + def eval_model_fn(features, labels, mode, params): + del params + dummy_eval_metric_fn_tensors_dict = { + 'eval_tensor_1': features + labels, + 'eval_tensor_2': features - labels, + 'extra_tensor': features * 2 - labels, + } + return tpu_estimator.TPUEstimatorSpec( + mode=mode, + loss=constant_op.constant(_EXPECTED_LOSS), + eval_metrics=(_test_eval_metric_fn, + dummy_eval_metric_fn_tensors_dict)) + + with self.assertRaisesRegexp( + ValueError, + re.escape("Arguments ['extra_tensor'] are provided by evaluation " + 'tensors (second element of TPUEstimatorSpec.eval_metrics) ' + 'but they are not needed by metric_fn (first element of ' + 'TPUEstimatorSpec.eval_metrics).')): + self._verify_evaluation_result(eval_model_fn) + + def test_eval_tensors_as_list(self): + + @xla.estimator_model_fn + def eval_model_fn(features, labels, mode, params): + del params + dummy_eval_metric_fn_tensors = [features + labels, features - labels] + return tpu_estimator.TPUEstimatorSpec( + mode=mode, + loss=constant_op.constant(_EXPECTED_LOSS), + eval_metrics=(_test_eval_metric_fn, dummy_eval_metric_fn_tensors)) + + self._verify_evaluation_result(eval_model_fn) + + def test_eval_tensors_as_dict(self): + + @xla.estimator_model_fn + def eval_model_fn(features, labels, mode, params): + del params + dummy_eval_metric_fn_tensors_dict = { + 'eval_tensor_1': features + labels, + 'eval_tensor_2': features - labels, + } + return tpu_estimator.TPUEstimatorSpec( + mode=mode, + loss=constant_op.constant(_EXPECTED_LOSS), + eval_metrics=(_test_eval_metric_fn, + dummy_eval_metric_fn_tensors_dict)) + + self._verify_evaluation_result(eval_model_fn) + + def test_model_with_summary(self): + """Tests that summary ops are disabled.""" + + @xla.estimator_model_fn + def model_fn_with_summary(features, labels, mode, params): + del features, labels, params + loss = constant_op.constant(_EXPECTED_LOSS) + summary.scalar('loss_scalar_summary', loss) + summary.histogram('loss_histogram_summary', loss) + summary.image('loss_image_summary', loss) + return tpu_estimator.TPUEstimatorSpec(mode=mode, loss=loss) + + features, labels = make_dummy_features_labels() + estimator_spec = model_fn_with_summary( + features=features, labels=labels, mode=_EVAL, params={}) + + with self.test_session() as sess: + self.assertEqual(sess.run(estimator_spec.loss), _EXPECTED_LOSS) + + +class XlaDecoratorScaffoldTest(test.TestCase, parameterized.TestCase): + + def _make_scaffold_fn(self, mode): + + def _scaffold_fn_on_cpu(): + scaffold = training.Scaffold() + self.assertNotIn(mode, self.is_scaffold_fn_called) + self.is_scaffold_fn_called[mode] = True + return scaffold + + return _scaffold_fn_on_cpu + + def test_scaffold_fn_return_none(self): + + @xla.estimator_model_fn + def model_fn(features, labels, mode, params): + del features, labels, params + return tpu_estimator.TPUEstimatorSpec( + mode=mode, + loss=constant_op.constant(_EXPECTED_LOSS), + train_op=control_flow_ops.no_op(), + scaffold_fn=lambda: None) + + features, labels = make_dummy_features_labels() + with self.assertRaisesRegexp( + ValueError, + 'TPUEstimatorSpec.scaffold_fn returns None, which is not allowed'): + model_fn(features=features, labels=labels, mode=_TRAIN, params={}) + + @parameterized.named_parameters( + ('train_mode', _TRAIN), + ('eval_mode', _EVAL), + # TODO(ycao): Add predict_mode test after PREDICT mode is implemented. + ) + def test_scaffold_fn_in_mode(self, mode): + + @xla.estimator_model_fn + def model_fn(features, labels, mode, params): + del features, labels, params + return tpu_estimator.TPUEstimatorSpec( + mode=mode, + loss=constant_op.constant(_EXPECTED_LOSS), + train_op=control_flow_ops.no_op(), + scaffold_fn=self._make_scaffold_fn(mode)) + + features, labels = make_dummy_features_labels() + + self.is_scaffold_fn_called = {} + model_fn(features=features, labels=labels, mode=mode, params={}) + self.assertTrue(self.is_scaffold_fn_called[mode]) + + if __name__ == '__main__': test.main() diff --git a/tensorflow/contrib/constrained_optimization/BUILD b/tensorflow/contrib/constrained_optimization/BUILD index 619153df67c90cea5a5082a411972948bac5fe90..6ca36f354f1866f0e68c6f459fcf44733fab1cd1 100644 --- a/tensorflow/contrib/constrained_optimization/BUILD +++ b/tensorflow/contrib/constrained_optimization/BUILD @@ -42,6 +42,12 @@ 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", + "no_oss", + ], deps = [ ":constrained_optimization", "//tensorflow/python:client_testlib", diff --git a/tensorflow/contrib/constrained_optimization/python/constrained_minimization_problem.py b/tensorflow/contrib/constrained_optimization/python/constrained_minimization_problem.py index 41258edd90866ae9f644a02c42dfe2dc589da998..6926c0d03fe38ab2d62cc588950c7f5a49b2aba1 100644 --- a/tensorflow/contrib/constrained_optimization/python/constrained_minimization_problem.py +++ b/tensorflow/contrib/constrained_optimization/python/constrained_minimization_problem.py @@ -74,8 +74,8 @@ class ConstrainedMinimizationProblem(object): if (constraints_shape.ndims is None or proxy_constraints_shape.ndims is None or - any([ii is None for ii in constraints_shape.as_list()]) or - any([ii is None for ii in proxy_constraints_shape.as_list()])): + any(ii is None for ii in constraints_shape.as_list()) or + any(ii is None for ii in proxy_constraints_shape.as_list())): raise ValueError( "constraints and proxy_constraints must have fully-known shapes") if constraints_shape != proxy_constraints_shape: diff --git a/tensorflow/contrib/constrained_optimization/python/external_regret_optimizer.py b/tensorflow/contrib/constrained_optimization/python/external_regret_optimizer.py index 67f8ac2b9322f39b02c521f8b9cde3831c7889b8..fb0f849b33b0c5d28fff09eb5aac7f2c0d1adc0b 100644 --- a/tensorflow/contrib/constrained_optimization/python/external_regret_optimizer.py +++ b/tensorflow/contrib/constrained_optimization/python/external_regret_optimizer.py @@ -82,7 +82,7 @@ def _project_multipliers_wrt_euclidean_norm(multipliers, radius): raise ValueError( "multipliers must be one dimensional (instead is %d-dimensional)" % multipliers_shape.ndims) - dimension = multipliers_shape[0].value + dimension = multipliers_shape.dims[0].value if dimension is None: raise ValueError("multipliers must have fully-known shape") diff --git a/tensorflow/contrib/constrained_optimization/python/swap_regret_optimizer.py b/tensorflow/contrib/constrained_optimization/python/swap_regret_optimizer.py index a6cb1f62f059770c90bd1aeea391d841aed9aacf..14e6d8701124ba67cdff8140250b5078f6194693 100644 --- a/tensorflow/contrib/constrained_optimization/python/swap_regret_optimizer.py +++ b/tensorflow/contrib/constrained_optimization/python/swap_regret_optimizer.py @@ -156,7 +156,7 @@ def _project_stochastic_matrix_wrt_euclidean_norm(matrix): if matrix_shape[0] != matrix_shape[1]: raise ValueError("matrix must be square (instead has shape (%d,%d))" % (matrix_shape[0], matrix_shape[1])) - dimension = matrix_shape[0].value + dimension = matrix_shape.dims[0].value if dimension is None: raise ValueError("matrix must have fully-known shape") @@ -601,7 +601,7 @@ class MultiplicativeSwapRegretOptimizer(_SwapRegretOptimizer): assert state_shape is not None assert state_shape.ndims == 2 assert state_shape[0] == state_shape[1] - dimension = state_shape[0].value + dimension = state_shape.dims[0].value assert dimension is not None minimum_log_multiplier = standard_ops.log( diff --git a/tensorflow/contrib/copy_graph/python/__init__.py b/tensorflow/contrib/copy_graph/python/__init__.py index b9ff28eb0d7115ff5919c2f758f70ba388f5d4d2..5c1048e02a3104c958f7710ba97980d3353adbad 100644 --- a/tensorflow/contrib/copy_graph/python/__init__.py +++ b/tensorflow/contrib/copy_graph/python/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2015 The TensorFlow Authors. All Rights Reserved. +# 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. diff --git a/tensorflow/contrib/copy_graph/python/util/__init__.py b/tensorflow/contrib/copy_graph/python/util/__init__.py index b9ff28eb0d7115ff5919c2f758f70ba388f5d4d2..5c1048e02a3104c958f7710ba97980d3353adbad 100644 --- a/tensorflow/contrib/copy_graph/python/util/__init__.py +++ b/tensorflow/contrib/copy_graph/python/util/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2015 The TensorFlow Authors. All Rights Reserved. +# 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. diff --git a/tensorflow/contrib/crf/__init__.py b/tensorflow/contrib/crf/__init__.py index fe5e34d258fbc1508a0a85655f29c2c9bc8fa8b1..d53549048f33162ec89dfe957ca58a4bbb4e95c6 100644 --- a/tensorflow/contrib/crf/__init__.py +++ b/tensorflow/contrib/crf/__init__.py @@ -14,8 +14,6 @@ # ============================================================================== """Linear-chain CRF layer. -See the [CRF](https://tensorflow.org/api_guides/python/contrib.crf) guide. - @@crf_binary_score @@crf_decode @@crf_log_likelihood diff --git a/tensorflow/contrib/crf/python/ops/crf.py b/tensorflow/contrib/crf/python/ops/crf.py index 43bb43129bfe1cb1c66f4965476f9b7f849658ad..40e159b8fcbd1864284e208cb15d9ed96119f840 100644 --- a/tensorflow/contrib/crf/python/ops/crf.py +++ b/tensorflow/contrib/crf/python/ops/crf.py @@ -38,12 +38,12 @@ tf_unary_scores, tf_sequence_lengths, tf_transition_params, _ = session.run( [unary_scores, sequence_lengths, transition_params, train_op]) for tf_unary_scores_, tf_sequence_length_ in zip(tf_unary_scores, tf_sequence_lengths): -# Remove padding. -tf_unary_scores_ = tf_unary_scores_[:tf_sequence_length_] + # Remove padding. + tf_unary_scores_ = tf_unary_scores_[:tf_sequence_length_] -# Compute the highest score and its tag sequence. -tf_viterbi_sequence, tf_viterbi_score = tf.contrib.crf.viterbi_decode( - tf_unary_scores_, tf_transition_params) + # Compute the highest score and its tag sequence. + tf_viterbi_sequence, tf_viterbi_score = tf.contrib.crf.viterbi_decode( + tf_unary_scores_, tf_transition_params) """ from __future__ import absolute_import @@ -54,6 +54,7 @@ import numpy as np from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes +from tensorflow.python.framework import tensor_shape from tensorflow.python.layers import utils from tensorflow.python.ops import array_ops from tensorflow.python.ops import gen_array_ops @@ -107,8 +108,10 @@ def crf_sequence_score(inputs, tag_indices, sequence_lengths, return sequence_scores return utils.smart_cond( - pred=math_ops.equal(inputs.shape[1].value or array_ops.shape(inputs)[1], - 1), + pred=math_ops.equal( + tensor_shape.dimension_value( + inputs.shape[1]) or array_ops.shape(inputs)[1], + 1), true_fn=_single_seq_fn, false_fn=_multi_seq_fn) @@ -157,8 +160,10 @@ def crf_multitag_sequence_score(inputs, tag_bitmap, sequence_lengths, transition_params=transition_params) return utils.smart_cond( - pred=math_ops.equal(inputs.shape[1].value or array_ops.shape(inputs)[1], - 1), + pred=math_ops.equal( + tensor_shape.dimension_value( + inputs.shape[1]) or array_ops.shape(inputs)[1], + 1), true_fn=_single_seq_fn, false_fn=_multi_seq_fn) @@ -214,8 +219,10 @@ def crf_log_norm(inputs, sequence_lengths, transition_params): return log_norm return utils.smart_cond( - pred=math_ops.equal(inputs.shape[1].value or - array_ops.shape(inputs)[1], 1), + pred=math_ops.equal( + tensor_shape.dimension_value( + inputs.shape[1]) or array_ops.shape(inputs)[1], + 1), true_fn=_single_seq_fn, false_fn=_multi_seq_fn) @@ -240,7 +247,7 @@ def crf_log_likelihood(inputs, provided by the caller or created in this function. """ # Get shape information. - num_tags = inputs.get_shape()[2].value + num_tags = tensor_shape.dimension_value(inputs.shape[2]) # Get the transition matrix if not provided. if transition_params is None: @@ -342,7 +349,7 @@ class CrfForwardRnnCell(rnn_cell.RNNCell): for the broadcast summation occurring within the cell. """ self._transition_params = array_ops.expand_dims(transition_params, 0) - self._num_tags = transition_params.get_shape()[0].value + self._num_tags = tensor_shape.dimension_value(transition_params.shape[0]) @property def state_size(self): @@ -428,7 +435,7 @@ class CrfDecodeForwardRnnCell(rnn_cell.RNNCell): summation occurring within the cell. """ self._transition_params = array_ops.expand_dims(transition_params, 0) - self._num_tags = transition_params.get_shape()[0].value + self._num_tags = tensor_shape.dimension_value(transition_params.shape[0]) @property def state_size(self): @@ -540,7 +547,7 @@ def crf_decode(potentials, transition_params, sequence_length): # For simplicity, in shape comments, denote: # 'batch_size' by 'B', 'max_seq_len' by 'T' , 'num_tags' by 'O' (output). - num_tags = potentials.get_shape()[2].value + num_tags = tensor_shape.dimension_value(potentials.shape[2]) # Computes forward decoding. Get last score and backpointers. crf_fwd_cell = CrfDecodeForwardRnnCell(transition_params) @@ -583,7 +590,7 @@ def crf_decode(potentials, transition_params, sequence_length): return decode_tags, best_score return utils.smart_cond( - pred=math_ops.equal(potentials.shape[1].value or + pred=math_ops.equal(tensor_shape.dimension_value(potentials.shape[1]) or array_ops.shape(potentials)[1], 1), true_fn=_single_seq_fn, false_fn=_multi_seq_fn) diff --git a/tensorflow/contrib/cudnn_rnn/BUILD b/tensorflow/contrib/cudnn_rnn/BUILD index 8d9ff15868f8fe93ae0dfeba82d52f813ba84cef..8d35622e393e15a2f2dfea7c75ad2c9f48aa7150 100644 --- a/tensorflow/contrib/cudnn_rnn/BUILD +++ b/tensorflow/contrib/cudnn_rnn/BUILD @@ -42,10 +42,11 @@ tf_custom_op_py_library( cuda_py_test( name = "cudnn_rnn_ops_test", - size = "large", + size = "medium", srcs = ["python/kernel_tests/cudnn_rnn_ops_test.py"], additional_deps = [ ":cudnn_rnn_py", + "@absl_py//absl/testing:parameterized", "//tensorflow/core:protos_all_py", "//tensorflow/contrib/rnn:rnn_py", "//tensorflow/python/ops/losses:losses", @@ -61,9 +62,10 @@ cuda_py_test( "//tensorflow/python:training", "//tensorflow/python:variables", ], - shard_count = 6, + shard_count = 2, tags = [ "noasan", # http://b/62067814 + "requires-gpu-sm35", ], ) diff --git a/tensorflow/contrib/cudnn_rnn/__init__.py b/tensorflow/contrib/cudnn_rnn/__init__.py index 5d8c6191f8db9f96532aa78e4790a4665d3b4877..5320232268657fa73bcd3e86da49d6525e9b8db5 100644 --- a/tensorflow/contrib/cudnn_rnn/__init__.py +++ b/tensorflow/contrib/cudnn_rnn/__init__.py @@ -24,6 +24,10 @@ @@CudnnGRUSaveable @@CudnnRNNReluSaveable @@CudnnRNNTanhSaveable +@@CudnnParamsFormatConverterLSTM +@@CudnnParamsFormatConverterGRU +@@CudnnParamsFormatConverterTanh +@@CudnnParamsFormatConverterRelu """ from __future__ import absolute_import @@ -48,6 +52,10 @@ _allowed_symbols = [ "CudnnGRUSaveable", "CudnnRNNReluSaveable", "CudnnRNNTanhSaveable", + "CudnnParamsFormatConverterLSTM", + "CudnnParamsFormatConverterGRU", + "CudnnParamsFormatConverterTanh", + "CudnnParamsFormatConverterRelu", ] remove_undocumented(__name__, _allowed_symbols) diff --git a/tensorflow/contrib/cudnn_rnn/python/kernel_tests/cudnn_rnn_ops_test.py b/tensorflow/contrib/cudnn_rnn/python/kernel_tests/cudnn_rnn_ops_test.py index c59d3682d404e032d9f4bf81ef54ab456341cefa..f5219eb134d07c09b16a544f71d4c18986c19681 100644 --- a/tensorflow/contrib/cudnn_rnn/python/kernel_tests/cudnn_rnn_ops_test.py +++ b/tensorflow/contrib/cudnn_rnn/python/kernel_tests/cudnn_rnn_ops_test.py @@ -18,24 +18,30 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import collections import itertools import os import unittest +from absl.testing import parameterized import numpy as np from tensorflow.contrib.cudnn_rnn.python.ops import cudnn_rnn_ops from tensorflow.core.protobuf import saver_pb2 +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 from tensorflow.python.framework import random_seed from tensorflow.python.framework.test_util import TensorFlowTestCase from tensorflow.python.ops import array_ops -from tensorflow.python.ops import gradient_checker -from tensorflow.python.ops import math_ops +from tensorflow.python.ops import gradients_impl +from tensorflow.python.ops import init_ops from tensorflow.python.ops import random_ops +from tensorflow.python.ops import rnn +from tensorflow.python.ops import rnn_cell_impl from tensorflow.python.ops import state_ops +from tensorflow.python.ops import variable_scope from tensorflow.python.ops import variables from tensorflow.python.platform import googletest from tensorflow.python.platform import test @@ -56,710 +62,1080 @@ CUDNN_RNN_TANH_PARAMS_PER_LAYER = cudnn_rnn_ops.CUDNN_RNN_TANH_PARAMS_PER_LAYER CUDNN_RNN_RELU_PARAMS_PER_LAYER = cudnn_rnn_ops.CUDNN_RNN_RELU_PARAMS_PER_LAYER -def _CreateModel(rnn_mode, - num_layers, - num_units, - input_size, - input_mode="linear_input", - direction=cudnn_rnn_ops.CUDNN_RNN_UNIDIRECTION, - dtype=dtypes.float32, - dropout=0.): - del input_mode - if rnn_mode == cudnn_rnn_ops.CUDNN_LSTM: - model_fn = cudnn_rnn_ops.CudnnLSTM - elif rnn_mode == cudnn_rnn_ops.CUDNN_GRU: - model_fn = cudnn_rnn_ops.CudnnGRU - elif rnn_mode == cudnn_rnn_ops.CUDNN_RNN_TANH: - model_fn = cudnn_rnn_ops.CudnnRNNTanh - elif rnn_mode == cudnn_rnn_ops.CUDNN_RNN_RELU: - model_fn = cudnn_rnn_ops.CudnnRNNRelu +def RunLSTM(sess, + num_units, + input_size, + batch_size, + time, + num_layers=1, + variable_seq_lengths=False, + is_training=True, + dropout=0., + num_dirs=True, + dtype=dtypes.float32): + # TODO(jamesqin): add multi-layer tests. + # TODO(jamesqin): add multi-dir tests + assert num_layers == 1 + assert num_dirs == 1 + if is_training and not np.isclose(dropout, 0): + raise ValueError("dropout can not be 0. when test training.") + + # set graph level random seed and numpy random seed. + random_seed.set_random_seed(0) + np.random.seed(0) + + inputs = variable_scope.get_variable( + "inputs", + initializer=np.random.rand(time, batch_size, + input_size).astype(dtype.as_numpy_dtype), + dtype=dtype) + initial_h_op = variable_scope.get_variable( + "initial_h_op", + initializer=np.random.rand(batch_size, + num_units).astype(dtype.as_numpy_dtype), + dtype=dtype) + initial_c_op = variable_scope.get_variable( + "initial_c_op", + initializer=np.random.rand(batch_size, + num_units).astype(dtype.as_numpy_dtype), + dtype=dtype) + + if variable_seq_lengths: + lengths_v = np.random.randint(low=1, high=time + 1, size=batch_size) + lengths_v[0] = time # make sure the max sequence has 'time' elems + lengths = ops.convert_to_tensor(lengths_v.astype(np.int32)) + else: + lengths = None + + initializer = init_ops.random_uniform_initializer( + -0.01, 0.01, dtype=dtype, seed=19980904) + + with variable_scope.variable_scope("test", initializer=initializer): + w = variable_scope.get_variable( + "rnn/lstm_cell/kernel", + shape=[input_size + num_units, num_units * 4], + dtype=dtype) + b = variable_scope.get_variable( + "rnn/lstm_cell/bias", shape=[num_units * 4], dtype=dtype) + + # canonical lstm. must set forget_bias to 0. to align with cudnn lstm. + cell = rnn_cell_impl.LSTMCell(num_units, forget_bias=0., reuse=True) + outputs_op, state_tuple_op = rnn.dynamic_rnn( + cell, + inputs, + sequence_length=lengths, + initial_state=rnn_cell_impl.LSTMStateTuple( + h=initial_h_op, c=initial_c_op), + dtype=dtype, + time_major=True, + scope=None) + + # Convert to cudnn opaque param. + format_converter = cudnn_rnn_ops.CudnnParamsFormatConverterLSTM( + num_layers, num_units, input_size) + opaque_params = format_converter.tf_canonical_to_opaque([w, b]) + + cu_initial_h_op = array_ops.expand_dims(initial_h_op, axis=0) + cu_initial_c_op = array_ops.expand_dims(initial_c_op, axis=0) + cu_outputs_op, cu_h_op, cu_c_op = cudnn_rnn_ops._cudnn_rnn( + inputs, + cu_initial_h_op, + cu_initial_c_op, + opaque_params, + sequence_lengths=lengths, + dropout=dropout, + is_training=is_training, + rnn_mode=cudnn_rnn_ops.CUDNN_LSTM) + # Remove the trivial 1st dimension. + cu_state_tuple_op = rnn_cell_impl.LSTMStateTuple( + c=array_ops.squeeze(cu_c_op, axis=0), + h=array_ops.squeeze(cu_h_op, axis=0)) + + if is_training: + (inp_grad_op, hgrad_op, + cgrad_op, wgrad_op, bgrad_op) = gradients_impl.gradients( + outputs_op, [inputs, initial_h_op, initial_c_op, w, b]) + + (cu_inp_grad_op, cu_hgrad_op, + cu_cgrad_op, opaque_grad_op) = gradients_impl.gradients( + cu_outputs_op, + [inputs, cu_initial_h_op, cu_initial_c_op, opaque_params]) + # Remove the trivial 1st dimension + cu_hgrad_op = array_ops.squeeze(cu_hgrad_op, axis=0) + # Remove the trivial 1st dimension + cu_cgrad_op = array_ops.squeeze(cu_cgrad_op, axis=0) + + cu_wgrad_op, cu_bgrad_op = format_converter.opaque_to_tf_canonical( + opaque_grad_op) + cu_wgrad_op = cu_wgrad_op[0] + cu_bgrad_op = cu_bgrad_op[0] + # cudnn lstm has 2 biases each gate. When converting to tf canonical format, + # the two biases are summed into one. Thus here bias gradient should be + # halved when comparing with tf lstm. + cu_bgrad_op *= 0.5 + + init_op = variables.global_variables_initializer() + sess.run(init_op) + + if is_training: + outputs, state_tuple, inp_grad, state_grad, wgrad, bgrad = sess.run([ + outputs_op, state_tuple_op, inp_grad_op, + (hgrad_op, cgrad_op), wgrad_op, bgrad_op + ]) + (cu_outputs, cu_state_tuple, cu_inp_grad, cu_state_grad, cu_wgrad, + cu_bgrad) = sess.run([ + cu_outputs_op, cu_state_tuple_op, cu_inp_grad_op, + (cu_hgrad_op, cu_cgrad_op), cu_wgrad_op, cu_bgrad_op + ]) + + logging.vlog(1, "outputs: %s" % outputs) + logging.vlog(1, "cu_outputs: %s" % cu_outputs) + logging.vlog(1, "state_tuple: %s" % str(state_tuple)) + logging.vlog(1, "cu_state_tuple: %s" % str(cu_state_tuple)) + logging.vlog(1, "inp_grad: %s" % inp_grad) + logging.vlog(1, "cu_inp_grad: %s" % cu_inp_grad) + logging.vlog(1, "state_grad: %s" % str(state_grad)) + logging.vlog(1, "cu_state_grad: %s" % str(cu_state_grad)) + logging.vlog(1, "wgrad: %s" % str(wgrad)) + logging.vlog(1, "bgrad: %s" % str(bgrad)) + logging.vlog(1, "cu_wgrad: %s" % str(cu_wgrad)) + logging.vlog(1, "cu_bgrad: %s" % str(cu_bgrad)) + return (outputs, cu_outputs, state_tuple, cu_state_tuple, inp_grad, + cu_inp_grad, state_grad, cu_state_grad, wgrad, bgrad, cu_wgrad, + cu_bgrad) else: - raise ValueError("Invalid rnn_mode: %s" % rnn_mode) - return model_fn( - num_layers, - num_units, - input_size, - direction=direction, - dtype=dtype, - dropout=dropout) - - -def _CreateParamsSavable(params, - model, - base_variable_scope=None, - name="params_canonical"): - """Create a RNNParamsSaveable for the weight and bias parameters. + outputs, state_tuple = sess.run([outputs_op, state_tuple_op]) + cu_outputs, cu_state_tuple = sess.run([cu_outputs_op, cu_state_tuple_op]) + + logging.vlog(1, "outputs: %s" % outputs) + logging.vlog(1, "cu_outputs: %s" % cu_outputs) + logging.vlog(1, "state_tuple: %s" % str(state_tuple)) + logging.vlog(1, "cu_state_tuple: %s" % str(cu_state_tuple)) + return outputs, cu_outputs, state_tuple, cu_state_tuple + + +# Basic set of RNN configs to test. They can be further extended in relevant +# test (e.g. adding num_dirs). +NAMED_RNN_TESTCASES = ({ + "testcase_name": "xsmall", + "num_units": 1, + "input_size": 1, + "batch_size": 1, + "time": 1, + "num_layers": 1, +}, { + "testcase_name": "small", + "num_units": 4, + "input_size": 4, + "batch_size": 4, + "time": 4, + "num_layers": 1, +}, { + "testcase_name": "medium", + "num_units": 128, + "input_size": 64, + "batch_size": 8, + "time": 16, + "num_layers": 1, +}, { + "testcase_name": "large", + "num_units": 128, + "input_size": 128, + "batch_size": 16, + "time": 32, + "num_layers": 1, +}) + + +def ExpandNamedTestCases(inputs, *remove_keys, **extra_configs): + """Expands testcase with new config dimensions. + + Example: + inputs = ( + {'testcase_name': 'test1', 'gender': 'male'} + {'testcase_name': 'test2', 'gender': 'female'} + ) + remove_keys: empty + extra_configs = { + 'age': [40, 80] + 'height': [5, 6] + } + + Returns: + ( + {'testcase_name': 'test1_age_40_height_5','gender': 'male', 'age': + 40,'height': 5} + {'testcase_name': 'test1_age_40_height_6', 'gender': 'male', 'age': 40, + 'height': 6} + {'testcase_name': 'test1_age_80_height_5', 'gender': 'male', 'age': 80, + 'height': 5} + {'testcase_name': 'test1_age_80_height_6', 'gender': 'male', 'age': 80, + 'height': 6} + + {'testcase_name': 'test2_age_40_height_5', 'gender': 'female', 'age': + 40, + 'height': 5} + {'testcase_name': 'test2_age_40_height_6', 'gender': 'female', 'age': + 40, + 'height': 6} + {'testcase_name': 'test2_age_80_height_5', 'gender': 'female', 'age': + 80, + 'height': 5} + {'testcase_name': 'test2_age_80_height_6', 'gender': 'female', 'age': + 80, + 'height': 6} + ) Args: - params: a Variable for weight and bias parameters. - model: a CudnnRNN model. - base_variable_scope: a string, prefix of names of saved variables. - name: a string, name of the RNNParamsSaveable object. + inputs: A list of dictionary, each being a testcase. + *remove_keys: A list of keys into testcase which are not needed in new + testcases. + **extra_configs: A dict of new test dimension and applicable values in that + dimension. + Returns: - a RNNParamsSaveable object. + A list of dictionary with expanded test cases. """ - if model._rnn_mode == CUDNN_LSTM: - fn = cudnn_rnn_ops.CudnnLSTMSaveable - elif model._rnn_mode == CUDNN_GRU: - fn = cudnn_rnn_ops.CudnnGRUSaveable - elif model._rnn_mode == CUDNN_RNN_TANH: - fn = cudnn_rnn_ops.CudnnRNNTanhSaveable - elif model._rnn_mode == CUDNN_RNN_RELU: - fn = cudnn_rnn_ops.CudnnRNNReluSaveable - params_saveable = fn( - params, - model.num_layers, - model.num_units, - model.input_size, - model.input_mode, - model.direction, - scope=base_variable_scope, - name=name) - ops.add_to_collection(ops.GraphKeys.SAVEABLE_OBJECTS, params_saveable) - return params_saveable - - -def _MinLSTMParamSize(num_layers, - num_units, - input_size, - direction=cudnn_rnn_ops.CUDNN_RNN_UNIDIRECTION): - if direction == cudnn_rnn_ops.CUDNN_RNN_UNIDIRECTION: - first_layer_weights = 4 * num_units * (num_units + input_size) - higher_layer_weights = 8 * (num_layers - 1) * num_units * num_units - all_biases = 8 * num_layers * num_units - return first_layer_weights + higher_layer_weights + all_biases - elif direction == cudnn_rnn_ops.CUDNN_RNN_BIDIRECTION: - first_layer_weights = 4 * num_units * (num_units + input_size) - higher_layer_weights = (num_layers - 1) * ( - 4 * 2 * num_units * num_units + 4 * num_units**2) - all_biases = 8 * num_layers * num_units - return 2 * (first_layer_weights + higher_layer_weights + all_biases) - else: - raise ValueError("%s direction is not supported.") + res = [] + ordered_extra_configs = collections.OrderedDict(extra_configs) + keys = ordered_extra_configs.keys() + # A list of list of configs. + # The outer loop is iterating keys, the innner is values of one key. + combined_kv = [[(k, v) for v in ordered_extra_configs[k]] for k in keys] + logging.info("combined_kv: %s", combined_kv) + for inp in inputs: + # Each inp is a dict + for config in itertools.product(*combined_kv): + new_inp = dict(inp) + # config is a list in the form of [(k_i, v_j), (k_p, v_q), ...] + suffix = ["%s_%s" % (p[0], str(p[1])) for p in config] + suffix = "_".join(suffix) + new_inp["testcase_name"] += "_" + suffix + for k, v in config: + new_inp[k] = v + # Remove not used keys from the new test case. + if remove_keys: + if not isinstance(remove_keys, (list, tuple)): + remove_keys = [remove_keys] + for k in remove_keys: + new_inp.pop(k, None) + logging.info("new_inp: %s", new_inp) + res.append(new_inp) + # Dedup, necessary if `remove_keys` is set. + return [dict(t) for t in {tuple(d.items()) for d in res}] -class CudnnRNNTestSaveRestore(TensorFlowTestCase): - def _CompareWeights(self, lhs, rhs): - self.assertEqual(len(lhs), len(rhs)) - for lw, rw in zip(lhs, rhs): - self.assertAllEqual(lw, rw) +class CudnnLSTMTest(TensorFlowTestCase, parameterized.TestCase): - def _CompareBiases(self, lhs, rhs, rnn_mode, num_layers, direction): - self.assertEqual(len(lhs), len(rhs)) - if rnn_mode == CUDNN_LSTM: - num_params_per_layer = CUDNN_LSTM_PARAMS_PER_LAYER - elif rnn_mode == CUDNN_GRU: - num_params_per_layer = CUDNN_GRU_PARAMS_PER_LAYER - elif rnn_mode == CUDNN_RNN_TANH: - num_params_per_layer = CUDNN_RNN_TANH_PARAMS_PER_LAYER - else: - num_params_per_layer = CUDNN_RNN_RELU_PARAMS_PER_LAYER - num_dirs = 1 if direction == CUDNN_RNN_UNIDIRECTION else 2 - num_params_per_layer *= num_dirs - self.assertEqual(num_params_per_layer * num_layers, len(lhs)) - - for i in range(num_layers): - layer_lhs = lhs[i * num_params_per_layer: (i+1) * num_params_per_layer] - layer_rhs = rhs[i * num_params_per_layer: (i+1) * num_params_per_layer] - if direction == CUDNN_RNN_UNIDIRECTION: - self._CompareSingleLayerBiases(layer_lhs, layer_rhs) - else: - size = len(layer_lhs) - fw_lhs, bw_lhs = layer_lhs[:size//2], layer_lhs[size//2:] - fw_rhs, bw_rhs = layer_rhs[:size//2], layer_rhs[size//2:] - self._CompareSingleLayerBiases(fw_lhs, fw_rhs) - self._CompareSingleLayerBiases(bw_lhs, bw_rhs) - - def _CompareSingleLayerBiases(self, lhs, rhs): - self.assertEqual(len(lhs), len(rhs)) - - lf_lhs, rt_lhs = lhs[:len(lhs)//2], lhs[len(lhs)//2:] - lf_rhs, rt_rhs = rhs[:len(rhs)//2], rhs[len(rhs)//2:] - self.assertEqual(len(lf_lhs), len(rt_lhs)) - self.assertEqual(len(lf_rhs), len(rt_rhs)) - - sum_lhs, sum_rhs = [], [] - for lf, rt in zip(lf_lhs, rt_lhs): - sum_lhs.append(lf + rt) - for lf, rt in zip(lf_rhs, rt_rhs): - sum_rhs.append(lf + rt) - self.assertEqual(len(sum_lhs), len(sum_rhs)) - for lf, rt in zip(sum_lhs, sum_rhs): - self.assertAllEqual(lf, rt) + def _test_training_helper(self, + num_units, + input_size, + batch_size, + time, + num_layers, + dtype, + variable_seq_lengths, + rtol=3e-6, + atol=3e-6): + with self.session(use_gpu=True) as sess: + (outputs, cu_outputs, state_tuple, cu_state_tuple, inp_grad, cu_inp_grad, + state_grad, cu_state_grad, wgrad, bgrad, cu_wgrad, cu_bgrad) = RunLSTM( + sess, + num_units, + input_size, + batch_size, + time, + num_layers, + variable_seq_lengths=variable_seq_lengths) - def _testSaveRestoreVariable(self, rnn_mode, direction, dtype): - num_layers = 2 - num_units = 7 - input_size = 3 - with ops.Graph().as_default(): - model = _CreateModel( - rnn_mode, - num_layers=num_layers, - num_units=num_units, - input_size=input_size, - direction=direction, - dtype=dtype) - random_seed.set_random_seed(1234) - params_size_t = model.params_size() - params = variables.Variable( - random_ops.random_uniform([params_size_t], dtype=dtype), - dtype=dtype, - validate_shape=False) - saveable = _CreateParamsSavable(params, model) - weights, biases = saveable._OpaqueParamsToCanonical() - reset_params = state_ops.assign( - params, - array_ops.zeros([params_size_t], dtype=dtype), - validate_shape=False) - save_path = os.path.join(self.get_temp_dir(), - "save-restore-variable-test") - saver = saver_lib.Saver(write_version=saver_pb2.SaverDef.V2) - # Passing graph explicitly, otherwise an old sess would be reused. - with self.test_session( - use_gpu=True, graph=ops.get_default_graph()) as sess: - sess.run(variables.global_variables_initializer()) - val = saver.save(sess, save_path) - self.assertEqual(save_path, val) + self.assertAllClose(outputs, cu_outputs, rtol=rtol, atol=atol) + for s, cu_s in zip(state_tuple, cu_state_tuple): + self.assertAllClose(s, cu_s, rtol=rtol, atol=atol) + for sg, cu_sg in zip(state_grad, cu_state_grad): + self.assertAllClose(sg, cu_sg, rtol=rtol, atol=atol) + self.assertAllClose(inp_grad, cu_inp_grad, rtol=rtol, atol=atol) + self.assertAllClose(bgrad, cu_bgrad, rtol=rtol, atol=atol) + self.assertAllClose(wgrad, cu_wgrad, rtol=rtol, atol=atol) - weights_v, biases_v = sess.run([weights, biases]) + @parameterized.named_parameters( + ExpandNamedTestCases(NAMED_RNN_TESTCASES, **{ + "variable_seq_lengths": [True, False], + })) + @unittest.skipUnless(test.is_built_with_cuda(), + "Test only applicable when running on GPUs") + def test_training(self, num_units, input_size, batch_size, time, num_layers, + variable_seq_lengths): + if not context.context().num_gpus(): + self.skipTest("No GPUs found") + self._test_training_helper( + num_units, + input_size, + batch_size, + time, + num_layers, + dtypes.float32, + variable_seq_lengths=variable_seq_lengths) - sess.run(reset_params) - saver.restore(sess, save_path) - weights_v_restored, biases_v_restored = sess.run([weights, biases]) - - self._CompareWeights(weights_v, weights_v_restored) - self._CompareBiases(biases_v, biases_v_restored, rnn_mode, num_layers, - direction) - - def _testSaveRestoreTwoVariables(self, rnn_mode, direction, dtype): - num_layers = 2 - num_units = 7 - input_size = 3 - with ops.Graph().as_default(): - model = _CreateModel( - rnn_mode, - num_layers=num_layers, - num_units=num_units, - input_size=input_size, - direction=direction, - dtype=dtype) - random_seed.set_random_seed(1234) - params_size_t = model.params_size() - names = ["rnn_1", "rnn_2"] - param_vars = [ - variables.Variable( - random_ops.random_uniform([params_size_t], dtype=dtype), - dtype=dtype, - validate_shape=False) for name in names - ] - saveables = [] - for name, params in zip(names, param_vars): - saveables.append(_CreateParamsSavable(params, model, name, name)) - weights1, biases1 = saveables[0]._OpaqueParamsToCanonical() - weights2, biases2 = saveables[1]._OpaqueParamsToCanonical() - reset_params = [ - state_ops.assign( - params, - array_ops.zeros([params_size_t], dtype=dtype), - validate_shape=False) for params in param_vars - ] - save_path = os.path.join(self.get_temp_dir(), - "save-restore-variable-test") - saver = saver_lib.Saver(write_version=saver_pb2.SaverDef.V2) - # Passing graph explicitly, otherwise an old sess would be reused. - with self.test_session(use_gpu=True, - graph=ops.get_default_graph()) as sess: - sess.run(variables.global_variables_initializer()) - val = saver.save(sess, save_path) - self.assertEqual(save_path, val) - weights1_v, biases1_v = sess.run([weights1, biases1]) - weights2_v, biases2_v = sess.run([weights2, biases2]) - - sess.run(reset_params) - saver.restore(sess, save_path) - weights1_v_restored, biases1_v_restored = sess.run([weights1, biases1]) - weights2_v_restored, biases2_v_restored = sess.run([weights2, biases2]) - - self._CompareWeights(weights1_v, weights1_v_restored) - self._CompareWeights(weights2_v, weights2_v_restored) - self._CompareBiases(biases1_v, biases1_v_restored, rnn_mode, num_layers, - direction) - self._CompareBiases(biases2_v, biases2_v_restored, rnn_mode, num_layers, - direction) - - def _testSaveRestoreOutput(self, rnn_mode, direction, dtype): - with ops.Graph().as_default(): - num_layers = 2 - num_units = 7 - input_size = 7 - seq_length = 10 - batch_size = 5 - dir_count = 1 if direction == cudnn_rnn_ops.CUDNN_RNN_UNIDIRECTION else 2 - model = _CreateModel( - rnn_mode, + @parameterized.named_parameters( + ExpandNamedTestCases(NAMED_RNN_TESTCASES, **{ + "variable_seq_lengths": [True, False], + })) + @unittest.skipUnless(test.is_built_with_cuda(), + "Test only applicable when running on GPUs") + def test_training_fp16(self, num_units, input_size, batch_size, time, + num_layers, variable_seq_lengths): + if not context.context().num_gpus(): + self.skipTest("No GPUs found") + self._test_training_helper( + num_units, + input_size, + batch_size, + time, + num_layers, + dtypes.float16, + rtol=5e-3, + atol=5e-4, + variable_seq_lengths=variable_seq_lengths) + + @parameterized.named_parameters( + ExpandNamedTestCases(NAMED_RNN_TESTCASES, **{ + "variable_seq_lengths": [True, False], + })) + @unittest.skipUnless(test.is_built_with_cuda(), + "Test only applicable when running on GPUs") + def test_inference(self, num_units, input_size, batch_size, time, num_layers, + variable_seq_lengths): + if not context.context().num_gpus(): + self.skipTest("No GPUs found") + with self.session(use_gpu=True) as sess: + (outputs, cu_outputs, state_tuple, cu_state_tuple) = RunLSTM( + sess, + num_units, + input_size, + batch_size, + time, num_layers, + is_training=False, + variable_seq_lengths=variable_seq_lengths) + + self.assertAllClose(outputs, cu_outputs) + # h + self.assertAllClose(state_tuple.h, cu_state_tuple.h) + # c + self.assertAllClose(state_tuple.c, cu_state_tuple.c) + + @parameterized.named_parameters( + ExpandNamedTestCases(NAMED_RNN_TESTCASES, **{ + "variable_seq_lengths": [True, False], + })) + @unittest.skipUnless(test.is_built_with_cuda(), + "Test only applicable when running on GPUs") + def test_inference_fp16(self, num_units, input_size, batch_size, time, + num_layers, variable_seq_lengths): + if not context.context().num_gpus(): + self.skipTest("No GPUs found") + with self.session(use_gpu=True) as sess: + (outputs, cu_outputs, state_tuple, cu_state_tuple) = RunLSTM( + sess, num_units, input_size, - direction=direction, - dtype=dtype) - params_size_t = model.params_size() - params = variables.Variable( - array_ops.ones([params_size_t], dtype=dtype), - validate_shape=False, - dtype=dtype) - _CreateParamsSavable(params, model) - save_path = os.path.join(self.get_temp_dir(), "save-restore-output-test") - saver = saver_lib.Saver(write_version=saver_pb2.SaverDef.V2) + batch_size, + time, + num_layers, + is_training=False, + dtype=dtypes.float16, + variable_seq_lengths=variable_seq_lengths) - np.random.seed(1234) - has_input_c = (rnn_mode == cudnn_rnn_ops.CUDNN_LSTM) - input_data = constant_op.constant( - np.random.randn(seq_length, batch_size, input_size), dtype=dtype) - input_h = constant_op.constant( - np.random.randn(num_layers * dir_count, batch_size, num_units), - dtype=dtype) - if has_input_c: - input_c = constant_op.constant( - np.random.randn(num_layers * dir_count, batch_size, num_units), - dtype=dtype) - outputs = model( - input_data=input_data, - input_h=input_h, - input_c=input_c, - params=params, - is_training=False) - else: - outputs = model( - input_data=input_data, - input_h=input_h, - params=params, - is_training=False) - total_sum = sum(map(math_ops.reduce_sum, outputs)) - # Passing graph explicitly, otherwise an old sess would be reused. - with self.test_session( - use_gpu=True, graph=ops.get_default_graph()) as sess: - sess.run(variables.global_variables_initializer()) - total_sum_v = sess.run(total_sum) - val = saver.save(sess, save_path) - self.assertEqual(save_path, val) - # Passing graph explicitly, otherwise an old sess would be reused. - with self.test_session( - use_gpu=True, graph=ops.get_default_graph()) as sess: - reset_params = state_ops.assign( - params, - array_ops.zeros([params_size_t], dtype=dtype), - validate_shape=False) - sess.run(reset_params) - saver.restore(sess, save_path) - total_sum_v_restored = sess.run(total_sum) - self.assertAllClose(total_sum_v, total_sum_v_restored, atol=1e-5) + rtol, atol = 5e-3, 5e-4 + self.assertAllClose(outputs, cu_outputs, rtol=rtol, atol=atol) + # h + self.assertAllClose( + state_tuple.h, cu_state_tuple.h, rtol=rtol, atol=atol) + # c + self.assertAllClose( + state_tuple.c, cu_state_tuple.c, rtol=rtol, atol=atol) + + @parameterized.named_parameters( + ExpandNamedTestCases(NAMED_RNN_TESTCASES, **{ + "variable_seq_lengths": [True, False], + })) + @unittest.skipUnless(test.is_built_with_cuda(), + "Test only applicable when running on GPUs") + def test_inference_with_dropout(self, num_units, input_size, batch_size, time, + num_layers, variable_seq_lengths): + """Validates that dropout does not affect Cudnn Rnn inference.""" + if not context.context().num_gpus(): + self.skipTest("No GPUs found") + # Hand-picked dropouts are used below (0. and 1.) + with ops.Graph().as_default() as g: + with self.session(use_gpu=True, graph=g) as sess: + # 1st time w/o dropout. + (_, cu_outputs, _, cu_state_tuple) = RunLSTM( + sess, + num_units, + input_size, + batch_size, + time, + num_layers, + is_training=False, + dropout=0., + variable_seq_lengths=variable_seq_lengths) + + with ops.Graph().as_default() as g: + with self.session(use_gpu=True, graph=g) as sess: + (_, cu_outputs2, _, cu_state_tuple2) = RunLSTM( + sess, + num_units, + input_size, + batch_size, + time, + num_layers, + is_training=False, + dropout=1., + variable_seq_lengths=variable_seq_lengths) + + self.assertAllClose(cu_outputs, cu_outputs2) + # h + self.assertAllClose(cu_state_tuple.h, cu_state_tuple2.h) + # c + self.assertAllClose(cu_state_tuple.c, cu_state_tuple2.c) + + +def RunGRU(sess, + num_units, + input_size, + batch_size, + time, + num_layers=1, + is_training=True, + variable_seq_lengths=False, + dropout=0., + num_dirs=True, + dtype=dtypes.float32): + # TODO(jamesqin): add multi-layer tests. + # TODO(jamesqin): add multi-dir tests + assert num_layers == 1 + assert num_dirs == 1 + if is_training and not np.isclose(dropout, 0): + raise ValueError("dropout can not be 0. when test training.") + + # set graph level random seed and numpy random seed. + random_seed.set_random_seed(0) + np.random.seed(0) + + inputs = variable_scope.get_variable( + "inputs", + initializer=np.random.rand(time, batch_size, + input_size).astype(dtype.as_numpy_dtype), + dtype=dtype) + initial_h_op = variable_scope.get_variable( + "initial_h_op", + initializer=np.random.rand(batch_size, + num_units).astype(dtype.as_numpy_dtype), + dtype=dtype) + + if variable_seq_lengths: + lengths_v = np.random.randint(low=1, high=time + 1, size=batch_size) + lengths_v[0] = time # make sure the max sequence has 'time' elems + lengths = ops.convert_to_tensor(lengths_v.astype(np.int32)) + else: + lengths = None + + initializer = init_ops.random_uniform_initializer( + -0.01, 0.01, dtype=dtype, seed=19980904) + with variable_scope.variable_scope("test", initializer=initializer): + gate_kernel = variable_scope.get_variable( + "rnn/cudnn_compatible_gru_cell/gates/kernel", + shape=[input_size + num_units, num_units * 2], + dtype=dtype) + gate_bias = variable_scope.get_variable( + "rnn/cudnn_compatible_gru_cell/gates/bias", + shape=[num_units * 2], + dtype=dtype) + candidate_inp_kernel = variable_scope.get_variable( + "rnn/cudnn_compatible_gru_cell/candidate/input_projection/kernel", + shape=[input_size, num_units], + dtype=dtype) + candidate_inp_bias = variable_scope.get_variable( + "rnn/cudnn_compatible_gru_cell/candidate/input_projection/bias", + shape=[num_units], + dtype=dtype) + candidate_hid_kernel = variable_scope.get_variable( + "rnn/cudnn_compatible_gru_cell/candidate/hidden_projection/kernel", + shape=[num_units, num_units], + dtype=dtype) + candidate_hid_bias = variable_scope.get_variable( + "rnn/cudnn_compatible_gru_cell/candidate/hidden_projection/bias", + shape=[num_units], + dtype=dtype) + + cell = cudnn_rnn_ops.CudnnCompatibleGRUCell(num_units, reuse=True) + outputs_op, h_op = rnn.dynamic_rnn( + cell, + inputs, + sequence_length=lengths, + initial_state=initial_h_op, + dtype=dtype, + time_major=True, + scope=None) + + ws = [gate_kernel, candidate_inp_kernel, candidate_hid_kernel] + bs = [gate_bias, candidate_inp_bias, candidate_hid_bias] + # Convert to cudnn opaque param. + format_converter = cudnn_rnn_ops.CudnnParamsFormatConverterGRU( + num_layers, num_units, input_size) + opaque_params = format_converter.tf_canonical_to_opaque(ws + bs) + + + cu_initial_h_op = array_ops.expand_dims(initial_h_op, axis=0) + cu_outputs_op, cu_h_op, _ = cudnn_rnn_ops._cudnn_rnn( + inputs, + cu_initial_h_op, + array_ops.zeros_like(cu_initial_h_op), # not used + opaque_params, + sequence_lengths=lengths, + dropout=dropout, + is_training=is_training, + rnn_mode=cudnn_rnn_ops.CUDNN_GRU) + + if is_training: + (inp_grad_op, hgrad_op, gk_grad_op, cik_grad_op, chk_grad_op, gb_grad_op, + cib_grad_op, chb_grad_op) = gradients_impl.gradients( + outputs_op, [inputs, initial_h_op] + ws + bs) + + (cu_inp_grad_op, cu_hgrad_op, opaque_grad_op) = gradients_impl.gradients( + cu_outputs_op, [inputs, cu_initial_h_op, opaque_params]) + # Remove the trivial 1st dimension + cu_hgrad_op = array_ops.squeeze(cu_hgrad_op, axis=0) + + cu_wgrad_op, cu_bgrad_op = format_converter.opaque_to_tf_canonical( + opaque_grad_op) + (cu_gk_grad_op, cu_cik_grad_op, cu_chk_grad_op) = cu_wgrad_op + (cu_gb_grad_op, cu_cib_grad_op, cu_chb_grad_op) = cu_bgrad_op + # cudnn gru has 2 biases for reset and update gates. When converting to tf + # canonical format, the two biases are summed into one. Thus here relevant + # bias gradient should be halved before comparing with tf gru. + cu_gb_grad_op *= 0.5 + + init_op = variables.global_variables_initializer() + sess.run(init_op) + + if is_training: + outputs, h, inp_grad, hgrad, wgrad, bgrad = sess.run([ + outputs_op, h_op, inp_grad_op, hgrad_op, + (gk_grad_op, cik_grad_op, chk_grad_op), + (gb_grad_op, cib_grad_op, chb_grad_op) + ]) + (cu_outputs, cu_h, cu_inp_grad, cu_hgrad, cu_wgrad, cu_bgrad) = sess.run([ + cu_outputs_op, cu_h_op, cu_inp_grad_op, cu_hgrad_op, + (cu_gk_grad_op, cu_cik_grad_op, cu_chk_grad_op), + (cu_gb_grad_op, cu_cib_grad_op, cu_chb_grad_op) + ]) + # Remove the trivial 1st dimension + cu_h = np.squeeze(cu_h, axis=0) + + logging.vlog(1, "outputs: %s" % outputs) + logging.vlog(1, "cu_outputs: %s" % cu_outputs) + logging.vlog(1, "h: %s" % h) + logging.vlog(1, "cu_h: %s" % h) + logging.vlog(1, "inp_grad: %s" % inp_grad) + logging.vlog(1, "cu_inp_grad: %s" % cu_inp_grad) + logging.vlog(1, "hgrad: %s" % hgrad) + logging.vlog(1, "cu_hgrad: %s" % cu_hgrad) + logging.vlog(1, "wgrad: %s" % str(wgrad)) + logging.vlog(1, "bgrad: %s" % str(bgrad)) + logging.vlog(1, "cu_wgrad: %s" % str(cu_wgrad)) + logging.vlog(1, "cu_bgrad: %s" % str(cu_bgrad)) + return (outputs, cu_outputs, h, cu_h, inp_grad, cu_inp_grad, hgrad, + cu_hgrad, wgrad, bgrad, cu_wgrad, cu_bgrad) + else: + outputs, h = sess.run([outputs_op, h_op]) + cu_outputs, cu_h = sess.run([cu_outputs_op, cu_h_op]) + # Remove the trivial 1st dimension. + cu_h = np.squeeze(cu_h, axis=0) + + logging.vlog(1, "outputs: %s" % outputs) + logging.vlog(1, "cu_outputs: %s" % cu_outputs) + logging.vlog(1, "h: %s" % h) + logging.vlog(1, "cu_h: %s" % h) + return outputs, cu_outputs, h, cu_h + +class CudnnGRUTest(TensorFlowTestCase, parameterized.TestCase): + + def _test_training_helper(self, + num_units, + input_size, + batch_size, + time, + num_layers, + dtype, + variable_seq_lengths, + rtol=3e-6, + atol=3e-6): + with self.session(use_gpu=True) as sess: + (outputs, cu_outputs, h, cu_h, inp_grad, cu_inp_grad, hgrad, cu_hgrad, + wgrad, bgrad, cu_wgrad, cu_bgrad) = RunGRU( + sess, + num_units, + input_size, + batch_size, + time, + num_layers, + variable_seq_lengths=variable_seq_lengths) + + self.assertAllClose(outputs, cu_outputs, rtol=rtol, atol=atol) + self.assertAllClose(h, cu_h, rtol=rtol, atol=atol) + self.assertAllClose(hgrad, cu_hgrad, rtol=rtol, atol=atol) + self.assertAllClose(inp_grad, cu_inp_grad, rtol=rtol, atol=atol) + for bg, cu_bg in zip(bgrad, cu_bgrad): + self.assertAllClose(bg, cu_bg, rtol=rtol, atol=atol) + for wg, cu_wg in zip(wgrad, cu_wgrad): + self.assertAllClose(wg, cu_wg, rtol=rtol, atol=atol) + + @parameterized.named_parameters( + ExpandNamedTestCases(NAMED_RNN_TESTCASES, **{ + "variable_seq_lengths": [True, False], + })) @unittest.skipUnless(test.is_built_with_cuda(), "Test only applicable when running on GPUs") - def testSaveRestore(self): - rnn_modes = [ - cudnn_rnn_ops.CUDNN_LSTM, cudnn_rnn_ops.CUDNN_GRU, - cudnn_rnn_ops.CUDNN_RNN_TANH, cudnn_rnn_ops.CUDNN_RNN_RELU - ] - directions = [ - cudnn_rnn_ops.CUDNN_RNN_UNIDIRECTION, - cudnn_rnn_ops.CUDNN_RNN_BIDIRECTION - ] - dtype_list = [dtypes.float32, dtypes.float64] - for rnn_mode, direction, dtype in itertools.product(rnn_modes, directions, - dtype_list): - self._testSaveRestoreVariable(rnn_mode, direction, dtype) - self._testSaveRestoreTwoVariables(rnn_mode, direction, dtype) - self._testSaveRestoreOutput(rnn_mode, direction, dtype) - - -class CudnnRNNTestParamsSize(TensorFlowTestCase): - - def _testOneLSTMParamsSize(self, num_layers, num_units, input_size, - direction): - logging.info("Testing one lstm param size with config: %s", locals()) - min_params_size = _MinLSTMParamSize(num_layers, num_units, input_size, - direction) - model = _CreateModel( - cudnn_rnn_ops.CUDNN_LSTM, + def test_training(self, num_units, input_size, batch_size, time, num_layers, + variable_seq_lengths): + if not context.context().num_gpus(): + self.skipTest("No GPUs found") + self._test_training_helper( + num_units, + input_size, + batch_size, + time, num_layers, + dtypes.float32, + variable_seq_lengths=variable_seq_lengths) + + @parameterized.named_parameters( + ExpandNamedTestCases(NAMED_RNN_TESTCASES, **{ + "variable_seq_lengths": [True, False], + })) + @unittest.skipUnless(test.is_built_with_cuda(), + "Test only applicable when running on GPUs") + def test_training_fp16(self, num_units, input_size, batch_size, time, + num_layers, variable_seq_lengths): + if not context.context().num_gpus(): + self.skipTest("No GPUs found") + self._test_training_helper( num_units, input_size, - direction=direction) - params_size = model.params_size() - with self.test_session(use_gpu=True, graph=ops.get_default_graph()) as sess: - params_size_v = sess.run(params_size) - self.assertLessEqual(min_params_size, params_size_v) + batch_size, + time, + num_layers, + dtypes.float16, + rtol=5e-3, + atol=5e-4, + variable_seq_lengths=variable_seq_lengths) + + @parameterized.named_parameters( + ExpandNamedTestCases(NAMED_RNN_TESTCASES, **{ + "variable_seq_lengths": [True, False], + })) + @unittest.skipUnless(test.is_built_with_cuda(), + "Test only applicable when running on GPUs") + def test_inference(self, num_units, input_size, batch_size, time, num_layers, + variable_seq_lengths): + if not context.context().num_gpus(): + self.skipTest("No GPUs found") + with self.session(use_gpu=True) as sess: + (outputs, cu_outputs, h, cu_h) = RunGRU( + sess, + num_units, + input_size, + batch_size, + time, + num_layers, + is_training=False, + variable_seq_lengths=variable_seq_lengths) + self.assertAllClose(outputs, cu_outputs) + self.assertAllClose(h, cu_h) + @parameterized.named_parameters( + ExpandNamedTestCases(NAMED_RNN_TESTCASES, **{ + "variable_seq_lengths": [True, False], + })) @unittest.skipUnless(test.is_built_with_cuda(), "Test only applicable when running on GPUs") - def testLSTMParamsSize(self): - test_configs = [ - [4, 200, 200], - [4, 200, 300], - [4, 200, 100], - [1, 100, 200], - [2, 200, 100], - [3, 200, 400], - ] - directions = [ - cudnn_rnn_ops.CUDNN_RNN_UNIDIRECTION, - cudnn_rnn_ops.CUDNN_RNN_BIDIRECTION - ] - for (config, direction) in itertools.product(test_configs, directions): - num_layers, num_units, input_size = config - with ops.Graph().as_default(): - self._testOneLSTMParamsSize(num_layers, num_units, input_size, - direction) + def test_inference_fp16(self, num_units, input_size, batch_size, time, + num_layers, variable_seq_lengths): + if not context.context().num_gpus(): + self.skipTest("No GPUs found") + with self.session(use_gpu=True) as sess: + (outputs, cu_outputs, h, cu_h) = RunGRU( + sess, + num_units, + input_size, + batch_size, + time, + num_layers, + is_training=False, + dtype=dtypes.float16, + variable_seq_lengths=variable_seq_lengths) + + rtol, atol = 5e-3, 5e-4 + self.assertAllClose(outputs, cu_outputs, rtol=rtol, atol=atol) + self.assertAllClose(h, cu_h, rtol=rtol, atol=atol) + @parameterized.named_parameters( + ExpandNamedTestCases(NAMED_RNN_TESTCASES, **{ + "variable_seq_lengths": [True, False], + })) @unittest.skipUnless(test.is_built_with_cuda(), "Test only applicable when running on GPUs") - def testLSTMParamsSizeShape(self): - with self.assertRaisesRegexp( - ValueError, "Shape must be rank 0 but is rank 1"): - model = _CreateModel( - cudnn_rnn_ops.CUDNN_LSTM, - constant_op.constant([4]), 200, 200, - direction=cudnn_rnn_ops.CUDNN_RNN_UNIDIRECTION) - params_size = model.params_size() - with self.assertRaisesRegexp( - ValueError, "Shape must be rank 0 but is rank 1"): - model = _CreateModel( - cudnn_rnn_ops.CUDNN_LSTM, - 4, constant_op.constant([200]), 200, - direction=cudnn_rnn_ops.CUDNN_RNN_UNIDIRECTION) - params_size = model.params_size() - with self.assertRaisesRegexp( - ValueError, "Shape must be rank 0 but is rank 1"): - model = _CreateModel( + def test_inference_with_dropout(self, num_units, input_size, batch_size, time, + num_layers, variable_seq_lengths): + """Validates that dropout does not affect Cudnn Rnn inference.""" + # Hand-picked dropouts are used below (0. and 1.) + if not context.context().num_gpus(): + self.skipTest("No GPUs found") + with ops.Graph().as_default() as g: + with self.session(use_gpu=True, graph=g) as sess: + # 1st time w/o dropout. + (_, cu_outputs, _, cu_h) = RunGRU( + sess, + num_units, + input_size, + batch_size, + time, + num_layers, + is_training=False, + dropout=0., + variable_seq_lengths=variable_seq_lengths) + + with ops.Graph().as_default() as g: + with self.session(use_gpu=True, graph=g) as sess: + (_, cu_outputs2, _, cu_h2) = RunGRU( + sess, + num_units, + input_size, + batch_size, + time, + num_layers, + is_training=False, + dropout=1., + variable_seq_lengths=variable_seq_lengths) + + self.assertAllClose(cu_outputs, cu_outputs2) + self.assertAllClose(cu_h[0], cu_h2[0]) + + +class CudnnParamsFormatConverterTest(TensorFlowTestCase, + parameterized.TestCase): + """Class for testing various format converters.""" + + def _test_lstm_helper(self, num_units, input_size, num_layers, direction): + with self.session(use_gpu=True) as sess: + random_seed.set_random_seed(0) + np.random.seed(0) + + num_dirs = 1 if direction == cudnn_rnn_ops.CUDNN_RNN_UNIDIRECTION else 2 + format_converter = cudnn_rnn_ops.CudnnParamsFormatConverterLSTM( + num_layers, num_units, input_size, direction=direction) + + ws, bs = [], [] + for _ in range(num_layers * num_dirs): + w = constant_op.constant( + np.random.rand(input_size + num_units, 4 * num_units), + dtype=dtypes.float32) + b = constant_op.constant( + np.random.rand(4 * num_units), dtype=dtypes.float32) + ws.append(w) + bs.append(b) + + opaque_params = format_converter.tf_canonical_to_opaque(ws + bs) + opaque_params_size = cudnn_rnn_ops.cudnn_rnn_opaque_params_size( cudnn_rnn_ops.CUDNN_LSTM, - 4, 200, constant_op.constant([200]), - direction=cudnn_rnn_ops.CUDNN_RNN_UNIDIRECTION) - params_size = model.params_size() + num_layers, + num_units, + input_size, + direction=direction) + ws_r, bs_r = format_converter.opaque_to_tf_canonical(opaque_params) -class CudnnRNNTestInference(TensorFlowTestCase): + # Test tf_canonical_to_opaque() followed by opaque_to_tf_canonical() + # returns the original input. + ws, ws_r, bs, bs_r = sess.run([ws, ws_r, bs, bs_r]) + for w, w_r in zip(ws, ws_r): + self.assertAllClose(w, w_r) + for b, b_r in zip(bs, bs_r): + self.assertAllClose(b, b_r) - def _testOneSimpleInference(self, rnn_mode, num_layers, num_units, input_size, - batch_size, seq_length, dir_count, dropout, - expected, tolerance): - random_seed.set_random_seed(5678) - model = _CreateModel( - rnn_mode, - num_layers, - num_units, - input_size, - input_mode="auto_select", - direction=(cudnn_rnn_ops.CUDNN_RNN_UNIDIRECTION if dir_count == 1 - else cudnn_rnn_ops.CUDNN_RNN_BIDIRECTION), - dropout=dropout) - has_input_c = (rnn_mode == cudnn_rnn_ops.CUDNN_LSTM) - params_size_t = model.params_size() - input_data = array_ops.ones([seq_length, batch_size, input_size]) - input_h = array_ops.ones([num_layers * dir_count, batch_size, num_units]) - params = variables.Variable( - array_ops.ones([params_size_t]), validate_shape=False) - if has_input_c: - input_c = array_ops.ones([num_layers * dir_count, batch_size, num_units]) - output, output_h, output_c = model( - input_data=input_data, - input_h=input_h, - input_c=input_c, - params=params, - is_training=False) - else: - output, output_h = model( - input_data=input_data, - input_h=input_h, - params=params, - is_training=False) - output_sum = math_ops.reduce_sum(output) - output_h_sum = math_ops.reduce_sum(output_h) - total_sum = output_sum + output_h_sum - if has_input_c: - output_c_sum = math_ops.reduce_sum(output_c) - total_sum += output_c_sum - with self.test_session(use_gpu=True, graph=ops.get_default_graph()) as sess: - sess.run(variables.global_variables_initializer()) - total_sum_v = sess.run([total_sum]) + # Test opaque_params size lower bound + opaque_params_size_v = sess.run(opaque_params_size) + min_params_size = sum(x.size for x in ws) + np.sum(x.size for x in bs) + logging.info("min_parm_size: %d vs actual_opaque_param_size: %d", + min_params_size, opaque_params_size_v) + self.assertLessEqual(min_params_size, opaque_params_size_v) - self.assertAllClose( - total_sum_v[0], expected, atol=tolerance, rtol=tolerance) + @parameterized.named_parameters((c["testcase_name"], c["num_units"], + c["input_size"], c["num_layers"]) + for c in NAMED_RNN_TESTCASES) + @unittest.skipUnless(test.is_built_with_cuda(), + "Test only applicable when running on GPUs") + def test_lstm(self, num_units, input_size, num_layers): + if not context.context().num_gpus(): + self.skipTest("No GPUs found") + self._test_lstm_helper(num_units, input_size, num_layers, + cudnn_rnn_ops.CUDNN_RNN_UNIDIRECTION) + @parameterized.named_parameters((c["testcase_name"], c["num_units"], + c["input_size"], c["num_layers"]) + for c in NAMED_RNN_TESTCASES) @unittest.skipUnless(test.is_built_with_cuda(), "Test only applicable when running on GPUs") - def testSimpleInference(self): - test_configs = [ - { - "rnn_mode": cudnn_rnn_ops.CUDNN_LSTM, - "expected": 231833.22, - "tolerance": 1e-2, - "shape": { - "num_layers": 4, - "num_units": 200, - "input_size": 200, - "batch_size": 20, - "seq_length": 10, - "dir_count": 1, - }, - }, - { - "rnn_mode": cudnn_rnn_ops.CUDNN_GRU, - "expected": 56000, - "tolerance": 1e-2, - "shape": { - "num_layers": 4, - "num_units": 200, - "input_size": 200, - "batch_size": 20, - "seq_length": 10, - "dir_count": 1, - }, - }, - { - "rnn_mode": cudnn_rnn_ops.CUDNN_RNN_TANH, - "expected": 56000, - "tolerance": 1e-2, - "shape": { - "num_layers": 4, - "num_units": 200, - "input_size": 200, - "batch_size": 20, - "seq_length": 10, - "dir_count": 1, - }, - }, - { - "rnn_mode": cudnn_rnn_ops.CUDNN_RNN_RELU, - "expected": 130688, - "tolerance": 1e-2, - "shape": { - "num_layers": 2, - "num_units": 8, - "input_size": 4, - "batch_size": 4, - "seq_length": 2, - "dir_count": 1, - }, - }, - ] - # Cudnn scales result for dropout during training, therefore dropout has no - # impact for inference results. - # (lstm, gru, rnn_tanh are saturated in the test. rnn_relu case is most - # demonstrative of the dropout-invariant nature of CudnnRnn.) - dropouts = [0., 0.5, 1.] - for (config, dropout) in itertools.product(test_configs, dropouts): - rnn_mode = config["rnn_mode"] - expected = config["expected"] - tolerance = config["tolerance"] - shape = config["shape"] - with ops.Graph().as_default(): - self._testOneSimpleInference( - rnn_mode, shape["num_layers"], shape["num_units"], - shape["input_size"], shape["batch_size"], shape["seq_length"], - shape["dir_count"], dropout, expected, tolerance) - - -class CudnnRNNTestTraining(TensorFlowTestCase): - - def _testOneSimpleTraining(self, rnn_mode, num_layers, num_units, input_size, - batch_size, seq_length, dir_count, dropout, dtype, - delta, tolerance): - # Gradient checking runs two forward ops with almost the same input. Need to - # make sure the drop patterns across the two runs are the same. - logging.info("Training test with config: %s", locals()) - old_env_state = os.environ.get("TF_CUDNN_RESET_RND_GEN_STATE", str(False)) - os.environ["TF_CUDNN_RESET_RND_GEN_STATE"] = str(True) - has_input_c = (rnn_mode == cudnn_rnn_ops.CUDNN_LSTM) - random_seed.set_random_seed(5678) - direction = (cudnn_rnn_ops.CUDNN_RNN_UNIDIRECTION if dir_count == 1 - else cudnn_rnn_ops.CUDNN_RNN_BIDIRECTION) - model = _CreateModel( - rnn_mode, - num_layers, - num_units, - input_size, - direction=direction, - dtype=dtype, - dropout=dropout) - params_size_t = model.params_size() - input_data = variables.Variable( - random_ops.random_uniform( - [seq_length, batch_size, input_size], dtype=dtype), - dtype=dtype) - input_h = variables.Variable( - random_ops.random_uniform( - [num_layers * dir_count, batch_size, num_units], dtype=dtype), - dtype=dtype) - params = variables.Variable( - random_ops.random_uniform([params_size_t], dtype=dtype), - validate_shape=False, - dtype=dtype) - if has_input_c: - input_c = variables.Variable( - random_ops.random_uniform( - [num_layers * dir_count, batch_size, num_units], dtype=dtype), - dtype=dtype) - - output, output_h, output_c = model( - input_data=input_data, - input_h=input_h, - input_c=input_c, - params=params) - else: - output, output_h = model( - input_data=input_data, input_h=input_h, params=params) - output_sum = math_ops.reduce_sum(output) - output_h_sum = math_ops.reduce_sum(output_h) - total_sum = output_sum + output_h_sum - if has_input_c: - output_c_sum = math_ops.reduce_sum(output_c) - total_sum += output_c_sum - - with self.test_session(use_gpu=True, graph=ops.get_default_graph()) as sess: - params_size_v = sess.run(params_size_t) - inputs_and_shapes = [ - (input_data, [seq_length, batch_size, input_size]), - (input_h, [num_layers * dir_count, batch_size, num_units]), - (params, [params_size_v]), - ] - if has_input_c: - inputs_and_shapes.append( - (input_c, [num_layers * dir_count, batch_size, num_units]),) - sess.run(variables.global_variables_initializer()) - all_inputs = [entry[0] for entry in inputs_and_shapes] - all_shapes = [entry[1] for entry in inputs_and_shapes] - - err = gradient_checker.compute_gradient_error( - all_inputs, all_shapes, total_sum, [1], delta=delta) - - self.assertLess(err, tolerance) - os.environ["TF_CUDNN_RESET_RND_GEN_STATE"] = old_env_state + def test_lstm_bidi(self, num_units, input_size, num_layers): + if not context.context().num_gpus(): + self.skipTest("No GPUs found") + self._test_lstm_helper(num_units, input_size, num_layers, + cudnn_rnn_ops.CUDNN_RNN_BIDIRECTION) + + def _test_gru_helper(self, num_units, input_size, num_layers, direction): + with self.session(use_gpu=True) as sess: + random_seed.set_random_seed(0) + np.random.seed(0) + num_dirs = 1 if direction == cudnn_rnn_ops.CUDNN_RNN_UNIDIRECTION else 2 + format_converter = cudnn_rnn_ops.CudnnParamsFormatConverterGRU( + num_layers, num_units, input_size, direction=direction) + + ws, bs = [], [] + for _ in range(num_layers * num_dirs): + gate_kernel = constant_op.constant( + np.random.rand(input_size + num_units, num_units * 2), + dtype=dtypes.float32) + gate_bias = constant_op.constant( + np.random.rand(num_units * 2), dtype=dtypes.float32) + candidate_inp_kernel = constant_op.constant( + np.random.rand(input_size, num_units), dtype=dtypes.float32) + candidate_inp_bias = constant_op.constant( + np.random.rand(num_units), dtype=dtypes.float32) + candidate_hid_kernel = constant_op.constant( + np.random.rand(num_units, num_units), dtype=dtypes.float32) + candidate_hid_bias = constant_op.constant( + np.random.rand(num_units), dtype=dtypes.float32) + ws.extend([gate_kernel, candidate_inp_kernel, candidate_hid_kernel]) + bs.extend([gate_bias, candidate_inp_bias, candidate_hid_bias]) + + opaque_params = format_converter.tf_canonical_to_opaque(ws + bs) + opaque_params_size = cudnn_rnn_ops.cudnn_rnn_opaque_params_size( + cudnn_rnn_ops.CUDNN_GRU, + num_layers, + num_units, + input_size, + direction=direction) + + ws_r, bs_r = format_converter.opaque_to_tf_canonical(opaque_params) + + # Test tf_canonical_to_opaque() followed by opaque_to_tf_canonical() + # returns the original input. + ws, ws_r, bs, bs_r = sess.run([ws, ws_r, bs, bs_r]) + for w, w_r in zip(ws, ws_r): + self.assertAllClose(w, w_r) + for b, b_r in zip(bs, bs_r): + self.assertAllClose(b, b_r) + + # Test opaque_params size lower bound + opaque_params_size_v = sess.run(opaque_params_size) + min_params_size = sum(x.size for x in ws) + sum(x.size for x in bs) + logging.info("min_parm_size: %d vs actual_opaque_param_size: %d", + min_params_size, opaque_params_size_v) + self.assertLessEqual(min_params_size, opaque_params_size_v) + + @parameterized.named_parameters((c["testcase_name"], c["num_units"], + c["input_size"], c["num_layers"]) + for c in NAMED_RNN_TESTCASES) + @unittest.skipUnless(test.is_built_with_cuda(), + "Test only applicable when running on GPUs") + def test_gru(self, num_units, input_size, num_layers): + if not context.context().num_gpus(): + self.skipTest("No GPUs found") + self._test_gru_helper(num_units, input_size, num_layers, + cudnn_rnn_ops.CUDNN_RNN_UNIDIRECTION) + + @parameterized.named_parameters((c["testcase_name"], c["num_units"], + c["input_size"], c["num_layers"]) + for c in NAMED_RNN_TESTCASES) + @unittest.skipUnless(test.is_built_with_cuda(), + "Test only applicable when running on GPUs") + def test_gru_bidi(self, num_units, input_size, num_layers): + if not context.context().num_gpus(): + self.skipTest("No GPUs found") + self._test_gru_helper(num_units, input_size, num_layers, + cudnn_rnn_ops.CUDNN_RNN_BIDIRECTION) + + +class CudnnRnnSaveRestoreTest(TensorFlowTestCase, parameterized.TestCase): + """Class for testing various Cudnn Rnn SaveableObjects.""" + + def _create_opaque_param(self, + rnn_mode, + num_units, + input_size, + num_layers, + direction, + name=None): + param_size_t = cudnn_rnn_ops.cudnn_rnn_opaque_params_size( + rnn_mode, num_layers, num_units, input_size, direction=direction) + init_val = random_ops.random_uniform([param_size_t]) + return variable_scope.get_variable( + name or "opaque_param", initializer=init_val, validate_shape=False) + + def _create_saveable(self, opaque_param, rnn_mode, num_units, input_size, + num_layers, direction): + if rnn_mode == CUDNN_LSTM: + fn = cudnn_rnn_ops.CudnnLSTMSaveable + elif rnn_mode == CUDNN_GRU: + fn = cudnn_rnn_ops.CudnnGRUSaveable + elif rnn_mode == CUDNN_RNN_TANH: + fn = cudnn_rnn_ops.CudnnRNNTanhSaveable + elif rnn_mode == CUDNN_RNN_RELU: + fn = cudnn_rnn_ops.CudnnRNNReluSaveable + saveable = fn( + opaque_param, num_layers, num_units, input_size, direction=direction) + return saveable + + def _compare_weights(self, lhs, rhs): + self.assertLen(rhs, len(lhs)) + for lw, rw in zip(lhs, rhs): + self.assertAllEqual(lw, rw) + + def _compare_biases(self, lhs, rhs): + self.assertLen(rhs, len(lhs)) + for lf, rt in zip(lhs, rhs): + self.assertAllEqual(lf, rt) + + @parameterized.named_parameters( + ExpandNamedTestCases( + NAMED_RNN_TESTCASES, "time", "batch_size", **{ + "rnn_mode": [ + CUDNN_LSTM, CUDNN_GRU, CUDNN_RNN_RELU, CUDNN_RNN_TANH + ], + "direction": [CUDNN_RNN_UNIDIRECTION, CUDNN_RNN_BIDIRECTION] + })) @unittest.skipUnless(test.is_built_with_cuda(), "Test only applicable when running on GPUs") - def testSimpleTraining(self): - test_configs = [ - { - "rnn_mode": cudnn_rnn_ops.CUDNN_LSTM, - "dtype": dtypes.float64, - "delta": 1e-4, - "tolerance": 5e-6, - "shape": { - "num_layers": 2, - "num_units": 3, - "input_size": 4, - "batch_size": 3, - "seq_length": 4, - "dir_count": 1, - }, - }, - { - "rnn_mode": cudnn_rnn_ops.CUDNN_GRU, - "dtype": dtypes.float64, - "delta": 1e-4, - "tolerance": 5e-6, - "shape": { - "num_layers": 2, - "num_units": 3, - "input_size": 4, - "batch_size": 3, - "seq_length": 4, - "dir_count": 1, - }, - }, - { - "rnn_mode": cudnn_rnn_ops.CUDNN_RNN_TANH, - "dtype": dtypes.float64, - "delta": 1e-4, - "tolerance": 5e-6, - "shape": { - "num_layers": 2, - "num_units": 3, - "input_size": 4, - "batch_size": 3, - "seq_length": 4, - "dir_count": 1, - }, - }, - { - "rnn_mode": cudnn_rnn_ops.CUDNN_RNN_RELU, - "dtype": dtypes.float64, - "delta": 1e-4, - "tolerance": 5e-6, - "shape": { - "num_layers": 2, - "num_units": 3, - "input_size": 4, - "batch_size": 3, - "seq_length": 4, - "dir_count": 1, - }, - }, - { - "rnn_mode": cudnn_rnn_ops.CUDNN_LSTM, - "dtype": dtypes.float32, - "tolerance": 1.5e-2, - "shape": { - "num_layers": 2, - "num_units": 3, - "input_size": 4, - "batch_size": 3, - "seq_length": 4, - }, - }, - { - "rnn_mode": cudnn_rnn_ops.CUDNN_GRU, - "dtype": dtypes.float32, - "tolerance": 4e-3, - "shape": { - "num_layers": 2, - "num_units": 3, - "input_size": 4, - "batch_size": 3, - "seq_length": 4, - }, - }, - { - "rnn_mode": cudnn_rnn_ops.CUDNN_RNN_TANH, - "dtype": dtypes.float32, - "tolerance": 5e-3, - "shape": { - "num_layers": 2, - "num_units": 3, - "input_size": 4, - "batch_size": 3, - "seq_length": 4, - }, - }, - { - "rnn_mode": cudnn_rnn_ops.CUDNN_RNN_RELU, - "dtype": dtypes.float32, - "tolerance": 5e-1, - "shape": { - "num_layers": 2, - "num_units": 3, - "input_size": 4, - "batch_size": 3, - "seq_length": 4, - }, - }, - ] - dropouts = [0., 0.5, 1.] - dir_counts = [1] - for config, dropout, dir_count in itertools.product(test_configs, dropouts, - dir_counts): - rnn_mode = config["rnn_mode"] - dtype = config.get("dtype", dtypes.float32) - delta = config.get("delta", 1e-3) - tolerance = config["tolerance"] - shape = config["shape"] - with ops.Graph().as_default(): - self._testOneSimpleTraining(rnn_mode, shape["num_layers"], - shape["num_units"], shape["input_size"], - shape["batch_size"], shape["seq_length"], - dir_count, dropout, dtype, delta, tolerance) + def test_save_restore_variable(self, rnn_mode, num_units, input_size, + num_layers, direction): + # Verify the restored opaque param, once converted to tf_canonical format, + # is the same as the tf canonicals of the pre-restored param. + if not context.context().num_gpus(): + self.skipTest("No GPUs found") + with self.session(use_gpu=True) as sess: + opaque_param = self._create_opaque_param(rnn_mode, num_units, input_size, + num_layers, direction) + saveable = self._create_saveable(opaque_param, rnn_mode, num_units, + input_size, num_layers, direction) + ops.add_to_collection(ops.GraphKeys.SAVEABLE_OBJECTS, saveable) + weights_op, biases_op = saveable.format_converter.opaque_to_tf_canonical( + saveable._variables) + + save_path = os.path.join(self.get_temp_dir(), "save_restore_var_test") + saver = saver_lib.Saver(write_version=saver_pb2.SaverDef.V2) + + init_op = variables.global_variables_initializer() + reset_op = state_ops.assign(opaque_param, + array_ops.zeros_like(opaque_param)) + sess.run(init_op) + self.assertEqual(save_path, saver.save(sess, save_path)) + + # Get the tf canonical vals before reset-restore + weights, biases = sess.run([weights_op, biases_op]) + + # Reset the opaque param value + sess.run(reset_op) + # Assert reset happened. + weights_z, biases_z = sess.run([weights_op, biases_op]) + for w in weights_z: + self.assertAllClose(w, np.zeros_like(w)) + for b in biases_z: + self.assertAllClose(b, np.zeros_like(b)) + + # Restore opaque param value from checkpoint. + saver.restore(sess, save_path) + weights_r, biases_r = sess.run([weights_op, biases_op]) + self._compare_weights(weights, weights_r) + self._compare_biases(biases, biases_r) + + @parameterized.named_parameters( + ExpandNamedTestCases( + NAMED_RNN_TESTCASES, "time", "batch_size", **{ + "rnn_mode": [ + CUDNN_LSTM, CUDNN_GRU, CUDNN_RNN_RELU, CUDNN_RNN_TANH + ], + "direction": [CUDNN_RNN_UNIDIRECTION, CUDNN_RNN_BIDIRECTION] + })) + @unittest.skipUnless(test.is_built_with_cuda(), + "Test only applicable when running on GPUs") + def test_save_restore_multi_variables(self, rnn_mode, num_units, input_size, + num_layers, direction): + # Verify the restored opaque param, once converted to tf_canonical format, + # is the same as the tf canonicals of the pre-restored param. + if not context.context().num_gpus(): + self.skipTest("No GPUs found") + with self.session(use_gpu=True) as sess: + opaque_params = [] + saveables = [] + num_opaque_params = 2 + for i in range(num_opaque_params): + opaque_params.append( + self._create_opaque_param( + rnn_mode, + num_units, + input_size, + num_layers, + direction, + name="opaque_param_%d" % i)) + saveable = self._create_saveable(opaque_params[i], rnn_mode, num_units, + input_size, num_layers, direction) + ops.add_to_collection(ops.GraphKeys.SAVEABLE_OBJECTS, saveable) + saveables.append(saveable) + + weights_ops, biases_ops = [], [] + for i in range(num_opaque_params): + weights_op, biases_op = ( + saveables[i].format_converter.opaque_to_tf_canonical( + saveables[i]._variables)) + weights_ops.append(weights_op) + biases_ops.append(biases_op) + + save_path = os.path.join(self.get_temp_dir(), "save_restore_var_test") + saver = saver_lib.Saver(write_version=saver_pb2.SaverDef.V2) + + init_op = variables.global_variables_initializer() + reset_ops = [] + for i in range(num_opaque_params): + reset_ops.append( + state_ops.assign(opaque_params[i], + array_ops.zeros_like(opaque_params[i]))) + sess.run(init_op) + self.assertEqual(save_path, saver.save(sess, save_path)) + + # Get the tf canonical vals before reset-restore + for i in range(num_opaque_params): + weights, biases = sess.run([weights_ops[i], biases_ops[i]]) + + # Reset the opaque param value + sess.run(reset_ops[i]) + + # Assert reset happened. + weights_z, biases_z = sess.run([weights_ops[i], biases_ops[i]]) + for w in weights_z: + self.assertAllClose(w, np.zeros_like(w)) + for b in biases_z: + self.assertAllClose(b, np.zeros_like(b)) + + # Restore opaque param value from checkpoint. + saver.restore(sess, save_path) + weights_r, biases_r = sess.run([weights_ops[i], biases_ops[i]]) + self._compare_weights(weights, weights_r) + self._compare_biases(biases, biases_r) if __name__ == "__main__": 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 d3a1b33233ce1bae1e00a2bcfa4ff5d4f537f4ae..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 @@ -536,7 +536,9 @@ class CudnnRNNTestSaveRestore(test_util.TensorFlowTestCase): save_path = os.path.join(self.get_temp_dir(), "save-restore-variable-test") saver = saver_lib.Saver() - weights, biases = model.rnn.saveable._OpaqueParamsToCanonical() + weights, biases = ( + model.rnn.saveable.format_converter._opaque_to_cu_canonical( + model.rnn.saveable._variables)) opaque_params = rnn.trainable_variables[0] # CudnnTestModel() creates CudnnOpaqueParamsSaveable that helps saver save # Cudnn vars in canonical format. @@ -583,8 +585,12 @@ class CudnnRNNTestSaveRestore(test_util.TensorFlowTestCase): dtype=dtype) opaque_params = (model1.rnn.trainable_variables[0], model2.rnn.trainable_variables[0]) - weights1, biases1 = model1.rnn.saveable._OpaqueParamsToCanonical() - weights2, biases2 = model2.rnn.saveable._OpaqueParamsToCanonical() + saveable1 = model1.rnn.saveable + weights1, biases1 = saveable1.format_converter._opaque_to_cu_canonical( + saveable1._variables) + saveable2 = model1.rnn.saveable + weights2, biases2 = saveable2.format_converter._opaque_to_cu_canonical( + saveable2._variables) reset_params = [ state_ops.assign(params, array_ops.zeros_like(params, dtype=dtype)) @@ -1017,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) @@ -1039,8 +1045,8 @@ class CudnnRNNTestParamsSize(test_util.TensorFlowTestCase): # Min param size estimate = sum(weights.size) + sum(biases.size) min_params_size = ( - np.sum(map(np.prod, rnn.canonical_weight_shapes)) + - np.sum([sp[0] for sp in rnn.canonical_bias_shapes])) + sum(map(np.prod, rnn.canonical_weight_shapes)) + + sum(sp[0] for sp in rnn.canonical_bias_shapes)) opaque_params = rnn.trainable_variables[0] with self.test_session(use_gpu=True, graph=ops.get_default_graph()): diff --git a/tensorflow/contrib/cudnn_rnn/python/layers/__init__.py b/tensorflow/contrib/cudnn_rnn/python/layers/__init__.py index f09466b631f69d6234573dd5eafada650421c117..60229af374be869005139921483793156e5e7a05 100644 --- a/tensorflow/contrib/cudnn_rnn/python/layers/__init__.py +++ b/tensorflow/contrib/cudnn_rnn/python/layers/__init__.py @@ -27,5 +27,10 @@ from tensorflow.contrib.cudnn_rnn.python.ops.cudnn_rnn_ops import CudnnCompatibl from tensorflow.contrib.cudnn_rnn.python.ops.cudnn_rnn_ops import CudnnCompatibleLSTMCell from tensorflow.contrib.cudnn_rnn.python.ops.cudnn_rnn_ops import CudnnGRUSaveable from tensorflow.contrib.cudnn_rnn.python.ops.cudnn_rnn_ops import CudnnLSTMSaveable +from tensorflow.contrib.cudnn_rnn.python.ops.cudnn_rnn_ops import CudnnParamsFormatConverterGRU +from tensorflow.contrib.cudnn_rnn.python.ops.cudnn_rnn_ops import CudnnParamsFormatConverterLSTM +from tensorflow.contrib.cudnn_rnn.python.ops.cudnn_rnn_ops import CudnnParamsFormatConverterRelu +from tensorflow.contrib.cudnn_rnn.python.ops.cudnn_rnn_ops import CudnnParamsFormatConverterTanh from tensorflow.contrib.cudnn_rnn.python.ops.cudnn_rnn_ops import CudnnRNNReluSaveable from tensorflow.contrib.cudnn_rnn.python.ops.cudnn_rnn_ops import CudnnRNNTanhSaveable + diff --git a/tensorflow/contrib/cudnn_rnn/python/layers/cudnn_rnn.py b/tensorflow/contrib/cudnn_rnn/python/layers/cudnn_rnn.py index a324c6e7d76223aaa6514e695e4ff8444db455d0..86ad8ae8073714657c78badb1e0b4a6d8c8ed5f0 100644 --- a/tensorflow/contrib/cudnn_rnn/python/layers/cudnn_rnn.py +++ b/tensorflow/contrib/cudnn_rnn/python/layers/cudnn_rnn.py @@ -21,6 +21,7 @@ from tensorflow.contrib.cudnn_rnn.python.ops import cudnn_rnn_ops from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.framework import tensor_shape +from tensorflow.python.keras.engine import input_spec from tensorflow.python.layers import base as base_layer from tensorflow.python.ops import array_ops from tensorflow.python.ops import init_ops @@ -322,7 +323,7 @@ class _CudnnRNN(base_layer.Layer): raise ValueError("The last dimension of the inputs to `CudnnRNN` " "should be defined. Found `None`.") self._input_size = input_shape[-1].value - self.input_spec = base_layer.InputSpec(ndim=3, axes={-1: self._input_size}) + self.input_spec = input_spec.InputSpec(ndim=3, axes={-1: self._input_size}) self._set_scope(None) @@ -373,7 +374,11 @@ class _CudnnRNN(base_layer.Layer): "This cell does not yet support object-based saving. File a feature " "request if this limitation bothers you.") - def call(self, inputs, initial_state=None, training=True): + def call(self, + inputs, + initial_state=None, + sequence_lengths=None, + training=True): """Runs the forward step for the RNN model. Args: @@ -381,6 +386,9 @@ class _CudnnRNN(base_layer.Layer): initial_state: a tuple of tensor(s) of shape `[num_layers * num_dirs, batch_size, num_units]`. If not provided, use zero initial states. The tuple size is 2 for LSTM and 1 for other RNNs. + sequence_lengths: an int32 array representing the variable sequence + lengths in a batch. The size of the array has to equal the + batch_size. If not provided, the same sequence length will be assumed. training: whether this operation will be used in training or inference. Returns: output: a tensor of shape `[time_len, batch_size, num_dirs * num_units]`. @@ -388,11 +396,11 @@ class _CudnnRNN(base_layer.Layer): output_states: a tuple of tensor(s) of the same shape and structure as `initial_state`. Raises: - ValueError: initial_state is not a tuple. + TypeError: initial_state is not a tuple. """ if initial_state is not None and not isinstance(initial_state, tuple): - raise ValueError("Invalid initial_state type: %s, expecting tuple.", - type(initial_state)) + raise TypeError("Invalid initial_state type: %s, expecting tuple." % + initial_state) dtype = self.dtype inputs = ops.convert_to_tensor(inputs, dtype=dtype) @@ -410,7 +418,7 @@ class _CudnnRNN(base_layer.Layer): # For model that doesn't take input_c, replace with a dummy tensor. c = array_ops.constant([], dtype=dtype) outputs, (output_h, output_c) = self._forward(inputs, h, c, self.kernel, - training) + sequence_lengths, training) if self._rnn_mode == CUDNN_LSTM: return outputs, (output_h, output_c) else: @@ -474,7 +482,7 @@ class _CudnnRNN(base_layer.Layer): dropout=self._dropout, direction=self._direction) - def _forward(self, inputs, h, c, opaque_params, training): + def _forward(self, inputs, h, c, opaque_params, sequence_lengths, training): output, output_h, output_c = cudnn_rnn_ops._cudnn_rnn( # pylint:disable=protected-access inputs, h, @@ -482,6 +490,7 @@ class _CudnnRNN(base_layer.Layer): opaque_params, training, self._rnn_mode, + sequence_lengths=sequence_lengths, input_mode=self._input_mode, direction=self._direction, dropout=self._dropout, 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 2c92f31788378c2a9f01183bc04b035668b59b59..f36e8d5022bc7e3f8268a161089153e5510dffc6 100644 --- a/tensorflow/contrib/cudnn_rnn/python/ops/cudnn_rnn_ops.py +++ b/tensorflow/contrib/cudnn_rnn/python/ops/cudnn_rnn_ops.py @@ -74,7 +74,7 @@ class CudnnCompatibleLSTMCell(lstm_ops.LSTMBlockCell): class CudnnCompatibleGRUCell(rnn_cell_impl.GRUCell): - """Cudnn Compatible GRUCell. + r"""Cudnn Compatible GRUCell. A GRU impl akin to `tf.nn.rnn_cell.GRUCell` to use along with `tf.contrib.cudnn_rnn.CudnnGRU`. The latter's params can be used by @@ -177,172 +177,60 @@ class CudnnCompatibleGRUCell(rnn_cell_impl.GRUCell): return new_h, new_h -# TODO(yaozhang): make sure we only save the canonical version of params and -# don't save the platform-specific version to avoid potential race -# conditions where params is updated by both versions when being restored. -# Currently, checkpointing will function properly, despite that we save both -# versions, because Saver restores customized savables after Variables. -# However, it is good to not rely on this restoring order of Saver and to -# avoid unnecessary storage. Add a test to check only the canonical version is -# saved. -class CudnnOpaqueParamsSaveable(saver.BaseSaverBuilder.SaveableObject): - """Abstract SaveableObject implementation handling Cudnn opaque params.""" +class CudnnParamsFormatConverter(object): + """Abstract class that converts between params of Cudnn Rnn and TF Rnn.""" def __init__(self, - opaque_params, num_layers, num_units, input_size, input_mode=CUDNN_INPUT_LINEAR_MODE, - direction=CUDNN_RNN_UNIDIRECTION, - scope=None, - name="cudnn_rnn_saveable"): - """Creates a CudnnOpaqueParamsSaveable object. - - CudnnOpaqueParamsSaveable is saveable/restorable in a checkpoint file - and is used to save/restore the weights and biases parameters in a - canonical format which is directly consumable by platform-independent tf - RNN cells. Parameters are saved as tensors layer by layer with weight - tensors followed by bias tensors, and forward direction followed by - backward direction (if applicable). When restoring, a user could name - param_variables as desired, and restore weight and bias tensors to these - variables. - - For CudnnRNNRelu or CudnnRNNTanh, there are 2 tensors per weight and per - bias for each layer: tensor 0 is applied to the input from the previous - layer and tensor 1 to the recurrent input. - - For CudnnLSTM, there are 8 tensors per weight and per bias for each - layer: tensor 0-3 are applied to the input from the previous layer and - tensor 4-7 to the recurrent input. Tensor 0 and 4 are for the input gate; - tensor 1 and 5 the forget gate; tensor 2 and 6 the new memory gate; - tensor 3 and 7 the output gate. - - For CudnnGRU, there are 6 tensors per weight and per bias for each layer: - tensor 0-2 are applied to the input from the previous layer and - tensor 3-5 to the recurrent input. Tensor 0 and 3 are for the reset gate; - tensor 1 and 4 the update gate; tensor 2 and 5 the new memory gate. + direction=CUDNN_RNN_UNIDIRECTION): + """Constructor. Args: - opaque_params: a variable, Cudnn RNN opaque params. num_layers: the number of layers for the RNN model. num_units: the number of units within the RNN model. input_size: the size of the input, it could be different from the - num_units. + num_units. input_mode: indicate whether there is a linear projection between the - input and the actual computation before the first layer. It could be - 'linear_input', 'skip_input' or 'auto_select'. - 'linear_input' (default) always applies a linear projection of input - onto RNN hidden state. (standard RNN behavior). - 'skip_input' is only allowed when input_size == num_units; - 'auto_select' implies 'skip_input' when input_size == num_units; - otherwise, it implies 'linear_input'. + input and the actual computation before the first layer. It could be one + of 'linear_input', 'skip_input' or 'auto_select'. * 'linear_input' + (default) always applies a linear projection of input onto RNN hidden + state. (standard RNN behavior). * 'skip_input' is only allowed when + input_size == num_units; * 'auto_select' implies 'skip_input' when + input_size == num_units; otherwise, it implies 'linear_input'. direction: the direction model that the model operates. Could be either - 'unidirectional' or 'bidirectional' - scope: string of VariableScope, the scope of equivalent subgraph - consisting only platform-independent tf RNN cells. - name: the name of the CudnnOpaqueParamsSaveable object. + 'unidirectional' or 'bidirectional' """ - # Define in subclasses. self._num_layers = num_layers self._input_size = input_size self._num_units = num_units self._input_mode = input_mode self._direction = direction - if scope is not None: - scope_name = scope.name if isinstance(scope, vs.VariableScope) else scope - self._scope = scope_name or None - else: - self._scope = None - - self._variables = opaque_params self._num_dirs = 1 if self._direction == CUDNN_RNN_UNIDIRECTION else 2 self._num_params = ( self._num_params_per_layer * self._num_layers * self._num_dirs) - weights, biases = self._OpaqueParamsToCanonical() - (weights, weight_names), (biases, bias_names) = self._TransformCanonical( - weights, biases) - # We currently don't use slice_spec. It might be useful in a distributed - # setting where each parameter server node stores a slice of variable, - # instead of having the master pull all slices and then save them. - slice_spec = "" - params = weights + biases - self._weight_names = weight_names - self._bias_names = bias_names - self._param_names = weight_names + bias_names - prefixed_param_names = weight_names + bias_names - if self._scope: - prefixed_param_names = [ - "%s/%s" % (self._scope, pn) for pn in prefixed_param_names] - specs = [ - saver.BaseSaverBuilder.SaveSpec(param, slice_spec, param_name) - for param, param_name in zip(params, prefixed_param_names) - ] - super(CudnnOpaqueParamsSaveable, self).__init__( - array_ops.identity(self._variables), specs, name) - - def restore(self, restored_tensors, restored_shapes): - weights, biases = self._ReverseTransformCanonical(restored_tensors) - weights = [array_ops.reshape(w, [-1]) for w in weights] - opaque_params = self._CanonicalToOpaqueParams(weights, biases) - - return state_ops.assign( - self._variables, opaque_params, validate_shape=False) - - def _checkpointable_save(self, save_buffer): - weights, biases = self._OpaqueParamsToCanonical() - with ops.device("gpu:0"): - (weights, _), (biases, _) = self._TransformCanonical( - weights, biases) - for name, tensor in zip(self._param_names, weights + biases): - save_buffer[name] = array_ops.identity(tensor) - - def _checkpointable_restore(self, restore_buffer): - tensors = [array_ops.identity(restore_buffer[name]) - for name in self._param_names] - return self.restore( - restored_tensors=tensors, - restored_shapes=None # Unused - ) - - def _add_checkpointable_dependencies(self, checkpointable, dtype): - """Add canonical weight dependencies to `checkpointable`. - - When saving or restoring, converts to or from the opaque buffer - format. Weights are saved and loaded in the configuration expected by - cuDNN-compatible cells. - - Args: - checkpointable: An object inheriting from `CheckpointableBase` to add - dependencies too (typically the cuDNN `Layer`). - dtype: The dtype for the canonical parameter Tensors. - """ - split_dependencies = split_dependency.split_dependency( - component_names=self._param_names, - component_dtypes=(dtype,) * len(self._param_names), - fill_save_buffer_fn=self._checkpointable_save, - consume_restore_buffer_fn=self._checkpointable_restore) - self._checkpointable_track_params(checkpointable, split_dependencies) - - def _checkpointable_track_params(self, checkpointable, params): - """Tracks parameters in a canonical configuration.""" - return # NotImplementedError raised by the Layer. + def tf_canonical_to_opaque(self, tf_canonicals): + r"""Converts tf canonical weights to cudnn opaque param.""" + cu_weights, cu_biases = self._tf_canonical_to_cu_canonical(tf_canonicals) + cu_weights = [array_ops.reshape(w, [-1]) for w in cu_weights] + opaque_params = self._cu_canonical_to_opaque(cu_weights, cu_biases) + return opaque_params - def _TFCanonicalNamePrefix(self, layer, is_fwd=True): - if self._direction == CUDNN_RNN_UNIDIRECTION: - return "rnn/multi_rnn_cell/cell_%d/%s" % (layer, self._rnn_cell_name) - else: - if is_fwd: - return ("stack_bidirectional_rnn/cell_%d/bidirectional_rnn/fw/%s" % - (layer, self._rnn_cell_name)) - else: - return ("stack_bidirectional_rnn/cell_%d/bidirectional_rnn/bw/%s" % - (layer, self._rnn_cell_name)) + def opaque_to_tf_canonical(self, opaque_param): + r"""Converts cudnn opaque param to tf canonical weights.""" + cu_weights, cu_biases = self._opaque_to_cu_canonical(opaque_param) + weights, biases = self._cu_canonical_to_tf_canonical(cu_weights, cu_biases) + return weights, biases - def _OpaqueParamsToCanonical(self): + def _opaque_to_cu_canonical(self, opaque_param): """Converts opaque params to Cudnn canonical format. + Args: + opaque_param: An opaque tensor storing cudnn rnn params (weights and + biases). Returns: 2 list for weights and biases respectively. """ @@ -351,14 +239,14 @@ class CudnnOpaqueParamsSaveable(saver.BaseSaverBuilder.SaveableObject): num_layers=self._num_layers, num_units=self._num_units, input_size=self._input_size, - params=self._variables, + params=opaque_param, num_params=self._num_params, rnn_mode=self._rnn_mode, input_mode=self._input_mode, direction=self._direction) return (weights, biases) - def _CanonicalToOpaqueParams(self, cu_weights, cu_biases): + def _cu_canonical_to_opaque(self, cu_weights, cu_biases): """Converts from Cudnn canonical format to opaque params. Args: @@ -378,7 +266,7 @@ class CudnnOpaqueParamsSaveable(saver.BaseSaverBuilder.SaveableObject): input_mode=self._input_mode, direction=self._direction) - def _TransformCanonical(self, cu_weights, cu_biases): + def _cu_canonical_to_tf_canonical(self, cu_weights, cu_biases): r"""Transform from Cudnn canonical to tf canonical. The elements of argument lists are laid out in the following format: @@ -398,46 +286,43 @@ class CudnnOpaqueParamsSaveable(saver.BaseSaverBuilder.SaveableObject): cu_weights: a list of tensors of Cudnn canonical weights. cu_biases: a list of tensors of Cudnn canonical biases. Returns: - 2 tuples, one for weights and the other for bias. - Each tuple has two lists: the 1st for transformed tf canonical tensors - and the 2nd for the names of the tensors under which they are saved. + 1 tuple, tf canonical weights and biases. """ tf_weights, tf_biases = [], [] - tf_weights_names, tf_bias_names = [], [] layer_weights_num = self._num_params_per_layer * self._num_dirs layer_biases_num = layer_weights_num for i in range(self._num_layers): - layer_weights = cu_weights[i * layer_weights_num: - (i + 1) * layer_weights_num] + layer_weights = cu_weights[i * layer_weights_num:(i + 1) * + layer_weights_num] layer_biases = cu_biases[i * layer_biases_num:(i + 1) * layer_biases_num] if self._direction == CUDNN_RNN_UNIDIRECTION: - prefix = self._TFCanonicalNamePrefix(i) - self._TransformSingleLayerCanonical(layer_weights, layer_biases, prefix, - tf_weights, tf_weights_names, - tf_biases, tf_bias_names) + self._cu_canonical_to_tf_canonical_single_layer( + layer_weights, layer_biases, tf_weights, tf_biases) else: - fw_prefix = self._TFCanonicalNamePrefix(i, is_fwd=True) - bw_prefix = self._TFCanonicalNamePrefix(i, is_fwd=False) - fw_weights = layer_weights[:len(layer_weights) // 2] bw_weights = layer_weights[len(layer_weights) // 2:] fw_biases = layer_biases[:len(layer_biases) // 2] bw_biases = layer_biases[len(layer_biases) // 2:] - self._TransformSingleLayerCanonical(fw_weights, fw_biases, fw_prefix, - tf_weights, tf_weights_names, - tf_biases, tf_bias_names) - - self._TransformSingleLayerCanonical(bw_weights, bw_biases, bw_prefix, - tf_weights, tf_weights_names, - tf_biases, tf_bias_names) - return (tf_weights, tf_weights_names), (tf_biases, tf_bias_names) - - def _TransformSingleLayerCanonical(self, cu_weights, cu_biases, prefix, - tf_weights, tf_weights_names, tf_biases, - tf_bias_names): + self._cu_canonical_to_tf_canonical_single_layer( + fw_weights, + fw_biases, + tf_weights, + tf_biases, + ) + + self._cu_canonical_to_tf_canonical_single_layer( + bw_weights, + bw_biases, + tf_weights, + tf_biases, + ) + return (tf_weights, tf_biases) + + def _cu_canonical_to_tf_canonical_single_layer(self, cu_weights, cu_biases, + tf_weights, tf_biases): r"""Transform single layer Cudnn canonicals to tf canonicals. The elements of cu_weights, cu_biases are laid out in the following format: @@ -447,15 +332,12 @@ class CudnnOpaqueParamsSaveable(saver.BaseSaverBuilder.SaveableObject): Args: cu_weights: a list of tensors, single layer weights. cu_biases: a list of tensors, single layer biases. - prefix: the shared prefix of all tensor names. tf_weights: a list where transformed weights are stored. - tf_weights_names: a list where names of transformed weights are stored. tf_biases: a list where transformed biases are stored. - tf_bias_names: a list where names of transformed biases are stored. """ raise NotImplementedError("Abstract method") - def _ReverseTransformCanonical(self, tf_canonicals): + def _tf_canonical_to_cu_canonical(self, tf_canonicals): r"""Transform from tf canonical to Cudnn canonical. This is the reverse routine of _TransformCanonical(). @@ -502,30 +384,27 @@ class CudnnOpaqueParamsSaveable(saver.BaseSaverBuilder.SaveableObject): return cu_weights, cu_biases def _cudnn_to_tf_weights(self, *cu_weights): - r"""Stitching cudnn canonical weights to generate tf canonical weights.""" + r"""Stitches cudnn canonical weights to generate tf canonical weights.""" raise NotImplementedError("Abstract method") def _tf_to_cudnn_weights(self, layer, *tf_weights): - r"""Reverse the operations in StitchWeights().""" + r"""Reverses the operations in StitchWeights().""" raise NotImplementedError("Abstract method") def _cudnn_to_tf_biases(self, *biases): - r"""Stitching cudnn canonical biases to generate tf canonical biases.""" + r"""Stitches cudnn canonical biases to generate tf canonical biases.""" raise NotImplementedError("Abstract method") def _tf_to_cudnn_biases(self, *tf_biases): - r"""Reverse the operations in StitchBiases().""" + r"""Reverses the operations in StitchBiases().""" raise NotImplementedError("Abstract method") -class CudnnLSTMSaveable(CudnnOpaqueParamsSaveable): - """SaveableObject implementation handling Cudnn LSTM opaque params.""" - +class CudnnParamsFormatConverterLSTM(CudnnParamsFormatConverter): + """Helper class that converts between params of Cudnn and TF LSTM.""" _rnn_mode = CUDNN_LSTM _num_params_per_layer = CUDNN_LSTM_PARAMS_PER_LAYER - _rnn_cell_name = base_layer.to_snake_case(CudnnCompatibleLSTMCell.__name__) - def _cudnn_to_tf_gate_params(self, *cu_gate_order): i_g, f_g, c_g, o_g = cu_gate_order return [i_g, c_g, f_g, o_g] @@ -603,44 +482,16 @@ class CudnnLSTMSaveable(CudnnOpaqueParamsSaveable): # Return ifco order for Cudnn LSTM. return b_wi, b_wf, b_wc, b_wo, b_ri, b_rf, b_rc, b_ro - def _TransformSingleLayerCanonical(self, weights, biases, prefix, tf_weights, - tf_weights_names, tf_biases, - tf_bias_names): - (w,) = self._cudnn_to_tf_weights(*weights) - (b,) = self._cudnn_to_tf_biases(*biases) - + def _cu_canonical_to_tf_canonical_single_layer(self, cu_weights, cu_biases, + tf_weights, tf_biases): + (w,) = self._cudnn_to_tf_weights(*cu_weights) + (b,) = self._cudnn_to_tf_biases(*cu_biases) tf_weights.append(w) - tf_weights_names.append(prefix + "/kernel") - tf_biases.append(b) - tf_bias_names.append(prefix + "/bias") - - def _checkpointable_track_params(self, checkpointable, params): - """Track parameters for compatibility with CudnnCompatibleLSTMCell.""" - biases = [] - weights = [] - for name in self._weight_names: - weights.append(params[name]) - for name in self._bias_names: - biases.append(params[name]) - assert len(params) == len(weights) + len(biases) - if len(weights) == 1 and len(biases) == 1: - # For single-layer cells, allow substituting a cell with no MultiRNNCell - # wrapping. - kernel, = weights # pylint: disable=unbalanced-tuple-unpacking - bias, = biases # pylint: disable=unbalanced-tuple-unpacking - checkpointable._track_checkpointable(kernel, name="kernel") # pylint: disable=protected-access - 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() - checkpointable._track_checkpointable(cell, name="cell-%d" % cell_index) # pylint: disable=protected-access - cell.bias = bias - cell.kernel = kernel -class CudnnGRUSaveable(CudnnOpaqueParamsSaveable): - """SaveableObject implementation handling Cudnn GRU opaque params.""" +class CudnnParamsFormatConverterGRU(CudnnParamsFormatConverter): + """Helper class that converts between params of Cudnn and TF GRU.""" _rnn_mode = CUDNN_GRU _num_params_per_layer = CUDNN_GRU_PARAMS_PER_LAYER @@ -702,29 +553,18 @@ class CudnnGRUSaveable(CudnnOpaqueParamsSaveable): b_ri, b_rr = array_ops.split(br, 2, axis=0) return b_wi, b_wr, b_wh, b_ri, b_rr, b_rh - def _TransformSingleLayerCanonical(self, weights, biases, prefix, tf_weights, - tf_weights_names, tf_biases, - tf_bias_names): + def _cu_canonical_to_tf_canonical_single_layer(self, cu_weights, cu_biases, + tf_weights, tf_biases): # pylint: disable=invalid-name - W_ir, w_h, r_h = self._cudnn_to_tf_weights(*weights) - b_ir, b_wh, b_rh = self._cudnn_to_tf_biases(*biases) + W_ir, w_h, r_h = self._cudnn_to_tf_weights(*cu_weights) + b_ir, b_wh, b_rh = self._cudnn_to_tf_biases(*cu_biases) # pylint: enable=invalid-name - tf_weights.extend([W_ir, w_h, r_h]) - tf_weights_names.append(prefix + "/gates/kernel") - tf_weights_names.append(prefix + "/candidate/input_projection/kernel") - tf_weights_names.append(prefix + "/candidate/hidden_projection/kernel") - tf_biases.extend([b_ir, b_wh, b_rh]) - tf_bias_names.append(prefix + "/gates/bias") - tf_bias_names.append(prefix + "/candidate/input_projection/bias") - tf_bias_names.append(prefix + "/candidate/hidden_projection/bias") -class CudnnRNNSimpleSaveable(CudnnLSTMSaveable): - """SaveableObject implementation handling Cudnn RNN Tanh opaque params.""" - - _rnn_cell_name = base_layer.to_snake_case(rnn_cell_impl.BasicRNNCell.__name__) +class CudnnParamsFormatConverterBasic(CudnnParamsFormatConverterLSTM): + """Helper class that converts between params of Cudnn and TF Relu/Tanh RNN.""" def _cudnn_to_tf_weights(self, *cu_weights): r"""Stitching cudnn canonical weights to generate tf canonical weights.""" @@ -766,18 +606,270 @@ class CudnnRNNSimpleSaveable(CudnnLSTMSaveable): return b_i, b_h -class CudnnRNNTanhSaveable(CudnnRNNSimpleSaveable): - """SaveableObject implementation handling Cudnn RNN Tanh opaque params.""" +class CudnnParamsFormatConverterTanh(CudnnParamsFormatConverterBasic): + """Helper class that converts between params of Cudnn and TF Tanh RNN.""" _rnn_mode = CUDNN_RNN_TANH _num_params_per_layer = CUDNN_RNN_TANH_PARAMS_PER_LAYER -class CudnnRNNReluSaveable(CudnnRNNSimpleSaveable): - """SaveableObject implementation handling Cudnn RNN Relu opaque params.""" +class CudnnParamsFormatConverterRelu(CudnnParamsFormatConverterBasic): + """Helper class that converts between params of Cudnn and TF Relu RNN.""" _rnn_mode = CUDNN_RNN_RELU _num_params_per_layer = CUDNN_RNN_RELU_PARAMS_PER_LAYER +# TODO(yaozhang): make sure we only save the canonical version of params and +# don't save the platform-specific version to avoid potential race +# conditions where params is updated by both versions when being restored. +# Currently, checkpointing will function properly, despite that we save both +# versions, because Saver restores customized savables after Variables. +# However, it is good to not rely on this restoring order of Saver and to +# avoid unnecessary storage. Add a test to check only the canonical version is +# saved. +class CudnnOpaqueParamsSaveable(saver.BaseSaverBuilder.SaveableObject): + """Abstract SaveableObject implementation handling Cudnn opaque params.""" + + def __init__(self, + opaque_params, + num_layers, + num_units, + input_size, + input_mode=CUDNN_INPUT_LINEAR_MODE, + direction=CUDNN_RNN_UNIDIRECTION, + scope=None, + name="cudnn_rnn_saveable"): + """Creates a CudnnOpaqueParamsSaveable object. + + CudnnOpaqueParamsSaveable is saveable/restorable in a checkpoint file + and is used to save/restore the weights and biases parameters in a + canonical format which is directly consumable by platform-independent tf + RNN cells. Parameters are saved as tensors layer by layer with weight + tensors followed by bias tensors, and forward direction followed by + backward direction (if applicable). When restoring, a user could name + param_variables as desired, and restore weight and bias tensors to these + variables. + + For CudnnRNNRelu or CudnnRNNTanh, there are 2 tensors per weight and per + bias for each layer: tensor 0 is applied to the input from the previous + layer and tensor 1 to the recurrent input. + + For CudnnLSTM, there are 8 tensors per weight and per bias for each + layer: tensor 0-3 are applied to the input from the previous layer and + tensor 4-7 to the recurrent input. Tensor 0 and 4 are for the input gate; + tensor 1 and 5 the forget gate; tensor 2 and 6 the new memory gate; + tensor 3 and 7 the output gate. + + For CudnnGRU, there are 6 tensors per weight and per bias for each layer: + tensor 0-2 are applied to the input from the previous layer and + tensor 3-5 to the recurrent input. Tensor 0 and 3 are for the reset gate; + tensor 1 and 4 the update gate; tensor 2 and 5 the new memory gate. + + Args: + opaque_params: a variable, Cudnn RNN opaque params. + num_layers: the number of layers for the RNN model. + num_units: the number of units within the RNN model. + input_size: the size of the input, it could be different from the + num_units. + input_mode: indicate whether there is a linear projection between the + input and the actual computation before the first layer. It could be + 'linear_input', 'skip_input' or 'auto_select'. 'linear_input' (default) + always applies a linear projection of input onto RNN hidden state. + (standard RNN behavior). 'skip_input' is only allowed when input_size == + num_units; 'auto_select' implies 'skip_input' when input_size == + num_units; otherwise, it implies 'linear_input'. + direction: the direction model that the model operates. Could be either + 'unidirectional' or 'bidirectional' + scope: string of VariableScope, the scope of equivalent subgraph + consisting only platform-independent tf RNN cells. + name: the name of the CudnnOpaqueParamsSaveable object. + """ + # Define in subclasses. + self._num_layers = num_layers + self._input_size = input_size + self._num_units = num_units + self._input_mode = input_mode + self._direction = direction + if scope is not None: + scope_name = scope.name if isinstance(scope, vs.VariableScope) else scope + self._scope = scope_name or None + else: + self._scope = None + + self._variables = opaque_params + self._num_dirs = 1 if self._direction == CUDNN_RNN_UNIDIRECTION else 2 + # Defined in subclasses. + self._format_converter = None + + tf_weights, tf_biases = ( + self.format_converter.opaque_to_tf_canonical(self._variables)) + tf_weight_names, tf_bias_names = self._tf_canonical_names() + # We currently don't use slice_spec. It might be useful in a distributed + # setting where each parameter server node stores a slice of variable, + # instead of having the master pull all slices and then save them. + slice_spec = "" + params = tf_weights + tf_biases + self._weight_names = tf_weight_names + self._bias_names = tf_bias_names + self._param_names = tf_weight_names + tf_bias_names + prefixed_param_names = tf_weight_names + tf_bias_names + if self._scope: + prefixed_param_names = [ + "%s/%s" % (self._scope, pn) for pn in prefixed_param_names + ] + specs = [ + saver.BaseSaverBuilder.SaveSpec(param, slice_spec, param_name) + for param, param_name in zip(params, prefixed_param_names) + ] + super(CudnnOpaqueParamsSaveable, self).__init__( + array_ops.identity(self._variables), specs, name) + + @property + def format_converter(self): + if self._format_converter is None: + self._format_converter = self._format_converter_cls( + self._num_layers, self._num_units, self._input_size, self._input_mode, + self._direction) + return self._format_converter + + def restore(self, restored_tensors, restored_shapes): + opaque_params = self.format_converter.tf_canonical_to_opaque( + restored_tensors) + return state_ops.assign( + self._variables, opaque_params, validate_shape=False) + + def _checkpointable_save(self, save_buffer): + weights, biases = self.format_converter.opaque_to_tf_canonical( + self._variables) + for name, tensor in zip(self._param_names, weights + biases): + save_buffer[name] = array_ops.identity(tensor) + + def _checkpointable_restore(self, restore_buffer): + tensors = [ + array_ops.identity(restore_buffer[name]) for name in self._param_names + ] + return self.restore( + restored_tensors=tensors, + restored_shapes=None # Unused + ) + + def _add_checkpointable_dependencies(self, checkpointable, dtype): + """Add canonical weight dependencies to `checkpointable`. + + When saving or restoring, converts to or from the opaque buffer + format. Weights are saved and loaded in the configuration expected by + cuDNN-compatible cells. + + Args: + checkpointable: An object inheriting from `CheckpointableBase` to add + dependencies too (typically the cuDNN `Layer`). + dtype: The dtype for the canonical parameter Tensors. + """ + split_dependencies = split_dependency.split_dependency( + component_names=self._param_names, + component_dtypes=(dtype,) * len(self._param_names), + fill_save_buffer_fn=self._checkpointable_save, + consume_restore_buffer_fn=self._checkpointable_restore) + self._checkpointable_track_params(checkpointable, split_dependencies) + + def _checkpointable_track_params(self, checkpointable, params): + """Tracks parameters in a canonical configuration.""" + return # NotImplementedError raised by the Layer. + + def _tf_canonical_names(self): + tf_weights_names, tf_biases_names = [], [] + for i in range(self._num_layers): + if self._direction == CUDNN_RNN_UNIDIRECTION: + prefix = self._tf_canonical_name_prefix(i) + self._tf_canonical_names_single_layer(prefix, tf_weights_names, + tf_biases_names) + else: + fwd_prefix = self._tf_canonical_name_prefix(i, is_fwd=True) + bak_prefix = self._tf_canonical_name_prefix(i, is_fwd=False) + + self._tf_canonical_names_single_layer(fwd_prefix, tf_weights_names, + tf_biases_names) + self._tf_canonical_names_single_layer(bak_prefix, tf_weights_names, + tf_biases_names) + return tf_weights_names, tf_biases_names + + def _tf_canonical_name_prefix(self, layer, is_fwd=True): + if self._direction == CUDNN_RNN_UNIDIRECTION: + return "rnn/multi_rnn_cell/cell_%d/%s" % (layer, self._rnn_cell_name) + else: + if is_fwd: + return ("stack_bidirectional_rnn/cell_%d/bidirectional_rnn/fw/%s" % + (layer, self._rnn_cell_name)) + else: + return ("stack_bidirectional_rnn/cell_%d/bidirectional_rnn/bw/%s" % + (layer, self._rnn_cell_name)) + + def _tf_canonical_names_single_layer(self, prefix, tf_weights_names, + tf_biases_names): + raise NotImplementedError("Abstract method") + + +class CudnnLSTMSaveable(CudnnOpaqueParamsSaveable): + """SaveableObject implementation handling Cudnn LSTM opaque params.""" + + _format_converter_cls = CudnnParamsFormatConverterLSTM + _rnn_cell_name = base_layer.to_snake_case(CudnnCompatibleLSTMCell.__name__) + + def _tf_canonical_names_single_layer(self, prefix, tf_weights_names, + tf_bias_names): + tf_weights_names.append(prefix + "/kernel") + tf_bias_names.append(prefix + "/bias") + + def _checkpointable_track_params(self, checkpointable, params): + """Track parameters for compatibility with CudnnCompatibleLSTMCell.""" + biases = [] + weights = [] + for name in self._weight_names: + weights.append(params[name]) + for name in self._bias_names: + biases.append(params[name]) + assert len(params) == len(weights) + len(biases) + if len(weights) == 1 and len(biases) == 1: + # For single-layer cells, allow substituting a cell with no MultiRNNCell + # wrapping. + kernel, = weights # pylint: disable=unbalanced-tuple-unpacking + bias, = biases # pylint: disable=unbalanced-tuple-unpacking + checkpointable._track_checkpointable(kernel, name="kernel") # pylint: disable=protected-access + 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.AutoCheckpointable() + checkpointable._track_checkpointable(cell, name="cell-%d" % cell_index) # pylint: disable=protected-access + cell.bias = bias + cell.kernel = kernel + + +class CudnnGRUSaveable(CudnnOpaqueParamsSaveable): + """SaveableObject implementation handling Cudnn GRU opaque params.""" + + _format_converter_cls = CudnnParamsFormatConverterGRU + _rnn_cell_name = base_layer.to_snake_case(CudnnCompatibleGRUCell.__name__) + + def _tf_canonical_names_single_layer(self, prefix, tf_weights_names, + tf_bias_names): + tf_weights_names.append(prefix + "/gates/kernel") + tf_weights_names.append(prefix + "/candidate/input_projection/kernel") + tf_weights_names.append(prefix + "/candidate/hidden_projection/kernel") + + tf_bias_names.append(prefix + "/gates/bias") + tf_bias_names.append(prefix + "/candidate/input_projection/bias") + tf_bias_names.append(prefix + "/candidate/hidden_projection/bias") + + +class CudnnRNNTanhSaveable(CudnnLSTMSaveable): + _format_converter_cls = CudnnParamsFormatConverterTanh + _rnn_cell_name = base_layer.to_snake_case(rnn_cell_impl.BasicRNNCell.__name__) + + +class CudnnRNNReluSaveable(CudnnLSTMSaveable): + _format_converter_cls = CudnnParamsFormatConverterRelu + _rnn_cell_name = base_layer.to_snake_case(rnn_cell_impl.BasicRNNCell.__name__) + + _cudnn_rnn_common_doc_string = """ Cudnn RNN has an opaque parameter buffer that can be used for inference and training. But it is possible that the layout of the parameter buffers @@ -850,7 +942,7 @@ def _get_num_params(rnn_mode, num_layers, direction): elif rnn_mode == CUDNN_RNN_TANH: num_params_per_layer = CUDNN_RNN_TANH_PARAMS_PER_LAYER else: - raise ValueError("Invalid \'rnn_mode\': %s", rnn_mode) + raise ValueError("Invalid \'rnn_mode\': %s" % rnn_mode) num_params = num_layers * num_params_per_layer if direction != CUDNN_RNN_UNIDIRECTION: num_params *= 2 @@ -863,6 +955,7 @@ def _cudnn_rnn(inputs, params, is_training, rnn_mode, + sequence_lengths=None, input_mode=CUDNN_INPUT_LINEAR_MODE, direction=CUDNN_RNN_UNIDIRECTION, dropout=0., @@ -880,6 +973,10 @@ def _cudnn_rnn(inputs, params: the parameter buffer created for this model. is_training: whether this operation will be used in training or inference rnn_mode: one of ('lstm', 'gru', 'rnn_relu', 'rnn_tanh'). + sequence_lengths: an int32 array representing the variable sequence lengths + in a batch. The size of the array has to equal the batch_size. Default to + None, in which case sequences in the batch are assumed to have the same + length, which is inferred from inputs. input_mode: indicate whether there is a linear projection between the input and the actual computation before the first layer. It could be 'linear_input', 'skip_input' or 'auto_select'. @@ -918,7 +1015,10 @@ def _cudnn_rnn(inputs, "seed2": seed2, "name": name } - if use_cudnn_v2 is not "1": + if sequence_lengths is not None: + args["sequence_lengths"] = sequence_lengths + outputs, output_h, output_c, _, _ = gen_cudnn_rnn_ops.cudnn_rnnv3(**args) + elif use_cudnn_v2 != "1": outputs, output_h, output_c, _ = gen_cudnn_rnn_ops.cudnn_rnn(**args) else: outputs, output_h, output_c, _, _ = gen_cudnn_rnn_ops.cudnn_rnnv2(**args) @@ -930,6 +1030,7 @@ def cudnn_lstm(inputs, input_c, params, is_training, + sequence_lengths=None, input_mode=CUDNN_INPUT_LINEAR_MODE, direction=CUDNN_RNN_UNIDIRECTION, dropout=0., @@ -959,12 +1060,17 @@ def cudnn_lstm(inputs, dropout: whether to enable dropout. With it is 0, dropout is disabled. seed: the op seed used for initializing dropout. See `tf.set_random_seed` for behavior. + sequence_lengths: an int32 array representing the variable sequence lengths + in a batch. The size of the array has to equal the batch_size. Default to + None, in which case sequences in the batch are assumed to have the same + length, which is inferred from inputs. name: name of the operation. Returns: outputs, output_h, output_c """ return _cudnn_rnn(inputs, input_h, input_c, params, is_training, CUDNN_LSTM, - input_mode, direction, dropout, seed, name) + sequence_lengths, input_mode, direction, dropout, seed, + name) def _cudnn_rnn_no_input_c(inputs, @@ -972,6 +1078,7 @@ def _cudnn_rnn_no_input_c(inputs, params, is_training, rnn_mode, + sequence_lengths=None, input_mode=CUDNN_INPUT_LINEAR_MODE, direction=CUDNN_RNN_UNIDIRECTION, dropout=0., @@ -987,6 +1094,10 @@ def _cudnn_rnn_no_input_c(inputs, params: the parameter buffer created for this model. is_training: whether this operation will be used in training or inference rnn_mode: one of ('lstm', 'gru', 'rnn_relu', 'rnn_tanh'). + sequence_lengths: an int32 array representing the variable sequence lengths + in a batch. The size of the array has to equal the batch_size. Default to + None, in which case sequences in the batch are assumed to have the same + length, which is inferred from inputs. input_mode: indicate whether there is a linear projection between the input and the actual computation before the first layer. It could be 'linear_input', 'skip_input' or 'auto_select'. @@ -1006,8 +1117,8 @@ def _cudnn_rnn_no_input_c(inputs, """ input_c = array_ops.constant([], dtype=input_h.dtype) outputs, output_h, _ = _cudnn_rnn(inputs, input_h, input_c, params, - is_training, rnn_mode, input_mode, - direction, dropout, seed, name) + is_training, rnn_mode, sequence_lengths, + input_mode, direction, dropout, seed, name) return outputs, output_h @@ -1015,6 +1126,7 @@ def cudnn_gru(inputs, input_h, params, is_training, + sequence_lengths=None, input_mode=CUDNN_INPUT_LINEAR_MODE, direction=CUDNN_RNN_UNIDIRECTION, dropout=0., @@ -1037,6 +1149,10 @@ def cudnn_gru(inputs, 'skip_input' is only allowed when input_size == num_units; 'auto_select' implies 'skip_input' when input_size == num_units; otherwise, it implies 'linear_input'. + sequence_lengths: an int32 array representing the variable sequence lengths + in a batch. The size of the array has to equal the batch_size. Default to + None, in which case sequences in the batch are assumed to have the same + length, which is inferred from inputs. direction: the direction model that the model operates. Could be either 'unidirectional' or 'bidirectional' dropout: whether to enable dropout. With it is 0, dropout is disabled. @@ -1047,7 +1163,8 @@ def cudnn_gru(inputs, outputs, output_h """ return _cudnn_rnn_no_input_c(inputs, input_h, params, is_training, CUDNN_GRU, - input_mode, direction, dropout, seed, name) + sequence_lengths, input_mode, direction, dropout, + seed, name) def cudnn_rnn_relu(inputs, @@ -1058,6 +1175,7 @@ def cudnn_rnn_relu(inputs, direction=CUDNN_RNN_UNIDIRECTION, dropout=0., seed=0, + sequence_lengths=None, name=None): """Cudnn RNN Relu. @@ -1070,30 +1188,34 @@ def cudnn_rnn_relu(inputs, is_training: whether this operation will be used in training or inference input_mode: indicate whether there is a linear projection between the input and the actual computation before the first layer. It could be - 'linear_input', 'skip_input' or 'auto_select'. - 'linear_input' (default) always applies a linear projection of input - onto RNN hidden state. (standard RNN behavior). - 'skip_input' is only allowed when input_size == num_units; - 'auto_select' implies 'skip_input' when input_size == num_units; - otherwise, it implies 'linear_input'. + 'linear_input', 'skip_input' or 'auto_select'. 'linear_input' (default) + always applies a linear projection of input onto RNN hidden state. + (standard RNN behavior). 'skip_input' is only allowed when input_size == + num_units; 'auto_select' implies 'skip_input' when input_size == + num_units; otherwise, it implies 'linear_input'. direction: the direction model that the model operates. Could be either - 'unidirectional' or 'bidirectional' + 'unidirectional' or 'bidirectional' dropout: whether to enable dropout. With it is 0, dropout is disabled. seed: the op seed used for initializing dropout. See `tf.set_random_seed` - for behavior. + for behavior. + sequence_lengths: an int32 array representing the variable sequence lengths + in a batch. The size of the array has to equal the batch_size. If not + provided, the same sequence length will be assumed. name: name of the operation. + Returns: outputs, output_h """ return _cudnn_rnn_no_input_c(inputs, input_h, params, is_training, - CUDNN_RNN_RELU, input_mode, direction, dropout, - seed, name) + CUDNN_RNN_RELU, sequence_lengths, input_mode, + direction, dropout, seed, name) def cudnn_rnn_tanh(inputs, input_h, params, is_training, + sequence_lengths=None, input_mode=CUDNN_INPUT_LINEAR_MODE, direction=CUDNN_RNN_UNIDIRECTION, dropout=0., @@ -1116,6 +1238,10 @@ def cudnn_rnn_tanh(inputs, 'skip_input' is only allowed when input_size == num_units; 'auto_select' implies 'skip_input' when input_size == num_units; otherwise, it implies 'linear_input'. + sequence_lengths: an int32 array representing the variable sequence lengths + in a batch. The size of the array has to equal the batch_size. Default to + None, in which case sequences in the batch are assumed to have the same + length, which is inferred from inputs. direction: the direction model that the model operates. Could be either 'unidirectional' or 'bidirectional' dropout: whether to enable dropout. With it is 0, dropout is disabled. @@ -1126,8 +1252,8 @@ def cudnn_rnn_tanh(inputs, outputs, output_h """ return _cudnn_rnn_no_input_c(inputs, input_h, params, is_training, - CUDNN_RNN_TANH, input_mode, direction, dropout, - seed, name) + CUDNN_RNN_TANH, sequence_lengths, input_mode, + direction, dropout, seed, name) def cudnn_rnn_opaque_params_to_canonical(rnn_mode, @@ -1405,7 +1531,13 @@ class _CudnnRNN(object): input_mode=self._input_mode, direction=self._direction) - def __call__(self, input_data, input_h, input_c, params, is_training=True): + def __call__(self, + input_data, + input_h, + input_c, + params, + is_training=True, + sequence_lengths=None): """Runs the forward step for the RNN model. Args: @@ -1417,6 +1549,10 @@ class _CudnnRNN(object): A Tensor of the same shape as input_h. params: the parameter buffer created for this model. is_training: whether this operation will be used in training or inference. + sequence_lengths: an int32 array representing the variable sequence + lengths in a batch. The size of the array has to equal the batch_size. + Default to None, in which case sequences in the batch are assumed to + have the same length, which is inferred from inputs. Returns: output: the output sequence. output_h: the final state for h. @@ -1429,6 +1565,7 @@ class _CudnnRNN(object): params, is_training, self._rnn_mode, + sequence_lengths=sequence_lengths, input_mode=self._input_mode, direction=self._direction, dropout=self._dropout, @@ -1523,7 +1660,13 @@ class CudnnLSTM(_CudnnRNN): dropout=dropout, seed=seed) - def __call__(self, input_data, input_h, input_c, params, is_training=True): + def __call__(self, + input_data, + input_h, + input_c, + params, + sequence_lengths=None, + is_training=True): """Runs the forward step for the Cudnn LSTM model. Args: @@ -1534,6 +1677,10 @@ class CudnnLSTM(_CudnnRNN): input_c: the initial hidden state for c. A Tensor of the same shape as input_h. params: the parameter buffer created for this model. + sequence_lengths: an int32 array representing the variable sequence + lengths in a batch. The size of the array has to equal the batch_size. + Default to None, in which case sequences in the batch are assumed to + have the same length, which is inferred from inputs. is_training: whether this operation will be used in training or inference. Returns: output: the output sequence. @@ -1541,7 +1688,12 @@ class CudnnLSTM(_CudnnRNN): output_c: the final state for c. """ output, output_h, output_c = super(CudnnLSTM, self).__call__( - input_data, input_h, input_c, params, is_training=is_training) + input_data, + input_h, + input_c, + params, + sequence_lengths=sequence_lengths, + is_training=is_training) return (output, output_h, output_c) @@ -1582,7 +1734,7 @@ class _CudnnRNNNoInputC(_CudnnRNN): """ if direction not in (CUDNN_RNN_UNIDIRECTION, CUDNN_RNN_BIDIRECTION): - raise ValueError("Invalid direction: %s", direction) + raise ValueError("Invalid direction: %s" % direction) super(_CudnnRNNNoInputC, self).__init__( self._rnn_mode, @@ -1595,7 +1747,12 @@ class _CudnnRNNNoInputC(_CudnnRNN): dropout=dropout, seed=seed) - def __call__(self, input_data, input_h, params, is_training=True): + def __call__(self, + input_data, + input_h, + params, + sequence_lengths=None, + is_training=True): """Runs the forward step for the Cudnn LSTM model. Args: @@ -1604,6 +1761,10 @@ class _CudnnRNNNoInputC(_CudnnRNN): input_h: the initial hidden state for h. A Tensor of shape [num_layers, batch_size, num_units]. params: the parameter buffer created for this model. + sequence_lengths: an int32 array representing the variable sequence + lengths in a batch. The size of the array has to equal the batch_size. + Default to None, in which case sequences in the batch are assumed to + have the same length, which is inferred from inputs. is_training: whether this operation will be used in training or inference. Returns: output: the output sequence. @@ -1615,6 +1776,7 @@ class _CudnnRNNNoInputC(_CudnnRNN): params, is_training, self._rnn_mode, + sequence_lengths=sequence_lengths, input_mode=self._input_mode, direction=self._direction, dropout=self._dropout, 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 0456463a1928cf226010670b90a5d574579e0411..6c5f8c6b00975b3fba041271309a93cecd9f5057 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 @@ -46,7 +46,7 @@ class AssertElementShapeTest(test_base.DatasetTestBase): result = dataset.apply(batching.assert_element_shape(expected_shapes)) self.assertEqual(expected_shapes, result.output_shapes) - iterator = result.make_initializable_iterator() + iterator = dataset_ops.make_initializable_iterator(result) init_op = iterator.initializer get_next = iterator.get_next() with self.cached_session() as sess: @@ -88,7 +88,7 @@ class AssertElementShapeTest(test_base.DatasetTestBase): result = dataset.apply(batching.assert_element_shape(expected_shapes)) self.assertEqual(expected_shapes, result.output_shapes) - iterator = result.make_initializable_iterator() + iterator = dataset_ops.make_initializable_iterator(result) init_op = iterator.initializer get_next = iterator.get_next() with self.cached_session() as sess: @@ -115,9 +115,8 @@ class AssertElementShapeTest(test_base.DatasetTestBase): wrong_shapes = (tensor_shape.TensorShape(2), tensor_shape.TensorShape((3, 10))) - iterator = ( - dataset.apply(batching.assert_element_shape(wrong_shapes)) - .make_initializable_iterator()) + iterator = dataset_ops.make_initializable_iterator( + dataset.apply(batching.assert_element_shape(wrong_shapes))) init_op = iterator.initializer get_next = iterator.get_next() with self.cached_session() as sess: @@ -142,7 +141,7 @@ class AssertElementShapeTest(test_base.DatasetTestBase): tensor_shape.TensorShape((3, 4))) self.assertEqual(actual_shapes, result.output_shapes) - iterator = result.make_initializable_iterator() + iterator = dataset_ops.make_initializable_iterator(result) init_op = iterator.initializer get_next = iterator.get_next() with self.cached_session() as sess: @@ -184,7 +183,7 @@ class AssertElementShapeTest(test_base.DatasetTestBase): result = dataset.apply(batching.assert_element_shape(expected_shapes)) self.assertEqual(expected_shapes, result.output_shapes) - iterator = result.make_initializable_iterator() + iterator = dataset_ops.make_initializable_iterator(result) init_op = iterator.initializer get_next = iterator.get_next() with self.cached_session() as sess: @@ -211,9 +210,8 @@ class AssertElementShapeTest(test_base.DatasetTestBase): wrong_shapes = (tensor_shape.TensorShape(2), tensor_shape.TensorShape((None, 10))) - iterator = ( - dataset.apply(batching.assert_element_shape(wrong_shapes)) - .make_initializable_iterator()) + iterator = dataset_ops.make_initializable_iterator( + dataset.apply(batching.assert_element_shape(wrong_shapes))) init_op = iterator.initializer get_next = iterator.get_next() with self.cached_session() as sess: 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 d2a72272db159755ac2d741bcdbce9ec646d928e..b9840b1ff1a3df5a05db0e64f436637220f49f80 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 @@ -23,6 +23,7 @@ import shutil from tensorflow.contrib.data.python.ops import readers from tensorflow.python.data.kernel_tests import test_base +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 @@ -48,7 +49,7 @@ class LMDBDatasetTest(test_base.DatasetTestBase): num_repeats = 2 dataset = readers.LMDBDataset(filenames).repeat(num_repeats) - iterator = dataset.make_initializable_iterator() + iterator = dataset_ops.make_initializable_iterator(dataset) init_op = iterator.initializer get_next = iterator.get_next() 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 c5a786232252432481566e3cde23e9310df172cc..2527706709fae8e459aca3489324d4db3c784be6 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 @@ -63,13 +63,13 @@ class SlideDatasetTest(test_base.DatasetTestBase, parameterized.TestCase): # The pipeline is TensorSliceDataset -> MapDataset(square_3) -> # RepeatDataset(count) -> # _SlideDataset(window_size, window_shift, window_stride). - iterator = ( + iterator = dataset_ops.make_initializable_iterator( dataset_ops.Dataset.from_tensor_slices(components).map(_map_fn) .repeat(count).apply( sliding.sliding_window_batch( window_size=window_size_t, window_shift=window_shift_t, - window_stride=window_stride_t)).make_initializable_iterator()) + window_stride=window_stride_t))) init_op = iterator.initializer get_next = iterator.get_next() @@ -127,13 +127,13 @@ class SlideDatasetTest(test_base.DatasetTestBase, parameterized.TestCase): # The pipeline is TensorSliceDataset -> MapDataset(square_3) -> # RepeatDataset(count) -> _SlideDataset(window_size, stride, window_stride). - iterator = ( + iterator = dataset_ops.make_initializable_iterator( dataset_ops.Dataset.from_tensor_slices(components).map(_map_fn) .repeat(count).apply( sliding.sliding_window_batch( window_size=window_size_t, stride=stride_t, - window_stride=window_stride_t)).make_initializable_iterator()) + window_stride=window_stride_t))) init_op = iterator.initializer get_next = iterator.get_next() @@ -173,12 +173,12 @@ class SlideDatasetTest(test_base.DatasetTestBase, parameterized.TestCase): window_shift_t = array_ops.placeholder(dtypes.int64, shape=[]) window_stride_t = array_ops.placeholder(dtypes.int64, shape=[]) - iterator = ( + iterator = dataset_ops.make_initializable_iterator( dataset_ops.Dataset.range(10).map(lambda x: x).repeat(count_t).apply( sliding.sliding_window_batch( window_size=window_size_t, window_shift=window_shift_t, - window_stride=window_stride_t)).make_initializable_iterator()) + window_stride=window_stride_t))) init_op = iterator.initializer with self.cached_session() as sess: @@ -204,9 +204,9 @@ class SlideDatasetTest(test_base.DatasetTestBase, parameterized.TestCase): return sparse_tensor.SparseTensorValue( indices=[[0]], values=(i * [1]), dense_shape=[1]) - iterator = dataset_ops.Dataset.range(10).map(_sparse).apply( - sliding.sliding_window_batch( - window_size=5, window_shift=3)).make_initializable_iterator() + iterator = dataset_ops.make_initializable_iterator( + dataset_ops.Dataset.range(10).map(_sparse).apply( + sliding.sliding_window_batch(window_size=5, window_shift=3))) init_op = iterator.initializer get_next = iterator.get_next() @@ -233,9 +233,9 @@ class SlideDatasetTest(test_base.DatasetTestBase, parameterized.TestCase): values=array_ops.fill([math_ops.to_int32(i)], i), dense_shape=[i]) - iterator = dataset_ops.Dataset.range(10).map(_sparse).apply( - sliding.sliding_window_batch( - window_size=5, window_shift=3)).make_initializable_iterator() + iterator = dataset_ops.make_initializable_iterator( + dataset_ops.Dataset.range(10).map(_sparse).apply( + sliding.sliding_window_batch(window_size=5, window_shift=3))) init_op = iterator.initializer get_next = iterator.get_next() @@ -265,11 +265,10 @@ class SlideDatasetTest(test_base.DatasetTestBase, parameterized.TestCase): return sparse_tensor.SparseTensorValue( indices=[[0]], values=(i * [1]), dense_shape=[1]) - iterator = ( + iterator = dataset_ops.make_initializable_iterator( dataset_ops.Dataset.range(10).map(_sparse).apply( sliding.sliding_window_batch(window_size=4, window_shift=2)).apply( - sliding.sliding_window_batch(window_size=3, window_shift=1)) - .make_initializable_iterator()) + sliding.sliding_window_batch(window_size=3, window_shift=1))) init_op = iterator.initializer get_next = iterator.get_next() @@ -305,11 +304,10 @@ class SlideDatasetTest(test_base.DatasetTestBase, parameterized.TestCase): yield [4.0, 5.0, 6.0] yield [7.0, 8.0, 9.0, 10.0] - iterator = ( + iterator = dataset_ops.make_initializable_iterator( dataset_ops.Dataset.from_generator( generator, dtypes.float32, output_shapes=[None]).apply( - sliding.sliding_window_batch(window_size=3, window_shift=1)) - .make_initializable_iterator()) + sliding.sliding_window_batch(window_size=3, window_shift=1))) next_element = iterator.get_next() with self.cached_session() as sess: diff --git a/tensorflow/contrib/data/python/ops/BUILD b/tensorflow/contrib/data/python/ops/BUILD index 34dc2379d0cb38f8f6962fa42efe21b793bc8d65..0fb406f1167053a128646c5c692986b0ce016f1e 100644 --- a/tensorflow/contrib/data/python/ops/BUILD +++ b/tensorflow/contrib/data/python/ops/BUILD @@ -188,8 +188,7 @@ py_library( "//tensorflow/python:framework_ops", "//tensorflow/python:function", "//tensorflow/python/data/ops:dataset_ops", - "//tensorflow/python/data/util:nest", - "//tensorflow/python/data/util:sparse", + "//tensorflow/python/data/util:structure", ], ) diff --git a/tensorflow/contrib/data/python/ops/readers.py b/tensorflow/contrib/data/python/ops/readers.py index 4601376dff47e161962e92678883039c4b88bab7..c6bf5215c9406d03d2704e46903b3aa57e7e68d9 100644 --- a/tensorflow/contrib/data/python/ops/readers.py +++ b/tensorflow/contrib/data/python/ops/readers.py @@ -21,10 +21,9 @@ from tensorflow.python.data.experimental.ops import optimization from tensorflow.python.data.experimental.ops import readers from tensorflow.python.data.ops import dataset_ops from tensorflow.python.data.ops import readers as core_readers -from tensorflow.python.data.util import nest +from tensorflow.python.data.util import structure from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops -from tensorflow.python.framework import tensor_shape from tensorflow.python.ops import gen_experimental_dataset_ops from tensorflow.python.util import deprecation @@ -355,7 +354,7 @@ def read_batch_features(file_pattern, shuffle=randomize_input, num_epochs=num_epochs, shuffle_buffer_size=capacity) - iterator = dataset.make_one_shot_iterator() + iterator = dataset_ops.make_one_shot_iterator(dataset) outputs = iterator.get_next() return outputs @@ -379,37 +378,25 @@ class LMDBDataset(dataset_ops.DatasetSource): (key value) pairs sequentially. For example: ```python + tf.enable_eager_execution() + dataset = tf.contrib.lmdb.LMDBDataset("/foo/bar.mdb") - iterator = dataset.make_one_shot_iterator() - next_element = iterator.get_next() + # Prints the (key, value) pairs inside a lmdb file. - while True: - try: - print(sess.run(next_element)) - except tf.errors.OutOfRangeError: - break + for key, value in dataset: + print(key, value) ``` Args: filenames: A `tf.string` tensor containing one or more filenames. """ - super(LMDBDataset, self).__init__() self._filenames = ops.convert_to_tensor( filenames, dtype=dtypes.string, name="filenames") - - def _as_variant_tensor(self): - return gen_experimental_dataset_ops.experimental_lmdb_dataset( - self._filenames, - output_types=nest.flatten(self.output_types), - output_shapes=nest.flatten(self.output_shapes)) - - @property - def output_classes(self): - return ops.Tensor, ops.Tensor - - @property - def output_shapes(self): - return (tensor_shape.TensorShape([]), tensor_shape.TensorShape([])) + variant_tensor = gen_experimental_dataset_ops.experimental_lmdb_dataset( + self._filenames, **dataset_ops.flat_structure(self)) + super(LMDBDataset, self).__init__(variant_tensor) @property - def output_types(self): - return dtypes.string, dtypes.string + def _element_structure(self): + return structure.NestedStructure( + (structure.TensorStructure(dtypes.string, []), + structure.TensorStructure(dtypes.string, []))) diff --git a/tensorflow/contrib/data/python/ops/sliding.py b/tensorflow/contrib/data/python/ops/sliding.py index bcc383587c54bd89502313f9328bc06c49046a87..6708e01d08135a132b797e317cd2a241c3428f40 100644 --- a/tensorflow/contrib/data/python/ops/sliding.py +++ b/tensorflow/contrib/data/python/ops/sliding.py @@ -18,11 +18,10 @@ from __future__ import division from __future__ import print_function from tensorflow.python.data.ops import dataset_ops -from tensorflow.python.data.util import nest +from tensorflow.python.data.util import structure from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops -from tensorflow.python.framework import tensor_shape -from tensorflow.python.ops import gen_dataset_ops +from tensorflow.python.ops import gen_experimental_dataset_ops as ged_ops from tensorflow.python.util import deprecation @@ -31,7 +30,6 @@ class _SlideDataset(dataset_ops.UnaryDataset): def __init__(self, input_dataset, window_size, window_shift, window_stride): """See `sliding_window_batch` for details.""" - super(_SlideDataset, self).__init__(input_dataset) self._input_dataset = input_dataset self._window_size = ops.convert_to_tensor( window_size, dtype=dtypes.int64, name="window_stride") @@ -40,29 +38,21 @@ class _SlideDataset(dataset_ops.UnaryDataset): self._window_shift = ops.convert_to_tensor( window_shift, dtype=dtypes.int64, name="window_shift") - def _as_variant_tensor(self): - return gen_dataset_ops.slide_dataset( - self._input_dataset._as_variant_tensor(), # pylint: disable=protected-access + input_structure = structure.convert_legacy_structure( + input_dataset.output_types, input_dataset.output_shapes, + input_dataset.output_classes) + self._structure = input_structure._batch(None) # pylint: disable=protected-access + variant_tensor = ged_ops.experimental_sliding_window_dataset( + self._input_dataset._variant_tensor, # pylint: disable=protected-access window_size=self._window_size, window_shift=self._window_shift, window_stride=self._window_stride, **dataset_ops.flat_structure(self)) + super(_SlideDataset, self).__init__(input_dataset, variant_tensor) @property - def output_classes(self): - return self._input_dataset.output_classes - - @property - def output_shapes(self): - input_shapes = self._input_dataset.output_shapes - return nest.pack_sequence_as(input_shapes, [ - tensor_shape.vector(None).concatenate(s) - for s in nest.flatten(self._input_dataset.output_shapes) - ]) - - @property - def output_types(self): - return self._input_dataset.output_types + def _element_structure(self): + return self._structure @deprecation.deprecated_args( diff --git a/tensorflow/contrib/distribute/BUILD b/tensorflow/contrib/distribute/BUILD index a87a5624c88d1d0af10055261dad55937ed6aeb0..3ecd755d86f6be47910aebbdb46d335d165427d8 100644 --- a/tensorflow/contrib/distribute/BUILD +++ b/tensorflow/contrib/distribute/BUILD @@ -26,7 +26,6 @@ py_library( visibility = ["//tensorflow:internal"], deps = [ "//tensorflow/contrib/distribute/python:collective_all_reduce_strategy", - "//tensorflow/contrib/distribute/python:cross_tower_ops", "//tensorflow/contrib/distribute/python:mirrored_strategy", "//tensorflow/contrib/distribute/python:monitor", "//tensorflow/contrib/distribute/python:one_device_strategy", @@ -35,6 +34,7 @@ py_library( "//tensorflow/contrib/distribute/python:tpu_strategy", "//tensorflow/python:training", "//tensorflow/python:util", + "//tensorflow/python/distribute:cross_device_ops", "//tensorflow/python/distribute:distribute_config", "//tensorflow/python/distribute:distribute_coordinator", ], diff --git a/tensorflow/contrib/distribute/README.md b/tensorflow/contrib/distribute/README.md index f82453f3b5ea01b8bb64a70bd49f5e3e831bb4e2..dbcaf8185fb7a9d2bcf22376439c0ebd49accb1a 100644 --- a/tensorflow/contrib/distribute/README.md +++ b/tensorflow/contrib/distribute/README.md @@ -43,25 +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: - -```python -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 +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 + 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. @@ -69,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 @@ -90,8 +87,8 @@ Similarly, we can also call `evaluate` and `predict` as before using appropriate datasets. ```python -model.evaluate(eval_dataset) -model.predict(predict_dataset) +model.evaluate(eval_dataset, steps=1) +model.predict(predict_dataset, steps=1) ``` That's all you need to train your model with Keras on multiple GPUs with @@ -131,7 +128,7 @@ def model_fn(features, labels, mode): return tf.estimator.EstimatorSpec(mode, loss=loss) if mode == tf.estimator.ModeKeys.TRAIN: - train_op = tf.train.GradientDescentOptimizer(0.2).minimize(loss_fn()) + train_op = tf.train.GradientDescentOptimizer(0.2).minimize(loss) return tf.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op) ``` @@ -245,19 +242,17 @@ Let's use the same example for multi-worker. We'll start a cluster with 3 workers doing synchronous all-reduce training. In the following code snippet, we start multi-worker training using `tf.estimator.train_and_evaluate`: - ```python def model_main(): - estimator = ... distribution = tf.contrib.distribute.CollectiveAllReduceStrategy( num_gpus_per_worker=2) config = tf.estimator.RunConfig(train_distribute=distribution) + estimator = tf.estimator.Estimator(model_fn=model_fn, config=config) train_spec = tf.estimator.TrainSpec(input_fn=input_fn) eval_spec = tf.estimator.EvalSpec(input_fn=eval_input_fn) tf.estimator.train_and_evaluate(estimator, train_spec, eval_spec) ``` - **Note**: You don't have to set "TF\_CONFIG" manually if you use our provided Kubernetes template. @@ -324,13 +319,13 @@ start training. On your laptop, you can run ```python -estimator = ... distribution = tf.contrib.distribute.CollectiveAllReduceStrategy( num_gpus_per_worker=2) config = tf.estimator.RunConfig( experimental_distribute=tf.contrib.distribute.DistributeConfig( train_distribute=distribution, remote_cluster={"worker": ["host1:port", "host2:port", "host3:port"]})) +estimator = tf.estimator.Estimator(model_fn=model_fn, config=config) train_spec = tf.estimator.TrainSpec(input_fn=input_fn) eval_spec = tf.estimator.EvalSpec(input_fn=eval_input_fn) tf.estimator.train_and_evaluate(estimator, train_spec, eval_spec) diff --git a/tensorflow/contrib/distribute/__init__.py b/tensorflow/contrib/distribute/__init__.py index d98fc30670052acf350ca9a5796781045b1f0580..59d76f5d1c817d7f2cc8ad285b9fb517fe994a81 100644 --- a/tensorflow/contrib/distribute/__init__.py +++ b/tensorflow/contrib/distribute/__init__.py @@ -25,44 +25,50 @@ from __future__ import print_function # pylint: disable=unused-import,wildcard-import from tensorflow.contrib.distribute.python.collective_all_reduce_strategy import CollectiveAllReduceStrategy -from tensorflow.contrib.distribute.python.cross_tower_ops import * from tensorflow.contrib.distribute.python.mirrored_strategy import MirroredStrategy from tensorflow.contrib.distribute.python.monitor import Monitor from tensorflow.contrib.distribute.python.one_device_strategy import OneDeviceStrategy from tensorflow.contrib.distribute.python.parameter_server_strategy import ParameterServerStrategy from tensorflow.contrib.distribute.python.step_fn import * +from tensorflow.contrib.distribute.python.tpu_strategy import initialize_tpu_system 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 _allowed_symbols = [ - 'AllReduceCrossTowerOps', + 'AllReduceCrossDeviceOps', 'CollectiveAllReduceStrategy', - 'CrossTowerOps', + 'CrossDeviceOps', 'DistributeConfig', 'DistributionStrategy', + 'DistributionStrategyExtended', 'MirroredStrategy', 'Monitor', 'MultiWorkerAllReduce', 'OneDeviceStrategy', 'ParameterServerStrategy', - 'ReductionToOneDeviceCrossTowerOps', + 'ReductionToOneDeviceCrossDeviceOps', 'Step', 'StandardInputStep', 'StandardSingleLossStep', - 'TowerContext', + 'ReplicaContext', 'TPUStrategy', - 'get_cross_tower_context', + 'initialize_tpu_system', + 'get_cross_replica_context', 'get_distribution_strategy', 'get_loss_reduction', - 'get_tower_context', + 'get_replica_context', + 'get_strategy', 'has_distribution_strategy', - 'require_tower_context', + 'has_strategy', + 'in_cross_replica_context', + 'require_replica_context', 'run_standard_tensorflow_server', 'UpdateContext', ] diff --git a/tensorflow/contrib/distribute/python/BUILD b/tensorflow/contrib/distribute/python/BUILD index dc2964568b529e4bf73dc648493d03db3cc80c1a..fb8732d59169ee9c809519dd413a1173d1187e2b 100644 --- a/tensorflow/contrib/distribute/python/BUILD +++ b/tensorflow/contrib/distribute/python/BUILD @@ -1,5 +1,10 @@ # 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") + package( default_visibility = [ "//tensorflow:internal", @@ -10,51 +15,26 @@ licenses(["notice"]) # Apache 2.0 exports_files(["LICENSE"]) -load("//tensorflow:tensorflow.bzl", "py_test") -load("//tensorflow:tensorflow.bzl", "cuda_py_test") - # TODO(priyag): Figure out testonly issues that are preventing us from # including our tests in pip for now. -py_library( - name = "values", - srcs = ["values.py"], - visibility = ["//tensorflow:internal"], - deps = [ - ":input_ops", - "//tensorflow/python:array_ops", - "//tensorflow/python:control_flow_ops", - "//tensorflow/python:device_util", - "//tensorflow/python:distribute", - "//tensorflow/python:framework_ops", - "//tensorflow/python:resource_variable_ops", - "//tensorflow/python:training", - "//tensorflow/python:util", - "//tensorflow/python/data/ops:multi_device_iterator_ops", - "//tensorflow/python/eager:context", - "//tensorflow/python/training/checkpointable:base", - "@six_archive//:six", - ], -) - cuda_py_test( name = "values_test", srcs = ["values_test.py"], additional_deps = [ + ":combinations", ":mirrored_strategy", - ":multi_worker_test_base", - ":values", + "@absl_py//absl/testing:parameterized", "//tensorflow/core:protos_all_py", - "//tensorflow/python/data/ops:dataset_ops", - "//tensorflow/python:errors", "//tensorflow/python:array_ops", "//tensorflow/python:constant_op", "//tensorflow/python:framework_ops", "//tensorflow/python:framework_test_lib", "//tensorflow/python:training", "//tensorflow/python:variable_scope", + "//tensorflow/python/distribute:device_util", + "//tensorflow/python/distribute:values", "//tensorflow/python/eager:context", - "//tensorflow/python:device_util", "//tensorflow/python/eager:test", "//tensorflow/python/estimator:estimator_py", ], @@ -63,30 +43,34 @@ cuda_py_test( ], ) +cuda_py_test( + name = "input_lib_test", + srcs = ["input_lib_test.py"], + additional_deps = [ + ":combinations", + ":mirrored_strategy", + ":multi_worker_test_base", + "@absl_py//absl/testing:parameterized", + "//tensorflow/python:errors", + "//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( name = "mirrored_strategy", srcs = ["mirrored_strategy.py"], visibility = ["//tensorflow:internal"], deps = [ - ":cross_tower_ops", - ":shared_variable_creator", - ":values", - "//tensorflow/core:protos_all_py", - "//tensorflow/python:array_ops", - "//tensorflow/python:constant_op", - "//tensorflow/python:control_flow_ops", - "//tensorflow/python:device", - "//tensorflow/python:device_util", - "//tensorflow/python:distribute", - "//tensorflow/python:framework_ops", - "//tensorflow/python:pywrap_tensorflow", - "//tensorflow/python:training", - "//tensorflow/python:util", - "//tensorflow/python:variable_scope", - "//tensorflow/python:variables", - "//tensorflow/python/distribute:multi_worker_util", - "//tensorflow/python/eager:context", - "//tensorflow/python/eager:tape", + "//tensorflow/python/distribute:distribute_lib", + "//tensorflow/python/distribute:input_lib", + "//tensorflow/python/distribute:mirrored_strategy", ], ) @@ -95,17 +79,10 @@ py_library( srcs = ["parameter_server_strategy.py"], visibility = ["//tensorflow:internal"], deps = [ - ":cross_tower_ops", - ":mirrored_strategy", - ":values", - "//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:multi_worker_util", - "//tensorflow/python/eager:context", + "//tensorflow/python/distribute:distribute_lib", + "//tensorflow/python/distribute:parameter_server_strategy", + "//tensorflow/python/distribute:values", + "//tensorflow/python/distribute/cluster_resolver:cluster_resolver_lib", ], ) @@ -116,7 +93,7 @@ cuda_py_test( ":combinations", ":multi_worker_test_base", ":parameter_server_strategy", - ":values", + ":strategy_test_lib", "@absl_py//absl/testing:parameterized", "//tensorflow/core:protos_all_py", "//tensorflow/python:array_ops", @@ -127,10 +104,12 @@ cuda_py_test( "//tensorflow/python:gradients", "//tensorflow/python:layers", "//tensorflow/python:session", + "//tensorflow/python:tensor_util", "//tensorflow/python:training", "//tensorflow/python:variable_scope", "//tensorflow/python:variables", "//tensorflow/python/distribute:multi_worker_util", + "//tensorflow/python/distribute:values", "//tensorflow/python/eager:context", "//tensorflow/python/estimator:estimator_py", ], @@ -145,12 +124,15 @@ py_library( srcs = ["one_device_strategy.py"], visibility = ["//tensorflow:internal"], deps = [ - ":values", - "//tensorflow/contrib/eager/python:datasets", "//tensorflow/python:array_ops", - "//tensorflow/python:distribute", + "//tensorflow/python:dtypes", "//tensorflow/python:framework_ops", "//tensorflow/python:math_ops", + "//tensorflow/python/distribute:distribute_lib", + "//tensorflow/python/distribute:input_lib", + "//tensorflow/python/distribute:numpy_dataset", + "//tensorflow/python/distribute:reduce_util", + "//tensorflow/python/distribute:values", "//tensorflow/python/eager:context", "@six_archive//:six", ], @@ -161,16 +143,18 @@ py_library( srcs = ["collective_all_reduce_strategy.py"], visibility = ["//tensorflow:internal"], deps = [ - ":cross_tower_ops", - ":cross_tower_utils", ":mirrored_strategy", - ":values", "//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:numpy_dataset", + "//tensorflow/python/distribute:values", "//tensorflow/python/eager:context", ], ) @@ -187,14 +171,15 @@ py_library( "//tensorflow/core:protos_all_py", "//tensorflow/python:array_ops", "//tensorflow/python:constant_op", - "//tensorflow/python:distribute", "//tensorflow/python:framework_ops", "//tensorflow/python:layers", "//tensorflow/python:training", "//tensorflow/python:variables", + "//tensorflow/python/distribute:distribute_lib", "//tensorflow/python/eager:backprop", "//tensorflow/python/eager:context", "//tensorflow/python/eager:test", + "//third_party/py/numpy", ], ) @@ -212,10 +197,10 @@ py_library( ":tpu_strategy", "//tensorflow/contrib/cluster_resolver:cluster_resolver_pip", "//tensorflow/contrib/optimizer_v2:training", - "//tensorflow/python:distribute", "//tensorflow/python:framework_ops", "//tensorflow/python:training", "//tensorflow/python:util", + "//tensorflow/python/distribute:distribute_lib", "//tensorflow/python/eager:context", "@absl_py//absl/testing:parameterized", ], @@ -233,28 +218,6 @@ py_test( ], ) -py_test( - name = "mirrored_strategy_test", - srcs = ["mirrored_strategy_test.py"], - srcs_version = "PY2AND3", - tags = [ - "no_pip", - ], - deps = [ - ":mirrored_strategy", - ":multi_worker_test_base", - ":strategy_test_lib", - "//tensorflow/python:constant_op", - "//tensorflow/python:distribute", - "//tensorflow/python:framework_ops", - "//tensorflow/python:framework_test_lib", - "//tensorflow/python:training", - "//tensorflow/python:variable_scope", - "//tensorflow/python/eager:context", - "//tensorflow/python/eager:test", - ], -) - py_test( name = "one_device_strategy_test", srcs = ["one_device_strategy_test.py"], @@ -270,35 +233,32 @@ py_test( ], ) +# TODO(priyag): Rename this test to mirrored_strategy_test cuda_py_test( name = "mirrored_strategy_multigpu_test", srcs = ["mirrored_strategy_multigpu_test.py"], additional_deps = [ + ":combinations", ":mirrored_strategy", ":multi_worker_test_base", - ":values", ":strategy_test_lib", - "//tensorflow/python:distribute", "//tensorflow/core:protos_all_py", "//tensorflow/python:array_ops", "//tensorflow/python:constant_op", + "//tensorflow/python:framework_test_lib", "//tensorflow/python:layers", "//tensorflow/python:state_ops", "//tensorflow/python:variable_scope", - "//tensorflow/python:framework_test_lib", + "//tensorflow/python/distribute:distribute_lib", + "//tensorflow/python/distribute:values", "//tensorflow/python/eager:context", "//tensorflow/python/eager:test", ], + shard_count = 5, tags = [ "guitar", - "no_pip", "multi_and_single_gpu", - # Do not perform the extra analysis on this test, because it is already - # performed for the `:mirrored_strategy_test` target. - "no_oss", - "noasan", - "notap", - "notsan", + "no_pip", ], ) @@ -315,6 +275,7 @@ py_library( "//tensorflow/python:client_testlib", "//tensorflow/python:distributed_framework_test_lib", "//tensorflow/python:session", + "//tensorflow/python:util", "//tensorflow/python/estimator:estimator_py", "//third_party/py/numpy", ], @@ -336,12 +297,17 @@ py_library( visibility = ["//tensorflow:internal"], deps = [ ":one_device_strategy", - ":values", "//tensorflow/contrib/tpu:tpu_lib", "//tensorflow/python:constant_op", "//tensorflow/python:control_flow_ops", + "//tensorflow/python:dtypes", "//tensorflow/python:framework_ops", + "//tensorflow/python:tensor_util", "//tensorflow/python:util", + "//tensorflow/python/distribute:input_lib", + "//tensorflow/python/distribute:numpy_dataset", + "//tensorflow/python/distribute:reduce_util", + "//tensorflow/python/distribute:values", ], ) @@ -351,7 +317,6 @@ cuda_py_test( additional_deps = [ ":collective_all_reduce_strategy", ":combinations", - ":cross_tower_utils", ":multi_worker_test_base", ":strategy_test_lib", "@absl_py//absl/testing:parameterized", @@ -367,6 +332,7 @@ cuda_py_test( "//tensorflow/python:layers", "//tensorflow/python:variable_scope", "//tensorflow/python:variables", + "//tensorflow/python/distribute:cross_device_utils", "//tensorflow/python/eager:context", "//tensorflow/python/estimator:estimator_py", ], @@ -376,10 +342,14 @@ cuda_py_test( ], ) -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", + "no_pip", + ], deps = [ ":combinations", ":mirrored_strategy", @@ -399,18 +369,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"], @@ -466,15 +424,31 @@ cuda_py_test( ], tags = [ "multi_and_single_gpu", + "no_oss", # http://b/119349471 "no_pip", + "tf_integration_test", + ], +) + +cuda_py_test( + name = "keras_optimizer_v2_test", + srcs = ["keras_optimizer_v2_test.py"], + additional_deps = [ + ":keras_test_lib", + ], + tags = [ + "multi_and_single_gpu", + "no_oss", # http://b/119349471 + "no_pip", + "tf_integration_test", ], ) cuda_py_test( name = "estimator_training_test", - size = "enormous", srcs = ["estimator_training_test.py"], additional_deps = [ + ":collective_all_reduce_strategy", ":combinations", ":mirrored_strategy", ":multi_worker_test_base", @@ -482,7 +456,9 @@ cuda_py_test( "//third_party/py/numpy", "//tensorflow/contrib/optimizer_v2:training", "//tensorflow/python/data/ops:dataset_ops", - "//tensorflow/python/distribute", + "//tensorflow/python/distribute:distribute_config", + "//tensorflow/python/distribute:distribute_coordinator", + "//tensorflow/python/distribute:distribute_coordinator_context", "//tensorflow/python/eager:test", "//tensorflow/python/estimator:estimator_py", "//tensorflow/python/feature_column", @@ -490,9 +466,15 @@ cuda_py_test( "//tensorflow/python:platform", "//tensorflow/python:summary", ], + shard_count = 48, tags = [ "multi_and_single_gpu", "no_pip", + # TODO(b/118768923): Re-enable {a,m,t}san test. + "noasan", + "nomsan", + "notsan", + "no_oss", # http://b/119349471 ], ) @@ -509,10 +491,14 @@ 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", + "no_pip", + ], deps = [ ":combinations", ":single_loss_example", @@ -525,18 +511,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"], @@ -568,52 +542,16 @@ cuda_py_test( ], ) -py_library( - name = "shared_variable_creator", - srcs = ["shared_variable_creator.py"], - visibility = ["//tensorflow:internal"], -) - -py_test( - name = "shared_variable_creator_test", - srcs = ["shared_variable_creator_test.py"], - srcs_version = "PY2AND3", - deps = [ - ":shared_variable_creator", - "//tensorflow/python:framework_test_lib", - "//tensorflow/python:variable_scope", - "//tensorflow/python/eager:test", - ], -) - -py_library( - name = "cross_tower_utils", - srcs = ["cross_tower_utils.py"], - srcs_version = "PY2AND3", - deps = [ - ":values", - "//tensorflow/contrib/all_reduce:all_reduce_py", - "//tensorflow/contrib/nccl:nccl_py", - "//tensorflow/python:array_ops", - "//tensorflow/python:collective_ops", - "//tensorflow/python:device", - "//tensorflow/python:dtypes", - "//tensorflow/python:framework_ops", - "//tensorflow/python:gradients", - "//tensorflow/python:math_ops", - ], -) - cuda_py_test( - name = "cross_tower_utils_test", - srcs = ["cross_tower_utils_test.py"], + name = "cross_device_utils_test", + srcs = ["cross_device_utils_test.py"], additional_deps = [ ":combinations", - ":cross_tower_utils", "@absl_py//absl/testing:parameterized", "//tensorflow/python:constant_op", "//tensorflow/python:framework_ops", "//tensorflow/python:math_ops", + "//tensorflow/python/distribute:cross_device_utils", "//tensorflow/python/eager:context", "//tensorflow/python/eager:test", ], @@ -622,41 +560,20 @@ cuda_py_test( ], ) -py_library( - name = "cross_tower_ops", - srcs = ["cross_tower_ops.py"], - srcs_version = "PY2AND3", - deps = [ - ":cross_tower_utils", - ":values", - "//tensorflow/python:array_ops", - "//tensorflow/python:device_lib", - "//tensorflow/python:framework_ops", - "//tensorflow/python:math_ops", - "//tensorflow/python:platform", - "//tensorflow/python:resource_variable_ops", - "//tensorflow/python:training", - "//tensorflow/python:variable_scope", - "//tensorflow/python/eager:context", - "@six_archive//:six", - ], -) - cuda_py_test( - name = "cross_tower_ops_test", - size = "large", - srcs = ["cross_tower_ops_test.py"], + name = "cross_device_ops_test", + srcs = ["cross_device_ops_test.py"], additional_deps = [ ":combinations", - ":cross_tower_ops", ":multi_worker_test_base", ":mirrored_strategy", - ":values", "@absl_py//absl/testing:parameterized", "//tensorflow/python:array_ops", "//tensorflow/python:constant_op", "//tensorflow/python:framework_ops", "//tensorflow/python:math_ops", + "//tensorflow/python/distribute:cross_device_ops", + "//tensorflow/python/distribute:values", "//tensorflow/python/eager:context", "//tensorflow/python/eager:test", ], @@ -667,46 +584,80 @@ cuda_py_test( ) py_library( - name = "input_ops", - srcs = ["input_ops.py"], - visibility = ["//tensorflow:internal"], + name = "keras_test_lib", + testonly = 1, + srcs = [ + "keras_backward_compat_test.py", + "keras_test.py", + ], deps = [ - "//tensorflow/python:framework_ops", - "//tensorflow/python/data/util:nest", + ":combinations", + "//tensorflow/contrib/distribute/python:mirrored_strategy", + "//tensorflow/contrib/distribute/python:tpu_strategy", + "//tensorflow/python:client_testlib", + "//tensorflow/python:training", + "//tensorflow/python/eager:test", + "//tensorflow/python/estimator:estimator_py", + "//tensorflow/python/keras", + "//third_party/py/numpy", + "@absl_py//absl/testing:parameterized", ], ) -cuda_py_test( - name = "input_ops_test", - srcs = ["input_ops_test.py"], - additional_deps = [ - ":input_ops", - "//tensorflow/python/data/ops:dataset_ops", - "//tensorflow/contrib/data/python/ops:batching", - "//tensorflow/contrib/data/python/ops:interleave_ops", - "//tensorflow/python:errors", - "//tensorflow/python:client_testlib", - "//tensorflow/python:framework_ops", - "//tensorflow/python:framework_test_lib", - "//tensorflow/python:io_ops", - "//tensorflow/python/data/ops:readers", - "//tensorflow/python:util", +distribute_py_test( + name = "keras_test", + srcs = ["keras_test.py"], + main = "keras_test.py", + shard_count = 16, + tags = [ + "multi_and_single_gpu", + "no_oss", # TODO(b/117919883): Fix python error. + "no_pip", + "no_windows_gpu", + "notsan", + ], + deps = [ + ":keras_test_lib", ], +) + +# TODO(b/121200287): Remove this in 2.0 +distribute_py_test( + name = "keras_backward_compat_test", + srcs = ["keras_backward_compat_test.py"], + full_precision = True, + main = "keras_backward_compat_test.py", + shard_count = 16, 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_test_lib", + name = "keras_correctness_test_lib", testonly = 1, - srcs = ["keras_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", "//tensorflow/contrib/distribute/python:tpu_strategy", "//tensorflow/python:client_testlib", "//tensorflow/python:training", + "//tensorflow/python/eager:test", "//tensorflow/python/estimator:estimator_py", "//tensorflow/python/keras", "//third_party/py/numpy", @@ -714,47 +665,128 @@ py_library( ], ) -cuda_py_test( - name = "keras_test", - srcs = ["keras_test.py"], - additional_deps = [ - ":keras_test_lib", +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_pip", + "no_windows_gpu", + "notsan", ], - shard_count = 16, + deps = [ + ":keras_correctness_test_lib", + ], +) + +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_pip", "no_windows_gpu", "notsan", ], + deps = [ + ":keras_correctness_test_lib", + ], ) -py_library( - name = "metrics_v1_test_lib", - testonly = 1, - srcs = ["metrics_v1_test.py"], +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_pip", + "no_windows_gpu", + "notsan", + ], deps = [ - ":combinations", - "//tensorflow/contrib/data/python/ops:batching", - "//tensorflow/python:math_ops", - "//tensorflow/python:metrics", - "//tensorflow/python:variables", - "//tensorflow/python/data/ops:dataset_ops", - "//tensorflow/python/eager:test", - "@absl_py//absl/testing:parameterized", + ":keras_correctness_test_lib", ], ) -cuda_py_test( +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_pip", + "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. + "no_pip", + "no_windows_gpu", + "notsan", + ], + deps = [ + ":keras_correctness_test_lib", + ], +) + +distribute_py_test( name = "metrics_v1_test", srcs = ["metrics_v1_test.py"], - additional_deps = [ - ":metrics_v1_test_lib", - ], + main = "metrics_v1_test.py", tags = [ "multi_and_single_gpu", "no_pip", ], + deps = [ + ":combinations", + "//tensorflow/python:math_ops", + "//tensorflow/python:metrics", + "//tensorflow/python:variables", + "//tensorflow/python/data/ops:dataset_ops", + "//tensorflow/python/eager:test", + "@absl_py//absl/testing:parameterized", + ], ) cuda_py_test( @@ -793,3 +825,24 @@ cuda_py_test( "no_pip", ], ) + +tf_xla_py_test( + name = "checkpointing_test", + srcs = ["checkpointing_test.py"], + disabled_backends = [ + # Only makes sense on TPUs + "cpu", + "gpu", + "cpu_ondemand", + ], + tags = [ + "no_oss", + "no_pip", + ], + deps = [ + ":tpu_strategy", + "//tensorflow/compiler/tests:xla_test", + "//tensorflow/python/eager:test", + "//tensorflow/python/training/checkpointable:util", + ], +) diff --git a/tensorflow/contrib/distribute/python/checkpoint_utils_test.py b/tensorflow/contrib/distribute/python/checkpoint_utils_test.py index 865dba803f562e0ab98341dd8343e3c72b03d39b..3ef8b9574a36730dcc1d8fd42b6c7f364d84bbed 100644 --- a/tensorflow/contrib/distribute/python/checkpoint_utils_test.py +++ b/tensorflow/contrib/distribute/python/checkpoint_utils_test.py @@ -43,10 +43,12 @@ class CheckpointUtilsWithDistributionStrategyTest( distribution=[combinations.default_strategy, combinations.one_device_strategy, combinations.mirrored_strategy_with_gpu_and_cpu, - combinations.mirrored_strategy_with_two_gpus], - in_tower_mode=[True, False], + combinations.mirrored_strategy_with_two_gpus, + combinations.core_mirrored_strategy_with_gpu_and_cpu, + combinations.core_mirrored_strategy_with_two_gpus], + in_replica_mode=[True, False], mode=["graph"])) - def testInitFromCheckpoint(self, distribution, in_tower_mode): + 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( @@ -68,8 +70,8 @@ class CheckpointUtilsWithDistributionStrategyTest( self.assertAllEqual(v2_value, self.evaluate(v2)) with ops.Graph().as_default() as g, distribution.scope(): - if in_tower_mode: - distribution.call_for_each_tower(init_and_verify, g) + if in_replica_mode: + 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 new file mode 100644 index 0000000000000000000000000000000000000000..aa5b9f57b8a5bc12ee94399ec1fc5a55177a5b5d --- /dev/null +++ b/tensorflow/contrib/distribute/python/checkpointing_test.py @@ -0,0 +1,95 @@ +# 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. +# ============================================================================== + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import functools +import os + +from tensorflow.compiler.tests import xla_test +from tensorflow.contrib.distribute.python import tpu_strategy +from tensorflow.python.framework import constant_op +from tensorflow.python.framework import ops +from tensorflow.python.keras.engine import training +from tensorflow.python.keras.layers import core +from tensorflow.python.platform import test +from tensorflow.python.training import adam as adam_v1 +from tensorflow.python.training import checkpoint_management +from tensorflow.python.training import training_util +from tensorflow.python.training.checkpointable import tracking +from tensorflow.python.training.checkpointable import util as checkpointable_utils + + +class NonLayerCheckpointable(tracking.AutoCheckpointable): + + def __init__(self): + super(NonLayerCheckpointable, self).__init__() + self.a_variable = checkpointable_utils.add_variable( + self, name="a_variable", shape=[]) + + +class Subclassed(training.Model): + """A concrete Model for testing.""" + + def __init__(self): + super(Subclassed, self).__init__() + self._named_dense = core.Dense(1, use_bias=True) + self._second = core.Dense(1, use_bias=False) + # We can still track Checkpointables which aren't Layers. + self._non_layer = NonLayerCheckpointable() + + def call(self, values): + ret = self._second(self._named_dense(values)) + return ret + + +class TrainingCheckpointTests(xla_test.XLATestCase): + + def testEagerTPUDistributionStrategy(self): + self.skipTest("b/121387144") + num_training_steps = 10 + checkpoint_directory = self.get_temp_dir() + checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") + + def _train_fn(optimizer, model): + input_value = constant_op.constant([[3.]]) + optimizer.minimize( + functools.partial(model, input_value), + global_step=root.optimizer_step) + + for training_continuation in range(3): + strategy = tpu_strategy.TPUStrategy() + with strategy.scope(): + model = Subclassed() + optimizer = adam_v1.AdamOptimizer(0.001) + root = checkpointable_utils.Checkpoint( + optimizer=optimizer, model=model, + optimizer_step=training_util.get_or_create_global_step()) + root.restore(checkpoint_management.latest_checkpoint( + checkpoint_directory)) + + for _ in range(num_training_steps): + strategy.extended.call_for_each_replica( + functools.partial(_train_fn, optimizer, model)) + root.save(file_prefix=checkpoint_prefix) + self.assertEqual((training_continuation + 1) * num_training_steps, + root.optimizer_step.numpy()) + + +if __name__ == "__main__": + ops.enable_eager_execution() + test.main() diff --git a/tensorflow/contrib/distribute/python/collective_all_reduce_strategy.py b/tensorflow/contrib/distribute/python/collective_all_reduce_strategy.py index 744cb7572e53fb72e2275108c8b58f811f33e0f1..aa4d82b4d0c0dffc66115346d5f82a9d64bcfa56 100644 --- a/tensorflow/contrib/distribute/python/collective_all_reduce_strategy.py +++ b/tensorflow/contrib/distribute/python/collective_all_reduce_strategy.py @@ -18,13 +18,20 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -from tensorflow.contrib.distribute.python import cross_tower_ops as cross_tower_ops_lib -from tensorflow.contrib.distribute.python import cross_tower_utils +import copy + from tensorflow.contrib.distribute.python import mirrored_strategy -from tensorflow.contrib.distribute.python import values 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 distribute_lib +from tensorflow.python.distribute import input_lib from tensorflow.python.distribute import multi_worker_util +from tensorflow.python.distribute import numpy_dataset +from tensorflow.python.distribute import values from tensorflow.python.eager import context +from tensorflow.python.eager import tape from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import collective_ops @@ -32,7 +39,7 @@ from tensorflow.python.platform import tf_logging as logging # TODO(yuefengz): support in-graph replication. -class CollectiveAllReduceStrategy(mirrored_strategy.MirroredStrategy): +class CollectiveAllReduceStrategy(distribute_lib.DistributionStrategy): """Distribution strategy that uses collective ops for all-reduce. It is similar to the MirroredStrategy but it uses collective ops for @@ -53,8 +60,21 @@ class CollectiveAllReduceStrategy(mirrored_strategy.MirroredStrategy): num_gpus_per_worker: number of local GPUs or GPUs per worker, the default is 0 meaning CPU only. """ + super(CollectiveAllReduceStrategy, self).__init__( + CollectiveAllReduceExtended(self, num_gpus_per_worker)) + + +class CollectiveAllReduceExtended(mirrored_strategy.MirroredExtended): + """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.""" @@ -62,19 +82,21 @@ class CollectiveAllReduceStrategy(mirrored_strategy.MirroredStrategy): self._num_workers = 1 if num_gpus_per_worker: - local_devices = [ + local_devices = tuple( "/device:GPU:%d" % i for i in range(num_gpus_per_worker) - ] + ) else: - local_devices = ["/device:CPU:0"] - - self._collective_keys = cross_tower_utils.CollectiveKeys() - super(CollectiveAllReduceStrategy, self).__init__( - devices=local_devices, - cross_tower_ops=cross_tower_ops_lib.CollectiveAllReduce( - num_workers=1, - num_gpus_per_worker=num_gpus_per_worker, - collective_keys=self._collective_keys)) + local_devices = ("/device:CPU:0",) + self._worker_device = device_util.canonicalize("/device:CPU:0") + self._host_input_device = numpy_dataset.SingleDevice(self._worker_device) + + self._collective_keys = cross_device_utils.CollectiveKeys() + self._initialize_local(local_devices) + # TODO(yuefengz): remove num_gpus_per_worker from CollectiveAllReduce. + 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 @@ -89,13 +111,12 @@ class CollectiveAllReduceStrategy(mirrored_strategy.MirroredStrategy): 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"]: + 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 = len(cluster_spec.as_dict().get("worker", [])) + len( - cluster_spec.as_dict().get("chief", [])) + 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`.") @@ -103,22 +124,24 @@ class CollectiveAllReduceStrategy(mirrored_strategy.MirroredStrategy): self._is_chief = multi_worker_util.is_chief(cluster_spec, task_type, task_id) - worker_device = "/job:%s/task:%d" % (task_type, task_id) + self._worker_device = "/job:%s/task:%d" % (task_type, task_id) + self._host_input_device = numpy_dataset.SingleDevice(self._worker_device) if num_gpus_per_worker: - local_devices = [ - "%s/device:GPU:%d" % (worker_device, i) + local_devices = tuple( + "%s/device:GPU:%d" % (self._worker_device, i) for i in range(num_gpus_per_worker) - ] + ) else: - local_devices = [worker_device] + local_devices = (self._worker_device,) - self._collective_keys = cross_tower_utils.CollectiveKeys() - super(CollectiveAllReduceStrategy, self).__init__( - devices=local_devices, - cross_tower_ops=cross_tower_ops_lib.CollectiveAllReduce( - num_workers=self._num_workers, - num_gpus_per_worker=num_gpus_per_worker, - collective_keys=self._collective_keys)) + 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. @@ -136,17 +159,26 @@ class CollectiveAllReduceStrategy(mirrored_strategy.MirroredStrategy): def _create_variable(self, next_creator, *args, **kwargs): colocate_with = kwargs.pop("colocate_with", None) - devices = self._get_devices_from(colocate_with) - group_size = len(devices) * self._num_workers - group_key = self._collective_keys.get_group_key(self._devices) + if colocate_with is None: + device_map = self._device_map + logical_device = 0 # TODO(josh11b): Get logical device from scope here. + elif isinstance(colocate_with, numpy_dataset.SingleDevice): + with ops.device(colocate_with.device): + return next_creator(*args, **kwargs) + else: + device_map = colocate_with.device_map + logical_device = colocate_with.logical_device def _real_mirrored_creator(devices, *args, **kwargs): """Creates one MirroredVariable on the current worker.""" - index = {} unique_var_name = ops.get_default_graph().unique_name( kwargs["name"], mark_as_used=False).rstrip("/") + # pylint: disable=protected-access collective_instance_key = self._collective_keys.get_instance_key( key_id=unique_var_name) + # Only the first device participles in the broadcast of initial values. + group_key = self._collective_keys.get_group_key([devices[0]]) + group_size = self._num_workers if "initial_value" not in kwargs: raise ValueError("Initial value must be specified.") initial_value = kwargs["initial_value"] @@ -155,64 +187,97 @@ class CollectiveAllReduceStrategy(mirrored_strategy.MirroredStrategy): else: initial_value_fn = lambda: initial_value + value_list = [] for i, d in enumerate(devices): - with ops.device(d): - if i > 0: + with ops.init_scope(), ops.device(d): + if i == 0: + # 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 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) + + assert index == 0, index + if self._num_workers > 1: + if self._is_chief: + 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) + return initial_value + else: # Give replicas meaningful distinct names: - var0name = index[devices[0]].name.split(":")[0] - # We append a / to variable names created on towers with id > 0 to + 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) + # Variables on non-first replica get initial values from the + # variables created on the first device of each worker. + def _overridden_initial_value_fn(device=d, index=i): + assert index > 0 + with ops.device(device): + if context.executing_eagerly(): + return array_ops.identity(value_list[0].value()) + else: + return array_ops.identity(value_list[0].initial_value) kwargs["initial_value"] = _overridden_initial_value_fn - with context.context().device_policy(context.DEVICE_PLACEMENT_SILENT): - v = next_creator(*args, **kwargs) + # Don't record operations (e.g. other variable reads) during + # variable creation. + with tape.stop_recording(): + 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) - index[d] = v - return index + value_list.append(v) + return value_list # pylint: disable=protected-access return mirrored_strategy._create_mirrored_variable( - devices, _real_mirrored_creator, *args, **kwargs) + self._container_strategy(), device_map, logical_device, + _real_mirrored_creator, *args, **kwargs) + + def _make_dataset_iterator(self, dataset): + return input_lib.DatasetIterator(dataset, self._input_workers, + self._num_replicas_in_sync) - def distribute_dataset(self, dataset_fn): + def _make_input_fn_iterator( + self, + input_fn, + replication_mode=distribute_lib.InputReplicationMode.PER_WORKER): """Distributes the dataset to each local GPU.""" - # TODO(yuefengz): shard the dataset. - return values.PerDeviceDataset( - self._call_dataset_fn(dataset_fn), self._devices, True) - - def configure(self, - session_config=None, - cluster_spec=None, - task_type=None, - task_id=None): + 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: @@ -231,9 +296,28 @@ class CollectiveAllReduceStrategy(mirrored_strategy.MirroredStrategy): # 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 session_config or not self._cluster_spec: - return + if not self._cluster_spec: + return updated_config assert self._task_type assert self._task_id is not None @@ -241,34 +325,28 @@ class CollectiveAllReduceStrategy(mirrored_strategy.MirroredStrategy): # Collective group leader is needed for collective ops to coordinate # workers. if "chief" in self._cluster_spec.jobs: - session_config.experimental.collective_group_leader = ( + 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`.") - session_config.experimental.collective_group_leader = ( + updated_config.experimental.collective_group_leader = ( "/job:worker/replica:0/task:0") # The device filters prevent communication between workers. - del session_config.device_filters[:] - session_config.device_filters.append( + del updated_config.device_filters[:] + updated_config.device_filters.append( "/job:%s/task:%d" % (self._task_type, self._task_id)) - # The scoped_allocator_optimization is to optimize graphs for collective - # ops. - rewrite_options = session_config.graph_options.rewrite_options - rewrite_options.scoped_allocator_optimization = ( - rewriter_config_pb2.RewriterConfig.ON) - del rewrite_options.scoped_allocator_opts.enable_op[:] - rewrite_options.scoped_allocator_opts.enable_op.append("CollectiveReduce") + return updated_config @property - def between_graph(self): + def experimental_between_graph(self): return True @property - def should_init(self): + def experimental_should_init(self): return True @property @@ -280,6 +358,17 @@ class CollectiveAllReduceStrategy(mirrored_strategy.MirroredStrategy): return self._is_chief @property - def num_replicas_in_sync(self): - return len(self._devices) * self._num_workers + 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): + """`make_dataset_iterator` and `make_numpy_iterator` use global batch size. + `make_input_fn_iterator` assumes per-replica batching. + + Returns: + Boolean. + """ + return True 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 de2c03f1a4af074f937274ba187a7be463dc78be..9b6236fd9f89ec30c1234c846930a05d9c32e99d 100644 --- a/tensorflow/contrib/distribute/python/collective_all_reduce_strategy_test.py +++ b/tensorflow/contrib/distribute/python/collective_all_reduce_strategy_test.py @@ -23,13 +23,19 @@ import numpy as np from tensorflow.contrib.distribute.python import collective_all_reduce_strategy from tensorflow.contrib.distribute.python import combinations -from tensorflow.contrib.distribute.python import cross_tower_utils from tensorflow.contrib.distribute.python import multi_worker_test_base +from tensorflow.contrib.distribute.python import strategy_test_lib from tensorflow.core.protobuf import config_pb2 +from tensorflow.core.protobuf import rewriter_config_pb2 from tensorflow.python import keras +from tensorflow.python.data.ops import dataset_ops +from tensorflow.python.distribute import cross_device_utils +from tensorflow.python.distribute import reduce_util +from tensorflow.python.distribute import values from tensorflow.python.eager import context 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.layers import core from tensorflow.python.ops import array_ops @@ -51,11 +57,6 @@ class CollectiveAllReduceStrategyTestBase( collective_key_base = 0 def setUp(self): - self._run_options = config_pb2.RunOptions() - self._run_options.experimental.collective_graph_key = 6 - - self._sess_config = config_pb2.ConfigProto() - # We use a different key_base for each test so that collective keys won't be # reused. # TODO(yuefengz, tucker): enable it to reuse collective keys in different @@ -66,33 +67,38 @@ class CollectiveAllReduceStrategyTestBase( 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=self._sess_config, + session_config=session_config, cluster_spec=self._cluster_spec, task_type=task_type, task_id=task_id) - collective_keys = cross_tower_utils.CollectiveKeys( + collective_keys = cross_device_utils.CollectiveKeys( group_key_start=10 * num_gpus + CollectiveAllReduceStrategyTestBase.collective_key_base, instance_key_start=num_gpus * 100 + CollectiveAllReduceStrategyTestBase.collective_key_base, instance_key_with_id_start=num_gpus * 10000 + CollectiveAllReduceStrategyTestBase.collective_key_base) - distribution._collective_keys = collective_keys - distribution._cross_tower_ops._collective_keys = collective_keys + 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] + return distribution, 'grpc://' + self._cluster_spec[task_type][ + task_id], session_config else: - return distribution, '' + return distribution, '', session_config def _test_minimize_loss_graph(self, task_type, task_id, num_gpus): - d, master_target = self._get_test_object(task_type, task_id, num_gpus) + d, master_target, config = self._get_test_object(task_type, task_id, + num_gpus) with ops.Graph().as_default(), \ - self.test_session(config=self._sess_config, - target=master_target) as sess, \ + self.cached_session(config=config, + target=master_target) as sess, \ d.scope(): - l = core.Dense(1, use_bias=False, name='gpu_%d' % d._num_gpus_per_worker) + l = core.Dense(1, use_bias=False, + name='gpu_%d' % d.extended._num_gpus_per_worker) def loss_fn(x): y = array_ops.reshape(l(x), []) - constant_op.constant(1.) @@ -117,32 +123,31 @@ class CollectiveAllReduceStrategyTestBase( def step(): """Perform one optimization step.""" # Run forward & backward to get gradients, variables list. - g_v = d.call_for_each_tower(grad_fn, 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 = [] for g, v in g_v: - fetched = d.read_var(v) + fetched = d.extended.read_var(v) before_list.append(fetched) with ops.control_dependencies([fetched]): # TODO(yuefengz): support non-Mirrored variable as destinations. - g = d.reduce( - variable_scope.VariableAggregation.SUM, g, destinations=v) + g = d.extended.reduce_to( + reduce_util.ReduceOp.SUM, g, destinations=v) with ops.control_dependencies( - d.update(v, update, g, grouped=False)): - after_list.append(d.read_var(v)) + d.extended.update(v, update, args=(g,), group=False)): + after_list.append(d.extended.read_var(v)) return before_list, after_list before_out, after_out = step() - if context.num_gpus() < d._num_gpus_per_worker: + if context.num_gpus() < d.extended._num_gpus_per_worker: return True - sess.run( - variables.global_variables_initializer(), options=self._run_options) + sess.run(variables.global_variables_initializer()) for i in range(10): - b, a = sess.run((before_out, after_out), options=self._run_options) + b, a = sess.run((before_out, after_out)) if i == 0: before, = b after, = a @@ -154,7 +159,8 @@ class CollectiveAllReduceStrategyTestBase( return error_after < error_before def _test_complex_model(self, task_type, task_id, num_gpus): - d, master_target = self._get_test_object(task_type, task_id, num_gpus) + d, master_target, config = self._get_test_object(task_type, task_id, + num_gpus) def model_fn(): """Mnist model with synthetic input.""" @@ -186,6 +192,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, @@ -193,10 +200,10 @@ class CollectiveAllReduceStrategyTestBase( return train_op with ops.Graph().as_default(), \ - self.test_session(config=self._sess_config, - target=master_target) as sess: + self.cached_session(config=config, + target=master_target) as sess: with d.scope(): - train_op = d.call_for_each_tower(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()) @@ -204,11 +211,11 @@ class CollectiveAllReduceStrategyTestBase( return True def _test_variable_initialization(self, task_type, task_id, num_gpus): - distribution, master_target = self._get_test_object(task_type, task_id, - num_gpus) + distribution, master_target, config = self._get_test_object( + task_type, task_id, num_gpus) with ops.Graph().as_default(), \ - self.test_session(config=self._sess_config, - target=master_target) as sess, \ + self.cached_session(config=config, + target=master_target) as sess, \ distribution.scope(): def model_fn(): @@ -219,27 +226,56 @@ class CollectiveAllReduceStrategyTestBase( 1.0, 10.0, dtype=dtypes.float32)) return array_ops.identity(x) - x = distribution.call_for_each_tower(model_fn) - reduced_x = distribution.unwrap( - distribution.reduce( - variable_scope.VariableAggregation.MEAN, x, - destinations='/cpu:0'))[0] + x = distribution.extended.call_for_each_replica(model_fn) + reduced_x = distribution.reduce(reduce_util.ReduceOp.MEAN, x) x = distribution.unwrap(x)[0] - sess.run( - variables.global_variables_initializer(), options=self._run_options) + sess.run(variables.global_variables_initializer()) - x_value, reduced_x_value = sess.run( - [x, reduced_x], options=self._run_options) + x_value, reduced_x_value = sess.run([x, reduced_x]) self.assertTrue( np.allclose(x_value, reduced_x_value, atol=1e-5), msg=('x_value = %r, reduced_x_value = %r' % (x_value, 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): + distribution, master_target, config = self._get_test_object( + task_type, task_id, num_gpus) + devices = distribution.extended.worker_devices + + with ops.Graph().as_default(), \ + self.cached_session(config=config, + target=master_target) as sess: + iterator = distribution.make_input_fn_iterator(input_fn) + 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) + + with self.assertRaises(errors.OutOfRangeError): + next_element = iterator.get_next() + sess.run([values.select_replica(r, next_element) + for r in range(len(devices))]) + + # After re-initializing the iterator, should be able to iterate again. + 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) + class DistributedCollectiveAllReduceStrategyTest( - CollectiveAllReduceStrategyTestBase, parameterized.TestCase): + CollectiveAllReduceStrategyTestBase, + strategy_test_lib.DistributionTestBase, + parameterized.TestCase): @classmethod def setUpClass(cls): @@ -267,7 +303,7 @@ class DistributedCollectiveAllReduceStrategyTest( combinations.combine(mode=['graph'], num_gpus=[0, 1, 2], required_gpus=1)) def testVariableInitialization(self, num_gpus): if context.num_gpus() < num_gpus: - return + self.skipTest('Not enough GPUs') self._run_between_graph_clients( self._test_variable_initialization, self._cluster_spec, @@ -277,10 +313,56 @@ class DistributedCollectiveAllReduceStrategyTest( combinations.combine(mode=['graph'], num_gpus=[0, 1, 2], required_gpus=1)) def testComplexModel(self, num_gpus): if context.num_gpus() < num_gpus: - return + self.skipTest('Not enough GPUs') self._run_between_graph_clients( self._test_complex_model, self._cluster_spec, num_gpus=num_gpus) + # 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): + if context.num_gpus() < num_gpus: + self.skipTest('Not enough GPUs') + dataset_fn = lambda: dataset_ops.Dataset.range(100) + # 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, + 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) + + 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) + + 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) + + # Verify group leader + self.assertEqual('/job:worker/replica:0/task:0', + new_config.experimental.collective_group_leader) + + # Verify device filters. + self.assertEqual(['/job:worker/task:1'], new_config.device_filters) + + # Verify rewrite options. + new_rewrite_options = new_config.graph_options.rewrite_options + self.assertEqual(rewriter_config_pb2.RewriterConfig.ON, + new_rewrite_options.scoped_allocator_optimization) + self.assertEqual(['CollectiveReduce'], + new_rewrite_options.scoped_allocator_opts.enable_op) + class DistributedCollectiveAllReduceStrategyTestWithChief( CollectiveAllReduceStrategyTestBase, parameterized.TestCase): @@ -291,10 +373,6 @@ class DistributedCollectiveAllReduceStrategyTestWithChief( cls._cluster_spec = multi_worker_test_base.create_in_process_cluster( num_workers=3, num_ps=0, has_chief=True) - def setUp(self): - super(DistributedCollectiveAllReduceStrategyTestWithChief, self).setUp() - self._run_options.experimental.collective_graph_key = 7 - @combinations.generate( combinations.combine(mode=['graph'], num_gpus=[0, 1, 2], required_gpus=1)) def testMinimizeLossGraph(self, num_gpus): @@ -321,20 +399,89 @@ class DistributedCollectiveAllReduceStrategyTestWithChief( class LocalCollectiveAllReduceStrategy( - CollectiveAllReduceStrategyTestBase, parameterized.TestCase): + 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)) + def testMinimizeLoss(self, num_gpus): # Collective ops doesn't support strategy with one device. if context.num_gpus() < num_gpus: - return - self._test_minimize_loss_graph(None, None, num_gpus) + self.skipTest('Not enough GPUs') + if context.executing_eagerly(): + strategy, _, _ = self._get_test_object(None, None, num_gpus) + self._test_minimize_loss_eager(strategy) + else: + self._test_minimize_loss_graph(None, None, num_gpus) - 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)) + def testComplexModel(self, num_gpus): if context.num_gpus() < num_gpus: - return + self.skipTest('Not enough GPUs') self._test_complex_model(None, None, num_gpus) + @combinations.generate( + combinations.combine(mode=['graph', 'eager'], required_gpus=2)) + def testMakeInputFnIterator(self): + num_gpus = 2 + dataset_fn = lambda: dataset_ops.Dataset.range(5 * num_gpus) + 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, + 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) + + def testAllReduceSum(self): + if context.num_gpus() < 2: self.skipTest('Not enough GPUs') + distribution, target, config = self._get_test_object(None, None, num_gpus=2) + with self.cached_session(config=config, target=target): + self._test_all_reduce_sum(distribution) + + def testAllReduceSumGradients(self): + if context.num_gpus() < 2: self.skipTest('Not enough GPUs') + distribution, target, config = self._get_test_object(None, None, num_gpus=2) + with self.cached_session(config=config, target=target): + self._test_all_reduce_sum_gradients(distribution) + + def testAllReduceSumGradientTape(self): + if context.num_gpus() < 2: self.skipTest('Not enough GPUs') + distribution, target, config = self._get_test_object(None, None, num_gpus=2) + with self.cached_session(config=config, target=target): + self._test_all_reduce_sum_gradient_tape(distribution) + + def testAllReduceMean(self): + if context.num_gpus() < 2: self.skipTest('Not enough GPUs') + distribution, target, config = self._get_test_object(None, None, num_gpus=2) + with self.cached_session(config=config, target=target): + self._test_all_reduce_mean(distribution) + + def testAllReduceMeanGradients(self): + if context.num_gpus() < 2: self.skipTest('Not enough GPUs') + distribution, target, config = self._get_test_object(None, None, num_gpus=2) + with self.cached_session(config=config, target=target): + self._test_all_reduce_mean_gradients(distribution) + + def testAllReduceMeanGradientTape(self): + if context.num_gpus() < 2: self.skipTest('Not enough GPUs') + distribution, target, config = self._get_test_object(None, None, num_gpus=2) + with self.cached_session(config=config, target=target): + self._test_all_reduce_mean_gradient_tape(distribution) + + def testNumpyIterator(self): + num_gpus = 2 + if context.num_gpus() < num_gpus: + self.skipTest('Not enough GPUs') + strategy, _, _ = self._get_test_object(None, None, num_gpus) + self._test_numpy_iterator(strategy) + if __name__ == '__main__': test.main() diff --git a/tensorflow/contrib/distribute/python/combinations.py b/tensorflow/contrib/distribute/python/combinations.py index 63a163e76cdd99c73399c657cbe9bc3d010369d2..88170d40857eb2d915c836e1d6672a866f545035 100644 --- a/tensorflow/contrib/distribute/python/combinations.py +++ b/tensorflow/contrib/distribute/python/combinations.py @@ -46,18 +46,19 @@ import unittest from absl.testing import parameterized import six -from tensorflow.contrib.cluster_resolver import TPUClusterResolver +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 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.training import adagrad from tensorflow.python.training import adam -from tensorflow.python.training import distribution_strategy_context from tensorflow.python.training import gradient_descent from tensorflow.python.training import rmsprop from tensorflow.python.util import tf_inspect @@ -168,6 +169,8 @@ def _augment_with_special_arguments(test_method): if GPU_TEST: self.skipTest("Test that doesn't require GPUs.") elif context.num_gpus() < required_gpus: + # TODO(priyag): Consider allowing tests in graph mode using soft + # placement. self.skipTest( "{} GPUs are not available for this test. {} GPUs are available". format(required_gpus, context.num_gpus())) @@ -190,7 +193,7 @@ def _augment_with_special_arguments(test_method): kwargs_to_pass[arg] = kwargs[arg] if mode == "eager": - with ops.Graph().as_default(), context.eager_mode(): + with context.eager_mode(): if distribution: kwargs_to_pass["distribution"] = distribution.strategy test_method(**kwargs_to_pass) @@ -319,33 +322,86 @@ class NamedDistribution(object): return self._required_tpu +def _get_tpu_strategy_creator(steps_per_run, use_single_core=False, **kwargs): + def _create_tpu_strategy(): + resolver = cluster_resolver.TPUClusterResolver("") + 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 + + # pylint: disable=g-long-lambda default_strategy = NamedDistribution( "Default", - distribution_strategy_context._get_default_distribution_strategy, # pylint: disable=protected-access + distribution_strategy_context._get_default_strategy, # pylint: disable=protected-access required_gpus=None) one_device_strategy = NamedDistribution( "OneDeviceCPU", lambda: one_device_lib.OneDeviceStrategy("/cpu:0"), required_gpus=None) tpu_strategy = NamedDistribution( - "TPU", lambda: tpu_lib.TPUStrategy( - TPUClusterResolver(""), steps_per_run=2), + "TPU", _get_tpu_strategy_creator(steps_per_run=2), required_tpu=True) tpu_strategy_one_step = NamedDistribution( - "TPUOneStep", lambda: tpu_lib.TPUStrategy( - TPUClusterResolver(""), steps_per_run=1), + "TPUOneStep", _get_tpu_strategy_creator(steps_per_run=1), + required_tpu=True) +tpu_strategy_loop_on_device_one_core = NamedDistribution( + "TPULoopOnDeviceOneCore", _get_tpu_strategy_creator( + steps_per_run=2, use_single_core=True, + _disable_training_loop_on_host=True), + required_tpu=True) +tpu_strategy_one_step_loop_on_device_one_core = NamedDistribution( + "TPUOneStepLoopOnDeviceOneCore", _get_tpu_strategy_creator( + steps_per_run=1, use_single_core=True, + _disable_training_loop_on_host=True), + required_tpu=True) +# TODO(b/122327153): Remove below two NamedDistributions. +tpu_strategy_loop_on_device = NamedDistribution( + "TPULoopOnDevice", _get_tpu_strategy_creator( + steps_per_run=2, _disable_training_loop_on_host=True), required_tpu=True) -# Note that we disable prefetching for testing since prefetching makes -# the input non-deterministic. +tpu_strategy_one_step_loop_on_device = NamedDistribution( + "TPUOneStepLoopOnDevice", _get_tpu_strategy_creator( + steps_per_run=1, _disable_training_loop_on_host=True), + required_tpu=True) + +mirrored_strategy_with_one_cpu = NamedDistribution( + "Mirrored1CPU", + lambda: mirrored_lib.MirroredStrategy(["/cpu:0"])) +mirrored_strategy_with_one_gpu = NamedDistribution( + "Mirrored1GPU", + lambda: mirrored_lib.MirroredStrategy(["/gpu:0"]), + required_gpus=1) mirrored_strategy_with_gpu_and_cpu = NamedDistribution( "MirroredCPUAndGPU", - lambda: mirrored_lib.MirroredStrategy( - ["/gpu:0", "/cpu:0"], prefetch_on_device=False), + lambda: mirrored_lib.MirroredStrategy(["/gpu:0", "/cpu:0"]), required_gpus=1) mirrored_strategy_with_two_gpus = NamedDistribution( "Mirrored2GPUs", - lambda: mirrored_lib.MirroredStrategy( - ["/gpu:0", "/gpu:1"], prefetch_on_device=False), + lambda: mirrored_lib.MirroredStrategy(["/gpu:0", "/gpu:1"]), + required_gpus=2) +core_mirrored_strategy_with_one_cpu = NamedDistribution( + "CoreMirrored1CPU", + lambda: mirrored_lib.CoreMirroredStrategy(["/cpu:0"])) +core_mirrored_strategy_with_one_gpu = NamedDistribution( + "CoreMirrored1GPU", + lambda: mirrored_lib.CoreMirroredStrategy(["/gpu:0"]), + required_gpus=1) +core_mirrored_strategy_with_gpu_and_cpu = NamedDistribution( + "CoreMirroredCPUAndGPU", + lambda: mirrored_lib.CoreMirroredStrategy(["/gpu:0", "/cpu:0"]), + required_gpus=1) +core_mirrored_strategy_with_two_gpus = NamedDistribution( + "CoreMirrored2GPUs", + lambda: mirrored_lib.CoreMirroredStrategy(["/gpu:0", "/gpu:1"]), required_gpus=2) @@ -377,8 +433,11 @@ def distributions_and_v1_optimizers(): """A common set of combination with DistributionStrategies and Optimizers.""" return combine( distribution=[ - one_device_strategy, mirrored_strategy_with_gpu_and_cpu, - mirrored_strategy_with_two_gpus + one_device_strategy, + mirrored_strategy_with_gpu_and_cpu, + mirrored_strategy_with_two_gpus, + core_mirrored_strategy_with_gpu_and_cpu, + core_mirrored_strategy_with_two_gpus, ], optimizer_fn=optimizers_v1) @@ -387,7 +446,10 @@ def distributions_and_v2_optimizers(): """DistributionStrategies and V2 Optimizers.""" return combine( distribution=[ - one_device_strategy, mirrored_strategy_with_gpu_and_cpu, - mirrored_strategy_with_two_gpus + one_device_strategy, + mirrored_strategy_with_gpu_and_cpu, + mirrored_strategy_with_two_gpus, + core_mirrored_strategy_with_gpu_and_cpu, + core_mirrored_strategy_with_two_gpus, ], optimizer_fn=optimizers_v2) diff --git a/tensorflow/contrib/distribute/python/cross_device_ops_test.py b/tensorflow/contrib/distribute/python/cross_device_ops_test.py new file mode 100644 index 0000000000000000000000000000000000000000..54cce2988383fcf5e063726948fbbf62c7094ce5 --- /dev/null +++ b/tensorflow/contrib/distribute/python/cross_device_ops_test.py @@ -0,0 +1,590 @@ +# 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 CrossDeviceOps.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import itertools + +from absl.testing import parameterized +import numpy as np + +from tensorflow.contrib.distribute.python import combinations +from tensorflow.contrib.distribute.python import mirrored_strategy +from tensorflow.contrib.distribute.python import multi_worker_test_base +from tensorflow.core.protobuf import 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 reduce_util +from tensorflow.python.distribute import values as value_lib +from tensorflow.python.eager import context +from tensorflow.python.eager import test +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 math_ops + + +def _get_devices(devices): + if isinstance(devices, (tuple, list)): + return tuple(device_util.resolve(d) for d in devices) + elif isinstance(devices, value_lib.DistributedValues): + return devices.devices + return (device_util.resolve(devices),) + + +def _make_per_replica(values, devices, regroup=False): + devices = _get_devices(devices) + assert len(values) == len(devices) + + # We simulate the result of regroup called on PerReplica which strips the + # PerReplica wrapper if it has only one value. + if len(values) == 1 and regroup: + with ops.device(devices[0]): + placed_v = array_ops.identity(values[0]) + return placed_v + + index = [] + for d, v in zip(devices, values): + with ops.device(d): + placed_v = array_ops.identity(v) + index.append(placed_v) + return value_lib.PerReplica(value_lib.ReplicaDeviceMap(devices), index) + + +# pylint: disable=g-doc-args,g-doc-return-or-yield +def _fake_mirrored(value, devices): + """Create a faked Mirrored object for testing. + + All components of the returned Mirrored have the same objects, which is not + true in reality. + """ + devices = _get_devices(devices) + return value_lib.Mirrored(value_lib.ReplicaDeviceMap(devices), + [value] * len(devices)) + + +def _make_indexed_slices(values, indices, dense_shape, device): + with ops.device(device): + tensor = ops.IndexedSlices( + values=constant_op.constant(values), + indices=constant_op.constant(indices), + dense_shape=constant_op.constant(dense_shape)) + return tensor + + +def _make_mirrored_indexed_slices(devices, values, indices, dense_shape): + values = [_make_indexed_slices(values, indices, dense_shape, d) + for d in devices] + return value_lib.Mirrored(value_lib.ReplicaDeviceMap(devices), values) + + +_cpu_device = "/device:CPU:0" + + +class CrossDeviceOpsTestBase(test.TestCase, parameterized.TestCase): + + def _assert_indexed_slices_equal(self, left, right): + self.assertIsInstance(left, ops.IndexedSlices) + self.assertIsInstance(right, ops.IndexedSlices) + self.assertEqual(device_util.resolve(left.device), + device_util.resolve(right.device)) + self.assertAllEqual( + self.evaluate(ops.convert_to_tensor(left)), + self.evaluate(ops.convert_to_tensor(right))) + + def _assert_values_equal(self, left, right): + if isinstance(left, list): + for l, r in zip(left, right): + self._assert_values_equal(l, r) + else: + self.assertEqual(type(left), type(right)) + self.assertEqual(set(left.devices), set(right.devices)) + if isinstance(left.values[0], ops.IndexedSlices): + for d in left.devices: + self._assert_indexed_slices_equal(left.get(d), right.get(d)) + elif context.executing_eagerly(): + self.assertEqual([v.numpy() for v in left.values], + list(right.values)) + else: + with self.cached_session() as sess: + self.assertEqual( + sess.run(list(left.values)), list(right.values)) + + def _testReductionAndBroadcast(self, cross_device_ops, distribution): + devices = distribution.extended.worker_devices + + values = [constant_op.constant(float(d)) for d in range(len(devices))] + per_replica = _make_per_replica(values, devices) + mean = (len(devices) - 1.) / 2. + + values_2 = [constant_op.constant(d + 1.0) for d in range(len(devices))] + per_replica_2 = _make_per_replica(values_2, devices) + mean_2 = mean + 1. + + destination_mirrored = _fake_mirrored(1., devices) + destination_different = _fake_mirrored(1., _cpu_device) + destination_str = _cpu_device + + all_destinations = [ + destination_mirrored, destination_different, destination_str, + ] + + # test reduce() + for destinations in all_destinations: + self._assert_values_equal( + cross_device_ops.reduce( + reduce_util.ReduceOp.MEAN, + per_replica, + destinations=destinations), + _fake_mirrored(mean, destinations)) + self._assert_values_equal( + cross_device_ops.reduce( + reduce_util.ReduceOp.MEAN, + per_replica_2, + destinations=destinations), + _fake_mirrored(mean_2, destinations)) + self._assert_values_equal( + cross_device_ops.reduce( + reduce_util.ReduceOp.SUM, per_replica, + destinations=destinations), + _fake_mirrored(mean * len(devices), destinations)) + self._assert_values_equal( + cross_device_ops.reduce( + reduce_util.ReduceOp.SUM, + per_replica_2, + destinations=destinations), + _fake_mirrored(mean_2 * len(devices), destinations)) + + # test batch_reduce() + for d1, d2 in itertools.product(all_destinations, all_destinations): + self._assert_values_equal( + cross_device_ops.batch_reduce( + reduce_util.ReduceOp.MEAN, + [(per_replica, d1), (per_replica_2, d2)]), + [ + _fake_mirrored(mean, d1), + _fake_mirrored(mean_2, d2) + ]) + self._assert_values_equal( + cross_device_ops.batch_reduce( + reduce_util.ReduceOp.SUM, + [(per_replica, d1), (per_replica_2, d2)]), + [ + _fake_mirrored(mean * len(devices), d1), + _fake_mirrored(mean_2 * len(devices), d2) + ]) + + # test broadcast() + for destinations in all_destinations: + self._assert_values_equal( + cross_device_ops.broadcast(constant_op.constant(1.), destinations), + _fake_mirrored(1., destinations)) + + +class SingleWorkerCrossDeviceOpsTest(CrossDeviceOpsTestBase): + # TODO(yuefengz): decouple the num_gpus check from distribution in + # combinations module so that we can pass in devices instead of a distribution + # strategy. + reduction_to_one_combinations = combinations.combine( + cross_device_ops=[ + combinations.NamedObject( + "DefaultReductionToOneDeviceCrossDeviceOps", + cross_device_ops_lib.ReductionToOneDeviceCrossDeviceOps()), + combinations.NamedObject( + "ReductionToCPUDeviceCrossDeviceOps", + cross_device_ops_lib.ReductionToOneDeviceCrossDeviceOps( + reduce_to_device=_cpu_device)), + combinations.NamedObject( + "AccumulateNCrossDeviceOp", + cross_device_ops_lib.ReductionToOneDeviceCrossDeviceOps( + accumulation_fn=math_ops.accumulate_n)), + ], + distribution=[ + combinations.one_device_strategy, + combinations.mirrored_strategy_with_gpu_and_cpu, + combinations.mirrored_strategy_with_two_gpus, + combinations.core_mirrored_strategy_with_gpu_and_cpu, + combinations.core_mirrored_strategy_with_two_gpus + ], + mode=["graph", "eager"]) + allreduce_combinations = combinations.combine( + cross_device_ops=[ + 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( + "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], + mode=["graph", "eager"]) + + @combinations.generate(reduction_to_one_combinations + allreduce_combinations) + def testReductionAndBroadcast(self, cross_device_ops, distribution): + with distribution.scope(): + self._testReductionAndBroadcast(cross_device_ops, distribution) + + def testChooseAlgorithm(self): + device_links = [[1, 2, 3, 4], [0, 2, 3, 5], [0, 1, 3, 6], [0, 1, 2, 7], + [0, 5, 6, 7], [1, 4, 6, 7], [2, 4, 5, 7], [3, 4, 5, 6]] + result = cross_device_ops_lib._choose_all_reduce_algorithm(device_links) + self.assertIsInstance(result, cross_device_ops_lib.AllReduceCrossDeviceOps) + self.assertEqual(result._all_reduce_alg, "hierarchical_copy") + self.assertEqual(result._num_packs, 8) + + # if there are only 4 devices + device_links = [[1, 2, 3, 4], [0, 2, 3, 5], [0, 1, 3, 6], [0, 1, 2, 7]] + result = cross_device_ops_lib._choose_all_reduce_algorithm(device_links) + self.assertIsInstance(result, cross_device_ops_lib.AllReduceCrossDeviceOps) + self.assertEqual(result._all_reduce_alg, "nccl") + self.assertEqual(result._num_packs, 1) + + # if devices links contain each device itself + device_links = [[0, 1, 2, 3, 4], [0, 1, 2, 3, 5], [0, 1, 2, 3, 6], + [0, 1, 2, 3, 7], [0, 4, 5, 6, 7], [1, 4, 5, 6, 7], + [2, 4, 5, 6, 7], [3, 4, 5, 6, 7]] + result = cross_device_ops_lib._choose_all_reduce_algorithm(device_links) + self.assertIsInstance(result, cross_device_ops_lib.AllReduceCrossDeviceOps) + self.assertEqual(result._all_reduce_alg, "hierarchical_copy") + self.assertEqual(result._num_packs, 8) + + # if not dgx1-like links + device_links = [[0, 2, 3, 5], [0, 1, 3, 6], [0, 1, 2, 7], [0, 5, 6, 7], + [1, 4, 6, 7], [2, 4, 5, 7], [3, 4, 5, 6], [1, 2, 3, 4]] + result = cross_device_ops_lib._choose_all_reduce_algorithm(device_links) + self.assertIsInstance(result, cross_device_ops_lib.AllReduceCrossDeviceOps) + self.assertEqual(result._all_reduce_alg, "nccl") + self.assertEqual(result._num_packs, 1) + + @combinations.generate(combinations.combine( + mode=["graph", "eager"], + required_gpus=1)) + def testSimpleReduceWithIndexedSlices(self): + devices = ["/cpu:0", "/gpu:0"] + t0 = _make_indexed_slices([[1., 2.]], [1], [5, 2], devices[0]) + t1 = _make_indexed_slices([[3., 4.], [5., 6.]], [1, 3], [5, 2], devices[1]) + per_replica = value_lib.PerReplica( + value_lib.ReplicaDeviceMap(devices), (t0, t1)) + result = cross_device_ops_lib._simple_reduce( + per_replica, devices[0], math_ops.add_n, reduce_util.ReduceOp.SUM) + + # Test that the result is semantically equal to both the concatenated + # IndexedSlices with and without duplicate indices. + total_with_dups = _make_indexed_slices( + [[1., 2.], [3., 4.], [5., 6.]], [1, 1, 3], [5, 2], devices[0]) + total_without_dups = _make_indexed_slices( + [[4., 6.], [5., 6.]], [1, 3], [5, 2], devices[0]) + self._assert_indexed_slices_equal(total_with_dups, result) + self._assert_indexed_slices_equal(total_without_dups, result) + + @combinations.generate( + combinations.combine( + cross_device_ops_instance=[ + combinations.NamedObject( + "ReductionToOneDeviceCrossDeviceOps", + cross_device_ops_lib.ReductionToOneDeviceCrossDeviceOps()), + combinations.NamedObject( + "AllReduceCrossDeviceOps", + cross_device_ops_lib.AllReduceCrossDeviceOps()) + ], + reduce_op=[reduce_util.ReduceOp.SUM, reduce_util.ReduceOp.MEAN], + batch_reduce=[True, False], + mode=["graph", "eager"], + required_gpus=1)) + def testIndexedSlicesAllReduce(self, cross_device_ops_instance, reduce_op, + batch_reduce): + devices = ["/cpu:0", "/gpu:0"] + dense_shape = [5, 2] + t0 = _make_indexed_slices([[1., 2.]], [1], dense_shape, devices[0]) + t1 = _make_indexed_slices( + [[3., 4.], [5., 6.]], [1, 3], dense_shape, devices[1]) + per_replica = value_lib.PerReplica( + value_lib.ReplicaDeviceMap(devices), (t0, t1)) + + if batch_reduce: + result = cross_device_ops_instance.batch_reduce( + reduce_op, [(per_replica, per_replica)]) + else: + result = cross_device_ops_instance.reduce( + reduce_op, per_replica, per_replica) + + total_indices_with_dups = [1, 1, 3] + total_indices_without_dups = [1, 3] + + if reduce_op == reduce_util.ReduceOp.SUM: + total_values_with_dups = [[1., 2.], [3., 4.], [5., 6.]] + total_values_without_dups = [[4., 6.], [5., 6.]] + else: + assert reduce_op == reduce_util.ReduceOp.MEAN + total_values_with_dups = [[0.5, 1.], [1.5, 2.], [2.5, 3.]] + total_values_without_dups = [[2., 3.], [2.5, 3.]] + + total_mirrored_with_dups = _make_mirrored_indexed_slices( + devices, total_values_with_dups, total_indices_with_dups, dense_shape) + total_mirrored_without_dups = _make_mirrored_indexed_slices( + devices, total_values_without_dups, total_indices_without_dups, + dense_shape) + + # Test that the result is semantically equal to both the concatenated + # IndexedSlices, as well as when the duplicate indices are summed up. + if batch_reduce: + total_mirrored_with_dups = [total_mirrored_with_dups] + total_mirrored_without_dups = [total_mirrored_without_dups] + + self._assert_values_equal(total_mirrored_with_dups, result) + self._assert_values_equal(total_mirrored_without_dups, result) + + +class MultiWorkerCrossDeviceOpsTest(multi_worker_test_base.MultiWorkerTestBase, + CrossDeviceOpsTestBase): + + worker_devices = [ + "/job:worker/replica:0/task:0", "/job:worker/replica:0/task:1" + ] + multi_worker_allreduce_combinations = combinations.combine( + cross_device_ops=[ + combinations.NamedObject( + "MultiWorkerAllReduce", + cross_device_ops_lib.MultiWorkerAllReduce( + worker_devices, 2, ("pscpu/pscpu", 2, -1), 0, 0, 0)), + combinations.NamedObject( + "MultiWorkerAllReducePack", + cross_device_ops_lib.MultiWorkerAllReduce( + worker_devices, 2, ("pscpu/pscpu", 2, -1), 1, 0, 0)), + combinations.NamedObject( + "MultiWorkerAllReduceAggregation", + cross_device_ops_lib.MultiWorkerAllReduce( + worker_devices, 2, ("pscpu/pscpu", 2, -1), 0, 100, 10)), + combinations.NamedObject( + "MultiWorkerAllReduceMultipleSpecs", + cross_device_ops_lib.MultiWorkerAllReduce( + worker_devices, 2, [("pscpu/pscpu", 2, 100), + ("xring", 2, -1)], 0, 0, 0)), + ], + distribution=[ + combinations.NamedDistribution( + "MirroredCPU", + lambda: mirrored_strategy.MirroredStrategy(num_gpus_per_worker=0), + required_gpus=0), + combinations.NamedDistribution( + "Mirrored1GPU", + lambda: mirrored_strategy.MirroredStrategy(num_gpus_per_worker=1), + required_gpus=1), + combinations.NamedDistribution( + "Mirrored2GPUs", + lambda: mirrored_strategy.MirroredStrategy(num_gpus_per_worker=2), + required_gpus=2), + # pylint: disable=g-long-lambda + combinations.NamedDistribution( + "CoreMirroredCPU", + lambda: mirrored_strategy.CoreMirroredStrategy(["/device:CPU:0"]), + required_gpus=0), + combinations.NamedDistribution( + "CoreMirrored1GPU", + lambda: mirrored_strategy.CoreMirroredStrategy(["/device:GPU:0"]), + required_gpus=1), + combinations.NamedDistribution( + "CoreMirrored2GPUs", + lambda: mirrored_strategy.CoreMirroredStrategy( + ["/device:GPU:0", "/device:GPU:1"]), + required_gpus=2), + ], + mode=["graph"]) + + @combinations.generate(multi_worker_allreduce_combinations) + def testReductionAndBroadcast(self, cross_device_ops, distribution): + distribution.configure(cluster_spec={ + "worker": + ["/job:worker/replica:0/task:0", "/job:worker/replica:0/task:1"] + }) + with distribution.scope(): + self._testReductionAndBroadcast(cross_device_ops, distribution) + + +class MultiWorkerCollectiveAllReduceTest( + multi_worker_test_base.MultiWorkerTestBase, parameterized.TestCase): + + collective_key_base = 100000 + + @classmethod + def setUpClass(cls): + """Create a local cluster with 2 workers.""" + cls._cluster_spec = multi_worker_test_base.create_in_process_cluster( + num_workers=3, num_ps=0) + + def setUp(self): + super(MultiWorkerCollectiveAllReduceTest, self).setUp() + # Reusing keys are not supported well. So we have to give a different + # collective key base for different tests. + MultiWorkerCollectiveAllReduceTest.collective_key_base += 100000 + + def _get_test_objects(self, task_type, task_id, num_gpus=0, local_mode=False): + collective_keys = cross_device_utils.CollectiveKeys( + group_key_start=10 * num_gpus + + MultiWorkerCollectiveAllReduceTest.collective_key_base, + instance_key_start=num_gpus * 100 + + MultiWorkerCollectiveAllReduceTest.collective_key_base, + instance_key_with_id_start=num_gpus * 10000 + + MultiWorkerCollectiveAllReduceTest.collective_key_base) + if local_mode: + collective_all_reduce_ops = cross_device_ops_lib.CollectiveAllReduce( + 1, num_gpus, collective_keys=collective_keys) + if num_gpus: + devices = ["/device:GPU:%d" % i for i in range(num_gpus)] + else: + devices = ["/device:CPU:0"] + return collective_all_reduce_ops, devices, "" + else: + collective_all_reduce_ops = cross_device_ops_lib.CollectiveAllReduce( + 3, num_gpus, collective_keys=collective_keys) + if num_gpus: + devices = [ + "/job:%s/task:%d/device:GPU:%d" % (task_type, task_id, i) + for i in range(num_gpus) + ] + else: + devices = ["/job:%s/task:%d" % (task_type, task_id)] + return (collective_all_reduce_ops, devices, + "grpc://" + self._cluster_spec[task_type][task_id]) + + def _assert_values_equal(self, left, right, sess): + if isinstance(left, list): + for l, r in zip(left, right): + self._assert_values_equal(l, r, sess) + else: + self.assertEqual(type(left), type(right)) + self.assertEqual(set(left.devices), set(right.devices)) + + run_options = config_pb2.RunOptions() + run_options.experimental.collective_graph_key = 6 + + left_values = np.array( + sess.run(list(left.values), options=run_options)).flatten() + right_values = np.array(list(right.values)).flatten() + self.assertEqual(len(left_values), len(right_values)) + for l, r in zip(left_values, right_values): + self.assertEqual(l, r) + + def _test_reduction(self, task_type, task_id, num_gpus, local_mode=False): + collective_all_reduce, devices, master_target = self._get_test_objects( + task_type, task_id, num_gpus, local_mode=local_mode) + if local_mode: + num_workers = 1 + worker_device = None + else: + num_workers = len(self._cluster_spec.get("chief", [])) + len( + self._cluster_spec.get("worker", [])) + worker_device = "/job:%s/task:%d" % (task_type, task_id) + with ops.Graph().as_default(), \ + ops.device(worker_device), \ + self.cached_session(target=master_target) as sess: + # Collective ops doesn't support scalar tensors, so we have to construct + # 1-d tensors. + values = [constant_op.constant([float(d)]) for d in range(len(devices))] + per_replica = _make_per_replica(values, devices) + mean = np.array([(len(devices) - 1.) / 2.]) + + values_2 = [constant_op.constant([d + 1.0]) for d in range(len(devices))] + per_replica_2 = _make_per_replica(values_2, devices) + mean_2 = np.array([mean[0] + 1.]) + + destination_mirrored = _fake_mirrored(1., devices) + destination_different = _fake_mirrored(1., _cpu_device) + destination_str = _cpu_device + + all_destinations = [ + destination_different, destination_mirrored, destination_str + ] + + # test reduce() + for destinations in all_destinations: + self._assert_values_equal( + collective_all_reduce.reduce( + reduce_util.ReduceOp.MEAN, + per_replica, + destinations=destinations), + _fake_mirrored(mean, destinations), sess) + self._assert_values_equal( + collective_all_reduce.reduce( + reduce_util.ReduceOp.MEAN, + per_replica_2, + destinations=destinations), + _fake_mirrored(mean_2, destinations), sess) + self._assert_values_equal( + collective_all_reduce.reduce( + reduce_util.ReduceOp.SUM, + per_replica, + destinations=destinations), + _fake_mirrored(mean * len(devices) * num_workers, destinations), + sess) + self._assert_values_equal( + collective_all_reduce.reduce( + reduce_util.ReduceOp.SUM, + per_replica_2, + destinations=destinations), + _fake_mirrored(mean_2 * len(devices) * num_workers, destinations), + sess) + + # test batch_reduce() + for d1, d2 in itertools.product(all_destinations, all_destinations): + self._assert_values_equal( + collective_all_reduce.batch_reduce(reduce_util.ReduceOp.MEAN, + [(per_replica, d1), + (per_replica_2, d2)]), + [ + _fake_mirrored(mean, d1), + _fake_mirrored(mean_2, d2) + ], sess) + self._assert_values_equal( + collective_all_reduce.batch_reduce(reduce_util.ReduceOp.SUM, + [(per_replica, d1), + (per_replica_2, d2)]), + [ + _fake_mirrored(mean * len(devices) * num_workers, d1), + _fake_mirrored(mean_2 * len(devices) * num_workers, d2) + ], sess) + + return True + + @combinations.generate( + combinations.combine(mode=["graph"], num_gpus=[0, 1, 2], required_gpus=1)) + def testReductionDistributed(self, num_gpus): + if context.num_gpus() < num_gpus: + return + self._run_between_graph_clients(self._test_reduction, self._cluster_spec, + num_gpus) + + # Collective ops doesn't support strategy with one device. + def testReductionLocal(self, num_gpus=2): + if context.num_gpus() < num_gpus: + return + self._test_reduction(None, None, num_gpus, local_mode=True) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/distribute/python/cross_device_utils_test.py b/tensorflow/contrib/distribute/python/cross_device_utils_test.py new file mode 100644 index 0000000000000000000000000000000000000000..275aac2eeca575e927878d1ece63ce37ed38e8a0 --- /dev/null +++ b/tensorflow/contrib/distribute/python/cross_device_utils_test.py @@ -0,0 +1,142 @@ +# 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 cross_device_utils.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from absl.testing import parameterized + +from tensorflow.contrib.distribute.python import combinations +from tensorflow.python.distribute import cross_device_utils +from tensorflow.python.distribute import device_util +from tensorflow.python.distribute import values as value_lib +from tensorflow.python.eager import test +from tensorflow.python.framework import constant_op +from tensorflow.python.framework import ops +from tensorflow.python.framework import test_util +from tensorflow.python.ops import math_ops + + +class IndexedSlicesUtilsTest(test.TestCase, parameterized.TestCase): + + def _assert_values_equal(self, left, right): + self.assertAllEqual( + self.evaluate(ops.convert_to_tensor(left)), + self.evaluate(ops.convert_to_tensor(right))) + + @test_util.run_in_graph_and_eager_modes + def testAggregateTensors(self): + t0 = constant_op.constant([[1., 2.], [0, 0], [3., 4.]]) + t1 = constant_op.constant([[0., 0.], [5, 6], [7., 8.]]) + total = constant_op.constant([[1., 2.], [5, 6], [10., 12.]]) + result = cross_device_utils.aggregate_tensors_or_indexed_slices([t0, t1]) + self._assert_values_equal(total, result) + + @test_util.run_in_graph_and_eager_modes + def testAggregateIndexedSlices(self): + t0 = math_ops._as_indexed_slices( + constant_op.constant([[1., 2.], [0, 0], [3., 4.]])) + t1 = math_ops._as_indexed_slices( + constant_op.constant([[0., 0.], [5, 6], [7., 8.]])) + total = constant_op.constant([[1., 2.], [5, 6], [10., 12.]]) + result = cross_device_utils.aggregate_tensors_or_indexed_slices([t0, t1]) + self.assertIsInstance(result, ops.IndexedSlices) + self._assert_values_equal(total, result) + + @test_util.run_in_graph_and_eager_modes + def testDivideTensor(self): + t = constant_op.constant([[1., 2.], [0, 0], [3., 4.]]) + n = 2 + expected = constant_op.constant([[0.5, 1.], [0, 0], [1.5, 2.]]) + result = cross_device_utils.divide_by_n_tensors_or_indexed_slices(t, n) + self._assert_values_equal(expected, result) + + @test_util.run_in_graph_and_eager_modes + def testDivideIndexedSlices(self): + t = math_ops._as_indexed_slices( + constant_op.constant([[1., 2.], [0, 0], [3., 4.]])) + n = 2 + expected = constant_op.constant([[0.5, 1.], [0, 0], [1.5, 2.]]) + result = cross_device_utils.divide_by_n_tensors_or_indexed_slices(t, n) + self.assertIsInstance(result, ops.IndexedSlices) + self._assert_values_equal(expected, result) + + @test_util.run_in_graph_and_eager_modes + def testIsIndexedSlices(self): + t = math_ops._as_indexed_slices( + constant_op.constant([[1., 2.], [0, 0], [3., 4.]])) + self.assertTrue(cross_device_utils.contains_indexed_slices(t)) + + @test_util.run_in_graph_and_eager_modes + def testContainsIndexedSlices_List(self): + t0 = math_ops._as_indexed_slices( + constant_op.constant([[1., 2.], [0, 0], [3., 4.]])) + t1 = math_ops._as_indexed_slices( + constant_op.constant([[0., 0.], [5, 6], [7., 8.]])) + self.assertTrue(cross_device_utils.contains_indexed_slices([t0, t1])) + + @test_util.run_in_graph_and_eager_modes + def testContainsIndexedSlices_Tuple(self): + t0 = math_ops._as_indexed_slices( + constant_op.constant([[1., 2.], [0, 0], [3., 4.]])) + t1 = math_ops._as_indexed_slices( + constant_op.constant([[0., 0.], [5, 6], [7., 8.]])) + self.assertTrue(cross_device_utils.contains_indexed_slices((t0, t1))) + + @test_util.run_in_graph_and_eager_modes + def testContainsIndexedSlices_PerReplica(self): + t0 = math_ops._as_indexed_slices( + constant_op.constant([[1., 2.], [0, 0], [3., 4.]])) + t1 = math_ops._as_indexed_slices( + constant_op.constant([[0., 0.], [5, 6], [7., 8.]])) + device_map = value_lib.ReplicaDeviceMap(("/gpu:0", "/cpu:0")) + per_replica = value_lib.PerReplica(device_map, (t0, t1)) + self.assertTrue(cross_device_utils.contains_indexed_slices(per_replica)) + + @combinations.generate(combinations.combine( + mode=["graph", "eager"], + required_gpus=1)) + def testCopyTensor(self): + with ops.device("/cpu:0"): + t = constant_op.constant([[1., 2.], [0, 0], [3., 4.]]) + destination = "/gpu:0" + result = cross_device_utils.copy_tensor_or_indexed_slices_to_device( + t, destination) + + self._assert_values_equal(t, result) + self.assertEqual(device_util.resolve(destination), + device_util.resolve(result.device)) + + @combinations.generate(combinations.combine( + mode=["graph", "eager"], + required_gpus=1)) + def testCopyIndexedSlices(self): + with ops.device("/cpu:0"): + t = math_ops._as_indexed_slices( + constant_op.constant([[1., 2.], [0, 0], [3., 4.]])) + destination = "/gpu:0" + result = cross_device_utils.copy_tensor_or_indexed_slices_to_device( + t, destination) + + self.assertIsInstance(result, ops.IndexedSlices) + self._assert_values_equal(t, result) + self.assertEqual(device_util.resolve(destination), + device_util.resolve(result.device)) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/distribute/python/cross_tower_ops.py b/tensorflow/contrib/distribute/python/cross_tower_ops.py deleted file mode 100644 index e08ba9c2a668cd675defb025d7ad060e1338506b..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/distribute/python/cross_tower_ops.py +++ /dev/null @@ -1,959 +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. -# ============================================================================== -"""Classes for different algorithms of reduction and broadcasting.""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import collections -import six - -from tensorflow.contrib.distribute.python import cross_tower_utils -from tensorflow.contrib.distribute.python import values as value_lib -from tensorflow.python.client import device_lib -from tensorflow.python.eager import context -from tensorflow.python.framework import ops -from tensorflow.python.ops import array_ops -from tensorflow.python.ops import math_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_util - - -def check_destinations(destinations): - """Checks whether `destinations` is not empty. - - Args: - destinations: a DistributedValues, Variable, string or a list of strings. - - Returns: - Boolean which is True if `destinations` is not empty. - """ - # Calling bool() on a ResourceVariable is not allowed. - if isinstance(destinations, resource_variable_ops.ResourceVariable): - return bool(destinations.device) - return bool(destinations) - - -def validate_destinations(destinations): - if not isinstance( - destinations, - (value_lib.DistributedValues, resource_variable_ops.ResourceVariable, - value_lib.AggregatingVariable, six.string_types, list)): - raise ValueError("destinations must be one of a `DistributedValues` object," - " a tf.Variable object, a device string, a list of device " - "strings") - - if not check_destinations(destinations): - raise ValueError("destinations can not be empty") - - -def _make_tensor_into_per_device(input_tensor): - """Converts a single tensor into a PerDevice object.""" - if isinstance(input_tensor, (tuple, list)): - raise ValueError("Cannot convert `input_tensor` to a `PerDevice` object, " - "got %r but expected a object that is not a tuple or list." - % (input_tensor,)) - if isinstance(input_tensor, value_lib.PerDevice): - return input_tensor - - try: - device = input_tensor.device - except AttributeError: - raise ValueError("Cannot convert `input_tensor` to a `PerDevice` object " - "because it doesn't have device set.") - - return value_lib.PerDevice({device: input_tensor}) - - -def _normalize_value_destination_pairs(value_destination_pairs): - """Converts each tensor into a PerDevice object in the input list.""" - result = [] - if not isinstance(value_destination_pairs, (list, tuple)): - raise ValueError("`value_destination_pairs` should be a list or tuple") - for pair in value_destination_pairs: - if not isinstance(pair, tuple): - raise ValueError( - "Each element of `value_destination_pairs` should be a tuple.") - if len(pair) != 2: - raise ValueError("Each element of `value_destination_pairs` should be a " - "tuple of size 2.") - - per_device = _make_tensor_into_per_device(pair[0]) - result.append((per_device, pair[1])) - return result - - -def _validate_value_destination_pairs(value_destination_pairs): - # TODO(yuefengz): raise exceptions instead of returning False. - # pylint: disable=g-missing-docstring - if not value_destination_pairs: return False - if not isinstance(value_destination_pairs, (list, tuple)): return False - if not all([isinstance(pair, tuple) for pair in value_destination_pairs]): - return False - if not all([isinstance(v[0], value_lib.PerDevice) - for v in value_destination_pairs]): - return False - return True - - -# TODO(yuefengz): consider calling this function in the caller of CrossTowerOps. -def get_devices_from(destinations): - if isinstance(destinations, value_lib.DistributedValues): - return list(destinations.devices) - elif isinstance(destinations, (resource_variable_ops.ResourceVariable, - value_lib.AggregatingVariable)): - return [destinations.device] - elif isinstance(destinations, six.string_types): - return [device_util.resolve(destinations)] - elif isinstance(destinations, (list, tuple)): - return [device_util.resolve(destination) for destination in destinations] - else: - return [destinations.device] - - -def _devices_match(left, right): - return set(get_devices_from(left)) == set(get_devices_from(right)) - - -def _all_devices_match(value_destination_pairs): - if not all([_devices_match(v, d) for v, d in value_destination_pairs]): - return False - if not all([_devices_match(v, value_destination_pairs[0][0]) - for v, _ in value_destination_pairs[1:]]): - return False - return True - - -def _simple_broadcast(value, destinations): - index = {} - devices = get_devices_from(destinations) - for d in devices: - index[d] = cross_tower_utils.copy_tensor_or_indexed_slices_to_device( - value, d) - return value_lib.Mirrored(index) - - -def _simple_reduce(per_device_value, reduce_to_device, accumulation_fn, - aggregation): - # pylint: disable=g-missing-docstring - all_values = [] - count = 0 - for v in per_device_value._index.values(): # pylint: disable=protected-access - if isinstance(v, value_lib.MapOutput): - v_list = v.get() - if not v_list: - continue - count += len(v_list) - # Sum within each device before aggregating across devices. - # TODO(yuefengz): Check whether it helps to use accumulation_fn here. - v = cross_tower_utils.aggregate_tensors_or_indexed_slices( - v_list, math_ops.add_n) - else: - count += 1 - all_values.append(v) - if not all_values: - raise ValueError("`per_device_value` must be non-empty") - - with ops.device(reduce_to_device): - with context.context().device_policy(context.DEVICE_PLACEMENT_SILENT): - reduced = cross_tower_utils.aggregate_tensors_or_indexed_slices( - all_values, accumulation_fn) - if aggregation == vs.VariableAggregation.MEAN: - reduced = cross_tower_utils.divide_by_n_tensors_or_indexed_slices( - reduced, count) - elif aggregation != vs.VariableAggregation.SUM: - raise ValueError("`aggregation` must be VariableAggregation.SUM " - "or VariableAggregation.MEAN.") - return reduced - - -class CrossTowerOps(object): - """Base class for cross-tower reduction and broadcasting algorithms.""" - - def __init__(self): - pass - - def reduce(self, aggregation, per_device_value, destinations): - """Reduce `per_device_value` to `destinations`. - - It runs the reduction operation defined by `aggregation` and put the - result on `destinations`. - - Args: - aggregation: Indicates how a variable will be aggregated. Accepted values - are `tf.VariableAggregation.SUM`, `tf.VariableAggregation.MEAN`. - per_device_value: a PerDevice object or a tensor with device set. - destinations: the reduction destinations. - - Returns: - a Mirrored object. - - Raises: - ValueError: if per_device_value is not a PerDevice object. - """ - if not isinstance(per_device_value, value_lib.PerDevice): - per_device_value = _make_tensor_into_per_device(per_device_value) - - validate_destinations(destinations) - return self._reduce(aggregation, per_device_value, destinations) - - def batch_reduce(self, aggregation, value_destination_pairs): - """Reduce PerDevice objects in a batch. - - Reduce each first element in `value_destination_pairs` to each second - element which indicates the destinations. - - Args: - aggregation: Indicates how a variable will be aggregated. Accepted values - are `tf.VariableAggregation.SUM`, `tf.VariableAggregation.MEAN`. - value_destination_pairs: a list or a tuple of tuples of PerDevice objects - (or tensors with device set if there is one tower) and destinations. - - Returns: - a list of Mirrored objects. - - Raises: - ValueError: if `value_destination_pairs` is not a list or a tuple of - tuples of PerDevice objects and destinations - """ - if not _validate_value_destination_pairs(value_destination_pairs): - # If the first element of each pair is a tensor, we try to turn it into a - # PerDevice object. - value_destination_pairs = _normalize_value_destination_pairs( - value_destination_pairs) - - for _, d in value_destination_pairs: - validate_destinations(d) - - return self._batch_reduce(aggregation, value_destination_pairs) - - def broadcast(self, tensor, destinations): - """Broadcast the `tensor` to destinations. - - Args: - tensor: the tensor to broadcast. - destinations: the broadcast destinations. - - Returns: - a Mirrored object. - """ - validate_destinations(destinations) - return self._broadcast(tensor, destinations) - - def _reduce(self, aggregation, per_device_value, destinations): - raise NotImplementedError( - "_reduce method must be implemented in descendants.") - - def _batch_reduce(self, aggregation, value_destination_pairs): - raise NotImplementedError( - "_batch_reduce method must be implemented in descendants.") - - def _broadcast(self, tensor, destinations): - return _simple_broadcast(tensor, destinations) - - -class ReductionToOneDeviceCrossTowerOps(CrossTowerOps): - """Always do reduction to one device first and then do broadcasting. - - Batch reduction is done by reduction on each element one by one. - """ - - def __init__(self, reduce_to_device=None, accumulation_fn=math_ops.add_n): - """Constructor. - - Args: - reduce_to_device: the intermediate device to reduce to. If None, reduce - to the first device in `destinations` of the reduce() method. - accumulation_fn: a function that does accumulation. - """ - self.reduce_to_device = reduce_to_device - self.accumulation_fn = accumulation_fn - super(ReductionToOneDeviceCrossTowerOps, self).__init__() - - def _reduce(self, aggregation, per_device_value, destinations): - if check_destinations(destinations): - devices = get_devices_from(destinations) - else: - devices = get_devices_from(per_device_value) - reduce_to_device = self.reduce_to_device or devices[0] - reduced = _simple_reduce(per_device_value, reduce_to_device, - self.accumulation_fn, aggregation) - return self.broadcast(reduced, devices) - - def _batch_reduce(self, aggregation, value_destination_pairs): - return [ - self._reduce(aggregation, t, destinations=v) - for t, v in value_destination_pairs - ] - - -def _group_value_by_device(per_device_values): - """Group values into sublists by their devices. - - This grouping is needed to call the all-reduce library because it expects a - list of the following form: - [[(grad0_gpu0, v0_gpu0), (grad1_gpu0, v1_gpu0), (grad2_gpu0, v2_gpu0) ...], - [(grad0_gpu1, v0_gpu1), (grad1_gpu1, v1_gpu1), (grad2_gpu1, v2_gpu1) ...], - [(grad0_gpu2, v0_gpu2), (grad1_gpu0, v1_gpu2), (grad2_gpu0, v2_gpu2) ...], - ... - ] - - Args: - per_device_values: a list of PerDevice obejcts. - - Returns: - a list of lists, each sublist has components for its corresponding device of - PerDevice objects, paired with a None. - """ - destinations = per_device_values[0].devices - grouped = [[] for _ in range(len(destinations))] - for per_device_value in per_device_values: - # pylint: disable=protected-access - for i, v in enumerate(per_device_value._index.values()): - assert per_device_value.devices == destinations - grouped[i].append((v, None)) - return grouped - - -def _ungroup_and_make_mirrored(grouped_reduced, - destinations, - aggregation, - num_between_graph_workers=1): - """Ungroup results from all-reduce and make Mirrored objects. - - Each all-reduce result will be divided by the number of destinations before - Mirrored objects are created if aggregation is "mean". - - Args: - grouped_reduced: a list of lists, each sublist has components for each - device, paired with a None. It is the result from - cross_tower_utils.aggregate_gradients_using*. - destinations: a list of device strings for returned Mirrored objects. - aggregation: Indicates how a variable will be aggregated. Accepted values - are `tf.VariableAggregation.SUM`, `tf.VariableAggregation.MEAN`. - num_between_graph_workers: number of workers in the between-graph - replication. - - Returns: - a list of Mirrored objects. - """ - index = [{} for _ in range(len(grouped_reduced[0]))] - for d, per_device_reduced in enumerate(grouped_reduced): - for i, (v, _) in enumerate(per_device_reduced): - if aggregation == vs.VariableAggregation.MEAN: - index[i][destinations[d]] = v / ( - len(destinations) * num_between_graph_workers) - else: - index[i][destinations[d]] = v - return [value_lib.Mirrored(v) for v in index] - - -class ConcatAndSplitPacker(object): - """Concatenate and split tensors for reduction.""" - - def __init__(self, num_packs=1): - """Initialize the ConcatAndSplitPacker object. - - Args: - num_packs: specifies the number of split packs that will be - formed. - - Raises: - ValueError: if num_packs is not greater than 0. - """ - if num_packs <= 0: - raise ValueError("num_packs must be greater than zero.") - self.num_packs = num_packs - - def pack(self, grouped_grads_and_vars): - """Pack tensors.""" - self.grouped_grads_and_vars = grouped_grads_and_vars - self.all_tower_shapes = [] - self.all_tower_sizes = [] - - device_grad_packs = [] - for tower_grads_and_vars in grouped_grads_and_vars: - with ops.colocate_with(tower_grads_and_vars[0][0]): - # Flatten all the grads. - flat_grads = [ - array_ops.reshape(g, [-1]) for g, _ in tower_grads_and_vars - ] - # Remember the original shape of all the grads. - tower_shapes = [array_ops.shape(g) for g, _ in tower_grads_and_vars] - # Remember the original sizes of all the grads. - tower_sizes = [array_ops.size(g) for g, _ in tower_grads_and_vars] - # Concat all the flat grads into a big flat tensor. - concat_grads = array_ops.concat(flat_grads, 0) - - # Split the big tensor into num_splits packs. In cases where the - # total size is not divisible num_splits, the last pack gets - # more elements. - # TODO(zhengxq): it is also possible to optimize away all the concat - # as well. - num_splits = self.num_packs - - # The array_ops.size function will sometimes remove static shapes. So if - # all gradient shapes are defined, we use another method to get the - # total size. - # TODO(yuefengz): move this logic to array_ops.size. - if all([g.shape.is_fully_defined() for g, _ in tower_grads_and_vars]): - total_grad_size = sum( - [g.shape.num_elements() for g, _ in tower_grads_and_vars]) - else: - total_grad_size = array_ops.size(concat_grads) - - split_size = total_grad_size // num_splits - split_size_last = total_grad_size - split_size * (num_splits - 1) - split_sizes = [split_size] * (num_splits - 1) + [split_size_last] - grad_packs = array_ops.split(concat_grads, split_sizes) - - # Ready to aggregate the repacked gradients, with fake variables. - # TODO(zhengxq): It is hacky to have to use fake variables. - # We should remove the need for variables in - # aggregate_gradients_using*. - device_grad_packs.append(zip(grad_packs, [None] * num_splits)) - self.all_tower_shapes.append(tower_shapes) - self.all_tower_sizes.append(tower_sizes) - - return device_grad_packs - - def unpack(self, summed_device_grad_packs): - """Reverse the pack.""" - aggregated_device_grads = [] - for (summed_tower_grad_packs, - tower_grads_and_vars, tower_shapes, tower_sizes) in zip( - summed_device_grad_packs, self.grouped_grads_and_vars, - self.all_tower_shapes, self.all_tower_sizes): - # pylint: enable=line-too-long - # Reverse the packing operations in the previous steps. Form the - # summed gradients back into their original shapes. - with ops.colocate_with(summed_tower_grad_packs[0][0]): - # Form a list of the summed grad packs. - device_grad_packs = [g for g, _ in summed_tower_grad_packs] - - # Concat them back into a big flat tensor. - device_grads_concat = array_ops.concat(device_grad_packs, 0) - - # Split the tensors back into their original sizes. - grads_with_sizes = array_ops.split(device_grads_concat, tower_sizes) - - # Reshape the tensors back into their original shapes. - grads_with_shapes = [ - array_ops.reshape(grad, shape) - for shape, grad in zip(tower_shapes, grads_with_sizes) - ] - - # Form the list with the original list of variables. - summed_tower_grads = [ - (g, v) for g, (_, v) in zip(grads_with_shapes, tower_grads_and_vars) - ] - aggregated_device_grads.append(summed_tower_grads) - return aggregated_device_grads - - -class AggregateSmallTensorPacker(object): - """Concatenate small gradient tensors together for reduction.""" - - def __init__(self, - agg_small_grads_max_bytes=1048576, - agg_small_grads_max_group=16): - """Initialize the AggregateSmallTensorPacker object. - - Args: - agg_small_grads_max_bytes: largest tensor eligible for aggregation, - in number of bytes. - agg_small_grads_max_group: largest permitted aggregation of small - tensors. - - Raises: - ValueError: if `agg_small_grads_max_bytes` or `agg_small_grads_max_group` - is not greater than 0. - """ - if agg_small_grads_max_bytes <= 0 or agg_small_grads_max_group <= 0: - raise ValueError("agg_small_grads_max_bytes and agg_small_grads_max_group" - " should both be greater than zero.") - self.agg_small_grads_max_bytes = agg_small_grads_max_bytes - self.agg_small_grads_max_group = agg_small_grads_max_group - - def pack(self, grouped_grads_and_vars): - """Aggregate small tensors.""" - if (self.agg_small_grads_max_bytes > 0 and - self.agg_small_grads_max_group > 0): - tower_grads, self.packing = cross_tower_utils.pack_small_tensors( - grouped_grads_and_vars, - max_bytes=self.agg_small_grads_max_bytes, - max_group=self.agg_small_grads_max_group) - return tower_grads - - def unpack(self, summed_device_grad_packs): - """Reverse the aggregation process.""" - return cross_tower_utils.unpack_small_tensors(summed_device_grad_packs, - self.packing) - - -def _pack_tensors(device_grads, - num_packs=0, - agg_small_grads_max_bytes=0, - agg_small_grads_max_group=0): - """Pack tensors if specified.""" - if num_packs > 0: - tensor_packer = ConcatAndSplitPacker(num_packs) - device_grad_packs = tensor_packer.pack(device_grads) - elif agg_small_grads_max_bytes > 0 and agg_small_grads_max_group > 0: - tensor_packer = AggregateSmallTensorPacker(agg_small_grads_max_bytes, - agg_small_grads_max_group) - device_grad_packs = tensor_packer.pack(device_grads) - else: - tensor_packer = None - device_grad_packs = device_grads - return device_grad_packs, tensor_packer - - -def _unpack_tensors(reduced, tensor_packer=None): - """Unpack tensors if they are packed before all-reduce.""" - if tensor_packer: - return tensor_packer.unpack(reduced) - return reduced - - -class AllReduceCrossTowerOps(CrossTowerOps): - """Reduction using all reduce.""" - - def __init__(self, - all_reduce_alg="nccl", - num_packs=1, - agg_small_grads_max_bytes=0, - agg_small_grads_max_group=10): - """All-reduce implementation of CrossTowerOps. - - Before performing all-reduce, tensors will be repacked or aggregated for - more efficient cross-device transportation: - 1) If `num_packs` is non-zero, pack values into - `num_packs` splits. - 2) Otherwise, if `agg_small_grads_max_bytes` > 0 and - `agg_small_grads_max_group` > 0, aggregate values smaller than - `agg_small_grads_max_bytes` into groups with at most - `agg_small_grads_max_group` values. - 3) Otherwise, no repacking or grouping will happen. - - Args: - all_reduce_alg: the all-reduce algorithm to use, currently only "nccl" or - "hierarchical_copy" are supported. - num_packs: see above. - agg_small_grads_max_bytes: see above. - agg_small_grads_max_group: see above. - tensors. - """ - self._all_reduce_alg = all_reduce_alg - self._num_packs = num_packs - self._agg_small_grads_max_bytes = agg_small_grads_max_bytes - self._agg_small_grads_max_group = agg_small_grads_max_group - super(AllReduceCrossTowerOps, self).__init__() - - def _reduce(self, aggregation, per_device_value, destinations): - contains_indexed_slices = cross_tower_utils.contains_indexed_slices( - per_device_value) - if (_devices_match(per_device_value, destinations) - and not context.executing_eagerly() - and not contains_indexed_slices): - return self._batch_all_reduce(aggregation, [per_device_value])[0] - else: - if contains_indexed_slices: - logging.log_first_n( - logging.WARN, - "Efficient allreduce is not supported for IndexedSlices.", 10) - - if check_destinations(destinations): - devices = get_devices_from(destinations) - else: - devices = get_devices_from(per_device_value) - reduce_to_device = devices[0] - reduced = _simple_reduce(per_device_value, reduce_to_device, - math_ops.add_n, aggregation) - return self.broadcast(reduced, devices) - - def _batch_reduce(self, aggregation, value_destination_pairs): - all_devices_match = _all_devices_match(value_destination_pairs) - contains_indexed_slices = cross_tower_utils.contains_indexed_slices( - value_destination_pairs) - if (all_devices_match and not context.executing_eagerly() - and not contains_indexed_slices): - return self._batch_all_reduce(aggregation, - [v[0] for v in value_destination_pairs]) - else: - if not all_devices_match: - logging.log_first_n(logging.WARN, - "Efficient batch_reduce is not supported if " - "destinations are different.", - 10) - - return [ - self._reduce(aggregation, t, destinations=v) - for t, v in value_destination_pairs - ] - - def _batch_all_reduce(self, aggregation, per_device_values): - """All reduce algorithm in a batch.""" - logging.log_first_n( - logging.INFO, "batch_all_reduce invoked for batches size = %d with " - "algorithm = %s, num_packs = %d, agg_small_grads_max_bytes = %d and " - "agg_small_grads_max_group = %d" % - (len(per_device_values), self._all_reduce_alg, self._num_packs, - self._agg_small_grads_max_bytes, self._agg_small_grads_max_group), 10) - destinations = per_device_values[0].devices - grouped = _group_value_by_device(per_device_values) - - device_grad_packs, tensor_packer = _pack_tensors( - grouped, self._num_packs, self._agg_small_grads_max_bytes, - self._agg_small_grads_max_group) - - # The actual aggregation of the repacked gradients. Note that they are - # sharded among different aggregation trees. So it is important to strike - # the balance on num_splits. - if self._all_reduce_alg == "nccl": - # TODO(yuefengz): merge this into the all-reduce library. - reduced = cross_tower_utils.aggregate_gradients_using_nccl( - device_grad_packs) - else: - # TODO(yuefengz): check that gpu ids in `destinations` are in ascending - # order. - reduced = ( - cross_tower_utils.aggregate_gradients_using_hierarchical_copy( - destinations, device_grad_packs)) - - reduced = _unpack_tensors(reduced, tensor_packer) - return _ungroup_and_make_mirrored(reduced, per_device_values[0].devices, - aggregation) - - -AllReduceSpecTuple = collections.namedtuple("AllReduceSpecTuple", - "alg shards limit") - - -class MultiWorkerAllReduce(AllReduceCrossTowerOps): - """All-reduce algorithms for distributed TensorFlow.""" - - def __init__(self, - worker_devices, - num_gpus_per_worker, - all_reduce_spec=("pscpu/pscpu", 2, -1), - num_packs=0, - agg_small_grads_max_bytes=0, - agg_small_grads_max_group=10): - """Initialize the all-reduce algorithm. - - Args: - worker_devices: a list of device strings for workers participating in - all-reduce. - num_gpus_per_worker: number of GPU devices per worker. - all_reduce_spec: a tuple or a named tuple or a list of tuples specifying - the all-reduce algorithm. - 1. The first element of a tuple is the name of the all-reduce algorithm. - Valid algorithm names are: "nccl", "nccl/xring", "nccl/rechd", - "nccl/pscpu", "xring", "pscpu", "psgpu", "pscpu/pscpu". Algorithms with - a "/" are hierarchical, so two all-reduces are executed, the first one - aggregates tensors within a worker and the second aggregates across - workers. - 2. The second element of a tuple is the number of shards when doing - all-reduce. Let's say its values is M, each tensor after packing will be - split into M shards and then M parallel all-reduces would be performed - before finally they are concatenated backed into a complete tensor. - 3. The third element is the maximum size of tensors that will be - applicable for the algorithm specified by the first element. For - example, if all_reduce_spec=[("nccl", 2, 1024), ("pscpu/pscpu", 2, -1)], - tensors with size not larger than 1024 bytes will be applied a 2-shard - "nccl" all-reduce and other tensors will be applied a 2-shard - "pscpu/pscpu" algorithm. The third elements should be in increasing - order across tuples and end with -1 which indicates infinity. - num_packs: see AllReduceCrossTowerOps. - agg_small_grads_max_bytes: see AllReduceCrossTowerOps. - agg_small_grads_max_group: see AllReduceCrossTowerOps. - """ - self._worker_devices = worker_devices - self._num_gpus_per_worker = num_gpus_per_worker - super(MultiWorkerAllReduce, self).__init__( - num_packs=num_packs, - agg_small_grads_max_bytes=agg_small_grads_max_bytes, - agg_small_grads_max_group=agg_small_grads_max_group) - - def validate_and_complete_spec(spec): - """Validate and complete the all-reduce spec.""" - # TODO(yuefengz): support namedtuple. - if not isinstance(spec, tuple): - raise ValueError( - "A tuple is expected for all-reduce spec: %r" % all_reduce_spec) - if not spec or len(spec) > 3: - raise ValueError( - "Too many elements in the all-reduce spec tuple: %r" % spec) - if len(spec) == 1: - return AllReduceSpecTuple(spec[0], 1, -1) - elif len(spec) == 2: - return AllReduceSpecTuple(spec[0], spec[1], -1) - else: - return AllReduceSpecTuple(*spec) - - self._all_reduce_spec = [] - if isinstance(all_reduce_spec, six.string_types): - self._all_reduce_spec.append(AllReduceSpecTuple(all_reduce_spec, 1, -1)) - elif isinstance(all_reduce_spec, tuple): - self._all_reduce_spec.append(validate_and_complete_spec(all_reduce_spec)) - elif isinstance(all_reduce_spec, list): - self._all_reduce_spec = [ - validate_and_complete_spec(spec) for spec in all_reduce_spec - ] - - def _batch_all_reduce(self, aggregation, per_device_values): - """All reduce algorithm in a batch.""" - logging.log_first_n( - logging.INFO, - "distributed batch_all_reduce invoked for batches size = %d with " - "allreduce_spec = %r, num_packs = %d, agg_small_grads_max_bytes = %d " - "and agg_small_grads_max_group = %d" % - (len(per_device_values), self._all_reduce_spec, self._num_packs, - self._agg_small_grads_max_bytes, self._agg_small_grads_max_group), 10) - - destinations = sorted(per_device_values[0].devices) - device_grads = _group_value_by_device(per_device_values) - - # The all reduce library requires fully defined shapes. - # TODO(yuefengz): when tensor sharding is not needed, static shapes are not - # required as well. - for device_grad in device_grads: - for grad, _ in device_grad: - if not grad.shape.is_fully_defined(): - raise ValueError("Shape is unknown for node %r" % grad) - - remaining_grads = device_grads - aggregated_grads = [] - for spec_tuple in self._all_reduce_spec: - if spec_tuple.limit < 0: - this_grads = remaining_grads - remaining_grads = [] - else: - (this_grads, remaining_grads) = cross_tower_utils.split_grads_by_size( - spec_tuple.limit, remaining_grads) - if this_grads: - device_grad_packs, tensor_packer = _pack_tensors( - this_grads, self._num_packs, self._agg_small_grads_max_bytes, - self._agg_small_grads_max_group) - range_agg_grads = cross_tower_utils.sum_gradients_all_reduce( - self._worker_devices, device_grad_packs, len(self._worker_devices), - spec_tuple.alg, spec_tuple.shards, range(self._num_gpus_per_worker)) - range_agg_grads = _unpack_tensors(range_agg_grads, tensor_packer) - - if not aggregated_grads: - aggregated_grads = range_agg_grads - else: - assert len(aggregated_grads) == len(range_agg_grads) - for i in range(len(aggregated_grads)): - aggregated_grads[i] += range_agg_grads[i] - assert not remaining_grads - - return _ungroup_and_make_mirrored(aggregated_grads, destinations, - aggregation) - - -# TODO(yuefengz): support in-graph collective all-reduce. -class CollectiveAllReduce(CrossTowerOps): - """All-reduce cross tower ops using collective ops. - - In the between-graph replicated training, it will still do all-reduces across - all workers and then put results on the right destinations. - """ - - def __init__(self, - num_workers=1, - num_gpus_per_worker=0, - all_reduce_merge_scope=32, - collective_keys=None): - """Initializes the object. - - Args: - num_workers: number of workers in the between-graph replicated training. - num_gpus_per_worker: number of GPUs per worker. - all_reduce_merge_scope: size of groups into which to partition consecutive - gradients grouped under a common 'allreduce' name scope. This is useful - for some optimization of collective ops. - collective_keys: an optional CollectiveKey object. - """ - self._num_workers = num_workers - self._num_gpus_per_worker = num_gpus_per_worker - self._all_reduce_merge_scope = all_reduce_merge_scope - self._collective_keys = collective_keys or cross_tower_utils.CollectiveKeys( - ) - super(CollectiveAllReduce, self).__init__() - - # TODO(yuefengz, tucker): is indexed slices supported by collective ops? - def _reduce(self, aggregation, per_device_value, destinations): - if cross_tower_utils.contains_indexed_slices(per_device_value): - raise ValueError( - "`IndexSlices` is not supported for Collective All-Reduce.") - if context.executing_eagerly(): - raise ValueError( - "Eager execution is not supported for Collective All-Reduce") - - all_reduced = self._batch_all_reduce(aggregation, [per_device_value])[0] - if _devices_match(per_device_value, destinations): - return all_reduced - else: - index = {} - for d in get_devices_from(destinations): - # pylint: disable=protected-access - if d in all_reduced._index: - index[d] = all_reduced._index[d] - else: - with ops.control_dependencies(list( - all_reduced._index.values())), ops.device(d): - index[d] = array_ops.identity(list(all_reduced._index.values())[0]) - - return value_lib.Mirrored(index) - - def _batch_reduce(self, aggregation, value_destination_pairs): - if cross_tower_utils.contains_indexed_slices(value_destination_pairs): - raise ValueError( - "`IndexSlices` is not supported for Collective All-Reduce.") - if context.executing_eagerly(): - raise ValueError( - "Eager execution is not supported for Collective All-Reduce") - - all_devices_match = _all_devices_match(value_destination_pairs) - if all_devices_match: - return self._batch_all_reduce(aggregation, - [v[0] for v in value_destination_pairs]) - else: - if not all_devices_match: - logging.log_first_n( - logging.WARN, "Efficient batch_reduce is not supported if " - "destinations are different.", 10) - - return [ - self._reduce(aggregation, t, destinations=v) - for t, v in value_destination_pairs - ] - - def _batch_all_reduce(self, aggregation, per_device_values): - """All-reduce across all workers in a batch.""" - if context.executing_eagerly(): - raise ValueError( - "Eager execution with collective ops is not supported yet.") - - logging.log_first_n( - logging.INFO, "Collective All-reduce invoked with batches size = %d, " - "num_workers = %d" % (len(per_device_values), self._num_workers), 10) - - grouped_by_tower = _group_value_by_device(per_device_values) - - grouped_by_var = list(zip(*grouped_by_tower)) - # grouped_by_var is grouped by variables and takes the following format: - # [((grad0_gpu0, v0_gpu0), (grad0_gpu1, v0_gpu1), (grad0_gpu2, v0_gpu2) ..), - # ((grad1_gpu0, v1_gpu0), (grad1_gpu1, v1_gpu1), (grad1_gpu0, v1_gpu2) ..), - # ((grad2_gpu0, v2_gpu0), (grad2_gpu1, v2_gpu1), (grad2_gpu0, v2_gpu2) ..), - # ... - # ] - chunked_gv = [ - grouped_by_var[x:x + self._all_reduce_merge_scope] - for x in range(0, len(grouped_by_var), self._all_reduce_merge_scope) - ] - - reduced_gv_list = [] - for chunk in chunked_gv: - with ops.name_scope("allreduce"): - for grad_and_vars in chunk: - scaled_grads = [g for g, _ in grad_and_vars] - collective_reduced = cross_tower_utils.build_collective_reduce( - scaled_grads, self._num_workers, self._collective_keys, "Add", - "Id") - result = [] - for (_, v), g in zip(grad_and_vars, collective_reduced): - result.append([g, v]) - reduced_gv_list.append(result) - - new_tower_grads = [list(x) for x in zip(*reduced_gv_list)] - return _ungroup_and_make_mirrored( - new_tower_grads, - per_device_values[0].devices, - aggregation, - num_between_graph_workers=self._num_workers) - - -_dgx1_links = [[1, 2, 3, 4], [0, 2, 3, 5], [0, 1, 3, 6], [0, 1, 2, 7], - [0, 5, 6, 7], [1, 4, 6, 7], [2, 4, 5, 7], [3, 4, 5, 6]] - - -def _has_dgx1_like_links(gpu_links): - if not gpu_links: - return False - # TODO(yuefengz): figure out the right topology for hierarchial copy if - # number of gpus are less than 8. - if len(gpu_links) < 8: - return False - for i, (gpu_link, dgx1_link) in enumerate(zip(gpu_links, _dgx1_links)): - if (set(gpu_link) != set(dgx1_link) and - set(gpu_link) != set(dgx1_link + [i])): - return False - return True - - -def _choose_all_reduce_algorithm(device_links): - if _has_dgx1_like_links(device_links): - logging.info("Configured hierarchical_copy with num_packs=%d", - len(device_links)) - return AllReduceCrossTowerOps( - "hierarchical_copy", num_packs=len(device_links)) - else: - logging.info("Configured nccl all-reduce.") - return AllReduceCrossTowerOps("nccl", num_packs=1) - - -def choose_the_best(devices, session_config=None): - """Find the best subclass of CrossTowerOps given a tensorflow session. - - Args: - devices: a list of devices passed for distribute strategy. - session_config: a tensorflow session config or None. If None, it will make - deciesion based on all local devices. - - Returns: - a subclass of CrossTowerOps. - """ - requested_devices = set([device_util.canonicalize(d) for d in devices]) - machine_devices = device_lib.list_local_devices(session_config=session_config) - using_devices = [] - for d in machine_devices: - if device_util.canonicalize(d.name) in requested_devices: - using_devices.append(d) - else: - logging.info( - "Device is available but not used by distribute strategy: %s", d.name) - - if len(using_devices) != len(requested_devices): - logging.warning("Not all devices in distribute strategy are visible by " - "TensorFlow sessions.") - return ReductionToOneDeviceCrossTowerOps() - - if any([d.device_type.lower() != "gpu" for d in using_devices]): - logging.warning("Not all devices in DistributionStrategy are visible to " - "TensorFlow session.") - return ReductionToOneDeviceCrossTowerOps() - - device_links = [[] for _ in range(len(using_devices))] - for i, device in enumerate(using_devices): - for link in device.locality.links.link: - device_links[i].append(link.device_id) - - return _choose_all_reduce_algorithm(device_links) diff --git a/tensorflow/contrib/distribute/python/cross_tower_ops_test.py b/tensorflow/contrib/distribute/python/cross_tower_ops_test.py deleted file mode 100644 index 490371477a1b43551c4b4d8768c96d60e5f2c6d8..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/distribute/python/cross_tower_ops_test.py +++ /dev/null @@ -1,564 +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. -# ============================================================================== -"""Tests for CrossTowerOps.""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import itertools - -from absl.testing import parameterized -import numpy as np - -from tensorflow.contrib.distribute.python import combinations -from tensorflow.contrib.distribute.python import cross_tower_ops as cross_tower_ops_lib -from tensorflow.contrib.distribute.python import cross_tower_utils -from tensorflow.contrib.distribute.python import mirrored_strategy -from tensorflow.contrib.distribute.python import multi_worker_test_base -from tensorflow.contrib.distribute.python import values as value_lib -from tensorflow.core.protobuf import config_pb2 -from tensorflow.python.eager import context -from tensorflow.python.eager import test -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 math_ops -from tensorflow.python.ops import variable_scope as vs -from tensorflow.python.training import device_util - - -def _make_per_device(values, devices, regroup=False): - devices = cross_tower_ops_lib.get_devices_from(devices) - assert len(values) == len(devices) - - # We simulate the result of regroup called on PerDevice which strips the - # PerDevice wrapper if it has only one value. - if len(values) == 1 and regroup: - with ops.device(devices[0]): - placed_v = array_ops.identity(values[0]) - return placed_v - - index = {} - for d, v in zip(devices, values): - with ops.device(d): - placed_v = array_ops.identity(v) - index[d] = placed_v - return value_lib.PerDevice(index) - - -# pylint: disable=g-doc-args,g-doc-return-or-yield -def _fake_mirrored(value, devices): - """Create a faked Mirrored object for testing. - - All components of the returned Mirrored have the same objects, which is not - true in reality. - """ - devices = cross_tower_ops_lib.get_devices_from(devices) - return value_lib.Mirrored( - {d: v for d, v in zip(devices, [value] * len(devices))}) - - -def _make_indexed_slices(values, indices, dense_shape, device): - with ops.device(device): - tensor = ops.IndexedSlices( - values=constant_op.constant(values), - indices=constant_op.constant(indices), - dense_shape=constant_op.constant(dense_shape)) - return tensor - - -def _make_mirrored_indexed_slices(devices, values, indices, dense_shape): - return value_lib.Mirrored({ - d: _make_indexed_slices(values, indices, dense_shape, d) for d in devices - }) - - -_cpu_device = "/device:CPU:0" - - -class CrossTowerOpsTestBase(test.TestCase, parameterized.TestCase): - - def _assert_indexed_slices_equal(self, left, right): - self.assertIsInstance(left, ops.IndexedSlices) - self.assertIsInstance(right, ops.IndexedSlices) - self.assertEqual(device_util.resolve(left.device), - device_util.resolve(right.device)) - self.assertAllEqual( - self.evaluate(ops.convert_to_tensor(left)), - self.evaluate(ops.convert_to_tensor(right))) - - def _assert_values_equal(self, left, right): - if isinstance(left, list): - for l, r in zip(left, right): - self._assert_values_equal(l, r) - else: - self.assertEqual(type(left), type(right)) - self.assertEqual(set(left.devices), set(right.devices)) - if isinstance(list(left._index.values())[0], ops.IndexedSlices): - for (d, v) in left._index.items(): - self._assert_indexed_slices_equal(v, right._index[d]) - elif context.executing_eagerly(): - self.assertEqual([v.numpy() for v in left._index.values()], - list(right._index.values())) - else: - with self.test_session() as sess: - self.assertEqual( - sess.run(list(left._index.values())), list(right._index.values())) - - def _testReductionAndBroadcast(self, cross_tower_ops, distribution): - devices = distribution.worker_devices - - values = [constant_op.constant(float(d)) for d in range(len(devices))] - per_device = _make_per_device(values, devices) - mean = (len(devices) - 1.) / 2. - - values_2 = [constant_op.constant(d + 1.0) for d in range(len(devices))] - per_device_2 = _make_per_device(values_2, devices) - mean_2 = mean + 1. - - destination_mirrored = _fake_mirrored(1., devices) - destination_different = _fake_mirrored(1., _cpu_device) - destination_str = _cpu_device - destination_list = devices - - all_destinations = [ - destination_mirrored, destination_different, destination_str, - destination_list - ] - - # test reduce() - for destinations in all_destinations: - self._assert_values_equal( - cross_tower_ops.reduce( - vs.VariableAggregation.MEAN, - per_device, - destinations=destinations), - _fake_mirrored(mean, destinations)) - self._assert_values_equal( - cross_tower_ops.reduce( - vs.VariableAggregation.MEAN, - per_device_2, - destinations=destinations), - _fake_mirrored(mean_2, destinations)) - self._assert_values_equal( - cross_tower_ops.reduce( - vs.VariableAggregation.SUM, per_device, - destinations=destinations), - _fake_mirrored(mean * len(devices), destinations)) - self._assert_values_equal( - cross_tower_ops.reduce( - vs.VariableAggregation.SUM, - per_device_2, - destinations=destinations), - _fake_mirrored(mean_2 * len(devices), destinations)) - - # test batch_reduce() - for d1, d2 in itertools.product(all_destinations, all_destinations): - self._assert_values_equal( - cross_tower_ops.batch_reduce(vs.VariableAggregation.MEAN, - [(per_device, d1), (per_device_2, d2)]), - [ - _fake_mirrored(mean, d1), - _fake_mirrored(mean_2, d2) - ]) - self._assert_values_equal( - cross_tower_ops.batch_reduce(vs.VariableAggregation.SUM, - [(per_device, d1), (per_device_2, d2)]), - [ - _fake_mirrored(mean * len(devices), d1), - _fake_mirrored(mean_2 * len(devices), d2) - ]) - - # test broadcast() - for destinations in all_destinations: - self._assert_values_equal( - cross_tower_ops.broadcast(constant_op.constant(1.), destinations), - _fake_mirrored(1., destinations)) - - -class SingleWorkerCrossTowerOpsTest(CrossTowerOpsTestBase): - # TODO(yuefengz): decouple the num_gpus check from distribution in - # combinations module so that we can pass in devices instead of a distribution - # strategy. - reduction_to_one_combinations = combinations.combine( - cross_tower_ops=[ - combinations.NamedObject( - "DefaultReductionToOneDeviceCrossTowerOps", - cross_tower_ops_lib.ReductionToOneDeviceCrossTowerOps()), - combinations.NamedObject( - "ReductionToCPUDeviceCrossTowerOps", - cross_tower_ops_lib.ReductionToOneDeviceCrossTowerOps( - reduce_to_device=_cpu_device)), - combinations.NamedObject( - "AccumulateNCrossTowerOp", - cross_tower_ops_lib.ReductionToOneDeviceCrossTowerOps( - accumulation_fn=math_ops.accumulate_n)), - ], - distribution=[ - combinations.one_device_strategy, - combinations.mirrored_strategy_with_gpu_and_cpu, - combinations.mirrored_strategy_with_two_gpus - ], - mode=["graph", "eager"]) - allreduce_combinations = combinations.combine( - cross_tower_ops=[ - combinations.NamedObject( - "AllReduce", - cross_tower_ops_lib.AllReduceCrossTowerOps("nccl", 1, 0, 0)), - combinations.NamedObject( - "HierarchicalCopy", - cross_tower_ops_lib.AllReduceCrossTowerOps( - "hierarchical_copy", 8, 0, 0)), - combinations.NamedObject( - "AllReduceNoGradientRepacking", - cross_tower_ops_lib.AllReduceCrossTowerOps("nccl", 0, 0, 0)), - combinations.NamedObject( - "HierarchicalCopyAggregateSmallTensors", - cross_tower_ops_lib.AllReduceCrossTowerOps( - "hierarchical_copy", 0, 100, 10)) - ], - distribution=[combinations.mirrored_strategy_with_two_gpus], - mode=["graph", "eager"]) - - @combinations.generate(reduction_to_one_combinations + allreduce_combinations) - def testReductionAndBroadcast(self, cross_tower_ops, distribution): - with distribution.scope(): - self._testReductionAndBroadcast(cross_tower_ops, distribution) - - def testChooseAlgorithm(self): - device_links = [[1, 2, 3, 4], [0, 2, 3, 5], [0, 1, 3, 6], [0, 1, 2, 7], - [0, 5, 6, 7], [1, 4, 6, 7], [2, 4, 5, 7], [3, 4, 5, 6]] - result = cross_tower_ops_lib._choose_all_reduce_algorithm(device_links) - self.assertIsInstance(result, cross_tower_ops_lib.AllReduceCrossTowerOps) - self.assertEqual(result._all_reduce_alg, "hierarchical_copy") - self.assertEqual(result._num_packs, 8) - - # if there are only 4 devices - device_links = [[1, 2, 3, 4], [0, 2, 3, 5], [0, 1, 3, 6], [0, 1, 2, 7]] - result = cross_tower_ops_lib._choose_all_reduce_algorithm(device_links) - self.assertIsInstance(result, cross_tower_ops_lib.AllReduceCrossTowerOps) - self.assertEqual(result._all_reduce_alg, "nccl") - self.assertEqual(result._num_packs, 1) - - # if devices links contain each device itself - device_links = [[0, 1, 2, 3, 4], [0, 1, 2, 3, 5], [0, 1, 2, 3, 6], - [0, 1, 2, 3, 7], [0, 4, 5, 6, 7], [1, 4, 5, 6, 7], - [2, 4, 5, 6, 7], [3, 4, 5, 6, 7]] - result = cross_tower_ops_lib._choose_all_reduce_algorithm(device_links) - self.assertIsInstance(result, cross_tower_ops_lib.AllReduceCrossTowerOps) - self.assertEqual(result._all_reduce_alg, "hierarchical_copy") - self.assertEqual(result._num_packs, 8) - - # if not dgx1-like links - device_links = [[0, 2, 3, 5], [0, 1, 3, 6], [0, 1, 2, 7], [0, 5, 6, 7], - [1, 4, 6, 7], [2, 4, 5, 7], [3, 4, 5, 6], [1, 2, 3, 4]] - result = cross_tower_ops_lib._choose_all_reduce_algorithm(device_links) - self.assertIsInstance(result, cross_tower_ops_lib.AllReduceCrossTowerOps) - self.assertEqual(result._all_reduce_alg, "nccl") - self.assertEqual(result._num_packs, 1) - - @combinations.generate(combinations.combine( - mode=["graph", "eager"], - required_gpus=1)) - def testSimpleReduceWithIndexedSlices(self): - devices = ["/cpu:0", "/gpu:0"] - t0 = _make_indexed_slices([[1., 2.]], [1], [5, 2], devices[0]) - t1 = _make_indexed_slices([[3., 4.], [5., 6.]], [1, 3], [5, 2], devices[1]) - per_device = value_lib.PerDevice({devices[0]: t0, devices[1]: t1}) - result = cross_tower_ops_lib._simple_reduce( - per_device, devices[0], math_ops.add_n, vs.VariableAggregation.SUM) - - # Test that the result is semantically equal to both the concatenated - # IndexedSlices with and without duplicate indices. - total_with_dups = _make_indexed_slices( - [[1., 2.], [3., 4.], [5., 6.]], [1, 1, 3], [5, 2], devices[0]) - total_without_dups = _make_indexed_slices( - [[4., 6.], [5., 6.]], [1, 3], [5, 2], devices[0]) - self._assert_indexed_slices_equal(total_with_dups, result) - self._assert_indexed_slices_equal(total_without_dups, result) - - @combinations.generate( - combinations.combine( - cross_tower_ops_instance=[ - combinations.NamedObject( - "ReductionToOneDeviceCrossTowerOps", - cross_tower_ops_lib.ReductionToOneDeviceCrossTowerOps()), - combinations.NamedObject( - "AllReduceCrossTowerOps", - cross_tower_ops_lib.AllReduceCrossTowerOps()) - ], - aggregation=[vs.VariableAggregation.SUM, vs.VariableAggregation.MEAN], - batch_reduce=[True, False], - mode=["graph", "eager"], - required_gpus=1)) - def testIndexedSlicesAllReduce(self, cross_tower_ops_instance, aggregation, - batch_reduce): - devices = ["/cpu:0", "/gpu:0"] - dense_shape = [5, 2] - t0 = _make_indexed_slices([[1., 2.]], [1], dense_shape, devices[0]) - t1 = _make_indexed_slices( - [[3., 4.], [5., 6.]], [1, 3], dense_shape, devices[1]) - per_device = value_lib.PerDevice({devices[0]: t0, devices[1]: t1}) - - if batch_reduce: - result = cross_tower_ops_instance.batch_reduce(aggregation, - [(per_device, devices)]) - else: - result = cross_tower_ops_instance.reduce(aggregation, per_device, devices) - - total_indices_with_dups = [1, 1, 3] - total_indices_without_dups = [1, 3] - - if aggregation == vs.VariableAggregation.SUM: - total_values_with_dups = [[1., 2.], [3., 4.], [5., 6.]] - total_values_without_dups = [[4., 6.], [5., 6.]] - else: - assert aggregation == vs.VariableAggregation.MEAN - total_values_with_dups = [[0.5, 1.], [1.5, 2.], [2.5, 3.]] - total_values_without_dups = [[2., 3.], [2.5, 3.]] - - total_mirrored_with_dups = _make_mirrored_indexed_slices( - devices, total_values_with_dups, total_indices_with_dups, dense_shape) - total_mirrored_without_dups = _make_mirrored_indexed_slices( - devices, total_values_without_dups, total_indices_without_dups, - dense_shape) - - # Test that the result is semantically equal to both the concatenated - # IndexedSlices, as well as when the duplicate indices are summed up. - if batch_reduce: - total_mirrored_with_dups = [total_mirrored_with_dups] - total_mirrored_without_dups = [total_mirrored_without_dups] - - self._assert_values_equal(total_mirrored_with_dups, result) - self._assert_values_equal(total_mirrored_without_dups, result) - - -class MultiWorkerCrossTowerOpsTest(multi_worker_test_base.MultiWorkerTestBase, - CrossTowerOpsTestBase): - - worker_devices = [ - "/job:worker/replica:0/task:0", "/job:worker/replica:0/task:1" - ] - multi_worker_allreduce_combinations = combinations.combine( - cross_tower_ops=[ - combinations.NamedObject( - "MultiWorkerAllReduce", - cross_tower_ops_lib.MultiWorkerAllReduce( - worker_devices, 2, ("pscpu/pscpu", 2, -1), 0, 0, 0)), - combinations.NamedObject( - "MultiWorkerAllReducePack", - cross_tower_ops_lib.MultiWorkerAllReduce( - worker_devices, 2, ("pscpu/pscpu", 2, -1), 1, 0, 0)), - combinations.NamedObject( - "MultiWorkerAllReduceAggregation", - cross_tower_ops_lib.MultiWorkerAllReduce( - worker_devices, 2, ("pscpu/pscpu", 2, -1), 0, 100, 10)), - combinations.NamedObject( - "MultiWorkerAllReduceMultipleSpecs", - cross_tower_ops_lib.MultiWorkerAllReduce( - worker_devices, 2, [("pscpu/pscpu", 2, 100), - ("xring", 2, -1)], 0, 0, 0)), - ], - distribution=[ - combinations.NamedDistribution( - "MirroredCPU", - lambda: mirrored_strategy.MirroredStrategy(num_gpus=0), - required_gpus=0), - combinations.NamedDistribution( - "Mirrored1GPU", - lambda: mirrored_strategy.MirroredStrategy(num_gpus=1), - required_gpus=1), - combinations.NamedDistribution( - "Mirrored2GPUs", - lambda: mirrored_strategy.MirroredStrategy(num_gpus=2), - required_gpus=2), - ], - mode=["graph"]) - - @combinations.generate(multi_worker_allreduce_combinations) - def testReductionAndBroadcast(self, cross_tower_ops, distribution): - distribution.configure(cluster_spec={ - "worker": - ["/job:worker/replica:0/task:0", "/job:worker/replica:0/task:1"] - }) - with distribution.scope(): - self._testReductionAndBroadcast(cross_tower_ops, distribution) - - -class MultiWorkerCollectiveAllReduceTest( - multi_worker_test_base.MultiWorkerTestBase, parameterized.TestCase): - - collective_key_base = 100000 - - @classmethod - def setUpClass(cls): - """Create a local cluster with 2 workers.""" - cls._cluster_spec = multi_worker_test_base.create_in_process_cluster( - num_workers=3, num_ps=0) - - def setUp(self): - super(MultiWorkerCollectiveAllReduceTest, self).setUp() - # Reusing keys are not supported well. So we have to give a different - # collective key base for different tests. - MultiWorkerCollectiveAllReduceTest.collective_key_base += 100000 - - def _get_test_objects(self, task_type, task_id, num_gpus=0, local_mode=False): - collective_keys = cross_tower_utils.CollectiveKeys( - group_key_start=10 * num_gpus + - MultiWorkerCollectiveAllReduceTest.collective_key_base, - instance_key_start=num_gpus * 100 + - MultiWorkerCollectiveAllReduceTest.collective_key_base, - instance_key_with_id_start=num_gpus * 10000 + - MultiWorkerCollectiveAllReduceTest.collective_key_base) - if local_mode: - collective_all_reduce_ops = cross_tower_ops_lib.CollectiveAllReduce( - 1, num_gpus, collective_keys=collective_keys) - if num_gpus: - devices = ["/device:GPU:%d" % i for i in range(num_gpus)] - else: - devices = ["/device:CPU:0"] - return collective_all_reduce_ops, devices, "" - else: - collective_all_reduce_ops = cross_tower_ops_lib.CollectiveAllReduce( - 3, num_gpus, collective_keys=collective_keys) - if num_gpus: - devices = [ - "/job:%s/task:%d/device:GPU:%d" % (task_type, task_id, i) - for i in range(num_gpus) - ] - else: - devices = ["/job:%s/task:%d" % (task_type, task_id)] - return (collective_all_reduce_ops, devices, - "grpc://" + self._cluster_spec[task_type][task_id]) - - def _assert_values_equal(self, left, right, sess): - if isinstance(left, list): - for l, r in zip(left, right): - self._assert_values_equal(l, r, sess) - else: - self.assertEqual(type(left), type(right)) - self.assertEqual(set(left.devices), set(right.devices)) - - run_options = config_pb2.RunOptions() - run_options.experimental.collective_graph_key = 6 - - left_values = np.array( - sess.run(list(left._index.values()), options=run_options)).flatten() - right_values = np.array(list(right._index.values())).flatten() - self.assertEqual(len(left_values), len(right_values)) - for l, r in zip(left_values, right_values): - self.assertEqual(l, r) - - def _test_reduction(self, task_type, task_id, num_gpus, local_mode=False): - collective_all_reduce, devices, master_target = self._get_test_objects( - task_type, task_id, num_gpus, local_mode=local_mode) - if local_mode: - num_workers = 1 - worker_device = None - else: - num_workers = len(self._cluster_spec.get("chief", [])) + len( - self._cluster_spec.get("worker", [])) - worker_device = "/job:%s/task:%d" % (task_type, task_id) - with ops.Graph().as_default(), \ - ops.device(worker_device), \ - self.test_session(target=master_target) as sess: - # Collective ops doesn't support scalar tensors, so we have to construct - # 1-d tensors. - values = [constant_op.constant([float(d)]) for d in range(len(devices))] - per_device = _make_per_device(values, devices, regroup=True) - mean = np.array([(len(devices) - 1.) / 2.]) - - values_2 = [constant_op.constant([d + 1.0]) for d in range(len(devices))] - per_device_2 = _make_per_device(values_2, devices) - mean_2 = np.array([mean[0] + 1.]) - - destination_mirrored = _fake_mirrored(1., devices) - destination_different = _fake_mirrored(1., _cpu_device) - destination_str = _cpu_device - destination_list = devices - - all_destinations = [ - destination_different, destination_mirrored, destination_str, - destination_list - ] - - # test reduce() - for destinations in all_destinations: - self._assert_values_equal( - collective_all_reduce.reduce( - vs.VariableAggregation.MEAN, - per_device, - destinations=destinations), - _fake_mirrored(mean, destinations), sess) - self._assert_values_equal( - collective_all_reduce.reduce( - vs.VariableAggregation.MEAN, - per_device_2, - destinations=destinations), - _fake_mirrored(mean_2, destinations), sess) - self._assert_values_equal( - collective_all_reduce.reduce( - vs.VariableAggregation.SUM, - per_device, - destinations=destinations), - _fake_mirrored(mean * len(devices) * num_workers, destinations), - sess) - self._assert_values_equal( - collective_all_reduce.reduce( - vs.VariableAggregation.SUM, - per_device_2, - destinations=destinations), - _fake_mirrored(mean_2 * len(devices) * num_workers, destinations), - sess) - - # test batch_reduce() - for d1, d2 in itertools.product(all_destinations, all_destinations): - self._assert_values_equal( - collective_all_reduce.batch_reduce(vs.VariableAggregation.MEAN, - [(per_device, d1), - (per_device_2, d2)]), - [ - _fake_mirrored(mean, d1), - _fake_mirrored(mean_2, d2) - ], sess) - self._assert_values_equal( - collective_all_reduce.batch_reduce(vs.VariableAggregation.SUM, - [(per_device, d1), - (per_device_2, d2)]), - [ - _fake_mirrored(mean * len(devices) * num_workers, d1), - _fake_mirrored(mean_2 * len(devices) * num_workers, d2) - ], sess) - - return True - - @combinations.generate( - combinations.combine(mode=["graph"], num_gpus=[0, 1, 2], required_gpus=1)) - def testReductionDistributed(self, num_gpus): - if context.num_gpus() < num_gpus: - return - self._run_between_graph_clients(self._test_reduction, self._cluster_spec, - num_gpus) - - # Collective ops doesn't support strategy with one device. - def testReductionLocal(self, num_gpus=2): - if context.num_gpus() < num_gpus: - return - self._test_reduction(None, None, num_gpus, local_mode=True) - - -if __name__ == "__main__": - test.main() diff --git a/tensorflow/contrib/distribute/python/cross_tower_utils.py b/tensorflow/contrib/distribute/python/cross_tower_utils.py deleted file mode 100644 index 9fc1b8895516f64a956accd9290e7bf42ccef330..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/distribute/python/cross_tower_utils.py +++ /dev/null @@ -1,671 +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. -# ============================================================================== -"""Utilities for cross_tower_ops.""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import collections as pycoll -import threading - -from tensorflow.contrib import nccl -from tensorflow.contrib.all_reduce.python import all_reduce -from tensorflow.contrib.distribute.python import values as value_lib -from tensorflow.python.framework import device as pydev -from tensorflow.python.framework import dtypes -from tensorflow.python.framework import ops -from tensorflow.python.ops import array_ops -from tensorflow.python.ops import collective_ops -from tensorflow.python.ops import gradients_impl -from tensorflow.python.ops import math_ops - - -def aggregate_gradients_using_nccl(tower_grads): - """Aggregate gradients using nccl allreduce.""" - agg_all_g_and_v = [] - for single_g_and_v in zip(*tower_grads): - single_grads = [g for g, _ in single_g_and_v] - agg_grads = nccl.all_sum(single_grads) - agg_all_g_and_v.append( - [(g, v) for g, (_, v) in zip(agg_grads, single_g_and_v)]) - - agg_all_g_and_v = list(zip(*agg_all_g_and_v)) - - return agg_all_g_and_v - - -def aggregate_gradients_using_hierarchical_copy(avail_devices, tower_grads): - """Aggregate gradients using hierarchical copies. - - Args: - avail_devices: available GPU devices. - tower_grads: List of lists of (gradient, variable) tuples. The outer list - is over towers. The inner list is over individual gradients. - - Returns: - The list of (aggregated_gradient, variable), where the gradient has been - summed across all towers and the variable is chosen from the first tower. - """ - # This only works for DGX-1 type of machine topology - # Device peer to peer matrix - # DMA: 0 1 2 3 4 5 6 7 - # 0: Y Y Y Y Y N N N - # 1: Y Y Y Y N Y N N - # 2: Y Y Y Y N N Y N - # 3: Y Y Y Y N N N Y - # 4: Y N N N Y Y Y Y - # 5: N Y N N Y Y Y Y - # 6: N N Y N Y Y Y Y - # 7: N N N Y Y Y Y Y - agg_grads = [] - num_devices = len(avail_devices) - # In the special case of DGX-1 machine topology, the two groups have equal - # size. - group_size = num_devices // 2 - for i, single_grads in enumerate(zip(*tower_grads)): - group_0_main_device = i % num_devices - group_1_main_device = (group_0_main_device + group_size) % num_devices - if group_0_main_device < group_size: - group_0_begin = 0 - group_1_begin = group_size - else: - group_0_begin = group_size - group_1_begin = 0 - - # Aggregate the first group. - group_0_device_grads = single_grads[group_0_begin: - group_0_begin + group_size] - with ops.device(avail_devices[group_0_main_device]): - group_0_agg_grads, _ = aggregate_single_gradient_using_copy( - group_0_device_grads, False, False) - - # Aggregate the second group. - group_1_device_grads = single_grads[group_1_begin: - group_1_begin + group_size] - with ops.device(avail_devices[group_1_main_device]): - group_1_agg_grads, _ = aggregate_single_gradient_using_copy( - group_1_device_grads, False, False) - - # Aggregate between the groups. - with ops.device(avail_devices[group_0_main_device]): - (agg_total_grads, _), _ = aggregate_single_gradient_using_copy( - [group_0_agg_grads, group_1_agg_grads], False, False) - - # Broadcast the result back into the root of each group. - with ops.device(avail_devices[group_0_main_device]): - group_0_agg_grads_bcast = array_ops.identity(agg_total_grads) - with ops.device(avail_devices[group_1_main_device]): - group_1_agg_grads_bcast = array_ops.identity(agg_total_grads) - - agg_grads_bcast = [] - for j in range(len(single_grads)): - with ops.device(avail_devices[j]): - # Broadcast the result back to each member in the group from the root. - if (group_0_main_device < group_size) == (j < group_size): - src_device_grad = group_0_agg_grads_bcast - else: - src_device_grad = group_1_agg_grads_bcast - agg_grads_bcast.append(array_ops.identity(src_device_grad)) - - agg_grads.append( - [(g, v) for g, (_, v) in zip(agg_grads_bcast, single_grads)]) - - agg_grads = list(zip(*agg_grads)) - - return agg_grads - - -def aggregate_single_gradient_using_copy(grad_and_vars, use_mean, - check_inf_nan): - """Calculate the average gradient for a shared variable across all towers. - - Note that this function provides a synchronization point across all towers. - - Args: - grad_and_vars: A list or tuple of (gradient, variable) tuples. Each - (gradient, variable) pair within the outer list represents the gradient - of the variable calculated for a single tower, and the number of pairs - equals the number of towers. - use_mean: if True, mean is taken, else sum of gradients is taken. - check_inf_nan: check grads for nans and infs. - - Returns: - The tuple ([(average_gradient, variable),], has_nan_or_inf) where the - gradient has been averaged across all towers. The variable is chosen from - the first tower. The has_nan_or_inf indicates the grads has nan or inf. - """ - grads = [g for g, _ in grad_and_vars] - grad = math_ops.add_n(grads) - - if use_mean and len(grads) > 1: - grad = array_ops.multiply(grad, 1.0 / len(grads)) - - v = grad_and_vars[0][1] - if check_inf_nan: - has_nan_or_inf = array_ops.logical_not( - array_ops.reduce_all(array_ops.is_finite(grads))) - return (grad, v), has_nan_or_inf - else: - return (grad, v), None - - -def group_device_names(devices, group_size): - """Group device names into groups of group_size. - - Args: - devices: a list of canonical device strings. - group_size: integer which is equal to or greater than 1. - - Returns: - list of lists of devices, where each inner list is group_size long, - and each device appears at least once in an inner list. If - len(devices) % group_size == 0 then each device will appear exactly once. - - Raises: - ValueError: if group_size > len(devices) - """ - num_devices = len(devices) - if group_size > num_devices: - raise ValueError( - 'only %d devices, but group_size=%d' % (num_devices, group_size)) - num_groups = ( - num_devices // group_size + (1 if (num_devices % group_size != 0) else 0)) - groups = [[] for i in range(num_groups)] - for i in range(num_groups * group_size): - groups[i % num_groups].append(devices[i % num_devices]) - return groups - - -def split_grads_by_size(threshold_size, device_grads): - """Break gradients into two sets according to tensor size. - - Args: - threshold_size: int size cutoff for small vs large tensor. - device_grads: List of lists of (gradient, variable) tuples. The outer - list is over devices. The inner list is over individual gradients. - - Returns: - small_grads: Subset of device_grads where shape is <= threshold_size - elements. - large_grads: Subset of device_grads where shape is > threshold_size - elements. - """ - small_grads = [] - large_grads = [] - for dl in device_grads: - small_dl = [] - large_dl = [] - for (g, v) in dl: - tensor_size = g.get_shape().num_elements() - if tensor_size <= threshold_size: - small_dl.append([g, v]) - else: - large_dl.append([g, v]) - if small_dl: - small_grads.append(small_dl) - if large_dl: - large_grads.append(large_dl) - return small_grads, large_grads - - -# threading.Lock() and threading.local() cannot be pickled and therefore cannot -# be a field of CollectiveKeys. Right now _thread_local is not necessary to be -# an instance member of CollectiveKeys since we always create a new thread for -# each tower. -_lock = threading.Lock() -_thread_local = threading.local() - - -# TODO(yuefengz): use random key starts to avoid reusing keys? -class CollectiveKeys(object): - """Class that manages collective keys. - - We need to manage three different keys for collective: - - *Group key*: an integer key to identify the set of cooperative devices. - Collective ops work under the same set of devices must using the same group - key. - - *Instance key*: an integer key to identify the set of same counterpart of - tensors on different devices in a device group that need to be all-reduced. - - "Graph key": an integer key that is unique key graph. This is used to support - multiple graphs per client session. It must be non-zero and set in the - `config` argument of each call to `session.run`. - """ - - def __init__(self, - group_key_start=1, - instance_key_start=100, - instance_key_with_id_start=10000): - """Initializes the object. - - Args: - group_key_start: the starting integer of group key. - instance_key_start: the starting integer of instance key. - instance_key_with_id_start: the starting integer of instance key that is - recorded with an id. - """ - self._group_key = group_key_start - self._group_key_table = dict() - - # For instance keys with ids - self._instance_key_id_to_key_table = dict() - self._instance_key_with_id_counter = instance_key_with_id_start - - # For instance keys without ids - self._instance_key_start = instance_key_start - - def _get_thread_local_object(self): - # We make instance key without key ids thread local so that it will work - # with MirroredStrategy and distribute coordinator. - if not hasattr(_thread_local, 'instance_key'): - _thread_local.instance_key = self._instance_key_start - return _thread_local - - def get_group_key(self, devices): - """Returns a group key for the set of devices. - - Args: - devices: list of strings naming devices in a collective group. - - Returns: - int key uniquely identifying the set of device names. - """ - parsed = [pydev.DeviceSpec.from_string(d) for d in devices] - # In the between-graph replicated training, different workers need to get - # the same device key. So we remove the task_type and task_id from the - # devices. - # TODO(yuefengz): in the in-graph replicated training, we need to include - # task_type and task_id. - names = sorted(['%s:%d' % (d.device_type, d.device_index) for d in parsed]) - key_id = ','.join(names) - with _lock: - if key_id not in self._group_key_table: - new_key = self._group_key - self._group_key += 1 - self._group_key_table[key_id] = new_key - return self._group_key_table[key_id] - - def get_instance_key(self, key_id=None): - """Returns a new instance key for use in defining a collective op. - - Args: - key_id: optional string. If set, key will be recorded and the same key - will be returned when the same key_id is provided. If not, an increasing - instance key will be returned. - """ - if key_id: - with _lock: - if key_id not in self._instance_key_id_to_key_table: - self._instance_key_with_id_counter += 1 - self._instance_key_id_to_key_table[key_id] = ( - self._instance_key_with_id_counter) - return self._instance_key_id_to_key_table[key_id] - else: - v = self._get_thread_local_object().instance_key - self._get_thread_local_object().instance_key += 1 - return v - - -def build_collective_reduce(input_tensors, - num_workers, - collective_keys, - reduction_op='Add', - unary_op='Id'): - """Build a subgraph that does one full all-reduce, using the collective Op. - - Args: - input_tensors: tensors within a single worker graph that are to be reduced - together; must be one per device. - num_workers: total number of workers with identical independent graphs that - will be doing this same reduction. The reduction will actually include - the corresponding tensors at all these workers. - collective_keys: a CollectiveKeys object. - reduction_op: string naming the reduction op. - unary_op: string naming the unary final op. - - Returns: - An array of final tensors, one per device, computed by the full reduction. - - Raises: - ValueError: There must be at least two tensors over all the workers. - """ - group_size = len(input_tensors) * num_workers - if group_size < 2: - raise ValueError('num_workers * len(input_tensors) must be 2 or greater') - devices = [t.device for t in input_tensors] - num_devices = len(devices) - group_key = collective_keys.get_group_key(devices) - instance_key = collective_keys.get_instance_key() - out_tensors = [] - subdiv_offsets = [0] # TODO(tucker): maybe support non-default subdiv spec - for d in range(num_devices): - with ops.device(devices[d]): - reduce_op = collective_ops.all_reduce( - input_tensors[d], group_size, group_key, instance_key, reduction_op, - unary_op, subdiv_offsets) - out_tensors.append(reduce_op) - return out_tensors - - -def sum_grad_and_var_all_reduce(grad_and_vars, - num_workers, - alg, - gpu_indices, - aux_devices=None, - num_shards=1): - """Apply all-reduce algorithm over specified gradient tensors.""" - with ops.name_scope('allreduce'): - # Note that each grad_and_vars looks like the following: - # ((grad0_gpu0, var0_gpu0), ... , (grad0_gpuN, var0_gpuN)) - scaled_grads = [g for g, _ in grad_and_vars] - if alg == 'nccl': - summed_grads = nccl.all_sum(scaled_grads) - elif alg == 'xring': - summed_grads = all_reduce.build_ring_all_reduce( - scaled_grads, num_workers, num_shards, gpu_indices, math_ops.add) - elif alg == 'nccl/xring': - summed_grads = all_reduce.build_nccl_then_ring(scaled_grads, num_shards, - math_ops.add) - elif alg == 'nccl/rechd': - summed_grads = all_reduce.build_nccl_then_recursive_hd( - scaled_grads, math_ops.add) - elif alg == 'nccl/pscpu': - summed_grads = all_reduce.build_nccl_then_shuffle( - scaled_grads, aux_devices, math_ops.add, math_ops.add_n) - elif alg == 'pscpu/pscpu': - second_gather_devices = aux_devices[:num_shards] - summed_grads = all_reduce.build_shuffle_then_shuffle( - scaled_grads, aux_devices, second_gather_devices, math_ops.add_n) - elif alg in ['pscpu', 'psgpu']: - summed_grads = all_reduce.build_shuffle_all_reduce( - scaled_grads, aux_devices, math_ops.add_n) - else: - raise ValueError('unsupported all_reduce alg: ', alg) - - result = [] - for (_, v), g in zip(grad_and_vars, summed_grads): - result.append([g, v]) - return result - - -def sum_gradients_all_reduce(dev_prefixes, tower_grads, num_workers, alg, - num_shards, gpu_indices): - """Apply all-reduce algorithm over specified gradient tensors. - - Args: - dev_prefixes: list of prefix strings to use to generate PS device names. - tower_grads: the gradients to reduce. - num_workers: number of worker processes across entire job. - alg: the all-reduce algorithm to apply. - num_shards: alg-specific sharding factor. - gpu_indices: indices of local GPUs in order usable for ring-reduce. - - Returns: - list of reduced tensors - """ - alg_contains_shuffle = any([n in alg for n in ['pscpu', 'psgpu']]) - is_hierarchical = '/' in alg - if 'pscpu' in alg: - aux_devices = [prefix + '/cpu:0' for prefix in dev_prefixes] - elif 'psgpu' in alg: - aux_devices = [ - prefix + '/gpu:%d' % i - for i in range(len(gpu_indices)) - for prefix in dev_prefixes - ] - else: - aux_devices = ['/job:localhost/cpu:0'] - # Auxiliary devices for hierarchical all-reduces. - aux_device_groups = group_device_names( - aux_devices, num_shards if alg_contains_shuffle else 1) - group_index = 0 - reduced_gv_list = [] - for grad_and_vars in zip(*tower_grads): - reduced_gv_list.append( - sum_grad_and_var_all_reduce( - grad_and_vars, num_workers, alg, gpu_indices, aux_devices - if is_hierarchical else aux_device_groups[group_index], num_shards)) - group_index = (group_index + 1) % len(aux_device_groups) - new_tower_grads = [list(x) for x in zip(*reduced_gv_list)] - return new_tower_grads - - -def extract_ranges(index_list, range_size_limit=32): - """Extract consecutive ranges and singles from index_list. - - Args: - index_list: List of monotone increasing non-negative integers. - range_size_limit: Largest size range to return. If a larger - consecutive range exists, it will be returned as multiple - ranges. - - Returns: - (ranges, singles) where ranges is a list of [first, last] pairs of - consecutive elements in index_list, and singles is all of the - other elements, in original order. - """ - if not index_list: - return [], [] - first = index_list[0] - last = first - ranges = [] - singles = [] - for i in index_list[1:]: - if i == last + 1 and (last - first) <= range_size_limit: - last = i - else: - if last > first: - ranges.append([first, last]) - else: - singles.append(first) - first = i - last = i - if last > first: - ranges.append([first, last]) - else: - singles.append(first) - return ranges, singles - - -GradPackTuple = pycoll.namedtuple('GradPackTuple', 'indices vars shapes') - - -def pack_range(key, packing, grad_vars, rng): - """Form the concatenation of a specified range of gradient tensors. - - Args: - key: Value under which to store meta-data in packing that will be used - later to restore the grad_var list structure. - packing: Dict holding data describing packed ranges of small tensors. - grad_vars: List of (grad, var) pairs for one tower. - rng: A pair of integers giving the first, last indices of a consecutive - range of tensors to be packed. - - Returns: - A tensor that is the concatenation of all the specified small tensors. - """ - to_pack = grad_vars[rng[0]:rng[1] + 1] - members = [] - variables = [] - restore_shapes = [] - with ops.name_scope('pack'): - for g, v in to_pack: - variables.append(v) - restore_shapes.append(g.shape) - with ops.device(g.device): - members.append(array_ops.reshape(g, [-1])) - packing[key] = GradPackTuple( - indices=range(rng[0], rng[1] + 1), - vars=variables, - shapes=restore_shapes) - with ops.device(members[0].device): - return array_ops.concat(members, 0) - - -def unpack_grad_tuple(gv, gpt): - """Unpack a previously packed collection of gradient tensors. - - Args: - gv: A (grad, var) pair to be unpacked. - gpt: A GradPackTuple describing the packing operation that produced gv. - - Returns: - A list of (grad, var) pairs corresponding to the values that were - originally packed into gv, maybe following subsequent operations like - reduction. - """ - elt_widths = [x.num_elements() for x in gpt.shapes] - with ops.device(gv[0][0].device): - with ops.name_scope('unpack'): - splits = array_ops.split(gv[0], elt_widths) - unpacked_gv = [] - for idx, s in enumerate(splits): - unpacked_gv.append((array_ops.reshape(s, gpt.shapes[idx]), - gpt.vars[idx])) - return unpacked_gv - - -def pack_small_tensors(tower_grads, max_bytes=0, max_group=0): - """Concatenate small gradient tensors together for reduction. - - Args: - tower_grads: List of lists of (gradient, variable) tuples. - max_bytes: Int giving max number of bytes in a tensor that - may be considered small. - max_group: Int giving max number of small tensors that may be - concatenated into one new tensor. - - Returns: - new_tower_grads, packing where new_tower_grads is identical to - tower_grads except that all feasible small_tensors have been removed - from their places and concatenated into larger tensors that are - now in the front of the list for each tower, and packing contains - the data necessary to restore the tower_grads structure. - - Look through the first tower for gradients of the same type (float), - and small size, that are all sequential. For each such group, - replace by a new tensor that is a flattened concatenation. Note - that the corresponding variable will be absent, which doesn't matter - because it isn't used during all-reduce. - - Requires: - Every gv_list in towers must have isomorphic structure including identical - tensor sizes and types. - """ - small_indices = [] - large_indices = [] - for idx, (g, _) in enumerate(tower_grads[0]): - if g.dtype == dtypes.float32 and (4 * g.shape.num_elements()) <= max_bytes: - small_indices.append(idx) - else: - large_indices.append(idx) - small_ranges, small_singles = extract_ranges( - small_indices, range_size_limit=max_group) - large_indices = sorted(large_indices + small_singles) - num_gv = len(tower_grads[0]) - packing = {} - if small_ranges: - new_tower_grads = [] - for dev_idx, gv_list in enumerate(tower_grads): - assert len(gv_list) == num_gv - new_gv_list = [] - for r in small_ranges: - key = '%d:%d' % (dev_idx, len(new_gv_list)) - new_gv_list.append((pack_range(key, packing, gv_list, r), - 'packing_var_placeholder')) - for i in large_indices: - new_gv_list.append(gv_list[i]) - new_tower_grads.append(new_gv_list) - return new_tower_grads, packing - else: - return tower_grads, None - - -def unpack_small_tensors(tower_grads, packing): - """Undo the structure alterations to tower_grads done by pack_small_tensors. - - Args: - tower_grads: List of List of (grad, var) tuples. - packing: A dict generated by pack_small_tensors describing the changes - it made to tower_grads. - - Returns: - new_tower_grads: identical to tower_grads except that concatenations - of small tensors have been split apart and returned to their original - positions, paired with their original variables. - """ - if not packing: - return tower_grads - new_tower_grads = [] - num_devices = len(tower_grads) - num_packed = len(packing.keys()) // num_devices - for dev_idx, gv_list in enumerate(tower_grads): - gv_list = list(gv_list) - new_gv_list = gv_list[num_packed:] - for i in range(num_packed): - k = '%d:%d' % (dev_idx, i) - gpt = packing[k] - gv = unpack_grad_tuple(gv_list[i], gpt) - for gi, idx in enumerate(gpt.indices): - assert idx == gpt.indices[gi] - new_gv_list.insert(idx, gv[gi]) - new_tower_grads.append(new_gv_list) - return new_tower_grads - - -def aggregate_tensors_or_indexed_slices(values, accumulation_fn=math_ops.add_n): - """Aggregate tensors using `accumulation_fn` and IndexedSlices via concat.""" - if any(isinstance(v, ops.IndexedSlices) for v in values): - return gradients_impl._AggregateIndexedSlicesGradients(values) # pylint: disable=protected-access - else: - return accumulation_fn(values) - - -def divide_by_n_tensors_or_indexed_slices(value, n): - if isinstance(value, ops.IndexedSlices): - value = gradients_impl._HandleNestedIndexedSlices(value) # pylint: disable=protected-access - return ops.IndexedSlices( - value.values / n, value.indices, value.dense_shape) - else: - return value / n - - -def copy_tensor_or_indexed_slices_to_device(value, device): - with ops.device(device): - if isinstance(value, ops.IndexedSlices): - copied_values = array_ops.identity(value.values) - copied_indices = array_ops.identity(value.indices) - copied_shape = array_ops.identity(value.dense_shape) - result = ops.IndexedSlices(copied_values, copied_indices, copied_shape) - else: - result = array_ops.identity(value) - return result - - -def contains_indexed_slices(value): - """Check whether the value is `IndexedSlices` or contains `IndexedSlices`.""" - if isinstance(value, ops.IndexedSlices): - return True - elif isinstance(value, (list, tuple)) and value: - return any(contains_indexed_slices(v) for v in value) - elif isinstance(value, value_lib.DistributedValues): - return contains_indexed_slices(list(value._index.values())) # pylint: disable=protected-access - elif isinstance(value, value_lib.MapOutput): - return contains_indexed_slices(value.get()) - else: - return False diff --git a/tensorflow/contrib/distribute/python/cross_tower_utils_test.py b/tensorflow/contrib/distribute/python/cross_tower_utils_test.py deleted file mode 100644 index d25964fa41adc7b1c9164a4ffe49c4c5532f76ac..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/distribute/python/cross_tower_utils_test.py +++ /dev/null @@ -1,152 +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. -# ============================================================================== -"""Tests for cross_tower_utils.""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -from absl.testing import parameterized - -from tensorflow.contrib.distribute.python import combinations -from tensorflow.contrib.distribute.python import cross_tower_utils -from tensorflow.contrib.distribute.python import values as value_lib -from tensorflow.python.eager import test -from tensorflow.python.framework import constant_op -from tensorflow.python.framework import ops -from tensorflow.python.framework import test_util -from tensorflow.python.ops import math_ops -from tensorflow.python.training import device_util - - -class IndexedSlicesUtilsTest(test.TestCase, parameterized.TestCase): - - def _assert_values_equal(self, left, right): - self.assertAllEqual( - self.evaluate(ops.convert_to_tensor(left)), - self.evaluate(ops.convert_to_tensor(right))) - - @test_util.run_in_graph_and_eager_modes - def testAggregateTensors(self): - t0 = constant_op.constant([[1., 2.], [0, 0], [3., 4.]]) - t1 = constant_op.constant([[0., 0.], [5, 6], [7., 8.]]) - total = constant_op.constant([[1., 2.], [5, 6], [10., 12.]]) - result = cross_tower_utils.aggregate_tensors_or_indexed_slices([t0, t1]) - self._assert_values_equal(total, result) - - @test_util.run_in_graph_and_eager_modes - def testAggregateIndexedSlices(self): - t0 = math_ops._as_indexed_slices( - constant_op.constant([[1., 2.], [0, 0], [3., 4.]])) - t1 = math_ops._as_indexed_slices( - constant_op.constant([[0., 0.], [5, 6], [7., 8.]])) - total = constant_op.constant([[1., 2.], [5, 6], [10., 12.]]) - result = cross_tower_utils.aggregate_tensors_or_indexed_slices([t0, t1]) - self.assertIsInstance(result, ops.IndexedSlices) - self._assert_values_equal(total, result) - - @test_util.run_in_graph_and_eager_modes - def testDivideTensor(self): - t = constant_op.constant([[1., 2.], [0, 0], [3., 4.]]) - n = 2 - expected = constant_op.constant([[0.5, 1.], [0, 0], [1.5, 2.]]) - result = cross_tower_utils.divide_by_n_tensors_or_indexed_slices(t, n) - self._assert_values_equal(expected, result) - - @test_util.run_in_graph_and_eager_modes - def testDivideIndexedSlices(self): - t = math_ops._as_indexed_slices( - constant_op.constant([[1., 2.], [0, 0], [3., 4.]])) - n = 2 - expected = constant_op.constant([[0.5, 1.], [0, 0], [1.5, 2.]]) - result = cross_tower_utils.divide_by_n_tensors_or_indexed_slices(t, n) - self.assertIsInstance(result, ops.IndexedSlices) - self._assert_values_equal(expected, result) - - @test_util.run_in_graph_and_eager_modes - def testIsIndexedSlices(self): - t = math_ops._as_indexed_slices( - constant_op.constant([[1., 2.], [0, 0], [3., 4.]])) - self.assertTrue(cross_tower_utils.contains_indexed_slices(t)) - - @test_util.run_in_graph_and_eager_modes - def testContainsIndexedSlices_List(self): - t0 = math_ops._as_indexed_slices( - constant_op.constant([[1., 2.], [0, 0], [3., 4.]])) - t1 = math_ops._as_indexed_slices( - constant_op.constant([[0., 0.], [5, 6], [7., 8.]])) - self.assertTrue(cross_tower_utils.contains_indexed_slices([t0, t1])) - - @test_util.run_in_graph_and_eager_modes - def testContainsIndexedSlices_Tuple(self): - t0 = math_ops._as_indexed_slices( - constant_op.constant([[1., 2.], [0, 0], [3., 4.]])) - t1 = math_ops._as_indexed_slices( - constant_op.constant([[0., 0.], [5, 6], [7., 8.]])) - self.assertTrue(cross_tower_utils.contains_indexed_slices((t0, t1))) - - @test_util.run_in_graph_and_eager_modes - def testContainsIndexedSlices_PerDevice(self): - t0 = math_ops._as_indexed_slices( - constant_op.constant([[1., 2.], [0, 0], [3., 4.]])) - t1 = math_ops._as_indexed_slices( - constant_op.constant([[0., 0.], [5, 6], [7., 8.]])) - per_device = value_lib.PerDevice({"/gpu:0": t0, "/cpu:0": t1}) - self.assertTrue(cross_tower_utils.contains_indexed_slices(per_device)) - - @test_util.run_in_graph_and_eager_modes - def testContainsIndexedSlices_PerDeviceMapOutput(self): - t0 = math_ops._as_indexed_slices( - constant_op.constant([[1., 2.], [0, 0], [3., 4.]])) - t1 = math_ops._as_indexed_slices( - constant_op.constant([[0., 0.], [5, 6], [7., 8.]])) - per_device = value_lib.PerDevice({ - "/gpu:0": value_lib.MapOutput([t0]), - "/cpu:0": value_lib.MapOutput([t1])}) - self.assertTrue(cross_tower_utils.contains_indexed_slices(per_device)) - - @combinations.generate(combinations.combine( - mode=["graph", "eager"], - required_gpus=1)) - def testCopyTensor(self): - with ops.device("/cpu:0"): - t = constant_op.constant([[1., 2.], [0, 0], [3., 4.]]) - destination = "/gpu:0" - result = cross_tower_utils.copy_tensor_or_indexed_slices_to_device( - t, destination) - - self._assert_values_equal(t, result) - self.assertEqual(device_util.resolve(destination), - device_util.resolve(result.device)) - - @combinations.generate(combinations.combine( - mode=["graph", "eager"], - required_gpus=1)) - def testCopyIndexedSlices(self): - with ops.device("/cpu:0"): - t = math_ops._as_indexed_slices( - constant_op.constant([[1., 2.], [0, 0], [3., 4.]])) - destination = "/gpu:0" - result = cross_tower_utils.copy_tensor_or_indexed_slices_to_device( - t, destination) - - self.assertIsInstance(result, ops.IndexedSlices) - self._assert_values_equal(t, result) - self.assertEqual(device_util.resolve(destination), - device_util.resolve(result.device)) - - -if __name__ == "__main__": - test.main() diff --git a/tensorflow/contrib/distribute/python/estimator_integration_test.py b/tensorflow/contrib/distribute/python/estimator_integration_test.py index cc626c33bf8e282736f8e6e0c151e5a3d3f3244b..e17085628ba6d1dfc79839fd824801723f07a518 100644 --- a/tensorflow/contrib/distribute/python/estimator_integration_test.py +++ b/tensorflow/contrib/distribute/python/estimator_integration_test.py @@ -34,7 +34,7 @@ from tensorflow.python.estimator.canned import dnn_linear_combined from tensorflow.python.estimator.canned import prediction_keys from tensorflow.python.estimator.export import export from tensorflow.python.estimator.inputs import numpy_io -from tensorflow.python.feature_column import feature_column +from tensorflow.python.feature_column import feature_column_lib as feature_column from tensorflow.python.framework import ops from tensorflow.python.platform import gfile from tensorflow.python.summary.writer import writer_cache @@ -63,7 +63,9 @@ class DNNLinearCombinedClassifierIntegrationTest(test.TestCase, distribution=[ combinations.one_device_strategy, combinations.mirrored_strategy_with_gpu_and_cpu, - combinations.mirrored_strategy_with_two_gpus + combinations.mirrored_strategy_with_two_gpus, + combinations.core_mirrored_strategy_with_gpu_and_cpu, + combinations.core_mirrored_strategy_with_two_gpus ], use_train_and_evaluate=[True, False])) def test_complete_flow_with_mode(self, distribution, use_train_and_evaluate): @@ -75,12 +77,12 @@ class DNNLinearCombinedClassifierIntegrationTest(test.TestCase, train_input_fn = self.dataset_input_fn( x={'x': data}, y=data, - batch_size=batch_size // len(distribution.worker_devices), + batch_size=batch_size // distribution.num_replicas_in_sync, shuffle=True) eval_input_fn = self.dataset_input_fn( x={'x': data}, y=data, - batch_size=batch_size // len(distribution.worker_devices), + batch_size=batch_size // distribution.num_replicas_in_sync, shuffle=False) predict_input_fn = numpy_io.numpy_input_fn( x={'x': data}, batch_size=batch_size, shuffle=False) @@ -126,8 +128,8 @@ class DNNLinearCombinedClassifierIntegrationTest(test.TestCase, feature_spec = feature_column.make_parse_example_spec(feature_columns) serving_input_receiver_fn = export.build_parsing_serving_input_receiver_fn( feature_spec) - export_dir = estimator.export_savedmodel(tempfile.mkdtemp(), - serving_input_receiver_fn) + export_dir = estimator.export_saved_model(tempfile.mkdtemp(), + serving_input_receiver_fn) self.assertTrue(gfile.Exists(export_dir)) def tearDown(self): diff --git a/tensorflow/contrib/distribute/python/estimator_training_test.py b/tensorflow/contrib/distribute/python/estimator_training_test.py index 157618f72ff2ea6dde171e7edb62ccaf7e1de516..3f55a8a1c8b88d1b8e4031547fa3fbe519983630 100644 --- a/tensorflow/contrib/distribute/python/estimator_training_test.py +++ b/tensorflow/contrib/distribute/python/estimator_training_test.py @@ -18,15 +18,16 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import copy import glob import json import os import sys import tempfile -import threading from absl.testing import parameterized import numpy as np +from tensorflow.contrib.distribute.python import collective_all_reduce_strategy from tensorflow.contrib.distribute.python import combinations from tensorflow.contrib.distribute.python import mirrored_strategy from tensorflow.contrib.distribute.python import multi_worker_test_base @@ -43,11 +44,13 @@ from tensorflow.python.estimator import training as estimator_training from tensorflow.python.estimator.canned import dnn_linear_combined from tensorflow.python.estimator.canned import prediction_keys from tensorflow.python.estimator.export import export as export_lib -from tensorflow.python.feature_column import feature_column +from tensorflow.python.feature_column import feature_column_lib as feature_column from tensorflow.python.platform import gfile from tensorflow.python.platform import test from tensorflow.python.summary import summary_iterator from tensorflow.python.summary.writer import writer_cache +from tensorflow.python.training import session_manager + BATCH_SIZE = 10 LABEL_DIMENSION = 2 @@ -66,57 +69,19 @@ PS = dc._TaskType.PS original_run_std_server = dc._run_std_server -class MockOsEnv(dict): - - def __init__(self, *args): - self._thread_local = threading.local() - super(MockOsEnv, self).__init__(*args) - - def get(self, key, default): - if not hasattr(self._thread_local, "dict"): - self._thread_local.dict = dict() - if key == "TF_CONFIG": - return dict.get(self._thread_local.dict, key, default) - else: - return dict.get(self, key, default) - - def __getitem__(self, key): - if not hasattr(self._thread_local, "dict"): - self._thread_local.dict = dict() - if key == "TF_CONFIG": - return dict.__getitem__(self._thread_local.dict, key) - else: - return dict.__getitem__(self, key) - - def __setitem__(self, key, val): - if not hasattr(self._thread_local, "dict"): - self._thread_local.dict = dict() - if key == "TF_CONFIG": - return dict.__setitem__(self._thread_local.dict, key, val) - else: - return dict.__setitem__(self, key, val) - - -class DistributeCoordinatorIntegrationTest(test.TestCase, - parameterized.TestCase): +class DistributeCoordinatorIntegrationTest( + multi_worker_test_base.IndependentWorkerTestBase, parameterized.TestCase): @classmethod def setUpClass(cls): """Create a local cluster with 2 workers.""" + super(DistributeCoordinatorIntegrationTest, cls).setUpClass() cls._cluster_spec = multi_worker_test_base.create_in_process_cluster( num_workers=3, num_ps=2, has_eval=True) def setUp(self): self._model_dir = tempfile.mkdtemp() - self._mock_os_env = MockOsEnv() - self._mock_context = test.mock.patch.object(os, "environ", - self._mock_os_env) super(DistributeCoordinatorIntegrationTest, self).setUp() - self._mock_context.__enter__() - - def tearDown(self): - self._mock_context.__exit__(None, None, None) - super(DistributeCoordinatorIntegrationTest, self).tearDown() def dataset_input_fn(self, x, y, batch_size, shuffle): @@ -139,6 +104,8 @@ class DistributeCoordinatorIntegrationTest(test.TestCase, def _extract_loss_and_global_step(self, event_folder): """Returns the loss and global step in last event.""" event_paths = glob.glob(os.path.join(event_folder, "events*")) + self.assertNotEmpty( + event_paths, msg="Event file not found in dir %s" % event_folder) loss = None global_step_count = None @@ -189,7 +156,8 @@ class DistributeCoordinatorIntegrationTest(test.TestCase, def _complete_flow(self, train_distribute, eval_distribute, - remote_cluster=None): + remote_cluster=None, + use_train_and_evaluate=True): estimator = self._get_estimator(train_distribute, eval_distribute, remote_cluster) @@ -197,10 +165,10 @@ class DistributeCoordinatorIntegrationTest(test.TestCase, train_input_fn = self.dataset_input_fn( x={"x": DATA}, y=DATA, - batch_size=BATCH_SIZE // len(train_distribute.worker_devices), + batch_size=BATCH_SIZE // train_distribute.num_replicas_in_sync, shuffle=True) if eval_distribute: - eval_batch_size = BATCH_SIZE // len(eval_distribute.worker_devices) + eval_batch_size = BATCH_SIZE // eval_distribute.num_replicas_in_sync else: eval_batch_size = BATCH_SIZE eval_input_fn = self.dataset_input_fn( @@ -214,16 +182,37 @@ class DistributeCoordinatorIntegrationTest(test.TestCase, ] feature_columns = linear_feature_columns + dnn_feature_columns - estimator_training.train_and_evaluate( - estimator, - estimator_training.TrainSpec(train_input_fn, max_steps=MAX_STEPS), - estimator_training.EvalSpec( - name=EVAL_NAME, - input_fn=eval_input_fn, - steps=None, - exporters=self._get_exporter(EXPORTER_NAME, feature_columns), - start_delay_secs=0, - throttle_secs=1)) + eval_spec = estimator_training.EvalSpec( + name=EVAL_NAME, + input_fn=eval_input_fn, + steps=None, + exporters=self._get_exporter(EXPORTER_NAME, feature_columns), + start_delay_secs=0, + throttle_secs=1) + + if use_train_and_evaluate: + estimator_training.train_and_evaluate( + estimator, + estimator_training.TrainSpec(train_input_fn, max_steps=MAX_STEPS), + eval_spec) + else: + estimator.train(train_input_fn, max_steps=MAX_STEPS) + + latest_ckpt_path = estimator.latest_checkpoint() + metrics = estimator.evaluate(eval_input_fn, + checkpoint_path=latest_ckpt_path, + name=EVAL_NAME) + + # Export the eval result to files. + eval_result = estimator_training._EvalResult( + status=estimator_training._EvalStatus.EVALUATED, + metrics=metrics, + checkpoint_path=latest_ckpt_path) + evaluator = estimator_training._TrainingExecutor._Evaluator(estimator, + eval_spec, + None) + evaluator._export_eval_result(eval_result, True) + return estimator def _inspect_train_and_eval_events(self, estimator): @@ -259,32 +248,74 @@ class DistributeCoordinatorIntegrationTest(test.TestCase, ]) self.assertAllEqual((BATCH_SIZE, LABEL_DIMENSION), predicted_proba.shape) + def _get_strategy_object(self, strategy_cls): + if strategy_cls == mirrored_strategy.CoreMirroredStrategy: + return strategy_cls(mirrored_strategy.all_local_devices()) + else: + return strategy_cls(num_gpus_per_worker=context.num_gpus()) + @combinations.generate( combinations.combine( mode=["graph"], train_distribute_cls=[ + collective_all_reduce_strategy.CollectiveAllReduceStrategy, mirrored_strategy.MirroredStrategy, + mirrored_strategy.CoreMirroredStrategy, parameter_server_strategy.ParameterServerStrategy ], eval_distribute_cls=[ - None, mirrored_strategy.MirroredStrategy, - parameter_server_strategy.ParameterServerStrategy + None, + mirrored_strategy.MirroredStrategy, + mirrored_strategy.CoreMirroredStrategy, + parameter_server_strategy.ParameterServerStrategy, ], - required_gpus=1)) + required_gpus=[0, 1])) def test_complete_flow_standalone_client(self, train_distribute_cls, eval_distribute_cls): - try: - train_distribute = train_distribute_cls(num_gpus=context.num_gpus()) - except TypeError: - train_distribute = train_distribute_cls(num_gpus_per_worker=2) + train_distribute = self._get_strategy_object(train_distribute_cls) if eval_distribute_cls: - eval_distribute = eval_distribute_cls() + eval_distribute = self._get_strategy_object(eval_distribute_cls) else: eval_distribute = None + cluster_spec = copy.deepcopy(self._cluster_spec) + if (train_distribute_cls != + parameter_server_strategy.ParameterServerStrategy): + cluster_spec.pop("ps", None) + estimator = self._complete_flow(train_distribute, eval_distribute, + cluster_spec) + self._inspect_train_and_eval_events(estimator) + + @combinations.generate( + combinations.combine( + mode=["graph"], + train_distribute_cls=[ + mirrored_strategy.MirroredStrategy, + mirrored_strategy.CoreMirroredStrategy, + ], + eval_distribute_cls=[ + None, + mirrored_strategy.MirroredStrategy, + mirrored_strategy.CoreMirroredStrategy, + ], + required_gpus=[0, 1])) + def test_estimator_standalone_client(self, train_distribute_cls, + eval_distribute_cls): + train_distribute = self._get_strategy_object(train_distribute_cls) + + if eval_distribute_cls: + eval_distribute = self._get_strategy_object(eval_distribute_cls) + else: + eval_distribute = None + + # We use the whole cluster for evaluation. + cluster = copy.deepcopy(self._cluster_spec) + cluster.pop("evaluator", None) + estimator = self._complete_flow( - train_distribute, eval_distribute, remote_cluster=self._cluster_spec) + train_distribute, eval_distribute, remote_cluster=cluster, + use_train_and_evaluate=False) self._inspect_train_and_eval_events(estimator) def _mock_run_std_server(self, *args, **kwargs): @@ -294,80 +325,63 @@ class DistributeCoordinatorIntegrationTest(test.TestCase, self._barrier.wait() return ret - def _task_thread(self, train_distribute, eval_distribute, tf_config): - os.environ["TF_CONFIG"] = json.dumps(tf_config) + def _independent_worker_fn( + self, + train_distribute, + eval_distribute, + ): with test.mock.patch.object(dc, "_run_std_server", self._mock_run_std_server): self._complete_flow(train_distribute, eval_distribute) - def _run_task_in_thread(self, cluster_spec, task_type, task_id, - train_distribute, eval_distribute): - if task_type: - tf_config = { - "cluster": cluster_spec, - "task": { - "type": task_type, - "index": task_id - } - } - else: - tf_config = { - "cluster": cluster_spec, - "task": { - "type": task_type, - "index": task_id - } - } - t = threading.Thread( - target=self._task_thread, - args=(train_distribute, eval_distribute, tf_config)) - t.start() - return t - - def _run_multiple_tasks_in_threads(self, cluster_spec, train_distribute, - eval_distribute): - threads = {} - for task_type in cluster_spec.keys(): - threads[task_type] = [] - for task_id in range(len(cluster_spec[task_type])): - t = self._run_task_in_thread(cluster_spec, task_type, task_id, - train_distribute, eval_distribute) - threads[task_type].append(t) - return threads - @combinations.generate( combinations.combine( mode=["graph"], train_distribute_cls=[ + collective_all_reduce_strategy.CollectiveAllReduceStrategy, parameter_server_strategy.ParameterServerStrategy, ], eval_distribute_cls=[ None, mirrored_strategy.MirroredStrategy, - parameter_server_strategy.ParameterServerStrategy + mirrored_strategy.CoreMirroredStrategy, + parameter_server_strategy.ParameterServerStrategy, ], - required_gpus=1)) + required_gpus=[0, 1])) def test_complete_flow_indepedent_worker_between_graph( self, train_distribute_cls, eval_distribute_cls): - train_distribute = train_distribute_cls( - num_gpus_per_worker=context.num_gpus()) + if (context.num_gpus() < 2 and eval_distribute_cls == + collective_all_reduce_strategy.CollectiveAllReduceStrategy): + self.skipTest("`CollectiveAllReduceStrategy` needs at least two towers.") + + train_distribute = self._get_strategy_object(train_distribute_cls) if eval_distribute_cls: - eval_distribute = eval_distribute_cls() + eval_distribute = self._get_strategy_object(eval_distribute_cls) else: eval_distribute = None - cluster_spec = multi_worker_test_base.create_cluster_spec( - num_workers=3, num_ps=2, has_eval=True) - # 3 workers, 2 ps and 1 evaluator. - self._barrier = dc._Barrier(6) - - threads = self._run_multiple_tasks_in_threads( - cluster_spec, train_distribute, eval_distribute) + if (train_distribute_cls == parameter_server_strategy + .ParameterServerStrategy): + cluster_spec = multi_worker_test_base.create_cluster_spec( + num_workers=3, num_ps=2, has_eval=True) + # 3 workers, 2 ps and 1 evaluator. + self._barrier = dc._Barrier(6) + else: + cluster_spec = multi_worker_test_base.create_cluster_spec( + num_workers=3, num_ps=0, has_eval=True) + # 3 workers and 1 evaluator. + self._barrier = dc._Barrier(4) + + threads = self.run_multiple_tasks_in_threads(self._independent_worker_fn, + cluster_spec, train_distribute, + eval_distribute) + threads_to_join = [] for task_type, ts in threads.items(): if task_type == PS: continue for t in ts: - t.join() + threads_to_join.append(t) + self.join_independent_workers(threads_to_join) estimator = self._get_estimator(train_distribute, eval_distribute) self._inspect_train_and_eval_events(estimator) @@ -375,15 +389,22 @@ class DistributeCoordinatorIntegrationTest(test.TestCase, @combinations.generate( combinations.combine( mode=["graph"], - train_distribute_cls=[mirrored_strategy.MirroredStrategy], - eval_distribute_cls=[None, mirrored_strategy.MirroredStrategy], - required_gpus=1)) + train_distribute_cls=[ + mirrored_strategy.MirroredStrategy, + mirrored_strategy.CoreMirroredStrategy + ], + eval_distribute_cls=[ + None, + mirrored_strategy.MirroredStrategy, + mirrored_strategy.CoreMirroredStrategy + ], + required_gpus=[0, 1])) def test_complete_flow_indepedent_worker_in_graph(self, train_distribute_cls, eval_distribute_cls): - train_distribute = train_distribute_cls(num_gpus=context.num_gpus()) + train_distribute = self._get_strategy_object(train_distribute_cls) if eval_distribute_cls: - eval_distribute = eval_distribute_cls() + eval_distribute = self._get_strategy_object(eval_distribute_cls) else: eval_distribute = None @@ -391,10 +412,10 @@ class DistributeCoordinatorIntegrationTest(test.TestCase, num_workers=3, num_ps=0, has_eval=True) # 3 workers and 1 evaluator. self._barrier = dc._Barrier(4) - threads = self._run_multiple_tasks_in_threads( - cluster_spec, train_distribute, eval_distribute) - threads[WORKER][0].join() - threads[EVALUATOR][0].join() + threads = self.run_multiple_tasks_in_threads(self._independent_worker_fn, + cluster_spec, train_distribute, + eval_distribute) + self.join_independent_workers([threads[WORKER][0], threads[EVALUATOR][0]]) estimator = self._get_estimator(train_distribute, eval_distribute) self._inspect_train_and_eval_events(estimator) @@ -430,7 +451,8 @@ class RunConfigTest(test.TestCase): "os.environ", {"TF_CONFIG": json.dumps(TF_CONFIG_WITHOUT_TASK)}): run_config_lib.RunConfig( experimental_distribute=DistributeConfig( - train_distribute=mirrored_strategy.MirroredStrategy(num_gpus=2))) + train_distribute=mirrored_strategy.CoreMirroredStrategy( + ["/device:GPU:0", "/device:GPU:1"]))) def test_should_run_distribute_coordinator(self): """Tests that should_run_distribute_coordinator return a correct value.""" @@ -453,10 +475,12 @@ class RunConfigTest(test.TestCase): {"TF_CONFIG": json.dumps(TF_CONFIG_WITH_CHIEF)}): config_with_train_distribute = run_config_lib.RunConfig( experimental_distribute=DistributeConfig( - train_distribute=mirrored_strategy.MirroredStrategy(num_gpus=2))) + train_distribute=mirrored_strategy.CoreMirroredStrategy( + ["/device:GPU:0", "/device:GPU:1"]))) config_with_eval_distribute = run_config_lib.RunConfig( experimental_distribute=DistributeConfig( - eval_distribute=mirrored_strategy.MirroredStrategy(num_gpus=2))) + eval_distribute=mirrored_strategy.CoreMirroredStrategy( + ["/device:GPU:0", "/device:GPU:1"]))) self.assertTrue( dc_training.should_run_distribute_coordinator( config_with_train_distribute)) @@ -469,26 +493,27 @@ class RunConfigTest(test.TestCase): {"TF_CONFIG": json.dumps(TF_CONFIG_WITH_MASTER)}): config = run_config_lib.RunConfig( experimental_distribute=DistributeConfig( - train_distribute=mirrored_strategy.MirroredStrategy(num_gpus=2))) + train_distribute=mirrored_strategy.CoreMirroredStrategy( + ["/device:GPU:0", "/device:GPU:1"]))) self.assertFalse(dc_training.should_run_distribute_coordinator(config)) def test_init_run_config_duplicate_distribute(self): with self.assertRaises(ValueError): run_config_lib.RunConfig( - train_distribute=mirrored_strategy.MirroredStrategy(), + train_distribute=mirrored_strategy.CoreMirroredStrategy(), experimental_distribute=DistributeConfig( - train_distribute=mirrored_strategy.MirroredStrategy())) + train_distribute=mirrored_strategy.CoreMirroredStrategy())) with self.assertRaises(ValueError): run_config_lib.RunConfig( - eval_distribute=mirrored_strategy.MirroredStrategy(), + eval_distribute=mirrored_strategy.CoreMirroredStrategy(), experimental_distribute=DistributeConfig( - eval_distribute=mirrored_strategy.MirroredStrategy())) + eval_distribute=mirrored_strategy.CoreMirroredStrategy())) def test_init_run_config_none_distribute_coordinator_mode(self): # We don't use distribute coordinator for local training. config = run_config_lib.RunConfig( - train_distribute=mirrored_strategy.MirroredStrategy()) + train_distribute=mirrored_strategy.CoreMirroredStrategy()) dc_training.init_run_config(config, {}) self.assertIsNone(config._distribute_coordinator_mode) @@ -496,7 +521,7 @@ class RunConfigTest(test.TestCase): with test.mock.patch.dict("os.environ", {"TF_CONFIG": json.dumps(TF_CONFIG_WITH_MASTER)}): config = run_config_lib.RunConfig( - train_distribute=mirrored_strategy.MirroredStrategy()) + train_distribute=mirrored_strategy.CoreMirroredStrategy()) self.assertIsNone(config._distribute_coordinator_mode) # When `train_distribute` is not specified, don't use distribute @@ -512,7 +537,7 @@ class RunConfigTest(test.TestCase): with test.mock.patch.dict("os.environ", {"TF_CONFIG": json.dumps(TF_CONFIG_WITH_CHIEF)}): config = run_config_lib.RunConfig( - train_distribute=mirrored_strategy.MirroredStrategy()) + train_distribute=mirrored_strategy.CoreMirroredStrategy()) self.assertEqual(config._distribute_coordinator_mode, dc.CoordinatorMode.INDEPENDENT_WORKER) @@ -521,7 +546,7 @@ class RunConfigTest(test.TestCase): # `experimental.remote_cluster` is set use distribute coordinator with # STANDALONE_CLIENT mode. config = run_config_lib.RunConfig( - train_distribute=mirrored_strategy.MirroredStrategy(), + train_distribute=mirrored_strategy.CoreMirroredStrategy(), experimental_distribute=DistributeConfig( remote_cluster={"chief": ["fake_worker"]})) self.assertEqual(config._distribute_coordinator_mode, @@ -529,5 +554,15 @@ class RunConfigTest(test.TestCase): if __name__ == "__main__": + # Reduce `recovery_wait_secs` from 30 seconds so the test completes quickly. + orig_init = session_manager.SessionManager.__init__ + + def new_init(*args, **kwargs): + kwargs.pop("recovery_wait_secs", None) + kwargs["recovery_wait_secs"] = 0.5 + orig_init(*args, **kwargs) + + session_manager.SessionManager.__init__ = new_init + with test.mock.patch.object(sys, "exit", os._exit): test.main() 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 da7f8c548f94972b6ec0a67848e1520386d1e28b..1ce91ecaf22a80a53124c8f00fac05c6b4711ed9 100644 --- a/tensorflow/contrib/distribute/python/examples/keras_mnist.py +++ b/tensorflow/contrib/distribute/python/examples/keras_mnist.py @@ -20,18 +20,26 @@ from __future__ import print_function import tensorflow as tf +from tensorflow.python.distribute import mirrored_strategy +from tensorflow.python.keras.optimizer_v2 import rmsprop + + NUM_CLASSES = 10 -def get_input_datasets(): +def get_input_datasets(use_bfloat16=False): """Downloads the MNIST dataset and creates train and eval dataset objects. + Args: + use_bfloat16: Boolean to determine if input should be cast to bfloat16 + Returns: Train dataset, eval dataset and input shape. """ # input image dimensions img_rows, img_cols = 28, 28 + cast_dtype = tf.bfloat16 if use_bfloat16 else tf.float32 # the data, split between train and test sets (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data() @@ -57,12 +65,13 @@ def get_input_datasets(): # train dataset train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train)) train_ds = train_ds.repeat() - train_ds = train_ds.shuffle(100) + train_ds = train_ds.map(lambda x, y: (tf.cast(x, cast_dtype), y)) train_ds = train_ds.batch(64, drop_remainder=True) # eval dataset eval_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test)) eval_ds = eval_ds.repeat() + eval_ds = eval_ds.map(lambda x, y: (tf.cast(x, cast_dtype), y)) eval_ds = eval_ds.batch(64, drop_remainder=True) return train_ds, eval_ds, input_shape @@ -97,20 +106,24 @@ def main(_): # Build the train and eval datasets from the MNIST data. Also return the # input shape which is constructed based on the `image_data_format` # i.e channels_first or channels_last. + 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. - strategy = tf.contrib.distribute.MirroredStrategy() - - # 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=tf.train.RMSPropOptimizer(learning_rate=0.001), - metrics=['accuracy'], - distribute=strategy) + # TODO(priyag): Use `tf.distribute.MirroredStrategy` once available. + strategy = mirrored_strategy.MirroredStrategy(['/gpu:0', '/cpu:0']) + + # 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..11a3b5e91a52a6881d48a0aceadbd901a46e86b2 --- /dev/null +++ b/tensorflow/contrib/distribute/python/examples/mnist_eager_multigpu.py @@ -0,0 +1,151 @@ +# 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 + +import numpy as np +import tensorflow as tf + +tf.flags.DEFINE_integer("num_gpus", None, "How many GPUs should we run on?" + "Defaults to all available GPUs, otherwise CPU.") +tf.flags.DEFINE_integer("batch_size", 64, + "What should be the size of each batch?") +tf.flags.DEFINE_integer("num_epochs", 10, "How many epochs to run?") +tf.flags.DEFINE_float("learning_rate", 0.01, "Learning Rate") +tf.flags.DEFINE_float("momentum", 0.5, "SGD momentum") + +FLAGS = tf.flags.FLAGS +NUM_TRAIN_IMAGES = 60000 +NUM_TEST_IMAGES = 10000 + + +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_eager_execution() + + 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.train.MomentumOptimizer(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): + # Train + print("Starting epoch {}".format(epoch)) + train_iterator.initialize() + for _ in range(NUM_TRAIN_IMAGES // FLAGS.batch_size): + strategy.experimental_run(train_step, train_iterator) + 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() + for _ in range(NUM_TEST_IMAGES // FLAGS.batch_size): + strategy.experimental_run(test_step, test_iterator) + 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__": + tf.app.run() diff --git a/tensorflow/contrib/distribute/python/input_lib_test.py b/tensorflow/contrib/distribute/python/input_lib_test.py new file mode 100644 index 0000000000000000000000000000000000000000..10a58316ec5b3d9d968a88c5c39ff70c277daa65 --- /dev/null +++ b/tensorflow/contrib/distribute/python/input_lib_test.py @@ -0,0 +1,246 @@ +# 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 the input_lib library.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from absl.testing import parameterized + +from tensorflow.contrib.distribute.python import combinations +from tensorflow.contrib.distribute.python import multi_worker_test_base +from tensorflow.python.data.experimental.ops import batching +from tensorflow.python.data.ops import dataset_ops +from tensorflow.python.distribute import distribute_lib +from tensorflow.python.distribute import input_lib +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.ops import control_flow_ops +from tensorflow.python.util import nest + + +class InputIteratorTestBase(test.TestCase): + + def _test_iterator(self, input_type, dataset_fn, worker_device_pairs, + expected_values, sess=None, split_batch_by=None): + devices = nest.flatten([ds for _, ds in worker_device_pairs]) + device_map = values.ReplicaDeviceMap(devices) + input_workers = input_lib.InputWorkers(device_map, worker_device_pairs) + + if input_type == "input_fn": + input_contexts = [ + distribute_lib.InputContext() for _ in worker_device_pairs] + input_fn = lambda _: dataset_fn() + iterator = input_lib.InputFunctionIterator( + input_fn, input_workers, input_contexts) + else: + iterator = input_lib.DatasetIterator( + dataset_fn(), input_workers, split_batch_by) + + evaluate = lambda x: sess.run(x) if sess else self.evaluate(x) + + evaluate(control_flow_ops.group(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.assertAllEqual(expected_value, computed_value) + + with self.assertRaises(errors.OutOfRangeError): + next_element = iterator.get_next() + evaluate([values.select_replica(r, next_element) + for r in range(len(devices))]) + + # After re-initializing the iterator, should be able to iterate again. + evaluate(control_flow_ops.group(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.assertAllEqual(expected_value, computed_value) + + +class InputIteratorSingleWorkerTest(InputIteratorTestBase, + parameterized.TestCase): + + @combinations.generate(combinations.combine( + mode=["graph", "eager"], + input_type=["input_fn", "dataset"])) + def testOneDeviceCPU(self, input_type): + worker_device_pairs = [("", ["/device:CPU:0"])] + dataset_fn = lambda: dataset_ops.Dataset.range(10) + + expected_values = [[i] for i in range(10)] + + self._test_iterator(input_type, dataset_fn, worker_device_pairs, + expected_values) + + @combinations.generate(combinations.combine( + mode=["graph", "eager"], + input_type=["input_fn", "dataset"], + required_gpus=1)) + def testTwoDevicesOneGPUOneCPU(self, input_type): + worker_device_pairs = [("", ["/device:GPU:0", "/device:CPU:0"])] + dataset_fn = lambda: dataset_ops.Dataset.range(10) + + expected_values = [[i, i+1] for i in range(0, 10, 2)] + + self._test_iterator(input_type, dataset_fn, worker_device_pairs, + expected_values) + + @combinations.generate(combinations.combine( + mode=["graph", "eager"], + input_type=["input_fn", "dataset"], + required_gpus=1)) + def testTupleDataset(self, input_type): + worker_device_pairs = [("", ["/device:GPU:0", "/device:CPU:0"])] + def dataset_fn(): + dataset1 = dataset_ops.Dataset.range(10) + dataset2 = dataset_ops.Dataset.range(10).map(lambda x: x**2) + return 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(input_type, dataset_fn, worker_device_pairs, + expected_values) + + @combinations.generate(combinations.combine( + mode=["graph", "eager"], + input_type=["input_fn", "dataset"], + required_gpus=1)) + def testUnevenDatasetBatches(self, input_type): + worker_device_pairs = [("", ["/device:GPU:0", "/device:CPU:0"])] + dataset_fn = lambda: dataset_ops.Dataset.range(11) + + expected_values = [[i, i+1] for i in range(0, 10, 2)] + self._test_iterator(input_type, dataset_fn, worker_device_pairs, + expected_values) + + @combinations.generate(combinations.combine( + mode=["graph", "eager"], + input_type=["dataset"], + split_batch_by=[None, 2], + required_gpus=1)) + def testBatchSplitting(self, input_type, split_batch_by): + worker_device_pairs = [("", ["/device:GPU:0", "/device:CPU:0"])] + batch_size = 10 + dataset_fn = lambda: dataset_ops.Dataset.range(100).batch(batch_size) + + updated_batch_size = ( + batch_size // split_batch_by if split_batch_by else batch_size) + expected_values = [[range(i, i+updated_batch_size), + range(i+updated_batch_size, i+2*updated_batch_size)] + for i in range(0, 100, updated_batch_size*2)] + + self._test_iterator(input_type, dataset_fn, worker_device_pairs, + expected_values, sess=None, + split_batch_by=split_batch_by) + + +class InputIteratorMultiWorkerTest( + multi_worker_test_base.MultiWorkerTestBase, InputIteratorTestBase, + parameterized.TestCase): + + def _cpu_devices(self): + return [ + ("/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"])] + + def _cpu_and_one_gpu_devices(self): + return [ + ("/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" + ]) + ] + + @combinations.generate(combinations.combine( + mode=["graph"], + input_type=["input_fn", "dataset"])) + def testOneDevicePerWorker(self, input_type): + worker_devices = self._cpu_devices() + with context.graph_mode(), self.cached_session() as sess: + dataset_fn = lambda: dataset_ops.Dataset.range(4) + self._test_iterator(input_type, dataset_fn, worker_devices, + [[0, 0], [1, 1], [2, 2], [3, 3]], sess) + + @combinations.generate(combinations.combine( + mode=["graph"], + input_type=["input_fn", "dataset"], + required_gpus=1)) + def testTwoDevicesPerWorker(self, input_type): + worker_devices = self._cpu_and_one_gpu_devices() + with context.graph_mode(), self.cached_session() as sess: + dataset_fn = lambda: dataset_ops.Dataset.range(4) + self._test_iterator(input_type, dataset_fn, worker_devices, + [[0, 1, 0, 1], [2, 3, 2, 3]], sess) + + @combinations.generate(combinations.combine( + mode=["graph"], + input_type=["input_fn", "dataset"])) + def testTupleDataset(self, input_type): + worker_devices = self._cpu_devices() + with context.graph_mode(), self.cached_session() as sess: + def dataset_fn(): + dataset1 = dataset_ops.Dataset.range(4) + dataset2 = dataset_ops.Dataset.range(4).map(lambda x: x**2) + return dataset_ops.Dataset.zip((dataset1, dataset2)) + + expected_values = [[(i, i**2), (i, i**2)] for i in range(0, 4)] + self._test_iterator(input_type, dataset_fn, worker_devices, + expected_values, sess) + + +class SplitDatasetBatchTest(test.TestCase): + + def testBatchDataset(self): + dataset = dataset_ops.Dataset.range(100).batch(20) + split_batch_by = 2 + result_dataset = input_lib._split_dataset_batch(dataset, split_batch_by) + expected_values = [range(i, i+10) for i in range(0, 100, 10)] + result = [self.evaluate(el) for el in result_dataset] + self.assertAllEqual(expected_values, result) + + def testMapAndBatchDataset(self): + dataset = dataset_ops.Dataset.range(100) + dataset = dataset.apply(batching.map_and_batch(lambda x: x, 20)) + split_batch_by = 2 + result_dataset = input_lib._split_dataset_batch(dataset, split_batch_by) + expected_values = [range(i, i+10) for i in range(0, 100, 10)] + result = [self.evaluate(el) for el in result_dataset] + self.assertAllEqual(expected_values, result) + + def testPrefetchDataset(self): + dataset = dataset_ops.Dataset.range(100).batch(20).prefetch(1) + split_batch_by = 2 + result_dataset = input_lib._split_dataset_batch(dataset, split_batch_by) + expected_values = [range(i, i+10) for i in range(0, 100, 10)] + result = [self.evaluate(el) for el in result_dataset] + self.assertAllEqual(expected_values, result) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/distribute/python/input_ops.py b/tensorflow/contrib/distribute/python/input_ops.py deleted file mode 100644 index f07ec8234dfe87f2869cd7c2dd6a64c477712d15..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/distribute/python/input_ops.py +++ /dev/null @@ -1,140 +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. -# ============================================================================== -"""Input-pipeline utilities for Distribution strategies.""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -from tensorflow.python.data.ops import readers -from tensorflow.python.data.util import nest -from tensorflow.python.framework import ops -from tensorflow.python.ops import array_ops -from tensorflow.python.ops import math_ops -from tensorflow.python.platform import tf_logging - -# TODO(priyag): Any other reader datasets to consider here? -_READER_DATASET_OPS = [ - "TextLineDataset", - "TFRecordDataset", - "FixedLengthRecordDataset" -] - - -# pylint: disable=protected-access -def auto_shard_dataset(dataset, num_shards, index): - """Shard the input pipeline by sharding the underlying list of files. - - Args: - dataset: A `tf.data.Dataset` instance, typically the result of a bunch of - dataset transformations. - num_shards: A `tf.int64` scalar `tf.Tensor`, representing the number of - shards operating in parallel. Same usage as in `Dataset.shard`. - index: A `tf.int64` scalar `tf.Tensor`, representing the worker index. - Same usage as in `Dataset.shard`. - - Returns: - A modified `Dataset` obtained by updating the pipeline sharded by the - files. The input dataset will be returned if we cannot automatically - determine a good way to shard the input dataset. - """ - - # TODO(priyag): Clone datasets instead of updating in place, similar to the - # clone method for TFRecordDataset. - def _auto_shard_impl(dataset, found_reader_op): - """Recursive implementation of auto sharding.""" - - if not found_reader_op: - # TODO(priyag): Make this check more robust by enforcing some common - # property on reader datasets. - if (isinstance(dataset, readers.TextLineDataset) or - isinstance(dataset, readers.FixedLengthRecordDataset)): - filenames_tensor = dataset._filenames - num_files = array_ops.size(filenames_tensor) - sharded_filenames_tensor = array_ops.gather( - filenames_tensor, math_ops.range(index, num_files, num_shards)) - dataset._filenames = sharded_filenames_tensor - return dataset - elif isinstance(dataset, readers.TFRecordDataset): - # `TFRecordDataset` needs to be handled separately than other readers - # because it converts filenames to a dataset first. Also, we clone it - # instead of updating in place because it has special logic in the - # constructor. Eventually we will change all cases to clone datasets - # instead of updating in-place. - return dataset._clone( - filenames=dataset._filenames.shard(num_shards, index)) - elif hasattr(dataset, "_map_func"): - # TODO(priyag): Make this check more robust by enforcing some common - # property on all map/flatmap/interleave datasets. - map_func_def = dataset._map_func.definition - for node in map_func_def.node_def: - if node.op in _READER_DATASET_OPS: - found_reader_op = True - break - elif node.op == "FlatMapDataset": - # TODO(priyag): Should this check for other map datasets? Should it - # be recursive? It is too specific to implementation of - # TFRecordDataset right now. - nested_func_name = node.attr["f"].func.name - nested_func = ops.get_default_graph()._functions[nested_func_name] - for nested_node in nested_func.definition.node_def: - if nested_node.op in _READER_DATASET_OPS: - found_reader_op = True - break - if found_reader_op: - break - if found_reader_op: - dataset._input_dataset = _auto_shard_impl( - dataset._input_dataset, found_reader_op) - return dataset - - # TODO(priyag): Make _input_dataset(s) a common property of all datasets to - # make this check more robust. - if hasattr(dataset, "_input_dataset"): - dataset._input_dataset = _auto_shard_impl( - dataset._input_dataset, found_reader_op) - if hasattr(dataset, "_dataset_to_concatenate"): - # Special case for `ConcatentateDataset`. We want to shard all input - # datasets. - dataset._dataset_to_concatenate = _auto_shard_impl( - dataset._dataset_to_concatenate, found_reader_op) - return dataset - - if hasattr(dataset, "_datasets"): - # Special case for `ZipDataset`. - dataset._datasets = nest.pack_sequence_as(dataset._datasets, [ - _auto_shard_impl(ds, found_reader_op) - for ds in nest.flatten(dataset._datasets) - ]) - return dataset - - if not found_reader_op: - tf_logging.warn( - "Could not find a standard reader in the input pipeline" - "(one of TextLineDataset, TFRecordDataset, FixedLengthRecordDataset)." - "So auto-sharding is not done. Please verify correctness of " - "auto-sharding for your input.") - # TODO(yuefengz): maybe still shard it? - return dataset - - # TODO(priyag): What do we want to do if the number of filenames is - # uneven in the number of shards? By default, this will just return as - # many items it can before throwing OutOfRangeError. - # TODO(priyag): This will shard the filenames before any shuffling of the - # filename dataset. It might be desirable to shard after shuffling - # filenames? If so, how do we achieve that? - return dataset.shard(num_shards, index) - - return _auto_shard_impl(dataset=dataset, found_reader_op=False) diff --git a/tensorflow/contrib/distribute/python/input_ops_test.py b/tensorflow/contrib/distribute/python/input_ops_test.py deleted file mode 100644 index 559de97bb1f93f990ddaf775d9203d5a2d46aa99..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/distribute/python/input_ops_test.py +++ /dev/null @@ -1,249 +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. -# ============================================================================== -"""Tests for input pipeline modifications for distribution strategies.""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import os - -from tensorflow.contrib.distribute.python import input_ops -from tensorflow.python.data.ops import dataset_ops -from tensorflow.python.data.ops import readers -from tensorflow.python.framework import errors -from tensorflow.python.lib.io import python_io -from tensorflow.python.platform import test -from tensorflow.python.util import compat - - -class AutoShardDatasetTest(test.TestCase): - - def setUp(self): - super(AutoShardDatasetTest, self).setUp() - self._num_files = 10 - self._num_records = 4 - self._num_shards = 2 - self._shard_index = 0 - self._record_bytes = 10 - - def _record(self, r, f): - return compat.as_bytes("Record %d of file %d" % (r, f)) - - def _text_line(self, r, f): - return compat.as_bytes("Text line %d of file %d" % (r, f)) - - def _fixed_length_record(self, r, f): - return compat.as_bytes(str((r * f) % 10) * self._record_bytes) - - def _createTFRecordFiles(self): - filenames = [] - for i in range(self._num_files): - fn = os.path.join(self.get_temp_dir(), "tf_record.%d.txt" % i) - filenames.append(fn) - writer = python_io.TFRecordWriter(fn) - for j in range(self._num_records): - record = self._record(j, i) - writer.write(record) - writer.close() - return filenames - - def _createTextFiles(self): - filenames = [] - for i in range(self._num_files): - fn = os.path.join(self.get_temp_dir(), "text_line.%d.txt" % i) - filenames.append(fn) - contents = [] - for j in range(self._num_records): - contents.append(self._text_line(j, i)) - if j + 1 != self._num_records or i == 0: - contents.append(b"\r\n") - contents = b"".join(contents) - - with open(fn, "wb") as f: - f.write(contents) - return filenames - - def _createFixedLengthRecordFiles(self): - filenames = [] - for i in range(self._num_files): - fn = os.path.join(self.get_temp_dir(), "fixed_length_record.%d.txt" % i) - filenames.append(fn) - with open(fn, "wb") as f: - for j in range(self._num_records): - f.write(self._fixed_length_record(j, i)) - return filenames - - def _verifySimpleShardingOutput(self, dataset, record_fn): - iterator = dataset.make_one_shot_iterator() - next_element = iterator.get_next() - with self.cached_session() as sess: - for f in range(self._shard_index, self._num_files, self._num_shards): - for r in range(self._num_records): - self.assertAllEqual(record_fn(r, f), sess.run(next_element)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(next_element) - - def testTFRecordDataset(self): - dataset = readers.TFRecordDataset(self._createTFRecordFiles()) - dataset = input_ops.auto_shard_dataset( - dataset, self._num_shards, self._shard_index) - - self._verifySimpleShardingOutput(dataset, self._record) - - def testFlatMap(self): - dataset = dataset_ops.Dataset.from_tensor_slices( - self._createTFRecordFiles()) - dataset = dataset.flat_map(readers.TFRecordDataset) - dataset = input_ops.auto_shard_dataset( - dataset, self._num_shards, self._shard_index) - - self._verifySimpleShardingOutput(dataset, self._record) - - def testInterleave(self): - dataset = dataset_ops.Dataset.from_tensor_slices( - self._createTFRecordFiles()) - dataset = dataset.interleave( - readers.TFRecordDataset, cycle_length=4, block_length=self._num_records) - dataset = input_ops.auto_shard_dataset( - dataset, self._num_shards, self._shard_index) - - # Since block_length == num records in each file, the output will still - # contain records in order of files. - self._verifySimpleShardingOutput(dataset, self._record) - - def testListfiles(self): - filenames = self._createTFRecordFiles() - file_pattern = filenames[0].rsplit("/", 1)[0] + "/tf_record.*.txt" - dataset = dataset_ops.Dataset.list_files(file_pattern, shuffle=False) - dataset = dataset.flat_map(readers.TFRecordDataset) - dataset = input_ops.auto_shard_dataset( - dataset, self._num_shards, self._shard_index) - - iterator = dataset.make_one_shot_iterator() - next_element = iterator.get_next() - with self.cached_session() as sess: - actual, expected = [], [] - for f in range(self._shard_index, self._num_files, self._num_shards): - for r in range(self._num_records): - actual.append(sess.run(next_element)) - expected.append(self._record(r, f)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(next_element) - self.assertAllEqual(expected, actual) - - def testComplexPipeline(self): - # Setup a complex input pipeline. - batch_size = 2 - num_epochs = 5 - dataset = dataset_ops.Dataset.from_tensor_slices( - self._createTFRecordFiles()) - dataset = dataset.shuffle(buffer_size=self._num_files) - dataset = dataset.flat_map(readers.TFRecordDataset) - dataset = dataset.prefetch(buffer_size=batch_size) - dataset = dataset.shuffle(2 * self._num_files * self._num_records) - dataset = dataset.repeat(num_epochs) - dataset = dataset.map(lambda x: x) - dataset = dataset.batch(batch_size) - dataset = dataset.prefetch(buffer_size=None) - - # Auto shard. - dataset = input_ops.auto_shard_dataset( - dataset, self._num_shards, self._shard_index) - - # Verify output. - iterator = dataset.make_one_shot_iterator() - next_element = iterator.get_next() - with self.cached_session() as sess: - actual = [] - num_iterations = (self._num_files * self._num_records * num_epochs) // ( - self._num_shards * batch_size) - for _ in range(num_iterations): - actual.extend(sess.run(next_element)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(next_element) - - expected = [] - for f in range(0, self._num_files, self._num_shards): - for r in range(self._num_records): - expected.append(self._record(r, f)) - expected *= num_epochs - - self.assertAllEqual(sorted(expected), sorted(actual)) - - def testZip(self): - dataset1 = readers.TFRecordDataset(self._createTFRecordFiles()) - dataset2 = readers.TextLineDataset(self._createTextFiles()) - dataset = dataset_ops.Dataset.zip((dataset1, dataset2)) - dataset = input_ops.auto_shard_dataset( - dataset, self._num_shards, self._shard_index) - - record_fn = lambda r, f: (self._record(r, f), self._text_line(r, f)) - self._verifySimpleShardingOutput(dataset, record_fn) - - def testConcat(self): - dataset1 = readers.TFRecordDataset(self._createTFRecordFiles()) - dataset2 = readers.TextLineDataset(self._createTextFiles()) - dataset = dataset1.concatenate(dataset2) - dataset = input_ops.auto_shard_dataset( - dataset, self._num_shards, self._shard_index) - - iterator = dataset.make_one_shot_iterator() - next_element = iterator.get_next() - with self.cached_session() as sess: - for f in range(self._shard_index, self._num_files, self._num_shards): - for r in range(self._num_records): - self.assertAllEqual(self._record(r, f), sess.run(next_element)) - for f in range(self._shard_index, self._num_files, self._num_shards): - for r in range(self._num_records): - self.assertAllEqual(self._text_line(r, f), sess.run(next_element)) - with self.assertRaises(errors.OutOfRangeError): - sess.run(next_element) - - def testTextLineReader(self): - dataset = readers.TextLineDataset(self._createTextFiles()) - dataset = input_ops.auto_shard_dataset( - dataset, self._num_shards, self._shard_index) - - self._verifySimpleShardingOutput(dataset, self._text_line) - - def testTextLineReaderWithFlatMap(self): - dataset = dataset_ops.Dataset.from_tensor_slices(self._createTextFiles()) - dataset = dataset.flat_map(readers.TextLineDataset) - dataset = input_ops.auto_shard_dataset( - dataset, self._num_shards, self._shard_index) - - self._verifySimpleShardingOutput(dataset, self._text_line) - - def testFixedLengthReader(self): - dataset = readers.FixedLengthRecordDataset( - self._createFixedLengthRecordFiles(), self._record_bytes) - dataset = input_ops.auto_shard_dataset( - dataset, self._num_shards, self._shard_index) - - self._verifySimpleShardingOutput(dataset, self._fixed_length_record) - - def testFixedLengthReaderWithFlatMap(self): - dataset = dataset_ops.Dataset.from_tensor_slices( - self._createFixedLengthRecordFiles()) - dataset = dataset.flat_map( - lambda f: readers.FixedLengthRecordDataset(f, self._record_bytes)) - dataset = input_ops.auto_shard_dataset( - dataset, self._num_shards, self._shard_index) - - self._verifySimpleShardingOutput(dataset, self._fixed_length_record) - -if __name__ == "__main__": - test.main() diff --git a/tensorflow/contrib/distribute/python/keras_backward_compat_test.py b/tensorflow/contrib/distribute/python/keras_backward_compat_test.py new file mode 100644 index 0000000000000000000000000000000000000000..92de8e643e7588365c23dc8513e197c0869c9ecf --- /dev/null +++ b/tensorflow/contrib/distribute/python/keras_backward_compat_test.py @@ -0,0 +1,1071 @@ +# 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 using DistributionStrategy.""" +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.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.ops import dataset_ops +from tensorflow.python.distribute import distribute_lib +from tensorflow.python.eager import test +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.ops.parsing_ops import gen_parsing_ops +from tensorflow.python.training import gradient_descent +from tensorflow.python.training import rmsprop + +_RANDOM_SEED = 1337 +_TRAIN_SIZE = 200 +_INPUT_SIZE = (10,) +_NUM_CLASS = 2 + + +# TODO(anjalisridhar): Add a decorator that will allow us to run these tests as +# part of the tf.keras unit tests suite. +def simple_sequential_model(): + model = keras.models.Sequential() + model.add(keras.layers.Dense(16, activation='relu', input_shape=_INPUT_SIZE)) + model.add(keras.layers.Dropout(0.1)) + model.add(keras.layers.Dense(_NUM_CLASS, activation='softmax')) + return model + + +def simple_functional_model(): + a = keras.layers.Input(shape=_INPUT_SIZE) + b = keras.layers.Dense(16, activation='relu')(a) + b = keras.layers.Dropout(0.1)(b) + b = keras.layers.Dense(_NUM_CLASS, activation='softmax')(b) + model = keras.models.Model(inputs=[a], outputs=[b]) + 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') + input_m = keras.layers.Input(shape=(8,), dtype='string', name='input_m') + dense = keras.layers.Dense(8, name='dense_1') + + interm_a = dense(input_a) + # Read m + interm_m = keras.layers.Lambda(gen_parsing_ops.string_to_number)(input_m) + interm_s = keras.layers.Lambda(lambda k: k[0] * k[1])([interm_m, interm_a]) + interm_b = dense(input_b) + merged = keras.layers.concatenate([interm_s, interm_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, input_m], outputs=[output_c, output_d]) + model.compile( + loss='categorical_crossentropy', + optimizer=gradient_descent.GradientDescentOptimizer(0.001), + metrics={ + 'dense_2': 'categorical_accuracy', + 'dense_3': 'categorical_accuracy' + }) + return model + + +def get_ds_train_input_fn(): + np.random.seed(_RANDOM_SEED) + (x_train, y_train), _ = testing_utils.get_test_data( + train_samples=_TRAIN_SIZE, + test_samples=50, + input_shape=_INPUT_SIZE, + num_classes=_NUM_CLASS) + y_train = keras.utils.to_categorical(y_train) + + dataset = dataset_ops.Dataset.from_tensor_slices((x_train, y_train)) + dataset = dataset.batch(32) + return dataset + + +def get_ds_test_input_fn(): + np.random.seed(_RANDOM_SEED) + _, (x_test, y_test) = testing_utils.get_test_data( + train_samples=_TRAIN_SIZE, + test_samples=50, + input_shape=_INPUT_SIZE, + num_classes=_NUM_CLASS) + y_test = keras.utils.to_categorical(y_test) + + dataset = dataset_ops.Dataset.from_tensor_slices((x_test, y_test)) + dataset = dataset.batch(32) + return dataset + + +def get_multi_inputs_multi_outputs_data(): + (a_train, c_train), (a_test, c_test) = testing_utils.get_test_data( + train_samples=_TRAIN_SIZE, + test_samples=50, + input_shape=(16,), + num_classes=3, + random_seed=_RANDOM_SEED) + (b_train, d_train), (b_test, d_test) = testing_utils.get_test_data( + train_samples=_TRAIN_SIZE, + test_samples=50, + input_shape=(16,), + num_classes=2, + random_seed=_RANDOM_SEED) + (m_train, _), (m_test, _) = testing_utils.get_test_data( + train_samples=_TRAIN_SIZE, + test_samples=50, + input_shape=(8,), + num_classes=2, + random_seed=_RANDOM_SEED) + + c_train = keras.utils.to_categorical(c_train) + c_test = keras.utils.to_categorical(c_test) + d_train = keras.utils.to_categorical(d_train) + d_test = keras.utils.to_categorical(d_test) + + train_data = { + 'input_a': a_train, + 'input_b': b_train, + 'input_m': m_train, + 'output_c': c_train, + 'output_d': d_train + } + test_data = { + 'input_a': a_test, + 'input_b': b_test, + 'input_m': m_test, + 'output_c': c_test, + 'output_d': d_test + } + + return (train_data, test_data) + + +def batch_wrapper(dataset, batch_size, distribution, repeat=None): + if repeat: + dataset = dataset.repeat(repeat) + # TPUs currently require fully defined input shapes, drop_remainder ensures + # the input will have fully defined shapes. + if isinstance(distribution, tpu_strategy.TPUStrategy): + return dataset.batch(batch_size, drop_remainder=True) + else: + return dataset.batch(batch_size) + + +def get_model(): + x = keras.layers.Input(shape=(3,), name='input') + y = keras.layers.Dense(4, name='dense')(x) + model = keras.Model(x, y) + return model + + +def get_dataset(distribution): + inputs = np.zeros((10, 3), dtype=np.float32) + targets = np.zeros((10, 4), dtype=np.float32) + dataset = dataset_ops.Dataset.from_tensor_slices((inputs, targets)) + dataset = dataset.repeat(100) + dataset = batch_wrapper(dataset, 10, distribution) + return dataset + + +def get_predict_dataset(distribution): + inputs = np.zeros((10, 3), dtype=np.float32) + dataset = dataset_ops.Dataset.from_tensor_slices(inputs) + dataset = dataset.repeat(100) + dataset = batch_wrapper(dataset, 10, distribution) + return dataset + + +def multi_input_output_model(): + a = keras.layers.Input(shape=(3,), name='input_a') + b = keras.layers.Input(shape=(5,), name='input_b') + # TODO(anjalisridhar): Change the output dimension of the second Dense layer + # once the iterator output validation issue has been fixed. + dense_1 = keras.layers.Dense(7, name='dense_1') + dense_2 = keras.layers.Dense(7, name='dense_2') + c = dense_1(a) + d = dense_2(b) + e = keras.layers.Dropout(0.5, name='dropout')(c) + model = keras.models.Model([a, b], [d, e]) + return model + + +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.""" + training_epochs = 2 + global_batch_size = 64 + batch_size = global_batch_size + # TODO(b/118776054): Use global batch size for Keras/DS support. + use_per_core_batch_size = ( + with_distribution and + not distributed_training_utils.global_batch_size_supported( + with_distribution)) + if use_per_core_batch_size: + batch_size //= with_distribution.num_replicas_in_sync + + if use_numpy: + training_inputs = { + 'batch_size': batch_size, + 'x': x_train, + 'y': y_train, + 'epochs': training_epochs, + 'shuffle': False, + } + + if use_validation_data: + eval_inputs = None + training_inputs['validation_data'] = (x_train, y_train) + else: + eval_inputs = { + 'batch_size': batch_size, + 'x': x_train, + 'y': y_train, + } + predict_inputs = { + 'x': np.array(x_predict, dtype=np.float32), + } + else: + # For dataset inputs, we do not pass batch_size to + # keras.fit/evaluate/predict. The batch size is part of the dataset. + train_dataset = dataset_ops.Dataset.from_tensor_slices( + (x_train, y_train)) + x = batch_wrapper( + train_dataset, batch_size, with_distribution, repeat=training_epochs) + + training_inputs = { + 'batch_size': None, + 'x': x, + 'y': None, + 'epochs': training_epochs, + 'shuffle': False, + 'steps_per_epoch': len(x_train) // global_batch_size, + } + if use_validation_data: + eval_inputs = None # Remove the eval_inputs + eval_dataset = dataset_ops.Dataset.from_tensor_slices( + (x_train, y_train)) + x = batch_wrapper(eval_dataset, batch_size, with_distribution) + training_inputs['validation_data'] = x + training_inputs['validation_steps'] = 5 + else: + eval_inputs = { + 'batch_size': None, + 'x': x, + 'y': None, + 'steps': 20, + } + + predict_batch_size = len(x_predict) + if use_per_core_batch_size: + predict_batch_size //= with_distribution.num_replicas_in_sync + predict_dataset = dataset_ops.Dataset.from_tensor_slices(x_predict) + predict_dataset = batch_wrapper(predict_dataset, + predict_batch_size, with_distribution) + predict_inputs = { + 'steps': 1, + 'x': predict_dataset, + } + + return training_inputs, eval_inputs, predict_inputs + + +strategies_minus_tpu = [ + combinations.default_strategy, + combinations.one_device_strategy, + combinations.mirrored_strategy_with_gpu_and_cpu, + combinations.mirrored_strategy_with_two_gpus, + combinations.core_mirrored_strategy_with_gpu_and_cpu, + combinations.core_mirrored_strategy_with_two_gpus] + +tpu_strategies = [ + combinations.tpu_strategy, # steps_per_run=2 + combinations.tpu_strategy_one_step] + + +def strategy_minus_tpu_combinations(): + return combinations.combine( + distribution=strategies_minus_tpu, + mode=['graph', 'eager']) + + +def tpu_strategy_combinations(): + 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 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_and_input_combinations(): + return ( + combinations.times( + combinations.combine(distribution=strategies_minus_tpu), + combinations.combine(mode=['graph'], + use_numpy=[True, False], + use_validation_data=[True, False]) + + combinations.combine(mode=['eager'], + use_numpy=[False], + use_validation_data=[False])) + + combinations.times( + combinations.combine(distribution=tpu_strategies), + combinations.combine(mode=['graph'], + use_numpy=[True, False], + use_validation_data=[True, False]))) + + +def strategy_for_numpy_input_combinations(): + return combinations.combine( + distribution=strategies_minus_tpu + tpu_strategies, + mode=['graph']) + + +class TestDistributionStrategyWithNumpyArrays(test.TestCase, + parameterized.TestCase): + + @combinations.generate(strategy_for_numpy_input_combinations()) + def test_calling_model_with_numpy_arrays(self, distribution): + with self.cached_session(): + model = get_model() + + optimizer = gradient_descent.GradientDescentOptimizer(0.001) + loss = 'mse' + metrics = ['mae'] + model.compile(optimizer, loss, metrics=metrics, distribute=distribution) + + 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)) + + # 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) + + @combinations.generate(strategy_for_numpy_input_combinations()) + def test_calling_model_with_nested_numpy_arrays(self, distribution): + with self.cached_session(): + model = multi_input_output_model() + + optimizer = gradient_descent.GradientDescentOptimizer(learning_rate=0.001) + loss = 'mse' + model.compile(optimizer, loss, distribute=distribution) + + input_a_np = np.asarray(np.random.random((64, 3)), dtype=np.float32) + input_b_np = np.asarray(np.random.random((64, 5)), dtype=np.float32) + inputs = [input_a_np, input_b_np] + + output_d_np = np.asarray(np.random.random((64, 7)), dtype=np.float32) + output_e_np = np.asarray(np.random.random((64, 7)), dtype=np.float32) + targets = [output_d_np, output_e_np] + + # Call fit with validation data + model.fit(inputs, targets, epochs=1, batch_size=8, verbose=0) + + # 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) + + @combinations.generate(combinations.combine( + distribution=strategies_minus_tpu, mode=['graph'])) + def test_numpy_with_sample_weights(self, distribution): + model = get_model() + optimizer = rmsprop.RMSPropOptimizer(learning_rate=0.001) + loss = 'mse' + model.compile(optimizer, loss, distribute=distribution) + + inputs = np.zeros((20, 3), np.float32) + targets = np.zeros((20, 4), np.float32) + sample_weights = np.ones((20), np.float32) + + model.fit(inputs, targets, sample_weight=sample_weights, epochs=1, + steps_per_epoch=2, verbose=1) + + @combinations.generate(strategy_for_numpy_input_combinations()) + def test_flatten_predict_outputs(self, distribution): + with self.cached_session(): + model = multi_input_output_model() + + optimizer = gradient_descent.GradientDescentOptimizer(learning_rate=0.001) + loss = 'mse' + model.compile(optimizer, loss, distribute=distribution) + + # We take 6 input samples with each input having a dimension of 3 or 5. + input_a_np = np.asarray(np.random.random((6, 3)), dtype=np.float32) + input_b_np = np.asarray(np.random.random((6, 5)), dtype=np.float32) + inputs = [input_a_np, input_b_np] + + outs = model.predict(inputs, steps=1) + # `predict` a list that is equal in length to the number of model outputs. + # In this test our model has two outputs and each element of `outs` + # corresponds to all the samples of one of the model outputs. + self.assertLen(outs, 2) + # Each of the output samples have a dimension of 7. We should process all + # the available input samples(6). + self.assertAllEqual([6, 7], outs[0].shape) + self.assertAllEqual([6, 7], outs[1].shape) + + +class TestDistributionStrategyWithDatasets(test.TestCase, + parameterized.TestCase): + + @combinations.generate(all_strategy_combinations()) + def test_calling_model_on_same_dataset(self, distribution): + with self.cached_session(): + model = get_model() + + optimizer = gradient_descent.GradientDescentOptimizer(0.001) + loss = 'mse' + metrics = ['mae', keras.metrics.CategoricalAccuracy()] + model.compile(optimizer, loss, metrics=metrics, distribute=distribution) + + dataset = get_dataset(distribution) + + # Call fit with validation data + model.fit(dataset, epochs=1, steps_per_epoch=2, verbose=0, + validation_data=dataset, validation_steps=2) + model.fit(dataset, epochs=1, steps_per_epoch=2, verbose=0, + validation_data=dataset, validation_steps=2) + model.predict(get_predict_dataset(distribution), steps=2) + + @combinations.generate(all_strategy_combinations()) + def test_model_interleaved_eval_same_as_direct_eval(self, distribution): + with self.cached_session(): + user_controlled_model = get_model() + user_controlled_model.compile( + gradient_descent.GradientDescentOptimizer(0.001), + loss='mse', + metrics=['mae', keras.metrics.CategoricalAccuracy()], + distribute=distribution) + + interleaved_model = get_model() + interleaved_model.set_weights(user_controlled_model.get_weights()) + interleaved_model.compile( + gradient_descent.GradientDescentOptimizer(0.001), + loss='mse', + metrics=['mae', keras.metrics.CategoricalAccuracy()], + distribute=distribution) + + dataset = get_dataset(distribution) + + # Call fit with validation interleaved + interleaved_output = interleaved_model.fit( + dataset, epochs=2, steps_per_epoch=2, verbose=1, + validation_data=dataset, validation_steps=2, shuffle=False) + + # Manually control the validation running after each epoch. + user_controlled_output = [] + for _ in range(2): + user_controlled_model.fit( + dataset, epochs=1, steps_per_epoch=2, verbose=1, shuffle=False) + user_controlled_output.append( + user_controlled_model.evaluate(dataset, steps=2)) + + self.assertEqual(interleaved_output.history['val_loss'], + [x[0] for x in user_controlled_output]) + self.assertEqual(interleaved_output.history['val_mean_absolute_error'], + [x[1] for x in user_controlled_output]) + self.assertEqual(interleaved_output.history['val_categorical_accuracy'], + [x[2] for x in user_controlled_output]) + + # TODO(priyag): Enable this test for TPU. Currently tuples/dict don't work + # as clone_model's input_tensors argument only seems to accept list and not + # tuples or dict. + + @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_fit_with_tuple_and_dict_dataset_inputs(self, distribution): + with self.cached_session(): + model = multi_input_output_model() + + optimizer = gradient_descent.GradientDescentOptimizer(learning_rate=0.001) + loss = 'mse' + metrics = ['mae', keras.metrics.CategoricalAccuracy()] + model.compile(optimizer, loss, metrics=metrics, distribute=distribution) + + input_a_np = np.random.random((10, 3)) + input_b_np = np.random.random((10, 5)) + output_d_np = np.random.random((10, 7)) + output_e_np = np.random.random((10, 7)) + + # Test with tuples + dataset_tuple = dataset_ops.Dataset.from_tensor_slices(( + (input_a_np, input_b_np), (output_d_np, output_e_np))) + dataset_tuple = dataset_tuple.repeat(100) + dataset_tuple = dataset_tuple.batch(10) + + model.fit(dataset_tuple, epochs=1, steps_per_epoch=2, verbose=1) + + # Test with dict + dataset_dict = dataset_ops.Dataset.from_tensor_slices(( + {'input_a': input_a_np, 'input_b': input_b_np}, + (output_d_np, output_e_np))) + dataset_dict = dataset_dict.repeat(100) + dataset_dict = dataset_dict.batch(10) + + 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(self, distribution): + with self.cached_session(): + model = get_model() + + optimizer = gradient_descent.GradientDescentOptimizer(0.001) + loss = 'mse' + metrics = ['mae', keras.metrics.CategoricalAccuracy()] + model.compile(optimizer, loss, metrics=metrics, distribute=distribution) + + dataset = get_dataset(distribution) + + model.fit(dataset, epochs=1, steps_per_epoch=2, verbose=1) + model.evaluate(dataset, steps=2, verbose=1) + model.predict(get_predict_dataset(distribution), steps=2) + + @combinations.generate(strategy_and_optimizer_combinations()) + def test_fit_eval_and_predict_with_optimizer(self, distribution, optimizer): + with self.cached_session(): + model = get_model() + + loss = 'mse' + model.compile(optimizer(), loss, distribute=distribution) + + dataset = get_dataset(distribution) + + model.fit(dataset, epochs=1, steps_per_epoch=2, verbose=1) + model.evaluate(dataset, steps=2, verbose=1) + model.predict(get_predict_dataset(distribution), steps=2) + + @combinations.generate(strategy_minus_tpu_combinations()) + def test_dataset_with_sample_weights(self, distribution): + model = get_model() + optimizer = rmsprop.RMSPropOptimizer(learning_rate=0.001) + loss = 'mse' + model.compile(optimizer, loss, distribute=distribution) + + inputs = np.zeros((10, 3), np.float32) + targets = np.zeros((10, 4), np.float32) + sample_weights = np.ones((10), np.float32) + dataset = dataset_ops.Dataset.from_tensor_slices((inputs, targets, + sample_weights)) + dataset = dataset.repeat() + dataset = dataset.batch(10) + + model.fit(dataset, epochs=1, steps_per_epoch=2, verbose=1) + model.evaluate(dataset, steps=2, verbose=1) + model.predict(dataset, steps=2) + + @combinations.generate(combinations.combine( + distribution=[ + combinations.mirrored_strategy_with_gpu_and_cpu, + combinations.core_mirrored_strategy_with_gpu_and_cpu], + mode=['graph', 'eager'])) + # TODO(b/120943676, b/120957836): Re-enable once the validation code is + # restored. + def DISABLED_test_dataset_wrong_input_shape(self, distribution): + with self.cached_session(): + model = get_model() + + optimizer = rmsprop.RMSPropOptimizer(learning_rate=0.001) + loss = 'mse' + model.compile(optimizer, loss, distribute=distribution) + + # Wrong input shape + inputs = np.zeros((10, 5), dtype=np.float32) + targets = np.zeros((10, 4), dtype=np.float32) + dataset = dataset_ops.Dataset.from_tensor_slices((inputs, targets)) + dataset = dataset.repeat(100) + dataset = dataset.batch(10) + + with self.assertRaisesRegexp(ValueError, + 'expected input to have shape'): + model.fit(dataset, epochs=1, steps_per_epoch=2, verbose=0) + + @combinations.generate(combinations.combine( + distribution=[combinations.mirrored_strategy_with_gpu_and_cpu], + mode=['graph', 'eager'])) + # TODO(b/120943676, b/120957836): Re-enable once the validation code is + # restored. + def DISABLED_test_dataset_no_batch_input_validation(self, distribution): + with self.cached_session(): + model = get_model() + + optimizer = rmsprop.RMSPropOptimizer(learning_rate=0.001) + loss = 'mse' + model.compile(optimizer, loss, distribute=distribution) + + # User forgets to batch the dataset + inputs = np.zeros((10, 3), dtype=np.float32) + targets = np.zeros((10, 4), dtype=np.float32) + dataset = dataset_ops.Dataset.from_tensor_slices((inputs, targets)) + dataset = dataset.repeat(100) + + with self.assertRaisesRegexp(ValueError, 'expected input to have shape'): + model.fit(dataset, epochs=1, steps_per_epoch=2, verbose=0) + + @combinations.generate(combinations.combine( + distribution=[combinations.tpu_strategy_one_step], + mode=['graph'])) + def test_dataset_input_shape_fully_defined(self, distribution): + with self.cached_session(): + model = get_model() + + optimizer = rmsprop.RMSPropOptimizer(learning_rate=0.001) + loss = 'mse' + model.compile(optimizer, loss, distribute=distribution) + + dataset = get_dataset(distribution) + # Input shapes are not fully known. Batch dimension is unknown as we are + # not using the drop_remainder argument. + dataset = dataset.repeat(100).batch(10) + + with self.assertRaisesRegexp(ValueError, 'requires fully defined shapes'): + model.fit(dataset, epochs=1, steps_per_epoch=2, verbose=0) + + @combinations.generate(combinations.combine( + distribution=[ + 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'])) + def test_learning_phase_value(self, distribution): + # TODO(anjalisridhar): Modify this test to use Lambdas since we can compare + # meaningful values. Currently we don't pass the learning phase if the + # Lambda layer uses the learning phase. + with self.cached_session(): + x = keras.layers.Input(shape=(1,), name='input') + y = keras.layers.Dense(1, kernel_initializer='ones')(x) + z = keras.layers.Dropout(0.9999)(y) + model = keras.Model(x, z) + initial_weights = model.get_weights() + + optimizer = gradient_descent.GradientDescentOptimizer(0.005) + loss = 'mse' + metrics = ['acc'] + model.compile(optimizer, loss, metrics=metrics, distribute=distribution) + + batch_size = 8 + if isinstance(distribution, mirrored_strategy.CoreMirroredStrategy): + # CoreMirroredStrategy uses global batch size. + batch_size = 8 * distribution.num_replicas_in_sync + + inputs = np.ones((10, 1), dtype=np.float32) + targets = np.ones((10, 1), dtype=np.float32) + dataset = dataset_ops.Dataset.from_tensor_slices((inputs, targets)) + dataset = dataset.repeat().batch(batch_size) + hist = model.fit(dataset, epochs=1, steps_per_epoch=20, verbose=1) + self.assertAlmostEqual(hist.history['acc'][0], 0, 0) + + model.set_weights(initial_weights) + # TODO(psv/anjalisridhar): Enable these lines after we fix b/117431185. + # evaluate_output = model.evaluate(dataset, steps=20) + # self.assertAlmostEqual(evaluate_output[1], 1, 0) + + inputs = np.ones((10, 1), dtype=np.float32) + predict_dataset = dataset_ops.Dataset.from_tensor_slices(inputs) + + predict_dataset = predict_dataset.repeat().batch(batch_size) + output = model.predict(predict_dataset, steps=10) + # `predict` runs for 10 steps + ref_output = np.ones((160, 1), dtype=np.float32) + self.assertArrayNear(output, ref_output, 1e-1) + + @combinations.generate(strategy_minus_tpu_combinations()) + def testOptimizerWithCallbacks(self, distribution): + with self.cached_session(): + model = get_model() + + optimizer = gradient_descent_keras.SGD(0.01) + loss = 'mse' + model.compile(optimizer, loss, distribute=distribution) + + dataset = get_dataset(distribution) + + def schedule(_): + return 0.001 + + model.fit(dataset, epochs=1, steps_per_epoch=2, verbose=0, + callbacks=[keras.callbacks.LearningRateScheduler(schedule)]) + grouped_models = distribution.unwrap(model._distributed_model) + with distribution.scope(): + for m in grouped_models: + self.assertAllClose(0.001, keras.backend.get_value( + m.optimizer.lr), atol=1e-05, rtol=1e-05) + + +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_unsupported_features(self, distribution): + with self.cached_session(): + model = get_model() + + optimizer = gradient_descent.GradientDescentOptimizer(0.001) + loss = 'mse' + metrics = ['mae'] + model.compile(optimizer, loss, metrics=metrics, distribute=distribution) + + dataset = 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. + 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(): + model = get_model() + + optimizer = gradient_descent.GradientDescentOptimizer(0.001) + loss = 'mse' + metrics = ['mae'] + model.compile(optimizer, loss, metrics=metrics, distribute=distribution) + + dataset = 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]]]) + 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), + distribute=distribution) + 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(all_strategy_combinations()) + def test_batchnorm_correctness(self, distribution): + with self.cached_session(): + model = keras.models.Sequential() + norm = keras.layers.BatchNormalization(input_shape=(10,), momentum=0.8) + model.add(norm) + model.compile(loss='mse', + optimizer=gradient_descent.GradientDescentOptimizer(0.01), + distribute=distribution) + + # 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) + + +class TestDistributionStrategyCorrectness(test.TestCase, + parameterized.TestCase): + + @combinations.generate(all_strategy_combinations()) + def test_metric_correctness(self, distribution): + 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. + 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()], + distribute=distribution) + + batch_size = 64 + if not distributed_training_utils.global_batch_size_supported( + distribution): + batch_size //= distribution.num_replicas_in_sync + 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()) + def test_eval_metrics_correctness(self, distribution): + with self.cached_session(): + 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), + distribute=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 = 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): + with self.cached_session(): + 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, + } + + keras.backend.set_image_data_format('channels_last') + np.random.seed(_RANDOM_SEED) + random_seed.set_random_seed(_RANDOM_SEED) + + # Train, eval, and predict datasets are created with the same input numpy + # arrays. + # 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.]] + + # 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. In addition, we add few non-linear layers to make + # it non-trivial. + def _create_model(): + 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)) + return model + + model = _create_model() + initial_weights = model.get_weights() + del model # avoid accident usage. + + def fit_eval_and_predict(with_distribution=None): + model = _create_model() + # We have initialized the model to the same weight for the distribution + # and non-distribution run. + model.set_weights(initial_weights) + model.compile( + loss=keras.losses.mean_squared_error, + optimizer=gradient_descent_keras.SGD(0.5), + metrics=['mse'], + distribute=with_distribution) + + training_inputs, eval_inputs, predict_inputs = ( + get_correctness_test_inputs(use_numpy, use_validation_data, + with_distribution, + x_train, y_train, x_predict)) + + result = {} + result['training_history_1'] = model.fit(**training_inputs).history + + if eval_inputs is not None: + result['eval_result_1'] = model.evaluate(**eval_inputs) + + result['weights_1'] = model.get_weights() + result['predict_result_1'] = model.predict(**predict_inputs) + + # Train and eval again to mimic user's flow. + + result['training_history_2'] = model.fit(**training_inputs).history + + if eval_inputs is not None: + result['eval_result_2'] = model.evaluate(**eval_inputs) + + result['weights_2'] = model.get_weights() + + return result + + results_with_ds = fit_eval_and_predict(with_distribution=distribution) + results_without_ds = fit_eval_and_predict(with_distribution=None) + + # Verify that the weights, training history, eval results, predict outputs + # are the same within some limits of tolerance. + for key in results_with_ds: + if (key.startswith('training_history') and + isinstance(distribution, tpu_strategy.TPUStrategy) and + distribution.extended.steps_per_run > 1): + # TODO(b/119894254): Enable this test for all cases once the + # underlying bug is fixed. + continue + + tolerance = tol_table.get(key, default_tolerance) + + self.assertAllClose( + results_with_ds[key], + results_without_ds[key], + atol=tolerance, + rtol=tolerance, + msg='Fail to assert {}.'.format(key)) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/distribute/python/keras_correctness_test_base.py b/tensorflow/contrib/distribute/python/keras_correctness_test_base.py new file mode 100644 index 0000000000000000000000000000000000000000..ed77e4fc74f07fea72ddb583b24d7044917b29bd --- /dev/null +++ b/tensorflow/contrib/distribute/python/keras_correctness_test_base.py @@ -0,0 +1,475 @@ +# 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. +# ============================================================================== +"""Correctness tests for tf.keras using DistributionStrategy.""" +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.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.ops import dataset_ops +from tensorflow.python.distribute import distribute_lib +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 + +_RANDOM_SEED = 1337 +_EVAL_STEPS = 20 +_GLOBAL_BATCH_SIZE = 64 + +# Note: Please make sure the tests in this file are also covered in +# keras_backward_compat_test for features that are supported with both APIs. + + +all_strategies = [ + combinations.default_strategy, + combinations.one_device_strategy, + combinations.mirrored_strategy_with_gpu_and_cpu, + combinations.mirrored_strategy_with_two_gpus, + combinations.core_mirrored_strategy_with_gpu_and_cpu, + combinations.core_mirrored_strategy_with_two_gpus, + combinations.tpu_strategy, # steps_per_run=2 + combinations.tpu_strategy_one_step, +] + + +def eager_mode_test_configuration(): + return combinations.combine(mode='eager', + use_numpy=False, + use_validation_data=False) + + +def graph_mode_test_configuration(): + return combinations.combine(mode='graph', + use_numpy=[True, False], + use_validation_data=[True, False]) + + +def all_strategy_and_input_config_combinations(): + return ( + combinations.times( + combinations.combine(distribution=all_strategies), + 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. + """ + + strategies = [s for s in all_strategies + if not s.required_tpu and s.required_gpus is not None] + strategies.append(combinations.tpu_strategy_loop_on_device) + strategies.append(combinations.tpu_strategy_one_step_loop_on_device) + return strategies + + +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_loop_on_device, + combinations.tpu_strategy_one_step_loop_on_device] + + return ( + combinations.times( + combinations.combine(distribution=tpu_strategies), + graph_mode_test_configuration())) + + +class MaybeDistributionScope(object): + """Provides a context allowing no distribution strategy.""" + + def __init__(self, distribution): + self._distribution = distribution + self._scope = None + + def __enter__(self): + if self._distribution: + self._scope = self._distribution.scope() + self._scope.__enter__() + + def __exit__(self, exc_type, value, traceback): + if self._distribution: + self._scope.__exit__(exc_type, value, traceback) + self._scope = None + + +def batch_wrapper(dataset, batch_size, distribution, repeat=None): + if repeat: + dataset = dataset.repeat(repeat) + # TPUs currently require fully defined input shapes, drop_remainder ensures + # the input will have fully defined shapes. + if isinstance(distribution, tpu_strategy.TPUStrategy): + return dataset.batch(batch_size, drop_remainder=True) + else: + return dataset.batch(batch_size) + + +def get_batch_size(global_batch_size, distribution): + batch_size = global_batch_size + # TODO(b/118776054): Use global batch size for Keras/DS support. + use_per_core_batch_size = ( + distribution and + not distributed_training_utils.global_batch_size_supported(distribution)) + if use_per_core_batch_size: + batch_size //= distribution.num_replicas_in_sync + return batch_size + + +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.""" + training_epochs = 2 + global_batch_size = _GLOBAL_BATCH_SIZE + batch_size = get_batch_size(global_batch_size, with_distribution) + + if use_numpy: + training_inputs = { + 'batch_size': batch_size, + 'x': x_train, + 'y': y_train, + 'epochs': training_epochs, + 'shuffle': False, + } + + if use_validation_data: + eval_inputs = None + training_inputs['validation_data'] = (x_train, y_train) + else: + eval_inputs = { + 'batch_size': batch_size, + 'x': x_train, + 'y': y_train, + } + predict_inputs = { + 'x': np.array(x_predict, dtype=np.float32), + } + else: + if len(x_train) < _GLOBAL_BATCH_SIZE * _EVAL_STEPS: + # Currently, we cannot detech 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') + # For dataset inputs, we do not pass batch_size to + # keras.fit/evaluate/predict. The batch size is part of the dataset. + train_dataset = dataset_ops.Dataset.from_tensor_slices((x_train, y_train)) + x = batch_wrapper(train_dataset, batch_size, with_distribution, + repeat=training_epochs) + + training_inputs = { + 'batch_size': None, + 'x': x, + 'y': None, + 'epochs': training_epochs, + 'shuffle': False, + 'steps_per_epoch': len(x_train) // global_batch_size, + } + if use_validation_data: + eval_inputs = None # Remove the eval_inputs + eval_dataset = dataset_ops.Dataset.from_tensor_slices((x_train, y_train)) + x = batch_wrapper(eval_dataset, batch_size, with_distribution) + training_inputs['validation_data'] = x + training_inputs['validation_steps'] = 5 + else: + eval_inputs = { + 'batch_size': None, + 'x': x, + 'y': None, + 'steps': _EVAL_STEPS, + } + + predict_batch_size = get_batch_size(len(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) + predict_inputs = { + 'steps': 1, + 'x': predict_dataset, + } + + return training_inputs, eval_inputs, predict_inputs + + +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) + + result = {} + result['training_history_1'] = model.fit(**training_inputs).history + + if eval_inputs is not None: + result['eval_result_1'] = model.evaluate(**eval_inputs) + + result['weights_1'] = model.get_weights() + + if predict_inputs is not None: + # 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. + + result['training_history_2'] = model.fit(**training_inputs).history + + if eval_inputs is not None: + result['eval_result_2'] = model.evaluate(**eval_inputs) + + result['weights_2'] = model.get_weights() + + return result + + +def compare_results(results_with_ds, results_without_ds, distribution, + testcase): + """Compares results of model compiled with/without distribution strategy.""" + + default_tolerance = 1e-5 + 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 + isinstance(distribution, tpu_strategy.TPUStrategy) and + distribution.extended.steps_per_run > 1): + # TODO(b/119894254): Enable this test for all cases once the + # underlying bug is fixed. + continue + + tolerance = _get_compare_result_tolerance(key) + testcase.assertAllClose( + results_with_ds[key], + results_without_ds[key], + atol=tolerance, + rtol=tolerance, + 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 + + def on_batch_begin(self, batch, logs=None): + if self._update_freq and batch % self._update_freq != 0: + return + + # To avoid divergence, limit the value range. + lr = 0.001 * (batch % 10) + keras.backend.set_value(self.model.optimizer.lr, lr) + + +class TestDistributionStrategyCorrectnessBase(test.TestCase, + parameterized.TestCase): + """Model agnostic testing infra to test correctness of Keras models.""" + + 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 + + keras.backend.set_image_data_format('channels_last') + np.random.seed(_RANDOM_SEED) + random_seed.set_random_seed(_RANDOM_SEED) + + 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) + + 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 self.use_numpy: + self.skipTest('Numpy as inputs is not supported with strategy in eager.') + + 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(): + self.set_up_test_config(use_numpy, use_validation_data, with_batch_norm) + self.skip_unsupported_test_configuration(distribution) + + # Train, eval, and predict datasets are created with the same input numpy + # arrays. + 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 = self.get_model() + initial_weights = model.get_weights() + + def input_fn(dist): + return get_correctness_test_inputs( + 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=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=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 (self.with_batch_norm and + distribution.num_replicas_in_sync > 1): + with self.assertRaises(AssertionError): + compare_results(results_with_ds, results_without_ds, distribution, + testcase=self) + else: + compare_results(results_with_ds, results_without_ds, distribution, + testcase=self) + + def run_dynamic_lr_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() + 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 + # every step. So, to compare the CPU/TPU, we let the CPU to behave the + # same as TPU. + 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) + + training_inputs = { + 'batch_size': batch_size, + 'x': x_train, + 'y': y_train, + 'epochs': training_epochs, + 'shuffle': False, + 'callbacks': [LearningRateBatchScheduler(update_freq)], + 'validation_data': (x_train, y_train) + } + # In this test case, we do not care eval and predict. + eval_inputs, predict_inputs = None, None + return training_inputs, eval_inputs, predict_inputs + + results_with_ds = fit_eval_and_predict( + 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=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..7afacab0ddbed8d5b448c2ed2b983bfa18d11b80 --- /dev/null +++ b/tensorflow/contrib/distribute/python/keras_dnn_correctness_test.py @@ -0,0 +1,171 @@ +# 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')) + 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 = [[1.], [2.], [3.], [4.]] + 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..3913f9bc0cdfff6b562d5727ec33eb4d83f4a619 --- /dev/null +++ b/tensorflow/contrib/distribute/python/keras_embedding_model_correctness_test.py @@ -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. +# ============================================================================== +"""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) + + +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..f625664372dfb6814ccbe9539f6abe018d2a4447 --- /dev/null +++ b/tensorflow/contrib/distribute/python/keras_image_model_correctness_test.py @@ -0,0 +1,92 @@ +# 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.training 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))( + 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.GradientDescentOptimizer( + 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 new file mode 100644 index 0000000000000000000000000000000000000000..5349794334b7f6ea3d718343fa84c693dd3d7a3c --- /dev/null +++ b/tensorflow/contrib/distribute/python/keras_optimizer_v2_test.py @@ -0,0 +1,181 @@ +# 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 that show that DistributionStrategy works with canned Estimator.""" + +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.distribute.python import combinations +from tensorflow.python import keras +from tensorflow.python.distribute import distribution_strategy_context as ds_context +from tensorflow.python.framework import constant_op +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import ops +from tensorflow.python.keras.optimizer_v2 import adam +from tensorflow.python.keras.optimizer_v2 import gradient_descent +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 + + +def get_model(): + x = keras.layers.Input(shape=(3,), name='input') + y = keras.layers.Dense(4, name='dense')(x) + model = keras.Model(x, y) + return 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'])) + def testKerasOptimizerWithUnequalInput(self, distribution): + def create_fn(): + 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) + + devices = ['/device:GPU:0', '/device:CPU:0'] + with distribution.scope(): + (var, m, v, op, + counter) = distribution.extended.call_for_each_replica(create_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])])) + # 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) + # 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])])) + # 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'])) + def testOptimizerWithKerasModelAndNumpyArrays(self, distribution): + + with self.cached_session(): + with distribution.scope(): + model = get_model() + optimizer = gradient_descent.SGD(0.001) + loss = 'mse' + metrics = ['mae'] + model.compile(optimizer, loss, metrics=metrics) + + inputs = np.zeros((64, 3), dtype=np.float32) + targets = np.zeros((64, 4), dtype=np.float32) + + model.fit( + inputs, + targets, + epochs=1, + batch_size=2, + verbose=0, + validation_data=(inputs, targets)) + model.evaluate(inputs, targets) + model.predict(inputs) + + +def _replica_id(): + replica_id = ds_context.get_replica_context().replica_id_in_sync_group + if not isinstance(replica_id, ops.Tensor): + replica_id = constant_op.constant(replica_id) + return replica_id + + +if __name__ == '__main__': + test.main() 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..ab56c01d862354bd74330f769502692bd8a8b982 --- /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_loop_on_device_one_core, + combinations.tpu_strategy_one_step_loop_on_device_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 dfa38912897359036d5c333b69f54047f52a2f49..8a607dd070f859aca69ee857d5c5f091f107e0ca 100644 --- a/tensorflow/contrib/distribute/python/keras_test.py +++ b/tensorflow/contrib/distribute/python/keras_test.py @@ -24,9 +24,10 @@ import numpy as np 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.contrib.distribute.python import values 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.estimator import keras as keras_lib from tensorflow.python.estimator import run_config as run_config_lib from tensorflow.python.framework import constant_op @@ -34,19 +35,21 @@ 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 +from tensorflow.python.keras.optimizer_v2 import gradient_descent as gradient_descent_keras from tensorflow.python.ops.parsing_ops import gen_parsing_ops from tensorflow.python.platform import gfile -from tensorflow.python.platform import test from tensorflow.python.summary.writer import writer_cache from tensorflow.python.training import gradient_descent from tensorflow.python.training import rmsprop - _RANDOM_SEED = 1337 _TRAIN_SIZE = 200 _INPUT_SIZE = (10,) _NUM_CLASS = 2 +# Note: Please make sure the tests in this file are also covered in +# keras_backward_compat_test for features that are supported with both APIs. + # TODO(anjalisridhar): Add a decorator that will allow us to run these tests as # part of the tf.keras unit tests suite. @@ -164,7 +167,9 @@ def get_multi_inputs_multi_outputs_data(): return (train_data, test_data) -def batch_wrapper(dataset, batch_size, distribution): +def batch_wrapper(dataset, batch_size, distribution, repeat=None): + if repeat: + dataset = dataset.repeat(repeat) # TPUs currently require fully defined input shapes, drop_remainder ensures # the input will have fully defined shapes. if isinstance(distribution, tpu_strategy.TPUStrategy): @@ -197,30 +202,80 @@ def get_predict_dataset(distribution): return dataset -strategies = [combinations.default_strategy, - combinations.one_device_strategy, - combinations.mirrored_strategy_with_gpu_and_cpu, - combinations.mirrored_strategy_with_two_gpus, - combinations.tpu_strategy_one_step] +def multi_input_output_model(): + a = keras.layers.Input(shape=(3,), name='input_a') + b = keras.layers.Input(shape=(5,), name='input_b') + # TODO(anjalisridhar): Change the output dimension of the second Dense layer + # once the iterator output validation issue has been fixed. + dense_1 = keras.layers.Dense(7, name='dense_1') + dense_2 = keras.layers.Dense(7, name='dense_2') + c = dense_1(a) + d = dense_2(b) + e = keras.layers.Dropout(0.5, name='dropout')(c) + model = keras.models.Model([a, b], [d, e]) + return model + + +strategies_minus_tpu = [ + combinations.default_strategy, + combinations.one_device_strategy, + combinations.mirrored_strategy_with_gpu_and_cpu, + combinations.mirrored_strategy_with_two_gpus, + combinations.core_mirrored_strategy_with_gpu_and_cpu, + combinations.core_mirrored_strategy_with_two_gpus] + +tpu_strategies = [ + combinations.tpu_strategy, # steps_per_run=2 + combinations.tpu_strategy_one_step] + + +def strategy_minus_tpu_combinations(): + return combinations.combine( + distribution=strategies_minus_tpu, + mode=['graph', 'eager']) -def strategy_combinations(): +def tpu_strategy_combinations(): return combinations.combine( - distribution=strategies, + distribution=tpu_strategies, mode=['graph']) +def all_strategy_combinations(): + return strategy_minus_tpu_combinations() + tpu_strategy_combinations() + + +def all_strategy_combinations_minus_default(): + strategy_minus_default_combinations = combinations.combine( + distribution=[ + combinations.one_device_strategy, + combinations.mirrored_strategy_with_gpu_and_cpu, + combinations.mirrored_strategy_with_two_gpus, + combinations.core_mirrored_strategy_with_gpu_and_cpu, + combinations.core_mirrored_strategy_with_two_gpus], + mode=['graph', 'eager']) + return strategy_minus_default_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])) + + +def strategy_for_numpy_input_combinations(): return combinations.combine( - distribution=strategies, - optimizer=[combinations.adagrad_optimizer_v1_fn, - combinations.adam_optimizer_v1_fn, - combinations.gradient_descent_optimizer_v1_fn, - combinations.rmsprop_optimizer_v1_fn], + distribution=strategies_minus_tpu + tpu_strategies, mode=['graph']) -class TestEstimatorDistributionStrategy(test_util.TensorFlowTestCase): +class TestEstimatorDistributionStrategy(test_util.TensorFlowTestCase, + parameterized.TestCase): def setUp(self): self._base_dir = os.path.join(self.get_temp_dir(), @@ -228,17 +283,20 @@ class TestEstimatorDistributionStrategy(test_util.TensorFlowTestCase): gfile.MakeDirs(self._base_dir) self._config = run_config_lib.RunConfig( tf_random_seed=_RANDOM_SEED, model_dir=self._base_dir) - self._dist = mirrored_strategy.MirroredStrategy( - devices=['/device:GPU:0', '/device:GPU:1']) def tearDown(self): writer_cache.FileWriterCache.clear() if os.path.isdir(self._base_dir): gfile.DeleteRecursively(self._base_dir) - def test_train_functional_with_distribution_strategy(self): - dist = mirrored_strategy.MirroredStrategy( - devices=['/device:GPU:0', '/device:GPU:1']) + @combinations.generate(combinations.combine( + distribution=[ + 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'])) + def test_train_functional_with_distribution_strategy(self, distribution): keras_model = simple_functional_model() keras_model.compile( loss='categorical_crossentropy', @@ -246,8 +304,8 @@ class TestEstimatorDistributionStrategy(test_util.TensorFlowTestCase): optimizer=rmsprop.RMSPropOptimizer(learning_rate=0.01)) config = run_config_lib.RunConfig(tf_random_seed=_RANDOM_SEED, model_dir=self._base_dir, - train_distribute=dist, - eval_distribute=dist) + train_distribute=distribution, + eval_distribute=distribution) with self.cached_session(): est_keras = keras_lib.model_to_estimator( keras_model=keras_model, config=config) @@ -261,9 +319,14 @@ class TestEstimatorDistributionStrategy(test_util.TensorFlowTestCase): writer_cache.FileWriterCache.clear() gfile.DeleteRecursively(self._config.model_dir) - def test_train_sequential_with_distribution_strategy(self): - dist = mirrored_strategy.MirroredStrategy( - devices=['/device:GPU:0', '/device:GPU:1']) + @combinations.generate(combinations.combine( + distribution=[ + 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'])) + def test_train_sequential_with_distribution_strategy(self, distribution): keras_model = simple_sequential_model() keras_model.compile( loss='categorical_crossentropy', @@ -271,7 +334,7 @@ class TestEstimatorDistributionStrategy(test_util.TensorFlowTestCase): optimizer=rmsprop.RMSPropOptimizer(learning_rate=0.01)) config = run_config_lib.RunConfig(tf_random_seed=_RANDOM_SEED, model_dir=self._base_dir, - train_distribute=dist) + train_distribute=distribution) with self.cached_session(): est_keras = keras_lib.model_to_estimator( keras_model=keras_model, config=config) @@ -285,7 +348,12 @@ class TestEstimatorDistributionStrategy(test_util.TensorFlowTestCase): writer_cache.FileWriterCache.clear() gfile.DeleteRecursively(self._config.model_dir) - def test_multi_inputs_multi_outputs_with_input_fn_as_dict(self): + @combinations.generate(combinations.combine( + distribution=[ + combinations.mirrored_strategy_with_gpu_and_cpu, + combinations.core_mirrored_strategy_with_gpu_and_cpu], + mode=['graph'])) + def test_multi_inputs_multi_outputs_with_input_fn_as_dict(self, distribution): train_data, test_data = get_multi_inputs_multi_outputs_data() def train_input_fn(): @@ -315,14 +383,14 @@ class TestEstimatorDistributionStrategy(test_util.TensorFlowTestCase): output_dict)).batch(16) self.do_test_multi_inputs_multi_outputs_with_input_fn( - train_input_fn, eval_input_fn) + distribution, train_input_fn, eval_input_fn) - def do_test_multi_inputs_multi_outputs_with_input_fn(self, train_input_fn, - eval_input_fn): + def do_test_multi_inputs_multi_outputs_with_input_fn( + self, distribution, train_input_fn, eval_input_fn): config = run_config_lib.RunConfig( tf_random_seed=_RANDOM_SEED, model_dir=self._base_dir, - train_distribute=self._dist) + train_distribute=distribution) with self.cached_session(): model = multi_inputs_multi_outputs_model() est_keras = keras_lib.model_to_estimator(keras_model=model, config=config) @@ -332,9 +400,12 @@ class TestEstimatorDistributionStrategy(test_util.TensorFlowTestCase): eval_results = est_keras.evaluate(input_fn=eval_input_fn, steps=1) self.assertLess(eval_results['loss'], baseline_eval_results['loss']) - def test_keras_optimizer_with_distribution_strategy(self): - dist = mirrored_strategy.MirroredStrategy( - devices=['/device:GPU:0', '/device:GPU:1']) + @combinations.generate(combinations.combine( + distribution=[ + combinations.mirrored_strategy_with_gpu_and_cpu, + combinations.core_mirrored_strategy_with_gpu_and_cpu], + mode=['graph'])) + def test_keras_optimizer_with_distribution_strategy(self, distribution): keras_model = simple_sequential_model() keras_model.compile( loss='categorical_crossentropy', @@ -342,7 +413,7 @@ class TestEstimatorDistributionStrategy(test_util.TensorFlowTestCase): config = run_config_lib.RunConfig(tf_random_seed=_RANDOM_SEED, model_dir=self._base_dir, - train_distribute=dist) + train_distribute=distribution) with self.cached_session(): est_keras = keras_lib.model_to_estimator(keras_model=keras_model, config=config) @@ -358,146 +429,181 @@ class TestEstimatorDistributionStrategy(test_util.TensorFlowTestCase): class TestDistributionStrategyWithNumpyArrays(test.TestCase, parameterized.TestCase): - @combinations.generate(strategy_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) - - def test_calculating_batch_params(self): - # This verifies that we calculate the number of steps when the batch size - # is specified. - with self.cached_session(): - # 64 is the number of input samples. - inputs = np.zeros((64, 3), dtype=np.float32) - # The number of towers is equal to 3. - strategy = mirrored_strategy.MirroredStrategy(['/device:GPU:0', - '/device:CPU:0', - '/device:GPU:1']) - - with self.assertRaisesRegexp(ValueError, 'Please specify a batch_size ' - 'that is smaller than'): - # The batch size(128) is larger than the number of input - # samples(64). - distributed_training_utils.get_input_batch_params(inputs, - 128, - strategy) - - with self.assertRaisesRegexp(ValueError, 'is smaller than the number ' - 'of towers'): - # The batch size(32) * num_towers(3) is 96 which is greater than the - # number of input samples(64). - distributed_training_utils.get_input_batch_params(inputs, - 32, - strategy) - - # The number of towers now is equal to 2. - strategy = mirrored_strategy.MirroredStrategy(['/device:GPU:0', - '/device:CPU:0']) - # 32 is the batch size per tower. - steps = distributed_training_utils.get_input_batch_params(inputs, - 32, - strategy) - # The number of batches is the ratio of input samples(64) to - # batch size(32) which is 2. The number of steps(1) is the ratio of - # number of batches(2) to the number of towers(2). - self.assertEqual(steps, 1) + @combinations.generate(strategy_for_numpy_input_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 + replica_scale_factor = 1.0 + if not distributed_training_utils.global_batch_size_supported(distribution): + replica_scale_factor = distribution.num_replicas_in_sync - # 16 is the batch size per tower. - steps = distributed_training_utils.get_input_batch_params(inputs, - 16, - strategy) - # The number of batches is the ratio of input samples(64) to - # batch size(16) which is 4. The number of steps(2) is the ratio of - # number of batches(4) to the number of towers(2). + with self.cached_session(): + # Input samples of different sizes + input_20_samples = np.zeros((20, 3), dtype=np.float32) + input_63_samples = np.zeros((63, 3), dtype=np.float32) + input_64_samples = np.zeros((64, 3), dtype=np.float32) + + # Default global batch size 32 for input with 64 samples run in 2 steps + steps, batch_size = distributed_training_utils.get_input_params( + distribution, input_64_samples, steps=None, batch_size=None) + self.assertEqual(batch_size, 32 // replica_scale_factor) self.assertEqual(steps, 2) - def test_calculating_batch_size(self): - with self.cached_session(): - # 64 is the number of input samples. - inputs = np.zeros((64, 3), dtype=np.float32) - targets = np.zeros((64, 4), dtype=np.float32) - - model = get_model() - optimizer = gradient_descent.GradientDescentOptimizer(0.001) - loss = 'mse' - strategy = mirrored_strategy.MirroredStrategy(['/device:GPU:0', - '/device:CPU:0']) - strategy._require_static_shapes = True - - model.compile(optimizer, loss, distribute=strategy) - iterator = model._distribution_standardize_user_data(inputs, - targets, - batch_size=None, - check_steps=True, - steps_name='steps', - steps=3) - - # The global batch size(21) across all towers is the ratio of the input - # samples(64) to the steps(3). - # The batch size(10) per device is the ratio of the global batch size(21) - # to the number of towers(2). - # The global batch size and batch size are rounded integer values. - self.assertEqual(10, distributed_training_utils.get_batch_dimension( - iterator._iterator)) - - @combinations.generate(strategy_combinations()) - def test_calling_model_with_numpy_arrays(self, distribution): - with self.cached_session(): - model = get_model() + # Computed global batch size 20 is lower than 32 if we pass less samples. + steps, batch_size = distributed_training_utils.get_input_params( + distribution, input_20_samples, steps=None, batch_size=None) + self.assertEqual(batch_size, 20 // replica_scale_factor) + self.assertEqual(steps, 1) - optimizer = gradient_descent.GradientDescentOptimizer(0.001) - loss = 'mse' - metrics = ['mae'] - model.compile(optimizer, loss, metrics=metrics, distribute=distribution) + # Default global batch size 32 cannot be used with 63 samples. + with self.assertRaisesRegexp(ValueError, 'not divisible by batch size'): + distributed_training_utils.get_input_params( + distribution, input_63_samples, steps=None, batch_size=None) - inputs = np.zeros((64, 3), dtype=np.float32) - targets = np.zeros((64, 4), dtype=np.float32) + @combinations.generate(strategy_for_numpy_input_combinations()) + def test_calculating_input_params_with_steps_no_batch_size(self, + distribution): + # Calculate the per_replica_batch_size scaling factor for strategies + # that use per_core_batch_size + replica_scale_factor = 1.0 + if not distributed_training_utils.global_batch_size_supported(distribution): + replica_scale_factor = distribution.num_replicas_in_sync - # Call fit with validation data - model.fit(inputs, targets, epochs=1, batch_size=2, verbose=0, - validation_data=(inputs, targets)) + with self.cached_session(): + # Input samples of different sizes + input_63_samples = np.zeros((63, 3), dtype=np.float32) + input_64_samples = np.zeros((64, 3), dtype=np.float32) + + # Computed global batch size is correct for number of specified 1 step + steps, batch_size = distributed_training_utils.get_input_params( + distribution, input_64_samples, steps=1, batch_size=None) + self.assertEqual(batch_size, 64 // replica_scale_factor) + self.assertEqual(steps, 1) - # 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) + # Computed global batch size is correct for number of specified 2 steps + steps, batch_size = distributed_training_utils.get_input_params( + distribution, input_64_samples, steps=2, batch_size=None) + self.assertEqual(batch_size, 32 // replica_scale_factor) + self.assertEqual(steps, 2) - model.predict(inputs) - # with steps - model.predict(inputs, steps=2) - # with batch_size - model.predict(inputs, batch_size=8) + # All samples can not be consumed in specified number of steps + with self.assertRaisesRegexp(ValueError, 'not divisible by steps'): + distributed_training_utils.get_input_params( + distribution, input_63_samples, steps=2, batch_size=None) + + # This cases is different for different strategies due to the + # difference in supported batch size being global or per-replica. + if replica_scale_factor == 1: + # Computed global batch size is correct even if not sharadable + steps, batch_size = distributed_training_utils.get_input_params( + distribution, input_63_samples, steps=3, batch_size=None) + self.assertEqual(batch_size, 21) + self.assertEqual(steps, 3) + else: + # Computed global batch size can not be sharded across replicas + with self.assertRaisesRegexp(ValueError, 'could not be sharded evenly ' + 'across the sync replicas'): + distributed_training_utils.get_input_params( + distribution, input_63_samples, steps=1, batch_size=None) + + @combinations.generate(strategy_for_numpy_input_combinations()) + def test_calculating_input_params_no_steps_with_batch_size(self, + distribution): + # Calculate the per_replica_batch_size scaling factor for strategies + # that use per_core_batch_size + replica_scale_factor = 1.0 + if not distributed_training_utils.global_batch_size_supported(distribution): + replica_scale_factor = distribution.num_replicas_in_sync - @combinations.generate(strategy_combinations()) - def test_calling_model_with_nested_numpy_arrays(self, distribution): with self.cached_session(): - a = keras.layers.Input(shape=(3,), name='input_a') - b = keras.layers.Input(shape=(3,), name='input_b') + input_64_samples = np.zeros((64, 3), dtype=np.float32) + + # Computed steps is correct for specified batch size + steps, batch_size = distributed_training_utils.get_input_params( + distribution, input_64_samples, steps=None, batch_size=16) + self.assertEqual(batch_size, 16) + self.assertEqual(steps, 4 // replica_scale_factor) + + # Computed steps is correct for specified batch size + steps, batch_size = distributed_training_utils.get_input_params( + distribution, input_64_samples, steps=None, batch_size=32) + self.assertEqual(batch_size, 32) + self.assertEqual(steps, 2 // replica_scale_factor) + + # Number of samples is not divisible by the global batch size + with self.assertRaisesRegexp(ValueError, 'not divisible by batch size'): + distributed_training_utils.get_input_params( + distribution, input_64_samples, steps=None, batch_size=20) + + # Number of samples is not divisible by the global batch size + with self.assertRaisesRegexp(ValueError, 'not divisible by batch size'): + distributed_training_utils.get_input_params( + distribution, input_64_samples, steps=None, batch_size=3) + + @combinations.generate(strategy_for_numpy_input_combinations()) + def test_calculating_input_params_with_steps_with_batch_size(self, + distribution): + with self.cached_session(): + input_64_samples = np.zeros((64, 3), dtype=np.float32) - dense = keras.layers.Dense(4, name='dense') - c = dense(a) - d = dense(b) - e = keras.layers.Dropout(0.5, name='dropout')(c) + # No change to steps and batch size if both specified and feasible + steps, batch_size = distributed_training_utils.get_input_params( + distribution, input_64_samples, steps=5, batch_size=3) + self.assertEqual(batch_size, 3) + self.assertEqual(steps, 5) - model = keras.models.Model([a, b], [d, e]) + # Number of samples is less than global batch size * steps + with self.assertRaisesRegexp(ValueError, 'less than samples required'): + distributed_training_utils.get_input_params( + distribution, input_64_samples, steps=10, batch_size=13) - optimizer = gradient_descent.GradientDescentOptimizer(learning_rate=0.001) - loss = 'mse' - model.compile(optimizer, loss, distribute=distribution) + @combinations.generate(strategy_for_numpy_input_combinations()) + def test_calling_model_with_numpy_arrays(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((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)) + + # 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) + + @combinations.generate(strategy_for_numpy_input_combinations()) + def test_calling_model_with_nested_numpy_arrays(self, distribution): + with self.cached_session(): + with distribution.scope(): + model = multi_input_output_model() + optimizer = gradient_descent.GradientDescentOptimizer( + learning_rate=0.001) + loss = 'mse' + model.compile(optimizer, loss) input_a_np = np.asarray(np.random.random((64, 3)), dtype=np.float32) - input_b_np = np.asarray(np.random.random((64, 3)), dtype=np.float32) + input_b_np = np.asarray(np.random.random((64, 5)), dtype=np.float32) inputs = [input_a_np, input_b_np] - output_d_np = np.asarray(np.random.random((64, 4)), dtype=np.float32) - output_e_np = np.asarray(np.random.random((64, 4)), dtype=np.float32) + output_d_np = np.asarray(np.random.random((64, 7)), dtype=np.float32) + output_e_np = np.asarray(np.random.random((64, 7)), dtype=np.float32) targets = [output_d_np, output_e_np] # Call fit with validation data @@ -517,19 +623,61 @@ class TestDistributionStrategyWithNumpyArrays(test.TestCase, # with batch_size model.predict(inputs, batch_size=8) + @combinations.generate(combinations.combine( + distribution=strategies_minus_tpu, mode=['graph'])) + def test_numpy_with_sample_weights(self, distribution): + with self.cached_session(): + with distribution.scope(): + model = get_model() + optimizer = rmsprop.RMSPropOptimizer(learning_rate=0.001) + loss = 'mse' + model.compile(optimizer, loss) + + inputs = np.zeros((20, 3), np.float32) + targets = np.zeros((20, 4), np.float32) + sample_weights = np.ones((20), np.float32) + + model.fit(inputs, targets, sample_weight=sample_weights, epochs=1, + steps_per_epoch=2, verbose=1) + + @combinations.generate(strategy_for_numpy_input_combinations()) + def test_flatten_predict_outputs(self, distribution): + with self.cached_session(): + with distribution.scope(): + model = multi_input_output_model() + optimizer = gradient_descent.GradientDescentOptimizer( + learning_rate=0.001) + loss = 'mse' + model.compile(optimizer, loss) + + # We take 6 input samples with each input having a dimension of 3 or 5. + input_a_np = np.asarray(np.random.random((6, 3)), dtype=np.float32) + input_b_np = np.asarray(np.random.random((6, 5)), dtype=np.float32) + inputs = [input_a_np, input_b_np] + + outs = model.predict(inputs, steps=1) + # `predict` a list that is equal in length to the number of model outputs. + # In this test our model has two outputs and each element of `outs` + # corresponds to all the samples of one of the model outputs. + self.assertLen(outs, 2) + # Each of the output samples have a dimension of 7. We should process all + # the available input samples(6). + self.assertAllEqual([6, 7], outs[0].shape) + self.assertAllEqual([6, 7], outs[1].shape) + class TestDistributionStrategyWithDatasets(test.TestCase, parameterized.TestCase): - @combinations.generate(strategy_combinations()) + @combinations.generate(all_strategy_combinations()) def test_calling_model_on_same_dataset(self, distribution): with self.cached_session(): - model = get_model() - - optimizer = gradient_descent.GradientDescentOptimizer(0.001) - loss = 'mse' - metrics = ['mae', keras.metrics.CategoricalAccuracy()] - model.compile(optimizer, loss, metrics=metrics, distribute=distribution) + 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) dataset = get_dataset(distribution) @@ -540,32 +688,68 @@ class TestDistributionStrategyWithDatasets(test.TestCase, validation_data=dataset, validation_steps=2) model.predict(get_predict_dataset(distribution), steps=2) - # TODO(priyag): Enable this test for TPU. Currently tuples/dict don't work - # as clone_model's input_tensors argument only seems to accept list and not - # tuples or dict. - def test_fit_with_tuple_and_dict_dataset_inputs(self): + @combinations.generate(all_strategy_combinations()) + def test_model_interleaved_eval_same_as_direct_eval(self, distribution): with self.cached_session(): - a = keras.layers.Input(shape=(3,), name='input_a') - b = keras.layers.Input(shape=(3,), name='input_b') + with distribution.scope(): + user_controlled_model = get_model() + user_controlled_model.compile( + gradient_descent.GradientDescentOptimizer(0.001), + loss='mse', + metrics=['mae', keras.metrics.CategoricalAccuracy()]) + + interleaved_model = get_model() + interleaved_model.set_weights(user_controlled_model.get_weights()) + interleaved_model.compile( + gradient_descent.GradientDescentOptimizer(0.001), + loss='mse', + metrics=['mae', keras.metrics.CategoricalAccuracy()]) + + dataset = get_dataset(distribution) - dense = keras.layers.Dense(4, name='dense') - c = dense(a) - d = dense(b) - e = keras.layers.Dropout(0.5, name='dropout')(c) + # Call fit with validation interleaved + interleaved_output = interleaved_model.fit( + dataset, epochs=2, steps_per_epoch=2, verbose=1, + validation_data=dataset, validation_steps=2, shuffle=False) + + # Manually control the validation running after each epoch. + user_controlled_output = [] + for _ in range(2): + user_controlled_model.fit( + dataset, epochs=1, steps_per_epoch=2, verbose=1, shuffle=False) + user_controlled_output.append( + user_controlled_model.evaluate(dataset, steps=2)) + + self.assertEqual(interleaved_output.history['val_loss'], + [x[0] for x in user_controlled_output]) + self.assertEqual(interleaved_output.history['val_mean_absolute_error'], + [x[1] for x in user_controlled_output]) + self.assertEqual(interleaved_output.history['val_categorical_accuracy'], + [x[2] for x in user_controlled_output]) - model = keras.models.Model([a, b], [d, e]) + # TODO(priyag): Enable this test for TPU. Currently tuples/dict don't work + # as clone_model's input_tensors argument only seems to accept list and not + # tuples or dict. - optimizer = gradient_descent.GradientDescentOptimizer(learning_rate=0.001) - loss = 'mse' - metrics = ['mae', keras.metrics.CategoricalAccuracy()] - strategy = mirrored_strategy.MirroredStrategy(['/device:GPU:0', - '/device:CPU:0']) - model.compile(optimizer, loss, metrics=metrics, distribute=strategy) + @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_fit_with_tuple_and_dict_dataset_inputs(self, distribution): + with self.cached_session(): + with distribution.scope(): + model = multi_input_output_model() + optimizer = gradient_descent.GradientDescentOptimizer( + learning_rate=0.001) + loss = 'mse' + metrics = ['mae', keras.metrics.CategoricalAccuracy()] + model.compile(optimizer, loss, metrics=metrics) input_a_np = np.random.random((10, 3)) - input_b_np = np.random.random((10, 3)) - output_d_np = np.random.random((10, 4)) - output_e_np = np.random.random((10, 4)) + input_b_np = np.random.random((10, 5)) + output_d_np = np.random.random((10, 7)) + output_e_np = np.random.random((10, 7)) # Test with tuples dataset_tuple = dataset_ops.Dataset.from_tensor_slices(( @@ -584,15 +768,15 @@ class TestDistributionStrategyWithDatasets(test.TestCase, model.fit(dataset_dict, epochs=1, steps_per_epoch=2, verbose=1) - @combinations.generate(strategy_combinations()) + @combinations.generate(all_strategy_combinations()) def test_fit_eval_and_predict_methods_on_dataset(self, distribution): with self.cached_session(): - model = get_model() - - optimizer = gradient_descent.GradientDescentOptimizer(0.001) - loss = 'mse' - metrics = ['mae', keras.metrics.CategoricalAccuracy()] - model.compile(optimizer, loss, metrics=metrics, distribute=distribution) + 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) dataset = get_dataset(distribution) @@ -603,10 +787,10 @@ class TestDistributionStrategyWithDatasets(test.TestCase, @combinations.generate(strategy_and_optimizer_combinations()) def test_fit_eval_and_predict_with_optimizer(self, distribution, optimizer): with self.cached_session(): - model = get_model() - - loss = 'mse' - model.compile(optimizer(), loss, distribute=distribution) + with distribution.scope(): + model = get_model() + loss = 'mse' + model.compile(optimizer(), loss) dataset = get_dataset(distribution) @@ -614,35 +798,73 @@ class TestDistributionStrategyWithDatasets(test.TestCase, model.evaluate(dataset, steps=2, verbose=1) model.predict(get_predict_dataset(distribution), steps=2) - def test_dataset_input_shape_validation(self): + @combinations.generate(strategy_minus_tpu_combinations()) + def test_dataset_with_sample_weights(self, distribution): with self.cached_session(): - model = get_model() + with distribution.scope(): + model = get_model() + optimizer = rmsprop.RMSPropOptimizer(learning_rate=0.001) + loss = 'mse' + model.compile(optimizer, loss) + + inputs = np.zeros((10, 3), np.float32) + targets = np.zeros((10, 4), np.float32) + sample_weights = np.ones((10), np.float32) + dataset = dataset_ops.Dataset.from_tensor_slices((inputs, targets, + sample_weights)) + dataset = dataset.repeat() + dataset = dataset.batch(10) - optimizer = rmsprop.RMSPropOptimizer(learning_rate=0.001) - loss = 'mse' - strategy = mirrored_strategy.MirroredStrategy(['/device:GPU:1', - '/device:GPU:0']) + model.fit(dataset, epochs=1, steps_per_epoch=2, verbose=1) + model.evaluate(dataset, steps=2, verbose=1) + model.predict(dataset, steps=2) - model.compile(optimizer, loss, distribute=strategy) + @combinations.generate(combinations.combine( + distribution=[ + combinations.mirrored_strategy_with_gpu_and_cpu, + combinations.core_mirrored_strategy_with_gpu_and_cpu], + mode=['graph', 'eager'])) + # TODO(b/120943676, b/120957836): Re-enable once the validation code is + # restored. + def DISABLED_test_dataset_wrong_input_shape(self, distribution): + with self.cached_session(): + with distribution.scope(): + model = get_model() + optimizer = rmsprop.RMSPropOptimizer(learning_rate=0.001) + loss = 'mse' + model.compile(optimizer, loss) - # User forgets to batch the dataset - inputs = np.zeros((10, 3), dtype=np.float32) + # Wrong input shape + inputs = np.zeros((10, 5), dtype=np.float32) targets = np.zeros((10, 4), dtype=np.float32) dataset = dataset_ops.Dataset.from_tensor_slices((inputs, targets)) dataset = dataset.repeat(100) + dataset = dataset.batch(10) - with self.assertRaisesRegexp(ValueError, 'expected input to have shape'): + with self.assertRaisesRegexp(ValueError, + 'expected input to have shape'): model.fit(dataset, epochs=1, steps_per_epoch=2, verbose=0) - # Wrong input shape - inputs = np.zeros((10, 5), dtype=np.float32) + @combinations.generate(combinations.combine( + distribution=[combinations.mirrored_strategy_with_gpu_and_cpu], + mode=['graph', 'eager'])) + # TODO(b/120943676, b/120957836): Re-enable once the validation code is + # restored. + def DISABLED_test_dataset_no_batch_input_validation(self, distribution): + with self.cached_session(): + with distribution.scope(): + model = get_model() + optimizer = rmsprop.RMSPropOptimizer(learning_rate=0.001) + loss = 'mse' + model.compile(optimizer, loss) + + # User forgets to batch the dataset + inputs = np.zeros((10, 3), dtype=np.float32) targets = np.zeros((10, 4), dtype=np.float32) dataset = dataset_ops.Dataset.from_tensor_slices((inputs, targets)) dataset = dataset.repeat(100) - dataset = dataset.batch(10) - with self.assertRaisesRegexp(ValueError, - 'expected input to have shape'): + with self.assertRaisesRegexp(ValueError, 'expected input to have shape'): model.fit(dataset, epochs=1, steps_per_epoch=2, verbose=0) @combinations.generate(combinations.combine( @@ -650,11 +872,11 @@ class TestDistributionStrategyWithDatasets(test.TestCase, mode=['graph'])) def test_dataset_input_shape_fully_defined(self, distribution): with self.cached_session(): - model = get_model() - - optimizer = rmsprop.RMSPropOptimizer(learning_rate=0.001) - loss = 'mse' - model.compile(optimizer, loss, distribute=distribution) + with distribution.scope(): + model = get_model() + optimizer = rmsprop.RMSPropOptimizer(learning_rate=0.001) + loss = 'mse' + model.compile(optimizer, loss) dataset = get_dataset(distribution) # Input shapes are not fully known. Batch dimension is unknown as we are @@ -664,97 +886,141 @@ class TestDistributionStrategyWithDatasets(test.TestCase, with self.assertRaisesRegexp(ValueError, 'requires fully defined shapes'): model.fit(dataset, epochs=1, steps_per_epoch=2, verbose=0) - def test_learning_phase_value(self): + @combinations.generate(combinations.combine( + distribution=[ + 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'])) + def test_learning_phase_value(self, distribution): # TODO(anjalisridhar): Modify this test to use Lambdas since we can compare # meaningful values. Currently we don't pass the learning phase if the # Lambda layer uses the learning phase. with self.cached_session(): - x = keras.layers.Input(shape=(1,), name='input') - y = keras.layers.Dense(1, kernel_initializer='ones')(x) - z = keras.layers.Dropout(0.9999)(y) - model = keras.Model(x, z) - initial_weights = model.get_weights() - - optimizer = gradient_descent.GradientDescentOptimizer(0.005) - loss = 'mse' - metrics = ['acc'] - strategy = mirrored_strategy.MirroredStrategy( - ['/device:GPU:0', '/device:GPU:1']) - - model.compile(optimizer, loss, metrics=metrics, distribute=strategy) + with distribution.scope(): + x = keras.layers.Input(shape=(1,), name='input') + y = keras.layers.Dense(1, kernel_initializer='ones')(x) + z = keras.layers.Dropout(0.9999)(y) + model = keras.Model(x, z) + initial_weights = model.get_weights() + + optimizer = gradient_descent.GradientDescentOptimizer(0.005) + loss = 'mse' + metrics = ['acc'] + model.compile(optimizer, loss, metrics=metrics) + + batch_size = 8 + if isinstance(distribution, mirrored_strategy.CoreMirroredStrategy): + # CoreMirroredStrategy uses global batch size. + batch_size = 8 * distribution.num_replicas_in_sync inputs = np.ones((10, 1), dtype=np.float32) targets = np.ones((10, 1), dtype=np.float32) dataset = dataset_ops.Dataset.from_tensor_slices((inputs, targets)) - dataset = dataset.repeat().batch(8) + dataset = dataset.repeat().batch(batch_size) hist = model.fit(dataset, epochs=1, steps_per_epoch=20, verbose=1) self.assertAlmostEqual(hist.history['acc'][0], 0, 0) - model.set_weights(initial_weights) - evaluate_output = model.evaluate(dataset, steps=20) - self.assertAlmostEqual(evaluate_output[1], 1, 0) + with distribution.scope(): + model.set_weights(initial_weights) + # TODO(psv/anjalisridhar): Enable these lines after we fix b/117431185. + # evaluate_output = model.evaluate(dataset, steps=20) + # self.assertAlmostEqual(evaluate_output[1], 1, 0) inputs = np.ones((10, 1), dtype=np.float32) predict_dataset = dataset_ops.Dataset.from_tensor_slices(inputs) - predict_dataset = predict_dataset.repeat().batch(5) + + predict_dataset = predict_dataset.repeat().batch(batch_size) output = model.predict(predict_dataset, steps=10) - ref_output = np.ones((50, 1), dtype=np.float32) - self.assertArrayNear(output[0], ref_output, 1e-1) + # `predict` runs for 10 steps + ref_output = np.ones((160, 1), dtype=np.float32) + self.assertArrayNear(output, ref_output, 1e-1) + + @combinations.generate(all_strategy_combinations()) + def testOptimizerWithCallbacks(self, distribution): + with self.cached_session(): + with distribution.scope(): + model = get_model() + optimizer = gradient_descent_keras.SGD(0.01) + loss = 'mse' + model.compile(optimizer, loss) + + dataset = get_dataset(distribution) + + def schedule(_): + return 0.001 + + model.fit(dataset, epochs=1, steps_per_epoch=2, verbose=0, + callbacks=[keras.callbacks.LearningRateScheduler(schedule)]) + self.assertAllClose(0.001, keras.backend.get_value(model.optimizer.lr)) class TestDistributionStrategyErrorCases(test.TestCase, parameterized.TestCase): - def test_validating_dataset_input_tensors_with_shape_mismatch(self): + @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(): - strategy = mirrored_strategy.MirroredStrategy(['/device:GPU:0', - '/device:CPU:0']) a = constant_op.constant([1, 2], shape=(1, 2)) b = constant_op.constant([[1, 2], [1, 2]], shape=(2, 2)) - x = values.DistributedValues({'/device:CPU:0': a, '/device:GPU:0': b}) - y = values.DistributedValues({'/device:CPU:0': a, '/device:GPU:0': a}) - with strategy.scope(): - # 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:.+'): + 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( - strategy, x, y) + distribution, x, y) - def test_validating_dataset_input_tensors_with_dtype_mismatch(self): + @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(): - strategy = mirrored_strategy.MirroredStrategy(['/device:GPU:0', - '/device:CPU:0']) a = constant_op.constant([1, 2], shape=(1, 2), dtype=dtypes.int32) b = constant_op.constant([1, 2], shape=(1, 2), dtype=dtypes.float64) - x = values.DistributedValues({'/device:CPU:0': a, '/device:GPU:0': b}) - y = values.DistributedValues({'/device:CPU:0': a, '/device:GPU:0': a}) - with strategy.scope(): - # 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:.+'): + 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( - strategy, x, y) + distribution, x, y) - def test_unsupported_features(self): + @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(): - model = get_model() + with distribution.scope(): + model = get_model() + optimizer = gradient_descent.GradientDescentOptimizer(0.001) + loss = 'mse' + metrics = ['mae'] + model.compile(optimizer, loss, metrics=metrics) - optimizer = gradient_descent.GradientDescentOptimizer(0.001) - loss = 'mse' - metrics = ['mae'] - strategy = mirrored_strategy.MirroredStrategy(['/device:GPU:1', - '/device:GPU:0']) - - model.compile(optimizer, loss, metrics=metrics, distribute=strategy) - - dataset = get_dataset(strategy) + dataset = get_dataset(distribution) # Test with validation split with self.assertRaisesRegexp( @@ -768,8 +1034,8 @@ class TestDistributionStrategyErrorCases(test.TestCase, parameterized.TestCase): # Test with sample weight. sample_weight = np.random.random((10,)) with self.assertRaisesRegexp( - NotImplementedError, '`sample_weight` is currently not supported ' - 'when using DistributionStrategy.'): + ValueError, '`sample_weight` argument is not supported when input ' + '`x` is a dataset or a dataset iterator.'): model.fit( dataset, epochs=1, @@ -779,69 +1045,67 @@ class TestDistributionStrategyErrorCases(test.TestCase, parameterized.TestCase): # Test with not specifying the `steps` argument. with self.assertRaisesRegexp( - ValueError, 'you should specify the `steps_per_epoch` argument'): + ValueError, 'the `steps_per_epoch` argument'): model.fit(dataset, epochs=1, verbose=0) - with self.assertRaisesRegexp(ValueError, - 'you should specify the `steps` argument'): + with self.assertRaisesRegexp(ValueError, 'the `steps` argument'): model.evaluate(dataset, verbose=0) - with self.assertRaisesRegexp(ValueError, - 'you should specify the `steps` argument'): + with self.assertRaisesRegexp(ValueError, 'the `steps` argument'): model.predict(dataset, verbose=0) - def test_calling_with_unsupported_predefined_callbacks(self): + @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(): - model = get_model() - - optimizer = gradient_descent.GradientDescentOptimizer(0.001) - loss = 'mse' - metrics = ['mae'] - strategy = mirrored_strategy.MirroredStrategy(['/device:GPU:1', - '/device:GPU:0']) - model.compile(optimizer, loss, metrics=metrics, distribute=strategy) + with distribution.scope(): + model = get_model() + optimizer = gradient_descent.GradientDescentOptimizer(0.001) + loss = 'mse' + metrics = ['mae'] + model.compile(optimizer, loss, metrics=metrics) - dataset = get_dataset(strategy) + dataset = get_dataset(distribution) def schedule(_): return 0.001 with self.assertRaisesRegexp(ValueError, - 'LearningRateScheduler callback is not ' - 'supported with DistributionStrategy.'): + '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, - 'ReduceLROnPlateau callback is not ' - 'supported with DistributionStrategy.'): + 'You must specify a Keras Optimizer V2 when ' + 'using'): model.fit(dataset, epochs=1, steps_per_epoch=2, verbose=0, callbacks=[keras.callbacks.ReduceLROnPlateau()]) - with self.assertRaisesRegexp(ValueError, - 'histogram_freq in the TensorBoard callback ' - 'is not supported when using ' - 'DistributionStrategy.'): - model.fit(dataset, epochs=1, steps_per_epoch=2, verbose=0, - callbacks=[keras.callbacks.TensorBoard(histogram_freq=10)]) -class TestDistributionStrategyWithLossMasking(test.TestCase): +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. - def test_masking(self): + @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]]]) - 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'))) - strategy = mirrored_strategy.MirroredStrategy(['/device:GPU:1', - '/device:GPU:0']) - - model.compile(loss='mse', - optimizer=gradient_descent.GradientDescentOptimizer(0.01), - distribute=strategy) + 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) @@ -853,15 +1117,18 @@ class TestDistributionStrategyWithLossMasking(test.TestCase): class TestDistributionStrategyWithNormalizationLayer( test.TestCase, parameterized.TestCase): - @combinations.generate(strategy_combinations()) - def test_batchnorm_correctness(self, distribution): + @combinations.generate(combinations.times( + all_strategy_combinations(), + combinations.combine(fused=[True, False]))) + def test_batchnorm_correctness(self, distribution, fused): with self.cached_session(): - model = keras.models.Sequential() - norm = keras.layers.BatchNormalization(input_shape=(10,), momentum=0.8) - model.add(norm) - model.compile(loss='mse', - optimizer=gradient_descent.GradientDescentOptimizer(0.01), - distribute=distribution) + 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)) @@ -882,104 +1149,36 @@ class TestDistributionStrategyWithNormalizationLayer( np.testing.assert_allclose(out.std(), 1.0, atol=1e-1) -class TestDistributionStrategyCorrectness(test.TestCase, - parameterized.TestCase): +class TestDistributionStrategyValidation(test.TestCase, + parameterized.TestCase): - @combinations.generate(strategy_combinations()) - def test_metric_correctness(self, distribution): + @combinations.generate(all_strategy_combinations_minus_default()) + def test_layer_outside_scope(self, distribution): 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. - 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()], - distribute=distribution) - - batch_size = 64 - batch_size //= distribution.num_towers - 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=1, steps_per_epoch=10) - self.assertEqual(history.history['binary_accuracy'], [1.0]) - - @combinations.generate(strategy_combinations()) - def test_correctness(self, distribution): + 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(all_strategy_combinations_minus_default()) + def test_model_outside_scope(self, distribution): with self.cached_session(): - keras.backend.set_image_data_format('channels_last') - num_samples = 10000 - - # Train and predict datasets are created with the same input numpy arrays. - 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') - - # 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 = keras.Sequential() - model.add(keras.layers.Dense(1, input_shape=(1,))) - initial_weights = model.get_weights() - - def fit_and_predict(with_distribution=None): - model.set_weights(initial_weights) - model.compile( - loss=keras.losses.mean_squared_error, - optimizer=gradient_descent.GradientDescentOptimizer(0.5), - distribute=with_distribution) - - batch_size = 64 - if with_distribution: - batch_size //= with_distribution.num_towers - train_dataset = dataset_ops.Dataset.from_tensor_slices((x_train, - y_train)) - train_dataset = batch_wrapper(train_dataset, batch_size, distribution) - # We have initialized the model to the same weight for the distribution - # and non-distribution run. If you want to initialize the model to - # random weights for each run, you need to run the model through the - # entire dataset at least once to ensure that the weights converge to - # the same value. - model.fit(x=train_dataset, epochs=1, steps_per_epoch=10) - - weights = model.get_weights() - x_predict = [[1.], [2.], [3.], [4.]] - predict_batch_size = 4 - if with_distribution: - predict_batch_size //= with_distribution.num_towers - predict_dataset = dataset_ops.Dataset.from_tensor_slices(x_predict) - predict_dataset = batch_wrapper(predict_dataset, - predict_batch_size, distribution) - predict_result = model.predict(predict_dataset, steps=1) - predict_result = np.reshape(predict_result, (4, 1)) - - return weights, predict_result - - wts_with_ds, predict_with_ds = fit_and_predict( - with_distribution=distribution) - wts_without_ds, predict_without_ds = fit_and_predict( - with_distribution=None) - - # Verify that the weights are the same within some limits of tolerance. - np.testing.assert_allclose(wts_with_ds[0], wts_without_ds[0], rtol=1e-3) - # Verify that the predicted outputs are the same within some limits of - # tolerance. - np.testing.assert_allclose(predict_with_ds, predict_without_ds, rtol=1e-3) - - -# TODO(priyag): Add a test for TPUStrategy with steps_per_run > 1. + 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__': diff --git a/tensorflow/contrib/distribute/python/metrics_v1_test.py b/tensorflow/contrib/distribute/python/metrics_v1_test.py index 2c79a8bfd3cf5a70c6a940b19aa3b7268ce7d524..a663e809dd45ea099e1d8a08e681d07b05bee3c9 100644 --- a/tensorflow/contrib/distribute/python/metrics_v1_test.py +++ b/tensorflow/contrib/distribute/python/metrics_v1_test.py @@ -72,14 +72,14 @@ def _regression_dataset_fn(): "predictions": [1., .75, .25, 0.]}).repeat() -# TODO(priyag): Add TPU Strategy to this once metrics aggregate correctly using -# TowerLocalVariables on TPUs. Submit http://cl/208914352. def all_combinations(): return combinations.combine( distribution=[combinations.default_strategy, combinations.one_device_strategy, combinations.mirrored_strategy_with_gpu_and_cpu, - combinations.mirrored_strategy_with_two_gpus], + combinations.mirrored_strategy_with_two_gpus, + combinations.core_mirrored_strategy_with_gpu_and_cpu, + combinations.core_mirrored_strategy_with_two_gpus], mode=["graph"]) @@ -95,33 +95,32 @@ 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_tower( - metric_fn, inputs) + 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( - step_fn, iterator, iterations=distribution.steps_per_run) + 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"] # In each run, we run multiple steps, and each steps consumes as many - # batches as number of towers. + # batches as number of replicas. batches_per_update = ( - distribution.num_towers * distribution.steps_per_run) + distribution.num_replicas_in_sync * + distribution.extended.steps_per_run) else: - value, update = distribution.call_for_each_tower( - metric_fn, iterator.get_next()) + 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_towers" with "1". - batches_per_update = distribution.num_towers + # 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 @@ -135,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 3c4544a39ef85c18a34216ba7f3ac65b45216003..f06c9b75644b2890b7657f75e74e4e20a6f15705 100644 --- a/tensorflow/contrib/distribute/python/minimize_loss_test.py +++ b/tensorflow/contrib/distribute/python/minimize_loss_test.py @@ -22,10 +22,10 @@ from absl.testing import parameterized import numpy from tensorflow.contrib.distribute.python import combinations -from tensorflow.contrib.distribute.python import mirrored_strategy from tensorflow.contrib.distribute.python.single_loss_example import batchnorm_example from tensorflow.contrib.distribute.python.single_loss_example import minimize_loss_example from tensorflow.python.data.ops import dataset_ops +from tensorflow.python.distribute import reduce_util from tensorflow.python.eager import context from tensorflow.python.eager import test from tensorflow.python.framework import constant_op @@ -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( @@ -64,19 +61,18 @@ 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) - def step_fn(ctx, *inputs): + def step_fn(ctx, inputs): del ctx # Unused return distribution.group( - distribution.call_for_each_tower( - model_fn, *inputs, run_concurrently=layer.built)) + 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()) @@ -85,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) @@ -100,18 +93,18 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase): combinations.distributions_and_v1_optimizers(), combinations.combine(mode=["graph"], use_callable_loss=[True, False]) + combinations.combine(mode=["eager"], use_callable_loss=[True]))) - def testTrainNetworkByCallForEachTower(self, distribution, optimizer_fn, - use_callable_loss): + def testTrainNetworkByCallForEachReplica(self, distribution, optimizer_fn, + use_callable_loss): with distribution.scope(): 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_tower( - model_fn, iterator.get_next(), run_concurrently=layer.built)) + distribution.extended.call_for_each_replica( + model_fn, args=(iterator.get_next(),))) if not context.executing_eagerly(): with self.cached_session() as sess: @@ -153,34 +146,30 @@ 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, create_optimizer_inside_model_fn=True) - def step_fn(ctx, *inputs): + def step_fn(ctx, inputs): del ctx # Unused return distribution.group( - distribution.call_for_each_tower( - model_fn, *inputs, run_concurrently=layer.built)) + 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"], @@ -199,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( @@ -210,47 +199,40 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase): combinations.combine( mode=["graph", "eager"], # TODO(isaprykin): Allow False here. Currently subsequent - # towers will re-execute UPDATE_OPS of previous towers. - update_ops_in_cross_tower_mode=[True])) + + # replicas will re-execute UPDATE_OPS of previous replicas. + update_ops_in_cross_replica_mode=[True])) + combinations.combine( distribution=[combinations.tpu_strategy], optimizer_fn=combinations.optimizers_v1, mode=["graph"], - update_ops_in_cross_tower_mode=[False]))) + update_ops_in_cross_replica_mode=[False]))) def testTrainNetworkWithBatchNorm(self, distribution, optimizer_fn, momentum, - renorm, update_ops_in_cross_tower_mode): - """Verifies that moving mean updates are reduced across towers.""" + renorm, update_ops_in_cross_replica_mode): + """Verifies that moving mean updates are reduced across replicas.""" with distribution.scope(): - num_towers = len(distribution.worker_devices) + num_replicas = distribution.num_replicas_in_sync model_fn, dataset_fn, batchnorm = batchnorm_example( optimizer_fn, - batch_per_epoch=num_towers, + batch_per_epoch=num_replicas, momentum=momentum, renorm=renorm, - update_ops_in_tower_mode=not update_ops_in_cross_tower_mode) + update_ops_in_replica_mode=not update_ops_in_cross_replica_mode) - # Make sure prefetching is disabled since that makes the - # specific input on each device to be non deterministic, and - # this test relies on specific input being on each device. - if isinstance(distribution, mirrored_strategy.MirroredStrategy): - self.assertFalse(distribution._prefetch_on_device) - - def step_fn(ctx, *inputs): + def step_fn(ctx, inputs): del ctx # Unused fetches = distribution.unwrap( - distribution.call_for_each_tower( - model_fn, *inputs, run_concurrently=batchnorm.built)) - if update_ops_in_cross_tower_mode: - fetches += ops.get_collection(ops.GraphKeys.UPDATE_OPS) + 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()) @@ -260,24 +242,22 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase): def averaged_batch_mean(i): # Each batch has shape [16, 8] where the ith element in jth list is - # (8 * j + i + tower_id * 100). So the batch mean in each tower is - # (60 + i + tower_id * 100). So here comes its batch mean over all - # towers: - return 60. + i + (num_towers - 1.) / 2. * 100. + # (8 * j + i + replica_id * 100). So the batch mean in each replica is + # (60 + i + replica_id * 100). So here comes its batch mean over all + # replicas: + return 60. + i + (num_replicas - 1.) / 2. * 100. for _ in range(10): run_step() moving_means = self.evaluate(batchnorm.moving_mean) # We make sure that the moving_mean is updated as if the sample mean is - # calculated over all towers. + # calculated over all replicas. for i, expected_moving_mean in enumerate(expected_moving_means): expected_moving_means[i] -= (( 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( @@ -295,7 +275,9 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase): distribution=[ combinations.one_device_strategy, combinations.mirrored_strategy_with_gpu_and_cpu, - combinations.mirrored_strategy_with_two_gpus + combinations.mirrored_strategy_with_two_gpus, + combinations.core_mirrored_strategy_with_gpu_and_cpu, + combinations.core_mirrored_strategy_with_two_gpus ]), combinations.combine( mode=["graph"], use_callable_loss=[True, False]) + @@ -309,8 +291,8 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase): with distribution.scope(): all_vars = [] - def model_fn(x, y): - + def model_fn(inputs): + x, y = inputs def loss_fn(): # Use fixed initialization to make the steps deterministic. w = variable_scope.get_variable("w", initializer=[[2.]]) @@ -331,19 +313,18 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase): labels = dataset_ops.Dataset.from_tensors([[6.], [21.]]) return dataset_ops.Dataset.zip((features, labels)).repeat() - def step_fn(ctx, x, y): + def step_fn(ctx, inputs): del ctx # Unused return distribution.group( - distribution.call_for_each_tower( - model_fn, x, y, run_concurrently=False)) + 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()) @@ -352,7 +333,7 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase): run_step() v = all_vars[0] - self.assertTrue(all([v is vi for vi in all_vars[1:]])) + self.assertTrue(all(v is vi for vi in all_vars[1:])) weight = numpy.squeeze(self.evaluate(v)) # Our model is: # predict = x * w @@ -369,16 +350,15 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase): # So unreplicated the update to w with lr=0.2 is -0.2 * -106 = 21.2 # with sum loss reduction, or 10.6 with mean. if loss_reduction == losses_impl.Reduction.SUM: - # Note that the "distribution.num_towers" factor will go away once - # we split the input across towers, instead of pulling a complete - # batch of input per tower. - self.assertNear(weight, 2 + 21.2 * distribution.num_towers, 0.0001) + # Note that the "distribution.num_replicas_in_sync" factor will go away + # once we split the input across replicas, instead of pulling a complete + # batch of input per replica. + self.assertNear(weight, 2 + 21.2 * distribution.num_replicas_in_sync, + 0.0001) else: # 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,60 +392,59 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase): train_op = optimizer.minimize(loss_fn) loss = loss_fn() output_context.set_last_step_output( - name="tower_loss_agg", + name="replica_loss_reduced", output=loss, - aggregation=variables_lib.VariableAggregation.MEAN) + reduce_op=reduce_util.ReduceOp.MEAN) output_context.set_non_tensor_output(key1, value1) return (train_op, loss) - def step_fn(output_context, *inputs): - (train_op, loss) = distribution.call_for_each_tower( - model_fn, output_context, *inputs, run_concurrently=False) + def step_fn(output_context, inputs): + (train_op, loss) = distribution.extended.call_for_each_replica( + model_fn, args=(output_context, inputs)) output_context.set_last_step_output( - name="cross_tower_loss_agg", + name="cross_replica_loss_reduced", output=loss, - aggregation=variables_lib.VariableAggregation.MEAN) + reduce_op=reduce_util.ReduceOp.MEAN) output_context.set_last_step_output( - name="cross_tower_loss_noagg", + name="cross_replica_loss_not_reduced", 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) - # Initial values corresponding to aggregated losses are just single - # tensors. But for non aggregated losses, we need to have initial + # Initial values corresponding to reduced losses are just single + # tensors. But for non reduced losses, we need to have initial # values that are of the same structure as non reduced losses. In # MirroredStrategy, this will be a list of losses, in TPUStrategy # it will be single tensor. Using `broadcast` followed by `unwrap` # gives us the desired initial value structure. initial_loop_values = { - "tower_loss_agg": initial_loss(), - "cross_tower_loss_agg": initial_loss(), - "cross_tower_loss_noagg": + "replica_loss_reduced": initial_loss(), + "cross_replica_loss_reduced": initial_loss(), + "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) - self.assertEqual({key1: [value1]}, ctx.non_tensor_outputs) + self.assertEqual({key1: (value1,)}, ctx.non_tensor_outputs) self._verify_loss_output( initial_loss(), - loss_output=ctx.last_step_outputs["tower_loss_agg"], - aggregated=True, distribution=distribution) + loss_output=ctx.last_step_outputs["replica_loss_reduced"], + reduced=True, distribution=distribution) self._verify_loss_output( initial_loss(), - loss_output=ctx.last_step_outputs["cross_tower_loss_agg"], - aggregated=True, distribution=distribution) + loss_output=ctx.last_step_outputs["cross_replica_loss_reduced"], + reduced=True, distribution=distribution) self._verify_loss_output( initial_loss(), - loss_output=ctx.last_step_outputs["cross_tower_loss_noagg"], - aggregated=False, distribution=distribution) - return (ctx.run_op, ctx.last_step_outputs["tower_loss_agg"]) + loss_output=ctx.last_step_outputs["cross_replica_loss_not_reduced"], + 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()) @@ -478,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) @@ -488,18 +465,16 @@ class MinimizeLossStepTest(test.TestCase, parameterized.TestCase): error_is_not_increasing = all(y <= x for x, y in zip(error, error[1:])) self.assertTrue(error_is_not_increasing) - def _verify_loss_output(self, initial_loss, loss_output, aggregated, + def _verify_loss_output(self, initial_loss, loss_output, reduced, distribution): - if not aggregated: - self.assertEqual(distribution.num_towers, - len(distribution.unwrap(loss_output))) - loss_output = distribution.reduce( - aggregation=variables_lib.VariableAggregation.MEAN, - value=loss_output, destinations="/device:CPU:0") - - unwrapped_output = distribution.unwrap(loss_output) - self.assertEqual(1, len(unwrapped_output)) - loss_tensor = unwrapped_output[0] + if not reduced: + self.assertLen(distribution.unwrap(loss_output), + distribution.num_replicas_in_sync) + loss_tensor = distribution.reduce(reduce_util.ReduceOp.MEAN, loss_output) + else: + unwrapped_output = distribution.unwrap(loss_output) + self.assertLen(unwrapped_output, 1) + loss_tensor = unwrapped_output[0] self.assertEqual(initial_loss.dtype, loss_tensor.dtype) self.assertEqual(initial_loss.shape, loss_tensor.shape) diff --git a/tensorflow/contrib/distribute/python/mirrored_strategy.py b/tensorflow/contrib/distribute/python/mirrored_strategy.py index 2b128935a43637f84e4cbcc324130b931e9a8dae..5391e083fc9b3ed99cc64bbed11bdeb8dea07f93 100644 --- a/tensorflow/contrib/distribute/python/mirrored_strategy.py +++ b/tensorflow/contrib/distribute/python/mirrored_strategy.py @@ -12,301 +12,33 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Class MirroredStrategy implementing DistributionStrategy.""" +"""Contrib version of MirroredStrategy.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function -import contextlib -from functools import partial -import threading +from tensorflow.python.distribute import distribute_lib +from tensorflow.python.distribute import input_lib +from tensorflow.python.distribute import mirrored_strategy -from tensorflow.contrib.distribute.python import cross_tower_ops as cross_tower_ops_lib -from tensorflow.contrib.distribute.python import shared_variable_creator -from tensorflow.contrib.distribute.python import values -from tensorflow.python import pywrap_tensorflow -from tensorflow.python.distribute import multi_worker_util -from tensorflow.python.eager import context -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 ops -from tensorflow.python.ops import array_ops -from tensorflow.python.ops import control_flow_ops -from tensorflow.python.ops import variable_scope -from tensorflow.python.ops import variables as variables_lib -from tensorflow.python.training import coordinator -from tensorflow.python.training import device_util -from tensorflow.python.training import distribute as distribute_lib -from tensorflow.python.util import nest - -# TODO(josh11b): Replace asserts in this file with if ...: raise ... - - -@contextlib.contextmanager -def _enter_graph(g): - if context.executing_eagerly(): - with g.as_default(), context.eager_mode(): - yield - else: - with g.as_default(): - yield - - -def _cpu_device(device): - cpu_device = tf_device.DeviceSpec.from_string(device) - cpu_device.merge_from(tf_device.DeviceSpec(device_type="CPU", device_index=0)) - return cpu_device.to_string() - - -class _RequestedStop(Exception): - pass - - -# _call_for_each_tower and _reduce_non_distributed_value are not members of -# MirroredStrategy so that they are generally not allowed to use anything -# specific to MirroredStrategy and thus can be shared with other distribution -# strategies. - - -# TODO(yuefengz): maybe create a common class for those who need to call this -# _call_for_each_tower. -def _call_for_each_tower(distribution, fn, *args, **kwargs): - """Run `fn` in separate threads, once per tower/worker device. - - Args: - distribution: the DistributionStrategy object. - fn: function to run (will be run once per device, each in its own thread). - *args: positional arguments for `fn` - **kwargs: keyword arguments for `fn`. - `"run_concurrently"`: Boolean indicating whether executions of `fn` - can be run concurrently (under eager execution only), defaults to - `True`. - - Returns: - Merged return value of `fn` across all towers. - - Raises: - RuntimeError: If fn() calls get_tower_context().merge_call() a different - number of times from the available devices. - """ - run_concurrently = kwargs.pop("run_concurrently", True) - if not context.executing_eagerly(): - # Lots of TF library code isn't thread-safe in graph mode, and - # there is little to be gained by turning on multithreading when - # constructing a graph. - run_concurrently = False - # Needed for per-thread device, etc. contexts in graph mode. - ops.get_default_graph().switch_to_thread_local() - elif run_concurrently is None: - run_concurrently = True - - coord = coordinator.Coordinator(clean_stop_exception_types=(_RequestedStop,)) - - shared_variable_store = {} - - # TODO(isaprykin): Create these threads once instead of during every run() - # call. - threads = [] - for index, d in enumerate(distribution.worker_devices): - variable_creator_fn = shared_variable_creator.make_fn( - shared_variable_store, index) - t = MirroredStrategy._MirroredTowerThread( # pylint: disable=protected-access - distribution, coord, d, variable_creator_fn, fn, - *values.select_device(d, args), **values.select_device(d, kwargs)) - threads.append(t) - - for t in threads: - t.start() - - # When `fn` starts `should_run` event is set on _MirroredTowerThread - # (`MTT`) threads. The execution waits until - # `MTT.has_paused` is set, which indicates that either `fn` is - # complete or a `get_tower_context().merge_call()` is called. If `fn` is - # complete, then `MTT.done` is set to True. Otherwise, arguments - # of `get_tower_context().merge_call` from all paused threads are grouped - # and the `merge_fn` is performed. Results of the - # `get_tower_context().merge_call` are then set to `MTT.merge_result`. - # Each such `get_tower_context().merge_call` call returns the - # `MTT.merge_result` for that thread when `MTT.should_run` event - # is reset again. Execution of `fn` resumes. - - try: - with coord.stop_on_exception(): - all_done = False - while not all_done and not coord.should_stop(): - done = [] - if run_concurrently: - for t in threads: - t.should_run.set() - for t in threads: - t.has_paused.wait() - t.has_paused.clear() - if coord.should_stop(): - return None - done.append(t.done) - else: - for t in threads: - t.should_run.set() - t.has_paused.wait() - t.has_paused.clear() - if coord.should_stop(): - return None - done.append(t.done) - if coord.should_stop(): - return None - all_done = all(done) - if not all_done: - if any(done): - raise RuntimeError("Some towers made a different number of " - "tower_context().merge_call() calls.") - # get_tower_context().merge_call() case - merge_args = values.regroup({t.device: t.merge_args for t in threads}) - merge_kwargs = values.regroup( - {t.device: t.merge_kwargs for t in threads}) - # We capture the name_scope of the MTT when we call merge_fn - # to ensure that if we have opened a name scope in the MTT, - # it will be respected when executing the merge function. We only - # capture the name_scope from the first MTT and assume it is - # the same for all other MTTs. - mtt_captured_name_scope = threads[0].captured_name_scope - with ops.name_scope(mtt_captured_name_scope): - merge_result = threads[0].merge_fn(distribution, *merge_args, - **merge_kwargs) - for t in threads: - t.merge_result = values.select_device(t.device, merge_result) - finally: - for t in threads: - t.should_run.set() - coord.join(threads) - - return values.regroup({t.device: t.main_result for t in threads}) - - -def _reduce_non_distributed_value(distribution, aggregation, value, - destinations): - """Reduce a non-DistributedValue `value` to `destinations`.""" - if isinstance(value, values.DistributedValues): - raise ValueError("You are passing a `DistributedValue` to " - "`_reduce_non_distributed_value`, which is not allowed.") - - # If the same value is present on all towers then the PerDevice value will - # be a single value. We also handle the case when `value` is a single value - # and equal to 0. - if value == 0: - return 0 - # If the aggregation type is MEAN or ONLY_FIRST_TOWER, then this - # essentially means that the same value should be on all destinations. - if aggregation in ( - variable_scope.VariableAggregation.MEAN, - variable_scope.VariableAggregation.ONLY_FIRST_TOWER): - return value - - cross_tower_ops_lib.validate_destinations(destinations) - # We do not support an aggregation type of SUM if the value is the same across - # all towers. We call this as part of assign functions for MirroredVariables - # and summing up identical values across towers is not clearly defined. - if (len(distribution.worker_devices) != 1 or - not cross_tower_ops_lib.check_destinations(destinations)): - raise ValueError("A non-DistributedValues value %s cannot be reduced with " - "the given aggregation %s." % (value, aggregation)) - # TODO(anjalisridhar): Moves these methods to a device utility file? - devices = cross_tower_ops_lib.get_devices_from(destinations) - if len(devices) == 1: - with ops.device(devices[0]): - return array_ops.identity(value) - else: - value_updates = {} - for d in devices: - with ops.device(d): - value_updates[d] = array_ops.identity(value) - return values.Mirrored(value_updates) - - -def _create_mirrored_variable(devices, real_mirrored_creator, *args, **kwargs): # pylint: disable=g-missing-docstring - # Figure out what collections this variable should be added to. - # We'll add the MirroredVariable to those collections instead. - collections = kwargs.pop("collections", None) - if collections is None: - collections = [ops.GraphKeys.GLOBAL_VARIABLES] - kwargs["collections"] = [] - - # Get synchronization value - synchronization = kwargs.get("synchronization", - variable_scope.VariableSynchronization.ON_WRITE) - if synchronization == variable_scope.VariableSynchronization.NONE: - raise ValueError("`NONE` variable synchronization mode is not " - "supported with `Mirrored` distribution strategy. Please" - " change the `synchronization` for variable: " + - kwargs["name"]) - elif synchronization == variable_scope.VariableSynchronization.ON_READ: - # Variables that are to be synced on read are tower local. - is_tower_local = True - kwargs["trainable"] = False - elif (synchronization == variable_scope.VariableSynchronization.ON_WRITE or - synchronization == variable_scope.VariableSynchronization.AUTO): - # `AUTO` synchronization for `MirroredStrategy` is `ON_WRITE`. - is_tower_local = False - else: - raise ValueError("Invalid variable synchronization mode: " + - synchronization + " for variable: " + kwargs["name"]) - - # Get aggregation value - aggregation = kwargs.pop("aggregation", - variable_scope.VariableAggregation.NONE) - if aggregation not in ( - variable_scope.VariableAggregation.NONE, - variable_scope.VariableAggregation.SUM, - variable_scope.VariableAggregation.MEAN, - variable_scope.VariableAggregation.ONLY_FIRST_TOWER - ): - raise ValueError("Invalid variable aggregation mode: " + aggregation + - " for variable: " + kwargs["name"]) - - # Ignore user-specified caching device, not needed for mirrored variables. - kwargs.pop("caching_device", None) - - # TODO(josh11b,apassos): It would be better if variable initialization - # was never recorded on the tape instead of having to do this manually - # here. - with tape.stop_recording(): - index = real_mirrored_creator(devices, *args, **kwargs) - - if is_tower_local: - result = values.TowerLocalVariable(index, index[devices[0]], aggregation) - else: - result = values.MirroredVariable(index, index[devices[0]], 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 member variables - # to the TRAINABLE_VARIABLES collection, so we manually remove - # them and replace with the MirroredVariable. We can't set - # "trainable" to False for next_creator() since that causes functions - # like implicit_gradients to skip those variables. - if kwargs.get("trainable", True): - collections.append(ops.GraphKeys.TRAINABLE_VARIABLES) - l = g.get_collection_ref(ops.GraphKeys.TRAINABLE_VARIABLES) - for v in index.values(): - if v in l: - l.remove(v) - g.add_to_collections(collections, result) - elif ops.GraphKeys.GLOBAL_STEP in collections: - ops.add_to_collections(ops.GraphKeys.GLOBAL_STEP, result) - - return result +# pylint: disable=protected-access,invalid-name +_call_for_each_replica = mirrored_strategy._call_for_each_replica +_create_mirrored_variable = mirrored_strategy._create_mirrored_variable +all_local_devices = mirrored_strategy.all_local_devices +CoreMirroredStrategy = mirrored_strategy.MirroredStrategy +CoreMirroredExtended = mirrored_strategy.MirroredExtended +# pylint: enable=protected-access,invalid-name class MirroredStrategy(distribute_lib.DistributionStrategy): """Mirrors vars to distribute across multiple devices and machines. - This strategy uses one tower per device and sync replication for its multi-GPU - version. + *** contrib version *** + + This strategy uses one replica per device and sync replication for its + multi-GPU version. When `cluster_spec` is given by the `configure` method., it turns into the mulit-worker version that works on multiple workers with in-graph replication. @@ -314,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: @@ -330,12 +62,12 @@ class MirroredStrategy(distribute_lib.DistributionStrategy): index. They all do similar things except for one worker checkpointing model variables, writing summaries, etc. in addition to its ordinary work. - The multi-worker version of this class maps one tower to one device on a - worker. It mirrors all model variables on all towers. For example, if you have - two `worker`s and each `worker` has 4 GPUs, it will create 8 copies of the - model variables on these 8 GPUs. Then like in MirroredStrategy, each tower - performs their computation with their own copy of variables unless in - cross-tower model where variable or tensor reduction happens. + The multi-worker version of this class maps one replica to one device on a + worker. It mirrors all model variables on all replicas. For example, if you + have two `worker`s and each `worker` has 4 GPUs, it will create 8 copies of + the model variables on these 8 GPUs. Then like in MirroredStrategy, each + replica performs their computation with their own copy of variables unless in + cross-replica model where variable or tensor reduction happens. Args: devices: a list of device strings. @@ -345,497 +77,121 @@ class MirroredStrategy(distribute_lib.DistributionStrategy): num_gpus_per_worker: number of GPUs per worker. This is the same as `num_gpus` and only one of `num_gpus` and `num_gpus_per_worker` can be specified. - cross_tower_ops: optional, a descedant of `CrossTowerOps`. If this is not + cross_device_ops: optional, a descedant of `CrossDeviceOps`. If this is not set, the `configure` method will try to find the best one. - prefetch_on_device: optional boolean to specify whether to prefetch input - data to devices. auto_shard_dataset: whether to auto-shard the dataset when there are multiple workers. + cross_tower_ops: Deprecated alias for `cross_device_ops`. """ def __init__(self, devices=None, num_gpus=None, num_gpus_per_worker=None, - cross_tower_ops=None, - prefetch_on_device=None, - auto_shard_dataset=False): - super(MirroredStrategy, self).__init__() - - self._cross_tower_ops = cross_tower_ops - self._prefetch_on_device = prefetch_on_device - self._auto_shard_dataset = auto_shard_dataset - # Remember num GPUs which might be needed by `configure` method. + cross_device_ops=None, + auto_shard_dataset=False, + cross_tower_ops=None): + assert not (cross_device_ops and cross_tower_ops) if num_gpus is not None and num_gpus_per_worker is not None: raise ValueError( "You cannot specify both `num_gpus` and `num_gpus_per_worker`.") - if num_gpus is not None: - self._num_gpus = num_gpus - else: - self._num_gpus = num_gpus_per_worker - - self._initialize_local(self._num_gpus, devices) - - def _initialize_local(self, num_gpus, devices): - """Initializes the object for local training.""" - self._cluster_spec = None - # Convert `num_gpus` into `devices`, shouldn't specify both. - if devices is None: - if num_gpus is None: - num_gpus = context.num_gpus() - if num_gpus == 0: - devices = ["/device:CPU:0"] - else: - devices = ["/device:GPU:%d" % d for d in range(num_gpus)] - elif num_gpus is not None: - raise ValueError("Must only specify one of `devices` and `num_gpus`.") - self._num_gpus = num_gpus - # TODO(yuefengz): consider setting the default device. - - assert devices, "Must specify at least one device." - assert len(set(devices)) == len(devices), ( - "No duplicates allowed in `devices` argument.") - # TODO(josh11b): Require at least 2 devices? - self._devices = [device_util.resolve(d) for d in devices] - self._canonical_device_set = set(self._devices) - self._device_index = values.PerDevice({d: i for i, d in enumerate(devices)}) - - def _initialize_multi_worker(self, num_gpus, cluster_spec): - """Initializes the object for multi-worker training.""" - cluster_spec = multi_worker_util.normalize_cluster_spec(cluster_spec) - self._cluster_spec = cluster_spec - - self._workers = [] - for job in ["chief", "worker"]: - for task in range(len(cluster_spec.as_dict().get(job, []))): - self._workers.append("/job:%s/task:%d" % (job, task)) - if num_gpus is None: - raise ValueError("`num_gpus` is required if `cluster_spec` is given.") - if num_gpus > 0: - self._worker_device_map = { - worker: [ - device_util.canonicalize(worker + "/device:GPU:%d" % gpu) - for gpu in range(num_gpus) - ] for worker in self._workers - } - else: - self._worker_device_map = { - worker: [device_util.canonicalize(worker, "/device:CPU:0")] - for worker in self._workers - } - - devices = nest.flatten(self._worker_device_map) - - # Setting `_default_device` will add a device scope in the - # distribution.scope. We set the default device to the first worker. When - # users specify device under distribution.scope by - # with tf.device("/cpu:0"): - # ... - # their ops will end up on the cpu device of its first worker, e.g. - # "/job:worker/task:0/device:CPU:0". Note this is not used in tower mode. - self._default_device = self._workers[0] - - assert devices, "Must specify at least one device." - assert len(set(devices)) == len(devices), ( - "No duplicates allowed in `devices` argument.") - # TODO(josh11b): Require at least 2 devices? - self._devices = [device_util.resolve(d) for d in devices] - self._canonical_device_set = set(self._devices) - self._device_index = values.PerDevice( - {d: i for i, d in enumerate(devices)}) - - def _create_variable(self, next_creator, *args, **kwargs): - """Create a mirrored variable. See `DistributionStrategy.scope`.""" - colocate_with = kwargs.pop("colocate_with", None) - devices = self._get_devices_from(colocate_with) - - def _real_mirrored_creator(devices, *args, **kwargs): # pylint: disable=g-missing-docstring - index = {} - for i, d in enumerate(devices): - with ops.device(d): - if i > 0: - # Give replicas meaningful distinct names: - var0name = index[devices[0]].name.split(":")[0] - # We append a / to variable names created on towers with id > 0 to - # ensure that we ignore the name scope and instead use the given - # name as the absolute name of the variable. - kwargs["name"] = "%s/replica_%d/" % (var0name, i) - # Initialize replicas with the same value: - def initial_value_fn(device=d): - if context.executing_eagerly(): - init_value = index[devices[0]].value() - return array_ops.identity(init_value) - else: - with ops.device(device): - init_value = index[devices[0]].initial_value - return array_ops.identity(init_value) - kwargs["initial_value"] = initial_value_fn - with context.context().device_policy(context.DEVICE_PLACEMENT_SILENT): - # Don't record operations (e.g. other variable reads) during - # variable creation. - with tape.stop_recording(): - v = next_creator(*args, **kwargs) - assert not isinstance(v, values.DistributedVariable) - index[d] = v - return index - - return _create_mirrored_variable(devices, _real_mirrored_creator, *args, - **kwargs) - - def distribute_dataset(self, dataset_fn): - if self._cluster_spec: - return values.MultiWorkerDataset( - partial(self._call_dataset_fn, dataset_fn), self._worker_device_map, - self._prefetch_on_device, self._auto_shard_dataset) - else: - return values.PerDeviceDataset( - self._call_dataset_fn(dataset_fn), self._devices, - self._prefetch_on_device) - - # TODO(priyag): Deal with OutOfRange errors once b/111349762 is fixed. - def _run_steps_on_dataset(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) + num_gpus = num_gpus_per_worker + extended = MirroredExtended(self, devices, num_gpus, + cross_device_ops or cross_tower_ops, + 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.""" - ctx = values.MultiStepContext() - def body(i, *args): - """A wrapper around `fn` to create the while loop body.""" - del args - fn_inputs = iterator.get_next() - if not isinstance(fn_inputs, tuple): - fn_inputs = (fn_inputs,) - fn_result = fn(ctx, *fn_inputs) - for (name, output) in ctx.last_step_outputs.items(): - # Convert all outputs to tensors, potentially from `DistributedValues`. - ctx.last_step_outputs[name] = self.unwrap(output) - 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 - - 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) - - for (name, aggregation) in ctx._last_step_outputs_aggregations.items(): # pylint: disable=protected-access - output = last_step_tensor_outputs_dict[name] - # For outputs that have already been aggregated, wrap them in a Mirrored - # container, else in a PerDevice container. - if aggregation is variables_lib.VariableAggregation.NONE: - last_step_tensor_outputs_dict[name] = values.regroup( - {d: t for d, t in zip(self._devices, output)}, values.PerDevice) - else: - assert len(output) == 1 - last_step_tensor_outputs_dict[name] = output[0] - - ctx._set_last_step_outputs(last_step_tensor_outputs_dict) # pylint: disable=protected-access - return ctx - - def _broadcast(self, tensor, destinations): - # TODO(josh11b): In eager mode, use one thread per device, or async mode. - return self._get_cross_tower_ops().broadcast(tensor, destinations or - self._devices) - - def _call_for_each_tower(self, fn, *args, **kwargs): - return _call_for_each_tower(self, fn, *args, **kwargs) - - def map(self, map_over, fn, *args, **kwargs): - # TODO(josh11b): In eager mode, use one thread per device. - index = {} - for i, m in enumerate(map_over): - d = self._devices[i % len(self._devices)] - with ops.device(d): - l = index.get(d, []) - l.append(fn(m, - *values.select_device_mirrored(d, args), - **values.select_device_mirrored(d, kwargs))) - index[d] = l - # TODO(josh11b): Need a values.regroup equivalent that handles MapOutput - # in addition to PerDevice data. - return values.PerDevice({k: values.MapOutput(v) for k, v in index.items()}) - - def configure(self, - session_config=None, - cluster_spec=None, - task_type=None, - task_id=None): - del task_type, task_id - - if session_config: - session_config.isolate_session_state = True - - if cluster_spec: - self._initialize_multi_worker(self._num_gpus, cluster_spec) - - if self._cross_tower_ops is None: - if self._cluster_spec: - # It currently cannot detect the toplogy of remote workers. So we - # hard-code the multi-worker all-reduce algorithm for now. - if len(self._workers) == 1: - # The default is "nccl". - self._cross_tower_ops = cross_tower_ops_lib.AllReduceCrossTowerOps() - else: - # The default is hierarchical reduce and broadcast. - self._cross_tower_ops = cross_tower_ops_lib.MultiWorkerAllReduce( - self._workers, self._num_gpus) - else: - self._cross_tower_ops = cross_tower_ops_lib.choose_the_best( - self._devices, session_config=session_config) - - def _get_cross_tower_ops(self): - if self._cross_tower_ops is None: - self._cross_tower_ops = ( - cross_tower_ops_lib.ReductionToOneDeviceCrossTowerOps()) - return self._cross_tower_ops - - def _reduce(self, aggregation, value, destinations): - assert not isinstance(value, values.Mirrored) - if not isinstance(value, values.DistributedValues): - # This function handles reducing values that are not PerDevice or Mirrored - # values. For example, the same value could be present on all towers in - # which case `value` would be a single value or value could be 0. - return _reduce_non_distributed_value(self, aggregation, value, - destinations) - if aggregation == variable_scope.VariableAggregation.ONLY_FIRST_TOWER: - value = value.get(self._devices[0]) - if isinstance(value, (int, float)): - return value - return self.broadcast(value, destinations) - return self._get_cross_tower_ops().reduce( - aggregation, value, destinations=destinations) - - def _batch_reduce(self, aggregation, value_destination_pairs): - if aggregation == variable_scope.VariableAggregation.ONLY_FIRST_TOWER: - return [self.broadcast(v.get(self._devices[0]), d) - for v, d in value_destination_pairs] - return self._get_cross_tower_ops().batch_reduce(aggregation, - value_destination_pairs) - - def _update(self, var, options, fn, *args, **kwargs): - # TODO(josh11b): In eager mode, use one thread per device. - assert isinstance(var, values.DistributedVariable) - should_group = options.pop("grouped") - assert not options # Validate that we are processing all of the options. - updates = {} - for d, v in var._index.items(): # pylint: disable=protected-access - name = "update_%d" % self._device_index.get(d) - with ops.device(d), distribute_lib.UpdateContext(d), ops.name_scope(name): - # If args and kwargs are not mirrored, the value is returned as is. - updates[d] = fn(v, - *values.select_device_mirrored(d, args), - **values.select_device_mirrored(d, kwargs)) - return values.update_regroup(self, updates, should_group) - - def _update_non_slot(self, colocate_with, options, fn, *args, **kwargs): - assert isinstance(colocate_with, list) - should_group = options.pop("grouped") - assert not options # Validate that we are processing all of the options. - # TODO(josh11b): In eager mode, use one thread per device. - updates = {} - for d in colocate_with: - name = "update_%d" % self._device_index.get(d) - with ops.device(d), distribute_lib.UpdateContext(d), ops.name_scope(name): - updates[d] = fn(*values.select_device_mirrored(d, args), - **values.select_device_mirrored(d, kwargs)) - return values.update_regroup(self, updates, should_group) - - def read_var(self, tower_local_var): - """Read the aggregate value of a tower-local variable.""" - if isinstance(tower_local_var, values.TowerLocalVariable): - return tower_local_var._get_cross_tower() # pylint: disable=protected-access - assert isinstance(tower_local_var, values.Mirrored) - return array_ops.identity(tower_local_var.get()) - - def _unwrap(self, val): - if isinstance(val, values.DistributedValues): - # Return in a deterministic order. - if set(val.devices) == self._canonical_device_set: - return [val.get(device=d) for d in self._devices] - return [val.get(device=d) for d in sorted(val.devices)] - return [val] - - def value_container(self, val): - return values.value_container(val) - - @property - def is_single_tower(self): - return len(self._devices) == 1 - - @property - def num_towers(self): - return len(self._devices) - - @property - def num_replicas_in_sync(self): - return len(self._devices) + def __init__(self, + container_strategy, + devices=None, + num_gpus_per_worker=None, + cross_device_ops=None, + auto_shard_dataset=False): + if devices is None: + devices = mirrored_strategy.all_local_devices(num_gpus_per_worker) + elif num_gpus_per_worker is not None: + raise ValueError( + "Must only specify one of `devices` and `num_gpus_per_worker`.") + super(MirroredExtended, self).__init__(container_strategy, devices, + cross_device_ops) + self._auto_shard_dataset = auto_shard_dataset - def _worker_device_index(self): - return self._device_index + def _make_dataset_iterator(self, dataset): + """Make iterator from dataset without splitting the batch. - @property - def worker_devices(self): - # Make a copy to prevent users from accidentally mutating our copy. - return list(self._devices) + This implementation is different than the one in + `tf.distribute.MirroredStrategy` for purposes of backward compatibility. + We treat the incoming dataset's batch size as per replica batch size. - @property - def parameter_devices(self): - return list(self._devices) + Args: + dataset: `tf.data.Dataset` for input. + Returns: + An `InputIterator` which returns inputs for each step of the computation. + """ + return input_lib.DatasetIterator(dataset, self._input_workers) + # TODO(priyag): Delete this once all strategies use global batch size. @property - def between_graph(self): + def _global_batch_size(self): + """The contrib version of Mirrored strategy uses per-replica batch size.""" return False - - @property - def should_init(self): - return True - - @property - def should_checkpoint(self): - return True - - @property - def should_save_summary(self): - return True - - def non_slot_devices(self, var_list): - del var_list - return list(self._devices) - - def _get_devices_from(self, colocate_with=None): - if colocate_with is None: - return self._devices - else: - return cross_tower_ops_lib.get_devices_from(colocate_with) - - class _MirroredTowerThread(threading.Thread): - """A thread that runs() a function on a device.""" - - def __init__(self, dist, coord, device, variable_creator_fn, fn, *args, - **kwargs): - super(MirroredStrategy._MirroredTowerThread, self).__init__() # pylint: disable=protected-access - self.coord = coord - self.distribution = dist - self.device = device - self.tower_id = dist.worker_devices.index(device) - self.variable_creator_fn = variable_creator_fn - # State needed to run and return the results of `fn`. - self.main_fn = fn - self.main_args = args - self.main_kwargs = kwargs - self.main_result = None - self.done = False - # State needed to run the next merge_call() (if any) requested via - # TowerContext. - self.merge_fn = None - self.merge_args = None - self.merge_kwargs = None - self.merge_result = None - self.captured_name_scope = None - # We use a thread.Event for the main thread to signal when this - # thread should start running (`should_run`), and another for - # this thread to transfer control back to the main thread - # (`has_paused`, either when it gets to a - # `get_tower_context().merge_call` or when `fn` returns). In - # either case the event starts cleared, is signaled by calling - # set(). The receiving thread waits for the signal by calling - # wait() and then immediately clearing the event using clear(). - self.should_run = threading.Event() - self.has_paused = threading.Event() - # These fields have to do with inheriting various contexts from the - # parent thread: - # pylint: disable=protected-access - self.context_mode = context.context()._eager_context.mode - if not context.context()._context_handle: - context.context()._initialize_handle_and_devices() - self.context_device_policy = ( - pywrap_tensorflow.TFE_ContextGetDevicePlacementPolicy( - context.context()._context_handle)) - self.graph = ops.get_default_graph() - self._variable_creator_stack = self.graph._variable_creator_stack[:] - self._captured_var_scope = variable_scope.get_variable_scope() - # Adding a "/" at end lets us re-enter this scope later. - self._name_scope = self.graph.get_name_scope() - if self._name_scope: - self._name_scope += "/" - if self.tower_id > 0: - if not self._name_scope: - self._name_scope = "" - self._name_scope += "tower_%d/" % self.tower_id - - def run(self): - # pylint: disable=protected-access - self.graph._variable_creator_stack = self._variable_creator_stack - self.should_run.wait() - self.should_run.clear() - try: - if self.coord.should_stop(): - return - with self.coord.stop_on_exception(), \ - context.context()._mode(self.context_mode), \ - context.context().device_policy(self.context_device_policy), \ - _enter_graph(self.graph), \ - MirroredTowerContext(self.distribution, self.tower_id), \ - ops.device(self.device), \ - ops.name_scope(self._name_scope), \ - variable_scope.variable_scope( - self._captured_var_scope, reuse=self.tower_id > 0), \ - variable_scope.variable_creator_scope(self.variable_creator_fn): - self.main_result = self.main_fn(*self.main_args, **self.main_kwargs) - self.done = True - finally: - self.has_paused.set() - - -class MirroredTowerContext(distribute_lib.TowerContext): - """TowerContext used in MirroredStrategy.call_for_each_tower(). - - Opened in `_MirroredTowerThread`, to allow the user to invoke - `MirroredStrategy`'s specific implementation of `merge_call()`, - which works by delegating the function and its arguments to - the main thread (the one that invoked - `MirroredStrategy.call_for_each_tower()`). - """ - - def _merge_call(self, fn, *args, **kwargs): - """Delegate to the main thread to actually perform merge_call().""" - t = threading.current_thread() # a _MirroredTowerThread - t.merge_fn = fn - t.merge_args = args - t.merge_kwargs = kwargs - t.captured_name_scope = t.graph.get_name_scope() - # Adding a "/" at end lets us re-enter this scope later. - if t.captured_name_scope: - t.captured_name_scope += "/" - t.has_paused.set() - t.should_run.wait() - t.should_run.clear() - if t.coord.should_stop(): - raise _RequestedStop() - return t.merge_result - - @property - def device(self): - distribute_lib.require_tower_context(self) - return self._distribution_strategy.worker_devices[self._tower_id] diff --git a/tensorflow/contrib/distribute/python/mirrored_strategy_multigpu_test.py b/tensorflow/contrib/distribute/python/mirrored_strategy_multigpu_test.py index d32c04068cff07073dff9fa90b4b1ce5a0991bb1..d6337d106fced921b8bda0a2faac99c2a77fab8e 100644 --- a/tensorflow/contrib/distribute/python/mirrored_strategy_multigpu_test.py +++ b/tensorflow/contrib/distribute/python/mirrored_strategy_multigpu_test.py @@ -20,22 +20,27 @@ from __future__ import print_function import sys +from absl.testing import parameterized import numpy as np +from tensorflow.contrib.distribute.python import combinations from tensorflow.contrib.distribute.python import mirrored_strategy from tensorflow.contrib.distribute.python import multi_worker_test_base from tensorflow.contrib.distribute.python import strategy_test_lib -from tensorflow.contrib.distribute.python import values 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 distribution_strategy_context as ds_context +from tensorflow.python.distribute import reduce_util +from tensorflow.python.distribute import values from tensorflow.python.eager import backprop from tensorflow.python.eager import context from tensorflow.python.eager import 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 ops -from tensorflow.python.framework import test_util from tensorflow.python.keras.engine import training as keras_training from tensorflow.python.keras.layers import core as keras_core from tensorflow.python.layers import core @@ -46,8 +51,6 @@ from tensorflow.python.ops import rnn_cell_impl from tensorflow.python.ops import state_ops from tensorflow.python.ops import variable_scope from tensorflow.python.ops import variables -from tensorflow.python.training import device_util -from tensorflow.python.training import distribution_strategy_context from tensorflow.python.training import gradient_descent from tensorflow.python.training import optimizer as optimizer_lib from tensorflow.python.training import server_lib @@ -56,246 +59,299 @@ from tensorflow.python.training import server_lib GPU_TEST = "test_gpu" in sys.argv[0] -class MirroredTwoDeviceDistributionTest(strategy_test_lib.DistributionTestBase): +@combinations.generate(combinations.combine( + distribution=[ + 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"])) +class MirroredTwoDeviceDistributionTest( + strategy_test_lib.DistributionTestBase, + strategy_test_lib.TwoDeviceDistributionTestBase, + parameterized.TestCase): - def _get_distribution_strategy(self): - devices = ["/device:CPU:0", "/device:GPU:0"] - if GPU_TEST: - self.assertGreater(context.num_gpus(), 0) - if context.num_gpus() > 1: - devices = ["/device:GPU:0", "/device:GPU:1"] - print(self.id().split(".")[-1], "devices:", ", ".join(devices)) - return mirrored_strategy.MirroredStrategy(devices) + def testMinimizeLoss(self, distribution): + if context.executing_eagerly(): + self._test_minimize_loss_eager(distribution) + else: + self._test_minimize_loss_graph(distribution) - def testMinimizeLossEager(self): - if not GPU_TEST: - self.skipTest("Not GPU test") - self._test_minimize_loss_eager(self._get_distribution_strategy()) + def testReplicaId(self, distribution): + self._test_replica_id(distribution) - def testMinimizeLossGraph(self): - soft_placement = not GPU_TEST - print("testMinimizeLossGraph soft_placement:", soft_placement) - self._test_minimize_loss_graph( - self._get_distribution_strategy(), soft_placement=soft_placement) - - def testMapReduce(self): - if not GPU_TEST: - self.skipTest("Not GPU test") - self._test_map_reduce(self._get_distribution_strategy()) - - def testDeviceIndex(self): - if not GPU_TEST: - self.skipTest("Not GPU test") - self._test_device_index(self._get_distribution_strategy()) - - def testTowerId(self): - if not GPU_TEST: - self.skipTest("Not GPU test") - self._test_tower_id(self._get_distribution_strategy()) - - def testNumTowers(self): - if not GPU_TEST: - self.skipTest("Not GPU test") - self.assertEqual(2, self._get_distribution_strategy().num_towers) - - def testNumReplicasInSync(self): - if not GPU_TEST: - self.skipTest("Not GPU test") - self.assertEqual(2, self._get_distribution_strategy(). - num_replicas_in_sync) - - @test_util.run_in_graph_and_eager_modes - def testCallAndMergeExceptions(self): - if not GPU_TEST: - self.skipTest("Not GPU test") - self._test_call_and_merge_exceptions(self._get_distribution_strategy()) - - @test_util.run_in_graph_and_eager_modes - def testRunRegroupError(self): - - def run_fn(device_id): + def testNumReplicasInSync(self, distribution): + self.assertEqual(2, distribution.num_replicas_in_sync) + + def testCallAndMergeExceptions(self, distribution): + self._test_call_and_merge_exceptions(distribution) + + def testRunRegroupError(self, distribution): + def run_fn(): + replica_id = int(self.evaluate(_replica_id())) # Generates a list with different lengths on different devices. # Will fail in _regroup() (if more than one device). - return list(range(device_id)) - - dist = self._get_distribution_strategy() - with dist.scope(), self.assertRaises(AssertionError): - dist.call_for_each_tower(run_fn, dist.worker_device_index) - - @test_util.run_in_graph_and_eager_modes - def testReduceToCpu(self): - if not GPU_TEST: - self.skipTest("Not GPU test") - - def run_fn(device_id): - return device_id - - dist = self._get_distribution_strategy() - with dist.scope(): - result = dist.call_for_each_tower(run_fn, dist.worker_device_index) - reduced = dist.reduce( - variable_scope.VariableAggregation.SUM, - result, - destinations="/device:CPU:0") - unwrapped = dist.unwrap(reduced) - self.assertEqual(1, len(unwrapped)) - expected = sum(range(len(dist.worker_devices))) - self.assertEqual(expected, self.evaluate(unwrapped[0])) - - @test_util.run_in_graph_and_eager_modes - def testReduceOnlyFirstTowerUpdates(self): - if not GPU_TEST: - self.skipTest("Not GPU test") - - def run_fn(device_id): - return constant_op.constant(3 + 5 * device_id) - - dist = self._get_distribution_strategy() - with dist.scope(): - result = dist.call_for_each_tower(run_fn, dist.worker_device_index) - reduced = dist.reduce( - variable_scope.VariableAggregation.ONLY_FIRST_TOWER, - result, - destinations="/device:CPU:0") - unwrapped = dist.unwrap(reduced) - self.assertEqual(1, len(unwrapped)) - self.assertEqual(3, self.evaluate(unwrapped[0])) - - @test_util.run_in_graph_and_eager_modes() - def testReduceToMultipleDestinations(self): - if not GPU_TEST: - self.skipTest("Not GPU test") - - devices = ["/device:GPU:0"] - if GPU_TEST: - self.assertGreater(context.num_gpus(), 0) - print(self.id().split(".")[-1], "devices:", ", ".join(devices)) - - dist = mirrored_strategy.MirroredStrategy(devices) - with dist.scope(): - reduced = dist.reduce( - variable_scope.VariableAggregation.SUM, - 1.0, - destinations=["/device:CPU:0", "/device:GPU:0"]) - unwrapped = dist.unwrap(reduced) - self.assertEqual(2, len(unwrapped)) - self.assertEqual(1.0, self.evaluate(unwrapped[0])) + return list(range(replica_id)) + with distribution.scope(), self.assertRaises(AssertionError): + distribution.extended.call_for_each_replica(run_fn) -class MirroredStrategyVariableCreationTest(test.TestCase): + def testReduceToCpu(self, distribution): + with distribution.scope(): + result = distribution.extended.call_for_each_replica(_replica_id) + reduced = distribution.reduce(reduce_util.ReduceOp.SUM, result) + expected = sum(range(distribution.num_replicas_in_sync)) + self.assertEqual(expected, self.evaluate(reduced)) + + def testMakeInputFnIterator(self, distribution): + dataset_fn = lambda: dataset_ops.Dataset.range(10) + expected_values = [[i, i+1] for i in range(0, 10, 2)] + + input_fn = self._input_fn_to_test_input_context( + dataset_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) + + 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) - config = config_pb2.ConfigProto() - config.allow_soft_placement = True + 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( + distribution=[ + combinations.mirrored_strategy_with_one_cpu, + combinations.mirrored_strategy_with_one_gpu, + combinations.core_mirrored_strategy_with_one_cpu, + combinations.core_mirrored_strategy_with_one_gpu], + mode=["graph", "eager"]) + + +@combinations.generate(one_device_combinations()) +class MirroredOneDeviceDistributionTest( + strategy_test_lib.DistributionTestBase, + strategy_test_lib.OneDeviceDistributionTestBase, + parameterized.TestCase): + + def testMinimizeLoss(self, distribution): + if context.executing_eagerly(): + self._test_minimize_loss_eager(distribution) + else: + self._test_minimize_loss_graph(distribution) - def _skip_eager_if_gpus_less_than(self, num_gpus): - if context.num_gpus() < num_gpus and context.executing_eagerly(): - self.skipTest("Enough GPUs not available for this test in eager mode.") + def testReplicaId(self, distribution): + self._test_replica_id(distribution) - @test_util.run_in_graph_and_eager_modes(config=config) - def testSingleVariable(self): - self._skip_eager_if_gpus_less_than(1) + 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): + + @combinations.generate(combinations.combine( + distribution=[combinations.mirrored_strategy_with_gpu_and_cpu, + combinations.core_mirrored_strategy_with_gpu_and_cpu], + mode=["graph"])) + def testCreatorStacksAreThreadLocal(self, distribution): def model_fn(): - # This variable should be created only once across the threads because of - # special variable_creator functions used by `dist.call_for_each_tower`. - v = variable_scope.variable(1.0, name="foo") - distribution_strategy_context.get_tower_context().merge_call(lambda _: _) + replica_id_str = str(self.evaluate(_replica_id())) + + def thread_creator_fn(next_creator, *args, **kwargs): + return next_creator(*args, **kwargs) + ":thread_" + replica_id_str + + with variable_scope.variable_creator_scope(thread_creator_fn): + # Create a variable in this scope. + v = variable_scope.variable(1.0) + + # This will pause the current thread, and execute the other thread. + ds_context.get_replica_context().merge_call(lambda _: _) return v - dist = mirrored_strategy.MirroredStrategy( - ["/device:GPU:0", "/device:CPU:0"]) + def main_thread_creator(next_creator, *args, **kwargs): + # We are not using the underlying next_creator for test purposes. + del next_creator, args, kwargs + return "main_thread" + + with context.graph_mode(), \ + distribution.scope(), \ + variable_scope.variable_creator_scope(main_thread_creator): + result = distribution.extended.call_for_each_replica(model_fn) + result = distribution.unwrap(result) + expected = ("main_thread:thread_0", "main_thread:thread_1") + self.assertEqual(expected, result) + +@combinations.generate(combinations.combine( + distribution=[ + combinations.mirrored_strategy_with_gpu_and_cpu, + combinations.core_mirrored_strategy_with_gpu_and_cpu], + mode=["graph", "eager"])) +class MirroredStrategyCallForEachReplicaTest(test.TestCase): + + def testExecutingEagerlyOutsideFunction(self, distribution): + """Verify we preserve the value of executing_eagerly_outside_functions().""" + def model_fn(): + return ops.executing_eagerly_outside_functions() + + originally = ops.executing_eagerly_outside_functions() + with distribution.scope(): + in_scope = ops.executing_eagerly_outside_functions() + in_model_fn = distribution.extended.call_for_each_replica(model_fn) + unwrapped = distribution.unwrap(in_model_fn) + self.assertEqual(in_scope, unwrapped[0]) + self.assertEqual(in_scope, originally) + + # Verify this all again, but this time in a FuncGraph. + with func_graph.FuncGraph("fg").as_default(), distribution.scope(): + in_scope = ops.executing_eagerly_outside_functions() + in_model_fn = distribution.extended.call_for_each_replica(model_fn) + unwrapped = distribution.unwrap(in_model_fn) + self.assertEqual(in_scope, unwrapped[0]) + self.assertEqual(in_scope, originally) + + +@combinations.generate(combinations.combine( + distribution=[ + combinations.mirrored_strategy_with_gpu_and_cpu, + combinations.core_mirrored_strategy_with_gpu_and_cpu], + mode=["graph", "eager"])) +class MirroredStrategyVariableCreationTest(test.TestCase): - with dist.scope(): - result = dist.call_for_each_tower(model_fn, run_concurrently=False) - self.assertIsInstance(result, values.MirroredVariable) - self.assertEquals("foo:0", result.name) + # TODO(priyag): Modify more tests to use this helper and check more + # properties. + def _test_mv_properties(self, var, name, strategy): + self.assertIsInstance(var, values.MirroredVariable) + self.assertEqual(name, var.name) + 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) # pylint: disable=protected-access + + def testVariableInFuncGraph(self, distribution): + def model_fn(): + v = variable_scope.variable(2.0, name="bar") + ds_context.get_replica_context().merge_call(lambda _: _) + return v + + with func_graph.FuncGraph("fg").as_default(), distribution.scope(): + v1 = variable_scope.variable(1.0, name="foo") + v2 = distribution.extended.call_for_each_replica(model_fn) - @test_util.run_in_graph_and_eager_modes(config=config) - def testUnnamedVariable(self): - self._skip_eager_if_gpus_less_than(1) + self._test_mv_properties(v1, "foo:0", distribution) + self._test_mv_properties(v2, "bar:0", distribution) + def testSingleVariable(self, distribution): def model_fn(): - v = variable_scope.variable(1.0) - distribution_strategy_context.get_tower_context().merge_call(lambda _: _) + # This variable should be created only once across the threads because of + # special variable_creator functions used by + # `distribution.extended.call_for_each_replica`. + v = variable_scope.variable(1.0, name="foo") + ds_context.get_replica_context().merge_call(lambda _: _) return v - dist = mirrored_strategy.MirroredStrategy( - ["/device:GPU:0", "/device:CPU:0"]) + with distribution.scope(): + result = distribution.extended.call_for_each_replica(model_fn) + self._test_mv_properties(result, "foo:0", distribution) - with dist.scope(): - result = dist.call_for_each_tower(model_fn, run_concurrently=False) - self.assertIsInstance(result, values.MirroredVariable) - # Default name of "Variable" will be used. - self.assertEquals("Variable:0", result.name) + def testUnnamedVariable(self, distribution): + def model_fn(): + v = variable_scope.variable(1.0) + ds_context.get_replica_context().merge_call(lambda _: _) + return v - @test_util.run_in_graph_and_eager_modes(config=config) - def testMultipleVariables(self): - self._skip_eager_if_gpus_less_than(1) + with distribution.scope(): + result = distribution.extended.call_for_each_replica(model_fn) + self._test_mv_properties(result, "Variable:0", distribution) + def testMultipleVariables(self, distribution): def model_fn(): vs = [] for i in range(5): vs.append(variable_scope.variable(1.0, name="foo" + str(i))) - distribution_strategy_context.get_tower_context().merge_call(lambda _: _) + ds_context.get_replica_context().merge_call(lambda _: _) return vs - dist = mirrored_strategy.MirroredStrategy( - ["/device:GPU:0", "/device:CPU:0"]) - - with dist.scope(): - result = dist.call_for_each_tower(model_fn, run_concurrently=False) + with distribution.scope(): + result = distribution.extended.call_for_each_replica(model_fn) for i, v in enumerate(result): - self.assertIsInstance(v, values.MirroredVariable) - self.assertEquals("foo" + str(i) + ":0", v.name) - - @test_util.run_in_graph_and_eager_modes(config=config) - def testMultipleVariablesWithSameCanonicalName(self): - self._skip_eager_if_gpus_less_than(1) + self._test_mv_properties(v, "foo" + str(i) + ":0", distribution) + def testMultipleVariablesWithSameCanonicalName(self, distribution): def model_fn(): vs = [] vs.append(variable_scope.variable(1.0, name="foo/bar")) vs.append(variable_scope.variable(1.0, name="foo_1/bar")) vs.append(variable_scope.variable(1.0, name="foo_1/bar_1")) vs.append(variable_scope.variable(1.0, name="foo/bar_1")) - distribution_strategy_context.get_tower_context().merge_call(lambda _: _) + ds_context.get_replica_context().merge_call(lambda _: _) return vs - dist = mirrored_strategy.MirroredStrategy( - ["/device:GPU:0", "/device:CPU:0"]) - - with dist.scope(): - result = dist.call_for_each_tower(model_fn, run_concurrently=False) + with distribution.scope(): + result = distribution.extended.call_for_each_replica(model_fn) for v in result: self.assertIsInstance(v, values.MirroredVariable) - self.assertEquals(4, len(result)) - self.assertEquals("foo/bar:0", result[0].name) - self.assertEquals("foo_1/bar:0", result[1].name) - self.assertEquals("foo_1/bar_1:0", result[2].name) - self.assertEquals("foo/bar_1:0", result[3].name) - - @test_util.run_in_graph_and_eager_modes(config=config) - def testVariableWithSameCanonicalNameAcrossThreads(self): - self._skip_eager_if_gpus_less_than(1) - - def model_fn(device_id): - v = variable_scope.variable(1.0, name="foo_" + str(device_id)) - distribution_strategy_context.get_tower_context().merge_call(lambda _: _) - return v + self.assertEqual(4, len(result)) + self.assertEqual("foo/bar:0", result[0].name) + self.assertEqual("foo_1/bar:0", result[1].name) + self.assertEqual("foo_1/bar_1:0", result[2].name) + self.assertEqual("foo/bar_1:0", result[3].name) - dist = mirrored_strategy.MirroredStrategy( - ["/device:GPU:0", "/device:CPU:0"]) + def testVariableWithSameCanonicalNameAcrossThreads(self, distribution): + def model_fn(): + replica_id = self.evaluate(_replica_id()) + v = variable_scope.variable(1.0, name="foo_" + str(replica_id)) + ds_context.get_replica_context().merge_call(lambda _: _) + return v - with dist.scope(): - result = dist.call_for_each_tower( - model_fn, dist.worker_device_index, run_concurrently=False) + with distribution.scope(): + result = distribution.extended.call_for_each_replica(model_fn) self.assertIsInstance(result, values.MirroredVariable) # The resulting mirrored variable will use the name from the first device. - self.assertEquals("foo_0:0", result.name) + self.assertEqual("foo_0:0", result.name) - @test_util.run_in_graph_and_eager_modes(config=config) - def testWithLayers(self): - self._skip_eager_if_gpus_less_than(1) + def testWithLayers(self, distribution): def model_fn(features): with variable_scope.variable_scope("common"): layer1 = core.Dense(1) @@ -303,47 +359,35 @@ class MirroredStrategyVariableCreationTest(test.TestCase): layer2 = core.Dense(1) layer2(features) # This will pause the current thread, and execute the other thread. - distribution_strategy_context.get_tower_context().merge_call( - lambda _: _) + ds_context.get_replica_context().merge_call(lambda _: _) layer3 = core.Dense(1) layer3(features) return [(layer1.kernel, layer1.bias), (layer2.kernel, layer2.bias), (layer3.kernel, layer3.bias)] - dist = mirrored_strategy.MirroredStrategy( - ["/device:GPU:0", "/device:CPU:0"]) - ds = dist.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 dist.scope(): - result = dist.call_for_each_tower( - model_fn, features, run_concurrently=False) + with distribution.scope(): + result = distribution.extended.call_for_each_replica( + model_fn, args=(features,)) suffixes = ["", "_1", "_2"] for (kernel, bias), suffix in zip(result, suffixes): self.assertIsInstance(kernel, values.MirroredVariable) - self.assertEquals("common/dense" + suffix + "/kernel:0", kernel.name) + self.assertEqual("common/dense" + suffix + "/kernel:0", kernel.name) self.assertIsInstance(bias, values.MirroredVariable) - self.assertEquals("common/dense" + suffix + "/bias:0", bias.name) - - @test_util.run_in_graph_and_eager_modes(config=config) - def testWithVariableAndVariableScope(self): - self._skip_eager_if_gpus_less_than(1) + self.assertEqual("common/dense" + suffix + "/bias:0", bias.name) + def testWithVariableAndVariableScope(self, distribution): def model_fn(): v0 = variable_scope.variable(1.0, name="var0", aggregation=None) with variable_scope.variable_scope("common"): v1 = variable_scope.variable(1.0, name="var1") # This will pause the current thread, and execute the other thread. - distribution_strategy_context.get_tower_context().merge_call( - lambda _: _) + ds_context.get_replica_context().merge_call(lambda _: _) v2 = variable_scope.variable( 1.0, name="var2", @@ -357,37 +401,31 @@ class MirroredStrategyVariableCreationTest(test.TestCase): return v0, v1, v2, v3 - devices = ["/device:CPU:0", "/device:GPU:0"] - dist = mirrored_strategy.MirroredStrategy(devices) - with dist.scope(): + with distribution.scope(): v = variable_scope.variable(1.0, name="var-main0") - self.assertEquals("var-main0:0", v.name) + self.assertEqual("var-main0:0", v.name) - result = dist.call_for_each_tower(model_fn, run_concurrently=False) - self.assertEquals(4, len(result)) + result = distribution.extended.call_for_each_replica(model_fn) + self.assertEqual(4, len(result)) v0, v1, v2, v3 = result self.assertIsInstance(v0, values.MirroredVariable) - self.assertEquals("var0:0", v0.name) + self.assertEqual("var0:0", v0.name) self.assertIsInstance(v1, values.MirroredVariable) - self.assertEquals("common/var1:0", v1.name) - self.assertIsInstance(v2, values.TowerLocalVariable) - self.assertEquals("common/var2:0", v2.name) - self.assertEquals(variable_scope.VariableAggregation.SUM, v2.aggregation) + self.assertEqual("common/var1:0", v1.name) + self.assertIsInstance(v2, values.ReplicaLocalVariable) + self.assertEqual("common/var2:0", v2.name) + self.assertEqual(variable_scope.VariableAggregation.SUM, v2.aggregation) self.assertIsInstance(v3, values.MirroredVariable) - self.assertEquals("common/var3:0", v3.name) - self.assertEquals(variable_scope.VariableAggregation.MEAN, v3.aggregation) - - @test_util.run_in_graph_and_eager_modes(config=config) - def testWithGetVariableAndVariableScope(self): - self._skip_eager_if_gpus_less_than(1) + self.assertEqual("common/var3:0", v3.name) + self.assertEqual(variable_scope.VariableAggregation.MEAN, v3.aggregation) + def testWithGetVariableAndVariableScope(self, distribution): def model_fn(): v0 = variable_scope.get_variable("var0", [1]) with variable_scope.variable_scope("common"): v1 = variable_scope.get_variable("var1", [1]) # This will pause the current thread, and execute the other thread. - distribution_strategy_context.get_tower_context().merge_call( - lambda _: _) + ds_context.get_replica_context().merge_call(lambda _: _) v2 = variable_scope.get_variable( "var2", [1], synchronization=variable_scope.VariableSynchronization.ON_READ, @@ -399,35 +437,30 @@ class MirroredStrategyVariableCreationTest(test.TestCase): return v0, v1, v2, v3 - devices = ["/device:CPU:0", "/device:GPU:0"] - dist = mirrored_strategy.MirroredStrategy(devices) - with dist.scope(): + with distribution.scope(): with variable_scope.variable_scope("main"): v = variable_scope.get_variable("var-main0", [1]) - self.assertEquals("main/var-main0:0", v.name) + self.assertEqual("main/var-main0:0", v.name) - result = dist.call_for_each_tower(model_fn, run_concurrently=False) - self.assertEquals(4, len(result)) + result = distribution.extended.call_for_each_replica(model_fn) + self.assertEqual(4, len(result)) v0, v1, v2, v3 = result self.assertIsInstance(v0, values.MirroredVariable) - self.assertEquals("main/var0:0", v0.name) + self.assertEqual("main/var0:0", v0.name) self.assertIsInstance(v1, values.MirroredVariable) - self.assertEquals("main/common/var1:0", v1.name) - self.assertIsInstance(v2, values.TowerLocalVariable) - self.assertEquals("main/common/var2:0", v2.name) - self.assertEquals(variable_scope.VariableAggregation.SUM, - v2.aggregation) + self.assertEqual("main/common/var1:0", v1.name) + self.assertIsInstance(v2, values.ReplicaLocalVariable) + self.assertEqual("main/common/var2:0", v2.name) + self.assertEqual(variable_scope.VariableAggregation.SUM, + v2.aggregation) self.assertIsInstance(v3, values.MirroredVariable) - self.assertEquals("main/common/var3:0", v3.name) - self.assertEquals(variable_scope.VariableAggregation.MEAN, - v3.aggregation) - - @test_util.run_in_graph_and_eager_modes(config=config) - def testOnlyFirstTowerUpdatesVariables(self): - self._skip_eager_if_gpus_less_than(1) + self.assertEqual("main/common/var3:0", v3.name) + self.assertEqual(variable_scope.VariableAggregation.MEAN, + v3.aggregation) + def testOnlyFirstReplicaUpdatesVariables(self, distribution): def create_fn(): - aggregation = variable_scope.VariableAggregation.ONLY_FIRST_TOWER + aggregation = variable_scope.VariableAggregation.ONLY_FIRST_REPLICA v0 = variable_scope.variable( 2.0, name="on_read", @@ -441,71 +474,73 @@ class MirroredStrategyVariableCreationTest(test.TestCase): return v0, v1 devices = ["/device:GPU:0", "/device:CPU:0"] - dist = mirrored_strategy.MirroredStrategy(devices) - with dist.scope(): - v0, v1 = dist.call_for_each_tower(create_fn, run_concurrently=False) + with distribution.scope(): + v0, v1 = distribution.extended.call_for_each_replica(create_fn) self.evaluate(v0.initializer) self.assertEqual(2.0, self.evaluate(v0.get(devices[0]))) self.assertEqual(2.0, self.evaluate(v0.get(devices[1]))) - self.assertEqual(2.0, self.evaluate(dist.read_var(v0))) + self.assertEqual(2.0, self.evaluate(distribution.extended.read_var(v0))) self.evaluate(v1.initializer) self.assertEqual(3.0, self.evaluate(v1.get(devices[0]))) self.assertEqual(3.0, self.evaluate(v1.get(devices[1]))) - self.assertEqual(3.0, self.evaluate(dist.read_var(v1))) + self.assertEqual(3.0, self.evaluate(distribution.extended.read_var(v1))) + + def replica_id_plus_one(): + return math_ops.cast(_replica_id() + 1, dtype=dtypes.float32) # Update using the assign_add member function. - def update_member_fn(device_id): - update0 = v0.assign_add(5.0 * (device_id + 1)) - update1 = v1.assign_add(7.0 * (device_id + 1)) + def update_member_fn(): + update0 = v0.assign_add(5.0 * replica_id_plus_one()) + update1 = v1.assign_add(7.0 * replica_id_plus_one()) return update0, update1 - update0a, update1a = dist.call_for_each_tower( - update_member_fn, dist.worker_device_index, run_concurrently=False) + update0a, update1a = distribution.extended.call_for_each_replica( + update_member_fn) # Update "sync on read" variable. - self.evaluate(dist.group(update0a)) + self.evaluate(distribution.group(update0a)) self.assertEqual(2.0 + 5.0, self.evaluate(v0.get(devices[0]))) # Writes are not synchronized for "sync on read" variables, # so device[1] can end up with a different value. self.assertEqual(2.0 + 2*5.0, self.evaluate(v0.get(devices[1]))) # Always reads from device 0. - self.assertEqual(2.0 + 5.0, self.evaluate(dist.read_var(v0))) + self.assertEqual(2.0 + 5.0, self.evaluate( + distribution.extended.read_var(v0))) # Update "sync on write" variable. - self.evaluate(dist.group(update1a)) + self.evaluate(distribution.group(update1a)) self.assertEqual(3.0 + 7.0, self.evaluate(v1.get(devices[0]))) # Writes are synchronized for v1, only the argument to assign_add on # device[0] is used. self.assertEqual(3.0 + 7.0, self.evaluate(v1.get(devices[1]))) - self.assertEqual(3.0 + 7.0, self.evaluate(dist.read_var(v1))) + self.assertEqual(3.0 + 7.0, self.evaluate( + distribution.extended.read_var(v1))) # Update using state_ops.assign_add global function. - def update_state_ops_fn(device_id): - update0 = state_ops.assign_add(v0, 11.0 * (device_id + 1)) - update1 = state_ops.assign_add(v1, 13.0 * (device_id + 1)) + def update_state_ops_fn(): + update0 = state_ops.assign_add(v0, 11.0 * replica_id_plus_one()) + update1 = state_ops.assign_add(v1, 13.0 * replica_id_plus_one()) return update0, update1 - update0b, update1b = dist.call_for_each_tower( - update_state_ops_fn, dist.worker_device_index, run_concurrently=False) - self.evaluate(dist.group(update0b)) + update0b, update1b = distribution.extended.call_for_each_replica( + update_state_ops_fn) + self.evaluate(distribution.group(update0b)) # Update "sync on read" variable. self.assertEqual(2.0 + 5.0 + 11.0, self.evaluate(v0.get(devices[0]))) self.assertEqual(2.0 + 2*5.0 + 2*11.0, self.evaluate(v0.get(devices[1]))) - self.assertEqual(2.0 + 5.0 + 11.0, self.evaluate(dist.read_var(v0))) + self.assertEqual(2.0 + 5.0 + 11.0, self.evaluate( + distribution.extended.read_var(v0))) # Update "sync on write" variable. - self.evaluate(dist.group(update1b)) + self.evaluate(distribution.group(update1b)) self.assertEqual(3.0 + 7.0 + 13.0, self.evaluate(v1.get(devices[0]))) self.assertEqual(3.0 + 7.0 + 13.0, self.evaluate(v1.get(devices[1]))) - self.assertEqual(3.0 + 7.0 + 13.0, self.evaluate(dist.read_var(v1))) - - @test_util.run_in_graph_and_eager_modes(config=config) - def testNoneSynchronizationWithGetVariable(self): - self._skip_eager_if_gpus_less_than(1) - devices = ["/device:CPU:0", "/device:GPU:0"] - dist = mirrored_strategy.MirroredStrategy(devices) - with dist.scope(): + self.assertEqual(3.0 + 7.0 + 13.0, self.evaluate( + distribution.extended.read_var(v1))) + + def testNoneSynchronizationWithGetVariable(self, distribution): + with distribution.scope(): with self.assertRaisesRegexp( ValueError, "`NONE` variable synchronization mode is not " "supported with `Mirrored` distribution strategy. Please change " @@ -514,12 +549,8 @@ class MirroredStrategyVariableCreationTest(test.TestCase): "v", [1], synchronization=variable_scope.VariableSynchronization.NONE) - @test_util.run_in_graph_and_eager_modes(config=config) - def testNoneSynchronizationWithVariable(self): - self._skip_eager_if_gpus_less_than(1) - devices = ["/device:CPU:0", "/device:GPU:0"] - dist = mirrored_strategy.MirroredStrategy(devices) - with dist.scope(): + def testNoneSynchronizationWithVariable(self, distribution): + with distribution.scope(): with self.assertRaisesRegexp( ValueError, "`NONE` variable synchronization mode is not " "supported with `Mirrored` distribution strategy. Please change " @@ -529,23 +560,15 @@ class MirroredStrategyVariableCreationTest(test.TestCase): name="v", synchronization=variable_scope.VariableSynchronization.NONE) - @test_util.run_in_graph_and_eager_modes(config=config) - def testInvalidSynchronizationWithVariable(self): - self._skip_eager_if_gpus_less_than(1) - devices = ["/device:CPU:0", "/device:GPU:0"] - dist = mirrored_strategy.MirroredStrategy(devices) - with dist.scope(): + def testInvalidSynchronizationWithVariable(self, distribution): + with distribution.scope(): with self.assertRaisesRegexp( ValueError, "Invalid variable synchronization mode: Invalid for " "variable: v"): variable_scope.variable(1.0, name="v", synchronization="Invalid") - @test_util.run_in_graph_and_eager_modes(config=config) - def testInvalidAggregationWithGetVariable(self): - self._skip_eager_if_gpus_less_than(1) - devices = ["/device:CPU:0", "/device:GPU:0"] - dist = mirrored_strategy.MirroredStrategy(devices) - with dist.scope(): + def testInvalidAggregationWithGetVariable(self, distribution): + with distribution.scope(): with self.assertRaisesRegexp( ValueError, "Invalid variable aggregation mode: invalid for " "variable: v"): @@ -554,12 +577,8 @@ class MirroredStrategyVariableCreationTest(test.TestCase): synchronization=variable_scope.VariableSynchronization.ON_WRITE, aggregation="invalid") - @test_util.run_in_graph_and_eager_modes(config=config) - def testInvalidAggregationWithVariable(self): - self._skip_eager_if_gpus_less_than(1) - devices = ["/device:CPU:0", "/device:GPU:0"] - dist = mirrored_strategy.MirroredStrategy(devices) - with dist.scope(): + def testInvalidAggregationWithVariable(self, distribution): + with distribution.scope(): with self.assertRaisesRegexp( ValueError, "Invalid variable aggregation mode: invalid for " "variable: v"): @@ -569,53 +588,26 @@ class MirroredStrategyVariableCreationTest(test.TestCase): synchronization=variable_scope.VariableSynchronization.ON_WRITE, aggregation="invalid") - @test_util.run_in_graph_and_eager_modes(config=config) - def testThreeDevices(self): - self._skip_eager_if_gpus_less_than(2) - - def model_fn(): - v = variable_scope.variable(1.0, name="foo") - distribution_strategy_context.get_tower_context().merge_call(lambda _: _) - return v - - dist = mirrored_strategy.MirroredStrategy( - ["/device:GPU:0", "/device:GPU:1", "/device:CPU:0"]) - - with dist.scope(): - result = dist.call_for_each_tower(model_fn, run_concurrently=False) - self.assertIsInstance(result, values.MirroredVariable) - self.assertEquals("foo:0", result.name) - - @test_util.run_in_graph_and_eager_modes(config=config) - def testNonMatchingVariableCreation(self): - self._skip_eager_if_gpus_less_than(1) - + def testNonMatchingVariableCreation(self, distribution): def model_fn(name): v = variable_scope.variable(1.0, name=name) - distribution_strategy_context.get_tower_context().merge_call(lambda _: _) + ds_context.get_replica_context().merge_call(lambda _: _) return v - dist = mirrored_strategy.MirroredStrategy( - ["/device:GPU:0", "/device:CPU:0"]) - - with dist.scope(): - names = values.DistributedValues({ - "/device:CPU:0": "foo", - "/device:GPU:0": "bar" - }) + with distribution.scope(): + device_map = values.ReplicaDeviceMap(("/device:CPU:0", "/device:GPU:0")) + names = values.DistributedValues(device_map, ("foo", "bar")) with self.assertRaises(RuntimeError): - _ = dist.call_for_each_tower(model_fn, names, run_concurrently=False) - - @test_util.run_in_graph_and_eager_modes(config=config) - def testTowerLocalVariable(self): - self._skip_eager_if_gpus_less_than(1) + _ = distribution.extended.call_for_each_replica(model_fn, args=(names,)) + def testReplicaLocalVariable(self, distribution): all_v_sum = {} all_v_mean = {} components_sum = {} components_mean = {} - def model_fn(device_id): + def model_fn(): + replica_id = self.evaluate(_replica_id()) v_sum = variable_scope.variable( 1.0, synchronization=variable_scope.VariableSynchronization.ON_READ, @@ -624,29 +616,25 @@ class MirroredStrategyVariableCreationTest(test.TestCase): 4.0, synchronization=variable_scope.VariableSynchronization.ON_READ, aggregation=variable_scope.VariableAggregation.MEAN) - self.assertTrue(isinstance(v_sum, values.TowerLocalVariable)) - self.assertTrue(isinstance(v_mean, values.TowerLocalVariable)) - updates = [v_sum.assign_add(2.0 + device_id), - v_mean.assign(6.0 * device_id)] - all_v_sum[device_id] = v_sum - all_v_mean[device_id] = v_mean + self.assertTrue(isinstance(v_sum, values.ReplicaLocalVariable)) + self.assertTrue(isinstance(v_mean, values.ReplicaLocalVariable)) + updates = [v_sum.assign_add(2.0 + replica_id), + v_mean.assign(6.0 * replica_id)] + all_v_sum[replica_id] = v_sum + all_v_mean[replica_id] = v_mean c_sum = v_sum.get() c_mean = v_mean.get() - components_sum[device_id] = c_sum - components_mean[device_id] = c_mean + components_sum[replica_id] = c_sum + components_mean[replica_id] = c_mean self.assertIsNot(v_sum, c_sum) self.assertIsNot(v_mean, c_mean) return updates, v_sum, v_mean, c_sum, c_mean - dist = mirrored_strategy.MirroredStrategy( - ["/device:GPU:0", "/device:CPU:0"]) - - with dist.scope(): - # Create "sum" and "mean" versions of TowerLocalVariables. + with distribution.scope(): + # Create "sum" and "mean" versions of ReplicaLocalVariables. ret_ops, ret_v_sum, ret_v_mean, regrouped_sum, regrouped_mean = ( - dist.call_for_each_tower( - model_fn, dist.worker_device_index, run_concurrently=False)) - # Should see the same wrapping instance in all towers. + distribution.extended.call_for_each_replica(model_fn)) + # Should see the same wrapping instance in all replicas. self.assertIs(all_v_sum[0], ret_v_sum) self.assertIs(all_v_mean[0], ret_v_mean) self.assertIs(all_v_sum[0], all_v_sum[1]) @@ -660,10 +648,10 @@ class MirroredStrategyVariableCreationTest(test.TestCase): # Apply updates self.evaluate(variables.global_variables_initializer()) - self.evaluate([y for x in ret_ops for y in dist.unwrap(x)]) + self.evaluate([y for x in ret_ops for y in distribution.unwrap(x)]) expected_sum = 0.0 expected_mean = 0.0 - for i, d in enumerate(dist.worker_devices): + for i, d in enumerate(distribution.extended.worker_devices): # Should see different values on different devices. v_sum_value = self.evaluate(ret_v_sum.get(d).read_value()) v_mean_value = self.evaluate(ret_v_mean.get(d).read_value()) @@ -673,221 +661,244 @@ class MirroredStrategyVariableCreationTest(test.TestCase): expected = i * 6.0 self.assertEqual(expected, v_mean_value) expected_mean += expected - expected_mean /= len(dist.worker_devices) + expected_mean /= len(distribution.extended.worker_devices) # Without get(device), should return the value you get by - # applying the reduction across all towers (whether you use + # applying the reduction across all replicas (whether you use # read_var(), get(), or nothing). - self.assertEqual(expected_sum, self.evaluate(dist.read_var(ret_v_sum))) - self.assertEqual(expected_mean, self.evaluate(dist.read_var(ret_v_mean))) + self.assertEqual(expected_sum, self.evaluate( + distribution.extended.read_var(ret_v_sum))) + self.assertEqual(expected_mean, self.evaluate( + distribution.extended.read_var(ret_v_mean))) self.assertEqual(expected_sum, self.evaluate(ret_v_sum.get())) self.assertEqual(expected_mean, self.evaluate(ret_v_mean.get())) self.assertEqual(expected_sum, self.evaluate(ret_v_sum)) self.assertEqual(expected_mean, self.evaluate(ret_v_mean)) + # TODO(priyag): Update this test to work in eager mode as well. + def testDynamicRnnVariables(self, distribution): + def model_fn(): + inputs = constant_op.constant(2 * [2 * [[0.0, 1.0, 2.0, 3.0, 4.0]]]) + cell_fw = rnn_cell_impl.LSTMCell(300) + cell_bw = rnn_cell_impl.LSTMCell(300) + (outputs, _) = rnn.bidirectional_dynamic_rnn( + cell_fw, + cell_bw, + inputs, + dtype=dtypes.float32) + return outputs + + with context.graph_mode(), distribution.scope(): + result = distribution.extended.call_for_each_replica(model_fn) + # Two variables are created by the RNN layer. + self.assertEqual(2, len(result)) + for v in result: + self.assertIsInstance(v, values.DistributedValues) + _, v1 = distribution.unwrap(v) + self.assertStartsWith(v1._op.name, "replica_1/") + + def testReplicaLocalVariableUpdate(self, distribution): + def model_fn(): + v_sum = variable_scope.variable( + 1.0, + synchronization=variable_scope.VariableSynchronization.ON_READ, + aggregation=variable_scope.VariableAggregation.SUM) + self.assertTrue(isinstance(v_sum, values.ReplicaLocalVariable)) + return v_sum + + def update(var, value): + return var.assign(value) + + with distribution.scope(): + ret_v_sum = distribution.extended.call_for_each_replica(model_fn) + + # Initialize variables. + self.evaluate(variables.global_variables_initializer()) + # Assert that the aggregated value of the replica local vars is the sum + # of the individual values before running the update ops. + self.assertEqual(1.0, self.evaluate(ret_v_sum.get( + distribution.extended.worker_devices[0]).read_value())) + self.assertEqual(2.0, self.evaluate(ret_v_sum)) + + # Apply updates. + update_ops = distribution.extended.update( + ret_v_sum, update, args=(5.0,), group=False) + self.evaluate(update_ops) + # Assert that the aggregated value of the replica local vars is the sum + # of the individual values after running the update ops. + self.assertEqual(5.0, self.evaluate(ret_v_sum.get( + distribution.extended.worker_devices[0]).read_value())) + self.assertEqual(10.0, self.evaluate(ret_v_sum)) + + def testVarDistributeStrategy(self, distribution): + with distribution.scope(): + mirrored = variable_scope.variable(1.0) + replica_local = variable_scope.variable( + 1.0, + synchronization=variable_scope.VariableSynchronization.ON_READ) + self.assertIs(distribution, mirrored.distribute_strategy) + self.assertIs(distribution, replica_local.distribute_strategy) + + +@combinations.generate(combinations.combine( + distribution=[ + combinations.mirrored_strategy_with_gpu_and_cpu, + combinations.core_mirrored_strategy_with_gpu_and_cpu], + mode=["graph"])) +class MirroredStrategyNameScopeTest(test.TestCase): # NOTE(priyag): Names and name scopes are ignored in eager, hence we are not # testing this in eager mode. - def testNameScope(self): + def testNameScope(self, distribution): def model_fn(): with ops.name_scope("foo"): a = constant_op.constant(1.0, name="a") - distribution_strategy_context.get_tower_context().merge_call( - lambda _: _) + ds_context.get_replica_context().merge_call(lambda _: _) b = constant_op.constant(1.0, name="b") return a, b - dist = mirrored_strategy.MirroredStrategy( - ["/device:GPU:0", "/device:CPU:0"]) - - with context.graph_mode(), dist.scope(): + with context.graph_mode(), distribution.scope(): with ops.name_scope("main"): - result = dist.call_for_each_tower(model_fn, run_concurrently=False) - self.assertEquals(2, len(result)) + result = distribution.extended.call_for_each_replica(model_fn) + self.assertEqual(2, len(result)) for v, name in zip(result, ["a", "b"]): self.assertIsInstance(v, values.DistributedValues) - v0, v1 = dist.unwrap(v) - self.assertEquals("main/foo/" + name + ":0", v0.name) - self.assertEquals("main/tower_1/foo/" + name + ":0", v1.name) + v0, v1 = distribution.unwrap(v) + self.assertEqual("main/foo/" + name + ":0", v0.name) + self.assertEqual("main/replica_1/foo/" + name + ":0", v1.name) - def testWithDefaultName(self): + def testWithDefaultName(self, distribution): def model_fn(): with ops.name_scope(None, "foo"): a = constant_op.constant(1.0, name="a") - distribution_strategy_context.get_tower_context().merge_call( - lambda _: _) + ds_context.get_replica_context().merge_call(lambda _: _) b = constant_op.constant(2.0, name="b") return a, b - dist = mirrored_strategy.MirroredStrategy( - ["/device:GPU:0", "/device:CPU:0"]) - - with context.graph_mode(), dist.scope(): - result = dist.call_for_each_tower(model_fn, run_concurrently=False) - self.assertEquals(2, len(result)) + with context.graph_mode(), distribution.scope(): + result = distribution.extended.call_for_each_replica(model_fn) + self.assertEqual(2, len(result)) for v, name in zip(result, ["a", "b"]): self.assertIsInstance(v, values.DistributedValues) - v0, v1 = dist.unwrap(v) - self.assertEquals("foo/" + name + ":0", v0.name) - self.assertEquals("tower_1/foo/" + name + ":0", v1.name) + v0, v1 = distribution.unwrap(v) + self.assertEqual("foo/" + name + ":0", v0.name) + self.assertEqual("replica_1/foo/" + name + ":0", v1.name) # variable_scope.variable() respects name scopes when creating # variables. On the other hand variable_scope.get_variable() ignores name # scopes when creating variables. We test both methods of creating variables # to make sure that we have the same variable names in both cases. - def testNameScopeWithVariable(self): - def in_cross_tower(_): + def testNameScopeWithVariable(self, distribution): + def in_cross_replica(_): c = variable_scope.variable(1.0, name="c") return c def model_fn(): b = variable_scope.variable(1.0, name="b") with ops.name_scope("foo"): - c = distribution_strategy_context.get_tower_context().merge_call( - in_cross_tower) + c = ds_context.get_replica_context().merge_call(in_cross_replica) return b, c - dist = mirrored_strategy.MirroredStrategy( - ["/device:GPU:0", "/device:CPU:0"]) - - with context.graph_mode(), dist.scope(): + with context.graph_mode(), distribution.scope(): with ops.name_scope("main"): a = variable_scope.variable(1.0, name="a") - result = dist.call_for_each_tower(model_fn, run_concurrently=False) + result = distribution.extended.call_for_each_replica(model_fn) result_b = result[0] result_c = result[1] self.assertIsInstance(result_b, values.DistributedValues) self.assertIsInstance(result_c, values.DistributedValues) - a0, a1 = dist.unwrap(a) - b0, b1 = dist.unwrap(result_b) - c0, c1 = dist.unwrap(result_c) - self.assertEquals("main/a:0", a0.name) - self.assertEquals("main/a/replica_1:0", a1.name) - self.assertEquals("main/b:0", b0.name) - self.assertEquals("main/b/replica_1:0", b1.name) - self.assertEquals("main/foo/c:0", c0.name) - self.assertEquals("main/foo/c/replica_1:0", c1.name) - - def testNameScopeWithGetVariable(self): - def in_cross_tower(_): + a0, a1 = distribution.unwrap(a) + b0, b1 = distribution.unwrap(result_b) + c0, c1 = distribution.unwrap(result_c) + self.assertEqual("main/a:0", a0.name) + self.assertEqual("main/a/replica_1:0", a1.name) + self.assertEqual("main/b:0", b0.name) + self.assertEqual("main/b/replica_1:0", b1.name) + self.assertEqual("main/foo/c:0", c0.name) + self.assertEqual("main/foo/c/replica_1:0", c1.name) + + def testNameScopeWithGetVariable(self, distribution): + def in_cross_replica(_): c = variable_scope.get_variable("c", [1]) return c def model_fn(): b = variable_scope.get_variable("b", [1]) with ops.name_scope("foo"): - c = distribution_strategy_context.get_tower_context().merge_call( - in_cross_tower) + c = ds_context.get_replica_context().merge_call(in_cross_replica) return b, c - dist = mirrored_strategy.MirroredStrategy( - ["/device:GPU:0", "/device:CPU:0"]) - - with context.graph_mode(), dist.scope(): + with context.graph_mode(), distribution.scope(): with ops.name_scope("main"): a = variable_scope.get_variable("a", [1]) - result = dist.call_for_each_tower(model_fn, run_concurrently=False) + result = distribution.extended.call_for_each_replica(model_fn) result_b = result[0] result_c = result[1] self.assertIsInstance(result_b, values.DistributedValues) self.assertIsInstance(result_c, values.DistributedValues) - a0, a1 = dist.unwrap(a) - b0, b1 = dist.unwrap(result_b) - c0, c1 = dist.unwrap(result_c) - self.assertEquals("a:0", a0.name) - self.assertEquals("a/replica_1:0", a1.name) - self.assertEquals("b:0", b0.name) - self.assertEquals("b/replica_1:0", b1.name) - self.assertEquals("c:0", c0.name) - self.assertEquals("c/replica_1:0", c1.name) - - def testDynamicRnnVariables(self): + a0, a1 = distribution.unwrap(a) + b0, b1 = distribution.unwrap(result_b) + c0, c1 = distribution.unwrap(result_c) + self.assertEqual("a:0", a0.name) + self.assertEqual("a/replica_1:0", a1.name) + self.assertEqual("b:0", b0.name) + self.assertEqual("b/replica_1:0", b1.name) + self.assertEqual("c:0", c0.name) + self.assertEqual("c/replica_1:0", c1.name) + + +@combinations.generate( + combinations.combine( + distribution=[ + combinations.NamedDistribution( + "Mirrored3Devices", + # pylint: disable=g-long-lambda + lambda: mirrored_strategy.MirroredStrategy( + ["/device:GPU:0", "/device:GPU:1", "/device:CPU:0"]), + required_gpus=2), + combinations.NamedDistribution( + "CoreMirrored3Devices", + # pylint: disable=g-long-lambda + lambda: mirrored_strategy.CoreMirroredStrategy( + ["/device:GPU:0", "/device:GPU:1", "/device:CPU:0"]), + required_gpus=2) + ], + mode=["graph", "eager"])) +class MirroredThreeDeviceDistributionTest( + strategy_test_lib.DistributionTestBase, + parameterized.TestCase): + + def testThreeDevices(self, distribution): def model_fn(): - inputs = constant_op.constant(2 * [2 * [[0.0, 1.0, 2.0, 3.0, 4.0]]]) - cell_fw = rnn_cell_impl.LSTMCell(300) - cell_bw = rnn_cell_impl.LSTMCell(300) - (outputs, _) = rnn.bidirectional_dynamic_rnn( - cell_fw, - cell_bw, - inputs, - dtype=dtypes.float32) - return outputs - - dist = mirrored_strategy.MirroredStrategy( - ["/device:GPU:0", "/device:CPU:0"]) - - with context.graph_mode(), dist.scope(): - result = dist.call_for_each_tower(model_fn, run_concurrently=False) - # Two variables are created by the RNN layer. - self.assertEquals(2, len(result)) - for v in result: - self.assertIsInstance(v, values.DistributedValues) - _, v1 = dist.unwrap(v) - self.assertStartsWith(v1.name, "tower_1/") - - @test_util.run_in_graph_and_eager_modes(config=config) - def testTowerLocalVariableUpdate(self): - with context.graph_mode(): - - def model_fn(): - v_sum = variable_scope.variable( - 1.0, - synchronization=variable_scope.VariableSynchronization.ON_READ, - aggregation=variable_scope.VariableAggregation.SUM) - self.assertTrue(isinstance(v_sum, values.TowerLocalVariable)) - return v_sum - - dist = mirrored_strategy.MirroredStrategy( - ["/device:GPU:0", "/device:GPU:1"]) - - def update(var, value): - return var.assign(value) - - with dist.scope(): - ret_v_sum = dist.call_for_each_tower(model_fn, run_concurrently=False) - update_ops = dist.update(ret_v_sum, update, 5.0, grouped=False) - - # Initialize variables. - self.evaluate(variables.global_variables_initializer()) - # Assert that the aggregated value of the tower local vars is the sum of - # the individual values before running the update ops. - self.assertEquals(1.0, self.evaluate( - ret_v_sum.get(dist._devices[0]).read_value())) - self.assertEquals(2.0, self.evaluate(ret_v_sum)) + v = variable_scope.variable(1.0, name="foo") + ds_context.get_replica_context().merge_call(lambda _: _) + return v - # Apply updates. - self.evaluate(update_ops) - # Assert that the aggregated value of the tower local vars is the sum of - # the individual values after running the update ops. - self.assertEquals(5.0, self.evaluate( - ret_v_sum.get(dist._devices[0]).read_value())) - self.assertEquals(10.0, self.evaluate(ret_v_sum)) + with distribution.scope(): + result = distribution.extended.call_for_each_replica(model_fn) + self.assertIsInstance(result, values.MirroredVariable) + self.assertEqual("foo:0", result.name) +@combinations.generate(combinations.combine( + distribution=[ + combinations.mirrored_strategy_with_gpu_and_cpu, + combinations.core_mirrored_strategy_with_gpu_and_cpu], + mode=["graph", "eager"])) class MirroredVariableUpdateTest(test.TestCase): # The following tests check assign, assign_add and assign_sub on Mirrored - # variables in tower and cross tower context. - config = config_pb2.ConfigProto() - config.allow_soft_placement = True - - def _skip_eager_if_gpus_less_than(self, num_gpus): - if context.num_gpus() < num_gpus and context.executing_eagerly(): - self.skipTest("Enough GPUs not available for this test in eager mode.") + # variables in replica and cross replica context. - @test_util.run_in_graph_and_eager_modes(config=config) - def testAssignMirroredVarTowerContextWithoutAggregationType(self): + def testAssignMirroredVarReplicaContextWithoutAggregationType(self, + distribution): # Test that we always have an aggregation type set on the mirrored variable - # if we assign to it in tower mode. - self._skip_eager_if_gpus_less_than(1) + # if we assign to it in replica mode. def var_fn(): v = variable_scope.variable(1.0, name="foo") return v - dist = mirrored_strategy.MirroredStrategy( - ["/device:GPU:0", "/device:CPU:0"]) - - with dist.scope(): - mirrored_var = dist.call_for_each_tower(var_fn, run_concurrently=False) + with distribution.scope(): + mirrored_var = distribution.extended.call_for_each_replica(var_fn) self.assertIsInstance(mirrored_var, values.MirroredVariable) self.evaluate(variables.global_variables_initializer()) @@ -896,24 +907,20 @@ class MirroredVariableUpdateTest(test.TestCase): with self.assertRaisesRegexp( ValueError, "You must specify an aggregation method to update a " - "MirroredVariable in Tower Context."): - self.evaluate(dist.unwrap(dist.call_for_each_tower(model_fn))) + "MirroredVariable in Replica Context."): + self.evaluate(distribution.unwrap( + distribution.extended.call_for_each_replica(model_fn))) - @test_util.run_in_graph_and_eager_modes(config=config) - def testAssignMirroredVarTowerContextWithSum(self): - # Test that we don't reduce a non-per-device value with the "sum" + def testAssignMirroredVarReplicaContextWithSum(self, distribution): + # Test that we don't reduce a non-per-replica value with the "sum" # aggregation type. - self._skip_eager_if_gpus_less_than(1) def var_fn(): v = variable_scope.variable( 1.0, name="foo", aggregation=variable_scope.VariableAggregation.SUM) return v - dist = mirrored_strategy.MirroredStrategy( - ["/device:GPU:0", "/device:CPU:0"]) - - with dist.scope(): - mirrored_var = dist.call_for_each_tower(var_fn, run_concurrently=False) + with distribution.scope(): + mirrored_var = distribution.extended.call_for_each_replica(var_fn) self.assertIsInstance(mirrored_var, values.MirroredVariable) self.evaluate(variables.global_variables_initializer()) @@ -922,225 +929,184 @@ class MirroredVariableUpdateTest(test.TestCase): with self.assertRaisesRegexp( ValueError, "A non-DistributedValues value 5.0 cannot be reduced " - "with the given aggregation VariableAggregation.SUM."): - self.evaluate(dist.unwrap(dist.call_for_each_tower(model_fn))) + "with the given reduce op ReduceOp.SUM."): + self.evaluate(distribution.unwrap( + distribution.extended.call_for_each_replica(model_fn))) - @test_util.run_in_graph_and_eager_modes(config=config) - def testAssignMirroredVarCrossTowerContext(self): - self._skip_eager_if_gpus_less_than(1) + def testAssignMirroredVarCrossDeviceContext(self, distribution): def var_fn(): return variable_scope.variable(1.0, name="foo") - dist = mirrored_strategy.MirroredStrategy( - ["/device:GPU:0", "/device:CPU:0"]) - - with dist.scope(): - mirrored_var = dist.call_for_each_tower(var_fn, run_concurrently=False) + with distribution.scope(): + mirrored_var = distribution.extended.call_for_each_replica(var_fn) self.assertIsInstance(mirrored_var, values.MirroredVariable) self.evaluate(variables.global_variables_initializer()) - self.assertEquals(1.0, self.evaluate(mirrored_var)) + self.assertEqual(1.0, self.evaluate(mirrored_var)) mirrored_var_result = self.evaluate(mirrored_var.assign(6.0)) - self.assertEquals(6.0, mirrored_var_result) + self.assertEqual(6.0, mirrored_var_result) - @test_util.run_in_graph_and_eager_modes(config=config) - def testAssignMirroredVarTowerContext(self): - self._skip_eager_if_gpus_less_than(1) + def testAssignMirroredVarReplicaContext(self, distribution): def var_fn(): return variable_scope.variable( 1.0, name="foo", aggregation=variable_scope.VariableAggregation.MEAN) - dist = mirrored_strategy.MirroredStrategy( - ["/device:GPU:0", "/device:CPU:0"]) - - with dist.scope(): - mirrored_var = dist.call_for_each_tower(var_fn, run_concurrently=False) + with distribution.scope(): + mirrored_var = distribution.extended.call_for_each_replica(var_fn) self.assertIsInstance(mirrored_var, values.MirroredVariable) self.evaluate(variables.global_variables_initializer()) - self.assertEquals(1.0, self.evaluate(mirrored_var)) + self.assertEqual(1.0, self.evaluate(mirrored_var)) def model_fn(): value = math_ops.cast( - distribution_strategy_context.get_tower_context().tower_id, + ds_context.get_replica_context().replica_id_in_sync_group, mirrored_var.dtype) return mirrored_var.assign(value) - self.evaluate(dist.unwrap(dist.call_for_each_tower( - model_fn, run_concurrently=False))) - self.assertEquals(0.5, self.evaluate(mirrored_var)) + self.evaluate(distribution.unwrap( + distribution.extended.call_for_each_replica(model_fn))) + self.assertEqual(0.5, self.evaluate(mirrored_var)) - @test_util.run_in_graph_and_eager_modes(config=config) - def testAssignMirroredVarTowerContextWithSingleValue(self): - self._skip_eager_if_gpus_less_than(1) + def testAssignMirroredVarReplicaContextWithSingleValue(self, distribution): def var_fn(): return variable_scope.variable( 1.0, name="foo", aggregation=variable_scope.VariableAggregation.MEAN) - dist = mirrored_strategy.MirroredStrategy( - ["/device:GPU:0", "/device:CPU:0"]) - - with dist.scope(): - mirrored_var = dist.call_for_each_tower(var_fn, run_concurrently=False) + with distribution.scope(): + mirrored_var = distribution.extended.call_for_each_replica(var_fn) self.assertIsInstance(mirrored_var, values.MirroredVariable) self.evaluate(variables.global_variables_initializer()) - self.assertEquals(1.0, self.evaluate(mirrored_var)) + self.assertEqual(1.0, self.evaluate(mirrored_var)) def model_fn(): return mirrored_var.assign(5.0) - self.evaluate(dist.unwrap(dist.call_for_each_tower( - model_fn, run_concurrently=False))) - self.assertEquals(5.0, self.evaluate(mirrored_var)) + self.evaluate(distribution.unwrap( + distribution.extended.call_for_each_replica(model_fn))) + self.assertEqual(5.0, self.evaluate(mirrored_var)) - @test_util.run_in_graph_and_eager_modes(config=config) - def testAssignAddMirroredVarCrossTowerContext(self): - self._skip_eager_if_gpus_less_than(1) + def testAssignAddMirroredVarCrossDeviceContext(self, distribution): def var_fn(): return variable_scope.variable(1.0, name="foo") - dist = mirrored_strategy.MirroredStrategy( - ["/device:GPU:0", "/device:CPU:0"]) - - with dist.scope(): - mirrored_var = dist.call_for_each_tower(var_fn, run_concurrently=False) + with distribution.scope(): + mirrored_var = distribution.extended.call_for_each_replica(var_fn) self.assertIsInstance(mirrored_var, values.MirroredVariable) self.evaluate(variables.global_variables_initializer()) - self.assertEquals(1.0, self.evaluate(mirrored_var)) + self.assertEqual(1.0, self.evaluate(mirrored_var)) # read_value == True mirrored_var_result = self.evaluate( mirrored_var.assign_add(6.0, read_value=True)) - self.assertEquals(7.0, mirrored_var_result) - self.assertEquals(7.0, self.evaluate(mirrored_var.get("/device:CPU:0"))) - self.assertEquals(7.0, self.evaluate(mirrored_var.get("/device:GPU:0"))) + self.assertEqual(7.0, mirrored_var_result) + self.assertEqual(7.0, self.evaluate(mirrored_var.get("/device:CPU:0"))) + self.assertEqual(7.0, self.evaluate(mirrored_var.get("/device:GPU:0"))) # read_value == False self.evaluate(mirrored_var.assign_add(2.0, read_value=False)) - self.assertEquals(9.0, self.evaluate(mirrored_var.get("/device:CPU:0"))) - self.assertEquals(9.0, self.evaluate(mirrored_var.get("/device:GPU:0"))) + self.assertEqual(9.0, self.evaluate(mirrored_var.get("/device:CPU:0"))) + self.assertEqual(9.0, self.evaluate(mirrored_var.get("/device:GPU:0"))) - @test_util.run_in_graph_and_eager_modes(config=config) - def testAssignAddMirroredVarTowerContext(self): - self._skip_eager_if_gpus_less_than(1) + def testAssignAddMirroredVarReplicaContext(self, distribution): def var_fn(): return variable_scope.variable( 1.0, name="foo", aggregation=variable_scope.VariableAggregation.MEAN) - dist = mirrored_strategy.MirroredStrategy( - ["/device:GPU:0", "/device:CPU:0"]) - - with dist.scope(): - mirrored_var = dist.call_for_each_tower(var_fn, run_concurrently=False) + with distribution.scope(): + mirrored_var = distribution.extended.call_for_each_replica(var_fn) self.assertIsInstance(mirrored_var, values.MirroredVariable) self.evaluate(variables.global_variables_initializer()) - self.assertEquals(1.0, self.evaluate(mirrored_var)) + self.assertEqual(1.0, self.evaluate(mirrored_var)) def model_fn(): value = math_ops.cast( - distribution_strategy_context.get_tower_context().tower_id, + ds_context.get_replica_context().replica_id_in_sync_group, mirrored_var.dtype) return mirrored_var.assign_add(value) - self.evaluate(dist.unwrap(dist.call_for_each_tower( - model_fn, run_concurrently=False))) - self.assertEquals(1.5, self.evaluate(mirrored_var)) + self.evaluate(distribution.unwrap( + distribution.extended.call_for_each_replica(model_fn))) + self.assertEqual(1.5, self.evaluate(mirrored_var)) - @test_util.run_in_graph_and_eager_modes(config=config) - def testAssignAddMirroredVarTowerContextWithSingleValue(self): - self._skip_eager_if_gpus_less_than(1) + def testAssignAddMirroredVarReplicaContextWithSingleValue(self, distribution): def var_fn(): return variable_scope.variable( 1.0, name="foo", aggregation=variable_scope.VariableAggregation.MEAN) - dist = mirrored_strategy.MirroredStrategy( - ["/device:GPU:0", "/device:CPU:0"]) - - with dist.scope(): - mirrored_var = dist.call_for_each_tower(var_fn, run_concurrently=False) + with distribution.scope(): + mirrored_var = distribution.extended.call_for_each_replica(var_fn) self.assertIsInstance(mirrored_var, values.MirroredVariable) self.evaluate(variables.global_variables_initializer()) - self.assertEquals(1.0, self.evaluate(mirrored_var)) + self.assertEqual(1.0, self.evaluate(mirrored_var)) def model_fn(): return mirrored_var.assign_add(5.0) - self.evaluate(dist.unwrap(dist.call_for_each_tower( - model_fn, run_concurrently=False))) - self.assertEquals(6.0, self.evaluate(mirrored_var)) + self.evaluate(distribution.unwrap( + distribution.extended.call_for_each_replica(model_fn))) + self.assertEqual(6.0, self.evaluate(mirrored_var)) - @test_util.run_in_graph_and_eager_modes(config=config) - def testAssignSubMirroredVarCrossTowerContext(self): - self._skip_eager_if_gpus_less_than(1) + def testAssignSubMirroredVarCrossDeviceContext(self, distribution): def var_fn(): return variable_scope.variable(5.0, name="foo") - dist = mirrored_strategy.MirroredStrategy( - ["/device:GPU:0", "/device:CPU:0"]) - - with dist.scope(): - mirrored_var = dist.call_for_each_tower(var_fn, run_concurrently=False) + with distribution.scope(): + mirrored_var = distribution.extended.call_for_each_replica(var_fn) self.assertIsInstance(mirrored_var, values.MirroredVariable) self.evaluate(variables.global_variables_initializer()) - self.assertEquals(5.0, self.evaluate(mirrored_var)) + self.assertEqual(5.0, self.evaluate(mirrored_var)) mirrored_var_result = self.evaluate(mirrored_var.assign_sub(2.0)) - self.assertEquals(3.0, mirrored_var_result) - self.assertEquals(3.0, self.evaluate(mirrored_var.get("/device:GPU:0"))) - self.assertEquals(3.0, self.evaluate(mirrored_var.get("/device:CPU:0"))) + self.assertEqual(3.0, mirrored_var_result) + self.assertEqual(3.0, self.evaluate(mirrored_var.get("/device:GPU:0"))) + self.assertEqual(3.0, self.evaluate(mirrored_var.get("/device:CPU:0"))) - @test_util.run_in_graph_and_eager_modes(config=config) - def testAssignSubMirroredVarTowerContext(self): - self._skip_eager_if_gpus_less_than(1) + def testAssignSubMirroredVarReplicaContext(self, distribution): def var_fn(): return variable_scope.variable( 5.0, name="foo", aggregation=variable_scope.VariableAggregation.MEAN) - dist = mirrored_strategy.MirroredStrategy( - ["/device:GPU:0", "/device:CPU:0"]) - - with dist.scope(): - mirrored_var = dist.call_for_each_tower(var_fn, run_concurrently=False) + with distribution.scope(): + mirrored_var = distribution.extended.call_for_each_replica(var_fn) self.assertIsInstance(mirrored_var, values.MirroredVariable) self.evaluate(variables.global_variables_initializer()) - self.assertEquals(5.0, self.evaluate(mirrored_var)) + self.assertEqual(5.0, self.evaluate(mirrored_var)) def model_fn(): value = math_ops.cast( - distribution_strategy_context.get_tower_context().tower_id, + ds_context.get_replica_context().replica_id_in_sync_group, mirrored_var.dtype) return mirrored_var.assign_sub(value) - self.evaluate(dist.unwrap(dist.call_for_each_tower( - model_fn, run_concurrently=False))) - self.assertEquals(4.5, self.evaluate(mirrored_var)) + self.evaluate(distribution.unwrap( + distribution.extended.call_for_each_replica(model_fn))) + self.assertEqual(4.5, self.evaluate(mirrored_var)) - @test_util.run_in_graph_and_eager_modes(config=config) - def testAssignSubMirroredVarTowerContextWithSingleValue(self): - self._skip_eager_if_gpus_less_than(1) + def testAssignSubMirroredVarReplicaContextWithSingleValue(self, distribution): def var_fn(): return variable_scope.variable( 5.0, name="foo", aggregation=variable_scope.VariableAggregation.MEAN) - dist = mirrored_strategy.MirroredStrategy( - ["/device:GPU:0", "/device:CPU:0"]) - - with dist.scope(): - mirrored_var = dist.call_for_each_tower(var_fn, run_concurrently=False) + with distribution.scope(): + mirrored_var = distribution.extended.call_for_each_replica(var_fn) self.assertIsInstance(mirrored_var, values.MirroredVariable) self.evaluate(variables.global_variables_initializer()) - self.assertEquals(5.0, self.evaluate(mirrored_var)) + self.assertEqual(5.0, self.evaluate(mirrored_var)) def model_fn(): return mirrored_var.assign_sub(1.0) - self.evaluate(dist.unwrap(dist.call_for_each_tower( - model_fn, run_concurrently=False))) - self.assertEquals(4.0, self.evaluate(mirrored_var)) + self.evaluate(distribution.unwrap( + distribution.extended.call_for_each_replica(model_fn))) + self.assertEqual(4.0, self.evaluate(mirrored_var)) -class MirroredAndTowerLocalVariableInitializerTest(test.TestCase): - config = config_pb2.ConfigProto() - config.allow_soft_placement = True +@combinations.generate(combinations.combine( + distribution=[ + combinations.mirrored_strategy_with_gpu_and_cpu, + combinations.core_mirrored_strategy_with_gpu_and_cpu], + mode=["graph", "eager"])) +class MirroredAndReplicaLocalVariableInitializerTest(test.TestCase): - def testAssignMirroredVarInitializer(self): + def testAssignMirroredVarInitializer(self, distribution): # This test is not eager compatible since in eager variables are initialized # upon construction instead of once the initialization op is run. with context.graph_mode(): @@ -1148,17 +1114,14 @@ class MirroredAndTowerLocalVariableInitializerTest(test.TestCase): v = variable_scope.variable(1.0, name="foo") return v - dist = mirrored_strategy.MirroredStrategy( - ["/device:GPU:0", "/device:CPU:0"]) - - with dist.scope(): - mirrored_var = dist.call_for_each_tower(var_fn) + with distribution.scope(): + mirrored_var = distribution.extended.call_for_each_replica(var_fn) self.assertIsInstance(mirrored_var, values.MirroredVariable) self.assertFalse(self.evaluate(mirrored_var.is_initialized())) self.evaluate(mirrored_var.initializer) self.assertTrue(self.evaluate(mirrored_var.is_initialized())) - def testAssignTowerLocalVarInitializer(self): + def testAssignReplicaLocalVarInitializer(self, distribution): # This test is not eager compatible since in eager variables are initialized # upon construction instead of once the initialization op is run. with context.graph_mode(): @@ -1167,31 +1130,27 @@ class MirroredAndTowerLocalVariableInitializerTest(test.TestCase): 1.0, synchronization=variable_scope.VariableSynchronization.ON_READ, aggregation=variable_scope.VariableAggregation.SUM) - self.assertTrue(isinstance(v_sum, values.TowerLocalVariable)) + self.assertTrue(isinstance(v_sum, values.ReplicaLocalVariable)) return v_sum - dist = mirrored_strategy.MirroredStrategy( - ["/device:GPU:0", "/device:CPU:0"]) - - with dist.scope(): - tower_local_var = dist.call_for_each_tower(model_fn) - self.assertTrue(isinstance(tower_local_var, values.TowerLocalVariable)) - self.assertFalse(self.evaluate(tower_local_var.is_initialized())) - self.evaluate(tower_local_var.initializer) - self.assertTrue(self.evaluate(tower_local_var.is_initialized())) + with distribution.scope(): + replica_local_var = distribution.extended.call_for_each_replica( + model_fn) + self.assertTrue(isinstance(replica_local_var, + values.ReplicaLocalVariable)) + self.assertFalse(self.evaluate(replica_local_var.is_initialized())) + self.evaluate(replica_local_var.initializer) + self.assertTrue(self.evaluate(replica_local_var.is_initialized())) -class TowerLocalVariableAssignTest(test.TestCase): - config = config_pb2.ConfigProto() - config.allow_soft_placement = True +@combinations.generate(combinations.combine( + distribution=[ + combinations.mirrored_strategy_with_gpu_and_cpu, + combinations.core_mirrored_strategy_with_gpu_and_cpu], + mode=["graph", "eager"])) +class ReplicaLocalVariableAssignTest(test.TestCase): - def _skip_eager_if_gpus_less_than(self, num_gpus): - if context.num_gpus() < num_gpus and context.executing_eagerly(): - self.skipTest("Not enough GPUs available for this test in eager mode.") - - @test_util.run_in_graph_and_eager_modes(config=config) - def testAssignTowerLocalVarSumAggregation(self): - self._skip_eager_if_gpus_less_than(1) + def testAssignReplicaLocalVarSumAggregation(self, distribution): def model_fn(): v_sum = variable_scope.variable( 1.0, @@ -1199,30 +1158,27 @@ class TowerLocalVariableAssignTest(test.TestCase): aggregation=variable_scope.VariableAggregation.SUM) return v_sum - dist = mirrored_strategy.MirroredStrategy( - ["/device:GPU:0", "/device:CPU:0"]) - - with dist.scope(): - tower_local_var = dist.call_for_each_tower(model_fn, - run_concurrently=False) - self.assertTrue(isinstance(tower_local_var, values.TowerLocalVariable)) + with distribution.scope(): + replica_local_var = distribution.extended.call_for_each_replica(model_fn) + self.assertTrue(isinstance(replica_local_var, + values.ReplicaLocalVariable)) self.evaluate(variables.global_variables_initializer()) - # Each tower has a value of 1.0 assigned to it in tower context. + # Each replica has a value of 1.0 assigned to it in replica context. # When we read the value using `read_var` we should see the SUM of each of - # values on each of the towers. - self.assertEqual(2.0, self.evaluate(dist.read_var(tower_local_var))) - # Assigning 6.0 in cross tower context will assign a value of - # 6.0/num_towers to each tower. - tlv_ops = tower_local_var.assign(6.0) + # values on each of the replicas. + self.assertEqual(2.0, self.evaluate( + distribution.extended.read_var(replica_local_var))) + # Assigning 6.0 in cross replica context will assign a value of + # 6.0/num_replicas to each replica. + tlv_ops = replica_local_var.assign(6.0) self.evaluate(tlv_ops) - # On reading the tower local var we should get the assigned value back. - # The value on all the towers are added before being returned by + # On reading the replica local var we should get the assigned value back. + # The value on all the replicas are added before being returned by # `read_var`. - self.assertEqual(6.0, self.evaluate(dist.read_var(tower_local_var))) + self.assertEqual(6.0, self.evaluate( + distribution.extended.read_var(replica_local_var))) - @test_util.run_in_graph_and_eager_modes(config=config) - def testAssignTowerLocalVarMeanAggregation(self): - self._skip_eager_if_gpus_less_than(1) + def testAssignReplicaLocalVarMeanAggregation(self, distribution): def model_fn(): v_sum = variable_scope.variable( 1.0, @@ -1230,23 +1186,22 @@ class TowerLocalVariableAssignTest(test.TestCase): aggregation=variable_scope.VariableAggregation.MEAN) return v_sum - dist = mirrored_strategy.MirroredStrategy( - ["/device:GPU:0", "/device:CPU:0"]) - - with dist.scope(): - tower_local_var = dist.call_for_each_tower(model_fn, - run_concurrently=False) - self.assertTrue(isinstance(tower_local_var, values.TowerLocalVariable)) + with distribution.scope(): + replica_local_var = distribution.extended.call_for_each_replica(model_fn) + self.assertTrue(isinstance(replica_local_var, + values.ReplicaLocalVariable)) self.evaluate(variables.global_variables_initializer()) - # Each tower has a value of 1.0 assigned to it in tower context. + # Each replica has a value of 1.0 assigned to it in replica context. # When we read the value using `read_var` we should see the MEAN of values - # on all towers which is the value assigned in tower context. - self.assertEqual(1.0, self.evaluate(dist.read_var(tower_local_var))) - tlv_ops = tower_local_var.assign(6.0) + # on all replicas which is the value assigned in replica context. + self.assertEqual(1.0, self.evaluate( + distribution.extended.read_var(replica_local_var))) + tlv_ops = replica_local_var.assign(6.0) self.evaluate(tlv_ops) - # On reading the tower local var we should get the MEAN of all values + # On reading the replica local var we should get the MEAN of all values # which is equal to the value assigned. - self.assertEqual(6.0, self.evaluate(dist.read_var(tower_local_var))) + self.assertEqual(6.0, self.evaluate( + distribution.extended.read_var(replica_local_var))) class MockModel(object): @@ -1280,48 +1235,45 @@ class MiniModel(keras_training.Model): return self.fc(inputs) +@combinations.generate(combinations.combine( + distribution=[ + combinations.mirrored_strategy_with_gpu_and_cpu, + combinations.core_mirrored_strategy_with_gpu_and_cpu], + mode=["graph", "eager"])) class MirroredStrategyDefunTest(test.TestCase): - def _skip_eager_if_gpus_less_than(self, num_gpus): - if context.num_gpus() < num_gpus and context.executing_eagerly(): - self.skipTest("Not enough GPUs available for this test in eager mode.") - - def _call_and_check(self, model_fn, inputs, expected_result, defuns, - two_variables=False): + def _call_and_check(self, distribution, model_fn, inputs, expected_result, + defuns, two_variables=False): cpu_dev = device_util.canonicalize("CPU:0") gpu_dev = device_util.canonicalize("GPU:0") devices = [cpu_dev, gpu_dev] - dist = mirrored_strategy.MirroredStrategy(devices) - with dist.scope(): + with distribution.scope(): mock_model = MockModel(two_variables) self.evaluate(variables.global_variables_initializer()) - result = dist.call_for_each_tower(model_fn, mock_model, *inputs, - run_concurrently=False) - for device in devices: - device_result = values.select_device(device, result) - device_expected_result = values.select_device(device, expected_result) + result = distribution.extended.call_for_each_replica( + model_fn, args=[mock_model] + inputs) + for r in range(len(devices)): + device_result = values.select_replica(r, result) + device_expected_result = values.select_replica(r, expected_result) self.assertAllClose(device_expected_result, 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. - per_device_graph_functions = dist.call_for_each_tower( - defun.get_concrete_function, - mock_model, *inputs, run_concurrently=False) + per_replica_graph_functions = ( + distribution.extended.call_for_each_replica( + defun.get_concrete_function, args=[mock_model] + inputs)) for device in devices: - graph_function = per_device_graph_functions.get(device=device) + graph_function = per_replica_graph_functions.get(device=device) self.assertEqual(set(mock_model.variables), set(graph_function.graph.variables)) - @test_util.run_in_graph_and_eager_modes() - def testVariableInDefun(self): - self._skip_eager_if_gpus_less_than(1) - + def testVariableInDefun(self, distribution): @function.defun def times_two(mock_model): return mock_model() @@ -1329,12 +1281,9 @@ class MirroredStrategyDefunTest(test.TestCase): def model_fn(mock_model): return times_two(mock_model) - self._call_and_check(model_fn, [], 2.5, [times_two]) - - @test_util.run_in_graph_and_eager_modes() - def testVariableInNestedDefun(self): - self._skip_eager_if_gpus_less_than(1) + self._call_and_check(distribution, model_fn, [], 2.5, [times_two]) + def testVariableInNestedDefun(self, distribution): @function.defun def times_two(mock_model): return mock_model() @@ -1346,12 +1295,10 @@ class MirroredStrategyDefunTest(test.TestCase): def model_fn(mock_model): return two_x_plus_one(mock_model) - self._call_and_check(model_fn, [], 3.5, [times_two, two_x_plus_one]) - - @test_util.run_in_graph_and_eager_modes() - def testTwoVariablesInNestedDefun(self): - self._skip_eager_if_gpus_less_than(1) + self._call_and_check(distribution, model_fn, [], 3.5, + [times_two, two_x_plus_one]) + def testTwoVariablesInNestedDefun(self, distribution): @function.defun def fn1(mock_model): return mock_model() @@ -1363,12 +1310,10 @@ class MirroredStrategyDefunTest(test.TestCase): def model_fn(mock_model): return fn2(mock_model) - self._call_and_check(model_fn, [], 5.5, [fn1, fn2], two_variables=True) - - @test_util.run_in_graph_and_eager_modes() - def testGradientTapeOverNestedDefuns(self): - self._skip_eager_if_gpus_less_than(1) + self._call_and_check(distribution, model_fn, [], 5.5, [fn1, fn2], + two_variables=True) + def testGradientTapeOverNestedDefuns(self, distribution): @function.defun def fn1(mock_model): return mock_model() @@ -1384,32 +1329,21 @@ class MirroredStrategyDefunTest(test.TestCase): [v.get() for v in mock_model.variables]) return grads - self._call_and_check(model_fn, [], [2.0, 1.0], [fn1, fn2], + self._call_and_check(distribution, model_fn, [], [2.0, 1.0], [fn1, fn2], two_variables=True) - @test_util.run_in_graph_and_eager_modes() - def testPassPerDevice(self): - self._skip_eager_if_gpus_less_than(1) - + def testPassPerReplica(self, distribution): @function.defun def fn1(mock_model, factor): return mock_model(factor) - factors = values.PerDevice({"CPU:0": 5.0, "GPU:0": 3.0}) - expected_result = values.PerDevice({"CPU:0": 5.0 * 1.25, - "GPU:0": 3.0 * 1.25}) - self._call_and_check(fn1, [factors], expected_result, [fn1]) - - @test_util.run_in_graph_and_eager_modes() - def testTrain(self): - self._skip_eager_if_gpus_less_than(1) + device_map = values.ReplicaDeviceMap(("/device:CPU:0", "/device:GPU:0")) + factors = values.PerReplica(device_map, (5.0, 3.0)) + expected_result = values.PerReplica(device_map, (5.0 * 1.25, 3.0 * 1.25)) + self._call_and_check(distribution, fn1, [factors], expected_result, [fn1]) - cpu_dev = device_util.canonicalize("CPU:0") - gpu_dev = device_util.canonicalize("GPU:0") - devices = [cpu_dev, gpu_dev] - dist = mirrored_strategy.MirroredStrategy(devices) - - with dist.scope(): + def testTrain(self, distribution): + with distribution.scope(): mock_model = MiniModel() mock_model.call = function.defun(mock_model.call) @@ -1419,11 +1353,11 @@ class MirroredStrategyDefunTest(test.TestCase): gradients_fn = backprop.implicit_grad(loss_fn) gradients_fn = optimizer_lib.get_filtered_grad_fn(gradients_fn) - grads_and_vars = dist.call_for_each_tower( - gradients_fn, None, run_concurrently=False) + grads_and_vars = distribution.extended.call_for_each_replica( + gradients_fn, args=(None,)) optimizer = gradient_descent.GradientDescentOptimizer(0.25) - update_ops = optimizer._distributed_apply(dist, grads_and_vars) # pylint: disable=protected-access + update_ops = optimizer._distributed_apply(distribution, grads_and_vars) # pylint: disable=protected-access if not context.executing_eagerly(): self.evaluate(variables.global_variables_initializer()) @@ -1435,27 +1369,82 @@ class MirroredStrategyDefunTest(test.TestCase): self.assertAllEqual([0.5], updated_var_values[1]) +@combinations.generate( + combinations.combine( + distribution=[ + combinations.NamedDistribution( + "Mirrored", + # pylint: disable=g-long-lambda + lambda: mirrored_strategy.MirroredStrategy(num_gpus_per_worker= + context.num_gpus()), + required_gpus=1), + combinations.NamedDistribution( + "CoreMirrored", + # pylint: disable=g-long-lambda + lambda: mirrored_strategy.CoreMirroredStrategy( + mirrored_strategy.all_local_devices()), + required_gpus=1) + ], + mode=["graph"])) class MultiWorkerMirroredStrategyTest( multi_worker_test_base.MultiWorkerTestBase, strategy_test_lib.DistributionTestBase): - def _get_distribution_strategy(self): + def _configure_distribution_strategy(self, distribution): cluster_spec = server_lib.ClusterSpec({ "worker": ["/job:worker/task:0", "/job:worker/task:1"] }) - strategy = mirrored_strategy.MirroredStrategy(num_gpus=context.num_gpus()) - strategy.configure(cluster_spec=cluster_spec) - return strategy + distribution.configure(cluster_spec=cluster_spec) - def test_num_replicas_in_sync(self): - strategy = self._get_distribution_strategy() + def test_num_replicas_in_sync(self, distribution): + self._configure_distribution_strategy(distribution) # We calculate the total number of gpus across the workers(2) specified in # the cluster spec. - self.assertEqual(context.num_gpus() * 2, strategy.num_replicas_in_sync) - - def testMinimizeLossGraph(self): - self._test_minimize_loss_graph(self._get_distribution_strategy(), - learning_rate=0.05) + self.assertEqual(context.num_gpus() * 2, distribution.num_replicas_in_sync) + + def testMinimizeLossGraph(self, distribution): + self._configure_distribution_strategy(distribution) + self._test_minimize_loss_graph(distribution, learning_rate=0.05) + + def testDeviceScope(self, distribution): + """Test the device scope of multi-worker MirroredStrategy.""" + self._configure_distribution_strategy(distribution) + with distribution.scope(): + a = constant_op.constant(1.) + with ops.device("/cpu:0"): + b = constant_op.constant(1.) + self.assertEqual(a.device, "/job:worker/task:0") + self.assertEqual(b.device, "/job:worker/task:0/device:CPU:0") + + def testMakeInputFnIterator(self, distribution): + self._configure_distribution_strategy(distribution) + dataset_fn = lambda: dataset_ops.Dataset.range(100) + num_gpus = context.num_gpus() + num_workers = 2 + + expected_values = [[i+j for j in range(num_gpus)] * num_workers + for i in range(0, 100, num_gpus)] + + 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( + dataset_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) + + def testUpdateConfigProto(self, distribution): + distribution.configure(cluster_spec={"worker": ["fake1", "fake2"]}) + + config_proto = config_pb2.ConfigProto() + new_config = distribution.update_config_proto(config_proto) + + # Verify isolate_session_state + self.assertTrue(new_config.isolate_session_state) class MultiWorkerMirroredStrategyTestWithChief( @@ -1475,6 +1464,19 @@ class MultiWorkerMirroredStrategyTestWithChief( strategy.configure(cluster_spec=self._cluster_spec) self._test_minimize_loss_graph(strategy, learning_rate=0.05) + def testMinimizeLossGraphCoreMirroredStrategy(self): + strategy = mirrored_strategy.CoreMirroredStrategy( + mirrored_strategy.all_local_devices()) + strategy.configure(cluster_spec=self._cluster_spec) + self._test_minimize_loss_graph(strategy, learning_rate=0.05) + + +def _replica_id(): + replica_id = ds_context.get_replica_context().replica_id_in_sync_group + if not isinstance(replica_id, ops.Tensor): + replica_id = constant_op.constant(replica_id) + return replica_id + if __name__ == "__main__": test.main() diff --git a/tensorflow/contrib/distribute/python/mirrored_strategy_test.py b/tensorflow/contrib/distribute/python/mirrored_strategy_test.py deleted file mode 100644 index 969e1269560e52736d05e6b14ce320d9bd4fcac0..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/distribute/python/mirrored_strategy_test.py +++ /dev/null @@ -1,109 +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. -# ============================================================================== -"""Tests for class MirroredStrategy.""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -from tensorflow.contrib.distribute.python import mirrored_strategy -from tensorflow.contrib.distribute.python import strategy_test_lib -from tensorflow.python.eager import context -from tensorflow.python.eager import test -from tensorflow.python.framework import constant_op -from tensorflow.python.framework import ops -from tensorflow.python.framework import test_util -from tensorflow.python.ops import variable_scope -from tensorflow.python.training import distribution_strategy_context - - -class MirroredOneCPUDistributionTest(strategy_test_lib.DistributionTestBase): - - def _get_distribution_strategy(self): - return mirrored_strategy.MirroredStrategy(["/device:CPU:0"]) - - def testMinimizeLossEager(self): - self._test_minimize_loss_eager(self._get_distribution_strategy()) - - def testMinimizeLossGraph(self): - self._test_minimize_loss_graph(self._get_distribution_strategy()) - - def testMapReduce(self): - self._test_map_reduce(self._get_distribution_strategy()) - - def testDeviceIndex(self): - self._test_device_index(self._get_distribution_strategy()) - - def testTowerId(self): - self._test_tower_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()) - - -class VariableCreatorStackTest(test.TestCase): - - def testCreatorStacksAreThreadLocal(self): - devices = ["/device:CPU:0", "/device:GPU:0"] - dist = mirrored_strategy.MirroredStrategy(devices) - - def model_fn(device_id): - assert isinstance(device_id, int) - - def thread_creator_fn(next_creator, *args, **kwargs): - return next_creator(*args, **kwargs) + ":thread_" + str(device_id) - - with variable_scope.variable_creator_scope(thread_creator_fn): - # Create a variable in this scope. - v = variable_scope.variable(1.0) - - # This will pause the current thread, and execute the other thread. - distribution_strategy_context.get_tower_context().merge_call( - lambda _: _) - return v - - def main_thread_creator(next_creator, *args, **kwargs): - # We are not using the underlying next_creator for test purposes. - del next_creator, args, kwargs - return "main_thread" - - with context.graph_mode(), \ - dist.scope(), \ - variable_scope.variable_creator_scope(main_thread_creator): - result = dist.call_for_each_tower(model_fn, dist.worker_device_index) - result = dist.unwrap(result) - expected = ["main_thread:thread_0", "main_thread:thread_1"] - self.assertEquals(expected, result) - - -class MultiWorkerMirroredStrategyTest(test.TestCase): - - def testDeviceScope(self): - """Test the device scope of multi-worker MirroredStrategy.""" - with context.graph_mode(): - strategy = mirrored_strategy.MirroredStrategy(num_gpus=context.num_gpus()) - strategy.configure( - cluster_spec={"worker": ["/job:worker/task:0", "/job:worker/task:1"]}) - with strategy.scope(): - a = constant_op.constant(1.) - with ops.device("/cpu:0"): - b = constant_op.constant(1.) - self.assertEqual(a.device, "/job:worker/task:0") - self.assertEqual(b.device, "/job:worker/task:0/device:CPU:0") - - -if __name__ == "__main__": - test.main() 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/moving_averages_test.py b/tensorflow/contrib/distribute/python/moving_averages_test.py index 119352ad9195dc51201863f34aef19cb3289e635..c4622cdd2af2f6a9c936fe554bcc2eb76f805fdc 100644 --- a/tensorflow/contrib/distribute/python/moving_averages_test.py +++ b/tensorflow/contrib/distribute/python/moving_averages_test.py @@ -32,33 +32,34 @@ from tensorflow.python.training import moving_averages all_combinations = combinations.combine( distribution=[combinations.default_strategy, combinations.one_device_strategy, - combinations.mirrored_strategy_with_gpu_and_cpu], + combinations.mirrored_strategy_with_gpu_and_cpu, + combinations.core_mirrored_strategy_with_gpu_and_cpu], mode=["graph"]) class AssignMovingAveragesTest(test.TestCase, parameterized.TestCase): @combinations.generate(all_combinations) - def testTowerModeWithoutZeroDebias(self, distribution): - tower_id = [0] + def testReplicaModeWithoutZeroDebias(self, distribution): + replica_id = [0] - def tower_fn(): + def replica_fn(): var = variables.Variable([10.0, 11.0]) - val = constant_op.constant([1.0 + tower_id[0], 2.0 - tower_id[0]]) - tower_id[0] += 1 + val = constant_op.constant([1.0 + replica_id[0], 2.0 - replica_id[0]]) + replica_id[0] += 1 decay = 0.25 assign = moving_averages.assign_moving_average( var, val, decay, zero_debias=False) return var, assign with distribution.scope(), self.cached_session() as sess: - var, assign = distribution.call_for_each_tower(tower_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)) - # Mean of val across calls to tower_fn(). - average_val = [1.0 + 0.5 * (tower_id[0] - 1), - 2.0 - 0.5 * (tower_id[0] - 1)] + # Mean of val across calls to replica_fn(). + average_val = [1.0 + 0.5 * (replica_id[0] - 1), + 2.0 - 0.5 * (replica_id[0] - 1)] val_weight = 1.0 - 0.25 self.assertAllClose( [10.0 * 0.25 + average_val[0] * val_weight, @@ -66,34 +67,35 @@ class AssignMovingAveragesTest(test.TestCase, parameterized.TestCase): var.eval()) @combinations.generate(all_combinations) - def testTowerMode(self, distribution): - tower_id = [0] + def testReplicaMode(self, distribution): + replica_id = [0] - def tower_fn(): + def replica_fn(): var = variables.Variable([0.0, 0.0]) - val = constant_op.constant([1.0 + tower_id[0], 2.0 - tower_id[0]]) - tower_id[0] += 1 + val = constant_op.constant([1.0 + replica_id[0], 2.0 - replica_id[0]]) + replica_id[0] += 1 decay = 0.25 assign = moving_averages.assign_moving_average(var, val, decay) return var, assign.op with distribution.scope(), self.cached_session() as sess: - var, assign_op = distribution.call_for_each_tower(tower_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)) - # Mean of val across calls to tower_fn(). - average_val = [1.0 + 0.5 * (tower_id[0] - 1), - 2.0 - 0.5 * (tower_id[0] - 1)] + # Mean of val across calls to replica_fn(). + average_val = [1.0 + 0.5 * (replica_id[0] - 1), + 2.0 - 0.5 * (replica_id[0] - 1)] self.assertAllClose(average_val, var.eval()) @combinations.generate(all_combinations) - def testCrossTowerWithoutZeroDebias(self, distribution): + def testCrossDeviceWithoutZeroDebias(self, distribution): with distribution.scope(), self.cached_session() as sess: var = variables.Variable([10.0, 11.0]) val = constant_op.constant([1.0, 2.0]) decay = 0.25 - # NOTE(josh11b): We currently generate an error if val is a PerDevice value. + # NOTE(josh11b): We currently generate an error if val is a PerReplica + # value. assign = moving_averages.assign_moving_average( var, val, decay, zero_debias=False) @@ -116,12 +118,13 @@ class AssignMovingAveragesTest(test.TestCase, parameterized.TestCase): var.eval()) @combinations.generate(all_combinations) - def testCrossTower(self, distribution): + def testCrossDevice(self, distribution): with distribution.scope(), self.cached_session() as sess: var = variables.Variable([0.0, 0.0]) val = array_ops.placeholder(dtypes.float32) decay = 0.25 - # NOTE(josh11b): We currently generate an error if val is a PerDevice value. + # NOTE(josh11b): We currently generate an error if val is a PerReplica + # value. assign = moving_averages.assign_moving_average(var, val, decay) variables.global_variables_initializer().run() @@ -136,6 +139,27 @@ class AssignMovingAveragesTest(test.TestCase, parameterized.TestCase): (2.0 * 0.25 + 0.0) / (1.0 * 0.25 + 1.0)], var.eval()) + @combinations.generate(all_combinations) + def testAssignVariable(self, distribution): + + def replica_fn(): + var = variables.Variable([10.0, 11.0]) + # Here we expect to check the case when input value are variable. + val = variables.Variable([1., 2.]) + decay = 0.25 + assign = moving_averages.assign_moving_average( + var, val, decay, zero_debias=False) + return var, assign + + with distribution.scope(), self.cached_session() as sess: + 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)) + self.assertAllClose( + [10 * 0.25 + 1. * (1 - 0.25), 11 * 0.25 + 2. * (1 - 0.25)], + var.eval()) + if __name__ == "__main__": test.main() diff --git a/tensorflow/contrib/distribute/python/multi_worker_test_base.py b/tensorflow/contrib/distribute/python/multi_worker_test_base.py index 9f92ba7dde5fc2798201cef2238bcc4b20b698a8..b05aac431f65b4281d9ed9c2fa95c210d55f4008 100644 --- a/tensorflow/contrib/distribute/python/multi_worker_test_base.py +++ b/tensorflow/contrib/distribute/python/multi_worker_test_base.py @@ -18,8 +18,11 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import collections import contextlib import copy +import json +import os import threading import numpy as np @@ -37,9 +40,9 @@ from tensorflow.python.client import session from tensorflow.python.estimator import run_config from tensorflow.python.platform import test from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.training import coordinator from tensorflow.python.training import server_lib - ASSIGNED_PORTS = set() lock = threading.Lock() @@ -207,12 +210,10 @@ class MultiWorkerTestBase(test.TestCase): self._lock = threading.Lock() @contextlib.contextmanager - def test_session(self, graph=None, config=None, target=None): + def session(self, graph=None, config=None, target=None): """Create a test session with master target set to the testing cluster. - This overrides the base class' method, removes arguments that are not needed - by the multi-node case and creates a test session that connects to the local - testing cluster. + Creates a test session that connects to the local testing cluster. Args: graph: Optional graph to use during the returned session. @@ -224,9 +225,44 @@ class MultiWorkerTestBase(test.TestCase): A Session object that should be used as a context manager to surround the graph building and execution code in a test case. """ - if self.id().endswith('.test_session'): - self.skipTest('Not a test.') + config = self._create_config(config) + + if target is None: + target = self._default_target + with session.Session(graph=graph, config=config, target=target) as sess: + yield sess + + @contextlib.contextmanager + # TODO(b/117573461): Overwrite self.evaluate() to use this function. + def cached_session(self, graph=None, config=None, target=None): + """Create a test session with master target set to the testing cluster. + + Creates a test session that connects to the local testing cluster. + The session is only created once per test and then reused. + + Args: + graph: Optional graph to use during the returned session. + config: An optional config_pb2.ConfigProto to use to configure the + session. + target: the target of session to connect to. + + Yields: + A Session object that should be used as a context manager to surround + the graph building and execution code in a test case. Note that the + session will live until the end of the test. + """ + config = self._create_config(config) + if target is None: + target = self._default_target + if getattr(self._thread_local, 'cached_session', None) is None: + self._thread_local.cached_session = session.Session( + graph=None, config=config, target=target) + sess = self._thread_local.cached_session + with sess.graph.as_default(), sess.as_default(): + yield sess + + def _create_config(self, config): if config is None: config = config_pb2.ConfigProto(allow_soft_placement=True) else: @@ -237,18 +273,7 @@ class MultiWorkerTestBase(test.TestCase): config.graph_options.rewrite_options.constant_folding = ( rewriter_config_pb2.RewriterConfig.OFF) - if target is None: - target = self._default_target - if graph is None: - if getattr(self._thread_local, 'cached_session', None) is None: - self._thread_local.cached_session = session.Session( - graph=None, config=config, target=target) - sess = self._thread_local.cached_session - with sess.graph.as_default(), sess.as_default(): - yield sess - else: - with session.Session(graph=graph, config=config, target=target) as sess: - yield sess + return config def _run_client(self, client_fn, task_type, task_id, num_gpus, *args, **kwargs): @@ -281,3 +306,106 @@ class MultiWorkerTestBase(test.TestCase): for t in threads: t.join() self.assertEqual(self._result, len(threads)) + + +class MockOsEnv(collections.Mapping): + """A class that allows per-thread TF_CONFIG.""" + + def __init__(self, *args): + self._dict = dict() + self._thread_local = threading.local() + super(MockOsEnv, self).__init__(*args) + + def get(self, key, default=None): + if not hasattr(self._thread_local, 'dict'): + self._thread_local.dict = dict() + if key == 'TF_CONFIG': + return dict.get(self._thread_local.dict, key, default) + else: + return dict.get(self._dict, key, default) + + def __getitem__(self, key): + if not hasattr(self._thread_local, 'dict'): + self._thread_local.dict = dict() + if key == 'TF_CONFIG': + return dict.__getitem__(self._thread_local.dict, key) + else: + return dict.__getitem__(self._dict, key) + + def __setitem__(self, key, val): + if not hasattr(self._thread_local, 'dict'): + self._thread_local.dict = dict() + if key == 'TF_CONFIG': + return dict.__setitem__(self._thread_local.dict, key, val) + else: + return dict.__setitem__(self._dict, key, val) + + def __iter__(self): + if not hasattr(self._thread_local, 'dict'): + self._thread_local.dict = dict() + for x in self._thread_local.dict.items(): + yield x + for x in self._dict.items(): + yield x + + def __len__(self): + if not hasattr(self._thread_local, 'dict'): + self._thread_local.dict = dict() + return self._thread_local.dict.__len__() + self._dict.__len__() + + +class IndependentWorkerTestBase(test.TestCase): + """Testing infra for independent workers.""" + + def setUp(self): + self._mock_os_env = MockOsEnv() + self._mock_context = test.mock.patch.object(os, 'environ', + self._mock_os_env) + self._coord = coordinator.Coordinator() + super(IndependentWorkerTestBase, self).setUp() + self._mock_context.__enter__() + + def tearDown(self): + self._mock_context.__exit__(None, None, None) + super(IndependentWorkerTestBase, self).tearDown() + + def _task_thread(self, task_fn, tf_config, *args, **kwargs): + with self._coord.stop_on_exception(): + os.environ['TF_CONFIG'] = json.dumps(tf_config) + task_fn(*args, **kwargs) + + def _run_task_in_thread(self, task_fn, cluster_spec, task_type, task_id, + *args, **kwargs): + if task_type: + tf_config = { + 'cluster': cluster_spec, + 'task': { + 'type': task_type, + 'index': task_id + } + } + else: + tf_config = { + 'cluster': cluster_spec, + } + t = threading.Thread( + target=self._task_thread, + args=(task_fn, tf_config) + args, + kwargs=kwargs) + t.start() + return t + + def run_multiple_tasks_in_threads(self, task_fn, cluster_spec, *args, + **kwargs): + # The task_fn should create std_server by itself. + threads = {} + for task_type in cluster_spec.keys(): + threads[task_type] = [] + for task_id in range(len(cluster_spec[task_type])): + t = self._run_task_in_thread(task_fn, cluster_spec, task_type, task_id, + *args, **kwargs) + threads[task_type].append(t) + return threads + + def join_independent_workers(self, worker_threads): + self._coord.join(worker_threads) diff --git a/tensorflow/contrib/distribute/python/one_device_strategy.py b/tensorflow/contrib/distribute/python/one_device_strategy.py index f5259190485e701c190beb49220caff743f8fdcb..24d6a443fe15c9b9ff34b7e6d3a5bc5a2bb7abfb 100644 --- a/tensorflow/contrib/distribute/python/one_device_strategy.py +++ b/tensorflow/contrib/distribute/python/one_device_strategy.py @@ -18,16 +18,16 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import six - -from tensorflow.contrib.distribute.python import values +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 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.ops import math_ops -from tensorflow.python.ops import variable_scope as vs -from tensorflow.python.training import distribute as distribute_lib from tensorflow.python.util import nest @@ -40,51 +40,65 @@ class OneDeviceStrategy(distribute_lib.DistributionStrategy): # doing something that won't work with other DistributionStrategy # implementations? - def __init__(self, device, prefetch_on_device=None): - super(OneDeviceStrategy, self).__init__() + 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._prefetch_on_device = prefetch_on_device self._default_device = device + self._input_device = device_util.canonicalize("/device:CPU:0") + worker_device_pairs = [(self._input_device, [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) - if isinstance(colocate_with, six.string_types): - with ops.device(colocate_with): - return next_creator(*args, **kwargs) - if (isinstance(colocate_with, list) and len(colocate_with) == 1 and - isinstance(colocate_with[0], six.string_types)): - with ops.device(colocate_with[0]): - return next_creator(*args, **kwargs) with ops.colocate_with(colocate_with): return next_creator(*args, **kwargs) - def distribute_dataset(self, dataset_fn): - return values.PerDeviceDataset( - self._call_dataset_fn(dataset_fn), [self._device], - self._prefetch_on_device) + 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 _broadcast(self, tensor, destinations): + 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 _experimental_make_numpy_dataset(self, numpy_input, session): + return numpy_dataset.one_host_numpy_dataset( + numpy_input, self._input_device, session) + + def _broadcast_to(self, tensor, destinations): del destinations return tensor # TODO(priyag): Deal with OutOfRange errors once b/111349762 is fixed. - def _run_steps_on_dataset(self, fn, iterator, iterations, - initial_loop_values=None): + 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 = values.MultiStepContext() + ctx = input_lib.MultiStepContext() def body(i, *args): """A wrapper around `fn` to create the while loop body.""" del args - fn_inputs = iterator.get_next() - if not isinstance(fn_inputs, tuple): - fn_inputs = (fn_inputs,) - fn_result = fn(ctx, *fn_inputs) + 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 @@ -117,86 +131,82 @@ class OneDeviceStrategy(distribute_lib.DistributionStrategy): ctx._set_last_step_outputs(last_step_tensor_outputs_dict) # pylint: disable=protected-access return ctx - def _call_for_each_tower(self, fn, *args, **kwargs): - # We don't run `fn` in multiple threads in OneDeviceStrategy. - kwargs.pop("run_concurrently", None) - with ops.device(self._device), _OneDeviceTowerContext(self): + 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 map(self, map_over, fn, *args, **kwargs): - with ops.device(self._device): - return values.MapOutput([fn(m, *args, **kwargs) for m in map_over]) - - def _reduce(self, aggregation, value, destinations): - del destinations - if not isinstance(value, values.MapOutput): - return value - l = value.get() - assert l - with ops.device(self._device): - if aggregation == vs.VariableAggregation.SUM: - return math_ops.add_n(l) - elif aggregation == vs.VariableAggregation.MEAN: - return math_ops.add_n(l) / len(l) - else: - assert False + def _reduce_to(self, reduce_op, value, destinations): + del reduce_op, destinations + return value - def _update(self, var, options, fn, *args, **kwargs): + 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, options, fn, var, *args, **kwargs) + return self._update_non_slot(var, fn, (var,) + tuple(args), kwargs, group) - def _update_non_slot(self, colocate_with, options, fn, *args, **kwargs): + def _update_non_slot(self, colocate_with, fn, args, kwargs, group): del colocate_with - should_group = options.pop("grouped") - assert not options # Validate that we are processing all of the options. with ops.device(self._device), distribute_lib.UpdateContext(self._device): result = fn(*args, **kwargs) - if should_group: + if group: return result else: return nest.map_structure(self._unwrap, result) - def read_var(self, tower_local_var): - """Read the aggregate value of a tower-local variable.""" - return array_ops.identity(tower_local_var) + 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] + return (value,) def value_container(self, value): return value @property - def is_single_tower(self): - return True - - @property - def num_towers(self): + def _num_replicas_in_sync(self): return 1 @property def worker_devices(self): - return [self._device] + return (self._device,) @property def parameter_devices(self): - return [self._device] + return (self._device,) def non_slot_devices(self, var_list): del var_list - return [self._device] + return (self._device,) - def _worker_device_index(self): - return 0 + @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): + """Global and per-replica batching are equivalent for OneDeviceStrategy.""" + return True -class _OneDeviceTowerContext(distribute_lib.TowerContext): +class _OneDeviceReplicaContext(distribute_lib.ReplicaContext): + """ReplicaContext for OneDeviceStrategy.""" - def __init__(self, distribution_strategy): - distribute_lib.TowerContext.__init__( - self, distribution_strategy, tower_id=0) + 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 device(self): - return self._distribution_strategy.worker_devices[0] + def devices(self): + return self._strategy.extended.worker_devices diff --git a/tensorflow/contrib/distribute/python/one_device_strategy_test.py b/tensorflow/contrib/distribute/python/one_device_strategy_test.py index 4fdc0f72e6745b7ef25c591157955f214e0b2c79..f81466a6c75f1cf287cdb00917872f77383c615e 100644 --- a/tensorflow/contrib/distribute/python/one_device_strategy_test.py +++ b/tensorflow/contrib/distribute/python/one_device_strategy_test.py @@ -20,11 +20,14 @@ from __future__ import print_function from tensorflow.contrib.distribute.python import one_device_strategy from tensorflow.contrib.distribute.python import strategy_test_lib +from tensorflow.python.data.ops import dataset_ops from tensorflow.python.eager import test from tensorflow.python.framework import test_util -class OneDeviceStrategyTest(strategy_test_lib.DistributionTestBase): +class OneDeviceStrategyTest( + strategy_test_lib.DistributionTestBase, + strategy_test_lib.OneDeviceDistributionTestBase): def _get_distribution_strategy(self): return one_device_strategy.OneDeviceStrategy("/device:CPU:0") @@ -35,19 +38,49 @@ class OneDeviceStrategyTest(strategy_test_lib.DistributionTestBase): def testMinimizeLossGraph(self): self._test_minimize_loss_graph(self._get_distribution_strategy()) - def testMapReduce(self): - self._test_map_reduce(self._get_distribution_strategy()) - - def testDeviceIndex(self): - self._test_device_index(self._get_distribution_strategy()) - - def testTowerId(self): - self._test_tower_id(self._get_distribution_strategy()) + 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") + 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( + dataset_fn, + expected_num_replicas_in_sync=1, + expected_num_input_pipelines=1, + expected_input_pipeline_id=0) + iterator = d.make_input_fn_iterator(input_fn) + self._test_input_fn_iterator( + iterator, d.extended.worker_devices, expected_values) + + @test_util.run_in_graph_and_eager_modes + def testNumpyIterator(self): + self._test_numpy_iterator(self._get_distribution_strategy()) + + def testAllReduceSum(self): + self._test_all_reduce_sum(self._get_distribution_strategy()) + + def testAllReduceSumGradients(self): + self._test_all_reduce_sum_gradients(self._get_distribution_strategy()) + + def testAllReduceSumGradientTape(self): + self._test_all_reduce_sum_gradient_tape(self._get_distribution_strategy()) + + def testAllReduceMean(self): + self._test_all_reduce_mean(self._get_distribution_strategy()) + + def testAllReduceMeanGradients(self): + self._test_all_reduce_mean_gradients(self._get_distribution_strategy()) + + def testAllReduceMeanGradientTape(self): + self._test_all_reduce_mean_gradient_tape(self._get_distribution_strategy()) + if __name__ == "__main__": test.main() diff --git a/tensorflow/contrib/distribute/python/optimizer_v2_test.py b/tensorflow/contrib/distribute/python/optimizer_v2_test.py index 3064433129865703f574ba79e843880fc4390e74..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_tower( - model_fn, iterator.get_next(), run_concurrently=layer.built))) + 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 03b8564996611d67723abd373133d78391970365..e42bc50fdc4e5e93c998708b0790fdea7768faf2 100644 --- a/tensorflow/contrib/distribute/python/parameter_server_strategy.py +++ b/tensorflow/contrib/distribute/python/parameter_server_strategy.py @@ -18,32 +18,24 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -from tensorflow.contrib.distribute.python import cross_tower_ops as cross_tower_ops_lib -from tensorflow.contrib.distribute.python import mirrored_strategy -from tensorflow.contrib.distribute.python import values -from tensorflow.python.distribute import multi_worker_util -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.training import device_util -from tensorflow.python.training import distribute as distribute_lib -from tensorflow.python.util import nest +from tensorflow.python.distribute import distribute_lib +from tensorflow.python.distribute import input_lib +from tensorflow.python.distribute import parameter_server_strategy +from tensorflow.python.distribute.cluster_resolver import SimpleClusterResolver +from tensorflow.python.distribute.cluster_resolver import TFConfigClusterResolver -_LOCAL_CPU = "/device:CPU:0" -_LOCAL_GPU_0 = "/device:GPU:0" +# pylint: disable=protected-access,invalid-name,line-too-long +CoreParameterServerStrategy = parameter_server_strategy.ParameterServerStrategy +CoreParameterServerExtended = parameter_server_strategy.ParameterServerStrategyExtended + +# 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 @@ -61,16 +53,16 @@ class ParameterServerStrategy(distribute_lib.DistributionStrategy): for a particular worker. Note that each graph and worker is independent. This means that while each worker will synchronously compute a single gradient update across all GPUs, updates between workers proceed asynchronously. - Operations that occur only on the first tower (such as incrementing the global - step), will occur on the first tower *of every worker*. + Operations that occur only on the first replica (such as incrementing the + global step), will occur on the first replica *of every worker*. - It is expected to call `call_for_each_tower(fn, *args, **kwargs)` for any - operations which potentially can be replicated across towers (i.e. multiple + It is expected to call `call_for_each_replica(fn, ...)` for any + operations which potentially can be replicated across replicas (i.e. multiple GPUs) even if there is only CPU or one GPU. When defining the `fn`, extra caution needs to be taken: 1) Always use `tf.get_variable` instead of `tf.Variable` which is not able - to refer to the same variable on different towers. + to refer to the same variable on different replicas. 2) It is generally not recommended to open a device scope under the strategy's scope. A device scope (i.e. calling `tf.device`) will be merged with or @@ -78,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): @@ -94,381 +86,87 @@ class ParameterServerStrategy(distribute_lib.DistributionStrategy): ValueError: if `cluster_spec` is given but `task_type` or `task_id` is not. """ - super(ParameterServerStrategy, self).__init__() - self._num_gpus_per_worker = num_gpus_per_worker - self._initialize_local(num_gpus_per_worker) + super(ParameterServerStrategy, self).__init__( + ParameterServerExtended(self, num_gpus_per_worker)) - # We typically don't need to do all-reduce in this strategy. - self._cross_tower_ops = ( - cross_tower_ops_lib.ReductionToOneDeviceCrossTowerOps( - reduce_to_device=_LOCAL_CPU)) + # 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`. - def _initialize_multi_worker(self, num_gpus_per_worker, cluster_spec, - task_type, task_id): - """Initialize devices for multiple workers. + NOTE: The batch size of the `dataset` argument is treated differently for + this contrib version of `ParameterServerStrategy`. - It creates variable devices and compute devices. Variables and operations - will be assigned to them respectively. We have one compute device per tower. - The variable device is a device function or device string. The default - variable device assigns variables to parameter servers in a round-robin - fashion. + 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: - 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) - - self._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 - # tower. When there are GPUs, replicate operations on these GPUs. Otherwise, - # place operations on CPU. - if num_gpus_per_worker > 0: - self._compute_devices = [ - "%s/device:GPU:%d" % (self._worker_device, i) - for i in range(num_gpus_per_worker) - ] - else: - self._compute_devices = [self._worker_device] - - self._compute_devices = list( - map(device_util.resolve, self._compute_devices)) - self._canonical_compute_device_set = set(self._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=self._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 = 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 = self._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, compute_devices = %r, " - "variable_device = %r", cluster_spec.as_dict(), task_type, task_id, - num_ps_replicas, self._is_chief, self._compute_devices, - self._variable_device) - - def _initialize_local(self, num_gpus_per_worker): - """Initialize internal devices for local training.""" - # Define compute devices which is a list of device strings and one for each - # tower. When there are GPUs, replicate operations on these GPUs. Otherwise, - # place operations on CPU. - if num_gpus_per_worker > 0: - self._compute_devices = list( - map("/device:GPU:{}".format, range(num_gpus_per_worker))) - else: - self._compute_devices = [_LOCAL_CPU] - - self._compute_devices = list( - map(device_util.resolve, self._compute_devices)) - self._canonical_compute_device_set = set(self._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(list(self._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", self._compute_devices, self._variable_device) - - def distribute_dataset(self, dataset_fn): - """Distributes the dataset to each local GPU.""" - return values.PerDeviceDataset( - self._call_dataset_fn(dataset_fn), self._compute_devices, True) - - def _broadcast(self, tensor, destinations): - if not cross_tower_ops_lib.check_destinations(destinations): - destinations = self._compute_devices - return self._cross_tower_ops.broadcast(tensor, destinations) - - # 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_towers > 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_TOWER - ): - raise ValueError("Invalid variable aggregation mode: " + aggregation + - " for variable: " + kwargs["name"]) - - def var_creator(*args, **kwargs): - # 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(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_tower(self, fn, *args, **kwargs): - # pylint: disable=protected-access - return mirrored_strategy._call_for_each_tower(self, fn, *args, **kwargs) - - def _verify_destinations_not_different_worker(self, destinations): - if destinations is None: - return - for d in cross_tower_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._worker_device)) - - def _reduce(self, aggregation, value, destinations): - self._verify_destinations_not_different_worker(destinations) - if not isinstance(value, values.DistributedValues): - # pylint: disable=protected-access - return mirrored_strategy._reduce_non_distributed_value( - self, aggregation, value, destinations) - if aggregation == vs.VariableAggregation.ONLY_FIRST_TOWER: - return self.broadcast(value.get(self._compute_devices[0]), destinations) - return self._cross_tower_ops.reduce( - aggregation, value, destinations=destinations) - - def _batch_reduce(self, aggregation, value_destination_pairs): - if aggregation == vs.VariableAggregation.ONLY_FIRST_TOWER: - return [self.broadcast(v.get(self._compute_devices[0]), d) - for v, d in value_destination_pairs] - for _, destinations in value_destination_pairs: - self._verify_destinations_not_different_worker(destinations) - return self._cross_tower_ops.batch_reduce(aggregation, - value_destination_pairs) - - def _select_single_value(self, structured): - """Select any single values in `structured`.""" + return super(ParameterServerStrategy, self).make_dataset_iterator(dataset) - def _select_fn(x): # pylint: disable=g-missing-docstring - if isinstance(x, values.Mirrored): - if len(x.devices) == 1: - return list(x._index.values())[0] # pylint: disable=protected-access - 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.PerDevice): - raise ValueError( - "You cannot update variable with a PerDevice object %r when using " - "ParameterServerStrategy. You must specify a single value or a " - "Mirrored with a single value" % x) - else: - return x + # 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. - return nest.map_structure(_select_fn, structured) - - def _update(self, var, options, fn, *args, **kwargs): - 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) - should_group = options.pop("grouped") - assert not options # Validate that we are processing all of the options. - with ops.colocate_with(var), distribute_lib.UpdateContext(var.device): - result = fn(var, *self._select_single_value(args), - **self._select_single_value(kwargs)) - if should_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, options, fn, *args, **kwargs): - should_group = options.pop("grouped") - assert not options # Validate that we are processing all of the options. - with ops.device( - colocate_with.device), distribute_lib.UpdateContext(colocate_with): - result = fn(*args, **kwargs) - if should_group: - return result - else: - return nest.map_structure(self._unwrap, result) - - def _unwrap(self, val): - if isinstance(val, values.DistributedValues): - # Return in a deterministic order. - if set(val.devices) == self._canonical_compute_device_set: - return [val.get(device=d) for d in self._compute_devices] - return [val.get(device=d) for d in sorted(val.devices)] - return [val] - - def value_container(self, val): - return values.value_container(val) - - def read_var(self, var): - # No need to distinguish between normal variables and tower-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 not session_config or not self._cluster_spec: - return - - session_config.isolate_session_state = False - - assert self._cluster_spec - 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 - del session_config.device_filters[:] - session_config.device_filters.extend( - ["/job:%s/task:%d" % (self._task_type, self._task_id), "/job:ps"]) - - @property - def num_towers(self): - return len(self._compute_devices) - - @property - def num_replicas_in_sync(self): - return len(self._compute_devices) - - @property - def worker_devices(self): - # Make a copy to prevent users from accidentally mutating our copy. - return list(self._compute_devices) - - @property - def parameter_devices(self): - return list(self._parameter_devices) - - def non_slot_devices(self, var_list): - return min(var_list, key=lambda x: x.name) - - @property - def between_graph(self): - return True - - @property - def should_init(self): - return self._is_chief - - @property - def should_checkpoint(self): - return self._is_chief - + return super(ParameterServerStrategy, + self).experimental_make_numpy_iterator( + numpy_input, batch_size, num_epochs, shuffle, session) + + +class ParameterServerExtended(CoreParameterServerExtended): + """Implementation of ParameterServerStrategy.""" + + 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) + + 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 should_save_summary(self): - return self._is_chief + 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 91f1fae47e96635db43318cc5303da7c2b6e6994..89dcdbcfc2f1f9d8cd46db9ccf133be08ff89533 100644 --- a/tensorflow/contrib/distribute/python/parameter_server_strategy_test.py +++ b/tensorflow/contrib/distribute/python/parameter_server_strategy_test.py @@ -25,29 +25,99 @@ from absl.testing import parameterized from tensorflow.contrib.distribute.python import combinations from tensorflow.contrib.distribute.python import multi_worker_test_base from tensorflow.contrib.distribute.python import parameter_server_strategy -from tensorflow.contrib.distribute.python import values +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 from tensorflow.python.framework import constant_op +from tensorflow.python.framework import errors from tensorflow.python.framework import ops +from tensorflow.python.framework import tensor_util from tensorflow.python.layers import core from tensorflow.python.ops import array_ops from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import gradients +from tensorflow.python.ops import partitioned_variables +from tensorflow.python.ops import resource_variable_ops from tensorflow.python.ops import variable_scope from tensorflow.python.ops import variables from tensorflow.python.platform import test -from tensorflow.python.training import device_util -from tensorflow.python.training import distribution_strategy_context 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 PS = run_config.TaskType.PS +def _get_replica_id_integer(): + replica_id = ds_context.get_replica_context().replica_id_in_sync_group + if isinstance(replica_id, ops.Tensor): + replica_id = tensor_util.constant_value(replica_id) + 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): @@ -61,31 +131,33 @@ 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.test_session(target=self._default_target, - config=sess_config) as sess, \ + self.cached_session(target=self._default_target, + config=sess_config) as sess, \ d.scope(): - # Define a variable outside the call_for_each_tower scope. This is not - # recommended. + # Define a variable outside the call_for_each_replica scope. n = variable_scope.get_variable('n', initializer=10.0) self.assertEqual(n.device, '/job:ps/task:0') @@ -93,9 +165,8 @@ class ParameterServerStrategyTestBase( if num_gpus == 0: last_part_device = 'device:CPU:0' else: - last_part_device = ( - 'device:GPU:%d' % - distribution_strategy_context.get_tower_context().tower_id) + replica_id = _get_replica_id_integer() + last_part_device = ('device:GPU:%d' % replica_id) a = constant_op.constant(1.0) b = constant_op.constant(2.0) @@ -119,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) @@ -165,7 +236,7 @@ class ParameterServerStrategyTestBase( self.assertIn('/job:ps/', h.device) return y_add, z_add, f - y, z, f = d.call_for_each_tower(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) @@ -177,39 +248,88 @@ class ParameterServerStrategyTestBase( self.assertEqual(z_val, 43.0) self.assertEqual(f_val, 46.0) + def _test_device_assignment_distributed_enable_partitioner( + 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, + config=sess_config) as sess, \ + d.scope(): + + n = variable_scope.get_variable( + 'n', + initializer=constant_op.constant([10.0, 20.0]), + aggregation=variable_scope.VariableAggregation.SUM, + partitioner=partitioner) + + for part_id, var in enumerate(n): + self.assertEqual(var.device, '/job:ps/task:%d' % part_id) + + def model_fn(): + a = constant_op.constant([3.0, 5.0]) + # The device scope is ignored for variables but not for normal ops. + with ops.device('/job:worker/task:0'): + x = variable_scope.get_variable( + 'x', + initializer=constant_op.constant([10.0, 20.0]), + aggregation=variable_scope.VariableAggregation.SUM, + partitioner=partitioner) + x_add = x.assign_add(a, name='x_add') + # The variable x is on the task 1 since the device_function has been + # called once before the model_fn. + for part_id, var in enumerate(x): + self.assertEqual(var.device, '/job:ps/task:%d' % part_id) + self.assertEqual(var.device, x_add[part_id].device) + + return x_add + + x = d.extended.call_for_each_replica(model_fn) + + if context.num_gpus() >= 1: + variables.global_variables_initializer().run() + x_val = sess.run(x) + if num_gpus < 1: + self.assertEqual(x_val, [13.0, 25.0]) + else: + x_expect = [10.0 + 3 * num_gpus, 20.0 + 5 * num_gpus] + self.assertEqual(x_val, x_expect) + def _test_device_assignment_local(self, d, compute_device='CPU', variable_device='CPU', num_gpus=0): with ops.Graph().as_default(), \ - self.test_session(target=self._default_target, - config=self._sess_config) as sess, \ + self.cached_session(target=self._default_target, + config=self._sess_config) as sess, \ d.scope(): def model_fn(): if 'CPU' in compute_device: - tower_compute_device = '/device:CPU:0' + replica_compute_device = '/device:CPU:0' else: - tower_compute_device = ( - '/device:GPU:%d' % - distribution_strategy_context.get_tower_context().tower_id) - tower_compute_device = device_util.canonicalize(tower_compute_device) + replica_id = _get_replica_id_integer() + replica_compute_device = ('/device:GPU:%d' % replica_id) + replica_compute_device = device_util.canonicalize( + replica_compute_device) if 'CPU' in variable_device: - tower_variable_device = '/device:CPU:0' + replica_variable_device = '/device:CPU:0' else: - tower_variable_device = ( - '/device:GPU:%d' % - distribution_strategy_context.get_tower_context().tower_id) - tower_variable_device = device_util.canonicalize(tower_variable_device) + replica_id = _get_replica_id_integer() + replica_variable_device = ('/device:GPU:%d' % replica_id) + replica_variable_device = device_util.canonicalize( + replica_variable_device) a = constant_op.constant(1.0) b = constant_op.constant(2.0) c = a + b - self.assertEqual(a.device, tower_compute_device) - self.assertEqual(b.device, tower_compute_device) - self.assertEqual(c.device, tower_compute_device) + self.assertEqual(a.device, replica_compute_device) + self.assertEqual(b.device, replica_compute_device) + self.assertEqual(c.device, replica_compute_device) # The device scope is ignored for variables but not for normal ops. with ops.device('/device:GPU:2'): @@ -219,12 +339,12 @@ class ParameterServerStrategyTestBase( x_add = x.assign_add(c) e = a + c self.assertEqual( - device_util.canonicalize(x.device), tower_variable_device) + device_util.canonicalize(x.device), replica_variable_device) self.assertEqual(x_add.device, x.device) 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) @@ -232,7 +352,7 @@ class ParameterServerStrategyTestBase( # non-distributed values. y_add = y.assign_add(array_ops.identity(x_add)) self.assertEqual( - device_util.canonicalize(y.device), tower_variable_device) + device_util.canonicalize(y.device), replica_variable_device) self.assertEqual(y_add.device, y.device) self.assertEqual(y.device, x.device) @@ -240,7 +360,7 @@ class ParameterServerStrategyTestBase( 'z', initializer=10.0, aggregation=variable_scope.VariableAggregation.SUM) self.assertEqual( - device_util.canonicalize(z.device), tower_variable_device) + device_util.canonicalize(z.device), replica_variable_device) with ops.control_dependencies([y_add]): # We add an identity here to avoid complaints about summing @@ -248,7 +368,7 @@ class ParameterServerStrategyTestBase( z_add = z.assign_add(array_ops.identity(y)) with ops.control_dependencies([z_add]): f = z + c - self.assertEqual(f.device, tower_compute_device) + self.assertEqual(f.device, replica_compute_device) # The device scope would merge with the default worker device. with ops.device('/CPU:1'): @@ -261,13 +381,13 @@ class ParameterServerStrategyTestBase( u = variable_scope.get_variable('u', initializer=30.0) h = f + 1.0 self.assertEqual( - device_util.canonicalize(u.device), tower_variable_device) + device_util.canonicalize(u.device), replica_variable_device) self.assertEqual( device_util.canonicalize(x.device), device_util.canonicalize(h.device)) return y_add, z_add, f - y, z, f = d.call_for_each_tower(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) @@ -279,18 +399,22 @@ 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) - if hasattr(d, '_cluster_spec') and d._cluster_spec: - num_workers = len(d._cluster_spec.as_dict().get(WORKER)) - if 'chief' in d._cluster_spec.as_dict(): + 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(): num_workers += 1 else: num_workers = 1 with ops.Graph().as_default(), \ - self.test_session(target=master_target, - config=sess_config) as sess, \ + self.cached_session(target=master_target, + config=sess_config) as sess, \ d.scope(): def model_fn(): @@ -302,7 +426,7 @@ class ParameterServerStrategyTestBase( aggregation=variable_scope.VariableAggregation.SUM) z = variable_scope.get_variable( 'z', initializer=30.0, - aggregation=variable_scope.VariableAggregation.ONLY_FIRST_TOWER) + aggregation=variable_scope.VariableAggregation.ONLY_FIRST_REPLICA) # We explicitly make a constant tensor here to avoid complaints about # summing non-distributed values. @@ -314,10 +438,10 @@ 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_tower(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._num_gpus_per_worker: + if context.num_gpus() < d.extended._num_gpus_per_worker: return True if task_id == 0: @@ -342,24 +466,33 @@ class ParameterServerStrategyTestBase( self._finish_condition.release() x_val, y_val, z_val = sess.run([x, y, z]) - self.assertEqual(x_val, 10.0 + 1.0 * num_workers * d.num_towers) - self.assertEqual(y_val, 20.0 + 1.0 * num_workers * d.num_towers) + self.assertEqual(x_val, 10.0 + 1.0 * num_workers * d.num_replicas_in_sync) + self.assertEqual(y_val, 20.0 + 1.0 * num_workers * d.num_replicas_in_sync) self.assertEqual(z_val, 30.0 + 1.0 * num_workers) - return (x_val == 10.0 + 1.0 * num_workers * d.num_towers and - y_val == 20.0 + 1.0 * num_workers * d.num_towers and + return (x_val == 10.0 + 1.0 * num_workers * d.num_replicas_in_sync and + 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) - assert hasattr(d, '_cluster_spec') and d._cluster_spec - num_workers = len(d._cluster_spec.as_dict().get(WORKER)) - if CHIEF in d._cluster_spec.as_dict(): - num_workers += 1 + 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 + num_workers = len(d.extended._cluster_spec.as_dict().get(WORKER)) + if CHIEF in d.extended._cluster_spec.as_dict(): + num_workers += 1 + else: + # local + num_workers = 1 with ops.Graph().as_default(), \ - self.test_session(target=master_target, - config=sess_config) as sess, \ + self.cached_session(target=master_target, + config=sess_config) as sess, \ d.scope(): l = core.Dense(1, use_bias=False) @@ -386,28 +519,30 @@ class ParameterServerStrategyTestBase( def step(): """Perform one optimization step.""" # Run forward & backward to get gradients, variables list. - g_v = d.call_for_each_tower(grad_fn, 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 = [] for g, v in g_v: - fetched = d.read_var(v) + fetched = d.extended.read_var(v) before_list.append(fetched) with ops.control_dependencies([fetched]): # TODO(yuefengz): support non-Mirrored variable as destinations. - g = d.reduce( - variable_scope.VariableAggregation.SUM, g, destinations=v) + g = d.extended.reduce_to( + reduce_util.ReduceOp.SUM, g, destinations=v) with ops.control_dependencies( - d.update(v, update, g, grouped=False)): - after_list.append(d.read_var(v)) + d.extended.update(v, update, args=(g,), group=False)): + after_list.append(d.extended.read_var(v)) return before_list, after_list before_out, after_out = step() - if context.num_gpus() < d._num_gpus_per_worker: + if context.num_gpus() < d.extended._num_gpus_per_worker: return True - if multi_worker_util.is_chief(d._cluster_spec, task_type, task_id): + if (not task_type or + multi_worker_util.is_chief( + d.extended._cluster_spec, task_type, task_id)): variables.global_variables_initializer().run() # Workers waiting for chief worker's initializing variables. @@ -430,9 +565,49 @@ 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, + use_core_strategy=False): + distribution, master_target, config = self._get_test_objects( + task_type, task_id, num_gpus, use_core_strategy=use_core_strategy) + devices = distribution.extended.worker_devices -class ParameterServerStrategyTest(ParameterServerStrategyTestBase, - parameterized.TestCase): + with ops.Graph().as_default(), \ + self.cached_session(config=config, + target=master_target) as sess: + iterator = distribution.make_input_fn_iterator(input_fn) + 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) + + with self.assertRaises(errors.OutOfRangeError): + next_element = iterator.get_next() + sess.run([values.select_replica(r, next_element) + for r in range(len(devices))]) + + # After re-initializing the iterator, should be able to iterate again. + 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) + + +class ParameterServerStrategyTest( + ParameterServerStrategyTestBase, + strategy_test_lib.DistributionTestBase, + strategy_test_lib.TwoDeviceDistributionTestBase, + parameterized.TestCase): @classmethod def setUpClass(cls): @@ -440,50 +615,204 @@ 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) - def testSimpleBetweenGraph(self): - self._run_between_graph_clients(self._test_simple_increment, - self._cluster_spec, context.num_gpus()) + @combinations.generate( + 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, use_core_strategy=use_core_strategy) + + @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 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 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], 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, + use_core_strategy=[True, False])) + def testMakeInputFnIteratorDistributed(self, num_gpus, use_core_strategy): + if context.num_gpus() < num_gpus: + self.skipTest('Not enough GPUs') + dataset_fn = lambda: dataset_ops.Dataset.range(100) + 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, + 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, + use_core_strategy=use_core_strategy) + + @combinations.generate( + combinations.combine( + mode=['graph'], + num_gpus=[1, 2], + required_gpus=1, + use_core_strategy=[True, False])) + def testMakeInputFnIteratorLocal(self, num_gpus, use_core_strategy): + if context.num_gpus() < num_gpus: + self.skipTest('Not enough GPUs') + dataset_fn = lambda: dataset_ops.Dataset.range(100) + 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, + 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, + use_core_strategy=use_core_strategy) + + @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) + + @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 = strategy.update_config_proto(config_proto) + + # Verify device filters. + self.assertEqual(['/job:worker/task:1', '/job:ps'], + new_config.device_filters) + + # Verify isolate_session_state + self.assertFalse(new_config.isolate_session_state) + + @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 = 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, @@ -495,20 +824,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 testGlobalStepIsWrapped(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, @@ -517,6 +857,55 @@ 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(strategy, created_step.distribute_strategy) + + @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, + msg=('created_step %s type %s vs. get_step %s type %s' % + (id(created_step), created_step.__class__.__name__, + 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)) + # 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 f(): + with backprop.GradientTape() as tape: + v = variable_scope.get_variable('v', initializer=10.0) + _ = v * v + v, = tape.watched_variables() + w = strategy.extended.value_container(v) + self.assertIs(values.AggregatingVariable, type(w)) + + 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/single_loss_example.py b/tensorflow/contrib/distribute/python/single_loss_example.py index 09b351ffa4165656e2fc9666ab4b7725ef061f50..be724fb59a7efa18c43c4cb98649ced806f7bcb4 100644 --- a/tensorflow/contrib/distribute/python/single_loss_example.py +++ b/tensorflow/contrib/distribute/python/single_loss_example.py @@ -90,7 +90,7 @@ def batchnorm_example(optimizer_fn, batch_per_epoch=1, momentum=0.9, renorm=False, - update_ops_in_tower_mode=False): + update_ops_in_replica_mode=False): """Example of non-distribution-aware legacy code with batch normalization.""" def dataset_fn(): @@ -113,7 +113,7 @@ def batchnorm_example(optimizer_fn, y = batchnorm(x, training=True) with ops.control_dependencies( ops.get_collection(ops.GraphKeys.UPDATE_OPS) - if update_ops_in_tower_mode else []): + if update_ops_in_replica_mode else []): loss = math_ops.reduce_mean( math_ops.reduce_sum(layer(y)) - constant_op.constant(1.)) # `x` and `y` will be fetched by the gradient computation, but not `loss`. diff --git a/tensorflow/contrib/distribute/python/step_fn.py b/tensorflow/contrib/distribute/python/step_fn.py index 23bf36184fa61105b43c71ada6c343b30dce8376..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): @@ -90,29 +90,25 @@ class StandardSingleLossStep(StandardInputStep): super(StandardSingleLossStep, self).__init__(dataset_fn, distribution) self._loss_fn = loss_fn self._optimizer = optimizer - self._is_run_concurrently = False self._iterations_per_step = iterations_per_step def __call__(self): with self._distribution.scope(): - def step_fn(ctx, *inputs): + def step_fn(ctx, inputs): """Function to run one iteration with one input.""" 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_tower( - gradients_fn, - ctx, *inputs, - run_concurrently=self._is_run_concurrently) + 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. # Otherwise, multiple sets of mirrored variables are going to be # created. - self._is_run_concurrently = True return self._optimizer._distributed_apply( # pylint: disable=protected-access 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 fd280f5754b34170cdd6b948236138d0e77dd8bc..2e2ee92b6e20471f367895ea53c0864bb3d1dae7 100644 --- a/tensorflow/contrib/distribute/python/strategy_test_lib.py +++ b/tensorflow/contrib/distribute/python/strategy_test_lib.py @@ -18,17 +18,26 @@ 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 from tensorflow.python.eager import backprop from tensorflow.python.eager import context from tensorflow.python.eager import 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.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 -from tensorflow.python.training import distribution_strategy_context from tensorflow.python.training import optimizer @@ -36,45 +45,45 @@ class _TestException(Exception): pass -# May be the argument to either distribution.call_for_each_tower() or -# get_tower_context().merge_call() +# 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_tower() call, calls a -# get_tower_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(): - distribution_strategy_context.get_tower_context().merge_call( - _raise_exception_fn) + ds_context.get_replica_context().merge_call(_raise_exception_fn) -# Must be the argument to a get_tower_context().merge_call() call, calls -# dist.call_for_each_tower() with a function that raises an exception. +# Must be the argument to a get_replica_context().merge_call() call, calls +# dist.extended.call_for_each_replica() with a function that raises an +# exception. def _call_raises_fn(dist): - dist.call_for_each_tower(_raise_exception_fn) + dist.extended.call_for_each_replica(_raise_exception_fn) -# Must be the argument to a distribution.call_for_each_tower() call, -# calls a get_tower_context().merge_call() that calls a -# call_for_each_tower() 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 raises an exception. def _merge_call_raises_fn(): - distribution_strategy_context.get_tower_context().merge_call(_call_raises_fn) + ds_context.get_replica_context().merge_call(_call_raises_fn) -# Must be the argument to a get_tower_context().merge_call() call, calls -# dist.call_for_each_tower() with a function that calls a -# get_tower_context().merge_call() that raises an exception. +# Must be the argument to a get_replica_context().merge_call() call, calls +# 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_tower(_merge_raises_fn) + dist.extended.call_for_each_replica(_merge_raises_fn) -# Must be the argument to a distribution.call_for_each_tower() call, calls a -# get_tower_context().merge_call() that calls a call_for_each_tower() that -# calls a get_tower_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(): - distribution_strategy_context.get_tower_context().merge_call( - _call_merge_raises_fn) + ds_context.get_replica_context().merge_call(_call_merge_raises_fn) class DistributionTestBase(test.TestCase): @@ -103,21 +112,21 @@ class DistributionTestBase(test.TestCase): def step(): """Perform one optimization step.""" # Run forward & backward to get gradients, variables list. - g_v = d.call_for_each_tower(grad_fn, one, run_concurrently=l.built) + 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 = [] for g, v in g_v: - fetched = d.read_var(v) + fetched = d.extended.read_var(v) before_list.append(fetched) # control_dependencies irrelevant but harmless in eager execution with ops.control_dependencies([fetched]): - g = d.reduce( - variable_scope.VariableAggregation.SUM, g, destinations=v) - with ops.control_dependencies(d.update( - v, update, g, grouped=False)): - after_list.append(d.read_var(v)) + g = d.extended.reduce_to( + reduce_util.ReduceOp.SUM, g, destinations=v) + 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 for i in range(10): @@ -138,7 +147,7 @@ class DistributionTestBase(test.TestCase): config.gpu_options.per_process_gpu_memory_fraction = 0.3 with context.graph_mode(), \ ops.Graph().as_default(), \ - self.test_session(config=config) as sess, \ + self.cached_session(config=config) as sess, \ d.scope(): l = core.Dense(1, use_bias=False) @@ -159,20 +168,20 @@ class DistributionTestBase(test.TestCase): def step(): """Perform one optimization step.""" # Run forward & backward to get gradients, variables list. - g_v = d.call_for_each_tower(grad_fn, 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 = [] for g, v in g_v: - fetched = d.read_var(v) + fetched = d.extended.read_var(v) before_list.append(fetched) with ops.control_dependencies([fetched]): - g = d.reduce( - variable_scope.VariableAggregation.SUM, g, destinations=v) - with ops.control_dependencies(d.update( - v, update, g, grouped=False)): - after_list.append(d.read_var(v)) + g = d.extended.reduce_to( + reduce_util.ReduceOp.SUM, g, destinations=v) + 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 before_out, after_out = step() @@ -188,47 +197,291 @@ class DistributionTestBase(test.TestCase): # Error should go down self.assertLess(error_after, error_before) - def _test_map_reduce(self, d, in_graph=None): - with d.scope(): - map_in = [constant_op.constant(i) for i in range(10)] - map_out = d.map(map_in, lambda x, y: x * y, 2) - observed = d.reduce(variable_scope.VariableAggregation.SUM, map_out, - "/device:CPU:0") - expected = 90 # 2 * (0 + 1 + ... + 9) - self.assertEqual(expected, observed.numpy()) - - def _test_device_index(self, d): - with d.scope(): - expected_devices = [False] * len(d.worker_devices) - - def mark_devices_fn(device_id): - self.assertLess(device_id, len(d.worker_devices)) - self.assertFalse(expected_devices[device_id]) - expected_devices[device_id] = True - - d.call_for_each_tower(mark_devices_fn, d.worker_device_index) - self.assertAllEqual(expected_devices, [True] * len(d.worker_devices)) - - def _test_tower_id(self, d): + def _test_replica_id(self, d): with d.scope(): - expected_devices = [False] * len(d.worker_devices) + expected_devices = [False] * len(d.extended.worker_devices) def mark_devices_fn(): - tower_id = distribution_strategy_context.get_tower_context().tower_id - self.assertLess(tower_id, len(d.worker_devices)) - self.assertFalse(expected_devices[tower_id]) - expected_devices[tower_id] = True + replica_id = self.evaluate( + ds_context.get_replica_context().replica_id_in_sync_group) + self.assertLess(replica_id, len(d.extended.worker_devices)) + self.assertFalse(expected_devices[replica_id]) + expected_devices[replica_id] = True - d.call_for_each_tower(mark_devices_fn) - self.assertAllEqual(expected_devices, [True] * len(d.worker_devices)) + 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_tower(_raise_exception_fn) + dist.extended.call_for_each_replica(_raise_exception_fn) with self.assertRaises(_TestException): - dist.call_for_each_tower(_merge_raises_fn) + dist.extended.call_for_each_replica(_merge_raises_fn) with self.assertRaises(_TestException): - dist.call_for_each_tower(_merge_call_raises_fn) + dist.extended.call_for_each_replica(_merge_call_raises_fn) with self.assertRaises(_TestException): - dist.call_for_each_tower(_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, + expected_num_replicas_in_sync, + expected_num_input_pipelines, + expected_input_pipeline_id): + # Use a list of one element as counter so that it can be captured by the + # `_input_fn`. This counter is incremented by 1 each time an input_fn is + # called. We use this counter to check whether the `input_pipeline_id` + # matches the counter in the in-graph replication. + worker_id_counter = [0] + + def _input_fn(input_context): + """Input fn for testing.""" + self.assertIsNotNone(input_context) + self.assertEqual(expected_num_replicas_in_sync, + input_context.num_replicas_in_sync) + self.assertEqual(expected_num_input_pipelines, + input_context.num_input_pipelines) + if expected_input_pipeline_id is not None: + self.assertEqual(expected_input_pipeline_id, + input_context.input_pipeline_id) + else: + self.assertEqual(worker_id_counter[0], input_context.input_pipeline_id) + worker_id_counter[0] += 1 + + return dataset_fn() + + return _input_fn + + def _test_input_fn_iterator(self, iterator, devices, expected_values, + sess=None): + evaluate = lambda x: sess.run(x) if sess else self.evaluate(x) + 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) + + with self.assertRaises(errors.OutOfRangeError): + next_element = iterator.get_next() + evaluate( + [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()) + + 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(): + global_step = variable_scope.get_variable( + "global_step", + shape=[], + dtype=dtypes.int64, + initializer=init_ops.zeros_initializer(), + trainable=False, + aggregation=variables.VariableAggregation.ONLY_FIRST_REPLICA) + self.evaluate(variables.global_variables_initializer()) + + def model_fn(): + train_op = global_step.assign_add(1) + value = global_step.read_value() + return train_op, value + + 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 e4bd7428c2ca58ea6225faad2059899e9c1f2eb5..4387210062e42bb1ab7e2351008a45979224ff1a 100644 --- a/tensorflow/contrib/distribute/python/tpu_strategy.py +++ b/tensorflow/contrib/distribute/python/tpu_strategy.py @@ -21,28 +21,61 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -from tensorflow.contrib.distribute.python import cross_tower_ops as cross_tower_ops_lib -from tensorflow.contrib.distribute.python import one_device_strategy -from tensorflow.contrib.distribute.python import values +import collections +import copy + from tensorflow.contrib.tpu.python.ops import tpu_ops +from tensorflow.contrib.tpu.python.tpu import device_assignment as device_assignment_lib +from tensorflow.contrib.tpu.python.tpu import topology from tensorflow.contrib.tpu.python.tpu import tpu from tensorflow.contrib.tpu.python.tpu import tpu_system_metadata as tpu_system_metadata_lib from tensorflow.contrib.tpu.python.tpu import training_loop +from tensorflow.core.protobuf import config_pb2 +from tensorflow.python.client import session as session_lib +from tensorflow.python.distribute import cross_device_ops as cross_device_ops_lib +from tensorflow.python.distribute import device_util +from tensorflow.python.distribute import distribute_lib +from tensorflow.python.distribute import input_lib +from tensorflow.python.distribute import numpy_dataset +from tensorflow.python.distribute import reduce_util +from tensorflow.python.distribute import values +from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver as resolver_lib from tensorflow.python.eager import context from tensorflow.python.eager import tape from tensorflow.python.framework import constant_op +from tensorflow.python.framework import device as tf_device +from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops +from tensorflow.python.framework import tensor_util from tensorflow.python.ops import array_ops from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import variable_scope as vs -from tensorflow.python.ops import variables as variables_lib -from tensorflow.python.training import device_util -from tensorflow.python.training import distribute as distribute_lib +from tensorflow.python.platform import tf_logging as logging 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("") + 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: + 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): @@ -62,26 +95,27 @@ def get_tpu_system_metadata(tpu_cluster_resolver): # TODO(jhseu): Deduplicate with MirroredStrategy? -def _create_tpu_mirrored_variable(devices, real_mirrored_creator, *args, - **kwargs): # pylint: disable=g-missing-docstring +def _create_tpu_mirrored_variable( # pylint: disable=missing-docstring + strategy, device_map, logical_device, real_mirrored_creator, + *args, **kwargs): # Figure out what collections this variable should be added to. # We'll add the TPUMirroredVariable to those collections instead. - 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 # synchronization settings? # Get aggregation value - # TODO(jhseu): Support aggregation in a tower context. + # TODO(jhseu): Support aggregation in a replica context. aggregation = kwargs.pop("aggregation", vs.VariableAggregation.NONE) if aggregation not in [ vs.VariableAggregation.NONE, vs.VariableAggregation.SUM, vs.VariableAggregation.MEAN, - vs.VariableAggregation.ONLY_FIRST_TOWER, + vs.VariableAggregation.ONLY_FIRST_REPLICA, ]: raise ValueError("Invalid variable aggregation mode: {} for variable: {}" .format(aggregation, kwargs["name"])) @@ -93,8 +127,11 @@ def _create_tpu_mirrored_variable(devices, real_mirrored_creator, *args, # was never recorded on the tape instead of having to do this manually # here. with tape.stop_recording(): - index = real_mirrored_creator(devices, *args, **kwargs) - result = values.TPUMirroredVariable(index, index[devices[0]], aggregation) + devices = device_map.logical_to_actual_devices(logical_device) + value_list = real_mirrored_creator(devices, *args, **kwargs) + result = values.TPUMirroredVariable( + strategy, device_map, value_list, aggregation, + logical_device=logical_device) if not context.executing_eagerly(): g = ops.get_default_graph() @@ -104,19 +141,22 @@ def _create_tpu_mirrored_variable(devices, real_mirrored_creator, *args, # "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 index.values(): + for v in value_list: l.remove(v) - g.add_to_collections(collections, result) + g.add_to_collections(var_collections, result) return result -# TODO(jhseu): Stop inheriting from OneDeviceStrategy. -class TPUStrategy(one_device_strategy.OneDeviceStrategy): - """Experimental TPU distribution strategy implementation.""" +class TPUStrategy(distribute_lib.DistributionStrategy): + """TPU distribution strategy implementation.""" - def __init__(self, tpu_cluster_resolver, steps_per_run, num_cores=None): + def __init__(self, + tpu_cluster_resolver=None, + steps_per_run=None, + device_assignment=None, + **kwargs): """Initializes the TPUStrategy object. Args: @@ -127,35 +167,153 @@ class TPUStrategy(one_device_strategy.OneDeviceStrategy): metrics, summaries etc. This parameter is only used when Distribution Strategy is used with estimator or keras. - num_cores: Number of cores to use on the TPU. If None specified, then - auto-detect the cores and topology of the TPU system. + 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. + **kwargs: Additional experimental flags. Will be removed in future. """ - # TODO(sourabhbajaj): OneDeviceStrategy should be initialized with the - # master node fetched from the cluster resolver. - super(TPUStrategy, self).__init__("/device:CPU:0") + if len(kwargs) > 1: + raise ValueError("TPUStrategy constructor only takes one experimental " + "flag now") + elif len(kwargs) == 1 and "_disable_training_loop_on_host" not in kwargs: + raise ValueError("TPUStrategy constructor does not support arguments: " + "{}".format(kwargs)) + + super(TPUStrategy, self).__init__(TPUExtended( + self, tpu_cluster_resolver, steps_per_run, device_assignment, + kwargs.get("_disable_training_loop_on_host", False))) + + @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(): + raise NotImplementedError("Eager mode not supported in TPUStrategy.") + + if self.extended._disable_training_loop_on_host: # pylint: disable=protected-access + raise NotImplementedError( + "`experimental_run` is not compatible with " + "`_disable_training_loop_on_host=True`") + + if input_iterator is None: + inputs = [] + else: + inputs = input_iterator.get_next() + + result = [None] + def replicated_fn(replica_id, inputs): + """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(inputs) + return result[0] + + replicate_inputs = [] # By replica. + for i in range(self.num_replicas_in_sync): + replicate_inputs.append( + [constant_op.constant(i, dtype=dtypes.int32), + values.select_replica(i, inputs)]) + + with self.scope(): + replicate_outputs = tpu.replicate(replicated_fn, replicate_inputs) + + # Workaround for `tpu.replicate` behaviour when single `Tensor` returned. + replicate_outputs = [ + nest.pack_sequence_as(result[0], nest.flatten(replica_outputs)) + for replica_outputs in replicate_outputs] + + device_map = self.extended._device_map # pylint: disable=protected-access + return values.regroup(device_map, replicate_outputs) + + +class TPUExtended(distribute_lib.DistributionStrategyExtended): + """Implementation of TPUStrategy.""" + + def __init__(self, + container_strategy, + tpu_cluster_resolver=None, + steps_per_run=None, + device_assignment=None, + disable_training_loop_on_host=False): + super(TPUExtended, self).__init__(container_strategy) + + if tpu_cluster_resolver is None: + tpu_cluster_resolver = resolver_lib.TPUClusterResolver("") + + if steps_per_run is None: + # TODO(frankchn): Warn when we are being used by DS/Keras and this is + # not specified. + steps_per_run = 1 self._tpu_cluster_resolver = tpu_cluster_resolver self._tpu_metadata = get_tpu_system_metadata(self._tpu_cluster_resolver) - # TODO(sourabhbajaj): Change this from num_cores to metadata_override - self._num_cores_override = num_cores + self._device_assignment = device_assignment + self._disable_training_loop_on_host = disable_training_loop_on_host + + # 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. - device_map = {d.name: i for i, d in enumerate(self._tpu_metadata.devices) - if "device:TPU:" in d.name} - self._device_index = values.PerDevice(device_map) - self._tpu_devices = sorted(device_map.keys()) - # Only create variables for the number of towers we're running. - self._tpu_devices = self._tpu_devices[:self.num_towers] + self._device_index = { + d.name: i for i, d in enumerate(self._tpu_metadata.devices) + if "device:TPU:" in d.name + } + self._host_device = self.get_host_cpu_device(0) + self._tpu_devices = tuple(sorted(self._device_index.keys())) + # Only create variables for the number of replicas we're running. + self._tpu_devices = self._tpu_devices[:self._num_replicas_in_sync] + self._device_map = values.ReplicaDeviceMap(self._tpu_devices) + + # If the training loop is on the device, we must use the infeed, with input + # on the host. Otherwise, we preload the data onto the TPUs. + if disable_training_loop_on_host: + 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) + ] + self._input_workers = input_lib.InputWorkers( + input_device_map, worker_devices) + else: + input_worker_devices = collections.OrderedDict() + for tpu_device in self._tpu_devices: + host_device = _get_host_for_device(tpu_device) + input_worker_devices.setdefault(host_device, []) + input_worker_devices[host_device].append(tpu_device) + self._input_workers = input_lib.InputWorkers( + self._device_map, tuple(input_worker_devices.items())) # TODO(sourabhbajaj): Remove this once performance of running one step # at a time is comparable to multiple steps. self.steps_per_run = steps_per_run - self._require_static_shapes = True - def _get_enqueue_op_per_host(self, host_id, iterator, input_shapes, - iterations): + def _validate_colocate_with_variable(self, colocate_with_variable): + values.validate_colocate_tpu_variable(colocate_with_variable, self) + + def _get_enqueue_op_per_host(self, host_id, multi_worker_iterator, + input_shapes, iterations): """Create an enqueue op for a single host identified using host_id. The while_loop op returned will run `iterations` times and in each run @@ -163,7 +321,7 @@ class TPUStrategy(one_device_strategy.OneDeviceStrategy): Args: host_id: integer, id of the host to run the enqueue ops on. - iterator: `tf.data` iterator to read the input data. + multi_worker_iterator: MultiWorkerDataIterator to read the input data. input_shapes: shape of inputs to be enqueue on the queue. This is same as the value of `nest.flatten(iterator.output_shapes)`. iterations: integer, number of iterations to be run; determines the @@ -174,6 +332,10 @@ class TPUStrategy(one_device_strategy.OneDeviceStrategy): on the infeed queue from the host with id `host_id` for each device shard. """ host = self.get_host_cpu_device(host_id) + # TODO(sourabhbajaj): Possibly make changes to MultiWorkerDataset + # to work with TPU Prefetch so clean up this code. + iterator = ( + multi_worker_iterator.get_iterator(self.get_host(host_id))._iterator) # pylint: disable=protected-access def _infeed_enqueue_ops_fn(): """Enqueue ops for one iteration.""" @@ -182,7 +344,7 @@ class TPUStrategy(one_device_strategy.OneDeviceStrategy): enqueue_ops = [] with ops.device(host): - for _ in range(self.num_towers_per_host): + for _ in range(self.num_replicas_per_host): # Use control dependencies to ensure a deterministic ordering. with ops.control_dependencies(control_deps): inputs = nest.flatten(iterator.get_next()) @@ -211,44 +373,170 @@ class TPUStrategy(one_device_strategy.OneDeviceStrategy): return enqueue_op_per_host - def distribute_dataset(self, dataset_fn): - # TODO(priyag): Perhaps distribute across cores here. - return self._call_dataset_fn(dataset_fn) + def _make_dataset_iterator(self, dataset): + """Make iterators for each of the TPU hosts.""" + return input_lib.DatasetIterator(dataset, self._input_workers, + self._num_replicas_in_sync) + + def _make_input_fn_iterator( + self, + input_fn, + replication_mode=distribute_lib.InputReplicationMode.PER_WORKER): + input_contexts = [] + num_workers = self._input_workers.num_workers + for i in range(num_workers): + input_contexts.append(distribute_lib.InputContext( + num_input_pipelines=num_workers, + input_pipeline_id=i, + num_replicas_in_sync=self._num_replicas_in_sync)) + return input_lib.InputFunctionIterator( + input_fn, self._input_workers, input_contexts) + + def _experimental_make_numpy_dataset(self, numpy_input, session): + return numpy_dataset.one_host_numpy_dataset( + numpy_input, numpy_dataset.SingleDevice(self.get_host_cpu_device(0)), + session) # TODO(priyag): Deal with OutOfRange errors once b/111349762 is fixed. # TODO(sourabhbajaj): Remove the initial_loop_values parameter when we have # a mechanism to infer the outputs of `fn`. Pending b/110550782. - def _run_steps_on_dataset(self, fn, iterator, iterations, - initial_loop_values=None): + def _experimental_run_steps_on_iterator( + self, fn, multi_worker_iterator, iterations, initial_loop_values=None): + if self._disable_training_loop_on_host: + impl = self._run_steps_on_iterator_with_device_loop + else: + impl = self._run_steps_on_iterator_with_host_loop + + return impl( + fn=fn, multi_worker_iterator=multi_worker_iterator, + iterations=iterations, initial_loop_values=initial_loop_values) + + def _run_steps_on_iterator_with_host_loop( + self, fn, multi_worker_iterator, iterations, initial_loop_values=None): + output_shapes = multi_worker_iterator.output_shapes + shapes = nest.flatten(output_shapes) + if any(not s.is_fully_defined() for s in shapes): + raise ValueError( + "TPU currently requires fully defined shapes. Either use " + "set_shape() on the input tensors or use " + "dataset.batch(..., drop_remainder=True).") + + # Wrap `fn` for repeat. + if initial_loop_values is None: + initial_loop_values = {} + initial_loop_values = nest.flatten(initial_loop_values) + ctx = input_lib.MultiStepContext() + + def run_fn(inputs): + """Single step on the TPU device.""" + fn_result = fn(ctx, inputs) + flat_last_step_outputs = nest.flatten(ctx.last_step_outputs) + if flat_last_step_outputs: + with ops.control_dependencies([fn_result]): + return [array_ops.identity(f) for f in flat_last_step_outputs] + else: + return fn_result + + # We capture the control_flow_context at this point, before we run `fn` + # inside a while_loop and TPU replicate context. This is useful in cases + # where we might need to exit these contexts and get back to the outer + # context to do some things, for e.g. create an op which should be + # evaluated only once at the end of the loop on the host. One such usage + # is in creating metrics' value op. + self._outer_control_flow_context = ( + ops.get_default_graph()._get_control_flow_context()) # pylint: disable=protected-access + + def rewrite_fn(*args): + """The rewritten step fn running on TPU.""" + del args + + per_replica_inputs = multi_worker_iterator.get_next() + replicate_inputs = [] + for replica_id in range(self._num_replicas_in_sync): + select_replica = lambda x: values.select_replica(replica_id, x) # pylint: disable=cell-var-from-loop + replicate_inputs.append((nest.map_structure( + select_replica, per_replica_inputs),)) + + replicate_outputs = tpu.replicate(run_fn, replicate_inputs) - shapes = nest.flatten(iterator.output_shapes) - if any([not s.is_fully_defined() for s in shapes]): + # If run_fn has tensor outputs, tpu.replicate returns a list of list. We + # will flatten it in this case. If run_fn has no tensor outputs, + # tpu.replicate returns a list of no_ops, we will keep the output as it + # is. + if isinstance(replicate_outputs[0], list): + replicate_outputs = nest.flatten(replicate_outputs) + + return replicate_outputs + + # TODO(sourabhbajaj): The input to while loop should be based on the + # output type of the step_fn + assert isinstance(initial_loop_values, list) + initial_loop_values = initial_loop_values * self._num_replicas_in_sync + + # Put the while loop op on host 0. + with ops.device(self.get_host_cpu_device(0)): + replicate_outputs = training_loop.repeat(iterations, rewrite_fn, + initial_loop_values) + + del self._outer_control_flow_context + ctx.run_op = control_flow_ops.group(replicate_outputs) + + if isinstance(replicate_outputs, list): + # Filter out any ops from the outputs, typically this would be the case + # when there were no tensor outputs. + last_step_tensor_outputs = [ + x for x in replicate_outputs if not isinstance(x, ops.Operation) + ] + + # Outputs are currently of the structure (flattened) + # [output0_device0, output1_device0, output2_device0, + # output0_device1, output1_device1, output2_device1, + # ...] + # Convert this to the following structure instead: (grouped by output) + # [[output0_device0, output0_device1], + # [output1_device0, output1_device1], + # [output2_device0, output2_device1]] + output_num = len(last_step_tensor_outputs) // self._num_replicas_in_sync + last_step_tensor_outputs = [ + last_step_tensor_outputs[i::output_num] for i in range(output_num) + ] + else: + # no tensors returned. + last_step_tensor_outputs = [] + + _set_last_step_outputs(ctx, last_step_tensor_outputs) + return ctx + + def _run_steps_on_iterator_with_device_loop( + self, fn, multi_worker_iterator, iterations, initial_loop_values=None): + output_shapes = multi_worker_iterator.output_shapes + shapes = nest.flatten(output_shapes) + if any(not s.is_fully_defined() for s in shapes): raise ValueError( - 'TPU currently requires fully defined shapes. Either use ' - 'set_shape() on the input tensors or use ' - 'dataset.batch(..., drop_remainder=True).') - types = nest.flatten(iterator.output_types) + "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, iterator, shapes, iterations) + 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(iterator.output_shapes, dequeued) + return nest.pack_sequence_as(output_shapes, dequeued) # Wrap `fn` for repeat. if initial_loop_values is None: initial_loop_values = {} initial_loop_values = nest.flatten(initial_loop_values) - ctx = values.MultiStepContext() + ctx = input_lib.MultiStepContext() + def run_fn(*args, **kwargs): """Single step on the TPU device.""" del args, kwargs - fn_inputs = dequeue_fn() - if not isinstance(fn_inputs, tuple): - fn_inputs = (fn_inputs,) - fn_result = fn(ctx, *fn_inputs) + fn_result = fn(ctx, dequeue_fn()) flat_last_step_outputs = nest.flatten(ctx.last_step_outputs) if flat_last_step_outputs: with ops.control_dependencies([fn_result]): @@ -256,8 +544,6 @@ class TPUStrategy(one_device_strategy.OneDeviceStrategy): else: return fn_result - # TODO(sourabhbajaj): The input to while loop should be based on the output - # type of the step_fn def iterate_on_tpu(): return training_loop.repeat(iterations, run_fn, initial_loop_values) @@ -270,8 +556,9 @@ class TPUStrategy(one_device_strategy.OneDeviceStrategy): self._outer_control_flow_context = ( ops.get_default_graph()._get_control_flow_context()) # pylint: disable=protected-access - replicate_inputs = [[]] * self.num_towers + replicate_inputs = [[]] * self._num_replicas_in_sync replicate_outputs = tpu.replicate(iterate_on_tpu, replicate_inputs) + del self._outer_control_flow_context ctx.run_op = control_flow_ops.group(replicate_outputs, enqueue_ops) @@ -287,144 +574,121 @@ class TPUStrategy(one_device_strategy.OneDeviceStrategy): # [[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)] - - # 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, aggregation) in ctx._last_step_outputs_aggregations.items(): # pylint: disable=protected-access - output = last_step_tensor_outputs_dict[name] - # For outputs that have already been aggregated, take the first value - # from the list as each value should be the same. Else return the full - # list of values. - # TODO(josh11b): If aggregation is NONE, we should return a PerDevice value. - if aggregation is not variables_lib.VariableAggregation.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 + last_step_tensor_outputs = [list(x) for x in + zip(*last_step_tensor_outputs)] + _set_last_step_outputs(ctx, last_step_tensor_outputs) return ctx - def _call_for_each_tower(self, fn, *args, **kwargs): - # TODO(jhseu): Consider making it so call_for_each_tower implies that we're - # in a tpu.rewrite(), and update TPUMirroredVariable accordingly. - kwargs.pop('run_concurrently', None) - with one_device_strategy._OneDeviceTowerContext(self): # pylint: disable=protected-access + def _call_for_each_replica(self, fn, args, kwargs): + # TODO(jhseu): Consider making it so call_for_each_replica implies that + # we're in a tpu.rewrite(), and update TPUMirroredVariable accordingly. + with _TPUReplicaContext(self._container_strategy()): return fn(*args, **kwargs) - def 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 [tpu.shutdown_system()] + def _experimental_initialize_system(self): + """Experimental method added to be used by Estimator. - def _get_devices_from(self, colocate_with=None): - # TODO(jhseu): Change this when we support model parallelism. - return self._tpu_devices + This is a private method only to be used by Estimator. Other frameworks + should directly be calling `tf.contrib.distribute.initialize_tpu_system` + """ + initialize_tpu_system(self._tpu_cluster_resolver) def _create_variable(self, next_creator, *args, **kwargs): """Create a TPUMirroredVariable. See `DistributionStrategy.scope`.""" colocate_with = kwargs.pop("colocate_with", None) - devices = self._get_devices_from(colocate_with) + if colocate_with is None: + device_map = self._device_map + logical_device = 0 # TODO(josh11b): Get logical device from scope here. + elif isinstance(colocate_with, numpy_dataset.SingleDevice): + with ops.device(colocate_with.device): + return next_creator(*args, **kwargs) + else: + device_map = colocate_with.device_map + logical_device = colocate_with.logical_device def _real_mirrored_creator(devices, *args, **kwargs): # pylint: disable=g-missing-docstring - index = {} + value_list = [] for i, d in enumerate(devices): with ops.device(d): if i > 0: # Give replicas meaningful distinct names: - var0name = index[devices[0]].name.split(":")[0] - # We append a / to variable names created on towers with id > 0 to + var0name = value_list[0].name.split(":")[0] + # We append a / to variable names created on replicas with id > 0 to # ensure that we ignore the name scope and instead use the given # name as the absolute name of the variable. kwargs["name"] = "%s/replica_%d/" % (var0name, i) # Initialize replicas with the same value: if context.executing_eagerly(): kwargs["initial_value"] = array_ops.identity( - index[devices[0]].value()) + value_list[0].value()) else: def initial_value_fn(device=d): with ops.device(device): - return array_ops.identity(index[devices[0]].initial_value) + return array_ops.identity(value_list[0].initial_value) kwargs["initial_value"] = initial_value_fn with context.context().device_policy(context.DEVICE_PLACEMENT_SILENT): v = next_creator(*args, **kwargs) assert not isinstance(v, values.TPUMirroredVariable) - index[d] = v - return index + value_list.append(v) + return value_list - return _create_tpu_mirrored_variable(devices, _real_mirrored_creator, *args, - **kwargs) + return _create_tpu_mirrored_variable( + self._container_strategy(), device_map, logical_device, + _real_mirrored_creator, *args, **kwargs) - def _reduce(self, aggregation, value, destinations): + def _reduce_to(self, reduce_op, value, destinations): if values._enclosing_tpu_context() is not None: # pylint: disable=protected-access - if aggregation == vs.VariableAggregation.MEAN: + if reduce_op == reduce_util.ReduceOp.MEAN: # TODO(jhseu): Revisit once we support model-parallelism. - value *= (1. / self.num_towers) - elif aggregation != vs.VariableAggregation.SUM: + value *= (1. / self._num_replicas_in_sync) + elif reduce_op != reduce_util.ReduceOp.SUM: raise NotImplementedError( "Currently only support sum & mean in TPUStrategy.") return tpu_ops.cross_replica_sum(value) + if not isinstance(value, values.DistributedValues): + # This function handles reducing values that are not PerReplica or + # Mirrored values. For example, the same value could be present on all + # replicas in which case `value` would be a single value or value could + # be 0. + return cross_device_ops_lib.reduce_non_distributed_value( + reduce_op, self._device_map, value, destinations) + # 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_tower_ops_lib.get_devices_from(destinations) + devices = cross_device_ops_lib.get_devices_from(destinations) if len(devices) == 1: assert device_util.canonicalize(devices[0]) == device_util.canonicalize( - self.get_host_cpu_device(0)) + self._host_device) else: - raise ValueError('Multiple devices are not supported for TPUStrategy') + raise ValueError("Multiple devices are not supported for TPUStrategy") - if aggregation == vs.VariableAggregation.ONLY_FIRST_TOWER: - return value[0] output = math_ops.add_n(value) - if aggregation == vs.VariableAggregation.MEAN: + if reduce_op == reduce_util.ReduceOp.MEAN: return output * (1. / len(value)) return output - def _update(self, var, options, fn, *args, **kwargs): + def _update(self, var, fn, args, kwargs, group): assert isinstance(var, values.TPUMirroredVariable) - should_group = options.pop("grouped") - assert not options # Validate that we are processing all of the options. - if values._enclosing_tpu_context() is not None: # pylint: disable=protected-access - if should_group: + if group: return fn(var, *args, **kwargs) else: - return [fn(var, *args, **kwargs)] + return (fn(var, *args, **kwargs),) # Otherwise, we revert to MirroredStrategy behavior and update each variable # directly. - updates = {} - for d, v in var._index.items(): # pylint: disable=protected-access - name = "update_%d" % self._device_index.get(d) + updates = [] + for i, (d, v) in enumerate(zip(var.devices, var.values)): + name = "update_%d" % i with ops.device(d), distribute_lib.UpdateContext(d), ops.name_scope(name): # If args and kwargs are not mirrored, the value is returned as is. - updates[d] = fn(v, - *values.select_device_mirrored(d, args), - **values.select_device_mirrored(d, kwargs)) - return values.update_regroup(self, updates, should_group) - - # TODO(josh11b): Need to implement _update_non_slot()! + updates.append(fn(v, + *values.select_device_mirrored(d, args), + **values.select_device_mirrored(d, kwargs))) + return values.update_regroup(self, self._device_map, updates, group) def read_var(self, var): assert isinstance(var, values.TPUMirroredVariable) @@ -433,36 +697,63 @@ class TPUStrategy(one_device_strategy.OneDeviceStrategy): def _unwrap(self, val): if isinstance(val, values.DistributedValues): # Return in a deterministic order. - return [val.get(device=d) for d in sorted(val.devices)] + return tuple(val.get(device=d) for d in sorted(val.devices)) elif isinstance(val, list): # TODO(josh11b): We need to remove this case; per device values should - # be represented using a PerDevice wrapper instead of a list with + # be represented using a PerReplica wrapper instead of a list with # one entry per device. - return val - return [val] + return tuple(val) + elif isinstance(val, values.TPUMirroredVariable): + # pylint: disable=protected-access + if values._enclosing_tpu_context() is not None: + return (val,) + return val.values + return (val,) - @property - def num_towers(self): - return self._num_cores_override or self._tpu_metadata.num_cores + def value_container(self, value): + return value + + def _broadcast_to(self, tensor, destinations): + del destinations + return tensor @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_towers_per_host(self): - return self._tpu_metadata.num_of_cores_per_host + def num_replicas_per_host(self): + if self._device_assignment is None: + return self._tpu_metadata.num_of_cores_per_host + + # TODO(sourabhbajaj): Remove this method we use inputs and remove infeed + # as the computation of num_replicas_per_host is not a constant + # when using device_assignment. This is a temporary workaround to support + # StatefulRNN as everything is 1 in that case. + # This method needs to take host_id as input for correct computation. + max_models_per_host = (self._tpu_metadata.num_of_cores_per_host // + self._device_assignment.num_cores_per_replica) + models_per_host = min(self._device_assignment.num_replicas, + max_models_per_host) + return models_per_host * self._device_assignment.num_cores_per_replica @property - def num_replicas_in_sync(self): - return self.num_towers + def _num_replicas_in_sync(self): + if self._device_assignment is None: + return self._tpu_metadata.num_cores + return (self._device_assignment.num_replicas * + self._device_assignment.num_cores_per_replica) @property - def between_graph(self): + def experimental_between_graph(self): return False @property - def should_init(self): + def experimental_should_init(self): return True @property @@ -481,20 +772,104 @@ class TPUStrategy(one_device_strategy.OneDeviceStrategy): def parameter_devices(self): return self._tpu_devices + def non_slot_devices(self, var_list): + return self._host_device + + def _update_non_slot(self, colocate_with, fn, args, kwargs, group): + del colocate_with + with ops.device(self._host_device), distribute_lib.UpdateContext( + self._host_device): + result = fn(*args, **kwargs) + if group: + return result + else: + return nest.map_structure(self._unwrap, result) + + def get_host(self, host_id): + if self._tpu_cluster_resolver.get_master() in ("", "local"): + return "/replica:0/task:0" + job_name = self._tpu_cluster_resolver.get_job_name() or "tpu_worker" + return "/job:%s/task:%d" % (job_name, host_id) + def get_host_cpu_device(self, host_id): - if self._tpu_cluster_resolver.get_master() in ('', 'local'): - return '/replica:0/task:0/device:CPU:0' - job_name = self._tpu_cluster_resolver.get_job_name() or 'tpu_worker' - return '/job:%s/task:%d/device:CPU:0' % (job_name, host_id) - - def configure(self, - session_config=None, - cluster_spec=None, - task_type=None, - task_id=None): + return self.get_host(host_id) + "/device:CPU:0" + + def _configure(self, + session_config=None, + cluster_spec=None, + task_type=None, + task_id=None): del cluster_spec, task_type, task_id if session_config: - session_config.isolate_session_state = True - cluster_spec = self._tpu_cluster_resolver.cluster_spec() - if cluster_spec: - session_config.cluster_def.CopyFrom(cluster_spec.as_cluster_def()) + session_config.CopyFrom(self._update_config_proto(session_config)) + + def _update_config_proto(self, config_proto): + updated_config = copy.deepcopy(config_proto) + updated_config.isolate_session_state = True + cluster_spec = self._tpu_cluster_resolver.cluster_spec() + if cluster_spec: + updated_config.cluster_def.CopyFrom(cluster_spec.as_cluster_def()) + return updated_config + + # TODO(priyag): Delete this once all strategies use global batch size. + @property + def _global_batch_size(self): + """`make_dataset_iterator` and `make_numpy_iterator` use global batch size. + + `make_input_fn_iterator` assumes per-replica batching. + + Returns: + Boolean. + """ + return True + + +class _TPUReplicaContext(distribute_lib.ReplicaContext): + """Replication Context class for TPU Strategy.""" + + # TODO(sourabhbajaj): Call for each replica should be updating this. + # TODO(b/118385803): Always properly initialize replica_id. + def __init__(self, strategy, replica_id_in_sync_group=None): + if replica_id_in_sync_group is None: + replica_id_in_sync_group = constant_op.constant(0, dtypes.int32) + distribute_lib.ReplicaContext.__init__( + self, strategy, replica_id_in_sync_group=replica_id_in_sync_group) + + @property + def devices(self): + distribute_lib.require_replica_context(self) + ds = self._strategy + replica_id = tensor_util.constant_value(self._replica_id_in_sync_group) + + if replica_id is None: # Non-constant `Tensor` inside `tpu.replicate`. + # TODO(cjfj): Return other devices when model parallelism is supported. + return (tpu.core(0),) + else: + return (ds.extended.worker_devices[replica_id],) + + +def _get_host_for_device(device): + spec = tf_device.DeviceSpec.from_string(device) + return tf_device.DeviceSpec( + job=spec.job, replica=spec.replica, task=spec.task, + device_type="CPU", device_index=0).to_string() + + +def _set_last_step_outputs(ctx, last_step_tensor_outputs): + """Sets the last step outputs on the given context.""" + # Convert replicate_outputs to the original dict structure of + # last_step_outputs. + last_step_tensor_outputs_dict = nest.pack_sequence_as( + ctx.last_step_outputs, last_step_tensor_outputs) + + for name, reduce_op in ctx._last_step_outputs_reduce_ops.items(): # pylint: disable=protected-access + output = last_step_tensor_outputs_dict[name] + # For outputs that have already been reduced, take the first value + # from the list as each value should be the same. Else return the full + # list of values. + # TODO(josh11b): If reduce_op is NONE, we should return a PerReplica + # value. + if reduce_op is not None: + # TODO(priyag): Should this return the element or a list with 1 element + last_step_tensor_outputs_dict[name] = output[0] + ctx._set_last_step_outputs(last_step_tensor_outputs_dict) # pylint: disable=protected-access diff --git a/tensorflow/contrib/distribute/python/values.py b/tensorflow/contrib/distribute/python/values.py deleted file mode 100644 index c555dc8a71d0ce56caa95265b39a85578ee85545..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/distribute/python/values.py +++ /dev/null @@ -1,1643 +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. -# ============================================================================== -"""Various classes representing distributed values. - -See go/tf-distribution-strategy. -""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import collections -import contextlib -import weakref -import six - -from tensorflow.contrib.distribute.python import input_ops -from tensorflow.python.data.ops import multi_device_iterator_ops -from tensorflow.python.eager import context -from tensorflow.python.eager import tape -from tensorflow.python.framework import device as tf_device -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 gen_resource_variable_ops -from tensorflow.python.ops import math_ops -from tensorflow.python.ops import variable_scope as vs -from tensorflow.python.ops import variables as variables_lib -from tensorflow.python.training import device_util -from tensorflow.python.training import distribute as distribute_lib -from tensorflow.python.training import distribution_strategy_context -from tensorflow.python.training import saver -from tensorflow.python.training.checkpointable import base as checkpointable -from tensorflow.python.util import nest - - -# pylint: disable=line-too-long -# TODO(josh11b): Should device values be strings or DeviceSpec objects? -# Not sure DeviceSpec objects are usable as a dict key. -class DistributedValues(object): - """Holds a map from device to values. Either PerDevice or Mirrored.""" - - def __init__(self, index): - self._index = {device_util.canonicalize(key): value - for key, value in six.iteritems(index)} - - def get(self, device=None): - """Returns the value for the current device or raises a ValueError.""" - if device is None: - tower_context = distribution_strategy_context.get_tower_context() - if tower_context: - device = tower_context.device - else: - device = distribute_lib.get_update_device() - if device is None: - return self._get_cross_tower() - device = device_util.canonicalize(device) - try: - return self._index[device] - except KeyError as e: - six.raise_from( - ValueError("Device %s not found in %s (current device %s)" % - (device, self._index.keys(), device_util.current())), e) - - def on_device(self, device): - device = device_util.canonicalize(device) - return device in self._index - - @property - def devices(self): - return list(self._index.keys()) - - @property - def is_tensor_like(self): - for v in self._index.values(): - if not tensor_util.is_tensor(v): - return False - return True - - def __str__(self): - return "%s:%s" % (self.__class__.__name__, self._index) - - def __repr__(self): - return "%s(%r)" % (self.__class__.__name__, self._index) - - # TODO(josh11b): Possibly make an accessor for _index for use by - # DistributionStrategy implementations. - - -class DistributedDelegate(DistributedValues): - """A map from device to values; acts as the same type as the values.""" - - def __init__(self, index): - super(DistributedDelegate, self).__init__(index) - - def __getattr__(self, name): - return getattr(self.get(), name) - - # pylint: disable=multiple-statements - def __add__(self, o): return self.get() + o - def __radd__(self, o): return o + self.get() - def __sub__(self, o): return self.get() - o - def __rsub__(self, o): return o - self.get() - def __mul__(self, o): return self.get() * o - def __rmul__(self, o): return o * self.get() - def __truediv__(self, o): return self.get() / o - def __rtruediv__(self, o): return o / self.get() - def __floordiv__(self, o): return self.get() // o - def __rfloordiv__(self, o): return o // self.get() - def __mod__(self, o): return self.get() % o - def __rmod__(self, o): return o % self.get() - def __lt__(self, o): return self.get() < o - def __le__(self, o): return self.get() <= o - def __gt__(self, o): return self.get() > o - def __ge__(self, o): return self.get() >= o - def __and__(self, o): return self.get() & o - def __rand__(self, o): return o & self.get() - def __or__(self, o): return self.get() | o - def __ror__(self, o): return o | self.get() - def __xor__(self, o): return self.get() ^ o - def __rxor__(self, o): return o ^ self.get() - def __getitem__(self, o): return self.get()[o] - def __pow__(self, o, modulo=None): return pow(self.get(), o, modulo) - def __rpow__(self, o): return pow(o, self.get()) - def __invert__(self): return ~self.get() - def __neg__(self): return -self.get() - def __abs__(self): return abs(self.get()) - - def __div__(self, o): - try: - return self.get().__div__(o) - except AttributeError: - # See https://docs.python.org/3/library/constants.html#NotImplemented - return NotImplemented - - def __rdiv__(self, o): - try: - return self.get().__rdiv__(o) - except AttributeError: - # See https://docs.python.org/3/library/constants.html#NotImplemented - return NotImplemented - - def __matmul__(self, o): - try: - return self.get().__matmul__(o) - except AttributeError: - # See https://docs.python.org/3/library/constants.html#NotImplemented - return NotImplemented - - def __rmatmul__(self, o): - try: - return self.get().__rmatmul__(o) - except AttributeError: - # See https://docs.python.org/3/library/constants.html#NotImplemented - return NotImplemented - - # TODO(josh11b): Even more operator overloads. - - -class PerDevice(DistributedValues): - """Holds a map from device to unsynchronized values.""" - pass - - -# Note that unlike PerDevice, Mirrored values inherit from -# DistributedDelegate and so can be used directly in cross-tower mode. -class Mirrored(DistributedDelegate): - """Holds a map from device to values which are kept in sync.""" - - def _get_cross_tower(self): - device = device_util.canonicalize(device_util.current()) - if device in self._index: - return self._index[device] - return list(self._index.values())[0] - - def _as_graph_element(self): - obj = self.get() - # pylint: disable=protected-access - conv_fn = getattr(obj, "_as_graph_element", None) - if conv_fn and callable(conv_fn): - return conv_fn() - return obj - - -def _assign_on_device(device, variable, tensor): - with ops.device(device): - return variable.assign(array_ops.identity(tensor)) - - -DistributedVarOp = collections.namedtuple( - "DistributedVarOp", ["name", "graph", "type"]) - - -class DistributedVariable(DistributedDelegate): - """Holds a map from device to variables.""" - # TODO(josh11b): Support changing the set of variables if e.g. if new - # devices are joining or a device is to leave. - - def __init__(self, index): - # Child class must set self._primary_var before calling - # super(...).__init__(index). - self._common_name = self._primary_var.name.split(":")[0] - # Use a weakref to make it easy to map from the contained values - # to the container without introducing a reference cycle. - for v in six.itervalues(index): - v._distributed_container = weakref.ref(self) # pylint: disable=protected-access - # tf.keras keeps track of variables initialized using this attribute. When - # tf.keras gets the default session, it initializes all uninitialized vars. - # We need to make _keras_initialized a member of DistributedVariable because - # without this it will use `__getattr__` which will delegate to a component - # variable. - self._keras_initialized = False - # Typically, a `DistributedVariable`'s initializer is composed of the - # initializers of the components variables. However, in some cases, such as - # when restoring from a checkpoint, we may set the _initializer_op - # property on the entire `DistributedVariable`. - self._initializer_op = None - super(DistributedVariable, self).__init__(index) - - def is_initialized(self, name=None): - """Identifies if all the component variables are initialized. - - Args: - name: Name of the final `logical_and` op. - - Returns: - The op that evaluates to True or False depending on if all the - component variables are initialized. - """ - # We have to cast the self._index.values() to a `list` because when we - # use `model_to_estimator` to run tf.keras models, self._index.values() is - # of type `dict_values` and not `list`. - values_list = list(self._index.values()) - result = values_list[0].is_initialized() - # We iterate through the list of values except the last one to allow us to - # name the final `logical_and` op the same name that is passed by the user - # to the `is_initialized` op. For distributed variables, the - # `is_initialized` op is a `logical_and` op. - for v in values_list[1:-1]: - result = math_ops.logical_and(result, v.is_initialized()) - result = math_ops.logical_and(result, values_list[-1].is_initialized(), - name=name) - return result - - @property - def initializer(self): - if self._initializer_op: - init_op = self._initializer_op - else: - # return grouped ops of all the var initializations of component values of - # the mirrored variable - init_op = control_flow_ops.group( - [v.initializer for v in self._index.values()]) - return init_op - - @property - def graph(self): - return self._primary_var.graph - - @property - def _shared_name(self): - return self._common_name - - @property - def _unique_id(self): - return self._primary_var._unique_id # pylint: disable=protected-access - - @property - def name(self): - return self._primary_var.name - - @property - def dtype(self): - return self._primary_var.dtype - - @property - def shape(self): - return self._primary_var.shape - - def get_shape(self): - return self._primary_var.get_shape() - - def to_proto(self, export_scope=None): - return self._primary_var.to_proto(export_scope=export_scope) - - @property - def op(self): - # We want cross-tower code that does some var.op.X calls - # to work (even if the current device isn't in self.devices), but - # other uses of var.op in a cross-tower context to fail. - if distribution_strategy_context.get_cross_tower_context(): - return DistributedVarOp(self._primary_var.op.name, - self._primary_var.op.graph, - self._primary_var.op.type) - return self.get().op - - @property - def _in_graph_mode(self): - return self._primary_var._in_graph_mode # pylint: disable=protected-access - - def read_value(self): - return distribution_strategy_context.get_distribution_strategy().read_var( - self) - - def _should_act_as_resource_variable(self): - """Pass resource_variable_ops.is_resource_variable check.""" - pass - - -ops.register_dense_tensor_like_type(DistributedVariable) - - -class _MirroredSaveable(saver.BaseSaverBuilder.ResourceVariableSaveable): - """Class for defining how to restore a MirroredVariable.""" - - def __init__(self, mirrored_variable, primary_variable, name): - self._mirrored_variable = mirrored_variable - super(_MirroredSaveable, self).__init__(primary_variable, "", name) - - def restore(self, restored_tensors, restored_shapes): - """Restore the same value into all variables.""" - tensor, = restored_tensors - return control_flow_ops.group([ - _assign_on_device(d, v, tensor) - for d, v in six.iteritems(self._mirrored_variable._index)]) # pylint: disable=protected-access - - -class MirroredVariable(DistributedVariable, Mirrored, - checkpointable.CheckpointableBase): - """Holds a map from device to variables whose values are kept in sync.""" - - def __init__(self, index, primary_var, aggregation): - self._primary_var = primary_var - self._aggregation = aggregation - super(MirroredVariable, self).__init__(index) - - # The arguments to update() are automatically unwrapped so the update() - # function would normally see regular variables, not MirroredVariables. - # However, the update function can still operate on wrapped MirroredVariables - # through object members, captured arguments, etc. This is more likely in an - # update_non_slot() function (like OptimizerV2._finish), which can - # update several non-slot variables in one call. - def _assign_func(self, *args, **kwargs): - f = kwargs.pop("f") - if distribution_strategy_context.get_cross_tower_context(): - update_device = distribute_lib.get_update_device() - if update_device is not None: - # We are calling an assign function on the mirrored variable in an - # update context. - v = self.get(device=update_device) - return f(v, *args, **kwargs) - - # We are calling assign on the mirrored variable in cross tower context, - # use update to update the variable. - strategy = distribution_strategy_context.get_distribution_strategy() - return strategy.update(self, f, *args, **kwargs) - else: - _assert_tower_context() - # We are calling an assign function on the mirrored variable in tower - # context. - # We reduce the value we want to assign/add/sub. More details about how we - # handle the different use cases can be found in the _reduce method. - # We call the function on each of the mirrored variables with the reduced - # value. - if self._aggregation == vs.VariableAggregation.NONE: - raise ValueError("You must specify an aggregation method to update a " - "MirroredVariable in Tower Context.") - - def merge_fn(strategy, value, *other_args, **other_kwargs): - return strategy.update( - self, f, - strategy.reduce( - aggregation=self._aggregation, value=value, destinations=self), - *other_args, **other_kwargs) - - return distribution_strategy_context.get_tower_context().merge_call( - merge_fn, *args, **kwargs) - - def assign_sub(self, *args, **kwargs): - assign_sub_fn = lambda var, *a, **kw: var.assign_sub(*a, **kw) - return self._assign_func(f=assign_sub_fn, *args, **kwargs) - - def assign_add(self, *args, **kwargs): - assign_add_fn = lambda var, *a, **kw: var.assign_add(*a, **kw) - return self._assign_func(f=assign_add_fn, *args, **kwargs) - - def assign(self, *args, **kwargs): - assign_fn = lambda var, *a, **kw: var.assign(*a, **kw) - return self._assign_func(f=assign_fn, *args, **kwargs) - - @property - def aggregation(self): - return self._aggregation - - def _get_cross_tower(self): - device = device_util.canonicalize(device_util.current()) - if device in self._index: - return array_ops.identity(self._index[device]) - return array_ops.identity(self._primary_var) - - def _as_graph_element(self): - # pylint: disable=protected-access - if distribution_strategy_context.get_cross_tower_context(): - return self._primary_var._as_graph_element() - return self.get()._as_graph_element() - - def _gather_saveables_for_checkpoint(self): - """Overrides CheckpointableBase method. - - This allows both name-based and object-based save and restore of - MirroredVariables. - - Returns: - A dictionary mapping attribute names to `SaveableObject` factories. - """ - def _saveable_factory(name=self._common_name): - return _MirroredSaveable(self, self._primary_var, name) - return {checkpointable.VARIABLE_VALUE_KEY: _saveable_factory} - - -# Register a conversion function which reads the value of the variable, -# allowing instances of the class to be used as tensors. -def _tensor_conversion_mirrored(var, dtype=None, name=None, as_ref=False): - # Try to avoid assignments to and other mutations of MirroredVariable - # state except through a DistributionStrategy.update() call. - assert not as_ref - return ops.internal_convert_to_tensor( - var.get(), dtype=dtype, name=name, as_ref=as_ref) - - -ops.register_tensor_conversion_function(MirroredVariable, - _tensor_conversion_mirrored) - - -def _enclosing_tpu_context(): - # pylint: disable=protected-access - tpu_context = ops.get_default_graph()._get_control_flow_context() - # pylint: enable=protected-access - while tpu_context is not None and not isinstance( - tpu_context, control_flow_ops.XLAControlFlowContext): - tpu_context = tpu_context.outer_context - return tpu_context - - -# TODO(jhseu): Deduplicate code. We copy code because we don't want to -# inherit from DistributedDelegate. DistributedDelegate will not work in a -# tpu.replicate() because it assumes that you're in a device context where you -# can operate on a single version of the variable, but a tpu.replicate() -# operates on all variables and is replicated during a rewrite pass. -class TPUMirroredVariable(checkpointable.CheckpointableBase): - """Holds a map from device to TPU variables whose values are kept in sync.""" - - def __init__(self, index, primary_var, aggregation): - # Use a weakref to make it easy to map from the contained values - # to the container without introducing a reference cycle. - for v in six.itervalues(index): - v._mirrored_container = weakref.ref(self) # pylint: disable=protected-access - self._index = {device_util.canonicalize(key): value - for key, value in six.iteritems(index)} - self._primary_var = primary_var - self._common_name = self._primary_var.name.split(":")[0] - self._aggregation = aggregation - # Needed for GradientTape - self._trainable = self._primary_var.trainable - # Typically like `DistributedVariable`, a `TPUMirroredVariable`'s - # initializer is composed of the initializers of the components variables. - # However, in some cases, such as when restoring from a checkpoint, we may - # set the _initializer_op property on the entire `TPUMirroredVariable`. - self._initializer_op = None - - def _get(self, device=None): - """Returns the value for the current device or raises a ValueError.""" - if device is None: - tower_context = distribution_strategy_context.get_tower_context() - if tower_context: - device = tower_context.device - else: - device = distribute_lib.get_update_device() - if device is None: - return self._get_cross_tower() - device = device_util.canonicalize(device) - try: - return self._index[device] - except KeyError as e: - six.raise_from( - ValueError("Device %s not found in %s (current device %s)" % - (device, self._index.keys(), device_util.current())), e) - - # pylint: disable=multiple-statements - def __add__(self, o): return self.read_value() + o - def __radd__(self, o): return o + self.read_value() - def __sub__(self, o): return self.read_value() - o - def __rsub__(self, o): return o - self.read_value() - def __mul__(self, o): return self.read_value() * o - def __rmul__(self, o): return o * self.read_value() - def __truediv__(self, o): return self.read_value() / o - def __rtruediv__(self, o): return o / self.read_value() - def __floordiv__(self, o): return self.read_value() // o - def __rfloordiv__(self, o): return o // self.read_value() - def __mod__(self, o): return self.read_value() % o - def __rmod__(self, o): return o % self.read_value() - def __lt__(self, o): return self.read_value() < o - def __le__(self, o): return self.read_value() <= o - def __gt__(self, o): return self.read_value() > o - def __ge__(self, o): return self.read_value() >= o - def __and__(self, o): return self.read_value() & o - def __rand__(self, o): return o & self.read_value() - def __or__(self, o): return self.read_value() | o - def __ror__(self, o): return o | self.read_value() - def __xor__(self, o): return self.read_value() ^ o - def __rxor__(self, o): return o ^ self.read_value() - def __getitem__(self, o): return self.read_value()[o] - def __pow__(self, o, modulo=None): return pow(self.read_value(), o, modulo) - def __rpow__(self, o): return pow(o, self.read_value()) - def __invert__(self): return ~self.read_value() - def __neg__(self): return -self.read_value() - def __abs__(self): return abs(self.read_value()) - - def __div__(self, o): - try: - return self.read_value().__div__(o) - except AttributeError: - # See https://docs.python.org/3/library/constants.html#NotImplemented - return NotImplemented - - def __rdiv__(self, o): - try: - return self.read_value().__rdiv__(o) - except AttributeError: - # See https://docs.python.org/3/library/constants.html#NotImplemented - return NotImplemented - - def __matmul__(self, o): - try: - return self.read_value().__matmul__(o) - except AttributeError: - # See https://docs.python.org/3/library/constants.html#NotImplemented - return NotImplemented - - def __rmatmul__(self, o): - try: - return self.read_value().__rmatmul__(o) - except AttributeError: - # See https://docs.python.org/3/library/constants.html#NotImplemented - return NotImplemented - - @property - def handle(self): - # If we're in a tpu.rewrite(), return the replicated handle. - tpu_context = _enclosing_tpu_context() - if tpu_context is not None: - return tpu_context.get_replicated_var_handle( - self._common_name, nest.flatten(self._index)) - - device = distribute_lib.get_update_device() - if device is None: - return self._primary_var.handle - device = device_util.canonicalize(device) - try: - return self._index[device].handle - except KeyError as e: - six.raise_from( - ValueError("Device %s not found in %s (current device %s)" % - (device, self._index.keys(), device_util.current())), e) - - @property - def device(self): - return self._get().device - - # The arguments to update() are automatically unwrapped so the update() - # function would normally see regular variables, not MirroredVariables. - # However, the update function can still operate on wrapped MirroredVariables - # through object members, captured arguments, etc. This is more likely in an - # update_non_slot() function (like OptimizerV2._finish), which can - # update several non-slot variables in one call. - def _assign_func(self, *args, **kwargs): - if distribution_strategy_context.get_distribution_strategy().__class__.__name__ != "TPUStrategy": - raise ValueError("You may only assign to a TPUMirroredVariable within a " - "TPUStrategy.") - f = kwargs.pop("f") - if distribution_strategy_context.get_cross_tower_context(): - if _enclosing_tpu_context() is not None: - return distribution_strategy_context.get_distribution_strategy().update( - self, f, *args, **kwargs) - - update_device = distribute_lib.get_update_device() - # We are calling update on the mirrored variable in cross tower context. - if update_device is not None: - # We are calling an assign function on the mirrored variable in cross - # tower context. - v = self._get(device=update_device) - return f(v, *args, **kwargs) - - return distribution_strategy_context.get_distribution_strategy().update( - self, f, *args, **kwargs) - else: - _assert_tower_context() - # We are calling an assign function on the mirrored variable in tower - # context. - # We reduce the value we want to assign/add/sub. More details about how we - # handle the different use cases can be found in the _reduce method. - # We call the function on each of the mirrored variables with the reduced - # value. - if self._aggregation == vs.VariableAggregation.NONE: - raise ValueError("You must specify an aggregation method to update a " - "TPUMirroredVariable in Tower Context.") - - def merge_fn(strategy, value, *other_args, **other_kwargs): - return strategy.update( - self, f, - strategy.reduce( - aggregation=self._aggregation, value=value, destinations=self), - *other_args, **other_kwargs) - - return distribution_strategy_context.get_tower_context().merge_call( - merge_fn, *args, **kwargs) - - @contextlib.contextmanager - def _handle_graph(self, handle): - # Note: might have an eager tensor but not be executing eagerly when - # building functions. - if (context.executing_eagerly() or isinstance(handle, ops.EagerTensor) - or ops.has_default_graph()): - yield - else: - with handle.graph.as_default(): - yield - - @property - def trainable(self): - return self._trainable - - def _read_variable_op(self, parent_op=None): - if self.trainable: - tape.variable_accessed(self) - if parent_op is not None: - with ops.control_dependencies([parent_op]): - return gen_resource_variable_ops.read_variable_op( - self.handle, self.dtype) - - return gen_resource_variable_ops.read_variable_op( - self.handle, self.dtype) - - def read_value(self): - return self._read_variable_op() - - def assign_sub(self, *args, **kwargs): - def assign_sub_fn(var, delta, **kw): - name = kw.pop("name", None) - read_value = kw.pop("read_value", True) - with self._handle_graph(var.handle): - op = gen_resource_variable_ops.assign_sub_variable_op( - var.handle, ops.convert_to_tensor(delta, dtype=self.dtype), - name=name) - if read_value: - return self._read_variable_op(parent_op=op) - return op - - return self._assign_func(f=assign_sub_fn, *args, **kwargs) - - def assign_add(self, *args, **kwargs): - def assign_add_fn(var, delta, **kw): - name = kw.pop("name", None) - read_value = kw.pop("read_value", True) - with self._handle_graph(var.handle): - op = gen_resource_variable_ops.assign_add_variable_op( - var.handle, ops.convert_to_tensor(delta, dtype=self.dtype), - name=name) - if read_value: - return self._read_variable_op(parent_op=op) - return op - - return self._assign_func(f=assign_add_fn, *args, **kwargs) - - def assign(self, *args, **kwargs): - def assign_fn(var, value, **kw): - name = kw.pop("name", None) - read_value = kw.pop("read_value", True) - with self._handle_graph(var.handle): - op = gen_resource_variable_ops.assign_variable_op( - var.handle, ops.convert_to_tensor(value, dtype=self.dtype), - name=name) - if read_value: - return self._read_variable_op(parent_op=op) - return op - - return self._assign_func(f=assign_fn, *args, **kwargs) - - @property - def aggregation(self): - return self._aggregation - - @property - def constraint(self): - return None - - @property - def initializer(self): - if self._initializer_op: - init_op = self._initializer_op - else: - init_op = control_flow_ops.group( - [v.initializer for v in self._index.values()]) - return init_op - - @property - def graph(self): - return self._primary_var.graph - - @property - def _shared_name(self): - return self._common_name - - @property - def _unique_id(self): - return self._primary_var._unique_id # pylint: disable=protected-access - - @property - def name(self): - return self._primary_var.name - - @property - def dtype(self): - return self._primary_var.dtype - - @property - def shape(self): - return self._primary_var.shape - - def get_shape(self): - return self._primary_var.get_shape() - - def to_proto(self, export_scope=None): - return self._primary_var.to_proto(export_scope=export_scope) - - def _get_cross_tower(self): - device = device_util.canonicalize(device_util.current()) - if device in self._index: - return self._index[device] - return self._primary_var - - def _as_graph_element(self): - # pylint: disable=protected-access - if distribution_strategy_context.get_cross_tower_context(): - return self._primary_var._as_graph_element() - return self._read_variable_op() - - def _gather_saveables_for_checkpoint(self): - """Overrides CheckpointableBase method. - - This allows both name-based and object-based save and restore of - MirroredVariables. - - Returns: - A dictionary mapping attribute names to `SaveableObject` factories. - """ - def _saveable_factory(name=self._common_name): - return _MirroredSaveable(self, self._primary_var, name) - return {checkpointable.VARIABLE_VALUE_KEY: _saveable_factory} - - def _should_act_as_resource_variable(self): - """Pass resource_variable_ops.is_resource_variable check.""" - pass - - # Needed to pass ResourceVariable checks. - @property - def op(self): - return self._primary_var.op - - @property - def _in_graph_mode(self): - return self._primary_var._in_graph_mode # pylint: disable=protected-access - - def _dense_var_to_tensor(self, dtype=None, name=None, as_ref=False): - """Converts a variable to a tensor.""" - # pylint: disable=protected-access - if _enclosing_tpu_context() is None: - return self._get()._dense_var_to_tensor(dtype, name, as_ref) - # pylint: enable=protected-access - if dtype is not None and dtype != self.dtype: - raise NotImplementedError - if as_ref: - return self.handle - else: - return self.read_value() - - def is_initialized(self, name=None): - """Identifies if all the component variables are initialized. - - Args: - name: Name of the final `logical_and` op. - - Returns: - The op that evaluates to True or False depending on if all the - component variables are initialized. - """ - # TODO(jhseu): Do we need TPU context implementation? - - # We have to cast the self._index.values() to a `list` because when we - # use `model_to_estimator` to run tf.keras models, self._index.values() is - # of type `dict_values` and not `list`. - values_list = nest.flatten(self._index) - result = values_list[0].is_initialized() - # We iterate through the list of values except the last one to allow us to - # name the final `logical_and` op the same name that is passed by the user - # to the `is_initialized` op. For distributed variables, the - # `is_initialized` op is a `logical_and` op. - for v in values_list[1:-1]: - result = math_ops.logical_and(result, v.is_initialized()) - result = math_ops.logical_and(result, values_list[-1].is_initialized(), - name=name) - return result - - -# Register a conversion function which reads the value of the variable, -# allowing instances of the class to be used as tensors. -def _tensor_conversion_tpu_mirrored(var, dtype=None, name=None, as_ref=False): - return var._dense_var_to_tensor(dtype=dtype, name=name, as_ref=as_ref) # pylint: disable=protected-access - - -ops.register_tensor_conversion_function(TPUMirroredVariable, - _tensor_conversion_tpu_mirrored) -ops.register_dense_tensor_like_type(TPUMirroredVariable) - - -class _TowerLocalSaveable(saver.BaseSaverBuilder.SaveableObject): - """Class for defining how to restore a TowerLocalVariable.""" - - def __init__(self, tower_local_variable, name): - self._tower_local_variable = tower_local_variable - # We use a callable so that we don't have to evaluate this expression - # in the case where we are trying to restore instead of save. - def tensor(): - return distribution_strategy_context.get_distribution_strategy().read_var( - tower_local_variable) - spec = saver.BaseSaverBuilder.SaveSpec( - tensor=tensor, - slice_spec="", - name=name, - dtype=tower_local_variable.dtype) - super(_TowerLocalSaveable, self).__init__(tensor, [spec], name) - - def restore(self, restored_tensors, restored_shapes): - """Restore the same value into all variables.""" - tensor, = restored_tensors - return self._tower_local_variable.assign(tensor) - - -def _assert_tower_context(): - if not distribution_strategy_context.get_tower_context(): - raise RuntimeError( - "Tower-local variables may only be assigned in a tower context.") - - -class TowerLocalVariable(DistributedVariable, PerDevice, - checkpointable.CheckpointableBase): - """Holds a map from device to variables whose values are reduced on save.""" - - def __init__(self, index, primary_var, aggregation): - self._primary_var = primary_var - self._aggregation = aggregation - super(TowerLocalVariable, self).__init__(index) - - def assign_sub(self, *args, **kwargs): - _assert_tower_context() - return self.get().assign_sub(*args, **kwargs) - - def assign_add(self, *args, **kwargs): - _assert_tower_context() - return self.get().assign_add(*args, **kwargs) - - def assign(self, *args, **kwargs): - if distribution_strategy_context.get_cross_tower_context(): - # To preserve the sum across save and restore, we have to divide the - # total across all devices when restoring a variable that was summed - # when saving. - tensor = args[0] - if self._aggregation == vs.VariableAggregation.SUM: - tensor *= 1. / len(self.devices) - return control_flow_ops.group( - [_assign_on_device(d, v, tensor) - for d, v in six.iteritems(self._index)]) - else: - _assert_tower_context() - return self.get().assign(*args, **kwargs) - - @property - def aggregation(self): - return self._aggregation - - def _get_cross_tower(self): - if self._aggregation == vs.VariableAggregation.ONLY_FIRST_TOWER: - return self._primary_var - all_components = tuple(self._index.values()) - # TODO(josh11b): Use a strategy-specific method. - total = math_ops.add_n(all_components) - if self._aggregation == vs.VariableAggregation.MEAN: - return total * (1./ len(all_components)) - return total - - def _as_graph_element(self): - # pylint: disable=protected-access - if distribution_strategy_context.get_cross_tower_context(): - return self._get_cross_tower() - return self.get()._as_graph_element() - - def _gather_saveables_for_checkpoint(self): - """Overrides CheckpointableBase method. - - This allows both name-based and object-based save and restore of - TowerLocalVariables. - - Returns: - A dictionary mapping attribute names to `SaveableObject` factories. - """ - def _saveable_factory(name=self._common_name): - return _TowerLocalSaveable(self, name) - return {checkpointable.VARIABLE_VALUE_KEY: _saveable_factory} - - -# Register a conversion function for TowerLocalVariable which allows as_ref to -# be true. -def _tensor_conversion_tower_local(var, dtype=None, name=None, as_ref=False): - return ops.internal_convert_to_tensor( - var.get(), dtype=dtype, name=name, as_ref=as_ref) - - -ops.register_tensor_conversion_function(TowerLocalVariable, - _tensor_conversion_tower_local) - - -def _devices_match(d1, d2): - return device_util.canonicalize(d1) == device_util.canonicalize(d2) - - -def regroup(per_device, wrap_class=PerDevice): - """Makes device->nest map into a nest of PerDevice/Mirrored values.""" - items = list(per_device.items()) - assert items - v0 = items[0][1] # First value - - if isinstance(v0, list): - for _, v in items[1:]: - assert isinstance(v, list) - assert len(v) == len(v0), ("len(v) == %d, len(v0) == %d, v: %s, v0: %s" % - (len(v), len(v0), v, v0)) - return [regroup({k: v[i] for k, v in items}, wrap_class) - for i in range(len(v0))] - - if isinstance(v0, tuple): - for _, v in items[1:]: - assert isinstance(v, tuple) - assert len(v) == len(v0) - regrouped_tuple = tuple(regroup({k: v[i] for k, v in items}, wrap_class) - for i in range(len(v0))) - if hasattr(v0, "_fields"): - # This tuple is in fact a namedtuple! Create a new namedtuple instance - # and initialize it with the regrouped values: - assert hasattr(type(v0), "_make") - return type(v0)._make(regrouped_tuple) - else: - return regrouped_tuple - - if isinstance(v0, dict): - v0keys = set(v0.keys()) - for _, v in items[1:]: - assert isinstance(v, dict) - assert set(v.keys()) == v0keys - return {key: regroup({k: v[key] for k, v in items}, wrap_class) - for key in v0keys} - - # If exactly the same object across all devices, return it unwrapped. - same_id = True - for _, v in items[1:]: - if v is not v0: - same_id = False - break - # Consider three cases where same_id is true: - # * If v0 is a DistributedVariable (a MirroredVariable or - # TowerLocalVariable, and same_id means it is the same across all - # devices), we want to return it. We check DistributedVariable - # specifically since it can look like it has a - # _distributed_container member since its members do. - # * If v0 is a member of a distributed variable, in which case - # hasattr(v0, "_distributed_container") is true, we want to - # return the DistributedVariable that contains it using the - # _distributed_container logic below. This case can trigger - # same_id when there is only one device. - # * In any other situation, same_id means we return v0. - if same_id and (isinstance(v0, DistributedVariable) or - not hasattr(v0, "_distributed_container")): - return v0 - - # Detect the case where each device has a parallel component of the - # same MirroredVariable (or TowerLocalVariable). In this case we - # want to return the containing MirroredVariable, after a bunch of - # sanity checking. In particular, each component should have the - # same container, and the devices of the variables should match the - # keys of the per-device dictionary. - if hasattr(v0, "_distributed_container"): - # pylint: disable=protected-access - assert not isinstance(v0, MirroredVariable), ( - "ids = %s, items = %s" % ([id(v[1]) for v in items], items)) - assert _devices_match(v0.device, items[0][0]), ( - "v0.device = %s, items = %s" % (v0.device, items)) - distributed_container = v0._distributed_container() - assert distributed_container is not None - for d, v in items[1:]: - assert _devices_match(v.device, d), ( - "v.device = %s, d = %s, items = %s" % (v.device, d, items)) - assert distributed_container is v._distributed_container() - return distributed_container - # pylint: enable=protected-access - - return wrap_class(per_device) - - -def select_device(device, structured): - """Specialize a nest of regular & per-device values for one device.""" - def _get(x): - return x.get(device) if isinstance(x, DistributedValues) else x - - return nest.map_structure(_get, structured) - - -def select_device_mirrored(device, structured): - """Specialize a nest of regular & mirrored values for one device.""" - def _get_mirrored(x): - if isinstance(x, DistributedValues): - if not isinstance(x, Mirrored): - raise TypeError( - "Expected value to be mirrored across towers: %s in %s." % - (x, structured)) - return x.get(device) - else: - return x - - return nest.map_structure(_get_mirrored, structured) - - -def update_regroup(strategy, updates, should_group): - """Regroup for an update, with dependencies to ensure all updates execute.""" - regrouped = regroup(updates, Mirrored) - if not should_group: - return nest.map_structure(strategy.unwrap, regrouped) - grouped_flat = [] - for u in nest.flatten(regrouped): - if isinstance(u, DistributedValues): - g = strategy.group(u) - if u.is_tensor_like: - # Make sure we run all updates. Without this, something like - # session.run(strategy.update(...)) may only update one tower. - index = {} - for d in u.devices: - with ops.device(d), ops.control_dependencies([g]): - index[d] = array_ops.identity(u.get(d)) - g = Mirrored(index) - else: - g = u - grouped_flat.append(g) - return nest.pack_sequence_as(regrouped, grouped_flat) - - -class PerDeviceDataIterator(object): - """An iterator (like `tf.data.Iterator`) into a `PerDeviceDataset`.""" - - def __init__(self, iterator, devices, prefetch_on_device=None): - self._iterator = iterator - self._devices = devices - self._prefetch_on_device = prefetch_on_device - - @property - def initializer(self): - return self._iterator.initializer - - def get_next(self, name=None): - """Scatter the input across devices.""" - if self._prefetch_on_device: - data_list = self._iterator.get_next() - index = dict(zip(self._devices, data_list)) - else: - batch = self._iterator.get_next(name=name) - index = {} - def get_ith(i): - return lambda x: x[i] - - for i, d in enumerate(self._devices): - index[d] = nest.map_structure(get_ith(i), batch) - if context.executing_eagerly(): - with ops.device(d): - index[d] = nest.map_structure(array_ops.identity, index[d]) - - return regroup(index) - - -class PerDeviceDataset(object): - """Like `tf.data.Dataset` split devices, producing `PerDevice` data.""" - - def __init__(self, dataset, devices, prefetch_on_device=None): - self._devices = devices - - # Default to using prefetching in graph mode, unless specified. - # TODO(rohanj): Enable prefetching in eager mode. - self._prefetch_on_device = prefetch_on_device - if self._prefetch_on_device is None: - self._prefetch_on_device = not context.executing_eagerly() - assert not (self._prefetch_on_device and context.executing_eagerly()), ( - "Prefetching is only supported in graph mode currently") - - self._dataset = dataset - if not self._prefetch_on_device: - # TODO(priyag): If dropping remainder is not appropriate, find another - # approach to distributing the dataset when not possible to divide evenly. - # Possibly not an issue when we start using PartitionedDataset. - self._dataset = dataset.batch(len(devices), drop_remainder=True) - - def make_one_shot_iterator(self): - """Get a one time use iterator for the distributed PerDeviceDataset.""" - # Graph mode with one shot iterator is disabled. - if not context.executing_eagerly(): - raise ValueError("Cannot create a one shot iterator. Please use " - "`make_initializable_iterator()` instead.") - # Eager mode prefetching would error out in constructor. Only remaining - # case is non-prefetching in eager mode. We delegate to - # PerDeviceDataIterator to handle that case. - dataset_iterator = self._dataset.make_one_shot_iterator() - return PerDeviceDataIterator( - dataset_iterator, self._devices, prefetch_on_device=False) - - def make_initializable_iterator(self): - """Get an initializable iterator for the distributed PerDeviceDataset.""" - # Eager mode generates already initialized iterators. Hence we cannot create - # an initializable iterator. - if context.executing_eagerly(): - raise ValueError("Cannot create initializable iterator in Eager mode. " - "Please use `make_one_shot_iterator` instead.") - if self._prefetch_on_device: - dataset_iterator = multi_device_iterator_ops.MultiDeviceIterator( - self._dataset, self._devices) - else: - dataset_iterator = self._dataset.make_initializable_iterator() - return PerDeviceDataIterator( - dataset_iterator, - self._devices, - prefetch_on_device=self._prefetch_on_device) - - -class MultiWorkerDataIterator(object): - """An iterator (like `tf.data.Iterator`) into a `MultiWorkerDataset`.""" - - def __init__(self, iterators, worker_device_map): - """Initialize the MultiWorkerDataIterator object. - - Args: - iterators: a dict mapping from each worker to an iterator for - that worker. - worker_device_map: a dict mapping from each worker's devices to a list of - devices that belong to this worker. - - Raises: - ValueError: if iterators and worker_device_map are not compatible. - """ - self._iterators = iterators - self._worker_device_map = worker_device_map - if set(self._iterators) != set(self._worker_device_map): - raise ValueError("iterators and worker_device_map are not compatible.") - - @property - def initializer(self): - return control_flow_ops.group( - [iterator.initializer for iterator in self._iterators.values()]) - - def get_next(self, name=None): - """Scatter the input across hosts and devices.""" - index = {} - for worker, iterator in six.iteritems(self._iterators): - if name is not None: - d = tf_device.DeviceSpec.from_string(worker) - new_name = "%s_%s_%d" % (name, d.job, d.task) - else: - new_name = None - with ops.device(worker): - data_per_worker = iterator.get_next(name=new_name) - - worker_devices = self._worker_device_map[worker] - # Ungroup these per-device value so as to get a flat map from devices to - # values. - for d in worker_devices: - v = select_device(d, data_per_worker) - if d in index: - raise ValueError("Duplicated devices in worker_device_map: %r" % v) - index[d] = v - - return regroup(index) - - -class MultiWorkerDataset(object): - """Like a `tf.data.Dataset` that distributes data to different workers. - - Each worker gets one shard of the input dataset. It is currently not working - in - eager mode. - """ - - def __init__(self, dataset_fn, worker_device_map, prefetch_on_device=None, - auto_shard=False): - """Initialize the MultiWorkerDataset object. - - Args: - dataset_fn: a function that returns a `tf.data.Dataset`. - worker_device_map: a dict mapping from each worker to a list of devices - that belong to this worker. - prefetch_on_device: whether to prefetch to devices. - auto_shard: whether to auto-shard the dataset. - """ - self._worker_device_map = worker_device_map - self._datasets = {} - # TODO(yuefengz, priyag): support different set of jobs for input - # processing. - for i, (worker, worker_devices) in enumerate( - six.iteritems(worker_device_map)): - with ops.device(worker): - worker_input = dataset_fn() - if auto_shard: - worker_input = input_ops.auto_shard_dataset( - worker_input, len(worker_device_map), i) - self._datasets[worker] = PerDeviceDataset( - worker_input, worker_devices, prefetch_on_device=prefetch_on_device) - - def make_one_shot_iterator(self): - iterators = {} - for worker, dataset in six.iteritems(self._datasets): - with ops.device(worker): - iterators[worker] = dataset.make_one_shot_iterator() - return MultiWorkerDataIterator(iterators, self._worker_device_map) - - def make_initializable_iterator(self): - iterators = {} - for worker, dataset in six.iteritems(self._datasets): - with ops.device(worker): - iterators[worker] = dataset.make_initializable_iterator() - return MultiWorkerDataIterator(iterators, self._worker_device_map) - - -class _PerKey(object): - """Holds data associated by keys.""" - - def __init__(self, *index): - # pylint: disable=protected-access - self._index = list(index) - - def get(self, iteration): - return array_ops.gather(self._index, iteration) - - def get_shape(self): - return self._index[-1][-1].get_shape() - - def get_dtype(self): - return self._index[-1][-1].dtype - - def __str__(self): - return "%s:%s" % (self.__class__.__name__, self._index) - - def __repr__(self): - return "%s(%r)" % (self.__class__.__name__, self._index) - - -class PerIteration(_PerKey): - """Holds input for multiple iterations at once.""" - - def __init__(self, *index): - # pylint: disable=protected-access - super(PerIteration, self).__init__(*[batch._index for batch in index]) - - -class Batches(_PerKey): - pass - - -class MultiIterator(object): - """Iterator that returns results of multiple get_next()s.""" - - def __init__(self, dataset_iterator, iterations, batches_per_iteration): - self._dataset_iterator = dataset_iterator - self._iterations = iterations - self._batches_per_iteration = batches_per_iteration - - def get_next(self, name=None): - """Return PerIteration with `iterations x batches_per_iteration` inputs.""" - data = [] - for _ in range(self._batches_per_iteration): - batch = [] - for _ in range(self._iterations): - batch.append(self._dataset_iterator.get_next(name=name)) - data.append(batch) - - # Here is an example. Suppose each get_next returns a tuple of two tensors. - # For 3 `iterations` and 2 `batches_per_iteration`, the `data` is: - # [[(a,z), (b,y), (c,x)], [(A,Z), (B,Y), (C,X)]] - # - # After the first `map_structure` it gets transformed to: - # [(Batches(a, A), Batches(z, Z)), - # (Batches(b, B), Batches(y, Y)), - # (Batches(c, C), Batches(x, X))] - # - # After the second `map_structure` it gets transformed to a tuple of: - # (PerIteration([Batches(a, A), Batches(b, B), Batches(c, C)]), - # PerIteration([Batches(z, Z), Batches(y, Y), Batches(x, X)])) - - data = nest.map_structure(Batches, *data) - data = nest.map_structure(PerIteration, *data) - - return data - - @property - def initializer(self): - return self._dataset_iterator.initializer - - -class PerIterationDataset(object): - """A dataset that returns MultiIterators.""" - - def __init__(self, dataset, iterations, batches_per_iteration): - self._dataset = dataset - self._iterations = iterations - self._batches_per_iteration = batches_per_iteration - - def make_one_shot_iterator(self): - iterator = self._dataset.make_one_shot_iterator() - return MultiIterator(iterator, self._iterations, - self._batches_per_iteration) - - def make_initializable_iterator(self): - iterator = self._dataset.make_initializable_iterator() - return MultiIterator(iterator, self._iterations, - self._batches_per_iteration) - - -class MapOutput(object): - """Map can result in multiple outputs per device.""" - - def __init__(self, l): - self._l = l - - def get(self): - return self._l - - -class MultiStepContext(object): - """A context object that can be used to capture things when running steps. - - This context object is useful when running multiple steps at a time using the - `run_steps_on_dataset` API. For e.g. it allows the user's step function to - specify which outputs to emit at what frequency. Currently it supports - capturing output from the last step, as well as capturing non tensor outputs. - In the future it will be augmented to support other use cases such as output - each N steps. - """ - - def __init__(self): - """Initializes an output context. - - Returns: - A context object. - """ - self._last_step_outputs = {} - self._last_step_outputs_aggregations = {} - self._non_tensor_outputs = {} - - @property - def last_step_outputs(self): - """A dictionary consisting of outputs to be captured on last step. - - Keys in the dictionary are names of tensors to be captured, as specified - when `set_last_step_output` is called. - Values in the dictionary are the tensors themselves. If - `set_last_step_output` was called with an `aggregation` for this output, - then the value is the aggregated value. - - Returns: - A dictionary with last step outputs. - """ - return self._last_step_outputs - - def _set_last_step_outputs(self, outputs): - """Replace the entire dictionary of last step outputs.""" - if not isinstance(outputs, dict): - raise ValueError("Need a dictionary to set last_step_outputs.") - self._last_step_outputs = outputs - - def set_last_step_output(self, name, output, - aggregation=variables_lib.VariableAggregation.NONE): - """Set `output` with `name` to be outputted from the last step. - - Args: - name: String, name to identify the output. Doesn't need to match tensor - name. - output: The tensors that should be outputted with `name`. See below for - actual types supported. - aggregation: Aggregation method to use to aggregate outputs from multiple - towers. Required if `set_last_step_output` is called in a tower context. - Optional in cross_tower_context. - When present, the outputs from all the towers are aggregated using the - current distribution strategy's `reduce` method. Hence, the type of - `output` must be what's supported by the corresponding `reduce` method. - For e.g. if using MirroredStrategy and aggregation is set, output - must be a `PerDevice` value. - The aggregation method is also recorded in a dictionary - `_last_step_outputs_aggregations` for later interpreting of the - outputs as already reduced or not. - - """ - if distribution_strategy_context.get_cross_tower_context(): - self._last_step_outputs_aggregations[name] = aggregation - if aggregation is variables_lib.VariableAggregation.NONE: - self._last_step_outputs[name] = output - else: - distribution = distribution_strategy_context.get_distribution_strategy() - self._last_step_outputs[name] = distribution.reduce( - aggregation, output, destinations="/device:CPU:0") - else: - assert aggregation is not variables_lib.VariableAggregation.NONE - def merge_fn(distribution, value): - self._last_step_outputs[name] = distribution.reduce( - aggregation, value, destinations="/device:CPU:0") - # Setting this inside the `merge_fn` because all towers share the same - # context object, so it's more robust to set it only once (even if all - # the towers are trying to set the same value). - self._last_step_outputs_aggregations[name] = aggregation - - distribution_strategy_context.get_tower_context().merge_call( - merge_fn, output) - - @property - def non_tensor_outputs(self): - """A dictionary consisting of any non tensor outputs to be captured.""" - return self._non_tensor_outputs - - def set_non_tensor_output(self, name, output): - """Set `output` with `name` to be captured as a non tensor output.""" - if distribution_strategy_context.get_cross_tower_context(): - self._non_tensor_outputs[name] = output - else: - def merge_fn(distribution, value): - # NOTE(priyag): For non tensor outputs, we simply return all the values - # in a list as aggregation doesn't make sense on non tensors. - self._non_tensor_outputs[name] = distribution.unwrap(value) - distribution_strategy_context.get_tower_context().merge_call( - merge_fn, output) - - -def value_container(val): - """Returns the container that this per-device `value` belongs to. - - Args: - val: A value returned by `call_for_each_tower()` or a variable - created in `scope()`. - - Returns: - A container that `value` belongs to. - If value does not belong to any container (including the case of - container having been destroyed), returns the value itself. - """ - # pylint: disable=protected-access - if (hasattr(val, "_distributed_container") and - # DistributedVariable has _distributed_container defined - # but we don't want to return it. - not isinstance(val, DistributedVariable)): - container = val._distributed_container() - # pylint: disable=protected-access - if container is not None: - return container - return val - - -# TODO(josh11b): Descend from Variable. -class AggregatingVariable(checkpointable.CheckpointableBase): - """A wrapper around a variable that aggregates updates across towers.""" - - def __init__(self, v, aggregation): - self._v = v - # TODO(josh11b): Set v._distributed_container? - # v._distributed_container = weakref.ref(self) # pylint: disable=protected-access - self._aggregation = aggregation - - def get(self): - return self._v - - def __getattr__(self, name): - return getattr(self._v, name) - - def _assign_func(self, *args, **kwargs): - f = kwargs.pop("f") - if distribution_strategy_context.get_cross_tower_context(): - update_device = distribute_lib.get_update_device() - if update_device is not None: - # We are calling an assign function in an update context. - return f(self._v, *args, **kwargs) - - # We are calling an assign function in cross tower context, wrap it in an - # update call. - return distribution_strategy_context.get_distribution_strategy().update( - self, f, *args, **kwargs) - else: - assert distribution_strategy_context.get_tower_context() - # We are calling an assign function in tower context. - # We reduce the value we want to assign/add/sub. More details about how we - # handle the different use cases can be found in the _reduce method. - # We call the function with the reduced value. - if self._aggregation == vs.VariableAggregation.NONE: - raise ValueError("You must specify an aggregation method to update a " - "a variable in Tower Context.") - - def merge_fn(strategy, value, *other_args, **other_kwargs): - return strategy.update( - self, f, - strategy.reduce( - aggregation=self._aggregation, value=value, destinations=self), - *other_args, **other_kwargs) - - return distribution_strategy_context.get_tower_context().merge_call( - merge_fn, *args, **kwargs) - - def assign_sub(self, *args, **kwargs): - assign_sub_fn = lambda var, *a, **kw: var.assign_sub(*a, **kw) - return self._assign_func(f=assign_sub_fn, *args, **kwargs) - - def assign_add(self, *args, **kwargs): - assign_add_fn = lambda var, *a, **kw: var.assign_add(*a, **kw) - return self._assign_func(f=assign_add_fn, *args, **kwargs) - - def assign(self, *args, **kwargs): - assign_fn = lambda var, *a, **kw: var.assign(*a, **kw) - return self._assign_func(f=assign_fn, *args, **kwargs) - - @property - def aggregation(self): - return self._aggregation - - @property - def name(self): - return self._v.name - - @property - def dtype(self): - return self._v.dtype - - # TODO(josh11b): Test saving & restoring. - def _gather_saveables_for_checkpoint(self): - return {checkpointable.VARIABLE_VALUE_KEY: self._v} - - # pylint: disable=multiple-statements - def __add__(self, o): return self._v + o - def __radd__(self, o): return o + self._v - def __sub__(self, o): return self._v - o - def __rsub__(self, o): return o - self._v - def __mul__(self, o): return self._v * o - def __rmul__(self, o): return o * self._v - def __truediv__(self, o): return self._v / o - def __rtruediv__(self, o): return o / self._v - def __floordiv__(self, o): return self._v // o - def __rfloordiv__(self, o): return o // self._v - def __mod__(self, o): return self._v % o - def __rmod__(self, o): return o % self._v - def __lt__(self, o): return self._v < o - def __le__(self, o): return self._v <= o - def __gt__(self, o): return self._v > o - def __ge__(self, o): return self._v >= o - def __and__(self, o): return self._v & o - def __rand__(self, o): return o & self._v - def __or__(self, o): return self._v | o - def __ror__(self, o): return o | self._v - def __xor__(self, o): return self._v ^ o - def __rxor__(self, o): return o ^ self._v - def __getitem__(self, o): return self._v[o] - def __pow__(self, o, modulo=None): return pow(self._v, o, modulo) - def __rpow__(self, o): return pow(o, self._v) - def __invert__(self): return ~self._v - def __neg__(self): return -self._v - def __abs__(self): return abs(self._v) - - def __div__(self, o): - try: - return self._v.__div__(o) - except AttributeError: - # See https://docs.python.org/3/library/constants.html#NotImplemented - return NotImplemented - - def __rdiv__(self, o): - try: - return self._v.__rdiv__(o) - except AttributeError: - # See https://docs.python.org/3/library/constants.html#NotImplemented - return NotImplemented - - def __matmul__(self, o): - try: - return self._v.__matmul__(o) - except AttributeError: - # See https://docs.python.org/3/library/constants.html#NotImplemented - return NotImplemented - - def __rmatmul__(self, o): - try: - return self._v.__rmatmul__(o) - except AttributeError: - # See https://docs.python.org/3/library/constants.html#NotImplemented - return NotImplemented - - def __str__(self): - return str(self._v) - - def __repr__(self): - return repr(self._v) - - def _should_act_as_resource_variable(self): - """Pass resource_variable_ops.is_resource_variable check.""" - pass - - -# Register a conversion function which reads the value of the variable, -# allowing instances of the class to be used as tensors. -def _tensor_conversion_aggregate(var, dtype=None, name=None, as_ref=False): - return ops.internal_convert_to_tensor( - var.get(), dtype=dtype, name=name, as_ref=as_ref) - - -ops.register_tensor_conversion_function( - AggregatingVariable, _tensor_conversion_aggregate) -ops.register_dense_tensor_like_type(AggregatingVariable) diff --git a/tensorflow/contrib/distribute/python/values_test.py b/tensorflow/contrib/distribute/python/values_test.py index 7ef4776ac6d414470c1597358063f6e77960728f..51c58b0b2f3dc2ab63e22718825a471b8657f892 100644 --- a/tensorflow/contrib/distribute/python/values_test.py +++ b/tensorflow/contrib/distribute/python/values_test.py @@ -18,29 +18,24 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import collections import os +from absl.testing import parameterized -from tensorflow.contrib.distribute.python import mirrored_strategy -from tensorflow.contrib.distribute.python import multi_worker_test_base -from tensorflow.contrib.distribute.python import values +from tensorflow.contrib.distribute.python import combinations 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 values from tensorflow.python.eager import context from tensorflow.python.eager import test from tensorflow.python.estimator import model_fn as model_fn_lib from tensorflow.python.framework import constant_op -from tensorflow.python.framework import errors from tensorflow.python.framework import ops from tensorflow.python.framework import tensor_util from tensorflow.python.framework import test_util from tensorflow.python.ops import array_ops -from tensorflow.python.ops import random_ops from tensorflow.python.ops import variable_scope from tensorflow.python.ops import variables as variables_lib -from tensorflow.python.training import device_util from tensorflow.python.training import saver as saver_lib -from tensorflow.python.util import nest class DistributedValuesTest(test.TestCase): @@ -49,7 +44,8 @@ class DistributedValuesTest(test.TestCase): with ops.device("/device:CPU:0"): one = constant_op.constant(1) two = constant_op.constant(2) - v = values.DistributedValues({"/device:CPU:0": one, "/device:GPU:0": two}) + device_map = values.ReplicaDeviceMap(("/device:CPU:0", "/device:GPU:0")) + v = values.DistributedValues(device_map, (one, two)) self.assertEqual(two, v.get("/device:GPU:0")) self.assertEqual(one, v.get()) with self.assertRaises(ValueError): @@ -61,24 +57,26 @@ class DistributedValuesTest(test.TestCase): ops.device("/device:CPU:0"): one = constant_op.constant(1) two = constant_op.constant(2) - v = values.DistributedValues({"/device:CPU:0": one, "/device:GPU:0": two}) + device_map = values.ReplicaDeviceMap(("/device:CPU:0", "/device:GPU:0")) + v = values.DistributedValues(device_map, (one, two)) self.assertEqual(two, v.get("/device:GPU:0")) self.assertEqual(one, v.get()) with self.assertRaises(ValueError): self.assertIsNone(v.get("/device:GPU:2")) def testCanonicalization(self): - canonical_cpu = ["/job:localhost/replica:0/task:0/device:CPU:0"] - v = values.DistributedValues({"": 42}) - self.assertEqual(canonical_cpu, list(v._index.keys())) - v = values.DistributedValues({"/device:CPU:0": 42}) - self.assertEqual(canonical_cpu, list(v._index.keys())) - v = values.DistributedValues({"/cpu:0": 42}) - self.assertEqual(canonical_cpu, list(v._index.keys())) - v = values.DistributedValues({"/CPU:0": 42}) - self.assertEqual(canonical_cpu, list(v._index.keys())) + canonical_cpu = ("/job:localhost/replica:0/task:0/device:CPU:0",) + v = values.DistributedValues(values.SingleDeviceMap(""), (42,)) + self.assertEqual(canonical_cpu, v.devices) + v = values.DistributedValues(values.SingleDeviceMap("/device:CPU:0"), (42,)) + self.assertEqual(canonical_cpu, v.devices) + v = values.DistributedValues(values.SingleDeviceMap("/cpu:0"), (42,)) + self.assertEqual(canonical_cpu, v.devices) + v = values.DistributedValues(values.SingleDeviceMap("/CPU:0"), (42,)) + self.assertEqual(canonical_cpu, v.devices) with self.assertRaises(AssertionError): - v = values.DistributedValues({"/device:cpu:0": 42}) + v = values.DistributedValues( + values.SingleDeviceMap("/device:cpu:0"), (42,)) def testIsTensorLike(self): with context.graph_mode(), \ @@ -86,7 +84,8 @@ class DistributedValuesTest(test.TestCase): ops.device("/device:CPU:0"): one = constant_op.constant(1) two = constant_op.constant(2) - v = values.DistributedValues({"/device:CPU:0": one, "/device:GPU:0": two}) + device_map = values.ReplicaDeviceMap(("/device:CPU:0", "/device:GPU:0")) + v = values.DistributedValues(device_map, (one, two)) self.assertEqual(two, v.get("/device:GPU:0")) self.assertEqual(one, v.get()) self.assertTrue(v.is_tensor_like) @@ -98,7 +97,8 @@ class DistributedValuesTest(test.TestCase): ops.device("/device:CPU:0"): one = constant_op.constant(1) two = 2.0 - v = values.DistributedValues({"/device:CPU:0": one, "/device:GPU:0": two}) + device_map = values.ReplicaDeviceMap(("/device:CPU:0", "/device:GPU:0")) + v = values.DistributedValues(device_map, (one, two)) self.assertEqual(two, v.get("/device:GPU:0")) self.assertEqual(one, v.get()) self.assertFalse(v.is_tensor_like) @@ -116,8 +116,8 @@ class DistributedDelegateTest(test.TestCase): def __init__(self, x): self.x = x - v = values.DistributedDelegate( - {"/device:CPU:0": Foo(7), "/device:GPU:0": Foo(8)}) + device_map = values.ReplicaDeviceMap(("/device:CPU:0", "/device:GPU:0")) + v = values.DistributedDelegate(device_map, (Foo(7), Foo(8))) self.assertEqual(7, v.x) with self.assertRaises(AttributeError): _ = v.y @@ -125,7 +125,8 @@ class DistributedDelegateTest(test.TestCase): @test_util.run_in_graph_and_eager_modes def testOperatorOverride(self): with ops.device("/device:CPU:0"): - v = values.DistributedDelegate({"/device:CPU:0": 7, "/device:GPU:0": 8}) + device_map = values.ReplicaDeviceMap(("/device:CPU:0", "/device:GPU:0")) + v = values.DistributedDelegate(device_map, (7, 8)) # v should act like int(7). self.assertEqual(8, v + 1) self.assertEqual(10, 3 + v) @@ -176,24 +177,23 @@ def _nested_value(d): def _make_mirrored(): v = [] - index = {} devices = ["/device:GPU:0", "/device:CPU:0"] for d, n, init in zip(devices, ["v", "v/replica"], [1., 2.]): with ops.device(d): v.append(variable_scope.get_variable( name=n, initializer=init, use_resource=True)) - index[d] = v[-1] - mirrored = values.MirroredVariable(index, v[0], + device_map = values.ReplicaDeviceMap(devices) + mirrored = values.MirroredVariable(None, device_map, v, variable_scope.VariableAggregation.SUM) - return v, devices, mirrored + return v, device_map, mirrored class RegroupAndSelectDeviceTest(test.TestCase): - def _is_per_device(self, result, expected, klass=values.PerDevice): + def _is_per_replica(self, result, expected, klass=values.PerReplica): self.assertIsInstance(result, klass) # We canonicalize the devices to match the device strings returned - # by PerDevice, which also does device string canonicalization. + # by PerReplica, which also does device string canonicalization. devices = [device_util.canonicalize(_device_str(i)) for i in range(len(expected))] self.assertEqual(set(devices), set(result.devices)) @@ -202,28 +202,29 @@ class RegroupAndSelectDeviceTest(test.TestCase): self.assertEqual(expected[i], result.get(_device_str(i))) def testNested(self): - result = values.regroup({_device_str(0): _nested_value("1"), - _device_str(1): _nested_value("2")}) + device_map = values.ReplicaDeviceMap((_device_str(0), _device_str(1))) + result = values.regroup(device_map, + (_nested_value("1"), _nested_value("2"))) self.assertIsInstance(result, tuple) self.assertEqual(3, len(result)) - self._is_per_device(result[0], ["a1", "a2"]) - self._is_per_device(result[2], ["h1", "h2"]) + self._is_per_replica(result[0], ["a1", "a2"]) + self._is_per_replica(result[2], ["h1", "h2"]) self.assertIsInstance(result[1], list) self.assertEqual(3, len(result[1])) - self._is_per_device(result[1][0], ["b1", "b2"]) - self._is_per_device(result[1][2], ["g1", "g2"]) + self._is_per_replica(result[1][0], ["b1", "b2"]) + self._is_per_replica(result[1][2], ["g1", "g2"]) self.assertIsInstance(result[1][1], dict) self.assertEqual(set(["c", "e"]), set(result[1][1].keys())) - self._is_per_device(result[1][1]["c"], ["d1", "d2"]) - self._is_per_device(result[1][1]["e"], ["f1", "f2"]) + self._is_per_replica(result[1][1]["c"], ["d1", "d2"]) + self._is_per_replica(result[1][1]["e"], ["f1", "f2"]) - # Also test that we can undo the merge using select_device() + # Also test that we can undo the merge using select_replica() self.assertEqual(_nested_value("1"), - values.select_device(_device_str(0), result)) + values.select_replica(0, result)) self.assertEqual(_nested_value("2"), - values.select_device(_device_str(1), result)) + values.select_replica(1, result)) # select_device_mirrored() should fail due to non-mirrored values with self.assertRaises(TypeError): values.select_device_mirrored(_device_str(0), result) @@ -233,29 +234,30 @@ class RegroupAndSelectDeviceTest(test.TestCase): def testWrapClass(self): # Normally a mirrored value would be the same across devices, but # for a test it is convenient to be able to tell the values apart. - result = values.regroup({_device_str(0): _nested_value("1"), - _device_str(1): _nested_value("2")}, + device_map = values.ReplicaDeviceMap((_device_str(0), _device_str(1))) + result = values.regroup(device_map, + (_nested_value("1"), _nested_value("2")), values.Mirrored) self.assertIsInstance(result, tuple) self.assertEqual(3, len(result)) - self._is_per_device(result[0], ["a1", "a2"], values.Mirrored) - self._is_per_device(result[2], ["h1", "h2"], values.Mirrored) + self._is_per_replica(result[0], ["a1", "a2"], values.Mirrored) + self._is_per_replica(result[2], ["h1", "h2"], values.Mirrored) self.assertIsInstance(result[1], list) self.assertEqual(3, len(result[1])) - self._is_per_device(result[1][0], ["b1", "b2"], values.Mirrored) - self._is_per_device(result[1][2], ["g1", "g2"], values.Mirrored) + self._is_per_replica(result[1][0], ["b1", "b2"], values.Mirrored) + self._is_per_replica(result[1][2], ["g1", "g2"], values.Mirrored) self.assertIsInstance(result[1][1], dict) self.assertEqual(set(["c", "e"]), set(result[1][1].keys())) - self._is_per_device(result[1][1]["c"], ["d1", "d2"], values.Mirrored) - self._is_per_device(result[1][1]["e"], ["f1", "f2"], values.Mirrored) + self._is_per_replica(result[1][1]["c"], ["d1", "d2"], values.Mirrored) + self._is_per_replica(result[1][1]["e"], ["f1", "f2"], values.Mirrored) - # Also test that we can undo the merge using select_device() + # Also test that we can undo the merge using select_replica() self.assertEqual(_nested_value("1"), - values.select_device(_device_str(0), result)) + values.select_replica(0, result)) self.assertEqual(_nested_value("2"), - values.select_device(_device_str(1), result)) + values.select_replica(1, result)) # Values are marked as mirrored, so select_device_mirrored() is allowed. self.assertEqual(_nested_value("1"), values.select_device_mirrored(_device_str(0), result)) @@ -265,332 +267,86 @@ class RegroupAndSelectDeviceTest(test.TestCase): def testMirroredContainer(self): if context.num_gpus() < 1 and context.executing_eagerly(): self.skipTest("A GPU is not available for this test in eager mode.") - v, devices, mirrored = _make_mirrored() - result = values.regroup(dict(zip(devices, v))) + v, device_map, mirrored = _make_mirrored() + result = values.regroup(device_map, v) self.assertIs(mirrored, result) def testSameId(self): foo = object() - result = values.regroup({_device_str(0): ("a", foo), - _device_str(1): ("b", foo)}) + device_map = values.ReplicaDeviceMap((_device_str(0), _device_str(1))) + result = values.regroup(device_map, (("a", foo), ("b", foo))) self.assertIsInstance(result, tuple) self.assertEqual(2, len(result)) - self._is_per_device(result[0], ["a", "b"]) + self._is_per_replica(result[0], ["a", "b"]) self.assertIs(foo, result[1]) - # Test select_device(), should undo the merge done by regroup(). - result_0 = values.select_device(_device_str(0), result) + # Test select_replica(), should undo the merge done by regroup(). + result_0 = values.select_replica(0, result) self.assertIsInstance(result_0, tuple) self.assertEqual(2, len(result_0)) self.assertEqual("a", result_0[0]) self.assertIs(foo, result_0[1]) - result_1 = values.select_device(_device_str(1), result) + result_1 = values.select_replica(1, result) self.assertIsInstance(result_1, tuple) self.assertEqual(2, len(result_1)) self.assertEqual("b", result_1[0]) self.assertIs(foo, result_1[1]) def testOneDevice(self): - result = values.regroup({_device_str(0): _nested_value("1")}) - # On one device regroup() and select_device() are basically identity. + device_map = values.ReplicaDeviceMap((_device_str(0),)) + result = values.regroup(device_map, (_nested_value("1"),)) + # On one device regroup() and select_replica() are basically identity. self.assertEqual(_nested_value("1"), result) self.assertEqual(_nested_value("1"), - values.select_device(_device_str(0), result)) + values.select_replica(0, result)) # The one exception has to do with MirroredVariables. d = "/device:CPU:0" with ops.device(d): v = variable_scope.get_variable( name="v", initializer=1., use_resource=True) - index = {d: v} - mirrored = values.MirroredVariable(index, v, + device_map = values.ReplicaDeviceMap((d,)) + mirrored = values.MirroredVariable(None, device_map, (v,), variable_scope.VariableAggregation.SUM) - result = values.regroup(index) + result = values.regroup(device_map, (v,)) self.assertIs(mirrored, result) def testNamedTupleEstimatorSpec(self): with context.graph_mode(), ops.Graph().as_default(): - created_estimator_specs = {} - to_regroup = {} + devices = [] + created_estimator_specs = [] for device_id in range(3): spec = model_fn_lib.EstimatorSpec( mode=model_fn_lib.ModeKeys.TRAIN, loss=constant_op.constant(device_id / 2), train_op=array_ops.identity(constant_op.constant(device_id))) - created_estimator_specs[device_id] = spec - to_regroup[_device_str(device_id)] = spec + devices.append(_device_str(device_id)) + created_estimator_specs.append(spec) - merged_estimator_spec = values.regroup(to_regroup) + device_map = values.ReplicaDeviceMap(devices) + merged_estimator_spec = values.regroup( + device_map, created_estimator_specs) self.assertTrue( isinstance(merged_estimator_spec, model_fn_lib.EstimatorSpec)) - self.assertEquals(model_fn_lib.ModeKeys.TRAIN, merged_estimator_spec.mode) + self.assertEqual(model_fn_lib.ModeKeys.TRAIN, merged_estimator_spec.mode) for device_id in range(3): d = _device_str(device_id) - self.assertEquals(created_estimator_specs[device_id].loss, - merged_estimator_spec.loss.get(d)) - self.assertEquals(created_estimator_specs[device_id].train_op, - merged_estimator_spec.train_op.get(d)) + self.assertEqual(created_estimator_specs[device_id].loss, + merged_estimator_spec.loss.get(d)) + self.assertEqual(created_estimator_specs[device_id].train_op, + merged_estimator_spec.train_op.get(d)) # Scaffold is populated by `EstimatorSpec.__new__`. - self.assertEquals(created_estimator_specs[device_id].scaffold, - merged_estimator_spec.scaffold.get(d)) - # Also test that we can undo the merge using select_device() - self.assertEquals(created_estimator_specs[device_id], - values.select_device(_device_str(device_id), + self.assertEqual(created_estimator_specs[device_id].scaffold, + merged_estimator_spec.scaffold.get(d)) + # Also test that we can undo the merge using select_replica() + self.assertEqual(created_estimator_specs[device_id], + values.select_replica(device_id, merged_estimator_spec)) -class PerDeviceDatasetTest(test.TestCase): - - config = config_pb2.ConfigProto() - config.allow_soft_placement = True - - def _test_iterator_no_prefetch(self, devices, dataset, expected_values): - per_device_dataset = values.PerDeviceDataset( - dataset, devices, prefetch_on_device=False) - if context.executing_eagerly(): - iterator = per_device_dataset.make_one_shot_iterator() - else: - iterator = per_device_dataset.make_initializable_iterator() - self.evaluate([iterator.initializer]) - - for expected_value in expected_values: - next_element = iterator.get_next() - actual = self.evaluate([ - values.select_device(d, next_element) for d in devices]) - self.assertEqual(expected_value, actual) - - with self.assertRaises(errors.OutOfRangeError): - next_element = iterator.get_next() - self.evaluate([ - values.select_device(d, next_element) for d in devices]) - - def _test_iterator_with_prefetch(self, devices, dataset, expected_values): - if not context.executing_eagerly(): - per_device_dataset = values.PerDeviceDataset( - dataset, devices, prefetch_on_device=True) - iterator = per_device_dataset.make_initializable_iterator() - self.evaluate([iterator.initializer]) - - for expected_value in expected_values: - next_element = iterator.get_next() - computed_value = self.evaluate( - [values.select_device(d, next_element) for d in devices]) - self.assertEqual(expected_value, computed_value) - - with self.assertRaises(errors.OutOfRangeError): - next_element = iterator.get_next() - self.evaluate([ - values.select_device(d, next_element) for d in devices]) - - def _test_iterator(self, devices, dataset, expected_values): - self._test_iterator_no_prefetch(devices, dataset, expected_values) - self._test_iterator_with_prefetch(devices, dataset, expected_values) - - @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,))) - - per_device_dataset = values.PerDeviceDataset( - dataset, devices, prefetch_on_device=False) - iterator = per_device_dataset.make_initializable_iterator() - - self.evaluate(iterator.initializer) - next_element = iterator.get_next() - 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, iterator, devices, expected_values): - next_element = iterator.get_next() - for device in devices: - v = values.select_device(device, 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: - actual = self.evaluate( - [values.select_device(d, next_element) for d in devices]) - self.assertEqual(expected_value, actual) - - with self.assertRaises(errors.OutOfRangeError): - self.evaluate([values.select_device(d, next_element) for d in devices]) - - def _test_dataset(self, dataset_fn, worker_device_map, devices, - expected_values): - multi_worker_dataset = values.MultiWorkerDataset( - dataset_fn, worker_device_map, prefetch_on_device=False) - multi_worker_iterator = multi_worker_dataset.make_one_shot_iterator() - self._test_iterator(multi_worker_iterator, devices, expected_values) - - def _cpu_devices(self): - worker_device_map = collections.OrderedDict( - [("/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_device_map, devices - - def _cpu_and_one_gpu_devices(self): - # The worker_device_map doesn't have to be a OrderDict object, this is just - # to simplify the testing so that we can pass expected values as a list - # instead of a dict. - worker_device_map = collections.OrderedDict( - [("/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_device_map, devices - - def testDataDistributionOneDevicePerWorker(self): - self.skipTest("Temporarily disabled.") - worker_device_map, devices = self._cpu_devices() - with context.graph_mode(): - dataset_fn = lambda: dataset_ops.Dataset.range(8) - self._test_dataset(dataset_fn, worker_device_map, devices, - [[0, 1], [2, 3], [4, 5], [6, 7]]) - - def testDataDistributionTwoDevicePerWorker(self): - self.skipTest("Temporarily disabled.") - if context.num_gpus() < 1: - self.skipTest("A GPU is not available for this test.") - worker_device_map, 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_device_map, devices, - [[0, 2, 1, 3], [4, 6, 5, 7]]) - - def testTupleDataset(self): - self.skipTest("Temporarily disabled.") - worker_device_map, 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 + 1, (i + 1)**2)] for i in range(0, 8, 2) - ] - self._test_dataset(dataset_fn, worker_device_map, devices, - expected_values) - - def testInitializableIterator(self): - self.skipTest("Temporarily disabled.") - worker_device_map, devices = self._cpu_devices() - with context.graph_mode(): - dataset_fn = lambda: dataset_ops.Dataset.range(8) - multi_worker_dataset = values.MultiWorkerDataset( - dataset_fn, worker_device_map, prefetch_on_device=False) - multi_worker_iterator = multi_worker_dataset.make_initializable_iterator() - - self.evaluate(multi_worker_iterator.initializer) - self._test_iterator(multi_worker_iterator, devices, - [[0, 1], [2, 3], [4, 5], [6, 7]]) - - # After re-initializing the iterator, should be able to iterate again. - self.evaluate(multi_worker_iterator.initializer) - self._test_iterator(multi_worker_iterator, devices, - [[0, 1], [2, 3], [4, 5], [6, 7]]) - - def testValueErrorForIterator(self): - self.skipTest("Temporarily disabled.") - # Incompatiable arguments. - with self.assertRaises(ValueError): - values.MultiWorkerDataIterator({"w1": None}, {"w1": "d1", "w2": "d2"}) - - # Test duplicated devices under same worker. - worker_device_map, _ = self._cpu_devices() - worker_device_map["/job:worker/replica:0/task:0"].append( - "/job:worker/replica:0/task:0/device:CPU:0") - with context.graph_mode(): - dataset_fn = lambda: dataset_ops.Dataset.range(8) - multi_worker_dataset = values.MultiWorkerDataset( - dataset_fn, worker_device_map, prefetch_on_device=False) - multi_worker_iterator = multi_worker_dataset.make_initializable_iterator() - with self.assertRaises(ValueError): - multi_worker_iterator.get_next() - - -class MirroredVariableTest(test.TestCase): +class MirroredVariableTest(test.TestCase, parameterized.TestCase): config = config_pb2.ConfigProto() config.allow_soft_placement = True @@ -602,21 +358,21 @@ class MirroredVariableTest(test.TestCase): v, _, mirrored = _make_mirrored() - self.assertEquals(v[0].name, mirrored.name) - self.assertEquals(v[0].dtype, mirrored.dtype) - self.assertEquals(v[0].shape, mirrored.shape) + self.assertEqual(v[0].name, mirrored.name) + self.assertEqual(v[0].dtype, mirrored.dtype) + self.assertEqual(v[0].shape, mirrored.shape) @test_util.run_in_graph_and_eager_modes(config=config) def testVariableOnAnotherDevice(self): v = variable_scope.get_variable( name="v", initializer=[1.], use_resource=True) - index = {"/job:foo/device:CPU:0": v} - mirrored = values.MirroredVariable(index, v, + device_map = values.ReplicaDeviceMap(("/job:foo/device:CPU:0",)) + mirrored = values.MirroredVariable(None, device_map, (v,), variable_scope.VariableAggregation.MEAN) - self.assertEquals(v.name, mirrored.name) - self.assertEquals(v.dtype, mirrored.dtype) - self.assertEquals(v.shape, mirrored.shape) + self.assertEqual(v.name, mirrored.name) + self.assertEqual(v.dtype, mirrored.dtype) + self.assertEqual(v.shape, mirrored.shape) def _assign_mirrored(self, devices, v, new): for d, var, n in zip(devices, v, new): @@ -639,7 +395,8 @@ class MirroredVariableTest(test.TestCase): self.skipTest("A GPU is not available for this test in eager mode.") with self.cached_session(config=self.config) as sess: - v, devices, mirrored = _make_mirrored() + v, device_map, mirrored = _make_mirrored() + devices = device_map.all_devices # Overwrite the initial values. self._assign_mirrored(devices, v, [3., 4.]) @@ -657,7 +414,8 @@ class MirroredVariableTest(test.TestCase): def _save_mirrored(self): """Save variables with mirroring, returns save_path.""" with self.session(graph=ops.Graph()) as sess: - v, devices, mirrored = _make_mirrored() + v, device_map, mirrored = _make_mirrored() + devices = device_map.all_devices # Overwrite the initial values. self._assign_mirrored(devices, v, [3., 4.]) @@ -702,7 +460,8 @@ class MirroredVariableTest(test.TestCase): def _restore_mirrored(self, save_path): """Restore to variables with mirroring in a fresh graph.""" with self.session(graph=ops.Graph()) as sess: - v, devices, mirrored = _make_mirrored() + v, device_map, mirrored = _make_mirrored() + devices = device_map.all_devices # Overwrite the initial values. self._assign_mirrored(devices, v, [7., 8.]) @@ -736,40 +495,38 @@ class MirroredVariableTest(test.TestCase): save_path = self._save_normal() self._restore_mirrored(save_path) - @test_util.run_in_graph_and_eager_modes(config=config) - def testFetchAMirroredVariable(self): - if context.num_gpus() < 1 or context.executing_eagerly(): - self.skipTest("A GPU is not available for this test or it's eager mode.") - - with self.session( - graph=ops.Graph()) as sess, mirrored_strategy.MirroredStrategy( - ["/device:GPU:0"]).scope(): + @combinations.generate(combinations.combine( + distribution=[ + combinations.mirrored_strategy_with_one_gpu, + combinations.core_mirrored_strategy_with_one_gpu], + mode=["graph"])) + def testFetchAMirroredVariable(self, distribution): + with self.session(graph=ops.Graph()) as sess, distribution.scope(): with ops.device("/device:GPU:0"): v = variable_scope.get_variable( name="v", initializer=1., use_resource=True) - mirrored = values.MirroredVariable({ - "/device:GPU:0": v - }, v, variable_scope.VariableAggregation.MEAN) + mirrored = values.MirroredVariable( + distribution, values.ReplicaDeviceMap(("/device:GPU:0",)), (v,), + variable_scope.VariableAggregation.MEAN) sess.run(variables_lib.global_variables_initializer()) sess.run({"complicated": mirrored}) -_devices = ["/device:GPU:0", "/device:CPU:0"] +_devices = ("/device:GPU:0", "/device:CPU:0") -def _make_tower_local(method): +def _make_replica_local(method, strategy=None): + device_map = values.ReplicaDeviceMap(_devices) v = [] - index = {} for d, n, init in zip(_devices, ["v", "v/replica"], [1., 2.]): with ops.device(d): v.append(variable_scope.get_variable( name=n, initializer=init, use_resource=True)) - index[d] = v[-1] - tower_local = values.TowerLocalVariable(index, v[0], method) - return v, tower_local + replica_local = values.ReplicaLocalVariable(strategy, device_map, v, method) + return v, replica_local -class TowerLocalVariableTest(test.TestCase): +class ReplicaLocalVariablePropertiesTest(test.TestCase): config = config_pb2.ConfigProto() config.allow_soft_placement = True @@ -778,30 +535,51 @@ class TowerLocalVariableTest(test.TestCase): def testProperties(self): if context.num_gpus() < 1 and context.executing_eagerly(): self.skipTest("A GPU is not available for this test in eager mode.") + v, replica_local = _make_replica_local( + variable_scope.VariableAggregation.SUM) - v, tower_local = _make_tower_local(variable_scope.VariableAggregation.SUM) - - self.assertEquals(v[0].name, tower_local.name) - self.assertEquals(v[0].dtype, tower_local.dtype) - self.assertEquals(v[0].shape, tower_local.shape) - self.assertEquals(variable_scope.VariableAggregation.SUM, - tower_local.aggregation) + self.assertEqual(v[0].name, replica_local.name) + self.assertEqual(v[0].dtype, replica_local.dtype) + self.assertEqual(v[0].shape, replica_local.shape) + self.assertEqual(variable_scope.VariableAggregation.SUM, + replica_local.aggregation) @test_util.run_in_graph_and_eager_modes(config=config) def testVariableOnAnotherDevice(self): v = variable_scope.get_variable( name="v", initializer=[1.], use_resource=True) - index = {"/job:foo/device:CPU:0": v} - tower_local = values.TowerLocalVariable( - index, v, variable_scope.VariableAggregation.MEAN) + device_map = values.ReplicaDeviceMap(("/job:foo/device:CPU:0",)) + replica_local = values.ReplicaLocalVariable( + None, device_map, (v,), variable_scope.VariableAggregation.MEAN) - self.assertEquals(v.name, tower_local.name) - self.assertEquals(v.dtype, tower_local.dtype) - self.assertEquals(v.shape, tower_local.shape) - self.assertEquals(variable_scope.VariableAggregation.MEAN, - tower_local.aggregation) + self.assertEqual(v.name, replica_local.name) + self.assertEqual(v.dtype, replica_local.dtype) + self.assertEqual(v.shape, replica_local.shape) + self.assertEqual(variable_scope.VariableAggregation.MEAN, + replica_local.aggregation) - def _assign_tower_local(self, devices, v, new): + def testTensorConversion(self): + with context.graph_mode(): + _, replica_local = _make_replica_local( + variable_scope.VariableAggregation.SUM) + converted = ops.internal_convert_to_tensor(replica_local, as_ref=False) + self.assertIsInstance(converted, ops.Tensor) + self.assertEqual(converted.dtype, replica_local.dtype) + + converted = ops.internal_convert_to_tensor(replica_local, as_ref=True) + # Resources variable are converted to tensors as well when as_ref is True. + self.assertIsInstance(converted, ops.Tensor) + self.assertEqual(converted.dtype, replica_local.dtype) + + +@combinations.generate(combinations.combine( + distribution=[ + combinations.mirrored_strategy_with_gpu_and_cpu, + combinations.core_mirrored_strategy_with_gpu_and_cpu], + mode=["graph", "eager"])) +class ReplicaLocalVariableTest(test.TestCase, parameterized.TestCase): + + def _assign_replica_local(self, devices, v, new): for d, var, n in zip(devices, v, new): with ops.device(d): self.evaluate(var.assign(n)) @@ -816,86 +594,79 @@ class TowerLocalVariableTest(test.TestCase): save_path, _ = self._save_return_saver(sess, var) return save_path - def _dist_scope(self): - return mirrored_strategy.MirroredStrategy(_devices).scope() - - @test_util.run_in_graph_and_eager_modes(config=config) - def testSaveAndRestoreTowerLocalSumOneGraph(self): - if context.num_gpus() < 1 and context.executing_eagerly(): - self.skipTest("A GPU is not available for this test in eager mode.") - - with self.cached_session(config=self.config) as sess: - v, tower_local = _make_tower_local(variable_scope.VariableAggregation.SUM) + def testSaveAndRestoreReplicaLocalSumOneGraph(self, distribution): + with self.cached_session() as sess: + v, replica_local = _make_replica_local( + variable_scope.VariableAggregation.SUM, distribution) # Overwrite the initial values. - self._assign_tower_local(_devices, v, [3., 4.]) + self._assign_replica_local(_devices, v, [3., 4.]) - with self._dist_scope(): + with distribution.scope(): # Saves the current value of v[0] + v[1], 7. - save_path, saver = self._save_return_saver(sess, tower_local) + save_path, saver = self._save_return_saver(sess, replica_local) # Change the values between save and restore. - self._assign_tower_local(_devices, v, [5., 6.]) + self._assign_replica_local(_devices, v, [5., 6.]) # Restores the saved value of 7. which gets divided equally # between the variables. saver.restore(sess, save_path) self.assertEqual([3.5, 3.5], self.evaluate([v[0], v[1]])) - @test_util.run_in_graph_and_eager_modes(config=config) - def testSaveAndRestoreTowerLocalMeanOneGraph(self): + def testSaveAndRestoreReplicaLocalMeanOneGraph(self, distribution): if context.num_gpus() < 1 and context.executing_eagerly(): self.skipTest("A GPU is not available for this test in eager mode.") - with self.cached_session(config=self.config) as sess: - v, tower_local = _make_tower_local( - variable_scope.VariableAggregation.MEAN) + with self.cached_session() as sess: + v, replica_local = _make_replica_local( + variable_scope.VariableAggregation.MEAN, distribution) # Overwrite the initial values. - self._assign_tower_local(_devices, v, [3., 4.]) + self._assign_replica_local(_devices, v, [3., 4.]) - with self._dist_scope(): + with distribution.scope(): # Saves the current value of (v[0] + v[1])/2, 3.5. - save_path, saver = self._save_return_saver(sess, tower_local) + save_path, saver = self._save_return_saver(sess, replica_local) # Change the values between save and restore. - self._assign_tower_local(_devices, v, [5., 6.]) + self._assign_replica_local(_devices, v, [5., 6.]) # Restores the saved value of 3.5 to both variables. saver.restore(sess, save_path) self.assertEqual([3.5, 3.5], self.evaluate([v[0], v[1]])) - def _save_tower_local_mean(self): + def _save_replica_local_mean(self, distribution): """Save variables with mirroring, returns save_path.""" with self.session(graph=ops.Graph()) as sess: - v, tower_local = _make_tower_local( - variable_scope.VariableAggregation.MEAN) + v, replica_local = _make_replica_local( + variable_scope.VariableAggregation.MEAN, distribution) # Overwrite the initial values. - self._assign_tower_local(_devices, v, [3., 4.]) + self._assign_replica_local(_devices, v, [3., 4.]) - with self._dist_scope(): + with distribution.scope(): # Saves the current value of (v[0] + v[1])/2, 3.5 - save_path = self._save(sess, tower_local) + save_path = self._save(sess, replica_local) # Change the values between save and restore. - self._assign_tower_local(_devices, v, [5., 6.]) + self._assign_replica_local(_devices, v, [5., 6.]) return save_path - def _save_tower_local_sum(self): + def _save_replica_local_sum(self, distribution): """Save variables with mirroring, returns save_path.""" with self.session(graph=ops.Graph()) as sess: - v, tower_local = _make_tower_local("sum") + v, replica_local = _make_replica_local("sum", distribution) # Overwrite the initial values. - self._assign_tower_local(_devices, v, [1.5, 2.]) + self._assign_replica_local(_devices, v, [1.5, 2.]) - with self._dist_scope(): + with distribution.scope(): # Saves the current value of v[0] + v[1], 3.5 - save_path = self._save(sess, tower_local) + save_path = self._save(sess, replica_local) # Change the values between save and restore. - self._assign_tower_local(_devices, v, [5., 6.]) + self._assign_replica_local(_devices, v, [5., 6.]) return save_path def _save_normal(self): @@ -928,94 +699,59 @@ class TowerLocalVariableTest(test.TestCase): saver.restore(sess, save_path) self.assertEqual(3.5, self.evaluate(var)) - def _restore_tower_local_mean(self, save_path): + def _restore_replica_local_mean(self, save_path, distribution): """Restore to variables with mirroring in a fresh graph.""" with self.session(graph=ops.Graph()) as sess: - v, tower_local = _make_tower_local( - variable_scope.VariableAggregation.MEAN) + v, replica_local = _make_replica_local( + variable_scope.VariableAggregation.MEAN, distribution) # Overwrite the initial values. - self._assign_tower_local(_devices, v, [7., 8.]) + self._assign_replica_local(_devices, v, [7., 8.]) - with self._dist_scope(): + with distribution.scope(): # Restores the saved value of 3.5 to both variables. - saver = saver_lib.Saver(var_list=[tower_local]) + saver = saver_lib.Saver(var_list=[replica_local]) saver.restore(sess, save_path) self.assertEqual([3.5, 3.5], self.evaluate([v[0], v[1]])) - def _restore_tower_local_sum(self, save_path): + def _restore_replica_local_sum(self, save_path, distribution): """Restore to variables with mirroring in a fresh graph.""" with self.session(graph=ops.Graph()) as sess: - v, tower_local = _make_tower_local(variable_scope.VariableAggregation.SUM) + v, replica_local = _make_replica_local( + variable_scope.VariableAggregation.SUM, distribution) # Overwrite the initial values. - self._assign_tower_local(_devices, v, [7., 8.]) + self._assign_replica_local(_devices, v, [7., 8.]) - with self._dist_scope(): + with distribution.scope(): # Restores the saved value of 3.5 to both variables. - saver = saver_lib.Saver(var_list=[tower_local]) + saver = saver_lib.Saver(var_list=[replica_local]) saver.restore(sess, save_path) self.assertEqual([1.75, 1.75], self.evaluate([v[0], v[1]])) - @test_util.run_in_graph_and_eager_modes(config=config) - def testSaveTowerLocalRestoreTowerLocalMean(self): - if context.num_gpus() < 1 and context.executing_eagerly(): - self.skipTest("A GPU is not available for this test in eager mode.") - - save_path = self._save_tower_local_mean() - self._restore_tower_local_mean(save_path) + def testSaveReplicaLocalRestoreReplicaLocalMean(self, distribution): + save_path = self._save_replica_local_mean(distribution) + self._restore_replica_local_mean(save_path, distribution) - @test_util.run_in_graph_and_eager_modes(config=config) - def testSaveTowerLocalRestoreTowerLocalSum(self): - if context.num_gpus() < 1 and context.executing_eagerly(): - self.skipTest("A GPU is not available for this test in eager mode.") + def testSaveReplicaLocalRestoreReplicaLocalSum(self, distribution): + save_path = self._save_replica_local_sum(distribution) + self._restore_replica_local_sum(save_path, distribution) - save_path = self._save_tower_local_sum() - self._restore_tower_local_sum(save_path) - - @test_util.run_in_graph_and_eager_modes(config=config) - def testSaveTowerLocalMeanRestoreNormal(self): - if context.num_gpus() < 1 and context.executing_eagerly(): - self.skipTest("A GPU is not available for this test in eager mode.") - - save_path = self._save_tower_local_mean() + def testSaveReplicaLocalMeanRestoreNormal(self, distribution): + save_path = self._save_replica_local_mean(distribution) self._restore_normal(save_path) - @test_util.run_in_graph_and_eager_modes(config=config) - def testSaveTowerLocalSumRestoreNormal(self): - if context.num_gpus() < 1 and context.executing_eagerly(): - self.skipTest("A GPU is not available for this test in eager mode.") - - save_path = self._save_tower_local_sum() + def testSaveReplicaLocalSumRestoreNormal(self, distribution): + save_path = self._save_replica_local_sum(distribution) self._restore_normal(save_path) - @test_util.run_in_graph_and_eager_modes(config=config) - def testSaveNormalRestoreTowerLocalMean(self): - if context.num_gpus() < 1 and context.executing_eagerly(): - self.skipTest("A GPU is not available for this test in eager mode.") - + def testSaveNormalRestoreReplicaLocalMean(self, distribution): save_path = self._save_normal() - self._restore_tower_local_mean(save_path) - - @test_util.run_in_graph_and_eager_modes(config=config) - def testSaveNormalRestoreTowerLocalSum(self): - if context.num_gpus() < 1 and context.executing_eagerly(): - self.skipTest("A GPU is not available for this test in eager mode.") + self._restore_replica_local_mean(save_path, distribution) + def testSaveNormalRestoreReplicaLocalSum(self, distribution): save_path = self._save_normal() - self._restore_tower_local_sum(save_path) - - def testTensorConversion(self): - with context.graph_mode(): - _, tower_local = _make_tower_local(variable_scope.VariableAggregation.SUM) - converted = ops.internal_convert_to_tensor(tower_local, as_ref=False) - self.assertIsInstance(converted, ops.Tensor) - self.assertEqual(converted.dtype, tower_local.dtype) - - converted = ops.internal_convert_to_tensor(tower_local, as_ref=True) - # Resources variable are converted to tensors as well when as_ref is True. - self.assertIsInstance(converted, ops.Tensor) - self.assertEqual(converted.dtype, tower_local.dtype) + self._restore_replica_local_sum(save_path, distribution) if __name__ == "__main__": diff --git a/tensorflow/contrib/distribute/python/warm_starting_util_test.py b/tensorflow/contrib/distribute/python/warm_starting_util_test.py index 5d57d144c1c16a08280970ecd89eb54f7cf1ffd4..b0bcf9b17456c938204a4892451928daf90b6743 100644 --- a/tensorflow/contrib/distribute/python/warm_starting_util_test.py +++ b/tensorflow/contrib/distribute/python/warm_starting_util_test.py @@ -44,7 +44,9 @@ class WarmStartingUtilWithDistributionStrategyTest( distribution=[combinations.default_strategy, combinations.one_device_strategy, combinations.mirrored_strategy_with_gpu_and_cpu, - combinations.mirrored_strategy_with_two_gpus], + combinations.mirrored_strategy_with_two_gpus, + combinations.core_mirrored_strategy_with_gpu_and_cpu, + combinations.core_mirrored_strategy_with_two_gpus], save_with_distribution=[True, False], restore_with_distribution=[True, False], mode=["graph"])) diff --git a/tensorflow/contrib/distributions/BUILD b/tensorflow/contrib/distributions/BUILD index 60f6b90edcb71f04bca29b90744db201e83cd545..3079175015a9aee1625404902070df8f13b2089c 100644 --- a/tensorflow/contrib/distributions/BUILD +++ b/tensorflow/contrib/distributions/BUILD @@ -72,7 +72,6 @@ py_library( "//tensorflow/python:nn", "//tensorflow/python:nn_ops", "//tensorflow/python:random_ops", - "//tensorflow/python:spectral_ops", "//tensorflow/python:state_ops", "//tensorflow/python:tensor_util", "//tensorflow/python:util", @@ -80,6 +79,7 @@ py_library( "//tensorflow/python:variables", "//tensorflow/python/ops/distributions", "//tensorflow/python/ops/linalg", + "//tensorflow/python/ops/signal", "//third_party/py/numpy", "@six_archive//:six", ], diff --git a/tensorflow/contrib/distributions/python/kernel_tests/distribution_test.py b/tensorflow/contrib/distributions/python/kernel_tests/distribution_test.py index 9b9b3ce2dd9d42286d2d9657d5f00de8445261f0..99cb105d66885fd5cf8cb6a3f87e2fe82a5bf4d0 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/distribution_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/distribution_test.py @@ -250,13 +250,22 @@ class DistributionTest(test.TestCase): mvn_dynamic = tfd.MultivariateNormalDiag( loc=array_ops.placeholder(shape=[None, 3], dtype=dtypes.float32), name="MVN2") - self.assertEqual( - ("tfp.distributions.MultivariateNormalDiag(" - "\"MVN2/\", " - "batch_shape=(?,), " # Partially known. - "event_shape=(3,), " - "dtype=float32)"), - str(mvn_dynamic)) + if mvn_dynamic.batch_shape._v2_behavior: + self.assertEqual( + ("tfp.distributions.MultivariateNormalDiag(" + "\"MVN2/\", " + "batch_shape=(None,), " # Partially known. + "event_shape=(3,), " + "dtype=float32)"), + str(mvn_dynamic)) + else: + self.assertEqual( + ("tfp.distributions.MultivariateNormalDiag(" + "\"MVN2/\", " + "batch_shape=(?,), " # Partially known. + "event_shape=(3,), " + "dtype=float32)"), + str(mvn_dynamic)) def testReprWorksCorrectlyScalar(self): normal = tfd.Normal(loc=np.float16(0), scale=np.float16(1)) @@ -300,13 +309,22 @@ class DistributionTest(test.TestCase): mvn_dynamic = tfd.MultivariateNormalDiag( loc=array_ops.placeholder(shape=[None, 3], dtype=dtypes.float32), name="MVN2") - self.assertEqual( - (""), - repr(mvn_dynamic)) + if mvn_dynamic.batch_shape._v2_behavior: + self.assertEqual( + (""), + repr(mvn_dynamic)) + else: + self.assertEqual( + (""), + repr(mvn_dynamic)) if __name__ == "__main__": diff --git a/tensorflow/contrib/distributions/python/kernel_tests/normal_conjugate_posteriors_test.py b/tensorflow/contrib/distributions/python/kernel_tests/normal_conjugate_posteriors_test.py index 29eeaf43c5185ce5519d4a1211f66e99ce61c6ab..ab3c07172a68255f4e387e071ac2f8341e93b90c 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/normal_conjugate_posteriors_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/normal_conjugate_posteriors_test.py @@ -82,7 +82,7 @@ class NormalTest(test.TestCase): x = constant_op.constant( [[-2.5, 2.5, 4.0, 0.0, -1.0, 2.0], [2.5, -2.5, -4.0, 0.0, 1.0, -2.0]], dtype=dtypes.float32) - s = math_ops.reduce_sum(x, reduction_indices=[1]) + s = math_ops.reduce_sum(x, axis=[1]) x = array_ops.transpose(x) # Reshape to shape (6, 2) n = constant_op.constant([6] * 2) prior = distributions.Normal(loc=mu0, scale=sigma0) diff --git a/tensorflow/contrib/distributions/python/kernel_tests/wishart_test.py b/tensorflow/contrib/distributions/python/kernel_tests/wishart_test.py index a60056c444a3fe7262939c5b3c73673f9a7c1469..cdee30bbc42e661952a9c757d7a30ebcd393f794 100644 --- a/tensorflow/contrib/distributions/python/kernel_tests/wishart_test.py +++ b/tensorflow/contrib/distributions/python/kernel_tests/wishart_test.py @@ -147,14 +147,13 @@ class WishartCholeskyTest(test.TestCase): x = chol_w.sample(10000, seed=42) self.assertAllEqual((10000, 3, 3), x.get_shape()) - moment1_estimate = math_ops.reduce_mean(x, reduction_indices=[0]).eval() + moment1_estimate = math_ops.reduce_mean(x, axis=[0]).eval() self.assertAllClose(chol_w.mean().eval(), moment1_estimate, rtol=0.05) # The Variance estimate uses the squares rather than outer-products # because Wishart.Variance is the diagonal of the Wishart covariance # matrix. - variance_estimate = (math_ops.reduce_mean( - math_ops.square(x), reduction_indices=[0]) - + variance_estimate = (math_ops.reduce_mean(math_ops.square(x), axis=[0]) - math_ops.square(moment1_estimate)).eval() self.assertAllClose( chol_w.variance().eval(), variance_estimate, rtol=0.05) diff --git a/tensorflow/contrib/distributions/python/ops/batch_reshape.py b/tensorflow/contrib/distributions/python/ops/batch_reshape.py index 612376efb7f43b0dfcd3ffeb5437f2a419f66f4d..d450379088813caeac6f3dca72fae99c5f886b5a 100644 --- a/tensorflow/contrib/distributions/python/ops/batch_reshape.py +++ b/tensorflow/contrib/distributions/python/ops/batch_reshape.py @@ -429,5 +429,6 @@ def validate_init_args_statically(distribution, batch_shape): if batch_shape_static.dims is not None: if any( - dim.value is not None and dim.value < 1 for dim in batch_shape_static): + dim.value is not None and + dim.value < 1 for dim in batch_shape_static.dims): raise ValueError("`batch_shape` elements must be >=-1.") diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/cholesky_outer_product.py b/tensorflow/contrib/distributions/python/ops/bijectors/cholesky_outer_product.py index 3e1e4fc82971b71792d193ea8518dd402e4a4d9d..2358ef5976b2f21c77130c71d5214a463d17bf0e 100644 --- a/tensorflow/contrib/distributions/python/ops/bijectors/cholesky_outer_product.py +++ b/tensorflow/contrib/distributions/python/ops/bijectors/cholesky_outer_product.py @@ -168,11 +168,11 @@ class CholeskyOuterProduct(bijector.Bijector): [is_matrix, is_square, is_positive_definite], x) # Create a vector equal to: [p, p-1, ..., 2, 1]. - if x.get_shape().ndims is None or x.get_shape()[-1].value is None: + if x.get_shape().ndims is None or x.get_shape().dims[-1].value is None: p_int = array_ops.shape(x)[-1] p_float = math_ops.cast(p_int, dtype=x.dtype) else: - p_int = x.get_shape()[-1].value + p_int = x.get_shape().dims[-1].value p_float = np.array(p_int, dtype=x.dtype.as_numpy_dtype) exponents = math_ops.linspace(p_float, 1., p_int) diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/fill_triangular.py b/tensorflow/contrib/distributions/python/ops/bijectors/fill_triangular.py index 31a9ca27e519bc312813668bf621a875838f12a0..7ae98878986eb10570b5e93a4a57d6bad6b38c0c 100644 --- a/tensorflow/contrib/distributions/python/ops/bijectors/fill_triangular.py +++ b/tensorflow/contrib/distributions/python/ops/bijectors/fill_triangular.py @@ -21,6 +21,7 @@ from __future__ import print_function import numpy as np 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 check_ops from tensorflow.python.ops import math_ops @@ -104,7 +105,8 @@ class FillTriangular(bijector.Bijector): return array_ops.zeros_like(y[..., 0, 0]) def _forward_event_shape(self, input_shape): - batch_shape, d = input_shape[:-1], input_shape[-1].value + batch_shape, d = (input_shape[:-1], + tensor_shape.dimension_value(input_shape[-1])) if d is None: n = None else: @@ -113,8 +115,8 @@ class FillTriangular(bijector.Bijector): def _inverse_event_shape(self, output_shape): batch_shape, n1, n2 = (output_shape[:-2], - output_shape[-2].value, - output_shape[-1].value) + tensor_shape.dimension_value(output_shape[-2]), + tensor_shape.dimension_value(output_shape[-1])) if n1 is None or n2 is None: m = None elif n1 != n2: diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/masked_autoregressive.py b/tensorflow/contrib/distributions/python/ops/bijectors/masked_autoregressive.py index 3b3d8ee6f2dc595983fd5e283d0435e8a227f2ba..c30de1f989a7b83fba1f69a83b96b8f45dea02c6 100644 --- a/tensorflow/contrib/distributions/python/ops/bijectors/masked_autoregressive.py +++ b/tensorflow/contrib/distributions/python/ops/bijectors/masked_autoregressive.py @@ -23,6 +23,7 @@ import numpy as np from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops +from tensorflow.python.framework import tensor_shape from tensorflow.python.layers import core as layers from tensorflow.python.ops import array_ops from tensorflow.python.ops import clip_ops @@ -237,7 +238,8 @@ class MaskedAutoregressiveFlow(bijector.Bijector): def _forward(self, x): if self._unroll_loop: - event_size = x.shape.with_rank_at_least(1)[-1].value + event_size = tensor_shape.dimension_value( + x.shape.with_rank_at_least(1)[-1]) if event_size is None: raise ValueError( "The final dimension of `x` must be known at graph construction " @@ -260,7 +262,8 @@ class MaskedAutoregressiveFlow(bijector.Bijector): # the graph compiler of the maximum number of steps. If not, # static_event_size will be None, and the maximum_iterations argument will # have no effect. - static_event_size = x.shape.with_rank_at_least(1)[-1].value + static_event_size = tensor_shape.dimension_value( + x.shape.with_rank_at_least(1)[-1]) y0 = array_ops.zeros_like(x, name="y0") # call the template once to ensure creation _ = self._shift_and_log_scale_fn(y0) @@ -405,7 +408,8 @@ def masked_dense(inputs, Conference on Machine Learning_, 2015. https://arxiv.org/abs/1502.03509 """ # TODO(b/67594795): Better support of dynamic shape. - input_depth = inputs.shape.with_rank_at_least(1)[-1].value + input_depth = tensor_shape.dimension_value( + inputs.shape.with_rank_at_least(1)[-1]) if input_depth is None: raise NotImplementedError( "Rightmost dimension must be known prior to graph execution.") @@ -520,7 +524,8 @@ def masked_autoregressive_default_template( def _fn(x): """MADE parameterized via `masked_autoregressive_default_template`.""" # TODO(b/67594795): Better support of dynamic shape. - input_depth = x.shape.with_rank_at_least(1)[-1].value + input_depth = tensor_shape.dimension_value( + x.shape.with_rank_at_least(1)[-1]) if input_depth is None: raise NotImplementedError( "Rightmost dimension must be known prior to graph execution.") diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/real_nvp.py b/tensorflow/contrib/distributions/python/ops/bijectors/real_nvp.py index 0bcb08cdea7142b82af3116245306a11773ef93c..17e9b8dec9f009415a9a26c3b043afacc2c4ec72 100644 --- a/tensorflow/contrib/distributions/python/ops/bijectors/real_nvp.py +++ b/tensorflow/contrib/distributions/python/ops/bijectors/real_nvp.py @@ -20,6 +20,7 @@ from __future__ import print_function from tensorflow.python.framework import constant_op from tensorflow.python.framework import ops +from tensorflow.python.framework import tensor_shape from tensorflow.python.layers import core as layers from tensorflow.python.ops import array_ops from tensorflow.python.ops import math_ops @@ -96,16 +97,18 @@ class RealNVP(bijector.Bijector): # A common choice for a normalizing flow is to use a Gaussian for the base # distribution. (However, any continuous distribution would work.) E.g., + num_dims = 3 + num_samples = 1 nvp = tfd.TransformedDistribution( - distribution=tfd.MultivariateNormalDiag(loc=[0., 0., 0.])), + distribution=tfd.MultivariateNormalDiag(loc=np.zeros(num_dims)), bijector=tfb.RealNVP( num_masked=2, shift_and_log_scale_fn=tfb.real_nvp_default_template( hidden_layers=[512, 512]))) - x = nvp.sample() + x = nvp.sample(num_samples) nvp.log_prob(x) - nvp.log_prob(0.) + nvp.log_prob(np.zeros([num_samples, num_dims])) ``` For more examples, see [Jang (2018)][3]. @@ -183,7 +186,8 @@ class RealNVP(bijector.Bijector): def _cache_input_depth(self, x): if self._input_depth is None: - self._input_depth = x.shape.with_rank_at_least(1)[-1].value + self._input_depth = tensor_shape.dimension_value( + x.shape.with_rank_at_least(1)[-1]) if self._input_depth is None: raise NotImplementedError( "Rightmost dimension must be known prior to graph execution.") diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/reshape.py b/tensorflow/contrib/distributions/python/ops/bijectors/reshape.py index 71ac29038fc12e7d046df8624c6e3e5bb97d3d8f..ec203e171730a1ef6de6b72c6d96c52d4010d7e6 100644 --- a/tensorflow/contrib/distributions/python/ops/bijectors/reshape.py +++ b/tensorflow/contrib/distributions/python/ops/bijectors/reshape.py @@ -23,6 +23,7 @@ import numpy as np from tensorflow.python.framework import constant_op 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.ops import array_ops from tensorflow.python.ops import check_ops @@ -46,7 +47,7 @@ __all__ = [ "instead of `tf.contrib.distributions`.", warn_once=True) def _static_ndims_from_shape(shape): - return shape.shape.with_rank_at_least(1)[0].value + return tensor_shape.dimension_value(shape.shape.with_rank_at_least(1)[0]) @deprecation.deprecated( diff --git a/tensorflow/contrib/distributions/python/ops/bijectors/softmax_centered.py b/tensorflow/contrib/distributions/python/ops/bijectors/softmax_centered.py index 20ee0d340833d5c5275e2ab52a89dcdf7198add1..74765f19e584c5de07c6aee4a36ec4e85438f862 100644 --- a/tensorflow/contrib/distributions/python/ops/bijectors/softmax_centered.py +++ b/tensorflow/contrib/distributions/python/ops/bijectors/softmax_centered.py @@ -110,7 +110,7 @@ class SoftmaxCentered(bijector.Bijector): # Set shape hints. if x.shape.ndims is not None: - shape = x.shape[:-1].concatenate(x.shape[-1] + 1) + shape = x.shape[:-1].concatenate(x.shape.dims[-1] + 1) y.shape.assert_is_compatible_with(shape) y.set_shape(shape) @@ -135,7 +135,7 @@ class SoftmaxCentered(bijector.Bijector): # Set shape hints. if y.shape.ndims is not None: - shape = y.shape[:-1].concatenate(y.shape[-1] - 1) + shape = y.shape[:-1].concatenate(y.shape.dims[-1] - 1) x.shape.assert_is_compatible_with(shape) x.set_shape(shape) @@ -168,7 +168,7 @@ class SoftmaxCentered(bijector.Bijector): # log_normalization = 1 + reduce_sum(exp(logits)) # -log_normalization + reduce_sum(logits - log_normalization) log_normalization = nn_ops.softplus( - math_ops.reduce_logsumexp(x, axis=-1, keep_dims=True)) + math_ops.reduce_logsumexp(x, axis=-1, keepdims=True)) return array_ops.squeeze( (-log_normalization + math_ops.reduce_sum( x - log_normalization, axis=-1, keepdims=True)), axis=-1) diff --git a/tensorflow/contrib/distributions/python/ops/distribution_util.py b/tensorflow/contrib/distributions/python/ops/distribution_util.py index b4ad33cf6dbf073419a27f378c8eefdba97c5af7..1415f85e5cb5598e99c4d6b8e6c6a2d254503db0 100644 --- a/tensorflow/contrib/distributions/python/ops/distribution_util.py +++ b/tensorflow/contrib/distributions/python/ops/distribution_util.py @@ -21,6 +21,7 @@ from __future__ import print_function from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.framework import smart_cond +from tensorflow.python.framework import tensor_shape from tensorflow.python.framework import tensor_util from tensorflow.python.ops import array_ops from tensorflow.python.ops import check_ops @@ -315,7 +316,7 @@ def shapes_from_loc_and_scale(loc, scale, name="shapes_from_loc_and_scale"): # Static check that event shapes match. if loc is not None: - loc_event_size = loc.get_shape()[-1].value + loc_event_size = tensor_shape.dimension_value(loc.get_shape()[-1]) if loc_event_size is not None and event_size_const is not None: if loc_event_size != 1 and loc_event_size != event_size_const: raise ValueError( diff --git a/tensorflow/contrib/distributions/python/ops/independent.py b/tensorflow/contrib/distributions/python/ops/independent.py index e1cfff3c66a2bcbc98af8a257dbdea2d916270e2..cf15deebb78b6c92865c34f61d806bc9e9ab3ee1 100644 --- a/tensorflow/contrib/distributions/python/ops/independent.py +++ b/tensorflow/contrib/distributions/python/ops/independent.py @@ -166,8 +166,10 @@ class Independent(distribution_lib.Distribution): def _batch_shape_tensor(self): with ops.control_dependencies(self._runtime_assertions): batch_shape = self.distribution.batch_shape_tensor() - batch_ndims = (batch_shape.shape[0].value - if batch_shape.shape.with_rank_at_least(1)[0].value + dim0 = tensor_shape.dimension_value( + batch_shape.shape.with_rank_at_least(1)[0]) + batch_ndims = (dim0 + if dim0 is not None else array_ops.shape(batch_shape)[0]) return batch_shape[:batch_ndims - self.reinterpreted_batch_ndims] @@ -182,8 +184,10 @@ class Independent(distribution_lib.Distribution): def _event_shape_tensor(self): with ops.control_dependencies(self._runtime_assertions): batch_shape = self.distribution.batch_shape_tensor() - batch_ndims = (batch_shape.shape[0].value - if batch_shape.shape.with_rank_at_least(1)[0].value + dim0 = tensor_shape.dimension_value( + batch_shape.shape.with_rank_at_least(1)[0]) + batch_ndims = (dim0 + if dim0 is not None else array_ops.shape(batch_shape)[0]) return array_ops.concat([ batch_shape[batch_ndims - self.reinterpreted_batch_ndims:], @@ -239,9 +243,11 @@ class Independent(distribution_lib.Distribution): static_reinterpreted_batch_ndims, batch_ndims)) elif validate_args: batch_shape = distribution.batch_shape_tensor() + dim0 = tensor_shape.dimension_value( + batch_shape.shape.with_rank_at_least(1)[0]) batch_ndims = ( - batch_shape.shape[0].value - if batch_shape.shape.with_rank_at_least(1)[0].value is not None + dim0 + if dim0 is not None else array_ops.shape(batch_shape)[0]) assertions.append(check_ops.assert_less_equal( reinterpreted_batch_ndims, batch_ndims, diff --git a/tensorflow/contrib/distributions/python/ops/mixture_same_family.py b/tensorflow/contrib/distributions/python/ops/mixture_same_family.py index f4d394ff29f072907a019afb52bd8dc5d244e955..f34317f5abfed1c71b516c5fde42baca614d7f9b 100644 --- a/tensorflow/contrib/distributions/python/ops/mixture_same_family.py +++ b/tensorflow/contrib/distributions/python/ops/mixture_same_family.py @@ -22,6 +22,7 @@ import numpy as np from tensorflow.contrib.distributions.python.ops import distribution_util as distribution_utils 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 @@ -147,8 +148,9 @@ class MixtureSameFamily(distribution.Distribution): self._runtime_assertions = [] s = components_distribution.event_shape_tensor() - self._event_ndims = (s.shape[0].value - if s.shape.with_rank_at_least(1)[0].value is not None + s_dim0 = tensor_shape.dimension_value(s.shape[0]) + self._event_ndims = (s_dim0 + if s_dim0 is not None else array_ops.shape(s)[0]) if not mixture_distribution.dtype.is_integer: @@ -186,8 +188,10 @@ class MixtureSameFamily(distribution.Distribution): "`mixture_distribution.batch_shape` is not " "compatible with `components_distribution.batch_shape`"))] - km = mixture_distribution.logits.shape.with_rank_at_least(1)[-1].value - kc = components_distribution.batch_shape.with_rank_at_least(1)[-1].value + km = tensor_shape.dimension_value( + mixture_distribution.logits.shape.with_rank_at_least(1)[-1]) + kc = tensor_shape.dimension_value( + components_distribution.batch_shape.with_rank_at_least(1)[-1]) if km is not None and kc is not None and km != kc: raise ValueError("`mixture_distribution components` ({}) does not " "equal `components_distribution.batch_shape[-1]` " diff --git a/tensorflow/contrib/distributions/python/ops/sample_stats.py b/tensorflow/contrib/distributions/python/ops/sample_stats.py index aa680a92be64cf0f099acd335369f2a1610c5953..19e99e03803e7f4cdfdb023feb04daaba68eceed 100644 --- a/tensorflow/contrib/distributions/python/ops/sample_stats.py +++ b/tensorflow/contrib/distributions/python/ops/sample_stats.py @@ -29,8 +29,8 @@ from tensorflow.python.ops import clip_ops from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import nn_ops -from tensorflow.python.ops import spectral_ops from tensorflow.python.ops.distributions import util +from tensorflow.python.ops.signal import fft_ops __all__ = [ "auto_correlation", @@ -157,11 +157,11 @@ def auto_correlation( dtype.real_dtype.as_numpy_dtype(0.)) # Autocorrelation is IFFT of power-spectral density (up to some scaling). - fft_x_rotated_pad = spectral_ops.fft(x_rotated_pad) + fft_x_rotated_pad = fft_ops.fft(x_rotated_pad) spectral_density = fft_x_rotated_pad * math_ops.conj(fft_x_rotated_pad) # shifted_product is R[m] from above detailed explanation. # It is the inner product sum_n X[n] * Conj(X[n - m]). - shifted_product = spectral_ops.ifft(spectral_density) + shifted_product = fft_ops.ifft(spectral_density) # Cast back to real-valued if x was real to begin with. shifted_product = math_ops.cast(shifted_product, dtype) @@ -300,7 +300,7 @@ def percentile(x, raise ValueError("Argument 'interpolation' must be in %s. Found %s" % (allowed_interpolations, interpolation)) - with ops.name_scope(name, [x, q]): + with ops.name_scope(name, values=[x, q]): x = ops.convert_to_tensor(x, name="x") # Double is needed here and below, else we get the wrong index if the array # is huge along axis. diff --git a/tensorflow/contrib/distributions/python/ops/vector_diffeomixture.py b/tensorflow/contrib/distributions/python/ops/vector_diffeomixture.py index a3d178357b79b9d0d15c738603d5019321eff112..a648d61ac8dd5c1d368cf41505b85827dfeb63e1 100644 --- a/tensorflow/contrib/distributions/python/ops/vector_diffeomixture.py +++ b/tensorflow/contrib/distributions/python/ops/vector_diffeomixture.py @@ -183,7 +183,7 @@ def quadrature_scheme_softmaxnormal_quantiles( def _get_final_shape(qs): """Helper to build `TensorShape`.""" bs = dist.batch_shape.with_rank_at_least(1) - num_components = bs[-1].value + num_components = tensor_shape.dimension_value(bs[-1]) if num_components is not None: num_components += 1 tail = tensor_shape.TensorShape([num_components, qs]) @@ -791,7 +791,7 @@ class VectorDiffeomixture(distribution_lib.Distribution): def _expand_mix_distribution_probs(self): p = self.mixture_distribution.probs # [B, deg] - deg = p.shape.with_rank_at_least(1)[-1].value + deg = tensor_shape.dimension_value(p.shape.with_rank_at_least(1)[-1]) if deg is None: deg = array_ops.shape(p)[-1] event_ndims = self.event_shape.ndims @@ -831,10 +831,12 @@ def maybe_check_quadrature_param(param, name, validate_args): # TODO(jvdillon): Remove once we support k-mixtures. if param.shape.with_rank_at_least(1)[-1] is not None: - if param.shape[-1].value != 1: + if tensor_shape.dimension_value(param.shape[-1]) != 1: raise NotImplementedError("Currently only bimixtures are supported; " "{}.shape[-1]={} is not 1.".format( - name, param.shape[-1].value)) + name, + tensor_shape.dimension_value( + param.shape[-1]))) elif validate_args: assertions.append(check_ops.assert_equal( array_ops.shape(param)[-1], 1, @@ -905,7 +907,7 @@ def interpolate_loc(grid, loc): if len(loc) != 2: raise NotImplementedError("Currently only bimixtures are supported; " "len(scale)={} is not 2.".format(len(loc))) - deg = grid.shape.with_rank_at_least(1)[-1].value + deg = tensor_shape.dimension_value(grid.shape.with_rank_at_least(1)[-1]) if deg is None: raise ValueError("Num quadrature grid points must be known prior " "to graph execution.") @@ -939,7 +941,7 @@ def interpolate_scale(grid, scale): if len(scale) != 2: raise NotImplementedError("Currently only bimixtures are supported; " "len(scale)={} is not 2.".format(len(scale))) - deg = grid.shape.with_rank_at_least(1)[-1].value + deg = tensor_shape.dimension_value(grid.shape.with_rank_at_least(1)[-1]) if deg is None: raise ValueError("Num quadrature grid points must be known prior " "to graph execution.") diff --git a/tensorflow/contrib/distributions/python/ops/wishart.py b/tensorflow/contrib/distributions/python/ops/wishart.py index ee2fc58864d4ac528ebae3d681d2e4922fb60771..2d83f0c13f14a8e5d1eee4fa1436bd05991e934e 100644 --- a/tensorflow/contrib/distributions/python/ops/wishart.py +++ b/tensorflow/contrib/distributions/python/ops/wishart.py @@ -136,13 +136,13 @@ class _WishartLinearOperator(distribution.Distribution): contrib_tensor_util.assert_same_float_dtype( (self._df, self._scale_operator)) if (self._scale_operator.shape.ndims is None or - self._scale_operator.shape[-1].value is None): + self._scale_operator.shape.dims[-1].value is None): self._dimension = math_ops.cast( self._scale_operator.domain_dimension_tensor(), dtype=self._scale_operator.dtype, name="dimension") else: self._dimension = ops.convert_to_tensor( - self._scale_operator.shape[-1].value, + self._scale_operator.shape.dims[-1].value, dtype=self._scale_operator.dtype, name="dimension") df_val = tensor_util.constant_value(self._df) dim_val = tensor_util.constant_value(self._dimension) diff --git a/tensorflow/contrib/eager/python/BUILD b/tensorflow/contrib/eager/python/BUILD index 77052a75a70bec1162feb2b126d247924b3a2e36..8966a9befcd3db4a3f397b319e80f37f84ad236b 100644 --- a/tensorflow/contrib/eager/python/BUILD +++ b/tensorflow/contrib/eager/python/BUILD @@ -15,7 +15,6 @@ py_library( ":metrics", ":network", ":parameter_server", - ":remote", ":saver", "//tensorflow/python:framework_ops", "//tensorflow/python:framework_test_lib", @@ -31,6 +30,7 @@ py_library( "//tensorflow/python/eager:def_function", "//tensorflow/python/eager:execution_callbacks", "//tensorflow/python/eager:function", + "//tensorflow/python/eager:remote", ], ) @@ -238,24 +238,12 @@ py_test( ], ) -py_library( - name = "remote", - srcs = ["remote.py"], - srcs_version = "PY2AND3", - visibility = ["//tensorflow:internal"], - deps = [ - "//tensorflow/core:protos_all_py", - "//tensorflow/python:platform", - "//tensorflow/python/eager:context", - ], -) - cuda_py_test( name = "remote_test", srcs = ["remote_test.py"], additional_deps = [ ":parameter_server", - ":remote", + "//tensorflow/python/eager:remote", "//tensorflow/contrib/eager/python:tfe", "//tensorflow/python:array_ops", "//tensorflow/python:client", diff --git a/tensorflow/contrib/eager/python/datasets.py b/tensorflow/contrib/eager/python/datasets.py index 3aed121233be1268531495a2fa83fd323412e1fd..34614b86a75b93ab93cf844c645c211b1329c6d5 100644 --- a/tensorflow/contrib/eager/python/datasets.py +++ b/tensorflow/contrib/eager/python/datasets.py @@ -52,12 +52,6 @@ class Iterator(iterator_ops.EagerIterator): TypeError: If `dataset` is an unsupported type. RuntimeError: When invoked without eager execution enabled. """ - if isinstance(dataset, prefetching_ops._PrefetchToDeviceDataset): # pylint: disable=protected-access - raise TypeError( - "`tf.data.experimental.prefetch_to_device()` is not compatible with " - "`tf.contrib.eager.Iterator`. Use `for ... in dataset:` to iterate " - "over the dataset instead.") - if not context.context().device_spec.device_type: is_remote_device = False else: diff --git a/tensorflow/contrib/eager/python/datasets_test.py b/tensorflow/contrib/eager/python/datasets_test.py index 6a508fc6ba98740c4d441a064dc8a3e2b321f585..78ab155896cfeda4dd259a8529f4b1f77a12cf0b 100644 --- a/tensorflow/contrib/eager/python/datasets_test.py +++ b/tensorflow/contrib/eager/python/datasets_test.py @@ -26,7 +26,6 @@ import numpy as np from tensorflow.contrib import lookup from tensorflow.contrib.eager.python import datasets from tensorflow.python.data import Dataset -from tensorflow.python.data.experimental.ops import prefetching_ops from tensorflow.python.data.experimental.ops import threadpool from tensorflow.python.data.experimental.ops import unique from tensorflow.python.eager import test @@ -201,25 +200,6 @@ class IteratorTest(test.TestCase): y = math_ops.add(x, x) self.assertAllEqual([0., 2.], y.numpy()) - def testGpuDefinedDataset(self): - with ops.device(test.gpu_device_name()): - ds = Dataset.from_tensors([0., 1.]) - for x in ds: - y = math_ops.add(x, x) - self.assertAllEqual([0., 2.], y.numpy()) - - def testTensorsExplicitPrefetchToDevice(self): - ds = Dataset.from_tensor_slices([0., 1.]) - ds = ds.apply(prefetching_ops.prefetch_to_device(test.gpu_device_name())) - - with self.assertRaisesRegexp(TypeError, 'prefetch_to_device'): - datasets.Iterator(ds) - - for i, x in enumerate(ds): - with ops.device(test.gpu_device_name()): - x = math_ops.add(x, x) - self.assertEqual(float(i) + float(i), x.numpy()) - def testOverrideThreadPool(self): def get_thread_id(_): diff --git a/tensorflow/contrib/eager/python/evaluator.py b/tensorflow/contrib/eager/python/evaluator.py index 7949a3f6da293abdd85512209242bae76ab4d816..51443d24829bdc31a41813e0ff50ad7102422112 100644 --- a/tensorflow/contrib/eager/python/evaluator.py +++ b/tensorflow/contrib/eager/python/evaluator.py @@ -22,6 +22,7 @@ import six from tensorflow.contrib.eager.python import datasets from tensorflow.contrib.eager.python import metrics +from tensorflow.python.data.ops import dataset_ops from tensorflow.python.eager import context from tensorflow.python.eager import function from tensorflow.python.framework import errors_impl @@ -164,8 +165,8 @@ class Evaluator(object): self.__call__(example, *args, **kwargs) return self.all_metric_results(summary_logdir) # Graph construction - call_op = self.__call__(dataset.make_one_shot_iterator().get_next(), *args, - **kwargs) + call_op = self.__call__( + dataset_ops.make_one_shot_iterator(dataset).get_next(), *args, **kwargs) init_op = self.init_variables() results_op = self.all_metric_results(summary_logdir) return (init_op, call_op, results_op) 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 2dc196f550a10367066730f6f042c4ed69533ec3..fbb5daf230bb79f08a3d071062ddc0e8507ab324 100644 --- a/tensorflow/contrib/eager/python/examples/densenet/BUILD +++ b/tensorflow/contrib/eager/python/examples/densenet/BUILD @@ -3,11 +3,19 @@ licenses(["notice"]) # Apache 2.0 package(default_visibility = ["//tensorflow:internal"]) load("//tensorflow:tensorflow.bzl", "cuda_py_test") +load("//tensorflow:tensorflow.bzl", "py_binary") 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", @@ -16,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/densenet/densenet_graph_test.py b/tensorflow/contrib/eager/python/examples/densenet/densenet_graph_test.py index 4b3cb624bc947a1d1956eff6accb6d4da3bf3b87..24f6b007b526b29157011f3b1e9abdbd50bacc8e 100644 --- a/tensorflow/contrib/eager/python/examples/densenet/densenet_graph_test.py +++ b/tensorflow/contrib/eager/python/examples/densenet/densenet_graph_test.py @@ -119,7 +119,8 @@ class DensenetBenchmark(tf.test.Benchmark): with tf.Graph().as_default(): np_images, np_labels = random_batch(batch_size) dataset = tf.data.Dataset.from_tensors((np_images, np_labels)).repeat() - (images, labels) = dataset.make_one_shot_iterator().get_next() + (images, labels) = tf.compat.v1.data.make_one_shot_iterator( + dataset).get_next() model = densenet.DenseNet(self.depth, self.growth_rate, self.num_blocks, self.output_classes, 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/gan/mnist_graph_test.py b/tensorflow/contrib/eager/python/examples/gan/mnist_graph_test.py index 12b39b0cde49d4c017acfa74572c725036c54eff..e73841fbf724e05eaa3be90cc8650f795d3e1ccf 100644 --- a/tensorflow/contrib/eager/python/examples/gan/mnist_graph_test.py +++ b/tensorflow/contrib/eager/python/examples/gan/mnist_graph_test.py @@ -42,7 +42,8 @@ class MnistGraphGanBenchmark(tf.test.Benchmark): # Generate some random data. images_data = np.random.randn(batch_size, 784).astype(np.float32) dataset = tf.data.Dataset.from_tensors(images_data) - images = dataset.repeat().make_one_shot_iterator().get_next() + images = tf.compat.v1.data.make_one_shot_iterator( + dataset.repeat()).get_next() # Create the models and optimizers generator = mnist.Generator(data_format()) diff --git a/tensorflow/contrib/eager/python/examples/generative_examples/cvae.ipynb b/tensorflow/contrib/eager/python/examples/generative_examples/cvae.ipynb index ca27a85a229d41a85fa26ecdc982da478fe9e202..1a08cc0fd06516be4af5c2b0b46a3ffcf9101e95 100644 --- a/tensorflow/contrib/eager/python/examples/generative_examples/cvae.ipynb +++ b/tensorflow/contrib/eager/python/examples/generative_examples/cvae.ipynb @@ -470,7 +470,7 @@ "\n", " if epoch % 1 == 0:\n", " loss = tfe.metrics.Mean()\n", - " for test_x in test_dataset.make_one_shot_iterator():\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", diff --git a/tensorflow/contrib/eager/python/examples/generative_examples/dcgan.ipynb b/tensorflow/contrib/eager/python/examples/generative_examples/dcgan.ipynb index 5621d6a358e8969ea1a6663c1c770987de41ce0c..78fcd397087fd1fd64aebed7ac3b5c6b2f45c450 100644 --- a/tensorflow/contrib/eager/python/examples/generative_examples/dcgan.ipynb +++ b/tensorflow/contrib/eager/python/examples/generative_examples/dcgan.ipynb @@ -1,324 +1,405 @@ { + "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", + "**Copyright 2018 The TensorFlow Authors**.\n", "\n", "Licensed under the Apache License, Version 2.0 (the \"License\").\n", "\n", - "# DCGAN: An example with tf.keras and eager\n", + "# Generating Handwritten Digits with DCGAN\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/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/tensorflow/tree/master/tensorflow/contrib/eager/python/examples/generative_examples/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" + "
\n", + "\n", + " Run in Google Colab \n", + "\n", + "View source on GitHub
" ] }, { - "cell_type": "markdown", "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", + "\n", + ">>[Learn more about GANs](#scrollTo=k6qC-SbjK0yW)\n", + "\n" + ] + }, + { + "metadata": { + "colab_type": "text", + "id": "2MbKJY38Puy9" + }, + "cell_type": "markdown", "source": [ - "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). To do so, we use Deep Convolutional Generative Adverserial Networks ([DCGAN](https://arxiv.org/pdf/1511.06434.pdf)).\n", + "## 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", - "This model takes about ~30 seconds per epoch (using tf.contrib.eager.defun to create graph functions) to train on a single Tesla K80 on Colab, as of July 2018.\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", - "Below is the output generated after training the generator and discriminator models for 150 epochs.\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)" ] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "u_2z-B3piVsw" + "id": "u_2z-B3piVsw", + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ - "# to generate gifs\n", + "# Install imgeio in order to generate an animated gif showing the image generating process\n", "!pip install imageio" - ] + ], + "execution_count": 0, + "outputs": [] }, { - "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "e1_Y75QXJS6h" }, + "cell_type": "markdown", "source": [ - "## Import TensorFlow and enable eager execution" + "### Import TensorFlow and enable eager execution" ] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "YfIk2es3hJEd" + "id": "YfIk2es3hJEd", + "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 as tf\n", "tf.enable_eager_execution()\n", "\n", - "import os\n", - "import time\n", - "import numpy as np\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 imageio\n", + "import time\n", + "\n", "from IPython import display" - ] + ], + "execution_count": 0, + "outputs": [] }, { - "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "iYn4MdZnKCey" }, + "cell_type": "markdown", "source": [ - "## Load the dataset\n", + "### Load the dataset\n", "\n", - "We are going to use the MNIST dataset to train the generator and the discriminator. The generator will then generate handwritten digits." + "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." ] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "a4fYMGxGhrna" + "id": "a4fYMGxGhrna", + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ "(train_images, train_labels), (_, _) = tf.keras.datasets.mnist.load_data()" - ] + ], + "execution_count": 0, + "outputs": [] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "NFC2ghIdiZYE" + "id": "NFC2ghIdiZYE", + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ "train_images = train_images.reshape(train_images.shape[0], 28, 28, 1).astype('float32')\n", - "# We are normalizing the images to the range of [-1, 1]\n", - "train_images = (train_images - 127.5) / 127.5" - ] + "train_images = (train_images - 127.5) / 127.5 # Normalize the images to [-1, 1]" + ], + "execution_count": 0, + "outputs": [] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "S4PIDhoDLbsZ" + "id": "S4PIDhoDLbsZ", + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ "BUFFER_SIZE = 60000\n", "BATCH_SIZE = 256" - ] + ], + "execution_count": 0, + "outputs": [] }, { - "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "PIGN6ouoQxt3" }, + "cell_type": "markdown", "source": [ - "## Use tf.data to create batches and shuffle the dataset" + "### Use tf.data to create batches and shuffle the dataset" ] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "-yKCCQOoJ7cn" + "id": "-yKCCQOoJ7cn", + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ "train_dataset = tf.data.Dataset.from_tensor_slices(train_images).shuffle(BUFFER_SIZE).batch(BATCH_SIZE)" - ] + ], + "execution_count": 0, + "outputs": [] }, { - "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "THY-sZMiQ4UV" }, + "cell_type": "markdown", "source": [ - "## Write the generator and discriminator models\n", + "## Create the models\n", "\n", - "* **Generator** \n", - " * It is responsible for **creating convincing images that are good enough to fool the discriminator**.\n", - " * It consists of Conv2DTranspose (Upsampling) layers. We start with a fully connected layer and upsample the image 2 times so as to reach the desired image size (mnist image size) which is (28, 28, 1). \n", - " * We use **leaky relu** activation except for the **last layer** which uses **tanh** activation.\n", - " \n", - "* **Discriminator**\n", - " * **The discriminator is responsible for classifying the fake images from the real images.**\n", - " * In other words, the discriminator is given generated images (from the generator) and the real MNIST images. The job of the discriminator is to classify these images into fake (generated) and real (MNIST images).\n", - " * **Basically the generator should be good enough to fool the discriminator that the generated images are real**." + "We will use tf.keras [Sequential API](https://www.tensorflow.org/guide/keras#sequential_model) to define the generator and discriminator models." ] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, + "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", - "id": "VGLbvBEmjK0a" + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ - "class Generator(tf.keras.Model):\n", - " def __init__(self):\n", - " super(Generator, self).__init__()\n", - " self.fc1 = tf.keras.layers.Dense(7*7*64, use_bias=False)\n", - " self.batchnorm1 = tf.keras.layers.BatchNormalization()\n", - " \n", - " self.conv1 = tf.keras.layers.Conv2DTranspose(64, (5, 5), strides=(1, 1), padding='same', use_bias=False)\n", - " self.batchnorm2 = tf.keras.layers.BatchNormalization()\n", - " \n", - " self.conv2 = tf.keras.layers.Conv2DTranspose(32, (5, 5), strides=(2, 2), padding='same', use_bias=False)\n", - " self.batchnorm3 = tf.keras.layers.BatchNormalization()\n", + "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", - " self.conv3 = tf.keras.layers.Conv2DTranspose(1, (5, 5), strides=(2, 2), padding='same', use_bias=False)\n", - "\n", - " def call(self, x, training=True):\n", - " x = self.fc1(x)\n", - " x = self.batchnorm1(x, training=training)\n", - " x = tf.nn.relu(x)\n", - "\n", - " x = tf.reshape(x, shape=(-1, 7, 7, 64))\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", - " x = self.conv1(x)\n", - " x = self.batchnorm2(x, training=training)\n", - " x = tf.nn.relu(x)\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", - " x = self.conv2(x)\n", - " x = self.batchnorm3(x, training=training)\n", - " x = tf.nn.relu(x)\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", - " x = tf.nn.tanh(self.conv3(x)) \n", - " return x" + "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": {}, "colab_type": "code", - "id": "bkOfJxk5j5Hi" - }, - "outputs": [], - "source": [ - "class Discriminator(tf.keras.Model):\n", - " def __init__(self):\n", - " super(Discriminator, self).__init__()\n", - " self.conv1 = tf.keras.layers.Conv2D(64, (5, 5), strides=(2, 2), padding='same')\n", - " self.conv2 = tf.keras.layers.Conv2D(128, (5, 5), strides=(2, 2), padding='same')\n", - " self.dropout = tf.keras.layers.Dropout(0.3)\n", - " self.flatten = tf.keras.layers.Flatten()\n", - " self.fc1 = tf.keras.layers.Dense(1)\n", - "\n", - " def call(self, x, training=True):\n", - " x = tf.nn.leaky_relu(self.conv1(x))\n", - " x = self.dropout(x, training=training)\n", - " x = tf.nn.leaky_relu(self.conv2(x))\n", - " x = self.dropout(x, training=training)\n", - " x = self.flatten(x)\n", - " x = self.fc1(x)\n", - " return x" + "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" ] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, - "colab_type": "code", - "id": "gDkA05NE6QMs" + "colab_type": "text", + "id": "Jd-3GCUEiKtv" }, - "outputs": [], + "cell_type": "markdown", "source": [ - "generator = Generator()\n", - "discriminator = Discriminator()" + "### 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." ] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "k1HpMSLImuRi" + "id": "90BIcCKcDMxz", + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ - "# Defun gives 10 secs/epoch performance boost\n", - "generator.call = tf.contrib.eager.defun(generator.call)\n", - "discriminator.call = tf.contrib.eager.defun(discriminator.call)" - ] + "def generator_loss(generated_output):\n", + " return tf.losses.sigmoid_cross_entropy(tf.ones_like(generated_output), generated_output)" + ], + "execution_count": 0, + "outputs": [] }, { - "cell_type": "markdown", "metadata": { "colab_type": "text", - "id": "0FMYgY_mPfTi" + "id": "PKY_iPSPNWoj" }, + "cell_type": "markdown", "source": [ - "## Define the loss functions and the optimizer\n", - "\n", - "* **Discriminator loss**\n", - " * The discriminator loss function takes 2 inputs; **real images, generated images**\n", - " * real_loss is a sigmoid cross entropy loss of the **real images** and an **array of ones (since these are the real images)**\n", - " * generated_loss is a sigmoid cross entropy loss of the **generated images** and an **array of zeros (since these are the fake images)**\n", - " * Then the total_loss is the sum of real_loss and the generated_loss\n", - " \n", - "* **Generator loss**\n", - " * It is a sigmoid cross entropy loss of the generated images and an **array of ones**\n", - " \n", + "### Discriminator loss\n", "\n", - "* The discriminator and the generator optimizers are different since we will train them separately." + "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." ] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "wkMNfBWlT-PV" + "id": "wkMNfBWlT-PV", + "colab": {} }, - "outputs": [], + "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\n", - " # our generated examples to look like it\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", @@ -327,55 +408,51 @@ " total_loss = real_loss + generated_loss\n", "\n", " return total_loss" - ] + ], + "execution_count": 0, + "outputs": [] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, - "colab_type": "code", - "id": "90BIcCKcDMxz" + "colab_type": "text", + "id": "MgIc7i0th_Iu" }, - "outputs": [], + "cell_type": "markdown", "source": [ - "def generator_loss(generated_output):\n", - " return tf.losses.sigmoid_cross_entropy(tf.ones_like(generated_output), generated_output)" + "The discriminator and the generator optimizers are different since we will train two networks separately." ] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "iWCn_PVdEJZ7" + "id": "iWCn_PVdEJZ7", + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ - "discriminator_optimizer = tf.train.AdamOptimizer(1e-4)\n", - "generator_optimizer = tf.train.AdamOptimizer(1e-4)" - ] + "generator_optimizer = tf.train.AdamOptimizer(1e-4)\n", + "discriminator_optimizer = tf.train.AdamOptimizer(1e-4)" + ], + "execution_count": 0, + "outputs": [] }, { - "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "mWtinsGDPJlV" }, + "cell_type": "markdown", "source": [ - "## Checkpoints (Object-based saving)" + "**Checkpoints (Object-based saving)**" ] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "CA1w-7s2POEy" + "id": "CA1w-7s2POEy", + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ "checkpoint_dir = './training_checkpoints'\n", "checkpoint_prefix = os.path.join(checkpoint_dir, \"ckpt\")\n", @@ -383,93 +460,85 @@ " 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": "Rw1fkAczTQYh" + "id": "Ff6oN6PZX27n" }, + "cell_type": "markdown", "source": [ - "## Training\n", - "\n", - "* We start by iterating over the dataset\n", - "* The generator is given **noise as an input** which when passed through the generator model will output a image looking like a handwritten digit\n", - "* The discriminator is given the **real MNIST images as well as the generated images (from the generator)**.\n", - "* Next, we calculate the generator and the discriminator loss.\n", - "* Then, we calculate the gradients of loss with respect to both the generator and the discriminator variables (inputs) and apply those to the optimizer.\n", - "\n", - "## Generate Images\n", - "\n", - "* After training, its time to generate some images!\n", - "* We start by creating noise array as an input to the generator\n", - "* The generator will then convert the noise into handwritten images.\n", - "* Last step is to plot the predictions and **voila!**" + "**Define training parameters**" ] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "NS2GWywBbAWo" + "id": "NS2GWywBbAWo", + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ - "EPOCHS = 150\n", + "EPOCHS = 50\n", "noise_dim = 100\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 of the gan.\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": [] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, - "colab_type": "code", - "id": "RmdVsmvhPxyy" + "colab_type": "text", + "id": "jylSonrqSWfi" }, - "outputs": [], + "cell_type": "markdown", "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", + "**Define training method**\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()" + "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." ] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, + "id": "3t5ibNo05jCB", "colab_type": "code", - "id": "2M7LmLtGEMQJ" + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ - "def train(dataset, epochs, noise_dim): \n", - " for epoch in range(epochs):\n", - " start = time.time()\n", - " \n", - " for images in dataset:\n", - " # generating noise from a uniform distribution\n", + "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", @@ -477,7 +546,7 @@ " \n", " real_output = discriminator(images, training=True)\n", " generated_output = discriminator(generated_images, training=True)\n", - " \n", + " \n", " gen_loss = generator_loss(generated_output)\n", " disc_loss = discriminator_loss(real_output, generated_output)\n", " \n", @@ -485,12 +554,54 @@ " 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))\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", - " \n", - " if epoch % 1 == 0:\n", - " display.clear_output(wait=True)\n", - " generate_and_save_images(generator,\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", @@ -505,111 +616,167 @@ " 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": {}, - "colab_type": "code", - "id": "Ly3UN0SLLY2l" + "colab_type": "text", + "id": "dZrd4CdjR-Fp" }, - "outputs": [], + "cell_type": "markdown", "source": [ - "train(train_dataset, EPOCHS, noise_dim)" + "## 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." ] }, { - "cell_type": "markdown", + "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" + "**Restore the latest checkpoint**" ] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "XhXsd0srPo8c" + "id": "XhXsd0srPo8c", + "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": "markdown", "metadata": { "colab_type": "text", "id": "P4M_vIbUi7c0" }, + "cell_type": "markdown", "source": [ - "## Display an image using the epoch number" + "## Generated images \n" ] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, - "colab_type": "code", - "id": "WfO5wCdclHGL" + "colab_type": "text", + "id": "mLskt7EfXAjr" }, - "outputs": [], + "cell_type": "markdown", "source": [ - "def display_image(epoch_no):\n", - " return PIL.Image.open('image_at_epoch_{:04d}.png'.format(epoch_no))" + "\n", + "After training, its time to generate some images! \n", + "The last step is to plot the generated images and voila!\n" ] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "5x3q9_Oe5q0A" + "id": "WfO5wCdclHGL", + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ - "display_image(EPOCHS)" - ] + "# 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": [] }, { - "cell_type": "markdown", "metadata": { - "colab_type": "text", - "id": "NywiH3nL8guF" + "colab_type": "code", + "id": "5x3q9_Oe5q0A", + "colab": {} }, + "cell_type": "code", "source": [ - "## Generate a GIF of all the saved images." - ] + "display_image(EPOCHS)" + ], + "execution_count": 0, + "outputs": [] }, { - "cell_type": "markdown", "metadata": { "colab_type": "text", - "id": "xmO0Dmu2WICn" + "id": "NywiH3nL8guF" }, + "cell_type": "markdown", "source": [ - "\u003c!-- TODO(markdaoust): Remove the hack when Ipython version is updated --\u003e\n" + "**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." ] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "IGKQgENQ8lEI" + "id": "IGKQgENQ8lEI", + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ "with imageio.get_writer('dcgan.gif', mode='I') as writer:\n", " filenames = glob.glob('image*.png')\n", @@ -617,7 +784,7 @@ " last = -1\n", " for i,filename in enumerate(filenames):\n", " frame = 2*(i**0.5)\n", - " if round(frame) \u003e round(last):\n", + " if round(frame) > round(last):\n", " last = frame\n", " else:\n", " continue\n", @@ -628,67 +795,84 @@ " \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." ] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "uV0yiKpzNP1b" + "id": "uV0yiKpzNP1b", + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ "display.Image(filename=\"dcgan.gif.png\")" - ] + ], + "execution_count": 0, + "outputs": [] }, { - "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "6EEG-wePkmJQ" }, + "cell_type": "markdown", "source": [ - "To downlod the animation from Colab uncomment the code below:" + "**Download the animated gif**\n", + "\n", + "Uncomment the code below to download an animated gif from Colab." ] }, { - "cell_type": "code", - "execution_count": 0, "metadata": { - "colab": {}, "colab_type": "code", - "id": "4UJjSnIMOzOJ" + "id": "4UJjSnIMOzOJ", + "colab": {} }, - "outputs": [], + "cell_type": "code", "source": [ "#from google.colab import files\n", "#files.download('dcgan.gif')" - ] - } - ], - "metadata": { - "accelerator": "GPU", - "colab": { - "collapsed_sections": [], - "name": "dcgan.ipynb", - "private_outputs": true, - "provenance": [ - { - "file_id": "1eb0NOTQapkYs3X0v-zL1x5_LFKgDISnp", - "timestamp": 1527173385672 - } ], - "toc_visible": true, - "version": "0.3.2" + "execution_count": 0, + "outputs": [] }, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" + { + "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" + ] } - }, - "nbformat": 4, - "nbformat_minor": 0 -} + ] +} \ No newline at end of file diff --git a/tensorflow/contrib/eager/python/examples/generative_examples/gans_diagram.png b/tensorflow/contrib/eager/python/examples/generative_examples/gans_diagram.png new file mode 100644 index 0000000000000000000000000000000000000000..b715bd83ef117641c6429e0ac173dbe9b8d5fd88 Binary files /dev/null and b/tensorflow/contrib/eager/python/examples/generative_examples/gans_diagram.png 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 3acecd283cda83992bab0c37cf0b8037ed2cf27a..12c5eff2b4aa901bdab52bf545e95b1e4dce7468 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,1184 +1,1174 @@ { - "nbformat": 4, - "nbformat_minor": 0, - "metadata": { + "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": { - "name": "image_captioning_with_attention.ipynb", - "version": "0.3.2", - "views": {}, - "default_view": {}, - "provenance": [ - { - "file_id": "1HI8OK2sMjcx9CTWVn0122QAHOuXaOaMg", - "timestamp": 1530222436922 - } - ], - "private_outputs": true, - "collapsed_sections": [], - "toc_visible": true + "autoexec": { + "startup": false, + "wait_interval": 0 + } }, - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "accelerator": "GPU" + "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" + ] }, - "cells": [ - { - "metadata": { - "id": "K2s1A9eLRPEj", - "colab_type": "text" - }, - "cell_type": "markdown", - "source": [ - "##### Copyright 2018 The TensorFlow Authors.\n", - "\n", - "Licensed under the Apache License, Version 2.0 (the \"License\").\n" - ] - }, - { - "metadata": { - "id": "Cffg2i257iMS", - "colab_type": "text" - }, - "cell_type": "markdown", - "source": [ - "# Image Captioning with Attention\n", - "\n", - "
\n", - "\n", - " Run in Google Colab \n", - "\n", - "View source on GitHub
" - ] - }, - { - "metadata": { - "id": "QASbY_HGo4Lq", - "colab_type": "text" - }, - "cell_type": "markdown", - "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" - ] - }, - { - "metadata": { - "id": "U8l4RJ0XRPEm", - "colab_type": "code", - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - } - }, - "cell_type": "code", - "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" - ], - "execution_count": 0, - "outputs": [] - }, - { - "metadata": { - "id": "b6qbGw8MRPE5", - "colab_type": "text" - }, - "cell_type": "markdown", - "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." - ] - }, - { - "metadata": { - "id": "krQuPYTtRPE7", - "colab_type": "code", - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - } - }, - "cell_type": "code", - "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/'" - ], - "execution_count": 0, - "outputs": [] - }, - { - "metadata": { - "id": "aANEzb5WwSzg", - "colab_type": "text" - }, - "cell_type": "markdown", - "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." - ] - }, - { - "metadata": { - "id": "4G3b8x8_RPFD", - "colab_type": "code", - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - } - }, - "cell_type": "code", - "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]" - ], - "execution_count": 0, - "outputs": [] - }, - { - "metadata": { - "id": "mPBMgK34RPFL", - "colab_type": "code", - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - } - }, - "cell_type": "code", - "source": [ - "len(train_captions), len(all_captions)" - ], - "execution_count": 0, - "outputs": [] - }, - { - "metadata": { - "id": "8cSW4u-ORPFQ", - "colab_type": "text" - }, - "cell_type": "markdown", - "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)." - ] - }, - { - "metadata": { - "id": "zXR0217aRPFR", - "colab_type": "code", - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - } - }, - "cell_type": "code", - "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" - ], - "execution_count": 0, - "outputs": [] - }, - { - "metadata": { - "id": "MDvIu4sXRPFV", - "colab_type": "text" - }, - "cell_type": "markdown", - "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." - ] - }, - { - "metadata": { - "id": "RD3vW4SsRPFW", - "colab_type": "code", - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - } - }, - "cell_type": "code", - "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)" - ], - "execution_count": 0, - "outputs": [] - }, - { - "metadata": { - "id": "rERqlR3WRPGO", - "colab_type": "text" - }, - "cell_type": "markdown", - "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):```." - ] - }, - { - "metadata": { - "id": "Dx_fvbVgRPGQ", - "colab_type": "code", - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - } - }, - "cell_type": "code", - "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())" - ], - "execution_count": 0, - "outputs": [] - }, - { - "metadata": { - "id": "nyqH3zFwRPFi", - "colab_type": "text" - }, - "cell_type": "markdown", - "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. " - ] - }, - { - "metadata": { - "id": "HZfK8RhQRPFj", - "colab_type": "code", - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - } - }, - "cell_type": "code", - "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)" - ], - "execution_count": 0, - "outputs": [] - }, - { - "metadata": { - "id": "oJGE34aiRPFo", - "colab_type": "code", - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - } - }, - "cell_type": "code", - "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)" - ], - "execution_count": 0, - "outputs": [] + { + "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 + } }, - { - "metadata": { - "id": "8Q44tNQVRPFt", - "colab_type": "code", - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - } - }, - "cell_type": "code", - "source": [ - "tokenizer.word_index = {key:value for key, value in tokenizer.word_index.items() if value <= top_k}\n", - "# putting token in the word2idx dictionary\n", - "tokenizer.word_index[tokenizer.oov_token] = top_k + 1\n", - "tokenizer.word_index[''] = 0" - ], - "execution_count": 0, - "outputs": [] + "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 + } }, - { - "metadata": { - "id": "0fpJb5ojRPFv", - "colab_type": "code", - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - } - }, - "cell_type": "code", - "source": [ - "# creating the tokenized vectors\n", - "train_seqs = tokenizer.texts_to_sequences(train_captions)" - ], - "execution_count": 0, - "outputs": [] + "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 + } }, - { - "metadata": { - "id": "olQArbgbRPF1", - "colab_type": "code", - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - } - }, - "cell_type": "code", - "source": [ - "# creating a reverse mapping (index -> word)\n", - "index_word = {value:key for key, value in tokenizer.word_index.items()}" - ], - "execution_count": 0, - "outputs": [] + "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 + } }, - { - "metadata": { - "id": "AidglIZVRPF4", - "colab_type": "code", - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - } - }, - "cell_type": "code", - "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')" - ], - "execution_count": 0, - "outputs": [] + "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 + } }, - { - "metadata": { - "id": "gL0wkttkRPGA", - "colab_type": "code", - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - } - }, - "cell_type": "code", - "source": [ - "# calculating the max_length \n", - "# used to store the attention weights\n", - "max_length = calc_max_length(train_seqs)" - ], - "execution_count": 0, - "outputs": [] + "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 + } }, - { - "metadata": { - "id": "M3CD75nDpvTI", - "colab_type": "text" - }, - "cell_type": "markdown", - "source": [ - "## Split the data into training and testing" - ] + "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 + } }, - { - "metadata": { - "id": "iS7DDMszRPGF", - "colab_type": "code", - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - } - }, - "cell_type": "code", - "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)" - ], - "execution_count": 0, - "outputs": [] + "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 + } }, - { - "metadata": { - "id": "XmViPkRFRPGH", - "colab_type": "code", - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - } - }, - "cell_type": "code", - "source": [ - "len(img_name_train), len(cap_train), len(img_name_val), len(cap_val)" - ], - "execution_count": 0, - "outputs": [] + "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 + } }, - { - "metadata": { - "id": "uEWM9xrYcg45", - "colab_type": "text" - }, - "cell_type": "markdown", - "source": [ - "## Our images and captions are ready! Next, let's create a tf.data dataset to use for training our model.\n", - "\n" - ] + "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 + } }, - { - "metadata": { - "id": "Q3TnZ1ToRPGV", - "colab_type": "code", - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - } - }, - "cell_type": "code", - "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" - ], - "execution_count": 0, - "outputs": [] + "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 + } }, - { - "metadata": { - "id": "SmZS2N0bXG3T", - "colab_type": "code", - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - } - }, - "cell_type": "code", - "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" - ], - "execution_count": 0, - "outputs": [] + "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 + } }, - { - "metadata": { - "id": "FDF_Nm3tRPGZ", - "colab_type": "code", - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - } - }, - "cell_type": "code", - "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)" - ], - "execution_count": 0, - "outputs": [] + "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 + } }, - { - "metadata": { - "id": "nrvoDphgRPGd", - "colab_type": "text" - }, - "cell_type": "markdown", - "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." - ] + "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 + } }, - { - "metadata": { - "id": "AAppCGLKRPGd", - "colab_type": "code", - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - } - }, - "cell_type": "code", - "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')" - ], - "execution_count": 0, - "outputs": [] + "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 + } }, - { - "metadata": { - "id": "ja2LFTMSdeV3", - "colab_type": "code", - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - } - }, - "cell_type": "code", - "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" - ], - "execution_count": 0, - "outputs": [] + "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 + } }, - { - "metadata": { - "id": "AZ7R1RxHRPGf", - "colab_type": "code", - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - } - }, - "cell_type": "code", - "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" - ], - "execution_count": 0, - "outputs": [] + "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 + } }, - { - "metadata": { - "id": "V9UbGQmERPGi", - "colab_type": "code", - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - } - }, - "cell_type": "code", - "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))" - ], - "execution_count": 0, - "outputs": [] + "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 + } }, - { - "metadata": { - "id": "Qs_Sr03wRPGk", - "colab_type": "code", - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - } - }, - "cell_type": "code", - "source": [ - "encoder = CNN_Encoder(embedding_dim)\n", - "decoder = RNN_Decoder(embedding_dim, units, vocab_size)" - ], - "execution_count": 0, - "outputs": [] + "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 + } }, - { - "metadata": { - "id": "-bYN7xA0RPGl", - "colab_type": "code", - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - } - }, - "cell_type": "code", - "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_)" - ], - "execution_count": 0, - "outputs": [] + "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 + } }, - { - "metadata": { - "id": "PHod7t72RPGn", - "colab_type": "text" - }, - "cell_type": "markdown", - "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" - ] + "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 + } }, - { - "metadata": { - "id": "Vt4WZ5mhJE-E", - "colab_type": "code", - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - } - }, - "cell_type": "code", - "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 = []" - ], - "execution_count": 0, - "outputs": [] + "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 + } }, - { - "metadata": { - "id": "UlA4VIQpRPGo", - "colab_type": "code", - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - } - }, - "cell_type": "code", - "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))" - ], - "execution_count": 0, - "outputs": [] + "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 + } }, - { - "metadata": { - "id": "1Wm83G-ZBPcC", - "colab_type": "code", - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - } - }, - "cell_type": "code", - "source": [ - "plt.plot(loss_plot)\n", - "plt.xlabel('Epochs')\n", - "plt.ylabel('Loss')\n", - "plt.title('Loss Plot')\n", - "plt.show()" - ], - "execution_count": 0, - "outputs": [] + "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 + } }, - { - "metadata": { - "id": "xGvOcLQKghXN", - "colab_type": "text" - }, - "cell_type": "markdown", - "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." - ] + "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 + } }, - { - "metadata": { - "id": "RCWpDtyNRPGs", - "colab_type": "code", - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - } - }, - "cell_type": "code", - "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(index_word[predicted_id])\n", - "\n", - " if 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" - ], - "execution_count": 0, - "outputs": [] + "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 + } }, - { - "metadata": { - "id": "fD_y7PD6RPGt", - "colab_type": "code", - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - } - }, - "cell_type": "code", - "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()" - ], - "execution_count": 0, - "outputs": [] + "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 + } }, - { - "metadata": { - "id": "io7ws3ReRPGv", - "colab_type": "code", - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - } - }, - "cell_type": "code", - "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([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])" - ], - "execution_count": 0, - "outputs": [] + "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 + } }, - { - "metadata": { - "id": "Rprk3HEvZuxb", - "colab_type": "text" - }, - "cell_type": "markdown", - "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" - ] + "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 + } }, - { - "metadata": { - "id": "9Psd1quzaAWg", - "colab_type": "code", - "colab": { - "autoexec": { - "startup": false, - "wait_interval": 0 - } - } - }, - "cell_type": "code", - "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)" - ], - "execution_count": 0, - "outputs": [] + "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 + } }, + "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": [ { - "metadata": { - "id": "VJZXyJco6uLO", - "colab_type": "text" - }, - "cell_type": "markdown", - "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." - ] + "file_id": "1HI8OK2sMjcx9CTWVn0122QAHOuXaOaMg", + "timestamp": 1530222436922 } - ] + ], + "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 } 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_graph_test.py b/tensorflow/contrib/eager/python/examples/linear_regression/linear_regression_graph_test.py index 557ad42752144243ae3da61b955b31398cba846e..d412b25b368260b81256fd58034330b884261b2b 100644 --- a/tensorflow/contrib/eager/python/examples/linear_regression/linear_regression_graph_test.py +++ b/tensorflow/contrib/eager/python/examples/linear_regression/linear_regression_graph_test.py @@ -36,7 +36,7 @@ class GraphLinearRegressionBenchmark(tf.test.Benchmark): noise_level=0.01, batch_size=batch_size, num_batches=num_batches) - iterator = dataset.make_initializable_iterator() + iterator = tf.compat.v1.data.make_initializable_iterator(dataset) x, y = iterator.get_next() model = linear_regression.LinearModel() 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 480777d948769b56ac1cc3be2052fe48459e98d6..66d52a74943d0d81fde05ce51b019558b327978d 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 @@ -768,7 +768,7 @@ }, "outputs": [], "source": [ - "translate('hace mucho frio aqui.', encoder, decoder, inp_lang, targ_lang, max_length_inp, max_length_targ)" + "translate(u'hace mucho frio aqui.', encoder, decoder, inp_lang, targ_lang, max_length_inp, max_length_targ)" ] }, { @@ -781,7 +781,7 @@ }, "outputs": [], "source": [ - "translate('esta es mi vida.', encoder, decoder, inp_lang, targ_lang, max_length_inp, max_length_targ)" + "translate(u'esta es mi vida.', encoder, decoder, inp_lang, targ_lang, max_length_inp, max_length_targ)" ] }, { @@ -794,7 +794,7 @@ }, "outputs": [], "source": [ - "translate('¿todavia estan en casa?', encoder, decoder, inp_lang, targ_lang, max_length_inp, max_length_targ)" + "translate(u'todavia estan en casa?', encoder, decoder, inp_lang, targ_lang, max_length_inp, max_length_targ)" ] }, { @@ -808,7 +808,7 @@ "outputs": [], "source": [ "# wrong translation\n", - "translate('trata de averiguarlo.', encoder, decoder, inp_lang, targ_lang, max_length_inp, max_length_targ)" + "translate(u'trata de averiguarlo.', encoder, decoder, inp_lang, targ_lang, max_length_inp, max_length_targ)" ] }, { diff --git a/tensorflow/contrib/eager/python/examples/resnet50/BUILD b/tensorflow/contrib/eager/python/examples/resnet50/BUILD index 68a84d5fbb4f13e4ebe0d71e3f5caebe97e2101c..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,11 +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", @@ -47,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/resnet50/resnet50_graph_test.py b/tensorflow/contrib/eager/python/examples/resnet50/resnet50_graph_test.py index f3bb978875e226f58d6a00e09154191673a97415..fb7975d8fe867711cff31d627788a2d62a520aa9 100644 --- a/tensorflow/contrib/eager/python/examples/resnet50/resnet50_graph_test.py +++ b/tensorflow/contrib/eager/python/examples/resnet50/resnet50_graph_test.py @@ -142,7 +142,8 @@ class ResNet50Benchmarks(tf.test.Benchmark): with tf.Graph().as_default(): np_images, np_labels = random_batch(batch_size) dataset = tf.data.Dataset.from_tensors((np_images, np_labels)).repeat() - (images, labels) = dataset.make_one_shot_iterator().get_next() + images, labels = tf.compat.v1.data.make_one_shot_iterator( + dataset).get_next() model = resnet50.ResNet50(data_format()) logits = model(images, training=True) 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/revnet/main.py b/tensorflow/contrib/eager/python/examples/revnet/main.py index b702e91f92220c2a9003a1b82411131332012a9e..9585f3565f83af724b6336e466d3671443ba2361 100644 --- a/tensorflow/contrib/eager/python/examples/revnet/main.py +++ b/tensorflow/contrib/eager/python/examples/revnet/main.py @@ -72,14 +72,11 @@ def main(_): train_one_iter(model, x, y, optimizer, global_step=global_step) if global_step.numpy() % config.log_every == 0: - it_test = ds_test.make_one_shot_iterator() - acc_test, loss_test = evaluate(model, it_test) + acc_test, loss_test = evaluate(model, ds_test) if FLAGS.validate: - it_train = ds_train_one_shot.make_one_shot_iterator() - it_validation = ds_validation.make_one_shot_iterator() - acc_train, loss_train = evaluate(model, it_train) - acc_validation, loss_validation = evaluate(model, it_validation) + acc_train, loss_train = evaluate(model, ds_train_one_shot) + acc_validation, loss_validation = evaluate(model, ds_validation) print("Iter {}, " "training set accuracy {:.4f}, loss {:.4f}; " "validation set accuracy {:.4f}, loss {:.4f}; " @@ -218,11 +215,11 @@ def train_one_iter(model, inputs, labels, optimizer, global_step=None): return logits, loss -def evaluate(model, iterator): +def evaluate(model, dataset): """Compute accuracy with the given dataset iterator.""" mean_loss = tfe.metrics.Mean() accuracy = tfe.metrics.Accuracy() - for x, y in iterator: + for x, y in dataset: logits, _ = model(x, training=False) loss = model.compute_loss(logits=logits, labels=y) accuracy( diff --git a/tensorflow/contrib/eager/python/examples/revnet/main_estimator.py b/tensorflow/contrib/eager/python/examples/revnet/main_estimator.py index 3a17eb30da3b989acb0b33f2fcb730da76546c18..125adbb9de6e4febbb4284bfe3a31f257e2e8037 100644 --- a/tensorflow/contrib/eager/python/examples/revnet/main_estimator.py +++ b/tensorflow/contrib/eager/python/examples/revnet/main_estimator.py @@ -173,7 +173,7 @@ def main(_): input_fn = tf.estimator.export.build_raw_serving_input_receiver_fn({ "image": inputs }) - revnet_estimator.export_savedmodel(FLAGS.model_dir, input_fn) + revnet_estimator.export_saved_model(FLAGS.model_dir, input_fn) if __name__ == "__main__": diff --git a/tensorflow/contrib/eager/python/examples/revnet/main_estimator_tpu.py b/tensorflow/contrib/eager/python/examples/revnet/main_estimator_tpu.py index 8520cf5b71af503be35d5415707a283fb363a476..b0676916a8da276704de741a50f40cd7d9525228 100644 --- a/tensorflow/contrib/eager/python/examples/revnet/main_estimator_tpu.py +++ b/tensorflow/contrib/eager/python/examples/revnet/main_estimator_tpu.py @@ -307,7 +307,7 @@ def main(_): # The guide to serve an exported TensorFlow model is at: # https://www.tensorflow.org/serving/serving_basic tf.logging.info("Starting to export model.") - revnet_classifier.export_savedmodel( + revnet_classifier.export_saved_model( export_dir_base=FLAGS.export_dir, serving_input_receiver_fn=imagenet_input.image_serving_input_fn) diff --git a/tensorflow/contrib/eager/python/examples/revnet/revnet.py b/tensorflow/contrib/eager/python/examples/revnet/revnet.py index 1f2cb14972f0b92d29489adff8f94e790e1ec4ed..7406787ba438345dc485c50e347e40597b2037f5 100644 --- a/tensorflow/contrib/eager/python/examples/revnet/revnet.py +++ b/tensorflow/contrib/eager/python/examples/revnet/revnet.py @@ -96,6 +96,7 @@ class RevNet(tf.keras.Model): def call(self, inputs, training=True): """Forward pass.""" + saved_hidden = None if training: saved_hidden = [inputs] 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_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/rnn_ptb/rnn_ptb_graph_test.py b/tensorflow/contrib/eager/python/examples/rnn_ptb/rnn_ptb_graph_test.py index 63b5c4c54d13e9c2448ec1f572ca1389f2443bef..770484abed96e540cf75cc5368a1410c31a8d2d0 100644 --- a/tensorflow/contrib/eager/python/examples/rnn_ptb/rnn_ptb_graph_test.py +++ b/tensorflow/contrib/eager/python/examples/rnn_ptb/rnn_ptb_graph_test.py @@ -82,7 +82,7 @@ class PTBBenchmark(tf.test.Benchmark): tf.ones( [PTBBenchmark.SEQ_LEN, PTBBenchmark.BATCH_SIZE], dtype=tf.int64)).repeat(num_iters + num_warmup) - inputs = dataset.make_one_shot_iterator().get_next() + inputs = tf.compat.v1.data.make_one_shot_iterator(dataset).get_next() with tf.device(tf.test.gpu_device_name()): outputs = model(inputs, training=True) @@ -124,7 +124,8 @@ class PTBBenchmark(tf.test.Benchmark): dtype=tf.int64)).repeat(num_iters + num_warmup) # inputs and labels have the same shape dataset = tf.data.Dataset.zip((dataset, dataset)) - (inputs, labels) = dataset.make_one_shot_iterator().get_next() + (inputs, labels) = tf.compat.v1.data.make_one_shot_iterator( + dataset).get_next() with tf.device(tf.test.gpu_device_name()): optimizer = tf.train.GradientDescentOptimizer(learning_rate=1.0) 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 930e62b68096b468846a01b9674c669a8b8e9a53..c8d9266672a8b87d32338ea7c4f74fb40d41c767 100644 --- a/tensorflow/contrib/eager/python/metrics_impl.py +++ b/tensorflow/contrib/eager/python/metrics_impl.py @@ -24,6 +24,7 @@ from tensorflow.python.eager import context from tensorflow.python.eager import function from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops +from tensorflow.python.framework import smart_cond from tensorflow.python.ops import array_ops from tensorflow.python.ops import check_ops from tensorflow.python.ops import control_flow_ops @@ -36,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: @@ -347,16 +348,17 @@ class Mean(Metric): Raises: ValueError: if the optional argument is not bool """ - # Convert the boolean to tensor for tf.cond, if it is not. + # Convert the boolean to tensor for tf.cond, if it is not. if not isinstance(write_summary, ops.Tensor): write_summary = ops.convert_to_tensor(write_summary) t = self.numer / self.denom def write_summary_f(): summary_ops.scalar(name=self.name, tensor=t) return t - control_flow_ops.cond(write_summary, + smart_cond.smart_cond(write_summary, write_summary_f, - lambda: t) + lambda: t, + name="") return t @@ -487,6 +489,8 @@ class BinaryAccuracy(Mean): message="Shapes of labels and predictions are unequal") predictions = ops.convert_to_tensor(predictions) predictions = predictions > self.threshold + # Convert labels to bool to match predictions. + labels = math_ops.cast(labels, dtypes.bool) matches = math_ops.equal(labels, predictions) matches = math_ops.cast(matches, self.dtype) super(BinaryAccuracy, self).call(matches, weights=weights) diff --git a/tensorflow/contrib/eager/python/metrics_test.py b/tensorflow/contrib/eager/python/metrics_test.py index 9d2d172752c7f3f3ee6eaa11ab8952313a4a3543..39e5957f5d1760613f2c33607c0bdb163040efb4 100644 --- a/tensorflow/contrib/eager/python/metrics_test.py +++ b/tensorflow/contrib/eager/python/metrics_test.py @@ -49,18 +49,6 @@ class MetricsTest(test.TestCase): self.assertEqual(dtypes.float64, m.dtype) self.assertEqual(dtypes.float64, m.result().dtype) - def testSummaryArg(self): - m = metrics.Mean() - m([1, 10, 100]) - m(1000) - m([10000.0, 100000.0]) - self.assertEqual(111111.0/6, m.result(write_summary=True).numpy()) - self.assertEqual(111111.0/6, m.result(write_summary=False).numpy()) - with self.assertRaises(ValueError): - m.result(write_summary=5) - with self.assertRaises(ValueError): - m.result(write_summary=[True]) - def testVariableCollections(self): with context.graph_mode(), ops.Graph().as_default(): m = metrics.Mean() diff --git a/tensorflow/contrib/eager/python/network.py b/tensorflow/contrib/eager/python/network.py index f801d9a47b2f831a48d9b6335c69612c1356d800..5cc0c4f23d9d641ff1452c7cc9c1fcde612a33a2 100644 --- a/tensorflow/contrib/eager/python/network.py +++ b/tensorflow/contrib/eager/python/network.py @@ -24,7 +24,7 @@ import weakref from tensorflow.python.eager import context from tensorflow.python.framework import ops -from tensorflow.python.keras.engine import base_layer as keras_base_layer +from tensorflow.python.keras.engine import base_layer_utils from tensorflow.python.layers import base from tensorflow.python.ops import variable_scope from tensorflow.python.platform import tf_logging as logging @@ -220,7 +220,7 @@ class Network(base.Layer): avoid_names = parent_network._owned_layers name_uid_map = parent_network._sub_layer_name_uids else: - name_uid_map = keras_base_layer.get_default_graph_uid_map() + name_uid_map = base_layer_utils.get_default_graph_uid_map() # Figure out which names we have to avoid based on which variable scope # we're nested in. strip_name = self._default_parent_variable_scope.name diff --git a/tensorflow/contrib/eager/python/parameter_server.py b/tensorflow/contrib/eager/python/parameter_server.py index 3a9e7b027ed68935f2bc0ddbd27a1821a663850d..7803a6799bb64441fab881bf6ca986d5cf3851a8 100644 --- a/tensorflow/contrib/eager/python/parameter_server.py +++ b/tensorflow/contrib/eager/python/parameter_server.py @@ -56,12 +56,7 @@ def _eager_safe_variable_handle(shape, dtype, shared_name, name, graph_mode): # shape inference doesn't run in eager mode we copy this data here for when # the handle is captured by an eager mode function. # pylint: disable=protected-access - if ops._USE_C_SHAPES: - handle._handle_data = resource_variable_ops.get_resource_handle_data(h) - else: - if h._handle_data is None: - ops.set_shape_and_handle_data_for_outputs(h.op) - handle._handle_data = h._handle_data + handle._handle_data = resource_variable_ops.get_resource_handle_data(h) # pylint: enable=protected-access # Clean up op->graph->op reference cycles. ops.dismantle_graph(graph) diff --git a/tensorflow/contrib/eager/python/remote_test.py b/tensorflow/contrib/eager/python/remote_test.py index 3926de15e71c9917f88fc3f58740b8c75354ab26..f540d9b37b69c7be3b0662b07bd6e9cb8220fadc 100644 --- a/tensorflow/contrib/eager/python/remote_test.py +++ b/tensorflow/contrib/eager/python/remote_test.py @@ -24,12 +24,12 @@ import os import numpy as np from tensorflow.contrib.eager.python import parameter_server -from tensorflow.contrib.eager.python import remote from tensorflow.core.protobuf import cluster_pb2 from tensorflow.core.protobuf import tensorflow_server_pb2 from tensorflow.python.eager import backprop from tensorflow.python.eager import context from tensorflow.python.eager import function +from tensorflow.python.eager import remote from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import math_ops diff --git a/tensorflow/contrib/eager/python/saver.py b/tensorflow/contrib/eager/python/saver.py index f9c716360c5755ee1902b576545d776725f9966f..1d0d6c6c14ce4a8e454206e0be9fea4724f09192 100644 --- a/tensorflow/contrib/eager/python/saver.py +++ b/tensorflow/contrib/eager/python/saver.py @@ -115,6 +115,11 @@ def restore_variables_on_create(save_path, map_func=None): class Saver(object): """A tf.train.Saver adapter for use when eager execution is enabled. + + `Saver`'s name-based checkpointing strategy is fragile. Please switch to + `tf.train.Checkpoint` or `tf.keras.Model.save_weights`, which perform a more + robust object-based saving. These APIs will load checkpoints written by + `Saver`. """ def __init__(self, var_list): diff --git a/tensorflow/contrib/eager/python/tfe.py b/tensorflow/contrib/eager/python/tfe.py index 33c988fd9065e7fbe7b9aeb85cad82eb3c119f76..b82e1bb71bce9a28d7bbbf961cc6d5e25dd18acf 100644 --- a/tensorflow/contrib/eager/python/tfe.py +++ b/tensorflow/contrib/eager/python/tfe.py @@ -41,6 +41,8 @@ To use, at program startup, call `tf.enable_eager_execution()`. @@add_execution_callback @@clear_execution_callbacks +@@errstate +@@ExecutionCallback @@inf_callback @@inf_nan_callback @@nan_callback @@ -97,7 +99,6 @@ from tensorflow.contrib.eager.python.network import Network from tensorflow.contrib.eager.python.network import Sequential from tensorflow.contrib.eager.python.network import save_network_checkpoint from tensorflow.contrib.eager.python.network import restore_network_checkpoint -from tensorflow.contrib.eager.python.remote import connect_to_remote_host from tensorflow.contrib.eager.python.saver import get_optimizer_variables from tensorflow.contrib.eager.python.saver import restore_variables_on_create from tensorflow.contrib.eager.python.saver import Saver @@ -119,10 +120,13 @@ from tensorflow.python.eager.context import set_server_def from tensorflow.python.eager.def_function import function from tensorflow.python.eager.execution_callbacks import add_execution_callback from tensorflow.python.eager.execution_callbacks import clear_execution_callbacks +from tensorflow.python.eager.execution_callbacks import errstate +from tensorflow.python.eager.execution_callbacks import ExecutionCallback from tensorflow.python.eager.execution_callbacks import inf_callback from tensorflow.python.eager.execution_callbacks import inf_nan_callback from tensorflow.python.eager.execution_callbacks import nan_callback from tensorflow.python.eager.execution_callbacks import seterr +from tensorflow.python.eager.remote import connect_to_remote_host from tensorflow.python.framework.tensor_spec import TensorSpec from tensorflow.python.framework.ops import enable_eager_execution from tensorflow.python.framework.ops import enable_eager_execution_internal as enable_remote_eager_execution @@ -134,7 +138,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.tracking import AutoCheckpointable as Checkpointable from tensorflow.python.training.checkpointable.util import CheckpointableSaver 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 4454abfb9667f824b9de0100bb81bae24ad5f7a6..8c35dddb5a515aa09cc70c173a9f0605e8567e82 100644 --- a/tensorflow/contrib/eager/python/tfe_test.py +++ b/tensorflow/contrib/eager/python/tfe_test.py @@ -87,8 +87,8 @@ class TFETest(test_util.TensorFlowTestCase): x += 1. # Without a device context, heuristics are used to place ops. # In this case, ops.reduce_mean runs on the GPU. - reduction_indices = range(x.shape.ndims) - m = math_ops.reduce_mean(x, reduction_indices) + axis = range(x.shape.ndims) + m = math_ops.reduce_mean(x, axis) # m is on GPU, bring it back to CPU and compare. self.assertEqual(3.5, m.cpu().numpy()) diff --git a/tensorflow/contrib/estimator/BUILD b/tensorflow/contrib/estimator/BUILD index 8b99158b3068a331a201bd7354a75f69c4b0afb5..a888379f13e79d1c246d4cd6d19a225c065692a2 100644 --- a/tensorflow/contrib/estimator/BUILD +++ b/tensorflow/contrib/estimator/BUILD @@ -15,10 +15,7 @@ py_library( srcs = ["__init__.py"], srcs_version = "PY2AND3", deps = [ - ":baseline", ":boosted_trees", - ":dnn", - ":dnn_linear_combined", ":dnn_with_layer_annotations", ":early_stopping", ":expect_tensorflow_estimator_installed", @@ -27,7 +24,6 @@ py_library( ":extenders", ":head", ":hooks", - ":linear", ":logit_fns", ":multi_head", ":replicate_model_fn", @@ -38,17 +34,6 @@ py_library( ], ) -py_library( - name = "baseline", - srcs = ["python/estimator/baseline.py"], - srcs_version = "PY2AND3", - deps = [ - ":expect_tensorflow_estimator_installed", - "//tensorflow/python/estimator", - "//tensorflow/python/estimator:baseline", - ], -) - py_library( name = "boosted_trees", srcs = ["python/estimator/boosted_trees.py"], @@ -60,18 +45,6 @@ py_library( ], ) -py_library( - name = "dnn", - srcs = ["python/estimator/dnn.py"], - srcs_version = "PY2AND3", - deps = [ - ":expect_tensorflow_estimator_installed", - "//tensorflow:tensorflow_py_no_contrib", - "//tensorflow/python/estimator", - "//tensorflow/python/estimator:dnn", - ], -) - py_library( name = "dnn_with_layer_annotations", srcs = ["python/estimator/dnn_with_layer_annotations.py"], @@ -86,18 +59,6 @@ py_library( ], ) -py_library( - name = "dnn_linear_combined", - srcs = ["python/estimator/dnn_linear_combined.py"], - srcs_version = "PY2AND3", - deps = [ - ":expect_tensorflow_estimator_installed", - "//tensorflow:tensorflow_py_no_contrib", - "//tensorflow/python/estimator", - "//tensorflow/python/estimator:dnn_linear_combined", - ], -) - py_library( name = "extenders", srcs = [ @@ -169,17 +130,6 @@ py_library( ], ) -py_library( - name = "linear", - srcs = ["python/estimator/linear.py"], - srcs_version = "PY2AND3", - deps = [ - ":expect_tensorflow_estimator_installed", - "//tensorflow/python/estimator", - "//tensorflow/python/estimator:linear", - ], -) - py_library( name = "logit_fns", srcs = [ diff --git a/tensorflow/contrib/estimator/__init__.py b/tensorflow/contrib/estimator/__init__.py index d9457df4aa3052c0e2e76178c836917daedb6893..7d61247e7ef26d3777843cd3be20684583e9058c 100644 --- a/tensorflow/contrib/estimator/__init__.py +++ b/tensorflow/contrib/estimator/__init__.py @@ -58,10 +58,6 @@ _allowed_symbols = [ 'multi_label_head', 'poisson_regression_head', 'regression_head', - 'BaselineEstimator', - 'DNNEstimator', - 'DNNLinearCombinedEstimator', - 'LinearEstimator', 'boosted_trees_classifier_train_in_memory', 'boosted_trees_regressor_train_in_memory', 'call_logit_fn', diff --git a/tensorflow/contrib/estimator/python/estimator/baseline.py b/tensorflow/contrib/estimator/python/estimator/baseline.py deleted file mode 100644 index fcd320091523cd89c940b9b07777b339ca2212a3..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/estimator/python/estimator/baseline.py +++ /dev/null @@ -1,32 +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. -# ============================================================================== -"""baseline python module. - -Importing from tensorflow.python.estimator is unsupported -and will soon break! -""" -# pylint: disable=unused-import,g-bad-import-order,g-import-not-at-top,wildcard-import - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -from tensorflow_estimator.contrib.estimator.python.estimator import baseline - -# Include attrs that start with single underscore. -_HAS_DYNAMIC_ATTRIBUTES = True -baseline.__all__ = [s for s in dir(baseline) if not s.startswith('__')] - -from tensorflow_estimator.contrib.estimator.python.estimator.baseline import * diff --git a/tensorflow/contrib/estimator/python/estimator/dnn.py b/tensorflow/contrib/estimator/python/estimator/dnn.py deleted file mode 100644 index 10f657df8de64cc96f0cf04f434a77df66629dca..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/estimator/python/estimator/dnn.py +++ /dev/null @@ -1,32 +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. -# ============================================================================== -"""dnn python module. - -Importing from tensorflow.python.estimator is unsupported -and will soon break! -""" -# pylint: disable=unused-import,g-bad-import-order,g-import-not-at-top,wildcard-import - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -from tensorflow_estimator.contrib.estimator.python.estimator import dnn - -# Include attrs that start with single underscore. -_HAS_DYNAMIC_ATTRIBUTES = True -dnn.__all__ = [s for s in dir(dnn) if not s.startswith('__')] - -from tensorflow_estimator.contrib.estimator.python.estimator.dnn import * diff --git a/tensorflow/contrib/estimator/python/estimator/dnn_linear_combined.py b/tensorflow/contrib/estimator/python/estimator/dnn_linear_combined.py deleted file mode 100644 index 7894418c4a16063da88710b43bbbbeb191fc1a2d..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/estimator/python/estimator/dnn_linear_combined.py +++ /dev/null @@ -1,34 +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. -# ============================================================================== -"""dnn_linear_combined python module. - -Importing from tensorflow.python.estimator is unsupported -and will soon break! -""" -# pylint: disable=unused-import,g-bad-import-order,g-import-not-at-top,wildcard-import - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -from tensorflow_estimator.contrib.estimator.python.estimator import dnn_linear_combined - -# Include attrs that start with single underscore. -_HAS_DYNAMIC_ATTRIBUTES = True -dnn_linear_combined.__all__ = [ - s for s in dir(dnn_linear_combined) if not s.startswith('__') -] - -from tensorflow_estimator.contrib.estimator.python.estimator.dnn_linear_combined import * diff --git a/tensorflow/contrib/estimator/python/estimator/linear.py b/tensorflow/contrib/estimator/python/estimator/linear.py deleted file mode 100644 index b6a4444f66c2dd9ce104053613997af1f9c543eb..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/estimator/python/estimator/linear.py +++ /dev/null @@ -1,32 +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. -# ============================================================================== -"""linear python module. - -Importing from tensorflow.python.estimator is unsupported -and will soon break! -""" -# pylint: disable=unused-import,g-bad-import-order,g-import-not-at-top,wildcard-import - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -from tensorflow_estimator.contrib.estimator.python.estimator import linear - -# Include attrs that start with single underscore. -_HAS_DYNAMIC_ATTRIBUTES = True -linear.__all__ = [s for s in dir(linear) if not s.startswith('__')] - -from tensorflow_estimator.contrib.estimator.python.estimator.linear import * diff --git a/tensorflow/contrib/factorization/BUILD b/tensorflow/contrib/factorization/BUILD index e344d7a23b55134612aab430b50cf065bd1095e4..cb86efb8da72f168b54f04773289a6fe421282b1 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", 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/clustering_ops.cc b/tensorflow/contrib/factorization/kernels/clustering_ops.cc deleted file mode 100644 index 025534d540bb82cdb87bb2977d08dfa4f02f1bc8..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/factorization/kernels/clustering_ops.cc +++ /dev/null @@ -1,576 +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. -// ============================================================================== - -#define EIGEN_USE_THREADS - -#include -#include -#include -#include -#include -#include - -#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/lib/core/blocking_counter.h" -#include "tensorflow/core/lib/core/errors.h" -#include "tensorflow/core/lib/core/threadpool.h" -#include "tensorflow/core/lib/gtl/top_n.h" -#include "tensorflow/core/lib/random/philox_random.h" -#include "tensorflow/core/lib/random/simple_philox.h" -#include "tensorflow/core/platform/byte_order.h" -#include "tensorflow/core/platform/cpu_info.h" -#include "tensorflow/core/platform/logging.h" - -namespace tensorflow { -namespace { -using errors::InvalidArgument; - -template -using RowMajorMatrix = - Eigen::Matrix; - -using MatrixXfRowMajor = RowMajorMatrix; -using MatrixXi64RowMajor = RowMajorMatrix; - -// Ideally this should be computed by dividing L3 cache size by the number of -// physical CPUs. Since there isn't a portable method to do this, we are using -// a conservative estimate here. -const int64 kDefaultL3CachePerCpu = 1 << 20; - -// These values were determined by performing a parameter sweep on the -// NearestNeighborsOp benchmark. -const int64 kNearestNeighborsCentersMaxBlockSize = 1024; -const int64 kNearestNeighborsPointsMinBlockSize = 16; - -// Returns the smallest multiple of a that is not smaller than b. -int64 NextMultiple(int64 a, int64 b) { - const int64 remainder = b % a; - return remainder == 0 ? b : (b + a - remainder); -} - -// Returns a / b rounded up to the next higher integer. -int64 CeilOfRatio(int64 a, int64 b) { return (a + b - 1) / b; } - -} // namespace - -// Implementation of K-means++ initialization. Samples points iteratively in -// proportion to the squared distances from selected points. -// TODO(ands): Add support for other distance metrics. -class KmeansPlusPlusInitializationOp : public OpKernel { - public: - explicit KmeansPlusPlusInitializationOp(OpKernelConstruction* context) - : OpKernel(context) { - OP_REQUIRES_OK(context, - context->MatchSignature( - {DT_FLOAT, DT_INT64, DT_INT64, DT_INT64}, {DT_FLOAT})); - } - - void Compute(OpKernelContext* context) override { - const Tensor& points_tensor = context->input(0); - const Tensor& num_to_sample_tensor = context->input(1); - const Tensor& seed_tensor = context->input(2); - const Tensor& num_retries_per_sample_tensor = context->input(3); - - OP_REQUIRES(context, TensorShapeUtils::IsMatrix(points_tensor.shape()), - InvalidArgument("Input points should be a matrix.")); - OP_REQUIRES(context, - TensorShapeUtils::IsScalar(num_to_sample_tensor.shape()), - InvalidArgument("Input num_to_sample should be a scalar.")); - OP_REQUIRES(context, TensorShapeUtils::IsScalar(seed_tensor.shape()), - InvalidArgument("Input seed should be a scalar.")); - OP_REQUIRES( - context, - TensorShapeUtils::IsScalar(num_retries_per_sample_tensor.shape()), - InvalidArgument("Input num_retries_per_sample should be a scalar.")); - - const int64 num_points = points_tensor.dim_size(0); - const int64 point_dimensions = points_tensor.dim_size(1); - const int64 num_to_sample = num_to_sample_tensor.scalar()(); - const int64 seed = seed_tensor.scalar()(); - const int64 num_retries_per_sample = [&]() { - const int64 value = num_retries_per_sample_tensor.scalar()(); - return value >= 0 ? value - : 2 + static_cast(std::log(num_to_sample)); - }(); - - OP_REQUIRES(context, num_points > 0, - InvalidArgument("Expected points.rows() > 0.")); - OP_REQUIRES(context, num_to_sample > 0, - InvalidArgument("Expected num_to_sample > 0.")); - OP_REQUIRES(context, num_to_sample <= num_points, - InvalidArgument("Expected num_to_sample <= points.rows(). ", - num_to_sample, " vs ", num_points, ".")); - - Tensor* output_sampled_points_tensor; - OP_REQUIRES_OK(context, - context->allocate_output( - 0, TensorShape({num_to_sample, point_dimensions}), - &output_sampled_points_tensor)); - - const Eigen::Map points( - points_tensor.matrix().data(), num_points, point_dimensions); - const Eigen::VectorXf points_half_squared_norm = - 0.5 * points.rowwise().squaredNorm(); - - Eigen::Map sampled_points( - output_sampled_points_tensor->matrix().data(), num_to_sample, - point_dimensions); - std::unordered_set sampled_indices; - - random::PhiloxRandom random(seed); - random::SimplePhilox rng(&random); - - auto add_one_point = [&](int64 from, int64 to) { - from = std::min(from, num_points - 1); - sampled_points.row(to) = points.row(from); - sampled_indices.insert(from); - }; - - // Distances from all points to nearest selected point. Initialize with - // distances to first selected point. - Eigen::VectorXf min_distances(num_points); - min_distances.fill(std::numeric_limits::infinity()); - Eigen::VectorXf min_distances_cumsum(num_points); - - auto draw_one_sample = [&]() -> int64 { - if (sampled_indices.empty()) return rng.Uniform64(num_points); - int64 index = 0; - do { - // If v is drawn from Uniform[0, distances.sum()), then - // Prob[cumsum(distances)(i - 1) <= v < cumsum(distances)(i)] is - // proportional to distances(i). - index = std::upper_bound( - min_distances_cumsum.data(), - min_distances_cumsum.data() + num_points, - rng.RandFloat() * min_distances_cumsum(num_points - 1)) - - min_distances_cumsum.data(); - } while (sampled_indices.find(index) != sampled_indices.end()); - return index; - }; - - auto sample_one_point = [&]() { - const int64 sampled_index = draw_one_sample(); - min_distances = min_distances.cwiseMin(GetHalfSquaredDistancesToY( - points, points_half_squared_norm, points.row(sampled_index), - points_half_squared_norm(sampled_index))); - return sampled_index; - }; - - auto sample_one_point_with_retries = [&]() { - Eigen::VectorXf best_new_min_distances(num_points); - float best_potential = std::numeric_limits::infinity(); - int64 best_sampled_index = 0; - for (int i = 1 + num_retries_per_sample; i > 0; --i) { - const int64 sampled_index = draw_one_sample(); - Eigen::VectorXf new_min_distances = - min_distances.cwiseMin(GetHalfSquaredDistancesToY( - points, points_half_squared_norm, points.row(sampled_index), - points_half_squared_norm(sampled_index))); - const float potential = new_min_distances.sum(); - if (potential < best_potential) { - best_potential = potential; - best_sampled_index = sampled_index; - best_new_min_distances.swap(new_min_distances); - } - } - min_distances.swap(best_new_min_distances); - return best_sampled_index; - }; - - for (int64 i = 0; i < num_to_sample; ++i) { - if (i > 0) { - std::partial_sum(min_distances.data(), - min_distances.data() + num_points, - min_distances_cumsum.data()); - } - int64 next = num_retries_per_sample == 0 - ? sample_one_point() - : sample_one_point_with_retries(); - add_one_point(next, i); - } - } - - private: - // Returns a column vector with the i-th element set to half the squared - // euclidean distance between the i-th row of xs, and y. Precomputed norms for - // each row of xs and y must be provided for efficiency. - // TODO(ands): Parallelize this for large xs. - static Eigen::VectorXf GetHalfSquaredDistancesToY( - const Eigen::Ref& xs, - const Eigen::Ref& xs_half_squared_norm, - const Eigen::Ref& y, - float y_half_squared_norm) { - // Squared distance between points xs_i and y is: - // || xs_i ||^2 - 2 + || y ||^2 - return (xs_half_squared_norm - xs * y.transpose()).array() + - y_half_squared_norm; - } -}; - -REGISTER_KERNEL_BUILDER(Name("KmeansPlusPlusInitialization").Device(DEVICE_CPU), - KmeansPlusPlusInitializationOp); - -// Implementation of one single Markov Chain for the k-MC^2 algorithm -class KMC2ChainInitializationOp : public OpKernel { - public: - explicit KMC2ChainInitializationOp(OpKernelConstruction* context) - : OpKernel(context) { - OP_REQUIRES_OK(context, - context->MatchSignature({DT_FLOAT, DT_INT64}, {DT_INT64})); - } - - void Compute(OpKernelContext* context) override { - const Tensor& distances_tensor = context->input(0); - const Tensor& seed_tensor = context->input(1); - OP_REQUIRES(context, TensorShapeUtils::IsVector(distances_tensor.shape()), - InvalidArgument("Input distances should be a vector.")); - OP_REQUIRES(context, TensorShapeUtils::IsScalar(seed_tensor.shape()), - InvalidArgument("Input seed should be a scalar.")); - const int64 num_points = distances_tensor.dim_size(0); - const int64 seed = seed_tensor.scalar()(); - OP_REQUIRES(context, num_points > 0, - InvalidArgument("Expected distances_tensor.size() > 0.")); - - random::PhiloxRandom random(seed); - random::SimplePhilox rng(&random); - - auto distances = distances_tensor.flat(); - // Set the initial state of the Markov chain to be the first candidate. - int64 selected_index = 0; - float selected_distance = distances(selected_index); - // Build a Markov chain of length num_points. - for (int64 i = 1; i < num_points; ++i) { - const float candidate_distance = distances(i); - // Set the next state of the Markov chain to be the candidate with - // probability min(1, candidate_distance/selected_distance). - if (candidate_distance > rng.RandFloat() * selected_distance) { - selected_index = i; - selected_distance = candidate_distance; - } - } - - Tensor* output_sampled_index_tensor; - OP_REQUIRES_OK(context, - context->allocate_output(0, TensorShape({}), - &output_sampled_index_tensor)); - auto output = output_sampled_index_tensor->scalar(); - // Return the last state of the Markov chain as the new center. - output() = selected_index; - } -}; - -REGISTER_KERNEL_BUILDER(Name("KMC2ChainInitialization").Device(DEVICE_CPU), - KMC2ChainInitializationOp); - -// Operator for computing the nearest neighbors for a set of points. -class NearestNeighborsOp : public OpKernel { - public: - explicit NearestNeighborsOp(OpKernelConstruction* context) - : OpKernel(context) { - OP_REQUIRES_OK(context, - context->MatchSignature({DT_FLOAT, DT_FLOAT, DT_INT64}, - {DT_INT64, DT_FLOAT})); - } - - void Compute(OpKernelContext* context) override { - const Tensor& points_tensor = context->input(0); - const Tensor& centers_tensor = context->input(1); - const Tensor& k_tensor = context->input(2); - - OP_REQUIRES(context, TensorShapeUtils::IsMatrix(points_tensor.shape()), - InvalidArgument("Input points should be a matrix.")); - OP_REQUIRES(context, TensorShapeUtils::IsMatrix(centers_tensor.shape()), - InvalidArgument("Input centers should be a matrix.")); - OP_REQUIRES(context, TensorShapeUtils::IsScalar(k_tensor.shape()), - InvalidArgument("Input k should be a scalar.")); - - const int64 num_points = points_tensor.dim_size(0); - const int64 point_dimensions = points_tensor.dim_size(1); - const int64 num_centers = centers_tensor.dim_size(0); - const int64 center_dimensions = centers_tensor.dim_size(1); - - OP_REQUIRES(context, num_points > 0, - InvalidArgument("Expected points.rows() > 0.")); - OP_REQUIRES( - context, point_dimensions == center_dimensions, - InvalidArgument("Expected point_dimensions == center_dimensions: ", - point_dimensions, " vs ", center_dimensions, ".")); - - const Eigen::Map points( - points_tensor.matrix().data(), num_points, point_dimensions); - const Eigen::Map centers( - centers_tensor.matrix().data(), num_centers, center_dimensions); - const int64 k = std::min(num_centers, k_tensor.scalar()()); - - Tensor* output_nearest_center_indices_tensor; - Tensor* output_nearest_center_distances_tensor; - OP_REQUIRES_OK(context, context->allocate_output( - 0, TensorShape({num_points, k}), - &output_nearest_center_indices_tensor)); - OP_REQUIRES_OK(context, context->allocate_output( - 1, TensorShape({num_points, k}), - &output_nearest_center_distances_tensor)); - - if (k == 0) return; - - Eigen::Map nearest_center_indices( - output_nearest_center_indices_tensor->matrix().data(), - num_points, k); - Eigen::Map nearest_center_distances( - output_nearest_center_distances_tensor->matrix().data(), - num_points, k); - - const Eigen::VectorXf centers_half_squared_norm = - 0.5 * centers.rowwise().squaredNorm(); - - // The distance computation is sharded to take advantage of multiple cores - // and to allow intermediate values to reside in L3 cache. This is done by - // sharding the points and centers as follows: - // - // 1. Centers are sharded such that each block of centers has at most - // kNearestNeighborsCentersMaxBlockSize rows. - // 2. Points are sharded, and each block of points is multiplied with each - // block of centers. The block size of points is chosen such that the - // point coordinates (point_dimensions) and the matrix of distances to - // each center in one block -- the intermediate data -- fits in L3 cache. - // 3. After performing each block-block distance computation, the results - // are reduced to a set of k nearest centers as soon as possible. This - // decreases total memory I/O. - auto worker_threads = *(context->device()->tensorflow_cpu_worker_threads()); - const int64 num_threads = worker_threads.num_threads; - // This kernel might be configured to use fewer than the total number of - // available CPUs on the host machine. To avoid destructive interference - // with other jobs running on the host machine, we must only use a fraction - // of total available L3 cache. Unfortunately, we cannot query the host - // machine to get the number of physical CPUs. So, we use a fixed per-CPU - // budget and scale it by the number of CPUs available to this operation. - const int64 total_memory_budget = - kDefaultL3CachePerCpu * port::NumSchedulableCPUs(); - // Compute the number of blocks into which rows of points must be split so - // that the distance matrix and the block of points can fit in cache. One - // row of points will yield a vector of distances to each center in a block. - const int64 bytes_per_row = - (std::min(kNearestNeighborsCentersMaxBlockSize, - num_centers) /* centers in a block */ - + point_dimensions /* coordinates of one point */) * - sizeof(float); - // The memory needed for storing the centers being processed. This is shared - // by all workers. Adding slack to the number of threads to avoid incorrect - // cache eviction when a new block of centers is loaded. - const int64 bytes_for_centers = - std::min(num_centers, - (num_threads + 2) * kNearestNeighborsCentersMaxBlockSize) * - point_dimensions * sizeof(float); - // The memory budget available for workers to store their distance matrices. - const int64 available_memory_budget = - total_memory_budget - bytes_for_centers; - // That memory budget is shared by all threads. - const int64 rows_per_block = - std::max(kNearestNeighborsPointsMinBlockSize, - available_memory_budget / num_threads / bytes_per_row); - // Divide rows into almost uniformly-sized units of work that are small - // enough for the memory budget (rows_per_block). Round up to a multiple of - // the number of threads. - const int64 num_units = - NextMultiple(num_threads, CeilOfRatio(num_points, rows_per_block)); - auto work = [&](int64 start, int64 limit) { - for (; start < limit; ++start) { - const int64 start_row = num_points * start / num_units; - const int64 limit_row = num_points * (start + 1) / num_units; - CHECK_LE(limit_row, num_points); - const int64 num_rows = limit_row - start_row; - auto points_shard = points.middleRows(start_row, num_rows); - const Eigen::VectorXf points_half_squared_norm = - 0.5 * points_shard.rowwise().squaredNorm(); - auto nearest_center_indices_shard = - nearest_center_indices.middleRows(start_row, num_rows); - auto nearest_center_distances_shard = - nearest_center_distances.middleRows(start_row, num_rows); - FindKNearestCenters(k, points_shard, points_half_squared_norm, centers, - centers_half_squared_norm, - nearest_center_indices_shard, - nearest_center_distances_shard); - } - }; - - const int64 units_per_thread = num_units / num_threads; - BlockingCounter counter(num_threads - 1); - for (int64 i = 1; i < num_threads; ++i) { - const int64 start = i * units_per_thread; - const int64 limit = start + units_per_thread; - worker_threads.workers->Schedule([work, &counter, start, limit]() { - work(start, limit); - counter.DecrementCount(); - }); - } - work(0, units_per_thread); - counter.Wait(); - } - - private: - static void FindKNearestCenters( - int64 k, const Eigen::Ref& points, - const Eigen::Ref& points_half_squared_norm, - const Eigen::Ref& centers, - const Eigen::Ref& centers_half_squared_norm, - const Eigen::Ref& nearest_center_indices, - const Eigen::Ref& nearest_center_distances) { - CHECK_LE(k, centers.rows()); - if (centers.rows() <= kNearestNeighborsCentersMaxBlockSize) { - FindKNearestCentersOneBlock(k, points, points_half_squared_norm, centers, - centers_half_squared_norm, - nearest_center_indices, - nearest_center_distances); - } else { - FindKNearestCentersBlockwise(k, points, points_half_squared_norm, centers, - centers_half_squared_norm, - nearest_center_indices, - nearest_center_distances); - } - } - - static void FindKNearestCentersOneBlock( - int64 k, const Eigen::Ref& points, - const Eigen::Ref& points_half_squared_norm, - const Eigen::Ref& centers, - const Eigen::Ref& centers_half_squared_norm, - Eigen::Ref nearest_center_indices, - Eigen::Ref nearest_center_distances) { - CHECK_LE(k, centers.rows()); - const int64 num_points = points.rows(); - const MatrixXfRowMajor inner_product = points * centers.transpose(); - // Find nearest neighbors. - if (k == 1) { - for (int i = 0; i < num_points; ++i) { - int64 index; - nearest_center_distances(i, 0) = - 2.0 * - (points_half_squared_norm(i) + - (centers_half_squared_norm.transpose() - inner_product.row(i)) - .minCoeff(&index)); - nearest_center_indices(i, 0) = index; - } - } else { - // Select k nearest centers for each point. - using Center = std::pair; - const int64 num_centers = centers.rows(); - gtl::TopN> selector(k); - std::unique_ptr> nearest_centers; - for (int i = 0; i < num_points; ++i) { - selector.reserve(num_centers); - for (int j = 0; j < num_centers; ++j) { - const float partial_distance = - centers_half_squared_norm(j) - inner_product(i, j); - selector.push(Center(partial_distance, j)); - } - nearest_centers.reset(selector.Extract()); - selector.Reset(); - const float point_half_squared_norm = points_half_squared_norm(i); - for (int j = 0; j < k; ++j) { - const Center& center = (*nearest_centers)[j]; - nearest_center_distances(i, j) = - 2.0 * (point_half_squared_norm + center.first); - nearest_center_indices(i, j) = center.second; - } - } - } - } - - static void FindKNearestCentersBlockwise( - int64 k, const Eigen::Ref& points, - const Eigen::Ref& points_half_squared_norm, - const Eigen::Ref& centers, - const Eigen::Ref& centers_half_squared_norm, - Eigen::Ref nearest_center_indices, - Eigen::Ref nearest_center_distances) { - const int64 num_points = points.rows(); - const int64 num_centers = centers.rows(); - CHECK_LE(k, num_centers); - CHECK_GT(num_centers, kNearestNeighborsCentersMaxBlockSize); - // Store nearest neighbors with first block of centers directly into the - // output matrices. - int64 out_k = std::min(k, kNearestNeighborsCentersMaxBlockSize); - FindKNearestCentersOneBlock( - out_k, points, points_half_squared_norm, - centers.topRows(kNearestNeighborsCentersMaxBlockSize), - centers_half_squared_norm.head(kNearestNeighborsCentersMaxBlockSize), - nearest_center_indices, nearest_center_distances); - // Iteratively compute nearest neighbors with other blocks of centers, and - // update the output matrices. - MatrixXi64RowMajor block_nearest_center_indices(num_points, k); - MatrixXfRowMajor block_nearest_center_distances(num_points, k); - Eigen::Matrix merged_indices(k); - Eigen::Matrix merged_distances(k); - for (int64 centers_start = kNearestNeighborsCentersMaxBlockSize; - centers_start < num_centers; - centers_start += kNearestNeighborsCentersMaxBlockSize) { - const int64 centers_block_size = std::min( - kNearestNeighborsCentersMaxBlockSize, num_centers - centers_start); - const int64 block_k = std::min(k, centers_block_size); - FindKNearestCentersOneBlock( - block_k, points, points_half_squared_norm, - centers.middleRows(centers_start, centers_block_size), - centers_half_squared_norm.segment(centers_start, centers_block_size), - block_nearest_center_indices, block_nearest_center_distances); - if (k == 1) { - for (int i = 0; i < num_points; ++i) { - if (block_nearest_center_distances(i, 0) < - nearest_center_distances(i, 0)) { - nearest_center_indices(i, 0) = - block_nearest_center_indices(i, 0) + centers_start; - nearest_center_distances(i, 0) = - block_nearest_center_distances(i, 0); - } - } - } else { - for (int i = 0; i < num_points; ++i) { - // Merge and accumulate top-k list from block_nearest_center_indices - // into nearest_center_indices. - for (int64 j_out = 0, j_block = 0, j_merged = 0; - (j_out < out_k || j_block < block_k) && j_merged < k; - ++j_merged) { - const float distance_out = - j_out < out_k ? nearest_center_distances(i, j_out) - : std::numeric_limits::infinity(); - const float distance_block = - j_block < block_k ? block_nearest_center_distances(i, j_block) - : std::numeric_limits::infinity(); - if (distance_out <= distance_block) { - merged_indices(j_merged) = nearest_center_indices(i, j_out); - merged_distances(j_merged) = distance_out; - ++j_out; - } else { - merged_indices(j_merged) = - block_nearest_center_indices(i, j_block) + centers_start; - merged_distances(j_merged) = distance_block; - ++j_block; - } - } - nearest_center_indices.row(i) = merged_indices; - nearest_center_distances.row(i) = merged_distances; - out_k = std::min(k, out_k + block_k); - } - } - } - } -}; - -REGISTER_KERNEL_BUILDER(Name("NearestNeighbors").Device(DEVICE_CPU), - NearestNeighborsOp); - -} // namespace tensorflow 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/kmeans.py b/tensorflow/contrib/factorization/python/ops/kmeans.py index f384d761a8430074f022c973d7ec3d46cd90f70b..3eb396a29ccdc0478384f9fa122465731740a30d 100644 --- a/tensorflow/contrib/factorization/python/ops/kmeans.py +++ b/tensorflow/contrib/factorization/python/ops/kmeans.py @@ -26,7 +26,7 @@ from tensorflow.contrib.factorization.python.ops import clustering_ops from tensorflow.python.estimator import estimator from tensorflow.python.estimator import model_fn as model_fn_lib from tensorflow.python.estimator.export import export_output -from tensorflow.python.feature_column import feature_column as fc +from tensorflow.python.feature_column import feature_column_lib as fc from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import control_flow_ops diff --git a/tensorflow/contrib/factorization/python/ops/kmeans_test.py b/tensorflow/contrib/factorization/python/ops/kmeans_test.py index 1ab5418fe4659cb0068ee8c3ca1442f6f723ee76..2f7cd131d3ed20df307ed231cce2ecb50ecfbceb 100644 --- a/tensorflow/contrib/factorization/python/ops/kmeans_test.py +++ b/tensorflow/contrib/factorization/python/ops/kmeans_test.py @@ -27,7 +27,7 @@ from sklearn.cluster import KMeans as SklearnKMeans # pylint: disable=g-import-not-at-top from tensorflow.contrib.factorization.python.ops import kmeans as kmeans_lib from tensorflow.python.estimator import run_config -from tensorflow.python.feature_column import feature_column as fc +from tensorflow.python.feature_column import feature_column_lib as fc from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops diff --git a/tensorflow/contrib/feature_column/BUILD b/tensorflow/contrib/feature_column/BUILD index a926ffd5982116a21dc7a0fd1ff957d4ecc6bf94..8fc5f1cfe7800653ef1e43c6d40d1a66e34f2106 100644 --- a/tensorflow/contrib/feature_column/BUILD +++ b/tensorflow/contrib/feature_column/BUILD @@ -6,7 +6,7 @@ package( licenses(["notice"]) # Apache 2.0 -load("//tensorflow:tensorflow.bzl", "py_test") +load("//tensorflow:tensorflow.bzl", "tf_py_test") py_library( name = "feature_column_py", @@ -14,6 +14,7 @@ py_library( srcs_version = "PY2AND3", deps = [ ":sequence_feature_column", + ":sequence_feature_column_v2", "//tensorflow/python:util", ], ) @@ -32,17 +33,17 @@ py_library( "//tensorflow/python:sparse_ops", "//tensorflow/python:tensor_shape", "//tensorflow/python:variable_scope", - "//tensorflow/python/feature_column", + "//tensorflow/python/feature_column:feature_column_py", ], ) -py_test( +tf_py_test( name = "sequence_feature_column_test", srcs = ["python/feature_column/sequence_feature_column_test.py"], - srcs_version = "PY2AND3", - tags = ["no_pip"], - deps = [ + additional_deps = [ ":sequence_feature_column", + "@absl_py//absl/testing:parameterized", + "//third_party/py/numpy", "//tensorflow/python:client_testlib", "//tensorflow/python:dtypes", "//tensorflow/python:errors", @@ -51,25 +52,80 @@ py_test( "//tensorflow/python:parsing_ops", "//tensorflow/python:sparse_tensor", "//tensorflow/python:training", + "//tensorflow/python/feature_column:feature_column_py", + ], + tags = ["no_pip"], +) + +tf_py_test( + name = "sequence_feature_column_integration_test", + srcs = ["python/feature_column/sequence_feature_column_integration_test.py"], + additional_deps = [ + ":sequence_feature_column", + "//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", + ], + tags = ["no_pip"], +) + +py_library( + name = "sequence_feature_column_v2", + srcs = ["python/feature_column/sequence_feature_column_v2.py"], + srcs_version = "PY2AND3", + deps = [ + "//tensorflow/python:array_ops", + "//tensorflow/python:check_ops", + "//tensorflow/python:dtypes", + "//tensorflow/python:framework_ops", + "//tensorflow/python:math_ops", + "//tensorflow/python:parsing_ops", + "//tensorflow/python:sparse_ops", + "//tensorflow/python:tensor_shape", + "//tensorflow/python:variable_scope", "//tensorflow/python/feature_column", - "//third_party/py/numpy", + "//tensorflow/python/feature_column:feature_column_py", + ], +) + +tf_py_test( + name = "sequence_feature_column_v2_test", + srcs = ["python/feature_column/sequence_feature_column_v2_test.py"], + additional_deps = [ + ":sequence_feature_column_v2", "@absl_py//absl/testing:parameterized", + "//third_party/py/numpy", + "//tensorflow/python:client_testlib", + "//tensorflow/python:dtypes", + "//tensorflow/python:errors", + "//tensorflow/python:framework_ops", + "//tensorflow/python:math_ops", + "//tensorflow/python:parsing_ops", + "//tensorflow/python:sparse_tensor", + "//tensorflow/python:training", + "//tensorflow/python/feature_column:feature_column_py", + "//tensorflow/python/feature_column:feature_column_v2_test", ], + tags = ["no_pip"], ) py_test( - name = "sequence_feature_column_integration_test", - srcs = ["python/feature_column/sequence_feature_column_integration_test.py"], + 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", + ":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", + "//tensorflow/python/feature_column:feature_column_py", "//tensorflow/python/keras:layers", ], ) diff --git a/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column.py b/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column.py index dd6da35ed009c07ad3819e7860a283c7837c1f83..9b3a5c58aaa9498257fc971ac60b97f31d5185d8 100644 --- a/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column.py +++ b/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column.py @@ -222,10 +222,8 @@ def sequence_categorical_column_with_identity( ValueError: if `default_value` is not in range `[0, num_buckets)`. """ return fc._SequenceCategoricalColumn( - fc.categorical_column_with_identity( - key=key, - num_buckets=num_buckets, - default_value=default_value)) + fc._categorical_column_with_identity( + key=key, num_buckets=num_buckets, default_value=default_value)) def sequence_categorical_column_with_hash_bucket( @@ -265,10 +263,8 @@ def sequence_categorical_column_with_hash_bucket( ValueError: `dtype` is neither string nor integer. """ return fc._SequenceCategoricalColumn( - fc.categorical_column_with_hash_bucket( - key=key, - hash_bucket_size=hash_bucket_size, - dtype=dtype)) + fc._categorical_column_with_hash_bucket( + key=key, hash_bucket_size=hash_bucket_size, dtype=dtype)) def sequence_categorical_column_with_vocabulary_file( @@ -324,7 +320,7 @@ def sequence_categorical_column_with_vocabulary_file( ValueError: `dtype` is neither string nor integer. """ return fc._SequenceCategoricalColumn( - fc.categorical_column_with_vocabulary_file( + fc._categorical_column_with_vocabulary_file( key=key, vocabulary_file=vocabulary_file, vocabulary_size=vocabulary_size, @@ -384,7 +380,7 @@ def sequence_categorical_column_with_vocabulary_list( ValueError: if `dtype` is not integer or string. """ return fc._SequenceCategoricalColumn( - fc.categorical_column_with_vocabulary_list( + fc._categorical_column_with_vocabulary_list( key=key, vocabulary_list=vocabulary_list, dtype=dtype, diff --git a/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column_integration_test.py b/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column_integration_test.py index d8ca363627eace15e039679545366648df174c33..bcc25b8de895a769f9e11b207c2092e23d029b1f 100644 --- a/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column_integration_test.py +++ b/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column_integration_test.py @@ -53,19 +53,20 @@ class SequenceFeatureColumnIntegrationTest(test.TestCase): return example def _build_feature_columns(self): - col = fc.categorical_column_with_identity( - 'int_ctx', num_buckets=100) + col = fc._categorical_column_with_identity('int_ctx', num_buckets=100) ctx_cols = [ - fc.embedding_column(col, dimension=10), - fc.numeric_column('float_ctx')] + 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)] + fc._embedding_column(identity_col, dimension=10), + fc._embedding_column(bucket_col, dimension=20) + ] return ctx_cols, seq_cols @@ -148,8 +149,8 @@ class SequenceExampleParsingTest(test.TestCase): """ example = _make_sequence_example() columns = [ - fc.categorical_column_with_identity('int_ctx', num_buckets=100), - fc.numeric_column('float_ctx'), + 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( diff --git a/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column_test.py b/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column_test.py index 2163af0b43864c96483df529f07881f2f985a80e..d5f74028298ee7015f5b2e3aaee7d9330c1acac1 100644 --- a/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column_test.py +++ b/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column_test.py @@ -24,6 +24,7 @@ import numpy as np from tensorflow.contrib.feature_column.python.feature_column import sequence_feature_column as sfc from tensorflow.python.feature_column import feature_column as fc +from tensorflow.python.feature_column import feature_column_lib as fc_lib from tensorflow.python.feature_column.feature_column import _LazyBuilder from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors @@ -109,13 +110,15 @@ 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.embedding_column( - categorical_column_a, dimension=embedding_dimension_a, + 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.embedding_column( - categorical_column_b, dimension=embedding_dimension_b, + 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( @@ -148,10 +151,9 @@ class SequenceInputLayerTest(test.TestCase, parameterized.TestCase): values=(2, 0, 1), dense_shape=(2, 2)) - categorical_column_a = fc.categorical_column_with_identity( + categorical_column_a = fc._categorical_column_with_identity( key='aaa', num_buckets=vocabulary_size) - embedding_column_a = fc.embedding_column( - categorical_column_a, dimension=2) + embedding_column_a = fc._embedding_column(categorical_column_a, dimension=2) with self.assertRaisesRegexp( ValueError, @@ -206,7 +208,7 @@ 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_lib.shared_embedding_columns( [categorical_column_b, categorical_column_a], dimension=embedding_dimension, initializer=_get_initializer(embedding_dimension, embedding_values)) @@ -244,11 +246,11 @@ class SequenceInputLayerTest(test.TestCase, parameterized.TestCase): values=(2, 0, 1), dense_shape=(2, 2)) - categorical_column_a = fc.categorical_column_with_identity( + categorical_column_a = fc._categorical_column_with_identity( key='aaa', num_buckets=vocabulary_size) - categorical_column_b = fc.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_lib.shared_embedding_columns( [categorical_column_a, categorical_column_b], dimension=2) with self.assertRaisesRegexp( @@ -315,10 +317,10 @@ 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.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.indicator_column(categorical_column_b) + indicator_column_b = fc._indicator_column(categorical_column_b) input_layer, sequence_length = sfc.sequence_input_layer( features={ 'aaa': sparse_input_a, @@ -342,9 +344,9 @@ class SequenceInputLayerTest(test.TestCase, parameterized.TestCase): values=(2, 0, 1), dense_shape=(2, 2)) - categorical_column_a = fc.categorical_column_with_identity( + categorical_column_a = fc._categorical_column_with_identity( key='aaa', num_buckets=vocabulary_size) - indicator_column_a = fc.indicator_column(categorical_column_a) + indicator_column_a = fc._indicator_column(categorical_column_a) with self.assertRaisesRegexp( ValueError, @@ -530,7 +532,7 @@ 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.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]) @@ -616,8 +618,7 @@ class InputLayerTest(test.TestCase): categorical_column_a = sfc.sequence_categorical_column_with_identity( key='aaa', num_buckets=vocabulary_size) - embedding_column_a = fc.embedding_column( - categorical_column_a, dimension=2) + embedding_column_a = fc._embedding_column(categorical_column_a, dimension=2) with self.assertRaisesRegexp( ValueError, @@ -639,7 +640,7 @@ class InputLayerTest(test.TestCase): categorical_column_a = sfc.sequence_categorical_column_with_identity( key='aaa', num_buckets=vocabulary_size) - indicator_column_a = fc.indicator_column(categorical_column_a) + indicator_column_a = fc._indicator_column(categorical_column_a) with self.assertRaisesRegexp( ValueError, @@ -918,8 +919,9 @@ class SequenceEmbeddingColumnTest( categorical_column = sfc.sequence_categorical_column_with_identity( key='aaa', num_buckets=vocabulary_size) - embedding_column = fc.embedding_column( - categorical_column, dimension=embedding_dimension, + embedding_column = fc._embedding_column( + categorical_column, + dimension=embedding_dimension, initializer=_initializer) embedding_lookup, _ = embedding_column._get_sequence_dense_tensor( @@ -956,8 +958,7 @@ class SequenceEmbeddingColumnTest( categorical_column = sfc.sequence_categorical_column_with_identity( key='aaa', num_buckets=vocabulary_size) - embedding_column = fc.embedding_column( - categorical_column, dimension=2) + embedding_column = fc._embedding_column(categorical_column, dimension=2) _, sequence_length = embedding_column._get_sequence_dense_tensor( _LazyBuilder({'aaa': inputs})) @@ -984,8 +985,7 @@ class SequenceEmbeddingColumnTest( categorical_column = sfc.sequence_categorical_column_with_identity( key='aaa', num_buckets=vocabulary_size) - embedding_column = fc.embedding_column( - categorical_column, dimension=2) + embedding_column = fc._embedding_column(categorical_column, dimension=2) _, sequence_length = embedding_column._get_sequence_dense_tensor( _LazyBuilder({'aaa': sparse_input})) @@ -1055,7 +1055,7 @@ class SequenceSharedEmbeddingColumnTest(test.TestCase): key='aaa', num_buckets=vocabulary_size) categorical_column_b = sfc.sequence_categorical_column_with_identity( key='bbb', num_buckets=vocabulary_size) - shared_embedding_columns = fc.shared_embedding_columns( + shared_embedding_columns = fc_lib.shared_embedding_columns( [categorical_column_a, categorical_column_b], dimension=embedding_dimension, initializer=_initializer) @@ -1101,7 +1101,7 @@ class SequenceSharedEmbeddingColumnTest(test.TestCase): expected_sequence_length_b = [2, 1] categorical_column_b = sfc.sequence_categorical_column_with_identity( key='bbb', num_buckets=vocabulary_size) - shared_embedding_columns = fc.shared_embedding_columns( + shared_embedding_columns = fc_lib.shared_embedding_columns( [categorical_column_a, categorical_column_b], dimension=2) sequence_length_a = shared_embedding_columns[0]._get_sequence_dense_tensor( @@ -1152,7 +1152,7 @@ class SequenceSharedEmbeddingColumnTest(test.TestCase): categorical_column_b = sfc.sequence_categorical_column_with_identity( key='bbb', num_buckets=vocabulary_size) - shared_embedding_columns = fc.shared_embedding_columns( + shared_embedding_columns = fc_lib.shared_embedding_columns( [categorical_column_a, categorical_column_b], dimension=2) sequence_length_a = shared_embedding_columns[0]._get_sequence_dense_tensor( @@ -1218,7 +1218,7 @@ class SequenceIndicatorColumnTest(test.TestCase, parameterized.TestCase): categorical_column = sfc.sequence_categorical_column_with_identity( key='aaa', num_buckets=vocabulary_size) - indicator_column = fc.indicator_column(categorical_column) + indicator_column = fc._indicator_column(categorical_column) indicator_tensor, _ = indicator_column._get_sequence_dense_tensor( _LazyBuilder({'aaa': inputs})) @@ -1250,7 +1250,7 @@ class SequenceIndicatorColumnTest(test.TestCase, parameterized.TestCase): categorical_column = sfc.sequence_categorical_column_with_identity( key='aaa', num_buckets=vocabulary_size) - indicator_column = fc.indicator_column(categorical_column) + indicator_column = fc._indicator_column(categorical_column) _, sequence_length = indicator_column._get_sequence_dense_tensor( _LazyBuilder({'aaa': inputs})) @@ -1277,7 +1277,7 @@ class SequenceIndicatorColumnTest(test.TestCase, parameterized.TestCase): categorical_column = sfc.sequence_categorical_column_with_identity( key='aaa', num_buckets=vocabulary_size) - indicator_column = fc.indicator_column(categorical_column) + indicator_column = fc._indicator_column(categorical_column) _, sequence_length = indicator_column._get_sequence_dense_tensor( _LazyBuilder({'aaa': sparse_input})) 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 new file mode 100644 index 0000000000000000000000000000000000000000..2f4bda194a41242167e0abfcaeac5044f6026f85 --- /dev/null +++ b/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column_v2.py @@ -0,0 +1,582 @@ +# 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. +# ============================================================================== +"""This API defines FeatureColumn for sequential input. + +NOTE: This API is a work in progress and will likely be changing frequently. +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + + +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 +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 + +# pylint: disable=protected-access + + +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. + + 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]`. + + 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] + + 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.keras.layers.SimpleRNNCell(hidden_size) + rnn_layer = tf.keras.layers.RNN(rnn_cell) + outputs, state = rnn_layer(sequence_input, mask=sequence_length_mask) + ``` + """ + + def __init__( + self, + feature_columns, + trainable=True, + name=None, + **kwargs): + """"Constructs a SequenceFeatures layer. + + 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) + + 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 = [] + + 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. + output_tensors.append(self._process_dense_tensor(column, dense_tensor)) + sequence_lengths.append(sequence_length) + + # 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 self._verify_and_concat_tensors(output_tensors), sequence_length + + +def concatenate_context_input(context_input, sequence_input): + """Replicates `context_input` across all timesteps of `sequence_input`. + + Expands dimension 1 of `context_input` then tiles it `sequence_length` times. + This value is appended to `sequence_input` on dimension 2 and the result is + returned. + + Args: + context_input: A `Tensor` of dtype `float32` and shape `[batch_size, d1]`. + sequence_input: A `Tensor` of dtype `float32` and shape `[batch_size, + padded_length, d0]`. + + Returns: + A `Tensor` of dtype `float32` and shape `[batch_size, padded_length, + d0 + d1]`. + + Raises: + ValueError: If `sequence_input` does not have rank 3 or `context_input` does + not have rank 2. + """ + seq_rank_check = check_ops.assert_rank( + sequence_input, + 3, + message='sequence_input must have rank 3', + data=[array_ops.shape(sequence_input)]) + seq_type_check = check_ops.assert_type( + sequence_input, + dtypes.float32, + message='sequence_input must have dtype float32; got {}.'.format( + sequence_input.dtype)) + ctx_rank_check = check_ops.assert_rank( + context_input, + 2, + message='context_input must have rank 2', + data=[array_ops.shape(context_input)]) + ctx_type_check = check_ops.assert_type( + context_input, + dtypes.float32, + message='context_input must have dtype float32; got {}.'.format( + context_input.dtype)) + with ops.control_dependencies( + [seq_rank_check, seq_type_check, ctx_rank_check, ctx_type_check]): + padded_length = array_ops.shape(sequence_input)[1] + tiled_context_input = array_ops.tile( + array_ops.expand_dims(context_input, 1), + array_ops.concat([[1], [padded_length], [1]], 0)) + return array_ops.concat([sequence_input, tiled_context_input], 2) + + +def sequence_categorical_column_with_identity( + key, num_buckets, default_value=None): + """Returns a feature column that represents sequences of integers. + + Pass this to `embedding_column` or `indicator_column` to convert sequence + categorical data into dense representation for input to sequence NN, such as + RNN. + + Example: + + ```python + watches = sequence_categorical_column_with_identity( + 'watches', num_buckets=1000) + watches_embedding = embedding_column(watches, dimension=10) + columns = [watches_embedding] + + features = tf.parse_example(..., features=make_parse_example_spec(columns)) + sequence_feature_layer = SequenceFeatures(columns) + sequence_input, sequence_length = sequence_feature_layer(features) + sequence_length_mask = tf.sequence_mask(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: + key: A unique string identifying the input feature. + num_buckets: Range of inputs. Namely, inputs are expected to be in the + range `[0, num_buckets)`. + default_value: If `None`, this column's graph operations will fail for + out-of-range inputs. Otherwise, this value must be in the range + `[0, num_buckets)`, and will replace out-of-range inputs. + + Returns: + A `SequenceCategoricalColumn`. + + Raises: + ValueError: if `num_buckets` is less than one. + ValueError: if `default_value` is not in range `[0, num_buckets)`. + """ + return fc.SequenceCategoricalColumn( + fc.categorical_column_with_identity( + key=key, + num_buckets=num_buckets, + default_value=default_value)) + + +def sequence_categorical_column_with_hash_bucket( + key, hash_bucket_size, dtype=dtypes.string): + """A sequence of categorical terms where ids are set by hashing. + + Pass this to `embedding_column` or `indicator_column` to convert sequence + categorical data into dense representation for input to sequence NN, such as + RNN. + + Example: + + ```python + tokens = sequence_categorical_column_with_hash_bucket( + 'tokens', hash_bucket_size=1000) + tokens_embedding = embedding_column(tokens, dimension=10) + columns = [tokens_embedding] + + features = tf.parse_example(..., features=make_parse_example_spec(columns)) + sequence_feature_layer = SequenceFeatures(columns) + sequence_input, sequence_length = sequence_feature_layer(features) + sequence_length_mask = tf.sequence_mask(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: + key: A unique string identifying the input feature. + hash_bucket_size: An int > 1. The number of buckets. + dtype: The type of features. Only string and integer types are supported. + + Returns: + A `SequenceCategoricalColumn`. + + Raises: + ValueError: `hash_bucket_size` is not greater than 1. + ValueError: `dtype` is neither string nor integer. + """ + return fc.SequenceCategoricalColumn( + fc.categorical_column_with_hash_bucket( + key=key, + hash_bucket_size=hash_bucket_size, + dtype=dtype)) + + +def sequence_categorical_column_with_vocabulary_file( + key, vocabulary_file, vocabulary_size=None, num_oov_buckets=0, + default_value=None, dtype=dtypes.string): + """A sequence of categorical terms where ids use a vocabulary file. + + Pass this to `embedding_column` or `indicator_column` to convert sequence + categorical data into dense representation for input to sequence NN, such as + RNN. + + Example: + + ```python + states = sequence_categorical_column_with_vocabulary_file( + key='states', vocabulary_file='/us/states.txt', vocabulary_size=50, + num_oov_buckets=5) + states_embedding = embedding_column(states, dimension=10) + columns = [states_embedding] + + features = tf.parse_example(..., features=make_parse_example_spec(columns)) + sequence_feature_layer = SequenceFeatures(columns) + sequence_input, sequence_length = sequence_feature_layer(features) + sequence_length_mask = tf.sequence_mask(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: + key: A unique string identifying the input feature. + vocabulary_file: The vocabulary file name. + vocabulary_size: Number of the elements in the vocabulary. This must be no + greater than length of `vocabulary_file`, if less than length, later + values are ignored. If None, it is set to the length of `vocabulary_file`. + num_oov_buckets: Non-negative integer, the number of out-of-vocabulary + buckets. All out-of-vocabulary inputs will be assigned IDs in the range + `[vocabulary_size, vocabulary_size+num_oov_buckets)` based on a hash of + the input value. A positive `num_oov_buckets` can not be specified with + `default_value`. + default_value: The integer ID value to return for out-of-vocabulary feature + values, defaults to `-1`. This can not be specified with a positive + `num_oov_buckets`. + dtype: The type of features. Only string and integer types are supported. + + Returns: + A `SequenceCategoricalColumn`. + + Raises: + ValueError: `vocabulary_file` is missing or cannot be opened. + ValueError: `vocabulary_size` is missing or < 1. + ValueError: `num_oov_buckets` is a negative integer. + ValueError: `num_oov_buckets` and `default_value` are both specified. + ValueError: `dtype` is neither string nor integer. + """ + return fc.SequenceCategoricalColumn( + fc.categorical_column_with_vocabulary_file( + key=key, + vocabulary_file=vocabulary_file, + vocabulary_size=vocabulary_size, + num_oov_buckets=num_oov_buckets, + default_value=default_value, + dtype=dtype)) + + +def sequence_categorical_column_with_vocabulary_list( + key, vocabulary_list, dtype=None, default_value=-1, num_oov_buckets=0): + """A sequence of categorical terms where ids use an in-memory list. + + Pass this to `embedding_column` or `indicator_column` to convert sequence + categorical data into dense representation for input to sequence NN, such as + RNN. + + Example: + + ```python + colors = sequence_categorical_column_with_vocabulary_list( + key='colors', vocabulary_list=('R', 'G', 'B', 'Y'), + num_oov_buckets=2) + colors_embedding = embedding_column(colors, dimension=3) + columns = [colors_embedding] + + features = tf.parse_example(..., features=make_parse_example_spec(columns)) + sequence_feature_layer = SequenceFeatures(columns) + sequence_input, sequence_length = sequence_feature_layer(features) + sequence_length_mask = tf.sequence_mask(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: + key: A unique string identifying the input feature. + vocabulary_list: An ordered iterable defining the vocabulary. Each feature + is mapped to the index of its value (if present) in `vocabulary_list`. + Must be castable to `dtype`. + dtype: The type of features. Only string and integer types are supported. + If `None`, it will be inferred from `vocabulary_list`. + default_value: The integer ID value to return for out-of-vocabulary feature + values, defaults to `-1`. This can not be specified with a positive + `num_oov_buckets`. + num_oov_buckets: Non-negative integer, the number of out-of-vocabulary + buckets. All out-of-vocabulary inputs will be assigned IDs in the range + `[len(vocabulary_list), len(vocabulary_list)+num_oov_buckets)` based on a + hash of the input value. A positive `num_oov_buckets` can not be specified + with `default_value`. + + Returns: + A `SequenceCategoricalColumn`. + + Raises: + ValueError: if `vocabulary_list` is empty, or contains duplicate keys. + ValueError: `num_oov_buckets` is a negative integer. + ValueError: `num_oov_buckets` and `default_value` are both specified. + ValueError: if `dtype` is not integer or string. + """ + return fc.SequenceCategoricalColumn( + fc.categorical_column_with_vocabulary_list( + key=key, + vocabulary_list=vocabulary_list, + dtype=dtype, + default_value=default_value, + num_oov_buckets=num_oov_buckets)) + + +def sequence_numeric_column( + key, + shape=(1,), + default_value=0., + dtype=dtypes.float32, + normalizer_fn=None): + """Returns a feature column that represents sequences of numeric data. + + Example: + + ```python + temperature = sequence_numeric_column('temperature') + columns = [temperature] + + features = tf.parse_example(..., features=make_parse_example_spec(columns)) + sequence_feature_layer = SequenceFeatures(columns) + sequence_input, sequence_length = sequence_feature_layer(features) + sequence_length_mask = tf.sequence_mask(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: + key: A unique string identifying the input features. + shape: The shape of the input data per sequence id. E.g. if `shape=(2,)`, + each example must contain `2 * sequence_length` values. + default_value: A single value compatible with `dtype` that is used for + padding the sparse data into a dense `Tensor`. + dtype: The type of values. + normalizer_fn: If not `None`, a function that can be used to normalize the + value of the tensor after `default_value` is applied for parsing. + Normalizer function takes the input `Tensor` as its argument, and returns + the output `Tensor`. (e.g. lambda x: (x - 3.0) / 4.2). Please note that + even though the most common use case of this function is normalization, it + can be used for any kind of Tensorflow transformations. + + Returns: + A `SequenceNumericColumn`. + + Raises: + TypeError: if any dimension in shape is not an int. + ValueError: if any dimension in shape is not a positive integer. + ValueError: if `dtype` is not convertible to `tf.float32`. + """ + 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)) + if normalizer_fn is not None and not callable(normalizer_fn): + raise TypeError( + 'normalizer_fn must be a callable. Given: {}'.format(normalizer_fn)) + + return SequenceNumericColumn( + key, + shape=shape, + default_value=default_value, + dtype=dtype, + normalizer_fn=normalizer_fn) + + +def _assert_all_equal_and_return(tensors, name=None): + """Asserts that all tensors are equal and returns the first one.""" + with ops.name_scope(name, 'assert_all_equal', values=tensors): + if len(tensors) == 1: + return tensors[0] + assert_equal_ops = [] + for t in tensors[1:]: + assert_equal_ops.append(check_ops.assert_equal(tensors[0], t)) + with ops.control_dependencies(assert_equal_ops): + return array_ops.identity(tensors[0]) + + +class SequenceNumericColumn( + fc.SequenceDenseColumn, + collections.namedtuple( + 'SequenceNumericColumn', + ('key', 'shape', 'default_value', 'dtype', 'normalizer_fn'))): + """Represents sequences of numeric data.""" + + @property + def _is_v2_column(self): + return True + + @property + def name(self): + """See `FeatureColumn` base class.""" + return self.key + + @property + def parse_example_spec(self): + """See `FeatureColumn` base class.""" + return {self.key: parsing_ops.VarLenFeature(self.dtype)} + + def transform_feature(self, transformation_cache, state_manager): + """See `FeatureColumn` base class. + + In this case, we apply the `normalizer_fn` to the input tensor. + + Args: + transformation_cache: A `FeatureTransformationCache` object to access + features. + state_manager: A `StateManager` to create / access resources such as + lookup tables. + + Returns: + Normalized input tensor. + """ + input_tensor = transformation_cache.get(self.key, state_manager) + if self.normalizer_fn is not None: + input_tensor = self.normalizer_fn(input_tensor) + return input_tensor + + @property + def variable_shape(self): + """Returns a `TensorShape` representing the shape of sequence input.""" + return tensor_shape.TensorShape(self.shape) + + def get_sequence_dense_tensor(self, transformation_cache, state_manager): + """Returns a `TensorSequenceLengthPair`. + + Args: + transformation_cache: A `FeatureTransformationCache` object to access + features. + state_manager: A `StateManager` to create / access resources such as + lookup tables. + """ + sp_tensor = transformation_cache.get(self, state_manager) + dense_tensor = sparse_ops.sparse_tensor_to_dense( + sp_tensor, default_value=self.default_value) + # Reshape into [batch_size, T, variable_shape]. + dense_shape = array_ops.concat( + [array_ops.shape(dense_tensor)[:1], [-1], self.variable_shape], + axis=0) + dense_tensor = array_ops.reshape(dense_tensor, shape=dense_shape) + + # Get the number of timesteps per example + # 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. + 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) + + return fc.SequenceDenseColumn.TensorSequenceLengthPair( + dense_tensor=dense_tensor, sequence_length=seq_length) + + # TODO(b/119409767): Implement parents, _{get,from}_config. + @property + def parents(self): + """See 'FeatureColumn` base class.""" + raise NotImplementedError() + + def _get_config(self): + """See 'FeatureColumn` base class.""" + raise NotImplementedError() + + @classmethod + def _from_config(cls, config, custom_objects=None, columns_by_name=None): + """See 'FeatureColumn` base class.""" + raise NotImplementedError() + +# pylint: enable=protected-access 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 new file mode 100644 index 0000000000000000000000000000000000000000..a1feaddcc00d5fac86dca3138dfa1c6314bb6a8b --- /dev/null +++ b/tensorflow/contrib/feature_column/python/feature_column/sequence_feature_column_v2_test.py @@ -0,0 +1,1525 @@ +# 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 sequential_feature_column.""" + +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.feature_column.python.feature_column import sequence_feature_column_v2 as sfc +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 +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 SequenceFeaturesTest(test.TestCase, parameterized.TestCase): + + @parameterized.named_parameters( + {'testcase_name': '2D', + 'sparse_input_args_a': { + # example 0, ids [2] + # example 1, ids [0, 1] + 'indices': ((0, 0), (1, 0), (1, 1)), + 'values': (2, 0, 1), + 'dense_shape': (2, 2)}, + 'sparse_input_args_b': { + # example 0, ids [1] + # example 1, ids [2, 0] + 'indices': ((0, 0), (1, 0), (1, 1)), + 'values': (1, 2, 0), + 'dense_shape': (2, 2)}, + 'expected_input_layer': [ + # example 0, ids_a [2], ids_b [1] + [[5., 6., 14., 15., 16.], [0., 0., 0., 0., 0.]], + # example 1, ids_a [0, 1], ids_b [2, 0] + [[1., 2., 17., 18., 19.], [3., 4., 11., 12., 13.]],], + 'expected_sequence_length': [1, 2]}, + {'testcase_name': '3D', + 'sparse_input_args_a': { + # feature 0, ids [[2], [0, 1]] + # feature 1, ids [[0, 0], [1]] + 'indices': ( + (0, 0, 0), (0, 1, 0), (0, 1, 1), + (1, 0, 0), (1, 0, 1), (1, 1, 0)), + 'values': (2, 0, 1, 0, 0, 1), + 'dense_shape': (2, 2, 2)}, + 'sparse_input_args_b': { + # feature 0, ids [[1, 1], [1]] + # feature 1, ids [[2], [0]] + 'indices': ((0, 0, 0), (0, 0, 1), (0, 1, 0), (1, 0, 0), (1, 1, 0)), + 'values': (1, 1, 1, 2, 0), + 'dense_shape': (2, 2, 2)}, + 'expected_input_layer': [ + # feature 0, [a: 2, -, b: 1, 1], [a: 0, 1, b: 1, -] + [[5., 6., 14., 15., 16.], [2., 3., 14., 15., 16.]], + # feature 1, [a: 0, 0, b: 2, -], [a: 1, -, b: 0, -] + [[1., 2., 17., 18., 19.], [3., 4., 11., 12., 13.]]], + 'expected_sequence_length': [2, 2]}, + ) + def test_embedding_column( + self, sparse_input_args_a, sparse_input_args_b, expected_input_layer, + expected_sequence_length): + + sparse_input_a = sparse_tensor.SparseTensorValue(**sparse_input_args_a) + sparse_input_b = sparse_tensor.SparseTensorValue(**sparse_input_args_b) + vocabulary_size = 3 + embedding_dimension_a = 2 + embedding_values_a = ( + (1., 2.), # id 0 + (3., 4.), # id 1 + (5., 6.) # id 2 + ) + embedding_dimension_b = 3 + embedding_values_b = ( + (11., 12., 13.), # id 0 + (14., 15., 16.), # id 1 + (17., 18., 19.) # id 2 + ) + def _get_initializer(embedding_dimension, embedding_values): + def _initializer(shape, dtype, partition_info): + self.assertAllEqual((vocabulary_size, embedding_dimension), shape) + self.assertEqual(dtypes.float32, dtype) + self.assertIsNone(partition_info) + return embedding_values + return _initializer + + categorical_column_a = sfc.sequence_categorical_column_with_identity( + key='aaa', num_buckets=vocabulary_size) + 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.embedding_column( + categorical_column_b, + dimension=embedding_dimension_b, + initializer=_get_initializer(embedding_dimension_b, embedding_values_b)) + + # 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_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)) + self.assertAllEqual(embedding_values_b, global_vars[1].eval(session=sess)) + self.assertAllEqual(expected_input_layer, input_layer.eval(session=sess)) + self.assertAllEqual( + expected_sequence_length, sequence_length.eval(session=sess)) + + def test_embedding_column_with_non_sequence_categorical(self): + """Tests that error is raised for non-sequence embedding column.""" + vocabulary_size = 3 + sparse_input = sparse_tensor.SparseTensorValue( + # example 0, ids [2] + # example 1, ids [0, 1] + indices=((0, 0), (1, 0), (1, 1)), + values=(2, 0, 1), + dense_shape=(2, 2)) + + categorical_column_a = fc.categorical_column_with_identity( + key='aaa', num_buckets=vocabulary_size) + 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 SequenceFeatures\.'): + sequence_input_layer = sfc.SequenceFeatures([embedding_column_a]) + _, _ = sequence_input_layer({'aaa': sparse_input}) + + def test_shared_embedding_column(self): + vocabulary_size = 3 + sparse_input_a = sparse_tensor.SparseTensorValue( + # example 0, ids [2] + # example 1, ids [0, 1] + indices=((0, 0), (1, 0), (1, 1)), + values=(2, 0, 1), + dense_shape=(2, 2)) + sparse_input_b = sparse_tensor.SparseTensorValue( + # example 0, ids [1] + # example 1, ids [2, 0] + indices=((0, 0), (1, 0), (1, 1)), + values=(1, 2, 0), + dense_shape=(2, 2)) + + embedding_dimension = 2 + embedding_values = ( + (1., 2.), # id 0 + (3., 4.), # id 1 + (5., 6.) # id 2 + ) + + def _get_initializer(embedding_dimension, embedding_values): + + def _initializer(shape, dtype, partition_info): + self.assertAllEqual((vocabulary_size, embedding_dimension), shape) + self.assertEqual(dtypes.float32, dtype) + self.assertIsNone(partition_info) + return embedding_values + + return _initializer + + expected_input_layer = [ + # example 0, ids_a [2], ids_b [1] + [[5., 6., 3., 4.], [0., 0., 0., 0.]], + # example 1, ids_a [0, 1], ids_b [2, 0] + [[1., 2., 5., 6.], [3., 4., 1., 2.]], + ] + expected_sequence_length = [1, 2] + + categorical_column_a = sfc.sequence_categorical_column_with_identity( + key='aaa', num_buckets=vocabulary_size) + 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_v2( + [categorical_column_b, categorical_column_a], + dimension=embedding_dimension, + initializer=_get_initializer(embedding_dimension, embedding_values)) + + 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( + ('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)) + self.assertAllEqual(expected_input_layer, input_layer.eval(session=sess)) + self.assertAllEqual( + expected_sequence_length, sequence_length.eval(session=sess)) + + def test_shared_embedding_column_with_non_sequence_categorical(self): + """Tests that error is raised for non-sequence shared embedding column.""" + vocabulary_size = 3 + sparse_input_a = sparse_tensor.SparseTensorValue( + # example 0, ids [2] + # example 1, ids [0, 1] + indices=((0, 0), (1, 0), (1, 1)), + values=(2, 0, 1), + dense_shape=(2, 2)) + sparse_input_b = sparse_tensor.SparseTensorValue( + # example 0, ids [2] + # example 1, ids [0, 1] + indices=((0, 0), (1, 0), (1, 1)), + values=(2, 0, 1), + dense_shape=(2, 2)) + + categorical_column_a = fc.categorical_column_with_identity( + key='aaa', num_buckets=vocabulary_size) + categorical_column_b = fc.categorical_column_with_identity( + key='bbb', num_buckets=vocabulary_size) + 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 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', + 'sparse_input_args_a': { + # example 0, ids [2] + # example 1, ids [0, 1] + 'indices': ((0, 0), (1, 0), (1, 1)), + 'values': (2, 0, 1), + 'dense_shape': (2, 2)}, + 'sparse_input_args_b': { + # example 0, ids [1] + # example 1, ids [1, 0] + 'indices': ((0, 0), (1, 0), (1, 1)), + 'values': (1, 1, 0), + 'dense_shape': (2, 2)}, + 'expected_input_layer': [ + # example 0, ids_a [2], ids_b [1] + [[0., 0., 1., 0., 1.], [0., 0., 0., 0., 0.]], + # example 1, ids_a [0, 1], ids_b [1, 0] + [[1., 0., 0., 0., 1.], [0., 1., 0., 1., 0.]]], + 'expected_sequence_length': [1, 2]}, + {'testcase_name': '3D', + 'sparse_input_args_a': { + # feature 0, ids [[2], [0, 1]] + # feature 1, ids [[0, 0], [1]] + 'indices': ( + (0, 0, 0), (0, 1, 0), (0, 1, 1), + (1, 0, 0), (1, 0, 1), (1, 1, 0)), + 'values': (2, 0, 1, 0, 0, 1), + 'dense_shape': (2, 2, 2)}, + 'sparse_input_args_b': { + # feature 0, ids [[1, 1], [1]] + # feature 1, ids [[1], [0]] + 'indices': ((0, 0, 0), (0, 0, 1), (0, 1, 0), (1, 0, 0), (1, 1, 0)), + 'values': (1, 1, 1, 1, 0), + 'dense_shape': (2, 2, 2)}, + 'expected_input_layer': [ + # feature 0, [a: 2, -, b: 1, 1], [a: 0, 1, b: 1, -] + [[0., 0., 1., 0., 2.], [1., 1., 0., 0., 1.]], + # feature 1, [a: 0, 0, b: 1, -], [a: 1, -, b: 0, -] + [[2., 0., 0., 0., 1.], [0., 1., 0., 1., 0.]]], + 'expected_sequence_length': [2, 2]}, + ) + def test_indicator_column( + self, sparse_input_args_a, sparse_input_args_b, expected_input_layer, + expected_sequence_length): + sparse_input_a = sparse_tensor.SparseTensorValue(**sparse_input_args_a) + sparse_input_b = sparse_tensor.SparseTensorValue(**sparse_input_args_b) + + vocabulary_size_a = 3 + vocabulary_size_b = 2 + + categorical_column_a = sfc.sequence_categorical_column_with_identity( + key='aaa', num_buckets=vocabulary_size_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.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)) + self.assertAllEqual( + expected_sequence_length, sequence_length.eval(session=sess)) + + def test_indicator_column_with_non_sequence_categorical(self): + """Tests that error is raised for non-sequence categorical column.""" + vocabulary_size = 3 + sparse_input = sparse_tensor.SparseTensorValue( + # example 0, ids [2] + # example 1, ids [0, 1] + indices=((0, 0), (1, 0), (1, 1)), + values=(2, 0, 1), + dense_shape=(2, 2)) + + categorical_column_a = fc.categorical_column_with_identity( + key='aaa', num_buckets=vocabulary_size) + 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 SequenceFeatures\.'): + sequence_input_layer = sfc.SequenceFeatures([indicator_column_a]) + _, _ = sequence_input_layer({'aaa': sparse_input}) + + @parameterized.named_parameters( + {'testcase_name': '2D', + 'sparse_input_args': { + # example 0, values [0., 1] + # example 1, [10.] + 'indices': ((0, 0), (0, 1), (1, 0)), + 'values': (0., 1., 10.), + 'dense_shape': (2, 2)}, + 'expected_input_layer': [ + [[0.], [1.]], + [[10.], [0.]]], + 'expected_sequence_length': [2, 1]}, + {'testcase_name': '3D', + 'sparse_input_args': { + # 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.), + 'dense_shape': (2, 2, 2)}, + 'expected_input_layer': [ + [[20.], [3.], [5.], [0.]], + [[3.], [0.], [8.], [0.]]], + 'expected_sequence_length': [2, 2]}, + ) + def test_numeric_column( + self, sparse_input_args, expected_input_layer, expected_sequence_length): + sparse_input = sparse_tensor.SparseTensorValue(**sparse_input_args) + + numeric_column = sfc.sequence_numeric_column('aaa') + + 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)) + self.assertAllEqual( + expected_sequence_length, sequence_length.eval(session=sess)) + + @parameterized.named_parameters( + {'testcase_name': '2D', + 'sparse_input_args': { + # example 0, values [0., 1., 2., 3., 4., 5., 6., 7.] + # example 1, [10., 11., 12., 13.] + 'indices': ((0, 0), (0, 1), (0, 2), (0, 3), (0, 4), (0, 5), (0, 6), + (0, 7), (1, 0), (1, 1), (1, 2), (1, 3)), + 'values': (0., 1., 2., 3., 4., 5., 6., 7., 10., 11., 12., 13.), + 'dense_shape': (2, 8)}, + 'expected_input_layer': [ + # The output of numeric_column._get_dense_tensor should be flattened. + [[0., 1., 2., 3.], [4., 5., 6., 7.]], + [[10., 11., 12., 13.], [0., 0., 0., 0.]]], + 'expected_sequence_length': [2, 1]}, + {'testcase_name': '3D', + 'sparse_input_args': { + # example 0, values [[0., 1., 2., 3.]], [[4., 5., 6., 7.]] + # example 1, [[10., 11., 12., 13.], []] + 'indices': ((0, 0, 0), (0, 0, 1), (0, 0, 2), (0, 0, 3), + (0, 1, 0), (0, 1, 1), (0, 1, 2), (0, 1, 3), + (1, 0, 0), (1, 0, 1), (1, 0, 2), (1, 0, 3)), + 'values': (0., 1., 2., 3., 4., 5., 6., 7., 10., 11., 12., 13.), + 'dense_shape': (2, 2, 4)}, + 'expected_input_layer': [ + # The output of numeric_column._get_dense_tensor should be flattened. + [[0., 1., 2., 3.], [4., 5., 6., 7.]], + [[10., 11., 12., 13.], [0., 0., 0., 0.]]], + 'expected_sequence_length': [2, 1]}, + ) + def test_numeric_column_multi_dim( + self, sparse_input_args, expected_input_layer, expected_sequence_length): + """Tests SequenceFeatures for multi-dimensional numeric_column.""" + sparse_input = sparse_tensor.SparseTensorValue(**sparse_input_args) + + numeric_column = sfc.sequence_numeric_column('aaa', shape=(2, 2)) + + 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)) + self.assertAllEqual( + expected_sequence_length, sequence_length.eval(session=sess)) + + def test_sequence_length_not_equal(self): + """Tests that an error is raised when sequence lengths are not equal.""" + # Input a with sequence_length = [2, 1] + sparse_input_a = sparse_tensor.SparseTensorValue( + indices=((0, 0), (0, 1), (1, 0)), + values=(0., 1., 10.), + dense_shape=(2, 2)) + # Input b with sequence_length = [1, 1] + sparse_input_b = sparse_tensor.SparseTensorValue( + indices=((0, 0), (1, 0)), + values=(1., 10.), + dense_shape=(2, 2)) + numeric_column_a = sfc.sequence_numeric_column('aaa') + numeric_column_b = sfc.sequence_numeric_column('bbb') + + 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_features/aaa/sequence_length:0\) = \] \[2 1\] ' + r'\[y \(sequence_features/bbb/sequence_length:0\) = \] \[1 1\]'): + sess.run(sequence_length) + + @parameterized.named_parameters( + {'testcase_name': '2D', + 'sparse_input_args': { + # example 0, values [[[0., 1.], [2., 3.]], [[4., 5.], [6., 7.]]] + # example 1, [[[10., 11.], [12., 13.]]] + 'indices': ((0, 0), (0, 1), (0, 2), (0, 3), (0, 4), (0, 5), (0, 6), + (0, 7), (1, 0), (1, 1), (1, 2), (1, 3)), + 'values': (0., 1., 2., 3., 4., 5., 6., 7., 10., 11., 12., 13.), + 'dense_shape': (2, 8)}, + 'expected_shape': [2, 2, 4]}, + {'testcase_name': '3D', + 'sparse_input_args': { + # example 0, values [[0., 1., 2., 3.]], [[4., 5., 6., 7.]] + # example 1, [[10., 11., 12., 13.], []] + 'indices': ((0, 0, 0), (0, 0, 1), (0, 0, 2), (0, 0, 3), + (0, 1, 0), (0, 1, 1), (0, 1, 2), (0, 1, 2), + (1, 0, 0), (1, 0, 1), (1, 0, 2), (1, 0, 3)), + 'values': (0., 1., 2., 3., 4., 5., 6., 7., 10., 11., 12., 13.), + 'dense_shape': (2, 2, 4)}, + 'expected_shape': [2, 2, 4]}, + ) + def test_static_shape_from_tensors_numeric( + 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.sequence_numeric_column('aaa', shape=(2, 2)) + + 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) + + @parameterized.named_parameters( + {'testcase_name': '2D', + 'sparse_input_args': { + # example 0, ids [2] + # example 1, ids [0, 1] + # example 2, ids [] + # example 3, ids [1] + 'indices': ((0, 0), (1, 0), (1, 1), (3, 0)), + 'values': (2, 0, 1, 1), + 'dense_shape': (4, 2)}, + 'expected_shape': [4, 2, 3]}, + {'testcase_name': '3D', + 'sparse_input_args': { + # example 0, ids [[2]] + # example 1, ids [[0, 1], [2]] + # example 2, ids [] + # example 3, ids [[1], [0, 2]] + 'indices': ((0, 0, 0), (1, 0, 0), (1, 0, 1), (1, 1, 0), + (3, 0, 0), (3, 1, 0), (3, 1, 1)), + 'values': (2, 0, 1, 2, 1, 0, 2), + 'dense_shape': (4, 2, 2)}, + 'expected_shape': [4, 2, 3]} + ) + def test_static_shape_from_tensors_indicator( + 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) + categorical_column = sfc.sequence_categorical_column_with_identity( + key='aaa', num_buckets=3) + indicator_column = fc.indicator_column(categorical_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.""" + + def test_concatenate_context_input(self): + seq_input = ops.convert_to_tensor(np.arange(12).reshape(2, 3, 2)) + context_input = ops.convert_to_tensor(np.arange(10).reshape(2, 5)) + seq_input = math_ops.cast(seq_input, dtype=dtypes.float32) + context_input = math_ops.cast(context_input, dtype=dtypes.float32) + input_layer = sfc.concatenate_context_input(context_input, seq_input) + + expected = np.array([ + [[0, 1, 0, 1, 2, 3, 4], [2, 3, 0, 1, 2, 3, 4], [4, 5, 0, 1, 2, 3, 4]], + [[6, 7, 5, 6, 7, 8, 9], [8, 9, 5, 6, 7, 8, 9], [10, 11, 5, 6, 7, 8, 9]] + ], dtype=np.float32) + with monitored_session.MonitoredSession() as sess: + output = sess.run(input_layer) + self.assertAllEqual(expected, output) + + @parameterized.named_parameters( + {'testcase_name': 'rank_lt_3', + 'seq_input_arg': np.arange(100).reshape(10, 10)}, + {'testcase_name': 'rank_gt_3', + 'seq_input_arg': np.arange(100).reshape(5, 5, 2, 2)} + ) + def test_sequence_input_throws_error(self, seq_input_arg): + seq_input = ops.convert_to_tensor(seq_input_arg) + context_input = ops.convert_to_tensor(np.arange(100).reshape(10, 10)) + seq_input = math_ops.cast(seq_input, dtype=dtypes.float32) + context_input = math_ops.cast(context_input, dtype=dtypes.float32) + with self.assertRaisesRegexp(ValueError, 'sequence_input must have rank 3'): + sfc.concatenate_context_input(context_input, seq_input) + + @parameterized.named_parameters( + {'testcase_name': 'rank_lt_2', + 'context_input_arg': np.arange(100)}, + {'testcase_name': 'rank_gt_2', + 'context_input_arg': np.arange(100).reshape(5, 5, 4)} + ) + def test_context_input_throws_error(self, context_input_arg): + context_input = ops.convert_to_tensor(context_input_arg) + seq_input = ops.convert_to_tensor(np.arange(100).reshape(5, 5, 4)) + seq_input = math_ops.cast(seq_input, dtype=dtypes.float32) + context_input = math_ops.cast(context_input, dtype=dtypes.float32) + with self.assertRaisesRegexp(ValueError, 'context_input must have rank 2'): + sfc.concatenate_context_input(context_input, seq_input) + + def test_integer_seq_input_throws_error(self): + seq_input = ops.convert_to_tensor(np.arange(100).reshape(5, 5, 4)) + context_input = ops.convert_to_tensor(np.arange(100).reshape(10, 10)) + context_input = math_ops.cast(context_input, dtype=dtypes.float32) + with self.assertRaisesRegexp( + TypeError, 'sequence_input must have dtype float32'): + sfc.concatenate_context_input(context_input, seq_input) + + def test_integer_context_input_throws_error(self): + seq_input = ops.convert_to_tensor(np.arange(100).reshape(5, 5, 4)) + context_input = ops.convert_to_tensor(np.arange(100).reshape(10, 10)) + seq_input = math_ops.cast(seq_input, dtype=dtypes.float32) + with self.assertRaisesRegexp( + TypeError, 'context_input must have dtype float32'): + sfc.concatenate_context_input(context_input, seq_input) + + +class DenseFeaturesTest(test.TestCase): + """Tests DenseFeatures with sequence feature columns.""" + + def test_embedding_column(self): + """Tests that error is raised for sequence embedding column.""" + vocabulary_size = 3 + sparse_input = sparse_tensor.SparseTensorValue( + # example 0, ids [2] + # example 1, ids [0, 1] + indices=((0, 0), (1, 0), (1, 1)), + values=(2, 0, 1), + dense_shape=(2, 2)) + + categorical_column_a = sfc.sequence_categorical_column_with_identity( + key='aaa', num_buckets=vocabulary_size) + 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\.'): + 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.""" + vocabulary_size = 3 + sparse_input = sparse_tensor.SparseTensorValue( + # example 0, ids [2] + # example 1, ids [0, 1] + indices=((0, 0), (1, 0), (1, 1)), + values=(2, 0, 1), + dense_shape=(2, 2)) + + categorical_column_a = sfc.sequence_categorical_column_with_identity( + key='aaa', num_buckets=vocabulary_size) + 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\.'): + input_layer = fc.DenseFeatures([indicator_column_a]) + _ = input_layer({'aaa': sparse_input}) + + +def _assert_sparse_tensor_value(test_case, expected, actual): + _assert_sparse_tensor_indices_shape(test_case, expected, actual) + + test_case.assertEqual( + np.array(expected.values).dtype, np.array(actual.values).dtype) + test_case.assertAllEqual(expected.values, actual.values) + + +def _assert_sparse_tensor_indices_shape(test_case, expected, actual): + test_case.assertEqual(np.int64, np.array(actual.indices).dtype) + test_case.assertAllEqual(expected.indices, actual.indices) + + test_case.assertEqual(np.int64, np.array(actual.dense_shape).dtype) + test_case.assertAllEqual(expected.dense_shape, actual.dense_shape) + + +def _get_sequence_dense_tensor(column, features): + return column.get_sequence_dense_tensor( + fc.FeatureTransformationCache(features), None) + + +def _get_sequence_dense_tensor_state(column, features): + state_manager = _TestStateManager() + column.create_state(state_manager) + return column.get_sequence_dense_tensor( + fc.FeatureTransformationCache(features), state_manager) + + +def _get_sparse_tensors(column, features): + return column.get_sparse_tensors( + fc.FeatureTransformationCache(features), None) + + +class SequenceCategoricalColumnWithIdentityTest( + test.TestCase, parameterized.TestCase): + + @parameterized.named_parameters( + {'testcase_name': '2D', + 'inputs_args': { + 'indices': ((0, 0), (1, 0), (1, 1)), + 'values': (1, 2, 0), + 'dense_shape': (2, 2)}, + 'expected_args': { + 'indices': ((0, 0, 0), (1, 0, 0), (1, 1, 0)), + 'values': np.array((1, 2, 0), dtype=np.int64), + 'dense_shape': (2, 2, 1)}}, + {'testcase_name': '3D', + 'inputs_args': { + 'indices': ((0, 0, 2), (1, 0, 0), (1, 2, 0)), + 'values': (6, 7, 8), + 'dense_shape': (2, 2, 2)}, + 'expected_args': { + 'indices': ((0, 0, 2), (1, 0, 0), (1, 2, 0)), + 'values': (6, 7, 8), + 'dense_shape': (2, 2, 2)}} + ) + def test_get_sparse_tensors(self, inputs_args, expected_args): + inputs = sparse_tensor.SparseTensorValue(**inputs_args) + expected = sparse_tensor.SparseTensorValue(**expected_args) + column = sfc.sequence_categorical_column_with_identity('aaa', num_buckets=9) + + id_weight_pair = _get_sparse_tensors(column, {'aaa': inputs}) + + self.assertIsNone(id_weight_pair.weight_tensor) + with monitored_session.MonitoredSession() as sess: + _assert_sparse_tensor_value( + self, expected, id_weight_pair.id_tensor.eval(session=sess)) + + +class SequenceCategoricalColumnWithHashBucketTest( + test.TestCase, parameterized.TestCase): + + @parameterized.named_parameters( + {'testcase_name': '2D', + 'inputs_args': { + 'indices': ((0, 0), (1, 0), (1, 1)), + 'values': ('omar', 'stringer', 'marlo'), + 'dense_shape': (2, 2)}, + 'expected_args': { + 'indices': ((0, 0, 0), (1, 0, 0), (1, 1, 0)), + # Ignored to avoid hash dependence in test. + 'values': np.array((0, 0, 0), dtype=np.int64), + 'dense_shape': (2, 2, 1)}}, + {'testcase_name': '3D', + 'inputs_args': { + 'indices': ((0, 0, 2), (1, 0, 0), (1, 2, 0)), + 'values': ('omar', 'stringer', 'marlo'), + 'dense_shape': (2, 2, 2)}, + 'expected_args': { + 'indices': ((0, 0, 2), (1, 0, 0), (1, 2, 0)), + # Ignored to avoid hash dependence in test. + 'values': np.array((0, 0, 0), dtype=np.int64), + 'dense_shape': (2, 2, 2)}} + ) + def test_get_sparse_tensors(self, inputs_args, expected_args): + inputs = sparse_tensor.SparseTensorValue(**inputs_args) + expected = sparse_tensor.SparseTensorValue(**expected_args) + column = sfc.sequence_categorical_column_with_hash_bucket( + 'aaa', hash_bucket_size=10) + + id_weight_pair = _get_sparse_tensors(column, {'aaa': inputs}) + + self.assertIsNone(id_weight_pair.weight_tensor) + with monitored_session.MonitoredSession() as sess: + _assert_sparse_tensor_indices_shape( + self, expected, id_weight_pair.id_tensor.eval(session=sess)) + + +class SequenceCategoricalColumnWithVocabularyFileTest( + test.TestCase, parameterized.TestCase): + + def _write_vocab(self, vocab_strings, file_name): + vocab_file = os.path.join(self.get_temp_dir(), file_name) + with open(vocab_file, 'w') as f: + f.write('\n'.join(vocab_strings)) + return vocab_file + + def setUp(self): + super(SequenceCategoricalColumnWithVocabularyFileTest, self).setUp() + + vocab_strings = ['omar', 'stringer', 'marlo'] + self._wire_vocabulary_file_name = self._write_vocab(vocab_strings, + 'wire_vocabulary.txt') + self._wire_vocabulary_size = 3 + + @parameterized.named_parameters( + {'testcase_name': '2D', + 'inputs_args': { + 'indices': ((0, 0), (1, 0), (1, 1)), + 'values': ('marlo', 'skywalker', 'omar'), + 'dense_shape': (2, 2)}, + 'expected_args': { + 'indices': ((0, 0, 0), (1, 0, 0), (1, 1, 0)), + 'values': np.array((2, -1, 0), dtype=np.int64), + 'dense_shape': (2, 2, 1)}}, + {'testcase_name': '3D', + 'inputs_args': { + 'indices': ((0, 0, 2), (1, 0, 0), (1, 2, 0)), + 'values': ('omar', 'skywalker', 'marlo'), + 'dense_shape': (2, 2, 2)}, + 'expected_args': { + 'indices': ((0, 0, 2), (1, 0, 0), (1, 2, 0)), + 'values': np.array((0, -1, 2), dtype=np.int64), + 'dense_shape': (2, 2, 2)}} + ) + def test_get_sparse_tensors(self, inputs_args, expected_args): + inputs = sparse_tensor.SparseTensorValue(**inputs_args) + expected = sparse_tensor.SparseTensorValue(**expected_args) + column = sfc.sequence_categorical_column_with_vocabulary_file( + key='aaa', + vocabulary_file=self._wire_vocabulary_file_name, + vocabulary_size=self._wire_vocabulary_size) + + id_weight_pair = _get_sparse_tensors(column, {'aaa': inputs}) + + self.assertIsNone(id_weight_pair.weight_tensor) + with monitored_session.MonitoredSession() as sess: + _assert_sparse_tensor_value( + self, expected, id_weight_pair.id_tensor.eval(session=sess)) + + def test_get_sparse_tensors_dynamic_zero_length(self): + """Tests _get_sparse_tensors with a dynamic sequence length.""" + inputs = sparse_tensor.SparseTensorValue( + indices=np.zeros((0, 2)), values=[], dense_shape=(2, 0)) + expected = sparse_tensor.SparseTensorValue( + indices=np.zeros((0, 3)), + values=np.array((), dtype=np.int64), + dense_shape=(2, 0, 1)) + column = sfc.sequence_categorical_column_with_vocabulary_file( + key='aaa', + vocabulary_file=self._wire_vocabulary_file_name, + vocabulary_size=self._wire_vocabulary_size) + input_placeholder_shape = list(inputs.dense_shape) + # Make second dimension (sequence length) dynamic. + input_placeholder_shape[1] = None + input_placeholder = array_ops.sparse_placeholder( + dtypes.string, shape=input_placeholder_shape) + id_weight_pair = _get_sparse_tensors(column, {'aaa': input_placeholder}) + + self.assertIsNone(id_weight_pair.weight_tensor) + with monitored_session.MonitoredSession() as sess: + result = id_weight_pair.id_tensor.eval( + session=sess, feed_dict={input_placeholder: inputs}) + _assert_sparse_tensor_value( + self, expected, result) + + +class SequenceCategoricalColumnWithVocabularyListTest( + test.TestCase, parameterized.TestCase): + + @parameterized.named_parameters( + {'testcase_name': '2D', + 'inputs_args': { + 'indices': ((0, 0), (1, 0), (1, 1)), + 'values': ('marlo', 'skywalker', 'omar'), + 'dense_shape': (2, 2)}, + 'expected_args': { + 'indices': ((0, 0, 0), (1, 0, 0), (1, 1, 0)), + 'values': np.array((2, -1, 0), dtype=np.int64), + 'dense_shape': (2, 2, 1)}}, + {'testcase_name': '3D', + 'inputs_args': { + 'indices': ((0, 0, 2), (1, 0, 0), (1, 2, 0)), + 'values': ('omar', 'skywalker', 'marlo'), + 'dense_shape': (2, 2, 2)}, + 'expected_args': { + 'indices': ((0, 0, 2), (1, 0, 0), (1, 2, 0)), + 'values': np.array((0, -1, 2), dtype=np.int64), + 'dense_shape': (2, 2, 2)}} + ) + def test_get_sparse_tensors(self, inputs_args, expected_args): + inputs = sparse_tensor.SparseTensorValue(**inputs_args) + expected = sparse_tensor.SparseTensorValue(**expected_args) + column = sfc.sequence_categorical_column_with_vocabulary_list( + key='aaa', + vocabulary_list=('omar', 'stringer', 'marlo')) + + id_weight_pair = _get_sparse_tensors(column, {'aaa': inputs}) + + self.assertIsNone(id_weight_pair.weight_tensor) + with monitored_session.MonitoredSession() as sess: + _assert_sparse_tensor_value( + self, expected, id_weight_pair.id_tensor.eval(session=sess)) + + +class SequenceEmbeddingColumnTest( + test.TestCase, parameterized.TestCase): + + @parameterized.named_parameters( + {'testcase_name': '2D', + 'inputs_args': { + # example 0, ids [2] + # example 1, ids [0, 1] + # example 2, ids [] + # example 3, ids [1] + 'indices': ((0, 0), (1, 0), (1, 1), (3, 0)), + 'values': (2, 0, 1, 1), + 'dense_shape': (4, 2)}, + 'expected': [ + # example 0, ids [2] + [[7., 11.], [0., 0.]], + # example 1, ids [0, 1] + [[1., 2.], [3., 5.]], + # example 2, ids [] + [[0., 0.], [0., 0.]], + # example 3, ids [1] + [[3., 5.], [0., 0.]]]}, + {'testcase_name': '3D', + 'inputs_args': { + # example 0, ids [[2]] + # example 1, ids [[0, 1], [2]] + # example 2, ids [] + # example 3, ids [[1], [0, 2]] + 'indices': ((0, 0, 0), (1, 0, 0), (1, 0, 1), (1, 1, 0), + (3, 0, 0), (3, 1, 0), (3, 1, 1)), + 'values': (2, 0, 1, 2, 1, 0, 2), + 'dense_shape': (4, 2, 2)}, + 'expected': [ + # example 0, ids [[2]] + [[7., 11.], [0., 0.]], + # example 1, ids [[0, 1], [2]] + [[2, 3.5], [7., 11.]], + # example 2, ids [] + [[0., 0.], [0., 0.]], + # example 3, ids [[1], [0, 2]] + [[3., 5.], [4., 6.5]]]} + ) + def test_get_sequence_dense_tensor(self, inputs_args, expected): + inputs = sparse_tensor.SparseTensorValue(**inputs_args) + vocabulary_size = 3 + embedding_dimension = 2 + embedding_values = ( + (1., 2.), # id 0 + (3., 5.), # id 1 + (7., 11.) # id 2 + ) + def _initializer(shape, dtype, partition_info): + self.assertAllEqual((vocabulary_size, embedding_dimension), shape) + self.assertEqual(dtypes.float32, dtype) + self.assertIsNone(partition_info) + return embedding_values + + categorical_column = sfc.sequence_categorical_column_with_identity( + key='aaa', num_buckets=vocabulary_size) + embedding_column = fc.embedding_column( + categorical_column, dimension=embedding_dimension, + initializer=_initializer) + + embedding_lookup, _ = _get_sequence_dense_tensor_state( + embedding_column, {'aaa': inputs}) + + global_vars = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES) + 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)) + self.assertAllEqual(expected, embedding_lookup.eval(session=sess)) + + @parameterized.named_parameters( + {'testcase_name': '2D', + 'inputs_args': { + # example 0, ids [2] + # example 1, ids [0, 1] + 'indices': ((0, 0), (1, 0), (1, 1)), + 'values': (2, 0, 1), + 'dense_shape': (2, 2)}, + 'expected_sequence_length': [1, 2]}, + {'testcase_name': '3D', + 'inputs_args': { + # example 0, ids [[2]] + # example 1, ids [[0, 1], [2]] + 'indices': ((0, 0, 0), (1, 0, 0), (1, 0, 1), (1, 1, 0)), + 'values': (2, 0, 1, 2), + 'dense_shape': (2, 2, 2)}, + 'expected_sequence_length': [1, 2]} + ) + def test_sequence_length(self, inputs_args, expected_sequence_length): + inputs = sparse_tensor.SparseTensorValue(**inputs_args) + vocabulary_size = 3 + + categorical_column = sfc.sequence_categorical_column_with_identity( + key='aaa', num_buckets=vocabulary_size) + embedding_column = fc.embedding_column( + categorical_column, dimension=2) + + _, sequence_length = _get_sequence_dense_tensor_state( + embedding_column, {'aaa': inputs}) + + with monitored_session.MonitoredSession() as sess: + sequence_length = sess.run(sequence_length) + self.assertAllEqual(expected_sequence_length, sequence_length) + self.assertEqual(np.int64, sequence_length.dtype) + + def test_sequence_length_with_empty_rows(self): + """Tests _sequence_length when some examples do not have ids.""" + vocabulary_size = 3 + sparse_input = sparse_tensor.SparseTensorValue( + # example 0, ids [] + # example 1, ids [2] + # example 2, ids [0, 1] + # example 3, ids [] + # example 4, ids [1] + # example 5, ids [] + indices=((1, 0), (2, 0), (2, 1), (4, 0)), + values=(2, 0, 1, 1), + dense_shape=(6, 2)) + expected_sequence_length = [0, 1, 2, 0, 1, 0] + + categorical_column = sfc.sequence_categorical_column_with_identity( + key='aaa', num_buckets=vocabulary_size) + embedding_column = fc.embedding_column( + categorical_column, dimension=2) + + _, sequence_length = _get_sequence_dense_tensor_state( + embedding_column, {'aaa': sparse_input}) + + with monitored_session.MonitoredSession() as sess: + self.assertAllEqual( + expected_sequence_length, sequence_length.eval(session=sess)) + + +class SequenceSharedEmbeddingColumnTest(test.TestCase): + + def test_get_sequence_dense_tensor(self): + vocabulary_size = 3 + embedding_dimension = 2 + embedding_values = ( + (1., 2.), # id 0 + (3., 5.), # id 1 + (7., 11.) # id 2 + ) + + def _initializer(shape, dtype, partition_info): + self.assertAllEqual((vocabulary_size, embedding_dimension), shape) + self.assertEqual(dtypes.float32, dtype) + self.assertIsNone(partition_info) + return embedding_values + + sparse_input_a = sparse_tensor.SparseTensorValue( + # example 0, ids [2] + # example 1, ids [0, 1] + # example 2, ids [] + # example 3, ids [1] + indices=((0, 0), (1, 0), (1, 1), (3, 0)), + values=(2, 0, 1, 1), + dense_shape=(4, 2)) + sparse_input_b = sparse_tensor.SparseTensorValue( + # example 0, ids [1] + # example 1, ids [0, 2] + # example 2, ids [0] + # example 3, ids [] + indices=((0, 0), (1, 0), (1, 1), (2, 0)), + values=(1, 0, 2, 0), + dense_shape=(4, 2)) + + expected_lookups_a = [ + # example 0, ids [2] + [[7., 11.], [0., 0.]], + # example 1, ids [0, 1] + [[1., 2.], [3., 5.]], + # example 2, ids [] + [[0., 0.], [0., 0.]], + # example 3, ids [1] + [[3., 5.], [0., 0.]], + ] + + expected_lookups_b = [ + # example 0, ids [1] + [[3., 5.], [0., 0.]], + # example 1, ids [0, 2] + [[1., 2.], [7., 11.]], + # example 2, ids [0] + [[1., 2.], [0., 0.]], + # example 3, ids [] + [[0., 0.], [0., 0.]], + ] + + categorical_column_a = sfc.sequence_categorical_column_with_identity( + key='aaa', num_buckets=vocabulary_size) + categorical_column_b = sfc.sequence_categorical_column_with_identity( + key='bbb', num_buckets=vocabulary_size) + shared_embedding_columns = fc.shared_embedding_columns_v2( + [categorical_column_a, categorical_column_b], + dimension=embedding_dimension, + initializer=_initializer) + + embedding_lookup_a = _get_sequence_dense_tensor( + shared_embedding_columns[0], {'aaa': sparse_input_a})[0] + embedding_lookup_b = _get_sequence_dense_tensor( + shared_embedding_columns[1], {'bbb': sparse_input_b})[0] + + global_vars = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES) + self.assertItemsEqual(('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)) + self.assertAllEqual( + expected_lookups_a, embedding_lookup_a.eval(session=sess)) + self.assertAllEqual( + expected_lookups_b, embedding_lookup_b.eval(session=sess)) + + def test_sequence_length(self): + vocabulary_size = 3 + + sparse_input_a = sparse_tensor.SparseTensorValue( + # example 0, ids [2] + # example 1, ids [0, 1] + indices=((0, 0), (1, 0), (1, 1)), + values=(2, 0, 1), + dense_shape=(2, 2)) + expected_sequence_length_a = [1, 2] + categorical_column_a = sfc.sequence_categorical_column_with_identity( + key='aaa', num_buckets=vocabulary_size) + + sparse_input_b = sparse_tensor.SparseTensorValue( + # example 0, ids [0, 2] + # example 1, ids [1] + indices=((0, 0), (0, 1), (1, 0)), + values=(0, 2, 1), + dense_shape=(2, 2)) + expected_sequence_length_b = [2, 1] + categorical_column_b = sfc.sequence_categorical_column_with_identity( + key='bbb', num_buckets=vocabulary_size) + shared_embedding_columns = fc.shared_embedding_columns_v2( + [categorical_column_a, categorical_column_b], dimension=2) + + sequence_length_a = _get_sequence_dense_tensor( + shared_embedding_columns[0], {'aaa': sparse_input_a})[1] + sequence_length_b = _get_sequence_dense_tensor( + shared_embedding_columns[1], {'bbb': sparse_input_b})[1] + + with monitored_session.MonitoredSession() as sess: + sequence_length_a = sess.run(sequence_length_a) + self.assertAllEqual(expected_sequence_length_a, sequence_length_a) + self.assertEqual(np.int64, sequence_length_a.dtype) + sequence_length_b = sess.run(sequence_length_b) + self.assertAllEqual(expected_sequence_length_b, sequence_length_b) + self.assertEqual(np.int64, sequence_length_b.dtype) + + def test_sequence_length_with_empty_rows(self): + """Tests _sequence_length when some examples do not have ids.""" + vocabulary_size = 3 + sparse_input_a = sparse_tensor.SparseTensorValue( + # example 0, ids [] + # example 1, ids [2] + # example 2, ids [0, 1] + # example 3, ids [] + # example 4, ids [1] + # example 5, ids [] + indices=((1, 0), (2, 0), (2, 1), (4, 0)), + values=(2, 0, 1, 1), + dense_shape=(6, 2)) + expected_sequence_length_a = [0, 1, 2, 0, 1, 0] + categorical_column_a = sfc.sequence_categorical_column_with_identity( + key='aaa', num_buckets=vocabulary_size) + + sparse_input_b = sparse_tensor.SparseTensorValue( + # example 0, ids [2] + # example 1, ids [] + # example 2, ids [] + # example 3, ids [] + # example 4, ids [1] + # example 5, ids [0, 1] + indices=((0, 0), (4, 0), (5, 0), (5, 1)), + values=(2, 1, 0, 1), + dense_shape=(6, 2)) + expected_sequence_length_b = [1, 0, 0, 0, 1, 2] + categorical_column_b = sfc.sequence_categorical_column_with_identity( + key='bbb', num_buckets=vocabulary_size) + + shared_embedding_columns = fc.shared_embedding_columns_v2( + [categorical_column_a, categorical_column_b], dimension=2) + + sequence_length_a = _get_sequence_dense_tensor( + shared_embedding_columns[0], {'aaa': sparse_input_a})[1] + sequence_length_b = _get_sequence_dense_tensor( + shared_embedding_columns[1], {'bbb': sparse_input_b})[1] + + with monitored_session.MonitoredSession() as sess: + self.assertAllEqual( + expected_sequence_length_a, sequence_length_a.eval(session=sess)) + self.assertAllEqual( + expected_sequence_length_b, sequence_length_b.eval(session=sess)) + + +class SequenceIndicatorColumnTest(test.TestCase, parameterized.TestCase): + + @parameterized.named_parameters( + {'testcase_name': '2D', + 'inputs_args': { + # example 0, ids [2] + # example 1, ids [0, 1] + # example 2, ids [] + # example 3, ids [1] + 'indices': ((0, 0), (1, 0), (1, 1), (3, 0)), + 'values': (2, 0, 1, 1), + 'dense_shape': (4, 2)}, + 'expected': [ + # example 0, ids [2] + [[0., 0., 1.], [0., 0., 0.]], + # example 1, ids [0, 1] + [[1., 0., 0.], [0., 1., 0.]], + # example 2, ids [] + [[0., 0., 0.], [0., 0., 0.]], + # example 3, ids [1] + [[0., 1., 0.], [0., 0., 0.]]]}, + {'testcase_name': '3D', + 'inputs_args': { + # example 0, ids [[2]] + # example 1, ids [[0, 1], [2]] + # example 2, ids [] + # example 3, ids [[1], [2, 2]] + 'indices': ((0, 0, 0), (1, 0, 0), (1, 0, 1), (1, 1, 0), + (3, 0, 0), (3, 1, 0), (3, 1, 1)), + 'values': (2, 0, 1, 2, 1, 2, 2), + 'dense_shape': (4, 2, 2)}, + 'expected': [ + # example 0, ids [[2]] + [[0., 0., 1.], [0., 0., 0.]], + # example 1, ids [[0, 1], [2]] + [[1., 1., 0.], [0., 0., 1.]], + # example 2, ids [] + [[0., 0., 0.], [0., 0., 0.]], + # example 3, ids [[1], [2, 2]] + [[0., 1., 0.], [0., 0., 2.]]]} + ) + def test_get_sequence_dense_tensor(self, inputs_args, expected): + inputs = sparse_tensor.SparseTensorValue(**inputs_args) + vocabulary_size = 3 + + categorical_column = sfc.sequence_categorical_column_with_identity( + key='aaa', num_buckets=vocabulary_size) + indicator_column = fc.indicator_column(categorical_column) + + indicator_tensor, _ = _get_sequence_dense_tensor( + indicator_column, {'aaa': inputs}) + + with monitored_session.MonitoredSession() as sess: + self.assertAllEqual(expected, indicator_tensor.eval(session=sess)) + + @parameterized.named_parameters( + {'testcase_name': '2D', + 'inputs_args': { + # example 0, ids [2] + # example 1, ids [0, 1] + 'indices': ((0, 0), (1, 0), (1, 1)), + 'values': (2, 0, 1), + 'dense_shape': (2, 2)}, + 'expected_sequence_length': [1, 2]}, + {'testcase_name': '3D', + 'inputs_args': { + # example 0, ids [[2]] + # example 1, ids [[0, 1], [2]] + 'indices': ((0, 0, 0), (1, 0, 0), (1, 0, 1), (1, 1, 0)), + 'values': (2, 0, 1, 2), + 'dense_shape': (2, 2, 2)}, + 'expected_sequence_length': [1, 2]} + ) + def test_sequence_length(self, inputs_args, expected_sequence_length): + inputs = sparse_tensor.SparseTensorValue(**inputs_args) + vocabulary_size = 3 + + categorical_column = sfc.sequence_categorical_column_with_identity( + key='aaa', num_buckets=vocabulary_size) + indicator_column = fc.indicator_column(categorical_column) + + _, sequence_length = _get_sequence_dense_tensor( + indicator_column, {'aaa': inputs}) + + with monitored_session.MonitoredSession() as sess: + sequence_length = sess.run(sequence_length) + self.assertAllEqual(expected_sequence_length, sequence_length) + self.assertEqual(np.int64, sequence_length.dtype) + + def test_sequence_length_with_empty_rows(self): + """Tests _sequence_length when some examples do not have ids.""" + vocabulary_size = 3 + sparse_input = sparse_tensor.SparseTensorValue( + # example 0, ids [] + # example 1, ids [2] + # example 2, ids [0, 1] + # example 3, ids [] + # example 4, ids [1] + # example 5, ids [] + indices=((1, 0), (2, 0), (2, 1), (4, 0)), + values=(2, 0, 1, 1), + dense_shape=(6, 2)) + expected_sequence_length = [0, 1, 2, 0, 1, 0] + + categorical_column = sfc.sequence_categorical_column_with_identity( + key='aaa', num_buckets=vocabulary_size) + indicator_column = fc.indicator_column(categorical_column) + + _, sequence_length = _get_sequence_dense_tensor( + indicator_column, {'aaa': sparse_input}) + + with monitored_session.MonitoredSession() as sess: + self.assertAllEqual( + expected_sequence_length, sequence_length.eval(session=sess)) + + +class SequenceNumericColumnTest(test.TestCase, parameterized.TestCase): + + def test_defaults(self): + a = sfc.sequence_numeric_column('aaa') + self.assertEqual('aaa', a.key) + self.assertEqual('aaa', a.name) + self.assertEqual((1,), a.shape) + self.assertEqual(0., a.default_value) + self.assertEqual(dtypes.float32, a.dtype) + self.assertIsNone(a.normalizer_fn) + + def test_shape_saved_as_tuple(self): + a = sfc.sequence_numeric_column('aaa', shape=[1, 2]) + self.assertEqual((1, 2), a.shape) + + def test_shape_must_be_positive_integer(self): + with self.assertRaisesRegexp(TypeError, 'shape dimensions must be integer'): + sfc.sequence_numeric_column('aaa', shape=[1.0]) + + with self.assertRaisesRegexp( + ValueError, 'shape dimensions must be greater than 0'): + sfc.sequence_numeric_column('aaa', shape=[0]) + + def test_dtype_is_convertible_to_float(self): + with self.assertRaisesRegexp( + ValueError, 'dtype must be convertible to float'): + sfc.sequence_numeric_column('aaa', dtype=dtypes.string) + + def test_normalizer_fn_must_be_callable(self): + with self.assertRaisesRegexp(TypeError, 'must be a callable'): + sfc.sequence_numeric_column('aaa', normalizer_fn='NotACallable') + + @parameterized.named_parameters( + {'testcase_name': '2D', + 'inputs_args': { + # example 0, values [0., 1] + # example 1, [10.] + 'indices': ((0, 0), (0, 1), (1, 0)), + 'values': (0., 1., 10.), + 'dense_shape': (2, 2)}, + 'expected': [ + [[0.], [1.]], + [[10.], [0.]]]}, + {'testcase_name': '3D', + 'inputs_args': { + # 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.), + 'dense_shape': (2, 2, 2)}, + 'expected': [ + [[20.], [3.], [5.], [0.]], + [[3.], [0.], [8.], [0.]]]}, + ) + def test_get_sequence_dense_tensor(self, inputs_args, expected): + inputs = sparse_tensor.SparseTensorValue(**inputs_args) + numeric_column = sfc.sequence_numeric_column('aaa') + + dense_tensor, _ = _get_sequence_dense_tensor( + numeric_column, {'aaa': inputs}) + with monitored_session.MonitoredSession() as sess: + self.assertAllEqual(expected, dense_tensor.eval(session=sess)) + + def test_get_sequence_dense_tensor_with_normalizer_fn(self): + + def _increment_two(input_sparse_tensor): + return sparse_ops.sparse_add( + input_sparse_tensor, + sparse_tensor.SparseTensor(((0, 0), (1, 1)), (2.0, 2.0), (2, 2)) + ) + + sparse_input = sparse_tensor.SparseTensorValue( + # example 0, values [[0.], [1]] + # example 1, [[10.]] + indices=((0, 0), (0, 1), (1, 0)), + values=(0., 1., 10.), + dense_shape=(2, 2)) + + # Before _increment_two: + # [[0.], [1.]], + # [[10.], [0.]], + # After _increment_two: + # [[2.], [1.]], + # [[10.], [2.]], + expected_dense_tensor = [ + [[2.], [1.]], + [[10.], [2.]], + ] + numeric_column = sfc.sequence_numeric_column( + 'aaa', normalizer_fn=_increment_two) + + dense_tensor, _ = _get_sequence_dense_tensor( + numeric_column, {'aaa': sparse_input}) + + with monitored_session.MonitoredSession() as sess: + self.assertAllEqual( + expected_dense_tensor, dense_tensor.eval(session=sess)) + + @parameterized.named_parameters( + {'testcase_name': '2D', + 'sparse_input_args': { + # example 0, values [[[0., 1.], [2., 3.]], [[4., 5.], [6., 7.]]] + # example 1, [[[10., 11.], [12., 13.]]] + 'indices': ((0, 0), (0, 1), (0, 2), (0, 3), (0, 4), (0, 5), (0, 6), + (0, 7), (1, 0), (1, 1), (1, 2), (1, 3)), + 'values': (0., 1., 2., 3., 4., 5., 6., 7., 10., 11., 12., 13.), + 'dense_shape': (2, 8)}, + 'expected_dense_tensor': [ + [[[0., 1.], [2., 3.]], [[4., 5.], [6., 7.]]], + [[[10., 11.], [12., 13.]], [[0., 0.], [0., 0.]]]]}, + {'testcase_name': '3D', + 'sparse_input_args': { + 'indices': ((0, 0, 0), (0, 0, 2), (0, 0, 4), (0, 0, 6), + (0, 1, 0), (0, 1, 2), (0, 1, 4), (0, 1, 6), + (1, 0, 0), (1, 0, 2), (1, 0, 4), (1, 0, 6)), + 'values': (0., 1., 2., 3., 4., 5., 6., 7., 10., 11., 12., 13.), + 'dense_shape': (2, 2, 8)}, + 'expected_dense_tensor': [ + [[[0., 0.], [1., 0.]], [[2., 0.], [3., 0.]], + [[4., 0.], [5., 0.]], [[6., 0.], [7., 0.]]], + [[[10., 0.], [11., 0.]], [[12., 0.], [13., 0.]], + [[0., 0.], [0., 0.]], [[0., 0.], [0., 0.]]]]}, + ) + def test_get_dense_tensor_multi_dim( + self, sparse_input_args, expected_dense_tensor): + """Tests get_sequence_dense_tensor for multi-dim numeric_column.""" + sparse_input = sparse_tensor.SparseTensorValue(**sparse_input_args) + numeric_column = sfc.sequence_numeric_column('aaa', shape=(2, 2)) + + dense_tensor, _ = _get_sequence_dense_tensor( + numeric_column, {'aaa': sparse_input}) + + with monitored_session.MonitoredSession() as sess: + self.assertAllEqual( + expected_dense_tensor, dense_tensor.eval(session=sess)) + + @parameterized.named_parameters( + {'testcase_name': '2D', + 'inputs_args': { + # example 0, ids [2] + # example 1, ids [0, 1] + 'indices': ((0, 0), (1, 0), (1, 1)), + 'values': (2., 0., 1.), + 'dense_shape': (2, 2)}, + 'expected_sequence_length': [1, 2], + 'shape': (1,)}, + {'testcase_name': '3D', + 'inputs_args': { + # example 0, ids [[2]] + # example 1, ids [[0, 1], [2]] + 'indices': ((0, 0, 0), (1, 0, 0), (1, 0, 1), (1, 1, 0)), + 'values': (2., 0., 1., 2.), + 'dense_shape': (2, 2, 2)}, + 'expected_sequence_length': [1, 2], + 'shape': (1,)}, + {'testcase_name': '2D_with_shape', + 'inputs_args': { + # example 0, ids [2] + # example 1, ids [0, 1] + 'indices': ((0, 0), (1, 0), (1, 1)), + 'values': (2., 0., 1.), + 'dense_shape': (2, 2)}, + 'expected_sequence_length': [1, 1], + 'shape': (2,)}, + {'testcase_name': '3D_with_shape', + 'inputs_args': { + # example 0, ids [[2]] + # example 1, ids [[0, 1], [2]] + 'indices': ((0, 0, 0), (1, 0, 0), (1, 0, 1), (1, 1, 0)), + 'values': (2., 0., 1., 2.), + 'dense_shape': (2, 2, 2)}, + 'expected_sequence_length': [1, 2], + 'shape': (2,)}, + ) + def test_sequence_length(self, inputs_args, expected_sequence_length, shape): + inputs = sparse_tensor.SparseTensorValue(**inputs_args) + numeric_column = sfc.sequence_numeric_column('aaa', shape=shape) + + _, sequence_length = _get_sequence_dense_tensor( + numeric_column, {'aaa': inputs}) + + with monitored_session.MonitoredSession() as sess: + sequence_length = sess.run(sequence_length) + self.assertAllEqual(expected_sequence_length, sequence_length) + self.assertEqual(np.int64, sequence_length.dtype) + + def test_sequence_length_with_empty_rows(self): + """Tests _sequence_length when some examples do not have ids.""" + sparse_input = sparse_tensor.SparseTensorValue( + # example 0, values [] + # example 1, values [[0.], [1.]] + # example 2, [[2.]] + # example 3, values [] + # example 4, [[3.]] + # example 5, values [] + indices=((1, 0), (1, 1), (2, 0), (4, 0)), + values=(0., 1., 2., 3.), + dense_shape=(6, 2)) + expected_sequence_length = [0, 2, 1, 0, 1, 0] + numeric_column = sfc.sequence_numeric_column('aaa') + + _, sequence_length = _get_sequence_dense_tensor( + numeric_column, {'aaa': sparse_input}) + + with monitored_session.MonitoredSession() as sess: + self.assertAllEqual( + expected_sequence_length, sequence_length.eval(session=sess)) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/framework/BUILD b/tensorflow/contrib/framework/BUILD index 249debbdf6dff412a5be6cb1032fc4a3567c7d0b..3f6dbe0cbdeeae5e2107755f80bcfe5f7fc310e4 100644 --- a/tensorflow/contrib/framework/BUILD +++ b/tensorflow/contrib/framework/BUILD @@ -1,15 +1,16 @@ # Description: # contains parts of TensorFlow that are experimental or unstable and which are not supported. -licenses(["notice"]) # Apache 2.0 - -exports_files(["LICENSE"]) - package(default_visibility = [ "//learning/brain:__subpackages__", "//tensorflow:__subpackages__", + "//tensorflow_model_optimization:__subpackages__", ]) +licenses(["notice"]) # Apache 2.0 + +exports_files(["LICENSE"]) + load("//tensorflow:tensorflow.bzl", "py_test") load("//tensorflow:tensorflow.bzl", "tf_custom_op_library") load("//tensorflow:tensorflow.bzl", "tf_gen_op_wrapper_py") @@ -46,6 +47,13 @@ tf_custom_op_py_library( ":variable_ops_op_lib", ], srcs_version = "PY2AND3", + visibility = [ + "//learning/brain:__subpackages__", + "//tensorflow:__subpackages__", + "//tensorflow_estimator:__subpackages__", + "//tensorflow_model_optimization:__subpackages__", + "//video/youtube/personalization:__subpackages__", + ], deps = [ ":gen_variable_ops", "//tensorflow/contrib/util:util_py", @@ -65,6 +73,7 @@ tf_custom_op_py_library( "//tensorflow/python:resource_variable_ops", "//tensorflow/python:script_ops", "//tensorflow/python:smart_cond", + "//tensorflow/python:sort_ops", "//tensorflow/python:sparse_tensor", "//tensorflow/python:state_ops", "//tensorflow/python:state_ops_gen", @@ -310,17 +319,3 @@ py_test( "//third_party/py/numpy", ], ) - -py_test( - name = "sort_ops_test", - size = "medium", - srcs = ["python/ops/sort_ops_test.py"], - srcs_version = "PY2AND3", - deps = [ - ":framework_py", - "//tensorflow/python:array_ops", - "//tensorflow/python:client_testlib", - "//tensorflow/python:random_ops", - "//third_party/py/numpy", - ], -) diff --git a/tensorflow/contrib/framework/__init__.py b/tensorflow/contrib/framework/__init__.py index 95f5ba90aba6ff8d3f1f5b93bde2211ddf1c231b..3784631dcbfbeb215b6c695e4b6f1bbd02fa708c 100644 --- a/tensorflow/contrib/framework/__init__.py +++ b/tensorflow/contrib/framework/__init__.py @@ -15,10 +15,6 @@ """Framework utilities. -See the -[Contrib Framework](https://tensorflow.org/api_guides/python/contrib.framework) -guide. - @@assert_same_float_dtype @@assert_scalar @@assert_scalar_int @@ -134,17 +130,21 @@ _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', '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/framework/experimental_test.py b/tensorflow/contrib/framework/python/framework/experimental_test.py index cfdc7df7d8fd4c1406bf447a79038ac33b11e047..00e04b83ac45a83e54eee7a6e4e146fb683c3d98 100644 --- a/tensorflow/contrib/framework/python/framework/experimental_test.py +++ b/tensorflow/contrib/framework/python/framework/experimental_test.py @@ -44,17 +44,18 @@ class ExperimentalTest(test.TestCase): # Assert function docs are properly updated. self.assertEqual("_fn", _fn.__name__) - self.assertEqual("fn doc. (experimental)" - "\n" - "\nTHIS FUNCTION IS EXPERIMENTAL. It may change or " - "be removed at any time, and without warning." - "\n" - "\nArgs:" - "\n arg0: Arg 0." - "\n arg1: Arg 1." - "\n" - "\nReturns:" - "\n Sum of args.", _fn.__doc__) + self.assertEqual( + "fn doc. (experimental)" + "\n" + "\nWarning: THIS FUNCTION IS EXPERIMENTAL. It may change " + "or be removed at any time, and without warning." + "\n" + "\nArgs:" + "\n arg0: Arg 0." + "\n arg1: Arg 1." + "\n" + "\nReturns:" + "\n Sum of args.", _fn.__doc__) # Assert calling new fn issues log warning. self.assertEqual(3, _fn(1, 2)) diff --git a/tensorflow/contrib/framework/python/framework/tensor_util_test.py b/tensorflow/contrib/framework/python/framework/tensor_util_test.py index 9b0b9b1e1bf51db9332806097c2b3ae14d0587ad..05788d2e820a6cc1ec67578f0d1b19448b674d2f 100644 --- a/tensorflow/contrib/framework/python/framework/tensor_util_test.py +++ b/tensorflow/contrib/framework/python/framework/tensor_util_test.py @@ -218,7 +218,6 @@ class WithShapeTest(test.TestCase): self.assertRaisesRegexp(errors_impl.OpError, "Wrong shape", tensor_2x2.eval, {tensor_no_shape: [42.0]}) - @test_util.enable_c_shapes def test_with_shape_partial(self): with self.cached_session(): tensor_partial_shape = array_ops.placeholder(dtypes.float32) diff --git a/tensorflow/contrib/framework/python/ops/sort_ops.py b/tensorflow/contrib/framework/python/ops/sort_ops.py index 1921a77c1e96ee3531d1ed0f98e41c27c9d427ac..42184a4e55e292f7921702e3f8909ae54f717702 100644 --- a/tensorflow/contrib/framework/python/ops/sort_ops.py +++ b/tensorflow/contrib/framework/python/ops/sort_ops.py @@ -22,173 +22,7 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import numpy as np +from tensorflow.python.ops import sort_ops -from tensorflow.python.framework import constant_op -from tensorflow.python.framework import ops as framework_ops -from tensorflow.python.framework import tensor_util -from tensorflow.python.ops import array_ops -from tensorflow.python.ops import math_ops -from tensorflow.python.ops import nn_ops - - -def sort(values, axis=-1, direction='ASCENDING', name=None): - """Sorts a tensor. - - Args: - values: 1-D or higher numeric `Tensor`. - axis: The axis along which to sort. The default is -1, which sorts the last - axis. - direction: The direction in which to sort the values (`'ASCENDING'` or - `'DESCENDING'`). - name: Optional name for the operation. - - Returns: - A `Tensor` with the same dtype and shape as `values`, with the elements - sorted along the given `axis`. - - Raises: - ValueError: If axis is not a constant scalar, or the direction is invalid. - """ - with framework_ops.name_scope(name, 'sort'): - return _sort_or_argsort(values, axis, direction, return_argsort=False) - - -def argsort(values, axis=-1, direction='ASCENDING', stable=False, name=None): - """Returns the indices of a tensor that give its sorted order along an axis. - - For a 1D tensor, `tf.gather(values, tf.argsort(values))` is equivalent to - `tf.sort(values)`. For higher dimensions, the output has the same shape as - `values`, but along the given axis, values represent the index of the sorted - element in that slice of the tensor at the given position. - - Args: - values: 1-D or higher numeric `Tensor`. - axis: The axis along which to sort. The default is -1, which sorts the last - axis. - direction: The direction in which to sort the values (`'ASCENDING'` or - `'DESCENDING'`). - stable: If True, equal elements in the original tensor will not be - re-ordered in the returned order. Unstable sort is not yet implemented, - but will eventually be the default for performance reasons. If you - require a stable order, pass `stable=True` for forwards compatibility. - name: Optional name for the operation. - - Returns: - An int32 `Tensor` with the same shape as `values`. The indices that would - sort each slice of the given `values` along the given `axis`. - - Raises: - ValueError: If axis is not a constant scalar, or the direction is invalid. - """ - del stable # Unused. - with framework_ops.name_scope(name, 'argsort'): - return _sort_or_argsort(values, axis, direction, return_argsort=True) - - -def _sort_or_argsort(values, axis, direction, return_argsort): - """Internal sort/argsort implementation. - - Args: - values: The input values. - axis: The axis along which to sort. - direction: 'ASCENDING' or 'DESCENDING'. - return_argsort: Whether to return the argsort result. - - Returns: - Either the sorted values, or the indices of the sorted values in the - original tensor. See the `sort` and `argsort` docstrings. - - Raises: - ValueError: If axis is not a constant scalar, or the direction is invalid. - """ - if direction not in _SORT_IMPL: - raise ValueError('%s should be one of %s' % - (direction, ', '.join(sorted(_SORT_IMPL.keys())))) - # Axis must be an integer, not a Tensor. - axis = framework_ops.convert_to_tensor(axis, name='axis') - axis_static = tensor_util.constant_value(axis) - if axis.shape.ndims != 0 or axis_static is None: - raise ValueError('axis must be a constant scalar') - axis_static = int(axis_static) # Avoids NumPy casting error - - values = framework_ops.convert_to_tensor(values, name='values') - - return _SORT_IMPL[direction](values, axis_static, return_argsort) - - -def _descending_sort(values, axis, return_argsort=False): - """Sorts values in reverse using `top_k`. - - Args: - values: Tensor of numeric values. - axis: Index of the axis which values should be sorted along. - return_argsort: If False, return the sorted values. If True, return the - indices that would sort the values. - - Returns: - The sorted values. - """ - k = array_ops.shape(values)[axis] - rank = array_ops.rank(values) - static_rank = values.shape.ndims - # Fast path: sorting the last axis. - if axis == -1 or axis + 1 == values.get_shape().ndims: - top_k_input = values - transposition = None - else: - # Otherwise, transpose the array. Swap axes `axis` and `rank - 1`. - if axis < 0: - # Calculate the actual axis index if counting from the end. Use the static - # rank if available, or else make the axis back into a tensor. - axis += static_rank or rank - if static_rank is not None: - # Prefer to calculate the transposition array in NumPy and make it a - # constant. - transposition = constant_op.constant( - np.r_[ - # Axes up to axis are unchanged. - np.arange(axis), - # Swap axis and rank - 1. - [static_rank - 1], - # Axes in [axis + 1, rank - 1) are unchanged. - np.arange(axis + 1, static_rank - 1), - # Swap axis and rank - 1. - [axis]], - name='transposition') - else: - # Generate the transposition array from the tensors. - transposition = array_ops.concat( - [ - # Axes up to axis are unchanged. - math_ops.range(axis), - # Swap axis and rank - 1. - [rank - 1], - # Axes in [axis + 1, rank - 1) are unchanged. - math_ops.range(axis + 1, rank - 1), - # Swap axis and rank - 1. - [axis] - ], - axis=0) - top_k_input = array_ops.transpose(values, transposition) - - values, indices = nn_ops.top_k(top_k_input, k) - return_value = indices if return_argsort else values - if transposition is not None: - # transposition contains a single cycle of length 2 (swapping 2 elements), - # so it is an involution (it is its own inverse). - return_value = array_ops.transpose(return_value, transposition) - return return_value - - -def _ascending_sort(values, axis, return_argsort=False): - # Negate the values to get the ascending order from descending sort. - values_or_indices = _descending_sort(-values, axis, return_argsort) - # If not argsort, negate the values again. - return values_or_indices if return_argsort else -values_or_indices - - -_SORT_IMPL = { - 'ASCENDING': _ascending_sort, - 'DESCENDING': _descending_sort, -} +sort = sort_ops.sort +argsort = sort_ops.argsort 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 93b1aaa85e88e00c1b12a388321a4d6fb10f1611..b6b75ffa248d66cc4cb49339f193d486f05a6a4a 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" @@ -522,7 +522,7 @@ void LaunchFusedConv2DBiasActivationOp:: auto bias_ptr = AsDeviceMemory(bias.template flat().data(), bias.template flat().size()); - static int64 ConvolveScratchSize = GetCudnnWorkspaceLimit( + static int64 ConvolveScratchSize = GetDnnWorkspaceLimit( // default value is in bytes despite the name of the environment variable "TF_CUDNN_WORKSPACE_LIMIT_IN_MB", 1LL << 32 // 4GB ); @@ -570,7 +570,7 @@ void LaunchFusedConv2DBiasActivationOp:: for (auto profile_algorithm : algorithms) { // TODO(zhengxq): profile each algorithm multiple times to better // accuracy. - CudnnScratchAllocator scratch_allocator(ConvolveScratchSize, ctx); + DnnScratchAllocator scratch_allocator(ConvolveScratchSize, ctx); dnn::ProfileResult profile_result; bool cudnn_launch_status = stream @@ -609,7 +609,7 @@ void LaunchFusedConv2DBiasActivationOp:: algorithm_config); } - CudnnScratchAllocator scratch_allocator(ConvolveScratchSize, ctx); + DnnScratchAllocator scratch_allocator(ConvolveScratchSize, ctx); bool cudnn_launch_status = stream ->ThenFusedConvolveWithAlgorithm( diff --git a/tensorflow/contrib/gan/BUILD b/tensorflow/contrib/gan/BUILD index 9d0e6e1335d0be3477b78abce94999122672ff05..db0868fb2c43464a811b3d6dfcd96480ba2463ee 100644 --- a/tensorflow/contrib/gan/BUILD +++ b/tensorflow/contrib/gan/BUILD @@ -1,12 +1,14 @@ -# Files for using TFGAN framework. -package(default_visibility = ["//tensorflow:__subpackages__"]) +# Files for using TF-GAN framework. +load("//tensorflow:tensorflow.bzl", "py_test") + +package(default_visibility = [ + "//tensorflow:__subpackages__", +]) licenses(["notice"]) # Apache 2.0 exports_files(["LICENSE"]) -load("//tensorflow:tensorflow.bzl", "py_test") - py_library( name = "gan", srcs = [ @@ -49,7 +51,6 @@ py_library( "//tensorflow/python:training", "//tensorflow/python:training_util", "//tensorflow/python:variable_scope", - "//tensorflow/python/ops/distributions", "//tensorflow/python/ops/losses", ], ) @@ -105,7 +106,9 @@ py_library( deps = [ ":gan_estimator", ":head", + ":latent_gan_estimator", ":stargan_estimator", + ":tpu_gan_estimator", "//tensorflow/python:util", ], ) @@ -129,6 +132,7 @@ py_library( ":clip_weights", ":conditioning_utils", ":random_tensor_pool", + ":spectral_normalization", ":virtual_batchnorm", "//tensorflow/python:util", ], @@ -142,16 +146,15 @@ py_library( "//tensorflow/contrib/framework:framework_py", "//tensorflow/python:array_ops", "//tensorflow/python:clip_ops", + "//tensorflow/python:dtypes", "//tensorflow/python:framework_ops", - "//tensorflow/python:gradients", + "//tensorflow/python:gradients_impl", "//tensorflow/python:math_ops", "//tensorflow/python:random_ops", "//tensorflow/python:summary", "//tensorflow/python:tensor_util", "//tensorflow/python:variable_scope", - "//tensorflow/python/ops/distributions", "//tensorflow/python/ops/losses", - "//third_party/py/numpy", ], ) @@ -519,15 +522,19 @@ py_test( "//tensorflow/python:array_ops", "//tensorflow/python:client_testlib", "//tensorflow/python:dtypes", + "//tensorflow/python:errors", "//tensorflow/python:framework_ops", "//tensorflow/python:math_ops", "//tensorflow/python:metrics", "//tensorflow/python:parsing_ops", "//tensorflow/python:summary", + "//tensorflow/python:tensor_shape", "//tensorflow/python:training", "//tensorflow/python:training_util", "//tensorflow/python:variable_scope", - "//tensorflow/python/estimator:estimator_py", + "//tensorflow/python/estimator", + "//tensorflow/python/estimator:model_fn", + "//tensorflow/python/estimator:numpy_io", "//third_party/py/numpy", "@absl_py//absl/testing:parameterized", "@six_archive//:six", @@ -563,28 +570,114 @@ py_test( deps = [ ":namedtuples", ":stargan_estimator", - ":tuple_losses", "//tensorflow/contrib/layers:layers_py", - "//tensorflow/contrib/learn", - "//tensorflow/core:protos_all_py", "//tensorflow/python:array_ops", "//tensorflow/python:client_testlib", - "//tensorflow/python:dtypes", "//tensorflow/python:framework_ops", "//tensorflow/python:math_ops", "//tensorflow/python:metrics", - "//tensorflow/python:parsing_ops", "//tensorflow/python:summary", "//tensorflow/python:training", "//tensorflow/python:training_util", "//tensorflow/python:variable_scope", - "//tensorflow/python/estimator:estimator_py", + "//tensorflow/python/estimator:model_fn", + "//tensorflow/python/estimator:numpy_io", + "//third_party/py/numpy", + "@absl_py//absl/testing:parameterized", + "@six_archive//:six", + ], +) + +py_library( + name = "tpu_gan_estimator", + srcs = [ + "python/estimator/python/tpu_gan_estimator.py", + "python/estimator/python/tpu_gan_estimator_impl.py", + ], + srcs_version = "PY2AND3", + deps = [ + ":gan_estimator", + ":namedtuples", + ":train", + "//tensorflow/contrib/tpu:tpu_estimator", + "//tensorflow/contrib/tpu:tpu_lib", + "//tensorflow/contrib/training:training_py", + "//tensorflow/python:control_flow_ops", + "//tensorflow/python:framework_ops", + "//tensorflow/python:metrics", + "//tensorflow/python:util", + "//tensorflow/python/estimator:model_fn", + "//tensorflow/python/ops/losses", + ], +) + +py_test( + name = "tpu_gan_estimator_test", + srcs = ["python/estimator/python/tpu_gan_estimator_test.py"], + shard_count = 11, + srcs_version = "PY2AND3", + tags = ["notsan"], + deps = [ + ":namedtuples", + ":tpu_gan_estimator", + ":tuple_losses", + "//tensorflow/contrib/layers:layers_py", + "//tensorflow/contrib/tpu:tpu_estimator", + "//tensorflow/contrib/tpu:tpu_lib", + "//tensorflow/python:array_ops", + "//tensorflow/python:client_testlib", + "//tensorflow/python:errors", + "//tensorflow/python:framework_ops", + "//tensorflow/python:metrics", + "//tensorflow/python:summary", + "//tensorflow/python:tensor_shape", + "//tensorflow/python:training", + "//tensorflow/python:training_util", + "//tensorflow/python:variable_scope", + "//tensorflow/python/data/ops:dataset_ops", + "//tensorflow/python/estimator", + "//tensorflow/python/estimator:model_fn", "//third_party/py/numpy", "@absl_py//absl/testing:parameterized", "@six_archive//:six", ], ) +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 = [ @@ -619,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 9ab86329eaf0e6fd426aef1f552f4e27c2ad65de..4eac4e80cdacd779fdbedef19e4a654196f0caf1 100644 --- a/tensorflow/contrib/gan/README.md +++ b/tensorflow/contrib/gan/README.md @@ -1,14 +1,15 @@ -# TensorFlow-GAN (TFGAN) + +# TensorFlow-GAN (TF-GAN) -TFGAN is a lightweight library for training and evaluating Generative +TF-GAN is a lightweight library for training and evaluating Generative Adversarial Networks (GANs). This technique allows you to train a network (called the 'generator') to sample from a distribution, without having to explicitly model the distribution and without writing an explicit loss. For example, the generator could learn to draw samples from the distribution of natural images. For more details on this technique, see ['Generative Adversarial Networks'](https://arxiv.org/abs/1406.2661) by -Goodfellow et al. See [tensorflow/models](https://github.com/tensorflow/models/tree/master/research/gan/) for examples, and [this tutorial](https://github.com/tensorflow/models/tree/master/research/gan/tutorial.ipynb) for an +Goodfellow et al. See [tensorflow/models](https://github.com/tensorflow/models/tree/master/research/gan/) for examples, and [this tutorial](http://https://github.com/tensorflow/models/tree/master/research/gan/tutorial.ipynb) for an introduction. #### Usage @@ -17,27 +18,27 @@ import tensorflow as tf tfgan = tf.contrib.gan ``` -## Why TFGAN? +## Why TF-GAN? * Easily train generator and discriminator networks with well-tested, flexible [library calls](https://www.tensorflow.org/code/tensorflow/contrib/gan/python/train.py). You can -mix TFGAN, native TF, and other custom frameworks +mix TF-GAN, native TF, and other custom frameworks * Use already implemented [GAN losses and penalties](https://www.tensorflow.org/code/tensorflow/contrib/gan/python/losses/python/losses_impl.py) (ex Wasserstein loss, gradient penalty, mutual information penalty, etc) * [Monitor and visualize](https://www.tensorflow.org/code/tensorflow/contrib/gan/python/eval/python/summaries_impl.py) GAN progress during training, and [evaluate](https://www.tensorflow.org/code/tensorflow/contrib/gan/python/eval/python/classifier_metrics_impl.py) them * Use already-implemented [tricks](https://www.tensorflow.org/code/tensorflow/contrib/gan/python/features/python/) to stabilize and improve training * Develop based on examples of [common GAN setups](https://github.com/tensorflow/models/tree/master/research/gan/) -* Use the TFGAN-backed [GANEstimator](https://www.tensorflow.org/code/tensorflow/contrib/gan/python/estimator/python/gan_estimator_impl.py) to easily train a GAN model -* Improvements in TFGAN infrastructure will automatically benefit your TFGAN project +* Use the TF-GAN-backed [GANEstimator](https://www.tensorflow.org/code/tensorflow/contrib/gan/python/estimator/python/gan_estimator_impl.py) to easily train a GAN model +* Improvements in TF-GAN infrastructure will automatically benefit your TF-GAN project * Stay up-to-date with research as we add more algorithms -## What are the TFGAN components? +## What are the TF-GAN components? -TFGAN is composed of several parts which were design to exist independently. +TF-GAN is composed of several parts which were design to exist independently. These include the following main pieces (explained in detail below). * [core](https://www.tensorflow.org/code/tensorflow/contrib/gan/python/train.py): provides the main infrastructure needed to train a GAN. Training occurs in four phases, and each phase can be completed by custom-code or by using a - TFGAN library call. + TF-GAN library call. * [features](https://www.tensorflow.org/code/tensorflow/contrib/gan/python/features/python/): Many common GAN operations and normalization techniques are implemented for @@ -56,14 +57,14 @@ These include the following main pieces (explained in detail below). generative models. * [examples](https://github.com/tensorflow/models/tree/master/research/gan/) - and [tutorial](https://github.com/tensorflow/models/tree/master/research/gan/tutorial.ipynb): See examples of how to use TFGAN to make - GAN training easier, or use the more complicated examples to jumpstart your + 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 -Training in TFGAN typically consists of the following steps: +Training in TF-GAN typically consists of the following steps: 1. Specify the input to your networks. 1. Set up your generator and discriminator using a `GANModel`. @@ -71,12 +72,12 @@ Training in TFGAN typically consists of the following steps: 1. Create your train ops using a `GANTrainOps`. 1. Run your train ops. -At each stage, you can either use TFGAN's convenience functions, or you can +At each stage, you can either use TF-GAN's convenience functions, or you can perform the step manually for fine-grained control. We provide examples below. There are various types of GAN setups. For instance, you can train a generator to sample unconditionally from a learned distribution, or you can condition on -extra information such as a class label. TFGAN is compatible with many setups, +extra information such as a class label. TF-GAN is compatible with many setups, and we demonstrate a few below: ### Examples @@ -254,9 +255,9 @@ with variable_scope.variable_scope(dis_scope, reuse=True): discriminator_real_outputs = discriminator_fn(images) generator_variables = variables_lib.get_trainable_variables(gen_scope) discriminator_variables = variables_lib.get_trainable_variables(dis_scope) -# Depending on what TFGAN features you use, you don't always need to supply +# Depending on what TF-GAN features you use, you don't always need to supply # every `GANModel` field. At a minimum, you need to include the discriminator -# outputs and variables if you want to use TFGAN to construct losses. +# outputs and variables if you want to use TF-GAN to construct losses. gan_model = tfgan.GANModel( generator_inputs, generated_data, diff --git a/tensorflow/contrib/gan/__init__.py b/tensorflow/contrib/gan/__init__.py index f1946c7f925660eae3aaa650c437e03da1f33d6c..1e6000898f7b8a53ad3f6fa12deebd54bf3a57ff 100644 --- a/tensorflow/contrib/gan/__init__.py +++ b/tensorflow/contrib/gan/__init__.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""TFGAN is a lightweight library for training and evaluating GANs. +"""TF-GAN is a lightweight library for training and evaluating GANs. In addition to providing the infrastructure for easily training and evaluating GANS, this library contains modules for a TFGAN-backed Estimator, @@ -24,7 +24,7 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -# Collapse TFGAN into a tiered namespace. +# Collapse TF-GAN into a tiered namespace. from tensorflow.contrib.gan.python import estimator from tensorflow.contrib.gan.python import eval # pylint:disable=redefined-builtin from tensorflow.contrib.gan.python import features diff --git a/tensorflow/contrib/gan/python/estimator/__init__.py b/tensorflow/contrib/gan/python/estimator/__init__.py index 99d38011ba677f03e198a431634fbb2ce349f912..430266555b723e6ca39dccffc1442dbef5d4a385 100644 --- a/tensorflow/contrib/gan/python/estimator/__init__.py +++ b/tensorflow/contrib/gan/python/estimator/__init__.py @@ -12,10 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""TFGAN estimator module. +"""TF-GAN estimator module. GANEstimator provides all the infrastructure support of a TensorFlow Estimator -with the feature support of TFGAN. +with the feature support of TF-GAN. """ from __future__ import absolute_import @@ -26,18 +26,25 @@ 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 from tensorflow.python.util.all_util import remove_undocumented -_allowed_symbols = [ +_allowed_symbols = ([ 'gan_estimator', 'stargan_estimator', + 'tpu_gan_estimator', + 'latent_gan_estimator', 'head', -] + gan_estimator.__all__ + stargan_estimator.__all__ + head.__all__ +] + gan_estimator.__all__ + stargan_estimator.__all__ + head.__all__ + + tpu_gan_estimator.__all__ + latent_gan_estimator.__all__) remove_undocumented(__name__, _allowed_symbols) diff --git a/tensorflow/contrib/gan/python/estimator/python/gan_estimator_impl.py b/tensorflow/contrib/gan/python/estimator/python/gan_estimator_impl.py index 219cc199d79eca8c263859ae46bbb1ce0b4442b3..dd904611d1a3bb78de8316d5ed29ab0f800f29a9 100644 --- a/tensorflow/contrib/gan/python/estimator/python/gan_estimator_impl.py +++ b/tensorflow/contrib/gan/python/estimator/python/gan_estimator_impl.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""A TFGAN-backed GAN Estimator.""" +"""A TF-GAN-backed GAN Estimator.""" from __future__ import absolute_import from __future__ import division @@ -56,10 +56,10 @@ _summary_type_map = { class GANEstimator(estimator.Estimator): """An estimator for Generative Adversarial Networks (GANs). - This Estimator is backed by TFGAN. The network functions follow the TFGAN API - except for one exception: if either `generator_fn` or `discriminator_fn` have - an argument called `mode`, then the tf.Estimator mode is passed in for that - argument. This helps with operations like batch normalization, which have + This Estimator is backed by TF-GAN. The network functions follow the TF-GAN + API except for one exception: if either `generator_fn` or `discriminator_fn` + have an argument called `mode`, then the tf.Estimator mode is passed in for + that argument. This helps with operations like batch normalization, which have different train and evaluation behavior. Example: @@ -68,7 +68,7 @@ class GANEstimator(estimator.Estimator): import tensorflow as tf tfgan = tf.contrib.gan - # See TFGAN's `train.py` for a description of the generator and + # See TF-GAN's `train.py` for a description of the generator and # discriminator API. def generator_fn(generator_inputs): ... @@ -113,7 +113,8 @@ class GANEstimator(estimator.Estimator): add_summaries=None, use_loss_summaries=True, config=None, - warm_start_from=None): + warm_start_from=None, + is_chief=True): """Initializes a GANEstimator instance. Args: @@ -122,13 +123,13 @@ class GANEstimator(estimator.Estimator): to continue training a previously saved model. generator_fn: A python function that takes a Tensor, Tensor list, or Tensor dictionary as inputs and returns the outputs of the GAN - generator. See `TFGAN` for more details and examples. Additionally, if + generator. See `TF-GAN` for more details and examples. Additionally, if it has an argument called `mode`, the Estimator's `mode` will be passed in (ex TRAIN, EVAL, PREDICT). This is useful for things like batch normalization. discriminator_fn: A python function that takes the output of `generator_fn` or real data in the GAN setup, and `generator_inputs`. - Outputs a Tensor in the range [-inf, inf]. See `TFGAN` for more details + Outputs a Tensor in the range [-inf, inf]. See `TF-GAN` for more details and examples. generator_loss_fn: The loss function on the generator. Takes a `GANModel` tuple. @@ -154,6 +155,8 @@ class GANEstimator(estimator.Estimator): config: `RunConfig` object to configure the runtime settings. warm_start_from: A filepath to a checkpoint or saved model, or a WarmStartSettings object to configure initialization. + is_chief: Whether or not this Estimator is running on a chief or worker. + Needs to be set appropriately if using SyncReplicasOptimizers. Raises: ValueError: If loss functions aren't callable. @@ -187,7 +190,7 @@ class GANEstimator(estimator.Estimator): return _get_estimator_spec( mode, gan_model, generator_loss_fn, discriminator_loss_fn, get_eval_metric_ops_fn, generator_optimizer, discriminator_optimizer, - get_hooks_fn, use_loss_summaries) + get_hooks_fn, use_loss_summaries, is_chief) super(GANEstimator, self).__init__( model_fn=_model_fn, model_dir=model_dir, config=config, @@ -215,7 +218,7 @@ def _get_gan_model( def _get_estimator_spec( mode, gan_model, generator_loss_fn, discriminator_loss_fn, get_eval_metric_ops_fn, generator_optimizer, discriminator_optimizer, - get_hooks_fn=None, use_loss_summaries=True): + get_hooks_fn=None, use_loss_summaries=True, is_chief=True): """Get the EstimatorSpec for the current mode.""" if mode == model_fn_lib.ModeKeys.PREDICT: estimator_spec = model_fn_lib.EstimatorSpec( @@ -230,13 +233,14 @@ def _get_estimator_spec( estimator_spec = _get_eval_estimator_spec( gan_model, gan_loss, get_eval_metric_ops_fn) else: # model_fn_lib.ModeKeys.TRAIN: - gopt = (generator_optimizer() if callable(generator_optimizer) else - generator_optimizer) - dopt = (discriminator_optimizer() if callable(discriminator_optimizer) - else discriminator_optimizer) + if callable(generator_optimizer): + generator_optimizer = generator_optimizer() + if callable(discriminator_optimizer): + discriminator_optimizer = discriminator_optimizer() get_hooks_fn = get_hooks_fn or tfgan_train.get_sequential_train_hooks() estimator_spec = _get_train_estimator_spec( - gan_model, gan_loss, gopt, dopt, get_hooks_fn) + gan_model, gan_loss, generator_optimizer, discriminator_optimizer, + get_hooks_fn, is_chief=is_chief) return estimator_spec @@ -321,11 +325,11 @@ def _get_eval_estimator_spec(gan_model, gan_loss, get_eval_metric_ops_fn=None, def _get_train_estimator_spec( gan_model, gan_loss, generator_optimizer, discriminator_optimizer, - get_hooks_fn, train_op_fn=tfgan_train.gan_train_ops): + get_hooks_fn, train_op_fn=tfgan_train.gan_train_ops, is_chief=True): """Return an EstimatorSpec for the train case.""" scalar_loss = gan_loss.generator_loss + gan_loss.discriminator_loss train_ops = train_op_fn(gan_model, gan_loss, generator_optimizer, - discriminator_optimizer) + discriminator_optimizer, is_chief=is_chief) training_hooks = get_hooks_fn(train_ops) return model_fn_lib.EstimatorSpec( loss=scalar_loss, diff --git a/tensorflow/contrib/gan/python/estimator/python/gan_estimator_test.py b/tensorflow/contrib/gan/python/estimator/python/gan_estimator_test.py index cfc867f0831986ef517f14fee0ed9d4773bb5cb6..5b9c54e43a16adf457d5ed0e7e73dcd168ab0d67 100644 --- a/tensorflow/contrib/gan/python/estimator/python/gan_estimator_test.py +++ b/tensorflow/contrib/gan/python/estimator/python/gan_estimator_test.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Tests for TFGAN's estimator.py.""" +"""Tests for TF-GAN's estimator.py.""" from __future__ import absolute_import from __future__ import division @@ -37,6 +37,7 @@ from tensorflow.python.estimator.estimator import WarmStartSettings from tensorflow.python.estimator.inputs import numpy_io from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops +from tensorflow.python.framework import tensor_shape from tensorflow.python.framework.errors_impl import NotFoundError from tensorflow.python.ops import array_ops from tensorflow.python.ops import math_ops @@ -47,6 +48,7 @@ from tensorflow.python.platform import test from tensorflow.python.summary.writer import writer_cache from tensorflow.python.training import input as input_lib from tensorflow.python.training import learning_rate_decay +from tensorflow.python.training import sync_replicas_optimizer from tensorflow.python.training import training from tensorflow.python.training import training_util @@ -54,7 +56,8 @@ from tensorflow.python.training import training_util def generator_fn(noise_dict, mode): del mode noise = noise_dict['x'] - return layers.fully_connected(noise, noise.shape[1].value) + return layers.fully_connected(noise, tensor_shape.dimension_value( + noise.shape[1])) def discriminator_fn(data, unused_conditioning, mode): @@ -72,15 +75,15 @@ class GetGANModelTest(test.TestCase, parameterized.TestCase): def test_get_gan_model(self, mode): with ops.Graph().as_default(): generator_inputs = {'x': array_ops.ones([3, 4])} - real_data = (array_ops.zeros([3, 4]) if - mode != model_fn_lib.ModeKeys.PREDICT else None) + is_predict = mode == model_fn_lib.ModeKeys.PREDICT + real_data = array_ops.zeros([3, 4]) if not is_predict else None gan_model = estimator._get_gan_model( mode, generator_fn, discriminator_fn, real_data, generator_inputs, add_summaries=False) self.assertEqual(generator_inputs, gan_model.generator_inputs) self.assertIsNotNone(gan_model.generated_data) - self.assertEqual(2, len(gan_model.generator_variables)) # 1 FC layer + self.assertLen(gan_model.generator_variables, 2) # 1 FC layer self.assertIsNotNone(gan_model.generator_fn) if mode == model_fn_lib.ModeKeys.PREDICT: self.assertIsNone(gan_model.real_data) @@ -93,7 +96,7 @@ class GetGANModelTest(test.TestCase, parameterized.TestCase): self.assertIsNotNone(gan_model.real_data) self.assertIsNotNone(gan_model.discriminator_real_outputs) self.assertIsNotNone(gan_model.discriminator_gen_outputs) - self.assertEqual(2, len(gan_model.discriminator_variables)) # 1 FC layer + self.assertLen(gan_model.discriminator_variables, 2) # 1 FC layer self.assertIsNotNone(gan_model.discriminator_scope) self.assertIsNotNone(gan_model.discriminator_fn) @@ -119,6 +122,7 @@ def get_dummy_gan_model(): def dummy_loss_fn(gan_model, add_summaries=True): + del add_summaries return math_ops.reduce_sum(gan_model.discriminator_real_outputs - gan_model.discriminator_gen_outputs) @@ -135,6 +139,7 @@ class GetEstimatorSpecTest(test.TestCase, parameterized.TestCase): @classmethod def setUpClass(cls): + super(GetEstimatorSpecTest, cls).setUpClass() cls._generator_optimizer = training.GradientDescentOptimizer(1.0) cls._discriminator_optimizer = training.GradientDescentOptimizer(1.0) @@ -166,8 +171,36 @@ class GetEstimatorSpecTest(test.TestCase, parameterized.TestCase): self.assertShapeEqual(np.array(0), spec.loss) # must be a scalar self.assertIsNotNone(spec.eval_metric_ops) + def test_get_sync_estimator_spec(self): + """Make sure spec is loaded with sync hooks for sync opts.""" + + def get_sync_optimizer(): + return sync_replicas_optimizer.SyncReplicasOptimizer( + training.GradientDescentOptimizer(learning_rate=1.0), + replicas_to_aggregate=1) + + with ops.Graph().as_default(): + self._gan_model = get_dummy_gan_model() + g_opt = get_sync_optimizer() + d_opt = get_sync_optimizer() + + spec = estimator._get_estimator_spec( + model_fn_lib.ModeKeys.TRAIN, + self._gan_model, + generator_loss_fn=dummy_loss_fn, + discriminator_loss_fn=dummy_loss_fn, + get_eval_metric_ops_fn=get_metrics, + generator_optimizer=g_opt, + discriminator_optimizer=d_opt) + + self.assertLen(spec.training_hooks, 4) + sync_opts = [ + hook._sync_optimizer for hook in spec.training_hooks if + isinstance(hook, sync_replicas_optimizer._SyncReplicasOptimizerHook)] + self.assertLen(sync_opts, 2) + self.assertSetEqual(frozenset(sync_opts), frozenset((g_opt, d_opt))) + -# TODO(joelshor): Add pandas test. class GANEstimatorIntegrationTest(test.TestCase): def setUp(self): @@ -198,11 +231,11 @@ class GANEstimatorIntegrationTest(test.TestCase): get_eval_metric_ops_fn=get_metrics, model_dir=self._model_dir) - # TRAIN + # Train. num_steps = 10 est.train(train_input_fn, steps=num_steps) - # EVALUTE + # Evaluate. scores = est.evaluate(eval_input_fn) self.assertEqual(num_steps, scores[ops.GraphKeys.GLOBAL_STEP]) self.assertIn('loss', six.iterkeys(scores)) @@ -210,7 +243,7 @@ class GANEstimatorIntegrationTest(test.TestCase): scores['loss']) self.assertIn('mse_custom_metric', six.iterkeys(scores)) - # PREDICT + # Predict. predictions = np.array([x for x in est.predict(predict_input_fn)]) self.assertAllEqual(prediction_size, predictions.shape) diff --git a/tensorflow/contrib/gan/python/estimator/python/head_impl.py b/tensorflow/contrib/gan/python/estimator/python/head_impl.py index 1a0ee6dfc498eb6dc8c97411589d9e35bc352062..cbe990b476c3b17ce61e0826b17d10976fea43c7 100644 --- a/tensorflow/contrib/gan/python/estimator/python/head_impl.py +++ b/tensorflow/contrib/gan/python/estimator/python/head_impl.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""A TFGAN-backed GAN Estimator.""" +"""A TF-GAN-backed GAN Estimator.""" from __future__ import absolute_import from __future__ import division diff --git a/tensorflow/contrib/gan/python/estimator/python/head_test.py b/tensorflow/contrib/gan/python/estimator/python/head_test.py index 8205bc889dc01c8680e2139393d65723280cfbd0..5b50234a0e33cd297b176f142b358338966b6758 100644 --- a/tensorflow/contrib/gan/python/estimator/python/head_test.py +++ b/tensorflow/contrib/gan/python/estimator/python/head_test.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Tests for TFGAN's head.py.""" +"""Tests for TF-GAN's head.py.""" from __future__ import absolute_import from __future__ import division 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/stargan_estimator_impl.py b/tensorflow/contrib/gan/python/estimator/python/stargan_estimator_impl.py index f60e16bc04662b33bc0bb22b5acc8c7fcc7a03ba..2a485e7d47ff10cf34c1b44f8dcc6b1f33c9a05f 100644 --- a/tensorflow/contrib/gan/python/estimator/python/stargan_estimator_impl.py +++ b/tensorflow/contrib/gan/python/estimator/python/stargan_estimator_impl.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""A TFGAN-backed StarGAN Estimator.""" +"""A TF-GAN-backed StarGAN Estimator.""" from __future__ import absolute_import from __future__ import division diff --git a/tensorflow/contrib/gan/python/estimator/python/stargan_estimator_test.py b/tensorflow/contrib/gan/python/estimator/python/stargan_estimator_test.py index 2ec7938c7c4051842c7e982b54c1213b6e841b79..c00ff4399748a77f88d9753df7592bf3859d754e 100644 --- a/tensorflow/contrib/gan/python/estimator/python/stargan_estimator_test.py +++ b/tensorflow/contrib/gan/python/estimator/python/stargan_estimator_test.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Tests for TFGAN's stargan_estimator.py.""" +"""Tests for TF-GAN's stargan_estimator.py.""" from __future__ import absolute_import from __future__ import division @@ -80,7 +80,7 @@ class StarGetGANModelTest(test.TestCase, parameterized.TestCase): self.assertEqual(input_data, gan_model.input_data) self.assertIsNotNone(gan_model.generated_data) self.assertIsNotNone(gan_model.generated_data_domain_target) - self.assertEqual(1, len(gan_model.generator_variables)) + self.assertLen(gan_model.generator_variables, 1) self.assertIsNotNone(gan_model.generator_scope) self.assertIsNotNone(gan_model.generator_fn) if mode == model_fn_lib.ModeKeys.PREDICT: @@ -109,7 +109,7 @@ class StarGetGANModelTest(test.TestCase, parameterized.TestCase): gan_model.discriminator_input_data_domain_predication) self.assertIsNotNone( gan_model.discriminator_generated_data_domain_predication) - self.assertEqual(2, len(gan_model.discriminator_variables)) # 1 FC layer + self.assertLen(gan_model.discriminator_variables, 2) # 1 FC layer self.assertIsNotNone(gan_model.discriminator_scope) self.assertIsNotNone(gan_model.discriminator_fn) @@ -163,6 +163,7 @@ class GetEstimatorSpecTest(test.TestCase, parameterized.TestCase): @classmethod def setUpClass(cls): + super(GetEstimatorSpecTest, cls).setUpClass() cls._generator_optimizer = training.GradientDescentOptimizer(1.0) cls._discriminator_optimizer = training.GradientDescentOptimizer(1.0) diff --git a/tensorflow/contrib/gan/python/estimator/python/tpu_gan_estimator.py b/tensorflow/contrib/gan/python/estimator/python/tpu_gan_estimator.py new file mode 100644 index 0000000000000000000000000000000000000000..deb381f7be3f9545ed918813ee55aede946f22d4 --- /dev/null +++ b/tensorflow/contrib/gan/python/estimator/python/tpu_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 `TPUGANEstimator`.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.gan.python.estimator.python import tpu_gan_estimator_impl +# pylint: disable=wildcard-import +from tensorflow.contrib.gan.python.estimator.python.tpu_gan_estimator_impl import * +# pylint: enable=wildcard-import +from tensorflow.python.util.all_util import remove_undocumented + +__all__ = tpu_gan_estimator_impl.__all__ +remove_undocumented(__name__, __all__) diff --git a/tensorflow/contrib/gan/python/estimator/python/tpu_gan_estimator_impl.py b/tensorflow/contrib/gan/python/estimator/python/tpu_gan_estimator_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..8f2a22c78a304c7cc66ef069a235483e9279b3b2 --- /dev/null +++ b/tensorflow/contrib/gan/python/estimator/python/tpu_gan_estimator_impl.py @@ -0,0 +1,423 @@ +# 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. +# ============================================================================== +"""A TF-GAN-backed GAN Estimator that works on TPU.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.gan.python import namedtuples as tfgan_tuples +from tensorflow.contrib.gan.python import train as tfgan_train +from tensorflow.contrib.gan.python.estimator.python import gan_estimator_impl as gan_estimator_lib +from tensorflow.contrib.tpu.python.tpu import tpu_estimator +from tensorflow.contrib.tpu.python.tpu import tpu_optimizer +from tensorflow.contrib.training.python.training import training +from tensorflow.python.estimator import model_fn as model_fn_lib +from tensorflow.python.framework import ops +from tensorflow.python.ops import control_flow_ops +from tensorflow.python.ops import metrics as metrics_lib +from tensorflow.python.ops.losses import losses + +__all__ = [ + 'TPUGANEstimator', +] + + +class TPUGANEstimator(tpu_estimator.TPUEstimator): + """An estimator for Generative Adversarial Networks (GANs) on TPU. + + This Estimator is backed by TFGAN. It is similar to `tfgan.GANEstimator`, + but works on TPU. + + Example: + + ```python + import tensorflow as tf + tfgan = tf.contrib.gan + + # See TFGAN's `train.py` for a description of the generator and + # discriminator API. + def generator_fn(generator_inputs): + ... + return generated_data + + def discriminator_fn(data, conditioning): + ... + return logits + + # Create GAN estimator. + config = tpu_config.RunConfig(model_dir='/my/dir') + gan_estimator = tfgan.estimator.TPUGANEstimator( + generator_fn=generator_fn, + discriminator_fn=discriminator_fn, + generator_loss_fn=tfgan.losses.wasserstein_generator_loss, + discriminator_loss_fn=tfgan.losses.wasserstein_discriminator_loss, + generator_optimizer=tf.train.AdamOptimizer(0.1, 0.5), + discriminator_optimizer=tf.train.AdamOptimizer(0.1, 0.5), + train_batch_size=4, + config=config) + + # Train estimator. + gan_estimator.train(train_input_fn, train_steps) + + # Evaluate resulting estimator. + gan_estimator.evaluate(eval_input_fn, eval_steps) + + # Generate samples from generator. + predictions = np.array([ + x['generated_data'] for x in gan_estimator.predict(predict_input_fn)]) + ``` + """ + + def __init__(self, + # Arguments to construct the `model_fn`. + generator_fn=None, + discriminator_fn=None, + generator_loss_fn=None, + discriminator_loss_fn=None, + generator_optimizer=None, + discriminator_optimizer=None, + get_eval_metric_ops_fn=None, + add_summaries=None, + joint_train=False, + gan_train_steps=tfgan_tuples.GANTrainSteps(1, 1), + # TPUEstimator options. + model_dir=None, + config=None, + params=None, + use_tpu=True, + train_batch_size=None, + eval_batch_size=None, + predict_batch_size=None, + batch_axis=None, + eval_on_tpu=True, + export_to_tpu=True, + warm_start_from=None): + """Initializes a TPUGANEstimator instance. + + Args: + generator_fn: A python function that takes a Tensor, Tensor list, or + Tensor dictionary as inputs and returns the outputs of the GAN + generator. See `TFGAN` for more details and examples. Additionally, if + it has an argument called `mode`, the Estimator's `mode` will be passed + in (ex TRAIN, EVAL, PREDICT). This is useful for things like batch + normalization. + discriminator_fn: A python function that takes the output of + `generator_fn` or real data in the GAN setup, and `generator_inputs`. + Outputs a Tensor in the range [-inf, inf]. See `TFGAN` for more details + and examples. + generator_loss_fn: The loss function on the generator. Takes a `GANModel` + tuple. + discriminator_loss_fn: The loss function on the discriminator. Takes a + `GANModel` tuple. + generator_optimizer: The optimizer for generator updates, or a function + that takes no arguments and returns an optimizer. This function will + be called when the default graph is the `GANEstimator`'s graph, so + utilities like `tf.contrib.framework.get_or_create_global_step` will + work. + discriminator_optimizer: Same as `generator_optimizer`, but for the + discriminator updates. + get_eval_metric_ops_fn: A function that takes a list of arguments and + returns a dict of metric results keyed by name. The output of this + function is passed into `tf.estimator.EstimatorSpec` during evaluation. + The arguments must be: + * generator_inputs + * generated_data + * real_data + * discriminator_real_outputs + * discriminator_gen_outputs + add_summaries: `None`, a single `SummaryType`, or a list of `SummaryType`. + This is ignored for jobs that run on TPU, such as the train job if + `use_tpu` is `True` or the eval job if `eval_on_tpu` is `True`. + joint_train: A Python boolean. If `True`, jointly train the generator and + the discriminator. If `False`, sequentially train them. See `train.py` + in TFGAN for more details on the differences between the two GAN + training methods. + gan_train_steps: A `tfgan.GANTrainSteps` named tuple describing the ratio + of generator to discriminator steps. For now, only supports 1:1 + training. + model_dir: Same as `TPUEstimator`: Directory to save model parameters, + graph and etc. This can also be used to load checkpoints from the + directory into a estimator to continue training a previously saved + model. If `None`, the model_dir in `config` will be used if set. If both + are set, they must be same. If both are `None`, a temporary directory + will be used. + config: Same as `TPUEstimator`: An `tpu_config.RunConfig` configuration + object. Cannot be `None`. + params: Same as `TPUEstimator`: An optional `dict` of hyper parameters + that will be passed into `input_fn` and `model_fn`. Keys are names of + parameters, values are basic python types. There are reserved keys for + `TPUEstimator`, including 'batch_size'. + use_tpu: Same as `TPUEstimator`: A bool indicating whether TPU support is + enabled. Currently, TPU training and evaluation respect this bit, but + eval_on_tpu can override execution of eval. See below. Predict still + happens on CPU. + train_batch_size: Same as `TPUEstimator`: An int representing the global + training batch size. TPUEstimator transforms this global batch size to a + per-shard batch size, as params['batch_size'], when calling `input_fn` + and `model_fn`. Cannot be `None` if `use_tpu` is `True`. Must be + divisible by total number of replicas. + eval_batch_size: Same as `TPUEstimator`: An int representing evaluation + batch size. Must be divisible by total number of replicas. + predict_batch_size: Same as `TPUEstimator`: An int representing the + prediction batch size. Must be divisible by total number of replicas. + batch_axis: Same as `TPUEstimator`: A python tuple of int values + describing how each tensor produced by the Estimator `input_fn` should + be split across the TPU compute shards. For example, if your input_fn + produced (images, labels) where the images tensor is in `HWCN` format, + your shard dimensions would be [3, 0], where 3 corresponds to the `N` + dimension of your images Tensor, and 0 corresponds to the dimension + along which to split the labels to match up with the corresponding + images. If None is supplied, and per_host_input_for_training is True, + batches will be sharded based on the major dimension. If + tpu_config.per_host_input_for_training is False or `PER_HOST_V2`, + batch_axis is ignored. + eval_on_tpu: Same as `TPUEstimator`: 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: Same as `TPUEstimator`: If True, `export_savedmodel()` + exports a metagraph for serving on TPU besides the one on CPU. + warm_start_from: Same as `TPUEstimator`: Optional string filepath to a + checkpoint or SavedModel to warm-start from, or a + `tf.estimator.WarmStartSettings` object to fully configure + warm-starting. If the string filepath is provided instead of a + `WarmStartSettings`, then all variables are warm-started, and it is + assumed that vocabularies and Tensor names are unchanged. + + Raises: + ValueError: If loss functions aren't callable. + ValueError: If `gan_train_steps` isn't a `tfgan_tuples.GANTrainSteps` + tuple. + ValueError: If `gan_train_steps` isn't 1:1 training. + """ + if not callable(generator_loss_fn): + raise ValueError('generator_loss_fn must be callable.') + if not callable(discriminator_loss_fn): + raise ValueError('discriminator_loss_fn must be callable.') + if not isinstance(gan_train_steps, tfgan_tuples.GANTrainSteps): + raise ValueError( + '`gan_train_steps` must be `tfgan_tuples.GANTrainSteps`. Instead, ' + 'was type: %s' % type(gan_train_steps)) + if (gan_train_steps.generator_train_steps != 1 or + gan_train_steps.discriminator_train_steps != 1): + raise ValueError('Estimator currently only supports 1:1 training.') + + if use_tpu: + generator_optimizer = _maybe_make_cross_shard_optimizer( + generator_optimizer) + discriminator_optimizer = _maybe_make_cross_shard_optimizer( + discriminator_optimizer) + + def _model_fn(features, labels, mode, params): + """GANEstimator model function.""" + del params # unused + if mode not in [model_fn_lib.ModeKeys.TRAIN, model_fn_lib.ModeKeys.EVAL, + model_fn_lib.ModeKeys.PREDICT]: + raise ValueError('Mode not recognized: %s' % mode) + real_data = labels # rename inputs for clarity + generator_inputs = features # rename inputs for clarity + + # Make GANModel, which encapsulates the GAN model architectures. + # TODO(joelshor): Switch TF-GAN over to TPU-compatible summaries, then + # remove `add_summaries` logic below. + is_on_tpu = _is_on_tpu(mode, use_tpu, eval_on_tpu) + gan_model = gan_estimator_lib._get_gan_model( # pylint:disable=protected-access + mode, generator_fn, discriminator_fn, real_data, generator_inputs, + add_summaries=None if is_on_tpu else add_summaries) + + # Make the TPUEstimatorSpec, which incorporates the GANModel, losses, eval + # metrics, and optimizers (if required). + estimator_spec = _get_estimator_spec( + mode, gan_model, generator_loss_fn, discriminator_loss_fn, + get_eval_metric_ops_fn, generator_optimizer, discriminator_optimizer, + joint_train, is_on_tpu, gan_train_steps) + assert isinstance(estimator_spec, tpu_estimator.TPUEstimatorSpec) + return estimator_spec + + super(TPUGANEstimator, self).__init__( + model_fn=_model_fn, + model_dir=model_dir, + config=config, + params=params, + use_tpu=use_tpu, + train_batch_size=train_batch_size, + eval_batch_size=eval_batch_size, + predict_batch_size=predict_batch_size, + batch_axis=batch_axis, + eval_on_tpu=eval_on_tpu, + export_to_tpu=export_to_tpu, + warm_start_from=warm_start_from) + + +def _is_on_tpu(mode, use_tpu, eval_on_tpu): + if mode == model_fn_lib.ModeKeys.TRAIN: + return use_tpu + elif mode == model_fn_lib.ModeKeys.EVAL: + return eval_on_tpu + else: + return False + + +def _get_estimator_spec( + mode, gan_model, generator_loss_fn, discriminator_loss_fn, + get_eval_metric_ops_fn, generator_optimizer, discriminator_optimizer, + joint_train, is_on_tpu, gan_train_steps): + """Get the TPUEstimatorSpec for the current mode.""" + if mode == model_fn_lib.ModeKeys.PREDICT: + estimator_spec = tpu_estimator.TPUEstimatorSpec( + mode=mode, predictions={'generated_data': gan_model.generated_data}) + elif mode == model_fn_lib.ModeKeys.EVAL: + gan_loss = tfgan_tuples.GANLoss( + generator_loss=generator_loss_fn( + gan_model, add_summaries=not is_on_tpu), + discriminator_loss=discriminator_loss_fn( + gan_model, add_summaries=not is_on_tpu)) + # Eval losses for metrics must preserve batch dimension. + gan_loss_no_reduction = tfgan_tuples.GANLoss( + generator_loss=generator_loss_fn( + gan_model, add_summaries=False, reduction=losses.Reduction.NONE), + discriminator_loss=discriminator_loss_fn( + gan_model, add_summaries=False, reduction=losses.Reduction.NONE)) + estimator_spec = _get_eval_estimator_spec( + gan_model, gan_loss, gan_loss_no_reduction, get_eval_metric_ops_fn) + else: # model_fn_lib.ModeKeys.TRAIN: + gan_loss = tfgan_tuples.GANLoss( + generator_loss=generator_loss_fn( + gan_model, add_summaries=not is_on_tpu), + discriminator_loss=discriminator_loss_fn( + gan_model, add_summaries=not is_on_tpu)) + + # Construct optimizers if arguments were callable. For TPUs, they must be + # `CrossShardOptimizer`. + g_callable = callable(generator_optimizer) + gopt = generator_optimizer() if g_callable else generator_optimizer + d_callable = callable(discriminator_optimizer) + dopt = discriminator_optimizer() if d_callable else discriminator_optimizer + + estimator_spec = _get_train_estimator_spec( + gan_model, gan_loss, gopt, dopt, joint_train, gan_train_steps) + + return estimator_spec + + +def _get_eval_estimator_spec(gan_model, gan_loss, gan_loss_no_reduction, + get_eval_metric_ops_fn): + """Return an TPUEstimatorSpec for the eval case.""" + # Make the metric function and tensor names. + if get_eval_metric_ops_fn is not None: + def metric_fn( + generator_inputs, generated_data, real_data, discriminator_real_outputs, + discriminator_gen_outputs, generator_loss, discriminator_loss): + """`metric_fn` used in TPUEstimator to calculate metrics.""" + eval_metric_ops = { + 'generator_loss': metrics_lib.mean(generator_loss), + 'discriminator_loss': metrics_lib.mean(discriminator_loss), + } + custom_eval_metric_ops = get_eval_metric_ops_fn( + generator_inputs, generated_data, real_data, + discriminator_real_outputs, discriminator_gen_outputs) + if not isinstance(custom_eval_metric_ops, dict): + raise TypeError('`get_eval_metric_ops_fn` must return a dict, ' + 'received: {}'.format(custom_eval_metric_ops)) + eval_metric_ops.update(custom_eval_metric_ops) + return eval_metric_ops + tensors = { + 'generator_loss': gan_loss_no_reduction.generator_loss, + 'discriminator_loss': gan_loss_no_reduction.discriminator_loss, + 'generator_inputs': gan_model.generator_inputs, + 'generated_data': gan_model.generated_data, + 'real_data': gan_model.real_data, + 'discriminator_real_outputs': gan_model.discriminator_real_outputs, + 'discriminator_gen_outputs': gan_model.discriminator_gen_outputs, + } + else: + def metric_fn(generator_loss, discriminator_loss): + return { + 'generator_loss': metrics_lib.mean(generator_loss), + 'discriminator_loss': metrics_lib.mean(discriminator_loss), + } + tensors = { + 'generator_loss': gan_loss_no_reduction.generator_loss, + 'discriminator_loss': gan_loss_no_reduction.discriminator_loss, + } + + scalar_loss = gan_loss.generator_loss + gan_loss.discriminator_loss + return tpu_estimator.TPUEstimatorSpec( + mode=model_fn_lib.ModeKeys.EVAL, + predictions=gan_model.generated_data, + loss=scalar_loss, + eval_metrics=(metric_fn, tensors)) + + +def _get_train_estimator_spec( + gan_model, gan_loss, generator_optimizer, discriminator_optimizer, + joint_train, gan_train_steps): + """Return a TPUEstimatorSpec for the train case.""" + scalar_loss = gan_loss.generator_loss + gan_loss.discriminator_loss + + # Get generator and discriminator update ops. We split them so that update + # ops aren't accidentally run multiple times. For now, throw an error if + # there are update ops that aren't associated with either the generator or + # the discriminator. Might modify the `kwargs` dictionary. + gen_update_ops, dis_update_ops = tfgan_train._get_update_ops( # pylint:disable=protected-access + {}, gan_model.generator_scope.name, gan_model.discriminator_scope.name) + + def gen_train_op(): + with ops.name_scope('generator_train'): + return training.create_train_op( + total_loss=gan_loss.generator_loss, + optimizer=generator_optimizer, + variables_to_train=gan_model.generator_variables, + update_ops=gen_update_ops) + def dis_train_op(): + with ops.name_scope('discriminator_train'): + return training.create_train_op( + total_loss=gan_loss.discriminator_loss, + optimizer=discriminator_optimizer, + variables_to_train=gan_model.discriminator_variables, + update_ops=dis_update_ops) + + # Either optimize the generator and discriminator sequentially or jointly. + tpu_train_op = _combine_train_ops(gen_train_op, dis_train_op, joint_train, + gan_train_steps) + + return tpu_estimator.TPUEstimatorSpec( + loss=scalar_loss, + mode=model_fn_lib.ModeKeys.TRAIN, + train_op=tpu_train_op) + + +# TODO(joelshor): Add support for multiple D / G steps. +def _combine_train_ops(gen_train_op, dis_train_op, joint_train, + gan_train_steps): + """Combine generator and discriminator train ops into a single op.""" + del gan_train_steps + if joint_train: + tpu_train_op = control_flow_ops.group(gen_train_op(), dis_train_op(), + name='joint_train') + else: + with ops.control_dependencies([dis_train_op()]): + tpu_train_op = gen_train_op() + + return tpu_train_op + + +def _maybe_make_cross_shard_optimizer(opt): + if callable(opt): + if not isinstance(opt(), tpu_optimizer.CrossShardOptimizer): + return lambda: tpu_optimizer.CrossShardOptimizer(opt()) + elif not isinstance(opt, tpu_optimizer.CrossShardOptimizer): + return tpu_optimizer.CrossShardOptimizer(opt) + return opt 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 new file mode 100644 index 0000000000000000000000000000000000000000..0d9e6489bdd1d89cc49bfedc2eed784999c31d2b --- /dev/null +++ b/tensorflow/contrib/gan/python/estimator/python/tpu_gan_estimator_test.py @@ -0,0 +1,319 @@ +# 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 TF-GAN's TPU Estimator.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import shutil +import tempfile + +from absl.testing import parameterized +import numpy as np +import six + +from tensorflow.contrib import layers +from tensorflow.contrib.gan.python import namedtuples as tfgan_tuples +from tensorflow.contrib.gan.python.estimator.python import tpu_gan_estimator_impl as estimator +from tensorflow.contrib.gan.python.losses.python import tuple_losses as losses +from tensorflow.contrib.tpu.python.tpu import tpu_config +from tensorflow.contrib.tpu.python.tpu import tpu_estimator +from tensorflow.contrib.tpu.python.tpu import tpu_optimizer +from tensorflow.python.data.ops import dataset_ops +from tensorflow.python.estimator import model_fn as model_fn_lib +from tensorflow.python.estimator.estimator import WarmStartSettings +from tensorflow.python.framework import ops +from tensorflow.python.framework import tensor_shape +from tensorflow.python.framework.errors_impl import NotFoundError +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import metrics as metrics_lib +from tensorflow.python.ops import variable_scope +from tensorflow.python.platform import flags +from tensorflow.python.platform import test +from tensorflow.python.summary.writer import writer_cache +from tensorflow.python.training import learning_rate_decay +from tensorflow.python.training import training +from tensorflow.python.training import training_util + +FLAGS = flags.FLAGS + +flags.DEFINE_bool('use_tpu', False, 'Whether to run test on TPU or not.') + + +def generator_fn(noise, mode): + del mode + return layers.fully_connected(noise, tensor_shape.dimension_value( + noise.shape[1])) + + +def discriminator_fn(data, unused_conditioning, mode): + del unused_conditioning, mode + return layers.fully_connected(data, 1) + + +def get_dummy_gan_model(): + # TODO(joelshor): Find a better way of creating a variable scope. + with variable_scope.variable_scope('generator') as gen_scope: + gen_var = variable_scope.get_variable('dummy_var', initializer=0.0) + with variable_scope.variable_scope('discriminator') as dis_scope: + dis_var = variable_scope.get_variable('dummy_var', initializer=0.0) + return tfgan_tuples.GANModel( + generator_inputs=None, + generated_data=array_ops.ones([3, 4]), + generator_variables=[gen_var], + generator_scope=gen_scope, + generator_fn=None, + real_data=array_ops.zeros([3, 4]), + discriminator_real_outputs=array_ops.ones([1, 2, 3]) * dis_var, + discriminator_gen_outputs=array_ops.ones([1, 2, 3]) * gen_var * dis_var, + discriminator_variables=[dis_var], + discriminator_scope=dis_scope, + discriminator_fn=None) + + +def get_metrics(generator_inputs, generated_data, real_data, + discriminator_real_outputs, discriminator_gen_outputs): + del generator_inputs, discriminator_real_outputs, discriminator_gen_outputs + return { + 'mse_custom_metric': metrics_lib.mean_squared_error( + real_data, generated_data) + } + + +class GetTPUEstimatorSpecTest(test.TestCase, parameterized.TestCase): + """Tests that the EstimatorSpec is constructed appropriately.""" + + @classmethod + def setUpClass(cls): + super(GetTPUEstimatorSpecTest, cls).setUpClass() + cls._generator_optimizer = tpu_optimizer.CrossShardOptimizer( + training.GradientDescentOptimizer(1.0)) + cls._discriminator_optimizer = tpu_optimizer.CrossShardOptimizer( + training.GradientDescentOptimizer(1.0)) + + @parameterized.named_parameters( + ('joint_train', model_fn_lib.ModeKeys.TRAIN, True), + ('train_sequential', model_fn_lib.ModeKeys.TRAIN, False), + ('eval', model_fn_lib.ModeKeys.EVAL, None), + ('predict', model_fn_lib.ModeKeys.PREDICT, None)) + def test_get_estimator_spec(self, mode, joint_train): + with ops.Graph().as_default(): + self._gan_model = get_dummy_gan_model() + spec = estimator._get_estimator_spec( + mode, + self._gan_model, + generator_loss_fn=losses.wasserstein_generator_loss, + discriminator_loss_fn=losses.wasserstein_discriminator_loss, + get_eval_metric_ops_fn=get_metrics, + generator_optimizer=self._generator_optimizer, + discriminator_optimizer=self._discriminator_optimizer, + joint_train=joint_train, + is_on_tpu=FLAGS.use_tpu, + gan_train_steps=tfgan_tuples.GANTrainSteps(1, 1)) + + self.assertIsInstance(spec, tpu_estimator.TPUEstimatorSpec) + self.assertEqual(mode, spec.mode) + if mode == model_fn_lib.ModeKeys.PREDICT: + self.assertEqual({'generated_data': self._gan_model.generated_data}, + spec.predictions) + elif mode == model_fn_lib.ModeKeys.TRAIN: + self.assertShapeEqual(np.array(0), spec.loss) # must be a scalar + self.assertIsNotNone(spec.train_op) + self.assertIsNotNone(spec.training_hooks) + elif mode == model_fn_lib.ModeKeys.EVAL: + self.assertEqual(self._gan_model.generated_data, spec.predictions) + self.assertShapeEqual(np.array(0), spec.loss) # must be a scalar + self.assertIsNotNone(spec.eval_metrics) + + +class TPUGANEstimatorIntegrationTest(test.TestCase, parameterized.TestCase): + + def setUp(self): + super(TPUGANEstimatorIntegrationTest, self).setUp() + self._model_dir = tempfile.mkdtemp() + self._config = tpu_config.RunConfig(model_dir=self._model_dir) + + def tearDown(self): + super(TPUGANEstimatorIntegrationTest, self).tearDown() + if self._model_dir: + writer_cache.FileWriterCache.clear() + shutil.rmtree(self._model_dir) + + def _test_complete_flow( + self, train_input_fn, eval_input_fn, predict_input_fn, prediction_size, + lr_decay=False, joint_train=True): + def make_opt(): + gstep = training_util.get_or_create_global_step() + lr = learning_rate_decay.exponential_decay(1.0, gstep, 10, 0.9) + return training.GradientDescentOptimizer(lr) + + gopt = make_opt if lr_decay else training.GradientDescentOptimizer(1.0) + dopt = make_opt if lr_decay else training.GradientDescentOptimizer(1.0) + est = estimator.TPUGANEstimator( + generator_fn=generator_fn, + discriminator_fn=discriminator_fn, + generator_loss_fn=losses.wasserstein_generator_loss, + discriminator_loss_fn=losses.wasserstein_discriminator_loss, + generator_optimizer=gopt, + discriminator_optimizer=dopt, + joint_train=joint_train, + get_eval_metric_ops_fn=get_metrics, + train_batch_size=4, + eval_batch_size=10, + predict_batch_size=8, + use_tpu=FLAGS.use_tpu, + config=self._config) + + # Train. + num_steps_train = 10 + est.train(train_input_fn, steps=num_steps_train) + + # Evaluate. + num_steps_eval = 2 + scores = est.evaluate(eval_input_fn, steps=num_steps_eval) + self.assertIn(ops.GraphKeys.GLOBAL_STEP, six.iterkeys(scores)) + self.assertIn('loss', six.iterkeys(scores)) + self.assertEqual(scores['discriminator_loss'] + scores['generator_loss'], + scores['loss']) + self.assertIn('mse_custom_metric', six.iterkeys(scores)) + + # Predict. + predictions = np.array([x['generated_data'] for x in + est.predict(predict_input_fn)]) + self.assertAllEqual(prediction_size, predictions.shape) + + @parameterized.named_parameters( + ('joint_train', True, False, False), + ('train_sequential', False, False, False), + ('lr_decay', False, True, False), + ('train_sequential_ds', False, False, True)) + def test_numpy_input_fn(self, joint_train, lr_decay, return_ds): + """Tests complete flow with numpy_input_fn.""" + input_dim = 4 + def train_input_fn(params): + data = np.zeros([input_dim], dtype=np.float32) + ds = (dataset_ops.Dataset + .from_tensors((data, data)) + .repeat() + .batch(params['batch_size'], drop_remainder=True)) + if return_ds: + return ds + else: + x, y = ds.make_one_shot_iterator().get_next() + return x, y + def eval_input_fn(params): + data = np.zeros([input_dim], dtype=np.float32) + ds = (dataset_ops.Dataset + .from_tensors((data, data)) + .repeat() + .batch(params['batch_size'], drop_remainder=True)) + if return_ds: + return ds + else: + x, y = ds.make_one_shot_iterator().get_next() + return x, y + predict_size = 10 + def predict_input_fn(params): + del params # unused + data = np.zeros([input_dim], dtype=np.float32) + ds = (dataset_ops.Dataset + .from_tensors(data) + .repeat(predict_size) + .batch(1, drop_remainder=True)) + return ds + + self._test_complete_flow( + train_input_fn=train_input_fn, + eval_input_fn=eval_input_fn, + predict_input_fn=predict_input_fn, + prediction_size=[predict_size, input_dim], + lr_decay=lr_decay, + joint_train=joint_train) + + +class TPUGANEstimatorWarmStartTest(test.TestCase): + + def setUp(self): + self._model_dir = self.get_temp_dir() + self._config = tpu_config.RunConfig(model_dir=self._model_dir) + self.new_variable_name = 'new_var' + self.new_variable_value = [1.0, 2.0, 3.0] + + def tearDown(self): + writer_cache.FileWriterCache.clear() + + def _test_warm_start(self, warm_start_from=None): + """Tests whether WarmStartSettings work as intended.""" + def generator_with_new_variable(noise_dict, mode): + variable_scope.get_variable(name=self.new_variable_name, + initializer=self.new_variable_value, + trainable=True) + return generator_fn(noise_dict, mode) + + est = estimator.TPUGANEstimator( + generator_fn=generator_fn, + discriminator_fn=discriminator_fn, + generator_loss_fn=losses.wasserstein_generator_loss, + discriminator_loss_fn=losses.wasserstein_discriminator_loss, + generator_optimizer=training.GradientDescentOptimizer(1.0), + discriminator_optimizer=training.GradientDescentOptimizer(1.0), + train_batch_size=4, + use_tpu=FLAGS.use_tpu, + config=self._config) + + def train_input_fn(params): + data = np.zeros([params['batch_size'], 4], dtype=np.float32) + return data, data + + est.train(train_input_fn, steps=1) + + est_warm = estimator.TPUGANEstimator( + generator_fn=generator_with_new_variable, + discriminator_fn=discriminator_fn, + generator_loss_fn=losses.wasserstein_generator_loss, + discriminator_loss_fn=losses.wasserstein_discriminator_loss, + generator_optimizer=training.GradientDescentOptimizer(1.0), + discriminator_optimizer=training.GradientDescentOptimizer(1.0), + config=tpu_config.RunConfig( + model_dir=None if warm_start_from else self._model_dir), + train_batch_size=4, + use_tpu=FLAGS.use_tpu, + warm_start_from=warm_start_from) + + est_warm.train(train_input_fn, steps=1) + + return est_warm + + def test_warm_start_error(self): + """Test if exception when reloading different estimators.""" + with self.assertRaises(NotFoundError): + self._test_warm_start() + + def test_warm_start_success(self): + """Test if GANEstimator allows explicit warm start variable assignment.""" + # Regex matches all variable names in ckpt except for new_var. + var_regex = '^(?!.*%s.*)' % self.new_variable_name + warmstart = WarmStartSettings(ckpt_to_initialize_from=self._model_dir, + vars_to_warm_start=var_regex) + est_warm = self._test_warm_start(warm_start_from=warmstart) + full_variable_name = 'Generator/%s' % self.new_variable_name + self.assertIn(full_variable_name, est_warm.get_variable_names()) + equal_vals = np.array_equal(est_warm.get_variable_value(full_variable_name), + self.new_variable_value) + self.assertTrue(equal_vals) + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/gan/python/eval/__init__.py b/tensorflow/contrib/gan/python/eval/__init__.py index f86b8513053a45f9830411f7df2c32d1f36a97b2..92e9abf8a35de1999eb800e169f32220fe47f8cd 100644 --- a/tensorflow/contrib/gan/python/eval/__init__.py +++ b/tensorflow/contrib/gan/python/eval/__init__.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""TFGAN evaluation module. +"""TF-GAN evaluation module. This module supports techniques such as Inception Score, Frechet Inception distance, and Sliced Wasserstein distance. diff --git a/tensorflow/contrib/gan/python/eval/python/classifier_metrics.py b/tensorflow/contrib/gan/python/eval/python/classifier_metrics.py index 1c872626a957279132772ae27df7a66a2564e9a5..a52e899114b62cb29752f72aa59f142f4a428aa1 100644 --- a/tensorflow/contrib/gan/python/eval/python/classifier_metrics.py +++ b/tensorflow/contrib/gan/python/eval/python/classifier_metrics.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Model evaluation tools for TFGAN.""" +"""Model evaluation tools for TF-GAN.""" from __future__ import absolute_import from __future__ import division 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 a71ee53311c1c057a5b41be0331bf56ce1a82f74..31f0d34ed68a6adc25cca102236079d0f66615cb 100644 --- a/tensorflow/contrib/gan/python/eval/python/classifier_metrics_impl.py +++ b/tensorflow/contrib/gan/python/eval/python/classifier_metrics_impl.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Model evaluation tools for TFGAN. +"""Model evaluation tools for TF-GAN. These methods come from https://arxiv.org/abs/1606.03498, https://arxiv.org/abs/1706.08500, and https://arxiv.org/abs/1801.01401. @@ -387,7 +387,7 @@ def classifier_score_from_logits(logits): # Use maximum precision for best results. logits_dtype = logits.dtype if logits_dtype != dtypes.float64: - logits = math_ops.to_double(logits) + logits = math_ops.cast(logits, dtypes.float64) p = nn_ops.softmax(logits) q = math_ops.reduce_mean(p, axis=0) @@ -562,8 +562,8 @@ def mean_only_frechet_classifier_distance_from_activations( activations_dtype = real_activations.dtype if activations_dtype != dtypes.float64: - real_activations = math_ops.to_double(real_activations) - generated_activations = math_ops.to_double(generated_activations) + real_activations = math_ops.cast(real_activations, dtypes.float64) + generated_activations = math_ops.cast(generated_activations, dtypes.float64) # Compute means of activations. m = math_ops.reduce_mean(real_activations, 0) @@ -623,8 +623,8 @@ def diagonal_only_frechet_classifier_distance_from_activations( activations_dtype = real_activations.dtype if activations_dtype != dtypes.float64: - real_activations = math_ops.to_double(real_activations) - generated_activations = math_ops.to_double(generated_activations) + real_activations = math_ops.cast(real_activations, dtypes.float64) + generated_activations = math_ops.cast(generated_activations, dtypes.float64) # Compute mean and covariance matrices of activations. m, var = nn_impl.moments(real_activations, axes=[0]) @@ -698,15 +698,16 @@ def frechet_classifier_distance_from_activations(real_activations, activations_dtype = real_activations.dtype if activations_dtype != dtypes.float64: - real_activations = math_ops.to_double(real_activations) - generated_activations = math_ops.to_double(generated_activations) + real_activations = math_ops.cast(real_activations, dtypes.float64) + generated_activations = math_ops.cast(generated_activations, dtypes.float64) # Compute mean and covariance matrices of activations. m = math_ops.reduce_mean(real_activations, 0) m_w = math_ops.reduce_mean(generated_activations, 0) - num_examples_real = math_ops.to_double(array_ops.shape(real_activations)[0]) - num_examples_generated = math_ops.to_double( - array_ops.shape(generated_activations)[0]) + num_examples_real = math_ops.cast( + array_ops.shape(real_activations)[0], dtypes.float64) + num_examples_generated = math_ops.cast( + array_ops.shape(generated_activations)[0], dtypes.float64) # sigma = (1 / (n - 1)) * (X - mu) (X - mu)^T real_centered = real_activations - m @@ -794,9 +795,9 @@ def kernel_classifier_distance(real_images, on a classifier. num_classifier_batches: Number of batches to split images in to in order to efficiently run them through the classifier network. - max_estimator_block_size: integer, default 1024. The distance estimator - splits samples into blocks for computational efficiency. Larger values are - more computationally expensive but decrease the variance of the distance + max_block_size: integer, default 1024. The distance estimator splits samples + into blocks for computational efficiency. Larger values are more + computationally expensive but decrease the variance of the distance estimate. dtype: if not None, coerce activations to this dtype before computations. @@ -871,9 +872,9 @@ def kernel_classifier_distance_and_std(real_images, on a classifier. num_classifier_batches: Number of batches to split images in to in order to efficiently run them through the classifier network. - max_estimator_block_size: integer, default 1024. The distance estimator - splits samples into blocks for computational efficiency. Larger values are - more computationally expensive but decrease the variance of the distance + max_block_size: integer, default 1024. The distance estimator splits samples + into blocks for computational efficiency. Larger values are more + computationally expensive but decrease the variance of the distance estimate. Having a smaller block size also gives a better estimate of the standard error. dtype: if not None, coerce activations to this dtype before computations. @@ -910,7 +911,7 @@ def kernel_classifier_distance_and_std(real_images, gen_a = array_ops.concat(array_ops.unstack(gen_a), 0) return kernel_classifier_distance_and_std_from_activations( - real_a, gen_a, max_block_size=max_block_size) + real_a, gen_a, max_block_size, dtype) kernel_inception_distance_and_std = functools.partial( @@ -967,14 +968,14 @@ def kernel_classifier_distance_from_activations(real_activations, into blocks for computational efficiency. Larger values are more computationally expensive but decrease the variance of the distance estimate. - dtype: if not None, coerce activations to this dtype before computations. + dtype: If not None, coerce activations to this dtype before computations. Returns: The Kernel Inception Distance. A floating-point scalar of the same type as the output of the activations. """ return kernel_classifier_distance_and_std_from_activations( - real_activations, generated_activations, max_block_size=max_block_size)[0] + real_activations, generated_activations, max_block_size, dtype)[0] def kernel_classifier_distance_and_std_from_activations(real_activations, @@ -1029,7 +1030,7 @@ def kernel_classifier_distance_and_std_from_activations(real_activations, computationally expensive but decrease the variance of the distance estimate. Having a smaller block size also gives a better estimate of the standard error. - dtype: if not None, coerce activations to this dtype before computations. + dtype: If not None, coerce activations to this dtype before computations. Returns: The Kernel Inception Distance. A floating-point scalar of the same type @@ -1080,7 +1081,7 @@ def kernel_classifier_distance_and_std_from_activations(real_activations, dim = math_ops.cast(real_activations.shape[1], dtype) def compute_kid_block(i): - 'Compute the ith block of the KID estimate.' + """Computes the ith block of the KID estimate.""" r_s = inds_r[i] r_e = inds_r[i + 1] r = real_activations[r_s:r_e] 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 dbff1d2a367e10adc607dafb4c571bb3607a3963..bd17571a0535a3c8e9dfee24a8da16eb2e72f165 100644 --- a/tensorflow/contrib/gan/python/eval/python/classifier_metrics_test.py +++ b/tensorflow/contrib/gan/python/eval/python/classifier_metrics_test.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Tests for TFGAN classifier_metrics.""" +"""Tests for TF-GAN classifier_metrics.""" from __future__ import absolute_import from __future__ import division @@ -234,7 +234,7 @@ class ClassifierMetricsTest(test.TestCase, parameterized.TestCase): else: logits = classifier_metrics.run_inception(img, _get_dummy_graphdef()) - self.assertTrue(isinstance(logits, ops.Tensor)) + self.assertIsInstance(logits, ops.Tensor) logits.shape.assert_is_compatible_with([batch_size, 1001]) # Check that none of the model variables are trainable. @@ -258,7 +258,7 @@ class ClassifierMetricsTest(test.TestCase, parameterized.TestCase): img, _get_dummy_graphdef(), output_tensor=classifier_metrics.INCEPTION_FINAL_POOL) - self.assertTrue(isinstance(pool, ops.Tensor)) + self.assertIsInstance(pool, ops.Tensor) pool.shape.assert_is_compatible_with([batch_size, 2048]) # Check that none of the model variables are trainable. @@ -276,8 +276,8 @@ class ClassifierMetricsTest(test.TestCase, parameterized.TestCase): classifier_metrics.INCEPTION_FINAL_POOL ]) - self.assertTrue(isinstance(logits, ops.Tensor)) - self.assertTrue(isinstance(pool, ops.Tensor)) + self.assertIsInstance(logits, ops.Tensor) + self.assertIsInstance(pool, ops.Tensor) logits.shape.assert_is_compatible_with([batch_size, 1001]) pool.shape.assert_is_compatible_with([batch_size, 2048]) @@ -290,7 +290,7 @@ class ClassifierMetricsTest(test.TestCase, parameterized.TestCase): classifier_metrics.inception_score, array_ops.zeros([6, 299, 299, 3]), num_batches=3) - self.assertTrue(isinstance(score, ops.Tensor)) + self.assertIsInstance(score, ops.Tensor) score.shape.assert_has_rank(0) # Check that none of the model variables are trainable. @@ -302,7 +302,7 @@ class ClassifierMetricsTest(test.TestCase, parameterized.TestCase): distance = _run_with_mock( classifier_metrics.frechet_inception_distance, img, img) - self.assertTrue(isinstance(distance, ops.Tensor)) + self.assertIsInstance(distance, ops.Tensor) distance.shape.assert_has_rank(0) # Check that none of the model variables are trainable. @@ -314,7 +314,7 @@ class ClassifierMetricsTest(test.TestCase, parameterized.TestCase): distance = _run_with_mock(classifier_metrics.kernel_inception_distance, img, img) - self.assertTrue(isinstance(distance, ops.Tensor)) + self.assertIsInstance(distance, ops.Tensor) distance.shape.assert_has_rank(0) # Check that none of the model variables are trainable. diff --git a/tensorflow/contrib/gan/python/eval/python/sliced_wasserstein.py b/tensorflow/contrib/gan/python/eval/python/sliced_wasserstein.py index 523968bed91f1021ae629bf52c405cf5c2d7b917..326fcb3cdbf2eda66207f134cd2926f09a216a99 100644 --- a/tensorflow/contrib/gan/python/eval/python/sliced_wasserstein.py +++ b/tensorflow/contrib/gan/python/eval/python/sliced_wasserstein.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Model evaluation tools for TFGAN.""" +"""Model evaluation tools for TF-GAN.""" from __future__ import absolute_import from __future__ import division diff --git a/tensorflow/contrib/gan/python/eval/python/summaries.py b/tensorflow/contrib/gan/python/eval/python/summaries.py index ecfdb39499b1e824e02415c0db1de3157e4f3216..1b202dfc97304ddc7ced42d65366aaf419439392 100644 --- a/tensorflow/contrib/gan/python/eval/python/summaries.py +++ b/tensorflow/contrib/gan/python/eval/python/summaries.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Common TFGAN summaries.""" +"""Common TF-GAN summaries.""" from __future__ import absolute_import from __future__ import division diff --git a/tensorflow/contrib/gan/python/eval/python/summaries_impl.py b/tensorflow/contrib/gan/python/eval/python/summaries_impl.py index f9995bb19d0d09eaf6fd96d039b0bba1d3a7055c..9f448d3a1602c503093214201bdc96fc9bee85b5 100644 --- a/tensorflow/contrib/gan/python/eval/python/summaries_impl.py +++ b/tensorflow/contrib/gan/python/eval/python/summaries_impl.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Common TFGAN summaries.""" +"""Common TF-GAN summaries.""" from __future__ import absolute_import from __future__ import division diff --git a/tensorflow/contrib/gan/python/eval/python/summaries_test.py b/tensorflow/contrib/gan/python/eval/python/summaries_test.py index 54a6f8d4d9086ad7fc8db31032677628561e48e8..53fc7cb8ede698c2d8590c7fd3016a884cef9be9 100644 --- a/tensorflow/contrib/gan/python/eval/python/summaries_test.py +++ b/tensorflow/contrib/gan/python/eval/python/summaries_test.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Tests for TFGAN summaries.""" +"""Tests for TF-GAN summaries.""" from __future__ import absolute_import from __future__ import division 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/conditioning_utils_impl.py b/tensorflow/contrib/gan/python/features/python/conditioning_utils_impl.py index e2594faf85bcf91cbe09f266e4d4211d20bdee17..364fa4eb461c62784803f0c309e3b7c5855df199 100644 --- a/tensorflow/contrib/gan/python/features/python/conditioning_utils_impl.py +++ b/tensorflow/contrib/gan/python/features/python/conditioning_utils_impl.py @@ -64,6 +64,9 @@ def condition_tensor(tensor, conditioning): """ tensor.shape[1:].assert_is_fully_defined() num_features = tensor.shape[1:].num_elements() + if conditioning.shape.ndims < 2: + raise ValueError('conditioning must be at least 2D, but saw shape: %s' + % conditioning.shape) mapped_conditioning = layers.linear( layers.flatten(conditioning), num_features) diff --git a/tensorflow/contrib/gan/python/features/python/conditioning_utils_test.py b/tensorflow/contrib/gan/python/features/python/conditioning_utils_test.py index 0aad769793761be69ee9d1e3416e44c7b3d8cea0..f5c7d53cf2c9aa08ba0074950983ef3ecd90168b 100644 --- a/tensorflow/contrib/gan/python/features/python/conditioning_utils_test.py +++ b/tensorflow/contrib/gan/python/features/python/conditioning_utils_test.py @@ -45,7 +45,7 @@ class ConditioningUtilsTest(test.TestCase): array_ops.placeholder(dtypes.float32, (5, None)), array_ops.placeholder(dtypes.float32, (5, 1))) - with self.assertRaisesRegexp(ValueError, 'expected min_ndim=2'): + with self.assertRaisesRegexp(ValueError, 'at least 2D'): conditioning_utils.condition_tensor( array_ops.placeholder(dtypes.float32, (5, 2)), array_ops.placeholder(dtypes.float32, (5))) 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/features/python/virtual_batchnorm_impl.py b/tensorflow/contrib/gan/python/features/python/virtual_batchnorm_impl.py index 650eab97a3952e9aec2b489fffcc83c3bc49f2dd..f5c448db41c67adb4edd2634dd63a1840180df70 100644 --- a/tensorflow/contrib/gan/python/features/python/virtual_batchnorm_impl.py +++ b/tensorflow/contrib/gan/python/features/python/virtual_batchnorm_impl.py @@ -200,7 +200,7 @@ class VBN(object): del reduction_axes[axis] self._broadcast_shape = [1] * len(input_shape) - self._broadcast_shape[axis] = input_shape[axis].value + self._broadcast_shape[axis] = input_shape.dims[axis] self._example_reduction_axes = list(range(ndims)) del self._example_reduction_axes[max(axis, self._batch_axis)] diff --git a/tensorflow/contrib/gan/python/losses/python/losses_impl.py b/tensorflow/contrib/gan/python/losses/python/losses_impl.py index 8bc4db8424f661bba65675f0cd1c2fc33696eda9..1f1ae2df4d6def618e86aced3296ac89c836eab7 100644 --- a/tensorflow/contrib/gan/python/losses/python/losses_impl.py +++ b/tensorflow/contrib/gan/python/losses/python/losses_impl.py @@ -28,7 +28,7 @@ wasserstein_gradient_penalty All losses must be able to accept 1D or 2D Tensors, so as to be compatible with patchGAN style losses (https://arxiv.org/abs/1611.07004). -To make these losses usable in the TFGAN framework, please create a tuple +To make these losses usable in the TF-GAN framework, please create a tuple version of the losses with `losses_utils.py`. """ @@ -36,9 +36,9 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import numpy as np from tensorflow.contrib.framework.python.ops import variables as contrib_variables_lib +from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.framework import tensor_util from tensorflow.python.ops import array_ops @@ -47,7 +47,6 @@ from tensorflow.python.ops import gradients_impl from tensorflow.python.ops import math_ops from tensorflow.python.ops import random_ops from tensorflow.python.ops import variable_scope -from tensorflow.python.ops.distributions import distribution as ds from tensorflow.python.ops.losses import losses from tensorflow.python.ops.losses import util from tensorflow.python.summary import summary @@ -71,6 +70,10 @@ __all__ = [ ] +def _to_float(tensor): + return math_ops.cast(tensor, dtypes.float32) + + # Wasserstein losses from `Wasserstein GAN` (https://arxiv.org/abs/1701.07875). def wasserstein_generator_loss( discriminator_gen_outputs, @@ -100,7 +103,7 @@ def wasserstein_generator_loss( """ with ops.name_scope(scope, 'generator_wasserstein_loss', ( discriminator_gen_outputs, weights)) as scope: - discriminator_gen_outputs = math_ops.to_float(discriminator_gen_outputs) + discriminator_gen_outputs = _to_float(discriminator_gen_outputs) loss = - discriminator_gen_outputs loss = losses.compute_weighted_loss( @@ -146,8 +149,8 @@ def wasserstein_discriminator_loss( with ops.name_scope(scope, 'discriminator_wasserstein_loss', ( discriminator_real_outputs, discriminator_gen_outputs, real_weights, generated_weights)) as scope: - discriminator_real_outputs = math_ops.to_float(discriminator_real_outputs) - discriminator_gen_outputs = math_ops.to_float(discriminator_gen_outputs) + discriminator_real_outputs = _to_float(discriminator_real_outputs) + discriminator_gen_outputs = _to_float(discriminator_gen_outputs) discriminator_real_outputs.shape.assert_is_compatible_with( discriminator_gen_outputs.shape) @@ -322,7 +325,7 @@ def wasserstein_gradient_penalty( generated_data: Output of the generator. generator_inputs: Exact argument to pass to the generator, which is used as optional conditioning to the discriminator. - discriminator_fn: A discriminator function that conforms to TFGAN API. + discriminator_fn: A discriminator function that conforms to TF-GAN API. discriminator_scope: If not `None`, reuse discriminators from this scope. epsilon: A small positive number added for numerical stability when computing the gradient norm. @@ -355,7 +358,8 @@ def wasserstein_gradient_penalty( raise ValueError('`generated_data` can\'t have unknown rank.') differences = generated_data - real_data - batch_size = differences.shape[0].value or array_ops.shape(differences)[0] + batch_size = differences.shape.dims[0].value or array_ops.shape( + differences)[0] alpha_shape = [batch_size] + [1] * (differences.shape.ndims - 1) alpha = random_ops.random_uniform(shape=alpha_shape) interpolates = real_data + (alpha * differences) @@ -648,7 +652,7 @@ def least_squares_generator_loss( """ with ops.name_scope(scope, 'lsq_generator_loss', (discriminator_gen_outputs, real_label)) as scope: - discriminator_gen_outputs = math_ops.to_float(discriminator_gen_outputs) + discriminator_gen_outputs = _to_float(discriminator_gen_outputs) loss = math_ops.squared_difference( discriminator_gen_outputs, real_label) / 2.0 loss = losses.compute_weighted_loss( @@ -703,8 +707,8 @@ def least_squares_discriminator_loss( """ with ops.name_scope(scope, 'lsq_discriminator_loss', (discriminator_gen_outputs, real_label)) as scope: - discriminator_real_outputs = math_ops.to_float(discriminator_real_outputs) - discriminator_gen_outputs = math_ops.to_float(discriminator_gen_outputs) + discriminator_real_outputs = _to_float(discriminator_real_outputs) + discriminator_gen_outputs = _to_float(discriminator_gen_outputs) discriminator_real_outputs.shape.assert_is_compatible_with( discriminator_gen_outputs.shape) @@ -739,11 +743,16 @@ def least_squares_discriminator_loss( def _validate_distributions(distributions): if not isinstance(distributions, (list, tuple)): raise ValueError('`distributions` must be a list or tuple. Instead, ' - 'found %s.', type(distributions)) + 'found %s.' % type(distributions)) for x in distributions: - if not isinstance(x, ds.Distribution): + # We used to check with `isinstance(x, tf.distributions.Distribution)`. + # However, distributions have migrated to `tfp.distributions.Distribution`, + # which is a new code repo, so we can't check this way anymore until + # TF-GAN is migrated to a new repo as well. + # This new check is not sufficient, but is a useful heuristic for now. + if not callable(getattr(x, 'log_prob', None)): raise ValueError('`distributions` must be a list of `Distributions`. ' - 'Instead, found %s.', type(x)) + 'Instead, found %s.' % type(x)) def _validate_information_penalty_inputs( @@ -816,7 +825,7 @@ def _numerically_stable_global_norm(tensor_list): Returns: A scalar tensor with the global norm. """ - if np.all([x is None for x in tensor_list]): + if all(x is None for x in tensor_list): return 0.0 list_max = math_ops.reduce_max([math_ops.reduce_max(math_ops.abs(x)) for x in diff --git a/tensorflow/contrib/gan/python/losses/python/tuple_losses_impl.py b/tensorflow/contrib/gan/python/losses/python/tuple_losses_impl.py index 221c70c38bd432a6be7f6cda9c6700aa2255821f..76e57df7f646547037b3461ac44f7ee5b971406c 100644 --- a/tensorflow/contrib/gan/python/losses/python/tuple_losses_impl.py +++ b/tensorflow/contrib/gan/python/losses/python/tuple_losses_impl.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""TFGAN utilities for loss functions that accept GANModel namedtuples. +"""TF-GAN utilities for loss functions that accept GANModel namedtuples. The losses and penalties in this file all correspond to losses in `losses_impl.py`. Losses in that file take individual arguments, whereas in this diff --git a/tensorflow/contrib/gan/python/namedtuples.py b/tensorflow/contrib/gan/python/namedtuples.py index b9ac1bf15138c7e7d15ab3ebdac605d84921b6e5..73dfee4fdeec87cf0bac5eb675fd02a64a9ad7f5 100644 --- a/tensorflow/contrib/gan/python/namedtuples.py +++ b/tensorflow/contrib/gan/python/namedtuples.py @@ -12,10 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Named tuples for TFGAN. +"""Named tuples for TF-GAN. -TFGAN training occurs in four steps, and each step communicates with the next -step via one of these named tuples. At each step, you can either use a TFGAN +TF-GAN training occurs in four steps, and each step communicates with the next +step via one of these named tuples. At each step, you can either use a TF-GAN helper function in `train.py`, or you can manually construct a tuple. """ @@ -213,7 +213,8 @@ class GANTrainOps( collections.namedtuple('GANTrainOps', ( 'generator_train_op', 'discriminator_train_op', - 'global_step_inc_op' + 'global_step_inc_op', + 'train_hooks' ))): """GANTrainOps contains the training ops. @@ -221,8 +222,17 @@ class GANTrainOps( generator_train_op: Op that performs a generator update step. discriminator_train_op: Op that performs a discriminator update step. global_step_inc_op: Op that increments the shared global step. + train_hooks: a list or tuple containing hooks related to training that need + to be populated when training ops are instantiated. Used primarily for + sync hooks. """ + def __new__(cls, generator_train_op, discriminator_train_op, + global_step_inc_op, train_hooks=()): + return super(GANTrainOps, cls).__new__(cls, generator_train_op, + discriminator_train_op, + global_step_inc_op, train_hooks) + class GANTrainSteps( collections.namedtuple('GANTrainSteps', ( diff --git a/tensorflow/contrib/gan/python/train.py b/tensorflow/contrib/gan/python/train.py index 9e5aea1498a7e9d47480af18cad9f80ede84c0f9..f36a5d346e0f27fbbc480e876380db51ed559c09 100644 --- a/tensorflow/contrib/gan/python/train.py +++ b/tensorflow/contrib/gan/python/train.py @@ -12,12 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""The TFGAN project provides a lightweight GAN training/testing framework. +"""The TF-GAN project provides a lightweight GAN training/testing framework. This file contains the core helper functions to create and train a GAN model. See the README or examples in `tensorflow_models` for details on how to use. -TFGAN training occurs in four steps: +TF-GAN training occurs in four steps: 1) Create a model 2) Add a loss 3) Create train ops @@ -45,7 +45,6 @@ from tensorflow.python.ops import init_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import random_ops from tensorflow.python.ops import variable_scope -from tensorflow.python.ops.distributions import distribution as ds from tensorflow.python.ops.losses import losses from tensorflow.python.summary import summary from tensorflow.python.training import session_run_hook @@ -115,7 +114,7 @@ def gan_model( discriminator_gen_outputs = discriminator_fn(generated_data, generator_inputs) with variable_scope.variable_scope(dis_scope, reuse=True): - real_data = ops.convert_to_tensor(real_data) + real_data = _convert_tensor_or_l_or_d(real_data) discriminator_real_outputs = discriminator_fn(real_data, generator_inputs) if check_shapes: @@ -646,9 +645,10 @@ def gan_loss( type(model)) # Optionally create pooled model. - pooled_model = ( - _tensor_pool_adjusted_model(model, tensor_pool_fn) - if tensor_pool_fn else model) + if tensor_pool_fn: + pooled_model = _tensor_pool_adjusted_model(model, tensor_pool_fn) + else: + pooled_model = model # Create standard losses. gen_loss = generator_loss_fn(model, add_summaries=add_summaries) @@ -666,10 +666,11 @@ def gan_loss( if _use_aux_loss(mutual_information_penalty_weight): gen_info_loss = tfgan_losses.mutual_information_penalty( model, add_summaries=add_summaries) - dis_info_loss = ( - gen_info_loss - if tensor_pool_fn is None else tfgan_losses.mutual_information_penalty( - pooled_model, add_summaries=add_summaries)) + if tensor_pool_fn is None: + dis_info_loss = gen_info_loss + else: + dis_info_loss = tfgan_losses.mutual_information_penalty( + pooled_model, add_summaries=add_summaries) gen_loss += mutual_information_penalty_weight * gen_info_loss dis_loss += mutual_information_penalty_weight * dis_info_loss if _use_aux_loss(aux_cond_generator_weight): @@ -925,11 +926,12 @@ def gan_train_ops( generator_optimizer, discriminator_optimizer, check_for_unused_update_ops=True, + is_chief=True, # Optional args to pass directly to the `create_train_op`. **kwargs): """Returns GAN train ops. - The highest-level call in TFGAN. It is composed of functions that can also + The highest-level call in TF-GAN. It is composed of functions that can also be called, should a user require more control over some part of the GAN training process. @@ -940,6 +942,8 @@ def gan_train_ops( discriminator_optimizer: The optimizer for the discriminator updates. check_for_unused_update_ops: If `True`, throws an exception if there are update ops outside of the generator or discriminator scopes. + is_chief: Specifies whether or not the training is being run by the primary + replica during replica training. **kwargs: Keyword args to pass directly to `training.create_train_op` for both the generator and discriminator train op. @@ -981,6 +985,9 @@ def gan_train_ops( kwargs, model.generator_scope.name, model.discriminator_scope.name, check_for_unused_update_ops) + # Get the sync hooks if these are needed. + sync_hooks = [] + generator_global_step = None if isinstance(generator_optimizer, sync_replicas_optimizer.SyncReplicasOptimizer): @@ -996,6 +1003,7 @@ def gan_train_ops( trainable=False, collections=[ops.GraphKeys.GLOBAL_VARIABLES]) gen_update_ops += [generator_global_step.assign(global_step)] + sync_hooks.append(generator_optimizer.make_session_run_hook(is_chief)) with ops.name_scope('generator_train'): gen_train_op = training.create_train_op( total_loss=loss.generator_loss, @@ -1017,6 +1025,7 @@ def gan_train_ops( trainable=False, collections=[ops.GraphKeys.GLOBAL_VARIABLES]) dis_update_ops += [discriminator_global_step.assign(global_step)] + sync_hooks.append(discriminator_optimizer.make_session_run_hook(is_chief)) with ops.name_scope('discriminator_train'): disc_train_op = training.create_train_op( total_loss=loss.discriminator_loss, @@ -1026,7 +1035,8 @@ def gan_train_ops( update_ops=dis_update_ops, **kwargs) - return namedtuples.GANTrainOps(gen_train_op, disc_train_op, global_step_inc) + return namedtuples.GANTrainOps(gen_train_op, disc_train_op, global_step_inc, + sync_hooks) # TODO(joelshor): Implement a dynamic GAN train loop, as in `Real-Time Adaptive @@ -1067,13 +1077,24 @@ def get_sequential_train_hooks(train_steps=namedtuples.GANTrainSteps(1, 1)): train_steps.generator_train_steps) discriminator_hook = RunTrainOpsHook(train_ops.discriminator_train_op, train_steps.discriminator_train_steps) - return [generator_hook, discriminator_hook] + return [generator_hook, discriminator_hook] + list(train_ops.train_hooks) return get_hooks +def _num_joint_steps(train_steps): + g_steps = train_steps.generator_train_steps + d_steps = train_steps.discriminator_train_steps + # Get the number of each type of step that should be run. + num_d_and_g_steps = min(g_steps, d_steps) + num_g_steps = g_steps - num_d_and_g_steps + num_d_steps = d_steps - num_d_and_g_steps + + return num_d_and_g_steps, num_g_steps, num_d_steps + + def get_joint_train_hooks(train_steps=namedtuples.GANTrainSteps(1, 1)): - """Returns a hooks function for sequential GAN training. + """Returns a hooks function for joint GAN training. When using these train hooks, IT IS RECOMMENDED TO USE `use_locking=True` ON ALL OPTIMIZERS TO AVOID RACE CONDITIONS. @@ -1106,12 +1127,7 @@ def get_joint_train_hooks(train_steps=namedtuples.GANTrainSteps(1, 1)): Returns: A function that takes a GANTrainOps tuple and returns a list of hooks. """ - g_steps = train_steps.generator_train_steps - d_steps = train_steps.discriminator_train_steps - # Get the number of each type of step that should be run. - num_d_and_g_steps = min(g_steps, d_steps) - num_g_steps = g_steps - num_d_and_g_steps - num_d_steps = d_steps - num_d_and_g_steps + num_d_and_g_steps, num_g_steps, num_d_steps = _num_joint_steps(train_steps) def get_hooks(train_ops): g_op = train_ops.generator_train_op @@ -1121,7 +1137,7 @@ def get_joint_train_hooks(train_steps=namedtuples.GANTrainSteps(1, 1)): g_hook = RunTrainOpsHook(g_op, num_g_steps) d_hook = RunTrainOpsHook(d_op, num_d_steps) - return [joint_hook, g_hook, d_hook] + return [joint_hook, g_hook, d_hook] + list(train_ops.train_hooks) return get_hooks @@ -1264,10 +1280,6 @@ def _validate_distributions(distributions_l, noise_l): if not isinstance(distributions_l, (tuple, list)): raise ValueError('`predicted_distributions` must be a list. Instead, found ' '%s.' % type(distributions_l)) - for dist in distributions_l: - if not isinstance(dist, ds.Distribution): - raise ValueError('Every element in `predicted_distributions` must be a ' - '`tf.Distribution`. Instead, found %s.' % type(dist)) if len(distributions_l) != len(noise_l): raise ValueError('Length of `predicted_distributions` %i must be the same ' 'as the length of structured noise %i.' % diff --git a/tensorflow/contrib/gan/python/train_test.py b/tensorflow/contrib/gan/python/train_test.py index 64d670619905a427a84bee4b661228abca591fae..841f25cd7f1852767776eed2dcbf2522d8b0743b 100644 --- a/tensorflow/contrib/gan/python/train_test.py +++ b/tensorflow/contrib/gan/python/train_test.py @@ -519,7 +519,7 @@ class GANLossTest(test.TestCase, parameterized.TestCase): """Test output type.""" loss = train.gan_loss(get_gan_model_fn(), add_summaries=True) self.assertIsInstance(loss, namedtuples.GANLoss) - self.assertGreater(len(ops.get_collection(ops.GraphKeys.SUMMARIES)), 0) + self.assertNotEmpty(ops.get_collection(ops.GraphKeys.SUMMARIES)) @parameterized.named_parameters( ('cyclegan', create_cyclegan_model), @@ -528,7 +528,7 @@ class GANLossTest(test.TestCase, parameterized.TestCase): def test_cyclegan_output_type(self, get_gan_model_fn): loss = train.cyclegan_loss(get_gan_model_fn(), add_summaries=True) self.assertIsInstance(loss, namedtuples.CycleGANLoss) - self.assertGreater(len(ops.get_collection(ops.GraphKeys.SUMMARIES)), 0) + self.assertNotEmpty(ops.get_collection(ops.GraphKeys.SUMMARIES)) @parameterized.named_parameters( ('gan', create_gan_model, False), @@ -759,7 +759,7 @@ class TensorPoolAdjusteModelTest(test.TestCase): # For [pool_size, ?), the pool is full, tensor2 must be equal to some # historical values of tensor1 (which is previously stored in the # pool). - self.assertTrue(any([(v == t2).all() for v in history_values])) + self.assertTrue(any((v == t2).all() for v in history_values)) def _make_new_model_and_check(self, model, pool_size): pool_fn = lambda x: random_tensor_pool.tensor_pool(x, pool_size=pool_size) @@ -836,6 +836,9 @@ class GANTrainOpsTest(test.TestCase, parameterized.TestCase): self.assertIsInstance(train_ops, namedtuples.GANTrainOps) + # Make sure there are no training hooks populated accidentally. + self.assertEmpty(train_ops.train_hooks) + # TODO(joelshor): Add a test to check that custom update op is run. @parameterized.named_parameters( ('gan', create_gan_model, False), @@ -923,8 +926,15 @@ class GANTrainOpsTest(test.TestCase, parameterized.TestCase): model, loss, generator_optimizer=g_opt, discriminator_optimizer=d_opt) self.assertIsInstance(train_ops, namedtuples.GANTrainOps) # No new trainable variables should have been added. - self.assertEqual(num_trainable_vars, - len(variables_lib.get_trainable_variables())) + self.assertLen(variables_lib.get_trainable_variables(), num_trainable_vars) + + # Sync hooks should be populated in the GANTrainOps. + self.assertLen(train_ops.train_hooks, 2) + for hook in train_ops.train_hooks: + self.assertIsInstance( + hook, sync_replicas_optimizer._SyncReplicasOptimizerHook) + sync_opts = [hook._sync_optimizer for hook in train_ops.train_hooks] + self.assertSetEqual(frozenset(sync_opts), frozenset((g_opt, d_opt))) g_sync_init_op = g_opt.get_init_tokens_op(num_tokens=1) d_sync_init_op = d_opt.get_init_tokens_op(num_tokens=1) @@ -959,6 +969,32 @@ class GANTrainOpsTest(test.TestCase, parameterized.TestCase): coord.request_stop() coord.join(g_threads + d_threads) + @parameterized.named_parameters( + ('is_chief', True), + ('is_not_chief', False), + ) + def test_is_chief_in_train_hooks(self, is_chief): + """Make sure is_chief is propagated correctly to sync hooks.""" + model = create_gan_model() + loss = train.gan_loss(model) + g_opt = get_sync_optimizer() + d_opt = get_sync_optimizer() + train_ops = train.gan_train_ops( + model, + loss, + g_opt, + d_opt, + is_chief=is_chief, + summarize_gradients=True, + colocate_gradients_with_ops=True) + + self.assertLen(train_ops.train_hooks, 2) + for hook in train_ops.train_hooks: + self.assertIsInstance( + hook, sync_replicas_optimizer._SyncReplicasOptimizerHook) + is_chief_list = [hook._is_chief for hook in train_ops.train_hooks] + self.assertListEqual(is_chief_list, [is_chief, is_chief]) + class GANTrainTest(test.TestCase, parameterized.TestCase): """Tests for `gan_train`.""" @@ -1036,6 +1072,44 @@ class GANTrainTest(test.TestCase, parameterized.TestCase): self.assertTrue(np.isscalar(final_loss)) self.assertEqual(17.0, final_loss) + @parameterized.named_parameters( + ('gan', create_gan_model), + ('callable_gan', create_callable_gan_model), + ('infogan', create_infogan_model), + ('callable_infogan', create_callable_infogan_model), + ('acgan', create_acgan_model), + ('callable_acgan', create_callable_acgan_model), + ) + def test_train_hooks_exist_in_get_hooks_fn(self, create_gan_model_fn): + model = create_gan_model_fn() + loss = train.gan_loss(model) + + g_opt = get_sync_optimizer() + d_opt = get_sync_optimizer() + train_ops = train.gan_train_ops( + model, + loss, + g_opt, + d_opt, + summarize_gradients=True, + colocate_gradients_with_ops=True) + + sequential_train_hooks = train.get_sequential_train_hooks()(train_ops) + self.assertLen(sequential_train_hooks, 4) + sync_opts = [ + hook._sync_optimizer for hook in sequential_train_hooks if + isinstance(hook, sync_replicas_optimizer._SyncReplicasOptimizerHook)] + self.assertLen(sync_opts, 2) + self.assertSetEqual(frozenset(sync_opts), frozenset((g_opt, d_opt))) + + joint_train_hooks = train.get_joint_train_hooks()(train_ops) + self.assertLen(joint_train_hooks, 5) + sync_opts = [ + hook._sync_optimizer for hook in joint_train_hooks if + isinstance(hook, sync_replicas_optimizer._SyncReplicasOptimizerHook)] + self.assertLen(sync_opts, 2) + self.assertSetEqual(frozenset(sync_opts), frozenset((g_opt, d_opt))) + class PatchGANTest(test.TestCase, parameterized.TestCase): """Tests that functions work on PatchGAN style output.""" diff --git a/tensorflow/contrib/gdr/BUILD b/tensorflow/contrib/gdr/BUILD index e534fdc17749974ebe713c2730682bea6d7a85e4..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", @@ -37,7 +32,7 @@ tf_proto_library_cc( ], ) -tf_cuda_library( +cc_library( name = "gdr_memory_manager", srcs = ["gdr_memory_manager.cc"], hdrs = ["gdr_memory_manager.h"], @@ -58,7 +53,7 @@ tf_cuda_library( ], ) -tf_cuda_library( +cc_library( name = "gdr_worker", srcs = ["gdr_worker.cc"], hdrs = ["gdr_worker.h"], @@ -66,7 +61,6 @@ tf_cuda_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.proto b/tensorflow/contrib/gdr/gdr.proto index c0b89245b150bfa49cb527d25b6e1f324f353b25..bd438787c3374be6ead4f6233101fd1f548643ea 100644 --- a/tensorflow/contrib/gdr/gdr.proto +++ b/tensorflow/contrib/gdr/gdr.proto @@ -9,5 +9,4 @@ message RemoteMemoryRegion { uint64 addr = 3; uint32 rkey = 4; uint32 tensor_key = 5; - uint64 checksum = 6; } 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..b84710d26eb8a64bf2f86b9f920551a8a8dbb233 --- /dev/null +++ b/tensorflow/contrib/gdr/gdr_collective_executor_mgr.cc @@ -0,0 +1,160 @@ +/* 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, dev_to_dev_stream_index, + 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 3549cedb70a6104ff3d3829d1b94cb5f08c5119c..7321e973191c4cc45f88735c6be7f2f67fe71c39 100644 --- a/tensorflow/contrib/gdr/gdr_memory_manager.cc +++ b/tensorflow/contrib/gdr/gdr_memory_manager.cc @@ -26,17 +26,14 @@ limitations under the License. #include #include #include -#include #include "tensorflow/contrib/gdr/gdr.pb.h" #include "tensorflow/core/common_runtime/device.h" #include "tensorflow/core/common_runtime/dma_helper.h" -#include "tensorflow/core/common_runtime/process_state.h" -#if GOOGLE_CUDA #include "tensorflow/core/common_runtime/gpu/gpu_process_state.h" -#include "tensorflow/core/common_runtime/gpu/gpu_util.h" -#endif // GOOGLE_CUDA +#include "tensorflow/core/common_runtime/process_state.h" #include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/lib/random/random.h" #include "tensorflow/core/platform/macros.h" #include "tensorflow/core/platform/mutex.h" #include "tensorflow/core/platform/numa.h" @@ -76,15 +73,14 @@ 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)) { if (value < 0) { - LOG(INFO) << "Successful NUMA node read from SysFS had negative value (" - << value - << "), but there must be at least one NUMA node" - ", so returning NUMA node zero"; return port::kNUMANoAffinity; } LOG(INFO) << "NUMA node for device: " << device->name << " is " << value; @@ -114,7 +110,7 @@ class GdrMemoryManager : public RemoteMemoryManager { public: GdrMemoryManager(const string& host, const string& port); - virtual ~GdrMemoryManager(); + virtual ~GdrMemoryManager() {} virtual Status Init() override; @@ -140,7 +136,7 @@ class GdrMemoryManager : public RemoteMemoryManager { return ptr < reinterpret_cast(other->addr) + other->length; } - ibv_mr* FindMemoryRegion(void* addr, size_t length); + ibv_mr* FindMemoryRegion(const Tensor* tensor); void InsertMemoryRegion(void* addr, size_t length, const std::string& allocator_name); @@ -152,7 +148,6 @@ class GdrMemoryManager : public RemoteMemoryManager { const string port_; RdmaEndpointPtr listening_; std::atomic stopped_; - int epfd_; int numa_node_; // Server side endpoints @@ -163,15 +158,19 @@ class GdrMemoryManager : public RemoteMemoryManager { std::atomic next_key_; // Server side on-the-fly tensor buffers - mutex server_mu_; - std::map tensor_buffers_ - GUARDED_BY(server_mu_); + mutex buf_mu_; + std::map tensor_buffers_ GUARDED_BY(buf_mu_); // Client side endpoints mutex client_mu_; std::map, RdmaEndpointPtr> clients_ GUARDED_BY(client_mu_); + // Client side callbacks + mutex callback_mu_; + std::map tensor_callbacks_ + GUARDED_BY(callback_mu_); + // Managed memory regions mutex alloc_mu_; std::vector mrs_ GUARDED_BY(alloc_mu_); @@ -184,16 +183,9 @@ GdrMemoryManager::GdrMemoryManager(const string& host, const string& port) port_(port), listening_(nullptr, EndpointDeleter), stopped_(true), - next_key_(0) {} - -GdrMemoryManager::~GdrMemoryManager() { close(epfd_); } + next_key_(static_cast(random::New64())) {} Status GdrMemoryManager::Init() { - epfd_ = epoll_create1(0); - if (epfd_ == -1) { - return errors::Unavailable(strerror(errno), ": ", "epoll_create"); - } - rdma_addrinfo* addrinfo; rdma_addrinfo hints = {}; hints.ai_port_space = RDMA_PS_TCP; @@ -206,7 +198,7 @@ Status GdrMemoryManager::Init() { ibv_qp_init_attr init_attr = {}; init_attr.qp_type = IBV_QPT_RC; - init_attr.cap.max_recv_wr = 32; + init_attr.cap.max_recv_wr = 1024; init_attr.cap.max_send_wr = 1; init_attr.cap.max_recv_sge = 1; init_attr.cap.max_send_sge = 1; @@ -239,14 +231,6 @@ Status GdrMemoryManager::Init() { "cannot set server to non-blocking mode"); } - epoll_event event = {}; - event.events = EPOLLIN | EPOLLPRI; - event.data.ptr = listening_.get(); - if (epoll_ctl(epfd_, EPOLL_CTL_ADD, listening_->channel->fd, &event)) { - return errors::Unavailable(strerror(errno), ": ", - "cannot add server to epoll"); - } - numa_node_ = TryToReadNumaNode(listening_->verbs->device); SubAllocator::Visitor alloc_visitor = [this](void* ptr, int numa_node, @@ -265,23 +249,23 @@ Status GdrMemoryManager::Init() { ProcessState::singleton()->AddCPUFreeVisitor(free_visitor); LOG(INFO) << "Instrumenting CPU allocator(s)"; -#if GOOGLE_CUDA - if (IsGDRAvailable()) { - int bus_id = numa_node_ + 1; + for (int numa_idx = 0; numa_idx < port::NUMANumNodes(); ++numa_idx) { + GPUProcessState::singleton()->AddCUDAHostAllocVisitor(numa_idx, + alloc_visitor); + GPUProcessState::singleton()->AddCUDAHostFreeVisitor(numa_idx, + free_visitor); + } + if (IsGDRAvailable()) { SubAllocator::Visitor cuda_alloc_visitor = [this](void* ptr, int gpu_id, size_t num_bytes) { VLOG(2) << "Registering RDMA capable memory region on GPU " << gpu_id; InsertMemoryRegion(ptr, num_bytes, strings::StrCat("GPU:", gpu_id)); }; - GPUProcessState::singleton()->AddGPUAllocVisitor(bus_id, + GPUProcessState::singleton()->AddGPUAllocVisitor(numa_node_, cuda_alloc_visitor); - GPUProcessState::singleton()->AddCUDAHostAllocVisitor(bus_id, - alloc_visitor); - GPUProcessState::singleton()->AddCUDAHostFreeVisitor(bus_id, free_visitor); - LOG(INFO) << "Instrumenting GPU allocator(s) with bus_id " << bus_id; + LOG(INFO) << "Instrumenting GPU allocator for NUMA " << numa_node_; } -#endif // GOOGLE_CUDA return Status::OK(); } @@ -289,95 +273,90 @@ Status GdrMemoryManager::Init() { void GdrMemoryManager::Run() { stopped_ = false; while (!stopped_) { - epoll_event events[32]; - int ret = epoll_wait(epfd_, events, 32, 1); - if (ret == -1) { - LOG(ERROR) << "epoll_wait: " << strerror(errno); - return; - } - for (int i = 0; i < ret; i++) { - rdma_cm_id* id = static_cast(events[i].data.ptr); - if (id == listening_.get()) { - // Accept incoming connections - if (!rdma_get_request(listening_.get(), &id)) { - if (!rdma_accept(id, nullptr)) { - LOG(INFO) << "Accepted new RDMA connection"; - if (ibv_req_notify_cq(id->recv_cq, 0)) { - LOG(ERROR) << strerror(errno) << ": ibv_req_notify_cq failed"; - EndpointDeleter(id); - continue; - } - for (int i = 0; i < 32; i++) { - if (rdma_post_recvv(id, nullptr, nullptr, 0)) { - LOG(ERROR) << strerror(errno) << ": rdma_post_recvv failed"; - EndpointDeleter(id); - continue; - } - } - int flags = fcntl(id->recv_cq_channel->fd, F_GETFL, 0); - if (fcntl(id->recv_cq_channel->fd, F_SETFL, flags | O_NONBLOCK)) { - LOG(ERROR) << strerror(errno) - << ": cannot set server_client to non-blocking mode"; - EndpointDeleter(id); - continue; - } - epoll_event event = {}; - event.events = EPOLLIN | EPOLLPRI; - event.data.ptr = id; - if (epoll_ctl(epfd_, EPOLL_CTL_ADD, id->recv_cq_channel->fd, - &event)) { - LOG(ERROR) << strerror(errno) - << ": cannot add server client to epoll"; - EndpointDeleter(id); - continue; - } - server_clients_.push_back({id, EndpointDeleter}); + rdma_cm_id* id = nullptr; + // Accept incoming connections + if (!rdma_get_request(listening_.get(), &id)) { + if (!rdma_accept(id, nullptr)) { + LOG(INFO) << "Accepted new RDMA connection"; + for (int i = 0; i < 1024; i++) { + if (rdma_post_recvv(id, nullptr, nullptr, 0)) { + LOG(ERROR) << strerror(errno) << ": rdma_post_recvv failed"; + EndpointDeleter(id); + continue; } } - } else { - // Polling work completions - ibv_cq* cq; - void* context; - if (!ibv_get_cq_event(id->recv_cq_channel, &cq, &context)) { - ibv_ack_cq_events(id->recv_cq, 1); - if (ibv_req_notify_cq(id->recv_cq, 0)) { - LOG(ERROR) << strerror(errno) << ": ibv_req_notify_cq failed"; - continue; + server_clients_.push_back({id, EndpointDeleter}); + } + } + // Polling server side work completions + for (const auto& client : server_clients_) { + ibv_wc wc[32]; + int ret = ibv_poll_cq(client->recv_cq, 32, wc); + if (ret < 0) { + LOG(ERROR) << "ibv_poll_cq failed"; + continue; + } + for (int i = 0; i < ret; i++) { + if (wc[i].opcode != IBV_WC_RECV_RDMA_WITH_IMM) { + LOG(ERROR) << "Received unknown operation " << wc[i].opcode; + } + if (wc[i].status != 0) { + LOG(ERROR) << ibv_wc_status_str(wc[i].status); + } + TensorKey tensor_key = ntohl(wc[i].imm_data); + + if (rdma_post_recvv(client.get(), nullptr, nullptr, 0)) { + perror("rdma_post_recvv"); + LOG(ERROR) << "rdma_post_recvv failed"; + } + + mutex_lock l(buf_mu_); + auto iter = tensor_buffers_.find(tensor_key); + if (iter == std::end(tensor_buffers_)) { + LOG(ERROR) << "Cannot find tensor buffer for tensor key " + << tensor_key; + } else { + const TensorBuffer* buffer = iter->second; + buffer->Unref(); + tensor_buffers_.erase(iter); + } + } + } + // Polling client side work completions + if (client_mu_.try_lock()) { + for (const auto& client : clients_) { + ibv_wc wc[32]; + int ret = ibv_poll_cq(client.second->send_cq, 32, wc); + for (int i = 0; i < ret; i++) { + Status s; + if (wc[i].status) { + s = errors::Unavailable(ibv_wc_status_str(wc[i].status)); + } else { + s = Status::OK(); } - ibv_wc wc[32]; - int ret = ibv_poll_cq(id->recv_cq, 32, wc); - if (ret < 0) { - LOG(ERROR) << "ibv_poll_cq failed"; - continue; + TensorKey key = wc[i].wr_id; + + ibv_send_wr wr = {}; + wr.opcode = IBV_WR_RDMA_WRITE_WITH_IMM; + wr.imm_data = htonl(key); + ibv_send_wr* bad_wr; + if (ibv_post_send(client.second->qp, &wr, &bad_wr)) { + LOG(ERROR) << strerror(errno) + << ": ibv_post_send failed for tensor_key " << key; } - for (int i = 0; i < ret; i++) { - if (wc[i].opcode != IBV_WC_RECV_RDMA_WITH_IMM) { - LOG(ERROR) << "Received unknown operation " << wc[i].opcode; - } - if (wc[i].status != 0) { - LOG(ERROR) << ibv_wc_status_str(wc[i].status); - } - TensorKey tensor_key = ntohl(wc[i].imm_data); - { - mutex_lock l(server_mu_); - auto iter = tensor_buffers_.find(tensor_key); - if (iter == std::end(tensor_buffers_)) { - LOG(ERROR) << "Cannot find tensor buffer for tensor key " - << tensor_key; - } else { - const TensorBuffer* buffer = iter->second; - buffer->Unref(); - tensor_buffers_.erase(iter); - } - } - if (rdma_post_recvv(id, nullptr, nullptr, 0)) { - perror("rdma_post_recvv"); - LOG(ERROR) << "rdma_post_recvv failed"; - continue; - } + + mutex_lock l(callback_mu_); + auto iter = tensor_callbacks_.find(key); + if (iter != std::end(tensor_callbacks_)) { + iter->second(s); + tensor_callbacks_.erase(iter); + } else { + LOG(WARNING) << "Cannot find client callback with tensor key " + << key; } } } + client_mu_.unlock(); } } } @@ -388,116 +367,58 @@ void GdrMemoryManager::TransportOptionsFromTensor( ::google::protobuf::Any* mutable_transport_options, const Tensor& tensor, Device* device, DeviceContext* device_context, bool on_host, StatusCallback done) { - auto buffer = DMAHelper::buffer(&tensor); - void* addr = buffer->data(); - size_t length = buffer->size(); - if (length == 0) { - done(errors::Unavailable("Cannot register tensor buffer of size 0")); - return; - } - - ibv_mr* mr = FindMemoryRegion(addr, length); - -#if GOOGLE_CUDA - if (device->tensorflow_gpu_device_info() && !on_host) { - Allocator* alloc = GPUProcessState::singleton()->GetCUDAHostAllocator(0); - Tensor* host_copy = new Tensor(alloc, tensor.dtype(), tensor.shape()); - GPUUtil::CopyGPUTensorToCPU( - device, device_context, &tensor, host_copy, - [done, host_copy, mutable_transport_options, this](const Status& s) { - if (!s.ok()) { - done(s); - delete host_copy; - return; - } - auto buffer = DMAHelper::buffer(host_copy); - void* addr = buffer->data(); - size_t length = buffer->size(); - ibv_mr* mr = FindMemoryRegion(addr, length); - - if (mr == nullptr) { - done(errors::Unavailable("Cannot find pinned memory region")); - delete host_copy; - return; - } - - buffer->Ref(); - TensorKey tensor_key = next_key_++; - { - mutex_lock l(server_mu_); - tensor_buffers_.insert(std::make_pair(tensor_key, buffer)); - } + ibv_mr* mr = FindMemoryRegion(&tensor); + const TensorBuffer* buffer = DMAHelper::buffer(&tensor); - uint64_t checksum = 0; - if (VLOG_IS_ON(2)) { - checksum = GPUUtil::Checksum(*host_copy); - } - - RemoteMemoryRegion remote_mr; - remote_mr.set_host(host_); - remote_mr.set_port(port_); - remote_mr.set_addr(reinterpret_cast(addr)); - remote_mr.set_rkey(mr->rkey); - remote_mr.set_tensor_key(tensor_key); - remote_mr.set_checksum(checksum); - mutable_transport_options->PackFrom(remote_mr); - - done(Status::OK()); - delete host_copy; - }); - return; - } -#endif + Tensor* copy = nullptr; if (mr == nullptr) { - Allocator* alloc = ProcessState::singleton()->GetCPUAllocator(numa_node_); - Tensor host_copy(alloc, tensor.dtype(), tensor.shape()); - - std::memcpy(DMAHelper::buffer(&host_copy)->data(), buffer->data(), length); - VLOG(2) << "Copying " << length << " bytes unpinned tensor buffer"; - - buffer = DMAHelper::buffer(&host_copy); - addr = buffer->data(); - length = buffer->size(); - - mr = FindMemoryRegion(addr, length); + AllocatorAttributes alloc_attrs; + alloc_attrs.set_gpu_compatible(true); + alloc_attrs.set_nic_compatible(true); + alloc_attrs.set_on_host(true); + Allocator* alloc = device->GetAllocator(alloc_attrs); + copy = new Tensor(alloc, tensor.dtype(), tensor.shape()); + + mr = FindMemoryRegion(copy); + buffer = DMAHelper::buffer(copy); if (mr == nullptr) { done(errors::Unavailable("Cannot find pinned memory region")); + delete copy; return; } - - buffer->Ref(); - } else { - buffer->Ref(); } TensorKey tensor_key = next_key_++; + buffer->Ref(); { - mutex_lock l(server_mu_); + mutex_lock l(buf_mu_); tensor_buffers_.insert(std::make_pair(tensor_key, buffer)); } - uint64_t checksum = 0; - if (VLOG_IS_ON(2)) { -#ifdef GOOGLE_CUDA - if (device->tensorflow_gpu_device_info() && !on_host) { - checksum = GPUUtil::Checksum(device, device_context, tensor); - } else { - checksum = GPUUtil::Checksum(tensor); - } -#endif - } - RemoteMemoryRegion remote_mr; remote_mr.set_host(host_); remote_mr.set_port(port_); - remote_mr.set_addr(reinterpret_cast(addr)); + remote_mr.set_addr(reinterpret_cast(buffer->data())); remote_mr.set_rkey(mr->rkey); remote_mr.set_tensor_key(tensor_key); - remote_mr.set_checksum(checksum); mutable_transport_options->PackFrom(remote_mr); - done(Status::OK()); + if (copy && device->tensorflow_gpu_device_info() && !on_host) { + device_context->CopyDeviceTensorToCPU(&tensor, "" /* tensor_name */, device, + copy, [done, copy](const Status& s) { + done(s); + delete copy; + }); + return; + } else if (copy) { + std::memcpy(buffer->data(), DMAHelper::buffer(&tensor)->data(), + buffer->size()); + done(Status::OK()); + delete copy; // OK to delete; we have reffed the underlying TensorBuffer + } else { + done(Status::OK()); + } } void GdrMemoryManager::TensorFromTransportOptions( @@ -510,42 +431,10 @@ void GdrMemoryManager::TensorFromTransportOptions( return; } - auto buffer = DMAHelper::buffer(tensor); - void* addr = buffer->data(); - size_t length = buffer->size(); - ibv_mr* mr = FindMemoryRegion(addr, length); - - Tensor host_copy; -#if GOOGLE_CUDA - if (mr == nullptr && !on_host) { - Allocator* alloc = - GPUProcessState::singleton()->GetCUDAHostAllocator(numa_node_); - host_copy = Tensor(alloc, tensor->dtype(), tensor->shape()); - buffer = DMAHelper::buffer(&host_copy); - addr = buffer->data(); - length = buffer->size(); - mr = FindMemoryRegion(addr, length); - } -#endif // GOOGLE_CUDA - - if (mr == nullptr) { - Allocator* alloc = ProcessState::singleton()->GetCPUAllocator(numa_node_); - host_copy = Tensor(alloc, tensor->dtype(), tensor->shape()); - - buffer = DMAHelper::buffer(&host_copy); - addr = buffer->data(); - length = buffer->size(); - - mr = FindMemoryRegion(addr, length); - if (mr == nullptr) { - done(errors::Unavailable("Cannot find pinned memory region")); - return; - } - } - - decltype(clients_)::iterator iter; - bool success; + rdma_cm_id* id = nullptr; { + decltype(clients_)::iterator iter; + bool success; mutex_lock l(client_mu_); std::tie(iter, success) = clients_.insert( std::make_pair(std::make_pair(remote_mr.host(), remote_mr.port()), @@ -558,93 +447,94 @@ void GdrMemoryManager::TensorFromTransportOptions( return; } } + id = iter->second.get(); } - rdma_cm_id* id = iter->second.get(); - uint64_t start = Env::Default()->NowMicros(); + ibv_mr* mr = FindMemoryRegion(tensor); + const TensorBuffer* buffer = DMAHelper::buffer(tensor); - if (rdma_post_read(id, nullptr, buffer->data(), buffer->size(), mr, 0, - remote_mr.addr(), remote_mr.rkey())) { - done(errors::Unavailable(strerror(errno), ": ", "rdma_post_read failed")); - return; - } - - ibv_send_wr wr = {}; - wr.opcode = IBV_WR_RDMA_WRITE_WITH_IMM; - wr.imm_data = htonl(remote_mr.tensor_key()); - wr.send_flags = IBV_SEND_SIGNALED; - ibv_send_wr* bad_wr; - if (ibv_post_send(id->qp, &wr, &bad_wr)) { - done(errors::Unavailable(strerror(errno), ": ", "ibv_post_send failed")); - return; - } - - ibv_wc wc = {}; - int ret; - while ((ret = ibv_poll_cq(id->send_cq, 1, &wc)) == 0) - ; - if (ret < 0 || wc.status) { - done(errors::Unavailable(ibv_wc_status_str(wc.status))); - return; - } + const Tensor* copy = nullptr; -#if GOOGLE_CUDA - if (device->tensorflow_gpu_device_info() && !on_host && - host_copy.NumElements() > 0) { - uint64_t checksum = 0; - if (VLOG_IS_ON(2)) { - checksum = GPUUtil::Checksum(host_copy); - CHECK(checksum == remote_mr.checksum()) - << "Checksum mismatch: " << checksum << "!=" << remote_mr.checksum(); + if (mr == nullptr) { + AllocatorAttributes alloc_attrs; + alloc_attrs.set_gpu_compatible(true); + alloc_attrs.set_nic_compatible(true); + alloc_attrs.set_on_host(true); + Allocator* alloc = device->GetAllocator(alloc_attrs); + copy = new Tensor(alloc, tensor->dtype(), tensor->shape()); + + mr = FindMemoryRegion(copy); + buffer = DMAHelper::buffer(copy); + if (mr == nullptr) { + done(errors::Unavailable("Cannot find pinned memory region")); + delete copy; + return; } - Tensor* ref = new Tensor; - std::swap(host_copy, *ref); - GPUUtil::CopyCPUTensorToGPU( - ref, device_context, device, tensor, - [ref, done, buffer, remote_mr, start](const Status& s) { - if (!s.ok()) { - done(s); - delete ref; - return; - } - uint64_t end = Env::Default()->NowMicros(); - - VLOG(2) << "RDMA from remote memory region " << remote_mr.rkey() - << " of size " << buffer->size() << " with tensor key " - << remote_mr.tensor_key() << " took " << (end - start) - << " micros"; - done(Status::OK()); - delete ref; - }); - return; } -#endif // GOOGLE_CUDA - if ((on_host || !device->tensorflow_gpu_device_info()) && - host_copy.NumElements() > 0) { - std::memcpy(DMAHelper::buffer(tensor)->data(), addr, length); - VLOG(2) << "Copying " << length << " bytes unpinned tensor buffer"; - } + uint64_t start = Env::Default()->NowMicros(); - uint64_t end = Env::Default()->NowMicros(); + TensorKey tensor_key = remote_mr.tensor_key(); - VLOG(2) << "RDMA from remote memory region " << remote_mr.rkey() - << " of size " << buffer->size() << " with tensor key " - << remote_mr.tensor_key() << " took " << (end - start) << " micros"; + StatusCallback callback = [done, copy, device, device_context, on_host, + tensor, start, tensor_key](const Status& s) { + if (!s.ok()) { + done(s); + if (copy) { + delete copy; + } + return; + } - uint64_t checksum = 0; - if (VLOG_IS_ON(2)) { -#ifdef GOOGLE_CUDA - if (device->tensorflow_gpu_device_info() && !on_host) { - checksum = GPUUtil::Checksum(device, device_context, *tensor); + VLOG(2) << "RDMA of tensor " << tensor_key << " of size " + << DMAHelper::buffer(tensor)->size() << " took " + << (Env::Default()->NowMicros() - start) << " micros"; + + if (copy && device->tensorflow_gpu_device_info() && !on_host) { + device_context->CopyCPUTensorToDevice(copy, device, tensor, + [done, copy](const Status& s) { + done(s); + delete copy; + }); + } else if (copy) { + std::memcpy(DMAHelper::buffer(tensor)->data(), + DMAHelper::buffer(copy)->data(), + DMAHelper::buffer(copy)->size()); + done(s); + delete copy; } else { - checksum = GPUUtil::Checksum(*tensor); + done(s); + } + }; + + { + mutex_lock l(callback_mu_); + if (tensor_callbacks_.find(tensor_key) == std::end(tensor_callbacks_)) { + tensor_callbacks_.insert(std::make_pair(tensor_key, std::move(callback))); + } else { + done(errors::Unavailable("Received duplicated tensor key")); + if (copy) { + delete copy; + } + return; + } + } + + if (rdma_post_read(id, reinterpret_cast(tensor_key), buffer->data(), + buffer->size(), mr, IBV_SEND_SIGNALED, remote_mr.addr(), + remote_mr.rkey())) { + done(errors::Unavailable(strerror(errno), ": ", "rdma_post_read failed")); + { + mutex_lock l(callback_mu_); + auto iter = tensor_callbacks_.find(tensor_key); + if (iter != std::end(tensor_callbacks_)) { + tensor_callbacks_.erase(iter); + } + } + if (copy) { + delete copy; } - CHECK(checksum == remote_mr.checksum()) - << "Checksum mismatch: " << checksum << "!=" << remote_mr.checksum(); -#endif } - done(Status::OK()); } Status GdrMemoryManager::CreateEndpoint(const string& host, const string& port, @@ -661,7 +551,7 @@ Status GdrMemoryManager::CreateEndpoint(const string& host, const string& port, ibv_qp_init_attr init_attr = {}; init_attr.qp_type = IBV_QPT_RC; init_attr.cap.max_recv_wr = 1; - init_attr.cap.max_send_wr = 32; + init_attr.cap.max_send_wr = 1024; init_attr.cap.max_recv_sge = 1; init_attr.cap.max_send_sge = 1; @@ -685,8 +575,8 @@ Status GdrMemoryManager::CreateEndpoint(const string& host, const string& port, return Status::OK(); } -ibv_mr* GdrMemoryManager::FindMemoryRegion(void* addr, size_t length) { - if (length == 0) return nullptr; +ibv_mr* GdrMemoryManager::FindMemoryRegion(const Tensor* tensor) { + const void* addr = DMAHelper::buffer(tensor)->data(); mutex_lock l(alloc_mu_); auto iter = std::upper_bound(mrs_.begin(), mrs_.end(), addr, &Comparator); if (iter == std::end(mrs_) || iter->get()->addr > addr) { diff --git a/tensorflow/contrib/gdr/gdr_rendezvous_mgr.cc b/tensorflow/contrib/gdr/gdr_rendezvous_mgr.cc index 94f522c04e5a09ed2d9355fa675125c340407923..5f8c300155770ed03ad12a9fa5ac74456edaf024 100644 --- a/tensorflow/contrib/gdr/gdr_rendezvous_mgr.cc +++ b/tensorflow/contrib/gdr/gdr_rendezvous_mgr.cc @@ -58,11 +58,9 @@ class GdrRecvTensorCall : public BaseRecvTensorCall { resp_.InitAlloc(dst_device_, recv_args_.alloc_attrs); StatusCallback cb = [this, recv_done](const Status& s) { bool dma_ok = resp_.metadata().has_transport_options(); - if (s.ok() && tensor().TotalBytes() > 0 && (!is_dead()) && dma_ok) { + if (s.ok() && tensor().TotalBytes() > 1024 && (!is_dead()) && dma_ok) { auto transport_options = resp_.metadata().transport_options(); - const bool on_host = - (dst_device_->tensorflow_gpu_device_info() == nullptr) || - recv_args_.alloc_attrs.on_host(); + const bool on_host = recv_args_.alloc_attrs.on_host(); remote_memory_manager_->TensorFromTransportOptions( const_cast(&tensor()), transport_options, dst_device_, recv_args_.device_context, on_host, @@ -70,9 +68,6 @@ class GdrRecvTensorCall : public BaseRecvTensorCall { if (!s.ok()) { mutex_lock l(mu_); status_.Update(s); - LOG(ERROR) << "Cannot find pinned memory region from allocator " - << dst_device_->GetAllocator(recv_args_.alloc_attrs) - ->Name(); } recv_done(); }); @@ -170,6 +165,14 @@ class GdrRemoteRendezvous : public BaseRemoteRendezvous { // Record "call" in active_ so that it can be aborted cleanly. RegisterCall(call); + // RendezvousMgr already aborted, shouldn't send RPC call any more + if (!call->status().ok()) { + done(call->status(), Args(), Args(), Tensor(), false); + session()->worker_cache->ReleaseWorker(src_worker, rwi); + delete call; + return; + } + // Start "call". Ref(); call->Start([this, call, src_worker, rwi, done]() { diff --git a/tensorflow/contrib/gdr/gdr_server_lib.cc b/tensorflow/contrib/gdr/gdr_server_lib.cc index 9025c992a4467f521d6d8d514e6a5e92f5492947..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 { @@ -52,14 +54,39 @@ Status GdrServer::Init() { [this](const WorkerEnv* env) { return new GdrRendezvousMgr(env, remote_memory_manager_.get()); }; - WorkerCreationFunction worker_func = [this](WorkerEnv* env) { + WorkerCreationFunction worker_func = [this](WorkerEnv* env, + const ConfigProto& config) { return std::unique_ptr( - new GdrWorker(env, remote_memory_manager_.get())); + new GdrWorker(env, config, remote_memory_manager_.get())); }; - TF_RETURN_IF_ERROR( - GrpcServer::Init(nullptr, rendezvous_mgr_func, nullptr, worker_func)); + 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 remote_memory_manager_->Init(); + 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() { @@ -73,9 +100,8 @@ Status GdrServer::Start() { } Status GdrServer::Stop() { - TF_RETURN_IF_ERROR(GrpcServer::Stop()); remote_memory_manager_->Stop(); - return Status::OK(); + return GrpcServer::Stop(); } Status GdrServer::Join() { diff --git a/tensorflow/contrib/gdr/gdr_worker.cc b/tensorflow/contrib/gdr/gdr_worker.cc index ce1d8d2d73000559f03046aceacb169890ecc1b6..1204b8ca501a8f99ea6abd6c047ab2d91350bae1 100644 --- a/tensorflow/contrib/gdr/gdr_worker.cc +++ b/tensorflow/contrib/gdr/gdr_worker.cc @@ -15,12 +15,10 @@ 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" -#if GOOGLE_CUDA -#include "tensorflow/core/common_runtime/gpu/gpu_util.h" -#endif // GOOGLE_CUDA #include "tensorflow/core/common_runtime/process_util.h" #include "tensorflow/core/common_runtime/step_stats_collector.h" #include "tensorflow/core/distributed_runtime/graph_mgr.h" @@ -32,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" @@ -39,17 +38,17 @@ limitations under the License. namespace tensorflow { -GdrWorker::GdrWorker(WorkerEnv* worker_env, +GdrWorker::GdrWorker(WorkerEnv* worker_env, const ConfigProto& config, RemoteMemoryManager* remote_memory_manager) - : GrpcWorker(worker_env), + : 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); @@ -78,7 +77,7 @@ void GdrWorker::GrpcRecvTensorAsync(CallOptions* opts, const bool dma_ok = request->dma_ok(); env_->rendezvous_mgr->RecvLocalAsync( step_id, parsed, - [this, opts, response, done, src_dev, dma_ok]( + [this, opts, response, done, src_dev, request, dma_ok]( const Status& status, const Rendezvous::Args& send_args, const Rendezvous::Args&, const Tensor& val, const bool is_dead) { opts->ClearCancelCallback(); @@ -89,10 +88,8 @@ void GdrWorker::GrpcRecvTensorAsync(CallOptions* opts, // 3) the tensor has the on_host allocation attribute, // i.e. it's in CPU RAM *independent of its assigned // device type*. - const bool on_host = - (src_dev->tensorflow_gpu_device_info() == nullptr) || - send_args.alloc_attrs.on_host(); - if (val.TotalBytes() > 0 && (!is_dead) && + const bool on_host = send_args.alloc_attrs.on_host(); + if (val.TotalBytes() > 1024 && (!is_dead) && DMAHelper::CanUseDMA(&val) && dma_ok) { // DMA cases. RecvTensorResponse* proto = new RecvTensorResponse; @@ -117,8 +114,7 @@ void GdrWorker::GrpcRecvTensorAsync(CallOptions* opts, } else { // Non-DMA cases. if (src_dev->tensorflow_gpu_device_info() && (!on_host)) { -#if GOOGLE_CUDA - const DeviceContext* send_dev_context = send_args.device_context; + DeviceContext* send_dev_context = send_args.device_context; AllocatorAttributes alloc_attrs; alloc_attrs.set_gpu_compatible(true); alloc_attrs.set_on_host(true); @@ -127,7 +123,8 @@ void GdrWorker::GrpcRecvTensorAsync(CallOptions* opts, CHECK(send_dev_context) << "send dev name: " << src_dev->name() << " gpu_info: " << src_dev->tensorflow_gpu_device_info(); - // "val" is on a GPU. Uses GPUUtil to fill the response proto. + // "val" is on an accelerator device. Uses the device_context to + // fill the copy on host. StatusCallback copy_ready = [response, done, copy, is_dead](const Status& s) { // The value is now ready to be returned on the wire. @@ -136,11 +133,8 @@ void GdrWorker::GrpcRecvTensorAsync(CallOptions* opts, delete copy; }; - GPUUtil::CopyGPUTensorToCPU(src_dev, send_dev_context, &val, copy, - copy_ready); -#else - done(errors::Internal("No GPU device in process")); -#endif // GOOGLE_CUDA + send_dev_context->CopyDeviceTensorToCPU( + &val, request->rendezvous_key(), src_dev, copy, copy_ready); } else { grpc::EncodeTensorToByteBuffer(is_dead, val, response); done(Status::OK()); @@ -153,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 65105ed997300aa77202301cdd8dddacb0309880..9a85cfd4263ad86f6579eedce95969c2829ff62c 100644 --- a/tensorflow/contrib/gdr/gdr_worker.h +++ b/tensorflow/contrib/gdr/gdr_worker.h @@ -25,7 +25,8 @@ namespace tensorflow { class GdrWorker : public GrpcWorker { public: - GdrWorker(WorkerEnv* env, RemoteMemoryManager* remote_memory_manager); + GdrWorker(WorkerEnv* env, const ConfigProto& config, + RemoteMemoryManager* remote_memory_manager); // Serve the RecvTensorRequest but omit the tensor content and transmit it // out-of-band using GPU Direct RDMA whenever possible. @@ -37,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/grid_rnn/python/kernel_tests/grid_rnn_test.py b/tensorflow/contrib/grid_rnn/python/kernel_tests/grid_rnn_test.py index 27aed091c249caa6e50748419a93f3579e6632a4..363a3c9b4a59eb8f769b16a11a3ade0643358c64 100644 --- a/tensorflow/contrib/grid_rnn/python/kernel_tests/grid_rnn_test.py +++ b/tensorflow/contrib/grid_rnn/python/kernel_tests/grid_rnn_test.py @@ -696,8 +696,8 @@ class GridRNNCellTest(test.TestCase): for out, inp in zip(outputs, inputs): self.assertEqual(len(out), 1) - self.assertTrue(out[0].get_shape()[0].value is None) - self.assertEqual(out[0].get_shape()[1], num_units) + self.assertTrue(out[0].get_shape().dims[0].value is None) + self.assertEqual(out[0].get_shape().dims[1], num_units) self.assertEqual(out[0].dtype, inp.dtype) with self.cached_session() as sess: diff --git a/tensorflow/contrib/grid_rnn/python/ops/grid_rnn_cell.py b/tensorflow/contrib/grid_rnn/python/ops/grid_rnn_cell.py index bcd2a34c4e791a2ab66a439109145d6b78c14e22..5f3af43a474a12787a111f8674b2b7bf0bb2481a 100644 --- a/tensorflow/contrib/grid_rnn/python/ops/grid_rnn_cell.py +++ b/tensorflow/contrib/grid_rnn/python/ops/grid_rnn_cell.py @@ -21,6 +21,7 @@ from __future__ import print_function from collections import namedtuple import functools +from tensorflow.python.framework import tensor_shape from tensorflow.python.ops import array_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import nn @@ -281,7 +282,8 @@ class GridRNNCell(rnn.RNNCell): """ conf = self._config - if (inputs is not None and inputs.get_shape().with_rank(2)[1].value > 0 and + if (inputs is not None and + tensor_shape.dimension_value(inputs.shape.with_rank(2)[1]) > 0 and conf.inputs): if isinstance(inputs, tuple): if len(conf.inputs) != len(inputs): @@ -291,7 +293,8 @@ class GridRNNCell(rnn.RNNCell): else: input_splits = array_ops.split( value=inputs, num_or_size_splits=len(conf.inputs), axis=1) - input_sz = input_splits[0].get_shape().with_rank(2)[1].value + input_sz = tensor_shape.dimension_value( + input_splits[0].shape.with_rank(2)[1]) for i, j in enumerate(conf.inputs): input_project_m = vs.get_variable( diff --git a/tensorflow/contrib/hadoop/python/kernel_tests/hadoop_test.py b/tensorflow/contrib/hadoop/python/kernel_tests/hadoop_test.py index f7f1189bb93c611719186a697c40f371644f63a2..bc941ae9f23eaa5c46fcca95b9aba0ac0d87960a 100644 --- a/tensorflow/contrib/hadoop/python/kernel_tests/hadoop_test.py +++ b/tensorflow/contrib/hadoop/python/kernel_tests/hadoop_test.py @@ -21,6 +21,7 @@ from __future__ import print_function import os from tensorflow.contrib.hadoop.python.ops import hadoop_dataset_ops +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 @@ -47,7 +48,7 @@ class SequenceFileDatasetTest(test.TestCase): dataset = hadoop_dataset_ops.SequenceFileDataset(filenames).repeat( num_repeats) - iterator = dataset.make_initializable_iterator() + iterator = dataset_ops.make_initializable_iterator(dataset) init_op = iterator.initializer get_next = iterator.get_next() diff --git a/tensorflow/contrib/hadoop/python/ops/hadoop_dataset_ops.py b/tensorflow/contrib/hadoop/python/ops/hadoop_dataset_ops.py index bf398b838dfaaff6fdaf33a6cd7086ef13e43a3e..71eac729a8a81c2f59f9ed5d7f42fb7b1c3e1b5c 100644 --- a/tensorflow/contrib/hadoop/python/ops/hadoop_dataset_ops.py +++ b/tensorflow/contrib/hadoop/python/ops/hadoop_dataset_ops.py @@ -20,15 +20,19 @@ from __future__ import print_function from tensorflow.contrib.hadoop.python.ops import gen_dataset_ops from tensorflow.contrib.hadoop.python.ops import hadoop_op_loader # pylint: disable=unused-import from tensorflow.python.data.ops import dataset_ops -from tensorflow.python.data.util import nest +from tensorflow.python.data.util import structure from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops -from tensorflow.python.framework import tensor_shape +from tensorflow.python.util import deprecation class SequenceFileDataset(dataset_ops.DatasetSource): """A Sequence File Dataset that reads the sequence file.""" + @deprecation.deprecated( + None, + "tf.contrib.hadoop will be removed in 2.0, the support for Apache Hadoop " + "will continue to be provided through the tensorflow/io GitHub project.") def __init__(self, filenames): """Create a `SequenceFileDataset`. @@ -40,36 +44,25 @@ class SequenceFileDataset(dataset_ops.DatasetSource): For example: ```python + tf.enable_eager_execution() + dataset = tf.contrib.hadoop.SequenceFileDataset("/foo/bar.seq") - iterator = dataset.make_one_shot_iterator() - next_element = iterator.get_next() # Prints the (key, value) pairs inside a hadoop sequence file. - while True: - try: - print(sess.run(next_element)) - except tf.errors.OutOfRangeError: - break + for key, value in dataset: + print(key, value) ``` Args: filenames: A `tf.string` tensor containing one or more filenames. """ - super(SequenceFileDataset, self).__init__() self._filenames = ops.convert_to_tensor( filenames, dtype=dtypes.string, name="filenames") - - def _as_variant_tensor(self): - return gen_dataset_ops.sequence_file_dataset( - self._filenames, nest.flatten(self.output_types)) - - @property - def output_classes(self): - return ops.Tensor, ops.Tensor - - @property - def output_shapes(self): - return (tensor_shape.TensorShape([]), tensor_shape.TensorShape([])) + variant_tensor = gen_dataset_ops.sequence_file_dataset( + self._filenames, self._element_structure._flat_types) # pylint: disable=protected-access + super(SequenceFileDataset, self).__init__(variant_tensor) @property - def output_types(self): - return dtypes.string, dtypes.string + def _element_structure(self): + return structure.NestedStructure( + (structure.TensorStructure(dtypes.string, []), + structure.TensorStructure(dtypes.string, []))) 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/BUILD b/tensorflow/contrib/ignite/BUILD index 9393b702d11a2ef84586f712d30c26fe2a8972bb..2698b83a56a1121fa30f5b05ffa027b4dfd4ba95 100644 --- a/tensorflow/contrib/ignite/BUILD +++ b/tensorflow/contrib/ignite/BUILD @@ -22,48 +22,92 @@ py_library( srcs_version = "PY2AND3", deps = [ ":dataset_ops", + ":igfs_ops", ], ) tf_custom_op_library( - name = "_dataset_ops.so", - srcs = ["ops/dataset_ops.cc"], - deps = [":dataset_kernels"], + name = "_ignite_ops.so", + srcs = [ + "kernels/igfs/igfs.h", + "ops/dataset_ops.cc", + "ops/igfs_ops.cc", + ], + deps = [ + ":dataset_kernels", + ":igfs_kernels", + ], ) tf_gen_op_libs( op_lib_names = ["dataset_ops"], ) +tf_gen_op_libs( + op_lib_names = ["igfs_ops"], + deps = [":igfs_kernels"], +) + cc_library( - name = "dataset_kernels", + name = "ignite_client", srcs = [ - "kernels/ignite_dataset_ops.cc", - "kernels/ignite_client.h", - "kernels/ignite_byte_swapper.h", - "kernels/ignite_plain_client.h", - "kernels/ignite_ssl_wrapper.h", - "kernels/ignite_ssl_wrapper.cc", - "kernels/ignite_binary_object_parser.h", - "kernels/ignite_binary_object_parser.cc", - "kernels/ignite_dataset.h", - "kernels/ignite_dataset.cc", - "kernels/ignite_dataset_iterator.h", - "kernels/ignite_dataset_iterator.cc", + "kernels/client/ignite_client.h", + "kernels/client/ignite_byte_swapper.h", + "kernels/client/ignite_plain_client.h", + "kernels/client/ignite_ssl_wrapper.h", + "kernels/client/ignite_ssl_wrapper.cc", ] + if_not_windows([ - "kernels/ignite_plain_client_unix.cc", + "kernels/client/ignite_plain_client_unix.cc", ]) + if_windows([ - "kernels/ignite_plain_client_windows.cc", + "kernels/client/ignite_plain_client_windows.cc", ]), copts = if_windows([ "-DWIN32_LEAN_AND_MEAN", ]), deps = [ "//tensorflow/core:framework_headers_lib", - "//third_party/eigen3", "@boringssl//:ssl", "@protobuf_archive//:protobuf_headers", ], +) + +cc_library( + name = "dataset_kernels", + srcs = [ + "kernels/dataset/ignite_binary_object_parser.cc", + "kernels/dataset/ignite_binary_object_parser.h", + "kernels/dataset/ignite_dataset.cc", + "kernels/dataset/ignite_dataset.h", + "kernels/dataset/ignite_dataset_iterator.cc", + "kernels/dataset/ignite_dataset_iterator.h", + "kernels/dataset/ignite_dataset_ops.cc", + ], + deps = [ + ":ignite_client", + "//tensorflow/core:framework_headers_lib", + "//third_party/eigen3", + "@protobuf_archive//:protobuf_headers", + ], + alwayslink = 1, +) + +cc_library( + name = "igfs_kernels", + srcs = [ + "kernels/igfs/igfs.cc", + "kernels/igfs/igfs.h", + "kernels/igfs/igfs_client.cc", + "kernels/igfs/igfs_client.h", + "kernels/igfs/igfs_extended_tcp_client.cc", + "kernels/igfs/igfs_extended_tcp_client.h", + "kernels/igfs/igfs_messages.cc", + "kernels/igfs/igfs_messages.h", + "kernels/igfs/igfs_random_access_file.cc", + "kernels/igfs/igfs_random_access_file.h", + "kernels/igfs/igfs_writable_file.cc", + "kernels/igfs/igfs_writable_file.h", + ], + deps = [":ignite_client"], alwayslink = 1, ) @@ -82,10 +126,29 @@ py_library( ], ) +py_library( + name = "igfs_ops", + srcs = [ + "python/ops/igfs_ops.py", + ], + srcs_version = "PY2AND3", + deps = [ + ":igfs_op_loader", + "//tensorflow/python:util", + "//tensorflow/python/data/util:nest", + ], +) + tf_gen_op_wrapper_py( name = "gen_dataset_ops", out = "python/ops/gen_dataset_ops.py", - deps = ["//tensorflow/contrib/ignite:dataset_ops_op_lib"], + deps = [":dataset_ops_op_lib"], +) + +tf_gen_op_wrapper_py( + name = "gen_igfs_ops", + out = "python/ops/gen_igfs_ops.py", + deps = [":igfs_ops_op_lib"], ) tf_kernel_library( @@ -97,13 +160,22 @@ tf_kernel_library( alwayslink = 1, ) +tf_kernel_library( + name = "igfs_ops_kernels", + deps = [ + ":igfs_kernels", + "//tensorflow/core:framework", + ], + alwayslink = 1, +) + tf_custom_op_py_library( name = "ignite_op_loader", srcs = ["python/ops/ignite_op_loader.py"], - dso = ["//tensorflow/contrib/ignite:_dataset_ops.so"], + dso = [":_ignite_ops.so"], kernels = [ ":dataset_ops_kernels", - "//tensorflow/contrib/ignite:dataset_ops_op_lib", + ":dataset_ops_op_lib", ], srcs_version = "PY2AND3", deps = [ @@ -113,6 +185,22 @@ tf_custom_op_py_library( ], ) +tf_custom_op_py_library( + name = "igfs_op_loader", + srcs = ["python/ops/igfs_op_loader.py"], + dso = [":_ignite_ops.so"], + kernels = [ + ":igfs_ops_kernels", + ":igfs_ops_op_lib", + ], + srcs_version = "PY2AND3", + deps = [ + ":gen_igfs_ops", + "//tensorflow/contrib/util:util_py", + "//tensorflow/python:platform", + ], +) + # The Apache Ignite servers have to setup before the test and tear down # after the test manually. The docker engine has to be installed. # @@ -122,8 +210,11 @@ tf_custom_op_py_library( # To tear down Apache Ignite servers: # $ bash ./python/tests/stop_ignite.sh tf_py_test( - name = "ignite_dataset_test", - srcs = ["python/tests/ignite_dataset_test.py"], + name = "ignite_test", + srcs = [ + "python/tests/igfs_test.py", + "python/tests/ignite_dataset_test.py", + ], additional_deps = [ ":ignite", "//tensorflow/python:client_testlib", diff --git a/tensorflow/contrib/ignite/README.md b/tensorflow/contrib/ignite/README.md index 55c89d27996318dabb29bb15372411005301ebd9..8e49e619e50c0fe6a128c2dcaefde17b365ce518 100644 --- a/tensorflow/contrib/ignite/README.md +++ b/tensorflow/contrib/ignite/README.md @@ -1,23 +1,37 @@ -# Ignite Dataset - -- [Overview](#overview) -- [Features](#features) - * [Distributed In-Memory Datasource](#distributed-in-memory-datasource) - * [Structured Objects](#structured-objects) - * [Distributed Training](#distributed-training) - * [SSL Connection](#ssl-connection) - * [Windows Support](#windows-support) -- [Try it out](#try-it-out) -- [Limitations](#limitations) +# Apache Ignite Integration + +- [Overview](#overview) +- [Features](#features) + * [Distributed In-Memory Datasource](#distributed-in-memory-datasource) + * [Structured Objects](#structured-objects) + * [Distributed Training](#distributed-training) + * [Distributed File System](#distributed-file-system) + * [SSL Connection](#ssl-connection) + * [Windows Support](#windows-support) +- [Try it out](#try-it-out) + * [Ignite Dataset](#ignite-dataset) + * [IGFS](#igfs) +- [Limitations](#limitations) ## Overview -[Apache Ignite](https://ignite.apache.org/) is a memory-centric distributed database, caching, and processing platform for -transactional, analytical, and streaming workloads, delivering in-memory speeds at petabyte scale. This contrib package contains an integration between Apache Ignite and TensorFlow. The integration is based on [tf.data](https://www.tensorflow.org/api_docs/python/tf/data) from TensorFlow side and [Binary Client Protocol](https://apacheignite.readme.io/v2.6/docs/binary-client-protocol) from Apache Ignite side. It allows to use Apache Ignite as a data source for neural network training, inference and all other computations supported by TensorFlow. +[Apache Ignite](https://ignite.apache.org/) is a memory-centric distributed +database, caching, and processing platform for transactional, analytical, and +streaming workloads, delivering in-memory speeds at petabyte scale. This contrib +package contains an integration between Apache Ignite and TensorFlow. The +integration is based on +[tf.data](https://www.tensorflow.org/api_docs/python/tf/data) from TensorFlow +side and +[Binary Client Protocol](https://apacheignite.readme.io/v2.6/docs/binary-client-protocol) +from Apache Ignite side. It allows to use Apache Ignite as a data source for +neural network training, inference and all other computations supported by +TensorFlow. Another part of this module is an integration with distributed file +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 @@ -41,14 +55,12 @@ 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="SQL_PUBLIC_KITTEN_CACHE") ->>> iterator = dataset.make_one_shot_iterator() ->>> next_obj = iterator.get_next() >>> ->>> with tf.Session() as sess: ->>> for _ in range(3): ->>> print(sess.run(next_obj)) +>>> for element in dataset: +>>> print(element) {'key': 1, 'val': {'NAME': b'WARM KITTY'}} {'key': 2, 'val': {'NAME': b'SOFT KITTY'}} @@ -61,23 +73,22 @@ 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") ->>> iterator = dataset.make_one_shot_iterator() ->>> next_obj = iterator.get_next() >>> ->>> with tf.Session() as sess: ->>> print(sess.run(next_obj)) +>>> for element in dataset.take(1): +>>> print(element) { - 'key': 'kitten.png', + 'key': 'kitten.png', 'val': { 'metadata': { 'file_name': b'kitten.png', 'label': b'little ball of fur', - width: 800, + width: 800, height: 600 - }, + }, 'pixels': [0, 0, 0, 0, ..., 0] } } @@ -87,13 +98,11 @@ jdbc:ignite:thin://localhost/> INSERT INTO KITTEN_CACHE VALUES (3, 'LITTLE BALL ```python >>> import tensorflow as tf >>> from tensorflow.contrib.ignite import IgniteDataset ->>> +>>> >>> dataset = IgniteDataset(cache_name="IMAGES").map(lambda obj: obj['val']['pixels']) ->>> iterator = dataset.make_one_shot_iterator() ->>> next_obj = iterator.get_next() >>> ->>> with tf.Session() as sess: ->>> print(sess.run(next_obj)) +>>> for element in dataset: +>>> print(element) [0, 0, 0, 0, ..., 0] ``` @@ -108,23 +117,31 @@ 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 >>> from tensorflow.contrib.ignite import IgniteDataset ->>> +>>> >>> dataset = IgniteDataset("IMAGES") >>> >>> # Compute gradients locally on every worker node. ->>> gradients = [] +>>> gradients = [] >>> for i in range(5): >>> with tf.device("/job:WORKER/task:%d" % i): ->>> device_iterator = dataset.make_one_shot_iterator() +>>> device_iterator = tf.compat.v1.data.make_one_shot_iterator(dataset) >>> device_next_obj = device_iterator.get_next() >>> gradient = compute_gradient(device_next_obj) ->>> gradients.append(gradient) ->>> +>>> gradients.append(gradient) +>>> >>> # Aggregate them on master node. >>> result_gradient = tf.reduce_sum(gradients) >>> @@ -132,18 +149,43 @@ Ignite Dataset allows using these two aspects of distributed neural network trai >>> print(sess.run(result_gradient)) ``` -High-level TensorFlow API for [distributed training](https://www.tensorflow.org/api_docs/python/tf/contrib/distribute/DistributionStrategy) is supported as well. +High-level TensorFlow API for [distributed training](https://www.tensorflow.org/api_docs/python/tf/contrib/distribute/DistributionStrategy) is supported as well. + +### Distributed File System + +In addition to database functionality Apache Ignite provides a distributed file +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 +[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. ### 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 ->>> ->>> dataset = IgniteDataset(cache_name="IMAGES", certfile="client.pem", cert_password="password", username="ignite", password="ignite") ->>> ... +>>> +>>> dataset = IgniteDataset(cache_name="IMAGES", + certfile="client.pem", + cert_password="password", + username="ignite", + password="ignite") ``` ### Windows Support @@ -152,7 +194,16 @@ Ignite Dataset is fully compatible with Windows. You can use it as part of Tenso ## Try it out -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 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: +Following examples will help you to easily start working with this module. + +### Ignite Dataset + +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 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: ``` docker run -it -p 10800:10800 dmitrievanthony/ignite-with-mnist @@ -162,6 +213,35 @@ 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") +### 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 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/). +You need to start this container on your machine: + +``` +docker run -it -p 10500:10500 dmitrievanthony/ignite-with-igfs +``` + +After that you will be able to work with it following way: + +```python +>>> import tensorflow as tf +>>> import tensorflow.contrib.ignite.python.ops.igfs_ops +>>> +>>> with tf.gfile.Open("igfs:///hello.txt", mode='w') as w: +>>> w.write("Hello, world!") +>>> +>>> with tf.gfile.Open("igfs:///hello.txt", mode='r') as r: +>>> print(r.read()) + +Hello, world! +``` + ## Limitations Presently, Ignite Dataset works with assumption that all objects in the cache have the same structure (homogeneous objects) and the cache contains at least one object. Another limitation concerns structured objects, Ignite Dataset does not support UUID, Maps and Object arrays that might be parts of an object structure. diff --git a/tensorflow/contrib/ignite/kernels/client/ignite_byte_swapper.h b/tensorflow/contrib/ignite/kernels/client/ignite_byte_swapper.h new file mode 100644 index 0000000000000000000000000000000000000000..aac950fcc2aaf016959bbda876ac93df4baea417 --- /dev/null +++ b/tensorflow/contrib/ignite/kernels/client/ignite_byte_swapper.h @@ -0,0 +1,125 @@ +/* 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_IGNITE_KERNELS_CLIENT_IGNITE_BYTE_SWAPPER_H_ +#define TENSORFLOW_CONTRIB_IGNITE_KERNELS_CLIENT_IGNITE_BYTE_SWAPPER_H_ + +#include +#include "tensorflow/core/platform/byte_order.h" + +namespace tensorflow { + +class ByteSwapper { + public: + ByteSwapper(bool big_endian) { swap_ = big_endian == port::kLittleEndian; } + + void SwapIfRequiredInt16(int16_t *x) const { + if (swap_) { + Swap16(x); + } + } + + void SwapIfRequiredUnsignedInt16(uint16_t *x) const { + if (swap_) { + Swap16(reinterpret_cast(x)); + } + } + + void SwapIfRequiredInt32(int32_t *x) const { + if (swap_) { + Swap32(x); + } + } + + void SwapIfRequiredFloat(float *x) const { + if (swap_) { + Swap32(reinterpret_cast(x)); + } + } + + void SwapIfRequiredInt64(int64_t *x) const { + if (swap_) { + Swap64(x); + } + } + + void SwapIfRequiredDouble(double *x) const { + if (swap_) { + Swap64(reinterpret_cast(x)); + } + } + + void SwapIfRequiredInt16Arr(int16_t *x, int32_t length) const { + if (swap_) { + for (int32_t i = 0; i < length; i++) Swap16(&x[i]); + } + } + + void SwapIfRequiredUnsignedInt16Arr(uint16_t *x, int32_t length) const { + if (swap_) { + for (int32_t i = 0; i < length; i++) + Swap16(reinterpret_cast(&x[i])); + } + } + + void SwapIfRequiredInt32Arr(int32_t *x, int32_t length) const { + if (swap_) { + for (int32_t i = 0; i < length; i++) Swap32(&x[i]); + } + } + + void SwapIfRequiredFloatArr(float *x, int32_t length) const { + if (swap_) { + for (int32_t i = 0; i < length; i++) + Swap32(reinterpret_cast(&x[i])); + } + } + + void SwapIfRequiredInt64Arr(int64_t *x, int32_t length) const { + if (swap_) { + for (int32_t i = 0; i < length; i++) Swap64(&x[i]); + } + } + + void SwapIfRequiredDoubleArr(double *x, int32_t length) const { + if (swap_) { + for (int32_t i = 0; i < length; i++) + Swap64(reinterpret_cast(&x[i])); + } + } + + private: + void Swap16(int16_t *x) const { + *x = ((*x & 0xFF) << 8) | ((*x >> 8) & 0xFF); + } + + void Swap32(int32_t *x) const { + *x = ((*x & 0xFF) << 24) | (((*x >> 8) & 0xFF) << 16) | + (((*x >> 16) & 0xFF) << 8) | ((*x >> 24) & 0xFF); + } + + void Swap64(int64_t *x) const { + *x = ((*x & 0xFF) << 56) | (((*x >> 8) & 0xFF) << 48) | + (((*x >> 16) & 0xFF) << 40) | (((*x >> 24) & 0xFF) << 32) | + (((*x >> 32) & 0xFF) << 24) | (((*x >> 40) & 0xFF) << 16) | + (((*x >> 48) & 0xFF) << 8) | ((*x >> 56) & 0xFF); + } + + bool swap_; +}; + +} // namespace tensorflow + +#endif // TENSORFLOW_CONTRIB_IGNITE_KERNELS_CLIENT_IGNITE_BYTE_SWAPPER_H_ diff --git a/tensorflow/contrib/ignite/kernels/client/ignite_client.h b/tensorflow/contrib/ignite/kernels/client/ignite_client.h new file mode 100644 index 0000000000000000000000000000000000000000..0da80769260d065c4ac6601c0e5cd7050b6b61cb --- /dev/null +++ b/tensorflow/contrib/ignite/kernels/client/ignite_client.h @@ -0,0 +1,84 @@ +/* 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_IGNITE_KERNELS_CLIENT_IGNITE_CLIENT_H_ +#define TENSORFLOW_CONTRIB_IGNITE_KERNELS_CLIENT_IGNITE_CLIENT_H_ + +#include "tensorflow/contrib/ignite/kernels/client/ignite_byte_swapper.h" +#include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/lib/core/status.h" + +namespace tensorflow { + +class Client { + public: + Client(bool big_endian) : byte_swapper_(ByteSwapper(big_endian)) {} + virtual Status Connect() = 0; + virtual Status Disconnect() = 0; + virtual bool IsConnected() = 0; + virtual int GetSocketDescriptor() = 0; + virtual Status ReadData(uint8_t *buf, const int32_t length) = 0; + virtual Status WriteData(const uint8_t *buf, const int32_t length) = 0; + + Status ReadByte(uint8_t *data) { return ReadData(data, 1); } + + Status ReadShort(int16_t *data) { + TF_RETURN_IF_ERROR(ReadData((uint8_t *)data, 2)); + byte_swapper_.SwapIfRequiredInt16(data); + + return Status::OK(); + } + + Status ReadInt(int32_t *data) { + TF_RETURN_IF_ERROR(ReadData((uint8_t *)data, 4)); + byte_swapper_.SwapIfRequiredInt32(data); + + return Status::OK(); + } + + Status ReadLong(int64_t *data) { + TF_RETURN_IF_ERROR(ReadData((uint8_t *)data, 8)); + byte_swapper_.SwapIfRequiredInt64(data); + + return Status::OK(); + } + + Status WriteByte(const uint8_t data) { return WriteData(&data, 1); } + + Status WriteShort(const int16_t data) { + int16_t tmp = data; + byte_swapper_.SwapIfRequiredInt16(&tmp); + return WriteData((uint8_t *)&tmp, 2); + } + + Status WriteInt(const int32_t data) { + int32_t tmp = data; + byte_swapper_.SwapIfRequiredInt32(&tmp); + return WriteData((uint8_t *)&tmp, 4); + } + + Status WriteLong(const int64_t data) { + int64_t tmp = data; + byte_swapper_.SwapIfRequiredInt64(&tmp); + return WriteData((uint8_t *)&tmp, 8); + } + + private: + const ByteSwapper byte_swapper_; +}; + +} // namespace tensorflow + +#endif // TENSORFLOW_CONTRIB_IGNITE_KERNELS_CLIENT_IGNITE_CLIENT_H_ diff --git a/tensorflow/contrib/ignite/kernels/ignite_plain_client.h b/tensorflow/contrib/ignite/kernels/client/ignite_plain_client.h similarity index 80% rename from tensorflow/contrib/ignite/kernels/ignite_plain_client.h rename to tensorflow/contrib/ignite/kernels/client/ignite_plain_client.h index 75424c19ee4b7df5378aa23cb41db1752e8d0651..546583246042855d179ebbb18b7dca711063b3f4 100644 --- a/tensorflow/contrib/ignite/kernels/ignite_plain_client.h +++ b/tensorflow/contrib/ignite/kernels/client/ignite_plain_client.h @@ -13,10 +13,10 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGNITE_PLAIN_CLIENT_H_ -#define TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGNITE_PLAIN_CLIENT_H_ +#ifndef TENSORFLOW_CONTRIB_IGNITE_KERNELS_CLIENT_IGNITE_PLAIN_CLIENT_H_ +#define TENSORFLOW_CONTRIB_IGNITE_KERNELS_CLIENT_IGNITE_PLAIN_CLIENT_H_ -#include "tensorflow/contrib/ignite/kernels/ignite_client.h" +#include "tensorflow/contrib/ignite/kernels/client/ignite_client.h" namespace tensorflow { @@ -40,4 +40,4 @@ class PlainClient : public Client { } // namespace tensorflow -#endif // TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGNITE_PLAIN_CLIENT_H_ +#endif // TENSORFLOW_CONTRIB_IGNITE_KERNELS_CLIENT_IGNITE_PLAIN_CLIENT_H_ diff --git a/tensorflow/contrib/ignite/kernels/ignite_plain_client_unix.cc b/tensorflow/contrib/ignite/kernels/client/ignite_plain_client_unix.cc similarity index 97% rename from tensorflow/contrib/ignite/kernels/ignite_plain_client_unix.cc rename to tensorflow/contrib/ignite/kernels/client/ignite_plain_client_unix.cc index cf672942c61e1239332711db12e62088737c4f41..54efb5b61761708a28dd031b8321ffba9a53ffa9 100644 --- a/tensorflow/contrib/ignite/kernels/ignite_plain_client_unix.cc +++ b/tensorflow/contrib/ignite/kernels/client/ignite_plain_client_unix.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/contrib/ignite/kernels/ignite_plain_client.h" +#include "tensorflow/contrib/ignite/kernels/client/ignite_plain_client.h" #include #include diff --git a/tensorflow/contrib/ignite/kernels/ignite_plain_client_windows.cc b/tensorflow/contrib/ignite/kernels/client/ignite_plain_client_windows.cc similarity index 98% rename from tensorflow/contrib/ignite/kernels/ignite_plain_client_windows.cc rename to tensorflow/contrib/ignite/kernels/client/ignite_plain_client_windows.cc index dad5aace5fabe1df58bb9579bf578f4c35324315..a99a3ada558e51c13ed47eb72911eb5862e71a60 100644 --- a/tensorflow/contrib/ignite/kernels/ignite_plain_client_windows.cc +++ b/tensorflow/contrib/ignite/kernels/client/ignite_plain_client_windows.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/contrib/ignite/kernels/ignite_plain_client.h" +#include "tensorflow/contrib/ignite/kernels/client/ignite_plain_client.h" #define WIN32_LEAN_AND_MEAN #include diff --git a/tensorflow/contrib/ignite/kernels/ignite_ssl_wrapper.cc b/tensorflow/contrib/ignite/kernels/client/ignite_ssl_wrapper.cc similarity index 98% rename from tensorflow/contrib/ignite/kernels/ignite_ssl_wrapper.cc rename to tensorflow/contrib/ignite/kernels/client/ignite_ssl_wrapper.cc index ceb479b0846574a35d86002ebb9c3e8e1d3687ac..8f09c24a3bedda524264f30282a0ad019d515540 100644 --- a/tensorflow/contrib/ignite/kernels/ignite_ssl_wrapper.cc +++ b/tensorflow/contrib/ignite/kernels/client/ignite_ssl_wrapper.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/contrib/ignite/kernels/ignite_ssl_wrapper.h" +#include "tensorflow/contrib/ignite/kernels/client/ignite_ssl_wrapper.h" #include #include diff --git a/tensorflow/contrib/ignite/kernels/ignite_ssl_wrapper.h b/tensorflow/contrib/ignite/kernels/client/ignite_ssl_wrapper.h similarity index 82% rename from tensorflow/contrib/ignite/kernels/ignite_ssl_wrapper.h rename to tensorflow/contrib/ignite/kernels/client/ignite_ssl_wrapper.h index 0406644bbaab3de816540ce85e84b489ea9fff12..543e03d1efc3ff186c9db399af18f7aa8ad2c450 100644 --- a/tensorflow/contrib/ignite/kernels/ignite_ssl_wrapper.h +++ b/tensorflow/contrib/ignite/kernels/client/ignite_ssl_wrapper.h @@ -13,10 +13,10 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGNITE_SSL_WRAPPER_H_ -#define TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGNITE_SSL_WRAPPER_H_ +#ifndef TENSORFLOW_CONTRIB_IGNITE_KERNELS_CLIENT_IGNITE_SSL_WRAPPER_H_ +#define TENSORFLOW_CONTRIB_IGNITE_KERNELS_CLIENT_IGNITE_SSL_WRAPPER_H_ -#include "tensorflow/contrib/ignite/kernels/ignite_client.h" +#include "tensorflow/contrib/ignite/kernels/client/ignite_client.h" #include @@ -48,4 +48,4 @@ class SslWrapper : public Client { } // namespace tensorflow -#endif // TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGNITE_SSL_WRAPPER_H_ +#endif // TENSORFLOW_CONTRIB_IGNITE_KERNELS_CLIENT_IGNITE_SSL_WRAPPER_H_ diff --git a/tensorflow/contrib/ignite/kernels/ignite_binary_object_parser.cc b/tensorflow/contrib/ignite/kernels/dataset/ignite_binary_object_parser.cc similarity index 99% rename from tensorflow/contrib/ignite/kernels/ignite_binary_object_parser.cc rename to tensorflow/contrib/ignite/kernels/dataset/ignite_binary_object_parser.cc index 2c8a7d44b07b43f788bcbc0850b5162cc14dd951..4218ec05f2c3486dd91e2188b674e01d6aadaa2b 100644 --- a/tensorflow/contrib/ignite/kernels/ignite_binary_object_parser.cc +++ b/tensorflow/contrib/ignite/kernels/dataset/ignite_binary_object_parser.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/contrib/ignite/kernels/ignite_binary_object_parser.h" +#include "tensorflow/contrib/ignite/kernels/dataset/ignite_binary_object_parser.h" #include "tensorflow/core/framework/types.h" #include "tensorflow/core/lib/core/errors.h" diff --git a/tensorflow/contrib/ignite/kernels/ignite_binary_object_parser.h b/tensorflow/contrib/ignite/kernels/dataset/ignite_binary_object_parser.h similarity index 87% rename from tensorflow/contrib/ignite/kernels/ignite_binary_object_parser.h rename to tensorflow/contrib/ignite/kernels/dataset/ignite_binary_object_parser.h index eb1f856643a790de6acaa82d4b8ad894fd364376..3e8a1a19623fab3e027db16228e0228e8ec4989a 100644 --- a/tensorflow/contrib/ignite/kernels/ignite_binary_object_parser.h +++ b/tensorflow/contrib/ignite/kernels/dataset/ignite_binary_object_parser.h @@ -13,11 +13,11 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGNITE_BINARY_OBJECT_PARSER_H_ -#define TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGNITE_BINARY_OBJECT_PARSER_H_ +#ifndef TENSORFLOW_CONTRIB_IGNITE_KERNELS_DATASET_IGNITE_BINARY_OBJECT_PARSER_H_ +#define TENSORFLOW_CONTRIB_IGNITE_KERNELS_DATASET_IGNITE_BINARY_OBJECT_PARSER_H_ #include -#include "tensorflow/contrib/ignite/kernels/ignite_byte_swapper.h" +#include "tensorflow/contrib/ignite/kernels/client/ignite_byte_swapper.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/lib/core/status.h" @@ -78,4 +78,4 @@ enum ObjectType { } // namespace tensorflow -#endif // TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGNITE_BINARY_OBJECT_PARSER_H_ +#endif // TENSORFLOW_CONTRIB_IGNITE_KERNELS_DATASET_IGNITE_BINARY_OBJECT_PARSER_H_ diff --git a/tensorflow/contrib/ignite/kernels/ignite_dataset.cc b/tensorflow/contrib/ignite/kernels/dataset/ignite_dataset.cc similarity index 97% rename from tensorflow/contrib/ignite/kernels/ignite_dataset.cc rename to tensorflow/contrib/ignite/kernels/dataset/ignite_dataset.cc index c4a7d3c513a796c9d95b371bedc609fd75188817..ace96e7b09fcf314757367baed66f622b294e43c 100644 --- a/tensorflow/contrib/ignite/kernels/ignite_dataset.cc +++ b/tensorflow/contrib/ignite/kernels/dataset/ignite_dataset.cc @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/contrib/ignite/kernels/ignite_dataset_iterator.h" +#include "tensorflow/contrib/ignite/kernels/dataset/ignite_dataset_iterator.h" #include "tensorflow/core/platform/logging.h" namespace tensorflow { diff --git a/tensorflow/contrib/ignite/kernels/ignite_dataset.h b/tensorflow/contrib/ignite/kernels/dataset/ignite_dataset.h similarity index 91% rename from tensorflow/contrib/ignite/kernels/ignite_dataset.h rename to tensorflow/contrib/ignite/kernels/dataset/ignite_dataset.h index 66bfdf2e2a168e59cd2fec8e2ac5b8fd482d5c15..db3bafb11f2a0047c22ece6d2bc1722afaa5ffdf 100644 --- a/tensorflow/contrib/ignite/kernels/ignite_dataset.h +++ b/tensorflow/contrib/ignite/kernels/dataset/ignite_dataset.h @@ -13,8 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGNITE_DATASET_H_ -#define TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGNITE_DATASET_H_ +#ifndef TENSORFLOW_CONTRIB_IGNITE_KERNELS_DATASET_IGNITE_DATASET_H_ +#define TENSORFLOW_CONTRIB_IGNITE_KERNELS_DATASET_IGNITE_DATASET_H_ #include "tensorflow/core/framework/dataset.h" @@ -60,4 +60,4 @@ class IgniteDataset : public DatasetBase { } // namespace tensorflow -#endif // TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGNITE_DATASET_H_ +#endif // TENSORFLOW_CONTRIB_IGNITE_KERNELS_DATASET_IGNITE_DATASET_H_ diff --git a/tensorflow/contrib/ignite/kernels/ignite_dataset_iterator.cc b/tensorflow/contrib/ignite/kernels/dataset/ignite_dataset_iterator.cc similarity index 98% rename from tensorflow/contrib/ignite/kernels/ignite_dataset_iterator.cc rename to tensorflow/contrib/ignite/kernels/dataset/ignite_dataset_iterator.cc index 5da9127aa6a3a4bc16347e6890cc1ba44406c0d5..ce8972f1e7fd59235556cb9514011f0b836077de 100644 --- a/tensorflow/contrib/ignite/kernels/ignite_dataset_iterator.cc +++ b/tensorflow/contrib/ignite/kernels/dataset/ignite_dataset_iterator.cc @@ -13,10 +13,10 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/contrib/ignite/kernels/ignite_dataset_iterator.h" +#include "tensorflow/contrib/ignite/kernels/dataset/ignite_dataset_iterator.h" -#include "tensorflow/contrib/ignite/kernels/ignite_plain_client.h" -#include "tensorflow/contrib/ignite/kernels/ignite_ssl_wrapper.h" +#include "tensorflow/contrib/ignite/kernels/client/ignite_plain_client.h" +#include "tensorflow/contrib/ignite/kernels/client/ignite_ssl_wrapper.h" #include "tensorflow/core/lib/gtl/cleanup.h" #include "tensorflow/core/platform/logging.h" diff --git a/tensorflow/contrib/ignite/kernels/ignite_dataset_iterator.h b/tensorflow/contrib/ignite/kernels/dataset/ignite_dataset_iterator.h similarity index 87% rename from tensorflow/contrib/ignite/kernels/ignite_dataset_iterator.h rename to tensorflow/contrib/ignite/kernels/dataset/ignite_dataset_iterator.h index c499e2c9ccfac5c15db08c8fd8b26c37aa0404f3..5868c2cb67f9d5c91654db8cf4bb4bbc072fc1ac 100644 --- a/tensorflow/contrib/ignite/kernels/ignite_dataset_iterator.h +++ b/tensorflow/contrib/ignite/kernels/dataset/ignite_dataset_iterator.h @@ -13,12 +13,12 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#ifndef TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGNITE_DATASET_ITERATOR_H_ -#define TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGNITE_DATASET_ITERATOR_H_ +#ifndef TENSORFLOW_CONTRIB_IGNITE_KERNELS_DATASET_IGNITE_DATASET_ITERATOR_H_ +#define TENSORFLOW_CONTRIB_IGNITE_KERNELS_DATASET_IGNITE_DATASET_ITERATOR_H_ -#include "tensorflow/contrib/ignite/kernels/ignite_binary_object_parser.h" -#include "tensorflow/contrib/ignite/kernels/ignite_client.h" -#include "tensorflow/contrib/ignite/kernels/ignite_dataset.h" +#include "tensorflow/contrib/ignite/kernels/client/ignite_client.h" +#include "tensorflow/contrib/ignite/kernels/dataset/ignite_binary_object_parser.h" +#include "tensorflow/contrib/ignite/kernels/dataset/ignite_dataset.h" #include "tensorflow/core/platform/mutex.h" namespace tensorflow { @@ -96,4 +96,4 @@ constexpr int32_t kMinResLength = 12; } // namespace tensorflow -#endif // TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGNITE_DATASET_ITERATOR_H_ +#endif // TENSORFLOW_CONTRIB_IGNITE_KERNELS_DATASET_IGNITE_DATASET_ITERATOR_H_ diff --git a/tensorflow/contrib/ignite/kernels/ignite_dataset_ops.cc b/tensorflow/contrib/ignite/kernels/dataset/ignite_dataset_ops.cc similarity index 97% rename from tensorflow/contrib/ignite/kernels/ignite_dataset_ops.cc rename to tensorflow/contrib/ignite/kernels/dataset/ignite_dataset_ops.cc index f75b1c5ff55ca9ee493148ff79c2edd4b15ac42a..f2108775e29b53765138dcd971bec89d7a10ce40 100644 --- a/tensorflow/contrib/ignite/kernels/ignite_dataset_ops.cc +++ b/tensorflow/contrib/ignite/kernels/dataset/ignite_dataset_ops.cc @@ -15,8 +15,8 @@ limitations under the License. #include -#include "tensorflow/contrib/ignite/kernels/ignite_binary_object_parser.h" -#include "tensorflow/contrib/ignite/kernels/ignite_dataset.h" +#include "tensorflow/contrib/ignite/kernels/dataset/ignite_binary_object_parser.h" +#include "tensorflow/contrib/ignite/kernels/dataset/ignite_dataset.h" #include "tensorflow/core/framework/dataset.h" #include "tensorflow/core/lib/strings/numbers.h" diff --git a/tensorflow/contrib/ignite/kernels/igfs/igfs.cc b/tensorflow/contrib/ignite/kernels/igfs/igfs.cc new file mode 100644 index 0000000000000000000000000000000000000000..ae2dbcc2cf5d0ae7e09a26a199dc0c3c80fe22c1 --- /dev/null +++ b/tensorflow/contrib/ignite/kernels/igfs/igfs.cc @@ -0,0 +1,331 @@ +/* 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/core/lib/io/path.h" +#include "tensorflow/core/platform/env.h" +#include "tensorflow/core/platform/file_system.h" +#include "tensorflow/core/platform/file_system_helper.h" + +#include "tensorflow/contrib/ignite/kernels/igfs/igfs.h" +#include "tensorflow/contrib/ignite/kernels/igfs/igfs_client.h" +#include "tensorflow/contrib/ignite/kernels/igfs/igfs_random_access_file.h" +#include "tensorflow/contrib/ignite/kernels/igfs/igfs_writable_file.h" + +namespace tensorflow { + +static string GetEnvOrElse(const string &env, string default_value) { + const char *env_c_str = env.c_str(); + return getenv(env_c_str) != nullptr ? getenv(env_c_str) : default_value; +} + +static string MakeRelative(const string &a, const string &b) { + string max = a; + string min = b; + bool first = b.size() > a.size(); + + if (first) { + max = b; + min = a; + } + + auto r = mismatch(min.begin(), min.end(), max.begin()); + return string((first ? r.first : r.second), first ? min.end() : max.end()); +} + +string IGFS::TranslateName(const string &name) const { + StringPiece scheme, namenode, path; + io::ParseURI(name, &scheme, &namenode, &path); + return string(path.data(), path.length()); +} + +IGFS::IGFS() + : host_(GetEnvOrElse("IGFS_HOST", "localhost")), + port_([] { + int port; + if (strings::safe_strto32(GetEnvOrElse("IGFS_PORT", "10500").c_str(), + &port)) { + return port; + } else { + LOG(WARNING) + << "IGFS_PORT environment variable had an invalid value: " + << getenv("IGFS_PORT") << "\nUsing default port 10500."; + return 10500; + } + }()), + fs_name_(GetEnvOrElse("IGFS_FS_NAME", "default_fs")) { + LOG(INFO) << "IGFS created [host=" << host_ << ", port=" << port_ + << ", fs_name=" << fs_name_ << "]"; +} + +IGFS::~IGFS() { + LOG(INFO) << "IGFS destroyed [host=" << host_ << ", port=" << port_ + << ", fs_name=" << fs_name_ << "]"; +} + +Status IGFS::NewRandomAccessFile(const string &file_name, + std::unique_ptr *result) { + std::unique_ptr client = CreateClient(); + string path = TranslateName(file_name); + + CtrlResponse handshake_response(true); + TF_RETURN_IF_ERROR(client->Handshake(&handshake_response)); + + CtrlResponse open_read_response(true); + TF_RETURN_IF_ERROR(client->OpenRead(&open_read_response, path)); + + int64 resource_id = open_read_response.res.stream_id; + result->reset(new IGFSRandomAccessFile(path, resource_id, std::move(client))); + + LOG(INFO) << "New random access file completed successfully [file_name=" + << file_name << "]"; + + return Status::OK(); +} + +Status IGFS::NewWritableFile(const string &file_name, + std::unique_ptr *result) { + std::unique_ptr client = CreateClient(); + string path = TranslateName(file_name); + + CtrlResponse handshake_response(true); + TF_RETURN_IF_ERROR(client->Handshake(&handshake_response)); + + CtrlResponse exists_response(false); + TF_RETURN_IF_ERROR(client->Exists(&exists_response, path)); + + if (exists_response.res.exists) { + CtrlResponse del_response(false); + TF_RETURN_IF_ERROR(client->Delete(&del_response, path, false)); + } + + CtrlResponse open_create_resp(false); + TF_RETURN_IF_ERROR(client->OpenCreate(&open_create_resp, path)); + + int64 resource_id = open_create_resp.res.stream_id; + result->reset(new IGFSWritableFile(path, resource_id, std::move(client))); + + LOG(INFO) << "New writable file completed successfully [file_name=" + << file_name << "]"; + + return Status::OK(); +} + +Status IGFS::NewAppendableFile(const string &file_name, + std::unique_ptr *result) { + std::unique_ptr client = CreateClient(); + + CtrlResponse handshake_response(true); + TF_RETURN_IF_ERROR(client->Handshake(&handshake_response)); + + CtrlResponse exists_response(false); + TF_RETURN_IF_ERROR(client->Exists(&exists_response, file_name)); + + if (exists_response.res.exists) { + CtrlResponse del_response(false); + TF_RETURN_IF_ERROR(client->Delete(&del_response, file_name, false)); + } + + CtrlResponse open_append_resp(false); + TF_RETURN_IF_ERROR(client->OpenAppend(&open_append_resp, file_name)); + + result->reset(new IGFSWritableFile(TranslateName(file_name), + open_append_resp.res.stream_id, + std::move(client))); + + LOG(INFO) << "New appendable file completed successfully [file_name=" + << file_name << "]"; + + return Status::OK(); +} + +Status IGFS::NewReadOnlyMemoryRegionFromFile( + const string &file_name, std::unique_ptr *result) { + return errors::Unimplemented("IGFS does not support ReadOnlyMemoryRegion"); +} + +Status IGFS::FileExists(const string &file_name) { + std::unique_ptr client = CreateClient(); + const string path = TranslateName(file_name); + + CtrlResponse handshake_response(true); + TF_RETURN_IF_ERROR(client->Handshake(&handshake_response)); + + CtrlResponse exists_response(false); + TF_RETURN_IF_ERROR(client->Exists(&exists_response, path)); + + if (!exists_response.res.exists) + return errors::NotFound("File ", path, " not found"); + + LOG(INFO) << "File exists completed successfully [file_name=" << file_name + << "]"; + + return Status::OK(); +} + +Status IGFS::GetChildren(const string &file_name, std::vector *result) { + std::unique_ptr client = CreateClient(); + string path = TranslateName(file_name); + path = path + "/"; + + CtrlResponse handshake_response(true); + TF_RETURN_IF_ERROR(client->Handshake(&handshake_response)); + + CtrlResponse list_paths_response(false); + TF_RETURN_IF_ERROR(client->ListPaths(&list_paths_response, path)); + + *result = std::vector(); + std::vector entries = list_paths_response.res.entries; + + for (IGFSPath &value : entries) + result->push_back(MakeRelative(value.path, path)); + + LOG(INFO) << "Get children completed successfully [file_name=" << file_name + << "]"; + + return Status::OK(); +} + +Status IGFS::GetMatchingPaths(const string &pattern, + std::vector *results) { + return internal::GetMatchingPaths(this, Env::Default(), pattern, results); +} + +Status IGFS::DeleteFile(const string &file_name) { + std::unique_ptr client = CreateClient(); + string path = TranslateName(file_name); + + CtrlResponse handshake_response(true); + TF_RETURN_IF_ERROR(client->Handshake(&handshake_response)); + + CtrlResponse del_response(false); + TF_RETURN_IF_ERROR(client->Delete(&del_response, path, false)); + + if (!del_response.res.exists) + return errors::NotFound("File ", path, " not found"); + + LOG(INFO) << "Delete file completed successfully [file_name=" << file_name + << "]"; + + return Status::OK(); +} + +Status IGFS::CreateDir(const string &file_name) { + std::unique_ptr client = CreateClient(); + const string path = TranslateName(file_name); + + CtrlResponse handshake_response(true); + TF_RETURN_IF_ERROR(client->Handshake(&handshake_response)); + + CtrlResponse mkdir_response(false); + TF_RETURN_IF_ERROR(client->MkDir(&mkdir_response, path)); + + if (!mkdir_response.res.successful) + return errors::Unknown("Can't create directory ", path); + + LOG(INFO) << "Create dir completed successful [file_name=" << file_name + << "]"; + + return Status::OK(); +} + +Status IGFS::DeleteDir(const string &file_name) { + std::unique_ptr client = CreateClient(); + string path = TranslateName(file_name); + + CtrlResponse handshake_response(true); + TF_RETURN_IF_ERROR(client->Handshake(&handshake_response)); + + CtrlResponse list_files_response(false); + TF_RETURN_IF_ERROR(client->ListFiles(&list_files_response, path)); + + if (!list_files_response.res.entries.empty()) { + return errors::FailedPrecondition("Can't delete a non-empty directory"); + } else { + CtrlResponse del_response(false); + TF_RETURN_IF_ERROR(client->Delete(&del_response, path, true)); + } + + LOG(INFO) << "Delete dir completed successful [file_name=" << file_name + << "]"; + + return Status::OK(); +} + +Status IGFS::GetFileSize(const string &file_name, uint64 *size) { + std::unique_ptr client = CreateClient(); + string path = TranslateName(file_name); + + CtrlResponse handshake_response(true); + TF_RETURN_IF_ERROR(client->Handshake(&handshake_response)); + + CtrlResponse info_response(false); + TF_RETURN_IF_ERROR(client->Info(&info_response, path)); + + *size = info_response.res.file_info.length; + + LOG(INFO) << "Get file size completed successful [file_name=" << file_name + << "]"; + + return Status::OK(); +} + +Status IGFS::RenameFile(const string &src, const string &dst) { + std::unique_ptr client = CreateClient(); + string src_path = TranslateName(src); + string dst_path = TranslateName(dst); + + if (FileExists(dst).ok()) DeleteFile(dst); + + CtrlResponse handshake_response(true); + TF_RETURN_IF_ERROR(client->Handshake(&handshake_response)); + + CtrlResponse rename_response(false); + TF_RETURN_IF_ERROR(client->Rename(&rename_response, src_path, dst_path)); + + if (!rename_response.res.successful) + return errors::NotFound("File ", src_path, " not found"); + + LOG(INFO) << "Rename file completed successful [src=" << src + << ", dst=" << dst << "]"; + + return Status::OK(); +} + +Status IGFS::Stat(const string &file_name, FileStatistics *stats) { + std::unique_ptr client = CreateClient(); + string path = TranslateName(file_name); + + CtrlResponse handshake_response(true); + TF_RETURN_IF_ERROR(client->Handshake(&handshake_response)); + + CtrlResponse info_response(false); + TF_RETURN_IF_ERROR(client->Info(&info_response, path)); + + IGFSFile info = info_response.res.file_info; + + *stats = FileStatistics(info.length, info.modification_time * 1000000, + (info.flags & 0x1) != 0); + + LOG(INFO) << "Stat completed successful [file_name=" << file_name << "]"; + + return Status::OK(); +} + +std::unique_ptr IGFS::CreateClient() const { + return std::unique_ptr( + new IGFSClient(host_, port_, fs_name_, "")); +} + +} // namespace tensorflow diff --git a/tensorflow/contrib/ignite/kernels/igfs/igfs.h b/tensorflow/contrib/ignite/kernels/igfs/igfs.h new file mode 100644 index 0000000000000000000000000000000000000000..4c347e937f75e8eea108811e6a3189412e22a982 --- /dev/null +++ b/tensorflow/contrib/ignite/kernels/igfs/igfs.h @@ -0,0 +1,60 @@ +/* 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_IGNITE_KERNELS_IGFS_IGFS_H_ +#define TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGFS_IGFS_H_ + +#include "tensorflow/contrib/ignite/kernels/igfs/igfs_client.h" +#include "tensorflow/core/platform/file_system.h" + +namespace tensorflow { + +class IGFS : public FileSystem { + public: + IGFS(); + ~IGFS(); + Status NewRandomAccessFile( + const string& file_name, + std::unique_ptr* result) override; + Status NewWritableFile(const string& fname, + std::unique_ptr* result) override; + Status NewAppendableFile(const string& fname, + std::unique_ptr* result) override; + Status NewReadOnlyMemoryRegionFromFile( + const string& fname, + std::unique_ptr* result) override; + Status FileExists(const string& fname) override; + Status GetChildren(const string& dir, std::vector* result) override; + Status GetMatchingPaths(const string& pattern, + std::vector* results) override; + Status DeleteFile(const string& fname) override; + Status CreateDir(const string& name) override; + Status DeleteDir(const string& name) override; + Status GetFileSize(const string& fname, uint64* size) override; + Status RenameFile(const string& src, const string& target) override; + Status Stat(const string& fname, FileStatistics* stat) override; + string TranslateName(const string& name) const override; + + private: + std::unique_ptr CreateClient() const; + + const string host_; + const int port_; + const string fs_name_; +}; + +} // namespace tensorflow + +#endif // TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGFS_IGFS_H_ diff --git a/tensorflow/contrib/ignite/kernels/igfs/igfs_client.cc b/tensorflow/contrib/ignite/kernels/igfs/igfs_client.cc new file mode 100644 index 0000000000000000000000000000000000000000..3f97c34fdd8b026a04506fd0ef9f3cc74129a9da --- /dev/null +++ b/tensorflow/contrib/ignite/kernels/igfs/igfs_client.cc @@ -0,0 +1,43 @@ +/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "tensorflow/contrib/ignite/kernels/igfs/igfs_client.h" + +namespace tensorflow { + +IGFSClient::IGFSClient(const string &host, int port, const string &fs_name, + const string &user_name) + : fs_name_(fs_name), + user_name_(user_name), + client_(ExtendedTCPClient(host, port, true)) { + client_.Connect(); +} + +IGFSClient::~IGFSClient() { client_.Disconnect(); } + +Status IGFSClient::SendRequestGetResponse(const Request &request, + Response *response) { + TF_RETURN_IF_ERROR(request.Write(&client_)); + client_.reset(); + + if (response != nullptr) { + TF_RETURN_IF_ERROR(response->Read(&client_)); + client_.reset(); + } + + return Status::OK(); +} + +} // namespace tensorflow diff --git a/tensorflow/contrib/ignite/kernels/igfs/igfs_client.h b/tensorflow/contrib/ignite/kernels/igfs/igfs_client.h new file mode 100644 index 0000000000000000000000000000000000000000..bbec7b000779be8772e850a556affffa1b3b6803 --- /dev/null +++ b/tensorflow/contrib/ignite/kernels/igfs/igfs_client.h @@ -0,0 +1,102 @@ +/* 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_IGNITE_KERNELS_IGFS_IGFS_CLIENT_H_ +#define TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGFS_IGFS_CLIENT_H_ + +#include "tensorflow/contrib/ignite/kernels/igfs/igfs_messages.h" + +namespace tensorflow { + +class IGFSClient { + public: + IGFSClient(const string &host, int port, const string &fs_name, + const string &user_name); + ~IGFSClient(); + + Status Handshake(CtrlResponse *res) { + return SendRequestGetResponse(HandshakeRequest(fs_name_, {}), res); + } + + Status ListFiles(CtrlResponse *res, const string &path) { + return SendRequestGetResponse(ListFilesRequest(user_name_, path), res); + } + + Status ListPaths(CtrlResponse *res, const string &path) { + return SendRequestGetResponse(ListPathsRequest(user_name_, path), res); + } + + Status Info(CtrlResponse *res, const string &path) { + return SendRequestGetResponse(InfoRequest(user_name_, path), res); + } + + Status OpenCreate(CtrlResponse *res, const string &path) { + return SendRequestGetResponse(OpenCreateRequest(user_name_, path), res); + } + + Status OpenAppend(CtrlResponse *res, const string &path) { + return SendRequestGetResponse(OpenAppendRequest(user_name_, path), res); + } + + Status OpenRead(CtrlResponse *res, const string &path) { + return SendRequestGetResponse(OpenReadRequest(user_name_, path), res); + } + + Status Exists(CtrlResponse *res, const string &path) { + return SendRequestGetResponse(ExistsRequest(user_name_, path), res); + } + + Status MkDir(CtrlResponse *res, const string &path) { + return SendRequestGetResponse(MakeDirectoriesRequest(user_name_, path), + res); + } + + Status Delete(CtrlResponse *res, const string &path, + bool recursive) { + return SendRequestGetResponse(DeleteRequest(user_name_, path, recursive), + res); + } + + Status WriteBlock(int64_t stream_id, const uint8_t *data, int32_t len) { + return SendRequestGetResponse(WriteBlockRequest(stream_id, data, len), + nullptr); + } + + Status ReadBlock(ReadBlockCtrlResponse *res, int64_t stream_id, int64_t pos, + int32_t length) { + return SendRequestGetResponse(ReadBlockRequest(stream_id, pos, length), + res); + } + + Status Close(CtrlResponse *res, int64_t stream_id) { + return SendRequestGetResponse(CloseRequest(stream_id), res); + } + + Status Rename(CtrlResponse *res, const string &source, + const string &dest) { + return SendRequestGetResponse(RenameRequest(user_name_, source, dest), res); + } + + private: + Status SendRequestGetResponse(const Request &request, Response *response); + + const string fs_name_; + const string user_name_; + ExtendedTCPClient client_; +}; + +} // namespace tensorflow + +#endif // TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGFS_IGFS_CLIENT_H_ diff --git a/tensorflow/contrib/ignite/kernels/igfs/igfs_extended_tcp_client.cc b/tensorflow/contrib/ignite/kernels/igfs/igfs_extended_tcp_client.cc new file mode 100644 index 0000000000000000000000000000000000000000..ea63436546d8b244b921206f9577c91b6578a775 --- /dev/null +++ b/tensorflow/contrib/ignite/kernels/igfs/igfs_extended_tcp_client.cc @@ -0,0 +1,144 @@ +/* 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/ignite/kernels/igfs/igfs_extended_tcp_client.h" + +namespace tensorflow { + +ExtendedTCPClient::ExtendedTCPClient(const string &host, int port, + bool big_endian) + : PlainClient(host, port, big_endian), pos_(0) {} + +Status ExtendedTCPClient::ReadData(uint8_t *buf, const int32_t length) { + TF_RETURN_IF_ERROR(PlainClient::ReadData(buf, length)); + pos_ += length; + + return Status::OK(); +} + +Status ExtendedTCPClient::WriteData(const uint8_t *buf, const int32_t length) { + TF_RETURN_IF_ERROR(PlainClient::WriteData(buf, length)); + pos_ += length; + + return Status::OK(); +} + +Status ExtendedTCPClient::Ignore(int n) { + uint8_t buf[n]; + return ReadData(buf, n); +} + +Status ExtendedTCPClient::SkipToPos(int target_pos) { + return Ignore(std::max(0, target_pos - pos_)); +} + +Status ExtendedTCPClient::ReadBool(bool *res) { + uint8_t buf = 0; + TF_RETURN_IF_ERROR(ReadData(&buf, 1)); + *res = buf != 0; + + return Status::OK(); +} + +Status ExtendedTCPClient::ReadNullableString(string *res) { + bool is_empty = false; + TF_RETURN_IF_ERROR(ReadBool(&is_empty)); + + if (!is_empty) { + TF_RETURN_IF_ERROR(ReadString(res)); + } + + return Status::OK(); +} + +Status ExtendedTCPClient::ReadString(string *res) { + int16_t length; + TF_RETURN_IF_ERROR(ReadShort(&length)); + + uint8_t *buf = new uint8_t[length]; + Status status = ReadData(buf, length); + + if (status.ok()) res->assign(reinterpret_cast(buf), length); + + delete[] buf; + return status; +} + +Status ExtendedTCPClient::ReadStringMap(std::map *res) { + int size; + TF_RETURN_IF_ERROR(ReadInt(&size)); + + for (int i = 0; i < size; i++) { + string key; + string val; + TF_RETURN_IF_ERROR(ReadString(&key)); + TF_RETURN_IF_ERROR(ReadString(&val)); + + res->insert(std::pair(std::move(key), std::move(val))); + } + + return Status::OK(); +} + +Status ExtendedTCPClient::WriteSize(std::map::size_type s) { + return WriteInt(s); +} + +Status ExtendedTCPClient::FillWithZerosUntil(int n) { + int to_skip = std::max(0, n - pos_); + + for (int i = 0; i < to_skip; i++) { + TF_RETURN_IF_ERROR(WriteByte(0)); + } + + return Status::OK(); +} + +Status ExtendedTCPClient::WriteBool(bool val) { + return WriteByte((char)(val ? 1 : 0)); +} + +Status ExtendedTCPClient::WriteString(string str) { + if (!str.empty()) { + TF_RETURN_IF_ERROR(WriteBool(false)); + size_t l = str.length(); + if (l > std::numeric_limits::max()) + return errors::InvalidArgument("String is too long"); + + TF_RETURN_IF_ERROR(WriteShort(l)); + TF_RETURN_IF_ERROR(WriteData(reinterpret_cast(str.c_str()), + str.length())); + } else { + TF_RETURN_IF_ERROR(WriteBool(true)); + } + + return Status::OK(); +} + +Status ExtendedTCPClient::WriteStringMap(std::map map) { + std::map::size_type size = map.size(); + TF_RETURN_IF_ERROR(WriteSize(size)); + + for (auto &x : map) { + TF_RETURN_IF_ERROR(WriteString(x.first)); + TF_RETURN_IF_ERROR(WriteString(x.second)); + } + + return Status::OK(); +} + +void ExtendedTCPClient::reset() { pos_ = 0; } + +} // namespace tensorflow diff --git a/tensorflow/contrib/ignite/kernels/igfs/igfs_extended_tcp_client.h b/tensorflow/contrib/ignite/kernels/igfs/igfs_extended_tcp_client.h new file mode 100644 index 0000000000000000000000000000000000000000..c5de342fd0c20cf5b01b647756797631b8a3f203 --- /dev/null +++ b/tensorflow/contrib/ignite/kernels/igfs/igfs_extended_tcp_client.h @@ -0,0 +1,47 @@ +/* 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_IGNITE_KERNELS_IGFS_IGFS_EXTENDED_TCP_CLIENT_H_ +#define TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGFS_IGFS_EXTENDED_TCP_CLIENT_H_ + +#include "tensorflow/contrib/ignite/kernels/client/ignite_plain_client.h" + +namespace tensorflow { + +class ExtendedTCPClient : public PlainClient { + public: + ExtendedTCPClient(const string &host, int port, bool big_endian); + Status ReadData(uint8_t *buf, const int32_t length) override; + Status WriteData(const uint8_t *buf, const int32_t length) override; + Status Ignore(int n); + Status SkipToPos(int target_pos); + Status ReadBool(bool *res); + Status ReadNullableString(string *res); + Status ReadString(string *res); + Status ReadStringMap(std::map *res); + Status WriteSize(std::map::size_type s); + Status FillWithZerosUntil(int n); + Status WriteBool(bool val); + Status WriteString(string str); + Status WriteStringMap(std::map map); + void reset(); + + private: + int pos_; +}; + +} // namespace tensorflow + +#endif // TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGFS_IGFS_EXTENDED_TCP_CLIENT_H_ diff --git a/tensorflow/contrib/ignite/kernels/igfs/igfs_messages.cc b/tensorflow/contrib/ignite/kernels/igfs/igfs_messages.cc new file mode 100644 index 0000000000000000000000000000000000000000..9c63f40f35fa53bc51c44f574df50ad0c79fba91 --- /dev/null +++ b/tensorflow/contrib/ignite/kernels/igfs/igfs_messages.cc @@ -0,0 +1,344 @@ +/* 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/ignite/kernels/igfs/igfs_messages.h" + +namespace tensorflow { + +Status IGFSPath::Read(ExtendedTCPClient *client) { + return client->ReadNullableString(&path); +} + +Status IGFSFile::Read(ExtendedTCPClient *client) { + int32_t block_size; + int64_t group_block_size; + std::map properties = {}; + int64_t access_time; + + bool has_path; + TF_RETURN_IF_ERROR(client->ReadBool(&has_path)); + if (has_path) { + IGFSPath path = {}; + TF_RETURN_IF_ERROR(path.Read(client)); + } + + TF_RETURN_IF_ERROR(client->ReadInt(&block_size)); + TF_RETURN_IF_ERROR(client->ReadLong(&group_block_size)); + TF_RETURN_IF_ERROR(client->ReadLong(&length)); + TF_RETURN_IF_ERROR(client->ReadStringMap(&properties)); + TF_RETURN_IF_ERROR(client->ReadLong(&access_time)); + TF_RETURN_IF_ERROR(client->ReadLong(&modification_time)); + TF_RETURN_IF_ERROR(client->ReadByte(&flags)); + + return Status::OK(); +} + +Request::Request(int32_t command_id) : command_id_(command_id) {} + +Status Request::Write(ExtendedTCPClient *client) const { + TF_RETURN_IF_ERROR(client->WriteByte(0)); + TF_RETURN_IF_ERROR(client->FillWithZerosUntil(8)); + TF_RETURN_IF_ERROR(client->WriteInt(command_id_)); + TF_RETURN_IF_ERROR(client->FillWithZerosUntil(24)); + + return Status::OK(); +} + +Status Response::Read(ExtendedTCPClient *client) { + TF_RETURN_IF_ERROR(client->Ignore(1)); + TF_RETURN_IF_ERROR(client->SkipToPos(8)); + TF_RETURN_IF_ERROR(client->ReadInt(&req_id)); + TF_RETURN_IF_ERROR(client->SkipToPos(24)); + TF_RETURN_IF_ERROR(client->ReadInt(&res_type)); + + bool has_error; + TF_RETURN_IF_ERROR(client->ReadBool(&has_error)); + + if (has_error) { + int32_t error_code; + string error_msg; + TF_RETURN_IF_ERROR(client->ReadString(&error_msg)); + TF_RETURN_IF_ERROR(client->ReadInt(&error_code)); + + return errors::Unknown("Error [code=", error_code, ", message=\"", + error_msg, "\"]"); + } + + TF_RETURN_IF_ERROR(client->SkipToPos(header_size_ + 5)); + TF_RETURN_IF_ERROR(client->ReadInt(&length)); + TF_RETURN_IF_ERROR(client->SkipToPos(header_size_ + response_header_size_)); + + return Status::OK(); +} + +PathCtrlRequest::PathCtrlRequest(int32_t command_id_, const string &user_name, + const string &path, + const string &destination_path, bool flag, + bool collocate, + const std::map &properties) + : Request(command_id_), + user_name_(user_name), + path_(path), + destination_path_(destination_path), + flag_(flag), + collocate_(collocate), + props_(properties) {} + +Status PathCtrlRequest::Write(ExtendedTCPClient *client) const { + TF_RETURN_IF_ERROR(Request::Write(client)); + + TF_RETURN_IF_ERROR(client->WriteString(user_name_)); + TF_RETURN_IF_ERROR(WritePath(client, path_)); + TF_RETURN_IF_ERROR(WritePath(client, destination_path_)); + TF_RETURN_IF_ERROR(client->WriteBool(flag_)); + TF_RETURN_IF_ERROR(client->WriteBool(collocate_)); + TF_RETURN_IF_ERROR(client->WriteStringMap(props_)); + + return Status::OK(); +} + +Status PathCtrlRequest::WritePath(ExtendedTCPClient *client, + const string &path) const { + TF_RETURN_IF_ERROR(client->WriteBool(!path.empty())); + if (!path.empty()) TF_RETURN_IF_ERROR(client->WriteString(path)); + + return Status::OK(); +} + +Status StreamCtrlRequest::Write(ExtendedTCPClient *client) const { + TF_RETURN_IF_ERROR(client->WriteByte(0)); + TF_RETURN_IF_ERROR(client->FillWithZerosUntil(8)); + TF_RETURN_IF_ERROR(client->WriteInt(command_id_)); + TF_RETURN_IF_ERROR(client->WriteLong(stream_id_)); + TF_RETURN_IF_ERROR(client->WriteInt(length_)); + + return Status::OK(); +} + +StreamCtrlRequest::StreamCtrlRequest(int32_t command_id_, int64_t stream_id, + int32_t length) + : Request(command_id_), stream_id_(stream_id), length_(length) {} + +DeleteRequest::DeleteRequest(const string &user_name, const string &path, + bool flag) + : PathCtrlRequest(DELETE_ID, user_name, path, {}, flag, true, {}) {} + +Status DeleteResponse::Read(ExtendedTCPClient *client) { + TF_RETURN_IF_ERROR(client->ReadBool(&exists)); + + return Status::OK(); +} + +ExistsRequest::ExistsRequest(const string &user_name, const string &path) + : PathCtrlRequest(EXISTS_ID, user_name, path, {}, false, true, {}) {} + +Status ExistsResponse::Read(ExtendedTCPClient *client) { + TF_RETURN_IF_ERROR(client->ReadBool(&exists)); + + return Status::OK(); +} + +HandshakeRequest::HandshakeRequest(const string &fs_name, const string &log_dir) + : Request(HANDSHAKE_ID), fs_name_(fs_name), log_dir_(log_dir) {} + +Status HandshakeRequest::Write(ExtendedTCPClient *client) const { + TF_RETURN_IF_ERROR(Request::Write(client)); + + TF_RETURN_IF_ERROR(client->WriteString(fs_name_)); + TF_RETURN_IF_ERROR(client->WriteString(log_dir_)); + + return Status::OK(); +} + +Status HandshakeResponse::Read(ExtendedTCPClient *client) { + int64_t block_size; + bool sampling; + + TF_RETURN_IF_ERROR(client->ReadNullableString(&fs_name)); + TF_RETURN_IF_ERROR(client->ReadLong(&block_size)); + + bool has_sampling_; + TF_RETURN_IF_ERROR(client->ReadBool(&has_sampling_)); + + if (has_sampling_) { + TF_RETURN_IF_ERROR(client->ReadBool(&sampling)); + } + + return Status::OK(); +} + +ListRequest::ListRequest(int32_t command_id_, const string &user_name, + const string &path) + : PathCtrlRequest(command_id_, user_name, path, {}, false, true, {}) {} + +ListFilesRequest::ListFilesRequest(const string &user_name, const string &path) + : ListRequest(LIST_FILES_ID, user_name, path) {} + +ListPathsRequest::ListPathsRequest(const string &user_name, const string &path) + : ListRequest(LIST_PATHS_ID, user_name, path) {} + +OpenCreateRequest::OpenCreateRequest(const string &user_name, + const string &path) + : PathCtrlRequest(OPEN_CREATE_ID, user_name, path, {}, false, true, {}) {} + +Status OpenCreateRequest::Write(ExtendedTCPClient *client) const { + TF_RETURN_IF_ERROR(PathCtrlRequest::Write(client)); + + TF_RETURN_IF_ERROR(client->WriteInt(replication_)); + TF_RETURN_IF_ERROR(client->WriteLong(blockSize_)); + + return Status::OK(); +} + +Status OpenCreateResponse::Read(ExtendedTCPClient *client) { + TF_RETURN_IF_ERROR(client->ReadLong(&stream_id)); + + return Status::OK(); +} + +OpenAppendRequest::OpenAppendRequest(const string &user_name, + const string &path) + : PathCtrlRequest(OPEN_APPEND_ID, user_name, path, {}, false, true, {}) {} + +Status OpenAppendRequest::Write(ExtendedTCPClient *client) const { + TF_RETURN_IF_ERROR(PathCtrlRequest::Write(client)); + + return Status::OK(); +} + +Status OpenAppendResponse::Read(ExtendedTCPClient *client) { + TF_RETURN_IF_ERROR(client->ReadLong(&stream_id)); + + return Status::OK(); +} + +OpenReadRequest::OpenReadRequest(const string &user_name, const string &path, + bool flag, + int32_t sequential_reads_before_prefetch) + : PathCtrlRequest(OPEN_READ_ID, user_name, path, {}, flag, true, {}), + sequential_reads_before_prefetch_(sequential_reads_before_prefetch) {} + +OpenReadRequest::OpenReadRequest(const string &user_name, const string &path) + : OpenReadRequest(user_name, path, false, 0) {} + +Status OpenReadRequest::Write(ExtendedTCPClient *client) const { + TF_RETURN_IF_ERROR(PathCtrlRequest::Write(client)); + + if (flag_) { + TF_RETURN_IF_ERROR(client->WriteInt(sequential_reads_before_prefetch_)); + } + + return Status::OK(); +} + +Status OpenReadResponse::Read(ExtendedTCPClient *client) { + TF_RETURN_IF_ERROR(client->ReadLong(&stream_id)); + TF_RETURN_IF_ERROR(client->ReadLong(&length)); + + return Status::OK(); +} + +InfoRequest::InfoRequest(const string &user_name, const string &path) + : PathCtrlRequest(INFO_ID, user_name, path, {}, false, true, {}) {} + +Status InfoResponse::Read(ExtendedTCPClient *client) { + file_info = IGFSFile(); + TF_RETURN_IF_ERROR(file_info.Read(client)); + + return Status::OK(); +} + +MakeDirectoriesRequest::MakeDirectoriesRequest(const string &user_name, + const string &path) + : PathCtrlRequest(MKDIR_ID, user_name, path, {}, false, true, {}) {} + +Status MakeDirectoriesResponse::Read(ExtendedTCPClient *client) { + TF_RETURN_IF_ERROR(client->ReadBool(&successful)); + + return Status::OK(); +} + +CloseRequest::CloseRequest(int64_t streamId) + : StreamCtrlRequest(CLOSE_ID, streamId, 0) {} + +Status CloseResponse::Read(ExtendedTCPClient *client) { + TF_RETURN_IF_ERROR(client->ReadBool(&successful)); + + return Status::OK(); +} + +ReadBlockRequest::ReadBlockRequest(int64_t stream_id, int64_t pos, + int32_t length) + : StreamCtrlRequest(READ_BLOCK_ID, stream_id, length), pos(pos) {} + +Status ReadBlockRequest::Write(ExtendedTCPClient *client) const { + TF_RETURN_IF_ERROR(StreamCtrlRequest::Write(client)); + + TF_RETURN_IF_ERROR(client->WriteLong(pos)); + + return Status::OK(); +} + +Status ReadBlockResponse::Read(ExtendedTCPClient *client, int32_t length, + uint8_t *dst) { + TF_RETURN_IF_ERROR(client->ReadData(dst, length)); + successfully_read = length; + + return Status::OK(); +} + +Status ReadBlockResponse::Read(ExtendedTCPClient *client) { + return Status::OK(); +} + +std::streamsize ReadBlockResponse::GetSuccessfullyRead() { + return successfully_read; +} + +ReadBlockCtrlResponse::ReadBlockCtrlResponse(uint8_t *dst) + : CtrlResponse(false), dst(dst) {} + +Status ReadBlockCtrlResponse::Read(ExtendedTCPClient *client) { + TF_RETURN_IF_ERROR(Response::Read(client)); + + res = ReadBlockResponse(); + TF_RETURN_IF_ERROR(res.Read(client, length, dst)); + + return Status::OK(); +} + +WriteBlockRequest::WriteBlockRequest(int64_t stream_id, const uint8_t *data, + int32_t length) + : StreamCtrlRequest(WRITE_BLOCK_ID, stream_id, length), data(data) {} + +Status WriteBlockRequest::Write(ExtendedTCPClient *client) const { + TF_RETURN_IF_ERROR(StreamCtrlRequest::Write(client)); + TF_RETURN_IF_ERROR(client->WriteData((uint8_t *)data, length_)); + + return Status::OK(); +} + +RenameRequest::RenameRequest(const string &user_name, const string &path, + const string &destination_path) + : PathCtrlRequest(RENAME_ID, user_name, path, destination_path, false, true, + {}) {} + +Status RenameResponse::Read(ExtendedTCPClient *client) { + TF_RETURN_IF_ERROR(client->ReadBool(&successful)); + + return Status::OK(); +} + +} // namespace tensorflow diff --git a/tensorflow/contrib/ignite/kernels/igfs/igfs_messages.h b/tensorflow/contrib/ignite/kernels/igfs/igfs_messages.h new file mode 100644 index 0000000000000000000000000000000000000000..44a2928a2b2b48849c7ba4454e0e7848c2217b3b --- /dev/null +++ b/tensorflow/contrib/ignite/kernels/igfs/igfs_messages.h @@ -0,0 +1,356 @@ +/* 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_IGNITE_KERNELS_IGFS_IGFS_MESSAGES_H_ +#define TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGFS_IGFS_MESSAGES_H_ + +#include "tensorflow/contrib/ignite/kernels/igfs/igfs_extended_tcp_client.h" + +namespace tensorflow { + +enum CommandId { + HANDSHAKE_ID = 0, + EXISTS_ID = 2, + INFO_ID = 3, + RENAME_ID = 6, + DELETE_ID = 7, + MKDIR_ID = 8, + LIST_PATHS_ID = 9, + LIST_FILES_ID = 10, + OPEN_READ_ID = 13, + OPEN_APPEND_ID = 14, + OPEN_CREATE_ID = 15, + CLOSE_ID = 16, + READ_BLOCK_ID = 17, + WRITE_BLOCK_ID = 18, +}; + +class IGFSPath { + public: + Status Read(ExtendedTCPClient *client); + + string path; +}; + +class IGFSFile { + public: + Status Read(ExtendedTCPClient *client); + + int64_t length; + int64_t modification_time; + uint8_t flags; +}; + +class Request { + public: + Request(int32_t command_id); + virtual Status Write(ExtendedTCPClient *client) const; + + protected: + const int32_t command_id_; +}; + +class Response { + public: + virtual Status Read(ExtendedTCPClient *client); + + int32_t res_type; + int32_t req_id; + int32_t length; + + protected: + static const int32_t header_size_ = 24; + static const int32_t response_header_size_ = 9; +}; + +class PathCtrlRequest : public Request { + public: + PathCtrlRequest(int32_t command_id, const string &user_name, + const string &path, const string &destination_path, bool flag, + bool collocate, const std::map &properties); + Status Write(ExtendedTCPClient *client) const override; + + protected: + Status WritePath(ExtendedTCPClient *client, const string &path) const; + + const string user_name_; + const string path_; + const string destination_path_; + const bool flag_; + const bool collocate_; + const std::map props_; +}; + +class StreamCtrlRequest : public Request { + public: + StreamCtrlRequest(int32_t command_id, int64_t stream_id, int32_t length); + Status Write(ExtendedTCPClient *client) const override; + + protected: + int64_t stream_id_; + int32_t length_; +}; + +template +class CtrlResponse : public Response { + public: + CtrlResponse(bool optional) : optional_(optional) {} + Status Read(ExtendedTCPClient *client) override { + TF_RETURN_IF_ERROR(Response::Read(client)); + + if (optional_) { + TF_RETURN_IF_ERROR(client->ReadBool(&has_content)); + + if (!has_content) return Status::OK(); + } + + res = R(); + has_content = true; + TF_RETURN_IF_ERROR(res.Read(client)); + + return Status::OK(); + } + + R res; + bool has_content; + + private: + bool optional_; +}; + +template +class ListResponse { + public: + Status Read(ExtendedTCPClient *client) { + int32_t len; + TF_RETURN_IF_ERROR(client->ReadInt(&len)); + + entries.clear(); + + for (int32_t i = 0; i < len; i++) { + T f = {}; + TF_RETURN_IF_ERROR(f.Read(client)); + entries.push_back(f); + } + + return Status::OK(); + } + + std::vector entries; +}; + +class DeleteRequest : public PathCtrlRequest { + public: + DeleteRequest(const string &user_name, const string &path, bool flag); +}; + +class DeleteResponse { + public: + Status Read(ExtendedTCPClient *client); + + bool exists; +}; + +class ExistsRequest : public PathCtrlRequest { + public: + explicit ExistsRequest(const string &user_name, const string &path); +}; + +class ExistsResponse { + public: + Status Read(ExtendedTCPClient *client); + + bool exists; +}; + +class HandshakeRequest : public Request { + public: + HandshakeRequest(const string &fs_name, const string &log_dir); + Status Write(ExtendedTCPClient *client) const override; + + private: + string fs_name_; + string log_dir_; +}; + +class HandshakeResponse { + public: + Status Read(ExtendedTCPClient *client); + + string fs_name; +}; + +class ListRequest : public PathCtrlRequest { + public: + explicit ListRequest(int32_t command_id, const string &user_name, + const string &path); +}; + +class ListFilesRequest : public ListRequest { + public: + ListFilesRequest(const string &user_name, const string &path); +}; + +class ListFilesResponse : public ListResponse {}; + +class ListPathsRequest : public ListRequest { + public: + ListPathsRequest(const string &user_name, const string &path); +}; + +class ListPathsResponse : public ListResponse {}; + +class OpenCreateRequest : public PathCtrlRequest { + public: + OpenCreateRequest(const string &user_name, const string &path); + Status Write(ExtendedTCPClient *client) const override; + + private: + int32_t replication_; + int64_t blockSize_; +}; + +class OpenCreateResponse { + public: + Status Read(ExtendedTCPClient *client); + + int64_t stream_id; +}; + +class OpenAppendRequest : public PathCtrlRequest { + public: + explicit OpenAppendRequest(const string &user_name, const string &path); + Status Write(ExtendedTCPClient *client) const override; +}; + +class OpenAppendResponse { + public: + Status Read(ExtendedTCPClient *client); + + int64_t stream_id; +}; + +class OpenReadRequest : public PathCtrlRequest { + public: + OpenReadRequest(const string &user_name, const string &path, bool flag, + int32_t seqReadsBeforePrefetch); + OpenReadRequest(const string &user_name, const string &path); + Status Write(ExtendedTCPClient *client) const override; + + protected: + /** Sequential reads before prefetch. */ + int32_t sequential_reads_before_prefetch_; +}; + +class OpenReadResponse { + public: + Status Read(ExtendedTCPClient *client); + + int64_t stream_id; + int64_t length; +}; + +class InfoRequest : public PathCtrlRequest { + public: + InfoRequest(const string &user_name, const string &path); +}; + +class InfoResponse { + public: + Status Read(ExtendedTCPClient *client); + + IGFSFile file_info; +}; + +class MakeDirectoriesRequest : public PathCtrlRequest { + public: + MakeDirectoriesRequest(const string &userName, const string &path); +}; + +class MakeDirectoriesResponse { + public: + Status Read(ExtendedTCPClient *client); + + bool successful; +}; + +/** Stream control requests. **/ + +class CloseRequest : public StreamCtrlRequest { + public: + explicit CloseRequest(int64_t stream_id); +}; + +class CloseResponse { + public: + Status Read(ExtendedTCPClient *client); + + bool successful; +}; + +class ReadBlockRequest : public StreamCtrlRequest { + public: + ReadBlockRequest(int64_t stream_id, int64_t pos, int32_t length); + Status Write(ExtendedTCPClient *client) const override; + + private: + int64_t pos; +}; + +class ReadBlockResponse { + public: + Status Read(ExtendedTCPClient *client, int32_t length, uint8_t *dst); + Status Read(ExtendedTCPClient *client); + std::streamsize GetSuccessfullyRead(); + + private: + int32_t length; + std::streamsize successfully_read; +}; + +class ReadBlockCtrlResponse : public CtrlResponse { + public: + ReadBlockCtrlResponse(uint8_t *dst); + Status Read(ExtendedTCPClient *client) override; + + private: + uint8_t *dst; +}; + +class WriteBlockRequest : public StreamCtrlRequest { + public: + WriteBlockRequest(int64_t stream_id, const uint8_t *data, int32_t length); + Status Write(ExtendedTCPClient *client) const override; + + private: + const uint8_t *data; +}; + +class RenameRequest : public PathCtrlRequest { + public: + RenameRequest(const string &user_name, const string &path, + const string &destination_path); +}; + +class RenameResponse { + public: + Status Read(ExtendedTCPClient *client); + + bool successful; +}; + +} // namespace tensorflow + +#endif // TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGFS_IGFS_MESSAGES_H_ diff --git a/tensorflow/contrib/ignite/kernels/igfs/igfs_random_access_file.cc b/tensorflow/contrib/ignite/kernels/igfs/igfs_random_access_file.cc new file mode 100644 index 0000000000000000000000000000000000000000..a4c898f14e6d298e65f563f4493a822172c40851 --- /dev/null +++ b/tensorflow/contrib/ignite/kernels/igfs/igfs_random_access_file.cc @@ -0,0 +1,48 @@ +/* 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/ignite/kernels/igfs/igfs_random_access_file.h" +#include "tensorflow/contrib/ignite/kernels/igfs/igfs_messages.h" + +namespace tensorflow { + +IGFSRandomAccessFile::IGFSRandomAccessFile(const string &file_name, + int64_t resource_id, + std::unique_ptr &&client) + : file_name_(file_name), + resource_id_(resource_id), + client_(std::move(client)) {} + +IGFSRandomAccessFile::~IGFSRandomAccessFile() { + CtrlResponse close_response = {false}; + Status status = client_->Close(&close_response, resource_id_); + + if (!status.ok()) LOG(ERROR) << status.ToString(); +} + +Status IGFSRandomAccessFile::Read(uint64 offset, size_t n, StringPiece *result, + char *scratch) const { + ReadBlockCtrlResponse response = ReadBlockCtrlResponse((uint8_t *)scratch); + TF_RETURN_IF_ERROR(client_->ReadBlock(&response, resource_id_, offset, n)); + + std::streamsize sz = response.res.GetSuccessfullyRead(); + if (sz == 0) return errors::OutOfRange("End of file"); + + *result = StringPiece(scratch, sz); + + return Status::OK(); +} + +} // namespace tensorflow diff --git a/tensorflow/contrib/ignite/kernels/igfs/igfs_random_access_file.h b/tensorflow/contrib/ignite/kernels/igfs/igfs_random_access_file.h new file mode 100644 index 0000000000000000000000000000000000000000..b21369ff8a3b19774bcc743f93a5ec4ae1c9b49a --- /dev/null +++ b/tensorflow/contrib/ignite/kernels/igfs/igfs_random_access_file.h @@ -0,0 +1,40 @@ +/* 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_IGNITE_KERNELS_IGFS_IGFS_RANDOM_ACCESS_FILE_H_ +#define TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGFS_IGFS_RANDOM_ACCESS_FILE_H_ + +#include "tensorflow/contrib/ignite/kernels/igfs/igfs_client.h" +#include "tensorflow/core/platform/file_system.h" + +namespace tensorflow { + +class IGFSRandomAccessFile : public RandomAccessFile { + public: + IGFSRandomAccessFile(const string &file_name, int64_t resource_id, + std::unique_ptr &&client); + ~IGFSRandomAccessFile() override; + Status Read(uint64 offset, size_t n, StringPiece *result, + char *scratch) const override; + + private: + const string file_name_; + const int64_t resource_id_; + std::unique_ptr client_; +}; + +} // namespace tensorflow + +#endif // TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGFS_IGFS_RANDOM_ACCESS_FILE_H_ diff --git a/tensorflow/contrib/ignite/kernels/igfs/igfs_writable_file.cc b/tensorflow/contrib/ignite/kernels/igfs/igfs_writable_file.cc new file mode 100644 index 0000000000000000000000000000000000000000..c15ecb7deeb0cf5a8a040e0d1e4b70c732729474 --- /dev/null +++ b/tensorflow/contrib/ignite/kernels/igfs/igfs_writable_file.cc @@ -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. +==============================================================================*/ + +#include "tensorflow/contrib/ignite/kernels/igfs/igfs_writable_file.h" +#include "tensorflow/contrib/ignite/kernels/igfs/igfs_messages.h" + +namespace tensorflow { + +IGFSWritableFile::IGFSWritableFile(const string &file_name, int64_t resource_id, + std::unique_ptr &&client) + : file_name_(file_name), + resource_id_(resource_id), + client_(std::move(client)) {} + +IGFSWritableFile::~IGFSWritableFile() { + if (resource_id_ >= 0) { + CtrlResponse close_response = {false}; + + Status status = client_->Close(&close_response, resource_id_); + if (!status.ok()) LOG(ERROR) << status.ToString(); + } +} + +Status IGFSWritableFile::Append(StringPiece data) { + return client_->WriteBlock(resource_id_, (uint8_t *)data.data(), data.size()); +} + +Status IGFSWritableFile::Close() { + int64_t resource_to_be_closed = resource_id_; + resource_id_ = -1; + + CtrlResponse close_response = {false}; + return client_->Close(&close_response, resource_to_be_closed); +} + +Status IGFSWritableFile::Flush() { return Sync(); } + +Status IGFSWritableFile::Sync() { + CtrlResponse close_response = {false}; + TF_RETURN_IF_ERROR(client_->Close(&close_response, resource_id_)); + + CtrlResponse open_append_resp(false); + TF_RETURN_IF_ERROR(client_->OpenAppend(&open_append_resp, file_name_)); + + resource_id_ = open_append_resp.res.stream_id; + + return Status::OK(); +} + +} // namespace tensorflow diff --git a/tensorflow/contrib/ignite/kernels/igfs/igfs_writable_file.h b/tensorflow/contrib/ignite/kernels/igfs/igfs_writable_file.h new file mode 100644 index 0000000000000000000000000000000000000000..b406db17e0e350e2cef610bb05c40f658e100140 --- /dev/null +++ b/tensorflow/contrib/ignite/kernels/igfs/igfs_writable_file.h @@ -0,0 +1,42 @@ +/* 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_IGNITE_KERNELS_IGFS_IGFS_WRITABLE_FILE_H_ +#define TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGFS_IGFS_WRITABLE_FILE_H_ + +#include "tensorflow/contrib/ignite/kernels/igfs/igfs_client.h" +#include "tensorflow/core/platform/file_system.h" + +namespace tensorflow { + +class IGFSWritableFile : public WritableFile { + public: + IGFSWritableFile(const string &file_name, int64_t resource_id, + std::unique_ptr &&client); + ~IGFSWritableFile() override; + Status Append(StringPiece data) override; + Status Close() override; + Status Flush() override; + Status Sync() override; + + private: + const string file_name_; + int64_t resource_id_; + std::unique_ptr client_; +}; + +} // namespace tensorflow + +#endif // TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGFS_IGFS_WRITABLE_FILE_H_ diff --git a/tensorflow/contrib/ignite/kernels/ignite_byte_swapper.h b/tensorflow/contrib/ignite/kernels/ignite_byte_swapper.h deleted file mode 100644 index 46df3e39dc4ec6dd4ef5730a184264eaa9fc5872..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/ignite/kernels/ignite_byte_swapper.h +++ /dev/null @@ -1,126 +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_IGNITE_KERNELS_IGNITE_BYTE_SWAPPER_H_ -#define TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGNITE_BYTE_SWAPPER_H_ - -#include -#include "tensorflow/core/platform/byte_order.h" - -namespace tensorflow { - -class ByteSwapper { - public: - ByteSwapper(bool big_endian) { swap_ = big_endian == port::kLittleEndian; } - - inline void SwapIfRequiredInt16(int16_t *x) const { - if (swap_) { - Swap16(x); - } - } - - inline void SwapIfRequiredUnsignedInt16(uint16_t *x) const { - if (swap_) { - Swap16(reinterpret_cast(x)); - } - } - - inline void SwapIfRequiredInt32(int32_t *x) const { - if (swap_) { - Swap32(x); - } - } - - inline void SwapIfRequiredFloat(float *x) const { - if (swap_) { - Swap32(reinterpret_cast(x)); - } - } - - inline void SwapIfRequiredInt64(int64_t *x) const { - if (swap_) { - Swap64(x); - } - } - - inline void SwapIfRequiredDouble(double *x) const { - if (swap_) { - Swap64(reinterpret_cast(x)); - } - } - - inline void SwapIfRequiredInt16Arr(int16_t *x, int32_t length) const { - if (swap_) { - for (int32_t i = 0; i < length; i++) Swap16(&x[i]); - } - } - - inline void SwapIfRequiredUnsignedInt16Arr(uint16_t *x, - int32_t length) const { - if (swap_) { - for (int32_t i = 0; i < length; i++) - Swap16(reinterpret_cast(&x[i])); - } - } - - inline void SwapIfRequiredInt32Arr(int32_t *x, int32_t length) const { - if (swap_) { - for (int32_t i = 0; i < length; i++) Swap32(&x[i]); - } - } - - inline void SwapIfRequiredFloatArr(float *x, int32_t length) const { - if (swap_) { - for (int32_t i = 0; i < length; i++) - Swap32(reinterpret_cast(&x[i])); - } - } - - inline void SwapIfRequiredInt64Arr(int64_t *x, int32_t length) const { - if (swap_) { - for (int32_t i = 0; i < length; i++) Swap64(&x[i]); - } - } - - inline void SwapIfRequiredDoubleArr(double *x, int32_t length) const { - if (swap_) { - for (int32_t i = 0; i < length; i++) - Swap64(reinterpret_cast(&x[i])); - } - } - - private: - inline void Swap16(int16_t *x) const { - *x = ((*x & 0xFF) << 8) | ((*x >> 8) & 0xFF); - } - - inline void Swap32(int32_t *x) const { - *x = ((*x & 0xFF) << 24) | (((*x >> 8) & 0xFF) << 16) | - (((*x >> 16) & 0xFF) << 8) | ((*x >> 24) & 0xFF); - } - - inline void Swap64(int64_t *x) const { - *x = ((*x & 0xFF) << 56) | (((*x >> 8) & 0xFF) << 48) | - (((*x >> 16) & 0xFF) << 40) | (((*x >> 24) & 0xFF) << 32) | - (((*x >> 32) & 0xFF) << 24) | (((*x >> 40) & 0xFF) << 16) | - (((*x >> 48) & 0xFF) << 8) | ((*x >> 56) & 0xFF); - } - - bool swap_; -}; - -} // namespace tensorflow - -#endif // TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGNITE_BYTE_SWAPPER_H_ diff --git a/tensorflow/contrib/ignite/kernels/ignite_client.h b/tensorflow/contrib/ignite/kernels/ignite_client.h deleted file mode 100644 index 459b50b48fd95ad105bccaca4076160e0ef152ee..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/ignite/kernels/ignite_client.h +++ /dev/null @@ -1,84 +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_IGNITE_KERNELS_IGNITE_CLIENT_H_ -#define TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGNITE_CLIENT_H_ - -#include "tensorflow/contrib/ignite/kernels/ignite_byte_swapper.h" -#include "tensorflow/core/lib/core/errors.h" -#include "tensorflow/core/lib/core/status.h" - -namespace tensorflow { - -class Client { - public: - Client(bool big_endian) : byte_swapper_(ByteSwapper(big_endian)) {} - virtual Status Connect() = 0; - virtual Status Disconnect() = 0; - virtual bool IsConnected() = 0; - virtual int GetSocketDescriptor() = 0; - virtual Status ReadData(uint8_t *buf, const int32_t length) = 0; - virtual Status WriteData(const uint8_t *buf, const int32_t length) = 0; - - inline Status ReadByte(uint8_t *data) { return ReadData(data, 1); } - - inline Status ReadShort(int16_t *data) { - TF_RETURN_IF_ERROR(ReadData((uint8_t *)data, 2)); - byte_swapper_.SwapIfRequiredInt16(data); - - return Status::OK(); - } - - inline Status ReadInt(int32_t *data) { - TF_RETURN_IF_ERROR(ReadData((uint8_t *)data, 4)); - byte_swapper_.SwapIfRequiredInt32(data); - - return Status::OK(); - } - - inline Status ReadLong(int64_t *data) { - TF_RETURN_IF_ERROR(ReadData((uint8_t *)data, 8)); - byte_swapper_.SwapIfRequiredInt64(data); - - return Status::OK(); - } - - inline Status WriteByte(const uint8_t data) { return WriteData(&data, 1); } - - inline Status WriteShort(const int16_t data) { - int16_t tmp = data; - byte_swapper_.SwapIfRequiredInt16(&tmp); - return WriteData((uint8_t *)&tmp, 2); - } - - inline Status WriteInt(const int32_t data) { - int32_t tmp = data; - byte_swapper_.SwapIfRequiredInt32(&tmp); - return WriteData((uint8_t *)&tmp, 4); - } - - inline Status WriteLong(const int64_t data) { - int64_t tmp = data; - byte_swapper_.SwapIfRequiredInt64(&tmp); - return WriteData((uint8_t *)&tmp, 8); - } - - private: - const ByteSwapper byte_swapper_; -}; - -} // namespace tensorflow - -#endif // TENSORFLOW_CONTRIB_IGNITE_KERNELS_IGNITE_CLIENT_H_ diff --git a/tensorflow/contrib/ignite/ops/igfs_ops.cc b/tensorflow/contrib/ignite/ops/igfs_ops.cc new file mode 100644 index 0000000000000000000000000000000000000000..473bddff08b339d3b76a33d40fe34486acdbe151 --- /dev/null +++ b/tensorflow/contrib/ignite/ops/igfs_ops.cc @@ -0,0 +1,24 @@ +/* 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/core/platform/env.h" + +#include "tensorflow/contrib/ignite/kernels/igfs/igfs.h" + +namespace tensorflow { + +REGISTER_FILE_SYSTEM("igfs", IGFS); + +} // namespace tensorflow diff --git a/tensorflow/contrib/ignite/python/ops/igfs_op_loader.py b/tensorflow/contrib/ignite/python/ops/igfs_op_loader.py new file mode 100644 index 0000000000000000000000000000000000000000..8e1d6707d6400a7cd84016150d20973809aca20e --- /dev/null +++ b/tensorflow/contrib/ignite/python/ops/igfs_op_loader.py @@ -0,0 +1,24 @@ +# 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. +# ============================================================================== +"""Python helper for loading IGFS ops and kernels.""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.util import loader +from tensorflow.python.platform import resource_loader + +_dataset_ops = loader.load_op_library( + resource_loader.get_path_to_datafile("../../_ignite_ops.so")) diff --git a/tensorflow/contrib/ignite/python/ops/igfs_ops.py b/tensorflow/contrib/ignite/python/ops/igfs_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..12b973b707730f6ba5b057b74a46b27d8f973ede --- /dev/null +++ b/tensorflow/contrib/ignite/python/ops/igfs_ops.py @@ -0,0 +1,40 @@ +# 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. +# ============================================================================== +"""Ignite File System for checkpointing and communication with TensorBoard. + +Apache Ignite is a memory-centric distributed database, caching, and +processing platform for transactional, analytical, and streaming workloads, +delivering in-memory speeds at petabyte scale. In addition to database +functionality Apache Ignite provides a distributed file 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. +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import os + +from tensorflow.contrib.ignite.python.ops import ignite_op_loader # pylint: disable=unused-import +from tensorflow.python.framework import load_library +from tensorflow.python.platform import resource_loader + +file_system_library = os.path.join(resource_loader.get_data_files_path(), + "../../_ignite_ops.so") +load_library.load_file_system_library(file_system_library) diff --git a/tensorflow/contrib/ignite/python/ops/ignite_dataset_ops.py b/tensorflow/contrib/ignite/python/ops/ignite_dataset_ops.py index 288d4853207176b215cd8a0cdcbfb2de5791ecb8..66e654ca636a5a051c6f9cd35bf9001dfbcbf7f4 100644 --- a/tensorflow/contrib/ignite/python/ops/ignite_dataset_ops.py +++ b/tensorflow/contrib/ignite/python/ops/ignite_dataset_ops.py @@ -22,19 +22,21 @@ import socket import ssl import struct +import six + from tensorflow.contrib.ignite.python.ops import gen_dataset_ops from tensorflow.contrib.ignite.python.ops import ignite_op_loader # pylint: disable=unused-import from tensorflow.python.data.ops import dataset_ops +from tensorflow.python.data.util import structure from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.framework import tensor_shape +from tensorflow.python.util import deprecation +@six.add_metaclass(abc.ABCMeta) class Readable(object): - """Readable abstract class that exposes methods to do reading-related - - operations. - """ + """Abstract class that exposes methods to do reading-related operations.""" @abc.abstractmethod def __init__(self): @@ -224,10 +226,7 @@ types = { class TypeTreeNode(object): - """TypeTreeNode class exposes methods to format object tree structure - - data. - """ + """TypeTreeNode class exposes methods to format object tree structure data.""" def __init__(self, name, type_id, fields=None, permutation=None): """Constructs a new instance of TypeTreeNode. @@ -689,18 +688,22 @@ class IgniteClient(TcpClient): class IgniteDataset(dataset_ops.DatasetSource): - """Apache Ignite is a memory-centric distributed database, caching, and - - processing platform for transactional, analytical, and streaming workloads, - delivering in-memory speeds at petabyte scale. This contrib package - contains an integration between Apache Ignite and TensorFlow. The - integration is based on tf.data from TensorFlow side and Binary Client - Protocol from Apache Ignite side. It allows to use Apache Ignite as a - datasource for neural network training, inference and all other + """Apache Ignite is a memory-centric distributed database. + + It acts as a caching and processing platform for transactional, analytical, + and streaming workloads, delivering in-memory speeds at petabyte scale. + This contrib package contains an integration between Apache Ignite and + TensorFlow. The integration is based on tf.data from TensorFlow side and + Binary Client Protocol from Apache Ignite side. It allows to use Apache + Ignite as a datasource for neural network training, inference and all other computations supported by TensorFlow. Ignite Dataset is based on Apache Ignite Binary Client Protocol. """ + @deprecation.deprecated( + None, + "tf.contrib.ignite will be removed in 2.0, the support for Apache Ignite " + "will continue to be provided through the tensorflow/io GitHub project.") def __init__(self, cache_name, host="localhost", @@ -753,6 +756,9 @@ class IgniteDataset(dataset_ops.DatasetSource): self.cache_type.to_permutation(), dtype=dtypes.int32, name="permutation") + self._structure = structure.convert_legacy_structure( + self.cache_type.to_output_types(), self.cache_type.to_output_shapes(), + self.cache_type.to_output_classes()) def _as_variant_tensor(self): return gen_dataset_ops.ignite_dataset(self.cache_name, self.host, self.port, @@ -760,13 +766,5 @@ class IgniteDataset(dataset_ops.DatasetSource): self.schema, self.permutation) @property - def output_classes(self): - return self.cache_type.to_output_classes() - - @property - def output_shapes(self): - return self.cache_type.to_output_shapes() - - @property - def output_types(self): - return self.cache_type.to_output_types() + def _element_structure(self): + return self._structure diff --git a/tensorflow/contrib/ignite/python/ops/ignite_op_loader.py b/tensorflow/contrib/ignite/python/ops/ignite_op_loader.py index c9af7386cf0a26ed1a950130aa36caa7fb831fd0..e450e2d84ba31a7de925fdb78fc972a592c6ad8c 100644 --- a/tensorflow/contrib/ignite/python/ops/ignite_op_loader.py +++ b/tensorflow/contrib/ignite/python/ops/ignite_op_loader.py @@ -21,4 +21,4 @@ from tensorflow.contrib.util import loader from tensorflow.python.platform import resource_loader _dataset_ops = loader.load_op_library( - resource_loader.get_path_to_datafile("../../_dataset_ops.so")) + resource_loader.get_path_to_datafile("../../_ignite_ops.so")) diff --git a/tensorflow/contrib/ignite/python/tests/bin/start-igfs.sh b/tensorflow/contrib/ignite/python/tests/bin/start-igfs.sh new file mode 100755 index 0000000000000000000000000000000000000000..5e39e16c05290f6b5786421670c69a3bd1e27add --- /dev/null +++ b/tensorflow/contrib/ignite/python/tests/bin/start-igfs.sh @@ -0,0 +1,20 @@ +#!/usr/bin/env bash +# 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. +# ============================================================================== + +nohup apache-ignite-fabric/bin/ignite.sh /data/config/ignite-config-igfs.xml & +sleep 5 # Wait Apache Ignite to be started + +tail -f nohup.out diff --git a/tensorflow/contrib/ignite/python/tests/config/ignite-config-igfs.xml b/tensorflow/contrib/ignite/python/tests/config/ignite-config-igfs.xml new file mode 100644 index 0000000000000000000000000000000000000000..5d81bf33226cad0d5cc0ea1fb5c5b55672494976 --- /dev/null +++ b/tensorflow/contrib/ignite/python/tests/config/ignite-config-igfs.xml @@ -0,0 +1,55 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + 127.0.0.1 + + + + + + + + + diff --git a/tensorflow/contrib/ignite/python/tests/igfs_test.py b/tensorflow/contrib/ignite/python/tests/igfs_test.py new file mode 100644 index 0000000000000000000000000000000000000000..cacfc568942e20200b7daf10599dde513a4a0a68 --- /dev/null +++ b/tensorflow/contrib/ignite/python/tests/igfs_test.py @@ -0,0 +1,215 @@ +# 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 IGFS.""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import tensorflow.contrib.ignite.python.ops.igfs_ops # pylint: disable=unused-import +from tensorflow.python.platform import gfile +from tensorflow.python.platform import test + + +class IGFSTest(test.TestCase): + """The Apache Ignite servers have to setup before the test and tear down + + after the test manually. The docker engine has to be installed. + + To setup Apache Ignite servers: + $ bash start_ignite.sh + + To tear down Apache Ignite servers: + $ bash stop_ignite.sh + """ + + def test_create_file(self): + """Test create file. + + """ + # Setup and check preconditions. + file_name = "igfs:///test_create_file/1" + self.assertFalse(gfile.Exists(file_name)) + # Create file. + with gfile.Open(file_name, mode="w") as w: + w.write("") + # Check that file was created. + self.assertTrue(gfile.Exists(file_name)) + + def test_write_read_file(self): + """Test write/read file. + + """ + # Setup and check preconditions. + file_name = "igfs:///test_write_read_file/1" + rows = 10000 + self.assertFalse(gfile.Exists(file_name)) + # Write data. + with gfile.Open(file_name, mode="w") as w: + for i in range(rows): + w.write("This is row\n") + # Read data. + with gfile.Open(file_name, mode="r") as r: + lines = r.readlines() + # Check that data is equal. + self.assertEqual(rows, len(lines)) + for i in range(rows): + self.assertEqual("This is row\n", lines[i]) + + def test_delete_recursively(self): + """Test delete recursively. + + """ + # Setup and check preconditions. + dir_name = "igfs:///test_delete_recursively/" + file_name = "igfs:///test_delete_recursively/1" + self.assertFalse(gfile.Exists(dir_name)) + self.assertFalse(gfile.Exists(file_name)) + gfile.MkDir(dir_name) + with gfile.Open(file_name, mode="w") as w: + w.write("") + self.assertTrue(gfile.Exists(dir_name)) + self.assertTrue(gfile.Exists(file_name)) + # Delete directory recursively. + gfile.DeleteRecursively(dir_name) + # Check that directory was deleted. + self.assertFalse(gfile.Exists(dir_name)) + self.assertFalse(gfile.Exists(file_name)) + + def test_copy(self): + """Test copy. + + """ + # Setup and check preconditions. + src_file_name = "igfs:///test_copy/1" + dst_file_name = "igfs:///test_copy/2" + self.assertFalse(gfile.Exists(src_file_name)) + self.assertFalse(gfile.Exists(dst_file_name)) + with gfile.Open(src_file_name, mode="w") as w: + w.write("42") + self.assertTrue(gfile.Exists(src_file_name)) + self.assertFalse(gfile.Exists(dst_file_name)) + # Copy file. + gfile.Copy(src_file_name, dst_file_name) + # Check that files are identical. + self.assertTrue(gfile.Exists(src_file_name)) + self.assertTrue(gfile.Exists(dst_file_name)) + with gfile.Open(dst_file_name, mode="r") as r: + data = r.read() + self.assertEqual("42", data) + + def test_is_directory(self): + """Test is directory. + + """ + # Setup and check preconditions. + dir_name = "igfs:///test_is_directory/1" + file_name = "igfs:///test_is_directory/2" + with gfile.Open(file_name, mode="w") as w: + w.write("") + gfile.MkDir(dir_name) + # Check that directory is a directory. + self.assertTrue(gfile.IsDirectory(dir_name)) + # Check that file is not a directory. + self.assertFalse(gfile.IsDirectory(file_name)) + + def test_list_directory(self): + """Test list directory. + + """ + # Setup and check preconditions. + dir_name = "igfs:///test_list_directory/" + file_names = [ + "igfs:///test_list_directory/1", "igfs:///test_list_directory/2/3" + ] + ch_dir_names = [ + "igfs:///test_list_directory/4", + ] + for file_name in file_names: + with gfile.Open(file_name, mode="w") as w: + w.write("") + for ch_dir_name in ch_dir_names: + gfile.MkDir(ch_dir_name) + ls_expected_result = file_names + ch_dir_names + # Get list of files in directory. + ls_result = gfile.ListDirectory(dir_name) + # Check that list of files is correct. + self.assertEqual(len(ls_expected_result), len(ls_result)) + for e in ["1", "2", "4"]: + self.assertTrue(e in ls_result) + + def test_make_dirs(self): + """Test make dirs. + + """ + # Setup and check preconditions. + dir_name = "igfs:///test_make_dirs/" + self.assertFalse(gfile.Exists(dir_name)) + # Make directory. + gfile.MkDir(dir_name) + # Check that directory was created. + self.assertTrue(gfile.Exists(dir_name)) + + def test_remove(self): + """Test remove. + + """ + # Setup and check preconditions. + file_name = "igfs:///test_remove/1" + self.assertFalse(gfile.Exists(file_name)) + with gfile.Open(file_name, mode="w") as w: + w.write("") + self.assertTrue(gfile.Exists(file_name)) + # Remove file. + gfile.Remove(file_name) + # Check that file was removed. + self.assertFalse(gfile.Exists(file_name)) + + def test_rename_file(self): + """Test rename file. + + """ + # Setup and check preconditions. + src_file_name = "igfs:///test_rename_file/1" + dst_file_name = "igfs:///test_rename_file/2" + with gfile.Open(src_file_name, mode="w") as w: + w.write("42") + self.assertTrue(gfile.Exists(src_file_name)) + # Rename file. + gfile.Rename(src_file_name, dst_file_name) + # Check that only new name of file is available. + self.assertFalse(gfile.Exists(src_file_name)) + self.assertTrue(gfile.Exists(dst_file_name)) + with gfile.Open(dst_file_name, mode="r") as r: + data = r.read() + self.assertEqual("42", data) + + def test_rename_dir(self): + """Test rename dir. + + """ + # Setup and check preconditions. + src_dir_name = "igfs:///test_rename_dir/1" + dst_dir_name = "igfs:///test_rename_dir/2" + gfile.MkDir(src_dir_name) + # Rename directory. + gfile.Rename(src_dir_name, dst_dir_name) + # Check that only new name of directory is available. + self.assertFalse(gfile.Exists(src_dir_name)) + self.assertTrue(gfile.Exists(dst_dir_name)) + self.assertTrue(gfile.IsDirectory(dst_dir_name)) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/ignite/python/tests/ignite_dataset_test.py b/tensorflow/contrib/ignite/python/tests/ignite_dataset_test.py index ef29b5f14a4b2fea2400ec4d56a7ad2cf44cf2cb..ff5d4c458c859fd8e5e3ae65ee41a454d55d6538 100644 --- a/tensorflow/contrib/ignite/python/tests/ignite_dataset_test.py +++ b/tensorflow/contrib/ignite/python/tests/ignite_dataset_test.py @@ -21,6 +21,7 @@ import os 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 @@ -65,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.make_one_shot_iterator() + it = dataset_ops.make_one_shot_iterator(dataset) ne = it.get_next() with session.Session() as sess: diff --git a/tensorflow/contrib/ignite/python/tests/start_ignite.sh b/tensorflow/contrib/ignite/python/tests/start_ignite.sh index a67bd44f2fb0d654ba07f022a5070c68df8e2ede..112e0dea844620de600e277bff3685dd7c42c49c 100755 --- a/tensorflow/contrib/ignite/python/tests/start_ignite.sh +++ b/tensorflow/contrib/ignite/python/tests/start_ignite.sh @@ -20,3 +20,7 @@ SCRIPT_PATH="$( cd "$(dirname "$0")" ; pwd -P )" # Start Apache Ignite with plain client listener. docker run -itd --name ignite-plain -p 42300:10800 \ -v ${SCRIPT_PATH}:/data apacheignite/ignite:${IGNITE_VERSION} /data/bin/start-plain.sh + +# Start Apache Ignite with IGFS. +docker run -itd --name ignite-igfs -p 10500:10500 \ +-v ${SCRIPT_PATH}:/data apacheignite/ignite:${IGNITE_VERSION} /data/bin/start-igfs.sh \ No newline at end of file diff --git a/tensorflow/contrib/ignite/python/tests/stop_ignite.sh b/tensorflow/contrib/ignite/python/tests/stop_ignite.sh index 8f03dbd1ede61f548d3de9d9738f97667e75df3c..35b0f32d1b3e1373a231ff23f2b40c8ccc417baf 100755 --- a/tensorflow/contrib/ignite/python/tests/stop_ignite.sh +++ b/tensorflow/contrib/ignite/python/tests/stop_ignite.sh @@ -15,5 +15,4 @@ # ============================================================================== docker rm -f ignite-plain -docker rm -f ignite-ssl -docker rm -f ignite-ssl-auth +docker rm -f ignite-igfs \ No newline at end of file diff --git a/tensorflow/contrib/image/kernels/adjust_hsv_in_yiq_op.cc b/tensorflow/contrib/image/kernels/adjust_hsv_in_yiq_op.cc index 478b716d88321101c971789f36c0ff8ecd3f418e..108da04494685f06f9afc26a26a5dadcdd99b0ff 100644 --- a/tensorflow/contrib/image/kernels/adjust_hsv_in_yiq_op.cc +++ b/tensorflow/contrib/image/kernels/adjust_hsv_in_yiq_op.cc @@ -115,7 +115,7 @@ class AdjustHsvInYiqOp : public AdjustHsvInYiqOpBase { *context->device()->tensorflow_cpu_worker_threads(); Shard(worker_threads.num_threads, worker_threads.workers, channel_count, kCostPerChannel, - [channel_count, &input_data, &output_data, &tranformation_matrix]( + [&input_data, &output_data, &tranformation_matrix]( int64 start_channel, int64 end_channel) { // Applying projection matrix to input RGB vectors. const float* p = input_data.data() + start_channel * kChannelSize; diff --git a/tensorflow/contrib/image/python/kernel_tests/dense_image_warp_test.py b/tensorflow/contrib/image/python/kernel_tests/dense_image_warp_test.py index 24b790977dfdb675ff7bf0a119a08e243a30d3aa..ae9c7a611945e1445c933d74b9944054b3f0e0a4 100644 --- a/tensorflow/contrib/image/python/kernel_tests/dense_image_warp_test.py +++ b/tensorflow/contrib/image/python/kernel_tests/dense_image_warp_test.py @@ -24,7 +24,7 @@ from tensorflow.contrib.image.python.ops import dense_image_warp 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.ops import array_ops from tensorflow.python.ops import gradients @@ -259,7 +259,7 @@ class DenseImageWarpTest(test_util.TensorFlowTestCase): shape = [1, 2, 1, 1] msg = 'Should have raised an exception for invalid image size' - with self.assertRaises(ValueError, msg=msg): + with self.assertRaises(errors.InvalidArgumentError, msg=msg): self.check_interpolation_correctness(shape, 'float32', 'float32') diff --git a/tensorflow/contrib/image/python/kernel_tests/image_ops_test.py b/tensorflow/contrib/image/python/kernel_tests/image_ops_test.py index 4997c31a7fc7f4243d03b22fc9c01fb13a2a25a4..ba5cdfebf92c07e496ed588848d5859ff6a5bff2 100644 --- a/tensorflow/contrib/image/python/kernel_tests/image_ops_test.py +++ b/tensorflow/contrib/image/python/kernel_tests/image_ops_test.py @@ -281,6 +281,13 @@ class ImageOpsTest(test_util.TensorFlowTestCase): value.eval(), np.array([[4, 4], [4, 4]]).astype(dtype.as_numpy_dtype())) + @test_util.run_in_graph_and_eager_modes + def test_transform_eager(self): + image = constant_op.constant([[1., 2.], [3., 4.]]) + value = image_ops.transform(image, [1] * 8) + with self.test_session(use_gpu=True): + self.assertAllEqual(self.evaluate(value), np.array([[4, 4], [4, 4]])) + class BipartiteMatchTest(test_util.TensorFlowTestCase): diff --git a/tensorflow/contrib/image/python/ops/dense_image_warp.py b/tensorflow/contrib/image/python/ops/dense_image_warp.py index f9b219ada492466919c615d8978e462e6c619d33..f7ced440720209cb05dfcd79395c51517f9de0d5 100644 --- a/tensorflow/contrib/image/python/ops/dense_image_warp.py +++ b/tensorflow/contrib/image/python/ops/dense_image_warp.py @@ -24,6 +24,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 check_ops from tensorflow.python.ops import math_ops @@ -60,28 +61,38 @@ def _interpolate_bilinear(grid, msg = 'Grid must be 4 dimensional. Received size: ' raise ValueError(msg + str(grid.get_shape())) - batch_size, height, width, channels = shape + batch_size, height, width, channels = (array_ops.shape(grid)[0], + array_ops.shape(grid)[1], + array_ops.shape(grid)[2], + array_ops.shape(grid)[3]) + + shape = [batch_size, height, width, channels] query_type = query_points.dtype grid_type = grid.dtype - if (len(query_points.get_shape()) != 3 or - query_points.get_shape()[2].value != 2): - msg = ('Query points must be 3 dimensional and size 2 in dim 2. Received ' - 'size: ') - raise ValueError(msg + str(query_points.get_shape())) - - _, num_queries, _ = query_points.get_shape().as_list() - - if height < 2 or width < 2: - msg = 'Grid must be at least batch_size x 2 x 2 in size. Received size: ' - raise ValueError(msg + str(grid.get_shape())) - - alphas = [] - floors = [] - ceils = [] - - index_order = [0, 1] if indexing == 'ij' else [1, 0] - unstacked_query_points = array_ops.unstack(query_points, axis=2) + with ops.control_dependencies([ + check_ops.assert_equal( + len(query_points.get_shape()), + 3, + message='Query points must be 3 dimensional.'), + check_ops.assert_equal( + array_ops.shape(query_points)[2], + 2, + message='Query points must be size 2 in dim 2.') + ]): + num_queries = array_ops.shape(query_points)[1] + + with ops.control_dependencies([ + check_ops.assert_greater_equal( + height, 2, message='Grid height must be at least 2.'), + check_ops.assert_greater_equal( + width, 2, message='Grid width must be at least 2.') + ]): + alphas = [] + floors = [] + ceils = [] + index_order = [0, 1] if indexing == 'ij' else [1, 0] + unstacked_query_points = array_ops.unstack(query_points, axis=2) for dim in index_order: with ops.name_scope('dim-' + str(dim)): @@ -112,16 +123,18 @@ def _interpolate_bilinear(grid, alpha = array_ops.expand_dims(alpha, 2) alphas.append(alpha) - if batch_size * height * width > np.iinfo(np.int32).max / 8: - error_msg = """The image size or batch size is sufficiently large - that the linearized addresses used by array_ops.gather - may exceed the int32 limit.""" - raise ValueError(error_msg) - - flattened_grid = array_ops.reshape(grid, - [batch_size * height * width, channels]) - batch_offsets = array_ops.reshape( - math_ops.range(batch_size) * height * width, [batch_size, 1]) + with ops.control_dependencies([ + check_ops.assert_less_equal( + math_ops.cast(batch_size * height * width, dtype=dtypes.float32), + np.iinfo(np.int32).max / 8, + message="""The image size or batch size is sufficiently large + that the linearized addresses used by array_ops.gather + may exceed the int32 limit.""") + ]): + flattened_grid = array_ops.reshape( + grid, [batch_size * height * width, channels]) + batch_offsets = array_ops.reshape( + math_ops.range(batch_size) * height * width, [batch_size, 1]) # This wraps array_ops.gather. We reshape the image data such that the # batch, y, and x coordinates are pulled into the first dimension. @@ -182,7 +195,11 @@ def dense_image_warp(image, flow, name='dense_image_warp'): of dimensions. """ with ops.name_scope(name): - batch_size, height, width, channels = image.get_shape().as_list() + batch_size, height, width, channels = (array_ops.shape(image)[0], + array_ops.shape(image)[1], + array_ops.shape(image)[2], + array_ops.shape(image)[3]) + # The flow is defined on the image grid. Turn the flow into a list of query # points in the grid space. grid_x, grid_y = array_ops.meshgrid( diff --git a/tensorflow/contrib/image/python/ops/image_ops.py b/tensorflow/contrib/image/python/ops/image_ops.py index d4fb99a017faebe30384d739f22f4ff5fa986bc4..b25a6f7b5742917a032946fe03a0dab20e7dc1ad 100644 --- a/tensorflow/contrib/image/python/ops/image_ops.py +++ b/tensorflow/contrib/image/python/ops/image_ops.py @@ -17,6 +17,7 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +from tensorflow.python.eager import context from tensorflow.contrib.image.ops import gen_image_ops from tensorflow.contrib.util import loader from tensorflow.python.framework import common_shapes @@ -271,8 +272,11 @@ def transform(images, raise TypeError("Images should have rank between 2 and 4.") if output_shape is None: - output_shape = tensor_util.constant_value( - array_ops.shape(images)[1:3]) or array_ops.shape(images)[1:3] + output_shape = array_ops.shape(images)[1:3] + if not context.executing_eagerly(): + output_shape_value = tensor_util.constant_value(output_shape) + if output_shape_value is not None: + output_shape = output_shape_value output_shape = ops.convert_to_tensor( output_shape, dtypes.int32, name="output_shape") diff --git a/tensorflow/contrib/image/python/ops/interpolate_spline.py b/tensorflow/contrib/image/python/ops/interpolate_spline.py index f0b408faa3320741cf83b3aaec0f40030f906578..3a444d26c2a45278261b684da4fa4ac249d0d5cd 100644 --- a/tensorflow/contrib/image/python/ops/interpolate_spline.py +++ b/tensorflow/contrib/image/python/ops/interpolate_spline.py @@ -18,6 +18,7 @@ from __future__ import division from __future__ import print_function 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 linalg_ops from tensorflow.python.ops import math_ops @@ -100,12 +101,12 @@ def _solve_interpolation(train_points, train_values, order, b, n, _ = array_ops.unstack(array_ops.shape(train_points), num=3) d = train_points.shape[-1] - if d.value is None: + if tensor_shape.dimension_value(d) is None: raise ValueError('The dimensionality of the input points (d) must be ' 'statically-inferrable.') k = train_values.shape[-1] - if k.value is None: + if tensor_shape.dimension_value(k) is None: raise ValueError('The dimensionality of the output values (k) must be ' 'statically-inferrable.') diff --git a/tensorflow/contrib/image/python/ops/sparse_image_warp.py b/tensorflow/contrib/image/python/ops/sparse_image_warp.py index 1ea8f705b7e6f522281de6384de0d42efab6a406..51449ff5e938946844fd4245215008a257e8b045 100644 --- a/tensorflow/contrib/image/python/ops/sparse_image_warp.py +++ b/tensorflow/contrib/image/python/ops/sparse_image_warp.py @@ -24,6 +24,7 @@ from tensorflow.contrib.image.python.ops import interpolate_spline from tensorflow.python.framework import constant_op from tensorflow.python.framework import ops +from tensorflow.python.framework import tensor_shape from tensorflow.python.ops import array_ops @@ -76,7 +77,7 @@ def _add_zero_flow_controls_at_boundary(control_point_locations, merged_control_point_flows: augmented set of control point flows """ - batch_size = control_point_locations.get_shape()[0].value + batch_size = tensor_shape.dimension_value(control_point_locations.shape[0]) boundary_point_locations = _get_boundary_locations(image_height, image_width, boundary_points_per_edge) diff --git a/tensorflow/contrib/integrate/python/ops/odes.py b/tensorflow/contrib/integrate/python/ops/odes.py index 7b7ac4f347e30d20eb2f4889e0cae5669c975e4f..b7d77130bd03ba05aca3ab94ccf94eb2ed5d9347 100644 --- a/tensorflow/contrib/integrate/python/ops/odes.py +++ b/tensorflow/contrib/integrate/python/ops/odes.py @@ -540,7 +540,8 @@ def odeint(func, **options) -class _FixedGridIntegrator(six.with_metaclass(abc.ABCMeta)): +@six.add_metaclass(abc.ABCMeta) +class _FixedGridIntegrator(object): """Base class for fixed-grid ODE integrators.""" def integrate(self, evol_func, y0, time_grid, dt_grid, steps_on_intervals): diff --git a/tensorflow/contrib/kafka/python/ops/kafka_dataset_ops.py b/tensorflow/contrib/kafka/python/ops/kafka_dataset_ops.py index 7129f09e8b42e48a9c768fd4a66cde3d4da9d31d..b399e1b6c2ac47db205b5d8bbc81875ef5c08a31 100644 --- a/tensorflow/contrib/kafka/python/ops/kafka_dataset_ops.py +++ b/tensorflow/contrib/kafka/python/ops/kafka_dataset_ops.py @@ -20,15 +20,20 @@ from __future__ import print_function from tensorflow.contrib.kafka.python.ops import gen_dataset_ops from tensorflow.contrib.kafka.python.ops import kafka_op_loader # pylint: disable=unused-import from tensorflow.python.data.ops import dataset_ops +from tensorflow.python.data.util import structure from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops -from tensorflow.python.framework import tensor_shape +from tensorflow.python.util import deprecation class KafkaDataset(dataset_ops.DatasetSource): """A Kafka Dataset that consumes the message. """ + @deprecation.deprecated( + None, + "tf.contrib.kafka will be removed in 2.0, the support for Apache Kafka " + "will continue to be provided through the tensorflow/io GitHub project.") def __init__(self, topics, servers="localhost", @@ -63,13 +68,5 @@ class KafkaDataset(dataset_ops.DatasetSource): self._group, self._eof, self._timeout) @property - def output_classes(self): - return ops.Tensor - - @property - def output_shapes(self): - return tensor_shape.scalar() - - @property - def output_types(self): - return dtypes.string + def _element_structure(self): + return structure.TensorStructure(dtypes.string, []) diff --git a/tensorflow/contrib/keras/api/keras/layers/__init__.py b/tensorflow/contrib/keras/api/keras/layers/__init__.py index 3327a9f9a613bfb56e6a25af0fe1c0ca18609035..9e19884df852c0fd259a55aef56c62b4189cd1da 100644 --- a/tensorflow/contrib/keras/api/keras/layers/__init__.py +++ b/tensorflow/contrib/keras/api/keras/layers/__init__.py @@ -20,7 +20,7 @@ from __future__ import print_function # Generic layers. # pylint: disable=g-bad-import-order -from tensorflow.python.keras.engine.base_layer import InputSpec +from tensorflow.python.keras.engine.input_spec import InputSpec from tensorflow.python.keras.engine.base_layer import Layer from tensorflow.python.keras.engine.input_layer import Input from tensorflow.python.keras.engine.input_layer import InputLayer diff --git a/tensorflow/contrib/keras/api/keras/utils/__init__.py b/tensorflow/contrib/keras/api/keras/utils/__init__.py index 47cd01b924fb43e8a83836c58f8ced61e9e88268..3b9fa1b230b837a350d521c4165053c187786201 100644 --- a/tensorflow/contrib/keras/api/keras/utils/__init__.py +++ b/tensorflow/contrib/keras/api/keras/utils/__init__.py @@ -30,6 +30,7 @@ from tensorflow.python.keras.utils.generic_utils import Progbar from tensorflow.python.keras.utils.generic_utils import serialize_keras_object from tensorflow.python.keras.utils.io_utils import HDF5Matrix from tensorflow.python.keras.utils.layer_utils import convert_all_kernels_in_model +from tensorflow.python.keras.utils.losses_utils import squeeze_or_expand_dimensions from tensorflow.python.keras.utils.np_utils import normalize from tensorflow.python.keras.utils.np_utils import to_categorical from tensorflow.python.keras.utils.vis_utils import plot_model diff --git a/tensorflow/contrib/kernel_methods/python/kernel_estimators.py b/tensorflow/contrib/kernel_methods/python/kernel_estimators.py index de7530231db4ea4f50996a67eb8c0d6936db9dd3..1626e55b9b3bc82bd96703bfab765ac6ad81f462 100644 --- a/tensorflow/contrib/kernel_methods/python/kernel_estimators.py +++ b/tensorflow/contrib/kernel_methods/python/kernel_estimators.py @@ -90,7 +90,7 @@ def _update_features_and_columns(features, feature_columns, mapped_column_name = column_name + "_MAPPED" # Construct new feature columns based on provided kernel_mappers. column_kernel_mappers = kernel_mappers_dict[feature_column] - new_dim = sum([mapper.output_dim for mapper in column_kernel_mappers]) + new_dim = sum(mapper.output_dim for mapper in column_kernel_mappers) mapped_columns.add( layers.feature_column.real_valued_column(mapped_column_name, new_dim)) 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/kernel_methods/python/mappers/dense_kernel_mapper.py b/tensorflow/contrib/kernel_methods/python/mappers/dense_kernel_mapper.py index db38b471520e1922392e7aaf8ee66d7f304248c9..04ecdbfdb6625766eb87c1527592e616c5cdfbf9 100644 --- a/tensorflow/contrib/kernel_methods/python/mappers/dense_kernel_mapper.py +++ b/tensorflow/contrib/kernel_methods/python/mappers/dense_kernel_mapper.py @@ -35,7 +35,6 @@ class DenseKernelMapper(object): This class is abstract. Users should not create instances of this class. """ - __metaclass__ = abc.ABCMeta @abc.abstractmethod def map(self, input_tensor): diff --git a/tensorflow/contrib/kfac/README.md b/tensorflow/contrib/kfac/README.md index 42b91d031375b8edb7e4f364ac91ffb74ef1f54b..19daffea6c7e4486499388314d0aaaa611e94218 100644 --- a/tensorflow/contrib/kfac/README.md +++ b/tensorflow/contrib/kfac/README.md @@ -1,3 +1,3 @@ # K-FAC: Kronecker-Factored Approximate Curvature -## KFAC moved to third_party/tensorflow_kfac. +## KFAC moved to https://github.com/tensorflow/kfac. diff --git a/tensorflow/contrib/kinesis/python/ops/kinesis_dataset_ops.py b/tensorflow/contrib/kinesis/python/ops/kinesis_dataset_ops.py index 75806dbbeb1819bb0a6965bbc384e02df9895210..9479afb180df7bb4a08d6aafa4fc3bf63489d9f3 100644 --- a/tensorflow/contrib/kinesis/python/ops/kinesis_dataset_ops.py +++ b/tensorflow/contrib/kinesis/python/ops/kinesis_dataset_ops.py @@ -20,9 +20,10 @@ from __future__ import print_function from tensorflow.contrib.kinesis.python.ops import gen_dataset_ops from tensorflow.contrib.kinesis.python.ops import kinesis_op_loader # pylint: disable=unused-import from tensorflow.python.data.ops import dataset_ops +from tensorflow.python.data.util import structure from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops -from tensorflow.python.framework import tensor_shape +from tensorflow.python.util import deprecation class KinesisDataset(dataset_ops.DatasetSource): @@ -34,15 +35,12 @@ class KinesisDataset(dataset_ops.DatasetSource): For example, we can construct and use the KinesisDataset as follows: ```python + tf.enable_eager_execution() + dataset = tf.contrib.kinesis.KinesisDataset( "kinesis_stream_name", read_indefinitely=False) - next = dataset.make_one_shot_iterator().get_next() - with tf.Session() as sess: - while True: - try: - print(sess.run(nxt)) - except tf.errors.OutOfRangeError: - break + for element in dataset: + print(element) ``` Since Kinesis is a data streaming service, data may not be available @@ -53,6 +51,10 @@ class KinesisDataset(dataset_ops.DatasetSource): is returned immediately instead. """ + @deprecation.deprecated( + None, + "tf.contrib.kinesis will be removed in 2.0, the support for Kinesis " + "will continue to be provided through the tensorflow/io GitHub project.") def __init__(self, stream, shard="", @@ -69,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( @@ -78,19 +79,12 @@ 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( self._stream, self._shard, self._read_indefinitely, self._interval) @property - def output_classes(self): - return ops.Tensor - - @property - def output_shapes(self): - return tensor_shape.scalar() - - @property - def output_types(self): - return dtypes.string + def _element_structure(self): + return structure.TensorStructure(dtypes.string, []) diff --git a/tensorflow/contrib/labeled_tensor/BUILD b/tensorflow/contrib/labeled_tensor/BUILD index c8812d4b23f94102d093db878a709b090a3318d6..588f15b867c1fedbadd5a5d945d870a356549468 100644 --- a/tensorflow/contrib/labeled_tensor/BUILD +++ b/tensorflow/contrib/labeled_tensor/BUILD @@ -70,7 +70,10 @@ py_test( "python/ops/core_test.py", ], srcs_version = "PY2AND3", - tags = ["no_windows"], # TODO: needs investigation on Windows + tags = [ + "no_windows", # TODO: needs investigation on Windows + "noasan", # TODO(b/119323169) + ], deps = [ ":_typecheck", ":core", diff --git a/tensorflow/contrib/layers/BUILD b/tensorflow/contrib/layers/BUILD index b4fe8cac74cb7d29b9646b6b968ccf37b3d6ea7a..69d5496f8aebb9b89c5d79f80a1a439f556093d7 100644 --- a/tensorflow/contrib/layers/BUILD +++ b/tensorflow/contrib/layers/BUILD @@ -1,15 +1,16 @@ # Description: # contains parts of TensorFlow that are experimental or unstable and which are not supported. -licenses(["notice"]) # Apache 2.0 - -exports_files(["LICENSE"]) - package(default_visibility = [ "//learning/brain:__subpackages__", "//tensorflow:__subpackages__", + "//tensorflow_model_optimization:__subpackages__", ]) +licenses(["notice"]) # Apache 2.0 + +exports_files(["LICENSE"]) + load("//tensorflow:tensorflow.bzl", "cuda_py_test") load("//tensorflow:tensorflow.bzl", "tf_custom_op_py_library") load("//tensorflow:tensorflow.bzl", "py_test") @@ -77,6 +78,12 @@ tf_custom_op_py_library( ":sparse_feature_cross_op_op_lib", ], srcs_version = "PY2AND3", + visibility = [ + "//learning/brain:__subpackages__", + "//tensorflow:__subpackages__", + "//tensorflow_model_optimization:__subpackages__", + "//video/youtube/personalization:__subpackages__", + ], deps = [ ":sparse_feature_cross_op", "//tensorflow/contrib/framework:framework_py", @@ -252,7 +259,7 @@ py_test( "//tensorflow/python:training", "//tensorflow/python:variable_scope", "//tensorflow/python:variables", - "//tensorflow/python/feature_column", + "//tensorflow/python/feature_column:feature_column_py", "//third_party/py/numpy", ], ) @@ -276,7 +283,7 @@ py_test( "//tensorflow/python:sparse_tensor", "//tensorflow/python:variable_scope", "//tensorflow/python:variables", - "//tensorflow/python/feature_column", + "//tensorflow/python/feature_column:feature_column_py", "//third_party/py/numpy", ], ) diff --git a/tensorflow/contrib/layers/__init__.py b/tensorflow/contrib/layers/__init__.py index af8e673f5906ad972408d30f23f2e8ba7e031a00..32f3006b749e3b34572a8d642054c0ec4c4664b0 100644 --- a/tensorflow/contrib/layers/__init__.py +++ b/tensorflow/contrib/layers/__init__.py @@ -14,10 +14,6 @@ # ============================================================================== """Ops for building neural network layers, regularizers, summaries, etc. -See the -[Contrib Layers](https://tensorflow.org/api_guides/python/contrib.layers) -guide. - @@avg_pool2d @@avg_pool3d @@batch_norm diff --git a/tensorflow/contrib/layers/python/layers/embedding_ops.py b/tensorflow/contrib/layers/python/layers/embedding_ops.py index 60e1d85ea9c08a51763fdaf08853f8d9b67347e5..429d696daf0baf85f6a60aa4d299b513d90c5925 100644 --- a/tensorflow/contrib/layers/python/layers/embedding_ops.py +++ b/tensorflow/contrib/layers/python/layers/embedding_ops.py @@ -124,7 +124,8 @@ def safe_embedding_lookup_sparse(embedding_weights, sparse_weights]) as scope: # Reshape higher-rank sparse ids and weights to linear segment ids. original_shape = sparse_ids.dense_shape - original_rank_dim = sparse_ids.dense_shape.get_shape()[0] + original_rank_dim = tensor_shape.Dimension(tensor_shape.dimension_value( + sparse_ids.dense_shape.get_shape()[0])) original_rank = ( array_ops.size(original_shape) if original_rank_dim.value is None @@ -349,7 +350,7 @@ def _sampled_scattered_embedding_lookup( shape = params[p].get_shape() shape.assert_has_rank(1) shape.assert_is_fully_defined() - partition_sizes.append(shape[0].value) + partition_sizes.append(tensor_shape.dimension_value(shape[0])) num_params = sum(partition_sizes) # Total number of parameters. # Assert the size of each partition. @@ -779,16 +780,16 @@ def _embedding_lookup_with_distributed_aggregation(params, # Compute num_total_ids as the sum of dim-0 of params, then assign to # partitions based on a constant number of ids per partition. Optimize # if we already know the full shape statically. - dim_0_size = params[0].get_shape()[0] + dim_0_size = params[0].get_shape().dims[0] for p in xrange(1, np): - dim_0_size += params[p].get_shape()[0] + dim_0_size += params[p].get_shape().dims[0] if dim_0_size.value: - num_total_ids = constant_op.constant(dim_0_size.value, flat_ids.dtype) + num_total_ids = constant_op.constant(dim_0_size, flat_ids.dtype) else: dim_0_sizes = [] for p in xrange(np): - if params[p].get_shape()[0].value is not None: - dim_0_sizes.append(params[p].get_shape()[0].value) + if params[p].get_shape().dims[0].value is not None: + dim_0_sizes.append(params[p].get_shape().dims[0].value) else: with ops.colocate_with(params[p]): dim_0_sizes.append(array_ops.shape(params[p])[0]) diff --git a/tensorflow/contrib/layers/python/layers/embedding_ops_test.py b/tensorflow/contrib/layers/python/layers/embedding_ops_test.py index 124515e5a6474f2cc1038830346e27277c6ceea7..295c721fceda6aaaf8672525ceed560308db6af7 100644 --- a/tensorflow/contrib/layers/python/layers/embedding_ops_test.py +++ b/tensorflow/contrib/layers/python/layers/embedding_ops_test.py @@ -21,6 +21,7 @@ from __future__ import print_function import itertools import math +import sys import numpy as np @@ -36,6 +37,7 @@ from tensorflow.python.ops import gradient_checker from tensorflow.python.ops import init_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import partitioned_variables +from tensorflow.python.ops import variable_scope from tensorflow.python.platform import test from tensorflow.python.util import compat @@ -48,11 +50,13 @@ class SafeEmbeddingLookupSparseTest(test.TestCase): assert num_shards > 0 assert num_shards <= vocab_size - embedding_weights = partitioned_variables.create_partitioned_variables( + initializer = init_ops.truncated_normal_initializer( + mean=0.0, stddev=1.0 / math.sqrt(vocab_size), dtype=dtypes.float32) + embedding_weights = list(variable_scope.get_variable( + "embedding_weights", shape=[vocab_size, embed_dim], - slicing=[num_shards, 1], - initializer=init_ops.truncated_normal_initializer( - mean=0.0, stddev=1.0 / math.sqrt(vocab_size), dtype=dtypes.float32)) + partitioner=partitioned_variables.fixed_size_partitioner(num_shards), + initializer=initializer)) for w in embedding_weights: w.initializer.run() embedding_weights = [w.eval() for w in embedding_weights] @@ -256,6 +260,13 @@ class SafeEmbeddingLookupSparseTest(test.TestCase): embedding_weights, sparse_ids, sparse_weights) +# pylint: disable=invalid-name +def local_variable_scope(): + """Create a variable scope named like the caller function.""" + return variable_scope.variable_scope(sys._getframe(1).f_code.co_name) +# pylint: enable=invalid-name + + class ScatteredEmbeddingLookupTest(test.TestCase): def setUp(self): @@ -266,17 +277,18 @@ class ScatteredEmbeddingLookupTest(test.TestCase): assert num_shards > 0 assert num_shards <= size - embedding_weights = partitioned_variables.create_partitioned_variables( + embedding_weights = list(variable_scope.get_variable( + "embedding_weights", shape=[size], - slicing=[num_shards], + partitioner=partitioned_variables.fixed_size_partitioner(num_shards), initializer=init_ops.truncated_normal_initializer( - mean=0.0, stddev=1.0, dtype=dtypes.float32)) + mean=0.0, stddev=1.0, dtype=dtypes.float32))) for w in embedding_weights: w.initializer.run() return embedding_weights def test_scattered_embedding_consistency(self): - with self.cached_session(): + with self.cached_session(), local_variable_scope(): embedding_weights = self._random_weights() values = constant_op.constant(["foo", "foo"]) @@ -288,7 +300,7 @@ class ScatteredEmbeddingLookupTest(test.TestCase): embedding_lookup_result[1]) def test_scattered_embedding_multiple_partition(self): - with self.cached_session(): + with self.cached_session(), local_variable_scope(): embedding_weights = self._random_weights(num_shards=7) values = constant_op.constant([4, 4, 5]) @@ -304,7 +316,7 @@ class ScatteredEmbeddingLookupTest(test.TestCase): self.assertGreater(embedding_diff, 0) def test_scattered_embedding_coverage(self): - with self.cached_session(): + with self.cached_session(), local_variable_scope(): size = 8 embedding_weights = self._random_weights(size=size, num_shards=3) values = constant_op.constant(["foo"]) @@ -316,7 +328,7 @@ class ScatteredEmbeddingLookupTest(test.TestCase): self.assertEqual(len(np.unique(embedding_lookup_result[0])), size) def test_scattered_embedding_multi_dimension(self): - with self.cached_session(): + with self.cached_session(), local_variable_scope(): embedding_weights = self._random_weights() values = constant_op.constant([["foo", "bar", "bar"], ["bar", "bar", "foo"]]) @@ -329,7 +341,7 @@ class ScatteredEmbeddingLookupTest(test.TestCase): embedding_lookup_result[1][2]) def test_scattered_embedding_lookup_sparse(self): - with self.cached_session(): + with self.cached_session(), local_variable_scope(): embedding_weights = self._random_weights(num_shards=3) sparse_tensor = sparse_tensor_lib.SparseTensor( values=["foo", "bar", "foo", "bar"], @@ -358,7 +370,7 @@ class ScatteredEmbeddingLookupTest(test.TestCase): embeds = np.random.randn(n_embed, d_embed) idx = np.random.randint(0, n_embed, idx_shape) - with self.cached_session(): + with self.cached_session(), local_variable_scope(): embedded_np = embeds[idx] embedded_tf = embedding_ops.embedding_lookup_unique(embeds, idx).eval() @@ -370,7 +382,7 @@ class ScatteredEmbeddingLookupTest(test.TestCase): idx = np.random.randint(0, 5, 10) idx2d = np.random.randint(0, 5, (10, 2)) - with self.cached_session(): + with self.cached_session(), local_variable_scope(): embedded_np = embeds[idx] embedded_np2d = embeds[idx2d] embedded_tf = embedding_ops.embedding_lookup_unique(embeds, idx).eval() @@ -398,17 +410,18 @@ class SampledScatteredEmbeddingLookupTest(test.TestCase): assert num_shards > 0 assert num_shards <= size - embedding_weights = partitioned_variables.create_partitioned_variables( + embedding_weights = list(variable_scope.get_variable( + "embedding_weights", shape=[size], - slicing=[num_shards], + partitioner=partitioned_variables.fixed_size_partitioner(num_shards), initializer=init_ops.truncated_normal_initializer( - mean=0.0, stddev=1.0, dtype=dtypes.float32)) + mean=0.0, stddev=1.0, dtype=dtypes.float32))) for w in embedding_weights: w.initializer.run() return embedding_weights def test_hashed_embedding_consistency(self): - with self.cached_session(): + with self.cached_session(), local_variable_scope(): embedding_weights = self._random_weights() values = constant_op.constant(["foo", "foo"]) # The first three sampled_candidates are equal, so the first three @@ -429,7 +442,7 @@ class SampledScatteredEmbeddingLookupTest(test.TestCase): embedding_lookup_result[1][3]) def test_hashed_embedding_multi_dimension(self): - with self.cached_session(): + with self.cached_session(), local_variable_scope(): embedding_weights = self._random_weights() values = constant_op.constant([["foo", "bar", "bar"], ["bar", "bar", "foo"]]) @@ -691,7 +704,6 @@ class EmbeddingLookupSparseWithDistributedAggregationTest(test.TestCase): index += num_val return grouped_vals - @test_util.enable_c_shapes def testEmbeddingLookupSparse(self): vocab_size = 13 batch_size = 10 diff --git a/tensorflow/contrib/layers/python/layers/encoders.py b/tensorflow/contrib/layers/python/layers/encoders.py index f42112206d0db9d2e42bd4cff19f6a6533951d46..3671633c8d795034b13cb55fd6db87c453e9fa12 100644 --- a/tensorflow/contrib/layers/python/layers/encoders.py +++ b/tensorflow/contrib/layers/python/layers/encoders.py @@ -84,8 +84,7 @@ def bow_encoder(ids, if isinstance(ids, sparse_tensor.SparseTensor): raise TypeError('ids are expected to be dense Tensor, got: %s', ids) return math_ops.reduce_mean( - embedding_ops.embedding_lookup(embeddings, ids), - reduction_indices=1) + embedding_ops.embedding_lookup(embeddings, ids), axis=1) def embed_sequence(ids, diff --git a/tensorflow/contrib/layers/python/layers/feature_column.py b/tensorflow/contrib/layers/python/layers/feature_column.py index 53c8ae5d0893641c79a7f24851a10afc44a2144a..00d819ed0e9fe3a5644105a571beda100204631e 100644 --- a/tensorflow/contrib/layers/python/layers/feature_column.py +++ b/tensorflow/contrib/layers/python/layers/feature_column.py @@ -194,6 +194,7 @@ class _DeepEmbeddingLookupArguments( pass +@six.add_metaclass(abc.ABCMeta) class _FeatureColumn(object): """Represents a feature column abstraction. @@ -205,7 +206,6 @@ class _FeatureColumn(object): Following classes (_SparseColumn, _RealValuedColumn, ...) are concrete instances. """ - __metaclass__ = abc.ABCMeta @abc.abstractproperty @deprecation.deprecated( @@ -1015,8 +1015,7 @@ class _OneHotColumn( dense_id_tensor, depth=self.length, on_value=1.0, off_value=0.0) # Reduce to get a multi-hot per example. - return math_ops.reduce_sum( - one_hot_id_tensor, reduction_indices=[output_rank - 1]) + return math_ops.reduce_sum(one_hot_id_tensor, axis=[output_rank - 1]) @property def _variable_shape(self): 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 6fb4b9ff3534cab34c84de5d13fea7aff756556d..00e41026d0038409ace178e6affd2c1cdc812122 100644 --- a/tensorflow/contrib/layers/python/layers/feature_column_ops_test.py +++ b/tensorflow/contrib/layers/python/layers/feature_column_ops_test.py @@ -27,7 +27,7 @@ from tensorflow.contrib.layers.python.layers import feature_column from tensorflow.contrib.layers.python.layers import feature_column_ops from tensorflow.core.example import example_pb2 from tensorflow.core.example import feature_pb2 -from tensorflow.python.feature_column import feature_column as fc_core +from tensorflow.python.feature_column import feature_column_lib as fc_core from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops @@ -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/feature_column_test.py b/tensorflow/contrib/layers/python/layers/feature_column_test.py index d90d6ecf7f671a40a7ff2b066b6782c7421f9887..cab8da808b6413518ff4864cb0b03a42809260f1 100644 --- a/tensorflow/contrib/layers/python/layers/feature_column_test.py +++ b/tensorflow/contrib/layers/python/layers/feature_column_test.py @@ -27,7 +27,7 @@ import numpy as np from tensorflow.contrib.layers.python.layers import feature_column as fc from tensorflow.contrib.layers.python.layers import feature_column_ops -from tensorflow.python.feature_column import feature_column as fc_core +from tensorflow.python.feature_column import feature_column_lib as fc_core from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import sparse_tensor as sparse_tensor_lib diff --git a/tensorflow/contrib/layers/python/layers/layers.py b/tensorflow/contrib/layers/python/layers/layers.py index a82d4c19510df2c6dcc3cdac2a808823795149a9..403b522ce45ac6ad98a321378626b87aaa7738aa 100644 --- a/tensorflow/contrib/layers/python/layers/layers.py +++ b/tensorflow/contrib/layers/python/layers/layers.py @@ -35,6 +35,7 @@ from tensorflow.python.framework import function from tensorflow.python.framework import ops from tensorflow.python.framework import sparse_tensor from tensorflow.python.framework import tensor_shape +from tensorflow.python.keras.engine import input_spec from tensorflow.python.layers import base from tensorflow.python.layers import convolutional as convolutional_layers from tensorflow.python.layers import core as core_layers @@ -274,7 +275,7 @@ def _fused_batch_norm(inputs, ' Expected 2 or 4 but got %d' % (inputs.name, original_rank)) if original_rank == 2: - channels = inputs.get_shape()[-1].value + channels = inputs.get_shape().dims[-1].value if channels is None: raise ValueError('`C` dimension must be known but is None') new_shape = [-1, 1, 1, channels] @@ -692,7 +693,7 @@ def batch_norm(inputs, # explicitly reshape the params to params_shape_broadcast when computing # the moments and the batch normalization. params_shape_broadcast = list( - [1, inputs_shape[1].value] + [1 for _ in range(2, inputs_rank)]) + [1, inputs_shape.dims[1].value] + [1 for _ in range(2, inputs_rank)]) else: moments_axes = list(range(inputs_rank - 1)) params_shape = inputs_shape[-1:] @@ -890,7 +891,7 @@ def bias_add(inputs, elif inputs_rank != 4 and data_format == DATA_FORMAT_NCHW: raise ValueError('Data format NCHW only supports 4D Tensor') axis = 1 if data_format == DATA_FORMAT_NCHW else -1 - num_features = inputs_shape[axis].value + num_features = inputs_shape.dims[axis].value if num_features is None: raise ValueError('`C` dimension must be known but is None') biases_collections = utils.get_variable_collections(variables_collections, @@ -1823,8 +1824,8 @@ def fully_connected(inputs, ValueError: If x has rank less than 2 or if its last dimension is not set. """ if not isinstance(num_outputs, six.integer_types): - raise ValueError('num_outputs should be int or long, got %s.' % - (num_outputs,)) + raise ValueError('num_outputs type should be one of %s, got %s.' % ( + list(six.integer_types), type(num_outputs))) layer_variable_getter = _build_variable_getter({ 'bias': 'biases', @@ -1958,7 +1959,7 @@ class GDN(base.Layer): self._reparam_offset = reparam_offset self.data_format = data_format self._channel_axis() # trigger ValueError early - self.input_spec = base.InputSpec(min_ndim=3, max_ndim=5) + self.input_spec = input_spec.InputSpec(min_ndim=3, max_ndim=5) def _channel_axis(self): try: @@ -2010,12 +2011,12 @@ class GDN(base.Layer): def build(self, input_shape): channel_axis = self._channel_axis() input_shape = tensor_shape.TensorShape(input_shape) - num_channels = input_shape[channel_axis].value + num_channels = input_shape.dims[channel_axis].value if num_channels is None: raise ValueError('The channel dimension of the inputs to `GDN` ' 'must be defined.') self._input_rank = input_shape.ndims - self.input_spec = base.InputSpec( + self.input_spec = input_spec.InputSpec( ndim=input_shape.ndims, axes={ channel_axis: num_channels }) @@ -2100,7 +2101,7 @@ class GDN(base.Layer): input_shape = tensor_shape.TensorShape(input_shape) if not 3 <= input_shape.ndim <= 5: raise ValueError('`input_shape` must be of rank 3 to 5, inclusive.') - if input_shape[channel_axis].value is None: + if input_shape.dims[channel_axis].value is None: raise ValueError( 'The channel dimension of `input_shape` must be defined.') return input_shape @@ -2951,7 +2952,7 @@ def spatial_softmax(features, num_channels, height, width = static_shape[1], shape[2], shape[3] else: raise ValueError('data_format has to be either NCHW or NHWC.') - if num_channels.value is None: + if tensor_shape.dimension_value(num_channels) is None: raise ValueError('The num_channels dimension of the inputs to ' '`spatial_softmax` should be defined. Found `None`.') @@ -2994,9 +2995,11 @@ def spatial_softmax(features, expected_y = math_ops.reduce_sum( pos_y * softmax_attention, [1], keepdims=True) expected_xy = array_ops.concat([expected_x, expected_y], 1) - feature_keypoints = array_ops.reshape(expected_xy, - [-1, num_channels.value * 2]) - feature_keypoints.set_shape([None, num_channels.value * 2]) + feature_keypoints = array_ops.reshape( + expected_xy, + [-1, tensor_shape.dimension_value(num_channels) * 2]) + feature_keypoints.set_shape( + [None, tensor_shape.dimension_value(num_channels) * 2]) return feature_keypoints diff --git a/tensorflow/contrib/layers/python/layers/layers_test.py b/tensorflow/contrib/layers/python/layers/layers_test.py index 8ead6336a08db4dd52edf0d3372db5a50f860e2b..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()) @@ -1459,13 +1459,6 @@ class DropoutTest(test.TestCase): class FlattenTest(test.TestCase): - def testInvalidRank(self): - with ops.Graph().as_default() as g, self.session(g): - inputs = array_ops.placeholder(dtype=dtypes.float32) - inputs.set_shape(tensor_shape.TensorShape((5,))) - with self.assertRaisesRegexp(ValueError, 'incompatible with the layer'): - _layers.flatten(inputs) - def testUnknownLastDim(self): with ops.Graph().as_default() as g, self.session(g): inputs = array_ops.placeholder(dtype=dtypes.float32) @@ -1502,6 +1495,12 @@ class FlattenTest(test.TestCase): images.get_shape().num_elements()) self.assertEqual(output.get_shape()[0], images.get_shape()[0]) + def testFlatten0D(self): + with self.cached_session(): + scalars = random_ops.random_uniform((5,), seed=1, name='scalars') + output = _layers.flatten(scalars) + self.assertEqual(output.shape, (5, 1)) + def testFlattenBatchSize(self): height, width = 3, 3 with self.cached_session() as sess: @@ -3811,7 +3810,7 @@ class UnitNormTests(test.TestCase): image = random_ops.random_uniform((height, width, 3)) output = _layers.unit_norm(image, dim=dim, epsilon=1e-6) norms = math_ops.sqrt( - math_ops.reduce_sum(math_ops.square(output), reduction_indices=dim)) + math_ops.reduce_sum(math_ops.square(output), axis=dim)) shape = [height, width, 3] del shape[dim] @@ -3847,7 +3846,7 @@ class UnitNormTests(test.TestCase): image = array_ops.placeholder(dtypes.float32, (None, None, 3)) output = _layers.unit_norm(image, dim=dim, epsilon=1e-6) norms = math_ops.sqrt( - math_ops.reduce_sum(math_ops.square(output), reduction_indices=dim)) + math_ops.reduce_sum(math_ops.square(output), axis=dim)) with self.cached_session(): actual = norms.eval({image: placeholder_value}) 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/regularizers_test.py b/tensorflow/contrib/layers/python/layers/regularizers_test.py index 51faba30c74d64c54d3d2b11d2a11195cca6b759..5cb00b76847430be8ade9f4e4fc8f7372035485a 100644 --- a/tensorflow/contrib/layers/python/layers/regularizers_test.py +++ b/tensorflow/contrib/layers/python/layers/regularizers_test.py @@ -141,7 +141,7 @@ class RegularizerTest(test.TestCase): dummy_regularizer = lambda x: math_ops.reduce_sum(2 * x) array_weights_list = [[1.5], [2, 3, 4.2], [10, 42, 666.6]] tensor_weights_list = [constant_op.constant(x) for x in array_weights_list] - expected = sum([2 * x for l in array_weights_list for x in l]) + expected = sum(2 * x for l in array_weights_list for x in l) with self.cached_session(): result = regularizers.apply_regularization(dummy_regularizer, tensor_weights_list) diff --git a/tensorflow/contrib/learn/BUILD b/tensorflow/contrib/learn/BUILD index 61185f65a9bd294003515456f891de0a68661a82..4749371248ee89a033912132986d7f76c85dbaa6 100644 --- a/tensorflow/contrib/learn/BUILD +++ b/tensorflow/contrib/learn/BUILD @@ -24,6 +24,11 @@ py_library( exclude = ["python/learn/**/*_test.py"], ), srcs_version = "PY2AND3", + visibility = [ + "//learning/brain:__subpackages__", + "//tensorflow:__subpackages__", + "//video/youtube/personalization:__subpackages__", + ], # This library should not depend on sklearn, even though some of the code # refers to it. (The code handles the presence of sklearn conditionally.) deps = [ @@ -269,6 +274,7 @@ py_test( name = "estimator_test", size = "medium", srcs = ["python/learn/estimators/estimator_test.py"], + shard_count = 2, srcs_version = "PY2AND3", tags = [ "manual", @@ -351,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 d516bffc5e0327a3400068b35de5503e5a925a54..b2d3a6273abba7e3a893f30bbdd4f8b2662bd54a 100644 --- a/tensorflow/contrib/learn/README.md +++ b/tensorflow/contrib/learn/README.md @@ -7,12 +7,11 @@ warnings. A high-level overview is below. ## Canned Estimators -Many canned estimators (subclasses of `Estimator`) have equivalents in core: +Many canned estimators (subclasses of `Estimator`) have equivalents in core +exposed under `tf.estimator`: `DNNClassifier`, `DNNRegressor`, `DNNEstimator`, `LinearClassifier`, -`LinearRegressor`, `DNNLinearCombinedClassifier` and -`DNNLinearCombinedRegressor`. They are exposed under `tf.estimator`. -`DNNEstimator`, `LinearEstimator` and `DNNLinearCombinedEstimator` -are exposed under `tf.contrib.estimator`. +`LinearRegressor`, `LinearEstimator`, `DNNLinearCombinedClassifier`, +`DNNLinearCombinedRegressor` and `DNNLinearCombinedEstimator`. To migrate to the new api, users need to take the following steps: @@ -45,7 +44,7 @@ To migrate to the new api, users need to take the following steps: `tf.contrib.learn` classifiers and regressors supported labels with shape `[batch_size]`. * If you pass custom metrics from the `evaluate()` method call, use - `tf.contrib.estimator.add_metrics`. + `tf.estimator.add_metrics`. * Replace your `serving_input_fn` with a `serving_input_receiver_fn`. Note this should be entirely distinct from your training `input_fn`, so if you previously had one `input_fn` with different "modes", you should now factor @@ -63,10 +62,10 @@ Some remaining estimators/classes: with a custom `model_fn`, or with `DNNEstimator`. * `StateSavingRnnEstimator`: Consider a custom `model_fn`. * SVM: Consider a custom `model_fn`. -* `LinearComposableModel` and `DNNComposableModel`: Not supported. +* `LinearComposableModel` and `DNNComposableModel`: Not supported. Consider `tf.contrib.estimator.DNNEstimator`, or write a custom model_fn. * `MetricSpec`: Deprecated. For adding custom metrics to canned Estimators, use - `tf.contrib.estimator.add_metrics`. + `tf.estimator.add_metrics`. ## Estimator `tf.contrib.learn.Estimator` is migrated to `tf.estimator.Estimator`. @@ -112,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/__init__.py b/tensorflow/contrib/learn/__init__.py index 28a6f5aed99b1443ebcc9c391ec332e0febbb04b..7bf2ac62d76d67f0eb131f8f57c5c063955424fa 100644 --- a/tensorflow/contrib/learn/__init__.py +++ b/tensorflow/contrib/learn/__init__.py @@ -19,9 +19,6 @@ This module and all its submodules are deprecated. See [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) for migration instructions. -See the [Contrib Learn](https://tensorflow.org/api_guides/python/contrib.learn) -guide. - @@BaseEstimator @@Estimator @@Trainable diff --git a/tensorflow/contrib/learn/python/learn/estimators/_sklearn.py b/tensorflow/contrib/learn/python/learn/estimators/_sklearn.py index 1f0e4663d060a3850e2002b27f809fde1db47e48..4c206839300b1c6b14b324b3d1ec2d70f7eca903 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/_sklearn.py +++ b/tensorflow/contrib/learn/python/learn/estimators/_sklearn.py @@ -194,7 +194,7 @@ if TRY_IMPORT_SKLEARN: # pylint: disable=g-import-not-at-top,g-multiple-import,unused-import from sklearn.base import BaseEstimator, ClassifierMixin, RegressorMixin, TransformerMixin from sklearn.metrics import accuracy_score, log_loss, mean_squared_error - from sklearn.cross_validation import train_test_split + from sklearn.model_selection import train_test_split try: from sklearn.exceptions import NotFittedError except ImportError: diff --git a/tensorflow/contrib/learn/python/learn/estimators/dnn.py b/tensorflow/contrib/learn/python/learn/estimators/dnn.py index eabebb7e881558471c343c0573cc9a8f4a425312..10fbd60ba2df4c3f84169bf04f249d67dc14573f 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/dnn.py +++ b/tensorflow/contrib/learn/python/learn/estimators/dnn.py @@ -28,7 +28,6 @@ import six from tensorflow.contrib import layers from tensorflow.contrib.framework import deprecated from tensorflow.contrib.framework import deprecated_arg_values -from tensorflow.python.training import training_util from tensorflow.contrib.layers.python.layers import feature_column from tensorflow.contrib.layers.python.layers import optimizers from tensorflow.contrib.learn.python.learn import metric_spec @@ -38,11 +37,12 @@ from tensorflow.contrib.learn.python.learn.estimators import head as head_lib from tensorflow.contrib.learn.python.learn.estimators import model_fn from tensorflow.contrib.learn.python.learn.estimators import prediction_key from tensorflow.contrib.learn.python.learn.utils import export -from tensorflow.python.feature_column import feature_column as fc_core +from tensorflow.python.feature_column import feature_column_lib as fc_core from tensorflow.python.ops import nn from tensorflow.python.ops import partitioned_variables from tensorflow.python.ops import variable_scope from tensorflow.python.summary import summary +from tensorflow.python.training import training_util # The default learning rate of 0.05 is a historical artifact of the initial # implementation, but seems a reasonable choice. @@ -150,10 +150,10 @@ def _dnn_model_fn(features, labels, mode, params, config=None): "input_from_feature_columns", values=tuple(six.itervalues(features)), partitioner=input_layer_partitioner) as input_layer_scope: - if all([ + if all( isinstance(fc, feature_column._FeatureColumn) # pylint: disable=protected-access for fc in feature_columns - ]): + ): net = layers.input_from_feature_columns( columns_to_tensors=features, feature_columns=feature_columns, diff --git a/tensorflow/contrib/learn/python/learn/estimators/dnn_linear_combined.py b/tensorflow/contrib/learn/python/learn/estimators/dnn_linear_combined.py index 3d85533d92d17095bae9a69f229171e1bf61ba10..2ade6b7b6ce2678ec8df7c98ffaa5636ae9d4b1d 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/dnn_linear_combined.py +++ b/tensorflow/contrib/learn/python/learn/estimators/dnn_linear_combined.py @@ -38,7 +38,7 @@ from tensorflow.contrib.learn.python.learn.estimators import head as head_lib from tensorflow.contrib.learn.python.learn.estimators import model_fn from tensorflow.contrib.learn.python.learn.estimators import prediction_key from tensorflow.contrib.learn.python.learn.utils import export -from tensorflow.python.feature_column import feature_column as fc_core +from tensorflow.python.feature_column import feature_column_lib as fc_core from tensorflow.python.framework import ops from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import nn @@ -236,10 +236,10 @@ def _dnn_linear_combined_model_fn(features, labels, mode, params, config=None): "input_from_feature_columns", values=tuple(six.itervalues(features)), partitioner=input_layer_partitioner) as dnn_input_scope: - if all([ + if all( isinstance(fc, feature_column_lib._FeatureColumn) # pylint: disable=protected-access for fc in dnn_feature_columns - ]): + ): net = layers.input_from_feature_columns( columns_to_tensors=features, feature_columns=dnn_feature_columns, @@ -292,8 +292,8 @@ def _dnn_linear_combined_model_fn(features, labels, mode, params, config=None): linear_parent_scope, values=tuple(six.itervalues(features)), partitioner=linear_partitioner) as scope: - if all([isinstance(fc, feature_column_lib._FeatureColumn) # pylint: disable=protected-access - for fc in linear_feature_columns]): + if all(isinstance(fc, feature_column_lib._FeatureColumn) # pylint: disable=protected-access + for fc in linear_feature_columns): if joint_linear_weights: linear_logits, _, _ = layers.joint_weighted_sum_from_feature_columns( columns_to_tensors=features, diff --git a/tensorflow/contrib/learn/python/learn/estimators/dnn_linear_combined_test.py b/tensorflow/contrib/learn/python/learn/estimators/dnn_linear_combined_test.py index 4e65c180d8bee9ab8fe9b1fbf32edc229c31af09..d46a873bfaa297e7f6242aa56e9d0bf0eb551867 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/dnn_linear_combined_test.py +++ b/tensorflow/contrib/learn/python/learn/estimators/dnn_linear_combined_test.py @@ -36,7 +36,7 @@ from tensorflow.contrib.learn.python.learn.estimators import run_config from tensorflow.contrib.learn.python.learn.estimators import test_data from tensorflow.contrib.learn.python.learn.metric_spec import MetricSpec from tensorflow.contrib.metrics.python.ops import metric_ops -from tensorflow.python.feature_column import feature_column as fc_core +from tensorflow.python.feature_column import feature_column_lib as fc_core from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops diff --git a/tensorflow/contrib/learn/python/learn/estimators/dnn_test.py b/tensorflow/contrib/learn/python/learn/estimators/dnn_test.py index 2bd57597c2e9444b51b1dacfbe4180b443c95a3d..ee25cebd484f1e831fe8b6d3aa7290da7558adee 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/dnn_test.py +++ b/tensorflow/contrib/learn/python/learn/estimators/dnn_test.py @@ -38,7 +38,7 @@ from tensorflow.contrib.learn.python.learn.estimators import run_config from tensorflow.contrib.learn.python.learn.estimators import test_data from tensorflow.contrib.learn.python.learn.metric_spec import MetricSpec from tensorflow.contrib.metrics.python.ops import metric_ops -from tensorflow.python.feature_column import feature_column as fc_core +from tensorflow.python.feature_column import feature_column_lib as fc_core from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import sparse_tensor 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 1d8a59281a4934ad063362cba064e6cb3abff5a2..28c4964527bb034c8c6b1642366c6c82c1a72201 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 @@ -668,7 +668,7 @@ class DynamicRNNEstimatorLearningTest(test.TestCase): sequences = centers + noise inputs = array_ops.expand_dims(sequences, 2) - labels = math_ops.reduce_mean(sequences, reduction_indices=[1]) + labels = math_ops.reduce_mean(sequences, axis=[1]) return {'inputs': inputs}, labels return input_fn @@ -722,8 +722,8 @@ class DynamicRNNEstimatorLearningTest(test.TestCase): inputs = array_ops.expand_dims(math_ops.to_float(random_sequence), 2) labels = math_ops.to_int32( array_ops.squeeze( - math_ops.reduce_sum( - inputs, reduction_indices=[1]) > (sequence_length / 2.0))) + math_ops.reduce_sum(inputs, axis=[1]) > ( + sequence_length / 2.0))) return {'inputs': inputs}, labels return input_fn diff --git a/tensorflow/contrib/learn/python/learn/estimators/estimator.py b/tensorflow/contrib/learn/python/learn/estimators/estimator.py index 3efceab3375d3a1801c87122c98920cc523a3aca..9132b2209bce8005b323d058d6d176784a84b2d1 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/estimator.py +++ b/tensorflow/contrib/learn/python/learn/estimators/estimator.py @@ -404,7 +404,6 @@ class BaseEstimator(sklearn.BaseEstimator, evaluable.Evaluable, Users should not instantiate or subclass this class. Instead, use an `Estimator`. """ - __metaclass__ = abc.ABCMeta # Note that for Google users, this is overridden with # learn_runner.EstimatorConfig. @@ -1067,11 +1066,11 @@ class BaseEstimator(sklearn.BaseEstimator, evaluable.Evaluable, chief_hooks = [] if (self._config.save_checkpoints_secs or self._config.save_checkpoints_steps): - saver_hook_exists = any([ + saver_hook_exists = any( isinstance(h, basic_session_run_hooks.CheckpointSaverHook) for h in (all_hooks + model_fn_ops.training_hooks + chief_hooks + model_fn_ops.training_chief_hooks) - ]) + ) if not saver_hook_exists: chief_hooks = [ basic_session_run_hooks.CheckpointSaverHook( @@ -1494,7 +1493,7 @@ class Estimator(BaseEstimator): # pylint: disable=protected-access class SKCompat(sklearn.BaseEstimator): """Scikit learn wrapper for TensorFlow Learn Estimator. - + THIS CLASS IS DEPRECATED. See [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) for general migration instructions. diff --git a/tensorflow/contrib/learn/python/learn/estimators/estimator_test.py b/tensorflow/contrib/learn/python/learn/estimators/estimator_test.py index 9e5aaf3118dfed4ce64dd244a915860b5a2eef44..8a461a0bd7ba457fcf830769f23c6ca2860a2732 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/estimator_test.py +++ b/tensorflow/contrib/learn/python/learn/estimators/estimator_test.py @@ -1368,7 +1368,7 @@ class ReplicaDeviceSetterTest(test.TestCase): table = lookup.MutableHashTable(dtypes.string, dtypes.int64, default_val) input_string = constant_op.constant(['brain', 'salad', 'tank']) output = table.lookup(input_string) - self.assertDeviceEqual('/job:ps/task:0', table._table_ref.device) + self.assertDeviceEqual('/job:ps/task:0', table.resource_handle.device) self.assertDeviceEqual('/job:ps/task:0', output.device) def testMutableHashTableIsLocal(self): @@ -1378,7 +1378,7 @@ class ReplicaDeviceSetterTest(test.TestCase): table = lookup.MutableHashTable(dtypes.string, dtypes.int64, default_val) input_string = constant_op.constant(['brain', 'salad', 'tank']) output = table.lookup(input_string) - self.assertDeviceEqual('', table._table_ref.device) + self.assertDeviceEqual('', table.resource_handle.device) self.assertDeviceEqual('', output.device) def testTaskIsSetOnWorkerWhenJobNameIsSet(self): diff --git a/tensorflow/contrib/learn/python/learn/estimators/head.py b/tensorflow/contrib/learn/python/learn/estimators/head.py index c6f79e00d5a5a584b0c5f8201a2576f02106a5b4..c1b97d8b49613ea49d9813954da3b7a63d3ba04c 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/head.py +++ b/tensorflow/contrib/learn/python/learn/estimators/head.py @@ -55,6 +55,7 @@ from tensorflow.python.util import tf_inspect from tensorflow.python.util.deprecation import deprecated +@six.add_metaclass(abc.ABCMeta) class Head(object): """Interface for the head/top of a model. @@ -132,7 +133,6 @@ class Head(object): ... update train_op and hooks in ModelFnOps and return ``` """ - __metaclass__ = abc.ABCMeta @abc.abstractproperty def logits_dimension(self): @@ -504,7 +504,6 @@ def no_op_train_fn(loss): class _SingleHead(Head): """Interface for a single head/top of a model.""" - __metaclass__ = abc.ABCMeta def __init__( self, problem_type, logits_dimension, label_name=None, diff --git a/tensorflow/contrib/learn/python/learn/estimators/linear.py b/tensorflow/contrib/learn/python/learn/estimators/linear.py index e100bc7a1e7be4896e9ab1c965775b5185b38897..9ee8d8004bf26224dd96a98bad109720c44d04f7 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/linear.py +++ b/tensorflow/contrib/learn/python/learn/estimators/linear.py @@ -37,7 +37,7 @@ from tensorflow.contrib.learn.python.learn.estimators import head as head_lib from tensorflow.contrib.learn.python.learn.estimators import prediction_key from tensorflow.contrib.learn.python.learn.utils import export from tensorflow.contrib.linear_optimizer.python import sdca_optimizer -from tensorflow.python.feature_column import feature_column as fc_core +from tensorflow.python.feature_column import feature_column_lib as fc_core from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.framework import sparse_tensor @@ -155,8 +155,8 @@ def _linear_model_fn(features, labels, mode, params, config=None): parent_scope, values=tuple(six.itervalues(features)), partitioner=partitioner) as scope: - if all([isinstance(fc, feature_column._FeatureColumn) # pylint: disable=protected-access - for fc in feature_columns]): + if all(isinstance(fc, feature_column._FeatureColumn) # pylint: disable=protected-access + for fc in feature_columns): if joint_weights: layer_fn = layers.joint_weighted_sum_from_feature_columns else: diff --git a/tensorflow/contrib/learn/python/learn/estimators/linear_test.py b/tensorflow/contrib/learn/python/learn/estimators/linear_test.py index 597ca4e86dbf66c86182f14a2a364b662d52fb0a..dfc76bfde6c0109f98093232b6f223d6938007f9 100644 --- a/tensorflow/contrib/learn/python/learn/estimators/linear_test.py +++ b/tensorflow/contrib/learn/python/learn/estimators/linear_test.py @@ -37,7 +37,7 @@ from tensorflow.contrib.learn.python.learn.estimators import test_data from tensorflow.contrib.learn.python.learn.metric_spec import MetricSpec from tensorflow.contrib.linear_optimizer.python import sdca_optimizer as sdca_optimizer_lib from tensorflow.contrib.metrics.python.ops import metric_ops -from tensorflow.python.feature_column import feature_column as fc_core +from tensorflow.python.feature_column import feature_column_lib as fc_core from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import sparse_tensor @@ -1745,7 +1745,7 @@ class LinearRegressorTest(test.TestCase): 'place_holder': constant_op.constant([[0.0]] * num_examples), }, constant_op.constant( - [[1 if i % 4 is 0 else 0] for i in range(num_examples)]) + [[1 if i % 4 == 0 else 0] for i in range(num_examples)]) place_holder = feature_column_lib.real_valued_column('place_holder') sdca_optimizer = sdca_optimizer_lib.SDCAOptimizer( diff --git a/tensorflow/contrib/learn/python/learn/evaluable.py b/tensorflow/contrib/learn/python/learn/evaluable.py index 10881ca885599bc81386e15f814a2687d907f63b..5dedf548f73d27bf543dcfd9885490b8b7c9ac96 100644 --- a/tensorflow/contrib/learn/python/learn/evaluable.py +++ b/tensorflow/contrib/learn/python/learn/evaluable.py @@ -25,7 +25,10 @@ from __future__ import print_function import abc +import six + +@six.add_metaclass(abc.ABCMeta) class Evaluable(object): """Interface for objects that are evaluatable by, e.g., `Experiment`. @@ -33,7 +36,6 @@ class Evaluable(object): [contrib/learn/README.md](https://www.tensorflow.org/code/tensorflow/contrib/learn/README.md) for general migration instructions. """ - __metaclass__ = abc.ABCMeta @abc.abstractproperty def model_dir(self): 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/learn_io/numpy_io.py b/tensorflow/contrib/learn/python/learn/learn_io/numpy_io.py index 29552d24f1eaa0d85a99c8b09f69d007e7e4fe9f..59a67636ae275c5ca1df21685770baa7a960d667 100644 --- a/tensorflow/contrib/learn/python/learn/learn_io/numpy_io.py +++ b/tensorflow/contrib/learn/python/learn/learn_io/numpy_io.py @@ -27,7 +27,7 @@ from tensorflow.python.estimator.inputs.numpy_io import numpy_input_fn as core_n from tensorflow.python.util.deprecation import deprecated -@deprecated(None, 'Use tf.estimator.inputs.numpy_input_fn.') +@deprecated(None, 'Use tf.compat.v1.estimator.inputs.numpy_input_fn.') def numpy_input_fn(x, y=None, batch_size=128, diff --git a/tensorflow/contrib/learn/python/learn/learn_io/pandas_io.py b/tensorflow/contrib/learn/python/learn/learn_io/pandas_io.py index b4ef055f5ae484ec704ad42efcf2c00c4a7a4f56..e9df7258a358d9543f2bb386518d900bd6ddef74 100644 --- a/tensorflow/contrib/learn/python/learn/learn_io/pandas_io.py +++ b/tensorflow/contrib/learn/python/learn/learn_io/pandas_io.py @@ -53,7 +53,7 @@ PANDAS_DTYPES = { } -@deprecated(None, 'Please use tf.estimator.inputs.pandas_input_fn') +@deprecated(None, 'Please use tf.compat.v1.estimator.inputs.pandas_input_fn') def pandas_input_fn(x, y=None, batch_size=128, diff --git a/tensorflow/contrib/learn/python/learn/trainable.py b/tensorflow/contrib/learn/python/learn/trainable.py index a1a3f20dcd8cb5ff7baa559ac41d5e5c40780511..1ea9e5d67a95dfc3ba57151085051ec7aea14226 100644 --- a/tensorflow/contrib/learn/python/learn/trainable.py +++ b/tensorflow/contrib/learn/python/learn/trainable.py @@ -25,13 +25,15 @@ from __future__ import print_function import abc +import six + +@six.add_metaclass(abc.ABCMeta) class Trainable(object): """Interface for objects that are trainable by, e.g., `Experiment`. THIS CLASS IS DEPRECATED. """ - __metaclass__ = abc.ABCMeta @abc.abstractmethod def fit(self, 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/linear_optimizer/python/kernel_tests/sdca_ops_test.py b/tensorflow/contrib/linear_optimizer/python/kernel_tests/sdca_ops_test.py index 8466dc36d13e223aed4f1dfe8e39a6f91c99fa55..d49834dc860a8b4341ddd3720fde52281f7474f7 100644 --- a/tensorflow/contrib/linear_optimizer/python/kernel_tests/sdca_ops_test.py +++ b/tensorflow/contrib/linear_optimizer/python/kernel_tests/sdca_ops_test.py @@ -12,7 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Tests for SdcaModel.""" +"""Tests for SdcaModel (deprecated). + +This module and all its submodules are deprecated. To UPDATE or USE linear +optimizers, please check its latest version in core: +tensorflow_estimator/python/estimator/canned/linear_optimizer/. +""" from __future__ import absolute_import from __future__ import division diff --git a/tensorflow/contrib/linear_optimizer/python/ops/sdca_ops.py b/tensorflow/contrib/linear_optimizer/python/ops/sdca_ops.py index 94ff1dd5b01b5a24d1deb7053553b9df48709c7c..c056a12fa5307a7e9ac4cf30e1386ddfd5cd7d75 100644 --- a/tensorflow/contrib/linear_optimizer/python/ops/sdca_ops.py +++ b/tensorflow/contrib/linear_optimizer/python/ops/sdca_ops.py @@ -12,7 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Proximal stochastic dual coordinate ascent optimizer for linear models.""" +# pylint: disable=line-too-long +"""Proximal stochastic dual coordinate ascent optimizer for linear models (deprecated). + +This module and all its submodules are deprecated. To UPDATE or USE linear +optimizers, please check its latest version in core: +tensorflow_estimator/python/estimator/canned/linear_optimizer/. +""" +# pylint: enable=line-too-long from __future__ import absolute_import from __future__ import division from __future__ import print_function @@ -26,6 +33,7 @@ from tensorflow.python.compat import compat from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops +from tensorflow.python.framework import tensor_shape from tensorflow.python.framework.ops import internal_convert_to_tensor from tensorflow.python.framework.ops import name_scope from tensorflow.python.ops import array_ops @@ -39,6 +47,7 @@ from tensorflow.python.ops import variables as var_ops from tensorflow.python.ops.nn import log_poisson_loss from tensorflow.python.ops.nn import sigmoid_cross_entropy_with_logits from tensorflow.python.summary import summary +from tensorflow.python.util import deprecation __all__ = ['SdcaModel'] @@ -47,7 +56,7 @@ __all__ = ['SdcaModel'] class SdcaModel(object): """Stochastic dual coordinate ascent solver for linear models. - Loss functions supported: + Loss functions supported: * Binary logistic loss * Squared loss @@ -108,6 +117,10 @@ class SdcaModel(object): ``` """ + @deprecation.deprecated( + None, 'This class is deprecated. To UPDATE or USE linear optimizers, ' + 'please check its latest version in core: ' + 'tensorflow_estimator/python/estimator/canned/linear_optimizer/.') def __init__(self, examples, variables, options): """Create a new sdca optimizer.""" @@ -427,14 +440,15 @@ class SdcaModel(object): dim_0_size = self._get_first_dimension_size_statically( w, num_partitions) - if dim_0_size.value: - num_total_ids = constant_op.constant(dim_0_size.value, - flat_ids.dtype) + if tensor_shape.dimension_value(dim_0_size): + num_total_ids = constant_op.constant( + tensor_shape.dimension_value(dim_0_size), + flat_ids.dtype) else: dim_0_sizes = [] for p in range(num_partitions): - if w[p].get_shape()[0].value is not None: - dim_0_sizes.append(w[p].get_shape()[0].value) + if tensor_shape.dimension_value(w[p].shape[0]) is not None: + dim_0_sizes.append(tensor_shape.dimension_value(w[p].shape[0])) else: with ops.colocate_with(w[p]): dim_0_sizes.append(array_ops.shape(w[p])[0]) diff --git a/tensorflow/contrib/linear_optimizer/python/ops/sharded_mutable_dense_hashtable.py b/tensorflow/contrib/linear_optimizer/python/ops/sharded_mutable_dense_hashtable.py index 44a869f7c2745c594b6a4ea69a2a9e6f1b4f780a..a28394964a12013c43d85701b5a0ab5c559afd62 100644 --- a/tensorflow/contrib/linear_optimizer/python/ops/sharded_mutable_dense_hashtable.py +++ b/tensorflow/contrib/linear_optimizer/python/ops/sharded_mutable_dense_hashtable.py @@ -12,7 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Sharded mutable dense hash table.""" +"""Sharded mutable dense hash table (deprecated). + +This module and all its submodules are deprecated. To UPDATE or USE linear +optimizers, please check its latest version in core: +tensorflow_estimator/python/estimator/canned/linear_optimizer/. +""" from __future__ import absolute_import from __future__ import division @@ -28,9 +33,12 @@ from tensorflow.python.ops import array_ops from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import data_flow_ops from tensorflow.python.ops import math_ops +from tensorflow.python.util import deprecation -class ShardedMutableDenseHashTable(lookup.LookupInterface): +# TODO(rohanj): This should subclass Checkpointable and implement +# _gather_saveables_for_checkpoint. +class ShardedMutableDenseHashTable(object): """A sharded version of MutableDenseHashTable. It is designed to be interface compatible with LookupInterface and @@ -43,6 +51,10 @@ class ShardedMutableDenseHashTable(lookup.LookupInterface): # TODO(andreasst): consider moving this to lookup module + @deprecation.deprecated( + None, 'This class is deprecated. To UPDATE or USE linear optimizers, ' + 'please check its latest version in core: ' + 'tensorflow_estimator/python/estimator/canned/linear_optimizer/.') def __init__(self, key_dtype, value_dtype, @@ -52,9 +64,10 @@ class ShardedMutableDenseHashTable(lookup.LookupInterface): num_shards=1, checkpoint=True, name='ShardedMutableHashTable'): + self._key_dtype = key_dtype + self._value_dtype = value_dtype with ops.name_scope(name, 'sharded_mutable_hash_table') as scope: - super(ShardedMutableDenseHashTable, self).__init__(key_dtype, - value_dtype, scope) + self._table_name = scope table_shards = [] for i in range(num_shards): table_shards.append( @@ -72,6 +85,10 @@ class ShardedMutableDenseHashTable(lookup.LookupInterface): self._value_shape = self._table_shards[0]._value_shape # pylint: enable=protected-access + @property + def name(self): + return self._table_name + @property def _num_shards(self): return len(self._table_shards) @@ -92,7 +109,7 @@ class ShardedMutableDenseHashTable(lookup.LookupInterface): if key_shape.ndims > 1: # If keys are a matrix (i.e. a single key is a vector), we use the first # element of each key vector to determine the shard. - keys = array_ops.slice(keys, [0, 0], [key_shape[0].value, 1]) + keys = array_ops.slice(keys, [0, 0], [key_shape.dims[0].value, 1]) keys = array_ops.reshape(keys, [-1]) indices = math_ops.mod(math_ops.abs(keys), self._num_shards) return math_ops.cast(indices, dtypes.int32) @@ -106,6 +123,7 @@ class ShardedMutableDenseHashTable(lookup.LookupInterface): keys.get_shape()) def lookup(self, keys, name=None): + """Looks up `keys` in a table, outputs the corresponding values.""" if keys.dtype.base_dtype != self._key_dtype: raise TypeError('Signature mismatch. Keys must be dtype %s, got %s.' % (self._key_dtype, keys.dtype)) @@ -134,6 +152,7 @@ class ShardedMutableDenseHashTable(lookup.LookupInterface): return result def insert(self, keys, values, name=None): + """Inserts `keys` in a table.""" self._check_keys(keys) num_shards = self._num_shards if num_shards == 1: diff --git a/tensorflow/contrib/linear_optimizer/python/ops/sharded_mutable_dense_hashtable_test.py b/tensorflow/contrib/linear_optimizer/python/ops/sharded_mutable_dense_hashtable_test.py index 2b56d0fa3a8b8564b7c73a62bd99cc900d6f5c54..2d1457f9e4cc576da696be191e718814dd9ff4e5 100644 --- a/tensorflow/contrib/linear_optimizer/python/ops/sharded_mutable_dense_hashtable_test.py +++ b/tensorflow/contrib/linear_optimizer/python/ops/sharded_mutable_dense_hashtable_test.py @@ -12,7 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Tests for sharded_mutable_dense_hashtable.py.""" +"""Tests for sharded_mutable_dense_hashtable.py (deprecated). + +This module and all its submodules are deprecated. To UPDATE or USE linear +optimizers, please check its latest version in core: +tensorflow_estimator/python/estimator/canned/linear_optimizer/. +""" from __future__ import absolute_import from __future__ import division diff --git a/tensorflow/contrib/linear_optimizer/python/ops/sparse_feature_column.py b/tensorflow/contrib/linear_optimizer/python/ops/sparse_feature_column.py index 003795233ff2b28e33fc10388ef25efb63c43bb0..64730f8eed1ff9bfcd4a980dceb28abb98e39f73 100644 --- a/tensorflow/contrib/linear_optimizer/python/ops/sparse_feature_column.py +++ b/tensorflow/contrib/linear_optimizer/python/ops/sparse_feature_column.py @@ -12,7 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Sparse feature column.""" +"""Sparse feature column (deprecated). + +This module and all its submodules are deprecated. To UPDATE or USE linear +optimizers, please check its latest version in core: +tensorflow_estimator/python/estimator/canned/linear_optimizer/. +""" from __future__ import absolute_import from __future__ import division @@ -21,6 +26,7 @@ from __future__ import print_function from tensorflow.python.framework import dtypes from tensorflow.python.framework.ops import internal_convert_to_tensor from tensorflow.python.framework.ops import name_scope +from tensorflow.python.util import deprecation class SparseFeatureColumn(object): @@ -68,6 +74,10 @@ class SparseFeatureColumn(object): @@feature_values """ + @deprecation.deprecated( + None, 'This class is deprecated. To UPDATE or USE linear optimizers, ' + 'please check its latest version in core: ' + 'tensorflow_estimator/python/estimator/canned/linear_optimizer/.') def __init__(self, example_indices, feature_indices, feature_values): """Creates a `SparseFeatureColumn` representation. diff --git a/tensorflow/contrib/linear_optimizer/python/ops/sparse_feature_column_test.py b/tensorflow/contrib/linear_optimizer/python/ops/sparse_feature_column_test.py index 51c4f68543da2f563481cc2d35b556796616cf9d..0ae780e1a100c7dadde7196803f2ae0d4bcb2334 100644 --- a/tensorflow/contrib/linear_optimizer/python/ops/sparse_feature_column_test.py +++ b/tensorflow/contrib/linear_optimizer/python/ops/sparse_feature_column_test.py @@ -12,7 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Tests for sparse_feature_column.py.""" +"""Tests for sparse_feature_column.py (deprecated). + +This module and all its submodules are deprecated. To UPDATE or USE linear +optimizers, please check its latest version in core: +tensorflow_estimator/python/estimator/canned/linear_optimizer/. +""" from __future__ import absolute_import from __future__ import division diff --git a/tensorflow/contrib/linear_optimizer/python/sdca_estimator_test.py b/tensorflow/contrib/linear_optimizer/python/sdca_estimator_test.py index 647667188238dc18b137eaad98356a79b3a549b4..7a5354222f103aa0f45adc513079e420bbbfd30c 100644 --- a/tensorflow/contrib/linear_optimizer/python/sdca_estimator_test.py +++ b/tensorflow/contrib/linear_optimizer/python/sdca_estimator_test.py @@ -524,7 +524,7 @@ class SDCALinearRegressorTest(test.TestCase): # LinearClassifier requires at least one column. 'place_holder': constant_op.constant([[0.0]] * num_examples), - }, constant_op.constant([[1 if i % 4 is 0 else 0] + }, constant_op.constant([[1 if i % 4 == 0 else 0] for i in range(num_examples)]) with self._single_threaded_test_session(): diff --git a/tensorflow/contrib/lite/BUILD b/tensorflow/contrib/lite/BUILD deleted file mode 100644 index 787a85644c35c807df84f74cbce06f80fd0b004d..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/BUILD +++ /dev/null @@ -1,358 +0,0 @@ -package(default_visibility = [ - "//visibility:public", -]) - -licenses(["notice"]) # Apache 2.0 - -load("//tensorflow:tensorflow.bzl", "tf_cc_test") -load("//tensorflow/contrib/lite:build_def.bzl", "tflite_copts", "gen_selected_ops") - -exports_files(glob([ - "testdata/*.bin", - "testdata/*.pb", - "models/testdata/*", -])) - -config_setting( - name = "mips", - values = { - "cpu": "mips", - }, -) - -config_setting( - name = "mips64", - values = { - "cpu": "mips64", - }, -) - -# Enables inclusion of TensorFlow kernels via the TF Lite Flex delegate. -# WARNING: This build flag is experimental and subject to change. -config_setting( - name = "with_tflite_flex", - define_values = {"with_tflite_flex": "true"}, - visibility = ["//visibility:public"], -) - -cc_library( - name = "schema_fbs_version", - hdrs = ["version.h"], -) - -cc_library( - name = "arena_planner", - srcs = ["arena_planner.cc"], - hdrs = ["arena_planner.h"], - deps = [ - ":graph_info", - ":memory_planner", - ":simple_memory_arena", - "//tensorflow/contrib/lite/c:c_api_internal", - ], -) - -cc_test( - name = "arena_planner_test", - size = "small", - srcs = ["arena_planner_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable", - ], - deps = [ - ":arena_planner", - "//tensorflow/contrib/lite/testing:util", - "//tensorflow/core:framework", - "//tensorflow/core:lib", - "@com_google_googletest//:gtest", - ], -) - -# Main library. No ops are included here. -# TODO(aselle): Resolve problems preventing C99 usage. -cc_library( - name = "context", - hdrs = ["context.h"], - deps = ["//tensorflow/contrib/lite/c:c_api_internal"], -) - -cc_library( - name = "graph_info", - hdrs = ["graph_info.h"], - deps = ["//tensorflow/contrib/lite/c:c_api_internal"], -) - -cc_library( - name = "memory_planner", - hdrs = ["memory_planner.h"], - deps = ["//tensorflow/contrib/lite/c:c_api_internal"], -) - -cc_library( - name = "simple_memory_arena", - srcs = ["simple_memory_arena.cc"], - hdrs = ["simple_memory_arena.h"], - deps = ["//tensorflow/contrib/lite/c:c_api_internal"], -) - -cc_library( - name = "builtin_op_data", - hdrs = [ - "builtin_op_data.h", - ], - deps = ["//tensorflow/contrib/lite/c:c_api_internal"], -) - -cc_library( - name = "kernel_api", - hdrs = [ - "builtin_op_data.h", - "builtin_ops.h", - "context.h", - "context_util.h", - ], -) - -exports_files(["builtin_ops.h"]) - -cc_library( - name = "string", - hdrs = [ - "string.h", - ], - deps = [ - "//tensorflow/core:lib_platform", - ], -) - -# TODO(ahentz): investigate dependency on gemm_support requiring usage of tf_copts. -cc_library( - name = "framework", - srcs = [ - "allocation.cc", - "graph_info.cc", - "interpreter.cc", - "model.cc", - "mutable_op_resolver.cc", - "optional_debug_tools.cc", - "stderr_reporter.cc", - ] + select({ - "//tensorflow:android": [ - "nnapi_delegate.cc", - "mmap_allocation.cc", - ], - "//tensorflow:windows": [ - "nnapi_delegate_disabled.cc", - "mmap_allocation_disabled.cc", - ], - "//conditions:default": [ - "nnapi_delegate_disabled.cc", - "mmap_allocation.cc", - ], - }), - hdrs = [ - "allocation.h", - "context.h", - "context_util.h", - "error_reporter.h", - "graph_info.h", - "interpreter.h", - "model.h", - "mutable_op_resolver.h", - "nnapi_delegate.h", - "op_resolver.h", - "optional_debug_tools.h", - "stderr_reporter.h", - ], - copts = tflite_copts(), - linkopts = [ - ] + select({ - "//tensorflow:android": [ - "-llog", - ], - "//conditions:default": [ - ], - }), - deps = [ - ":arena_planner", - ":graph_info", - ":memory_planner", - ":schema_fbs_version", - ":simple_memory_arena", - ":string", - ":util", - "//tensorflow/contrib/lite/c:c_api_internal", - "//tensorflow/contrib/lite/core/api", - "//tensorflow/contrib/lite/kernels:eigen_support", - "//tensorflow/contrib/lite/kernels:gemm_support", - "//tensorflow/contrib/lite/nnapi:nnapi_lib", - "//tensorflow/contrib/lite/profiling:profiler", - "//tensorflow/contrib/lite/schema:schema_fbs", - ] + select({ - ":with_tflite_flex": [ - "//tensorflow/contrib/lite/delegates/flex:delegate", - ], - "//conditions:default": [], - }), -) - -cc_library( - name = "string_util", - srcs = ["string_util.cc"], - hdrs = ["string_util.h"], - deps = [ - ":framework", - ":string", - ], -) - -cc_test( - name = "string_util_test", - size = "small", - srcs = ["string_util_test.cc"], - deps = [ - ":framework", - ":string_util", - "//tensorflow/contrib/lite/testing:util", - "@com_google_googletest//:gtest", - ], -) - -# Test main interpreter -cc_test( - name = "interpreter_test", - size = "small", - srcs = ["interpreter_test.cc"], - deps = [ - ":framework", - ":string_util", - "//tensorflow/contrib/lite/c:c_api_internal", - "//tensorflow/contrib/lite/core/api", - "//tensorflow/contrib/lite/kernels:builtin_ops", - "//tensorflow/contrib/lite/kernels:kernel_util", - "//tensorflow/contrib/lite/kernels/internal:tensor_utils", - "//tensorflow/contrib/lite/schema:schema_fbs", - "//tensorflow/contrib/lite/testing:util", - "@com_google_googletest//:gtest", - ], -) - -# Test graph utils -cc_test( - name = "graph_info_test", - size = "small", - srcs = ["graph_info_test.cc"], - tags = ["no_oss"], - deps = [ - ":framework", - ":string_util", - "//tensorflow/contrib/lite/testing:util", - "@com_google_googletest//:gtest", - ], -) - -# Test arena allocator -cc_test( - name = "simple_memory_arena_test", - size = "small", - srcs = ["simple_memory_arena_test.cc"], - deps = [ - ":simple_memory_arena", - "//tensorflow/contrib/lite/testing:util", - "@com_google_googletest//:gtest", - ], -) - -# Test model framework. -cc_test( - name = "model_test", - size = "small", - srcs = ["model_test.cc"], - data = [ - "testdata/0_subgraphs.bin", - "testdata/2_subgraphs.bin", - "testdata/empty_model.bin", - "testdata/multi_add_flex.bin", - "testdata/test_model.bin", - "testdata/test_model_broken.bin", - ], - deps = [ - ":framework", - "//tensorflow/contrib/lite/c:c_api_internal", - "//tensorflow/contrib/lite/core/api", - "//tensorflow/contrib/lite/kernels:builtin_ops", - "//tensorflow/contrib/lite/testing:util", - "@com_google_googletest//:gtest", - ], -) - -# Test model framework with the flex library linked into the target. -tf_cc_test( - name = "model_flex_test", - size = "small", - srcs = ["model_flex_test.cc"], - data = [ - "testdata/multi_add_flex.bin", - ], - tags = ["no_windows"], # TODO(b/116667551): No weak symbols with MSVC. - deps = [ - ":framework", - "//tensorflow/contrib/lite/core/api", - "//tensorflow/contrib/lite/delegates/flex:delegate", - "//tensorflow/contrib/lite/kernels:builtin_ops", - "//tensorflow/contrib/lite/testing:util", - "@com_google_googletest//:gtest", - ], -) - -# Test OpResolver. -cc_test( - name = "mutable_op_resolver_test", - size = "small", - srcs = ["mutable_op_resolver_test.cc"], - tags = ["no_oss"], - deps = [ - ":framework", - "//tensorflow/contrib/lite/testing:util", - "@com_google_googletest//:gtest", - ], -) - -cc_library( - name = "util", - srcs = ["util.cc"], - hdrs = ["util.h"], - deps = [ - "//tensorflow/contrib/lite/c:c_api_internal", - ], -) - -cc_test( - name = "util_test", - size = "small", - srcs = ["util_test.cc"], - tags = ["no_oss"], - deps = [ - ":util", - "//tensorflow/contrib/lite/testing:util", - "@com_google_googletest//:gtest", - ], -) - -# Test the serialization of a model with optional tensors. - -# Model tests - -#cc_library( -# name = "models_test_utils", -# testonly = 1, -# hdrs = ["models/test_utils.h"], -# deps = select({ -# "//tensorflow:android": [], -# "//conditions:default": [ -# "@com_google_absl//absl/strings", -# "//tensorflow/core:test", -# ], -# }), -#) diff --git a/tensorflow/contrib/lite/README.md b/tensorflow/contrib/lite/README.md deleted file mode 100644 index a4b3d83efe09358cb8e7a5f673a96f28faa84d08..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/README.md +++ /dev/null @@ -1,8 +0,0 @@ -# TensorFlow Lite - -TensorFlow Lite is TensorFlow's lightweight solution for mobile and embedded -devices. It enables low-latency inference of on-device machine learning models -with a small binary size and fast performance supporting hardware acceleration. - -See the documentation: https://www.tensorflow.org/lite/ -Documentation edits can be made here: [tensorflow/contrib/lite/g3doc](./g3doc/) diff --git a/tensorflow/contrib/lite/builtin_op_data.h b/tensorflow/contrib/lite/builtin_op_data.h deleted file mode 100644 index 30901bd0fae9510ebea288288941218d6994d888..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/builtin_op_data.h +++ /dev/null @@ -1,22 +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. -==============================================================================*/ -// Compatibility shim for new location of interface definitions. - -#ifndef TENSORFLOW_CONTRIB_LITE_BUILTIN_OP_DATA_H_ -#define TENSORFLOW_CONTRIB_LITE_BUILTIN_OP_DATA_H_ - -#include "tensorflow/contrib/lite/c/builtin_op_data.h" - -#endif // TENSORFLOW_CONTRIB_LITE_BUILTIN_OP_DATA_H_ diff --git a/tensorflow/contrib/lite/c/BUILD b/tensorflow/contrib/lite/c/BUILD deleted file mode 100644 index 663eb63cad0da0781cc2d07d1b78242bea2ee3c8..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/c/BUILD +++ /dev/null @@ -1,39 +0,0 @@ -package( - default_visibility = ["//visibility:public"], -) - -licenses(["notice"]) # Apache 2.0 - -cc_library( - name = "c_api_internal", - srcs = ["c_api_internal.c"], - hdrs = [ - "builtin_op_data.h", - "c_api_internal.h", - ], - visibility = [ - "//tensorflow/contrib/lite:__subpackages__", - ], -) - -# Test the C extension API code. -cc_test( - name = "c_api_internal_test", - size = "small", - srcs = ["c_api_internal_test.cc"], - deps = [ - ":c_api_internal", - "@com_google_googletest//:gtest", - ], -) - -cc_test( - name = "builtin_op_data_test", - size = "small", - srcs = ["builtin_op_data_test.cc"], - copts = ["-Wno-unused-variable"], - deps = [ - ":c_api_internal", - "@com_google_googletest//:gtest", - ], -) diff --git a/tensorflow/contrib/lite/c/builtin_op_data.h b/tensorflow/contrib/lite/c/builtin_op_data.h deleted file mode 100644 index 5a5f3ad61c1c9753cffb34e8f7cc0e005a8c971f..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/c/builtin_op_data.h +++ /dev/null @@ -1,331 +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_CONTRIB_LITE_C_BUILTIN_OP_DATA_H_ -#define TENSORFLOW_CONTRIB_LITE_C_BUILTIN_OP_DATA_H_ - -#include - -#include "tensorflow/contrib/lite/c/c_api_internal.h" - -#ifdef __cplusplus -extern "C" { -#endif // __cplusplus - -// TODO(aselle): Consider using "if this then that" for testing. - -// IMPORTANT: All new members of structs must be added at the end to ensure -// backwards compatibility. - -// Possible padding types (for convolutions) -typedef enum { - kTfLitePaddingUnknown = 0, - kTfLitePaddingSame, - kTfLitePaddingValid, -} TfLitePadding; - -typedef struct { - int width; - int height; -} TfLitePaddingValues; - -// Possible fused activation functions. -// TODO(aselle): rename to TfLiteActivation -typedef enum { - kTfLiteActNone = 0, - kTfLiteActRelu, - kTfLiteActRelu1, - kTfLiteActRelu6, - kTfLiteActTanh, - kTfLiteActSignBit, - kTfLiteActSigmoid, -} TfLiteFusedActivation; - -typedef struct { - TfLitePadding padding; - int stride_width; - int stride_height; - int dilation_width_factor; - int dilation_height_factor; - TfLiteFusedActivation activation; -} TfLiteConvParams; - -typedef struct { - TfLitePadding padding; - int stride_width; - int stride_height; - int filter_width; - int filter_height; - TfLiteFusedActivation activation; - struct { - TfLitePaddingValues padding; - } computed; -} TfLitePoolParams; - -typedef struct { - // Parameters for DepthwiseConv version 1 or above. - TfLitePadding padding; - int stride_width; - int stride_height; - int depth_multiplier; - TfLiteFusedActivation activation; - // Parameters for DepthwiseConv version 2 or above. - int dilation_width_factor; - int dilation_height_factor; -} TfLiteDepthwiseConvParams; - -typedef struct { - int rank; - TfLiteFusedActivation activation; -} TfLiteSVDFParams; - -typedef struct { - TfLiteFusedActivation activation; -} TfLiteRNNParams; - -typedef struct { - bool time_major; - TfLiteFusedActivation activation; -} TfLiteSequenceRNNParams; - -typedef struct { - bool time_major; - TfLiteFusedActivation activation; - bool merge_outputs; -} TfLiteBidirectionalSequenceRNNParams; - -typedef enum { - kTfLiteFullyConnectedWeightsFormatDefault = 0, - kTfLiteFullyConnectedWeightsFormatShuffled4x16Int8 = 1, -} TfLiteFullyConnectedWeightsFormat; - -typedef struct { - // Parameters for FullyConnected version 1 or above. - TfLiteFusedActivation activation; - - // Parameters for FullyConnected version 2 or above. - TfLiteFullyConnectedWeightsFormat weights_format; -} TfLiteFullyConnectedParams; - -typedef enum { - kTfLiteLshProjectionUnknown = 0, - kTfLiteLshProjectionSparse = 1, - kTfLiteLshProjectionDense = 2, -} TfLiteLSHProjectionType; - -typedef struct { - TfLiteLSHProjectionType type; -} TfLiteLSHProjectionParams; - -typedef struct { - float beta; -} TfLiteSoftmaxParams; - -typedef struct { - int axis; - TfLiteFusedActivation activation; -} TfLiteConcatenationParams; - -typedef struct { - TfLiteFusedActivation activation; -} TfLiteAddParams; - -typedef struct { -} TfLiteSpaceToBatchNDParams; - -typedef struct { -} TfLiteBatchToSpaceNDParams; - -typedef struct { - TfLiteFusedActivation activation; -} TfLiteMulParams; - -typedef struct { - TfLiteFusedActivation activation; -} TfLiteSubParams; - -typedef struct { - TfLiteFusedActivation activation; -} TfLiteDivParams; - -typedef struct { - TfLiteFusedActivation activation; -} TfLiteL2NormParams; - -typedef struct { - int radius; - float bias; - float alpha; - float beta; -} TfLiteLocalResponseNormParams; - -typedef enum { - kTfLiteLSTMFullKernel = 0, - kTfLiteLSTMBasicKernel -} TfLiteLSTMKernelType; - -typedef struct { - // Parameters for LSTM version 1. - TfLiteFusedActivation activation; - float cell_clip; - float proj_clip; - - // Parameters for LSTM version 2. - // kTfLiteLSTMBasicKernel is only supported in version 2 or above. - TfLiteLSTMKernelType kernel_type; -} TfLiteLSTMParams; - -typedef struct { - // Parameters needed for the underlying LSTM. - TfLiteFusedActivation activation; - float cell_clip; - float proj_clip; - - // If set to true then the first dimension is time, otherwise batch. - bool time_major; -} TfLiteUnidirectionalSequenceLSTMParams; - -typedef struct { - // Parameters for the LSTM kernel. - TfLiteFusedActivation activation; - float cell_clip; - float proj_clip; - - // If true, store the outputs of both directions in the first output. - bool merge_outputs; -} TfLiteBidirectionalSequenceLSTMParams; - -typedef struct { - bool align_corners; -} TfLiteResizeBilinearParams; - -typedef struct { -} TfLitePadParams; - -typedef struct { -} TfLitePadV2Params; - -typedef struct { - // TODO(ahentz): We can't have dynamic data in this struct, at least not yet. - // For now we will fix the maximum possible number of dimensions. - int shape[8]; - int num_dimensions; -} TfLiteReshapeParams; - -typedef struct { - int ngram_size; - int max_skip_size; - bool include_all_ngrams; -} TfLiteSkipGramParams; - -typedef struct { - int block_size; -} TfLiteSpaceToDepthParams; - -typedef struct { - TfLiteType in_data_type; - TfLiteType out_data_type; -} TfLiteCastParams; - -typedef enum { - kTfLiteCombinerTypeSum = 0, - kTfLiteCombinerTypeMean = 1, - kTfLiteCombinerTypeSqrtn = 2, -} TfLiteCombinerType; - -typedef struct { - TfLiteCombinerType combiner; -} TfLiteEmbeddingLookupSparseParams; - -typedef struct { - int axis; -} TfLiteGatherParams; - -typedef struct { -} TfLiteTransposeParams; - -typedef struct { - bool keep_dims; -} TfLiteReducerParams; - -typedef struct { - int num_splits; -} TfLiteSplitParams; - -typedef struct { - // TODO(ahentz): We can't have dynamic data in this struct, at least not yet. - // For now we will fix the maximum possible number of dimensions. - int squeeze_dims[8]; - int num_squeeze_dims; -} TfLiteSqueezeParams; - -typedef struct { - int begin_mask; - int end_mask; - int ellipsis_mask; - int new_axis_mask; - int shrink_axis_mask; -} TfLiteStridedSliceParams; - -typedef struct { - TfLiteType output_type; -} TfLiteArgMaxParams; - -typedef struct { - TfLiteType output_type; -} TfLiteArgMinParams; - -typedef struct { - TfLitePadding padding; - int stride_width; - int stride_height; -} TfLiteTransposeConvParams; - -typedef struct { - bool validate_indices; -} TfLiteSparseToDenseParams; - -typedef struct { - TfLiteType out_type; -} TfLiteShapeParams; - -typedef struct { - // Parameters supported by version 1: - float min; - float max; - int num_bits; - - // Parameters supported by version 2: - bool narrow_range; -} TfLiteFakeQuantParams; - -typedef struct { - int values_count; - int axis; -} TfLitePackParams; - -typedef struct { - int axis; -} TfLiteOneHotParams; - -typedef struct { - int num; - int axis; -} TfLiteUnpackParams; - -#ifdef __cplusplus -} // extern "C" -#endif // __cplusplus - -#endif // TENSORFLOW_CONTRIB_LITE_C_BUILTIN_OP_DATA_H_ diff --git a/tensorflow/contrib/lite/c/c_api_internal.c b/tensorflow/contrib/lite/c/c_api_internal.c deleted file mode 100644 index 8be37945ca2a5ddf3c8cedc5a3ae5e34da8a4b9b..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/c/c_api_internal.c +++ /dev/null @@ -1,114 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/contrib/lite/c/c_api_internal.h" -#ifndef TF_LITE_STATIC_MEMORY -#include -#include -#include -#endif // TF_LITE_STATIC_MEMORY - -int TfLiteIntArrayGetSizeInBytes(int size) { - static TfLiteIntArray dummy; - return sizeof(dummy) + sizeof(dummy.data[0]) * size; -} - -int TfLiteIntArrayEqual(TfLiteIntArray* a, TfLiteIntArray* b) { - if (a == b) return 1; - if (a == NULL || b == NULL) return 0; - return TfLiteIntArrayEqualsArray(a, b->size, b->data); -} - -int TfLiteIntArrayEqualsArray(TfLiteIntArray* a, int b_size, int b_data[]) { - if (a == NULL) return (b_size == 0); - if (a->size != b_size) return 0; - int i = 0; - for (; i < a->size; i++) - if (a->data[i] != b_data[i]) return 0; - return 1; -} - -#ifndef TF_LITE_STATIC_MEMORY - -TfLiteIntArray* TfLiteIntArrayCreate(int size) { - TfLiteIntArray* ret = - (TfLiteIntArray*)malloc(TfLiteIntArrayGetSizeInBytes(size)); - ret->size = size; - return ret; -} - -void TfLiteIntArrayPrint(const char* s, TfLiteIntArray* a) { - printf("%s: length=%d [", s, a->size); - if (a->size) printf("%d", a->data[0]); - int i = 1; - for (; i < a->size; i++) { - printf(" %d", a->data[i]); - } - printf("]\n"); -} - -TfLiteIntArray* TfLiteIntArrayCopy(TfLiteIntArray* src) { - if (!src) return NULL; - TfLiteIntArray* ret = TfLiteIntArrayCreate(src->size); - if (ret) { - memcpy(ret->data, src->data, src->size * sizeof(int)); - } - return ret; -} - -void TfLiteIntArrayFree(TfLiteIntArray* a) { free(a); } - -void TfLiteTensorDataFree(TfLiteTensor* t) { - if (t->allocation_type == kTfLiteDynamic && t->data.raw) { - free(t->data.raw); - } - t->data.raw = NULL; -} - -void TfLiteTensorFree(TfLiteTensor* t) { - TfLiteTensorDataFree(t); - if (t->dims) TfLiteIntArrayFree(t->dims); - t->dims = NULL; -} - -void TfLiteTensorReset(TfLiteType type, const char* name, TfLiteIntArray* dims, - TfLiteQuantizationParams quantization, char* buffer, - size_t size, TfLiteAllocationType allocation_type, - const void* allocation, bool is_variable, - TfLiteTensor* tensor) { - TfLiteTensorFree(tensor); - tensor->type = type; - tensor->name = name; - tensor->dims = dims; - tensor->params = quantization; - tensor->data.raw = buffer; - tensor->bytes = size; - tensor->allocation_type = allocation_type; - tensor->allocation = allocation; - tensor->is_variable = is_variable; -} - -void TfLiteTensorRealloc(size_t num_bytes, TfLiteTensor* tensor) { - if (tensor->allocation_type != kTfLiteDynamic) { - return; - } - if (!tensor->data.raw) { - tensor->data.raw = malloc(num_bytes); - } else if (num_bytes > tensor->bytes) { - tensor->data.raw = realloc(tensor->data.raw, num_bytes); - } - tensor->bytes = num_bytes; -} -#endif // TF_LITE_STATIC_MEMORY diff --git a/tensorflow/contrib/lite/c/c_api_internal.h b/tensorflow/contrib/lite/c/c_api_internal.h deleted file mode 100644 index fdc9ff634a19d348ab2dfae60d94722619dfec06..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/c/c_api_internal.h +++ /dev/null @@ -1,499 +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. -==============================================================================*/ -// This file defines a C API for implementing operations in tflite. -// These operations can be defined using c++ but the interface between -// the interpreter and the operations are C. -// -// Summary of abstractions -// TF_LITE_ENSURE - Self-sufficient error checking -// TfLiteStatus - Status reporting -// TfLiteIntArray - stores tensor shapes (dims), -// TfLiteContext - allows an op to access the tensors -// TfLiteTensor - tensor (a multidimensional array) -// TfLiteNode - a single node or operation -// TfLiteRegistration - the implementation of a conceptual operation. -// -// Some abstractions in this file are created and managed by Interpreter. -#ifndef TENSORFLOW_CONTRIB_LITE_C_C_API_INTERNAL_H_ -#define TENSORFLOW_CONTRIB_LITE_C_C_API_INTERNAL_H_ - -#include -#include -#include - -#ifdef __cplusplus -extern "C" { -#endif // __cplusplus - -typedef enum { kTfLiteOk = 0, kTfLiteError = 1 } TfLiteStatus; - -// The list of external context types known to TF Lite. This list exists solely -// to avoid conflicts and to ensure ops can share the external contexts they -// need. Access to the external contexts is controled by one of the -// corresponding support files. -typedef enum { - kTfLiteEigenContext = 0, // include eigen_support.h to use. - kTfLiteGemmLowpContext = 1, // include gemm_support.h to use. - kTfLiteEdgeTpuContext = 2, // Placeholder for Edge TPU support. - kTfLiteMaxExternalContexts = 3 -} TfLiteExternalContextType; - -// An external context is a collection of information unrelated to the TF Lite -// framework, but useful to a subset of the ops. TF Lite knows very little -// about about the actual contexts, but it keeps a list of them, and is able to -// refresh them if configurations like the number of recommended threads -// change. -typedef struct { - TfLiteExternalContextType type; - TfLiteStatus (*Refresh)(struct TfLiteContext* context); -} TfLiteExternalContext; - -// Forward declare so GetNode can use this is in Context. -typedef struct _TfLiteRegistration TfLiteRegistration; -typedef struct _TfLiteDelegate TfLiteDelegate; - -#define kOptionalTensor (-1) - -// Fixed size list of integers. Used for dimensions and inputs/outputs tensor -// indices -typedef struct { - int size; -// gcc 6.1+ have a bug where flexible members aren't properly handled -// https://github.com/google/re2/commit/b94b7cd42e9f02673cd748c1ac1d16db4052514c -#if !defined(__clang__) && defined(__GNUC__) && __GNUC__ == 6 && \ - __GNUC_MINOR__ >= 1 - int data[0]; -#else - int data[]; -#endif -} TfLiteIntArray; - -// Given the size (number of elements) in a TfLiteIntArray, calculate its size -// in bytes. -int TfLiteIntArrayGetSizeInBytes(int size); - -// Create a array of a given `size` (uninitialized entries). -// This returns a pointer, that you must free using TfLiteIntArrayFree(). -TfLiteIntArray* TfLiteIntArrayCreate(int size); - -// Check if two intarrays are equal. Returns 1 if they are equal, 0 otherwise. -int TfLiteIntArrayEqual(TfLiteIntArray* a, TfLiteIntArray* b); - -// Check if an intarray equals an array. Returns 1 if equals, 0 otherwise. -int TfLiteIntArrayEqualsArray(TfLiteIntArray* a, int b_size, int b_data[]); - -// Create a copy of an array passed as `src`. -// You are expected to free memory with TfLiteIntArrayFree -TfLiteIntArray* TfLiteIntArrayCopy(TfLiteIntArray* src); - -// Free memory of array `v`. -void TfLiteIntArrayFree(TfLiteIntArray* v); - -// Since we must not depend on any libraries, define a minimal subset of -// error macros while avoiding names that have pre-conceived meanings like -// assert and check. - -// Check whether value is true, and if not return kTfLiteError from -// the current function (and report the error string msg). -#define TF_LITE_ENSURE_MSG(context, value, msg) \ - do { \ - if (!(value)) { \ - (context)->ReportError((context), __FILE__ " " msg); \ - return kTfLiteError; \ - } \ - } while (0) - -// Check whether the value `a` is true, and if not return kTfLiteError from -// the current function, while also reporting the location of the error. -#define TF_LITE_ENSURE(context, a) \ - do { \ - if (!(a)) { \ - (context)->ReportError((context), "%s:%d %s was not true.", __FILE__, \ - __LINE__, #a); \ - return kTfLiteError; \ - } \ - } while (0) - -#define TF_LITE_ENSURE_STATUS(a) \ - do { \ - if ((a) != kTfLiteOk) { \ - return kTfLiteError; \ - } \ - } while (0) - -// Check whether the value `a == b` is true, and if not return kTfLiteError from -// the current function, while also reporting the location of the error. -// `a` and `b` may be evaluated more than once, so no side effects or -// extremely expensive computations should be done. -#define TF_LITE_ENSURE_EQ(context, a, b) \ - do { \ - if ((a) != (b)) { \ - (context)->ReportError((context), "%s:%d %s != %s (%d != %d)", __FILE__, \ - __LINE__, #a, #b, (a), (b)); \ - return kTfLiteError; \ - } \ - } while (0) - -#define TF_LITE_ENSURE_OK(context, status) \ - do { \ - if ((status) != kTfLiteOk) { \ - return kTfLiteError; \ - } \ - } while (0) - -// Single-precision complex data type compatible with the C99 definition. -typedef struct { - float re, im; // real and imaginary parts, respectively. -} TfLiteComplex64; - -// Types supported by tensor -typedef enum { - kTfLiteNoType = 0, - kTfLiteFloat32 = 1, - kTfLiteInt32 = 2, - kTfLiteUInt8 = 3, - kTfLiteInt64 = 4, - kTfLiteString = 5, - kTfLiteBool = 6, - kTfLiteInt16 = 7, - kTfLiteComplex64 = 8, -} TfLiteType; - -// Parameters for asymmetric quantization. Quantized values can be converted -// back to float using: -// real_value = scale * (quantized_value - zero_point); -typedef struct { - float scale; - int32_t zero_point; -} TfLiteQuantizationParams; - -// A union of pointers that points to memory for a given tensor. -typedef union { - int* i32; - int64_t* i64; - float* f; - char* raw; - const char* raw_const; - uint8_t* uint8; - bool* b; - int16_t* i16; - TfLiteComplex64* c64; -} TfLitePtrUnion; - -// Memory allocation strategies. kTfLiteMmapRo is for read-only memory-mapped -// data (or data externally allocated). kTfLiteArenaRw is arena allocated -// data. kTfLiteDynamic is for tensors that are allocated during evaluation. -typedef enum { - kTfLiteMemNone = 0, - kTfLiteMmapRo, - kTfLiteArenaRw, - kTfLiteArenaRwPersistent, - kTfLiteDynamic, -} TfLiteAllocationType; - -// The delegates should use zero or positive integers to represent handles. -// -1 is reserved from unallocated status. -typedef int TfLiteBufferHandle; -const TfLiteBufferHandle kTfLiteNullBufferHandle = -1; - -// An tensor in the interpreter system which is a wrapper around a buffer of -// data including a dimensionality (or NULL if not currently defined). -typedef struct { - // The data type specification for data stored in `data`. This affects - // what member of `data` union should be used. - TfLiteType type; - // A union of data pointers. The appropriate type should be used for a typed - // tensor based on `type`. - TfLitePtrUnion data; - // A pointer to a structure representing the dimensionality interpretation - // that the buffer should have. NOTE: the product of elements of `dims` - // and the element datatype size should be equal to `bytes` below. - TfLiteIntArray* dims; - // Quantization information. - TfLiteQuantizationParams params; - // How memory is mapped - // kTfLiteMmapRo: Memory mapped read only. - // i.e. weights - // kTfLiteArenaRw: Arena allocated read write memory - // (i.e. temporaries, outputs). - TfLiteAllocationType allocation_type; - // The number of bytes required to store the data of this Tensor. I.e. - // (bytes of each element) * dims[0] * ... * dims[n-1]. For example, if - // type is kTfLiteFloat32 and dims = {3, 2} then - // bytes = sizeof(float) * 3 * 2 = 4 * 3 * 2 = 24. - size_t bytes; - - // An opaque pointer to a tflite::MMapAllocation - const void* allocation; - - // Null-terminated name of this tensor. - const char* name; - - // The delegate which knows how to handle `buffer_handle`. - // WARNING: This is an experimental interface that is subject to change. - TfLiteDelegate* delegate; - - // An integer buffer handle that can be handled by `delegate`. - // The value is valid only when delegate is not null. - // WARNING: This is an experimental interface that is subject to change. - TfLiteBufferHandle buffer_handle; - - // If the delegate uses its own buffer (e.g. GPU memory), the delegate is - // responsible to set data_is_stale to true. - // `delegate->CopyFromBufferHandle` can be called to copy the data from - // delegate buffer. - // WARNING: This is an // experimental interface that is subject to change. - bool data_is_stale; - - // True if the tensor is a variable. - bool is_variable; -} TfLiteTensor; - -// Free data memory of tensor `t`; -void TfLiteTensorDataFree(TfLiteTensor* t); - -// Free memory of tensor `t`; -void TfLiteTensorFree(TfLiteTensor* t); - -// Set all of a tensor's fields (and free any previously allocated data). -void TfLiteTensorReset(TfLiteType type, const char* name, TfLiteIntArray* dims, - TfLiteQuantizationParams quantization, char* buffer, - size_t size, TfLiteAllocationType allocation_type, - const void* allocation, bool is_variable, - TfLiteTensor* tensor); - -// Resize the allocated data of a (dynamic) tensor. Tensors with allocation -// types other than kTfLiteDynamic will be ignored. -void TfLiteTensorRealloc(size_t num_bytes, TfLiteTensor* tensor); - -// A structure representing an instance of a node. -// This structure only exhibits the inputs, outputs and user defined data, not -// other features like the type. -typedef struct { - // Inputs to this node expressed as indices into the simulator's tensors. - TfLiteIntArray* inputs; - - // Outputs to this node expressed as indices into the simulator's tensors. - TfLiteIntArray* outputs; - - // Temporary tensors uses during the computations. This usually contains no - // tensors, but ops are allowed to change that if they need scratch space of - // any sort. - TfLiteIntArray* temporaries; - - // Opaque data provided by the node implementer through `Registration.init`. - void* user_data; - - // Opaque data provided to the node if the node is a builtin. This is usually - // a structure defined in builtin_op_data.h - void* builtin_data; - - // Custom initial data. This is the opaque data provided in the flatbuffer. - // WARNING: This is an experimental interface that is subject to change. - const void* custom_initial_data; - int custom_initial_data_size; - - // The pointer to the delegate. This is non-null only when the node is - // created by calling `interpreter.ModifyGraphWithDelegate`. - // WARNING: This is an experimental interface that is subject to change. - TfLiteDelegate* delegate; -} TfLiteNode; - -typedef struct TfLiteContext { - // Number of tensors in the context. - size_t tensors_size; - - // The execution plan contains a list of the node indices in execution - // order. execution_plan->size is the current number of nodes. And, - // execution_plan->data[0] is the first node that needs to be run. - // TfLiteDelegates can traverse the current execution plan by iterating - // through each member of this array and using GetNodeAndRegistration() to - // access details about a node. i.e. - // TfLiteIntArray* execution_plan; - // TF_LITE_ENSURE_STATUS(context->GetExecutionPlan(context, &execution_plan)); - // for (int exec_index = 0; exec_index < execution_plan->size; exec_index++) { - // int node_index = execution_plan->data[exec_index]; - // TfLiteNode* node; - // TfLiteRegistration* reg; - // context->GetNodeAndRegistration(context, node_index, &node, ®); - // } - // WARNING: This is an experimental interface that is subject to change. - TfLiteStatus (*GetExecutionPlan)(struct TfLiteContext* context, - TfLiteIntArray** execution_plan); - - // An array of tensors in the interpreter context (of length `tensors_size`) - TfLiteTensor* tensors; - - // opaque full context ptr (an opaque c++ data structure) - void* impl_; - - // Request memory pointer be resized. Updates dimensions on the tensor. - // NOTE: ResizeTensor takes ownership of newSize. - TfLiteStatus (*ResizeTensor)(struct TfLiteContext*, TfLiteTensor* tensor, - TfLiteIntArray* new_size); - // Request that a error be reported with format string msg. - void (*ReportError)(struct TfLiteContext*, const char* msg, ...); - - // Add `tensors_to_add` tensors, preserving pre-existing Tensor entries. If - // non-null, the value pointed to by `first_new_tensor_index` will be set to - // the index of the first new tensor. - TfLiteStatus (*AddTensors)(struct TfLiteContext*, int tensors_to_add, - int* first_new_tensor_index); - - // Get a Tensor node by node_index. - // WARNING: This is an experimental interface that is subject to change. - TfLiteStatus (*GetNodeAndRegistration)(struct TfLiteContext*, int node_index, - TfLiteNode** node, - TfLiteRegistration** registration); - - // Replace ops with one or more stub delegate operations. This function - // does not take ownership of `nodes_to_replace`. - TfLiteStatus (*ReplaceSubgraphsWithDelegateKernels)( - struct TfLiteContext*, TfLiteRegistration registration, - const TfLiteIntArray* nodes_to_replace, TfLiteDelegate* delegate); - - // Number of threads that are recommended to subsystems like gemmlowp and - // eigen. - int recommended_num_threads; - - // Access external contexts by type. - // WARNING: This is an experimental interface that is subject to change. - TfLiteExternalContext* (*GetExternalContext)(struct TfLiteContext*, - TfLiteExternalContextType); - // Set the value of a external context. Does not take ownership of the - // pointer. - // WARNING: This is an experimental interface that is subject to change. - void (*SetExternalContext)(struct TfLiteContext*, TfLiteExternalContextType, - TfLiteExternalContext*); - - // Flag for allowing float16 precision for FP32 calculation. - // default: false. - // WARNING: This is an experimental API and subject to change. - bool allow_fp32_relax_to_fp16; -} TfLiteContext; - -typedef struct _TfLiteRegistration { - // Initializes the op from serialized data. - // If a built-in op: - // `buffer` is the op's params data (TfLiteLSTMParams*). - // `length` is zero. - // If custom op: - // `buffer` is the op's `custom_options`. - // `length` is the size of the buffer. - // - // Returns a type-punned (i.e. void*) opaque data (e.g. a primitive pointer - // or an instance of a struct). - // - // The returned pointer will be stored with the node in the `user_data` field, - // accessible within prepare and invoke functions below. - // NOTE: if the data is already in the desired format, simply implement this - // function to return `nullptr` and implement the free function to be a no-op. - void* (*init)(TfLiteContext* context, const char* buffer, size_t length); - - // The pointer `buffer` is the data previously returned by an init invocation. - void (*free)(TfLiteContext* context, void* buffer); - - // prepare is called when the inputs this node depends on have been resized. - // context->ResizeTensor() can be called to request output tensors to be - // resized. - // - // Returns kTfLiteOk on success. - TfLiteStatus (*prepare)(TfLiteContext* context, TfLiteNode* node); - - // Execute the node (should read node->inputs and output to node->outputs). - // Returns kTfLiteOk on success. - TfLiteStatus (*invoke)(TfLiteContext* context, TfLiteNode* node); - - // profiling_string is called during summarization of profiling information - // in order to group executions together. Providing a value here will cause a - // given op to appear multiple times is the profiling report. This is - // particularly useful for custom ops that can perform significantly - // different calculations depending on their `user-data`. - const char* (*profiling_string)(const TfLiteContext* context, - const TfLiteNode* node); - - // Builtin codes. If this kernel refers to a builtin this is the code - // of the builtin. This is so we can do marshaling to other frameworks like - // NN API. - // Note: It is the responsibility of the registration binder to set this - // properly. - int32_t builtin_code; - - // Custom op name. If the op is a builtin, this will be null. - // Note: It is the responsibility of the registration binder to set this - // properly. - // WARNING: This is an experimental interface that is subject to change. - const char* custom_name; - - // The version of the op. - // Note: It is the responsibility of the registration binder to set this - // properly. - int version; -} TfLiteRegistration; - -// WARNING: This is an experimental interface that is subject to change. -typedef struct _TfLiteDelegate { - // Data that delegate needs to identify itself. This data is owned by the - // delegate. The delegate is owned in the user code, so the delegate is - // responsible for doing this when it is destroyed. - void* data_; - - // Invoked by ModifyGraphWithDelegate. This prepare is called, giving the - // delegate a view of the current graph through TfLiteContext*. It typically - // will look at the nodes and call ReplaceSubgraphsWithDelegateKernels() - // to ask the TensorFlow lite runtime to create macro-nodes to represent - // delegated subgraphs of the original graph. - TfLiteStatus (*Prepare)(TfLiteContext* context, TfLiteDelegate* delegate); - - // Copy the data from delegate buffer handle to raw memory. - // This can be null if the delegate doesn't use its own buffer. - TfLiteStatus (*CopyFromBufferHandle)(TfLiteContext* context, - TfLiteDelegate* delegate, - TfLiteBufferHandle buffer_handle, - void* data, size_t size); - - // Copy the data from raw memory to delegate buffer handle. - // This can be null if the delegate doesn't use its own buffer. - TfLiteStatus (*CopyToBufferHandle)(TfLiteContext* context, - TfLiteDelegate* delegate, - TfLiteBufferHandle buffer_handle, - void* data, size_t size); - - // Free the Delegate Buffer Handle. Note: This only frees the handle, but - // this doesn't release the underlying resource (e.g. textures). The - // resources are either owned by application layer or the delegate. - // This can be null if the delegate doesn't use its own buffer. - void (*FreeBufferHandle)(TfLiteContext* context, TfLiteDelegate* delegate, - TfLiteBufferHandle* handle); -} TfLiteDelegate; - -// WARNING: This is an experimental interface that is subject to change. -// -// Currently, TfLiteDelegateParams has to be allocated in a way that it's -// trivially destructable. It will be stored as `builtin_data` field in -// `TfLiteNode` of the delegate node. -// -// See also the `CreateDelegateParams` function in `interpreter.cc` details. -typedef struct { - TfLiteDelegate* delegate; - TfLiteIntArray* nodes_to_replace; - TfLiteIntArray* input_tensors; - TfLiteIntArray* output_tensors; -} TfLiteDelegateParams; - -#ifdef __cplusplus -} // extern "C" -#endif // __cplusplus -#endif // TENSORFLOW_CONTRIB_LITE_C_C_API_INTERNAL_H_ diff --git a/tensorflow/contrib/lite/c/c_api_internal_test.cc b/tensorflow/contrib/lite/c/c_api_internal_test.cc deleted file mode 100644 index af398f32075b46e2ea487d49448f13435c4b5768..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/c/c_api_internal_test.cc +++ /dev/null @@ -1,73 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/contrib/lite/c/c_api_internal.h" -#include - -namespace tflite { - -// NOTE: this tests only the TfLiteIntArray part of context. -// most of c_api_internal.h is provided in the context of using it with -// interpreter.h and interpreter.cc, so interpreter_test.cc tests context -// structures more thoroughly. - -TEST(IntArray, TestIntArrayCreate) { - TfLiteIntArray* a = TfLiteIntArrayCreate(0); - TfLiteIntArray* b = TfLiteIntArrayCreate(3); - TfLiteIntArrayFree(a); - TfLiteIntArrayFree(b); -} - -TEST(IntArray, TestIntArrayCopy) { - TfLiteIntArray* a = TfLiteIntArrayCreate(2); - a->data[0] = 22; - a->data[1] = 24; - TfLiteIntArray* b = TfLiteIntArrayCopy(a); - ASSERT_NE(a, b); - ASSERT_EQ(a->size, b->size); - ASSERT_EQ(a->data[0], b->data[0]); - ASSERT_EQ(a->data[1], b->data[1]); - TfLiteIntArrayFree(a); - TfLiteIntArrayFree(b); -} - -TEST(IntArray, TestIntArrayEqual) { - TfLiteIntArray* a = TfLiteIntArrayCreate(1); - a->data[0] = 1; - TfLiteIntArray* b = TfLiteIntArrayCreate(2); - b->data[0] = 5; - b->data[1] = 6; - TfLiteIntArray* c = TfLiteIntArrayCreate(2); - c->data[0] = 5; - c->data[1] = 6; - TfLiteIntArray* d = TfLiteIntArrayCreate(2); - d->data[0] = 6; - d->data[1] = 6; - ASSERT_FALSE(TfLiteIntArrayEqual(a, b)); - ASSERT_TRUE(TfLiteIntArrayEqual(b, c)); - ASSERT_TRUE(TfLiteIntArrayEqual(b, b)); - ASSERT_FALSE(TfLiteIntArrayEqual(c, d)); - TfLiteIntArrayFree(a); - TfLiteIntArrayFree(b); - TfLiteIntArrayFree(c); - TfLiteIntArrayFree(d); -} - -} // namespace tflite - -int main(int argc, char** argv) { - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/contrib/lite/core/api/BUILD b/tensorflow/contrib/lite/core/api/BUILD deleted file mode 100644 index e4500534f348f15b47d3c3868461237e68fc3ac3..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/core/api/BUILD +++ /dev/null @@ -1,57 +0,0 @@ -package( - default_visibility = ["//visibility:public"], -) - -licenses(["notice"]) # Apache 2.0 - -load("//tensorflow/contrib/lite:build_def.bzl", "tflite_copts") - -cc_library( - name = "api", - srcs = [ - "error_reporter.cc", - "flatbuffer_conversions.cc", - "op_resolver.cc", - ], - hdrs = [ - "error_reporter.h", - "flatbuffer_conversions.h", - "op_resolver.h", - ], - copts = tflite_copts(), - deps = [ - "//tensorflow/contrib/lite/c:c_api_internal", - "//tensorflow/contrib/lite/schema:schema_fbs", - ], -) - -cc_test( - name = "error_reporter_test", - size = "small", - srcs = ["error_reporter_test.cc"], - deps = [ - ":api", - "@com_google_googletest//:gtest", - ], -) - -cc_test( - name = "op_resolver_test", - size = "small", - srcs = ["op_resolver_test.cc"], - deps = [ - ":api", - "@com_google_googletest//:gtest", - ], -) - -cc_test( - name = "flatbuffer_conversions_test", - size = "small", - srcs = ["flatbuffer_conversions_test.cc"], - deps = [ - ":api", - "//tensorflow/contrib/lite/c:c_api_internal", - "@com_google_googletest//:gtest", - ], -) diff --git a/tensorflow/contrib/lite/core/api/error_reporter.h b/tensorflow/contrib/lite/core/api/error_reporter.h deleted file mode 100644 index a2f780b003fc213d28ba29d7783dbfe99088cccc..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/core/api/error_reporter.h +++ /dev/null @@ -1,45 +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_CONTRIB_LITE_CORE_API_ERROR_REPORTER_H_ -#define TENSORFLOW_CONTRIB_LITE_CORE_API_ERROR_REPORTER_H_ - -#include - -namespace tflite { - -// A functor that reports error to supporting system. Invoked similar to -// printf. -// -// Usage: -// ErrorReporter foo; -// foo.Report("test %d", 5); -// or -// va_list args; -// foo.Report("test %d", args); // where args is va_list -// -// Subclass ErrorReporter to provide another reporting destination. -// For example, if you have a GUI program, you might redirect to a buffer -// that drives a GUI error log box. -class ErrorReporter { - public: - virtual ~ErrorReporter() {} - virtual int Report(const char* format, va_list args) = 0; - int Report(const char* format, ...); - int ReportError(void*, const char* format, ...); -}; - -} // namespace tflite - -#endif // TENSORFLOW_CONTRIB_LITE_CORE_API_ERROR_REPORTER_H_ diff --git a/tensorflow/contrib/lite/core/api/flatbuffer_conversions_test.cc b/tensorflow/contrib/lite/core/api/flatbuffer_conversions_test.cc deleted file mode 100644 index 8ae94e1d330c1958b857cff0b44c38108f153550..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/core/api/flatbuffer_conversions_test.cc +++ /dev/null @@ -1,124 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/contrib/lite/core/api/flatbuffer_conversions.h" - -#include - -#include -#include "tensorflow/contrib/lite/c/builtin_op_data.h" - -namespace tflite { -namespace { - -class MockErrorReporter : public ErrorReporter { - public: - MockErrorReporter() : buffer_size_(0) {} - int Report(const char* format, va_list args) override { - buffer_size_ = vsnprintf(buffer_, kBufferSize, format, args); - return buffer_size_; - } - char* GetBuffer() { return buffer_; } - int GetBufferSize() { return buffer_size_; } - - private: - static constexpr int kBufferSize = 256; - char buffer_[kBufferSize]; - int buffer_size_; -}; - -// Used to determine how the op data parsing function creates its working space. -class MockDataAllocator : public BuiltinDataAllocator { - public: - MockDataAllocator() : is_allocated_(false) {} - void* Allocate(size_t size) override { - EXPECT_FALSE(is_allocated_); - const int max_size = kBufferSize; - EXPECT_LE(size, max_size); - is_allocated_ = true; - return buffer_; - } - void Deallocate(void* data) override { is_allocated_ = false; } - - private: - static constexpr int kBufferSize = 1024; - char buffer_[kBufferSize]; - bool is_allocated_; -}; - -} // namespace - -TEST(FlatbufferConversions, TestParseOpDataConv) { - MockErrorReporter mock_reporter; - ErrorReporter* reporter = &mock_reporter; - MockDataAllocator mock_allocator; - - flatbuffers::FlatBufferBuilder builder; - flatbuffers::Offset conv_options = - CreateConv2DOptions(builder, Padding_SAME, 1, 2, - ActivationFunctionType_RELU, 3, 4) - .Union(); - flatbuffers::Offset conv_offset = CreateOperatorDirect( - builder, 0, nullptr, nullptr, BuiltinOptions_Conv2DOptions, conv_options, - nullptr, CustomOptionsFormat_FLEXBUFFERS, nullptr); - builder.Finish(conv_offset); - void* conv_pointer = builder.GetBufferPointer(); - const Operator* conv_op = flatbuffers::GetRoot(conv_pointer); - void* output_data = nullptr; - EXPECT_EQ(kTfLiteOk, ParseOpData(conv_op, BuiltinOperator_CONV_2D, reporter, - &mock_allocator, &output_data)); - EXPECT_NE(nullptr, output_data); - TfLiteConvParams* params = reinterpret_cast(output_data); - EXPECT_EQ(kTfLitePaddingSame, params->padding); - EXPECT_EQ(1, params->stride_width); - EXPECT_EQ(2, params->stride_height); - EXPECT_EQ(kTfLiteActRelu, params->activation); - EXPECT_EQ(3, params->dilation_width_factor); - EXPECT_EQ(4, params->dilation_height_factor); -} - -TEST(FlatbufferConversions, TestParseOpDataCustom) { - MockErrorReporter mock_reporter; - ErrorReporter* reporter = &mock_reporter; - MockDataAllocator mock_allocator; - - flatbuffers::FlatBufferBuilder builder; - flatbuffers::Offset null_options; - flatbuffers::Offset custom_offset = CreateOperatorDirect( - builder, 0, nullptr, nullptr, BuiltinOptions_NONE, null_options, nullptr, - CustomOptionsFormat_FLEXBUFFERS, nullptr); - builder.Finish(custom_offset); - void* custom_pointer = builder.GetBufferPointer(); - const Operator* custom_op = flatbuffers::GetRoot(custom_pointer); - void* output_data = nullptr; - EXPECT_EQ(kTfLiteOk, ParseOpData(custom_op, BuiltinOperator_CUSTOM, reporter, - &mock_allocator, &output_data)); - EXPECT_EQ(nullptr, output_data); -} - -TEST(FlatbufferConversions, TestConvertTensorType) { - MockErrorReporter mock_reporter; - ErrorReporter* reporter = &mock_reporter; - TfLiteType type; - EXPECT_EQ(kTfLiteOk, ConvertTensorType(TensorType_FLOAT32, &type, reporter)); - EXPECT_EQ(kTfLiteFloat32, type); -} - -} // namespace tflite - -int main(int argc, char** argv) { - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/contrib/lite/core/api/op_resolver.h b/tensorflow/contrib/lite/core/api/op_resolver.h deleted file mode 100644 index 5f5e6b27363b525094659719da9decf49dbeac45..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/core/api/op_resolver.h +++ /dev/null @@ -1,47 +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_LITE_CORE_API_OP_RESOLVER_H_ -#define TENSORFLOW_CONTRIB_LITE_CORE_API_OP_RESOLVER_H_ - -#include "tensorflow/contrib/lite/c/c_api_internal.h" -#include "tensorflow/contrib/lite/core/api/error_reporter.h" -#include "tensorflow/contrib/lite/schema/schema_generated.h" - -namespace tflite { - -// Abstract interface that returns TfLiteRegistrations given op codes or custom -// op names. This is the mechanism that ops being referenced in the flatbuffer -// model are mapped to executable function pointers (TfLiteRegistrations). -class OpResolver { - public: - // Finds the op registration for a builtin operator by enum code. - virtual const TfLiteRegistration* FindOp(tflite::BuiltinOperator op, - int version) const = 0; - // Finds the op registration of a custom operator by op name. - virtual const TfLiteRegistration* FindOp(const char* op, - int version) const = 0; - virtual ~OpResolver() {} -}; - -// Handles the logic for converting between an OperatorCode structure extracted -// from a flatbuffer and information about a registered operator implementation. -TfLiteStatus GetRegistrationFromOpCode(const OperatorCode* opcode, - const OpResolver& op_resolver, - ErrorReporter* error_reporter, - const TfLiteRegistration** registration); - -} // namespace tflite - -#endif // TENSORFLOW_CONTRIB_LITE_CORE_API_OP_RESOLVER_H_ diff --git a/tensorflow/contrib/lite/delegates/flex/BUILD b/tensorflow/contrib/lite/delegates/flex/BUILD deleted file mode 100644 index 2f866eaecb801695d800565e195f959d55a88201..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/delegates/flex/BUILD +++ /dev/null @@ -1,232 +0,0 @@ -# -# This is a TF Lite delegate that is powered by TensorFlow's Eager. -# -package(default_visibility = [ - "//visibility:private", -]) - -licenses(["notice"]) # Apache 2.0 - -load("//tensorflow:tensorflow.bzl", "tf_cc_test") - -cc_library( - name = "buffer_map", - srcs = ["buffer_map.cc"], - hdrs = ["buffer_map.h"], - deps = [ - ":util", - "//tensorflow/c:c_api_internal", - "//tensorflow/contrib/lite/c:c_api_internal", - "//tensorflow/contrib/lite:kernel_api", - ] + select({ - "//tensorflow:android": [ - "//tensorflow/core:android_tensorflow_lib_lite", - ], - "//conditions:default": [ - "//tensorflow/core:framework", - "//tensorflow/core:protos_all_cc", - ], - }), -) - -tf_cc_test( - name = "buffer_map_test", - size = "small", - srcs = ["buffer_map_test.cc"], - deps = [ - ":buffer_map", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite:util", - "//tensorflow/contrib/lite/testing:util", - "@com_google_googletest//:gtest", - ], -) - -# Delegate implementation that pulls in the standard set of TensorFlow ops and -# kernels. -cc_library( - name = "delegate", - hdrs = [ - "delegate.h", - ], - visibility = ["//visibility:public"], - deps = [ - ":delegate_only_runtime", - ] + select({ - "//tensorflow:android": [ - "//tensorflow/core:android_tensorflow_lib", - ], - "//conditions:default": [ - "//tensorflow/core:tensorflow", - ], - }), - alwayslink = 1, -) - -# Delegate implementation that does *not* pull in the standard set of TensorFlow -# ops and kernels. -cc_library( - name = "delegate_only_runtime", - srcs = [ - "delegate.cc", - ], - hdrs = [ - "delegate.h", - ], - visibility = ["//visibility:public"], - deps = [ - ":buffer_map", - ":delegate_data", - ":kernel", - ":util", - "//tensorflow/contrib/lite/c:c_api_internal", - "//tensorflow/contrib/lite:kernel_api", - "//tensorflow/contrib/lite:util", - ] + select({ - "//tensorflow:android": [ - "//tensorflow/core:android_tensorflow_lib_lite", - ], - "//conditions:default": [ - "//tensorflow/core:lib", - ], - }), - alwayslink = 1, -) - -tf_cc_test( - name = "delegate_test", - size = "small", - srcs = ["delegate_test.cc"], - deps = [ - ":delegate", - ":test_util", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -cc_library( - name = "delegate_data", - srcs = ["delegate_data.cc"], - hdrs = ["delegate_data.h"], - deps = [ - ":buffer_map", - "//tensorflow/core/common_runtime/eager:context", - ] + select({ - "//tensorflow:android": [ - "//tensorflow/core:android_tensorflow_lib_lite", - ], - "//conditions:default": [ - "//tensorflow/core:core_cpu", - "//tensorflow/core:lib", - ], - }), -) - -tf_cc_test( - name = "delegate_data_test", - size = "small", - srcs = ["delegate_data_test.cc"], - deps = [ - ":delegate_data", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite:util", - "//tensorflow/contrib/lite/c:c_api_internal", - "//tensorflow/contrib/lite/testing:util", - "@com_google_googletest//:gtest", - ], -) - -cc_library( - name = "kernel", - srcs = ["kernel.cc"], - hdrs = ["kernel.h"], - deps = [ - ":delegate_data", - ":util", - "@flatbuffers", - "//tensorflow/contrib/lite/c:c_api_internal", - "//tensorflow/contrib/lite:kernel_api", - "//tensorflow/contrib/lite:string", - "//tensorflow/contrib/lite/kernels:kernel_util", - "//tensorflow/core/common_runtime/eager:context", - "//tensorflow/core/common_runtime/eager:execute", - "//tensorflow/core/common_runtime/eager:tensor_handle", - ] + select({ - # TODO(b/111881878): The android_tensorflow_lib target pulls in the full - # set of core TensorFlow kernels. We may want to revisit this dependency - # to allow selective registration via build targets. - "//tensorflow:android": [ - "//tensorflow/core:android_tensorflow_lib_lite", - ], - "//conditions:default": [ - "//tensorflow/core:lib", - "//tensorflow/core:protos_all_cc", - "//tensorflow/core:framework", - ], - }), -) - -tf_cc_test( - name = "kernel_test", - size = "small", - srcs = ["kernel_test.cc"], - deps = [ - ":delegate_data", - ":kernel", - ":test_util", - "@com_google_googletest//:gtest", - ] + select({ - "//tensorflow:android": [ - "//tensorflow/core:android_tensorflow_lib", - ], - "//conditions:default": [ - "//tensorflow/core:tensorflow", - ], - }), -) - -cc_library( - name = "test_util", - testonly = True, - srcs = ["test_util.cc"], - hdrs = ["test_util.h"], - deps = [ - "//tensorflow/c:c_api_internal", - "//tensorflow/contrib/lite:string", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_absl//absl/memory", - "@flatbuffers", - ], -) - -cc_library( - name = "util", - srcs = ["util.cc"], - hdrs = ["util.h"], - deps = [ - "//tensorflow/c:c_api_internal", - "//tensorflow/contrib/lite/c:c_api_internal", - "//tensorflow/contrib/lite:kernel_api", - ] + select({ - "//tensorflow:android": [ - "//tensorflow/core:android_tensorflow_lib_lite", - ], - "//conditions:default": [ - "//tensorflow/core:lib", - "//tensorflow/core:framework", - ], - }), -) - -tf_cc_test( - name = "util_test", - size = "small", - srcs = ["util_test.cc"], - deps = [ - ":util", - "//tensorflow/contrib/lite:string", - "//tensorflow/contrib/lite/testing:util", - "@com_google_googletest//:gtest", - ], -) diff --git a/tensorflow/contrib/lite/delegates/flex/buffer_map.cc b/tensorflow/contrib/lite/delegates/flex/buffer_map.cc deleted file mode 100644 index 63e39196d96a176eca105e7b11107ab52fe528dd..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/delegates/flex/buffer_map.cc +++ /dev/null @@ -1,111 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include "tensorflow/contrib/lite/delegates/flex/buffer_map.h" - -#include "tensorflow/c/c_api_internal.h" -#include "tensorflow/contrib/lite/delegates/flex/util.h" -#include "tensorflow/core/framework/allocation_description.pb.h" -#include "tensorflow/core/framework/log_memory.h" - -namespace tflite { -namespace flex { -namespace { -// A tensor buffer that is allocated, deallocated and populated by TF Lite. -class TfLiteTensorBuffer : public tensorflow::TensorBuffer { - public: - explicit TfLiteTensorBuffer(const TfLiteTensor* tensor) { - len_ = tensor->bytes; - // TODO(ahentz): if we can guarantee that TF Lite allocated tensors with - // the same alignment as TensorFlow (EIGEN_MAX_ALIGN_BYTES), then we can - // potentially eliminate the copy below. - data_ = - tensorflow::cpu_allocator()->AllocateRaw(EIGEN_MAX_ALIGN_BYTES, len_); - if (data_ != nullptr) { - if (tensorflow::LogMemory::IsEnabled()) { - tensorflow::LogMemory::RecordRawAllocation( - "TfLiteTensorBuffer_New", - tensorflow::LogMemory::EXTERNAL_TENSOR_ALLOCATION_STEP_ID, len_, - data_, tensorflow::cpu_allocator()); - } - std::memcpy(data_, tensor->data.raw, tensor->bytes); - } - } - - ~TfLiteTensorBuffer() override { - if (tensorflow::LogMemory::IsEnabled() && data_ != nullptr) { - tensorflow::LogMemory::RecordRawDeallocation( - "TfLiteTensorBuffer_Delete", - tensorflow::LogMemory::EXTERNAL_TENSOR_ALLOCATION_STEP_ID, data_, - tensorflow::cpu_allocator(), false); - } - tensorflow::cpu_allocator()->DeallocateRaw(data_); - } - - void* data() const override { return data_; } - size_t size() const override { return len_; } - - TensorBuffer* root_buffer() override { return this; } - void FillAllocationDescription( - tensorflow::AllocationDescription* proto) const override { - tensorflow::int64 rb = size(); - proto->set_requested_bytes(rb); - proto->set_allocator_name(tensorflow::cpu_allocator()->Name()); - } - - // Prevents input forwarding from mutating this buffer. - bool OwnsMemory() const override { return false; } - - private: - void* data_; - size_t len_; -}; -} // namespace - -BufferMap::BufferMap() {} - -BufferMap::~BufferMap() {} - -bool BufferMap::HasTensor(int tensor_index) const { - return id_to_tensor_.count(tensor_index) != 0; -} - -tensorflow::Tensor BufferMap::GetTensor(int tensor_index) const { - return id_to_tensor_.at(tensor_index); -} - -void BufferMap::SetFromTfLite(int tensor_index, const TfLiteTensor* tensor) { - tensorflow::TensorShape shape; - int num_dims = tensor->dims->size; - for (int i = 0; i < num_dims; ++i) { - shape.AddDim(tensor->dims->data[i]); - } - // TODO(ahentz): we assume this is a new tensor and allocate a new buffer - // for it. This is not always the best approach. For example, this might - // be a reallocation after resizing tensors. In that case we would be - // preferable to somehow reuse the buffer. - auto* buf = new TfLiteTensorBuffer(tensor); - tensorflow::Tensor t = tensorflow::TensorCApi::MakeTensor( - GetTensorFlowDataType(tensor->type), shape, buf); - buf->Unref(); - - SetFromTensorFlow(tensor_index, std::move(t)); -} - -void BufferMap::SetFromTensorFlow(int tensor_index, tensorflow::Tensor tensor) { - id_to_tensor_[tensor_index] = std::move(tensor); -} - -} // namespace flex -} // namespace tflite diff --git a/tensorflow/contrib/lite/delegates/flex/buffer_map.h b/tensorflow/contrib/lite/delegates/flex/buffer_map.h deleted file mode 100644 index 4ce886568a55773971bc0543ec973ec84c0aac1b..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/delegates/flex/buffer_map.h +++ /dev/null @@ -1,61 +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_LITE_DELEGATES_FLEX_BUFFER_MAP_H_ -#define TENSORFLOW_CONTRIB_LITE_DELEGATES_FLEX_BUFFER_MAP_H_ - -#include - -#include "tensorflow/contrib/lite/c/c_api_internal.h" -#include "tensorflow/core/framework/tensor.h" - -namespace tflite { -namespace flex { - -// Maps a TF Lite tensor index into a TensorFlow tensor. -// -// The TF Lite interpreter assigns integer indices to each of its tensors, but -// the Flex delegate deals in terms of TensorFlow tensors. This class maps -// from indices to tensors and allows the creation of new tensors to be -// associated with a given index. -class BufferMap { - public: - BufferMap(); - ~BufferMap(); - - // Returns true if the given 'tensor_index' has a corresponding - // tensorflow::Tensor. - bool HasTensor(int tensor_index) const; - - // Returns the tensorflow::Tensor associated with the given 'tensor_index'. - // Precondition: HasTensor() is true. - tensorflow::Tensor GetTensor(int tensor_index) const; - - // Associates the given tensorflow::Tensor with the given 'tensor_index'. - // Note that tensorflow Tensors share data buffers, so this method is only a - // shallow copy. - void SetFromTensorFlow(int tensor_index, tensorflow::Tensor tensor); - - // Same as above but creates a new tensorflow::Tensor with a copy of the - // given TfLiteTensor's data. - void SetFromTfLite(int tensor_index, const TfLiteTensor* tensor); - - private: - std::map id_to_tensor_; -}; - -} // namespace flex -} // namespace tflite - -#endif // TENSORFLOW_CONTRIB_LITE_DELEGATES_FLEX_BUFFER_MAP_H_ diff --git a/tensorflow/contrib/lite/delegates/flex/buffer_map_test.cc b/tensorflow/contrib/lite/delegates/flex/buffer_map_test.cc deleted file mode 100644 index bb80e25e8076bb95782e4137945ad1c7cd178aee..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/delegates/flex/buffer_map_test.cc +++ /dev/null @@ -1,174 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include "tensorflow/contrib/lite/delegates/flex/buffer_map.h" - -#include -#include -#include "tensorflow/contrib/lite/interpreter.h" -#include "tensorflow/contrib/lite/testing/util.h" -#include "tensorflow/contrib/lite/util.h" - -namespace tflite { -namespace flex { -namespace { - -using ::testing::ElementsAre; - -// A bit of RAII to simplify handling of TfLiteTensors in the tests. -using UniqueTfLiteTensor = - std::unique_ptr>; - -template -UniqueTfLiteTensor MakeLiteTensor(const std::vector& shape, - const std::vector& data) { - auto tensor = UniqueTfLiteTensor(new TfLiteTensor, [](TfLiteTensor* t) { - TfLiteTensorDataFree(t); - TfLiteIntArrayFree(t->dims); - delete t; - }); - tensor->allocation_type = kTfLiteDynamic; - tensor->type = typeToTfLiteType(); - tensor->dims = ConvertVectorToTfLiteIntArray(shape); - tensor->data.raw = nullptr; - TfLiteTensorRealloc(data.size() * sizeof(T), tensor.get()); - memcpy(tensor->data.raw, data.data(), data.size() * sizeof(T)); - return tensor; -} - -template -tensorflow::Tensor MakeTensor(const std::vector& shape, - const std::vector& data) { - BufferMap buffer_map; // BufferMap is the easiest way to build the tensor. - UniqueTfLiteTensor t1 = MakeLiteTensor(shape, data); - buffer_map.SetFromTfLite(0, t1.get()); - return buffer_map.GetTensor(0); -} - -std::vector GetTensorShape(const tensorflow::Tensor& t) { - std::vector shape(t.dims()); - for (int i = 0; i < t.dims(); ++i) { - shape[i] = t.dim_size(i); - } - return shape; -} - -template -std::vector GetTensorData(const tensorflow::Tensor& t) { - const T* data = t.flat().data(); - return std::vector(data, data + t.NumElements()); -} - -TEST(BufferMapTest, EmptyBuffer) { - BufferMap buffer_map; - EXPECT_FALSE(buffer_map.HasTensor(0)); -} - -TEST(BufferMapTest, SetFromTfLite) { - BufferMap buffer_map; - - UniqueTfLiteTensor t = - MakeLiteTensor({1, 2, 1, 3}, {0, 0, 0, 0.123f, 0, 0}); - buffer_map.SetFromTfLite(0, t.get()); - ASSERT_TRUE(buffer_map.HasTensor(0)); - - EXPECT_THAT(GetTensorData(buffer_map.GetTensor(0)), - ElementsAre(0, 0, 0, 0.123f, 0, 0)); - - // Also check details of the tensor. - tensorflow::Tensor out_tensor = buffer_map.GetTensor(0); - ASSERT_EQ(out_tensor.dtype(), tensorflow::DT_FLOAT); - ASSERT_EQ(out_tensor.NumElements(), 6); - ASSERT_THAT(GetTensorShape(out_tensor), ElementsAre(1, 2, 1, 3)); -} - -TEST(BufferMapTest, SetFromTfLiteTwice) { - UniqueTfLiteTensor t1 = - MakeLiteTensor({1, 2, 1, 3}, {0, 0, 0, 0.123f, 0, 0}); - UniqueTfLiteTensor t2 = - MakeLiteTensor({1, 2, 4}, {0, 0, 0, 3, 0, 0, 1, 2}); - - BufferMap buffer_map; - buffer_map.SetFromTfLite(0, t1.get()); - buffer_map.SetFromTfLite(0, t2.get()); - - EXPECT_THAT(GetTensorData(buffer_map.GetTensor(0)), - ElementsAre(0, 0, 0, 3, 0, 0, 1, 2)); -} - -TEST(BufferMapTest, SetFromTensorFlow) { - tensorflow::Tensor t1 = - MakeTensor({1, 2, 1, 3}, {0, 0, 0, 0.123f, 0, 0}); - - BufferMap buffer_map; - buffer_map.SetFromTensorFlow(0, t1); - - EXPECT_THAT(GetTensorData(buffer_map.GetTensor(0)), - ElementsAre(0, 0, 0, 0.123f, 0, 0)); - - // Also check details of the tensor. - tensorflow::Tensor out_tensor = buffer_map.GetTensor(0); - ASSERT_EQ(out_tensor.dtype(), tensorflow::DT_FLOAT); - ASSERT_EQ(out_tensor.NumElements(), 6); - ASSERT_THAT(GetTensorShape(out_tensor), ElementsAre(1, 2, 1, 3)); -} - -TEST(BufferMapTest, SetFromTensorFlowTwice) { - tensorflow::Tensor t1 = - MakeTensor({1, 2, 1, 3}, {0, 0, 0, 0.123f, 0, 0}); - tensorflow::Tensor t2 = MakeTensor({1, 2, 4}, {0, 0, 0, 3, 0, 0, 1, 2}); - BufferMap buffer_map; - buffer_map.SetFromTensorFlow(0, t1); - buffer_map.SetFromTensorFlow(0, t2); - - EXPECT_THAT(GetTensorData(buffer_map.GetTensor(0)), - ElementsAre(0, 0, 0, 3, 0, 0, 1, 2)); -} - -TEST(BufferMapTest, TfLiteOverwritesTensorFlow) { - tensorflow::Tensor t1 = - MakeTensor({1, 2, 1, 3}, {0, 0, 0, 0.123f, 0, 0}); - UniqueTfLiteTensor t2 = - MakeLiteTensor({1, 2, 4}, {0, 0, 0, 3, 0, 0, 1, 2}); - - BufferMap buffer_map; - buffer_map.SetFromTensorFlow(0, t1); - buffer_map.SetFromTfLite(0, t2.get()); - - EXPECT_THAT(GetTensorData(buffer_map.GetTensor(0)), - ElementsAre(0, 0, 0, 3, 0, 0, 1, 2)); -} - -TEST(BufferMapTest, TensorFlowOverwritesTfLite) { - tensorflow::Tensor t1 = - MakeTensor({1, 2, 1, 3}, {0, 0, 0, 0.123f, 0, 0}); - UniqueTfLiteTensor t2 = - MakeLiteTensor({1, 2, 4}, {0, 0, 0, 3, 0, 0, 1, 2}); - BufferMap buffer_map; - buffer_map.SetFromTfLite(0, t2.get()); - buffer_map.SetFromTensorFlow(0, t1); - - EXPECT_THAT(GetTensorData(buffer_map.GetTensor(0)), - ElementsAre(0, 0, 0, 0.123f, 0, 0)); -} - -} // namespace -} // namespace flex -} // namespace tflite - -int main(int argc, char** argv) { - ::tflite::LogToStderr(); - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/contrib/lite/delegates/flex/delegate.cc b/tensorflow/contrib/lite/delegates/flex/delegate.cc deleted file mode 100644 index c72b0cf51383897ce3afec0c39ed6bfe178d88c1..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/delegates/flex/delegate.cc +++ /dev/null @@ -1,117 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include "tensorflow/contrib/lite/delegates/flex/delegate.h" - -#include - -#include "tensorflow/contrib/lite/context_util.h" -#include "tensorflow/contrib/lite/delegates/flex/buffer_map.h" -#include "tensorflow/contrib/lite/delegates/flex/kernel.h" -#include "tensorflow/contrib/lite/delegates/flex/util.h" -#include "tensorflow/contrib/lite/util.h" -#include "tensorflow/core/lib/core/status.h" - -namespace tflite { -namespace flex { -namespace delegate { - -TfLiteStatus Prepare(TfLiteContext* context, TfLiteDelegate* delegate) { - // Get the nodes in the current execution plan. Interpreter owns this array. - TfLiteIntArray* plan; - TF_LITE_ENSURE_STATUS(context->GetExecutionPlan(context, &plan)); - - // Add all custom ops starting with "Flex" to list of supported nodes. - std::vector supported_nodes; - for (int node_index : TfLiteIntArrayView(plan)) { - TfLiteNode* node; - TfLiteRegistration* registration; - TF_LITE_ENSURE_STATUS(context->GetNodeAndRegistration( - context, node_index, &node, ®istration)); - - if (IsFlexOp(registration->custom_name)) { - supported_nodes.push_back(node_index); - } - } - - // Request TFLite to partition the graph and make kernels for each independent - // subgraph. - TfLiteIntArray* size_and_nodes = - ConvertVectorToTfLiteIntArray(supported_nodes); - context->ReplaceSubgraphsWithDelegateKernels(context, GetKernel(), - size_and_nodes, delegate); - TfLiteIntArrayFree(size_and_nodes); - return kTfLiteOk; -} - -TfLiteStatus CopyFromBufferHandle(TfLiteContext* context, - TfLiteDelegate* delegate, - TfLiteBufferHandle buffer_handle, void* data, - size_t size) { - BufferMap* buffer_map = - reinterpret_cast(delegate->data_)->GetBufferMap(context); - - if (!buffer_map->HasTensor(buffer_handle)) { - context->ReportError(context, "Invalid tensor index %d.", buffer_handle); - return kTfLiteError; - } - - tensorflow::Tensor t = buffer_map->GetTensor(buffer_handle); - tensorflow::StringPiece t_data = t.tensor_data(); - - if (size != t_data.size()) { - context->ReportError( - context, "Not enough space to store TensorFlow's aligned buffer."); - return kTfLiteError; - } - - memcpy(data, t_data.data(), t_data.size()); - return kTfLiteOk; -} - -} // namespace delegate -} // namespace flex - -// Corresponding weak declaration found in lite/model.cc. -std::unique_ptr -AcquireFlexDelegate() { - return std::unique_ptr( - tflite::FlexDelegate::Create().release(), [](TfLiteDelegate* delegate) { - delete reinterpret_cast(delegate); - }); -} - -std::unique_ptr FlexDelegate::Create() { - std::unique_ptr delegate_data; - if (!flex::DelegateData::Create(&delegate_data).ok()) { - fprintf(stderr, "Unable to initialize TensorFlow context.\n"); - return nullptr; - } - - return std::unique_ptr( - new FlexDelegate(std::move(delegate_data))); -} - -FlexDelegate::FlexDelegate(std::unique_ptr delegate_data) - : TfLiteDelegate{ - /*data_=*/delegate_data.get(), - /*nullptr,*/ &flex::delegate::Prepare, - /*CopyFromBufferHandle=*/&flex::delegate::CopyFromBufferHandle, - /*CopyToBufferHandle=*/nullptr, - /*FreeBufferHandle=*/nullptr}, - delegate_data_(std::move(delegate_data)) {} - -FlexDelegate::~FlexDelegate() {} - -} // namespace tflite diff --git a/tensorflow/contrib/lite/delegates/flex/delegate_data.cc b/tensorflow/contrib/lite/delegates/flex/delegate_data.cc deleted file mode 100644 index 8f985f770cfba9fc6a7184cfdb0a35e9e6c754af..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/delegates/flex/delegate_data.cc +++ /dev/null @@ -1,47 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include "tensorflow/contrib/lite/delegates/flex/delegate_data.h" - -#include "tensorflow/core/common_runtime/device_factory.h" -#include "tensorflow/core/lib/core/status.h" - -namespace tflite { -namespace flex { -tensorflow::Status DelegateData::Create(std::unique_ptr* data) { - std::vector devices; - - TF_RETURN_IF_ERROR(tensorflow::DeviceFactory::AddDevices( - tensorflow::SessionOptions(), "/job:localhost/replica:0/task:0", - &devices)); - - std::unique_ptr device_mgr( - new tensorflow::DeviceMgr(devices)); - // Note that Rendezvous is ref-counted so it will be automatically deleted. - tensorflow::Rendezvous* rendezvous = - new tensorflow::IntraProcessRendezvous(device_mgr.get()); - data->reset(new DelegateData(new tensorflow::EagerContext( - tensorflow::SessionOptions(), - tensorflow::ContextDevicePlacementPolicy::DEVICE_PLACEMENT_SILENT, - /*async=*/false, std::move(device_mgr), rendezvous))); - return tensorflow::Status(); -} - -DelegateData::DelegateData(tensorflow::EagerContext* eager_context) - : eager_context_(eager_context) {} - -DelegateData::~DelegateData() {} - -} // namespace flex -} // namespace tflite diff --git a/tensorflow/contrib/lite/delegates/flex/delegate_data.h b/tensorflow/contrib/lite/delegates/flex/delegate_data.h deleted file mode 100644 index 8d75f0b0efe758074d035f0ebcf0f5f12602323b..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/delegates/flex/delegate_data.h +++ /dev/null @@ -1,52 +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_LITE_DELEGATES_FLEX_DELEGATE_DATA_H_ -#define TENSORFLOW_CONTRIB_LITE_DELEGATES_FLEX_DELEGATE_DATA_H_ - -#include "tensorflow/contrib/lite/delegates/flex/buffer_map.h" -#include "tensorflow/core/common_runtime/eager/context.h" - -namespace tflite { -namespace flex { - -// Data kept by the Flex delegate for the lifetime of an Interpreter. -class DelegateData { - public: - // Create a new DelegateData, initialized with a newly-created EagerContext. - static tensorflow::Status Create(std::unique_ptr* data); - - ~DelegateData(); - - // The EagerContext that is required for execution of Flex Ops. - tensorflow::EagerContext* GetEagerContext() { return eager_context_.get(); } - - // Map from TF Lite tensor index to TensorFlow tensor for a given context. - BufferMap* GetBufferMap(const TfLiteContext* context) { - return &buffer_map_[context]; - } - - private: - explicit DelegateData(tensorflow::EagerContext* eager_context); - - std::unique_ptr eager_context_; - // TODO(b/112439500): Clean up stale BufferMap instances after adding the - // necessary cleanup hook from a TfLiteContext to a TfLiteDelegate. - std::unordered_map buffer_map_; -}; - -} // namespace flex -} // namespace tflite - -#endif // TENSORFLOW_CONTRIB_LITE_DELEGATES_FLEX_DELEGATE_DATA_H_ diff --git a/tensorflow/contrib/lite/delegates/flex/delegate_data_test.cc b/tensorflow/contrib/lite/delegates/flex/delegate_data_test.cc deleted file mode 100644 index 30b10f435a23785f88e2645714a414501bc2fab9..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/delegates/flex/delegate_data_test.cc +++ /dev/null @@ -1,49 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include "tensorflow/contrib/lite/delegates/flex/delegate_data.h" - -#include -#include -#include "tensorflow/contrib/lite/c/c_api_internal.h" -#include "tensorflow/contrib/lite/testing/util.h" - -namespace tflite { -namespace flex { -namespace { - -TEST(DelegateDataTest, Basic) { - std::unique_ptr data; - // We only check for success because it is hard to make initialization fail. - // It only happens if we manage to not link the CPU device factory into the - // binary. - EXPECT_TRUE(DelegateData::Create(&data).ok()); - - TfLiteContext dummy_context1 = {}; - TfLiteContext dummy_context2 = {}; - EXPECT_NE(data->GetEagerContext(), nullptr); - EXPECT_NE(data->GetBufferMap(&dummy_context1), nullptr); - EXPECT_NE(data->GetBufferMap(&dummy_context1), - data->GetBufferMap(&dummy_context2)); -} - -} // namespace -} // namespace flex -} // namespace tflite - -int main(int argc, char** argv) { - ::tflite::LogToStderr(); - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/contrib/lite/delegates/flex/kernel.cc b/tensorflow/contrib/lite/delegates/flex/kernel.cc deleted file mode 100644 index e4f1aea990da97da08a3e5adf2dd70307b20fe88..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/delegates/flex/kernel.cc +++ /dev/null @@ -1,299 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include "tensorflow/contrib/lite/delegates/flex/kernel.h" - -#include "flatbuffers/flexbuffers.h" // TF:flatbuffers -#include "tensorflow/contrib/lite/builtin_ops.h" -#include "tensorflow/contrib/lite/c/c_api_internal.h" -#include "tensorflow/contrib/lite/context_util.h" -#include "tensorflow/contrib/lite/delegates/flex/delegate_data.h" -#include "tensorflow/contrib/lite/delegates/flex/util.h" -#include "tensorflow/contrib/lite/kernels/kernel_util.h" -#include "tensorflow/contrib/lite/string.h" -#include "tensorflow/core/common_runtime/eager/context.h" -#include "tensorflow/core/common_runtime/eager/execute.h" -#include "tensorflow/core/common_runtime/eager/tensor_handle.h" -#include "tensorflow/core/framework/node_def.pb.h" -#include "tensorflow/core/framework/node_def_util.h" - -// Note: this is part of TF Lite's Flex delegation code which is to be -// completed soon. - -// This is the TF Lite op that is created by the flex delegate to handle -// execution of a supported subgraph. The usual flow is that the delegate -// informs the interpreter of supported nodes in a graph, and each supported -// subgraph is replaced with one instance of this kernel. -// -// The kernel is initialized with TfLiteDelegateParams from which we retrieve -// the global EagerContext and BufferMap, as well as a list of inputs and -// outputs to the subgraph. Those are used to build the OpData, with a list of -// TensorFlow Ops that should be executed in order (which we call an OpNode). -// -// For each node included in the subgraph, we query the interpreter and -// retrieve the associated NodeDef, which is then used to configure the -// corresponding TensorFlow/Eager Op. - -namespace tflite { -namespace flex { -namespace kernel { - -// Controls the lifetime of tensor handles in a vector. -class VectorOfHandles { - public: - explicit VectorOfHandles(int num_elements) : vector_(num_elements, nullptr) {} - - ~VectorOfHandles() { - for (auto* handle : vector_) { - if (handle) handle->Unref(); - } - } - - tensorflow::gtl::InlinedVector* GetVector() { - return &vector_; - } - - tensorflow::TensorHandle* GetHandle(int index) { return vector_[index]; } - - private: - tensorflow::gtl::InlinedVector vector_; -}; - -// Executes the TensorFlow op given by 'op_name', with the attributes specified -// in 'nodedef'. Inputs and outputs are given as indices into the 'buffer_map'. -tensorflow::Status ExecuteFlexOp(tensorflow::EagerContext* eager_context, - BufferMap* buffer_map, const string& op_name, - const tensorflow::NodeDef& nodedef, - const std::vector& inputs, - const std::vector& outputs) { - const tensorflow::AttrTypeMap* attr_types; - TF_RETURN_WITH_CONTEXT_IF_ERROR( - tensorflow::AttrTypeMapForOp(op_name.c_str(), &attr_types), - " (while processing attributes of '", op_name, "')"); - - tensorflow::EagerOperation op(eager_context, op_name.c_str(), attr_types); - for (const auto& attr : nodedef.attr()) { - op.MutableAttrs()->Set(attr.first, attr.second); - } - - for (int input_index : inputs) { - if (!buffer_map->HasTensor(input_index)) { - return tensorflow::errors::Internal( - "Cannot read from invalid tensor index ", input_index); - } - auto* handle = new tensorflow::TensorHandle( - buffer_map->GetTensor(input_index), nullptr, nullptr, nullptr); - op.AddInput(handle); - handle->Unref(); - } - - int num_retvals = outputs.size(); - VectorOfHandles retvals(num_retvals); - TF_RETURN_WITH_CONTEXT_IF_ERROR( - EagerExecute(&op, retvals.GetVector(), &num_retvals), - " (while executing '", op_name, "' via Eager)"); - - if (num_retvals != outputs.size()) { - return tensorflow::errors::Internal( - "Unexpected number of outputs from EagerExecute"); - } - - for (int i = 0; i < num_retvals; ++i) { - const tensorflow::Tensor* tensor = nullptr; - TF_RETURN_IF_ERROR(retvals.GetHandle(i)->Tensor(&tensor)); - buffer_map->SetFromTensorFlow(outputs[i], *tensor); - } - - return tensorflow::Status::OK(); -} - -// A single node within the larger 'op'. Note that this kernel executes many -// TensorFlow ops within a single TF Lite op. -struct OpNode { - // The name of the TensorFlow op to execute. - string name; - // The corresponding NodeDef, containing the attributes for the op. - tensorflow::NodeDef nodedef; - // List of inputs, as TF Lite tensor indices. - std::vector inputs; - // List of outputs, as TF Lite tensor indices. - std::vector outputs; -}; - -// The Larger 'op', which contains all the nodes in a supported subgraph. -struct OpData { - tensorflow::EagerContext* eager_context; - BufferMap* buffer_map; - std::vector nodes; - std::vector subgraph_inputs; - std::vector subgraph_outputs; -}; - -void* Init(TfLiteContext* context, const char* buffer, size_t length) { - auto* op_data = new OpData; - - const TfLiteDelegateParams* params = - reinterpret_cast(buffer); - CHECK(params); - CHECK(params->delegate); - CHECK(params->delegate->data_); - op_data->eager_context = - reinterpret_cast(params->delegate->data_) - ->GetEagerContext(); - op_data->buffer_map = reinterpret_cast(params->delegate->data_) - ->GetBufferMap(context); - - CHECK(params->output_tensors); - for (auto tensor_index : TfLiteIntArrayView(params->output_tensors)) { - op_data->subgraph_outputs.push_back(tensor_index); - } - - CHECK(params->input_tensors); - for (auto tensor_index : TfLiteIntArrayView(params->input_tensors)) { - op_data->subgraph_inputs.push_back(tensor_index); - } - - CHECK(params->nodes_to_replace); - for (auto node_index : TfLiteIntArrayView(params->nodes_to_replace)) { - TfLiteNode* node; - TfLiteRegistration* reg; - context->GetNodeAndRegistration(context, node_index, &node, ®); - - op_data->nodes.push_back(OpNode()); - OpNode& node_data = op_data->nodes.back(); - - node_data.name = ""; - if (node->custom_initial_data) { - // The flexbuffer contains a vector where the first elements is the - // op name and the second is a serialized NodeDef. - const flexbuffers::Vector& v = - flexbuffers::GetRoot( - reinterpret_cast(node->custom_initial_data), - node->custom_initial_data_size) - .AsVector(); - - node_data.name = v[0].AsString().str(); - if (!node_data.nodedef.ParseFromString(v[1].AsString().str())) { - // We will just leave the nodedef empty and error out in Eval(). - node_data.nodedef.Clear(); - } - } - - // Fill NodeDef with defaults if it's a valid op. - const tensorflow::OpRegistrationData* op_reg_data; - auto tf_status = tensorflow::OpRegistry::Global()->LookUp( - node_data.nodedef.op(), &op_reg_data); - if (tf_status.ok()) { - AddDefaultsToNodeDef(op_reg_data->op_def, &node_data.nodedef); - } - - for (auto input_index : TfLiteIntArrayView(node->inputs)) { - node_data.inputs.push_back(input_index); - } - for (auto output_index : TfLiteIntArrayView(node->outputs)) { - node_data.outputs.push_back(output_index); - } - } - - return op_data; -} - -void Free(TfLiteContext* context, void* buffer) { - delete reinterpret_cast(buffer); -} - -TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { - const auto* op_data = reinterpret_cast(node->user_data); - TF_LITE_ENSURE_MSG( - context, op_data->eager_context != nullptr, - "Failed to initialize eager context. This often happens when a CPU " - "device has not been registered, presumably because some symbols from " - "tensorflow/core:core_cpu_impl were not linked into the binary."); - - // Whenever we find a constant tensor, insert it in the buffer map. - BufferMap* buffer_map = op_data->buffer_map; - for (auto tensor_index : op_data->subgraph_inputs) { - TfLiteTensor* tensor = &context->tensors[tensor_index]; - if (IsConstantTensor(tensor)) { - if (!buffer_map->HasTensor(tensor_index)) { - buffer_map->SetFromTfLite(tensor_index, tensor); - } - } - } - - // All output tensors are allocated by TensorFlow/Eager, so we - // mark them as kTfLiteDynamic. - for (auto tensor_index : op_data->subgraph_outputs) { - SetTensorToDynamic(&context->tensors[tensor_index]); - } - - return kTfLiteOk; -} - -TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { - const auto* op_data = reinterpret_cast(node->user_data); - BufferMap* buffer_map = op_data->buffer_map; - tensorflow::EagerContext* eager_context = op_data->eager_context; - - // Insert a tensor in the buffer map for all inputs that are not constant. - // Constants were handled in Prepare() already. - for (auto tensor_index : op_data->subgraph_inputs) { - TfLiteTensor* tensor = &context->tensors[tensor_index]; - if (!IsConstantTensor(tensor)) { - buffer_map->SetFromTfLite(tensor_index, tensor); - } - } - - // Execute the TensorFlow Ops sequentially. - for (const auto& node_data : op_data->nodes) { - if (node_data.nodedef.op().empty()) { - context->ReportError(context, "Invalid NodeDef in Flex op '%s'", - node_data.name.c_str()); - return kTfLiteError; - } - auto status = - ExecuteFlexOp(eager_context, buffer_map, node_data.name, - node_data.nodedef, node_data.inputs, node_data.outputs); - TF_LITE_ENSURE_OK(context, ConvertStatus(context, status)); - } - - for (auto tensor_index : op_data->subgraph_outputs) { - if (!buffer_map->HasTensor(tensor_index)) { - context->ReportError(context, "Cannot write to invalid tensor index %d", - tensor_index); - return kTfLiteError; - } - - TfLiteTensor* tensor = &context->tensors[tensor_index]; - TF_LITE_ENSURE_OK( - context, - CopyShapeAndType(context, buffer_map->GetTensor(tensor_index), tensor)); - tensor->buffer_handle = tensor_index; - tensor->data_is_stale = true; - } - - return kTfLiteOk; -} - -} // namespace kernel - -TfLiteRegistration GetKernel() { - TfLiteRegistration registration{&kernel::Init, &kernel::Free, - &kernel::Prepare, &kernel::Eval, - nullptr, kTfLiteBuiltinDelegate}; - return registration; -} - -} // namespace flex -} // namespace tflite diff --git a/tensorflow/contrib/lite/delegates/flex/kernel_test.cc b/tensorflow/contrib/lite/delegates/flex/kernel_test.cc deleted file mode 100644 index 94a6f8b61ad28144f6b8d0d462338ab4176af168..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/delegates/flex/kernel_test.cc +++ /dev/null @@ -1,230 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include "tensorflow/contrib/lite/delegates/flex/kernel.h" - -#include -#include -#include "tensorflow/contrib/lite/delegates/flex/delegate_data.h" -#include "tensorflow/contrib/lite/delegates/flex/test_util.h" - -namespace tflite { -namespace flex { -namespace { - -using ::testing::ContainsRegex; -using ::testing::ElementsAre; - -TfLiteStatus GenericPrepare(TfLiteContext* context, TfLiteDelegate* delegate, - const std::vector& supported_nodes) { - TfLiteIntArray* size_and_nodes = - ConvertVectorToTfLiteIntArray(supported_nodes); - TF_LITE_ENSURE_STATUS(context->ReplaceSubgraphsWithDelegateKernels( - context, flex::GetKernel(), size_and_nodes, delegate)); - TfLiteIntArrayFree(size_and_nodes); - return kTfLiteOk; -} - -class KernelTest : public testing::FlexModelTest { - public: - KernelTest() { - CHECK(DelegateData::Create(&delegate_data_).ok()); - interpreter_.reset(new Interpreter(&error_reporter_)); - } - - ~KernelTest() override { - // The data needs to be released before the interpreter because the - // interpreter references the data. - delegate_data_.reset(); - interpreter_.reset(); - } - - template - void ConfigureDelegate(T prepare_function) { - delegate_.data_ = delegate_data_.get(); - delegate_.FreeBufferHandle = nullptr; - delegate_.Prepare = prepare_function; - delegate_.CopyFromBufferHandle = [](TfLiteContext* context, - TfLiteDelegate* delegate, - TfLiteBufferHandle buffer_handle, - void* data, size_t size) { - auto* delegate_data = reinterpret_cast(delegate->data_); - tensorflow::StringPiece values = delegate_data->GetBufferMap(context) - ->GetTensor(buffer_handle) - .tensor_data(); - memcpy(data, values.data(), values.size()); - return kTfLiteOk; - }; - CHECK(interpreter_->ModifyGraphWithDelegate( - &delegate_, /*allow_dynamic_tensors=*/true) == kTfLiteOk); - } - - private: - std::unique_ptr delegate_data_; - TfLiteDelegate delegate_; -}; - -TEST_F(KernelTest, FullGraph) { - // Define the graph. - AddTensors(9, {0, 3}, {8}, kTfLiteFloat32, {3}); - - AddTfOp(testing::kUnpack, {0}, {1, 2}); - AddTfOp(testing::kUnpack, {3}, {4, 5}); - AddTfOp(testing::kAdd, {1, 4}, {6}); - AddTfOp(testing::kAdd, {2, 5}, {7}); - AddTfOp(testing::kMul, {6, 7}, {8}); - - // Apply Delegate. - ConfigureDelegate([](TfLiteContext* context, TfLiteDelegate* delegate) { - return GenericPrepare(context, delegate, {0, 1, 2, 3, 4}); - }); - - // Define inputs. - SetShape(0, {2, 2, 1}); - SetValues(0, {1.1f, 2.2f, 3.3f, 4.4f}); - SetShape(3, {2, 2, 1}); - SetValues(3, {1.1f, 2.2f, 3.3f, 4.4f}); - - ASSERT_TRUE(Invoke()); - - ASSERT_THAT(GetShape(8), ElementsAre(2, 1)); - ASSERT_THAT(GetValues(8), ElementsAre(14.52f, 38.72f)); -} - -TEST_F(KernelTest, BadTensorFlowOp) { - AddTensors(2, {0}, {1}, kTfLiteFloat32, {3}); - AddTfOp(testing::kNonExistent, {0}, {1}); - - ConfigureDelegate([](TfLiteContext* context, TfLiteDelegate* delegate) { - return GenericPrepare(context, delegate, {0}); - }); - - SetShape(0, {2, 2, 1}); - SetValues(0, {1.1f, 2.2f, 3.3f, 4.4f}); - - ASSERT_FALSE(Invoke()); - ASSERT_THAT(error_reporter().error_messages(), - ContainsRegex("while processing attributes of 'NonExistentOp'")); -} - -TEST_F(KernelTest, BadNumberOfOutputs) { - AddTensors(3, {0}, {1, 2}, kTfLiteFloat32, {3}); - AddTfOp(testing::kIdentity, {0}, {1, 2}); - - ConfigureDelegate([](TfLiteContext* context, TfLiteDelegate* delegate) { - return GenericPrepare(context, delegate, {0}); - }); - - SetShape(0, {2, 2, 1}); - SetValues(0, {1.1f, 2.2f, 3.3f, 4.4f}); - - ASSERT_FALSE(Invoke()); - ASSERT_THAT(error_reporter().error_messages(), - ContainsRegex("Unexpected number of outputs")); -} - -TEST_F(KernelTest, IncompatibleNodeDef) { - AddTensors(2, {0}, {1}, kTfLiteFloat32, {3}); - - // Cast is a TF op, but we don't add the proper nodedef to it in AddTfOp. - AddTfOp(testing::kIncompatibleNodeDef, {0}, {1}); - - ConfigureDelegate([](TfLiteContext* context, TfLiteDelegate* delegate) { - return GenericPrepare(context, delegate, {0}); - }); - - SetShape(0, {2, 2, 1}); - SetValues(0, {1.1f, 2.2f, 3.3f, 4.4f}); - - ASSERT_FALSE(Invoke()); - ASSERT_THAT(error_reporter().error_messages(), - ContainsRegex("while executing 'Cast' via Eager")); -} - -TEST_F(KernelTest, WrongSetOfNodes) { - AddTensors(4, {0}, {3}, kTfLiteFloat32, {3}); - AddTfOp(testing::kUnpack, {0}, {1, 2}); - AddTfLiteMulOp({1, 2}, {3}); - - // Specify that testing::kMul (#1) is supported when it actually isn't. - ConfigureDelegate([](TfLiteContext* context, TfLiteDelegate* delegate) { - return GenericPrepare(context, delegate, {0, 1}); - }); - - SetShape(0, {2, 2, 1}); - SetValues(0, {1.1f, 2.2f, 3.3f, 4.4f}); - - ASSERT_FALSE(Invoke()); - ASSERT_THAT(error_reporter().error_messages(), - ContainsRegex("Invalid NodeDef in Flex op")); -} - -TEST_F(KernelTest, MixedGraph) { - AddTensors(9, {0, 3}, {8}, kTfLiteFloat32, {3}); - - AddTfOp(testing::kUnpack, {0}, {1, 2}); - AddTfOp(testing::kUnpack, {3}, {4, 5}); - AddTfOp(testing::kAdd, {1, 4}, {6}); - AddTfOp(testing::kAdd, {2, 5}, {7}); - AddTfLiteMulOp({6, 7}, {8}); - - ConfigureDelegate([](TfLiteContext* context, TfLiteDelegate* delegate) { - return GenericPrepare(context, delegate, {0, 1, 2, 3}); - }); - - SetShape(0, {2, 2, 1}); - SetValues(0, {1.1f, 2.2f, 3.3f, 4.4f}); - SetShape(3, {2, 2, 1}); - SetValues(3, {1.1f, 2.2f, 3.3f, 4.4f}); - - ASSERT_TRUE(Invoke()); - - ASSERT_THAT(GetShape(8), ElementsAre(2, 1)); - ASSERT_THAT(GetValues(8), ElementsAre(14.52f, 38.72f)); -} - -TEST_F(KernelTest, SplitGraph) { - AddTensors(10, {0}, {9}, kTfLiteFloat32, {3}); - - AddTfOp(testing::kUnpack, {0}, {1, 2}); - AddTfOp(testing::kAdd, {1, 2}, {3}); - AddTfOp(testing::kUnpack, {3}, {4, 5}); - - AddTfLiteMulOp({4, 5}, {6}); - - AddTfOp(testing::kUnpack, {6}, {7, 8}); - AddTfOp(testing::kAdd, {7, 8}, {9}); - - ConfigureDelegate([](TfLiteContext* context, TfLiteDelegate* delegate) { - return GenericPrepare(context, delegate, {0, 1, 2, 4, 5}); - }); - - SetShape(0, {2, 2, 2, 1}); - SetValues(0, {3.0f, 1.0f, 0.5f, -1.0f, 0.0f, 1.0f, 1.5f, 3.0f}); - - ASSERT_TRUE(Invoke()); - - ASSERT_THAT(GetShape(9), ElementsAre(1)); - ASSERT_THAT(GetValues(9), ElementsAre(10.0f)); -} - -} // namespace -} // namespace flex -} // namespace tflite - -int main(int argc, char** argv) { - ::tflite::LogToStderr(); - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/contrib/lite/delegates/flex/test_util.cc b/tensorflow/contrib/lite/delegates/flex/test_util.cc deleted file mode 100644 index 69c336a01a57416bb331a897faba03ad75a38f95..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/delegates/flex/test_util.cc +++ /dev/null @@ -1,157 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/contrib/lite/delegates/flex/test_util.h" - -#include "absl/memory/memory.h" -#include "flatbuffers/flexbuffers.h" // TF:flatbuffers -#include "tensorflow/contrib/lite/string.h" - -namespace tflite { -namespace flex { -namespace testing { - -bool FlexModelTest::Invoke() { return interpreter_->Invoke() == kTfLiteOk; } - -void FlexModelTest::SetShape(int tensor_index, const std::vector& values) { - ASSERT_EQ(interpreter_->ResizeInputTensor(tensor_index, values), kTfLiteOk); - ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk); -} - -std::vector FlexModelTest::GetShape(int tensor_index) { - std::vector result; - auto* dims = interpreter_->tensor(tensor_index)->dims; - result.reserve(dims->size); - for (int i = 0; i < dims->size; ++i) { - result.push_back(dims->data[i]); - } - return result; -} - -TfLiteType FlexModelTest::GetType(int tensor_index) { - return interpreter_->tensor(tensor_index)->type; -} - -void FlexModelTest::AddTensors(int num_tensors, const std::vector& inputs, - const std::vector& outputs, TfLiteType type, - const std::vector& dims) { - interpreter_->AddTensors(num_tensors); - for (int i = 0; i < num_tensors; ++i) { - TfLiteQuantizationParams quant; - // Suppress explicit output type specification to ensure type inference - // works properly. - if (std::find(outputs.begin(), outputs.end(), i) != outputs.end()) { - type = kTfLiteFloat32; - } - CHECK_EQ(interpreter_->SetTensorParametersReadWrite(i, type, - /*name=*/"", - /*dims=*/dims, quant), - kTfLiteOk); - } - - CHECK_EQ(interpreter_->SetInputs(inputs), kTfLiteOk); - CHECK_EQ(interpreter_->SetOutputs(outputs), kTfLiteOk); -} - -void FlexModelTest::AddTfLiteMulOp(const std::vector& inputs, - const std::vector& outputs) { - static TfLiteRegistration reg = {nullptr, nullptr, nullptr, nullptr}; - reg.builtin_code = BuiltinOperator_MUL; - reg.prepare = [](TfLiteContext* context, TfLiteNode* node) { - auto* i0 = &context->tensors[node->inputs->data[0]]; - auto* o = &context->tensors[node->outputs->data[0]]; - return context->ResizeTensor(context, o, TfLiteIntArrayCopy(i0->dims)); - }; - reg.invoke = [](TfLiteContext* context, TfLiteNode* node) { - auto* i0 = &context->tensors[node->inputs->data[0]]; - auto* i1 = &context->tensors[node->inputs->data[1]]; - auto* o = &context->tensors[node->outputs->data[0]]; - for (int i = 0; i < o->bytes / sizeof(float); ++i) { - o->data.f[i] = i0->data.f[i] * i1->data.f[i]; - } - return kTfLiteOk; - }; - - CHECK_EQ(interpreter_->AddNodeWithParameters(inputs, outputs, nullptr, 0, - nullptr, ®), - kTfLiteOk); -} - -void FlexModelTest::AddTfOp(TfOpType op, const std::vector& inputs, - const std::vector& outputs) { - auto attr = [](const string& key, const string& value) { - return " attr{ key: '" + key + "' value {" + value + "}}"; - }; - - // Crude type attribution, will need fleshing out as more tests are added. - // TODO(b/113613439): Use nodedef string utilities to properly handle - // all types. - string type_attribute = attr("T", "type: DT_FLOAT"); - if (interpreter_->tensor(inputs[0])->type == kTfLiteInt32) { - type_attribute = attr("T", "type: DT_INT32"); - } - - if (op == kUnpack) { - string attributes = - type_attribute + attr("num", "i: 2") + attr("axis", "i: 0"); - AddTfOp("FlexUnpack", "Unpack", attributes, inputs, outputs); - } else if (op == kIdentity) { - string attributes = type_attribute; - AddTfOp("FlexIdentity", "Identity", attributes, inputs, outputs); - } else if (op == kAdd) { - string attributes = type_attribute; - AddTfOp("FlexAdd", "Add", attributes, inputs, outputs); - } else if (op == kMul) { - string attributes = type_attribute; - AddTfOp("FlexMul", "Mul", attributes, inputs, outputs); - } else if (op == kNonExistent) { - AddTfOp("NonExistentOp", "NonExistentOp", "", inputs, outputs); - } else if (op == kIncompatibleNodeDef) { - // "Cast" op is created without attributes - making it incompatible. - AddTfOp("FlexCast", "Cast", "", inputs, outputs); - } -} - -void FlexModelTest::AddTfOp(const char* tflite_name, const string& tf_name, - const string& nodedef_str, - const std::vector& inputs, - const std::vector& outputs) { - static TfLiteRegistration reg = {nullptr, nullptr, nullptr, nullptr}; - reg.builtin_code = BuiltinOperator_CUSTOM; - reg.custom_name = tflite_name; - - tensorflow::NodeDef nodedef; - CHECK(tensorflow::protobuf::TextFormat::ParseFromString( - nodedef_str + " op: '" + tf_name + "'", &nodedef)); - string serialized_nodedef; - CHECK(nodedef.SerializeToString(&serialized_nodedef)); - flexbuffers::Builder fbb; - fbb.Vector([&]() { - fbb.String(nodedef.op()); - fbb.String(serialized_nodedef); - }); - fbb.Finish(); - - flexbuffers_.push_back(fbb.GetBuffer()); - auto& buffer = flexbuffers_.back(); - CHECK_EQ(interpreter_->AddNodeWithParameters( - inputs, outputs, reinterpret_cast(buffer.data()), - buffer.size(), nullptr, ®), - kTfLiteOk); -} - -} // namespace testing -} // namespace flex -} // namespace tflite diff --git a/tensorflow/contrib/lite/delegates/flex/test_util.h b/tensorflow/contrib/lite/delegates/flex/test_util.h deleted file mode 100644 index a8c81b90a3b8dc49ae058adb172456fe4d6e7172..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/delegates/flex/test_util.h +++ /dev/null @@ -1,119 +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_LITE_DELEGATES_FLEX_TEST_UTIL_H_ -#define TENSORFLOW_CONTRIB_LITE_DELEGATES_FLEX_TEST_UTIL_H_ - -#include "tensorflow/c/c_api_internal.h" -#include "tensorflow/contrib/lite/kernels/test_util.h" - -namespace tflite { -namespace flex { -namespace testing { - -enum TfOpType { - kUnpack, - kIdentity, - kAdd, - kMul, - // Represents an op that does not exist in TensorFlow. - kNonExistent, - // Represents an valid TensorFlow op where the NodeDef is incompatible. - kIncompatibleNodeDef, -}; - -// This class creates models with TF and TFLite ops. In order to use this class -// to test the Flex delegate, implement a function that calls -// interpreter->ModifyGraphWithDelegate. -class FlexModelTest : public ::testing::Test { - public: - FlexModelTest() {} - ~FlexModelTest() {} - - bool Invoke(); - - // Sets the (typed) tensor's values at the given index. - template - void SetTypedValues(int tensor_index, const std::vector& values) { - memcpy(interpreter_->typed_tensor(tensor_index), values.data(), - values.size() * sizeof(T)); - } - - // Returns the (typed) tensor's values at the given index. - template - std::vector GetTypedValues(int tensor_index) { - const TfLiteTensor* t = interpreter_->tensor(tensor_index); - const T* tdata = interpreter_->typed_tensor(tensor_index); - return std::vector(tdata, tdata + t->bytes / sizeof(T)); - } - - // Sets the tensor's values at the given index. - void SetValues(int tensor_index, const std::vector& values) { - SetTypedValues(tensor_index, values); - } - - // Returns the tensor's values at the given index. - std::vector GetValues(int tensor_index) { - return GetTypedValues(tensor_index); - } - - // Sets the tensor's shape at the given index. - void SetShape(int tensor_index, const std::vector& values); - - // Returns the tensor's shape at the given index. - std::vector GetShape(int tensor_index); - - // Returns the tensor's type at the given index. - TfLiteType GetType(int tensor_index); - - const TestErrorReporter& error_reporter() const { return error_reporter_; } - - // Adds `num_tensor` tensors to the model. `inputs` contains the indices of - // the input tensors and `outputs` contains the indices of the output - // tensors. All tensors are set to have `type` and `dims`. - void AddTensors(int num_tensors, const std::vector& inputs, - const std::vector& outputs, TfLiteType type, - const std::vector& dims); - - // Adds a TFLite Mul op. `inputs` contains the indices of the input tensors - // and `outputs` contains the indices of the output tensors. - void AddTfLiteMulOp(const std::vector& inputs, - const std::vector& outputs); - - // Adds a TensorFlow op. `inputs` contains the indices of the - // input tensors and `outputs` contains the indices of the output tensors. - // This function is limited to the set of ops defined in TfOpType. - void AddTfOp(TfOpType op, const std::vector& inputs, - const std::vector& outputs); - - protected: - std::unique_ptr interpreter_; - TestErrorReporter error_reporter_; - - private: - // Helper method to add a TensorFlow op. tflite_names needs to start with - // "Flex" in order to work with the Flex delegate. - void AddTfOp(const char* tflite_name, const string& tf_name, - const string& nodedef_str, const std::vector& inputs, - const std::vector& outputs); - - std::vector> flexbuffers_; -}; - -} // namespace testing -} // namespace flex -} // namespace tflite - -#endif // TENSORFLOW_CONTRIB_LITE_DELEGATES_FLEX_TEST_UTIL_H_ diff --git a/tensorflow/contrib/lite/delegates/flex/util.cc b/tensorflow/contrib/lite/delegates/flex/util.cc deleted file mode 100644 index 829bc388bf4f613e82600edfc7363d0774d49878..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/delegates/flex/util.cc +++ /dev/null @@ -1,104 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include "tensorflow/contrib/lite/delegates/flex/util.h" - -namespace tflite { -namespace flex { - -TfLiteStatus ConvertStatus(TfLiteContext* context, - const tensorflow::Status& status) { - if (!status.ok()) { - context->ReportError(context, "%s", status.error_message().c_str()); - return kTfLiteError; - } - return kTfLiteOk; -} - -TfLiteStatus CopyShapeAndType(TfLiteContext* context, - const tensorflow::Tensor& src, - TfLiteTensor* tensor) { - tensor->type = GetTensorFlowLiteType(static_cast(src.dtype())); - if (tensor->type == kTfLiteNoType) { - context->ReportError(context, - "TF Lite does not support TensorFlow data type: %s", - DataTypeString(src.dtype()).c_str()); - return kTfLiteError; - } - - int num_dims = src.dims(); - TfLiteIntArray* shape = TfLiteIntArrayCreate(num_dims); - for (int j = 0; j < num_dims; ++j) { - // We need to cast from TensorFlow's int64 to TF Lite's int32. Let's - // make sure there's no overflow. - if (src.dim_size(j) >= std::numeric_limits::max()) { - context->ReportError(context, - "Dimension value in TensorFlow shape is larger than " - "supported by TF Lite"); - TfLiteIntArrayFree(shape); - return kTfLiteError; - } - shape->data[j] = static_cast(src.dim_size(j)); - } - return context->ResizeTensor(context, tensor, shape); -} - -TF_DataType GetTensorFlowDataType(TfLiteType type) { - switch (type) { - case kTfLiteNoType: - return TF_FLOAT; - case kTfLiteFloat32: - return TF_FLOAT; - case kTfLiteInt16: - return TF_INT16; - case kTfLiteInt32: - return TF_INT32; - case kTfLiteUInt8: - return TF_UINT8; - case kTfLiteInt64: - return TF_INT64; - case kTfLiteComplex64: - return TF_COMPLEX64; - case kTfLiteString: - return TF_STRING; - case kTfLiteBool: - return TF_BOOL; - } -} - -TfLiteType GetTensorFlowLiteType(TF_DataType type) { - switch (type) { - case TF_FLOAT: - return kTfLiteFloat32; - case TF_INT16: - return kTfLiteInt16; - case TF_INT32: - return kTfLiteInt32; - case TF_UINT8: - return kTfLiteUInt8; - case TF_INT64: - return kTfLiteInt64; - case TF_COMPLEX64: - return kTfLiteComplex64; - case TF_STRING: - return kTfLiteString; - case TF_BOOL: - return kTfLiteBool; - default: - return kTfLiteNoType; - } -} - -} // namespace flex -} // namespace tflite diff --git a/tensorflow/contrib/lite/delegates/flex/util.h b/tensorflow/contrib/lite/delegates/flex/util.h deleted file mode 100644 index 7f910e7316e67363a6e54389f1d0cc94b3e009a0..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/delegates/flex/util.h +++ /dev/null @@ -1,47 +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_LITE_DELEGATES_FLEX_UTIL_H_ -#define TENSORFLOW_CONTRIB_LITE_DELEGATES_FLEX_UTIL_H_ - -#include "tensorflow/c/c_api_internal.h" -#include "tensorflow/contrib/lite/c/c_api_internal.h" -#include "tensorflow/core/framework/tensor.h" -#include "tensorflow/core/lib/core/status.h" - -namespace tflite { -namespace flex { - -// Converts a tensorflow:Status into a TfLiteStatus. If the original status -// represented an error, reports it using the given 'context'. -TfLiteStatus ConvertStatus(TfLiteContext* context, - const tensorflow::Status& status); - -// Copies the given shape and type of the TensorFlow 'src' tensor into a TF Lite -// 'tensor'. Logs an error and returns kTfLiteError if the shape or type can't -// be converted. -TfLiteStatus CopyShapeAndType(TfLiteContext* context, - const tensorflow::Tensor& src, - TfLiteTensor* tensor); - -// Returns the TF C API Data type that corresponds to the given TfLiteType. -TF_DataType GetTensorFlowDataType(TfLiteType type); - -// Returns the TfLiteType that corresponds to the given TF C API Data type. -TfLiteType GetTensorFlowLiteType(TF_DataType); - -} // namespace flex -} // namespace tflite - -#endif // TENSORFLOW_CONTRIB_LITE_DELEGATES_FLEX_UTIL_H_ diff --git a/tensorflow/contrib/lite/delegates/flex/util_test.cc b/tensorflow/contrib/lite/delegates/flex/util_test.cc deleted file mode 100644 index 5f049e7b0a0c1f7be28d33b532157c6f9211c7c1..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/delegates/flex/util_test.cc +++ /dev/null @@ -1,142 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include "tensorflow/contrib/lite/delegates/flex/util.h" - -#include - -#include -#include -#include "tensorflow/contrib/lite/string.h" -#include "tensorflow/contrib/lite/testing/util.h" - -namespace tflite { -namespace flex { -namespace { - -using tensorflow::DT_FLOAT; -using tensorflow::DT_INT32; -using tensorflow::Tensor; -using ::testing::ElementsAre; - -struct TestContext : public TfLiteContext { - string error; - std::vector new_size; -}; - -void ReportError(TfLiteContext* context, const char* format, ...) { - TestContext* c = static_cast(context); - const size_t kBufferSize = 1024; - char temp_buffer[kBufferSize]; - - va_list args; - va_start(args, format); - vsnprintf(temp_buffer, kBufferSize, format, args); - va_end(args); - - c->error = temp_buffer; -} - -TfLiteStatus ResizeTensor(TfLiteContext* context, TfLiteTensor* tensor, - TfLiteIntArray* new_size) { - TestContext* c = static_cast(context); - c->new_size.clear(); - for (int i = 0; i < new_size->size; ++i) { - c->new_size.push_back(new_size->data[i]); - } - TfLiteIntArrayFree(new_size); - return kTfLiteOk; -} - -TEST(UtilTest, ConvertStatus) { - TestContext context; - context.ReportError = ReportError; - - EXPECT_EQ(ConvertStatus(&context, tensorflow::errors::Internal("Some Error")), - kTfLiteError); - EXPECT_EQ(context.error, "Some Error"); - - context.error.clear(); - EXPECT_EQ(ConvertStatus(&context, tensorflow::Status()), kTfLiteOk); - EXPECT_TRUE(context.error.empty()); -} - -TEST(UtilTest, CopyShapeAndType) { - TestContext context; - context.ReportError = ReportError; - context.ResizeTensor = ResizeTensor; - - TfLiteTensor dst; - - EXPECT_EQ(CopyShapeAndType(&context, Tensor(), &dst), kTfLiteOk); - EXPECT_THAT(context.new_size, ElementsAre(0)); - EXPECT_EQ(dst.type, kTfLiteFloat32); - - EXPECT_EQ(CopyShapeAndType(&context, Tensor(DT_FLOAT, {1, 2}), &dst), - kTfLiteOk); - EXPECT_THAT(context.new_size, ElementsAre(1, 2)); - EXPECT_EQ(dst.type, kTfLiteFloat32); - - EXPECT_EQ(CopyShapeAndType(&context, Tensor(DT_INT32, {1, 2}), &dst), - kTfLiteOk); - EXPECT_THAT(context.new_size, ElementsAre(1, 2)); - EXPECT_EQ(dst.type, kTfLiteInt32); - - EXPECT_EQ(CopyShapeAndType(&context, Tensor(DT_FLOAT, {1LL << 44, 2}), &dst), - kTfLiteError); - EXPECT_EQ(context.error, - "Dimension value in TensorFlow shape is larger than supported by " - "TF Lite"); - - EXPECT_EQ( - CopyShapeAndType(&context, Tensor(tensorflow::DT_HALF, {1, 2}), &dst), - kTfLiteError); - EXPECT_EQ(context.error, - "TF Lite does not support TensorFlow data type: half"); -} - -TEST(UtilTest, TypeConversionsFromTFLite) { - EXPECT_EQ(TF_FLOAT, GetTensorFlowDataType(kTfLiteNoType)); - EXPECT_EQ(TF_FLOAT, GetTensorFlowDataType(kTfLiteFloat32)); - EXPECT_EQ(TF_INT16, GetTensorFlowDataType(kTfLiteInt16)); - EXPECT_EQ(TF_INT32, GetTensorFlowDataType(kTfLiteInt32)); - EXPECT_EQ(TF_UINT8, GetTensorFlowDataType(kTfLiteUInt8)); - EXPECT_EQ(TF_INT64, GetTensorFlowDataType(kTfLiteInt64)); - EXPECT_EQ(TF_COMPLEX64, GetTensorFlowDataType(kTfLiteComplex64)); - EXPECT_EQ(TF_STRING, GetTensorFlowDataType(kTfLiteString)); - EXPECT_EQ(TF_BOOL, GetTensorFlowDataType(kTfLiteBool)); -} - -TEST(UtilTest, TypeConversionsFromTensorFlow) { - EXPECT_EQ(kTfLiteFloat32, GetTensorFlowLiteType(TF_FLOAT)); - EXPECT_EQ(kTfLiteInt16, GetTensorFlowLiteType(TF_INT16)); - EXPECT_EQ(kTfLiteInt32, GetTensorFlowLiteType(TF_INT32)); - EXPECT_EQ(kTfLiteUInt8, GetTensorFlowLiteType(TF_UINT8)); - EXPECT_EQ(kTfLiteInt64, GetTensorFlowLiteType(TF_INT64)); - EXPECT_EQ(kTfLiteComplex64, GetTensorFlowLiteType(TF_COMPLEX64)); - EXPECT_EQ(kTfLiteString, GetTensorFlowLiteType(TF_STRING)); - EXPECT_EQ(kTfLiteBool, GetTensorFlowLiteType(TF_BOOL)); - EXPECT_EQ(kTfLiteNoType, GetTensorFlowLiteType(TF_RESOURCE)); - EXPECT_EQ(kTfLiteNoType, GetTensorFlowLiteType(TF_VARIANT)); -} - -} // namespace -} // namespace flex -} // namespace tflite - -int main(int argc, char** argv) { - ::tflite::LogToStderr(); - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/contrib/lite/delegates/nnapi/BUILD b/tensorflow/contrib/lite/delegates/nnapi/BUILD deleted file mode 100644 index 4e7b2948fb920c3aaf9a6f4a9cdff7c476911e7a..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/delegates/nnapi/BUILD +++ /dev/null @@ -1,37 +0,0 @@ -package(default_visibility = [ - "//visibility:public", -]) - -load("//tensorflow:tensorflow.bzl", "tf_cc_test") - -licenses(["notice"]) # Apache 2.0 - -cc_library( - name = "nnapi_delegate", - srcs = ["nnapi_delegate.cc"], - hdrs = ["nnapi_delegate.h"], - deps = [ - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite:kernel_api", - "//tensorflow/contrib/lite/c:c_api_internal", - "//tensorflow/contrib/lite/kernels:kernel_util", - "//tensorflow/contrib/lite/nnapi:nnapi_lib", - ], -) - -tf_cc_test( - name = "nnapi_delegate_test", - size = "small", - srcs = ["nnapi_delegate_test.cc"], - tags = [ - "no_oss", - "noasan", # TODO(b/112326936): re-enable for asan once fixed. - ], - deps = [ - ":nnapi_delegate", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/c:c_api_internal", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) diff --git a/tensorflow/contrib/lite/delegates/nnapi/nnapi_delegate.cc b/tensorflow/contrib/lite/delegates/nnapi/nnapi_delegate.cc deleted file mode 100644 index d85e576284fac87519d7f4bb4bd76fe2619b59d5..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/delegates/nnapi/nnapi_delegate.cc +++ /dev/null @@ -1,1220 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include -#include -#include -#include - -#include "tensorflow/contrib/lite/allocation.h" -#include "tensorflow/contrib/lite/builtin_op_data.h" -#include "tensorflow/contrib/lite/builtin_ops.h" -#include "tensorflow/contrib/lite/c/c_api_internal.h" -#include "tensorflow/contrib/lite/context_util.h" -#include "tensorflow/contrib/lite/delegates/nnapi/nnapi_delegate.h" -#include "tensorflow/contrib/lite/kernels/kernel_util.h" -#include "tensorflow/contrib/lite/nnapi/NeuralNetworksShim.h" - -#ifdef __ANDROID__ -#include -#include -#include -#endif - -namespace tflite { -namespace { - -// TODO(b/80621585): Consider printing error string, but don't for now to -// minimize binary size. -#define CHECK_NN(context, code) \ - if (code != ANEURALNETWORKS_NO_ERROR) { \ - context->ReportError(context, "NN API returned error (%d).\n", code); \ - return kTfLiteError; \ - } - -namespace { -int32_t GetAndroidSdkVersion() { -#ifdef __ANDROID__ - const char* sdkProp = "ro.build.version.sdk"; - char sdkVersion[PROP_VALUE_MAX]; - int length = __system_property_get(sdkProp, sdkVersion); - if (length != 0) { - for (int i = 0; i < length; ++i) { - int digit = sdkVersion[i] - '0'; - if (digit < 0 || digit > 9) { - // Non-numeric SDK version, assume it's higher then expected; - return std::numeric_limits::max(); - } - } - return atoi(sdkVersion); - } -#endif // __ANDROID__ - return 0; -} - -constexpr int32_t kMinSdkVersionForNNAPI = 27; -constexpr int32_t kMinSdkVersionForNNAPI11 = 28; -static const int32_t kAndroidSdkVersion = GetAndroidSdkVersion(); - -} // namespace - -// RAII NN API Model Destructor for use with std::unique_ptr -struct NNFreeModel { - void operator()(ANeuralNetworksModel* model) { - ANeuralNetworksModel_free(model); - } -}; -// RAII NN API Compilation Destructor for use with std::unique_ptr -struct NNFreeCompilation { - void operator()(ANeuralNetworksCompilation* model) { - ANeuralNetworksCompilation_free(model); - } -}; - -// Manage NNAPI shared memory handle -class NNMemory { - public: - NNMemory(const char* name, size_t size) { -#ifdef __ANDROID__ - byte_size_ = size; - fd_ = ASharedMemory_create(name, size); - data_ptr_ = reinterpret_cast( - mmap(nullptr, size, PROT_READ | PROT_WRITE, MAP_SHARED, fd_, 0)); - ANeuralNetworksMemory_createFromFd(size, PROT_READ | PROT_WRITE, fd_, 0, - &nn_memory_handle_); -#endif - } - - ~NNMemory() { -#ifdef __ANDROID__ - if (data_ptr_) { - munmap(data_ptr_, byte_size_); - } - if (nn_memory_handle_) { - ANeuralNetworksMemory_free(nn_memory_handle_); - } - if (fd_ > 0) close(fd_); -#endif - } - - ANeuralNetworksMemory* get_handle() { return nn_memory_handle_; } - uint8_t* get_data_ptr() { return data_ptr_; } - - private: -#ifdef __ANDROID__ - int fd_ = 0; - size_t byte_size_ = 0; -#endif - uint8_t* data_ptr_ = nullptr; - ANeuralNetworksMemory* nn_memory_handle_ = nullptr; -}; // namespace - -// Track tensor indices to NN API tensor indices mapping. -class OperandMapping { - public: - // Given a TFLite index return the ANN index. If it doesn't exist - // return -1. - int lite_index_to_ann(int index) const { - if (index < lite_tensor_to_ann_tensor_.size()) - return lite_tensor_to_ann_tensor_[index]; - else - return -1; - } - - // NN API uses non tensor operands instead of structs. This creates one - // and returns the index. It uses a std::vector and resizes it as needed - // keeping -1 to unmapped values. Intermediate tensors likely will not - // be mapped. - int add_new_non_tensor_operand() { return next_ann_tensor_index_++; } - - // Add a new mapping from `tflite_index` and return the NN API tensor index. - int add_new_ann_tensor_index(int tflite_index) { - if (tflite_index >= lite_tensor_to_ann_tensor_.size()) { - lite_tensor_to_ann_tensor_.resize(tflite_index + 1, -1); - } - int new_tensor_index = next_ann_tensor_index_++; - lite_tensor_to_ann_tensor_[tflite_index] = new_tensor_index; - return new_tensor_index; - } - - private: - // Next index of ann tensor - int next_ann_tensor_index_ = 0; - - // Mapping from lite index. Use a std::vector for speed and code size - // rather than a map. - std::vector lite_tensor_to_ann_tensor_; -}; - -// Abstract builder for building an op in the NN API graph. This handles -// the disparity between TFLite and NN API operand types. NN API has singular -// operands for both tensors and parameters, and TFLite separates the two. -class NNAPIOpBuilder { - public: - NNAPIOpBuilder(TfLiteContext* context, OperandMapping* tensor_mapping, - ANeuralNetworksModel* nn_model) - : context_(context), - operand_mapping_(tensor_mapping), - nn_model_(nn_model) {} - - TfLiteStatus AddScalarInt32Operand(int32_t value) { - return AddScalarOperand(value, ANEURALNETWORKS_INT32); - } - - TfLiteStatus AddScalarFloat32Operand(float value) { - return AddScalarOperand(value, ANEURALNETWORKS_FLOAT32); - } - - TfLiteStatus AddVectorInt32Operand(const int32_t* values, - uint32_t num_values) { - return AddVectorOperand(values, num_values, - ANEURALNETWORKS_TENSOR_INT32); - } - - TfLiteStatus AddVectorFloat32Operand(const float* values, - uint32_t num_values) { - return AddVectorOperand(values, num_values, - ANEURALNETWORKS_TENSOR_FLOAT32); - } - - TfLiteStatus AddPoolingParams(void* data) { - auto builtin = reinterpret_cast(data); - AddScalarInt32Operand(builtin->padding); - AddScalarInt32Operand(builtin->stride_width); - AddScalarInt32Operand(builtin->stride_height); - AddScalarInt32Operand(builtin->filter_width); - AddScalarInt32Operand(builtin->filter_height); - AddScalarInt32Operand(builtin->activation); - return kTfLiteOk; - } - - TfLiteStatus AddTensorInput(int tensor_index) { - int ann_index; - TF_LITE_ENSURE_STATUS(AddTensor(tensor_index, &ann_index)); - augmented_inputs_.push_back(ann_index); - return kTfLiteOk; - } - - TfLiteStatus AddTensorOutput(int tensor_index) { - int ann_index; - TF_LITE_ENSURE_STATUS(AddTensor(tensor_index, &ann_index)); - augmented_outputs_.push_back(ann_index); - return kTfLiteOk; - } - - TfLiteStatus AddAdditionalFloat32OutputTensor(uint32_t dimension_count) { - std::vector dims(dimension_count, 0); - ANeuralNetworksOperandType operand_type{ - .type = ANEURALNETWORKS_TENSOR_FLOAT32, - .dimensionCount = dimension_count, - .dimensions = dims.data()}; - CHECK_NN(context_, - ANeuralNetworksModel_addOperand(nn_model_, &operand_type)); - int ann_operand = operand_mapping_->add_new_non_tensor_operand(); - augmented_outputs_.push_back(ann_operand); - return kTfLiteOk; - } - - TfLiteStatus AddStateFloat32Tensor(int tensor_index, - int* ann_tensor_index_out) { - TfLiteTensor* tensor = &context_->tensors[tensor_index]; - int ann_index = operand_mapping_->add_new_non_tensor_operand(); - - ANeuralNetworksOperandType operand_type{ - ANEURALNETWORKS_TENSOR_FLOAT32, - static_cast(tensor->dims->size), - reinterpret_cast(tensor->dims->data), tensor->params.scale, - tensor->params.zero_point}; - CHECK_NN(context_, - ANeuralNetworksModel_addOperand(nn_model_, &operand_type)); - augmented_outputs_.push_back(ann_index); - - *ann_tensor_index_out = ann_index; - return kTfLiteOk; - } - - // Adds a new NN API tensor that shadows the TF Lite tensor `tensor_index`. - // This returns the NN API tensor index corresponding to the created tensor. - // If another caller previously created a NN API tensor for `tensor_index` - // then the existing one is returned. - TfLiteStatus AddTensor(int tensor_index, int* ann_tensor_index_out) { - int ann_tensor_index = operand_mapping_->lite_index_to_ann(tensor_index); - if (ann_tensor_index != -1) { - *ann_tensor_index_out = ann_tensor_index; - return kTfLiteOk; - } - // Allocate a new tensor index - ann_tensor_index = operand_mapping_->add_new_ann_tensor_index(tensor_index); - - // Parameters needed for new type. - int32_t nn_type = 0; - float scale = 0.0f; - int32_t zeroPoint = 0; - TfLiteTensor* tensor = &context_->tensors[tensor_index]; - switch (tensor->type) { - case kTfLiteNoType: - // Tensors added during initialization of Ops don't have a type yet and - // should not be registered with the NNAPI. - *ann_tensor_index_out = -1; - return kTfLiteOk; - case kTfLiteFloat32: - nn_type = ANEURALNETWORKS_TENSOR_FLOAT32; - break; - case kTfLiteUInt8: - nn_type = ANEURALNETWORKS_TENSOR_QUANT8_ASYMM; - scale = tensor->params.scale; - zeroPoint = tensor->params.zero_point; - if (scale == 0) { - // TENSOR_QUANT8_ASYMM with zero scale is not valid in NNAPI. - scale = 1; - } - break; - case kTfLiteInt32: - nn_type = ANEURALNETWORKS_TENSOR_INT32; - scale = tensor->params.scale; - zeroPoint = tensor->params.zero_point; - break; - default: - context_->ReportError(context_, "Logic error in NN API Delegate.\n"); - return kTfLiteError; - } - - ANeuralNetworksOperandType operand_type{ - nn_type, static_cast(tensor->dims->size), - reinterpret_cast(tensor->dims->data), scale, zeroPoint}; - CHECK_NN(context_, - ANeuralNetworksModel_addOperand(nn_model_, &operand_type)); - - if (tensor->allocation_type == kTfLiteMmapRo) { - // TODO(b/80630405): Use NNAPIAllocation. - CHECK_NN(context_, ANeuralNetworksModel_setOperandValue( - nn_model_, ann_tensor_index, tensor->data.raw, - tensor->bytes)); - } - - *ann_tensor_index_out = ann_tensor_index; - return kTfLiteOk; - } - - // Finish emitting the op (of type `type`) into the NN API. - TfLiteStatus FinalizeAddOperation(ANeuralNetworksOperationType type) { - // Actually add a NN API operation - CHECK_NN(context_, ANeuralNetworksModel_addOperation( - nn_model_, type, - static_cast(augmented_inputs_.size()), - augmented_inputs_.data(), - static_cast(augmented_outputs_.size()), - augmented_outputs_.data())); - augmented_inputs_.clear(); - augmented_outputs_.clear(); - return kTfLiteOk; - } - - private: - template - TfLiteStatus AddScalarOperand(T value, int32_t nn_type) { - ANeuralNetworksOperandType operand_type{.type = nn_type}; - CHECK_NN(context_, - ANeuralNetworksModel_addOperand(nn_model_, &operand_type)); - int ann_operand = operand_mapping_->add_new_non_tensor_operand(); - CHECK_NN(context_, ANeuralNetworksModel_setOperandValue( - nn_model_, ann_operand, &value, sizeof(T))); - augmented_inputs_.push_back(ann_operand); - return kTfLiteOk; - } - - template - TfLiteStatus AddVectorOperand(const T* values, uint32_t num_values, - int32_t nn_type) { - ANeuralNetworksOperandType operand_type{ - .type = nn_type, .dimensionCount = 1, .dimensions = &num_values}; - CHECK_NN(context_, - ANeuralNetworksModel_addOperand(nn_model_, &operand_type)); - int ann_operand = operand_mapping_->add_new_non_tensor_operand(); - CHECK_NN(context_, - ANeuralNetworksModel_setOperandValue( - nn_model_, ann_operand, values, sizeof(T) * num_values)); - augmented_inputs_.push_back(ann_operand); - return kTfLiteOk; - } - - // TfLiteContext for error handling. Must be named context for macros to - // work. - TfLiteContext* context_; - - // Tracks relationship between indices - OperandMapping* operand_mapping_; - - // The model - ANeuralNetworksModel* nn_model_; - - // Inputs and outputs for the current op. These are augmented in the sense - // that NN API uses operands for all arguments, not just tensors, unlike - // TensorFlow lite. - std::vector augmented_inputs_; - std::vector augmented_outputs_; -}; - -struct NNAPIOpMappingArgs { - TfLiteContext* context; - NNAPIOpBuilder* builder; - TfLiteNode* node; - std::vector* model_state_outputs; - std::vector* model_state_tfl_inputs; -}; - -// The kernel that represents the subgraph of TF Lite being run on NN API. -class NNAPIDelegateKernel { - public: - NNAPIDelegateKernel() = default; - - typedef ANeuralNetworksOperationType (*MappingFn)( - const NNAPIOpMappingArgs& mapping_args); - - // Return a function that knows how to translate a node into its operands - // when called. You can use this function to see if a node is supported - // (i.e. that MappingFn is not nullptr). - MappingFn Map(TfLiteContext* context, int builtin_code, int version, - TfLiteNode* node) { - switch (builtin_code) { - case kTfLiteBuiltinAdd: - if (version == 1) { - return [](const NNAPIOpMappingArgs& mapping_args) - -> ANeuralNetworksOperationType { - auto builtin = reinterpret_cast( - mapping_args.node->builtin_data); - mapping_args.builder->AddScalarInt32Operand(builtin->activation); - return ANEURALNETWORKS_ADD; - }; - } else { - return nullptr; - } - break; - case kTfLiteBuiltinMul: - if (version == 1) { - return [](const NNAPIOpMappingArgs& mapping_args) - -> ANeuralNetworksOperationType { - auto builtin = reinterpret_cast( - mapping_args.node->builtin_data); - mapping_args.builder->AddScalarInt32Operand(builtin->activation); - return ANEURALNETWORKS_MUL; - }; - } else { - return nullptr; - } - break; - case kTfLiteBuiltinAveragePool2d: - if (version == 1) { - return [](const NNAPIOpMappingArgs& mapping_args) - -> ANeuralNetworksOperationType { - mapping_args.builder->AddPoolingParams( - mapping_args.node->builtin_data); - return ANEURALNETWORKS_AVERAGE_POOL_2D; - }; - } else { - return nullptr; - } - break; - case kTfLiteBuiltinMaxPool2d: - if (version == 1) { - return [](const NNAPIOpMappingArgs& mapping_args) - -> ANeuralNetworksOperationType { - mapping_args.builder->AddPoolingParams( - mapping_args.node->builtin_data); - return ANEURALNETWORKS_MAX_POOL_2D; - }; - } else { - return nullptr; - } - break; - case kTfLiteBuiltinL2Pool2d: - if (version == 1) { - return [](const NNAPIOpMappingArgs& mapping_args) - -> ANeuralNetworksOperationType { - mapping_args.builder->AddPoolingParams( - mapping_args.node->builtin_data); - return ANEURALNETWORKS_L2_POOL_2D; - }; - } else { - return nullptr; - } - break; - case kTfLiteBuiltinConv2d: - if (version == 1) { - auto builtin = - reinterpret_cast(node->builtin_data); - if (builtin->dilation_width_factor != 1 || - builtin->dilation_height_factor != 1 || node->inputs->size != 3) { - // NNAPI does not support dilated Conv2D. - return nullptr; - } - return [](const NNAPIOpMappingArgs& mapping_args) - -> ANeuralNetworksOperationType { - auto builtin = reinterpret_cast( - mapping_args.node->builtin_data); - mapping_args.builder->AddScalarInt32Operand(builtin->padding); - mapping_args.builder->AddScalarInt32Operand(builtin->stride_width); - mapping_args.builder->AddScalarInt32Operand(builtin->stride_height); - mapping_args.builder->AddScalarInt32Operand(builtin->activation); - return ANEURALNETWORKS_CONV_2D; - }; - } else { - return nullptr; - } - break; - case kTfLiteBuiltinDepthwiseConv2d: - if (version == 1) { - return [](const NNAPIOpMappingArgs& mapping_args) - -> ANeuralNetworksOperationType { - auto builtin = reinterpret_cast( - mapping_args.node->builtin_data); - mapping_args.builder->AddScalarInt32Operand(builtin->padding); - mapping_args.builder->AddScalarInt32Operand(builtin->stride_width); - mapping_args.builder->AddScalarInt32Operand(builtin->stride_height); - mapping_args.builder->AddScalarInt32Operand( - builtin->depth_multiplier); - mapping_args.builder->AddScalarInt32Operand(builtin->activation); - return ANEURALNETWORKS_DEPTHWISE_CONV_2D; - }; - } else { - return nullptr; - } - break; - case kTfLiteBuiltinFullyConnected: - if (version == 1) { - return [](const NNAPIOpMappingArgs& mapping_args) - -> ANeuralNetworksOperationType { - auto builtin = reinterpret_cast( - mapping_args.node->builtin_data); - mapping_args.builder->AddScalarInt32Operand(builtin->activation); - return ANEURALNETWORKS_FULLY_CONNECTED; - }; - } else { - return nullptr; - } - break; - case kTfLiteBuiltinSoftmax: - if (version == 1) { - return [](const NNAPIOpMappingArgs& mapping_args) - -> ANeuralNetworksOperationType { - auto builtin = reinterpret_cast( - mapping_args.node->builtin_data); - mapping_args.builder->AddScalarFloat32Operand(builtin->beta); - return ANEURALNETWORKS_SOFTMAX; - }; - } else { - return nullptr; - } - break; - case kTfLiteBuiltinReshape: - if (version == 1 && node->inputs->size == 2) { - return [](const NNAPIOpMappingArgs& mapping_args) - -> ANeuralNetworksOperationType { - return ANEURALNETWORKS_RESHAPE; - }; - } else { - return nullptr; - } - break; - case kTfLiteBuiltinSqueeze: - if (version == 1 && kAndroidSdkVersion >= kMinSdkVersionForNNAPI11) { - return [](const NNAPIOpMappingArgs& mapping_args) - -> ANeuralNetworksOperationType { - auto builtin = reinterpret_cast( - mapping_args.node->builtin_data); - // Note that we add the squeeze dimensions even if the dimensions - // were unspecified (empty), as NNAPI requires the operand. - mapping_args.builder->AddVectorInt32Operand( - builtin->squeeze_dims, - static_cast(builtin->num_squeeze_dims)); - return ANEURALNETWORKS_SQUEEZE; - }; - } else { - return nullptr; - } - case kTfLiteBuiltinL2Normalization: { - auto builtin = - reinterpret_cast(node->builtin_data); - if (builtin->activation != kTfLiteActNone) { - // NNAPI does not support activations - return nullptr; - } - return [](const NNAPIOpMappingArgs& mapping_args) - -> ANeuralNetworksOperationType { - return ANEURALNETWORKS_L2_NORMALIZATION; - }; - } - case kTfLiteBuiltinLocalResponseNormalization: - if (version == 1) { - return [](const NNAPIOpMappingArgs& mapping_args) - -> ANeuralNetworksOperationType { - auto builtin = reinterpret_cast( - mapping_args.node->builtin_data); - mapping_args.builder->AddScalarInt32Operand(builtin->radius); - mapping_args.builder->AddScalarFloat32Operand(builtin->bias); - mapping_args.builder->AddScalarFloat32Operand(builtin->alpha); - mapping_args.builder->AddScalarFloat32Operand(builtin->beta); - return ANEURALNETWORKS_LOCAL_RESPONSE_NORMALIZATION; - }; - } else { - // TODO(miaowang): clean-up code and return early in the unsupported - // case. - return nullptr; - } - break; - case kTfLiteBuiltinLshProjection: - if (version == 1) { - // NNAPI does not support sparse projection correctly (b/111751836). - if (reinterpret_cast(node->builtin_data) - ->type == kTfLiteLshProjectionSparse) { - return nullptr; - } - return [](const NNAPIOpMappingArgs& mapping_args) - -> ANeuralNetworksOperationType { - auto builtin = reinterpret_cast( - mapping_args.node->builtin_data); - mapping_args.builder->AddScalarInt32Operand(builtin->type); - return ANEURALNETWORKS_LSH_PROJECTION; - }; - } else { - return nullptr; - } - break; - case kTfLiteBuiltinConcatenation: - if (version == 1 && - reinterpret_cast(node->builtin_data) - ->activation == kTfLiteActNone) { - if (context->tensors[node->inputs->data[0]].type == kTfLiteUInt8) { - // NNAPI only support concatenating quantized tensor of the same - // scale and offset. - auto first_param = context->tensors[node->inputs->data[0]].params; - for (int i = 0; i < node->inputs->size; i++) { - auto curr_param = context->tensors[node->inputs->data[i]].params; - if (curr_param.scale != first_param.scale || - curr_param.zero_point != first_param.zero_point) { - return nullptr; - } - } - } - return [](const NNAPIOpMappingArgs& mapping_args) - -> ANeuralNetworksOperationType { - auto builtin = reinterpret_cast( - mapping_args.node->builtin_data); - mapping_args.builder->AddScalarInt32Operand(builtin->axis); - return ANEURALNETWORKS_CONCATENATION; - }; - } else { - return nullptr; - } - break; - case kTfLiteBuiltinDequantize: - if (version == 1) { - return [](const NNAPIOpMappingArgs& mapping_args) - -> ANeuralNetworksOperationType { - return ANEURALNETWORKS_DEQUANTIZE; - }; - } else { - return nullptr; - } - break; - case kTfLiteBuiltinFloor: - if (version == 1) { - return [](const NNAPIOpMappingArgs& mapping_args) - -> ANeuralNetworksOperationType { - return ANEURALNETWORKS_FLOOR; - }; - } else { - return nullptr; - } - break; - case kTfLiteBuiltinRelu: - if (version == 1) { - return [](const NNAPIOpMappingArgs& mapping_args) - -> ANeuralNetworksOperationType { - return ANEURALNETWORKS_RELU; - }; - } else { - return nullptr; - } - break; - case kTfLiteBuiltinReluN1To1: - if (version == 1) { - return [](const NNAPIOpMappingArgs& mapping_args) - -> ANeuralNetworksOperationType { - return ANEURALNETWORKS_RELU1; - }; - } else { - return nullptr; - } - break; - case kTfLiteBuiltinRelu6: - if (version == 1) { - return [](const NNAPIOpMappingArgs& mapping_args) - -> ANeuralNetworksOperationType { - return ANEURALNETWORKS_RELU6; - }; - } else { - return nullptr; - } - break; - case kTfLiteBuiltinLogistic: - if (version == 1) { - return [](const NNAPIOpMappingArgs& mapping_args) - -> ANeuralNetworksOperationType { - return ANEURALNETWORKS_LOGISTIC; - }; - } else { - return nullptr; - } - break; - case kTfLiteBuiltinTanh: - // TODO(miaowang): add additional checks for the parameters. - if (version == 1 && - context->tensors[node->inputs->data[0]].type == kTfLiteFloat32) { - // NNAPI only support float tanh. - return [](const NNAPIOpMappingArgs& mapping_args) - -> ANeuralNetworksOperationType { - return ANEURALNETWORKS_TANH; - }; - } else { - return nullptr; - } - break; - case kTfLiteBuiltinSub: - if (version == 1 && kAndroidSdkVersion >= kMinSdkVersionForNNAPI11 && - context->tensors[node->inputs->data[0]].type == kTfLiteFloat32) { - // NNAPI only support float sub. - return [](const NNAPIOpMappingArgs& mapping_args) - -> ANeuralNetworksOperationType { - auto builtin = reinterpret_cast( - mapping_args.node->builtin_data); - mapping_args.builder->AddScalarInt32Operand(builtin->activation); - return ANEURALNETWORKS_SUB; - }; - } else { - return nullptr; - } - break; - case kTfLiteBuiltinDiv: - if (version == 1 && kAndroidSdkVersion >= kMinSdkVersionForNNAPI11 && - context->tensors[node->inputs->data[0]].type == kTfLiteFloat32) { - // NNAPI only support float div. - return [](const NNAPIOpMappingArgs& mapping_args) - -> ANeuralNetworksOperationType { - auto builtin = reinterpret_cast( - mapping_args.node->builtin_data); - mapping_args.builder->AddScalarInt32Operand(builtin->activation); - return ANEURALNETWORKS_DIV; - }; - } else { - return nullptr; - } - break; - case kTfLiteBuiltinPad: - if (version == 1 && kAndroidSdkVersion >= kMinSdkVersionForNNAPI11 && - node->inputs->size == 2 && - context->tensors[node->inputs->data[0]].type == kTfLiteFloat32) { - // NNAPI does not support specifying the padding value. - // NNAPI pads physical zero for quantized tensors, so only delegate - // float pad to NNAPI. - return [](const NNAPIOpMappingArgs& mapping_args) - -> ANeuralNetworksOperationType { - return ANEURALNETWORKS_PAD; - }; - } else { - return nullptr; - } - break; - case kTfLiteBuiltinSpaceToBatchNd: - if (version == 1 && kAndroidSdkVersion >= kMinSdkVersionForNNAPI11) { - return [](const NNAPIOpMappingArgs& mapping_args) - -> ANeuralNetworksOperationType { - return ANEURALNETWORKS_SPACE_TO_BATCH_ND; - }; - } else { - return nullptr; - } - break; - case kTfLiteBuiltinStridedSlice: - if (version == 1 && kAndroidSdkVersion >= kMinSdkVersionForNNAPI11) { - return [](const NNAPIOpMappingArgs& mapping_args) - -> ANeuralNetworksOperationType { - auto builtin = reinterpret_cast( - mapping_args.node->builtin_data); - mapping_args.builder->AddScalarInt32Operand(builtin->begin_mask); - mapping_args.builder->AddScalarInt32Operand(builtin->end_mask); - mapping_args.builder->AddScalarInt32Operand( - builtin->shrink_axis_mask); - return ANEURALNETWORKS_STRIDED_SLICE; - }; - } else { - return nullptr; - } - break; - case kTfLiteBuiltinTranspose: - // Note that the permutation input tensor value dictates the output - // dimensions. - // TODO(b/110888333): Support dynamically-sized tensors in delegates. - if ((version == 1) && - (kAndroidSdkVersion >= kMinSdkVersionForNNAPI11) && - (node->inputs->size > 1) && - (context->tensors[node->inputs->data[1]].allocation_type == - kTfLiteMmapRo)) { - return [](const NNAPIOpMappingArgs& mapping_args) - -> ANeuralNetworksOperationType { - return ANEURALNETWORKS_TRANSPOSE; - }; - } else { - return nullptr; - } - break; - case kTfLiteBuiltinRnn: - // NNAPI only support float32 weights. - if (version == 1 && node->inputs->size == 5 && - context->tensors[node->inputs->data[/*kWeightsTensor*/ 1]].type == - kTfLiteFloat32) { - return [](const NNAPIOpMappingArgs& mapping_args) - -> ANeuralNetworksOperationType { - // NNAPI need both state_in and state_out. - int ann_index; - mapping_args.builder->AddStateFloat32Tensor( - mapping_args.node->inputs->data[/*kHiddenStateTensor*/ 4], - &ann_index); - mapping_args.model_state_outputs->push_back(ann_index); - mapping_args.model_state_tfl_inputs->push_back( - mapping_args.node->inputs->data[/*kHiddenStateTensor*/ 4]); - auto builtin = reinterpret_cast( - mapping_args.node->builtin_data); - mapping_args.builder->AddScalarInt32Operand(builtin->activation); - return ANEURALNETWORKS_RNN; - }; - } else { - return nullptr; - } - break; - case kTfLiteBuiltinSvdf: - // NNAPI only support float32 weights. - if (version == 1 && node->inputs->size == 5 && - context->tensors[node->inputs->data[/*kWeightsFeatureTensor*/ 1]] - .type == kTfLiteFloat32) { - return [](const NNAPIOpMappingArgs& mapping_args) - -> ANeuralNetworksOperationType { - // NNAPI need both state_in and state_out. - int ann_index; - mapping_args.builder->AddStateFloat32Tensor( - mapping_args.node->inputs - ->data[/*kInputActivationStateTensor*/ 4], - &ann_index); - mapping_args.model_state_outputs->push_back(ann_index); - mapping_args.model_state_tfl_inputs->push_back( - mapping_args.node->inputs - ->data[/*kInputActivationStateTensor*/ 4]); - - auto builtin = reinterpret_cast( - mapping_args.node->builtin_data); - mapping_args.builder->AddScalarInt32Operand(builtin->rank); - mapping_args.builder->AddScalarInt32Operand(builtin->activation); - return ANEURALNETWORKS_SVDF; - }; - } else { - return nullptr; - } - break; - case kTfLiteBuiltinLstm: - // NNAPI only support float32 weights. - // TODO(miaowang): add loggings to indicate why the op is rejected. - if (version == 1 && node->inputs->size == 20 && - context->tensors[node->inputs - ->data[/*kInputToOutputWeightsTensor*/ 4]] - .type == kTfLiteFloat32) { - return [](const NNAPIOpMappingArgs& mapping_args) - -> ANeuralNetworksOperationType { - auto builtin = reinterpret_cast( - mapping_args.node->builtin_data); - mapping_args.builder->AddScalarInt32Operand(builtin->activation); - mapping_args.builder->AddScalarFloat32Operand(builtin->cell_clip); - mapping_args.builder->AddScalarFloat32Operand(builtin->proj_clip); - - // Current NNAPI implementation requires the sratch_buffer as - // output. - mapping_args.builder->AddAdditionalFloat32OutputTensor(2); - - // NNAPI need both state_in and state_out for cell_state and - // output_state. - int ann_index; - mapping_args.builder->AddStateFloat32Tensor( - mapping_args.node->inputs - ->data[/*kInputActivationStateTensor*/ 18], - &ann_index); - mapping_args.model_state_outputs->push_back(ann_index); - mapping_args.model_state_tfl_inputs->push_back( - mapping_args.node->inputs - ->data[/*kInputActivationStateTensor*/ 18]); - mapping_args.builder->AddStateFloat32Tensor( - mapping_args.node->inputs->data[/*kInputCellStateTensor*/ 19], - &ann_index); - mapping_args.model_state_outputs->push_back(ann_index); - mapping_args.model_state_tfl_inputs->push_back( - mapping_args.node->inputs->data[/*kInputCellStateTensor*/ 19]); - - return ANEURALNETWORKS_LSTM; - }; - } else { - return nullptr; - } - break; - case kTfLiteBuiltinMean: - // NNAPI does not support generating a scalar as output for MEAN. - if (version == 1 && kAndroidSdkVersion >= kMinSdkVersionForNNAPI11 && - context->tensors[node->inputs->data[0]].type == kTfLiteFloat32 && - context->tensors[node->outputs->data[0]].dims->size > 0) { - return [](const NNAPIOpMappingArgs& mapping_args) - -> ANeuralNetworksOperationType { - auto builtin = reinterpret_cast( - mapping_args.node->builtin_data); - int32_t keep_dims = 0; - if (builtin->keep_dims) keep_dims = 1; - mapping_args.builder->AddScalarInt32Operand(keep_dims); - return ANEURALNETWORKS_MEAN; - }; - } else { - return nullptr; - } - case kTfLiteBuiltinEmbeddingLookup: - // NNAPI only support float32 values. - if (version == 1 && - context->tensors[node->inputs->data[1]].type == kTfLiteFloat32) { - return [](const NNAPIOpMappingArgs& mapping_args) - -> ANeuralNetworksOperationType { - return ANEURALNETWORKS_EMBEDDING_LOOKUP; - }; - } else { - return nullptr; - } - break; - case kTfLiteBuiltinHashtableLookup: - // NNAPI only support float32 output. - if (version == 1 && - context->tensors[node->outputs->data[0]].type == kTfLiteFloat32) { - return [](const NNAPIOpMappingArgs& mapping_args) - -> ANeuralNetworksOperationType { - return ANEURALNETWORKS_HASHTABLE_LOOKUP; - }; - } else { - return nullptr; - } - break; - default: - return nullptr; - } - } - - // Initialize the kernel (a NN model). - TfLiteStatus Init(TfLiteContext* context, - const TfLiteDelegateParams* params) { - for (auto node_index : TfLiteIntArrayView(params->nodes_to_replace)) { - nodes_.push_back(node_index); - } - - if (!nn_model_) { - ANeuralNetworksModel* model; - CHECK_NN(context, ANeuralNetworksModel_create(&model)); - nn_model_.reset(model); - - TF_LITE_ENSURE_STATUS( - BuildGraph(context, params->input_tensors, params->output_tensors)); - } - - if (!nn_compilation_) { - ANeuralNetworksCompilation* compilation; - CHECK_NN(context, ANeuralNetworksCompilation_create(nn_model_.get(), - &compilation)); - CHECK_NN(context, ANeuralNetworksCompilation_finish(compilation)); - nn_compilation_.reset(compilation); - } - return kTfLiteOk; - } - - TfLiteStatus Invoke(TfLiteContext* context, TfLiteNode* node) { - ANeuralNetworksExecution* execution = nullptr; - CHECK_NN(context, ANeuralNetworksExecution_create(nn_compilation_.get(), - &execution)); - - // Set the input tensor buffers. Note: we access tflite tensors using - // absolute indices but NN api indices inputs by relative indices. - int relative_input_index = 0; - - size_t input_offset = 0; - for (auto absolute_input_index : TfLiteIntArrayView(node->inputs)) { - if (absolute_input_index == kOptionalTensor) { - continue; - } - TfLiteTensor* tensor = &context->tensors[absolute_input_index]; - // TODO(miaowang): make sure the delegation works with dequantized weights - // as intermediate tensors. - if (tensor->allocation_type != kTfLiteMmapRo) { - // copy data to pre-allocated shared memory. - memcpy(nn_input_memory_->get_data_ptr() + input_offset, - tensor->data.raw, tensor->bytes); - CHECK_NN(context, ANeuralNetworksExecution_setInputFromMemory( - execution, relative_input_index, nullptr, - nn_input_memory_->get_handle(), input_offset, - tensor->bytes)); - input_offset += tensor->bytes; - relative_input_index++; - } - } - - // Set the output tensor buffers. - int relative_output_index = 0; - size_t output_offset = 0; - for (auto output_index : TfLiteIntArrayView(node->outputs)) { - TfLiteTensor* tensor = &context->tensors[output_index]; - CHECK_NN(context, ANeuralNetworksExecution_setOutputFromMemory( - execution, relative_output_index, nullptr, - nn_output_memory_->get_handle(), output_offset, - tensor->bytes)); - output_offset += tensor->bytes; - relative_output_index++; - } - - // The state_out of previous invocation need to be mapped to state_in of - // current invocation. - for (size_t i = 0; i < model_state_tfl_inputs_.size(); i++) { - int state_tensor_idx = model_state_tfl_inputs_[i]; - TfLiteTensor* tensor = &context->tensors[state_tensor_idx]; - // Here we are using a deep copy for state_in tensors so that we are not - // reading and writing into the same buffer during a invocation. - // TODO(110369471): using double shared buffer to minimize the copies. - CHECK_NN(context, ANeuralNetworksExecution_setOutput( - execution, relative_output_index, nullptr, - tensor->data.raw, tensor->bytes)); - relative_output_index++; - } - // Invoke ANN in blocking fashion. - ANeuralNetworksEvent* event = nullptr; - CHECK_NN(context, ANeuralNetworksExecution_startCompute(execution, &event)); - CHECK_NN(context, ANeuralNetworksEvent_wait(event)); - ANeuralNetworksEvent_free(event); - ANeuralNetworksExecution_free(execution); - - // copy results from shared memory to the destination. - output_offset = 0; - for (auto output_index : TfLiteIntArrayView(node->outputs)) { - TfLiteTensor* tensor = &context->tensors[output_index]; - memcpy(tensor->data.raw, - nn_output_memory_->get_data_ptr() + output_offset, tensor->bytes); - output_offset += tensor->bytes; - } - - return kTfLiteOk; - } - - private: - // ANN API state. - std::unique_ptr nn_model_; - std::unique_ptr - nn_compilation_; - // Node indices that this delegate is responsible for. Indices here - // indexes into the nodes array in the TfLiteContext. - std::vector nodes_; - // Track indices we use - OperandMapping operand_mapping_; - - std::vector model_state_outputs_; - std::vector model_state_tfl_inputs_; - - std::unique_ptr nn_input_memory_; - std::unique_ptr nn_output_memory_; - - TfLiteStatus AddOpsAndTensors(TfLiteContext* context) { - // The operand builder allows creating a single op. We create it at this - // reduced power position rather than in the for loop to avoid reallocating - // the vectors. - NNAPIOpBuilder builder(context, &operand_mapping_, nn_model_.get()); - // Add Tensors - // allocate outside to avoid realloc - for (auto node_index : nodes_) { - // Obtain the op and registration. - TfLiteNode* node; - TfLiteRegistration* reg; - context->GetNodeAndRegistration(context, node_index, &node, ®); - // Map inputs to NN API tensor indices. - for (auto input_index : TfLiteIntArrayView(node->inputs)) { - if (input_index == kOptionalTensor && - (reg->builtin_code == kTfLiteBuiltinLstm || - reg->builtin_code == kTfLiteBuiltinSvdf)) { - // properly handle the optional tensor for LSTM and SVDF. - // currently only support float32. - // TODO(miaowang): make sure this is also able to handle quantized - // tensor when supported by NNAPI. - TF_LITE_ENSURE_STATUS(builder.AddVectorFloat32Operand(nullptr, 0)); - } else { - TF_LITE_ENSURE_STATUS(builder.AddTensorInput(input_index)); - } - } - // Get op type and operands - int nn_op_type = Map(context, reg->builtin_code, reg->version, node)( - {context, &builder, node, &model_state_outputs_, - &model_state_tfl_inputs_}); - // Map outputs to NN API tensor indices. - for (auto output_index : TfLiteIntArrayView(node->outputs)) { - TF_LITE_ENSURE_STATUS(builder.AddTensorOutput(output_index)); - } - - builder.FinalizeAddOperation(nn_op_type); - } - return kTfLiteOk; - } - - TfLiteStatus BuildGraph(TfLiteContext* context, - const TfLiteIntArray* input_tensors, - const TfLiteIntArray* output_tensors) { - // Build the ops and tensors. - TF_LITE_ENSURE_STATUS(AddOpsAndTensors(context)); - // Map input and output tensor indices to ANN - std::vector inputs; - inputs.reserve(input_tensors->size); - std::vector outputs; - outputs.reserve(output_tensors->size); - - size_t total_input_byte_size = 0; - // Make the TensorFlow lite inputs and outputs to ann_indices. - for (int i : TfLiteIntArrayView(input_tensors)) { - // Constant tensors are not NNAPI inputs. - if (i != kOptionalTensor && - context->tensors[i].allocation_type != kTfLiteMmapRo) { - inputs.push_back(operand_mapping_.lite_index_to_ann(i)); - total_input_byte_size += context->tensors[i].bytes; - } - } - - size_t total_output_byte_size = 0; - for (int i : TfLiteIntArrayView(output_tensors)) { - outputs.push_back(operand_mapping_.lite_index_to_ann(i)); - total_output_byte_size += context->tensors[i].bytes; - } - - // Add state output tensors as model inputs - for (int i : model_state_outputs_) { - outputs.push_back(i); - } - - // Tell ANN to declare inputs/outputs - CHECK_NN(context, ANeuralNetworksModel_identifyInputsAndOutputs( - nn_model_.get(), inputs.size(), inputs.data(), - outputs.size(), outputs.data())); - - // Set relaxed computation mode for fp32 if possible. - if (kAndroidSdkVersion >= kMinSdkVersionForNNAPI11) { - CHECK_NN(context, - ANeuralNetworksModel_relaxComputationFloat32toFloat16( - nn_model_.get(), context->allow_fp32_relax_to_fp16)); - } - - // Finalize the model - CHECK_NN(context, ANeuralNetworksModel_finish(nn_model_.get())); - - // Create shared memory pool for inputs and outputs. - nn_input_memory_.reset(new NNMemory("input_pool", total_input_byte_size)); - nn_output_memory_.reset( - new NNMemory("output_pool", total_output_byte_size)); - - return kTfLiteOk; - } -}; - -} // namespace - -// Return a NN API Delegate struct that can check for support of ops. -TfLiteDelegate* NnApiDelegate() { - static TfLiteDelegate delegate = { - .data_ = nullptr, - .Prepare = [](TfLiteContext* context, - TfLiteDelegate* delegate) -> TfLiteStatus { - // Do not check nodes_ if NN API is unavailable. - if (kAndroidSdkVersion < kMinSdkVersionForNNAPI || !NNAPIExists()) { - return kTfLiteOk; - } - - std::vector supported_nodes(1); - // We don't care about all nodes_, we only care about ones in the - // current plan. - TfLiteIntArray* plan; - TF_LITE_ENSURE_STATUS(context->GetExecutionPlan(context, &plan)); - int total_supported_nodes = 0; - - // Check for every node if it is supported - // TODO(b/80625235): Fix this to do more careful checking of versioning. - for (int node_index : TfLiteIntArrayView(plan)) { - TfLiteNode* node; - TfLiteRegistration* registration; - TF_LITE_ENSURE_STATUS(context->GetNodeAndRegistration( - context, node_index, &node, ®istration)); - NNAPIDelegateKernel dummy_kernel; - if (dummy_kernel.Map(context, registration->builtin_code, - registration->version, node)) { - supported_nodes.push_back(node_index); - } - total_supported_nodes += 1; - } - // Put the size at the beginning of the array. - supported_nodes[0] = supported_nodes.size() - 1; - - // NN API Delegate Registration (the pseudo kernel that will invoke NN - // API subgraphs) - static const TfLiteRegistration nnapi_delegate_kernel = { - .init = [](TfLiteContext* context, const char* buffer, - size_t length) -> void* { - const TfLiteDelegateParams* params = - reinterpret_cast(buffer); - NNAPIDelegateKernel* kernel_state = new NNAPIDelegateKernel; - kernel_state->Init(context, params); - return kernel_state; - }, - - .free = [](TfLiteContext* context, void* buffer) -> void { - delete reinterpret_cast(buffer); - }, - - .prepare = [](TfLiteContext* context, - TfLiteNode* node) -> TfLiteStatus { - // Since the underlying resize happened ahead of delegation - // worked. This does nothing. - return kTfLiteOk; - }, - - .invoke = [](TfLiteContext* context, - TfLiteNode* node) -> TfLiteStatus { - NNAPIDelegateKernel* state = - reinterpret_cast(node->user_data); - return state->Invoke(context, node); - }, - - .builtin_code = kTfLiteBuiltinDelegate, - }; - - // Request TFLite to partition the graph and make kernels - // for each independent subgraph a new nnapi_delegate_kernel. - context->ReplaceSubgraphsWithDelegateKernels( - context, nnapi_delegate_kernel, - reinterpret_cast(supported_nodes.data()), - delegate); - return kTfLiteOk; - }}; - - return &delegate; -} - -} // namespace tflite diff --git a/tensorflow/contrib/lite/delegates/nnapi/nnapi_delegate.h b/tensorflow/contrib/lite/delegates/nnapi/nnapi_delegate.h deleted file mode 100644 index 4852b7697432c30c1258e790b97ce2563e7f9711..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/delegates/nnapi/nnapi_delegate.h +++ /dev/null @@ -1,31 +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_CONTRIB_LITE_DELEGATES_NNAPI_NNAPI_DELEGATE_H_ -#define TENSORFLOW_CONTRIB_LITE_DELEGATES_NNAPI_NNAPI_DELEGATE_H_ - -#include "tensorflow/contrib/lite/c/c_api_internal.h" - -namespace tflite { - -// Return a delegate that can be used to use the NN API. -// e.g. -// NnApiDelegate* delegate = NnApiDelegate(); -// interpreter->ModifyGraphWithDelegate(&delegate); -// NnApiDelegate() returns a singleton, so you should not free this -// pointer or worry about its lifetime. -TfLiteDelegate* NnApiDelegate(); -} // namespace tflite - -#endif // TENSORFLOW_CONTRIB_LITE_DELEGATES_NNAPI_NNAPI_DELEGATE_H_ diff --git a/tensorflow/contrib/lite/error_reporter.h b/tensorflow/contrib/lite/error_reporter.h deleted file mode 100644 index 5c20eedc255ca6f7578873593aa86759fbeb490b..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/error_reporter.h +++ /dev/null @@ -1,22 +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. -==============================================================================*/ -// Compatibility shim for moved header location. -#ifndef TENSORFLOW_CONTRIB_LITE_ERROR_REPORTER_H_ -#define TENSORFLOW_CONTRIB_LITE_ERROR_REPORTER_H_ - -#include "tensorflow/contrib/lite/core/api/error_reporter.h" -#include "tensorflow/contrib/lite/stderr_reporter.h" - -#endif // TENSORFLOW_CONTRIB_LITE_ERROR_REPORTER_H_ diff --git a/tensorflow/contrib/lite/examples/android/BUILD b/tensorflow/contrib/lite/examples/android/BUILD deleted file mode 100644 index d180cb478566a9e5df24b2e67445f24a2f623215..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/examples/android/BUILD +++ /dev/null @@ -1,61 +0,0 @@ -# Description: -# TensorFlow camera demo app for Android. - -load("@build_bazel_rules_android//android:rules.bzl", "android_binary") - -package(default_visibility = ["//visibility:public"]) - -licenses(["notice"]) # Apache 2.0 - -exports_files(["LICENSE"]) - -# Build the demo native demo lib from the original directory to reduce code -# reuse. Note that the Java counterparts (ObjectTracker.java and -# ImageUtils.java) are still duplicated. -cc_library( - name = "tensorflow_native_libs", - srcs = [ - "//tensorflow/examples/android:libtensorflow_demo.so", - ], - tags = [ - "manual", - "notap", - ], -) - -android_binary( - name = "tflite_demo", - srcs = glob([ - "app/src/main/java/**/*.java", - ]), - aapt_version = "aapt", - # Package assets from assets dir as well as all model targets. - # Remove undesired models (and corresponding Activities in source) - # to reduce APK size. - assets = [ - "//tensorflow/contrib/lite/examples/android/app/src/main/assets:labels_mobilenet_quant_v1_224.txt", - "@tflite_mobilenet//:mobilenet_quant_v1_224.tflite", - "@tflite_conv_actions_frozen//:conv_actions_frozen.tflite", - "//tensorflow/contrib/lite/examples/android/app/src/main/assets:conv_actions_labels.txt", - "@tflite_mobilenet_ssd//:mobilenet_ssd.tflite", - "@tflite_mobilenet_ssd_quant//:detect.tflite", - "//tensorflow/contrib/lite/examples/android/app/src/main/assets:box_priors.txt", - "//tensorflow/contrib/lite/examples/android/app/src/main/assets:coco_labels_list.txt", - ], - assets_dir = "", - custom_package = "org.tensorflow.lite.demo", - inline_constants = 1, - manifest = "app/src/main/AndroidManifest.xml", - nocompress_extensions = [ - ".tflite", - ], - resource_files = glob(["app/src/main/res/**"]), - tags = [ - "manual", - "notap", - ], - deps = [ - ":tensorflow_native_libs", - "//tensorflow/contrib/lite/java:tensorflowlite", - ], -) diff --git a/tensorflow/contrib/lite/examples/android/app/README.md b/tensorflow/contrib/lite/examples/android/app/README.md deleted file mode 100644 index 7347147f997540e67c2c713b597dc90d933c5cb8..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/examples/android/app/README.md +++ /dev/null @@ -1,54 +0,0 @@ -# TF Lite Android App Example - -A simple Android example that demonstrates image classification and object -detection using the camera, as well as speech recognition using the microphone. - -## Building in Android Studio with TensorFlow Lite AAR from JCenter. -The build.gradle is configured to use TensorFlow Lite's nightly build. - -If you see a build error related to compatibility with Tensorflow Lite's Java -API (example: method X is undefined for type Interpreter), there has likely been -a backwards compatible change to the API. You will need to pull new app code -that's compatible with the nightly build and may need to first wait a few days -for our external and internal code to merge. - -## Building from Source with Bazel - -1. Follow the [Bazel steps for the TF Demo App](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/examples/android#bazel): - - 1. [Install Bazel and Android Prerequisites](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/examples/android#install-bazel-and-android-prerequisites). - It's easiest with Android Studio. - - - You'll need at least SDK version 23. - - Make sure to install the latest version of Bazel. Some distributions - ship with Bazel 0.5.4, which is too old. - - Bazel requires Android Build Tools `26.0.1` or higher. - - You also need to install the Android Support Repository, available - through Android Studio under `Android SDK Manager -> SDK Tools -> - Android Support Repository`. - - 2. [Edit your `WORKSPACE`](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/examples/android#edit-workspace) - to add SDK and NDK targets. - - NOTE: As long as you have the SDK and NDK installed, the `./configure` - script will create these rules for you. Answer "Yes" when the script asks - to automatically configure the `./WORKSPACE`. - - - Make sure the `api_level` in `WORKSPACE` is set to an SDK version that - you have installed. - - By default, Android Studio will install the SDK to `~/Android/Sdk` and - the NDK to `~/Android/Sdk/ndk-bundle`. - -2. Build this demo app with Bazel. The demo needs C++11. We configure the fat_apk_cpu flag to package support for 4 hardware variants. You may replace it with --config=android_arm64 on a 64-bit device and --config=android_arm for 32-bit device: - - ```shell - bazel build -c opt --cxxopt='--std=c++11' --fat_apk_cpu=x86,x86_64,arm64-v8a,armeabi-v7a \ - //tensorflow/contrib/lite/examples/android:tflite_demo - ``` - -3. Install the demo on a - [debug-enabled device](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/examples/android#install): - - ```shell - adb install bazel-bin/tensorflow/contrib/lite/examples/android/tflite_demo.apk - ``` diff --git a/tensorflow/contrib/lite/examples/android/app/build.gradle b/tensorflow/contrib/lite/examples/android/app/build.gradle deleted file mode 100644 index 35e78878526a4956448cdd81eb848cf73c105754..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/examples/android/app/build.gradle +++ /dev/null @@ -1,54 +0,0 @@ -apply plugin: 'com.android.application' - -android { - compileSdkVersion 26 - buildToolsVersion '26.0.2' - defaultConfig { - applicationId "org.tensorflow.lite.demo" - minSdkVersion 15 - targetSdkVersion 26 - versionCode 1 - versionName "1.0" - - // Remove this block. - jackOptions { - enabled true - } - } - lintOptions { - abortOnError false - } - buildTypes { - release { - minifyEnabled false - proguardFiles getDefaultProguardFile('proguard-android.txt'), 'proguard-rules.pro' - } - } - aaptOptions { - noCompress "tflite" - } - - compileOptions { - sourceCompatibility JavaVersion.VERSION_1_8 - targetCompatibility JavaVersion.VERSION_1_8 - } -} - -repositories { - maven { - url 'https://google.bintray.com/tensorflow' - } -} - -// import DownloadModels task -project.ext.ASSET_DIR = projectDir.toString() + '/src/main/assets' -project.ext.TMP_DIR = project.buildDir.toString() + '/downloads' - -// Download default models; if you wish to use your own models then -// place them in the "assets" directory and comment out this line. -apply from: "download-models.gradle" - -dependencies { - compile fileTree(dir: 'libs', include: ['*.jar']) - compile 'org.tensorflow:tensorflow-lite:0.0.0-nightly' -} diff --git a/tensorflow/contrib/lite/examples/android/app/download-models.gradle b/tensorflow/contrib/lite/examples/android/app/download-models.gradle deleted file mode 100644 index c100e37c16f38a65f7b1f64a3f6e3eaa1477e8eb..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/examples/android/app/download-models.gradle +++ /dev/null @@ -1,74 +0,0 @@ -/* - * download-models.gradle - * Downloads model files from ${MODEL_URL} into application's asset folder - * Input: - * project.ext.TMP_DIR: absolute path to hold downloaded zip files - * project.ext.ASSET_DIR: absolute path to save unzipped model files - * Output: - * 3 model files will be downloaded into given folder of ext.ASSET_DIR - */ -// hard coded model files -// LINT.IfChange - -def models = ['conv_actions_tflite.zip', - 'mobilenet_ssd_tflite_v1.zip', - 'mobilenet_v1_224_android_quant_2017_11_08.zip', - 'coco_ssd_mobilenet_v1_1.0_quant_2018_06_29.zip'] -// LINT.ThenChange(//tensorflow/contrib/lite/examples/android/BUILD) - -// Root URL for model archives -def MODEL_URL = 'https://storage.googleapis.com/download.tensorflow.org/models/tflite' - -buildscript { - repositories { - jcenter() - } - dependencies { - classpath 'de.undercouch:gradle-download-task:3.2.0' - } -} - -import de.undercouch.gradle.tasks.download.Download -task downloadFile(type: Download){ - for (f in models) { - def modelUrl = MODEL_URL + "/" + f - println "Downloading ${f} from ${modelUrl}" - src modelUrl - } - - dest new File(project.ext.TMP_DIR) - overwrite true -} - -task extractModels(type: Copy) { - for (f in models) { - def localFile = f.split("/")[-1] - from zipTree(project.ext.TMP_DIR + '/' + localFile) - } - - into file(project.ext.ASSET_DIR) - fileMode 0644 - exclude '**/LICENSE' - - def needDownload = false - for (f in models) { - def localFile = f.split("/")[-1] - if (!(new File(project.ext.TMP_DIR + '/' + localFile)).exists()) { - needDownload = true - } - } - - if (needDownload) { - dependsOn downloadFile - } -} - -tasks.whenTaskAdded { task -> - if (task.name == 'assembleDebug') { - task.dependsOn 'extractModels' - } - if (task.name == 'assembleRelease') { - task.dependsOn 'extractModels' - } -} - diff --git a/tensorflow/contrib/lite/examples/android/build.gradle b/tensorflow/contrib/lite/examples/android/build.gradle deleted file mode 100644 index 66a62a921a7f492df30b3de2e5dc4b68fc84f1d9..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/examples/android/build.gradle +++ /dev/null @@ -1,24 +0,0 @@ -// Top-level build file where you can add configuration options common to all sub-projects/modules. - -buildscript { - repositories { - jcenter() - } - dependencies { - classpath 'com.android.tools.build:gradle:3.0.1' - - // NOTE: Do not place your application dependencies here; they belong - // in the individual module build.gradle files - } -} - -allprojects { - repositories { - google() - jcenter() - } -} - -task clean(type: Delete) { - delete rootProject.buildDir -} diff --git a/tensorflow/contrib/lite/examples/ios/camera/CameraExampleViewController.h b/tensorflow/contrib/lite/examples/ios/camera/CameraExampleViewController.h deleted file mode 100644 index fb5800e86d365b56f1b52147c3f9cc8d7211f8c3..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/examples/ios/camera/CameraExampleViewController.h +++ /dev/null @@ -1,48 +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. - -#import -#import - -#include - -#include "tensorflow/contrib/lite/kernels/register.h" -#include "tensorflow/contrib/lite/model.h" - -@interface CameraExampleViewController - : UIViewController { - IBOutlet UIView* previewView; - AVCaptureVideoPreviewLayer* previewLayer; - AVCaptureVideoDataOutput* videoDataOutput; - dispatch_queue_t videoDataOutputQueue; - UIView* flashView; - BOOL isUsingFrontFacingCamera; - NSMutableDictionary* oldPredictionValues; - NSMutableArray* labelLayers; - AVCaptureSession* session; - - std::vector labels; - std::unique_ptr model; - tflite::ops::builtin::BuiltinOpResolver resolver; - std::unique_ptr interpreter; - - double total_latency; - int total_count; -} -@property(strong, nonatomic) CATextLayer* predictionTextLayer; - -- (IBAction)takePicture:(id)sender; -- (IBAction)switchCameras:(id)sender; - -@end diff --git a/tensorflow/contrib/lite/examples/ios/camera/Podfile b/tensorflow/contrib/lite/examples/ios/camera/Podfile deleted file mode 100644 index f460693122af8353286ea7069d5db873fedfc9b3..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/examples/ios/camera/Podfile +++ /dev/null @@ -1,5 +0,0 @@ -platform :ios, '8.0' -inhibit_all_warnings! - -target 'tflite_camera_example' - pod 'TensorFlowLite', '1.10.1' diff --git a/tensorflow/contrib/lite/examples/ios/camera/tflite_camera_example.xcodeproj/project.pbxproj b/tensorflow/contrib/lite/examples/ios/camera/tflite_camera_example.xcodeproj/project.pbxproj deleted file mode 100644 index 98d3b5bb8ad45bf34f6996b3361291896a451a6f..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/examples/ios/camera/tflite_camera_example.xcodeproj/project.pbxproj +++ /dev/null @@ -1,405 +0,0 @@ -// !$*UTF8*$! -{ - archiveVersion = 1; - classes = { - }; - objectVersion = 46; - objects = { - -/* Begin PBXBuildFile section */ - 1C3C9DCC1ED3AB4200B8B5FA /* main.mm in Sources */ = {isa = PBXBuildFile; fileRef = 1C3C9DCA1ED3AB4200B8B5FA /* main.mm */; }; - 1C99111C1ED3B0E600A6BFB9 /* MainStoryboard_iPhone.storyboard in Resources */ = {isa = PBXBuildFile; fileRef = 1C99111B1ED3B0E600A6BFB9 /* MainStoryboard_iPhone.storyboard */; }; - 1CA5EB931ED3ABFB00247A34 /* CoreMedia.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 1CA5EB921ED3ABFB00247A34 /* CoreMedia.framework */; }; - 1CB47D491ED3AD1700DF7666 /* AVFoundation.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 1CB47D481ED3AD1700DF7666 /* AVFoundation.framework */; }; - 1CDB2D491ED3A9CD007929E9 /* CameraExampleAppDelegate.m in Sources */ = {isa = PBXBuildFile; fileRef = 1CDB2D431ED3A9CD007929E9 /* CameraExampleAppDelegate.m */; }; - 1CDB2D4A1ED3A9CD007929E9 /* CameraExampleViewController.mm in Sources */ = {isa = PBXBuildFile; fileRef = 1CDB2D451ED3A9CD007929E9 /* CameraExampleViewController.mm */; }; - 1CDB2D4E1ED3AA35007929E9 /* Info.plist in Resources */ = {isa = PBXBuildFile; fileRef = 1CDB2D4D1ED3AA35007929E9 /* Info.plist */; }; - 54DC6C3C5F734F3A58069F0C /* libPods-tflite_camera_example.a in Frameworks */ = {isa = PBXBuildFile; fileRef = 3BA8BF92C84895BFE59D8236 /* libPods-tflite_camera_example.a */; }; - AC1F82661FBA3CBD0052BA77 /* labels.txt in Resources */ = {isa = PBXBuildFile; fileRef = AC1F82641FBA3CBD0052BA77 /* labels.txt */; }; - ACA1A4CA1FBB6C28009B8D86 /* mobilenet_quant_v1_224.tflite in Resources */ = {isa = PBXBuildFile; fileRef = ACA1A4C91FBB6C28009B8D86 /* mobilenet_quant_v1_224.tflite */; }; -/* End PBXBuildFile section */ - -/* Begin PBXFileReference section */ - 1C0D73481ECCC41B008C1DAB /* CoreImage.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = CoreImage.framework; path = System/Library/Frameworks/CoreImage.framework; sourceTree = SDKROOT; }; - 1C0D734A1ECCC460008C1DAB /* CoreGraphics.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = CoreGraphics.framework; path = System/Library/Frameworks/CoreGraphics.framework; sourceTree = SDKROOT; }; - 1C3C9DCA1ED3AB4200B8B5FA /* main.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = main.mm; sourceTree = ""; }; - 1C564C0D1ED3A92E00087306 /* tflite_camera_example.app */ = {isa = PBXFileReference; explicitFileType = wrapper.application; includeInIndex = 0; path = tflite_camera_example.app; sourceTree = BUILT_PRODUCTS_DIR; }; - 1C99111B1ED3B0E600A6BFB9 /* MainStoryboard_iPhone.storyboard */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = file.storyboard; path = MainStoryboard_iPhone.storyboard; sourceTree = ""; }; - 1CA45FFE1ECCC356002FA6A4 /* UIKit.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = UIKit.framework; path = System/Library/Frameworks/UIKit.framework; sourceTree = SDKROOT; }; - 1CA5EB921ED3ABFB00247A34 /* CoreMedia.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = CoreMedia.framework; path = System/Library/Frameworks/CoreMedia.framework; sourceTree = SDKROOT; }; - 1CB47D481ED3AD1700DF7666 /* AVFoundation.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = AVFoundation.framework; path = System/Library/Frameworks/AVFoundation.framework; sourceTree = SDKROOT; }; - 1CDB2D421ED3A9CD007929E9 /* CameraExampleAppDelegate.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = CameraExampleAppDelegate.h; sourceTree = ""; }; - 1CDB2D431ED3A9CD007929E9 /* CameraExampleAppDelegate.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = CameraExampleAppDelegate.m; sourceTree = ""; }; - 1CDB2D441ED3A9CD007929E9 /* CameraExampleViewController.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = CameraExampleViewController.h; sourceTree = ""; }; - 1CDB2D451ED3A9CD007929E9 /* CameraExampleViewController.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = CameraExampleViewController.mm; sourceTree = ""; }; - 1CDB2D4D1ED3AA35007929E9 /* Info.plist */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = text.plist.xml; path = Info.plist; sourceTree = ""; }; - 3BA8BF92C84895BFE59D8236 /* libPods-tflite_camera_example.a */ = {isa = PBXFileReference; explicitFileType = archive.ar; includeInIndex = 0; path = "libPods-tflite_camera_example.a"; sourceTree = BUILT_PRODUCTS_DIR; }; - 3BC5BE4BBD09374D3E98F082 /* Pods-tflite_camera_example.debug.xcconfig */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = text.xcconfig; name = "Pods-tflite_camera_example.debug.xcconfig"; path = "Pods/Target Support Files/Pods-tflite_camera_example/Pods-tflite_camera_example.debug.xcconfig"; sourceTree = ""; }; - 55ED318E8D29C8AFEF03DF1E /* Pods-tflite_camera_example.release.xcconfig */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = text.xcconfig; name = "Pods-tflite_camera_example.release.xcconfig"; path = "Pods/Target Support Files/Pods-tflite_camera_example/Pods-tflite_camera_example.release.xcconfig"; sourceTree = ""; }; - AC1F82641FBA3CBD0052BA77 /* labels.txt */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = text; path = labels.txt; sourceTree = ""; }; - ACA1A4C91FBB6C28009B8D86 /* mobilenet_quant_v1_224.tflite */ = {isa = PBXFileReference; lastKnownFileType = file; path = mobilenet_quant_v1_224.tflite; sourceTree = ""; }; -/* End PBXFileReference section */ - -/* Begin PBXFrameworksBuildPhase section */ - 1C564C0A1ED3A92E00087306 /* Frameworks */ = { - isa = PBXFrameworksBuildPhase; - buildActionMask = 2147483647; - files = ( - 1CB47D491ED3AD1700DF7666 /* AVFoundation.framework in Frameworks */, - 1CA5EB931ED3ABFB00247A34 /* CoreMedia.framework in Frameworks */, - 54DC6C3C5F734F3A58069F0C /* libPods-tflite_camera_example.a in Frameworks */, - ); - runOnlyForDeploymentPostprocessing = 0; - }; -/* End PBXFrameworksBuildPhase section */ - -/* Begin PBXGroup section */ - 24D7686C331131624F4454A0 /* Frameworks */ = { - isa = PBXGroup; - children = ( - 1CB47D481ED3AD1700DF7666 /* AVFoundation.framework */, - 1CA5EB921ED3ABFB00247A34 /* CoreMedia.framework */, - 1C0D734A1ECCC460008C1DAB /* CoreGraphics.framework */, - 1C0D73481ECCC41B008C1DAB /* CoreImage.framework */, - 1CA45FFE1ECCC356002FA6A4 /* UIKit.framework */, - 3BA8BF92C84895BFE59D8236 /* libPods-tflite_camera_example.a */, - ); - name = Frameworks; - sourceTree = ""; - }; - 3E9FC355632FB928EA23BEED /* Pods */ = { - isa = PBXGroup; - children = ( - 3BC5BE4BBD09374D3E98F082 /* Pods-tflite_camera_example.debug.xcconfig */, - 55ED318E8D29C8AFEF03DF1E /* Pods-tflite_camera_example.release.xcconfig */, - ); - name = Pods; - sourceTree = ""; - }; - 591157921CF4011C00C31E3A = { - isa = PBXGroup; - children = ( - 1C99111B1ED3B0E600A6BFB9 /* MainStoryboard_iPhone.storyboard */, - 1C3C9DCA1ED3AB4200B8B5FA /* main.mm */, - 1CDB2D4D1ED3AA35007929E9 /* Info.plist */, - 1CDB2D421ED3A9CD007929E9 /* CameraExampleAppDelegate.h */, - 1CDB2D431ED3A9CD007929E9 /* CameraExampleAppDelegate.m */, - 1CDB2D441ED3A9CD007929E9 /* CameraExampleViewController.h */, - 1CDB2D451ED3A9CD007929E9 /* CameraExampleViewController.mm */, - 59A3CFF31CF4E68100C4259F /* data */, - 5911579C1CF4011C00C31E3A /* Products */, - 3E9FC355632FB928EA23BEED /* Pods */, - 24D7686C331131624F4454A0 /* Frameworks */, - ); - sourceTree = ""; - }; - 5911579C1CF4011C00C31E3A /* Products */ = { - isa = PBXGroup; - children = ( - 1C564C0D1ED3A92E00087306 /* tflite_camera_example.app */, - ); - name = Products; - sourceTree = ""; - }; - 59A3CFF31CF4E68100C4259F /* data */ = { - isa = PBXGroup; - children = ( - ACA1A4C91FBB6C28009B8D86 /* mobilenet_quant_v1_224.tflite */, - AC1F82641FBA3CBD0052BA77 /* labels.txt */, - ); - path = data; - sourceTree = ""; - }; -/* End PBXGroup section */ - -/* Begin PBXNativeTarget section */ - 1C564C0C1ED3A92E00087306 /* tflite_camera_example */ = { - isa = PBXNativeTarget; - buildConfigurationList = 1C564C351ED3A92E00087306 /* Build configuration list for PBXNativeTarget "tflite_camera_example" */; - buildPhases = ( - 66DAEAAEE9EF6550C3A061E0 /* [CP] Check Pods Manifest.lock */, - 1C564C091ED3A92E00087306 /* Sources */, - 1C564C0A1ED3A92E00087306 /* Frameworks */, - 1C564C0B1ED3A92E00087306 /* Resources */, - 00E875C3B066535AE6B77101 /* [CP] Embed Pods Frameworks */, - 5C2D02120E3E5E09567AA946 /* [CP] Copy Pods Resources */, - ); - buildRules = ( - ); - dependencies = ( - ); - name = tflite_camera_example; - productName = tflite_camera_example; - productReference = 1C564C0D1ED3A92E00087306 /* tflite_camera_example.app */; - productType = "com.apple.product-type.application"; - }; -/* End PBXNativeTarget section */ - -/* Begin PBXProject section */ - 591157931CF4011C00C31E3A /* Project object */ = { - isa = PBXProject; - attributes = { - LastSwiftUpdateCheck = 0830; - LastUpgradeCheck = 0830; - ORGANIZATIONNAME = Google; - TargetAttributes = { - 1C564C0C1ED3A92E00087306 = { - CreatedOnToolsVersion = 8.3.2; - DevelopmentTeam = EQHXZ8M8AV; - ProvisioningStyle = Automatic; - }; - }; - }; - buildConfigurationList = 591157961CF4011C00C31E3A /* Build configuration list for PBXProject "tflite_camera_example" */; - compatibilityVersion = "Xcode 3.2"; - developmentRegion = English; - hasScannedForEncodings = 0; - knownRegions = ( - en, - Base, - ); - mainGroup = 591157921CF4011C00C31E3A; - productRefGroup = 5911579C1CF4011C00C31E3A /* Products */; - projectDirPath = ""; - projectRoot = ""; - targets = ( - 1C564C0C1ED3A92E00087306 /* tflite_camera_example */, - ); - }; -/* End PBXProject section */ - -/* Begin PBXResourcesBuildPhase section */ - 1C564C0B1ED3A92E00087306 /* Resources */ = { - isa = PBXResourcesBuildPhase; - buildActionMask = 2147483647; - files = ( - ACA1A4CA1FBB6C28009B8D86 /* mobilenet_quant_v1_224.tflite in Resources */, - 1C99111C1ED3B0E600A6BFB9 /* MainStoryboard_iPhone.storyboard in Resources */, - 1CDB2D4E1ED3AA35007929E9 /* Info.plist in Resources */, - AC1F82661FBA3CBD0052BA77 /* labels.txt in Resources */, - ); - runOnlyForDeploymentPostprocessing = 0; - }; -/* End PBXResourcesBuildPhase section */ - -/* Begin PBXShellScriptBuildPhase section */ - 00E875C3B066535AE6B77101 /* [CP] Embed Pods Frameworks */ = { - isa = PBXShellScriptBuildPhase; - buildActionMask = 2147483647; - files = ( - ); - inputPaths = ( - ); - name = "[CP] Embed Pods Frameworks"; - outputPaths = ( - ); - runOnlyForDeploymentPostprocessing = 0; - shellPath = /bin/sh; - shellScript = "\"${SRCROOT}/Pods/Target Support Files/Pods-tflite_camera_example/Pods-tflite_camera_example-frameworks.sh\"\n"; - showEnvVarsInLog = 0; - }; - 5C2D02120E3E5E09567AA946 /* [CP] Copy Pods Resources */ = { - isa = PBXShellScriptBuildPhase; - buildActionMask = 2147483647; - files = ( - ); - inputPaths = ( - ); - name = "[CP] Copy Pods Resources"; - outputPaths = ( - ); - runOnlyForDeploymentPostprocessing = 0; - shellPath = /bin/sh; - shellScript = "\"${SRCROOT}/Pods/Target Support Files/Pods-tflite_camera_example/Pods-tflite_camera_example-resources.sh\"\n"; - showEnvVarsInLog = 0; - }; - 66DAEAAEE9EF6550C3A061E0 /* [CP] Check Pods Manifest.lock */ = { - isa = PBXShellScriptBuildPhase; - buildActionMask = 2147483647; - files = ( - ); - inputPaths = ( - "${PODS_PODFILE_DIR_PATH}/Podfile.lock", - "${PODS_ROOT}/Manifest.lock", - ); - name = "[CP] Check Pods Manifest.lock"; - outputPaths = ( - "$(DERIVED_FILE_DIR)/Pods-tflite_camera_example-checkManifestLockResult.txt", - ); - runOnlyForDeploymentPostprocessing = 0; - shellPath = /bin/sh; - shellScript = "diff \"${PODS_PODFILE_DIR_PATH}/Podfile.lock\" \"${PODS_ROOT}/Manifest.lock\" > /dev/null\nif [ $? != 0 ] ; then\n # print error to STDERR\n echo \"error: The sandbox is not in sync with the Podfile.lock. Run 'pod install' or update your CocoaPods installation.\" >&2\n exit 1\nfi\n# This output is used by Xcode 'outputs' to avoid re-running this script phase.\necho \"SUCCESS\" > \"${SCRIPT_OUTPUT_FILE_0}\"\n"; - showEnvVarsInLog = 0; - }; -/* End PBXShellScriptBuildPhase section */ - -/* Begin PBXSourcesBuildPhase section */ - 1C564C091ED3A92E00087306 /* Sources */ = { - isa = PBXSourcesBuildPhase; - buildActionMask = 2147483647; - files = ( - 1CDB2D4A1ED3A9CD007929E9 /* CameraExampleViewController.mm in Sources */, - 1CDB2D491ED3A9CD007929E9 /* CameraExampleAppDelegate.m in Sources */, - 1C3C9DCC1ED3AB4200B8B5FA /* main.mm in Sources */, - ); - runOnlyForDeploymentPostprocessing = 0; - }; -/* End PBXSourcesBuildPhase section */ - -/* Begin XCBuildConfiguration section */ - 1C564C361ED3A92E00087306 /* Debug */ = { - isa = XCBuildConfiguration; - baseConfigurationReference = 3BC5BE4BBD09374D3E98F082 /* Pods-tflite_camera_example.debug.xcconfig */; - buildSettings = { - ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon; - CLANG_ANALYZER_NONNULL = YES; - CLANG_ANALYZER_NUMBER_OBJECT_CONVERSION = YES_AGGRESSIVE; - CLANG_WARN_DOCUMENTATION_COMMENTS = YES; - DEVELOPMENT_TEAM = EQHXZ8M8AV; - INFOPLIST_FILE = Info.plist; - IPHONEOS_DEPLOYMENT_TARGET = 10.3; - LD_RUNPATH_SEARCH_PATHS = "$(inherited) @executable_path/Frameworks"; - PRODUCT_BUNDLE_IDENTIFIER = "com.pf.tf-camera-example"; - PRODUCT_NAME = "$(TARGET_NAME)"; - SWIFT_ACTIVE_COMPILATION_CONDITIONS = DEBUG; - SWIFT_OPTIMIZATION_LEVEL = "-Onone"; - SWIFT_VERSION = 3.0; - }; - name = Debug; - }; - 1C564C371ED3A92E00087306 /* Release */ = { - isa = XCBuildConfiguration; - baseConfigurationReference = 55ED318E8D29C8AFEF03DF1E /* Pods-tflite_camera_example.release.xcconfig */; - buildSettings = { - ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon; - CLANG_ANALYZER_NONNULL = YES; - CLANG_ANALYZER_NUMBER_OBJECT_CONVERSION = YES_AGGRESSIVE; - CLANG_WARN_DOCUMENTATION_COMMENTS = YES; - DEVELOPMENT_TEAM = EQHXZ8M8AV; - INFOPLIST_FILE = Info.plist; - IPHONEOS_DEPLOYMENT_TARGET = 10.3; - LD_RUNPATH_SEARCH_PATHS = "$(inherited) @executable_path/Frameworks"; - PRODUCT_BUNDLE_IDENTIFIER = "com.pf.tf-camera-example"; - PRODUCT_NAME = "$(TARGET_NAME)"; - SWIFT_OPTIMIZATION_LEVEL = "-Owholemodule"; - SWIFT_VERSION = 3.0; - }; - name = Release; - }; - 591157B01CF4011D00C31E3A /* Debug */ = { - isa = XCBuildConfiguration; - buildSettings = { - ALWAYS_SEARCH_USER_PATHS = NO; - CLANG_CXX_LANGUAGE_STANDARD = "gnu++0x"; - CLANG_CXX_LIBRARY = "libc++"; - CLANG_ENABLE_MODULES = YES; - CLANG_ENABLE_OBJC_ARC = YES; - CLANG_WARN_BOOL_CONVERSION = YES; - CLANG_WARN_CONSTANT_CONVERSION = YES; - CLANG_WARN_DIRECT_OBJC_ISA_USAGE = YES_ERROR; - CLANG_WARN_EMPTY_BODY = YES; - CLANG_WARN_ENUM_CONVERSION = YES; - CLANG_WARN_INFINITE_RECURSION = YES; - CLANG_WARN_INT_CONVERSION = YES; - CLANG_WARN_OBJC_ROOT_CLASS = YES_ERROR; - CLANG_WARN_SUSPICIOUS_MOVE = YES; - CLANG_WARN_UNREACHABLE_CODE = YES; - CLANG_WARN__DUPLICATE_METHOD_MATCH = YES; - "CODE_SIGN_IDENTITY[sdk=iphoneos*]" = "iPhone Developer"; - COPY_PHASE_STRIP = NO; - DEBUG_INFORMATION_FORMAT = dwarf; - ENABLE_STRICT_OBJC_MSGSEND = YES; - ENABLE_TESTABILITY = YES; - GCC_C_LANGUAGE_STANDARD = gnu99; - GCC_DYNAMIC_NO_PIC = NO; - GCC_NO_COMMON_BLOCKS = YES; - GCC_OPTIMIZATION_LEVEL = 0; - GCC_PREPROCESSOR_DEFINITIONS = ( - "DEBUG=1", - "$(inherited)", - ); - GCC_WARN_64_TO_32_BIT_CONVERSION = YES; - GCC_WARN_ABOUT_RETURN_TYPE = YES_ERROR; - GCC_WARN_UNDECLARED_SELECTOR = YES; - GCC_WARN_UNINITIALIZED_AUTOS = YES_AGGRESSIVE; - GCC_WARN_UNUSED_FUNCTION = YES; - GCC_WARN_UNUSED_VARIABLE = YES; - HEADER_SEARCH_PATHS = ( - "$(inherited)", - ); - IPHONEOS_DEPLOYMENT_TARGET = 8.0; - MTL_ENABLE_DEBUG_INFO = YES; - ONLY_ACTIVE_ARCH = YES; - SDKROOT = iphoneos; - TARGETED_DEVICE_FAMILY = "1,2"; - }; - name = Debug; - }; - 591157B11CF4011D00C31E3A /* Release */ = { - isa = XCBuildConfiguration; - buildSettings = { - ALWAYS_SEARCH_USER_PATHS = NO; - CLANG_CXX_LANGUAGE_STANDARD = "gnu++0x"; - CLANG_CXX_LIBRARY = "libc++"; - CLANG_ENABLE_MODULES = YES; - CLANG_ENABLE_OBJC_ARC = YES; - CLANG_WARN_BOOL_CONVERSION = YES; - CLANG_WARN_CONSTANT_CONVERSION = YES; - CLANG_WARN_DIRECT_OBJC_ISA_USAGE = YES_ERROR; - CLANG_WARN_EMPTY_BODY = YES; - CLANG_WARN_ENUM_CONVERSION = YES; - CLANG_WARN_INFINITE_RECURSION = YES; - CLANG_WARN_INT_CONVERSION = YES; - CLANG_WARN_OBJC_ROOT_CLASS = YES_ERROR; - CLANG_WARN_SUSPICIOUS_MOVE = YES; - CLANG_WARN_UNREACHABLE_CODE = YES; - CLANG_WARN__DUPLICATE_METHOD_MATCH = YES; - "CODE_SIGN_IDENTITY[sdk=iphoneos*]" = "iPhone Developer"; - COPY_PHASE_STRIP = NO; - DEBUG_INFORMATION_FORMAT = "dwarf-with-dsym"; - ENABLE_NS_ASSERTIONS = NO; - ENABLE_STRICT_OBJC_MSGSEND = YES; - GCC_C_LANGUAGE_STANDARD = gnu99; - GCC_NO_COMMON_BLOCKS = YES; - GCC_WARN_64_TO_32_BIT_CONVERSION = YES; - GCC_WARN_ABOUT_RETURN_TYPE = YES_ERROR; - GCC_WARN_UNDECLARED_SELECTOR = YES; - GCC_WARN_UNINITIALIZED_AUTOS = YES_AGGRESSIVE; - GCC_WARN_UNUSED_FUNCTION = YES; - GCC_WARN_UNUSED_VARIABLE = YES; - HEADER_SEARCH_PATHS = ( - "$(inherited)", - ); - IPHONEOS_DEPLOYMENT_TARGET = 8.0; - MTL_ENABLE_DEBUG_INFO = NO; - SDKROOT = iphoneos; - TARGETED_DEVICE_FAMILY = "1,2"; - VALIDATE_PRODUCT = YES; - }; - name = Release; - }; -/* End XCBuildConfiguration section */ - -/* Begin XCConfigurationList section */ - 1C564C351ED3A92E00087306 /* Build configuration list for PBXNativeTarget "tflite_camera_example" */ = { - isa = XCConfigurationList; - buildConfigurations = ( - 1C564C361ED3A92E00087306 /* Debug */, - 1C564C371ED3A92E00087306 /* Release */, - ); - defaultConfigurationIsVisible = 0; - defaultConfigurationName = Release; - }; - 591157961CF4011C00C31E3A /* Build configuration list for PBXProject "tflite_camera_example" */ = { - isa = XCConfigurationList; - buildConfigurations = ( - 591157B01CF4011D00C31E3A /* Debug */, - 591157B11CF4011D00C31E3A /* Release */, - ); - defaultConfigurationIsVisible = 0; - defaultConfigurationName = Release; - }; -/* End XCConfigurationList section */ - }; - rootObject = 591157931CF4011C00C31E3A /* Project object */; -} diff --git a/tensorflow/contrib/lite/examples/ios/download_models.sh b/tensorflow/contrib/lite/examples/ios/download_models.sh deleted file mode 100755 index ccd163758c5830dc9367e023dcb3a604e07ca5db..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/examples/ios/download_models.sh +++ /dev/null @@ -1,57 +0,0 @@ -#!/bin/bash -# 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. -# ============================================================================== - -set -ex - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -MODELS_URL="https://storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_v1_1.0_224_ios_lite_float_2017_11_08.zip" -QUANTIZED_MODELS_URL="https://storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_v1_224_android_quant_2017_11_08.zip" -DOWNLOADS_DIR=$(mktemp -d) - -cd $SCRIPT_DIR - -download_and_extract() { - local usage="Usage: download_and_extract URL DIR" - local url="${1:?${usage}}" - local dir="${2:?${usage}}" - echo "downloading ${url}" >&2 - mkdir -p "${dir}" - tempdir=$(mktemp -d) - tempdir2=$(mktemp -d) - - curl -L ${url} > ${tempdir}/zipped.zip - unzip ${tempdir}/zipped.zip -d ${tempdir2} - - # If the zip file contains nested directories, extract the files from the - # inner directory. - if ls ${tempdir2}/*/* 1> /dev/null 2>&1; then - # unzip has no strip components, so unzip to a temp dir, and move the - # files we want from the tempdir to destination. - cp -R ${tempdir2}/*/* ${dir}/ - else - cp -R ${tempdir2}/* ${dir}/ - fi - rm -rf ${tempdir2} ${tempdir} -} - -download_and_extract "${MODELS_URL}" "${DOWNLOADS_DIR}/models" -download_and_extract "${QUANTIZED_MODELS_URL}" "${DOWNLOADS_DIR}/quantized_models" - -file ${DOWNLOADS_DIR}/models - -cp ${DOWNLOADS_DIR}/models/models/* simple/data/ -cp ${DOWNLOADS_DIR}/quantized_models/* camera/data/ - diff --git a/tensorflow/contrib/lite/examples/ios/simple/Podfile b/tensorflow/contrib/lite/examples/ios/simple/Podfile deleted file mode 100644 index ddb77088d9f16fb55e8060a91504ebc44dd0b73e..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/examples/ios/simple/Podfile +++ /dev/null @@ -1,5 +0,0 @@ -platform :ios, '8.0' -inhibit_all_warnings! - -target 'tflite_simple_example' - pod 'TensorFlowLite', '1.10.1' diff --git a/tensorflow/contrib/lite/examples/label_image/BUILD b/tensorflow/contrib/lite/examples/label_image/BUILD deleted file mode 100644 index fc55a78019b4a12b24231034a7e4b912869389f2..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/examples/label_image/BUILD +++ /dev/null @@ -1,71 +0,0 @@ -# Description: -# TensorFlow Lite Example Label Image. - -package(default_visibility = ["//visibility:public"]) - -licenses(["notice"]) # Apache 2.0 - -load("//tensorflow:tensorflow.bzl", "tf_cc_binary") -load("//tensorflow/contrib/lite:build_def.bzl", "tflite_linkopts") - -exports_files(glob([ - "testdata/*.bmp", -])) - -tf_cc_binary( - name = "label_image", - srcs = [ - "get_top_n.h", - "get_top_n_impl.h", - "label_image.cc", - ], - linkopts = tflite_linkopts() + select({ - "//tensorflow:android": [ - "-pie", # Android 5.0 and later supports only PIE - "-lm", # some builtin ops, e.g., tanh, need -lm - ], - "//conditions:default": [], - }), - deps = [ - ":bitmap_helpers", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite:string_util", - "//tensorflow/contrib/lite/kernels:builtin_ops", - ], -) - -cc_library( - name = "bitmap_helpers", - srcs = ["bitmap_helpers.cc"], - hdrs = [ - "bitmap_helpers.h", - "bitmap_helpers_impl.h", - "label_image.h", - ], - deps = [ - "//tensorflow/contrib/lite:builtin_op_data", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite:schema_fbs_version", - "//tensorflow/contrib/lite:string", - "//tensorflow/contrib/lite:string_util", - "//tensorflow/contrib/lite/kernels:builtin_ops", - "//tensorflow/contrib/lite/schema:schema_fbs", - ], -) - -cc_test( - name = "label_image_test", - srcs = [ - "get_top_n.h", - "get_top_n_impl.h", - "label_image_test.cc", - ], - data = [ - "testdata/grace_hopper.bmp", - ], - tags = ["no_oss"], - deps = [ - ":bitmap_helpers", - "@com_google_googletest//:gtest", - ], -) diff --git a/tensorflow/contrib/lite/examples/label_image/label_image.md b/tensorflow/contrib/lite/examples/label_image/label_image.md deleted file mode 100644 index 9ce32cf101897f2d41cd14a485aeb432344928a0..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/examples/label_image/label_image.md +++ /dev/null @@ -1,78 +0,0 @@ -label_image for TensorFlow Lite inspired by TensorFlow's label_image. - -To build label_image for Android, run $TENSORFLOW_ROOT/configure -and set Android NDK or configure NDK setting in -$TENSORFLOW_ROOT/WORKSPACE first. - -To build it for android ARMv8: -``` -> bazel build --config monolithic --cxxopt=-std=c++11 \ - --crosstool_top=//external:android/crosstool \ - --host_crosstool_top=@bazel_tools//tools/cpp:toolchain \ - --cpu=arm64-v8a \ - //tensorflow/contrib/lite/examples/label_image:label_image -``` -or -``` -> bazel build --config android_arm64 --config monolithic --cxxopt=-std=c++11 \ - //tensorflow/contrib/lite/examples/label_image:label_image -``` - -To build it for android arm-v7a: -``` -> bazel build --config monolithic --cxxopt=-std=c++11 \ - --crosstool_top=//external:android/crosstool \ - --host_crosstool_top=@bazel_tools//tools/cpp:toolchain \ - --cpu=armeabi-v7a \ - //tensorflow/contrib/lite/examples/label_image:label_image -``` -or -``` -> bazel build --config android_arm --config monolithic --cxxopt=-std=c++11 \ - //tensorflow/contrib/lite/examples/label_image:label_image -``` - -Build it for desktop machines (tested on Ubuntu and OS X) -``` -> bazel build --config opt --cxxopt=-std=c++11 //tensorflow/contrib/lite/examples/label_image:label_image -``` -To run it. Prepare `./mobilenet_quant_v1_224.tflite`, `./grace_hopper.bmp`, and `./labels.txt`. - -Run it: -``` -> ./label_image -Loaded model ./mobilenet_quant_v1_224.tflite -resolved reporter -invoked -average time: 100.986 ms -0.439216: 653 military uniform -0.372549: 458 bow tie -0.0705882: 466 bulletproof vest -0.0235294: 514 cornet -0.0196078: 835 suit -``` -Run `interpreter->Invoker()` 100 times: -``` -> ./label_image -c 100 -Loaded model ./mobilenet_quant_v1_224.tflite -resolved reporter -invoked -average time: 33.4694 ms -... -``` - -Run a floating point (`mobilenet_v1_1.0_224.tflite`) model, -``` -> ./label_image -f 1 -m mobilenet_v1_1.0_224.tflite -Loaded model mobilenet_v1_1.0_224.tflite -resolved reporter -invoked -average time: 263.493 ms -0.88615: 653 military uniform -0.0422316: 440 bearskin -0.0109948: 466 bulletproof vest -0.0105327: 401 academic gown -0.00947104: 723 ping-pong bal -``` - -See the source code for other command line options. diff --git a/tensorflow/contrib/lite/examples/minimal/BUILD b/tensorflow/contrib/lite/examples/minimal/BUILD deleted file mode 100644 index b403628d6c457ce3fb67eac3675fd7bb9187deab..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/examples/minimal/BUILD +++ /dev/null @@ -1,27 +0,0 @@ -# Description: -# TensorFlow Lite minimal example. - -package(default_visibility = ["//visibility:public"]) - -licenses(["notice"]) # Apache 2.0 - -load("//tensorflow:tensorflow.bzl", "tf_cc_binary") -load("//tensorflow/contrib/lite:build_def.bzl", "tflite_linkopts") - -tf_cc_binary( - name = "minimal", - srcs = [ - "minimal.cc", - ], - linkopts = tflite_linkopts() + select({ - "//tensorflow:android": [ - "-pie", # Android 5.0 and later supports only PIE - "-lm", # some builtin ops, e.g., tanh, need -lm - ], - "//conditions:default": [], - }), - deps = [ - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:builtin_ops", - ], -) diff --git a/tensorflow/contrib/lite/examples/python/BUILD b/tensorflow/contrib/lite/examples/python/BUILD deleted file mode 100644 index d337c3ddc43a23e50a5afdab93b16c0f61ccd538..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/examples/python/BUILD +++ /dev/null @@ -1,13 +0,0 @@ -licenses(["notice"]) # Apache 2.0 - -package(default_visibility = ["//tensorflow:internal"]) - -py_binary( - name = "label_image", - srcs = ["label_image.py"], - main = "label_image.py", - srcs_version = "PY2AND3", - deps = [ - "//tensorflow/contrib/lite/python:lite", - ], -) diff --git a/tensorflow/contrib/lite/examples/python/label_image.md b/tensorflow/contrib/lite/examples/python/label_image.md deleted file mode 100644 index e81192a96c142f2b3e7e85d160166fdd37ccdc53..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/examples/python/label_image.md +++ /dev/null @@ -1,50 +0,0 @@ - -With model, input image (grace_hopper.bmp), and labels file (labels.txt) -in /tmp. - -The example input image and labels file are from TensorFlow repo and -MobileNet V1 model files. - -``` -curl https://raw.githubusercontent.com/tensorflow/tensorflow/master/tensorflow/contrib/lite/examples/label_image/testdata/grace_hopper.bmp > /tmp/grace_hopper.bmp - -curl https://storage.googleapis.com/download.tensorflow.org/models/mobilenet_v1_1.0_224_frozen.tgz | tar xzv -C /tmp mobilenet_v1_1.0_224/labels.txt -mv /tmp/mobilenet_v1_1.0_224/labels.txt /tmp/ - -``` - -Run - -``` -curl http://download.tensorflow.org/models/mobilenet_v1_2018_02_22/mobilenet_v1_1.0_224_quant.tgz | tar xzv -C /tmp -bazel run --config opt //tensorflow/contrib/lite/examples/python:label_image -``` - -We can get results like - -``` -0.470588: military uniform -0.337255: Windsor tie -0.047059: bow tie -0.031373: mortarboard -0.019608: suit -``` - -Run - -``` -curl http://download.tensorflow.org/models/mobilenet_v1_2018_02_22/mobilenet_v1_1.0_224.tgz | tar xzv -C /tmp -bazel run --config opt //tensorflow/contrib/lite/examples/python:label_image \ --- --model_file /tmp/mobilenet_v1_1.0_224.tflite -``` - -We can get results like -``` -0.728693: military uniform -0.116163: Windsor tie -0.035517: bow tie -0.014874: mortarboard -0.011758: bolo tie -``` - -Check [models](../../g3doc/models.md) for models hosted by Google. diff --git a/tensorflow/contrib/lite/experimental/c/BUILD b/tensorflow/contrib/lite/experimental/c/BUILD deleted file mode 100644 index 52e71619def71a0c2130539afe8e7d00e7a24894..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/experimental/c/BUILD +++ /dev/null @@ -1,111 +0,0 @@ -package(default_visibility = ["//visibility:private"]) - -package_group( - name = "experimental", - packages = [ - "//tensorflow/contrib/lite/experimental/...", - ], -) - -licenses(["notice"]) # Apache 2.0 - -load( - "//tensorflow/contrib/lite:build_def.bzl", - "tflite_cc_shared_object", - "tflite_copts", - "tflite_jni_binary", -) - -tflite_cc_shared_object( - name = "libtensorflowlite_c.so", - linkopts = select({ - "//tensorflow:darwin": [ - "-Wl,-exported_symbols_list", # This line must be directly followed by the exported_symbols.lds file - "$(location //tensorflow/contrib/lite/experimental/c:exported_symbols.lds)", - "-Wl,-install_name,@rpath/libtensorflowlite_c.so", - ], - "//tensorflow:windows": [], - "//conditions:default": [ - "-z defs", - "-Wl,--version-script", # This line must be directly followed by the version_script.lds file - "$(location //tensorflow/contrib/lite/experimental/c:version_script.lds)", - ], - }), - deps = [ - ":c_api", - ":c_api_experimental", - ":exported_symbols.lds", - ":version_script.lds", - ], -) - -cc_library( - name = "c_api_internal", - srcs = ["c_api.h"], - hdrs = ["c_api_internal.h"], - copts = tflite_copts(), - visibility = [ - "//tensorflow/contrib/lite/experimental/c:__subpackages__", - ], - deps = [ - "//tensorflow/contrib/lite:context", - "//tensorflow/contrib/lite:framework", - ], -) - -cc_library( - name = "c_api", - srcs = ["c_api.cc"], - hdrs = ["c_api.h"], - copts = tflite_copts(), - visibility = [ - ":experimental", - ], - deps = [ - ":c_api_internal", - "//tensorflow/contrib/lite:context", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite:schema_fbs_version", - "//tensorflow/contrib/lite/kernels:builtin_ops", - ], -) - -cc_library( - name = "c_api_experimental", - srcs = ["c_api_experimental.cc"], - hdrs = ["c_api_experimental.h"], - copts = tflite_copts(), - deps = [ - ":c_api", - ":c_api_internal", - "//tensorflow/contrib/lite:kernel_api", - ], -) - -cc_test( - name = "c_api_test", - size = "small", - srcs = ["c_api_test.cc"], - data = ["//tensorflow/contrib/lite:testdata/add.bin"], - deps = [ - ":c_api", - "//tensorflow/contrib/lite:context", - "//tensorflow/contrib/lite:kernel_api", - "//tensorflow/contrib/lite/testing:util", - "@com_google_googletest//:gtest", - ], -) - -cc_test( - name = "c_api_experimental_test", - size = "small", - srcs = ["c_api_experimental_test.cc"], - data = ["//tensorflow/contrib/lite:testdata/add.bin"], - deps = [ - ":c_api", - ":c_api_experimental", - "//tensorflow/contrib/lite:kernel_api", - "//tensorflow/contrib/lite/testing:util", - "@com_google_googletest//:gtest", - ], -) diff --git a/tensorflow/contrib/lite/experimental/c/c_api_experimental.cc b/tensorflow/contrib/lite/experimental/c/c_api_experimental.cc deleted file mode 100644 index 29f8701f53407dc47adfaca8c85c86210e4cb09a..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/experimental/c/c_api_experimental.cc +++ /dev/null @@ -1,46 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/contrib/lite/experimental/c/c_api_experimental.h" - -#include "tensorflow/contrib/lite/experimental/c/c_api_internal.h" - -#ifdef __cplusplus -extern "C" { -#endif // __cplusplus - -TFL_Status TFL_InterpreterResetVariableTensors(TFL_Interpreter* interpreter) { - return interpreter->impl->ResetVariableTensors(); -} - -void TFL_InterpreterOptionsAddBuiltinOp(TFL_InterpreterOptions* options, - TFL_BuiltinOperator op, - const TFL_Registration* registration, - int32_t min_version, - int32_t max_version) { - options->op_resolver.AddBuiltin(static_cast(op), - registration, min_version, max_version); -} - -void TFL_InterpreterOptionsAddCustomOp(TFL_InterpreterOptions* options, - const char* name, - const TFL_Registration* registration, - int min_version, int max_version) { - options->op_resolver.AddCustom(name, registration, min_version, max_version); -} - -#ifdef __cplusplus -} // extern "C" -#endif // __cplusplus diff --git a/tensorflow/contrib/lite/experimental/c/c_api_experimental.h b/tensorflow/contrib/lite/experimental/c/c_api_experimental.h deleted file mode 100644 index fca5d92f77caff987f6a70c3a8fd03849bce1165..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/experimental/c/c_api_experimental.h +++ /dev/null @@ -1,57 +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_LITE_EXPERIMENTAL_C_C_API_EXPERIMENTAL_H_ -#define TENSORFLOW_CONTRIB_LITE_EXPERIMENTAL_C_C_API_EXPERIMENTAL_H_ - -#include "tensorflow/contrib/lite/builtin_ops.h" -#include "tensorflow/contrib/lite/experimental/c/c_api.h" - -#ifdef __cplusplus -extern "C" { -#endif // __cplusplus - -typedef TfLiteBuiltinOperator TFL_BuiltinOperator; - -// Resets all variable tensors to zero. -TFL_CAPI_EXPORT extern TFL_Status TFL_InterpreterResetVariableTensors( - TFL_Interpreter* interpreter); - -// Adds an op registration for a builtin operator. -// -// NOTE: The interpreter will make a copy of `registration` internally, so the -// caller should ensure that its contents (function pointers, etc...) remain -// valid for the duration of the interpreter's lifetime. A common practice is -// making the provided TFL_Registration instance static. -void TFL_InterpreterOptionsAddBuiltinOp(TFL_InterpreterOptions* options, - TFL_BuiltinOperator op, - const TFL_Registration* registration, - int min_version, int max_version); - -// Adds an op registration for a custom operator. -// -// NOTE: The interpreter will make a copy of `registration` internally, so the -// caller should ensure that its contents (function pointers, etc...) remain -// valid for the duration of the interpreter's lifetime. A common practice is -// making the provided TFL_Registration instance static. -void TFL_InterpreterOptionsAddCustomOp(TFL_InterpreterOptions* options, - const char* name, - const TFL_Registration* registration, - int min_version, int max_version); - -#ifdef __cplusplus -} // extern "C" -#endif // __cplusplus - -#endif // TENSORFLOW_CONTRIB_LITE_EXPERIMENTAL_C_C_API_EXPERIMENTAL_H_ diff --git a/tensorflow/contrib/lite/experimental/c/c_api_experimental_test.cc b/tensorflow/contrib/lite/experimental/c/c_api_experimental_test.cc deleted file mode 100644 index 1b1bedb75470638d4b3cfac92819e18b8fe6e65a..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/experimental/c/c_api_experimental_test.cc +++ /dev/null @@ -1,61 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/contrib/lite/experimental/c/c_api_experimental.h" - -#include -#include "tensorflow/contrib/lite/builtin_ops.h" -#include "tensorflow/contrib/lite/experimental/c/c_api.h" -#include "tensorflow/contrib/lite/testing/util.h" - -namespace { - -TfLiteRegistration* GetDummyRegistration() { - static TfLiteRegistration registration = { - .init = nullptr, - .free = nullptr, - .prepare = nullptr, - .invoke = [](TfLiteContext*, TfLiteNode*) { return kTfLiteOk; }, - }; - return ®istration; -} - -TEST(CApiExperimentalSimple, Smoke) { - TFL_Model* model = TFL_NewModelFromFile( - "tensorflow/contrib/lite/testdata/add.bin"); - ASSERT_NE(model, nullptr); - - TFL_InterpreterOptions* options = TFL_NewInterpreterOptions(); - TFL_InterpreterOptionsAddBuiltinOp(options, kTfLiteBuiltinAdd, - GetDummyRegistration(), 1, 1); - - TFL_Interpreter* interpreter = TFL_NewInterpreter(model, options); - ASSERT_NE(interpreter, nullptr); - ASSERT_EQ(TFL_InterpreterAllocateTensors(interpreter), kTfLiteOk); - EXPECT_EQ(TFL_InterpreterResetVariableTensors(interpreter), kTfLiteOk); - EXPECT_EQ(TFL_InterpreterInvoke(interpreter), kTfLiteOk); - - TFL_DeleteInterpreter(interpreter); - TFL_DeleteInterpreterOptions(options); - TFL_DeleteModel(model); -} - -} // namespace - -int main(int argc, char** argv) { - ::tflite::LogToStderr(); - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/contrib/lite/experimental/c/c_api_internal.h b/tensorflow/contrib/lite/experimental/c/c_api_internal.h deleted file mode 100644 index da3af3cad4c54865cfe778b79538e5800c284985..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/experimental/c/c_api_internal.h +++ /dev/null @@ -1,61 +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_LITE_EXPERIMENTAL_C_C_API_INTERNAL_H_ -#define TENSORFLOW_CONTRIB_LITE_EXPERIMENTAL_C_C_API_INTERNAL_H_ - -#include "tensorflow/contrib/lite/experimental/c/c_api.h" - -#include "tensorflow/contrib/lite/interpreter.h" -#include "tensorflow/contrib/lite/model.h" -#include "tensorflow/contrib/lite/op_resolver.h" - -// Internal structures used by the C API. These are likely to change and should -// not be depended on. -// -// NOTE: This header does not follow C conventions and does not define a C API. -// It is effectively an (internal) implementation detail of the C API. - -struct TFL_Model { - // Sharing is safe as FlatBufferModel is const. - std::shared_ptr impl; -}; - -struct TFL_InterpreterOptions { - enum { - kDefaultNumThreads = -1, - }; - int num_threads = kDefaultNumThreads; - - tflite::MutableOpResolver op_resolver; - - void (*error_reporter)(void* user_data, const char* format, - va_list args) = nullptr; - void* error_reporter_user_data = nullptr; -}; - -struct TFL_Interpreter { - // Taking a reference to the (const) model data avoids lifetime-related issues - // and complexity with the TFL_Model's existence. - std::shared_ptr model; - - // The interpreter does not take ownership of the provided ErrorReporter - // instance, so we ensure its validity here. Note that the interpreter may use - // the reporter in its destructor, so it should be declared first. - std::unique_ptr optional_error_reporter; - - std::unique_ptr impl; -}; - -#endif // TENSORFLOW_CONTRIB_LITE_EXPERIMENTAL_C_C_API_INTERNAL_H_ diff --git a/tensorflow/contrib/lite/experimental/c/c_api_test.cc b/tensorflow/contrib/lite/experimental/c/c_api_test.cc deleted file mode 100644 index 48a3714ec345a6f4bc4be8ebe937471a91c60218..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/experimental/c/c_api_test.cc +++ /dev/null @@ -1,125 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include - -#include "tensorflow/contrib/lite/experimental/c/c_api.h" - -#include -#include "tensorflow/contrib/lite/context.h" -#include "tensorflow/contrib/lite/testing/util.h" - -namespace { - -TEST(CApiSimple, Smoke) { - TFL_Model* model = TFL_NewModelFromFile( - "tensorflow/contrib/lite/testdata/add.bin"); - ASSERT_NE(model, nullptr); - - TFL_InterpreterOptions* options = TFL_NewInterpreterOptions(); - ASSERT_NE(options, nullptr); - TFL_InterpreterOptionsSetNumThreads(options, 2); - - TFL_Interpreter* interpreter = TFL_NewInterpreter(model, options); - ASSERT_NE(interpreter, nullptr); - - // The options/model can be deleted immediately after interpreter creation. - TFL_DeleteInterpreterOptions(options); - TFL_DeleteModel(model); - - ASSERT_EQ(TFL_InterpreterAllocateTensors(interpreter), kTfLiteOk); - ASSERT_EQ(TFL_InterpreterGetInputTensorCount(interpreter), 1); - ASSERT_EQ(TFL_InterpreterGetOutputTensorCount(interpreter), 1); - - std::array input_dims = {2}; - ASSERT_EQ(TFL_InterpreterResizeInputTensor(interpreter, 0, input_dims.data(), - input_dims.size()), - kTfLiteOk); - ASSERT_EQ(TFL_InterpreterAllocateTensors(interpreter), kTfLiteOk); - - TFL_Tensor* input_tensor = TFL_InterpreterGetInputTensor(interpreter, 0); - ASSERT_NE(input_tensor, nullptr); - EXPECT_EQ(TFL_TensorType(input_tensor), kTfLiteFloat32); - EXPECT_EQ(TFL_TensorNumDims(input_tensor), 1); - EXPECT_EQ(TFL_TensorDim(input_tensor, 0), 2); - EXPECT_EQ(TFL_TensorByteSize(input_tensor), sizeof(float) * 2); - EXPECT_NE(TFL_TensorData(input_tensor), nullptr); - EXPECT_STREQ(TFL_TensorName(input_tensor), "input"); - - std::array input = {1.f, 3.f}; - ASSERT_EQ(TFL_TensorCopyFromBuffer(input_tensor, input.data(), - input.size() * sizeof(float)), - kTfLiteOk); - - ASSERT_EQ(TFL_InterpreterInvoke(interpreter), kTfLiteOk); - - const TFL_Tensor* output_tensor = - TFL_InterpreterGetOutputTensor(interpreter, 0); - ASSERT_NE(output_tensor, nullptr); - EXPECT_EQ(TFL_TensorType(output_tensor), kTfLiteFloat32); - EXPECT_EQ(TFL_TensorNumDims(output_tensor), 1); - EXPECT_EQ(TFL_TensorDim(output_tensor, 0), 2); - EXPECT_EQ(TFL_TensorByteSize(output_tensor), sizeof(float) * 2); - EXPECT_NE(TFL_TensorData(output_tensor), nullptr); - EXPECT_STREQ(TFL_TensorName(output_tensor), "output"); - - std::array output; - ASSERT_EQ(TFL_TensorCopyToBuffer(output_tensor, output.data(), - output.size() * sizeof(float)), - kTfLiteOk); - EXPECT_EQ(output[0], 3.f); - EXPECT_EQ(output[1], 9.f); - - TFL_DeleteInterpreter(interpreter); -} - -TEST(CApiSimple, ErrorReporter) { - TFL_Model* model = TFL_NewModelFromFile( - "tensorflow/contrib/lite/testdata/add.bin"); - TFL_InterpreterOptions* options = TFL_NewInterpreterOptions(); - - // Install a custom error reporter into the interpreter by way of options. - tflite::TestErrorReporter reporter; - TFL_InterpreterOptionsSetErrorReporter( - options, - [](void* user_data, const char* format, va_list args) { - reinterpret_cast(user_data)->Report(format, - args); - }, - &reporter); - TFL_Interpreter* interpreter = TFL_NewInterpreter(model, options); - - // The options/model can be deleted immediately after interpreter creation. - TFL_DeleteInterpreterOptions(options); - TFL_DeleteModel(model); - - // Invoke the interpreter before tensor allocation. - EXPECT_EQ(TFL_InterpreterInvoke(interpreter), kTfLiteError); - - // The error should propagate to the custom error reporter. - EXPECT_EQ(reporter.error_messages(), - "Invoke called on model that is not ready."); - EXPECT_EQ(reporter.num_calls(), 1); - - TFL_DeleteInterpreter(interpreter); -} - -} // namespace - -int main(int argc, char** argv) { - ::tflite::LogToStderr(); - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/contrib/lite/experimental/examples/lstm/BUILD b/tensorflow/contrib/lite/experimental/examples/lstm/BUILD deleted file mode 100644 index 2125f218ca877f94ec9f4d98928b6a1c8f2576eb..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/experimental/examples/lstm/BUILD +++ /dev/null @@ -1,40 +0,0 @@ -licenses(["notice"]) # Apache 2.0 - -package(default_visibility = ["//tensorflow:internal"]) - -load("//tensorflow:tensorflow.bzl", "py_test") - -py_library( - name = "tflite_lstm", - srcs = ["tflite_lstm.py"], - srcs_version = "PY2AND3", - visibility = ["//visibility:public"], - deps = [ - "//tensorflow:tensorflow_py", - "//tensorflow/contrib/lite/python:lite", - "//tensorflow/python:framework", - "@six_archive//:six", - ], -) - -py_test( - name = "unidirectional_sequence_lstm_test", - size = "large", - srcs = ["unidirectional_sequence_lstm_test.py"], - srcs_version = "PY2AND3", - tags = [ - "no_oss", - "no_pip", - ], - deps = [ - ":tflite_lstm", - "//tensorflow:tensorflow_py", - "//tensorflow/contrib/lite/python:lite", - "//tensorflow/examples/tutorials/mnist:input_data", - "//tensorflow/python:framework_test_lib", - "//tensorflow/python:platform", - "//tensorflow/python/tools:optimize_for_inference", - "//third_party/py/numpy", - "@six_archive//:six", - ], -) diff --git a/tensorflow/contrib/lite/experimental/examples/unity/TensorFlowLitePlugin/README.md b/tensorflow/contrib/lite/experimental/examples/unity/TensorFlowLitePlugin/README.md deleted file mode 100644 index f480c49cd050de2192e9673f72c9e4d5c3c6ceff..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/experimental/examples/unity/TensorFlowLitePlugin/README.md +++ /dev/null @@ -1,29 +0,0 @@ -# TF Lite Experimental Unity Plugin - -This directory contains an experimental sample Unity (2017) Plugin, based on -the experimental TF Lite C API. The sample demonstrates running inference within -Unity by way of a C# `Interpreter` wrapper. - -Note that the native TF Lite plugin(s) *must* be built before using the Unity -Plugin, and placed in Assets/TensorFlowLite/SDK/Plugins/. For the editor (note -that this has only been tested on Linux; the syntax may differ on Mac/Windows): - -```sh -bazel build -c opt --cxxopt=--std=c++11 \ - //tensorflow/contrib/lite/experimental/c:libtensorflowlite_c.so -``` - -and for Android: - -```sh -bazel build -c opt --cxxopt=--std=c++11 \ - --crosstool_top=//external:android/crosstool \ - --host_crosstool_top=@bazel_tools//tools/cpp:toolchain \ - --cpu=armeabi-v7a \ - //tensorflow/contrib/lite/experimental/c:libtensorflowlite_c.so -``` - -If you encounter issues with native plugin discovery on Mac ("Darwin") -platforms, try renaming `libtensorflowlite_c.so` to `tensorflowlite_c.bundle`. -Similarly, on Windows you'll likely need to rename `libtensorflowlite_c.so` to -`tensorflowlite_c.dll`. diff --git a/tensorflow/contrib/lite/experimental/kernels/BUILD b/tensorflow/contrib/lite/experimental/kernels/BUILD deleted file mode 100644 index 4786cc62f93dc0a27efa02c2b436820867ab95f5..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/experimental/kernels/BUILD +++ /dev/null @@ -1,85 +0,0 @@ -package(default_visibility = [ - "//visibility:public", -]) - -licenses(["notice"]) # Apache 2.0 - -load("//tensorflow/contrib/lite:build_def.bzl", "tflite_copts") -load("//tensorflow:tensorflow.bzl", "tf_cc_test") - -# ctc support classes imported directly from TensorFlow. -cc_library( - name = "ctc_utils", - hdrs = [ - "ctc_beam_entry.h", - "ctc_beam_scorer.h", - "ctc_beam_search.h", - "ctc_decoder.h", - "ctc_loss_util.h", - ], - deps = [ - ":top_n", - "//tensorflow/contrib/lite/kernels/internal:types", - "//third_party/eigen3", - ], -) - -# top_n support classes imported directly from TensorFlow. -cc_library( - name = "top_n", - hdrs = [ - "top_n.h", - ], - deps = [ - "//tensorflow/contrib/lite/kernels/internal:types", - ], -) - -cc_library( - name = "experimental_ops", - srcs = [ - "ctc_beam_search_decoder.cc", - ], - # Suppress warnings that are introduced by Eigen Tensor. - copts = tflite_copts() + [ - "-Wno-error=reorder", - ] + select({ - "//tensorflow:ios": ["-Wno-error=invalid-partial-specialization"], - "//conditions:default": [ - ], - }), - deps = [ - ":ctc_utils", - "//tensorflow/contrib/lite:builtin_op_data", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite:string_util", - "//tensorflow/contrib/lite/c:c_api_internal", - "//tensorflow/contrib/lite/kernels:builtin_ops", - "//tensorflow/contrib/lite/kernels:gemm_support", - "//tensorflow/contrib/lite/kernels:kernel_util", - "//tensorflow/contrib/lite/kernels:op_macros", - "//tensorflow/contrib/lite/kernels/internal:kernel_utils", - "//tensorflow/contrib/lite/kernels/internal:optimized", - "//tensorflow/contrib/lite/kernels/internal:optimized_base", - "//tensorflow/contrib/lite/kernels/internal:quantization_util", - "//tensorflow/contrib/lite/kernels/internal:reference_base", - "//tensorflow/contrib/lite/kernels/internal:tensor", - "//tensorflow/contrib/lite/kernels/internal:tensor_utils", - "@flatbuffers", - ], -) - -tf_cc_test( - name = "ctc_beam_search_decoder_test", - size = "small", - srcs = ["ctc_beam_search_decoder_test.cc"], - tags = ["tflite_not_portable_ios"], - deps = [ - ":experimental_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:builtin_ops", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - "@flatbuffers", - ], -) diff --git a/tensorflow/contrib/lite/experimental/micro/BUILD b/tensorflow/contrib/lite/experimental/micro/BUILD deleted file mode 100644 index df1036bc8b9cc84f4b63ae2a771e3aa8f8989060..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/experimental/micro/BUILD +++ /dev/null @@ -1,76 +0,0 @@ -package( - default_visibility = ["//visibility:public"], -) - -licenses(["notice"]) # Apache 2.0 - -load( - "//tensorflow/contrib/lite/experimental/micro/testing:micro_test.bzl", - "tflite_micro_cc_test", -) - -cc_library( - name = "micro_framework", - srcs = [ - "micro_error_reporter.cc", - "micro_interpreter.cc", - "micro_mutable_op_resolver.cc", - "simple_tensor_allocator.cc", - ], - hdrs = [ - "compatibility.h", - "micro_error_reporter.h", - "micro_interpreter.h", - "micro_mutable_op_resolver.h", - "simple_tensor_allocator.h", - ], - deps = [ - "//tensorflow/contrib/lite:schema_fbs_version", - "//tensorflow/contrib/lite/c:c_api_internal", - "//tensorflow/contrib/lite/core/api", - "//tensorflow/contrib/lite/schema:schema_fbs", - ], -) - -tflite_micro_cc_test( - name = "micro_error_reporter_test", - srcs = [ - "micro_error_reporter_test.cc", - ], - deps = [ - ":micro_framework", - ], -) - -tflite_micro_cc_test( - name = "micro_mutable_op_resolver_test", - srcs = [ - "micro_mutable_op_resolver_test.cc", - ], - deps = [ - ":micro_framework", - "//tensorflow/contrib/lite/experimental/micro/testing:micro_test", - ], -) - -tflite_micro_cc_test( - name = "micro_interpreter_test", - srcs = [ - "micro_interpreter_test.cc", - ], - deps = [ - ":micro_framework", - "//tensorflow/contrib/lite/experimental/micro/testing:micro_test", - ], -) - -tflite_micro_cc_test( - name = "simple_tensor_allocator_test", - srcs = [ - "simple_tensor_allocator_test.cc", - ], - deps = [ - ":micro_framework", - "//tensorflow/contrib/lite/experimental/micro/testing:micro_test", - ], -) diff --git a/tensorflow/contrib/lite/experimental/micro/README.md b/tensorflow/contrib/lite/experimental/micro/README.md deleted file mode 100644 index e03703f4967ba969df0d512e188f0ca7b0ac4d9f..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/experimental/micro/README.md +++ /dev/null @@ -1,128 +0,0 @@ -# TensorFlow Lite for Microcontrollers - -This an experimental port of TensorFlow Lite aimed at micro controllers and other devices with only kilobytes of memory. It doesn't require any operating system support, any standard C or C++ libraries, or dynamic memory allocation, so it's designed to be portable even to 'bare metal' systems. The core runtime fits in 16KB on a Cortex M3, and with enough operators to run a speech keyword detection model, takes up a total of 22KB. - -The design goals are for the framework to be: - -- **Readable**: We want embedded software engineers to be able to understand what's required to run ML inference without having to study research papers. We've tried to keep the code base small, modular, and have reference implementations of all operations to help with this. - -- **Easy to modify**: We know that there are a lot of different platforms and requirements in the embedded world, and we don't expect to cover all of them in one framework. Instead, we're hoping that it can be a good starting point for developers to build on top of to meet their own needs. For example, we tried to make it easy to replace the implementations of key computational operators that are often crucial for performance, without having to touch the data flow and other runtime code. We want it to make more sense to use our workflow to handle things like model import and less-important operations, and customize the parts that matter, rather than having to reimplement everything in your own engine. - -- **Well-tested**: If you're modifying code, you need to know if your changes are correct. Having an easy way to test lets you develop much faster. To help there, we've written tests for all the components, and we've made sure that the tests can be run on almost any platform, with no dependencies apart from the ability to log text to a debug console somewhere. We also provide an easy way to run all the tests on-device as part of an automated test framework, and we use qemu/Renode emulation so that tests can be run even without physical devices present. - -- **Easy to integrate**: We want to be as open a system as possible, and use the best code available for each platform. To do that, we're going to rely on projects like [CMSIS-NN](https://www.keil.com/pack/doc/CMSIS/NN/html/index.html), [uTensor](https://github.com/uTensor/uTensor), and other vendor libraries to handle as much performance-critical code as possible. We know that there are an increasing number of options to accelerate neural networks on microcontrollers, so we're aiming to be a good host for deploying those hardware technologies too. - -- **Compatible**: We're using the same file schema, interpreter API, and kernel interface as regular TensorFlow Lite, so we leverage the large existing set of tools, documentation, and examples for the project. The biggest barrier to deploying ML models is getting them from a training environment into a form that's easy to run inference on, so we see reusing this rich ecosystem as being crucial to being easily usable. We also hope to integrate this experimental work back into the main codebase in the future. - -To meet those goals, we've made some tradeoffs: - -- **Simple C++**: To help with readability, our code is written in a modern version of C++, but we generally treat it as a "better C", rather relying on more complex features such as template meta-programming. As mentioned earlier, we avoid any use of dynamic memory allocation (new/delete) or the standard C/C++ libraries, so we believe this should still be fairly portable. It does mean that some older devices with C-only toolchains won't be supported, but we're hoping that the reference operator implementations (which are simple C-like functions) can still be useful in those cases. The interfaces are also designed to be C-only, so it should be possible to integrate the resulting library with pure C projects. - -- **Interpreted**: Code generation is a popular pattern for embedded code, because it gives standalone code that's easy to modify and step through, but we've chosen to go with an interpreted approach. In our internal microcontroller work we've found that using an extremely stripped-down interpreter with almost no dependencies gives us a lot of the same advantages, but is easier to maintain. For example, when new updates come out for the underlying library, you can just merge your local modifications in a single step, rather than having to regenerate new code and then patch in any changes you subsequently made. The coarse granularity of the interpreted primitives means that each operation call typically takes hundreds of thousands of instruction cycles at least, so we don't see noticeable performance gains from avoiding what's essentially a single switch statement at the interpreter level to call each operation. We're still working on improving the packaging though, for example we're considering having the ability to snapshot all the source files and headers used for a particular model, being able to compile the code and data together as a library, and then access it through a minimal set of C interface calls which hide the underlying complexity. - -- **Flatbuffers**: We represent our models using [the standard flatbuffer schema used by the rest of TensorFlow Lite](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/lite/schema/schema.fbs), with the difference that we always keep it in read-only program memory (typically flash) rather than relying on having a file system to read it from. This is a good fit because flatbuffer's serialized format is designed to be mapped into memory without requiring any extra memory allocations or modifications to access it. All of the functions to read model values work directly on the serialized bytes, and large sections of data like weights are directly accessible as sequential C-style arrays of their data type, with no strides or unpacking needed. We do get a lot of value from using flatbuffers, but there is a cost in complexity. The flat buffer library code is all inline [inside the main headers](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/lite/schema/schema_generated.h), but it isn't straightforward to inspect their implementations, and the model data structures aren't easy to comprehend from the debugger. The header for the schema itself also has to be periodically updated when new information is added to the file format, though we try to handle that transparently for most developers by checking in a pre-generated version. - -- **Code Duplication**: Some of the code in this prototype largely duplicates the logic in other parts of the TensorFlow Lite code base, for example the operator wrappers. We've tried to keep share as much as we can between the two interpreters, but there are some assumptions built into the original runtime that make this difficult. We'll be working on modularizing the main interpreter so that we can move to an entirely shared system. - -This initial preview release is designed to get early feedback, and is not intended to be a final product. It only includes enough operations to run a simple keyword recognition model, and the implementations are not optimized. We're hoping this will be a good way to get feedback and collaborate to improve the framework. - -## Getting Started - -Building requires a Linux or OS X machine. - - - Open a terminal - - Download the TensorFlow source with `git clone https://github.com/tensorflow` - - Enter the source root directory by running `cd tensorflow` - - Download the dependencies by running `tensorflow/contrib/lite/experimental/micro/tools/make/download_dependencies.sh`. This may take a few minutes - - Build and test the library with `make -f tensorflow/contrib/lite/experimental/micro/tools/make/Makefile test` - -You should see a series of compilation steps, followed by `~~~ALL TESTS -PASSED~~~` for the various tests of the code that it will run. If there's an -error, you should get an informative message from make about what went wrong. - -These tests are all built as simple binaries with few dependencies, so you can run them manually. For example, here's how to run the depthwise convolution test, and its output: - -``` -tensorflow/contrib/lite/experimental/micro/tools/make/gen/linux_x86_64/bin/tensorflow/contrib/lite/experimental/micro/kernels/depthwise_conv_test - -Testing SimpleTest -Testing SimpleTestQuantized -Testing SimpleTestRelu -Testing SimpleTestReluQuantized -4/4 tests passed -~ALL TESTS PASSED~~~ -``` - -Looking at the [depthwise_conv_test.cc](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/lite/experimental/micro/kernels/depthwise_conv_test.cc) code, you'll see a sequence that looks like this: - -``` -... -TF_LITE_MICRO_TESTS_BEGIN - -TF_LITE_MICRO_TEST(SimpleTest) { -... -} -... -TF_LITE_MICRO_TESTS_END -``` - -These macros work a lot like -[the Google test framework](https://github.com/google/googletest), but they -don't require any dependencies and just write results to stderr, rather than -aborting the program. If all the tests pass, then `~~~ALL TESTS PASSED~~~` is -output, and the test harness that runs the binary during the make process knows -that everything ran correctly. If there's an error, the lack of the expected -string lets the harness know that the test failed. - -So, why are we running tests in this complicated way? So far, we've been building binaries that run locally on the Mac OS or Linux machine you're building on, but this approach becomes important when we're targeting simple micro controller devices. - -## Building for the "Blue Pill" STM32F103 - -The goal of this library is to enable machine learning on resource-constrained micro controllers and DSPs, and as part of that we've targeted the ["Blue Pill" STM32F103-compatible development board](https://github.com/google/googletest) as a cheap and popular platform. It only has 20KB of RAM and 64KB of flash, so it's a good device to ensure we can run efficiently on small chips. - -It's fairly easy to [buy and wire up a physical board](https://github.com/google/stm32_bare_lib#wiring-up-your-blue-pill), but even if you don't have an actual device, the [Renode project](https://renode.io/) makes it easy to run a faithful emulation on your desktop machine. You'll need [Docker](https://www.docker.com/) installed, but once you have that set up, try running the following command: - -`make -f tensorflow/contrib/lite/experimental/micro/tools/make/Makefile TARGET=bluepill test` - -You should see a similar set of outputs as you did in the previous section, with the addition of some extra Docker logging messages. These are because we're using Docker to run the Renode micro controller emulation tool, and the tests themselves are being run on a simulated STM32F103 device. The communication channels between an embedded device and the host are quite limited, so the test harness looks at the output of the debug log to see if tests have passed, just as it did in the previous section. This makes it a very flexible way to run cross-platform tests, even when a platform has no operating system facilities, as long as it can output debugging text logs. - -To understand what's happening here, try running the same depthwise convolution test, but through the emulated device test harness, with the following command: - -``` -tensorflow/contrib/lite/experimental/micro/testing/test_bluepill_binary.sh \ -tensorflow/contrib/lite/experimental/micro/tools/make/gen/bluepill_cortex-m3/bin/tensorflow/contrib/lite/experimental/micro/kernels/depthwise_conv_test \ -'~~~ALL TESTS PASSED~~~' - -``` - -You should see output that looks something like this: - -``` -Sending build context to Docker daemon 21.5kB -Step 1/2 : FROM antmicro/renode:latest - ---> 1b670a243e8f -Step 2/2 : LABEL maintainer="Pete Warden " - ---> Using cache - ---> 3afcd410846d -Successfully built 3afcd410846d -Successfully tagged renode_bluepill:latest -LOGS: -... -03:27:32.4340 [INFO] machine-0: Machine started. -03:27:32.4790 [DEBUG] cpu.uartSemihosting: [+0.22s host +0s virt 0s virt from start] Testing SimpleTest -03:27:32.4812 [DEBUG] cpu.uartSemihosting: [+2.21ms host +0s virt 0s virt from start] Testing SimpleTestQuantized -03:27:32.4833 [DEBUG] cpu.uartSemihosting: [+2.14ms host +0s virt 0s virt from start] Testing SimpleTestRelu -03:27:32.4834 [DEBUG] cpu.uartSemihosting: [+0.18ms host +0s virt 0s virt from start] Testing SimpleTestReluQuantized -03:27:32.4838 [DEBUG] cpu.uartSemihosting: [+0.4ms host +0s virt 0s virt from start] 4/4 tests passed -03:27:32.4839 [DEBUG] cpu.uartSemihosting: [+41µs host +0s virt 0s virt from start] ~~~ALL TESTS PASSED~~~ -03:27:32.4839 [DEBUG] cpu.uartSemihosting: [+5µs host +0s virt 0s virt from start] -... -tensorflow/contrib/lite/experimental/micro/tools/make/gen/bluepill_cortex-m3/bin/tensorflow/contrib/lite/experimental/micro/kernels/depthwise_conv_test: PASS -``` - -There's a lot of output here, but you should be able to see that the same tests -that were covered when we ran locally on the development machine show up in the -debug logs here, along with the magic string `~~~ALL TESTS PASSED~~~`. This is -the exact same code as before, just compiled and run on the STM32F103 rather -than your desktop. We hope that the simplicity of this testing approach will -help make adding support for new platforms as easy as possible. diff --git a/tensorflow/contrib/lite/experimental/micro/compatibility.h b/tensorflow/contrib/lite/experimental/micro/compatibility.h deleted file mode 100644 index 4f0fd9f3120a5db74cdfb84e7b17a0f3656520bc..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/experimental/micro/compatibility.h +++ /dev/null @@ -1,32 +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_LITE_EXPERIMENTAL_MICRO_COMPATIBILITY_H_ -#define TENSORFLOW_CONTRIB_LITE_EXPERIMENTAL_MICRO_COMPATIBILITY_H_ - -// C++ will automatically create class-specific delete operators for virtual -// objects, which by default call the global delete function. For embedded -// applications we want to avoid this, and won't be calling new/delete on these -// objects, so we need to override the default implementation with one that does -// nothing to avoid linking in ::delete(). -// This macro needs to be included in all subclasses of a virtual base class in -// the private section. -#ifdef TF_LITE_STATIC_MEMORY -#define TF_LITE_REMOVE_VIRTUAL_DELETE \ - void operator delete(void* p) {} -#else -#define TF_LITE_REMOVE_VIRTUAL_DELETE -#endif - -#endif // TENSORFLOW_CONTRIB_LITE_EXPERIMENTAL_MICRO_COMPATIBILITY_H_ diff --git a/tensorflow/contrib/lite/experimental/micro/examples/micro_speech/BUILD b/tensorflow/contrib/lite/experimental/micro/examples/micro_speech/BUILD deleted file mode 100644 index 626f733540264c6fa13ab82557b822690b2d5b8f..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/experimental/micro/examples/micro_speech/BUILD +++ /dev/null @@ -1,35 +0,0 @@ -# Description: -# TensorFlow Lite microcontroller example. - -package(default_visibility = ["//visibility:public"]) - -licenses(["notice"]) # Apache 2.0 - -load( - "//tensorflow/contrib/lite/experimental/micro/testing:micro_test.bzl", - "tflite_micro_cc_test", -) - -tflite_micro_cc_test( - name = "micro_speech_test", - srcs = [ - "micro_speech_test.cc", - "no_features_data.cc", - "no_features_data.h", - "tiny_conv_model_data.cc", - "tiny_conv_model_data.h", - "yes_features_data.cc", - "yes_features_data.h", - ], - tags = [ - "nomsan", - ], - deps = [ - "//tensorflow/contrib/lite:schema_fbs_version", - "//tensorflow/contrib/lite/experimental/micro:micro_framework", - "//tensorflow/contrib/lite/experimental/micro/kernels:all_ops_resolver", - "//tensorflow/contrib/lite/experimental/micro/kernels:micro_ops", - "//tensorflow/contrib/lite/experimental/micro/testing:micro_test", - "//tensorflow/contrib/lite/schema:schema_fbs", - ], -) diff --git a/tensorflow/contrib/lite/experimental/micro/examples/micro_speech/README.md b/tensorflow/contrib/lite/experimental/micro/examples/micro_speech/README.md deleted file mode 100644 index 438a432356be5c3cc9bfd08de5bd4d6f797c7014..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/experimental/micro/examples/micro_speech/README.md +++ /dev/null @@ -1,103 +0,0 @@ -# Micro Speech Example - -This examples shows how you can use TensorFlow Lite to run a 20 kilobyte neural network model to recognize keywords in speech. It's designed to run on systems with very small amounts of memory such as microcontrollers and DSPs. The code itself also has a small footprint (for example around 22 kilobytes on a Cortex M3) and only uses about 10 kilobytes of RAM for working memory, so it's able to run on systems like an STM32F103 with only 20 kilobytes of total SRAM and 64 kilobytes of Flash. - -## Table of Contents - - * [Getting Started](#getting-started) - * [Getting Started on a Microcontroller](#getting-started-on-a-microcontroller) - * [Calculating the Input to the Neural Network](#calculating-the-input-to-the-neural-network) - * [Creating Your Own Model](#creating-your-own-model) - -## Getting Started - -To compile and test this example on a desktop Linux or MacOS machine, download [the TensorFlow source code](https://github.com/tensorflow/tensorflow), `cd` into the source directory from a terminal, and then retrieve the support libraries you need by running: - -``` -tensorflow/contrib/lite/experimental/micro/tools/make/download_dependencies.sh -``` - -This will take a few minutes, and downloads frameworks the code uses like [CMSIS](https://developer.arm.com/embedded/cmsis) and [flatbuffers](https://google.github.io/flatbuffers/). Once that process has finished, run: - -``` -make -f tensorflow/contrib/lite/experimental/micro/tools/make/Makefile test_micro_speech -``` - -You should see a series of files get compiled, followed by some logging output from a test, which should conclude with "~~~ALL TESTS PASSED~~~". If you see this, it means that a small program has been built and run that loads a trained TensorFlow model, runs some example inputs through it, and got the expected outputs. This particular test runs spectrograms generated from recordings of people saying "Yes" and "No", and checks that the network correctly identifies them. - -To understand how TensorFlow Lite does this, you can look at the `TestInvoke()` function in [micro_speech_test.cc](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite/experimental/micro/examples/micro_speech/micro_speech_test.cc). It's a fairly small amount of code, creating an interpreter, getting a handle to a model that's been compiled into the program, and then invoking the interpreter with the model and sample inputs. - -## Getting Started on a Microcontroller - -Once you have downloaded the dependencies and got the x86/Linux build working, you can try building a version for the STM32F103 'bluepill' device. The following command will build the test and then run it on an emulator, assuming you have Docker installed: - -``` -make -f tensorflow/contrib/lite/experimental/micro/tools/make/Makefile TARGET=bluepill test_micro_speech -``` - -If you have a real device [(see here for how to set one up)](https://github.com/google/stm32_bare_lib/tree/master/README.md) you can then convert the ELF file into a a `.bin` format executable to load onto it by running: - -``` -arm-none-eabi-objcopy \ -tensorflow/contrib/lite/experimental/micro/tools/make/gen/bluepill_cortex-m3/bin/micro_speech_test \ -tensorflow/contrib/lite/experimental/micro/tools/make/gen/bluepill_cortex-m3/bin/micro_speech_test.bin \ ---output binary -``` - -## Calculating the Input to the Neural Network - -The TensorFlow Lite model doesn't take in raw audio sample data. Instead it works with spectrograms, which are two dimensional arrays that are made up of slices of frequency information, each taken from a different time window. This test uses spectrograms that have been pre-calculated from one-second WAV files in the test data set. In a complete application these spectrograms would be calculated at runtime from microphone inputs, but the code for doing that is not yet included in this sample code. - -The recipe for creating the spectrogram data is that each frequency slice is created by running an FFT across a 30ms section of the audio sample data. The input samples are treated as being between -1 and +1 as real values (encoded as -32,768 and 32,767 in 16-bit signed integer samples). This results in an FFT with 256 entries. Every sequence of six entries is averaged together, giving a total of 43 frequency buckets in the final slice. The results are stored as unsigned eight-bit values, where 0 represents a real number of zero, and 255 represents 127.5 as a real number. Each adjacent frequency entry is stored in ascending memory order (frequency bucket 0 at data[0], bucket 1 at data [1], etc). The window for the frequency analysis is then moved forward by 20ms, and the process repeated, storing the results in the next memory row (for example bucket 0 in this moved window would be in data[43 + 0], etc). This process happens 49 times in total, producing a single channel image that is 43 pixels wide, and 49 rows high. Here's an illustration of the process: - -![spectrogram diagram](https://storage.googleapis.com/download.tensorflow.org/example_images/spectrogram_diagram.png) - - -The test data files have been generated by running the following commands: - -``` -bazel run tensorflow/examples/speech_commands:wav_to_features -- \ ---input_wav=${HOME}/speech_commands_test_set_v0.02/yes/f2e59fea_nohash_1.wav \ ---output_c_file=yes_features_data.cc \ ---window_stride=20 --preprocess=average --quantize=1 - -bazel run tensorflow/examples/speech_commands:wav_to_features -- \ ---input_wav=${HOME}/speech_commands_test_set_v0.02/no/f9643d42_nohash_4.wav \ ---output_c_file=no_features_data.cc \ ---window_stride=20 --preprocess=average --quantize=1 -``` - -## Creating Your Own Model - -The neural network model used in this example was built using the [TensorFlow speech commands tutorial](https://www.tensorflow.org/tutorials/sequences/audio_recognition). If you would like to create your own, you can start by training a model with this command: - -``` -bazel run -c opt --copt=-mavx2 --copt=-mfma \ -tensorflow/examples/speech_commands:train -- \ ---model_architecture=tiny_conv --window_stride=20 --preprocess=average \ ---wanted_words="yes,no" --silence_percentage=25 --unknown_percentage=25 --quantize=1 -``` - -If you see a compiling error on older machines, try leaving out the `--copt` arguments, they are just there to accelerate training on chips that support the extensions. The training process is likely to take a couple of hours. Once it has completed, the next step is to freeze the variables: - -``` -bazel run tensorflow/examples/speech_commands:freeze -- \ ---model_architecture=tiny_conv --window_stride=20 --preprocess=average \ ---wanted_words="yes,no" --quantize=1 --output_file=/tmp/tiny_conv.pb -``` - -The next step is to create a TensorFlow Lite file from the frozen graph: - -``` -bazel run tensorflow/contrib/lite/toco:toco -- \ ---input_file=/tmp/tiny_conv.pb --output_file=/tmp/tiny_conv.tflite \ ---input_shapes=1,49,43,1 --input_arrays=Reshape_1 --output_arrays='labels_softmax' \ ---inference_type=QUANTIZED_UINT8 --mean_values=0 --std_values=2 \ ---change_concat_input_ranges=false -``` - -Finally, convert the file into a C source file that can be compiled into an embedded system: - -``` -xxd -i /tmp/tiny_conv.tflite > /tmp/tiny_conv_model_data.cc -``` diff --git a/tensorflow/contrib/lite/experimental/micro/examples/micro_speech/no_features_data.h b/tensorflow/contrib/lite/experimental/micro/examples/micro_speech/no_features_data.h deleted file mode 100644 index b53d0a202b75eab7db82107f2c71c504a85f881e..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/experimental/micro/examples/micro_speech/no_features_data.h +++ /dev/null @@ -1,23 +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_LITE_EXPERIMENTAL_MICRO_EXAMPLES_MICRO_SPEECH_NO_FEATURES_DATA_H_ -#define TENSORFLOW_CONTRIB_LITE_EXPERIMENTAL_MICRO_EXAMPLES_MICRO_SPEECH_NO_FEATURES_DATA_H_ - -extern const int g_no_f9643d42_nohash_4_width; -extern const int g_no_f9643d42_nohash_4_height; -extern const unsigned char g_no_f9643d42_nohash_4_data[]; - -#endif // TENSORFLOW_CONTRIB_LITE_EXPERIMENTAL_MICRO_EXAMPLES_MICRO_SPEECH_NO_FEATURES_DATA_H_ diff --git a/tensorflow/contrib/lite/experimental/micro/examples/micro_speech/yes_features_data.h b/tensorflow/contrib/lite/experimental/micro/examples/micro_speech/yes_features_data.h deleted file mode 100644 index 33ac2308624235fc380782cd61e6a0247b81b093..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/experimental/micro/examples/micro_speech/yes_features_data.h +++ /dev/null @@ -1,23 +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_LITE_EXPERIMENTAL_MICRO_EXAMPLES_MICRO_SPEECH_YES_FEATURES_DATA_H_ -#define TENSORFLOW_CONTRIB_LITE_EXPERIMENTAL_MICRO_EXAMPLES_MICRO_SPEECH_YES_FEATURES_DATA_H_ - -extern const int g_yes_f2e59fea_nohash_1_width; -extern const int g_yes_f2e59fea_nohash_1_height; -extern const unsigned char g_yes_f2e59fea_nohash_1_data[]; - -#endif // TENSORFLOW_CONTRIB_LITE_EXPERIMENTAL_MICRO_EXAMPLES_MICRO_SPEECH_YES_FEATURES_DATA_H_ diff --git a/tensorflow/contrib/lite/experimental/micro/kernels/BUILD b/tensorflow/contrib/lite/experimental/micro/kernels/BUILD deleted file mode 100644 index a012f950e6f58f082d0a7c9ac0b4cd9018bcf40b..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/experimental/micro/kernels/BUILD +++ /dev/null @@ -1,107 +0,0 @@ -package(default_visibility = [ - "//visibility:public", -]) - -licenses(["notice"]) # Apache 2.0 - -load("//tensorflow/contrib/lite:build_def.bzl", "tflite_copts") -load( - "//tensorflow/contrib/lite/experimental/micro/testing:micro_test.bzl", - "tflite_micro_cc_test", -) - -cc_library( - name = "micro_ops", - srcs = [ - "depthwise_conv.cc", - "fully_connected.cc", - "softmax.cc", - ], - hdrs = [ - ], - copts = tflite_copts(), - deps = [ - "//tensorflow/contrib/lite/c:c_api_internal", - "//tensorflow/contrib/lite/experimental/micro:micro_framework", - "//tensorflow/contrib/lite/kernels:kernel_util", - "//tensorflow/contrib/lite/kernels:op_macros", - "//tensorflow/contrib/lite/kernels:padding", - "//tensorflow/contrib/lite/kernels/internal:quantization_util", - "//tensorflow/contrib/lite/kernels/internal:reference_base", - "//tensorflow/contrib/lite/kernels/internal:tensor", - ], -) - -cc_library( - name = "all_ops_resolver", - srcs = [ - "all_ops_resolver.cc", - ], - hdrs = [ - "all_ops_resolver.h", - ], - copts = tflite_copts(), - deps = [ - ":micro_ops", - "//tensorflow/contrib/lite/c:c_api_internal", - "//tensorflow/contrib/lite/experimental/micro:micro_framework", - ], -) - -cc_library( - name = "test_utils", - srcs = [ - ], - hdrs = [ - "test_utils.h", - ], - copts = tflite_copts(), - deps = [ - "//tensorflow/contrib/lite/c:c_api_internal", - "//tensorflow/contrib/lite/core/api", - "//tensorflow/contrib/lite/experimental/micro:micro_framework", - "//tensorflow/contrib/lite/experimental/micro/testing:micro_test", - ], -) - -tflite_micro_cc_test( - name = "depthwise_conv_test", - srcs = [ - "depthwise_conv_test.cc", - ], - deps = [ - ":all_ops_resolver", - ":test_utils", - "//tensorflow/contrib/lite/c:c_api_internal", - "//tensorflow/contrib/lite/experimental/micro:micro_framework", - "//tensorflow/contrib/lite/experimental/micro/testing:micro_test", - ], -) - -tflite_micro_cc_test( - name = "fully_connected_test", - srcs = [ - "fully_connected_test.cc", - ], - deps = [ - ":all_ops_resolver", - ":test_utils", - "//tensorflow/contrib/lite/c:c_api_internal", - "//tensorflow/contrib/lite/experimental/micro:micro_framework", - "//tensorflow/contrib/lite/experimental/micro/testing:micro_test", - ], -) - -tflite_micro_cc_test( - name = "softmax_test", - srcs = [ - "softmax_test.cc", - ], - deps = [ - ":all_ops_resolver", - ":test_utils", - "//tensorflow/contrib/lite/c:c_api_internal", - "//tensorflow/contrib/lite/experimental/micro:micro_framework", - "//tensorflow/contrib/lite/experimental/micro/testing:micro_test", - ], -) diff --git a/tensorflow/contrib/lite/experimental/micro/kernels/all_ops_resolver.h b/tensorflow/contrib/lite/experimental/micro/kernels/all_ops_resolver.h deleted file mode 100644 index f836064a3f63443ff577e7ac7a8b791cbb2c24c5..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/experimental/micro/kernels/all_ops_resolver.h +++ /dev/null @@ -1,34 +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_LITE_EXPERIMENTAL_MICRO_KERNELS_ALL_OPS_RESOLVER_H_ -#define TENSORFLOW_CONTRIB_LITE_EXPERIMENTAL_MICRO_KERNELS_ALL_OPS_RESOLVER_H_ - -#include "tensorflow/contrib/lite/experimental/micro/compatibility.h" -#include "tensorflow/contrib/lite/experimental/micro/micro_mutable_op_resolver.h" - -namespace tflite { -namespace ops { -namespace micro { - -class AllOpsResolver : public MicroMutableOpResolver { - public: - AllOpsResolver(); - - private: - TF_LITE_REMOVE_VIRTUAL_DELETE -}; - -} // namespace micro -} // namespace ops -} // namespace tflite - -#endif // TENSORFLOW_CONTRIB_LITE_EXPERIMENTAL_MICRO_KERNELS_ALL_OPS_RESOLVER_H_ diff --git a/tensorflow/contrib/lite/experimental/micro/kernels/depthwise_conv.cc b/tensorflow/contrib/lite/experimental/micro/kernels/depthwise_conv.cc deleted file mode 100644 index 4f17263181982afdaa1941194b88d58f0ef0ca74..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/experimental/micro/kernels/depthwise_conv.cc +++ /dev/null @@ -1,208 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/contrib/lite/c/builtin_op_data.h" -#include "tensorflow/contrib/lite/c/c_api_internal.h" -#include "tensorflow/contrib/lite/kernels/internal/common.h" -#include "tensorflow/contrib/lite/kernels/internal/quantization_util.h" -#include "tensorflow/contrib/lite/kernels/internal/tensor_ctypes.h" -#include "tensorflow/contrib/lite/kernels/kernel_util.h" -#include "tensorflow/contrib/lite/kernels/padding.h" - -#include "tensorflow/contrib/lite/kernels/internal/reference/depthwiseconv_float.h" -#include "tensorflow/contrib/lite/kernels/internal/reference/depthwiseconv_uint8.h" - -namespace tflite { -namespace ops { -namespace micro { -namespace depthwise_conv { -namespace { - -constexpr int kInputTensor = 0; -constexpr int kFilterTensor = 1; -constexpr int kBiasTensor = 2; -constexpr int kOutputTensor = 0; - -struct OpData { - TfLitePaddingValues padding; - // The scaling factor from input to output (aka the 'real multiplier') can - // be represented as a fixed point multiplier plus a left shift. - int32_t output_multiplier; - int output_shift; - // The range of the fused activation layer. For example for kNone and - // uint8_t these would be 0 and 255. - int32_t output_activation_min; - int32_t output_activation_max; -}; - -TfLiteStatus CalculateOpData(TfLiteContext* context, TfLiteNode* node, - TfLiteDepthwiseConvParams* params, int width, - int height, int filter_width, int filter_height, - int out_width, int out_height, - const TfLiteType data_type, OpData* data) { - data->padding.height = ComputePadding(params->stride_height, 1, height, - filter_height, out_height); - data->padding.width = - ComputePadding(params->stride_width, 1, width, filter_width, out_width); - - // Note that quantized inference requires that all tensors have their - // parameters set. This is usually done during quantized training. - if (data_type != kTfLiteFloat32) { - const TfLiteTensor* input = GetInput(context, node, kInputTensor); - const TfLiteTensor* filter = GetInput(context, node, kFilterTensor); - const TfLiteTensor* bias = - GetOptionalInputTensor(context, node, kBiasTensor); - TfLiteTensor* output = GetOutput(context, node, kOutputTensor); - - double real_multiplier = 0.0; - TF_LITE_ENSURE_STATUS(GetQuantizedConvolutionMultipler( - context, input, filter, bias, output, &real_multiplier)); - int exponent; - QuantizeMultiplier(real_multiplier, &data->output_multiplier, &exponent); - data->output_shift = -exponent; - CalculateActivationRangeUint8(params->activation, output, - &data->output_activation_min, - &data->output_activation_max); - } - return kTfLiteOk; -} - -} // namespace - -void* Init(TfLiteContext* context, const char* buffer, size_t length) { - return nullptr; -} - -void Free(TfLiteContext* context, void* buffer) {} - -TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { - return kTfLiteOk; -} - -void EvalFloat(TfLiteContext* context, TfLiteNode* node, - TfLiteDepthwiseConvParams* params, OpData* data, - const TfLiteTensor* input, const TfLiteTensor* filter, - const TfLiteTensor* bias, TfLiteTensor* output) { - float output_activation_min, output_activation_max; - CalculateActivationRange(params->activation, &output_activation_min, - &output_activation_max); - - tflite::DepthwiseParams op_params; - // Padding type is ignored, but still set. - op_params.padding_type = PaddingType::kSame; - op_params.padding_values.width = data->padding.width; - op_params.padding_values.height = data->padding.height; - op_params.stride_width = params->stride_width; - op_params.stride_height = params->stride_height; - op_params.dilation_width_factor = 1; - op_params.dilation_height_factor = 1; - op_params.depth_multiplier = params->depth_multiplier; - op_params.float_activation_min = output_activation_min; - op_params.float_activation_max = output_activation_max; - - tflite::reference_ops::DepthwiseConv( - op_params, GetTensorShape(input), GetTensorData(input), - GetTensorShape(filter), GetTensorData(filter), - GetTensorShape(bias), GetTensorData(bias), GetTensorShape(output), - GetTensorData(output)); -} - -void EvalQuantized(TfLiteContext* context, TfLiteNode* node, - TfLiteDepthwiseConvParams* params, OpData* data, - const TfLiteTensor* input, const TfLiteTensor* filter, - const TfLiteTensor* bias, TfLiteTensor* output) { - const int32_t input_offset = -input->params.zero_point; - const int32_t filter_offset = -filter->params.zero_point; - const int32_t output_offset = output->params.zero_point; - - tflite::DepthwiseParams op_params; - // Padding type is ignored, but still set. - op_params.padding_type = PaddingType::kSame; - op_params.padding_values.width = data->padding.width; - op_params.padding_values.height = data->padding.height; - op_params.stride_width = params->stride_width; - op_params.stride_height = params->stride_height; - op_params.dilation_width_factor = 1; - op_params.dilation_height_factor = 1; - op_params.depth_multiplier = params->depth_multiplier; - op_params.quantized_activation_min = data->output_activation_min; - op_params.quantized_activation_max = data->output_activation_max; - op_params.input_offset = input_offset; - op_params.weights_offset = filter_offset; - op_params.output_offset = output_offset; - op_params.output_multiplier = data->output_multiplier; - // Legacy ops used mixed left and right shifts. Now all are +ve-means-left. - op_params.output_shift = -data->output_shift; - - tflite::reference_ops::DepthwiseConv( - op_params, GetTensorShape(input), GetTensorData(input), - GetTensorShape(filter), GetTensorData(filter), - GetTensorShape(bias), GetTensorData(bias), - GetTensorShape(output), GetTensorData(output)); -} - -TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { - auto* params = - reinterpret_cast(node->builtin_data); - - TfLiteTensor* output = GetOutput(context, node, kOutputTensor); - const TfLiteTensor* input = GetInput(context, node, kInputTensor); - const TfLiteTensor* filter = GetInput(context, node, kFilterTensor); - const TfLiteTensor* bias = - (NumInputs(node) == 3) ? GetInput(context, node, kBiasTensor) : nullptr; - - const TfLiteType data_type = input->type; - int width = SizeOfDimension(input, 2); - int height = SizeOfDimension(input, 1); - int filter_width = SizeOfDimension(filter, 2); - int filter_height = SizeOfDimension(filter, 1); - int out_width = ComputeOutSize(params->padding, width, filter_width, - params->stride_width); - int out_height = ComputeOutSize(params->padding, height, filter_height, - params->stride_height); - OpData local_data_object; - OpData* data = &local_data_object; - TF_LITE_ENSURE_STATUS(CalculateOpData(context, node, params, width, height, - filter_width, filter_height, out_width, - out_height, data_type, data)); - - // TODO(aselle): Consider whether float conv and quantized conv should be - // separate ops to avoid dispatch overhead here. - switch (input->type) { // Already know in/out types are same. - case kTfLiteFloat32: - EvalFloat(context, node, params, data, input, filter, bias, output); - break; - case kTfLiteUInt8: - EvalQuantized(context, node, params, data, input, filter, bias, output); - break; - default: - context->ReportError(context, "Type %d not currently supported.", - input->type); - return kTfLiteError; - } - return kTfLiteOk; -} - -} // namespace depthwise_conv - -TfLiteRegistration* Register_DEPTHWISE_CONV_2D() { - static TfLiteRegistration r = {depthwise_conv::Init, depthwise_conv::Free, - depthwise_conv::Prepare, depthwise_conv::Eval}; - return &r; -} - -} // namespace micro -} // namespace ops -} // namespace tflite diff --git a/tensorflow/contrib/lite/experimental/micro/kernels/depthwise_conv_test.cc b/tensorflow/contrib/lite/experimental/micro/kernels/depthwise_conv_test.cc deleted file mode 100644 index 169899c471dd44399b4d8a479cecbbbd78ba1215..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/experimental/micro/kernels/depthwise_conv_test.cc +++ /dev/null @@ -1,406 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/contrib/lite/c/builtin_op_data.h" -#include "tensorflow/contrib/lite/c/c_api_internal.h" -#include "tensorflow/contrib/lite/experimental/micro/kernels/all_ops_resolver.h" -#include "tensorflow/contrib/lite/experimental/micro/kernels/test_utils.h" -#include "tensorflow/contrib/lite/experimental/micro/simple_tensor_allocator.h" -#include "tensorflow/contrib/lite/experimental/micro/testing/micro_test.h" - -namespace tflite { -namespace testing { -namespace { - -void TestDepthwiseConvFloat(std::initializer_list input_dims_data, - std::initializer_list input_data, - std::initializer_list filter_dims_data, - std::initializer_list filter_data, - std::initializer_list bias_dims_data, - std::initializer_list bias_data, - std::initializer_list expected_output_data, - std::initializer_list output_dims_data, - TfLiteFusedActivation activation, - float* output_data) { - TfLiteIntArray* input_dims = IntArrayFromInitializer(input_dims_data); - TfLiteIntArray* filter_dims = IntArrayFromInitializer(filter_dims_data); - TfLiteIntArray* bias_dims = IntArrayFromInitializer(bias_dims_data); - TfLiteIntArray* output_dims = IntArrayFromInitializer(output_dims_data); - const int output_dims_count = ElementCount(*output_dims); - - constexpr int inputs_size = 3; - constexpr int outputs_size = 1; - constexpr int tensors_size = inputs_size + outputs_size; - TfLiteTensor tensors[tensors_size] = { - CreateFloatTensor(input_data, input_dims, "input_tensor"), - CreateFloatTensor(filter_data, filter_dims, "filter_tensor"), - CreateFloatTensor(bias_data, bias_dims, "bias_tensor"), - CreateFloatTensor(output_data, output_dims, "output_tensor"), - }; - - TfLiteContext context; - PopulateContext(tensors, tensors_size, &context); - - ::tflite::ops::micro::AllOpsResolver resolver; - const TfLiteRegistration* registration = - resolver.FindOp(tflite::BuiltinOperator_DEPTHWISE_CONV_2D, 1); - TF_LITE_MICRO_EXPECT_NE(nullptr, registration); - - int input_depth = input_dims->data[3]; - int output_depth = filter_dims->data[3]; - int depth_mul = output_depth / input_depth; - TfLiteDepthwiseConvParams builtin_data = { - kTfLitePaddingValid, 1, 1, depth_mul, activation, - }; - const char* init_data = reinterpret_cast(&builtin_data); - size_t init_data_size = 0; - void* user_data = nullptr; - if (registration->init) { - user_data = registration->init(&context, init_data, init_data_size); - } - int inputs_array_data[] = {3, 0, 1, 2}; - TfLiteIntArray* inputs_array = IntArrayFromInts(inputs_array_data); - int outputs_array_data[] = {1, 3}; - TfLiteIntArray* outputs_array = IntArrayFromInts(outputs_array_data); - int temporaries_array_data[] = {0}; - TfLiteIntArray* temporaries_array = IntArrayFromInts(temporaries_array_data); - - TfLiteNode node; - node.inputs = inputs_array; - node.outputs = outputs_array; - node.temporaries = temporaries_array; - node.user_data = user_data; - node.builtin_data = reinterpret_cast(&builtin_data); - node.custom_initial_data = nullptr; - node.custom_initial_data_size = 0; - node.delegate = nullptr; - if (registration->prepare) { - TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, registration->prepare(&context, &node)); - } - TF_LITE_MICRO_EXPECT_NE(nullptr, registration->invoke); - TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, registration->invoke(&context, &node)); - if (registration->free) { - registration->free(&context, user_data); - } - for (int i = 0; i < output_dims_count; ++i) { - TF_LITE_MICRO_EXPECT_NEAR(expected_output_data.begin()[i], output_data[i], - 1e-5f); - } -} - -void TestDepthwiseConvQuantized( - std::initializer_list input_dims_data, - std::initializer_list input_data, float input_min, float input_max, - std::initializer_list filter_dims_data, - std::initializer_list filter_data, float filter_min, - float filter_max, std::initializer_list bias_dims_data, - std::initializer_list bias_data, float bias_min, float bias_max, - std::initializer_list expected_output_data, - std::initializer_list output_dims_data, float output_min, - float output_max, TfLiteFusedActivation activation, uint8_t* output_data) { - TfLiteIntArray* input_dims = IntArrayFromInitializer(input_dims_data); - TfLiteIntArray* filter_dims = IntArrayFromInitializer(filter_dims_data); - TfLiteIntArray* bias_dims = IntArrayFromInitializer(bias_dims_data); - TfLiteIntArray* output_dims = IntArrayFromInitializer(output_dims_data); - const int output_dims_count = ElementCount(*output_dims); - - constexpr int inputs_size = 3; - constexpr int outputs_size = 1; - constexpr int tensors_size = inputs_size + outputs_size; - TfLiteTensor tensors[tensors_size] = { - CreateQuantizedTensor(input_data, input_dims, "input_tensor", input_min, - input_max), - CreateQuantizedTensor(filter_data, filter_dims, "filter_tensor", - filter_min, filter_max), - CreateQuantized32Tensor(bias_data, bias_dims, "bias_tensor", bias_min, - bias_max), - CreateQuantizedTensor(output_data, output_dims, "output_tensor", - output_min, output_max), - }; - - TfLiteContext context; - PopulateContext(tensors, tensors_size, &context); - - ::tflite::ops::micro::AllOpsResolver resolver; - const TfLiteRegistration* registration = - resolver.FindOp(tflite::BuiltinOperator_DEPTHWISE_CONV_2D, 1); - TF_LITE_MICRO_EXPECT_NE(nullptr, registration); - - int input_depth = input_dims->data[3]; - int output_depth = filter_dims->data[3]; - int depth_mul = output_depth / input_depth; - TfLiteDepthwiseConvParams builtin_data = { - kTfLitePaddingValid, 1, 1, depth_mul, activation, - }; - const char* init_data = reinterpret_cast(&builtin_data); - size_t init_data_size = 0; - void* user_data = nullptr; - if (registration->init) { - user_data = registration->init(&context, init_data, init_data_size); - } - - int inputs_array_data[] = {3, 0, 1, 2}; - TfLiteIntArray* inputs_array = IntArrayFromInts(inputs_array_data); - int outputs_array_data[] = {1, 3}; - TfLiteIntArray* outputs_array = IntArrayFromInts(outputs_array_data); - int temporaries_array_data[] = {0}; - TfLiteIntArray* temporaries_array = IntArrayFromInts(temporaries_array_data); - - TfLiteNode node; - node.inputs = inputs_array; - node.outputs = outputs_array; - node.temporaries = temporaries_array; - node.user_data = user_data; - node.builtin_data = reinterpret_cast(&builtin_data); - node.custom_initial_data = nullptr; - node.custom_initial_data_size = 0; - node.delegate = nullptr; - - if (registration->prepare) { - TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, registration->prepare(&context, &node)); - } - TF_LITE_MICRO_EXPECT_NE(nullptr, registration->invoke); - TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, registration->invoke(&context, &node)); - if (registration->free) { - registration->free(&context, user_data); - } - for (int i = 0; i < output_dims_count; ++i) { - TF_LITE_MICRO_EXPECT_EQ(expected_output_data.begin()[i], output_data[i]); - } -} - -} // namespace -} // namespace testing -} // namespace tflite - -TF_LITE_MICRO_TESTS_BEGIN - -TF_LITE_MICRO_TEST(SimpleTest) { - const int output_dims_count = 8; - float output_data[output_dims_count]; - tflite::testing::TestDepthwiseConvFloat( // - {4, 1, 3, 2, 2}, // Input shape. - { - 1, 2, 7, 8, // Input values. - 3, 4, 9, 10, // - 5, 6, 11, 12, // - }, - {4, 1, 2, 2, 4}, // Filters shape. - { - 1, 2, 3, 4, // Filters values. - -9, 10, -11, 12, // - 5, 6, 7, 8, // - 13, -14, 15, -16, // - }, - {1, 4}, // Bias shape. - { - 1, 2, 3, 4, // Bias values. - }, - { - 71, -34, 99, -20, // Expected results. - 91, -26, 127, -4, // - }, - {4, 1, 2, 1, 4}, // Output shape. - kTfLiteActNone, output_data); -} - -TF_LITE_MICRO_TEST(SimpleTestQuantized) { - using tflite::testing::F2Q; - using tflite::testing::F2Q32; - - const float input_min = -63.5f; - const float input_max = 64.0f; - const float filter_min = -63.5f; - const float filter_max = 64.0f; - const float bias_min = 0.0f; - const float bias_max = 64.0f * (1 << 24); - const float output_min = -127.0f; - const float output_max = 128.0f; - const int output_dims_count = 8; - uint8_t output_data[output_dims_count]; - - tflite::testing::TestDepthwiseConvQuantized( // - {4, 1, 3, 2, 2}, // Input shape. - { - // Input values. - F2Q(1, input_min, input_max), - F2Q(2, input_min, input_max), - F2Q(7, input_min, input_max), - F2Q(8, input_min, input_max), - F2Q(3, input_min, input_max), - F2Q(4, input_min, input_max), - F2Q(9, input_min, input_max), - F2Q(10, input_min, input_max), - F2Q(5, input_min, input_max), - F2Q(6, input_min, input_max), - F2Q(11, input_min, input_max), - F2Q(12, input_min, input_max), - }, - input_min, input_max, // Input quantization range. - {4, 1, 2, 2, 4}, // Filter shape. - { - // Filter values. - F2Q(1, filter_min, filter_max), - F2Q(2, filter_min, filter_max), - F2Q(3, filter_min, filter_max), - F2Q(4, filter_min, filter_max), - F2Q(-9, filter_min, filter_max), - F2Q(10, filter_min, filter_max), - F2Q(-11, filter_min, filter_max), - F2Q(12, filter_min, filter_max), - F2Q(5, filter_min, filter_max), - F2Q(6, filter_min, filter_max), - F2Q(7, filter_min, filter_max), - F2Q(8, filter_min, filter_max), - F2Q(13, filter_min, filter_max), - F2Q(-14, filter_min, filter_max), - F2Q(15, filter_min, filter_max), - F2Q(-16, filter_min, filter_max), - }, - filter_min, filter_max, // Filter quantization range. - {1, 4}, // Bias shape. - { - // Bias values. - F2Q32(1, bias_min, bias_max), - F2Q32(2, bias_min, bias_max), - F2Q32(3, bias_min, bias_max), - F2Q32(4, bias_min, bias_max), - }, - bias_min, bias_max, // Bias quantization range. - { - // Expected results. - F2Q(71, output_min, output_max), - F2Q(-34, output_min, output_max), - F2Q(99, output_min, output_max), - F2Q(-20, output_min, output_max), - F2Q(91, output_min, output_max), - F2Q(-26, output_min, output_max), - F2Q(127, output_min, output_max), - F2Q(-4, output_min, output_max), - }, - {4, 1, 2, 1, 4}, // Output shape. - output_min, output_max, // Output quantization range. - kTfLiteActNone, output_data); -} - -TF_LITE_MICRO_TEST(SimpleTestRelu) { - const int output_dims_count = 8; - float output_data[output_dims_count]; - tflite::testing::TestDepthwiseConvFloat( // - {4, 1, 3, 2, 2}, // Input shape. - { - 1, 2, 7, 8, // Input values. - 3, 4, 9, 10, // - 5, 6, 11, 12, // - }, - {4, 1, 2, 2, 4}, // Filters shape. - { - 1, 2, 3, 4, // Filters values. - -9, 10, -11, 12, // - 5, 6, 7, 8, // - 13, -14, 15, -16, // - }, - {1, 4}, // Bias shape. - { - 1, 2, 3, 4, // Bias values. - }, - { - 71, 0, 99, 0, // Expected results. - 91, 0, 127, 0, // - }, - {4, 1, 2, 1, 4}, // Output shape. - kTfLiteActRelu, output_data); -} - -TF_LITE_MICRO_TEST(SimpleTestReluQuantized) { - using tflite::testing::F2Q; - using tflite::testing::F2Q32; - - const float input_min = -63.5f; - const float input_max = 64.0f; - const float filter_min = -63.5f; - const float filter_max = 64.0f; - const float bias_min = 0.0f; - const float bias_max = 64.0f * (1 << 24); - const float output_min = -127.0f; - const float output_max = 128.0f; - const int output_dims_count = 8; - uint8_t output_data[output_dims_count]; - - tflite::testing::TestDepthwiseConvQuantized( // - {4, 1, 3, 2, 2}, // Input shape. - { - // Input values. - F2Q(1, input_min, input_max), - F2Q(2, input_min, input_max), - F2Q(7, input_min, input_max), - F2Q(8, input_min, input_max), - F2Q(3, input_min, input_max), - F2Q(4, input_min, input_max), - F2Q(9, input_min, input_max), - F2Q(10, input_min, input_max), - F2Q(5, input_min, input_max), - F2Q(6, input_min, input_max), - F2Q(11, input_min, input_max), - F2Q(12, input_min, input_max), - }, - input_min, input_max, // Input quantization range. - {4, 1, 2, 2, 4}, // Filter shape. - { - // Filter values. - F2Q(1, filter_min, filter_max), - F2Q(2, filter_min, filter_max), - F2Q(3, filter_min, filter_max), - F2Q(4, filter_min, filter_max), - F2Q(-9, filter_min, filter_max), - F2Q(10, filter_min, filter_max), - F2Q(-11, filter_min, filter_max), - F2Q(12, filter_min, filter_max), - F2Q(5, filter_min, filter_max), - F2Q(6, filter_min, filter_max), - F2Q(7, filter_min, filter_max), - F2Q(8, filter_min, filter_max), - F2Q(13, filter_min, filter_max), - F2Q(-14, filter_min, filter_max), - F2Q(15, filter_min, filter_max), - F2Q(-16, filter_min, filter_max), - }, - filter_min, filter_max, // Filter quantization range. - {1, 4}, // Bias shape. - { - // Bias values. - F2Q32(1, bias_min, bias_max), - F2Q32(2, bias_min, bias_max), - F2Q32(3, bias_min, bias_max), - F2Q32(4, bias_min, bias_max), - }, - bias_min, bias_max, // Bias quantization range. - { - // Expected results. - F2Q(71, output_min, output_max), - F2Q(0, output_min, output_max), - F2Q(99, output_min, output_max), - F2Q(0, output_min, output_max), - F2Q(91, output_min, output_max), - F2Q(0, output_min, output_max), - F2Q(127, output_min, output_max), - F2Q(0, output_min, output_max), - }, - {4, 1, 2, 1, 4}, // Output shape. - output_min, output_max, // Output quantization range. - kTfLiteActRelu, output_data); -} - -TF_LITE_MICRO_TESTS_END diff --git a/tensorflow/contrib/lite/experimental/micro/kernels/fully_connected.cc b/tensorflow/contrib/lite/experimental/micro/kernels/fully_connected.cc deleted file mode 100644 index 1e9e54cafb8c91af1b42d6d23396495ecad6e602..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/experimental/micro/kernels/fully_connected.cc +++ /dev/null @@ -1,184 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/contrib/lite/kernels/internal/reference/fully_connected.h" -#include "tensorflow/contrib/lite/c/builtin_op_data.h" -#include "tensorflow/contrib/lite/c/c_api_internal.h" -#include "tensorflow/contrib/lite/kernels/internal/common.h" -#include "tensorflow/contrib/lite/kernels/internal/quantization_util.h" -#include "tensorflow/contrib/lite/kernels/internal/tensor_ctypes.h" -#include "tensorflow/contrib/lite/kernels/kernel_util.h" - -namespace tflite { -namespace ops { -namespace micro { -namespace fully_connected { -namespace { - -struct OpData { - // The scaling factor from input to output (aka the 'real multiplier') can - // be represented as a fixed point multiplier plus a left shift. - int32_t output_multiplier; - int output_shift; - // The range of the fused activation layer. For example for kNone and - // uint8_t these would be 0 and 255. - int32_t output_activation_min; - int32_t output_activation_max; - // The index of the temporary tensor where the quantized inputs are cached. - int input_quantized_index; -}; - -constexpr int kInputTensor = 0; -constexpr int kWeightsTensor = 1; -constexpr int kBiasTensor = 2; -constexpr int kOutputTensor = 0; - -TfLiteStatus CalculateOpData(TfLiteContext* context, - TfLiteFullyConnectedParams* params, - TfLiteType data_type, const TfLiteTensor* input, - const TfLiteTensor* filter, - const TfLiteTensor* bias, TfLiteTensor* output, - OpData* data) { - TfLiteStatus status = kTfLiteOk; - if (data_type != kTfLiteFloat32) { - double real_multiplier = 0.0; - TF_LITE_ENSURE_STATUS(GetQuantizedConvolutionMultipler( - context, input, filter, bias, output, &real_multiplier)); - int exponent; - QuantizeMultiplier(real_multiplier, &data->output_multiplier, &exponent); - data->output_shift = -exponent; - TF_LITE_ENSURE_STATUS(CalculateActivationRangeQuantized( - context, params->activation, output, &data->output_activation_min, - &data->output_activation_max)); - } - return status; -} - -} // namespace - -void* Init(TfLiteContext* context, const char* buffer, size_t length) { - return nullptr; -} - -void Free(TfLiteContext* context, void* buffer) {} - -TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { - return kTfLiteOk; -} - -TfLiteStatus EvalQuantized(TfLiteContext* context, TfLiteNode* node, - TfLiteFullyConnectedParams* params, OpData* data, - const TfLiteTensor* input, - const TfLiteTensor* filter, const TfLiteTensor* bias, - TfLiteTensor* output) { - const int32_t input_offset = -input->params.zero_point; - const int32_t filter_offset = -filter->params.zero_point; - const int32_t output_offset = output->params.zero_point; - - tflite::FullyConnectedParams op_params; - op_params.input_offset = input_offset; - op_params.weights_offset = filter_offset; - op_params.output_offset = output_offset; - op_params.output_multiplier = data->output_multiplier; - // Legacy ops used mixed left and right shifts. Now all are +ve-means-left. - op_params.output_shift = -data->output_shift; - op_params.quantized_activation_min = data->output_activation_min; - op_params.quantized_activation_max = data->output_activation_max; - -#define TF_LITE_FULLY_CONNECTED(output_data_type) \ - reference_ops::FullyConnected( \ - op_params, GetTensorShape(input), GetTensorData(input), \ - GetTensorShape(filter), GetTensorData(filter), \ - GetTensorShape(bias), GetTensorData(bias), \ - GetTensorShape(output), GetTensorData(output), \ - nullptr) - switch (output->type) { - case kTfLiteUInt8: - TF_LITE_FULLY_CONNECTED(uint8_t); - break; - case kTfLiteInt16: - TF_LITE_FULLY_CONNECTED(int16_t); - break; - default: - context->ReportError( - context, - "Quantized FullyConnected expects output data type uint8 or int16"); - return kTfLiteError; - } - - return kTfLiteOk; -} - -TfLiteStatus EvalFloat(TfLiteContext* context, TfLiteNode* node, - TfLiteFullyConnectedParams* params, OpData* data, - const TfLiteTensor* input, const TfLiteTensor* filter, - const TfLiteTensor* bias, TfLiteTensor* output) { - float output_activation_min, output_activation_max; - CalculateActivationRange(params->activation, &output_activation_min, - &output_activation_max); - tflite::FullyConnectedParams op_params; - op_params.float_activation_min = output_activation_min; - op_params.float_activation_max = output_activation_max; - tflite::reference_ops::FullyConnected( - op_params, GetTensorShape(input), GetTensorData(input), - GetTensorShape(filter), GetTensorData(filter), - GetTensorShape(bias), GetTensorData(bias), GetTensorShape(output), - GetTensorData(output)); - return kTfLiteOk; -} - -TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { - auto* params = - reinterpret_cast(node->builtin_data); - - const TfLiteTensor* input = GetInput(context, node, kInputTensor); - const TfLiteTensor* filter = GetInput(context, node, kWeightsTensor); - const TfLiteTensor* bias = GetOptionalInputTensor(context, node, kBiasTensor); - TfLiteTensor* output = GetOutput(context, node, kOutputTensor); - - TfLiteType data_type = input->type; - OpData local_data_object; - OpData* data = &local_data_object; - TF_LITE_ENSURE_STATUS(CalculateOpData(context, params, data_type, input, - filter, bias, output, data)); - - switch (filter->type) { // Already know in/out types are same. - case kTfLiteFloat32: - return EvalFloat(context, node, params, data, input, filter, bias, - output); - case kTfLiteUInt8: - return EvalQuantized(context, node, params, data, input, filter, bias, - output); - - default: - context->ReportError(context, "Type %d not currently supported.", - filter->type); - return kTfLiteError; - } - return kTfLiteOk; -} - -} // namespace fully_connected - -TfLiteRegistration* Register_FULLY_CONNECTED() { - static TfLiteRegistration r = {fully_connected::Init, fully_connected::Free, - fully_connected::Prepare, - fully_connected::Eval}; - return &r; -} - -} // namespace micro -} // namespace ops -} // namespace tflite diff --git a/tensorflow/contrib/lite/experimental/micro/kernels/fully_connected_test.cc b/tensorflow/contrib/lite/experimental/micro/kernels/fully_connected_test.cc deleted file mode 100644 index b42bf4c3bca75572dbf8e1907e7fb94be24d41bd..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/experimental/micro/kernels/fully_connected_test.cc +++ /dev/null @@ -1,643 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/contrib/lite/c/builtin_op_data.h" -#include "tensorflow/contrib/lite/c/c_api_internal.h" -#include "tensorflow/contrib/lite/experimental/micro/kernels/all_ops_resolver.h" -#include "tensorflow/contrib/lite/experimental/micro/kernels/test_utils.h" -#include "tensorflow/contrib/lite/experimental/micro/simple_tensor_allocator.h" -#include "tensorflow/contrib/lite/experimental/micro/testing/micro_test.h" - -namespace tflite { -namespace testing { -namespace { - -void TestFullyConnectedFloat(std::initializer_list input_dims_data, - std::initializer_list input_data, - std::initializer_list weights_dims_data, - std::initializer_list weights_data, - std::initializer_list bias_dims_data, - std::initializer_list bias_data, - std::initializer_list expected_output_data, - std::initializer_list output_dims_data, - TfLiteFusedActivation activation, - float* output_data) { - TfLiteIntArray* input_dims = IntArrayFromInitializer(input_dims_data); - TfLiteIntArray* weights_dims = IntArrayFromInitializer(weights_dims_data); - TfLiteIntArray* bias_dims = IntArrayFromInitializer(bias_dims_data); - TfLiteIntArray* output_dims = IntArrayFromInitializer(output_dims_data); - const int output_dims_count = ElementCount(*output_dims); - - constexpr int inputs_size = 3; - constexpr int outputs_size = 1; - constexpr int tensors_size = inputs_size + outputs_size; - TfLiteTensor tensors[tensors_size] = { - CreateFloatTensor(input_data, input_dims, "input_tensor"), - CreateFloatTensor(weights_data, weights_dims, "weights_tensor"), - CreateFloatTensor(bias_data, bias_dims, "bias_tensor"), - CreateFloatTensor(output_data, output_dims, "output_tensor"), - }; - - TfLiteContext context; - PopulateContext(tensors, tensors_size, &context); - - ::tflite::ops::micro::AllOpsResolver resolver; - const TfLiteRegistration* registration = - resolver.FindOp(tflite::BuiltinOperator_FULLY_CONNECTED, 1); - TF_LITE_MICRO_EXPECT_NE(nullptr, registration); - - TfLiteFullyConnectedParams builtin_data = { - activation, - kTfLiteFullyConnectedWeightsFormatDefault, - }; - const char* init_data = reinterpret_cast(&builtin_data); - size_t init_data_size = 0; - void* user_data = nullptr; - if (registration->init) { - user_data = registration->init(&context, init_data, init_data_size); - } - int inputs_array_data[] = {3, 0, 1, 2}; - TfLiteIntArray* inputs_array = IntArrayFromInts(inputs_array_data); - int outputs_array_data[] = {1, 3}; - TfLiteIntArray* outputs_array = IntArrayFromInts(outputs_array_data); - int temporaries_array_data[] = {0}; - TfLiteIntArray* temporaries_array = IntArrayFromInts(temporaries_array_data); - - TfLiteNode node; - node.inputs = inputs_array; - node.outputs = outputs_array; - node.temporaries = temporaries_array; - node.user_data = user_data; - node.builtin_data = reinterpret_cast(&builtin_data); - node.custom_initial_data = nullptr; - node.custom_initial_data_size = 0; - node.delegate = nullptr; - if (registration->prepare) { - TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, registration->prepare(&context, &node)); - } - TF_LITE_MICRO_EXPECT_NE(nullptr, registration->invoke); - TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, registration->invoke(&context, &node)); - if (registration->free) { - registration->free(&context, user_data); - } - for (int i = 0; i < output_dims_count; ++i) { - TF_LITE_MICRO_EXPECT_NEAR(expected_output_data.begin()[i], output_data[i], - 1e-5f); - } -} - -void TestFullyConnectedQuantized( - std::initializer_list input_dims_data, - std::initializer_list input_data, float input_min, float input_max, - std::initializer_list weights_dims_data, - std::initializer_list weights_data, float weights_min, - float weights_max, std::initializer_list bias_dims_data, - std::initializer_list bias_data, float bias_min, float bias_max, - std::initializer_list expected_output_data, - std::initializer_list output_dims_data, float output_min, - float output_max, TfLiteFusedActivation activation, uint8_t* output_data) { - TfLiteIntArray* input_dims = IntArrayFromInitializer(input_dims_data); - TfLiteIntArray* weights_dims = IntArrayFromInitializer(weights_dims_data); - TfLiteIntArray* bias_dims = IntArrayFromInitializer(bias_dims_data); - TfLiteIntArray* output_dims = IntArrayFromInitializer(output_dims_data); - const int output_dims_count = ElementCount(*output_dims); - - constexpr int inputs_size = 3; - constexpr int outputs_size = 1; - constexpr int tensors_size = inputs_size + outputs_size; - TfLiteTensor tensors[tensors_size] = { - CreateQuantizedTensor(input_data, input_dims, "input_tensor", input_min, - input_max), - CreateQuantizedTensor(weights_data, weights_dims, "weights_tensor", - weights_min, weights_max), - CreateQuantized32Tensor(bias_data, bias_dims, "bias_tensor", bias_min, - bias_max), - CreateQuantizedTensor(output_data, output_dims, "output_tensor", - output_min, output_max), - }; - - TfLiteContext context; - PopulateContext(tensors, tensors_size, &context); - - ::tflite::ops::micro::AllOpsResolver resolver; - const TfLiteRegistration* registration = - resolver.FindOp(tflite::BuiltinOperator_FULLY_CONNECTED, 1); - TF_LITE_MICRO_EXPECT_NE(nullptr, registration); - - TfLiteFullyConnectedParams builtin_data = { - activation, - kTfLiteFullyConnectedWeightsFormatDefault, - }; - const char* init_data = reinterpret_cast(&builtin_data); - size_t init_data_size = 0; - void* user_data = nullptr; - if (registration->init) { - user_data = registration->init(&context, init_data, init_data_size); - } - - int inputs_array_data[] = {3, 0, 1, 2}; - TfLiteIntArray* inputs_array = IntArrayFromInts(inputs_array_data); - int outputs_array_data[] = {1, 3}; - TfLiteIntArray* outputs_array = IntArrayFromInts(outputs_array_data); - int temporaries_array_data[] = {0}; - TfLiteIntArray* temporaries_array = IntArrayFromInts(temporaries_array_data); - - TfLiteNode node; - node.inputs = inputs_array; - node.outputs = outputs_array; - node.temporaries = temporaries_array; - node.user_data = user_data; - node.builtin_data = reinterpret_cast(&builtin_data); - node.custom_initial_data = nullptr; - node.custom_initial_data_size = 0; - node.delegate = nullptr; - - if (registration->prepare) { - TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, registration->prepare(&context, &node)); - } - TF_LITE_MICRO_EXPECT_NE(nullptr, registration->invoke); - TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, registration->invoke(&context, &node)); - if (registration->free) { - registration->free(&context, user_data); - } - for (int i = 0; i < output_dims_count; ++i) { - TF_LITE_MICRO_EXPECT_EQ(expected_output_data.begin()[i], output_data[i]); - } -} - -} // namespace -} // namespace testing -} // namespace tflite - -TF_LITE_MICRO_TESTS_BEGIN - -TF_LITE_MICRO_TEST(SimpleTest) { - const int output_dims_count = 6; - float output_data[output_dims_count]; - tflite::testing::TestFullyConnectedFloat( // - {2, 2, 10}, // Input shape. - { - 1, 2, 3, 4, 5, 6, 7, 8, -9, -10, // b = 0 - 1, 2, 3, 4, 5, 6, 7, -8, 9, -10, // b = 1 - }, - {2, 3, 10}, // Weights shape. - { - 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, // u = 0 - 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, // u = 1 - 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, // u = 2 - }, - {1, 3}, // Bias shape. - { - 1, 2, 3, // Bias values. - }, - { - 24, 25, 26, 58, 59, 60, // Expected results. - }, - {2, 2, 3}, // Output shape. - kTfLiteActNone, output_data); -} - -TF_LITE_MICRO_TEST(SimpleTest2) { - const int output_dims_count = 6; - float output_data[output_dims_count]; - tflite::testing::TestFullyConnectedFloat( // - {2, 2, 2}, // Input shape. - { - 1, 2, // b = 0 - 2, 1, // b = 1 - }, - {2, 1, 2}, // Weights shape. - { - 2, 4, // u = 0 - }, - {1, 1}, // Bias shape. - { - 1, // Bias values. - }, - { - 11, 9, // Expected results. - }, - {2, 2, 1}, // Output shape. - kTfLiteActNone, output_data); -} - -TF_LITE_MICRO_TEST(SimpleTestRelu) { - const int output_dims_count = 6; - float output_data[output_dims_count]; - tflite::testing::TestFullyConnectedFloat( // - {2, 2, 10}, // Input shape. - { - 1, 2, 3, 4, 5, 6, 7, 8, -9, -10, // b = 0 - 1, 2, 3, 4, 5, 6, 7, -8, 9, -10, // b = 1 - }, - {2, 3, 10}, // Weights shape. - { - 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, // u = 0 - -1, -2, -3, -4, -5, -6, -7, -8, -9, -10, // u = 1 - 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, // u = 2 - }, - {1, 3}, // Bias shape. - { - 1, -2, 3, // Bias values. - }, - { - 24, 0, 26, 58, 0, 60, // Expected results. - }, - {2, 2, 3}, // Output shape. - kTfLiteActRelu, output_data); -} - -TF_LITE_MICRO_TEST(SimpleTestQuantized) { - using tflite::testing::F2Q; - using tflite::testing::F2Q32; - - const float input_min = -63.5f; - const float input_max = 64.0f; - const float weights_min = -63.5f; - const float weights_max = 64.0f; - const float bias_min = 0.0f; - const float bias_max = 64.0f * (1 << 24); - const float output_min = -127.0f; - const float output_max = 128.0f; - const int output_dims_count = 6; - uint8_t output_data[output_dims_count]; - tflite::testing::TestFullyConnectedQuantized( // - {2, 2, 10}, // Input shape. - { - // Input values. - F2Q(1, input_min, input_max), F2Q(2, input_min, input_max), - F2Q(3, input_min, input_max), F2Q(4, input_min, input_max), - F2Q(5, input_min, input_max), F2Q(6, input_min, input_max), - F2Q(7, input_min, input_max), F2Q(8, input_min, input_max), - F2Q(-9, input_min, input_max), F2Q(-10, input_min, input_max), - F2Q(1, input_min, input_max), F2Q(2, input_min, input_max), - F2Q(3, input_min, input_max), F2Q(4, input_min, input_max), - F2Q(5, input_min, input_max), F2Q(6, input_min, input_max), - F2Q(7, input_min, input_max), F2Q(-8, input_min, input_max), - F2Q(9, input_min, input_max), F2Q(-10, input_min, input_max), - }, - input_min, input_max, // Input quantization range. - {2, 3, 10}, // Weights shape. - { - // Weight values. - F2Q(1, weights_min, weights_max), F2Q(2, weights_min, weights_max), - F2Q(3, weights_min, weights_max), F2Q(4, weights_min, weights_max), - F2Q(5, weights_min, weights_max), F2Q(6, weights_min, weights_max), - F2Q(7, weights_min, weights_max), F2Q(8, weights_min, weights_max), - F2Q(9, weights_min, weights_max), F2Q(10, weights_min, weights_max), - F2Q(1, weights_min, weights_max), F2Q(2, weights_min, weights_max), - F2Q(3, weights_min, weights_max), F2Q(4, weights_min, weights_max), - F2Q(5, weights_min, weights_max), F2Q(6, weights_min, weights_max), - F2Q(7, weights_min, weights_max), F2Q(8, weights_min, weights_max), - F2Q(9, weights_min, weights_max), F2Q(10, weights_min, weights_max), - F2Q(1, weights_min, weights_max), F2Q(2, weights_min, weights_max), - F2Q(3, weights_min, weights_max), F2Q(4, weights_min, weights_max), - F2Q(5, weights_min, weights_max), F2Q(6, weights_min, weights_max), - F2Q(7, weights_min, weights_max), F2Q(8, weights_min, weights_max), - F2Q(9, weights_min, weights_max), F2Q(10, weights_min, weights_max), - }, - weights_min, weights_max, // Weights quantization range. - {1, 3}, // Bias shape. - { - F2Q32(1, bias_min, bias_max), - F2Q32(2, bias_min, bias_max), - F2Q32(3, bias_min, bias_max), - }, - bias_min, bias_max, // Bias quantization range. - { - // Expected results. - F2Q(24, output_min, output_max), - F2Q(25, output_min, output_max), - F2Q(26, output_min, output_max), - F2Q(58, output_min, output_max), - F2Q(59, output_min, output_max), - F2Q(60, output_min, output_max), - }, - {2, 2, 3}, // Output shape. - output_min, output_max, // Output quantization range. - kTfLiteActNone, output_data); -} - -TF_LITE_MICRO_TEST(SimpleTestQuantizedRelu) { - using tflite::testing::F2Q; - using tflite::testing::F2Q32; - - const float input_min = -63.5f; - const float input_max = 64.0f; - const float weights_min = -63.5f; - const float weights_max = 64.0f; - const float bias_min = 0.0f; - const float bias_max = 64.0f * (1 << 24); - const float output_min = -127.0f; - const float output_max = 128.0f; - const int output_dims_count = 6; - uint8_t output_data[output_dims_count]; - tflite::testing::TestFullyConnectedQuantized( // - {2, 2, 10}, // Input shape. - { - // Input values. - F2Q(1, input_min, input_max), F2Q(2, input_min, input_max), - F2Q(3, input_min, input_max), F2Q(4, input_min, input_max), - F2Q(5, input_min, input_max), F2Q(6, input_min, input_max), - F2Q(7, input_min, input_max), F2Q(8, input_min, input_max), - F2Q(-9, input_min, input_max), F2Q(-10, input_min, input_max), - F2Q(1, input_min, input_max), F2Q(2, input_min, input_max), - F2Q(3, input_min, input_max), F2Q(4, input_min, input_max), - F2Q(5, input_min, input_max), F2Q(6, input_min, input_max), - F2Q(7, input_min, input_max), F2Q(-8, input_min, input_max), - F2Q(9, input_min, input_max), F2Q(-10, input_min, input_max), - }, - input_min, input_max, // Input quantization range. - {2, 3, 10}, // Weights shape. - { - // Weight values. - F2Q(1, weights_min, weights_max), F2Q(2, weights_min, weights_max), - F2Q(3, weights_min, weights_max), F2Q(4, weights_min, weights_max), - F2Q(5, weights_min, weights_max), F2Q(6, weights_min, weights_max), - F2Q(7, weights_min, weights_max), F2Q(8, weights_min, weights_max), - F2Q(9, weights_min, weights_max), F2Q(10, weights_min, weights_max), - F2Q(-1, weights_min, weights_max), F2Q(-2, weights_min, weights_max), - F2Q(-3, weights_min, weights_max), F2Q(-4, weights_min, weights_max), - F2Q(-5, weights_min, weights_max), F2Q(-6, weights_min, weights_max), - F2Q(-7, weights_min, weights_max), F2Q(-8, weights_min, weights_max), - F2Q(-9, weights_min, weights_max), F2Q(-10, weights_min, weights_max), - F2Q(1, weights_min, weights_max), F2Q(2, weights_min, weights_max), - F2Q(3, weights_min, weights_max), F2Q(4, weights_min, weights_max), - F2Q(5, weights_min, weights_max), F2Q(6, weights_min, weights_max), - F2Q(7, weights_min, weights_max), F2Q(8, weights_min, weights_max), - F2Q(9, weights_min, weights_max), F2Q(10, weights_min, weights_max), - }, - weights_min, weights_max, // Weights quantization range. - {1, 3}, // Bias shape. - { - F2Q32(1, bias_min, bias_max), - F2Q32(0, bias_min, bias_max), - F2Q32(3, bias_min, bias_max), - }, - bias_min, bias_max, // Bias quantization range. - { - // Expected results. - F2Q(24, output_min, output_max), - F2Q(0, output_min, output_max), - F2Q(26, output_min, output_max), - F2Q(58, output_min, output_max), - F2Q(0, output_min, output_max), - F2Q(60, output_min, output_max), - }, - {2, 2, 3}, // Output shape. - output_min, output_max, // Output quantization range. - kTfLiteActRelu, output_data); -} - -TF_LITE_MICRO_TEST(SimpleTestQuantizedOutputMultiplierGreaterThan1) { - using tflite::testing::F2Q; - using tflite::testing::F2Q32; - - const float input_min = -127.0f; - const float input_max = 128.0f; - const float weights_min = -127.0f; - const float weights_max = 128.0f; - const float bias_min = 0.0f; - const float bias_max = 256.0f * (1 << 24); - const float output_min = -63.5f; - const float output_max = 64.0f; - const int output_dims_count = 6; - uint8_t output_data[output_dims_count]; - tflite::testing::TestFullyConnectedQuantized( // - {2, 2, 10}, // Input shape. - { - // Input values. - F2Q(1, input_min, input_max), F2Q(2, input_min, input_max), - F2Q(3, input_min, input_max), F2Q(4, input_min, input_max), - F2Q(5, input_min, input_max), F2Q(6, input_min, input_max), - F2Q(7, input_min, input_max), F2Q(8, input_min, input_max), - F2Q(-9, input_min, input_max), F2Q(-10, input_min, input_max), - F2Q(1, input_min, input_max), F2Q(2, input_min, input_max), - F2Q(3, input_min, input_max), F2Q(4, input_min, input_max), - F2Q(5, input_min, input_max), F2Q(6, input_min, input_max), - F2Q(7, input_min, input_max), F2Q(-8, input_min, input_max), - F2Q(9, input_min, input_max), F2Q(-10, input_min, input_max), - }, - input_min, input_max, // Input quantization range. - {2, 3, 10}, // Weights shape. - { - // Weight values. - F2Q(1, weights_min, weights_max), F2Q(2, weights_min, weights_max), - F2Q(3, weights_min, weights_max), F2Q(4, weights_min, weights_max), - F2Q(5, weights_min, weights_max), F2Q(6, weights_min, weights_max), - F2Q(7, weights_min, weights_max), F2Q(8, weights_min, weights_max), - F2Q(9, weights_min, weights_max), F2Q(10, weights_min, weights_max), - F2Q(1, weights_min, weights_max), F2Q(2, weights_min, weights_max), - F2Q(3, weights_min, weights_max), F2Q(4, weights_min, weights_max), - F2Q(5, weights_min, weights_max), F2Q(6, weights_min, weights_max), - F2Q(7, weights_min, weights_max), F2Q(8, weights_min, weights_max), - F2Q(9, weights_min, weights_max), F2Q(10, weights_min, weights_max), - F2Q(1, weights_min, weights_max), F2Q(2, weights_min, weights_max), - F2Q(3, weights_min, weights_max), F2Q(4, weights_min, weights_max), - F2Q(5, weights_min, weights_max), F2Q(6, weights_min, weights_max), - F2Q(7, weights_min, weights_max), F2Q(8, weights_min, weights_max), - F2Q(9, weights_min, weights_max), F2Q(10, weights_min, weights_max), - }, - weights_min, weights_max, // Weights quantization range. - {1, 3}, // Bias shape. - { - F2Q32(1, bias_min, bias_max), - F2Q32(2, bias_min, bias_max), - F2Q32(3, bias_min, bias_max), - }, - bias_min, bias_max, // Bias quantization range. - { - // Expected results. - F2Q(24, output_min, output_max), - F2Q(25, output_min, output_max), - F2Q(26, output_min, output_max), - F2Q(58, output_min, output_max), - F2Q(59, output_min, output_max), - F2Q(60, output_min, output_max), - }, - {2, 2, 3}, // Output shape. - output_min, output_max, // Output quantization range. - kTfLiteActNone, output_data); -} - -TF_LITE_MICRO_TEST(SimpleTest4DInput) { - const int output_dims_count = 6; - float output_data[output_dims_count]; - tflite::testing::TestFullyConnectedFloat( // - {4, 1, 1, 5, 1}, // Input shape. - { - 1, 2, 3, 4, 5, 6, 7, 8, -9, -10, // b = 0 - 1, 2, 3, 4, 5, 6, 7, -8, 9, -10, // b = 1 - }, - {2, 3, 10}, // Weights shape. - { - 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, // u = 0 - 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, // u = 1 - 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, // u = 2 - }, - {1, 3}, // Bias shape. - { - 1, 2, 3, // Bias values. - }, - { - 24, 25, 26, 58, 59, 60, // Expected results. - }, - {2, 2, 3}, // Output shape. - kTfLiteActNone, output_data); -} - -TF_LITE_MICRO_TEST(SimpleTest4DInputQuantized) { - using tflite::testing::F2Q; - using tflite::testing::F2Q32; - - const float input_min = -63.5f; - const float input_max = 64.0f; - const float weights_min = -63.5f; - const float weights_max = 64.0f; - const float bias_min = 0.0f; - const float bias_max = 64.0f * (1 << 24); - const float output_min = -127.0f; - const float output_max = 128.0f; - const int output_dims_count = 6; - uint8_t output_data[output_dims_count]; - tflite::testing::TestFullyConnectedQuantized( // - {4, 1, 1, 5, 1}, // Input shape. - { - // Input values. - F2Q(1, input_min, input_max), F2Q(2, input_min, input_max), - F2Q(3, input_min, input_max), F2Q(4, input_min, input_max), - F2Q(5, input_min, input_max), F2Q(6, input_min, input_max), - F2Q(7, input_min, input_max), F2Q(8, input_min, input_max), - F2Q(-9, input_min, input_max), F2Q(-10, input_min, input_max), - F2Q(1, input_min, input_max), F2Q(2, input_min, input_max), - F2Q(3, input_min, input_max), F2Q(4, input_min, input_max), - F2Q(5, input_min, input_max), F2Q(6, input_min, input_max), - F2Q(7, input_min, input_max), F2Q(-8, input_min, input_max), - F2Q(9, input_min, input_max), F2Q(-10, input_min, input_max), - }, - input_min, input_max, // Input quantization range. - {2, 3, 10}, // Weights shape. - { - // Weight values. - F2Q(1, weights_min, weights_max), F2Q(2, weights_min, weights_max), - F2Q(3, weights_min, weights_max), F2Q(4, weights_min, weights_max), - F2Q(5, weights_min, weights_max), F2Q(6, weights_min, weights_max), - F2Q(7, weights_min, weights_max), F2Q(8, weights_min, weights_max), - F2Q(9, weights_min, weights_max), F2Q(10, weights_min, weights_max), - F2Q(1, weights_min, weights_max), F2Q(2, weights_min, weights_max), - F2Q(3, weights_min, weights_max), F2Q(4, weights_min, weights_max), - F2Q(5, weights_min, weights_max), F2Q(6, weights_min, weights_max), - F2Q(7, weights_min, weights_max), F2Q(8, weights_min, weights_max), - F2Q(9, weights_min, weights_max), F2Q(10, weights_min, weights_max), - F2Q(1, weights_min, weights_max), F2Q(2, weights_min, weights_max), - F2Q(3, weights_min, weights_max), F2Q(4, weights_min, weights_max), - F2Q(5, weights_min, weights_max), F2Q(6, weights_min, weights_max), - F2Q(7, weights_min, weights_max), F2Q(8, weights_min, weights_max), - F2Q(9, weights_min, weights_max), F2Q(10, weights_min, weights_max), - }, - weights_min, weights_max, // Weights quantization range. - {1, 3}, // Bias shape. - { - F2Q32(1, bias_min, bias_max), - F2Q32(2, bias_min, bias_max), - F2Q32(3, bias_min, bias_max), - }, - bias_min, bias_max, // Bias quantization range. - { - // Expected results. - F2Q(24, output_min, output_max), - F2Q(25, output_min, output_max), - F2Q(26, output_min, output_max), - F2Q(58, output_min, output_max), - F2Q(59, output_min, output_max), - F2Q(60, output_min, output_max), - }, - {2, 2, 3}, // Output shape. - output_min, output_max, // Output quantization range. - kTfLiteActNone, output_data); -} - -TF_LITE_MICRO_TEST(SimpleTest4DInputQuantizedOutputMultiplierGreaterThan1) { - using tflite::testing::F2Q; - using tflite::testing::F2Q32; - - const float input_min = -127.0f; - const float input_max = 128.0f; - const float weights_min = -127.0f; - const float weights_max = 128.0f; - const float bias_min = 0.0f; - const float bias_max = 256.0f * (1 << 24); - const float output_min = -63.5f; - const float output_max = 64.0f; - const int output_dims_count = 6; - uint8_t output_data[output_dims_count]; - tflite::testing::TestFullyConnectedQuantized( // - {4, 1, 1, 5, 1}, // Input shape. - { - // Input values. - F2Q(1, input_min, input_max), F2Q(2, input_min, input_max), - F2Q(3, input_min, input_max), F2Q(4, input_min, input_max), - F2Q(5, input_min, input_max), F2Q(6, input_min, input_max), - F2Q(7, input_min, input_max), F2Q(8, input_min, input_max), - F2Q(-9, input_min, input_max), F2Q(-10, input_min, input_max), - F2Q(1, input_min, input_max), F2Q(2, input_min, input_max), - F2Q(3, input_min, input_max), F2Q(4, input_min, input_max), - F2Q(5, input_min, input_max), F2Q(6, input_min, input_max), - F2Q(7, input_min, input_max), F2Q(-8, input_min, input_max), - F2Q(9, input_min, input_max), F2Q(-10, input_min, input_max), - }, - input_min, input_max, // Input quantization range. - {2, 3, 10}, // Weights shape. - { - // Weight values. - F2Q(1, weights_min, weights_max), F2Q(2, weights_min, weights_max), - F2Q(3, weights_min, weights_max), F2Q(4, weights_min, weights_max), - F2Q(5, weights_min, weights_max), F2Q(6, weights_min, weights_max), - F2Q(7, weights_min, weights_max), F2Q(8, weights_min, weights_max), - F2Q(9, weights_min, weights_max), F2Q(10, weights_min, weights_max), - F2Q(1, weights_min, weights_max), F2Q(2, weights_min, weights_max), - F2Q(3, weights_min, weights_max), F2Q(4, weights_min, weights_max), - F2Q(5, weights_min, weights_max), F2Q(6, weights_min, weights_max), - F2Q(7, weights_min, weights_max), F2Q(8, weights_min, weights_max), - F2Q(9, weights_min, weights_max), F2Q(10, weights_min, weights_max), - F2Q(1, weights_min, weights_max), F2Q(2, weights_min, weights_max), - F2Q(3, weights_min, weights_max), F2Q(4, weights_min, weights_max), - F2Q(5, weights_min, weights_max), F2Q(6, weights_min, weights_max), - F2Q(7, weights_min, weights_max), F2Q(8, weights_min, weights_max), - F2Q(9, weights_min, weights_max), F2Q(10, weights_min, weights_max), - }, - weights_min, weights_max, // Weights quantization range. - {1, 3}, // Bias shape. - { - F2Q32(1, bias_min, bias_max), - F2Q32(2, bias_min, bias_max), - F2Q32(3, bias_min, bias_max), - }, - bias_min, bias_max, // Bias quantization range. - { - // Expected results. - F2Q(24, output_min, output_max), - F2Q(25, output_min, output_max), - F2Q(26, output_min, output_max), - F2Q(58, output_min, output_max), - F2Q(59, output_min, output_max), - F2Q(60, output_min, output_max), - }, - {2, 2, 3}, // Output shape. - output_min, output_max, // Output quantization range. - kTfLiteActNone, output_data); -} - -TF_LITE_MICRO_TESTS_END diff --git a/tensorflow/contrib/lite/experimental/micro/kernels/softmax_test.cc b/tensorflow/contrib/lite/experimental/micro/kernels/softmax_test.cc deleted file mode 100644 index 694456d8ace5182578f9b59c2de8bbad0447b4ee..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/experimental/micro/kernels/softmax_test.cc +++ /dev/null @@ -1,220 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/contrib/lite/c/builtin_op_data.h" -#include "tensorflow/contrib/lite/c/c_api_internal.h" -#include "tensorflow/contrib/lite/experimental/micro/kernels/all_ops_resolver.h" -#include "tensorflow/contrib/lite/experimental/micro/kernels/test_utils.h" -#include "tensorflow/contrib/lite/experimental/micro/simple_tensor_allocator.h" -#include "tensorflow/contrib/lite/experimental/micro/testing/micro_test.h" - -namespace tflite { -namespace testing { -namespace { - -void TestSoftmaxFloat(std::initializer_list input_dims_data, - std::initializer_list input_data, - std::initializer_list expected_output_data, - std::initializer_list output_dims_data, - float* output_data) { - TfLiteIntArray* input_dims = IntArrayFromInitializer(input_dims_data); - TfLiteIntArray* output_dims = IntArrayFromInitializer(output_dims_data); - const int output_dims_count = ElementCount(*output_dims); - - constexpr int inputs_size = 2; - constexpr int outputs_size = 1; - constexpr int tensors_size = inputs_size + outputs_size; - TfLiteTensor tensors[tensors_size] = { - CreateFloatTensor(input_data, input_dims, "input_tensor"), - CreateFloatTensor(output_data, output_dims, "output_tensor"), - }; - - TfLiteContext context; - PopulateContext(tensors, tensors_size, &context); - - ::tflite::ops::micro::AllOpsResolver resolver; - const TfLiteRegistration* registration = - resolver.FindOp(tflite::BuiltinOperator_SOFTMAX, 1); - TF_LITE_MICRO_EXPECT_NE(nullptr, registration); - - TfLiteSoftmaxParams builtin_data = {1.0f}; - const char* init_data = reinterpret_cast(&builtin_data); - size_t init_data_size = 0; - void* user_data = nullptr; - if (registration->init) { - user_data = registration->init(&context, init_data, init_data_size); - } - int inputs_array_data[] = {1, 0}; - TfLiteIntArray* inputs_array = IntArrayFromInts(inputs_array_data); - int outputs_array_data[] = {1, 1}; - TfLiteIntArray* outputs_array = IntArrayFromInts(outputs_array_data); - int temporaries_array_data[] = {0}; - TfLiteIntArray* temporaries_array = IntArrayFromInts(temporaries_array_data); - - TfLiteNode node; - node.inputs = inputs_array; - node.outputs = outputs_array; - node.temporaries = temporaries_array; - node.user_data = user_data; - node.builtin_data = reinterpret_cast(&builtin_data); - node.custom_initial_data = nullptr; - node.custom_initial_data_size = 0; - node.delegate = nullptr; - if (registration->prepare) { - TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, registration->prepare(&context, &node)); - } - TF_LITE_MICRO_EXPECT_NE(nullptr, registration->invoke); - TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, registration->invoke(&context, &node)); - if (registration->free) { - registration->free(&context, user_data); - } - for (int i = 0; i < output_dims_count; ++i) { - TF_LITE_MICRO_EXPECT_NEAR(expected_output_data.begin()[i], output_data[i], - 1e-5f); - } -} - -void TestSoftmaxQuantized(std::initializer_list input_dims_data, - std::initializer_list input_data, - float input_min, float input_max, - std::initializer_list expected_output_data, - std::initializer_list output_dims_data, - float output_min, float output_max, - uint8_t* output_data) { - TfLiteIntArray* input_dims = IntArrayFromInitializer(input_dims_data); - TfLiteIntArray* output_dims = IntArrayFromInitializer(output_dims_data); - const int output_dims_count = ElementCount(*output_dims); - - constexpr int inputs_size = 1; - constexpr int outputs_size = 1; - constexpr int tensors_size = inputs_size + outputs_size; - TfLiteTensor tensors[tensors_size] = { - CreateQuantizedTensor(input_data, input_dims, "input_tensor", input_min, - input_max), - CreateQuantizedTensor(output_data, output_dims, "output_tensor", - output_min, output_max), - }; - - TfLiteContext context; - PopulateContext(tensors, tensors_size, &context); - - ::tflite::ops::micro::AllOpsResolver resolver; - const TfLiteRegistration* registration = - resolver.FindOp(tflite::BuiltinOperator_SOFTMAX, 1); - TF_LITE_MICRO_EXPECT_NE(nullptr, registration); - - TfLiteSoftmaxParams builtin_data = {1.0f}; - const char* init_data = reinterpret_cast(&builtin_data); - size_t init_data_size = 0; - void* user_data = nullptr; - if (registration->init) { - user_data = registration->init(&context, init_data, init_data_size); - } - - int inputs_array_data[] = {1, 0}; - TfLiteIntArray* inputs_array = IntArrayFromInts(inputs_array_data); - int outputs_array_data[] = {1, 1}; - TfLiteIntArray* outputs_array = IntArrayFromInts(outputs_array_data); - int temporaries_array_data[] = {0}; - TfLiteIntArray* temporaries_array = IntArrayFromInts(temporaries_array_data); - - TfLiteNode node; - node.inputs = inputs_array; - node.outputs = outputs_array; - node.temporaries = temporaries_array; - node.user_data = user_data; - node.builtin_data = reinterpret_cast(&builtin_data); - node.custom_initial_data = nullptr; - node.custom_initial_data_size = 0; - node.delegate = nullptr; - - if (registration->prepare) { - TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, registration->prepare(&context, &node)); - } - TF_LITE_MICRO_EXPECT_NE(nullptr, registration->invoke); - TF_LITE_MICRO_EXPECT_EQ(kTfLiteOk, registration->invoke(&context, &node)); - if (registration->free) { - registration->free(&context, user_data); - } - for (int i = 0; i < output_dims_count; ++i) { - TF_LITE_MICRO_EXPECT_EQ(expected_output_data.begin()[i], output_data[i]); - } -} - -} // namespace -} // namespace testing -} // namespace tflite - -TF_LITE_MICRO_TESTS_BEGIN - -TF_LITE_MICRO_TEST(SimpleTest) { - const int output_dims_count = 10; - float output_data[output_dims_count]; - tflite::testing::TestSoftmaxFloat( // - {2, 2, 5}, // Input shape. - { - 1.0, 2.0, 3.0, 4.0, 5.0, // b = 0 - -1.0, -2.0, -3.0, -4.0, -5.0, // b = 0 - }, - { - // Expected results. - 0.011656231, - 0.031684921, - 0.086128544, - 0.234121657, - 0.636408647, - 0.636408647, - 0.234121657, - 0.086128544, - 0.031684921, - 0.011656231, - }, - {2, 2, 5}, // Output shape. - output_data); -} - -TF_LITE_MICRO_TEST(SimpleTestQuantized) { - using tflite::testing::F2Q; - - const float input_min = -63.5f; - const float input_max = 64.0f; - const float output_min = 0.0f; - const float output_max = (255.0f / 256.0f); - const int output_dims_count = 5; - uint8_t output_data[output_dims_count]; - tflite::testing::TestSoftmaxQuantized( // - {2, 1, 5}, // Input shape. - { - F2Q(1.0, input_min, input_max), - F2Q(2.0, input_min, input_max), - F2Q(3.0, input_min, input_max), - F2Q(4.0, input_min, input_max), - F2Q(5.0, input_min, input_max), - }, - input_min, input_max, // Input quantized range. - { - // Expected results. - F2Q(0.011656231, output_min, output_max), - F2Q(0.031684921, output_min, output_max), - F2Q(0.086128544, output_min, output_max), - F2Q(0.234121657, output_min, output_max), - F2Q(0.636408647, output_min, output_max), - }, - {2, 1, 5}, // Output shape. - output_min, output_max, // Output quantized range. - output_data); -} - -TF_LITE_MICRO_TESTS_END diff --git a/tensorflow/contrib/lite/experimental/micro/kernels/test_utils.h b/tensorflow/contrib/lite/experimental/micro/kernels/test_utils.h deleted file mode 100644 index 789a48ece8bd68544649fb05548355cb796ccabb..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/experimental/micro/kernels/test_utils.h +++ /dev/null @@ -1,170 +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_LITE_EXPERIMENTAL_MICRO_KERNELS_TEST_UTILS_H_ -#define TENSORFLOW_CONTRIB_LITE_EXPERIMENTAL_MICRO_KERNELS_TEST_UTILS_H_ - -#include -#include -#include - -#include "tensorflow/contrib/lite/c/builtin_op_data.h" -#include "tensorflow/contrib/lite/c/c_api_internal.h" -#include "tensorflow/contrib/lite/core/api/error_reporter.h" -#include "tensorflow/contrib/lite/experimental/micro/kernels/test_utils.h" -#include "tensorflow/contrib/lite/experimental/micro/testing/micro_test.h" - -namespace tflite { -namespace testing { - -// How many elements are in the array with this shape. -inline int ElementCount(const TfLiteIntArray& dims) { - int result = 1; - for (int i = 0; i < dims.size; ++i) { - result *= dims.data[i]; - } - return result; -} - -// Wrapper to forward kernel errors to the interpreter's error reporter. -inline void ReportOpError(struct TfLiteContext* context, const char* format, - ...) { - ErrorReporter* error_reporter = static_cast(context->impl_); - va_list args; - va_start(args, format); - error_reporter->Report(format, args); - va_end(args); -} - -// Derives the quantization scaling factor from a min and max range. -template -inline float ScaleFromMinMax(const float min, const float max) { - return (max - min) / ((std::numeric_limits::max() * 1.0) - - std::numeric_limits::min()); -} - -// Derives the quantization zero point from a min and max range. -template -inline int ZeroPointFromMinMax(const float min, const float max) { - return static_cast((-min / ScaleFromMinMax(min, max)) + 0.5f); -} - -// Converts a float value into an unsigned eight-bit quantized value. -inline uint8_t F2Q(const float value, const float min, const float max) { - int32_t result = ZeroPointFromMinMax(min, max) + - (value / ScaleFromMinMax(min, max)) + 0.5f; - if (result < 0) { - result = 0; - } - if (result > 256) { - result = 256; - } - return result; -} - -// Converts a float value into a signed thirty-two-bit quantized value. -inline uint8_t F2Q32(const float value, const float min, const float max) { - return static_cast((value - ZeroPointFromMinMax(min, max)) / - ScaleFromMinMax(min, max)); -} - -inline void PopulateContext(TfLiteTensor* tensors, int tensors_size, - TfLiteContext* context) { - context->tensors_size = tensors_size; - context->tensors = tensors; - context->impl_ = static_cast(micro_test::reporter); - context->GetExecutionPlan = nullptr; - context->ResizeTensor = nullptr; - context->ReportError = ReportOpError; - context->AddTensors = nullptr; - context->GetNodeAndRegistration = nullptr; - context->ReplaceSubgraphsWithDelegateKernels = nullptr; - context->recommended_num_threads = 1; - context->GetExternalContext = nullptr; - context->SetExternalContext = nullptr; -} - -inline TfLiteIntArray* IntArrayFromInts(const int* int_array) { - return const_cast( - reinterpret_cast(int_array)); -} - -inline TfLiteIntArray* IntArrayFromInitializer( - std::initializer_list int_initializer) { - return IntArrayFromInts(int_initializer.begin()); -} - -inline TfLiteTensor CreateFloatTensor(const float* data, TfLiteIntArray* dims, - const char* name) { - const size_t bytes = ElementCount(*dims) * sizeof(float); - return { - kTfLiteFloat32, {const_cast(reinterpret_cast(data))}, - dims, {}, - kTfLiteMemNone, bytes, - nullptr, name}; -} - -inline TfLiteTensor CreateFloatTensor(std::initializer_list data, - TfLiteIntArray* dims, const char* name) { - return CreateFloatTensor(data.begin(), dims, name); -} - -inline TfLiteTensor CreateQuantizedTensor(const uint8_t* data, - TfLiteIntArray* dims, - const char* name, float min, - float max) { - const size_t bytes = ElementCount(*dims) * sizeof(uint8_t); - const TfLiteQuantizationParams q_params = { - ScaleFromMinMax(min, max), - ZeroPointFromMinMax(min, max)}; - return { - kTfLiteUInt8, {const_cast(reinterpret_cast(data))}, - dims, q_params, - kTfLiteMemNone, bytes, - nullptr, name}; -} - -inline TfLiteTensor CreateQuantizedTensor(std::initializer_list data, - TfLiteIntArray* dims, - const char* name, float min, - float max) { - return CreateQuantizedTensor(data.begin(), dims, name, min, max); -} - -inline TfLiteTensor CreateQuantized32Tensor(const int32_t* data, - TfLiteIntArray* dims, - const char* name, float min, - float max) { - const size_t bytes = ElementCount(*dims) * sizeof(int32_t); - const TfLiteQuantizationParams q_params = { - ScaleFromMinMax(min, max), - ZeroPointFromMinMax(min, max)}; - return { - kTfLiteUInt8, {const_cast(reinterpret_cast(data))}, - dims, q_params, - kTfLiteMemNone, bytes, - nullptr, name}; -} - -inline TfLiteTensor CreateQuantized32Tensor(std::initializer_list data, - TfLiteIntArray* dims, - const char* name, float min, - float max) { - return CreateQuantized32Tensor(data.begin(), dims, name, min, max); -} - -} // namespace testing -} // namespace tflite - -#endif // TENSORFLOW_CONTRIB_LITE_EXPERIMENTAL_MICRO_KERNELS_TEST_UTILS_H_ diff --git a/tensorflow/contrib/lite/experimental/micro/micro_error_reporter.h b/tensorflow/contrib/lite/experimental/micro/micro_error_reporter.h deleted file mode 100644 index 33e54f7990af6cff4f8706d2889c335087581af4..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/experimental/micro/micro_error_reporter.h +++ /dev/null @@ -1,34 +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_LITE_EXPERIMENTAL_MICRO_MICRO_ERROR_REPORTER_H_ -#define TENSORFLOW_CONTRIB_LITE_EXPERIMENTAL_MICRO_MICRO_ERROR_REPORTER_H_ - -#include "tensorflow/contrib/lite/core/api/error_reporter.h" -#include "tensorflow/contrib/lite/experimental/micro/compatibility.h" - -namespace tflite { - -class MicroErrorReporter : public ErrorReporter { - public: - ~MicroErrorReporter() {} - int Report(const char* format, va_list args) override; - - private: - TF_LITE_REMOVE_VIRTUAL_DELETE -}; - -} // namespace tflite - -#endif // TENSORFLOW_CONTRIB_LITE_EXPERIMENTAL_MICRO_MICRO_ERROR_REPORTER_H_ diff --git a/tensorflow/contrib/lite/experimental/micro/testing/BUILD b/tensorflow/contrib/lite/experimental/micro/testing/BUILD deleted file mode 100644 index 0d23be5712ad1bc6d81cc467cce8c9927caece3d..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/experimental/micro/testing/BUILD +++ /dev/null @@ -1,17 +0,0 @@ -package( - default_visibility = ["//visibility:public"], -) - -licenses(["notice"]) # Apache 2.0 - -exports_files(["test_linux_binary.sh"]) - -cc_library( - name = "micro_test", - hdrs = [ - "micro_test.h", - ], - deps = [ - "//tensorflow/contrib/lite/experimental/micro:micro_framework", - ], -) diff --git a/tensorflow/contrib/lite/experimental/micro/tools/make/Makefile b/tensorflow/contrib/lite/experimental/micro/tools/make/Makefile deleted file mode 100644 index 3f749e53ef1aa995247f16cba059c369e27757c9..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/experimental/micro/tools/make/Makefile +++ /dev/null @@ -1,168 +0,0 @@ -MAKEFILE_DIR := tensorflow/contrib/lite/experimental/micro/tools/make - -# Try to figure out the host system -HOST_OS := -ifeq ($(OS),Windows_NT) - HOST_OS = windows -else - UNAME_S := $(shell uname -s) - ifeq ($(UNAME_S),Linux) - HOST_OS := linux - endif - ifeq ($(UNAME_S),Darwin) - HOST_OS := osx - endif -endif - -HOST_ARCH := $(shell if [[ $(shell uname -m) =~ i[345678]86 ]]; then echo x86_32; else echo $(shell uname -m); fi) - -# Override these on the make command line to target a specific architecture. For example: -# make -f tensorflow/contrib/lite/Makefile TARGET=rpi TARGET_ARCH=armv7l -TARGET := $(HOST_OS) -TARGET_ARCH := $(HOST_ARCH) - -INCLUDES := \ --I. \ --I$(MAKEFILE_DIR)/../../../../../ \ --I$(MAKEFILE_DIR)/../../../../../../ \ --I$(MAKEFILE_DIR)/downloads/ \ --I$(MAKEFILE_DIR)/downloads/gemmlowp \ --I$(MAKEFILE_DIR)/downloads/flatbuffers/include \ --I$(OBJDIR) -# This is at the end so any globally-installed frameworks like protobuf don't -# override local versions in the source tree. -INCLUDES += -I/usr/local/include - -TEST_SCRIPT := tensorflow/contrib/lite/experimental/micro/testing/test_linux_binary.sh - -MICROLITE_LIBS := -lm - -# There are no rules for compiling objects for the host system (since we don't -# generate things like the protobuf compiler that require that), so all of -# these settings are for the target compiler. -CXXFLAGS := -O3 -DNDEBUG -CXXFLAGS += --std=c++11 -g -DTF_LITE_STATIC_MEMORY -CCFLAGS := -DNDEBUG -g -DTF_LITE_STATIC_MEMORY -LDOPTS := -L/usr/local/lib -ARFLAGS := -r -TARGET_TOOLCHAIN_PREFIX := -CC_PREFIX := - -# This library is the main target for this makefile. It will contain a minimal -# runtime that can be linked in to other programs. -MICROLITE_LIB_NAME := libtensorflow-microlite.a - -# Test binary for the microcontroller speech model. -MICRO_SPEECH_TEST_SRCS := \ -tensorflow/contrib/lite/experimental/micro/examples/micro_speech/micro_speech_test.cc \ -tensorflow/contrib/lite/experimental/micro/examples/micro_speech/tiny_conv_model_data.cc \ -tensorflow/contrib/lite/experimental/micro/examples/micro_speech/no_features_data.cc \ -tensorflow/contrib/lite/experimental/micro/examples/micro_speech/yes_features_data.cc - -MICROLITE_TEST_SRCS := \ -$(wildcard tensorflow/contrib/lite/experimental/micro/*test.cc) \ -$(wildcard tensorflow/contrib/lite/experimental/micro/kernels/*test.cc) - -MICROLITE_CC_BASE_SRCS := \ -$(wildcard tensorflow/contrib/lite/experimental/micro/*.cc) \ -$(wildcard tensorflow/contrib/lite/experimental/micro/kernels/*.cc) \ -tensorflow/contrib/lite/c/c_api_internal.c \ -tensorflow/contrib/lite/core/api/error_reporter.cc \ -tensorflow/contrib/lite/core/api/flatbuffer_conversions.cc \ -tensorflow/contrib/lite/core/api/op_resolver.cc \ -tensorflow/contrib/lite/kernels/kernel_util.cc \ -tensorflow/contrib/lite/kernels/internal/quantization_util.cc -MICROLITE_CC_SRCS := $(filter-out $(MICROLITE_TEST_SRCS), $(MICROLITE_CC_BASE_SRCS)) - -# These target-specific makefiles should modify or replace options like -# CXXFLAGS or LIBS to work for a specific targetted architecture. All logic -# based on platforms or architectures should happen within these files, to -# keep this main makefile focused on the sources and dependencies. -include $(wildcard $(MAKEFILE_DIR)/targets/*_makefile.inc) - -ALL_SRCS := \ - $(MICRO_SPEECH_TEST_SRCS) \ - $(MICROLITE_CC_SRCS) \ - $(MICROLITE_TEST_SRCS) - -# Where compiled objects are stored. -GENDIR := $(MAKEFILE_DIR)/gen/$(TARGET)_$(TARGET_ARCH)/ -OBJDIR := $(GENDIR)obj/ -BINDIR := $(GENDIR)bin/ -LIBDIR := $(GENDIR)lib/ - -MICROLITE_LIB_PATH := $(LIBDIR)$(MICROLITE_LIB_NAME) - -MICRO_SPEECH_TEST_BINARY := $(BINDIR)micro_speech_test - -CXX := $(CC_PREFIX)${TARGET_TOOLCHAIN_PREFIX}g++ -CC := $(CC_PREFIX)${TARGET_TOOLCHAIN_PREFIX}gcc -AR := $(CC_PREFIX)${TARGET_TOOLCHAIN_PREFIX}ar - -MICRO_SPEECH_TEST_OBJS := $(addprefix $(OBJDIR), \ -$(patsubst %.cc,%.o,$(patsubst %.c,%.o,$(MICRO_SPEECH_TEST_SRCS)))) - -MICROLITE_LIB_OBJS := $(addprefix $(OBJDIR), \ -$(patsubst %.cc,%.o,$(patsubst %.c,%.o,$(MICROLITE_CC_SRCS)))) - -MICROLITE_TEST_TARGETS := $(addprefix $(BINDIR), \ -$(patsubst %_test.cc,%.test_target,$(MICROLITE_TEST_SRCS))) - -# For normal manually-created TensorFlow C++ source files. -$(OBJDIR)%.o: %.cc - @mkdir -p $(dir $@) - $(CXX) $(CXXFLAGS) $(INCLUDES) -c $< -o $@ - -# For normal manually-created TensorFlow C source files. -$(OBJDIR)%.o: %.c - @mkdir -p $(dir $@) - $(CC) $(CCFLAGS) $(INCLUDES) -c $< -o $@ - -# The target that's compiled if there's no command-line arguments. -all: $(MICROLITE_LIB_PATH) $(MICRO_SPEECH_TEST_BINARY) - -microlite: $(MICROLITE_LIB_PATH) - -# Hack for generating schema file bypassing flatbuffer parsing -tensorflow/contrib/lite/schema/schema_generated.h: - @cp -u tensorflow/contrib/lite/schema/schema_generated.h.OPENSOURCE tensorflow/contrib/lite/schema/schema_generated.h - -# Gathers together all the objects we've compiled into a single '.a' archive. -$(MICROLITE_LIB_PATH): tensorflow/contrib/lite/schema/schema_generated.h $(MICROLITE_LIB_OBJS) - @mkdir -p $(dir $@) - $(AR) $(ARFLAGS) $(MICROLITE_LIB_PATH) $(MICROLITE_LIB_OBJS) - -$(MICRO_SPEECH_TEST_BINARY): $(MICRO_SPEECH_TEST_OBJS) $(MICROLITE_LIB_PATH) - @mkdir -p $(dir $@) - $(CXX) $(CXXFLAGS) $(INCLUDES) \ - -o $(MICRO_SPEECH_TEST_BINARY) $(MICRO_SPEECH_TEST_OBJS) \ - $(LIBFLAGS) $(MICROLITE_LIB_PATH) $(LDFLAGS) $(MICROLITE_LIBS) - -micro_speech_test: $(MICRO_SPEECH_TEST_BINARY) -micro_speech_test_bin: $(MICRO_SPEECH_TEST_BINARY).bin - -test_micro_speech: $(MICRO_SPEECH_TEST_BINARY) - $(TEST_SCRIPT) $(MICRO_SPEECH_TEST_BINARY) '~~~ALL TESTS PASSED~~~' - -$(BINDIR)%_test : $(OBJDIR)%_test.o $(MICROLITE_LIB_PATH) - @mkdir -p $(dir $@) - $(CXX) $(CXXFLAGS) $(INCLUDES) \ - -o $@ $< \ - $(LIBFLAGS) $(MICROLITE_LIB_PATH) $(LDFLAGS) $(MICROLITE_LIBS) - -$(BINDIR)%.test_target: $(BINDIR)%_test - $(TEST_SCRIPT) $< '~~~ALL TESTS PASSED~~~' - -$(info $(MICROLITE_TEST_TARGETS)) - -test: test_micro_speech $(MICROLITE_TEST_TARGETS) - -# Gets rid of all generated files. -clean: - rm -rf $(MAKEFILE_DIR)/gen - -$(DEPDIR)/%.d: ; -.PRECIOUS: $(DEPDIR)/%.d -.PRECIOUS: $(BINDIR)%_test - --include $(patsubst %,$(DEPDIR)/%.d,$(basename $(ALL_SRCS))) diff --git a/tensorflow/contrib/lite/experimental/micro/tools/make/download_dependencies.sh b/tensorflow/contrib/lite/experimental/micro/tools/make/download_dependencies.sh deleted file mode 100755 index 4c2ff8545dbdcc426bf62aaeb07ca22d8b17cc69..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/experimental/micro/tools/make/download_dependencies.sh +++ /dev/null @@ -1,73 +0,0 @@ -#!/bin/bash -# 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. -# ============================================================================== - -set -e - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -cd "$SCRIPT_DIR/../../../../../../.." - -DOWNLOADS_DIR=tensorflow/contrib/lite/experimental/micro/tools/make/downloads -BZL_FILE_PATH=tensorflow/workspace.bzl - -# Ensure it is being run from repo root -if [ ! -f $BZL_FILE_PATH ]; then - echo "Could not find ${BZL_FILE_PATH}": - echo "Likely you are not running this from the root directory of the repository."; - exit 1; -fi - -GEMMLOWP_URL="https://github.com/google/gemmlowp/archive/719139ce755a0f31cbf1c37f7f98adcc7fc9f425.zip" -FLATBUFFERS_URL="https://github.com/google/flatbuffers/archive/1f5eae5d6a135ff6811724f6c57f911d1f46bb15.tar.gz" -CMSIS_URL="https://github.com/ARM-software/CMSIS_5/archive/5.4.0.zip" -STM32_BARE_LIB_URL="https://github.com/google/stm32_bare_lib/archive/50e0da307a2821bb54af1f57b969e6b76cb89d32.zip" - -download_and_extract() { - local usage="Usage: download_and_extract URL DIR" - local url="${1:?${usage}}" - local dir="${2:?${usage}}" - echo "downloading ${url}" >&2 - mkdir -p "${dir}" - if [[ "${url}" == *gz ]]; then - curl -Ls "${url}" | tar -C "${dir}" --strip-components=1 -xz - elif [[ "${url}" == *zip ]]; then - tempdir=$(mktemp -d) - tempdir2=$(mktemp -d) - - curl -L ${url} > ${tempdir}/zipped.zip - unzip ${tempdir}/zipped.zip -d ${tempdir2} - - # If the zip file contains nested directories, extract the files from the - # inner directory. - if ls ${tempdir2}/*/* 1> /dev/null 2>&1; then - # unzip has no strip components, so unzip to a temp dir, and move the - # files we want from the tempdir to destination. - cp -R ${tempdir2}/*/* ${dir}/ - else - cp -R ${tempdir2}/* ${dir}/ - fi - rm -rf ${tempdir2} ${tempdir} - fi - - # Delete any potential BUILD files, which would interfere with Bazel builds. - find "${dir}" -type f -name '*BUILD' -delete -} - -download_and_extract "${GEMMLOWP_URL}" "${DOWNLOADS_DIR}/gemmlowp" -download_and_extract "${FLATBUFFERS_URL}" "${DOWNLOADS_DIR}/flatbuffers" -download_and_extract "${CMSIS_URL}" "${DOWNLOADS_DIR}/cmsis" -download_and_extract "${STM32_BARE_LIB_URL}" "${DOWNLOADS_DIR}/stm32_bare_lib" - -echo "download_dependencies.sh completed successfully." >&2 diff --git a/tensorflow/contrib/lite/experimental/microfrontend/lib/BUILD b/tensorflow/contrib/lite/experimental/microfrontend/lib/BUILD deleted file mode 100644 index 3fd4b9fe82f7959fd86df7696950a9d3ae205042..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/experimental/microfrontend/lib/BUILD +++ /dev/null @@ -1,188 +0,0 @@ -# Library for generating feature vectors from audio data - -package( - default_visibility = ["//visibility:private"], -) - -licenses(["notice"]) # Apache 2.0 - -cc_library( - name = "bits", - hdrs = ["bits.h"], -) - -cc_library( - name = "fft", - srcs = [ - "fft.c", - "fft_util.c", - ], - hdrs = [ - "fft.h", - "fft_util.h", - ], - deps = ["@kissfft//:kiss_fftr_16"], -) - -cc_library( - name = "filterbank", - srcs = [ - "filterbank.c", - "filterbank_util.c", - ], - hdrs = [ - "filterbank.h", - "filterbank_util.h", - ], - deps = [ - ":bits", - ":fft", - ], -) - -cc_library( - name = "frontend", - srcs = [ - "frontend.c", - "frontend_util.c", - ], - hdrs = [ - "frontend.h", - "frontend_util.h", - ], - deps = [ - ":bits", - ":fft", - ":filterbank", - ":log_scale", - ":noise_reduction", - ":pcan_gain_control", - ":window", - ], -) - -cc_library( - name = "log_scale", - srcs = [ - "log_lut.c", - "log_scale.c", - "log_scale_util.c", - ], - hdrs = [ - "log_lut.h", - "log_scale.h", - "log_scale_util.h", - ], - deps = [ - ":bits", - ], -) - -cc_library( - name = "noise_reduction", - srcs = [ - "noise_reduction.c", - "noise_reduction_util.c", - ], - hdrs = [ - "noise_reduction.h", - "noise_reduction_util.h", - ], -) - -cc_library( - name = "pcan_gain_control", - srcs = [ - "pcan_gain_control.c", - "pcan_gain_control_util.c", - ], - hdrs = [ - "pcan_gain_control.h", - "pcan_gain_control_util.h", - ], - deps = [ - ":bits", - ], -) - -cc_library( - name = "window", - srcs = [ - "window.c", - "window_util.c", - ], - hdrs = [ - "window.h", - "window_util.h", - ], -) - -cc_test( - name = "fft_test", - size = "small", - srcs = ["fft_test.cc"], - deps = [ - ":fft", - "@com_google_googletest//:gtest_main", - ], -) - -cc_test( - name = "filterbank_test", - size = "small", - srcs = ["filterbank_test.cc"], - deps = [ - ":filterbank", - "@com_google_googletest//:gtest_main", - ], -) - -cc_test( - name = "frontend_test", - size = "small", - srcs = ["frontend_test.cc"], - deps = [ - ":frontend", - "@com_google_googletest//:gtest_main", - ], -) - -cc_test( - name = "log_scale_test", - size = "small", - srcs = ["log_scale_test.cc"], - deps = [ - ":log_scale", - "@com_google_googletest//:gtest_main", - ], -) - -cc_test( - name = "noise_reduction_test", - size = "small", - srcs = ["noise_reduction_test.cc"], - deps = [ - ":noise_reduction", - "@com_google_googletest//:gtest_main", - ], -) - -cc_test( - name = "pcan_gain_control_test", - size = "small", - srcs = ["pcan_gain_control_test.cc"], - deps = [ - ":pcan_gain_control", - "@com_google_googletest//:gtest_main", - ], -) - -cc_test( - name = "window_test", - size = "small", - srcs = ["window_test.cc"], - deps = [ - ":window", - "@com_google_googletest//:gtest_main", - ], -) diff --git a/tensorflow/contrib/lite/experimental/microfrontend/lib/frontend.h b/tensorflow/contrib/lite/experimental/microfrontend/lib/frontend.h deleted file mode 100644 index 71ae81024cb3aee6248de17712d7051665aed97c..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/experimental/microfrontend/lib/frontend.h +++ /dev/null @@ -1,64 +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_LITE_EXPERIMENTAL_MICROFRONTEND_LIB_FRONTEND_H_ -#define TENSORFLOW_CONTRIB_LITE_EXPERIMENTAL_MICROFRONTEND_LIB_FRONTEND_H_ - -#include -#include - -#include "tensorflow/contrib/lite/experimental/microfrontend/lib/fft.h" -#include "tensorflow/contrib/lite/experimental/microfrontend/lib/filterbank.h" -#include "tensorflow/contrib/lite/experimental/microfrontend/lib/log_scale.h" -#include "tensorflow/contrib/lite/experimental/microfrontend/lib/noise_reduction.h" -#include "tensorflow/contrib/lite/experimental/microfrontend/lib/pcan_gain_control.h" -#include "tensorflow/contrib/lite/experimental/microfrontend/lib/window.h" - -#ifdef __cplusplus -extern "C" { -#endif - -struct FrontendState { - struct WindowState window; - struct FftState fft; - struct FilterbankState filterbank; - struct NoiseReductionState noise_reduction; - struct PcanGainControlState pcan_gain_control; - struct LogScaleState log_scale; -}; - -struct FrontendOutput { - const uint16_t* values; - size_t size; -}; - -// Main entry point to processing frontend samples. Updates num_samples_read to -// contain the number of samples that have been consumed from the input array. -// Returns a struct containing the generated output. If not enough samples were -// added to generate a feature vector, the returned size will be 0 and the -// values pointer will be NULL. Note that the output pointer will be invalidated -// as soon as FrontendProcessSamples is called again, so copy the contents -// elsewhere if you need to use them later. -struct FrontendOutput FrontendProcessSamples(struct FrontendState* state, - const int16_t* samples, - size_t num_samples, - size_t* num_samples_read); - -void FrontendReset(struct FrontendState* state); - -#ifdef __cplusplus -} // extern "C" -#endif - -#endif // TENSORFLOW_CONTRIB_LITE_EXPERIMENTAL_MICROFRONTEND_LIB_FRONTEND_H_ diff --git a/tensorflow/contrib/lite/experimental/microfrontend/lib/frontend_io.h b/tensorflow/contrib/lite/experimental/microfrontend/lib/frontend_io.h deleted file mode 100644 index 4f45577caeab7aaff9210355d9d4d811dedfdf16..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/experimental/microfrontend/lib/frontend_io.h +++ /dev/null @@ -1,31 +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_LITE_EXPERIMENTAL_MICROFRONTEND_LIB_FRONTEND_IO_H_ -#define TENSORFLOW_CONTRIB_LITE_EXPERIMENTAL_MICROFRONTEND_LIB_FRONTEND_IO_H_ - -#include "tensorflow/contrib/lite/experimental/microfrontend/lib/frontend.h" - -#ifdef __cplusplus -extern "C" { -#endif - -int WriteFrontendStateMemmap(const char* header, const char* source, - const struct FrontendState* state); - -#ifdef __cplusplus -} // extern "C" -#endif - -#endif // TENSORFLOW_CONTRIB_LITE_EXPERIMENTAL_MICROFRONTEND_LIB_FRONTEND_IO_H_ diff --git a/tensorflow/contrib/lite/experimental/microfrontend/lib/frontend_util.h b/tensorflow/contrib/lite/experimental/microfrontend/lib/frontend_util.h deleted file mode 100644 index a958b610eae689192aa96eacefe654292350fcd7..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/experimental/microfrontend/lib/frontend_util.h +++ /dev/null @@ -1,52 +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_LITE_EXPERIMENTAL_MICROFRONTEND_LIB_FRONTEND_UTIL_H_ -#define TENSORFLOW_CONTRIB_LITE_EXPERIMENTAL_MICROFRONTEND_LIB_FRONTEND_UTIL_H_ - -#include "tensorflow/contrib/lite/experimental/microfrontend/lib/fft_util.h" -#include "tensorflow/contrib/lite/experimental/microfrontend/lib/filterbank_util.h" -#include "tensorflow/contrib/lite/experimental/microfrontend/lib/frontend.h" -#include "tensorflow/contrib/lite/experimental/microfrontend/lib/log_scale_util.h" -#include "tensorflow/contrib/lite/experimental/microfrontend/lib/noise_reduction_util.h" -#include "tensorflow/contrib/lite/experimental/microfrontend/lib/pcan_gain_control_util.h" -#include "tensorflow/contrib/lite/experimental/microfrontend/lib/window_util.h" - -#ifdef __cplusplus -extern "C" { -#endif - -struct FrontendConfig { - struct WindowConfig window; - struct FilterbankConfig filterbank; - struct NoiseReductionConfig noise_reduction; - struct PcanGainControlConfig pcan_gain_control; - struct LogScaleConfig log_scale; -}; - -// Fills the frontendConfig with "sane" defaults. -void FrontendFillConfigWithDefaults(struct FrontendConfig* config); - -// Allocates any buffers. -int FrontendPopulateState(const struct FrontendConfig* config, - struct FrontendState* state, int sample_rate); - -// Frees any allocated buffers. -void FrontendFreeStateContents(struct FrontendState* state); - -#ifdef __cplusplus -} // extern "C" -#endif - -#endif // TENSORFLOW_CONTRIB_LITE_EXPERIMENTAL_MICROFRONTEND_LIB_FRONTEND_UTIL_H_ diff --git a/tensorflow/contrib/lite/experimental/microfrontend/lib/log_scale_io.h b/tensorflow/contrib/lite/experimental/microfrontend/lib/log_scale_io.h deleted file mode 100644 index 5444303b2445ac72d1069eff5c50db8622d4b536..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/experimental/microfrontend/lib/log_scale_io.h +++ /dev/null @@ -1,33 +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_LITE_EXPERIMENTAL_MICROFRONTEND_LIB_LOG_SCALE_IO_H_ -#define TENSORFLOW_CONTRIB_LITE_EXPERIMENTAL_MICROFRONTEND_LIB_LOG_SCALE_IO_H_ - -#include - -#include "tensorflow/contrib/lite/experimental/microfrontend/lib/log_scale.h" - -#ifdef __cplusplus -extern "C" { -#endif - -void LogScaleWriteMemmap(FILE* fp, const struct LogScaleState* state, - const char* variable); - -#ifdef __cplusplus -} // extern "C" -#endif - -#endif // TENSORFLOW_CONTRIB_LITE_EXPERIMENTAL_MICROFRONTEND_LIB_LOG_SCALE_IO_H_ diff --git a/tensorflow/contrib/lite/experimental/writer/BUILD b/tensorflow/contrib/lite/experimental/writer/BUILD deleted file mode 100644 index 82d39c00abd27d9931131317e9750bbf7face981..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/experimental/writer/BUILD +++ /dev/null @@ -1,66 +0,0 @@ -package(default_visibility = [ - "//visibility:public", -]) - -licenses(["notice"]) # Apache 2.0 - -cc_binary( - name = "option_writer_generator", - srcs = ["option_writer_generator.cc"], - deps = [ - "//tensorflow/contrib/lite/schema:schema_fbs_with_reflection", - "@flatbuffers", - ], -) - -cc_library( - name = "writer_lib", - srcs = [ - "enum_mapping.h", - "writer_lib.cc", - ], - hdrs = [ - "writer_lib.h", - ], - data = [ - ":option_writer_gen", - ], - textual_hdrs = ["option_writer_generated.h"], - deps = [ - "//tensorflow/contrib/lite:builtin_op_data", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite:schema_fbs_version", - "//tensorflow/contrib/lite/kernels:builtin_ops", - "//tensorflow/contrib/lite/schema:schema_fbs_with_reflection", - ], -) - -cc_binary( - name = "writer", - srcs = ["writer.cc"], - deps = [ - ":writer_lib", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:builtin_ops", - ], -) - -cc_test( - name = "writer_lib_test", - size = "small", - srcs = ["writer_lib_test.cc"], - deps = [ - ":writer_lib", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:builtin_ops", - "//tensorflow/contrib/lite/testing:util", - "@com_google_googletest//:gtest", - ], -) - -genrule( - name = "option_writer_gen", - outs = ["option_writer_generated.h"], - cmd = "$(location :option_writer_generator) $(@)", - tools = [":option_writer_generator"], -) diff --git a/tensorflow/contrib/lite/g3doc/_book.yaml b/tensorflow/contrib/lite/g3doc/_book.yaml deleted file mode 100644 index 05c65441c3db4e74b6e7834437fa9cd0633af636..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/g3doc/_book.yaml +++ /dev/null @@ -1,73 +0,0 @@ -upper_tabs: -# Tabs left of dropdown menu -- include: /_upper_tabs_left.yaml -- include: /versions/_upper_tabs_versions.yaml -# Dropdown menu -- name: Ecosystem - path: /ecosystem - is_default: true - menu: - - include: /ecosystem/_menu_toc.yaml - lower_tabs: - # Subsite tabs - other: - - name: Guide - contents: - - title: Overview - path: /lite/overview - - title: Developer guide - path: /lite/devguide - - title: Android demo app - path: /lite/demo_android - - title: iOS demo app - path: /lite/demo_ios - - title: Performance - path: /lite/performance - - break: true - - title: TensorFlow Lite APIs - path: /lite/apis - - title: Custom operators - path: /lite/custom_operators - - title: TensorFlow Lite ops versioning - path: /lite/ops_versioning - - title: TensorFlow Lite compatibility guide - path: /lite/tf_ops_compatibility - - title: List of hosted models - path: /lite/models - - title: TensorFlow Lite for iOS - path: /lite/ios - - title: TensorFlow Lite for Raspberry Pi - path: /lite/rpi - - - heading: TF Lite converter - - title: Overview - path: /lite/convert/ - - title: Python API guide - path: /lite/convert/python_api - - title: Command line examples - path: /lite/convert/cmdline_examples - - title: Command line reference - path: /lite/convert/cmdline_reference - - - title: TF Mobile - style: accordion - status: deprecated - section: - - title: Overview - path: /lite/tfmobile/ - - title: Building TensorFlow on Android - path: /lite/tfmobile/android_build - - title: Building TensorFlow on IOS - path: /lite/tfmobile/ios_build - - title: Integrating TensorFlow libraries - path: /lite/tfmobile/linking_libs - - title: Preparing models for mobile deployment - path: /lite/tfmobile/prepare_models - - title: Optimizing for mobile - path: /lite/tfmobile/optimizing - - - name: API - skip_translation: true - contents: - - title: API - path: /api_docs/python/tf/contrib/lite diff --git a/tensorflow/contrib/lite/g3doc/_index.yaml b/tensorflow/contrib/lite/g3doc/_index.yaml deleted file mode 100644 index 44ee6ba7505d421e46c8806ea5ca0ed4bc07f147..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/g3doc/_index.yaml +++ /dev/null @@ -1,218 +0,0 @@ -project_path: /lite/_project.yaml -book_path: /lite/_book.yaml -description: -landing_page: - custom_css_path: /site-assets/css/style.css - rows: - - heading: TensorFlow Lite is for mobile and embedded devices. - description: > -

- TensorFlow Lite is the official solution for running machine learning - models on mobile and embedded devices. It enables on‑device machine - learning inference with low latency and a small binary size on Android, - iOS, and other operating systems. -

- - - - classname: tfo-landing-row-heading tfo-landing-row-heading-list - heading: Many benefits - description: > - On-device ML inference is difficult because of the many constraints—TensorFlow Lite can solve these: - items: - - list: - - heading: Performance - description: > - TF Lite is fast with no noticeable accuracy loss—see the metrics. - icon: - icon_name: lens - foreground: theme - - heading: Portability - description: > - Android, - iOS, and more specialized IoT devices. - icon: - icon_name: lens - foreground: theme - - list: - - heading: Low latency - description: > - Optimized float- and fixed-point CPU kernels, op‑fusing, and more. - icon: - icon_name: lens - foreground: theme - - heading: Acceleration - description: > - Integration with GPU and internal/external accelerators. - icon: - icon_name: lens - foreground: theme - - list: - - heading: Small model size - description: > - Controlled dependencies, quantization, - and op registration. - icon: - icon_name: lens - foreground: theme - - heading: Tooling - description: > - Conversion, compression, benchmarking, power-consumption, and more. - icon: - icon_name: lens - foreground: theme - - - classname: devsite-landing-row-logos tfo-landing-row-heading - heading: Companies using TensorFlow Lite - items: - - custom_image: - path: ./images/landing-page/photos_logo.png - path: https://www.photos.google.com - - custom_image: - path: ./images/landing-page/gboard_logo.png - path: https://play.google.com/store/apps/details?id=com.google.android.inputmethod.latin&hl=en_US - - custom_image: - path: ./images/landing-page/gmail_logo.png - path: https://www.google.com/gmail/ - - custom_image: - path: ./images/landing-page/assistant_logo.png - path: https://assistant.google.com/ - - - classname: devsite-landing-row-logos - items: - - custom_image: - path: ./images/landing-page/vsco_logo.png - path: https://vsco.co - - custom_image: - path: ./images/landing-page/shazam_logo.png - path: https://www.shazam.com/ - - custom_image: - path: ./images/landing-page/nest_logo.png - path: https://nest.com/ - - custom_image: - path: ./images/landing-page/loseit_logo.png - path: https://www.loseit.com/ - - - classname: devsite-landing-row-no-image-background devsite-landing-row-67 - background: grey - items: - - description: > - “TensorFlow Lite helped us introduce machine learning and AI into our - app in an easy and streamlined way. We could reduce the size of our - models while keeping the accuracy high. This helped us create an amazing - fishing experience for our users by allowing them to identify any fish - species with just a photo.” - image_path: ./images/landing-page/fishbrain_logo_big.png - - - heading: How it works - items: - - heading: Build - icon: - icon_name: build - description: > - Build a new model or retrain an existing one, such as using transfer learning. - buttons: - - label: Read the developer guide - path: /lite/devguide - classname: button button-primary tfo-button-primary - - heading: Convert - icon: - icon_name: autorenew - description: > - Convert a TensorFlow model into a compressed flat buffer with the - TensorFlow Lite Converter. - buttons: - - label: Read the converter guide - path: /lite/convert/ - classname: button button-primary tfo-button-primary - - heading: Deploy - icon: - icon_name: bolt - description: > - Take the compressed .tflite file and load it into a mobile - or embedded device.
- See the tutorials below to build an app. - - - heading: Build your first TensorFlow Lite app - background: grey - items: - - classname: tfo-landing-row-item-inset-white - heading: Get started - description: > - - - classname: tfo-landing-row-item-inset-white - heading: Share your TensorFlow Lite story - description: > - We love to hear what you're working on—it may even get highlighted on - our social media! Tell us. - - - classname: devsite-landing-row-no-image-background devsite-landing-row-67 - items: - - description: > -

- “The release of TensorFlow Lite has allowed us to deploy an engaging - real-time experience to our users that eliminates the requirement - for a data connection. TensorFlow Lite’s ability to compress and - optimize the TensorFlow graph for mobile deployment has been - transformative in expanding the capabilities of Snap It. -

-

- Through TensorFlow Lite, our users can now enjoy a state of the - art, computer-vision-based food logging experience without worrying - about signal strength. We look forward to future collaborations - with the TensorFlow Lite team.” -

- image_path: ./images/landing-page/loseit_logo_big.png - - - classname: devsite-landing-row-cards - background: grey - heading: Updates - items: - - heading: Introducing the Model Optimization Toolkit - image_path: /ecosystem/images/tf-logo-card-16x9.png - path: https://medium.com/tensorflow/introducing-the-model-optimization-toolkit-for-tensorflow-254aca1ba0a3 - buttons: - - label: Read on TensorFlow blog - path: https://medium.com/tensorflow/introducing-the-model-optimization-toolkit-for-tensorflow-254aca1ba0a3 - - heading: East Africa Cassava App - image_path: ./images/landing-page/detect_crop_disease_in_africa.png - path: https://heartbeat.fritz.ai/community-spotlight-nuru-a-mobile-app-by-plantvillage-to-detect-crop-disease-in-africa-28d142bf63d5 - buttons: - - label: Read more - path: https://heartbeat.fritz.ai/community-spotlight-nuru-a-mobile-app-by-plantvillage-to-detect-crop-disease-in-africa-28d142bf63d5 - - heading: Using TensorFlow Lite on Android - image_path: /ecosystem/images/tf-logo-card-16x9.png - path: https://medium.com/tensorflow/using-tensorflow-lite-on-android-9bbc9cb7d69d - buttons: - - label: Read on TensorFlow blog - path: https://medium.com/tensorflow/using-tensorflow-lite-on-android-9bbc9cb7d69d - - - classname: devsite-landing-row-cards - background: grey - items: - - heading: TensorFlow Lite at the Dev Summit - youtube_id: FAMfy7izB6A - buttons: - - label: Watch the video - path: https://www.youtube.com/watch?v=FAMfy7izB6A - - heading: TensorFlow Lite on GitHub - image_path: /ecosystem/images/github-card-16x9.png - path: https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite - buttons: - - label: View on GitHub - path: https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite - - classname: devsite-landing-row-item-hidden diff --git a/tensorflow/contrib/lite/g3doc/convert/index.md b/tensorflow/contrib/lite/g3doc/convert/index.md deleted file mode 100644 index bc92a1c1a11a6f3808e44f37d04704ece1627fc3..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/g3doc/convert/index.md +++ /dev/null @@ -1,19 +0,0 @@ -# TensorFlow Lite Converter - -The TensorFlow Lite Converter takes a TensorFlow graph file and creates a graph -file used by the TensorFlow Lite interpreter. - -## From model training to device deployment - -After a TensorFlow model is trained, the TensorFlow Lite converter uses that -model to generate a TensorFlow Lite [FlatBuffer](https://google.github.io/flatbuffers/) -file (`.tflite`). The converter supports as input: -[SavedModels](https://www.tensorflow.org/guide/saved_model#using_savedmodel_with_estimators), -frozen graphs (models generated by -[freeze_graph.py](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/tools/freeze_graph.py)), -and `tf.keras` models. The TensorFlow Lite `FlatBuffer` file is deployed to a -client device (generally a mobile or embedded device), and the TensorFlow Lite -interpreter uses the compressed model for on-device inference. This conversion -process is shown in the diagram below: - -![TFLite converter workflow](../images/convert/workflow.svg) diff --git a/tensorflow/contrib/lite/g3doc/demo_ios.md b/tensorflow/contrib/lite/g3doc/demo_ios.md deleted file mode 100644 index 7579ad84a049ec592aafb16ce95a4b703ac78c5a..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/g3doc/demo_ios.md +++ /dev/null @@ -1,69 +0,0 @@ - -# iOS Demo App - -The TensorFlow Lite demo is a camera app that continuously classifies whatever -it sees from your device's back camera, using a quantized MobileNet model. These -instructions walk you through building and running the demo on an iOS device. - -## Prerequisites - -* You must have [Xcode](https://developer.apple.com/xcode/) installed and have a - valid Apple Developer ID, and have an iOS device set up and linked to your - developer account with all of the appropriate certificates. For these - instructions, we assume that you have already been able to build and deploy an - app to an iOS device with your current developer environment. - -* The demo app requires a camera and must be executed on a real iOS device. You - can build it and run with the iPhone Simulator but it won't have any camera - information to classify. - -* You don't need to build the entire TensorFlow library to run the demo, but you - will need to clone the TensorFlow repository if you haven't already: - - git clone https://github.com/tensorflow/tensorflow - -* You'll also need the Xcode command-line tools: - - xcode-select --install - - If this is a new install, you will need to run the Xcode application once to - agree to the license before continuing. - -## Building the iOS Demo App - -1. Install CocoaPods if you don't have it: - - sudo gem install cocoapods - -2. Download the model files used by the demo app (this is done from inside the - cloned directory): - - sh tensorflow/contrib/lite/examples/ios/download_models.sh - -3. Install the pod to generate the workspace file: - - cd tensorflow/contrib/lite/examples/ios/camera - pod install - - If you have installed this pod before and that command doesn't work, try - - pod update - - At the end of this step you should have a file called - `tflite_camera_example.xcworkspace`. - -4. Open the project in Xcode by typing this on the command line: - - open tflite_camera_example.xcworkspace - - This launches Xcode if it isn't open already and opens the - `tflite_camera_example` project. - -5. Build and run the app in Xcode. - - Note that as mentioned earlier, you must already have a device set up and - linked to your Apple Developer account in order to deploy the app on a - device. - -You'll have to grant permissions for the app to use the device's camera. Point -the camera at various objects and enjoy seeing how the model classifies things! diff --git a/tensorflow/contrib/lite/g3doc/devguide.md b/tensorflow/contrib/lite/g3doc/devguide.md deleted file mode 100644 index 0eed5160009c07727f0c2985ebe963efc7bb9d8e..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/g3doc/devguide.md +++ /dev/null @@ -1,232 +0,0 @@ -# TF Lite Developer Guide - -Using a TensorFlow Lite model in your mobile app requires multiple -considerations: you must choose a pre-trained or custom model, convert the model -to a TensorFLow Lite format, and finally, integrate the model in your app. - -## 1. Choose a model - -Depending on the use case, you can choose one of the popular open-sourced models, -such as *InceptionV3* or *MobileNets*, and re-train these models with a custom -data set or even build your own custom model. - -### Use a pre-trained model - -[MobileNets](https://research.googleblog.com/2017/06/mobilenets-open-source-models-for.html) -is a family of mobile-first computer vision models for TensorFlow designed to -effectively maximize accuracy, while taking into consideration the restricted -resources for on-device or embedded applications. MobileNets are small, -low-latency, low-power models parameterized to meet the resource constraints for -a variety of uses. They can be used for classification, detection, embeddings, and -segmentation—similar to other popular large scale models, such as -[Inception](https://arxiv.org/pdf/1602.07261.pdf). Google provides 16 pre-trained -[ImageNet](http://www.image-net.org/challenges/LSVRC/) classification checkpoints -for MobileNets that can be used in mobile projects of all sizes. - -[Inception-v3](https://arxiv.org/abs/1512.00567) is an image recognition model -that achieves fairly high accuracy recognizing general objects with 1000 classes, -for example, "Zebra", "Dalmatian", and "Dishwasher". The model extracts general -features from input images using a convolutional neural network and classifies -them based on those features with fully-connected and softmax layers. - -[On Device Smart Reply](https://research.googleblog.com/2017/02/on-device-machine-intelligence.html) -is an on-device model that provides one-touch replies for incoming text messages -by suggesting contextually relevant messages. The model is built specifically for -memory constrained devices, such as watches and phones, and has been successfully -used in Smart Replies on Android Wear. Currently, this model is Android-specific. - -These pre-trained models are [available for download](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/lite/g3doc/models.md) - -### Re-train Inception-V3 or MobileNet for a custom data set - -These pre-trained models were trained on the *ImageNet* data set which contains -1000 predefined classes. If these classes are not sufficient for your use case, -the model will need to be re-trained. This technique is called -*transfer learning* and starts with a model that has been already trained on a -problem, then retrains the model on a similar problem. Deep learning from -scratch can take days, but transfer learning is fairly quick. In order to do -this, you need to generate a custom data set labeled with the relevant classes. - -The [TensorFlow for Poets](https://codelabs.developers.google.com/codelabs/tensorflow-for-poets/) -codelab walks through the re-training process step-by-step. The code supports -both floating point and quantized inference. - -### Train a custom model - -A developer may choose to train a custom model using Tensorflow (see the -[TensorFlow tutorials](../tutorials/) for examples of building and training -models). If you have already written a model, the first step is to export this -to a `tf.GraphDef` file. This is required because some formats do not store the -model structure outside the code, and we must communicate with other parts of the -framework. See -[Exporting the Inference Graph](https://github.com/tensorflow/models/blob/master/research/slim/README.md) -to create .pb file for the custom model. - -TensorFlow Lite currently supports a subset of TensorFlow operators. Refer to the -[TensorFlow Lite & TensorFlow Compatibility Guide](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite/g3doc/tf_ops_compatibility.md) -for supported operators and their usage. This set of operators will continue to -grow in future Tensorflow Lite releases. - - -## 2. Convert the model format - -The model generated (or downloaded) in the previous step is a *standard* -Tensorflow model and you should now have a .pb or .pbtxt `tf.GraphDef` file. -Models generated with transfer learning (re-training) or custom models must be -converted—but, we must first freeze the graph to convert the model to the -Tensorflow Lite format. This process uses several model formats: - -* `tf.GraphDef` (.pb) —A protobuf that represents the TensorFlow training or - computation graph. It contains operators, tensors, and variables definitions. -* *CheckPoint* (.ckpt) —Serialized variables from a TensorFlow graph. Since this - does not contain a graph structure, it cannot be interpreted by itself. -* `FrozenGraphDef` —A subclass of `GraphDef` that does not contain - variables. A `GraphDef` can be converted to a `FrozenGraphDef` by taking a - CheckPoint and a `GraphDef`, and converting each variable into a constant - using the value retrieved from the CheckPoint. -* `SavedModel` —A `GraphDef` and CheckPoint with a signature that labels - input and output arguments to a model. A `GraphDef` and CheckPoint can be - extracted from a `SavedModel`. -* *TensorFlow Lite model* (.tflite) —A serialized - [FlatBuffer](https://google.github.io/flatbuffers/) that contains TensorFlow - Lite operators and tensors for the TensorFlow Lite interpreter, similar to a - `FrozenGraphDef`. - -### Freeze Graph - -To use the `GraphDef` .pb file with TensorFlow Lite, you must have checkpoints -that contain trained weight parameters. The .pb file only contains the structure -of the graph. The process of merging the checkpoint values with the graph -structure is called *freezing the graph*. - -You should have a checkpoints folder or download them for a pre-trained model -(for example, -[MobileNets](https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet_v1.md)). - -To freeze the graph, use the following command (changing the arguments): - -``` -freeze_graph --input_graph=/tmp/mobilenet_v1_224.pb \ - --input_checkpoint=/tmp/checkpoints/mobilenet-10202.ckpt \ - --input_binary=true \ - --output_graph=/tmp/frozen_mobilenet_v1_224.pb \ - --output_node_names=MobileNetV1/Predictions/Reshape_1 -``` - -The `input_binary` flag must be enabled so the protobuf is read and written in -a binary format. Set the `input_graph` and `input_checkpoint` files. - -The `output_node_names` may not be obvious outside of the code that built the -model. The easiest way to find them is to visualize the graph, either with -[TensorBoard](https://codelabs.developers.google.com/codelabs/tensorflow-for-poets-2/#3) -or `graphviz`. - -The frozen `GraphDef` is now ready for conversion to the `FlatBuffer` format -(.tflite) for use on Android or iOS devices. For Android, the Tensorflow -Optimizing Converter tool supports both float and quantized models. To convert -the frozen `GraphDef` to the .tflite format: - -``` -toco --input_file=$(pwd)/mobilenet_v1_1.0_224/frozen_graph.pb \ - --input_format=TENSORFLOW_GRAPHDEF \ - --output_format=TFLITE \ - --output_file=/tmp/mobilenet_v1_1.0_224.tflite \ - --inference_type=FLOAT \ - --input_type=FLOAT \ - --input_arrays=input \ - --output_arrays=MobilenetV1/Predictions/Reshape_1 \ - --input_shapes=1,224,224,3 -``` - -The `input_file` argument should reference the frozen `GraphDef` file -containing the model architecture. The [frozen_graph.pb](https://storage.googleapis.com/download.tensorflow.org/models/mobilenet_v1_1.0_224_frozen.tgz) -file used here is available for download. `output_file` is where the TensorFlow -Lite model will get generated. The `input_type` and `inference_type` -arguments should be set to `FLOAT`, unless converting a -quantized model. -Setting the `input_array`, `output_array`, and `input_shape` arguments are not as -straightforward. The easiest way to find these values is to explore the graph -using Tensorboard. Reuse the arguments for specifying the output nodes for -inference in the `freeze_graph` step. - -It is also possible to use the Tensorflow Optimizing Converter with protobufs -from either Python or from the command line (see the -[toco_from_protos.py](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite/toco/python/toco_from_protos.py) -example). This allows you to integrate the conversion step into the model design -workflow, ensuring the model is easily convertible to a mobile inference graph. -For example: - -```python -import tensorflow as tf - -img = tf.placeholder(name="img", dtype=tf.float32, shape=(1, 64, 64, 3)) -val = img + tf.constant([1., 2., 3.]) + tf.constant([1., 4., 4.]) -out = tf.identity(val, name="out") - -with tf.Session() as sess: - tflite_model = tf.contrib.lite.toco_convert(sess.graph_def, [img], [out]) - open("converteds_model.tflite", "wb").write(tflite_model) -``` - -For usage, see the Tensorflow Optimizing Converter -[command-line examples](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite/toco/g3doc/cmdline_examples.md). - -Refer to the -[Ops compatibility guide](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite/g3doc/tf_ops_compatibility.md) -for troubleshooting help, and if that doesn't help, please -[file an issue](https://github.com/tensorflow/tensorflow/issues). - -The [development repo](https://github.com/tensorflow/tensorflow) contains a tool -to visualize TensorFlow Lite models after conversion. To build the -[visualize.py](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/lite/tools/visualize.py) -tool: - -```sh -bazel run tensorflow/contrib/lite/tools:visualize -- model.tflite model_viz.html -``` - -This generates an interactive HTML page listing subgraphs, operations, and a -graph visualization. - - -## 3. Use the TensorFlow Lite model for inference in a mobile app - -After completing the prior steps, you should now have a `.tflite` model file. - -### Android - -Since Android apps are written in Java and the core TensorFlow library is in C++, -a JNI library is provided as an interface. This is only meant for inference—it -provides the ability to load a graph, set up inputs, and run the model to -calculate outputs. - -The open source Android demo app uses the JNI interface and is available -[on GitHub](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite/java/demo/app). -You can also download a -[prebuilt APK](http://download.tensorflow.org/deps/tflite/TfLiteCameraDemo.apk). -See the Android demo guide for details. - -The Android mobile guide has instructions for -installing TensorFlow on Android and setting up `bazel` and Android Studio. - -### iOS - -To integrate a TensorFlow model in an iOS app, see the -[TensorFlow Lite for iOS](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite/g3doc/ios.md) -guide and iOS demo guide. - -#### Core ML support - -Core ML is a machine learning framework used in Apple products. In addition to -using Tensorflow Lite models directly in your applications, you can convert -trained Tensorflow models to the -[CoreML](https://developer.apple.com/machine-learning/) format for use on Apple -devices. To use the converter, refer to the -[Tensorflow-CoreML converter documentation](https://github.com/tf-coreml/tf-coreml). - -### Raspberry Pi - -Compile Tensorflow Lite for a Raspberry Pi by following the -[RPi build instructions](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/lite/g3doc/rpi.md) -This compiles a static library file (`.a`) used to build your app. There are -plans for Python bindings and a demo app. diff --git a/tensorflow/contrib/lite/g3doc/images/convert/sample_after.png b/tensorflow/contrib/lite/g3doc/images/convert/sample_after.png deleted file mode 100644 index 6c451f97903f7f70a9f28dee8abf6daeb7ec5693..0000000000000000000000000000000000000000 Binary files a/tensorflow/contrib/lite/g3doc/images/convert/sample_after.png and /dev/null differ diff --git a/tensorflow/contrib/lite/g3doc/images/convert/sample_before.png b/tensorflow/contrib/lite/g3doc/images/convert/sample_before.png deleted file mode 100644 index e5317ef295062e79c66430512ef1c45925858ce0..0000000000000000000000000000000000000000 Binary files a/tensorflow/contrib/lite/g3doc/images/convert/sample_before.png and /dev/null differ diff --git a/tensorflow/contrib/lite/g3doc/ios.md b/tensorflow/contrib/lite/g3doc/ios.md deleted file mode 100644 index 3b9fcca8117dc1859d075ae5f048cfc9f0d988a3..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/g3doc/ios.md +++ /dev/null @@ -1,86 +0,0 @@ - -# Build TensorFlow Lite for iOS - -This document describes how to build TensorFlow Lite iOS library. If you just -want to use it, the easiest way is using the TensorFlow Lite CocoaPod releases. -See [TensorFlow Lite iOS Demo](demo_ios.md) for examples. - - -## Building - -To create a universal iOS library for TensorFlow Lite, you need to build it -using Xcode's command line tools on a MacOS machine. If you have not already, -you will need to install Xcode 8 or later and the tools using `xcode-select`: - -```bash -xcode-select --install -``` - -If this is a new install, you will need to run XCode once to agree to the -license before continuing. - -(You will also need to have [Homebrew](http://brew.sh/) installed.) - -Then install -[automake](https://en.wikipedia.org/wiki/Automake)/[libtool](https://en.wikipedia.org/wiki/GNU_Libtool): - -```bash -brew install automake -brew install libtool -``` -If you get an error where either automake or libtool install but do not link correctly, you'll first need to: -```bash -sudo chown -R $(whoami) /usr/local/* -``` -Then follow the instructions to perform the linking: -```bash -brew link automake -brew link libtool -``` - -Then you need to run a shell script to download the dependencies you need: - -```bash -tensorflow/contrib/lite/tools/make/download_dependencies.sh -``` - -This will fetch copies of libraries and data from the web and install them in -`tensorflow/contrib/lite/downloads`. - -With all of the dependencies set up, you can now build the library for all five -supported architectures on iOS: - -```bash -tensorflow/contrib/lite/tools/make/build_ios_universal_lib.sh -``` - -Under the hood this uses a makefile in `tensorflow/contrib/lite` to build the -different versions of the library, followed by a call to `lipo` to bundle them -into a universal file containing armv7, armv7s, arm64, i386, and x86_64 -architectures. The resulting library is in -`tensorflow/contrib/lite/tools/make/gen/lib/libtensorflow-lite.a`. - -If you get an error such as `no such file or directory: 'x86_64'` when running -`build_ios_universal_lib.sh`: open Xcode > Preferences > Locations, and ensure -a value is selected in the "Command Line Tools" dropdown. - -## Using in your own application - -You'll need to update various settings in your app to link against TensorFlow -Lite. You can view them in the example project at -`tensorflow/contrib/lite/examples/ios/simple/simple.xcodeproj` but here's a full -rundown: - -- You'll need to add the library at - `tensorflow/contrib/lite/gen/lib/libtensorflow-lite.a` to your linking build - stage, and in Search Paths add `tensorflow/contrib/lite/gen/lib` to the - Library Search Paths setting. - -- The _Header Search_ paths needs to contain: - - - the root folder of tensorflow, - - `tensorflow/contrib/lite/downloads` - - `tensorflow/contrib/lite/downloads/flatbuffers/include` - -- C++11 support (or later) should be enabled by setting `C++ Language Dialect` - to `GNU++11` (or `GNU++14`), and `C++ Standard Library` to `libc++`. diff --git a/tensorflow/contrib/lite/g3doc/overview.md b/tensorflow/contrib/lite/g3doc/overview.md deleted file mode 100644 index 9d035a69211d7ced913e6d16061c6ad8ca912e64..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/g3doc/overview.md +++ /dev/null @@ -1,202 +0,0 @@ - -# Introduction to TensorFlow Lite - -TensorFlow Lite is TensorFlow’s lightweight solution for mobile and embedded -devices. It enables on-device machine learning inference with low latency and a -small binary size. TensorFlow Lite also supports hardware acceleration with the -[Android Neural Networks -API](https://developer.android.com/ndk/guides/neuralnetworks/index.html). - -TensorFlow Lite uses many techniques for achieving low latency such as -optimizing the kernels for mobile apps, pre-fused activations, and quantized -kernels that allow smaller and faster (fixed-point math) models. - -Most of our TensorFlow Lite documentation is [on -GitHub](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite) -for the time being. - -## What does TensorFlow Lite contain? - -TensorFlow Lite supports a set of core operators, both quantized and -float, which have been tuned for mobile platforms. They incorporate pre-fused -activations and biases to further enhance performance and quantized -accuracy. Additionally, TensorFlow Lite also supports using custom operations in -models. - -TensorFlow Lite defines a new model file format, based on -[FlatBuffers](https://google.github.io/flatbuffers/). FlatBuffers is an -efficient open-source cross-platform serialization library. It is similar to -[protocol buffers](https://developers.google.com/protocol-buffers/?hl=en), but -the primary difference is that FlatBuffers does not need a parsing/unpacking -step to a secondary representation before you can access data, often coupled -with per-object memory allocation. Also, the code footprint of FlatBuffers is an -order of magnitude smaller than protocol buffers. - -TensorFlow Lite has a new mobile-optimized interpreter, which has the key goals -of keeping apps lean and fast. The interpreter uses a static graph ordering and -a custom (less-dynamic) memory allocator to ensure minimal load, initialization, -and execution latency. - -TensorFlow Lite provides an interface to leverage hardware acceleration, if -available on the device. It does so via the -[Android Neural Networks API](https://developer.android.com/ndk/guides/neuralnetworks/index.html), -available on Android 8.1 (API level 27) and higher. - -## Why do we need a new mobile-specific library? - -Machine Learning is changing the computing paradigm, and we see an emerging -trend of new use cases on mobile and embedded devices. Consumer expectations are -also trending toward natural, human-like interactions with their devices, driven -by the camera and voice interaction models. - -There are several factors which are fueling interest in this domain: - -- Innovation at the silicon layer is enabling new possibilities for hardware - acceleration, and frameworks such as the Android Neural Networks API make it - easy to leverage these. - -- Recent advances in real-time computer-vision and spoken language understanding - have led to mobile-optimized benchmark models being open sourced - (e.g. MobileNets, SqueezeNet). - -- Widely-available smart appliances create new possibilities for - on-device intelligence. - -- Interest in stronger user data privacy paradigms where user data does not need - to leave the mobile device. - -- Ability to serve ‘offline’ use cases, where the device does not need to be - connected to a network. - -We believe the next wave of machine learning applications will have significant -processing on mobile and embedded devices. - -## TensorFlow Lite highlights - -TensorFlow Lite provides: - -- A set of core operators, both quantized and float, many of which have been - tuned for mobile platforms. These can be used to create and run custom - models. Developers can also write their own custom operators and use them in - models. - -- A new [FlatBuffers](https://google.github.io/flatbuffers/)-based - model file format. - -- On-device interpreter with kernels optimized for faster execution on mobile. - -- TensorFlow converter to convert TensorFlow-trained models to the TensorFlow - Lite format. - -- Smaller in size: TensorFlow Lite is smaller than 300KB when all supported - operators are linked and less than 200KB when using only the operators needed - for supporting InceptionV3 and Mobilenet. - -- **Pre-tested models:** - - All of the following models are guaranteed to work out of the box: - - - Inception V3, a popular model for detecting the dominant objects - present in an image. - - - [MobileNets](https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet_v1.md), - a family of mobile-first computer vision models designed to effectively - maximize accuracy while being mindful of the restricted resources for an - on-device or embedded application. They are small, low-latency, low-power - models parameterized to meet the resource constraints of a variety of use - cases. They can be built upon for classification, detection, embeddings - and segmentation. MobileNet models are smaller but [lower in - accuracy](https://research.googleblog.com/2017/06/mobilenets-open-source-models-for.html) - than Inception V3. - - - On Device Smart Reply, an on-device model which provides one-touch - replies for an incoming text message by suggesting contextually relevant - messages. The model was built specifically for memory constrained devices - such as watches & phones and it has been successfully used to surface - [Smart Replies on Android - Wear](https://research.googleblog.com/2017/02/on-device-machine-intelligence.html) - to all first-party and third-party apps. - - Also see the complete list of - [TensorFlow Lite's supported models](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/lite/g3doc/models.md), - including the model sizes, performance numbers, and downloadable model files. - -- Quantized versions of the MobileNet model, which runs faster than the - non-quantized (float) version on CPU. - -- New Android demo app to illustrate the use of TensorFlow Lite with a quantized - MobileNet model for object classification. - -- Java and C++ API support - - -## Getting Started - -We recommend you try out TensorFlow Lite with the pre-tested models indicated -above. If you have an existing model, you will need to test whether your model -is compatible with both the converter and the supported operator set. To test -your model, see the -[documentation on GitHub](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite). - -### Retrain Inception-V3 or MobileNet for a custom data set - -The pre-trained models mentioned above have been trained on the ImageNet data -set, which consists of 1000 predefined classes. If those classes are not -relevant or useful for your use case, you will need to retrain those -models. This technique is called transfer learning, which starts with a model -that has been already trained on a problem and will then be retrained on a -similar problem. Deep learning from scratch can take days, but transfer learning -can be done fairly quickly. In order to do this, you'll need to generate your -custom data set labeled with the relevant classes. - -The [TensorFlow for Poets](https://codelabs.developers.google.com/codelabs/tensorflow-for-poets/) -codelab walks through this process step-by-step. The retraining code supports -retraining for both floating point and quantized inference. - -## TensorFlow Lite Architecture - -The following diagram shows the architectural design of TensorFlow Lite: - -TensorFlow Lite architecture diagram - -Starting with a trained TensorFlow model on disk, you'll convert that model to -the TensorFlow Lite file format (`.tflite`) using the TensorFlow Lite -Converter. Then you can use that converted file in your mobile application. - -Deploying the TensorFlow Lite model file uses: - -- Java API: A convenience wrapper around the C++ API on Android. - -- C++ API: Loads the TensorFlow Lite Model File and invokes the Interpreter. The - same library is available on both Android and iOS. - -- Interpreter: Executes the model using a set of kernels. The interpreter - supports selective kernel loading; without kernels it is only 100KB, and 300KB - with all the kernels loaded. This is a significant reduction from the 1.5M - required by TensorFlow Mobile. - -- On select Android devices, the Interpreter will use the Android Neural - Networks API for hardware acceleration, or default to CPU execution if none - are available. - -You can also implement custom kernels using the C++ API that can be used by the -Interpreter. - -## Future Work - -In future releases, TensorFlow Lite will support more models and built-in -operators, contain performance improvements for both fixed point and floating -point models, improvements to the tools to enable easier developer workflows and -support for other smaller devices and more. As we continue development, we hope -that TensorFlow Lite will greatly simplify the developer experience of targeting -a model for small devices. - -Future plans include using specialized machine learning hardware to get the best -possible performance for a particular model on a particular device. - -## Next Steps - -The TensorFlow Lite [GitHub repository](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite). -contains additional docs, code samples, and demo applications. diff --git a/tensorflow/contrib/lite/g3doc/performance.md b/tensorflow/contrib/lite/g3doc/performance.md deleted file mode 100644 index ed114527166da79dba2d92c3ffad78e9885f9e94..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/g3doc/performance.md +++ /dev/null @@ -1,45 +0,0 @@ - -# Performance best practices - -Mobile and embedded devices have limited computational resources and it is important to keep your application resource efficient. We have compiled a list of best practices and strategies you can use to optimize your model and application when using Tensorflow Lite. - -## Choose the best model for the task -Depending on the task you will need to make a tradeoff between model complexity and size. If your task requires high accuracy then you may need a large and complex model. Some tasks may work with a less precise model, for these tasks it is better to use a smaller but less precise model. Smaller models not only use less disk space and memory but are generally faster and more energy efficient. For example, graphs below show accuracy and latency tradeoff for some common image classification models. - -![accuracy vs model size](images/performance/model_size_vs_accuracy.png "Accuracy vs Model size") - - -![latency vs model size](images/performance/model_size_vs_latency.png "Latency vs Model size") - -One example of models optimized for mobile devices are [MobileNets](https://arxiv.org/abs/1704.04861), which are optimized for mobile vision applications. Tensorflow Lite [models page](models.md) lists several other models that have been optimized specifically for mobile and embedded devices. - -You can retrain the listed models on your own dataset by using transfer learning. Check out our transfer learning tutorial for -[image classification](https://codelabs.developers.google.com/codelabs/tensorflow-for-poets/#0) and - [object detection](https://medium.com/tensorflow/training-and-serving-a-realtime-mobile-object-detector-in-30-minutes-with-cloud-tpus-b78971cf1193). - - -## Profile your model -Once you have selected a candidate model that is right for your task, it is a good practice to profile and benchmark your model. Tensorflow Lite [benchmarking tool](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite/tools/benchmark) has a built-in profiler that shows per operator profiling statistics. This can help in understanding performance bottlenecks and which operators dominate the computation time. - -## Profile and optimize operators in the graph -If a particular operator appears frequently in the model and based on profiling you find the operator consuming the most amount of time, you can look into optimizing the operator. - This scenario should be rare as Tensorflow Lite has optimized versions for most ops. However you may be able to write a faster version of a custom op, if you know the constraints in which the operator is executed. Check out our [custom operator documentation](custom_operators.md). - -## Quantize your model -If your model uses floating point weights or activations then it may be possible to reduce the size of model up to ~4x by using quantization and other model optimizations. Check out our [model optimization toolkit](https://www.tensorflow.org/performance/model_optimization) for details about optimizing your model. - -## Tweak the number of threads -Tensorflow Lite supports multi-threaded kernels for many operators. You can increase the number of threads and speed up execution of operators. Increasing the number of threads will however make your model use more resources and power. For some applications latency may be more important than energy efficiency. You can increase the number of threads by setting the number of [interpreter](https://github.com/tensorflow/tensorflow/blob/1084594657a5d139102ac794f84d1427a710e39a/tensorflow/contrib/lite/interpreter.h#L337) threads. Multi-threaded execution however comes at the cost of increased performance variability depending on what else is been executed concurrently. This is particularly the case for mobile apps. For example, isolated tests may show 2x speed up vs single-threaded but if another app is executing at the same time may result in worst performance than single-threaded. - -## Eliminate redundant copies -If your application is not careful, there can be redundant copies when feeding the input to the model and reading output from the model. Make sure to eliminate redundant copies. If you are using higher level APIs like Java API, make sure to carefully check the documentation for performance caveats. For example, the Java API is a lot faster if ByteBuffers are used as [inputs](https://github.com/tensorflow/tensorflow/blob/6305a6d83552ba6a472cd72398b60d9241467f1f/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/Interpreter.java#L151). - -## Profile your application with platform specific tools -Platform specific tools like [Android profiler](https://developer.android.com/studio/profile/android-profiler) and [Instruments](https://help.apple.com/instruments/mac/current/) provide a wealth of profiling information that can be used to debug your app. Sometimes the performance bug may be not in the model but in parts of application code that interact with the model. Make sure to familiarize yourself with platform specific profiling tools and best practices for your platform. - -## Evaluate whether your model benefits from using hardware accelerators available on the device -Tensorflow Lite is working on adding support for accelerators like GPU and provides acceleration through [Neural Networks API](https://developer.android.com/ndk/guides/neuralnetworks/) on Android. -You can utilize these hardware accelerator backends to improve the speed and efficiency of your model. To enable Neural Networks API call [UseNNAPI](https://github.com/tensorflow/tensorflow/blob/6305a6d83552ba6a472cd72398b60d9241467f1f/tensorflow/contrib/lite/interpreter.h#L334) on the interpreter instance. - -## Need more help -The Tensorflow team is happy to help diagnose and address specific performance issues you may be facing. Please file a bug on [github](https://github.com/tensorflow/tensorflow/issues) with details of the issue. diff --git a/tensorflow/contrib/lite/g3doc/performance_benchmarks.md b/tensorflow/contrib/lite/g3doc/performance_benchmarks.md deleted file mode 100644 index 28cb6aba6ec61d12d86e078e47665833df8afec7..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/g3doc/performance_benchmarks.md +++ /dev/null @@ -1,174 +0,0 @@ - -# Performance - -This document lists TensorFlow Lite performance benchmarks when running well -known models on some Android and iOS devices. - -These performance benchmark numbers were generated with the -[Android TFLite benchmark binary](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite/tools/benchmark) -and the [iOS benchmark app](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite/tools/benchmark/ios). - -# Android performance benchmarks - -For Android benchmarks, the CPU affinity is set to use big cores on the device to -reduce variance (see [details](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite/tools/benchmark#reducing-variance-between-runs-on-android)). - -It assumes that models were download and unzipped to the -`/data/local/tmp/tflite_models` directory. The benchmark binary is built -using [these instructions](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite/tools/benchmark#on-android) -and assumed in the `/data/local/tmp` directory. - -To run the benchmark: - -``` -adb shell taskset ${CPU_MASK} /data/local/tmp/benchmark_model \ - --num_threads=1 \ - --graph=/data/local/tmp/tflite_models/${GRAPH} \ - --warmup_runs=1 \ - --num_runs=50 \ - --use_nnapi=false -``` - -Here, `${GRAPH}` is the name of model and `${CPU_MASK}` is the CPU affinity -chosen according to the following table: - -Device | CPU_MASK | --------| ---------- -Pixel 2 | f0 | -Pixel xl | 0c | - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
Model NameDevice Mean inference time (std dev)
- Mobilenet_1.0_224(float) - Pixel 2 166.5 ms (2.6 ms)
Pixel xl 122.9 ms (1.8 ms)
- Mobilenet_1.0_224 (quant) - Pixel 2 69.5 ms (0.9 ms)
Pixel xl 78.9 ms (2.2 ms)
- NASNet mobile - Pixel 2 273.8 ms (3.5 ms)
Pixel xl 210.8 ms (4.2 ms)
- SqueezeNet - Pixel 2 234.0 ms (2.1 ms)
Pixel xl 158.0 ms (2.1 ms)
- Inception_ResNet_V2 - Pixel 2 2846.0 ms (15.0 ms)
Pixel xl 1973.0 ms (15.0 ms)
- Inception_V4 - Pixel 2 3180.0 ms (11.7 ms)
Pixel xl 2262.0 ms (21.0 ms)
- -# iOS benchmarks - -To run iOS benchmarks, the [benchmark -app](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite/tools/benchmark/ios) -was modified to include the appropriate model and `benchmark_params.json` was -modified to set `num_threads` to 1. - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
Model NameDevice Mean inference time (std dev)
- Mobilenet_1.0_224(float) - iPhone 8 32.2 ms (0.8 ms)
- Mobilenet_1.0_224 (quant) - iPhone 8 24.4 ms (0.8 ms)
- NASNet mobile - iPhone 8 60.3 ms (0.6 ms)
- SqueezeNet - iPhone 8 44.3 (0.7 ms)
- Inception_ResNet_V2 - iPhone 8562.4 ms (18.2 ms)
- Inception_V4 - iPhone 8 661.0 ms (29.2 ms)
diff --git a/tensorflow/contrib/lite/graph_info.cc b/tensorflow/contrib/lite/graph_info.cc deleted file mode 100644 index e60ed2c2463cb621015ba725ca030e8d8c02f3c7..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/graph_info.cc +++ /dev/null @@ -1,224 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include "tensorflow/contrib/lite/graph_info.h" -#include - -namespace tflite { - -namespace { - -// Provide a range iterable wrapper for TfLiteIntArray* (C lists that TfLite -// C api uses. Can't use the google array_view, since we can't depend on even -// absl for embedded device reasons. -// TODO(aselle): Move this into central utilities. -class TfLiteIntArrayView { - public: - // Construct a view of a TfLiteIntArray*. Note, `int_array` should be non-null - // and this view does not take ownership of it. - explicit TfLiteIntArrayView(const TfLiteIntArray* int_array) - : int_array_(int_array) {} - - typedef const int* const_iterator; - const_iterator begin() const { return int_array_->data; } - const_iterator end() const { return &int_array_->data[int_array_->size]; } - - TfLiteIntArrayView(const TfLiteIntArrayView&) = default; - TfLiteIntArrayView& operator=(const TfLiteIntArrayView& rhs) = default; - - private: - const TfLiteIntArray* int_array_; -}; - -// Helper class that actually performs partitioning by subgraph. -// Outputs to a provided `subgraphs` structure. -// -// Example usage: -// PartitionGraphIntoIndependentSubgraphsImpl partitioner( -// info, nodes_to_part, subgraphs); -// partitioner.Partition(); -class PartitionGraphIntoIndependentSubgraphsImpl { - public: - PartitionGraphIntoIndependentSubgraphsImpl( - const GraphInfo* info, const TfLiteIntArray* nodes_to_partition, - std::vector* subgraphs) - : info_(info), - subgraphs_(subgraphs), - node_type_(info->num_nodes(), Subgraph::kTfNonPartition) { - // Populate the node_type_ map. - for (auto node_index : TfLiteIntArrayView(nodes_to_partition)) { - node_type_[node_index] = Subgraph::kTfPartition; - } - } - - // Actually partition the graph. - void Partition() { - // Initialize here to make Partition() re-entrant. - subgraphs_->clear(); - tensor_epochs_.clear(); - tensor_epochs_.resize(info_->num_tensors(), kEpochAlwaysReady); - node_epochs_.clear(); - node_epochs_.resize(info_->num_nodes(), kEpochNotReady); - // Set computed tensors to be kEpochNotReady (initializer set everything to - // AlwaysReady). - for (int node_index = 0; node_index < info_->num_nodes(); node_index++) { - const TfLiteNode& node = info_->node(node_index); - for (int output_tensor_index : TfLiteIntArrayView(node.outputs)) { - tensor_epochs_[output_tensor_index] = kEpochNotReady; - } - } - - // Do a graph traversal where each iteration in the loop is an epoch - // that corresponds to a subgraph that only contains nodes that are of - // the same node_type_. - while (true) { - BuildSubgraph(); - if (subgraphs_->back().nodes.empty()) { - subgraphs_->pop_back(); - break; - } - } - - // Mark model outputs as subgraph outputs. All the rest have already been - // identified. - for (int output_index : info_->outputs()) { - int output_epoch = tensor_epochs_[output_index]; - Subgraph& output_subgraph = (*subgraphs_)[output_epoch]; - output_subgraph.output_tensors.push_back(output_index); - } - // Make sure every subgraph's inputs and outputs are unique. Since the - // list of inputs and outputs is generated in a way that produces - // duplicates. - for (Subgraph& subgraph : *subgraphs_) { - // Sort and uniquefy using standard library algorithms. - auto uniquefy = [](std::vector* items) { - std::sort(items->begin(), items->end()); - auto last = std::unique(items->begin(), items->end()); - items->erase(last, items->end()); - }; - uniquefy(&subgraph.input_tensors); - uniquefy(&subgraph.output_tensors); - } - } - - private: - // Special integer values needed for tensor_epochs_ and node_epochs_. - enum { - // The node or tensor is not ready to be assigned an epoch. e.g. a node's - // inputs have not all been assigned epochs. - kEpochNotReady = -1, - // Used for tensor_epochs_. This means that the tensor is always ready. - // e.g. an input to the whole model or a constant that has no dependencies. - kEpochAlwaysReady = -2 - }; - - // Updates the node `node_index` and returns true if it is assigned to an - // epoch. False is returned if the node is already set to an epoch, its inputs - // are not all assigned to epochs, or if it cannot be assigned to the current - // epoch since the epoch's node_type doesn't match. - bool UpdateNode(int node_index) { - const TfLiteNode& node = info_->node(node_index); - Subgraph& current_subgraph = subgraphs_->back(); - int current_epoch = subgraphs_->size() - 1; - // Check if node is already done. - if (node_epochs_[node_index] != kEpochNotReady) { - return false; - } - // See if all dependencies of this node are already assigned to a - // subgraph. - for (int input_tensor_index : TfLiteIntArrayView(node.inputs)) { - if (tensor_epochs_[input_tensor_index] == kEpochNotReady) { - return false; - } - } - // When we are starting a new epoch, the first ready node defines - // the type of that epoch. - if (current_subgraph.type == Subgraph::kTfUnexplored) { - current_subgraph.type = node_type_[node_index]; - } - // The node gets assigned to this epoch if it is the same type as - // the epoch's assigned type. Note, if this is the current ready - // node encountered during this epoch, this condition will be - // automatically true. - if (current_subgraph.type == node_type_[node_index]) { - node_epochs_[node_index] = current_epoch; - current_subgraph.nodes.push_back(node_index); - // All outputs of this node now are assigned to this epoch as - // well. - for (int output_tensor_index : TfLiteIntArrayView(node.outputs)) { - tensor_epochs_[output_tensor_index] = current_epoch; - } - // Look at our inputs one more time to update that tensor's - // epochs' outputs - for (int input_tensor_index : TfLiteIntArrayView(node.inputs)) { - int input_epoch = tensor_epochs_[input_tensor_index]; - int node_epoch = current_epoch; - if (input_epoch != node_epoch) { - current_subgraph.input_tensors.push_back(input_tensor_index); - // Set inputs to be outputs of the subgraph where they reside. - // the if condition makes sure inputs to the whole computation - // are not included (i.e. those initialized to -2 above). - if (input_epoch >= 0) { - Subgraph& input_subgraph = (*subgraphs_)[input_epoch]; - input_subgraph.output_tensors.push_back(input_tensor_index); - } - } - } - return true; - } else { - return false; - } - } - - // Completely populates the current subgraph by doing graph traversal - void BuildSubgraph() { - subgraphs_->emplace_back(Subgraph()); - // loop until no more nodes can be updated. - while (true) { - bool did_something = false; - for (int node_index = 0; node_index < info_->num_nodes(); node_index++) { - if (UpdateNode(node_index)) { - did_something = true; - } - } - if (!did_something) return; - } - } - - // Temporary data needed for partitioning. - const GraphInfo* info_; - // List of subgraphs to populate - std::vector* subgraphs_; - std::vector node_type_; - // Maps from tensor index to the epoch in which it is assigned. Also special - // negative values of kEpochNotAssigned if not assigned, kEpochNotReady if it - // is an input or constant. - std::vector tensor_epochs_; - // Maps from tensor index to the epoch in which it is assigned. Also special - // negative values of kEpochNotAssigned if not assigned. - std::vector node_epochs_; -}; - -} // namespace - -TfLiteStatus PartitionGraphIntoIndependentSubgraphs( - const GraphInfo* info, const TfLiteIntArray* nodes_to_partition, - std::vector* subgraphs) { - PartitionGraphIntoIndependentSubgraphsImpl(info, nodes_to_partition, - subgraphs) - .Partition(); - return kTfLiteOk; -} - -} // namespace tflite diff --git a/tensorflow/contrib/lite/graph_info.h b/tensorflow/contrib/lite/graph_info.h deleted file mode 100644 index 8ee83827bb3fdf59b88d8304ad781cae98140b75..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/graph_info.h +++ /dev/null @@ -1,82 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#ifndef TENSORFLOW_CONTRIB_LITE_GRAPH_INFO_H_ -#define TENSORFLOW_CONTRIB_LITE_GRAPH_INFO_H_ - -#include - -#include "tensorflow/contrib/lite/c/c_api_internal.h" - -namespace tflite { - -// Basic information about an inference graph, where execution nodes -// are connected via tensors. -class GraphInfo { - public: - virtual ~GraphInfo() {} - - // Total number of tensors in the graph. - virtual size_t num_tensors() const = 0; - - // Returns a tensor given its index which is expected to be between 0 and - // num_tensors(). - virtual TfLiteTensor* tensor(size_t index) = 0; - - // Total number of nodes in the graph. - virtual size_t num_nodes() const = 0; - - // Returns a node given its index which is expected to be between 0 and - // num_nodes(). - virtual const TfLiteNode& node(size_t index) const = 0; - - // Returns the indices of the input tensors. - virtual const std::vector& inputs() const = 0; - - // Returns the indices of the output tensors. - virtual const std::vector& outputs() const = 0; - - // Returns the indices of the variable tensors. - virtual const std::vector& variables() const = 0; -}; - -// Represents a subgraph of a TensorFlow Lite graph. -struct Subgraph { - enum Type { - kTfUnexplored = 0, // temporarily used during creation - kTfPartition, - kTfNonPartition - }; - Type type = kTfUnexplored; - // Nodes within the subgraph - std::vector nodes; - // Tensors that stride output from another subgraph that this depends on, - // or global inputs to the TensorFlow Lite full graph. - std::vector input_tensors; - // Outputs that are consumed by other subgraphs or are global output tensors. - // All output tensors of the nodes in the subgraph that do not appear in this - // list are intermediate results that can be potentially elided. - std::vector output_tensors; -}; - -// Partitions a list of node indices `nodes_to_partition` into subgraphs. -// Each subgraph is in dependency order (i.e. all members of the subgraph). -// `subgraphs` is assumed to be empty. -TfLiteStatus PartitionGraphIntoIndependentSubgraphs( - const GraphInfo* info, const TfLiteIntArray* nodes_to_partition, - std::vector* subgraphs); - -} // namespace tflite - -#endif // TENSORFLOW_CONTRIB_LITE_GRAPH_INFO_H_ diff --git a/tensorflow/contrib/lite/interpreter.cc b/tensorflow/contrib/lite/interpreter.cc deleted file mode 100644 index c72e7bf33ebbaac09916ffda6faf4b812d702ea8..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/interpreter.cc +++ /dev/null @@ -1,1018 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/contrib/lite/interpreter.h" - -#include -#include -#include -#include - -#include "tensorflow/contrib/lite/arena_planner.h" -#include "tensorflow/contrib/lite/c/c_api_internal.h" -#include "tensorflow/contrib/lite/context_util.h" -#include "tensorflow/contrib/lite/core/api/error_reporter.h" -#include "tensorflow/contrib/lite/graph_info.h" -#include "tensorflow/contrib/lite/memory_planner.h" -#include "tensorflow/contrib/lite/nnapi_delegate.h" -#include "tensorflow/contrib/lite/profiling/profiler.h" -#include "tensorflow/contrib/lite/schema/schema_generated.h" -#include "tensorflow/contrib/lite/util.h" - -namespace tflite { -namespace { - -TfLiteStatus ReportOpError(TfLiteContext* context, const TfLiteNode& node, - const TfLiteRegistration& registration, - int node_index, const char* message) { - context->ReportError( - context, "Node number %d (%s) %s.\n", node_index, - registration.custom_name - ? registration.custom_name - : EnumNameBuiltinOperator( - static_cast(registration.builtin_code)), - message); - return kTfLiteError; -} - -// Stub method which returns kTfLiteError when the function is forbidden. -// We're registrating this function to several different function to save -// compiled binary size. Please note the restrictions: -// * The type of first parameter have to be `TfLiteContext*`. -// * All paramteters must be trivailly destructible. (E.g. No C++ class) -TfLiteStatus ForbiddenContextFunction(TfLiteContext* context, ...) { - context->ReportError(context, - "The function is forbidden if not calling in delegate."); - return kTfLiteError; -} - -// Set the ForbiddenContextFunction to a compatible function pointer. -template -void SetForbiddenContextFunction(FunctionType* func) { - *func = reinterpret_cast(ForbiddenContextFunction); -} - -// Returns true if at least one tensor in the given list is kTfLiteDynamic. -template -bool HasDynamicTensorImpl(const TfLiteContext& context, - const TensorIntArray& int_array) { - for (int i : int_array) { - const TfLiteTensor& tensor = context.tensors[i]; - if (tensor.allocation_type == kTfLiteDynamic) { - return true; - } - } - return false; -} - -} // namespace - -// A trivial implementation of GraphInfo around the Interpreter. -// NOTE: this interpreter info represents the subset of the -// graph that is executed according to execution plan. Thus, -// the indices are execution plan indices rather than raw node -// indices. -class InterpreterInfo : public GraphInfo { - public: - explicit InterpreterInfo(Interpreter* interpreter) - : interpreter_(interpreter) {} - - size_t num_tensors() const override { return interpreter_->tensors_size(); } - TfLiteTensor* tensor(size_t index) override { - return interpreter_->tensor(index); - } - size_t num_nodes() const override { - return interpreter_->execution_plan().size(); - } - const TfLiteNode& node(size_t index) const override { - int node_index = interpreter_->execution_plan()[index]; - return interpreter_->node_and_registration(node_index)->first; - } - const std::vector& inputs() const override { - return interpreter_->inputs(); - } - const std::vector& outputs() const override { - return interpreter_->outputs(); - } - const std::vector& variables() const override { - return interpreter_->variables(); - } - - public: - Interpreter* interpreter_; -}; - -Interpreter::Interpreter(ErrorReporter* error_reporter) - : error_reporter_(error_reporter ? error_reporter - : DefaultErrorReporter()) { - context_.impl_ = static_cast(this); - context_.ResizeTensor = ResizeTensor; - context_.ReportError = ReportError; - context_.AddTensors = AddTensors; - context_.tensors = nullptr; - context_.tensors_size = 0; - context_.allow_fp32_relax_to_fp16 = false; - context_.recommended_num_threads = -1; - context_.GetExternalContext = GetExternalContext; - context_.SetExternalContext = SetExternalContext; - - // Invalid to call these these except from TfLiteDelegate - SwitchToKernelContext(); - - // Reserve some space for the tensors to avoid excessive resizing. - tensors_.reserve(kTensorsReservedCapacity); - nodes_and_registration_.reserve(kTensorsReservedCapacity); - next_execution_plan_index_to_prepare_ = 0; - - for (int i = 0; i < kTfLiteMaxExternalContexts; ++i) { - external_contexts_[i] = nullptr; - } - - UseNNAPI(false); -} - -Interpreter::~Interpreter() { - for (auto& nodeAndReg : nodes_and_registration_) { - TfLiteNode& node = nodeAndReg.first; - TfLiteIntArrayFree(node.inputs); - TfLiteIntArrayFree(node.outputs); - TfLiteIntArrayFree(node.temporaries); - if (node.builtin_data) free(node.builtin_data); - OpFree(nodeAndReg.second, node.user_data); - node.builtin_data = nullptr; - } - - for (int i = 0; i < context_.tensors_size; i++) { - TfLiteTensor* tensor = &context_.tensors[i]; - if (tensor->buffer_handle != kTfLiteNullBufferHandle && - tensor->delegate->FreeBufferHandle != nullptr) { - tensor->delegate->FreeBufferHandle(&context_, tensor->delegate, - &tensor->buffer_handle); - } - TfLiteTensorFree(tensor); - } -} - -TfLiteStatus Interpreter::ReplaceSubgraphsWithDelegateKernels( - TfLiteContext* context, TfLiteRegistration registration, - const TfLiteIntArray* nodes_to_replace, TfLiteDelegate* delegate) { - return static_cast(context->impl_) - ->ReplaceSubgraphsWithDelegateKernels(registration, nodes_to_replace, - delegate); -} - -namespace { - -// Copy a std::vector to an existing TfLiteIntArray. -// This is a low-level data manipulation function, and it's caller's -// responsibility to ensure TfLiteIntArray has enough size. -void CopyVectorToTfLiteIntArray(const std::vector& vec, - TfLiteIntArray* arr) { - arr->size = vec.size(); - memcpy(arr->data, vec.data(), sizeof(int) * arr->size); -} - -// This function allocates a continuous memory space that contains a -// TfLiteDelegateParams followed by a several TfLiteIntArray. -// When calling `free` at TfLiteDelegateParams*, all the allocated space -// will be freed together. -// -// +-----------------------------------+ -// | TfLiteDelegateParams | -// | TfLiteDelegate* delegate; | -// | TfLiteIntArray* nodes_to_replace; |--\ -// | TfLiteIntArray* input_tensors; |--+--\ -// | TfLiteIntArray* output_tensors; |--+--+--\ -// +-----------------------------------+ | | | -// | TfLiteIntArray (variable size) |<-/ | | -// +-----------------------------------+ | | -// | TfLiteIntArray (variable size) |<----/ | -// +-----------------------------------+ | -// | TfLiteIntArray (variable size) |<-------/ -// +-----------------------------------+ -TfLiteDelegateParams* CreateDelegateParams(TfLiteDelegate* delegate, - const Subgraph& subgraph) { - // Step 1: Calculate the allocation size. - int allocation_size = sizeof(TfLiteDelegateParams); - - int nodes_to_replace_size = - TfLiteIntArrayGetSizeInBytes(subgraph.nodes.size()); - allocation_size += nodes_to_replace_size; - - int input_tensors_size = - TfLiteIntArrayGetSizeInBytes(subgraph.input_tensors.size()); - allocation_size += input_tensors_size; - - int output_tensors_size = - TfLiteIntArrayGetSizeInBytes(subgraph.output_tensors.size()); - allocation_size += output_tensors_size; - - // Step 2: Allocate the memory. - // Use `char*` for conveniently step through the allocated space by bytes. - char* allocation = reinterpret_cast(malloc(allocation_size)); - - // Step 3: Fill all data structures structures. - TfLiteDelegateParams* params = - reinterpret_cast(allocation); - params->delegate = delegate; - allocation += sizeof(TfLiteDelegateParams); - - params->nodes_to_replace = reinterpret_cast(allocation); - CopyVectorToTfLiteIntArray(subgraph.nodes, params->nodes_to_replace); - allocation += nodes_to_replace_size; - - params->input_tensors = reinterpret_cast(allocation); - CopyVectorToTfLiteIntArray(subgraph.input_tensors, params->input_tensors); - allocation += input_tensors_size; - - params->output_tensors = reinterpret_cast(allocation); - CopyVectorToTfLiteIntArray(subgraph.output_tensors, params->output_tensors); - allocation += output_tensors_size; - - return params; -} - -} // namespace - -TfLiteStatus Interpreter::ReplaceSubgraphsWithDelegateKernels( - TfLiteRegistration registration, const TfLiteIntArray* nodes_to_replace, - TfLiteDelegate* delegate) { - // Annotate the registration as DELEGATE op. - registration.builtin_code = BuiltinOperator_DELEGATE; - - // Analyze the graph to find all independent subgraphs that are either - // fully not-this-delegate or this-delegate computation. - InterpreterInfo info(this); - std::vector subgraphs; - PartitionGraphIntoIndependentSubgraphs(&info, nodes_to_replace, &subgraphs); - - execution_plan_.clear(); - for (auto& subgraph : subgraphs) { - // Subgraphs calimed by the delegate should have a "macro" op created, the - // other subgraphs (kTfNonPartition) just have their nodes added back to - // the execution plan. - switch (subgraph.type) { - case Subgraph::kTfNonPartition: - for (auto it = subgraph.nodes.begin(); it != subgraph.nodes.end(); - ++it) { - execution_plan_.push_back(*it); - } - break; - case Subgraph::kTfPartition: { - int node_index; - - TfLiteDelegateParams* params = CreateDelegateParams(delegate, subgraph); - TF_LITE_ENSURE_STATUS(AddNodeWithParameters( - subgraph.input_tensors, subgraph.output_tensors, nullptr, 0, params, - ®istration, &node_index)); - - // Initialize the output tensors's delegate-related fields. - for (int tensor_index : subgraph.output_tensors) { - TfLiteTensor* tensor = &tensors_[tensor_index]; - TF_LITE_ENSURE(&context_, tensor->delegate == nullptr || - tensor->delegate == delegate); - tensor->delegate = delegate; - } - - // Associate the node with the delegate. - TfLiteNode* node = &nodes_and_registration_[node_index].first; - node->delegate = delegate; - } break; - case Subgraph::kTfUnexplored: - return kTfLiteError; - break; - } - } - return kTfLiteOk; -} - -TfLiteExternalContext* Interpreter::GetExternalContext( - TfLiteExternalContextType type) { - if (type >= 0 && type < kTfLiteMaxExternalContexts) { - return external_contexts_[type]; - } - return nullptr; -} - -TfLiteExternalContext* Interpreter::GetExternalContext( - struct TfLiteContext* context, TfLiteExternalContextType type) { - return static_cast(context->impl_)->GetExternalContext(type); -} - -void Interpreter::SetExternalContext(TfLiteExternalContextType type, - TfLiteExternalContext* ctx) { - if (type >= 0 && type < kTfLiteMaxExternalContexts) { - external_contexts_[type] = ctx; - } -} - -void Interpreter::SetExternalContext(struct TfLiteContext* context, - TfLiteExternalContextType type, - TfLiteExternalContext* ctx) { - return static_cast(context->impl_) - ->SetExternalContext(type, ctx); -} - -// Gets an TfLiteIntArray* representing the execution plan. The interpreter owns -// this memory and it is only guaranteed to exist during the invocation of the -// delegate prepare. -TfLiteStatus Interpreter::GetExecutionPlan(TfLiteIntArray** execution_plan) { - // TODO(aselle): Do not make a copy here - plan_cache_.reset(TfLiteIntArrayCreate(execution_plan_.size())); - *execution_plan = plan_cache_.get(); - static_assert(sizeof(plan_cache_->data[0]) == sizeof(execution_plan_[0]), - "TfLiteIntArray and execution_plan do not contain same type."); - std::memcpy(plan_cache_->data, execution_plan_.data(), - sizeof(plan_cache_->data[0]) * execution_plan_.size()); - return kTfLiteOk; -} - -// WARNING: This is an experimental interface that is subject to change. -// Entry point for C node plugin API to get the execution plan -TfLiteStatus Interpreter::GetExecutionPlan(struct TfLiteContext* context, - TfLiteIntArray** execution_plan) { - return static_cast(context->impl_) - ->GetExecutionPlan(execution_plan); -} - -TfLiteStatus Interpreter::SetInputs(std::vector inputs) { - TF_LITE_ENSURE_OK(&context_, - CheckTensorIndices("inputs", inputs.data(), inputs.size())); - inputs_ = std::move(inputs); - return kTfLiteOk; -} - -TfLiteStatus Interpreter::SetOutputs(std::vector outputs) { - TF_LITE_ENSURE_OK( - &context_, CheckTensorIndices("outputs", outputs.data(), outputs.size())); - outputs_ = std::move(outputs); - return kTfLiteOk; -} - -TfLiteStatus Interpreter::SetVariables(std::vector variables) { - TF_LITE_ENSURE_OK(&context_, CheckTensorIndices("variables", variables.data(), - variables.size())); - variables_ = std::move(variables); - return kTfLiteOk; -} - -TfLiteStatus Interpreter::CheckTensorIndices(const char* label, - const int* indices, int length) { - // Making sure kOptionalTensor is not re-defined to something other than -1. - static_assert(kOptionalTensor == -1, "kOptionalTensor should be defined -1"); - - for (int i = 0; i < length; i++) { - int index = indices[i]; - // Continue if index == kOptionalTensor before additional comparisons below, - // size_t(-1) is always >= context_tensors_size. - if (index == kOptionalTensor) { - continue; - } - if (index < 0 || static_cast(index) >= context_.tensors_size) { - ReportError(&context_, "Invalid tensor index %d in %s\n", index, label); - consistent_ = false; - return kTfLiteError; - } - } - return kTfLiteOk; -} - -TfLiteStatus Interpreter::BytesRequired(TfLiteType type, const int* dims, - size_t dims_size, size_t* bytes) { - // TODO(aselle): Check for overflow here using overflow.h in TensorFlow - // MultiplyWithoutOverflow. - TF_LITE_ENSURE(&context_, bytes != nullptr); - size_t count = 1; - for (int k = 0; k < dims_size; k++) count *= dims[k]; - switch (type) { - case kTfLiteFloat32: - *bytes = sizeof(float) * count; - break; - case kTfLiteInt16: - *bytes = sizeof(int16_t) * count; - break; - case kTfLiteInt32: - *bytes = sizeof(int32_t) * count; - break; - case kTfLiteUInt8: - *bytes = sizeof(uint8_t) * count; - break; - case kTfLiteInt64: - *bytes = sizeof(int64_t) * count; - break; - case kTfLiteBool: - *bytes = sizeof(bool) * count; - break; - case kTfLiteComplex64: - *bytes = sizeof(std::complex) * count; - break; - default: - ReportError(&context_, - "Only float32, int16, int32, int64, uint8, bool, complex64 " - "supported currently."); - return kTfLiteError; - } - return kTfLiteOk; -} - -TfLiteStatus Interpreter::AllocateTensors() { - if (!consistent_) { - ReportError(&context_, "AllocateTensors() called on inconsistent model."); - return kTfLiteError; - } - - // Explicit (re)allocation is necessary if nodes have been changed or tensors - // have been resized. For inputs marked as dynamic, we can't short-circuit the - // allocation as the client may have done the resize manually. - if (state_ != kStateUninvokable && !HasDynamicTensorImpl(context_, inputs_)) { - return kTfLiteOk; - } - - next_execution_plan_index_to_prepare_ = 0; - if (memory_planner_) { - TF_LITE_ENSURE_STATUS(memory_planner_->ResetAllocations()); - } - - TF_LITE_ENSURE_STATUS(PrepareOpsAndTensors()); - - state_ = kStateInvokable; - - // Reset the variable tensors to zero after (re)allocating the tensors. - // Developers shouldn't rely on the side effect of this function to reset - // variable tesnsors. They should call `ResetVariableTensors` directly - // instead. - ResetVariableTensors(); - - return kTfLiteOk; -} - -// TODO(ycling): Support non-zero default values. -TfLiteStatus Interpreter::ResetVariableTensors() { - for (auto& tensor : tensors_) { - if (!tensor.is_variable) { - continue; - } - - // Variable tensors have to be `kTfLiteArenaRwPersistent`, and must be - // allocated after the initial `PrepareOpsAndTensors()` is called. - TF_LITE_ENSURE_EQ(&context_, tensor.allocation_type, - kTfLiteArenaRwPersistent); - TF_LITE_ENSURE(&context_, tensor.data.raw != nullptr); - - memset(tensor.data.raw, 0, tensor.bytes); - } - return kTfLiteOk; -} - -void Interpreter::ReserveNodes(int count) { - nodes_and_registration_.reserve(count); -} - -TfLiteStatus Interpreter::AddNodeWithParameters( - const std::vector& inputs, const std::vector& outputs, - const char* init_data, size_t init_data_size, void* builtin_data, - const TfLiteRegistration* registration, int* node_index) { - if (state_ == kStateInvokableAndImmutable) { - ReportError(&context_, - "AddNodeWithParameters is disallowed when graph is immutable."); - return kTfLiteError; - } - state_ = kStateUninvokable; - - std::unique_ptr builtin_data_deleter(builtin_data, - free); - - TF_LITE_ENSURE_OK(&context_, CheckTensorIndices("node inputs", inputs.data(), - inputs.size())); - TF_LITE_ENSURE_OK( - &context_, - CheckTensorIndices("node outputs", outputs.data(), outputs.size())); - - int new_node_index = nodes_and_registration_.size(); - if (node_index) *node_index = new_node_index; - nodes_and_registration_.resize(nodes_and_registration_.size() + 1); - auto& node_and_reg = nodes_and_registration_.back(); - TfLiteNode& node = node_and_reg.first; - if (node.inputs) TfLiteIntArrayFree(node.inputs); - if (node.outputs) TfLiteIntArrayFree(node.outputs); - if (node.temporaries) TfLiteIntArrayFree(node.temporaries); - - // NOTE, here we are not using move semantics yet, since our internal - // representation isn't std::vector, but in the future we would like to avoid - // copies, so we want the interface to take r-value references now. - node.inputs = ConvertVectorToTfLiteIntArray(inputs); - node.outputs = ConvertVectorToTfLiteIntArray(outputs); - node.temporaries = TfLiteIntArrayCreate(0); - if (init_data) { - node.user_data = OpInit(*registration, init_data, init_data_size); - } else { - node.user_data = - OpInit(*registration, - reinterpret_cast(builtin_data_deleter.get()), 0); - } - - node.builtin_data = builtin_data_deleter.release(); - // TODO(ycling): Filling `custom_initial_data` and `custom_initial_data_size` - // properly for nodes generated by ReplaceSubgraphsWithDelegateKernels. - - if (registration->builtin_code == BuiltinOperator_CUSTOM) { - // When it's a CUSTOM op, the `custom_options` field in the Flatbuffer - // `Operator` table is passed in. - node.custom_initial_data = init_data; - node.custom_initial_data_size = init_data_size; - } else { - node.custom_initial_data = nullptr; - node.custom_initial_data_size = 0; - } - - node.delegate = nullptr; - node_and_reg.second = *registration; - execution_plan_.push_back(new_node_index); - return kTfLiteOk; -} - -TfLiteStatus Interpreter::ResizeInputTensor(int tensor_index, - const std::vector& dims) { - if (state_ == kStateInvokableAndImmutable) { - ReportError(&context_, - "ResizeInputTensor is disallowed when graph is immutable."); - return kTfLiteError; - } - - // TODO(aselle): All bounds checks can be implemented as one-sided bounds - // checks by casting to unsigned for efficiency. Profile before doing this. - TF_LITE_ENSURE(&context_, - tensor_index < context_.tensors_size && tensor_index >= 0); - TfLiteTensor* tensor = &context_.tensors[tensor_index]; - - // Short-circuit the state change if the dimensions don't change, avoiding - // unnecessary (re)allocations. - if (EqualArrayAndTfLiteIntArray(tensor->dims, dims.size(), dims.data())) { - return kTfLiteOk; - } - - state_ = kStateUninvokable; - return ResizeTensorImpl(tensor, ConvertVectorToTfLiteIntArray(dims)); -} - -bool HasDynamicTensor(const TfLiteContext& context, - const TfLiteIntArray* int_array) { - return HasDynamicTensorImpl(context, TfLiteIntArrayView{int_array}); -} - -TfLiteStatus Interpreter::PrepareOpsStartingAt( - int first_execution_plan_index, int* last_execution_plan_index_prepared) { - for (int execution_plan_index = first_execution_plan_index; - execution_plan_index < execution_plan_.size(); execution_plan_index++) { - int node_index = execution_plan_[execution_plan_index]; - TfLiteNode& node = nodes_and_registration_[node_index].first; - const TfLiteRegistration& registration = - nodes_and_registration_[node_index].second; - EnsureTensorsVectorCapacity(); - if (OpPrepare(registration, &node) == kTfLiteError) { - return ReportOpError(&context_, node, registration, node_index, - "failed to prepare"); - } - - *last_execution_plan_index_prepared = execution_plan_index; - - // Discontinue if the node has dynamic outputs. Note that we don't - // stop for dynamic temporary tensors since they won't affect the - // sizes of other tensors in the graph. - if (HasDynamicTensor(context_, node.outputs)) { - break; - } - } - return kTfLiteOk; -} - -TfLiteStatus Interpreter::PrepareOpsAndTensors() { - if (!memory_planner_) { - memory_planner_.reset(new ArenaPlanner( - &context_, std::unique_ptr(new InterpreterInfo(this)), - /*preserve_inputs=*/true, /*preserve_intermediates*/ false)); - memory_planner_->PlanAllocations(); - } - - int last_exec_plan_index_prepared = 0; - - TF_LITE_ENSURE_STATUS(PrepareOpsStartingAt( - next_execution_plan_index_to_prepare_, &last_exec_plan_index_prepared)); - TF_LITE_ENSURE_STATUS(memory_planner_->ExecuteAllocations( - next_execution_plan_index_to_prepare_, last_exec_plan_index_prepared)); - - next_execution_plan_index_to_prepare_ = last_exec_plan_index_prepared + 1; - return kTfLiteOk; -} - -TfLiteStatus Interpreter::Invoke() { - if (!consistent_) { - ReportError(&context_, "Invoke called on model that is not consistent."); - return kTfLiteError; - } - if (state_ == kStateUninvokable) { - ReportError(&context_, "Invoke called on model that is not ready."); - return kTfLiteError; - } - - TfLiteStatus status = kTfLiteOk; - if (nnapi_delegate_) { - if (next_execution_plan_index_to_prepare_ == execution_plan_.size()) { - TF_LITE_ENSURE_OK(&context_, nnapi_delegate_->Invoke(this)); - return kTfLiteOk; - } else { - // TODO(aselle): In the future, we would like this to be an - // automatic tflite CPU fallback. - ReportError(&context_, - "NNAPI was requested, but dependent sized tensors " - "being used.\n"); - return kTfLiteError; - } - } - - // Invocations are always done in node order. - // Note that calling Invoke repeatedly will cause the original memory plan to - // be reused, unless either ResizeInputTensor() or AllocateTensors() has been - // called. - // TODO(b/71913981): we should force recalculation in the presence of dynamic - // tensors, because they may have new value which in turn may affect shapes - // and allocations. - for (int execution_plan_index = 0; - execution_plan_index < execution_plan_.size(); execution_plan_index++) { - if (execution_plan_index == next_execution_plan_index_to_prepare_) { - TF_LITE_ENSURE_STATUS(PrepareOpsAndTensors()); - TF_LITE_ENSURE(&context_, next_execution_plan_index_to_prepare_ >= - execution_plan_index); - } - int node_index = execution_plan_[execution_plan_index]; - TfLiteNode& node = nodes_and_registration_[node_index].first; - const TfLiteRegistration& registration = - nodes_and_registration_[node_index].second; - SCOPED_OPERATOR_PROFILE(profiler_, node_index); - - // TODO(ycling): This is an extra loop through inputs to check if the data - // need to be copied from Delegate buffer to raw memory, which is often not - // needed. We may want to cache this in prepare to know if this needs to be - // done for a node or not. - for (int i = 0; i < node.inputs->size; ++i) { - int tensor_index = node.inputs->data[i]; - if (tensor_index == kOptionalTensor) { - continue; - } - TfLiteTensor* tensor = &tensors_[tensor_index]; - if (tensor->delegate && tensor->delegate != node.delegate && - tensor->data_is_stale) { - EnsureTensorDataIsReadable(tensor_index); - } - } - - EnsureTensorsVectorCapacity(); - tensor_resized_since_op_invoke_ = false; - if (OpInvoke(registration, &node) == kTfLiteError) { - status = ReportOpError(&context_, node, registration, node_index, - "failed to invoke"); - } - - // Force execution prep for downstream ops if the latest op triggered the - // resize of a dynamic tensor. - if (tensor_resized_since_op_invoke_ && - HasDynamicTensor(context_, node.outputs)) { - next_execution_plan_index_to_prepare_ = execution_plan_index + 1; - } - } - - if (!allow_buffer_handle_output_) { - for (int tensor_index : outputs_) { - EnsureTensorDataIsReadable(tensor_index); - } - } - - return status; -} - -TfLiteStatus Interpreter::ResizeTensor(TfLiteContext* context, - TfLiteTensor* tensor, - TfLiteIntArray* new_size) { - // Note here that context->impl_ is recovering the this pointer for an - // instance of Interpreter to call into the member function ResizeTensorImpl - // (this function is static). - return static_cast(context->impl_) - ->ResizeTensorImpl(tensor, new_size); -} - -void Interpreter::ReportErrorImpl(const char* format, va_list args) { - error_reporter_->Report(format, args); -} - -void Interpreter::ReportError(TfLiteContext* context, const char* format, ...) { - va_list args; - va_start(args, format); - auto* f = static_cast(context->impl_); - // Note here that context->impl_ is recovering the this pointer for an - // instance of Interpreter to call into the member function ReportErrorImpl - // (this function is static). - f->ReportErrorImpl(format, args); - va_end(args); -} - -TfLiteStatus Interpreter::AddTensors(int tensors_to_add, - int* first_new_tensor_index) { - int base_index = tensors_.size(); - if (first_new_tensor_index) *first_new_tensor_index = base_index; - tensors_.resize(tensors_.size() + tensors_to_add); - for (int i = base_index; i < tensors_.size(); i++) { - memset(&tensors_[i], 0, sizeof(tensors_[i])); - tensors_[i].buffer_handle = kTfLiteNullBufferHandle; - } - context_.tensors = tensors_.data(); - context_.tensors_size = tensors_.size(); - return kTfLiteOk; -} - -TfLiteStatus Interpreter::AddTensors(TfLiteContext* context, int tensors_to_add, - int* first_new_tensor_index) { - // Note here that context->impl_ is recovering the this pointer for an - // instance of Interpreter to call into the member function AddTensors - // (this function is static). - return static_cast(context->impl_) - ->AddTensors(tensors_to_add, first_new_tensor_index); -} - -TfLiteStatus Interpreter::GetNodeAndRegistration( - int node_index, TfLiteNode** node, TfLiteRegistration** registration) { - TF_LITE_ENSURE(&context_, node_index < nodes_size() && node_index >= 0); - TF_LITE_ENSURE(&context_, node != nullptr && registration != nullptr); - *node = &nodes_and_registration_[node_index].first; - *registration = &nodes_and_registration_[node_index].second; - return kTfLiteOk; -} - -TfLiteStatus Interpreter::GetNodeAndRegistration( - struct TfLiteContext* context, int node_index, TfLiteNode** node, - TfLiteRegistration** registration) { - return static_cast(context->impl_) - ->GetNodeAndRegistration(node_index, node, registration); -} - -TfLiteStatus Interpreter::SetTensorParametersReadOnly( - int tensor_index, TfLiteType type, const char* name, const size_t rank, - const int* dims, TfLiteQuantizationParams quantization, const char* buffer, - size_t bytes, const Allocation* allocation) { - if (state_ == kStateInvokableAndImmutable) { - ReportError( - &context_, - "SetTensorParametersReadOnly is disallowed when graph is immutable."); - return kTfLiteError; - } - - TF_LITE_ENSURE(&context_, - tensor_index < context_.tensors_size && tensor_index >= 0); - // For most tensors we know exactly how much memory is necessary so we can - // ensure the buffer is large enough. However, we need to skip string tensors - // because their sizes change with the contents of the individual strings. - if (type != kTfLiteString) { - size_t required_bytes; - TF_LITE_ENSURE_OK(&context_, - BytesRequired(type, dims, rank, &required_bytes)); - TF_LITE_ENSURE_EQ(&context_, required_bytes, bytes); - } - - TfLiteTensor& tensor = context_.tensors[tensor_index]; - if (type == tensor.type && - EqualArrayAndTfLiteIntArray(tensor.dims, rank, dims)) { - // Fast path which does not invalidate the invokable property. - TfLiteTensorDataFree(&tensor); - tensor.data.raw = const_cast(buffer); - if (!tensor.dims) tensor.dims = ConvertArrayToTfLiteIntArray(rank, dims); - tensor.params = quantization; - tensor.allocation_type = kTfLiteMmapRo; - tensor.allocation = allocation; - } else { - state_ = kStateUninvokable; - TfLiteTensorReset(type, name, ConvertArrayToTfLiteIntArray(rank, dims), - quantization, const_cast(buffer), bytes, - kTfLiteMmapRo, allocation, false, &tensor); - } - return kTfLiteOk; -} - -// Set description of inputs/outputs/data/fptrs for node `node_index`. -// This variant assumes an external buffer has been allocated of size -// bytes. The lifetime of buffer must be ensured to be greater or equal -// to Interpreter. -TfLiteStatus Interpreter::SetTensorParametersReadWrite( - int tensor_index, TfLiteType type, const char* name, const size_t rank, - const int* dims, TfLiteQuantizationParams quantization, bool is_variable) { - if (state_ == kStateInvokableAndImmutable) { - ReportError( - &context_, - "SetTensorParametersReadWrite is disallowed when graph is immutable."); - return kTfLiteError; - } - TF_LITE_ENSURE(&context_, - tensor_index < context_.tensors_size && tensor_index >= 0); - size_t required_bytes = 0; - if (type != kTfLiteString) { - // These types will be allocated in our arena so we need to record how - // many bytes we will need based on the dimensions. String tensors are - // allocated dynamically and we can't know ahead of time how much space - // they will require. - TF_LITE_ENSURE_OK(&context_, - BytesRequired(type, dims, rank, &required_bytes)); - } - - TfLiteAllocationType allocation_type = kTfLiteArenaRw; - if (type == kTfLiteString) { - if (is_variable) { - // We don't have a real use case for string variable tensor. - ReportError(&context_, "String variable tensor isn't supported."); - return kTfLiteError; - } - allocation_type = kTfLiteDynamic; - } else if (is_variable) { - allocation_type = kTfLiteArenaRwPersistent; - } - - TfLiteTensorReset(type, name, ConvertArrayToTfLiteIntArray(rank, dims), - quantization, - /*buffer=*/nullptr, required_bytes, allocation_type, - nullptr, is_variable, &context_.tensors[tensor_index]); - return kTfLiteOk; -} - -TfLiteStatus Interpreter::SetExecutionPlan(const std::vector& new_plan) { - for (int node_index : new_plan) { - TF_LITE_ENSURE(&context_, node_index >= 0 && node_index < nodes_size()); - } - execution_plan_ = new_plan; - return kTfLiteOk; -} - -TfLiteStatus Interpreter::ResizeTensorImpl(TfLiteTensor* tensor, - TfLiteIntArray* new_size) { - // Note that in theory we could resize kTfLiteArenaRwPersistent tensors too. - if (tensor->allocation_type == kTfLiteArenaRw || - tensor->allocation_type == kTfLiteDynamic || - tensor->allocation_type == kTfLiteArenaRwPersistent) { - tensor_resized_since_op_invoke_ |= - TfLiteIntArrayEqual(tensor->dims, new_size) == 0; - if (tensor->type != kTfLiteString) { - size_t bytesRequired; - TfLiteStatus status = BytesRequired(tensor->type, new_size->data, - new_size->size, &bytesRequired); - if (status != kTfLiteOk) { - TfLiteIntArrayFree(new_size); - return kTfLiteError; - } - - // Realloc space for kTfLiteDynamic tensors. - TfLiteTensorRealloc(bytesRequired, tensor); - tensor->bytes = bytesRequired; - } - if (tensor->dims) TfLiteIntArrayFree(tensor->dims); - tensor->dims = new_size; - - if (tensor->allocation_type != kTfLiteDynamic) { - tensor->data.raw = nullptr; - } - } else { - // kTfLiteMmapRo tensors are stored in the flatbuffer and are therefore - // of fixed size. - TfLiteIntArrayFree(new_size); - ReportError(&context_, "Attempting to resize a fixed-size tensor."); - return kTfLiteError; - } - return kTfLiteOk; -} - -void Interpreter::UseNNAPI(bool enable) { - // TODO(aselle): This is a workaround for finding if NNAPI exists. - // We also need to make sure getLibraryHandle() is renamed to be NNAPI - // prefixed. - if (!NNAPIDelegate::IsSupported()) enable = false; - if (!enable) { - nnapi_delegate_.reset(); - } else if (!nnapi_delegate_) { - nnapi_delegate_.reset(new NNAPIDelegate); - } -} - -void Interpreter::SetNumThreads(int num_threads) { - context_.recommended_num_threads = num_threads; - - for (int i = 0; i < kTfLiteMaxExternalContexts; ++i) { - auto* c = external_contexts_[i]; - if (c && c->Refresh) { - c->Refresh(&context_); - } - } -} - -void Interpreter::SwitchToDelegateContext() { - context_.GetNodeAndRegistration = GetNodeAndRegistration; - context_.ReplaceSubgraphsWithDelegateKernels = - ReplaceSubgraphsWithDelegateKernels; - context_.GetExecutionPlan = GetExecutionPlan; -} - -void Interpreter::SwitchToKernelContext() { - SetForbiddenContextFunction(&context_.GetNodeAndRegistration); - SetForbiddenContextFunction(&context_.ReplaceSubgraphsWithDelegateKernels); - SetForbiddenContextFunction(&context_.GetExecutionPlan); -} - -TfLiteStatus Interpreter::ModifyGraphWithDelegate(TfLiteDelegate* delegate, - bool allow_dynamic_tensors) { - if (!allow_dynamic_tensors) { - int last_execution_plan_index_prepared; - TF_LITE_ENSURE_OK(&context_, PrepareOpsStartingAt( - 0, &last_execution_plan_index_prepared)); - - bool has_dynamic_tensors = true; - // Dynamic tensors exist if not all nodes can be prepared. - if (last_execution_plan_index_prepared + 1 == execution_plan_.size()) { - // If all the nodes can be prepared, check if the last node has dynamic - // tensors. - int node_index = execution_plan_[last_execution_plan_index_prepared]; - TfLiteNode& node = nodes_and_registration_[node_index].first; - if (!HasDynamicTensor(context_, node.outputs)) { - has_dynamic_tensors = false; - } - } - if (has_dynamic_tensors) { - ReportError( - &context_, - "Attempting to use a delegate that only supports static-sized " - "tensors with a graph that has dynamic-sized tensors."); - return kTfLiteError; - } - } - - // TODO(aselle): Consider if it is worth storing pointers to delegates. - // Setup additional context interface. - SwitchToDelegateContext(); - - TfLiteStatus status = delegate->Prepare(&context_, delegate); - - // Remove additional context info. - SwitchToKernelContext(); - - TF_LITE_ENSURE_OK(&context_, status); - - if (!allow_dynamic_tensors) { - // Reset the state to force tensor/op reallocation. - state_ = kStateUninvokable; - TF_LITE_ENSURE_OK(&context_, AllocateTensors()); - TF_LITE_ENSURE_EQ(&context_, state_, kStateInvokable); - // After using a delegate which doesn't support dynamic tensors, make the - // entire graph immutable. - state_ = kStateInvokableAndImmutable; - } - - return status; -} - -TfLiteStatus Interpreter::SetBufferHandle(int tensor_index, - TfLiteBufferHandle buffer_handle, - TfLiteDelegate* delegate) { - TF_LITE_ENSURE(&context_, tensor_index < tensors_size()); - TfLiteTensor* tensor = &tensors_[tensor_index]; - - TF_LITE_ENSURE(&context_, - tensor->delegate == nullptr || tensor->delegate == delegate); - tensor->delegate = delegate; - if (tensor->buffer_handle != kTfLiteNullBufferHandle) { - TF_LITE_ENSURE(&context_, tensor->delegate->FreeBufferHandle != nullptr); - tensor->delegate->FreeBufferHandle(&context_, tensor->delegate, - &tensor->buffer_handle); - } - tensor->buffer_handle = buffer_handle; - - return kTfLiteOk; -} - -TfLiteStatus Interpreter::GetBufferHandle(int tensor_index, - TfLiteBufferHandle* buffer_handle, - TfLiteDelegate** delegate) { - TF_LITE_ENSURE(&context_, tensor_index < tensors_size()); - TfLiteTensor* tensor = &tensors_[tensor_index]; - - *delegate = tensor->delegate; - *buffer_handle = tensor->buffer_handle; - - return kTfLiteOk; -} - -} // namespace tflite diff --git a/tensorflow/contrib/lite/interpreter.h b/tensorflow/contrib/lite/interpreter.h deleted file mode 100644 index 651a97e9dc84350569514528ae5635ec040d607f..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/interpreter.h +++ /dev/null @@ -1,695 +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. -==============================================================================*/ -// Main abstraction controlling the tflite interpreter. -// See context.h for the API for defining operations (TfLiteRegistration). -#ifndef TENSORFLOW_CONTRIB_LITE_INTERPRETER_H_ -#define TENSORFLOW_CONTRIB_LITE_INTERPRETER_H_ - -#include -#include -#include -#include - -#include "tensorflow/contrib/lite/allocation.h" -#include "tensorflow/contrib/lite/c/c_api_internal.h" -#include "tensorflow/contrib/lite/core/api/error_reporter.h" -#include "tensorflow/contrib/lite/memory_planner.h" -#include "tensorflow/contrib/lite/profiling/profiler.h" -#include "tensorflow/contrib/lite/stderr_reporter.h" - -namespace tflite { - -// Map statically from a c++ type to a TfLiteType (used below for safe casts). -template -constexpr TfLiteType typeToTfLiteType() { - return kTfLiteNoType; -} -template <> -constexpr TfLiteType typeToTfLiteType() { - return kTfLiteInt32; -} -template <> -constexpr TfLiteType typeToTfLiteType() { - return kTfLiteInt16; -} -template <> -constexpr TfLiteType typeToTfLiteType() { - return kTfLiteInt64; -} -template <> -constexpr TfLiteType typeToTfLiteType() { - return kTfLiteFloat32; -} -template <> -constexpr TfLiteType typeToTfLiteType() { - return kTfLiteUInt8; -} -template <> -constexpr TfLiteType typeToTfLiteType() { - return kTfLiteBool; -} -template <> -constexpr TfLiteType typeToTfLiteType>() { - return kTfLiteComplex64; -} -template <> -constexpr TfLiteType typeToTfLiteType() { - return kTfLiteString; -} - -// Forward declare since NNAPIDelegate uses Interpreter. -class NNAPIDelegate; - -// An interpreter for a graph of nodes that input and output from tensors. -// Each node of the graph processes a set of input tensors and produces a -// set of output Tensors. All inputs/output tensors are referenced by index. -// -// Usage: -// -// -- Create basic model -// Interpreter foo(2, 1); -// foo.SetTensorParametersReadWrite(0, ...); -// foo.SetTensorParametersReadOnly(1, ...); -// foo.SetNodeParameters(0, ...) -// -// -- Resize input array to 1 length. -// foo.ResizeInputTensor(0, 1); -// foo.AllocateTensors(); -// -- Install array data -// foo.typed_tensor(0)[0] = 3; -// foo.Invoke(); -// foo.typed_tensor(0)[0] = 4; -// foo.Invoke(); -// -- Resize input array and set data. -// foo.ResizeInputTensor(0, 2); -// foo.AllocateTensors(); -// foo.typed_tensor(0)[0] = 4; -// foo.typed_tensor(0)[1] = 8; -// foo.Invoke(); -// - -struct TfLiteIntArrayDeleter { - void operator()(TfLiteIntArray* a) { - if (a) TfLiteIntArrayFree(a); - } -}; - -class Interpreter { - public: - // Instantiate an interpreter. All errors associated with reading and - // processing this model will be forwarded to the error_reporter object. - // - // Note, if error_reporter is nullptr, then a default StderrReporter is - // used. Ownership of 'error_reporter' remains with the caller. - explicit Interpreter(ErrorReporter* error_reporter = DefaultErrorReporter()); - - ~Interpreter(); - - Interpreter(const Interpreter&) = delete; - Interpreter& operator=(const Interpreter&) = delete; - - // Functions to build interpreter - - // Provide a list of tensor indexes that are inputs to the model. - // Each index is bound check and this modifies the consistent_ flag of the - // interpreter. - TfLiteStatus SetInputs(std::vector inputs); - - // Provide a list of tensor indexes that are outputs to the model - // Each index is bound check and this modifies the consistent_ flag of the - // interpreter. - TfLiteStatus SetOutputs(std::vector outputs); - - // Provide a list of tensor indexes that are variable tensors. - // Each index is bound check and this modifies the consistent_ flag of the - // interpreter. - TfLiteStatus SetVariables(std::vector variables); - - // Ensure the internal node storage memory allocates at least `count` - // spots for node. NOTE, this doesn't actually add operators. This is an - // efficiency optimization that is subject to change. - void ReserveNodes(int count); - - // Adds a node with the given parameters and returns the index of the new - // node in `node_index` (optionally). Interpreter will take ownership of - // `builtin_data` and destroy it with `free`. Ownership of 'init_data' - // remains with the caller. - TfLiteStatus AddNodeWithParameters(const std::vector& inputs, - const std::vector& outputs, - const char* init_data, - size_t init_data_size, void* builtin_data, - const TfLiteRegistration* registration, - int* node_index = nullptr); - - // Adds `tensors_to_add` tensors, preserving pre-existing Tensor entries. - // The value pointed to by `first_new_tensor_index` will be set to the - // index of the first new tensor if `first_new_tensor_index` is non-null. - TfLiteStatus AddTensors(int tensors_to_add, - int* first_new_tensor_index = nullptr); - - // Set description of inputs/outputs/data/fptrs for node `node_index`. - // This variant assumes an external buffer has been allocated of size - // bytes. The lifetime of buffer must be ensured to be greater or equal - // to Interpreter. - inline TfLiteStatus SetTensorParametersReadOnly( - int tensor_index, TfLiteType type, const char* name, - const std::vector& dims, TfLiteQuantizationParams quantization, - const char* buffer, size_t bytes, - const Allocation* allocation = nullptr) { - return SetTensorParametersReadOnly(tensor_index, type, name, dims.size(), - dims.data(), quantization, buffer, bytes, - allocation); - } - - TfLiteStatus SetTensorParametersReadOnly( - int tensor_index, TfLiteType type, const char* name, const size_t rank, - const int* dims, TfLiteQuantizationParams quantization, - const char* buffer, size_t bytes, const Allocation* allocation = nullptr); - - // Set description of inputs/outputs/data/fptrs for node `node_index`. - // This variant assumes an external buffer has been allocated of size - // bytes. The lifetime of buffer must be ensured to be greater or equal - // to Interpreter. - inline TfLiteStatus SetTensorParametersReadWrite( - int tensor_index, TfLiteType type, const char* name, - const std::vector& dims, TfLiteQuantizationParams quantization, - bool is_variable = false) { - return SetTensorParametersReadWrite(tensor_index, type, name, dims.size(), - dims.data(), quantization, is_variable); - } - TfLiteStatus SetTensorParametersReadWrite( - int tensor_index, TfLiteType type, const char* name, const size_t rank, - const int* dims, TfLiteQuantizationParams quantization, - bool is_variable = false); - - // Functions to access tensor data - - // Read only access to list of inputs. - const std::vector& inputs() const { return inputs_; } - - // Return the name of a given input. The given index must be between 0 and - // inputs().size(). - const char* GetInputName(int index) const { - return context_.tensors[inputs_[index]].name; - } - - // Read only access to list of outputs. - const std::vector& outputs() const { return outputs_; } - - // Read only access to list of variable tensors. - const std::vector& variables() const { return variables_; } - - // Return the name of a given output. The given index must be between 0 and - // outputs().size(). - const char* GetOutputName(int index) const { - return context_.tensors[outputs_[index]].name; - } - - // Return the number of tensors in the model. - size_t tensors_size() const { return context_.tensors_size; } - - // Return the number of ops in the model. - size_t nodes_size() const { return nodes_and_registration_.size(); } - - // WARNING: Experimental interface, subject to change - const std::vector& execution_plan() const { return execution_plan_; } - - // WARNING: Experimental interface, subject to change - // Overrides execution plan. This bounds checks indices sent in. - TfLiteStatus SetExecutionPlan(const std::vector& new_plan); - - // Get a mutable tensor data structure. - // TODO(aselle): Create a safe ArrayHandle interface to avoid exposing this - // read/write access to structure - TfLiteTensor* tensor(int tensor_index) { - if (tensor_index >= context_.tensors_size || tensor_index < 0) - return nullptr; - return &context_.tensors[tensor_index]; - } - - // Get an immutable tensor data structure. - const TfLiteTensor* tensor(int tensor_index) const { - if (tensor_index >= context_.tensors_size || tensor_index < 0) - return nullptr; - return &context_.tensors[tensor_index]; - } - - // Get a pointer to an operation and registration data structure if in bounds. - const std::pair* node_and_registration( - int node_index) const { - if (node_index >= nodes_and_registration_.size() || node_index < 0) - return nullptr; - return &nodes_and_registration_[node_index]; - } - - // Perform a checked cast to the appropriate tensor type (mutable pointer - // version). - template - T* typed_tensor(int tensor_index) { - if (TfLiteTensor* tensor_ptr = tensor(tensor_index)) { - if (tensor_ptr->type == typeToTfLiteType()) { - return reinterpret_cast(tensor_ptr->data.raw); - } - } - return nullptr; - } - - // Perform a checked cast to the appropriate tensor type (immutable pointer - // version). - template - const T* typed_tensor(int tensor_index) const { - if (const TfLiteTensor* tensor_ptr = tensor(tensor_index)) { - if (tensor_ptr->type == typeToTfLiteType()) { - return reinterpret_cast(tensor_ptr->data.raw); - } - } - return nullptr; - } - - // Return a mutable pointer into the data of a given input tensor. The given - // index must be between 0 and inputs().size(). - template - T* typed_input_tensor(int index) { - return typed_tensor(inputs_[index]); - } - - // Return an immutable pointer into the data of a given input tensor. The - // given index must be between 0 and inputs().size(). - template - const T* typed_input_tensor(int index) const { - return typed_tensor(inputs_[index]); - } - - // Return a mutable pointer into the data of a given output tensor. The given - // index must be between 0 and outputs().size(). - template - T* typed_output_tensor(int index) { - return typed_tensor(outputs_[index]); - } - - // Return an immutable pointer into the data of a given output tensor. The - // given index must be between 0 and outputs().size(). - template - const T* typed_output_tensor(int index) const { - return typed_tensor(outputs_[index]); - } - - // Change the dimensionality of a given tensor. Note, this is only acceptable - // for tensor indices that are inputs. - // Returns status of failure or success. - // TODO(aselle): Consider implementing ArraySlice equivalent to make this - // more adept at accepting data without an extra copy. Use absl::ArraySlice - // if our partners determine that dependency is acceptable. - TfLiteStatus ResizeInputTensor(int tensor_index, - const std::vector& dims); - - // Update allocations for all tensors. This will redim dependent tensors using - // the input tensor dimensionality as given. This is relatively expensive. - // If you know that your sizes are not changing, you need not call this. - - // Returns status of success or failure. - TfLiteStatus AllocateTensors(); - - // Invoke the interpreter (run the whole graph in dependency order). - // - // NOTE: It is possible that the interpreter is not in a ready state - // to evaluate (i.e. if a ResizeTensor() has been performed without an - // AllocateTensors(). - // Returns status of success or failure. - TfLiteStatus Invoke(); - - // Enable or disable the NN API (true to enable) - void UseNNAPI(bool enable); - - // Set the number of threads available to the interpreter. - void SetNumThreads(int num_threads); - - // Allow float16 precision for FP32 calculation when possible. - // default: not allow. - // WARNING: This is an experimental API and subject to change. - void SetAllowFp16PrecisionForFp32(bool allow) { - context_.allow_fp32_relax_to_fp16 = allow; - } - - // Get the half precision flag. - // WARNING: This is an experimental API and subject to change. - bool GetAllowFp16PrecisionForFp32() const { - return context_.allow_fp32_relax_to_fp16; - } - - // Owning handle to a TfLiteDelegate instance. - using TfLiteDelegatePtr = - std::unique_ptr; - - // Allow a delegate to look at the graph and modify the graph to handle - // parts of the graph themselves. After this is called, the graph may - // contain new nodes that replace 1 more nodes. - // WARNING: This is an experimental API and subject to change. - TfLiteStatus ModifyGraphWithDelegate(TfLiteDelegate* delegate, - bool allow_dynamic_tensors = false); - - // Ensure the data in `tensor.data` is readable. In case delegate is used, - // it might require to copy the data from delegate buffer to raw memory. - // WARNING: This is an experimental API and subject to change. - TfLiteStatus EnsureTensorDataIsReadable(int tensor_index) { - TF_LITE_ENSURE(&context_, tensor_index < tensors_size()); - TfLiteTensor* tensor = &tensors_[tensor_index]; - if (tensor->data_is_stale) { - TF_LITE_ENSURE(&context_, tensor->delegate != nullptr); - TF_LITE_ENSURE(&context_, - tensor->buffer_handle != kTfLiteNullBufferHandle); - // This can be null if the delegate doesn't use its own buffer. - TF_LITE_ENSURE(&context_, - tensor->delegate->CopyFromBufferHandle != nullptr); - tensor->delegate->CopyFromBufferHandle(&context_, tensor->delegate, - tensor->buffer_handle, - tensor->data.raw, tensor->bytes); - tensor->data_is_stale = false; - } - return kTfLiteOk; - } - - // Set the delegate buffer handle to a tensor. It can be called in the - // following cases: - // 1. Set the buffer handle to a tensor that's not being written by a - // delegate. For example, feeding an OpenGL texture as the input of the - // inference graph. - // 2. Set the buffer handle to a tensor that uses the same delegate. - // For example, set an OpenGL texture as the output of inference, while - // the node which produces output is an OpenGL delegate node. - // WARNING: This is an experimental API and subject to change. - TfLiteStatus SetBufferHandle(int tensor_index, - TfLiteBufferHandle buffer_handle, - TfLiteDelegate* delegate); - - // Get the delegate buffer handle, and the delegate which can process the - // buffer handle. - // WARNING: This is an experimental API and subject to change. - TfLiteStatus GetBufferHandle(int tensor_index, - TfLiteBufferHandle* buffer_handle, - TfLiteDelegate** delegate); - - void SetProfiler(profiling::Profiler* profiler) { profiler_ = profiler; } - - profiling::Profiler* GetProfiler() { return profiler_; } - - // The default capacity of `tensors_` vector. - static constexpr int kTensorsReservedCapacity = 128; - // The capacity headroom of `tensors_` vector before calling ops' - // `prepare` and `invoke` function. In these functions, it's guaranteed - // allocating up to `kTensorsCapacityHeadroom` more tensors won't invalidate - // pointers to existing tensors. - static constexpr int kTensorsCapacityHeadroom = 16; - - // Set if buffer handle output is allowed. - // - // When using hardware delegation, Interpreter will make the data of output - // tensors available in `tensor->data` by default. If the application can - // consume the buffer handle directly (e.g. reading output from OpenGL - // texture), it can set this flag to false, so Interpreter won't copy the data - // from buffer handle to CPU memory. - // WARNING: This is an experimental API and subject to change. - void SetAllowBufferHandleOutput(bool allow_buffer_handle_output) { - allow_buffer_handle_output_ = allow_buffer_handle_output; - } - - // Reset all variable tensors to the default value. - // If a variable tensor doesn't have a buffer, reset it to zero. - // TODO(b/115961645): Implement - If a variable tensor has a buffer, reset it - // to the value of the buffer. - // WARNING: This is an experimental API and subject to change. - TfLiteStatus ResetVariableTensors(); - - // Retrieve an operator's description of its work, for profiling purposes. - const char* OpProfilingString(const TfLiteRegistration& op_reg, - const TfLiteNode* node) const { - if (op_reg.profiling_string == nullptr) return nullptr; - return op_reg.profiling_string(&context_, node); - } - - // Set the value of an external context. - void SetExternalContext(TfLiteExternalContextType type, - TfLiteExternalContext* ctx); - - private: - friend class InterpreterBuilder; - friend class InterpreterTest; - - // Prevent 'context_' from accessing functions that are only available to - // delegated kernels. - void SwitchToKernelContext(); - - // Add delegate-only functions to 'context_'. - void SwitchToDelegateContext(); - - // Give 'op_reg' a chance to initialize itself using the contents of - // 'buffer'. - void* OpInit(const TfLiteRegistration& op_reg, const char* buffer, - size_t length) { - if (op_reg.init == nullptr) return nullptr; - return op_reg.init(&context_, buffer, length); - } - - // Let 'op_reg' release any memory it might have allocated via 'OpInit'. - void OpFree(const TfLiteRegistration& op_reg, void* buffer) { - if (op_reg.free == nullptr) return; - if (buffer) { - op_reg.free(&context_, buffer); - } - } - - // Prepare the given 'node' for execution. - TfLiteStatus OpPrepare(const TfLiteRegistration& op_reg, TfLiteNode* node) { - if (op_reg.prepare == nullptr) return kTfLiteOk; - return op_reg.prepare(&context_, node); - } - - // Invoke the operator represented by 'node'. - TfLiteStatus OpInvoke(const TfLiteRegistration& op_reg, TfLiteNode* node) { - if (op_reg.invoke == nullptr) return kTfLiteError; - return op_reg.invoke(&context_, node); - } - - // Call OpPrepare() for as many ops as possible, allocating memory for their - // tensors. If an op containing dynamic tensors is found, preparation will be - // postponed until this function is called again. This allows the interpreter - // to wait until Invoke() to resolve the sizes of dynamic tensors. - TfLiteStatus PrepareOpsAndTensors(); - - // Call OpPrepare() for all ops starting at 'first_node'. Stop when a - // dynamic tensors is found or all ops have been prepared. Fill - // 'last_node_prepared' with the id of the op containing dynamic tensors, or - // the last in the graph. - TfLiteStatus PrepareOpsStartingAt(int first_execution_plan_index, - int* last_execution_plan_index_prepared); - - // Tensors needed by the interpreter. Use `AddTensors` to add more blank - // tensor entries. Note, `tensors_.data()` needs to be synchronized to the - // `context_` whenever this std::vector is reallocated. Currently this - // only happens in `AddTensors()`. - std::vector tensors_; - - // Check if an array of tensor indices are valid with respect to the Tensor - // array. - // NOTE: this changes consistent_ to be false if indices are out of bounds. - TfLiteStatus CheckTensorIndices(const char* label, const int* indices, - int length); - - // Compute the number of bytes required to represent a tensor with dimensions - // specified by the array dims (of length dims_size). Returns the status code - // and bytes. - TfLiteStatus BytesRequired(TfLiteType type, const int* dims, size_t dims_size, - size_t* bytes); - - // Request an tensor be resized implementation. If the given tensor is of - // type kTfLiteDynamic it will also be allocated new memory. - TfLiteStatus ResizeTensorImpl(TfLiteTensor* tensor, TfLiteIntArray* new_size); - - // Report a detailed error string (will be printed to stderr). - // TODO(aselle): allow user of class to provide alternative destinations. - void ReportErrorImpl(const char* format, va_list args); - - // Entry point for C node plugin API to request an tensor be resized. - static TfLiteStatus ResizeTensor(TfLiteContext* context, TfLiteTensor* tensor, - TfLiteIntArray* new_size); - // Entry point for C node plugin API to report an error. - static void ReportError(TfLiteContext* context, const char* format, ...); - - // Entry point for C node plugin API to add new tensors. - static TfLiteStatus AddTensors(TfLiteContext* context, int tensors_to_add, - int* first_new_tensor_index); - - // WARNING: This is an experimental API and subject to change. - // Entry point for C API ReplaceSubgraphsWithDelegateKernels - static TfLiteStatus ReplaceSubgraphsWithDelegateKernels( - TfLiteContext* context, TfLiteRegistration registration, - const TfLiteIntArray* nodes_to_replace, TfLiteDelegate* delegate); - - // Update the execution graph to replace some of the nodes with stub - // nodes. Specifically any node index that has `nodes[index]==1` will be - // slated for replacement with a delegate kernel specified by registration. - // Ownership of 'nodes_to_replace' and 'delegate' remains with the caller. - // WARNING: This is an experimental interface that is subject to change. - TfLiteStatus ReplaceSubgraphsWithDelegateKernels( - TfLiteRegistration registration, const TfLiteIntArray* nodes_to_replace, - TfLiteDelegate* delegate); - - // WARNING: This is an experimental interface that is subject to change. - // Gets the internal pointer to a TensorFlow lite node by node_index. - TfLiteStatus GetNodeAndRegistration(int node_index, TfLiteNode** node, - TfLiteRegistration** registration); - - // WARNING: This is an experimental interface that is subject to change. - // Entry point for C node plugin API to get a node by index. - static TfLiteStatus GetNodeAndRegistration(struct TfLiteContext*, - int node_index, TfLiteNode** node, - TfLiteRegistration** registration); - - // WARNING: This is an experimental interface that is subject to change. - // Gets an TfLiteIntArray* representing the execution plan. The interpreter - // owns this memory and it is only guaranteed to exist during the invocation - // of the delegate prepare. - TfLiteStatus GetExecutionPlan(TfLiteIntArray** execution_plan); - - // WARNING: This is an experimental interface that is subject to change. - // Entry point for C node plugin API to get the execution plan. - static TfLiteStatus GetExecutionPlan(struct TfLiteContext* context, - TfLiteIntArray** execution_plan); - - // Retrieve an existing external context by type. - TfLiteExternalContext* GetExternalContext(TfLiteExternalContextType type); - static TfLiteExternalContext* GetExternalContext( - struct TfLiteContext* context, TfLiteExternalContextType type); - - // Set the value of an external context. - static void SetExternalContext(struct TfLiteContext* context, - TfLiteExternalContextType type, - TfLiteExternalContext* ctx); - - // Variant of the public ModifyGraphWithDelegate method that additionally - // Assumes ownership of the provided delegate. - // WARNING: This is an experimental API and subject to change. - TfLiteStatus ModifyGraphWithDelegate(TfLiteDelegatePtr delegate, - bool allow_dynamic_tensors = false) { - // Note that we retain ownership of the delegate even if graph modification - // fails, as delegate use will be in an indeterminate state at that point. - owned_delegates_.push_back(std::move(delegate)); - return ModifyGraphWithDelegate(owned_delegates_.back().get(), - allow_dynamic_tensors); - } - - // Ensures that `tensors_` has at least `kTensorsCapacityHeadroom` extra - // capacity. Calling this function may invalidate existing pointers to - // tensors. After calling this function, adding `kTensorsCapacityHeadroom` - // more tensors won't invalidate the pointer to existing tensors. - void EnsureTensorsVectorCapacity() { - const size_t required_capacity = tensors_size() + kTensorsCapacityHeadroom; - if (required_capacity > tensors_.capacity()) { - tensors_.reserve(required_capacity); - context_.tensors = tensors_.data(); - } - } - - // The state of the Interpreter. - enum State { - // The interpreter isn't ready to be invoked. - // `AllocateTensor` need to be called to enter an invokable state. - kStateUninvokable = 0, - // The interpreter is ready to be invoked. - kStateInvokable, - // The interpreter is ready to be invoked, and graph can't be further - // modified. The interpreter will enter this state when calling - // `ModifyGraphWithDelegate` with `allow_dynamic_tensors=false`. - kStateInvokableAndImmutable, - }; - State state_ = kStateUninvokable; - - // A pure C data structure used to communicate with the pure C plugin - // interface. To avoid copying tensor metadata, this is also the definitive - // structure to store tensors. - TfLiteContext context_; - - // Node inputs/outputs are stored in TfLiteNode and TfLiteRegistration stores - // function pointers to actual implementation. - std::vector> - nodes_and_registration_; - - // Whether the model is consistent. That is to say if the inputs and outputs - // of every node and the global inputs and outputs are valid indexes into - // the tensor array. - bool consistent_ = true; - - // Array of indices representing the tensors that are inputs to the - // interpreter. - std::vector inputs_; - - // Array of indices representing the tensors that are outputs to the - // interpreter. - std::vector outputs_; - - // Array of indices representing the tensors that are variable tensors. - std::vector variables_; - - // The error reporter delegate that tflite will forward queries errors to. - ErrorReporter* error_reporter_; - - // Index of the next node to prepare. - // During Invoke(), Interpreter will allocate input tensors first, which are - // known to be fixed size. Then it will allocate outputs from nodes as many - // as possible. When there is a node that produces dynamic sized tensor. - // Interpreter will stop allocating tensors, set the value of next allocate - // node id, and execute the node to generate the output tensor before continue - // to allocate successors. This process repeats until all nodes are executed. - // NOTE: this relies on the order of nodes that is in topological order. - int next_execution_plan_index_to_prepare_; - - // WARNING: This is an experimental interface that is subject to change. - // This is a list of node indices (to index into nodes_and_registration). - // This represents a valid topological sort (dependency ordered) execution - // plan. In particular, it is valid for this ordering to contain only a - // subset of the node indices. - std::vector execution_plan_; - - // In the future, we'd like a TfLiteIntArray compatible representation. - // TODO(aselle): replace execution_plan_ with this. - std::unique_ptr plan_cache_; - - // Whether to delegate to NN API - std::unique_ptr nnapi_delegate_; - - // List of delegates that have been installed and are owned by this - // interpreter instance. Useful if client delegate ownership is burdensome. - // WARNING: This is an experimental API and subject to change. - // TODO(b/116667551): Use TfLiteExternalContext for storing state. - std::vector owned_delegates_; - - std::unique_ptr memory_planner_; - - bool allow_buffer_handle_output_ = false; - - // Tracking bit for whether a tensor was resized in the course of an op - // invocation. This is a useful hint to ensure that dynamic tensor outputs - // trigger downstream reallocation after op invocation. - bool tensor_resized_since_op_invoke_ = false; - - // Profiler for this interpreter instance. - profiling::Profiler* profiler_ = nullptr; - - // List of active external contexts. - TfLiteExternalContext* external_contexts_[kTfLiteMaxExternalContexts]; -}; - -} // namespace tflite -#endif // TENSORFLOW_CONTRIB_LITE_INTERPRETER_H_ diff --git a/tensorflow/contrib/lite/java/BUILD b/tensorflow/contrib/lite/java/BUILD deleted file mode 100644 index cab8d5277f2d3f539e7a69f15ebda20821b19a3b..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/java/BUILD +++ /dev/null @@ -1,229 +0,0 @@ -# Description: -# TensorFlow Lite Java API. - -package(default_visibility = [ - "//tensorflow/contrib/lite/java/ovic:__pkg__", -]) - -licenses(["notice"]) # Apache 2.0 - -load("//tensorflow/java:build_defs.bzl", "JAVACOPTS") -load("//tensorflow/contrib/lite:build_def.bzl", "tflite_jni_binary") -load("//tensorflow/contrib/lite/java:aar_with_jni.bzl", "aar_with_jni") - -JAVA_SRCS = glob([ - "src/main/java/org/tensorflow/lite/*.java", -]) - -# Building tensorflow-lite.aar including 4 variants of .so -# To build an aar for release, run below command: -# bazel build --cxxopt='--std=c++11' -c opt --fat_apk_cpu=x86,x86_64,arm64-v8a,armeabi-v7a \ -# tensorflow/contrib/lite/java:tensorflow-lite -aar_with_jni( - name = "tensorflow-lite", - android_library = ":tensorflowlite", -) - -# EXPERIMENTAL: AAR target that supports TensorFlow op execution with TFLite. -aar_with_jni( - name = "tensorflow-lite-flex", - android_library = ":tensorflowlite_flex", -) - -android_library( - name = "tensorflowlite", - srcs = JAVA_SRCS, - manifest = "AndroidManifest.xml", - visibility = ["//visibility:public"], - deps = [ - ":tensorflowlite_native", - "@org_checkerframework_qual", - ], -) - -# EXPERIMENTAL: Android target that supports TensorFlow op execution with TFLite. -android_library( - name = "tensorflowlite_flex", - srcs = JAVA_SRCS, - manifest = "AndroidManifest.xml", - visibility = ["//visibility:public"], - deps = [ - ":tensorflowlite_native_flex", - "@org_checkerframework_qual", - ], -) - -android_library( - name = "tensorflowlite_java", - srcs = JAVA_SRCS, - visibility = ["//visibility:public"], - deps = [ - "@org_checkerframework_qual", - ], -) - -java_library( - name = "tensorflowlitelib", - srcs = JAVA_SRCS, - javacopts = JAVACOPTS, - visibility = ["//visibility:public"], - deps = [ - ":libtensorflowlite_jni.so", - "@org_checkerframework_qual", - ], -) - -# EXPERIMENTAL: Java target that supports TensorFlow op execution with TFLite. -java_library( - name = "tensorflowlitelib_flex", - srcs = JAVA_SRCS, - javacopts = JAVACOPTS, - visibility = ["//visibility:public"], - deps = [ - ":libtensorflowlite_flex_jni.so", - "@org_checkerframework_qual", - ], -) - -java_test( - name = "TensorFlowLiteTest", - size = "small", - srcs = ["src/test/java/org/tensorflow/lite/TensorFlowLiteTest.java"], - javacopts = JAVACOPTS, - tags = ["no_oss"], - test_class = "org.tensorflow.lite.TensorFlowLiteTest", - deps = [ - ":tensorflowlitelib", - "@com_google_truth", - "@junit", - ], -) - -java_test( - name = "DataTypeTest", - size = "small", - srcs = ["src/test/java/org/tensorflow/lite/DataTypeTest.java"], - javacopts = JAVACOPTS, - tags = ["no_oss"], - test_class = "org.tensorflow.lite.DataTypeTest", - deps = [ - ":tensorflowlitelib", - "@com_google_truth", - "@junit", - ], -) - -java_test( - name = "NativeInterpreterWrapperTest", - size = "small", - srcs = ["src/test/java/org/tensorflow/lite/NativeInterpreterWrapperTest.java"], - data = [ - "src/testdata/add.bin", - "src/testdata/int32.bin", - "src/testdata/int64.bin", - "src/testdata/invalid_model.bin", - "src/testdata/uint8.bin", - "src/testdata/with_custom_op.lite", - ], - javacopts = JAVACOPTS, - tags = ["no_oss"], - test_class = "org.tensorflow.lite.NativeInterpreterWrapperTest", - deps = [ - ":tensorflowlitelib", - "@com_google_truth", - "@junit", - ], -) - -# TODO: generate large models at runtime, instead of storing them. -java_test( - name = "InterpreterTest", - size = "small", - srcs = ["src/test/java/org/tensorflow/lite/InterpreterTest.java"], - data = [ - "src/testdata/add.bin", - "src/testdata/mobilenet.tflite.bin", - "//tensorflow/contrib/lite:testdata/multi_add_flex.bin", - ], - javacopts = JAVACOPTS, - tags = ["no_oss"], - test_class = "org.tensorflow.lite.InterpreterTest", - visibility = ["//visibility:private"], - deps = [ - ":tensorflowlitelib", - "@com_google_truth", - "@junit", - ], -) - -java_test( - name = "InterpreterFlexTest", - size = "small", - srcs = ["src/test/java/org/tensorflow/lite/InterpreterFlexTest.java"], - data = [ - "//tensorflow/contrib/lite:testdata/multi_add_flex.bin", - ], - javacopts = JAVACOPTS, - tags = ["no_oss"], - test_class = "org.tensorflow.lite.InterpreterFlexTest", - visibility = ["//visibility:private"], - deps = [ - ":tensorflowlitelib_flex", - "@com_google_truth", - "@junit", - ], -) - -java_test( - name = "TensorTest", - size = "small", - srcs = ["src/test/java/org/tensorflow/lite/TensorTest.java"], - data = [ - "src/testdata/add.bin", - ], - javacopts = JAVACOPTS, - tags = ["no_oss"], - test_class = "org.tensorflow.lite.TensorTest", - deps = [ - ":tensorflowlitelib", - "@com_google_truth", - "@junit", - ], -) - -filegroup( - name = "libtensorflowlite_jni", - srcs = select({ - "//conditions:default": [":libtensorflowlite_jni.so"], - }), - visibility = ["//visibility:public"], -) - -cc_library( - name = "tensorflowlite_native", - srcs = ["libtensorflowlite_jni.so"], - visibility = ["//visibility:public"], -) - -cc_library( - name = "tensorflowlite_native_flex", - srcs = ["libtensorflowlite_flex_jni.so"], - visibility = ["//visibility:public"], -) - -tflite_jni_binary( - name = "libtensorflowlite_jni.so", - deps = [ - "//tensorflow/contrib/lite/java/src/main/native", - ], -) - -# EXPERIMENTAL: Native target that supports TensorFlow op execution with TFLite. -tflite_jni_binary( - name = "libtensorflowlite_flex_jni.so", - deps = [ - "//tensorflow/contrib/lite/delegates/flex:delegate", - "//tensorflow/contrib/lite/java/src/main/native", - "//tensorflow/contrib/lite/java/src/main/native:init_tensorflow", - ], -) diff --git a/tensorflow/contrib/lite/java/demo/README.md b/tensorflow/contrib/lite/java/demo/README.md deleted file mode 100644 index c04b2a61942430108891c612ae410d04d373c840..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/java/demo/README.md +++ /dev/null @@ -1,53 +0,0 @@ -# TF Lite Android Image Classifier App Example - -A simple Android example that demonstrates image classification using the camera. - -## Building in Android Studio with TensorFlow Lite AAR from JCenter. -The build.gradle is configured to use TensorFlow Lite's nightly build. - -If you see a build error related to compatibility with Tensorflow Lite's Java API (example: method X is -undefined for type Interpreter), there has likely been a backwards compatible -change to the API. You will need to pull new app code that's compatible with the -nightly build and may need to first wait a few days for our external and internal -code to merge. - -## Building from Source with Bazel - -1. Follow the [Bazel steps for the TF Demo App](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/examples/android#bazel): - - 1. [Install Bazel and Android Prerequisites](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/examples/android#install-bazel-and-android-prerequisites). - It's easiest with Android Studio. - - - You'll need at least SDK version 23. - - Make sure to install the latest version of Bazel. Some distributions - ship with Bazel 0.5.4, which is too old. - - Bazel requires Android Build Tools `26.0.1` or higher. - - You also need to install the Android Support Repository, available - through Android Studio under `Android SDK Manager -> SDK Tools -> - Android Support Repository`. - - 2. [Edit your `WORKSPACE`](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/examples/android#edit-workspace) - to add SDK and NDK targets. - - NOTE: As long as you have the SDK and NDK installed, the `./configure` - script will create these rules for you. Answer "Yes" when the script asks - to automatically configure the `./WORKSPACE`. - - - Make sure the `api_level` in `WORKSPACE` is set to an SDK version that - you have installed. - - By default, Android Studio will install the SDK to `~/Android/Sdk` and - the NDK to `~/Android/Sdk/ndk-bundle`. - -2. Build the app with Bazel. The demo needs C++11: - - ```shell - bazel build -c opt --cxxopt='--std=c++11' \ - //tensorflow/contrib/lite/java/demo/app/src/main:TfLiteCameraDemo - ``` - -3. Install the demo on a - [debug-enabled device](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/examples/android#install): - - ```shell - adb install bazel-bin/tensorflow/contrib/lite/java/demo/app/src/main/TfLiteCameraDemo.apk - ``` diff --git a/tensorflow/contrib/lite/java/demo/app/build.gradle b/tensorflow/contrib/lite/java/demo/app/build.gradle deleted file mode 100644 index 05301ebf88c12cc95f71d5efd74062d76e598e1d..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/java/demo/app/build.gradle +++ /dev/null @@ -1,89 +0,0 @@ -apply plugin: 'com.android.application' - -android { - compileSdkVersion 26 - buildToolsVersion "26.0.1" - defaultConfig { - applicationId "android.example.com.tflitecamerademo" - // Required by Camera2 API. - minSdkVersion 21 - targetSdkVersion 26 - versionCode 1 - versionName "1.0" - - // Remove this block. - jackOptions { - enabled true - } - } - lintOptions { - abortOnError false - } - buildTypes { - release { - minifyEnabled false - proguardFiles getDefaultProguardFile('proguard-android.txt'), 'proguard-rules.pro' - } - } - aaptOptions { - noCompress "tflite" - } - - compileOptions { - sourceCompatibility JavaVersion.VERSION_1_8 - targetCompatibility JavaVersion.VERSION_1_8 - } -} - -repositories { - maven { - url 'https://google.bintray.com/tensorflow' - } -} - -dependencies { - compile fileTree(dir: 'libs', include: ['*.jar']) - compile 'com.android.support:appcompat-v7:25.2.0' - compile 'com.android.support.constraint:constraint-layout:1.0.2' - compile 'com.android.support:design:25.2.0' - compile 'com.android.support:support-annotations:25.3.1' - compile 'com.android.support:support-v13:25.2.0' - - compile 'org.tensorflow:tensorflow-lite:0.0.0-nightly' -} - -def modelDownloadUrl = "https://storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_v1_224_android_quant_2017_11_08.zip" -def localCache = "build/intermediates/mobilenet_v1_224_android_quant_2017_11_08.zip" -def targetFolder = "src/main/assets" - -task downloadModel(type: DownloadUrlTask) { - doFirst { - println "Downloading ${modelDownloadUrl}" - } - sourceUrl = "${modelDownloadUrl}" - target = file("${localCache}") -} - -task unzipModel(type: Copy, dependsOn: 'downloadModel') { - doFirst { - println "Unzipping ${localCache}" - } - from zipTree("${localCache}") - into "${targetFolder}" -} - -// Ensure the model file is downloaded and extracted before every build -preBuild.dependsOn unzipModel - -class DownloadUrlTask extends DefaultTask { - @Input - String sourceUrl - - @OutputFile - File target - - @TaskAction - void download() { - ant.get(src: sourceUrl, dest: target) - } -} diff --git a/tensorflow/contrib/lite/java/demo/app/src/main/BUILD b/tensorflow/contrib/lite/java/demo/app/src/main/BUILD deleted file mode 100644 index 5ad738389eb8bc1d875fc888c1336fb3fa140eee..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/java/demo/app/src/main/BUILD +++ /dev/null @@ -1,32 +0,0 @@ -load("@build_bazel_rules_android//android:rules.bzl", "android_binary") - -package(default_visibility = ["//visibility:private"]) - -licenses(["notice"]) # Apache 2.0 - -android_binary( - name = "TfLiteCameraDemo", - srcs = glob(["java/**/*.java"]), - aapt_version = "aapt", - assets = [ - "//tensorflow/contrib/lite/java/demo/app/src/main/assets:labels_mobilenet_quant_v1_224.txt", - "@tflite_mobilenet//:mobilenet_quant_v1_224.tflite", - ], - assets_dir = "", - custom_package = "com.example.android.tflitecamerademo", - manifest = "AndroidManifest.xml", - nocompress_extensions = [ - ".tflite", - ], - resource_files = glob(["res/**"]), - # In some platforms we don't have an Android SDK/NDK and this target - # can't be built. We need to prevent the build system from trying to - # use the target in that case. - tags = ["manual"], - deps = [ - "//tensorflow/contrib/lite/java:tensorflowlite", - "//tensorflow/contrib/lite/java/src/testhelper/java/org/tensorflow/lite:testhelper", - "@androidsdk//com.android.support:support-v13-25.2.0", - "@androidsdk//com.android.support:support-v4-25.2.0", - ], -) diff --git a/tensorflow/contrib/lite/java/demo/app/src/main/res/layout-land/fragment_camera2_basic.xml b/tensorflow/contrib/lite/java/demo/app/src/main/res/layout-land/fragment_camera2_basic.xml deleted file mode 100644 index ef8a9e08450d72e392815756606f5ef8301cdd58..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/java/demo/app/src/main/res/layout-land/fragment_camera2_basic.xml +++ /dev/null @@ -1,67 +0,0 @@ - - - - - - - - - - - - - - - - - - diff --git a/tensorflow/contrib/lite/java/demo/app/src/main/res/layout-v26/fragment_camera2_basic.xml b/tensorflow/contrib/lite/java/demo/app/src/main/res/layout-v26/fragment_camera2_basic.xml deleted file mode 100644 index ddb099a950c2f83d7b2867f8f35d96885229536d..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/java/demo/app/src/main/res/layout-v26/fragment_camera2_basic.xml +++ /dev/null @@ -1,95 +0,0 @@ - - - - - - - - - - - - - - - - - - - - diff --git a/tensorflow/contrib/lite/java/demo/app/src/main/res/layout/fragment_camera2_basic.xml b/tensorflow/contrib/lite/java/demo/app/src/main/res/layout/fragment_camera2_basic.xml deleted file mode 100644 index e567009a424ed77384bee193c47d4f4d253f5767..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/java/demo/app/src/main/res/layout/fragment_camera2_basic.xml +++ /dev/null @@ -1,94 +0,0 @@ - - - - - - - - - - - - - - - - - - - diff --git a/tensorflow/contrib/lite/java/demo/build.gradle b/tensorflow/contrib/lite/java/demo/build.gradle deleted file mode 100644 index b78a0b86c939620b6f05483ce45c4d3ef0ef595e..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/java/demo/build.gradle +++ /dev/null @@ -1,23 +0,0 @@ -// Top-level build file where you can add configuration options common to all sub-projects/modules. - -buildscript { - repositories { - jcenter() - } - dependencies { - classpath 'com.android.tools.build:gradle:2.3.1' - - // NOTE: Do not place your application dependencies here; they belong - // in the individual module build.gradle files - } -} - -allprojects { - repositories { - jcenter() - } -} - -task clean(type: Delete) { - delete rootProject.buildDir -} diff --git a/tensorflow/contrib/lite/java/demo/gradle/wrapper/gradle-wrapper.properties b/tensorflow/contrib/lite/java/demo/gradle/wrapper/gradle-wrapper.properties deleted file mode 100644 index fa7a38a0e43eecd1e7292dd49efa79a5d0742e2a..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/java/demo/gradle/wrapper/gradle-wrapper.properties +++ /dev/null @@ -1,6 +0,0 @@ -#Thu Sep 28 09:01:41 PDT 2017 -distributionBase=GRADLE_USER_HOME -distributionPath=wrapper/dists -zipStoreBase=GRADLE_USER_HOME -zipStorePath=wrapper/dists -distributionUrl=https\://services.gradle.org/distributions/gradle-3.3-all.zip diff --git a/tensorflow/contrib/lite/java/ovic/BUILD b/tensorflow/contrib/lite/java/ovic/BUILD deleted file mode 100644 index ea9b9ed4b66a601981f4c402f7f8a4f6749e07fd..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/java/ovic/BUILD +++ /dev/null @@ -1,132 +0,0 @@ -# Description: -# OVIC Benchmarker Java API. - -load("@build_bazel_rules_android//android:rules.bzl", "android_library") - -package(default_visibility = ["//visibility:public"]) - -licenses(["notice"]) # Apache 2.0 - -load("//tensorflow/java:build_defs.bzl", "JAVACOPTS") - -# Build targets for OVIC classification. -java_test( - name = "OvicClassifierTest", - size = "medium", - srcs = ["src/test/java/org/tensorflow/ovic/OvicClassifierTest.java"], - data = [ - "//tensorflow/contrib/lite/java/ovic/src/testdata:labels.txt", - "//tensorflow/contrib/lite/java/ovic/src/testdata:ovic_testdata", - ], - javacopts = JAVACOPTS, - tags = ["no_oss"], - test_class = "org.tensorflow.ovic.OvicClassifierTest", - visibility = ["//visibility:public"], - deps = [ - "//tensorflow/contrib/lite/java/ovic:ovicbenchmarkerlib_java", - "@com_google_truth", - "@junit", - ], -) - -java_binary( - name = "ovic_validator", - srcs = ["src/main/java/org/tensorflow/ovic/OvicValidator.java"], - data = [ - "//tensorflow/contrib/lite/java/ovic/src/testdata:labels.txt", - ], - main_class = "org.tensorflow.ovic.OvicValidator", - tags = ["no_oss"], - deps = [ - "//tensorflow/contrib/lite/java/ovic:ovicbenchmarkerlib_java", - ], -) - -android_library( - name = "ovicbenchmarkerlib", - srcs = [ - "src/main/java/org/tensorflow/ovic/OvicBenchmarker.java", - "src/main/java/org/tensorflow/ovic/OvicClassificationResult.java", - "src/main/java/org/tensorflow/ovic/OvicClassifier.java", - "src/main/java/org/tensorflow/ovic/OvicClassifierBenchmarker.java", - ], - manifest = "//tensorflow/contrib/lite/java:AndroidManifest.xml", - tags = ["no_oss"], - deps = [ - "//tensorflow/contrib/lite/java:tensorflowlite", - "//tensorflow/contrib/lite/java/src/testhelper/java/org/tensorflow/lite:testhelper", - "@org_checkerframework_qual", - ], -) - -java_library( - name = "ovicbenchmarkerlib_java", - srcs = [ - "src/main/java/org/tensorflow/ovic/OvicClassificationResult.java", - "src/main/java/org/tensorflow/ovic/OvicClassifier.java", - ], - javacopts = JAVACOPTS, - tags = ["no_oss"], - deps = [ - "//tensorflow/contrib/lite/java:libtensorflowlite_jni.so", - "//tensorflow/contrib/lite/java:tensorflowlite_java", - "//tensorflow/contrib/lite/java/src/main/native", - "//tensorflow/contrib/lite/java/src/testhelper/java/org/tensorflow/lite:testhelper", - "@org_checkerframework_qual", - ], -) - -# Build targets for OVIC detection. -java_test( - name = "OvicDetectorTest", - size = "medium", - srcs = ["src/test/java/org/tensorflow/ovic/OvicDetectorTest.java"], - data = [ - "//tensorflow/contrib/lite/java/ovic/src/testdata:coco_labels.txt", - "//tensorflow/contrib/lite/java/ovic/src/testdata:ovic_testdata", - "@tflite_mobilenet_ssd_quant//:detect.tflite", - ], - javacopts = JAVACOPTS, - tags = ["no_oss"], - test_class = "org.tensorflow.ovic.OvicDetectorTest", - visibility = ["//visibility:public"], - deps = [ - "//tensorflow/contrib/lite/java/ovic:ovicdetectionbenchmarkerlib_java", - "@com_google_truth", - "@junit", - ], -) - -android_library( - name = "ovicdetectionbenchmarkerlib", - srcs = [ - "src/main/java/org/tensorflow/ovic/BoundingBox.java", - "src/main/java/org/tensorflow/ovic/OvicBenchmarker.java", - "src/main/java/org/tensorflow/ovic/OvicDetectionResult.java", - "src/main/java/org/tensorflow/ovic/OvicDetector.java", - "src/main/java/org/tensorflow/ovic/OvicDetectorBenchmarker.java", - ], - manifest = "//tensorflow/contrib/lite/java:AndroidManifest.xml", - deps = [ - "//tensorflow/contrib/lite/java:tensorflowlite", - "//tensorflow/contrib/lite/java/src/testhelper/java/org/tensorflow/lite:testhelper", - "@org_checkerframework_qual", - ], -) - -java_library( - name = "ovicdetectionbenchmarkerlib_java", - srcs = [ - "src/main/java/org/tensorflow/ovic/BoundingBox.java", - "src/main/java/org/tensorflow/ovic/OvicDetectionResult.java", - "src/main/java/org/tensorflow/ovic/OvicDetector.java", - ], - javacopts = JAVACOPTS, - deps = [ - "//tensorflow/contrib/lite/java:libtensorflowlite_jni.so", - "//tensorflow/contrib/lite/java:tensorflowlite_java", - "//tensorflow/contrib/lite/java/src/main/native", - "//tensorflow/contrib/lite/java/src/testhelper/java/org/tensorflow/lite:testhelper", - "@org_checkerframework_qual", - ], -) diff --git a/tensorflow/contrib/lite/java/ovic/README.md b/tensorflow/contrib/lite/java/ovic/README.md deleted file mode 100644 index df77bfaab3251c0ebe2e377e84d11965fdb821dd..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/java/ovic/README.md +++ /dev/null @@ -1,158 +0,0 @@ -# Benchmarker for LPIRC Workshop at CVPR 2018 - -This folder contains building code for track one of the [Low Power ImageNet Recognition Challenge workshop at CVPR 2018.](https://rebootingcomputing.ieee.org/home/sitemap/14-lpirc/80-low-power-image-recognition-challenge-lpirc-2018) - -## Pre-requisite - -Follow the steps [here](https://www.tensorflow.org/lite/demo_android) to install Tensorflow, Bazel, and the Android NDK and SDK. - -## Test the benchmarker: - -The testing utilities helps the developers (you) to make sure that your submissions in TfLite format will be processed as expected in the competition's benchmarking system. - -Note: for now the tests only provides correctness checks, i.e. classifier predicts the correct category on the test image, but no on-device latency measurements. To test the latency measurement functionality, the tests will print the latency running on a desktop computer, which is not indicative of the on-device run-time. -We are releasing an benchmarker Apk that would allow developers to measure latency on their own devices. - -### Obtain the sample models - -The test data (models and images) should be downloaded automatically for you by Bazel. In case they are not, you can manually install them as below. - -Note: all commands should be called from your tensorflow installation folder (under this folder you should find `tensorflow/contrib/lite`). - - -* Download the [testdata package](https://storage.googleapis.com/download.tensorflow.org/data/ovic.zip): - -```sh -curl -L https://storage.googleapis.com/download.tensorflow.org/data/ovic.zip -o /tmp/ovic.zip -``` - -* Unzip the package into the testdata folder: - -```sh -unzip -j /tmp/ovic.zip -d tensorflow/contrib/lite/java/ovic/src/testdata/ -``` - -### Run tests - -You can run test with Bazel as below. This helps to ensure that the installation is correct. - -```sh -bazel test --cxxopt=--std=c++11 //tensorflow/contrib/lite/java/ovic:OvicClassifierTest --cxxopt=-Wno-all --test_output=all -``` - -### Test your submissions - -Once you have a submission that follows the instructions from the [competition site](https://rebootingcomputing.ieee.org/home/sitemap/14-lpirc/80-low-power-image-recognition-challenge-lpirc-2018), you can verify it in two ways: - -#### Validate using randomly generated images - -You can call the validator binary below to verify that your model fits the format requirements. This often helps you to catch size mismatches (e.g. output should be [1, 1001] instead of [1,1,1,1001]). Let say the submission file is located at `/path/to/my_model.lite`, then call: - -```sh -bazel build --cxxopt=--std=c++11 //tensorflow/contrib/lite/java/ovic:ovic_validator --cxxopt=-Wno-all -bazel-bin/tensorflow/contrib/lite/java/ovic/ovic_validator /path/to/my_model.lite -``` - -Successful validation should print the following message to terminal: - -``` -Successfully validated /path/to/my_model.lite. - -``` - -#### Test that the model produces sensible outcomes - -You can go a step further to verify that the model produces results as expected. This helps you catch bugs during TOCO conversion (e.g. using the wrong mean and std values). - -* Move your submission to the testdata folder: - -```sh -cp /path/to/my_model.lite tensorflow/contrib/lite/java/ovic/src/testdata/ -``` - -* Resize the test image to the resolutions that are expected by your submission: - -The test images can be found at `tensorflow/contrib/lite/java/ovic/src/testdata/test_image_*.jpg`. You may reuse these images if your image resolutions are 128x128 or 224x224. - -* Add your model and test image to the BUILD rule at `tensorflow/contrib/lite/java/ovic/src/testdata/BUILD`: - -```JSON -filegroup( - name = "ovic_testdata", - srcs = [ - "@tflite_ovic_testdata//:float_model.lite", - "@tflite_ovic_testdata//:low_res_model.lite", - "@tflite_ovic_testdata//:quantized_model.lite", - "@tflite_ovic_testdata//:test_image_128.jpg", - "@tflite_ovic_testdata//:test_image_224.jpg" - "my_model.lite", # <--- Your submission. - "my_test_image.jpg", # <--- Your test image. - ], - ... -``` - -* Modify `OvicClassifierTest.java` to test your model. - -Change `TEST_IMAGE_PATH` to `my_test_image.jpg`. Change either `FLOAT_MODEL_PATH` or `QUANTIZED_MODEL_PATH` to `my_model.lite` depending on whether your model runs inference in float or [8-bit](https://www.tensorflow.org/performance/quantization). - -Now you can run the bazel tests to catch any runtime issues with the submission. - -Note: Please make sure that your submission passes the test. If a submission fails to pass the test it will not be processed by the submission server. - -## Measure on-device latency - -We provide two ways to measure the on-device latency of your submission. The first is through our competition server, which is reliable and repeatable, but is limited to a few trials per day. The second is through the benchmarker Apk, which requires a device and may not be as accurate as the server, but has a fast turn-around and no access limitations. We recommend that the participants use the benchmarker apk for early development, and reserve the competition server for evaluating promising submissions. - -### Running the benchmarker app - -Make sure that you have followed instructions in [Test your submissions](#test-your-submissions) to add your model to the testdata folder and to the corresponding build rules. - -Modify `tensorflow/contrib/lite/java/ovic/demo/app/OvicBenchmarkerActivity.java`: - -* Add your model to the benchmarker apk by changing `MODEL_PATH` and `TEST_IMAGE_PATH` below to your submission and test image. - -``` - private static final String TEST_IMAGE_PATH = "my_test_image.jpg"; - private static final String MODEL_PATH = "my_model.lite"; -``` - -* Adjust the benchmark parameters when needed: - -You can chnage the length of each experiment, and the processor affinity below. `BIG_CORE_MASK` is an integer whose binary encoding represents the set of used cores. This number is phone-specific. For example, Pixel 2 has 8 cores: the 4 little cores are represented by the 4 less significant bits, and the 4 big cores by the 4 more significant bits. Therefore a mask value of 16, or in binary `00010000`, represents using only the first big core. The mask 32, or in binary `00100000` uses the second big core and should deliver identical results as the mask 16 because the big cores are interchangeable. - -``` - /** Wall time for each benchmarking experiment. */ - private static final double WALL_TIME = 3000; - /** Maximum number of iterations in each benchmarking experiment. */ - private static final int MAX_ITERATIONS = 100; - /** Mask for binding to a single big core. Pixel 1 (4), Pixel 2 (16). */ - private static final int BIG_CORE_MASK = 16; -``` - -Note: You'll need ROOT access to the phone to change processor affinity. - -* Build and install the app. - -``` -bazel build -c opt --cxxopt=--std=c++11 --cxxopt=-Wno-all //tensorflow/contrib/lite/java/ovic/demo/app:ovic_benchmarker_binary -adb install -r bazel-bin/tensorflow/contrib/lite/java/ovic/demo/app/ovic_benchmarker_binary.apk -``` - -Start the app and click the `Start` button in dark green. The button should turn bright green, signaling that the experiment is running. The benchmarking results will be displayed after about the `WALL_TIME` you specified above. For example: - -``` -my_model.lite: Average latency=158.6ms after 20 runs. -``` - -### Sample latencies - -Note: the benchmarking results can be quite different depending on the background processes running on the phone. A few things that help stabilize the app's readings are placing the phone on a cooling plate, restarting the phone, and shutting down internet access. - -| Model | Pixel 1 latency (ms) | Pixel 2 latency (ms) | -| -------------------- |:---------------------:| --------------------:| -| float_model.lite | 120 | 155 | -| quantized_model.lite | 85 | 74 | -| low_res_model.lite | 4.2 | 4.0 | - -Since Pixel 2 has excellent support for 8-bit quantized models, we strongly recommend you to check out the [quantization training tutorial](https://www.tensorflow.org/performance/quantization). - diff --git a/tensorflow/contrib/lite/java/ovic/demo/app/BUILD b/tensorflow/contrib/lite/java/ovic/demo/app/BUILD deleted file mode 100644 index f567358ea33966ea8fdb422749662e22111c5fcc..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/java/ovic/demo/app/BUILD +++ /dev/null @@ -1,34 +0,0 @@ -load("@build_bazel_rules_android//android:rules.bzl", "android_binary") - -# Sample app for OVIC benchmarking. -licenses(["notice"]) # Apache 2.0 - -android_binary( - name = "ovic_benchmarker_binary", - srcs = [ - "OvicBenchmarkerActivity.java", - ], - aapt_version = "aapt", - assets = [ - "//tensorflow/contrib/lite/java/ovic/src/testdata:coco_labels.txt", - "//tensorflow/contrib/lite/java/ovic/src/testdata:labels.txt", - "//tensorflow/contrib/lite/java/ovic/src/testdata:ovic_testdata", - "@tflite_mobilenet_ssd_quant//:detect.tflite", - ], - assets_dir = "", - custom_package = "ovic.demo.app", - manifest = "AndroidManifest.xml", - nocompress_extensions = [ - ".lite", - ".tflite", - ], - resource_files = glob(["res/**"]), - tags = ["manual"], - deps = [ - "//tensorflow/contrib/lite/java:tensorflowlite", - "//tensorflow/contrib/lite/java/ovic:ovicbenchmarkerlib", - "//tensorflow/contrib/lite/java/ovic:ovicdetectionbenchmarkerlib", - "@androidsdk//com.android.support:support-v13-25.2.0", - "@androidsdk//com.android.support:support-v4-25.2.0", - ], -) diff --git a/tensorflow/contrib/lite/java/ovic/demo/app/build.gradle b/tensorflow/contrib/lite/java/ovic/demo/app/build.gradle deleted file mode 100644 index 4f3a6cdb2f8fe58008c9315bf08f4d328e720073..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/java/ovic/demo/app/build.gradle +++ /dev/null @@ -1,52 +0,0 @@ -apply plugin: 'com.android.application' - -android { - compileSdkVersion 26 - buildToolsVersion "26.0.1" - defaultConfig { - applicationId "android.example.com.ovicbenchmarker" - minSdkVersion 15 - targetSdkVersion 26 - versionCode 1 - versionName "1.0" - - // Remove this block. - jackOptions { - enabled true - } - } - lintOptions { - abortOnError false - } - buildTypes { - release { - minifyEnabled false - proguardFiles getDefaultProguardFile('proguard-android.txt'), 'proguard-rules.pro' - } - } - aaptOptions { - noCompress "lite", "tflite" - } - - compileOptions { - sourceCompatibility JavaVersion.VERSION_1_8 - targetCompatibility JavaVersion.VERSION_1_8 - } -} - -repositories { - maven { - url 'https://google.bintray.com/tensorflow' - } -} - -dependencies { - compile fileTree(dir: 'libs', include: ['*.jar']) - compile 'com.android.support:appcompat-v7:25.2.0' - compile 'com.android.support.constraint:constraint-layout:1.0.2' - compile 'com.android.support:design:25.2.0' - compile 'com.android.support:support-annotations:25.3.1' - compile 'com.android.support:support-v13:25.2.0' - - compile 'org.tensorflow:tensorflow-lite:+' -} diff --git a/tensorflow/contrib/lite/java/ovic/demo/build.gradle b/tensorflow/contrib/lite/java/ovic/demo/build.gradle deleted file mode 100644 index b78a0b86c939620b6f05483ce45c4d3ef0ef595e..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/java/ovic/demo/build.gradle +++ /dev/null @@ -1,23 +0,0 @@ -// Top-level build file where you can add configuration options common to all sub-projects/modules. - -buildscript { - repositories { - jcenter() - } - dependencies { - classpath 'com.android.tools.build:gradle:2.3.1' - - // NOTE: Do not place your application dependencies here; they belong - // in the individual module build.gradle files - } -} - -allprojects { - repositories { - jcenter() - } -} - -task clean(type: Delete) { - delete rootProject.buildDir -} diff --git a/tensorflow/contrib/lite/java/ovic/demo/gradle/wrapper/gradle-wrapper.properties b/tensorflow/contrib/lite/java/ovic/demo/gradle/wrapper/gradle-wrapper.properties deleted file mode 100644 index fa7a38a0e43eecd1e7292dd49efa79a5d0742e2a..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/java/ovic/demo/gradle/wrapper/gradle-wrapper.properties +++ /dev/null @@ -1,6 +0,0 @@ -#Thu Sep 28 09:01:41 PDT 2017 -distributionBase=GRADLE_USER_HOME -distributionPath=wrapper/dists -zipStoreBase=GRADLE_USER_HOME -zipStorePath=wrapper/dists -distributionUrl=https\://services.gradle.org/distributions/gradle-3.3-all.zip diff --git a/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicValidator.java b/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicValidator.java deleted file mode 100644 index baa14baf920fee0b0a2feecee7e65ef5a9e96f95..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/java/ovic/src/main/java/org/tensorflow/ovic/OvicValidator.java +++ /dev/null @@ -1,94 +0,0 @@ -/*Copyright 2018 Google LLC - -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 - - https://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. -==============================================================================*/ -package org.tensorflow.ovic; - -import java.io.File; -import java.io.FileInputStream; -import java.io.IOException; -import java.io.InputStream; -import java.io.PrintStream; -import java.nio.ByteBuffer; -import java.nio.ByteOrder; -import java.nio.MappedByteBuffer; -import java.nio.channels.FileChannel; -import java.util.Random; - -/** Validate a submission model. */ -public class OvicValidator { - private static void printUsage(PrintStream s) { - s.println("Java program that validates a submission model."); - s.println(); - s.println("Usage: ovic_validator "); - s.println(); - s.println("Where:"); - s.println(" is the model in TfLite format;"); - } - - public static void main(String[] args) { - if (args.length != 1) { - printUsage(System.err); - System.exit(1); - } - final String labelPath = - "tensorflow/contrib/lite/java/ovic/src/testdata/labels.txt"; - - final String modelFile = args[0]; - try { - File labelsfile = new File(labelPath); - InputStream labelsInputStream = new FileInputStream(labelsfile); - MappedByteBuffer model = loadModelFile(modelFile); - OvicClassifier classifier = new OvicClassifier(labelsInputStream, model); - ByteBuffer imgData = createByteBufferForClassifier(classifier); - OvicClassificationResult testResult = classifier.classifyByteBuffer(imgData); - if (testResult.topKClasses.isEmpty()) { - throw new RuntimeException("Failed to return top K predictions."); - } - System.out.printf("Successfully validated %s.%n", modelFile); - } catch (Exception e) { - System.out.println(e.getMessage()); - System.out.printf("Failed to validate %s.%n", modelFile); - } - } - - private static ByteBuffer createByteBufferForClassifier(OvicClassifier classifier) { - if (classifier == null) { - throw new RuntimeException("Cannot create image buffer with the classifier."); - } - int[] inputDims = classifier.getInputDims(); - int imgHeight = inputDims[1]; - int imgWidth = inputDims[2]; - ByteBuffer imgData = ByteBuffer.allocateDirect(imgHeight * imgWidth * 3); - imgData.order(ByteOrder.nativeOrder()); - Random rand = new Random(); - for (int y = 0; y < imgHeight; y++) { - for (int x = 0; x < imgWidth; x++) { - int val = rand.nextInt(); - imgData.put((byte) ((val >> 16) & 0xFF)); - imgData.put((byte) ((val >> 8) & 0xFF)); - imgData.put((byte) (val & 0xFF)); - } - } - return imgData; - } - - private static MappedByteBuffer loadModelFile(String modelFilePath) throws IOException { - File modelfile = new File(modelFilePath); - FileInputStream inputStream = new FileInputStream(modelfile); - FileChannel fileChannel = inputStream.getChannel(); - long startOffset = 0L; - long declaredLength = fileChannel.size(); - return fileChannel.map(FileChannel.MapMode.READ_ONLY, startOffset, declaredLength); - } -} diff --git a/tensorflow/contrib/lite/java/ovic/src/testdata/BUILD b/tensorflow/contrib/lite/java/ovic/src/testdata/BUILD deleted file mode 100644 index 051aa2204efd37fbcb12fb8ae67195780ffffad6..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/java/ovic/src/testdata/BUILD +++ /dev/null @@ -1,22 +0,0 @@ -# Testdata for OVIC benchmarker demo App and tests. -licenses(["notice"]) # Apache 2.0 - -filegroup( - name = "ovic_testdata", - srcs = [ - "@tflite_ovic_testdata//:float_model.lite", - "@tflite_ovic_testdata//:low_res_model.lite", - "@tflite_ovic_testdata//:quantized_model.lite", - "@tflite_ovic_testdata//:test_image_128.jpg", - "@tflite_ovic_testdata//:test_image_224.jpg", - ], - visibility = ["//visibility:public"], -) - -exports_files( - [ - "labels.txt", - "coco_labels.txt", - ], - visibility = ["//visibility:public"], -) diff --git a/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/Interpreter.java b/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/Interpreter.java deleted file mode 100644 index 5cc6e754f3380331a85714c2ed69f8c8e49ba4dc..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/Interpreter.java +++ /dev/null @@ -1,343 +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. -==============================================================================*/ - -package org.tensorflow.lite; - -import java.io.File; -import java.nio.ByteBuffer; -import java.nio.MappedByteBuffer; -import java.util.HashMap; -import java.util.Map; -import org.checkerframework.checker.nullness.qual.NonNull; - -/** - * Driver class to drive model inference with TensorFlow Lite. - * - *

A {@code Interpreter} encapsulates a pre-trained TensorFlow Lite model, in which operations - * are executed for model inference. - * - *

For example, if a model takes only one input and returns only one output: - * - * 

{@code
- * try (Interpreter interpreter = new Interpreter(file_of_a_tensorflowlite_model)) {
- *   interpreter.run(input, output);
- * }
- * }
- * - *

If a model takes multiple inputs or outputs: - * - * 

{@code
- * Object[] inputs = {input0, input1, ...};
- * Map map_of_indices_to_outputs = new HashMap<>();
- * float[][][] ith_output = new float[3][2][4];
- * map_of_indices_to_outputs.put(i, ith_output);
- * try (Interpreter interpreter = new Interpreter(file_of_a_tensorflowlite_model)) {
- *   interpreter.runForMultipleInputsOutputs(inputs, map_of_indices_to_outputs);
- * }
- * }
- * - *

Orders of inputs and outputs are determined when converting TensorFlow model to TensorFlowLite - * model with Toco. - * - *

WARNING:Instances of a {@code Interpreter} is not thread-safe. A {@code - * Interpreter} owns resources that must be explicitly freed by invoking {@link #close()} - */ -public final class Interpreter implements AutoCloseable { - - /** An options class for controlling runtime interpreter behavior. */ - public static class Options { - public Options() {} - - /** - * Sets the number of threads to be used for ops that support multi-threading. Defaults to a - * platform-dependent value. - */ - public Options setNumThreads(int numThreads) { - this.numThreads = numThreads; - return this; - } - - /** Sets whether to use NN API (if available) for op execution. Defaults to false (disabled). */ - public Options setUseNNAPI(boolean useNNAPI) { - this.useNNAPI = useNNAPI; - return this; - } - - /** - * Sets whether to allow float16 precision for FP32 calculation when possible. Defaults to false - * (disallow). - * WARNING: This is an experimental API and subject to change. - */ - public Options setAllowFp16PrecisionForFp32(boolean allow) { - this.allowFp16PrecisionForFp32 = allow; - return this; - } - - int numThreads = -1; - boolean useNNAPI = false; - boolean allowFp16PrecisionForFp32 = false; - } - - /** - * Initializes a {@code Interpreter} - * - * @param modelFile: a File of a pre-trained TF Lite model. - */ - public Interpreter(@NonNull File modelFile) { - this(modelFile, /*options = */ null); - } - - /** - * Initializes a {@code Interpreter} and specifies the number of threads used for inference. - * - * @param modelFile: a file of a pre-trained TF Lite model - * @param numThreads: number of threads to use for inference - * @deprecated Prefer using the {@link #Interpreter(File,Options)} constructor. This method will - * be removed in a future release. - */ - @Deprecated - public Interpreter(@NonNull File modelFile, int numThreads) { - this(modelFile, new Options().setNumThreads(numThreads)); - } - - /** - * Initializes a {@code Interpreter} and specifies the number of threads used for inference. - * - * @param modelFile: a file of a pre-trained TF Lite model - * @param options: a set of options for customizing interpreter behavior - */ - public Interpreter(@NonNull File modelFile, Options options) { - wrapper = new NativeInterpreterWrapper(modelFile.getAbsolutePath(), options); - } - - /** - * Initializes a {@code Interpreter} with a {@code ByteBuffer} of a model file. - * - *

The ByteBuffer should not be modified after the construction of a {@code Interpreter}. The - * {@code ByteBuffer} can be either a {@code MappedByteBuffer} that memory-maps a model file, or a - * direct {@code ByteBuffer} of nativeOrder() that contains the bytes content of a model. - */ - public Interpreter(@NonNull ByteBuffer byteBuffer) { - this(byteBuffer, /* options= */ null); - } - - /** - * Initializes a {@code Interpreter} with a {@code ByteBuffer} of a model file and specifies the - * number of threads used for inference. - * - *

The ByteBuffer should not be modified after the construction of a {@code Interpreter}. The - * {@code ByteBuffer} can be either a {@code MappedByteBuffer} that memory-maps a model file, or a - * direct {@code ByteBuffer} of nativeOrder() that contains the bytes content of a model. - * - * @deprecated Prefer using the {@link #Interpreter(ByteBuffer,Options)} constructor. This method - * will be removed in a future release. - */ - @Deprecated - public Interpreter(@NonNull ByteBuffer byteBuffer, int numThreads) { - this(byteBuffer, new Options().setNumThreads(numThreads)); - } - - /** - * Initializes a {@code Interpreter} with a {@code MappedByteBuffer} to the model file. - * - *

The {@code MappedByteBuffer} should remain unchanged after the construction of a {@code - * Interpreter}. - * - * @deprecated Prefer using the {@link #Interpreter(ByteBuffer,Options)} constructor. This method - * will be removed in a future release. - */ - @Deprecated - public Interpreter(@NonNull MappedByteBuffer mappedByteBuffer) { - this(mappedByteBuffer, /* options= */ null); - } - - /** - * Initializes a {@code Interpreter} with a {@code ByteBuffer} of a model file and a set of custom - * {@link #Options}. - * - *

The ByteBuffer should not be modified after the construction of a {@code Interpreter}. The - * {@code ByteBuffer} can be either a {@code MappedByteBuffer} that memory-maps a model file, or a - * direct {@code ByteBuffer} of nativeOrder() that contains the bytes content of a model. - */ - public Interpreter(@NonNull ByteBuffer byteBuffer, Options options) { - wrapper = new NativeInterpreterWrapper(byteBuffer, options); - } - - /** - * Runs model inference if the model takes only one input, and provides only one output. - * - *

Warning: The API runs much faster if {@link ByteBuffer} is used as input data type. Please - * consider using {@link ByteBuffer} to feed input data for better performance. - * - * @param input an array or multidimensional array, or a {@link ByteBuffer} of primitive types - * including int, float, long, and byte. {@link ByteBuffer} is the preferred way to pass large - * input data. When {@link ByteBuffer} is used, its content should remain unchanged until - * model inference is done. - * @param output a multidimensional array of output data, or a {@link ByteBuffer} of primitive - * types including int, float, long, and byte. - */ - public void run(@NonNull Object input, @NonNull Object output) { - Object[] inputs = {input}; - Map outputs = new HashMap<>(); - outputs.put(0, output); - runForMultipleInputsOutputs(inputs, outputs); - } - - /** - * Runs model inference if the model takes multiple inputs, or returns multiple outputs. - * - *

Warning: The API runs much faster if {@link ByteBuffer} is used as input data type. Please - * consider using {@link ByteBuffer} to feed input data for better performance. - * - * @param inputs an array of input data. The inputs should be in the same order as inputs of the - * model. Each input can be an array or multidimensional array, or a {@link ByteBuffer} of - * primitive types including int, float, long, and byte. {@link ByteBuffer} is the preferred - * way to pass large input data. When {@link ByteBuffer} is used, its content should remain - * unchanged until model inference is done. - * @param outputs a map mapping output indices to multidimensional arrays of output data or {@link - * ByteBuffer}s of primitive types including int, float, long, and byte. It only needs to keep - * entries for the outputs to be used. - */ - public void runForMultipleInputsOutputs( - @NonNull Object[] inputs, @NonNull Map outputs) { - checkNotClosed(); - wrapper.run(inputs, outputs); - } - - /** - * Resizes idx-th input of the native model to the given dims. - * - *

IllegalArgumentException will be thrown if it fails to resize. - */ - public void resizeInput(int idx, @NonNull int[] dims) { - checkNotClosed(); - wrapper.resizeInput(idx, dims); - } - - /** Gets the number of input tensors. */ - public int getInputTensorCount() { - checkNotClosed(); - return wrapper.getInputTensorCount(); - } - - /** - * Gets index of an input given the op name of the input. - * - *

IllegalArgumentException will be thrown if the op name does not exist in the model file used - * to initialize the {@link Interpreter}. - */ - public int getInputIndex(String opName) { - checkNotClosed(); - return wrapper.getInputIndex(opName); - } - - /** - * Gets the Tensor associated with the provdied input index. - * - *

IllegalArgumentException will be thrown if the provided index is invalid. - */ - public Tensor getInputTensor(int inputIndex) { - checkNotClosed(); - return wrapper.getInputTensor(inputIndex); - } - - /** Gets the number of output Tensors. */ - public int getOutputTensorCount() { - checkNotClosed(); - return wrapper.getOutputTensorCount(); - } - - /** - * Gets index of an output given the op name of the output. - * - *

IllegalArgumentException will be thrown if the op name does not exist in the model file used - * to initialize the {@link Interpreter}. - */ - public int getOutputIndex(String opName) { - checkNotClosed(); - return wrapper.getOutputIndex(opName); - } - - /** - * Gets the Tensor associated with the provdied output index. - * - *

IllegalArgumentException will be thrown if the provided index is invalid. - */ - public Tensor getOutputTensor(int outputIndex) { - checkNotClosed(); - return wrapper.getOutputTensor(outputIndex); - } - - /** - * Returns native inference timing. - * - *

IllegalArgumentException will be thrown if the model is not initialized by the {@link - * Interpreter}. - */ - public Long getLastNativeInferenceDurationNanoseconds() { - checkNotClosed(); - return wrapper.getLastNativeInferenceDurationNanoseconds(); - } - - /** - * Turns on/off Android NNAPI for hardware acceleration when it is available. - * - * @deprecated Prefer using {@link Options#setUseNNAPI(boolean)} directly for enabling NN API. - * This method will be removed in a future release. - */ - @Deprecated - public void setUseNNAPI(boolean useNNAPI) { - checkNotClosed(); - wrapper.setUseNNAPI(useNNAPI); - } - - /** - * Sets the number of threads to be used for ops that support multi-threading. - * - * @deprecated Prefer using {@link Options#setNumThreads(int)} directly for controlling thread - * multi-threading. This method will be removed in a future release. - */ - @Deprecated - public void setNumThreads(int numThreads) { - checkNotClosed(); - wrapper.setNumThreads(numThreads); - } - - /** Release resources associated with the {@code Interpreter}. */ - @Override - public void close() { - if (wrapper != null) { - wrapper.close(); - wrapper = null; - } - } - - @Override - protected void finalize() throws Throwable { - try { - close(); - } finally { - super.finalize(); - } - } - - private void checkNotClosed() { - if (wrapper == null) { - throw new IllegalStateException("Internal error: The Interpreter has already been closed."); - } - } - - NativeInterpreterWrapper wrapper; -} diff --git a/tensorflow/contrib/lite/java/src/main/native/BUILD b/tensorflow/contrib/lite/java/src/main/native/BUILD deleted file mode 100644 index f91345f369fe118839cd6e28032b36c346008f58..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/java/src/main/native/BUILD +++ /dev/null @@ -1,127 +0,0 @@ -# Description: -# Java Native Interface (JNI) library intended for implementing the -# TensorFlow Lite Java API using the TensorFlow Lite CC library. - -package(default_visibility = ["//visibility:public"]) - -load("//tensorflow/contrib/lite:build_def.bzl", "tflite_copts") - -licenses(["notice"]) # Apache 2.0 - -cc_library( - name = "native_framework_only", - srcs = [ - "exception_jni.cc", - "nativeinterpreterwrapper_jni.cc", - "tensor_jni.cc", - "tensorflow_lite_jni.cc", - ] + select({ - # The Android toolchain makes "jni.h" available in the include path. - # For non-Android toolchains, generate jni.h and jni_md.h. - "//tensorflow:android": [], - "//conditions:default": [ - ":jni.h", - ":jni_md.h", - ], - }), - hdrs = [ - "exception_jni.h", - "nativeinterpreterwrapper_jni.h", - "tensor_jni.h", - "tensorflow_lite_jni.h", - ], - copts = tflite_copts(), - includes = select({ - "//tensorflow:android": [], - "//conditions:default": ["."], - }), - linkopts = [ - "-lm", - "-ldl", - ], - deps = [ - "//tensorflow/contrib/lite:context", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite:schema_fbs_version", - ], - alwayslink = 1, -) - -# Silly rules to make -# #include -# in the source headers work -# (in combination with the "includes" attribute of the tf_cuda_library rule -# above. Not needed when using the Android toolchain). -# -# Inspired from: -# https://github.com/bazelbuild/bazel/blob/f99a0543f8d97339d32075c7176b79f35be84606/src/main/native/BUILD -# but hopefully there is a simpler alternative to this. -genrule( - name = "copy_jni_h", - srcs = ["@bazel_tools//tools/jdk:jni_header"], - outs = ["jni.h"], - cmd = "cp -f $< $@", -) - -genrule( - name = "copy_jni_md_h", - srcs = select({ - "//tensorflow:darwin": ["@bazel_tools//tools/jdk:jni_md_header-darwin"], - "//conditions:default": ["@bazel_tools//tools/jdk:jni_md_header-linux"], - }), - outs = ["jni_md.h"], - cmd = "cp -f $< $@", -) - -cc_library( - name = "init_tensorflow", - srcs = [ - "init_tensorflow_jni.cc", - ] + select({ - # The Android toolchain makes "jni.h" available in the include path. - # For non-Android toolchains, generate jni.h and jni_md.h. - "//tensorflow:android": [], - "//conditions:default": [ - ":jni.h", - ":jni_md.h", - ], - }), - hdrs = [ - "init_tensorflow_jni.h", - ], - copts = tflite_copts(), - includes = select({ - "//tensorflow:android": [], - "//conditions:default": ["."], - }), - linkopts = [ - "-lm", - "-ldl", - ], - deps = [ - "//tensorflow/contrib/lite/testing:init_tensorflow", - ], - alwayslink = 1, -) - -# This includes all ops. If you want a smaller binary, you should copy and -# modify builtin_ops_jni.cc. You should then link your binary against both -# ":native_framework_only" and your own version of ":native_builtin_ops". -cc_library( - name = "native", - srcs = [ - "builtin_ops_jni.cc", - ], - copts = tflite_copts(), - deps = [ - ":native_framework_only", - "//tensorflow/contrib/lite/kernels:builtin_ops", - ], - alwayslink = 1, -) - -exports_files( - [ - "version_script.lds", - ], -) diff --git a/tensorflow/contrib/lite/java/src/main/native/tensor_jni.cc b/tensorflow/contrib/lite/java/src/main/native/tensor_jni.cc deleted file mode 100644 index d3378f5f145deef375b38777fa27046993e15a6c..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/java/src/main/native/tensor_jni.cc +++ /dev/null @@ -1,318 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/contrib/lite/java/src/main/native/tensor_jni.h" -#include -#include -#include "tensorflow/contrib/lite/interpreter.h" -#include "tensorflow/contrib/lite/java/src/main/native/exception_jni.h" - -namespace { - -// Convenience handle for obtaining a TfLiteTensor given an interpreter and -// tensor index. -// -// Historically, the Java Tensor class used a TfLiteTensor pointer as its native -// handle. However, this approach isn't generally safe, as the interpreter may -// invalidate all TfLiteTensor* handles during inference or allocation. -class TensorHandle { - public: - TensorHandle(tflite::Interpreter* interpreter, int tensor_index) - : interpreter_(interpreter), tensor_index_(tensor_index) {} - - TfLiteTensor* tensor() const { return interpreter_->tensor(tensor_index_); } - - private: - tflite::Interpreter* const interpreter_; - const int tensor_index_; -}; - -TfLiteTensor* GetTensorFromHandle(JNIEnv* env, jlong handle) { - if (handle == 0) { - throwException(env, kIllegalArgumentException, - "Internal error: Invalid handle to TfLiteTensor."); - return nullptr; - } - return reinterpret_cast(handle)->tensor(); -} - -size_t elementByteSize(TfLiteType data_type) { - // The code in this file makes the assumption that the - // TensorFlow TF_DataTypes and the Java primitive types - // have the same byte sizes. Validate that: - switch (data_type) { - case kTfLiteFloat32: - static_assert(sizeof(jfloat) == 4, - "Interal error: Java float not compatible with " - "kTfLiteFloat"); - return 4; - case kTfLiteInt32: - static_assert(sizeof(jint) == 4, - "Interal error: Java int not compatible with kTfLiteInt"); - return 4; - case kTfLiteUInt8: - static_assert(sizeof(jbyte) == 1, - "Interal error: Java byte not compatible with " - "kTfLiteUInt8"); - return 1; - case kTfLiteInt64: - static_assert(sizeof(jlong) == 8, - "Interal error: Java long not compatible with " - "kTfLiteInt64"); - return 8; - default: - return 0; - } -} - -size_t writeOneDimensionalArray(JNIEnv* env, jobject object, TfLiteType type, - void* dst, size_t dst_size) { - jarray array = static_cast(object); - const int num_elements = env->GetArrayLength(array); - size_t to_copy = num_elements * elementByteSize(type); - if (to_copy > dst_size) { - throwException(env, kIllegalStateException, - "Internal error: cannot write Java array of %d bytes to " - "Tensor of %d bytes", - to_copy, dst_size); - return 0; - } - switch (type) { - case kTfLiteFloat32: { - jfloatArray float_array = static_cast(array); - jfloat* float_dst = static_cast(dst); - env->GetFloatArrayRegion(float_array, 0, num_elements, float_dst); - return to_copy; - } - case kTfLiteInt32: { - jintArray int_array = static_cast(array); - jint* int_dst = static_cast(dst); - env->GetIntArrayRegion(int_array, 0, num_elements, int_dst); - return to_copy; - } - case kTfLiteInt64: { - jlongArray long_array = static_cast(array); - jlong* long_dst = static_cast(dst); - env->GetLongArrayRegion(long_array, 0, num_elements, long_dst); - return to_copy; - } - case kTfLiteUInt8: { - jbyteArray byte_array = static_cast(array); - jbyte* byte_dst = static_cast(dst); - env->GetByteArrayRegion(byte_array, 0, num_elements, byte_dst); - return to_copy; - } - default: { - throwException(env, kUnsupportedOperationException, - "DataType error: TensorFlowLite currently supports float " - "(32 bits), int (32 bits), byte (8 bits), and long " - "(64 bits), support for other types (DataType %d in this " - "case) will be added in the future", - kTfLiteFloat32, type); - return 0; - } - } -} - -size_t readOneDimensionalArray(JNIEnv* env, TfLiteType data_type, - const void* src, size_t src_size, jarray dst) { - const int len = env->GetArrayLength(dst); - const size_t size = len * elementByteSize(data_type); - if (size > src_size) { - throwException( - env, kIllegalStateException, - "Internal error: cannot fill a Java array of %d bytes with a Tensor of " - "%d bytes", - size, src_size); - return 0; - } - switch (data_type) { - case kTfLiteFloat32: { - jfloatArray float_array = static_cast(dst); - env->SetFloatArrayRegion(float_array, 0, len, - static_cast(src)); - return size; - } - case kTfLiteInt32: { - jintArray int_array = static_cast(dst); - env->SetIntArrayRegion(int_array, 0, len, static_cast(src)); - return size; - } - case kTfLiteInt64: { - jlongArray long_array = static_cast(dst); - env->SetLongArrayRegion(long_array, 0, len, - static_cast(src)); - return size; - } - case kTfLiteUInt8: { - jbyteArray byte_array = static_cast(dst); - env->SetByteArrayRegion(byte_array, 0, len, - static_cast(src)); - return size; - } - default: { - throwException(env, kIllegalStateException, - "DataType error: invalid DataType(%d)", data_type); - } - } - return 0; -} - -size_t readMultiDimensionalArray(JNIEnv* env, TfLiteType data_type, char* src, - size_t src_size, int dims_left, jarray dst) { - if (dims_left == 1) { - return readOneDimensionalArray(env, data_type, src, src_size, dst); - } else { - jobjectArray ndarray = static_cast(dst); - int len = env->GetArrayLength(ndarray); - size_t size = 0; - for (int i = 0; i < len; ++i) { - jarray row = static_cast(env->GetObjectArrayElement(ndarray, i)); - size += readMultiDimensionalArray(env, data_type, src + size, - src_size - size, dims_left - 1, row); - env->DeleteLocalRef(row); - if (env->ExceptionCheck()) return size; - } - return size; - } -} - -size_t writeMultiDimensionalArray(JNIEnv* env, jobject src, TfLiteType type, - int dims_left, char** dst, int dst_size) { - if (dims_left <= 1) { - return writeOneDimensionalArray(env, src, type, *dst, dst_size); - } else { - jobjectArray ndarray = static_cast(src); - int len = env->GetArrayLength(ndarray); - size_t sz = 0; - for (int i = 0; i < len; ++i) { - jobject row = env->GetObjectArrayElement(ndarray, i); - char* next_dst = *dst + sz; - sz += writeMultiDimensionalArray(env, row, type, dims_left - 1, &next_dst, - dst_size - sz); - env->DeleteLocalRef(row); - if (env->ExceptionCheck()) return sz; - } - return sz; - } -} - -} // namespace - -JNIEXPORT jlong JNICALL Java_org_tensorflow_lite_Tensor_create( - JNIEnv* env, jclass clazz, jlong interpreter_handle, jint tensor_index) { - tflite::Interpreter* interpreter = - reinterpret_cast(interpreter_handle); - return reinterpret_cast(new TensorHandle(interpreter, tensor_index)); -} - -JNIEXPORT void JNICALL Java_org_tensorflow_lite_Tensor_delete(JNIEnv* env, - jclass clazz, - jlong handle) { - delete reinterpret_cast(handle); -} - -JNIEXPORT jobject JNICALL Java_org_tensorflow_lite_Tensor_buffer(JNIEnv* env, - jclass clazz, - jlong handle) { - TfLiteTensor* tensor = GetTensorFromHandle(env, handle); - if (tensor == nullptr) return nullptr; - if (tensor->data.raw == nullptr) { - throwException(env, kIllegalArgumentException, - "Internal error: Tensor hasn't been allocated."); - return nullptr; - } - return env->NewDirectByteBuffer(static_cast(tensor->data.raw), - static_cast(tensor->bytes)); -} - -JNIEXPORT void JNICALL Java_org_tensorflow_lite_Tensor_writeDirectBuffer( - JNIEnv* env, jclass clazz, jlong handle, jobject src) { - TfLiteTensor* tensor = GetTensorFromHandle(env, handle); - if (tensor == nullptr) return; - - char* src_data_raw = static_cast(env->GetDirectBufferAddress(src)); - if (!src_data_raw) { - throwException(env, kIllegalArgumentException, - "Input ByteBuffer is not a direct buffer"); - return; - } - - tensor->data.raw = src_data_raw; -} - -JNIEXPORT void JNICALL -Java_org_tensorflow_lite_Tensor_readMultiDimensionalArray(JNIEnv* env, - jclass clazz, - jlong handle, - jobject value) { - TfLiteTensor* tensor = GetTensorFromHandle(env, handle); - if (tensor == nullptr) return; - int num_dims = tensor->dims->size; - if (num_dims == 0) { - throwException(env, kIllegalArgumentException, - "Internal error: Cannot copy empty/scalar Tensors."); - return; - } - readMultiDimensionalArray(env, tensor->type, tensor->data.raw, tensor->bytes, - num_dims, static_cast(value)); -} - -JNIEXPORT void JNICALL -Java_org_tensorflow_lite_Tensor_writeMultiDimensionalArray(JNIEnv* env, - jclass clazz, - jlong handle, - jobject src) { - TfLiteTensor* tensor = GetTensorFromHandle(env, handle); - if (tensor == nullptr) return; - if (tensor->data.raw == nullptr) { - throwException(env, kIllegalArgumentException, - "Internal error: Target Tensor hasn't been allocated."); - return; - } - if (tensor->dims->size == 0) { - throwException(env, kIllegalArgumentException, - "Internal error: Cannot copy empty/scalar Tensors."); - return; - } - writeMultiDimensionalArray(env, src, tensor->type, tensor->dims->size, - &tensor->data.raw, tensor->bytes); -} - -JNIEXPORT jint JNICALL Java_org_tensorflow_lite_Tensor_dtype(JNIEnv* env, - jclass clazz, - jlong handle) { - TfLiteTensor* tensor = GetTensorFromHandle(env, handle); - if (tensor == nullptr) return 0; - return static_cast(tensor->type); -} - -JNIEXPORT jintArray JNICALL -Java_org_tensorflow_lite_Tensor_shape(JNIEnv* env, jclass clazz, jlong handle) { - TfLiteTensor* tensor = GetTensorFromHandle(env, handle); - if (tensor == nullptr) return nullptr; - int num_dims = tensor->dims->size; - jintArray result = env->NewIntArray(num_dims); - env->SetIntArrayRegion(result, 0, num_dims, tensor->dims->data); - return result; -} - -JNIEXPORT jint JNICALL Java_org_tensorflow_lite_Tensor_numBytes(JNIEnv* env, - jclass clazz, - jlong handle) { - const TfLiteTensor* tensor = GetTensorFromHandle(env, handle); - if (tensor == nullptr) return 0; - return static_cast(tensor->bytes); -} diff --git a/tensorflow/contrib/lite/java/src/test/java/org/tensorflow/lite/InterpreterFlexTest.java b/tensorflow/contrib/lite/java/src/test/java/org/tensorflow/lite/InterpreterFlexTest.java deleted file mode 100644 index 3b3d9f0e7fc0706c35045b85b316bcb16296cd90..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/java/src/test/java/org/tensorflow/lite/InterpreterFlexTest.java +++ /dev/null @@ -1,50 +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. -==============================================================================*/ - -package org.tensorflow.lite; - -import static com.google.common.truth.Truth.assertThat; - -import java.io.File; -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; - -/** - * Unit tests for {@link org.tensorflow.lite.Interpreter} that validate execution with models that - * have TensorFlow ops. - */ -@RunWith(JUnit4.class) -public final class InterpreterFlexTest { - - private static final File FLEX_MODEL_FILE = - new File("tensorflow/contrib/lite/testdata/multi_add_flex.bin"); - - /** Smoke test validating that flex model loading works when the flex delegate is linked. */ - @Test - public void testFlexModel() throws Exception { - try (Interpreter interpreter = new Interpreter(FLEX_MODEL_FILE)) { - assertThat(interpreter.getInputTensorCount()).isEqualTo(4); - assertThat(interpreter.getInputTensor(0).dataType()).isEqualTo(DataType.FLOAT32); - assertThat(interpreter.getOutputTensorCount()).isEqualTo(4); - assertThat(interpreter.getOutputTensor(0).dataType()).isEqualTo(DataType.FLOAT32); - interpreter.run(new float[1], new float[1]); - } - } - - static { - TensorFlowLite.initTensorFlow(); - } -} diff --git a/tensorflow/contrib/lite/java/src/test/java/org/tensorflow/lite/InterpreterTest.java b/tensorflow/contrib/lite/java/src/test/java/org/tensorflow/lite/InterpreterTest.java deleted file mode 100644 index f8b73c7cf3bc2dcd814dd19924a4d5597a7249b7..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/java/src/test/java/org/tensorflow/lite/InterpreterTest.java +++ /dev/null @@ -1,362 +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. -==============================================================================*/ - -package org.tensorflow.lite; - -import static com.google.common.truth.Truth.assertThat; -import static org.junit.Assert.fail; - -import java.io.File; -import java.nio.ByteBuffer; -import java.nio.ByteOrder; -import java.nio.MappedByteBuffer; -import java.nio.channels.FileChannel; -import java.nio.file.Files; -import java.nio.file.Path; -import java.nio.file.StandardOpenOption; -import java.util.EnumSet; -import java.util.HashMap; -import java.util.Map; -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; - -/** Unit tests for {@link org.tensorflow.lite.Interpreter}. */ -@RunWith(JUnit4.class) -public final class InterpreterTest { - - private static final File MODEL_FILE = - new File("tensorflow/contrib/lite/java/src/testdata/add.bin"); - - private static final File MOBILENET_MODEL_FILE = - new File("tensorflow/contrib/lite/java/src/testdata/mobilenet.tflite.bin"); - - private static final File FLEX_MODEL_FILE = - new File("tensorflow/contrib/lite/testdata/multi_add_flex.bin"); - - @Test - public void testInterpreter() throws Exception { - Interpreter interpreter = new Interpreter(MODEL_FILE); - assertThat(interpreter).isNotNull(); - assertThat(interpreter.getInputTensorCount()).isEqualTo(1); - assertThat(interpreter.getInputTensor(0).dataType()).isEqualTo(DataType.FLOAT32); - assertThat(interpreter.getOutputTensorCount()).isEqualTo(1); - assertThat(interpreter.getOutputTensor(0).dataType()).isEqualTo(DataType.FLOAT32); - interpreter.close(); - } - - @Test - public void testInterpreterWithOptions() throws Exception { - Interpreter interpreter = - new Interpreter(MODEL_FILE, new Interpreter.Options().setNumThreads(2).setUseNNAPI(true)); - assertThat(interpreter).isNotNull(); - assertThat(interpreter.getInputTensorCount()).isEqualTo(1); - assertThat(interpreter.getInputTensor(0).dataType()).isEqualTo(DataType.FLOAT32); - assertThat(interpreter.getOutputTensorCount()).isEqualTo(1); - assertThat(interpreter.getOutputTensor(0).dataType()).isEqualTo(DataType.FLOAT32); - interpreter.close(); - } - - @Test - public void testRunWithMappedByteBufferModel() throws Exception { - Path path = MODEL_FILE.toPath(); - FileChannel fileChannel = - (FileChannel) Files.newByteChannel(path, EnumSet.of(StandardOpenOption.READ)); - ByteBuffer mappedByteBuffer = - fileChannel.map(FileChannel.MapMode.READ_ONLY, 0, fileChannel.size()); - Interpreter interpreter = new Interpreter(mappedByteBuffer); - float[] oneD = {1.23f, 6.54f, 7.81f}; - float[][] twoD = {oneD, oneD, oneD, oneD, oneD, oneD, oneD, oneD}; - float[][][] threeD = {twoD, twoD, twoD, twoD, twoD, twoD, twoD, twoD}; - float[][][][] fourD = {threeD, threeD}; - float[][][][] parsedOutputs = new float[2][8][8][3]; - interpreter.run(fourD, parsedOutputs); - float[] outputOneD = parsedOutputs[0][0][0]; - float[] expected = {3.69f, 19.62f, 23.43f}; - assertThat(outputOneD).usingTolerance(0.1f).containsExactly(expected).inOrder(); - interpreter.close(); - fileChannel.close(); - } - - @Test - public void testRunWithDirectByteBufferModel() throws Exception { - Path path = MODEL_FILE.toPath(); - FileChannel fileChannel = - (FileChannel) Files.newByteChannel(path, EnumSet.of(StandardOpenOption.READ)); - ByteBuffer byteBuffer = ByteBuffer.allocateDirect((int) fileChannel.size()); - byteBuffer.order(ByteOrder.nativeOrder()); - fileChannel.read(byteBuffer); - Interpreter interpreter = new Interpreter(byteBuffer); - float[] oneD = {1.23f, 6.54f, 7.81f}; - float[][] twoD = {oneD, oneD, oneD, oneD, oneD, oneD, oneD, oneD}; - float[][][] threeD = {twoD, twoD, twoD, twoD, twoD, twoD, twoD, twoD}; - float[][][][] fourD = {threeD, threeD}; - float[][][][] parsedOutputs = new float[2][8][8][3]; - interpreter.run(fourD, parsedOutputs); - float[] outputOneD = parsedOutputs[0][0][0]; - float[] expected = {3.69f, 19.62f, 23.43f}; - assertThat(outputOneD).usingTolerance(0.1f).containsExactly(expected).inOrder(); - interpreter.close(); - fileChannel.close(); - } - - @Test - public void testRunWithInvalidByteBufferModel() throws Exception { - Path path = MODEL_FILE.toPath(); - FileChannel fileChannel = - (FileChannel) Files.newByteChannel(path, EnumSet.of(StandardOpenOption.READ)); - ByteBuffer byteBuffer = ByteBuffer.allocate((int) fileChannel.size()); - byteBuffer.order(ByteOrder.nativeOrder()); - fileChannel.read(byteBuffer); - try { - new Interpreter(byteBuffer); - fail(); - } catch (IllegalArgumentException e) { - assertThat(e) - .hasMessageThat() - .contains( - "Model ByteBuffer should be either a MappedByteBuffer" - + " of the model file, or a direct ByteBuffer using ByteOrder.nativeOrder()"); - } - fileChannel.close(); - } - - @Test - public void testRun() { - Interpreter interpreter = new Interpreter(MODEL_FILE); - Float[] oneD = {1.23f, 6.54f, 7.81f}; - Float[][] twoD = {oneD, oneD, oneD, oneD, oneD, oneD, oneD, oneD}; - Float[][][] threeD = {twoD, twoD, twoD, twoD, twoD, twoD, twoD, twoD}; - Float[][][][] fourD = {threeD, threeD}; - Float[][][][] parsedOutputs = new Float[2][8][8][3]; - try { - interpreter.run(fourD, parsedOutputs); - fail(); - } catch (IllegalArgumentException e) { - assertThat(e).hasMessageThat().contains("cannot resolve DataType of [[[[Ljava.lang.Float;"); - } - interpreter.close(); - } - - @Test - public void testRunWithBoxedInputs() { - Interpreter interpreter = new Interpreter(MODEL_FILE); - float[] oneD = {1.23f, 6.54f, 7.81f}; - float[][] twoD = {oneD, oneD, oneD, oneD, oneD, oneD, oneD, oneD}; - float[][][] threeD = {twoD, twoD, twoD, twoD, twoD, twoD, twoD, twoD}; - float[][][][] fourD = {threeD, threeD}; - float[][][][] parsedOutputs = new float[2][8][8][3]; - interpreter.run(fourD, parsedOutputs); - float[] outputOneD = parsedOutputs[0][0][0]; - float[] expected = {3.69f, 19.62f, 23.43f}; - assertThat(outputOneD).usingTolerance(0.1f).containsExactly(expected).inOrder(); - interpreter.close(); - } - - @Test - public void testRunForMultipleInputsOutputs() { - Interpreter interpreter = new Interpreter(MODEL_FILE); - float[] oneD = {1.23f, 6.54f, 7.81f}; - float[][] twoD = {oneD, oneD, oneD, oneD, oneD, oneD, oneD, oneD}; - float[][][] threeD = {twoD, twoD, twoD, twoD, twoD, twoD, twoD, twoD}; - float[][][][] fourD = {threeD, threeD}; - Object[] inputs = {fourD}; - float[][][][] parsedOutputs = new float[2][8][8][3]; - Map outputs = new HashMap<>(); - outputs.put(0, parsedOutputs); - interpreter.runForMultipleInputsOutputs(inputs, outputs); - float[] outputOneD = parsedOutputs[0][0][0]; - float[] expected = {3.69f, 19.62f, 23.43f}; - assertThat(outputOneD).usingTolerance(0.1f).containsExactly(expected).inOrder(); - interpreter.close(); - } - - @Test - public void testRunWithByteBufferOutput() { - float[] oneD = {1.23f, 6.54f, 7.81f}; - float[][] twoD = {oneD, oneD, oneD, oneD, oneD, oneD, oneD, oneD}; - float[][][] threeD = {twoD, twoD, twoD, twoD, twoD, twoD, twoD, twoD}; - float[][][][] fourD = {threeD, threeD}; - ByteBuffer parsedOutput = - ByteBuffer.allocateDirect(2 * 8 * 8 * 3 * 4).order(ByteOrder.nativeOrder()); - try (Interpreter interpreter = new Interpreter(MODEL_FILE)) { - interpreter.run(fourD, parsedOutput); - } - float[] outputOneD = { - parsedOutput.getFloat(0), parsedOutput.getFloat(4), parsedOutput.getFloat(8) - }; - float[] expected = {3.69f, 19.62f, 23.43f}; - assertThat(outputOneD).usingTolerance(0.1f).containsExactly(expected).inOrder(); - } - - @Test - public void testResizeInput() { - try (Interpreter interpreter = new Interpreter(MODEL_FILE)) { - int[] inputDims = {1}; - interpreter.resizeInput(0, inputDims); - assertThat(interpreter.getInputTensor(0).shape()).isEqualTo(inputDims); - ByteBuffer input = ByteBuffer.allocateDirect(4).order(ByteOrder.nativeOrder()); - ByteBuffer output = ByteBuffer.allocateDirect(4).order(ByteOrder.nativeOrder()); - interpreter.run(input, output); - assertThat(interpreter.getOutputTensor(0).shape()).isEqualTo(inputDims); - } - } - - @Test - public void testMobilenetRun() { - // Create a gray image. - float[][][][] img = new float[1][224][224][3]; - for (int i = 0; i < 224; ++i) { - for (int j = 0; j < 224; ++j) { - img[0][i][j][0] = 0.5f; - img[0][i][j][1] = 0.5f; - img[0][i][j][2] = 0.5f; - } - } - - // Allocate memory to receive the output values. - float[][] labels = new float[1][1001]; - - Interpreter interpreter = new Interpreter(MOBILENET_MODEL_FILE); - interpreter.run(img, labels); - assertThat(interpreter.getInputTensor(0).shape()).isEqualTo(new int[] {1, 224, 224, 3}); - assertThat(interpreter.getOutputTensor(0).shape()).isEqualTo(new int[] {1, 1001}); - interpreter.close(); - - assertThat(labels[0]) - .usingExactEquality() - .containsNoneOf(new float[] {Float.NaN, Float.NEGATIVE_INFINITY, Float.POSITIVE_INFINITY}); - } - - @Test - public void testRunWithWrongInputType() { - Interpreter interpreter = new Interpreter(MODEL_FILE); - int[] oneD = {4, 3, 9}; - int[][] twoD = {oneD, oneD, oneD, oneD, oneD, oneD, oneD, oneD}; - int[][][] threeD = {twoD, twoD, twoD, twoD, twoD, twoD, twoD, twoD}; - int[][][][] fourD = {threeD, threeD}; - float[][][][] parsedOutputs = new float[2][8][8][3]; - try { - interpreter.run(fourD, parsedOutputs); - fail(); - } catch (IllegalArgumentException e) { - assertThat(e) - .hasMessageThat() - .contains( - "Cannot convert between a TensorFlowLite tensor with type " - + "FLOAT32 and a Java object of type [[[[I (which is compatible with the" - + " TensorFlowLite type INT32)"); - } - interpreter.close(); - } - - @Test - public void testRunWithWrongOutputType() { - Interpreter interpreter = new Interpreter(MODEL_FILE); - float[] oneD = {1.23f, 6.54f, 7.81f}; - float[][] twoD = {oneD, oneD, oneD, oneD, oneD, oneD, oneD, oneD}; - float[][][] threeD = {twoD, twoD, twoD, twoD, twoD, twoD, twoD, twoD}; - float[][][][] fourD = {threeD, threeD}; - int[][][][] parsedOutputs = new int[2][8][8][3]; - try { - interpreter.run(fourD, parsedOutputs); - fail(); - } catch (IllegalArgumentException e) { - assertThat(e) - .hasMessageThat() - .contains( - "Cannot convert between a TensorFlowLite tensor with type " - + "FLOAT32 and a Java object of type [[[[I (which is compatible with the" - + " TensorFlowLite type INT32)"); - } - interpreter.close(); - } - - @Test - public void testGetInputIndex() { - Interpreter interpreter = new Interpreter(MOBILENET_MODEL_FILE); - try { - interpreter.getInputIndex("WrongInputName"); - fail(); - } catch (IllegalArgumentException e) { - assertThat(e) - .hasMessageThat() - .contains( - "'WrongInputName' is not a valid name for any input. Names of inputs and their " - + "indexes are {input=0}"); - } - int index = interpreter.getInputIndex("input"); - assertThat(index).isEqualTo(0); - } - - @Test - public void testGetOutputIndex() { - Interpreter interpreter = new Interpreter(MOBILENET_MODEL_FILE); - try { - interpreter.getOutputIndex("WrongOutputName"); - fail(); - } catch (IllegalArgumentException e) { - assertThat(e) - .hasMessageThat() - .contains( - "'WrongOutputName' is not a valid name for any output. Names of outputs and their" - + " indexes are {MobilenetV1/Predictions/Softmax=0}"); - } - int index = interpreter.getOutputIndex("MobilenetV1/Predictions/Softmax"); - assertThat(index).isEqualTo(0); - } - - @Test - public void testTurnOnNNAPI() throws Exception { - Path path = MODEL_FILE.toPath(); - FileChannel fileChannel = - (FileChannel) Files.newByteChannel(path, EnumSet.of(StandardOpenOption.READ)); - MappedByteBuffer mappedByteBuffer = - fileChannel.map(FileChannel.MapMode.READ_ONLY, 0, fileChannel.size()); - Interpreter interpreter = - new Interpreter( - mappedByteBuffer, - new Interpreter.Options().setUseNNAPI(true).setAllowFp16PrecisionForFp32(true)); - float[] oneD = {1.23f, 6.54f, 7.81f}; - float[][] twoD = {oneD, oneD, oneD, oneD, oneD, oneD, oneD, oneD}; - float[][][] threeD = {twoD, twoD, twoD, twoD, twoD, twoD, twoD, twoD}; - float[][][][] fourD = {threeD, threeD}; - float[][][][] parsedOutputs = new float[2][8][8][3]; - interpreter.run(fourD, parsedOutputs); - float[] outputOneD = parsedOutputs[0][0][0]; - float[] expected = {3.69f, 19.62f, 23.43f}; - assertThat(outputOneD).usingTolerance(0.1f).containsExactly(expected).inOrder(); - interpreter.close(); - fileChannel.close(); - } - - @Test - public void testRedundantClose() throws Exception { - Interpreter interpreter = new Interpreter(MODEL_FILE); - interpreter.close(); - interpreter.close(); - } - - /** Smoke test validating that flex model loading fails when the flex delegate is not linked. */ - @Test - public void testFlexModel() throws Exception { - try { - new Interpreter(FLEX_MODEL_FILE); - fail(); - } catch (IllegalStateException e) { - // Expected failure. - } - } -} diff --git a/tensorflow/contrib/lite/java/src/testhelper/java/org/tensorflow/lite/BUILD b/tensorflow/contrib/lite/java/src/testhelper/java/org/tensorflow/lite/BUILD deleted file mode 100644 index af1d99ef41e6413d8ef2c6f478aaa8f9e3931ff8..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/java/src/testhelper/java/org/tensorflow/lite/BUILD +++ /dev/null @@ -1,20 +0,0 @@ -# Description: -# Internal helper function to test TF Lite API. - -load("@build_bazel_rules_android//android:rules.bzl", "android_library") - -package(default_visibility = ["//visibility:public"]) - -licenses(["notice"]) # Apache 2.0 - -android_library( - name = "testhelper", - srcs = glob( - [ - "*.java", - ], - ), - deps = [ - "//tensorflow/contrib/lite/java:tensorflowlite_java", - ], -) diff --git a/tensorflow/contrib/lite/kernels/BUILD b/tensorflow/contrib/lite/kernels/BUILD deleted file mode 100644 index f20bb420a0240b01af9248cd5c5ba60f5e329f55..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/BUILD +++ /dev/null @@ -1,1346 +0,0 @@ -package(default_visibility = [ - "//visibility:public", -]) - -licenses(["notice"]) # Apache 2.0 - -load("//tensorflow/contrib/lite:build_def.bzl", "tflite_copts") -load("//tensorflow/contrib/lite:special_rules.bzl", "tflite_portable_test_suite") -load("//tensorflow:tensorflow.bzl", "tf_cc_test", "tf_opts_nortti_if_android") - -# Suppress warnings that are introduced by Eigen Tensor. -EXTRA_EIGEN_COPTS = select({ - "//tensorflow:ios": [ - "-Wno-error=invalid-partial-specialization", - "-Wno-error=reorder", - ], - "//tensorflow:windows": [ - "/DEIGEN_HAS_C99_MATH", - "/DEIGEN_AVOID_STL_ARRAY", - ], - "//conditions:default": ["-Wno-error=reorder"], -}) - -tf_cc_test( - name = "optional_tensor_test", - size = "small", - srcs = ["optional_tensor_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -cc_library( - name = "test_util", - testonly = 1, - srcs = ["test_util.cc"], - hdrs = ["test_util.h"], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite:schema_fbs_version", - "//tensorflow/contrib/lite:string_util", - "//tensorflow/contrib/lite/kernels/internal:tensor_utils", - "//tensorflow/contrib/lite/testing:util", - "//tensorflow/core:tflite_portable_logging", - "@com_google_googletest//:gtest", - ], -) - -cc_library( - name = "eigen_support", - srcs = [ - "eigen_support.cc", - ], - hdrs = [ - "eigen_support.h", - ], - copts = tflite_copts() + EXTRA_EIGEN_COPTS, - deps = [ - ":op_macros", - "//tensorflow/contrib/lite:arena_planner", - "//tensorflow/contrib/lite/c:c_api_internal", - "//tensorflow/contrib/lite/kernels/internal:optimized", - ], -) - -cc_library( - name = "gemm_support", - srcs = [ - "gemm_support.cc", - ], - hdrs = [ - "gemm_support.h", - ], - copts = tflite_copts(), - deps = [ - ":op_macros", - "//tensorflow/contrib/lite/c:c_api_internal", - "@gemmlowp", - ], -) - -cc_library( - name = "activation_functor", - hdrs = [ - "activation_functor.h", - ], - deps = [ - "//tensorflow/contrib/lite/c:c_api_internal", - ], -) - -cc_library( - name = "op_macros", - hdrs = [ - "op_macros.h", - ], -) - -cc_library( - name = "kernel_util", - srcs = [ - "kernel_util.cc", - ], - hdrs = [ - "kernel_util.h", - ], - deps = [ - "//tensorflow/contrib/lite/c:c_api_internal", - "//tensorflow/contrib/lite/kernels/internal:round", - "//tensorflow/contrib/lite/kernels/internal:types", - ], -) - -tf_cc_test( - name = "kernel_util_test", - size = "small", - srcs = ["kernel_util_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":kernel_util", - "//tensorflow/contrib/lite/testing:util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "test_util_test", - size = "small", - srcs = ["test_util_test.cc"], - tags = ["no_oss"], - deps = [ - ":test_util", - "//tensorflow/contrib/lite/testing:util", - "@com_google_googletest//:gtest", - ], -) - -cc_library( - name = "padding", - srcs = [], - hdrs = ["padding.h"], - deps = [ - "//tensorflow/contrib/lite/c:c_api_internal", - ], -) - -cc_library( - name = "builtin_op_kernels", - srcs = [ - "activations.cc", - "add.cc", - "arg_min_max.cc", - "audio_spectrogram.cc", - "basic_rnn.cc", - "batch_to_space_nd.cc", - "bidirectional_sequence_lstm.cc", - "bidirectional_sequence_rnn.cc", - "cast.cc", - "comparisons.cc", - "concatenation.cc", - "conv.cc", - "depthwise_conv.cc", - "dequantize.cc", - "detection_postprocess.cc", - "div.cc", - "elementwise.cc", - "embedding_lookup.cc", - "embedding_lookup_sparse.cc", - "exp.cc", - "expand_dims.cc", - "fake_quant.cc", - "floor.cc", - "floor_div.cc", - "fully_connected.cc", - "gather.cc", - "hashtable_lookup.cc", - "l2norm.cc", - "layer_norm_lstm.cc", - "local_response_norm.cc", - "logical.cc", - "lsh_projection.cc", - "lstm.cc", - "maximum_minimum.cc", - "mfcc.cc", - "mul.cc", - "neg.cc", - "one_hot.cc", - "pack.cc", - "pad.cc", - "pooling.cc", - "pow.cc", - "reduce.cc", - "relu1.cc", - "reshape.cc", - "resize_bilinear.cc", - "select.cc", - "shape.cc", - "skip_gram.cc", - "slice.cc", - "space_to_batch_nd.cc", - "space_to_depth.cc", - "sparse_output_fully_connected.cc", - "sparse_to_dense.cc", - "split.cc", - "squeeze.cc", - "strided_slice.cc", - "sub.cc", - "svdf.cc", - "tile.cc", - "topk_v2.cc", - "transpose.cc", - "transpose_conv.cc", - "unidirectional_sequence_lstm.cc", - "unidirectional_sequence_rnn.cc", - "unpack.cc", - "zeros_like.cc", - ], - hdrs = [ - ], - copts = tflite_copts() + tf_opts_nortti_if_android() + EXTRA_EIGEN_COPTS, - visibility = ["//visibility:private"], - deps = [ - ":activation_functor", - ":eigen_support", - ":kernel_util", - ":lstm_eval", - ":op_macros", - ":padding", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite:string_util", - "//tensorflow/contrib/lite/c:c_api_internal", - "//tensorflow/contrib/lite/kernels:gemm_support", - "//tensorflow/contrib/lite/kernels/internal:audio_utils", - "//tensorflow/contrib/lite/kernels/internal:kernel_utils", - "//tensorflow/contrib/lite/kernels/internal:optimized", - "//tensorflow/contrib/lite/kernels/internal:optimized_base", - "//tensorflow/contrib/lite/kernels/internal:quantization_util", - "//tensorflow/contrib/lite/kernels/internal:reference_base", - "//tensorflow/contrib/lite/kernels/internal:tensor", - "//tensorflow/contrib/lite/kernels/internal:tensor_utils", - "@farmhash_archive//:farmhash", - "@flatbuffers", - ], -) - -cc_library( - name = "lstm_eval", - srcs = ["lstm_eval.cc"], - hdrs = ["lstm_eval.h"], - deps = [ - ":op_macros", - "//tensorflow/contrib/lite/c:c_api_internal", - "//tensorflow/contrib/lite/kernels/internal:kernel_utils", - "//tensorflow/contrib/lite/kernels/internal:tensor_utils", - ], -) - -cc_library( - name = "builtin_ops", - srcs = ["register.cc"], - hdrs = ["register.h"], - deps = [ - ":builtin_op_kernels", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite:util", - "//tensorflow/contrib/lite/c:c_api_internal", - ], -) - -tf_cc_test( - name = "audio_spectrogram_test", - size = "small", - srcs = ["audio_spectrogram_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - "@flatbuffers", - ], -) - -tf_cc_test( - name = "mfcc_test", - size = "small", - srcs = ["mfcc_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - "@flatbuffers", - ], -) - -tf_cc_test( - name = "detection_postprocess_test", - size = "small", - srcs = ["detection_postprocess_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - "@flatbuffers", - ], -) - -tf_cc_test( - name = "relu1_test", - size = "small", - srcs = ["relu1_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - "@flatbuffers", - ], -) - -tf_cc_test( - name = "sparse_output_fully_connected_test", - size = "small", - srcs = ["sparse_output_fully_connected_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - "@flatbuffers", - ], -) - -tf_cc_test( - name = "activations_test", - size = "small", - srcs = ["activations_test.cc"], - tags = ["tflite_not_portable_ios"], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "add_test", - size = "small", - srcs = ["add_test.cc"], - tags = ["tflite_not_portable_ios"], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "arg_min_max_test", - size = "small", - srcs = ["arg_min_max_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "div_test", - size = "small", - srcs = ["div_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "sub_test", - size = "small", - srcs = ["sub_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "transpose_test", - size = "small", - srcs = ["transpose_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "//tensorflow/contrib/lite/kernels/internal:reference", - "//tensorflow/contrib/lite/kernels/internal:reference_base", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "space_to_batch_nd_test", - size = "small", - srcs = ["space_to_batch_nd_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "batch_to_space_nd_test", - size = "small", - srcs = ["batch_to_space_nd_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "cast_test", - size = "small", - srcs = ["cast_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "concatenation_test", - size = "small", - srcs = ["concatenation_test.cc"], - tags = ["tflite_not_portable_ios"], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "conv_test", - size = "small", - srcs = ["conv_test.cc"], - tags = ["tflite_not_portable_ios"], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_absl//absl/memory", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "depthwise_conv_test", - size = "small", - srcs = ["depthwise_conv_test.cc"], - tags = ["tflite_not_portable_ios"], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_absl//absl/memory", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "dequantize_test", - size = "small", - srcs = ["dequantize_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_absl//absl/memory", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "basic_rnn_test", - size = "small", - srcs = ["basic_rnn_test.cc"], - tags = ["tflite_not_portable_ios"], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "bidirectional_sequence_lstm_test", - size = "small", - srcs = ["bidirectional_sequence_lstm_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "//tensorflow/contrib/lite/schema:schema_fbs", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "floor_test", - size = "small", - srcs = ["floor_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "elementwise_test", - size = "small", - srcs = ["elementwise_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "unidirectional_sequence_lstm_test", - size = "small", - srcs = ["unidirectional_sequence_lstm_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "bidirectional_sequence_rnn_test", - size = "small", - srcs = ["bidirectional_sequence_rnn_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "unidirectional_sequence_rnn_test", - size = "small", - srcs = ["unidirectional_sequence_rnn_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "l2norm_test", - size = "small", - srcs = ["l2norm_test.cc"], - tags = ["tflite_not_portable_ios"], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "exp_test", - size = "small", - srcs = ["exp_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "fake_quant_test", - size = "small", - srcs = ["fake_quant_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "maximum_minimum_test", - size = "small", - srcs = ["maximum_minimum_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "reduce_test", - size = "small", - srcs = ["reduce_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "mul_test", - size = "small", - srcs = ["mul_test.cc"], - tags = ["tflite_not_portable_ios"], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "pad_test", - size = "small", - srcs = ["pad_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "reshape_test", - size = "small", - srcs = ["reshape_test.cc"], - tags = ["tflite_not_portable_ios"], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "gather_test", - size = "small", - srcs = ["gather_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/c:c_api_internal", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "topk_v2_test", - size = "small", - srcs = ["topk_v2_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/c:c_api_internal", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "resize_bilinear_test", - size = "small", - srcs = ["resize_bilinear_test.cc"], - tags = ["tflite_not_portable_ios"], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "svdf_test", - size = "small", - srcs = ["svdf_test.cc"], - tags = ["tflite_not_portable_ios"], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "embedding_lookup_test", - size = "small", - srcs = ["embedding_lookup_test.cc"], - tags = ["tflite_not_portable_ios"], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "embedding_lookup_sparse_test", - size = "small", - srcs = ["embedding_lookup_sparse_test.cc"], - tags = ["tflite_not_portable_ios"], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "fully_connected_test", - size = "small", - srcs = ["fully_connected_test.cc"], - tags = ["tflite_not_portable_ios"], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "//tensorflow/contrib/lite/kernels/internal:tensor_utils", - "@com_google_absl//absl/memory", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "local_response_norm_test", - size = "small", - srcs = ["local_response_norm_test.cc"], - tags = ["tflite_not_portable_ios"], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "pooling_test", - size = "small", - srcs = ["pooling_test.cc"], - tags = ["tflite_not_portable_ios"], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "softmax_test", - size = "small", - srcs = ["softmax_test.cc"], - tags = ["tflite_not_portable_ios"], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "//tensorflow/contrib/lite/kernels/internal:reference_base", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "log_softmax_test", - size = "small", - srcs = ["log_softmax_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "//tensorflow/contrib/lite/kernels/internal:reference_base", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "lsh_projection_test", - size = "small", - srcs = ["lsh_projection_test.cc"], - tags = ["tflite_not_portable_ios"], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "hashtable_lookup_test", - size = "small", - srcs = ["hashtable_lookup_test.cc"], - tags = ["tflite_not_portable_ios"], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite:string_util", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "layer_norm_lstm_test", - size = "small", - srcs = ["layer_norm_lstm_test.cc"], - tags = ["tflite_not_portable_ios"], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - "@flatbuffers", - ], -) - -tf_cc_test( - name = "lstm_test", - size = "small", - srcs = ["lstm_test.cc"], - tags = ["tflite_not_portable_ios"], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "skip_gram_test", - size = "small", - srcs = ["skip_gram_test.cc"], - tags = ["tflite_not_portable_ios"], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite:string_util", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "space_to_depth_test", - size = "small", - srcs = ["space_to_depth_test.cc"], - tags = ["tflite_not_portable_ios"], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "split_test", - size = "small", - srcs = ["split_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "squeeze_test", - size = "small", - srcs = ["squeeze_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "strided_slice_test", - size = "small", - srcs = ["strided_slice_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "tile_test", - size = "small", - srcs = ["tile_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/c:c_api_internal", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "comparisons_test", - size = "small", - srcs = [ - "comparisons_test.cc", - ], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "neg_test", - size = "small", - srcs = ["neg_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "select_test", - size = "small", - srcs = [ - "select_test.cc", - ], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "slice_test", - size = "small", - srcs = [ - "slice_test.cc", - ], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "transpose_conv_test", - size = "small", - srcs = ["transpose_conv_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "expand_dims_test", - size = "small", - srcs = ["expand_dims_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/c:c_api_internal", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "sparse_to_dense_test", - size = "small", - srcs = ["sparse_to_dense_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/c:c_api_internal", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "shape_test", - size = "small", - srcs = ["shape_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/c:c_api_internal", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "pow_test", - size = "small", - srcs = ["pow_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/c:c_api_internal", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "pack_test", - size = "small", - srcs = ["pack_test.cc"], - tags = ["tflite_not_portable_ios"], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/c:c_api_internal", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "one_hot_test", - size = "small", - srcs = ["one_hot_test.cc"], - tags = ["tflite_not_portable_ios"], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "logical_test", - size = "small", - srcs = ["logical_test.cc"], - tags = ["tflite_not_portable_ios"], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/c:c_api_internal", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "unpack_test", - size = "small", - srcs = ["unpack_test.cc"], - tags = ["tflite_not_portable_ios"], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:builtin_op_data", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "floor_div_test", - size = "small", - srcs = ["floor_div_test.cc"], - tags = ["tflite_not_portable_ios"], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:builtin_op_data", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -tf_cc_test( - name = "zeros_like_test", - size = "small", - srcs = ["zeros_like_test.cc"], - tags = ["tflite_not_portable_ios"], - deps = [ - ":builtin_ops", - "//tensorflow/contrib/lite:builtin_op_data", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -filegroup( - name = "all_files", - srcs = glob( - ["**/*"], - exclude = [ - "**/METADATA", - "**/OWNERS", - ], - ), - visibility = ["//tensorflow:__subpackages__"], -) - -tflite_portable_test_suite() diff --git a/tensorflow/contrib/lite/kernels/activations.cc b/tensorflow/contrib/lite/kernels/activations.cc deleted file mode 100644 index 9aed4f09b82cc0ac70c68a4da46706a6244084aa..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/activations.cc +++ /dev/null @@ -1,726 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include -#include -#include -#include -#include -#include - -#include "tensorflow/contrib/lite/c/builtin_op_data.h" -#include "tensorflow/contrib/lite/c/c_api_internal.h" -#include "tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h" -#include "tensorflow/contrib/lite/kernels/internal/quantization_util.h" -#include "tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h" -#include "tensorflow/contrib/lite/kernels/internal/tensor.h" -#include "tensorflow/contrib/lite/kernels/kernel_util.h" -#include "tensorflow/contrib/lite/kernels/op_macros.h" - -namespace tflite { -namespace ops { -namespace builtin { -namespace activations { - -struct OpData { - int32_t input_multiplier = 0; - int input_left_shift = 0; - int32_t input_range_radius = 0; - int diff_min = 0; -}; - -struct LogSoftmaxOpData : public OpData { - int32_t reverse_scaling_divisor = 0; - int32_t reverse_scaling_right_shift = 0; -}; - -void* Init(TfLiteContext* context, const char* buffer, size_t length) { - // This is a builtin op, so we don't use the contents in 'buffer', if any. - // Instead, we allocate a new object to carry information from Prepare() to - // Eval(). - return new OpData; -} - -void* LogSoftmaxInit(TfLiteContext* context, const char* buffer, - size_t length) { - return new LogSoftmaxOpData; -} - -void Free(TfLiteContext* context, void* buffer) { - delete reinterpret_cast(buffer); -} - -void LogSoftmaxFree(TfLiteContext* context, void* buffer) { - delete reinterpret_cast(buffer); -} - -TfLiteStatus GenericPrepare(TfLiteContext* context, TfLiteNode* node) { - TF_LITE_ENSURE_EQ(context, NumInputs(node), 1); - TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); - const TfLiteTensor* input = GetInput(context, node, 0); - TfLiteTensor* output = GetOutput(context, node, 0); - TF_LITE_ENSURE_EQ(context, input->type, output->type); - - return context->ResizeTensor(context, output, - TfLiteIntArrayCopy(input->dims)); -} - -TfLiteStatus TanhPrepare(TfLiteContext* context, TfLiteNode* node) { - OpData* data = reinterpret_cast(node->user_data); - - TF_LITE_ENSURE_EQ(context, NumInputs(node), 1); - TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); - const TfLiteTensor* input = GetInput(context, node, 0); - TfLiteTensor* output = GetOutput(context, node, 0); - TF_LITE_ENSURE_EQ(context, input->type, output->type); - - if (input->type == kTfLiteUInt8) { - static constexpr int kInputIntegerBits = 4; - - const double input_real_multiplier = - input->params.scale * - static_cast(1 << (31 - kInputIntegerBits)); - - QuantizeMultiplierGreaterThanOne(input_real_multiplier, - &data->input_multiplier, - &data->input_left_shift); - data->input_range_radius = - CalculateInputRadius(kInputIntegerBits, data->input_left_shift); - } else if (input->type == kTfLiteInt16) { - static constexpr int kInputIntegerBits = 3; - static constexpr int kOutputFractionalBits = 15; - - // These operators are implemented in fixed-point arithmetic, - // which intrinsically wants symmetric ranges (zero_point==0) - // and power-of-two scales (power-of-two is abbreviated below as POT). - // While more general support would be possible by means of rescaling, - // that would add some overhead and some loss of accuracy and wouldn't - // be used at the moment as current quantized LSTM applications are - // happy with symmetric, power-of-two-scales quantization. So we just - // implement that narrow case only for now. - - TF_LITE_ENSURE_EQ(context, input->params.zero_point, 0); - TF_LITE_ENSURE_EQ(context, output->params.zero_point, 0); - - int input_scale_log2_rounded; - TF_LITE_ENSURE(context, - CheckedLog2(input->params.scale, &input_scale_log2_rounded)); - - int output_scale_log2_rounded; - TF_LITE_ENSURE( - context, CheckedLog2(output->params.scale, &output_scale_log2_rounded)); - TF_LITE_ENSURE_EQ(context, output_scale_log2_rounded, - -kOutputFractionalBits); - - data->input_left_shift = - (15 - kInputIntegerBits) + input_scale_log2_rounded; - // Support for shifts is limited until we have a parameterized version of - // SaturatingRoundingMultiplyByPOT(). - TF_LITE_ENSURE(context, data->input_left_shift >= 0); - TF_LITE_ENSURE(context, data->input_left_shift <= 1); - } - - return context->ResizeTensor(context, output, - TfLiteIntArrayCopy(input->dims)); -} - -TfLiteStatus SigmoidPrepare(TfLiteContext* context, TfLiteNode* node) { - OpData* data = reinterpret_cast(node->user_data); - - TF_LITE_ENSURE_EQ(context, NumInputs(node), 1); - TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); - const TfLiteTensor* input = GetInput(context, node, 0); - TfLiteTensor* output = GetOutput(context, node, 0); - TF_LITE_ENSURE_EQ(context, input->type, output->type); - - if (input->type == kTfLiteUInt8) { - TF_LITE_ENSURE_EQ(context, output->params.zero_point, 0); - TF_LITE_ENSURE(context, output->params.scale == 1. / 256); - - static constexpr int kInputIntegerBits = 4; - - const double input_real_multiplier = - input->params.scale * - static_cast(1 << (31 - kInputIntegerBits)); - - QuantizeMultiplierGreaterThanOne(input_real_multiplier, - &data->input_multiplier, - &data->input_left_shift); - data->input_range_radius = - CalculateInputRadius(kInputIntegerBits, data->input_left_shift); - } else if (input->type == kTfLiteInt16) { - static constexpr int kInputIntegerBits = 3; - static constexpr int kOutputFractionalBits = 15; - - // See comments in TanhPrepare about requiring zero_point==0 - // and a power-of-two ("POT") scale. - - TF_LITE_ENSURE_EQ(context, input->params.zero_point, 0); - TF_LITE_ENSURE_EQ(context, output->params.zero_point, 0); - - int input_scale_log2_rounded; - TF_LITE_ENSURE(context, - CheckedLog2(input->params.scale, &input_scale_log2_rounded)); - - int output_scale_log2_rounded; - TF_LITE_ENSURE( - context, CheckedLog2(output->params.scale, &output_scale_log2_rounded)); - TF_LITE_ENSURE_EQ(context, output_scale_log2_rounded, - -kOutputFractionalBits); - - data->input_left_shift = - (15 - kInputIntegerBits) + input_scale_log2_rounded; - // The int16 logistic implementation does not support shifting of the input. - TF_LITE_ENSURE_EQ(context, data->input_left_shift, 0); - } - - return context->ResizeTensor(context, output, - TfLiteIntArrayCopy(input->dims)); -} - -TfLiteStatus SoftmaxPrepare(TfLiteContext* context, TfLiteNode* node) { - auto* params = reinterpret_cast(node->builtin_data); - OpData* data = reinterpret_cast(node->user_data); - - TF_LITE_ENSURE_EQ(context, NumInputs(node), 1); - TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); - const TfLiteTensor* input = GetInput(context, node, 0); - TfLiteTensor* output = GetOutput(context, node, 0); - TF_LITE_ENSURE_EQ(context, input->type, output->type); - - const int num_dims = NumDimensions(input); - TF_LITE_ENSURE(context, num_dims >= 1 && num_dims <= 4); - - if (input->type == kTfLiteUInt8) { - TF_LITE_ENSURE_EQ(context, output->params.zero_point, 0); - TF_LITE_ENSURE(context, output->params.scale == 1. / 256); - - static const int kScaledDiffIntegerBits = 5; - - tflite::PreprocessSoftmaxScaling( - params->beta, input->params.scale, kScaledDiffIntegerBits, - &data->input_multiplier, &data->input_left_shift); - data->diff_min = -1.0 * tflite::CalculateInputRadius( - kScaledDiffIntegerBits, data->input_left_shift); - } - - return context->ResizeTensor(context, output, - TfLiteIntArrayCopy(input->dims)); -} - -TfLiteStatus LogSoftmaxPrepare(TfLiteContext* context, TfLiteNode* node) { - LogSoftmaxOpData* data = reinterpret_cast(node->user_data); - - TF_LITE_ENSURE_EQ(context, NumInputs(node), 1); - TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); - const TfLiteTensor* input = GetInput(context, node, 0); - TfLiteTensor* output = GetOutput(context, node, 0); - TF_LITE_ENSURE_EQ(context, input->type, output->type); - - if (input->type == kTfLiteUInt8) { - TF_LITE_ENSURE_EQ(context, output->params.zero_point, 255); - TF_LITE_ENSURE_EQ(context, output->params.scale, 16.0 / 256); - - static const double kBeta = 1.0; - static const int kScaledDiffIntegerBits = 5; - tflite::PreprocessLogSoftmaxScalingExp( - kBeta, input->params.scale, kScaledDiffIntegerBits, - &data->input_multiplier, &data->input_left_shift, - &data->reverse_scaling_divisor, &data->reverse_scaling_right_shift); - data->reverse_scaling_right_shift *= -1; - data->diff_min = -1.0 * tflite::CalculateInputRadius( - kScaledDiffIntegerBits, data->input_left_shift); - } - - return context->ResizeTensor(context, output, - TfLiteIntArrayCopy(input->dims)); -} - -TfLiteStatus PreluPrepare(TfLiteContext* context, TfLiteNode* node) { - TF_LITE_ENSURE_EQ(context, NumInputs(node), 2); - TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); - const TfLiteTensor* input = GetInput(context, node, 0); - TfLiteTensor* output = GetOutput(context, node, 0); - const TfLiteTensor* alpha = GetInput(context, node, 1); - - // Currently only Float32 is supported - // TODO(ycling): Support other data types. - TF_LITE_ENSURE_EQ(context, input->type, kTfLiteFloat32); - TF_LITE_ENSURE_EQ(context, alpha->type, kTfLiteFloat32); - output->type = input->type; - - // PRelu (parameteric Relu) shares the same alpha value on "shared axis". - // This means it's always required to "broadcast" alpha values in PRelu. - TfLiteIntArray* output_size = nullptr; - TF_LITE_ENSURE_OK( - context, CalculateShapeForBroadcast(context, input, alpha, &output_size)); - - TF_LITE_ENSURE_OK(context, - context->ResizeTensor(context, output, output_size)); - // After broadcasting, the output shape should always be the same as the - // input shape. - TF_LITE_ENSURE(context, HaveSameShapes(input, output)); - - return kTfLiteOk; -} - -TfLiteStatus ReluEval(TfLiteContext* context, TfLiteNode* node) { - const TfLiteTensor* input = GetInput(context, node, 0); - TfLiteTensor* output = GetOutput(context, node, 0); - switch (input->type) { - case kTfLiteFloat32: { - size_t elements = input->bytes / sizeof(float); - float* in = input->data.f; - float* in_end = in + elements; - float* out = output->data.f; - for (; in < in_end; in++, out++) *out = std::max(0.f, *in); - return kTfLiteOk; - } break; - default: - context->ReportError(context, "Only float32 supported currently, got %d.", - input->type); - return kTfLiteError; - } -} - -TfLiteStatus Relu1Eval(TfLiteContext* context, TfLiteNode* node) { - const TfLiteTensor* input = GetInput(context, node, 0); - TfLiteTensor* output = GetOutput(context, node, 0); - switch (input->type) { - case kTfLiteFloat32: { - size_t elements = input->bytes / sizeof(float); - float* in = input->data.f; - float* in_end = in + elements; - float* out = output->data.f; - for (; in < in_end; in++, out++) { - *out = std::min(std::max(-1.f, *in), 1.f); - } - return kTfLiteOk; - } break; - default: - context->ReportError(context, "Only float32 supported currently, got %d.", - input->type); - return kTfLiteError; - } -} - -TfLiteStatus Relu6Eval(TfLiteContext* context, TfLiteNode* node) { - const TfLiteTensor* input = GetInput(context, node, 0); - TfLiteTensor* output = GetOutput(context, node, 0); - switch (input->type) { - case kTfLiteFloat32: { - size_t elements = input->bytes / sizeof(float); - float* in = input->data.f; - float* in_end = in + elements; - float* out = output->data.f; - for (; in < in_end; in++, out++) *out = std::min(std::max(0.f, *in), 6.f); - return kTfLiteOk; - } break; - default: - context->ReportError(context, "Only float32 supported currently, got %d.", - input->type); - return kTfLiteError; - } -} - -TfLiteStatus TanhEval(TfLiteContext* context, TfLiteNode* node) { - OpData* data = reinterpret_cast(node->user_data); - const TfLiteTensor* input = GetInput(context, node, 0); - TfLiteTensor* output = GetOutput(context, node, 0); - switch (input->type) { - case kTfLiteFloat32: { - size_t elements = input->bytes / sizeof(float); - float* in = input->data.f; - float* in_end = in + elements; - float* out = output->data.f; - for (; in < in_end; in++, out++) *out = std::tanh(*in); - return kTfLiteOk; - } break; - case kTfLiteInt16: { - TanhParams params; - params.input_left_shift = data->input_left_shift; - optimized_ops::Tanh(params, GetTensorShape(input), - GetTensorData(input), GetTensorShape(output), - GetTensorData(output)); - return kTfLiteOk; - } break; - case kTfLiteUInt8: { - TanhParams params; - params.input_zero_point = input->params.zero_point; - params.input_range_radius = data->input_range_radius; - params.input_multiplier = data->input_multiplier; - params.input_left_shift = data->input_left_shift; - optimized_ops::Tanh(params, GetTensorShape(input), - GetTensorData(input), GetTensorShape(output), - GetTensorData(output)); - return kTfLiteOk; - } break; - default: - context->ReportError(context, "Only float32 supported currently, got %d.", - input->type); - return kTfLiteError; - } -} - -// Sigmoid is also know as "Logistic". -TfLiteStatus SigmoidEval(TfLiteContext* context, TfLiteNode* node) { - OpData* data = reinterpret_cast(node->user_data); - - const TfLiteTensor* input = GetInput(context, node, 0); - TfLiteTensor* output = GetOutput(context, node, 0); - switch (input->type) { - case kTfLiteFloat32: { - size_t elements = input->bytes / sizeof(float); - float* in = input->data.f; - float* in_end = in + elements; - float* out = output->data.f; - for (; in < in_end; in++, out++) *out = 1.f / (1.f + std::exp(-*in)); - break; - } - case kTfLiteInt16: { - LogisticParams params; - optimized_ops::Logistic( - params, GetTensorShape(input), GetTensorData(input), - GetTensorShape(output), GetTensorData(output)); - break; - } - case kTfLiteUInt8: { - LogisticParams params; - params.input_zero_point = input->params.zero_point; - params.input_range_radius = data->input_range_radius; - params.input_multiplier = data->input_multiplier; - params.input_left_shift = data->input_left_shift; - optimized_ops::Logistic( - params, GetTensorShape(input), GetTensorData(input), - GetTensorShape(output), GetTensorData(output)); - break; - } - default: - context->ReportError(context, "Only float32 supported currently, got %d.", - input->type); - return kTfLiteError; - } - return kTfLiteOk; -} - -// Performs softmax along the input of size (input_size * batch_size). -void Softmax(const float* in, const int input_size, const int batch_size, - const float beta, float* out) { - TF_LITE_ASSERT(input_size > 0); - - // For each batch - for (int b = 0; b < batch_size; b++) { - // Find the max coeff. - float max_coeff = in[0]; - for (int i = 1; i < input_size; i++) { - if (in[i] > max_coeff) max_coeff = in[i]; - } - - // Compute the normalized sum of exps. - float exp_sum = 0.0; - for (int i = 0; i < input_size; i++) { - out[i] = std::exp((in[i] - max_coeff) * beta); - exp_sum += out[i]; - } - - // Divide by the sum of exps. - float reciprocal_sum_exp = 1.f / exp_sum; - for (int i = 0; i < input_size; i++) { - out[i] *= reciprocal_sum_exp; - } - - // Advance in and out pointers for the next batch. - in += input_size; - out += input_size; - } -} - -// Takes a 1D tensor and performs softmax along it. -void Softmax1DFloat(const TfLiteTensor* input, TfLiteTensor* output, - TfLiteSoftmaxParams* params) { - const int input_size = input->dims->data[0]; - Softmax(input->data.f, input_size, 1, params->beta, output->data.f); -} - -// Takes a 2D tensor and perform softmax along the last dimension. -void Softmax2DFloat(const TfLiteTensor* input, TfLiteTensor* output, - TfLiteSoftmaxParams* params) { - const int batch_size = input->dims->data[0]; - const int input_size = input->dims->data[1]; - Softmax(input->data.f, input_size, batch_size, params->beta, output->data.f); -} - -// Takes a 3D tensor and perform softmax along the last dimension. -void Softmax3DFloat(const TfLiteTensor* input, TfLiteTensor* output, - TfLiteSoftmaxParams* params) { - const int batch_size = input->dims->data[0]; - const int intermediate_size = input->dims->data[1]; - const int input_size = input->dims->data[2]; - SoftmaxParams op_params; - op_params.beta = params->beta; - optimized_ops::Softmax( - op_params, GetTensorShape({batch_size, intermediate_size, 1, input_size}), - GetTensorData(input), - GetTensorShape({batch_size, intermediate_size, 1, input_size}), - GetTensorData(output)); -} - -void Softmax1DQuantized(const TfLiteTensor* input, TfLiteTensor* output, - TfLiteSoftmaxParams* params, OpData* data) { - // TODO(ahentz): this is arguably a dirty trick. Since the implementation - // always traverses the last dimension of a 4D tensor, we will pretend our 1D - // tensor is 4D in a special way. We will convert a (Y) shape into a (1, - // 1, 1, Y) shape. - const int input_size = input->dims->data[0]; - SoftmaxParams op_params; - op_params.input_multiplier = data->input_multiplier; - op_params.input_left_shift = data->input_left_shift; - op_params.diff_min = data->diff_min; - optimized_ops::Softmax(op_params, GetTensorShape({1, 1, 1, input_size}), - GetTensorData(input), - GetTensorShape({1, 1, 1, input_size}), - GetTensorData(output)); -} -void Softmax2DQuantized(const TfLiteTensor* input, TfLiteTensor* output, - TfLiteSoftmaxParams* params, OpData* data) { - // TODO(ahentz): this is arguably a dirty trick. Since the implementation - // always traverses the last dimension of a 4D tensor, we will pretend our 2D - // tensor is 4D in a special way. We will convert a (X, Y) shape into a (X, - // 1, 1, Y) shape. - const int batch_size = input->dims->data[0]; - const int input_size = input->dims->data[1]; - SoftmaxParams op_params; - op_params.input_multiplier = data->input_multiplier; - op_params.input_left_shift = data->input_left_shift; - op_params.diff_min = data->diff_min; - optimized_ops::Softmax(op_params, - GetTensorShape({batch_size, 1, 1, input_size}), - GetTensorData(input), - GetTensorShape({batch_size, 1, 1, input_size}), - GetTensorData(output)); -} - -void Softmax3DQuantized(const TfLiteTensor* input, TfLiteTensor* output, - TfLiteSoftmaxParams* params, OpData* data) { - const int batch_size = input->dims->data[0]; - const int intermediate_size = input->dims->data[1]; - const int input_size = input->dims->data[2]; - SoftmaxParams op_params; - op_params.input_multiplier = data->input_multiplier; - op_params.input_left_shift = data->input_left_shift; - op_params.diff_min = data->diff_min; - optimized_ops::Softmax( - op_params, GetTensorShape({batch_size, intermediate_size, 1, input_size}), - GetTensorData(input), - GetTensorShape({batch_size, intermediate_size, 1, input_size}), - GetTensorData(output)); -} - -// Takes a 4D tensor and perform softmax along the forth dimension. -void Softmax4DFloat(const TfLiteTensor* input, TfLiteTensor* output, - TfLiteSoftmaxParams* params) { - SoftmaxParams op_params; - op_params.beta = params->beta; - optimized_ops::Softmax(op_params, GetTensorShape(input), - GetTensorData(input), GetTensorShape(output), - GetTensorData(output)); -} - -void Softmax4DQuantized(const TfLiteTensor* input, TfLiteTensor* output, - TfLiteSoftmaxParams* params, OpData* data) { - SoftmaxParams op_params; - op_params.input_multiplier = data->input_multiplier; - op_params.input_left_shift = data->input_left_shift; - op_params.diff_min = data->diff_min; - optimized_ops::Softmax(op_params, GetTensorShape(input), - GetTensorData(input), GetTensorShape(output), - GetTensorData(output)); -} - -TfLiteStatus SoftmaxEval(TfLiteContext* context, TfLiteNode* node) { - auto* params = reinterpret_cast(node->builtin_data); - OpData* data = reinterpret_cast(node->user_data); - - const TfLiteTensor* input = GetInput(context, node, 0); - TfLiteTensor* output = GetOutput(context, node, 0); - - // TODO(ahentz): consider an implementation that works for many (all?) - // dimensions. - switch (input->type) { - case kTfLiteFloat32: { - if (NumDimensions(input) == 1) { - Softmax1DFloat(input, output, params); - return kTfLiteOk; - } - if (NumDimensions(input) == 2) { - Softmax2DFloat(input, output, params); - return kTfLiteOk; - } - if (NumDimensions(input) == 3) { - Softmax3DFloat(input, output, params); - return kTfLiteOk; - } - if (NumDimensions(input) == 4) { - Softmax4DFloat(input, output, params); - return kTfLiteOk; - } - context->ReportError( - context, "Only 1D, 2D and 4D tensors supported currently, got %dD.", - NumDimensions(input)); - return kTfLiteError; - } - case kTfLiteUInt8: { - if (NumDimensions(input) == 1) { - Softmax1DQuantized(input, output, params, data); - return kTfLiteOk; - } - if (NumDimensions(input) == 2) { - Softmax2DQuantized(input, output, params, data); - return kTfLiteOk; - } - if (NumDimensions(input) == 3) { - Softmax3DQuantized(input, output, params, data); - return kTfLiteOk; - } - if (NumDimensions(input) == 4) { - Softmax4DQuantized(input, output, params, data); - return kTfLiteOk; - } - context->ReportError( - context, "Only 2D and 4D tensors supported currently, got %dD.", - NumDimensions(input)); - return kTfLiteError; - } - default: - context->ReportError( - context, "Only float32 and uint8_t supported currently, got %d.", - input->type); - return kTfLiteError; - } -} - -TfLiteStatus LogSoftmaxEval(TfLiteContext* context, TfLiteNode* node) { - const LogSoftmaxOpData* data = - reinterpret_cast(node->user_data); - const TfLiteTensor* input = GetInput(context, node, 0); - TfLiteTensor* output = GetOutput(context, node, 0); - switch (input->type) { - case kTfLiteFloat32: { - SoftmaxParams op_params; - optimized_ops::LogSoftmax( - op_params, GetTensorShape(input), GetTensorData(input), - GetTensorShape(output), GetTensorData(output)); - return kTfLiteOk; - } - case kTfLiteUInt8: { - SoftmaxParams op_params; - op_params.input_multiplier = data->input_multiplier; - op_params.input_left_shift = data->input_left_shift; - op_params.reverse_scaling_divisor = data->reverse_scaling_divisor; - op_params.reverse_scaling_right_shift = data->reverse_scaling_right_shift; - op_params.diff_min = data->diff_min; - optimized_ops::LogSoftmax( - op_params, GetTensorShape(input), GetTensorData(input), - GetTensorShape(output), GetTensorData(output)); - return kTfLiteOk; - } - default: - context->ReportError(context, "Only float32 supported currently., got %d", - input->type); - return kTfLiteError; - } -} - -template -T ApplyPrelu(T input, T alpha) { - return input >= 0.0 ? input : input * alpha; -} - -TfLiteStatus PreluEval(TfLiteContext* context, TfLiteNode* node) { - const TfLiteTensor* input = GetInput(context, node, 0); - const TfLiteTensor* alpha = GetInput(context, node, 1); - TfLiteTensor* output = GetOutput(context, node, 0); - if (input->type != kTfLiteFloat32) { - context->ReportError(context, "Only float32 supported currently, got %d.", - input->type); - return kTfLiteError; - } - reference_ops::BroadcastBinaryFunction4DSlow( - GetTensorShape(input), GetTensorData(input), GetTensorShape(alpha), - GetTensorData(alpha), GetTensorShape(output), - GetTensorData(output), ApplyPrelu); - return kTfLiteOk; -} - -} // namespace activations - -TfLiteRegistration* Register_RELU() { - static TfLiteRegistration r = {/*init=*/nullptr, /*free=*/nullptr, - activations::GenericPrepare, - activations::ReluEval}; - return &r; -} - -TfLiteRegistration* Register_RELU_N1_TO_1() { - static TfLiteRegistration r = {/*init=*/nullptr, /*free=*/nullptr, - activations::GenericPrepare, - activations::Relu1Eval}; - return &r; -} - -TfLiteRegistration* Register_RELU6() { - static TfLiteRegistration r = {/*init=*/nullptr, /*free=*/nullptr, - activations::GenericPrepare, - activations::Relu6Eval}; - return &r; -} - -TfLiteRegistration* Register_TANH() { - static TfLiteRegistration r = {activations::Init, activations::Free, - activations::TanhPrepare, - activations::TanhEval}; - return &r; -} - -TfLiteRegistration* Register_LOGISTIC() { - static TfLiteRegistration r = {activations::Init, activations::Free, - activations::SigmoidPrepare, - activations::SigmoidEval}; - return &r; -} - -TfLiteRegistration* Register_SOFTMAX() { - static TfLiteRegistration r = {activations::Init, activations::Free, - activations::SoftmaxPrepare, - activations::SoftmaxEval}; - return &r; -} - -TfLiteRegistration* Register_LOG_SOFTMAX() { - static TfLiteRegistration r = { - activations::LogSoftmaxInit, activations::LogSoftmaxFree, - activations::LogSoftmaxPrepare, activations::LogSoftmaxEval}; - return &r; -} - -TfLiteRegistration* Register_PRELU() { - static TfLiteRegistration r = {/*init=*/nullptr, /*free=*/nullptr, - activations::PreluPrepare, - activations::PreluEval}; - return &r; -} - -} // namespace builtin -} // namespace ops -} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/activations_test.cc b/tensorflow/contrib/lite/kernels/activations_test.cc deleted file mode 100644 index 9fa47e190a1dc797264e31979b9a6603ce8c5498..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/activations_test.cc +++ /dev/null @@ -1,616 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include -#include -#include "tensorflow/contrib/lite/interpreter.h" -#include "tensorflow/contrib/lite/kernels/register.h" -#include "tensorflow/contrib/lite/kernels/test_util.h" -#include "tensorflow/contrib/lite/model.h" - -namespace tflite { -namespace { - -using ::testing::ElementsAreArray; - -class BaseActivationsOpModel : public SingleOpModel { - public: - // Most activations don't take any options, so this constructor works for - // them. - BaseActivationsOpModel(BuiltinOperator type, TensorData input) { - input_ = AddInput(input); - if (input.type == TensorType_UINT8) { - output_ = AddOutput({input.type, {}, 0, 0, 1. / 256}); - } else { - output_ = AddOutput({input.type, {}}); - } - SetBuiltinOp(type, BuiltinOptions_NONE, 0); - BuildInterpreter({GetShape(input_)}); - } - - // A dedicated constructor for SOFTMAX, which does some options. - BaseActivationsOpModel(float softmax_beta, TensorData input) { - input_ = AddInput(input); - if (input.type == TensorType_UINT8) { - output_ = AddOutput({input.type, {}, 0, 0, 1. / 256}); - } else { - output_ = AddOutput({input.type, {}}); - } - SetBuiltinOp(BuiltinOperator_SOFTMAX, BuiltinOptions_SoftmaxOptions, - CreateSoftmaxOptions(builder_, softmax_beta).Union()); - BuildInterpreter({GetShape(input_)}); - } - - BaseActivationsOpModel(BuiltinOperator type, const TensorData &input, - const TensorData &output) { - input_ = AddInput(input); - output_ = AddOutput(output); - SetBuiltinOp(type, BuiltinOptions_NONE, 0); - BuildInterpreter({GetShape(input_)}); - } - - protected: - int input_; - int output_; -}; - -class FloatActivationsOpModel : public BaseActivationsOpModel { - public: - using BaseActivationsOpModel::BaseActivationsOpModel; - - void SetInput(std::initializer_list data) { - PopulateTensor(input_, data); - } - std::vector GetOutput() { return ExtractVector(output_); } -}; - -// Our fixed-point math function implementations have roughly 12 bits of -// accuracy, when specialized to 16-bit fixed-point arithmetic. -// That is purely an implementation compromise, it would have been possible -// to get closer to 16 bits of accuracy but that would be more expensive, -// and not needed for our purposes as ultimately the output is either -// immediately down-quantized to 8 bits, or will typically be at the output -// of the surrounding LSTM cell. -// So we can require roughly 2^-12 accuracy when the output is 16-bit, and -// we can more or less expect the full 2^-8 accuracy when the output is 8-bit. -// -// However, the representable output interval is often [-1, 1] (it has to be -// for tanh, and even for logistic, when we implement it in fixed-point, we -// typically have to do so on such a symmetric interval, e.g. ARM NEON only -// has signed fixed-point arithmetic (SQRDMULH)). As the width of [-1, 1] -// is 2, our representable values are often diluted by a factor of 2, whence -// the factor of 2 below. -const float kQuantizedTolerance = 2 * (1. / 256); -const float kQuantizedToleranceInt16 = 2 * (1. / 4096); - -class QuantizedActivationsOpModel : public BaseActivationsOpModel { - public: - using BaseActivationsOpModel::BaseActivationsOpModel; - - template - void SetInput(std::initializer_list data) { - QuantizeAndPopulate(input_, data); - } - template - - std::vector GetOutput() { - return ExtractVector(output_); - } - template - std::vector GetDequantizedOutput() { - return Dequantize(ExtractVector(output_), GetScale(output_), - GetZeroPoint(output_)); - } -}; - -TEST(FloatActivationsOpTest, Relu) { - FloatActivationsOpModel m(BuiltinOperator_RELU, - /*input=*/{TensorType_FLOAT32, {1, 2, 4, 1}}); - m.SetInput({ - 0, -6, 2, 4, // - 3, -2, 10, 1, // - }); - m.Invoke(); - EXPECT_THAT(m.GetOutput(), ElementsAreArray({ - 0, 0, 2, 4, // - 3, 0, 10, 1, // - })); -} - -TEST(FloatActivationsOpTest, Relu1) { - FloatActivationsOpModel m(BuiltinOperator_RELU_N1_TO_1, - /*input=*/{TensorType_FLOAT32, {1, 2, 4, 1}}); - m.SetInput({ - 0.0, -0.6, 0.2, -0.4, // - 0.3, -2.0, 1.1, -0.1, // - }); - m.Invoke(); - EXPECT_THAT(m.GetOutput(), ElementsAreArray({ - 0.0, -0.6, 0.2, -0.4, // - 0.3, -1.0, 1.0, -0.1, // - })); -} - -TEST(FloatActivationsOpTest, Relu6) { - FloatActivationsOpModel m(BuiltinOperator_RELU6, - /*input=*/{TensorType_FLOAT32, {1, 2, 4, 1}}); - m.SetInput({ - 0, -6, 2, 4, // - 3, -2, 10, 1, // - }); - m.Invoke(); - EXPECT_THAT(m.GetOutput(), ElementsAreArray({ - 0, 0, 2, 4, // - 3, 0, 6, 1, // - })); -} - -TEST(FloatActivationsOpTest, Tanh) { - FloatActivationsOpModel m(BuiltinOperator_TANH, - /*input=*/{TensorType_FLOAT32, {1, 2, 4, 1}}); - m.SetInput({ - 0, -6, 2, 4, // - 3, -2, 10, 1, // - }); - m.Invoke(); - EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear({ - 0, -0.9999877, 0.9640275, 0.999329, // - 0.99505475, -0.9640275, 1, 0.7615941, // - }))); -} - -TEST(QuantizedActivationsOpTest, Tanh) { - const float kMin = -1; - const float kMax = 127.f / 128.f; - QuantizedActivationsOpModel m( - BuiltinOperator_TANH, - /*input=*/{TensorType_UINT8, {1, 2, 4, 1}, 8 * kMin, 8 * kMax}, - /*output=*/{TensorType_UINT8, {1, 2, 4, 1}, kMin, kMax}); - m.SetInput({ - 0, -6, 2, 4, // - -4, -2, 8, 1, // - }); - m.Invoke(); - EXPECT_THAT(m.GetDequantizedOutput(), - ElementsAreArray(ArrayFloatNear( - { - 0.0, -0.999987, 0.964027, 0.999329, // - -0.999329, -0.96402, 0.99999, 0.76159, // - }, - kQuantizedTolerance))); - EXPECT_THAT(m.GetOutput(), - ElementsAreArray({128, 0, 251, 255, 0, 5, 255, 225})); -} - -TEST(QuantizedActivationsOpTest, TanhInt16) { - const float kMin = -1; - const float kMax = 32767.f / 32768.f; - QuantizedActivationsOpModel m( - BuiltinOperator_TANH, - /*input=*/{TensorType_INT16, {1, 2, 4, 1}, 8 * kMin, 8 * kMax}, - /*output=*/{TensorType_INT16, {1, 2, 4, 1}, kMin, kMax}); - m.SetInput({ - 0, -6, 2, 4, // - -4, -2, 8, 1, // - }); - m.Invoke(); - EXPECT_THAT(m.GetDequantizedOutput(), - ElementsAreArray(ArrayFloatNear( - { - 0.0, -0.999987, 0.964027, 0.999329, // - -0.999329, -0.96402, 0.99999, 0.76159, // - }, - kQuantizedToleranceInt16))); -} - -TEST(FloatActivationsOpTest, Sigmoid) { - FloatActivationsOpModel m(BuiltinOperator_LOGISTIC, - /*input=*/{TensorType_FLOAT32, {1, 2, 4, 1}}); - m.SetInput({ - 0, -6, 2, 4, // - 3, -2, 10, 1, // - }); - m.Invoke(); - EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear({ - 0.5, 0.002473, 0.880797, 0.982014, // - 0.952574, 0.119203, 0.999955, 0.731059, // - }))); -} - -TEST(QuantizedActivationsOpTest, Sigmoid) { - QuantizedActivationsOpModel m( - BuiltinOperator_LOGISTIC, - /*input=*/{TensorType_UINT8, {1, 2, 4, 1}, -10, 10}); - m.SetInput({ - 0, -6, 2, 4, // - 3, -2, 10, 1, // - }); - m.Invoke(); - EXPECT_THAT(m.GetDequantizedOutput(), - ElementsAreArray(ArrayFloatNear( - { - 0.5, 0.002473, 0.880797, 0.982014, // - 0.952574, 0.119203, 0.999955, 0.731059, // - }, - kQuantizedTolerance))); - EXPECT_THAT(m.GetOutput(), - ElementsAreArray({128, 1, 227, 251, 244, 32, 255, 188})); -} - -TEST(QuantizedActivationsOpTest, SigmoidInt16) { - const float kMin = -1; - const float kMax = 32767.f / 32768.f; - QuantizedActivationsOpModel m( - BuiltinOperator_LOGISTIC, - /*input=*/{TensorType_INT16, {1, 2, 4, 1}, 8 * kMin, 8 * kMax}, - /*output=*/{TensorType_INT16, {1, 2, 4, 1}, kMin, kMax}); - m.SetInput({ - 0, -6, 2, 4, // - 3, -2, 10, 1, // - }); - m.Invoke(); - EXPECT_THAT(m.GetDequantizedOutput(), - ElementsAreArray(ArrayFloatNear( - { - 0.5, 0.002473, 0.880797, 0.982014, // - 0.952574, 0.119203, 0.999955, 0.731059, // - }, - kQuantizedToleranceInt16))); -} - -TEST(FloatActivationsOpTest, Softmax4D) { - FloatActivationsOpModel m(0.1, - /*input=*/{TensorType_FLOAT32, {1, 2, 1, 4}}); - m.SetInput({ - 0, -6, 2, 4, // depth = 0 - 3, -2, 10, 1, // depth = 1 - }); - m.Invoke(); - EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear({ - .23463, .12877, .28658, .35003, // - .22528, .13664, .45365, .18443, // - }))); - - // Same input, but a different shape. - FloatActivationsOpModel m2(0.1, - /*input=*/{TensorType_FLOAT32, {4, 1, 1, 2}}); - m2.SetInput({ - 0, -6, // - 2, 4, // - 3, -2, // - 10, 1, // - }); - m2.Invoke(); - EXPECT_THAT(m2.GetOutput(), ElementsAreArray(ArrayFloatNear({ - 0.645656, 0.354344, // - 0.450166, 0.549834, // - 0.622459, 0.377541, // - 0.710949, 0.28905, // - }))); -} - -TEST(QuantizedActivationsOpTest, Softmax4D) { - QuantizedActivationsOpModel m( - 0.1, - /*input=*/{TensorType_UINT8, {1, 2, 1, 4}, -10, 10}); - m.SetInput({ - 0, -6, 2, 4, // depth = 0 - 3, -2, 10, 1, // depth = 1 - }); - m.Invoke(); - EXPECT_THAT(m.GetDequantizedOutput(), - ElementsAreArray(ArrayFloatNear( - { - .23463, .12877, .28658, .35003, // - .22528, .13664, .45365, .18443, // - }, - kQuantizedTolerance))); - - // Same input, but a different shape. - QuantizedActivationsOpModel m2( - 0.1, - /*input=*/{TensorType_UINT8, {4, 1, 1, 2}, -10, 10}); - m2.SetInput({ - 0, -6, // - 2, 4, // - 3, -2, // - 10, 1, // - }); - m2.Invoke(); - EXPECT_THAT(m2.GetDequantizedOutput(), - ElementsAreArray(ArrayFloatNear( - { - 0.645656, 0.354344, // - 0.450166, 0.549834, // - 0.622459, 0.377541, // - 0.710949, 0.28905, // - }, - kQuantizedTolerance))); -} - -TEST(FloatActivationsOpTest, Softmax3D) { - FloatActivationsOpModel m(0.1, - /*input=*/{TensorType_FLOAT32, {1, 2, 4}}); - m.SetInput({ - 0, -6, 2, 4, // depth = 0 - 3, -2, 10, 1, // depth = 1 - }); - m.Invoke(); - EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear({ - .23463, .12877, .28658, .35003, // - .22528, .13664, .45365, .18443, // - }))); - - // Same input, but a different shape. - FloatActivationsOpModel m2(0.1, - /*input=*/{TensorType_FLOAT32, {4, 1, 2}}); - m2.SetInput({ - 0, -6, // - 2, 4, // - 3, -2, // - 10, 1, // - }); - m2.Invoke(); - EXPECT_THAT(m2.GetOutput(), ElementsAreArray(ArrayFloatNear({ - 0.645656, 0.354344, // - 0.450166, 0.549834, // - 0.622459, 0.377541, // - 0.710949, 0.28905, // - }))); -} - -TEST(QuantizedActivationsOpTest, Softmax3D) { - QuantizedActivationsOpModel m( - 0.1, - /*input=*/{TensorType_UINT8, {1, 2, 4}, -10, 10}); - m.SetInput({ - 0, -6, 2, 4, // depth = 0 - 3, -2, 10, 1, // depth = 1 - }); - m.Invoke(); - EXPECT_THAT(m.GetDequantizedOutput(), - ElementsAreArray(ArrayFloatNear( - { - .23463, .12877, .28658, .35003, // - .22528, .13664, .45365, .18443, // - }, - kQuantizedTolerance))); - - // Same input, but a different shape. - QuantizedActivationsOpModel m2( - 0.1, - /*input=*/{TensorType_UINT8, {4, 1, 2}, -10, 10}); - m2.SetInput({ - 0, -6, // - 2, 4, // - 3, -2, // - 10, 1, // - }); - m2.Invoke(); - EXPECT_THAT(m2.GetDequantizedOutput(), - ElementsAreArray(ArrayFloatNear( - { - 0.645656, 0.354344, // - 0.450166, 0.549834, // - 0.622459, 0.377541, // - 0.710949, 0.28905, // - }, - kQuantizedTolerance))); -} - -TEST(FloatActivationsOpTest, Softmax1D) { - FloatActivationsOpModel m(0.1, - /*input=*/{TensorType_FLOAT32, {8}}); - m.SetInput({0, -6, 2, 4, 3, -2, 10, 1}); - m.Invoke(); - EXPECT_THAT( - m.GetOutput(), - ElementsAreArray(ArrayFloatNear( - {.09752, .05352, .11911, .14548, .13164, .07984, .26509, .10778}))); -} - -TEST(QuantizedActivationsOpTest, Softmax1D) { - QuantizedActivationsOpModel m(0.1, - /*input=*/{TensorType_UINT8, {8}, -10, 10}); - m.SetInput({0, -6, 2, 4, 3, -2, 10, 1}); - m.Invoke(); - EXPECT_THAT( - m.GetDequantizedOutput(), - ElementsAreArray(ArrayFloatNear({0.09766, 0.05469, 0.12109, 0.14453, - 0.13281, 0.07813, 0.26563, 0.10938}, - kQuantizedTolerance))); -} - -TEST(FloatActivationsOpTest, Softmax2D) { - FloatActivationsOpModel m(0.1, - /*input=*/{TensorType_FLOAT32, {2, 4}}); - m.SetInput({ - 0, -6, 2, 4, // - 3, -2, 10, 1, // - }); - m.Invoke(); - EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear({ - .23463, .12877, .28658, .35003, // - .22528, .13664, .45365, .18443, // - }))); - - // Same input, but a different shape. - FloatActivationsOpModel m2(0.1, - /*input=*/{TensorType_FLOAT32, {4, 2}}); - m2.SetInput({ - 0, -6, // - 2, 4, // - 3, -2, // - 10, 1, // - }); - m2.Invoke(); - EXPECT_THAT(m2.GetOutput(), ElementsAreArray(ArrayFloatNear({ - 0.645656, 0.354344, // - 0.450166, 0.549834, // - 0.622459, 0.377541, // - 0.710949, 0.28905, // - }))); -} - -TEST(QuantizedActivationsOpTest, Softmax2D) { - QuantizedActivationsOpModel m(0.1, - /*input=*/{TensorType_UINT8, {2, 4}, -10, 10}); - m.SetInput({ - 0, -6, 2, 4, // - 3, -2, 10, 1, // - }); - m.Invoke(); - EXPECT_THAT(m.GetDequantizedOutput(), - ElementsAreArray(ArrayFloatNear( - { - .23463, .12877, .28658, .35003, // - .22528, .13664, .45365, .18443, // - }, - kQuantizedTolerance))); - - // Same input, but a different shape. - QuantizedActivationsOpModel m2(0.1, - /*input=*/{TensorType_UINT8, {4, 2}, -10, 10}); - m2.SetInput({ - 0, -6, // - 2, 4, // - 3, -2, // - 10, 1, // - }); - m2.Invoke(); - EXPECT_THAT(m2.GetDequantizedOutput(), - ElementsAreArray(ArrayFloatNear( - { - 0.645656, 0.354344, // - 0.450166, 0.549834, // - 0.622459, 0.377541, // - 0.710949, 0.28905, // - }, - kQuantizedTolerance))); -} - -// This contains the same test values as the Softmax test, but reference answer -// generated via the following snippet of python: -// logits1 = tf.constant([[0, -6, 2, 4],[3, -2, 10, 1]], dtype=tf.float32) -// logits2 = tf.constant([[0,-6],[2,4],[3,-2],[10,1]], dtype=tf.float32) -// lsm1 = tf.nn.log_softmax(logits1) -// lsm2 = tf.nn.log_softmax(logits2) -// with tf.Session() as sess: -// print('lsm1', sess.run(lsm1)) -// print('lsm2', sess.run(lsm2)) - -TEST(FloatActivationsOpTest, LogSoftmax) { - FloatActivationsOpModel m(BuiltinOperator_LOG_SOFTMAX, - /*input=*/{TensorType_FLOAT32, {2, 4}}); - m.SetInput({ - 0, -6, 2, 4, // - 3, -2, 10, 1, // - }); - m.Invoke(); - EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear({ - -4.14297, -10.14297, -2.14297, -.142971, // - -7.00104, -12.00104, -.00104087, -9.00104, // - }))); - - // Same input, but a different shape. - FloatActivationsOpModel m2(BuiltinOperator_LOG_SOFTMAX, - /*input=*/{TensorType_FLOAT32, {4, 2}}); - m2.SetInput({ - 0, -6, // - 2, 4, // - 3, -2, // - 10, 1, // - }); - m2.Invoke(); - EXPECT_THAT(m2.GetOutput(), ElementsAreArray(ArrayFloatNear({ - -.00247565, -6.00247, // - -2.12692, -.126928, // - -.00671534, -5.00671, // - -.000123374, -9.00012, // - }))); -} - -TEST(QuantizedActivationsOpTest, LogSoftmax) { - const float kLogSoftmaxQuantizedTolerance = 16 / 256.0; - QuantizedActivationsOpModel m( - BuiltinOperator_LOG_SOFTMAX, - /*input=*/{TensorType_UINT8, {2, 4}, -10, 10}, - /*output=*/{TensorType_UINT8, {}, 0, 0, 16. / 256, 255}); - m.SetInput({ - 0, -6, 2, 4, // - 3, -2, 10, 1, // - }); - m.Invoke(); - EXPECT_THAT(m.GetDequantizedOutput(), - ElementsAreArray(ArrayFloatNear( - { - -4.14297, -10.14297, -2.14297, -.142971, // - -7.00104, -12.00104, -.00104087, -9.00104, // - }, - kLogSoftmaxQuantizedTolerance))); - EXPECT_THAT(m.GetOutput(), - ElementsAreArray({189, 93, 221, 253, 142, 63, 255, 111})); -} - -class PReluOpModel : public SingleOpModel { - public: - PReluOpModel(const TensorData& input, const TensorData& alpha) { - input_ = AddInput(input); - alpha_ = AddInput(alpha); - output_ = AddOutput(input); - SetBuiltinOp(BuiltinOperator_PRELU, BuiltinOptions_NONE, 0); - BuildInterpreter({GetShape(input_), GetShape(alpha_)}); - } - void SetInput(std::initializer_list data) { - PopulateTensor(input_, data); - } - void SetAlpha(std::initializer_list data) { - PopulateTensor(alpha_, data); - } - std::vector GetOutput() { return ExtractVector(output_); } - - protected: - int input_; - int alpha_; - int output_; -}; - -TEST(FloatActivationsOpTest, PRelu) { - PReluOpModel m({TensorType_FLOAT32, {1, 2, 2, 3}}, - {TensorType_FLOAT32, {1, 1, 3}}); - - m.SetInput({ - 0.0f, 0.0f, 0.0f, // Row 1, Column 1 - 1.0f, 1.0f, 1.0f, // Row 1, Column 2 - -1.0f, -1.0f, -1.0f, // Row 2, Column 1 - -2.0f, -2.0f, -2.0f, // Row 1, Column 2 - }); - m.SetAlpha({0.0f, 1.0f, 2.0f}); - m.Invoke(); - EXPECT_THAT(m.GetOutput(), ElementsAreArray({ - 0.0f, 0.0f, 0.0f, // Row 1, Column 1 - 1.0f, 1.0f, 1.0f, // Row 1, Column 2 - 0.0f, -1.0f, -2.0f, // Row 2, Column 1 - 0.0f, -2.0f, -4.0f, // Row 1, Column 2 - })); -} - -} // namespace -} // namespace tflite - -int main(int argc, char** argv) { - ::tflite::LogToStderr(); - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/contrib/lite/kernels/add_test.cc b/tensorflow/contrib/lite/kernels/add_test.cc deleted file mode 100644 index 261dd36ef0c517bd7880f79948b8ac9682f9bab4..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/add_test.cc +++ /dev/null @@ -1,307 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include -#include "tensorflow/contrib/lite/interpreter.h" -#include "tensorflow/contrib/lite/kernels/register.h" -#include "tensorflow/contrib/lite/kernels/test_util.h" -#include "tensorflow/contrib/lite/model.h" - -namespace tflite { -namespace { - -using ::testing::ElementsAreArray; - -class BaseAddOpModel : public SingleOpModel { - public: - BaseAddOpModel(const TensorData& input1, const TensorData& input2, - const TensorData& output, - ActivationFunctionType activation_type) { - input1_ = AddInput(input1); - input2_ = AddInput(input2); - output_ = AddOutput(output); - SetBuiltinOp(BuiltinOperator_ADD, BuiltinOptions_AddOptions, - CreateAddOptions(builder_, activation_type).Union()); - BuildInterpreter({GetShape(input1_), GetShape(input2_)}); - } - - int input1() { return input1_; } - int input2() { return input2_; } - - protected: - int input1_; - int input2_; - int output_; -}; - -class FloatAddOpModel : public BaseAddOpModel { - public: - using BaseAddOpModel::BaseAddOpModel; - - std::vector GetOutput() { return ExtractVector(output_); } -}; - -class IntegerAddOpModel : public BaseAddOpModel { - public: - using BaseAddOpModel::BaseAddOpModel; - - std::vector GetOutput() { return ExtractVector(output_); } -}; - -class QuantizedAddOpModel : public BaseAddOpModel { - public: - using BaseAddOpModel::BaseAddOpModel; - - std::vector GetDequantizedOutput() { - return Dequantize(ExtractVector(output_), - GetScale(output_), GetZeroPoint(output_)); - } - - std::vector GetDequantizedOutputInt16() { - return Dequantize(ExtractVector(output_), - GetScale(output_), GetZeroPoint(output_)); - } -}; - -// for quantized Add, the error shouldn't exceed 2*step -float GetTolerance(float min, float max) { - float kQuantizedStep = (max - min) / 255.0; - float kQuantizedTolerance = 2.0 * kQuantizedStep; - return kQuantizedTolerance; -} - -float GetToleranceInt16(float min, float max) { - float kQuantizedStep = (max - min) / 32767.f; - float kQuantizedTolerance = 2.0 * kQuantizedStep; - return kQuantizedTolerance; -} - -TEST(FloatAddOpModel, NoActivation) { - FloatAddOpModel m({TensorType_FLOAT32, {1, 2, 2, 1}}, - {TensorType_FLOAT32, {1, 2, 2, 1}}, - {TensorType_FLOAT32, {}}, ActivationFunctionType_NONE); - m.PopulateTensor(m.input1(), {-2.0, 0.2, 0.7, 0.8}); - m.PopulateTensor(m.input2(), {0.1, 0.2, 0.3, 0.5}); - m.Invoke(); - EXPECT_THAT(m.GetOutput(), ElementsAreArray({-1.9, 0.4, 1.0, 1.3})); -} - -TEST(FloatAddOpModel, ActivationRELU_N1_TO_1) { - FloatAddOpModel m( - {TensorType_FLOAT32, {1, 2, 2, 1}}, {TensorType_FLOAT32, {1, 2, 2, 1}}, - {TensorType_FLOAT32, {}}, ActivationFunctionType_RELU_N1_TO_1); - m.PopulateTensor(m.input1(), {-2.0, 0.2, 0.7, 0.8}); - m.PopulateTensor(m.input2(), {0.1, 0.2, 0.3, 0.5}); - m.Invoke(); - EXPECT_THAT(m.GetOutput(), ElementsAreArray({-1.0, 0.4, 1.0, 1.0})); -} - -TEST(FloatAddOpModel, VariousInputShapes) { - std::vector> test_shapes = { - {6}, {2, 3}, {2, 1, 3}, {1, 3, 1, 2}}; - for (int i = 0; i < test_shapes.size(); ++i) { - FloatAddOpModel m({TensorType_FLOAT32, test_shapes[i]}, - {TensorType_FLOAT32, test_shapes[i]}, - {TensorType_FLOAT32, {}}, ActivationFunctionType_NONE); - m.PopulateTensor(m.input1(), {-2.0, 0.2, 0.7, 0.8, 1.1, 2.0}); - m.PopulateTensor(m.input2(), {0.1, 0.2, 0.3, 0.5, 1.1, 0.1}); - m.Invoke(); - EXPECT_THAT(m.GetOutput(), - ElementsAreArray({-1.9, 0.4, 1.0, 1.3, 2.2, 2.1})) - << "With shape number " << i; - } -} - -TEST(FloatAddOpModel, WithBroadcast) { - std::vector> test_shapes = { - {6}, {2, 3}, {2, 1, 3}, {1, 3, 1, 2}}; - for (int i = 0; i < test_shapes.size(); ++i) { - FloatAddOpModel m({TensorType_FLOAT32, test_shapes[i]}, - {TensorType_FLOAT32, {}}, // always a scalar - {TensorType_FLOAT32, {}}, ActivationFunctionType_NONE); - m.PopulateTensor(m.input1(), {-2.0, 0.2, 0.7, 0.8, 1.1, 2.0}); - m.PopulateTensor(m.input2(), {0.1}); - m.Invoke(); - EXPECT_THAT( - m.GetOutput(), - ElementsAreArray(ArrayFloatNear({-1.9, 0.3, 0.8, 0.9, 1.2, 2.1}))) - << "With shape number " << i; - } -} - -TEST(IntegerAddOpModel, NoActivation) { - IntegerAddOpModel m({TensorType_INT32, {1, 2, 2, 1}}, - {TensorType_INT32, {1, 2, 2, 1}}, {TensorType_INT32, {}}, - ActivationFunctionType_NONE); - m.PopulateTensor(m.input1(), {-20, 2, 7, 8}); - m.PopulateTensor(m.input2(), {1, 2, 3, 5}); - m.Invoke(); - EXPECT_THAT(m.GetOutput(), ElementsAreArray({-19, 4, 10, 13})); -} - -TEST(IntegerAddOpModel, ActivationRELU_N1_TO_1) { - IntegerAddOpModel m({TensorType_INT32, {1, 2, 2, 1}}, - {TensorType_INT32, {1, 2, 2, 1}}, {TensorType_INT32, {}}, - ActivationFunctionType_RELU_N1_TO_1); - m.PopulateTensor(m.input1(), {-20, 2, 7, 8}); - m.PopulateTensor(m.input2(), {1, 2, 3, 5}); - m.Invoke(); - EXPECT_THAT(m.GetOutput(), ElementsAreArray({-1, 1, 1, 1})); -} - -TEST(IntegerAddOpModel, VariousInputShapes) { - std::vector> test_shapes = { - {6}, {2, 3}, {2, 1, 3}, {1, 3, 1, 2}}; - for (int i = 0; i < test_shapes.size(); ++i) { - IntegerAddOpModel m({TensorType_INT32, test_shapes[i]}, - {TensorType_INT32, test_shapes[i]}, - {TensorType_INT32, {}}, ActivationFunctionType_NONE); - m.PopulateTensor(m.input1(), {-20, 2, 7, 8, 11, 20}); - m.PopulateTensor(m.input2(), {1, 2, 3, 5, 11, 1}); - m.Invoke(); - EXPECT_THAT(m.GetOutput(), ElementsAreArray({-19, 04, 10, 13, 22, 21})) - << "With shape number " << i; - } -} - -TEST(IntegerAddOpModel, WithBroadcast) { - std::vector> test_shapes = { - {6}, {2, 3}, {2, 1, 3}, {1, 3, 1, 2}}; - for (int i = 0; i < test_shapes.size(); ++i) { - IntegerAddOpModel m({TensorType_INT32, test_shapes[i]}, - {TensorType_INT32, {}}, // always a scalar - {TensorType_INT32, {}}, ActivationFunctionType_NONE); - m.PopulateTensor(m.input1(), {-20, 2, 7, 8, 11, 20}); - m.PopulateTensor(m.input2(), {1}); - m.Invoke(); - EXPECT_THAT(m.GetOutput(), - ElementsAreArray(ArrayFloatNear({-19, 3, 8, 9, 12, 21}))) - << "With shape number " << i; - } -} - -TEST(QuantizedAddOpModel, QuantizedTestsNoActivation) { - float kQuantizedTolerance = GetTolerance(-1.0, 1.0); - std::vector> inputs1 = { - {0.1, 0.2, 0.3, 0.4}, {-0.8, 0.2, 0.4, 0.7}, {-0.8, 0.2, 0.7, 0.3}}; - std::vector> inputs2 = { - {0.6, 0.4, 0.3, 0.1}, {0.6, 0.4, 0.5, -0.8}, {0.6, 0.4, -0.8, 0.5}}; - std::vector> results = { - {0.7, 0.6, 0.6, 0.5}, {-0.2, 0.6, 0.9, -0.1}, {-0.2, 0.6, -0.1, 0.8}}; - for (int i = 0; i < inputs1.size(); ++i) { - QuantizedAddOpModel m({TensorType_UINT8, {1, 2, 2, 1}, -1.0, 1.0}, - {TensorType_UINT8, {1, 2, 2, 1}, -1.0, 1.0}, - {TensorType_UINT8, {}, -1.0, 1.0}, - ActivationFunctionType_NONE); - m.QuantizeAndPopulate(m.input1(), inputs1[i]); - m.QuantizeAndPopulate(m.input2(), inputs2[i]); - m.Invoke(); - EXPECT_THAT(m.GetDequantizedOutput(), ElementsAreArray(ArrayFloatNear( - results[i], kQuantizedTolerance))) - << "With test number " << i; - } -} - -TEST(QuantizedAddOpModel, QuantizedTestsNoActivationInt16) { - const float kMin = -1.f; - const float kMax = 32767.f / 32768.f; - float kQuantizedTolerance = GetToleranceInt16(kMin, kMax); - std::vector> inputs1 = { - {0.1, 0.2, 0.3, 0.4}, {-0.8, 0.2, 0.4, 0.7}, {-0.8, 0.2, 0.7, 0.3}}; - std::vector> inputs2 = { - {0.6, 0.4, 0.3, 0.1}, {0.6, 0.4, 0.5, -0.8}, {0.6, 0.4, -0.8, 0.5}}; - std::vector> results = { - {0.7, 0.6, 0.6, 0.5}, {-0.2, 0.6, 0.9, -0.1}, {-0.2, 0.6, -0.1, 0.8}}; - for (int i = 0; i < inputs1.size(); ++i) { - QuantizedAddOpModel m({TensorType_INT16, {1, 2, 2, 1}, kMin, kMax}, - {TensorType_INT16, {1, 2, 2, 1}, kMin, kMax}, - {TensorType_INT16, {}, kMin, kMax}, - ActivationFunctionType_NONE); - m.QuantizeAndPopulate(m.input1(), inputs1[i]); - m.QuantizeAndPopulate(m.input2(), inputs2[i]); - m.Invoke(); - EXPECT_THAT( - m.GetDequantizedOutputInt16(), - ElementsAreArray(ArrayFloatNear(results[i], kQuantizedTolerance))) - << "With test number " << i; - } -} - -TEST(QuantizedAddOpModel, QuantizedTestsActivationRELU_N1_TO_1) { - float kQuantizedTolerance = GetTolerance(-1.0, 1.0); - std::vector> inputs1 = {{-0.8, 0.2, 0.9, 0.7}, - {-0.8, 0.2, 0.7, 0.3}}; - std::vector> inputs2 = {{0.6, 0.4, 0.9, -0.8}, - {0.6, 0.4, -0.8, 0.5}}; - std::vector> results = {{-0.2, 0.6, 1.0, -0.1}, - {-0.2, 0.6, -0.1, 0.8}}; - for (int i = 0; i < inputs1.size(); ++i) { - QuantizedAddOpModel m({TensorType_UINT8, {1, 2, 2, 1}, -1.0, 1.0}, - {TensorType_UINT8, {1, 2, 2, 1}, -1.0, 1.0}, - {TensorType_UINT8, {}, -1.0, 1.0}, - ActivationFunctionType_RELU_N1_TO_1); - m.QuantizeAndPopulate(m.input1(), inputs1[i]); - m.QuantizeAndPopulate(m.input2(), inputs2[i]); - m.Invoke(); - EXPECT_THAT(m.GetDequantizedOutput(), ElementsAreArray(ArrayFloatNear( - results[i], kQuantizedTolerance))) - << "With test number " << i; - } -} - -TEST(QuantizedAddOpModel, QuantizedVariousInputShapes) { - float kQuantizedTolerance = GetTolerance(-3.0, 3.0); - std::vector> test_shapes = { - {6}, {2, 3}, {2, 1, 3}, {1, 3, 1, 2}}; - for (int i = 0; i < test_shapes.size(); ++i) { - QuantizedAddOpModel m({TensorType_UINT8, test_shapes[i], -3.0, 3.0}, - {TensorType_UINT8, test_shapes[i], -3.0, 3.0}, - {TensorType_UINT8, {}, -3.0, 3.0}, - ActivationFunctionType_NONE); - m.QuantizeAndPopulate(m.input1(), {-2.0, 0.2, 0.7, 0.8, 1.1, 2.0}); - m.QuantizeAndPopulate(m.input2(), {0.1, 0.3, 0.3, 0.5, 1.1, 0.1}); - m.Invoke(); - EXPECT_THAT(m.GetDequantizedOutput(), - ElementsAreArray(ArrayFloatNear({-1.9, 0.5, 1.0, 1.3, 2.2, 2.1}, - kQuantizedTolerance))) - << "With shape number " << i; - } -} - -TEST(QuantizedAddOpModel, QuantizedWithBroadcast) { - float kQuantizedTolerance = GetTolerance(-3.0, 3.0); - std::vector> test_shapes = { - {6}, {2, 3}, {2, 1, 3}, {1, 3, 1, 2}}; - for (int i = 0; i < test_shapes.size(); ++i) { - QuantizedAddOpModel m({TensorType_UINT8, test_shapes[i], -3.0, 3.0}, - {TensorType_UINT8, {}, -3.0, 3.0}, - {TensorType_UINT8, {}, -3.0, 3.0}, - ActivationFunctionType_NONE); - m.QuantizeAndPopulate(m.input1(), {-2.0, 0.2, 0.7, 0.8, 1.1, 2.0}); - m.QuantizeAndPopulate(m.input2(), {0.1}); - m.Invoke(); - EXPECT_THAT(m.GetDequantizedOutput(), - ElementsAreArray(ArrayFloatNear({-1.9, 0.3, 0.8, 0.9, 1.2, 2.1}, - kQuantizedTolerance))) - << "With shape number " << i; - } -} - -} // namespace -} // namespace tflite -int main(int argc, char** argv) { - ::tflite::LogToStderr(); - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/contrib/lite/kernels/batch_to_space_nd_test.cc b/tensorflow/contrib/lite/kernels/batch_to_space_nd_test.cc deleted file mode 100644 index 95b025c1b30cc627cf5858ec17f8ff7c57f7bd95..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/batch_to_space_nd_test.cc +++ /dev/null @@ -1,142 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include -#include "tensorflow/contrib/lite/interpreter.h" -#include "tensorflow/contrib/lite/kernels/register.h" -#include "tensorflow/contrib/lite/kernels/test_util.h" -#include "tensorflow/contrib/lite/model.h" - -namespace tflite { -namespace { - -using ::testing::ElementsAreArray; - -class BatchToSpaceNDOpModel : public SingleOpModel { - public: - void SetInput(std::initializer_list data) { - PopulateTensor(input_, data); - } - - void SetBlockShape(std::initializer_list data) { - PopulateTensor(block_shape_, data); - } - - void SetCrops(std::initializer_list data) { - PopulateTensor(crops_, data); - } - - std::vector GetOutput() { return ExtractVector(output_); } - std::vector GetOutputShape() { return GetTensorShape(output_); } - - protected: - int input_; - int block_shape_; - int crops_; - int output_; -}; - -// Tests case where block_shape and crops are const tensors. -// -// Example usage is as follows: -// BatchToSpaceNDOpConstModel m(input_shape, block_shape, crops); -// m.SetInput(input_data); -// m.Invoke(); -class BatchToSpaceNDOpConstModel : public BatchToSpaceNDOpModel { - public: - BatchToSpaceNDOpConstModel(std::initializer_list input_shape, - std::initializer_list block_shape, - std::initializer_list crops) { - input_ = AddInput(TensorType_FLOAT32); - block_shape_ = AddConstInput(TensorType_INT32, block_shape, {2}); - crops_ = AddConstInput(TensorType_INT32, crops, {2, 2}); - output_ = AddOutput(TensorType_FLOAT32); - - SetBuiltinOp(BuiltinOperator_BATCH_TO_SPACE_ND, - BuiltinOptions_BatchToSpaceNDOptions, - CreateBatchToSpaceNDOptions(builder_).Union()); - BuildInterpreter({input_shape}); - } -}; - -// Tests case where block_shape and crops are non-const tensors. -// -// Example usage is as follows: -// BatchToSpaceNDOpDynamicModel m(input_shape); -// m.SetInput(input_data); -// m.SetBlockShape(block_shape); -// m.SetPaddings(crops); -// m.Invoke(); -class BatchToSpaceNDOpDynamicModel : public BatchToSpaceNDOpModel { - public: - BatchToSpaceNDOpDynamicModel(std::initializer_list input_shape) { - input_ = AddInput(TensorType_FLOAT32); - block_shape_ = AddInput(TensorType_INT32); - crops_ = AddInput(TensorType_INT32); - output_ = AddOutput(TensorType_FLOAT32); - - SetBuiltinOp(BuiltinOperator_BATCH_TO_SPACE_ND, - BuiltinOptions_BatchToSpaceNDOptions, - CreateBatchToSpaceNDOptions(builder_).Union()); - BuildInterpreter({input_shape, {2}, {2, 2}}); - } -}; - -TEST(BatchToSpaceNDOpTest, SimpleConstTest) { - BatchToSpaceNDOpConstModel m({4, 2, 2, 1}, {2, 2}, {0, 0, 0, 0}); - m.SetInput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}); - m.Invoke(); - EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4, 4, 1})); - EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 5, 2, 6, 9, 13, 10, 14, 3, 7, - 4, 8, 11, 15, 12, 16})); -} - -TEST(BatchToSpaceNDOpTest, SimpleDynamicTest) { - BatchToSpaceNDOpDynamicModel m({4, 2, 2, 1}); - m.SetInput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}); - m.SetBlockShape({2, 2}); - m.SetCrops({0, 0, 0, 0}); - m.Invoke(); - EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4, 4, 1})); - EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 5, 2, 6, 9, 13, 10, 14, 3, 7, - 4, 8, 11, 15, 12, 16})); -} - -TEST(BatchToSpaceNDOpTest, InvalidShapeTest) { - EXPECT_DEATH(BatchToSpaceNDOpConstModel({3, 2, 2, 1}, {2, 2}, {0, 0, 0, 0}), - "Cannot allocate tensors"); -} - -TEST(BatchToSpaceNDOpTest, InvalidCropsConstTest) { - EXPECT_DEATH(BatchToSpaceNDOpConstModel({3, 2, 2, 1}, {2, 2}, {0, 0, 0, -1}), - "crops.3. >= 0 was not true."); -} - -TEST(BatchToSpaceNDOpTest, InvalidCropsDynamicTest) { - BatchToSpaceNDOpDynamicModel m({4, 2, 2, 1}); - m.SetInput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}); - m.SetBlockShape({2, 2}); - m.SetCrops({0, 0, -1, 0}); - EXPECT_DEATH(m.Invoke(), "crops.2. >= 0 was not true."); -} - -} // namespace -} // namespace tflite - -int main(int argc, char** argv) { - ::tflite::LogToStderr(); - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/contrib/lite/kernels/bidirectional_sequence_lstm_test.cc b/tensorflow/contrib/lite/kernels/bidirectional_sequence_lstm_test.cc deleted file mode 100644 index db98d6c49d42ac9991d2512de54a60c093cf5799..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/bidirectional_sequence_lstm_test.cc +++ /dev/null @@ -1,1890 +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. -==============================================================================*/ -// Unit test for TFLite Bidirectional LSTM op. - -#include -#include -#include -#include - -#include -#include -#include "tensorflow/contrib/lite/interpreter.h" -#include "tensorflow/contrib/lite/kernels/register.h" -#include "tensorflow/contrib/lite/kernels/test_util.h" -#include "tensorflow/contrib/lite/model.h" -#include "tensorflow/contrib/lite/schema/schema_generated.h" - -namespace tflite { -namespace { - -using ::testing::ElementsAreArray; - -class BidirectionalLSTMOpModel : public SingleOpModel { - public: - BidirectionalLSTMOpModel(int n_batch, int n_input, int n_cell, int n_output, - int sequence_length, bool use_cifg, - bool use_peephole, bool use_projection_weights, - bool use_projection_bias, bool merge_outputs, - float cell_clip, float proj_clip, - bool quantize_weights, - const std::vector>& input_shapes) - : n_batch_(n_batch), - n_input_(n_input), - n_fw_cell_(n_cell), - n_bw_cell_(n_cell), - n_fw_output_(n_output), - n_bw_output_(n_output), - sequence_length_(sequence_length), - quantize_weights_(quantize_weights) { - input_ = AddInput(TensorType_FLOAT32); - const auto weight_type = - quantize_weights_ ? TensorType_UINT8 : TensorType_FLOAT32; - - if (use_cifg) { - fw_input_to_input_weights_ = AddNullInput(); - } else { - fw_input_to_input_weights_ = AddInput(weight_type); - } - - fw_input_to_forget_weights_ = AddInput(weight_type); - fw_input_to_cell_weights_ = AddInput(weight_type); - fw_input_to_output_weights_ = AddInput(weight_type); - - if (use_cifg) { - fw_recurrent_to_input_weights_ = AddNullInput(); - } else { - fw_recurrent_to_input_weights_ = AddInput(weight_type); - } - - fw_recurrent_to_forget_weights_ = AddInput(weight_type); - fw_recurrent_to_cell_weights_ = AddInput(weight_type); - fw_recurrent_to_output_weights_ = AddInput(weight_type); - - if (use_peephole) { - if (use_cifg) { - fw_cell_to_input_weights_ = AddNullInput(); - } else { - fw_cell_to_input_weights_ = AddInput(weight_type); - } - fw_cell_to_forget_weights_ = AddInput(weight_type); - fw_cell_to_output_weights_ = AddInput(weight_type); - } else { - fw_cell_to_input_weights_ = AddNullInput(); - fw_cell_to_forget_weights_ = AddNullInput(); - fw_cell_to_output_weights_ = AddNullInput(); - } - - if (use_cifg) { - fw_input_gate_bias_ = AddNullInput(); - } else { - fw_input_gate_bias_ = AddInput(TensorType_FLOAT32); - } - fw_forget_gate_bias_ = AddInput(TensorType_FLOAT32); - fw_cell_bias_ = AddInput(TensorType_FLOAT32); - fw_output_gate_bias_ = AddInput(TensorType_FLOAT32); - - if (use_projection_weights) { - fw_projection_weights_ = AddInput(TensorType_FLOAT32); - if (use_projection_bias) { - fw_projection_bias_ = AddInput(TensorType_FLOAT32); - } else { - fw_projection_bias_ = AddNullInput(); - } - } else { - fw_projection_weights_ = AddNullInput(); - fw_projection_bias_ = AddNullInput(); - } - - if (use_cifg) { - bw_input_to_input_weights_ = AddNullInput(); - } else { - bw_input_to_input_weights_ = AddInput(weight_type); - } - - bw_input_to_forget_weights_ = AddInput(weight_type); - bw_input_to_cell_weights_ = AddInput(weight_type); - bw_input_to_output_weights_ = AddInput(weight_type); - - if (use_cifg) { - bw_recurrent_to_input_weights_ = AddNullInput(); - } else { - bw_recurrent_to_input_weights_ = AddInput(weight_type); - } - - bw_recurrent_to_forget_weights_ = AddInput(weight_type); - bw_recurrent_to_cell_weights_ = AddInput(weight_type); - bw_recurrent_to_output_weights_ = AddInput(weight_type); - - if (use_peephole) { - if (use_cifg) { - bw_cell_to_input_weights_ = AddNullInput(); - } else { - bw_cell_to_input_weights_ = AddInput(weight_type); - } - bw_cell_to_forget_weights_ = AddInput(weight_type); - bw_cell_to_output_weights_ = AddInput(weight_type); - } else { - bw_cell_to_input_weights_ = AddNullInput(); - bw_cell_to_forget_weights_ = AddNullInput(); - bw_cell_to_output_weights_ = AddNullInput(); - } - - if (use_cifg) { - bw_input_gate_bias_ = AddNullInput(); - } else { - bw_input_gate_bias_ = AddInput(TensorType_FLOAT32); - } - bw_forget_gate_bias_ = AddInput(TensorType_FLOAT32); - bw_cell_bias_ = AddInput(TensorType_FLOAT32); - bw_output_gate_bias_ = AddInput(TensorType_FLOAT32); - - if (use_projection_weights) { - bw_projection_weights_ = AddInput(weight_type); - if (use_projection_bias) { - bw_projection_bias_ = AddInput(TensorType_FLOAT32); - } else { - bw_projection_bias_ = AddNullInput(); - } - } else { - bw_projection_weights_ = AddNullInput(); - bw_projection_bias_ = AddNullInput(); - } - - // Adding the 2 input state tensors. - fw_input_activation_state_ = - AddInput(TensorData{TensorType_FLOAT32, {n_fw_output_ * n_batch_}}, - /*is_variable=*/true); - fw_input_cell_state_ = - AddInput(TensorData{TensorType_FLOAT32, {n_fw_cell_ * n_batch_}}, - /*is_variable=*/true); - - // Adding the 2 input state tensors. - bw_input_activation_state_ = - AddInput(TensorData{TensorType_FLOAT32, {n_bw_output_ * n_batch_}}, - /*is_variable=*/true); - bw_input_cell_state_ = - AddInput(TensorData{TensorType_FLOAT32, {n_bw_cell_ * n_batch_}}, - /*is_variable=*/true); - - fw_output_ = AddOutput(TensorType_FLOAT32); - - if (!merge_outputs) { - bw_output_ = AddOutput(TensorType_FLOAT32); - } - - aux_input_ = AddNullInput(); - fw_aux_input_to_input_weights_ = AddNullInput(); - fw_aux_input_to_forget_weights_ = AddNullInput(); - fw_aux_input_to_cell_weights_ = AddNullInput(); - fw_aux_input_to_output_weights_ = AddNullInput(); - bw_aux_input_to_input_weights_ = AddNullInput(); - bw_aux_input_to_forget_weights_ = AddNullInput(); - bw_aux_input_to_cell_weights_ = AddNullInput(); - bw_aux_input_to_output_weights_ = AddNullInput(); - - SetBuiltinOp(BuiltinOperator_BIDIRECTIONAL_SEQUENCE_LSTM, - BuiltinOptions_BidirectionalSequenceLSTMOptions, - CreateBidirectionalSequenceLSTMOptions( - builder_, ActivationFunctionType_TANH, cell_clip, - proj_clip, merge_outputs) - .Union()); - BuildInterpreter(input_shapes); - } - - void PopulateWeightTensor(int tensor_id, const std::vector& f) { - if (quantize_weights_) { - SymmetricQuantizeAndPopulate(tensor_id, f); - } else { - PopulateTensor(tensor_id, f); - } - } - - // Set weights in forward and backward cells to be the same. - void SetInputToInputWeights(const std::vector& f) { - PopulateWeightTensor(fw_input_to_input_weights_, f); - PopulateWeightTensor(bw_input_to_input_weights_, f); - } - - void SetInputToForgetWeights(const std::vector& f) { - PopulateWeightTensor(fw_input_to_forget_weights_, f); - PopulateWeightTensor(bw_input_to_forget_weights_, f); - } - - void SetInputToCellWeights(const std::vector& f) { - PopulateWeightTensor(fw_input_to_cell_weights_, f); - PopulateWeightTensor(bw_input_to_cell_weights_, f); - } - - void SetInputToOutputWeights(const std::vector& f) { - PopulateWeightTensor(fw_input_to_output_weights_, f); - PopulateWeightTensor(bw_input_to_output_weights_, f); - } - - void SetRecurrentToInputWeights(const std::vector& f) { - PopulateWeightTensor(fw_recurrent_to_input_weights_, f); - PopulateWeightTensor(bw_recurrent_to_input_weights_, f); - } - - void SetRecurrentToForgetWeights(const std::vector& f) { - PopulateWeightTensor(fw_recurrent_to_forget_weights_, f); - PopulateWeightTensor(bw_recurrent_to_forget_weights_, f); - } - - void SetRecurrentToCellWeights(const std::vector& f) { - PopulateWeightTensor(fw_recurrent_to_cell_weights_, f); - PopulateWeightTensor(bw_recurrent_to_cell_weights_, f); - } - - void SetRecurrentToOutputWeights(const std::vector& f) { - PopulateWeightTensor(fw_recurrent_to_output_weights_, f); - PopulateWeightTensor(bw_recurrent_to_output_weights_, f); - } - - void SetCellToInputWeights(const std::vector& f) { - PopulateWeightTensor(fw_cell_to_input_weights_, f); - PopulateWeightTensor(bw_cell_to_input_weights_, f); - } - - void SetCellToForgetWeights(const std::vector& f) { - PopulateWeightTensor(fw_cell_to_forget_weights_, f); - PopulateWeightTensor(bw_cell_to_forget_weights_, f); - } - - void SetCellToOutputWeights(const std::vector& f) { - PopulateWeightTensor(fw_cell_to_output_weights_, f); - PopulateWeightTensor(bw_cell_to_output_weights_, f); - } - - void SetInputGateBias(const std::vector& f) { - PopulateTensor(fw_input_gate_bias_, f); - PopulateTensor(bw_input_gate_bias_, f); - } - - void SetForgetGateBias(const std::vector& f) { - PopulateTensor(fw_forget_gate_bias_, f); - PopulateTensor(bw_forget_gate_bias_, f); - } - - void SetCellBias(const std::vector& f) { - PopulateTensor(fw_cell_bias_, f); - PopulateTensor(bw_cell_bias_, f); - } - - void SetOutputGateBias(const std::vector& f) { - PopulateTensor(fw_output_gate_bias_, f); - PopulateTensor(bw_output_gate_bias_, f); - } - - void SetProjectionWeights(const std::vector& f) { - PopulateWeightTensor(fw_projection_weights_, f); - PopulateWeightTensor(bw_projection_weights_, f); - } - - void SetProjectionBias(const std::vector& f) { - PopulateTensor(fw_projection_bias_, f); - PopulateTensor(bw_projection_bias_, f); - } - - void SetInput(int offset, float* begin, float* end) { - PopulateTensor(input_, offset, begin, end); - } - - std::vector GetFwOutput() { return ExtractVector(fw_output_); } - std::vector GetBwOutput() { return ExtractVector(bw_output_); } - - int num_inputs() { return n_input_; } - int num_fw_outputs() { return n_fw_output_; } - int num_bw_outputs() { return n_bw_output_; } - int num_fw_cells() { return n_fw_cell_; } - int num_bw_cells() { return n_bw_cell_; } - int num_batches() { return n_batch_; } - int sequence_length() { return sequence_length_; } - - private: - int input_; - int fw_input_to_input_weights_; - int fw_input_to_forget_weights_; - int fw_input_to_cell_weights_; - int fw_input_to_output_weights_; - - int fw_recurrent_to_input_weights_; - int fw_recurrent_to_forget_weights_; - int fw_recurrent_to_cell_weights_; - int fw_recurrent_to_output_weights_; - - int fw_cell_to_input_weights_; - int fw_cell_to_forget_weights_; - int fw_cell_to_output_weights_; - - int fw_input_gate_bias_; - int fw_forget_gate_bias_; - int fw_cell_bias_; - int fw_output_gate_bias_; - - int fw_projection_weights_; - int fw_projection_bias_; - - int bw_input_to_input_weights_; - int bw_input_to_forget_weights_; - int bw_input_to_cell_weights_; - int bw_input_to_output_weights_; - - int bw_recurrent_to_input_weights_; - int bw_recurrent_to_forget_weights_; - int bw_recurrent_to_cell_weights_; - int bw_recurrent_to_output_weights_; - - int bw_cell_to_input_weights_; - int bw_cell_to_forget_weights_; - int bw_cell_to_output_weights_; - - int bw_input_gate_bias_; - int bw_forget_gate_bias_; - int bw_cell_bias_; - int bw_output_gate_bias_; - - int bw_projection_weights_; - int bw_projection_bias_; - - int fw_input_activation_state_; - int fw_input_cell_state_; - int bw_input_activation_state_; - int bw_input_cell_state_; - - int fw_output_; - int bw_output_; - - int aux_input_; - int fw_aux_input_to_input_weights_; - int fw_aux_input_to_forget_weights_; - int fw_aux_input_to_cell_weights_; - int fw_aux_input_to_output_weights_; - int bw_aux_input_to_input_weights_; - int bw_aux_input_to_forget_weights_; - int bw_aux_input_to_cell_weights_; - int bw_aux_input_to_output_weights_; - - int n_batch_; - int n_input_; - int n_fw_cell_; - int n_bw_cell_; - int n_fw_output_; - int n_bw_output_; - int sequence_length_; - - bool quantize_weights_; -}; - -// Declare LSTMOpTest as a parameterized test, where the parameter is a boolean -// indicating whether to use quantization or not. -class LSTMOpTest : public ::testing::TestWithParam {}; - -INSTANTIATE_TEST_CASE_P(QuantizationOrNot, LSTMOpTest, ::testing::Bool()); - -TEST_P(LSTMOpTest, BlackBoxTestNoCifgNoPeepholeNoProjectionNoClipping) { - const int n_batch = 1; - const int n_input = 2; - // n_cell and n_output have the same size when there is no projection. - const int n_cell = 4; - const int n_output = 4; - const int sequence_length = 3; - const bool quantize_weights = GetParam(); - - BidirectionalLSTMOpModel lstm( - n_batch, n_input, n_cell, n_output, sequence_length, /*use_cifg=*/false, - /*use_peephole=*/false, /*use_projection_weights=*/false, - /*use_projection_bias=*/false, /*merge_outputs=*/false, /*cell_clip=*/0.0, - /*proj_clip=*/0.0, quantize_weights, - { - {sequence_length, n_batch, n_input}, // input tensor - - // Forward cell - {n_cell, n_input}, // input_to_input_weight tensor - {n_cell, n_input}, // input_to_forget_weight tensor - {n_cell, n_input}, // input_to_cell_weight tensor - {n_cell, n_input}, // input_to_output_weight tensor - - {n_cell, n_output}, // recurrent_to_input_weight tensor - {n_cell, n_output}, // recurrent_to_forget_weight tensor - {n_cell, n_output}, // recurrent_to_cell_weight tensor - {n_cell, n_output}, // recurrent_to_output_weight tensor - - {0}, // cell_to_input_weight tensor - {0}, // cell_to_forget_weight tensor - {0}, // cell_to_output_weight tensor - - {n_cell}, // input_gate_bias tensor - {n_cell}, // forget_gate_bias tensor - {n_cell}, // cell_bias tensor - {n_cell}, // output_gate_bias tensor - - {0, 0}, // projection_weight tensor - {0}, // projection_bias tensor - - // Backward cell - {n_cell, n_input}, // input_to_input_weight tensor - {n_cell, n_input}, // input_to_forget_weight tensor - {n_cell, n_input}, // input_to_cell_weight tensor - {n_cell, n_input}, // input_to_output_weight tensor - - {n_cell, n_output}, // recurrent_to_input_weight tensor - {n_cell, n_output}, // recurrent_to_forget_weight tensor - {n_cell, n_output}, // recurrent_to_cell_weight tensor - {n_cell, n_output}, // recurrent_to_output_weight tensor - - {0}, // cell_to_input_weight tensor - {0}, // cell_to_forget_weight tensor - {0}, // cell_to_output_weight tensor - - {n_cell}, // input_gate_bias tensor - {n_cell}, // forget_gate_bias tensor - {n_cell}, // cell_bias tensor - {n_cell}, // output_gate_bias tensor - - {0, 0}, // projection_weight tensor - {0}, // projection_bias tensor - - {n_batch, n_output}, // activation_state tensor - {n_batch, n_cell}, // cell_state tensor - - {n_batch, n_output}, // activation_state tensor - {n_batch, n_cell}, // cell_state tensor - - {n_batch, sequence_length, 0}, // aux_input tensor - {n_cell, 0}, // aux_fw_input_to_input tensor - {n_cell, 0}, // aux_fw_input_to_forget tensor - {n_cell, 0}, // aux_fw_input_to_cell tensor - {n_cell, 0}, // aux_fw_input_to_output tensor - {n_cell, 0}, // aux_bw_input_to_input tensor - {n_cell, 0}, // aux_bw_input_to_forget tensor - {n_cell, 0}, // aux_bw_input_to_cell tensor - {n_cell, 0}, // aux_bw_input_to_output tensor - }); - - lstm.SetInputToInputWeights({-0.45018822, -0.02338299, -0.0870589, - -0.34550029, 0.04266912, -0.15680569, - -0.34856534, 0.43890524}); - - lstm.SetInputToCellWeights({-0.50013041, 0.1370284, 0.11810488, 0.2013163, - -0.20583314, 0.44344562, 0.22077113, - -0.29909778}); - - lstm.SetInputToForgetWeights({0.09701663, 0.20334584, -0.50592935, - -0.31343272, -0.40032279, 0.44781327, - 0.01387155, -0.35593212}); - - lstm.SetInputToOutputWeights({-0.25065863, -0.28290087, 0.04613829, - 0.40525138, 0.44272184, 0.03897077, -0.1556896, - 0.19487578}); - - lstm.SetInputGateBias({0., 0., 0., 0.}); - - lstm.SetCellBias({0., 0., 0., 0.}); - - lstm.SetForgetGateBias({1., 1., 1., 1.}); - - lstm.SetOutputGateBias({0., 0., 0., 0.}); - - lstm.SetRecurrentToInputWeights( - {-0.0063535, -0.2042388, 0.31454784, -0.35746509, 0.28902304, 0.08183324, - -0.16555229, 0.02286911, -0.13566875, 0.03034258, 0.48091322, - -0.12528998, 0.24077177, -0.51332325, -0.33502164, 0.10629296}); - - lstm.SetRecurrentToCellWeights( - {-0.3407414, 0.24443203, -0.2078532, 0.26320225, 0.05695659, -0.00123841, - -0.4744786, -0.35869038, -0.06418842, -0.13502428, -0.501764, 0.22830659, - -0.46367589, 0.26016325, -0.03894562, -0.16368064}); - - lstm.SetRecurrentToForgetWeights( - {-0.48684245, -0.06655136, 0.42224967, 0.2112639, 0.27654213, 0.20864892, - -0.07646349, 0.45877004, 0.00141793, -0.14609534, 0.36447752, 0.09196436, - 0.28053468, 0.01560611, -0.20127171, -0.01140004}); - - lstm.SetRecurrentToOutputWeights( - {0.43385774, -0.17194885, 0.2718237, 0.09215671, 0.24107647, -0.39835793, - 0.18212086, 0.01301402, 0.48572797, -0.50656658, 0.20047462, -0.20607421, - -0.51818722, -0.15390486, 0.0468148, 0.39922136}); - - // Input should have n_input * sequence_length many values. - static float lstm_input[] = {2., 3., 3., 4., 1., 1.}; - static float lstm_fw_golden_output[] = { - -0.02973187, 0.1229473, 0.20885126, -0.15358765, - -0.03716109, 0.12507336, 0.41193449, -0.20860538, - -0.15053082, 0.09120187, 0.24278517, -0.12222792}; - static float lstm_bw_golden_output[] = { - -0.0806187, 0.139077, 0.400476, -0.197842, -0.0332076, 0.123838, - 0.309777, -0.17621, -0.0490733, 0.0739237, 0.067706, -0.0208124}; - - float* batch0_start = lstm_input; - float* batch0_end = batch0_start + lstm.num_inputs() * lstm.sequence_length(); - - lstm.SetInput(0, batch0_start, batch0_end); - - lstm.Invoke(); - - float* fw_golden_start = lstm_fw_golden_output; - float* fw_golden_end = - fw_golden_start + lstm.num_fw_outputs() * lstm.sequence_length(); - std::vector fw_expected; - fw_expected.insert(fw_expected.end(), fw_golden_start, fw_golden_end); - EXPECT_THAT(lstm.GetFwOutput(), - ElementsAreArray( - ArrayFloatNear(fw_expected, quantize_weights ? 1e-2 : 1e-5))); - - float* bw_golden_start = lstm_bw_golden_output; - float* bw_golden_end = - bw_golden_start + lstm.num_bw_outputs() * lstm.sequence_length(); - std::vector bw_expected; - bw_expected.insert(bw_expected.end(), bw_golden_start, bw_golden_end); - EXPECT_THAT(lstm.GetBwOutput(), - ElementsAreArray( - ArrayFloatNear(bw_expected, quantize_weights ? 1e-2 : 1e-5))); -} - -// Same as the previous test, yet with a single merged output tensor and n_batch -// of 2. -TEST_P(LSTMOpTest, BlackBoxTestMergedOutput) { - const int n_batch = 2; - const int n_input = 2; - // n_cell and n_output have the same size when there is no projection. - const int n_cell = 4; - const int n_output = 4; - const int sequence_length = 3; - const bool quantize_weights = GetParam(); - - BidirectionalLSTMOpModel lstm( - n_batch, n_input, n_cell, n_output, sequence_length, /*use_cifg=*/false, - /*use_peephole=*/false, /*use_projection_weights=*/false, - /*use_projection_bias=*/false, /*merge_outputs=*/true, /*cell_clip=*/0.0, - /*proj_clip=*/0.0, quantize_weights, - { - {sequence_length, n_batch, n_input}, // input tensor - - // Forward cell - {n_cell, n_input}, // input_to_input_weight tensor - {n_cell, n_input}, // input_to_forget_weight tensor - {n_cell, n_input}, // input_to_cell_weight tensor - {n_cell, n_input}, // input_to_output_weight tensor - - {n_cell, n_output}, // recurrent_to_input_weight tensor - {n_cell, n_output}, // recurrent_to_forget_weight tensor - {n_cell, n_output}, // recurrent_to_cell_weight tensor - {n_cell, n_output}, // recurrent_to_output_weight tensor - - {0}, // cell_to_input_weight tensor - {0}, // cell_to_forget_weight tensor - {0}, // cell_to_output_weight tensor - - {n_cell}, // input_gate_bias tensor - {n_cell}, // forget_gate_bias tensor - {n_cell}, // cell_bias tensor - {n_cell}, // output_gate_bias tensor - - {0, 0}, // projection_weight tensor - {0}, // projection_bias tensor - - // Backward cell - {n_cell, n_input}, // input_to_input_weight tensor - {n_cell, n_input}, // input_to_forget_weight tensor - {n_cell, n_input}, // input_to_cell_weight tensor - {n_cell, n_input}, // input_to_output_weight tensor - - {n_cell, n_output}, // recurrent_to_input_weight tensor - {n_cell, n_output}, // recurrent_to_forget_weight tensor - {n_cell, n_output}, // recurrent_to_cell_weight tensor - {n_cell, n_output}, // recurrent_to_output_weight tensor - - {0}, // cell_to_input_weight tensor - {0}, // cell_to_forget_weight tensor - {0}, // cell_to_output_weight tensor - - {n_cell}, // input_gate_bias tensor - {n_cell}, // forget_gate_bias tensor - {n_cell}, // cell_bias tensor - {n_cell}, // output_gate_bias tensor - - {0, 0}, // projection_weight tensor - {0}, // projection_bias tensor - - {n_batch, n_output}, // activation_state tensor - {n_batch, n_cell}, // cell_state tensor - - {n_batch, n_output}, // activation_state tensor - {n_batch, n_cell}, // cell_state tensor - - {n_batch, sequence_length, 0}, // aux_input tensor - {n_cell, 0}, // aux_fw_input_to_input tensor - {n_cell, 0}, // aux_fw_input_to_forget tensor - {n_cell, 0}, // aux_fw_input_to_cell tensor - {n_cell, 0}, // aux_fw_input_to_output tensor - {n_cell, 0}, // aux_bw_input_to_input tensor - {n_cell, 0}, // aux_bw_input_to_forget tensor - {n_cell, 0}, // aux_bw_input_to_cell tensor - {n_cell, 0}, // aux_bw_input_to_output tensor - }); - - lstm.SetInputToInputWeights({-0.45018822, -0.02338299, -0.0870589, - -0.34550029, 0.04266912, -0.15680569, - -0.34856534, 0.43890524}); - - lstm.SetInputToCellWeights({-0.50013041, 0.1370284, 0.11810488, 0.2013163, - -0.20583314, 0.44344562, 0.22077113, - -0.29909778}); - - lstm.SetInputToForgetWeights({0.09701663, 0.20334584, -0.50592935, - -0.31343272, -0.40032279, 0.44781327, - 0.01387155, -0.35593212}); - - lstm.SetInputToOutputWeights({-0.25065863, -0.28290087, 0.04613829, - 0.40525138, 0.44272184, 0.03897077, -0.1556896, - 0.19487578}); - - lstm.SetInputGateBias({0., 0., 0., 0.}); - - lstm.SetCellBias({0., 0., 0., 0.}); - - lstm.SetForgetGateBias({1., 1., 1., 1.}); - - lstm.SetOutputGateBias({0., 0., 0., 0.}); - - lstm.SetRecurrentToInputWeights( - {-0.0063535, -0.2042388, 0.31454784, -0.35746509, 0.28902304, 0.08183324, - -0.16555229, 0.02286911, -0.13566875, 0.03034258, 0.48091322, - -0.12528998, 0.24077177, -0.51332325, -0.33502164, 0.10629296}); - - lstm.SetRecurrentToCellWeights( - {-0.3407414, 0.24443203, -0.2078532, 0.26320225, 0.05695659, -0.00123841, - -0.4744786, -0.35869038, -0.06418842, -0.13502428, -0.501764, 0.22830659, - -0.46367589, 0.26016325, -0.03894562, -0.16368064}); - - lstm.SetRecurrentToForgetWeights( - {-0.48684245, -0.06655136, 0.42224967, 0.2112639, 0.27654213, 0.20864892, - -0.07646349, 0.45877004, 0.00141793, -0.14609534, 0.36447752, 0.09196436, - 0.28053468, 0.01560611, -0.20127171, -0.01140004}); - - lstm.SetRecurrentToOutputWeights( - {0.43385774, -0.17194885, 0.2718237, 0.09215671, 0.24107647, -0.39835793, - 0.18212086, 0.01301402, 0.48572797, -0.50656658, 0.20047462, -0.20607421, - -0.51818722, -0.15390486, 0.0468148, 0.39922136}); - - // Input should have n_input * sequence_length many values. - static float lstm_input[] = {2., 3., 2., 3., 3., 4., 3., 4., 1., 1., 1., 1.}; - static float lstm_fw_golden_output[] = { - -0.02973187, 0.1229473, 0.20885126, -0.15358765, -0.02973187, - 0.1229473, 0.20885126, -0.15358765, -0.03716109, 0.12507336, - 0.41193449, -0.20860538, -0.03716109, 0.12507336, 0.41193449, - -0.20860538, -0.15053082, 0.09120187, 0.24278517, -0.12222792, - -0.15053082, 0.09120187, 0.24278517, -0.12222792}; - static float lstm_bw_golden_output[] = { - -0.0806187, 0.139077, 0.400476, -0.197842, -0.0806187, 0.139077, - 0.400476, -0.197842, -0.0332076, 0.123838, 0.309777, -0.17621, - -0.0332076, 0.123838, 0.309777, -0.17621, -0.0490733, 0.0739237, - 0.067706, -0.0208124, -0.0490733, 0.0739237, 0.067706, -0.0208124}; - - float* batch0_start = lstm_input; - float* batch0_end = batch0_start + lstm.num_inputs() * lstm.num_batches() * - lstm.sequence_length(); - - lstm.SetInput(0, batch0_start, batch0_end); - - lstm.Invoke(); - - std::vector merged_expected; - for (int k = 0; k < lstm.sequence_length() * lstm.num_batches(); k++) { - merged_expected.insert( - merged_expected.end(), - lstm_fw_golden_output + k * lstm.num_fw_outputs(), - lstm_fw_golden_output + (k + 1) * lstm.num_fw_outputs()); - merged_expected.insert( - merged_expected.end(), - lstm_bw_golden_output + k * lstm.num_bw_outputs(), - lstm_bw_golden_output + (k + 1) * lstm.num_bw_outputs()); - } - EXPECT_THAT(lstm.GetFwOutput(), - ElementsAreArray(ArrayFloatNear(merged_expected, - quantize_weights ? 1e-2 : 1e-5))); -} - -TEST(LSTMOpTest, BlackBoxTestNoCifgNoPeepholeNoProjectionNoClippingReverse) { - const int n_batch = 1; - const int n_input = 2; - // n_cell and n_output have the same size when there is no projection. - const int n_cell = 4; - const int n_output = 4; - const int sequence_length = 3; - - BidirectionalLSTMOpModel lstm( - n_batch, n_input, n_cell, n_output, sequence_length, /*use_cifg=*/false, - /*use_peephole=*/false, /*use_projection_weights=*/false, - /*use_projection_bias=*/false, /*merge_outputs=*/false, /*cell_clip=*/0.0, - /*proj_clip=*/0.0, /*quantize_weights=*/false, - { - {sequence_length, n_batch, n_input}, // input tensor - - // Forward cell - {n_cell, n_input}, // input_to_input_weight tensor - {n_cell, n_input}, // input_to_forget_weight tensor - {n_cell, n_input}, // input_to_cell_weight tensor - {n_cell, n_input}, // input_to_output_weight tensor - - {n_cell, n_output}, // recurrent_to_input_weight tensor - {n_cell, n_output}, // recurrent_to_forget_weight tensor - {n_cell, n_output}, // recurrent_to_cell_weight tensor - {n_cell, n_output}, // recurrent_to_output_weight tensor - - {0}, // cell_to_input_weight tensor - {0}, // cell_to_forget_weight tensor - {0}, // cell_to_output_weight tensor - - {n_cell}, // input_gate_bias tensor - {n_cell}, // forget_gate_bias tensor - {n_cell}, // cell_bias tensor - {n_cell}, // output_gate_bias tensor - - {0, 0}, // projection_weight tensor - {0}, // projection_bias tensor - - // Backward cell - {n_cell, n_input}, // input_to_input_weight tensor - {n_cell, n_input}, // input_to_forget_weight tensor - {n_cell, n_input}, // input_to_cell_weight tensor - {n_cell, n_input}, // input_to_output_weight tensor - - {n_cell, n_output}, // recurrent_to_input_weight tensor - {n_cell, n_output}, // recurrent_to_forget_weight tensor - {n_cell, n_output}, // recurrent_to_cell_weight tensor - {n_cell, n_output}, // recurrent_to_output_weight tensor - - {0}, // cell_to_input_weight tensor - {0}, // cell_to_forget_weight tensor - {0}, // cell_to_output_weight tensor - - {n_cell}, // input_gate_bias tensor - {n_cell}, // forget_gate_bias tensor - {n_cell}, // cell_bias tensor - {n_cell}, // output_gate_bias tensor - - {0, 0}, // projection_weight tensor - {0}, // projection_bias tensor - - {n_batch, n_output}, // activation_state tensor - {n_batch, n_cell}, // cell_state tensor - - {n_batch, n_output}, // activation_state tensor - {n_batch, n_cell}, // cell_state tensor - - {n_batch, sequence_length, 0}, // aux_input tensor - {n_cell, 0}, // aux_fw_input_to_input tensor - {n_cell, 0}, // aux_fw_input_to_forget tensor - {n_cell, 0}, // aux_fw_input_to_cell tensor - {n_cell, 0}, // aux_fw_input_to_output tensor - {n_cell, 0}, // aux_bw_input_to_input tensor - {n_cell, 0}, // aux_bw_input_to_forget tensor - {n_cell, 0}, // aux_bw_input_to_cell tensor - {n_cell, 0}, // aux_bw_input_to_output tensor - }); - - lstm.SetInputToInputWeights({-0.45018822, -0.02338299, -0.0870589, - -0.34550029, 0.04266912, -0.15680569, - -0.34856534, 0.43890524}); - - lstm.SetInputToCellWeights({-0.50013041, 0.1370284, 0.11810488, 0.2013163, - -0.20583314, 0.44344562, 0.22077113, - -0.29909778}); - - lstm.SetInputToForgetWeights({0.09701663, 0.20334584, -0.50592935, - -0.31343272, -0.40032279, 0.44781327, - 0.01387155, -0.35593212}); - - lstm.SetInputToOutputWeights({-0.25065863, -0.28290087, 0.04613829, - 0.40525138, 0.44272184, 0.03897077, -0.1556896, - 0.19487578}); - - lstm.SetInputGateBias({0., 0., 0., 0.}); - - lstm.SetCellBias({0., 0., 0., 0.}); - - lstm.SetForgetGateBias({1., 1., 1., 1.}); - - lstm.SetOutputGateBias({0., 0., 0., 0.}); - - lstm.SetRecurrentToInputWeights( - {-0.0063535, -0.2042388, 0.31454784, -0.35746509, 0.28902304, 0.08183324, - -0.16555229, 0.02286911, -0.13566875, 0.03034258, 0.48091322, - -0.12528998, 0.24077177, -0.51332325, -0.33502164, 0.10629296}); - - lstm.SetRecurrentToCellWeights( - {-0.3407414, 0.24443203, -0.2078532, 0.26320225, 0.05695659, -0.00123841, - -0.4744786, -0.35869038, -0.06418842, -0.13502428, -0.501764, 0.22830659, - -0.46367589, 0.26016325, -0.03894562, -0.16368064}); - - lstm.SetRecurrentToForgetWeights( - {-0.48684245, -0.06655136, 0.42224967, 0.2112639, 0.27654213, 0.20864892, - -0.07646349, 0.45877004, 0.00141793, -0.14609534, 0.36447752, 0.09196436, - 0.28053468, 0.01560611, -0.20127171, -0.01140004}); - - lstm.SetRecurrentToOutputWeights( - {0.43385774, -0.17194885, 0.2718237, 0.09215671, 0.24107647, -0.39835793, - 0.18212086, 0.01301402, 0.48572797, -0.50656658, 0.20047462, -0.20607421, - -0.51818722, -0.15390486, 0.0468148, 0.39922136}); - - // Input should have n_input * sequence_length many values. - // Check reversed inputs. - static float lstm_input_reversed[] = {1., 1., 3., 4., 2., 3.}; - static float lstm_fw_golden_output[] = { - -0.02973187, 0.1229473, 0.20885126, -0.15358765, - -0.03716109, 0.12507336, 0.41193449, -0.20860538, - -0.15053082, 0.09120187, 0.24278517, -0.12222792}; - static float lstm_bw_golden_output[] = { - -0.0806187, 0.139077, 0.400476, -0.197842, -0.0332076, 0.123838, - 0.309777, -0.17621, -0.0490733, 0.0739237, 0.067706, -0.0208124}; - - float* batch0_start = lstm_input_reversed; - float* batch0_end = batch0_start + lstm.num_inputs() * lstm.sequence_length(); - - lstm.SetInput(0, batch0_start, batch0_end); - - lstm.Invoke(); - - std::vector fw_expected; - for (int s = 0; s < lstm.sequence_length(); s++) { - float* fw_golden_start = lstm_fw_golden_output + s * lstm.num_fw_outputs(); - float* fw_golden_end = fw_golden_start + lstm.num_fw_outputs(); - fw_expected.insert(fw_expected.begin(), fw_golden_start, fw_golden_end); - } - EXPECT_THAT(lstm.GetBwOutput(), - ElementsAreArray(ArrayFloatNear(fw_expected))); - - std::vector bw_expected; - for (int s = 0; s < lstm.sequence_length(); s++) { - float* bw_golden_start = lstm_bw_golden_output + s * lstm.num_bw_outputs(); - float* bw_golden_end = bw_golden_start + lstm.num_bw_outputs(); - bw_expected.insert(bw_expected.begin(), bw_golden_start, bw_golden_end); - } - EXPECT_THAT(lstm.GetFwOutput(), - ElementsAreArray(ArrayFloatNear(bw_expected))); -} - -TEST(LSTMOpTest, BlackBoxTestWithCifgWithPeepholeNoProjectionNoClipping) { - const int n_batch = 1; - const int n_input = 2; - // n_cell and n_output have the same size when there is no projection. - const int n_cell = 4; - const int n_output = 4; - const int sequence_length = 3; - - BidirectionalLSTMOpModel lstm( - n_batch, n_input, n_cell, n_output, sequence_length, /*use_cifg=*/true, - /*use_peephole=*/true, /*use_projection_weights=*/false, - /*use_projection_bias=*/false, /*merge_outputs=*/false, /*cell_clip=*/0.0, - /*proj_clip=*/0.0, /*quantize_weights=*/false, - { - {sequence_length, n_batch, n_input}, // input tensor - - {0, 0}, // input_to_input_weight tensor - {n_cell, n_input}, // input_to_forget_weight tensor - {n_cell, n_input}, // input_to_cell_weight tensor - {n_cell, n_input}, // input_to_output_weight tensor - - {0, 0}, // recurrent_to_input_weight tensor - {n_cell, n_output}, // recurrent_to_forget_weight tensor - {n_cell, n_output}, // recurrent_to_cell_weight tensor - {n_cell, n_output}, // recurrent_to_output_weight tensor - - {0}, // cell_to_input_weight tensor - {n_cell}, // cell_to_forget_weight tensor - {n_cell}, // cell_to_output_weight tensor - - {0}, // input_gate_bias tensor - {n_cell}, // forget_gate_bias tensor - {n_cell}, // cell_bias tensor - {n_cell}, // output_gate_bias tensor - - {0, 0}, // projection_weight tensor - {0}, // projection_bias tensor - - {0, 0}, // input_to_input_weight tensor - {n_cell, n_input}, // input_to_forget_weight tensor - {n_cell, n_input}, // input_to_cell_weight tensor - {n_cell, n_input}, // input_to_output_weight tensor - - {0, 0}, // recurrent_to_input_weight tensor - {n_cell, n_output}, // recurrent_to_forget_weight tensor - {n_cell, n_output}, // recurrent_to_cell_weight tensor - {n_cell, n_output}, // recurrent_to_output_weight tensor - - {0}, // cell_to_input_weight tensor - {n_cell}, // cell_to_forget_weight tensor - {n_cell}, // cell_to_output_weight tensor - - {0}, // input_gate_bias tensor - {n_cell}, // forget_gate_bias tensor - {n_cell}, // cell_bias tensor - {n_cell}, // output_gate_bias tensor - - {0, 0}, // projection_weight tensor - {0}, // projection_bias tensor - - {n_batch, n_output}, // activation_state tensor - {n_batch, n_cell}, // cell_state tensor - - {n_batch, n_output}, // activation_state tensor - {n_batch, n_cell}, // cell_state tensor - - {n_batch, sequence_length, 0}, // aux_input tensor - {n_cell, 0}, // aux_fw_input_to_input tensor - {n_cell, 0}, // aux_fw_input_to_forget tensor - {n_cell, 0}, // aux_fw_input_to_cell tensor - {n_cell, 0}, // aux_fw_input_to_output tensor - {n_cell, 0}, // aux_bw_input_to_input tensor - {n_cell, 0}, // aux_bw_input_to_forget tensor - {n_cell, 0}, // aux_bw_input_to_cell tensor - {n_cell, 0}, // aux_bw_input_to_output tensor - }); - - lstm.SetInputToCellWeights({-0.49770179, -0.27711356, -0.09624726, 0.05100781, - 0.04717243, 0.48944736, -0.38535351, - -0.17212132}); - - lstm.SetInputToForgetWeights({-0.55291498, -0.42866567, 0.13056988, - -0.3633365, -0.22755712, 0.28253698, 0.24407166, - 0.33826375}); - - lstm.SetInputToOutputWeights({0.10725588, -0.02335852, -0.55932593, - -0.09426838, -0.44257352, 0.54939759, - 0.01533556, 0.42751634}); - - lstm.SetCellBias({0., 0., 0., 0.}); - - lstm.SetForgetGateBias({1., 1., 1., 1.}); - - lstm.SetOutputGateBias({0., 0., 0., 0.}); - - lstm.SetRecurrentToCellWeights( - {0.54066205, -0.32668582, -0.43562764, -0.56094903, 0.42957711, - 0.01841056, -0.32764608, -0.33027974, -0.10826075, 0.20675004, - 0.19069612, -0.03026325, -0.54532051, 0.33003211, 0.44901288, - 0.21193194}); - - lstm.SetRecurrentToForgetWeights( - {-0.13832897, -0.0515101, -0.2359007, -0.16661474, -0.14340827, - 0.36986142, 0.23414481, 0.55899, 0.10798943, -0.41174671, 0.17751795, - -0.34484994, -0.35874045, -0.11352962, 0.27268326, 0.54058349}); - - lstm.SetRecurrentToOutputWeights( - {0.41613156, 0.42610586, -0.16495961, -0.5663873, 0.30579174, -0.05115908, - -0.33941799, 0.23364776, 0.11178309, 0.09481031, -0.26424935, 0.46261835, - 0.50248802, 0.26114327, -0.43736315, 0.33149987}); - - lstm.SetCellToForgetWeights( - {0.47485286, -0.51955009, -0.24458408, 0.31544167}); - lstm.SetCellToOutputWeights( - {-0.17135078, 0.82760304, 0.85573703, -0.77109635}); - - static float lstm_input[] = {2., 3., 3., 4., 1., 1.}; - static float lstm_fw_golden_output[] = { - -0.36444446, -0.00352185, 0.12886585, -0.05163646, - -0.42312205, -0.01218222, 0.24201041, -0.08124574, - -0.358325, -0.04621704, 0.21641694, -0.06471302}; - static float lstm_bw_golden_output[] = { - -0.401685, -0.0232794, 0.288642, -0.123074, -0.42915, -0.00871577, - 0.20912, -0.103567, -0.166398, -0.00486649, 0.0697471, -0.0537578}; - - float* batch0_start = lstm_input; - float* batch0_end = batch0_start + lstm.num_inputs() * lstm.sequence_length(); - - lstm.SetInput(0, batch0_start, batch0_end); - - lstm.Invoke(); - - float* fw_golden_start = lstm_fw_golden_output; - float* fw_golden_end = - fw_golden_start + lstm.num_fw_outputs() * lstm.sequence_length(); - std::vector fw_expected; - fw_expected.insert(fw_expected.end(), fw_golden_start, fw_golden_end); - EXPECT_THAT(lstm.GetFwOutput(), - ElementsAreArray(ArrayFloatNear(fw_expected))); - - float* bw_golden_start = lstm_bw_golden_output; - float* bw_golden_end = - bw_golden_start + lstm.num_bw_outputs() * lstm.sequence_length(); - std::vector bw_expected; - bw_expected.insert(bw_expected.end(), bw_golden_start, bw_golden_end); - EXPECT_THAT(lstm.GetBwOutput(), - ElementsAreArray(ArrayFloatNear(bw_expected))); -} - -TEST(LSTMOpTest, - BlackBoxTestWithCifgWithPeepholeNoProjectionNoClippingReversed) { - const int n_batch = 1; - const int n_input = 2; - // n_cell and n_output have the same size when there is no projection. - const int n_cell = 4; - const int n_output = 4; - const int sequence_length = 3; - - BidirectionalLSTMOpModel lstm( - n_batch, n_input, n_cell, n_output, sequence_length, /*use_cifg=*/true, - /*use_peephole=*/true, /*use_projection_weights=*/false, - /*use_projection_bias=*/false, /*merge_outputs=*/false, /*cell_clip=*/0.0, - /*proj_clip=*/0.0, /*quantize_weights=*/false, - { - {sequence_length, n_batch, n_input}, // input tensor - - {0, 0}, // input_to_input_weight tensor - {n_cell, n_input}, // input_to_forget_weight tensor - {n_cell, n_input}, // input_to_cell_weight tensor - {n_cell, n_input}, // input_to_output_weight tensor - - {0, 0}, // recurrent_to_input_weight tensor - {n_cell, n_output}, // recurrent_to_forget_weight tensor - {n_cell, n_output}, // recurrent_to_cell_weight tensor - {n_cell, n_output}, // recurrent_to_output_weight tensor - - {0}, // cell_to_input_weight tensor - {n_cell}, // cell_to_forget_weight tensor - {n_cell}, // cell_to_output_weight tensor - - {0}, // input_gate_bias tensor - {n_cell}, // forget_gate_bias tensor - {n_cell}, // cell_bias tensor - {n_cell}, // output_gate_bias tensor - - {0, 0}, // projection_weight tensor - {0}, // projection_bias tensor - - {0, 0}, // input_to_input_weight tensor - {n_cell, n_input}, // input_to_forget_weight tensor - {n_cell, n_input}, // input_to_cell_weight tensor - {n_cell, n_input}, // input_to_output_weight tensor - - {0, 0}, // recurrent_to_input_weight tensor - {n_cell, n_output}, // recurrent_to_forget_weight tensor - {n_cell, n_output}, // recurrent_to_cell_weight tensor - {n_cell, n_output}, // recurrent_to_output_weight tensor - - {0}, // cell_to_input_weight tensor - {n_cell}, // cell_to_forget_weight tensor - {n_cell}, // cell_to_output_weight tensor - - {0}, // input_gate_bias tensor - {n_cell}, // forget_gate_bias tensor - {n_cell}, // cell_bias tensor - {n_cell}, // output_gate_bias tensor - - {0, 0}, // projection_weight tensor - {0}, // projection_bias tensor - - {n_batch, n_output}, // activation_state tensor - {n_batch, n_cell}, // cell_state tensor - - {n_batch, n_output}, // activation_state tensor - {n_batch, n_cell}, // cell_state tensor - - {n_batch, sequence_length, 0}, // aux_input tensor - {n_cell, 0}, // aux_fw_input_to_input tensor - {n_cell, 0}, // aux_fw_input_to_forget tensor - {n_cell, 0}, // aux_fw_input_to_cell tensor - {n_cell, 0}, // aux_fw_input_to_output tensor - {n_cell, 0}, // aux_bw_input_to_input tensor - {n_cell, 0}, // aux_bw_input_to_forget tensor - {n_cell, 0}, // aux_bw_input_to_cell tensor - {n_cell, 0}, // aux_bw_input_to_output tensor - }); - - lstm.SetInputToCellWeights({-0.49770179, -0.27711356, -0.09624726, 0.05100781, - 0.04717243, 0.48944736, -0.38535351, - -0.17212132}); - - lstm.SetInputToForgetWeights({-0.55291498, -0.42866567, 0.13056988, - -0.3633365, -0.22755712, 0.28253698, 0.24407166, - 0.33826375}); - - lstm.SetInputToOutputWeights({0.10725588, -0.02335852, -0.55932593, - -0.09426838, -0.44257352, 0.54939759, - 0.01533556, 0.42751634}); - - lstm.SetCellBias({0., 0., 0., 0.}); - - lstm.SetForgetGateBias({1., 1., 1., 1.}); - - lstm.SetOutputGateBias({0., 0., 0., 0.}); - - lstm.SetRecurrentToCellWeights( - {0.54066205, -0.32668582, -0.43562764, -0.56094903, 0.42957711, - 0.01841056, -0.32764608, -0.33027974, -0.10826075, 0.20675004, - 0.19069612, -0.03026325, -0.54532051, 0.33003211, 0.44901288, - 0.21193194}); - - lstm.SetRecurrentToForgetWeights( - {-0.13832897, -0.0515101, -0.2359007, -0.16661474, -0.14340827, - 0.36986142, 0.23414481, 0.55899, 0.10798943, -0.41174671, 0.17751795, - -0.34484994, -0.35874045, -0.11352962, 0.27268326, 0.54058349}); - - lstm.SetRecurrentToOutputWeights( - {0.41613156, 0.42610586, -0.16495961, -0.5663873, 0.30579174, -0.05115908, - -0.33941799, 0.23364776, 0.11178309, 0.09481031, -0.26424935, 0.46261835, - 0.50248802, 0.26114327, -0.43736315, 0.33149987}); - - lstm.SetCellToForgetWeights( - {0.47485286, -0.51955009, -0.24458408, 0.31544167}); - lstm.SetCellToOutputWeights( - {-0.17135078, 0.82760304, 0.85573703, -0.77109635}); - - static float lstm_input_reversed[] = {1., 1., 3., 4., 2., 3.}; - static float lstm_fw_golden_output[] = { - -0.36444446, -0.00352185, 0.12886585, -0.05163646, - -0.42312205, -0.01218222, 0.24201041, -0.08124574, - -0.358325, -0.04621704, 0.21641694, -0.06471302}; - static float lstm_bw_golden_output[] = { - -0.401685, -0.0232794, 0.288642, -0.123074, -0.42915, -0.00871577, - 0.20912, -0.103567, -0.166398, -0.00486649, 0.0697471, -0.0537578}; - - float* batch0_start = lstm_input_reversed; - float* batch0_end = batch0_start + lstm.num_inputs() * lstm.sequence_length(); - - lstm.SetInput(0, batch0_start, batch0_end); - - lstm.Invoke(); - - std::vector fw_expected; - for (int s = 0; s < lstm.sequence_length(); s++) { - float* fw_golden_start = lstm_fw_golden_output + s * lstm.num_fw_outputs(); - float* fw_golden_end = fw_golden_start + lstm.num_fw_outputs(); - fw_expected.insert(fw_expected.begin(), fw_golden_start, fw_golden_end); - } - EXPECT_THAT(lstm.GetBwOutput(), - ElementsAreArray(ArrayFloatNear(fw_expected))); - - std::vector bw_expected; - for (int s = 0; s < lstm.sequence_length(); s++) { - float* bw_golden_start = lstm_bw_golden_output + s * lstm.num_bw_outputs(); - float* bw_golden_end = bw_golden_start + lstm.num_bw_outputs(); - bw_expected.insert(bw_expected.begin(), bw_golden_start, bw_golden_end); - } - EXPECT_THAT(lstm.GetFwOutput(), - ElementsAreArray(ArrayFloatNear(bw_expected))); -} - -TEST(LSTMOpTest, BlackBoxTestWithPeepholeWithProjectionNoClipping) { - const int n_batch = 2; - const int n_input = 5; - const int n_cell = 20; - const int n_output = 16; - const int sequence_length = 4; - - BidirectionalLSTMOpModel lstm( - n_batch, n_input, n_cell, n_output, sequence_length, /*use_cifg=*/false, - /*use_peephole=*/true, /*use_projection_weights=*/true, - /*use_projection_bias=*/false, /*merge_outputs=*/false, /*cell_clip=*/0.0, - /*proj_clip=*/0.0, /*quantize_weights=*/false, - { - {sequence_length, n_batch, n_input}, // input tensor - - {n_cell, n_input}, // input_to_input_weight tensor - {n_cell, n_input}, // input_to_forget_weight tensor - {n_cell, n_input}, // input_to_cell_weight tensor - {n_cell, n_input}, // input_to_output_weight tensor - - {n_cell, n_output}, // recurrent_to_input_weight tensor - {n_cell, n_output}, // recurrent_to_forget_weight tensor - {n_cell, n_output}, // recurrent_to_cell_weight tensor - {n_cell, n_output}, // recurrent_to_output_weight tensor - - {n_cell}, // cell_to_input_weight tensor - {n_cell}, // cell_to_forget_weight tensor - {n_cell}, // cell_to_output_weight tensor - - {n_cell}, // input_gate_bias tensor - {n_cell}, // forget_gate_bias tensor - {n_cell}, // cell_bias tensor - {n_cell}, // output_gate_bias tensor - - {n_output, n_cell}, // projection_weight tensor - {0}, // projection_bias tensor - - {n_cell, n_input}, // input_to_input_weight tensor - {n_cell, n_input}, // input_to_forget_weight tensor - {n_cell, n_input}, // input_to_cell_weight tensor - {n_cell, n_input}, // input_to_output_weight tensor - - {n_cell, n_output}, // recurrent_to_input_weight tensor - {n_cell, n_output}, // recurrent_to_forget_weight tensor - {n_cell, n_output}, // recurrent_to_cell_weight tensor - {n_cell, n_output}, // recurrent_to_output_weight tensor - - {n_cell}, // cell_to_input_weight tensor - {n_cell}, // cell_to_forget_weight tensor - {n_cell}, // cell_to_output_weight tensor - - {n_cell}, // input_gate_bias tensor - {n_cell}, // forget_gate_bias tensor - {n_cell}, // cell_bias tensor - {n_cell}, // output_gate_bias tensor - - {n_output, n_cell}, // projection_weight tensor - {0}, // projection_bias tensor - - {n_batch, n_output}, // activation_state tensor - {n_batch, n_cell}, // cell_state tensor - - {n_batch, n_output}, // activation_state tensor - {n_batch, n_cell}, // cell_state tensor - - {n_batch, sequence_length, 0}, // aux_input tensor - {n_cell, 0}, // aux_fw_input_to_input tensor - {n_cell, 0}, // aux_fw_input_to_forget tensor - {n_cell, 0}, // aux_fw_input_to_cell tensor - {n_cell, 0}, // aux_fw_input_to_output tensor - {n_cell, 0}, // aux_bw_input_to_input tensor - {n_cell, 0}, // aux_bw_input_to_forget tensor - {n_cell, 0}, // aux_bw_input_to_cell tensor - {n_cell, 0}, // aux_bw_input_to_output tensor - }); - - lstm.SetInputToInputWeights( - {0.021393683, 0.06124551, 0.046905167, -0.014657677, -0.03149463, - 0.09171803, 0.14647801, 0.10797193, -0.0057968358, 0.0019193048, - -0.2726754, 0.10154029, -0.018539885, 0.080349885, -0.10262385, - -0.022599787, -0.09121155, -0.008675967, -0.045206103, -0.0821282, - -0.008045952, 0.015478081, 0.055217247, 0.038719587, 0.044153627, - -0.06453243, 0.05031825, -0.046935108, -0.008164439, 0.014574226, - -0.1671009, -0.15519552, -0.16819797, -0.13971269, -0.11953059, - 0.25005487, -0.22790983, 0.009855087, -0.028140958, -0.11200698, - 0.11295408, -0.0035217577, 0.054485075, 0.05184695, 0.064711206, - 0.10989193, 0.11674786, 0.03490607, 0.07727357, 0.11390585, - -0.1863375, -0.1034451, -0.13945189, -0.049401227, -0.18767063, - 0.042483903, 0.14233552, 0.13832581, 0.18350165, 0.14545603, - -0.028545704, 0.024939531, 0.050929718, 0.0076203286, -0.0029723682, - -0.042484224, -0.11827596, -0.09171104, -0.10808628, -0.16327988, - -0.2273378, -0.0993647, -0.017155107, 0.0023917493, 0.049272764, - 0.0038534778, 0.054764505, 0.089753784, 0.06947234, 0.08014476, - -0.04544234, -0.0497073, -0.07135631, -0.048929106, -0.004042012, - -0.009284026, 0.018042054, 0.0036860977, -0.07427302, -0.11434604, - -0.018995456, 0.031487543, 0.012834908, 0.019977754, 0.044256654, - -0.39292613, -0.18519334, -0.11651281, -0.06809892, 0.011373677}); - - lstm.SetInputToForgetWeights( - {-0.0018401089, -0.004852237, 0.03698424, 0.014181704, 0.028273236, - -0.016726194, -0.05249759, -0.10204261, 0.00861066, -0.040979505, - -0.009899187, 0.01923892, -0.028177269, -0.08535103, -0.14585495, - 0.10662567, -0.01909731, -0.017883534, -0.0047269356, -0.045103323, - 0.0030784295, 0.076784775, 0.07463696, 0.094531395, 0.0814421, - -0.12257899, -0.033945758, -0.031303465, 0.045630626, 0.06843887, - -0.13492945, -0.012480007, -0.0811829, -0.07224499, -0.09628791, - 0.045100946, 0.0012300825, 0.013964662, 0.099372394, 0.02543059, - 0.06958324, 0.034257296, 0.0482646, 0.06267997, 0.052625068, - 0.12784666, 0.07077897, 0.025725935, 0.04165009, 0.07241905, - 0.018668644, -0.037377294, -0.06277783, -0.08833636, -0.040120605, - -0.011405586, -0.007808335, -0.010301386, -0.005102167, 0.027717464, - 0.05483423, 0.11449111, 0.11289652, 0.10939839, 0.13396506, - -0.08402166, -0.01901462, -0.044678304, -0.07720565, 0.014350063, - -0.11757958, -0.0652038, -0.08185733, -0.076754324, -0.092614375, - 0.10405491, 0.052960336, 0.035755895, 0.035839386, -0.012540553, - 0.036881298, 0.02913376, 0.03420159, 0.05448447, -0.054523353, - 0.02582715, 0.02327355, -0.011857179, -0.0011980024, -0.034641717, - -0.026125094, -0.17582615, -0.15923657, -0.27486774, -0.0006143371, - 0.0001771948, -8.470171e-05, 0.02651807, 0.045790765, 0.06956496}); - - lstm.SetInputToCellWeights( - {-0.04580283, -0.09549462, -0.032418985, -0.06454633, - -0.043528453, 0.043018587, -0.049152344, -0.12418144, - -0.078985475, -0.07596889, 0.019484362, -0.11434962, - -0.0074034138, -0.06314844, -0.092981495, 0.0062155537, - -0.025034338, -0.0028890965, 0.048929527, 0.06235075, - 0.10665918, -0.032036792, -0.08505916, -0.10843358, - -0.13002433, -0.036816437, -0.02130134, -0.016518239, - 0.0047691227, -0.0025825808, 0.066017866, 0.029991534, - -0.10652836, -0.1037554, -0.13056071, -0.03266643, - -0.033702414, -0.006473424, -0.04611692, 0.014419339, - -0.025174323, 0.0396852, 0.081777506, 0.06157468, - 0.10210095, -0.009658194, 0.046511717, 0.03603906, - 0.0069369148, 0.015960095, -0.06507666, 0.09551598, - 0.053568836, 0.06408714, 0.12835667, -0.008714329, - -0.20211966, -0.12093674, 0.029450472, 0.2849013, - -0.029227901, 0.1164364, -0.08560263, 0.09941786, - -0.036999565, -0.028842626, -0.0033637602, -0.017012902, - -0.09720865, -0.11193351, -0.029155117, -0.017936034, - -0.009768936, -0.04223324, -0.036159635, 0.06505112, - -0.021742892, -0.023377212, -0.07221364, -0.06430552, - 0.05453865, 0.091149814, 0.06387331, 0.007518393, - 0.055960953, 0.069779344, 0.046411168, 0.10509911, - 0.07463894, 0.0075130584, 0.012850982, 0.04555431, - 0.056955688, 0.06555285, 0.050801456, -0.009862683, - 0.00826772, -0.026555609, -0.0073611983, -0.0014897042}); - - lstm.SetInputToOutputWeights( - {-0.0998932, -0.07201956, -0.052803773, -0.15629593, -0.15001918, - -0.07650751, 0.02359855, -0.075155355, -0.08037709, -0.15093534, - 0.029517552, -0.04751393, 0.010350531, -0.02664851, -0.016839722, - -0.023121163, 0.0077019283, 0.012851257, -0.05040649, -0.0129761, - -0.021737747, -0.038305793, -0.06870586, -0.01481247, -0.001285394, - 0.10124236, 0.083122835, 0.053313006, -0.062235646, -0.075637154, - -0.027833903, 0.029774971, 0.1130802, 0.09218906, 0.09506135, - -0.086665764, -0.037162706, -0.038880914, -0.035832845, -0.014481564, - -0.09825003, -0.12048569, -0.097665586, -0.05287633, -0.0964047, - -0.11366429, 0.035777505, 0.13568819, 0.052451383, 0.050649304, - 0.05798951, -0.021852335, -0.099848844, 0.014740475, -0.078897946, - 0.04974699, 0.014160473, 0.06973932, 0.04964942, 0.033364646, - 0.08190124, 0.025535367, 0.050893165, 0.048514254, 0.06945813, - -0.078907564, -0.06707616, -0.11844508, -0.09986688, -0.07509403, - 0.06263226, 0.14925587, 0.20188436, 0.12098451, 0.14639415, - 0.0015017595, -0.014267382, -0.03417257, 0.012711468, 0.0028300495, - -0.024758482, -0.05098548, -0.0821182, 0.014225672, 0.021544158, - 0.08949725, 0.07505268, -0.0020780868, 0.04908258, 0.06476295, - -0.022907063, 0.027562456, 0.040185735, 0.019567577, -0.015598739, - -0.049097303, -0.017121866, -0.083368234, -0.02332002, -0.0840956}); - - lstm.SetInputGateBias( - {0.02234832, 0.14757581, 0.18176508, 0.10380666, 0.053110216, - -0.06928846, -0.13942584, -0.11816189, 0.19483899, 0.03652339, - -0.10250295, 0.036714908, -0.18426876, 0.036065217, 0.21810818, - 0.02383196, -0.043370757, 0.08690144, -0.04444982, 0.00030581196}); - - lstm.SetForgetGateBias({0.035185695, -0.042891346, -0.03032477, 0.23027696, - 0.11098921, 0.15378423, 0.09263801, 0.09790885, - 0.09508917, 0.061199076, 0.07665568, -0.015443159, - -0.03499149, 0.046190713, 0.08895977, 0.10899629, - 0.40694186, 0.06030037, 0.012413437, -0.06108739}); - - lstm.SetCellBias({-0.024379363, 0.0055531194, 0.23377132, 0.033463873, - -0.1483596, -0.10639995, -0.091433935, 0.058573797, - -0.06809782, -0.07889636, -0.043246906, -0.09829136, - -0.4279842, 0.034901652, 0.18797937, 0.0075234566, - 0.016178843, 0.1749513, 0.13975595, 0.92058027}); - - lstm.SetOutputGateBias( - {0.046159424, -0.0012809046, 0.03563469, 0.12648113, 0.027195795, - 0.35373217, -0.018957434, 0.008907322, -0.0762701, 0.12018895, - 0.04216877, 0.0022856654, 0.040952638, 0.3147856, 0.08225149, - -0.057416286, -0.14995944, -0.008040261, 0.13208859, 0.029760877}); - - lstm.SetRecurrentToInputWeights( - {-0.001374326, -0.078856036, 0.10672688, 0.029162422, - -0.11585556, 0.02557986, -0.13446963, -0.035785314, - -0.01244275, 0.025961924, -0.02337298, -0.044228926, - -0.055839065, -0.046598054, -0.010546039, -0.06900766, - 0.027239809, 0.022582639, -0.013296484, -0.05459212, - 0.08981, -0.045407712, 0.08682226, -0.06867011, - -0.14390695, -0.02916037, 0.000996957, 0.091420636, - 0.14283475, -0.07390571, -0.06402044, 0.062524505, - -0.093129106, 0.04860203, -0.08364217, -0.08119002, - 0.009352075, 0.22920375, 0.0016303885, 0.11583097, - -0.13732095, 0.012405723, -0.07551853, 0.06343048, - 0.12162708, -0.031923793, -0.014335606, 0.01790974, - -0.10650317, -0.0724401, 0.08554849, -0.05727212, - 0.06556731, -0.042729504, -0.043227166, 0.011683251, - -0.013082158, -0.029302018, -0.010899579, -0.062036745, - -0.022509435, -0.00964907, -0.01567329, 0.04260106, - -0.07787477, -0.11576462, 0.017356863, 0.048673786, - -0.017577527, -0.05527947, -0.082487635, -0.040137455, - -0.10820036, -0.04666372, 0.022746278, -0.07851417, - 0.01068115, 0.032956902, 0.022433773, 0.0026891115, - 0.08944216, -0.0685835, 0.010513544, 0.07228705, - 0.02032331, -0.059686817, -0.0005566496, -0.086984694, - 0.040414046, -0.1380399, 0.094208956, -0.05722982, - 0.012092817, -0.04989123, -0.086576, -0.003399834, - -0.04696032, -0.045747425, 0.10091314, 0.048676282, - -0.029037097, 0.031399418, -0.0040285117, 0.047237843, - 0.09504992, 0.041799378, -0.049185462, -0.031518843, - -0.10516937, 0.026374253, 0.10058866, -0.0033195973, - -0.041975245, 0.0073591834, 0.0033782164, -0.004325073, - -0.10167381, 0.042500053, -0.01447153, 0.06464186, - -0.017142897, 0.03312627, 0.009205989, 0.024138335, - -0.011337001, 0.035530265, -0.010912711, 0.0706555, - -0.005894094, 0.051841937, -0.1401738, -0.02351249, - 0.0365468, 0.07590991, 0.08838724, 0.021681072, - -0.10086113, 0.019608743, -0.06195883, 0.077335775, - 0.023646897, -0.095322326, 0.02233014, 0.09756986, - -0.048691444, -0.009579111, 0.07595467, 0.11480546, - -0.09801813, 0.019894179, 0.08502348, 0.004032281, - 0.037211012, 0.068537936, -0.048005626, -0.091520436, - -0.028379958, -0.01556313, 0.06554592, -0.045599163, - -0.01672207, -0.020169014, -0.011877351, -0.20212261, - 0.010889619, 0.0047078193, 0.038385306, 0.08540671, - -0.017140968, -0.0035865551, 0.016678626, 0.005633034, - 0.015963363, 0.00871737, 0.060130805, 0.028611384, - 0.10109069, -0.015060172, -0.07894427, 0.06401885, - 0.011584063, -0.024466386, 0.0047652307, -0.09041358, - 0.030737216, -0.0046374933, 0.14215417, -0.11823516, - 0.019899689, 0.006106124, -0.027092824, 0.0786356, - 0.05052217, -0.058925, -0.011402121, -0.024987547, - -0.0013661642, -0.06832946, -0.015667673, -0.1083353, - -0.00096863037, -0.06988685, -0.053350925, -0.027275559, - -0.033664223, -0.07978348, -0.025200296, -0.017207067, - -0.058403496, -0.055697463, 0.005798788, 0.12965427, - -0.062582195, 0.0013350133, -0.10482091, 0.0379771, - 0.072521195, -0.0029455067, -0.13797039, -0.03628521, - 0.013806405, -0.017858358, -0.01008298, -0.07700066, - -0.017081132, 0.019358726, 0.0027079724, 0.004635139, - 0.062634714, -0.02338735, -0.039547626, -0.02050681, - 0.03385117, -0.083611414, 0.002862572, -0.09421313, - 0.058618143, -0.08598433, 0.00972939, 0.023867095, - -0.053934585, -0.023203006, 0.07452513, -0.048767887, - -0.07314807, -0.056307215, -0.10433547, -0.06440842, - 0.04328182, 0.04389765, -0.020006588, -0.09076438, - -0.11652589, -0.021705797, 0.03345259, -0.010329105, - -0.025767034, 0.013057034, -0.07316461, -0.10145612, - 0.06358255, 0.18531723, 0.07759293, 0.12006465, - 0.1305557, 0.058638252, -0.03393652, 0.09622831, - -0.16253184, -2.4580743e-06, 0.079869635, -0.070196845, - -0.005644518, 0.06857898, -0.12598175, -0.035084512, - 0.03156317, -0.12794146, -0.031963028, 0.04692781, - 0.030070418, 0.0071660685, -0.095516115, -0.004643372, - 0.040170413, -0.062104587, -0.0037324072, 0.0554317, - 0.08184801, -0.019164372, 0.06791302, 0.034257166, - -0.10307039, 0.021943003, 0.046745934, 0.0790918, - -0.0265588, -0.007824208, 0.042546265, -0.00977924, - -0.0002440307, -0.017384544, -0.017990116, 0.12252321, - -0.014512694, -0.08251313, 0.08861942, 0.13589665, - 0.026351685, 0.012641483, 0.07466548, 0.044301085, - -0.045414884, -0.051112458, 0.03444247, -0.08502782, - -0.04106223, -0.028126027, 0.028473156, 0.10467447}); - - lstm.SetRecurrentToForgetWeights( - {-0.057784554, -0.026057621, -0.068447545, -0.022581743, - 0.14811787, 0.10826372, 0.09471067, 0.03987225, - -0.0039523416, 0.00030638507, 0.053185795, 0.10572994, - 0.08414449, -0.022036452, -0.00066928595, -0.09203576, - 0.032950465, -0.10985798, -0.023809856, 0.0021431844, - -0.02196096, -0.00326074, 0.00058621005, -0.074678116, - -0.06193199, 0.055729095, 0.03736828, 0.020123724, - 0.061878487, -0.04729229, 0.034919553, -0.07585433, - -0.04421272, -0.044019096, 0.085488975, 0.04058006, - -0.06890133, -0.030951202, -0.024628663, -0.07672815, - 0.034293607, 0.08556707, -0.05293577, -0.033561368, - -0.04899627, 0.0241671, 0.015736353, -0.095442444, - -0.029564252, 0.016493602, -0.035026584, 0.022337519, - -0.026871363, 0.004780428, 0.0077918363, -0.03601621, - 0.016435321, -0.03263031, -0.09543275, -0.047392778, - 0.013454138, 0.028934088, 0.01685226, -0.086110644, - -0.046250615, -0.01847454, 0.047608484, 0.07339695, - 0.034546845, -0.04881143, 0.009128804, -0.08802852, - 0.03761666, 0.008096139, -0.014454086, 0.014361001, - -0.023502491, -0.0011840804, -0.07607001, 0.001856849, - -0.06509276, -0.006021153, -0.08570962, -0.1451793, - 0.060212336, 0.055259194, 0.06974018, 0.049454916, - -0.027794661, -0.08077226, -0.016179763, 0.1169753, - 0.17213494, -0.0056326236, -0.053934924, -0.0124349, - -0.11520337, 0.05409887, 0.088759385, 0.0019655675, - 0.0042065294, 0.03881498, 0.019844765, 0.041858196, - -0.05695512, 0.047233116, 0.038937137, -0.06542224, - 0.014429736, -0.09719407, 0.13908425, -0.05379757, - 0.012321099, 0.082840554, -0.029899208, 0.044217527, - 0.059855383, 0.07711018, -0.045319796, 0.0948846, - -0.011724666, -0.0033288454, -0.033542685, -0.04764985, - -0.13873616, 0.040668588, 0.034832682, -0.015319203, - -0.018715994, 0.046002675, 0.0599172, -0.043107376, - 0.0294216, -0.002314414, -0.022424703, 0.0030315618, - 0.0014641669, 0.0029166266, -0.11878115, 0.013738511, - 0.12375372, -0.0006038222, 0.029104086, 0.087442465, - 0.052958444, 0.07558703, 0.04817258, 0.044462286, - -0.015213451, -0.08783778, -0.0561384, -0.003008196, - 0.047060397, -0.002058388, 0.03429439, -0.018839769, - 0.024734668, 0.024614193, -0.042046934, 0.09597743, - -0.0043254104, 0.04320769, 0.0064070094, -0.0019131786, - -0.02558259, -0.022822596, -0.023273505, -0.02464396, - -0.10991725, -0.006240552, 0.0074488563, 0.024044557, - 0.04383914, -0.046476185, 0.028658995, 0.060410924, - 0.050786525, 0.009452605, -0.0073054377, -0.024810238, - 0.0052906186, 0.0066939713, -0.0020913032, 0.014515517, - 0.015898481, 0.021362653, -0.030262267, 0.016587038, - -0.011442813, 0.041154444, -0.007631438, -0.03423484, - -0.010977775, 0.036152758, 0.0066366293, 0.11915515, - 0.02318443, -0.041350313, 0.021485701, -0.10906167, - -0.028218046, -0.00954771, 0.020531068, -0.11995105, - -0.03672871, 0.024019798, 0.014255957, -0.05221243, - -0.00661567, -0.04630967, 0.033188973, 0.10107534, - -0.014027541, 0.030796422, -0.10270911, -0.035999842, - 0.15443139, 0.07684145, 0.036571592, -0.035900835, - -0.0034699554, 0.06209149, 0.015920248, -0.031122351, - -0.03858649, 0.01849943, 0.13872518, 0.01503974, - 0.069941424, -0.06948533, -0.0088794185, 0.061282158, - -0.047401894, 0.03100163, -0.041533746, -0.10430945, - 0.044574402, -0.01425562, -0.024290353, 0.034563623, - 0.05866852, 0.023947537, -0.09445152, 0.035450947, - 0.02247216, -0.0042998926, 0.061146557, -0.10250651, - 0.020881841, -0.06747029, 0.10062043, -0.0023941975, - 0.03532124, -0.016341697, 0.09685456, -0.016764693, - 0.051808182, 0.05875331, -0.04536488, 0.001626336, - -0.028892258, -0.01048663, -0.009793449, -0.017093895, - 0.010987891, 0.02357273, -0.00010856845, 0.0099760275, - -0.001845119, -0.03551521, 0.0018358806, 0.05763657, - -0.01769146, 0.040995963, 0.02235177, -0.060430344, - 0.11475477, -0.023854522, 0.10071741, 0.0686208, - -0.014250481, 0.034261297, 0.047418304, 0.08562733, - -0.030519066, 0.0060542435, 0.014653856, -0.038836084, - 0.04096551, 0.032249358, -0.08355519, -0.026823482, - 0.056386515, -0.010401743, -0.028396193, 0.08507674, - 0.014410365, 0.020995233, 0.17040324, 0.11511526, - 0.02459721, 0.0066619175, 0.025853224, -0.023133837, - -0.081302024, 0.017264642, -0.009585969, 0.09491168, - -0.051313367, 0.054532815, -0.014298593, 0.10657464, - 0.007076659, 0.10964551, 0.0409152, 0.008275321, - -0.07283536, 0.07937492, 0.04192024, -0.1075027}); - - lstm.SetRecurrentToCellWeights( - {-0.037322544, 0.018592842, 0.0056175636, -0.06253426, - 0.055647098, -0.05713207, -0.05626563, 0.005559383, - 0.03375411, -0.025757805, -0.088049285, 0.06017052, - -0.06570978, 0.007384076, 0.035123326, -0.07920549, - 0.053676967, 0.044480428, -0.07663568, 0.0071805613, - 0.08089997, 0.05143358, 0.038261272, 0.03339287, - -0.027673481, 0.044746667, 0.028349208, 0.020090483, - -0.019443132, -0.030755889, -0.0040000007, 0.04465846, - -0.021585021, 0.0031670958, 0.0053199246, -0.056117613, - -0.10893326, 0.076739706, -0.08509834, -0.027997585, - 0.037871376, 0.01449768, -0.09002357, -0.06111149, - -0.046195522, 0.0422062, -0.005683705, -0.1253618, - -0.012925729, -0.04890792, 0.06985068, 0.037654128, - 0.03398274, -0.004781977, 0.007032333, -0.031787455, - 0.010868644, -0.031489216, 0.09525667, 0.013939797, - 0.0058680447, 0.0167067, 0.02668468, -0.04797466, - -0.048885044, -0.12722108, 0.035304096, 0.06554885, - 0.00972396, -0.039238118, -0.05159735, -0.11329045, - 0.1613692, -0.03750952, 0.06529313, -0.071974665, - -0.11769596, 0.015524369, -0.0013754242, -0.12446318, - 0.02786344, -0.014179351, 0.005264273, 0.14376344, - 0.015983658, 0.03406988, -0.06939408, 0.040699873, - 0.02111075, 0.09669095, 0.041345075, -0.08316494, - -0.07684199, -0.045768797, 0.032298047, -0.041805092, - 0.0119405, 0.0061010392, 0.12652606, 0.0064572375, - -0.024950314, 0.11574242, 0.04508852, -0.04335324, - 0.06760663, -0.027437469, 0.07216407, 0.06977076, - -0.05438599, 0.034033038, -0.028602652, 0.05346137, - 0.043184172, -0.037189785, 0.10420091, 0.00882477, - -0.054019816, -0.074273005, -0.030617684, -0.0028467078, - 0.024302477, -0.0038869337, 0.005332455, 0.0013399826, - 0.04361412, -0.007001822, 0.09631092, -0.06702025, - -0.042049985, -0.035070654, -0.04103342, -0.10273396, - 0.0544271, 0.037184782, -0.13150354, -0.0058036847, - -0.008264958, 0.042035464, 0.05891794, 0.029673764, - 0.0063542654, 0.044788733, 0.054816857, 0.062257513, - -0.00093483756, 0.048938446, -0.004952862, -0.007730018, - -0.04043371, -0.017094059, 0.07229206, -0.023670016, - -0.052195564, -0.025616996, -0.01520939, 0.045104615, - -0.007376126, 0.003533447, 0.006570588, 0.056037236, - 0.12436656, 0.051817212, 0.028532185, -0.08686856, - 0.11868599, 0.07663395, -0.07323171, 0.03463402, - -0.050708205, -0.04458982, -0.11590894, 0.021273347, - 0.1251325, -0.15313013, -0.12224372, 0.17228661, - 0.023029093, 0.086124025, 0.006445803, -0.03496501, - 0.028332196, 0.04449512, -0.042436164, -0.026587414, - -0.006041347, -0.09292539, -0.05678812, 0.03897832, - 0.09465633, 0.008115513, -0.02171956, 0.08304309, - 0.071401566, 0.019622514, 0.032163795, -0.004167056, - 0.02295182, 0.030739572, 0.056506045, 0.004612461, - 0.06524936, 0.059999723, 0.046395954, -0.0045512207, - -0.1335546, -0.030136576, 0.11584653, -0.014678886, - 0.0020118146, -0.09688814, -0.0790206, 0.039770417, - -0.0329582, 0.07922767, 0.029322514, 0.026405897, - 0.04207835, -0.07073373, 0.063781224, 0.0859677, - -0.10925287, -0.07011058, 0.048005477, 0.03438226, - -0.09606514, -0.006669445, -0.043381985, 0.04240257, - -0.06955775, -0.06769346, 0.043903265, -0.026784198, - -0.017840602, 0.024307009, -0.040079936, -0.019946516, - 0.045318738, -0.12233574, 0.026170589, 0.0074471775, - 0.15978073, 0.10185836, 0.10298046, -0.015476589, - -0.039390966, -0.072174534, 0.0739445, -0.1211869, - -0.0347889, -0.07943156, 0.014809798, -0.12412325, - -0.0030663363, 0.039695457, 0.0647603, -0.08291318, - -0.018529687, -0.004423833, 0.0037507233, 0.084633216, - -0.01514876, -0.056505352, -0.012800942, -0.06994386, - 0.012962922, -0.031234352, 0.07029052, 0.016418684, - 0.03618972, 0.055686004, -0.08663945, -0.017404709, - -0.054761406, 0.029065743, 0.052404847, 0.020238016, - 0.0048197987, -0.0214882, 0.07078733, 0.013016777, - 0.06262858, 0.009184685, 0.020785125, -0.043904778, - -0.0270329, -0.03299152, -0.060088247, -0.015162964, - -0.001828936, 0.12642565, -0.056757294, 0.013586685, - 0.09232601, -0.035886683, 0.06000002, 0.05229691, - -0.052580316, -0.082029596, -0.010794592, 0.012947712, - -0.036429964, -0.085508935, -0.13127148, -0.017744139, - 0.031502828, 0.036232427, -0.031581745, 0.023051167, - -0.05325106, -0.03421577, 0.028793324, -0.034633752, - -0.009881397, -0.043551125, -0.018609839, 0.0019097115, - -0.008799762, 0.056595087, 0.0022273948, 0.055752404}); - - lstm.SetRecurrentToOutputWeights({ - 0.025825322, -0.05813119, 0.09495884, -0.045984812, -0.01255415, - -0.0026479573, -0.08196161, -0.054914974, -0.0046604523, -0.029587349, - -0.044576716, -0.07480124, -0.082868785, 0.023254942, 0.027502948, - -0.0039728214, -0.08683098, -0.08116779, -0.014675607, -0.037924774, - -0.023314456, -0.007401714, -0.09255757, 0.029460307, -0.08829125, - -0.005139627, -0.08989442, -0.0555066, 0.13596267, -0.025062224, - -0.048351806, -0.03850004, 0.07266485, -0.022414139, 0.05940088, - 0.075114764, 0.09597592, -0.010211725, -0.0049794707, -0.011523867, - -0.025980417, 0.072999895, 0.11091378, -0.081685916, 0.014416728, - 0.043229222, 0.034178585, -0.07530371, 0.035837382, -0.085607, - -0.007721233, -0.03287832, -0.043848954, -0.06404588, -0.06632928, - -0.073643476, 0.008214239, -0.045984086, 0.039764922, 0.03474462, - 0.060612556, -0.080590084, 0.049127717, 0.04151091, -0.030063879, - 0.008801774, -0.023021035, -0.019558564, 0.05158114, -0.010947698, - -0.011825728, 0.0075720972, 0.0699727, -0.0039981045, 0.069350146, - 0.08799282, 0.016156472, 0.035502106, 0.11695009, 0.006217345, - 0.13392477, -0.037875112, 0.025745004, 0.08940699, -0.00924166, - 0.0046702605, -0.036598757, -0.08811812, 0.10522024, -0.032441203, - 0.008176899, -0.04454919, 0.07058152, 0.0067963637, 0.039206743, - 0.03259838, 0.03725492, -0.09515802, 0.013326398, -0.052055415, - -0.025676316, 0.03198509, -0.015951829, -0.058556724, 0.036879618, - 0.043357447, 0.028362012, -0.05908629, 0.0059240665, -0.04995891, - -0.019187413, 0.0276265, -0.01628143, 0.0025863599, 0.08800015, - 0.035250366, -0.022165963, -0.07328642, -0.009415526, -0.07455109, - 0.11690406, 0.0363299, 0.07411125, 0.042103454, -0.009660886, - 0.019076364, 0.018299393, -0.046004917, 0.08891175, 0.0431396, - -0.026327137, -0.051502608, 0.08979574, -0.051670972, 0.04940282, - -0.07491107, -0.021240504, 0.022596184, -0.034280192, 0.060163025, - -0.058211457, -0.051837247, -0.01349775, -0.04639988, -0.035936575, - -0.011681591, 0.064818054, 0.0073146066, -0.021745546, -0.043124277, - -0.06471268, -0.07053354, -0.029321948, -0.05330136, 0.016933719, - -0.053782392, 0.13747959, -0.1361751, -0.11569455, 0.0033329215, - 0.05693899, -0.053219706, 0.063698, 0.07977434, -0.07924483, - 0.06936997, 0.0034815092, -0.007305279, -0.037325785, -0.07251102, - -0.033633437, -0.08677009, 0.091591336, -0.14165086, 0.021752775, - 0.019683983, 0.0011612234, -0.058154266, 0.049996935, 0.0288841, - -0.0024567875, -0.14345716, 0.010955264, -0.10234828, 0.1183656, - -0.0010731248, -0.023590032, -0.072285876, -0.0724771, -0.026382286, - -0.0014920527, 0.042667855, 0.0018776858, 0.02986552, 0.009814309, - 0.0733756, 0.12289186, 0.018043943, -0.0458958, 0.049412545, - 0.033632483, 0.05495232, 0.036686596, -0.013781798, -0.010036754, - 0.02576849, -0.08307328, 0.010112348, 0.042521734, -0.05869831, - -0.071689695, 0.03876447, -0.13275425, -0.0352966, -0.023077697, - 0.10285965, 0.084736146, 0.15568255, -0.00040734606, 0.027835453, - -0.10292561, -0.032401145, 0.10053256, -0.026142767, -0.08271222, - -0.0030240538, -0.016368777, 0.1070414, 0.042672627, 0.013456989, - -0.0437609, -0.022309763, 0.11576483, 0.04108048, 0.061026827, - -0.0190714, -0.0869359, 0.037901703, 0.0610107, 0.07202949, - 0.01675338, 0.086139716, -0.08795751, -0.014898893, -0.023771819, - -0.01965048, 0.007955471, -0.043740474, 0.03346837, -0.10549954, - 0.090567775, 0.042013682, -0.03176985, 0.12569028, -0.02421228, - -0.029526481, 0.023851605, 0.031539805, 0.05292009, -0.02344001, - -0.07811758, -0.08834428, 0.10094801, 0.16594367, -0.06861939, - -0.021256343, -0.041093912, -0.06669611, 0.035498552, 0.021757556, - -0.09302526, -0.015403468, -0.06614931, -0.051798206, -0.013874718, - 0.03630673, 0.010412845, -0.08077351, 0.046185967, 0.0035662893, - 0.03541868, -0.094149634, -0.034814864, 0.003128424, -0.020674974, - -0.03944324, -0.008110165, -0.11113267, 0.08484226, 0.043586485, - 0.040582247, 0.0968012, -0.065249965, -0.028036479, 0.0050708856, - 0.0017462453, 0.0326779, 0.041296225, 0.09164146, -0.047743853, - -0.015952192, -0.034451712, 0.084197424, -0.05347844, -0.11768019, - 0.085926116, -0.08251791, -0.045081906, 0.0948852, 0.068401024, - 0.024856757, 0.06978981, -0.057309967, -0.012775832, -0.0032452994, - 0.01977615, -0.041040014, -0.024264973, 0.063464895, 0.05431621, - }); - - lstm.SetCellToInputWeights( - {0.040369894, 0.030746894, 0.24704495, 0.018586371, -0.037586458, - -0.15312155, -0.11812848, -0.11465643, 0.20259799, 0.11418174, - -0.10116027, -0.011334949, 0.12411352, -0.076769054, -0.052169047, - 0.21198851, -0.38871562, -0.09061183, -0.09683246, -0.21929175}); - - lstm.SetCellToForgetWeights( - {-0.01998659, -0.15568835, -0.24248174, -0.012770197, 0.041331276, - -0.072311886, -0.052123554, -0.0066330447, -0.043891653, 0.036225766, - -0.047248036, 0.021479502, 0.033189066, 0.11952997, -0.020432774, - 0.64658105, -0.06650122, -0.03467612, 0.095340036, 0.23647355}); - - lstm.SetCellToOutputWeights( - {0.08286371, -0.08261836, -0.51210177, 0.002913762, 0.17764764, - -0.5495371, -0.08460716, -0.24552552, 0.030037103, 0.04123544, - -0.11940523, 0.007358328, 0.1890978, 0.4833202, -0.34441817, - 0.36312827, -0.26375428, 0.1457655, -0.19724406, 0.15548733}); - - lstm.SetProjectionWeights( - {-0.009802181, 0.09401916, 0.0717386, -0.13895074, 0.09641832, - 0.060420845, 0.08539281, 0.054285463, 0.061395317, 0.034448683, - -0.042991187, 0.019801661, -0.16840284, -0.015726732, -0.23041931, - -0.024478018, -0.10959692, -0.013875541, 0.18600968, -0.061274476, - 0.0138165, -0.08160894, -0.07661644, 0.032372914, 0.16169067, - 0.22465782, -0.03993472, -0.004017731, 0.08633481, -0.28869787, - 0.08682067, 0.17240396, 0.014975425, 0.056431185, 0.031037588, - 0.16702051, 0.0077946745, 0.15140012, 0.29405436, 0.120285, - -0.188994, -0.027265169, 0.043389652, -0.022061434, 0.014777949, - -0.20203483, 0.094781205, 0.19100232, 0.13987629, -0.036132768, - -0.06426278, -0.05108664, 0.13221376, 0.009441198, -0.16715929, - 0.15859416, -0.040437475, 0.050779544, -0.022187516, 0.012166504, - 0.027685808, -0.07675938, -0.0055694645, -0.09444123, 0.0046453946, - 0.050794356, 0.10770313, -0.20790008, -0.07149004, -0.11425117, - 0.008225835, -0.035802525, 0.14374903, 0.15262283, 0.048710253, - 0.1847461, -0.007487823, 0.11000021, -0.09542012, 0.22619456, - -0.029149994, 0.08527916, 0.009043713, 0.0042746216, 0.016261552, - 0.022461696, 0.12689082, -0.043589946, -0.12035478, -0.08361797, - -0.050666027, -0.1248618, -0.1275799, -0.071875185, 0.07377272, - 0.09944291, -0.18897448, -0.1593054, -0.06526116, -0.040107165, - -0.004618631, -0.067624845, -0.007576253, 0.10727444, 0.041546922, - -0.20424393, 0.06907816, 0.050412357, 0.00724631, 0.039827548, - 0.12449835, 0.10747581, 0.13708383, 0.09134148, -0.12617786, - -0.06428341, 0.09956831, 0.1208086, -0.14676677, -0.0727722, - 0.1126304, 0.010139365, 0.015571211, -0.038128063, 0.022913318, - -0.042050496, 0.16842307, -0.060597885, 0.10531834, -0.06411776, - -0.07451711, -0.03410368, -0.13393489, 0.06534304, 0.003620307, - 0.04490757, 0.05970546, 0.05197996, 0.02839995, 0.10434969, - -0.013699693, -0.028353551, -0.07260381, 0.047201227, -0.024575593, - -0.036445823, 0.07155557, 0.009672501, -0.02328883, 0.009533515, - -0.03606021, -0.07421458, -0.028082801, -0.2678904, -0.13221288, - 0.18419984, -0.13012612, -0.014588381, -0.035059117, -0.04824723, - 0.07830115, -0.056184657, 0.03277091, 0.025466874, 0.14494097, - -0.12522776, -0.098633975, -0.10766018, -0.08317623, 0.08594209, - 0.07749552, 0.039474737, 0.1776665, -0.07409566, -0.0477268, - 0.29323658, 0.10801441, 0.1154011, 0.013952499, 0.10739139, - 0.10708251, -0.051456142, 0.0074137426, -0.10430189, 0.10034707, - 0.045594677, 0.0635285, -0.0715442, -0.089667566, -0.10811871, - 0.00026344223, 0.08298446, -0.009525053, 0.006585689, -0.24567553, - -0.09450807, 0.09648481, 0.026996298, -0.06419476, -0.04752702, - -0.11063944, -0.23441927, -0.17608605, -0.052156363, 0.067035615, - 0.19271925, -0.0032889997, -0.043264326, 0.09663576, -0.057112187, - -0.10100678, 0.0628376, 0.04447668, 0.017961001, -0.10094388, - -0.10190601, 0.18335468, 0.10494553, -0.052095775, -0.0026118709, - 0.10539724, -0.04383912, -0.042349473, 0.08438151, -0.1947263, - 0.02251204, 0.11216432, -0.10307853, 0.17351969, -0.039091777, - 0.08066188, -0.00561982, 0.12633002, 0.11335965, -0.0088127935, - -0.019777594, 0.06864014, -0.059751723, 0.016233567, -0.06894641, - -0.28651384, -0.004228674, 0.019708522, -0.16305895, -0.07468996, - -0.0855457, 0.099339016, -0.07580735, -0.13775392, 0.08434318, - 0.08330512, -0.12131499, 0.031935584, 0.09180414, -0.08876437, - -0.08049874, 0.008753825, 0.03498998, 0.030215185, 0.03907079, - 0.089751154, 0.029194152, -0.03337423, -0.019092513, 0.04331237, - 0.04299654, -0.036394123, -0.12915532, 0.09793732, 0.07512415, - -0.11319543, -0.032502122, 0.15661901, 0.07671967, -0.005491124, - -0.19379048, -0.218606, 0.21448623, 0.017840758, 0.1416943, - -0.07051762, 0.19488361, 0.02664691, -0.18104725, -0.09334311, - 0.15026465, -0.15493552, -0.057762887, -0.11604192, -0.262013, - -0.01391798, 0.012185008, 0.11156489, -0.07483202, 0.06693364, - -0.26151478, 0.046425626, 0.036540434, -0.16435726, 0.17338543, - -0.21401681, -0.11385144, -0.08283257, -0.069031075, 0.030635102, - 0.010969227, 0.11109743, 0.010919218, 0.027526086, 0.13519906, - 0.01891392, -0.046839405, -0.040167913, 0.017953383, -0.09700955, - 0.0061885654, -0.07000971, 0.026893595, -0.038844477, 0.14543656}); - - static float lstm_input[][20] = { - {// Batch0: 4 (input_sequence_size) * 5 (n_input) - 0.787926, 0.151646, 0.071352, 0.118426, 0.458058, 0.596268, 0.998386, - 0.568695, 0.864524, 0.571277, 0.073204, 0.296072, 0.743333, 0.069199, - 0.045348, 0.867394, 0.291279, 0.013714, 0.482521, 0.626339}, - - {// Batch1: 4 (input_sequence_size) * 5 (n_input) - 0.295743, 0.544053, 0.690064, 0.858138, 0.497181, 0.642421, 0.524260, - 0.134799, 0.003639, 0.162482, 0.640394, 0.930399, 0.050782, 0.432485, - 0.988078, 0.082922, 0.563329, 0.865614, 0.333232, 0.259916}}; - - static float lstm_fw_golden_output[][64] = { - {// Batch0: 4 (input_sequence_size) * 16 (n_output) - -0.00396806, 0.029352, -0.00279226, 0.0159977, -0.00835576, - -0.0211779, 0.0283512, -0.0114597, 0.00907307, -0.0244004, - -0.0152191, -0.0259063, 0.00914318, 0.00415118, 0.017147, - 0.0134203, -0.0166936, 0.0381209, 0.000889694, 0.0143363, - -0.0328911, -0.0234288, 0.0333051, -0.012229, 0.0110322, - -0.0457725, -0.000832209, -0.0202817, 0.0327257, 0.0121308, - 0.0155969, 0.0312091, -0.0213783, 0.0350169, 0.000324794, - 0.0276012, -0.0263374, -0.0371449, 0.0446149, -0.0205474, - 0.0103729, -0.0576349, -0.0150052, -0.0292043, 0.0376827, - 0.0136115, 0.0243435, 0.0354492, -0.0189322, 0.0464512, - -0.00251373, 0.0225745, -0.0308346, -0.0317124, 0.0460407, - -0.0189395, 0.0149363, -0.0530162, -0.0150767, -0.0340193, - 0.0286833, 0.00824207, 0.0264887, 0.0305169}, - {// Batch1: 4 (input_sequence_size) * 16 (n_output) - -0.013869, 0.0287268, -0.00334693, 0.00733398, -0.0287926, - -0.0186926, 0.0193662, -0.0115437, 0.00422612, -0.0345232, - 0.00223253, -0.00957321, 0.0210624, 0.013331, 0.0150954, - 0.02168, -0.0141913, 0.0322082, 0.00227024, 0.0260507, - -0.0188721, -0.0296489, 0.0399134, -0.0160509, 0.0116039, - -0.0447318, -0.0150515, -0.0277406, 0.0316596, 0.0118233, - 0.0214762, 0.0293641, -0.0204549, 0.0450315, -0.00117378, - 0.0167673, -0.0375007, -0.0238314, 0.038784, -0.0174034, - 0.0131743, -0.0506589, -0.0048447, -0.0240239, 0.0325789, - 0.00790065, 0.0220157, 0.0333314, -0.0264787, 0.0387855, - -0.000764675, 0.0217599, -0.037537, -0.0335206, 0.0431679, - -0.0211424, 0.010203, -0.062785, -0.00832363, -0.025181, - 0.0412031, 0.0118723, 0.0239643, 0.0394009}}; - - static float lstm_combined_golden_output[][64] = { - {-0.022014, 0.073544, -0.002235, 0.040068, -0.037136, -0.052788, - 0.075325, -0.029378, 0.024298, -0.07733, -0.030674, -0.060229, - 0.040599, 0.011608, 0.042005, 0.045977, -0.039225, 0.076294, - 0.000735, 0.032852, -0.069869, -0.053312, 0.073527, -0.028136, - 0.021585, -0.102679, -0.004327, -0.043304, 0.072861, 0.027077, - 0.034558, 0.068292, -0.036292, 0.069832, -0.003032, 0.053829, - -0.043821, -0.072713, 0.085029, -0.040374, 0.020014, -0.104521, - -0.034504, -0.059759, 0.062569, 0.025652, 0.049306, 0.061189, - -0.025146, 0.079643, -0.005188, 0.033080, -0.048079, -0.048082, - 0.069369, -0.028900, 0.024572, -0.077547, -0.022517, -0.054477, - 0.038857, 0.013336, 0.043234, 0.044788}, - {-0.039186, 0.070792, -0.005913, 0.02642, -0.068274, -0.05022, - 0.061444, -0.031241, 0.014996, -0.094544, -0.004146, -0.03464, - 0.058981, 0.026097, 0.039781, 0.058408, -0.031887, 0.069252, - 0.00576, 0.054062, -0.042801, -0.059974, 0.085272, -0.034453, - 0.026097, -0.0959, -0.031164, -0.058699, 0.06839, 0.020512, - 0.044727, 0.063609, -0.039863, 0.084819, -0.003909, 0.028666, - -0.075677, -0.045125, 0.070379, -0.033895, 0.022111, -0.097184, - -0.004921, -0.040851, 0.062316, 0.017435, 0.041437, 0.064568, - -0.039656, 0.060726, -0.003402, 0.036854, -0.056503, -0.058554, - 0.068588, -0.034879, 0.01352, -0.09962, -0.01434, -0.039505, - 0.065133, 0.024321, 0.038473, 0.062438}}; - - for (int i = 0; i < lstm.sequence_length(); i++) { - float* batch0_start = lstm_input[0] + i * lstm.num_inputs(); - float* batch0_end = batch0_start + lstm.num_inputs(); - - lstm.SetInput(2 * i * lstm.num_inputs(), batch0_start, batch0_end); - - float* batch1_start = lstm_input[1] + i * lstm.num_inputs(); - float* batch1_end = batch1_start + lstm.num_inputs(); - lstm.SetInput((2 * i + 1) * lstm.num_inputs(), batch1_start, batch1_end); - } - - lstm.Invoke(); - - std::vector expected; - for (int i = 0; i < lstm.sequence_length(); i++) { - float* golden_start_batch0 = - lstm_fw_golden_output[0] + i * lstm.num_fw_outputs(); - float* golden_end_batch0 = golden_start_batch0 + lstm.num_fw_outputs(); - float* golden_start_batch1 = - lstm_fw_golden_output[1] + i * lstm.num_fw_outputs(); - float* golden_end_batch1 = golden_start_batch1 + lstm.num_fw_outputs(); - expected.insert(expected.end(), golden_start_batch0, golden_end_batch0); - expected.insert(expected.end(), golden_start_batch1, golden_end_batch1); - } - EXPECT_THAT(lstm.GetFwOutput(), ElementsAreArray(ArrayFloatNear(expected))); - - // Check if the sum of forward backward matches the golden. - expected.clear(); - for (int i = 0; i < lstm.sequence_length(); i++) { - float* golden_start_batch0 = - lstm_combined_golden_output[0] + i * lstm.num_fw_outputs(); - float* golden_end_batch0 = golden_start_batch0 + lstm.num_fw_outputs(); - float* golden_start_batch1 = - lstm_combined_golden_output[1] + i * lstm.num_fw_outputs(); - float* golden_end_batch1 = golden_start_batch1 + lstm.num_fw_outputs(); - expected.insert(expected.end(), golden_start_batch0, golden_end_batch0); - expected.insert(expected.end(), golden_start_batch1, golden_end_batch1); - } - - std::vector combined; - for (int i = 0; i < lstm.GetFwOutput().size(); ++i) { - combined.push_back(lstm.GetFwOutput()[i] + lstm.GetBwOutput()[i]); - } - EXPECT_THAT(combined, ElementsAreArray(ArrayFloatNear(expected))); -} - -} // namespace -} // namespace tflite - -int main(int argc, char** argv) { - ::tflite::LogToStderr(); - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/contrib/lite/kernels/depthwise_conv.cc b/tensorflow/contrib/lite/kernels/depthwise_conv.cc deleted file mode 100644 index 19958844a1af876bf26251d5ef3ff249a087ffcc..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/depthwise_conv.cc +++ /dev/null @@ -1,317 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include -#include -#include -#include -#include -#include - -#include "tensorflow/contrib/lite/c/builtin_op_data.h" -#include "tensorflow/contrib/lite/c/c_api_internal.h" -#include "tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_float.h" -#include "tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_uint8.h" -#include "tensorflow/contrib/lite/kernels/internal/quantization_util.h" -#include "tensorflow/contrib/lite/kernels/internal/reference/depthwiseconv_float.h" -#include "tensorflow/contrib/lite/kernels/internal/reference/depthwiseconv_uint8.h" -#include "tensorflow/contrib/lite/kernels/internal/tensor.h" -#include "tensorflow/contrib/lite/kernels/kernel_util.h" -#include "tensorflow/contrib/lite/kernels/op_macros.h" -#include "tensorflow/contrib/lite/kernels/padding.h" - -namespace tflite { -namespace ops { -namespace builtin { -namespace depthwise_conv { - -constexpr int kInputTensor = 0; -constexpr int kFilterTensor = 1; -constexpr int kBiasTensor = 2; -constexpr int kOutputTensor = 0; - -// This file has three implementation of DepthwiseConv. -enum KernelType { - kReference, - kGenericOptimized, // Neon-free - kNeonOptimized, -}; - -struct OpData { - TfLitePaddingValues padding; - // The scaling factor from input to output (aka the 'real multiplier') can - // be represented as a fixed point multiplier plus a left shift. - int32_t output_multiplier; - int output_shift; - // The range of the fused activation layer. For example for kNone and - // uint8_t these would be 0 and 255. - int32_t output_activation_min; - int32_t output_activation_max; -}; - -void* Init(TfLiteContext* context, const char* buffer, size_t length) { - // This is a builtin op, so we don't use the contents in 'buffer', if any. - // Instead, we allocate a new object to carry information from Prepare() to - // Eval(). - return new OpData; -} - -void Free(TfLiteContext* context, void* buffer) { - delete reinterpret_cast(buffer); -} - -TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { - auto* params = - reinterpret_cast(node->builtin_data); - OpData* data = reinterpret_cast(node->user_data); - - // TODO(ahentz): use could use GetOptionalInputTensor() here, but we need to - // decide whether we are OK with optional tensors being completely absent, as - // opposed to having -1 as their index. - bool hasBias = NumInputs(node) == 3; - - TF_LITE_ENSURE(context, hasBias || NumInputs(node) == 2); - const TfLiteTensor* input = GetInput(context, node, kInputTensor); - const TfLiteTensor* filter = GetInput(context, node, kFilterTensor); - const TfLiteTensor* bias = nullptr; - - TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); - TfLiteTensor* output = GetOutput(context, node, kOutputTensor); - - TF_LITE_ENSURE_EQ(context, NumDimensions(input), 4); - TF_LITE_ENSURE_EQ(context, NumDimensions(filter), 4); - - // The parameter 'depth_multiplier' is redundant, so we check here to make - // sure it is consistent with the given dimensions. - TF_LITE_ENSURE_EQ(context, - params->depth_multiplier * SizeOfDimension(input, 3), - SizeOfDimension(filter, 3)); - - const TfLiteType data_type = input->type; - TF_LITE_ENSURE(context, - data_type == kTfLiteFloat32 || data_type == kTfLiteUInt8); - TF_LITE_ENSURE_EQ(context, output->type, data_type); - TF_LITE_ENSURE_EQ(context, filter->type, data_type); - - if (hasBias) { - bias = GetInput(context, node, kBiasTensor); - if (data_type == kTfLiteUInt8) { - TF_LITE_ENSURE_EQ(context, bias->type, kTfLiteInt32); - TF_LITE_ENSURE_EQ(context, bias->params.zero_point, 0); - } else { - TF_LITE_ENSURE_EQ(context, bias->type, data_type); - } - TF_LITE_ENSURE_EQ(context, NumDimensions(bias), 1); - TF_LITE_ENSURE_EQ(context, SizeOfDimension(filter, 3), - SizeOfDimension(bias, 0)); - } - - int channels_out = SizeOfDimension(filter, 3); - int width = SizeOfDimension(input, 2); - int height = SizeOfDimension(input, 1); - int filter_width = SizeOfDimension(filter, 2); - int filter_height = SizeOfDimension(filter, 1); - int batches = SizeOfDimension(input, 0); - - // Matching GetWindowedOutputSize in TensorFlow. - auto padding = params->padding; - auto compute_out_size = [padding](int image_size, int filter_size, int stride, - int dilation_rate) -> int { - int effective_filter_size = (filter_size - 1) * dilation_rate + 1; - return padding == kTfLitePaddingSame - ? (image_size + stride - 1) / stride - : padding == kTfLitePaddingValid - ? (image_size - effective_filter_size + stride) / stride - : 0; - }; - - int out_width = compute_out_size(width, filter_width, params->stride_width, - params->dilation_width_factor); - int out_height = - compute_out_size(height, filter_height, params->stride_height, - params->dilation_height_factor); - - data->padding.height = - ComputePadding(params->stride_height, params->dilation_height_factor, - height, filter_height, out_height); - data->padding.width = - ComputePadding(params->stride_width, params->dilation_width_factor, width, - filter_width, out_width); - - // Note that quantized inference requires that all tensors have their - // parameters set. This is usually done during quantized training. - if (data_type != kTfLiteFloat32) { - double real_multiplier = 0.0; - TF_LITE_ENSURE_STATUS(GetQuantizedConvolutionMultipler( - context, input, filter, bias, output, &real_multiplier)); - int exponent; - QuantizeMultiplier(real_multiplier, &data->output_multiplier, &exponent); - data->output_shift = -exponent; - CalculateActivationRangeUint8(params->activation, output, - &data->output_activation_min, - &data->output_activation_max); - } - - TfLiteIntArray* outputSize = TfLiteIntArrayCreate(4); - outputSize->data[0] = batches; - outputSize->data[1] = out_height; - outputSize->data[2] = out_width; - outputSize->data[3] = channels_out; - return context->ResizeTensor(context, output, outputSize); -} - -template -void EvalFloat(TfLiteContext* context, TfLiteNode* node, - TfLiteDepthwiseConvParams* params, OpData* data, - const TfLiteTensor* input, const TfLiteTensor* filter, - const TfLiteTensor* bias, TfLiteTensor* output) { - float output_activation_min, output_activation_max; - CalculateActivationRange(params->activation, &output_activation_min, - &output_activation_max); - - void (*depthwise_conv)(const DepthwiseParams&, const RuntimeShape&, - const float*, const RuntimeShape&, const float*, - const RuntimeShape&, const float*, const RuntimeShape&, - float*); - if (kernel_type == kReference) { - depthwise_conv = &reference_ops::DepthwiseConv; - } else { - depthwise_conv = &optimized_ops::DepthwiseConv; - } - - DepthwiseParams op_params; - op_params.padding_type = PaddingType::kSame; - op_params.padding_values.width = data->padding.width; - op_params.padding_values.height = data->padding.height; - op_params.stride_width = params->stride_width; - op_params.stride_height = params->stride_height; - op_params.dilation_width_factor = params->dilation_width_factor; - op_params.dilation_height_factor = params->dilation_height_factor; - op_params.depth_multiplier = params->depth_multiplier; - op_params.float_activation_min = output_activation_min; - op_params.float_activation_max = output_activation_max; - depthwise_conv(op_params, GetTensorShape(input), GetTensorData(input), - GetTensorShape(filter), GetTensorData(filter), - GetTensorShape(bias), GetTensorData(bias), - GetTensorShape(output), GetTensorData(output)); -} - -template -void EvalQuantized(TfLiteContext* context, TfLiteNode* node, - TfLiteDepthwiseConvParams* params, OpData* data, - const TfLiteTensor* input, const TfLiteTensor* filter, - const TfLiteTensor* bias, TfLiteTensor* output) { - auto input_offset = -input->params.zero_point; - auto filter_offset = -filter->params.zero_point; - auto output_offset = output->params.zero_point; - - void (*depthwise_conv)(const DepthwiseParams&, const RuntimeShape&, - const uint8*, const RuntimeShape&, const uint8*, - const RuntimeShape&, const int32*, const RuntimeShape&, - uint8*); - - if (kernel_type == kReference) { - depthwise_conv = &reference_ops::DepthwiseConv; - } else { - depthwise_conv = &optimized_ops::DepthwiseConv; - } - - DepthwiseParams op_params; - op_params.padding_type = PaddingType::kSame; - op_params.padding_values.width = data->padding.width; - op_params.padding_values.height = data->padding.height; - op_params.stride_width = params->stride_width; - op_params.stride_height = params->stride_height; - op_params.dilation_width_factor = params->dilation_width_factor; - op_params.dilation_height_factor = params->dilation_height_factor; - op_params.depth_multiplier = params->depth_multiplier; - op_params.input_offset = input_offset; - op_params.weights_offset = filter_offset; - op_params.output_offset = output_offset; - op_params.output_multiplier = data->output_multiplier; - op_params.output_shift = -data->output_shift; - op_params.quantized_activation_min = data->output_activation_min; - op_params.quantized_activation_max = data->output_activation_max; - depthwise_conv(op_params, GetTensorShape(input), - GetTensorData(input), GetTensorShape(filter), - GetTensorData(filter), GetTensorShape(bias), - GetTensorData(bias), GetTensorShape(output), - GetTensorData(output)); -} - -template -TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { - auto* params = - reinterpret_cast(node->builtin_data); - OpData* data = reinterpret_cast(node->user_data); - - TfLiteTensor* output = GetOutput(context, node, kOutputTensor); - const TfLiteTensor* input = GetInput(context, node, kInputTensor); - const TfLiteTensor* filter = GetInput(context, node, kFilterTensor); - const TfLiteTensor* bias = - (NumInputs(node) == 3) ? GetInput(context, node, kBiasTensor) : nullptr; - - // TODO(aselle): Consider whether float conv and quantized conv should be - // separate ops to avoid dispatch overhead here. - switch (input->type) { // Already know in/out types are same. - case kTfLiteFloat32: - EvalFloat(context, node, params, data, input, filter, bias, - output); - break; - case kTfLiteUInt8: - EvalQuantized(context, node, params, data, input, filter, - bias, output); - break; - default: - context->ReportError(context, "Type %d not currently supported.", - input->type); - return kTfLiteError; - } - return kTfLiteOk; -} - -} // namespace depthwise_conv - -TfLiteRegistration* Register_DEPTHWISE_CONVOLUTION_REF() { - static TfLiteRegistration r = { - depthwise_conv::Init, depthwise_conv::Free, depthwise_conv::Prepare, - depthwise_conv::Eval}; - return &r; -} - -TfLiteRegistration* Register_DEPTHWISE_CONVOLUTION_GENERIC_OPT() { - static TfLiteRegistration r = { - depthwise_conv::Init, depthwise_conv::Free, depthwise_conv::Prepare, - depthwise_conv::Eval}; - return &r; -} - -TfLiteRegistration* Register_DEPTHWISE_CONVOLUTION_NEON_OPT() { - static TfLiteRegistration r = { - depthwise_conv::Init, depthwise_conv::Free, depthwise_conv::Prepare, - depthwise_conv::Eval}; - return &r; -} - -TfLiteRegistration* Register_DEPTHWISE_CONV_2D() { -#ifdef USE_NEON - return Register_DEPTHWISE_CONVOLUTION_NEON_OPT(); -#else - return Register_DEPTHWISE_CONVOLUTION_GENERIC_OPT(); -#endif -} - -} // namespace builtin -} // namespace ops -} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/depthwise_conv_test.cc b/tensorflow/contrib/lite/kernels/depthwise_conv_test.cc deleted file mode 100644 index 4a33a0319d0dc3fd56cd3a173518d4fe49ace3ec..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/depthwise_conv_test.cc +++ /dev/null @@ -1,455 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include -#include -#include "absl/memory/memory.h" -#include "tensorflow/contrib/lite/interpreter.h" -#include "tensorflow/contrib/lite/kernels/register.h" -#include "tensorflow/contrib/lite/kernels/test_util.h" -#include "tensorflow/contrib/lite/model.h" - -namespace tflite { - -namespace ops { -namespace builtin { - -TfLiteRegistration* Register_DEPTHWISE_CONVOLUTION_REF(); -TfLiteRegistration* Register_DEPTHWISE_CONVOLUTION_GENERIC_OPT(); -TfLiteRegistration* Register_DEPTHWISE_CONVOLUTION_NEON_OPT(); - -} // namespace builtin -} // namespace ops - -namespace { - -using ::testing::ElementsAreArray; - -class BaseDepthwiseConvolutionOpModel : public SingleOpModel { - public: - // TODO(ahentz): Also test different activation types, bias, padding types, - // stride values. - BaseDepthwiseConvolutionOpModel(TfLiteRegistration* registration, - const TensorData& input, - const TensorData& filter, - const TensorData& output, - Padding padding_type, - int dilation_factor = 1) { - input_ = AddInput(input); - filter_ = AddInput(filter); - - int bias_size = GetShape(filter_)[3]; - if (input.type == TensorType_FLOAT32) { - bias_ = AddInput({TensorType_FLOAT32, {bias_size}}); - } else { - // This is a quantized version. The scale of 'bias' depends on the scales - // of input and filter. Supposedly this is correctly set during quantized - // training. - auto bias_scale = GetScale(input_) * GetScale(filter_); - TensorData bias{TensorType_INT32, {bias_size}, 0, 0, bias_scale}; - bias_ = AddInput(bias); - } - - output_ = AddOutput(output); - - int input_depth = GetShape(input_)[3]; - int output_depth = GetShape(filter_)[3]; - int depth_mul = output_depth / input_depth; - - SetBuiltinOp( - BuiltinOperator_DEPTHWISE_CONV_2D, - BuiltinOptions_DepthwiseConv2DOptions, - CreateDepthwiseConv2DOptions(builder_, padding_type, 1, 1, depth_mul, - ActivationFunctionType_NONE, - dilation_factor, dilation_factor) - .Union()); - - resolver_ = absl::make_unique( - BuiltinOperator_DEPTHWISE_CONV_2D, registration); - - BuildInterpreter({GetShape(input_), GetShape(filter_), GetShape(bias_)}); - } - - protected: - int input_; - int filter_; - int bias_; - int output_; -}; - -class DepthwiseConvolutionOpModel : public BaseDepthwiseConvolutionOpModel { - public: - using BaseDepthwiseConvolutionOpModel::BaseDepthwiseConvolutionOpModel; - - void SetFilter(std::initializer_list f) { PopulateTensor(filter_, f); } - - void SetBias(std::initializer_list f) { PopulateTensor(bias_, f); } - - void SetInput(std::initializer_list data) { - PopulateTensor(input_, data); - } - - std::vector GetOutput() { return ExtractVector(output_); } -}; - -const auto kKernelMap = new std::map({ - {"Reference", ops::builtin::Register_DEPTHWISE_CONVOLUTION_REF()}, - {"GenericOptimized", - ops::builtin::Register_DEPTHWISE_CONVOLUTION_GENERIC_OPT()}, - {"NeonOptimized", ops::builtin::Register_DEPTHWISE_CONVOLUTION_NEON_OPT()}, -}); - -class DepthwiseConvolutionOpTest : public SingleOpTest { - protected: - const std::map& GetKernelMap() override { - return *kKernelMap; - } -}; - -TEST_P(DepthwiseConvolutionOpTest, SimpleTest) { - DepthwiseConvolutionOpModel m(GetRegistration(), - {TensorType_FLOAT32, {1, 3, 2, 2}}, - {TensorType_FLOAT32, {1, 2, 2, 4}}, - {TensorType_FLOAT32, {}}, Padding_VALID); - - m.SetInput({ - 1, 2, 7, 8, // column 1 - 3, 4, 9, 10, // column 2 - 5, 6, 11, 12, // column 3 - }); - m.SetFilter({ - 1, 2, 3, 4, // - -9, 10, -11, 12, // - 5, 6, 7, 8, // - 13, -14, 15, -16, // - }); - m.SetBias({1, 2, 3, 4}); - - m.Invoke(); - - EXPECT_THAT(m.GetOutput(), ElementsAreArray({ - 71, -34, 99, -20, // - 91, -26, 127, -4, // - })); -} - -TEST_P(DepthwiseConvolutionOpTest, SimpleDilatedTestPaddingValid) { - const int depth = 1; - const int image_width = 9; - const int image_height = 9; - const int image_batch_count = 1; - const int filter_size = 3; - const int filter_count = 1; - const int dilation_factor = 3; - DepthwiseConvolutionOpModel m( - GetRegistration(), - {TensorType_FLOAT32, - {image_batch_count, image_height, image_width, depth}}, - {TensorType_FLOAT32, {depth, filter_size, filter_size, filter_count}}, - {TensorType_FLOAT32, {}}, Padding_VALID, dilation_factor); - - // The image matrix is: - // | 0 | 0 | 0 | 0 | 0 | 0 | 0 | 0 | 0 | - // | 0 | 0 | 0 | 0 | 0 | 0 | 0 | 0 | 0 | - // | 0 | 0 | 0 | 0 | 0 | 0 | 0 | 0 | 0 | - // | 0 | 0 | 0 | 1 | 1 | 1 | 0 | 0 | 0 | - // | 0 | 0 | 0 | 1 | 1 | 1 | 0 | 0 | 0 | - // | 0 | 0 | 0 | 1 | 1 | 1 | 0 | 0 | 0 | - // | 0 | 0 | 0 | 0 | 0 | 0 | 0 | 0 | 0 | - // | 0 | 0 | 0 | 0 | 0 | 0 | 0 | 0 | 0 | - // | 0 | 0 | 0 | 0 | 0 | 0 | 0 | 0 | 0 | - // clang-format off - m.SetInput({0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 1, 1, 1, 0, 0, 0, - 0, 0, 0, 1, 1, 1, 0, 0, 0, - 0, 0, 0, 1, 1, 1, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0}); - // clang-format on - // The filter matrix is: - // | 1 | 2 | 3 | - // | 4 | 5 | 6 | - // | 7 | 8 | 9 | - m.SetFilter({1, 2, 3, 4, 5, 6, 7, 8, 9}); - // No bias for this test. - m.SetBias({0}); - m.Invoke(); - - // Since the dilation rate is 3 this will reduce the size of the output from - // 10x10 to 3x3 of all 5s. Specifically: - // | 5 | 5 | 5 | - // | 5 | 5 | 5 | - // | 5 | 5 | 5 | - EXPECT_THAT(m.GetOutput(), ElementsAreArray({5, 5, 5, 5, 5, 5, 5, 5, 5})); -} - -TEST_P(DepthwiseConvolutionOpTest, SimpleDilatedTestPaddingSame) { - const int depth = 1; - const int image_width = 3; - const int image_height = 3; - const int image_batch_count = 1; - const int filter_size = 2; - const int filter_count = 1; - const int dilation_factor = 2; - DepthwiseConvolutionOpModel m( - GetRegistration(), - {TensorType_FLOAT32, - {image_batch_count, image_height, image_width, depth}}, - {TensorType_FLOAT32, {depth, filter_size, filter_size, filter_count}}, - {TensorType_FLOAT32, {}}, Padding_SAME, dilation_factor); - - // The image matrix is: - // | 1 | 1 | 1 | - // | 1 | 1 | 1 | - // | 1 | 1 | 1 | - m.SetInput({1, 1, 1, 1, 1, 1, 1, 1, 1}); - // The filter matrix is: - // | 1 | 2 | - // | 3 | 4 | - m.SetFilter({1, 2, 3, 4}); - // No bias for this test. - m.SetBias({0}); - m.Invoke(); - - // Output: - // | 4 | 7 | 3 | - // | 6 |10 | 4 | - // | 2 | 3 | 1 | - EXPECT_THAT(m.GetOutput(), ElementsAreArray({4, 7, 3, 6, 10, 4, 2, 3, 1})); -} - -class QuantizedDepthwiseConvolutionOpModel - : public BaseDepthwiseConvolutionOpModel { - public: - using BaseDepthwiseConvolutionOpModel::BaseDepthwiseConvolutionOpModel; - - void SetInput(std::initializer_list data) { - QuantizeAndPopulate(input_, data); - } - - void SetFilter(std::initializer_list data) { - QuantizeAndPopulate(filter_, data); - } - - void SetBias(std::initializer_list data) { - QuantizeAndPopulate(bias_, data); - } - - std::vector GetOutput() { return ExtractVector(output_); } - std::vector GetDequantizedOutput() { - return Dequantize(ExtractVector(output_), - GetScale(output_), GetZeroPoint(output_)); - } -}; - -class QuantizedDepthwiseConvolutionOpTest : public SingleOpTest { - protected: - const std::map& GetKernelMap() override { - return *kKernelMap; - } -}; - -// In this test we set the input and output scales so that the results match -// exactly the 'non-quantized' version. -TEST_P(QuantizedDepthwiseConvolutionOpTest, SimpleTestQuantized) { - QuantizedDepthwiseConvolutionOpModel m( - GetRegistration(), {TensorType_UINT8, {1, 3, 2, 2}, -63.5, 64}, - {TensorType_UINT8, {1, 2, 2, 4}, -63.5, 64}, - {TensorType_UINT8, {}, -127, 128}, Padding_VALID); - - m.SetInput({ - 1, 2, 7, 8, // column 1 - 3, 4, 9, 10, // column 2 - 5, 6, 11, 12, // column 3 - }); - m.SetFilter({ - 1, 2, 3, 4, // - -9, 10, -11, 12, // - 5, 6, 7, 8, // - 13, -14, 15, -16, // - }); - m.SetBias({1, 2, 3, 4}); - - m.Invoke(); - - EXPECT_THAT(m.GetDequantizedOutput(), ElementsAreArray(ArrayFloatNear( - { - 71, -34, 99, -20, // - 91, -26, 127, -4, // - }, - 1e-5))); - // For good measure, let's also verify the quantized values: - EXPECT_THAT(m.GetOutput(), ElementsAreArray({ - 198, 93, 226, 107, // - 218, 101, 254, 123, // - })); -} - -TEST_P(QuantizedDepthwiseConvolutionOpTest, - SimpleTestQuantizedFilterMultiplierGreaterThan1) { - QuantizedDepthwiseConvolutionOpModel quant_op( - GetRegistration(), {TensorType_UINT8, {1, 3, 2, 2}, -63.5, 64}, - {TensorType_UINT8, {1, 2, 2, 4}, -128.5, 128}, - {TensorType_UINT8, {}, -127, 128}, Padding_VALID); - DepthwiseConvolutionOpModel float_op(GetRegistration(), - {TensorType_FLOAT32, {1, 3, 2, 2}}, - {TensorType_FLOAT32, {1, 2, 2, 4}}, - {TensorType_FLOAT32, {}}, Padding_VALID); - - std::initializer_list input = { - 1, 2, 7, 8, // column 1 - 3, 4, 9, 10, // column 2 - 5, 6, 11, 12, // column 3 - }; - std::initializer_list filter = { - 1, 2, 3, 4, // - -9, 10, -11, 12, // - 5, 6, 7, 8, // - 13, -14, 15, -16, // - }; - std::initializer_list bias = {1, 2, 3, 4}; - - quant_op.SetInput(input); - quant_op.SetFilter(filter); - quant_op.SetBias(bias); - quant_op.Invoke(); - - float_op.SetInput(input); - float_op.SetFilter(filter); - float_op.SetBias(bias); - float_op.Invoke(); - - EXPECT_THAT(quant_op.GetDequantizedOutput(), - ElementsAreArray(ArrayFloatNear(float_op.GetOutput(), 1))); -} - -TEST_P(QuantizedDepthwiseConvolutionOpTest, SimpleDilatedTestPaddingValid) { - const int depth = 1; - const int image_width = 9; - const int image_height = 9; - const int image_batch_count = 1; - const int filter_size = 3; - const int filter_count = 1; - const int dilation_factor = 3; - QuantizedDepthwiseConvolutionOpModel m( - GetRegistration(), - {TensorType_UINT8, - {image_batch_count, image_height, image_width, depth}, - 0, - 255}, - {TensorType_UINT8, - {depth, filter_size, filter_size, filter_count}, - 0, - 255}, - {TensorType_UINT8, {}, 0, 255}, Padding_VALID, dilation_factor); - - // The image matrix is: - // | 0 | 0 | 0 | 0 | 0 | 0 | 0 | 0 | 0 | - // | 0 | 0 | 0 | 0 | 0 | 0 | 0 | 0 | 0 | - // | 0 | 0 | 0 | 0 | 0 | 0 | 0 | 0 | 0 | - // | 0 | 0 | 0 | 1 | 1 | 1 | 0 | 0 | 0 | - // | 0 | 0 | 0 | 1 | 1 | 1 | 0 | 0 | 0 | - // | 0 | 0 | 0 | 1 | 1 | 1 | 0 | 0 | 0 | - // | 0 | 0 | 0 | 0 | 0 | 0 | 0 | 0 | 0 | - // | 0 | 0 | 0 | 0 | 0 | 0 | 0 | 0 | 0 | - // | 0 | 0 | 0 | 0 | 0 | 0 | 0 | 0 | 0 | - // clang-format off - m.SetInput({0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 1, 1, 1, 0, 0, 0, - 0, 0, 0, 1, 1, 1, 0, 0, 0, - 0, 0, 0, 1, 1, 1, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0, - 0, 0, 0, 0, 0, 0, 0, 0, 0}); - // clang-format on - // The filter matrix is: - // | 1 | 2 | 3 | - // | 4 | 5 | 6 | - // | 7 | 8 | 9 | - m.SetFilter({1, 2, 3, 4, 5, 6, 7, 8, 9}); - // No bias for this test. - m.SetBias({0}); - m.Invoke(); - - // Since the dilation rate is 3 this will reduce the size of the output from - // 10x10 to 3x3 of all 5s. Specifically: - // | 5 | 5 | 5 | - // | 5 | 5 | 5 | - // | 5 | 5 | 5 | - EXPECT_THAT(m.GetDequantizedOutput(), - ElementsAreArray({5, 5, 5, 5, 5, 5, 5, 5, 5})); -} - -TEST_P(QuantizedDepthwiseConvolutionOpTest, SimpleDilatedTestPaddingSame) { - const int depth = 1; - const int image_width = 3; - const int image_height = 3; - const int image_batch_count = 1; - const int filter_size = 2; - const int filter_count = 1; - const int dilation_factor = 2; - QuantizedDepthwiseConvolutionOpModel m( - GetRegistration(), - {TensorType_UINT8, - {image_batch_count, image_height, image_width, depth}, - 0, - 255}, - {TensorType_UINT8, - {depth, filter_size, filter_size, filter_count}, - 0, - 255}, - {TensorType_UINT8, {}, 0, 255}, Padding_SAME, dilation_factor); - - // The image matrix is: - // | 1 | 1 | 1 | - // | 1 | 1 | 1 | - // | 1 | 1 | 1 | - m.SetInput({1, 1, 1, 1, 1, 1, 1, 1, 1}); - // The filter matrix is: - // | 1 | 2 | - // | 3 | 4 | - m.SetFilter({1, 2, 3, 4}); - // No bias for this test. - m.SetBias({0}); - m.Invoke(); - - // Output: - // | 4 | 7 | 3 | - // | 6 |10 | 4 | - // | 2 | 3 | 1 | - EXPECT_THAT(m.GetDequantizedOutput(), - ElementsAreArray({4, 7, 3, 6, 10, 4, 2, 3, 1})); -} - -INSTANTIATE_TEST_CASE_P( - DepthwiseConvolutionOpTest, DepthwiseConvolutionOpTest, - ::testing::ValuesIn(SingleOpTest::GetKernelTags(*kKernelMap))); - -INSTANTIATE_TEST_CASE_P( - QuantizedDepthwiseConvolutionOpTest, QuantizedDepthwiseConvolutionOpTest, - ::testing::ValuesIn(SingleOpTest::GetKernelTags(*kKernelMap))); - -} // namespace -} // namespace tflite - -int main(int argc, char** argv) { - ::tflite::LogToStderr(); - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/contrib/lite/kernels/dequantize.cc b/tensorflow/contrib/lite/kernels/dequantize.cc deleted file mode 100644 index 59bf64e0afabc44a984a9797cabbcfcde531f1f6..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/dequantize.cc +++ /dev/null @@ -1,107 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include -#include - -#include "tensorflow/contrib/lite/c/builtin_op_data.h" -#include "tensorflow/contrib/lite/c/c_api_internal.h" -#include "tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h" -#include "tensorflow/contrib/lite/kernels/internal/tensor.h" -#include "tensorflow/contrib/lite/kernels/kernel_util.h" -#include "tensorflow/contrib/lite/kernels/op_macros.h" - -namespace tflite { -namespace ops { -namespace builtin { -namespace dequantize { - -struct OpContext { - OpContext(TfLiteContext* context, TfLiteNode* node) { - input = GetInput(context, node, 0); - output = GetOutput(context, node, 0); - } - const TfLiteTensor* input; - TfLiteTensor* output; -}; - -struct OpData { - // This boolean value is only used when the input tensor is constant. - bool float_dequantized_weights_initialized; -}; - -void* Init(TfLiteContext* context, const char* buffer, size_t length) { - auto* op_data = new OpData(); - op_data->float_dequantized_weights_initialized = false; - return op_data; -} - -void Free(TfLiteContext* context, void* buffer) { - delete reinterpret_cast(buffer); -} - -TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { - TF_LITE_ENSURE_EQ(context, NumInputs(node), 1); - TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); - - OpContext op_context(context, node); - - TF_LITE_ENSURE(context, op_context.input->type == kTfLiteUInt8); - - op_context.output->type = kTfLiteFloat32; - // If the input tensor is constant, we can persist the dequantized value in - // the output tensor. Otherwise we run dequantize upon each eval. - if (IsConstantTensor(op_context.input)) { - op_context.output->allocation_type = kTfLiteArenaRwPersistent; - } - return context->ResizeTensor(context, op_context.output, - TfLiteIntArrayCopy(op_context.input->dims)); -} - -TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { - OpData* op_data = reinterpret_cast(node->user_data); - OpContext op_context(context, node); - if (IsConstantTensor(op_context.input) && - op_data->float_dequantized_weights_initialized) { - return kTfLiteOk; - } - - tflite::DequantizationParams op_params; - op_params.zero_point = op_context.input->params.zero_point; - op_params.scale = op_context.input->params.scale; - optimized_ops::Dequantize(op_params, GetTensorShape(op_context.input), - GetTensorData(op_context.input), - GetTensorShape(op_context.output), - GetTensorData(op_context.output)); - - if (IsConstantTensor(op_context.input)) { - op_data->float_dequantized_weights_initialized = true; - } - - return kTfLiteOk; -} - -} // namespace dequantize - -TfLiteRegistration* Register_DEQUANTIZE_OPT() { - static TfLiteRegistration r = {dequantize::Init, dequantize::Free, - dequantize::Prepare, dequantize::Eval}; - return &r; -} - -TfLiteRegistration* Register_DEQUANTIZE() { return Register_DEQUANTIZE_OPT(); } - -} // namespace builtin -} // namespace ops -} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/dequantize_test.cc b/tensorflow/contrib/lite/kernels/dequantize_test.cc deleted file mode 100644 index fcd74206177a0a97db168338e3619d4b95c052a9..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/dequantize_test.cc +++ /dev/null @@ -1,65 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include -#include "tensorflow/contrib/lite/interpreter.h" -#include "tensorflow/contrib/lite/kernels/register.h" -#include "tensorflow/contrib/lite/kernels/test_util.h" -#include "tensorflow/contrib/lite/model.h" - -namespace tflite { -namespace { - -using ::testing::ElementsAreArray; - -class DequantizeOpModel : public SingleOpModel { - public: - DequantizeOpModel(std::initializer_list shape, float min, float max) { - input_ = AddInput({TensorType_UINT8, shape, min, max}); - output_ = AddOutput({TensorType_FLOAT32, shape}); - SetBuiltinOp(BuiltinOperator_DEQUANTIZE, BuiltinOptions_DequantizeOptions, - CreateDequantizeOptions(builder_).Union()); - - BuildInterpreter({GetShape(input_)}); - } - - void SetInput(std::initializer_list data) { - PopulateTensor(input_, data); - } - - std::vector GetOutput() { return ExtractVector(output_); } - - private: - int input_; - int output_; -}; - -TEST(SplitOpTest, FourDimensional) { - DequantizeOpModel m({2, 5}, -63.5, 64); - - m.SetInput({0, 1, 2, 3, 4, 251, 252, 253, 254, 255}); - m.Invoke(); - EXPECT_THAT(m.GetOutput(), - ElementsAreArray(ArrayFloatNear( - {-63.5, -63, -62.5, -62, -61.5, 62, 62.5, 63, 63.5, 64}))); -} - -} // namespace -} // namespace tflite - -int main(int argc, char** argv) { - ::tflite::LogToStderr(); - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/contrib/lite/kernels/detection_postprocess_test.cc b/tensorflow/contrib/lite/kernels/detection_postprocess_test.cc deleted file mode 100644 index 1e8caebd820248e2e2bc031e08ba671b28084198..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/detection_postprocess_test.cc +++ /dev/null @@ -1,235 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include -#include -#include - -#include -#include "flatbuffers/flexbuffers.h" // TF:flatbuffers -#include "tensorflow/contrib/lite/interpreter.h" -#include "tensorflow/contrib/lite/kernels/register.h" -#include "tensorflow/contrib/lite/kernels/test_util.h" -#include "tensorflow/contrib/lite/model.h" - -namespace tflite { -namespace ops { -namespace custom { - -TfLiteRegistration* Register_DETECTION_POSTPROCESS(); - -namespace { - -using ::testing::ElementsAre; -using ::testing::ElementsAreArray; - -class BaseDetectionPostprocessOpModel : public SingleOpModel { - public: - BaseDetectionPostprocessOpModel(const TensorData& input1, - const TensorData& input2, - const TensorData& input3, - const TensorData& output1, - const TensorData& output2, - const TensorData& output3, - const TensorData& output4) { - input1_ = AddInput(input1); - input2_ = AddInput(input2); - input3_ = AddInput(input3); - output1_ = AddOutput(output1); - output2_ = AddOutput(output2); - output3_ = AddOutput(output3); - output4_ = AddOutput(output4); - - flexbuffers::Builder fbb; - fbb.Map([&]() { - fbb.Int("max_detections", 3); - fbb.Int("max_classes_per_detection", 1); - fbb.Float("nms_score_threshold", 0.0); - fbb.Float("nms_iou_threshold", 0.5); - fbb.Int("num_classes", 2); - fbb.Float("y_scale", 10.0); - fbb.Float("x_scale", 10.0); - fbb.Float("h_scale", 5.0); - fbb.Float("w_scale", 5.0); - }); - fbb.Finish(); - SetCustomOp("TFLite_Detection_PostProcess", fbb.GetBuffer(), - Register_DETECTION_POSTPROCESS); - BuildInterpreter({GetShape(input1_), GetShape(input2_), GetShape(input3_)}); - } - - int input1() { return input1_; } - int input2() { return input2_; } - int input3() { return input3_; } - - template - void SetInput1(std::initializer_list data) { - PopulateTensor(input1_, data); - } - - template - void SetInput2(std::initializer_list data) { - PopulateTensor(input2_, data); - } - - template - void SetInput3(std::initializer_list data) { - PopulateTensor(input3_, data); - } - - template - std::vector GetOutput1() { - return ExtractVector(output1_); - } - - template - std::vector GetOutput2() { - return ExtractVector(output2_); - } - - template - std::vector GetOutput3() { - return ExtractVector(output3_); - } - - template - std::vector GetOutput4() { - return ExtractVector(output4_); - } - - std::vector GetOutputShape1() { return GetTensorShape(output1_); } - std::vector GetOutputShape2() { return GetTensorShape(output2_); } - std::vector GetOutputShape3() { return GetTensorShape(output3_); } - std::vector GetOutputShape4() { return GetTensorShape(output4_); } - - protected: - int input1_; - int input2_; - int input3_; - int output1_; - int output2_; - int output3_; - int output4_; -}; - -TEST(DetectionPostprocessOpTest, FloatTest) { - BaseDetectionPostprocessOpModel m( - {TensorType_FLOAT32, {1, 6, 4}}, {TensorType_FLOAT32, {1, 6, 3}}, - {TensorType_FLOAT32, {6, 4}}, {TensorType_FLOAT32, {}}, - {TensorType_FLOAT32, {}}, {TensorType_FLOAT32, {}}, - {TensorType_FLOAT32, {}}); - - // six boxes in center-size encoding - m.SetInput1({0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, - 0.0, -1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, - 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0}); - // class scores - two classes with background - m.SetInput2({0., .9, .8, 0., .75, .72, 0., .6, .5, 0., .93, .95, 0., - .5, .4, 0., .3, .2}); - // six anchors in center-size encoding - m.SetInput3({0.5, 0.5, 1.0, 1.0, 0.5, 0.5, 1.0, 1.0, - 0.5, 0.5, 1.0, 1.0, 0.5, 10.5, 1.0, 1.0, - 0.5, 10.5, 1.0, 1.0, 0.5, 100.5, 1.0, 1.0}); - // Same boxes in box-corner encoding: - // { 0.0, 0.0, 1.0, 1.0, - // 0.0, 0.1, 1.0, 1.1, - // 0.0, -0.1, 1.0, 0.9, - // 0.0, 10.0, 1.0, 11.0, - // 0.0, 10.1, 1.0, 11.1, - // 0.0, 100.0, 1.0, 101.0} - m.Invoke(); - // detection_boxes - // in center-size - std::vector output_shape1 = m.GetOutputShape1(); - EXPECT_THAT(output_shape1, ElementsAre(1, 3, 4)); - EXPECT_THAT( - m.GetOutput1(), - ElementsAreArray(ArrayFloatNear( - {0.0, 10.0, 1.0, 11.0, 0.0, 0.0, 1.0, 1.0, 0.0, 100.0, 1.0, 101.0}, - 1e-1))); - // detection_classes - std::vector output_shape2 = m.GetOutputShape2(); - EXPECT_THAT(output_shape2, ElementsAre(1, 3)); - EXPECT_THAT(m.GetOutput2(), - ElementsAreArray(ArrayFloatNear({1, 0, 0}, 1e-1))); - // detection_scores - std::vector output_shape3 = m.GetOutputShape3(); - EXPECT_THAT(output_shape3, ElementsAre(1, 3)); - EXPECT_THAT(m.GetOutput3(), - ElementsAreArray(ArrayFloatNear({0.95, 0.9, 0.3}, 1e-1))); - // num_detections - std::vector output_shape4 = m.GetOutputShape4(); - EXPECT_THAT(output_shape4, ElementsAre(1)); - EXPECT_THAT(m.GetOutput4(), - ElementsAreArray(ArrayFloatNear({3.0}, 1e-1))); -} - -TEST(DetectionPostprocessOpTest, QuantizedTest) { - BaseDetectionPostprocessOpModel m( - {TensorType_UINT8, {1, 6, 4}, -1.0, 1.0}, - {TensorType_UINT8, {1, 6, 3}, 0.0, 1.0}, - {TensorType_UINT8, {6, 4}, 0.0, 100.5}, {TensorType_FLOAT32, {}}, - {TensorType_FLOAT32, {}}, {TensorType_FLOAT32, {}}, - {TensorType_FLOAT32, {}}); - // six boxes in center-size encoding - std::vector> inputs1 = { - {0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, -1.0, 0.0, 0.0, - 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0}}; - m.QuantizeAndPopulate(m.input1(), inputs1[0]); - // class scores - two classes with background - std::vector> inputs2 = { - {0., .9, .8, 0., .75, .72, 0., .6, .5, 0., .93, .95, 0., .5, .4, 0., .3, - .2}}; - m.QuantizeAndPopulate(m.input2(), inputs2[0]); - // six anchors in center-size encoding - std::vector> inputs3 = { - {0.5, 0.5, 1.0, 1.0, 0.5, 0.5, 1.0, 1.0, 0.5, 0.5, 1.0, 1.0, - 0.5, 10.5, 1.0, 1.0, 0.5, 10.5, 1.0, 1.0, 0.5, 100.5, 1.0, 1.0}}; - m.QuantizeAndPopulate(m.input3(), inputs3[0]); - m.Invoke(); - // detection_boxes - // in center-size - std::vector output_shape1 = m.GetOutputShape1(); - EXPECT_THAT(output_shape1, ElementsAre(1, 3, 4)); - EXPECT_THAT( - m.GetOutput1(), - ElementsAreArray(ArrayFloatNear( - {0.0, 10.0, 1.0, 11.0, 0.0, 0.0, 1.0, 1.0, 0.0, 100.0, 1.0, 101.0}, - 3e-1))); - // detection_classes - std::vector output_shape2 = m.GetOutputShape2(); - EXPECT_THAT(output_shape2, ElementsAre(1, 3)); - EXPECT_THAT(m.GetOutput2(), - ElementsAreArray(ArrayFloatNear({1, 0, 0}, 1e-1))); - // detection_scores - std::vector output_shape3 = m.GetOutputShape3(); - EXPECT_THAT(output_shape3, ElementsAre(1, 3)); - EXPECT_THAT(m.GetOutput3(), - ElementsAreArray(ArrayFloatNear({0.95, 0.9, 0.3}, 1e-1))); - // num_detections - std::vector output_shape4 = m.GetOutputShape4(); - EXPECT_THAT(output_shape4, ElementsAre(1)); - EXPECT_THAT(m.GetOutput4(), - ElementsAreArray(ArrayFloatNear({3.0}, 1e-1))); -} -} // namespace -} // namespace custom -} // namespace ops -} // namespace tflite - -int main(int argc, char** argv) { - ::tflite::LogToStderr(); - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/contrib/lite/kernels/embedding_lookup_test.cc b/tensorflow/contrib/lite/kernels/embedding_lookup_test.cc deleted file mode 100644 index 4a88d168c60203f10802e634def9b1d1316c9c6d..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/embedding_lookup_test.cc +++ /dev/null @@ -1,176 +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. -==============================================================================*/ -// Unit test for TFLite Lookup op. - -#include -#include -#include - -#include -#include -#include "tensorflow/contrib/lite/interpreter.h" -#include "tensorflow/contrib/lite/kernels/register.h" -#include "tensorflow/contrib/lite/kernels/test_util.h" -#include "tensorflow/contrib/lite/model.h" - -namespace tflite { -namespace { - -using ::testing::ElementsAreArray; - -class BaseEmbeddingLookupOpModel : public SingleOpModel { - public: - BaseEmbeddingLookupOpModel(std::initializer_list index_shape, - std::initializer_list weight_shape, - TensorType weight_type = TensorType_FLOAT32) { - input_ = AddInput(TensorType_INT32); - weight_ = AddInput(weight_type); - output_ = AddOutput(TensorType_FLOAT32); - SetBuiltinOp(BuiltinOperator_EMBEDDING_LOOKUP, BuiltinOptions_NONE, 0); - BuildInterpreter({index_shape, weight_shape}); - } - - void SetInput(std::initializer_list data) { - PopulateTensor(input_, data); - } - - std::vector GetOutput() { return ExtractVector(output_); } - - protected: - int input_; - int weight_; - int output_; -}; - -class EmbeddingLookupOpModel : public BaseEmbeddingLookupOpModel { - public: - using BaseEmbeddingLookupOpModel::BaseEmbeddingLookupOpModel; - - void Set3DWeightMatrix(const std::function& function) { - TfLiteTensor* tensor = interpreter_->tensor(weight_); - int rows = tensor->dims->data[0]; - int columns = tensor->dims->data[1]; - int features = tensor->dims->data[2]; - for (int i = 0; i < rows; i++) { - for (int j = 0; j < columns; j++) { - for (int k = 0; k < features; k++) { - tensor->data.f[(i * columns + j) * features + k] = function(i, j, k); - } - } - } - } -}; - -class HybridEmbeddingLookupOpModel : public BaseEmbeddingLookupOpModel { - public: - HybridEmbeddingLookupOpModel(std::initializer_list index_shape, - std::initializer_list weight_shape) - : BaseEmbeddingLookupOpModel(index_shape, weight_shape, - TensorType_UINT8) {} - - void SetWeight(std::initializer_list data) { - SymmetricQuantizeAndPopulate(weight_, data); - } -}; - -// TODO(ahentz): write more tests that exercise the details of the op, such as -// lookup errors and variable input shapes. -TEST(EmbeddingLookupOpTest, SimpleTest) { - EmbeddingLookupOpModel m({3}, {3, 2, 4}); - m.SetInput({1, 0, 2}); - m.Set3DWeightMatrix( - [](int i, int j, int k) { return i + j / 10.0f + k / 100.0f; }); - - m.Invoke(); - - EXPECT_THAT(m.GetOutput(), - ElementsAreArray(ArrayFloatNear({ - 1.00, 1.01, 1.02, 1.03, 1.10, 1.11, 1.12, 1.13, // Row 1 - 0.00, 0.01, 0.02, 0.03, 0.10, 0.11, 0.12, 0.13, // Row 0 - 2.00, 2.01, 2.02, 2.03, 2.10, 2.11, 2.12, 2.13, // Row 2 - }))); -} - -TEST(HybridEmbeddingLookupHybridOpTest, Simple2DTest) { - HybridEmbeddingLookupOpModel m({3}, {3, 8}); - m.SetInput({1, 0, 2}); - m.SetWeight({ - 0.00, 0.01, 0.02, 0.03, 0.10, 0.11, 0.12, 0.13, // Row 0 - 1.00, -1.01, 1.02, 1.03, 1.10, 1.11, 1.12, 1.13, // Row 1 - 2.00, 2.01, 2.02, 2.03, 2.10, 2.11, 2.12, 2.13, // Row 2 - }); - - m.Invoke(); - - EXPECT_THAT(m.GetOutput(), - ElementsAreArray(ArrayFloatNear( - { - 1.00, -1.01, 1.02, 1.03, 1.10, 1.11, 1.12, 1.13, // Row 1 - 0.00, 0.01, 0.02, 0.03, 0.10, 0.11, 0.12, 0.13, // Row 0 - 2.00, 2.01, 2.02, 2.03, 2.10, 2.11, 2.12, 2.13, // Row 2 - }, - 7.41e-03))); -} - -TEST(HybridEmbeddingLookupHybridOpTest, Simple3DTest) { - HybridEmbeddingLookupOpModel m({3}, {3, 2, 4}); - m.SetInput({1, 0, 2}); - m.SetWeight({ - 0.00, 0.01, 0.02, 0.03, 0.10, 0.11, 0.12, 0.13, // Row 0 - 1.00, -1.01, 1.02, 1.03, 1.10, 1.11, 1.12, 1.13, // Row 1 - 2.00, 2.01, 2.02, 2.03, 2.10, 2.11, 2.12, 2.13, // Row 2 - }); - - m.Invoke(); - - EXPECT_THAT(m.GetOutput(), - ElementsAreArray(ArrayFloatNear( - { - 1.00, -1.01, 1.02, 1.03, 1.10, 1.11, 1.12, 1.13, // Row 1 - 0.00, 0.01, 0.02, 0.03, 0.10, 0.11, 0.12, 0.13, // Row 0 - 2.00, 2.01, 2.02, 2.03, 2.10, 2.11, 2.12, 2.13, // Row 2 - }, - 7.41e-03))); -} - -TEST(HybridEmbeddingLookupHybridOpTest, Simple4DTest) { - HybridEmbeddingLookupOpModel m({3}, {3, 2, 2, 2}); - m.SetInput({1, 0, 2}); - m.SetWeight({ - 0.00, 0.01, 0.02, 0.03, 0.10, 0.11, 0.12, 0.13, // Row 0 - 1.00, -1.01, 1.02, 1.03, 1.10, 1.11, 1.12, 1.13, // Row 1 - 2.00, 2.01, 2.02, 2.03, 2.10, 2.11, 2.12, 2.13, // Row 2 - }); - - m.Invoke(); - - EXPECT_THAT(m.GetOutput(), - ElementsAreArray(ArrayFloatNear( - { - 1.00, -1.01, 1.02, 1.03, 1.10, 1.11, 1.12, 1.13, // Row 1 - 0.00, 0.01, 0.02, 0.03, 0.10, 0.11, 0.12, 0.13, // Row 0 - 2.00, 2.01, 2.02, 2.03, 2.10, 2.11, 2.12, 2.13, // Row 2 - }, - 7.41e-03))); -} - -} // namespace -} // namespace tflite - -int main(int argc, char** argv) { - ::tflite::LogToStderr(); - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/contrib/lite/kernels/floor.cc b/tensorflow/contrib/lite/kernels/floor.cc deleted file mode 100644 index 59ff77f35b8d3f1e4abb41687b2985cd75dd45a2..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/floor.cc +++ /dev/null @@ -1,59 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/contrib/lite/c/c_api_internal.h" -#include "tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h" -#include "tensorflow/contrib/lite/kernels/internal/tensor.h" -#include "tensorflow/contrib/lite/kernels/kernel_util.h" - -namespace tflite { -namespace ops { -namespace builtin { -namespace floor { - -constexpr int kInputTensor = 0; -constexpr int kOutputTensor = 0; - -TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { - const TfLiteTensor* input = GetInput(context, node, kInputTensor); - TfLiteTensor* output = GetOutput(context, node, kOutputTensor); - TF_LITE_ENSURE_EQ(context, NumInputs(node), 1); - TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); - TF_LITE_ENSURE_EQ(context, input->type, kTfLiteFloat32); - output->type = input->type; - TfLiteIntArray* output_size = TfLiteIntArrayCopy(input->dims); - return context->ResizeTensor(context, output, output_size); -} - -TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { - const TfLiteTensor* input = GetInput(context, node, kInputTensor); - TfLiteTensor* output = GetOutput(context, node, kOutputTensor); - - optimized_ops::Floor(GetTensorShape(input), GetTensorData(input), - GetTensorShape(output), GetTensorData(output)); - - return kTfLiteOk; -} -} // namespace floor - -TfLiteRegistration* Register_FLOOR() { - static TfLiteRegistration r = {/*init=*/nullptr, - /*free=*/nullptr, floor::Prepare, floor::Eval}; - return &r; -} - -} // namespace builtin -} // namespace ops -} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/fully_connected.cc b/tensorflow/contrib/lite/kernels/fully_connected.cc deleted file mode 100644 index cac556db33a6fad43ae51736f10df2435ee17152..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/fully_connected.cc +++ /dev/null @@ -1,508 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include -#include -#include -#include -#include -#include - -#include "tensorflow/contrib/lite/c/builtin_op_data.h" -#include "tensorflow/contrib/lite/c/c_api_internal.h" -#include "tensorflow/contrib/lite/kernels/activation_functor.h" -#include "tensorflow/contrib/lite/kernels/gemm_support.h" -#include "tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h" -#include "tensorflow/contrib/lite/kernels/internal/quantization_util.h" -#include "tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h" -#include "tensorflow/contrib/lite/kernels/internal/tensor.h" -#include "tensorflow/contrib/lite/kernels/internal/tensor_utils.h" -#include "tensorflow/contrib/lite/kernels/kernel_util.h" -#include "tensorflow/contrib/lite/kernels/op_macros.h" - -namespace tflite { -namespace ops { -namespace builtin { -namespace fully_connected { - -// This file has four implementations of FullyConnected -enum KernelType { - kReference, - kGenericOptimized, // Neon-free - kNeonOptimized, - kPie, // Used by the PIE team -}; - -struct OpData { - // The scaling factor from input to output (aka the 'real multiplier') can - // be represented as a fixed point multiplier plus a left shift. - int32_t output_multiplier; - int output_shift; - // The range of the fused activation layer. For example for kNone and - // uint8_t these would be 0 and 255. - int32_t output_activation_min; - int32_t output_activation_max; - // The index of the temporary tensor where the quantized inputs are cached. - int scratch_tensor_index; -}; - -constexpr int kInputTensor = 0; -constexpr int kWeightsTensor = 1; -constexpr int kBiasTensor = 2; -constexpr int kOutputTensor = 0; -constexpr int kShuffledInputWorkspaceTensor = 1; - -void* Init(TfLiteContext* context, const char* buffer, size_t length) { - // This is a builtin op, so we don't use the contents in 'buffer', if any. - // Instead, we allocate a new object to carry information from Prepare() to - // Eval(). - gemm_support::IncrementUsageCounter(context); - auto* op_data = new OpData(); - context->AddTensors(context, /*tensors_to_add=*/2, - &op_data->scratch_tensor_index); - return op_data; -} - -void Free(TfLiteContext* context, void* buffer) { - gemm_support::DecrementUsageCounter(context); - delete reinterpret_cast(buffer); -} - -TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { - auto* params = - reinterpret_cast(node->builtin_data); - OpData* data = reinterpret_cast(node->user_data); - - // Check we have all the inputs and outputs we need. - TF_LITE_ENSURE_EQ(context, node->inputs->size, 3); - // Shuffled formats need a workspace to store the shuffled input activations. - const int expected_outputs_count = - params->weights_format == kTfLiteFullyConnectedWeightsFormatDefault ? 1 - : 2; - TF_LITE_ENSURE_EQ(context, node->outputs->size, expected_outputs_count); - - const TfLiteTensor* input = GetInput(context, node, kInputTensor); - const TfLiteTensor* filter = GetInput(context, node, kWeightsTensor); - const TfLiteTensor* bias = GetOptionalInputTensor(context, node, kBiasTensor); - TfLiteTensor* output = GetOutput(context, node, kOutputTensor); - - // Check all the parameters of tensor match within themselves and match the - // input configuration. - int input_size = 1; - for (int i = 0; i < input->dims->size; i++) { - input_size *= input->dims->data[i]; - } - - TF_LITE_ENSURE_EQ(context, NumDimensions(filter), 2); - const int batch_size = input_size / filter->dims->data[1]; - const int num_units = filter->dims->data[0]; - - TF_LITE_ENSURE_EQ(context, input_size, batch_size * filter->dims->data[1]); - if (bias) { - TF_LITE_ENSURE_EQ(context, NumElements(bias), SizeOfDimension(filter, 0)); - } - - // Note that quantized inference requires that all tensors have their - // parameters set. This is usually done during quantized training. - TfLiteType data_type = input->type; - if (data_type != kTfLiteFloat32) { - double real_multiplier = 0.0; - TF_LITE_ENSURE_STATUS(GetQuantizedConvolutionMultipler( - context, input, filter, bias, output, &real_multiplier)); - int exponent; - QuantizeMultiplier(real_multiplier, &data->output_multiplier, &exponent); - data->output_shift = -exponent; - TF_LITE_ENSURE_STATUS(CalculateActivationRangeQuantized( - context, params->activation, output, &data->output_activation_min, - &data->output_activation_max)); - } - - // If we have to perform on-the-fly quantization (with quantized weights and - // float inputs) first we need to quantize the inputs. Allocate a temporary - // buffer to store the intermediate quantized values. - if (input->type == kTfLiteFloat32 && filter->type == kTfLiteUInt8) { - TfLiteIntArrayFree(node->temporaries); - node->temporaries = TfLiteIntArrayCreate(2); - node->temporaries->data[0] = data->scratch_tensor_index; - - TfLiteTensor* input_quantized = GetTemporary(context, node, /*index=*/0); - input_quantized->type = kTfLiteUInt8; - input_quantized->allocation_type = kTfLiteArenaRw; - - // TODO(raziel): add this logic to ResizeTensor. - if (!TfLiteIntArrayEqual(input_quantized->dims, input->dims)) { - TfLiteIntArray* input_quantized_size = TfLiteIntArrayCopy(input->dims); - TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, input_quantized, - input_quantized_size)); - } - node->temporaries->data[1] = data->scratch_tensor_index + 1; - TfLiteTensor* scaling_factors = GetTemporary(context, node, /*index=*/1); - scaling_factors->type = kTfLiteFloat32; - scaling_factors->allocation_type = kTfLiteArenaRw; - int scaling_dims[1] = {batch_size}; - if (!TfLiteIntArrayEqualsArray(scaling_factors->dims, 1, scaling_dims)) { - TfLiteIntArray* scaling_factors_size = TfLiteIntArrayCreate(1); - scaling_factors_size->data[0] = batch_size; - TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, scaling_factors, - scaling_factors_size)); - } - } - - // Resize output. - TfLiteIntArray* output_size_array = TfLiteIntArrayCreate(2); - output_size_array->data[0] = batch_size; - output_size_array->data[1] = num_units; - TF_LITE_ENSURE_OK(context, - context->ResizeTensor(context, output, output_size_array)); - return kTfLiteOk; -} - -TfLiteStatus EvalPie(TfLiteContext* context, TfLiteNode* node, - TfLiteFullyConnectedParams* params, OpData* data, - const TfLiteTensor* input, const TfLiteTensor* filter, - const TfLiteTensor* bias, TfLiteTensor* output) { - int total_input_size = 1; - for (int i = 0; i < input->dims->size; i++) { - total_input_size *= input->dims->data[i]; - } - - int input_size = filter->dims->data[1]; - const int batch_size = total_input_size / filter->dims->data[1]; - const int num_units = filter->dims->data[0]; - - // Output = bias if bias tensor exists. - if (bias) { - tensor_utils::VectorBatchVectorAssign(bias->data.f, num_units, batch_size, - output->data.f); - } else { - tensor_utils::ZeroVector(output->data.f, batch_size * num_units); - } - - // Compute output += weight * input - tensor_utils::MatrixBatchVectorMultiplyAccumulate( - filter->data.f, num_units, input_size, input->data.f, batch_size, - output->data.f, /*result_stride=*/1); - - // Apply activation function - tensor_utils::ApplyActivationToVector(output->data.f, batch_size * num_units, - params->activation, output->data.f); - - return kTfLiteOk; -} - -TfLiteStatus EvalHybrid(TfLiteContext* context, TfLiteNode* node, - TfLiteFullyConnectedParams* params, OpData* data, - const TfLiteTensor* input, const TfLiteTensor* filter, - const TfLiteTensor* bias, TfLiteTensor* input_quantized, - TfLiteTensor* scaling_factors, TfLiteTensor* output) { - // Check the types for this hybrid Op. - TF_LITE_ENSURE_EQ(context, input->type, kTfLiteFloat32); - TF_LITE_ENSURE_EQ(context, filter->type, kTfLiteUInt8); - TF_LITE_ENSURE_EQ(context, bias->type, kTfLiteFloat32); - TF_LITE_ENSURE_EQ(context, output->type, kTfLiteFloat32); - - int total_input_size = 1; - for (int i = 0; i < input->dims->size; i++) { - total_input_size *= input->dims->data[i]; - } - - const int input_size = filter->dims->data[1]; - const int batch_size = total_input_size / filter->dims->data[1]; - const int num_units = filter->dims->data[0]; - - // Output = bias if bias tensor exists. - if (bias) { - tensor_utils::VectorBatchVectorAssign(bias->data.f, num_units, batch_size, - output->data.f); - } else { - tensor_utils::ZeroVector(output->data.f, batch_size * num_units); - } - - // Save matrix multiplication computation for all zero input. - if (tensor_utils::IsZeroVector(input->data.f, total_input_size)) { - tensor_utils::ApplyActivationToVector(output->data.f, - batch_size * num_units, - params->activation, output->data.f); - return kTfLiteOk; - } - - // Quantize input from float to uint8 + quantization params (scaling factor). - float unused_min, unused_max; - float* scaling_factors_ptr = scaling_factors->data.f; - int8_t* quant_data = reinterpret_cast(input_quantized->data.uint8); - - // Quantize each batch independently. - for (int b = 0; b < batch_size; ++b) { - const int offset = b * input_size; - tensor_utils::SymmetricQuantizeFloats(input->data.f + offset, input_size, - quant_data + offset, &unused_min, - &unused_max, &scaling_factors_ptr[b]); - // Incorporate scaling of the filter. - scaling_factors_ptr[b] *= filter->params.scale; - } - - // Compute output += weight * quantized_input - tensor_utils::MatrixBatchVectorMultiplyAccumulate( - reinterpret_cast(filter->data.uint8), num_units, input_size, - quant_data, scaling_factors_ptr, batch_size, output->data.f, - /*result_stride=*/1); - - // Apply activation function to floats. - tensor_utils::ApplyActivationToVector(output->data.f, batch_size * num_units, - params->activation, output->data.f); - return kTfLiteOk; -} - -#define TF_LITE_MACRO_DISPATCH(macro_name, params, target_namespace) \ - if (params->activation == kTfLiteActNone) { \ - macro_name(target_namespace, kNone); \ - } \ - if (params->activation == kTfLiteActRelu) { \ - macro_name(target_namespace, kRelu); \ - } \ - if (params->activation == kTfLiteActRelu6) { \ - macro_name(target_namespace, kRelu6); \ - } - -template -TfLiteStatus EvalQuantized(TfLiteContext* context, TfLiteNode* node, - TfLiteFullyConnectedParams* params, OpData* data, - const TfLiteTensor* input, - const TfLiteTensor* filter, const TfLiteTensor* bias, - TfLiteTensor* output) { - gemmlowp::GemmContext* gemm_context = gemm_support::GetFromContext(context); - - int32_t input_offset = -input->params.zero_point; - int32_t filter_offset = -filter->params.zero_point; - int32_t output_offset = output->params.zero_point; -#define TF_LITE_FULLY_CONNECTED(type, output_data_type) \ - { \ - FullyConnectedParams op_params; \ - op_params.input_offset = input_offset; \ - op_params.weights_offset = filter_offset; \ - op_params.output_offset = output_offset; \ - op_params.output_multiplier = data->output_multiplier; \ - op_params.output_shift = -data->output_shift; \ - op_params.quantized_activation_min = data->output_activation_min; \ - op_params.quantized_activation_max = data->output_activation_max; \ - type::FullyConnected( \ - op_params, GetTensorShape(input), GetTensorData(input), \ - GetTensorShape(filter), GetTensorData(filter), \ - GetTensorShape(bias), GetTensorData(bias), \ - GetTensorShape(output), GetTensorData(output), \ - gemm_context); \ - } - if (kernel_type == kReference) { - switch (output->type) { - case kTfLiteUInt8: - TF_LITE_FULLY_CONNECTED(reference_ops, uint8_t); - break; - case kTfLiteInt16: - TF_LITE_FULLY_CONNECTED(reference_ops, int16_t); - break; - default: - context->ReportError( - context, - "Quantized FullyConnected expects output data type uint8 or int16"); - return kTfLiteError; - } - } else if (kernel_type == kPie && input->type == kTfLiteFloat32) { - // Pie currently only supports quantized models and float inputs/outputs. - TfLiteTensor* input_quantized = GetTemporary(context, node, /*index=*/0); - TfLiteTensor* scaling_factors = GetTemporary(context, node, /*index=*/1); - return EvalHybrid(context, node, params, data, input, filter, bias, - input_quantized, scaling_factors, output); - } else { - switch (output->type) { - case kTfLiteUInt8: - TF_LITE_FULLY_CONNECTED(optimized_ops, uint8_t); - break; - case kTfLiteInt16: - TF_LITE_FULLY_CONNECTED(optimized_ops, int16_t); - break; - default: - context->ReportError( - context, - "Quantized FullyConnected expects output data type uint8 or int16"); - return kTfLiteError; - } - } -#undef TF_LITE_FULLY_CONNECTED - - return kTfLiteOk; -} - -template -TfLiteStatus EvalShuffledQuantized(TfLiteContext* context, TfLiteNode* node, - TfLiteFullyConnectedParams* params, - OpData* data, const TfLiteTensor* input, - const TfLiteTensor* filter, - const TfLiteTensor* bias, - TfLiteTensor* output, - TfLiteTensor* shuffled_input_workspace) { - gemmlowp::GemmContext* gemm_context = gemm_support::GetFromContext(context); - - // TODO(b/110697972) decide more consistently if / how / where we want - // to perform this kind of runtime data type checks. - if (input->type != kTfLiteUInt8 || filter->type != kTfLiteUInt8 || - bias->type != kTfLiteInt32 || output->type != kTfLiteInt16 || - shuffled_input_workspace->type != kTfLiteUInt8) { - context->ReportError(context, "Unexpected data type"); - return kTfLiteError; - } - -#define TF_LITE_SHUFFLED_FULLY_CONNECTED(type) \ - { \ - FullyConnectedParams op_params; \ - op_params.output_multiplier = data->output_multiplier; \ - op_params.output_shift = -data->output_shift; \ - op_params.quantized_activation_min = data->output_activation_min; \ - op_params.quantized_activation_max = data->output_activation_max; \ - type::ShuffledFullyConnected( \ - op_params, GetTensorShape(input), GetTensorData(input), \ - GetTensorShape(filter), GetTensorData(filter), \ - GetTensorShape(bias), GetTensorData(bias), \ - GetTensorShape(output), GetTensorData(output), \ - GetTensorData(shuffled_input_workspace), gemm_context); \ - } - if (kernel_type == kReference) { - TF_LITE_SHUFFLED_FULLY_CONNECTED(reference_ops); - } else { - TF_LITE_SHUFFLED_FULLY_CONNECTED(optimized_ops); - } -#undef TF_LITE_SHUFFLED_FULLY_CONNECTED - - return kTfLiteOk; -} - -template -TfLiteStatus EvalFloat(TfLiteContext* context, TfLiteNode* node, - TfLiteFullyConnectedParams* params, OpData* data, - const TfLiteTensor* input, const TfLiteTensor* filter, - const TfLiteTensor* bias, TfLiteTensor* output) { - float output_activation_min, output_activation_max; - CalculateActivationRange(params->activation, &output_activation_min, - &output_activation_max); -#define TF_LITE_FULLY_CONNECTED(type) \ - { \ - FullyConnectedParams op_params; \ - op_params.float_activation_min = output_activation_min; \ - op_params.float_activation_max = output_activation_max; \ - type::FullyConnected(op_params, GetTensorShape(input), \ - GetTensorData(input), GetTensorShape(filter), \ - GetTensorData(filter), GetTensorShape(bias), \ - GetTensorData(bias), GetTensorShape(output), \ - GetTensorData(output)); \ - } - if (kernel_type == kReference) { - TF_LITE_FULLY_CONNECTED(reference_ops); - } else if (kernel_type == kPie) { - return EvalPie(context, node, params, data, input, filter, bias, output); - } else { - TF_LITE_FULLY_CONNECTED(optimized_ops); - } -#undef TF_LITE_FULLY_CONNECTED - - return kTfLiteOk; -} - -#undef TF_LITE_MACRO_DISPATCH - -template -TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { - auto* params = - reinterpret_cast(node->builtin_data); - OpData* data = reinterpret_cast(node->user_data); - - const TfLiteTensor* input = GetInput(context, node, kInputTensor); - const TfLiteTensor* filter = GetInput(context, node, kWeightsTensor); - const TfLiteTensor* bias = GetOptionalInputTensor(context, node, kBiasTensor); - TfLiteTensor* output = GetOutput(context, node, kOutputTensor); - - switch (filter->type) { // Already know in/out types are same. - case kTfLiteFloat32: - return EvalFloat(context, node, params, data, input, filter, - bias, output); - case kTfLiteUInt8: - if (params->weights_format == - kTfLiteFullyConnectedWeightsFormatShuffled4x16Int8) { - TfLiteTensor* shuffled_input_workspace = - GetOutput(context, node, kShuffledInputWorkspaceTensor); - return EvalShuffledQuantized(context, node, params, data, - input, filter, bias, output, - shuffled_input_workspace); - } else if (params->weights_format == - kTfLiteFullyConnectedWeightsFormatDefault) { - return EvalQuantized(context, node, params, data, input, - filter, bias, output); - } else { - context->ReportError(context, - "Unhandled fully-connected weights format"); - return kTfLiteError; - } - default: - context->ReportError(context, "Type %d not currently supported.", - filter->type); - return kTfLiteError; - } - return kTfLiteOk; -} - -} // namespace fully_connected - -TfLiteRegistration* Register_FULLY_CONNECTED_REF() { - static TfLiteRegistration r = { - fully_connected::Init, fully_connected::Free, fully_connected::Prepare, - fully_connected::Eval}; - return &r; -} - -TfLiteRegistration* Register_FULLY_CONNECTED_NEON_OPT() { - static TfLiteRegistration r = { - fully_connected::Init, fully_connected::Free, fully_connected::Prepare, - fully_connected::Eval}; - return &r; -} - -TfLiteRegistration* Register_FULLY_CONNECTED_GENERIC_OPT() { - static TfLiteRegistration r = { - fully_connected::Init, fully_connected::Free, fully_connected::Prepare, - fully_connected::Eval}; - return &r; -} - -TfLiteRegistration* Register_FULLY_CONNECTED_PIE() { - static TfLiteRegistration r = {fully_connected::Init, fully_connected::Free, - fully_connected::Prepare, - fully_connected::Eval}; - return &r; -} - -TfLiteRegistration* Register_FULLY_CONNECTED() { - // TODO(ahentz): We don't have a dedicated quantized version of the PIE - // kernel. For now, the quantized version just defer to the corresponding - // optimized MINI kernel. At some point we will allow different libraries to - // be built with different kernels, but for now we have to pick one here. - return Register_FULLY_CONNECTED_PIE(); -#ifdef USE_NEON - return Register_FULLY_CONNECTED_NEON_OPT(); -#else - return Register_FULLY_CONNECTED_GENERIC_OPT(); -#endif -} - -} // namespace builtin -} // namespace ops -} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/fully_connected_test.cc b/tensorflow/contrib/lite/kernels/fully_connected_test.cc deleted file mode 100644 index 08b43209466a1b85613ae41d5aa776194f992c60..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/fully_connected_test.cc +++ /dev/null @@ -1,766 +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. -==============================================================================*/ -// Unit test for TFLite FULLY_CONNECTED op. - -#include -#include -#include - -#include -#include -#include "absl/memory/memory.h" -#include "tensorflow/contrib/lite/interpreter.h" -#include "tensorflow/contrib/lite/kernels/internal/tensor_utils.h" -#include "tensorflow/contrib/lite/kernels/register.h" -#include "tensorflow/contrib/lite/kernels/test_util.h" -#include "tensorflow/contrib/lite/model.h" - -namespace tflite { - -namespace ops { -namespace builtin { - -TfLiteRegistration* Register_FULLY_CONNECTED_REF(); -TfLiteRegistration* Register_FULLY_CONNECTED_NEON_OPT(); -TfLiteRegistration* Register_FULLY_CONNECTED_GENERIC_OPT(); -TfLiteRegistration* Register_FULLY_CONNECTED_PIE(); - -} // namespace builtin -} // namespace ops - -namespace { - -using ::testing::ElementsAre; -using ::testing::ElementsAreArray; - -static float fully_connected_input[] = { - 0.503691, 0.196961, 0.521017, 0.554248, 0.288678, 0.792476, 0.561653, - 0.462230, 0.650736, 0.163132, 0.029658, 0.411544, 0.470539, 0.572390, - 0.538755, 0.212030, 0.264309, 0.193908, 0.777480, 0.745661, 0.423314, - 0.470804, 0.175501, 0.492225, 0.192743, 0.540183, 0.372514, 0.446550, - 0.498173, 0.126472, 0.132706, 0.001864, 0.323433, 0.653723, 0.556112, - 0.612111, 0.446199, 0.117765, 0.074341, 0.096935, 0.280897, 0.103999, - 0.508479, 0.751437, 0.676389, 0.047234, 0.963467, 0.940698, 0.241142, - 0.740947, 0.686359, 0.664456, 0.211751, 0.861860, 0.156681, 0.404494, - 0.402043, 0.529195, 0.851044, 0.900216, 0.655667, 0.983750, 0.902081, - 0.979100, 0.637473, 0.458193, 0.591211, 0.083671, 0.575958, 0.665552, - 0.180606, 0.856856, 0.769551, 0.689086, 0.608293, 0.445940, 0.736320, - 0.571760, 0.386637, 0.977461, 0.312707, 0.072996, 0.641918, 0.524458, - 0.934856, 0.798598, 0.928951, 0.336899, 0.327793, 0.779995, 0.237115, - 0.983460, 0.763746, 0.139196, 0.962560, 0.401218, 0.597389, 0.553771, - 0.484890, 0.173347, 0.219322, 0.665496, 0.030203, 0.988873, 0.354582, - 0.638496, 0.434813, 0.090902, 0.210256, 0.821450, 0.068363, 0.522962, - 0.894446, 0.710280, 0.047420, 0.829302, 0.508879, 0.976371, 0.166202, - 0.836672, 0.756367, 0.403317, 0.820132, 0.520112, 0.542513, 0.782691, - 0.921330, 0.139902}; - -static float fully_connected_golden_output[] = { - 0, 0.0732134, 0, 0, 0, 0.280859, - 0, 0.128927, 0, 0.0777251, 0, 0.270268, - 0.271435, 0.0173503, 0.335465, 0.235562, - - 0, 0.0745866, 0, 0.051611, 0, 0.253876, - 0, 0.0814873, 0, 0.104104, 0, 0.248529, - 0.264194, 0, 0.302973, 0.166252, - - 0, 0.0170409, 0, 0.0509851, 0, 0.212834, - 0, 0.0208326, 0, 0.129932, 0.203978, 0.103428, - 0.298051, 0, 0.332233, 0.00445903, - - 0, 0.125246, 0, 0.0735336, 0, 0.0910256, - 0, 0, 0, 0.18933, 0.378111, 0.0712443, - 0.277298, 0.0123414, 0.267454, 0, - - 0, 0.14687, 0, 0.155495, 0.0300215, 0.147256, - 0, 0, 0, 0.156412, 0.434914, 0.0461529, - 0.246508, 0, 0.363138, 0, - - 0, 0, 0, 0.0212949, 0, 0.301708, - 0, 0.35497, 0, 0.406223, 0.0260211, 0.049195, - 0.197161, 0, 0.37316, 0, - - 0, 0.221783, 0, 0, 0.0116515, 0.281945, - 0, 0, 0, 0, 0.285626, 0.181773, - 0.296401, 0.170452, 0.367135, 0.142597, - - 0, 0, 0, 0, 0, 0.418886, - 0, 0.291063, 0, 0.227541, 0.0424759, 0.27589, - 0.398286, 0.177146, 0.40359, 0.121452, - - 0, 0.0834884, 0, 0, 0, 0.287441, - 0, 0.0046838, 0, 0.0122087, 0, 0.217376, - 0.140183, 0.0948412, 0.436677, 0.0589876, - - 0, 0.0289969, 0, 0.0921397, 0, 0.396802, - 0, 0.0126157, 0, 0.0968433, 0, 0.172271, - 0.173295, 0.0664741, 0.53645, 0.00915603, - - 0, 0, 0, 0, 0, 0.147942, - 0, 0.263795, 0, 0.39782, 0, 0.382435, - 0.561072, 0.0579847, 0.145712, 0.13508, - - 0, 0, 0, 0.16382, 0, 0.322294, - 0, 0.163798, 0, 0.405211, 0.367953, 0.076852, - 0.342473, 0.0834118, 0.377537, 0, - - 0, 0.206, 0, 0, 0, 0.375769, - 0, 0, 0, 0, 0, 0.125165, - 0, 0.105591, 0.52055, 0.0536445, - - 0, 0.259261, 0, 0, 0, 0.247707, - 0, 0, 0, 0, 0, 0.215862, - 0.149153, 0.224678, 0.359519, 0.129419, - - 0, 0.17611, 0, 0.280895, 0, 0.576484, - 0, 0.000418848, 0, 0, 0, 0.151112, - 0.211902, 0, 0.566341, 0.106305, - - 0, 0.0246284, 0, 0, 0, 0.196267, - 0, 0.0248624, 0, 0.265635, 0, 0.436199, - 0.408079, 0.134514, 0.328489, 0.411368}; - -class BaseFullyConnectedOpModel : public SingleOpModel { - public: - // TODO(ahentz): test different activation types too. - BaseFullyConnectedOpModel( - TfLiteRegistration* registration, int units, int batches, - const TensorData& input, const TensorData& output = {TensorType_FLOAT32}, - ActivationFunctionType activation_func = ActivationFunctionType_RELU, - FullyConnectedOptionsWeightsFormat weights_format = - FullyConnectedOptionsWeightsFormat_DEFAULT) - : batches_(batches), units_(units) { - int total_input_size = 1; - for (int i = 0; i < input.shape.size(); ++i) { - total_input_size *= input.shape[i]; - } - input_size_ = total_input_size / batches_; - - input_ = AddInput(input); - weights_ = - AddInput({input.type, {units_, input_size_}, input.min, input.max}); - - if (input.type == TensorType_FLOAT32) { - bias_ = AddInput({TensorType_FLOAT32, {units_}}); - } else { - // This is a quantized version. The scale of 'bias' depends on the scales - // of input and filter. Supposedly this is correctly set during quantized - // training. - auto bias_scale = GetScale(input_) * GetScale(weights_); - TensorData bias{TensorType_INT32, {units_}, 0, 0, bias_scale}; - bias_ = AddInput(bias); - } - - output_ = AddOutput(output); - if (weights_format != FullyConnectedOptionsWeightsFormat_DEFAULT) { - AddOutput({TensorType_UINT8, input.shape}); - } - - SetBuiltinOp( - BuiltinOperator_FULLY_CONNECTED, BuiltinOptions_FullyConnectedOptions, - CreateFullyConnectedOptions(builder_, activation_func, weights_format) - .Union()); - resolver_ = absl::make_unique( - BuiltinOperator_FULLY_CONNECTED, registration); - BuildInterpreter({GetShape(input_), GetShape(weights_), GetShape(bias_)}); - } - - int input_size() { return input_size_; } - int num_units() { return units_; } - int num_batches() { return batches_; } - - protected: - int input_; - int weights_; - int bias_; - int output_; - - int batches_; - int units_; - int input_size_; -}; - -class FloatFullyConnectedOpModel : public BaseFullyConnectedOpModel { - public: - using BaseFullyConnectedOpModel::BaseFullyConnectedOpModel; - - void SetBias(const std::vector& f) { PopulateTensor(bias_, f); } - - void SetWeights(const std::vector& f) { PopulateTensor(weights_, f); } - - void SetInput(const std::vector& data) { - PopulateTensor(input_, data); - } - void SetInput(int offset, float* begin, float* end) { - PopulateTensor(input_, offset, begin, end); - } - - std::vector GetOutput() { return ExtractVector(output_); } -}; - -class QuantizedFullyConnectedOpModel : public BaseFullyConnectedOpModel { - public: - using BaseFullyConnectedOpModel::BaseFullyConnectedOpModel; - - void SetBias(const std::vector& data) { - QuantizeAndPopulate(bias_, data); - } - void SetWeights(const std::vector& data) { - QuantizeAndPopulate(weights_, data); - } - void ShuffleAndSetWeights(const std::vector& data, int input_depth, - int output_depth) { - std::vector shuffled_data(data.size()); - CHECK_EQ(input_depth % 16, 0); - CHECK_EQ(output_depth % 4, 0); - float* shuffled_data_ptr = shuffled_data.data(); - for (int block_o = 0; block_o < output_depth; block_o += 4) { - for (int block_i = 0; block_i < input_depth; block_i += 16) { - for (int o = 0; o < 4; o++) { - for (int i = 0; i < 16; i++) { - *shuffled_data_ptr++ = - data[(block_o + o) * input_depth + block_i + i]; - } - } - } - } - TfLiteTensor* t = interpreter_->tensor(weights_); - auto quantized_data = - Quantize(shuffled_data, t->params.scale, t->params.zero_point); - for (uint8_t& q : quantized_data) { - q ^= 0x80; - } - PopulateTensor(weights_, 0, quantized_data.data(), - quantized_data.data() + quantized_data.size()); - } - void SetInput(const std::vector& data) { - QuantizeAndPopulate(input_, data); - } - - template - std::vector GetOutput() { - return ExtractVector(output_); - } - - template - std::vector GetDequantizedOutput() { - return Dequantize(ExtractVector(output_), GetScale(output_), - GetZeroPoint(output_)); - } -}; - -// In the hybrid model the weights are quantized (to uint8). But the bias, -// input (and output) are expected to be in float precision. -class HybridFullyConnectedOpModel : public SingleOpModel { - public: - HybridFullyConnectedOpModel(int units, int batches, const TensorData& input, - const TensorData& weights, - const TensorData& output = {TensorType_FLOAT32}) - : batches_(batches), units_(units) { - int total_input_size = 1; - for (int i = 0; i < input.shape.size(); ++i) { - total_input_size *= input.shape[i]; - } - input_size_ = total_input_size / batches_; - - input_ = AddInput(input); - weights_ = AddInput(weights); - - TensorData bias{TensorType_FLOAT32, {units_}}; - bias_ = AddInput(bias); - - output_ = AddOutput(output); - - SetBuiltinOp( - BuiltinOperator_FULLY_CONNECTED, BuiltinOptions_FullyConnectedOptions, - CreateFullyConnectedOptions(builder_, ActivationFunctionType_RELU) - .Union()); - resolver_ = absl::make_unique( - BuiltinOperator_FULLY_CONNECTED, - ops::builtin::Register_FULLY_CONNECTED_PIE()); - BuildInterpreter({GetShape(input_), GetShape(weights_), GetShape(bias_)}); - } - void SetBias(const std::vector& f) { PopulateTensor(bias_, f); } - void SetWeights(const std::vector& data) { - SymmetricQuantizeAndPopulate(weights_, data); - } - - void SetInput(const std::vector& f) { PopulateTensor(input_, f); } - std::vector GetOutput() { return ExtractVector(output_); } - - int input_size() { return input_size_; } - int num_units() { return units_; } - int num_batches() { return batches_; } - - protected: - int input_; - int weights_; - int bias_; - int output_; - - int batches_; - int units_; - int input_size_; -}; - -const auto kKernelMap = new std::map({ - {"Reference", ops::builtin::Register_FULLY_CONNECTED_REF()}, - {"NeonOptimized", ops::builtin::Register_FULLY_CONNECTED_NEON_OPT()}, - {"GenericOptimized", ops::builtin::Register_FULLY_CONNECTED_GENERIC_OPT()}, - {"Pie", ops::builtin::Register_FULLY_CONNECTED_PIE()}, -}); - -class FloatFullyConnectedOpTest : public SingleOpTest { - protected: - const std::map& GetKernelMap() override { - return *kKernelMap; - } -}; - -const auto kKernelMapNoPie = new std::map({ - {"Reference", ops::builtin::Register_FULLY_CONNECTED_REF()}, - {"NeonOptimized", ops::builtin::Register_FULLY_CONNECTED_NEON_OPT()}, - {"GenericOptimized", ops::builtin::Register_FULLY_CONNECTED_GENERIC_OPT()}, -}); - -class QuantizedFullyConnectedOpTest : public SingleOpTest { - protected: - const std::map& GetKernelMap() override { - return *kKernelMapNoPie; - } -}; - -const auto kKernelMapPie = new std::map({ - {"Pie", ops::builtin::Register_FULLY_CONNECTED_PIE()}, -}); - -// Hybrid mode is used by the Pie quantized kernel. -class HybridFullyConnectedOpTest : public SingleOpTest { - protected: - const std::map& GetKernelMap() override { - return *kKernelMapPie; - } -}; - -// TODO(ahentz): add more small tests like this one, focused on making sure the -// calculations are correct. -TEST_P(FloatFullyConnectedOpTest, SimpleTest) { - FloatFullyConnectedOpModel m(GetRegistration(), /*units=*/3, /*batches=*/2, - /*input=*/{TensorType_FLOAT32, {2, 10}}); - m.SetWeights({ - 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, // u = 0 - 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, // u = 1 - 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, // u = 1 - }); - m.SetBias({1, 2, 3}); - - m.SetInput({ - 1, 2, 3, 4, 5, 6, 7, 8, -9, -10, // b = 0 - 1, 2, 3, 4, 5, 6, 7, -8, 9, -10, // b = 1 - }); - - m.Invoke(); - - EXPECT_THAT(m.GetOutput(), ElementsAre(24, 25, 26, 58, 59, 60)); -} - -TEST_P(FloatFullyConnectedOpTest, SimpleTest2) { - FloatFullyConnectedOpModel m(GetRegistration(), /*units=*/1, /*batches=*/2, - /*input=*/{TensorType_FLOAT32, {2, 2}}); - m.SetWeights({ - 2, 4, // u = 0 - }); - m.SetBias({1}); - - m.SetInput({ - 1, 2, // b = 0 - 2, 1, // b = 1 - }); - - m.Invoke(); - - EXPECT_THAT(m.GetOutput(), ElementsAre(11, 9)); -} - -TEST_P(QuantizedFullyConnectedOpTest, SimpleTestQuantized) { - QuantizedFullyConnectedOpModel m( - GetRegistration(), /*units=*/3, /*batches*/ 2, - /*input=*/{TensorType_UINT8, {2, 10}, -63.5, 64}, - /*output=*/{TensorType_UINT8, {}, -127, 128}); - - // input_product_scale < output_scale was not true. - m.SetWeights({ - 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, // u = 0 - 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, // u = 1 - 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, // u = 2 - }); - m.SetBias({1, 2, 3}); - - m.SetInput({ - 1, 2, 3, 4, 5, 6, 7, 8, -9, -10, // b = 0 - 1, 2, 3, 4, 5, 6, 7, -8, 9, -10, // b = 1 - }); - - m.Invoke(); - - EXPECT_THAT(m.GetDequantizedOutput(), - ElementsAreArray(ArrayFloatNear({ - 24, 25, 26, // - 58, 59, 60, // - }))); - EXPECT_THAT(m.GetOutput(), - ElementsAre(151, 152, 153, 185, 186, 187)); -} - -TEST_P(QuantizedFullyConnectedOpTest, - SimpleTestQuantizedOutputMultiplierGreaterThan1) { - // real_multiplier = 2. - QuantizedFullyConnectedOpModel m( - GetRegistration(), /*units=*/3, /*batches*/ 2, - /*input=*/{TensorType_UINT8, {2, 10}, -127, 128}, - /*output=*/{TensorType_UINT8, {}, -63.5, 64}); - - m.SetWeights({ - 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, // u = 0 - 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, // u = 1 - 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, // u = 2 - }); - m.SetBias({1, 2, 3}); - - m.SetInput({ - 1, 2, 3, 4, 5, 6, 7, 8, -9, -10, // b = 0 - 1, 2, 3, 4, 5, 6, 7, -8, 9, -10, // b = 1 - }); - - m.Invoke(); - - EXPECT_THAT(m.GetDequantizedOutput(), - ElementsAreArray(ArrayFloatNear({ - 24, 25, 26, // first batch - 58, 59, 60, // second batch - }))); - EXPECT_THAT(m.GetOutput(), - ElementsAre(175, 177, 179, 243, 245, 247)); -} - -void SimpleTestQuantizedInt16OutputCase( - TfLiteRegistration* registration, int input_depth, int output_depth, - int batches, FullyConnectedOptionsWeightsFormat weights_format) { - const uint8_t kWeightsZeroPoint = 128; - const float kWeightsScale = 1.f / 128.f; - const uint8_t kInputZeroPoint = 128; - const float kInputScale = 1.f / 128.f; - const float kInputMin = (0 - kInputZeroPoint) * kInputScale; - const float kInputMax = (255 - kInputZeroPoint) * kInputScale; - // Output ranges in [-8..8] encoded as int16 - const float kOutputScale = 8.f / 32768.f; - const float kOutputMin = -32768 * kOutputScale; - const float kOutputMax = 32767 * kOutputScale; - - QuantizedFullyConnectedOpModel m( - registration, output_depth, batches, - /*input=*/ - {TensorType_UINT8, {batches, input_depth}, kInputMin, kInputMax}, - /*output=*/{TensorType_INT16, {}, kOutputMin, kOutputMax}, - /*activation_func=*/ActivationFunctionType_NONE, weights_format); - - std::mt19937 random_engine; - std::uniform_int_distribution weights_dist; - - std::vector weights_data(input_depth * output_depth); - for (auto& w : weights_data) { - uint8_t q = weights_dist(random_engine); - w = (q - kWeightsZeroPoint) * kWeightsScale; - } - - // Based on weights_format, enforce any shape requirement for that format/path - // and set the (possibly shuffled) weights. - switch (weights_format) { - case FullyConnectedOptionsWeightsFormat_DEFAULT: - m.SetWeights(weights_data); - break; - case FullyConnectedOptionsWeightsFormat_SHUFFLED4x16INT8: - // The shuffled path currently supports only a restrictive subset of - // shapes, described by the following assertions: - CHECK_EQ(input_depth % 16, 0); - CHECK_EQ(output_depth % 4, 0); - CHECK(batches == 1 || batches == 4); - m.ShuffleAndSetWeights(weights_data, input_depth, output_depth); - break; - default: - LOG(FATAL) << "Unhandled weights format"; - } - - std::uniform_int_distribution input_dist; - std::vector input_data(input_depth * batches); - for (auto& i : input_data) { - uint8_t q = input_dist(random_engine); - i = (q - kInputZeroPoint) * kInputScale; - } - - std::vector bias_data(output_depth); - // As the output ranges in [-8, 8], it's reasonable to have bias values - // in [-1, 1], this won't result in too much saturation. - std::uniform_real_distribution bias_dist(-1.f, 1.f); - for (auto& b : bias_data) { - b = bias_dist(random_engine); - } - - m.SetBias(bias_data); - m.SetInput(input_data); - - m.Invoke(); - - std::vector expected_output_data(output_depth * batches); - for (int b = 0; b < batches; b++) { - for (int o = 0; o < output_depth; o++) { - float accum = bias_data[o]; - for (int i = 0; i < input_depth; i++) { - accum += - input_data[b * input_depth + i] * weights_data[o * input_depth + i]; - } - accum = std::min(accum, kOutputMax); - accum = std::max(accum, kOutputMin); - expected_output_data[b * output_depth + o] = accum; - } - } - - EXPECT_THAT(m.GetDequantizedOutput(), - ElementsAreArray(ArrayFloatNear(expected_output_data, 3e-4f))); -} - -TEST_P(QuantizedFullyConnectedOpTest, - SimpleTestQuantizedInt16OutputDefaultWeights) { - for (int input_depth : {1, 3, 10, 100}) { - for (int output_depth : {1, 3, 10, 100}) { - for (int batch : {1, 3, 10, 100}) { - SimpleTestQuantizedInt16OutputCase( - GetRegistration(), input_depth, output_depth, batch, - FullyConnectedOptionsWeightsFormat_DEFAULT); - } - } - } -} - -TEST_P(QuantizedFullyConnectedOpTest, - SimpleTestQuantizedInt16OutputShuffled4x16Int8Weights) { - // The shuffled weights block shape is 4x16. The shape of the weights matrix - // is: rows = output_depth, cols = input_depth. It must be a multiple of 4x16. - // This means that output_depth must be a multiple of 4, and input_deth must - // be a multiple of 16. - for (int input_depth_numblocks : {1, 3}) { - for (int output_depth_numblocks : {1, 3}) { - int input_depth = 16 * input_depth_numblocks; - int output_depth = 4 * output_depth_numblocks; - // The fast shuffled path is currently supporting only batch sizes of 1 - // and 4. The idea is that the whole point of that path is to go as fast - // as possible for small batch size, which requires fully specializing - // it for each batch size, and for larger batch sizes the generic - // gemmlowp-based implementation is fast enough. - for (int batch : {1, 4}) { - SimpleTestQuantizedInt16OutputCase( - GetRegistration(), input_depth, output_depth, batch, - FullyConnectedOptionsWeightsFormat_SHUFFLED4x16INT8); - } - } - } -} - -TEST(HybridFullyConnectedOpTest, SimpleTestQuantized) { - HybridFullyConnectedOpModel m( - /*units=*/3, /*batches=*/2, - /*input=*/{TensorType_FLOAT32, {2, 10}}, - /*weights=*/{TensorType_UINT8, {3, 10}, -63.5, 64}); // PIE - - m.SetWeights({ - 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, // u = 0 - 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, // u = 1 - 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, // u = 1 - }); - m.SetBias({1, 2, 3}); - - m.SetInput({ - 1, 2, 3, 4, 5, 6, 7, 8, -9, -10, // b = 0 - 1, 2, 3, 4, 5, 6, 7, -8, 9, -10, // b = 1 - }); - - m.Invoke(); - - EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear( - { - 24, 25, 26, // - 58, 59, 60, // - }, - /*max_abs_error=*/1.3f))); -} - -TEST_P(FloatFullyConnectedOpTest, SimpleTest4DInput) { - // Note that it is not required that the first dimension be the number of - // batches. All we care is that the input can be evenly distributed in - // batches. In this case, we need the input to have multiples of '2'. - FloatFullyConnectedOpModel m(GetRegistration(), - /*units=*/3, /*batches=*/2, - /*input=*/{TensorType_FLOAT32, {4, 1, 5, 1}}); - m.SetWeights({ - 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, // u = 0 - 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, // u = 1 - 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, // u = 1 - }); - m.SetBias({1, 2, 3}); - - m.SetInput({ - 1, 2, 3, 4, 5, 6, 7, 8, -9, -10, // first batch - 1, 2, 3, 4, 5, 6, 7, -8, 9, -10, // second batch - }); - - m.Invoke(); - - EXPECT_THAT(m.GetOutput(), ElementsAreArray({ - 24, 25, 26, // first batch - 58, 59, 60, // second batch - })); -} - -TEST_P(QuantizedFullyConnectedOpTest, SimpleTest4dInputQuantized) { - QuantizedFullyConnectedOpModel m( - GetRegistration(), /*units=*/3, /*batches=*/2, - /*input=*/{TensorType_UINT8, {4, 1, 5, 1}, -63.5, 64}, - /*output=*/{TensorType_UINT8, {}, -127, 128}); - - // input_product_scale < output_scale was not true. - m.SetWeights({ - 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, // u = 0 - 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, // u = 1 - 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, // u = 1 - }); - m.SetBias({1, 2, 3}); - - m.SetInput({ - 1, 2, 3, 4, 5, 6, 7, 8, -9, -10, // b = 0 - 1, 2, 3, 4, 5, 6, 7, -8, 9, -10, // b = 1 - }); - - m.Invoke(); - - EXPECT_THAT(m.GetDequantizedOutput(), - ElementsAreArray(ArrayFloatNear({ - 24, 25, 26, // - 58, 59, 60, // - }))); - EXPECT_THAT(m.GetOutput(), - ElementsAre(151, 152, 153, 185, 186, 187)); -} - -TEST_P(QuantizedFullyConnectedOpTest, - SimpleTest4dInputQuantizedOutputMultiplierGreaterThan1) { - // real_multiplier = 2. - QuantizedFullyConnectedOpModel m( - GetRegistration(), /*units=*/3, /*batches=*/2, - /*input=*/{TensorType_UINT8, {4, 1, 5, 1}, -127, 128}, - /*output=*/{TensorType_UINT8, {}, -63.5, 64}); - - m.SetWeights({ - 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, // u = 0 - 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, // u = 1 - 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, // u = 1 - }); - m.SetBias({1, 2, 3}); - - m.SetInput({ - 1, 2, 3, 4, 5, 6, 7, 8, -9, -10, // b = 0 - 1, 2, 3, 4, 5, 6, 7, -8, 9, -10, // b = 1 - }); - - m.Invoke(); - - EXPECT_THAT(m.GetDequantizedOutput(), - ElementsAreArray(ArrayFloatNear({ - 24, 25, 26, // first batch - 58, 59, 60, // second batch - }))); - EXPECT_THAT(m.GetOutput(), - ElementsAre(175, 177, 179, 243, 245, 247)); -} - -INSTANTIATE_TEST_CASE_P( - FloatFullyConnectedOpTest, FloatFullyConnectedOpTest, - ::testing::ValuesIn(SingleOpTest::GetKernelTags(*kKernelMap))); - -INSTANTIATE_TEST_CASE_P( - QuantizedFullyConnectedOpTest, QuantizedFullyConnectedOpTest, - ::testing::ValuesIn(SingleOpTest::GetKernelTags(*kKernelMapNoPie))); - -// TODO(ahentz): Reconsider this test. Having arbitrary weights makes it hard -// to debug errors and doesn't necessarily test all the important details. -TEST_P(FloatFullyConnectedOpTest, BlackBoxTest) { - FloatFullyConnectedOpModel m(GetRegistration(), /*units=*/16, /*batches=*/2, - /*input=*/{TensorType_FLOAT32, {2, 8}}); - m.SetWeights( - {0.091327, 0.103366, -0.316505, -0.083120, 0.149366, -0.196636, - -0.123672, 0.062800, 0.063031, 0.191670, -0.062001, -0.061504, - -0.275581, 0.059388, -0.118497, -0.079224, 0.109758, 0.008307, - -0.062657, -0.060962, -0.049782, -0.106719, -0.319482, -0.103650, - 0.266455, 0.051517, -0.123448, 0.322464, 0.043282, -0.173782, - -0.190381, 0.002013, 0.096086, 0.131157, 0.031164, 0.100638, - -0.312191, -0.080923, -0.101318, -0.116614, 0.142238, 0.086540, - -0.139154, 0.174268, -0.073161, 0.080072, 0.006874, 0.229382, - -0.104321, -0.176035, -0.208587, -0.001019, -0.162032, 0.080824, - -0.025021, 0.074460, -0.252595, -0.161750, -0.136403, 0.008308, - 0.005710, 0.096600, 0.289839, 0.218816, -0.304651, -0.070958, - 0.054598, 0.147113, -0.139112, -0.072798, -0.163335, -0.167863, - -0.128762, -0.035780, 0.117262, 0.017177, 0.263335, -0.176612, - 0.262961, -0.093654, -0.339283, 0.333071, 0.180827, 0.287583, - 0.066350, -0.197947, -0.114449, -0.236035, 0.103532, -0.034284, - 0.093299, -0.145361, 0.054001, 0.250570, 0.157010, -0.143480, - -0.139061, -0.048873, 0.067557, 0.139038, 0.324106, 0.227041, - 0.037793, -0.225747, -0.241619, 0.357835, 0.135762, -0.306764, - -0.125982, 0.091916, 0.266587, 0.030135, 0.265148, 0.141627, - 0.020120, 0.083815, -0.124556, -0.100124, -0.048159, 0.181172, - 0.302309, -0.041084, 0.146334, -0.061511, -0.232605, 0.281324, - 0.145408, -0.221897}); - m.SetBias({-0.160594, 0.205770, -0.078307, -0.077984, 0.001937, 0.015860, - 0.036810, 0.012346, 0.001028, 0.038551, 0.075415, 0.020804, - 0.048478, -0.032270, 0.175688, -0.085662}); - - const int input_sequence_size = sizeof(fully_connected_input) / - sizeof(float) / - (m.input_size() * m.num_batches()); - for (int i = 0; i < input_sequence_size; i++) { - // TODO(ahentz): This is what the original test was doing: two equal - // batches per invocation. We could instead use two different batches. - float* batch_start = fully_connected_input + i * m.input_size(); - float* batch_end = batch_start + m.input_size(); - m.SetInput(0, batch_start, batch_end); - m.SetInput(m.input_size(), batch_start, batch_end); - - m.Invoke(); - - float* golden_start = fully_connected_golden_output + i * m.num_units(); - float* golden_end = golden_start + m.num_units(); - std::vector expected; - expected.insert(expected.end(), golden_start, golden_end); - expected.insert(expected.end(), golden_start, golden_end); - - EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear(expected))); - } -} - -} // namespace -} // namespace tflite - -int main(int argc, char** argv) { - ::tflite::LogToStderr(); - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/contrib/lite/kernels/gather.cc b/tensorflow/contrib/lite/kernels/gather.cc deleted file mode 100644 index b5afeb1a7bd5528328bd5585d9696b3362cbe3a3..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/gather.cc +++ /dev/null @@ -1,135 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include -#include "tensorflow/contrib/lite/c/builtin_op_data.h" -#include "tensorflow/contrib/lite/c/c_api_internal.h" -#include "tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h" -#include "tensorflow/contrib/lite/kernels/internal/tensor.h" -#include "tensorflow/contrib/lite/kernels/kernel_util.h" -#include "tensorflow/contrib/lite/kernels/op_macros.h" -#include "tensorflow/contrib/lite/string_util.h" - -namespace tflite { -namespace ops { -namespace builtin { -namespace gather { -constexpr int kInputTensor = 0; -constexpr int kInputPositions = 1; -constexpr int kOutputTensor = 0; - -TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { - TF_LITE_ENSURE_EQ(context, NumInputs(node), 2); - TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); - - const auto* params = - reinterpret_cast(node->builtin_data); - const TfLiteTensor* input = GetInput(context, node, kInputTensor); - const TfLiteTensor* positions = GetInput(context, node, kInputPositions); - TfLiteTensor* output = GetOutput(context, node, kOutputTensor); - // Only INT32 positions are supported. - TF_LITE_ENSURE_EQ(context, positions->type, kTfLiteInt32); - // Assign to output the input type. - output->type = input->type; - // TODO(mgubin): Only default axis == 0 is supported. - TF_LITE_ENSURE_EQ(context, params->axis, 0); - // Check conditions for different types. - switch (input->type) { - case kTfLiteFloat32: - case kTfLiteUInt8: - case kTfLiteInt32: { - // Fully supported by reference_ops::Gather. - } break; - - case kTfLiteString: { - // Only 1D input is supported. - TF_LITE_ENSURE_EQ(context, NumDimensions(input), 1); - } break; - default: - context->ReportError( - context, "Only float32 and string types are supported, got %d", - input->type); - return kTfLiteError; - } - const int num_dimensions = - NumDimensions(input) + NumDimensions(positions) - 1; - TF_LITE_ENSURE(context, params->axis <= num_dimensions); - TfLiteIntArray* output_shape = TfLiteIntArrayCreate(num_dimensions); - int output_index = 0; - for (int i = 0; i < params->axis; ++i) { - output_shape->data[output_index++] = input->dims->data[i]; - } - for (int i = 0; i < positions->dims->size; ++i) { - output_shape->data[output_index++] = positions->dims->data[i]; - } - for (int i = params->axis + 1; i < input->dims->size; ++i) { - output_shape->data[output_index++] = input->dims->data[i]; - } - return context->ResizeTensor(context, output, output_shape); -} - -TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { - const TfLiteTensor* input = GetInput(context, node, kInputTensor); - const TfLiteTensor* positions = GetInput(context, node, kInputPositions); - TfLiteTensor* output = GetOutput(context, node, kOutputTensor); - const int input_rank = NumDimensions(input); -#define TF_LITE_GATHER(data_type, index_type) \ - { \ - tflite::GatherParams op_params; \ - op_params.input_rank = input_rank; \ - optimized_ops::Gather( \ - op_params, GetTensorShape(input), GetTensorData(input), \ - GetTensorShape(positions), GetTensorData(positions), \ - GetTensorShape(output), GetTensorData(output)); \ - } - switch (input->type) { - case kTfLiteFloat32: - TF_LITE_GATHER(float, int32_t); - break; - case kTfLiteUInt8: - TF_LITE_GATHER(uint8_t, int32_t); - break; - case kTfLiteInt32: - TF_LITE_GATHER(int32_t, int32_t); - break; - case kTfLiteString: { - // TODO(mgubin): Currently support only for 1D output tensors. - DynamicBuffer buffer; - const int32* indexes = positions->data.i32; - const int num_strings = GetStringCount(input); - for (int i = 0; i < positions->dims->data[0]; ++i) { - const int pos = indexes[i]; - TF_LITE_ENSURE(context, pos < num_strings); - const auto string_ref = GetString(input, pos); - buffer.AddString(string_ref.str, string_ref.len); - } - buffer.WriteToTensor(output); - } break; - default: - return kTfLiteError; - } -#undef TF_LITE_GATHER - return kTfLiteOk; -} -} // namespace gather - -TfLiteRegistration* Register_GATHER() { - static TfLiteRegistration r = {nullptr, nullptr, gather::Prepare, - gather::Eval}; - return &r; -} - -} // namespace builtin -} // namespace ops -} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/gather_test.cc b/tensorflow/contrib/lite/kernels/gather_test.cc deleted file mode 100644 index 1b48884e0907c67919f65680ab2f096481551eb7..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/gather_test.cc +++ /dev/null @@ -1,150 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include -#include "tensorflow/contrib/lite/c/builtin_op_data.h" -#include "tensorflow/contrib/lite/interpreter.h" -#include "tensorflow/contrib/lite/kernels/register.h" -#include "tensorflow/contrib/lite/kernels/test_util.h" -#include "tensorflow/contrib/lite/model.h" - -namespace tflite { -namespace { - -using ::testing::ElementsAreArray; - -class GatherOpModel : public SingleOpModel { - public: - GatherOpModel(std::initializer_list input_shape, TensorType input_type, - std::initializer_list positions_shape) { - input_ = AddInput(input_type); - positions_ = AddInput(TensorType_INT32); - output_ = AddOutput(input_type); - SetBuiltinOp(BuiltinOperator_GATHER, BuiltinOptions_GatherOptions, - CreateGatherOptions(builder_, 0).Union()); - BuildInterpreter({input_shape, positions_shape}); - } - - void SetInputFloat(std::initializer_list data) { - PopulateTensor(input_, data); - } - - void SetInputUint8(std::initializer_list data) { - PopulateTensor(input_, data); - } - - void SetInput(std::initializer_list data) { - PopulateStringTensor(input_, data); - } - - void SetPositions(std::initializer_list data) { - PopulateTensor(positions_, data); - } - - std::vector GetOutputFloat() { return ExtractVector(output_); } - std::vector GetOutputUint8() { - return ExtractVector(output_); - } - std::vector GetOutputString() { - return ExtractVector(output_); - } - std::vector GetOutputShape() { return GetTensorShape(output_); } - - protected: - int input_; - int positions_; - int output_; -}; - -TEST(GatherOpTest, Shuffle) { - GatherOpModel m({2, 2}, TensorType_FLOAT32, {2}); - m.SetInputFloat({-2.0, 0.2, 0.7, 0.8}); - m.SetPositions({1, 0}); - m.Invoke(); - EXPECT_THAT(m.GetOutputFloat(), - ElementsAreArray(ArrayFloatNear({0.7, 0.8, -2, 0.2}))); -} - -TEST(GatherOpTest, Test0DIndex) { - GatherOpModel m({2, 2}, TensorType_FLOAT32, {}); - m.SetInputFloat({-2.0, 0.2, 0.7, 0.8}); - m.SetPositions({1}); - m.Invoke(); - EXPECT_THAT(m.GetOutputFloat(), ElementsAreArray(ArrayFloatNear({0.7, 0.8}))); - EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2})); -} - -TEST(GatherOpTest, Test0DIndexWith0DResult) { - // 0D tensor is special case in current TFLite. Test it once to make sure - // existing workarounds are fine with it. - GatherOpModel m({3}, TensorType_FLOAT32, {}); - m.SetInputFloat({1.0, 2.0, 3.0}); - m.SetPositions({1}); - m.Invoke(); - EXPECT_THAT(m.GetOutputFloat(), ElementsAreArray(ArrayFloatNear({2.0}))); - EXPECT_TRUE(m.GetOutputShape().empty()); -} - -TEST(GatherOpTest, Test2DIndexWith2DResult) { - GatherOpModel m({3}, TensorType_FLOAT32, {1, 2}); - m.SetInputFloat({1.0, 2.0, 3.0}); - m.SetPositions({1, 0}); - m.Invoke(); - EXPECT_THAT(m.GetOutputFloat(), ElementsAreArray(ArrayFloatNear({2.0, 1.0}))); - EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 2})); -} - -TEST(FloatGatherOpTest, Duplicate) { - GatherOpModel m({1, 2, 2}, TensorType_FLOAT32, {2}); - m.SetInputFloat({-2.0, 0.2, 0.7, 0.8}); - m.SetPositions({0, 0}); - m.Invoke(); - EXPECT_THAT( - m.GetOutputFloat(), - ElementsAreArray(ArrayFloatNear({-2, 0.2, 0.7, 0.8, -2, 0.2, 0.7, 0.8}))); -} - -TEST(FloatGatherOpTest, Slice) { - GatherOpModel m({4, 1}, TensorType_FLOAT32, {2}); - m.SetInputFloat({-2.0, 0.2, 0.7, 0.8}); - m.SetPositions({1, 3}); - m.Invoke(); - EXPECT_THAT(m.GetOutputFloat(), ElementsAreArray(ArrayFloatNear({0.2, 0.8}))); -} - -TEST(Uint8tGatherOpTest, Shuffle) { - GatherOpModel m({2, 2}, TensorType_UINT8, {2}); - m.SetInputUint8({133, 134, 14, 15}); - m.SetPositions({1, 0}); - m.Invoke(); - - EXPECT_THAT(m.GetOutputUint8(), ElementsAreArray({14, 15, 133, 134})); -} - -TEST(GatherOpTest, SimpleString) { - GatherOpModel m({3}, TensorType_STRING, {2}); - m.SetInput({"A", "B", "C"}); - m.SetPositions({0, 2}); - m.Invoke(); - ASSERT_THAT(m.GetOutputShape(), ElementsAreArray({2})); - EXPECT_THAT(m.GetOutputString(), ElementsAreArray({"A", "C"})); -} -} // namespace -} // namespace tflite - -int main(int argc, char** argv) { - ::tflite::LogToStderr(); - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/contrib/lite/kernels/internal/BUILD b/tensorflow/contrib/lite/kernels/internal/BUILD deleted file mode 100644 index 5c9ca6e910ac9c0f0814b99b7d12206405bf1f28..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/internal/BUILD +++ /dev/null @@ -1,694 +0,0 @@ -package(default_visibility = [ - "//visibility:public", -]) - -licenses(["notice"]) # Apache 2.0 - -load("//tensorflow/contrib/lite:build_def.bzl", "tflite_copts") -load("//tensorflow/contrib/lite:special_rules.bzl", "tflite_portable_test_suite") - -tflite_deps_intel = [ - "@arm_neon_2_x86_sse", -] - -HARD_FP_FLAGS_IF_APPLICABLE = select({ - "//tensorflow:android_arm": ["-mfloat-abi=softfp"], - "//tensorflow:android_arm64": ["-mfloat-abi=softfp"], - "//tensorflow:android_armeabi": ["-mfloat-abi=softfp"], - "//conditions:default": [], -}) - -NEON_FLAGS_IF_APPLICABLE = select({ - ":arm": [ - "-O3", - "-mfpu=neon", - ], - ":armeabi-v7a": [ - "-O3", - "-mfpu=neon", - ], - ":armv7a": [ - "-O3", - "-mfpu=neon", - ], - "//conditions:default": [ - "-O3", - ], -}) - -cc_library( - name = "types", - srcs = [], - hdrs = [ - "compatibility.h", - "types.h", - ], - deps = [ - "//tensorflow/contrib/lite/kernels:op_macros", - "@com_google_absl//absl/base:core_headers", - ], -) - -cc_library( - name = "legacy_types", - srcs = [], - hdrs = [ - "compatibility.h", - "legacy_types.h", - "types.h", - ], - deps = [ - "//tensorflow/contrib/lite/kernels:op_macros", - "@com_google_absl//absl/base:core_headers", - ], -) - -config_setting( - name = "arm", - values = { - "cpu": "arm", - }, -) - -config_setting( - name = "arm64-v8a", - values = { - "cpu": "arm64-v8a", - }, -) - -config_setting( - name = "armv7a", - values = { - "cpu": "armv7a", - }, -) - -config_setting( - name = "armeabi-v7a", - values = { - "cpu": "armeabi-v7a", - }, -) - -config_setting( - name = "haswell", - values = { - "cpu": "haswell", - }, -) - -config_setting( - name = "ios_x86_64", - values = { - "cpu": "ios_x86_64", - }, -) - -config_setting( - name = "ios_armv7", - values = { - "cpu": "ios_armv7", - }, -) - -config_setting( - name = "ios_arm64", - values = { - "cpu": "ios_arm64", - }, -) - -config_setting( - name = "k8", - values = { - "cpu": "k8", - }, -) - -config_setting( - name = "x86", - values = { - "cpu": "x86", - }, -) - -config_setting( - name = "x86_64", - values = { - "cpu": "x86_64", - }, -) - -config_setting( - name = "darwin", - values = { - "cpu": "darwin", - }, -) - -config_setting( - name = "darwin_x86_64", - values = { - "cpu": "darwin_x86_64", - }, -) - -config_setting( - name = "freebsd", - values = { - "cpu": "freebsd", - }, -) - -cc_library( - name = "optimized_base", - srcs = [], - hdrs = [ - "common.h", - "optimized/depthwiseconv_float.h", - "optimized/depthwiseconv_uint8.h", - "optimized/depthwiseconv_uint8_3x3_filter.h", - "optimized/optimized_ops.h", - ], - copts = tflite_copts(), - deps = [ - ":quantization_util", - ":strided_slice_logic", - ":types", - ":reference_base", - ":round", - ":tensor_utils", - "//third_party/eigen3", - "@gemmlowp", - "//tensorflow/contrib/lite/c:c_api_internal", - ] + select({ - ":haswell": tflite_deps_intel, - ":ios_x86_64": tflite_deps_intel, - ":k8": tflite_deps_intel, - ":x86": tflite_deps_intel, - ":x86_64": tflite_deps_intel, - ":darwin": tflite_deps_intel, - ":darwin_x86_64": tflite_deps_intel, - ":freebsd": tflite_deps_intel, - "//conditions:default": [], - }), -) - -cc_library( - name = "legacy_optimized_base", - srcs = [], - hdrs = [ - "common.h", - "optimized/depthwiseconv_float.h", - "optimized/depthwiseconv_uint8.h", - "optimized/depthwiseconv_uint8_3x3_filter.h", - "optimized/legacy_optimized_ops.h", - "optimized/optimized_ops.h", - ], - copts = tflite_copts(), - deps = [ - ":quantization_util", - ":strided_slice_logic", - ":tensor_utils", - ":types", - ":legacy_types", - ":legacy_reference_base", - ":round", - "//third_party/eigen3", - "@gemmlowp", - "//tensorflow/contrib/lite/c:c_api_internal", - ] + select({ - ":haswell": tflite_deps_intel, - ":ios_x86_64": tflite_deps_intel, - ":k8": tflite_deps_intel, - ":x86": tflite_deps_intel, - ":x86_64": tflite_deps_intel, - ":darwin": tflite_deps_intel, - ":darwin_x86_64": tflite_deps_intel, - ":freebsd": tflite_deps_intel, - "//conditions:default": [], - }), -) - -cc_library( - name = "optimized", - hdrs = [ - "optimized/cblas_conv.h", - "optimized/cblas_reference.h", - "optimized/eigen_spatial_convolutions.h", - "optimized/eigen_tensor_reduced_instantiations_oss.h", - "optimized/multithreaded_conv.h", - # FIXME(petewarden) - This should be removed, since it's a header from the - # :tensor dependency below. - "tensor.h", - ], - deps = [ - ":optimized_base", - ":tensor", - ":types", - "//tensorflow/contrib/lite/c:c_api_internal", - "//third_party/eigen3", - ], -) - -cc_test( - name = "tensor_test", - srcs = ["tensor_test.cc"], - tags = ["no_oss"], - deps = [ - ":tensor", - "@com_google_googletest//:gtest", - ], -) - -cc_library( - name = "round", - srcs = [], - hdrs = ["round.h"], -) - -cc_library( - name = "quantization_util", - srcs = ["quantization_util.cc"], - hdrs = [ - "compatibility.h", - "quantization_util.h", - ], - deps = [ - ":round", - ":types", - "//tensorflow/contrib/lite/kernels:op_macros", - ], -) - -cc_test( - name = "quantization_util_test", - srcs = ["quantization_util_test.cc"], - tags = ["no_oss"], - deps = [ - ":quantization_util", - "@com_google_googletest//:gtest", - ], -) - -cc_library( - name = "strided_slice_logic", - srcs = [], - hdrs = [ - "strided_slice_logic.h", - ], - deps = [ - ":types", - ], -) - -cc_library( - name = "reference_base", - srcs = [], - hdrs = [ - "common.h", - "reference/depthwiseconv_float.h", - "reference/depthwiseconv_uint8.h", - "reference/fully_connected.h", - "reference/reference_ops.h", - "reference/softmax.h", - ], - deps = [ - ":quantization_util", - ":round", - ":strided_slice_logic", - ":types", - "@gemmlowp", - "//tensorflow/contrib/lite/c:c_api_internal", - "//tensorflow/contrib/lite/kernels:op_macros", - ] + select({ - ":haswell": tflite_deps_intel, - ":ios_x86_64": tflite_deps_intel, - ":k8": tflite_deps_intel, - ":x86": tflite_deps_intel, - ":x86_64": tflite_deps_intel, - ":darwin": tflite_deps_intel, - ":darwin_x86_64": tflite_deps_intel, - ":freebsd": tflite_deps_intel, - "//conditions:default": [], - }), -) - -cc_library( - name = "legacy_reference_base", - srcs = [], - hdrs = [ - "common.h", - "reference/depthwiseconv_float.h", - "reference/depthwiseconv_uint8.h", - "reference/fully_connected.h", - "reference/legacy_reference_ops.h", - "reference/reference_ops.h", - "reference/softmax.h", - ], - deps = [ - ":quantization_util", - ":round", - ":strided_slice_logic", - ":legacy_types", - ":types", - "@gemmlowp", - "//tensorflow/contrib/lite/c:c_api_internal", - "//tensorflow/contrib/lite/kernels:op_macros", - ] + select({ - ":haswell": tflite_deps_intel, - ":ios_x86_64": tflite_deps_intel, - ":k8": tflite_deps_intel, - ":x86": tflite_deps_intel, - ":x86_64": tflite_deps_intel, - ":darwin": tflite_deps_intel, - ":darwin_x86_64": tflite_deps_intel, - ":freebsd": tflite_deps_intel, - "//conditions:default": [], - }), -) - -cc_library( - name = "tensor", - hdrs = [ - "tensor.h", - "tensor_ctypes.h", - ], - deps = [ - ":types", - "//tensorflow/contrib/lite/c:c_api_internal", - ], -) - -# Deprecated version of :tensor, kept for backwards compatibility. -cc_library( - name = "reference", - hdrs = [ - "tensor.h", - "tensor_ctypes.h", - ], - deps = [ - ":types", - "//tensorflow/contrib/lite/c:c_api_internal", - ], -) - -cc_library( - name = "portable_tensor_utils", - srcs = [ - "reference/portable_tensor_utils.cc", - ], - hdrs = [ - "reference/portable_tensor_utils.h", - ], - deps = [ - ":round", - "//tensorflow/contrib/lite/c:c_api_internal", - "//tensorflow/contrib/lite/kernels:activation_functor", - "//tensorflow/contrib/lite/kernels:op_macros", - ], -) - -cc_library( - name = "neon_tensor_utils", - srcs = [ - "optimized/neon_tensor_utils.cc", - "reference/portable_tensor_utils.cc", - "reference/portable_tensor_utils.h", - ], - hdrs = [ - "common.h", - "compatibility.h", - "optimized/cpu_check.h", - "optimized/neon_tensor_utils.h", - "optimized/tensor_utils_impl.h", - ], - copts = NEON_FLAGS_IF_APPLICABLE + HARD_FP_FLAGS_IF_APPLICABLE, - deps = [ - ":cpu_check", - ":round", - ":types", - "//tensorflow/contrib/lite/c:c_api_internal", - "//tensorflow/contrib/lite/kernels:activation_functor", - "//tensorflow/contrib/lite/kernels:op_macros", - "@arm_neon_2_x86_sse", - "@gemmlowp", - ], -) - -cc_library( - name = "kernel_utils", - srcs = ["kernel_utils.cc"], - hdrs = ["kernel_utils.h"], - deps = [ - ":tensor_utils", - "//tensorflow/contrib/lite/c:c_api_internal", - ], -) - -# Audio support classes imported directly from TensorFlow. -cc_library( - name = "audio_utils", - srcs = [ - "mfcc.cc", - "mfcc_dct.cc", - "mfcc_mel_filterbank.cc", - "spectrogram.cc", - ], - hdrs = [ - "mfcc.h", - "mfcc_dct.h", - "mfcc_mel_filterbank.h", - "spectrogram.h", - ], - deps = [ - "//third_party/fft2d:fft2d_headers", - "@fft2d", - ], -) - -cc_library( - name = "tensor_utils", - srcs = [ - "tensor_utils.cc", - ], - hdrs = [ - "common.h", - "compatibility.h", - "optimized/cpu_check.h", - "optimized/neon_tensor_utils.h", - "optimized/tensor_utils_impl.h", - "reference/portable_tensor_utils.h", - "tensor_utils.h", - "types.h", - ], - copts = NEON_FLAGS_IF_APPLICABLE, - deps = [ - "@com_google_absl//absl/base:core_headers", - "//tensorflow/contrib/lite/c:c_api_internal", - "@arm_neon_2_x86_sse", - "//tensorflow/contrib/lite/kernels:op_macros", - "@gemmlowp", - ] + select({ - ":arm": [ - ":neon_tensor_utils", - ], - ":arm64-v8a": [ - ":neon_tensor_utils", - ], - ":armeabi-v7a": [ - ":neon_tensor_utils", - ], - ":armv7a": [ - ":neon_tensor_utils", - ], - ":haswell": [ - ":neon_tensor_utils", - ], - ":ios_armv7": [ - ":neon_tensor_utils", - ], - ":ios_arm64": [ - ":neon_tensor_utils", - ], - ":ios_x86_64": [ - ":neon_tensor_utils", - ], - ":x86_64": [ - ":neon_tensor_utils", - ], - ":x86": [ - ":neon_tensor_utils", - ], - ":k8": [ - ":neon_tensor_utils", - ], - ":darwin": [ - ":neon_tensor_utils", - ], - ":darwin_x86_64": [ - ":neon_tensor_utils", - ], - "//conditions:default": [ - ":portable_tensor_utils", - ], - }), -) - -cc_library( - name = "test_util", - srcs = ["test_util.cc"], - hdrs = ["test_util.h"], - deps = [ - ":types", - "//tensorflow/contrib/lite:string", - ], -) - -cc_test( - name = "tensor_utils_test", - srcs = ["tensor_utils_test.cc"], - copts = NEON_FLAGS_IF_APPLICABLE, - linkopts = select({ - "//tensorflow:android": [ - "-fPIE -pie", - ], - "//conditions:default": [], - }), - linkstatic = 1, - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":tensor_utils", - "//tensorflow/contrib/lite/c:c_api_internal", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest_main", - ], -) - -cc_test( - name = "depthwiseconv_float_test", - srcs = ["depthwiseconv_float_test.cc"], - tags = ["no_oss"], - deps = [ - ":optimized_base", - ":reference_base", - ":test_util", - ":types", - "@com_google_googletest//:gtest_main", - ], -) - -cc_test( - name = "depthwiseconv_quantized_test", - srcs = ["depthwiseconv_quantized_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable_ios", - ], - deps = [ - ":optimized_base", - ":reference_base", - ":test_util", - ":types", - "@com_google_googletest//:gtest_main", - ], -) - -cc_test( - name = "resize_bilinear_test", - srcs = ["resize_bilinear_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable", - ], - deps = [ - ":optimized_base", - ":reference_base", - ":test_util", - ":types", - "@com_google_googletest//:gtest_main", - ], -) - -cc_test( - name = "softmax_quantized_test", - timeout = "long", - srcs = [ - "softmax_quantized_test.cc", - ], - tags = ["no_oss"], - deps = [ - ":optimized_base", - ":quantization_util", - ":reference_base", - ":test_util", - "//tensorflow/contrib/lite:string", - "@com_google_googletest//:gtest_main", - ], -) - -cc_test( - name = "logsoftmax_quantized_test", - timeout = "long", - srcs = [ - "logsoftmax_quantized_test.cc", - ], - tags = [ - "no_oss", - "tflite_not_portable", - ], - deps = [ - ":optimized_base", - ":quantization_util", - ":reference_base", - ":test_util", - "//tensorflow/contrib/lite:string", - "@com_google_googletest//:gtest_main", - ], -) - -cc_test( - name = "log_quantized_test", - srcs = ["log_quantized_test.cc"], - tags = ["no_oss"], - deps = [ - ":optimized_base", - ":reference_base", - "//tensorflow/contrib/lite:string", - "@com_google_googletest//:gtest_main", - ], -) - -cc_library( - name = "cpu_check", - hdrs = [ - "optimized/cpu_check.h", - ], - deps = [ - ] + select( - { - "//tensorflow:android": [ - "@androidndk//:cpufeatures", - ], - "//conditions:default": [], - }, - ), -) - -cc_test( - name = "batch_to_space_nd_test", - srcs = ["batch_to_space_nd_test.cc"], - tags = ["no_oss"], - deps = [ - ":optimized_base", - "@com_google_googletest//:gtest_main", - ], -) - -exports_files(["optimized/eigen_tensor_reduced_instantiations_oss.h"]) - -tflite_portable_test_suite() diff --git a/tensorflow/contrib/lite/kernels/internal/batch_to_space_nd_test.cc b/tensorflow/contrib/lite/kernels/internal/batch_to_space_nd_test.cc deleted file mode 100644 index 5a2901ac8c297265e542cc30d3127fe774c19e78..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/internal/batch_to_space_nd_test.cc +++ /dev/null @@ -1,98 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include "tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h" - -#include - -namespace tflite { -namespace { - -// A light wrapper of GetIndexRange which returns a pair of start / end -// indices. -std::pair GetIndexRange(int spatial_index_dim, int block_shape_dim, - int input_dim, int output_dim) { - int index_start = 0; - int index_end = 0; - optimized_ops::GetIndexRange(spatial_index_dim, block_shape_dim, input_dim, - output_dim, &index_start, &index_end); - return {index_start, index_end}; -} - -TEST(BatchToSpaceNDTest, TestIndexRange) { - // Simple test case, no cropping. - EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/3, /*block_shape_dim=*/6, - /*input_dim=*/1, /*output_dim=*/6), - std::make_pair(0, 1)); - - // No cropping and input_dim > 1. - EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/2, /*block_shape_dim=*/6, - /*input_dim=*/5, /*output_dim=*/30), - std::make_pair(0, 5)); - - // With small cropping values (can be either at the beginning or at the end). - EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/0, /*block_shape_dim=*/2, - /*input_dim=*/3, /*output_dim=*/4), - std::make_pair(0, 2)); - - // With positive cropping values at the beginning. - EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/-2, /*block_shape_dim=*/2, - /*input_dim=*/3, /*output_dim=*/4), - std::make_pair(1, 3)); - - // Large crop at the beginning. - EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/-30, /*block_shape_dim=*/5, - /*input_dim=*/7, /*output_dim=*/5), - std::make_pair(6, 7)); - - EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/-26, /*block_shape_dim=*/5, - /*input_dim=*/7, /*output_dim=*/5), - std::make_pair(6, 7)); - - // Large crop at the end. - EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/0, /*block_shape_dim=*/5, - /*input_dim=*/7, /*output_dim=*/5), - std::make_pair(0, 1)); - - EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/4, /*block_shape_dim=*/5, - /*input_dim=*/7, /*output_dim=*/5), - std::make_pair(0, 1)); - - // Rounding up incorrectly will fail this test. - EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/3, /*block_shape_dim=*/5, - /*input_dim=*/7, /*output_dim=*/5), - std::make_pair(0, 1)); - - // Extreme cropping with output of a single spatial location. - // Valid position 1, when large crop at the end. - EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/0, /*block_shape_dim=*/5, - /*input_dim=*/7, /*output_dim=*/1), - std::make_pair(0, 1)); - - // Valid position 2, when large crop at the beginning. - EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/-30, /*block_shape_dim=*/5, - /*input_dim=*/7, /*output_dim=*/1), - std::make_pair(6, 7)); - - // Invalid positions. - EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/1, /*block_shape_dim=*/5, - /*input_dim=*/7, /*output_dim=*/1), - std::make_pair(0, 0)); - EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/-29, /*block_shape_dim=*/5, - /*input_dim=*/7, /*output_dim=*/1), - std::make_pair(6, 6)); -} - -} // namespace -} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/internal/common.h b/tensorflow/contrib/lite/kernels/internal/common.h deleted file mode 100644 index e67fee11b8d24d386d3b7c5efa4b07463fb8024a..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/internal/common.h +++ /dev/null @@ -1,269 +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_CONTRIB_LITE_KERNELS_INTERNAL_COMMON_H_ -#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_COMMON_H_ - -#ifndef ALLOW_SLOW_GENERIC_DEPTHWISECONV_FALLBACK -#ifdef GEMMLOWP_ALLOW_SLOW_SCALAR_FALLBACK -#define ALLOW_SLOW_GENERIC_DEPTHWISECONV_FALLBACK -#endif -#endif - -#ifndef USE_NEON -#if defined(__ARM_NEON__) || defined(__ARM_NEON) -#define USE_NEON -#include -#endif - -#if defined __GNUC__ && defined __SSE4_1__ -#define USE_NEON - -#define OPTIMIZED_OPS_H__IGNORE_DEPRECATED_DECLARATIONS -#pragma GCC diagnostic push -#pragma GCC diagnostic ignored "-Wdeprecated-declarations" -#pragma GCC diagnostic ignored "-Wattributes" - -#pragma GCC diagnostic push -#pragma GCC diagnostic ignored "-Wnarrowing" -#pragma GCC diagnostic ignored "-Wsequence-point" - -#include "NEON_2_SSE.h" - -#pragma GCC diagnostic pop -#endif -#endif - -#include "fixedpoint/fixedpoint.h" -#include "tensorflow/contrib/lite/kernels/internal/types.h" - -namespace tflite { - -inline void GetActivationMinMax(FusedActivationFunctionType ac, - float* output_activation_min, - float* output_activation_max) { - switch (ac) { - case FusedActivationFunctionType::kNone: - *output_activation_min = std::numeric_limits::lowest(); - *output_activation_max = std::numeric_limits::max(); - break; - case FusedActivationFunctionType::kRelu: - *output_activation_min = 0.f; - *output_activation_max = std::numeric_limits::max(); - break; - case FusedActivationFunctionType::kRelu1: - *output_activation_min = -1.f; - *output_activation_max = 1.f; - break; - case FusedActivationFunctionType::kRelu6: - *output_activation_min = 0.f; - *output_activation_max = 6.f; - break; - } -} - -inline float ActivationFunctionWithMinMax(float x, float output_activation_min, - float output_activation_max) { - return std::min(std::max(x, output_activation_min), output_activation_max); -} - -// Legacy function, left for compatibility only. -template -float ActivationFunction(float x) { - float output_activation_min, output_activation_max; - GetActivationMinMax(Ac, &output_activation_min, &output_activation_max); - return ActivationFunctionWithMinMax(x, output_activation_min, - output_activation_max); -} - -inline int32 MultiplyByQuantizedMultiplierSmallerThanOneExp( - int32 x, int32 quantized_multiplier, int left_shift) { - using gemmlowp::RoundingDivideByPOT; - using gemmlowp::SaturatingRoundingDoublingHighMul; - return RoundingDivideByPOT( - SaturatingRoundingDoublingHighMul(x, quantized_multiplier), -left_shift); -} - -inline int32 MultiplyByQuantizedMultiplierGreaterThanOne( - int32 x, int32 quantized_multiplier, int left_shift) { - using gemmlowp::SaturatingRoundingDoublingHighMul; - return SaturatingRoundingDoublingHighMul(x * (1 << left_shift), - quantized_multiplier); -} - -inline int32 MultiplyByQuantizedMultiplier(int32 x, int32 quantized_multiplier, - int shift) { - using gemmlowp::RoundingDivideByPOT; - using gemmlowp::SaturatingRoundingDoublingHighMul; - int left_shift = shift > 0 ? shift : 0; - int right_shift = shift > 0 ? 0 : -shift; - return RoundingDivideByPOT(SaturatingRoundingDoublingHighMul( - x * (1 << left_shift), quantized_multiplier), - right_shift); -} - -template -int CountLeadingZeros(T integer_input) { - static_assert(std::is_unsigned::value, - "Only unsigned integer types handled."); -#if defined(__GNUC__) - return integer_input ? __builtin_clz(integer_input) : 0; -#else - const T one_in_leading_positive = static_cast(1) - << (std::numeric_limits::digits - 1); - int leading_zeros = 0; - while (integer_input < one_in_leading_positive) { - integer_input <<= 1; - ++leading_zeros; - } - return leading_zeros; -#endif -} - -// DO NOT USE THIS STRUCT FOR NEW FUNCTIONALITY BEYOND IMPLEMENTING -// BROADCASTING. -// -// NdArrayDesc describes the shape and memory layout of an N-dimensional -// rectangular array of numbers. -// -// NdArrayDesc is basically identical to Dims defined in types.h. -// However, as Dims is to be deprecated, this class exists as an adaptor -// to enable simple unoptimized implementations of element-wise broadcasting -// operations. -template -struct NdArrayDesc { - // The "extent" of each dimension. Indices along dimension d must be in the - // half-open interval [0, extents[d]). - int extents[N]; - - // The number of *elements* (not bytes) between consecutive indices of each - // dimension. - int strides[N]; -}; - -// DO NOT USE THIS FUNCTION FOR NEW FUNCTIONALITY BEYOND IMPLEMENTING -// BROADCASTING. -// -// Same as Offset(), except takes as NdArrayDesc instead of Dims. -inline int SubscriptToIndex(const NdArrayDesc<4>& desc, int i0, int i1, int i2, - int i3) { - TFLITE_DCHECK(i0 >= 0 && i0 < desc.extents[0]); - TFLITE_DCHECK(i1 >= 0 && i1 < desc.extents[1]); - TFLITE_DCHECK(i2 >= 0 && i2 < desc.extents[2]); - TFLITE_DCHECK(i3 >= 0 && i3 < desc.extents[3]); - return i0 * desc.strides[0] + i1 * desc.strides[1] + i2 * desc.strides[2] + - i3 * desc.strides[3]; -} - -// Given the dimensions of the operands for an element-wise binary broadcast, -// adjusts them so that they can be directly iterated over with simple loops. -// Returns the adjusted dims as instances of NdArrayDesc in 'desc0_out' and -// 'desc1_out'. 'desc0_out' and 'desc1_out' cannot be nullptr. -// -// This function assumes that the two input shapes are compatible up to -// broadcasting and the shorter one has already been prepended with 1s to be the -// same length. E.g., if shape0 is (1, 16, 16, 64) and shape1 is (1, 64), -// shape1 must already have been prepended to be (1, 1, 1, 64). Recall that -// Dims refer to shapes in reverse order. In this case, input0_dims will be -// (64, 16, 16, 1) and input1_dims will be (64, 1, 1, 1). -// -// When two shapes are compatible up to broadcasting, for each dimension d, -// the input extents are either equal, or one of them is 1. -// -// This function performs the following for each dimension d: -// - If the extents are equal, then do nothing since the loop that walks over -// both of the input arrays is correct. -// - Otherwise, one (and only one) of the extents must be 1. Say extent0 is 1 -// and extent1 is e1. Then set extent0 to e1 and stride0 *to 0*. This allows -// array0 to be referenced *at any index* in dimension d and still access the -// same slice. -template -inline void NdArrayDescsForElementwiseBroadcast(const Dims& input0_dims, - const Dims& input1_dims, - NdArrayDesc* desc0_out, - NdArrayDesc* desc1_out) { - TFLITE_DCHECK(desc0_out != nullptr); - TFLITE_DCHECK(desc1_out != nullptr); - - // Copy dims to desc. - for (int i = 0; i < N; ++i) { - desc0_out->extents[i] = input0_dims.sizes[i]; - desc0_out->strides[i] = input0_dims.strides[i]; - desc1_out->extents[i] = input1_dims.sizes[i]; - desc1_out->strides[i] = input1_dims.strides[i]; - } - - // Walk over each dimension. If the extents are equal do nothing. - // Otherwise, set the desc with extent 1 to have extent equal to the other and - // stride 0. - for (int i = 0; i < N; ++i) { - const int extent0 = ArraySize(input0_dims, i); - const int extent1 = ArraySize(input1_dims, i); - if (extent0 != extent1) { - if (extent0 == 1) { - desc0_out->strides[i] = 0; - desc0_out->extents[i] = extent1; - } else { - TFLITE_DCHECK_EQ(extent1, 1); - desc1_out->strides[i] = 0; - desc1_out->extents[i] = extent0; - } - } - } -} - -template -inline void NdArrayDescsForElementwiseBroadcast( - const RuntimeShape& input0_shape, const RuntimeShape& input1_shape, - NdArrayDesc* desc0_out, NdArrayDesc* desc1_out) { - TFLITE_DCHECK(desc0_out != nullptr); - TFLITE_DCHECK(desc1_out != nullptr); - - auto extended_input0_shape = RuntimeShape::ExtendedShape(N, input0_shape); - auto extended_input1_shape = RuntimeShape::ExtendedShape(N, input1_shape); - - // Copy dims to desc, calculating strides. - int desc0_stride = 1; - int desc1_stride = 1; - for (int i = N - 1; i >= 0; --i) { - desc0_out->extents[i] = extended_input0_shape.Dims(i); - desc0_out->strides[i] = desc0_stride; - desc0_stride *= extended_input0_shape.Dims(i); - desc1_out->extents[i] = extended_input1_shape.Dims(i); - desc1_out->strides[i] = desc1_stride; - desc1_stride *= extended_input1_shape.Dims(i); - } - - // Walk over each dimension. If the extents are equal do nothing. - // Otherwise, set the desc with extent 1 to have extent equal to the other and - // stride 0. - for (int i = 0; i < N; ++i) { - const int extent0 = extended_input0_shape.Dims(i); - const int extent1 = extended_input1_shape.Dims(i); - if (extent0 != extent1) { - if (extent0 == 1) { - desc0_out->strides[i] = 0; - desc0_out->extents[i] = extent1; - } else { - TFLITE_DCHECK_EQ(extent1, 1); - desc1_out->strides[i] = 0; - desc1_out->extents[i] = extent0; - } - } - } -} - -} // namespace tflite - -#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_COMMON_H_ diff --git a/tensorflow/contrib/lite/kernels/internal/compatibility.h b/tensorflow/contrib/lite/kernels/internal/compatibility.h deleted file mode 100644 index 7c176e0fa1c8e8c8b6a094dbeb1025f2be091b3d..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/internal/compatibility.h +++ /dev/null @@ -1,110 +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_CONTRIB_LITE_KERNELS_INTERNAL_COMPATIBILITY_H_ -#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_COMPATIBILITY_H_ - -#include - -#include "tensorflow/contrib/lite/kernels/op_macros.h" - -#ifndef TFLITE_DCHECK -#define TFLITE_DCHECK(condition) (condition) ? (void)0 : TFLITE_ASSERT_FALSE -#endif - -#ifndef TFLITE_DCHECK_EQ -#define TFLITE_DCHECK_EQ(x, y) ((x) == (y)) ? (void)0 : TFLITE_ASSERT_FALSE -#endif - -#ifndef TFLITE_DCHECK_NE -#define TFLITE_DCHECK_NE(x, y) ((x) != (y)) ? (void)0 : TFLITE_ASSERT_FALSE -#endif - -#ifndef TFLITE_DCHECK_GE -#define TFLITE_DCHECK_GE(x, y) ((x) >= (y)) ? (void)0 : TFLITE_ASSERT_FALSE -#endif - -#ifndef TFLITE_DCHECK_GT -#define TFLITE_DCHECK_GT(x, y) ((x) > (y)) ? (void)0 : TFLITE_ASSERT_FALSE -#endif - -#ifndef TFLITE_DCHECK_LE -#define TFLITE_DCHECK_LE(x, y) ((x) <= (y)) ? (void)0 : TFLITE_ASSERT_FALSE -#endif - -#ifndef TFLITE_DCHECK_LT -#define TFLITE_DCHECK_LT(x, y) ((x) < (y)) ? (void)0 : TFLITE_ASSERT_FALSE -#endif - -// TODO(ahentz): Clean up: We should stick to the DCHECK versions. -#ifndef TFLITE_CHECK -#define TFLITE_CHECK(condition) (condition) ? (void)0 : TFLITE_ABORT -#endif - -#ifndef TFLITE_CHECK_EQ -#define TFLITE_CHECK_EQ(x, y) ((x) == (y)) ? (void)0 : TFLITE_ABORT -#endif - -#ifndef TFLITE_CHECK_NE -#define TFLITE_CHECK_NE(x, y) ((x) != (y)) ? (void)0 : TFLITE_ABORT -#endif - -#ifndef TFLITE_CHECK_GE -#define TFLITE_CHECK_GE(x, y) ((x) >= (y)) ? (void)0 : TFLITE_ABORT -#endif - -#ifndef TFLITE_CHECK_GT -#define TFLITE_CHECK_GT(x, y) ((x) > (y)) ? (void)0 : TFLITE_ABORT -#endif - -#ifndef TFLITE_CHECK_LE -#define TFLITE_CHECK_LE(x, y) ((x) <= (y)) ? (void)0 : TFLITE_ABORT -#endif - -#ifndef TFLITE_CHECK_LT -#define TFLITE_CHECK_LT(x, y) ((x) < (y)) ? (void)0 : TFLITE_ABORT -#endif - -// TODO(ahentz): Clean up. -using int8 = std::int8_t; -using uint8 = std::uint8_t; -using int16 = std::int16_t; -using uint16 = std::uint16_t; -using int32 = std::int32_t; -using uint32 = std::uint32_t; - -// TFLITE_DEPRECATED() -// -// Duplicated from absl/base/macros.h to avoid pulling in that library. -// Marks a deprecated class, struct, enum, function, method and variable -// declarations. The macro argument is used as a custom diagnostic message (e.g. -// suggestion of a better alternative). -// -// Example: -// -// class TFLITE_DEPRECATED("Use Bar instead") Foo {...}; -// TFLITE_DEPRECATED("Use Baz instead") void Bar() {...} -// -// Every usage of a deprecated entity will trigger a warning when compiled with -// clang's `-Wdeprecated-declarations` option. This option is turned off by -// default, but the warnings will be reported by clang-tidy. -#if defined(__clang__) && __cplusplus >= 201103L -#define TFLITE_DEPRECATED(message) __attribute__((deprecated(message))) -#endif - -#ifndef TFLITE_DEPRECATED -#define TFLITE_DEPRECATED(message) -#endif - -#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_COMPATIBILITY_H_ diff --git a/tensorflow/contrib/lite/kernels/internal/depthwiseconv_quantized_test.cc b/tensorflow/contrib/lite/kernels/internal/depthwiseconv_quantized_test.cc deleted file mode 100644 index 9414e109c302510a6fd434b410f0cbb575023e76..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/internal/depthwiseconv_quantized_test.cc +++ /dev/null @@ -1,349 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include -#include -#include -#include -#include -#include -#include -#include - -#include -#include "tensorflow/contrib/lite/kernels/internal/test_util.h" -#include "tensorflow/contrib/lite/kernels/internal/types.h" - -#define ALLOW_SLOW_GENERIC_DEPTHWISECONV_FALLBACK -#include "tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_uint8.h" -#include "tensorflow/contrib/lite/kernels/internal/reference/depthwiseconv_uint8.h" - -namespace tflite { -namespace { - -// Runs the DepthwiseConv and compares against the reference implementation. -template -int TestOneDepthwiseConvWithGivenOutputShift( - const std::uint8_t* input_data, const RuntimeShape& input_shape, - std::int32_t input_offset, const std::uint8_t* filter_data, - const RuntimeShape& filter_shape, std::int32_t filter_offset, - const std::int32_t* bias_data, const RuntimeShape& bias_shape, int stride, - int pad_width, int pad_height, int depth_multiplier, - std::int32_t output_offset, std::int32_t output_multiplier, - int output_shift, std::int32_t output_activation_min, - std::int32_t output_activation_max, const RuntimeShape& output_shape) { - const int output_buffer_size = output_shape.FlatSize(); - std::vector output_data(output_buffer_size); - std::vector reference_output_data(output_buffer_size); - - tflite::DepthwiseParams op_params; - op_params.padding_type = PaddingType::kSame; - op_params.padding_values.width = pad_width; - op_params.padding_values.height = pad_height; - op_params.stride_width = stride; - op_params.stride_height = stride; - op_params.dilation_width_factor = 1; - op_params.dilation_height_factor = 1; - op_params.depth_multiplier = depth_multiplier; - op_params.quantized_activation_min = output_activation_min; - op_params.quantized_activation_max = output_activation_max; - op_params.input_offset = input_offset; - op_params.weights_offset = filter_offset; - op_params.output_offset = output_offset; - op_params.output_multiplier = output_multiplier; - op_params.output_shift = -output_shift; - reference_ops::DepthwiseConv(op_params, input_shape, input_data, filter_shape, - filter_data, bias_shape, bias_data, output_shape, - reference_output_data.data()); - optimized_ops::DepthwiseConv(op_params, input_shape, input_data, filter_shape, - filter_data, bias_shape, bias_data, output_shape, - output_data.data()); - int saturated_min = 0; - int saturated_max = 0; - std::vector diff(output_buffer_size); - std::int64_t sum_diff = 0; - std::int64_t sum_abs_diff = 0; - for (int i = 0; i < output_buffer_size; i++) { - diff[i] = static_cast(output_data[i]) - - static_cast(reference_output_data[i]); - sum_diff += diff[i]; - sum_abs_diff += std::abs(diff[i]); - saturated_min += output_data[i] == output_activation_min; - saturated_max += output_data[i] == output_activation_max; - } - // These stats help understand test failures. - std::sort(std::begin(diff), std::end(diff)); - const int min_diff = diff.front(); - const int max_diff = diff.back(); - const int median_diff = diff[diff.size() / 2]; - const float mean_diff = static_cast(sum_diff) / output_buffer_size; - const float mean_abs_diff = - static_cast(sum_abs_diff) / output_buffer_size; - // Normally we should require bit-for-bit exact results. Unfortunately a bug - // in the Intel arm_neon_sse.h translation header that we use for x86 tests - // causes 1-bit inaccuracy in - // the vqrdmulh_n_s32 intrinsic, which causes off-by-1 errors in quantized - // DepthwiseConv ops. So we have to live with a few off-by-one errors for now, - // yet still ensure that no more than a small minority of values are wrong. - EXPECT_TRUE(std::abs(mean_diff) < 1e-5f && mean_abs_diff < 1e-5f && - std::abs(median_diff) == 0 && std::abs(min_diff) <= 1 && - std::abs(max_diff) <= 1); - if (saturated_min > 2 * saturated_max) { - return -1; - } - if (saturated_max > 2 * saturated_min) { - return 1; - } - return 0; -} - -// The point of this function is that we can't practically know which -// output_shift value to pass to test DepthwiseConv. It's not easy to guess (we -// could do some -// statistics for large size, but they would be fragile at smaller sizes), and -// guessing wrong would mean that all the values get saturated so the test -// becomes -// vacuous. So we just bisect our way to reasonable output_shift values. -template -void TestOneDepthwiseConvBisectOutputShift( - const std::uint8_t* input_data, const RuntimeShape& input_shape, - std::int32_t input_offset, const std::uint8_t* filter_data, - const RuntimeShape& filter_shape, std::int32_t filter_offset, - const std::int32_t* bias_data, const RuntimeShape& bias_shape, int stride, - int pad_width, int pad_height, int depth_multiplier, - std::int32_t output_offset, std::int32_t output_multiplier, - int output_activation_bisect_start, int output_activation_bisect_end, - std::int32_t output_activation_min, std::int32_t output_activation_max, - const RuntimeShape& output_shape) { - ASSERT_LT(output_activation_bisect_start, output_activation_bisect_end) - << "Bisection failed ?!?!"; - int output_shift_bisect_midpoint = - (output_activation_bisect_start + output_activation_bisect_end) / 2; - int bisect_result = TestOneDepthwiseConvWithGivenOutputShift( - input_data, input_shape, input_offset, filter_data, filter_shape, - filter_offset, bias_data, bias_shape, stride, pad_width, pad_height, - depth_multiplier, output_offset, output_multiplier, - output_shift_bisect_midpoint, output_activation_min, - output_activation_max, output_shape); - // At this point we know that the test succeeded (otherwise it would have - // aborted). - if (bisect_result == 0) { - // The result isn't particularly saturated on one or the other side. - // All good, we're done. - return; - } - if (output_activation_bisect_start == output_activation_bisect_end - 1) { - // There is still some saturation on one side, but the bisection is - // finished anyways. We're done; nothing more we can do about it. This - // happens - // in particular when using an activation with a narrow range. - return; - } - // Continue the bisection based on the present result. - int new_output_activation_bisect_start = bisect_result == 1 - ? output_shift_bisect_midpoint - : output_activation_bisect_start; - int new_output_activation_bisect_end = bisect_result == 1 - ? output_activation_bisect_end - : output_shift_bisect_midpoint; - TestOneDepthwiseConvBisectOutputShift( - input_data, input_shape, input_offset, filter_data, filter_shape, - filter_offset, bias_data, bias_shape, stride, pad_width, pad_height, - depth_multiplier, output_offset, output_multiplier, - new_output_activation_bisect_start, new_output_activation_bisect_end, - output_activation_min, output_activation_max, output_shape); -} - -template -void TestOneDepthwiseConv( - const std::uint8_t* input_data, const RuntimeShape& input_shape, - std::int32_t input_offset, const std::uint8_t* filter_data, - const RuntimeShape& filter_shape, std::int32_t filter_offset, - const std::int32_t* bias_data, const RuntimeShape& bias_shape, int stride, - int pad_width, int pad_height, int depth_multiplier, - std::int32_t output_offset, std::int32_t output_multiplier, - std::int32_t output_activation_min, std::int32_t output_activation_max, - const RuntimeShape& output_shape) { - TestOneDepthwiseConvBisectOutputShift( - input_data, input_shape, input_offset, filter_data, filter_shape, - filter_offset, bias_data, bias_shape, stride, pad_width, pad_height, - depth_multiplier, output_offset, output_multiplier, 0, 32, - output_activation_min, output_activation_max, output_shape); -} - -void TestOneDepthwiseConv( - FusedActivationFunctionType Ac, const std::uint8_t* input_data, - const RuntimeShape& input_shape, std::int32_t input_offset, - const std::uint8_t* filter_data, const RuntimeShape& filter_shape, - std::int32_t filter_offset, const std::int32_t* bias_data, - const RuntimeShape& bias_shape, int stride, int pad_width, int pad_height, - int depth_multiplier, std::int32_t output_offset, - std::int32_t output_multiplier, std::int32_t output_activation_min, - std::int32_t output_activation_max, const RuntimeShape& output_shape) { -#define TOCO_HANDLE_CASE(AC_TYPE) \ - if (AC_TYPE == Ac) { \ - TestOneDepthwiseConv( \ - input_data, input_shape, input_offset, filter_data, filter_shape, \ - filter_offset, bias_data, bias_shape, stride, pad_width, pad_height, \ - depth_multiplier, output_offset, output_multiplier, \ - output_activation_min, output_activation_max, output_shape); \ - return; \ - } - TOCO_HANDLE_CASE(FusedActivationFunctionType::kNone) - TOCO_HANDLE_CASE(FusedActivationFunctionType::kRelu) - TOCO_HANDLE_CASE(FusedActivationFunctionType::kRelu1) - TOCO_HANDLE_CASE(FusedActivationFunctionType::kRelu6) -#undef TOCO_HANDLE_CASE -} - -bool TryTestDepthwiseConv(int batch, int input_depth, int input_width, - int input_height, int filter_width, int filter_height, - int depth_multiplier, int stride, - int dilation_width_factor, int dilation_height_factor, - PaddingType padding_type) { - const int output_depth = input_depth * depth_multiplier; - // The optimized DepthwiseConv implementation currently uses a fixed-size - // accumulator buffer on the stack, with that size. This currently means - // that it does not support larger output depths. It CHECK's for it, - // so it's safe in the sense that if a larger output depth was encountered, - // it would explicitly fail. We just need to adjust our testing to that - // constraint. - const int kMaxSupportedOutputDepth = 1024; - if (output_depth > kMaxSupportedOutputDepth) { - return false; - } - const auto ac = RandomElement(std::vector( - {FusedActivationFunctionType::kNone, FusedActivationFunctionType::kRelu, - FusedActivationFunctionType::kRelu6, - FusedActivationFunctionType::kRelu1})); - int output_activation_min = 0; - int output_activation_max = 255; - if (ac != FusedActivationFunctionType::kNone && UniformRandomInt(0, 1)) { - output_activation_min = UniformRandomInt(0, 50); - output_activation_max = UniformRandomInt(200, 255); - } - const std::int32_t output_multiplier = - UniformRandomInt(1 << 29, std::numeric_limits::max()); - const std::int32_t input_offset = UniformRandomInt(-256, 0); - const std::int32_t filter_offset = UniformRandomInt(-256, 0); - const std::int32_t output_offset = UniformRandomInt(-256, 0); - RuntimeShape input_shape_inference( - {batch, input_height, input_width, input_depth}); - RuntimeShape output_shape_inference; - int pad_width, pad_height; - if (!ComputeConvSizes(input_shape_inference, output_depth, filter_width, - filter_height, stride, dilation_width_factor, - dilation_height_factor, padding_type, - &output_shape_inference, &pad_width, &pad_height)) { - return false; - } - RuntimeShape filter_shape_inference( - {1, filter_height, filter_width, output_depth}); - RuntimeShape bias_shape_inference({1, 1, 1, output_depth}); - const int input_buffer_size = input_shape_inference.FlatSize(); - const int filter_buffer_size = filter_shape_inference.FlatSize(); - std::vector input_data(input_buffer_size); - std::vector filter_data(filter_buffer_size); - std::vector bias_data(output_depth); - FillRandom(&input_data); - FillRandom(&filter_data); - FillRandom(&bias_data, -10000, 10000); - TestOneDepthwiseConv(ac, input_data.data(), input_shape_inference, - input_offset, filter_data.data(), filter_shape_inference, - filter_offset, bias_data.data(), bias_shape_inference, - stride, pad_width, pad_height, depth_multiplier, - output_offset, output_multiplier, output_activation_min, - output_activation_max, output_shape_inference); - return true; -} - -// This function picks some random DepthwiseConv params, which may or may not -// be legal. If they're not legal, it returns false. If they're legal, -// it runs the DepthwiseConv test and returns true. This allows the caller -// to loop until a test has been run. -bool TryTestOneDepthwiseConv() { - // We have to pick a lot of positive values, where we are particularly - // interested in small values because they are most likely to be special - // cases in optimized implementations, and secondarily because they allow - // tests to run fast, which means we can run more tests and get more - // coverage. - const int batch = ExponentialRandomPositiveInt(0.9f, 3, 20); - const int input_depth = ExponentialRandomPositiveInt(0.9f, 6, 50); - const int input_width = ExponentialRandomPositiveInt(0.9f, 20, 200); - const int input_height = ExponentialRandomPositiveInt(0.9f, 20, 200); - const int filter_width = ExponentialRandomPositiveInt(0.9f, 4, 10); - const int filter_height = ExponentialRandomPositiveInt(0.9f, 4, 10); - const int depth_multiplier = ExponentialRandomPositiveInt(0.8f, 6, 50); - const int stride = ExponentialRandomPositiveInt(0.9f, 3, 8); - const int dilation_width_factor = RandomElement(std::vector({1, 2, 4})); - const int dilation_height_factor = RandomElement(std::vector({1, 2, 4})); - const auto padding_type = - UniformRandomInt(0, 1) ? PaddingType::kSame : PaddingType::kValid; - - return TryTestDepthwiseConv(batch, input_depth, input_width, input_height, - filter_width, filter_height, depth_multiplier, - stride, dilation_width_factor, - dilation_height_factor, padding_type); -} - -// Tests parameters for the 3x3 filter kernel. -bool TryTestOneDepthwiseConv3x3Filter() { - const int batch = ExponentialRandomPositiveInt(0.9f, 3, 20); - const int input_depth = 8 * ExponentialRandomPositiveInt(0.9f, 10, 50); - const int input_width = ExponentialRandomPositiveInt(0.9f, 20, 200); - const int input_height = ExponentialRandomPositiveInt(0.9f, 20, 200); - const int filter_width = 3; - const int filter_height = 3; - const int depth_multiplier = 1; - const int stride = UniformRandomInt(1, 2); - // We don't support dilations in the 3x3 filter. - const int dilation_width_factor = 1; - const int dilation_height_factor = 1; - // Although the kernel supports only kValid padding, we test that kSame - // is using the correct code path. - const auto padding_type = - UniformRandomInt(0, 1) ? PaddingType::kSame : PaddingType::kValid; - - return TryTestDepthwiseConv(batch, input_depth, input_width, input_height, - filter_width, filter_height, depth_multiplier, - stride, dilation_width_factor, - dilation_height_factor, padding_type); -} - -void TestOneDepthwiseConv() { - while (!TryTestOneDepthwiseConv()) { - } -} - -void TestOneDepthwiseConv3x3Filter() { - while (!TryTestOneDepthwiseConv3x3Filter()) { - } -} - -TEST(TestDepthwiseConv, TestDepthwiseConv) { - const int kTestsToRun = 10 * 1000; - for (int i = 0; i < kTestsToRun; i++) { - TestOneDepthwiseConv(); - } -} - -TEST(TestDepthwiseConv3x3Filter, TestDepthwiseConv) { - const int kTestsToRun = 3 * 1000; - for (int i = 0; i < kTestsToRun; i++) { - TestOneDepthwiseConv3x3Filter(); - } -} - -} // namespace -} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/internal/legacy_types.h b/tensorflow/contrib/lite/kernels/internal/legacy_types.h deleted file mode 100644 index 2e4d3137f5c6acfdd7a0942510433323e5fee5ed..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/internal/legacy_types.h +++ /dev/null @@ -1,26 +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_LITE_KERNELS_INTERNAL_LEGACY_TYPES_H_ -#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_LEGACY_TYPES_H_ - -#include "tensorflow/contrib/lite/kernels/internal/types.h" - -namespace tflite { - -// TODO(b/116772710): Insert legacy Dims<> code in here. - -} // namespace tflite - -#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_LEGACY_TYPES_H_ diff --git a/tensorflow/contrib/lite/kernels/internal/mfcc.cc b/tensorflow/contrib/lite/kernels/internal/mfcc.cc deleted file mode 100644 index eafe0c7afee6fabd5a4a258aa5176e23f5e8d62a..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/internal/mfcc.cc +++ /dev/null @@ -1,65 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include - -#include "tensorflow/contrib/lite/kernels/internal/mfcc.h" - -namespace tflite { -namespace internal { - -const double kDefaultUpperFrequencyLimit = 4000; -const double kDefaultLowerFrequencyLimit = 20; -const double kFilterbankFloor = 1e-12; -const int kDefaultFilterbankChannelCount = 40; -const int kDefaultDCTCoefficientCount = 13; - -Mfcc::Mfcc() - : initialized_(false), - lower_frequency_limit_(kDefaultLowerFrequencyLimit), - upper_frequency_limit_(kDefaultUpperFrequencyLimit), - filterbank_channel_count_(kDefaultFilterbankChannelCount), - dct_coefficient_count_(kDefaultDCTCoefficientCount) {} - -bool Mfcc::Initialize(int input_length, double input_sample_rate) { - bool initialized = mel_filterbank_.Initialize( - input_length, input_sample_rate, filterbank_channel_count_, - lower_frequency_limit_, upper_frequency_limit_); - initialized &= - dct_.Initialize(filterbank_channel_count_, dct_coefficient_count_); - initialized_ = initialized; - return initialized; -} - -void Mfcc::Compute(const std::vector& spectrogram_frame, - std::vector* output) const { - if (!initialized_) { - // LOG(ERROR) << "Mfcc not initialized."; - return; - } - std::vector working; - mel_filterbank_.Compute(spectrogram_frame, &working); - for (int i = 0; i < working.size(); ++i) { - double val = working[i]; - if (val < kFilterbankFloor) { - val = kFilterbankFloor; - } - working[i] = log(val); - } - dct_.Compute(working, output); -} - -} // namespace internal -} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/internal/optimized/cblas_conv.h b/tensorflow/contrib/lite/kernels/internal/optimized/cblas_conv.h deleted file mode 100644 index 2d96da65c33bd4d1d132501dfaa49148f2c26484..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/internal/optimized/cblas_conv.h +++ /dev/null @@ -1,109 +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_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_CBLAS_CONV_H_ -#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_CBLAS_CONV_H_ - -// The Conv implementation based on CBLAS interface. This is only used on iOS -// for now, utilizing Apple's Accelerate framework. - -#if TFLITE_USE_APPLE_ACCELERATE_FOR_CONV -#include -#else -#include "tensorflow/contrib/lite/kernels/internal/optimized/cblas_reference.h" -#endif - -#include "tensorflow/contrib/lite/kernels/internal/optimized/multithreaded_conv.h" -#include "tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h" - -namespace tflite { -namespace cblas_ops { - -inline void Conv(const ConvParams& params, const RuntimeShape& input_shape, - const float* input_data, const RuntimeShape& filter_shape, - const float* filter_data, const RuntimeShape& bias_shape, - const float* bias_data, const RuntimeShape& output_shape, - float* output_data, const RuntimeShape& im2col_shape, - float* im2col_data) { - const int stride_width = params.stride_width; - const int stride_height = params.stride_height; - const int pad_width = params.padding_values.width; - const int pad_height = params.padding_values.height; - const int dilation_width_factor = params.dilation_width_factor; - const int dilation_height_factor = params.dilation_height_factor; - const float output_activation_min = params.float_activation_min; - const float output_activation_max = params.float_activation_max; - TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4); - TFLITE_DCHECK_EQ(filter_shape.DimensionsCount(), 4); - TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4); - gemmlowp::ScopedProfilingLabel label("Conv/cblas"); - - const float* gemm_input_data = nullptr; - const RuntimeShape* gemm_input_shape = nullptr; - const int filter_width = filter_shape.Dims(2); - const int filter_height = filter_shape.Dims(1); - const bool need_im2col = stride_width != 1 || stride_height != 1 || - filter_width != 1 || filter_height != 1; - if (need_im2col) { - TFLITE_DCHECK(im2col_data); - ConvParams op_params; - op_params.padding_type = PaddingType::kSame; - op_params.padding_values.width = pad_width; - op_params.padding_values.height = pad_height; - op_params.stride_width = stride_width; - op_params.stride_height = stride_height; - op_params.dilation_width_factor = dilation_width_factor; - op_params.dilation_height_factor = dilation_height_factor; - optimized_ops::Im2col(op_params, filter_height, filter_width, 0, - input_shape, input_data, im2col_shape, im2col_data); - - gemm_input_data = im2col_data; - gemm_input_shape = &im2col_shape; - } else { - TFLITE_DCHECK(!im2col_data); - gemm_input_data = input_data; - gemm_input_shape = &input_shape; - } - - // The following code computes matrix multiplication c = a * transponse(b) - // with CBLAS, where: - // * `a` is a matrix with dimensions (m, k). - // * `b` is a matrix with dimensions (n, k), so transpose(b) is (k, n). - // * `c` is a matrix with dimensions (m, n). - // The naming of variables are aligned with CBLAS specification here. - const float* a = gemm_input_data; - const float* b = filter_data; - float* c = output_data; - const int gemm_input_dims = gemm_input_shape->DimensionsCount(); - int m = FlatSizeSkipDim(*gemm_input_shape, gemm_input_dims - 1); - int n = output_shape.Dims(3); - int k = gemm_input_shape->Dims(gemm_input_dims - 1); - // The stride of matrix a, b and c respectively. - int stride_a = k; - int stride_b = k; - int stride_c = n; - - cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasTrans, m, n, k, 1.0f, a, - stride_a, b, stride_b, 0.0f, c, stride_c); - - optimized_ops::AddBiasAndEvalActivationFunction( - output_activation_min, output_activation_max, bias_shape, bias_data, - output_shape, output_data); -} - -} // namespace cblas_ops -} // namespace tflite - -#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_CBLAS_CONV_H_ diff --git a/tensorflow/contrib/lite/kernels/internal/optimized/cblas_reference.h b/tensorflow/contrib/lite/kernels/internal/optimized/cblas_reference.h deleted file mode 100644 index 6acc513805c9398c304f3e24175d3bd6c96938f6..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/internal/optimized/cblas_reference.h +++ /dev/null @@ -1,69 +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_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_CBLAS_REFERENCE_H_ -#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_CBLAS_REFERENCE_H_ - -#include "tensorflow/contrib/lite/kernels/internal/compatibility.h" - -// The reference implementation for a small subset of CBLAS interface. -// This is only used for testing CBLAS implementation, and should never be used -// in production code. - -namespace tflite { -namespace cblas_ops { - -// The following code follows the original CBLAS specification, and it might -// conflict with the TensorFlow naming convention. -// TODO(ycling): Find another way to test CBLAS with bazel, without writing -// a reference implementation by ourselves. -enum CBLAS_ORDER { CblasRowMajor = 0, CblasColMajor = 1 }; - -enum CBLAS_TRANSPOSE { CblasNoTrans = 0, CblasTrans = 1, CblasConjTrans = 2 }; - -// A reference implementation for matrix multiplication. -// The following code computes, c = a * transponse(b) matrix multiplication -// with CBLAS, where: -// * `a` is a matrix with dimensions (m, k). -// * `b` is a matrix with dimensions (n, k), so transpose(b) is (k, n). -// * `c` is a matrix with dimensions (m, n). -// The naming of variables is aligned with CBLAS specification here. -void cblas_sgemm(const enum CBLAS_ORDER order, - const enum CBLAS_TRANSPOSE trans_a, - const enum CBLAS_TRANSPOSE trans_b, const int m, const int n, - const int k, const float alpha, const float *a, - const int stride_a, const float *b, const int stride_b, - const float beta, float *c, const int stride_c) { - TFLITE_DCHECK(order == CblasRowMajor); - TFLITE_DCHECK(trans_a == CblasNoTrans); - TFLITE_DCHECK(trans_b == CblasTrans); - TFLITE_DCHECK(beta == 0.0f); - for (int row = 0; row < m; ++row) { - for (int col = 0; col < n; ++col) { - // If `beta` non-zero, multiple it with the original values in output. - // Otherwise, ignore the original value in output completely. - float value = 0.0f; - for (int idx = 0; idx < k; ++idx) { - value += alpha * a[stride_a * row + idx] * b[stride_b * col + idx]; - } - c[stride_c * row + col] = value; - } - } -} - -} // namespace cblas_ops -} // namespace tflite - -#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_CBLAS_REFERENCE_H_ diff --git a/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_float.h b/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_float.h deleted file mode 100644 index bcadfb2f8cdb1a3301b1f62c6b6d3a964dda0125..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_float.h +++ /dev/null @@ -1,1071 +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_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_DEPTHWISECONV_FLOAT_H_ -#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_DEPTHWISECONV_FLOAT_H_ - -#include "public/gemmlowp.h" -#include "tensorflow/contrib/lite/kernels/internal/common.h" -#include "tensorflow/contrib/lite/kernels/internal/types.h" - -namespace tflite { -namespace optimized_ops { - -// Implementation of float DepthwiseConv - -template -struct FloatDepthwiseConvKernel {}; - -#ifdef USE_NEON - -template <> -struct FloatDepthwiseConvKernel { - static void Run(int num_output_pixels, int input_depth, int depth_multiplier, - const float* input_ptr, int input_ptr_increment, - const float* filter_ptr, float* acc_buffer_ptr) { - // Load the filters - float32x4_t filter[2]; - for (int i = 0; i < 2; i++) { - filter[i] = vld1q_f32(filter_ptr + 4 * i); - } - int outp = 0; - // Handle 2 output pixels at a time. - for (; outp <= num_output_pixels - 2; outp += 2) { - // Load the inputs - float32x4_t input[4]; - for (int i = 0; i < 4; i++) { - input[i] = vld1q_f32(input_ptr + 4 * i); - } - input_ptr += 16; - // Load the accumulators from acc_buffer - float32x4_t acc[4]; - for (int i = 0; i < 4; i++) { - acc[i] = vld1q_f32(acc_buffer_ptr + 4 * i); - } - // Multiply-accumulate - acc[0] = vmlaq_f32(acc[0], input[0], filter[0]); - acc[1] = vmlaq_f32(acc[1], input[1], filter[1]); - acc[2] = vmlaq_f32(acc[2], input[2], filter[0]); - acc[3] = vmlaq_f32(acc[3], input[3], filter[1]); - // Store the accumulators back to acc_buffer - for (int i = 0; i < 4; i++) { - vst1q_f32(acc_buffer_ptr + 4 * i, acc[i]); - } - acc_buffer_ptr += 16; - } - // Handle one output pixel at a time. - for (; outp < num_output_pixels; outp++) { - // Load the inputs - float32x4_t input[2]; - for (int i = 0; i < 2; i++) { - input[i] = vld1q_f32(input_ptr + 4 * i); - } - input_ptr += 8; - // Load the accumulators from acc_buffer - float32x4_t acc[2]; - for (int i = 0; i < 2; i++) { - acc[i] = vld1q_f32(acc_buffer_ptr + 4 * i); - } - // Multiply-accumulate - for (int i = 0; i < 2; i++) { - acc[i] = vmlaq_f32(acc[i], input[i], filter[i]); - } - // Store the accumulators back to acc_buffer - for (int i = 0; i < 2; i++) { - vst1q_f32(acc_buffer_ptr + 4 * i, acc[i]); - } - acc_buffer_ptr += 8; - } - } -}; - -template <> -struct FloatDepthwiseConvKernel { - static void Run(int num_output_pixels, int input_depth, int depth_multiplier, - const float* input_ptr, int input_ptr_increment, - const float* filter_ptr, float* acc_buffer_ptr) { - const float32x2_t filters = vld1_f32(filter_ptr); - const float32x4_t filters_dup2 = vcombine_f32(filters, filters); - int outp = 0; - // Handle 8 output pixels at a time. - for (; outp <= num_output_pixels - 8; outp += 8) { - // Load the inputs - float32x4_t input[4]; - for (int i = 0; i < 4; i++) { - input[i] = vld1q_f32(input_ptr + 4 * i); - } - input_ptr += 16; - // Load the accumulators from acc_buffer - float32x4_t acc[4]; - for (int i = 0; i < 4; i++) { - acc[i] = vld1q_f32(acc_buffer_ptr + 4 * i); - } - // Multiply-accumulate - for (int i = 0; i < 4; i++) { - acc[i] = vmlaq_f32(acc[i], input[i], filters_dup2); - } - // Store the accumulators back to acc_buffer - for (int i = 0; i < 4; i++) { - vst1q_f32(acc_buffer_ptr + 4 * i, acc[i]); - } - acc_buffer_ptr += 16; - } - // Handle 4 output pixels at a time. - for (; outp <= num_output_pixels - 4; outp += 4) { - // Load the inputs - float32x4_t input[2]; - for (int i = 0; i < 2; i++) { - input[i] = vld1q_f32(input_ptr + 4 * i); - } - input_ptr += 8; - // Load the accumulators from acc_buffer - float32x4_t acc[2]; - for (int i = 0; i < 2; i++) { - acc[i] = vld1q_f32(acc_buffer_ptr + 4 * i); - } - // Multiply-accumulate - for (int i = 0; i < 2; i++) { - acc[i] = vmlaq_f32(acc[i], input[i], filters_dup2); - } - // Store the accumulators back to acc_buffer - for (int i = 0; i < 2; i++) { - vst1q_f32(acc_buffer_ptr + 4 * i, acc[i]); - } - acc_buffer_ptr += 8; - } - // Handle 2 output pixels at a time. - for (; outp <= num_output_pixels - 2; outp += 2) { - // Load the inputs - const float32x4_t input = vld1q_f32(input_ptr); - input_ptr += 4; - // Load the accumulators from acc_buffer - float32x4_t acc = vld1q_f32(acc_buffer_ptr); - // Multiply-accumulate - acc = vmlaq_f32(acc, input, filters_dup2); - // Store the accumulators back to acc_buffer - vst1q_f32(acc_buffer_ptr, acc); - acc_buffer_ptr += 4; - } - // Handle 1 output pixel at a time - for (; outp < num_output_pixels; outp++) { - // Load the inputs - const float32x2_t input = vld1_f32(input_ptr); - input_ptr += 2; - // Load the accumulators from acc_buffer - float32x2_t acc = vld1_f32(acc_buffer_ptr); - // Multiply-accumulate - acc = vmla_f32(acc, input, filters); - // Store the accumulators back to acc_buffer - vst1_f32(acc_buffer_ptr, acc); - acc_buffer_ptr += 2; - } - } -}; - -template <> -struct FloatDepthwiseConvKernel { - static void Run(int num_output_pixels, int input_depth, int depth_multiplier, - const float* input_ptr, int input_ptr_increment, - const float* filter_ptr, float* acc_buffer_ptr) { - // Handle one output pixel at a time. - for (int outp = 0; outp < num_output_pixels; outp++) { - const float* local_filter_ptr = filter_ptr; - const float* local_input_ptr = input_ptr; - int ic = 0; - // Handle 16 input channels at a time. - for (; ic <= input_depth - 16; ic += 16) { - // Load the filters - float32x4_t filter_0 = vld1q_f32(local_filter_ptr + 4 * 0); - float32x4_t filter_1 = vld1q_f32(local_filter_ptr + 4 * 1); - float32x4_t filter_2 = vld1q_f32(local_filter_ptr + 4 * 2); - float32x4_t filter_3 = vld1q_f32(local_filter_ptr + 4 * 3); - local_filter_ptr += 16; - // Load the inputs - float32x4_t input_0 = vld1q_f32(local_input_ptr + 4 * 0); - float32x4_t input_1 = vld1q_f32(local_input_ptr + 4 * 1); - float32x4_t input_2 = vld1q_f32(local_input_ptr + 4 * 2); - float32x4_t input_3 = vld1q_f32(local_input_ptr + 4 * 3); - local_input_ptr += 16; - // Load the accumulators from acc_buffer - float32x4_t acc_0 = vld1q_f32(acc_buffer_ptr + 4 * 0); - float32x4_t acc_1 = vld1q_f32(acc_buffer_ptr + 4 * 1); - float32x4_t acc_2 = vld1q_f32(acc_buffer_ptr + 4 * 2); - float32x4_t acc_3 = vld1q_f32(acc_buffer_ptr + 4 * 3); - // Multiply-accumulate - acc_0 = vmlaq_f32(acc_0, input_0, filter_0); - acc_1 = vmlaq_f32(acc_1, input_1, filter_1); - acc_2 = vmlaq_f32(acc_2, input_2, filter_2); - acc_3 = vmlaq_f32(acc_3, input_3, filter_3); - // Store the accumulators back to acc_buffer - vst1q_f32(acc_buffer_ptr + 4 * 0, acc_0); - vst1q_f32(acc_buffer_ptr + 4 * 1, acc_1); - vst1q_f32(acc_buffer_ptr + 4 * 2, acc_2); - vst1q_f32(acc_buffer_ptr + 4 * 3, acc_3); - acc_buffer_ptr += 16; - } - // Handle 4 input channels at a time. - for (; ic <= input_depth - 4; ic += 4) { - // Load the filters - float32x4_t filter; - filter = vld1q_f32(local_filter_ptr); - local_filter_ptr += 4; - // Load the inputs - float32x4_t input; - input = vld1q_f32(local_input_ptr); - local_input_ptr += 4; - // Load the accumulators from acc_buffer - float32x4_t acc; - acc = vld1q_f32(acc_buffer_ptr); - // Multiply-accumulate - acc = vmlaq_f32(acc, input, filter); - // Store the accumulators back to acc_buffer - vst1q_f32(acc_buffer_ptr, acc); - acc_buffer_ptr += 4; - } - // Handle one input channel at a time. - for (; ic < input_depth; ic++) { - const float input_val = *local_input_ptr++; - const float filter_val = *local_filter_ptr++; - *acc_buffer_ptr++ += filter_val * input_val; - } - input_ptr += input_ptr_increment; - } - } -}; - -template <> -struct FloatDepthwiseConvKernel { - static void Run(int num_output_pixels, int input_depth, int depth_multiplier, - const float* input_ptr, int input_ptr_increment, - const float* filter_ptr, float* acc_buffer_ptr) { - // Handle one output pixel at a time. - for (int outp = 0; outp < num_output_pixels; outp++) { - const float* local_filter_ptr = filter_ptr; - const float* local_input_ptr = input_ptr; - int ic = 0; - // Handle 2 input channels at a time. - for (; ic <= input_depth - 2; ic += 2) { - // Load the filters - float32x4_t filter[4]; - for (int i = 0; i < 4; i++) { - filter[i] = vld1q_f32(local_filter_ptr + 4 * i); - } - local_filter_ptr += 16; - // Load the inputs - const float32x2_t input = vld1_f32(local_input_ptr); - local_input_ptr += 2; - // Load the accumulators from acc_buffer - float32x4_t acc[4]; - for (int i = 0; i < 4; i++) { - acc[i] = vld1q_f32(acc_buffer_ptr + 4 * i); - } - // Multiply-accumulate - acc[0] = vmlaq_lane_f32(acc[0], filter[0], input, 0); - acc[1] = vmlaq_lane_f32(acc[1], filter[1], input, 0); - acc[2] = vmlaq_lane_f32(acc[2], filter[2], input, 1); - acc[3] = vmlaq_lane_f32(acc[3], filter[3], input, 1); - // Store the accumulators back to acc_buffer - for (int i = 0; i < 4; i++) { - vst1q_f32(acc_buffer_ptr + 4 * i, acc[i]); - } - acc_buffer_ptr += 16; - } - // Handle one input channel at a time. - for (; ic < input_depth; ic++) { - // Load the filters - float32x4_t filter[2]; - for (int i = 0; i < 2; i++) { - filter[i] = vld1q_f32(local_filter_ptr + 4 * i); - } - local_filter_ptr += 8; - // Load the inputs - const float input_val = *local_input_ptr++; - // Load the accumulators from acc_buffer - float32x4_t acc[2]; - for (int i = 0; i < 2; i++) { - acc[i] = vld1q_f32(acc_buffer_ptr + 4 * i); - } - // Multiply-accumulate - for (int i = 0; i < 2; i++) { - acc[i] = vmlaq_n_f32(acc[i], filter[i], input_val); - } - // Store the accumulators back to acc_buffer - for (int i = 0; i < 2; i++) { - vst1q_f32(acc_buffer_ptr + 4 * i, acc[i]); - } - acc_buffer_ptr += 8; - } - input_ptr += input_ptr_increment; - } - } -}; - -// Note this implementation is very slow for input_depths < 8 -// (e.g. comparable to reference implementation) see, specializations for -// input_depth=3 below. -template <> -struct FloatDepthwiseConvKernel { - static void Run(int num_output_pixels, int input_depth, int depth_multiplier, - const float* input_ptr, int input_ptr_increment, - const float* filter_ptr, float* acc_buffer_ptr) { - // Handle one output pixel at a time. - for (int outp = 0; outp < num_output_pixels; outp++) { - const float* local_filter_ptr = filter_ptr; - const float* local_input_ptr = input_ptr; - int ic = 0; - // Handle 8 input channels at a time. - for (; ic <= input_depth - 8; ic += 8) { - // Load the filters - float32x4_t filter[4]; - for (int i = 0; i < 4; i++) { - filter[i] = vld1q_f32(local_filter_ptr + 4 * i); - } - local_filter_ptr += 16; - // Load the inputs - float32x4x2_t input_dup2[2]; - for (int i = 0; i < 2; i++) { - const float32x4_t input = vld1q_f32(local_input_ptr + 4 * i); - input_dup2[i] = vzipq_f32(input, input); - } - local_input_ptr += 8; - // Load the accumulators from acc_buffer - float32x4_t acc[4]; - for (int i = 0; i < 4; i++) { - acc[i] = vld1q_f32(acc_buffer_ptr + 4 * i); - } - // Multiply-accumulate - acc[0] = vmlaq_f32(acc[0], filter[0], input_dup2[0].val[0]); - acc[1] = vmlaq_f32(acc[1], filter[1], input_dup2[0].val[1]); - acc[2] = vmlaq_f32(acc[2], filter[2], input_dup2[1].val[0]); - acc[3] = vmlaq_f32(acc[3], filter[3], input_dup2[1].val[1]); - // Store the accumulators back to acc_buffer - for (int i = 0; i < 4; i++) { - vst1q_f32(acc_buffer_ptr + 4 * i, acc[i]); - } - acc_buffer_ptr += 16; - } - // Handle 4 input channels at a time. - for (; ic <= input_depth - 4; ic += 4) { - // Load the filters - float32x2_t filter[4]; - for (int i = 0; i < 4; i++) { - filter[i] = vld1_f32(local_filter_ptr + 2 * i); - } - local_filter_ptr += 8; - // Load the inputs - const float32x4_t input = vld1q_f32(local_input_ptr); - local_input_ptr += 4; - // Load the accumulators from acc_buffer - float32x2_t acc[4]; - for (int i = 0; i < 4; i++) { - acc[i] = vld1_f32(acc_buffer_ptr + 2 * i); - } - // Multiply-accumulate - acc[0] = vmla_lane_f32(acc[0], filter[0], vget_low_f32(input), 0); - acc[1] = vmla_lane_f32(acc[1], filter[1], vget_low_f32(input), 1); - acc[2] = vmla_lane_f32(acc[2], filter[2], vget_high_f32(input), 0); - acc[3] = vmla_lane_f32(acc[3], filter[3], vget_high_f32(input), 1); - // Store the accumulators back to acc_buffer - for (int i = 0; i < 4; i++) { - vst1_f32(acc_buffer_ptr + 2 * i, acc[i]); - } - acc_buffer_ptr += 8; - } - // Handle 2 input channels at a time. - for (; ic <= input_depth - 2; ic += 2) { - // Load the filters - const float32x4_t filter = vld1q_f32(local_filter_ptr); - local_filter_ptr += 4; - // Load the inputs - const float32x2_t input = vld1_f32(local_input_ptr); - local_input_ptr += 2; - // Load the accumulators from acc_buffer - float32x2_t acc[2]; - for (int i = 0; i < 2; i++) { - acc[i] = vld1_f32(acc_buffer_ptr + 2 * i); - } - // Multiply-accumulate - acc[0] = vmla_lane_f32(acc[0], vget_low_f32(filter), input, 0); - acc[1] = vmla_lane_f32(acc[1], vget_high_f32(filter), input, 1); - // Store the accumulators back to acc_buffer - for (int i = 0; i < 2; i++) { - vst1_f32(acc_buffer_ptr + 2 * i, acc[i]); - } - acc_buffer_ptr += 4; - } - // Handle one input channel at a time. - for (; ic < input_depth; ic++) { - // Load the inputs - const float input_val = *local_input_ptr++; - // Multiply-accumulate - for (int i = 0; i < 2; i++) { - acc_buffer_ptr[i] += local_filter_ptr[i] * input_val; - } - local_filter_ptr += 2; - acc_buffer_ptr += 2; - } - input_ptr += input_ptr_increment; - } - } -}; - -template <> -struct FloatDepthwiseConvKernel { - static void Run(int num_output_pixels, int input_depth, int depth_multiplier, - const float* input_ptr, int input_ptr_increment, - const float* filter_ptr, float* acc_buffer_ptr) { - // Load the filters - float32x2_t filter[3]; - for (int i = 0; i < 3; i++) { - filter[i] = vld1_f32(filter_ptr + 2 * i); - } - // Handle one output pixel at a time. - for (int outp = 0; outp < num_output_pixels; outp++) { - const float32x2_t input01 = vld1_f32(input_ptr); - const float32x2_t input2 = vld1_dup_f32(input_ptr + 2); - // Load the accumulators from acc_buffer - float32x2_t acc[3]; - for (int i = 0; i < 3; i++) { - acc[i] = vld1_f32(acc_buffer_ptr + 2 * i); - } - // Multiply-accumulate for each input channel there 2 outputs - acc[0] = vmla_lane_f32(acc[0], filter[0], input01, 0); - acc[1] = vmla_lane_f32(acc[1], filter[1], input01, 1); - acc[2] = vmla_lane_f32(acc[2], filter[2], input2, 0); - // Store the accumulators back to acc_buffer - for (int i = 0; i < 3; i++) { - vst1_f32(acc_buffer_ptr + 2 * i, acc[i]); - } - acc_buffer_ptr += 6; - input_ptr += input_ptr_increment; - } - } -}; - -template <> -struct FloatDepthwiseConvKernel { - static void Run(int num_output_pixels, int input_depth, int depth_multiplier, - const float* input_ptr, int input_ptr_increment, - const float* filter_ptr, float* acc_buffer_ptr) { - // Load the filters - float32x4_t filter[3]; - for (int i = 0; i < 3; i++) { - filter[i] = vld1q_f32(filter_ptr + 4 * i); - } - // Handle one output pixel at a time. - for (int outp = 0; outp < num_output_pixels; outp++) { - // NOTE: we only want 3 values, so we read it as two ops where - // the second op just duplicates the lane - const float32x2_t input01 = vld1_f32(input_ptr); - const float32x2_t input2 = vld1_dup_f32(input_ptr + 2); - // Load the accumulators from acc_buffer - float32x4_t acc[3]; - for (int i = 0; i < 3; i++) { - acc[i] = vld1q_f32(acc_buffer_ptr + 4 * i); - } - // Multiply-accumulate all outputs. - acc[0] = vmlaq_lane_f32(acc[0], filter[0], input01, 0); - acc[1] = vmlaq_lane_f32(acc[1], filter[1], input01, 1); - acc[2] = vmlaq_lane_f32(acc[2], filter[2], input2, 0); - // Store the accumulators back to acc_buffer - for (int i = 0; i < 3; i++) { - vst1q_f32(acc_buffer_ptr + 4 * i, acc[i]); - } - acc_buffer_ptr += 12; - input_ptr += input_ptr_increment; - } - } -}; - -template <> -struct FloatDepthwiseConvKernel { - static void Run(int num_output_pixels, int input_depth, int depth_multiplier, - const float* input_ptr, int input_ptr_increment, - const float* filter_ptr, float* acc_buffer_ptr) { - // Load the filters - float32x4_t filter[2]; - for (int i = 0; i < 2; i++) { - filter[i] = vld1q_f32(filter_ptr + 4 * i); - } - // Handle one output pixel at a time. - for (int outp = 0; outp < num_output_pixels; outp++) { - // Load the inputs - const float input_val = *input_ptr; - input_ptr += input_ptr_increment; - // Load the accumulators from acc_buffer - float32x4_t acc[2]; - for (int i = 0; i < 2; i++) { - acc[i] = vld1q_f32(acc_buffer_ptr + 4 * i); - } - // Multiply-accumulate - for (int i = 0; i < 2; i++) { - acc[i] = vmlaq_n_f32(acc[i], filter[i], input_val); - } - // Store the accumulators back to acc_buffer - for (int i = 0; i < 2; i++) { - vst1q_f32(acc_buffer_ptr + 4 * i, acc[i]); - } - acc_buffer_ptr += 8; - } - } -}; - -template <> -struct FloatDepthwiseConvKernel { - static void Run(int num_output_pixels, int input_depth, int depth_multiplier, - const float* input_ptr, int input_ptr_increment, - const float* filter_ptr, float* acc_buffer_ptr) { - // Load the filters - float32x4_t filter_0 = vld1q_f32(filter_ptr + 4 * 0); - float32x4_t filter_1 = vld1q_f32(filter_ptr + 4 * 1); - float32x4_t filter_2 = vld1q_f32(filter_ptr + 4 * 2); - float32x4_t filter_3 = vld1q_f32(filter_ptr + 4 * 3); - float32x4_t filter_4 = vld1q_f32(filter_ptr + 4 * 4); - float32x4_t filter_5 = vld1q_f32(filter_ptr + 4 * 5); - float32x4_t filter_6 = vld1q_f32(filter_ptr + 4 * 6); - float32x4_t filter_7 = vld1q_f32(filter_ptr + 4 * 7); - - // Handle one output pixel at a time. - for (int outp = 0; outp < num_output_pixels; outp++) { - // Load the inputs - const float input_val = *input_ptr; - input_ptr += input_ptr_increment; - // Load the accumulators from acc_buffer - float32x4_t acc_0 = vld1q_f32(acc_buffer_ptr + 4 * 0); - float32x4_t acc_1 = vld1q_f32(acc_buffer_ptr + 4 * 1); - float32x4_t acc_2 = vld1q_f32(acc_buffer_ptr + 4 * 2); - float32x4_t acc_3 = vld1q_f32(acc_buffer_ptr + 4 * 3); - float32x4_t acc_4 = vld1q_f32(acc_buffer_ptr + 4 * 4); - float32x4_t acc_5 = vld1q_f32(acc_buffer_ptr + 4 * 5); - float32x4_t acc_6 = vld1q_f32(acc_buffer_ptr + 4 * 6); - float32x4_t acc_7 = vld1q_f32(acc_buffer_ptr + 4 * 7); - // Multiply-accumulate - acc_0 = vmlaq_n_f32(acc_0, filter_0, input_val); - acc_1 = vmlaq_n_f32(acc_1, filter_1, input_val); - acc_2 = vmlaq_n_f32(acc_2, filter_2, input_val); - acc_3 = vmlaq_n_f32(acc_3, filter_3, input_val); - acc_4 = vmlaq_n_f32(acc_4, filter_4, input_val); - acc_5 = vmlaq_n_f32(acc_5, filter_5, input_val); - acc_6 = vmlaq_n_f32(acc_6, filter_6, input_val); - acc_7 = vmlaq_n_f32(acc_7, filter_7, input_val); - // Store the accumulators back to acc_buffer - vst1q_f32(acc_buffer_ptr + 4 * 0, acc_0); - vst1q_f32(acc_buffer_ptr + 4 * 1, acc_1); - vst1q_f32(acc_buffer_ptr + 4 * 2, acc_2); - vst1q_f32(acc_buffer_ptr + 4 * 3, acc_3); - vst1q_f32(acc_buffer_ptr + 4 * 4, acc_4); - vst1q_f32(acc_buffer_ptr + 4 * 5, acc_5); - vst1q_f32(acc_buffer_ptr + 4 * 6, acc_6); - vst1q_f32(acc_buffer_ptr + 4 * 7, acc_7); - acc_buffer_ptr += 32; - } - } -}; - -template <> -struct FloatDepthwiseConvKernel { - static void Run(int num_output_pixels, int input_depth, int depth_multiplier, - const float* input_ptr, int input_ptr_increment, - const float* filter_ptr, float* acc_buffer_ptr) { - // Load the filters - float32x4_t filter_0 = vld1q_f32(filter_ptr + 4 * 0); - float32x4_t filter_1 = vld1q_f32(filter_ptr + 4 * 1); - float32x4_t filter_2 = vld1q_f32(filter_ptr + 4 * 2); - float32x4_t filter_3 = vld1q_f32(filter_ptr + 4 * 3); - float32x4_t filter_4 = vld1q_f32(filter_ptr + 4 * 4); - - // Handle one output pixel at a time. - for (int outp = 0; outp < num_output_pixels; outp++) { - // Load the inputs - const float input_val = *input_ptr; - input_ptr += input_ptr_increment; - // Load the accumulators from acc_buffer - float32x4_t acc_0 = vld1q_f32(acc_buffer_ptr + 4 * 0); - float32x4_t acc_1 = vld1q_f32(acc_buffer_ptr + 4 * 1); - float32x4_t acc_2 = vld1q_f32(acc_buffer_ptr + 4 * 2); - float32x4_t acc_3 = vld1q_f32(acc_buffer_ptr + 4 * 3); - float32x4_t acc_4 = vld1q_f32(acc_buffer_ptr + 4 * 4); - // Multiply-accumulate - acc_0 = vmlaq_n_f32(acc_0, filter_0, input_val); - acc_1 = vmlaq_n_f32(acc_1, filter_1, input_val); - acc_2 = vmlaq_n_f32(acc_2, filter_2, input_val); - acc_3 = vmlaq_n_f32(acc_3, filter_3, input_val); - acc_4 = vmlaq_n_f32(acc_4, filter_4, input_val); - // Store the accumulators back to acc_buffer - vst1q_f32(acc_buffer_ptr + 4 * 0, acc_0); - vst1q_f32(acc_buffer_ptr + 4 * 1, acc_1); - vst1q_f32(acc_buffer_ptr + 4 * 2, acc_2); - vst1q_f32(acc_buffer_ptr + 4 * 3, acc_3); - vst1q_f32(acc_buffer_ptr + 4 * 4, acc_4); - acc_buffer_ptr += 20; - } - } -}; - -template <> -struct FloatDepthwiseConvKernel { - static void Run(int num_output_pixels, int input_depth, int depth_multiplier, - const float* input_ptr, int input_ptr_increment, - const float* filter_ptr, float* acc_buffer_ptr) { - // Handle one output pixel at a time. - for (int outp = 0; outp < num_output_pixels; outp++) { - const float* local_filter_ptr = filter_ptr; - const float* local_input_ptr = input_ptr; - for (int ic = 0; ic < input_depth; ic++) { - // Load the filters - float32x4_t filter[4]; - for (int i = 0; i < 4; i++) { - filter[i] = vld1q_f32(local_filter_ptr + 4 * i); - } - local_filter_ptr += 16; - // Load the inputs - const float input_val = *local_input_ptr++; - // Load the accumulators from acc_buffer - float32x4_t acc[4]; - for (int i = 0; i < 4; i++) { - acc[i] = vld1q_f32(acc_buffer_ptr + 4 * i); - } - // Multiply-accumulate - for (int i = 0; i < 4; i++) { - acc[i] = vmlaq_n_f32(acc[i], filter[i], input_val); - } - // Store the accumulators back to acc_buffer - for (int i = 0; i < 4; i++) { - vst1q_f32(acc_buffer_ptr + 4 * i, acc[i]); - } - acc_buffer_ptr += 16; - } - input_ptr += input_ptr_increment; - } - } -}; - -template <> -struct FloatDepthwiseConvKernel { - static void Run(int num_output_pixels, int input_depth, int depth_multiplier, - const float* input_ptr, int input_ptr_increment, - const float* filter_ptr, float* acc_buffer_ptr) { - // Load the filters - float32x4_t filter[2]; - for (int i = 0; i < 2; i++) { - filter[i] = vld1q_f32(filter_ptr + 4 * i); - } - // Handle one output pixel at a time. - for (int outp = 0; outp < num_output_pixels; outp++) { - // Load the inputs - float32x4_t input[2]; - for (int i = 0; i < 2; i++) { - input[i] = vld1q_f32(input_ptr + 4 * i); - } - // Load the accumulators from acc_buffer - float32x4_t acc[2]; - for (int i = 0; i < 2; i++) { - acc[i] = vld1q_f32(acc_buffer_ptr + 4 * i); - } - // Multiply-accumulate - for (int i = 0; i < 2; i++) { - acc[i] = vmlaq_f32(acc[i], input[i], filter[i]); - } - // Store the accumulators back to acc_buffer - for (int i = 0; i < 2; i++) { - vst1q_f32(acc_buffer_ptr + 4 * i, acc[i]); - } - acc_buffer_ptr += 8; - input_ptr += input_ptr_increment; - } - } -}; - -template <> -struct FloatDepthwiseConvKernel { - static void Run(int num_output_pixels, int input_depth, int depth_multiplier, - const float* input_ptr, int input_ptr_increment, - const float* filter_ptr, float* acc_buffer_ptr) { - float32x2_t filter = vld1_f32(filter_ptr); - float32x4_t filter_x4 = vcombine_f32(filter, filter); - int outp = 0; - - // Handle two output pixels at a time. - for (; outp <= num_output_pixels - 2; outp += 2) { - // Load the inputs - float32x2_t input_1 = vld1_f32(input_ptr); - input_ptr += input_ptr_increment; - float32x2_t input_2 = vld1_f32(input_ptr); - input_ptr += input_ptr_increment; - float32x4_t input = vcombine_f32(input_1, input_2); - - // Load the accumulators from acc_buffer - float32x4_t acc = vld1q_f32(acc_buffer_ptr); - - // Multiply-accumulate - acc = vmlaq_f32(acc, input, filter_x4); - - // Store the accumulators back to acc_buffer - vst1q_f32(acc_buffer_ptr, acc); - acc_buffer_ptr += 4; - } - // Handle one output pixel at a time. - for (; outp < num_output_pixels; outp++) { - // Load the inputs - float32x2_t input = vld1_f32(input_ptr); - input_ptr += input_ptr_increment; - - // Load the accumulators from acc_buffer - float32x2_t acc = vld1_f32(acc_buffer_ptr); - - // Multiply-accumulate - acc = vmla_f32(acc, input, filter); - - // Store the accumulators back to acc_buffer - vst1_f32(acc_buffer_ptr, acc); - acc_buffer_ptr += 2; - } - } -}; - -template <> -struct FloatDepthwiseConvKernel { - static void Run(int num_output_pixels, int input_depth, int depth_multiplier, - const float* input_ptr, int input_ptr_increment, - const float* filter_ptr, float* acc_buffer_ptr) { - float32x4_t filter = vld1q_f32(filter_ptr); - - // Handle one output pixel at a time. - for (int outp = 0; outp < num_output_pixels; outp++) { - // Load the inputs - float32x4_t input = vld1q_f32(input_ptr); - // Load the accumulators from acc_buffer - float32x4_t acc = vld1q_f32(acc_buffer_ptr); - // Multiply-accumulate - acc = vmlaq_f32(acc, input, filter); - // Store the accumulators back to acc_buffer - vst1q_f32(acc_buffer_ptr, acc); - acc_buffer_ptr += 4; - input_ptr += input_ptr_increment; - } - } -}; -#endif - -// Accumulates the effect of one row of the filter, on a segment of one row -// of the output, accessing the corresponding one row of the input. -template -void FloatDepthwiseConvAccumRow(int stride, int dilation_factor, - int input_depth, int input_width, - const float* input_data, int pad_width, - int depth_multiplier, int filter_width, - const float* filter_data, - int out_x_buffer_start, int out_x_buffer_end, - int output_depth, float* acc_buffer) { -#ifdef GEMMLOWP_PROFILING - gemmlowp::ScopedProfilingLabel label(__PRETTY_FUNCTION__); -#endif - // Sanity check parameters. This is important in particular to ensure - // that we keep the number of template instantiations minimal, so we don't - // increase binary size unnecessarily. - static_assert(kFixedDepthMultiplier || !kFixedInputDepth, ""); - static_assert(kFixedInputDepth || kAllowStrided, ""); - TFLITE_DCHECK(stride == 1 || kAllowStrided); - if (kFixedInputDepth) { - TFLITE_DCHECK_EQ(input_depth, kFixedInputDepth); - } - if (kFixedDepthMultiplier) { - TFLITE_DCHECK_EQ(depth_multiplier, kFixedDepthMultiplier); - } - TFLITE_DCHECK_EQ(output_depth, input_depth * depth_multiplier); - const int input_ptr_increment = stride * input_depth; - const float* filter_base_ptr = filter_data; - for (int filter_x = 0; filter_x < filter_width; ++filter_x) { - // For the current (filter_x, filter_y) point in the filter, - // compute the boundaries of the corresponding output row segment. - int out_x_loop_start_unclampled = 0; - int out_x_loop_end_unclampled = 0; - if (kAllowStrided) { - if (stride == 2) { - out_x_loop_start_unclampled = (pad_width - filter_x + 1) / 2; - out_x_loop_end_unclampled = - (pad_width + input_width - filter_x + 1) / 2; - } else if (stride == 4) { - out_x_loop_start_unclampled = (pad_width - filter_x + 3) / 4; - out_x_loop_end_unclampled = - (pad_width + input_width - filter_x + 3) / 4; - } else { - out_x_loop_start_unclampled = - (pad_width - filter_x + stride - 1) / stride; - out_x_loop_end_unclampled = - (pad_width + input_width - filter_x + stride - 1) / stride; - } - } else { - out_x_loop_start_unclampled = pad_width - filter_x; - out_x_loop_end_unclampled = pad_width + input_width - filter_x; - } - // The kernel will have to iterate on the segment of the - // output row that starts at out_x_loop_start and out_x_loop_end. - const int out_x_loop_start = - std::max(out_x_buffer_start, out_x_loop_start_unclampled); - const int out_x_loop_end = - std::min(out_x_buffer_end, out_x_loop_end_unclampled); - - float* acc_buffer_ptr = - acc_buffer + (out_x_loop_start - out_x_buffer_start) * output_depth; - const int in_x_origin = (out_x_loop_start * stride) - pad_width + filter_x; - const float* input_ptr = input_data + in_x_origin * input_depth; - const int num_output_pixels = out_x_loop_end - out_x_loop_start; - FloatDepthwiseConvKernel::Run(num_output_pixels, - input_depth, - depth_multiplier, - input_ptr, - input_ptr_increment, - filter_base_ptr, - acc_buffer_ptr); - filter_base_ptr += output_depth; - } -} - -// generic fallback of FloatDepthwiseConvAccumRow, portable, non-templatized. -inline void FloatDepthwiseConvAccumRowGeneric( - int stride, int dilation_factor, int input_depth, int input_width, - const float* input_data, int pad_width, int depth_multiplier, - int filter_width, const float* filter_data, int out_x_buffer_start, - int out_x_buffer_end, int output_depth, float* acc_buffer) { - gemmlowp::ScopedProfilingLabel label("DepthwiseConvAccumRowGeneric (slow)"); - const float* filter_base_ptr = filter_data; - for (int filter_x = 0; filter_x < filter_width; ++filter_x) { - const int out_x_loop_start = std::max( - out_x_buffer_start, - (pad_width - dilation_factor * filter_x + stride - 1) / stride); - const int out_x_loop_end = std::min( - out_x_buffer_end, - (pad_width + input_width - dilation_factor * filter_x + stride - 1) / - stride); - - float* acc_buffer_ptr = - acc_buffer + (out_x_loop_start - out_x_buffer_start) * output_depth; - const int in_x_origin = - (out_x_loop_start * stride) - pad_width + dilation_factor * filter_x; - const float* input_ptr = input_data + in_x_origin * input_depth; - const int input_ptr_increment = (stride - 1) * input_depth; - for (int out_x = out_x_loop_start; out_x < out_x_loop_end; out_x++) { - const float* filter_ptr = filter_base_ptr; - for (int ic = 0; ic < input_depth; ++ic) { - const float input_val = *input_ptr++; - for (int m = 0; m < depth_multiplier; m++) { - const float filter_val = *filter_ptr++; - *acc_buffer_ptr++ += filter_val * input_val; - } - } - input_ptr += input_ptr_increment; - } - filter_base_ptr += output_depth; - } -} - -// Initializes the accumulator buffer with bias values. -inline void DepthwiseConvInitAccBuffer(int num_output_pixels, int output_depth, - const float* bias_data, - float* acc_buffer) { - // TODO(benoitjacob): This might need optimized specializations - // for small output_depth values, if that ever becomes an important - // case (like it was for some quantized DepthwiseConv cases). - for (int i = 0; i < num_output_pixels; i++) { - memcpy(acc_buffer + i * output_depth, bias_data, - sizeof(acc_buffer[0]) * output_depth); - } -} - -inline void DepthwiseConv( - const DepthwiseParams& params, const RuntimeShape& input_shape, - const float* input_data, const RuntimeShape& filter_shape, - const float* filter_data, const RuntimeShape& bias_shape, - const float* bias_data, const RuntimeShape& output_shape, - float* output_data) { - gemmlowp::ScopedProfilingLabel label("DepthwiseConv"); - const int stride_width = params.stride_width; - const int stride_height = params.stride_height; - const int pad_width = params.padding_values.width; - const int pad_height = params.padding_values.height; - const int depth_multiplier = params.depth_multiplier; - const float output_activation_min = params.float_activation_min; - const float output_activation_max = params.float_activation_max; - const int dilation_width_factor = params.dilation_width_factor; - const int dilation_height_factor = params.dilation_height_factor; - TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4); - TFLITE_DCHECK_EQ(filter_shape.DimensionsCount(), 4); - TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4); - - const int batches = MatchingDim(input_shape, 0, output_shape, 0); - const int output_depth = MatchingDim(filter_shape, 3, output_shape, 3); - const int input_height = input_shape.Dims(1); - const int input_width = input_shape.Dims(2); - const int input_depth = input_shape.Dims(3); - const int filter_height = filter_shape.Dims(1); - const int filter_width = filter_shape.Dims(2); - const int output_height = output_shape.Dims(1); - const int output_width = output_shape.Dims(2); - TFLITE_DCHECK_EQ(output_depth, input_depth * depth_multiplier); - TFLITE_DCHECK_EQ(bias_shape.FlatSize(), output_depth); - - static const int kAccBufferMaxSize = 2048; - float acc_buffer[kAccBufferMaxSize]; - TFLITE_DCHECK_GE(kAccBufferMaxSize, output_depth); - const int kOutputPixelsInAccBuffer = kAccBufferMaxSize / output_depth; - const int kAccBufferActualSize = kOutputPixelsInAccBuffer * output_depth; - TFLITE_DCHECK_LE(kOutputPixelsInAccBuffer * output_depth, - kAccBufferActualSize); - TFLITE_DCHECK_LE(kAccBufferActualSize, kAccBufferMaxSize); - TFLITE_DCHECK_GE(kOutputPixelsInAccBuffer, 1); - - // row_accum_func will point to the core accumulation function to be used - // for this DepthwiseConv op. - using row_accum_func_t = decltype(&FloatDepthwiseConvAccumRowGeneric); - row_accum_func_t row_accum_func = nullptr; - -#define TFMINI_USE_DEPTHWISECONV_KERNEL(ALLOW_STRIDED, FIXED_INPUT_DEPTH, \ - FIXED_DEPTH_MULTIPLIER) \ - if (!row_accum_func && (stride_width == 1 || ALLOW_STRIDED) && \ - (input_depth == FIXED_INPUT_DEPTH || FIXED_INPUT_DEPTH == 0) && \ - depth_multiplier == FIXED_DEPTH_MULTIPLIER && \ - dilation_height_factor == 1 && dilation_width_factor == 1) { \ - row_accum_func = \ - FloatDepthwiseConvAccumRow; \ - } - -#ifdef USE_NEON - // We go over our list of kernels by decreasing order of preference - // for the cases where multiple kernels could apply. - - // Start with the fastest kernels: AllowStrided=false, fixed input depth. - - TFMINI_USE_DEPTHWISECONV_KERNEL(false, 8, 1) - TFMINI_USE_DEPTHWISECONV_KERNEL(false, 2, 1) - - // Next come the strided kernels: AllowStrided=true, fixed input depth. - // They are a bit less efficient, but allow stride!=1. - - TFMINI_USE_DEPTHWISECONV_KERNEL(true, 8, 1) - TFMINI_USE_DEPTHWISECONV_KERNEL(true, 1, 8) - TFMINI_USE_DEPTHWISECONV_KERNEL(true, 1, 20) - TFMINI_USE_DEPTHWISECONV_KERNEL(true, 1, 32) - TFMINI_USE_DEPTHWISECONV_KERNEL(true, 2, 1) - TFMINI_USE_DEPTHWISECONV_KERNEL(true, 3, 2) - TFMINI_USE_DEPTHWISECONV_KERNEL(true, 3, 4) - TFMINI_USE_DEPTHWISECONV_KERNEL(true, 4, 1) - - // Finally, the kernels allowing a variable input depth, - // these are the least efficient but most general kernels. - - TFMINI_USE_DEPTHWISECONV_KERNEL(true, 0, 1) - TFMINI_USE_DEPTHWISECONV_KERNEL(true, 0, 2) - TFMINI_USE_DEPTHWISECONV_KERNEL(true, 0, 8) - TFMINI_USE_DEPTHWISECONV_KERNEL(true, 0, 16) - -#endif // USE_NEON - -#undef TFMINI_USE_DEPTHWISECONV_KERNEL - - // No matching fast kernel found, use slow fallback. - if (!row_accum_func) { - row_accum_func = FloatDepthwiseConvAccumRowGeneric; - } - - const int input_height_stride = input_shape.Dims(3) * input_shape.Dims(2); - const int input_batch_stride = input_height_stride * input_shape.Dims(1); - const int filter_height_stride = filter_shape.Dims(3) * filter_shape.Dims(2); - - // Now that we have determined row_accum_func, we can start work. - float* output_ptr = output_data; - for (int b = 0; b < batches; ++b) { - for (int out_y = 0; out_y < output_height; ++out_y) { - const int in_y_origin = (out_y * stride_height) - pad_height; - const int filter_y_start = - std::max(0, (-in_y_origin + dilation_height_factor - 1) / - dilation_height_factor); - const int filter_y_end = - std::min(filter_height, - (input_height - in_y_origin + dilation_height_factor - 1) / - dilation_height_factor); - for (int out_x_buffer_start = 0; out_x_buffer_start < output_width; - out_x_buffer_start += kOutputPixelsInAccBuffer) { - const int out_x_buffer_end = std::min( - output_width, out_x_buffer_start + kOutputPixelsInAccBuffer); - // We call a 'pixel' a group of activation that share all but the - // 'depth'/'channel' coordinate. num_output_pixels is the number of - // output pixels that we will accumulate in this loop iteration. - const int num_output_pixels = out_x_buffer_end - out_x_buffer_start; - // Initialize our local accumulator with the bias values, so we don't - // have to add them later. - DepthwiseConvInitAccBuffer(num_output_pixels, output_depth, bias_data, - acc_buffer); - // Accumulation loop. Most of the time should be spent in here. - for (int filter_y = filter_y_start; filter_y < filter_y_end; - ++filter_y) { - const int in_y = in_y_origin + dilation_height_factor * filter_y; - row_accum_func( - stride_width, dilation_width_factor, input_depth, input_width, - input_data + in_y * input_height_stride + b * input_batch_stride, - pad_width, depth_multiplier, filter_width, - filter_data + filter_y * filter_height_stride, out_x_buffer_start, - out_x_buffer_end, output_depth, acc_buffer); - } - // Finished accumulating. Now store to destination. - const int num_output_values = output_depth * num_output_pixels; - int i = 0; -// TODO(benoitjacob) optimized code goes here -#ifdef USE_NEON - // Handle 16 values at a time - for (; i <= num_output_values - 16; i += 16) { - float32x4_t acc[4]; - for (int k = 0; k < 4; k++) { - acc[k] = vld1q_f32(acc_buffer + i + 4 * k); - } - for (int k = 0; k < 4; k++) { - acc[k] = vmaxq_f32( - vdupq_n_f32(output_activation_min), - vminq_f32(vdupq_n_f32(output_activation_max), acc[k])); - } - for (int k = 0; k < 4; k++) { - vst1q_f32(output_ptr + 4 * k, acc[k]); - } - output_ptr += 16; - } - // Handle 4 values at a time - for (; i <= num_output_values - 4; i += 4) { - float32x4_t acc = vld1q_f32(acc_buffer + i); - - acc = vmaxq_f32(vdupq_n_f32(output_activation_min), - vminq_f32(vdupq_n_f32(output_activation_max), acc)); - - vst1q_f32(output_ptr, acc); - output_ptr += 4; - } -#endif - // Handle leftover values, one by one. This is very slow. - for (; i < num_output_values; i++) { - float acc = acc_buffer[i]; - acc = std::max(output_activation_min, - std::min(output_activation_max, acc)); - - *output_ptr++ = acc; - } - } - } - } -} - -} // namespace optimized_ops -} // namespace tflite - -#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_DEPTHWISECONV_FLOAT_H_ diff --git a/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_uint8.h b/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_uint8.h deleted file mode 100644 index eff9cab4778a6865f93338489df28707fa3afae2..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_uint8.h +++ /dev/null @@ -1,1972 +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_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_DEPTHWISECONV_UINT8_H_ -#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_DEPTHWISECONV_UINT8_H_ - -#include "fixedpoint/fixedpoint.h" -#include "public/gemmlowp.h" -#include "tensorflow/contrib/lite/kernels/internal/common.h" -#include "tensorflow/contrib/lite/kernels/internal/optimized/depthwiseconv_uint8_3x3_filter.h" -#include "tensorflow/contrib/lite/kernels/internal/types.h" - -namespace tflite { -namespace optimized_ops { - -// Implementation of quantized DepthwiseConv - -template -struct QuantizedDepthwiseConvKernel {}; - -#ifdef USE_NEON -template <> -struct QuantizedDepthwiseConvKernel { - static void Run(int num_output_pixels, int input_depth, int depth_multiplier, - const uint8* input_ptr, int16 input_offset, - int input_ptr_increment, const uint8* filter_ptr, - int16 filter_offset, int32* acc_buffer_ptr) { - // Load the filters, add filter_offset. - uint8x8x2_t filter_u8; - filter_u8.val[0] = vld1_u8(filter_ptr); - filter_u8.val[1] = vld1_u8(filter_ptr + 8); - int16x8_t filter[2]; - for (int i = 0; i < 2; i++) { - filter[i] = vaddq_s16(vreinterpretq_s16_u16(vmovl_u8(filter_u8.val[i])), - vdupq_n_s16(filter_offset)); - } - // Handle one output pixel at a time. - for (int outp = 0; outp < num_output_pixels; outp++) { - // Load the accumulators from acc_buffer - int32x4x2_t acc[2]; - for (int i = 0; i < 2; i++) { - acc[i].val[0] = vld1q_s32(acc_buffer_ptr + 4 * i); - acc[i].val[1] = vld1q_s32(acc_buffer_ptr + 4 * i + 8); - } - // Load the inputs, add input_offset. - const uint8x8_t input_u8 = vld1_u8(input_ptr); - input_ptr += input_ptr_increment; - const int16x8_t input_s16 = vreinterpretq_s16_u16(vmovl_u8(input_u8)); - const int16x8_t input = vaddq_s16(input_s16, vdupq_n_s16(input_offset)); - // Duplicate the input values, 2-fold - const int16x8x2_t input_dup2 = vzipq_s16(input, input); - // Multiply-accumulate - for (int i = 0; i < 2; i++) { - acc[0].val[i] = vmlal_s16(acc[0].val[i], vget_low_s16(filter[i]), - vget_low_s16(input_dup2.val[i])); - acc[1].val[i] = vmlal_s16(acc[1].val[i], vget_high_s16(filter[i]), - vget_high_s16(input_dup2.val[i])); - } - // Store the accumulators back to acc_buffer - for (int i = 0; i < 2; i++) { - vst1q_s32(acc_buffer_ptr + 4 * i, acc[i].val[0]); - vst1q_s32(acc_buffer_ptr + 4 * i + 8, acc[i].val[1]); - } - acc_buffer_ptr += 16; - } - } -}; - -template <> -struct QuantizedDepthwiseConvKernel { - static void Run(int num_output_pixels, int input_depth, int depth_multiplier, - const uint8* input_ptr, int16 input_offset, - int input_ptr_increment, const uint8* filter_ptr, - int16 filter_offset, int32* acc_buffer_ptr) { - // Load the filters, add filter_offset. - const uint8x8_t filter_u8 = vld1_u8(filter_ptr); - const int16x8_t filter_s16 = vreinterpretq_s16_u16(vmovl_u8(filter_u8)); - const int16x8_t filter = vaddq_s16(filter_s16, vdupq_n_s16(filter_offset)); - - int outp = 0; - // Handle 2 output pixels at a time. - for (; outp <= num_output_pixels - 2; outp += 2) { - // Load the accumulators from acc_buffer. - int32x4_t acc[4]; - for (int i = 0; i < 4; i++) { - acc[i] = vld1q_s32(acc_buffer_ptr + 4 * i); - } - // Load the inputs, add input_offset. - uint8x8_t input_u8[2]; - for (int i = 0; i < 2; i++) { - input_u8[i] = vld1_u8(input_ptr + 8 * i); - } - input_ptr += 16; - int16x8_t input[2]; - for (int i = 0; i < 2; i++) { - input[i] = vreinterpretq_s16_u16(vmovl_u8(input_u8[i])); - } - for (int i = 0; i < 2; i++) { - input[i] = vaddq_s16(input[i], vdupq_n_s16(input_offset)); - } - // Multiply-accumulate. - acc[0] = vmlal_s16(acc[0], vget_low_s16(filter), vget_low_s16(input[0])); - acc[1] = - vmlal_s16(acc[1], vget_high_s16(filter), vget_high_s16(input[0])); - acc[2] = vmlal_s16(acc[2], vget_low_s16(filter), vget_low_s16(input[1])); - acc[3] = - vmlal_s16(acc[3], vget_high_s16(filter), vget_high_s16(input[1])); - // Store the accumulators back to acc_buffer - for (int i = 0; i < 4; i++) { - vst1q_s32(acc_buffer_ptr + 4 * i, acc[i]); - } - acc_buffer_ptr += 16; - } - // Handle 1 output pixel at a time. - for (; outp < num_output_pixels; outp++) { - // Load the accumulators from acc_buffer. - int32x4_t acc[2]; - acc[0] = vld1q_s32(acc_buffer_ptr); - acc[1] = vld1q_s32(acc_buffer_ptr + 4); - - // Load the inputs, add input_offset. - const uint8x8_t input_u8 = vld1_u8(input_ptr); - input_ptr += 8; - const int16x8_t input_s16 = vreinterpretq_s16_u16(vmovl_u8(input_u8)); - const int16x8_t input = vaddq_s16(input_s16, vdupq_n_s16(input_offset)); - // Multiply-accumulate. - acc[0] = vmlal_s16(acc[0], vget_low_s16(filter), vget_low_s16(input)); - acc[1] = vmlal_s16(acc[1], vget_high_s16(filter), vget_high_s16(input)); - // Store the accumulators back to acc_buffer - vst1q_s32(acc_buffer_ptr, acc[0]); - vst1q_s32(acc_buffer_ptr + 4, acc[1]); - acc_buffer_ptr += 8; - } - } -}; - -template <> -struct QuantizedDepthwiseConvKernel { - static void Run(int num_output_pixels, int input_depth, int depth_multiplier, - const uint8* input_ptr, int16 input_offset, - int input_ptr_increment, const uint8* filter_ptr, - int16 filter_offset, int32* acc_buffer_ptr) { - // Load the filters, add filter_offset. - const uint8x8_t filter_u8 = vld1_u8(filter_ptr); - const int16x8_t filter_s16 = vreinterpretq_s16_u16(vmovl_u8(filter_u8)); - const int16x8_t filter = vaddq_s16(filter_s16, vdupq_n_s16(filter_offset)); - - int outp = 0; - // Handle 2 output pixels at a time. - for (; outp <= num_output_pixels - 2; outp += 2) { - // Load the accumulators from acc_buffer - int32x4_t acc[4]; - for (int i = 0; i < 4; i++) { - acc[i] = vld1q_s32(acc_buffer_ptr + 4 * i); - } - // Load the inputs, add input_offset. - const uint8x8_t input_u8 = vld1_u8(input_ptr); - input_ptr += 8; - const int16x8_t input_s16 = vreinterpretq_s16_u16(vmovl_u8(input_u8)); - const int16x8_t input = vaddq_s16(input_s16, vdupq_n_s16(input_offset)); - // Duplicate the input values, 2-fold - const int16x8x2_t input_dup2 = vzipq_s16(input, input); - // Multiply-accumulate - for (int i = 0; i < 2; i++) { - acc[2 * i + 0] = vmlal_s16(acc[2 * i + 0], vget_low_s16(filter), - vget_low_s16(input_dup2.val[i])); - acc[2 * i + 1] = vmlal_s16(acc[2 * i + 1], vget_high_s16(filter), - vget_high_s16(input_dup2.val[i])); - } - // Store the accumulators back to acc_buffer - for (int i = 0; i < 4; i++) { - vst1q_s32(acc_buffer_ptr + 4 * i, acc[i]); - } - acc_buffer_ptr += 16; - } - // Handle one output pixel at a time. - for (; outp < num_output_pixels; outp++) { - // Load the accumulators from acc_buffer - int32x4_t acc[2]; - for (int i = 0; i < 2; i++) { - acc[i] = vld1q_s32(acc_buffer_ptr + 4 * i); - } - // Load the inputs, add input_offset. - uint8x8_t input_u8 = vdup_n_u8(0); - input_u8 = vset_lane_u8(input_ptr[0], input_u8, 0); - input_u8 = vset_lane_u8(input_ptr[1], input_u8, 1); - input_u8 = vset_lane_u8(input_ptr[2], input_u8, 2); - input_u8 = vset_lane_u8(input_ptr[3], input_u8, 3); - input_ptr += 4; - const int16x4_t input_s16 = - vreinterpret_s16_u16(vget_low_u16(vmovl_u8(input_u8))); - const int16x4_t input = vadd_s16(input_s16, vdup_n_s16(input_offset)); - // Duplicate the input values, 2-fold - const int16x4x2_t input_dup2 = vzip_s16(input, input); - // Multiply-accumulate - acc[0] = vmlal_s16(acc[0], vget_low_s16(filter), input_dup2.val[0]); - acc[1] = vmlal_s16(acc[1], vget_high_s16(filter), input_dup2.val[1]); - // Store the accumulators back to acc_buffer - for (int i = 0; i < 2; i++) { - vst1q_s32(acc_buffer_ptr + 4 * i, acc[i]); - } - acc_buffer_ptr += 8; - } - } -}; - -template <> -struct QuantizedDepthwiseConvKernel { - static void Run(int num_output_pixels, int input_depth, int depth_multiplier, - const uint8* input_ptr, int16 input_offset, - int input_ptr_increment, const uint8* filter_ptr, - int16 filter_offset, int32* acc_buffer_ptr) { - // Load the filters, add filter_offset. - int16x8_t filter[2]; - for (int i = 0; i < 2; i++) { - const uint8x8_t filter_u8 = vld1_u8(filter_ptr + 8 * i); - const int16x8_t filter_s16 = vreinterpretq_s16_u16(vmovl_u8(filter_u8)); - filter[i] = vaddq_s16(filter_s16, vdupq_n_s16(filter_offset)); - } - int outp = 0; - // Handle two output pixels at a time. - for (; outp <= num_output_pixels - 2; outp += 2) { - // Load the accumulators from acc_buffer. - int32x4_t acc[8]; - for (int i = 0; i < 8; i++) { - acc[i] = vld1q_s32(acc_buffer_ptr + 4 * i); - } - // Load the inputs, add input_offset. - uint8x8_t input_u8 = vdup_n_u8(0); - input_u8 = vset_lane_u8(input_ptr[0], input_u8, 0); - input_u8 = vset_lane_u8(input_ptr[1], input_u8, 1); - input_u8 = vset_lane_u8(input_ptr[2], input_u8, 2); - input_u8 = vset_lane_u8(input_ptr[3], input_u8, 3); - input_ptr += 4; - const int16x4_t input_s16 = - vreinterpret_s16_u16(vget_low_u16(vmovl_u8(input_u8))); - const int16x4_t input = vadd_s16(input_s16, vdup_n_s16(input_offset)); - // Multiply-accumulate. - acc[0] = vmlal_lane_s16(acc[0], vget_low_s16(filter[0]), input, 0); - acc[1] = vmlal_lane_s16(acc[1], vget_high_s16(filter[0]), input, 0); - acc[2] = vmlal_lane_s16(acc[2], vget_low_s16(filter[1]), input, 1); - acc[3] = vmlal_lane_s16(acc[3], vget_high_s16(filter[1]), input, 1); - acc[4] = vmlal_lane_s16(acc[4], vget_low_s16(filter[0]), input, 2); - acc[5] = vmlal_lane_s16(acc[5], vget_high_s16(filter[0]), input, 2); - acc[6] = vmlal_lane_s16(acc[6], vget_low_s16(filter[1]), input, 3); - acc[7] = vmlal_lane_s16(acc[7], vget_high_s16(filter[1]), input, 3); - // Store the accumulators back to acc_buffer. - for (int i = 0; i < 8; i++) { - vst1q_s32(acc_buffer_ptr + 4 * i, acc[i]); - } - acc_buffer_ptr += 32; - } - // Handle one output pixel at a time. - for (; outp < num_output_pixels; outp++) { - // Load the accumulators from acc_buffer. - int32x4_t acc[4]; - for (int i = 0; i < 4; i++) { - acc[i] = vld1q_s32(acc_buffer_ptr + 4 * i); - } - // Load the inputs, add input_offset. - uint8x8_t input_u8 = vdup_n_u8(0); - input_u8 = vset_lane_u8(input_ptr[0], input_u8, 0); - input_u8 = vset_lane_u8(input_ptr[1], input_u8, 1); - input_ptr += 2; - const int16x4_t input_s16 = - vreinterpret_s16_u16(vget_low_u16(vmovl_u8(input_u8))); - const int16x4_t input = vadd_s16(input_s16, vdup_n_s16(input_offset)); - - // Multiply-accumulate. - acc[0] = vmlal_lane_s16(acc[0], vget_low_s16(filter[0]), input, 0); - acc[1] = vmlal_lane_s16(acc[1], vget_high_s16(filter[0]), input, 0); - acc[2] = vmlal_lane_s16(acc[2], vget_low_s16(filter[1]), input, 1); - acc[3] = vmlal_lane_s16(acc[3], vget_high_s16(filter[1]), input, 1); - - // Store the accumulators back to acc_buffer. - for (int i = 0; i < 4; i++) { - vst1q_s32(acc_buffer_ptr + 4 * i, acc[i]); - } - acc_buffer_ptr += 16; - } - } -}; - -template <> -struct QuantizedDepthwiseConvKernel { - static void Run(int num_output_pixels, int input_depth, int depth_multiplier, - const uint8* input_ptr, int16 input_offset, - int input_ptr_increment, const uint8* filter_ptr, - int16 filter_offset, int32* acc_buffer_ptr) { - // Load the filters, add filter_offset. - uint8x8_t filter_u8 = vdup_n_u8(0); - filter_u8 = vset_lane_u8(filter_ptr[0], filter_u8, 0); - filter_u8 = vset_lane_u8(filter_ptr[1], filter_u8, 1); - filter_u8 = vset_lane_u8(filter_ptr[2], filter_u8, 2); - filter_u8 = vset_lane_u8(filter_ptr[3], filter_u8, 3); - const int16x4_t filter_s16 = - vreinterpret_s16_u16(vget_low_u16(vmovl_u8(filter_u8))); - const int16x4_t filter = vadd_s16(filter_s16, vdup_n_s16(filter_offset)); - - int outp = 0; - // Handle 4 output pixels at a time. - for (; outp <= num_output_pixels - 4; outp += 4) { - // Load the accumulators from acc_buffer - int32x4_t acc[4]; - for (int i = 0; i < 4; i++) { - acc[i] = vld1q_s32(acc_buffer_ptr + 4 * i); - } - - // Load the inputs, add input_offset. - const uint8x8_t input_u8 = vld1_u8(input_ptr); - input_ptr += 8; - const int16x8_t input_s16 = vreinterpretq_s16_u16(vmovl_u8(input_u8)); - const int16x8_t input = vaddq_s16(input_s16, vdupq_n_s16(input_offset)); - // Duplicate the input values, 2-fold - const int16x8x2_t input_dup2 = vzipq_s16(input, input); - // Multiply-accumulate - acc[0] = vmlal_s16(acc[0], filter, vget_low_s16(input_dup2.val[0])); - acc[1] = vmlal_s16(acc[1], filter, vget_high_s16(input_dup2.val[0])); - acc[2] = vmlal_s16(acc[2], filter, vget_low_s16(input_dup2.val[1])); - acc[3] = vmlal_s16(acc[3], filter, vget_high_s16(input_dup2.val[1])); - // Store the accumulators back to acc_buffer - for (int i = 0; i < 4; i++) { - vst1q_s32(acc_buffer_ptr + 4 * i, acc[i]); - } - acc_buffer_ptr += 16; - } - // Handle one output pixel at a time. - for (; outp < num_output_pixels; outp++) { - // Load the accumulators from acc_buffer - int32x4_t acc = vld1q_s32(acc_buffer_ptr); - - uint8x8_t input_u8 = vdup_n_u8(0); - input_u8 = vset_lane_u8(input_ptr[0], input_u8, 0); - input_u8 = vset_lane_u8(input_ptr[1], input_u8, 1); - input_ptr += 2; - const int16x4_t input_s16 = - vreinterpret_s16_u16(vget_low_u16(vmovl_u8(input_u8))); - const int16x4_t input = vadd_s16(input_s16, vdup_n_s16(input_offset)); - // Duplicate the input values, 2-fold - const int16x4_t input_dup2 = vzip_s16(input, input).val[0]; - // Multiply-accumulate - acc = vmlal_s16(acc, filter, input_dup2); - // Store the accumulators back to acc_buffer - vst1q_s32(acc_buffer_ptr, acc); - acc_buffer_ptr += 4; - } - } -}; - -template <> -struct QuantizedDepthwiseConvKernel { - static void Run(int num_output_pixels, int input_depth, int depth_multiplier, - const uint8* input_ptr, int16 input_offset, - int input_ptr_increment, const uint8* filter_ptr, - int16 filter_offset, int32* acc_buffer_ptr) { - // Load the filters, add filter_offset. - uint8x8_t filter_u8 = vdup_n_u8(0); - filter_u8 = vset_lane_u8(filter_ptr[0], filter_u8, 0); - filter_u8 = vset_lane_u8(filter_ptr[1], filter_u8, 1); - filter_u8 = vset_lane_u8(filter_ptr[0], filter_u8, 2); - filter_u8 = vset_lane_u8(filter_ptr[1], filter_u8, 3); - const int16x4_t filter_s16 = - vreinterpret_s16_u16(vget_low_u16(vmovl_u8(filter_u8))); - const int16x4_t filter = vadd_s16(filter_s16, vdup_n_s16(filter_offset)); - - int outp = 0; - // Handle 8 output pixels at a time. - for (; outp <= num_output_pixels - 8; outp += 8) { - // Load the accumulators from acc_buffer. - int32x4_t acc[4]; - for (int i = 0; i < 4; i++) { - acc[i] = vld1q_s32(acc_buffer_ptr + 4 * i); - } - // Load the inputs, add input_offset. - uint8x8_t input_u8[2]; - for (int i = 0; i < 2; i++) { - input_u8[i] = vld1_u8(input_ptr + 8 * i); - } - input_ptr += 16; - int16x8_t input[2]; - for (int i = 0; i < 2; i++) { - input[i] = vreinterpretq_s16_u16(vmovl_u8(input_u8[i])); - } - for (int i = 0; i < 2; i++) { - input[i] = vaddq_s16(input[i], vdupq_n_s16(input_offset)); - } - - // Multiply-accumulate. - acc[0] = vmlal_s16(acc[0], filter, vget_low_s16(input[0])); - acc[1] = vmlal_s16(acc[1], filter, vget_high_s16(input[0])); - acc[2] = vmlal_s16(acc[2], filter, vget_low_s16(input[1])); - acc[3] = vmlal_s16(acc[3], filter, vget_high_s16(input[1])); - // Store the accumulators back to acc_buffer. - for (int i = 0; i < 4; i++) { - vst1q_s32(acc_buffer_ptr + 4 * i, acc[i]); - } - acc_buffer_ptr += 16; - } - // Handle 4 output pixels at a time. - for (; outp <= num_output_pixels - 4; outp += 4) { - // Load the accumulators from acc_buffer. - int32x4_t acc[2]; - for (int i = 0; i < 2; i++) { - acc[i] = vld1q_s32(acc_buffer_ptr + 4 * i); - } - // Load the inputs, add input_offset. - const uint8x8_t input_u8 = vld1_u8(input_ptr); - input_ptr += 8; - const int16x8_t input_s16 = vreinterpretq_s16_u16(vmovl_u8(input_u8)); - const int16x8_t input = vaddq_s16(input_s16, vdupq_n_s16(input_offset)); - - // Multiply-accumulate. - acc[0] = vmlal_s16(acc[0], filter, vget_low_s16(input)); - acc[1] = vmlal_s16(acc[1], filter, vget_high_s16(input)); - // Store the accumulators back to acc_buffer. - for (int i = 0; i < 2; i++) { - vst1q_s32(acc_buffer_ptr + 4 * i, acc[i]); - } - acc_buffer_ptr += 8; - } - // Handle 2 output pixels at a time. - for (; outp <= num_output_pixels - 2; outp += 2) { - // Load the accumulators from acc_buffer. - int32x4_t acc = vld1q_s32(acc_buffer_ptr); - // Load the inputs, add input_offset. - uint8x8_t input_u8 = vdup_n_u8(0); - input_u8 = vset_lane_u8(input_ptr[0], input_u8, 0); - input_u8 = vset_lane_u8(input_ptr[1], input_u8, 1); - input_u8 = vset_lane_u8(input_ptr[2], input_u8, 2); - input_u8 = vset_lane_u8(input_ptr[3], input_u8, 3); - input_ptr += 4; - const int16x4_t input_s16 = - vreinterpret_s16_u16(vget_low_u16(vmovl_u8(input_u8))); - const int16x4_t input = vadd_s16(input_s16, vdup_n_s16(input_offset)); - - // Multiply-accumulate. - acc = vmlal_s16(acc, filter, input); - // Store the accumulators back to acc_buffer. - vst1q_s32(acc_buffer_ptr, acc); - acc_buffer_ptr += 4; - } - // Handle 1 output pixel at a time. - for (; outp < num_output_pixels; outp++) { - // Load the accumulators from acc_buffer. - int32x2_t acc = vld1_s32(acc_buffer_ptr); - // Load the inputs, add input_offset. - uint8x8_t input_u8 = vdup_n_u8(0); - input_u8 = vset_lane_u8(input_ptr[0], input_u8, 0); - input_u8 = vset_lane_u8(input_ptr[1], input_u8, 1); - input_ptr += 2; - const int16x4_t input_s16 = - vreinterpret_s16_u16(vget_low_u16(vmovl_u8(input_u8))); - const int16x4_t input = vadd_s16(input_s16, vdup_n_s16(input_offset)); - - // Multiply-accumulate. - acc = vget_low_s32(vmlal_s16(vcombine_s32(acc, acc), filter, input)); - // Store the accumulators back to acc_buffer. - vst1_s32(acc_buffer_ptr, acc); - acc_buffer_ptr += 2; - } - } -}; - -template <> -struct QuantizedDepthwiseConvKernel { - static void Run(int num_output_pixels, int input_depth, int depth_multiplier, - const uint8* input_ptr, int16 input_offset, - int input_ptr_increment, const uint8* filter_ptr, - int16 filter_offset, int32* acc_buffer_ptr) { - // Load the filters, add filter_offset. - uint8x8_t filter_u8 = vdup_n_u8(0); - filter_u8 = vset_lane_u8(filter_ptr[0], filter_u8, 0); - filter_u8 = vset_lane_u8(filter_ptr[1], filter_u8, 1); - filter_u8 = vset_lane_u8(filter_ptr[0], filter_u8, 2); - filter_u8 = vset_lane_u8(filter_ptr[1], filter_u8, 3); - const int16x4_t filter_s16 = - vreinterpret_s16_u16(vget_low_u16(vmovl_u8(filter_u8))); - const int16x4_t filter = vadd_s16(filter_s16, vdup_n_s16(filter_offset)); - - int outp = 0; - // Handle 8 output pixels at a time. - for (; outp <= num_output_pixels - 8; outp += 8) { - // Load the accumulators from acc_buffer - int32x4_t acc[4]; - for (int i = 0; i < 4; i++) { - acc[i] = vld1q_s32(acc_buffer_ptr + 4 * i); - } - - // Load the inputs, add input_offset. - const uint8x8_t input_u8 = vld1_u8(input_ptr); - input_ptr += 8; - const int16x8_t input_s16 = vreinterpretq_s16_u16(vmovl_u8(input_u8)); - const int16x8_t input = vaddq_s16(input_s16, vdupq_n_s16(input_offset)); - // Duplicate the input values, 2-fold - const int16x8x2_t input_dup2 = vzipq_s16(input, input); - // Multiply-accumulate - acc[0] = vmlal_s16(acc[0], filter, vget_low_s16(input_dup2.val[0])); - acc[1] = vmlal_s16(acc[1], filter, vget_high_s16(input_dup2.val[0])); - acc[2] = vmlal_s16(acc[2], filter, vget_low_s16(input_dup2.val[1])); - acc[3] = vmlal_s16(acc[3], filter, vget_high_s16(input_dup2.val[1])); - // Store the accumulators back to acc_buffer - for (int i = 0; i < 4; i++) { - vst1q_s32(acc_buffer_ptr + 4 * i, acc[i]); - } - acc_buffer_ptr += 16; - } - // Handle one output pixel at a time. - for (; outp < num_output_pixels; outp++) { - // Load the accumulators from acc_buffer - int32x2_t acc = vld1_s32(acc_buffer_ptr); - - // Load the inputs, add input_offset. - const uint32 input = *input_ptr++ + input_offset; - - // Multiply-accumulate - acc = vget_low_s32(vmlal_n_s16(vcombine_s32(acc, acc), filter, input)); - // Store the accumulators back to acc_buffer - vst1_s32(acc_buffer_ptr, acc); - acc_buffer_ptr += 2; - } - } -}; - -template <> -struct QuantizedDepthwiseConvKernel { - static void Run(int num_output_pixels, int input_depth, int depth_multiplier, - const uint8* input_ptr, int16 input_offset, - int input_ptr_increment, const uint8* filter_ptr, - int16 filter_offset, int32* acc_buffer_ptr) { - // Load the filters, add filter_offset. - uint8x8_t filter_u8 = vdup_n_u8(0); - filter_u8 = vset_lane_u8(filter_ptr[0], filter_u8, 0); - filter_u8 = vset_lane_u8(filter_ptr[1], filter_u8, 1); - filter_u8 = vset_lane_u8(filter_ptr[2], filter_u8, 2); - filter_u8 = vset_lane_u8(filter_ptr[3], filter_u8, 3); - const int16x4_t filter_s16 = - vreinterpret_s16_u16(vget_low_u16(vmovl_u8(filter_u8))); - const int16x4_t filter = vadd_s16(filter_s16, vdup_n_s16(filter_offset)); - - int outp = 0; - // Handle 8 output pixels at a time. - for (; outp <= num_output_pixels - 8; outp += 8) { - // Load the accumulators from acc_buffer - int32x4_t acc[8]; - for (int i = 0; i < 8; i++) { - acc[i] = vld1q_s32(acc_buffer_ptr + 4 * i); - } - - // Load the inputs, add input_offset. - uint8x8_t input_u8 = vld1_u8(input_ptr); - input_ptr += 8; - const int16x8_t input_s16 = vreinterpretq_s16_u16(vmovl_u8(input_u8)); - const int16x8_t input = vaddq_s16(input_s16, vdupq_n_s16(input_offset)); - - // Multiply-accumulate - acc[0] = vmlal_lane_s16(acc[0], filter, vget_low_s16(input), 0); - acc[1] = vmlal_lane_s16(acc[1], filter, vget_low_s16(input), 1); - acc[2] = vmlal_lane_s16(acc[2], filter, vget_low_s16(input), 2); - acc[3] = vmlal_lane_s16(acc[3], filter, vget_low_s16(input), 3); - acc[4] = vmlal_lane_s16(acc[4], filter, vget_high_s16(input), 0); - acc[5] = vmlal_lane_s16(acc[5], filter, vget_high_s16(input), 1); - acc[6] = vmlal_lane_s16(acc[6], filter, vget_high_s16(input), 2); - acc[7] = vmlal_lane_s16(acc[7], filter, vget_high_s16(input), 3); - - // Store the accumulators back to acc_buffer - for (int i = 0; i < 8; i++) { - vst1q_s32(acc_buffer_ptr + 4 * i, acc[i]); - } - acc_buffer_ptr += 32; - } - // Handle 4 output pixels at a time. - for (; outp <= num_output_pixels - 4; outp += 4) { - // Load the accumulators from acc_buffer - int32x4_t acc[4]; - for (int i = 0; i < 4; i++) { - acc[i] = vld1q_s32(acc_buffer_ptr + 4 * i); - } - - // Load the inputs, add input_offset. - uint8x8_t input_u8 = vdup_n_u8(0); - input_u8 = vset_lane_u8(input_ptr[0], input_u8, 0); - input_u8 = vset_lane_u8(input_ptr[1], input_u8, 1); - input_u8 = vset_lane_u8(input_ptr[2], input_u8, 2); - input_u8 = vset_lane_u8(input_ptr[3], input_u8, 3); - input_ptr += 4; - const int16x4_t input_s16 = - vreinterpret_s16_u16(vget_low_u16(vmovl_u8(input_u8))); - const int16x4_t input = vadd_s16(input_s16, vdup_n_s16(input_offset)); - - // Multiply-accumulate - acc[0] = vmlal_lane_s16(acc[0], filter, input, 0); - acc[1] = vmlal_lane_s16(acc[1], filter, input, 1); - acc[2] = vmlal_lane_s16(acc[2], filter, input, 2); - acc[3] = vmlal_lane_s16(acc[3], filter, input, 3); - - // Store the accumulators back to acc_buffer - for (int i = 0; i < 4; i++) { - vst1q_s32(acc_buffer_ptr + 4 * i, acc[i]); - } - acc_buffer_ptr += 16; - } - // Handle one output pixel at a time. - for (; outp < num_output_pixels; outp++) { - // Load the accumulators from acc_buffer - int32x4_t acc = vld1q_s32(acc_buffer_ptr); - - // Load the inputs, add input_offset. - const uint32 input = *input_ptr++ + input_offset; - - // Multiply-accumulate - acc = vmlal_n_s16(acc, filter, input); - // Store the accumulators back to acc_buffer - vst1q_s32(acc_buffer_ptr, acc); - acc_buffer_ptr += 4; - } - } -}; - -template <> -struct QuantizedDepthwiseConvKernel { - static void Run(int num_output_pixels, int input_depth, int depth_multiplier, - const uint8* input_ptr, int16 input_offset, - int input_ptr_increment, const uint8* filter_ptr, - int16 filter_offset, int32* acc_buffer_ptr) { - // Load the filters, add filter_offset. - uint8x8_t filter_u8 = vdup_n_u8(0); - filter_u8 = vset_lane_u8(filter_ptr[0], filter_u8, 0); - filter_u8 = vset_lane_u8(filter_ptr[1], filter_u8, 1); - filter_u8 = vset_lane_u8(filter_ptr[2], filter_u8, 2); - filter_u8 = vset_lane_u8(filter_ptr[3], filter_u8, 3); - const int16x4_t filter_s16 = - vreinterpret_s16_u16(vget_low_u16(vmovl_u8(filter_u8))); - const int16x4_t filter = vadd_s16(filter_s16, vdup_n_s16(filter_offset)); - - int outp = 0; - // Handle 4 output pixels at a time. - for (; outp <= num_output_pixels - 4; outp += 4) { - // Load the accumulators from acc_buffer - int32x4_t acc[4]; - for (int i = 0; i < 4; i++) { - acc[i] = vld1q_s32(acc_buffer_ptr + 4 * i); - } - // Load the inputs, add input_offset. - int16x8_t input[2]; - for (int i = 0; i < 2; i++) { - const uint8x8_t input_u8 = vld1_u8(input_ptr + 8 * i); - const int16x8_t input_s16 = vreinterpretq_s16_u16(vmovl_u8(input_u8)); - input[i] = vaddq_s16(input_s16, vdupq_n_s16(input_offset)); - } - input_ptr += 16; - // Multiply-accumulate - for (int i = 0; i < 2; i++) { - acc[2 * i + 0] = - vmlal_s16(acc[2 * i + 0], filter, vget_low_s16(input[i])); - acc[2 * i + 1] = - vmlal_s16(acc[2 * i + 1], filter, vget_high_s16(input[i])); - } - // Store the accumulators back to acc_buffer - for (int i = 0; i < 4; i++) { - vst1q_s32(acc_buffer_ptr + 4 * i, acc[i]); - } - acc_buffer_ptr += 16; - } - // Handle one output pixel at a time. - for (; outp < num_output_pixels; outp++) { - // Load the accumulators from acc_buffer - int32x4_t acc; - acc = vld1q_s32(acc_buffer_ptr); - - // Load the inputs, add input_offset. - uint8x8_t input_u8 = vdup_n_u8(0); - input_u8 = vset_lane_u8(input_ptr[0], input_u8, 0); - input_u8 = vset_lane_u8(input_ptr[1], input_u8, 1); - input_u8 = vset_lane_u8(input_ptr[2], input_u8, 2); - input_u8 = vset_lane_u8(input_ptr[3], input_u8, 3); - input_ptr += 4; - const int16x4_t input_s16 = - vreinterpret_s16_u16(vget_low_u16(vmovl_u8(input_u8))); - const int16x4_t input = vadd_s16(input_s16, vdup_n_s16(input_offset)); - // Multiply-accumulate - acc = vmlal_s16(acc, filter, input); - // Store the accumulators back to acc_buffer - vst1q_s32(acc_buffer_ptr, acc); - acc_buffer_ptr += 4; - } - } -}; - -template <> -struct QuantizedDepthwiseConvKernel { - static void Run(int num_output_pixels, int input_depth, int depth_multiplier, - const uint8* input_ptr, int16 input_offset, - int input_ptr_increment, const uint8* filter_ptr, - int16 filter_offset, int32* acc_buffer_ptr) { - // Load the filters, add filter_offset. - int16x8_t filter[2]; - for (int i = 0; i < 2; i++) { - const uint8x8_t filter_u8 = vld1_u8(filter_ptr + 8 * i); - const int16x8_t filter_s16 = vreinterpretq_s16_u16(vmovl_u8(filter_u8)); - filter[i] = vaddq_s16(filter_s16, vdupq_n_s16(filter_offset)); - } - - int outp = 0; - // Handle 2 output pixels at a time. - for (; outp <= num_output_pixels - 2; outp += 2) { - // Load the accumulators from acc_buffer - int32x4_t acc[8]; - for (int i = 0; i < 8; i++) { - acc[i] = vld1q_s32(acc_buffer_ptr + 4 * i); - } - - // Load the inputs, add input_offset. - uint8x8_t input_u8 = vld1_u8(input_ptr); - input_ptr += 8; - const int16x8_t input_s16 = vreinterpretq_s16_u16(vmovl_u8(input_u8)); - const int16x8_t input = vaddq_s16(input_s16, vdupq_n_s16(input_offset)); - - // Multiply-accumulate - acc[0] = vmlal_lane_s16(acc[0], vget_low_s16(filter[0]), - vget_low_s16(input), 0); - acc[1] = vmlal_lane_s16(acc[1], vget_high_s16(filter[0]), - vget_low_s16(input), 1); - acc[2] = vmlal_lane_s16(acc[2], vget_low_s16(filter[1]), - vget_low_s16(input), 2); - acc[3] = vmlal_lane_s16(acc[3], vget_high_s16(filter[1]), - vget_low_s16(input), 3); - acc[4] = vmlal_lane_s16(acc[4], vget_low_s16(filter[0]), - vget_high_s16(input), 0); - acc[5] = vmlal_lane_s16(acc[5], vget_high_s16(filter[0]), - vget_high_s16(input), 1); - acc[6] = vmlal_lane_s16(acc[6], vget_low_s16(filter[1]), - vget_high_s16(input), 2); - acc[7] = vmlal_lane_s16(acc[7], vget_high_s16(filter[1]), - vget_high_s16(input), 3); - // Store the accumulators back to acc_buffer - for (int i = 0; i < 8; i++) { - vst1q_s32(acc_buffer_ptr + 4 * i, acc[i]); - } - acc_buffer_ptr += 32; - } - // Handle one output pixel at a time. - for (; outp < num_output_pixels; outp++) { - // Load the accumulators from acc_buffer - int32x4_t acc[4]; - for (int i = 0; i < 4; i++) { - acc[i] = vld1q_s32(acc_buffer_ptr + 4 * i); - } - - // Load the inputs, add input_offset. - uint8x8_t input_u8 = vdup_n_u8(0); - input_u8 = vset_lane_u8(input_ptr[0], input_u8, 0); - input_u8 = vset_lane_u8(input_ptr[1], input_u8, 1); - input_u8 = vset_lane_u8(input_ptr[2], input_u8, 2); - input_u8 = vset_lane_u8(input_ptr[3], input_u8, 3); - input_ptr += 4; - const int16x4_t input_s16 = - vreinterpret_s16_u16(vget_low_u16(vmovl_u8(input_u8))); - const int16x4_t input = vadd_s16(input_s16, vdup_n_s16(input_offset)); - - // Multiply-accumulate - acc[0] = vmlal_lane_s16(acc[0], vget_low_s16(filter[0]), input, 0); - acc[1] = vmlal_lane_s16(acc[1], vget_high_s16(filter[0]), input, 1); - acc[2] = vmlal_lane_s16(acc[2], vget_low_s16(filter[1]), input, 2); - acc[3] = vmlal_lane_s16(acc[3], vget_high_s16(filter[1]), input, 3); - // Store the accumulators back to acc_buffer - for (int i = 0; i < 4; i++) { - vst1q_s32(acc_buffer_ptr + 4 * i, acc[i]); - } - acc_buffer_ptr += 16; - } - } -}; - -template <> -struct QuantizedDepthwiseConvKernel { - static void Run(int num_output_pixels, int input_depth, int depth_multiplier, - const uint8* input_ptr, int16 input_offset, - int input_ptr_increment, const uint8* filter_ptr, - int16 filter_offset, int32* acc_buffer_ptr) { - // We will have to duplicate bytes in a NEON register, 3-fold. - // We will do that by register-level table-look-up using VTBL instructions. - // Here we prepare the registers containing the table-lookup indices. - static const uint8 dup3_indices_array[3][8] = {{0, 0, 0, 1, 1, 1, 2, 2}, - {2, 3, 3, 3, 4, 4, 4, 5}, - {5, 5, 6, 6, 6, 7, 7, 7}}; - uint8x8_t dup3_indices[3]; - for (int i = 0; i < 3; i++) { - dup3_indices[i] = vld1_u8(dup3_indices_array[i]); - } - - // Handle one output pixel at a time. - for (int outp = 0; outp < num_output_pixels; outp++) { - const uint8* local_filter_ptr = filter_ptr; - const uint8* local_input_ptr = input_ptr; - int ic = 0; - // Handle 8 input channels at a time. - for (; ic <= input_depth - 8; ic += 8) { - // Load the filters, add filter_offset. - int16x8_t filter[3]; - uint8x8x3_t filter_u8; - filter_u8.val[0] = vld1_u8(local_filter_ptr); - filter_u8.val[1] = vld1_u8(local_filter_ptr + 8); - filter_u8.val[2] = vld1_u8(local_filter_ptr + 16); - local_filter_ptr += 24; - for (int i = 0; i < 3; i++) { - const int16x8_t filter_s16 = - vreinterpretq_s16_u16(vmovl_u8(filter_u8.val[i])); - filter[i] = vaddq_s16(filter_s16, vdupq_n_s16(filter_offset)); - } - // Load the inputs, duplicate 3-fold, add input_offset. - const uint8x8_t input_u8 = vld1_u8(local_input_ptr); - local_input_ptr += 8; - - uint8x8_t input_u8_dup3[3]; - for (int i = 0; i < 3; i++) { - input_u8_dup3[i] = vtbl1_u8(input_u8, dup3_indices[i]); - } - int16x8_t input_dup3[3]; - for (int i = 0; i < 3; i++) { - const int16x8_t input_s16_dup3 = - vreinterpretq_s16_u16(vmovl_u8(input_u8_dup3[i])); - input_dup3[i] = vaddq_s16(input_s16_dup3, vdupq_n_s16(input_offset)); - } - // Load the accumulators from acc_buffer - int32x4x3_t acc[2]; - for (int i = 0; i < 2; i++) { - acc[i].val[0] = vld1q_s32(acc_buffer_ptr + 4 * i); - acc[i].val[1] = vld1q_s32(acc_buffer_ptr + 4 * i + 8); - acc[i].val[2] = vld1q_s32(acc_buffer_ptr + 4 * i + 16); - } - // Multiply-accumulate - for (int j = 0; j < 3; j++) { - acc[0].val[j] = vmlal_s16(acc[0].val[j], vget_low_s16(input_dup3[j]), - vget_low_s16(filter[j])); - acc[1].val[j] = vmlal_s16(acc[1].val[j], vget_high_s16(input_dup3[j]), - vget_high_s16(filter[j])); - } - // Store the accumulators back to acc_buffer - for (int i = 0; i < 2; i++) { - vst1q_s32(acc_buffer_ptr + 4 * i, acc[i].val[0]); - vst1q_s32(acc_buffer_ptr + 4 * i + 8, acc[i].val[1]); - vst1q_s32(acc_buffer_ptr + 4 * i + 16, acc[i].val[2]); - } - acc_buffer_ptr += 24; - } - // Handle one input channel at a time. - for (; ic < input_depth; ic++) { - const int16 input_val = *local_input_ptr++ + input_offset; - for (int i = 0; i < 3; i++) { - const int16 filter_val = local_filter_ptr[i] + filter_offset; - *acc_buffer_ptr++ += static_cast(filter_val) * input_val; - } - local_filter_ptr += 3; - } - input_ptr += input_ptr_increment; - } - } -}; - -template <> -struct QuantizedDepthwiseConvKernel { - static void Run(int num_output_pixels, int input_depth, int depth_multiplier, - const uint8* input_ptr, int16 input_offset, - int input_ptr_increment, const uint8* filter_ptr, - int16 filter_offset, int32* acc_buffer_ptr) { - // Handle one output pixel at a time. - for (int outp = 0; outp < num_output_pixels; outp++) { - const uint8* local_filter_ptr = filter_ptr; - const uint8* local_input_ptr = input_ptr; - int ic = 0; - // Handle 8 input channels at a time. - for (; ic <= input_depth - 8; ic += 8) { - // Load the filters, add filter_offset. - int16x8_t filter[2]; - uint8x8x2_t filter_u8; - filter_u8.val[0] = vld1_u8(local_filter_ptr); - filter_u8.val[1] = vld1_u8(local_filter_ptr + 8); - local_filter_ptr += 16; - for (int i = 0; i < 2; i++) { - const int16x8_t filter_s16 = - vreinterpretq_s16_u16(vmovl_u8(filter_u8.val[i])); - filter[i] = vaddq_s16(filter_s16, vdupq_n_s16(filter_offset)); - } - // Load the inputs, add input_offset, duplicate 2-fold. - const uint8x8_t input_u8 = vld1_u8(local_input_ptr); - local_input_ptr += 8; - const int16x8_t input_s16 = vreinterpretq_s16_u16(vmovl_u8(input_u8)); - const int16x8_t input = vaddq_s16(input_s16, vdupq_n_s16(input_offset)); - const int16x8x2_t input_dup2 = vzipq_s16(input, input); - // Load the accumulators from acc_buffer. - int32x4x2_t acc[2]; - for (int i = 0; i < 2; i++) { - acc[i].val[0] = vld1q_s32(acc_buffer_ptr + 4 * i); - acc[i].val[1] = vld1q_s32(acc_buffer_ptr + 4 * i + 8); - } - // Multiply-accumulate. - for (int j = 0; j < 2; j++) { - acc[0].val[j] = vmlal_s16(acc[0].val[j], vget_low_s16(filter[j]), - vget_low_s16(input_dup2.val[j])); - acc[1].val[j] = vmlal_s16(acc[1].val[j], vget_high_s16(filter[j]), - vget_high_s16(input_dup2.val[j])); - } - // Store the accumulators back to acc_buffer. - for (int i = 0; i < 2; i++) { - vst1q_s32(acc_buffer_ptr + 4 * i, acc[i].val[0]); - vst1q_s32(acc_buffer_ptr + 4 * i + 8, acc[i].val[1]); - } - acc_buffer_ptr += 16; - } - // Handle one input channel at a time. - for (; ic < input_depth; ic++) { - // Load the inputs. - const int16 input_val = *local_input_ptr++ + input_offset; - for (int i = 0; i < 2; i++) { - const int16 filter_val = local_filter_ptr[i] + filter_offset; - *acc_buffer_ptr++ += static_cast(filter_val) * input_val; - } - local_filter_ptr += 2; - } - input_ptr += input_ptr_increment; - } - } -}; - -template <> -struct QuantizedDepthwiseConvKernel { - static void Run(int num_output_pixels, int input_depth, int depth_multiplier, - const uint8* input_ptr, int16 input_offset, - int input_ptr_increment, const uint8* filter_ptr, - int16 filter_offset, int32* acc_buffer_ptr) { - // Handle one output pixel at a time. - for (int outp = 0; outp < num_output_pixels; outp++) { - const uint8* local_filter_ptr = filter_ptr; - const uint8* local_input_ptr = input_ptr; - int ic = 0; - // Handle 16 input channels at a time. - for (; ic <= input_depth - 16; ic += 16) { - // Load the filters, add filter_offset. - uint8x8_t filter_u8_0 = vld1_u8(local_filter_ptr + 8 * 0); - uint8x8_t filter_u8_1 = vld1_u8(local_filter_ptr + 8 * 1); - local_filter_ptr += 16; - int16x8_t filter_0 = vreinterpretq_s16_u16(vmovl_u8(filter_u8_0)); - int16x8_t filter_1 = vreinterpretq_s16_u16(vmovl_u8(filter_u8_1)); - filter_0 = vaddq_s16(filter_0, vdupq_n_s16(filter_offset)); - filter_1 = vaddq_s16(filter_1, vdupq_n_s16(filter_offset)); - // Load the inputs, add input_offset. - uint8x8_t input_u8_0 = vld1_u8(local_input_ptr + 8 * 0); - uint8x8_t input_u8_1 = vld1_u8(local_input_ptr + 8 * 1); - local_input_ptr += 16; - int16x8_t input_0 = vreinterpretq_s16_u16(vmovl_u8(input_u8_0)); - int16x8_t input_1 = vreinterpretq_s16_u16(vmovl_u8(input_u8_1)); - input_0 = vaddq_s16(input_0, vdupq_n_s16(input_offset)); - input_1 = vaddq_s16(input_1, vdupq_n_s16(input_offset)); - // Load the accumulators from acc_buffer - int32x4_t acc_0 = vld1q_s32(acc_buffer_ptr + 4 * 0); - int32x4_t acc_1 = vld1q_s32(acc_buffer_ptr + 4 * 1); - int32x4_t acc_2 = vld1q_s32(acc_buffer_ptr + 4 * 2); - int32x4_t acc_3 = vld1q_s32(acc_buffer_ptr + 4 * 3); - acc_0 = vmlal_s16(acc_0, vget_low_s16(input_0), vget_low_s16(filter_0)); - acc_1 = - vmlal_s16(acc_1, vget_high_s16(input_0), vget_high_s16(filter_0)); - acc_2 = vmlal_s16(acc_2, vget_low_s16(input_1), vget_low_s16(filter_1)); - acc_3 = - vmlal_s16(acc_3, vget_high_s16(input_1), vget_high_s16(filter_1)); - // Store the accumulators back to acc_buffer - vst1q_s32(acc_buffer_ptr + 4 * 0, acc_0); - vst1q_s32(acc_buffer_ptr + 4 * 1, acc_1); - vst1q_s32(acc_buffer_ptr + 4 * 2, acc_2); - vst1q_s32(acc_buffer_ptr + 4 * 3, acc_3); - acc_buffer_ptr += 16; - } - // Handle 8 input channels at a time. - for (; ic <= input_depth - 8; ic += 8) { - // Load the filters, add filter_offset. - const uint8x8_t filter_u8 = vld1_u8(local_filter_ptr); - local_filter_ptr += 8; - const int16x8_t filter_s16 = vreinterpretq_s16_u16(vmovl_u8(filter_u8)); - const int16x8_t filter = - vaddq_s16(filter_s16, vdupq_n_s16(filter_offset)); - // Load the inputs, add input_offset. - const uint8x8_t input_u8 = vld1_u8(local_input_ptr); - local_input_ptr += 8; - const int16x8_t input_s16 = vreinterpretq_s16_u16(vmovl_u8(input_u8)); - const int16x8_t input = vaddq_s16(input_s16, vdupq_n_s16(input_offset)); - // Load the accumulators from acc_buffer - int32x4_t acc[2]; - for (int i = 0; i < 2; i++) { - acc[i] = vld1q_s32(acc_buffer_ptr + 4 * i); - } - // Multiply-accumulate - acc[0] = vmlal_s16(acc[0], vget_low_s16(input), vget_low_s16(filter)); - acc[1] = vmlal_s16(acc[1], vget_high_s16(input), vget_high_s16(filter)); - // Store the accumulators back to acc_buffer - for (int i = 0; i < 2; i++) { - vst1q_s32(acc_buffer_ptr + 4 * i, acc[i]); - } - acc_buffer_ptr += 8; - } - // Handle one input channel at a time. - for (; ic < input_depth; ic++) { - const int16 input_val = *local_input_ptr++ + input_offset; - const int16 filter_val = *local_filter_ptr++ + filter_offset; - *acc_buffer_ptr++ += static_cast(filter_val) * input_val; - } - input_ptr += input_ptr_increment; - } - } -}; - -template <> -struct QuantizedDepthwiseConvKernel { - static void Run(int num_output_pixels, int input_depth, int depth_multiplier, - const uint8* input_ptr, int16 input_offset, - int input_ptr_increment, const uint8* filter_ptr, - int16 filter_offset, int32* acc_buffer_ptr) { - // Load the filters, add filter_offset. - uint8x8_t filter_u8[2]; - for (int i = 0; i < 2; i++) { - filter_u8[i] = vld1_u8(filter_ptr + 8 * i); - } - int16x8_t filter[2]; - for (int i = 0; i < 2; i++) { - filter[i] = vreinterpretq_s16_u16(vmovl_u8(filter_u8[i])); - } - for (int i = 0; i < 2; i++) { - filter[i] = vaddq_s16(filter[i], vdupq_n_s16(filter_offset)); - } - // Handle one output pixel at a time. - for (int outp = 0; outp < num_output_pixels; outp++) { - // Load the inputs, add input_offset. - uint8x8_t input_u8[2]; - for (int i = 0; i < 2; i++) { - input_u8[i] = vld1_u8(input_ptr + 8 * i); - } - input_ptr += input_ptr_increment; - int16x8_t input[2]; - for (int i = 0; i < 2; i++) { - input[i] = vreinterpretq_s16_u16(vmovl_u8(input_u8[i])); - } - for (int i = 0; i < 2; i++) { - input[i] = vaddq_s16(input[i], vdupq_n_s16(input_offset)); - } - // Load the accumulators from acc_buffer - int32x4_t acc[4]; - for (int i = 0; i < 4; i++) { - acc[i] = vld1q_s32(acc_buffer_ptr + 4 * i); - } - // Multiply-accumulate - for (int i = 0; i < 2; i++) { - acc[2 * i + 0] = vmlal_s16(acc[2 * i + 0], vget_low_s16(input[i]), - vget_low_s16(filter[i])); - acc[2 * i + 1] = vmlal_s16(acc[2 * i + 1], vget_high_s16(input[i]), - vget_high_s16(filter[i])); - } - // Store the accumulators back to acc_buffer - for (int i = 0; i < 4; i++) { - vst1q_s32(acc_buffer_ptr + 4 * i, acc[i]); - } - acc_buffer_ptr += 16; - } - } -}; - -template <> -struct QuantizedDepthwiseConvKernel { - static void Run(int num_output_pixels, int input_depth, int depth_multiplier, - const uint8* input_ptr, int16 input_offset, - int input_ptr_increment, const uint8* filter_ptr, - int16 filter_offset, int32* acc_buffer_ptr) { - // Load the filters, add filter_offset. - const uint8x8_t filter_u8 = vld1_u8(filter_ptr); - const int16x8_t filter_s16 = vreinterpretq_s16_u16(vmovl_u8(filter_u8)); - const int16x8_t filter = vaddq_s16(filter_s16, vdupq_n_s16(filter_offset)); - // Handle one output pixel at a time. - for (int outp = 0; outp < num_output_pixels; outp++) { - // Load the inputs, add input_offset. - const uint8x8_t input_u8 = vld1_u8(input_ptr); - const int16x8_t input_s16 = vreinterpretq_s16_u16(vmovl_u8(input_u8)); - const int16x8_t input = vaddq_s16(input_s16, vdupq_n_s16(input_offset)); - // Load the accumulators from acc_buffer - int32x4_t acc[2]; - for (int i = 0; i < 2; i++) { - acc[i] = vld1q_s32(acc_buffer_ptr + 4 * i); - } - // Multiply-accumulate - acc[0] = vmlal_s16(acc[0], vget_low_s16(input), vget_low_s16(filter)); - acc[1] = vmlal_s16(acc[1], vget_high_s16(input), vget_high_s16(filter)); - // Store the accumulators back to acc_buffer - for (int i = 0; i < 2; i++) { - vst1q_s32(acc_buffer_ptr + 4 * i, acc[i]); - } - acc_buffer_ptr += 8; - input_ptr += input_ptr_increment; - } - } -}; - -template <> -struct QuantizedDepthwiseConvKernel { - static void Run(int num_output_pixels, int input_depth, int depth_multiplier, - const uint8* input_ptr, int16 input_offset, - int input_ptr_increment, const uint8* filter_ptr, - int16 filter_offset, int32* acc_buffer_ptr) { - // Load the filters, add filter_offset. - uint8x8_t filter_u8[2]; - for (int i = 0; i < 2; i++) { - filter_u8[i] = vld1_u8(filter_ptr + 8 * i); - } - int16x8_t filter[2]; - for (int i = 0; i < 2; i++) { - filter[i] = vreinterpretq_s16_u16(vmovl_u8(filter_u8[i])); - } - for (int i = 0; i < 2; i++) { - filter[i] = vaddq_s16(filter[i], vdupq_n_s16(filter_offset)); - } - // Handle one output pixel at a time. - for (int outp = 0; outp < num_output_pixels; outp++) { - uint8 input_u8 = *input_ptr; - input_ptr += input_ptr_increment; - int16 input = static_cast(input_u8 + input_offset); - // Load the accumulators from acc_buffer - int32x4_t acc[4]; - for (int i = 0; i < 4; i++) { - acc[i] = vld1q_s32(acc_buffer_ptr + 4 * i); - } - // Multiply-accumulate - for (int i = 0; i < 2; i++) { - acc[2 * i + 0] = - vmlal_n_s16(acc[2 * i + 0], vget_low_s16(filter[i]), input); - acc[2 * i + 1] = - vmlal_n_s16(acc[2 * i + 1], vget_high_s16(filter[i]), input); - } - // Store the accumulators back to acc_buffer - for (int i = 0; i < 4; i++) { - vst1q_s32(acc_buffer_ptr + 4 * i, acc[i]); - } - acc_buffer_ptr += 16; - } - } -}; - -template <> -struct QuantizedDepthwiseConvKernel { - static void Run(int num_output_pixels, int input_depth, int depth_multiplier, - const uint8* input_ptr, int16 input_offset, - int input_ptr_increment, const uint8* filter_ptr, - int16 filter_offset, int32* acc_buffer_ptr) { - // Load the filters, add filter_offset. - uint8x8_t filter_u8_0 = vld1_u8(filter_ptr + 8 * 0); - uint8x8_t filter_u8_1 = vld1_u8(filter_ptr + 8 * 1); - uint8x8_t filter_u8_2 = vld1_u8(filter_ptr + 8 * 2); - uint8x8_t filter_u8_3 = vld1_u8(filter_ptr + 8 * 3); - int16x8_t filter_0 = vreinterpretq_s16_u16(vmovl_u8(filter_u8_0)); - int16x8_t filter_1 = vreinterpretq_s16_u16(vmovl_u8(filter_u8_1)); - int16x8_t filter_2 = vreinterpretq_s16_u16(vmovl_u8(filter_u8_2)); - int16x8_t filter_3 = vreinterpretq_s16_u16(vmovl_u8(filter_u8_3)); - filter_0 = vaddq_s16(filter_0, vdupq_n_s16(filter_offset)); - filter_1 = vaddq_s16(filter_1, vdupq_n_s16(filter_offset)); - filter_2 = vaddq_s16(filter_2, vdupq_n_s16(filter_offset)); - filter_3 = vaddq_s16(filter_3, vdupq_n_s16(filter_offset)); - // Handle one output pixel at a time. - for (int outp = 0; outp < num_output_pixels; outp++) { - uint8 input_u8 = *input_ptr; - input_ptr += input_ptr_increment; - int16 input = static_cast(input_u8 + input_offset); - // Load the accumulators from acc_buffer - int32x4_t acc_0 = vld1q_s32(acc_buffer_ptr + 4 * 0); - int32x4_t acc_1 = vld1q_s32(acc_buffer_ptr + 4 * 1); - int32x4_t acc_2 = vld1q_s32(acc_buffer_ptr + 4 * 2); - int32x4_t acc_3 = vld1q_s32(acc_buffer_ptr + 4 * 3); - int32x4_t acc_4 = vld1q_s32(acc_buffer_ptr + 4 * 4); - int32x4_t acc_5 = vld1q_s32(acc_buffer_ptr + 4 * 5); - int32x4_t acc_6 = vld1q_s32(acc_buffer_ptr + 4 * 6); - int32x4_t acc_7 = vld1q_s32(acc_buffer_ptr + 4 * 7); - // Multiply-accumulate - acc_0 = vmlal_n_s16(acc_0, vget_low_s16(filter_0), input); - acc_1 = vmlal_n_s16(acc_1, vget_high_s16(filter_0), input); - acc_2 = vmlal_n_s16(acc_2, vget_low_s16(filter_1), input); - acc_3 = vmlal_n_s16(acc_3, vget_high_s16(filter_1), input); - acc_4 = vmlal_n_s16(acc_4, vget_low_s16(filter_2), input); - acc_5 = vmlal_n_s16(acc_5, vget_high_s16(filter_2), input); - acc_6 = vmlal_n_s16(acc_6, vget_low_s16(filter_3), input); - acc_7 = vmlal_n_s16(acc_7, vget_high_s16(filter_3), input); - // Store the accumulators back to acc_buffer - vst1q_s32(acc_buffer_ptr + 4 * 0, acc_0); - vst1q_s32(acc_buffer_ptr + 4 * 1, acc_1); - vst1q_s32(acc_buffer_ptr + 4 * 2, acc_2); - vst1q_s32(acc_buffer_ptr + 4 * 3, acc_3); - vst1q_s32(acc_buffer_ptr + 4 * 4, acc_4); - vst1q_s32(acc_buffer_ptr + 4 * 5, acc_5); - vst1q_s32(acc_buffer_ptr + 4 * 6, acc_6); - vst1q_s32(acc_buffer_ptr + 4 * 7, acc_7); - acc_buffer_ptr += 32; - } - } -}; - -template <> -struct QuantizedDepthwiseConvKernel { - static void Run(int num_output_pixels, int input_depth, int depth_multiplier, - const uint8* input_ptr, int16 input_offset, - int input_ptr_increment, const uint8* filter_ptr, - int16 filter_offset, int32* acc_buffer_ptr) { - // Load the filters, add filter_offset. - // NEON wants to load 8 bytes at a time, but 20 is not divisible by 8. - // We load the first 16 bytes into filter_u8_{0,1} as usual. - // Then we load the 8 last bytes into filter_u8_x (x for 'extra'). - // This is redundant: the first 4 bytes of filter_u8_x are the same - // as the last 4 bytes of filter_u8_x. - uint8x8_t filter_u8_0 = vld1_u8(filter_ptr + 8 * 0); - uint8x8_t filter_u8_1 = vld1_u8(filter_ptr + 8 * 1); - uint8x8_t filter_u8_x = vld1_u8(filter_ptr + 8 * 1 + 4); - int16x8_t filter_0 = vreinterpretq_s16_u16(vmovl_u8(filter_u8_0)); - int16x8_t filter_1 = vreinterpretq_s16_u16(vmovl_u8(filter_u8_1)); - int16x8_t filter_x = vreinterpretq_s16_u16(vmovl_u8(filter_u8_x)); - filter_0 = vaddq_s16(filter_0, vdupq_n_s16(filter_offset)); - filter_1 = vaddq_s16(filter_1, vdupq_n_s16(filter_offset)); - filter_x = vaddq_s16(filter_x, vdupq_n_s16(filter_offset)); - // Handle one output pixel at a time. - for (int outp = 0; outp < num_output_pixels; outp++) { - uint8 input_u8 = *input_ptr; - input_ptr += input_ptr_increment; - int16 input = static_cast(input_u8 + input_offset); - // Load the accumulators from acc_buffer - int32x4_t acc_0 = vld1q_s32(acc_buffer_ptr + 4 * 0); - int32x4_t acc_1 = vld1q_s32(acc_buffer_ptr + 4 * 1); - int32x4_t acc_2 = vld1q_s32(acc_buffer_ptr + 4 * 2); - int32x4_t acc_3 = vld1q_s32(acc_buffer_ptr + 4 * 3); - int32x4_t acc_4 = vld1q_s32(acc_buffer_ptr + 4 * 4); - // Multiply-accumulate - acc_0 = vmlal_n_s16(acc_0, vget_low_s16(filter_0), input); - acc_1 = vmlal_n_s16(acc_1, vget_high_s16(filter_0), input); - acc_2 = vmlal_n_s16(acc_2, vget_low_s16(filter_1), input); - acc_3 = vmlal_n_s16(acc_3, vget_high_s16(filter_1), input); - acc_4 = vmlal_n_s16(acc_4, vget_high_s16(filter_x), input); - // Store the accumulators back to acc_buffer - vst1q_s32(acc_buffer_ptr + 4 * 0, acc_0); - vst1q_s32(acc_buffer_ptr + 4 * 1, acc_1); - vst1q_s32(acc_buffer_ptr + 4 * 2, acc_2); - vst1q_s32(acc_buffer_ptr + 4 * 3, acc_3); - vst1q_s32(acc_buffer_ptr + 4 * 4, acc_4); - acc_buffer_ptr += 20; - } - } -}; - -template <> -struct QuantizedDepthwiseConvKernel { - static void Run(int num_output_pixels, int input_depth, int depth_multiplier, - const uint8* input_ptr, int16 input_offset, - int input_ptr_increment, const uint8* filter_ptr, - int16 filter_offset, int32* acc_buffer_ptr) { - // Load the filters, add filter_offset. - const uint8x8_t filter_u8 = vld1_u8(filter_ptr); - const int16x8_t filter = vaddq_s16( - vreinterpretq_s16_u16(vmovl_u8(filter_u8)), vdupq_n_s16(filter_offset)); - // Handle one output pixel at a time. - for (int outp = 0; outp < num_output_pixels; outp++) { - uint8 input_u8 = *input_ptr; - input_ptr += input_ptr_increment; - int16 input = static_cast(input_u8 + input_offset); - // Load the accumulators from acc_buffer - int32x4_t acc[2]; - for (int i = 0; i < 2; i++) { - acc[i] = vld1q_s32(acc_buffer_ptr + 4 * i); - } - // Multiply-accumulate - acc[0] = vmlal_n_s16(acc[0], vget_low_s16(filter), input); - acc[1] = vmlal_n_s16(acc[1], vget_high_s16(filter), input); - // Store the accumulators back to acc_buffer - for (int i = 0; i < 2; i++) { - vst1q_s32(acc_buffer_ptr + 4 * i, acc[i]); - } - acc_buffer_ptr += 8; - } - } -}; - -template <> -struct QuantizedDepthwiseConvKernel { - static void Run(int num_output_pixels, int input_depth, int depth_multiplier, - const uint8* input_ptr, int16 input_offset, - int input_ptr_increment, const uint8* filter_ptr, - int16 filter_offset, int32* acc_buffer_ptr) { - // Load the filters, add filter_offset. - uint8x8_t filter_u8 = vdup_n_u8(0); - filter_u8 = vset_lane_u8(filter_ptr[0], filter_u8, 0); - filter_u8 = vset_lane_u8(filter_ptr[1], filter_u8, 1); - filter_u8 = vset_lane_u8(filter_ptr[0], filter_u8, 2); - filter_u8 = vset_lane_u8(filter_ptr[1], filter_u8, 3); - const int16x4_t filter_s16 = - vreinterpret_s16_u16(vget_low_u16(vmovl_u8(filter_u8))); - const int16x4_t filter = vadd_s16(filter_s16, vdup_n_s16(filter_offset)); - - int outp = 0; - - // Handle 2 output pixels at a time. - for (; outp <= num_output_pixels - 2; outp += 2) { - // Load the accumulators from acc_buffer. - int32x4_t acc = vld1q_s32(acc_buffer_ptr); - // Load the inputs, add input_offset. - uint16x4_t input_u16 = vdup_n_u16(0); - input_u16 = vset_lane_u16((reinterpret_cast(input_ptr))[0], - input_u16, 0); - input_ptr += input_ptr_increment; - input_u16 = vset_lane_u16((reinterpret_cast(input_ptr))[0], - input_u16, 1); - input_ptr += input_ptr_increment; - const int16x4_t input_s16 = vreinterpret_s16_u16( - vget_low_u16(vmovl_u8(vreinterpret_u8_u16(input_u16)))); - const int16x4_t input = vadd_s16(input_s16, vdup_n_s16(input_offset)); - - // Multiply-accumulate. - acc = vmlal_s16(acc, filter, input); - // Store the accumulators back to acc_buffer. - vst1q_s32(acc_buffer_ptr, acc); - acc_buffer_ptr += 4; - } - - // Handle 1 output pixel at a time. - for (; outp < num_output_pixels; outp++) { - // Load the accumulators from acc_buffer. - int32x2_t acc = vld1_s32(acc_buffer_ptr); - // Load the inputs, add input_offset. - uint8x8_t input_u8 = vdup_n_u8(0); - input_u8 = vset_lane_u8(input_ptr[0], input_u8, 0); - input_u8 = vset_lane_u8(input_ptr[1], input_u8, 1); - input_ptr += input_ptr_increment; - const int16x4_t input_s16 = - vreinterpret_s16_u16(vget_low_u16(vmovl_u8(input_u8))); - const int16x4_t input = vadd_s16(input_s16, vdup_n_s16(input_offset)); - - // Multiply-accumulate. - acc = vget_low_s32(vmlal_s16(vcombine_s32(acc, acc), filter, input)); - // Store the accumulators back to acc_buffer. - vst1_s32(acc_buffer_ptr, acc); - acc_buffer_ptr += 2; - } - } -}; - -template <> -struct QuantizedDepthwiseConvKernel { - static void Run(int num_output_pixels, int input_depth, int depth_multiplier, - const uint8* input_ptr, int16 input_offset, - int input_ptr_increment, const uint8* filter_ptr, - int16 filter_offset, int32* acc_buffer_ptr) { - if (num_output_pixels <= 0) { - return; - } - - // Load the filters, add filter_offset. - uint8x8_t filter_u8 = vdup_n_u8(0); - filter_u8 = vset_lane_u8(filter_ptr[0], filter_u8, 0); - filter_u8 = vset_lane_u8(filter_ptr[1], filter_u8, 1); - filter_u8 = vset_lane_u8(filter_ptr[2], filter_u8, 2); - filter_u8 = vset_lane_u8(filter_ptr[3], filter_u8, 3); - const int16x4_t filter_s16 = - vreinterpret_s16_u16(vget_low_u16(vmovl_u8(filter_u8))); - const int16x4_t filter = vadd_s16(filter_s16, vdup_n_s16(filter_offset)); - - int outp = 0; - - // Handle one output pixel at a time until second to the last pixel. Second - // to the last because we read eight input pixels while only processing - // four. - for (; outp < num_output_pixels - 1; outp++) { - // Load the accumulators from acc_buffer - int32x4_t acc; - acc = vld1q_s32(acc_buffer_ptr); - - // Load the inputs, add input_offset. - uint8x8_t input_u8 = vld1_u8(input_ptr); - input_ptr += input_ptr_increment; - const int16x4_t input_s16 = - vreinterpret_s16_u16(vget_low_u16(vmovl_u8(input_u8))); - const int16x4_t input = vadd_s16(input_s16, vdup_n_s16(input_offset)); - // Multiply-accumulate - acc = vmlal_s16(acc, filter, input); - // Store the accumulators back to acc_buffer - vst1q_s32(acc_buffer_ptr, acc); - acc_buffer_ptr += 4; - } - - // Handle the last output pixel. - // Load the accumulators from acc_buffer - int32x4_t acc; - acc = vld1q_s32(acc_buffer_ptr); - - // Load the inputs, add input_offset. - uint8x8_t input_u8 = vdup_n_u8(0); - input_u8 = vset_lane_u8(input_ptr[0], input_u8, 0); - input_u8 = vset_lane_u8(input_ptr[1], input_u8, 1); - input_u8 = vset_lane_u8(input_ptr[2], input_u8, 2); - input_u8 = vset_lane_u8(input_ptr[3], input_u8, 3); - const int16x4_t input_s16 = - vreinterpret_s16_u16(vget_low_u16(vmovl_u8(input_u8))); - const int16x4_t input = vadd_s16(input_s16, vdup_n_s16(input_offset)); - // Multiply-accumulate - acc = vmlal_s16(acc, filter, input); - // Store the accumulators back to acc_buffer - vst1q_s32(acc_buffer_ptr, acc); - } -}; - -template <> -struct QuantizedDepthwiseConvKernel { - static void Run(int num_output_pixels, int input_depth, int depth_multiplier, - const uint8* input_ptr, int16 input_offset, - int input_ptr_increment, const uint8* filter_ptr, - int16 filter_offset, int32* acc_buffer_ptr) { - // Load the filters, add filter_offset. - uint8x8_t filter_u8_0 = vld1_u8(filter_ptr); - uint8x8_t filter_u8_1 = vld1_u8(filter_ptr + 4); - int16x8_t filter_s16_0 = vreinterpretq_s16_u16(vmovl_u8(filter_u8_0)); - int16x8_t filter_s16_1 = vreinterpretq_s16_u16(vmovl_u8(filter_u8_1)); - filter_s16_0 = vaddq_s16(filter_s16_0, vdupq_n_s16(filter_offset)); - filter_s16_1 = vaddq_s16(filter_s16_1, vdupq_n_s16(filter_offset)); - int16x4_t filter_0 = vget_low_s16(filter_s16_0); - int16x4_t filter_1 = vget_high_s16(filter_s16_0); - int16x4_t filter_2 = vget_high_s16(filter_s16_1); - - // Handle one output pixel at a time. - for (int outp = 0; outp < num_output_pixels; outp++) { - // Load the inputs, add input_offset. - uint8x8_t input_u8_0 = vld1_u8(input_ptr); - uint8x8_t input_u8_1 = vld1_u8(input_ptr + 4); - input_ptr += input_ptr_increment; - int16x8_t input_0 = vreinterpretq_s16_u16(vmovl_u8(input_u8_0)); - int16x8_t input_1 = vreinterpretq_s16_u16(vmovl_u8(input_u8_1)); - input_0 = vaddq_s16(input_0, vdupq_n_s16(input_offset)); - input_1 = vaddq_s16(input_1, vdupq_n_s16(input_offset)); - - // Load the accumulators from acc_buffer - int32x4_t acc_0 = vld1q_s32(acc_buffer_ptr + 4 * 0); - int32x4_t acc_1 = vld1q_s32(acc_buffer_ptr + 4 * 1); - int32x4_t acc_2 = vld1q_s32(acc_buffer_ptr + 4 * 2); - - // Multiply-accumulate - acc_0 = vmlal_s16(acc_0, vget_low_s16(input_0), filter_0); - acc_1 = vmlal_s16(acc_1, vget_high_s16(input_0), filter_1); - acc_2 = vmlal_s16(acc_2, vget_high_s16(input_1), filter_2); - - // Store the accumulators back to acc_buffer - vst1q_s32(acc_buffer_ptr + 4 * 0, acc_0); - vst1q_s32(acc_buffer_ptr + 4 * 1, acc_1); - vst1q_s32(acc_buffer_ptr + 4 * 2, acc_2); - - acc_buffer_ptr += 12; - } - } -}; -#endif - -// Accumulates the effect of one row of the filter, on a segment of one row -// of the output, accessing the corresponding one row of the input. -template -void QuantizedDepthwiseConvAccumRow(int stride, int dilation_factor, - int input_depth, int input_width, - const uint8* input_data, int16 input_offset, - int pad_width, int depth_multiplier, - int filter_width, const uint8* filter_data, - int16 filter_offset, int out_x_buffer_start, - int out_x_buffer_end, int output_depth, - int32* acc_buffer) { -#ifdef GEMMLOWP_PROFILING - gemmlowp::ScopedProfilingLabel label(__PRETTY_FUNCTION__); -#endif - // Sanity check parameters. This is important in particular to ensure - // that we keep the number of template instantiations minimal, so we don't - // increase binary size unnecessarily. - static_assert(kFixedDepthMultiplier || !kFixedInputDepth, ""); - static_assert(kFixedInputDepth || kAllowStrided, ""); - TFLITE_DCHECK(stride == 1 || kAllowStrided); - if (kFixedInputDepth) { - TFLITE_DCHECK_EQ(input_depth, kFixedInputDepth); - } - if (kFixedDepthMultiplier) { - TFLITE_DCHECK_EQ(depth_multiplier, kFixedDepthMultiplier); - } - TFLITE_DCHECK_EQ(output_depth, input_depth * depth_multiplier); - const int input_ptr_increment = stride * input_depth; - const uint8* filter_base_ptr = filter_data; - for (int filter_x = 0; filter_x < filter_width; ++filter_x) { - // For the current (filter_x, filter_y) point in the filter, - // compute the boundaries of the corresponding output row segment. - int out_x_loop_start_unclampled = 0; - int out_x_loop_end_unclampled = 0; - if (kAllowStrided) { - if (stride == 2) { - out_x_loop_start_unclampled = (pad_width - filter_x + 1) / 2; - out_x_loop_end_unclampled = - (pad_width + input_width - filter_x + 1) / 2; - } else if (stride == 4) { - out_x_loop_start_unclampled = (pad_width - filter_x + 3) / 4; - out_x_loop_end_unclampled = - (pad_width + input_width - filter_x + 3) / 4; - } else { - out_x_loop_start_unclampled = - (pad_width - filter_x + stride - 1) / stride; - out_x_loop_end_unclampled = - (pad_width + input_width - filter_x + stride - 1) / stride; - } - } else { - out_x_loop_start_unclampled = pad_width - filter_x; - out_x_loop_end_unclampled = pad_width + input_width - filter_x; - } - // The kernel will have to iterate on the segment of the - // output row that starts at out_x_loop_start and out_x_loop_end. - const int out_x_loop_start = - std::max(out_x_buffer_start, out_x_loop_start_unclampled); - const int out_x_loop_end = - std::min(out_x_buffer_end, out_x_loop_end_unclampled); - - int32* acc_buffer_ptr = - acc_buffer + (out_x_loop_start - out_x_buffer_start) * output_depth; - const int in_x_origin = (out_x_loop_start * stride) - pad_width + filter_x; - const uint8* input_ptr = input_data + in_x_origin * input_depth; - const int num_output_pixels = out_x_loop_end - out_x_loop_start; - QuantizedDepthwiseConvKernel< - kAllowStrided, kFixedInputDepth, - kFixedDepthMultiplier>::Run(num_output_pixels, input_depth, - depth_multiplier, input_ptr, input_offset, - input_ptr_increment, filter_base_ptr, - filter_offset, acc_buffer_ptr); - filter_base_ptr += output_depth; - } -} - -// generic fallback of DepthwiseConvAccumRow, portable, non-templatized. -inline void QuantizedDepthwiseConvAccumRowGeneric( - int stride, int dilation_factor, int input_depth, int input_width, - const uint8* input_data, int16 input_offset, int pad_width, - int depth_multiplier, int filter_width, const uint8* filter_data, - int16 filter_offset, int out_x_buffer_start, int out_x_buffer_end, - int output_depth, int32* acc_buffer) { - gemmlowp::ScopedProfilingLabel label("DepthwiseConvAccumRowGeneric (slow)"); - const uint8* filter_base_ptr = filter_data; - for (int filter_x = 0; filter_x < filter_width; ++filter_x) { - const int out_x_loop_start = std::max( - out_x_buffer_start, - (pad_width - dilation_factor * filter_x + stride - 1) / stride); - const int out_x_loop_end = std::min( - out_x_buffer_end, - (pad_width + input_width - dilation_factor * filter_x + stride - 1) / - stride); - - int32* acc_buffer_ptr = - acc_buffer + (out_x_loop_start - out_x_buffer_start) * output_depth; - const int in_x_origin = - (out_x_loop_start * stride) - pad_width + dilation_factor * filter_x; - const uint8* input_ptr = input_data + in_x_origin * input_depth; - const int input_ptr_increment = (stride - 1) * input_depth; - for (int out_x = out_x_loop_start; out_x < out_x_loop_end; out_x++) { - const uint8* filter_ptr = filter_base_ptr; - for (int ic = 0; ic < input_depth; ++ic) { - const int16 input_val = *input_ptr++ + input_offset; - for (int m = 0; m < depth_multiplier; m++) { - const int16 filter_val = *filter_ptr++ + filter_offset; - *acc_buffer_ptr++ += static_cast(filter_val) * input_val; - } - } - input_ptr += input_ptr_increment; - } - filter_base_ptr += output_depth; - } -} - -// Initializes the accumulator buffer with bias values. -inline void DepthwiseConvInitAccBuffer(int num_output_pixels, int output_depth, - const int32* bias_data, - int32* acc_buffer) { - int i = 0; -#ifdef USE_NEON - if (output_depth == 1) { - const int32x4_t b = vdupq_n_s32(bias_data[0]); - for (; i <= num_output_pixels - 16; i += 16) { - vst1q_s32(acc_buffer + i + 0, b); - vst1q_s32(acc_buffer + i + 4, b); - vst1q_s32(acc_buffer + i + 8, b); - vst1q_s32(acc_buffer + i + 12, b); - } - for (; i <= num_output_pixels - 4; i += 4) { - vst1q_s32(acc_buffer + i, b); - } - } else if (output_depth == 2) { - int32x4_t b = vdupq_n_s32(bias_data[0]); - b = vsetq_lane_s32(bias_data[1], b, 1); - b = vsetq_lane_s32(bias_data[1], b, 3); - for (; i <= num_output_pixels - 8; i += 8) { - vst1q_s32(acc_buffer + 2 * i + 0, b); - vst1q_s32(acc_buffer + 2 * i + 4, b); - vst1q_s32(acc_buffer + 2 * i + 8, b); - vst1q_s32(acc_buffer + 2 * i + 12, b); - } - for (; i <= num_output_pixels - 2; i += 2) { - vst1q_s32(acc_buffer + 2 * i, b); - } - } else if (output_depth == 4) { - const int32x4_t b = vld1q_s32(bias_data); - for (; i <= num_output_pixels - 4; i += 4) { - vst1q_s32(acc_buffer + 4 * i + 0, b); - vst1q_s32(acc_buffer + 4 * i + 4, b); - vst1q_s32(acc_buffer + 4 * i + 8, b); - vst1q_s32(acc_buffer + 4 * i + 12, b); - } - for (; i < num_output_pixels; i++) { - vst1q_s32(acc_buffer + 4 * i, b); - } - } else if (output_depth == 8) { - const int32x4_t b0 = vld1q_s32(bias_data); - const int32x4_t b1 = vld1q_s32(bias_data + 4); - for (; i <= num_output_pixels - 2; i += 2) { - vst1q_s32(acc_buffer + 8 * i + 0, b0); - vst1q_s32(acc_buffer + 8 * i + 4, b1); - vst1q_s32(acc_buffer + 8 * i + 8, b0); - vst1q_s32(acc_buffer + 8 * i + 12, b1); - } - for (; i < num_output_pixels; i++) { - vst1q_s32(acc_buffer + 8 * i + 0, b0); - vst1q_s32(acc_buffer + 8 * i + 4, b1); - } - } else if (output_depth == 16) { - const int32x4_t b0 = vld1q_s32(bias_data); - const int32x4_t b1 = vld1q_s32(bias_data + 4); - const int32x4_t b2 = vld1q_s32(bias_data + 8); - const int32x4_t b3 = vld1q_s32(bias_data + 12); - for (; i < num_output_pixels; i++) { - vst1q_s32(acc_buffer + 16 * i + 0, b0); - vst1q_s32(acc_buffer + 16 * i + 4, b1); - vst1q_s32(acc_buffer + 16 * i + 8, b2); - vst1q_s32(acc_buffer + 16 * i + 12, b3); - } - } -#endif - for (; i < num_output_pixels; i++) { - memcpy(acc_buffer + i * output_depth, bias_data, - sizeof(acc_buffer[0]) * output_depth); - } -} - -inline void DepthwiseConv( - const DepthwiseParams& params, const RuntimeShape& input_shape, - const uint8* input_data, const RuntimeShape& filter_shape, - const uint8* filter_data, const RuntimeShape& bias_shape, - const int32* bias_data, const RuntimeShape& output_shape, - uint8* output_data) { - gemmlowp::ScopedProfilingLabel label("DepthwiseConv/8bit"); - const int stride_width = params.stride_width; - const int stride_height = params.stride_height; - const int pad_width = params.padding_values.width; - const int pad_height = params.padding_values.height; - const int depth_multiplier = params.depth_multiplier; - const int32 output_activation_min = params.quantized_activation_min; - const int32 output_activation_max = params.quantized_activation_max; - const int32 input_offset = params.input_offset; - const int32 filter_offset = params.weights_offset; - const int32 output_offset = params.output_offset; - const int32 output_multiplier = params.output_multiplier; - const int output_shift = params.output_shift; - const int dilation_width_factor = params.dilation_width_factor; - const int dilation_height_factor = params.dilation_height_factor; - TFLITE_DCHECK_GE(dilation_width_factor, 1); - TFLITE_DCHECK_GE(dilation_height_factor, 1); - TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4); - TFLITE_DCHECK_EQ(filter_shape.DimensionsCount(), 4); - TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4); - TFLITE_DCHECK_LE(output_activation_min, output_activation_max); - const int batches = MatchingDim(input_shape, 0, output_shape, 0); - const int output_depth = MatchingDim(filter_shape, 3, output_shape, 3); - const int input_height = input_shape.Dims(1); - const int input_width = input_shape.Dims(2); - const int input_depth = input_shape.Dims(3); - const int filter_height = filter_shape.Dims(1); - const int filter_width = filter_shape.Dims(2); - const int output_height = output_shape.Dims(1); - const int output_width = output_shape.Dims(2); -#ifdef USE_NEON - const bool shift_left = (output_shift > 0); - const int32 multiplier_power_of_two = shift_left ? (1 << output_shift) : 1; -#endif - TFLITE_DCHECK_EQ(output_depth, input_depth * depth_multiplier); - TFLITE_DCHECK_EQ(bias_shape.FlatSize(), output_depth); - -// Enable for arm64 except for the Nvidia Linux 4 Tegra (L4T) running on -// Jetson TX-2. This compiler does not support the offsetof() macro. -#if defined(__aarch64__) && !defined(GOOGLE_L4T) - // Call kernel optimized for depthwise convolutions using 3x3 filters if - // parameters are supported. - if (Fast3x3FilterKernelSupported( - input_shape, filter_shape, stride_width, stride_height, - dilation_width_factor, dilation_height_factor, pad_width, pad_height, - depth_multiplier, output_shape, output_shift)) { - DepthwiseConv3x3Filter(params, input_shape, input_data, filter_shape, - filter_data, bias_shape, bias_data, output_shape, - output_data); - return; - } -#endif - - static const int kAccBufferMaxSize = 2048; - int32 acc_buffer[kAccBufferMaxSize]; - TFLITE_DCHECK_GE(kAccBufferMaxSize, output_depth); - const int kOutputPixelsInAccBuffer = kAccBufferMaxSize / output_depth; - const int kAccBufferActualSize = kOutputPixelsInAccBuffer * output_depth; - TFLITE_DCHECK_LE(kOutputPixelsInAccBuffer * output_depth, - kAccBufferActualSize); - TFLITE_DCHECK_LE(kAccBufferActualSize, kAccBufferMaxSize); - TFLITE_DCHECK_GE(kOutputPixelsInAccBuffer, 1); - - // row_accum_func will point to the core accumulation function to be used - // for this DepthwiseConv op. - using row_accum_func_t = decltype(&QuantizedDepthwiseConvAccumRowGeneric); - row_accum_func_t row_accum_func = nullptr; - -#define TFMINI_USE_DEPTHWISECONV_KERNEL(ALLOW_STRIDED, FIXED_INPUT_DEPTH, \ - FIXED_DEPTH_MULTIPLIER) \ - if (!row_accum_func && (stride_width == 1 || ALLOW_STRIDED) && \ - (input_depth == FIXED_INPUT_DEPTH || FIXED_INPUT_DEPTH == 0) && \ - depth_multiplier == FIXED_DEPTH_MULTIPLIER && \ - dilation_width_factor == 1 && dilation_height_factor == 1) { \ - row_accum_func = \ - QuantizedDepthwiseConvAccumRow; \ - } - -#ifdef USE_NEON - // We go over our list of kernels by decreasing order of preference - // for the cases where multiple kernels could apply. - - // Start with the fastest kernels: AllowStrided=false, fixed input depth. - - TFMINI_USE_DEPTHWISECONV_KERNEL(false, 1, 2) - TFMINI_USE_DEPTHWISECONV_KERNEL(false, 2, 2) - TFMINI_USE_DEPTHWISECONV_KERNEL(false, 4, 2) - TFMINI_USE_DEPTHWISECONV_KERNEL(false, 1, 4) - TFMINI_USE_DEPTHWISECONV_KERNEL(false, 4, 1) - TFMINI_USE_DEPTHWISECONV_KERNEL(false, 4, 4) - TFMINI_USE_DEPTHWISECONV_KERNEL(false, 8, 1) - TFMINI_USE_DEPTHWISECONV_KERNEL(false, 2, 8) - TFMINI_USE_DEPTHWISECONV_KERNEL(false, 2, 1) - TFMINI_USE_DEPTHWISECONV_KERNEL(false, 12, 1) - - // Next come the strided kernels: AllowStrided=true, fixed input depth. - // They are a bit less efficient, but allow stride!=1. - - TFMINI_USE_DEPTHWISECONV_KERNEL(true, 8, 2) - TFMINI_USE_DEPTHWISECONV_KERNEL(true, 16, 1) - TFMINI_USE_DEPTHWISECONV_KERNEL(true, 1, 16) - TFMINI_USE_DEPTHWISECONV_KERNEL(true, 1, 20) - TFMINI_USE_DEPTHWISECONV_KERNEL(true, 1, 32) - TFMINI_USE_DEPTHWISECONV_KERNEL(true, 1, 8) - TFMINI_USE_DEPTHWISECONV_KERNEL(true, 8, 1) - TFMINI_USE_DEPTHWISECONV_KERNEL(true, 2, 1) - TFMINI_USE_DEPTHWISECONV_KERNEL(true, 4, 1) - - // Finally, the kernels allowing a variable input depth, - // these are the least efficient but most general kernels. - - TFMINI_USE_DEPTHWISECONV_KERNEL(true, 0, 1) - TFMINI_USE_DEPTHWISECONV_KERNEL(true, 0, 2) - TFMINI_USE_DEPTHWISECONV_KERNEL(true, 0, 3) -#endif // USE_NEON - - // No matching fast kernel found, use slow fallback. - if (!row_accum_func) { - row_accum_func = QuantizedDepthwiseConvAccumRowGeneric; - } - -#undef TFMINI_USE_DEPTHWISECONV_KERNEL - - const int input_height_stride = input_shape.Dims(3) * input_shape.Dims(2); - const int input_batch_stride = input_height_stride * input_shape.Dims(1); - const int filter_height_stride = filter_shape.Dims(3) * filter_shape.Dims(2); - - // Now that we have determined row_accum_func, we can start work. - uint8* output_ptr = output_data; - for (int b = 0; b < batches; ++b) { - for (int out_y = 0; out_y < output_height; ++out_y) { - const int in_y_origin = (out_y * stride_height) - pad_height; - const int filter_y_start = - std::max(0, (-in_y_origin + dilation_height_factor - 1) / - dilation_height_factor); - const int filter_y_end = - std::min(filter_height, - (input_height - in_y_origin + dilation_height_factor - 1) / - dilation_height_factor); - for (int out_x_buffer_start = 0; out_x_buffer_start < output_width; - out_x_buffer_start += kOutputPixelsInAccBuffer) { - const int out_x_buffer_end = std::min( - output_width, out_x_buffer_start + kOutputPixelsInAccBuffer); - // We call a 'pixel' a group of activation that share all but the - // 'depth'/'channel' coordinate. num_output_pixels is the number of - // output pixels that we will accumulate in this loop iteration. - const int num_output_pixels = out_x_buffer_end - out_x_buffer_start; - // Initialize our local accumulator with the bias values, so we don't - // have to add them later. - DepthwiseConvInitAccBuffer(num_output_pixels, output_depth, bias_data, - acc_buffer); - // Accumulation loop. Most of the time should be spent in here. - for (int filter_y = filter_y_start; filter_y < filter_y_end; - ++filter_y) { - const int in_y = in_y_origin + dilation_height_factor * filter_y; - row_accum_func( - stride_width, dilation_width_factor, input_depth, input_width, - input_data + in_y * input_height_stride + b * input_batch_stride, - input_offset, pad_width, depth_multiplier, filter_width, - filter_data + filter_y * filter_height_stride, filter_offset, - out_x_buffer_start, out_x_buffer_end, output_depth, acc_buffer); - } - // Finished accumulating int32 values. Now need to convert them to - // the final 8bit form and store them. - gemmlowp::ScopedProfilingLabel label("downquantize+store"); - const int num_output_values = output_depth * num_output_pixels; - int i = 0; -#ifdef USE_NEON - using gemmlowp::RoundingDivideByPOT; - const int32x4_t output_offset_vec = vdupq_n_s32(output_offset); - const int32x4_t output_activation_min_vec = - vdupq_n_s32(output_activation_min); - const int32x4_t output_activation_max_vec = - vdupq_n_s32(output_activation_max); - // Handle 16 values at once. - // This allows us to issue 4 mutually independent int32 - // multiplications (vqrdmulh), which should alleviate most of their - // high latency. - for (; i <= num_output_values - 16; i += 16) { - int32x4_t acc[4]; - for (int j = 0; j < 4; j++) { - acc[j] = vld1q_s32(acc_buffer + i + 4 * j); - } - - if (!shift_left) { - // Fixed-point multiplication. - for (int j = 0; j < 4; j++) { - acc[j] = vqrdmulhq_n_s32(acc[j], output_multiplier); - } - for (int j = 0; j < 4; j++) { - acc[j] = RoundingDivideByPOT(acc[j], -output_shift); - } - } else { - // Fixed-point multiplication. - for (int j = 0; j < 4; j++) { - acc[j] = vmulq_n_s32(acc[j], multiplier_power_of_two); - acc[j] = vqrdmulhq_n_s32(acc[j], output_multiplier); - } - } - // Add the output offset. - for (int j = 0; j < 4; j++) { - acc[j] = vaddq_s32(acc[j], output_offset_vec); - } - // Apply the activation function. - for (int j = 0; j < 4; j++) { - acc[j] = vmaxq_s32(acc[j], output_activation_min_vec); - } - for (int j = 0; j < 4; j++) { - acc[j] = vminq_s32(acc[j], output_activation_max_vec); - } - // Saturating cast to uint8 and store to destination. - int16x4_t acc_s16[4]; - for (int j = 0; j < 4; j++) { - acc_s16[j] = vqmovn_s32(acc[j]); - } - const int16x8_t res_s16_0 = vcombine_s16(acc_s16[0], acc_s16[1]); - const int16x8_t res_s16_1 = vcombine_s16(acc_s16[2], acc_s16[3]); - const uint8x8_t res_u8_0 = vqmovun_s16(res_s16_0); - const uint8x8_t res_u8_1 = vqmovun_s16(res_s16_1); - vst1q_u8(output_ptr, vcombine_u8(res_u8_0, res_u8_1)); - output_ptr += 16; - } - // Handle 8 values at once. - // Not as good as 16 (now we're only issuing 2 mutually independent - // vqrdmulh instructions, so we're probably paying for their high - // latency). - for (; i <= num_output_values - 8; i += 8) { - int32x4_t acc0 = vld1q_s32(acc_buffer + i); - int32x4_t acc1 = vld1q_s32(acc_buffer + i + 4); - if (!shift_left) { - // Fixed-point multiplication. - acc0 = vqrdmulhq_n_s32(acc0, output_multiplier); - acc1 = vqrdmulhq_n_s32(acc1, output_multiplier); - // Rounding right shift. - acc0 = RoundingDivideByPOT(acc0, -output_shift); - acc1 = RoundingDivideByPOT(acc1, -output_shift); - } else { - // Fixed-point multiplication. - acc0 = vmulq_n_s32(acc0, multiplier_power_of_two); - acc0 = vqrdmulhq_n_s32(acc0, output_multiplier); - - acc1 = vmulq_n_s32(acc1, multiplier_power_of_two); - acc1 = vqrdmulhq_n_s32(acc1, output_multiplier); - } - // Add the output offset. - acc0 = vaddq_s32(acc0, output_offset_vec); - acc1 = vaddq_s32(acc1, output_offset_vec); - // Apply the activation function. - acc0 = vmaxq_s32(acc0, output_activation_min_vec); - acc1 = vmaxq_s32(acc1, output_activation_min_vec); - acc0 = vminq_s32(acc0, output_activation_max_vec); - acc1 = vminq_s32(acc1, output_activation_max_vec); - // Saturating cast to uint8 and store to destination. - const int16x4_t acc0_s16 = vqmovn_s32(acc0); - const int16x4_t acc1_s16 = vqmovn_s32(acc1); - const int16x8_t res_s16 = vcombine_s16(acc0_s16, acc1_s16); - const uint8x8_t res_u8 = vqmovun_s16(res_s16); - vst1_u8(output_ptr, res_u8); - output_ptr += 8; - } - // Handle 4 values at once. Now we're paying the full price of the - // high latency of vqrdmulh. Also, storing only 4 bytes at the end - // (without any alignment) can only be done 1 byte at a time. - // Yet, that is still worth doing to minimize the amount of leftover - // that will have to go through the very slow scalar code. - for (; i <= num_output_values - 4; i += 4) { - int32x4_t acc = vld1q_s32(acc_buffer + i); - if (!shift_left) { - // Fixed-point multiplication. - acc = vqrdmulhq_n_s32(acc, output_multiplier); - // Rounding right shift. - acc = RoundingDivideByPOT(acc, -output_shift); - } else { - // Fixed-point multiplication. - acc = vmulq_n_s32(acc, multiplier_power_of_two); - acc = vqrdmulhq_n_s32(acc, output_multiplier); - } - // Add the output offset. - acc = vaddq_s32(acc, output_offset_vec); - // Apply the activation function. - acc = vmaxq_s32(acc, output_activation_min_vec); - acc = vminq_s32(acc, output_activation_max_vec); - // Saturating cast to uint8 and store to destination. - const int16x4_t acc_s16 = vqmovn_s32(acc); - const int16x8_t res_s16 = vcombine_s16(acc_s16, acc_s16); - const uint8x8_t res_u8 = vqmovun_s16(res_s16); - vst1_lane_u8(output_ptr + 0, res_u8, 0); - vst1_lane_u8(output_ptr + 1, res_u8, 1); - vst1_lane_u8(output_ptr + 2, res_u8, 2); - vst1_lane_u8(output_ptr + 3, res_u8, 3); - output_ptr += 4; - } -#endif // USE_NEON - - // Handle leftover values, one by one. This is very slow. - for (; i < num_output_values; i++) { - int32 acc = acc_buffer[i]; - acc = MultiplyByQuantizedMultiplier(acc, output_multiplier, - output_shift); - acc += output_offset; - acc = std::max(acc, output_activation_min); - acc = std::min(acc, output_activation_max); - *output_ptr++ = static_cast(acc); - } - } - } - } -} - -} // namespace optimized_ops -} // namespace tflite - -#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_OPTIMIZED_DEPTHWISECONV_UINT8_H_ diff --git a/tensorflow/contrib/lite/kernels/internal/quantization_util.cc b/tensorflow/contrib/lite/kernels/internal/quantization_util.cc deleted file mode 100644 index 544ef16ce18a36e52acb8813021800189150b13f..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/internal/quantization_util.cc +++ /dev/null @@ -1,369 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include -#include -#include - -#include "tensorflow/contrib/lite/kernels/internal/compatibility.h" -#include "tensorflow/contrib/lite/kernels/internal/quantization_util.h" -#include "tensorflow/contrib/lite/kernels/internal/round.h" - -namespace tflite { - -namespace { -// These constants are used to manipulate the binary representation of doubles. -// Double-precision binary64 floating point format is: -// Bit | 63 | 62-52 | 51-0 | -// | Sign | Exponent | Fraction | -// To avoid 64-bit integers as much as possible, I break this into high and -// low 32-bit chunks. High is: -// Bit | 31 | 30-20 | 19-0 | -// | Sign | Exponent | High Fraction | -// Low is: -// Bit | 31-0 | -// | Low Fraction | -// We then access the components through logical bit-wise operations to -// extract the parts needed, with the positions and masks derived from the -// layout shown above. -constexpr uint64_t kSignMask = 0x8000000000000000LL; -constexpr uint64_t kExponentMask = 0x7ff0000000000000LL; -constexpr int32_t kExponentShift = 52; -constexpr int32_t kExponentBias = 1023; -constexpr uint32_t kExponentIsBadNum = 0x7ff; -constexpr uint64_t kFractionMask = 0x000fffffffc00000LL; -constexpr uint32_t kFractionShift = 22; -constexpr uint32_t kFractionRoundingMask = 0x003fffff; -constexpr uint32_t kFractionRoundingThreshold = 0x00200000; -} // namespace - -void QuantizeMultiplier(double double_multiplier, int32_t* quantized_multiplier, - int* shift) { - if (double_multiplier == 0.) { - *quantized_multiplier = 0; - *shift = 0; - return; - } -#ifdef TFLITE_EMULATE_FLOAT - // If we're trying to avoid the use of floating-point instructions (for - // example on microcontrollers) then use an alternative implementation - // that only requires integer and bitwise operations. To enable this, you - // need to set the define during the build process for your platform. - int64_t q_fixed = IntegerFrExp(double_multiplier, shift); -#else // TFLITE_EMULATE_FLOAT - const double q = std::frexp(double_multiplier, shift); - auto q_fixed = static_cast(TfLiteRound(q * (1ll << 31))); -#endif // TFLITE_EMULATE_FLOAT - TFLITE_CHECK(q_fixed <= (1ll << 31)); - if (q_fixed == (1ll << 31)) { - q_fixed /= 2; - ++*shift; - } - TFLITE_CHECK_LE(q_fixed, std::numeric_limits::max()); - *quantized_multiplier = static_cast(q_fixed); -} - -void QuantizeMultiplierGreaterThanOne(double double_multiplier, - int32_t* quantized_multiplier, - int* left_shift) { - TFLITE_CHECK_GT(double_multiplier, 1.); - QuantizeMultiplier(double_multiplier, quantized_multiplier, left_shift); - TFLITE_CHECK_GE(*left_shift, 0); -} - -void QuantizeMultiplierSmallerThanOneExp(double double_multiplier, - int32_t* quantized_multiplier, - int* left_shift) { - TFLITE_CHECK_LT(double_multiplier, 1.); - TFLITE_CHECK_GT(double_multiplier, 0.); - int shift; - QuantizeMultiplier(double_multiplier, quantized_multiplier, &shift); - TFLITE_CHECK_LE(shift, 0); - *left_shift = shift; -} - -int64_t IntegerFrExp(double input, int* shift) { - // Make sure our assumptions about the double layout hold. - TFLITE_CHECK_EQ(8, sizeof(double)); - - // We want to access the bits of the input double value directly, which is - // tricky to do safely, so use a union to handle the casting. - union { - double double_value; - uint64_t double_as_uint; - } cast_union; - cast_union.double_value = input; - const uint64_t u = cast_union.double_as_uint; - - // If the bitfield is all zeros apart from the sign bit, this is a normalized - // zero value, so return standard values for this special case. - if ((u & ~kSignMask) == 0) { - *shift = 0; - return 0; - } - - // Deal with NaNs and Infs, which are always indicated with a fixed pattern in - // the exponent, and distinguished by whether the fractions are zero or - // non-zero. - const uint32_t exponent_part = ((u & kExponentMask) >> kExponentShift); - if (exponent_part == kExponentIsBadNum) { - *shift = std::numeric_limits::max(); - if (u & kFractionMask) { - // NaN, so just return zero (with the exponent set to INT_MAX). - return 0; - } else { - // Infinity, so return +/- INT_MAX. - if (u & kSignMask) { - return std::numeric_limits::min(); - } else { - return std::numeric_limits::max(); - } - } - } - - // The shift is fairly easy to extract from the high bits of the double value, - // just by masking it out and applying a bias. The std::frexp() implementation - // always returns values between 0.5 and 1.0 though, whereas the exponent - // assumes 1.0 to 2.0 is the standard range, so I add on one to match that - // interface. - *shift = (exponent_part - kExponentBias) + 1; - - // There's an implicit high bit in the double format definition, so make sure - // we include that at the top, and then reconstruct the rest of the fractional - // value from the remaining fragments. - int64_t fraction = 0x40000000 + ((u & kFractionMask) >> kFractionShift); - - // We're cutting off some bits at the bottom, so to exactly match the standard - // frexp implementation here we'll apply rounding by adding one to the least - // significant bit of the result if the discarded portion is over half of the - // maximum. - if ((u & kFractionRoundingMask) > kFractionRoundingThreshold) { - fraction += 1; - } - // Negate the fraction if the sign bit was set. - if (u & kSignMask) { - fraction *= -1; - } - - return fraction; -} - -double DoubleFromFractionAndShift(int64_t fraction, int shift) { - union { - double double_value; - uint64_t double_as_uint; - } result; - - // Detect NaNs and infinities. - if (shift == std::numeric_limits::max()) { - if (fraction == 0) { - return NAN; - } else if (fraction > 0) { - return INFINITY; - } else { - return -INFINITY; - } - } - - // Return a normalized zero for a zero fraction. - if (fraction == 0) { - result.double_as_uint = 0; - return result.double_value; - } - - bool is_negative = (fraction < 0); - int64_t encoded_fraction = is_negative ? -fraction : fraction; - int64_t encoded_shift = (shift - 1); - while (encoded_fraction < 0x40000000) { - encoded_fraction *= 2; - encoded_shift -= 1; - } - while (encoded_fraction > 0x80000000) { - encoded_fraction /= 2; - encoded_shift += 1; - } - encoded_fraction -= 0x40000000; - if (encoded_shift < -1022) { - encoded_shift = -1023; - } else if (encoded_shift > 1022) { - encoded_shift = 1023; - } - encoded_shift += kExponentBias; - uint64_t encoded_sign = is_negative ? kSignMask : 0; - result.double_as_uint = encoded_sign | (encoded_shift << kExponentShift) | - (encoded_fraction << kFractionShift); - return result.double_value; -} - -double IntegerDoubleMultiply(double a, double b) { - int a_shift; - const int64_t a_fraction = IntegerFrExp(a, &a_shift); - int b_shift; - const int64_t b_fraction = IntegerFrExp(b, &b_shift); - // Detect NaNs and infinities. - if (a_shift == std::numeric_limits::max() || - (b_shift == std::numeric_limits::max())) { - return NAN; - } - const int result_shift = a_shift + b_shift + 1; - const int64_t result_fraction = (a_fraction * b_fraction) >> 32; - return DoubleFromFractionAndShift(result_fraction, result_shift); -} - -int IntegerDoubleCompare(double a, double b) { - int a_shift; - const int64_t a_fraction = IntegerFrExp(a, &a_shift); - int b_shift; - const int64_t b_fraction = IntegerFrExp(b, &b_shift); - - // Detect NaNs and infinities. - if (a_shift == std::numeric_limits::max() || - (b_shift == std::numeric_limits::max())) { - return 1; - } - - if ((a_fraction == 0) && (b_fraction < 0)) { - return 1; - } else if ((a_fraction < 0) && (b_fraction == 0)) { - return -1; - } else if (a_shift < b_shift) { - return -1; - } else if (a_shift > b_shift) { - return 1; - } else if (a_fraction < b_fraction) { - return -1; - } else if (a_fraction > b_fraction) { - return 1; - } else { - return 0; - } -} - -void PreprocessSoftmaxScaling(double beta, double input_scale, - int input_integer_bits, - int32_t* quantized_multiplier, int* left_shift) { - // If the overall multiplier (input and beta) is large, then exp() of an - // input difference of 1 scaled by this will be large. In other words, we - // can cap the multiplier and know that, when it is used, the output will be - // (round to) zero wherever the input is not at the maximum value. - - // If the overall scale is less than one, and input_integer_bits=0, then the - // result is double equivalent of Q0.31 (actually with more precision). Thus - // this generates a Q(input_integer_bits).(31-input_integer_bits) - // representation. -#ifdef TFLITE_EMULATE_FLOAT - const double input_beta = IntegerDoubleMultiply(beta, input_scale); - int shift; - int64_t fraction = IntegerFrExp(input_beta, &shift); - shift += (31 - input_integer_bits); - double input_beta_real_multiplier = - DoubleFromFractionAndShift(fraction, shift); - if (IntegerDoubleCompare(input_beta_real_multiplier, (1ll << 31) - 1.0) > 0) { - input_beta_real_multiplier = (1ll << 31) - 1.0; - } -#else // TFLITE_EMULATE_FLOAT - const double input_beta_real_multiplier = std::min( - beta * input_scale * (1 << (31 - input_integer_bits)), (1ll << 31) - 1.0); -#endif // TFLITE_EMULATE_FLOAT - - QuantizeMultiplierGreaterThanOne(input_beta_real_multiplier, - quantized_multiplier, left_shift); -} - -void PreprocessLogSoftmaxScalingExp(double beta, double input_scale, - int input_integer_bits, - int32_t* quantized_multiplier, - int* left_shift, - int32_t* reverse_scaling_divisor, - int* reverse_scaling_left_shift) { - PreprocessSoftmaxScaling(beta, input_scale, input_integer_bits, - quantized_multiplier, left_shift); - - // Also calculate what amounts to the inverse scaling factor for the input. - const double real_reverse_scaling_divisor = - (1 << (31 - *left_shift)) / static_cast(*quantized_multiplier); - tflite::QuantizeMultiplierSmallerThanOneExp(real_reverse_scaling_divisor, - reverse_scaling_divisor, - reverse_scaling_left_shift); -} - -int CalculateInputRadius(int input_integer_bits, int input_left_shift) { -#ifdef TFLITE_EMULATE_FLOAT - int64_t result = (1 << input_integer_bits) - 1; - result <<= (31 - input_integer_bits); - result >>= input_left_shift; - return result; -#else // TFLITE_EMULATE_FLOAT - const double max_input_rescaled = 1.0 * ((1 << input_integer_bits) - 1) * - (1ll << (31 - input_integer_bits)) / - (1ll << input_left_shift); - // Tighten bound using floor. Suppose that we could use the exact value. - // After scaling the difference, the result would be at the maximum. Thus we - // must ensure that our value has lower magnitude. - return static_cast(std::floor(max_input_rescaled)); -#endif // TFLITE_EMULATE_FLOAT -} - -void NudgeQuantizationRange(const float min, const float max, - const int quant_min, const int quant_max, - float* nudged_min, float* nudged_max, - float* nudged_scale) { - // This code originates from tensorflow/core/kernels/fake_quant_ops_functor.h. - const float quant_min_float = static_cast(quant_min); - const float quant_max_float = static_cast(quant_max); - *nudged_scale = (max - min) / (quant_max_float - quant_min_float); - const float zero_point_from_min = quant_min_float - min / *nudged_scale; - uint16 nudged_zero_point; - if (zero_point_from_min < quant_min_float) { - nudged_zero_point = static_cast(quant_min); - } else if (zero_point_from_min > quant_max_float) { - nudged_zero_point = static_cast(quant_max); - } else { - nudged_zero_point = static_cast(TfLiteRound(zero_point_from_min)); - } - *nudged_min = (quant_min_float - nudged_zero_point) * (*nudged_scale); - *nudged_max = (quant_max_float - nudged_zero_point) * (*nudged_scale); -} - -void FakeQuantizeArray(const float nudged_scale, const float nudged_min, - const float nudged_max, const float* input_data, - float* output_data, const float size) { - // This code originates from tensorflow/core/kernels/fake_quant_ops_functor.h. - const float inv_nudged_scale = 1.0f / nudged_scale; - - for (int i = 0; i < size; i++) { - const float src_val = input_data[i]; - const float clamped = std::min(nudged_max, std::max(nudged_min, src_val)); - const float clamped_shifted = clamped - nudged_min; - const float dst_val = - TfLiteRound(clamped_shifted * inv_nudged_scale) * nudged_scale + - nudged_min; - output_data[i] = dst_val; - } -} - -bool CheckedLog2(const float x, int* log2_result) { - // Using TfLiteRound instead of std::round and std::log instead of - // std::log2 to work around these fuctions being missing in a toolchain - // used in some TensorFlow tests as of May 2018. - const float x_log2 = std::log(x) * (1.0f / std::log(2.0f)); - const float x_log2_rounded = TfLiteRound(x_log2); - const float x_log2_fracpart = x_log2 - x_log2_rounded; - - *log2_result = static_cast(x_log2_rounded); - return std::abs(x_log2_fracpart) < 1e-3; -} - -} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/internal/quantization_util.h b/tensorflow/contrib/lite/kernels/internal/quantization_util.h deleted file mode 100644 index d74a1bac97f86cba1e63e9141a9d00ded3c63c8f..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/internal/quantization_util.h +++ /dev/null @@ -1,280 +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_CONTRIB_LITE_KERNELS_INTERNAL_QUANTIZATION_UTIL_H_ -#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_QUANTIZATION_UTIL_H_ - -#include -#include -#include - -#include "tensorflow/contrib/lite/kernels/internal/compatibility.h" -#include "tensorflow/contrib/lite/kernels/internal/round.h" -#include "tensorflow/contrib/lite/kernels/internal/types.h" - -namespace tflite { - -// Given the min and max values of a float array, return -// reasonable quantization parameters to use for this array. -template -QuantizationParams ChooseQuantizationParams(double rmin, double rmax, - bool narrow_range) { - const T qmin = std::numeric_limits::min() + (narrow_range ? 1 : 0); - const T qmax = std::numeric_limits::max(); - const double qmin_double = qmin; - const double qmax_double = qmax; - // 0 should always be a representable value. Let's assume that the initial - // min,max range contains 0. - TFLITE_CHECK_LE(rmin, 0.); - TFLITE_CHECK_GE(rmax, 0.); - if (rmin == rmax) { - // Special case where the min,max range is a point. Should be {0}. - TFLITE_CHECK_EQ(rmin, 0.); - TFLITE_CHECK_EQ(rmax, 0.); - QuantizationParams quantization_params; - quantization_params.zero_point = 0; - quantization_params.scale = 0.; - return quantization_params; - } - - // General case. - // - // First determine the scale. - const double scale = (rmax - rmin) / (qmax_double - qmin_double); - - // Zero-point computation. - // First the initial floating-point computation. The zero-point can be - // determined from solving an affine equation for any known pair - // (real value, corresponding quantized value). - // We know two such pairs: (rmin, qmin) and (rmax, qmax). - // The arithmetic error on the zero point computed from either pair - // will be roughly machine_epsilon * (sum of absolute values of terms) - // so we want to use the variant that adds the smaller terms. - const double zero_point_from_min = qmin_double - rmin / scale; - const double zero_point_from_max = qmax_double - rmax / scale; - const double zero_point_from_min_error = - std::abs(qmin_double) + std::abs(rmin / scale); - const double zero_point_from_max_error = - std::abs(qmax_double) + std::abs(rmax / scale); - - const double zero_point_double = - zero_point_from_min_error < zero_point_from_max_error - ? zero_point_from_min - : zero_point_from_max; - - // Now we need to nudge the zero point to be an integer - // (our zero points are integer, and this is motivated by the requirement - // to be able to represent the real value "0" exactly as a quantized value, - // which is required in multiple places, for example in Im2col with SAME - // padding). - T nudged_zero_point = 0; - if (zero_point_double < qmin_double) { - nudged_zero_point = qmin; - } else if (zero_point_double > qmax_double) { - nudged_zero_point = qmax; - } else { - nudged_zero_point = static_cast(round(zero_point_double)); - } - // The zero point should always be in the range of quantized value, - // [qmin, qmax]. - TFLITE_CHECK_GE(nudged_zero_point, qmin); - TFLITE_CHECK_LE(nudged_zero_point, qmax); - - // Finally, store the result nudged quantization params. - QuantizationParams quantization_params; - quantization_params.zero_point = nudged_zero_point; - quantization_params.scale = scale; - return quantization_params; -} - -template -QuantizationParams ChooseQuantizationParams(double rmin, double rmax) { - return ChooseQuantizationParams(rmin, rmax, false); -} - -// Converts a floating-point number to an integer. For all inputs x where -// static_cast(x) is legal according to the C++ standard, the result -// is identical to that cast (i.e. the result is x with its fractional part -// truncated whenever that is representable as IntOut). -// -// static_cast would cause undefined behavior for the following cases, which -// have well-defined behavior for this function: -// -// 1. If x is NaN, the result is zero. -// -// 2. If the truncated form of x is above the representable range of IntOut, -// the result is std::numeric_limits::max(). -// -// 3. If the truncated form of x is below the representable range of IntOut, -// the result is std::numeric_limits::min(). -// -// Note that cases #2 and #3 cover infinities as well as finite numbers. -// -// The range of FloatIn must include the range of IntOut, otherwise -// the results are undefined. -// TODO(sfeuz): Replace by absl::SafeCast once available. -template -IntOut SafeCast(FloatIn x) { - static_assert(!std::numeric_limits::is_integer, - "FloatIn is integer"); - static_assert(std::numeric_limits::is_integer, - "IntOut is not integer"); - static_assert(std::numeric_limits::radix == 2, "IntOut is base 2"); - - // Special case NaN, for which the logic below doesn't work. - if (std::isnan(x)) { - return 0; - } - - // Negative values all clip to zero for unsigned results. - if (!std::numeric_limits::is_signed && x < 0) { - return 0; - } - - // Handle infinities. - if (std::isinf(x)) { - return x < 0 ? std::numeric_limits::min() - : std::numeric_limits::max(); - } - - // Set exp such that x == f * 2^exp for some f with |f| in [0.5, 1.0), - // unless x is zero in which case exp == 0. Note that this implies that the - // magnitude of x is strictly less than 2^exp. - int exp = 0; - std::frexp(x, &exp); - - // Let N be the number of non-sign bits in the representation of IntOut. If - // the magnitude of x is strictly less than 2^N, the truncated version of x - // is representable as IntOut. The only representable integer for which this - // is not the case is kMin for signed types (i.e. -2^N), but that is covered - // by the fall-through below. - if (exp <= std::numeric_limits::digits) { - return x; - } - - // Handle numbers with magnitude >= 2^N. - return x < 0 ? std::numeric_limits::min() - : std::numeric_limits::max(); -} - -// Decompose a double multiplier into a Q0.31 int32 representation of its -// significand, and shift representation of NEGATIVE its exponent --- -// this is intended as a RIGHT-shift. -// -// Restricted to the case where the multiplier < 1 (and non-negative). -void QuantizeMultiplierSmallerThanOneExp(double double_multiplier, - int32_t* quantized_multiplier, - int* left_shift); - -// Decompose a double multiplier into a Q0.31 int32 representation of its -// significand, and shift representation of its exponent. -// -// Restricted to the case where the multiplier > 1. -void QuantizeMultiplierGreaterThanOne(double double_multiplier, - int32_t* quantized_multiplier, - int* left_shift); - -// Decompose a double multiplier into a Q0.31 int32 representation of its -// significand, and shift representation of its exponent. -// -// Handles an arbitrary positive multiplier. The 'shift' output-value is -// basically the 'floating-point exponent' of the multiplier: -// Negative for a right-shift (when the multiplier is <1), positive for a -// left-shift (when the multiplier is >1) -void QuantizeMultiplier(double double_multiplier, int32_t* quantized_multiplier, - int* shift); - -// Splits a double input value into a returned fraction, and a shift value from -// the exponent, using only bitwise and integer operations to support -// microcontrollers and other environments without floating-point support. -// -// This is designed to be a replacement for how std::frexp() is used within the -// QuantizeMultiplier() function, and so has a different signature than the -// standard version, returning a 64-bit integer rather than a double. This -// result has a maximum value of 1<<31, with the fraction expressed as a -// proportion of that maximum. -// -// std::frexp() returns NaNs and infinities unmodified, but since we're -// returning integers that can't represent those values, instead we return -// a shift of std::numeric_limits::max() for all bad numbers, with an int64 -// result of 0 for NaNs, std:numeric_limits::max() for +INFINITY, and -// std::numeric_limits::min() for -INFINITY. Denormalized inputs will -// result in return values that end up truncating some bits at the end, -// reflecting the loss of precision inherent in denormalization. -int64_t IntegerFrExp(double input, int* shift); - -// Converts an integer fraction in the format produced by IntegerFrExp (where -// 0x40000000 is 1.0) and an exponent shift (between -1022 and +1022) into an -// IEEE binary64 double format result. The implementation uses only integer and -// bitwise operators, so no floating point hardware support or emulation is -// needed. This is here so quantized operations can run non-time-critical -// preparation calculations on microcontrollers and other platforms without -// float support. -double DoubleFromFractionAndShift(int64_t fraction, int shift); - -// Performs a multiplication of two numbers in double format, using only integer -// and bitwise instructions. This is aimed at supporting housekeeping functions -// for quantized operations on microcontrollers without floating-point hardware. -double IntegerDoubleMultiply(double a, double b); - -// Returns -1 if a is less than b, 0 if a and b are equal, and +1 if a is -// greater than b. It is implemented using only integer and logical instructions -// so that it can be easily run on microcontrollers for quantized operations. -int IntegerDoubleCompare(double a, double b); - -// This first creates a multiplier in a double equivalent of -// Q(input_integer_bits).(31-input_integer_bits) representation, with extra -// precision in the double's fractional bits. It then splits the result into -// significand and exponent. -void PreprocessSoftmaxScaling(double beta, double input_scale, - int input_integer_bits, - int32_t* quantized_multiplier, int* left_shift); -// Like PreprocessSoftmaxScaling, but inverse scaling factors also calculated. -void PreprocessLogSoftmaxScalingExp(double beta, double input_scale, - int input_integer_bits, - int32_t* quantized_multiplier, - int* left_shift, - int32_t* reverse_scaling_divisor, - int* reverse_scaling_left_shift); -// Calculate the largest input that will result in a within-bounds intermediate -// result within MultiplyByQuantizedMultiplierGreaterThanOne. In other words, -// it must not overflow before we reduce the value by multiplication by the -// input multiplier. The negative radius is used as the minimum difference in -// Softmax. -int CalculateInputRadius(int input_integer_bits, int input_left_shift); - -// Nudges a min/max quantization range to ensure zero is zero. -// Gymnastics with nudged zero point is to ensure that real zero maps to -// an integer, which is required for e.g. zero-padding in convolutional layers. -// Outputs nudged_min, nudged_max, nudged_scale. -void NudgeQuantizationRange(const float min, const float max, - const int quant_min, const int quant_max, - float* nudged_min, float* nudged_max, - float* nudged_scale); - -// Fake quantizes (quantizes and dequantizes) input_data using the scale, -// nudged_min, and nudged_max from NudgeQuantizationRange. This matches the code -// in TensorFlow's FakeQuantizeWithMinMaxVarsFunctor. -void FakeQuantizeArray(const float nudged_scale, const float nudged_min, - const float nudged_max, const float* input_data, - float* output_data, const float size); - -// If x is approximately a power of two (with any positive or negative -// exponent), stores that exponent (i.e. log2(x)) in *log2_result, otherwise -// returns false. -bool CheckedLog2(const float x, int* log2_result); - -} // namespace tflite - -#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_QUANTIZATION_UTIL_H_ diff --git a/tensorflow/contrib/lite/kernels/internal/reference/depthwiseconv_float.h b/tensorflow/contrib/lite/kernels/internal/reference/depthwiseconv_float.h deleted file mode 100644 index 11224270a4b17f4299703eaaa0dfd49b42b2a321..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/internal/reference/depthwiseconv_float.h +++ /dev/null @@ -1,100 +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_CONTRIB_LITE_KERNELS_INTERNAL_REFERENCE_DEPTHWISECONV_FLOAT_H_ -#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_REFERENCE_DEPTHWISECONV_FLOAT_H_ - -#include "tensorflow/contrib/lite/kernels/internal/common.h" -#include "tensorflow/contrib/lite/kernels/internal/compatibility.h" -#include "tensorflow/contrib/lite/kernels/internal/types.h" - -namespace tflite { -namespace reference_ops { - -inline void DepthwiseConv( - const DepthwiseParams& params, const RuntimeShape& input_shape, - const float* input_data, const RuntimeShape& filter_shape, - const float* filter_data, const RuntimeShape& bias_shape, - const float* bias_data, const RuntimeShape& output_shape, - float* output_data) { - const int stride_width = params.stride_width; - const int stride_height = params.stride_height; - const int dilation_width_factor = params.dilation_width_factor; - const int dilation_height_factor = params.dilation_height_factor; - const int pad_width = params.padding_values.width; - const int pad_height = params.padding_values.height; - const int depth_multiplier = params.depth_multiplier; - const float output_activation_min = params.float_activation_min; - const float output_activation_max = params.float_activation_max; - TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4); - TFLITE_DCHECK_EQ(filter_shape.DimensionsCount(), 4); - TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4); - - const int batches = MatchingDim(input_shape, 0, output_shape, 0); - const int output_depth = MatchingDim(filter_shape, 3, output_shape, 3); - const int input_height = input_shape.Dims(1); - const int input_width = input_shape.Dims(2); - const int input_depth = input_shape.Dims(3); - const int filter_height = filter_shape.Dims(1); - const int filter_width = filter_shape.Dims(2); - const int output_height = output_shape.Dims(1); - const int output_width = output_shape.Dims(2); - TFLITE_DCHECK_EQ(output_depth, input_depth * depth_multiplier); - TFLITE_DCHECK_EQ(bias_shape.FlatSize(), output_depth); - - for (int b = 0; b < batches; ++b) { - for (int out_y = 0; out_y < output_height; ++out_y) { - for (int out_x = 0; out_x < output_width; ++out_x) { - for (int ic = 0; ic < input_depth; ++ic) { - for (int m = 0; m < depth_multiplier; m++) { - const int oc = m + ic * depth_multiplier; - const int in_x_origin = (out_x * stride_width) - pad_width; - const int in_y_origin = (out_y * stride_height) - pad_height; - float total = 0.f; - for (int filter_y = 0; filter_y < filter_height; ++filter_y) { - for (int filter_x = 0; filter_x < filter_width; ++filter_x) { - const int in_x = in_x_origin + dilation_width_factor * filter_x; - const int in_y = - in_y_origin + dilation_height_factor * filter_y; - // If the location is outside the bounds of the input image, - // use zero as a default value. - if ((in_x >= 0) && (in_x < input_width) && (in_y >= 0) && - (in_y < input_height)) { - float input_value = - input_data[Offset(input_shape, b, in_y, in_x, ic)]; - float filter_value = filter_data[Offset( - filter_shape, 0, filter_y, filter_x, oc)]; - total += (input_value * filter_value); - } - } - } - float bias_value = 0.0f; - if (bias_data) { - bias_value = bias_data[oc]; - } - output_data[Offset(output_shape, b, out_y, out_x, oc)] = - ActivationFunctionWithMinMax(total + bias_value, - output_activation_min, - output_activation_max); - } - } - } - } - } -} - -} // end namespace reference_ops -} // end namespace tflite - -#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_REFERENCE_DEPTHWISECONV_FLOAT_H_ diff --git a/tensorflow/contrib/lite/kernels/internal/reference/depthwiseconv_uint8.h b/tensorflow/contrib/lite/kernels/internal/reference/depthwiseconv_uint8.h deleted file mode 100644 index eab28e6c84c77fd9a12a97203e05b79c3ab5fb31..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/internal/reference/depthwiseconv_uint8.h +++ /dev/null @@ -1,112 +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_CONTRIB_LITE_KERNELS_INTERNAL_REFERENCE_DEPTHWISECONV_UINT8_H_ -#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_REFERENCE_DEPTHWISECONV_UINT8_H_ - -#include - -#include "fixedpoint/fixedpoint.h" -#include "tensorflow/contrib/lite/kernels/internal/common.h" -#include "tensorflow/contrib/lite/kernels/internal/compatibility.h" -#include "tensorflow/contrib/lite/kernels/internal/types.h" - -namespace tflite { -namespace reference_ops { - -inline void DepthwiseConv( - const DepthwiseParams& params, const RuntimeShape& input_shape, - const uint8* input_data, const RuntimeShape& filter_shape, - const uint8* filter_data, const RuntimeShape& bias_shape, - const int32* bias_data, const RuntimeShape& output_shape, - uint8* output_data) { - const int stride_width = params.stride_width; - const int stride_height = params.stride_height; - const int dilation_width_factor = params.dilation_width_factor; - const int dilation_height_factor = params.dilation_height_factor; - const int pad_width = params.padding_values.width; - const int pad_height = params.padding_values.height; - const int depth_multiplier = params.depth_multiplier; - const int32 output_activation_min = params.quantized_activation_min; - const int32 output_activation_max = params.quantized_activation_max; - const int32 input_offset = params.input_offset; - const int32 filter_offset = params.weights_offset; - const int32 output_offset = params.output_offset; - const int32 output_multiplier = params.output_multiplier; - const int output_shift = params.output_shift; - TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4); - TFLITE_DCHECK_EQ(filter_shape.DimensionsCount(), 4); - TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4); - - TFLITE_DCHECK_LE(output_activation_min, output_activation_max); - const int batches = MatchingDim(input_shape, 0, output_shape, 0); - const int output_depth = MatchingDim(filter_shape, 3, output_shape, 3); - const int input_height = input_shape.Dims(1); - const int input_width = input_shape.Dims(2); - const int input_depth = input_shape.Dims(3); - const int filter_height = filter_shape.Dims(1); - const int filter_width = filter_shape.Dims(2); - const int output_height = output_shape.Dims(1); - const int output_width = output_shape.Dims(2); - TFLITE_DCHECK_EQ(output_depth, input_depth * depth_multiplier); - TFLITE_DCHECK_EQ(bias_shape.FlatSize(), output_depth); - - for (int b = 0; b < batches; ++b) { - for (int out_y = 0; out_y < output_height; ++out_y) { - for (int out_x = 0; out_x < output_width; ++out_x) { - for (int ic = 0; ic < input_depth; ++ic) { - for (int m = 0; m < depth_multiplier; m++) { - const int oc = m + ic * depth_multiplier; - const int in_x_origin = (out_x * stride_width) - pad_width; - const int in_y_origin = (out_y * stride_height) - pad_height; - int32 acc = 0; - for (int filter_y = 0; filter_y < filter_height; ++filter_y) { - for (int filter_x = 0; filter_x < filter_width; ++filter_x) { - const int in_x = in_x_origin + dilation_width_factor * filter_x; - const int in_y = - in_y_origin + dilation_height_factor * filter_y; - // If the location is outside the bounds of the input image, - // use zero as a default value. - if ((in_x >= 0) && (in_x < input_width) && (in_y >= 0) && - (in_y < input_height)) { - int32 input_val = - input_data[Offset(input_shape, b, in_y, in_x, ic)]; - int32 filter_val = filter_data[Offset( - filter_shape, 0, filter_y, filter_x, oc)]; - acc += - (filter_val + filter_offset) * (input_val + input_offset); - } - } - } - if (bias_data) { - acc += bias_data[oc]; - } - acc = MultiplyByQuantizedMultiplier(acc, output_multiplier, - output_shift); - acc += output_offset; - acc = std::max(acc, output_activation_min); - acc = std::min(acc, output_activation_max); - output_data[Offset(output_shape, b, out_y, out_x, oc)] = - static_cast(acc); - } - } - } - } - } -} - -} // end namespace reference_ops -} // end namespace tflite - -#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_REFERENCE_DEPTHWISECONV_UINT8_H_ diff --git a/tensorflow/contrib/lite/kernels/internal/reference/softmax.h b/tensorflow/contrib/lite/kernels/internal/reference/softmax.h deleted file mode 100644 index 7d442961349e349b3101d8a1798ee7dd388426d3..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/internal/reference/softmax.h +++ /dev/null @@ -1,179 +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_CONTRIB_LITE_KERNELS_INTERNAL_REFERENCE_SOFTMAX_H_ -#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_REFERENCE_SOFTMAX_H_ - -#include "fixedpoint/fixedpoint.h" -#include "tensorflow/contrib/lite/kernels/internal/common.h" -#include "tensorflow/contrib/lite/kernels/internal/quantization_util.h" -#include "tensorflow/contrib/lite/kernels/internal/round.h" -#include "tensorflow/contrib/lite/kernels/internal/types.h" -#include "tensorflow/contrib/lite/kernels/op_macros.h" - -namespace tflite { -namespace reference_ops { - -inline void Softmax(const SoftmaxParams& params, - const RuntimeShape& input_shape, const float* input_data, - const RuntimeShape& output_shape, float* output_data) { - const int trailing_dim = input_shape.DimensionsCount() - 1; - const int outer_size = - MatchingFlatSizeSkipDim(input_shape, trailing_dim, output_shape); - const int depth = - MatchingDim(input_shape, trailing_dim, output_shape, trailing_dim); - - for (int i = 0; i < outer_size; ++i) { - // Find max element value which we'll use to ensure numerical stability - // taking advantage of the following equality: - // exp(x[i])/sum(exp(x[i])) == exp(x[i]+C)/sum(exp(x[i]+C)) - float max = std::numeric_limits::lowest(); - for (int c = 0; c < depth; ++c) { - max = std::max(max, input_data[i * depth + c]); - } - - // Compute sum. - float sum = 0.f; - for (int c = 0; c < depth; ++c) { - sum += std::exp((input_data[i * depth + c] - max) * params.beta); - } - - // Compute result. - for (int c = 0; c < depth; ++c) { - output_data[i * depth + c] = - std::exp((input_data[i * depth + c] - max) * params.beta) / sum; - } - } -} - -inline void Softmax(const SoftmaxParams& params, - const RuntimeShape& input_shape, const uint8* input_data, - const RuntimeShape& output_shape, uint8* output_data) { - const int32 input_beta_multiplier = params.input_multiplier; - const int32 input_beta_left_shift = params.input_left_shift; - const int diff_min = params.diff_min; - // The representation chosen for the input to the exp() function is Q5.26. - // We need to leave extra space since values that we skip might be as large as - // -32 before multiplying by input_beta_multiplier, and therefore as large as - // -16 afterwards. Note that exp(-8) is definitely not insignificant to - // accumulation, but exp(-16) definitely is. - static const int kScaledDiffIntegerBits = 5; - static const int kAccumulationIntegerBits = 12; - using FixedPointScaledDiff = - gemmlowp::FixedPoint; - using FixedPointAccum = gemmlowp::FixedPoint; - using FixedPoint0 = gemmlowp::FixedPoint; - - const int trailing_dim = input_shape.DimensionsCount() - 1; - const int outer_size = - MatchingFlatSizeSkipDim(input_shape, trailing_dim, output_shape); - const int depth = - MatchingDim(input_shape, trailing_dim, output_shape, trailing_dim); - - for (int i = 0; i < outer_size; ++i) { - uint8 max_in_row = 0; - for (int c = 0; c < depth; ++c) { - max_in_row = std::max(max_in_row, input_data[i * depth + c]); - } - - FixedPointAccum sum_of_exps = FixedPointAccum::Zero(); - for (int c = 0; c < depth; ++c) { - int32 input_diff = - static_cast(input_data[i * depth + c]) - max_in_row; - if (input_diff >= diff_min) { - const int32 input_diff_rescaled = - MultiplyByQuantizedMultiplierGreaterThanOne( - input_diff, input_beta_multiplier, input_beta_left_shift); - const FixedPointScaledDiff scaled_diff_f8 = - FixedPointScaledDiff::FromRaw(input_diff_rescaled); - sum_of_exps = sum_of_exps + gemmlowp::Rescale( - exp_on_negative_values(scaled_diff_f8)); - } - } - - int32 fixed_sum_of_exps = sum_of_exps.raw(); - int headroom_plus_one = - CountLeadingZeros(static_cast(fixed_sum_of_exps)); - // This is the number of bits to the left of the binary point above 1.0. - // Consider fixed_sum_of_exps=1.25. In that case shifted_scale=0.8 and - // no later adjustment will be needed. - int num_bits_over_unit = kAccumulationIntegerBits - headroom_plus_one; - int32 shifted_sum_minus_one = static_cast( - (static_cast(fixed_sum_of_exps) << headroom_plus_one) - - (static_cast(1) << 31)); - - FixedPoint0 shifted_scale = gemmlowp::one_over_one_plus_x_for_x_in_0_1( - FixedPoint0::FromRaw(shifted_sum_minus_one)); - - for (int c = 0; c < depth; ++c) { - int32 input_diff = - static_cast(input_data[i * depth + c]) - max_in_row; - if (input_diff >= diff_min) { - const int32 input_diff_rescaled = - MultiplyByQuantizedMultiplierGreaterThanOne( - input_diff, input_beta_multiplier, input_beta_left_shift); - const FixedPointScaledDiff scaled_diff_f8 = - FixedPointScaledDiff::FromRaw(input_diff_rescaled); - - FixedPoint0 exp_in_0 = exp_on_negative_values(scaled_diff_f8); - int32 unsat_output = gemmlowp::RoundingDivideByPOT( - (shifted_scale * exp_in_0).raw(), num_bits_over_unit + 31 - 8); - - output_data[i * depth + c] = static_cast( - std::max(std::min(unsat_output, static_cast(255)), - static_cast(0))); - - } else { - output_data[i * depth + c] = 0; - } - } - } -} - -// Performs softmax along the input of size (input_size * batch_size). -inline void Softmax(const float* in, const int input_size, const int batch_size, - const float beta, float* out) { - // TF_LITE_ASSERT(input_size > 0); - - // For each batch - for (int b = 0; b < batch_size; b++) { - // Find the max coeff. - float max_coeff = in[0]; - for (int i = 1; i < input_size; i++) { - if (in[i] > max_coeff) max_coeff = in[i]; - } - - // Compute the normalized sum of exps. - float exp_sum = 0.0; - for (int i = 0; i < input_size; i++) { - out[i] = std::exp((in[i] - max_coeff) * beta); - exp_sum += out[i]; - } - - // Divide by the sum of exps. - float reciprocal_sum_exp = 1.f / exp_sum; - for (int i = 0; i < input_size; i++) { - out[i] *= reciprocal_sum_exp; - } - - // Advance in and out pointers for the next batch. - in += input_size; - out += input_size; - } -} - -} // namespace reference_ops -} // namespace tflite - -#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_REFERENCE_SOFTMAX_H_ diff --git a/tensorflow/contrib/lite/kernels/internal/resize_bilinear_test.cc b/tensorflow/contrib/lite/kernels/internal/resize_bilinear_test.cc deleted file mode 100644 index 15df31f75a69b9c0076eb4978e06707b5966417d..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/internal/resize_bilinear_test.cc +++ /dev/null @@ -1,138 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include -#include -#include - -#include -#include "tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h" -#include "tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h" -#include "tensorflow/contrib/lite/kernels/internal/test_util.h" -#include "tensorflow/contrib/lite/kernels/internal/types.h" - -namespace tflite { -namespace { -template -void TestOneResizeBilinear(int batch, int depth, int input_width, - int input_height, int output_width, - int output_height, float error_threshold) { - RuntimeShape input_dims_inference({batch, input_height, input_width, depth}); - RuntimeShape output_dims_inference( - {batch, output_height, output_width, depth}); - - const int input_buffer_size = input_dims_inference.FlatSize(); - const int output_buffer_size = output_dims_inference.FlatSize(); - - std::vector input_data(input_buffer_size, 0); - std::vector reference_output_data(output_buffer_size, 0); - // Initialize the output data with something other than zero, so we can catch - // issue with kernels failing to initialize the output. - std::vector output_data(output_buffer_size, 3); - - const T min_amplitude = static_cast(0); - const T max_amplitude = static_cast(255); - FillRandom(&input_data, min_amplitude, max_amplitude); - - RuntimeShape output_size_dims({1, 1, 1, 2}); - std::vector output_size_data = {output_height, output_width}; - - tflite::ResizeBilinearParams op_params; - op_params.align_corners = false; - - reference_ops::ResizeBilinear(op_params, input_dims_inference, - input_data.data(), output_size_dims, - output_size_data.data(), output_dims_inference, - reference_output_data.data()); - optimized_ops::ResizeBilinear( - op_params, input_dims_inference, input_data.data(), output_size_dims, - output_size_data.data(), output_dims_inference, output_data.data()); - - double sum_diff = 0; - float max_abs_val = 0; - for (int i = 0; i < output_buffer_size; i++) { - sum_diff += std::abs(static_cast(output_data[i]) - - static_cast(reference_output_data[i])); - max_abs_val = std::max( - max_abs_val, std::abs(static_cast(reference_output_data[i]))); - } - - if (sum_diff != 0.f) { - const float mean_diff = static_cast(sum_diff / output_buffer_size); - const float relative_error = std::abs(mean_diff) / max_abs_val; - ASSERT_LT(relative_error, error_threshold); - } -} - -TEST(ResizeBilinear, TestResizeBilinear8Bit) { - const int kTestsToRun = 100 * 1000; - for (int i = 0; i < kTestsToRun; i++) { - const int batch = ExponentialRandomPositiveInt(0.9f, 3, 20); - const int depth = ExponentialRandomPositiveInt(0.9f, 6, 50); - const int input_width = ExponentialRandomPositiveInt(0.9f, 20, 200); - const int input_height = ExponentialRandomPositiveInt(0.9f, 20, 200); - const int output_width = ExponentialRandomPositiveInt(0.9f, 20, 200); - const int output_height = ExponentialRandomPositiveInt(0.9f, 20, 200); - - TestOneResizeBilinear(batch, depth, input_width, input_height, - output_width, output_height, 0.025); - } -} - -TEST(ResizeBilinear2x2, TestResizeBilinear8Bit) { - const int kTestsToRun = 100 * 1000; - for (int i = 0; i < kTestsToRun; i++) { - const int batch = ExponentialRandomPositiveInt(0.9f, 3, 20); - const int depth = ExponentialRandomPositiveInt(0.9f, 6, 50); - const int input_width = ExponentialRandomPositiveInt(0.9f, 20, 200); - const int input_height = ExponentialRandomPositiveInt(0.9f, 20, 200); - const int output_width = input_width * 2; - const int output_height = input_height * 2; - - TestOneResizeBilinear(batch, depth, input_width, input_height, - output_width, output_height, 1e-5); - } -} - -TEST(ResizeBilinear, TestResizeBilinear) { - const int kTestsToRun = 100 * 1000; - for (int i = 0; i < kTestsToRun; i++) { - const int batch = ExponentialRandomPositiveInt(0.9f, 3, 20); - const int depth = ExponentialRandomPositiveInt(0.9f, 6, 50); - const int input_width = ExponentialRandomPositiveInt(0.9f, 20, 200); - const int input_height = ExponentialRandomPositiveInt(0.9f, 20, 200); - const int output_width = ExponentialRandomPositiveInt(0.9f, 20, 200); - const int output_height = ExponentialRandomPositiveInt(0.9f, 20, 200); - - TestOneResizeBilinear(batch, depth, input_width, input_height, - output_width, output_height, 1e-5); - } -} - -TEST(ResizeBilinear2x2, TestResizeBilinear) { - const int kTestsToRun = 100 * 1000; - for (int i = 0; i < kTestsToRun; i++) { - const int batch = ExponentialRandomPositiveInt(0.9f, 3, 20); - const int depth = ExponentialRandomPositiveInt(0.9f, 6, 50); - const int input_width = ExponentialRandomPositiveInt(0.9f, 20, 200); - const int input_height = ExponentialRandomPositiveInt(0.9f, 20, 200); - const int output_width = input_width * 2; - const int output_height = input_height * 2; - - TestOneResizeBilinear(batch, depth, input_width, input_height, - output_width, output_height, 1e-5); - } -} -} // namespace -} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/internal/tensor_utils.cc b/tensorflow/contrib/lite/kernels/internal/tensor_utils.cc deleted file mode 100644 index f4181b18a8f46fd9bef4b81a210a6b8134a4e9d0..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/internal/tensor_utils.cc +++ /dev/null @@ -1,28 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include "tensorflow/contrib/lite/kernels/internal/tensor_utils.h" -#include "tensorflow/contrib/lite/kernels/internal/common.h" - -#ifndef USE_NEON -#if defined(__ARM_NEON__) || defined(__ARM_NEON) -#define USE_NEON -#endif // defined(__ARM_NEON__) || defined(__ARM_NEON) -#endif // USE_NEON - -#ifdef USE_NEON -#include "tensorflow/contrib/lite/kernels/internal/optimized/neon_tensor_utils.h" -#else -#include "tensorflow/contrib/lite/kernels/internal/reference/portable_tensor_utils.h" -#endif // USE_NEON diff --git a/tensorflow/contrib/lite/kernels/internal/tensor_utils_test.cc b/tensorflow/contrib/lite/kernels/internal/tensor_utils_test.cc deleted file mode 100644 index 6458af714b8c714f7132dc17adf4eca20ece3e37..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/internal/tensor_utils_test.cc +++ /dev/null @@ -1,806 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include "tensorflow/contrib/lite/kernels/internal/tensor_utils.h" -#include -#include "tensorflow/contrib/lite/c/builtin_op_data.h" -#include "tensorflow/contrib/lite/kernels/test_util.h" - -namespace tflite { -namespace tensor_utils { - -TEST(uKernels, ClipTest) { - constexpr int kVectorSize = 10; - constexpr float kAbsLimit = 2.0; - static float input[kVectorSize] = {0.0, -0.5, 1.0, -1.5, 2.0, - -2.5, 3.0, -3.5, 4.0, -4.5}; - std::vector output(kVectorSize); - ClipVector(input, kVectorSize, kAbsLimit, output.data()); - EXPECT_THAT(output, - ElementsAreArray(ArrayFloatNear( - {0.0, -0.5, 1.0, -1.5, 2.0, -2.0, 2.0, -2.0, 2.0, -2.0}))); -} - -TEST(uKernels, VectorScalarMultiply) { - constexpr int kVectorSize = 29; - static int8_t input[kVectorSize]; - for (int i = 0; i < 29; ++i) { - input[i] = static_cast(i - 14); - } - const float scale = 0.1f; - std::vector output(kVectorSize, 0.0f); - VectorScalarMultiply(input, kVectorSize, scale, output.data()); - EXPECT_THAT(output, - ElementsAreArray(ArrayFloatNear( - {-1.4, -1.3, -1.2, -1.1, -1.0, -0.9, -0.8, -0.7, -0.6, -0.5, - -0.4, -0.3, -0.2, -0.1, 0, 0.1, 0.2, 0.3, 0.4, 0.5, - 0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2, 1.3, 1.4}))); -} - -TEST(uKernels, IsZeroTest) { - constexpr int kVectorSize = 21; - static float zeros[kVectorSize] = {0.0}; - EXPECT_TRUE(IsZeroVector(zeros, kVectorSize)); - - static float nonzeros[kVectorSize] = { - 1e-6, 1e-7, 1e-8, 1e-9, 1e-10, 1e-11, 1e-12, - 1e-13, 1e-14, 1e-15, 1e-16, 1e-17, 1e-18, 1e-19, - 1e-20, 1e-21, 1e-22, 1e-23, 1e-24, 1e-25, 1e-26}; - EXPECT_FALSE(IsZeroVector(nonzeros, kVectorSize)); -} - -TEST(uKernels, GeneratedIsZeroTest) { - constexpr int kVectorSize = 39; - std::vector input(kVectorSize); - ZeroVector(input.data(), kVectorSize); - EXPECT_TRUE(IsZeroVector(input.data(), kVectorSize)); -} - -TEST(uKernels, SymmetricQuantizeFloatsTest) { - constexpr int kVectorSize = 9; - static float input[kVectorSize] = {-640, -635.0, -630, 10.0, 2.0, - -5.0, -10.0, 0.0, 1000.0}; - - int8_t output[kVectorSize]; - float min, max, scaling_factor; - SymmetricQuantizeFloats(input, kVectorSize, output, &min, &max, - &scaling_factor); - - EXPECT_EQ(min, -640); - EXPECT_EQ(max, 1000); - // EQ won't work due to fpoint. - EXPECT_NEAR(scaling_factor, 1000 / 127.0, 1e-6); - EXPECT_THAT(output, - testing::ElementsAreArray({-81, -81, -80, 1, 0, -1, -1, 0, 127})); -} - -TEST(uKernels, SymmetricQuantizeFloatsAllZerosTest) { - constexpr int kVectorSize = 9; - static float input[kVectorSize] = {0, 0, 0, 0, 0, 0, 0, 0, 0}; - - int8_t output[kVectorSize]; - float min, max, scaling_factor; - SymmetricQuantizeFloats(input, kVectorSize, output, &min, &max, - &scaling_factor); - - EXPECT_EQ(min, 0); - EXPECT_EQ(max, 0); - EXPECT_EQ(scaling_factor, 1); - EXPECT_THAT(output, testing::ElementsAreArray({0, 0, 0, 0, 0, 0, 0, 0, 0})); -} - -TEST(uKernels, SymmetricQuantizeFloatsAllAlmostZeroTest) { - constexpr int kVectorSize = 9; - static float input[kVectorSize] = {-1e-5, 3e-5, -7e-6, -9e-5, 1e-6, - 4e-5, 9e-6, 2e-4, 0}; - - int8_t output[kVectorSize]; - float min, max, scaling_factor; - SymmetricQuantizeFloats(input, kVectorSize, output, &min, &max, - &scaling_factor); - - EXPECT_NEAR(min, -9e-05, 1e-6); - EXPECT_NEAR(max, 0.0002, 1e-6); - EXPECT_NEAR(scaling_factor, 1.57e-6, 1e-6); - EXPECT_THAT(output, - testing::ElementsAreArray({-6, 19, -4, -57, 1, 25, 6, 127, 0})); -} - -TEST(uKernels, MatrixBatchVectorMultiplyAccumulateTest) { - constexpr int kRow = 3; - constexpr int kCol = 4; - constexpr int kBatch = 2; - static float matrix[kRow * kCol] = {1.0, 2.0, 3.0, 4.0, // - -1.0, -2.0, -3.0, -4.0, // - 1.0, -2.0, 3.0, -4.0}; - static float vector[kCol * kBatch] = {1.0, -1.0, 1.0, -1.0, // - 2.0, -2.0, 2.0, -2.0}; - std::vector output(kRow * kBatch); - std::fill(output.begin(), output.end(), 3.0); - MatrixBatchVectorMultiplyAccumulate(matrix, kRow, kCol, vector, kBatch, - output.data(), /*result_stride=*/1); - EXPECT_THAT(output, ElementsAreArray(ArrayFloatNear({1., 5., 13., // - -1., 7., 23.}))); - - std::vector output_with_stride2(kRow * kBatch * 2); - std::fill(output_with_stride2.begin(), output_with_stride2.end(), 3.0); - MatrixBatchVectorMultiplyAccumulate(matrix, kRow, kCol, vector, kBatch, - output_with_stride2.data(), - /*result_stride=*/2); - EXPECT_THAT(output_with_stride2, - ElementsAreArray(ArrayFloatNear({1., 3., 5., 3., 13., 3., // - -1., 3., 7., 3., 23., 3.}))); -} - -#ifdef __ANDROID__ -TEST(uKernels, MatrixBatchVectorMultiplyAccumulateSymmetricQuantizedTest) { - // Note we use 29 columns as this exercises all the neon kernel: the - // 16-block SIMD code, the 8-block postamble, and the leftover postamble. - const int a_rows = 4, a_cols = 29; - const int kWeightsPerUint32 = 4; - const float a_float_data[] = { - /* 1st row */ - 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9, 10.1, 11.11, 12.12, 13.13, - 14.14, 15.15, 16.16, 17.17, 18.18, 19.19, 20.2, 21.21, 22.22, 23.23, - 24.24, 25.25, 26.26, 27.27, 28.28, 0, - /* 2nd row */ - -1.1, -2.2, -3.3, -4.4, -5.5, -6.6, -7.7, -8.8, -9.9, -10.1, -11.11, - -12.12, -13.13, -14.14, -15.15, -16.16, -17.17, -18.18, -19.19, -20.2, - -21.21, -22.22, -23.23, -24.24, -25.25, -26.26, -27.27, -28.28, 0, - /* 3rd row */ - 1.1, -2.2, 3.3, -4.4, 5.5, -6.6, 7.7, -8.8, 9.9, -10.1, 11.11, -12.12, - 13.13, -14.14, 15.15, -16.16, 17.17, -18.18, 19.19, -20.2, 21.21, -22.22, - 23.23, -24.24, 25.25, -26.26, 27.27, -28.28, 0, - /* 4th row */ - -1.1, 2.2, -3.3, 4.4, -5.5, 6.6, -7.7, 8.8, -9.9, 10.1, -11.11, 12.12, - -13.13, 14.14, -15.15, 16.16, -17.17, 18.18, -19.19, 20.2, -21.21, 22.22, - -23.23, 24.24, -25.25, 26.26, -27.27, 28.28, 0}; - - int8_t* a_int8_data = reinterpret_cast( - aligned_malloc(a_rows * a_cols, kWeightsPerUint32)); - float a_min, a_max; - float scaling_factor_a; - SymmetricQuantizeFloats(a_float_data, a_rows * a_cols, a_int8_data, &a_min, - &a_max, &scaling_factor_a); - const int8_t expected_a_int8_data[] = { - /* 1st row */ - 5, - 10, - 15, - 20, - 25, - 30, - 35, - 40, - 44, - 45, - 50, - 54, - 59, - 64, - 68, - 73, - 77, - 82, - 86, - 91, - 95, - 100, - 104, - 109, - 113, - 118, - 122, - 127, - 0, - /* 2nd row */ - -5, - -10, - -15, - -20, - -25, - -30, - -35, - -40, - -44, - -45, - -50, - -54, - -59, - -64, - -68, - -73, - -77, - -82, - -86, - -91, - -95, - -100, - -104, - -109, - -113, - -118, - -122, - -127, - 0, - /* 3rd row */ - 5, - -10, - 15, - -20, - 25, - -30, - 35, - -40, - 44, - -45, - 50, - -54, - 59, - -64, - 68, - -73, - 77, - -82, - 86, - -91, - 95, - -100, - 104, - -109, - 113, - -118, - 122, - -127, - 0, - /* 4th row */ - -5, - 10, - -15, - 20, - -25, - 30, - -35, - 40, - -44, - 45, - -50, - 54, - -59, - 64, - -68, - 73, - -77, - 82, - -86, - 91, - -95, - 100, - -104, - 109, - -113, - 118, - -122, - 127, - 0, - }; - for (int i = 0; i < a_rows * a_cols; ++i) { - EXPECT_EQ(expected_a_int8_data[i], a_int8_data[i]); - } - - const int b_rows = 29, b_cols = 1, batches = 2; - const float b_float_data[] = { - /* batch 1 */ - 1.0, - -1.0, - 1.0, - -1.0, - 1.0, - -1.0, - 1.0, - -1.0, - 1.0, - -1.0, - 1.0, - -1.0, - 1.0, - -1.0, - 1.0, - -1.0, - 1.0, - -1.0, - 1.0, - -1.0, - 1.0, - -1.0, - 1.0, - -1.0, - 1.0, - -1.0, - 1.0, - -1.0, - 1.0, - /* batch 2 */ - 2.5, - -2.1, - 3.0, - -1.3, - 1.3, - -1.1, - 2.0, - -1.7, - 1.9, - -1.5, - 0.5, - -0.7, - 0.8, - -0.3, - 2.8, - -2.8, - 1.1, - -2.3, - 1.9, - -1.9, - 2.1, - -0.5, - 2.4, - -0.1, - 1.0, - -2.5, - 0.7, - -1.9, - 0.2, - }; - - // Quantized values of B: - int8_t b_int8_data[b_rows * b_cols * batches]; - float b_min, b_max; - float scaling_factor_b[batches]; - SymmetricQuantizeFloats(b_float_data, b_rows * b_cols, b_int8_data, &b_min, - &b_max, &scaling_factor_b[0]); - SymmetricQuantizeFloats(&b_float_data[b_rows * b_cols], b_rows * b_cols, - &b_int8_data[b_rows * b_cols], &b_min, &b_max, - &scaling_factor_b[1]); - - const int8_t expected_b_int8_data[] = { - /* batch 1 */ - 127, - -127, - 127, - -127, - 127, - -127, - 127, - -127, - 127, - -127, - 127, - -127, - 127, - -127, - 127, - -127, - 127, - -127, - 127, - -127, - 127, - -127, - 127, - -127, - 127, - -127, - 127, - -127, - 127, - /* batch 2 */ - 106, - -89, - 127, - -55, - 55, - -47, - 85, - -72, - 80, - -64, - 21, - -30, - 34, - -13, - 119, - -119, - 47, - -97, - 80, - -80, - 89, - -21, - 102, - -4, - 42, - -106, - 30, - -80, - 8, - }; - for (int i = 0; i < b_rows * b_cols * batches; ++i) { - EXPECT_EQ(expected_b_int8_data[i], b_int8_data[i]); - } - - // Full float operation results in: - // -13.69, 13.69, 414.11, -414.11 - // -6.325, 6.325, 631.263, -631.263 - float c_float_data[a_rows * b_cols * batches]; - for (int i = 0; i < a_rows * b_cols * batches; ++i) { - c_float_data[i] = 0.0; - } - - // Testing product. - const float scaling_factor_c[2] = { - scaling_factor_a * scaling_factor_b[0], - scaling_factor_a * scaling_factor_b[1], - }; - MatrixBatchVectorMultiplyAccumulate(a_int8_data, a_rows, a_cols, b_int8_data, - scaling_factor_c, batches, c_float_data, - /*result_stride=*/1); - - // Assert we obtain the expected recovered float values. - const float expected_c_float_data[] = { - -14.474, 14.474, 414.402, -414.402, -6.92228, 6.92228, 632.042, -632.042, - }; - for (int i = 0; i < a_rows * b_cols * batches; ++i) { - EXPECT_NEAR(expected_c_float_data[i], c_float_data[i], 0.001); - } - - aligned_free(a_int8_data); -} -#endif // __ANDROID__ - -TEST(uKernels, VectorVectorCwiseProductTest) { - constexpr int kVectorSize = 10; - static float input1[kVectorSize] = {0.0, -0.5, 1.0, -1.5, 2.0, - -2.5, 3.0, -3.5, 4.0, -4.5}; - static float input2[kVectorSize] = {0.1, -0.1, 0.1, -0.1, 0.1, - -0.1, 0.1, -0.1, 0.1, -0.1}; - std::vector output(kVectorSize); - VectorVectorCwiseProduct(input1, input2, kVectorSize, output.data()); - EXPECT_THAT(output, - ElementsAreArray(ArrayFloatNear( - {0.0, 0.05, 0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45}))); -} - -TEST(uKernels, VectorVectorCwiseProductAccumulateTest) { - constexpr int kVectorSize = 10; - static float input1[kVectorSize] = {0.0, -0.5, 1.0, -1.5, 2.0, - -2.5, 3.0, -3.5, 4.0, -4.5}; - static float input2[kVectorSize] = {0.1, -0.1, 0.1, -0.1, 0.1, - -0.1, 0.1, -0.1, 0.1, -0.1}; - std::vector output(kVectorSize); - std::fill(output.begin(), output.end(), 1.0); - VectorVectorCwiseProductAccumulate(input1, input2, kVectorSize, - output.data()); - EXPECT_THAT(output, - ElementsAreArray(ArrayFloatNear( - {1.0, 1.05, 1.1, 1.15, 1.2, 1.25, 1.3, 1.35, 1.4, 1.45}))); -} - -TEST(uKernels, VectorBatchVectorAddTest) { - constexpr int kVectorSize = 3; - constexpr int kBatchSize = 2; - static float input[kVectorSize] = {0.0, -0.5, 1.0}; - std::vector output = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0}; - VectorBatchVectorAdd(input, kVectorSize, kBatchSize, output.data()); - EXPECT_THAT(output, - testing::ElementsAreArray({1.0, 1.5, 4.0, 4.0, 4.5, 7.0})); -} - -TEST(uKernels, VectorBatchVectorAssignTest) { - constexpr int kVectorSize = 5; - constexpr int kBatchSize = 3; - static float input[kVectorSize] = {0.0, -0.5, 1.0, -1.5, 2.0}; - std::vector output(kVectorSize * kBatchSize); - VectorBatchVectorAssign(input, kVectorSize, kBatchSize, output.data()); - EXPECT_THAT(output, ElementsAreArray(ArrayFloatNear( - {0.0, -0.5, 1.0, -1.5, 2.0, 0.0, -0.5, 1.0, -1.5, 2.0, - 0.0, -0.5, 1.0, -1.5, 2.0}))); -} - -TEST(uKernels, ApplySigmoidToVectorTest) { - constexpr int kVectorSize = 5; - static float input[kVectorSize] = {0.0, -0.5, 1.0, -1.5, 2.0}; - std::vector output(kVectorSize); - ApplySigmoidToVector(input, kVectorSize, output.data()); - EXPECT_THAT(output, ElementsAreArray(ArrayFloatNear( - {0.5, 0.377541, 0.731059, 0.182426, 0.880797}))); -} - -TEST(uKernels, ApplyActivationToVectorTest) { - constexpr int kVectorSize = 5; - static float input[kVectorSize] = {0.0, -0.5, 1.0, -1.5, 2.0}; - std::vector output(kVectorSize); - ApplyActivationToVector(input, kVectorSize, kTfLiteActRelu, output.data()); - EXPECT_THAT(output, - ElementsAreArray(ArrayFloatNear({0.0, 0.0, 1.0, 0.0, 2.0}))); - - ApplyActivationToVector(input, kVectorSize, kTfLiteActTanh, output.data()); - EXPECT_THAT(output, ElementsAreArray(ArrayFloatNear( - {0.0, -0.462117, 0.761594, -0.905148, 0.964028}))); -} - -TEST(uKernels, CopyVectorTest) { - constexpr int kVectorSize = 5; - static float input[kVectorSize] = {0.0, -0.5, 1.0, -1.5, 2.0}; - std::vector output(kVectorSize); - CopyVector(input, kVectorSize, output.data()); - EXPECT_THAT(output, - ElementsAreArray(ArrayFloatNear({0.0, -0.5, 1.0, -1.5, 2.0}))); -} - -TEST(uKernels, Sub1VectorTest) { - constexpr int kVectorSize = 5; - static float input[kVectorSize] = {0.0, -0.5, 1.0, -1.5, 2.0}; - std::vector output(kVectorSize); - Sub1Vector(input, kVectorSize, output.data()); - EXPECT_THAT(output, - ElementsAreArray(ArrayFloatNear({1.0, 1.5, 0.0, 2.5, -1.0}))); -} - -TEST(uKernels, ZeroVectorTest) { - constexpr int kVectorSize = 5; - std::vector output(kVectorSize); - ZeroVector(output.data(), kVectorSize); - EXPECT_THAT(output, - ElementsAreArray(ArrayFloatNear({0.0, 0.0, 0.0, 0.0, 0.0}))); -} - -TEST(uKernels, VectorBatchVectorCwiseProductAccumulate) { - constexpr int kVectorSize = 29; - constexpr int kBatchSize = 4; - static float input[kVectorSize] = { - 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9, 10.1, - 11.11, 12.12, 13.13, 14.14, 15.15, 16.16, 17.17, 18.18, 19.19, 20.2, - 21.21, 22.22, 23.23, 24.24, 25.25, 26.26, 27.27, 28.28, 0}; - std::vector output = { - /* batch 0 */ - 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9, 10.1, 11.11, 12.12, 13.13, - 14.14, 15.15, 16.16, 17.17, 18.18, 19.19, 20.2, 21.21, 22.22, 23.23, - 24.24, 25.25, 26.26, 27.27, 28.28, 0, - /* batch 1 */ - -1.1, -2.2, -3.3, -4.4, -5.5, -6.6, -7.7, -8.8, -9.9, -10.1, -11.11, - -12.12, -13.13, -14.14, -15.15, -16.16, -17.17, -18.18, -19.19, -20.2, - -21.21, -22.22, -23.23, -24.24, -25.25, -26.26, -27.27, -28.28, 0, - /* batch 2 */ - 1.1, -2.2, 3.3, -4.4, 5.5, -6.6, 7.7, -8.8, 9.9, -10.1, 11.11, -12.12, - 13.13, -14.14, 15.15, -16.16, 17.17, -18.18, 19.19, -20.2, 21.21, -22.22, - 23.23, -24.24, 25.25, -26.26, 27.27, -28.28, 0, - /* batch 3 */ - -1.1, 2.2, -3.3, 4.4, -5.5, 6.6, -7.7, 8.8, -9.9, 10.1, -11.11, 12.12, - -13.13, 14.14, -15.15, 16.16, -17.17, 18.18, -19.19, 20.2, -21.21, 22.22, - -23.23, 24.24, -25.25, 26.26, -27.27, 28.28, 0}; - VectorBatchVectorCwiseProductAccumulate(input, kVectorSize, output.data(), - kBatchSize, output.data()); - - // Expect output = input * output + output. - const std::vector expected_output = { - /* batch 0 */ - 2.310000, 7.040000, 14.190000, 23.760000, 35.750000, 50.159996, 66.989998, - 86.240005, 107.909996, 112.110008, 134.542084, 159.014389, 185.526901, - 214.079605, 244.672485, 277.305603, 311.978912, 348.692413, 387.446136, - 428.240051, 471.074066, 515.948364, 562.862854, 611.817566, 662.812500, - 715.847595, 770.922974, 828.038452, 0.000000, - /* batch 1 */ - -2.310000, -7.040000, -14.190000, -23.760000, -35.750000, -50.159996, - -66.989998, -86.240005, -107.909996, -112.110008, -134.542084, - -159.014389, -185.526901, -214.079605, -244.672485, -277.305603, - -311.978912, -348.692413, -387.446136, -428.240051, -471.074066, - -515.948364, -562.862854, -611.817566, -662.812500, -715.847595, - -770.922974, -828.038452, 0.000000, - /* batch 2 */ - 2.310000, -7.040000, 14.190000, -23.760000, 35.750000, -50.159996, - 66.989998, -86.240005, 107.909996, -112.110008, 134.542084, -159.014389, - 185.526901, -214.079605, 244.672485, -277.305603, 311.978912, -348.692413, - 387.446136, -428.240051, 471.074066, -515.948364, 562.862854, -611.817566, - 662.812500, -715.847595, 770.922974, -828.038452, 0.000000, - /* batch 3 */ - -2.310000, 7.040000, -14.190000, 23.760000, -35.750000, 50.159996, - -66.989998, 86.240005, -107.909996, 112.110008, -134.542084, 159.014389, - -185.526901, 214.079605, -244.672485, 277.305603, -311.978912, 348.692413, - -387.446136, 428.240051, -471.074066, 515.948364, -562.862854, 611.817566, - -662.812500, 715.847595, -770.922974, 828.038452, 0.000000}; - EXPECT_THAT(output, testing::ElementsAreArray(expected_output)); -} - -TEST(uKernels, VectorBatchVectorCwiseProductNoAccumulate) { - constexpr int kVectorSize = 29; - constexpr int kBatchSize = 4; - static float input[kVectorSize] = { - 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9, 10.1, - 11.11, 12.12, 13.13, 14.14, 15.15, 16.16, 17.17, 18.18, 19.19, 20.2, - 21.21, 22.22, 23.23, 24.24, 25.25, 26.26, 27.27, 28.28, 0}; - std::vector output = { - /* batch 0 */ - 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9, 10.1, 11.11, 12.12, 13.13, - 14.14, 15.15, 16.16, 17.17, 18.18, 19.19, 20.2, 21.21, 22.22, 23.23, - 24.24, 25.25, 26.26, 27.27, 28.28, 0, - /* batch 1 */ - -1.1, -2.2, -3.3, -4.4, -5.5, -6.6, -7.7, -8.8, -9.9, -10.1, -11.11, - -12.12, -13.13, -14.14, -15.15, -16.16, -17.17, -18.18, -19.19, -20.2, - -21.21, -22.22, -23.23, -24.24, -25.25, -26.26, -27.27, -28.28, 0, - /* batch 2 */ - 1.1, -2.2, 3.3, -4.4, 5.5, -6.6, 7.7, -8.8, 9.9, -10.1, 11.11, -12.12, - 13.13, -14.14, 15.15, -16.16, 17.17, -18.18, 19.19, -20.2, 21.21, -22.22, - 23.23, -24.24, 25.25, -26.26, 27.27, -28.28, 0, - /* batch 3 */ - -1.1, 2.2, -3.3, 4.4, -5.5, 6.6, -7.7, 8.8, -9.9, 10.1, -11.11, 12.12, - -13.13, 14.14, -15.15, 16.16, -17.17, 18.18, -19.19, 20.2, -21.21, 22.22, - -23.23, 24.24, -25.25, 26.26, -27.27, 28.28, 0}; - VectorBatchVectorCwiseProduct(input, kVectorSize, output.data(), kBatchSize, - output.data()); - - // Expect output = input * output + output. - const std::vector expected_output = { - /* batch 0 */ - 1.210000, 4.840000, 10.889999, 19.360001, 30.250000, 43.559998, 59.289997, - 77.440002, 98.009995, 102.010010, 123.432091, 146.894394, 172.396896, - 199.939606, 229.522491, 261.145599, 294.808899, 330.512421, 368.256134, - 408.040039, 449.864075, 493.728363, 539.632874, 587.577576, 637.562500, - 689.587585, 743.652954, 799.758423, 0.000000, - /* batch 1 */ - -1.210000, -4.840000, -10.889999, -19.360001, -30.250000, -43.559998, - -59.289997, -77.440002, -98.009995, -102.010010, -123.432091, -146.894394, - -172.396896, -199.939606, -229.522491, -261.145599, -294.808899, - -330.512421, -368.256134, -408.040039, -449.864075, -493.728363, - -539.632874, -587.577576, -637.562500, -689.587585, -743.652954, - -799.758423, 0.000000, - /* batch 2 */ - 1.210000, -4.840000, 10.889999, -19.360001, 30.250000, -43.559998, - 59.289997, -77.440002, 98.009995, -102.010010, 123.432091, -146.894394, - 172.396896, -199.939606, 229.522491, -261.145599, 294.808899, -330.512421, - 368.256134, -408.040039, 449.864075, -493.728363, 539.632874, -587.577576, - 637.562500, -689.587585, 743.652954, -799.758423, 0.000000, - /* batch 3 */ - -1.210000, 4.840000, -10.889999, 19.360001, -30.250000, 43.559998, - -59.289997, 77.440002, -98.009995, 102.010010, -123.432091, 146.894394, - -172.396896, 199.939606, -229.522491, 261.145599, -294.808899, 330.512421, - -368.256134, 408.040039, -449.864075, 493.728363, -539.632874, 587.577576, - -637.562500, 689.587585, -743.652954, 799.758423, 0.000000}; - EXPECT_THAT(output, testing::ElementsAreArray(expected_output)); -} - -TEST(uKernels, BatchVectorBatchVectorDotProductTest) { - constexpr int kVectorSize = 5; - constexpr int kBatch = 2; - static float input1[kVectorSize * kBatch] = {0.0, -0.5, 1.0, -1.5, 2.0, - -2.5, 3.0, -3.5, 4.0, -4.5}; - static float input2[kVectorSize * kBatch] = {0.1, -0.1, 0.1, -0.1, 0.1, - -0.1, 0.1, -0.1, 0.1, -0.1}; - std::vector output(kBatch); - BatchVectorBatchVectorDotProduct(input1, input2, kVectorSize, kBatch, - output.data(), /*result_stride=*/1); - EXPECT_THAT(output, ElementsAreArray(ArrayFloatNear({0.5, 1.75}))); -} - -TEST(uKernels, VectorShiftLeftTest) { - constexpr int kVectorSize = 5; - static float input[kVectorSize] = {0.0, -0.5, 1.0, -1.5, 2.0}; - std::vector result(kVectorSize); - VectorShiftLeft(input, kVectorSize, 3.0); - result.assign(input, input + kVectorSize); - EXPECT_THAT(result, - ElementsAreArray(ArrayFloatNear({-0.5, 1.0, -1.5, 2.0, 3.0}))); -} - -TEST(uKernels, ReductionSumVectorTest) { - constexpr int kInputVectorSize = 10; - constexpr int kOutputVectorSize1 = 5; - constexpr int kReductionSize1 = 2; - static float input[kInputVectorSize] = {0.0, -0.5, 1.0, -1.5, 2.0, - 0.0, -0.5, 1.0, 1.0, 2.0}; - std::vector result1(kOutputVectorSize1); - ReductionSumVector(input, result1.data(), kOutputVectorSize1, - kReductionSize1); - EXPECT_THAT(result1, - ElementsAreArray(ArrayFloatNear({-0.5, -0.5, 2.0, 0.5, 3.0}))); - - constexpr int kOutputVectorSize2 = 2; - constexpr int kReductionSize2 = 5; - std::vector result2(kOutputVectorSize2); - ReductionSumVector(input, result2.data(), kOutputVectorSize2, - kReductionSize2); - EXPECT_THAT(result2, ElementsAreArray(ArrayFloatNear({1.0, 3.5}))); -} - -TEST(uKernels, MeanStddevNormalizationNoneZeroInput) { - constexpr int kVectorSize = 4; - constexpr int kBatchSize = 2; - constexpr float kNormalizationEpsilon = 1e-8; - - // None-zero input. - static float input[kVectorSize * kBatchSize] = { - 0.1, 0.2, 0.3, 0.4, // batch 0 - 0.9, 1.0, 1.1, 1.2, // batch 1 - }; - std::vector output(kVectorSize * kBatchSize); - MeanStddevNormalization(input, output.data(), kVectorSize, kBatchSize, - kNormalizationEpsilon); - const std::vector expected_output = { - -1.34164071, -0.447213531, 0.44721365, 1.34164071, // batch 0 - -1.34163153, -0.447210163, 0.447211236, 1.3416326, // batch 1 - }; - EXPECT_THAT(output, testing::ElementsAreArray(expected_output)); -} - -TEST(uKernels, MeanStddevNormalizationAllZeroInput) { - constexpr int kVectorSize = 4; - constexpr int kBatchSize = 2; - constexpr float kNormalizationEpsilon = 1e-8; - - // Zero input. - static float input[kVectorSize * kBatchSize] = { - 0.0, 0.0, 0.0, 0.0, // batch 0 - 0.0, 0.0, 0.0, 0.0, // batch 1 - }; - std::vector output(kVectorSize * kBatchSize); - MeanStddevNormalization(input, output.data(), kVectorSize, kBatchSize, - kNormalizationEpsilon); - const std::vector expected_output = { - 0.0, 0.0, 0.0, 0.0, // batch 0 - 0.0, 0.0, 0.0, 0.0, // batch 1 - }; - EXPECT_THAT(output, testing::ElementsAreArray(expected_output)); -} - -TEST(uKernels, MeanStddevNormalizationMixed) { - constexpr int kVectorSize = 4; - constexpr int kBatchSize = 2; - constexpr float kNormalizationEpsilon = 1e-8; - - // Mix of zero and non-zero input. - static float input[kVectorSize * kBatchSize] = { - 0.0, 0.0, 0.0, 0.0, // batch 0 - 0.1, 0.2, 0.3, 0.4, // batch 1 - }; - std::vector output(kVectorSize * kBatchSize); - MeanStddevNormalization(input, output.data(), kVectorSize, kBatchSize, - kNormalizationEpsilon); - const std::vector expected_output = { - 0.0, 0.0, 0.0, 0.0, // batch 0 - -1.34164071, -0.447213531, 0.44721365, 1.34164071, // batch 1 - }; - EXPECT_THAT(output, testing::ElementsAreArray(expected_output)); -} - -TEST(uKernels, MeanStddevNormalizationSmallValue) { - constexpr int kVectorSize = 4; - constexpr int kBatchSize = 2; - constexpr float kNormalizationEpsilon = 1e-8; - - // Mix of zero and non-zero input. - static float input[kVectorSize * kBatchSize] = { - 3e-5, -7e-6, -9e-5, 1e-6, // batch 0 - 4e-5, 9e-6, 2e-4, 0.0, // batch 1 - }; - std::vector output(kVectorSize * kBatchSize); - MeanStddevNormalization(input, output.data(), kVectorSize, kBatchSize, - kNormalizationEpsilon); - const std::vector expected_output = { - 1.04231524, 0.212946132, -1.64753067, 0.392269224, // batch 0 - -0.275023013, -0.658201098, 1.70267045, -0.769446373, // batch 1 - }; - EXPECT_THAT(output, testing::ElementsAreArray(expected_output)); -} - -} // namespace tensor_utils -} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/internal/test_util.cc b/tensorflow/contrib/lite/kernels/internal/test_util.cc deleted file mode 100644 index 75d568ae3aaf9b186ffda0a1415f75ffb3e8c46b..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/internal/test_util.cc +++ /dev/null @@ -1,107 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include "tensorflow/contrib/lite/kernels/internal/test_util.h" - -#include -#include - -namespace tflite { - -// this is a copied from an internal function in propagate_fixed_sizes.cc -bool ComputeConvSizes(const RuntimeShape& input_shape, int output_depth, - int filter_width, int filter_height, int stride, - int dilation_width_factor, int dilation_height_factor, - PaddingType padding_type, RuntimeShape* output_shape, - int* pad_width, int* pad_height) { - const int input_width = input_shape.Dims(2); - const int input_height = input_shape.Dims(1); - const int batch = input_shape.Dims(0); - - int dilated_filter_width = dilation_width_factor * (filter_width - 1) + 1; - int dilated_filter_height = dilation_height_factor * (filter_height - 1) + 1; - - int output_height = 0; - int output_width = 0; - if (padding_type == PaddingType::kValid) { - output_height = (input_height + stride - dilated_filter_height) / stride; - output_width = (input_width + stride - dilated_filter_width) / stride; - } else if (padding_type == PaddingType::kSame) { - output_height = (input_height + stride - 1) / stride; - output_width = (input_width + stride - 1) / stride; - } else { - return false; - } - - if (output_width <= 0 || output_height <= 0) { - return false; - } - - *pad_height = std::max( - 0, ((output_height - 1) * stride + dilated_filter_height - input_height) / - 2); - *pad_width = std::max( - 0, - ((output_width - 1) * stride + dilated_filter_width - input_width) / 2); - - output_shape->BuildFrom({batch, output_height, output_width, output_depth}); - return true; -} - -std::mt19937& RandomEngine() { - static std::mt19937 engine; - return engine; -} - -int UniformRandomInt(int min, int max) { - std::uniform_int_distribution dist(min, max); - return dist(RandomEngine()); -} - -float UniformRandomFloat(float min, float max) { - std::uniform_real_distribution dist(min, max); - return dist(RandomEngine()); -} - -int ExponentialRandomPositiveInt(float percentile, int percentile_val, - int max_val) { - const float lambda = - -std::log(1.f - percentile) / static_cast(percentile_val); - std::exponential_distribution dist(lambda); - float val; - do { - val = dist(RandomEngine()); - } while (!val || !std::isfinite(val) || val > max_val); - return static_cast(std::ceil(val)); -} - -float ExponentialRandomPositiveFloat(float percentile, float percentile_val, - float max_val) { - const float lambda = - -std::log(1.f - percentile) / static_cast(percentile_val); - std::exponential_distribution dist(lambda); - float val; - do { - val = dist(RandomEngine()); - } while (!std::isfinite(val) || val > max_val); - return val; -} - -void FillRandom(std::vector* vec, float min, float max) { - std::uniform_real_distribution dist(min, max); - auto gen = std::bind(dist, RandomEngine()); - std::generate(std::begin(*vec), std::end(*vec), gen); -} - -} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/internal/test_util.h b/tensorflow/contrib/lite/kernels/internal/test_util.h deleted file mode 100644 index e4a383bedfc034a5398a3a1a78082e49c38f4afe..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/internal/test_util.h +++ /dev/null @@ -1,103 +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_LITE_KERNELS_INTERNAL_TEST_UTIL_H_ -#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_TEST_UTIL_H_ - -#include -#include -#include -#include -#include -#include - -#include "tensorflow/contrib/lite/kernels/internal/types.h" - -namespace tflite { - -// Computes output and padding dimensions. -bool ComputeConvSizes(const RuntimeShape& input_shape, int output_depth, - int filter_width, int filter_height, int stride, - int dilation_width_factor, int dilation_height_factor, - PaddingType padding_type, RuntimeShape* output_shape, - int* pad_width, int* pad_height); - -// Returns a mt19937 random engine. -std::mt19937& RandomEngine(); - -// Returns a random integer uniformly distributed between |min| and |max|. -int UniformRandomInt(int min, int max); - -// Returns a random float uniformly distributed between |min| and |max|. -float UniformRandomFloat(float min, float max); - -// Returns a random element in |v|. -template -const T& RandomElement(const std::vector& v) { - return v[UniformRandomInt(0, v.size() - 1)]; -} - -// Returns a random exponentially distributed integer. -int ExponentialRandomPositiveInt(float percentile, int percentile_val, - int max_val); - -// Returns a random exponentially distributed float. -float ExponentialRandomPositiveFloat(float percentile, float percentile_val, - float max_val); - -// Fills a vector with random floats between |min| and |max|. -void FillRandom(std::vector* vec, float min, float max); - -// Fills a vector with random numbers between |min| and |max|. -template -void FillRandom(std::vector* vec, T min, T max) { - std::uniform_int_distribution dist(min, max); - auto gen = std::bind(dist, RandomEngine()); - std::generate(std::begin(*vec), std::end(*vec), gen); -} - -// Fills a vector with random numbers. -template -void FillRandom(std::vector* vec) { - FillRandom(vec, std::numeric_limits::min(), std::numeric_limits::max()); -} - -template -void FillRandom(typename std::vector::iterator begin_it, - typename std::vector::iterator end_it, T min, T max) { - std::uniform_int_distribution dist(min, max); - auto gen = std::bind(dist, RandomEngine()); - std::generate(begin_it, end_it, gen); -} - -// Fill with a "skyscraper" pattern, in which there is a central section (across -// the depth) with higher values than the surround. -template -void FillRandomSkyscraper(std::vector* vec, int depth, - double middle_proportion, uint8 middle_min, - uint8 sides_max) { - for (auto base_it = std::begin(*vec); base_it != std::end(*vec); - base_it += depth) { - auto left_it = base_it + std::ceil(0.5 * depth * (1.0 - middle_proportion)); - auto right_it = - base_it + std::ceil(0.5 * depth * (1.0 + middle_proportion)); - FillRandom(base_it, left_it, std::numeric_limits::min(), sides_max); - FillRandom(left_it, right_it, middle_min, std::numeric_limits::max()); - FillRandom(right_it, base_it + depth, std::numeric_limits::min(), - sides_max); - } -} - -} // namespace tflite -#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_TEST_UTIL_H_ diff --git a/tensorflow/contrib/lite/kernels/internal/types.h b/tensorflow/contrib/lite/kernels/internal/types.h deleted file mode 100644 index a5913143b9ab834d13a3aef3056f68c9e4ad6b76..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/internal/types.h +++ /dev/null @@ -1,1023 +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_LITE_KERNELS_INTERNAL_TYPES_H_ -#define TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_TYPES_H_ - -#include -#include - -#include "tensorflow/contrib/lite/kernels/internal/compatibility.h" - -namespace tflite { - -enum class FusedActivationFunctionType : uint8 { kNone, kRelu6, kRelu1, kRelu }; -enum class PaddingType : uint8 { kNone, kSame, kValid }; - -struct PaddingValues { - int16 width; - int16 height; -}; - -// This enumeration allows for non-default formats for the weights array -// of a fully-connected operator, allowing the use of special optimized -// runtime paths. -enum class FullyConnectedWeightsFormat : uint8 { - // Default format (flat 2D layout, the inner contiguous dimension - // is input_depth, the outer non-contiguous dimension is output_depth) - kDefault, - // Summary: optimized layout for fast CPU runtime implementation, - // aimed specifically at ARM CPUs at the moment, and specialized for - // 8-bit quantized layers. - // - // The use case we're concerned with here is: 8-bit quantization, - // large weights matrix that doesn't fit in cache (e.g. 4096x2048 in - // a key application that drove this), very small batch size (e.g. 1 -- 4). - // - // Even with 8-bit quantization of weights, the performance of memory - // accesses to the weights can become the dominant issue when - // the batch size is small, so each weight value is used in only a few - // arithmetic ops, i.e. the fully-connected node has a low arithmetic - // intensity. The specific issues that arise are of three kinds: - // (1) One may, ideally, max out DRAM bandwidth, i.e. be truly memory - // bound. That's the "good" issue to run into. - // (2) One may run into sub-optimal pre-fetching: the data hasn't been - // prefetched into the cache by the time we need it. - // (3) One may run into cache aliasing: multiple values that are - // pre-fetched, alias each other in the L1 cache (which typically - // has only 4-way set associativity in ARM CPUs) and thus evict - // each other before we get to using them. - // - // The point of this shuffling is to avoid issues (2) and (3) so that - // we get as fast as possible given only the hard constraint (1). - // This is achieved by turning the difficulty into a solution: the - // difficulty, that each value loaded from memory is used only in - // one kernel iteration, making this operation memory-intensive, hints at - // the solution, of shuffling the weights so that they are stored in the - // exact order as the kernel needs to load them, so that the memory - // accesses made by the kernel are trivial. This solves (2) because the - // trivial memory access pattern allows the CPU's automatic prefetching - // to perform very well (no need even for preload instructions), and this - // solves (3) because the values being loaded concurrently are now - // contiguous in the address space, thus don't alias each other in the cache. - // - // On ARM, we typically want our kernel to process a 4x16 block of weights - // at a time, because: - // - 16 is the number of bytes in a NEON register. - // - 4 is how many rows we need to handle concurrently in the kernel in - // order to have sufficient mutual independence of instructions to - // maximize arithmetic throughput. - // - // Finally, the 'Int8' part in the name refers to the fact that this - // weights format has each weights value encoded as a signed int8 value, - // even if the data type of the weights buffer is uint8. This is intended - // to save runtime kernels the effort to have to XOR the top bit of these - // bytes before using them in signed arithmetic, see this file for more - // explanations on the 'signed int8 trick' in matrix multiplication kernels: - // - // tensorflow/contrib/lite/toco/graph_transformations/ensure_uint8_weights_safe_for_fast_int8_kernels.cc - // - kShuffled4x16Int8, -}; - -// Quantization parameters, determining the mapping of quantized values -// to real values (i.e. determining how quantized values are mathematically -// interpreted). -// -// The correspondence is as follows: -// -// real_value = scale * (quantized_value - zero_point); -// -// In other words, zero_point designates which quantized value corresponds to -// the real 0 value, and scale designates the difference between the real values -// corresponding to consecutive quantized values differing by 1. -struct QuantizationParams { - int32 zero_point = 0; - double scale = 0.0; -}; - -inline bool operator==(const QuantizationParams& qp1, - const QuantizationParams& qp2) { - return qp1.zero_point == qp2.zero_point && qp1.scale == qp2.scale; -} - -template -struct Dims { - int sizes[N]; - int strides[N]; -}; - -class RuntimeShape { - public: - // Shapes with dimensions up to 4 are stored directly in the structure, while - // larger shapes are separately allocated. - static constexpr int kMaxSmallSize = 4; - - RuntimeShape& operator=(RuntimeShape const&) = delete; - - RuntimeShape() : size_(0) {} - - explicit RuntimeShape(int dimensions_count) : size_(dimensions_count) { - if (dimensions_count > kMaxSmallSize) { -#ifdef TF_LITE_STATIC_MEMORY - TFLITE_CHECK(false && "No shape resizing supported on this platform"); -#else // TF_LITE_STATIC_MEMORY - dims_pointer_ = new int32[dimensions_count]; -#endif // TF_LITE_STATIC_MEMORY - } - } - - RuntimeShape(int shape_size, int32 value) : size_(0) { - Resize(shape_size); - for (int i = 0; i < shape_size; ++i) { - SetDim(i, value); - } - } - - RuntimeShape(int dimensions_count, const int32* dims_data) : size_(0) { - ReplaceWith(dimensions_count, dims_data); - } - - RuntimeShape(const std::initializer_list init_list) : size_(0) { - BuildFrom(init_list); - } - - // Avoid using this constructor. We should be able to delete it when C++17 - // rolls out. - RuntimeShape(RuntimeShape const& other) : size_(other.DimensionsCount()) { - if (size_ > kMaxSmallSize) { - dims_pointer_ = new int32[size_]; - } - std::memcpy(DimsData(), other.DimsData(), sizeof(int32) * size_); - } - - bool operator==(const RuntimeShape& comp) const { - return this->size_ == comp.size_ && - std::memcmp(DimsData(), comp.DimsData(), size_ * sizeof(int32)) == 0; - } - - ~RuntimeShape() { - if (size_ > kMaxSmallSize) { -#ifdef TF_LITE_STATIC_MEMORY - TFLITE_CHECK(false && "No shape resizing supported on this platform"); -#else // TF_LITE_STATIC_MEMORY - delete[] dims_pointer_; -#endif // TF_LITE_STATIC_MEMORY - } - } - - inline int32 DimensionsCount() const { return size_; } - inline int32 Dims(int i) const { - TFLITE_DCHECK_GE(i, 0); - TFLITE_DCHECK_LT(i, size_); - return size_ > kMaxSmallSize ? dims_pointer_[i] : dims_[i]; - } - inline void SetDim(int i, int32 val) { - TFLITE_DCHECK_GE(i, 0); - TFLITE_DCHECK_LT(i, size_); - if (size_ > kMaxSmallSize) { - dims_pointer_[i] = val; - } else { - dims_[i] = val; - } - } - - inline int32* DimsData() { - return size_ > kMaxSmallSize ? dims_pointer_ : dims_; - } - inline const int32* DimsData() const { - return size_ > kMaxSmallSize ? dims_pointer_ : dims_; - } - // The caller must ensure that the shape is no bigger than 4-D. - inline const int32* DimsDataUpTo4D() const { return dims_; } - - inline void Resize(int dimensions_count) { - if (size_ > kMaxSmallSize) { -#ifdef TF_LITE_STATIC_MEMORY - TFLITE_CHECK(false && "No shape resizing supported on this platform"); -#else // TF_LITE_STATIC_MEMORY - delete[] dims_pointer_; -#endif // TF_LITE_STATIC_MEMORY - } - size_ = dimensions_count; - if (dimensions_count > kMaxSmallSize) { -#ifdef TF_LITE_STATIC_MEMORY - TFLITE_CHECK(false && "No shape resizing supported on this platform"); -#else // TF_LITE_STATIC_MEMORY - dims_pointer_ = new int32[dimensions_count]; -#endif // TF_LITE_STATIC_MEMORY - } - } - - inline void ReplaceWith(int dimensions_count, const int32* dims_data) { - Resize(dimensions_count); - int32* dst_dims = DimsData(); - std::memcpy(dst_dims, dims_data, dimensions_count * sizeof(int32)); - } - - template - inline void BuildFrom(const T& src_iterable) { - const int dimensions_count = - std::distance(src_iterable.begin(), src_iterable.end()); - Resize(dimensions_count); - int32* data = DimsData(); - for (auto it : src_iterable) { - *data = it; - ++data; - } - } - - // This will probably be factored out. Old code made substantial use of 4-D - // shapes, and so this function is used to extend smaller shapes. Note that - // (a) as Dims<4>-dependent code is eliminated, the reliance on this should be - // reduced, and (b) some kernels are stricly 4-D, but then the shapes of their - // inputs should already be 4-D, so this function should not be needed. - inline static RuntimeShape ExtendedShape(int new_shape_size, - const RuntimeShape& shape) { - return RuntimeShape(new_shape_size, shape, 1); - } - - inline void BuildFrom(const std::initializer_list init_list) { - BuildFrom>(init_list); - } - - // Returns the total count of elements, that is the size when flattened into a - // vector. - inline int FlatSize() const { - int buffer_size = 1; - const int* dims_data = DimsData(); - for (int i = 0; i < size_; i++) { - const int dim = dims_data[i]; - TFLITE_DCHECK_GE(dim, 1); - buffer_size *= dim; - } - return buffer_size; - } - - bool operator!=(const RuntimeShape& comp) const { return !((*this) == comp); } - - private: - // For use only by ExtendedShape(), written to guarantee (return-value) copy - // elision in C++17. - // This creates a shape padded to the desired size with the specified value. - RuntimeShape(int new_shape_size, const RuntimeShape& shape, int pad_value) - : size_(0) { - // If the following check fails, it is likely because a 4D-only kernel is - // being used with an array of larger dimension count. - TFLITE_CHECK_GE(new_shape_size, shape.DimensionsCount()); - Resize(new_shape_size); - const int size_increase = new_shape_size - shape.DimensionsCount(); - for (int i = 0; i < size_increase; ++i) { - SetDim(i, pad_value); - } - std::memcpy(DimsData() + size_increase, shape.DimsData(), - sizeof(int32) * shape.DimensionsCount()); - } - - int32 size_; - union { - int32 dims_[kMaxSmallSize]; - int32* dims_pointer_; - }; -}; - -// Converts inference-style shape to legacy tflite::Dims<4>. -inline tflite::Dims<4> ToRuntimeDims(const tflite::RuntimeShape& array_shape) { - tflite::Dims<4> result; - const int dimensions_count = array_shape.DimensionsCount(); - TFLITE_CHECK_LE(dimensions_count, 4); - int cum_prod = 1; - for (int i = 0; i < 4; i++) { - const int new_dim = - (i < dimensions_count) ? array_shape.Dims(dimensions_count - 1 - i) : 1; - result.sizes[i] = new_dim; - result.strides[i] = cum_prod; - cum_prod *= new_dim; - } - return result; -} - -// TODO(b/80418076): Move to legacy ops file, update invocations. -inline RuntimeShape DimsToShape(const tflite::Dims<4>& dims) { - return RuntimeShape( - {dims.sizes[3], dims.sizes[2], dims.sizes[1], dims.sizes[0]}); -} - -// Gets next index to iterate through a multidimensional array. -inline bool NextIndex(const int num_dims, const int* dims, int* current) { - if (num_dims == 0) { - return false; - } - TFLITE_DCHECK(dims != nullptr); - TFLITE_DCHECK(current != nullptr); - int carry = 1; - for (int idx = num_dims - 1; idx >= 0; --idx) { - int current_val = current[idx] + carry; - TFLITE_DCHECK_GE(dims[idx], current_val); - if (dims[idx] == current_val) { - current[idx] = 0; - } else { - current[idx] = current_val; - carry = 0; - break; - } - } - return (carry == 0); -} - -// Gets offset of index if reducing on axis. When reducing, the flattened offset -// will not change, if the input index changes on the given axis. For example, -// if you have a 3D tensor and you are reducing to 2D by eliminating axis 0, -// then index (0, 1, 2) and index (1, 1, 2) will map to the same flattened -// offset. -// TODO(kanlig): uses Dims to represent dimensions. -inline size_t ReducedOutputOffset(const int num_dims, const int* dims, - const int* index, const int num_axis, - const int* axis) { - if (num_dims == 0) { - return 0; - } - TFLITE_DCHECK(dims != nullptr); - TFLITE_DCHECK(index != nullptr); - size_t offset = 0; - for (int idx = 0; idx < num_dims; ++idx) { - // if we need to skip this axis - bool is_axis = false; - if (axis != nullptr) { - for (int axis_idx = 0; axis_idx < num_axis; ++axis_idx) { - if (idx == axis[axis_idx]) { - is_axis = true; - break; - } - } - } - if (!is_axis) { - offset = offset * static_cast(dims[idx]) + - static_cast(index[idx]); - } - } - return offset; -} - -inline int Offset(const RuntimeShape& shape, int i0, int i1, int i2, int i3) { - TFLITE_DCHECK_EQ(shape.DimensionsCount(), 4); - const int* dims_data = shape.DimsDataUpTo4D(); - TFLITE_DCHECK(i0 >= 0 && i0 < dims_data[0]); - TFLITE_DCHECK(i1 >= 0 && i1 < dims_data[1]); - TFLITE_DCHECK(i2 >= 0 && i2 < dims_data[2]); - TFLITE_DCHECK(i3 >= 0 && i3 < dims_data[3]); - return ((i0 * dims_data[1] + i1) * dims_data[2] + i2) * dims_data[3] + i3; -} - -inline int Offset(const Dims<4>& dims, int i0, int i1, int i2, int i3) { - TFLITE_DCHECK(i0 >= 0 && i0 < dims.sizes[0]); - TFLITE_DCHECK(i1 >= 0 && i1 < dims.sizes[1]); - TFLITE_DCHECK(i2 >= 0 && i2 < dims.sizes[2]); - TFLITE_DCHECK(i3 >= 0 && i3 < dims.sizes[3]); - return i0 * dims.strides[0] + i1 * dims.strides[1] + i2 * dims.strides[2] + - i3 * dims.strides[3]; -} - -inline int Offset(const Dims<4>& dims, int* index) { - return Offset(dims, index[0], index[1], index[2], index[3]); -} - -inline int Offset(const RuntimeShape& shape, int* index) { - return Offset(shape, index[0], index[1], index[2], index[3]); -} - -// Get array size, DCHECKing that the dim index is in range. -// -// Note that this will be phased out with Dims<4>, since RuntimeShape::Dims() -// already performs this check. -template -int ArraySize(const Dims& array, int index) { - TFLITE_DCHECK(index >= 0 && index < N); - return array.sizes[index]; -} - -// Get common array size, DCHECKing that they all agree. -template -int MatchingArraySize(const ArrayType1& array1, int index1, - const ArrayType2& array2, int index2) { - TFLITE_DCHECK_EQ(ArraySize(array1, index1), ArraySize(array2, index2)); - return ArraySize(array1, index1); -} - -template -int MatchingArraySize(const ArrayType1& array1, int index1, - const ArrayType2& array2, int index2, Args... args) { - TFLITE_DCHECK_EQ(ArraySize(array1, index1), ArraySize(array2, index2)); - return MatchingArraySize(array1, index1, args...); -} - -// Get common shape dim, DCHECKing that they all agree. -inline int MatchingDim(const RuntimeShape& shape1, int index1, - const RuntimeShape& shape2, int index2) { - TFLITE_DCHECK_EQ(shape1.Dims(index1), shape2.Dims(index2)); - return shape1.Dims(index1); -} - -template -int MatchingDim(const RuntimeShape& shape1, int index1, - const RuntimeShape& shape2, int index2, Args... args) { - TFLITE_DCHECK_EQ(shape1.Dims(index1), shape2.Dims(index2)); - return MatchingDim(shape1, index1, args...); -} - -// Will be phased out with Dims<4>, replaced by RuntimeShape::FlatSize(). -template -inline int FlatSize(const Dims& dims) { - int flat_size = 1; - for (int i = 0; i < N; ++i) { - flat_size *= dims.sizes[i]; - } - return flat_size; -} - -TFLITE_DEPRECATED("Prefer FlatSize.") -inline int RequiredBufferSizeForDims(const Dims<4>& dims) { - return FlatSize(dims); -} - -// Flat size calculation, checking that dimensions match with one or more other -// arrays. -inline int MatchingFlatSize(const RuntimeShape& shape, - const RuntimeShape& check_shape_0) { - TFLITE_DCHECK_EQ(shape.DimensionsCount(), check_shape_0.DimensionsCount()); - const int dims_count = shape.DimensionsCount(); - for (int i = 0; i < dims_count; ++i) { - TFLITE_DCHECK_EQ(shape.Dims(i), check_shape_0.Dims(i)); - } - return shape.FlatSize(); -} - -inline int MatchingFlatSize(const RuntimeShape& shape, - const RuntimeShape& check_shape_0, - const RuntimeShape& check_shape_1) { - TFLITE_DCHECK_EQ(shape.DimensionsCount(), check_shape_0.DimensionsCount()); - const int dims_count = shape.DimensionsCount(); - for (int i = 0; i < dims_count; ++i) { - TFLITE_DCHECK_EQ(shape.Dims(i), check_shape_0.Dims(i)); - } - return MatchingFlatSize(shape, check_shape_1); -} - -inline int MatchingFlatSize(const RuntimeShape& shape, - const RuntimeShape& check_shape_0, - const RuntimeShape& check_shape_1, - const RuntimeShape& check_shape_2) { - TFLITE_DCHECK_EQ(shape.DimensionsCount(), check_shape_0.DimensionsCount()); - const int dims_count = shape.DimensionsCount(); - for (int i = 0; i < dims_count; ++i) { - TFLITE_DCHECK_EQ(shape.Dims(i), check_shape_0.Dims(i)); - } - return MatchingFlatSize(shape, check_shape_1, check_shape_2); -} - -inline int MatchingFlatSize(const RuntimeShape& shape, - const RuntimeShape& check_shape_0, - const RuntimeShape& check_shape_1, - const RuntimeShape& check_shape_2, - const RuntimeShape& check_shape_3) { - TFLITE_DCHECK_EQ(shape.DimensionsCount(), check_shape_0.DimensionsCount()); - const int dims_count = shape.DimensionsCount(); - for (int i = 0; i < dims_count; ++i) { - TFLITE_DCHECK_EQ(shape.Dims(i), check_shape_0.Dims(i)); - } - return MatchingFlatSize(shape, check_shape_1, check_shape_2, check_shape_3); -} - -// Flat size calculation, checking that dimensions match with one or more other -// arrays. -template -inline int MatchingFlatSize(const Dims& dims, const Dims& check_dims_0) { - for (int i = 0; i < N; ++i) { - TFLITE_DCHECK_EQ(ArraySize(dims, i), ArraySize(check_dims_0, i)); - } - return FlatSize(dims); -} - -template -inline int MatchingFlatSize(const Dims& dims, const Dims& check_dims_0, - const Dims& check_dims_1) { - for (int i = 0; i < N; ++i) { - TFLITE_DCHECK_EQ(ArraySize(dims, i), ArraySize(check_dims_0, i)); - } - return MatchingFlatSize(dims, check_dims_1); -} - -template -inline int MatchingFlatSize(const Dims& dims, const Dims& check_dims_0, - const Dims& check_dims_1, - const Dims& check_dims_2) { - for (int i = 0; i < N; ++i) { - TFLITE_DCHECK_EQ(ArraySize(dims, i), ArraySize(check_dims_0, i)); - } - return MatchingFlatSize(dims, check_dims_1, check_dims_2); -} - -template -inline int MatchingFlatSize(const Dims& dims, const Dims& check_dims_0, - const Dims& check_dims_1, - const Dims& check_dims_2, - const Dims& check_dims_3) { - for (int i = 0; i < N; ++i) { - TFLITE_DCHECK_EQ(ArraySize(dims, i), ArraySize(check_dims_0, i)); - } - return MatchingFlatSize(dims, check_dims_1, check_dims_2, check_dims_3); -} - -// Data is required to be contiguous, and so many operators can use either the -// full array flat size or the flat size with one dimension skipped (commonly -// the depth). -template -inline int FlatSizeSkipDim(const Dims& dims, int skip_dim) { - TFLITE_DCHECK(skip_dim >= 0 && skip_dim < N); - int flat_size = 1; - for (int i = 0; i < N; ++i) { - flat_size *= (i == skip_dim) ? 1 : dims.sizes[i]; - } - return flat_size; -} - -// A combination of MatchingFlatSize() and FlatSizeSkipDim(). -template -inline int MatchingFlatSizeSkipDim(const Dims& dims, int skip_dim, - const Dims& check_dims_0) { - for (int i = 0; i < N; ++i) { - if (i != skip_dim) { - TFLITE_DCHECK_EQ(ArraySize(dims, i), ArraySize(check_dims_0, i)); - } - } - return FlatSizeSkipDim(dims, skip_dim); -} - -template -inline int MatchingFlatSizeSkipDim(const Dims& dims, int skip_dim, - const Dims& check_dims_0, - const Dims& check_dims_1) { - for (int i = 0; i < N; ++i) { - if (i != skip_dim) { - TFLITE_DCHECK_EQ(ArraySize(dims, i), ArraySize(check_dims_0, i)); - } - } - return MatchingFlatSizeSkipDim(dims, skip_dim, check_dims_1); -} - -template -inline int MatchingFlatSizeSkipDim(const Dims& dims, int skip_dim, - const Dims& check_dims_0, - const Dims& check_dims_1, - const Dims& check_dims_2) { - for (int i = 0; i < N; ++i) { - if (i != skip_dim) { - TFLITE_DCHECK_EQ(ArraySize(dims, i), ArraySize(check_dims_0, i)); - } - } - return MatchingFlatSizeSkipDim(dims, skip_dim, check_dims_1, check_dims_2); -} - -template -inline int MatchingFlatSizeSkipDim(const Dims& dims, int skip_dim, - const Dims& check_dims_0, - const Dims& check_dims_1, - const Dims& check_dims_2, - const Dims& check_dims_3) { - for (int i = 0; i < N; ++i) { - if (i != skip_dim) { - TFLITE_DCHECK_EQ(ArraySize(dims, i), ArraySize(check_dims_0, i)); - } - } - return MatchingFlatSizeSkipDim(dims, skip_dim, check_dims_1, check_dims_2, - check_dims_3); -} - -// Data is required to be contiguous, and so many operators can use either the -// full array flat size or the flat size with one dimension skipped (commonly -// the depth). -inline int FlatSizeSkipDim(const RuntimeShape& shape, int skip_dim) { - const int dims_count = shape.DimensionsCount(); - TFLITE_DCHECK(skip_dim >= 0 && skip_dim < dims_count); - const auto* dims_data = shape.DimsData(); - int flat_size = 1; - for (int i = 0; i < dims_count; ++i) { - flat_size *= (i == skip_dim) ? 1 : dims_data[i]; - } - return flat_size; -} - -// A combination of MatchingFlatSize() and FlatSizeSkipDim(). -inline int MatchingFlatSizeSkipDim(const RuntimeShape& shape, int skip_dim, - const RuntimeShape& check_shape_0) { - const int dims_count = shape.DimensionsCount(); - for (int i = 0; i < dims_count; ++i) { - if (i != skip_dim) { - TFLITE_DCHECK_EQ(shape.Dims(i), check_shape_0.Dims(i)); - } - } - return FlatSizeSkipDim(shape, skip_dim); -} - -inline int MatchingFlatSizeSkipDim(const RuntimeShape& shape, int skip_dim, - const RuntimeShape& check_shape_0, - const RuntimeShape& check_shape_1) { - const int dims_count = shape.DimensionsCount(); - for (int i = 0; i < dims_count; ++i) { - if (i != skip_dim) { - TFLITE_DCHECK_EQ(shape.Dims(i), check_shape_0.Dims(i)); - } - } - return MatchingFlatSizeSkipDim(shape, skip_dim, check_shape_1); -} - -inline int MatchingFlatSizeSkipDim(const RuntimeShape& shape, int skip_dim, - const RuntimeShape& check_shape_0, - const RuntimeShape& check_shape_1, - const RuntimeShape& check_shape_2) { - const int dims_count = shape.DimensionsCount(); - for (int i = 0; i < dims_count; ++i) { - if (i != skip_dim) { - TFLITE_DCHECK_EQ(shape.Dims(i), check_shape_0.Dims(i)); - } - } - return MatchingFlatSizeSkipDim(shape, skip_dim, check_shape_1, check_shape_2); -} - -inline int MatchingFlatSizeSkipDim(const RuntimeShape& shape, int skip_dim, - const RuntimeShape& check_shape_0, - const RuntimeShape& check_shape_1, - const RuntimeShape& check_shape_2, - const RuntimeShape& check_shape_3) { - const int dims_count = shape.DimensionsCount(); - for (int i = 0; i < dims_count; ++i) { - if (i != skip_dim) { - TFLITE_DCHECK_EQ(shape.Dims(i), check_shape_0.Dims(i)); - } - } - return MatchingFlatSizeSkipDim(shape, skip_dim, check_shape_1, check_shape_2, - check_shape_3); -} - -template -bool IsPackedWithoutStrides(const Dims& dims) { - int expected_stride = 1; - for (int d = 0; d < N; d++) { - if (dims.strides[d] != expected_stride) return false; - expected_stride *= dims.sizes[d]; - } - return true; -} - -template -void ComputeStrides(Dims* dims) { - dims->strides[0] = 1; - for (int d = 1; d < N; d++) { - dims->strides[d] = dims->strides[d - 1] * dims->sizes[d - 1]; - } -} - -enum class BroadcastableOpCategory : uint8 { - kNone, - kNonBroadcast, // Matching input shapes. - kFirstInputBroadcastsFast, // Fivefold nested loops. - kSecondInputBroadcastsFast, // Fivefold nested loops. - kGenericBroadcast, // Fall-back. -}; - -struct MinMax { - float min; - float max; -}; -static_assert(sizeof(MinMax) == 8, ""); - -struct ActivationParams { - FusedActivationFunctionType activation_type; - // uint8, etc, activation params. - int32 quantized_activation_min; - int32 quantized_activation_max; -}; - -// For Add, Sub, Mul ops. -struct ArithmeticParams { - // Shape dependent / common to data / op types. - BroadcastableOpCategory broadcast_category; - // uint8 inference params. - int32 input1_offset; - int32 input2_offset; - int32 output_offset; - int32 output_multiplier; - int output_shift; - // Add / Sub, not Mul, uint8 inference params. - int left_shift; - int32 input1_multiplier; - int input1_shift; - int32 input2_multiplier; - int input2_shift; - // uint8, etc, activation params. - int32 quantized_activation_min; - int32 quantized_activation_max; - // float activation params. - float float_activation_min; - float float_activation_max; - - // Processed output dimensions. - // Let input "a" be the one that broadcasts in the faster-changing dimension. - // Then, after coalescing, for shapes {a0, a1, a2, a3, a4} and - // {b0, b1, b2, b3, b4}, - // broadcast_shape[4] = b0 = a0. - // broadcast_shape[3] = b1; a1 = 1. - // broadcast_shape[2] = b2 = a2. - // broadcast_shape[1] = a3; b3 = 1. - // broadcast_shape[0] = b4 = a4. - int broadcast_shape[5]; -}; - -struct ConcatenationParams { - int8 axis; - const int32* input_zeropoint; - const float* input_scale; - uint16 inputs_count; - int32 output_zeropoint; - float output_scale; -}; - -struct ComparisonParams { - // uint8 inference params. - int left_shift; - int32 input1_offset; - int32 input1_multiplier; - int input1_shift; - int32 input2_offset; - int32 input2_multiplier; - int input2_shift; - // Shape dependent / common to inference types. - bool is_broadcast; -}; - -struct ConvParams { - PaddingType padding_type; - PaddingValues padding_values; - // TODO(starka): This was just "stride", so check that width+height is OK. - int16 stride_width; - int16 stride_height; - int16 dilation_width_factor; - int16 dilation_height_factor; - // uint8 inference params. - // TODO(b/65838351): Use smaller types if appropriate. - int32 input_offset; - int32 weights_offset; - int32 output_offset; - int32 output_multiplier; - int output_shift; - // uint8, etc, activation params. - int32 quantized_activation_min; - int32 quantized_activation_max; - // float activation params. - float float_activation_min; - float float_activation_max; -}; - -struct DepthToSpaceParams { - int32 block_size; -}; - -struct DepthwiseParams { - PaddingType padding_type; - PaddingValues padding_values; - int16 stride_width; - int16 stride_height; - int16 dilation_width_factor; - int16 dilation_height_factor; - int16 depth_multiplier; - // uint8 inference params. - // TODO(b/65838351): Use smaller types if appropriate. - int32 input_offset; - int32 weights_offset; - int32 output_offset; - int32 output_multiplier; - int output_shift; - // uint8, etc, activation params. - int32 quantized_activation_min; - int32 quantized_activation_max; - // float activation params. - float float_activation_min; - float float_activation_max; -}; - -struct DequantizationParams { - double scale; - int32 zero_point; -}; - -struct FakeQuantParams { - MinMax minmax; - int32 num_bits; -}; - -struct FullyConnectedParams { - // uint8 inference params. - // TODO(b/65838351): Use smaller types if appropriate. - int32 input_offset; - int32 weights_offset; - int32 output_offset; - int32 output_multiplier; - int output_shift; - // uint8, etc, activation params. - int32 quantized_activation_min; - int32 quantized_activation_max; - // float activation params. - float float_activation_min; - float float_activation_max; - FullyConnectedWeightsFormat weights_format; -}; - -struct GatherParams { - int16 input_rank; - int16 axis; -}; - -struct L2NormalizationParams { - // uint8 inference params. - int32 input_zero_point; -}; - -struct LocalResponseNormalizationParams { - int32 range; - double bias; - double alpha; - double beta; -}; - -struct LogisticParams { - // uint8 inference params. - int32 input_zero_point; - int32 input_range_radius; - int32 input_multiplier; - int input_left_shift; -}; - -struct LstmCellParams { - int32 weights_zero_point; - int32 accum_multiplier; - int accum_shift; - int state_integer_bits; -}; - -struct MeanParams { - int8 axis_count; - int16 axis[4]; -}; - -struct PackParams { - int8 axis; - const int32* input_zeropoint; - const float* input_scale; - uint16 inputs_count; - int32 output_zeropoint; - float output_scale; -}; - -struct PadParams { - int8 left_padding_count; - int32 left_padding[4]; - int8 right_padding_count; - int32 right_padding[4]; -}; - -struct PoolParams { - FusedActivationFunctionType activation; - PaddingType padding_type; - PaddingValues padding_values; - int stride_height; - int stride_width; - int filter_height; - int filter_width; - // uint8, etc, activation params. - int32 quantized_activation_min; - int32 quantized_activation_max; - // float activation params. - float float_activation_min; - float float_activation_max; -}; - -struct ReshapeParams { - int8 shape_count; - int32 shape[4]; -}; - -struct ResizeBilinearParams { - bool align_corners; -}; - -struct SliceParams { - int8 begin_count; - int32 begin[4]; - int8 size_count; - int32 size[4]; -}; - -struct SoftmaxParams { - // beta is not really used (not a Tensorflow parameter) and not implemented - // for LogSoftmax. - double beta; - // uint8 inference params. Used even when beta defaults to 1.0. - int32 input_multiplier; - int32 input_left_shift; - // Reverse scaling is only used by LogSoftmax. - int32 reverse_scaling_divisor; - int32 reverse_scaling_right_shift; - int diff_min; -}; - -struct SpaceToBatchParams { - // "Zero" padding for uint8 means padding with the output offset. - int32 output_offset; -}; - -struct SpaceToDepthParams { - int32 block_size; -}; - -struct SplitParams { - // Graphs that split into, say, 2000 nodes are encountered. The indices in - // OperatorEdges are of type uint16. - uint16 num_split; - int16 axis; -}; - -struct SqueezeParams { - int8 squeeze_dims_count; - int32 squeeze_dims[4]; -}; - -struct StridedSliceParams { - int8 start_indices_count; - int16 start_indices[4]; - int8 stop_indices_count; - int16 stop_indices[4]; - int8 strides_count; - int16 strides[4]; - - int16 begin_mask; - int16 ellipsis_mask; - int16 end_mask; - int16 new_axis_mask; - int16 shrink_axis_mask; -}; - -struct TanhParams { - int32 input_zero_point; - int32 input_range_radius; - int32 input_multiplier; - int input_left_shift; -}; - -struct TransposeParams { - int8 perm_count; - int32 perm[4]; -}; - -struct UnpackParams { - uint16 num_split; - int16 axis; -}; - -template -inline void SetActivationParams(float min, float max, P* params) { - params->float_activation_min = min; - params->float_activation_max = max; -} - -template -inline void SetActivationParams(int32 min, int32 max, P* params) { - params->quantized_activation_min = min; - params->quantized_activation_max = max; -} - -template -inline void GetActivationParams(const P& params, int32* min, int32* max) { - *min = params.quantized_activation_min; - *max = params.quantized_activation_max; -} - -template -inline void GetActivationParams(const P& params, float* min, float* max) { - *min = params.float_activation_min; - *max = params.float_activation_max; -} - -} // namespace tflite - -#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_INTERNAL_TYPES_H_ diff --git a/tensorflow/contrib/lite/kernels/layer_norm_lstm_test.cc b/tensorflow/contrib/lite/kernels/layer_norm_lstm_test.cc deleted file mode 100644 index 1535f750f94e725061265f76209fbaa213558346..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/layer_norm_lstm_test.cc +++ /dev/null @@ -1,662 +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. -==============================================================================*/ -// Unit test for TFLite Layer Norm LSTM op. - -#include -#include - -#include -#include -#include "flatbuffers/flexbuffers.h" // TF:flatbuffers -#include "tensorflow/contrib/lite/interpreter.h" -#include "tensorflow/contrib/lite/kernels/register.h" -#include "tensorflow/contrib/lite/kernels/test_util.h" -#include "tensorflow/contrib/lite/model.h" - -namespace tflite { -namespace ops { -namespace custom { - -TfLiteRegistration* Register_LAYER_NORM_LSTM(); - -namespace { - -using ::testing::ElementsAreArray; - -class LayerNormLSTMOpModel : public SingleOpModel { - public: - LayerNormLSTMOpModel(int n_batch, int n_input, int n_cell, int n_output, - bool use_cifg, bool use_peephole, - bool use_projection_weights, bool use_projection_bias, - float cell_clip, float proj_clip, - const std::vector>& input_shapes, - const TensorType& weight_type = TensorType_FLOAT32) - : n_batch_(n_batch), - n_input_(n_input), - n_cell_(n_cell), - n_output_(n_output) { - input_ = AddInput(TensorType_FLOAT32); - - if (use_cifg) { - input_to_input_weights_ = AddNullInput(); - } else { - input_to_input_weights_ = AddInput(weight_type); - } - - input_to_forget_weights_ = AddInput(weight_type); - input_to_cell_weights_ = AddInput(weight_type); - input_to_output_weights_ = AddInput(weight_type); - - if (use_cifg) { - recurrent_to_input_weights_ = AddNullInput(); - } else { - recurrent_to_input_weights_ = AddInput(weight_type); - } - - recurrent_to_forget_weights_ = AddInput(weight_type); - recurrent_to_cell_weights_ = AddInput(weight_type); - recurrent_to_output_weights_ = AddInput(weight_type); - - if (use_peephole) { - if (use_cifg) { - cell_to_input_weights_ = AddNullInput(); - } else { - cell_to_input_weights_ = AddInput(weight_type); - } - cell_to_forget_weights_ = AddInput(weight_type); - cell_to_output_weights_ = AddInput(weight_type); - } else { - cell_to_input_weights_ = AddNullInput(); - cell_to_forget_weights_ = AddNullInput(); - cell_to_output_weights_ = AddNullInput(); - } - - input_layer_norm_weights_ = AddInput(TensorType_FLOAT32); - forget_layer_norm_weights_ = AddInput(TensorType_FLOAT32); - cell_layer_norm_weights_ = AddInput(TensorType_FLOAT32); - output_layer_norm_weights_ = AddInput(TensorType_FLOAT32); - - if (use_cifg) { - input_gate_bias_ = AddNullInput(); - } else { - input_gate_bias_ = AddInput(TensorType_FLOAT32); - } - forget_gate_bias_ = AddInput(TensorType_FLOAT32); - cell_bias_ = AddInput(TensorType_FLOAT32); - output_gate_bias_ = AddInput(TensorType_FLOAT32); - - if (use_projection_weights) { - projection_weights_ = AddInput(weight_type); - if (use_projection_bias) { - projection_bias_ = AddInput(TensorType_FLOAT32); - } else { - projection_bias_ = AddNullInput(); - } - } else { - projection_weights_ = AddNullInput(); - projection_bias_ = AddNullInput(); - } - - // Adding the 2 state tensors. - output_state_ = - AddInput(TensorData{TensorType_FLOAT32, {n_output_ * n_batch_}}, true); - cell_state_ = - AddInput(TensorData{TensorType_FLOAT32, {n_cell_ * n_batch_}}, true); - - output_ = AddOutput(TensorType_FLOAT32); - - // Set up and pass in custom options using flexbuffer. - flexbuffers::Builder fbb; - fbb.Map([&]() { - fbb.Int("cell_clip", cell_clip); - fbb.Int("proj_clip", proj_clip); - fbb.String("fused_activation_function", "TANH"); - }); - fbb.Finish(); - SetCustomOp("LAYER_NORM_LSTM", fbb.GetBuffer(), Register_LAYER_NORM_LSTM); - BuildInterpreter(input_shapes); - } - - void SetInputToInputWeights(std::vector f) { - PopulateTensor(input_to_input_weights_, f); - } - - void SetInputToForgetWeights(std::vector f) { - PopulateTensor(input_to_forget_weights_, f); - } - - void SetInputToCellWeights(std::vector f) { - PopulateTensor(input_to_cell_weights_, f); - } - - void SetInputToOutputWeights(std::vector f) { - PopulateTensor(input_to_output_weights_, f); - } - - void SetRecurrentToInputWeights(std::vector f) { - PopulateTensor(recurrent_to_input_weights_, f); - } - - void SetRecurrentToForgetWeights(std::vector f) { - PopulateTensor(recurrent_to_forget_weights_, f); - } - - void SetRecurrentToCellWeights(std::vector f) { - PopulateTensor(recurrent_to_cell_weights_, f); - } - - void SetRecurrentToOutputWeights(std::vector f) { - PopulateTensor(recurrent_to_output_weights_, f); - } - - void SetCellToInputWeights(std::vector f) { - PopulateTensor(cell_to_input_weights_, f); - } - - void SetCellToForgetWeights(std::vector f) { - PopulateTensor(cell_to_forget_weights_, f); - } - - void SetCellToOutputWeights(std::vector f) { - PopulateTensor(cell_to_output_weights_, f); - } - - void SetInputLayerNormWeights(std::vector f) { - PopulateTensor(input_layer_norm_weights_, f); - } - - void SetForgetLayerNormWeights(std::vector f) { - PopulateTensor(forget_layer_norm_weights_, f); - } - - void SetCellLayerNormWeights(std::vector f) { - PopulateTensor(cell_layer_norm_weights_, f); - } - - void SetOutputLayerNormWeights(std::vector f) { - PopulateTensor(output_layer_norm_weights_, f); - } - - void SetInputGateBias(std::vector f) { - PopulateTensor(input_gate_bias_, f); - } - - void SetForgetGateBias(std::vector f) { - PopulateTensor(forget_gate_bias_, f); - } - - void SetCellBias(std::vector f) { PopulateTensor(cell_bias_, f); } - - void SetOutputGateBias(std::vector f) { - PopulateTensor(output_gate_bias_, f); - } - - void SetProjectionWeights(std::vector f) { - PopulateTensor(projection_weights_, f); - } - - void SetProjectionBias(std::vector f) { - PopulateTensor(projection_bias_, f); - } - - void SetInput(int offset, const float* begin, const float* end) { - PopulateTensor(input_, offset, const_cast(begin), - const_cast(end)); - } - - std::vector GetOutput() { return ExtractVector(output_); } - - int num_inputs() { return n_input_; } - int num_outputs() { return n_output_; } - int num_cells() { return n_cell_; } - int num_batches() { return n_batch_; } - - protected: - int input_; - int input_to_input_weights_; - int input_to_forget_weights_; - int input_to_cell_weights_; - int input_to_output_weights_; - - int recurrent_to_input_weights_; - int recurrent_to_forget_weights_; - int recurrent_to_cell_weights_; - int recurrent_to_output_weights_; - - int cell_to_input_weights_; - int cell_to_forget_weights_; - int cell_to_output_weights_; - - int input_layer_norm_weights_; - int forget_layer_norm_weights_; - int cell_layer_norm_weights_; - int output_layer_norm_weights_; - - int input_gate_bias_; - int forget_gate_bias_; - int cell_bias_; - int output_gate_bias_; - - int projection_weights_; - int projection_bias_; - - int output_state_; - int cell_state_; - - int output_; - - int n_batch_; - int n_input_; - int n_cell_; - int n_output_; -}; - -class HybridLayerNormLSTMOpModel : public LayerNormLSTMOpModel { - public: - HybridLayerNormLSTMOpModel(int n_batch, int n_input, int n_cell, int n_output, - bool use_cifg, bool use_peephole, - bool use_projection_weights, - bool use_projection_bias, float cell_clip, - float proj_clip, - const std::vector>& input_shapes) - : LayerNormLSTMOpModel(n_batch, n_input, n_cell, n_output, use_cifg, - use_peephole, use_projection_weights, - use_projection_bias, cell_clip, proj_clip, - input_shapes, TensorType_UINT8) {} - - void SetInputToInputWeights(std::vector f) { - SymmetricQuantizeAndPopulate(input_to_input_weights_, f); - } - - void SetInputToForgetWeights(std::vector f) { - SymmetricQuantizeAndPopulate(input_to_forget_weights_, f); - } - - void SetInputToCellWeights(std::vector f) { - SymmetricQuantizeAndPopulate(input_to_cell_weights_, f); - } - - void SetInputToOutputWeights(std::vector f) { - SymmetricQuantizeAndPopulate(input_to_output_weights_, f); - } - - void SetRecurrentToInputWeights(std::vector f) { - SymmetricQuantizeAndPopulate(recurrent_to_input_weights_, f); - } - - void SetRecurrentToForgetWeights(std::vector f) { - SymmetricQuantizeAndPopulate(recurrent_to_forget_weights_, f); - } - - void SetRecurrentToCellWeights(std::vector f) { - SymmetricQuantizeAndPopulate(recurrent_to_cell_weights_, f); - } - - void SetRecurrentToOutputWeights(std::vector f) { - SymmetricQuantizeAndPopulate(recurrent_to_output_weights_, f); - } - - void SetCellToInputWeights(std::vector f) { - SymmetricQuantizeAndPopulate(cell_to_input_weights_, f); - } - - void SetCellToForgetWeights(std::vector f) { - SymmetricQuantizeAndPopulate(cell_to_forget_weights_, f); - } - - void SetCellToOutputWeights(std::vector f) { - SymmetricQuantizeAndPopulate(cell_to_output_weights_, f); - } - - void SetInputLayerNormWeights(std::vector f) { - PopulateTensor(input_layer_norm_weights_, f); - } - - void SetForgetLayerNormWeights(std::vector f) { - PopulateTensor(forget_layer_norm_weights_, f); - } - - void SetCellLayerNormWeights(std::vector f) { - PopulateTensor(cell_layer_norm_weights_, f); - } - - void SetOutputLayerNormWeights(std::vector f) { - PopulateTensor(output_layer_norm_weights_, f); - } - - void SetProjectionWeights(std::vector f) { - SymmetricQuantizeAndPopulate(projection_weights_, f); - } -}; - -class BaseLayerNormLstmTest : public ::testing::Test { - protected: - // Weights of the Layer Norm LSTM model. Some are optional. - std::vector input_to_input_weights_; - std::vector input_to_cell_weights_; - std::vector input_to_forget_weights_; - std::vector input_to_output_weights_; - std::vector input_gate_bias_; - std::vector cell_gate_bias_; - std::vector forget_gate_bias_; - std::vector output_gate_bias_; - std::vector recurrent_to_input_weights_; - std::vector recurrent_to_cell_weights_; - std::vector recurrent_to_forget_weights_; - std::vector recurrent_to_output_weights_; - std::vector cell_to_input_weights_; - std::vector cell_to_forget_weights_; - std::vector cell_to_output_weights_; - std::vector input_layer_norm_weights_; - std::vector forget_layer_norm_weights_; - std::vector cell_layer_norm_weights_; - std::vector output_layer_norm_weights_; - std::vector projection_weights_; - - // Layer Norm LSTM input is stored as num_batch x num_inputs vector. - std::vector> layer_norm_lstm_input_; - - // Compares output up to tolerance to the result of the layer_norm_lstm given - // the input. - void VerifyGoldens(const std::vector>& input, - const std::vector>& output, - LayerNormLSTMOpModel* layer_norm_lstm, - float tolerance = 1e-5) { - const int num_batches = input.size(); - EXPECT_GT(num_batches, 0); - const int num_inputs = layer_norm_lstm->num_inputs(); - EXPECT_GT(num_inputs, 0); - const int input_sequence_size = input[0].size() / num_inputs; - EXPECT_GT(input_sequence_size, 0); - for (int i = 0; i < input_sequence_size; ++i) { - for (int b = 0; b < num_batches; ++b) { - const float* batch_start = input[b].data() + i * num_inputs; - const float* batch_end = batch_start + num_inputs; - - layer_norm_lstm->SetInput(b * layer_norm_lstm->num_inputs(), - batch_start, batch_end); - } - - layer_norm_lstm->Invoke(); - - const int num_outputs = layer_norm_lstm->num_outputs(); - std::vector expected; - for (int b = 0; b < num_batches; ++b) { - const float* golden_start_batch = output[b].data() + i * num_outputs; - const float* golden_end_batch = golden_start_batch + num_outputs; - expected.insert(expected.end(), golden_start_batch, golden_end_batch); - } - EXPECT_THAT(layer_norm_lstm->GetOutput(), - ElementsAreArray(ArrayFloatNear(expected, tolerance))); - } - } -}; - -class NoCifgPeepholeProjectionNoClippingLayerNormLstmTest - : public BaseLayerNormLstmTest { - void SetUp() override { - input_to_input_weights_ = {0.5, 0.6, 0.7, -0.8, -0.9, 0.1, 0.2, - 0.3, -0.4, 0.5, -0.8, 0.7, -0.6, 0.5, - -0.4, -0.5, -0.4, -0.3, -0.2, -0.1}; - - input_to_forget_weights_ = {-0.6, -0.1, 0.3, 0.2, 0.9, -0.5, -0.2, - -0.4, 0.3, -0.8, -0.4, 0.3, -0.5, -0.4, - -0.6, 0.3, -0.4, -0.6, -0.5, -0.5}; - - input_to_cell_weights_ = {-0.4, -0.3, -0.2, -0.1, -0.5, 0.5, -0.2, - -0.3, -0.2, -0.6, 0.6, -0.1, -0.4, -0.3, - -0.7, 0.7, -0.9, -0.5, 0.8, 0.6}; - - input_to_output_weights_ = {-0.8, -0.4, -0.2, -0.9, -0.1, -0.7, 0.3, - -0.3, -0.8, -0.2, 0.6, -0.2, 0.4, -0.7, - -0.3, -0.5, 0.1, 0.5, -0.6, -0.4}; - - input_gate_bias_ = {0.03, 0.15, 0.22, 0.38}; - - forget_gate_bias_ = {0.1, -0.3, -0.2, 0.1}; - - cell_gate_bias_ = {-0.05, 0.72, 0.25, 0.08}; - - output_gate_bias_ = {0.05, -0.01, 0.2, 0.1}; - - recurrent_to_input_weights_ = {-0.2, -0.3, 0.4, 0.1, -0.5, 0.9, - -0.2, -0.3, -0.7, 0.05, -0.2, -0.6}; - - recurrent_to_cell_weights_ = {-0.3, 0.2, 0.1, -0.3, 0.8, -0.08, - -0.2, 0.3, 0.8, -0.6, -0.1, 0.2}; - - recurrent_to_forget_weights_ = {-0.5, -0.3, -0.5, -0.2, 0.6, 0.4, - 0.9, 0.3, -0.1, 0.2, 0.5, 0.2}; - - recurrent_to_output_weights_ = {0.3, -0.1, 0.1, -0.2, -0.5, -0.7, - -0.2, -0.6, -0.1, -0.4, -0.7, -0.2}; - - cell_to_input_weights_ = {0.05, 0.1, 0.25, 0.15}; - - cell_to_forget_weights_ = {-0.02, -0.15, -0.25, -0.03}; - - cell_to_output_weights_ = {0.1, -0.1, -0.5, 0.05}; - - input_layer_norm_weights_ = {0.1, 0.2, 0.3, 0.5}; - forget_layer_norm_weights_ = {0.2, 0.2, 0.4, 0.3}; - cell_layer_norm_weights_ = {0.7, 0.2, 0.3, 0.8}; - output_layer_norm_weights_ = {0.6, 0.2, 0.2, 0.5}; - - projection_weights_ = {-0.1, 0.2, 0.01, -0.2, 0.1, 0.5, - 0.3, 0.08, 0.07, 0.2, -0.4, 0.2}; - - layer_norm_lstm_input_ = { - {// Batch0: 3 (input_sequence_size) * 5 (n_input) - 0.7, 0.8, 0.1, 0.2, 0.3, // seq 0 - 0.8, 0.1, 0.2, 0.4, 0.5, // seq 1 - 0.2, 0.7, 0.7, 0.1, 0.7}, // seq 2 - - {// Batch1: 3 (input_sequence_size) * 5 (n_input) - 0.3, 0.2, 0.9, 0.8, 0.1, // seq 0 - 0.1, 0.5, 0.2, 0.4, 0.2, // seq 1 - 0.6, 0.9, 0.2, 0.5, 0.7}, // seq 2 - }; - } -}; - -TEST_F(NoCifgPeepholeProjectionNoClippingLayerNormLstmTest, - LayerNormLstmBlackBoxTest) { - const int n_batch = 2; - const int n_input = 5; - const int n_cell = 4; - const int n_output = 3; - const float ceil_clip = 0.0; - const float proj_clip = 0.0; - - LayerNormLSTMOpModel layer_norm_lstm( - n_batch, n_input, n_cell, n_output, - /*use_cifg=*/false, /*use_peephole=*/true, - /*use_projection_weights=*/true, - /*use_projection_bias=*/false, ceil_clip, proj_clip, - { - {n_batch, n_input}, // input tensor - - {n_cell, n_input}, // input_to_input_weight tensor - {n_cell, n_input}, // input_to_forget_weight tensor - {n_cell, n_input}, // input_to_cell_weight tensor - {n_cell, n_input}, // input_to_output_weight tensor - - {n_cell, n_output}, // recurrent_to_input_weight tensor - {n_cell, n_output}, // recurrent_to_forget_weight tensor - {n_cell, n_output}, // recurrent_to_cell_weight tensor - {n_cell, n_output}, // recurrent_to_output_weight tensor - - {n_cell}, // cell_to_input_weight tensor - {n_cell}, // cell_to_forget_weight tensor - {n_cell}, // cell_to_output_weight tensor - - {n_cell}, // input_layer_norm_weight tensor - {n_cell}, // forget_layer_norm_weight tensor - {n_cell}, // cell_layer_norm_weight tensor - {n_cell}, // output_layer_norm_weight tensor - - {n_cell}, // input_gate_bias tensor - {n_cell}, // forget_gate_bias tensor - {n_cell}, // cell_bias tensor - {n_cell}, // output_gate_bias tensor - - {n_output, n_cell}, // projection_weight tensor - {0}, // projection_bias tensor - }); - - layer_norm_lstm.SetInputToInputWeights(input_to_input_weights_); - layer_norm_lstm.SetInputToCellWeights(input_to_cell_weights_); - layer_norm_lstm.SetInputToForgetWeights(input_to_forget_weights_); - layer_norm_lstm.SetInputToOutputWeights(input_to_output_weights_); - - layer_norm_lstm.SetInputGateBias(input_gate_bias_); - layer_norm_lstm.SetCellBias(cell_gate_bias_); - layer_norm_lstm.SetForgetGateBias(forget_gate_bias_); - layer_norm_lstm.SetOutputGateBias(output_gate_bias_); - - layer_norm_lstm.SetRecurrentToInputWeights(recurrent_to_input_weights_); - layer_norm_lstm.SetRecurrentToCellWeights(recurrent_to_cell_weights_); - layer_norm_lstm.SetRecurrentToForgetWeights(recurrent_to_forget_weights_); - layer_norm_lstm.SetRecurrentToOutputWeights(recurrent_to_output_weights_); - - layer_norm_lstm.SetCellToInputWeights(cell_to_input_weights_); - layer_norm_lstm.SetCellToForgetWeights(cell_to_forget_weights_); - layer_norm_lstm.SetCellToOutputWeights(cell_to_output_weights_); - - layer_norm_lstm.SetInputLayerNormWeights(input_layer_norm_weights_); - layer_norm_lstm.SetForgetLayerNormWeights(forget_layer_norm_weights_); - layer_norm_lstm.SetCellLayerNormWeights(cell_layer_norm_weights_); - layer_norm_lstm.SetOutputLayerNormWeights(output_layer_norm_weights_); - - layer_norm_lstm.SetProjectionWeights(projection_weights_); - - // Verify the final output. - const std::vector> layer_norm_lstm_golden_output = { - { - // Batch0: 3 (input_sequence_size) * 3 (n_output) - 0.0244077, 0.128027, -0.00170918, // seq 0 - 0.0137642, 0.140751, 0.0395835, // seq 1 - -0.00459231, 0.155278, 0.0837377, // seq 2 - }, - { - // Batch1: 3 (input_sequence_size) * 3 (n_output) - -0.00692428, 0.0848741, 0.063445, // seq 0 - -0.00403912, 0.139963, 0.072681, // seq 1 - 0.00752706, 0.161903, 0.0561371, // seq 2 - }}; - - VerifyGoldens(layer_norm_lstm_input_, layer_norm_lstm_golden_output, - &layer_norm_lstm); -} - -TEST_F(NoCifgPeepholeProjectionNoClippingLayerNormLstmTest, - HybridLayerNormLstmBlackBoxTest) { - const int n_batch = 2; - const int n_input = 5; - const int n_cell = 4; - const int n_output = 3; - const float ceil_clip = 0.0; - const float proj_clip = 0.0; - - HybridLayerNormLSTMOpModel layer_norm_lstm( - n_batch, n_input, n_cell, n_output, - /*use_cifg=*/false, /*use_peephole=*/true, - /*use_projection_weights=*/true, - /*use_projection_bias=*/false, ceil_clip, proj_clip, - { - {n_batch, n_input}, // input tensor - - {n_cell, n_input}, // input_to_input_weight tensor - {n_cell, n_input}, // input_to_forget_weight tensor - {n_cell, n_input}, // input_to_cell_weight tensor - {n_cell, n_input}, // input_to_output_weight tensor - - {n_cell, n_output}, // recurrent_to_input_weight tensor - {n_cell, n_output}, // recurrent_to_forget_weight tensor - {n_cell, n_output}, // recurrent_to_cell_weight tensor - {n_cell, n_output}, // recurrent_to_output_weight tensor - - {n_cell}, // cell_to_input_weight tensor - {n_cell}, // cell_to_forget_weight tensor - {n_cell}, // cell_to_output_weight tensor - - {n_cell}, // input_layer_norm_weight tensor - {n_cell}, // forget_layer_norm_weight tensor - {n_cell}, // cell_layer_norm_weight tensor - {n_cell}, // output_layer_norm_weight tensor - - {n_cell}, // input_gate_bias tensor - {n_cell}, // forget_gate_bias tensor - {n_cell}, // cell_bias tensor - {n_cell}, // output_gate_bias tensor - - {n_output, n_cell}, // projection_weight tensor - {0}, // projection_bias tensor - }); - - layer_norm_lstm.SetInputToInputWeights(input_to_input_weights_); - layer_norm_lstm.SetInputToCellWeights(input_to_cell_weights_); - layer_norm_lstm.SetInputToForgetWeights(input_to_forget_weights_); - layer_norm_lstm.SetInputToOutputWeights(input_to_output_weights_); - - layer_norm_lstm.SetInputGateBias(input_gate_bias_); - layer_norm_lstm.SetCellBias(cell_gate_bias_); - layer_norm_lstm.SetForgetGateBias(forget_gate_bias_); - layer_norm_lstm.SetOutputGateBias(output_gate_bias_); - - layer_norm_lstm.SetRecurrentToInputWeights(recurrent_to_input_weights_); - layer_norm_lstm.SetRecurrentToCellWeights(recurrent_to_cell_weights_); - layer_norm_lstm.SetRecurrentToForgetWeights(recurrent_to_forget_weights_); - layer_norm_lstm.SetRecurrentToOutputWeights(recurrent_to_output_weights_); - - layer_norm_lstm.SetCellToInputWeights(cell_to_input_weights_); - layer_norm_lstm.SetCellToForgetWeights(cell_to_forget_weights_); - layer_norm_lstm.SetCellToOutputWeights(cell_to_output_weights_); - - layer_norm_lstm.SetInputLayerNormWeights(input_layer_norm_weights_); - layer_norm_lstm.SetForgetLayerNormWeights(forget_layer_norm_weights_); - layer_norm_lstm.SetCellLayerNormWeights(cell_layer_norm_weights_); - layer_norm_lstm.SetOutputLayerNormWeights(output_layer_norm_weights_); - - layer_norm_lstm.SetProjectionWeights(projection_weights_); - - const std::vector> layer_norm_lstm_golden_output = { - { - // Batch0: 3 (input_sequence_size) * 3 (n_output) - 0.0244576, 0.127847, -0.00181765, // seq 0 - 0.0137518, 0.140892, 0.0402234, // seq 1 - -0.0048839, 0.155096, 0.0840309, // seq 2 - }, - { - // Batch1: 3 (input_sequence_size) * 3 (n_output) - -0.00728636, 0.0843957, 0.0634786, // seq 0 - -0.00448382, 0.139278, 0.0737372, // seq 1 - 0.00734616, 0.161793, 0.0560238, // seq 2 - }}; - - VerifyGoldens(layer_norm_lstm_input_, layer_norm_lstm_golden_output, - &layer_norm_lstm); -} - -} // namespace -} // namespace custom -} // namespace ops -} // namespace tflite - -int main(int argc, char** argv) { - ::tflite::LogToStderr(); - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/contrib/lite/kernels/lstm_eval.cc b/tensorflow/contrib/lite/kernels/lstm_eval.cc deleted file mode 100644 index f2ba7b46d9b053358f71b076fbdfc17a3e423f38..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/lstm_eval.cc +++ /dev/null @@ -1,1139 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include "tensorflow/contrib/lite/kernels/lstm_eval.h" - -#include - -#include "tensorflow/contrib/lite/kernels/internal/kernel_utils.h" -#include "tensorflow/contrib/lite/kernels/internal/tensor_utils.h" -#include "tensorflow/contrib/lite/kernels/op_macros.h" - -namespace tflite { -namespace ops { -namespace builtin { -namespace lstm_eval { - -namespace { - -// Performs an LSTM batch inference step for input specified by input_ptr_batch. -// The LSTM cell is specified by the pointers to its weights (*_weights_ptr) and -// biases (*_bias_ptr), and buffers (*_scratch), along with additional -// parameters: -// - params: various LSTM params including activation, clipping, etc., -// - n_batch: size of batch, -// - n_cell: number of cells (or units), -// - n_input: the input size, -// - n_output: the output size. -// - output_batch_leading_dim: the leading dimension of the output buffer. -// -// The pointers to the cell and output state and the output are updated. -// -// The pointers with the suffix "_batch" point to data aligned in batch_major -// order, and each step processes batch_size many inputs from input_ptr_batch, -// and updates batch_size many cell and output states. -// -// The output_batch_dim is output.shape[-1], i.e. the outermost dimension of the -// output tensor, and in most cases will be equal to n_output. It is usually not -// when we want to store the LSTM output into a slice of the output tensor, e.g. -// for bidirectional LSTMs with merge_outputs. In this case, the batched -// operations cannot be used since they assume that the batched outputs are -// contiguous, and we manually loop over the batched outputs. -inline void LstmStepWithAuxInput( - const float* input_ptr_batch, const float* input_to_input_weights_ptr, - const float* input_to_forget_weights_ptr, - const float* input_to_cell_weights_ptr, - const float* input_to_output_weights_ptr, const float* aux_input_ptr_batch, - const float* aux_input_to_input_weights_ptr, - const float* aux_input_to_forget_weights_ptr, - const float* aux_input_to_cell_weights_ptr, - const float* aux_input_to_output_weights_ptr, - const float* recurrent_to_input_weights_ptr, - const float* recurrent_to_forget_weights_ptr, - const float* recurrent_to_cell_weights_ptr, - const float* recurrent_to_output_weights_ptr, - const float* cell_to_input_weights_ptr, - const float* cell_to_forget_weights_ptr, - const float* cell_to_output_weights_ptr, const float* input_gate_bias_ptr, - const float* forget_gate_bias_ptr, const float* cell_bias_ptr, - const float* output_gate_bias_ptr, const float* projection_weights_ptr, - const float* projection_bias_ptr, const TfLiteLSTMParams* params, - int n_batch, int n_cell, int n_input, int n_aux_input, int n_output, - int output_batch_leading_dim, float* output_state_ptr, - float* cell_state_ptr, float* input_gate_scratch, - float* forget_gate_scratch, float* cell_scratch, float* output_gate_scratch, - float* output_ptr_batch) { - // Since we have already checked that weights are all there or none, we can - // check the existense of only one to the get the condition. - const bool use_cifg = (input_to_input_weights_ptr == nullptr); - const bool use_peephole = (cell_to_output_weights_ptr != nullptr); - // Initialize scratch buffers with bias. - if (!use_cifg) { - tensor_utils::VectorBatchVectorAssign(input_gate_bias_ptr, n_cell, n_batch, - input_gate_scratch); - } - tensor_utils::VectorBatchVectorAssign(forget_gate_bias_ptr, n_cell, n_batch, - forget_gate_scratch); - tensor_utils::VectorBatchVectorAssign(cell_bias_ptr, n_cell, n_batch, - cell_scratch); - tensor_utils::VectorBatchVectorAssign(output_gate_bias_ptr, n_cell, n_batch, - output_gate_scratch); - - // For each batch and cell: compute input_weight * input. - if (!use_cifg) { - tensor_utils::MatrixBatchVectorMultiplyAccumulate( - input_to_input_weights_ptr, n_cell, n_input, input_ptr_batch, n_batch, - input_gate_scratch, /*result_stride=*/1); - } - - tensor_utils::MatrixBatchVectorMultiplyAccumulate( - input_to_forget_weights_ptr, n_cell, n_input, input_ptr_batch, n_batch, - forget_gate_scratch, /*result_stride=*/1); - tensor_utils::MatrixBatchVectorMultiplyAccumulate( - input_to_cell_weights_ptr, n_cell, n_input, input_ptr_batch, n_batch, - cell_scratch, /*result_stride=*/1); - tensor_utils::MatrixBatchVectorMultiplyAccumulate( - input_to_output_weights_ptr, n_cell, n_input, input_ptr_batch, n_batch, - output_gate_scratch, /*result_stride=*/1); - - // If auxiliary input is available then compute aux_input_weight * aux_input - if (aux_input_ptr_batch != nullptr) { - if (!use_cifg) { - tensor_utils::MatrixBatchVectorMultiplyAccumulate( - aux_input_to_input_weights_ptr, n_cell, n_aux_input, - aux_input_ptr_batch, n_batch, input_gate_scratch, - /*result_stride=*/1); - } - - tensor_utils::MatrixBatchVectorMultiplyAccumulate( - aux_input_to_forget_weights_ptr, n_cell, n_aux_input, - aux_input_ptr_batch, n_batch, forget_gate_scratch, /*result_stride=*/1); - tensor_utils::MatrixBatchVectorMultiplyAccumulate( - aux_input_to_cell_weights_ptr, n_cell, n_aux_input, aux_input_ptr_batch, - n_batch, cell_scratch, /*result_stride=*/1); - tensor_utils::MatrixBatchVectorMultiplyAccumulate( - aux_input_to_output_weights_ptr, n_cell, n_aux_input, - aux_input_ptr_batch, n_batch, output_gate_scratch, /*result_stride=*/1); - } - - // For each batch and cell: compute recurrent_weight * output_state. - if (!use_cifg) { - tensor_utils::MatrixBatchVectorMultiplyAccumulate( - recurrent_to_input_weights_ptr, n_cell, n_output, output_state_ptr, - n_batch, input_gate_scratch, /*result_stride=*/1); - } - tensor_utils::MatrixBatchVectorMultiplyAccumulate( - recurrent_to_forget_weights_ptr, n_cell, n_output, output_state_ptr, - n_batch, forget_gate_scratch, - /*result_stride=*/1); - tensor_utils::MatrixBatchVectorMultiplyAccumulate( - recurrent_to_cell_weights_ptr, n_cell, n_output, output_state_ptr, - n_batch, cell_scratch, /*result_stride=*/1); - tensor_utils::MatrixBatchVectorMultiplyAccumulate( - recurrent_to_output_weights_ptr, n_cell, n_output, output_state_ptr, - n_batch, output_gate_scratch, - /*result_stride=*/1); - - // For each batch and cell: update input gate. - if (!use_cifg) { - if (use_peephole) { - tensor_utils::VectorBatchVectorCwiseProductAccumulate( - cell_to_input_weights_ptr, n_cell, cell_state_ptr, n_batch, - input_gate_scratch); - } - tensor_utils::ApplySigmoidToVector(input_gate_scratch, n_cell * n_batch, - input_gate_scratch); - } - - // For each batch and cell: update forget gate. - if (use_peephole) { - tensor_utils::VectorBatchVectorCwiseProductAccumulate( - cell_to_forget_weights_ptr, n_cell, cell_state_ptr, n_batch, - forget_gate_scratch); - } - tensor_utils::ApplySigmoidToVector(forget_gate_scratch, n_cell * n_batch, - forget_gate_scratch); - - // For each batch and cell: update the cell. - tensor_utils::VectorVectorCwiseProduct(forget_gate_scratch, cell_state_ptr, - n_batch * n_cell, cell_state_ptr); - tensor_utils::ApplyActivationToVector(cell_scratch, n_batch * n_cell, - params->activation, cell_scratch); - if (use_cifg) { - tensor_utils::Sub1Vector(forget_gate_scratch, n_batch * n_cell, - forget_gate_scratch); - tensor_utils::VectorVectorCwiseProductAccumulate( - cell_scratch, forget_gate_scratch, n_batch * n_cell, cell_state_ptr); - } else { - tensor_utils::VectorVectorCwiseProductAccumulate( - cell_scratch, input_gate_scratch, n_batch * n_cell, cell_state_ptr); - } - if (params->cell_clip > 0.0) { - tensor_utils::ClipVector(cell_state_ptr, n_batch * n_cell, - params->cell_clip, cell_state_ptr); - } - - // For each batch and cell: update the output gate. - if (use_peephole) { - tensor_utils::VectorBatchVectorCwiseProductAccumulate( - cell_to_output_weights_ptr, n_cell, cell_state_ptr, n_batch, - output_gate_scratch); - } - tensor_utils::ApplySigmoidToVector(output_gate_scratch, n_batch * n_cell, - output_gate_scratch); - tensor_utils::ApplyActivationToVector(cell_state_ptr, n_batch * n_cell, - params->activation, cell_scratch); - tensor_utils::VectorVectorCwiseProduct(output_gate_scratch, cell_scratch, - n_batch * n_cell, output_gate_scratch); - - const bool use_projection_weight = (projection_weights_ptr != nullptr); - const bool use_projection_bias = (projection_bias_ptr != nullptr); - - // For each batch: update the projection and output_state. Note that since - // the output batch rows may not be contiguous (output_batch_leading_dim != - // n_output), we unroll the batched operations where this is the case. - if (output_batch_leading_dim == n_output) { - if (use_projection_weight) { - if (use_projection_bias) { - tensor_utils::VectorBatchVectorAssign(projection_bias_ptr, n_output, - n_batch, output_ptr_batch); - } else { - tensor_utils::ZeroVector(output_ptr_batch, n_batch * n_output); - } - tensor_utils::MatrixBatchVectorMultiplyAccumulate( - projection_weights_ptr, n_output, n_cell, output_gate_scratch, - n_batch, output_ptr_batch, /*result_stride=*/1); - if (params->proj_clip > 0.0) { - tensor_utils::ClipVector(output_ptr_batch, n_batch * n_output, - params->proj_clip, output_ptr_batch); - } - } else { - tensor_utils::CopyVector(output_gate_scratch, n_batch * n_output, - output_ptr_batch); - } - tensor_utils::CopyVector(output_ptr_batch, n_batch * n_output, - output_state_ptr); - } else { - if (use_projection_weight) { - if (use_projection_bias) { - for (int k = 0; k < n_batch; k++) { - tensor_utils::CopyVector( - projection_bias_ptr, n_output, - output_ptr_batch + k * output_batch_leading_dim); - } - } else { - for (int k = 0; k < n_batch; k++) { - tensor_utils::ZeroVector( - output_ptr_batch + k * output_batch_leading_dim, n_output); - } - } - for (int k = 0; k < n_batch; k++) { - tensor_utils::MatrixBatchVectorMultiplyAccumulate( - projection_weights_ptr, n_output, n_cell, - output_gate_scratch + k * n_cell, - /*n_batch=*/1, output_ptr_batch + k * output_batch_leading_dim, - /*result_stride=*/1); - if (params->proj_clip > 0.0) { - tensor_utils::ClipVector( - output_ptr_batch + k * output_batch_leading_dim, n_output, - params->proj_clip, - output_ptr_batch + k * output_batch_leading_dim); - } - } - } else { - for (int k = 0; k < n_batch; k++) { - tensor_utils::CopyVector( - output_gate_scratch + k * n_output, n_output, - output_ptr_batch + k * output_batch_leading_dim); - } - } - for (int k = 0; k < n_batch; k++) { - tensor_utils::CopyVector(output_ptr_batch + k * output_batch_leading_dim, - n_output, output_state_ptr + k * n_output); - } - } -} - -// Same as above but with quantized weight matrices. In detail: -// Input of size 'n_batch * n_input': -// input_ptr_batch -// -// LSTM weights: -// Quantized input weights of size 'n_cell * n_input': -// input_to_input_weights - optional (can be nullptr) -// input_to_forget_weights -// input_to_cell_weights -// input_to_input_weights -// Quantized recurrent weights of size 'n_cell * n_output': -// recurrent_to_input_weights - optional -// recurrent_to_forget_weights -// recurrent_to_cell_weights -// recurrent_to_input_weights -// Quantized peephole weights of size 'n_cell', representing diagonal matrices. -// cell_to_input_weights - optional -// cell_to_cell_weights - optional -// cell_to_output_weights - optional -// Quantized projection weights of size 'n_output * n_cell' -// projection_weights_ptr - optional -// Weight scales (scalars) for each of the weights above. -// input_to_input_weights_scale - optional -// input_to_forget_weights_scale -// input_to_cell_weights_scale -// input_to_output_weights_scale -// recurrent_to_input_weights_scale - optional -// recurrent_to_forget_weights_scale -// recurrent_to_cell_weights_scale -// recurrent_to_output_weights_scale -// cell_to_input_weights_scale, -// cell_to_forget_weights_scale, -// cell_to_output_weights_scale, -// projection_weights_scale - optional -// Gate biases of size 'n_cell': -// input_gate_bias_ptr - optional -// forget_gate_bias_ptr -// cell_gate_bias_ptr -// output_gate_bias_ptr -// -// Temporary pre-allocated storage for quantized values: -// quantized_input_ptr_batch (same size as input_ptr_batch) -// quantized_output_state_ptr (same size as output_state_ptr) -// quantized_cell_state_ptr (same size as cell_state_ptr) -// Temporary pre-allocated storage for recovered values: -// recovered_cell_weights (same size as cell_to_*_weights) -// -// Outputs: -// output_state_ptr - size 'n_batch * n_output' -// cell_state_ptr - size 'n_batch * n_cell' -// output_ptr_batch - size 'n_batch * output_batch_leading_dim' -inline void LstmStepWithAuxInput( - const float* input_ptr_batch, const int8_t* input_to_input_weights_ptr, - float input_to_input_weights_scale, - const int8_t* input_to_forget_weights_ptr, - float input_to_forget_weights_scale, - const int8_t* input_to_cell_weights_ptr, float input_to_cell_weights_scale, - const int8_t* input_to_output_weights_ptr, - float input_to_output_weights_scale, const float* aux_input_ptr_batch, - const int8_t* aux_input_to_input_weights_ptr, - float aux_input_to_input_weights_scale, - const int8_t* aux_input_to_forget_weights_ptr, - float aux_input_to_forget_weights_scale, - const int8_t* aux_input_to_cell_weights_ptr, - float aux_input_to_cell_weights_scale, - const int8_t* aux_input_to_output_weights_ptr, - float aux_input_to_output_weights_scale, - const int8_t* recurrent_to_input_weights_ptr, - float recurrent_to_input_weights_scale, - const int8_t* recurrent_to_forget_weights_ptr, - float recurrent_to_forget_weights_scale, - const int8_t* recurrent_to_cell_weights_ptr, - float recurrent_to_cell_weights_scale, - const int8_t* recurrent_to_output_weights_ptr, - float recurrent_to_output_weights_scale, - const int8_t* cell_to_input_weights_ptr, float cell_to_input_weights_scale, - const int8_t* cell_to_forget_weights_ptr, - float cell_to_forget_weights_scale, - const int8_t* cell_to_output_weights_ptr, - float cell_to_output_weights_scale, const float* input_gate_bias_ptr, - const float* forget_gate_bias_ptr, const float* cell_bias_ptr, - const float* output_gate_bias_ptr, const int8_t* projection_weights_ptr, - float projection_weights_scale, const float* projection_bias_ptr, - const TfLiteLSTMParams* params, int n_batch, int n_cell, int n_input, - int n_aux_input, int n_output, int output_batch_leading_dim, - float* input_gate_scratch, float* forget_gate_scratch, float* cell_scratch, - float* output_gate_scratch, float* scaling_factors, - float* product_scaling_factors, float* recovered_cell_weights, - int8_t* quantized_input_ptr_batch, int8_t* quantized_aux_input_ptr_batch, - int8_t* quantized_output_state_ptr, int8_t* quantized_cell_state_ptr, - float* output_state_ptr, float* cell_state_ptr, float* output_ptr_batch) { - // Since we have already checked that weights are all there or none, we - // can check the existense of only one to the get the condition. - const bool use_cifg = (input_to_input_weights_ptr == nullptr); - const bool use_peephole = (cell_to_output_weights_ptr != nullptr); - // Initialize scratch buffers with bias. - if (!use_cifg) { - tensor_utils::VectorBatchVectorAssign(input_gate_bias_ptr, n_cell, n_batch, - input_gate_scratch); - } - tensor_utils::VectorBatchVectorAssign(forget_gate_bias_ptr, n_cell, n_batch, - forget_gate_scratch); - tensor_utils::VectorBatchVectorAssign(cell_bias_ptr, n_cell, n_batch, - cell_scratch); - tensor_utils::VectorBatchVectorAssign(output_gate_bias_ptr, n_cell, n_batch, - output_gate_scratch); - - if (!tensor_utils::IsZeroVector(input_ptr_batch, n_batch * n_input)) { - // Save quantization and matmul computation for all zero input. - float unused_min, unused_max; - for (int b = 0; b < n_batch; ++b) { - const int offset = b * n_input; - tensor_utils::SymmetricQuantizeFloats( - input_ptr_batch + offset, n_input, quantized_input_ptr_batch + offset, - &unused_min, &unused_max, &scaling_factors[b]); - } - // For each batch and cell: compute input_weight * input. - if (!use_cifg) { - for (int b = 0; b < n_batch; ++b) { - product_scaling_factors[b] = - scaling_factors[b] * input_to_input_weights_scale; - } - tensor_utils::MatrixBatchVectorMultiplyAccumulate( - input_to_input_weights_ptr, n_cell, n_input, - quantized_input_ptr_batch, product_scaling_factors, n_batch, - input_gate_scratch, /*result_stride=*/1); - } - - for (int b = 0; b < n_batch; ++b) { - product_scaling_factors[b] = - scaling_factors[b] * input_to_forget_weights_scale; - } - tensor_utils::MatrixBatchVectorMultiplyAccumulate( - input_to_forget_weights_ptr, n_cell, n_input, quantized_input_ptr_batch, - product_scaling_factors, n_batch, forget_gate_scratch, - /*result_stride=*/1); - - for (int b = 0; b < n_batch; ++b) { - product_scaling_factors[b] = - scaling_factors[b] * input_to_cell_weights_scale; - } - tensor_utils::MatrixBatchVectorMultiplyAccumulate( - input_to_cell_weights_ptr, n_cell, n_input, quantized_input_ptr_batch, - product_scaling_factors, n_batch, cell_scratch, /*result_stride=*/1); - - for (int b = 0; b < n_batch; ++b) { - product_scaling_factors[b] = - scaling_factors[b] * input_to_output_weights_scale; - } - tensor_utils::MatrixBatchVectorMultiplyAccumulate( - input_to_output_weights_ptr, n_cell, n_input, quantized_input_ptr_batch, - product_scaling_factors, n_batch, output_gate_scratch, - /*result_stride=*/1); - } - - if (aux_input_ptr_batch != nullptr && - !tensor_utils::IsZeroVector(aux_input_ptr_batch, n_batch * n_input)) { - // Save quantization and matmul computation for all zero input. - float unused_min, unused_max; - for (int b = 0; b < n_batch; ++b) { - const int offset = b * n_input; - tensor_utils::SymmetricQuantizeFloats( - aux_input_ptr_batch + offset, n_input, - quantized_aux_input_ptr_batch + offset, &unused_min, &unused_max, - &scaling_factors[b]); - } - // For each batch and cell: compute input_weight * input. - if (!use_cifg) { - for (int b = 0; b < n_batch; ++b) { - product_scaling_factors[b] = - scaling_factors[b] * aux_input_to_input_weights_scale; - } - tensor_utils::MatrixBatchVectorMultiplyAccumulate( - aux_input_to_input_weights_ptr, n_cell, n_input, - quantized_aux_input_ptr_batch, product_scaling_factors, n_batch, - input_gate_scratch, /*result_stride=*/1); - } - - for (int b = 0; b < n_batch; ++b) { - product_scaling_factors[b] = - scaling_factors[b] * aux_input_to_forget_weights_scale; - } - tensor_utils::MatrixBatchVectorMultiplyAccumulate( - aux_input_to_forget_weights_ptr, n_cell, n_input, - quantized_aux_input_ptr_batch, product_scaling_factors, n_batch, - forget_gate_scratch, /*result_stride=*/1); - - for (int b = 0; b < n_batch; ++b) { - product_scaling_factors[b] = - scaling_factors[b] * aux_input_to_cell_weights_scale; - } - tensor_utils::MatrixBatchVectorMultiplyAccumulate( - aux_input_to_cell_weights_ptr, n_cell, n_input, - quantized_aux_input_ptr_batch, product_scaling_factors, n_batch, - cell_scratch, /*result_stride=*/1); - - for (int b = 0; b < n_batch; ++b) { - product_scaling_factors[b] = - scaling_factors[b] * aux_input_to_output_weights_scale; - } - tensor_utils::MatrixBatchVectorMultiplyAccumulate( - aux_input_to_output_weights_ptr, n_cell, n_input, - quantized_aux_input_ptr_batch, product_scaling_factors, n_batch, - output_gate_scratch, /*result_stride=*/1); - } - - if (!tensor_utils::IsZeroVector(output_state_ptr, n_batch * n_output)) { - // Save quantization and matmul computation for all zero input. - float unused_min, unused_max; - for (int b = 0; b < n_batch; ++b) { - const int offset = b * n_output; - tensor_utils::SymmetricQuantizeFloats(output_state_ptr + offset, n_output, - quantized_output_state_ptr + offset, - &unused_min, &unused_max, - &scaling_factors[b]); - } - // For each batch and cell: compute recurrent_weight * output_state. - if (!use_cifg) { - for (int b = 0; b < n_batch; ++b) { - product_scaling_factors[b] = - scaling_factors[b] * recurrent_to_input_weights_scale; - } - tensor_utils::MatrixBatchVectorMultiplyAccumulate( - recurrent_to_input_weights_ptr, n_cell, n_output, - quantized_output_state_ptr, product_scaling_factors, n_batch, - input_gate_scratch, /*result_stride=*/1); - } - - for (int b = 0; b < n_batch; ++b) { - product_scaling_factors[b] = - scaling_factors[b] * recurrent_to_forget_weights_scale; - } - tensor_utils::MatrixBatchVectorMultiplyAccumulate( - recurrent_to_forget_weights_ptr, n_cell, n_output, - quantized_output_state_ptr, product_scaling_factors, n_batch, - forget_gate_scratch, /*result_stride=*/1); - - for (int b = 0; b < n_batch; ++b) { - product_scaling_factors[b] = - scaling_factors[b] * recurrent_to_cell_weights_scale; - } - tensor_utils::MatrixBatchVectorMultiplyAccumulate( - recurrent_to_cell_weights_ptr, n_cell, n_output, - quantized_output_state_ptr, product_scaling_factors, n_batch, - cell_scratch, /*result_stride=*/1); - - for (int b = 0; b < n_batch; ++b) { - product_scaling_factors[b] = - scaling_factors[b] * recurrent_to_output_weights_scale; - } - tensor_utils::MatrixBatchVectorMultiplyAccumulate( - recurrent_to_output_weights_ptr, n_cell, n_output, - quantized_output_state_ptr, product_scaling_factors, n_batch, - output_gate_scratch, /*result_stride=*/1); - } - - // Save quantization and matmul computation for all zero input. - bool is_cell_state_all_zeros = - tensor_utils::IsZeroVector(cell_state_ptr, n_batch * n_cell); - - // For each batch and cell: update input gate. - if (!use_cifg) { - if (use_peephole && !is_cell_state_all_zeros) { - tensor_utils::VectorScalarMultiply(cell_to_input_weights_ptr, n_cell, - cell_to_input_weights_scale, - recovered_cell_weights); - tensor_utils::VectorBatchVectorCwiseProductAccumulate( - recovered_cell_weights, n_cell, cell_state_ptr, n_batch, - input_gate_scratch); - } - tensor_utils::ApplySigmoidToVector(input_gate_scratch, n_cell * n_batch, - input_gate_scratch); - } - - // For each batch and cell: update forget gate. - if (use_peephole && !is_cell_state_all_zeros) { - tensor_utils::VectorScalarMultiply(cell_to_forget_weights_ptr, n_cell, - cell_to_forget_weights_scale, - recovered_cell_weights); - tensor_utils::VectorBatchVectorCwiseProductAccumulate( - recovered_cell_weights, n_cell, cell_state_ptr, n_batch, - forget_gate_scratch); - } - tensor_utils::ApplySigmoidToVector(forget_gate_scratch, n_cell * n_batch, - forget_gate_scratch); - - // For each batch and cell: update the cell. - tensor_utils::VectorVectorCwiseProduct(forget_gate_scratch, cell_state_ptr, - n_batch * n_cell, cell_state_ptr); - tensor_utils::ApplyActivationToVector(cell_scratch, n_batch * n_cell, - params->activation, cell_scratch); - if (use_cifg) { - tensor_utils::Sub1Vector(forget_gate_scratch, n_batch * n_cell, - forget_gate_scratch); - tensor_utils::VectorVectorCwiseProductAccumulate( - cell_scratch, forget_gate_scratch, n_batch * n_cell, cell_state_ptr); - } else { - tensor_utils::VectorVectorCwiseProductAccumulate( - cell_scratch, input_gate_scratch, n_batch * n_cell, cell_state_ptr); - } - if (params->cell_clip > 0.0) { - tensor_utils::ClipVector(cell_state_ptr, n_batch * n_cell, - params->cell_clip, cell_state_ptr); - } - - is_cell_state_all_zeros = - tensor_utils::IsZeroVector(cell_state_ptr, n_batch * n_cell); - // For each batch and cell: update the output gate. - if (use_peephole && !is_cell_state_all_zeros) { - tensor_utils::VectorScalarMultiply(cell_to_output_weights_ptr, n_cell, - cell_to_output_weights_scale, - recovered_cell_weights); - tensor_utils::VectorBatchVectorCwiseProductAccumulate( - recovered_cell_weights, n_cell, cell_state_ptr, n_batch, - output_gate_scratch); - } - tensor_utils::ApplySigmoidToVector(output_gate_scratch, n_batch * n_cell, - output_gate_scratch); - tensor_utils::ApplyActivationToVector(cell_state_ptr, n_batch * n_cell, - params->activation, cell_scratch); - tensor_utils::VectorVectorCwiseProduct(output_gate_scratch, cell_scratch, - n_batch * n_cell, output_gate_scratch); - - const bool use_projection_weight = (projection_weights_ptr != nullptr); - const bool use_projection_bias = (projection_bias_ptr != nullptr); - - // For each batch: update the projection and output_state. Note that since - // the output batch rows may not be contiguous (output_batch_leading_dim != - // n_output), we unroll the batched operations where this is the case. - if (output_batch_leading_dim == n_output) { - if (use_projection_weight) { - if (use_projection_bias) { - tensor_utils::VectorBatchVectorAssign(projection_bias_ptr, n_output, - n_batch, output_ptr_batch); - } else { - tensor_utils::ZeroVector(output_ptr_batch, n_batch * n_output); - } - if (!tensor_utils::IsZeroVector(output_gate_scratch, n_batch * n_cell)) { - // Save quantization and matmul computation for all zero input. - float unused_min, unused_max; - for (int b = 0; b < n_batch; ++b) { - const int offset = b * n_cell; - tensor_utils::SymmetricQuantizeFloats( - output_gate_scratch + offset, n_cell, - quantized_cell_state_ptr + offset, &unused_min, &unused_max, - &scaling_factors[b]); - } - for (int b = 0; b < n_batch; ++b) { - product_scaling_factors[b] = - scaling_factors[b] * projection_weights_scale; - } - tensor_utils::MatrixBatchVectorMultiplyAccumulate( - projection_weights_ptr, n_output, n_cell, quantized_cell_state_ptr, - product_scaling_factors, n_batch, output_ptr_batch, - /*result_stride=*/1); - } - if (params->proj_clip > 0.0) { - tensor_utils::ClipVector(output_ptr_batch, n_batch * n_output, - params->proj_clip, output_ptr_batch); - } - } else { - tensor_utils::CopyVector(output_gate_scratch, n_batch * n_output, - output_ptr_batch); - } - tensor_utils::CopyVector(output_ptr_batch, n_batch * n_output, - output_state_ptr); - } else { - if (use_projection_weight) { - if (use_projection_bias) { - for (int k = 0; k < n_batch; k++) { - tensor_utils::CopyVector( - projection_bias_ptr, n_output, - output_ptr_batch + k * output_batch_leading_dim); - } - } else { - for (int k = 0; k < n_batch; k++) { - tensor_utils::ZeroVector( - output_ptr_batch + k * output_batch_leading_dim, n_output); - } - } - if (!tensor_utils::IsZeroVector(output_gate_scratch, n_batch * n_cell)) { - // Save quantization and matmul computation for all zero input. - float unused_min, unused_max; - for (int b = 0; b < n_batch; ++b) { - const int offset = b * n_cell; - tensor_utils::SymmetricQuantizeFloats( - output_gate_scratch + offset, n_cell, - quantized_cell_state_ptr + offset, &unused_min, &unused_max, - &scaling_factors[b]); - } - for (int b = 0; b < n_batch; ++b) { - product_scaling_factors[b] = - scaling_factors[b] * projection_weights_scale; - } - for (int k = 0; k < n_batch; k++) { - tensor_utils::MatrixBatchVectorMultiplyAccumulate( - projection_weights_ptr, n_output, n_cell, - quantized_cell_state_ptr + k * n_cell, - &product_scaling_factors[k], - /*n_batch=*/1, output_ptr_batch + k * output_batch_leading_dim, - /*result_stride=*/1); - } - } - if (params->proj_clip > 0.0) { - for (int k = 0; k < n_batch; k++) { - tensor_utils::ClipVector( - output_ptr_batch + k * output_batch_leading_dim, n_output, - params->proj_clip, - output_ptr_batch + k * output_batch_leading_dim); - } - } - } else { - for (int k = 0; k < n_batch; k++) { - tensor_utils::CopyVector( - output_gate_scratch + k * n_output, n_output, - output_ptr_batch + k * output_batch_leading_dim); - } - } - for (int k = 0; k < n_batch; k++) { - tensor_utils::CopyVector(output_ptr_batch + k * output_batch_leading_dim, - n_output, output_state_ptr + k * n_output); - } - } -} -} // namespace - -TfLiteStatus EvalFloat( - const TfLiteTensor* input, const TfLiteTensor* input_to_input_weights, - const TfLiteTensor* input_to_forget_weights, - const TfLiteTensor* input_to_cell_weights, - const TfLiteTensor* input_to_output_weights, - const TfLiteTensor* recurrent_to_input_weights, - const TfLiteTensor* recurrent_to_forget_weights, - const TfLiteTensor* recurrent_to_cell_weights, - const TfLiteTensor* recurrent_to_output_weights, - const TfLiteTensor* cell_to_input_weights, - const TfLiteTensor* cell_to_forget_weights, - const TfLiteTensor* cell_to_output_weights, const TfLiteTensor* aux_input, - const TfLiteTensor* aux_input_to_input_weights, - const TfLiteTensor* aux_input_to_forget_weights, - const TfLiteTensor* aux_input_to_cell_weights, - const TfLiteTensor* aux_input_to_output_weights, - const TfLiteTensor* input_gate_bias, const TfLiteTensor* forget_gate_bias, - const TfLiteTensor* cell_bias, const TfLiteTensor* output_gate_bias, - const TfLiteTensor* projection_weights, const TfLiteTensor* projection_bias, - const TfLiteLSTMParams* params, bool forward_sequence, bool time_major, - int output_offset, TfLiteTensor* scratch_buffer, - TfLiteTensor* activation_state, TfLiteTensor* cell_state, - TfLiteTensor* output) { - TF_LITE_ASSERT(input->dims->size >= 2 && input->dims->size <= 3); - int max_time, n_batch; - if (input->dims->size == 3) { - max_time = (time_major) ? input->dims->data[0] : input->dims->data[1]; - n_batch = (time_major) ? input->dims->data[1] : input->dims->data[0]; - } else { - max_time = 1; - n_batch = input->dims->data[0]; - } - const int n_input = input->dims->data[input->dims->size - 1]; - const int aux_input_size = - (aux_input) ? aux_input->dims->data[aux_input->dims->size - 1] : 0; - - // n_cell and n_output will be the same size when there is no projection. - const int n_cell = input_to_output_weights->dims->data[0]; - const int n_output = recurrent_to_output_weights->dims->data[1]; - - // Since we have already checked that weights are all there or none, we can - // check the existense of only one to the get the condition. - const bool use_cifg = (input_to_input_weights == nullptr); - const bool use_peephole = (cell_to_output_weights != nullptr); - - // Index the scratch buffers pointers to the global scratch buffer. - float* input_gate_scratch = nullptr; - float* cell_scratch = nullptr; - float* forget_gate_scratch = nullptr; - float* output_gate_scratch = nullptr; - if (use_cifg) { - cell_scratch = scratch_buffer->data.f; - forget_gate_scratch = scratch_buffer->data.f + n_cell * n_batch; - output_gate_scratch = scratch_buffer->data.f + 2 * n_cell * n_batch; - } else { - input_gate_scratch = scratch_buffer->data.f; - cell_scratch = scratch_buffer->data.f + n_cell * n_batch; - forget_gate_scratch = scratch_buffer->data.f + 2 * n_cell * n_batch; - output_gate_scratch = scratch_buffer->data.f + 3 * n_cell * n_batch; - } - - // Check optional tensors, the respective pointers can be null. - const float* input_to_input_weights_ptr = - (use_cifg) ? nullptr : input_to_input_weights->data.f; - const float* recurrent_to_input_weights_ptr = - (use_cifg) ? nullptr : recurrent_to_input_weights->data.f; - const float* input_gate_bias_ptr = - (use_cifg) ? nullptr : input_gate_bias->data.f; - const float* cell_to_input_weights_ptr = - (use_peephole && !use_cifg) ? cell_to_input_weights->data.f : nullptr; - const float* cell_to_forget_weights_ptr = - (use_peephole) ? cell_to_forget_weights->data.f : nullptr; - const float* cell_to_output_weights_ptr = - (use_peephole) ? cell_to_output_weights->data.f : nullptr; - const float* projection_weights_ptr = - (projection_weights == nullptr) ? nullptr : projection_weights->data.f; - const float* projection_bias_ptr = - (projection_bias == nullptr) ? nullptr : projection_bias->data.f; - - float* aux_input_ptr = nullptr; - float* aux_input_to_input_weights_ptr = nullptr; - float* aux_input_to_forget_weights_ptr = nullptr; - float* aux_input_to_cell_weights_ptr = nullptr; - float* aux_input_to_output_weights_ptr = nullptr; - if (aux_input_size > 0) { - if (!use_cifg) { - aux_input_to_input_weights_ptr = aux_input_to_input_weights->data.f; - } - aux_input_to_forget_weights_ptr = aux_input_to_forget_weights->data.f; - aux_input_to_cell_weights_ptr = aux_input_to_cell_weights->data.f; - aux_input_to_output_weights_ptr = aux_input_to_output_weights->data.f; - } - - const int output_batch_leading_dim = - output->dims->data[output->dims->size - 1]; - if (time_major) { - // Loop through the sequence. - const int input_step = n_batch * n_input; - const int output_step = n_batch * output_batch_leading_dim; - for (int t = 0; t < max_time; t++) { - // If this is the forward_sequence, step forward, otherwise step - // backwards. - const int t_rel = forward_sequence ? t : max_time - t - 1; - const float* input_ptr = input->data.f + t_rel * input_step; - if (aux_input) { - aux_input_ptr = aux_input->data.f + t_rel * input_step; - } - float* output_ptr_time = - output->data.f + t_rel * output_step + output_offset; - - LstmStepWithAuxInput( - input_ptr, input_to_input_weights_ptr, - input_to_forget_weights->data.f, input_to_cell_weights->data.f, - input_to_output_weights->data.f, aux_input_ptr, - aux_input_to_input_weights_ptr, aux_input_to_forget_weights_ptr, - aux_input_to_cell_weights_ptr, aux_input_to_output_weights_ptr, - recurrent_to_input_weights_ptr, recurrent_to_forget_weights->data.f, - recurrent_to_cell_weights->data.f, - recurrent_to_output_weights->data.f, cell_to_input_weights_ptr, - cell_to_forget_weights_ptr, cell_to_output_weights_ptr, - input_gate_bias_ptr, forget_gate_bias->data.f, cell_bias->data.f, - output_gate_bias->data.f, projection_weights_ptr, projection_bias_ptr, - params, n_batch, n_cell, n_input, aux_input_size, n_output, - output_batch_leading_dim, activation_state->data.f, - cell_state->data.f, input_gate_scratch, forget_gate_scratch, - cell_scratch, output_gate_scratch, output_ptr_time); - } - } else { - for (int b = 0; b < n_batch; b++) { - const int input_step = n_input; - const int output_step = output_batch_leading_dim; - for (int t = 0; t < max_time; t++) { - // If this is the forward_sequence, step forward, otherwise step - // backwards. - const int t_rel = forward_sequence ? t : max_time - t - 1; - const float* input_ptr = input->data.f + t_rel * input_step; - if (aux_input) { - aux_input_ptr = aux_input->data.f + t_rel * input_step; - } - float* output_ptr_time = - output->data.f + t_rel * output_step + output_offset; - - LstmStepWithAuxInput( - input_ptr, input_to_input_weights_ptr, - input_to_forget_weights->data.f, input_to_cell_weights->data.f, - input_to_output_weights->data.f, aux_input_ptr, - aux_input_to_input_weights_ptr, aux_input_to_forget_weights_ptr, - aux_input_to_cell_weights_ptr, aux_input_to_output_weights_ptr, - recurrent_to_input_weights_ptr, recurrent_to_forget_weights->data.f, - recurrent_to_cell_weights->data.f, - recurrent_to_output_weights->data.f, cell_to_input_weights_ptr, - cell_to_forget_weights_ptr, cell_to_output_weights_ptr, - input_gate_bias_ptr, forget_gate_bias->data.f, cell_bias->data.f, - output_gate_bias->data.f, projection_weights_ptr, - projection_bias_ptr, params, /*n_batch=*/1, n_cell, n_input, - aux_input_size, n_output, output_batch_leading_dim, - activation_state->data.f, cell_state->data.f, input_gate_scratch, - forget_gate_scratch, cell_scratch, output_gate_scratch, - output_ptr_time); - } - } - } - return kTfLiteOk; -} - -TfLiteStatus EvalHybrid( - const TfLiteTensor* input, const TfLiteTensor* input_to_input_weights, - const TfLiteTensor* input_to_forget_weights, - const TfLiteTensor* input_to_cell_weights, - const TfLiteTensor* input_to_output_weights, - const TfLiteTensor* recurrent_to_input_weights, - const TfLiteTensor* recurrent_to_forget_weights, - const TfLiteTensor* recurrent_to_cell_weights, - const TfLiteTensor* recurrent_to_output_weights, - const TfLiteTensor* cell_to_input_weights, - const TfLiteTensor* cell_to_forget_weights, - const TfLiteTensor* cell_to_output_weights, const TfLiteTensor* aux_input, - const TfLiteTensor* aux_input_to_input_weights, - const TfLiteTensor* aux_input_to_forget_weights, - const TfLiteTensor* aux_input_to_cell_weights, - const TfLiteTensor* aux_input_to_output_weights, - const TfLiteTensor* input_gate_bias, const TfLiteTensor* forget_gate_bias, - const TfLiteTensor* cell_bias, const TfLiteTensor* output_gate_bias, - const TfLiteTensor* projection_weights, const TfLiteTensor* projection_bias, - const TfLiteLSTMParams* params, bool forward_sequence, bool time_major, - int output_offset, TfLiteTensor* scratch_buffer, - TfLiteTensor* scaling_factors, TfLiteTensor* prod_scaling_factors, - TfLiteTensor* recovered_cell_weights, TfLiteTensor* input_quantized, - TfLiteTensor* aux_input_quantized, TfLiteTensor* output_state_quantized, - TfLiteTensor* cell_state_quantized, TfLiteTensor* output_state, - TfLiteTensor* cell_state, TfLiteTensor* output) { - TF_LITE_ASSERT(input->dims->size >= 2 && input->dims->size <= 3); - const int n_input = input->dims->data[input->dims->size - 1]; - int max_time, n_batch; - if (input->dims->size == 2) { - max_time = 1; - n_batch = input->dims->data[0]; - } else { - max_time = (time_major) ? input->dims->data[0] : input->dims->data[1]; - n_batch = (time_major) ? input->dims->data[1] : input->dims->data[0]; - } - const int aux_input_size = - (aux_input) ? aux_input->dims->data[aux_input->dims->size - 1] : 0; - // n_cell and n_output will be the same size when there is no projection. - const int n_cell = input_to_output_weights->dims->data[0]; - const int n_output = recurrent_to_output_weights->dims->data[1]; - - // Since we have already checked that weights are all there or none, we can - // check the existence of only one to get the condition. - const bool use_cifg = (input_to_input_weights == nullptr); - const bool use_peephole = (cell_to_output_weights != nullptr); - - float* input_gate_scratch = nullptr; - float* cell_scratch = nullptr; - float* forget_gate_scratch = nullptr; - float* output_gate_scratch = nullptr; - if (use_cifg) { - cell_scratch = scratch_buffer->data.f; - forget_gate_scratch = scratch_buffer->data.f + n_cell * n_batch; - output_gate_scratch = scratch_buffer->data.f + 2 * n_cell * n_batch; - } else { - input_gate_scratch = scratch_buffer->data.f; - cell_scratch = scratch_buffer->data.f + n_cell * n_batch; - forget_gate_scratch = scratch_buffer->data.f + 2 * n_cell * n_batch; - output_gate_scratch = scratch_buffer->data.f + 3 * n_cell * n_batch; - } - - // Check optional tensors, the respective pointers can be null. - int8_t* input_to_input_weights_ptr = nullptr; - float input_to_input_weights_scale = 1.0f; - int8_t* recurrent_to_input_weights_ptr = nullptr; - float recurrent_to_input_weights_scale = 1.0f; - float* input_gate_bias_ptr = nullptr; - if (!use_cifg) { - input_to_input_weights_ptr = - reinterpret_cast(input_to_input_weights->data.uint8); - recurrent_to_input_weights_ptr = - reinterpret_cast(recurrent_to_input_weights->data.uint8); - input_gate_bias_ptr = input_gate_bias->data.f; - input_to_input_weights_scale = input_to_input_weights->params.scale; - recurrent_to_input_weights_scale = recurrent_to_input_weights->params.scale; - } - - int8_t* cell_to_input_weights_ptr = nullptr; - int8_t* cell_to_forget_weights_ptr = nullptr; - int8_t* cell_to_output_weights_ptr = nullptr; - float cell_to_input_weights_scale = 1.0f; - float cell_to_forget_weights_scale = 1.0f; - float cell_to_output_weights_scale = 1.0f; - if (use_peephole) { - if (!use_cifg) { - cell_to_input_weights_ptr = - reinterpret_cast(cell_to_input_weights->data.uint8); - cell_to_input_weights_scale = cell_to_input_weights->params.scale; - } - cell_to_forget_weights_ptr = - reinterpret_cast(cell_to_forget_weights->data.uint8); - cell_to_output_weights_ptr = - reinterpret_cast(cell_to_output_weights->data.uint8); - cell_to_forget_weights_scale = cell_to_forget_weights->params.scale; - cell_to_output_weights_scale = cell_to_output_weights->params.scale; - } - - const int8_t* projection_weights_ptr = - (projection_weights == nullptr) - ? nullptr - : reinterpret_cast(projection_weights->data.uint8); - const float projection_weights_scale = - (projection_weights == nullptr) ? 1.0f : projection_weights->params.scale; - const float* projection_bias_ptr = - (projection_bias == nullptr) ? nullptr : projection_bias->data.f; - - // Required tensors, pointers are non-null. - const int8_t* input_to_forget_weights_ptr = - reinterpret_cast(input_to_forget_weights->data.uint8); - const float input_to_forget_weights_scale = - input_to_forget_weights->params.scale; - const int8_t* input_to_cell_weights_ptr = - reinterpret_cast(input_to_cell_weights->data.uint8); - const float input_to_cell_weights_scale = input_to_cell_weights->params.scale; - const int8_t* input_to_output_weights_ptr = - reinterpret_cast(input_to_output_weights->data.uint8); - const float input_to_output_weights_scale = - input_to_output_weights->params.scale; - const int8_t* recurrent_to_forget_weights_ptr = - reinterpret_cast(recurrent_to_forget_weights->data.uint8); - const float recurrent_to_forget_weights_scale = - recurrent_to_forget_weights->params.scale; - const int8_t* recurrent_to_cell_weights_ptr = - reinterpret_cast(recurrent_to_cell_weights->data.uint8); - const float recurrent_to_cell_weights_scale = - recurrent_to_cell_weights->params.scale; - const int8_t* recurrent_to_output_weights_ptr = - reinterpret_cast(recurrent_to_output_weights->data.uint8); - const float recurrent_to_output_weights_scale = - recurrent_to_output_weights->params.scale; - const float* forget_gate_bias_ptr = forget_gate_bias->data.f; - const float* cell_bias_ptr = cell_bias->data.f; - const float* output_gate_bias_ptr = output_gate_bias->data.f; - - float* output_state_ptr = output_state->data.f; - float* cell_state_ptr = cell_state->data.f; - - // Temporary storage for quantized values and scaling factors. - int8_t* quantized_input_ptr = - reinterpret_cast(input_quantized->data.uint8); - int8_t* quantized_aux_input_ptr = - (aux_input_quantized == nullptr) - ? nullptr - : reinterpret_cast(aux_input_quantized->data.uint8); - int8_t* quantized_output_state_ptr = - reinterpret_cast(output_state_quantized->data.uint8); - int8_t* quantized_cell_state_ptr = - reinterpret_cast(cell_state_quantized->data.uint8); - float* scaling_factors_ptr = scaling_factors->data.f; - float* prod_scaling_factors_ptr = prod_scaling_factors->data.f; - float* recovered_cell_weights_ptr = recovered_cell_weights->data.f; - - // Auxiliary input and weights. - float* aux_input_ptr = nullptr; - int8_t* aux_input_to_input_weights_ptr = nullptr; - int8_t* aux_input_to_forget_weights_ptr = nullptr; - int8_t* aux_input_to_cell_weights_ptr = nullptr; - int8_t* aux_input_to_output_weights_ptr = nullptr; - float aux_input_to_input_weights_scale = 0.0f; - float aux_input_to_forget_weights_scale = 0.0f; - float aux_input_to_cell_weights_scale = 0.0f; - float aux_input_to_output_weights_scale = 0.0f; - if (aux_input_size > 0) { - if (!use_cifg) { - aux_input_to_input_weights_ptr = - reinterpret_cast(aux_input_to_input_weights->data.uint8); - } - aux_input_to_forget_weights_ptr = - reinterpret_cast(aux_input_to_forget_weights->data.uint8); - aux_input_to_cell_weights_ptr = - reinterpret_cast(aux_input_to_cell_weights->data.uint8); - aux_input_to_output_weights_ptr = - reinterpret_cast(aux_input_to_output_weights->data.uint8); - if (!use_cifg) { - aux_input_to_input_weights_scale = - aux_input_to_input_weights->params.scale; - } - aux_input_to_forget_weights_scale = - aux_input_to_forget_weights->params.scale; - aux_input_to_cell_weights_scale = aux_input_to_cell_weights->params.scale; - aux_input_to_output_weights_scale = - aux_input_to_output_weights->params.scale; - } - - const int output_batch_leading_dim = - output->dims->data[output->dims->size - 1]; - if (time_major) { - // Feed the sequence into the LSTM step-by-step. - const int input_step = n_batch * n_input; - const int output_step = n_batch * output_batch_leading_dim; - for (int t = 0; t < max_time; t++) { - // If this is the forward_sequence, step forward, otherwise step - // backwards. - const int t_rel = forward_sequence ? t : max_time - t - 1; - const float* input_ptr = input->data.f + t_rel * input_step; - if (aux_input) { - aux_input_ptr = aux_input->data.f + t_rel * input_step; - } - float* output_ptr = output->data.f + t_rel * output_step + output_offset; - - LstmStepWithAuxInput( - input_ptr, input_to_input_weights_ptr, input_to_input_weights_scale, - input_to_forget_weights_ptr, input_to_forget_weights_scale, - input_to_cell_weights_ptr, input_to_cell_weights_scale, - input_to_output_weights_ptr, input_to_output_weights_scale, - aux_input_ptr, aux_input_to_input_weights_ptr, - aux_input_to_input_weights_scale, aux_input_to_forget_weights_ptr, - aux_input_to_forget_weights_scale, aux_input_to_cell_weights_ptr, - aux_input_to_cell_weights_scale, aux_input_to_output_weights_ptr, - aux_input_to_output_weights_scale, recurrent_to_input_weights_ptr, - recurrent_to_input_weights_scale, recurrent_to_forget_weights_ptr, - recurrent_to_forget_weights_scale, recurrent_to_cell_weights_ptr, - recurrent_to_cell_weights_scale, recurrent_to_output_weights_ptr, - recurrent_to_output_weights_scale, cell_to_input_weights_ptr, - cell_to_input_weights_scale, cell_to_forget_weights_ptr, - cell_to_forget_weights_scale, cell_to_output_weights_ptr, - cell_to_output_weights_scale, input_gate_bias_ptr, - forget_gate_bias_ptr, cell_bias_ptr, output_gate_bias_ptr, - projection_weights_ptr, projection_weights_scale, projection_bias_ptr, - params, n_batch, n_cell, n_input, aux_input_size, n_output, - output_batch_leading_dim, input_gate_scratch, forget_gate_scratch, - cell_scratch, output_gate_scratch, scaling_factors_ptr, - prod_scaling_factors_ptr, recovered_cell_weights_ptr, - quantized_input_ptr, quantized_aux_input_ptr, - quantized_output_state_ptr, quantized_cell_state_ptr, - output_state_ptr, cell_state_ptr, output_ptr); - } - } else { - for (int b = 0; b < n_batch; b++) { - const int input_step = n_input; - const int output_step = output_batch_leading_dim; - for (int t = 0; t < max_time; t++) { - // If this is the forward_sequence, step forward, otherwise step - // backwards. - const int t_rel = forward_sequence ? t : max_time - t - 1; - const float* input_ptr = input->data.f + t_rel * input_step; - if (aux_input) { - aux_input_ptr = aux_input->data.f + t_rel * input_step; - } - float* output_ptr = - output->data.f + t_rel * output_step + output_offset; - - LstmStepWithAuxInput( - input_ptr, input_to_input_weights_ptr, input_to_input_weights_scale, - input_to_forget_weights_ptr, input_to_forget_weights_scale, - input_to_cell_weights_ptr, input_to_cell_weights_scale, - input_to_output_weights_ptr, input_to_output_weights_scale, - aux_input_ptr, aux_input_to_input_weights_ptr, - aux_input_to_input_weights_scale, aux_input_to_forget_weights_ptr, - aux_input_to_forget_weights_scale, aux_input_to_cell_weights_ptr, - aux_input_to_cell_weights_scale, aux_input_to_output_weights_ptr, - aux_input_to_output_weights_scale, recurrent_to_input_weights_ptr, - recurrent_to_input_weights_scale, recurrent_to_forget_weights_ptr, - recurrent_to_forget_weights_scale, recurrent_to_cell_weights_ptr, - recurrent_to_cell_weights_scale, recurrent_to_output_weights_ptr, - recurrent_to_output_weights_scale, cell_to_input_weights_ptr, - cell_to_input_weights_scale, cell_to_forget_weights_ptr, - cell_to_forget_weights_scale, cell_to_output_weights_ptr, - cell_to_output_weights_scale, input_gate_bias_ptr, - forget_gate_bias_ptr, cell_bias_ptr, output_gate_bias_ptr, - projection_weights_ptr, projection_weights_scale, - projection_bias_ptr, params, n_batch, n_cell, n_input, - aux_input_size, n_output, output_batch_leading_dim, - input_gate_scratch, forget_gate_scratch, cell_scratch, - output_gate_scratch, scaling_factors_ptr, prod_scaling_factors_ptr, - recovered_cell_weights_ptr, quantized_input_ptr, - quantized_aux_input_ptr, quantized_output_state_ptr, - quantized_cell_state_ptr, output_state_ptr, cell_state_ptr, - output_ptr); - } - } - } - - return kTfLiteOk; -} - -} // namespace lstm_eval -} // namespace builtin -} // namespace ops -} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/lstm_test.cc b/tensorflow/contrib/lite/kernels/lstm_test.cc deleted file mode 100644 index f8947db724217421aba637ae884bedb5d412b2db..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/lstm_test.cc +++ /dev/null @@ -1,1402 +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. -==============================================================================*/ -// Unit test for TFLite LSTM op. -// -// TODO(alanchiao): add unit test with invalid input dimensions for this and its -// variants. - -#include -#include - -#include -#include -#include "tensorflow/contrib/lite/interpreter.h" -#include "tensorflow/contrib/lite/kernels/register.h" -#include "tensorflow/contrib/lite/kernels/test_util.h" -#include "tensorflow/contrib/lite/model.h" - -namespace tflite { -namespace { - -using ::testing::ElementsAreArray; - -class LSTMOpModel : public SingleOpModel { - public: - LSTMOpModel(int n_batch, int n_input, int n_cell, int n_output, bool use_cifg, - bool use_peephole, bool use_projection_weights, - bool use_projection_bias, float cell_clip, float proj_clip, - const std::vector>& input_shapes, - const TensorType& weight_type = TensorType_FLOAT32) - : n_batch_(n_batch), - n_input_(n_input), - n_cell_(n_cell), - n_output_(n_output) { - input_ = AddInput(TensorType_FLOAT32); - - if (use_cifg) { - input_to_input_weights_ = AddNullInput(); - } else { - input_to_input_weights_ = AddInput(weight_type); - } - - input_to_forget_weights_ = AddInput(weight_type); - input_to_cell_weights_ = AddInput(weight_type); - input_to_output_weights_ = AddInput(weight_type); - - if (use_cifg) { - recurrent_to_input_weights_ = AddNullInput(); - } else { - recurrent_to_input_weights_ = AddInput(weight_type); - } - - recurrent_to_forget_weights_ = AddInput(weight_type); - recurrent_to_cell_weights_ = AddInput(weight_type); - recurrent_to_output_weights_ = AddInput(weight_type); - - if (use_peephole) { - if (use_cifg) { - cell_to_input_weights_ = AddNullInput(); - } else { - cell_to_input_weights_ = AddInput(weight_type); - } - cell_to_forget_weights_ = AddInput(weight_type); - cell_to_output_weights_ = AddInput(weight_type); - } else { - cell_to_input_weights_ = AddNullInput(); - cell_to_forget_weights_ = AddNullInput(); - cell_to_output_weights_ = AddNullInput(); - } - - if (use_cifg) { - input_gate_bias_ = AddNullInput(); - } else { - input_gate_bias_ = AddInput(TensorType_FLOAT32); - } - forget_gate_bias_ = AddInput(TensorType_FLOAT32); - cell_bias_ = AddInput(TensorType_FLOAT32); - output_gate_bias_ = AddInput(TensorType_FLOAT32); - - if (use_projection_weights) { - projection_weights_ = AddInput(weight_type); - if (use_projection_bias) { - projection_bias_ = AddInput(TensorType_FLOAT32); - } else { - projection_bias_ = AddNullInput(); - } - } else { - projection_weights_ = AddNullInput(); - projection_bias_ = AddNullInput(); - } - - // Adding the 2 input state tensors. - input_activation_state_ = - AddInput(TensorData{TensorType_FLOAT32, {n_output_ * n_batch_}}, true); - input_cell_state_ = - AddInput(TensorData{TensorType_FLOAT32, {n_cell_ * n_batch_}}, true); - - output_ = AddOutput(TensorType_FLOAT32); - - SetBuiltinOp(BuiltinOperator_LSTM, BuiltinOptions_LSTMOptions, - CreateLSTMOptions(builder_, ActivationFunctionType_TANH, - cell_clip, proj_clip) - .Union()); - - BuildInterpreter(input_shapes); - } - - void SetInputToInputWeights(std::vector f) { - PopulateTensor(input_to_input_weights_, f); - } - - void SetInputToForgetWeights(std::vector f) { - PopulateTensor(input_to_forget_weights_, f); - } - - void SetInputToCellWeights(std::vector f) { - PopulateTensor(input_to_cell_weights_, f); - } - - void SetInputToOutputWeights(std::vector f) { - PopulateTensor(input_to_output_weights_, f); - } - - void SetRecurrentToInputWeights(std::vector f) { - PopulateTensor(recurrent_to_input_weights_, f); - } - - void SetRecurrentToForgetWeights(std::vector f) { - PopulateTensor(recurrent_to_forget_weights_, f); - } - - void SetRecurrentToCellWeights(std::vector f) { - PopulateTensor(recurrent_to_cell_weights_, f); - } - - void SetRecurrentToOutputWeights(std::vector f) { - PopulateTensor(recurrent_to_output_weights_, f); - } - - void SetCellToInputWeights(std::vector f) { - PopulateTensor(cell_to_input_weights_, f); - } - - void SetCellToForgetWeights(std::vector f) { - PopulateTensor(cell_to_forget_weights_, f); - } - - void SetCellToOutputWeights(std::vector f) { - PopulateTensor(cell_to_output_weights_, f); - } - - void SetInputGateBias(std::vector f) { - PopulateTensor(input_gate_bias_, f); - } - - void SetForgetGateBias(std::vector f) { - PopulateTensor(forget_gate_bias_, f); - } - - void SetCellBias(std::vector f) { PopulateTensor(cell_bias_, f); } - - void SetOutputGateBias(std::vector f) { - PopulateTensor(output_gate_bias_, f); - } - - void SetProjectionWeights(std::vector f) { - PopulateTensor(projection_weights_, f); - } - - void SetProjectionBias(std::vector f) { - PopulateTensor(projection_bias_, f); - } - - void SetInput(int offset, const float* begin, const float* end) { - PopulateTensor(input_, offset, const_cast(begin), - const_cast(end)); - } - - std::vector GetOutput() { return ExtractVector(output_); } - - int num_inputs() { return n_input_; } - int num_outputs() { return n_output_; } - int num_cells() { return n_cell_; } - int num_batches() { return n_batch_; } - - protected: - int input_; - int input_to_input_weights_; - int input_to_forget_weights_; - int input_to_cell_weights_; - int input_to_output_weights_; - - int recurrent_to_input_weights_; - int recurrent_to_forget_weights_; - int recurrent_to_cell_weights_; - int recurrent_to_output_weights_; - - int cell_to_input_weights_; - int cell_to_forget_weights_; - int cell_to_output_weights_; - - int input_gate_bias_; - int forget_gate_bias_; - int cell_bias_; - int output_gate_bias_; - - int projection_weights_; - int projection_bias_; - int input_activation_state_; - int input_cell_state_; - - int output_; - int output_state_; - int cell_state_; - - int n_batch_; - int n_input_; - int n_cell_; - int n_output_; -}; - -class HybridLSTMOpModel : public LSTMOpModel { - public: - HybridLSTMOpModel(int n_batch, int n_input, int n_cell, int n_output, - bool use_cifg, bool use_peephole, - bool use_projection_weights, bool use_projection_bias, - float cell_clip, float proj_clip, - const std::vector>& input_shapes) - : LSTMOpModel(n_batch, n_input, n_cell, n_output, use_cifg, use_peephole, - use_projection_weights, use_projection_bias, cell_clip, - proj_clip, input_shapes, TensorType_UINT8) {} - - void SetInputToInputWeights(std::vector f) { - SymmetricQuantizeAndPopulate(input_to_input_weights_, f); - } - - void SetInputToForgetWeights(std::vector f) { - SymmetricQuantizeAndPopulate(input_to_forget_weights_, f); - } - - void SetInputToCellWeights(std::vector f) { - SymmetricQuantizeAndPopulate(input_to_cell_weights_, f); - } - - void SetInputToOutputWeights(std::vector f) { - SymmetricQuantizeAndPopulate(input_to_output_weights_, f); - } - - void SetRecurrentToInputWeights(std::vector f) { - SymmetricQuantizeAndPopulate(recurrent_to_input_weights_, f); - } - - void SetRecurrentToForgetWeights(std::vector f) { - SymmetricQuantizeAndPopulate(recurrent_to_forget_weights_, f); - } - - void SetRecurrentToCellWeights(std::vector f) { - SymmetricQuantizeAndPopulate(recurrent_to_cell_weights_, f); - } - - void SetRecurrentToOutputWeights(std::vector f) { - SymmetricQuantizeAndPopulate(recurrent_to_output_weights_, f); - } - - void SetCellToInputWeights(std::vector f) { - SymmetricQuantizeAndPopulate(cell_to_input_weights_, f); - } - - void SetCellToForgetWeights(std::vector f) { - SymmetricQuantizeAndPopulate(cell_to_forget_weights_, f); - } - - void SetCellToOutputWeights(std::vector f) { - SymmetricQuantizeAndPopulate(cell_to_output_weights_, f); - } - - void SetProjectionWeights(std::vector f) { - SymmetricQuantizeAndPopulate(projection_weights_, f); - } -}; - -class BaseLstmTest : public ::testing::Test { - protected: - // Weights of the LSTM model. Some are optional. - std::vector input_to_input_weights_; - std::vector input_to_cell_weights_; - std::vector input_to_forget_weights_; - std::vector input_to_output_weights_; - std::vector input_gate_bias_; - std::vector cell_gate_bias_; - std::vector forget_gate_bias_; - std::vector output_gate_bias_; - std::vector recurrent_to_input_weights_; - std::vector recurrent_to_cell_weights_; - std::vector recurrent_to_forget_weights_; - std::vector recurrent_to_output_weights_; - std::vector cell_to_input_weights_; - std::vector cell_to_forget_weights_; - std::vector cell_to_output_weights_; - std::vector projection_weights_; - - // LSTM input is stored as num_batch x num_inputs vector. - std::vector> lstm_input_; - // LSTM output is stored as num_batch x num_outputs vector. - std::vector> lstm_golden_output_; - - // Compares output up to tolerance to the result of the lstm given the input. - void VerifyGoldens(const std::vector>& input, - const std::vector>& output, - LSTMOpModel* lstm, float tolerance = 1e-5) { - const int num_batches = input.size(); - EXPECT_GT(num_batches, 0); - const int num_inputs = lstm->num_inputs(); - EXPECT_GT(num_inputs, 0); - const int input_sequence_size = input[0].size() / num_inputs; - EXPECT_GT(input_sequence_size, 0); - for (int i = 0; i < input_sequence_size; ++i) { - for (int b = 0; b < num_batches; ++b) { - const float* batch_start = input[b].data() + i * num_inputs; - const float* batch_end = batch_start + num_inputs; - - lstm->SetInput(b * lstm->num_inputs(), batch_start, batch_end); - } - - lstm->Invoke(); - - const int num_outputs = lstm->num_outputs(); - std::vector expected; - for (int b = 0; b < num_batches; ++b) { - const float* golden_start_batch = output[b].data() + i * num_outputs; - const float* golden_end_batch = golden_start_batch + num_outputs; - expected.insert(expected.end(), golden_start_batch, golden_end_batch); - } - EXPECT_THAT(lstm->GetOutput(), - ElementsAreArray(ArrayFloatNear(expected, tolerance))); - } - } -}; - -class NoCifgNoPeepholeNoProjectionNoClippingLstmTest : public BaseLstmTest { - void SetUp() override { - input_to_input_weights_ = {-0.45018822, -0.02338299, -0.0870589, - -0.34550029, 0.04266912, -0.15680569, - -0.34856534, 0.43890524}; - input_to_cell_weights_ = {-0.50013041, 0.1370284, 0.11810488, 0.2013163, - -0.20583314, 0.44344562, 0.22077113, -0.29909778}; - input_to_forget_weights_ = {0.09701663, 0.20334584, -0.50592935, - -0.31343272, -0.40032279, 0.44781327, - 0.01387155, -0.35593212}; - input_to_output_weights_ = {-0.25065863, -0.28290087, 0.04613829, - 0.40525138, 0.44272184, 0.03897077, - -0.1556896, 0.19487578}; - input_gate_bias_ = {0., 0., 0., 0.}; - cell_gate_bias_ = {0., 0., 0., 0.}; - forget_gate_bias_ = {1., 1., 1., 1.}; - output_gate_bias_ = {0., 0., 0., 0.}; - - recurrent_to_input_weights_ = { - -0.0063535, -0.2042388, 0.31454784, -0.35746509, - 0.28902304, 0.08183324, -0.16555229, 0.02286911, - -0.13566875, 0.03034258, 0.48091322, -0.12528998, - 0.24077177, -0.51332325, -0.33502164, 0.10629296}; - - recurrent_to_cell_weights_ = { - -0.3407414, 0.24443203, -0.2078532, 0.26320225, - 0.05695659, -0.00123841, -0.4744786, -0.35869038, - -0.06418842, -0.13502428, -0.501764, 0.22830659, - -0.46367589, 0.26016325, -0.03894562, -0.16368064}; - - recurrent_to_forget_weights_ = { - -0.48684245, -0.06655136, 0.42224967, 0.2112639, - 0.27654213, 0.20864892, -0.07646349, 0.45877004, - 0.00141793, -0.14609534, 0.36447752, 0.09196436, - 0.28053468, 0.01560611, -0.20127171, -0.01140004}; - - recurrent_to_output_weights_ = { - 0.43385774, -0.17194885, 0.2718237, 0.09215671, - 0.24107647, -0.39835793, 0.18212086, 0.01301402, - 0.48572797, -0.50656658, 0.20047462, -0.20607421, - -0.51818722, -0.15390486, 0.0468148, 0.39922136}; - - lstm_input_ = {{2., 3., 3., 4., 1., 1.}}; - lstm_golden_output_ = {{-0.02973187, 0.1229473, 0.20885126, -0.15358765, - -0.03716109, 0.12507336, 0.41193449, -0.20860538, - -0.15053082, 0.09120187, 0.24278517, -0.12222792}}; - } -}; - -TEST_F(NoCifgNoPeepholeNoProjectionNoClippingLstmTest, LstmBlackBoxTest) { - const int n_batch = 1; - const int n_input = 2; - // n_cell and n_output have the same size when there is no projection. - const int n_cell = 4; - const int n_output = 4; - - LSTMOpModel lstm(n_batch, n_input, n_cell, n_output, - /*use_cifg=*/false, /*use_peephole=*/false, - /*use_projection_weights=*/false, - /*use_projection_bias=*/false, - /*cell_clip=*/0.0, /*proj_clip=*/0.0, - { - {n_batch, n_input}, // input tensor - - {n_cell, n_input}, // input_to_input_weight tensor - {n_cell, n_input}, // input_to_forget_weight tensor - {n_cell, n_input}, // input_to_cell_weight tensor - {n_cell, n_input}, // input_to_output_weight tensor - - {n_cell, n_output}, // recurrent_to_input_weight_tensor - {n_cell, n_output}, // recurrent_to_forget_weight_tensor - {n_cell, n_output}, // recurrent_to_cell_weight_tensor - {n_cell, n_output}, // recurrent_to_output_weight_tensor - - {0}, // cell_to_input_weight tensor - {0}, // cell_to_forget_weight tensor - {0}, // cell_to_output_weight tensor - - {n_cell}, // input_gate_bias tensor - {n_cell}, // forget_gate_bias tensor - {n_cell}, // cell_bias tensor - {n_cell}, // output_gate_bias tensor - - {0, 0}, // projection_weight tensor - {0}, // projection_bias tensor - }); - - lstm.SetInputToInputWeights(input_to_input_weights_); - lstm.SetInputToCellWeights(input_to_cell_weights_); - lstm.SetInputToForgetWeights(input_to_forget_weights_); - lstm.SetInputToOutputWeights(input_to_output_weights_); - - lstm.SetInputGateBias(input_gate_bias_); - lstm.SetCellBias(cell_gate_bias_); - lstm.SetForgetGateBias(forget_gate_bias_); - lstm.SetOutputGateBias(output_gate_bias_); - - lstm.SetRecurrentToInputWeights(recurrent_to_input_weights_); - lstm.SetRecurrentToCellWeights(recurrent_to_cell_weights_); - lstm.SetRecurrentToForgetWeights(recurrent_to_forget_weights_); - lstm.SetRecurrentToOutputWeights(recurrent_to_output_weights_); - - VerifyGoldens(lstm_input_, lstm_golden_output_, &lstm); -} - -TEST_F(NoCifgNoPeepholeNoProjectionNoClippingLstmTest, HybridLstmBlackBoxTest) { - const int n_batch = 1; - const int n_input = 2; - // n_cell and n_output have the same size when there is no projection. - const int n_cell = 4; - const int n_output = 4; - - HybridLSTMOpModel lstm( - n_batch, n_input, n_cell, n_output, - /*use_cifg=*/false, /*use_peephole=*/false, - /*use_projection_weights=*/false, - /*use_projection_bias=*/false, /*cell_clip=*/0.0, /*proj_clip=*/0.0, - { - {n_batch, n_input}, // input tensor - - {n_cell, n_input}, // input_to_input_weight tensor - {n_cell, n_input}, // input_to_forget_weight tensor - {n_cell, n_input}, // input_to_cell_weight tensor - {n_cell, n_input}, // input_to_output_weight tensor - - {n_cell, n_output}, // recurrent_to_input_weight tensor - {n_cell, n_output}, // recurrent_to_forget_weight tensor - {n_cell, n_output}, // recurrent_to_cell_weight tensor - {n_cell, n_output}, // recurrent_to_output_weight tensor - - {0}, // cell_to_input_weight tensor - {0}, // cell_to_forget_weight tensor - {0}, // cell_to_output_weight tensor - - {n_cell}, // input_gate_bias tensor - {n_cell}, // forget_gate_bias tensor - {n_cell}, // cell_bias tensor - {n_cell}, // output_gate_bias tensor - - {0, 0}, // projection_weight tensor - {0}, // projection_bias tensor - }); - - lstm.SetInputToInputWeights(input_to_input_weights_); - lstm.SetInputToCellWeights(input_to_cell_weights_); - lstm.SetInputToForgetWeights(input_to_forget_weights_); - lstm.SetInputToOutputWeights(input_to_output_weights_); - - lstm.SetInputGateBias(input_gate_bias_); - lstm.SetCellBias(cell_gate_bias_); - lstm.SetForgetGateBias(forget_gate_bias_); - lstm.SetOutputGateBias(output_gate_bias_); - - lstm.SetRecurrentToInputWeights(recurrent_to_input_weights_); - lstm.SetRecurrentToCellWeights(recurrent_to_cell_weights_); - lstm.SetRecurrentToForgetWeights(recurrent_to_forget_weights_); - lstm.SetRecurrentToOutputWeights(recurrent_to_output_weights_); - - VerifyGoldens(lstm_input_, lstm_golden_output_, &lstm, - /*tolerance=*/0.0157651); -} - -class CifgNoPeepholeNoProjectionNoClippingLstmTest : public BaseLstmTest { - void SetUp() override { - input_to_cell_weights_ = {-0.49770179, -0.27711356, -0.09624726, - 0.05100781, 0.04717243, 0.48944736, - -0.38535351, -0.17212132}; - - input_to_forget_weights_ = {-0.55291498, -0.42866567, 0.13056988, - -0.3633365, -0.22755712, 0.28253698, - 0.24407166, 0.33826375}; - - input_to_output_weights_ = {0.10725588, -0.02335852, -0.55932593, - -0.09426838, -0.44257352, 0.54939759, - 0.01533556, 0.42751634}; - cell_gate_bias_ = {0., 0., 0., 0.}; - forget_gate_bias_ = {1., 1., 1., 1.}; - output_gate_bias_ = {0., 0., 0., 0.}; - - recurrent_to_cell_weights_ = { - 0.54066205, -0.32668582, -0.43562764, -0.56094903, - 0.42957711, 0.01841056, -0.32764608, -0.33027974, - -0.10826075, 0.20675004, 0.19069612, -0.03026325, - -0.54532051, 0.33003211, 0.44901288, 0.21193194}; - - recurrent_to_forget_weights_ = { - -0.13832897, -0.0515101, -0.2359007, -0.16661474, - -0.14340827, 0.36986142, 0.23414481, 0.55899, - 0.10798943, -0.41174671, 0.17751795, -0.34484994, - -0.35874045, -0.11352962, 0.27268326, 0.54058349}; - - recurrent_to_output_weights_ = { - 0.41613156, 0.42610586, -0.16495961, -0.5663873, - 0.30579174, -0.05115908, -0.33941799, 0.23364776, - 0.11178309, 0.09481031, -0.26424935, 0.46261835, - 0.50248802, 0.26114327, -0.43736315, 0.33149987}; - - cell_to_forget_weights_ = {0.47485286, -0.51955009, -0.24458408, - 0.31544167}; - cell_to_output_weights_ = {-0.17135078, 0.82760304, 0.85573703, - -0.77109635}; - - lstm_input_ = {{2., 3., 3., 4., 1., 1.}}; - lstm_golden_output_ = {{-0.36444446, -0.00352185, 0.12886585, -0.05163646, - -0.42312205, -0.01218222, 0.24201041, -0.08124574, - -0.358325, -0.04621704, 0.21641694, -0.06471302}}; - } -}; - -TEST_F(CifgNoPeepholeNoProjectionNoClippingLstmTest, LstmBlackBoxTest) { - const int n_batch = 1; - const int n_input = 2; - // n_cell and n_output have the same size when there is no projection. - const int n_cell = 4; - const int n_output = 4; - - LSTMOpModel lstm(n_batch, n_input, n_cell, n_output, - /*use_cifg=*/true, /*use_peephole=*/true, - /*use_projection_weights=*/false, - /*use_projection_bias=*/false, - /*cell_clip=*/0.0, /*proj_clip=*/0.0, - { - {n_batch, n_input}, // input tensor - - {0, 0}, // input_to_input_weight tensor - {n_cell, n_input}, // input_to_forget_weight tensor - {n_cell, n_input}, // input_to_cell_weight tensor - {n_cell, n_input}, // input_to_output_weight tensor - - {0, 0}, // recurrent_to_input_weight tensor - {n_cell, n_output}, // recurrent_to_forget_weight tensor - {n_cell, n_output}, // recurrent_to_cell_weight tensor - {n_cell, n_output}, // recurrent_to_output_weight tensor - - {0}, // cell_to_input_weight tensor - {n_cell}, // cell_to_forget_weight tensor - {n_cell}, // cell_to_output_weight tensor - - {0}, // input_gate_bias tensor - {n_cell}, // forget_gate_bias tensor - {n_cell}, // cell_bias tensor - {n_cell}, // output_gate_bias tensor - - {0, 0}, // projection_weight tensor - {0}, // projection_bias tensor - }); - - lstm.SetInputToCellWeights(input_to_cell_weights_); - lstm.SetInputToForgetWeights(input_to_forget_weights_); - lstm.SetInputToOutputWeights(input_to_output_weights_); - - lstm.SetCellBias(cell_gate_bias_); - lstm.SetForgetGateBias(forget_gate_bias_); - lstm.SetOutputGateBias(output_gate_bias_); - - lstm.SetRecurrentToCellWeights(recurrent_to_cell_weights_); - lstm.SetRecurrentToForgetWeights(recurrent_to_forget_weights_); - lstm.SetRecurrentToOutputWeights(recurrent_to_output_weights_); - - lstm.SetCellToForgetWeights(cell_to_forget_weights_); - lstm.SetCellToOutputWeights(cell_to_output_weights_); - - VerifyGoldens(lstm_input_, lstm_golden_output_, &lstm); -} - -TEST_F(CifgNoPeepholeNoProjectionNoClippingLstmTest, HybridLstmBlackBoxTest) { - const int n_batch = 1; - const int n_input = 2; - // n_cell and n_output have the same size when there is no projection. - const int n_cell = 4; - const int n_output = 4; - - HybridLSTMOpModel lstm( - n_batch, n_input, n_cell, n_output, - /*use_cifg=*/true, /*use_peephole=*/true, - /*use_projection_weights=*/false, - /*use_projection_bias=*/false, - /*cell_clip=*/0.0, /*proj_clip=*/0.0, - { - {n_batch, n_input}, // input tensor - - {0, 0}, // input_to_input_weight tensor - {n_cell, n_input}, // input_to_forget_weight tensor - {n_cell, n_input}, // input_to_cell_weight tensor - {n_cell, n_input}, // input_to_output_weight tensor - - {0, 0}, // recurrent_to_input_weight tensor - {n_cell, n_output}, // recurrent_to_forget_weight tensor - {n_cell, n_output}, // recurrent_to_cell_weight tensor - {n_cell, n_output}, // recurrent_to_output_weight tensor - - {0}, // cell_to_input_weight tensor - {n_cell}, // cell_to_forget_weight tensor - {n_cell}, // cell_to_output_weight tensor - - {0}, // input_gate_bias tensor - {n_cell}, // forget_gate_bias tensor - {n_cell}, // cell_bias tensor - {n_cell}, // output_gate_bias tensor - - {0, 0}, // projection_weight tensor - {0}, // projection_bias tensor - }); - - lstm.SetInputToCellWeights(input_to_cell_weights_); - lstm.SetInputToForgetWeights(input_to_forget_weights_); - lstm.SetInputToOutputWeights(input_to_output_weights_); - - lstm.SetCellBias(cell_gate_bias_); - lstm.SetForgetGateBias(forget_gate_bias_); - lstm.SetOutputGateBias(output_gate_bias_); - - lstm.SetRecurrentToCellWeights(recurrent_to_cell_weights_); - lstm.SetRecurrentToForgetWeights(recurrent_to_forget_weights_); - lstm.SetRecurrentToOutputWeights(recurrent_to_output_weights_); - - lstm.SetCellToForgetWeights(cell_to_forget_weights_); - lstm.SetCellToOutputWeights(cell_to_output_weights_); - - VerifyGoldens(lstm_input_, lstm_golden_output_, &lstm, /*tolerance=*/0.03573); -} - -class NoCifgPeepholeProjectionNoClippingLstmTest : public BaseLstmTest { - void SetUp() override { - input_to_input_weights_ = { - 0.021393683, 0.06124551, 0.046905167, -0.014657677, -0.03149463, - 0.09171803, 0.14647801, 0.10797193, -0.0057968358, 0.0019193048, - -0.2726754, 0.10154029, -0.018539885, 0.080349885, -0.10262385, - -0.022599787, -0.09121155, -0.008675967, -0.045206103, -0.0821282, - -0.008045952, 0.015478081, 0.055217247, 0.038719587, 0.044153627, - -0.06453243, 0.05031825, -0.046935108, -0.008164439, 0.014574226, - -0.1671009, -0.15519552, -0.16819797, -0.13971269, -0.11953059, - 0.25005487, -0.22790983, 0.009855087, -0.028140958, -0.11200698, - 0.11295408, -0.0035217577, 0.054485075, 0.05184695, 0.064711206, - 0.10989193, 0.11674786, 0.03490607, 0.07727357, 0.11390585, - -0.1863375, -0.1034451, -0.13945189, -0.049401227, -0.18767063, - 0.042483903, 0.14233552, 0.13832581, 0.18350165, 0.14545603, - -0.028545704, 0.024939531, 0.050929718, 0.0076203286, -0.0029723682, - -0.042484224, -0.11827596, -0.09171104, -0.10808628, -0.16327988, - -0.2273378, -0.0993647, -0.017155107, 0.0023917493, 0.049272764, - 0.0038534778, 0.054764505, 0.089753784, 0.06947234, 0.08014476, - -0.04544234, -0.0497073, -0.07135631, -0.048929106, -0.004042012, - -0.009284026, 0.018042054, 0.0036860977, -0.07427302, -0.11434604, - -0.018995456, 0.031487543, 0.012834908, 0.019977754, 0.044256654, - -0.39292613, -0.18519334, -0.11651281, -0.06809892, 0.011373677}; - - input_to_forget_weights_ = { - -0.0018401089, -0.004852237, 0.03698424, 0.014181704, - 0.028273236, -0.016726194, -0.05249759, -0.10204261, - 0.00861066, -0.040979505, -0.009899187, 0.01923892, - -0.028177269, -0.08535103, -0.14585495, 0.10662567, - -0.01909731, -0.017883534, -0.0047269356, -0.045103323, - 0.0030784295, 0.076784775, 0.07463696, 0.094531395, - 0.0814421, -0.12257899, -0.033945758, -0.031303465, - 0.045630626, 0.06843887, -0.13492945, -0.012480007, - -0.0811829, -0.07224499, -0.09628791, 0.045100946, - 0.0012300825, 0.013964662, 0.099372394, 0.02543059, - 0.06958324, 0.034257296, 0.0482646, 0.06267997, - 0.052625068, 0.12784666, 0.07077897, 0.025725935, - 0.04165009, 0.07241905, 0.018668644, -0.037377294, - -0.06277783, -0.08833636, -0.040120605, -0.011405586, - -0.007808335, -0.010301386, -0.005102167, 0.027717464, - 0.05483423, 0.11449111, 0.11289652, 0.10939839, - 0.13396506, -0.08402166, -0.01901462, -0.044678304, - -0.07720565, 0.014350063, -0.11757958, -0.0652038, - -0.08185733, -0.076754324, -0.092614375, 0.10405491, - 0.052960336, 0.035755895, 0.035839386, -0.012540553, - 0.036881298, 0.02913376, 0.03420159, 0.05448447, - -0.054523353, 0.02582715, 0.02327355, -0.011857179, - -0.0011980024, -0.034641717, -0.026125094, -0.17582615, - -0.15923657, -0.27486774, -0.0006143371, 0.0001771948, - -8.470171e-05, 0.02651807, 0.045790765, 0.06956496}; - - input_to_cell_weights_ = { - -0.04580283, -0.09549462, -0.032418985, -0.06454633, - -0.043528453, 0.043018587, -0.049152344, -0.12418144, - -0.078985475, -0.07596889, 0.019484362, -0.11434962, - -0.0074034138, -0.06314844, -0.092981495, 0.0062155537, - -0.025034338, -0.0028890965, 0.048929527, 0.06235075, - 0.10665918, -0.032036792, -0.08505916, -0.10843358, - -0.13002433, -0.036816437, -0.02130134, -0.016518239, - 0.0047691227, -0.0025825808, 0.066017866, 0.029991534, - -0.10652836, -0.1037554, -0.13056071, -0.03266643, - -0.033702414, -0.006473424, -0.04611692, 0.014419339, - -0.025174323, 0.0396852, 0.081777506, 0.06157468, - 0.10210095, -0.009658194, 0.046511717, 0.03603906, - 0.0069369148, 0.015960095, -0.06507666, 0.09551598, - 0.053568836, 0.06408714, 0.12835667, -0.008714329, - -0.20211966, -0.12093674, 0.029450472, 0.2849013, - -0.029227901, 0.1164364, -0.08560263, 0.09941786, - -0.036999565, -0.028842626, -0.0033637602, -0.017012902, - -0.09720865, -0.11193351, -0.029155117, -0.017936034, - -0.009768936, -0.04223324, -0.036159635, 0.06505112, - -0.021742892, -0.023377212, -0.07221364, -0.06430552, - 0.05453865, 0.091149814, 0.06387331, 0.007518393, - 0.055960953, 0.069779344, 0.046411168, 0.10509911, - 0.07463894, 0.0075130584, 0.012850982, 0.04555431, - 0.056955688, 0.06555285, 0.050801456, -0.009862683, - 0.00826772, -0.026555609, -0.0073611983, -0.0014897042}; - - input_to_output_weights_ = { - -0.0998932, -0.07201956, -0.052803773, -0.15629593, -0.15001918, - -0.07650751, 0.02359855, -0.075155355, -0.08037709, -0.15093534, - 0.029517552, -0.04751393, 0.010350531, -0.02664851, -0.016839722, - -0.023121163, 0.0077019283, 0.012851257, -0.05040649, -0.0129761, - -0.021737747, -0.038305793, -0.06870586, -0.01481247, -0.001285394, - 0.10124236, 0.083122835, 0.053313006, -0.062235646, -0.075637154, - -0.027833903, 0.029774971, 0.1130802, 0.09218906, 0.09506135, - -0.086665764, -0.037162706, -0.038880914, -0.035832845, -0.014481564, - -0.09825003, -0.12048569, -0.097665586, -0.05287633, -0.0964047, - -0.11366429, 0.035777505, 0.13568819, 0.052451383, 0.050649304, - 0.05798951, -0.021852335, -0.099848844, 0.014740475, -0.078897946, - 0.04974699, 0.014160473, 0.06973932, 0.04964942, 0.033364646, - 0.08190124, 0.025535367, 0.050893165, 0.048514254, 0.06945813, - -0.078907564, -0.06707616, -0.11844508, -0.09986688, -0.07509403, - 0.06263226, 0.14925587, 0.20188436, 0.12098451, 0.14639415, - 0.0015017595, -0.014267382, -0.03417257, 0.012711468, 0.0028300495, - -0.024758482, -0.05098548, -0.0821182, 0.014225672, 0.021544158, - 0.08949725, 0.07505268, -0.0020780868, 0.04908258, 0.06476295, - -0.022907063, 0.027562456, 0.040185735, 0.019567577, -0.015598739, - -0.049097303, -0.017121866, -0.083368234, -0.02332002, -0.0840956}; - - input_gate_bias_ = {0.02234832, 0.14757581, 0.18176508, 0.10380666, - 0.053110216, -0.06928846, -0.13942584, -0.11816189, - 0.19483899, 0.03652339, -0.10250295, 0.036714908, - -0.18426876, 0.036065217, 0.21810818, 0.02383196, - -0.043370757, 0.08690144, -0.04444982, 0.00030581196}; - - forget_gate_bias_ = {0.035185695, -0.042891346, -0.03032477, 0.23027696, - 0.11098921, 0.15378423, 0.09263801, 0.09790885, - 0.09508917, 0.061199076, 0.07665568, -0.015443159, - -0.03499149, 0.046190713, 0.08895977, 0.10899629, - 0.40694186, 0.06030037, 0.012413437, -0.06108739}; - - cell_gate_bias_ = {-0.024379363, 0.0055531194, 0.23377132, 0.033463873, - -0.1483596, -0.10639995, -0.091433935, 0.058573797, - -0.06809782, -0.07889636, -0.043246906, -0.09829136, - -0.4279842, 0.034901652, 0.18797937, 0.0075234566, - 0.016178843, 0.1749513, 0.13975595, 0.92058027}; - - output_gate_bias_ = {0.046159424, -0.0012809046, 0.03563469, 0.12648113, - 0.027195795, 0.35373217, -0.018957434, 0.008907322, - -0.0762701, 0.12018895, 0.04216877, 0.0022856654, - 0.040952638, 0.3147856, 0.08225149, -0.057416286, - -0.14995944, -0.008040261, 0.13208859, 0.029760877}; - - recurrent_to_input_weights_ = { - -0.001374326, -0.078856036, 0.10672688, 0.029162422, - -0.11585556, 0.02557986, -0.13446963, -0.035785314, - -0.01244275, 0.025961924, -0.02337298, -0.044228926, - -0.055839065, -0.046598054, -0.010546039, -0.06900766, - 0.027239809, 0.022582639, -0.013296484, -0.05459212, - 0.08981, -0.045407712, 0.08682226, -0.06867011, - -0.14390695, -0.02916037, 0.000996957, 0.091420636, - 0.14283475, -0.07390571, -0.06402044, 0.062524505, - -0.093129106, 0.04860203, -0.08364217, -0.08119002, - 0.009352075, 0.22920375, 0.0016303885, 0.11583097, - -0.13732095, 0.012405723, -0.07551853, 0.06343048, - 0.12162708, -0.031923793, -0.014335606, 0.01790974, - -0.10650317, -0.0724401, 0.08554849, -0.05727212, - 0.06556731, -0.042729504, -0.043227166, 0.011683251, - -0.013082158, -0.029302018, -0.010899579, -0.062036745, - -0.022509435, -0.00964907, -0.01567329, 0.04260106, - -0.07787477, -0.11576462, 0.017356863, 0.048673786, - -0.017577527, -0.05527947, -0.082487635, -0.040137455, - -0.10820036, -0.04666372, 0.022746278, -0.07851417, - 0.01068115, 0.032956902, 0.022433773, 0.0026891115, - 0.08944216, -0.0685835, 0.010513544, 0.07228705, - 0.02032331, -0.059686817, -0.0005566496, -0.086984694, - 0.040414046, -0.1380399, 0.094208956, -0.05722982, - 0.012092817, -0.04989123, -0.086576, -0.003399834, - -0.04696032, -0.045747425, 0.10091314, 0.048676282, - -0.029037097, 0.031399418, -0.0040285117, 0.047237843, - 0.09504992, 0.041799378, -0.049185462, -0.031518843, - -0.10516937, 0.026374253, 0.10058866, -0.0033195973, - -0.041975245, 0.0073591834, 0.0033782164, -0.004325073, - -0.10167381, 0.042500053, -0.01447153, 0.06464186, - -0.017142897, 0.03312627, 0.009205989, 0.024138335, - -0.011337001, 0.035530265, -0.010912711, 0.0706555, - -0.005894094, 0.051841937, -0.1401738, -0.02351249, - 0.0365468, 0.07590991, 0.08838724, 0.021681072, - -0.10086113, 0.019608743, -0.06195883, 0.077335775, - 0.023646897, -0.095322326, 0.02233014, 0.09756986, - -0.048691444, -0.009579111, 0.07595467, 0.11480546, - -0.09801813, 0.019894179, 0.08502348, 0.004032281, - 0.037211012, 0.068537936, -0.048005626, -0.091520436, - -0.028379958, -0.01556313, 0.06554592, -0.045599163, - -0.01672207, -0.020169014, -0.011877351, -0.20212261, - 0.010889619, 0.0047078193, 0.038385306, 0.08540671, - -0.017140968, -0.0035865551, 0.016678626, 0.005633034, - 0.015963363, 0.00871737, 0.060130805, 0.028611384, - 0.10109069, -0.015060172, -0.07894427, 0.06401885, - 0.011584063, -0.024466386, 0.0047652307, -0.09041358, - 0.030737216, -0.0046374933, 0.14215417, -0.11823516, - 0.019899689, 0.006106124, -0.027092824, 0.0786356, - 0.05052217, -0.058925, -0.011402121, -0.024987547, - -0.0013661642, -0.06832946, -0.015667673, -0.1083353, - -0.00096863037, -0.06988685, -0.053350925, -0.027275559, - -0.033664223, -0.07978348, -0.025200296, -0.017207067, - -0.058403496, -0.055697463, 0.005798788, 0.12965427, - -0.062582195, 0.0013350133, -0.10482091, 0.0379771, - 0.072521195, -0.0029455067, -0.13797039, -0.03628521, - 0.013806405, -0.017858358, -0.01008298, -0.07700066, - -0.017081132, 0.019358726, 0.0027079724, 0.004635139, - 0.062634714, -0.02338735, -0.039547626, -0.02050681, - 0.03385117, -0.083611414, 0.002862572, -0.09421313, - 0.058618143, -0.08598433, 0.00972939, 0.023867095, - -0.053934585, -0.023203006, 0.07452513, -0.048767887, - -0.07314807, -0.056307215, -0.10433547, -0.06440842, - 0.04328182, 0.04389765, -0.020006588, -0.09076438, - -0.11652589, -0.021705797, 0.03345259, -0.010329105, - -0.025767034, 0.013057034, -0.07316461, -0.10145612, - 0.06358255, 0.18531723, 0.07759293, 0.12006465, - 0.1305557, 0.058638252, -0.03393652, 0.09622831, - -0.16253184, -2.4580743e-06, 0.079869635, -0.070196845, - -0.005644518, 0.06857898, -0.12598175, -0.035084512, - 0.03156317, -0.12794146, -0.031963028, 0.04692781, - 0.030070418, 0.0071660685, -0.095516115, -0.004643372, - 0.040170413, -0.062104587, -0.0037324072, 0.0554317, - 0.08184801, -0.019164372, 0.06791302, 0.034257166, - -0.10307039, 0.021943003, 0.046745934, 0.0790918, - -0.0265588, -0.007824208, 0.042546265, -0.00977924, - -0.0002440307, -0.017384544, -0.017990116, 0.12252321, - -0.014512694, -0.08251313, 0.08861942, 0.13589665, - 0.026351685, 0.012641483, 0.07466548, 0.044301085, - -0.045414884, -0.051112458, 0.03444247, -0.08502782, - -0.04106223, -0.028126027, 0.028473156, 0.10467447}; - - recurrent_to_cell_weights_ = { - -0.037322544, 0.018592842, 0.0056175636, -0.06253426, - 0.055647098, -0.05713207, -0.05626563, 0.005559383, - 0.03375411, -0.025757805, -0.088049285, 0.06017052, - -0.06570978, 0.007384076, 0.035123326, -0.07920549, - 0.053676967, 0.044480428, -0.07663568, 0.0071805613, - 0.08089997, 0.05143358, 0.038261272, 0.03339287, - -0.027673481, 0.044746667, 0.028349208, 0.020090483, - -0.019443132, -0.030755889, -0.0040000007, 0.04465846, - -0.021585021, 0.0031670958, 0.0053199246, -0.056117613, - -0.10893326, 0.076739706, -0.08509834, -0.027997585, - 0.037871376, 0.01449768, -0.09002357, -0.06111149, - -0.046195522, 0.0422062, -0.005683705, -0.1253618, - -0.012925729, -0.04890792, 0.06985068, 0.037654128, - 0.03398274, -0.004781977, 0.007032333, -0.031787455, - 0.010868644, -0.031489216, 0.09525667, 0.013939797, - 0.0058680447, 0.0167067, 0.02668468, -0.04797466, - -0.048885044, -0.12722108, 0.035304096, 0.06554885, - 0.00972396, -0.039238118, -0.05159735, -0.11329045, - 0.1613692, -0.03750952, 0.06529313, -0.071974665, - -0.11769596, 0.015524369, -0.0013754242, -0.12446318, - 0.02786344, -0.014179351, 0.005264273, 0.14376344, - 0.015983658, 0.03406988, -0.06939408, 0.040699873, - 0.02111075, 0.09669095, 0.041345075, -0.08316494, - -0.07684199, -0.045768797, 0.032298047, -0.041805092, - 0.0119405, 0.0061010392, 0.12652606, 0.0064572375, - -0.024950314, 0.11574242, 0.04508852, -0.04335324, - 0.06760663, -0.027437469, 0.07216407, 0.06977076, - -0.05438599, 0.034033038, -0.028602652, 0.05346137, - 0.043184172, -0.037189785, 0.10420091, 0.00882477, - -0.054019816, -0.074273005, -0.030617684, -0.0028467078, - 0.024302477, -0.0038869337, 0.005332455, 0.0013399826, - 0.04361412, -0.007001822, 0.09631092, -0.06702025, - -0.042049985, -0.035070654, -0.04103342, -0.10273396, - 0.0544271, 0.037184782, -0.13150354, -0.0058036847, - -0.008264958, 0.042035464, 0.05891794, 0.029673764, - 0.0063542654, 0.044788733, 0.054816857, 0.062257513, - -0.00093483756, 0.048938446, -0.004952862, -0.007730018, - -0.04043371, -0.017094059, 0.07229206, -0.023670016, - -0.052195564, -0.025616996, -0.01520939, 0.045104615, - -0.007376126, 0.003533447, 0.006570588, 0.056037236, - 0.12436656, 0.051817212, 0.028532185, -0.08686856, - 0.11868599, 0.07663395, -0.07323171, 0.03463402, - -0.050708205, -0.04458982, -0.11590894, 0.021273347, - 0.1251325, -0.15313013, -0.12224372, 0.17228661, - 0.023029093, 0.086124025, 0.006445803, -0.03496501, - 0.028332196, 0.04449512, -0.042436164, -0.026587414, - -0.006041347, -0.09292539, -0.05678812, 0.03897832, - 0.09465633, 0.008115513, -0.02171956, 0.08304309, - 0.071401566, 0.019622514, 0.032163795, -0.004167056, - 0.02295182, 0.030739572, 0.056506045, 0.004612461, - 0.06524936, 0.059999723, 0.046395954, -0.0045512207, - -0.1335546, -0.030136576, 0.11584653, -0.014678886, - 0.0020118146, -0.09688814, -0.0790206, 0.039770417, - -0.0329582, 0.07922767, 0.029322514, 0.026405897, - 0.04207835, -0.07073373, 0.063781224, 0.0859677, - -0.10925287, -0.07011058, 0.048005477, 0.03438226, - -0.09606514, -0.006669445, -0.043381985, 0.04240257, - -0.06955775, -0.06769346, 0.043903265, -0.026784198, - -0.017840602, 0.024307009, -0.040079936, -0.019946516, - 0.045318738, -0.12233574, 0.026170589, 0.0074471775, - 0.15978073, 0.10185836, 0.10298046, -0.015476589, - -0.039390966, -0.072174534, 0.0739445, -0.1211869, - -0.0347889, -0.07943156, 0.014809798, -0.12412325, - -0.0030663363, 0.039695457, 0.0647603, -0.08291318, - -0.018529687, -0.004423833, 0.0037507233, 0.084633216, - -0.01514876, -0.056505352, -0.012800942, -0.06994386, - 0.012962922, -0.031234352, 0.07029052, 0.016418684, - 0.03618972, 0.055686004, -0.08663945, -0.017404709, - -0.054761406, 0.029065743, 0.052404847, 0.020238016, - 0.0048197987, -0.0214882, 0.07078733, 0.013016777, - 0.06262858, 0.009184685, 0.020785125, -0.043904778, - -0.0270329, -0.03299152, -0.060088247, -0.015162964, - -0.001828936, 0.12642565, -0.056757294, 0.013586685, - 0.09232601, -0.035886683, 0.06000002, 0.05229691, - -0.052580316, -0.082029596, -0.010794592, 0.012947712, - -0.036429964, -0.085508935, -0.13127148, -0.017744139, - 0.031502828, 0.036232427, -0.031581745, 0.023051167, - -0.05325106, -0.03421577, 0.028793324, -0.034633752, - -0.009881397, -0.043551125, -0.018609839, 0.0019097115, - -0.008799762, 0.056595087, 0.0022273948, 0.055752404}; - - recurrent_to_forget_weights_ = { - -0.057784554, -0.026057621, -0.068447545, -0.022581743, - 0.14811787, 0.10826372, 0.09471067, 0.03987225, - -0.0039523416, 0.00030638507, 0.053185795, 0.10572994, - 0.08414449, -0.022036452, -0.00066928595, -0.09203576, - 0.032950465, -0.10985798, -0.023809856, 0.0021431844, - -0.02196096, -0.00326074, 0.00058621005, -0.074678116, - -0.06193199, 0.055729095, 0.03736828, 0.020123724, - 0.061878487, -0.04729229, 0.034919553, -0.07585433, - -0.04421272, -0.044019096, 0.085488975, 0.04058006, - -0.06890133, -0.030951202, -0.024628663, -0.07672815, - 0.034293607, 0.08556707, -0.05293577, -0.033561368, - -0.04899627, 0.0241671, 0.015736353, -0.095442444, - -0.029564252, 0.016493602, -0.035026584, 0.022337519, - -0.026871363, 0.004780428, 0.0077918363, -0.03601621, - 0.016435321, -0.03263031, -0.09543275, -0.047392778, - 0.013454138, 0.028934088, 0.01685226, -0.086110644, - -0.046250615, -0.01847454, 0.047608484, 0.07339695, - 0.034546845, -0.04881143, 0.009128804, -0.08802852, - 0.03761666, 0.008096139, -0.014454086, 0.014361001, - -0.023502491, -0.0011840804, -0.07607001, 0.001856849, - -0.06509276, -0.006021153, -0.08570962, -0.1451793, - 0.060212336, 0.055259194, 0.06974018, 0.049454916, - -0.027794661, -0.08077226, -0.016179763, 0.1169753, - 0.17213494, -0.0056326236, -0.053934924, -0.0124349, - -0.11520337, 0.05409887, 0.088759385, 0.0019655675, - 0.0042065294, 0.03881498, 0.019844765, 0.041858196, - -0.05695512, 0.047233116, 0.038937137, -0.06542224, - 0.014429736, -0.09719407, 0.13908425, -0.05379757, - 0.012321099, 0.082840554, -0.029899208, 0.044217527, - 0.059855383, 0.07711018, -0.045319796, 0.0948846, - -0.011724666, -0.0033288454, -0.033542685, -0.04764985, - -0.13873616, 0.040668588, 0.034832682, -0.015319203, - -0.018715994, 0.046002675, 0.0599172, -0.043107376, - 0.0294216, -0.002314414, -0.022424703, 0.0030315618, - 0.0014641669, 0.0029166266, -0.11878115, 0.013738511, - 0.12375372, -0.0006038222, 0.029104086, 0.087442465, - 0.052958444, 0.07558703, 0.04817258, 0.044462286, - -0.015213451, -0.08783778, -0.0561384, -0.003008196, - 0.047060397, -0.002058388, 0.03429439, -0.018839769, - 0.024734668, 0.024614193, -0.042046934, 0.09597743, - -0.0043254104, 0.04320769, 0.0064070094, -0.0019131786, - -0.02558259, -0.022822596, -0.023273505, -0.02464396, - -0.10991725, -0.006240552, 0.0074488563, 0.024044557, - 0.04383914, -0.046476185, 0.028658995, 0.060410924, - 0.050786525, 0.009452605, -0.0073054377, -0.024810238, - 0.0052906186, 0.0066939713, -0.0020913032, 0.014515517, - 0.015898481, 0.021362653, -0.030262267, 0.016587038, - -0.011442813, 0.041154444, -0.007631438, -0.03423484, - -0.010977775, 0.036152758, 0.0066366293, 0.11915515, - 0.02318443, -0.041350313, 0.021485701, -0.10906167, - -0.028218046, -0.00954771, 0.020531068, -0.11995105, - -0.03672871, 0.024019798, 0.014255957, -0.05221243, - -0.00661567, -0.04630967, 0.033188973, 0.10107534, - -0.014027541, 0.030796422, -0.10270911, -0.035999842, - 0.15443139, 0.07684145, 0.036571592, -0.035900835, - -0.0034699554, 0.06209149, 0.015920248, -0.031122351, - -0.03858649, 0.01849943, 0.13872518, 0.01503974, - 0.069941424, -0.06948533, -0.0088794185, 0.061282158, - -0.047401894, 0.03100163, -0.041533746, -0.10430945, - 0.044574402, -0.01425562, -0.024290353, 0.034563623, - 0.05866852, 0.023947537, -0.09445152, 0.035450947, - 0.02247216, -0.0042998926, 0.061146557, -0.10250651, - 0.020881841, -0.06747029, 0.10062043, -0.0023941975, - 0.03532124, -0.016341697, 0.09685456, -0.016764693, - 0.051808182, 0.05875331, -0.04536488, 0.001626336, - -0.028892258, -0.01048663, -0.009793449, -0.017093895, - 0.010987891, 0.02357273, -0.00010856845, 0.0099760275, - -0.001845119, -0.03551521, 0.0018358806, 0.05763657, - -0.01769146, 0.040995963, 0.02235177, -0.060430344, - 0.11475477, -0.023854522, 0.10071741, 0.0686208, - -0.014250481, 0.034261297, 0.047418304, 0.08562733, - -0.030519066, 0.0060542435, 0.014653856, -0.038836084, - 0.04096551, 0.032249358, -0.08355519, -0.026823482, - 0.056386515, -0.010401743, -0.028396193, 0.08507674, - 0.014410365, 0.020995233, 0.17040324, 0.11511526, - 0.02459721, 0.0066619175, 0.025853224, -0.023133837, - -0.081302024, 0.017264642, -0.009585969, 0.09491168, - -0.051313367, 0.054532815, -0.014298593, 0.10657464, - 0.007076659, 0.10964551, 0.0409152, 0.008275321, - -0.07283536, 0.07937492, 0.04192024, -0.1075027}; - - recurrent_to_output_weights_ = { - 0.025825322, -0.05813119, 0.09495884, -0.045984812, - -0.01255415, -0.0026479573, -0.08196161, -0.054914974, - -0.0046604523, -0.029587349, -0.044576716, -0.07480124, - -0.082868785, 0.023254942, 0.027502948, -0.0039728214, - -0.08683098, -0.08116779, -0.014675607, -0.037924774, - -0.023314456, -0.007401714, -0.09255757, 0.029460307, - -0.08829125, -0.005139627, -0.08989442, -0.0555066, - 0.13596267, -0.025062224, -0.048351806, -0.03850004, - 0.07266485, -0.022414139, 0.05940088, 0.075114764, - 0.09597592, -0.010211725, -0.0049794707, -0.011523867, - -0.025980417, 0.072999895, 0.11091378, -0.081685916, - 0.014416728, 0.043229222, 0.034178585, -0.07530371, - 0.035837382, -0.085607, -0.007721233, -0.03287832, - -0.043848954, -0.06404588, -0.06632928, -0.073643476, - 0.008214239, -0.045984086, 0.039764922, 0.03474462, - 0.060612556, -0.080590084, 0.049127717, 0.04151091, - -0.030063879, 0.008801774, -0.023021035, -0.019558564, - 0.05158114, -0.010947698, -0.011825728, 0.0075720972, - 0.0699727, -0.0039981045, 0.069350146, 0.08799282, - 0.016156472, 0.035502106, 0.11695009, 0.006217345, - 0.13392477, -0.037875112, 0.025745004, 0.08940699, - -0.00924166, 0.0046702605, -0.036598757, -0.08811812, - 0.10522024, -0.032441203, 0.008176899, -0.04454919, - 0.07058152, 0.0067963637, 0.039206743, 0.03259838, - 0.03725492, -0.09515802, 0.013326398, -0.052055415, - -0.025676316, 0.03198509, -0.015951829, -0.058556724, - 0.036879618, 0.043357447, 0.028362012, -0.05908629, - 0.0059240665, -0.04995891, -0.019187413, 0.0276265, - -0.01628143, 0.0025863599, 0.08800015, 0.035250366, - -0.022165963, -0.07328642, -0.009415526, -0.07455109, - 0.11690406, 0.0363299, 0.07411125, 0.042103454, - -0.009660886, 0.019076364, 0.018299393, -0.046004917, - 0.08891175, 0.0431396, -0.026327137, -0.051502608, - 0.08979574, -0.051670972, 0.04940282, -0.07491107, - -0.021240504, 0.022596184, -0.034280192, 0.060163025, - -0.058211457, -0.051837247, -0.01349775, -0.04639988, - -0.035936575, -0.011681591, 0.064818054, 0.0073146066, - -0.021745546, -0.043124277, -0.06471268, -0.07053354, - -0.029321948, -0.05330136, 0.016933719, -0.053782392, - 0.13747959, -0.1361751, -0.11569455, 0.0033329215, - 0.05693899, -0.053219706, 0.063698, 0.07977434, - -0.07924483, 0.06936997, 0.0034815092, -0.007305279, - -0.037325785, -0.07251102, -0.033633437, -0.08677009, - 0.091591336, -0.14165086, 0.021752775, 0.019683983, - 0.0011612234, -0.058154266, 0.049996935, 0.0288841, - -0.0024567875, -0.14345716, 0.010955264, -0.10234828, - 0.1183656, -0.0010731248, -0.023590032, -0.072285876, - -0.0724771, -0.026382286, -0.0014920527, 0.042667855, - 0.0018776858, 0.02986552, 0.009814309, 0.0733756, - 0.12289186, 0.018043943, -0.0458958, 0.049412545, - 0.033632483, 0.05495232, 0.036686596, -0.013781798, - -0.010036754, 0.02576849, -0.08307328, 0.010112348, - 0.042521734, -0.05869831, -0.071689695, 0.03876447, - -0.13275425, -0.0352966, -0.023077697, 0.10285965, - 0.084736146, 0.15568255, -0.00040734606, 0.027835453, - -0.10292561, -0.032401145, 0.10053256, -0.026142767, - -0.08271222, -0.0030240538, -0.016368777, 0.1070414, - 0.042672627, 0.013456989, -0.0437609, -0.022309763, - 0.11576483, 0.04108048, 0.061026827, -0.0190714, - -0.0869359, 0.037901703, 0.0610107, 0.07202949, - 0.01675338, 0.086139716, -0.08795751, -0.014898893, - -0.023771819, -0.01965048, 0.007955471, -0.043740474, - 0.03346837, -0.10549954, 0.090567775, 0.042013682, - -0.03176985, 0.12569028, -0.02421228, -0.029526481, - 0.023851605, 0.031539805, 0.05292009, -0.02344001, - -0.07811758, -0.08834428, 0.10094801, 0.16594367, - -0.06861939, -0.021256343, -0.041093912, -0.06669611, - 0.035498552, 0.021757556, -0.09302526, -0.015403468, - -0.06614931, -0.051798206, -0.013874718, 0.03630673, - 0.010412845, -0.08077351, 0.046185967, 0.0035662893, - 0.03541868, -0.094149634, -0.034814864, 0.003128424, - -0.020674974, -0.03944324, -0.008110165, -0.11113267, - 0.08484226, 0.043586485, 0.040582247, 0.0968012, - -0.065249965, -0.028036479, 0.0050708856, 0.0017462453, - 0.0326779, 0.041296225, 0.09164146, -0.047743853, - -0.015952192, -0.034451712, 0.084197424, -0.05347844, - -0.11768019, 0.085926116, -0.08251791, -0.045081906, - 0.0948852, 0.068401024, 0.024856757, 0.06978981, - -0.057309967, -0.012775832, -0.0032452994, 0.01977615, - -0.041040014, -0.024264973, 0.063464895, 0.05431621, - }; - - cell_to_input_weights_ = { - 0.040369894, 0.030746894, 0.24704495, 0.018586371, -0.037586458, - -0.15312155, -0.11812848, -0.11465643, 0.20259799, 0.11418174, - -0.10116027, -0.011334949, 0.12411352, -0.076769054, -0.052169047, - 0.21198851, -0.38871562, -0.09061183, -0.09683246, -0.21929175}; - - cell_to_forget_weights_ = { - -0.01998659, -0.15568835, -0.24248174, -0.012770197, 0.041331276, - -0.072311886, -0.052123554, -0.0066330447, -0.043891653, 0.036225766, - -0.047248036, 0.021479502, 0.033189066, 0.11952997, -0.020432774, - 0.64658105, -0.06650122, -0.03467612, 0.095340036, 0.23647355}; - - cell_to_output_weights_ = { - 0.08286371, -0.08261836, -0.51210177, 0.002913762, 0.17764764, - -0.5495371, -0.08460716, -0.24552552, 0.030037103, 0.04123544, - -0.11940523, 0.007358328, 0.1890978, 0.4833202, -0.34441817, - 0.36312827, -0.26375428, 0.1457655, -0.19724406, 0.15548733}; - - projection_weights_ = { - -0.009802181, 0.09401916, 0.0717386, -0.13895074, - 0.09641832, 0.060420845, 0.08539281, 0.054285463, - 0.061395317, 0.034448683, -0.042991187, 0.019801661, - -0.16840284, -0.015726732, -0.23041931, -0.024478018, - -0.10959692, -0.013875541, 0.18600968, -0.061274476, - 0.0138165, -0.08160894, -0.07661644, 0.032372914, - 0.16169067, 0.22465782, -0.03993472, -0.004017731, - 0.08633481, -0.28869787, 0.08682067, 0.17240396, - 0.014975425, 0.056431185, 0.031037588, 0.16702051, - 0.0077946745, 0.15140012, 0.29405436, 0.120285, - -0.188994, -0.027265169, 0.043389652, -0.022061434, - 0.014777949, -0.20203483, 0.094781205, 0.19100232, - 0.13987629, -0.036132768, -0.06426278, -0.05108664, - 0.13221376, 0.009441198, -0.16715929, 0.15859416, - -0.040437475, 0.050779544, -0.022187516, 0.012166504, - 0.027685808, -0.07675938, -0.0055694645, -0.09444123, - 0.0046453946, 0.050794356, 0.10770313, -0.20790008, - -0.07149004, -0.11425117, 0.008225835, -0.035802525, - 0.14374903, 0.15262283, 0.048710253, 0.1847461, - -0.007487823, 0.11000021, -0.09542012, 0.22619456, - -0.029149994, 0.08527916, 0.009043713, 0.0042746216, - 0.016261552, 0.022461696, 0.12689082, -0.043589946, - -0.12035478, -0.08361797, -0.050666027, -0.1248618, - -0.1275799, -0.071875185, 0.07377272, 0.09944291, - -0.18897448, -0.1593054, -0.06526116, -0.040107165, - -0.004618631, -0.067624845, -0.007576253, 0.10727444, - 0.041546922, -0.20424393, 0.06907816, 0.050412357, - 0.00724631, 0.039827548, 0.12449835, 0.10747581, - 0.13708383, 0.09134148, -0.12617786, -0.06428341, - 0.09956831, 0.1208086, -0.14676677, -0.0727722, - 0.1126304, 0.010139365, 0.015571211, -0.038128063, - 0.022913318, -0.042050496, 0.16842307, -0.060597885, - 0.10531834, -0.06411776, -0.07451711, -0.03410368, - -0.13393489, 0.06534304, 0.003620307, 0.04490757, - 0.05970546, 0.05197996, 0.02839995, 0.10434969, - -0.013699693, -0.028353551, -0.07260381, 0.047201227, - -0.024575593, -0.036445823, 0.07155557, 0.009672501, - -0.02328883, 0.009533515, -0.03606021, -0.07421458, - -0.028082801, -0.2678904, -0.13221288, 0.18419984, - -0.13012612, -0.014588381, -0.035059117, -0.04824723, - 0.07830115, -0.056184657, 0.03277091, 0.025466874, - 0.14494097, -0.12522776, -0.098633975, -0.10766018, - -0.08317623, 0.08594209, 0.07749552, 0.039474737, - 0.1776665, -0.07409566, -0.0477268, 0.29323658, - 0.10801441, 0.1154011, 0.013952499, 0.10739139, - 0.10708251, -0.051456142, 0.0074137426, -0.10430189, - 0.10034707, 0.045594677, 0.0635285, -0.0715442, - -0.089667566, -0.10811871, 0.00026344223, 0.08298446, - -0.009525053, 0.006585689, -0.24567553, -0.09450807, - 0.09648481, 0.026996298, -0.06419476, -0.04752702, - -0.11063944, -0.23441927, -0.17608605, -0.052156363, - 0.067035615, 0.19271925, -0.0032889997, -0.043264326, - 0.09663576, -0.057112187, -0.10100678, 0.0628376, - 0.04447668, 0.017961001, -0.10094388, -0.10190601, - 0.18335468, 0.10494553, -0.052095775, -0.0026118709, - 0.10539724, -0.04383912, -0.042349473, 0.08438151, - -0.1947263, 0.02251204, 0.11216432, -0.10307853, - 0.17351969, -0.039091777, 0.08066188, -0.00561982, - 0.12633002, 0.11335965, -0.0088127935, -0.019777594, - 0.06864014, -0.059751723, 0.016233567, -0.06894641, - -0.28651384, -0.004228674, 0.019708522, -0.16305895, - -0.07468996, -0.0855457, 0.099339016, -0.07580735, - -0.13775392, 0.08434318, 0.08330512, -0.12131499, - 0.031935584, 0.09180414, -0.08876437, -0.08049874, - 0.008753825, 0.03498998, 0.030215185, 0.03907079, - 0.089751154, 0.029194152, -0.03337423, -0.019092513, - 0.04331237, 0.04299654, -0.036394123, -0.12915532, - 0.09793732, 0.07512415, -0.11319543, -0.032502122, - 0.15661901, 0.07671967, -0.005491124, -0.19379048, - -0.218606, 0.21448623, 0.017840758, 0.1416943, - -0.07051762, 0.19488361, 0.02664691, -0.18104725, - -0.09334311, 0.15026465, -0.15493552, -0.057762887, - -0.11604192, -0.262013, -0.01391798, 0.012185008, - 0.11156489, -0.07483202, 0.06693364, -0.26151478, - 0.046425626, 0.036540434, -0.16435726, 0.17338543, - -0.21401681, -0.11385144, -0.08283257, -0.069031075, - 0.030635102, 0.010969227, 0.11109743, 0.010919218, - 0.027526086, 0.13519906, 0.01891392, -0.046839405, - -0.040167913, 0.017953383, -0.09700955, 0.0061885654, - -0.07000971, 0.026893595, -0.038844477, 0.14543656}; - - lstm_input_ = { - {// Batch0: 4 (input_sequence_size) * 5 (n_input) - 0.787926, 0.151646, 0.071352, 0.118426, 0.458058, // step 0 - 0.596268, 0.998386, 0.568695, 0.864524, 0.571277, // step 1 - 0.073204, 0.296072, 0.743333, 0.069199, 0.045348, // step 2 - 0.867394, 0.291279, 0.013714, 0.482521, 0.626339}, // step 3 - - {// Batch1: 4 (input_sequence_size) * 5 (n_input) - 0.295743, 0.544053, 0.690064, 0.858138, 0.497181, // step 0 - 0.642421, 0.524260, 0.134799, 0.003639, 0.162482, // step 1 - 0.640394, 0.930399, 0.050782, 0.432485, 0.988078, // step 2 - 0.082922, 0.563329, 0.865614, 0.333232, 0.259916} // step 3 - }; - - lstm_golden_output_ = { - {// Batch0: 4 (input_sequence_size) * 16 (n_output) - -0.00396806, 0.029352, -0.00279226, 0.0159977, -0.00835576, - -0.0211779, 0.0283512, -0.0114597, 0.00907307, -0.0244004, - -0.0152191, -0.0259063, 0.00914318, 0.00415118, 0.017147, - 0.0134203, -0.0166936, 0.0381209, 0.000889694, 0.0143363, - -0.0328911, -0.0234288, 0.0333051, -0.012229, 0.0110322, - -0.0457725, -0.000832209, -0.0202817, 0.0327257, 0.0121308, - 0.0155969, 0.0312091, -0.0213783, 0.0350169, 0.000324794, - 0.0276012, -0.0263374, -0.0371449, 0.0446149, -0.0205474, - 0.0103729, -0.0576349, -0.0150052, -0.0292043, 0.0376827, - 0.0136115, 0.0243435, 0.0354492, -0.0189322, 0.0464512, - -0.00251373, 0.0225745, -0.0308346, -0.0317124, 0.0460407, - -0.0189395, 0.0149363, -0.0530162, -0.0150767, -0.0340193, - 0.0286833, 0.00824207, 0.0264887, 0.0305169}, - {// Batch1: 4 (input_sequence_size) * 16 (n_output) - -0.013869, 0.0287268, -0.00334693, 0.00733398, -0.0287926, - -0.0186926, 0.0193662, -0.0115437, 0.00422612, -0.0345232, - 0.00223253, -0.00957321, 0.0210624, 0.013331, 0.0150954, - 0.02168, -0.0141913, 0.0322082, 0.00227024, 0.0260507, - -0.0188721, -0.0296489, 0.0399134, -0.0160509, 0.0116039, - -0.0447318, -0.0150515, -0.0277406, 0.0316596, 0.0118233, - 0.0214762, 0.0293641, -0.0204549, 0.0450315, -0.00117378, - 0.0167673, -0.0375007, -0.0238314, 0.038784, -0.0174034, - 0.0131743, -0.0506589, -0.0048447, -0.0240239, 0.0325789, - 0.00790065, 0.0220157, 0.0333314, -0.0264787, 0.0387855, - -0.000764675, 0.0217599, -0.037537, -0.0335206, 0.0431679, - -0.0211424, 0.010203, -0.062785, -0.00832363, -0.025181, - 0.0412031, 0.0118723, 0.0239643, 0.0394009}}; - } -}; - -TEST_F(NoCifgPeepholeProjectionNoClippingLstmTest, LstmBlackBoxTest) { - const int n_batch = 2; - const int n_input = 5; - const int n_cell = 20; - const int n_output = 16; - - LSTMOpModel lstm(n_batch, n_input, n_cell, n_output, - /*use_cifg=*/false, /*use_peephole=*/true, - /*use_projection_weights=*/true, - /*use_projection_bias=*/false, - /*cell_clip=*/0.0, /*proj_clip=*/0.0, - { - {n_batch, n_input}, // input tensor - - {n_cell, n_input}, // input_to_input_weight tensor - {n_cell, n_input}, // input_to_forget_weight tensor - {n_cell, n_input}, // input_to_cell_weight tensor - {n_cell, n_input}, // input_to_output_weight tensor - - {n_cell, n_output}, // recurrent_to_input_weight tensor - {n_cell, n_output}, // recurrent_to_forget_weight tensor - {n_cell, n_output}, // recurrent_to_cell_weight tensor - {n_cell, n_output}, // recurrent_to_output_weight tensor - - {n_cell}, // cell_to_input_weight tensor - {n_cell}, // cell_to_forget_weight tensor - {n_cell}, // cell_to_output_weight tensor - - {n_cell}, // input_gate_bias tensor - {n_cell}, // forget_gate_bias tensor - {n_cell}, // cell_bias tensor - {n_cell}, // output_gate_bias tensor - - {n_output, n_cell}, // projection_weight tensor - {0}, // projection_bias tensor - }); - - lstm.SetInputToInputWeights(input_to_input_weights_); - lstm.SetInputToCellWeights(input_to_cell_weights_); - lstm.SetInputToForgetWeights(input_to_forget_weights_); - lstm.SetInputToOutputWeights(input_to_output_weights_); - - lstm.SetInputGateBias(input_gate_bias_); - lstm.SetCellBias(cell_gate_bias_); - lstm.SetForgetGateBias(forget_gate_bias_); - lstm.SetOutputGateBias(output_gate_bias_); - - lstm.SetRecurrentToInputWeights(recurrent_to_input_weights_); - lstm.SetRecurrentToCellWeights(recurrent_to_cell_weights_); - lstm.SetRecurrentToForgetWeights(recurrent_to_forget_weights_); - lstm.SetRecurrentToOutputWeights(recurrent_to_output_weights_); - - lstm.SetCellToInputWeights(cell_to_input_weights_); - lstm.SetCellToForgetWeights(cell_to_forget_weights_); - lstm.SetCellToOutputWeights(cell_to_output_weights_); - - lstm.SetProjectionWeights(projection_weights_); - - VerifyGoldens(lstm_input_, lstm_golden_output_, &lstm); -} - -TEST_F(NoCifgPeepholeProjectionNoClippingLstmTest, HybridLstmBlackBoxTest) { - const int n_batch = 2; - const int n_input = 5; - const int n_cell = 20; - const int n_output = 16; - - HybridLSTMOpModel lstm( - n_batch, n_input, n_cell, n_output, - /*use_cifg=*/false, /*use_peephole=*/true, - /*use_projection_weights=*/true, - /*use_projection_bias=*/false, - /*cell_clip=*/0.0, /*proj_clip=*/0.0, - { - {n_batch, n_input}, // input tensor - - {n_cell, n_input}, // input_to_input_weight tensor - {n_cell, n_input}, // input_to_forget_weight tensor - {n_cell, n_input}, // input_to_cell_weight tensor - {n_cell, n_input}, // input_to_output_weight tensor - - {n_cell, n_output}, // recurrent_to_input_weight tensor - {n_cell, n_output}, // recurrent_to_forget_weight tensor - {n_cell, n_output}, // recurrent_to_cell_weight tensor - {n_cell, n_output}, // recurrent_to_output_weight tensor - - {n_cell}, // cell_to_input_weight tensor - {n_cell}, // cell_to_forget_weight tensor - {n_cell}, // cell_to_output_weight tensor - - {n_cell}, // input_gate_bias tensor - {n_cell}, // forget_gate_bias tensor - {n_cell}, // cell_bias tensor - {n_cell}, // output_gate_bias tensor - - {n_output, n_cell}, // projection_weight tensor - {0}, // projection_bias tensor - }); - - lstm.SetInputToInputWeights(input_to_input_weights_); - lstm.SetInputToCellWeights(input_to_cell_weights_); - lstm.SetInputToForgetWeights(input_to_forget_weights_); - lstm.SetInputToOutputWeights(input_to_output_weights_); - - lstm.SetInputGateBias(input_gate_bias_); - lstm.SetCellBias(cell_gate_bias_); - lstm.SetForgetGateBias(forget_gate_bias_); - lstm.SetOutputGateBias(output_gate_bias_); - - lstm.SetRecurrentToInputWeights(recurrent_to_input_weights_); - lstm.SetRecurrentToCellWeights(recurrent_to_cell_weights_); - lstm.SetRecurrentToForgetWeights(recurrent_to_forget_weights_); - lstm.SetRecurrentToOutputWeights(recurrent_to_output_weights_); - - lstm.SetCellToInputWeights(cell_to_input_weights_); - lstm.SetCellToForgetWeights(cell_to_forget_weights_); - lstm.SetCellToOutputWeights(cell_to_output_weights_); - - lstm.SetProjectionWeights(projection_weights_); - - VerifyGoldens(lstm_input_, lstm_golden_output_, &lstm, /*tolerance=*/0.00467); -} - -} // namespace -} // namespace tflite - -int main(int argc, char** argv) { - ::tflite::LogToStderr(); - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/contrib/lite/kernels/mfcc.cc b/tensorflow/contrib/lite/kernels/mfcc.cc deleted file mode 100644 index 5153ce5634c33e829c3742e4d11a22a18f0d2f79..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/mfcc.cc +++ /dev/null @@ -1,154 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include "tensorflow/contrib/lite/kernels/internal/mfcc.h" -#include "flatbuffers/flexbuffers.h" // TF:flatbuffers -#include "tensorflow/contrib/lite/c/builtin_op_data.h" -#include "tensorflow/contrib/lite/c/c_api_internal.h" -#include "tensorflow/contrib/lite/kernels/internal/mfcc_dct.h" -#include "tensorflow/contrib/lite/kernels/internal/mfcc_mel_filterbank.h" -#include "tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h" -#include "tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h" -#include "tensorflow/contrib/lite/kernels/internal/tensor.h" -#include "tensorflow/contrib/lite/kernels/kernel_util.h" -#include "tensorflow/contrib/lite/kernels/op_macros.h" - -namespace tflite { -namespace ops { -namespace custom { -namespace mfcc { - -enum KernelType { - kReference, -}; - -typedef struct { - float upper_frequency_limit; - float lower_frequency_limit; - int filterbank_channel_count; - int dct_coefficient_count; -} TfLiteMfccParams; - -constexpr int kInputTensorWav = 0; -constexpr int kInputTensorRate = 1; -constexpr int kOutputTensor = 0; - -void* Init(TfLiteContext* context, const char* buffer, size_t length) { - auto* data = new TfLiteMfccParams; - - const uint8_t* buffer_t = reinterpret_cast(buffer); - - const flexbuffers::Map& m = flexbuffers::GetRoot(buffer_t, length).AsMap(); - data->upper_frequency_limit = m["upper_frequency_limit"].AsInt64(); - data->lower_frequency_limit = m["lower_frequency_limit"].AsInt64(); - data->filterbank_channel_count = m["filterbank_channel_count"].AsInt64(); - data->dct_coefficient_count = m["dct_coefficient_count"].AsInt64(); - return data; -} - -void Free(TfLiteContext* context, void* buffer) { - delete reinterpret_cast(buffer); -} - -TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { - auto* params = reinterpret_cast(node->user_data); - - TF_LITE_ENSURE_EQ(context, NumInputs(node), 2); - TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); - - const TfLiteTensor* inputWav = GetInput(context, node, kInputTensorWav); - const TfLiteTensor* inputRate = GetInput(context, node, kInputTensorRate); - TfLiteTensor* output = GetOutput(context, node, kOutputTensor); - - TF_LITE_ENSURE_EQ(context, NumDimensions(inputWav), 3); - TF_LITE_ENSURE_EQ(context, NumDimensions(inputRate), 1); - - TF_LITE_ENSURE_EQ(context, output->type, kTfLiteFloat32); - TF_LITE_ENSURE_EQ(context, inputWav->type, output->type); - - TfLiteIntArray* output_size = TfLiteIntArrayCreate(3); - output_size->data[0] = inputWav->dims->data[0]; - output_size->data[1] = inputWav->dims->data[1]; - output_size->data[2] = params->dct_coefficient_count; - - return context->ResizeTensor(context, output, output_size); -} - -// Input is a single squared-magnitude spectrogram frame. The input spectrum -// is converted to linear magnitude and weighted into bands using a -// triangular mel filterbank, and a discrete cosine transform (DCT) of the -// values is taken. Output is populated with the lowest dct_coefficient_count -// of these values. -template -TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { - auto* params = reinterpret_cast(node->user_data); - - const TfLiteTensor* inputWav = GetInput(context, node, kInputTensorWav); - const TfLiteTensor* inputRate = GetInput(context, node, kInputTensorRate); - TfLiteTensor* output = GetOutput(context, node, kOutputTensor); - - const int32 sample_rate = *GetTensorData(inputRate); - - const int spectrogram_channels = inputWav->dims->data[2]; - const int spectrogram_samples = inputWav->dims->data[1]; - const int audio_channels = inputWav->dims->data[0]; - - internal::Mfcc mfcc; - mfcc.set_upper_frequency_limit(params->upper_frequency_limit); - mfcc.set_lower_frequency_limit(params->lower_frequency_limit); - mfcc.set_filterbank_channel_count(params->filterbank_channel_count); - mfcc.set_dct_coefficient_count(params->dct_coefficient_count); - - mfcc.Initialize(spectrogram_channels, sample_rate); - - const float* spectrogram_flat = GetTensorData(inputWav); - float* output_flat = GetTensorData(output); - - for (int audio_channel = 0; audio_channel < audio_channels; ++audio_channel) { - for (int spectrogram_sample = 0; spectrogram_sample < spectrogram_samples; - ++spectrogram_sample) { - const float* sample_data = - spectrogram_flat + - (audio_channel * spectrogram_samples * spectrogram_channels) + - (spectrogram_sample * spectrogram_channels); - std::vector mfcc_input(sample_data, - sample_data + spectrogram_channels); - std::vector mfcc_output; - mfcc.Compute(mfcc_input, &mfcc_output); - TF_LITE_ENSURE_EQ(context, params->dct_coefficient_count, - mfcc_output.size()); - float* output_data = output_flat + - (audio_channel * spectrogram_samples * - params->dct_coefficient_count) + - (spectrogram_sample * params->dct_coefficient_count); - for (int i = 0; i < params->dct_coefficient_count; ++i) { - output_data[i] = mfcc_output[i]; - } - } - } - - return kTfLiteOk; -} - -} // namespace mfcc - -TfLiteRegistration* Register_MFCC() { - static TfLiteRegistration r = {mfcc::Init, mfcc::Free, mfcc::Prepare, - mfcc::Eval}; - return &r; -} - -} // namespace custom -} // namespace ops -} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/pack_test.cc b/tensorflow/contrib/lite/kernels/pack_test.cc deleted file mode 100644 index c70dbd2764b615530a9587b521a3616eece92cb6..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/pack_test.cc +++ /dev/null @@ -1,154 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include -#include "tensorflow/contrib/lite/interpreter.h" -#include "tensorflow/contrib/lite/kernels/register.h" -#include "tensorflow/contrib/lite/kernels/test_util.h" -#include "tensorflow/contrib/lite/model.h" - -namespace tflite { -namespace { - -using ::testing::ElementsAre; -using ::testing::ElementsAreArray; - -template -class PackOpModel : public SingleOpModel { - public: - PackOpModel(const TensorData& input_template, int axis, int values_count) { - std::vector> all_input_shapes; - for (int i = 0; i < values_count; ++i) { - all_input_shapes.push_back(input_template.shape); - AddInput(input_template); - } - output_ = AddOutput({input_template.type, /*shape=*/{}, input_template.min, - input_template.max}); - SetBuiltinOp(BuiltinOperator_PACK, BuiltinOptions_PackOptions, - CreatePackOptions(builder_, values_count, axis).Union()); - BuildInterpreter(all_input_shapes); - } - - void SetInput(int index, std::initializer_list data) { - PopulateTensor(index, data); - } - - std::vector GetOutput() { return ExtractVector(output_); } - std::vector GetOutputShape() { return GetTensorShape(output_); } - - private: - int output_; -}; - -// float32 tests. -TEST(PackOpTest, FloatThreeInputs) { - PackOpModel model({TensorType_FLOAT32, {2}}, 0, 3); - model.SetInput(0, {1, 4}); - model.SetInput(1, {2, 5}); - model.SetInput(2, {3, 6}); - model.Invoke(); - EXPECT_THAT(model.GetOutputShape(), ElementsAre(3, 2)); - EXPECT_THAT(model.GetOutput(), ElementsAreArray({1, 4, 2, 5, 3, 6})); -} - -TEST(PackOpTest, FloatThreeInputsDifferentAxis) { - PackOpModel model({TensorType_FLOAT32, {2}}, 1, 3); - model.SetInput(0, {1, 4}); - model.SetInput(1, {2, 5}); - model.SetInput(2, {3, 6}); - model.Invoke(); - EXPECT_THAT(model.GetOutputShape(), ElementsAre(2, 3)); - EXPECT_THAT(model.GetOutput(), ElementsAreArray({1, 2, 3, 4, 5, 6})); -} - -TEST(PackOpTest, FloatMultilDimensions) { - PackOpModel model({TensorType_FLOAT32, {2, 3}}, 1, 2); - model.SetInput(0, {1, 2, 3, 4, 5, 6}); - model.SetInput(1, {7, 8, 9, 10, 11, 12}); - model.Invoke(); - EXPECT_THAT(model.GetOutputShape(), ElementsAre(2, 2, 3)); - EXPECT_THAT(model.GetOutput(), - ElementsAreArray({1, 2, 3, 7, 8, 9, 4, 5, 6, 10, 11, 12})); -} - -// int32 tests. -TEST(PackOpTest, Int32ThreeInputs) { - PackOpModel model({TensorType_INT32, {2}}, 0, 3); - model.SetInput(0, {1, 4}); - model.SetInput(1, {2, 5}); - model.SetInput(2, {3, 6}); - model.Invoke(); - EXPECT_THAT(model.GetOutputShape(), ElementsAre(3, 2)); - EXPECT_THAT(model.GetOutput(), ElementsAreArray({1, 4, 2, 5, 3, 6})); -} - -TEST(PackOpTest, Int32ThreeInputsDifferentAxis) { - PackOpModel model({TensorType_INT32, {2}}, 1, 3); - model.SetInput(0, {1, 4}); - model.SetInput(1, {2, 5}); - model.SetInput(2, {3, 6}); - model.Invoke(); - EXPECT_THAT(model.GetOutputShape(), ElementsAre(2, 3)); - EXPECT_THAT(model.GetOutput(), ElementsAreArray({1, 2, 3, 4, 5, 6})); -} - -TEST(PackOpTest, Int32MultilDimensions) { - PackOpModel model({TensorType_INT32, {2, 3}}, 1, 2); - model.SetInput(0, {1, 2, 3, 4, 5, 6}); - model.SetInput(1, {7, 8, 9, 10, 11, 12}); - model.Invoke(); - EXPECT_THAT(model.GetOutputShape(), ElementsAre(2, 2, 3)); - EXPECT_THAT(model.GetOutput(), - ElementsAreArray({1, 2, 3, 7, 8, 9, 4, 5, 6, 10, 11, 12})); -} - -// uint8 -TEST(PackOpTest, Uint8ThreeInputs) { - PackOpModel model({TensorType_UINT8, {2}}, 0, 3); - model.SetInput(0, {1, 4}); - model.SetInput(1, {2, 5}); - model.SetInput(2, {3, 6}); - model.Invoke(); - EXPECT_THAT(model.GetOutputShape(), ElementsAre(3, 2)); - EXPECT_THAT(model.GetOutput(), ElementsAreArray({1, 4, 2, 5, 3, 6})); -} - -TEST(PackOpTest, Uint8ThreeInputsDifferentAxis) { - PackOpModel model({TensorType_UINT8, {2}}, 1, 3); - model.SetInput(0, {1, 4}); - model.SetInput(1, {2, 5}); - model.SetInput(2, {3, 6}); - model.Invoke(); - EXPECT_THAT(model.GetOutputShape(), ElementsAre(2, 3)); - EXPECT_THAT(model.GetOutput(), ElementsAreArray({1, 2, 3, 4, 5, 6})); -} - -TEST(PackOpTest, Uint8MultilDimensions) { - PackOpModel model({TensorType_UINT8, {2, 3}}, 1, 2); - model.SetInput(0, {1, 2, 3, 4, 5, 6}); - model.SetInput(1, {7, 8, 9, 10, 11, 12}); - model.Invoke(); - EXPECT_THAT(model.GetOutputShape(), ElementsAre(2, 2, 3)); - EXPECT_THAT(model.GetOutput(), - ElementsAreArray({1, 2, 3, 7, 8, 9, 4, 5, 6, 10, 11, 12})); -} - -} // namespace -} // namespace tflite - -int main(int argc, char** argv) { - ::tflite::LogToStderr(); - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/contrib/lite/kernels/pooling_test.cc b/tensorflow/contrib/lite/kernels/pooling_test.cc deleted file mode 100644 index 01c91b2ba905e249c36af19f175c68a7e7f17f6d..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/pooling_test.cc +++ /dev/null @@ -1,161 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include -#include -#include "tensorflow/contrib/lite/interpreter.h" -#include "tensorflow/contrib/lite/kernels/register.h" -#include "tensorflow/contrib/lite/kernels/test_util.h" -#include "tensorflow/contrib/lite/model.h" - -namespace tflite { -namespace { - -using ::testing::ElementsAreArray; - -class BasePoolingOpModel : public SingleOpModel { - public: - // TODO(ahentz): Also test different activation types, bias, padding types, - // stride values. - BasePoolingOpModel(BuiltinOperator type, const TensorData& input, - int filter_width, int filter_height, - const TensorData& output) { - input_ = AddInput(input); - output_ = AddOutput(output); - - SetBuiltinOp( - type, BuiltinOptions_Pool2DOptions, - CreatePool2DOptions(builder_, Padding_VALID, 2, 2, filter_width, - filter_height, ActivationFunctionType_NONE) - .Union()); - - BuildInterpreter({GetShape(input_)}); - } - - protected: - int input_; - int output_; -}; - -class FloatPoolingOpModel : public BasePoolingOpModel { - public: - using BasePoolingOpModel::BasePoolingOpModel; - - void SetInput(std::initializer_list data) { - PopulateTensor(input_, data); - } - - std::vector GetOutput() { return ExtractVector(output_); } -}; - -class QuantizedPoolingOpModel : public BasePoolingOpModel { - public: - using BasePoolingOpModel::BasePoolingOpModel; - - void SetInput(std::initializer_list data) { - QuantizeAndPopulate(input_, data); - } - - std::vector GetOutput() { return ExtractVector(output_); } - std::vector GetDequantizedOutput() { - return Dequantize(ExtractVector(output_), - GetScale(output_), GetZeroPoint(output_)); - } -}; - -TEST(FloatPoolingOpTest, AveragePool) { - FloatPoolingOpModel m(BuiltinOperator_AVERAGE_POOL_2D, - /*input=*/{TensorType_FLOAT32, {1, 2, 4, 1}}, - /*filter_width=*/2, /*filter_height=*/2, - /*output=*/{TensorType_FLOAT32, {}}); - m.SetInput({ - 0, 6, 2, 4, // - 3, 2, 10, 7, // - }); - m.Invoke(); - EXPECT_THAT(m.GetOutput(), ElementsAreArray({2.75, 5.75})); -} - -TEST(QuantizedPoolingOpTest, AveragePool) { - // Choose the input ranges carefully so that the dequantized output matches - // the results of the float model above. - QuantizedPoolingOpModel m( - BuiltinOperator_AVERAGE_POOL_2D, - /*input=*/{TensorType_UINT8, {1, 2, 4, 1}, 0, 15.9375}, - /*filter_width=*/2, /*filter_height=*/2, - /*output=*/{TensorType_UINT8, {}, 0, 15.9375}); - m.SetInput({ - 0, 6, 2, 4, // - 3, 2, 10, 7, // - }); - m.Invoke(); - - EXPECT_THAT(m.GetDequantizedOutput(), - ElementsAreArray(ArrayFloatNear({2.75, 5.75}))); - EXPECT_THAT(m.GetOutput(), ElementsAreArray({44, 92})); -} - -TEST(FloatPoolingOpTest, MaxPool) { - FloatPoolingOpModel m(BuiltinOperator_MAX_POOL_2D, - /*input=*/{TensorType_FLOAT32, {1, 2, 4, 1}}, - /*filter_width=*/2, /*filter_height=*/2, - /*output=*/{TensorType_FLOAT32, {}}); - m.SetInput({ - 0, 6, 2, 4, // - 3, 2, 10, 7, // - }); - m.Invoke(); - EXPECT_THAT(m.GetOutput(), ElementsAreArray({6, 10})); -} - -TEST(QuantizedPoolingOpTest, MaxPool) { - // Choose the input ranges carefully so that the dequantized output matches - // the results of the float model above. - QuantizedPoolingOpModel m( - BuiltinOperator_MAX_POOL_2D, - /*input=*/{TensorType_UINT8, {1, 2, 4, 1}, 0, 15.9375}, - /*filter_width=*/2, /*filter_height=*/2, - /*output=*/{TensorType_UINT8, {}, 0, 15.9375}); - m.SetInput({ - 0, 6, 2, 4, // - 3, 2, 10, 7, // - }); - m.Invoke(); - - EXPECT_THAT(m.GetDequantizedOutput(), - ElementsAreArray(ArrayFloatNear({6, 10}))); - EXPECT_THAT(m.GetOutput(), ElementsAreArray({96, 160})); -} - -TEST(FloatPoolingOpTest, L2Pool) { - FloatPoolingOpModel m(BuiltinOperator_L2_POOL_2D, - /*input=*/{TensorType_FLOAT32, {1, 2, 4, 1}}, - /*filter_width=*/2, /*filter_height=*/2, - /*output=*/{TensorType_FLOAT32, {}}); - m.SetInput({ - 0, 6, 2, 4, // - 3, 2, 10, 7, // - }); - m.Invoke(); - EXPECT_THAT(m.GetOutput(), ElementsAreArray({3.5, 6.5})); -} - -} // namespace -} // namespace tflite - -int main(int argc, char** argv) { - ::tflite::LogToStderr(); - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/contrib/lite/kernels/reshape.cc b/tensorflow/contrib/lite/kernels/reshape.cc deleted file mode 100644 index f41147b2d6433addc63538fbd8c4338d749535d3..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/reshape.cc +++ /dev/null @@ -1,133 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include -#include "tensorflow/contrib/lite/c/builtin_op_data.h" -#include "tensorflow/contrib/lite/c/c_api_internal.h" -#include "tensorflow/contrib/lite/kernels/internal/tensor.h" -#include "tensorflow/contrib/lite/kernels/kernel_util.h" -#include "tensorflow/contrib/lite/kernels/op_macros.h" - -namespace tflite { -namespace ops { -namespace builtin { -namespace reshape { - -constexpr int kInputTensor = 0; -constexpr int kShapeTensor = 1; -constexpr int kOutputTensor = 0; - -TfLiteStatus ResizeOutput(TfLiteContext* context, TfLiteNode* node, - TfLiteIntArray* output_shape) { - const TfLiteTensor* input = GetInput(context, node, kInputTensor); - TfLiteTensor* output = GetOutput(context, node, kOutputTensor); - - // Tensorflow's Reshape allows one of the shape components to have the - // special -1 value, meaning it will be calculated automatically based on the - // input. Here we calculate what that dimension should be so that the number - // of output elements in the same as the number of input elements. - int num_input_elements = NumElements(input); - - int num_output_elements = 1; - int stretch_dim = -1; - for (int i = 0; i < output_shape->size; ++i) { - int value = output_shape->data[i]; - if (value == -1) { - TF_LITE_ENSURE_EQ(context, stretch_dim, -1); - stretch_dim = i; - } else { - num_output_elements *= value; - } - } - if (stretch_dim != -1) { - output_shape->data[stretch_dim] = num_input_elements / num_output_elements; - num_output_elements *= output_shape->data[stretch_dim]; - } - - TF_LITE_ENSURE_EQ(context, num_input_elements, num_output_elements); - return context->ResizeTensor(context, output, output_shape); -} - -TfLiteStatus ResizeOutputWithShapeTensor(TfLiteContext* context, - TfLiteNode* node) { - const TfLiteTensor* shape = GetInput(context, node, kShapeTensor); - - TfLiteIntArray* output_shape = TfLiteIntArrayCreate(shape->dims->data[0]); - for (int i = 0; i < output_shape->size; ++i) { - output_shape->data[i] = shape->data.i32[i]; - } - return ResizeOutput(context, node, output_shape); -} - -TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { - auto* params = reinterpret_cast(node->builtin_data); - - TF_LITE_ENSURE(context, NumInputs(node) == 1 || NumInputs(node) == 2); - TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); - - // Attempt to use shape tensor if it exists. - if (NumInputs(node) == 2) { - const TfLiteTensor* shape = GetInput(context, node, kShapeTensor); - // Check if the shape tensor is valid. - if (shape->dims->size == 1 && shape->type == kTfLiteInt32) { - // Set the output tensor as dynamic if the shape isn't constnat. - if (!IsConstantTensor(shape)) { - TfLiteTensor* output = GetOutput(context, node, kOutputTensor); - SetTensorToDynamic(output); - return kTfLiteOk; - } - // Shape is constant. Resize now. - return ResizeOutputWithShapeTensor(context, node); - } - } - // The function is returned above this line if the shape tensor is usable. - // Now fallback to the shape parameter in `TfLiteReshapeParams`. - int num_dimensions = params->num_dimensions; - if (num_dimensions == 1 && params->shape[0] == 0) { - // Legacy tflite models use a shape parameter of [0] to indicate scalars, - // so adjust accordingly. TODO(b/111614235): Allow zero-sized buffers during - // toco conversion. - num_dimensions = 0; - } - TfLiteIntArray* output_shape = TfLiteIntArrayCreate(num_dimensions); - for (int i = 0; i < num_dimensions; ++i) { - output_shape->data[i] = params->shape[i]; - } - return ResizeOutput(context, node, output_shape); -} - -TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { - const TfLiteTensor* input = GetInput(context, node, kInputTensor); - TfLiteTensor* output = GetOutput(context, node, kOutputTensor); - - if (IsDynamicTensor(output)) { - TF_LITE_ENSURE_OK(context, ResizeOutputWithShapeTensor(context, node)); - } - - memcpy(output->data.raw, input->data.raw, input->bytes); - - return kTfLiteOk; -} - -} // namespace reshape - -TfLiteRegistration* Register_RESHAPE() { - static TfLiteRegistration r = {nullptr, nullptr, reshape::Prepare, - reshape::Eval}; - return &r; -} - -} // namespace builtin -} // namespace ops -} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/reshape_test.cc b/tensorflow/contrib/lite/kernels/reshape_test.cc deleted file mode 100644 index 52d71350d3ba9a27bf9a8df7a194161c4fb7f87c..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/reshape_test.cc +++ /dev/null @@ -1,122 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include -#include "tensorflow/contrib/lite/interpreter.h" -#include "tensorflow/contrib/lite/kernels/register.h" -#include "tensorflow/contrib/lite/kernels/test_util.h" -#include "tensorflow/contrib/lite/model.h" - -namespace tflite { -namespace { - -using ::testing::ElementsAreArray; -using ::testing::IsEmpty; - -class ReshapeOpModel : public SingleOpModel { - public: - ReshapeOpModel(std::initializer_list input_shape, - std::initializer_list new_shape, - bool use_shape_input_tensor = false) { - input_ = AddInput(TensorType_FLOAT32); - output_ = AddOutput(TensorType_FLOAT32); - int shape_input_tensor = - use_shape_input_tensor ? AddInput(TensorType_INT32) : -1; - SetBuiltinOp( - BuiltinOperator_RESHAPE, BuiltinOptions_ReshapeOptions, - CreateReshapeOptions(builder_, builder_.CreateVector(new_shape)) - .Union()); - if (use_shape_input_tensor) { - BuildInterpreter({input_shape, GetShape(shape_input_tensor)}); - PopulateTensor(shape_input_tensor, new_shape); - } else { - BuildInterpreter({input_shape}); - } - } - - void SetInput(std::initializer_list data) { - PopulateTensor(input_, data); - } - std::vector GetOutput() { return ExtractVector(output_); } - std::vector GetOutputShape() { return GetTensorShape(output_); } - - private: - int input_; - int output_; -}; - -TEST(ReshapeOpTest, MismatchedDimensions) { - EXPECT_DEATH(ReshapeOpModel({1, 2, 4, 1}, {2, 1}), - "num_input_elements != num_output_elements"); -} - -TEST(ReshapeOpTest, TooManyDimensions) { - EXPECT_DEATH( - ReshapeOpModel({1, 2, 3, 4, 5, 6, 7, 8, 9}, {1, 2, 3, 4, 5, 6, 7, 8, 9}), - "Found too many dimensions"); -} - -TEST(ReshapeOpTest, TooManySpecialDimensions) { - EXPECT_DEATH(ReshapeOpModel({1, 2, 4, 1}, {-1, -1, 2, 4}), - "stretch_dim != -1"); -} - -TEST(ReshapeOpTest, SimpleTest) { - ReshapeOpModel m({1, 2, 4, 1}, {2, 2, 2}); - m.SetInput({1, 2, 3, 4, 5, 6, 7, 8}); - m.Invoke(); - EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 2, 3, 4, 5, 6, 7, 8})); - EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2, 2, 2})); -} - -TEST(ReshapeOpTest, ShapeTensorInput) { - ReshapeOpModel m({1, 2, 4, 1}, {2, 2, 2}, /*use_shape_input_tensor=*/true); - m.SetInput({1, 2, 3, 4, 5, 6, 7, 8}); - m.Invoke(); - EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 2, 3, 4, 5, 6, 7, 8})); - EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2, 2, 2})); -} - -TEST(ReshapeOpTest, WithStretchDimension) { - ReshapeOpModel m({1, 2, 4, 1}, {2, 1, -1}); - m.SetInput({1, 2, 3, 4, 5, 6, 7, 8}); - m.Invoke(); - EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 2, 3, 4, 5, 6, 7, 8})); - EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2, 1, 4})); -} - -TEST(ReshapeOpTest, ScalarOutput) { - ReshapeOpModel m({1}, {}); - m.SetInput({3}); - m.Invoke(); - EXPECT_THAT(m.GetOutput(), ElementsAreArray({3})); - EXPECT_THAT(m.GetOutputShape(), IsEmpty()); -} - -TEST(ReshapeOpTest, LegacyScalarOutput) { - ReshapeOpModel m({1}, {0}); - m.SetInput({3}); - m.Invoke(); - EXPECT_THAT(m.GetOutput(), ElementsAreArray({3})); - EXPECT_THAT(m.GetOutputShape(), IsEmpty()); -} - -} // namespace -} // namespace tflite - -int main(int argc, char** argv) { - ::tflite::LogToStderr(); - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/contrib/lite/kernels/resize_bilinear_test.cc b/tensorflow/contrib/lite/kernels/resize_bilinear_test.cc deleted file mode 100644 index f4289105f7931ae572f219a61b5479287aff926a..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/resize_bilinear_test.cc +++ /dev/null @@ -1,314 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include -#include "tensorflow/contrib/lite/interpreter.h" -#include "tensorflow/contrib/lite/kernels/register.h" -#include "tensorflow/contrib/lite/kernels/test_util.h" -#include "tensorflow/contrib/lite/model.h" - -namespace tflite { -namespace { - -using ::testing::ElementsAreArray; -using uint8 = std::uint8_t; - -class ResizeBilinearOpModel : public SingleOpModel { - public: - ResizeBilinearOpModel(const TensorData& input, - std::initializer_list size_data = {}) { - bool const_size = size_data.size() != 0; - input_ = AddInput(input); - if (const_size) { - size_ = AddConstInput(TensorType_INT32, size_data, {2}); - } else { - size_ = AddInput({TensorType_INT32, {2}}); - } - output_ = AddOutput(input.type); - SetBuiltinOp(BuiltinOperator_RESIZE_BILINEAR, - BuiltinOptions_ResizeBilinearOptions, - CreateResizeBilinearOptions(builder_).Union()); - if (const_size) { - BuildInterpreter({GetShape(input_)}); - } else { - BuildInterpreter({GetShape(input_), GetShape(size_)}); - } - } - - template - void SetInput(std::initializer_list data) { - PopulateTensor(input_, data); - } - void SetSize(std::initializer_list data) { PopulateTensor(size_, data); } - - template - std::vector GetOutput() { - return ExtractVector(output_); - } - - private: - int input_; - int size_; - int output_; -}; - -TEST(ResizeBilinearOpTest, HorizontalResize) { - ResizeBilinearOpModel m({TensorType_FLOAT32, {1, 1, 2, 1}}); - m.SetInput({3, 6}); - m.SetSize({1, 3}); - m.Invoke(); - EXPECT_THAT(m.GetOutput(), - ElementsAreArray(ArrayFloatNear({3, 5, 6}))); - - ResizeBilinearOpModel const_m({TensorType_FLOAT32, {1, 1, 2, 1}}, {1, 3}); - const_m.SetInput({3, 6}); - const_m.Invoke(); - EXPECT_THAT(const_m.GetOutput(), - ElementsAreArray(ArrayFloatNear({3, 5, 6}))); -} - -TEST(ResizeBilinearOpTest, HorizontalResize8Bit) { - ResizeBilinearOpModel m({TensorType_UINT8, {1, 1, 2, 1}}); - m.SetInput({3, 6}); - m.SetSize({1, 3}); - m.Invoke(); - EXPECT_THAT(m.GetOutput(), - ElementsAreArray(ArrayFloatNear({3, 5, 6}))); - - ResizeBilinearOpModel const_m({TensorType_UINT8, {1, 1, 2, 1}}, {1, 3}); - const_m.SetInput({3, 6}); - const_m.Invoke(); - EXPECT_THAT(const_m.GetOutput(), - ElementsAreArray(ArrayFloatNear({3, 5, 6}))); -} - -TEST(ResizeBilinearOpTest, VerticalResize) { - ResizeBilinearOpModel m({TensorType_FLOAT32, {1, 2, 1, 1}}); - m.SetInput({3, 9}); - m.SetSize({3, 1}); - m.Invoke(); - EXPECT_THAT(m.GetOutput(), - ElementsAreArray(ArrayFloatNear({3, 7, 9}))); - - ResizeBilinearOpModel const_m({TensorType_FLOAT32, {1, 2, 1, 1}}, {3, 1}); - const_m.SetInput({3, 9}); - const_m.Invoke(); - EXPECT_THAT(const_m.GetOutput(), - ElementsAreArray(ArrayFloatNear({3, 7, 9}))); -} - -TEST(ResizeBilinearOpTest, VerticalResize8Bit) { - ResizeBilinearOpModel m({TensorType_UINT8, {1, 2, 1, 1}}); - m.SetInput({3, 9}); - m.SetSize({3, 1}); - m.Invoke(); - EXPECT_THAT(m.GetOutput(), - ElementsAreArray(ArrayFloatNear({3, 7, 9}))); - - ResizeBilinearOpModel const_m({TensorType_UINT8, {1, 2, 1, 1}}, {3, 1}); - const_m.SetInput({3, 9}); - const_m.Invoke(); - EXPECT_THAT(const_m.GetOutput(), - ElementsAreArray(ArrayFloatNear({3, 7, 9}))); -} - -TEST(ResizeBilinearOpTest, TwoDimensionalResize) { - ResizeBilinearOpModel m({TensorType_FLOAT32, {1, 2, 2, 1}}); - m.SetInput({ - 3, 6, // - 9, 12 // - }); - m.SetSize({3, 3}); - m.Invoke(); - EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear({ - 3, 5, 6, // - 7, 9, 10, // - 9, 11, 12, // - }))); - - ResizeBilinearOpModel const_m({TensorType_FLOAT32, {1, 2, 2, 1}}, {3, 3}); - const_m.SetInput({ - 3, 6, // - 9, 12 // - }); - const_m.Invoke(); - EXPECT_THAT(const_m.GetOutput(), ElementsAreArray(ArrayFloatNear({ - 3, 5, 6, // - 7, 9, 10, // - 9, 11, 12, // - }))); -} - -TEST(ResizeBilinearOpTest, TwoDimensionalResize8Bit) { - ResizeBilinearOpModel m({TensorType_UINT8, {1, 2, 2, 1}}); - m.SetInput({ - 3, 6, // - 9, 12 // - }); - m.SetSize({3, 3}); - m.Invoke(); - EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear({ - 3, 5, 6, // - 7, 9, 10, // - 9, 11, 12, // - }))); - - ResizeBilinearOpModel const_m({TensorType_UINT8, {1, 2, 2, 1}}, {3, 3}); - const_m.SetInput({ - 3, 6, // - 9, 12 // - }); - const_m.Invoke(); - EXPECT_THAT(const_m.GetOutput(), ElementsAreArray(ArrayFloatNear({ - 3, 5, 6, // - 7, 9, 10, // - 9, 11, 12, // - }))); -} - -TEST(ResizeBilinearOpTest, TwoDimensionalResizeWithTwoBatches) { - ResizeBilinearOpModel m({TensorType_FLOAT32, {2, 2, 2, 1}}); - m.SetInput({ - 3, 6, // - 9, 12, // - 4, 10, // - 10, 16 // - }); - m.SetSize({3, 3}); - m.Invoke(); - EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear({ - 3, 5, 6, // - 7, 9, 10, // - 9, 11, 12, // - 4, 8, 10, // - 8, 12, 14, // - 10, 14, 16, // - }))); - - ResizeBilinearOpModel const_m({TensorType_FLOAT32, {2, 2, 2, 1}}, {3, 3}); - const_m.SetInput({ - 3, 6, // - 9, 12, // - 4, 10, // - 10, 16 // - }); - const_m.Invoke(); - EXPECT_THAT(const_m.GetOutput(), ElementsAreArray(ArrayFloatNear({ - 3, 5, 6, // - 7, 9, 10, // - 9, 11, 12, // - 4, 8, 10, // - 8, 12, 14, // - 10, 14, 16, // - }))); -} - -TEST(ResizeBilinearOpTest, ThreeDimensionalResize) { - ResizeBilinearOpModel m({TensorType_FLOAT32, {1, 2, 2, 2}}); - m.SetInput({ - 3, 4, 6, 10, // - 9, 10, 12, 16, // - }); - m.SetSize({3, 3}); - m.Invoke(); - EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear({ - 3, 4, 5, 8, 6, 10, // - 7, 8, 9, 12, 10, 14, // - 9, 10, 11, 14, 12, 16, // - }))); - - ResizeBilinearOpModel const_m({TensorType_FLOAT32, {1, 2, 2, 2}}, {3, 3}); - const_m.SetInput({ - 3, 4, 6, 10, // - 9, 10, 12, 16, // - }); - const_m.Invoke(); - EXPECT_THAT(const_m.GetOutput(), ElementsAreArray(ArrayFloatNear({ - 3, 4, 5, 8, 6, 10, // - 7, 8, 9, 12, 10, 14, // - 9, 10, 11, 14, 12, 16, // - }))); -} - -TEST(ResizeBilinearOpTest, TwoDimensionalResizeWithTwoBatches8Bit) { - ResizeBilinearOpModel m({TensorType_UINT8, {2, 2, 2, 1}}); - m.SetInput({ - 3, 6, // - 9, 12, // - 4, 10, // - 12, 16 // - }); - m.SetSize({3, 3}); - m.Invoke(); - EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear({ - 3, 5, 6, // - 7, 9, 10, // - 9, 11, 12, // - 4, 8, 10, // - 9, 12, 14, // - 12, 14, 16, // - }))); - - ResizeBilinearOpModel const_m({TensorType_UINT8, {2, 2, 2, 1}}, {3, 3}); - const_m.SetInput({ - 3, 6, // - 9, 12, // - 4, 10, // - 12, 16 // - }); - const_m.Invoke(); - EXPECT_THAT(const_m.GetOutput(), ElementsAreArray(ArrayFloatNear({ - 3, 5, 6, // - 7, 9, 10, // - 9, 11, 12, // - 4, 8, 10, // - 9, 12, 14, // - 12, 14, 16, // - }))); -} - -TEST(ResizeBilinearOpTest, ThreeDimensionalResize8Bit) { - ResizeBilinearOpModel m({TensorType_UINT8, {1, 2, 2, 2}}); - m.SetInput({ - 3, 4, 6, 10, // - 10, 12, 14, 16, // - }); - m.SetSize({3, 3}); - m.Invoke(); - EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear({ - 3, 4, 5, 8, 6, 10, // - 7, 9, 10, 12, 11, 14, // - 10, 12, 12, 14, 14, 16, // - }))); - - ResizeBilinearOpModel const_m({TensorType_UINT8, {1, 2, 2, 2}}, {3, 3}); - const_m.SetInput({ - 3, 4, 6, 10, // - 10, 12, 14, 16, // - }); - const_m.Invoke(); - EXPECT_THAT(const_m.GetOutput(), ElementsAreArray(ArrayFloatNear({ - 3, 4, 5, 8, 6, 10, // - 7, 9, 10, 12, 11, 14, // - 10, 12, 12, 14, 14, 16, // - }))); -} -} // namespace -} // namespace tflite - -int main(int argc, char** argv) { - ::tflite::LogToStderr(); - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/contrib/lite/kernels/shape.cc b/tensorflow/contrib/lite/kernels/shape.cc deleted file mode 100644 index 66d4c9e5c1a430b621d873012b6ba392afae4157..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/shape.cc +++ /dev/null @@ -1,93 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include "tensorflow/contrib/lite/c/builtin_op_data.h" -#include "tensorflow/contrib/lite/c/c_api_internal.h" -#include "tensorflow/contrib/lite/kernels/internal/tensor.h" -#include "tensorflow/contrib/lite/kernels/kernel_util.h" -#include "tensorflow/contrib/lite/kernels/op_macros.h" - -namespace tflite { -namespace ops { -namespace builtin { -namespace shape { - -constexpr int kInputTensor = 0; -constexpr int kOutputTensor = 0; - -template -void ExtractShape(const TfLiteTensor* input, OutType* output_data) { - for (int i = 0; i < NumDimensions(input); ++i) { - output_data[i] = SizeOfDimension(input, i); - } -} - -TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { - TF_LITE_ENSURE_EQ(context, NumInputs(node), 1); - TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); - - const TfLiteTensor* input = GetInput(context, node, kInputTensor); - TfLiteTensor* output = GetOutput(context, node, kOutputTensor); - - auto* params = reinterpret_cast(node->builtin_data); - switch (params->out_type) { - case kTfLiteInt32: - output->type = kTfLiteInt32; - break; - case kTfLiteInt64: - output->type = kTfLiteInt64; - break; - default: - context->ReportError(context, "Unknown shape output data type: %d", - params->out_type); - return kTfLiteError; - } - - // Shape always produces a 1-dimensional output tensor, where each output - // element is the length of the corresponding input tensor's dimension. - TfLiteIntArray* output_size = TfLiteIntArrayCreate(1); - output_size->data[0] = NumDimensions(input); - return context->ResizeTensor(context, output, output_size); -} - -TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { - const TfLiteTensor* input = GetInput(context, node, kInputTensor); - TfLiteTensor* output = GetOutput(context, node, kOutputTensor); - TFLITE_DCHECK_EQ(NumDimensions(output), 1); - TFLITE_DCHECK_EQ(SizeOfDimension(output, 0), NumDimensions(input)); - - switch (output->type) { - case kTfLiteInt32: - ExtractShape(input, GetTensorData(output)); - break; - case kTfLiteInt64: - ExtractShape(input, GetTensorData(output)); - break; - default: - return kTfLiteError; - } - - return kTfLiteOk; -} - -} // namespace shape - -TfLiteRegistration* Register_SHAPE() { - static TfLiteRegistration r = {nullptr, nullptr, shape::Prepare, shape::Eval}; - return &r; -} - -} // namespace builtin -} // namespace ops -} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/shape_test.cc b/tensorflow/contrib/lite/kernels/shape_test.cc deleted file mode 100644 index 27b48f4e992a8f02d56815bd1bd9074f5b41f400..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/shape_test.cc +++ /dev/null @@ -1,95 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include - -#include -#include "tensorflow/contrib/lite/interpreter.h" -#include "tensorflow/contrib/lite/kernels/register.h" -#include "tensorflow/contrib/lite/kernels/test_util.h" -#include "tensorflow/contrib/lite/model.h" - -namespace tflite { -namespace { - -using ::testing::ElementsAreArray; - -template -class ShapeOpModel : public SingleOpModel { - public: - ShapeOpModel(std::initializer_list input_shape, TensorType input_type, - TensorType output_type) { - input_ = AddInput(input_type); - output_ = AddOutput(output_type); - SetBuiltinOp(BuiltinOperator_SHAPE, BuiltinOptions_ShapeOptions, - CreateShapeOptions(builder_, output_type).Union()); - BuildInterpreter({input_shape}); - } - - TfLiteStatus InvokeWithResult() { return interpreter_->Invoke(); } - - int input() { return input_; } - - int32_t GetOutputSize() { return GetTensorSize(output_); } - std::vector GetOutput() { return ExtractVector(output_); } - std::vector GetOutputShape() { return GetTensorShape(output_); } - - private: - int input_; - int output_; -}; - -TEST(ShapeOpTest, OutTypeInt) { - ShapeOpModel model({1, 3, 1, 3, 5}, TensorType_FLOAT32, - TensorType_INT32); - model.Invoke(); - - EXPECT_THAT(model.GetOutput(), ElementsAreArray({1, 3, 1, 3, 5})); - EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({5})); -} - -TEST(ShapeOpTest, OutTypeInt64) { - ShapeOpModel model({1, 3, 1, 3, 5}, TensorType_FLOAT32, - TensorType_INT64); - model.Invoke(); - - EXPECT_THAT(model.GetOutput(), ElementsAreArray({1, 3, 1, 3, 5})); - EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({5})); -} - -TEST(ShapeOpTest, ScalarTensor) { - ShapeOpModel model({}, TensorType_FLOAT32, TensorType_INT32); - model.Invoke(); - - EXPECT_EQ(model.GetOutputSize(), 0); - EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({0})); -} - -TEST(ShapeOpTest, EmptyTensor) { - ShapeOpModel model({1, 0}, TensorType_FLOAT32, TensorType_INT32); - model.Invoke(); - - EXPECT_THAT(model.GetOutput(), ElementsAreArray({1, 0})); - EXPECT_THAT(model.GetOutputShape(), ElementsAreArray({2})); -} - -} // namespace -} // namespace tflite - -int main(int argc, char** argv) { - ::tflite::LogToStderr(); - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/contrib/lite/kernels/softmax_test.cc b/tensorflow/contrib/lite/kernels/softmax_test.cc deleted file mode 100644 index bd66980226cee0cfd3cf3e81476c60db3d58951c..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/softmax_test.cc +++ /dev/null @@ -1,144 +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. -==============================================================================*/ -// Unit test for TFLite SOFTMAX op. - -#include -#include -#include - -#include -#include -#include "tensorflow/contrib/lite/interpreter.h" -#include "tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h" -#include "tensorflow/contrib/lite/kernels/register.h" -#include "tensorflow/contrib/lite/kernels/test_util.h" -#include "tensorflow/contrib/lite/model.h" - -namespace tflite { -namespace { - -class SoftmaxOpModel : public SingleOpModel { - public: - SoftmaxOpModel(int batches, int size, float beta) - : batches_(batches), input_size_(size), beta_(beta) { - input_ = AddInput(TensorType_FLOAT32); - output_ = AddOutput(TensorType_FLOAT32); - SetBuiltinOp(BuiltinOperator_SOFTMAX, BuiltinOptions_SoftmaxOptions, - CreateSoftmaxOptions(builder_, beta_).Union()); - BuildInterpreter({{batches_, input_size_}}); - } - - void SetInput(std::initializer_list data) { - PopulateTensor(input_, data); - } - - void SetInput(int offset, float* begin, float* end) { - PopulateTensor(input_, offset, begin, end); - } - - std::vector GetOutput() { return ExtractVector(output_); } - - private: - int input_; - int output_; - - int batches_; - int input_size_; - float beta_; -}; - -TEST(SoftmaxOpTest, SimpleTest) { - SoftmaxOpModel m(/*batches=*/2, /*size=*/5, /*beta=*/1.0); - m.SetInput({ - 1.0, 2.0, 3.0, 4.0, 5.0, // b = 0 - -1.0, -2.0, -3.0, -4.0, -5.0, // b = 0 - }); - - m.Invoke(); - - EXPECT_THAT( - m.GetOutput(), - ElementsAreArray(ArrayFloatNear( - {0.011656231, 0.031684921, 0.086128544, 0.234121657, 0.636408647, - 0.636408647, 0.234121657, 0.086128544, 0.031684921, 0.011656231}, - 1e-6))); -} - -TEST(SoftmaxOpTest, CompareWithTFminiBetaEq1) { - const int batch_size = 2; - const int input_size = 5; - const float beta = 1.0; - static float input_buffer[] = { - 1.0, 2.0, 3.0, 4.0, 5.0, // b = 0 - -1.0, -2.0, -3.0, -4.0, -5.0, // b = 1 - }; - - SoftmaxOpModel m(batch_size, input_size, beta); - - m.SetInput(0, input_buffer, input_buffer + input_size * batch_size); - - m.Invoke(); - - std::unique_ptr output_buffer(new float[input_size * batch_size]); - auto input_shape = RuntimeShape({batch_size, 1, 1, input_size}); - SoftmaxParams params; - params.beta = beta; - tflite::reference_ops::Softmax(params, input_shape, input_buffer, input_shape, - output_buffer.get()); - - std::vector expected; - expected.insert(expected.end(), output_buffer.get(), - output_buffer.get() + input_size * batch_size); - - EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear(expected, 1e-6))); -} - -TEST(SoftmaxOpTest, CompareWithTFminiBetaNotEq1) { - const int batch_size = 2; - const int input_size = 5; - const float beta = 0.5; - static float input_buffer[] = { - 1.0, 2.0, 3.0, 4.0, 5.0, // b = 0 - -1.0, -2.0, -3.0, -4.0, -5.0, // b = 1 - }; - - SoftmaxOpModel m(batch_size, input_size, beta); - - m.SetInput(0, input_buffer, input_buffer + input_size * batch_size); - - m.Invoke(); - - std::unique_ptr output_buffer(new float[input_size * batch_size]); - auto input_shape = RuntimeShape({batch_size, 1, 1, input_size}); - SoftmaxParams params; - params.beta = beta; - tflite::reference_ops::Softmax(params, input_shape, input_buffer, input_shape, - output_buffer.get()); - - std::vector expected; - expected.insert(expected.end(), output_buffer.get(), - output_buffer.get() + input_size * batch_size); - - EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear(expected, 1e-6))); -} - -} // namespace -} // namespace tflite - -int main(int argc, char** argv) { - ::tflite::LogToStderr(); - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/contrib/lite/kernels/sparse_output_fully_connected.cc b/tensorflow/contrib/lite/kernels/sparse_output_fully_connected.cc deleted file mode 100644 index 66daf5e84a05670189aa57c9fc195af5c7c83294..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/sparse_output_fully_connected.cc +++ /dev/null @@ -1,243 +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. -==============================================================================*/ -// SparseOutputFullyConnected is a fully connected layer that uses a single -// row in the weights and bias via a lookup. -#include "tensorflow/contrib/lite/context.h" -#include "tensorflow/contrib/lite/kernels/internal/tensor.h" -#include "tensorflow/contrib/lite/kernels/internal/tensor_utils.h" -#include "tensorflow/contrib/lite/kernels/kernel_util.h" - -namespace tflite { -namespace ops { -namespace custom { -namespace sparse_output_fully_connected { - -// Input tensors of size {n_batch, n_input} -constexpr int kInputTensor = 0; -// Auxiliary input tensor of size { 1 } -constexpr int kInputLookupTensor = 1; - -// Weights tensor of size { n_embeddings , n_input } -constexpr int kWeightsTensor = 2; -// Bias tensor of size { n_embeddings } -constexpr int kBiasTensor = 3; - -// Output tensor. -constexpr int kOutputTensor = 0; - -// Temporary tensors. -enum TemporaryTensor { - kInputQuantized = 0, - kScalingFactors = 1, - kNumTemporaryTensors = 2 -}; - -// Struct to hold op data. -struct OpData { - int scratch_tensor_index; -}; - -void* Init(TfLiteContext* context, const char* buffer, size_t length) { - auto* data = new OpData; - context->AddTensors(context, /*tensors_to_add=*/kNumTemporaryTensors, - &data->scratch_tensor_index); - return data; -} - -void Free(TfLiteContext* context, void* buffer) { - delete reinterpret_cast(buffer); -} - -TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { - OpData* op_data = reinterpret_cast(node->user_data); - - TF_LITE_ENSURE_EQ(context, node->inputs->size, 4); - TF_LITE_ENSURE_EQ(context, node->outputs->size, 1); - - const TfLiteTensor* input = GetInput(context, node, kInputTensor); - TF_LITE_ENSURE_EQ(context, input->type, kTfLiteFloat32); - TF_LITE_ENSURE_EQ(context, NumDimensions(input), 2); - const int n_batch = SizeOfDimension(input, 0); - const int n_input = SizeOfDimension(input, 1); - - const TfLiteTensor* lookup = GetInput(context, node, kInputLookupTensor); - TF_LITE_ENSURE_EQ(context, lookup->type, kTfLiteInt32); - TF_LITE_ENSURE_EQ(context, NumDimensions(lookup), 1); - // Only support single lookup. - TF_LITE_ENSURE_EQ(context, SizeOfDimension(lookup, 0), 1); - - const TfLiteTensor* weights = GetInput(context, node, kWeightsTensor); - TF_LITE_ENSURE_EQ(context, NumDimensions(weights), 2); - TF_LITE_ENSURE_EQ(context, SizeOfDimension(weights, 1), n_input); - - const TfLiteTensor* bias = GetInput(context, node, kBiasTensor); - TF_LITE_ENSURE_EQ(context, NumElements(bias), SizeOfDimension(weights, 0)); - - const bool is_hybrid_op = - (weights->type == kTfLiteUInt8 && input->type == kTfLiteFloat32); - - // Resize output. - TfLiteTensor* output = GetOutput(context, node, kOutputTensor); - TfLiteIntArray* output_size_array = TfLiteIntArrayCreate(1); - output_size_array->data[0] = 1; - TF_LITE_ENSURE_OK(context, - context->ResizeTensor(context, output, output_size_array)); - - if (is_hybrid_op) { - TfLiteIntArrayFree(node->temporaries); - node->temporaries = TfLiteIntArrayCreate(kNumTemporaryTensors); - - // Allocate temporary tensors to store quantized values of input. - node->temporaries->data[kInputQuantized] = op_data->scratch_tensor_index; - TfLiteTensor* input_quantized = - GetTemporary(context, node, /*index=*/kInputQuantized); - input_quantized->type = kTfLiteUInt8; - input_quantized->allocation_type = kTfLiteArenaRw; - if (!TfLiteIntArrayEqual(input_quantized->dims, input->dims)) { - TfLiteIntArray* input_quantized_size = TfLiteIntArrayCopy(input->dims); - TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, input_quantized, - input_quantized_size)); - } - - // Tell interpreter to allocate temporary tensors to store scaling factors. - node->temporaries->data[kScalingFactors] = - op_data->scratch_tensor_index + kScalingFactors; - TfLiteTensor* scaling_factors = - GetTemporary(context, node, /*index=*/kScalingFactors); - scaling_factors->type = kTfLiteFloat32; - scaling_factors->allocation_type = kTfLiteArenaRw; - int scaling_dims[1] = {n_batch}; - if (!TfLiteIntArrayEqualsArray(scaling_factors->dims, 1, scaling_dims)) { - TfLiteIntArray* scaling_factors_size = TfLiteIntArrayCreate(1); - scaling_factors_size->data[0] = n_batch; - TF_LITE_ENSURE_OK(context, context->ResizeTensor(context, scaling_factors, - scaling_factors_size)); - } - } - return kTfLiteOk; -} - -TfLiteStatus EvalFloat(const TfLiteTensor* input, const TfLiteTensor* lookup, - const TfLiteTensor* weights, const TfLiteTensor* bias, - TfLiteTensor* output) { - const int n_batch = SizeOfDimension(input, 0); - const int n_input = SizeOfDimension(input, 1); - - const float* input_ptr_batch = input->data.f; - - // Initialize pointer to right row according to lookup value. - int32 lookup_index = lookup->data.i32[0]; - const float* weights_ptr = weights->data.f + lookup_index * n_input; - - // Initialize output to bias. - if (bias) { - float* bias_ptr = bias->data.f + lookup_index; - tensor_utils::VectorBatchVectorAssign(bias_ptr, 1, n_batch, output->data.f); - } else { - tensor_utils::ZeroVector(output->data.f, n_batch * 1); - } - - tensor_utils::MatrixBatchVectorMultiplyAccumulate( - weights_ptr, /*m_rows=*/1, n_input, input_ptr_batch, n_batch, - output->data.f, /*result_stride=*/1); - - return kTfLiteOk; -} - -TfLiteStatus EvalHybrid(const TfLiteTensor* input, const TfLiteTensor* lookup, - const TfLiteTensor* weights, const TfLiteTensor* bias, - TfLiteTensor* scaling_factors, - TfLiteTensor* input_quantized, TfLiteTensor* output) { - const int n_batch = SizeOfDimension(input, 0); - const int n_input = SizeOfDimension(input, 1); - - const float* input_ptr_batch = input->data.f; - // Initialize the pointer to storage for quantized values and - // scaling factors. - int8_t* quantized_input_ptr_batch = - reinterpret_cast(input_quantized->data.uint8); - float* scaling_factors_ptr = scaling_factors->data.f; - - // Initialize pointer to right row according to lookup value. - int32 lookup_index = lookup->data.i32[0]; - int8_t* weights_ptr = - reinterpret_cast(weights->data.uint8) + lookup_index * n_input; - - // Initialize output to bias. - if (bias) { - float* bias_ptr = bias->data.f + lookup_index; - tensor_utils::VectorBatchVectorAssign(bias_ptr, 1, n_batch, output->data.f); - } else { - tensor_utils::ZeroVector(output->data.f, n_batch * 1); - } - - if (!tensor_utils::IsZeroVector(input_ptr_batch, n_batch * n_input)) { - // Quantize input from float to int8. - float unused_min, unused_max; - for (int b = 0; b < n_batch; ++b) { - const int offset = b * n_input; - tensor_utils::SymmetricQuantizeFloats( - input_ptr_batch + offset, n_input, quantized_input_ptr_batch + offset, - &unused_min, &unused_max, &scaling_factors_ptr[b]); - scaling_factors_ptr[b] *= weights->params.scale; - } - - tensor_utils::MatrixBatchVectorMultiplyAccumulate( - weights_ptr, /*m_rows=*/1, n_input, quantized_input_ptr_batch, - scaling_factors_ptr, n_batch, output->data.f, /*result_stride=*/1); - } - return kTfLiteOk; -} - -TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { - const TfLiteTensor* input = GetInput(context, node, kInputTensor); - const TfLiteTensor* lookup = GetInput(context, node, kInputLookupTensor); - const TfLiteTensor* weights = GetInput(context, node, kWeightsTensor); - const TfLiteTensor* bias = GetInput(context, node, kBiasTensor); - TfLiteTensor* output = GetOutput(context, node, kOutputTensor); - - switch (weights->type) { - case kTfLiteFloat32: { - return EvalFloat(input, lookup, weights, bias, output); - } - case kTfLiteUInt8: { - TfLiteTensor* input_quantized = - GetTemporary(context, node, /*index=*/kInputQuantized); - TfLiteTensor* scaling_factors = - GetTemporary(context, node, /*index=*/kScalingFactors); - return EvalHybrid(input, lookup, weights, bias, scaling_factors, - input_quantized, output); - } - default: - context->ReportError(context, "Type %d is not currently supported.", - weights->type); - return kTfLiteError; - } - return kTfLiteOk; -} - -} // namespace sparse_output_fully_connected - -TfLiteRegistration* Register_SPARSE_OUTPUT_FULLY_CONNECTED() { - static TfLiteRegistration r = {sparse_output_fully_connected::Init, - sparse_output_fully_connected::Free, - sparse_output_fully_connected::Prepare, - sparse_output_fully_connected::Eval}; - return &r; -} - -} // namespace custom -} // namespace ops -} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/sparse_output_fully_connected_test.cc b/tensorflow/contrib/lite/kernels/sparse_output_fully_connected_test.cc deleted file mode 100644 index 365986a5c177ee58604138f279ca1186bacc742e..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/sparse_output_fully_connected_test.cc +++ /dev/null @@ -1,158 +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. -==============================================================================*/ -// Unit test for TFLite sparse output fully connected op. -#include -#include -#include - -#include -#include "flatbuffers/flexbuffers.h" // TF:flatbuffers -#include "tensorflow/contrib/lite/kernels/register.h" -#include "tensorflow/contrib/lite/kernels/test_util.h" - -namespace tflite { - -namespace ops { -namespace custom { - -TfLiteRegistration* Register_SPARSE_OUTPUT_FULLY_CONNECTED(); - -namespace { - -using ::testing::ElementsAreArray; - -class BaseSparseOutputFullyConnectedOpModel : public SingleOpModel { - public: - BaseSparseOutputFullyConnectedOpModel(const TensorData& input, - const TensorData& weights, - const TensorData& output = { - TensorType_FLOAT32}) { - input_ = AddInput(input); - lookup_ = AddInput({TensorType_INT32, {1}}); - weights_ = AddInput(weights); - int bias_size = GetShape(weights_)[0]; - bias_ = AddInput({TensorType_FLOAT32, {bias_size}}); - output_ = AddOutput(output); - - // Create empty (required) options map. - flexbuffers::Builder fbb; - fbb.Map([&]() {}); - fbb.Finish(); - - SetCustomOp("SPARSE_OUTPUT_FULLY_CONNECTED", fbb.GetBuffer(), - Register_SPARSE_OUTPUT_FULLY_CONNECTED); - BuildInterpreter({GetShape(input_), GetShape(lookup_), GetShape(weights_), - GetShape(bias_)}); - } - - void SetInput(const std::vector& data) { - PopulateTensor(input_, data); - } - - void SetLookup(const std::vector& f) { PopulateTensor(lookup_, f); } - - void SetBias(const std::vector& f) { PopulateTensor(bias_, f); } - - std::vector GetOutput() { return ExtractVector(output_); } - - protected: - int input_; - int lookup_; - int weights_; - int bias_; - int output_; -}; - -class FloatSparseOutputFullyConnectedOpModel - : public BaseSparseOutputFullyConnectedOpModel { - public: - using BaseSparseOutputFullyConnectedOpModel:: - BaseSparseOutputFullyConnectedOpModel; - - void SetWeights(const std::vector& f) { PopulateTensor(weights_, f); } -}; - -class HybridSparseOutputFullyConnectedOpModel - : public BaseSparseOutputFullyConnectedOpModel { - public: - using BaseSparseOutputFullyConnectedOpModel:: - BaseSparseOutputFullyConnectedOpModel; - - void SetWeights(const std::vector& f) { - SymmetricQuantizeAndPopulate(weights_, f); - } -}; - -TEST(SparseOutputFullyConnectedOpTest, SimpleTestFloat) { - FloatSparseOutputFullyConnectedOpModel m({TensorType_FLOAT32, {1, 5}}, - {TensorType_FLOAT32, {3, 5}}, - {TensorType_FLOAT32, {}}); - - m.SetInput({-1.0, 0.0, 1.0, 2.0, 3.0}); - - m.SetLookup({2}); - - m.SetWeights({ - -1.0, 0.0, 1.0, 2.0, 3.0, // - 0.0, 1.0, 2.0, 3.0, 4.0, // - 1.0, 2.0, 3.0, 4.0, 5.0, // - }); - - m.SetBias({1.0, 2.0, 3.0}); - - m.Invoke(); - - EXPECT_THAT(m.GetOutput(), ElementsAreArray({28})); -} - -TEST(SparseOutputFullyConnectedOpTest, SimpleTestHybrid) { - HybridSparseOutputFullyConnectedOpModel m({TensorType_FLOAT32, {1, 5}}, - {TensorType_UINT8, {3, 5}}, - {TensorType_FLOAT32, {}}); - - m.SetInput({-1.0, 0.0, 1.0, 2.0, 3.0}); - - m.SetLookup({2}); - - m.SetWeights({ - -1.0, 0.0, 1.0, 2.0, 3.0, // - 0.0, 1.0, 2.0, 3.0, 4.0, // - 1.0, 2.0, 3.0, 4.0, 5.0, // - }); - - m.SetBias({1.0, 2.0, 3.0}); - - m.Invoke(); - - // We get 28.0552 instead of 28. - // - // Input -> -42, 0, 42, 85, 127 with scale factor of 127/3. - // Looked up weights -> 25, 51, 76, 102, 127 with scale factor of 127/5. - // - // (-42 * 25 + 0 * 51 + 42 * 76 + 85 * 102 + 127 * 127) * (3*5/127^2) + 3.0 - // gives us the expected result. - EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear({28}, 0.0553))); -} - -} // namespace -} // namespace custom -} // namespace ops -} // namespace tflite - -int main(int argc, char** argv) { - ::tflite::LogToStderr(); - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/contrib/lite/kernels/split.cc b/tensorflow/contrib/lite/kernels/split.cc deleted file mode 100644 index dab887bf9ccac0ff43cb5f7bd11033657aaf1fd2..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/split.cc +++ /dev/null @@ -1,165 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include -#include -#include "tensorflow/contrib/lite/c/builtin_op_data.h" -#include "tensorflow/contrib/lite/c/c_api_internal.h" -#include "tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h" -#include "tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h" -#include "tensorflow/contrib/lite/kernels/internal/tensor.h" -#include "tensorflow/contrib/lite/kernels/kernel_util.h" -#include "tensorflow/contrib/lite/kernels/op_macros.h" - -namespace tflite { -namespace ops { -namespace builtin { -namespace split { - -struct OpContext { - OpContext(TfLiteContext* context, TfLiteNode* node) { - params = reinterpret_cast(node->builtin_data); - axis = GetInput(context, node, 0); - input = GetInput(context, node, 1); - } - TfLiteSplitParams* params; - const TfLiteTensor* axis; - const TfLiteTensor* input; -}; - -TfLiteStatus UseDynamicOutputTensors(TfLiteContext* context, TfLiteNode* node) { - for (int i = 0; i < NumOutputs(node); ++i) { - SetTensorToDynamic(GetOutput(context, node, i)); - } - return kTfLiteOk; -} - -TfLiteStatus ResizeOutputTensors(TfLiteContext* context, TfLiteNode* node, - const TfLiteTensor* axis, - const TfLiteTensor* input, int num_splits) { - int axis_value = GetTensorData(axis)[0]; - if (axis_value < 0) { - axis_value += NumDimensions(input); - } - - const int input_size = SizeOfDimension(input, axis_value); - TF_LITE_ENSURE_MSG(context, input_size % num_splits == 0, - "Not an even split"); - const int slice_size = input_size / num_splits; - - for (int i = 0; i < NumOutputs(node); ++i) { - TfLiteIntArray* output_dims = TfLiteIntArrayCopy(input->dims); - output_dims->data[axis_value] = slice_size; - TfLiteTensor* output = GetOutput(context, node, i); - TF_LITE_ENSURE_STATUS(context->ResizeTensor(context, output, output_dims)); - } - - return kTfLiteOk; -} - -TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { - TF_LITE_ENSURE_EQ(context, NumInputs(node), 2); - - OpContext op_context(context, node); - - TF_LITE_ENSURE_EQ(context, NumOutputs(node), op_context.params->num_splits); - - auto input_type = op_context.input->type; - TF_LITE_ENSURE(context, input_type == kTfLiteFloat32 || - input_type == kTfLiteUInt8 || - input_type == kTfLiteInt16); - for (int i = 0; i < NumOutputs(node); ++i) { - GetOutput(context, node, i)->type = input_type; - } - - // If we know the contents of the 'axis' tensor, resize all outputs. - // Otherwise, wait until Eval(). - if (IsConstantTensor(op_context.axis)) { - return ResizeOutputTensors(context, node, op_context.axis, op_context.input, - op_context.params->num_splits); - } else { - return UseDynamicOutputTensors(context, node); - } -} - -TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { - OpContext op_context(context, node); - - // When the 'axis' tensor is non-const we can't resize output tensors in - // Prepare(), and we have to do it now. - if (!IsConstantTensor(op_context.axis)) { - TF_LITE_ENSURE_OK( - context, - ResizeOutputTensors(context, node, op_context.axis, op_context.input, - op_context.params->num_splits)); - } - - int axis_value = GetTensorData(op_context.axis)[0]; - if (axis_value < 0) { - axis_value += NumDimensions(op_context.input); - } - - // TODO(ahentz): Our usage of VectorOfTensors could be optimized by - // calculating it in Prepare, unless we defer shape calculation. - // TODO(ahentz): We can improve the optimized_ops version to handle other - // cases too. -#define TF_LITE_SPLIT(scalar) \ - VectorOfTensors all_outputs(*context, *node->outputs); \ - tflite::SplitParams op_params; \ - op_params.num_split = NumOutputs(node); \ - op_params.axis = axis_value; \ - if (axis_value == 0) { \ - optimized_ops::Split(op_params, GetTensorShape(op_context.input), \ - GetTensorData(op_context.input), \ - all_outputs.shapes(), all_outputs.data()); \ - } else { \ - reference_ops::Split(op_params, GetTensorShape(op_context.input), \ - GetTensorData(op_context.input), \ - all_outputs.shapes(), all_outputs.data()); \ - } - switch (op_context.input->type) { - case kTfLiteFloat32: { - TF_LITE_SPLIT(float); - break; - } - case kTfLiteUInt8: { - TF_LITE_SPLIT(uint8_t); - break; - } - case kTfLiteInt16: { - TF_LITE_SPLIT(int16_t); - break; - } - default: - context->ReportError( - context, - "Only float32, uint8 and int16 are currently supported, got %d.", - op_context.input->type); - return kTfLiteError; - } -#undef TF_LITE_SPLIT - - return kTfLiteOk; -} - -} // namespace split - -TfLiteRegistration* Register_SPLIT() { - static TfLiteRegistration r = {nullptr, nullptr, split::Prepare, split::Eval}; - return &r; -} - -} // namespace builtin -} // namespace ops -} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/split_test.cc b/tensorflow/contrib/lite/kernels/split_test.cc deleted file mode 100644 index 61a0759c6475795c06a9b55d3586d2b818f298b2..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/split_test.cc +++ /dev/null @@ -1,147 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include -#include "tensorflow/contrib/lite/interpreter.h" -#include "tensorflow/contrib/lite/kernels/register.h" -#include "tensorflow/contrib/lite/kernels/test_util.h" -#include "tensorflow/contrib/lite/model.h" - -namespace tflite { -namespace { - -using ::testing::ElementsAreArray; - -constexpr int kAxisIsATensor = -1000; - -class SplitOpModel : public SingleOpModel { - public: - SplitOpModel(const TensorData& input, int num_splits, - int axis = kAxisIsATensor) { - if (axis == kAxisIsATensor) { - axis_ = AddInput({TensorType_INT32, {1}}); - } else { - axis_ = AddConstInput(TensorType_INT32, {axis}, {1}); - } - input_ = AddInput(input); - for (int i = 0; i < num_splits; ++i) { - outputs_.push_back(AddOutput(input.type)); - } - SetBuiltinOp(BuiltinOperator_SPLIT, BuiltinOptions_SplitOptions, - CreateSplitOptions(builder_, num_splits).Union()); - if (axis == kAxisIsATensor) { - BuildInterpreter({GetShape(axis_), GetShape(input_)}); - } else { - BuildInterpreter({{}, GetShape(input_)}); - } - } - - void SetInput(std::initializer_list data) { - PopulateTensor(input_, data); - } - void SetAxis(int axis) { PopulateTensor(axis_, {axis}); } - - std::vector GetOutput(int i) { - return ExtractVector(outputs_[i]); - } - std::vector GetOutputShape(int i) { return GetTensorShape(outputs_[i]); } - - private: - int input_; - int axis_; - std::vector outputs_; -}; - -using TensorValues = std::initializer_list; - -void Check(int axis, int num_splits, std::initializer_list input_shape, - std::initializer_list output_shape, - const TensorValues& input_data, - const std::vector& output_data) { - auto debug = [&](int i) { - std::stringstream ss; - ss << "for output tensor " << i << " axis=" << axis - << " and num_splits=" << num_splits; - return ss.str(); - }; - SplitOpModel m({TensorType_FLOAT32, input_shape}, num_splits); - m.SetInput(input_data); - m.SetAxis(axis); - m.Invoke(); - for (int i = 0; i < num_splits; ++i) { - EXPECT_THAT(m.GetOutput(i), ElementsAreArray(output_data[i])) << debug(i); - EXPECT_THAT(m.GetOutputShape(i), ElementsAreArray(output_shape)) - << debug(i); - } - - SplitOpModel const_m({TensorType_FLOAT32, input_shape}, num_splits, axis); - const_m.SetInput(input_data); - const_m.Invoke(); - for (int i = 0; i < num_splits; ++i) { - EXPECT_THAT(const_m.GetOutput(i), ElementsAreArray(output_data[i])) - << debug(i); - EXPECT_THAT(const_m.GetOutputShape(i), ElementsAreArray(output_shape)) - << debug(i); - } -} - -TEST(SplitOpTest, FourDimensional) { - Check(/*axis=*/0, /*num_splits=*/2, {2, 2, 2, 2}, {1, 2, 2, 2}, - {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}, - { - {1, 2, 3, 4, 5, 6, 7, 8}, - {9, 10, 11, 12, 13, 14, 15, 16}, - }); - Check(/*axis=*/1, /*num_splits=*/2, {2, 2, 2, 2}, {2, 1, 2, 2}, - {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}, - { - {1, 2, 3, 4, 9, 10, 11, 12}, - {5, 6, 7, 8, 13, 14, 15, 16}, - }); - Check(/*axis=*/2, /*num_splits=*/2, {2, 2, 2, 2}, {2, 2, 1, 2}, - {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}, - { - {1, 2, 5, 6, 9, 10, 13, 14}, - {3, 4, 7, 8, 11, 12, 15, 16}, - }); - Check(/*axis=*/3, /*num_splits=*/2, {2, 2, 2, 2}, {2, 2, 2, 1}, - {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}, - { - {1, 3, 5, 7, 9, 11, 13, 15}, - {2, 4, 6, 8, 10, 12, 14, 16}, - }); -} - -TEST(SplitOpTest, OneDimensional) { - Check(/*axis=*/0, /*num_splits=*/8, {8}, {1}, {1, 2, 3, 4, 5, 6, 7, 8}, - {{1}, {2}, {3}, {4}, {5}, {6}, {7}, {8}}); -} - -TEST(SplitOpTest, NegativeAxis) { - Check(/*axis=*/-4, /*num_splits=*/2, {2, 2, 2, 2}, {1, 2, 2, 2}, - {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}, - { - {1, 2, 3, 4, 5, 6, 7, 8}, - {9, 10, 11, 12, 13, 14, 15, 16}, - }); -} - -} // namespace -} // namespace tflite - -int main(int argc, char** argv) { - ::tflite::LogToStderr(); - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/contrib/lite/kernels/test_util.cc b/tensorflow/contrib/lite/kernels/test_util.cc deleted file mode 100644 index 05a7c23ba10ef717ee3debf0a6316885d4612746..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/test_util.cc +++ /dev/null @@ -1,153 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include "tensorflow/contrib/lite/kernels/test_util.h" - -#include "tensorflow/contrib/lite/version.h" -#include "tensorflow/core/platform/logging.h" - -namespace tflite { - -using ::testing::FloatNear; -using ::testing::Matcher; - -std::vector> ArrayFloatNear(const std::vector& values, - float max_abs_error) { - std::vector> matchers; - matchers.reserve(values.size()); - for (const float& v : values) { - matchers.emplace_back(FloatNear(v, max_abs_error)); - } - return matchers; -} - -int SingleOpModel::AddInput(const TensorData& t, bool is_variable) { - int id = AddTensor(t, {}, is_variable); - inputs_.push_back(id); - return id; -} - -int SingleOpModel::AddNullInput() { - int id = kOptionalTensor; - inputs_.push_back(id); - return id; -} - -int SingleOpModel::AddOutput(const TensorData& t) { - int id = AddTensor(t, {}); - outputs_.push_back(id); - return id; -} - -void SingleOpModel::SetBuiltinOp(BuiltinOperator type, - BuiltinOptions builtin_options_type, - flatbuffers::Offset builtin_options) { - opcodes_.push_back(CreateOperatorCode(builder_, type, 0)); - operators_.push_back(CreateOperator( - builder_, /*opcode_index=*/0, builder_.CreateVector(inputs_), - builder_.CreateVector(outputs_), builtin_options_type, - builtin_options, - /*custom_options=*/0, CustomOptionsFormat_FLEXBUFFERS)); -} - -void SingleOpModel::SetCustomOp( - const string& name, const std::vector& custom_option, - const std::function& registration) { - custom_registrations_[name] = registration; - opcodes_.push_back( - CreateOperatorCodeDirect(builder_, BuiltinOperator_CUSTOM, name.data())); - operators_.push_back(CreateOperator( - builder_, /*opcode_index=*/0, builder_.CreateVector(inputs_), - builder_.CreateVector(outputs_), BuiltinOptions_NONE, 0, - builder_.CreateVector(custom_option), - CustomOptionsFormat_FLEXBUFFERS)); -} - -void SingleOpModel::BuildInterpreter(std::vector> input_shapes, - bool allow_fp32_relax_to_fp16) { - auto opcodes = builder_.CreateVector(opcodes_); - auto operators = builder_.CreateVector(operators_); - auto tensors = builder_.CreateVector(tensors_); - auto inputs = builder_.CreateVector(inputs_); - auto outputs = builder_.CreateVector(outputs_); - // Create a single subgraph - std::vector> subgraphs; - auto subgraph = CreateSubGraph(builder_, tensors, inputs, outputs, operators); - subgraphs.push_back(subgraph); - auto subgraphs_flatbuffer = builder_.CreateVector(subgraphs); - - auto buffers = builder_.CreateVector(buffers_); - auto description = builder_.CreateString("programmatic model"); - builder_.Finish(CreateModel(builder_, TFLITE_SCHEMA_VERSION, opcodes, - subgraphs_flatbuffer, description, buffers)); - - auto* model = GetModel(builder_.GetBufferPointer()); - - if (!resolver_) { - auto resolver = new ops::builtin::BuiltinOpResolver(); - for (const auto& reg : custom_registrations_) { - resolver->AddCustom(reg.first.data(), reg.second()); - } - resolver_ = std::unique_ptr(resolver); - } - CHECK(InterpreterBuilder(model, *resolver_)(&interpreter_) == kTfLiteOk); - - CHECK(interpreter_ != nullptr); - - int i = 0; - for (const auto& shape : input_shapes) { - int input_idx = interpreter_->inputs()[i++]; - if (input_idx == kOptionalTensor) continue; - if (shape.empty()) continue; - CHECK(interpreter_->ResizeInputTensor(input_idx, shape) == kTfLiteOk); - } - - interpreter_->SetAllowFp16PrecisionForFp32(allow_fp32_relax_to_fp16); - - // Modify delegate with function. - if (apply_delegate_fn_) { - apply_delegate_fn_(interpreter_.get()); - } - - CHECK(interpreter_->AllocateTensors() == kTfLiteOk) - << "Cannot allocate tensors"; - interpreter_->ResetVariableTensors(); -} - -void SingleOpModel::Invoke() { CHECK(interpreter_->Invoke() == kTfLiteOk); } - -int32_t SingleOpModel::GetTensorSize(int index) const { - TfLiteTensor* t = interpreter_->tensor(index); - CHECK(t); - int total_size = 1; - for (int i = 0; i < t->dims->size; ++i) { - total_size *= t->dims->data[i]; - } - return total_size; -} - -template <> -std::vector SingleOpModel::ExtractVector(int index) { - TfLiteTensor* tensor_ptr = interpreter_->tensor(index); - CHECK(tensor_ptr != nullptr); - const int num_strings = GetStringCount(tensor_ptr); - std::vector result; - result.reserve(num_strings); - for (int i = 0; i < num_strings; ++i) { - const auto str = GetString(tensor_ptr, i); - result.emplace_back(str.str, str.len); - } - return result; -} -} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/test_util.h b/tensorflow/contrib/lite/kernels/test_util.h deleted file mode 100644 index 670120219ff4e17e72276a34c98ef4ceb142c361..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/test_util.h +++ /dev/null @@ -1,386 +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_CONTRIB_LITE_KERNELS_TEST_UTIL_H_ -#define TENSORFLOW_CONTRIB_LITE_KERNELS_TEST_UTIL_H_ - -#include - -#include -#include - -#include "tensorflow/contrib/lite/interpreter.h" -#include "tensorflow/contrib/lite/kernels/internal/tensor_utils.h" -#include "tensorflow/contrib/lite/kernels/register.h" -#include "tensorflow/contrib/lite/model.h" -#include "tensorflow/contrib/lite/string_util.h" -#include "tensorflow/contrib/lite/testing/util.h" -#include "tensorflow/core/platform/logging.h" - -namespace tflite { - -// A gmock matcher that check that elements of a float vector match to a given -// tolerance. -std::vector<::testing::Matcher> ArrayFloatNear( - const std::vector& values, float max_abs_error = 1e-5); - -template -inline std::vector Quantize(const std::vector& data, float scale, - int32_t zero_point) { - std::vector q; - for (float f : data) { - q.push_back(static_cast(std::max( - std::numeric_limits::min(), - std::min(std::numeric_limits::max(), - std::round(zero_point + (f / scale)))))); - } - return q; -} - -template -inline std::vector Dequantize(const std::vector& data, float scale, - int32_t zero_point) { - std::vector f; - for (T q : data) { - f.push_back(scale * (q - zero_point)); - } - return f; -} - -// A test model that contains a single operator. All operator inputs and -// output are external to the model, so the tests can directly access them. -// Typical usage: -// SingleOpModel m; -// int a = m.AddInput({TensorType_FLOAT32, a_shape}); -// int b = m.AddInput({TensorType_FLOAT32, b_shape}); -// int c = m.AddOutput({TensorType_FLOAT32, {}}); -// m.SetBuiltinOp(...); -// m.BuildInterpreter({GetShape(a), GetShape(b)}); -// m.PopulateTensor(a, {...}); -// m.PopulateTensor(b, {...}); -// m.Invoke(); -// EXPECT_THAT(m.ExtractVector(c), ArrayFloatNear({...})); -// - -// A helper struct to construct test tensors. This is particularly useful for -// quantized tensor which must have their scale and zero_point defined before -// the actual data is known. This mimics what happens in practice: quantization -// parameters are calculated during training. -struct TensorData { - TensorType type; - std::vector shape; - float min; - float max; - float scale; - int32_t zero_point; -}; - -class SingleOpResolver : public OpResolver { - public: - SingleOpResolver(const BuiltinOperator op, TfLiteRegistration* registration) - : op_(op), registration_(*registration) { - registration_.builtin_code = static_cast(op); - registration_.version = 1; - } - const TfLiteRegistration* FindOp(BuiltinOperator op, - int version) const override { - if (op == op_) { - return ®istration_; - } - return nullptr; - } - const TfLiteRegistration* FindOp(const char* op, int version) const override { - return nullptr; - } - - private: - const BuiltinOperator op_; - TfLiteRegistration registration_; -}; - -class SingleOpModel { - public: - SingleOpModel() {} - ~SingleOpModel() {} - - // Set a function callback that is run right after graph is prepared - // that allows applying external delegates. This is useful for testing - // other runtimes like NN API or GPU. - void SetApplyDelegate(std::function apply_delegate_fn) { - apply_delegate_fn_ = apply_delegate_fn; - } - - // Copying or assignment is disallowed to simplify ownership semantics. - SingleOpModel(const SingleOpModel&) = delete; - SingleOpModel& operator=(const SingleOpModel&) = delete; - - // Add a TensorType input tensor and return its index. - int AddInput(TensorType type, bool is_variable = false) { - return AddInput(TensorData{type}, is_variable); - } - int AddInput(const TensorData& t, bool is_variable = false); - - // Templated version of AddConstInput(). - template - int AddConstInput(TensorType type, std::initializer_list data, - std::initializer_list shape) { - int id = AddTensor(TensorData{type, shape}, data); - inputs_.push_back(id); - return id; - } - - // Add a null input tensor (optional input) and return kOptionalTensor. - int AddNullInput(); - - // Add a TensorType output tensor and return its index. - int AddOutput(TensorType type) { return AddOutput(TensorData{type}); } - int AddOutput(const TensorData& t); - - template - void QuantizeAndPopulate(int index, const std::vector& data) { - TfLiteTensor* t = interpreter_->tensor(index); - auto q = Quantize(data, t->params.scale, t->params.zero_point); - PopulateTensor(index, 0, q.data(), q.data() + q.size()); - } - - void SymmetricQuantizeAndPopulate(int index, const std::vector& data) { - TfLiteTensor* t = interpreter_->tensor(index); - const int length = data.size(); - std::vector q(length); - float min, max, scaling_factor; - tensor_utils::SymmetricQuantizeFloats(data.data(), length, q.data(), &min, - &max, &scaling_factor); - // Update quantization params. - t->params.scale = scaling_factor; - t->params.zero_point = 0; - PopulateTensor(index, /*offset=*/0, reinterpret_cast(q.data()), - reinterpret_cast(q.data() + q.size())); - } - - const std::vector& GetShape(int id) { return tensor_data_.at(id).shape; } - - float GetScale(int id) { return tensor_data_.at(id).scale; } - int32_t GetZeroPoint(int id) { return tensor_data_.at(id).zero_point; } - - // Define the operator in this model. - void SetBuiltinOp(BuiltinOperator type, BuiltinOptions builtin_options_type, - flatbuffers::Offset builtin_options); - void SetCustomOp(const string& name, - const std::vector& custom_option, - const std::function& registeration); - - // Build the interpreter for this model. Also, resize and allocate all - // tensors given the shapes of the inputs. - void BuildInterpreter(std::vector> input_shapes, - bool allow_fp32_relax_to_fp16 = false); - - void Invoke(); - - void PopulateStringTensor(int index, const std::vector& content) { - auto tensor = interpreter_->tensor(index); - DynamicBuffer buf; - for (const string& s : content) { - buf.AddString(s.data(), s.length()); - } - buf.WriteToTensor(tensor); - } - - // Populate the tensor given its index. - // TODO(b/110696148) clean up and merge with vector-taking variant below. - template - void PopulateTensor(int index, const std::initializer_list& data) { - T* v = interpreter_->typed_tensor(index); - CHECK(v) << "No tensor with index '" << index << "'."; - for (T f : data) { - *v = f; - ++v; - } - } - - // Populate the tensor given its index. - // TODO(b/110696148) clean up and merge with initializer_list-taking variant - // above. - template - void PopulateTensor(int index, const std::vector& data) { - T* v = interpreter_->typed_tensor(index); - CHECK(v) << "No tensor with index '" << index << "'."; - for (T f : data) { - *v = f; - ++v; - } - } - - // Partially populate the tensor, starting at the given offset. - template - void PopulateTensor(int index, int offset, T* begin, T* end) { - T* v = interpreter_->typed_tensor(index); - memcpy(v + offset, begin, (end - begin) * sizeof(T)); - } - - // Return a vector with the flattened contents of a tensor. - template - std::vector ExtractVector(int index) { - T* v = interpreter_->typed_tensor(index); - CHECK(v); - return std::vector(v, v + GetTensorSize(index)); - } - - std::vector GetTensorShape(int index) { - std::vector result; - TfLiteTensor* t = interpreter_->tensor(index); - for (int i = 0; i < t->dims->size; ++i) { - result.push_back(t->dims->data[i]); - } - return result; - } - - void SetResolver(std::unique_ptr resolver) { - resolver_ = std::move(resolver); - } - - protected: - int32_t GetTensorSize(int index) const; - - flatbuffers::FlatBufferBuilder builder_; - std::unique_ptr interpreter_; - std::unique_ptr resolver_; - - private: - // TODO(gavinbelson): sync this method with - // //tensorflow/contrib/lite/kernels/internal/quantization_util.h?l=31 - template - std::pair QuantizationParams(float f_min, float f_max) { - // These are required by many quantized operations. - CHECK_LE(f_min, 0); - CHECK_GE(f_max, 0); - T q_min = std::numeric_limits::min(); - T q_max = std::numeric_limits::max(); - float range = q_max - q_min; - float scale = (f_max - f_min) / range; - int32_t zero_point = std::min( - q_max, - std::max(q_min, static_cast(std::round(q_min - f_min / scale)))); - return {scale, zero_point}; - } - - template - int AddTensor(TensorData t, std::initializer_list data, - bool is_variable = false) { - int id = tensors_.size(); - - // This is slightly different depending on whether we are adding a - // quantized or a regular tensor. - bool is_quantized = (t.min != 0 || t.max != 0 || t.scale != 0); - - flatbuffers::Offset q_params = 0; - - if (is_quantized) { - if (t.min != 0 || t.max != 0) { - if (t.type == TensorType_UINT8) { - std::tie(t.scale, t.zero_point) = - QuantizationParams(t.min, t.max); - } else if (t.type == TensorType_INT32) { - std::tie(t.scale, t.zero_point) = - QuantizationParams(t.min, t.max); - } else if (t.type == TensorType_INT16) { - std::tie(t.scale, t.zero_point) = - QuantizationParams(t.min, t.max); - } else { - LOG(FATAL) << "No support for the requested quantized type"; - } - t.min = 0; - t.max = 0; - } - - q_params = CreateQuantizationParameters( - builder_, /*min=*/0, /*max=*/0, - builder_.CreateVector({t.scale}), - builder_.CreateVector({t.zero_point})); - } - - int buffer_id = 0; - if (data.size()) { - // Initialize buffers list with empty buffer to allow for non-const - // tensors. - if (buffers_.empty()) { - buffers_.push_back(CreateBuffer(builder_, builder_.CreateVector({}))); - } - - // Add data as a Buffer to buffers list. - buffer_id = buffers_.size(); - auto data_buffer = - builder_.CreateVector(reinterpret_cast(data.begin()), - sizeof(T) * data.size()); - buffers_.push_back(CreateBuffer(builder_, data_buffer)); - } - - tensors_.push_back(CreateTensor(builder_, - builder_.CreateVector(t.shape), t.type, - /*buffer=*/buffer_id, - /*name=*/0, q_params, is_variable)); - - tensor_data_[id] = t; - - return id; - } - - std::map tensor_data_; - std::vector inputs_; - std::vector outputs_; - std::vector> tensors_; - std::vector> opcodes_; - std::vector> operators_; - std::vector> buffers_; - std::map> custom_registrations_; - // A function pointer that gets called after the interpreter is created but - // before evaluation happens. This is useful for applying a delegate. - std::function apply_delegate_fn_; -}; - -// Base class for single op unit tests. -// The tests are parameterized to test multiple kernels for a single op. -// The parameters are strings like "optimized" and "reference" to have better -// readability in test reports. -// -// To use this class: -// * Define a constant map from strings to TfLiteRegistration. -// * Implement a test class that inherits SingleOpTest. -// * Instantiate the test cases with SingleOpTest::GetKernelTags helper -// function. -// * Call GetRegistration to get the TfLiteRegistration to be used before -// building the interpreter. -class SingleOpTest : public ::testing::TestWithParam { - public: - static std::vector GetKernelTags( - const std::map& kernel_map) { - std::vector tags; - for (auto it : kernel_map) { - tags.push_back(it.first); - } - return tags; - } - - protected: - virtual const std::map& GetKernelMap() = 0; - TfLiteRegistration* GetRegistration() { - return GetKernelMap().at(GetParam()); - } -}; - -// Strings have a special implementation that is in test_util.cc -template <> -std::vector SingleOpModel::ExtractVector(int index); -} // namespace tflite - -#endif // TENSORFLOW_CONTRIB_LITE_KERNELS_TEST_UTIL_H_ diff --git a/tensorflow/contrib/lite/kernels/tile_test.cc b/tensorflow/contrib/lite/kernels/tile_test.cc deleted file mode 100644 index e73ca7b7504f6fe891f310d181b0039893f18852..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/tile_test.cc +++ /dev/null @@ -1,256 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include -#include "tensorflow/contrib/lite/c/builtin_op_data.h" -#include "tensorflow/contrib/lite/interpreter.h" -#include "tensorflow/contrib/lite/kernels/register.h" -#include "tensorflow/contrib/lite/kernels/test_util.h" -#include "tensorflow/contrib/lite/model.h" - -namespace tflite { -namespace { - -using ::testing::ElementsAreArray; -class TileOpModel : public SingleOpModel { - public: - TileOpModel(std::initializer_list input_shape, TensorType input_type, - TensorType multiply_type) { - input_ = AddInput(input_type); - multipliers_ = AddInput(TensorType_INT32); - output_ = AddOutput(input_type); - SetBuiltinOp(BuiltinOperator_TILE, BuiltinOptions_TileOptions, 0); - BuildInterpreter({input_shape, {static_cast(input_shape.size())}}); - } - - void SetInputFloat(std::initializer_list data) { - PopulateTensor(input_, data); - } - - void SetInputUInt8(std::initializer_list data) { - PopulateTensor(input_, data); - } - - void SetInputInt32(std::initializer_list data) { - PopulateTensor(input_, data); - } - - void SetInputInt64(std::initializer_list data) { - PopulateTensor(input_, data); - } - - void SetMultipliers(std::initializer_list data) { - PopulateTensor(multipliers_, data); - } - - std::vector GetOutputFloat() { return ExtractVector(output_); } - - std::vector GetOutputUInt8() { return ExtractVector(output_); } - - std::vector GetOutputInt32() { return ExtractVector(output_); } - - std::vector GetOutputInt64() { - return ExtractVector(output_); - } - - std::vector GetOutputShape() { return GetTensorShape(output_); } - - protected: - int input_; - int multipliers_; - int output_; -}; - -TEST(TileTest, Float32Vector) { - TileOpModel m({3}, TensorType_FLOAT32, TensorType_INT32); - m.SetInputFloat({1.f, 2.f, 3.f}); - m.SetMultipliers({2}); - m.Invoke(); - EXPECT_THAT(m.GetOutputFloat(), - ElementsAreArray({1.f, 2.f, 3.f, 1.f, 2.f, 3.f})); -} - -TEST(TileTest, Float32Matrix) { - TileOpModel m({2, 3}, TensorType_FLOAT32, TensorType_INT32); - m.SetInputFloat({ - 11.f, - 12.f, - 13.f, - 21.f, - 22.f, - 23.f, - }); - m.SetMultipliers({2, 1}); - m.Invoke(); - EXPECT_THAT(m.GetOutputFloat(), ElementsAreArray({ - 11.f, - 12.f, - 13.f, - 21.f, - 22.f, - 23.f, - 11.f, - 12.f, - 13.f, - 21.f, - 22.f, - 23.f, - })); - EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({4, 3})); -} - -TEST(TileTest, Float32HighDimension) { - TileOpModel m({1, 2, 3}, TensorType_FLOAT32, TensorType_INT32); - m.SetInputFloat({ - 11.f, - 12.f, - 13.f, - 21.f, - 22.f, - 23.f, - }); - m.SetMultipliers({2, 3, 1}); - m.Invoke(); - EXPECT_THAT( - m.GetOutputFloat(), - ElementsAreArray({11.f, 12.f, 13.f, 21.f, 22.f, 23.f, 11.f, 12.f, 13.f, - 21.f, 22.f, 23.f, 11.f, 12.f, 13.f, 21.f, 22.f, 23.f, - 11.f, 12.f, 13.f, 21.f, 22.f, 23.f, 11.f, 12.f, 13.f, - 21.f, 22.f, 23.f, 11.f, 12.f, 13.f, 21.f, 22.f, 23.f})); - EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2, 6, 3})); -} - -TEST(TileTest, Uint8Matrix) { - TileOpModel m({2, 3}, TensorType_UINT8, TensorType_INT32); - m.SetInputUInt8({ - 11, - 12, - 13, - 21, - 22, - 23, - }); - m.SetMultipliers({2, 1}); - m.Invoke(); - EXPECT_THAT(m.GetOutputUInt8(), ElementsAreArray({ - 11, - 12, - 13, - 21, - 22, - 23, - 11, - 12, - 13, - 21, - 22, - 23, - })); - EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({4, 3})); -} - -TEST(TileTest, Int32Matrix) { - TileOpModel m({2, 3}, TensorType_INT32, TensorType_INT32); - m.SetInputInt32({ - 11, - 12, - 13, - 21, - 22, - 23, - }); - m.SetMultipliers({2, 1}); - m.Invoke(); - EXPECT_THAT(m.GetOutputInt32(), ElementsAreArray({ - 11, - 12, - 13, - 21, - 22, - 23, - 11, - 12, - 13, - 21, - 22, - 23, - })); - EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({4, 3})); -} - -TEST(TileTest, Int64Matrix) { - TileOpModel m({2, 3}, TensorType_INT64, TensorType_INT32); - m.SetInputInt64({ - 11, - 12, - 13, - 21, - 22, - 23, - }); - m.SetMultipliers({2, 1}); - m.Invoke(); - EXPECT_THAT(m.GetOutputInt64(), ElementsAreArray({ - 11, - 12, - 13, - 21, - 22, - 23, - 11, - 12, - 13, - 21, - 22, - 23, - })); - EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({4, 3})); -} - -TEST(TileTest, Int64Matrix64Multipliers) { - TileOpModel m({2, 3}, TensorType_INT64, TensorType_INT64); - m.SetInputInt64({ - 11, - 12, - 13, - 21, - 22, - 23, - }); - m.SetMultipliers({2, 1}); - m.Invoke(); - EXPECT_THAT(m.GetOutputInt64(), ElementsAreArray({ - 11, - 12, - 13, - 21, - 22, - 23, - 11, - 12, - 13, - 21, - 22, - 23, - })); - EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({4, 3})); -} -} // namespace -} // namespace tflite - -int main(int argc, char** argv) { - ::tflite::LogToStderr(); - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/contrib/lite/kernels/transpose_conv.cc b/tensorflow/contrib/lite/kernels/transpose_conv.cc deleted file mode 100644 index 1c4a5ee91d038cb222820c4d35fc713f8f41cb63..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/transpose_conv.cc +++ /dev/null @@ -1,158 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include -#include -#include -#include -#include -#include - -#include "tensorflow/contrib/lite/c/builtin_op_data.h" -#include "tensorflow/contrib/lite/c/c_api_internal.h" -#include "tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h" -#include "tensorflow/contrib/lite/kernels/internal/tensor.h" -#include "tensorflow/contrib/lite/kernels/kernel_util.h" -#include "tensorflow/contrib/lite/kernels/op_macros.h" -#include "tensorflow/contrib/lite/kernels/padding.h" - -namespace tflite { -namespace ops { -namespace builtin { -namespace transpose_conv { - -constexpr int kOutputShapeTensor = 0; -constexpr int kWeightsTensor = 1; -constexpr int kDataInputTensor = 2; -constexpr int kOutputTensor = 0; - -TfLiteStatus ResizeOutputShape(TfLiteContext* context, - const TfLiteTensor* output_shape, - TfLiteTensor* output) { - // Currently only support int32 for output shape. - if (output_shape->type != kTfLiteInt32) { - context->ReportError(context, "Output shape is %d, not int32.", - output_shape->type); - return kTfLiteError; - } - const int output_dimensions = NumElements(output_shape); - TfLiteIntArray* output_shape_array = TfLiteIntArrayCreate(output_dimensions); - for (int i = 0; i < output_dimensions; ++i) { - output_shape_array->data[i] = GetTensorData(output_shape)[i]; - } - - return context->ResizeTensor(context, output, output_shape_array); -} - -TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) { - TF_LITE_ENSURE_EQ(context, NumInputs(node), 3); - TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1); - - const TfLiteTensor* output_shape = - GetInput(context, node, kOutputShapeTensor); - const TfLiteTensor* weights = GetInput(context, node, kWeightsTensor); - const TfLiteTensor* input = GetInput(context, node, kDataInputTensor); - TfLiteTensor* output = GetOutput(context, node, kOutputTensor); - - TF_LITE_ENSURE_EQ(context, NumDimensions(output_shape), 1); - TF_LITE_ENSURE_EQ(context, NumDimensions(input), 4); - TF_LITE_ENSURE_EQ(context, NumDimensions(weights), 4); - - // Currently only supports float32. - const TfLiteType data_type = input->type; - TF_LITE_ENSURE(context, data_type == kTfLiteFloat32); - TF_LITE_ENSURE_EQ(context, output->type, data_type); - TF_LITE_ENSURE_EQ(context, weights->type, data_type); - - // Ensure that weights and inputs have the same channel dimension. - // Note: TOCO will reorder weights in the following format: OHWI. - TF_LITE_ENSURE_EQ(context, SizeOfDimension(input, 3), - SizeOfDimension(weights, 3)); - - if (!IsConstantTensor(output_shape)) { - SetTensorToDynamic(output); - return kTfLiteOk; - } - return ResizeOutputShape(context, output_shape, output); -} - -TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) { - const TfLiteTensor* output_shape = - GetInput(context, node, kOutputShapeTensor); - const TfLiteTensor* weights = GetInput(context, node, kWeightsTensor); - const TfLiteTensor* input = GetInput(context, node, kDataInputTensor); - TfLiteTensor* output = GetOutput(context, node, kOutputTensor); - - const auto* params = - reinterpret_cast(node->builtin_data); - - if (IsDynamicTensor(output)) { - TF_LITE_ENSURE_OK(context, - ResizeOutputShape(context, output_shape, output)); - } - - // Get height and width of the output image. - const int width = SizeOfDimension(output, 2); - const int height = SizeOfDimension(output, 1); - const int filter_width = SizeOfDimension(weights, 1); - const int filter_height = SizeOfDimension(weights, 2); - - const int stride_width = params->stride_width; - const int stride_height = params->stride_height; - - const TfLitePaddingValues& padding_size = - ComputePaddingHeightWidth(stride_height, stride_width, 1, height, width, - filter_height, filter_width, params->padding); - - // Currently only support float32. - switch (input->type) { - case kTfLiteFloat32: { - tflite::ConvParams op_params; - op_params.padding_type = PaddingType::kSame; - op_params.padding_values.width = padding_size.width; - op_params.padding_values.height = padding_size.height; - op_params.stride_width = stride_width; - op_params.stride_height = stride_height; - - reference_ops::TransposeConv( - op_params, GetTensorShape(input), GetTensorData(input), - GetTensorShape(weights), GetTensorData(weights), - GetTensorShape(output), GetTensorData(output), - // Last two args specify im2col which reference_ops ignores. - // (Note this does not lead to a performance regression, as the - // previous optimized version was just a copy of the reference code.) - // TODO(b/110208176): Allocate im2col tensors and switch to - // optimized_ops. - GetTensorShape(output), GetTensorData(output)); - break; - } - default: - context->ReportError(context, "Type %d, not currently supported.", - input->type); - return kTfLiteError; - } - return kTfLiteOk; -} - -} // namespace transpose_conv - -TfLiteRegistration* Register_TRANSPOSE_CONV() { - static TfLiteRegistration r = {nullptr, nullptr, transpose_conv::Prepare, - transpose_conv::Eval}; - return &r; -} - -} // namespace builtin -} // namespace ops -} // namespace tflite diff --git a/tensorflow/contrib/lite/kernels/transpose_conv_test.cc b/tensorflow/contrib/lite/kernels/transpose_conv_test.cc deleted file mode 100644 index 55df8971806ed0baae9f5bcaebd24fb8065ec300..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/kernels/transpose_conv_test.cc +++ /dev/null @@ -1,222 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include -#include -#include "tensorflow/contrib/lite/interpreter.h" -#include "tensorflow/contrib/lite/kernels/register.h" -#include "tensorflow/contrib/lite/kernels/test_util.h" -#include "tensorflow/contrib/lite/model.h" - -namespace tflite { -namespace { - -using ::testing::ElementsAreArray; - -class TransposeConvOpModel : public SingleOpModel { - public: - TransposeConvOpModel(std::initializer_list input_shape, - std::initializer_list filter_shape, Padding padding, - int stride_w, int stride_h) { - output_shape_ = AddInput(TensorType_INT32); - filter_ = AddInput(TensorType_FLOAT32); - input_ = AddInput(TensorType_FLOAT32); - output_ = AddOutput(TensorType_FLOAT32); - SetBuiltinOp( - BuiltinOperator_TRANSPOSE_CONV, BuiltinOptions_TransposeConvOptions, - CreateTransposeConvOptions(builder_, padding, stride_w, stride_h) - .Union()); - BuildInterpreter({{4}, filter_shape, input_shape}); - } - - int output_shape() { return output_shape_; } - int filter() { return filter_; } - int input() { return input_; } - - std::vector GetOutput() { return ExtractVector(output_); } - std::vector GetOutputShape() { return GetTensorShape(output_); } - - private: - int output_shape_; - int filter_; - int input_; - int output_; -}; - -// Test case: -// output = tf.nn.conv2d_backprop_input( -// tf.constant([ 1, 4, 4, 1 ]), -// tf.constant(np.arange(1, 10), shape=[ 3, 3, 1, 1 ], dtype=tf.float32), -// tf.constant(np.arange(1, 17), shape=[ 1, 4, 4, 1 ], dtype=tf.float32), -// [1, 1, 1, 1 ], -// "SAME") -TEST(TransposeConvOpModelTest, SimpleTest) { - TransposeConvOpModel m({1, 4, 4, 1}, {1, 3, 3, 1}, Padding_SAME, 1, 1); - m.PopulateTensor(m.output_shape(), {1, 4, 4, 1}); - m.PopulateTensor(m.filter(), {1, 2, 3, 4, 5, 6, 7, 8, 9}); - m.PopulateTensor( - m.input(), {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}); - m.Invoke(); - - EXPECT_THAT(m.GetOutput(), - ElementsAreArray({29, 62, 83, 75, 99, 192, 237, 198, 207, 372, - 417, 330, 263, 446, 485, 365})); - EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4, 4, 1})); -} - -// Test case: -// filter = tf.constant(np.arange(1, 19), -// shape=[ 3, 3, 1, 2 ], -// dtype=tf.float32) -// output = tf.nn.conv2d_backprop_input( -// tf.constant([ 1, 4, 4, 1 ]), -// filter, -// tf.constant(np.arange(1, 33), shape=[ 1, 4, 4, 2 ], dtype=tf.float32), -// [1, 1, 1, 1 ], -// "SAME") -// And filter value is derived by: -// filter = tf.reshape(tf.transpose(filter, perm=[3, 0, 1, 2]), shape=[18, 1]) -TEST(TransposeConvOpModelTest, TwoFiltersTest) { - TransposeConvOpModel m({1, 4, 4, 2}, {1, 3, 3, 2}, Padding_SAME, 1, 1); - m.PopulateTensor(m.output_shape(), {1, 4, 4, 1}); - m.PopulateTensor(m.filter(), {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, - 13, 14, 15, 16, 17, 18}); - m.PopulateTensor( - m.input(), - {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, - 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32}); - m.Invoke(); - - EXPECT_THAT(m.GetOutput(), - ElementsAreArray({184, 412, 568, 528, 678, 1347, 1689, 1434, 1494, - 2715, 3057, 2442, 1968, 3352, 3652, 2760})); - EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4, 4, 1})); -} - -// Test case: -// filter = tf.constant(np.arange(1, 19), -// shape=[ 3, 3, 1, 2 ], -// dtype=tf.float32) -// output = tf.nn.conv2d_backprop_input( -// tf.constant([ 1, 6, 6, 1 ]), -// filter, -// tf.constant(np.arange(1, 33), shape=[ 1, 4, 4, 2 ], dtype=tf.float32), -// [1, 1, 1, 1 ], -// "VALID") -// And filter value is derived by: -// filter = tf.reshape(tf.transpose(filter, perm=[3, 0, 1, 2]), shape=[1, 18]) -TEST(TransposeConvOpModelTest, PaddingValidTest) { - TransposeConvOpModel m({1, 4, 4, 2}, {1, 3, 3, 2}, Padding_VALID, 1, 1); - m.PopulateTensor(m.output_shape(), {1, 6, 6, 1}); - m.PopulateTensor(m.filter(), {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, - 13, 14, 15, 16, 17, 18}); - m.PopulateTensor( - m.input(), - {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, - 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32}); - m.Invoke(); - - EXPECT_THAT( - m.GetOutput(), - ElementsAreArray({5, 22, 59, 101, 114, 83, 52, 184, 412, - 568, 528, 344, 237, 678, 1347, 1689, 1434, 879, - 597, 1494, 2715, 3057, 2442, 1431, 856, 1968, 3352, - 3652, 2760, 1548, 689, 1534, 2543, 2729, 2010, 1103})); - EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 6, 6, 1})); -} - -// Test case: -// filter = tf.constant(np.arange(1, 10), -// shape=[ 3, 3, 1, 1 ], -// dtype=tf.float32) -// output = tf.nn.conv2d_backprop_input( -// tf.constant([ 1, 5, 5, 1 ]), -// filter, -// tf.constant(np.arange(1, 5), shape=[ 1, 2, 2, 1 ], dtype=tf.float32), -// [1, 2, 2, 1 ], -// "VALID") -TEST(TransposeConvOpModelTest, StrideValidTest) { - TransposeConvOpModel m({1, 2, 2, 1}, {1, 3, 3, 1}, Padding_VALID, 2, 2); - m.PopulateTensor(m.output_shape(), {1, 5, 5, 1}); - m.PopulateTensor(m.filter(), {1, 2, 3, 4, 5, 6, 7, 8, 9}); - m.PopulateTensor(m.input(), {1, 2, 3, 4}); - m.Invoke(); - - EXPECT_THAT( - m.GetOutput(), - ElementsAreArray({1, 2, 5, 4, 6, 4, 5, 14, 10, 12, 10, 14, 36, - 24, 30, 12, 15, 34, 20, 24, 21, 24, 55, 32, 36})); - EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 5, 5, 1})); -} - -// Test case: -// filter = tf.constant(np.arange(1, 19), -// shape=[ 3, 3, 2, 1 ], -// dtype=tf.float32) -// output = tf.nn.conv2d_backprop_input( -// tf.constant([ 1, 5, 5, 2 ]), -// filter, -// tf.constant(np.arange(1, 5), shape=[ 1, 2, 2, 1 ], dtype=tf.float32), -// [1, 2, 2, 1 ], -// "VALID") -TEST(TransposeConvOpModelTest, MultiChannelTest) { - TransposeConvOpModel m({1, 2, 2, 1}, {2, 3, 3, 1}, Padding_VALID, 2, 2); - m.PopulateTensor(m.output_shape(), {1, 5, 5, 2}); - m.PopulateTensor(m.filter(), {1, 3, 5, 7, 9, 11, 13, 15, 17, 2, 4, 6, - 8, 10, 12, 14, 16, 18}); - m.PopulateTensor(m.input(), {1, 2, 3, 4}); - m.Invoke(); - - EXPECT_THAT( - m.GetOutput(), - ElementsAreArray({1, 2, 3, 4, 7, 10, 6, 8, 10, 12, 7, 8, 9, - 10, 25, 28, 18, 20, 22, 24, 16, 20, 24, 28, 62, 72, - 42, 48, 54, 60, 21, 24, 27, 30, 61, 68, 36, 40, 44, - 48, 39, 42, 45, 48, 103, 110, 60, 64, 68, 72})); - EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 5, 5, 2})); -} - -// Test case: -// filter = tf.constant(np.random.randint(1, 10, size=9), -// shape=[ 3, 3, 1, 1 ], -// dtype=tf.float32) -// output = tf.nn.conv2d_backprop_input( -// tf.constant([ 1, 3, 4, 1 ]), -// filter, -// tf.constant([323, 521], shape=[ 1, 1, 2, 1], dtype=tf.float32), -// [1, 3, 3, 1 ], -// "SAME") -// And filter value is derived by: -// filter = tf.reshape(tf.transpose(filter, perm=[3, 0, 1, 2]), shape=[-1]) -TEST(TransposeConvOpModelTest, AccuracyTest) { - TransposeConvOpModel m({1, 1, 2, 1}, {1, 3, 3, 1}, Padding_SAME, 3, 3); - m.PopulateTensor(m.output_shape(), {1, 3, 4, 1}); - m.PopulateTensor(m.filter(), {9, 5, 6, 9, 8, 5, 3, 1, 4}); - m.PopulateTensor(m.input(), {323, 521}); - m.Invoke(); - - EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear( - {1615., 1938., 4689., 2605., 2584., 1615., - 4689., 4168., 323., 1292., 1563., 521.}))); - EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 3, 4, 1})); -} - -} // namespace -} // namespace tflite - -int main(int argc, char** argv) { - ::tflite::LogToStderr(); - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/contrib/lite/lib_package/create_ios_frameworks.sh b/tensorflow/contrib/lite/lib_package/create_ios_frameworks.sh deleted file mode 100755 index 6195426d6d441e858fbe225c132b409ac0a0be32..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/lib_package/create_ios_frameworks.sh +++ /dev/null @@ -1,82 +0,0 @@ -#!/bin/bash -x -# 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. -# ============================================================================== - -# TODO(ycling): Refactoring - Move this script into `tools/make`. -set -e - -echo "Starting" -TFLITE_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/.." - -TMP_DIR=$(mktemp -d) -echo "Package dir: " $TMP_DIR -FW_DIR=$TMP_DIR/tensorflow_lite_ios_frameworks -FW_DIR_TFLITE=$FW_DIR/tensorflow_lite.framework -FW_DIR_TFLITE_HDRS=$FW_DIR_TFLITE/Headers - -echo "Creating target Headers directories" -mkdir -p $FW_DIR_TFLITE_HDRS - -echo "Headers, populating: TensorFlow Lite" -cd $TFLITE_DIR/../../.. - -find tensorflow/contrib/lite -name '*.h' \ - -not -path 'tensorflow/contrib/lite/tools/*' \ - -not -path 'tensorflow/contrib/lite/examples/*' \ - -not -path 'tensorflow/contrib/lite/gen/*' \ - -not -path 'tensorflow/contrib/lite/toco/*' \ - -not -path 'tensorflow/contrib/lite/nnapi/*' \ - -not -path 'tensorflow/contrib/lite/java/*' \ - | tar -cf $FW_DIR_TFLITE_HDRS/tmp.tar -T - -cd $FW_DIR_TFLITE_HDRS -tar xf tmp.tar -rm -f tmp.tar - -echo "Headers, populating: Flatbuffer" -cd $TFLITE_DIR/tools/make/downloads/flatbuffers/include/ -find . -name '*.h' | tar -cf $FW_DIR_TFLITE_HDRS/tmp.tar -T - -cd $FW_DIR_TFLITE_HDRS -tar xf tmp.tar -rm -f tmp.tar - -cd $TFLITE_DIR/../../.. -echo "Generate master LICENSE file and copy to target" -bazel build //tensorflow/tools/lib_package:clicenses_generate -cp $TFLITE_DIR/../../../bazel-genfiles/tensorflow/tools/lib_package/include/tensorflow/c/LICENSE \ - $FW_DIR_TFLITE - -echo "Copying static libraries" -cp $TFLITE_DIR/tools/make/gen/lib/libtensorflow-lite.a \ - $FW_DIR_TFLITE/tensorflow_lite - -# This is required, otherwise they interfere with the documentation of the -# pod at cocoapods.org. -echo "Remove all README files" -cd $FW_DIR_TFLITE_HDRS -find . -type f -name README\* -exec rm -f {} \; -find . -type f -name readme\* -exec rm -f {} \; - -TARGET_GEN_LOCATION="$TFLITE_DIR/gen/ios_frameworks" -echo "Moving results to target: " $TARGET_GEN_LOCATION -cd $FW_DIR -zip -q -r tensorflow_lite.framework.zip tensorflow_lite.framework -x .DS_Store -rm -rf $TARGET_GEN_LOCATION -mkdir -p $TARGET_GEN_LOCATION -cp -r tensorflow_lite.framework.zip $TARGET_GEN_LOCATION - -echo "Cleaning up" -rm -rf $TMP_DIR - -echo "Finished" diff --git a/tensorflow/contrib/lite/model.h b/tensorflow/contrib/lite/model.h deleted file mode 100644 index 9505824dcc933b60b89b9b98adcb1f685278b9da..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/model.h +++ /dev/null @@ -1,189 +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. -==============================================================================*/ -// Deserialization infrastructure for tflite. Provides functionality -// to go from a serialized tflite model in flatbuffer format to an -// interpreter. -// -// using namespace tflite; -// StderrReporter error_reporter; -// auto model = FlatBufferModel::BuildFromFile("interesting_model.tflite", -// &error_reporter); -// MyOpResolver resolver; // You need to subclass OpResolver to provide -// // implementations. -// InterpreterBuilder builder(*model, resolver); -// std::unique_ptr interpreter; -// if(builder(&interpreter) == kTfLiteOk) { -// .. run model inference with interpreter -// } -// -// OpResolver must be defined to provide your kernel implementations to the -// interpreter. This is environment specific and may consist of just the builtin -// ops, or some custom operators you defined to extend tflite. -#ifndef TENSORFLOW_CONTRIB_LITE_MODEL_H_ -#define TENSORFLOW_CONTRIB_LITE_MODEL_H_ - -#include -#include "tensorflow/contrib/lite/core/api/error_reporter.h" -#include "tensorflow/contrib/lite/core/api/op_resolver.h" -#include "tensorflow/contrib/lite/interpreter.h" -#include "tensorflow/contrib/lite/mutable_op_resolver.h" -#include "tensorflow/contrib/lite/schema/schema_generated.h" - -namespace tflite { - -// Abstract interface that verifies whether a given model is legit. -// It facilitates the use-case to verify and build a model without loading it -// twice. -class TfLiteVerifier { - public: - // Returns true if the model is legit. - virtual bool Verify(const char* data, int length, - ErrorReporter* reporter) = 0; - virtual ~TfLiteVerifier() {} -}; - -// An RAII object that represents a read-only tflite model, copied from disk, -// or mmapped. This uses flatbuffers as the serialization format. -class FlatBufferModel { - public: - // Builds a model based on a file. - // Caller retains ownership of `error_reporter` and must ensure its lifetime - // is longer than the FlatBufferModel instance. - // Returns a nullptr in case of failure. - static std::unique_ptr BuildFromFile( - const char* filename, - ErrorReporter* error_reporter = DefaultErrorReporter()); - - // Verifies whether the content of the file is legit, then builds a model - // based on the file. - // Caller retains ownership of `error_reporter` and must ensure its lifetime - // is longer than the FlatBufferModel instance. - // Returns a nullptr in case of failure. - static std::unique_ptr VerifyAndBuildFromFile( - const char* filename, TfLiteVerifier* verifier = nullptr, - ErrorReporter* error_reporter = DefaultErrorReporter()); - - // Builds a model based on a pre-loaded flatbuffer. The caller retains - // ownership of the buffer and should keep it alive until the returned object - // is destroyed. Caller retains ownership of `error_reporter` and must ensure - // its lifetime is longer than the FlatBufferModel instance. - // Returns a nullptr in case of failure. - static std::unique_ptr BuildFromBuffer( - const char* buffer, size_t buffer_size, - ErrorReporter* error_reporter = DefaultErrorReporter()); - - // Builds a model directly from a flatbuffer pointer. The caller retains - // ownership of the buffer and should keep it alive until the returned object - // is destroyed. Caller retains ownership of `error_reporter` and must ensure - // its lifetime is longer than the FlatBufferModel instance. - // Returns a nullptr in case of failure. - static std::unique_ptr BuildFromModel( - const tflite::Model* model_spec, - ErrorReporter* error_reporter = DefaultErrorReporter()); - - // Releases memory or unmaps mmaped memory. - ~FlatBufferModel(); - - // Copying or assignment is disallowed to simplify ownership semantics. - FlatBufferModel(const FlatBufferModel&) = delete; - FlatBufferModel& operator=(const FlatBufferModel&) = delete; - - bool initialized() const { return model_ != nullptr; } - const tflite::Model* operator->() const { return model_; } - const tflite::Model* GetModel() const { return model_; } - ErrorReporter* error_reporter() const { return error_reporter_; } - const Allocation* allocation() const { return allocation_; } - - // Returns true if the model identifier is correct (otherwise false and - // reports an error). - bool CheckModelIdentifier() const; - - private: - // Loads a model from a given allocation. FlatBufferModel will take over the - // ownership of `allocation`, and delete it in destructor. The ownership of - // `error_reporter`remains with the caller and must have lifetime at least - // as much as FlatBufferModel. This is to allow multiple models to use the - // same ErrorReporter instance. - FlatBufferModel(Allocation* allocation, - ErrorReporter* error_reporter = DefaultErrorReporter()); - - // Loads a model from Model flatbuffer. The `model` has to remain alive and - // unchanged until the end of this flatbuffermodel's lifetime. - FlatBufferModel(const Model* model, ErrorReporter* error_reporter); - - // Flatbuffer traverser pointer. (Model* is a pointer that is within the - // allocated memory of the data allocated by allocation's internals. - const tflite::Model* model_ = nullptr; - // The error reporter to use for model errors and subsequent errors when - // the interpreter is created - ErrorReporter* error_reporter_; - // The allocator used for holding memory of the model. - Allocation* allocation_ = nullptr; -}; - -// Build an interpreter capable of interpreting `model`. -// -// model: a scoped model whose lifetime must be at least as long as -// the interpreter. In principle multiple interpreters can be made from -// a single model. -// op_resolver: An instance that implements the Resolver interface which maps -// custom op names and builtin op codes to op registrations. -// reportError: a functor that is called to report errors that handles -// printf var arg semantics. The lifetime of the reportError object must -// be greater than or equal to the Interpreter created by operator(). -// -// Returns a kTfLiteOk when successful and sets interpreter to a valid -// Interpreter. Note: the user must ensure the model lifetime is at least as -// long as interpreter's lifetime. -class InterpreterBuilder { - public: - InterpreterBuilder(const FlatBufferModel& model, - const OpResolver& op_resolver); - // Builds an interpreter given only the raw flatbuffer Model object (instead - // of a FlatBufferModel). Mostly used for testing. - // If `error_reporter` is null, then DefaultErrorReporter() is used. - InterpreterBuilder(const ::tflite::Model* model, - const OpResolver& op_resolver, - ErrorReporter* error_reporter = DefaultErrorReporter()); - ~InterpreterBuilder(); - InterpreterBuilder(const InterpreterBuilder&) = delete; - InterpreterBuilder& operator=(const InterpreterBuilder&) = delete; - TfLiteStatus operator()(std::unique_ptr* interpreter); - TfLiteStatus operator()(std::unique_ptr* interpreter, - int num_threads); - - private: - TfLiteStatus BuildLocalIndexToRegistrationMapping(); - TfLiteStatus ParseNodes( - const flatbuffers::Vector>* operators, - Interpreter* interpreter); - TfLiteStatus ParseTensors( - const flatbuffers::Vector>* buffers, - const flatbuffers::Vector>* tensors, - Interpreter* interpreter); - TfLiteStatus ApplyDelegates(Interpreter* interpreter); - - const ::tflite::Model* model_; - const OpResolver& op_resolver_; - ErrorReporter* error_reporter_; - - std::vector flatbuffer_op_index_to_registration_; - std::vector flatbuffer_op_index_to_registration_types_; - const Allocation* allocation_ = nullptr; -}; - -} // namespace tflite - -#endif // TENSORFLOW_CONTRIB_LITE_MODEL_H_ diff --git a/tensorflow/contrib/lite/model_test.cc b/tensorflow/contrib/lite/model_test.cc deleted file mode 100644 index b969bea5dcff2f5347ba0aa90f649ba3de89702b..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/model_test.cc +++ /dev/null @@ -1,327 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include -#include -#include -#include -#include -#include -#include - -#include "tensorflow/contrib/lite/model.h" - -#include -#include "tensorflow/contrib/lite/core/api/error_reporter.h" -#include "tensorflow/contrib/lite/kernels/register.h" -#include "tensorflow/contrib/lite/testing/util.h" - -// Comparison for TfLiteRegistration. Since TfLiteRegistration is a C object, -// we must declare this in global namespace, so argument-dependent operator -// lookup works. -inline bool operator==(const TfLiteRegistration& a, - const TfLiteRegistration& b) { - return a.invoke == b.invoke && a.init == b.init && a.prepare == b.prepare && - a.free == b.free; -} - -namespace tflite { - -// Provide a dummy operation that does nothing. -namespace { -void* dummy_init(TfLiteContext*, const char*, size_t) { return nullptr; } -void dummy_free(TfLiteContext*, void*) {} -TfLiteStatus dummy_resize(TfLiteContext*, TfLiteNode*) { return kTfLiteOk; } -TfLiteStatus dummy_invoke(TfLiteContext*, TfLiteNode*) { return kTfLiteOk; } -TfLiteRegistration dummy_reg = {dummy_init, dummy_free, dummy_resize, - dummy_invoke}; -} // namespace - -// Provide a trivial resolver that returns a constant value no matter what -// op is asked for. -class TrivialResolver : public OpResolver { - public: - explicit TrivialResolver(TfLiteRegistration* constant_return = nullptr) - : constant_return_(constant_return) {} - // Find the op registration of a custom operator by op name. - const TfLiteRegistration* FindOp(tflite::BuiltinOperator op, - int version) const override { - return constant_return_; - } - // Find the op registration of a custom operator by op name. - const TfLiteRegistration* FindOp(const char* op, int version) const override { - return constant_return_; - } - - private: - TfLiteRegistration* constant_return_; -}; - -TEST(BasicFlatBufferModel, TestNonExistantFiles) { - ASSERT_TRUE(!FlatBufferModel::BuildFromFile("/tmp/tflite_model_1234")); -} - -// Make sure a model with nothing in it loads properly. -TEST(BasicFlatBufferModel, TestEmptyModelsAndNullDestination) { - auto model = FlatBufferModel::BuildFromFile( - "tensorflow/contrib/lite/testdata/empty_model.bin"); - ASSERT_TRUE(model); - // Now try to build it into a model. - std::unique_ptr interpreter; - ASSERT_EQ(InterpreterBuilder(*model, TrivialResolver())(&interpreter), - kTfLiteOk); - ASSERT_NE(interpreter, nullptr); - ASSERT_NE(InterpreterBuilder(*model, TrivialResolver())(nullptr), kTfLiteOk); -} - -// Make sure currently unsupported # of subgraphs are checked -// TODO(aselle): Replace this test when multiple subgraphs are supported. -TEST(BasicFlatBufferModel, TestZeroAndMultipleSubgraphs) { - auto m1 = FlatBufferModel::BuildFromFile( - "tensorflow/contrib/lite/testdata/0_subgraphs.bin"); - ASSERT_TRUE(m1); - std::unique_ptr interpreter1; - ASSERT_NE(InterpreterBuilder(*m1, TrivialResolver())(&interpreter1), - kTfLiteOk); - - auto m2 = FlatBufferModel::BuildFromFile( - "tensorflow/contrib/lite/testdata/2_subgraphs.bin"); - ASSERT_TRUE(m2); - std::unique_ptr interpreter2; - ASSERT_NE(InterpreterBuilder(*m2, TrivialResolver())(&interpreter2), - kTfLiteOk); -} - -// Test what happens if we cannot bind any of the ops. -TEST(BasicFlatBufferModel, TestModelWithoutNullRegistrations) { - auto model = FlatBufferModel::BuildFromFile( - "tensorflow/contrib/lite/testdata/test_model.bin"); - ASSERT_TRUE(model); - // Check that we get an error code and interpreter pointer is reset. - std::unique_ptr interpreter(new Interpreter); - ASSERT_NE(InterpreterBuilder(*model, TrivialResolver(nullptr))(&interpreter), - kTfLiteOk); - ASSERT_EQ(interpreter, nullptr); -} - -// Make sure model is read to interpreter propelrly -TEST(BasicFlatBufferModel, TestModelInInterpreter) { - auto model = FlatBufferModel::BuildFromFile( - "tensorflow/contrib/lite/testdata/test_model.bin"); - ASSERT_TRUE(model); - // Check that we get an error code and interpreter pointer is reset. - std::unique_ptr interpreter(new Interpreter); - ASSERT_EQ( - InterpreterBuilder(*model, TrivialResolver(&dummy_reg))(&interpreter), - kTfLiteOk); - ASSERT_NE(interpreter, nullptr); - ASSERT_EQ(interpreter->tensors_size(), 4); - ASSERT_EQ(interpreter->nodes_size(), 2); - std::vector inputs = {0, 1}; - std::vector outputs = {2, 3}; - ASSERT_EQ(interpreter->inputs(), inputs); - ASSERT_EQ(interpreter->outputs(), outputs); - - EXPECT_EQ(std::string(interpreter->GetInputName(0)), "input0"); - EXPECT_EQ(std::string(interpreter->GetInputName(1)), "input1"); - EXPECT_EQ(std::string(interpreter->GetOutputName(0)), "out1"); - EXPECT_EQ(std::string(interpreter->GetOutputName(1)), "out2"); - - // Make sure all input tensors are correct - TfLiteTensor* i0 = interpreter->tensor(0); - ASSERT_EQ(i0->type, kTfLiteFloat32); - ASSERT_NE(i0->data.raw, nullptr); // mmapped - ASSERT_EQ(i0->allocation_type, kTfLiteMmapRo); - TfLiteTensor* i1 = interpreter->tensor(1); - ASSERT_EQ(i1->type, kTfLiteFloat32); - ASSERT_EQ(i1->data.raw, nullptr); - ASSERT_EQ(i1->allocation_type, kTfLiteArenaRw); - TfLiteTensor* o0 = interpreter->tensor(2); - ASSERT_EQ(o0->type, kTfLiteFloat32); - ASSERT_EQ(o0->data.raw, nullptr); - ASSERT_EQ(o0->allocation_type, kTfLiteArenaRw); - TfLiteTensor* o1 = interpreter->tensor(3); - ASSERT_EQ(o1->type, kTfLiteFloat32); - ASSERT_EQ(o1->data.raw, nullptr); - ASSERT_EQ(o1->allocation_type, kTfLiteArenaRw); - - // Check op 0 which has inputs {0, 1} outputs {2}. - { - const std::pair* node_and_reg0 = - interpreter->node_and_registration(0); - ASSERT_NE(node_and_reg0, nullptr); - const TfLiteNode& node0 = node_and_reg0->first; - const TfLiteRegistration& reg0 = node_and_reg0->second; - TfLiteIntArray* desired_inputs = TfLiteIntArrayCreate(2); - desired_inputs->data[0] = 0; - desired_inputs->data[1] = 1; - TfLiteIntArray* desired_outputs = TfLiteIntArrayCreate(1); - desired_outputs->data[0] = 2; - ASSERT_TRUE(TfLiteIntArrayEqual(node0.inputs, desired_inputs)); - ASSERT_TRUE(TfLiteIntArrayEqual(node0.outputs, desired_outputs)); - TfLiteIntArrayFree(desired_inputs); - TfLiteIntArrayFree(desired_outputs); - ASSERT_EQ(reg0, dummy_reg); - } - - // Check op 1 which has inputs {2} outputs {3}. - { - const std::pair* node_and_reg1 = - interpreter->node_and_registration(1); - ASSERT_NE(node_and_reg1, nullptr); - const TfLiteNode& node1 = node_and_reg1->first; - const TfLiteRegistration& reg1 = node_and_reg1->second; - TfLiteIntArray* desired_inputs = TfLiteIntArrayCreate(1); - TfLiteIntArray* desired_outputs = TfLiteIntArrayCreate(1); - desired_inputs->data[0] = 2; - desired_outputs->data[0] = 3; - ASSERT_TRUE(TfLiteIntArrayEqual(node1.inputs, desired_inputs)); - ASSERT_TRUE(TfLiteIntArrayEqual(node1.outputs, desired_outputs)); - TfLiteIntArrayFree(desired_inputs); - TfLiteIntArrayFree(desired_outputs); - ASSERT_EQ(reg1, dummy_reg); - } -} - -// Test that loading a model with TensorFlow ops fails when the flex delegate is -// not linked into the target. -TEST(FlexModel, FailureWithoutFlexDelegate) { - auto model = FlatBufferModel::BuildFromFile( - "tensorflow/contrib/lite/testdata/multi_add_flex.bin"); - ASSERT_TRUE(model); - - // Note that creation will succeed when using the BuiltinOpResolver, but - // unless the appropriate delegate is linked into the target or the client - // explicitly installs the delegate, execution will fail. - std::unique_ptr interpreter; - ASSERT_EQ(InterpreterBuilder(*model, - ops::builtin::BuiltinOpResolver{})(&interpreter), - kTfLiteOk); - ASSERT_TRUE(interpreter); - - // As the flex ops weren't resolved implicitly by the flex delegate, runtime - // allocation and execution will fail. - ASSERT_EQ(interpreter->AllocateTensors(), kTfLiteError); -} - -// This tests on a flatbuffer that defines a shape of 2 to be a memory mapped -// buffer. But the buffer is provided to be only 1 element. -TEST(BasicFlatBufferModel, TestBrokenMmap) { - ASSERT_FALSE(FlatBufferModel::BuildFromFile( - "tensorflow/contrib/lite/testdata/test_model_broken.bin")); -} - -TEST(BasicFlatBufferModel, TestNullModel) { - // Check that we get an error code and interpreter pointer is reset. - std::unique_ptr interpreter(new Interpreter); - ASSERT_NE( - InterpreterBuilder(nullptr, TrivialResolver(&dummy_reg))(&interpreter), - kTfLiteOk); - ASSERT_EQ(interpreter.get(), nullptr); -} - -// Mocks the verifier by setting the result in ctor. -class FakeVerifier : public tflite::TfLiteVerifier { - public: - explicit FakeVerifier(bool result) : result_(result) {} - bool Verify(const char* data, int length, - tflite::ErrorReporter* reporter) override { - return result_; - } - - private: - bool result_; -}; - -TEST(BasicFlatBufferModel, TestWithTrueVerifier) { - FakeVerifier verifier(true); - ASSERT_TRUE(FlatBufferModel::VerifyAndBuildFromFile( - "tensorflow/contrib/lite/testdata/test_model.bin", - &verifier)); -} - -TEST(BasicFlatBufferModel, TestWithFalseVerifier) { - FakeVerifier verifier(false); - ASSERT_FALSE(FlatBufferModel::VerifyAndBuildFromFile( - "tensorflow/contrib/lite/testdata/test_model.bin", - &verifier)); -} - -TEST(BasicFlatBufferModel, TestWithNullVerifier) { - ASSERT_TRUE(FlatBufferModel::VerifyAndBuildFromFile( - "tensorflow/contrib/lite/testdata/test_model.bin", nullptr)); -} - -// This makes sure the ErrorReporter is marshalled from FlatBufferModel to -// the Interpreter. -TEST(BasicFlatBufferModel, TestCustomErrorReporter) { - TestErrorReporter reporter; - auto model = FlatBufferModel::BuildFromFile( - "tensorflow/contrib/lite/testdata/empty_model.bin", - &reporter); - ASSERT_TRUE(model); - - std::unique_ptr interpreter; - TrivialResolver resolver; - InterpreterBuilder(*model, resolver)(&interpreter); - ASSERT_NE(interpreter->Invoke(), kTfLiteOk); - ASSERT_EQ(reporter.num_calls(), 1); -} - -// This makes sure the ErrorReporter is marshalled from FlatBufferModel to -// the Interpreter. -TEST(BasicFlatBufferModel, TestNullErrorReporter) { - auto model = FlatBufferModel::BuildFromFile( - "tensorflow/contrib/lite/testdata/empty_model.bin", nullptr); - ASSERT_TRUE(model); - - std::unique_ptr interpreter; - TrivialResolver resolver; - InterpreterBuilder(*model, resolver)(&interpreter); - ASSERT_NE(interpreter->Invoke(), kTfLiteOk); -} - -// Test that loading model directly from a Model flatbuffer works. -TEST(BasicFlatBufferModel, TestBuildFromModel) { - TestErrorReporter reporter; - FileCopyAllocation model_allocation( - "tensorflow/contrib/lite/testdata/test_model.bin", &reporter); - ASSERT_TRUE(model_allocation.valid()); - ::flatbuffers::Verifier verifier( - reinterpret_cast(model_allocation.base()), - model_allocation.bytes()); - ASSERT_TRUE(VerifyModelBuffer(verifier)); - const Model* model_fb = ::tflite::GetModel(model_allocation.base()); - - auto model = FlatBufferModel::BuildFromModel(model_fb); - ASSERT_TRUE(model); - - std::unique_ptr interpreter; - ASSERT_EQ( - InterpreterBuilder(*model, TrivialResolver(&dummy_reg))(&interpreter), - kTfLiteOk); - ASSERT_NE(interpreter, nullptr); -} - -// TODO(aselle): Add tests for serialization of builtin op data types. -// These tests will occur with the evaluation tests of individual operators, -// not here. - -} // namespace tflite - -int main(int argc, char** argv) { - ::tflite::LogToStderr(); - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/contrib/lite/models/BUILD b/tensorflow/contrib/lite/models/BUILD deleted file mode 100644 index efa47b06fa7f06cc6312535713ec582af4705d85..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/models/BUILD +++ /dev/null @@ -1,14 +0,0 @@ -# Model tests -package( - default_visibility = ["//visibility:public"], -) - -licenses(["notice"]) # Apache 2.0 - -exports_files(["LICENSE"]) - -load("//tensorflow/contrib/lite:build_def.bzl", "tflite_copts") - -exports_files(glob([ - "testdata/*", -])) diff --git a/tensorflow/contrib/lite/models/smartreply/BUILD b/tensorflow/contrib/lite/models/smartreply/BUILD deleted file mode 100644 index 9d88c396ba69948e3ae285c913a4499a1409b93a..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/models/smartreply/BUILD +++ /dev/null @@ -1,89 +0,0 @@ -package(default_visibility = ["//visibility:public"]) - -load("//tensorflow/contrib/lite:build_def.bzl", "tflite_copts", "gen_selected_ops") - -licenses(["notice"]) # Apache 2.0 - -gen_selected_ops( - name = "smartreply_ops", - model = "@tflite_smartreply//:smartreply.tflite", -) - -cc_library( - name = "custom_ops", - srcs = [ - "ops/extract_feature.cc", - "ops/normalize.cc", - "ops/predict.cc", - ":smartreply_ops", - ], - copts = tflite_copts(), - deps = [ - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite:string_util", - "//tensorflow/contrib/lite/kernels:builtin_ops", - "@com_google_absl//absl/strings", - "@com_googlesource_code_re2//:re2", - "@farmhash_archive//:farmhash", - ], -) - -cc_library( - name = "predictor_lib", - srcs = ["predictor.cc"], - hdrs = ["predictor.h"], - copts = tflite_copts(), - deps = [ - ":custom_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite:string_util", - "//tensorflow/contrib/lite/kernels:builtin_ops", - "@com_google_absl//absl/strings", - "@com_googlesource_code_re2//:re2", - ], -) - -cc_test( - name = "extract_feature_op_test", - size = "small", - srcs = ["ops/extract_feature_test.cc"], - tags = ["no_oss"], - deps = [ - ":custom_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:builtin_ops", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - "@farmhash_archive//:farmhash", - ], -) - -cc_test( - name = "normalize_op_test", - size = "small", - srcs = ["ops/normalize_test.cc"], - tags = ["no_oss"], - deps = [ - ":custom_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite:string_util", - "//tensorflow/contrib/lite/kernels:builtin_ops", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) - -cc_test( - name = "predict_op_test", - size = "small", - srcs = ["ops/predict_test.cc"], - tags = ["no_oss"], - deps = [ - ":custom_ops", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite:string_util", - "//tensorflow/contrib/lite/kernels:builtin_ops", - "//tensorflow/contrib/lite/kernels:test_util", - "@com_google_googletest//:gtest", - ], -) diff --git a/tensorflow/contrib/lite/models/smartreply/demo/app/src/main/BUILD b/tensorflow/contrib/lite/models/smartreply/demo/app/src/main/BUILD deleted file mode 100644 index 2e5033dab1356e2dbf2eef2b8c14e1ac7fc7566c..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/models/smartreply/demo/app/src/main/BUILD +++ /dev/null @@ -1,68 +0,0 @@ -load("@build_bazel_rules_android//android:rules.bzl", "android_binary") - -package(default_visibility = ["//visibility:public"]) - -licenses(["notice"]) # Apache 2.0 - -load( - "//tensorflow/contrib/lite:build_def.bzl", - "tflite_copts", - "tflite_jni_binary", -) - -filegroup( - name = "assets", - srcs = [ - "@tflite_smartreply//:model_files", - ], -) - -android_binary( - name = "SmartReplyDemo", - srcs = glob(["java/**/*.java"]), - aapt_version = "aapt", - assets = [":assets"], - assets_dir = "", - custom_package = "com.example.android.smartreply", - manifest = "AndroidManifest.xml", - nocompress_extensions = [ - ".tflite", - ], - resource_files = glob(["res/**"]), - tags = ["manual"], - deps = [ - ":smartreply_runtime", - "@androidsdk//com.android.support:support-v13-25.2.0", - "@androidsdk//com.android.support:support-v4-25.2.0", - ], -) - -cc_library( - name = "smartreply_runtime", - srcs = ["libsmartreply_jni.so"], - visibility = ["//visibility:public"], -) - -tflite_jni_binary( - name = "libsmartreply_jni.so", - deps = [ - ":smartreply_jni_lib", - ], -) - -cc_library( - name = "smartreply_jni_lib", - srcs = [ - "smartreply_jni.cc", - ], - copts = tflite_copts(), - linkopts = [ - "-lm", - "-ldl", - ], - deps = [ - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/models/smartreply:predictor_lib", - ], - alwayslink = 1, -) diff --git a/tensorflow/contrib/lite/models/smartreply/g3doc/README.md b/tensorflow/contrib/lite/models/smartreply/g3doc/README.md deleted file mode 100644 index a6d75648b3f3da98afd85daad6c2234e73a802e8..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/models/smartreply/g3doc/README.md +++ /dev/null @@ -1,146 +0,0 @@ -# Smart Reply Model - -## What is On-Device Smart Reply Model? - -Smart Replies are contextually relevant, one-touch responses that help the user -to reply to an incoming text message (or email) efficiently and effortlessly. -Smart Replies have been highly successful across several Google products -including -[Gmail](https://www.blog.google/products/gmail/save-time-with-smart-reply-in-gmail/), -[Inbox](https://www.blog.google/products/gmail/computer-respond-to-this-email/) -and -[Allo](https://blog.google/products/allo/google-allo-smarter-messaging-app/). - -The On-device Smart Reply model is targeted towards text chat use cases. It has -a completely different architecture from its cloud-based counterparts, and is -built specifically for memory constraints devices such as phones & watches. It -has been successfully used to provide [Smart Replies on Android -Wear](https://research.googleblog.com/2017/02/on-device-machine-intelligence.html) -to all first- & third-party apps. - -The on-device model comes with several benefits. It is: - -* **Faster**: The model resides on the device and does not require internet - connectivity. Thus, the inference is very fast and has an average latency of - only a few milliseconds. -* **Resource efficient**: The model has a small memory footprint on - the device. -* **Privacy-friendly**: The user data never leaves the device and this - eliminates any privacy restrictions. - -A caveat, though, is that the on-device model has lower triggering rate than its -cloud counterparts (triggering rate is the percentage of times the model -suggests a response for an incoming message). - -## When to use this Model? - -The On-Device Smart Reply model is aimed towards improving the messaging -experience for day-to-day conversational chat messages. We recommend using this -model for similar use cases. Some sample messages on which the model does well -are provided in this [tsv -file](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite/models/testdata/smartreply_samples.tsv) -for reference. The file format is: - -``` - {incoming_message smart_reply1 [smart_reply2] [smart_reply3]} -``` - -For the current model, we see a triggering rate of about 30-40% for messages -which are similar to those provided in the tsv file above. - -In case the model does not trigger any response, the system falls back to -suggesting replies from a fixed back-off set that was compiled from popular -response intents observed in chat conversations. Some of the fallback responses -are `Ok, Yes, No, 👍, ☺`. - -The model can only be used for inference at this time (i.e. it cannot be custom -trained). If you are interested to know how the model was trained, please refer -to this [blog -post](https://research.googleblog.com/2017/02/on-device-machine-intelligence.html) -and [research paper](https://arxiv.org/pdf/1708.00630). - -## How to use this Model? - -We have provided a pre-built demo APK that you can download, install and test on -your phone ([demo APK -here](http://download.tensorflow.org/deps/tflite/SmartReplyDemo.apk)). - -The On-Device Smart Reply demo App works in the following way: - -1. Android app links to the JNI binary with a predictor library. - -2. In the predictor library, `GetSegmentPredictions` is called with a list of input - strings. - - 2.1 The input string can be 1-3 most recent messages of the conversations in - form of string vector. The model will run on these input sentences and - provide Smart Replies corresponding to them. - - 2.2 The function performs some preprocessing on input data which includes: - - * Sentence splitting: The input message will be split into sentences if - message has more than one sentence. Eg: a message like “How are you? - Want to grab lunch?” will be broken down into 2 different sentences. - * Normalization: The individual sentences will be normalized by converting - them into lower cases, removing unnecessary punctuations, etc. Eg: “how - are you????” will be converted to “how are you?” (refer for NORMALIZE op - for more details). - - The input string content will be converted to tensors. - - 2.3 The function then runs the prediction model on the input tensors. - - 2.4 The function also performs some post-processing which includes - aggregating the model predictions for the input sentences from 2.2 and - returning the appropriate responses. - -3. Finally, it gets response(s) from `std::vector`, and - returns back to Android app. Responses are sorted in descending order of - confidence score. - -## Ops and Functionality Supported - -Following are the ops supported for using On-Device Smart Reply model: - -* **NORMALIZE** - - This is a custom op which normalizes the sentences by: - - * Converting all sentences into lower case. - * Removing unnecessary punctuations (eg: “how are you????” → “how are - you?”). - * Expanding sentences wherever necessary (eg: “ I’m home” → “I am home”). - -* **SKIP_GRAM** - - This is an op inside TensorFlow Lite that converts sentences into a list of - skip grams. The configurable parameters are `ngram_size` and - `max_skip_size`. For the model provided, the values for these parameters are - set to 3 & 2 respectively. - -* **EXTRACT_FEATURES** - - This is a custom op that hashes skip grams to features represented as - integers. Longer skip-grams are allocated higher weights. - -* **LSH_PROJECTION** - - This is an op inside TensorFlow Lite that projects input features to a - corresponding bit vector space using Locality Sensitive Hashing (LSH). - -* **PREDICT** - - This is a custom op that runs the input features through the projection - model (details [here](https://arxiv.org/pdf/1708.00630.pdf)), computes the - appropriate response labels along with weights for the projected features, - and aggregates the response labels and weights together. - -* **HASHTABLE_LOOKUP** - - This is an op inside TensorFlow Lite that uses label id from predict op and - looks up the response text from the given label id. - -## Further Information - -* Open source code - [here](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite/models/smartreply/). diff --git a/tensorflow/contrib/lite/models/testdata/g3doc/README.md b/tensorflow/contrib/lite/models/testdata/g3doc/README.md deleted file mode 100644 index 1c47e00aae2a0e76ba04004a2fc3cc02ec4536f7..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/models/testdata/g3doc/README.md +++ /dev/null @@ -1,143 +0,0 @@ -## Speech Model Tests - -Sample test data has been provided for speech related models in Tensorflow Lite -to help users working with speech models to verify and test their models. - -For the hotword, speaker-id and automatic speech recognition sample models, the -architecture assumes that the models receive their input from a speech -pre-processing module. The speech pre-processing module receives the audio -signal and produces features for the encoder neural network and uses some -typical signal processing algorithms, like FFT and spectral subtraction, and -ultimately produces a log-mel filterbank (the log of the triangular mel filters -applied to the power spectra). The text-to-speech model assumes that the inputs -are linguistic features describing characteristics of phonemes, syllables, -words, phrases, and sentence. The outputs are acoustic features including -mel-cepstral coefficients, log fundamental frequency, and band aperiodicity. -The pre-processing modules for these models are not provided in the open source -version of TensorFlow Lite. - -The following sections describe the architecture of the sample models at a high -level: - -### Hotword Model - -The hotword model is the neural network model we use for keyphrase/hotword -spotting (i.e. "okgoogle" detection). It is the entry point for voice -interaction (e.g. Google search app on Android devices or Google Home, etc.). -The speech hotword model block diagram is shown in Figure below. It has an input -size of 40 (float), an output size of 7 (float), one Svdf layer, and four fully -connected layers with the corresponding parameters as shown in figure below. - -![hotword_model](hotword.svg "Hotword model") - -### Speaker-id Model - -The speaker-id model is the neural network model we use for speaker -verification. It runs after the hotword triggers. The speech speaker-id model -block diagram is shown in Figure below. It has an input size of 80 (float), an -output size of 64 (float), three Lstm layers, and one fully connected layers -with the corresponding parameters as shown in figure below. - -![speakerid_model](speakerid.svg "Speaker-id model") - -### Text-to-speech (TTS) Model - -The text-to-speech model is the neural network model used to generate speech -from text. The speech text-to-speech model’s block diagram is shown -in Figure below. It has and input size of 334 (float), an output size of 196 -(float), two fully connected layers, three Lstm layers, and one recurrent layer -with the corresponding parameters as shown in the figure. - -![tts_model](tts.svg "TTS model") - -### Automatic Speech Recognizer (ASR) Acoustic Model (AM) - -The acoustic model for automatic speech recognition is the neural network model -for matching phonemes to the input audio features. It generates posterior -probabilities of phonemes from speech frontend features (log-mel filterbanks). -It has an input size of 320 (float), an output size of 42 (float), five LSTM -layers and one fully connected layers with a Softmax activation function, with -the corresponding parameters as shown in the figure. - -![asr_am_model](asr_am.svg "ASR AM model") - -### Automatic Speech Recognizer (ASR) Language Model (LM) - -The language model for automatic speech recognition is the neural network model -for predicting the probability of a word given previous words in a sentence. -It generates posterior probabilities of the next word based from a sequence of -words. The words are encoded as indices in a fixed size dictionary. -The model has two inputs both of size one (integer): the current word index and -next word index, an output size of one (float): the log probability. It consists -of three embedding layer, three LSTM layers, followed by a multiplication, a -fully connected layers and an addition. -The corresponding parameters as shown in the figure. - -![asr_lm_model](asr_lm.svg "ASR LM model") - -### Endpointer Model - -The endpointer model is the neural network model for predicting end of speech -in an utterance. More precisely, it generates posterior probabilities of various -events that allow detection of speech start and end events. -It has an input size of 40 (float) which are speech frontend features -(log-mel filterbanks), and an output size of four corresponding to: -speech, intermediate non-speech, initial non-speech, and final non-speech. -The model consists of a convolutional layer, followed by a fully-connected -layer, two LSTM layers, and two additional fully-connected layers. -The corresponding parameters as shown in the figure. -![endpointer_model](endpointer.svg "Endpointer model") - - -## Speech models test input/output generation - -As mentioned above the input to models are generated from a pre-processing -module (output of a log-mel filterbank, or linguistic features), and the outputs -are generated by running the equivalent TensorFlow model by feeding them the -same input. - -## Link to the open source code - -### Models: - -[Speech hotword model (Svdf -rank=1)](https://storage.googleapis.com/download.tensorflow.org/models/tflite/speech_hotword_model_rank1_2017_11_14.tflite) - -[Speech hotword model (Svdf -rank=2)](https://storage.googleapis.com/download.tensorflow.org/models/tflite/speech_hotword_model_rank2_2017_11_14.tflite) - -[Speaker-id -model](https://storage.googleapis.com/download.tensorflow.org/models/tflite/speech_speakerid_model_2017_11_14.tflite) - -[TTS -model](https://storage.googleapis.com/download.tensorflow.org/models/tflite/speech_tts_model_2017_11_14.tflite) - -[ASR AM -model](https://storage.googleapis.com/download.tensorflow.org/models/tflite/speech_terse_am_model_2017_11_14.tflite) - -### Test benches - -[Speech hotword model -test](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite/models/speech_hotword_model_test.cc) - -[Speaker-id model -test](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite/models/speech_speakerid_model_test.cc) - -[TTS model -test](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite/models/speech_tts_model_test.cc) - -[ASR AM model -test](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite/models/speech_asr_am_model_test.cc) - -[ASR LM model -test](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite/models/speech_asr_lm_model_test.cc) - -[Endpointer model -test](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite/models/speech_endpointer_model_test.cc) - -## Android Support -The models have been tested on Android phones, using the following tests: - -[Hotword] (https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite/android/BUILD?rcl=172930882&l=25) - -[Speaker-id] (https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite/android/BUILD?rcl=172930882&l=36) diff --git a/tensorflow/contrib/lite/nnapi/BUILD b/tensorflow/contrib/lite/nnapi/BUILD deleted file mode 100644 index 467a2b7a7bc9a40135428240585cd2c2a133cf9f..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/nnapi/BUILD +++ /dev/null @@ -1,13 +0,0 @@ -licenses(["notice"]) # Apache 2.0 - -package(default_visibility = [ - "//visibility:public", -]) - -cc_library( - name = "nnapi_lib", - hdrs = [ - "NeuralNetworksShim.h", - ], - linkopts = ["-ldl"], -) diff --git a/tensorflow/contrib/lite/nnapi/NeuralNetworksShim.h b/tensorflow/contrib/lite/nnapi/NeuralNetworksShim.h deleted file mode 100644 index eccf4aefb6372b71c3b87dc0cdea24fec22ff625..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/nnapi/NeuralNetworksShim.h +++ /dev/null @@ -1,1012 +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_CONTRIB_LITE_NNAPI_NEURALNETWORKSSHIM_H_ -#define TENSORFLOW_CONTRIB_LITE_NNAPI_NEURALNETWORKSSHIM_H_ - -#include -#include -#include -#include - -// helpers - -#define NNAPI_LOG(format, ...) fprintf(stderr, format "\n", __VA_ARGS__); -#define LOAD_FUNCTION(name) \ - static name##_fn fn = reinterpret_cast(loadFunction(#name)); -#define EXECUTE_FUNCTION(...) \ - if (fn != nullptr) { \ - fn(__VA_ARGS__); \ - } -#define EXECUTE_FUNCTION_RETURN(...) return fn != nullptr ? fn(__VA_ARGS__) : 0; - -inline void* loadLibrary(const char* name) { - // TODO: change RTLD_LOCAL? Assumes there can be multiple instances of nn - // api RT - void* handle = nullptr; -#ifdef __ANDROID__ - handle = dlopen(name, RTLD_LAZY | RTLD_LOCAL); - if (handle == nullptr) { - NNAPI_LOG("nnapi error: unable to open library %s", name); - } -#endif - return handle; -} - -typedef int (*ASharedMemory_create_fn)(const char* name, size_t size); - -// ASharedMemory_create was added in Android 8.0, so safe to use with NNAPI -// which was added in 8.1. -inline int ASharedMemory_create(const char* name, size_t size) { - static void* handle = loadLibrary("libandroid.so"); - static ASharedMemory_create_fn fn = - handle != nullptr ? reinterpret_cast( - dlsym(handle, "ASharedMemory_create")) - : nullptr; - return fn(name, size); -} - -inline void* getLibraryHandle() { - static void* handle = loadLibrary("libneuralnetworks.so"); - return handle; -} - -inline void* loadFunction(const char* name) { - void* fn = nullptr; - if (getLibraryHandle() != nullptr) { - fn = dlsym(getLibraryHandle(), name); - } - if (fn == nullptr) { - NNAPI_LOG("nnapi error: unable to open function %s", name); - } - return fn; -} - -inline bool NNAPIExists() { - static bool nnapi_is_available = getLibraryHandle(); - return nnapi_is_available; -} - -// NN api types based on NNAPI header file -// https://developer.android.com/ndk/reference/group/neural-networks - -/** - * Operand types. - * - * The type of operands that can be added to a model. - * - * Although we define many types, most operators accept just a few - * types. Most used are ANEURALNETWORKS_TENSOR_FLOAT32, - * ANEURALNETWORKS_TENSOR_QUANT8_ASYMM, and ANEURALNETWORKS_INT32. - */ -enum { - ANEURALNETWORKS_FLOAT32 = 0, - ANEURALNETWORKS_INT32 = 1, - ANEURALNETWORKS_UINT32 = 2, - ANEURALNETWORKS_TENSOR_FLOAT32 = 3, - ANEURALNETWORKS_TENSOR_INT32 = 4, - ANEURALNETWORKS_TENSOR_QUANT8_ASYMM = 5, -}; - -/** - * Operation types. - * - * The type of operations that can be added to a model. - */ -enum { - ANEURALNETWORKS_ADD = 0, - ANEURALNETWORKS_AVERAGE_POOL_2D = 1, - ANEURALNETWORKS_CONCATENATION = 2, - ANEURALNETWORKS_CONV_2D = 3, - ANEURALNETWORKS_DEPTHWISE_CONV_2D = 4, - ANEURALNETWORKS_DEPTH_TO_SPACE = 5, - ANEURALNETWORKS_DEQUANTIZE = 6, - ANEURALNETWORKS_EMBEDDING_LOOKUP = 7, - ANEURALNETWORKS_FLOOR = 8, - ANEURALNETWORKS_FULLY_CONNECTED = 9, - ANEURALNETWORKS_HASHTABLE_LOOKUP = 10, - ANEURALNETWORKS_L2_NORMALIZATION = 11, - ANEURALNETWORKS_L2_POOL_2D = 12, - ANEURALNETWORKS_LOCAL_RESPONSE_NORMALIZATION = 13, - ANEURALNETWORKS_LOGISTIC = 14, - ANEURALNETWORKS_LSH_PROJECTION = 15, - ANEURALNETWORKS_LSTM = 16, - ANEURALNETWORKS_MAX_POOL_2D = 17, - ANEURALNETWORKS_MUL = 18, - ANEURALNETWORKS_RELU = 19, - ANEURALNETWORKS_RELU1 = 20, - ANEURALNETWORKS_RELU6 = 21, - ANEURALNETWORKS_RESHAPE = 22, - ANEURALNETWORKS_RESIZE_BILINEAR = 23, - ANEURALNETWORKS_RNN = 24, - ANEURALNETWORKS_SOFTMAX = 25, - ANEURALNETWORKS_SPACE_TO_DEPTH = 26, - ANEURALNETWORKS_SVDF = 27, - ANEURALNETWORKS_TANH = 28, - ANEURALNETWORKS_BATCH_TO_SPACE_ND = 29, - ANEURALNETWORKS_DIV = 30, - ANEURALNETWORKS_MEAN = 31, - ANEURALNETWORKS_PAD = 32, - ANEURALNETWORKS_SPACE_TO_BATCH_ND = 33, - ANEURALNETWORKS_SQUEEZE = 34, - ANEURALNETWORKS_STRIDED_SLICE = 35, - ANEURALNETWORKS_SUB = 36, - ANEURALNETWORKS_TRANSPOSE = 37, -}; - -/** - * Fused activation function types. - * - */ -enum { - ANEURALNETWORKS_FUSED_NONE = 0, - ANEURALNETWORKS_FUSED_RELU = 1, - ANEURALNETWORKS_FUSED_RELU1 = 2, - ANEURALNETWORKS_FUSED_RELU6 = 3, -}; - -/** - * Execution preferences. - */ -enum { - ANEURALNETWORKS_PREFER_LOW_POWER = 0, - ANEURALNETWORKS_PREFER_FAST_SINGLE_ANSWER = 1, - ANEURALNETWORKS_PREFER_SUSTAINED_SPEED = 2, -}; - -/** - * Result codes. - */ -enum { - ANEURALNETWORKS_NO_ERROR = 0, - ANEURALNETWORKS_OUT_OF_MEMORY = 1, - ANEURALNETWORKS_INCOMPLETE = 2, - ANEURALNETWORKS_UNEXPECTED_NULL = 3, - ANEURALNETWORKS_BAD_DATA = 4, - ANEURALNETWORKS_OP_FAILED = 5, - ANEURALNETWORKS_UNMAPPABLE = 5, - ANEURALNETWORKS_BAD_STATE = 6, -}; - -/** - * Implicit padding algorithms. - */ -enum { - ANEURALNETWORKS_PADDING_SAME = 1, - ANEURALNETWORKS_PADDING_VALID = 2, -}; - -/** - * ANeuralNetworksMemory is an opaque type that represents memory. - * - * This type is used to represent shared memory, memory mapped files, - * and similar memories. - * - * By using shared memory, a program can efficiently communicate to the - * runtime and drivers the tensors that define a model. See - * {@link ANeuralNetworksModel_setOperandValueFromMemory}. An application - * should typically create one shared memory object that contains every tensor - * needed to define a model. {@link ANeuralNetworksMemory_createFromFd} can be - * used to create shared memory from a file handle. {@link - * ANeuralNetworksMemory_createShared} can be used to directly created shared - * memory. - * - * Memory objects can also be used to specify the input and output arguments of - * an execution. See {@link ANeuralNetworksExecution_setInputFromMemory} - * and {@link ANeuralNetworksExecution_setOutputFromMemory}. - */ -typedef struct ANeuralNetworksMemory ANeuralNetworksMemory; - -/** - * ANeuralNetworksModel is an opaque type that contains a description of the - * mathematical operations that constitute the model. - * - *

The model will be built by calling

    - *
  • {@link ANeuralNetworksModel_create},
    • - *
  • {@link ANeuralNetworksModel_addOperation},
    • - *
  • {@link ANeuralNetworksModel_addOperand},
    • - *
- * - * A model is completed by calling {@link ANeuralNetworksModel_finish}. - * A model is destroyed by calling {@link ANeuralNetworksModel_free}. - * - *

It is the application's responsibility to make sure that only one thread - * modifies a model at a given time. It is however safe for more than one - * thread to use the model once {@link ANeuralNetworksModel_finish} has - * returned.

- * - *

It is also the application's responsibility to ensure that there are no - * other uses of the model after calling {@link ANeuralNetworksModel_free}. This - * includes any compilation or execution object created using the model.

- */ -typedef struct ANeuralNetworksModel ANeuralNetworksModel; - -/** - * ANeuralNetworksCompilation is an opaque type that can be used to compile - * a machine learning model. - * - *

To use:

    - *
  • Create a new compilation instance by calling the - * {@link ANeuralNetworksCompilation_create} function.
    • - *
  • Perform the compilation with {@link - * ANeuralNetworksCompilation_start}.
  • Wait for the compilation to - * complete with {@link ANeuralNetworksCompilation_wait}.
  • Use the - * compilation as many times as needed with {@link - * ANeuralNetworksExecution_create}.
  • Destroy the compilation with - * {@link ANeuralNetworksCompilation_free} once all executions using the - * compilation have completed.

- * - *

A compilation cannot be modified once {@link - * ANeuralNetworksCompilation_start} has been called on it.

- * - *

It is the application's responsibility to make sure that only one thread - * modifies a compilation at a given time. It is however safe for more than one - * thread to use {@link ANeuralNetworksCompilation_wait} at the same time. - * It is also safe for multiple threads to use a compilation object once - * {@link ANeuralNetworksCompilation_wait} has completed.

- * - *

It is also the application's responsibility to ensure that there are no - * other uses of the compilation after calling {@link - * ANeuralNetworksCompilation_free}. This includes any execution object created - * using the compilation.

- */ -typedef struct ANeuralNetworksCompilation ANeuralNetworksCompilation; - -/** - * ANeuralNetworksExecution is an opaque type that can be used to apply a - * machine learning model to a set of inputs. - * - *

To use:

    - *
  • Create a new execution instance by calling the - * {@link ANeuralNetworksExecution_create} function.
    • - *
  • Associate data to the model inputs with - * {@link ANeuralNetworksExecution_setInput} or - * {@link ANeuralNetworksExecution_setInputFromMemory}.
    • - *
  • Associate output buffers to the model outputs with - * {@link ANeuralNetworksExecution_setOutput} or - * {@link ANeuralNetworksExecution_setOutputFromMemory}.
    • - *
  • Apply the model with {@link - * ANeuralNetworksExecution_startCompute}.
  • Wait for the execution to - * complete with {@link ANeuralNetworksExecution_wait}.
  • Destroy the - * execution with - * {@link ANeuralNetworksExecution_free}.

- * - *

An execution cannot be modified once {@link - * ANeuralNetworksExecution_start} has been called on it.

- * - *

An execution can be applied to a model with - * {@link ANeuralNetworksExecution_startCompute} only once. Create new - * executions to do new evaluations of the model.

- * - *

It is the application's responsibility to make sure that only one thread - * modifies an execution at a given time. It is however safe for more than one - * thread to use {@link ANeuralNetworksExecution_wait} at the same time.

- * - *

It is also the application's responsibility to ensure that there are no - * other uses of the request after calling {@link - * ANeuralNetworksRequest_free}.

- */ -typedef struct ANeuralNetworksExecution ANeuralNetworksExecution; - -/** - * ANeuralNetworksOperandType describes the type of an operand. - * This structure is used to describe both scalars and tensors. - */ -typedef struct ANeuralNetworksOperandType { - /** The data type, e.g ANEURALNETWORKS_INT8. */ - int32_t type; - /** The number of dimensions. It should be 0 for scalars. */ - uint32_t dimensionCount; - /** The dimensions of the tensor. It should be nullptr for scalars. */ - const uint32_t* dimensions; - /** These two fields are only used for quantized tensors. - * They should be zero for scalars and non-fixed point tensors. - * The dequantized value of each entry is (value - offset) * scale. - */ - float scale; - int32_t zeroPoint; -} ANeuralNetworksOperandType; - -/** - * ANeuralNetworksEvent is an opaque type that represents an event - * that will be signaled once an execution completes. - */ -typedef struct ANeuralNetworksEvent ANeuralNetworksEvent; - -typedef int32_t ANeuralNetworksOperationType; - -// nn api function types - -typedef int (*ANeuralNetworksMemory_createFromFd_fn)( - size_t size, int protect, int fd, size_t offset, - ANeuralNetworksMemory** memory); - -typedef void (*ANeuralNetworksMemory_free_fn)(ANeuralNetworksMemory* memory); - -typedef int (*ANeuralNetworksModel_create_fn)(ANeuralNetworksModel** model); - -typedef int (*ANeuralNetworksModel_finish_fn)(ANeuralNetworksModel* model); - -typedef void (*ANeuralNetworksModel_free_fn)(ANeuralNetworksModel* model); - -typedef int (*ANeuralNetworksCompilation_create_fn)( - ANeuralNetworksModel* model, ANeuralNetworksCompilation** compilation); - -typedef void (*ANeuralNetworksCompilation_free_fn)( - ANeuralNetworksCompilation* compilation); - -typedef int (*ANeuralNetworksCompilation_setPreference_fn)( - ANeuralNetworksCompilation* compilation, int32_t preference); - -typedef int (*ANeuralNetworksCompilation_finish_fn)( - ANeuralNetworksCompilation* compilation); - -typedef int (*ANeuralNetworksModel_addOperand_fn)( - ANeuralNetworksModel* model, const ANeuralNetworksOperandType* type); - -typedef int (*ANeuralNetworksModel_setOperandValue_fn)( - ANeuralNetworksModel* model, int32_t index, const void* buffer, - size_t length); - -typedef int (*ANeuralNetworksModel_setOperandValueFromMemory_fn)( - ANeuralNetworksModel* model, int32_t index, - const ANeuralNetworksMemory* memory, size_t offset, size_t length); - -typedef int (*ANeuralNetworksModel_addOperation_fn)( - ANeuralNetworksModel* model, ANeuralNetworksOperationType type, - uint32_t inputCount, const uint32_t* inputs, uint32_t outputCount, - const uint32_t* outputs); - -typedef int (*ANeuralNetworksModel_identifyInputsAndOutputs_fn)( - ANeuralNetworksModel* model, uint32_t inputCount, const uint32_t* inputs, - uint32_t outputCount, const uint32_t* outputs); - -typedef int (*ANeuralNetworksModel_relaxComputationFloat32toFloat16_fn)( - ANeuralNetworksModel* model, bool allow); - -typedef int (*ANeuralNetworksExecution_create_fn)( - ANeuralNetworksCompilation* compilation, - ANeuralNetworksExecution** execution); - -typedef void (*ANeuralNetworksExecution_free_fn)( - ANeuralNetworksExecution* execution); - -typedef int (*ANeuralNetworksExecution_setInput_fn)( - ANeuralNetworksExecution* execution, int32_t index, - const ANeuralNetworksOperandType* type, const void* buffer, size_t length); - -typedef int (*ANeuralNetworksExecution_setInputFromMemory_fn)( - ANeuralNetworksExecution* execution, int32_t index, - const ANeuralNetworksOperandType* type, const ANeuralNetworksMemory* memory, - size_t offset, size_t length); - -typedef int (*ANeuralNetworksExecution_setOutput_fn)( - ANeuralNetworksExecution* execution, int32_t index, - const ANeuralNetworksOperandType* type, void* buffer, size_t length); - -typedef int (*ANeuralNetworksExecution_setOutputFromMemory_fn)( - ANeuralNetworksExecution* execution, int32_t index, - const ANeuralNetworksOperandType* type, const ANeuralNetworksMemory* memory, - size_t offset, size_t length); - -typedef int (*ANeuralNetworksExecution_startCompute_fn)( - ANeuralNetworksExecution* execution, ANeuralNetworksEvent** event); - -typedef int (*ANeuralNetworksEvent_wait_fn)(ANeuralNetworksEvent* event); - -typedef void (*ANeuralNetworksEvent_free_fn)(ANeuralNetworksEvent* event); - -/** - * Creates a shared memory object from a file descriptor. - * - * The shared memory is backed by a file descriptor via mmap. - * See {@link ANeuralNetworksMemory} for a description on how to use - * this shared memory. - * - * @param size The requested size in bytes. - * Must not be larger than the file size. - * @param prot The desired memory protection for the mapping. - * It is either PROT_NONE or the bitwise OR of one or - * more of the following flags: PROT_READ, PROT_WRITE. - * @param fd The requested file descriptor. - * The file descriptor has to be mmap-able. The file - * descriptor will be duplicated. - * @param offset The offset to the beginning of the file of the area to map. - * The offset has to be aligned to a page size. - * @param memory The memory object to be created. - * Set to NULL if unsuccessful. - * - * @return ANEURALNETWORKS_NO_ERROR if the request completed normally. - */ -inline int ANeuralNetworksMemory_createFromFd(size_t size, int protect, int fd, - size_t offset, - ANeuralNetworksMemory** memory) { - LOAD_FUNCTION(ANeuralNetworksMemory_createFromFd); - EXECUTE_FUNCTION_RETURN(size, protect, fd, offset, memory); -} - -/** - * Delete a memory object. - * - * Destroys the object used by the run time to keep track of the memory. - * This will free the underlying actual memory if no other code has open - * handles to this memory. - * - * @param memory The memory object to be freed. - */ -inline void ANeuralNetworksMemory_free(ANeuralNetworksMemory* memory) { - LOAD_FUNCTION(ANeuralNetworksMemory_free); - EXECUTE_FUNCTION(memory); -} - -/** - * Create an empty {@link ANeuralNetworksModel}. - * - *

This only creates the object. Computation is performed once - * {@link ANeuralNetworksExecution_startCompute} is invoked. - * - * The model should be constructed with calls to - * {@link ANeuralNetworksModel_addOperation} and - * {@link ANeuralNetworksModel_addOperand} - * - *

{@link ANeuralNetworksModel_finish} should be called once the model - * has been fully constructed.

- * - *

{@link ANeuralNetworksModel_free} should be called once the model - * is no longer needed.

- * - * @param model The {@link ANeuralNetworksModel} to be created. - * Set to NULL if unsuccessful. - * - * @return ANEURALNETWORKS_NO_ERROR if successful. - */ -inline int ANeuralNetworksModel_create(ANeuralNetworksModel** model) { - LOAD_FUNCTION(ANeuralNetworksModel_create); - EXECUTE_FUNCTION_RETURN(model); -} - -/** - * Destroy a model. - * - * The model need not have been finished by a call to - * {@link ANeuralNetworksModel_finish}. - * - * See {@link ANeuralNetworksModel} for information on multithreaded usage. - * - * @param model The model to be destroyed. Passing NULL is acceptable and - * results in no operation. - */ -inline void ANeuralNetworksModel_free(ANeuralNetworksModel* model) { - LOAD_FUNCTION(ANeuralNetworksModel_free); - EXECUTE_FUNCTION(model); -} - -/** - * Indicate that we have finished modifying a model. Required before - * calling {@link ANeuralNetworksCompilation_compile}. - * - * An application is responsible to make sure that no other thread uses - * the model at the same time. - * - * See {@link ANeuralNetworksModel} for information on multithreaded usage. - * - * @param model The model to be finished. - * - * @return ANEURALNETWORKS_NO_ERROR if successful. - */ -inline int ANeuralNetworksModel_finish(ANeuralNetworksModel* model) { - LOAD_FUNCTION(ANeuralNetworksModel_finish); - EXECUTE_FUNCTION_RETURN(model); -} - -/** - * Add an operand to a model. - * - * The order in which the operands are added is important. The first one added - * to a model will have the index value 0, the second 1, etc. These indexes are - * used as operand identifiers in {@link ANeuralNetworksModel_addOperation}, - * {@link ANeuralNetworksExecution_setInput}, - * {@link ANeuralNetworksExecution_setInputFromMemory}, - * {@link ANeuralNetworksExecution_setOutput}, - * {@link ANeuralNetworksExecution_setOutputFromMemory} and - * {@link ANeuralNetworksExecution_setOperandValue}. - * - * To build a model that can accommodate inputs of various sizes, as you may - * want to do for a CNN, set the size of the dimensions that will vary at run - * time to 0. If you do so, provide the full dimensions when calling - * {@link ANeuralNetworksExecution_setInput} or {@link - * ANeuralNetworksExecution_setInputFromMemory}. - * - * Attempting to modify a model once {@link ANeuralNetworksModel_finish} has - * been called will return an error. - * - * See {@link ANeuralNetworksModel} for information on multithreaded usage. - * - * @param model The model to be modified. - * @param type The {@link ANeuralNetworksOperandType} that describes the shape - * of the operand. - * - * @return ANEURALNETWORKS_NO_ERROR if successful. - */ -inline int ANeuralNetworksModel_addOperand( - ANeuralNetworksModel* model, const ANeuralNetworksOperandType* type) { - LOAD_FUNCTION(ANeuralNetworksModel_addOperand); - EXECUTE_FUNCTION_RETURN(model, type); -} - -/** - * Sets an operand to a constant value. - * - * For scalar values, the content of buffer is copied into the model. - * - * For tensor values, a pointer to the buffer is stored within the model. - * The application is responsible for not changing the content of this region - * until all executions using this model have completed. As the data may - * be copied during processing, modifying the data after this call yields - * undefined results. - * - * Attempting to modify a model once {@link ANeuralNetworksModel_finish} has - * been called will return an error. - * - * See {@link ANeuralNetworksModel} for information on multithreaded usage. - * - * @param model The model to be modified. - * @param index The index of the model operand we're setting. - * @param buffer A pointer to the data to use. - * @param length The size in bytes of the data value. - * - * @return ANEURALNETWORKS_NO_ERROR if successful. - */ -inline int ANeuralNetworksModel_setOperandValue(ANeuralNetworksModel* model, - int32_t index, - const void* buffer, - size_t length) { - LOAD_FUNCTION(ANeuralNetworksModel_setOperandValue); - EXECUTE_FUNCTION_RETURN(model, index, buffer, length); -} - -/** - * Sets an operand to a value stored in a memory object. - * - * The content of the memory is not copied. A reference to that memory is stored - * inside the model. The application is responsible for not changing the content - * of the memory region until all executions using this model have completed. - * As the data may be copied during processing, modifying the data after this - * call yields undefined results. - * - * Attempting to modify a model once {@link ANeuralNetworksModel_finish} has - * been called will return an error. - * - * See {@link ANeuralNetworksModel} for information on multithreaded usage. - * - * @param model The model to be modified. - * @param index The index of the model operand we're setting. - * @param buffer A pointer to the data to use. - * @param memory The memory containing the data. - * @param offset This specifies the location of the data within the memory. - * The offset is in bytes from the start of memory. - * @param length The size in bytes of the data value. - * - * @return ANEURALNETWORKS_NO_ERROR if successful. - */ -inline int ANeuralNetworksModel_setOperandValueFromMemory( - ANeuralNetworksModel* model, int32_t index, - const ANeuralNetworksMemory* memory, size_t offset, size_t length) { - LOAD_FUNCTION(ANeuralNetworksModel_setOperandValueFromMemory); - EXECUTE_FUNCTION_RETURN(model, index, memory, offset, length); -} - -/** - * Add an operation to a model. - * - * @param model The model to be modified. - * @param type The type of the operation. - * @param inputCount The number of entries in the inputs array. - * @param inputs An array of indexes identifying each operand. - * @param outputCount The number of entries in the outputs array. - * @param outputs An array of indexes identifying each operand. - * - * The operands specified by inputs and outputs must have been - * previously added by calls to {@link ANeuralNetworksModel_addOperand}. - * - * Attempting to modify a model once {@link ANeuralNetworksModel_finish} has - * been called will return an error. - * - * See {@link ANeuralNetworksModel} for information on multithreaded usage. - * - * @return ANEURALNETWORKS_NO_ERROR if successful. - */ -inline int ANeuralNetworksModel_addOperation(ANeuralNetworksModel* model, - ANeuralNetworksOperationType type, - uint32_t inputCount, - const uint32_t* inputs, - uint32_t outputCount, - const uint32_t* outputs) { - LOAD_FUNCTION(ANeuralNetworksModel_addOperation); - EXECUTE_FUNCTION_RETURN(model, type, inputCount, inputs, outputCount, - outputs); -} - -/** - * Specifies which operands will be the model's inputs and outputs. - * - * An operand cannot be used for both input and output. Doing so will - * return an error. - * - * @param model The model to be modified. - * @param inputCount The number of entries in the inputs array. - * @param inputs An array of indexes identifying the input operands. - * @param outputCount The number of entries in the outputs array. - * @param outputs An array of indexes identifying the output operands. - * - * The operands specified by inputs and outputs must have been - * previously added by calls to {@link ANeuralNetworksModel_addOperand}. - * - * Attempting to modify a model once {@link ANeuralNetworksModel_finish} has - * been called will return an error. - * - * See {@link ANeuralNetworksModel} for information on multithreaded usage. - * - */ -inline int ANeuralNetworksModel_identifyInputsAndOutputs( - ANeuralNetworksModel* model, uint32_t inputCount, const uint32_t* inputs, - uint32_t outputCount, const uint32_t* outputs) { - LOAD_FUNCTION(ANeuralNetworksModel_identifyInputsAndOutputs); - EXECUTE_FUNCTION_RETURN(model, inputCount, inputs, outputCount, outputs); -} - -/** - * Specifies whether {@link ANEURALNETWORKS_TENSOR_FLOAT32} is allowed to be - * calculated with range and/or precision as low as that of the IEEE 754 16-bit - * floating-point format. By default, {@link ANEURALNETWORKS_TENSOR_FLOAT32} - * must be calculated using at least the range and precision of the IEEE 754 - * 32-bit floating-point format. - * - * @param model The model to be modified. - * @param allow 'true' indicates {@link ANEURALNETWORKS_TENSOR_FLOAT32} may be - * calculated with range and/or precision as low as that of the - * IEEE 754 16-bit floating point format. 'false' indicates - * {@link ANEURALNETWORKS_TENSOR_FLOAT32} must be calculated using - * at least the range and precision of the IEEE 754 32-bit floating - * point format. - * - * Attempting to modify a model once {@link ANeuralNetworksModel_finish} has - * been called will return an error. - * - * Available since API level 28. - * - * See {@link ANeuralNetworksModel} for information on multithreaded usage. - */ -inline int ANeuralNetworksModel_relaxComputationFloat32toFloat16( - ANeuralNetworksModel* model, bool allow) { - LOAD_FUNCTION(ANeuralNetworksModel_relaxComputationFloat32toFloat16); - EXECUTE_FUNCTION_RETURN(model, allow); -} - -/** - * Create a {@link ANeuralNetworksCompilation} to compile the given model. - * This only creates the object. Compilation is only performed once - * {@link ANeuralNetworksCompilation_start} is invoked. - * - *

The provided model must outlive the compilation.

- * - * The model must already have been finished by a call to - * {@link ANeuralNetworksModel_finish}. - * - * See {@link ANeuralNetworksCompilation} for information on multithreaded - * usage. - * - * @param model The {@link ANeuralNetworksModel} to be compiled. - * @param compilation The newly created object or NULL if unsuccessful. - * - * @return ANEURALNETWORKS_NO_ERROR if successful, ANEURALNETWORKS_BAD_DATA - * if the model is invalid. - */ -inline int ANeuralNetworksCompilation_create( - ANeuralNetworksModel* model, ANeuralNetworksCompilation** compilation) { - LOAD_FUNCTION(ANeuralNetworksCompilation_create); - EXECUTE_FUNCTION_RETURN(model, compilation); -} - -/** - * Destroy a compilation. - * - *

If called on a compilation for which - * {@link ANeuralNetworksCompilation_start} has been called, the - * function will return immediately but will mark the compilation to be deleted - * once the compilation completes. The {@link ANeuralNetworksCompilation_wait} - * will return ERROR_DELETED. - * - * See {@link ANeuralNetworksCompilation} for information on multithreaded - * usage. - * - * @param compilation The compilation to be destroyed. Passing NULL is - * acceptable and results in no operation. - */ -inline void ANeuralNetworksCompilation_free( - ANeuralNetworksCompilation* compilation) { - LOAD_FUNCTION(ANeuralNetworksCompilation_free); - EXECUTE_FUNCTION(compilation); -} - -/** - * Sets the execution preference. - * - *

Provides guidance to the runtime when trade-offs are possible.

- * - * See {@link ANeuralNetworksCompilation} for information on multithreaded - * usage. - * - * @param compilation The compilation to be modified. - * @param preference Either {@link PREFER_LOW_POWER}, - * {@link PREFER_SINGLE_FAST_ANSWER}, or - * {@link PREFER_SUSTAINED_SPEED}. - * - * @return ANEURALNETWORKS_NO_ERROR if successful. - */ -inline int ANeuralNetworksCompilation_setPreference( - ANeuralNetworksCompilation* compilation, int32_t preference) { - LOAD_FUNCTION(ANeuralNetworksCompilation_setPreference); - EXECUTE_FUNCTION_RETURN(compilation, preference); -} - -/** - * Waits until the compilation completes. - * - * More than one thread can wait on a compilation. When the compilation - * completes, all threads will be released. - * - * See {@link ANeuralNetworksCompilation} for information on multithreaded - * usage. - * - * @return ANEURALNETWORKS_NO_ERROR if the compilation completed normally. - */ -inline int ANeuralNetworksCompilation_finish( - ANeuralNetworksCompilation* compilation) { - LOAD_FUNCTION(ANeuralNetworksCompilation_finish); - EXECUTE_FUNCTION_RETURN(compilation); -} -/** - * Create a {@link ANeuralNetworksExecution} to apply the given compilation. - * This only creates the object. Computation is only performed once - * {@link ANeuralNetworksExecution_startCompute} is invoked. - * - *

The provided compilation must outlive the execution.

- * - * See {@link ANeuralNetworksExecution} for information on multithreaded usage. - * - * @param compilation The {@link ANeuralNetworksCompilation} to be evaluated. - * @param execution The newly created object or NULL if unsuccessful. - * - * @return ANEURALNETWORKS_NO_ERROR if successful, ANEURALNETWORKS_BAD_DATA - * if the compilation is invalid. - */ -inline int ANeuralNetworksExecution_create( - ANeuralNetworksCompilation* compilation, - ANeuralNetworksExecution** execution) { - LOAD_FUNCTION(ANeuralNetworksExecution_create); - EXECUTE_FUNCTION_RETURN(compilation, execution); -} - -/** - * Destroy an execution. - * - *

If called on an execution for which - * {@link ANeuralNetworksExecution_startCompute} has been called, the - * function will return immediately but will mark the execution to be deleted - * once the computation completes. The {link ANeuralNetworksExecution_wait} - * will return ANEURALNETWORKS_ERROR_DELETED. - * - * See {@link ANeuralNetworksExecution} for information on multithreaded usage. - * - * @param execution The execution to be destroyed. Passing NULL is acceptable - * and results in no operation. - */ -inline void ANeuralNetworksExecution_free(ANeuralNetworksExecution* execution) { - LOAD_FUNCTION(ANeuralNetworksExecution_free); - EXECUTE_FUNCTION(execution); -} - -/** - * Associate a user buffer with an input of the model of the - * {@link ANeuralNetworksExecution}. - * - *

The provided buffer must outlive the execution.

- * - * See {@link ANeuralNetworksExecution} for information on multithreaded usage. - * - * @param execution The execution to be modified. - * @param index The index of the input argument we are setting. It is - * an index into the lists passed to - * {@link ANeuralNetworksModel_identifyInputsAndOutputs}. It is not - * the index associated with {@link - * ANeuralNetworksModel_addOperand}. - * @param type The type of the operand. This should be used to specify the - * dimensions that were set to 0 when the operand was added to the - * model. All other properties of the type must be the same as - * specified in the model. If the type is the same as specified - * when the model was built, NULL can be passed. - * @param buffer The buffer containing the data. - * @param length The length in bytes of the buffer. - * - * @return ANEURALNETWORKS_NO_ERROR if successful, ANEURALNETWORKS_BAD_DATA if - * the name is not recognized or the buffer is too small for the input. - */ -inline int ANeuralNetworksExecution_setInput( - ANeuralNetworksExecution* execution, int32_t index, - const ANeuralNetworksOperandType* type, const void* buffer, size_t length) { - LOAD_FUNCTION(ANeuralNetworksExecution_setInput); - EXECUTE_FUNCTION_RETURN(execution, index, type, buffer, length); -} - -/** - * Associate part of a memory object with an input of the model of the - * {@link ANeuralNetworksExecution}. - * - *

The provided memory must outlive the execution.

- * - * See {@link ANeuralNetworksExecution} for information on multithreaded usage. - * - * @param execution The execution to be modified. - * @param index The index of the input argument we are setting. It is - * an index into the lists passed to - * {@link ANeuralNetworksModel_identifyInputsAndOutputs}. It is not - * the index associated with {@link - * ANeuralNetworksModel_addOperand}. - * @param type The type of the operand. This can be used to specify the - * dimensions that were set to 0 when the operand was added to the - * model. All other values must be the same as specified in the - * model. If the type is the same as specified when the model - * was built, NULL can be passed. - * @param memory The memory containing the data. - * @param offset This specifies the location of the data within the memory. - * The offset is in bytes from the start of memory. - * @param length The size in bytes of the data value. - * - * @return ANEURALNETWORKS_NO_ERROR if successful, ANEURALNETWORKS_BAD_DATA if - * the name is not recognized or the buffer is too small for the input. - */ -inline int ANeuralNetworksExecution_setInputFromMemory( - ANeuralNetworksExecution* execution, int32_t index, - const ANeuralNetworksOperandType* type, const ANeuralNetworksMemory* memory, - size_t offset, size_t length) { - LOAD_FUNCTION(ANeuralNetworksExecution_setInputFromMemory); - EXECUTE_FUNCTION_RETURN(execution, index, type, memory, offset, length); -} - -/** - * Associate a user buffer with an output of the model of the - * {@link ANeuralNetworksExecution}. - * - *

The provided buffer must outlive the execution.

- * - * See {@link ANeuralNetworksExecution} for information on multithreaded usage. - * - * @param execution The execution to be modified. - * @param index The index of the output argument we are setting. It is - * an index into the lists passed to - * {@link ANeuralNetworksModel_identifyInputsAndOutputs}. It is not - * the index associated with {@link - * ANeuralNetworksModel_addOperand}. - * @param type The type of the operand. This can be used to specify the - * dimensions that were set to 0 when the operand was added to the - * model. All other values must be the same as specified in the - * model. If the type is the same as specified when the model - * was built, NULL can be passed. - * @param buffer The buffer where the data is to be written. - * @param length The length in bytes of the buffer. - * - * @return ANEURALNETWORKS_NO_ERROR if successful, ANEURALNETWORKS_BAD_DATA if - * the name is not recognized or the buffer is too small for the output. - */ -inline int ANeuralNetworksExecution_setOutput( - ANeuralNetworksExecution* execution, int32_t index, - const ANeuralNetworksOperandType* type, void* buffer, size_t length) { - LOAD_FUNCTION(ANeuralNetworksExecution_setOutput); - EXECUTE_FUNCTION_RETURN(execution, index, type, buffer, length); -} - -/** - * Associate part of a memory object with an output of the model of the - * {@link ANeuralNetworksExecution}. - * - *

The provided memory must outlive the execution.

- * - * See {@link ANeuralNetworksExecution} for information on multithreaded usage. - * - * @param execution The execution to be modified. - * @param index The index of the output argument we are setting. It is - * an index into the lists passed to - * {@link ANeuralNetworksModel_identifyInputsAndOutputs}. It is not - * the index associated with {@link - * ANeuralNetworksModel_addOperand}. - * @param type The type of the operand. This can be used to specify the - * dimensions that were set to 0 when the operand was added to the - * model. All other values must be the same as specified in the - * model. If the type is the same as specified when the model - * was built, NULL can be passed. - * @param memory The memory where the data is to be stored. - * @param offset This specifies the location of the data within the memory. - * The offset is in bytes from the start of memory. - * @param length The length in bytes of the data value. - * - * @return ANEURALNETWORKS_NO_ERROR if successful, ANEURALNETWORKS_BAD_DATA if - * the name is not recognized or the buffer is too small for the output. - */ -inline int ANeuralNetworksExecution_setOutputFromMemory( - ANeuralNetworksExecution* execution, int32_t index, - const ANeuralNetworksOperandType* type, const ANeuralNetworksMemory* memory, - size_t offset, size_t length) { - LOAD_FUNCTION(ANeuralNetworksExecution_setOutputFromMemory); - EXECUTE_FUNCTION_RETURN(execution, index, type, memory, offset, length); -} - -/** - * Schedule evaluation of the execution. - * - *

Schedules evaluation of the execution. Once the model has been - * applied and the outputs are ready to be consumed, the execution will be - * signaled. Use {@link ANeuralNetworksExecution_wait} to wait for that signal. - *

- * - * Multiple executions can be scheduled and evaluated concurrently, and - * compilations can be performed concurrently with executions. The runtime makes - * no guarantee on the ordering of the completion of compilations and - * executions. If it's important to the application, the application should - * enforce the ordering by using {@link ANeuralNetworksCompilation_wait} and - * {@link ANeuralNetworksExecution_wait}. - * - * ANeuralNetworksExecution_wait must be called to recuperate the resources used - * by the execution. - * - * See {@link ANeuralNetworksExecution} for information on multithreaded usage. - * - * @param execution The execution to be scheduled and executed. - * - * @return ANEURALNETWORKS_NO_ERROR if successful. - */ -inline int ANeuralNetworksExecution_startCompute( - ANeuralNetworksExecution* execution, ANeuralNetworksEvent** event) { - LOAD_FUNCTION(ANeuralNetworksExecution_startCompute); - EXECUTE_FUNCTION_RETURN(execution, event); -} - -/** - * Waits until the execution completes. - * - * More than one thread can wait on an event. When the execution completes, - * all threads will be released. - * - * See {@link ANeuralNetworksExecution} for information on multithreaded usage. - * - * @return ANEURALNETWORKS_NO_ERROR if the execution completed normally. - */ -inline int ANeuralNetworksEvent_wait(ANeuralNetworksEvent* event) { - LOAD_FUNCTION(ANeuralNetworksEvent_wait); - EXECUTE_FUNCTION_RETURN(event); -} - -/** - * Destroys the event. - * - * See {@link ANeuralNetworksExecution} for information on multithreaded usage. - */ -inline void ANeuralNetworksEvent_free(ANeuralNetworksEvent* event) { - LOAD_FUNCTION(ANeuralNetworksEvent_free); - EXECUTE_FUNCTION(event); -} - -/**/ - -#endif // TENSORFLOW_CONTRIB_LITE_NNAPI_NEURALNETWORKSSHIM_H_ diff --git a/tensorflow/contrib/lite/nnapi_delegate.cc b/tensorflow/contrib/lite/nnapi_delegate.cc deleted file mode 100644 index 31f233352000117d00356a9c323b0e933234b189..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/nnapi_delegate.cc +++ /dev/null @@ -1,857 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/contrib/lite/nnapi_delegate.h" -#include -#include -#include -#include -#include "tensorflow/contrib/lite/c/builtin_op_data.h" -#include "tensorflow/contrib/lite/core/api/error_reporter.h" -#include "tensorflow/contrib/lite/model.h" -#include "tensorflow/contrib/lite/nnapi/NeuralNetworksShim.h" - -#ifdef __ANDROID__ -#include -#include -#endif - -namespace tflite { - -void logError(const char* format, ...) { - // stderr is convenient for native tests, but is not captured for apps - va_list args_for_stderr; - va_start(args_for_stderr, format); - vfprintf(stderr, format, args_for_stderr); - va_end(args_for_stderr); - fprintf(stderr, "\n"); - fflush(stderr); -#ifdef __ANDROID__ - // produce logcat output for general consumption - va_list args_for_log; - va_start(args_for_log, format); - __android_log_vprint(ANDROID_LOG_ERROR, "tflite", format, args_for_log); - va_end(args_for_log); -#endif -} - -#define FATAL(...) \ - logError(__VA_ARGS__); \ - exit(1); - -// TODO(aselle): Change the error model to use status codes. -#define CHECK_TFLITE_SUCCESS(x) \ - if (x != kTfLiteOk) { \ - FATAL("Aborting since tflite returned failure nnapi_delegate.cc:%d.", \ - __LINE__); \ - } - -#define CHECK_NN(x) \ - if (x != ANEURALNETWORKS_NO_ERROR) { \ - FATAL("Aborting since NNAPI returned failure nnapi_delegate.cc:%d", \ - __LINE__); \ - } - -#define RETURN_ERROR_IF_TFLITE_FAILED(x) \ - if (x != kTfLiteOk) { \ - logError( \ - "Returning error since TFLite returned failure nnapi_delegate.cc:%d.", \ - __LINE__); \ - return kTfLiteError; \ - } - -#define RETURN_ERROR_IF_NN_FAILED(x) \ - if (x != ANEURALNETWORKS_NO_ERROR) { \ - logError( \ - "Returning error since NNAPI returned failure nnapi_delegate.cc:%d.", \ - __LINE__); \ - return kTfLiteError; \ - } - -// Tracking of NNAPI operand ids -static const int64_t kOperandIdNotSet = -1; -static const int64_t kOperandNotNeeded = -2; - -namespace { - -int32_t GetAndroidSdkVersion() { -#ifdef __ANDROID__ - const char* sdkProp = "ro.build.version.sdk"; - char sdkVersion[PROP_VALUE_MAX]; - int length = __system_property_get(sdkProp, sdkVersion); - if (length != 0) { - for (int i = 0; i < length; ++i) { - int digit = sdkVersion[i] - '0'; - if (digit < 0 || digit > 9) { - // Non-numeric SDK version, assume it's higher then expected; - return 0xFFFF; - } - } - return atoi(sdkVersion); - } - FATAL("No %s prop", sdkProp); -#endif // __ANDROID__ - return 0; -} - -int32_t GetAndroidSdkVersionCached() { - static int32_t androidSdkVersion = GetAndroidSdkVersion(); - return androidSdkVersion; -} - -} // namespace - -NNAPIAllocation::NNAPIAllocation(const char* filename, - ErrorReporter* error_reporter) - : MMAPAllocation(filename, error_reporter) { - if (mmapped_buffer_ != MAP_FAILED) - CHECK_NN(ANeuralNetworksMemory_createFromFd(buffer_size_bytes_, PROT_READ, - mmap_fd_, 0, &handle_)); -} - -NNAPIAllocation::~NNAPIAllocation() { - if (handle_) { - ANeuralNetworksMemory_free(handle_); - } -} - -NNAPIDelegate::~NNAPIDelegate() { - if (nn_compiled_model_) { - ANeuralNetworksCompilation_free(nn_compiled_model_); - nn_compiled_model_ = nullptr; - } - if (nn_model_) { - ANeuralNetworksModel_free(nn_model_); - nn_model_ = nullptr; - // TODO(aselle): Is this thread-safe and callable multiple times? - } - // ANeuralNetworksShutdown(); -} - -// Adds the tensors of the interpreter to the NN API model. -TfLiteStatus addTensorOperands(tflite::Interpreter* interpreter, - ANeuralNetworksModel* nn_model, - uint32_t* no_of_operands_added, - std::vector* nnapi_ids) { - uint32_t next_id = 0; - for (size_t i = 0; i < interpreter->tensors_size(); i++) { - // Skip temporaries and RNN back-edges. - if ((*nnapi_ids)[i] == kOperandNotNeeded) continue; - - (*nnapi_ids)[i] = int64_t(next_id); - - int32_t nn_type = 0; - // NNAPI requires 32-bit float scale to be zero, tflite doesn't care - float scale = 0.0f; - int32_t zeroPoint = 0; - TfLiteTensor* tensor = interpreter->tensor(i); - switch (tensor->type) { - case kTfLiteNoType: - // Tensors added during initialization of Ops don't have a type yet and - // should not be registered with the NNAPI. - continue; - case kTfLiteFloat32: - nn_type = ANEURALNETWORKS_TENSOR_FLOAT32; - break; - case kTfLiteUInt8: - nn_type = ANEURALNETWORKS_TENSOR_QUANT8_ASYMM; - scale = tensor->params.scale; - zeroPoint = tensor->params.zero_point; - break; - case kTfLiteInt32: - nn_type = ANEURALNETWORKS_TENSOR_INT32; - scale = tensor->params.scale; - zeroPoint = tensor->params.zero_point; - break; - default: - logError("Unsupported tensor type %d", tensor->type); - return kTfLiteError; - } - if (tensor->dims->size == 0) { - logError("NNAPI doesn't support tensors with rank 0 (index %d name %s)", - i, tensor->name); - return kTfLiteError; - } - if (tensor->dims->size > 4) { - logError("NNAPI doesn't support tensors with rank > 4 (index %d name %s)", - i, tensor->name); - return kTfLiteError; - } - // TODO(aselle): Note, many of these are intermediate results. Do I need - // to ever specify these sizes. I am currently below doing setValue - // on all of them, but I shouldn't in the future. - // Answer(jeanluc): If all the operators can set the dimension correctly, - // you won't need to. - ANeuralNetworksOperandType operand_type{ - nn_type, static_cast(tensor->dims->size), - reinterpret_cast(tensor->dims->data), scale, zeroPoint}; - RETURN_ERROR_IF_NN_FAILED( - ANeuralNetworksModel_addOperand(nn_model, &operand_type)); - // TODO(aselle): Based on Michael's suggestion, limiting this to read - // only memory - if (tensor->allocation_type == kTfLiteMmapRo) { - if (const NNAPIAllocation* alloc = dynamic_cast( - static_cast(tensor->allocation))) { - RETURN_ERROR_IF_NN_FAILED( - ANeuralNetworksModel_setOperandValueFromMemory( - nn_model, next_id, alloc->memory(), - alloc->offset(tensor->data.raw), tensor->bytes)); - } else { - RETURN_ERROR_IF_NN_FAILED(ANeuralNetworksModel_setOperandValue( - nn_model, next_id, tensor->data.raw, tensor->bytes)); - } - } else if (tensor->bytes == 0) { - // These size 0 tensors are optional tensors reserved. - RETURN_ERROR_IF_NN_FAILED( - ANeuralNetworksModel_setOperandValue(nn_model, next_id, nullptr, 0)); - } - - ++next_id; - } - *no_of_operands_added = next_id; - return kTfLiteOk; -} - -void MapAndAddTensorIds(const int* from_ids_buf, size_t from_ids_count, - std::vector* into, - const std::vector& map) { - for (size_t i = 0; i < from_ids_count; i++) { - int from_id = from_ids_buf[i]; - if (from_id == kOptionalTensor) { - into->push_back(from_id); - } else { - into->push_back(map[from_id]); - } - } -} - -// Adds the operations and their parameters to the NN API model. -// 'next-id' is the operand ID of the next operand of the model. -TfLiteStatus AddOpsAndParams( - tflite::Interpreter* interpreter, ANeuralNetworksModel* nn_model, - uint32_t next_id, std::vector* model_state_inputs, - std::vector* model_state_outputs, - const std::vector& tensor_id_to_nnapi_id) { - for (size_t i = 0; i < interpreter->nodes_size(); i++) { - const auto* node_and_registration = interpreter->node_and_registration(i); - const TfLiteNode& node = node_and_registration->first; - const TfLiteRegistration& registration = node_and_registration->second; - tflite::BuiltinOperator builtin = - static_cast(registration.builtin_code); - - // Add the parameters. - std::vector augmented_inputs, augmented_outputs; - MapAndAddTensorIds(node.inputs->data, node.inputs->size, &augmented_inputs, - tensor_id_to_nnapi_id); - MapAndAddTensorIds(node.outputs->data, node.outputs->size, - &augmented_outputs, tensor_id_to_nnapi_id); - - auto add_scalar_int32 = [&nn_model, &augmented_inputs, - &next_id](int value) { - ANeuralNetworksOperandType operand_type{.type = ANEURALNETWORKS_INT32}; - CHECK_NN(ANeuralNetworksModel_addOperand(nn_model, &operand_type)) - CHECK_NN(ANeuralNetworksModel_setOperandValue(nn_model, next_id, &value, - sizeof(int32_t))) - augmented_inputs.push_back(next_id++); - }; - - auto add_scalar_float32 = [&nn_model, &augmented_inputs, - &next_id](float value) { - ANeuralNetworksOperandType operand_type{.type = ANEURALNETWORKS_FLOAT32}; - CHECK_NN(ANeuralNetworksModel_addOperand(nn_model, &operand_type)) - CHECK_NN(ANeuralNetworksModel_setOperandValue(nn_model, next_id, &value, - sizeof(float))) - augmented_inputs.push_back(next_id++); - }; - - auto add_vector_int32 = [&](const int* values, uint32_t num_values) { - ANeuralNetworksOperandType operand_type{ - .type = ANEURALNETWORKS_TENSOR_INT32, - .dimensionCount = 1, - .dimensions = &num_values}; - CHECK_NN(ANeuralNetworksModel_addOperand(nn_model, &operand_type)) - CHECK_NN(ANeuralNetworksModel_setOperandValue( - nn_model, next_id, values, sizeof(int32_t) * num_values)); - augmented_inputs.push_back(next_id++); - }; - - // Handle state tensors of RNN, LSTM, SVDF. - // For each state_out tensor, a corresponding state_in operand needs to be - // created for NNAPI. - auto duplicate_state_tensor_float32 = - [interpreter, &nn_model, &next_id, &augmented_inputs, - &model_state_inputs, &model_state_outputs](int tensor_id) { - const TfLiteTensor* tensor = interpreter->tensor(tensor_id); - ANeuralNetworksOperandType operand_type{ - ANEURALNETWORKS_TENSOR_FLOAT32, - static_cast(tensor->dims->size), - reinterpret_cast(tensor->dims->data), - tensor->params.scale, tensor->params.zero_point}; - CHECK_NN(ANeuralNetworksModel_addOperand(nn_model, &operand_type)); - augmented_inputs.push_back(next_id); - model_state_inputs->push_back(next_id); - model_state_outputs->push_back(tensor_id); - next_id++; - }; - auto check_and_add_activation = [&add_scalar_int32](int activation) { - if (activation > kTfLiteActRelu6) { - logError("NNAPI only supports RELU, RELU1 and RELU6 activations"); - return kTfLiteError; - } - add_scalar_int32(activation); - return kTfLiteOk; - }; - - auto add_add_params = [&add_scalar_int32](void* data) { - auto* builtin = reinterpret_cast(data); - if (builtin->activation > kTfLiteActRelu6) { - logError("NNAPI only supports RELU, RELU1 and RELU6 activations"); - return kTfLiteError; - } - add_scalar_int32(builtin->activation); - return kTfLiteOk; - }; - - auto add_pooling_params = [&add_scalar_int32, - &check_and_add_activation](void* data) { - auto builtin = reinterpret_cast(data); - add_scalar_int32(builtin->padding); - add_scalar_int32(builtin->stride_width); - add_scalar_int32(builtin->stride_height); - add_scalar_int32(builtin->filter_width); - add_scalar_int32(builtin->filter_height); - return check_and_add_activation(builtin->activation); - }; - - auto add_convolution_params = [&add_scalar_int32, - &check_and_add_activation](void* data) { - auto builtin = reinterpret_cast(data); - add_scalar_int32(builtin->padding); - add_scalar_int32(builtin->stride_width); - add_scalar_int32(builtin->stride_height); - return check_and_add_activation(builtin->activation); - }; - - auto add_depthwise_conv_params = [&add_scalar_int32, - &check_and_add_activation](void* data) { - auto builtin = reinterpret_cast(data); - add_scalar_int32(builtin->padding); - add_scalar_int32(builtin->stride_width); - add_scalar_int32(builtin->stride_height); - add_scalar_int32(builtin->depth_multiplier); - return check_and_add_activation(builtin->activation); - }; - - auto add_fully_connected_params = [&check_and_add_activation](void* data) { - auto builtin = reinterpret_cast(data); - return check_and_add_activation(builtin->activation); - }; - - auto add_concatenation_params = [&add_scalar_int32](void* data) { - auto builtin = reinterpret_cast(data); - add_scalar_int32(builtin->axis); - if (builtin->activation != kTfLiteActNone) { - logError("Concatenation does not support fused activation in NNAPI"); - return kTfLiteError; - } - return kTfLiteOk; - }; - - auto add_softmax_params = [&add_scalar_float32](void* data) { - auto builtin = reinterpret_cast(data); - add_scalar_float32(builtin->beta); - }; - - auto add_space_to_depth_params = [&add_scalar_int32](void* data) { - auto builtin = reinterpret_cast(data); - add_scalar_int32(builtin->block_size); - }; - - auto add_lstm_params = [&add_scalar_int32, - &add_scalar_float32](void* data) { - auto builtin = reinterpret_cast(data); - add_scalar_int32(builtin->activation); - add_scalar_float32(builtin->cell_clip); - add_scalar_float32(builtin->proj_clip); - }; - - // LSTM in NNAPI requires scratch tensor as an output operand. - auto add_lstm_scratch_tensor_float32 = [interpreter, &node, &nn_model, - &next_id, &augmented_outputs]() { - if (node.temporaries->size == 0) return; - int scratch_buffer_index = node.temporaries->data[0]; - const TfLiteTensor* tensor = interpreter->tensor(scratch_buffer_index); - ANeuralNetworksOperandType operand_type{ - ANEURALNETWORKS_TENSOR_FLOAT32, - static_cast(tensor->dims->size), - reinterpret_cast(tensor->dims->data), tensor->params.scale, - tensor->params.zero_point}; - CHECK_NN(ANeuralNetworksModel_addOperand(nn_model, &operand_type)); - augmented_outputs.insert(augmented_outputs.begin(), next_id++); - }; - - auto add_mean_params = [&add_scalar_int32](void* data) { - auto builtin = reinterpret_cast(data); - add_scalar_int32(builtin->keep_dims); - }; - - auto add_svdf_params = [&add_scalar_int32](void* data) { - auto builtin = reinterpret_cast(data); - add_scalar_int32(builtin->rank); - add_scalar_int32(builtin->activation); - }; - - auto add_rnn_params = [&add_scalar_int32](void* data) { - auto builtin = reinterpret_cast(data); - add_scalar_int32(builtin->activation); - }; - - auto add_squeeze_params = [&](void* data) { - const auto* builtin = reinterpret_cast(data); - // Note that we add the squeeze dimensions even if the dimensions were - // unspecified (empty), as NNAPI requires the operand. - add_vector_int32(builtin->squeeze_dims, - static_cast(builtin->num_squeeze_dims)); - }; - - // Handle optional input tensors. - auto add_optional_tensors = [&nn_model, &augmented_inputs, - &next_id](int nn_type) { - for (size_t idx = 0; idx < augmented_inputs.size(); idx++) { - if (augmented_inputs[idx] == kOptionalTensor) { - const std::vector dim = {0, 0}; - ANeuralNetworksOperandType operand_type{nn_type, 2, dim.data(), 0, 0}; - CHECK_NN(ANeuralNetworksModel_addOperand(nn_model, &operand_type)) - CHECK_NN(ANeuralNetworksModel_setOperandValue(nn_model, next_id, - nullptr, 0)) - augmented_inputs[idx] = next_id++; - } - } - }; - - int nnapi_version = 10; - ANeuralNetworksOperationType nn_op_type; - - switch (builtin) { - case tflite::BuiltinOperator_ADD: - nn_op_type = ANEURALNETWORKS_ADD; - RETURN_ERROR_IF_TFLITE_FAILED(add_add_params(node.builtin_data)); - break; - case tflite::BuiltinOperator_MUL: - nn_op_type = ANEURALNETWORKS_MUL; - RETURN_ERROR_IF_TFLITE_FAILED(add_add_params(node.builtin_data)); - break; - case tflite::BuiltinOperator_AVERAGE_POOL_2D: - RETURN_ERROR_IF_TFLITE_FAILED(add_pooling_params(node.builtin_data)); - nn_op_type = ANEURALNETWORKS_AVERAGE_POOL_2D; - break; - case tflite::BuiltinOperator_MAX_POOL_2D: - RETURN_ERROR_IF_TFLITE_FAILED(add_pooling_params(node.builtin_data)); - nn_op_type = ANEURALNETWORKS_MAX_POOL_2D; - break; - case tflite::BuiltinOperator_L2_POOL_2D: - RETURN_ERROR_IF_TFLITE_FAILED(add_pooling_params(node.builtin_data)); - nn_op_type = ANEURALNETWORKS_L2_POOL_2D; - break; - case tflite::BuiltinOperator_CONV_2D: { - auto builtin = reinterpret_cast(node.builtin_data); - if (builtin->dilation_width_factor != 1 || - builtin->dilation_height_factor != 1 || node.inputs->size != 3) { - logError("NNAPI does not support dilated Conv2D."); - return kTfLiteError; - } - } - RETURN_ERROR_IF_TFLITE_FAILED( - add_convolution_params(node.builtin_data)); - nn_op_type = ANEURALNETWORKS_CONV_2D; - break; - case tflite::BuiltinOperator_RELU: - nn_op_type = ANEURALNETWORKS_RELU; - break; - case tflite::BuiltinOperator_RELU6: - nn_op_type = ANEURALNETWORKS_RELU6; - break; - case tflite::BuiltinOperator_TANH: - nn_op_type = ANEURALNETWORKS_TANH; - break; - case tflite::BuiltinOperator_FLOOR: - nn_op_type = ANEURALNETWORKS_FLOOR; - break; - case tflite::BuiltinOperator_LOGISTIC: - nn_op_type = ANEURALNETWORKS_LOGISTIC; - break; - case tflite::BuiltinOperator_DEPTHWISE_CONV_2D: - RETURN_ERROR_IF_TFLITE_FAILED( - add_depthwise_conv_params(node.builtin_data)); - nn_op_type = ANEURALNETWORKS_DEPTHWISE_CONV_2D; - break; - case tflite::BuiltinOperator_CONCATENATION: - RETURN_ERROR_IF_TFLITE_FAILED( - add_concatenation_params(node.builtin_data)); - nn_op_type = ANEURALNETWORKS_CONCATENATION; - break; - case tflite::BuiltinOperator_SOFTMAX: - add_softmax_params(node.builtin_data); - nn_op_type = ANEURALNETWORKS_SOFTMAX; - break; - case tflite::BuiltinOperator_FULLY_CONNECTED: - RETURN_ERROR_IF_TFLITE_FAILED( - add_fully_connected_params(node.builtin_data)); - nn_op_type = ANEURALNETWORKS_FULLY_CONNECTED; - break; - case tflite::BuiltinOperator_RESHAPE: - if (node.inputs->size != 2) { - logError("NNAPI only supports 2-input RESHAPE"); - return kTfLiteError; - } - nn_op_type = ANEURALNETWORKS_RESHAPE; - // add_reshape_params(node.builtin_data); - break; - case tflite::BuiltinOperator_SPACE_TO_DEPTH: - add_space_to_depth_params(node.builtin_data); - nn_op_type = ANEURALNETWORKS_SPACE_TO_DEPTH; - break; - case tflite::BuiltinOperator_LSTM: { - if (node.inputs->size + /* no of params */ 3 != 21) { - logError("NNAPI only supports 21-input LSTMs"); - return kTfLiteError; - } - duplicate_state_tensor_float32( - node.outputs->data[/*kOutputStateTensor*/ 0]); - duplicate_state_tensor_float32( - node.outputs->data[/*kCellStateTensor*/ 1]); - add_lstm_params(node.builtin_data); - add_lstm_scratch_tensor_float32(); - add_optional_tensors(ANEURALNETWORKS_TENSOR_FLOAT32); - nn_op_type = ANEURALNETWORKS_LSTM; - break; - } - case tflite::BuiltinOperator_SVDF: { - duplicate_state_tensor_float32(node.outputs->data[/*kStateTensor*/ 0]); - add_svdf_params(node.builtin_data); - nn_op_type = ANEURALNETWORKS_SVDF; - break; - } - case tflite::BuiltinOperator_RNN: { - duplicate_state_tensor_float32( - node.outputs->data[/*kHiddenStateTensor*/ 0]); - add_rnn_params(node.builtin_data); - nn_op_type = ANEURALNETWORKS_RNN; - break; - } - case tflite::BuiltinOperator_EMBEDDING_LOOKUP: - nn_op_type = ANEURALNETWORKS_EMBEDDING_LOOKUP; - break; - case tflite::BuiltinOperator_PAD: - nnapi_version = 11; // require NNAPI 1.1 - nn_op_type = ANEURALNETWORKS_PAD; - break; - case tflite::BuiltinOperator_MEAN: - nnapi_version = 11; // require NNAPI 1.1 - add_mean_params(node.builtin_data); - nn_op_type = ANEURALNETWORKS_MEAN; - break; - case tflite::BuiltinOperator_DIV: - nnapi_version = 11; // require NNAPI 1.1 - nn_op_type = ANEURALNETWORKS_DIV; - RETURN_ERROR_IF_TFLITE_FAILED(check_and_add_activation( - reinterpret_cast(node.builtin_data)->activation)); - break; - case tflite::BuiltinOperator_SUB: - nnapi_version = 11; // require NNAPI 1.1 - nn_op_type = ANEURALNETWORKS_SUB; - RETURN_ERROR_IF_TFLITE_FAILED(check_and_add_activation( - reinterpret_cast(node.builtin_data)->activation)); - break; - case tflite::BuiltinOperator_SQUEEZE: - nnapi_version = 11; // requires NNAPI 1.1 - add_squeeze_params(node.builtin_data); - nn_op_type = ANEURALNETWORKS_SQUEEZE; - break; - case tflite::BuiltinOperator_TRANSPOSE: - // The permutation input tensor value dictates the output dimensions. - // TODO(b/110888333): Support dynamically-sized tensors in delegates. - if ((node.inputs->size > 1) && - (interpreter->tensor(node.inputs->data[1])->allocation_type != - kTfLiteMmapRo)) { - logError("NNAPI does not yet support dynamic tensors."); - return kTfLiteError; - } - nnapi_version = 11; // require NNAPI 1.1 - nn_op_type = ANEURALNETWORKS_TRANSPOSE; - break; - case tflite::BuiltinOperator_L2_NORMALIZATION: - nn_op_type = ANEURALNETWORKS_L2_NORMALIZATION; - if (reinterpret_cast(node.builtin_data) - ->activation != kTfLiteActNone) { - logError( - "NNAPI does not support L2Normalization with fused activations"); - return kTfLiteError; - } - if ((node.inputs->size > 0) && - (interpreter->tensor(node.inputs->data[0])->dims->size != 4)) { - logError("NNAPI only supports input rank 4 for L2Normalization"); - return kTfLiteError; - } - break; - case tflite::BuiltinOperator_HASHTABLE_LOOKUP: - if (interpreter->tensor(node.outputs->data[0])->type != - kTfLiteFloat32) { - logError("NNAPI only support HASHTABLE_LOOKUP with float32 output", - builtin); - return kTfLiteError; - } - nn_op_type = ANEURALNETWORKS_HASHTABLE_LOOKUP; - break; - case tflite::BuiltinOperator_CONCAT_EMBEDDINGS: - case tflite::BuiltinOperator_LSH_PROJECTION: - case tflite::BuiltinOperator_BIDIRECTIONAL_SEQUENCE_RNN: - case tflite::BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_RNN: - case tflite::BuiltinOperator_EMBEDDING_LOOKUP_SPARSE: - case tflite::BuiltinOperator_BIDIRECTIONAL_SEQUENCE_LSTM: - case tflite::BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_LSTM: - case tflite::BuiltinOperator_LOCAL_RESPONSE_NORMALIZATION: - case tflite::BuiltinOperator_PADV2: - case tflite::BuiltinOperator_RESIZE_BILINEAR: - case tflite::BuiltinOperator_CALL: - case tflite::BuiltinOperator_SKIP_GRAM: - case tflite::BuiltinOperator_RELU_N1_TO_1: - case tflite::BuiltinOperator_GATHER: - case tflite::BuiltinOperator_SPACE_TO_BATCH_ND: - case tflite::BuiltinOperator_BATCH_TO_SPACE_ND: - case tflite::BuiltinOperator_TOPK_V2: - case tflite::BuiltinOperator_SPLIT: - case tflite::BuiltinOperator_STRIDED_SLICE: - case tflite::BuiltinOperator_EXP: - case tflite::BuiltinOperator_LOG_SOFTMAX: - case tflite::BuiltinOperator_DEQUANTIZE: - case tflite::BuiltinOperator_DELEGATE: - case tflite::BuiltinOperator_CAST: - case tflite::BuiltinOperator_PRELU: - case tflite::BuiltinOperator_MAXIMUM: - case tflite::BuiltinOperator_MINIMUM: - case tflite::BuiltinOperator_ARG_MAX: - case tflite::BuiltinOperator_ARG_MIN: - case tflite::BuiltinOperator_GREATER: - case tflite::BuiltinOperator_GREATER_EQUAL: - case tflite::BuiltinOperator_LESS: - case tflite::BuiltinOperator_LESS_EQUAL: - case tflite::BuiltinOperator_NEG: - case tflite::BuiltinOperator_SELECT: - case tflite::BuiltinOperator_SLICE: - case tflite::BuiltinOperator_SIN: - case tflite::BuiltinOperator_LOG: - case tflite::BuiltinOperator_TRANSPOSE_CONV: - case tflite::BuiltinOperator_TILE: - case tflite::BuiltinOperator_EXPAND_DIMS: - case tflite::BuiltinOperator_SPARSE_TO_DENSE: - case tflite::BuiltinOperator_EQUAL: - case tflite::BuiltinOperator_NOT_EQUAL: - case tflite::BuiltinOperator_SUM: - case tflite::BuiltinOperator_REDUCE_MAX: - case tflite::BuiltinOperator_REDUCE_MIN: - case tflite::BuiltinOperator_REDUCE_PROD: - case tflite::BuiltinOperator_SQRT: - case tflite::BuiltinOperator_RSQRT: - case tflite::BuiltinOperator_SHAPE: - case tflite::BuiltinOperator_POW: - case tflite::BuiltinOperator_FAKE_QUANT: - case tflite::BuiltinOperator_PACK: - case tflite::BuiltinOperator_LOGICAL_OR: - case tflite::BuiltinOperator_ONE_HOT: - case tflite::BuiltinOperator_LOGICAL_AND: - case tflite::BuiltinOperator_LOGICAL_NOT: - case tflite::BuiltinOperator_UNPACK: - case tflite::BuiltinOperator_FLOOR_DIV: - case tflite::BuiltinOperator_REDUCE_ANY: - case tflite::BuiltinOperator_SQUARE: - case tflite::BuiltinOperator_ZEROS_LIKE: - case tflite::BuiltinOperator_FILL: - case tflite::BuiltinOperator_FLOOR_MOD: - case tflite::BuiltinOperator_RANGE: - logError("Op code %d is currently not delegated to NNAPI", builtin); - return kTfLiteError; - break; - case tflite::BuiltinOperator_CUSTOM: - logError("Custom operations are not supported when using NNAPI."); - return kTfLiteError; - break; - } - - if (nnapi_version == 11 && GetAndroidSdkVersionCached() < 28) { - logError("Op %d needs NNAPI1.1", builtin); - return kTfLiteError; - } - - // Add the operation. - RETURN_ERROR_IF_NN_FAILED(ANeuralNetworksModel_addOperation( - nn_model, nn_op_type, static_cast(augmented_inputs.size()), - augmented_inputs.data(), - static_cast(augmented_outputs.size()), - reinterpret_cast(augmented_outputs.data()))); - } - return kTfLiteOk; -} - -TfLiteStatus NNAPIDelegate::BuildGraph(Interpreter* interpreter) { - if (nn_model_ && nn_compiled_model_) return model_status_; - - // TODO(aselle): This is not correct. need to handle resize invalidation. - if (!nn_model_) { - CHECK_NN(ANeuralNetworksModel_create(&nn_model_)); - - // Find which tensors should be added to NNAPI. TFLite has temporaries - // and RNN back-edges which are are not valid for NNAPI. We look through all - // inputs and outputs and mark the mapping in tensor_id_to_nnapi_id with - // kOperandIdNotSet. addTensorOperands will replace those with the - // corresponding NNAPI operand ids and skip kOperandNotNeeded entries. - std::vector tensor_id_to_nnapi_id(interpreter->tensors_size(), - kOperandNotNeeded); - auto set_ids_to_not_set = [&tensor_id_to_nnapi_id](const int* buf, - size_t count) { - for (int j = 0; j < count; j++) { - auto tensor_id = buf[j]; - if (tensor_id != kOptionalTensor) { - tensor_id_to_nnapi_id[tensor_id] = kOperandIdNotSet; - } - } - }; - for (size_t i = 0; i < interpreter->nodes_size(); i++) { - const auto* node_and_registration = interpreter->node_and_registration(i); - const TfLiteNode& node = node_and_registration->first; - set_ids_to_not_set(node.inputs->data, node.inputs->size); - set_ids_to_not_set(node.outputs->data, node.outputs->size); - } - set_ids_to_not_set(interpreter->inputs().data(), - interpreter->inputs().size()); - set_ids_to_not_set(interpreter->outputs().data(), - interpreter->outputs().size()); - - uint32_t next_id = 0; - RETURN_ERROR_IF_TFLITE_FAILED(addTensorOperands( - interpreter, nn_model_, &next_id, &tensor_id_to_nnapi_id)); - RETURN_ERROR_IF_TFLITE_FAILED( - AddOpsAndParams(interpreter, nn_model_, next_id, &model_states_inputs_, - &model_states_outputs_, tensor_id_to_nnapi_id)); - - std::vector augmented_inputs; - MapAndAddTensorIds(interpreter->inputs().data(), - interpreter->inputs().size(), &augmented_inputs, - tensor_id_to_nnapi_id); - augmented_inputs.insert(augmented_inputs.end(), - model_states_inputs_.begin(), - model_states_inputs_.end()); - std::vector augmented_outputs; - MapAndAddTensorIds(interpreter->outputs().data(), - interpreter->outputs().size(), &augmented_outputs, - tensor_id_to_nnapi_id); - MapAndAddTensorIds(model_states_outputs_.data(), - model_states_outputs_.size(), &augmented_outputs, - tensor_id_to_nnapi_id); - - CHECK_NN(ANeuralNetworksModel_identifyInputsAndOutputs( - nn_model_, static_cast(augmented_inputs.size()), - reinterpret_cast(augmented_inputs.data()), - static_cast(augmented_outputs.size()), - reinterpret_cast(augmented_outputs.data()))); - - if (GetAndroidSdkVersionCached() >= 28) { - CHECK_NN(ANeuralNetworksModel_relaxComputationFloat32toFloat16( - nn_model_, interpreter->GetAllowFp16PrecisionForFp32())); - } - CHECK_NN(ANeuralNetworksModel_finish(nn_model_)); - } - if (!nn_compiled_model_) { - CHECK_NN(ANeuralNetworksCompilation_create(nn_model_, &nn_compiled_model_)); - CHECK_NN(ANeuralNetworksCompilation_finish(nn_compiled_model_)); - } - return kTfLiteOk; -} - -TfLiteStatus NNAPIDelegate::Invoke(Interpreter* interpreter) { - if (!nn_model_) { - model_status_ = BuildGraph(interpreter); - if (model_status_ != kTfLiteOk) { - logError("Failed to build graph for NNAPI"); - } - } - if (model_status_ != kTfLiteOk) { - return model_status_; - } - - ANeuralNetworksExecution* execution = nullptr; - CHECK_NN(ANeuralNetworksExecution_create(nn_compiled_model_, &execution)); - - // Currently perform deep copy of input buffer - for (size_t i = 0; i < interpreter->inputs().size(); i++) { - int input = interpreter->inputs()[i]; - // TODO(aselle): Is this what we want or do we want input instead? - // TODO(aselle): This should be called setInputValue maybe to be cons. - TfLiteTensor* tensor = interpreter->tensor(input); - CHECK_NN(ANeuralNetworksExecution_setInput( - execution, i, nullptr, tensor->data.raw, tensor->bytes)); - } - - // Tell nn api where to place final data. - for (size_t i = 0; i < interpreter->outputs().size(); i++) { - int output = interpreter->outputs()[i]; - TfLiteTensor* tensor = interpreter->tensor(output); - CHECK_NN(ANeuralNetworksExecution_setOutput( - execution, i, nullptr, tensor->data.raw, tensor->bytes)); - } - - // The state_out of previous invocation need to be mapped to state_in of - // current invocation. - for (size_t i = 0; i < model_states_outputs_.size(); i++) { - int state_tensor_idx = model_states_outputs_[i]; - TfLiteTensor* tensor = interpreter->tensor(state_tensor_idx); - // Here we are using a deep copy for state_in tensors so that we are not - // reading and writing into the same buffer during a invocation. - // TODO(miaowang): using double shared buffer to minimize the copies. - CHECK_NN(ANeuralNetworksExecution_setInput( - execution, i + interpreter->inputs().size(), nullptr, tensor->data.raw, - tensor->bytes)); - // Tell NNAPI where to output the state_out. - CHECK_NN(ANeuralNetworksExecution_setOutput( - execution, i + interpreter->outputs().size(), nullptr, tensor->data.raw, - tensor->bytes)); - } - - // Currently use blocking compute. - ANeuralNetworksEvent* event = nullptr; - CHECK_NN(ANeuralNetworksExecution_startCompute(execution, &event)); - CHECK_NN(ANeuralNetworksEvent_wait(event)); - ANeuralNetworksEvent_free(event); - ANeuralNetworksExecution_free(execution); - -#if 0 - printf("From the NN API:\n"); - TfLiteTensor* tensor = interpreter->tensor(interpreter->outputs()[0]); - if (float* data = - interpreter->typed_tensor(interpreter->outputs()[0])) { - size_t num = tensor->bytes / sizeof(float); - for (float* p = data; p < data + num; p++) { - printf(" %f", *p); - } - printf("\n"); - } -#endif - - return kTfLiteOk; -} - -bool NNAPIDelegate::IsSupported() { return NNAPIExists(); } - -} // namespace tflite diff --git a/tensorflow/contrib/lite/nnapi_delegate.h b/tensorflow/contrib/lite/nnapi_delegate.h deleted file mode 100644 index 22359d557e61e3ca3a977803276f1c67a2229c22..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/nnapi_delegate.h +++ /dev/null @@ -1,80 +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_CONTRIB_LITE_NNAPI_DELEGATE_H_ -#define TENSORFLOW_CONTRIB_LITE_NNAPI_DELEGATE_H_ - -#include "tensorflow/contrib/lite/allocation.h" -#include "tensorflow/contrib/lite/c/c_api_internal.h" -#include "tensorflow/contrib/lite/core/api/error_reporter.h" -#include "tensorflow/contrib/lite/interpreter.h" - -class ANeuralNetworksModel; -class ANeuralNetworksMemory; -class ANeuralNetworksCompilation; - -namespace tflite { - -class NNAPIAllocation : public MMAPAllocation { - public: - NNAPIAllocation(const char* filename, ErrorReporter* error_reporter); - ~NNAPIAllocation(); - - size_t offset(const void* ptr) const { - auto signed_offset = reinterpret_cast(ptr) - - reinterpret_cast(mmapped_buffer_); - - return static_cast(signed_offset); - } - - ANeuralNetworksMemory* memory() const { return handle_; } - bool valid() const override { return handle_ != nullptr; } - - private: - mutable ANeuralNetworksMemory* handle_ = nullptr; -}; - -class NNAPIDelegate { - public: - ~NNAPIDelegate(); - - // Convert a tflite graph to NNAPI - TfLiteStatus BuildGraph(Interpreter* interpreter); - - // Run - TfLiteStatus Invoke(Interpreter* interpreter); - - // Whether the current platform supports NNAPI delegation. - static bool IsSupported(); - - private: - // The NN API model handle - ANeuralNetworksModel* nn_model_ = nullptr; - // The NN API compilation handle - ANeuralNetworksCompilation* nn_compiled_model_ = nullptr; - // Model status - TfLiteStatus model_status_ = kTfLiteOk; - - // List of state tensors for LSTM, RNN, SVDF. - // NN API does not allow ops to maintain states across multiple - // invocations. We need to manually create state input tensors from - // corresponding state output tensors of TFLite operations, and map them - // correctly. - std::vector model_states_inputs_; // holds NNAPI operand ids - std::vector model_states_outputs_; // holds TFLite tensor ids -}; - -} // namespace tflite - -#endif // TENSORFLOW_CONTRIB_LITE_NNAPI_DELEGATE_H_ diff --git a/tensorflow/contrib/lite/op_resolver.h b/tensorflow/contrib/lite/op_resolver.h deleted file mode 100644 index e93134cbdecd58cb11e6be4d777549b7c63f6595..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/op_resolver.h +++ /dev/null @@ -1,22 +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. -==============================================================================*/ -// Compatibility shim for moved header location. -#ifndef TENSORFLOW_CONTRIB_LITE_OP_RESOLVER_H_ -#define TENSORFLOW_CONTRIB_LITE_OP_RESOLVER_H_ - -#include "tensorflow/contrib/lite/core/api/op_resolver.h" -#include "tensorflow/contrib/lite/mutable_op_resolver.h" - -#endif // TENSORFLOW_CONTRIB_LITE_OP_RESOLVER_H_ diff --git a/tensorflow/contrib/lite/profiling/BUILD b/tensorflow/contrib/lite/profiling/BUILD deleted file mode 100644 index 1172722f7a70771af73eb07571349e431755471c..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/profiling/BUILD +++ /dev/null @@ -1,84 +0,0 @@ -package(default_visibility = ["//visibility:public"]) - -licenses(["notice"]) # Apache 2.0 - -load("//tensorflow/contrib/lite:build_def.bzl", "tflite_copts") - -common_copts = [ - "-Wall", -] + tflite_copts() - -cc_library( - name = "profiler", - hdrs = ["profiler.h"], - copts = common_copts, - deps = [":profile_buffer"], -) - -cc_test( - name = "profiler_test", - srcs = ["profiler_test.cc"], - copts = ["-DTFLITE_PROFILING_ENABLED"], - defines = ["TFLITE_PROFILING_ENABLED"], - deps = [ - ":profiler", - "//tensorflow/contrib/lite/testing:util", - "@com_google_googletest//:gtest", - ], -) - -cc_library( - name = "profile_buffer", - hdrs = ["profile_buffer.h"], - copts = common_copts, - deps = [":time"], -) - -cc_library( - name = "time", - srcs = ["time.cc"], - hdrs = ["time.h"], - copts = common_copts, -) - -cc_library( - name = "profile_summarizer", - srcs = ["profile_summarizer.cc"], - hdrs = ["profile_summarizer.h"], - copts = common_copts, - deps = [ - ":profiler", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/schema:schema_fbs", - "//tensorflow/core:stats_calculator_portable", - ], -) - -cc_test( - name = "profile_summarizer_test", - srcs = ["profile_summarizer_test.cc"], - copts = common_copts, - tags = ["no_oss"], - deps = [ - ":profile_summarizer", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite:schema_fbs_version", - "//tensorflow/contrib/lite/kernels:builtin_ops", - "//tensorflow/contrib/lite/kernels:kernel_util", - "//tensorflow/contrib/lite/kernels:test_util", - "//tensorflow/contrib/lite/testing:util", - "@com_google_googletest//:gtest", - ], -) - -cc_test( - name = "profile_buffer_test", - srcs = ["profile_buffer_test.cc"], - copts = ["-DTFLITE_PROFILING_ENABLED"], - defines = ["TFLITE_PROFILING_ENABLED"], - deps = [ - ":profile_buffer", - "//tensorflow/contrib/lite/testing:util", - "@com_google_googletest//:gtest", - ], -) diff --git a/tensorflow/contrib/lite/profiling/time.cc b/tensorflow/contrib/lite/profiling/time.cc deleted file mode 100644 index 875ddb02bcfc30f4c2ef543fe1c15bec467e5410..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/profiling/time.cc +++ /dev/null @@ -1,47 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include "tensorflow/contrib/lite/profiling/time.h" - -#if defined(_MSC_VER) -#include // NOLINT(build/c++11) -#else -#include -#endif - -namespace tflite { -namespace profiling { -namespace time { - -#if defined(_MSC_VER) - -uint64_t NowMicros() { - return std::chrono::duration_cast( - std::chrono::system_clock::now().time_since_epoch()) - .count(); -} - -#else - -uint64_t NowMicros() { - struct timeval tv; - gettimeofday(&tv, nullptr); - return static_cast(tv.tv_sec) * 1000000 + tv.tv_usec; -} - -#endif // defined(_MSC_VER) - -} // namespace time -} // namespace profiling -} // namespace tflite diff --git a/tensorflow/contrib/lite/python/BUILD b/tensorflow/contrib/lite/python/BUILD index 916788f21500046a33c88016b7d13d7a46430fbe..893ddd78231c8a0d819cbe5776e6873bdab57355 100644 --- a/tensorflow/contrib/lite/python/BUILD +++ b/tensorflow/contrib/lite/python/BUILD @@ -1,191 +1,12 @@ -licenses(["notice"]) # Apache 2.0 - -package(default_visibility = ["//tensorflow:internal"]) - -load("//tensorflow:tensorflow.bzl", "py_test") - -filegroup( - name = "interpreter_test_data", - srcs = glob(["**/testdata/*"]), - visibility = ["//tensorflow:__subpackages__"], -) - -py_library( - name = "interpreter", - srcs = [ - "interpreter.py", - ], - srcs_version = "PY2AND3", - visibility = ["//visibility:public"], - deps = [ - "//tensorflow/contrib/lite/python/interpreter_wrapper:tensorflow_wrap_interpreter_wrapper", - "//tensorflow/python:util", - "//third_party/py/numpy", - ], -) - -py_test( - name = "interpreter_test", - srcs = ["interpreter_test.py"], - data = [":interpreter_test_data"], - srcs_version = "PY2AND3", - tags = ["no_oss"], - deps = [ - ":interpreter", - "//tensorflow/python:client_testlib", - "//tensorflow/python:framework_test_lib", - "//tensorflow/python:platform", - "//third_party/py/numpy", - ], -) - -py_binary( - name = "tflite_convert", - srcs = ["tflite_convert.py"], - srcs_version = "PY2AND3", - visibility = ["//visibility:public"], - deps = [ - ":lite", - ], -) +licenses(["notice"]) +# DO NOT USE THIS TARGET. TensorFlow Lite has moved to tensorflow/lite. py_library( name = "lite", - srcs = ["lite.py"], - srcs_version = "PY2AND3", - visibility = ["//visibility:public"], - deps = [ - ":convert", - ":convert_saved_model", - ":interpreter", - ":lite_constants", - ":op_hint", - "//tensorflow/python:graph_util", - "//tensorflow/python/saved_model:constants", - "//tensorflow/python/saved_model:loader", - "//tensorflow/python/tools:freeze_graph_lib", - ], -) - -py_test( - name = "lite_test", - srcs = ["lite_test.py"], - data = ["@tflite_mobilenet_ssd_quant_protobuf//:tflite_graph.pb"], - srcs_version = "PY2AND3", - tags = [ - "no_oss", - "no_windows", - ], - deps = [ - ":lite", - ], -) - -py_library( - name = "lite_constants", - srcs = ["lite_constants.py"], - srcs_version = "PY2AND3", - deps = [ - "//tensorflow/contrib/lite/toco:toco_flags_proto_py", - ], -) - -py_library( - name = "convert", - srcs = ["convert.py"], + srcs = ["__init__.py"], srcs_version = "PY2AND3", visibility = ["//visibility:public"], deps = [ - ":lite_constants", - "//tensorflow/contrib/lite/toco:model_flags_proto_py", - "//tensorflow/contrib/lite/toco:toco_flags_proto_py", - "//tensorflow/contrib/lite/toco/python:tensorflow_wrap_toco", - "//tensorflow/contrib/lite/toco/python:toco_from_protos", - "//tensorflow/python:platform", - ], -) - -py_library( - name = "op_hint", - srcs = ["op_hint.py"], - srcs_version = "PY2AND3", - visibility = ["//visibility:public"], - deps = [ - "//tensorflow/contrib/framework:framework_py", - "//tensorflow/contrib/graph_editor:graph_editor_py", - "//tensorflow/core:protos_all_py", - "//tensorflow/python:framework", - "//tensorflow/python:platform", - "//tensorflow/python:util", - ], -) - -py_test( - name = "convert_test", - srcs = ["convert_test.py"], - srcs_version = "PY2AND3", - tags = [ - "no-internal-py3", - "no_oss", - ], - deps = [ - ":convert", - ":interpreter", - ":op_hint", - "//tensorflow/python:array_ops", - "//tensorflow/python:client_testlib", - "//tensorflow/python:dtypes", - "//tensorflow/python:platform_test", - "//tensorflow/python:session", - ], -) - -py_library( - name = "convert_saved_model", - srcs = ["convert_saved_model.py"], - srcs_version = "PY2AND3", - visibility = ["//tensorflow/contrib/lite:__subpackages__"], - deps = [ - ":convert", - "//tensorflow/contrib/saved_model:saved_model_py", - "//tensorflow/python:graph_util", - "//tensorflow/python:platform", - "//tensorflow/python/tools:freeze_graph_lib", - ], -) - -py_binary( - name = "create_custom_op", - srcs = ["create_custom_op.py"], - srcs_version = "PY2AND3", - visibility = ["//visibility:public"], - deps = [ - "//tensorflow/contrib/framework:framework_py", - "//tensorflow/core:protos_all_py", - "//tensorflow/python:platform", - "@absl_py//absl/flags", - ], -) - -py_test( - name = "convert_saved_model_test", - srcs = ["convert_saved_model_test.py"], - srcs_version = "PY2AND3", - tags = [ - "no_oss", - "no_windows", - ], - visibility = ["//visibility:public"], - deps = [ - ":convert_saved_model", - "//tensorflow/python:client_testlib", - "//tensorflow/python:layers", - "//tensorflow/python:nn", - "//tensorflow/python:platform_test", - "//tensorflow/python:session", - "//tensorflow/python/estimator:estimator_py", - "//tensorflow/python/keras", - "//tensorflow/python/ops/losses", - "//tensorflow/python/saved_model", + "//tensorflow/lite/python:lite", ], ) diff --git a/tensorflow/contrib/lite/python/__init__.py b/tensorflow/contrib/lite/python/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..27b1ffb251e76469092eb613d3c381718d8dc4fd --- /dev/null +++ b/tensorflow/contrib/lite/python/__init__.py @@ -0,0 +1,26 @@ +# 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. +# ============================================================================== +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow import lite + +import warnings as _warnings + +WARNING = ("WARNING: TF Lite has moved from tf.contrib.lite to tf.lite. Please " + "update your imports. This will be a breaking error in TensorFlow " + "version 2.0.") +_warnings.warn(WARNING, PendingDeprecationWarning) diff --git a/tensorflow/contrib/lite/python/convert_saved_model_test.py b/tensorflow/contrib/lite/python/convert_saved_model_test.py deleted file mode 100644 index 92c4ebb2465c2abaa1cefd020e69b2f7ad6a54a5..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/python/convert_saved_model_test.py +++ /dev/null @@ -1,459 +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. -# ============================================================================== -"""TFLite SavedModel conversion test cases. - - - Tests converting simple SavedModel graph to TFLite FlatBuffer. - - Tests converting simple SavedModel graph to frozen graph. - - Tests converting MNIST SavedModel to TFLite FlatBuffer. -""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import os -from tensorflow.contrib.lite.python import convert_saved_model -from tensorflow.python import keras -from tensorflow.python.client import session -from tensorflow.python.estimator import estimator_lib as estimator -from tensorflow.python.framework import dtypes -from tensorflow.python.framework import ops -from tensorflow.python.framework import tensor_shape -from tensorflow.python.framework import test_util -from tensorflow.python.layers import layers -from tensorflow.python.ops import array_ops -from tensorflow.python.ops import math_ops -from tensorflow.python.ops import nn -from tensorflow.python.ops import random_ops -from tensorflow.python.ops.losses import losses -from tensorflow.python.platform import test -from tensorflow.python.saved_model import saved_model -from tensorflow.python.saved_model import signature_constants -from tensorflow.python.saved_model import tag_constants -from tensorflow.python.training import training as train - - -class TensorFunctionsTest(test_util.TensorFlowTestCase): - - def testGetTensorsValid(self): - in_tensor = array_ops.placeholder( - shape=[1, 16, 16, 3], dtype=dtypes.float32) - _ = in_tensor + in_tensor - sess = session.Session() - - tensors = convert_saved_model.get_tensors_from_tensor_names( - sess.graph, ["Placeholder"]) - self.assertEqual("Placeholder:0", tensors[0].name) - - def testGetTensorsInvalid(self): - in_tensor = array_ops.placeholder( - shape=[1, 16, 16, 3], dtype=dtypes.float32) - _ = in_tensor + in_tensor - sess = session.Session() - - with self.assertRaises(ValueError) as error: - convert_saved_model.get_tensors_from_tensor_names(sess.graph, - ["invalid-input"]) - self.assertEqual("Invalid tensors 'invalid-input' were found.", - str(error.exception)) - - def testSetTensorShapeValid(self): - tensor = array_ops.placeholder(shape=[None, 3, 5], dtype=dtypes.float32) - self.assertEqual([None, 3, 5], tensor.shape.as_list()) - - convert_saved_model.set_tensor_shapes([tensor], {"Placeholder": [5, 3, 5]}) - self.assertEqual([5, 3, 5], tensor.shape.as_list()) - - def testSetTensorShapeNoneValid(self): - tensor = array_ops.placeholder(dtype=dtypes.float32) - self.assertEqual(None, tensor.shape) - - convert_saved_model.set_tensor_shapes([tensor], {"Placeholder": [1, 3, 5]}) - self.assertEqual([1, 3, 5], tensor.shape.as_list()) - - def testSetTensorShapeInvalid(self): - tensor = array_ops.placeholder(shape=[None, 3, 5], dtype=dtypes.float32) - self.assertEqual([None, 3, 5], tensor.shape.as_list()) - - convert_saved_model.set_tensor_shapes([tensor], - {"invalid-input": [5, 3, 5]}) - self.assertEqual([None, 3, 5], tensor.shape.as_list()) - - def testSetTensorShapeEmpty(self): - tensor = array_ops.placeholder(shape=[None, 3, 5], dtype=dtypes.float32) - self.assertEqual([None, 3, 5], tensor.shape.as_list()) - - convert_saved_model.set_tensor_shapes([tensor], {}) - self.assertEqual([None, 3, 5], tensor.shape.as_list()) - - -class FreezeSavedModelTest(test_util.TensorFlowTestCase): - - def _createSimpleSavedModel(self, shape): - """Create a simple SavedModel on the fly.""" - saved_model_dir = os.path.join(self.get_temp_dir(), "simple_savedmodel") - with session.Session() as sess: - in_tensor = array_ops.placeholder(shape=shape, dtype=dtypes.float32) - out_tensor = in_tensor + in_tensor - inputs = {"x": in_tensor} - outputs = {"y": out_tensor} - saved_model.simple_save(sess, saved_model_dir, inputs, outputs) - return saved_model_dir - - def _createSavedModelTwoInputArrays(self, shape): - """Create a simple SavedModel.""" - saved_model_dir = os.path.join(self.get_temp_dir(), "simple_savedmodel") - with session.Session() as sess: - in_tensor_1 = array_ops.placeholder( - shape=shape, dtype=dtypes.float32, name="inputB") - in_tensor_2 = array_ops.placeholder( - shape=shape, dtype=dtypes.float32, name="inputA") - out_tensor = in_tensor_1 + in_tensor_2 - inputs = {"x": in_tensor_1, "y": in_tensor_2} - outputs = {"z": out_tensor} - saved_model.simple_save(sess, saved_model_dir, inputs, outputs) - return saved_model_dir - - def _getArrayNames(self, tensors): - return [tensor.name for tensor in tensors] - - def _getArrayShapes(self, tensors): - dims = [] - for tensor in tensors: - dim_tensor = [] - for dim in tensor.shape: - if isinstance(dim, tensor_shape.Dimension): - dim_tensor.append(dim.value) - else: - dim_tensor.append(dim) - dims.append(dim_tensor) - return dims - - def _convertSavedModel(self, - saved_model_dir, - input_arrays=None, - input_shapes=None, - output_arrays=None, - tag_set=None, - signature_key=None): - if tag_set is None: - tag_set = set([tag_constants.SERVING]) - if signature_key is None: - signature_key = signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY - graph_def, in_tensors, out_tensors = convert_saved_model.freeze_saved_model( - saved_model_dir=saved_model_dir, - input_arrays=input_arrays, - input_shapes=input_shapes, - output_arrays=output_arrays, - tag_set=tag_set, - signature_key=signature_key) - return graph_def, in_tensors, out_tensors - - def testSimpleSavedModel(self): - """Test a SavedModel.""" - saved_model_dir = self._createSimpleSavedModel(shape=[1, 16, 16, 3]) - _, in_tensors, out_tensors = self._convertSavedModel(saved_model_dir) - - self.assertEqual(self._getArrayNames(out_tensors), ["add:0"]) - self.assertEqual(self._getArrayNames(in_tensors), ["Placeholder:0"]) - self.assertEqual(self._getArrayShapes(in_tensors), [[1, 16, 16, 3]]) - - def testSimpleSavedModelWithNoneBatchSizeInShape(self): - """Test a SavedModel with None in input tensor's shape.""" - saved_model_dir = self._createSimpleSavedModel(shape=[None, 16, 16, 3]) - _, in_tensors, out_tensors = self._convertSavedModel(saved_model_dir) - - self.assertEqual(self._getArrayNames(out_tensors), ["add:0"]) - self.assertEqual(self._getArrayNames(in_tensors), ["Placeholder:0"]) - self.assertEqual(self._getArrayShapes(in_tensors), [[None, 16, 16, 3]]) - - def testSimpleSavedModelWithInvalidSignatureKey(self): - """Test a SavedModel that fails due to an invalid signature_key.""" - saved_model_dir = self._createSimpleSavedModel(shape=[1, 16, 16, 3]) - with self.assertRaises(ValueError) as error: - self._convertSavedModel(saved_model_dir, signature_key="invalid-key") - self.assertEqual( - "No 'invalid-key' in the SavedModel's SignatureDefs. " - "Possible values are 'serving_default'.", str(error.exception)) - - def testSimpleSavedModelWithInvalidOutputArray(self): - """Test a SavedModel that fails due to invalid output arrays.""" - saved_model_dir = self._createSimpleSavedModel(shape=[1, 16, 16, 3]) - with self.assertRaises(ValueError) as error: - self._convertSavedModel(saved_model_dir, output_arrays=["invalid-output"]) - self.assertEqual("Invalid tensors 'invalid-output' were found.", - str(error.exception)) - - def testSimpleSavedModelWithWrongInputArrays(self): - """Test a SavedModel that fails due to invalid input arrays.""" - saved_model_dir = self._createSimpleSavedModel(shape=[1, 16, 16, 3]) - - # Check invalid input_arrays. - with self.assertRaises(ValueError) as error: - self._convertSavedModel(saved_model_dir, input_arrays=["invalid-input"]) - self.assertEqual("Invalid tensors 'invalid-input' were found.", - str(error.exception)) - - # Check valid and invalid input_arrays. - with self.assertRaises(ValueError) as error: - self._convertSavedModel( - saved_model_dir, input_arrays=["Placeholder", "invalid-input"]) - self.assertEqual("Invalid tensors 'invalid-input' were found.", - str(error.exception)) - - def testSimpleSavedModelWithCorrectArrays(self): - """Test a SavedModel with correct input_arrays and output_arrays.""" - saved_model_dir = self._createSimpleSavedModel(shape=[None, 16, 16, 3]) - _, in_tensors, out_tensors = self._convertSavedModel( - saved_model_dir=saved_model_dir, - input_arrays=["Placeholder"], - output_arrays=["add"]) - - self.assertEqual(self._getArrayNames(out_tensors), ["add:0"]) - self.assertEqual(self._getArrayNames(in_tensors), ["Placeholder:0"]) - self.assertEqual(self._getArrayShapes(in_tensors), [[None, 16, 16, 3]]) - - def testSimpleSavedModelWithCorrectInputArrays(self): - """Test a SavedModel with correct input_arrays and input_shapes.""" - saved_model_dir = self._createSimpleSavedModel(shape=[1, 16, 16, 3]) - _, in_tensors, out_tensors = self._convertSavedModel( - saved_model_dir=saved_model_dir, - input_arrays=["Placeholder"], - input_shapes={"Placeholder": [1, 16, 16, 3]}) - - self.assertEqual(self._getArrayNames(out_tensors), ["add:0"]) - self.assertEqual(self._getArrayNames(in_tensors), ["Placeholder:0"]) - self.assertEqual(self._getArrayShapes(in_tensors), [[1, 16, 16, 3]]) - - def testTwoInputArrays(self): - """Test a simple SavedModel.""" - saved_model_dir = self._createSavedModelTwoInputArrays(shape=[1, 16, 16, 3]) - - _, in_tensors, out_tensors = self._convertSavedModel( - saved_model_dir=saved_model_dir, input_arrays=["inputB", "inputA"]) - - self.assertEqual(self._getArrayNames(out_tensors), ["add:0"]) - self.assertEqual(self._getArrayNames(in_tensors), ["inputA:0", "inputB:0"]) - self.assertEqual( - self._getArrayShapes(in_tensors), [[1, 16, 16, 3], [1, 16, 16, 3]]) - - def testSubsetInputArrays(self): - """Test a SavedModel with a subset of the input array names of the model.""" - saved_model_dir = self._createSavedModelTwoInputArrays(shape=[1, 16, 16, 3]) - - # Check case where input shape is given. - _, in_tensors, out_tensors = self._convertSavedModel( - saved_model_dir=saved_model_dir, - input_arrays=["inputA"], - input_shapes={"inputA": [1, 16, 16, 3]}) - - self.assertEqual(self._getArrayNames(out_tensors), ["add:0"]) - self.assertEqual(self._getArrayNames(in_tensors), ["inputA:0"]) - self.assertEqual(self._getArrayShapes(in_tensors), [[1, 16, 16, 3]]) - - # Check case where input shape is None. - _, in_tensors, out_tensors = self._convertSavedModel( - saved_model_dir=saved_model_dir, input_arrays=["inputA"]) - - self.assertEqual(self._getArrayNames(out_tensors), ["add:0"]) - self.assertEqual(self._getArrayNames(in_tensors), ["inputA:0"]) - self.assertEqual(self._getArrayShapes(in_tensors), [[1, 16, 16, 3]]) - - def testMultipleMetaGraphDef(self): - """Test saved model with multiple MetaGraphDefs.""" - saved_model_dir = os.path.join(self.get_temp_dir(), "savedmodel_two_mgd") - builder = saved_model.builder.SavedModelBuilder(saved_model_dir) - with session.Session(graph=ops.Graph()) as sess: - # MetaGraphDef 1 - in_tensor = array_ops.placeholder(shape=[1, 28, 28], dtype=dtypes.float32) - out_tensor = in_tensor + in_tensor - sig_input_tensor = saved_model.utils.build_tensor_info(in_tensor) - sig_input_tensor_signature = {"x": sig_input_tensor} - sig_output_tensor = saved_model.utils.build_tensor_info(out_tensor) - sig_output_tensor_signature = {"y": sig_output_tensor} - predict_signature_def = ( - saved_model.signature_def_utils.build_signature_def( - sig_input_tensor_signature, sig_output_tensor_signature, - saved_model.signature_constants.PREDICT_METHOD_NAME)) - signature_def_map = { - saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: - predict_signature_def - } - builder.add_meta_graph_and_variables( - sess, - tags=[saved_model.tag_constants.SERVING, "additional_test_tag"], - signature_def_map=signature_def_map) - - # MetaGraphDef 2 - builder.add_meta_graph(tags=["tflite"]) - builder.save(True) - - # Convert to tflite - _, in_tensors, out_tensors = self._convertSavedModel( - saved_model_dir=saved_model_dir, - tag_set=set([saved_model.tag_constants.SERVING, "additional_test_tag"])) - - self.assertEqual(self._getArrayNames(out_tensors), ["add:0"]) - self.assertEqual(self._getArrayNames(in_tensors), ["Placeholder:0"]) - self.assertEqual(self._getArrayShapes(in_tensors), [[1, 28, 28]]) - - -class Model(keras.Model): - """Model to recognize digits in the MNIST dataset. - - Train and export SavedModel, used for testOnflyTrainMnistSavedModel - - Network structure is equivalent to: - https://github.com/tensorflow/tensorflow/blob/r1.5/tensorflow/examples/tutorials/mnist/mnist_deep.py - and - https://github.com/tensorflow/models/blob/master/tutorials/image/mnist/convolutional.py - - But written as a ops.keras.Model using the layers API. - """ - - def __init__(self, data_format): - """Creates a model for classifying a hand-written digit. - - Args: - data_format: Either "channels_first" or "channels_last". - "channels_first" is typically faster on GPUs while "channels_last" is - typically faster on CPUs. See - https://www.tensorflow.org/performance/performance_guide#data_formats - """ - super(Model, self).__init__() - self._input_shape = [-1, 28, 28, 1] - - self.conv1 = layers.Conv2D( - 32, 5, padding="same", data_format=data_format, activation=nn.relu) - self.conv2 = layers.Conv2D( - 64, 5, padding="same", data_format=data_format, activation=nn.relu) - self.fc1 = layers.Dense(1024, activation=nn.relu) - self.fc2 = layers.Dense(10) - self.dropout = layers.Dropout(0.4) - self.max_pool2d = layers.MaxPooling2D( - (2, 2), (2, 2), padding="same", data_format=data_format) - - def __call__(self, inputs, training): - """Add operations to classify a batch of input images. - - Args: - inputs: A Tensor representing a batch of input images. - training: A boolean. Set to True to add operations required only when - training the classifier. - - Returns: - A logits Tensor with shape [, 10]. - """ - y = array_ops.reshape(inputs, self._input_shape) - y = self.conv1(y) - y = self.max_pool2d(y) - y = self.conv2(y) - y = self.max_pool2d(y) - y = layers.flatten(y) - y = self.fc1(y) - y = self.dropout(y, training=training) - return self.fc2(y) - - -def model_fn(features, labels, mode, params): - """The model_fn argument for creating an Estimator.""" - model = Model(params["data_format"]) - image = features - if isinstance(image, dict): - image = features["image"] - - if mode == estimator.ModeKeys.PREDICT: - logits = model(image, training=False) - predictions = { - "classes": math_ops.argmax(logits, axis=1), - "probabilities": nn.softmax(logits), - } - return estimator.EstimatorSpec( - mode=estimator.ModeKeys.PREDICT, - predictions=predictions, - export_outputs={ - "classify": estimator.export.PredictOutput(predictions) - }) - - elif mode == estimator.ModeKeys.TRAIN: - optimizer = train.AdamOptimizer(learning_rate=1e-4) - - logits = model(image, training=True) - loss = losses.sparse_softmax_cross_entropy(labels=labels, logits=logits) - return estimator.EstimatorSpec( - mode=estimator.ModeKeys.TRAIN, - loss=loss, - train_op=optimizer.minimize(loss, train.get_or_create_global_step())) - - elif mode == estimator.ModeKeys.EVAL: - logits = model(image, training=False) - loss = losses.sparse_softmax_cross_entropy(labels=labels, logits=logits) - return estimator.EstimatorSpec( - mode=estimator.ModeKeys.EVAL, - loss=loss, - eval_metric_ops={ - "accuracy": - ops.metrics.accuracy( - labels=labels, predictions=math_ops.argmax(logits, axis=1)), - }) - - -def dummy_input_fn(): - image = random_ops.random_uniform([100, 784]) - labels = random_ops.random_uniform([100, 1], maxval=9, dtype=dtypes.int32) - return image, labels - - -class FreezeSavedModelTestTrainGraph(test_util.TensorFlowTestCase): - - def testTrainedMnistSavedModel(self): - """Test mnist SavedModel, trained with dummy data and small steps.""" - # Build classifier - classifier = estimator.Estimator( - model_fn=model_fn, - params={ - "data_format": "channels_last" # tflite format - }) - - # Train and pred for serving - classifier.train(input_fn=dummy_input_fn, steps=2) - image = array_ops.placeholder(dtypes.float32, [None, 28, 28]) - pred_input_fn = estimator.export.build_raw_serving_input_receiver_fn({ - "image": image, - }) - - # Export SavedModel - saved_model_dir = os.path.join(self.get_temp_dir(), "mnist_savedmodel") - classifier.export_savedmodel(saved_model_dir, pred_input_fn) - - # Convert to tflite and test output - saved_model_name = os.listdir(saved_model_dir)[0] - saved_model_final_dir = os.path.join(saved_model_dir, saved_model_name) - - # TODO(zhixianyan): no need to limit output_arrays to `Softmax' - # once b/74205001 fixed and argmax implemented in tflite. - result = convert_saved_model.freeze_saved_model( - saved_model_dir=saved_model_final_dir, - input_arrays=None, - input_shapes=None, - output_arrays=["Softmax"], - tag_set=set([tag_constants.SERVING]), - signature_key=signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY) - - self.assertTrue(result) - - -if __name__ == "__main__": - test.main() diff --git a/tensorflow/contrib/lite/python/interpreter_wrapper/BUILD b/tensorflow/contrib/lite/python/interpreter_wrapper/BUILD deleted file mode 100644 index 69ee95c320b72b68052c6f76f32c1493707f34b1..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/python/interpreter_wrapper/BUILD +++ /dev/null @@ -1,31 +0,0 @@ -package( - default_visibility = ["//visibility:public"], -) - -licenses(["notice"]) # Apache 2.0 - -load("//tensorflow:tensorflow.bzl", "tf_py_wrap_cc") - -cc_library( - name = "interpreter_wrapper_lib", - srcs = ["interpreter_wrapper.cc"], - hdrs = ["interpreter_wrapper.h"], - deps = [ - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:builtin_ops", - "//third_party/py/numpy:headers", - "//third_party/python_runtime:headers", - "@com_google_absl//absl/memory", - ], -) - -tf_py_wrap_cc( - name = "tensorflow_wrap_interpreter_wrapper", - srcs = [ - "interpreter_wrapper.i", - ], - deps = [ - ":interpreter_wrapper_lib", - "//third_party/python_runtime:headers", - ], -) diff --git a/tensorflow/contrib/lite/python/lite_constants.py b/tensorflow/contrib/lite/python/lite_constants.py deleted file mode 100644 index 195d7a732f337676937c7af5137d4dea84989c03..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/python/lite_constants.py +++ /dev/null @@ -1,53 +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. -# ============================================================================== -"""Constants for TFLite.""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -from tensorflow.contrib.lite.toco import toco_flags_pb2 as _toco_flags_pb2 -from tensorflow.contrib.lite.toco import types_pb2 as _types_pb2 -from tensorflow.python.util.all_util import remove_undocumented - -# Enum types from the protobuf promoted to the API -FLOAT = _types_pb2.FLOAT -INT32 = _types_pb2.INT32 -INT64 = _types_pb2.INT64 -STRING = _types_pb2.STRING -QUANTIZED_UINT8 = _types_pb2.QUANTIZED_UINT8 -TENSORFLOW_GRAPHDEF = _toco_flags_pb2.TENSORFLOW_GRAPHDEF -TFLITE = _toco_flags_pb2.TFLITE -GRAPHVIZ_DOT = _toco_flags_pb2.GRAPHVIZ_DOT - -# Currently the default mode of operation is to shell to another python process -# to protect against crashes. However, it breaks some dependent targets because -# it forces us to depend on an external py_binary. The experimental API doesn't -# have that drawback. -EXPERIMENTAL_USE_TOCO_API_DIRECTLY = False - - -_allowed_symbols = [ - "FLOAT", - "INT32", - "INT64", - "STRING", - "QUANTIZED_UINT8", - "TENSORFLOW_GRAPHDEF", - "TFLITE", - "GRAPHVIZ_DOT", - "EXPERIMENTAL_USE_TOCO_API_DIRECTLY", -] -remove_undocumented(__name__, _allowed_symbols) diff --git a/tensorflow/contrib/lite/schema/BUILD b/tensorflow/contrib/lite/schema/BUILD deleted file mode 100644 index d892466c7a1d9c953644bd4e91a468a2e9702bde..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/schema/BUILD +++ /dev/null @@ -1,99 +0,0 @@ -package(default_visibility = [ - "//visibility:public", -]) - -licenses(["notice"]) # Apache 2.0 - -load("//tensorflow:tensorflow.bzl", "py_test") -load("//tensorflow/contrib/lite:special_rules.bzl", "tflite_portable_test_suite") - -py_binary( - name = "upgrade_schema", - srcs = [ - "upgrade_schema.py", - ], - data = [ - "schema_v0.fbs", - "schema_v1.fbs", - "schema_v2.fbs", - "schema_v3.fbs", - "@flatbuffers//:flatc", - ], - deps = [ - "//tensorflow:tensorflow_py", - "//tensorflow/python:platform", - ], -) - -# TODO(wvo): re-enable this test once latest FlatBuffers has landed. - -py_test( - name = "upgrade_schema_test", - size = "small", - srcs = ["upgrade_schema_test.py"], - srcs_version = "PY2AND3", - tags = [ - "manual", - "no_oss", - "no_pip", - "notap", - ], - deps = [ - ":upgrade_schema", - "//tensorflow/python:client_testlib", - "//tensorflow/python:framework_test_lib", - ], -) - -exports_files([ - "schema_v0.fbs", - "schema_v1.fbs", - "schema_v2.fbs", - "schema_v3.fbs", -]) - -load("@flatbuffers//:build_defs.bzl", "flatbuffer_cc_library") - -# Generic schema for inference on device. -flatbuffer_cc_library( - name = "schema_fbs", - srcs = ["schema.fbs"], -) - -# Generic schema for inference on device (but with reflections makes bigger). -flatbuffer_cc_library( - name = "schema_fbs_with_reflection", - srcs = ["schema.fbs"], - flatc_args = [ - "--reflect-types", - "--reflect-names", - "--no-union-value-namespacing", - "--gen-object-api", - ], - gen_reflections = True, - out_prefix = "reflection/", -) - -# Schema test to make sure we don't introduce backward incompatible changes -# to schemas. -cc_test( - name = "flatbuffer_compatibility_test", - size = "small", - srcs = ["flatbuffer_compatibility_test.cc"], - data = [ - "schema.fbs", - "schema_v3.fbs", - ], - tags = [ - "no_oss", - "tflite_not_portable_android", - "tflite_not_portable_ios", - ], - deps = [ - "//tensorflow/core:lib_platform", - "@com_google_googletest//:gtest", - "@flatbuffers//:flatc_library", - ], -) - -tflite_portable_test_suite() diff --git a/tensorflow/contrib/lite/schema/builtin_ops_header/BUILD b/tensorflow/contrib/lite/schema/builtin_ops_header/BUILD deleted file mode 100644 index 4a627761daf45b0fddd7b99e8a9c3d0d0ed2ee5e..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/schema/builtin_ops_header/BUILD +++ /dev/null @@ -1,45 +0,0 @@ -package(default_visibility = [ - "//visibility:public", -]) - -licenses(["notice"]) # Apache 2.0 - -cc_library( - name = "generator", - srcs = ["generator.cc"], - hdrs = ["generator.h"], - deps = [ - "//tensorflow/contrib/lite/schema:schema_fbs", - ], -) - -cc_binary( - name = "generate", - srcs = ["generate.cc"], - deps = [ - ":generator", - ], -) - -cc_test( - name = "generator_test", - srcs = ["generator_test.cc"], - tags = ["no_oss"], - deps = [ - ":generator", - "@com_google_googletest//:gtest", - ], -) - -cc_test( - name = "consistency_test", - srcs = ["consistency_test.cc"], - data = [ - "//tensorflow/contrib/lite:builtin_ops.h", - ], - tags = ["no_oss"], - deps = [ - ":generator", - "@com_google_googletest//:gtest", - ], -) diff --git a/tensorflow/contrib/lite/schema/builtin_ops_header/README.md b/tensorflow/contrib/lite/schema/builtin_ops_header/README.md deleted file mode 100644 index f20d4f664e62fdd52e55339e45b9603307a2b671..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/schema/builtin_ops_header/README.md +++ /dev/null @@ -1,12 +0,0 @@ -# Builtin Ops Header Generator. - -This directory contains a code generator to generate a pure C header for -builtin op definition. - -Whenever you add a new builtin op, please execute: - -```sh -bazel run \ - //tensorflow/contrib/lite/schema/builtin_ops_header:generate > \ - tensorflow/contrib/lite/builtin_ops.h -``` diff --git a/tensorflow/contrib/lite/special_rules.bzl b/tensorflow/contrib/lite/special_rules.bzl deleted file mode 100644 index 54083c49182c707620cbd231b957405cfe24be92..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/special_rules.bzl +++ /dev/null @@ -1,6 +0,0 @@ -"""External versions of build rules that differ outside of Google.""" - -def tflite_portable_test_suite(**kwargs): - """This is a no-op outside of Google.""" - _ignore = [kwargs] - pass diff --git a/tensorflow/contrib/lite/string_util_test.cc b/tensorflow/contrib/lite/string_util_test.cc deleted file mode 100644 index a583a9184be91b0b51ac3719bf734a1a6cf563ca..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/string_util_test.cc +++ /dev/null @@ -1,118 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include "tensorflow/contrib/lite/string_util.h" - -#include -#include "tensorflow/contrib/lite/c/c_api_internal.h" -#include "tensorflow/contrib/lite/interpreter.h" -#include "tensorflow/contrib/lite/testing/util.h" - -namespace tflite { - -TEST(StringUtil, TestStringUtil) { - Interpreter interpreter; - interpreter.AddTensors(3); - - TfLiteTensor* t0 = interpreter.tensor(0); - t0->type = kTfLiteString; - t0->allocation_type = kTfLiteDynamic; - - TfLiteTensor* t1 = interpreter.tensor(1); - t1->type = kTfLiteString; - t1->allocation_type = kTfLiteDynamic; - - char data[] = {1, 0, 0, 0, 12, 0, 0, 0, 15, 0, 0, 0, 'X', 'Y', 'Z'}; - - interpreter.SetTensorParametersReadOnly(2, kTfLiteString, "", {1}, {}, data, - 15); - TfLiteTensor* t2 = interpreter.tensor(2); - interpreter.AllocateTensors(); - - char s0[] = "ABC"; - string s1 = "DEFG"; - char s2[] = ""; - - // Write strings to tensors - DynamicBuffer buf0; - buf0.AddString(s0, 3); - DynamicBuffer buf1; - buf1.AddString(s1.data(), s1.length()); - buf0.AddString(s2, 0); - buf0.WriteToTensor(t0); - buf1.WriteToTensor(t1); - - // Read strings from tensors. - ASSERT_EQ(GetStringCount(t0), 2); - StringRef str_ref; - str_ref = GetString(t0, 0); - ASSERT_EQ(string(str_ref.str, str_ref.len), "ABC"); - str_ref = GetString(t0, 1); - ASSERT_EQ(string(str_ref.str, str_ref.len), ""); - ASSERT_EQ(t0->bytes, 19); - - ASSERT_EQ(GetStringCount(t1), 1); - str_ref = GetString(t1, 0); - ASSERT_EQ(string(str_ref.str, str_ref.len), "DEFG"); - ASSERT_EQ(t1->bytes, 16); - - ASSERT_EQ(GetStringCount(t2), 1); - str_ref = GetString(t2, 0); - ASSERT_EQ(string(str_ref.str, str_ref.len), "XYZ"); - ASSERT_EQ(t2->bytes, 15); -} - -TEST(StringUtil, TestAddJoinedString) { - Interpreter interpreter; - interpreter.AddTensors(1); - TfLiteTensor* t0 = interpreter.tensor(0); - t0->type = kTfLiteString; - t0->allocation_type = kTfLiteDynamic; - - char s0[] = "ABC"; - char s1[] = "DEFG"; - char s2[] = ""; - char s3[] = "XYZ"; - - DynamicBuffer buf; - buf.AddJoinedString({{s0, 3}, {s1, 4}, {s2, 0}, {s3, 3}}, ' '); - buf.WriteToTensor(t0); - - ASSERT_EQ(GetStringCount(t0), 1); - StringRef str_ref; - str_ref = GetString(t0, 0); - ASSERT_EQ(string(str_ref.str, str_ref.len), "ABC DEFG XYZ"); - ASSERT_EQ(t0->bytes, 25); -} - -TEST(StringUtil, TestEmptyList) { - Interpreter interpreter; - interpreter.AddTensors(1); - TfLiteTensor* t0 = interpreter.tensor(0); - t0->type = kTfLiteString; - t0->allocation_type = kTfLiteDynamic; - DynamicBuffer buf; - buf.WriteToTensor(t0); - - ASSERT_EQ(GetStringCount(t0), 0); - ASSERT_EQ(t0->bytes, 8); -} - -} // namespace tflite - -int main(int argc, char** argv) { - ::tflite::LogToStderr(); - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/contrib/lite/testdata/add.bin b/tensorflow/contrib/lite/testdata/add.bin deleted file mode 100644 index aef0fe3d82c9d92dc444076d3b46e05af1923f46..0000000000000000000000000000000000000000 Binary files a/tensorflow/contrib/lite/testdata/add.bin and /dev/null differ diff --git a/tensorflow/contrib/lite/testdata/multi_add.json b/tensorflow/contrib/lite/testdata/multi_add.json deleted file mode 100644 index 97b931dba8b1050ecf91939d1d9dcea5e0ea56fb..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/testdata/multi_add.json +++ /dev/null @@ -1,46 +0,0 @@ -{ - "version": 1, - "operator_codes": [ - { - "builtin_code": "ADD" - } - ], - "subgraphs": [ - { - "tensors": [ - { "shape": [ 1, 8, 8, 3 ], "name": "a" }, - { "shape": [ 1, 8, 8, 3 ], "name": "b" }, - { "shape": [ 1, 8, 8, 3 ], "name": "c" }, - { "shape": [ 1, 8, 8, 3 ], "name": "d" }, - { "shape": [ 1, 8, 8, 3 ], "name": "i" }, - { "shape": [ 1, 8, 8, 3 ], "name": "x" }, - { "shape": [ 1, 8, 8, 3 ], "name": "y" } - ], - "inputs": [ 0, 1, 2, 3 ], - "outputs": [ 5, 6 ], - "operators": [ - { - "inputs": [ 1, 2 ], - "outputs": [ 4 ], - "builtin_options_type": "AddOptions", - "builtin_options": { - } - }, - { - "inputs": [ 0, 4 ], - "outputs": [ 5 ], - "builtin_options_type": "AddOptions", - "builtin_options": { - } - }, - { - "inputs": [ 3, 4 ], - "outputs": [ 6 ], - "builtin_options_type": "AddOptions", - "builtin_options": { - } - } - ] - } - ] -} diff --git a/tensorflow/contrib/lite/testing/BUILD b/tensorflow/contrib/lite/testing/BUILD deleted file mode 100644 index 891d44d2b60c713532d7e0e1b2c347ab891eb718..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/testing/BUILD +++ /dev/null @@ -1,390 +0,0 @@ -package(default_visibility = [ - "//visibility:public", -]) - -licenses(["notice"]) # Apache 2.0 - -load( - "//tensorflow/contrib/lite:build_def.bzl", - "gen_zip_test", - "generated_test_models_all", -) -load("//tensorflow/contrib/lite:special_rules.bzl", "tflite_portable_test_suite") -load( - "//tensorflow:tensorflow.bzl", - "tf_cc_test", - "py_test", -) - -[gen_zip_test( - name = "zip_test_%s" % test_name, - size = "large", - srcs = ["generated_examples_zip_test.cc"], - args = args + select({ - "//tensorflow:android": [], - "//conditions:default": [ - "--zip_file_path=$(location :zip_%s)" % test_name, - # TODO(angerson) We may be able to add an external unzip binary instead - # of relying on an existing one for OSS builds. - "--unzip_binary_path=/usr/bin/unzip", - ], - }), - conversion_mode = conversion_mode, - data = [ - ":zip_%s" % test_name, - ], - shard_count = 20, - tags = tags + [ - "gen_zip_test", - "no_oss", - "tflite_not_portable_intentional", - ], - test_name = test_name, - deps = [ - ":parse_testdata_lib", - ":tflite_driver", - ":util", - "@com_google_googletest//:gtest", - "@com_googlesource_code_re2//:re2", - "//tensorflow/contrib/lite:builtin_op_data", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:builtin_ops", - ] + select({ - "//conditions:default": [ - "//tensorflow/core:framework_internal", - "//tensorflow/core:lib", - "//tensorflow/core:test", - ], - "//tensorflow:android": [ - "//tensorflow/core:android_tensorflow_lib", - "//tensorflow/core:android_tensorflow_test_lib", - ], - }), -) for conversion_mode, test_name, tags, args in generated_test_models_all()] - -test_suite( - name = "generated_zip_tests", - tags = [ - "gen_zip_test", - ], -) - -py_binary( - name = "generate_examples", - srcs = ["generate_examples.py"], - data = [ - "//tensorflow/contrib/lite/toco", - ], - srcs_version = "PY2AND3", - deps = [ - ":generate_examples_report", - "//tensorflow:tensorflow_py", - "//tensorflow/python:graph_util", - "//third_party/py/numpy", - "@six_archive//:six", - ], -) - -py_library( - name = "generate_examples_report", - srcs = ["generate_examples_report.py"], - srcs_version = "PY2AND3", -) - -cc_library( - name = "parse_testdata_lib", - srcs = ["parse_testdata.cc"], - hdrs = ["parse_testdata.h"], - deps = [ - ":message", - ":split", - ":test_runner", - "//tensorflow/contrib/lite:framework", - ], -) - -cc_library( - name = "message", - srcs = ["message.cc"], - hdrs = ["message.h"], - deps = [":tokenize"], -) - -cc_test( - name = "message_test", - srcs = ["message_test.cc"], - deps = [ - ":message", - "@com_google_googletest//:gtest_main", - ], -) - -cc_library( - name = "split", - srcs = ["split.cc"], - hdrs = ["split.h"], - deps = [ - "//tensorflow/contrib/lite:string", - ], -) - -cc_test( - name = "split_test", - size = "small", - srcs = ["split_test.cc"], - deps = [ - ":split", - "@com_google_googletest//:gtest_main", - ], -) - -cc_library( - name = "join", - hdrs = ["join.h"], - deps = ["//tensorflow/contrib/lite:string"], -) - -cc_test( - name = "join_test", - size = "small", - srcs = ["join_test.cc"], - deps = [ - ":join", - "@com_google_googletest//:gtest_main", - ], -) - -cc_library( - name = "tflite_driver", - srcs = ["tflite_driver.cc"], - hdrs = ["tflite_driver.h"], - deps = [ - ":split", - ":test_runner", - "//tensorflow/contrib/lite:builtin_op_data", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/delegates/flex:delegate", - "//tensorflow/contrib/lite/kernels:builtin_ops", - ], -) - -tf_cc_test( - name = "tflite_driver_test", - size = "small", - srcs = ["tflite_driver_test.cc"], - data = ["//tensorflow/contrib/lite:testdata/multi_add.bin"], - tags = [ - "tflite_not_portable_android", - "tflite_not_portable_ios", - ], - deps = [ - ":tflite_driver", - "@com_google_googletest//:gtest_main", - ], -) - -cc_library( - name = "tokenize", - srcs = ["tokenize.cc"], - hdrs = ["tokenize.h"], - deps = [ - "//tensorflow/contrib/lite:string", - ], -) - -cc_test( - name = "tokenize_test", - srcs = ["tokenize_test.cc"], - deps = [ - ":tokenize", - "@com_google_googletest//:gtest_main", - ], -) - -cc_library( - name = "test_runner", - hdrs = ["test_runner.h"], - deps = [ - "//tensorflow/contrib/lite:string", - ], -) - -cc_library( - name = "util", - hdrs = ["util.h"], - deps = [ - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite:string", - "//tensorflow/contrib/lite/core/api", - ], -) - -cc_test( - name = "test_runner_test", - srcs = ["test_runner_test.cc"], - deps = [ - ":test_runner", - "@com_google_googletest//:gtest_main", - ], -) - -cc_binary( - name = "nnapi_example", - srcs = ["nnapi_example.cc"], - deps = [ - ":parse_testdata_lib", - ":tflite_driver", - "//tensorflow/contrib/lite/nnapi:nnapi_lib", - ], -) - -cc_library( - name = "tf_driver", - srcs = ["tf_driver.cc"], - hdrs = ["tf_driver.h"], - deps = [ - ":join", - ":split", - ":test_runner", - "//tensorflow/core:core_cpu", - "//tensorflow/core:framework", - "//tensorflow/core:lib", - "//tensorflow/core:protos_all_cc", - "//tensorflow/core:tensorflow", - ], -) - -cc_test( - name = "tf_driver_test", - size = "small", - srcs = ["tf_driver_test.cc"], - data = ["//tensorflow/contrib/lite:testdata/multi_add.pb"], - tags = [ - "no_oss", - "tflite_not_portable", - ], - deps = [ - ":tf_driver", - "@com_google_googletest//:gtest_main", - ], -) - -cc_library( - name = "generate_testspec", - srcs = ["generate_testspec.cc"], - hdrs = ["generate_testspec.h"], - deps = [ - ":join", - ":split", - ":tf_driver", - "//tensorflow/contrib/lite:string", - "//tensorflow/core:framework", - ], -) - -cc_test( - name = "generate_testspec_test", - size = "small", - srcs = ["generate_testspec_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable", - ], - deps = [ - ":generate_testspec", - "@com_google_googletest//:gtest_main", - ], -) - -cc_library( - name = "init_tensorflow", - srcs = [ - "init_tensorflow.cc", - ], - hdrs = [ - "init_tensorflow.h", - ], - visibility = [ - "//tensorflow/contrib/lite/java/src/main/native:__subpackages__", - "//tensorflow/contrib/lite/testing:__subpackages__", - "//tensorflow/contrib/lite/tools/benchmark:__subpackages__", - ], - deps = select({ - "//conditions:default": [ - "//tensorflow/core:lib", - ], - "//tensorflow:android": [ - "//tensorflow/core:android_tensorflow_lib", - ], - }), -) - -cc_library( - name = "tflite_diff_util", - srcs = ["tflite_diff_util.cc"], - hdrs = ["tflite_diff_util.h"], - deps = [ - ":generate_testspec", - ":parse_testdata_lib", - ":tflite_driver", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite:string", - ], -) - -cc_library( - name = "tflite_diff_flags", - hdrs = ["tflite_diff_flags.h"], - deps = [ - ":split", - ":tflite_diff_util", - ] + select({ - "//conditions:default": [ - "//tensorflow/core:framework_internal", - "//tensorflow/core:lib", - ], - "//tensorflow:android": [ - "//tensorflow/core:android_tensorflow_lib", - ], - }), -) - -tf_cc_test( - name = "tflite_diff_example_test", - size = "medium", - srcs = ["tflite_diff_example_test.cc"], - args = [ - "--tensorflow_model=third_party/tensorflow/contrib/lite/testdata/multi_add.pb", - "--tflite_model=third_party/tensorflow/contrib/lite/testdata/multi_add.bin", - "--input_layer=a,b,c,d", - "--input_layer_type=float,float,float,float", - "--input_layer_shape=1,3,4,3:1,3,4,3:1,3,4,3:1,3,4,3", - "--output_layer=x,y", - ], - data = [ - "//tensorflow/contrib/lite:testdata/multi_add.bin", - "//tensorflow/contrib/lite:testdata/multi_add.pb", - ], - tags = [ - "no_cuda_on_cpu_tap", - "no_oss", # needs test data - "tflite_not_portable", - ], - deps = [ - ":init_tensorflow", - ":tflite_diff_flags", - ":tflite_diff_util", - ], -) - -cc_binary( - name = "tflite_diff", - srcs = ["tflite_diff_example_test.cc"], - deps = [ - ":init_tensorflow", - ":tflite_diff_flags", - ":tflite_diff_util", - ], -) - -tflite_portable_test_suite() diff --git a/tensorflow/contrib/lite/testing/join.h b/tensorflow/contrib/lite/testing/join.h deleted file mode 100644 index 4be19ad7569c3333b6647b91adbc6e77ff088f10..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/testing/join.h +++ /dev/null @@ -1,59 +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_CONTRIB_LITE_TESTING_JOIN_H_ -#define TENSORFLOW_CONTRIB_LITE_TESTING_JOIN_H_ - -#include -#include - -#include "tensorflow/contrib/lite/string.h" - -namespace tflite { -namespace testing { - -// Join a list of data separated by delimiter. -template -string Join(T* data, size_t len, const string& delimiter) { - if (len == 0 || data == nullptr) { - return ""; - } - std::stringstream result; - result << data[0]; - for (int i = 1; i < len; i++) { - result << delimiter << data[i]; - } - return result.str(); -} - -// Join a list of uint8 data separated by a delimiter. Cast data to int before -// placing it in the string to prevent values from being treated like chars. -template <> -inline string Join(uint8_t* data, size_t len, - const string& delimiter) { - if (len == 0 || data == nullptr) { - return ""; - } - std::stringstream result; - result << static_cast(data[0]); - for (int i = 1; i < len; i++) { - result << delimiter << static_cast(data[i]); - } - return result.str(); -} - -} // namespace testing -} // namespace tflite - -#endif // TENSORFLOW_CONTRIB_LITE_TESTING_JOIN_H_ diff --git a/tensorflow/contrib/lite/testing/model_coverage/BUILD b/tensorflow/contrib/lite/testing/model_coverage/BUILD deleted file mode 100644 index c8359bab064b7c487a5b0e2303e76bc348b11ce1..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/testing/model_coverage/BUILD +++ /dev/null @@ -1,33 +0,0 @@ -package(default_visibility = [ - "//tensorflow/contrib/lite:__subpackages__", -]) - -licenses(["notice"]) # Apache 2.0 - -py_binary( - name = "model_coverage_lib", - srcs = ["model_coverage_lib.py"], - srcs_version = "PY2AND3", - tags = ["no_pip"], - deps = [ - "//tensorflow/contrib/lite/python:lite", - "//tensorflow/python:platform", - ], -) - -py_test( - name = "model_coverage_lib_test", - srcs = ["model_coverage_lib_test.py"], - srcs_version = "PY2AND3", - tags = [ - "manual", - "no_oss", - "no_pip", - "no_windows", - "notap", - ], - deps = [ - ":model_coverage_lib", - "//tensorflow/python:client_testlib", - ], -) diff --git a/tensorflow/contrib/lite/testing/split.cc b/tensorflow/contrib/lite/testing/split.cc deleted file mode 100644 index 5836f4ff049b70c00d22524a3bf3327074281f3a..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/testing/split.cc +++ /dev/null @@ -1,42 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include "tensorflow/contrib/lite/testing/split.h" - -namespace tflite { -namespace testing { - -std::vector> SplitToPos(const string& s, - const string& delimiter) { - std::vector> fields; - if (delimiter.length() == 0) { - fields.emplace_back(0, s.length()); - return fields; - } - size_t pos = 0; - size_t start = 0; - while ((pos = s.find(delimiter, start)) != string::npos) { - if (pos != start) { - fields.emplace_back(start, pos); - } - start = pos + delimiter.length(); - } - if (start != s.length()) { - fields.emplace_back(start, s.length()); - } - return fields; -} - -} // namespace testing -} // namespace tflite diff --git a/tensorflow/contrib/lite/testing/split_test.cc b/tensorflow/contrib/lite/testing/split_test.cc deleted file mode 100644 index 76b918cbcd83ef43c52057b84bcc2a8f4ff6b8f7..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/testing/split_test.cc +++ /dev/null @@ -1,62 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include "tensorflow/contrib/lite/testing/split.h" - -#include -#include - -namespace tflite { -namespace testing { -namespace { - -using ::testing::ElementsAre; -using ::testing::Pair; - -TEST(SplitTest, SplitToPos) { - EXPECT_THAT(SplitToPos("test;:1-2-3 ;: test", ";:"), - ElementsAre(Pair(0, 4), Pair(6, 12), Pair(14, 19))); - EXPECT_THAT(SplitToPos("test;:1-2-3 ;: test", ":"), - ElementsAre(Pair(0, 5), Pair(6, 13), Pair(14, 19))); - EXPECT_THAT(SplitToPos("test", ":"), ElementsAre(Pair(0, 4))); - EXPECT_THAT(SplitToPos("test ", ":"), ElementsAre(Pair(0, 5))); - EXPECT_THAT(SplitToPos("", ":"), ElementsAre()); - EXPECT_THAT(SplitToPos("test ", ""), ElementsAre(Pair(0, 5))); - EXPECT_THAT(SplitToPos("::::", ":"), ElementsAre()); -} - -TEST(SplitTest, SplitString) { - EXPECT_THAT(Split("A;B;C", ";"), ElementsAre("A", "B", "C")); -} - -TEST(SplitTest, SplitFloat) { - EXPECT_THAT(Split("1.0 B 1e-5", " "), ElementsAre(1.0, 0.0, 1e-5)); -} - -TEST(SplitTest, SplitInt) { - EXPECT_THAT(Split("1,-1,258", ","), ElementsAre(1, -1, 258)); -} - -TEST(SplitTest, SplitUint8) { - EXPECT_THAT(Split("1,-1,258", ","), ElementsAre(1, 255, 2)); -} - -TEST(SplitTest, SplitBool) { - EXPECT_THAT(Split("1, 0, 0, 1", ","), - ElementsAre(true, false, false, true)); -} - -} // namespace -} // namespace testing -} // namespace tflite diff --git a/tensorflow/contrib/lite/testing/tf_driver.cc b/tensorflow/contrib/lite/testing/tf_driver.cc deleted file mode 100644 index 30381ba028352e32a4220231eda45204889c05fb..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/testing/tf_driver.cc +++ /dev/null @@ -1,189 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include "tensorflow/contrib/lite/testing/tf_driver.h" - -#include -#include - -#include "tensorflow/contrib/lite/testing/join.h" -#include "tensorflow/contrib/lite/testing/split.h" -#include "tensorflow/core/lib/gtl/array_slice.h" - -namespace tflite { -namespace testing { - -namespace { - -tensorflow::Tensor CreateTensor(const tensorflow::DataType type, - const std::vector& dim) { - tensorflow::TensorShape shape{tensorflow::gtl::ArraySlice{ - reinterpret_cast(dim.data()), dim.size()}}; - return {type, shape}; -} - -template -void FillTensorWithData(tensorflow::Tensor* tensor, const string& csv_values) { - auto data = tensor->flat(); - - const auto& values = testing::Split(csv_values, ","); - for (int i = 0; i < values.size(); i++) { - data(i) = values[i]; - } -} - -template -void FillTensorWithZeros(tensorflow::Tensor* tensor) { - auto data = tensor->flat(); - for (int i = 0; i < tensor->NumElements(); i++) { - data(i) = 0; - } -} - -template -string TensorDataToCsvString(const tensorflow::Tensor& tensor) { - const auto& data = tensor.flat(); - return Join(data.data(), data.size(), ","); -} - -} // namespace - -TfDriver::TfDriver(const std::vector& input_layer, - const std::vector& input_layer_type, - const std::vector& input_layer_shape, - const std::vector& output_layer) - : input_names_(input_layer), output_names_(output_layer) { - CHECK_EQ(input_layer.size(), input_layer_type.size()); - CHECK_EQ(input_layer.size(), input_layer_shape.size()); - - input_ids_.resize(input_layer.size()); - input_tensors_.reserve(input_layer.size()); - input_types_.resize(input_layer.size()); - input_shapes_.resize(input_layer.size()); - for (int i = 0; i < input_layer.size(); i++) { - input_ids_[i] = i; - input_tensors_[input_layer[i]] = {}; - CHECK(DataTypeFromString(input_layer_type[i], &input_types_[i])); - input_shapes_[i] = Split(input_layer_shape[i], ","); - } - - output_ids_.resize(output_layer.size()); - output_tensors_.reserve(output_layer.size()); - for (int i = 0; i < output_layer.size(); i++) { - output_ids_[i] = i; - } -} - -void TfDriver::LoadModel(const string& bin_file_path) { - if (!IsValid()) return; - std::ifstream model(bin_file_path); - if (model.fail()) { - Invalidate("Failed to find the model " + bin_file_path); - return; - } - - tensorflow::GraphDef graphdef; - if (!graphdef.ParseFromIstream(&model)) { - Invalidate("Failed to parse tensorflow graphdef"); - return; - } - - tensorflow::SessionOptions options; - session_.reset(tensorflow::NewSession(options)); - auto status = session_->Create(graphdef); - if (!status.ok()) { - Invalidate("Failed to create session. " + status.error_message()); - } -} - -void TfDriver::SetInput(int id, const string& csv_values) { - if (!IsValid()) return; - - auto tensor = CreateTensor(input_types_[id], input_shapes_[id]); - switch (input_types_[id]) { - case tensorflow::DT_FLOAT: { - FillTensorWithData(&tensor, csv_values); - break; - } - case tensorflow::DT_INT32: { - FillTensorWithData(&tensor, csv_values); - break; - } - case tensorflow::DT_UINT8: { - FillTensorWithData(&tensor, csv_values); - break; - } - default: - fprintf(stderr, "Unsupported type %d in SetInput\n", input_types_[id]); - Invalidate("Unsupported tensor data type"); - return; - } - input_tensors_[input_names_[id]] = tensor; -} - -void TfDriver::ResetTensor(int id) { - if (!IsValid()) return; - auto tensor = input_tensors_[input_names_[id]]; - switch (input_types_[id]) { - case tensorflow::DT_FLOAT: { - FillTensorWithZeros(&tensor); - break; - } - case tensorflow::DT_INT32: { - FillTensorWithZeros(&tensor); - break; - } - default: - fprintf(stderr, "Unsupported type %d in ResetTensor\n", input_types_[id]); - Invalidate("Unsupported tensor data type"); - return; - } -} - -void TfDriver::ReshapeTensor(int id, const string& csv_values) { - input_shapes_[id] = Split(csv_values, ","); - input_tensors_[input_names_[id]] = - CreateTensor(input_types_[id], input_shapes_[id]); - ResetTensor(id); -} - -string TfDriver::ReadOutput(int id) { - if (!IsValid()) return ""; - switch (output_tensors_[id].dtype()) { - case tensorflow::DT_FLOAT: - return TensorDataToCsvString(output_tensors_[id]); - case tensorflow::DT_INT32: - return TensorDataToCsvString(output_tensors_[id]); - case tensorflow::DT_UINT8: - return TensorDataToCsvString(output_tensors_[id]); - default: - fprintf(stderr, "Unsupported type %d in ResetTensor\n", input_types_[id]); - Invalidate("Unsupported tensor data type"); - return ""; - } -} - -void TfDriver::Invoke() { - if (!IsValid()) return; - auto status = session_->Run({input_tensors_.begin(), input_tensors_.end()}, - output_names_, {}, &output_tensors_); - if (!status.ok()) { - Invalidate( - "Failed to run input data on graph. Make sure the correct value is " - "defined for the input and output arrays."); - } -} - -} // namespace testing -} // namespace tflite diff --git a/tensorflow/contrib/lite/testing/tf_driver_test.cc b/tensorflow/contrib/lite/testing/tf_driver_test.cc deleted file mode 100644 index c0faa4676adc3e846ad398bb203b77b99a2ba360..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/testing/tf_driver_test.cc +++ /dev/null @@ -1,56 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include "tensorflow/contrib/lite/testing/tf_driver.h" - -#include -#include - -namespace tflite { -namespace testing { -namespace { - -using ::testing::ElementsAre; - -TEST(TfDriverTest, SimpleTest) { - std::unique_ptr runner( - new TfDriver({"a", "b", "c", "d"}, {"float", "float", "float", "float"}, - {"1,8,8,3", "1,8,8,3", "1,8,8,3", "1,8,8,3"}, {"x", "y"})); - - runner->LoadModel( - "third_party/tensorflow/contrib/lite/testdata/multi_add.pb"); - EXPECT_TRUE(runner->IsValid()) << runner->GetErrorMessage(); - - ASSERT_THAT(runner->GetInputs(), ElementsAre(0, 1, 2, 3)); - ASSERT_THAT(runner->GetOutputs(), ElementsAre(0, 1)); - - for (int i : {0, 1, 2, 3}) { - runner->ReshapeTensor(i, "1,2,2,1"); - } - ASSERT_TRUE(runner->IsValid()); - - runner->SetInput(0, "0.1,0.2,0.3,0.4"); - runner->SetInput(1, "0.001,0.002,0.003,0.004"); - runner->SetInput(2, "0.001,0.002,0.003,0.004"); - runner->SetInput(3, "0.01,0.02,0.03,0.04"); - runner->ResetTensor(2); - runner->Invoke(); - - ASSERT_EQ(runner->ReadOutput(0), "0.101,0.202,0.303,0.404"); - ASSERT_EQ(runner->ReadOutput(1), "0.011,0.022,0.033,0.044"); -} - -} // namespace -} // namespace testing -} // namespace tflite diff --git a/tensorflow/contrib/lite/testing/tflite_driver.cc b/tensorflow/contrib/lite/testing/tflite_driver.cc deleted file mode 100644 index ef49e6f8bc30a63144521571046d9dcbd22df22e..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/testing/tflite_driver.cc +++ /dev/null @@ -1,308 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include "tensorflow/contrib/lite/testing/tflite_driver.h" - -#include - -#include "tensorflow/contrib/lite/builtin_op_data.h" -#include "tensorflow/contrib/lite/delegates/flex/delegate.h" -#include "tensorflow/contrib/lite/testing/split.h" - -namespace tflite { -namespace testing { - -namespace { - -// Returns the value in the given position in a tensor. -template -T Value(const TfLitePtrUnion& data, int index); -template <> -float Value(const TfLitePtrUnion& data, int index) { - return data.f[index]; -} -template <> -int32_t Value(const TfLitePtrUnion& data, int index) { - return data.i32[index]; -} -template <> -int64_t Value(const TfLitePtrUnion& data, int index) { - return data.i64[index]; -} -template <> -uint8_t Value(const TfLitePtrUnion& data, int index) { - return data.uint8[index]; -} -template <> -bool Value(const TfLitePtrUnion& data, int index) { - return data.b[index]; -} - -template -void SetTensorData(const std::vector& values, TfLitePtrUnion* data) { - T* input_ptr = reinterpret_cast(data->raw); - for (const T& v : values) { - *input_ptr = v; - ++input_ptr; - } -} - -} // namespace - -class TfLiteDriver::Expectation { - public: - Expectation() { - data_.raw = nullptr; - num_elements_ = 0; - } - ~Expectation() { delete[] data_.raw; } - template - void SetData(const string& csv_values) { - const auto& values = testing::Split(csv_values, ","); - num_elements_ = values.size(); - data_.raw = new char[num_elements_ * sizeof(T)]; - SetTensorData(values, &data_); - } - - bool Check(bool verbose, const TfLiteTensor& tensor) { - switch (tensor.type) { - case kTfLiteFloat32: - return TypedCheck(verbose, tensor); - case kTfLiteInt32: - return TypedCheck(verbose, tensor); - case kTfLiteInt64: - return TypedCheck(verbose, tensor); - case kTfLiteUInt8: - return TypedCheck(verbose, tensor); - case kTfLiteBool: - return TypedCheck(verbose, tensor); - default: - fprintf(stderr, "Unsupported type %d in Check\n", tensor.type); - return false; - } - } - - private: - template - bool TypedCheck(bool verbose, const TfLiteTensor& tensor) { - // TODO(ahentz): must find a way to configure the tolerance. - constexpr double kRelativeThreshold = 1e-2f; - constexpr double kAbsoluteThreshold = 1e-4f; - - size_t tensor_size = tensor.bytes / sizeof(T); - - if (tensor_size != num_elements_) { - std::cerr << "Expected a tensor with " << num_elements_ - << " elements, got " << tensor_size << std::endl; - return false; - } - - bool good_output = true; - for (int i = 0; i < tensor_size; ++i) { - float computed = Value(tensor.data, i); - float reference = Value(data_, i); - float diff = std::abs(computed - reference); - bool error_is_large = false; - // For very small numbers, try absolute error, otherwise go with - // relative. - if (std::abs(reference) < kRelativeThreshold) { - error_is_large = (diff > kAbsoluteThreshold); - } else { - error_is_large = (diff > kRelativeThreshold * std::abs(reference)); - } - if (error_is_large) { - good_output = false; - if (verbose) { - std::cerr << " index " << i << ": got " << computed - << ", but expected " << reference << std::endl; - } - } - } - return good_output; - } - - TfLitePtrUnion data_; - size_t num_elements_; -}; - -TfLiteDriver::TfLiteDriver(bool use_nnapi, const string& delegate_name) - : use_nnapi_(use_nnapi) { - if (delegate_name == "FLEX") { - delegate_ = FlexDelegate::Create(); - } -} - -TfLiteDriver::~TfLiteDriver() {} - -void TfLiteDriver::AllocateTensors() { - if (must_allocate_tensors_) { - if (interpreter_->AllocateTensors() != kTfLiteOk) { - Invalidate("Failed to allocate tensors"); - return; - } - ResetLSTMStateTensors(); - must_allocate_tensors_ = false; - } -} - -void TfLiteDriver::LoadModel(const string& bin_file_path) { - if (!IsValid()) return; - - model_ = FlatBufferModel::BuildFromFile(GetFullPath(bin_file_path).c_str()); - if (!model_) { - Invalidate("Failed to mmap model " + bin_file_path); - return; - } - ops::builtin::BuiltinOpResolver builtins; - InterpreterBuilder(*model_, builtins)(&interpreter_); - if (!interpreter_) { - Invalidate("Failed build interpreter"); - return; - } - interpreter_->UseNNAPI(use_nnapi_); - - if (delegate_) { - if (interpreter_->ModifyGraphWithDelegate(delegate_.get(), - /*allow_dynamic_tensors=*/true) != - kTfLiteOk) { - Invalidate("Unable to the build graph using the delegate"); - return; - } - } - - must_allocate_tensors_ = true; -} - -void TfLiteDriver::ResetTensor(int id) { - if (!IsValid()) return; - auto* tensor = interpreter_->tensor(id); - memset(tensor->data.raw, 0, tensor->bytes); -} - -void TfLiteDriver::ReshapeTensor(int id, const string& csv_values) { - if (!IsValid()) return; - if (interpreter_->ResizeInputTensor( - id, testing::Split(csv_values, ",")) != kTfLiteOk) { - Invalidate("Failed to resize input tensor " + std::to_string(id)); - return; - } - must_allocate_tensors_ = true; -} - -void TfLiteDriver::SetInput(int id, const string& csv_values) { - if (!IsValid()) return; - auto* tensor = interpreter_->tensor(id); - switch (tensor->type) { - case kTfLiteFloat32: { - const auto& values = testing::Split(csv_values, ","); - if (!CheckSizes(tensor->bytes, values.size())) return; - SetTensorData(values, &tensor->data); - break; - } - case kTfLiteInt32: { - const auto& values = testing::Split(csv_values, ","); - if (!CheckSizes(tensor->bytes, values.size())) return; - SetTensorData(values, &tensor->data); - break; - } - case kTfLiteInt64: { - const auto& values = testing::Split(csv_values, ","); - if (!CheckSizes(tensor->bytes, values.size())) return; - SetTensorData(values, &tensor->data); - break; - } - case kTfLiteUInt8: { - const auto& values = testing::Split(csv_values, ","); - if (!CheckSizes(tensor->bytes, values.size())) return; - SetTensorData(values, &tensor->data); - break; - } - case kTfLiteBool: { - const auto& values = testing::Split(csv_values, ","); - if (!CheckSizes(tensor->bytes, values.size())) return; - SetTensorData(values, &tensor->data); - break; - } - default: - fprintf(stderr, "Unsupported type %d in SetInput\n", tensor->type); - Invalidate("Unsupported tensor data type"); - return; - } -} - -void TfLiteDriver::SetExpectation(int id, const string& csv_values) { - if (!IsValid()) return; - auto* tensor = interpreter_->tensor(id); - if (expected_output_.count(id) != 0) { - fprintf(stderr, "Overridden expectation for tensor %d\n", id); - Invalidate("Overridden expectation"); - } - expected_output_[id].reset(new Expectation); - switch (tensor->type) { - case kTfLiteFloat32: - expected_output_[id]->SetData(csv_values); - break; - case kTfLiteInt32: - expected_output_[id]->SetData(csv_values); - break; - case kTfLiteInt64: - expected_output_[id]->SetData(csv_values); - break; - case kTfLiteUInt8: - expected_output_[id]->SetData(csv_values); - break; - case kTfLiteBool: - expected_output_[id]->SetData(csv_values); - break; - default: - fprintf(stderr, "Unsupported type %d in SetExpectation\n", tensor->type); - Invalidate("Unsupported tensor data type"); - return; - } -} - -void TfLiteDriver::Invoke() { - if (!IsValid()) return; - if (interpreter_->Invoke() != kTfLiteOk) { - Invalidate("Failed to invoke interpreter"); - } -} - -bool TfLiteDriver::CheckResults() { - if (!IsValid()) return false; - bool success = true; - for (const auto& p : expected_output_) { - int id = p.first; - auto* tensor = interpreter_->tensor(id); - if (!p.second->Check(/*verbose=*/false, *tensor)) { - // Do not invalidate anything here. Instead, simply output the - // differences and return false. Invalidating would prevent all - // subsequent invocations from running.. - std::cerr << "There were errors in invocation '" << GetInvocationId() - << "', output tensor '" << id << "':" << std::endl; - p.second->Check(/*verbose=*/true, *tensor); - success = false; - SetOverallSuccess(false); - } - } - expected_output_.clear(); - return success; -} - -void TfLiteDriver::ResetLSTMStateTensors() { - interpreter_->ResetVariableTensors(); -} - -} // namespace testing -} // namespace tflite diff --git a/tensorflow/contrib/lite/testing/tflite_driver.h b/tensorflow/contrib/lite/testing/tflite_driver.h deleted file mode 100644 index dc2a4e58773a9e069aa4420907c068039252c418..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/testing/tflite_driver.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_CONTRIB_LITE_TESTING_TFLITE_DRIVER_H_ -#define TENSORFLOW_CONTRIB_LITE_TESTING_TFLITE_DRIVER_H_ - -#include - -#include "tensorflow/contrib/lite/delegates/flex/delegate.h" -#include "tensorflow/contrib/lite/interpreter.h" -#include "tensorflow/contrib/lite/kernels/register.h" -#include "tensorflow/contrib/lite/model.h" -#include "tensorflow/contrib/lite/testing/test_runner.h" - -namespace tflite { -namespace testing { - -// A test runner that feeds inputs into TF Lite and verifies its outputs. -class TfLiteDriver : public TestRunner { - public: - explicit TfLiteDriver(bool use_nnapi, const string& delegate = ""); - ~TfLiteDriver() override; - - void LoadModel(const string& bin_file_path) override; - const std::vector& GetInputs() override { - return interpreter_->inputs(); - } - const std::vector& GetOutputs() override { - return interpreter_->outputs(); - } - void ReshapeTensor(int id, const string& csv_values) override; - void AllocateTensors() override; - void ResetTensor(int id) override; - void SetInput(int id, const string& csv_values) override; - void SetExpectation(int id, const string& csv_values) override; - void Invoke() override; - bool CheckResults() override; - string ReadOutput(int id) override { return "no-op"; } - - private: - void ResetLSTMStateTensors(); - - class Expectation; - - std::unique_ptr delegate_; - bool use_nnapi_ = false; - std::unique_ptr model_; - std::unique_ptr interpreter_; - std::map> expected_output_; - bool must_allocate_tensors_ = true; -}; - -} // namespace testing -} // namespace tflite - -#endif // TENSORFLOW_CONTRIB_LITE_TESTING_TFLITE_DRIVER_H_ diff --git a/tensorflow/contrib/lite/testing/tflite_driver_test.cc b/tensorflow/contrib/lite/testing/tflite_driver_test.cc deleted file mode 100644 index 37010c468f250fdf4ef958b23a38aa38b7a533db..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/testing/tflite_driver_test.cc +++ /dev/null @@ -1,61 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include "tensorflow/contrib/lite/testing/tflite_driver.h" - -#include -#include - -namespace tflite { -namespace testing { -namespace { - -using ::testing::ElementsAre; - -TEST(TfliteDriverTest, SimpleTest) { - std::unique_ptr runner(new TfLiteDriver(/*use_nnapi=*/false)); - - runner->SetModelBaseDir("tensorflow/contrib/lite"); - runner->LoadModel("testdata/multi_add.bin"); - ASSERT_TRUE(runner->IsValid()); - - ASSERT_THAT(runner->GetInputs(), ElementsAre(0, 1, 2, 3)); - ASSERT_THAT(runner->GetOutputs(), ElementsAre(5, 6)); - - for (int i : {0, 1, 2, 3}) { - runner->ReshapeTensor(i, "1,2,2,1"); - } - ASSERT_TRUE(runner->IsValid()); - - runner->AllocateTensors(); - - runner->SetInput(0, "0.1,0.2,0.3,0.4"); - runner->SetInput(1, "0.001,0.002,0.003,0.004"); - runner->SetInput(2, "0.001,0.002,0.003,0.004"); - runner->SetInput(3, "0.01,0.02,0.03,0.04"); - - runner->ResetTensor(2); - - runner->SetExpectation(5, "0.101,0.202,0.303,0.404"); - runner->SetExpectation(6, "0.011,0.022,0.033,0.044"); - - runner->Invoke(); - ASSERT_TRUE(runner->IsValid()); - - ASSERT_TRUE(runner->CheckResults()); -} - -} // namespace -} // namespace testing -} // namespace tflite diff --git a/tensorflow/contrib/lite/testing/util.h b/tensorflow/contrib/lite/testing/util.h deleted file mode 100644 index 925791d3908dc569a05f7c6b632448266c08c48f..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/testing/util.h +++ /dev/null @@ -1,59 +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_CONTRIB_LITE_TESTING_UTIL_H_ -#define TENSORFLOW_CONTRIB_LITE_TESTING_UTIL_H_ - -#include - -#include "tensorflow/contrib/lite/core/api/error_reporter.h" -#include "tensorflow/contrib/lite/string.h" - -namespace tflite { - -// An ErrorReporter that collects error message in a string, in addition -// to printing to stderr. -class TestErrorReporter : public ErrorReporter { - public: - int Report(const char* format, va_list args) override { - char buffer[1024]; - int size = vsnprintf(buffer, sizeof(buffer), format, args); - fprintf(stderr, "%s", buffer); - error_messages_ += buffer; - num_calls_++; - return size; - } - - void Reset() { - num_calls_ = 0; - error_messages_.clear(); - } - - int num_calls() const { return num_calls_; } - const string& error_messages() const { return error_messages_; } - - private: - int num_calls_ = 0; - string error_messages_; -}; - -inline void LogToStderr() { -#ifdef PLATFORM_GOOGLE - FLAGS_logtostderr = true; -#endif -} - -} // namespace tflite - -#endif // TENSORFLOW_CONTRIB_LITE_TESTING_UTIL_H_ diff --git a/tensorflow/contrib/lite/toco/BUILD b/tensorflow/contrib/lite/toco/BUILD deleted file mode 100644 index 96b88b60fc650981bc880c309a38836d694b3ad0..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/toco/BUILD +++ /dev/null @@ -1,430 +0,0 @@ -package(default_visibility = ["//visibility:public"]) - -licenses(["notice"]) # Apache 2.0 - -load( - "//tensorflow/core:platform/default/build_config.bzl", - "tf_proto_library_cc", - "tf_proto_library_py", -) -load( - "//tensorflow:tensorflow.bzl", - "tf_cc_binary", - "tf_cc_test", - "tf_copts", -) - -tf_proto_library_cc( - name = "types_proto", - srcs = ["types.proto"], - visibility = ["//visibility:public"], -) - -tf_proto_library_cc( - name = "toco_flags_proto", - srcs = ["toco_flags.proto"], - protodeps = [":types_proto"], - visibility = ["//visibility:public"], -) - -tf_proto_library_cc( - name = "model_flags_proto", - srcs = ["model_flags.proto"], - protodeps = [":types_proto"], - visibility = ["//visibility:public"], -) - -tf_proto_library_py( - name = "types_proto", - srcs = [ - "types.proto", - ], - visibility = ["//visibility:public"], -) - -tf_proto_library_py( - name = "toco_flags_proto", - srcs = [ - "toco_flags.proto", - ], - protodeps = [":types_proto"], - visibility = ["//visibility:public"], -) - -tf_proto_library_py( - name = "model_flags_proto", - srcs = [ - "model_flags.proto", - ], - protodeps = [":types_proto"], - visibility = ["//visibility:public"], -) - -cc_library( - name = "tensorflow_core_cc_protos_all", - deps = ["//tensorflow/core:protos_all_cc"], -) - -cc_library( - name = "runtime", - hdrs = [ - "runtime/common.h", - "runtime/types.h", - ], - linkstatic = 1, - visibility = ["//visibility:public"], - deps = [ - "//tensorflow/contrib/lite/kernels/internal:reference_base", - "//tensorflow/contrib/lite/kernels/internal:types", - ], -) - -# :model offers the core data structures representing a model (a.k.a. "graph") -# for tooling purposes (not needed at inference runtime). -# That includes the top-level Model structure, and the lower-level Operator, -# Array, Buffer structures, etc. -cc_library( - name = "model", - hdrs = [ - "model.h", - ], - visibility = ["//visibility:public"], - deps = [ - ":model_flags_proto_cc", - ":runtime", - ":toco_port", - "//tensorflow/core:lib", - "@com_google_absl//absl/types:optional", - ], -) - -cc_library( - name = "toco_graphviz_dump_options", - srcs = [ - "toco_graphviz_dump_options.cc", - ], - hdrs = [ - "toco_graphviz_dump_options.h", - ], - visibility = ["//visibility:public"], -) - -cc_library( - name = "toco_cmdline_flags", - srcs = [ - "toco_cmdline_flags.cc", - ], - hdrs = [ - "toco_cmdline_flags.h", - ], - visibility = ["//visibility:public"], - deps = [ - ":model_cmdline_flags", - ":toco_flags_proto_cc", - ":toco_port", - ":types_proto_cc", - "//tensorflow/core:framework_internal", - "//tensorflow/core:lib", - "@com_google_absl//absl/strings", - "@com_google_absl//absl/types:optional", - ], -) - -cc_library( - name = "model_cmdline_flags", - srcs = [ - "model_cmdline_flags.cc", - ], - hdrs = [ - "args.h", - "model_cmdline_flags.h", - ], - visibility = ["//visibility:public"], - deps = [ - ":model_flags_proto_cc", - ":toco_graphviz_dump_options", - ":toco_port", - ":types_proto_cc", - "//tensorflow/core:framework_internal", - "//tensorflow/core:lib", - "@com_google_absl//absl/strings", - ], -) - -cc_library( - name = "toco_port", - srcs = [ - "toco_port.cc", - ], - hdrs = [ - "format_port.h", - "toco_port.h", - "toco_types.h", - ], - deps = [ - # Placeholder for internal file dependency. - "@protobuf_archive//:protobuf_headers", - "//tensorflow/core:framework_lite", - "//tensorflow/core:lib", - "//tensorflow/core:lib_internal", - ], -) - -cc_library( - name = "graph_transformations", - srcs = [ - "graph_transformations/convert_expanddims_to_reshape.cc", - "graph_transformations/convert_pure_conv_to_depthwise.cc", - "graph_transformations/convert_reorder_axes.cc", - "graph_transformations/convert_squeeze_to_reshape.cc", - "graph_transformations/convert_trivial_addn_to_add.cc", - "graph_transformations/convert_trivial_pack_to_reshape.cc", - "graph_transformations/convert_trivial_tile_to_concat.cc", - "graph_transformations/convert_trivial_transpose_to_reshape.cc", - "graph_transformations/create_im2col_arrays.cc", - "graph_transformations/dequantize.cc", - "graph_transformations/drop_fake_quant.cc", - "graph_transformations/drop_im2col_arrays.cc", - "graph_transformations/ensure_bias_vectors.cc", - "graph_transformations/ensure_uint8_weights_safe_for_fast_int8_kernels.cc", - "graph_transformations/fuse_activation_functions.cc", - "graph_transformations/fuse_binary_into_following_affine.cc", - "graph_transformations/fuse_binary_into_preceding_affine.cc", - "graph_transformations/fuse_broadcast_into_following_binary.cc", - "graph_transformations/graph_transformations.cc", - "graph_transformations/hardcode_min_max.cc", - "graph_transformations/identify_dilated_conv.cc", - "graph_transformations/identify_l2_normalization.cc", - "graph_transformations/identify_l2_pool.cc", - "graph_transformations/identify_lstm.cc", - "graph_transformations/identify_lstm_merge_inputs.cc", - "graph_transformations/identify_lstm_split_inputs.cc", - "graph_transformations/identify_prelu.cc", - "graph_transformations/identify_relu1.cc", - "graph_transformations/lstm_utils.cc", - "graph_transformations/make_initial_dequantize_operator.cc", - "graph_transformations/merge_reshape_into_preceding_transpose.cc", - "graph_transformations/move_binary_operator_before_reshape.cc", - "graph_transformations/propagate_activation_function_into_constants.cc", - "graph_transformations/propagate_array_data_types.cc", - "graph_transformations/propagate_default_min_max.cc", - "graph_transformations/propagate_fake_quant_num_bits.cc", - "graph_transformations/propagate_fixed_sizes.cc", - "graph_transformations/quantization_util.cc", - "graph_transformations/quantization_util.h", - "graph_transformations/quantize.cc", - "graph_transformations/read_array_minmax_and_narrow_range_from_fake_quant.cc", - "graph_transformations/remove_final_dequantize_op.cc", - "graph_transformations/remove_tensorflow_assert.cc", - "graph_transformations/remove_tensorflow_identity.cc", - "graph_transformations/remove_trivial_binary.cc", - "graph_transformations/remove_trivial_concatenation.cc", - "graph_transformations/remove_trivial_concatenation_input.cc", - "graph_transformations/remove_trivial_fake_quant.cc", - "graph_transformations/remove_trivial_passthrough.cc", - "graph_transformations/remove_trivial_passthrough.h", - "graph_transformations/remove_trivial_quantized_activation_func.cc", - "graph_transformations/remove_trivial_quantized_min_max.cc", - "graph_transformations/remove_trivial_reshape.cc", - "graph_transformations/remove_trivial_slice.cc", - "graph_transformations/remove_unused_op.cc", - "graph_transformations/reorder_elementwise_unary.cc", - "graph_transformations/reorder_reshape_transpose.cc", - "graph_transformations/resolve_batch_normalization.cc", - "graph_transformations/resolve_batch_to_space_nd_attributes.cc", - "graph_transformations/resolve_constant_binary.cc", - "graph_transformations/resolve_constant_concatenation.cc", - "graph_transformations/resolve_constant_fake_quant.cc", - "graph_transformations/resolve_constant_fill.cc", - "graph_transformations/resolve_constant_gather.cc", - "graph_transformations/resolve_constant_pack.cc", - "graph_transformations/resolve_constant_random_uniform.cc", - "graph_transformations/resolve_constant_range.cc", - "graph_transformations/resolve_constant_reshape.cc", - "graph_transformations/resolve_constant_select.cc", - "graph_transformations/resolve_constant_shape_or_rank.cc", - "graph_transformations/resolve_constant_slice.cc", - "graph_transformations/resolve_constant_strided_slice.cc", - "graph_transformations/resolve_constant_tile.cc", - "graph_transformations/resolve_constant_transpose.cc", - "graph_transformations/resolve_constant_unary.cc", - "graph_transformations/resolve_fake_quant_args_from_vars.cc", - "graph_transformations/resolve_gather_attributes.cc", - "graph_transformations/resolve_multiply_by_zero.cc", - "graph_transformations/resolve_pad_attributes.cc", - "graph_transformations/resolve_padv2_attributes.cc", - "graph_transformations/resolve_reduce_attributes.cc", - "graph_transformations/resolve_reorder_axes.cc", - "graph_transformations/resolve_reshape_attributes.cc", - "graph_transformations/resolve_slice_attributes.cc", - "graph_transformations/resolve_space_to_batch_nd_attributes.cc", - "graph_transformations/resolve_squeeze_attributes.cc", - "graph_transformations/resolve_strided_slice_attributes.cc", - "graph_transformations/resolve_tensorflow_concat.cc", - "graph_transformations/resolve_tensorflow_matmul.cc", - "graph_transformations/resolve_tensorflow_merge.cc", - "graph_transformations/resolve_tensorflow_switch.cc", - "graph_transformations/resolve_transpose_attributes.cc", - "graph_transformations/shuffle_fc_weights.cc", - "graph_transformations/unfuse_activation_functions.cc", - "graph_transformations/unpartition_embedding_lookup.cc", - "graph_transformations/unroll_batch_matmul.cc", - ], - hdrs = [ - "graph_transformations/graph_transformations.h", - "graph_transformations/lstm_utils.h", - ], - visibility = ["//visibility:public"], - deps = [ - ":model", - ":model_flags_proto_cc", - ":runtime", - ":toco_port", - ":tooling_util", - "//tensorflow/contrib/lite/kernels/internal:quantization_util", - "//tensorflow/contrib/lite/kernels/internal:strided_slice_logic", - "//tensorflow/core:lib", - "@com_google_absl//absl/memory", - "@com_google_absl//absl/strings", - ], -) - -# :toco_tooling is the library providing the offline tooling functionality -# exposed by the :toco command-line tool. -cc_library( - name = "toco_tooling", - srcs = [ - "allocate_transient_arrays.cc", - "export_tensorflow.cc", - "import_tensorflow.cc", - "tensorflow_util.cc", - "toco_tooling.cc", - ], - hdrs = [ - "allocate_transient_arrays.h", - "export_tensorflow.h", - "import_tensorflow.h", - "tensorflow_util.h", - "toco_tooling.h", - ], - copts = tf_copts() + select({ - "//tensorflow:darwin": ["-DTOCO_SUPPORT_PORTABLE_PROTOS=0"], - "//conditions:default": [], - }), - visibility = ["//visibility:public"], - deps = [ - ":graph_transformations", - ":model", - ":model_flags_proto_cc", - ":types_proto_cc", - ":runtime", - ":toco_graphviz_dump_options", - ":toco_flags_proto_cc", - ":toco_port", - ":tooling_util", - "@protobuf_archive//:protobuf_headers", - "@com_google_absl//absl/memory", - "@com_google_absl//absl/strings", - "//tensorflow/contrib/lite/toco/tensorflow_graph_matching:resolve_cluster", - "//tensorflow/contrib/lite/toco/tflite:export", - "//tensorflow/contrib/lite/toco/tflite:import", - "//tensorflow/core:core_cpu_lib", - "//tensorflow/core:framework", - "//tensorflow/core:lib", - "//tensorflow/core:protos_all_cc", - ] + select({ - # Placeholder for internal darwin rule. - "//conditions:default": [], - }), -) - -tf_cc_test( - name = "import_tensorflow_test", - srcs = ["import_tensorflow_test.cc"], - tags = ["no_oss"], - deps = [ - ":toco_tooling", - "//tensorflow/core:framework", - "//tensorflow/core:graph", - "//tensorflow/core:lib", - "//tensorflow/core:ops", - "//tensorflow/core:protos_all_cc", - "@com_google_googletest//:gtest_main", - ], -) - -cc_library( - name = "tooling_util", - srcs = [ - "dump_graphviz.cc", - "tooling_util.cc", - ], - hdrs = [ - "dump_graphviz.h", - "tooling_util.h", - ], - copts = tf_copts(), - visibility = ["//visibility:public"], - deps = [ - ":model", - ":model_flags_proto_cc", - ":runtime", - ":toco_flags_proto_cc", - ":toco_graphviz_dump_options", - ":toco_port", - ":types_proto_cc", - "//tensorflow/contrib/lite/kernels/internal:types", - "//tensorflow/core:lib", - "@com_google_absl//absl/strings", - "@com_googlesource_code_re2//:re2", - "@protobuf_archive//:protobuf_headers", - ], -) - -tf_cc_test( - name = "tooling_util_test", - srcs = ["tooling_util_test.cc"], - tags = ["no_oss"], - deps = [ - ":model", - ":tooling_util", - "//tensorflow/core:lib", - "@com_google_googletest//:gtest_main", - ], -) - -# :toco is the main public command-line tool exposing the functionality -# of the :toco_tooling library. -tf_cc_binary( - name = "toco", - srcs = ["toco.cc"], - visibility = ["//visibility:public"], - deps = [ - ":model", - ":model_cmdline_flags", - ":model_flags_proto_cc", - ":toco_cmdline_flags", - ":toco_flags_proto_cc", - ":toco_port", - ":toco_tooling", - ":types_proto_cc", - "@com_google_absl//absl/strings", - "//tensorflow/core:lib", - # We cannot embed the core:ops dependency directly into :toco_tooling as - # it can conflict with downstream deps when toco is used as a library. - "//tensorflow/core:ops", - ], -) - -tf_cc_test( - name = "toco_port_test", - srcs = ["toco_port_test.cc"], - data = [ - "toco_port_test.cc", - ], - tags = ["no_oss"], - deps = [ - ":toco_port", - "@com_google_googletest//:gtest_main", - ], -) diff --git a/tensorflow/contrib/lite/toco/README.md b/tensorflow/contrib/lite/toco/README.md deleted file mode 100644 index 91f6f618a376ff4df7c51dfd285152229f4757cc..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/toco/README.md +++ /dev/null @@ -1,29 +0,0 @@ -# TensorFlow Lite Converter - -The TensorFlow Lite Converter converts TensorFlow graphs into -TensorFlow Lite graphs. There are additional usages that are also detailed in -the usage documentation. - -## Usage documentation - -Usage information is given in these documents: - -* [Command-line glossary](g3doc/cmdline_reference.md) -* [Command-line examples](g3doc/cmdline_examples.md) -* [Python API examples](g3doc/python_api.md) - -## Where the converter fits in the TensorFlow landscape - -Once an application developer has a trained TensorFlow model, the TensorFlow -Lite Converter will accept -that model and generate a TensorFlow Lite -[FlatBuffer](https://google.github.io/flatbuffers/) file. The converter currently supports -[SavedModels](https://www.tensorflow.org/guide/saved_model#using_savedmodel_with_estimators), -frozen graphs (models generated via -[freeze_graph.py](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/tools/freeze_graph.py)), -and `tf.Keras` model files. The TensorFlow Lite FlatBuffer file can be shipped -to client devices, generally mobile devices, where the TensorFlow Lite -interpreter handles them on-device. This flow is represented in the diagram -below. - -![drawing](g3doc/toco_landscape.svg) diff --git a/tensorflow/contrib/lite/toco/g3doc/README.md b/tensorflow/contrib/lite/toco/g3doc/README.md deleted file mode 100644 index 2153b6cc6360a7a0e0375600c83b0c0945d3b326..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/toco/g3doc/README.md +++ /dev/null @@ -1,3 +0,0 @@ -# TOCO - -These files have moved to [../../g3doc/tflite_convert](../../g3doc/tflite_convert) diff --git a/tensorflow/contrib/lite/toco/graph_transformations/dequantize.cc b/tensorflow/contrib/lite/toco/graph_transformations/dequantize.cc deleted file mode 100644 index 2119174950b1ade67c66a47e65b9808c216a7c54..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/toco/graph_transformations/dequantize.cc +++ /dev/null @@ -1,230 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include -#include -#include -#include - -#include "tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h" -#include "tensorflow/contrib/lite/toco/graph_transformations/remove_trivial_passthrough.h" -#include "tensorflow/contrib/lite/toco/model.h" -#include "tensorflow/contrib/lite/toco/tooling_util.h" -#include "tensorflow/core/platform/logging.h" - -namespace toco { - -namespace { - -template -void DequantizeBuffer(Array* array) { - const auto old_data = array->GetBuffer().data; - array->buffer = nullptr; - array->data_type = ArrayDataType::kFloat; - auto& new_data = array->GetMutableBuffer().data; - new_data.resize(old_data.size()); - const auto& qparams = array->GetQuantizationParams(); - for (int i = 0; i < old_data.size(); i++) { - new_data[i] = qparams.scale * (old_data[i] - qparams.zero_point); - } -} - -std::vector>::iterator FindFirstOpWithInput( - Model* model, const string& array_name) { - for (auto it = model->operators.begin(); it != model->operators.end(); ++it) { - for (const auto& input : it->get()->inputs) { - if (input == array_name) { - return it; - } - } - } - return model->operators.end(); -} - -void ClearArrayQuantizationParams(const string& array_name, Model* model) { - auto* array = &model->GetArray(array_name); - CHECK(array->quantization_params); - for (auto& input_array : *model->flags.mutable_input_arrays()) { - if (input_array.name() == array_name) { - auto& qparams = *array->quantization_params; - const double new_std_value = 1. / qparams.scale; - const double new_mean_value = qparams.zero_point; - if (input_array.has_std_value()) { - CHECK_LE(std::abs(new_std_value - input_array.std_value()), 0.001); - } else { - input_array.set_std_value(new_std_value); - } - if (input_array.has_mean_value()) { - CHECK_LE(std::abs(new_mean_value - input_array.mean_value()), 0.001); - } else { - input_array.set_mean_value(new_mean_value); - } - } - } - array->quantization_params = nullptr; -} - -bool DequantizeArray(const string& array_name, - GraphTransformation* transformation, Model* model) { - auto* array = &model->GetArray(array_name); - if (!array->quantization_params) { - return false; - } - transformation->AddMessageF("Dequantizing array: %s", array_name); - - // Dequantize any buffer - if (array->buffer) { - if (array->data_type == ArrayDataType::kUint8) { - DequantizeBuffer(array); - } else if (array->data_type == ArrayDataType::kInt32) { - DequantizeBuffer(array); - } else { - LOG(FATAL) << "Unhandled data type"; - } - CHECK(array->data_type == ArrayDataType::kFloat); - CHECK(array->buffer->type == ArrayDataType::kFloat); - - // Clear quantization params, officially makes this a non-quantized array. - ClearArrayQuantizationParams(array_name, model); - return true; - } else { - array->data_type = ArrayDataType::kFloat; - } - - // Clear quantization params, officially makes this a non-quantized array. - ClearArrayQuantizationParams(array_name, model); - - if (array->buffer) { - return true; - } - - auto* op_outputting_array = GetOpWithOutput(*model, array_name); - if (op_outputting_array) { - if (op_outputting_array->type == OperatorType::kReshape) { - return true; - } - } - - // If there was no minmax info, we can return now. Indeed, - // the below only serves to create a FakeQuant node, but some arrays are - // quantized without MinMax (see the CHECK above) and that corresponds to - // places where a FakeQuant node is actually not wanted, because the - // quantization params are meant to be inferred in another way (e.g. bias - // vector for a Conv op, see their special-casing in quantize.cc). - if (!array->minmax) { - return true; - } - - // Determine whether to insert a FakeQuant before or after - // this array. - bool must_insert_fakequant_before = false; - bool must_insert_fakequant_after = false; - if (IsInputArray(*model, array_name)) { - must_insert_fakequant_after = true; - } - for (const string& output_array : model->flags.output_arrays()) { - if (array_name == output_array) { - must_insert_fakequant_before = true; - } - } - for (const auto& rnn_state : model->flags.rnn_states()) { - if (array_name == rnn_state.state_array()) { - must_insert_fakequant_after = true; - } - if (array_name == rnn_state.back_edge_source_array()) { - must_insert_fakequant_before = true; - } - } - CHECK(!(must_insert_fakequant_before && must_insert_fakequant_after)); - - // Create and insert the FakeQuant node - auto* fakequant_op = new FakeQuantOperator; - model->operators.emplace(FindFirstOpWithInput(model, array_name), - fakequant_op); - const string& new_array_name = AvailableArrayName(*model, array_name); - auto& new_array = model->GetOrCreateArray(new_array_name); - new_array.data_type = ArrayDataType::kFloat; - new_array.copy_shape(array->shape()); - new_array.GetOrCreateMinMax() = array->GetMinMax(); - fakequant_op->minmax.reset(new MinMax); - *fakequant_op->minmax = array->GetMinMax(); - fakequant_op->narrow_range = array->narrow_range; - if (must_insert_fakequant_before) { - for (const auto& op : model->operators) { - for (string& output : op->outputs) { - if (output == array_name) { - output = new_array_name; - } - } - } - fakequant_op->inputs = {new_array_name}; - fakequant_op->outputs = {array_name}; - } else { - for (const auto& op : model->operators) { - for (string& input : op->inputs) { - if (input == array_name) { - input = new_array_name; - } - } - } - fakequant_op->inputs = {array_name}; - fakequant_op->outputs = {new_array_name}; - } - return true; -} - -} // namespace - -::tensorflow::Status Dequantize::Run(Model* model, std::size_t op_index, - bool* modified) { - *modified = false; - const auto op_it = model->operators.begin() + op_index; - auto* op = op_it->get(); - - if (op->type == OperatorType::kDequantize) { - auto& input_array = model->GetArray(op->inputs[0]); - if (input_array.data_type == ArrayDataType::kFloat) { - return ::tensorflow::Status::OK(); - } - if (input_array.final_data_type != ArrayDataType::kFloat) { - return ::tensorflow::Status::OK(); - } - input_array.data_type = ArrayDataType::kFloat; - input_array.quantization_params = nullptr; - auto& output_array = model->GetArray(op->outputs[0]); - output_array.data_type = ArrayDataType::kFloat; - output_array.quantization_params = nullptr; - *modified = RemoveTrivialPassthroughOp(this, model, op_index); - return ::tensorflow::Status::OK(); - } - - std::vector arrays; - for (const string& input : op->inputs) { - arrays.push_back(input); - } - for (const string& output : op->outputs) { - arrays.push_back(output); - } - bool changed = false; - for (const string& array : arrays) { - if (!model->IsOptionalArray(array)) { - changed |= DequantizeArray(array, this, model); - } - } - - *modified = changed; - return ::tensorflow::Status::OK(); -} - -} // namespace toco diff --git a/tensorflow/contrib/lite/toco/graph_transformations/quantization_util.cc b/tensorflow/contrib/lite/toco/graph_transformations/quantization_util.cc deleted file mode 100644 index 44733391f5a1d9ebf9a24f4f31b425a35354e1fc..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/toco/graph_transformations/quantization_util.cc +++ /dev/null @@ -1,276 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include - -#include "tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h" -#include "tensorflow/contrib/lite/toco/graph_transformations/quantization_util.h" -#include "tensorflow/contrib/lite/toco/model.h" -#include "tensorflow/contrib/lite/toco/tooling_util.h" -#include "tensorflow/core/platform/logging.h" - -namespace toco { - -bool InferQuantizedDataTypeFromFakeQuant( - const FakeQuantOperator& op, ArrayDataType* out_quantized_data_type) { - if (op.num_bits <= 8) { - *out_quantized_data_type = ArrayDataType::kUint8; - return true; - } else if (op.num_bits <= 16) { - *out_quantized_data_type = ArrayDataType::kInt16; - return true; - } else { - *out_quantized_data_type = ArrayDataType::kNone; - return false; - } -} - -bool GetQuantizedDataTypeNumericalRange(ArrayDataType data_type, - double* out_min_value, - double* out_max_value) { - switch (data_type) { - case ArrayDataType::kUint8: - *out_min_value = 0; - *out_max_value = 255; - return true; - case ArrayDataType::kInt16: - *out_min_value = -32768; - *out_max_value = 32767; - return true; - default: - return false; - } -} - -ArrayDataType GetQuantizedDataType(const Array& array, - ArrayDataType default_type) { - switch (array.final_data_type) { - case ArrayDataType::kInt8: - case ArrayDataType::kUint8: - case ArrayDataType::kInt16: - case ArrayDataType::kUint16: - case ArrayDataType::kInt32: - case ArrayDataType::kUint32: - case ArrayDataType::kInt64: - case ArrayDataType::kUint64: - return array.final_data_type; - case ArrayDataType::kFloat: - case ArrayDataType::kNone: - return default_type; - default: - LOG(FATAL) << "Unhandled final quantization type " - << static_cast(array.final_data_type); - } -} - -template -void ChooseQuantizationParamsForArrayAndQuantizedDataType( - const Array& array, QuantizationParams* quantization_params) { - *quantization_params = ::tflite::ChooseQuantizationParams>( - array.minmax->min, array.minmax->max, array.narrow_range); -} - -void ChooseQuantizationParamsForArrayAndQuantizedDataType( - const Array& array, ArrayDataType quantized_data_type, - QuantizationParams* quantization_params) { - switch (quantized_data_type) { - case ArrayDataType::kInt8: - ChooseQuantizationParamsForArrayAndQuantizedDataType< - ArrayDataType::kInt8>(array, quantization_params); - break; - case ArrayDataType::kUint8: - ChooseQuantizationParamsForArrayAndQuantizedDataType< - ArrayDataType::kUint8>(array, quantization_params); - break; - case ArrayDataType::kInt16: - ChooseQuantizationParamsForArrayAndQuantizedDataType< - ArrayDataType::kInt16>(array, quantization_params); - break; - case ArrayDataType::kUint16: - ChooseQuantizationParamsForArrayAndQuantizedDataType< - ArrayDataType::kUint16>(array, quantization_params); - break; - case ArrayDataType::kInt32: - ChooseQuantizationParamsForArrayAndQuantizedDataType< - ArrayDataType::kInt32>(array, quantization_params); - break; - case ArrayDataType::kUint32: - ChooseQuantizationParamsForArrayAndQuantizedDataType< - ArrayDataType::kUint32>(array, quantization_params); - break; - case ArrayDataType::kInt64: - ChooseQuantizationParamsForArrayAndQuantizedDataType< - ArrayDataType::kInt64>(array, quantization_params); - break; - case ArrayDataType::kUint64: - ChooseQuantizationParamsForArrayAndQuantizedDataType< - ArrayDataType::kUint64>(array, quantization_params); - break; - case ArrayDataType::kFloat: - case ArrayDataType::kNone: - default: - LOG(FATAL) << "Unhandled final quantization type " - << static_cast(quantized_data_type); - } -} - -namespace { - -template -std::unique_ptr QuantizeBuffer( - const Array& array, const QuantizationParams& quantization_params) { - const GenericBuffer& buffer = *array.buffer; - const auto inverse_scale = 1. / quantization_params.scale; - CHECK(buffer.type == ArrayDataType::kFloat); - const auto& float_buffer = - static_cast&>(buffer); - auto* quantized_buffer = new Buffer; - quantized_buffer->data.resize(float_buffer.data.size()); - for (std::size_t i = 0; i < float_buffer.data.size(); i++) { - const float src_val = float_buffer.data[i]; - double scaled_val; // Astonishingly, using 'float' degrades accuracy just - // enough to make a few tests fail! - if (quantization_params.scale == 0) { - CHECK_EQ(src_val, 0) << "The quantization scale for this array is 0, " - << "so all its values should be 0."; - scaled_val = quantization_params.zero_point; - } else { - scaled_val = quantization_params.zero_point + inverse_scale * src_val; - } - auto integer_val = tflite::SafeCast>(std::round(scaled_val)); - // In addition to its effect on the choice of quantization params upstream - // of here, narrow_range also means nudge the min quantized value by +1, - // so e.g. uint8 values get constrained to [1, 255]. - if (integer_val == std::numeric_limits>::min() && - array.narrow_range) { - integer_val++; - } - quantized_buffer->data[i] = integer_val; - } - return std::unique_ptr(quantized_buffer); -} - -template -void QuantizeArray(GraphTransformation* transformation, Model* model, - const string& name, - const QuantizationParams& quantization_params) { - auto& array = model->GetArray(name); - CHECK(array.data_type == ArrayDataType::kFloat); - CHECK(!array.quantization_params); - array.GetOrCreateQuantizationParams() = quantization_params; - if (array.buffer) { - array.buffer = QuantizeBuffer(array, quantization_params); - } - array.data_type = A; - array.final_data_type = A; - transformation->AddMessageF( - "Quantized array %s to %s zero_point=%g, scale=%g", name, - ArrayDataTypeName(array.data_type), quantization_params.zero_point, - quantization_params.scale); -} - -} // namespace - -void QuantizeArray(GraphTransformation* transformation, Model* model, - const string& name, ArrayDataType quantized_data_type, - const QuantizationParams& quantization_params) { - ArrayDataType adjusted_data_type = quantized_data_type; - auto& array = model->GetArray(name); - if (array.final_data_type == ArrayDataType::kInt16) { - adjusted_data_type = array.final_data_type; - } - - switch (adjusted_data_type) { - case ArrayDataType::kUint8: - return QuantizeArray(transformation, model, name, - quantization_params); - case ArrayDataType::kInt16: - return QuantizeArray(transformation, model, name, - quantization_params); - case ArrayDataType::kInt32: - return QuantizeArray(transformation, model, name, - quantization_params); - default: - LOG(FATAL) << "Unhandled case."; - } -} - -bool IsArrayQuantizedRangeSubset(GraphTransformation* transformation, - const Array& array, double clamp_min, - double clamp_max) { - ArrayDataType quantized_data_type = - GetQuantizedDataType(array, array.data_type); - if (quantized_data_type == ArrayDataType::kNone || - quantized_data_type == ArrayDataType::kFloat) { - // The array is not (or never will be) quantized. - return false; - } - - QuantizationParams quantization_params; - if (!array.quantization_params) { - if (!array.minmax) { - transformation->AddMessageF("No quantization params and no minmax"); - return false; - } else { - // Work around cases where we are asking for this prior to the Quantize - // transformation having added the quantization_params. - ChooseQuantizationParamsForArrayAndQuantizedDataType( - array, quantized_data_type, &quantization_params); - transformation->AddMessageF( - "No quantization params - infering from data type %s with minmax " - "%g,%g as zero_point=%g, scale=%g", - ArrayDataTypeName(quantized_data_type), array.minmax->min, - array.minmax->max, quantization_params.zero_point, - quantization_params.scale); - } - } else { - quantization_params = array.GetQuantizationParams(); - } - - double quantized_min, quantized_max; - CHECK(GetQuantizedDataTypeNumericalRange(quantized_data_type, &quantized_min, - &quantized_max)) - << "Type is not quantized"; - - bool has_nontrivial_min_bound = false; - bool has_nontrivial_max_bound = false; - - double lowest_representable_output = - (quantized_min - quantization_params.zero_point) * - quantization_params.scale; - if (lowest_representable_output < clamp_min) { - has_nontrivial_min_bound = true; - transformation->AddMessageF( - "Quantized activation function is not trivial: " - "the lowest representable output value %g" - " less than the clamp min bound %g.", - lowest_representable_output, clamp_min); - } - - double highest_representable_output = - (quantized_max - quantization_params.zero_point) * - quantization_params.scale; - if (highest_representable_output > clamp_max) { - has_nontrivial_max_bound = true; - transformation->AddMessageF( - "Quantized activation function is not trivial: " - "the highest representable output value %g" - " is greater than the clamp max bound %g.", - highest_representable_output, clamp_max); - } - - return !has_nontrivial_min_bound && !has_nontrivial_max_bound; -} - -} // namespace toco diff --git a/tensorflow/contrib/lite/toco/graph_transformations/quantization_util.h b/tensorflow/contrib/lite/toco/graph_transformations/quantization_util.h deleted file mode 100644 index cf093c6f17b45839156dae0d06ca2fc7e5e2f3c6..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/toco/graph_transformations/quantization_util.h +++ /dev/null @@ -1,63 +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_LITE_TOCO_GRAPH_TRANSFORMATIONS_QUANTIZATION_UTIL_H_ -#define TENSORFLOW_CONTRIB_LITE_TOCO_GRAPH_TRANSFORMATIONS_QUANTIZATION_UTIL_H_ - -#include "tensorflow/contrib/lite/kernels/internal/quantization_util.h" -#include "tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h" -#include "tensorflow/contrib/lite/toco/model.h" - -namespace toco { - -// Gets the target quantized data type of an array based on the fake quant op. -// For example, if the num_bits is 8 the data type will be kUint8. -bool InferQuantizedDataTypeFromFakeQuant( - const FakeQuantOperator& op, ArrayDataType* out_quantized_data_type); - -// Gets the min/max numerical range for the given quantized data type. -// For example, kUint8 will return [0,255]. -// Returns true if the ranges were set and false if the type is not quantized. -bool GetQuantizedDataTypeNumericalRange(ArrayDataType data_type, - double* out_min_value, - double* out_max_value); - -// Returns the quantized data type of an array, falling back to the provided -// default data type. -ArrayDataType GetQuantizedDataType(const Array& array, - ArrayDataType default_type); - -// Chooses the quantization params for a given array and a given target -// quantized data type (which may not be the array's current data type). -void ChooseQuantizationParamsForArrayAndQuantizedDataType( - const Array& array, ArrayDataType quantized_data_type, - QuantizationParams* quantization_params); - -// Quantizes an array by setting its data type and (if constant) quantizing -// all values in the array. -void QuantizeArray(GraphTransformation* transformation, Model* model, - const string& name, ArrayDataType quantized_data_type, - const QuantizationParams& quantization_params); - -// Returns true if the given array, when quantized, contains only values between -// the provided clamp min/max. -// Either clamp_min or clamp_max may be +/-infinity to indicate that the value -// is unbounded on that side. -bool IsArrayQuantizedRangeSubset(GraphTransformation* transformation, - const Array& array, double clamp_min, - double clamp_max); - -} // namespace toco - -#endif // TENSORFLOW_CONTRIB_LITE_TOCO_GRAPH_TRANSFORMATIONS_QUANTIZATION_UTIL_H_ diff --git a/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_range.cc b/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_range.cc deleted file mode 100644 index 069d4dafaabd41e79adc27083f8c14e164e22ff3..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/toco/graph_transformations/resolve_constant_range.cc +++ /dev/null @@ -1,111 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include "tensorflow/contrib/lite/toco/graph_transformations/graph_transformations.h" -#include "tensorflow/contrib/lite/toco/model.h" -#include "tensorflow/contrib/lite/toco/tooling_util.h" -#include "tensorflow/core/platform/logging.h" - -namespace toco { - -::tensorflow::Status ResolveConstantRange::Run(Model* model, - std::size_t op_index, - bool* modified) { - *modified = false; - const auto it = model->operators.begin() + op_index; - auto* base_op = it->get(); - if (base_op->type != OperatorType::kRange) { - return ::tensorflow::Status::OK(); - } - auto* op = static_cast(base_op); - - CHECK_EQ(op->inputs.size(), 3); - const auto& start_array = model->GetArray(op->inputs[0]); - if (!start_array.has_shape()) { - // Yield until all input dims have been resolved. - return ::tensorflow::Status::OK(); - } - const auto& limit_array = model->GetArray(op->inputs[1]); - if (!limit_array.has_shape()) { - // Yield until all input dims have been resolved. - return ::tensorflow::Status::OK(); - } - const auto& delta_array = model->GetArray(op->inputs[2]); - if (!delta_array.has_shape()) { - // Yield until all input dims have been resolved. - return ::tensorflow::Status::OK(); - } - - for (const auto& input : op->inputs) { - if (!IsConstantParameterArray(*model, input)) { - // yield if any input is mutable - return ::tensorflow::Status::OK(); - } - } - - CHECK_EQ(op->outputs.size(), 1); - auto& output_array = model->GetArray(op->outputs[0]); - if (output_array.data_type == ArrayDataType::kNone) { - // Yield until the output type has been set by PropagateArrayDataTypes - return ::tensorflow::Status::OK(); - } - - CHECK_EQ(RequiredBufferSizeForShape(start_array.shape()), 1) - << "Range op inputs must be scalar."; - CHECK_EQ(RequiredBufferSizeForShape(limit_array.shape()), 1) - << "Range op inputs must be scalar."; - CHECK_EQ(RequiredBufferSizeForShape(delta_array.shape()), 1) - << "Range op inputs must be scalar."; - - CHECK(start_array.data_type == ArrayDataType::kInt32) - << "Range op inputs must be int32."; - CHECK(limit_array.data_type == ArrayDataType::kInt32) - << "Range op inputs must be int32."; - CHECK(delta_array.data_type == ArrayDataType::kInt32) - << "Range op inputs must be int32."; - - // Compute buffer contents - int start = start_array.GetBuffer().data[0]; - int limit = limit_array.GetBuffer().data[0]; - int delta = delta_array.GetBuffer().data[0]; - auto& buffer = output_array.GetMutableBuffer(); - buffer.data.clear(); - for (int32 val = start; val < limit; val += delta) { - buffer.data.push_back(val); - } - CHECK_EQ(floor((limit - start) / delta), buffer.data.size()); - CHECK_EQ(buffer.data.size(), output_array.shape().dims()[0]); - - // Delete the input array if no longer used - if (IsDiscardableArray(*model, op->inputs[0]) && - CountOpsWithInput(*model, op->inputs[0]) == 1) { - model->EraseArray(op->inputs[0]); - } - if (IsDiscardableArray(*model, op->inputs[1]) && - CountOpsWithInput(*model, op->inputs[1]) == 1) { - model->EraseArray(op->inputs[1]); - } - if (IsDiscardableArray(*model, op->inputs[2]) && - CountOpsWithInput(*model, op->inputs[2]) == 1) { - model->EraseArray(op->inputs[2]); - } - - // Delete the operator - model->operators.erase(it); - - *modified = true; - return ::tensorflow::Status::OK(); -} - -} // namespace toco diff --git a/tensorflow/contrib/lite/toco/graph_transformations/tests/BUILD b/tensorflow/contrib/lite/toco/graph_transformations/tests/BUILD deleted file mode 100644 index 6f1be298caaf110f7ef7113bfeb930c96b0ec9de..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/toco/graph_transformations/tests/BUILD +++ /dev/null @@ -1,45 +0,0 @@ -package(default_visibility = ["//visibility:public"]) - -licenses(["notice"]) # Apache 2.0 - -load( - "//tensorflow:tensorflow.bzl", - "tf_cc_test", -) - -tf_cc_test( - name = "lstm_utils_test", - srcs = ["lstm_utils_test.cc"], - tags = ["no_oss"], - deps = [ - "//tensorflow/contrib/lite/toco:graph_transformations", - "//tensorflow/contrib/lite/toco:model", - "//tensorflow/contrib/lite/toco:tooling_util", - "@com_google_googletest//:gtest_main", - ], -) - -tf_cc_test( - name = "resolve_constant_concatenation_test", - srcs = ["resolve_constant_concatenation_test.cc"], - tags = ["no_oss"], - deps = [ - "//tensorflow/contrib/lite/toco:graph_transformations", - "//tensorflow/contrib/lite/toco:model", - "//tensorflow/contrib/lite/toco:tooling_util", - "@com_google_googletest//:gtest_main", - ], -) - -tf_cc_test( - name = "resolve_constant_unary_test", - srcs = ["resolve_constant_unary_test.cc"], - tags = ["no_oss"], - deps = [ - "//tensorflow/contrib/lite/toco:graph_transformations", - "//tensorflow/contrib/lite/toco:model", - "//tensorflow/contrib/lite/toco:tooling_util", - "@com_google_absl//absl/memory", - "@com_google_googletest//:gtest_main", - ], -) diff --git a/tensorflow/contrib/lite/toco/import_tensorflow_test.cc b/tensorflow/contrib/lite/toco/import_tensorflow_test.cc deleted file mode 100644 index 0767221b83cb066583dcd63a118015649d25d248..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/toco/import_tensorflow_test.cc +++ /dev/null @@ -1,334 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include "tensorflow/contrib/lite/toco/import_tensorflow.h" - -#include -#include -#include "tensorflow/core/framework/attr_value.pb.h" -#include "tensorflow/core/framework/attr_value_util.h" -#include "tensorflow/core/framework/node_def.pb.h" -#include "tensorflow/core/framework/node_def_builder.h" -#include "tensorflow/core/framework/tensor.pb.h" -#include "tensorflow/core/framework/tensor_shape.pb.h" -#include "tensorflow/core/lib/core/status.h" - -namespace toco { - -using tensorflow::AttrValue; -using tensorflow::DT_BOOL; -using tensorflow::DT_FLOAT; -using tensorflow::DT_INT32; -using tensorflow::DT_INT64; -using tensorflow::DT_QUINT8; -using tensorflow::DT_STRING; -using tensorflow::NodeDef; -using tensorflow::Status; - -namespace internal { -using ConverterType = tensorflow::Status (*)( - const NodeDef& node, const TensorFlowImportFlags& tf_import_flags, - Model* model); -using ConverterMapType = std::unordered_map; - -ConverterMapType GetTensorFlowNodeConverterMap(); -Status ImportTensorFlowNode(const NodeDef&, const TensorFlowImportFlags&, - Model*, const ConverterMapType&); -} // namespace internal - -namespace { - -Status ImportNode(const NodeDef& node, Model* model) { - const auto converter = internal::GetTensorFlowNodeConverterMap(); - return internal::ImportTensorFlowNode(node, TensorFlowImportFlags(), model, - converter); -} - -Status ImportFlexNode(const NodeDef& node, Model* model) { - // Empty converter => all nodes are flex nodes. - const auto converter = internal::ConverterMapType(); - return internal::ImportTensorFlowNode(node, TensorFlowImportFlags(), model, - converter); -} - -Status ImportNode(const NodeDef& node) { - Model model; - return ImportNode(node, &model); -} - -NodeDef BuildNode( - const std::string& op, - const std::vector>& output_shapes) { - NodeDef node; - node.set_op(op); - node.set_name("Node1"); - node.add_input(); - node.set_input(0, "Node0"); - - AttrValue::ListValue* shapes = - (*node.mutable_attr())["_output_shapes"].mutable_list(); - for (const auto& output_shape : output_shapes) { - tensorflow::TensorShapeProto* shape = shapes->add_shape(); - for (int64_t output_shape_dim : output_shape) { - auto shape_dim = shape->add_dim(); - shape_dim->set_size(output_shape_dim); - } - } - - return node; -} - -class ShapeImportTest : public ::testing::TestWithParam { - protected: - ShapeImportTest() {} - - void BuildConstNode(std::initializer_list shape, - tensorflow::DataType dtype, int64_t num_elements, - NodeDef* node) { - node->set_op("Const"); - node->set_name("Node1"); - - // An attribute describing the type of this const node. - AttrValue dtype_attr; - SetAttrValue(dtype, &dtype_attr); - (*node->mutable_attr())["dtype"] = dtype_attr; - - // An attribute describing the content of this const node. - tensorflow::TensorProto t; - t.set_dtype(dtype); - auto* s = t.mutable_tensor_shape(); - for (auto d : shape) { - s->add_dim()->set_size(d); - } - - // TODO(ahentz): also need to test via tensor_content() - switch (dtype) { - case DT_FLOAT: - for (int64_t i = 0; i < num_elements; ++i) { - t.add_float_val(i / 10000.0); - } - break; - case DT_INT32: - for (int64_t i = 0; i < num_elements; ++i) { - t.add_int_val(i % std::numeric_limits::max()); - } - break; - case DT_QUINT8: - for (int64_t i = 0; i < num_elements; ++i) { - t.add_int_val(i % std::numeric_limits::max()); - } - break; - case DT_INT64: - for (int64_t i = 0; i < num_elements; ++i) { - t.add_int64_val(i); - } - break; - case DT_STRING: - break; - case DT_BOOL: - for (int64_t i = 0; i < num_elements; ++i) { - t.add_bool_val(i % 2); - } - break; - default: - break; - } - - AttrValue value_attr; - SetAttrValue(t, &value_attr); - (*node->mutable_attr())["value"] = value_attr; - } -}; - -class TypeImportTest : public ::testing::TestWithParam< - std::pair> { - protected: - TypeImportTest() {} - - void BuildUnaryNode(const std::string& op_name, tensorflow::DataType dtype, - NodeDef* node) { - node->set_op(op_name); - node->set_name("Node1"); - - node->add_input(); - node->set_input(0, "Node0"); - - AttrValue dtype_attr; - SetAttrValue(dtype, &dtype_attr); - (*node->mutable_attr())["T"] = dtype_attr; - } -}; - -std::vector TestTypes() { - return {DT_FLOAT, DT_INT32, DT_INT64, DT_BOOL, DT_QUINT8}; -} - -TEST_P(ShapeImportTest, ShapeElementIsNegative) { - NodeDef node; - BuildConstNode({1, -2, 10}, GetParam(), 0, &node); - auto status = ImportNode(node); - EXPECT_EQ( - status.error_message(), - "Tensor shape should not include negative values\n\t (while processing " - "node 'Node1')"); -} -INSTANTIATE_TEST_CASE_P(ShapeElementIsNegative, ShapeImportTest, - ::testing::ValuesIn(TestTypes())); - -TEST_P(ShapeImportTest, ShapeElementTooLarge) { - NodeDef node; - BuildConstNode({3000000000}, GetParam(), 0, &node); - auto status = ImportNode(node); - EXPECT_EQ(status.error_message(), - "Shape element overflows\n\t (while processing node 'Node1')"); -} -INSTANTIATE_TEST_CASE_P(ShapeElementTooLarge, ShapeImportTest, - ::testing::ValuesIn(TestTypes())); - -TEST_P(ShapeImportTest, ShapeTooLarge) { - NodeDef node; - BuildConstNode({1000000, 2000000, 2000000, 2000000}, GetParam(), 0, &node); - auto status = ImportNode(node); - EXPECT_EQ(status.error_message(), - "Tensor shape is too large\n\t (while processing node 'Node1')"); -} -INSTANTIATE_TEST_CASE_P(ShapeTooLarge, ShapeImportTest, - ::testing::ValuesIn(TestTypes())); - -TEST_P(ShapeImportTest, ValidShapeButZeroElements) { - NodeDef node; - BuildConstNode({1, 2, 2, 2}, GetParam(), 0, &node); - auto status = ImportNode(node); - EXPECT_THAT(status.error_message(), - ::testing::MatchesRegex( - "Neither input_content .0. nor .*_val .0. have the right " - "dimensions .8. for this .* tensor\n\t .while processing " - "node 'Node1'.")); -} -INSTANTIATE_TEST_CASE_P(ValidShapeButZeroElements, ShapeImportTest, - ::testing::ValuesIn(TestTypes())); - -std::vector> UnaryTestTypes() { - return {{DT_FLOAT, ArrayDataType::kFloat}, - {DT_INT32, ArrayDataType::kInt32}, - {DT_INT64, ArrayDataType::kInt64}}; -} - -TEST_P(TypeImportTest, BasicTypeInference) { - NodeDef node; - BuildUnaryNode("Atan", GetParam().first, &node); - - Model model; - EXPECT_TRUE(ImportNode(node, &model).ok()); - - ASSERT_THAT(model.operators.size(), ::testing::Ge(1)); - ASSERT_EQ(model.operators[0]->type, OperatorType::kUnsupported); - const TensorFlowUnsupportedOperator* op = - static_cast( - model.operators[0].get()); - ASSERT_THAT(op->output_data_types, ::testing::ElementsAre(GetParam().second)); -} -INSTANTIATE_TEST_CASE_P(BasicTypeInference, TypeImportTest, - ::testing::ValuesIn(UnaryTestTypes())); - -TEST(ImportTest, TypeInferenceWithFixedOutputType) { - // Create an op that has a fixed output type (bool). - Model model; - EXPECT_TRUE(ImportNode(BuildNode("IsFinite", {{1, 2}, {2, 3}}), &model).ok()); - ASSERT_THAT(model.operators.size(), ::testing::Ge(1)); - ASSERT_EQ(model.operators[0]->type, OperatorType::kUnsupported); - const TensorFlowUnsupportedOperator* op = - static_cast( - model.operators[0].get()); - - // The static output type should be indicated in the imported op. - ASSERT_THAT(op->output_data_types, - ::testing::ElementsAre(ArrayDataType::kBool)); -} - -TEST(ImportTest, FailedTypeInference) { - // Create a unary op with no Type ("T") annotation. - NodeDef node; - node.set_op("Atan"); - node.set_name("Node1"); - node.add_input(); - node.set_input(0, "Node0"); - - Model model; - EXPECT_TRUE(ImportNode(node, &model).ok()); - - ASSERT_THAT(model.operators.size(), ::testing::Ge(1)); - ASSERT_EQ(model.operators[0]->type, OperatorType::kUnsupported); - const TensorFlowUnsupportedOperator* op = - static_cast( - model.operators[0].get()); - ASSERT_TRUE(op->output_data_types.empty()); -} - -TEST(ImportTest, UnsupportedOpWithOutputShapes) { - // Create an unsupported op with output shapes. - Model model; - EXPECT_TRUE(ImportNode(BuildNode("Atan", {{1, 2}, {2, 3}}), &model).ok()); - ASSERT_THAT(model.operators.size(), ::testing::Ge(1)); - ASSERT_EQ(model.operators[0]->type, OperatorType::kUnsupported); - const TensorFlowUnsupportedOperator* op = - static_cast( - model.operators[0].get()); - - // The output shapes should be imported. - ASSERT_EQ(op->output_shapes.size(), 2); - ASSERT_THAT(op->output_shapes[0].dims(), ::testing::ElementsAre(1, 2)); - ASSERT_THAT(op->output_shapes[1].dims(), ::testing::ElementsAre(2, 3)); -} - -TEST(ImportTest, UnsupportedOpWithWildcardOutputShapes) { - // Create an unsupported op with wildcard output shapes. - Model model; - EXPECT_TRUE(ImportNode(BuildNode("Atan", {{-1, 2}}), &model).ok()); - ASSERT_THAT(model.operators.size(), ::testing::Ge(1)); - ASSERT_EQ(model.operators[0]->type, OperatorType::kUnsupported); - const TensorFlowUnsupportedOperator* op = - static_cast( - model.operators[0].get()); - - // Wildcard shapes aren't yet supported. - ASSERT_TRUE(op->output_shapes.empty()); -} - -TEST(ImportTest, UnsupportedOpWithMultipleOutputs) { - NodeDef node = BuildNode("Unpack", {}); - - // Unpack's OpDef has a single output which gets multiplied based on the - // "num" attribute of the NodeDef. - AttrValue value_attr; - SetAttrValue(3, &value_attr); // 3 outputs. - (*node.mutable_attr())["num"] = value_attr; - - Model model; - EXPECT_TRUE(ImportFlexNode(node, &model).ok()); - - ASSERT_THAT(model.operators.size(), ::testing::Ge(1)); - ASSERT_EQ(model.operators[0]->type, OperatorType::kUnsupported); - const TensorFlowUnsupportedOperator* op = - static_cast( - model.operators[0].get()); - - ASSERT_EQ(op->outputs.size(), 3); - ASSERT_EQ(op->outputs[0], "Node1"); - ASSERT_EQ(op->outputs[1], "Node1:1"); - ASSERT_EQ(op->outputs[2], "Node1:2"); -} - -} // namespace -} // namespace toco diff --git a/tensorflow/contrib/lite/toco/model.h b/tensorflow/contrib/lite/toco/model.h deleted file mode 100644 index f3b84430dbdceba65711c4d04c24829ea36f250e..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/toco/model.h +++ /dev/null @@ -1,2131 +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_CONTRIB_LITE_TOCO_MODEL_H_ -#define TENSORFLOW_CONTRIB_LITE_TOCO_MODEL_H_ - -#include -#include -#include -#include -#include -#include -#include - -#include "absl/types/optional.h" -#include "tensorflow/contrib/lite/toco/model_flags.pb.h" -#include "tensorflow/contrib/lite/toco/runtime/types.h" -#include "tensorflow/contrib/lite/toco/toco_port.h" -#include "tensorflow/contrib/lite/toco/toco_types.h" -#include "tensorflow/core/platform/logging.h" - -namespace toco { - -using tflite::QuantizationParams; - -enum class OperatorType : uint8 { - kNone, - // General-purpose neural network operators. - kAdd, - kAddN, - kAveragePool, - kBatchMatMul, - kBatchNormalization, - kConv, - kConcatenation, - kDepthwiseConv, - kDepthToSpace, - kSpaceToDepth, - kDequantize, - kDiv, - kExp, - kExpandDims, - kFill, - kFloorDiv, - kFloorMod, - kFullyConnected, - kL2Normalization, - kL2Pool, - kLstmCell, - kUnidirectionalSequenceLstm, - kLocalResponseNormalization, - kLog, - kLogistic, - kMaxPool, - kFakeQuant, - kMul, - kOneHot, - kRandomUniform, - kRange, - kRank, - kRelu, - kRelu1, - kRelu6, - kPRelu, - kSoftmax, - kLogSoftmax, - kSub, - kTanh, - kTransposeConv, - kCast, - kFloor, - kGather, - kResizeBilinear, - kSin, - kSpaceToBatchND, - kPack, - kBatchToSpaceND, - kPad, - kPadV2, - kReduceProd, // Reduction product - kStridedSlice, - kSlice, - kSqueeze, - kMean, - kArgMax, - // The SVDF Op is a decomposition of a densely connected Op into - // low rank filters. For details: - // https://research.google.com/pubs/pub43813.html - kSvdf, - // Special operators used for importing TensorFlow nodes. - // The general intent is to have some graph transformation either - // drop them or rewrite them as general-purpose operators. - kAll, - kAssert, - kConcat, - kConcatV2, - kGreater, - kGreaterEqual, - kIdentity, - kLess, - kLessEqual, - kReduceMax, // Reduction Max - kMaximum, // Element-wise Maximum - kReduceMin, // Reduction Min - kMinimum, // Element-wise Minimum - kMatMul, - kMerge, - kNeg, - kReshape, - kRsqrt, - kShape, - kSplit, - kSqrt, - kSquare, - kSum, - kSwitch, - kTile, - kTranspose, - kTopK_V2, - kDynamicPartition, - kDynamicStitch, - // An unsupported TF operation. It's only needed to be able to represent TF - // graph internally and is expected to be dropped by graph transformations. - kUnsupported, - // Finally, TensorFlow uses different conventions for axes ordering, - // see AxesOrder, and this cannot always be resolved at the time of importing - // nodes, as TensorFlow parameters may be constant-expression subgraphs - // instead of being given as plain constant arrays. So we need to insert - // special nodes in the graph to shuffle axes. - kReorderAxes, - kSelect, - kSparseToDense, - kEqual, - kNotEqual, - kPow, - kArgMin, - kAny, - kLogicalAnd, - kLogicalNot, - kLogicalOr, - kCTCBeamSearchDecoder, - kUnpack, - kZerosLike, -}; - -// Helper to deal with TensorFlow arrays using a different ordering of -// dimensions -// ("axes") than our own. -// TODO(benoitjacob): Ultimately, we shouldn't have any "ordering" of axes, -// we should have associative arrays mapping symbolic axes identifiers (like -// "output_depth") to dimensions. We would then not need this anymore. -enum class AxesOrder { - kOneAxis, // one-dimensional array, one unique axis. - kCR, // column-major matrix storage order. Our standard. - kRC, // row-major matrix storage order. TensorFlow default. - kOHWI, // Our standard for conv weights - kHWIO, // TensorFlow conv weights - k1HWO, // Our standard for DepthwiseConv weights - kHWIM, // TensorFlow DepthwiseConv weights - kNHWC, // TensorFlow activations - kHWOI, // TensorFlow back-prop conv weights -}; - -// The type of the scalars in an array. -// Note that the type does not by itself tell whether the values in the array -// are non-quantized (can be accessed directly) or quantized (must be -// interpreted in conjunction with QuantizationParams). -// -// In practice though: -// float values are never quantized -// uint8 values are always quantized -// int32 values are sometimes quantized (depending on whether -// QuantizationParams are present). -// complex values are never quantized -// other types are never quantized at the moment. -// -// kNone means that we don't know the data type yet, or that we don't care -// because we'll be dropping the array anyway (e.g. some exotic array types -// may be involved only in debug-only subgraphs that we may not be interested -// in actually supporting). -enum class ArrayDataType : uint8 { - kNone, // 0 - kBool, - kFloat, - kInt8, - kUint8, - kInt16, // 5 - kUint16, - kInt32, - kUint32, - kInt64, - kUint64, // 10 - kString, - kComplex64, -}; - -// Compile-time logic to map ArrayDataType to the corresponding C++ scalar type -template -struct DataTypeImpl {}; -template <> -struct DataTypeImpl { - typedef int Type; -}; -template <> -struct DataTypeImpl { - typedef bool Type; -}; -template <> -struct DataTypeImpl { - typedef float Type; -}; -template <> -struct DataTypeImpl { - typedef int8 Type; -}; -template <> -struct DataTypeImpl { - typedef uint8 Type; -}; -template <> -struct DataTypeImpl { - typedef int16 Type; -}; -template <> -struct DataTypeImpl { - typedef uint16 Type; -}; -template <> -struct DataTypeImpl { - typedef int32 Type; -}; -template <> -struct DataTypeImpl { - typedef uint32 Type; -}; -template <> -struct DataTypeImpl { - typedef int64 Type; -}; -template <> -struct DataTypeImpl { - typedef uint64 Type; -}; -template <> -struct DataTypeImpl { - typedef string Type; -}; -template <> -struct DataTypeImpl { - typedef std::complex Type; -}; - -template -using DataType = typename DataTypeImpl::Type; - -// Base class for type-specific buffer types. -struct GenericBuffer { - // Non-default-constructible: only ArrayDataType-specific subclass - // objects may be constructed. - GenericBuffer() = delete; - // Non-copyable-or-movable: we should only store pointers-to-Buffer - // in containers, not Operators themselves, so there should be no - // copy or move. - GenericBuffer(const GenericBuffer&) = delete; - GenericBuffer(const GenericBuffer&&) = delete; - - // We need a virtual destructor so we can store pointers-to-Buffer - // in containers and have the containers call the right subclass destructor. - virtual ~GenericBuffer() {} - - virtual int Length() const = 0; - - const ArrayDataType type; - - protected: - // Constructor used by subclasses for specific ArrayDataType's. - explicit GenericBuffer(ArrayDataType t) : type(t) {} -}; - -// Type-specific buffer, containing type-specific storage. -template -struct Buffer : GenericBuffer { - Buffer() : GenericBuffer(A) {} - - int Length() const override { return data.size(); } - - std::vector> data; -}; - -class Shape { - public: - // For Shape, we stick to half-way encapsulation for now: - // we hide the raw dims_ member, but expose it raw by accessors - // because from some brainstorming, it's not at all easy to - // anticipate which flavor of more hermetic encapsulation would - // actually buy us future-proof-ness without being needlessly - // cumbersome. - Shape() {} - Shape(std::initializer_list dim_list) : dims_(dim_list) {} - - void ReplaceDims(std::initializer_list dim_list) { - dims_ = std::vector(dim_list); - } - - const std::vector& dims() const { return dims_; } - std::vector* mutable_dims() { return &dims_; } - const int dimensions_count() const { return dims_.size(); } - - // We still have that one convenience accessor to avoid - // the awkward double bracket issue: shape.dims()[i]. - int dims(int i) const { - // Always check for out-of-bounds accesses, even in optimized builds where - // standard assertions are disabled. Out-of-bounds access here is a common - // occurrence. - CHECK_GE(i, 0); - CHECK_GT(dims_.size(), i); - return dims_[i]; - } - - bool operator==(const Shape& comp) const { - return (this->dims_ == comp.dims()); - } - - bool operator!=(const Shape& comp) const { return !((*this) == comp); } - - private: - std::vector dims_; -}; - -// Base class for all operator classes. -struct Operator { - // Non-default-constructible: only OperatorType-specific subclass - // objects may be constructed. - Operator() = delete; - // Non-copyable-or-movable: we should only store pointers-to-Operator - // in containers, not Operators themselves, so there should be no - // copy or move. - Operator(const Operator&) = delete; - Operator(const Operator&&) = delete; - - // We need a virtual destructor so we can store pointers-to-Operator - // in containers and have the containers call the right subclass destructor. - virtual ~Operator() {} - - // The specific type of operator. Corresponds 1:1 to subclasses. - const OperatorType type; - - // The activation function that may be fused into this operator, - // or None if no activation function is fused. - FusedActivationFunctionType fused_activation_function; - - // Input arrays: either activation arrays or constant array parameters. - // We refer to them by their name, not by their address; the mapping of - // names to addresses is given by the Model, which owns both Operator's and - // Array's. Thus, an Operator on its own doesn't contain much information, - // it is meant to be used in conjunction with the Model that owns it. - std::vector inputs; - - // Output activation arrays. Same comments as for inputs apply here too. - std::vector outputs; - - // If true, the array has more outputs than are listed in the 'outputs' - // member. These need to be resolved by some graph transformation. - // This flag is only here to indicate that an operator should not be - // discarded as unused, even if from its 'outputs' member alone it - // looks unused. - bool unresolved_outputs = false; - - // A serialized tensorflow::NodeDef string. - // The field is filled only when importing from TensorFlow. - // It's guaranteed to be filled for `TensorFlowUnsupportedOperator`. - // It's not guaranteed to be filled for other ops. Ops created by graph - // transformations won't have TensorFlow NodeDef. - string tensorflow_node_def; - - protected: - // Constructor used by subclasses for specific OperatorType's. - explicit Operator(OperatorType t) - : type(t), - fused_activation_function(FusedActivationFunctionType::kNone) {} -}; - -// Padding types for Conv-like operators. This is how padding is typically -// specified in model files. But for inference, we will need to resolve this -// to a FixedPadding, see below. -enum class PaddingType { kNone, kSame, kValid }; - -// Padding as resolved for a specific layer shape, as needed for inference. -// For a given layer shape, a given padding type will resolve to a choice of -// a number of padding rows and columns, which we call the padding height and -// width respectively. -struct FixedPadding { - int width = 0; - int height = 0; -}; - -// "Universal" padding struct containing both a generic PaddingType (as -// represented in a model file), and a FixedPadding (as needed for inference). -// The latter is resolved during the PropagateFixedSizes pass. -struct Padding { - FixedPadding& GetOrCreateFixedPadding() { - if (!fixed) { - FixedPadding* ptr = new FixedPadding; - fixed = std::unique_ptr(ptr); - } - return *fixed; - } - - Padding() : type(PaddingType::kNone) {} - PaddingType type; - std::unique_ptr fixed; -}; - -// "Convolutional" layer, as represented in model files. -// -// Inputs: -// inputs[0]: required: the input activations array -// inputs[1]: required: the Conv weights -// inputs[2]: optional: the bias vector, specifying the biases for each output -// channel. -// -// Outputs: -// outputs[0]: required: the output activations array -// outputs[1]: optional: the intermediate array of im2col-replicated input -// activations. Present when targeting implementations -// of Conv layers as Im2col+GEMM. -// -// TensorFlow equivalent: Conv2D -struct ConvOperator : Operator { - ConvOperator() : Operator(OperatorType::kConv) {} - Padding padding; - int stride_width = 0; - int stride_height = 0; - // A dilation_rate of 0 is invalid and this field is an optional attribute. - // Thus initializing it to 1 to allow default conv behavior when the - // attribute is not present. - int dilation_width_factor = 1; - int dilation_height_factor = 1; -}; - -// CTCBeamSearchDecoder operator: -// -// Inputs: -// inputs[0]: required: the logits. -// inputs[1]: required: sequence length. -// inputs[2]: optional: beam width. -// inputs[3]: optional: top paths. -// inputs[4]: optional: merge repeated. -// -// Outputs: -// outputs[0]: deocoded. -// outputs[1]: log probability. -// -// TensorFlow equivalent: CTCBeamSearchDecoder -struct CTCBeamSearchDecoderOperator : Operator { - CTCBeamSearchDecoderOperator() - : Operator(OperatorType::kCTCBeamSearchDecoder) {} - int beam_width; - int top_paths; - bool merge_repeated = true; -}; - -// Depthwise-separable convolution operator. -// -// Inputs: -// inputs[0]: required: the input activations array -// inputs[1]: required: the DepthwiseConv weights -// inputs[2]: optional: the bias vector, specifying the biases for each output -// channel. -// -// TensorFlow equivalent: DepthwiseConv2dNative -struct DepthwiseConvOperator : Operator { - DepthwiseConvOperator() : Operator(OperatorType::kDepthwiseConv) {} - Padding padding; - int stride_height = 0; - int stride_width = 0; - int depth_multiplier = 0; - // A dilation_rate of 0 is invalid and this field is an optional attribute. - // Thus initializing it to 1 to allow default conv behavior when the - // attribute is not present. - int dilation_width_factor = 1; - int dilation_height_factor = 1; -}; - -// Depth-to-space transform operator. -// -// Inputs: -// inputs[0]: required: the input activations array -// -// TensorFlow equivalent: DepthToSpace -struct DepthToSpaceOperator : Operator { - DepthToSpaceOperator() : Operator(OperatorType::kDepthToSpace) {} - int block_size = 0; -}; - -// Space-to-depth transform operator. -// -// Inputs: -// inputs[0]: required: the input activations array -// -// TensorFlow equivalent: SpaceToDepth -struct SpaceToDepthOperator : Operator { - SpaceToDepthOperator() : Operator(OperatorType::kSpaceToDepth) {} - int block_size = 0; -}; - -// Fully-connected operator. -// -// Inputs: -// inputs[0]: required: the input activations array -// inputs[1]: required: the FullyConnected weights -// inputs[2]: optional: the bias vector, specifying the biases for each output -// channel. -// -// TensorFlow equivalent: a pair consisting of a Reshape node reshaping the -// input activations as a matrix, followed by a MatMul node. -struct FullyConnectedOperator : Operator { - FullyConnectedOperator() : Operator(OperatorType::kFullyConnected) {} - FullyConnectedWeightsFormat weights_format = - FullyConnectedWeightsFormat::kDefault; -}; - -// Dequantization operator, converting a quantized array of integers with -// quantization parameters specifying how these integers correspond to real -// numbers -// (see QuantizationParams) to an output activations array of floating-point -// values. -// -// In floating-point image models, there is typically a Dequantization operator -// at the very beginning, converting the input image RGB data, consisting of -// uint8 integer values, to floating-point input activations. That is where -// image model parameters such as "mean_value" and "std_value" are typically -// handled. -// -// This is the only operator type that converts from quantized to -// floating-point, -// and there is at the moment no operator type at all to convert from -// floating-point -// to quantized. Every other operator does either float->float or -// quantized->quantized. -// -// Inputs: -// inputs[0]: required: the input quantized activations array -// -// TensorFlow equivalent: Dequantize -struct DequantizeOperator : Operator { - DequantizeOperator() : Operator(OperatorType::kDequantize) {} -}; - -// Batch-normalization operator. -// -// We only support batch-normalization using pre-learned moments, so this is -// just -// computing (input - mean) * multiplier + offset. As such, this can be -// expressed as a combination of Add and Mul nodes, and indeed this is how -// we break it down during tooling for the purpose of fusing it into -// other operators. -// -// Inputs: -// inputs[0]: required: the input activations array -// inputs[1]: required: the learned mean array -// inputs[2]: required: the learned multiplier array -// inputs[3]: required: the learned offset array -// -// TensorFlow equivalent: a combination of Add and Mul nodes -struct BatchNormalizationOperator : Operator { - BatchNormalizationOperator() - : Operator(OperatorType::kBatchNormalization), - global_normalization(false) {} - bool global_normalization; -}; - -// L2-normalization operator. -// -// Inputs: -// inputs[0]: required: the input activations array -// -// TensorFlow equivalent: none. In TensorFlow, L2 normalization is implemented -// by a sub-graph of operators implementing L2-normalization -// from lower-level arithmetic nodes; during tooling, we identify such -// sub-graphs -// and replace them by L2NormalizationOperator's. See IdentifyL2Normalization. -struct L2NormalizationOperator : Operator { - L2NormalizationOperator() : Operator(OperatorType::kL2Normalization) {} -}; - -// LSTM Cell operator. -// -// Inputs: -// inputs[0]: required: the input data array -// inputs[1]: required: the previous output activations array -// inputs[2]: required: the learned weights array -// inputs[3]: required: the learned biases array -// inputs[4]: required: the previous output state -// outputs[0]: required: the output activations array -// outputs[1]: required: the new state array -// -// TensorFlow equivalent: none. In TensorFlow, an LSTM is implemented -// with a sub-graph of lower-level arithmetic nodes; during tooling, we identify -// such sub-graphs and replace them with LstmCells. See IdentifyLstmCell(). -struct LstmCellOperator : Operator { - enum Inputs { - DATA_INPUT = 0, - PREV_ACTIV_INPUT = 1, - WEIGHTS_INPUT = 2, - BIASES_INPUT = 3, - PREV_STATE_INPUT = 4, - NUM_INPUTS = 5 - }; - enum Outputs { - ACTIV_OUTPUT = 0, - STATE_OUTPUT = 1, - CONCAT_TEMP = 2, - ACTIV_TEMP = 3, - NUM_OUTPUTS = 4 - }; - enum KernelType { - KERNEL_BASIC = 0, - KERNEL_FULL = 1, - }; - - LstmCellOperator() - : Operator(OperatorType::kLstmCell), kernel_type(KERNEL_BASIC) {} - - KernelType kernel_type; -}; - -struct UnidirectionalSequenceLstmOperator : Operator { - UnidirectionalSequenceLstmOperator() - : Operator(OperatorType::kUnidirectionalSequenceLstm) {} -}; - -// Element-wise multiplication operator. -// -// Inputs: -// inputs[0]: required: the left-hand side array -// inputs[1]: required: the right-hand side array -// -// TensorFlow equivalent: Mul -struct MulOperator : Operator { - MulOperator() : Operator(OperatorType::kMul) {} -}; - -// Element-wise Relu operator: -// x -> max(0, x) -// -// Inputs: -// inputs[0]: required: the input array -// -// TensorFlow equivalent: Relu -struct ReluOperator : Operator { - ReluOperator() : Operator(OperatorType::kRelu) {} -}; - -// Element-wise Relu1 operator: -// x -> min(max(x, -1), 1) -// -// Inputs: -// inputs[0]: required: the input array -// -// TensorFlow equivalent: none. We can construct the operator with Minimum -// and Maximum operations -struct Relu1Operator : Operator { - Relu1Operator() : Operator(OperatorType::kRelu1) {} -}; - -// Element-wise Relu6 operator: -// x -> max(0, min(6, x)) -// -// Inputs: -// inputs[0]: required: the input array -// -// TensorFlow equivalent: Relu6 -struct Relu6Operator : Operator { - Relu6Operator() : Operator(OperatorType::kRelu6) {} -}; - -// PRelu -// f(x) = alpha * x for x < 0, f(x) = x for x >= 0. -// -// Inputs: -// inputs[0]: required: the input array -// inputs[1]: required: the alpha array -// -// Equivalent to keras.layers.PReLU. -struct PReluOperator : Operator { - PReluOperator() : Operator(OperatorType::kPRelu) {} -}; - -// Element-wise Logistic operator: -// x -> Logistic(x) = 1 / (1 + exp(-x)) -// -// Inputs: -// inputs[0]: required: the input array -// -// TensorFlow equivalent: Sigmoid -struct LogisticOperator : Operator { - LogisticOperator() : Operator(OperatorType::kLogistic) {} -}; - -// Element-wise natural log operator: -// x -> ln(x) -// -// Inputs: -// inputs[0]: required: the input array -// -// TensorFlow equivalent: Log -struct LogOperator : Operator { - LogOperator() : Operator(OperatorType::kLog) {} -}; - -// Element-wise Tanh operator: -// x -> Tanh(x) = (exp(x) - exp(-x)) / (exp(x) + exp(-x)) -// -// Inputs: -// inputs[0]: required: the input array -// -// TensorFlow equivalent: Tanh -struct TanhOperator : Operator { - TanhOperator() : Operator(OperatorType::kTanh) {} -}; - -// Element-wise Sin operator: -// x -> Sin(x) = sin(x) -// -// Inputs: -// inputs[0]: required: the input array -// -// TensorFlow equivalent: Sin -struct SinOperator : Operator { - SinOperator() : Operator(OperatorType::kSin) {} -}; - -// Element-wise addition operator. -// -// Inputs: -// inputs[0]: required: the left-hand side array -// inputs[1]: required: the right-hand side array -// -// TensorFlow equivalent: Add -struct AddOperator : Operator { - AddOperator() : Operator(OperatorType::kAdd) {} -}; - -// Element-wise addition operator for N inputs. -// -// Inputs: -// inputs[i]: The i-th array to add together to form the output. -// -// TensorFlow equivalent: AddN -struct AddNOperator : Operator { - AddNOperator() : Operator(OperatorType::kAddN) {} -}; - -// Concatenation operator: concatenates its inputs -// along the axis. -// -// Inputs: this operator accepts any number >= 1 of inputs. -// inputs[i]: the i-th array to concatenate. -// -// TensorFlow equivalent: Concat. -struct ConcatenationOperator : Operator { - ConcatenationOperator() : Operator(OperatorType::kConcatenation) {} - int axis = 0; -}; - -// Reordering dimensions. Used only during tooling to transform graphs from -// the TensorFlow format. -// -// Inputs: -// inputs[0]: required: the input array -// -// TensorFlow equivalent: none. This is only useful to convert between formats. -struct ReorderAxesOperator : Operator { - ReorderAxesOperator() : Operator(OperatorType::kReorderAxes) {} - AxesOrder input_axes_order; - AxesOrder output_axes_order; -}; - -// Average-pooling operator. -// -// Inputs: -// inputs[0]: required: the input array -// -// TensorFlow equivalent: AveragePool -struct AveragePoolOperator : Operator { - AveragePoolOperator() : Operator(OperatorType::kAveragePool) {} - Padding padding; - int stride_height = 0; - int stride_width = 0; - int kheight = 0; - int kwidth = 0; -}; - -// Local response normalization operator. -// -// Inputs: -// inputs[0]: required: the input array -// -// TensorFlow equivalent: LRN -struct LocalResponseNormalizationOperator : Operator { - LocalResponseNormalizationOperator() - : Operator(OperatorType::kLocalResponseNormalization) {} - - int range = 0; - float bias = 0.f; - float alpha = 0.f; - float beta = 0.f; -}; - -// Max-pooling operator. -// -// Inputs: -// inputs[0]: required: the input array -// -// TensorFlow equivalent: MaxPool -struct MaxPoolOperator : Operator { - MaxPoolOperator() : Operator(OperatorType::kMaxPool) {} - Padding padding; - int stride_height = 0; - int stride_width = 0; - int kheight = 0; - int kwidth = 0; -}; - -// L2-pooling operator. -// -// Inputs: -// inputs[0]: required: the input array -// -// TensorFlow equivalent: none. Can be shimmed by squaring+avgpool+sqrt. -struct L2PoolOperator : Operator { - L2PoolOperator() : Operator(OperatorType::kL2Pool) {} - Padding padding; - int stride_height = 0; - int stride_width = 0; - int kheight = 0; - int kwidth = 0; -}; - -// The expected [min, max] range of values in a given array. -// Used for quantization only. -// This information typically comes from special nodes found in quantized -// models, see FakeQuantOperator, and is used during quantization to resolve -// actual quantization parameters (see QuantizationParams). -struct MinMax { - double min = 0.; - double max = 0.; -}; - -inline bool operator==(const MinMax& m1, const MinMax& m2) { - return m1.min == m2.min && m1.max == m2.max; -} - -// Fake-quantization operator. This does two things: -// - Annotate its input and output arrays with MinMax information, -// - Arithmetic-wise, this operator rounds incoming activation values -// to the nearest representable value on the scale of 256 -// values from the min to the max value dictated by its MinMax info. -// -// Inputs: -// inputs[0]: required: the input array -// inputs[1]: optional: the 'min' value, if it has not yet been resolved -// to a constant. -// inputs[2]: optional: the 'max' value, if it has not yet been resolved -// to a constant. -// -// TensorFlow equivalent: FakeQuantWithMinMaxVars, FakeQuantWithMinMaxArgs. -struct FakeQuantOperator : Operator { - FakeQuantOperator() : Operator(OperatorType::kFakeQuant) {} - std::unique_ptr minmax; - int num_bits = 8; - bool narrow_range = false; -}; - -// Element-wise division operator. -// -// Inputs: -// inputs[0]: required: the left-hand side array -// inputs[1]: required: the right-hand side array -// -// TensorFlow equivalent: Div -struct DivOperator : Operator { - DivOperator() : Operator(OperatorType::kDiv) {} -}; - -// Element-wise identity (x->x) operator. -// -// Inputs: -// inputs[0]: required: the input array -// -// TensorFlow equivalent: Identity -struct TensorFlowIdentityOperator : Operator { - TensorFlowIdentityOperator() : Operator(OperatorType::kIdentity) {} -}; - -// Batch matrix multiplication operator. This comes from the (deprecated) -// tf.batch_matmul or a tf.matmul that has rank 3. dims(0) is the batch count -// and it can be trivially unrolled into a series of matmuls on each element. -// -// Inputs: -// inputs[0]: required: the left-hand side matrix -// inputs[1]: required: the right-hand side matrix -// -// TensorFlow equivalent: MatMul -struct BatchMatMulOperator : Operator { - BatchMatMulOperator() : Operator(OperatorType::kBatchMatMul) {} -}; - -// General matrix multiplication operator. We don't want to support general -// matrix multiplication at inference time, so we resolve it during tooling -// to more specific operator types, namely, FullyConnected. -// -// Inputs: -// inputs[0]: required: the left-hand side matrix -// inputs[1]: required: the right-hand side matrix -// -// TensorFlow equivalent: MatMul -struct TensorFlowMatMulOperator : Operator { - TensorFlowMatMulOperator() : Operator(OperatorType::kMatMul) {} - bool transpose_a = false; - bool transpose_b = false; -}; - -// Padding operator. Pads a tensor with zeros. -// -// Inputs: -// inputs[0]: required: the input array -// inputs[1]: required: the padding array -// -// This operation pads a `input` with zeros according to the `paddings` you -// specify. `paddings` is an integer tensor with shape `[Dn, 2]`, where n is the -// rank of `input`. For each dimension D of `input`, `paddings[D, 0]` indicates -// how many zeros to add before the contents of `input` in that dimension, and -// `paddings[D, 1]` indicates how many zeros to add after the contents of -// `input` in that dimension. -// -// TensorFlow equivalent: Pad -struct PadOperator : Operator { - PadOperator() : Operator(OperatorType::kPad) {} - - std::vector left_padding; - std::vector right_padding; -}; - -// PaddingV2 operator. Pads a tensor with the given constant value. -// -// Inputs: -// inputs[0]: required: the input array -// inputs[1]: required: the padding array -// inputs[2]: required: the scalar constant_values -// -// This operation pads input according to the paddings and constant_values you -// specify. paddings is an integer tensor with shape [Dn, 2], where n is the -// rank of input. For each dimension D of input, paddings[D, 0] indicates how -// many padding values to add before the contents of input in that dimension, -// and paddings[D, 1] indicates how many padding values to add after the -// contents of input in that dimension. constant_values is a scalar tensor of -// the same type as input that indicates the value to use for padding input. -// -// TensorFlow equivalent: PadV2 -struct PadV2Operator : Operator { - PadV2Operator() : Operator(OperatorType::kPadV2) {} - - std::vector left_padding; - std::vector right_padding; -}; - -// Strided slice operator. -// -// Inputs: -// inputs[0]: required: the input array -// inputs[1]: required: the begin array -// inputs[2]: required: the end array -// inputs[3]: optional: the strides array -// -// TensorFlow equivalent: StridedSlice -struct StridedSliceOperator : Operator { - StridedSliceOperator() : Operator(OperatorType::kStridedSlice) {} - - std::vector start_indices; - std::vector stop_indices; - std::vector strides; - - int begin_mask; - int ellipsis_mask; - int end_mask; - int new_axis_mask; - int shrink_axis_mask; - - StridedSliceOperator(const StridedSliceOperator& other) - : Operator(OperatorType::kStridedSlice) { - inputs = other.inputs; - outputs = other.outputs; - - start_indices = other.start_indices; - stop_indices = other.stop_indices; - strides = other.strides; - - begin_mask = other.begin_mask; - ellipsis_mask = other.ellipsis_mask; - end_mask = other.end_mask; - new_axis_mask = other.new_axis_mask; - shrink_axis_mask = other.shrink_axis_mask; - } - - void PadIndices(int dim_count) { - // Add indices and mask bits to fully include extra dimensions - CHECK_GE(dim_count, start_indices.size()); - CHECK_EQ(start_indices.size(), stop_indices.size()); - CHECK_EQ(stop_indices.size(), strides.size()); - - for (int i = start_indices.size(); i < dim_count; i++) { - start_indices.push_back(0); - stop_indices.push_back(0); - strides.push_back(1); - begin_mask |= 1 << i; - end_mask |= 1 << i; - } - } - - void ReverseIndices() { - CHECK_EQ(start_indices.size(), stop_indices.size()); - CHECK_EQ(stop_indices.size(), strides.size()); - - std::reverse(start_indices.begin(), start_indices.end()); - std::reverse(stop_indices.begin(), stop_indices.end()); - std::reverse(strides.begin(), strides.end()); - - begin_mask = toco::port::ReverseBits32(static_cast(begin_mask)) >> - (32 - start_indices.size()); - ellipsis_mask = - toco::port::ReverseBits32(static_cast(ellipsis_mask)) >> - (32 - start_indices.size()); - end_mask = toco::port::ReverseBits32(static_cast(end_mask)) >> - (32 - start_indices.size()); - new_axis_mask = - toco::port::ReverseBits32(static_cast(new_axis_mask)) >> - (32 - start_indices.size()); - shrink_axis_mask = - toco::port::ReverseBits32(static_cast(shrink_axis_mask)) >> - (32 - start_indices.size()); - } -}; - -// Reshaping operator, reshaping its input array to a two-dimensional shape -// (a "matrix"). This is used in the TensorFlow format, in conjunction with -// MatMul nodes, to implement fully-connected layers. -// -// Inputs: -// inputs[0]: required: the input array -// -// TensorFlow equivalent: Reshape --- except that we only support a special case -// here, where the output shape is a matrix (2D) shape. -struct TensorFlowReshapeOperator : Operator { - TensorFlowReshapeOperator() : Operator(OperatorType::kReshape) {} - std::vector shape; -}; - -// Removes dimensions of size 1 from the shape of a tensor. -// https://www.tensorflow.org/api_docs/python/tf/squeeze -// -// Inputs: -// inputs[0]: required: the input array -// -// TensorFlow equivalent: Squeeze -struct SqueezeOperator : Operator { - SqueezeOperator() : Operator(OperatorType::kSqueeze) {} - - std::vector squeeze_dims; -}; - -// Inputs: -// inputs[0]: required: the output shape -// inputs[1]: required: the weights -// inputs[2]: required: the input activations array -// NOTE: The input activations is NOT the first input. -// -// -// Outputs: -// outputs[0]: required: the output activations array -// -// TensorFlow equivalent: Conv2DBackpropInput -struct TransposeConvOperator : Operator { - enum Inputs { - OUTPUT_SHAPE = 0, - WEIGHTS = 1, - DATA_INPUT = 2, - }; - - TransposeConvOperator() : Operator(OperatorType::kTransposeConv) {} - Padding padding; - int stride_width = 0; - int stride_height = 0; - // Dilation is possible with transpose convolution, but Tensorflow does not - // currently support it, so we omit it. -}; - -// Given a tensor input, this operation calculates element-wise exponential -// (y = e^x). -// -// Inputs: -// inputs[0]: required: input tensor -// -// TensorFlow equivalent: Exp -struct ExpOperator : Operator { - ExpOperator() : Operator(OperatorType::kExp) {} -}; - -// Given a tensor input, this operation inserts a dimension of 1 at the -// dimension index axis of input's shape. The dimension index axis starts at -// zero; if you specify a negative number for axis it is counted backward from -// the end. -// -// Inputs: -// inputs[0]: required: input tensor -// inputs[1]: required: 0-D (scalar). Specifies the dimension index at which -// to expand the shape of input -// -// TensorFlow equivalent: ExpandDims -struct ExpandDimsOperator : Operator { - ExpandDimsOperator() : Operator(OperatorType::kExpandDims) {} -}; - -// Ceates a tensor of shape dims and fills it with the given scalar value. -// Output type will be the same as the given scalar value. -// -// Inputs: -// inputs[0]: required: 1-D (int32) - the shape of the output tensor -// inputs[1]: required: 0-D (scalar) - value to fill the tensor with -// -// TensorFlow equivalent: Fill -struct FillOperator : Operator { - FillOperator() : Operator(OperatorType::kFill) {} -}; - -// Element-wise floor division operator. -// -// Inputs: -// inputs[0]: required: the left-hand side array -// inputs[1]: required: the right-hand side array -// -// TensorFlow equivalent: FloorDiv -struct FloorDivOperator : Operator { - FloorDivOperator() : Operator(OperatorType::kFloorDiv) {} -}; - -// Element-wise floor mod operator. -// -// Inputs: -// inputs[0]: required: the left-hand side array -// inputs[1]: required: the right-hand side array -// -// TensorFlow equivalent: FloorMod -struct FloorModOperator : Operator { - FloorModOperator() : Operator(OperatorType::kFloorMod) {} -}; - -struct RandomUniformOperator : Operator { - RandomUniformOperator() : Operator(OperatorType::kRandomUniform) {} - ArrayDataType dtype = ArrayDataType::kNone; - int64 seed; - int64 seed2; -}; - -// Creates a sequence of numbers that begins at start and extends by increments -// of delta up to but not including limit. -// -// The dtype of the resulting tensor is inferred from the inputs unless it is -// provided explicitly. -// -// Inputs: -// inputs[0]: required: the start -// inputs[1]: required: the limit -// inputs[2]: required: the delta -// -// TensorFlow equivalent: Range -struct RangeOperator : Operator { - RangeOperator() : Operator(OperatorType::kRange) {} - ArrayDataType dtype = ArrayDataType::kNone; -}; - -// Rank operator. Extracts the rank of the tensor. -// -// Inputs: -// inputs[0]: required: the input array -// -// This operation outputs a 0-D integer tensor representing the rank of -// the input. -// -// TensorFlow equivalent: Rank. We currently assume that the output is int32 -// and not int64. The output type could be stored herein. -struct RankOperator : Operator { - RankOperator() : Operator(OperatorType::kRank) {} -}; - -// Element-wise negation (-x) operator. -// -// Inputs: -// inputs[0]: required: the input array -// -// TensorFlow equivalent: Neg -struct NegOperator : Operator { - NegOperator() : Operator(OperatorType::kNeg) {} -}; - -// Element-wise select operator choosing elements from inputs[1] or input[2] -// -// Inputs: -// inputs[0]: required: boolean mask per index -// inputs[1]: required: tensor of values if true -// inputs[2]: required: tensor of values if false -// -// TensorFlow equivalent: Select -struct SelectOperator : Operator { - SelectOperator() : Operator(OperatorType::kSelect) {} -}; - -// Element-wise reciprocal-square-root (x^-0.5) operator. -// -// Inputs: -// inputs[0]: required: the input array -// -// TensorFlow equivalent: Rsqrt -struct TensorFlowRsqrtOperator : Operator { - TensorFlowRsqrtOperator() : Operator(OperatorType::kRsqrt) {} -}; - -// Stacks a list of rank-R tensors into one rank-(R+1) tensor. -// -// Packs the list of tensors in values into a tensor with rank one higher than -// each tensor in values, by packing them along the axis dimension. Given a list -// of length N of tensors of shape (A, B, C);. -// -// Inputs: this operator accepts any number >= 1 of inputs. -// inputs[i]: the i-th array to merge. -// -// TensorFlow equivalent: Pack -struct PackOperator : Operator { - PackOperator() : Operator(OperatorType::kPack) {} - int values_count; - int axis = 0; - ArrayDataType dtype = ArrayDataType::kNone; -}; - -// Shape operator. Extracts the shape of the tensor. -// -// Inputs: -// inputs[0]: required: the input array -// -// This operation outputs a 1-D integer tensor representing the shape of -// the input. -// -// TensorFlow equivalent: Shape. -struct TensorFlowShapeOperator : Operator { - TensorFlowShapeOperator() : Operator(OperatorType::kShape) {} - ArrayDataType output_data_type = ArrayDataType::kInt32; -}; - -// Element-wise square-root (x^0.5) operator. -// -// Inputs: -// inputs[0]: required: the input array -// -// TensorFlow equivalent: Sqrt -struct TensorFlowSqrtOperator : Operator { - TensorFlowSqrtOperator() : Operator(OperatorType::kSqrt) {} -}; - -// Element-wise square (x*x) operator. -// -// Inputs: -// inputs[0]: required: the input array -// -// TensorFlow equivalent: Square -struct TensorFlowSquareOperator : Operator { - TensorFlowSquareOperator() : Operator(OperatorType::kSquare) {} -}; - -// Transposes a tensor. -// -// By default, this operation performs a regular matrix transpose on 2-D input -// tensors. -// -// Inputs: -// inputs[0]: required: the input array -// -// TensorFlow equivalent: Transpose -struct TransposeOperator : Operator { - TransposeOperator() : Operator(OperatorType::kTranspose) {} - std::vector perm; -}; - -// Element-wise subtraction operator. -// -// Inputs: -// inputs[0]: required: the left-hand side array -// inputs[1]: required: the right-hand side array -// -// TensorFlow equivalent: Sub -struct SubOperator : Operator { - SubOperator() : Operator(OperatorType::kSub) {} -}; - -// Sum reduction: computes the sum of all of entries across the axes. -// -// Inputs: -// inputs[0]: required: the input array -// -// TensorFlow equivalent: Sum -struct TensorFlowSumOperator : Operator { - TensorFlowSumOperator() : Operator(OperatorType::kSum) {} - std::vector axis; - bool keep_dims = false; -}; - -// Prod reduction: computes the product of all of entries across the axes. -// -// Inputs: -// inputs[0]: required: the input array -// -// TensorFlow equivalent: Prod -struct TensorFlowProdOperator : Operator { - TensorFlowProdOperator() : Operator(OperatorType::kReduceProd) {} - std::vector axis; - bool keep_dims = false; -}; - -// TensorFlow Tile equivalent. Refer to TensorFlow documentation for details. -// -// Inputs: -// inputs[0]: required: the input array -// inputs[1]: required: int array with length of rank(input[0]) -struct TensorFlowTileOperator : Operator { - TensorFlowTileOperator() : Operator(OperatorType::kTile) {} -}; - -// TensorFlow Slice equivalent. Refer to TensorFlow documentation for details. -struct SliceOperator : Operator { - SliceOperator() : Operator(OperatorType::kSlice) {} - - std::vector begin; - std::vector size; -}; - -// TensorFlow Split equivalent. Refer to TensorFlow documentation for details. -// Not fully supported, just a placeholder to handle TensorFlow graphs and -// support graph transformations to other operator types by matching sub-graphs. -struct TensorFlowSplitOperator : Operator { - TensorFlowSplitOperator() : Operator(OperatorType::kSplit) {} - int num_split = 0; -}; - -// TensorFlow Concat equivalent. Refer to TensorFlow documentation for details. -// Not fully supported, just a placeholder to handle TensorFlow graphs and -// support graph transformations to other operator types by matching sub-graphs. -// Concretely, once the concat dim becomes known, if it is the depth -// dimension then we can change this op into a DepthConcatenation op. -// Otherwise, we hope for some other graph transformation to drop this node. -struct TensorFlowConcatOperator : Operator { - TensorFlowConcatOperator() : Operator(OperatorType::kConcat) {} -}; - -// TensorFlow ConcatV2 equivalent. Refer to TensorFlow documentation for -// details. -// Not fully supported, just a placeholder to handle TensorFlow graphs and -// support graph transformations to other operator types by matching sub-graphs. -// Concretely, once the concat dim becomes known, if it is the depth -// dimension then we can change this op into a DepthConcatenation op. -// Otherwise, we hope for some other graph transformation to drop this node. -struct TensorFlowConcatV2Operator : Operator { - TensorFlowConcatV2Operator() : Operator(OperatorType::kConcatV2) {} -}; - -// TensorFlow Merge equivalent. Refer to TensorFlow documentation for details. -// -// Inputs: this operator accepts any number >= 1 of inputs. -// inputs[i]: the i-th array to merge. -// -// It is expected that graph transformations will drop all but exactly one -// of the inputs, at which point the Merge node will be equivalent to an -// Identity node forwarding the remaining input. -// -// Note: We do not currently support runtime control flow: we only support -// control flow that can be resolved at tooling time (independently of input -// activations). -struct TensorFlowMergeOperator : Operator { - TensorFlowMergeOperator() : Operator(OperatorType::kMerge) {} -}; - -// TensorFlow Switch equivalent. Refer to TensorFlow documentation for details. -// -// Inputs: -// inputs[0]: required: the input array -// inputs[1]: required: the boolean predicate, given as an array of size 1 -// and of type kBool, will determine which output gets selected. -// -// Outputs: a TensorFlow Switch node always has exactly two outputs. Depending -// on the boolean value that the input predicate resolves to (see note below), -// one or the other of the outputs will be 'selected': the input array will be -// forwarded to the 'selected output' as if by a Identity node, while the other -// output will be discarded, and any graph edge connecting that discarded output -// will be dropped. The rule for selecting outputs is as follows: -// outputs[0] will be selected if the input predicate resolves to 'true'. -// outputs[1] will be selected if the input predicate resolves to 'false'. -// -// Note: We do not currently support runtime control flow: we only support -// control flow that can be resolved at tooling time (independently of input -// activations). -struct TensorFlowSwitchOperator : Operator { - TensorFlowSwitchOperator() : Operator(OperatorType::kSwitch) {} -}; - -// TensorFlow All equivalent. Refer to TensorFlow documentation for details. -// Not fully supported, just a placeholder to handle TensorFlow graphs and -// support graph transformations to other operator types by matching sub-graphs. -// Typically, this is only used as an input to an Assert node, so can be -// removed as an unused node as we drop Assert nodes. -struct TensorFlowAllOperator : Operator { - TensorFlowAllOperator() : Operator(OperatorType::kAll) {} -}; - -// TensorFlow Assert equivalent. Refer to TensorFlow documentation for details. -// Not fully supported, just a placeholder to handle TensorFlow graphs and -// support graph transformations to other operator types by matching sub-graphs. -// Typically, we just drop Assert nodes. -struct TensorFlowAssertOperator : Operator { - TensorFlowAssertOperator() : Operator(OperatorType::kAssert) {} -}; - -// TensorFlow Less equivalent. Refer to TensorFlow documentation for details. -// Not fully supported, just a placeholder to handle TensorFlow graphs and -// support graph transformations to other operator types by matching sub-graphs. -// Typically, this is only used as an input to an Assert node, so can be -// removed as an unused node as we drop Assert nodes. -struct TensorFlowLessOperator : Operator { - TensorFlowLessOperator() : Operator(OperatorType::kLess) {} -}; - -// TensorFlow LessEqual equivalent. Refer to TensorFlow documentation for -// details. -// Not fully supported, just a placeholder to handle TensorFlow graphs and -// support graph transformations to other operator types by matching sub-graphs. -// Typically, this is only used as an input to an Assert node, so can be -// removed as an unused node as we drop Assert nodes. -struct TensorFlowLessEqualOperator : Operator { - TensorFlowLessEqualOperator() : Operator(OperatorType::kLessEqual) {} -}; - -// TensorFlow Less equivalent. Refer to TensorFlow documentation for details. -// Not fully supported, just a placeholder to handle TensorFlow graphs and -// support graph transformations to other operator types by matching sub-graphs. -// Typically, this is only used as an input to an Assert node, so can be -// removed as an unused node as we drop Assert nodes. -struct TensorFlowGreaterOperator : Operator { - TensorFlowGreaterOperator() : Operator(OperatorType::kGreater) {} -}; - -// TensorFlow GreaterEqual equivalent. Refer to TensorFlow documentation for -// details. -// Not fully supported, just a placeholder to handle TensorFlow graphs and -// support graph transformations to other operator types by matching sub-graphs. -// Typically, this is only used as an input to an Assert node, so can be -// removed as an unused node as we drop Assert nodes. -struct TensorFlowGreaterEqualOperator : Operator { - TensorFlowGreaterEqualOperator() : Operator(OperatorType::kGreaterEqual) {} -}; - -// TensorFlow Equal equivalent. Refer to TensorFlow documentation for -// details. -// Not fully supported, just a placeholder to handle TensorFlow graphs and -// support graph transformations to other operator types by matching sub-graphs. -// Typically, this is only used as an input to an Assert node, so can be -// removed as an unused node as we drop Assert nodes. -struct TensorFlowEqualOperator : Operator { - TensorFlowEqualOperator() : Operator(OperatorType::kEqual) {} -}; - -// TensorFlow Not Equal equivalent. Refer to TensorFlow documentation for -// details. -struct TensorFlowNotEqualOperator : Operator { - TensorFlowNotEqualOperator() : Operator(OperatorType::kNotEqual) {} -}; - -// Max reduction: computes the max of all of entries across the axes. -// -// Inputs: -// inputs[0]: required: the input array -// -// TensorFlow equivalent: Max -struct TensorFlowMaxOperator : Operator { - TensorFlowMaxOperator() : Operator(OperatorType::kReduceMax) {} - std::vector axis; - bool keep_dims = false; -}; - -// Min reduction: computes the min of all of entries across the axes. -// -// Inputs: -// inputs[0]: required: the input array -// -// TensorFlow equivalent: Min -struct TensorFlowMinOperator : Operator { - TensorFlowMinOperator() : Operator(OperatorType::kReduceMin) {} - std::vector axis; - bool keep_dims = false; -}; - -// Element-wise maximum operator. Currently it only supports scalar as -// the second operand. -// -// Inputs: -// inputs[0]: required: the left-hand side array -// inputs[1]: required: the right-hand side array -// -// TensorFlow equivalent: Maximum -struct TensorFlowMaximumOperator : Operator { - TensorFlowMaximumOperator() : Operator(OperatorType::kMaximum) {} -}; - -// Element-wise minimum operator. Currently it only supports scalar as -// the second operand. -// -// Inputs: -// inputs[0]: required: the left-hand side array -// inputs[1]: required: the right-hand side array -// -// TensorFlow equivalent: Minimum -struct TensorFlowMinimumOperator : Operator { - TensorFlowMinimumOperator() : Operator(OperatorType::kMinimum) {} -}; - -// General TF operation, unsupported by tf.mini. Expected to be dropped by -// graph transformations. -struct TensorFlowUnsupportedOperator : Operator { - TensorFlowUnsupportedOperator() : Operator(OperatorType::kUnsupported) {} - - // The original TF operation type. Used for diagnostic purposes. - string tensorflow_op; - // A boolean indicating if the unsupported op should be treated as quantized. - bool quantized = false; - // A boolean indicating if the unsupported op output should allow float values - // in quantized mode. - bool support_output_type_float_in_quantized_op = false; - // Output data types - std::vector output_data_types; - // Output shapes. - std::vector output_shapes; -}; - -// Softmax activation function. -// -// Inputs: -// inputs[0]: required: the input array -// -// TensorFlow equivalent: Softmax -struct SoftmaxOperator : Operator { - SoftmaxOperator() : Operator(OperatorType::kSoftmax) {} - float beta = 0.f; -}; - -// LogSoftmax activation function. -// -// Inputs: -// inputs[0]: required: the logits input array -// -// TensorFlow equivalent: LogSoftmax -struct LogSoftmaxOperator : Operator { - LogSoftmaxOperator() : Operator(OperatorType::kLogSoftmax) {} - - // LogSoftmax can in principal have very large negative output, depending on - // the input size. However, input x_i that is less than x_max-10 is - // accumulated as exp(x_i-x_max), which is truncated to zero. - // - // Since we effectively disregard smallish inputs in the normalizing factor, - // we also drop them in the output (set to minimum output), and in doing so - // make better use of the quantization range / resolution. - static constexpr float kOutputRangeMin = -16.0; -}; - -// Cast operator. -// -// Inputs: -// inputs[0]: required: the input array -// -// TensorFlow equivalent: Cast -struct CastOperator : Operator { - CastOperator() : Operator(OperatorType::kCast) {} - ArrayDataType src_data_type = ArrayDataType::kNone; - ArrayDataType dst_data_type = ArrayDataType::kNone; -}; - -// Floor operator. -// -// Inputs: -// inputs[0]: required: the input array -// -// TensorFlow equivalent: Floor -struct FloorOperator : Operator { - FloorOperator() : Operator(OperatorType::kFloor) {} -}; - -// Gather operator. It gathers slices from params according to indices. -// Only 1-D indices are supported at the moment. -// -// Inputs: -// inputs[0]: required: the params array -// inputs[1]: required: the indices to gather -// inputs[2]: optional: axis -// -// TensorFlow equivalent: Gather -struct GatherOperator : Operator { - GatherOperator() : Operator(OperatorType::kGather) {} - // Axis is populated explicitly or implicitly from the axis input by - // ResolveGatherAttributes. An empty axis indicates that the axis has not yet - // be resolved. - absl::optional axis; - int input_rank = 0; -}; - -// ArgMax operator. It returns the index of the maximum value along axis. -// -// Inputs: -// inputs[0]: required: the input tensor -// -// TensorFlow equivalent: ArgMax -struct ArgMaxOperator : Operator { - ArgMaxOperator() : Operator(OperatorType::kArgMax) {} - ArrayDataType output_data_type = ArrayDataType::kInt64; -}; - -// ArgMin operator. It returns the index of the minimum value along axis. -// -// Inputs: -// inputs[0]: required: the input tensor -// -// TensorFlow equivalent: ArgMin -struct ArgMinOperator : Operator { - ArgMinOperator() : Operator(OperatorType::kArgMin) {} - ArrayDataType output_data_type = ArrayDataType::kInt64; -}; - -// ResizeBilinear operator. It resizes input images with bilinear interpolation. -// It does not support align_corners at the moment. -// -// Inputs: -// inputs[0]: required: the input array -// inputs[1]: required: the new image size -// -// TensorFlow equivalent: ResizeBilinear -struct ResizeBilinearOperator : Operator { - ResizeBilinearOperator() : Operator(OperatorType::kResizeBilinear) {} - - bool align_corners = false; -}; - -// SpaceToBatchND operator. It divides spatial dimensions into a grid of -// blocks and interleaves these blocks with the batch dimension. Currently, -// only 2-d blocks are supported. -// -// Inputs: -// inputs[0]: required: the input array -// inputs[1]: required: the block shape -// inputs[2]: required: the paddings -// -// TensorFlow equivalent: SpaceToBatchND -struct SpaceToBatchNDOperator : Operator { - SpaceToBatchNDOperator() : Operator(OperatorType::kSpaceToBatchND) {} - - std::vector block_shape; - std::vector before_paddings; - std::vector after_paddings; -}; - -// BatchToSpaceND operator. Rearranges data from batch into blocks of -// spatial data. Currently, only 2-d blocks are supported. -// -// Inputs: -// inputs[0]: required: the input array -// inputs[1]: required: the block shape -// inputs[2]: required: the crops -// -// TensorFlow equivalent: BatchToSpaceND -struct BatchToSpaceNDOperator : Operator { - BatchToSpaceNDOperator() : Operator(OperatorType::kBatchToSpaceND) {} - - std::vector block_shape; - std::vector before_crops; - std::vector after_crops; -}; - -// Mean operator. -// -// Inputs: -// inputs[0]: required: the input array -// -// TensorFlow equivalent: Mean -struct MeanOperator : Operator { - MeanOperator() : Operator(OperatorType::kMean) {} - - std::vector axis; - bool keep_dims = false; -}; - -// Svdf operator: -// -// Inputs: -// inputs[0]: required: the input array -// inputs[1]: required: weights_feature -// inputs[2]: required: weights_time -// inputs[3]: optional: bias -struct SvdfOperator : Operator { - SvdfOperator() : Operator(OperatorType::kSvdf) {} - int rank; -}; - -// TopKV2 operator. -// -// Inputs: -// input tensor and top_k scalar. -struct TopKV2Operator : Operator { - TopKV2Operator() : Operator(OperatorType::kTopK_V2) {} -}; - -// DynamicPartition operator: -// -// Inputs: -// inputs[0]: required: data. -// inputs[1]: required: partitions. -// -// TensorFlow equivalent: DynamicPartition -struct DynamicPartitionOperator : Operator { - DynamicPartitionOperator() : Operator(OperatorType::kDynamicPartition) {} - int num_partitions; -}; - -// DynamicStitch operator: -// -// Inputs: -// inputs[0,N): required: indices. -// inputs[N,2N): required: data. -// -// TensorFlow equivalent: DynamicStitch/ParallelDynamicStitch -struct DynamicStitchOperator : Operator { - DynamicStitchOperator() : Operator(OperatorType::kDynamicStitch) {} - int num_partitions; -}; - -// SparseToDense operator: -// -// Inputs: -// Inputs[0]: required: sparse_indices. -// Inputs[1]: required: output_shape. -// Inputs[2]: required: sparse_values. -// -// TensorFlow equivalent: SparseToDense. -struct SparseToDenseOperator : Operator { - SparseToDenseOperator() : Operator(OperatorType::kSparseToDense) {} - bool validate_indices; -}; - -// Pow operator: -// -// Inputs: -// Inputs[0]: required: A tensor. -// Inputs[1]: required: A tensor. -// -// TensorFlow equivalent: Pow. -struct PowOperator : Operator { - PowOperator() : Operator(OperatorType::kPow) {} -}; - -// Any operator: -// -// Inputs: -// Inputs[0]: required: A boolean input tensor. -// Inputs[1]: required: reduction_indices. -// -// TensorFlow equivalent: tf.reduce_any. -struct TensorFlowAnyOperator : Operator { - TensorFlowAnyOperator() : Operator(OperatorType::kAny) {} - std::vector axis; - bool keep_dims = false; -}; - -// LogicalAnd operator: -// -// Inputs: -// Inputs[0]: required: A boolean tensor. -// Inputs[1]: required: A boolean tensor. -// -// TensorFlow equivalent: tf.logical_and. -struct LogicalAndOperator : Operator { - LogicalAndOperator() : Operator(OperatorType::kLogicalAnd) {} -}; - -// LogicalNot operator: -// -// Inputs: -// Inputs[0]: required: A boolean tensor. -// -// TensorFlow equivalent: tf.logical_not. -struct LogicalNotOperator : Operator { - LogicalNotOperator() : Operator(OperatorType::kLogicalNot) {} -}; - -// OneHot operator: -// -// Inputs: -// Inputs[0]: required: indices. -// Inputs[1]: required: depth. -// Inputs[2]: required: on_value. -// Inputs[3]: required: off_value. -// -// TensorFlow equivalent: OneHot. -struct OneHotOperator : Operator { - enum Inputs { - INDICES_INPUT = 0, - DEPTH_INPUT = 1, - ON_VALUE_INPUT = 2, - OFF_VALUE_INPUT = 3, - }; - - OneHotOperator() : Operator(OperatorType::kOneHot) {} - int axis = -1; -}; - -// LogicalOr operator: -// -// Inputs: -// Inputs[0]: required: A Bool tensor. -// Inputs[1]: required: A Bool tensor. -// -// TensorFlow equivalent: LogicalOr. -struct LogicalOrOperator : Operator { - LogicalOrOperator() : Operator(OperatorType::kLogicalOr) {} -}; - -// Unpack operator: -// -// Inputs: -// Inputs[0]: required: A boolean input tensor. -// Inputs[1]: required: reduction_indices. -// -// TensorFlow equivalent: tf.unstack. -struct UnpackOperator : Operator { - UnpackOperator() : Operator(OperatorType::kUnpack) {} - int num; - int axis; - ArrayDataType dtype = ArrayDataType::kNone; -}; - -// ZerosLike operator: -// -// Inputs: -// inputs[0]: required: the input array -// -// TensorFlow equivalent: tf.zeros_like -struct TensorFlowZerosLikeOperator : Operator { - TensorFlowZerosLikeOperator() : Operator(OperatorType::kZerosLike) {} -}; - -// Alloc's are used for transient arrays only. An Alloc specifies which interval -// of the "transient_data" workspace buffer passed to inference functions, is to -// be used for the transient array at hand. The 'start' and 'end' values are -// offsets from the start of the workspace buffer, expressed in bytes. -struct Alloc { - int64 start = 0; - int64 end = 0; -}; - -inline bool operator<(const Alloc& a, const Alloc& b) { - return a.start < b.start; -} - -// Array represents an array (either a constant parameter array or an -// activations array) in a Model. -struct Array { - template - const Buffer& GetBuffer() const { - DCHECK(buffer); - DCHECK(buffer->type == A); - return *static_cast*>(buffer.get()); - } - template - Buffer& GetMutableBuffer() { - if (!buffer) { - Buffer* ptr = new Buffer; - buffer = std::unique_ptr(ptr); - } - DCHECK(buffer); - DCHECK(buffer->type == A); - return *static_cast*>(buffer.get()); - } - Alloc& GetOrCreateAlloc() { - if (!alloc) { - alloc = std::unique_ptr(new Alloc); - } - return *alloc; - } - MinMax& GetOrCreateMinMax() { - if (!minmax) { - minmax = std::unique_ptr(new MinMax); - } - return *minmax; - } - MinMax& GetMinMax() const { - DCHECK(minmax); - return *minmax; - } - QuantizationParams& GetOrCreateQuantizationParams() { - if (!quantization_params) { - quantization_params = - std::unique_ptr(new QuantizationParams); - } - return *quantization_params; - } - QuantizationParams& GetQuantizationParams() const { - DCHECK(quantization_params); - return *quantization_params; - } - - // The data type of the actual elements of this array, that is: - // - If there is a buffer (see 'buffer' member), it must be of the same - // type. - // - If there is no buffer, meaning that this is a runtime (i.e. activations) - // array, then this specifies the type of elements that there will be - // at runtime. - // - // Note that this only specifies the storage type of elements; this does - // not specify whether these are to be treated as 'real' or 'quantized' - // values. - // That is decided by whether the 'quantization_params' member is null. - ArrayDataType data_type = ArrayDataType::kNone; - // The final value that data_type should have at the end of graph - // transformations - ArrayDataType final_data_type = ArrayDataType::kNone; - // The dimensions of this array --- this specifies both sizes and strides - // (the storage layout). - // - // Issues with shape handling that remain include: - // - No way to distinguish between 0-dimensional dims and missing dims. - // - No way to describe dims that may be runtime-variable. - // - Addressing of dims by integer index differs in different graph formats - // (TensorFlow vs. other frameworks vs. what we have informally grown - // within toco). - // This is currently quite messy; see ReorderAxesOperator which is how we - // bridge some of these discrepancies at the moment. This is overdue for - // a redesign; I'm thinking that it would be nice to have more flexible - // dims that allow mapping 1:1, cleanly, dims as they are in various - // formats, - // then explicitly convert between different conventions. - - // Proto-style accessors - bool has_shape() const { return array_shape != nullptr; } - const Shape& shape() const { - CHECK(has_shape()); - return *array_shape; - } - Shape* mutable_shape() { - if (!array_shape) { - array_shape.reset(new Shape); - } - return array_shape.get(); - } - void copy_shape(const Shape& src_shape) { *mutable_shape() = src_shape; } - void clear_shape() { array_shape = nullptr; } - - // The constant buffer backing this array. This is non-null if and only if - // this is a constant parameter array. Conversely, this is null for - // activations arrays. - // - // Note that this buffer is pure storage. In the case of quantized values, - // it only stores the quantized values, it does not know by itself about the - // quantization parameters necessary to interprete these values, that is - // in the separate 'quantization_params' field. In fact, this 'buffer' field - // does no even know whether values are quantized. It only has a data_type, - // which must equal the 'data_type' member here, and which only describes - // the storage type of element, does not tell whether they are quantized i.e. - // whether they are to be interpreted with quantization_params. - std::unique_ptr buffer; - // Only for activation arrays (i.e. when 'buffer' is null). - // Only for code generation. - // - // Describes the allocation of this array within the workspace buffer - // allocated - // for all transient arrays. - std::unique_ptr alloc; - // Describes the [min, max] range of values - // to be assumed when determining quantization_params. - // - // Only used for quantization. In fact, only used for determining - // quantization_params. - // - // Used for both constant arrays (those having a 'buffer') and non-constant - // arrays (activations). Indeed, it is important to use the same min-max range - // as was used during training, even if that min-max range is slightly wrong - // w.r.t. actual buffer elements. Doing otherwise would defeat the point of - // re-training for quantization. - std::unique_ptr minmax; - // Quantization parameters. The non-null-ness of this pointer is what - // defines whether this array is quantized or not. - // - // If this is non-null, then these quantization parameters are to be used - // to assign a meaning as real numbers to the elements of this array. - std::unique_ptr quantization_params; - // narrow_range is a detail of how toco handles FakeQuant operators with - // narrow_range, see - // https://www.tensorflow.org/api_docs/python/tf/fake_quant_with_min_max_vars - // - // For more context about what that is useful for, see the big comment in - // graph_transformations/ensure_uint8_weights_safe_for_fast_int8_kernels.cc - // - // The narrow_range flag applies only to quantized arrays, and changes - // their quantization in the following way when it is set to 'true': - // 1. The computation of {zero_point, scale} from {min, max} needs to be - // amended so that the real min value will get quantized to - // (min_quantized_value + 1) instead of just (min_quantized_value). - // E.g. for uint8 quantization, the real min value should get quantized to - // the uint8 value 1, not 0. - // 2. Quantized values should get clamped to the interval - // [min_quantized_value + 1, max_value]. Equivalently, the - // min_quantized_value should get nudged to (min_quantized_value + 1). - // The reason why 1. does not imply 2. is that real values may not belong to - // the stated [min, max] interval. Concretely, weights recorded at the last - // learning step may not fall in the [min, max] interval recorded over - // previous learning steps, as the values evolve across learning steps. - // - // Rationale why this is directly a field on Array: - // - This can't be just a field on FakeQuantOperator, because - // FakeQuantOperators are gone (DropFakeQuant) before we get to using that - // information (Quantize). We need a place to store that bit in the interim. - // - This can't be in QuantizationParams because we need to record this - // ahead of quantization, and QuantizationParams are only created during - // quantization. - // - This could be in MinMax, but that would be an abuse of what MinMax is - // about, and would break existing code that assumes that a MinMax is just - // a min and a max. Unlike MinMax which is agnostic as to the quantized - // data type, narrow_range refers to values in the quantized data type. - bool narrow_range = false; - - private: - std::unique_ptr array_shape; -}; - -// Our Model struct, represents an entire model (our "top-level" struct). -// Owns everything. -class Model { - public: - using ArrayMap = std::unordered_map>; - - bool HasArray(const string& name) const { return arrays.count(name) > 0; } - Array& GetArray(const string& name) const { - DCHECK(HasArray(name)) << "Array not found: " << name; - return *arrays.at(name); - } - Array& GetOrCreateArray(const string& name) { - // Make sure name is not used by an optional array - DCHECK(!optional_arrays.count(name)); - if (!HasArray(name)) { - Array* ptr = new Array; - arrays[name] = std::unique_ptr(ptr); - } - Array& result = GetArray(name); - return result; - } - void CreateOptionalArray(const string& name) { - DCHECK(!arrays.count(name) && !optional_arrays.count(name)); - optional_arrays.insert(name); - } - bool IsOptionalArray(const string& name) const { - return optional_arrays.count(name); - } - - // Note that this invalidates all array iterators. - void EraseArray(const string& name) { arrays.erase(name); } - void EraseArrays(std::function discardable) { - for (auto it = arrays.begin(); it != arrays.end();) { - if (discardable(it->first)) { - it = arrays.erase(it); - } else { - ++it; - } - } - } - const ArrayMap& GetArrayMap() const { return arrays; } - ArrayMap& GetMutableArrayMap() { return arrays; } - - int64 ArithmeticOpsCount() const { return ops_count; } - - // Optional arrays are used for optional tensors, - // these tensors do not have data, but with reserved names as op inputs. - std::set optional_arrays; - - // The list of operators. Notice how it's a list of unique_ptr's, implying - // that the Model is what owns Operator's and keeps them alive. - std::vector> operators; - - // Generic flags, a place where we combine information passed to us via - // command-line parameters (e.g. --input_width=N) with information that - // we may or may not find in the input model file. - ModelFlags flags; - // For code-generation only: required size of the transient_data buffer - std::size_t transient_data_size = 0; - // For code-generation only: required alignment of the transient_data buffer - std::size_t transient_data_alignment = 0; - // Arithmetic operations performed in the model. - int64 ops_count = 0; - - private: - // The associative array mapping names to Array's. - // Notice how it's a container of unique_ptr's, implying - // that the Model is what owns Array's and keeps them alive. - // The Operator's refer to these Array's by their name strings, not by their - // addresses. See Operator::inputs, Operator::outputs. - std::unordered_map> arrays; -}; -} // namespace toco - -#endif // TENSORFLOW_CONTRIB_LITE_TOCO_MODEL_H_ diff --git a/tensorflow/contrib/lite/toco/python/BUILD b/tensorflow/contrib/lite/toco/python/BUILD deleted file mode 100644 index cf97ba7084d48e55a1874e77d3817aa721de7de9..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/toco/python/BUILD +++ /dev/null @@ -1,61 +0,0 @@ -package(default_visibility = ["//visibility:public"]) - -licenses(["notice"]) # Apache 2.0 - -load("//tensorflow:tensorflow.bzl", "tf_py_wrap_cc") -load("//tensorflow:tensorflow.bzl", "tf_py_test") -load("//tensorflow:tensorflow.bzl", "py_binary") - -cc_library( - name = "toco_python_api", - srcs = ["toco_python_api.cc"], - hdrs = ["toco_python_api.h"], - deps = [ - "//tensorflow/contrib/lite/toco:model_flags_proto_cc", - "//tensorflow/contrib/lite/toco:toco_flags_proto_cc", - "//tensorflow/contrib/lite/toco:toco_graphviz_dump_options", - "//tensorflow/contrib/lite/toco:toco_port", - "//tensorflow/contrib/lite/toco:toco_tooling", - "//tensorflow/core:lib", - "//third_party/python_runtime:headers", - ], -) - -tf_py_wrap_cc( - name = "tensorflow_wrap_toco", - srcs = ["toco.i"], - deps = [ - ":toco_python_api", - "//tensorflow/contrib/lite/toco:model_flags_proto_cc", - "//tensorflow/contrib/lite/toco:toco_flags_proto_cc", - "//third_party/python_runtime:headers", - "@com_google_absl//absl/strings", - ], -) - -py_binary( - name = "toco_from_protos", - srcs = ["toco_from_protos.py"], - srcs_version = "PY2AND3", - deps = [ - ":tensorflow_wrap_toco", - "//tensorflow/python:platform", - ], -) - -tf_py_test( - name = "toco_from_protos_test", - srcs = ["toco_from_protos_test.py"], - additional_deps = [ - "//tensorflow:tensorflow_py", - "//tensorflow/contrib/lite/toco:model_flags_proto_py", - "//tensorflow/contrib/lite/toco:toco_flags_proto_py", - ], - data = [ - ":toco_from_protos", - ], - tags = [ - "no_oss", - "no_pip", - ], -) diff --git a/tensorflow/contrib/lite/toco/runtime/common.h b/tensorflow/contrib/lite/toco/runtime/common.h deleted file mode 100644 index 3c6828840c4a963a4a68774ec5d559b7f80baf22..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/toco/runtime/common.h +++ /dev/null @@ -1,26 +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_CONTRIB_LITE_TOCO_RUNTIME_COMMON_H_ -#define TENSORFLOW_CONTRIB_LITE_TOCO_RUNTIME_COMMON_H_ - -#ifndef ALLOW_SLOW_GENERIC_DEPTHWISECONV_FALLBACK -#ifdef GEMMLOWP_ALLOW_SLOW_SCALAR_FALLBACK -#define ALLOW_SLOW_GENERIC_DEPTHWISECONV_FALLBACK -#endif -#endif - -#include "tensorflow/contrib/lite/kernels/internal/common.h" - -#endif // TENSORFLOW_CONTRIB_LITE_TOCO_RUNTIME_COMMON_H_ diff --git a/tensorflow/contrib/lite/toco/runtime/types.h b/tensorflow/contrib/lite/toco/runtime/types.h deleted file mode 100644 index 207f2c1706ef4cc12572e381c38f61a504ece232..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/toco/runtime/types.h +++ /dev/null @@ -1,33 +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_CONTRIB_LITE_TOCO_RUNTIME_TYPES_H_ -#define TENSORFLOW_CONTRIB_LITE_TOCO_RUNTIME_TYPES_H_ - -#include "tensorflow/contrib/lite/kernels/internal/common.h" -#include "tensorflow/contrib/lite/kernels/internal/compatibility.h" -#include "tensorflow/contrib/lite/kernels/internal/types.h" - -namespace toco { - -// TODO(ahentz): These are just stopgaps for now, untils we move all -// the code over to tflite. -using tflite::Dims; -using tflite::FullyConnectedWeightsFormat; -using tflite::FusedActivationFunctionType; -using tflite::RequiredBufferSizeForDims; - -} // namespace toco - -#endif // TENSORFLOW_CONTRIB_LITE_TOCO_RUNTIME_TYPES_H_ diff --git a/tensorflow/contrib/lite/toco/tensorflow_graph_matching/BUILD b/tensorflow/contrib/lite/toco/tensorflow_graph_matching/BUILD deleted file mode 100644 index ea1fc2827ead7e7442bbf7f569e3ea88c3b0de57..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/toco/tensorflow_graph_matching/BUILD +++ /dev/null @@ -1,91 +0,0 @@ -package(default_visibility = ["//visibility:public"]) - -licenses(["notice"]) # Apache 2.0 - -load( - "//tensorflow:tensorflow.bzl", - "tf_cc_test", -) - -cc_library( - name = "cluster_utils", - srcs = [ - "cluster_utils.cc", - ], - hdrs = [ - "cluster_utils.h", - ], - deps = [ - "//tensorflow/contrib/lite/toco:toco_port", - ], -) - -cc_library( - name = "cluster", - srcs = [ - "cluster.cc", - ], - hdrs = [ - "cluster.h", - ], - deps = [ - ":cluster_utils", - "//tensorflow/contrib/lite/toco:model", - "//tensorflow/contrib/lite/toco:tooling_util", - "//tensorflow/core:protos_all_cc", - ], -) - -cc_library( - name = "resolve_svdf", - srcs = [ - "resolve_svdf.cc", - ], - hdrs = [ - "resolve_svdf.h", - ], - visibility = ["//visibility:public"], - deps = [ - ":cluster", - ":cluster_utils", - "//tensorflow/contrib/lite/toco:model", - "//tensorflow/contrib/lite/toco:toco_port", - "//tensorflow/contrib/lite/toco:tooling_util", - "//tensorflow/core:lib", - "//tensorflow/core:protos_all_cc", - "@protobuf_archive//:protobuf_headers", - ], -) - -tf_cc_test( - name = "resolve_svdf_test", - srcs = ["resolve_svdf_test.cc"], - tags = ["no_oss"], - deps = [ - ":cluster", - ":cluster_utils", - ":resolve_cluster", - ":resolve_svdf", - "//tensorflow/core:lib", - "//tensorflow/core:protos_all_cc", - "@com_google_googletest//:gtest_main", - ], -) - -cc_library( - name = "resolve_cluster", - srcs = [ - "resolve_cluster.cc", - ], - hdrs = [ - "resolve_cluster.h", - ], - visibility = ["//visibility:public"], - deps = [ - ":cluster", - ":cluster_utils", - ":resolve_svdf", - "//tensorflow/contrib/lite/toco:tooling_util", - "//tensorflow/core:protos_all_cc", - ], -) diff --git a/tensorflow/contrib/lite/toco/tflite/BUILD b/tensorflow/contrib/lite/toco/tflite/BUILD deleted file mode 100644 index fcb628fec8f38e065b7267c209c831da037804aa..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/toco/tflite/BUILD +++ /dev/null @@ -1,152 +0,0 @@ -package( - # To suppress build cleaner error about inclusion of schema_generate.h. - features = ["-layering_check"], -) - -licenses(["notice"]) # Apache 2.0 - -load( - "//tensorflow:tensorflow.bzl", - "tf_cc_test", -) - -cc_library( - name = "operator", - srcs = [ - "operator.cc", - ], - hdrs = [ - "builtin_operator.h", - "custom_operator.h", - "operator.h", - "simple_operator.h", - ], - deps = [ - ":types", - "//tensorflow/contrib/lite/schema:schema_fbs", - "//tensorflow/contrib/lite/toco:graph_transformations", - "//tensorflow/contrib/lite/toco:model", - "//tensorflow/core:framework", - "//tensorflow/core:protos_all_cc", - "//tensorflow/core:ptr_util", - "@com_google_absl//absl/memory", - "@flatbuffers", - ], -) - -tf_cc_test( - name = "operator_test", - srcs = [ - "operator_test.cc", - ], - tags = ["no_oss"], - deps = [ - ":operator", - "//tensorflow/contrib/lite/toco:tooling_util", - "//tensorflow/core:ops", - "//tensorflow/core:protos_all_cc", - "@com_google_googletest//:gtest_main", - "@flatbuffers", - ], -) - -cc_library( - name = "types", - srcs = [ - "types.cc", - ], - hdrs = [ - "types.h", - ], - deps = [ - "//tensorflow/contrib/lite:string_util", - "//tensorflow/contrib/lite/schema:schema_fbs", - "//tensorflow/contrib/lite/toco:model", - ], -) - -tf_cc_test( - name = "types_test", - srcs = [ - "types_test.cc", - ], - tags = ["no_oss"], - deps = [ - ":types", - "//tensorflow/core:ops", - "@com_google_googletest//:gtest_main", - ], -) - -cc_library( - name = "export", - srcs = [ - "export.cc", - ], - hdrs = [ - "export.h", - ], - visibility = ["//visibility:public"], - deps = [ - ":operator", - ":types", - "//tensorflow/contrib/lite:schema_fbs_version", - "//tensorflow/contrib/lite/schema:schema_fbs", - "//tensorflow/contrib/lite/toco:model", - "//tensorflow/contrib/lite/toco:tooling_util", - "//tensorflow/contrib/lite/tools/optimize:quantize_weights", - "@com_google_absl//absl/strings", - "@flatbuffers", - ], -) - -tf_cc_test( - name = "export_test", - srcs = [ - "export_test.cc", - ], - tags = ["no_oss"], - deps = [ - ":export", - "//tensorflow/contrib/lite/schema:schema_fbs", - "//tensorflow/core:ops", - "@com_google_googletest//:gtest_main", - ], -) - -cc_library( - name = "import", - srcs = [ - "import.cc", - ], - hdrs = [ - "import.h", - ], - visibility = ["//visibility:public"], - deps = [ - ":operator", - ":types", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/schema:schema_fbs", - "//tensorflow/contrib/lite/toco:model", - "//tensorflow/contrib/lite/toco:tooling_util", - "//tensorflow/contrib/lite/tools:verifier", - "@flatbuffers", - ], -) - -tf_cc_test( - name = "import_test", - srcs = [ - "import_test.cc", - ], - tags = ["no_oss"], - deps = [ - ":import", - "//tensorflow/contrib/lite:schema_fbs_version", - "//tensorflow/contrib/lite/schema:schema_fbs", - "//tensorflow/core:ops", - "@com_google_googletest//:gtest_main", - "@flatbuffers", - ], -) diff --git a/tensorflow/contrib/lite/toco/tflite/export.cc b/tensorflow/contrib/lite/toco/tflite/export.cc deleted file mode 100644 index b3824d2bc0b2b4f8474f535338819a8ded49b9ec..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/toco/tflite/export.cc +++ /dev/null @@ -1,552 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include "tensorflow/contrib/lite/toco/tflite/export.h" - -#include "flatbuffers/flexbuffers.h" -#include "absl/strings/str_join.h" -#include "tensorflow/contrib/lite/context.h" -#include "tensorflow/contrib/lite/schema/schema_generated.h" -#include "tensorflow/contrib/lite/toco/tflite/operator.h" -#include "tensorflow/contrib/lite/toco/tflite/types.h" -#include "tensorflow/contrib/lite/toco/tooling_util.h" -#include "tensorflow/contrib/lite/tools/optimize/quantize_weights.h" -#include "tensorflow/contrib/lite/version.h" - -namespace toco { - -namespace tflite { - -using flatbuffers::FlatBufferBuilder; -using flatbuffers::Offset; -using flatbuffers::Vector; -using ::tflite::Buffer; -using ::tflite::BuiltinOperator; -using ::tflite::BuiltinOperator_CUSTOM; -using ::tflite::BuiltinOperator_MAX; -using ::tflite::BuiltinOperator_MIN; -using ::tflite::CreateBuffer; -using ::tflite::CreateModel; -using ::tflite::CreateOperator; -using ::tflite::CreateTensor; -using ::tflite::Operator; -using ::tflite::OperatorCode; -using ::tflite::SubGraph; -using ::tflite::Tensor; - -namespace { - -// Check if a TensorFlow Op is a control flow op by its name. -bool IsControlFlowOp(const string& tensorflow_op) { - // Technically this is equalivent to `::tensorflow::Node::IsControlFlow()`. - // It requires to construct a `::tensorflow::Graph` to use that helper - // function, so we simply hardcode the list of control flow ops here. - if (tensorflow_op == "Switch" || tensorflow_op == "RefSwitch" || - tensorflow_op == "Merge" || tensorflow_op == "RefMerge" || - tensorflow_op == "Enter" || tensorflow_op == "RefEnter" || - tensorflow_op == "Exit" || tensorflow_op == "RefExit" || - tensorflow_op == "NextIteration" || tensorflow_op == "RefNextIteration") { - return true; - } - // TODO(ycling): Also check how to handle Variable ops and Assign ops. - return false; -} - -// Check if a TensorFlow Op is unsupportred by the Flex runtime. -bool IsUnsupportedFlexOp(const string& tensorflow_op) { - if (IsControlFlowOp(tensorflow_op)) { - return true; - } - // `HashTableV2` isn't supported for now since it requires an additinonal - // initialization step. - // TODO(b/117651199): Support `HashTableV2` with Flex runtime. - if (tensorflow_op == "HashTableV2") { - return true; - } - return false; -} - -// Map from operator name to TF Lite enum value, for all builtins. -const std::map& GetBuiltinOpsMap() { - static std::map* builtin_ops = nullptr; - if (builtin_ops == nullptr) { - builtin_ops = new std::map(); - - for (int i = BuiltinOperator_MIN; i <= BuiltinOperator_MAX; ++i) { - BuiltinOperator op = static_cast(i); - string name = EnumNameBuiltinOperator(op); - if (op != BuiltinOperator_CUSTOM && !name.empty()) { - (*builtin_ops)[name] = op; - } - } - } - return *builtin_ops; -} - -void WriteModelToString(const flatbuffers::FlatBufferBuilder& builder, - string* file_contents) { - const uint8_t* buffer = builder.GetBufferPointer(); - int size = builder.GetSize(); - *file_contents = string(reinterpret_cast(buffer), size); -} - -} // Anonymous namespace. - -namespace details { - -OperatorKey GetOperatorKey( - const ::toco::Operator& op, - const std::map>& ops_by_type, - bool allow_flex_ops) { - // Get the op name (by Toco definition). - string name = HelpfulOperatorTypeName(op); - - bool is_builtin = false; - OperatorKey key; - - const auto& builtin_ops = GetBuiltinOpsMap(); - if (ops_by_type.count(op.type) != 0) { - key.version = ops_by_type.at(op.type)->GetVersion(op); - name = ops_by_type.at(op.type)->name(); - is_builtin = (builtin_ops.count(name) > 0); - } - - if (is_builtin) { - // For TFLite supported builtin ops, find out its BuiltinOperator enum used - // in FlatBuffer. - key.type = builtin_ops.at(name); - return key; - } - - // The logic below is all for custom ops. - key.is_custom_op = true; - key.type = BuiltinOperator_CUSTOM; - - if (op.type == OperatorType::kUnsupported) { - const TensorFlowUnsupportedOperator& unsupported_op = - static_cast(op); - const auto tensorflow_op = unsupported_op.tensorflow_op; - - // TODO(b/113715895): When `allow_flex_ops` is on, for now there's no way - // to populate a regular custom op. We need to find a way to fix this. - if (ShouldExportAsFlexOp(allow_flex_ops, unsupported_op.tensorflow_op)) { - key.is_custom_op = true; - key.is_flex_op = true; - key.flex_tensorflow_op = tensorflow_op; - key.custom_code = - string(::tflite::kFlexCustomCodePrefix) + key.flex_tensorflow_op; - } else { - key.custom_code = tensorflow_op; - } - } else if (allow_flex_ops && !op.tensorflow_node_def.empty()) { - // For Toco-supported/TFLite-unsupported ops, if the TensorFlow NodeDef - // is retained in the Toco Operator, we produce a Flex op if Flex mode - // is enabled. - key.is_flex_op = true; - key.flex_tensorflow_op = name; - key.custom_code = - string(::tflite::kFlexCustomCodePrefix) + key.flex_tensorflow_op; - } else { - // If Flex is disabled or the original TensorFlow NodeDef isn't available, - // we produce a custom op. This gives developers a chance to implemenr - // custom ops. - key.custom_code = name; - } - - if (key.is_flex_op) { - if (IsUnsupportedFlexOp(key.flex_tensorflow_op)) { - key.is_unsupported_flex_op = true; - } - } - return key; -} - -void LoadTensorsMap(const Model& model, TensorsMap* tensors_map) { - // First find a list of unique array names. - std::set names; - for (const auto& array_pair : model.GetArrayMap()) { - names.insert(array_pair.first); - } - - // Now assign indices to them and fill in the map. - int index = 0; - for (const auto& name : names) { - (*tensors_map)[name] = index; - ++index; - } -} - -void LoadOperatorsMap( - const Model& model, OperatorsMap* operators_map, - const std::map>& ops_by_type, - bool allow_flex_ops) { - // First find a list of unique operator types. - std::set keys; - for (const auto& op : model.operators) { - keys.insert(GetOperatorKey(*op, ops_by_type, allow_flex_ops)); - } - // Now assign indices to them and fill in the map. - int index = 0; - for (const auto& key : keys) { - (*operators_map)[key] = index; - ++index; - } -} - -} // namespace details - -Offset>> ExportTensors( - const Model& model, const details::TensorsMap& tensors_map, - FlatBufferBuilder* builder, std::vector* buffers_to_write, - const std::set& variable_tensor_indices) { - // In the end we will need to produce a vector sorted by the indices of the - // tensors in the tensors_map. - std::map> ordered_tensors; - - for (const auto& array_pair : model.GetArrayMap()) { - const string& tensor_name = array_pair.first; - const toco::Array& array = *array_pair.second; - - int buffer_index = buffers_to_write->size(); - auto type = DataType::Serialize(array.data_type); - buffers_to_write->push_back(&array); - - std::vector shape; - if (array.has_shape()) { - for (int d : array.shape().dims()) { - shape.push_back(d); - } - } - - Offset> min; - Offset> max; - Offset> scale; - Offset> zero_point; - if (array.minmax) { - min = builder->CreateVector( - std::vector{static_cast(array.minmax->min)}); - max = builder->CreateVector( - std::vector{static_cast(array.minmax->max)}); - } - if (array.quantization_params) { - scale = builder->CreateVector(std::vector{ - static_cast(array.quantization_params->scale)}); - zero_point = builder->CreateVector( - std::vector{array.quantization_params->zero_point}); - } - auto q_param = ::tflite::CreateQuantizationParameters(*builder, min, max, - scale, zero_point); - - int index = tensors_map.at(tensor_name); - bool is_variable = - variable_tensor_indices.find(index) != variable_tensor_indices.end(); - ordered_tensors[index] = - CreateTensor(*builder, builder->CreateVector(shape), type, buffer_index, - builder->CreateString(tensor_name), q_param, is_variable); - } - - std::vector> tensor_vector; - tensor_vector.reserve(ordered_tensors.size()); - for (const auto& tensor : ordered_tensors) { - tensor_vector.push_back(tensor.second); - } - - return builder->CreateVector(tensor_vector); -} - -Offset> ExportInputTensors( - const Model& model, const details::TensorsMap& tensors_map, - FlatBufferBuilder* builder) { - std::vector inputs; - for (const auto& input : model.flags.input_arrays()) { - inputs.push_back(tensors_map.at(input.name())); - } - return builder->CreateVector(inputs); -} - -Offset> ExportOutputTensors( - const Model& model, const details::TensorsMap& tensors_map, - FlatBufferBuilder* builder) { - std::vector outputs; - for (const string& output : model.flags.output_arrays()) { - outputs.push_back(tensors_map.at(output)); - } - return builder->CreateVector(outputs); -} - -Offset>> ExportOperatorCodes( - const Model& model, - const std::map>& ops_by_type, - const details::OperatorsMap& operators_map, FlatBufferBuilder* builder, - const ExportParams& params) { - // Map from operator name to TF Lite enum value, for all builtins. - std::map builtin_ops; - for (int i = BuiltinOperator_MIN; i <= BuiltinOperator_MAX; ++i) { - BuiltinOperator op = static_cast(i); - string name = EnumNameBuiltinOperator(op); - if (op != BuiltinOperator_CUSTOM && !name.empty()) { - builtin_ops[name] = op; - } - } - - // We will need to produce a vector of codes in the same order as they - // appear in the operators_map. - std::map> ordered_opcodes; - - for (const auto& op : model.operators) { - const details::OperatorKey operator_key = - details::GetOperatorKey(*op, ops_by_type, params.allow_flex_ops); - int op_index = operators_map.at(operator_key); - - flatbuffers::Offset custom_code = 0; - if (!operator_key.custom_code.empty()) { - custom_code = builder->CreateString(operator_key.custom_code); - } - - ordered_opcodes[op_index] = CreateOperatorCode( - *builder, operator_key.type, custom_code, operator_key.version); - } - - std::vector> opcode_vector; - opcode_vector.reserve(ordered_opcodes.size()); - for (const auto& opcode : ordered_opcodes) { - opcode_vector.push_back(opcode.second); - } - - return builder->CreateVector(opcode_vector); -} - -Offset>> ExportOperators( - const Model& model, - const std::map>& ops_by_type, - const details::OperatorsMap& operators_map, - const details::TensorsMap& tensors_map, FlatBufferBuilder* builder, - std::set* variable_tensor_indices, const ExportParams& params) { - variable_tensor_indices->clear(); - - // The operators are in execution order, so we just follow tf.mini order. - std::vector> op_vector; - for (const auto& op : model.operators) { - std::vector inputs; - for (const string& input : op->inputs) { - // -1 is the ID for optional tensor in TFLite output - int id = model.IsOptionalArray(input) ? -1 : tensors_map.at(input); - inputs.push_back(id); - } - std::vector outputs; - for (const string& output : op->outputs) { - outputs.push_back(tensors_map.at(output)); - } - - const auto key = - details::GetOperatorKey(*op, ops_by_type, params.allow_flex_ops); - int op_index = operators_map.at(key); - - auto tflite_op_it = ops_by_type.find(op->type); - BaseOperator* tflite_op = tflite_op_it == ops_by_type.end() - ? nullptr - : tflite_op_it->second.get(); - - // This is a custom op unless we can find it in ops_by_type, and even then - // it could be a custom op (such as kUnsupported). - auto options = Options::Custom(0); - - std::vector mutating_input_variables; - if (tflite_op) { - options = tflite_op->Serialize(*op, builder); - mutating_input_variables = tflite_op->GetMutatingInputVariables(*op); - - if (!mutating_input_variables.empty()) { - for (int i = 0; i < op->inputs.size(); ++i) { - if (!mutating_input_variables[i]) { - continue; - } - int32_t variable_tensor_index = tensors_map.at(op->inputs[i]); - variable_tensor_indices->insert(variable_tensor_index); - } - } - } else if (key.is_flex_op && !op->tensorflow_node_def.empty()) { - auto fbb = WriteFlexOpOptions(op->tensorflow_node_def); - if (fbb) { - options = Options::Custom(builder->CreateVector(fbb->GetBuffer())); - } - } - // The only supported CustomOptionFormat is FLEXBUFFERS now. - op_vector.push_back(CreateOperator( - *builder, op_index, builder->CreateVector(inputs), - builder->CreateVector(outputs), options.type, options.builtin, - options.custom, ::tflite::CustomOptionsFormat_FLEXBUFFERS, - builder->CreateVector(mutating_input_variables))); - } - - return builder->CreateVector(op_vector); -} - -Offset>> ExportBuffers( - const Model& model, const std::vector& buffers_to_write, - FlatBufferBuilder* builder) { - std::vector> buffer_vector; - size_t index = 0; - for (const Array* array_ptr : buffers_to_write) { - const Array& array = *array_ptr; - Offset> data_buffer = DataBuffer::Serialize(array, builder); - buffer_vector.push_back(CreateBuffer(*builder, data_buffer)); - index++; - } - return builder->CreateVector(buffer_vector); -} - -tensorflow::Status Export(const Model& model, string* output_file_contents, - const ExportParams& params) { - const auto ops_by_type = BuildOperatorByTypeMap(params.allow_flex_ops); - return Export(model, output_file_contents, params, ops_by_type); -} - -tensorflow::Status Export( - const Model& model, string* output_file_contents, - const ExportParams& params, - const std::map>& ops_by_type) { - flatbuffers::FlatBufferBuilder builder(/*initial_size=*/10240); - - details::TensorsMap tensors_map; - details::LoadTensorsMap(model, &tensors_map); - - details::OperatorsMap operators_map; - details::LoadOperatorsMap(model, &operators_map, ops_by_type, - params.allow_flex_ops); - - std::vector buffers_to_write; - Array empty_array; - buffers_to_write.push_back(&empty_array); - - auto op_codes = - ExportOperatorCodes(model, ops_by_type, operators_map, &builder, params); - - for (const auto& op : model.operators) { - if (op->type == OperatorType::kFakeQuant) { - LOG(WARNING) << "FAKE_QUANT operation " << LogName(*op) - << " was not converted. If running quantized make sure you " - "are passing --inference_type=QUANTIZED_UINT8 and values " - "for --std_values and --mean_values."; - } - } - - std::set custom_ops; - std::set unsupported_flex_ops; - for (const auto& it : operators_map) { - const details::OperatorKey& key = it.first; - if (key.is_custom_op) { - custom_ops.insert(key.custom_code); - } - if (key.is_unsupported_flex_op) { - unsupported_flex_ops.insert(key.flex_tensorflow_op); - } - } - - if (!custom_ops.empty()) { - if (!params.allow_custom_ops) { - // Remove ExpandDims and ReorderAxes from unimplemented list unless they - // compose the list. Both ops are removed during graph transformations. - // However, if an op is unimplemented earlier in the model, the graph - // transformation is unable to run because the output shape is not - // defined. This causes unnecessary confusion during model conversion - // time. - std::set custom_ops_final; - for (const auto& op_type : custom_ops) { - if (op_type != "ReorderAxes" && op_type != "ExpandDims") { - custom_ops_final.insert(op_type); - } - } - if (custom_ops_final.empty()) { - custom_ops_final = custom_ops; - } - - return tensorflow::errors::InvalidArgument(absl::StrCat( - "Some of the operators in the model are not supported by the " - "standard TensorFlow Lite runtime. If you have a custom " - "implementation for them you can disable this error with " - "--allow_custom_ops, or by setting allow_custom_ops=True when " - "calling tf.contrib.lite.TFLiteConverter(). Here is a list of " - "operators for which you will need custom implementations: ", - absl::StrJoin(custom_ops_final, ", "), ".")); - } - - std::set unsupported_control_flow_ops; - // Check if unsupported ops contains control flow ops. It's impossible - // to implement these ops as custom ops at the moment. - for (const auto& op : custom_ops) { - if (IsControlFlowOp(op)) { - unsupported_control_flow_ops.insert(op); - } - } - if (!unsupported_control_flow_ops.empty()) { - return tensorflow::errors::InvalidArgument(absl::StrCat( - "TensorFlow Lite currently doesn't support control flow ops: ", - absl::StrJoin(unsupported_control_flow_ops, ", "), ".")); - } - } - - if (!unsupported_flex_ops.empty()) { - return tensorflow::errors::InvalidArgument( - absl::StrCat("Some of the operators in the model are not supported by " - "TensorFlow Flex runtime: ", - absl::StrJoin(unsupported_flex_ops, ", "), ".")); - } - - std::set variable_tensor_indices; - auto ops = ExportOperators(model, ops_by_type, operators_map, tensors_map, - &builder, &variable_tensor_indices, params); - - auto tensors = ExportTensors(model, tensors_map, &builder, &buffers_to_write, - variable_tensor_indices); - auto inputs = ExportInputTensors(model, tensors_map, &builder); - auto outputs = ExportOutputTensors(model, tensors_map, &builder); - - // TODO(aselle): add support to toco for multiple subgraphs. - auto subgraph = CreateSubGraph(builder, tensors, inputs, outputs, ops, - /* name */ 0); - std::vector> subgraphs = {subgraph}; - - auto buffers = ExportBuffers(model, buffers_to_write, &builder); - auto description = builder.CreateString("TOCO Converted."); - auto new_model_location = - CreateModel(builder, TFLITE_SCHEMA_VERSION, op_codes, - builder.CreateVector(subgraphs), description, buffers); - ::tflite::FinishModelBuffer(builder, new_model_location); - - if (params.quantize_weights) { - // Call the quantize_weights tool. - LOG(INFO) << "Quantizing TFLite model after conversion to flatbuffer. " - "dump_graphviz will only output the model before this " - "transformation. To visualize the output graph use " - "lite/tools/optimize.py."; - flatbuffers::FlatBufferBuilder q_builder(/*initial_size=*/10240); - const uint8_t* buffer = builder.GetBufferPointer(); - const ::tflite::Model* input_model = ::tflite::GetModel(buffer); - if (::tflite::optimize::QuantizeWeights(&q_builder, input_model) != - kTfLiteOk) { - return tensorflow::errors::InvalidArgument( - "Quantize weights transformation failed."); - } - WriteModelToString(q_builder, output_file_contents); - } else { - WriteModelToString(builder, output_file_contents); - } - - return tensorflow::Status(); -} - -} // namespace tflite - -} // namespace toco diff --git a/tensorflow/contrib/lite/toco/tflite/export.h b/tensorflow/contrib/lite/toco/tflite/export.h deleted file mode 100644 index 06b4cb1caa344c2c5072cd7282433fc96c910dc6..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/toco/tflite/export.h +++ /dev/null @@ -1,151 +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_CONTRIB_LITE_TOCO_TFLITE_EXPORT_H_ -#define TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_EXPORT_H_ - -#include "tensorflow/contrib/lite/toco/model.h" -#include "tensorflow/contrib/lite/toco/tflite/operator.h" -#include "tensorflow/contrib/lite/util.h" - -namespace toco { - -namespace tflite { - -// The parameters for exporting a TFLite model. -struct ExportParams { - bool allow_custom_ops = false; - bool allow_flex_ops = false; - bool quantize_weights = false; -}; - -// Transform the given tf.mini model into a TF Lite flatbuffer and deposit the -// result in the given string. -tensorflow::Status Export(const Model& model, string* output_file_contents, - const ExportParams& params); - -// Export API with custom TFLite operator mapping. -tensorflow::Status Export( - const Model& model, string* output_file_contents, - const ExportParams& params, - const std::map>& ops_by_type); - -// This is for backward-compatibility. -// TODO(ycling): Remove the deprecated entry functions. -inline void Export(const Model& model, bool allow_custom_ops, - bool quantize_weights, string* output_file_contents) { - ExportParams params; - params.allow_custom_ops = allow_custom_ops; - params.quantize_weights = quantize_weights; - auto status = Export(model, output_file_contents, params); - if (!status.ok()) LOG(QFATAL) << status.error_message(); -} - -// This is for backward-compatibility. -// TODO(ycling): Remove the deprecated entry functions. -inline void Export( - const Model& model, bool allow_custom_ops, bool quantize_weights, - string* output_file_contents, - const std::map>& ops_by_type) { - ExportParams params; - params.allow_custom_ops = allow_custom_ops; - params.quantize_weights = quantize_weights; - auto status = Export(model, output_file_contents, params, ops_by_type); - if (!status.ok()) LOG(QFATAL) << status.error_message(); -} - -// This is for backward-compatibility. -// TODO(ycling): Remove the deprecated entry functions. -inline void Export(const Model& model, string* output_file_contents) { - ExportParams params; - params.allow_custom_ops = true; - auto status = Export(model, output_file_contents, params); - if (!status.ok()) LOG(QFATAL) << status.error_message(); -} - -namespace details { - -// A maps from tensor name to its final position in the TF Lite buffer. -using TensorsMap = std::unordered_map; - -// A key to identify an operator. -// Only when `type` is `kUnsupported`, `custom_code` is filled to -// identify which operation is used. -struct OperatorKey { - OperatorKey() {} - OperatorKey(::tflite::BuiltinOperator type, const std::string& custom_code, - int version) - : type(type), custom_code(custom_code), version(version) {} - - // Only `type`, `custom_code` and `version` is used to compute hash and - // identity. - ::tflite::BuiltinOperator type = ::tflite::BuiltinOperator_CUSTOM; - std::string custom_code; - int version = 1; - - // The fields below are not used to compute hash and identity. - // TODO(ycling): Consider to change these fields to accessor functions. - bool is_custom_op = false; - bool is_flex_op = false; - bool is_unsupported_flex_op = false; - // The original TensorFlow op name for the flex op. Filled only when - // `is_flex_op` is true. - std::string flex_tensorflow_op; - - bool operator<(const OperatorKey& other) const { - if (type < other.type) return true; - else if (type > other.type) - return false; - else if (custom_code < other.custom_code) - return true; - else if (custom_code > other.custom_code) - return false; - else - return version < other.version; - } - - bool operator==(const OperatorKey& other) const { - return type == other.type && custom_code == other.custom_code && - version == other.version; - } - - struct Hash { - size_t operator()(const OperatorKey& key) const { - return ::tflite::CombineHashes( - {std::hash()(static_cast(key.type)), - std::hash()(key.custom_code), - std::hash()(key.version)}); - } - }; -}; - -OperatorKey GetOperatorKey( - const ::toco::Operator& op, - const std::map>& ops_by_type, - bool allow_flex_ops); - -// A maps from operator type to its final position in the TF Lite buffer. -using OperatorsMap = std::unordered_map; - -void LoadTensorsMap(const Model& model, TensorsMap* tensors_map); -void LoadOperatorsMap( - const Model& model, OperatorsMap* operators_map, - const std::map>& ops_by_type, - bool allow_flex_ops); - -} // namespace details -} // namespace tflite -} // namespace toco - -#endif // TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_EXPORT_H_ diff --git a/tensorflow/contrib/lite/toco/tflite/export_test.cc b/tensorflow/contrib/lite/toco/tflite/export_test.cc deleted file mode 100644 index 6a293901485d8fb4e1dba82094554966d54614d7..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/toco/tflite/export_test.cc +++ /dev/null @@ -1,556 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include "tensorflow/contrib/lite/toco/tflite/export.h" - -#include -#include -#include "tensorflow/contrib/lite/schema/schema_generated.h" -#include "tensorflow/contrib/lite/toco/tflite/builtin_operator.h" -#include "tensorflow/contrib/lite/toco/tflite/operator.h" -#include "tensorflow/contrib/lite/toco/tflite/types.h" -#include "tensorflow/core/framework/node_def.pb.h" - -namespace toco { -namespace tflite { -namespace { - -using ::testing::ElementsAre; - -class ExportTest : public ::testing::Test { - protected: - void ResetOperators() { input_model_.operators.clear(); } - void AddTensorsByName(std::initializer_list names) { - for (const string& name : names) { - input_model_.GetOrCreateArray(name); - } - } - void AddOperatorsByName(std::initializer_list names) { - for (const string& name : names) { - if (name == "Conv") { - auto* op = new ConvOperator; - op->padding.type = PaddingType::kSame; - input_model_.operators.emplace_back(op); - } else if (name == "Add") { - input_model_.operators.emplace_back(new AddOperator); - } else if (name == "Sub") { - input_model_.operators.emplace_back(new SubOperator); - } else { - auto* op = new TensorFlowUnsupportedOperator; - op->tensorflow_op = name; - input_model_.operators.emplace_back(op); - } - } - } - - void BuildQuantizableTestModel() { - input_model_.GetOrCreateArray("inputs"); - Array& weight_array = input_model_.GetOrCreateArray("weights"); - - // Make the buffer large enough for QuantizeWeights transformation to take - // effect. - int buf_size = 1296; - auto weight_buf = absl::make_unique(buf_size); - for (int i = 0; i < buf_size; i++) { - // Fill the array with some garbage values. - weight_buf[i] = static_cast(i % 128); - } - - weight_array.data_type = ArrayDataType::kFloat; - - // Initialize shape for the input array. - Shape* weight_array_shape = weight_array.mutable_shape(); - std::vector* weight_array_shape_dim = - weight_array_shape->mutable_dims(); - weight_array_shape_dim->resize(4, 6); - auto& weight_array_buffer = - weight_array.GetMutableBuffer(); - weight_array_buffer.data.resize(buf_size); - float* buf_ptr = - weight_array.GetMutableBuffer().data.data(); - std::copy(weight_buf.get(), weight_buf.get() + buf_size, buf_ptr); - - { - auto* op = new ConvOperator; - op->padding.type = PaddingType::kSame; - op->inputs = {"inputs", "weights"}; - input_model_.operators.emplace_back(op); - } - input_model_.operators.emplace_back(new AddOperator); - } - - std::vector ExportAndSummarizeOperators(const ExportParams& params) { - std::vector names; - - string result; - auto status = Export(input_model_, &result, params); - if (!status.ok()) { - LOG(INFO) << status.error_message(); - return names; - } - - auto* model = ::tflite::GetModel(result.data()); - - for (const ::tflite::OperatorCode* opcode : *model->operator_codes()) { - if (opcode->builtin_code() != ::tflite::BuiltinOperator_CUSTOM) { - names.push_back(string("builtin:") + ::tflite::EnumNameBuiltinOperator( - opcode->builtin_code())); - } else { - names.push_back(string("custom:") + opcode->custom_code()->c_str()); - } - } - - return names; - } - - std::vector ExportAndGetOperatorIndices( - const ExportParams& params) { - std::vector indices; - - string result; - if (!Export(input_model_, &result, params).ok()) return indices; - auto* model = ::tflite::GetModel(result.data()); - - auto operators = (*model->subgraphs())[0]->operators(); - for (const auto* op : *operators) { - indices.push_back(op->opcode_index()); - } - return indices; - } - - Model input_model_; -}; - -TEST_F(ExportTest, LoadTensorsMap) { - AddTensorsByName({"tensor_one", "tensor_two"}); - - details::TensorsMap tensors; - details::LoadTensorsMap(input_model_, &tensors); - EXPECT_EQ(0, tensors["tensor_one"]); - EXPECT_EQ(1, tensors["tensor_two"]); -} - -TEST_F(ExportTest, LoadOperatorsMap) { - AddOperatorsByName({"Conv", "Add", "MyCrazyOp", "Sub"}); - - details::OperatorsMap operators; - const auto ops_by_type = BuildOperatorByTypeMap(); - details::LoadOperatorsMap(input_model_, &operators, ops_by_type, false); - EXPECT_EQ( - 0, operators[details::OperatorKey(::tflite::BuiltinOperator_ADD, "", 1)]); - EXPECT_EQ(1, operators[details::OperatorKey(::tflite::BuiltinOperator_CONV_2D, - "", 1)]); - EXPECT_EQ(2, operators[details::OperatorKey(::tflite::BuiltinOperator_CUSTOM, - "MyCrazyOp", 1)]); - EXPECT_EQ( - 3, operators[details::OperatorKey(::tflite::BuiltinOperator_SUB, "", 1)]); -} - -TEST_F(ExportTest, Export) { - AddOperatorsByName({"Conv", "Add", "MyCrazyOp", "Sub"}); - - ExportParams params; - params.allow_custom_ops = true; - params.allow_flex_ops = false; - params.quantize_weights = false; - - EXPECT_THAT(ExportAndSummarizeOperators(params), - ElementsAre("builtin:ADD", "builtin:CONV_2D", "custom:MyCrazyOp", - "builtin:SUB")); - EXPECT_THAT(ExportAndGetOperatorIndices(params), ElementsAre(1, 0, 2, 3)); -} - -TEST_F(ExportTest, QuantizeWeights) { - // Sanity check for quantize_weights parameter. - BuildQuantizableTestModel(); - string unquantized_result; - Export(input_model_, true, /*quantize_weights*/ false, &unquantized_result); - - BuildQuantizableTestModel(); - string quantized_result; - Export(input_model_, true, /*quantize_weights*/ true, &quantized_result); - - // The quantized models should be smaller. - EXPECT_LT(quantized_result.size(), unquantized_result.size()); -} - -class OpSetsTest : public ExportTest { - public: - enum OpSet { kTfLiteBuiltins, kSelectTfOps, kCustomOps }; - - void SetAllowedOpSets(std::initializer_list sets) { - import_all_ops_as_unsupported_ = true; - params_.allow_custom_ops = false; - params_.allow_flex_ops = false; - params_.quantize_weights = false; - - for (OpSet i : sets) { - switch (i) { - case kTfLiteBuiltins: - import_all_ops_as_unsupported_ = false; - break; - case kSelectTfOps: - params_.allow_flex_ops = true; - break; - case kCustomOps: - params_.allow_custom_ops = true; - break; - } - } - } - - std::vector ImportExport(std::initializer_list op_names) { - ResetOperators(); - if (!import_all_ops_as_unsupported_) { - AddOperatorsByName(op_names); - } else { - for (const string& name : op_names) { - auto* op = new TensorFlowUnsupportedOperator; - op->tensorflow_op = name; - input_model_.operators.emplace_back(op); - } - } - return ExportAndSummarizeOperators(params_); - } - - private: - bool import_all_ops_as_unsupported_; - ExportParams params_; -}; - -TEST_F(OpSetsTest, BuiltinsOnly) { - // --target_op_set=TFLITE_BUILTINS - SetAllowedOpSets({kTfLiteBuiltins}); - EXPECT_THAT(ImportExport({"Add", "AdjustHue", "UnrollAndFold"}), - ElementsAre()); - EXPECT_THAT(ImportExport({"Add"}), ElementsAre("builtin:ADD")); - - // --target_op_set=TFLITE_BUILTINS --allow_custom_ops - SetAllowedOpSets({kTfLiteBuiltins, kCustomOps}); - EXPECT_THAT( - ImportExport({"Add", "AdjustHue", "UnrollAndFold"}), - ElementsAre("builtin:ADD", "custom:AdjustHue", "custom:UnrollAndFold")); -} - -TEST_F(OpSetsTest, TfSelectOnly) { - // --target_op_set=SELECT_TF_OPS - SetAllowedOpSets({kSelectTfOps}); - EXPECT_THAT(ImportExport({"Add", "AdjustHue", "UnrollAndFold"}), - ElementsAre()); - - // TODO(b/117845348): Add should be recognized as a flex op, which should be - // OK even if we don't specify --allow_custom_ops - // EXPECT_THAT(ImportExport({"Add"}), ElementsAre("custom:FlexAdd")); - EXPECT_THAT(ImportExport({"Add"}), ElementsAre()); - - // --target_op_set=SELECT_TF_OPS --allow_custom_ops - SetAllowedOpSets({kSelectTfOps, kCustomOps}); - EXPECT_THAT(ImportExport({"Add", "AdjustHue", "UnrollAndFold"}), - ElementsAre("custom:FlexAdd", "custom:FlexAdjustHue", - "custom:UnrollAndFold")); -} - -TEST_F(OpSetsTest, BuiltinsAndTfSelect) { - // --target_op_set=TFLITE_BUILTINS,SELECT_TF_OPS - SetAllowedOpSets({kTfLiteBuiltins, kSelectTfOps}); - EXPECT_THAT(ImportExport({"Add", "AdjustHue", "UnrollAndFold"}), - ElementsAre()); - - // TODO(b/117845348): AdjustHue should be recognized as a flex op, - // which should be OK even if we don't specify --allow_custom_ops - // EXPECT_THAT(ImportExport({"Add", "AdjustHue"}), - // ElementsAre("builtin:ADD", "custom:FlexAdjustHue")); - EXPECT_THAT(ImportExport({"Add", "AdjustHue"}), ElementsAre()); - - // --target_op_set=TFLITE_BUILTINS,SELECT_TF_OPS --allow_custom_ops - SetAllowedOpSets({kTfLiteBuiltins, kSelectTfOps, kCustomOps}); - EXPECT_THAT(ImportExport({"Add", "AdjustHue", "UnrollAndFold"}), - ElementsAre("builtin:ADD", "custom:FlexAdjustHue", - "custom:UnrollAndFold")); -} - -// This test is based on a hypothetical scenario that dilation is supported -// only in Conv version 2. So Toco populates version=1 when dialation -// parameters are all 1, and version=2 otehrwise. -class FakeConvolutionOperator - : public BuiltinOperator { - public: - FakeConvolutionOperator() - : BuiltinOperator(::tflite::BuiltinOperator_CONV_2D, - OperatorType::kConv) {} - - // Returning the op version according to the op parameters. - int GetVersion(const Operator& op) const override { - const TocoOperator& conv_op = static_cast(op); - if (conv_op.dilation_width_factor != 1 || - conv_op.dilation_height_factor != 1) { - // Version 2 if dilation is used. - return 2; - } - return 1; - } - - // Note: The read / write code doesn't need to be changed if we stick with - // the restrictions: - // * Only adding parameters at the bottom of the Flatbuffer tables. - // * When the default value of parameters are used, the op works consistently - // with the previous version. - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - auto padding = Padding::Serialize(op.padding.type); - auto activation_function = - ActivationFunction::Serialize(op.fused_activation_function); - return ::tflite::CreateConv2DOptions(*builder, padding, op.stride_width, - op.stride_height, activation_function, - op.dilation_width_factor, - op.dilation_height_factor); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->padding.type = Padding::Deserialize(options.padding()); - op->stride_width = options.stride_w(); - op->stride_height = options.stride_h(); - op->dilation_width_factor = options.dilation_w_factor(); - op->dilation_height_factor = options.dilation_h_factor(); - op->fused_activation_function = - ActivationFunction::Deserialize(options.fused_activation_function()); - } -}; - -class VersionedOpExportTest : public ::testing::Test { - protected: - void SetUp() override { - input_model_.GetOrCreateArray("input"); - input_model_.GetOrCreateArray("filter"); - input_model_.GetOrCreateArray("output"); - } - void AddConvOp(bool use_dialation) { - { - auto* op = new ConvOperator; - op->inputs.push_back("input"); - op->inputs.push_back("filter"); - op->inputs.push_back("output"); - - op->padding.type = PaddingType::kSame; - op->stride_width = 1; - op->stride_height = 1; - if (use_dialation) { - op->dilation_width_factor = 2; - op->dilation_height_factor = 2; - } else { - op->dilation_width_factor = 1; - op->dilation_height_factor = 1; - } - input_model_.operators.emplace_back(op); - } - } - - std::map> - BuildFakeOperatorByTypeMap() { - std::map> result; - result[OperatorType::kConv] = - std::unique_ptr(new FakeConvolutionOperator); - return result; - } - - Model input_model_; -}; - -TEST_F(VersionedOpExportTest, LoadOperatorsMapWithOpV1) { - AddConvOp(false); - - details::OperatorsMap operators; - const auto ops_by_type = BuildFakeOperatorByTypeMap(); - details::LoadOperatorsMap(input_model_, &operators, ops_by_type, false); - - EXPECT_EQ(1, operators.size()); - EXPECT_EQ(0, operators.at(details::OperatorKey( - ::tflite::BuiltinOperator_CONV_2D, "", 1))); -} - -TEST_F(VersionedOpExportTest, LoadOperatorsMapWithOpV2) { - AddConvOp(true); - - details::OperatorsMap operators; - const auto ops_by_type = BuildFakeOperatorByTypeMap(); - details::LoadOperatorsMap(input_model_, &operators, ops_by_type, false); - - EXPECT_EQ(1, operators.size()); - EXPECT_EQ(0, operators.at(details::OperatorKey( - ::tflite::BuiltinOperator_CONV_2D, "", 2))); -} - -TEST_F(VersionedOpExportTest, LoadOperatorsMapWithBothVersions) { - AddConvOp(false); - AddConvOp(true); - - details::OperatorsMap operators; - const auto ops_by_type = BuildFakeOperatorByTypeMap(); - details::LoadOperatorsMap(input_model_, &operators, ops_by_type, false); - - EXPECT_EQ(2, operators.size()); - EXPECT_EQ(0, operators.at(details::OperatorKey( - ::tflite::BuiltinOperator_CONV_2D, "", 1))); - EXPECT_EQ(1, operators.at(details::OperatorKey( - ::tflite::BuiltinOperator_CONV_2D, "", 2))); -} - -TEST_F(VersionedOpExportTest, Export) { - AddConvOp(false); - AddConvOp(true); - - string result; - const auto ops_by_type = BuildFakeOperatorByTypeMap(); - Export(input_model_, true, false, &result, ops_by_type); - - auto* model = ::tflite::GetModel(result.data()); - auto operator_codes = model->operator_codes(); - - // Verify that 2 operator codes are populdated. Both are CONV_2D but with - // different versions. - EXPECT_EQ(2, operator_codes->size()); - EXPECT_EQ(::tflite::BuiltinOperator_CONV_2D, - (*operator_codes)[0]->builtin_code()); - EXPECT_EQ(1, (*operator_codes)[0]->version()); - EXPECT_EQ(::tflite::BuiltinOperator_CONV_2D, - (*operator_codes)[1]->builtin_code()); - EXPECT_EQ(2, (*operator_codes)[1]->version()); - - // Verify that the 2 operators points to the correct indices of the operation - // codes. - auto operators = (*model->subgraphs())[0]->operators(); - EXPECT_EQ(2, operators->size()); - EXPECT_EQ(0, (*operators)[0]->opcode_index()); - EXPECT_EQ(1, (*operators)[1]->opcode_index()); -} - -TEST(OperatorKeyTest, TestBuiltinOp) { - auto op = absl::make_unique(); - - const auto ops_by_type = BuildOperatorByTypeMap(); - const auto key = details::GetOperatorKey(*op, ops_by_type, false); - - EXPECT_EQ(key.type, ::tflite::BuiltinOperator_CONV_2D); - EXPECT_EQ(key.custom_code, ""); - EXPECT_EQ(key.version, 1); -} - -TEST(OperatorKeyTest, TestCustomOp) { - auto op = absl::make_unique(); - op->tensorflow_op = "MyCrazyCustomOp"; - - const auto ops_by_type = BuildOperatorByTypeMap(); - const auto key = details::GetOperatorKey(*op, ops_by_type, false); - - EXPECT_EQ(key.type, ::tflite::BuiltinOperator_CUSTOM); - EXPECT_EQ(key.custom_code, "MyCrazyCustomOp"); - EXPECT_EQ(key.version, 1); -} - -TEST(OperatorKeyTest, TestFlexOp) { - auto op = absl::make_unique(); - op->tensorflow_op = "BatchMatMul"; - - const auto ops_by_type = BuildOperatorByTypeMap(); - { - const auto key = details::GetOperatorKey(*op, ops_by_type, false); - // It shouldn't be converted to Flex op if `allow_flex_op` is false. - EXPECT_EQ(key.type, ::tflite::BuiltinOperator_CUSTOM); - EXPECT_EQ(key.custom_code, "BatchMatMul"); - EXPECT_EQ(key.version, 1); - EXPECT_FALSE(key.is_flex_op); - } - - { - // Verify that the custom op name is prefixed by "Flex" and `is_flex_op` - // is true. - const auto key = details::GetOperatorKey(*op, ops_by_type, true); - EXPECT_EQ(key.type, ::tflite::BuiltinOperator_CUSTOM); - EXPECT_EQ(key.custom_code, "FlexBatchMatMul"); - EXPECT_EQ(key.version, 1); - EXPECT_TRUE(key.is_flex_op); - } -} - -TEST(OperatorKeyTest, TestFlexWithControlFlowOp) { - auto op = absl::make_unique(); - op->tensorflow_op = "Merge"; - - const auto ops_by_type = BuildOperatorByTypeMap(); - const auto key = details::GetOperatorKey(*op, ops_by_type, true); - - EXPECT_EQ(key.type, ::tflite::BuiltinOperator_CUSTOM); - EXPECT_EQ(key.custom_code, "FlexMerge"); - EXPECT_EQ(key.version, 1); - EXPECT_TRUE(key.is_flex_op); - // The control flow ops should be marked as unsupported. - EXPECT_TRUE(key.is_unsupported_flex_op); -} - -TEST(OperatorKeyTest, TestFlexWithUnsupportedOp) { - auto op = absl::make_unique(); - op->tensorflow_op = "HashTableV2"; - - const auto ops_by_type = BuildOperatorByTypeMap(); - const auto key = details::GetOperatorKey(*op, ops_by_type, true); - - EXPECT_EQ(key.type, ::tflite::BuiltinOperator_CUSTOM); - EXPECT_EQ(key.custom_code, "FlexHashTableV2"); - EXPECT_EQ(key.version, 1); - EXPECT_TRUE(key.is_flex_op); - // The control flow ops should be marked as unsupported. - EXPECT_TRUE(key.is_unsupported_flex_op); -} - -TEST(OperatorKeyTest, TestFlexWithPartiallySupportedOps) { - // Test Toco-supported/TFLite-unsupported operators. - // TODO(ycling): The test will be broken if Range is implemented in TFLite. - // Find a more robust way to test the fallback logic. - auto op = absl::make_unique(); - - const auto ops_by_type = BuildOperatorByTypeMap(); - - { - // If NodeDef isn't retained in the Toco op, a regular custom op - // will be exported. - const auto key = details::GetOperatorKey(*op, ops_by_type, true); - EXPECT_EQ(key.type, ::tflite::BuiltinOperator_CUSTOM); - EXPECT_EQ(key.custom_code, "Range"); - EXPECT_EQ(key.version, 1); - EXPECT_FALSE(key.is_flex_op); - } - - ::tensorflow::NodeDef node_def; - node_def.set_name("Range"); - node_def.set_op("Range"); - node_def.SerializeToString(&op->tensorflow_node_def); - - { - // If NodeDef is retained in the Toco op, a Flex op will be exported. - const auto key = details::GetOperatorKey(*op, ops_by_type, true); - EXPECT_EQ(key.type, ::tflite::BuiltinOperator_CUSTOM); - EXPECT_EQ(key.custom_code, "FlexRange"); - EXPECT_EQ(key.version, 1); - EXPECT_TRUE(key.is_flex_op); - } -} - -// TODO(ahentz): tests for tensors, inputs, outputs, opcodes and operators. - -} // namespace -} // namespace tflite -} // namespace toco diff --git a/tensorflow/contrib/lite/toco/tflite/operator.cc b/tensorflow/contrib/lite/toco/tflite/operator.cc deleted file mode 100644 index a4fce2aa9e48530e8ba87985a013ca6136d1eed6..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/toco/tflite/operator.cc +++ /dev/null @@ -1,1621 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include "tensorflow/contrib/lite/toco/tflite/operator.h" - -// TODO(ycling): Consider refactoring to extract the LSTM definition out of -// graph_transformation module. -#include "tensorflow/contrib/lite/toco/graph_transformations/lstm_utils.h" -#include "tensorflow/contrib/lite/toco/tflite/builtin_operator.h" -#include "tensorflow/contrib/lite/toco/tflite/custom_operator.h" -#include "tensorflow/contrib/lite/toco/tflite/simple_operator.h" -#include "tensorflow/contrib/lite/toco/tflite/types.h" -#include "tensorflow/core/framework/attr_value.pb.h" -#include "tensorflow/core/framework/node_def.pb.h" -#include "tensorflow/core/framework/op.h" -#include "tensorflow/core/framework/op_def.pb.h" -#include "tensorflow/core/util/ptr_util.h" - -namespace toco { - -namespace tflite { - -class AveragePool - : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - auto padding = Padding::Serialize(op.padding.type); - auto activation_function = - ActivationFunction::Serialize(op.fused_activation_function); - return ::tflite::CreatePool2DOptions(*builder, padding, op.stride_width, - op.stride_height, op.kwidth, - op.kheight, activation_function); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->padding.type = Padding::Deserialize(options.padding()); - op->stride_width = options.stride_w(); - op->stride_height = options.stride_h(); - op->kwidth = options.filter_width(); - op->kheight = options.filter_height(); - op->fused_activation_function = - ActivationFunction::Deserialize(options.fused_activation_function()); - } - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class Convolution - : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - auto padding = Padding::Serialize(op.padding.type); - auto activation_function = - ActivationFunction::Serialize(op.fused_activation_function); - return ::tflite::CreateConv2DOptions(*builder, padding, op.stride_width, - op.stride_height, activation_function, - op.dilation_width_factor, - op.dilation_height_factor); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->padding.type = Padding::Deserialize(options.padding()); - op->stride_width = options.stride_w(); - op->stride_height = options.stride_h(); - op->dilation_width_factor = options.dilation_w_factor(); - op->dilation_height_factor = options.dilation_h_factor(); - op->fused_activation_function = - ActivationFunction::Deserialize(options.fused_activation_function()); - } - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class DepthwiseConvolution - : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - auto padding = Padding::Serialize(op.padding.type); - auto activation_function = - ActivationFunction::Serialize(op.fused_activation_function); - return ::tflite::CreateDepthwiseConv2DOptions( - *builder, padding, op.stride_width, op.stride_height, - op.depth_multiplier, activation_function, op.dilation_width_factor, - op.dilation_height_factor); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->padding.type = Padding::Deserialize(options.padding()); - op->stride_width = options.stride_w(); - op->stride_height = options.stride_h(); - op->depth_multiplier = options.depth_multiplier(); - op->fused_activation_function = - ActivationFunction::Deserialize(options.fused_activation_function()); - op->dilation_width_factor = options.dilation_w_factor(); - op->dilation_height_factor = options.dilation_h_factor(); - } - - int GetVersion(const Operator& op) const override { - const auto& conv_op = static_cast(op); - if (conv_op.dilation_width_factor != 1 || - conv_op.dilation_height_factor != 1) { - return 2; - } - return 1; - } -}; - -class Add : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - auto activation_function = - ActivationFunction::Serialize(op.fused_activation_function); - return ::tflite::CreateAddOptions(*builder, activation_function); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->fused_activation_function = - ActivationFunction::Deserialize(options.fused_activation_function()); - } - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class SpaceToBatchND - : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - return ::tflite::CreateSpaceToBatchNDOptions(*builder); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override {} - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class Sub : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - auto activation_function = - ActivationFunction::Serialize(op.fused_activation_function); - return ::tflite::CreateSubOptions(*builder, activation_function); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->fused_activation_function = - ActivationFunction::Deserialize(options.fused_activation_function()); - } - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class Div : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - auto activation_function = - ActivationFunction::Serialize(op.fused_activation_function); - return ::tflite::CreateDivOptions(*builder, activation_function); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->fused_activation_function = - ActivationFunction::Deserialize(options.fused_activation_function()); - } - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class BatchToSpaceND - : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - return ::tflite::CreateBatchToSpaceNDOptions(*builder); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override {} - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class Cast : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - return ::tflite::CreateCastOptions(*builder, - DataType::Serialize(op.src_data_type), - DataType::Serialize(op.dst_data_type)); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->src_data_type = DataType::Deserialize(options.in_data_type()); - op->dst_data_type = DataType::Deserialize(options.out_data_type()); - } - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class Concatenation - : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - return ::tflite::CreateConcatenationOptions(*builder, op.axis); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->axis = options.axis(); - } - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class DepthToSpace : public CustomOperator { - public: - using CustomOperator::CustomOperator; - void WriteOptions(const TocoOperator& op, - flexbuffers::Builder* fbb) const override { - fbb->Int("block_size", op.block_size); - } - void ReadOptions(const flexbuffers::Map& m, TocoOperator* op) const override { - op->block_size = m["block_size"].AsInt64(); - } - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class FakeQuant - : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - return ::tflite::CreateFakeQuantOptions( - *builder, op.minmax->min, op.minmax->max, op.num_bits, op.narrow_range); - } - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - auto* minmax = new MinMax; - minmax->min = options.min(); - minmax->max = options.max(); - op->minmax.reset(minmax); - op->num_bits = options.num_bits(); - op->narrow_range = options.narrow_range(); - } - - int GetVersion(const Operator& op) const override { - const auto& fq_op = static_cast(op); - return fq_op.narrow_range ? 2 : 1; - } -}; - -class FullyConnected - : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - auto activation_function = - ActivationFunction::Serialize(op.fused_activation_function); - ::tflite::FullyConnectedOptionsWeightsFormat tflite_weights_format; - switch (op.weights_format) { - case FullyConnectedWeightsFormat::kDefault: - tflite_weights_format = - ::tflite::FullyConnectedOptionsWeightsFormat_DEFAULT; - break; - case FullyConnectedWeightsFormat::kShuffled4x16Int8: - tflite_weights_format = - ::tflite::FullyConnectedOptionsWeightsFormat_SHUFFLED4x16INT8; - break; - default: - LOG(ERROR) << "Unhandled FC weights format"; - tflite_weights_format = - ::tflite::FullyConnectedOptionsWeightsFormat_DEFAULT; - } - return ::tflite::CreateFullyConnectedOptions(*builder, activation_function, - tflite_weights_format); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->fused_activation_function = - ActivationFunction::Deserialize(options.fused_activation_function()); - switch (options.weights_format()) { - case ::tflite::FullyConnectedOptionsWeightsFormat_DEFAULT: - op->weights_format = FullyConnectedWeightsFormat::kDefault; - break; - case ::tflite::FullyConnectedOptionsWeightsFormat_SHUFFLED4x16INT8: - op->weights_format = FullyConnectedWeightsFormat::kShuffled4x16Int8; - break; - default: - LOG(ERROR) << "Unhandled FC weights format"; - op->weights_format = FullyConnectedWeightsFormat::kDefault; - } - } - - int GetVersion(const Operator& op) const override { - const auto& fc_op = static_cast(op); - return fc_op.weights_format == FullyConnectedWeightsFormat::kDefault ? 1 - : 2; - } -}; - -class Gather : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - int axis = op.axis ? op.axis.value() : 0; - return ::tflite::CreateGatherOptions(*builder, axis); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->axis = {options.axis()}; - } - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class Svdf : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - auto activation_function = - ActivationFunction::Serialize(op.fused_activation_function); - return ::tflite::CreateSVDFOptions(*builder, op.rank, activation_function); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->fused_activation_function = - ActivationFunction::Deserialize(options.fused_activation_function()); - op->rank = options.rank(); - } - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class L2Normalization - : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - auto activation_function = - ActivationFunction::Serialize(op.fused_activation_function); - return ::tflite::CreateL2NormOptions(*builder, activation_function); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->fused_activation_function = - ActivationFunction::Deserialize(options.fused_activation_function()); - } - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class L2Pool : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - auto padding = Padding::Serialize(op.padding.type); - auto activation_function = - ActivationFunction::Serialize(op.fused_activation_function); - return ::tflite::CreatePool2DOptions(*builder, padding, op.stride_width, - op.stride_height, op.kwidth, - op.kheight, activation_function); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->padding.type = Padding::Deserialize(options.padding()); - op->stride_width = options.stride_w(); - op->stride_height = options.stride_h(); - op->kwidth = options.filter_width(); - op->kheight = options.filter_height(); - op->fused_activation_function = - ActivationFunction::Deserialize(options.fused_activation_function()); - } - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class LocalResponseNormalization - : public BuiltinOperator< - LocalResponseNormalizationOperator, - ::tflite::LocalResponseNormalizationOptions, - ::tflite::BuiltinOptions_LocalResponseNormalizationOptions> { - public: - using BuiltinOperator::BuiltinOperator; - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - return ::tflite::CreateLocalResponseNormalizationOptions( - *builder, op.range, op.bias, op.alpha, op.beta); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->range = options.radius(); - op->bias = options.bias(); - op->alpha = options.alpha(); - op->beta = options.beta(); - } - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class MaxPool : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - auto padding = Padding::Serialize(op.padding.type); - auto activation_function = - ActivationFunction::Serialize(op.fused_activation_function); - return ::tflite::CreatePool2DOptions(*builder, padding, op.stride_width, - op.stride_height, op.kwidth, - op.kheight, activation_function); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->padding.type = Padding::Deserialize(options.padding()); - op->stride_width = options.stride_w(); - op->stride_height = options.stride_h(); - op->kwidth = options.filter_width(); - op->kheight = options.filter_height(); - op->fused_activation_function = - ActivationFunction::Deserialize(options.fused_activation_function()); - } - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class Mul : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - auto activation_function = - ActivationFunction::Serialize(op.fused_activation_function); - return ::tflite::CreateMulOptions(*builder, activation_function); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->fused_activation_function = - ActivationFunction::Deserialize(options.fused_activation_function()); - } - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class Pad : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - return ::tflite::CreatePadOptions(*builder); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override {} - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class Tile - : public BuiltinOperator { - using BuiltinOperator::BuiltinOperator; - - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - return ::tflite::CreateTileOptions(*builder); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override {} - int GetVersion(const Operator& op) const override { return 1; } -}; - -class PadV2 : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - return ::tflite::CreatePadV2Options(*builder); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override {} - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class Reshape - : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - return ::tflite::CreateReshapeOptions(*builder, - builder->CreateVector(op.shape)); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->shape.insert(op->shape.end(), options.new_shape()->begin(), - options.new_shape()->end()); - } - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class Softmax - : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - return ::tflite::CreateSoftmaxOptions(*builder, op.beta); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->beta = options.beta(); - } - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class SpaceToDepth - : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - return ::tflite::CreateSpaceToDepthOptions(*builder, op.block_size); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->block_size = options.block_size(); - } - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class Transpose - : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - return ::tflite::CreateTransposeOptions(*builder); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override {} - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class Lstm : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - ::tflite::LSTMKernelType kernel_type; - switch (op.kernel_type) { - case LstmCellOperator::KERNEL_BASIC: - kernel_type = ::tflite::LSTMKernelType_BASIC; - break; - case LstmCellOperator::KERNEL_FULL: - kernel_type = ::tflite::LSTMKernelType_FULL; - break; - } - - // Current toco converter only supports tanh, no clip. - return ::tflite::CreateLSTMOptions(*builder, /*fused_activation_function=*/ - ::tflite::ActivationFunctionType_TANH, - /*cell_clip=*/0.0, - /*proj_clip=*/0.0, kernel_type); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - // Only support tanh activation, so check that tflite type is tanh. - CHECK(options.fused_activation_function() == - ::tflite::ActivationFunctionType_TANH); - - switch (options.kernel_type()) { - case ::tflite::LSTMKernelType_BASIC: - op->kernel_type = LstmCellOperator::KERNEL_BASIC; - break; - case ::tflite::LSTMKernelType_FULL: - op->kernel_type = LstmCellOperator::KERNEL_FULL; - break; - } - } - - int GetVersion(const Operator& op) const override { - const auto& lstm_op = static_cast(op); - switch (lstm_op.kernel_type) { - case LstmCellOperator::KERNEL_FULL: - return 1; - case LstmCellOperator::KERNEL_BASIC: - return 2; - } - } - - std::vector GetMutatingInputVariables( - const Operator& op) const override { - const auto& lstm_op = static_cast(op); - - std::vector mutating_input_variables(op.inputs.size(), false); - switch (lstm_op.kernel_type) { - case LstmCellOperator::KERNEL_FULL: { - mutating_input_variables[kInputActivationStateTensor] = true; - mutating_input_variables[kInputCellStateTensor] = true; - break; - } - case LstmCellOperator::KERNEL_BASIC: { - mutating_input_variables[LstmCellOperator::PREV_ACTIV_INPUT] = true; - mutating_input_variables[LstmCellOperator::PREV_STATE_INPUT] = true; - break; - } - } - return mutating_input_variables; - } -}; - -class UnidirectionalSequenceLstm - : public BuiltinOperator< - UnidirectionalSequenceLstmOperator, - ::tflite::UnidirectionalSequenceLSTMOptions, - ::tflite::BuiltinOptions_UnidirectionalSequenceLSTMOptions> { - public: - using BuiltinOperator::BuiltinOperator; - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - // Current toco converter only supports tanh, no clip. - return ::tflite::CreateUnidirectionalSequenceLSTMOptions( - *builder, /*fused_activation_function=*/ - ::tflite::ActivationFunctionType_TANH, - /*cell_clip=*/0.0, - /*proj_clip=*/0.0, - /*time_major=*/true); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - // Only support tanh activation, so check that tflite type is tanh. - DCHECK(options.fused_activation_function() == - ::tflite::ActivationFunctionType_TANH); - } - - int GetVersion(const Operator& op) const override { return 1; } - - std::vector GetMutatingInputVariables( - const Operator& op) const override { - std::vector mutating_input_variables(op.inputs.size(), false); - mutating_input_variables[kInputActivationStateTensor] = true; - mutating_input_variables[kInputCellStateTensor] = true; - return mutating_input_variables; - } -}; - -class Mean : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - return ::tflite::CreateReducerOptions(*builder, op.keep_dims); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->keep_dims = options.keep_dims(); - } - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class Sum - : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - return ::tflite::CreateReducerOptions(*builder, op.keep_dims); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->keep_dims = options.keep_dims(); - } - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class ReduceMax - : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - return ::tflite::CreateReducerOptions(*builder, op.keep_dims); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->keep_dims = options.keep_dims(); - } - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class ReduceMin - : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - return ::tflite::CreateReducerOptions(*builder, op.keep_dims); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->keep_dims = options.keep_dims(); - } - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class ReduceProd - : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - return ::tflite::CreateReducerOptions(*builder, op.keep_dims); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->keep_dims = options.keep_dims(); - } - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class ReduceAny - : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - return ::tflite::CreateReducerOptions(*builder, op.keep_dims); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->keep_dims = options.keep_dims(); - } - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class ResizeBilinear - : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - return ::tflite::CreateResizeBilinearOptions(*builder, op.align_corners); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->align_corners = options.align_corners(); - } - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class Squeeze - : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - auto squeeze_dims = builder->CreateVector(op.squeeze_dims); - return ::tflite::CreateSqueezeOptions(*builder, squeeze_dims); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->squeeze_dims.insert(op->squeeze_dims.end(), - options.squeeze_dims()->begin(), - options.squeeze_dims()->end()); - } - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class Split - : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - return ::tflite::CreateSplitOptions(*builder, op.num_split); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->num_split = options.num_splits(); - } - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class StridedSlice - : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - return ::tflite::CreateStridedSliceOptions( - *builder, op.begin_mask, op.end_mask, op.ellipsis_mask, - op.new_axis_mask, op.shrink_axis_mask); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->begin_mask = options.begin_mask(); - op->end_mask = options.end_mask(); - op->ellipsis_mask = options.ellipsis_mask(); - op->new_axis_mask = options.new_axis_mask(); - op->shrink_axis_mask = options.shrink_axis_mask(); - } - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class TopK_V2 : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - return ::tflite::CreateTopKV2Options(*builder); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override {} - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class ArgMax : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - return ::tflite::CreateArgMaxOptions( - *builder, DataType::Serialize(op.output_data_type)); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->output_data_type = DataType::Deserialize(options.output_type()); - } - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class ArgMin : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - return ::tflite::CreateArgMinOptions( - *builder, DataType::Serialize(op.output_data_type)); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->output_data_type = DataType::Deserialize(options.output_type()); - } - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class TransposeConv - : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - auto padding = Padding::Serialize(op.padding.type); - return ::tflite::CreateTransposeConvOptions( - *builder, padding, op.stride_width, op.stride_height); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->padding.type = Padding::Deserialize(options.padding()); - op->stride_width = options.stride_w(); - op->stride_height = options.stride_h(); - } - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class SparseToDense - : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - return ::tflite::CreateSparseToDenseOptions(*builder, op.validate_indices); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->validate_indices = options.validate_indices(); - } - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class ExpandDims - : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - return ::tflite::CreateExpandDimsOptions(*builder); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override {} - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class Pack : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - return ::tflite::CreatePackOptions(*builder, op.values_count, op.axis); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->values_count = options.values_count(); - op->axis = options.axis(); - } - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class Shape - : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - return ::tflite::CreateShapeOptions( - *builder, DataType::Serialize(op.output_data_type)); - } - - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->output_data_type = DataType::Deserialize(options.out_type()); - } - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class OneHot : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - return ::tflite::CreateOneHotOptions(*builder, op.axis); - } - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->axis = options.axis(); - } - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class CTCBeamSearchDecoder - : public CustomOperator { - public: - using CustomOperator::CustomOperator; - - void WriteOptions(const TocoOperator& op, - flexbuffers::Builder* fbb) const override { - fbb->Int("beam_width", op.beam_width); - fbb->Int("top_paths", op.top_paths); - fbb->Bool("merge_repeated", op.merge_repeated); - } - - void ReadOptions(const flexbuffers::Map& m, TocoOperator* op) const override { - op->beam_width = m["beam_width"].AsInt32(); - op->top_paths = m["top_paths"].AsInt32(); - op->merge_repeated = m["merge_repeated"].AsBool(); - } - - int GetVersion(const Operator& op) const override { return 1; } -}; - -class Unpack : public BuiltinOperator { - public: - using BuiltinOperator::BuiltinOperator; - flatbuffers::Offset WriteOptions( - const TocoOperator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - return ::tflite::CreateUnpackOptions(*builder, op.num, op.axis); - } - void ReadOptions(const TfLiteOptions& options, - TocoOperator* op) const override { - op->num = options.num(); - op->axis = options.axis(); - } - - int GetVersion(const Operator& op) const override { return 1; } -}; - -std::unique_ptr WriteFlexOpOptions( - const string& tensorflow_node_def) { - auto fbb = absl::make_unique(); - - ::tensorflow::NodeDef node_def; - if (!node_def.ParseFromString(tensorflow_node_def)) { - LOG(ERROR) << "Failed to parse TensorFlow NodeDef"; - return {}; - } - - fbb->Vector([&]() { - fbb->String(node_def.op()); - fbb->String(tensorflow_node_def); - }); - fbb->Finish(); - LOG(INFO) << "Writing flex op: " << node_def.op(); - return std::unique_ptr(fbb.release()); -} - -class TensorFlowUnsupported : public BaseOperator { - public: - TensorFlowUnsupported(const string& name, OperatorType type, - bool allow_flex_ops) - : BaseOperator(name, type), allow_flex_ops_(allow_flex_ops) {} - - Options Serialize(const Operator& op, - flatbuffers::FlatBufferBuilder* builder) const override { - auto fbb = - WriteOptions(static_cast(op)); - if (fbb) { - return Options::Custom(builder->CreateVector(fbb->GetBuffer())); - } else { - return Options::Custom(0); - } - } - - std::unique_ptr Deserialize( - const BuiltinOptions* builtin_options, - const CustomOptions* custom_options) const override { - // Deserializing Flex ops doesn't work now. - // TODO(ycling): Revisit and decide if we should fix the flow for importing - // TFLite models with Flex ops. - auto op = absl::make_unique(); - if (custom_options) { - auto flexbuffer_map = - flexbuffers::GetRoot(custom_options->data(), custom_options->size()) - .AsMap(); - ReadOptions(flexbuffer_map, op.get()); - } - return std::unique_ptr(op.release()); - } - - std::unique_ptr WriteOptions( - const TensorFlowUnsupportedOperator& op) const { - if (allow_flex_ops_) { - return WriteFlexOpOptions(op.tensorflow_node_def); - } - auto fbb = absl::make_unique(); - - ::tensorflow::NodeDef node_def; - if (!node_def.ParseFromString(op.tensorflow_node_def)) { - LOG(ERROR) << "Failed to parse TensorFlow NodeDef"; - return std::unique_ptr(); - } - - if (ShouldExportAsFlexOp(allow_flex_ops_, node_def.op())) { - fbb->Vector([&]() { - fbb->String(node_def.op()); - fbb->String(op.tensorflow_node_def); - }); - fbb->Finish(); - LOG(INFO) << "Writing flex op: " << node_def.op(); - return std::unique_ptr(fbb.release()); - } - - bool has_valid_attr = false; - size_t map_start = fbb->StartMap(); - for (const auto& pair : node_def.attr()) { - const char* key = pair.first.c_str(); - const auto& attr = pair.second; - switch (attr.value_case()) { - case ::tensorflow::AttrValue::kS: - fbb->String(key, attr.s()); - has_valid_attr = true; - break; - case ::tensorflow::AttrValue::kI: - fbb->Int(key, attr.i()); - has_valid_attr = true; - break; - case ::tensorflow::AttrValue::kF: - fbb->Float(key, attr.f()); - has_valid_attr = true; - break; - case ::tensorflow::AttrValue::kB: - fbb->Bool(key, attr.b()); - has_valid_attr = true; - break; - case tensorflow::AttrValue::kList: - if (attr.list().i_size() > 0) { - auto start = fbb->StartVector(key); - for (const int64_t v : attr.list().i()) { - fbb->Add(v); - } - fbb->EndVector(start, /*typed=*/true, /*fixed=*/false); - has_valid_attr = true; - } else { - LOG(WARNING) - << "Ignoring unsupported type in list attribute with key '" - << key << "'"; - } - break; - default: - LOG(WARNING) << "Ignoring unsupported attribute type with key '" - << key << "'"; - break; - } - } - if (!has_valid_attr) { - return std::unique_ptr(); - } - fbb->EndMap(map_start); - fbb->Finish(); - return std::unique_ptr(fbb.release()); - } - -// TODO(wvo): hack to make this code compile with 2 different API versions. -// Please remove once OS/internal versions are in sync. -// See hardcoded values in the switch below. - - void ReadOptions(const flexbuffers::Map& m, - TensorFlowUnsupportedOperator* op) const { - ::tensorflow::NodeDef node_def; - auto attr = node_def.mutable_attr(); - - const auto& keys = m.Keys(); - for (size_t i = 0; i < keys.size(); ++i) { - const auto key = keys[i].AsKey(); - const auto& value = m[key]; - switch (value.GetType()) { - case 5: // flexbuffers::FBT_STRING: - (*attr)[key].set_s(value.AsString().c_str()); - break; - case 1: // flexbuffers::FBT_INT: - (*attr)[key].set_i(value.AsInt64()); - break; - case 3: // flexbuffers::FBT_FLOAT: - (*attr)[key].set_f(value.AsFloat()); - break; - case 26: // flexbuffers::FBT_BOOL: - (*attr)[key].set_b(value.AsBool()); - if (string(key) == "_output_quantized") { - op->quantized = value.AsBool(); - } - if (string(key) == "_support_output_type_float_in_quantized_op") { - op->support_output_type_float_in_quantized_op = value.AsBool(); - } - break; - case 11: { // flexbuffers::FBT_VECTOR_INT: { - auto* list = (*attr)[key].mutable_list(); - const auto& vector = value.AsTypedVector(); - for (size_t i = 0; i < vector.size(); i++) { - list->add_i(vector[i].AsInt64()); - } - break; - } - default: - LOG(WARNING) << "Ignoring unsupported attribute type with key '" - << key << "'"; - break; - } - } - node_def.SerializeToString(&op->tensorflow_node_def); - } - - int GetVersion(const Operator& op) const override { - // TODO(ycling): Deisng and implement a way to plumb the version of - // custom ops. - return 1; - } - - private: - const bool allow_flex_ops_; -}; - -namespace { -// Build a vector containing all the known operators. -std::vector> BuildOperatorList( - bool allow_flex_ops = false) { - std::vector> ops; - using tensorflow::MakeUnique; - // Builtin Operators. - ops.push_back( - MakeUnique(::tflite::BuiltinOperator_ADD, OperatorType::kAdd)); - ops.push_back( - MakeUnique
(::tflite::BuiltinOperator_DIV, OperatorType::kDiv)); - ops.push_back( - MakeUnique(::tflite::BuiltinOperator_SUB, OperatorType::kSub)); - ops.push_back(MakeUnique( - ::tflite::BuiltinOperator_AVERAGE_POOL_2D, OperatorType::kAveragePool)); - ops.push_back( - MakeUnique(::tflite::BuiltinOperator_SPACE_TO_BATCH_ND, - OperatorType::kSpaceToBatchND)); - ops.push_back( - MakeUnique(::tflite::BuiltinOperator_BATCH_TO_SPACE_ND, - OperatorType::kBatchToSpaceND)); - ops.push_back(MakeUnique( - ::tflite::BuiltinOperator_CONCATENATION, OperatorType::kConcatenation)); - ops.push_back(MakeUnique(::tflite::BuiltinOperator_CONV_2D, - OperatorType::kConv)); - ops.push_back(MakeUnique( - ::tflite::BuiltinOperator_DEPTHWISE_CONV_2D, - OperatorType::kDepthwiseConv)); - ops.push_back( - MakeUnique(::tflite::BuiltinOperator_FULLY_CONNECTED, - OperatorType::kFullyConnected)); - ops.push_back(MakeUnique(::tflite::BuiltinOperator_GATHER, - OperatorType::kGather)); - ops.push_back( - MakeUnique(::tflite::BuiltinOperator_L2_NORMALIZATION, - OperatorType::kL2Normalization)); - ops.push_back(MakeUnique(::tflite::BuiltinOperator_L2_POOL_2D, - OperatorType::kL2Pool)); - ops.push_back(MakeUnique( - ::tflite::BuiltinOperator_LOCAL_RESPONSE_NORMALIZATION, - OperatorType::kLocalResponseNormalization)); - ops.push_back(MakeUnique(::tflite::BuiltinOperator_MAX_POOL_2D, - OperatorType::kMaxPool)); - ops.push_back( - MakeUnique(::tflite::BuiltinOperator_MUL, OperatorType::kMul)); - ops.push_back( - MakeUnique(::tflite::BuiltinOperator_PAD, OperatorType::kPad)); - ops.push_back( - MakeUnique(::tflite::BuiltinOperator_PADV2, OperatorType::kPadV2)); - ops.push_back(MakeUnique(::tflite::BuiltinOperator_RESHAPE, - OperatorType::kReshape)); - ops.push_back(MakeUnique(::tflite::BuiltinOperator_SOFTMAX, - OperatorType::kSoftmax)); - ops.push_back(MakeUnique( - ::tflite::BuiltinOperator_SPACE_TO_DEPTH, OperatorType::kSpaceToDepth)); - ops.push_back( - MakeUnique(::tflite::BuiltinOperator_SVDF, OperatorType::kSvdf)); - ops.push_back(MakeUnique(::tflite::BuiltinOperator_TRANSPOSE, - OperatorType::kTranspose)); - ops.push_back( - MakeUnique(::tflite::BuiltinOperator_MEAN, OperatorType::kMean)); - ops.push_back( - MakeUnique(::tflite::BuiltinOperator_SUM, OperatorType::kSum)); - ops.push_back(MakeUnique(::tflite::BuiltinOperator_REDUCE_PROD, - OperatorType::kReduceProd)); - ops.push_back(MakeUnique(::tflite::BuiltinOperator_REDUCE_MAX, - OperatorType::kReduceMax)); - ops.push_back(MakeUnique(::tflite::BuiltinOperator_REDUCE_MIN, - OperatorType::kReduceMin)); - ops.push_back(MakeUnique(::tflite::BuiltinOperator_REDUCE_ANY, - OperatorType::kAny)); - ops.push_back( - MakeUnique(::tflite::BuiltinOperator_RESIZE_BILINEAR, - OperatorType::kResizeBilinear)); - ops.push_back(MakeUnique(::tflite::BuiltinOperator_SQUEEZE, - OperatorType::kSqueeze)); - ops.push_back( - MakeUnique(::tflite::BuiltinOperator_SPLIT, OperatorType::kSplit)); - ops.push_back(MakeUnique( - ::tflite::BuiltinOperator_STRIDED_SLICE, OperatorType::kStridedSlice)); - ops.push_back(MakeUnique(::tflite::BuiltinOperator_TOPK_V2, - OperatorType::kTopK_V2)); - ops.push_back(MakeUnique(::tflite::BuiltinOperator_LSTM, - OperatorType::kLstmCell)); - ops.push_back( - MakeUnique(::tflite::BuiltinOperator_CAST, OperatorType::kCast)); - ops.push_back(MakeUnique(::tflite::BuiltinOperator_ARG_MAX, - OperatorType::kArgMax)); - ops.push_back(MakeUnique(::tflite::BuiltinOperator_ARG_MIN, - OperatorType::kArgMin)); - ops.push_back( - MakeUnique(::tflite::BuiltinOperator_TILE, OperatorType::kTile)); - ops.push_back(MakeUnique(::tflite::BuiltinOperator_EXPAND_DIMS, - OperatorType::kExpandDims)); - ops.push_back(MakeUnique( - ::tflite::BuiltinOperator_TRANSPOSE_CONV, OperatorType::kTransposeConv)); - ops.push_back(MakeUnique( - ::tflite::BuiltinOperator_SPARSE_TO_DENSE, OperatorType::kSparseToDense)); - ops.push_back( - MakeUnique(::tflite::BuiltinOperator_SHAPE, OperatorType::kShape)); - ops.push_back(MakeUnique(::tflite::BuiltinOperator_FAKE_QUANT, - OperatorType::kFakeQuant)); - ops.push_back( - MakeUnique(::tflite::BuiltinOperator_PACK, OperatorType::kPack)); - ops.emplace_back(MakeUnique( - ::tflite::BuiltinOperator_UNIDIRECTIONAL_SEQUENCE_LSTM, - OperatorType::kUnidirectionalSequenceLstm)); - ops.push_back(MakeUnique(::tflite::BuiltinOperator_ONE_HOT, - OperatorType::kOneHot)); - ops.push_back(MakeUnique(::tflite::BuiltinOperator_UNPACK, - OperatorType::kUnpack)); - - // Custom Operators. - ops.push_back( - MakeUnique("DEPTH_TO_SPACE", OperatorType::kDepthToSpace)); - ops.push_back(MakeUnique( - "CTC_BEAM_SEARCH_DECODER", OperatorType::kCTCBeamSearchDecoder)); - ops.push_back(MakeUnique( - "TENSORFLOW_UNSUPPORTED", OperatorType::kUnsupported, allow_flex_ops)); - - // There operators are supported by Toco, but not by TF Lite, and has no - // attributes. - ops.push_back( - MakeUnique>("ADDN", OperatorType::kAddN)); - // Simple Operators. - ops.push_back(MakeUnique>( - "DEQUANTIZE", OperatorType::kDequantize)); - ops.push_back( - MakeUnique>("FLOOR", OperatorType::kFloor)); - ops.push_back( - MakeUnique>("RELU", OperatorType::kRelu)); - ops.push_back(MakeUnique>( - "RELU_N1_TO_1", OperatorType::kRelu1)); - ops.push_back( - MakeUnique>("RELU6", OperatorType::kRelu6)); - ops.push_back( - MakeUnique>("PRELU", OperatorType::kPRelu)); - ops.push_back(MakeUnique>( - "LOGISTIC", OperatorType::kLogistic)); - ops.push_back( - MakeUnique>("TANH", OperatorType::kTanh)); - ops.push_back( - MakeUnique>("EXP", OperatorType::kExp)); - ops.push_back(MakeUnique>( - "LOG_SOFTMAX", OperatorType::kLogSoftmax)); - ops.push_back(MakeUnique>( - "MAXIMUM", OperatorType::kMaximum)); // Element-wise Maximum - ops.push_back(MakeUnique>( - "MINIMUM", OperatorType::kMinimum)); // Element-wise Minimum - ops.push_back(MakeUnique>( - "GREATER", OperatorType::kGreater)); - ops.push_back(MakeUnique>( - "GREATER_EQUAL", OperatorType::kGreaterEqual)); - ops.push_back(MakeUnique>( - "LESS", OperatorType::kLess)); - ops.push_back(MakeUnique>( - "LESS_EQUAL", OperatorType::kLessEqual)); - ops.push_back(MakeUnique>( - "EQUAL", OperatorType::kEqual)); - ops.push_back(MakeUnique>( - "NOT_EQUAL", OperatorType::kNotEqual)); - ops.push_back( - MakeUnique>("NEG", OperatorType::kNeg)); - ops.push_back(MakeUnique>( - "SELECT", OperatorType::kSelect)); - ops.push_back( - MakeUnique>("SLICE", OperatorType::kSlice)); - ops.push_back( - MakeUnique>("POW", OperatorType::kPow)); - ops.push_back(MakeUnique>( - "LOGICAL_OR", OperatorType::kLogicalOr)); - ops.emplace_back(new SimpleOperator( - "LOGICAL_AND", OperatorType::kLogicalAnd)); - ops.emplace_back(new SimpleOperator( - "LOGICAL_NOT", OperatorType::kLogicalNot)); - ops.emplace_back(new SimpleOperator( - "FLOOR_DIV", OperatorType::kFloorDiv)); - // Element-wise operator - ops.push_back( - MakeUnique>("SIN", OperatorType::kSin)); - ops.push_back( - MakeUnique>("LOG", OperatorType::kLog)); - ops.push_back(MakeUnique>( - "SQRT", OperatorType::kSqrt)); - ops.push_back(MakeUnique>( - "RSQRT", OperatorType::kRsqrt)); - ops.push_back(MakeUnique>( - "SQUARE", OperatorType::kSquare)); - ops.push_back(MakeUnique>( - "ZEROS_LIKE", OperatorType::kZerosLike)); - - return ops; -} -} // namespace - -std::map> BuildOperatorByTypeMap( - bool allow_flex_ops) { - std::map> result; - - std::vector> ops = - BuildOperatorList(allow_flex_ops); - for (auto& op : ops) { - result[op->type()] = std::move(op); - } - - return result; -} - -std::map> BuildOperatorByNameMap( - bool allow_flex_ops) { - std::map> result; - - std::vector> ops = - BuildOperatorList(allow_flex_ops); - for (auto& op : ops) { - result[op->name()] = std::move(op); - } - - return result; -} - -bool ShouldExportAsFlexOp(bool allow_flex_ops, - const string& tensorflow_op_name) { - // If Flex ops aren't allow at all, simply return false. - if (!allow_flex_ops) { - return false; - } - // Check if we can find the `OpDef` for the TensorFlow op. If we can find - // it, export the op as an Flex op. Otherwise, export it as a regular custom - // op. - const tensorflow::OpDef* op_def = nullptr; - return tensorflow::OpRegistry::Global() - ->LookUpOpDef(tensorflow_op_name, &op_def) - .ok(); -} - -} // namespace tflite - -} // namespace toco diff --git a/tensorflow/contrib/lite/toco/tflite/types.h b/tensorflow/contrib/lite/toco/tflite/types.h deleted file mode 100644 index 3923756fc94e3175a6505740a96cce8d614c3990..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/toco/tflite/types.h +++ /dev/null @@ -1,58 +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_CONTRIB_LITE_TOCO_TFLITE_TYPES_H_ -#define TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_TYPES_H_ - -#include "tensorflow/contrib/lite/schema/schema_generated.h" -#include "tensorflow/contrib/lite/toco/model.h" - -namespace toco { - -namespace tflite { - -struct DataType { - static ::tflite::TensorType Serialize(ArrayDataType array_data_type); - static ArrayDataType Deserialize(int tensor_type); -}; - -struct DataBuffer { - using FlatBufferOffset = flatbuffers::Offset>; - - // Build the flatbuffer representation of a toco's Array and return the - // corresponding offset into the flatbuffer. Note that data from the array - // will be copied into the flatbuffer. - static FlatBufferOffset Serialize(const Array& array, - flatbuffers::FlatBufferBuilder* builder); - // Copy data from the given tensor into toco's Array. - static void Deserialize(const ::tflite::Tensor& tensor, - const ::tflite::Buffer& buffer, Array* array); -}; - -struct Padding { - static ::tflite::Padding Serialize(PaddingType padding_type); - static PaddingType Deserialize(int padding); -}; - -struct ActivationFunction { - static ::tflite::ActivationFunctionType Serialize( - FusedActivationFunctionType faf_type); - static FusedActivationFunctionType Deserialize(int activation_function); -}; - -} // namespace tflite - -} // namespace toco - -#endif // TENSORFLOW_CONTRIB_LITE_TOCO_TFLITE_TYPES_H_ diff --git a/tensorflow/contrib/lite/toco/toco.cc b/tensorflow/contrib/lite/toco/toco.cc deleted file mode 100644 index 0b460bd178a49cafefd3438b7ae1c38a07b2ab7c..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/toco/toco.cc +++ /dev/null @@ -1,128 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include -#include -#include - -#include "absl/strings/string_view.h" -#include "tensorflow/contrib/lite/toco/model.h" -#include "tensorflow/contrib/lite/toco/model_cmdline_flags.h" -#include "tensorflow/contrib/lite/toco/model_flags.pb.h" -#include "tensorflow/contrib/lite/toco/toco_cmdline_flags.h" -#include "tensorflow/contrib/lite/toco/toco_flags.pb.h" -#include "tensorflow/contrib/lite/toco/toco_port.h" -#include "tensorflow/contrib/lite/toco/toco_tooling.h" -#include "tensorflow/contrib/lite/toco/toco_types.h" -#include "tensorflow/core/platform/logging.h" - -namespace toco { -namespace { - -// Checks the permissions of the output file to ensure it is writeable. -void CheckOutputFilePermissions(const Arg& output_file) { - QCHECK(output_file.specified()) << "Missing required flag --output_file.\n"; - QCHECK(port::file::Writable(output_file.value()).ok()) - << "Specified output_file is not writable: " << output_file.value() - << ".\n"; -} - -// Checks the permissions of the frozen model file. -void CheckFrozenModelPermissions(const Arg& input_file) { - QCHECK(input_file.specified()) << "Missing required flag --input_file.\n"; - QCHECK(port::file::Exists(input_file.value(), port::file::Defaults()).ok()) - << "Specified input_file does not exist: " << input_file.value() << ".\n"; - QCHECK(port::file::Readable(input_file.value(), port::file::Defaults()).ok()) - << "Specified input_file exists, but is not readable: " - << input_file.value() << ".\n"; -} - -// Reads the contents of the GraphDef from either the frozen graph file or the -// SavedModel directory. If it reads the SavedModel directory, it updates the -// ModelFlags and TocoFlags accordingly. -void ReadInputData(const ParsedTocoFlags& parsed_toco_flags, - const ParsedModelFlags& parsed_model_flags, - TocoFlags* toco_flags, ModelFlags* model_flags, - string* graph_def_contents) { - port::CheckInitGoogleIsDone("InitGoogle is not done yet.\n"); - - // Ensure savedmodel_directory is not set. - QCHECK(!parsed_toco_flags.savedmodel_directory.specified()) - << "Use `tensorflow/contrib/lite/python/tflite_convert` script with " - << "SavedModel directories.\n"; - - // Checks the input file permissions and reads the contents. - CheckFrozenModelPermissions(parsed_toco_flags.input_file); - CHECK(port::file::GetContents(parsed_toco_flags.input_file.value(), - graph_def_contents, port::file::Defaults()) - .ok()); -} - -void ToolMain(const ParsedTocoFlags& parsed_toco_flags, - const ParsedModelFlags& parsed_model_flags) { - ModelFlags model_flags; - ReadModelFlagsFromCommandLineFlags(parsed_model_flags, &model_flags); - - TocoFlags toco_flags; - ReadTocoFlagsFromCommandLineFlags(parsed_toco_flags, &toco_flags); - - string graph_def_contents; - ReadInputData(parsed_toco_flags, parsed_model_flags, &toco_flags, - &model_flags, &graph_def_contents); - CheckOutputFilePermissions(parsed_toco_flags.output_file); - - std::unique_ptr model = - Import(toco_flags, model_flags, graph_def_contents); - Transform(toco_flags, model.get()); - string output_file_contents; - Export(toco_flags, *model, toco_flags.allow_custom_ops(), - &output_file_contents); - CHECK(port::file::SetContents(parsed_toco_flags.output_file.value(), - output_file_contents, port::file::Defaults()) - .ok()); -} - -} // namespace -} // namespace toco - -int main(int argc, char** argv) { - toco::string msg; - toco::ParsedTocoFlags parsed_toco_flags; - toco::ParsedModelFlags parsed_model_flags; - - // If no args were specified, give a help string to be helpful. - int* effective_argc = &argc; - char** effective_argv = argv; - if (argc == 1) { - // No arguments, so manufacture help argv. - static int dummy_argc = 2; - static char* dummy_argv[] = {argv[0], const_cast("--help")}; - effective_argc = &dummy_argc; - effective_argv = dummy_argv; - } - - // Parse toco flags and command flags in sequence, each one strips off args, - // giving InitGoogle a chance to handle all remaining arguments. - bool toco_success = toco::ParseTocoFlagsFromCommandLineFlags( - effective_argc, effective_argv, &msg, &parsed_toco_flags); - bool model_success = toco::ParseModelFlagsFromCommandLineFlags( - effective_argc, effective_argv, &msg, &parsed_model_flags); - if (!toco_success || !model_success || !msg.empty()) { - fprintf(stderr, "%s", msg.c_str()); - fflush(stderr); - return 1; - } - toco::port::InitGoogle(argv[0], effective_argc, &effective_argv, true); - toco::ToolMain(parsed_toco_flags, parsed_model_flags); -} diff --git a/tensorflow/contrib/lite/toco/toco_tooling.h b/tensorflow/contrib/lite/toco/toco_tooling.h deleted file mode 100644 index e731c149eef412d3048a1d5f84145ce6ff87208d..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/toco/toco_tooling.h +++ /dev/null @@ -1,50 +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_CONTRIB_LITE_TOCO_TOCO_TOOLING_H_ -#define TENSORFLOW_CONTRIB_LITE_TOCO_TOCO_TOOLING_H_ - -#include -#include - -#include "tensorflow/contrib/lite/toco/model.h" -#include "tensorflow/contrib/lite/toco/model_flags.pb.h" -#include "tensorflow/contrib/lite/toco/toco_flags.pb.h" - -namespace toco { - -// Imports the input file into a Model object. -std::unique_ptr Import(const TocoFlags& toco_flags, - const ModelFlags& model_flags, - const string& input_file_contents); - -// Transforms a Model. The resulting Model is ready to be passed -// to Export with the exact same toco_flags. -void Transform(const TocoFlags& toco_flags, Model* model); - -// Exports the Model, which must be of the 'lowered' form returned by -// Transform, to a file of the format given by -// toco_flags.output_format(). -void Export(const TocoFlags& toco_flags, const Model& model, - bool allow_custom_ops, string* output_file_contents); - -// This if for backward-compatibility with internal tools. -inline void Export(const TocoFlags& toco_flags, const Model& model, - string* output_file_contents) { - Export(toco_flags, model, true, output_file_contents); -} - -} // namespace toco - -#endif // TENSORFLOW_CONTRIB_LITE_TOCO_TOCO_TOOLING_H_ diff --git a/tensorflow/contrib/lite/tools/BUILD b/tensorflow/contrib/lite/tools/BUILD deleted file mode 100644 index 0b268264031f4f1e86b2956a75bde173a945ddf4..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/tools/BUILD +++ /dev/null @@ -1,98 +0,0 @@ -package(default_visibility = [ - "//visibility:public", -]) - -licenses(["notice"]) # Apache 2.0 - -load("//tensorflow/contrib/lite:special_rules.bzl", "tflite_portable_test_suite") -load("//tensorflow:tensorflow.bzl", "tf_cc_binary") - -common_copts = ["-Wall"] - -py_binary( - name = "visualize", - srcs = ["visualize.py"], - data = [ - "//tensorflow/contrib/lite/schema:schema.fbs", - "//tensorflow/python:platform", - "@flatbuffers//:flatc", - ], - srcs_version = "PY2AND3", -) - -tf_cc_binary( - name = "generate_op_registrations", - srcs = ["gen_op_registration_main.cc"], - deps = [ - "//tensorflow/contrib/lite/tools:gen_op_registration", - "//tensorflow/core:framework_internal", - "//tensorflow/core:lib", - "@com_google_absl//absl/strings", - ], -) - -cc_library( - name = "gen_op_registration", - srcs = ["gen_op_registration.cc"], - hdrs = ["gen_op_registration.h"], - deps = [ - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite:string", - "@com_googlesource_code_re2//:re2", - ], -) - -cc_test( - name = "gen_op_registration_test", - srcs = ["gen_op_registration_test.cc"], - data = [ - "//tensorflow/contrib/lite:testdata/0_subgraphs.bin", - "//tensorflow/contrib/lite:testdata/2_subgraphs.bin", - "//tensorflow/contrib/lite:testdata/empty_model.bin", - "//tensorflow/contrib/lite:testdata/test_model.bin", - "//tensorflow/contrib/lite:testdata/test_model_broken.bin", - ], - tags = [ - "no_oss", - "tflite_not_portable_android", - "tflite_not_portable_ios", - ], - deps = [ - ":gen_op_registration", - "@com_google_googletest//:gtest", - ], -) - -cc_library( - name = "verifier", - srcs = ["verifier.cc"], - hdrs = ["verifier.h"], - deps = [ - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite:schema_fbs_version", - "//tensorflow/contrib/lite:string_util", - "//tensorflow/contrib/lite/schema:schema_fbs", - ], -) - -cc_test( - name = "verifier_test", - size = "small", - srcs = ["verifier_test.cc"], - tags = [ - "no_oss", - "tflite_not_portable", - ], - deps = [ - ":verifier", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite:schema_fbs_version", - "//tensorflow/contrib/lite/schema:schema_fbs", - "//tensorflow/contrib/lite/testing:util", - "//tensorflow/core:framework_lite", - "@com_google_googletest//:gtest", - "@flatbuffers", - ], -) - -tflite_portable_test_suite() diff --git a/tensorflow/contrib/lite/tools/accuracy/BUILD b/tensorflow/contrib/lite/tools/accuracy/BUILD deleted file mode 100644 index 1b60d6a60d39ccb59613871d1f438b31c16fec7a..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/tools/accuracy/BUILD +++ /dev/null @@ -1,328 +0,0 @@ -package(default_visibility = [ - "//visibility:public", -]) - -licenses(["notice"]) # Apache 2.0 - -load("//tensorflow:tensorflow.bzl", "tf_cc_binary", "tf_cc_test") -load("//tensorflow/contrib/lite:build_def.bzl", "tflite_copts", "tflite_linkopts") -load("//tensorflow/contrib/lite:special_rules.bzl", "tflite_portable_test_suite") - -common_linkopts = tflite_linkopts() + select({ - "//conditions:default": [], - "//tensorflow:android": [ - "-pie", - "-llog", - ], -}) - -cc_library( - name = "utils", - srcs = ["utils.cc"], - hdrs = ["utils.h"], - copts = tflite_copts(), - deps = [ - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:builtin_ops", - ] + select( - { - "//tensorflow:android": [ - "//tensorflow/core:android_tensorflow_lib", - ], - "//conditions:default": [ - "//tensorflow/core:framework", - ], - }, - ), -) - -tf_cc_test( - name = "utils_test", - srcs = ["utils_test.cc"], - args = [ - "--test_model_file=$(location //tensorflow/contrib/lite:testdata/multi_add.bin)", - ], - data = ["//tensorflow/contrib/lite:testdata/multi_add.bin"], - linkopts = common_linkopts, - linkstatic = 1, - tags = [ - "tflite_not_portable_android", - "tflite_not_portable_ios", - ], - deps = [ - ":utils", - "@com_google_googletest//:gtest", - ] + select( - { - "//tensorflow:android": [ - "//tensorflow/core:android_tensorflow_lib", - "//tensorflow/core:android_tensorflow_test_lib", - ], - "//conditions:default": [ - "//tensorflow/core:framework_internal", - "//tensorflow/core:lib", - ], - }, - ), -) - -cc_library( - name = "run_tflite_model_op", - srcs = ["run_tflite_model_op.cc"], - copts = tflite_copts(), - deps = [ - ":utils", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels:builtin_ops", - ] + select( - { - "//tensorflow:android": [ - "//tensorflow/core:android_tensorflow_lib", - ], - "//conditions:default": [ - "//tensorflow/core:tensorflow", - "//tensorflow/core:protos_all_cc", - "//tensorflow/core:core_cpu", - "//tensorflow/core:framework", - "//tensorflow/core:lib", - "//tensorflow/core:ops", - ], - }, - ), - alwayslink = 1, -) - -cc_library( - name = "android_required_build_flags", - srcs = ["android_required_build_flags.cc"], - copts = tflite_copts(), -) - -tf_cc_test( - name = "run_tflite_model_op_test", - srcs = ["run_tflite_model_op_test.cc"], - args = [ - "--test_model_file=$(location //tensorflow/contrib/lite:testdata/multi_add.bin)", - ], - data = ["//tensorflow/contrib/lite:testdata/multi_add.bin"], - linkopts = common_linkopts, - linkstatic = 1, - tags = [ - "tflite_not_portable_android", - "tflite_not_portable_ios", - ], - deps = [ - "//tensorflow/cc:cc_ops", - "//tensorflow/cc:scope", - ":run_tflite_model_op", - ":android_required_build_flags", - "@com_google_googletest//:gtest", - ] + select( - { - "//tensorflow:android": [ - "//tensorflow/core:android_tensorflow_lib", - "//tensorflow/core:android_tensorflow_test_lib", - ], - "//conditions:default": [ - "//tensorflow/core:core_cpu", - "//tensorflow/core:framework", - "//tensorflow/core:framework_internal", - "//tensorflow/core:lib", - "//tensorflow/core:ops", - "//tensorflow/core:protos_all_cc", - "//tensorflow/core:tensorflow", - ], - }, - ), -) - -cc_library( - name = "stage", - hdrs = ["stage.h"], - copts = tflite_copts(), - deps = [ - "//tensorflow/cc:scope", - ], -) - -cc_library( - name = "file_reader_stage", - srcs = ["file_reader_stage.cc"], - hdrs = ["file_reader_stage.h"], - deps = [ - ":stage", - "//tensorflow/cc:cc_ops", - "//tensorflow/cc:scope", - ], -) - -tf_cc_test( - name = "file_reader_stage_test", - srcs = ["file_reader_stage_test.cc"], - linkopts = common_linkopts, - linkstatic = 1, - tags = ["tflite_not_portable_ios"], - deps = [ - ":file_reader_stage", - "@com_google_googletest//:gtest", - ] + select( - { - "//tensorflow:android": [ - "//tensorflow/core:android_tensorflow_lib", - "//tensorflow/core/kernels:android_whole_file_read_ops", - "//tensorflow/core:android_tensorflow_test_lib", - ], - "//conditions:default": [ - "//tensorflow/core:core_cpu", - "//tensorflow/core:tensorflow", - ], - }, - ), -) - -cc_library( - name = "run_tflite_model_stage", - srcs = ["run_tflite_model_stage.cc"], - hdrs = ["run_tflite_model_stage.h"], - copts = tflite_copts(), - deps = [ - ":run_tflite_model_op", - ":stage", - "//tensorflow/cc:cc_ops", - "//tensorflow/cc:scope", - ], -) - -cc_library( - name = "accuracy_eval_stage", - hdrs = ["accuracy_eval_stage.h"], - copts = tflite_copts(), - deps = [ - ] + select( - { - "//tensorflow:android": [ - "//tensorflow/core:android_tensorflow_lib", - ], - "//conditions:default": [ - "//tensorflow/core:framework", - ], - }, - ), -) - -cc_library( - name = "eval_pipeline", - srcs = ["eval_pipeline.cc"], - hdrs = ["eval_pipeline.h"], - copts = tflite_copts(), - deps = [ - ":accuracy_eval_stage", - ":stage", - ] + select( - { - "//tensorflow:android": [ - "//tensorflow/core:android_tensorflow_lib", - ], - "//conditions:default": [ - "//tensorflow/core:framework", - "//tensorflow/core:core_cpu", - ], - }, - ), -) - -tf_cc_test( - name = "eval_pipeline_test", - srcs = ["eval_pipeline_test.cc"], - linkopts = common_linkopts, - linkstatic = 1, - tags = ["tflite_not_portable_ios"], - deps = [ - ":eval_pipeline", - "//tensorflow/cc:cc_ops", - "@com_google_googletest//:gtest", - ] + select( - { - "//tensorflow:android": [ - "//tensorflow/core:android_tensorflow_lib", - "//tensorflow/core:android_tensorflow_test_lib", - ], - "//conditions:default": [ - "//tensorflow/core:framework", - "//tensorflow/core:core_cpu", - "//tensorflow/core:ops", - "//tensorflow/core:tensorflow", - ], - }, - ), -) - -cc_library( - name = "eval_pipeline_builder", - srcs = ["eval_pipeline_builder.cc"], - hdrs = ["eval_pipeline_builder.h"], - copts = tflite_copts(), - deps = [ - ":eval_pipeline", - ":accuracy_eval_stage", - ":stage", - "@com_google_absl//absl/memory", - "//tensorflow/cc:cc_ops", - ] + select( - { - "//tensorflow:android": [ - "//tensorflow/core:android_tensorflow_lib", - ], - "//conditions:default": [ - "//tensorflow/core:framework", - "//tensorflow/core:core_cpu", - "//tensorflow/core:ops", - "//tensorflow/core:tensorflow", - ], - }, - ), -) - -tf_cc_test( - name = "eval_pipeline_builder_test", - srcs = ["eval_pipeline_builder_test.cc"], - linkopts = common_linkopts, - linkstatic = 1, - tags = ["tflite_not_portable_ios"], - deps = [ - ":eval_pipeline_builder", - "//tensorflow/cc:cc_ops", - "@com_google_googletest//:gtest", - ] + select( - { - "//tensorflow:android": [ - "//tensorflow/core:android_tensorflow_lib", - "//tensorflow/core:android_tensorflow_test_lib", - ], - "//conditions:default": [ - "//tensorflow/core:framework", - "//tensorflow/core:core_cpu", - "//tensorflow/core:ops", - "//tensorflow/core:tensorflow", - ], - }, - ), -) - -cc_library( - name = "csv_writer", - hdrs = ["csv_writer.h"], - copts = tflite_copts(), - deps = select( - { - "//tensorflow:android": [ - "//tensorflow/core:android_tensorflow_lib", - ], - "//conditions:default": [ - "//tensorflow/core:lib", - ], - }, - ), -) - -tflite_portable_test_suite() diff --git a/tensorflow/contrib/lite/tools/accuracy/ilsvrc/BUILD b/tensorflow/contrib/lite/tools/accuracy/ilsvrc/BUILD deleted file mode 100644 index 98e2835b2ebd2f7918a939fb89aebec0fd54fb43..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/tools/accuracy/ilsvrc/BUILD +++ /dev/null @@ -1,182 +0,0 @@ -package(default_visibility = [ - "//visibility:public", -]) - -licenses(["notice"]) # Apache 2.0 - -load("//tensorflow:tensorflow.bzl", "tf_cc_binary", "tf_cc_test") -load("//tensorflow/contrib/lite:build_def.bzl", "tflite_copts", "tflite_linkopts") -load("//tensorflow/contrib/lite:special_rules.bzl", "tflite_portable_test_suite") - -common_linkopts = tflite_linkopts() + select({ - "//conditions:default": [], - "//tensorflow:android": [ - "-pie", - "-llog", - ], -}) - -cc_library( - name = "inception_preprocessing", - srcs = ["inception_preprocessing.cc"], - hdrs = ["inception_preprocessing.h"], - copts = tflite_copts(), - deps = [ - "//tensorflow/contrib/lite/tools/accuracy:android_required_build_flags", - "//tensorflow/contrib/lite/tools/accuracy:stage", - "//tensorflow/cc:cc_ops", - "//tensorflow/cc:scope", - ] + select( - { - "//tensorflow:android": [ - "//tensorflow/core:android_tensorflow_lib", - "//tensorflow/core/kernels:android_tensorflow_image_op", - ], - "//conditions:default": [ - "//tensorflow/core:tensorflow", - "//tensorflow/core:protos_all_cc", - "//tensorflow/core:core_cpu", - "//tensorflow/core:framework", - "//tensorflow/core:lib", - "//tensorflow/core:ops", - ], - }, - ), -) - -tf_cc_test( - name = "inception_preprocessing_test", - srcs = ["inception_preprocessing_test.cc"], - args = [ - "--test_image=$(location :testdata/grace_hopper.jpg)", - ], - data = [":testdata/grace_hopper.jpg"], - linkopts = common_linkopts, - linkstatic = 1, - tags = [ - "no_oss", # b/114307765 - "tflite_not_portable_android", - "tflite_not_portable_ios", - ], - deps = [ - ":inception_preprocessing", - "//tensorflow/contrib/lite/tools/accuracy:android_required_build_flags", - "@com_google_googletest//:gtest", - ] + select( - { - "//tensorflow:android": [ - "//tensorflow/core:android_tensorflow_lib", - "//tensorflow/core:android_tensorflow_test_lib", - ], - "//conditions:default": [ - "//tensorflow/core:core_cpu", - "//tensorflow/core:framework_internal", - "//tensorflow/core:lib", - ], - }, - ), -) - -cc_library( - name = "imagenet_topk_eval", - srcs = ["imagenet_topk_eval.cc"], - hdrs = ["imagenet_topk_eval.h"], - copts = tflite_copts(), - deps = [ - "//tensorflow/contrib/lite/tools/accuracy:accuracy_eval_stage", - ] + select( - { - "//tensorflow:android": [ - "//tensorflow/core:android_tensorflow_lib", - ], - "//conditions:default": [ - "//tensorflow/core:framework", - "//tensorflow/core:lib", - ], - }, - ), -) - -tf_cc_test( - name = "imagenet_topk_eval_test", - srcs = ["imagenet_topk_eval_test.cc"], - linkopts = common_linkopts, - linkstatic = 1, - tags = ["tflite_not_portable_ios"], - deps = [ - ":imagenet_topk_eval", - "@com_google_googletest//:gtest", - ] + select( - { - "//tensorflow:android": [ - "//tensorflow/core:android_tensorflow_lib", - "//tensorflow/core:android_tensorflow_test_lib", - ], - "//conditions:default": [ - "//tensorflow/core:framework", - ], - }, - ), -) - -cc_library( - name = "imagenet_model_evaluator", - srcs = ["imagenet_model_evaluator.cc"], - hdrs = ["imagenet_model_evaluator.h"], - copts = tflite_copts(), - deps = [ - ":imagenet_topk_eval", - ":inception_preprocessing", - "//tensorflow/contrib/lite/tools/accuracy:android_required_build_flags", - "//tensorflow/contrib/lite/tools/accuracy:eval_pipeline", - "//tensorflow/contrib/lite/tools/accuracy:eval_pipeline_builder", - "//tensorflow/contrib/lite/tools/accuracy:file_reader_stage", - "//tensorflow/contrib/lite/tools/accuracy:run_tflite_model_stage", - "//tensorflow/contrib/lite/tools/accuracy:utils", - "@com_google_absl//absl/memory", - "//tensorflow/cc:cc_ops", - "//tensorflow/cc:scope", - ] + select( - { - "//tensorflow:android": [ - "//tensorflow/core:android_tensorflow_lib", - "//tensorflow/core/kernels:android_whole_file_read_ops", - "//tensorflow/core/kernels:android_tensorflow_image_op", - ], - "//conditions:default": [ - "//tensorflow/core:tensorflow", - "//tensorflow/core:lib_internal", - "//tensorflow/core:framework_internal", - "//tensorflow/core:framework", - "//tensorflow/core:lib", - "//tensorflow/core:core_cpu", - ], - }, - ), -) - -tf_cc_binary( - name = "imagenet_accuracy_eval", - srcs = ["imagenet_accuracy_eval.cc"], - copts = tflite_copts(), - linkopts = common_linkopts, - deps = [ - ":imagenet_model_evaluator", - ":imagenet_topk_eval", - "@com_google_absl//absl/memory", - "//tensorflow/contrib/lite/tools/accuracy:android_required_build_flags", - "//tensorflow/contrib/lite/tools/accuracy:csv_writer", - ] + select( - { - "//tensorflow:android": [ - "//tensorflow/core:android_tensorflow_lib", - ], - "//conditions:default": [ - "//tensorflow/core:lib", - "//tensorflow/core:framework_internal", - ], - }, - ), -) - -tflite_portable_test_suite() diff --git a/tensorflow/contrib/lite/tools/accuracy/ilsvrc/README.md b/tensorflow/contrib/lite/tools/accuracy/ilsvrc/README.md deleted file mode 100644 index 362ea3ac34f60a93ec242bf11306c5798b982035..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/tools/accuracy/ilsvrc/README.md +++ /dev/null @@ -1,146 +0,0 @@ -## Accuracy evaluation for ILSVRC 2012 (Imagenet Large Scale Visual Recognition Challenge) image classification task - -This binary can evaluate the accuracy of TFLite models trained for the [ILSVRC 2012 image classification task] -(http://www.image-net.org/challenges/LSVRC/2012/). -The binary takes the path to validation images and labels as inputs. It outputs the accuracy after running the TFLite model on the validation sets. - -To run the binary download the ILSVRC 2012 devkit [see instructions](#downloading-ilsvrc) and run the [`generate_validation_ground_truth` script](#ground-truth-label-generation) to generate the ground truth labels. - -## Parameters -The binary takes the following parameters: - -* `model_file` : `string` \ - Path to the TFlite model file. - -* `ground_truth_images_path`: `string` \ - The path to the directory containing ground truth images. - -* `ground_truth_labels`: `string` \ - Path to ground truth labels file. This file should contain the same number of labels as the number images in the ground truth directory. The labels are assumed to be in the - same order as the sorted filename of images. See [ground truth label generation](#ground-truth-label-generation) - section for more information about how to generate labels for images. - -* `model_output_labels`: `string` \ - Path to the file containing labels, that is used to interpret the output of - the model. E.g. in case of mobilenets, this is the path to - `mobilenet_labels.txt` where each label is in the same order as the output - 1001 dimension tensor. - -* `output_path`: `string` \ - This is the path to the output file. The output is a CSV file that has top-10 accuracies in each row. Each line of output file is the cumulative accuracy after processing images in a sorted order. So first line is accuracy after processing the first image, second line is accuracy after procesing first two images. The last line of the file is accuracy after processing the entire validation set. - -and the following optional parameters: - -* `blacklist_file_path`: `string` \ - Path to blacklist file. This file contains the indices of images that are blacklisted for evaluation. 1762 images are blacklisted in ILSVRC dataset. For details please refer to readme.txt of ILSVRC2014 devkit. - -* `num_images`: `int` (default=0) \ - The number of images to process, if 0, all images in the directory are processed otherwise only num_images will be processed. - -* `num_threads`: `int` (default=4) \ - The number of threads to use for evaluation. - - -## Downloading ILSVRC -In order to use this tool to run evaluation on the full 50K ImageNet dataset, -download the data set from http://image-net.org/request. - -## Ground truth label generation -The ILSVRC 2012 devkit `validation_ground_truth.txt` contains IDs that correspond to synset of the image. -The accuracy binary however expects the ground truth labels to contain the actual name of -category instead of synset ids. A conversion script has been provided to convert the validation ground truth to -category labels. The `validation_ground_truth.txt` can be converted by the following steps: - -``` -ILSVRC_2012_DEVKIT_DIR=[set to path to ILSVRC 2012 devkit] -VALIDATION_LABELS=[set to path to output] - -python generate_validation_labels.py -- \ ---ilsvrc_devkit_dir=${ILSVRC_2012_DEVKIT_DIR} \ ---validation_labels_output=${VALIDATION_LABELS} -``` - -## Running the binary - -### On Android - -(0) Refer to https://github.com/tensorflow/tensorflow/tree/master/tensorflow/examples/android for configuring NDK and SDK. - -(1) Build using the following command: - -``` -bazel build -c opt \ - --config=android_arm \ - --config=monolithic \ - --cxxopt='--std=c++11' \ - --copt=-D__ANDROID_TYPES_FULL__ \ - --copt=-DSUPPORT_SELECTIVE_REGISTRATION \ - //tensorflow/contrib/lite/tools/accuracy/ilsvrc:imagenet_accuracy_eval -``` - -(2) Connect your phone. Push the binary to your phone with adb push - (make the directory if required): - -``` -adb push bazel-bin/tensorflow/contrib/lite/tools/accuracy/ilsvrc/imagenet_accuracy_eval /data/local/tmp -``` - -(3) Make the binary executable. - -``` -adb shell chmod +x /data/local/tmp/imagenet_accuracy_eval -``` - -(4) Push the TFLite model that you need to test. For example: - -``` -adb push mobilenet_quant_v1_224.tflite /data/local/tmp -``` - -(5) Push the imagenet images to device, make sure device has sufficient storage available before pushing the dataset: - -``` -adb shell mkdir /data/local/tmp/ilsvrc_images && \ -adb push ${IMAGENET_IMAGES_DIR} /data/local/tmp/ilsvrc_images -``` - -(6) Push the generated validation ground labels to device. - -``` -adb push ${VALIDATION_LABELS} /data/local/tmp/ilsvrc_validation_labels.txt -``` - -(7) Push the model labels text file to device. - -``` -adb push ${MODEL_LABELS_TXT} /data/local/tmp/model_output_labels.txt -``` - -(8) Run the binary. - -``` -adb shell /data/local/tmp/imagenet_accuracy_eval \ - --model_file=/data/local/tmp/mobilenet_quant_v1_224.tflite \ - --ground_truth_images_path=/data/local/tmp/ilsvrc_images \ - --ground_truth_labels=/data/local/tmp/ilsvrc_validation_labels.txt \ - --model_output_labels=/data/local/tmp/model_output_labels.txt \ - --output_file_path=/data/local/tmp/accuracy_output.txt \ - --num_images=0 # Run on all images. -``` - -### On Desktop - -(1) Build and run using the following command: - -``` -bazel run -c opt \ - --cxxopt='--std=c++11' \ - -- \ - //tensorflow/contrib/lite/tools/accuracy/ilsvrc:imagenet_accuracy_eval \ - --model_file=mobilenet_quant_v1_224.tflite \ - --ground_truth_images_path=${IMAGENET_IMAGES_DIR} \ - --ground_truth_labels=${VALIDATION_LABELS} \ - --model_output_labels=${MODEL_LABELS_TXT} \ - --output_file_path=/tmp/accuracy_output.txt \ - --num_images=0 # Run on all images. -``` diff --git a/tensorflow/contrib/lite/tools/accuracy/ilsvrc/inception_preprocessing.cc b/tensorflow/contrib/lite/tools/accuracy/ilsvrc/inception_preprocessing.cc deleted file mode 100644 index 7512b39c32f98faed9b41f829666bf1d4d145d82..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/tools/accuracy/ilsvrc/inception_preprocessing.cc +++ /dev/null @@ -1,80 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/contrib/lite/tools/accuracy/ilsvrc/inception_preprocessing.h" - -#include - -#include "tensorflow/cc/framework/scope.h" -#include "tensorflow/cc/ops/standard_ops.h" -#include "tensorflow/core/framework/graph.pb.h" -#include "tensorflow/core/framework/tensor.h" -#include "tensorflow/core/graph/graph_def_builder.h" -#include "tensorflow/core/lib/core/errors.h" -#include "tensorflow/core/platform/logging.h" -#include "tensorflow/core/platform/types.h" -#include "tensorflow/core/public/session.h" - -namespace tensorflow { -namespace metrics { - -namespace { -void CentralCropImage(const Scope& s, const tensorflow::Output& decoded_image, - double crop_fraction, tensorflow::Output* cropped_image) { - auto image_dims = ops::Slice(s, ops::Shape(s, decoded_image), {0}, {2}); - auto height_width = ops::Cast(s, image_dims, DT_DOUBLE); - auto cropped_begin = ops::Div( - s, ops::Sub(s, height_width, ops::Mul(s, height_width, crop_fraction)), - 2.0); - auto bbox_begin = ops::Cast(s, cropped_begin, DT_INT32); - auto bbox_size = ops::Sub(s, image_dims, ops::Mul(s, bbox_begin, 2)); - auto slice_begin = ops::Concat(s, {bbox_begin, Input({0})}, 0); - auto slice_size = ops::Concat(s, {bbox_size, {-1}}, 0); - *cropped_image = ops::Slice(s, decoded_image, slice_begin, slice_size); -} - -} // namespace - -void InceptionPreprocessingStage::AddToGraph(const Scope& scope, - const Input& input) { - if (!scope.ok()) return; - Scope s = scope.WithOpName(name()); - ops::DecodeJpeg::Attrs attrs; - attrs.channels_ = 3; - auto decoded_jpeg = ops::DecodeJpeg(s, input, attrs); - tensorflow::Output cropped_image; - CentralCropImage(s, decoded_jpeg, params_.cropping_fraction, &cropped_image); - auto dims_expander = ops::ExpandDims(s, cropped_image, 0); - auto resized_image = ops::ResizeBilinear( - s, dims_expander, - ops::Const(s.WithOpName("size"), {image_height_, image_width_})); - if (is_quantized_) { - this->stage_output_ = - ops::Cast(s.WithOpName(output_name()), resized_image, DT_UINT8); - } else { - auto squeezed_image = ops::Squeeze(s, resized_image); - auto normalized_image = - ops::Div(s, - ops::Sub(s, squeezed_image, - {params_.input_means[0], params_.input_means[1], - params_.input_means[2]}), - {params_.scale}); - this->stage_output_ = - ops::ExpandDims(s.WithOpName(output_name()), normalized_image, {0}); - } -} - -} // namespace metrics -} // namespace tensorflow diff --git a/tensorflow/contrib/lite/tools/accuracy/ilsvrc/inception_preprocessing.h b/tensorflow/contrib/lite/tools/accuracy/ilsvrc/inception_preprocessing.h deleted file mode 100644 index 15df71981756f6171b8e12bd9ed2a337c4867b64..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/tools/accuracy/ilsvrc/inception_preprocessing.h +++ /dev/null @@ -1,75 +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_LITE_TOOLS_ACCURACY_INCEPTION_PREPROCESSING_H_ -#define TENSORFLOW_CONTRIB_LITE_TOOLS_ACCURACY_INCEPTION_PREPROCESSING_H_ - -#include - -#include "tensorflow/contrib/lite/tools/accuracy/stage.h" -#include "tensorflow/core/framework/tensor.h" -#include "tensorflow/core/lib/core/status.h" - -namespace tensorflow { -namespace metrics { - -// A stage that does inception preprocessing. -// Inputs: A tensor containing bytes of a JPEG image. -// Outputs: A tensor containing rescaled and preprocessed image that has -// shape {1, image_height, image_width, 3}, where 3 is the number of channels. -class InceptionPreprocessingStage : public Stage { - public: - struct Params { - std::vector input_means; - float scale; - double cropping_fraction; - }; - - static Params DefaultParams() { - return {.input_means = {127.5, 127.5, 127.5}, - .scale = 127.5, - .cropping_fraction = 0.875}; - } - - // Creates a new preprocessing stage object with provided |image_width| - // |image_height| as the size of output image. - // If |is_quantized| is set to true then |params| is ignored since quantized - // images don't go through any preprocessing. - InceptionPreprocessingStage(int image_width, int image_height, - bool is_quantized, - Params params = DefaultParams()) - : image_width_(image_width), - image_height_(image_height), - is_quantized_(is_quantized), - params_(std::move(params)) {} - - string name() const override { return "stage_inception_preprocess"; } - string output_name() const override { - return "stage_inception_preprocess_output"; - } - - void AddToGraph(const Scope& scope, const Input& input) override; - - private: - int image_width_; - int image_height_; - bool is_quantized_; - Params params_; -}; - -} // namespace metrics -} // namespace tensorflow - -#endif // TENSORFLOW_CONTRIB_LITE_TOOLS_ACCURACY_INCEPTION_PREPROCESSING_H_ diff --git a/tensorflow/contrib/lite/tools/accuracy/ilsvrc/inception_preprocessing_test.cc b/tensorflow/contrib/lite/tools/accuracy/ilsvrc/inception_preprocessing_test.cc deleted file mode 100644 index 3587878ba3cadd13eb0af4c004f4f98184daf5de..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/tools/accuracy/ilsvrc/inception_preprocessing_test.cc +++ /dev/null @@ -1,123 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include -#include - -#include -#include "tensorflow/contrib/lite/tools/accuracy/ilsvrc/inception_preprocessing.h" -#include "tensorflow/core/platform/init_main.h" -#include "tensorflow/core/public/session.h" -#include "tensorflow/core/util/command_line_flags.h" - -namespace { -tensorflow::string* g_test_image_file = nullptr; -} // namespace - -namespace tensorflow { -namespace metrics { - -namespace { - -using tensorflow::Status; -using tensorflow::Tensor; - -Status GetContents(const string& filename, string* output) { - std::ifstream input(filename, std::ios::binary); - const int kBufferSize = 2048; - char buffer[kBufferSize]; - while (true) { - input.read(buffer, kBufferSize); - output->append(buffer, input.gcount()); - if (!input.good()) { - if (input.eof()) return Status::OK(); - return Status(tensorflow::error::ABORTED, "Failed to read file."); - } - } -} - -TEST(InceptionPreprocessingTest, TestImagePreprocessQuantized) { - ASSERT_TRUE(g_test_image_file != nullptr); - string image_contents; - string image_path = *g_test_image_file; - auto status = GetContents(image_path, &image_contents); - ASSERT_TRUE(status.ok()) << status.error_message(); - const int width = 224; - const int height = 224; - const bool is_quantized = true; - InceptionPreprocessingStage preprocess_stage(width, height, is_quantized); - Scope scope = Scope::NewRootScope(); - preprocess_stage.AddToGraph(scope, image_contents); - TF_CHECK_OK(scope.status()); - - GraphDef graph_def; - TF_CHECK_OK(scope.ToGraphDef(&graph_def)); - std::unique_ptr session(NewSession(SessionOptions())); - TF_CHECK_OK(session->Create(graph_def)); - std::vector outputs; - auto run_status = - session->Run({}, /*inputs*/ - {preprocess_stage.output_name()}, {}, /*target node names */ - &outputs); - TF_CHECK_OK(run_status); - EXPECT_EQ(1, outputs.size()); - EXPECT_EQ(DT_UINT8, outputs[0].dtype()); - EXPECT_TRUE(outputs[0].shape().IsSameSize({1, 224, 224, 3})); -} - -TEST(InceptionPreprocessingTest, TestImagePreprocessFloat) { - ASSERT_TRUE(g_test_image_file != nullptr); - string image_contents; - string image_path = *g_test_image_file; - auto status = GetContents(image_path, &image_contents); - ASSERT_TRUE(status.ok()) << status.error_message(); - const int width = 224; - const int height = 224; - const bool is_quantized = false; - InceptionPreprocessingStage preprocess_stage(width, height, is_quantized); - Scope scope = Scope::NewRootScope(); - preprocess_stage.AddToGraph(scope, image_contents); - TF_CHECK_OK(scope.status()); - - GraphDef graph_def; - TF_CHECK_OK(scope.ToGraphDef(&graph_def)); - std::unique_ptr session(NewSession(SessionOptions())); - TF_CHECK_OK(session->Create(graph_def)); - std::vector outputs; - auto run_status = - session->Run({}, /*inputs*/ - {preprocess_stage.output_name()}, {}, /*target node names */ - &outputs); - TF_CHECK_OK(run_status); - EXPECT_EQ(1, outputs.size()); - EXPECT_EQ(DT_FLOAT, outputs[0].dtype()); - EXPECT_TRUE(outputs[0].shape().IsSameSize({1, 224, 224, 3})); -} - -} // namespace -} // namespace metrics -} // namespace tensorflow - -int main(int argc, char** argv) { - g_test_image_file = new tensorflow::string(); - const std::vector flag_list = { - tensorflow::Flag("test_image", g_test_image_file, - "Path to image file for test."), - }; - const bool parse_result = tensorflow::Flags::Parse(&argc, argv, flag_list); - CHECK(parse_result) << "Required test_model_file"; - ::tensorflow::port::InitMain(argv[0], &argc, &argv); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/contrib/lite/tools/accuracy/utils.cc b/tensorflow/contrib/lite/tools/accuracy/utils.cc deleted file mode 100644 index f5493301fc4d781418cc5c7397bae02ecc155c56..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/tools/accuracy/utils.cc +++ /dev/null @@ -1,102 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/contrib/lite/tools/accuracy/utils.h" - -#include - -#include -#include -#include -#include - -#include "tensorflow/contrib/lite/interpreter.h" -#include "tensorflow/contrib/lite/kernels/register.h" -#include "tensorflow/contrib/lite/model.h" -#include "tensorflow/contrib/lite/op_resolver.h" - -namespace tensorflow { -namespace metrics { - -namespace utils { - -DataType GetTFDataType(TfLiteType tflite_type) { - switch (tflite_type) { - case kTfLiteFloat32: - return DT_FLOAT; - case kTfLiteUInt8: - return DT_UINT8; - default: - return DT_INVALID; - } -} - -TensorShape GetTFLiteTensorShape(const TfLiteTensor& tflite_tensor) { - TensorShape shape; - for (int i = 0; i < tflite_tensor.dims->size; i++) { - shape.AddDim(tflite_tensor.dims->data[i]); - } - return shape; -} - -Status ReadFileLines(const string& file_path, - std::vector* lines_output) { - if (!lines_output) { - return errors::InvalidArgument("Invalid output"); - } - std::vector lines; - std::ifstream stream(file_path, std::ios_base::in); - if (!stream) { - return errors::InvalidArgument("Unable to open file: ", file_path); - } - std::string line; - while (std::getline(stream, line)) { - lines_output->push_back(line); - } - return Status::OK(); -} - -Status GetTFliteModelInfo(const string& model_file_path, - ModelInfo* model_info) { - if (model_file_path.empty()) { - return errors::InvalidArgument("Invalid model file."); - } - struct stat stat_buf; - if (stat(model_file_path.c_str(), &stat_buf) != 0) { - int error_num = errno; - return errors::InvalidArgument("Invalid model file: ", model_file_path, - std::strerror(error_num)); - } - - std::unique_ptr model; - std::unique_ptr interpreter; - model = tflite::FlatBufferModel::BuildFromFile(model_file_path.data()); - tflite::ops::builtin::BuiltinOpResolver resolver; - - tflite::InterpreterBuilder(*model, resolver)(&interpreter); - if (!interpreter) { - return errors::InvalidArgument("Invalid model", model_file_path); - } - for (int i : interpreter->inputs()) { - TfLiteTensor* tensor = interpreter->tensor(i); - model_info->input_shapes.push_back(utils::GetTFLiteTensorShape(*tensor)); - model_info->input_types.push_back(utils::GetTFDataType(tensor->type)); - } - return Status::OK(); -} - -} // namespace utils -} // namespace metrics -} // namespace tensorflow diff --git a/tensorflow/contrib/lite/tools/accuracy/utils.h b/tensorflow/contrib/lite/tools/accuracy/utils.h deleted file mode 100644 index 37cbad4d51fd0ddf700b14ead037ae4aeed4d82a..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/tools/accuracy/utils.h +++ /dev/null @@ -1,46 +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_LITE_TOOLS_ACCURACY_UTILS_H_ -#define TENSORFLOW_CONTRIB_LITE_TOOLS_ACCURACY_UTILS_H_ - -#include -#include - -#include "tensorflow/contrib/lite/context.h" -#include "tensorflow/core/framework/tensor_shape.h" - -namespace tensorflow { -namespace metrics { - -namespace utils { - -struct ModelInfo { - std::vector input_shapes; - std::vector input_types; -}; - -Status GetTFliteModelInfo(const string& model_file_path, ModelInfo* model_info); - -DataType GetTFDataType(TfLiteType tflite_type); - -TensorShape GetTFLiteTensorShape(const TfLiteTensor& tflite_tensor); - -Status ReadFileLines(const string& file_path, - std::vector* lines_output); -} // namespace utils -} // namespace metrics -} // namespace tensorflow -#endif // TENSORFLOW_CONTRIB_LITE_TOOLS_ACCURACY_UTILS_H_ diff --git a/tensorflow/contrib/lite/tools/benchmark/BUILD b/tensorflow/contrib/lite/tools/benchmark/BUILD deleted file mode 100644 index a3822f4215a2c746af0b1671ab9945d47c21dc53..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/tools/benchmark/BUILD +++ /dev/null @@ -1,146 +0,0 @@ -package(default_visibility = [ - "//visibility:public", -]) - -licenses(["notice"]) # Apache 2.0 - -load("//tensorflow/contrib/lite:special_rules.bzl", "tflite_portable_test_suite") -load("//tensorflow/contrib/lite:build_def.bzl", "tflite_copts") -load("//tensorflow/contrib/lite:build_def.bzl", "tflite_linkopts") - -common_copts = ["-Wall"] + tflite_copts() - -cc_library( - name = "logging", - hdrs = ["logging.h"], - copts = common_copts, -) - -cc_binary( - name = "benchmark_model", - srcs = [ - "benchmark_main.cc", - ], - copts = common_copts, - linkopts = tflite_linkopts() + select({ - "//tensorflow:android": [ - "-pie", # Android 5.0 and later supports only PIE - "-lm", # some builtin ops, e.g., tanh, need -lm - ], - "//conditions:default": [], - }), - deps = [ - ":benchmark_tflite_model_lib", - ":logging", - ], -) - -cc_binary( - name = "benchmark_model_plus_flex", - srcs = [ - "benchmark_plus_flex_main.cc", - ], - copts = common_copts, - linkopts = tflite_linkopts() + select({ - "//tensorflow:android": [ - "-pie", # Android 5.0 and later supports only PIE - "-lm", # some builtin ops, e.g., tanh, need -lm - ], - "//conditions:default": [], - }), - deps = [ - ":benchmark_tflite_model_lib", - ":logging", - "//tensorflow/contrib/lite/delegates/flex:delegate", - "//tensorflow/contrib/lite/testing:init_tensorflow", - ], -) - -cc_test( - name = "benchmark_test", - srcs = ["benchmark_test.cc"], - args = [ - "--graph=$(location //tensorflow/contrib/lite:testdata/multi_add.bin)", - ], - data = ["//tensorflow/contrib/lite:testdata/multi_add.bin"], - tags = [ - "tflite_not_portable_android", - "tflite_not_portable_ios", - ], - deps = [ - ":benchmark_tflite_model_lib", - ":command_line_flags", - "//tensorflow/contrib/lite/testing:util", - "@com_google_googletest//:gtest", - ], -) - -cc_library( - name = "command_line_flags", - srcs = ["command_line_flags.cc"], - hdrs = ["command_line_flags.h"], - copts = common_copts, -) - -cc_test( - name = "command_line_flags_test", - srcs = ["command_line_flags_test.cc"], - copts = common_copts, - visibility = ["//visibility:private"], - deps = [ - ":command_line_flags", - "//tensorflow/contrib/lite/testing:util", - "@com_google_googletest//:gtest", - ], -) - -cc_library( - name = "benchmark_tflite_model_lib", - srcs = [ - "benchmark_tflite_model.cc", - "logging.h", - ], - hdrs = ["benchmark_tflite_model.h"], - copts = common_copts, - deps = [ - ":benchmark_model_lib", - ":logging", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite:string_util", - "//tensorflow/contrib/lite/kernels:builtin_ops", - "//tensorflow/contrib/lite/profiling:profile_summarizer", - ], -) - -cc_library( - name = "benchmark_params", - srcs = [ - "benchmark_params.cc", - ], - hdrs = ["benchmark_params.h"], - copts = common_copts, - deps = [":logging"], -) - -cc_library( - name = "benchmark_model_lib", - srcs = [ - "benchmark_model.cc", - ], - hdrs = ["benchmark_model.h"], - copts = common_copts, - deps = [ - ":benchmark_params", - ":command_line_flags", - ":logging", - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite:string_util", - "//tensorflow/contrib/lite/kernels:builtin_ops", - "//tensorflow/contrib/lite/profiling:profile_summarizer", - "//tensorflow/contrib/lite/profiling:profiler", - "//tensorflow/contrib/lite/profiling:time", - "//tensorflow/core:stats_calculator_portable", - ], -) - -tflite_portable_test_suite() diff --git a/tensorflow/contrib/lite/tools/benchmark/README.md b/tensorflow/contrib/lite/tools/benchmark/README.md deleted file mode 100644 index 8d997639fb7a363f911b1183dfb05d8138e4c531..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/tools/benchmark/README.md +++ /dev/null @@ -1,197 +0,0 @@ -# TFLite Model Benchmark Tool - -## Description - -A simple C++ binary to benchmark a TFLite model and its individual operators, -both on desktop machines and on Android. The binary takes a TFLite model, -generates random inputs and then repeatedly runs the model for specified number -of runs. Aggregrate latency statistics are reported after running the benchmark. - -The instructions below are for running the binary on Desktop and Android, -for iOS please use the -[iOS benchmark app](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite/tools/benchmark/ios). - -## Parameters - -The binary takes the following required parameters: - -* `graph`: `string` \ - The path to the TFLite model file. - -and the following optional parameters: - -* `num_threads`: `int` (default=1) \ - The number of threads to use for running TFLite interpreter. -* `warmup_runs`: `int` (default=1) \ - The number of warmup runs to do before starting the benchmark. -* `num_runs`: `int` (default=50) \ - The number of runs. Increase this to reduce variance. -* `run_delay`: `float` (default=-1.0) \ - The delay in seconds between subsequent benchmark runs. Non-positive values - mean use no delay. -* `use_nnapi`: `bool` (default=false) \ - Whether to use [Android NNAPI](https://developer.android.com/ndk/guides/neuralnetworks/). - This API is available on recent Android devices. - -## To build/install/run - -### On Android: - -(0) Refer to https://github.com/tensorflow/tensorflow/tree/master/tensorflow/examples/android to edit the `WORKSPACE` to configure the android NDK/SDK. - -(1) Build for your specific platform, e.g.: - -``` -bazel build -c opt \ - --config=android_arm \ - --cxxopt='--std=c++11' \ - tensorflow/contrib/lite/tools/benchmark:benchmark_model -``` - -(2) Connect your phone. Push the binary to your phone with adb push - (make the directory if required): - -``` -adb push bazel-bin/tensorflow/contrib/lite/tools/benchmark/benchmark_model /data/local/tmp -``` - -(3) Make the binary executable. - -``` -adb shell chmod +x /data/local/tmp/benchmark_model -``` - -(4) Push the compute graph that you need to test. For example: - -``` -adb push mobilenet_quant_v1_224.tflite /data/local/tmp -``` - -(5) Run the benchmark. For example: - -``` -adb shell /data/local/tmp/benchmark_model \ - --graph=/data/local/tmp/mobilenet_quant_v1_224.tflite \ - --num_threads=4 -``` - -### On desktop: -(1) build the binary - -``` -bazel build -c opt tensorflow/contrib/lite/tools/benchmark:benchmark_model -``` - -(2) Run on your compute graph, similar to the Android case but without the need of adb shell. -For example: - -``` -bazel-bin/tensorflow/contrib/lite/tools/benchmark/benchmark_model \ - --graph=mobilenet_quant_v1_224.tflite \ - --num_threads=4 -``` - -The MobileNet graph used as an example here may be downloaded from [here](https://storage.googleapis.com/download.tensorflow.org/models/tflite/mobilenet_v1_224_android_quant_2017_11_08.zip). - - -## Reducing variance between runs on Android. - -Most modern Android phones use [ARM big.LITTLE](https://en.wikipedia.org/wiki/ARM_big.LITTLE) -architecture where some cores are more power hungry but faster than other cores. -When running benchmarks on these phones there can be significant variance -between different runs of the benchmark. One way to reduce variance between runs -is to set the [CPU affinity](https://en.wikipedia.org/wiki/Processor_affinity) -before running the benchmark. On Android this can be done using the `taskset` -command. -E.g. for running the benchmark on big cores on Pixel 2 with a single thread one -can use the following command: - -``` -adb shell taskset f0 /data/local/tmp/benchmark_model \ - --graph=/data/local/tmp/mobilenet_quant_v1_224.tflite \ - --num_threads=1 -``` - -where `f0` is the affinity mask for big cores on Pixel 2. -Note: The affinity mask varies with the device. - -## Profiling model operators -The benchmark model binary also allows you to profile operators and give execution times of each operator. To do this, -compile the binary with a compiler flag that enables profiling to be compiled in. Pass **--copt=-DTFLITE_PROFILING_ENABLED** -to compile benchmark with profiling support. -For example, to compile with profiling support on Android, add this flag to the previous command: - -``` -bazel build -c opt \ - --config=android_arm \ - --cxxopt='--std=c++11' \ - --copt=-DTFLITE_PROFILING_ENABLED \ - tensorflow/contrib/lite/tools/benchmark:benchmark_model -``` -This compiles TFLite with profiling enabled, now you can run the benchmark binary like before. The binary will produce detailed statistics for each operation similar to those shown below: - -``` - -============================== Run Order ============================== - [node type] [start] [first] [avg ms] [%] [cdf%] [mem KB] [times called] [Name] - CONV_2D 0.000 4.269 4.269 0.107% 0.107% 0.000 0 [MobilenetV1/MobilenetV1/Conv2d_0/Relu6] - DEPTHWISE_CONV_2D 4.270 2.150 2.150 0.054% 0.161% 0.000 0 [MobilenetV1/MobilenetV1/Conv2d_1_depthwise/Relu6] - CONV_2D 6.421 6.107 6.107 0.153% 0.314% 0.000 0 [MobilenetV1/MobilenetV1/Conv2d_1_pointwise/Relu6] - DEPTHWISE_CONV_2D 12.528 1.366 1.366 0.034% 0.348% 0.000 0 [MobilenetV1/MobilenetV1/Conv2d_2_depthwise/Relu6] - CONV_2D 13.895 4.195 4.195 0.105% 0.454% 0.000 0 [MobilenetV1/MobilenetV1/Conv2d_2_pointwise/Relu6] - DEPTHWISE_CONV_2D 18.091 1.260 1.260 0.032% 0.485% 0.000 0 [MobilenetV1/MobilenetV1/Conv2d_3_depthwise/Relu6] - CONV_2D 19.352 6.652 6.652 0.167% 0.652% 0.000 0 [MobilenetV1/MobilenetV1/Conv2d_3_pointwise/Relu6] - DEPTHWISE_CONV_2D 26.005 0.698 0.698 0.018% 0.670% 0.000 0 [MobilenetV1/MobilenetV1/Conv2d_4_depthwise/Relu6] - CONV_2D 26.703 3.344 3.344 0.084% 0.754% 0.000 0 [MobilenetV1/MobilenetV1/Conv2d_4_pointwise/Relu6] - DEPTHWISE_CONV_2D 30.047 0.646 0.646 0.016% 0.770% 0.000 0 [MobilenetV1/MobilenetV1/Conv2d_5_depthwise/Relu6] - CONV_2D 30.694 5.800 5.800 0.145% 0.915% 0.000 0 [MobilenetV1/MobilenetV1/Conv2d_5_pointwise/Relu6] - DEPTHWISE_CONV_2D 36.495 0.331 0.331 0.008% 0.924% 0.000 0 [MobilenetV1/MobilenetV1/Conv2d_6_depthwise/Relu6] - CONV_2D 36.826 2.838 2.838 0.071% 0.995% 0.000 0 [MobilenetV1/MobilenetV1/Conv2d_6_pointwise/Relu6] - DEPTHWISE_CONV_2D 39.665 0.439 0.439 0.011% 1.006% 0.000 0 [MobilenetV1/MobilenetV1/Conv2d_7_depthwise/Relu6] - CONV_2D 40.105 5.293 5.293 0.133% 1.139% 0.000 0 [MobilenetV1/MobilenetV1/Conv2d_7_pointwise/Relu6] - DEPTHWISE_CONV_2D 45.399 0.352 0.352 0.009% 1.147% 0.000 0 [MobilenetV1/MobilenetV1/Conv2d_8_depthwise/Relu6] - CONV_2D 45.752 5.322 5.322 0.133% 1.281% 0.000 0 [MobilenetV1/MobilenetV1/Conv2d_8_pointwise/Relu6] - DEPTHWISE_CONV_2D 51.075 0.357 0.357 0.009% 1.290% 0.000 0 [MobilenetV1/MobilenetV1/Conv2d_9_depthwise/Relu6] - CONV_2D 51.432 5.693 5.693 0.143% 1.433% 0.000 0 [MobilenetV1/MobilenetV1/Conv2d_9_pointwise/Relu6] - DEPTHWISE_CONV_2D 57.126 0.366 0.366 0.009% 1.442% 0.000 0 [MobilenetV1/MobilenetV1/Conv2d_10_depthwise/Relu6] - CONV_2D 57.493 5.472 5.472 0.137% 1.579% 0.000 0 [MobilenetV1/MobilenetV1/Conv2d_10_pointwise/Relu6] - DEPTHWISE_CONV_2D 62.966 0.364 0.364 0.009% 1.588% 0.000 0 [MobilenetV1/MobilenetV1/Conv2d_11_depthwise/Relu6] - CONV_2D 63.330 5.404 5.404 0.136% 1.724% 0.000 0 [MobilenetV1/MobilenetV1/Conv2d_11_pointwise/Relu6] - DEPTHWISE_CONV_2D 68.735 0.155 0.155 0.004% 1.728% 0.000 0 [MobilenetV1/MobilenetV1/Conv2d_12_depthwise/Relu6] - CONV_2D 68.891 2.970 2.970 0.074% 1.802% 0.000 0 [MobilenetV1/MobilenetV1/Conv2d_12_pointwise/Relu6] - DEPTHWISE_CONV_2D 71.862 0.206 0.206 0.005% 1.807% 0.000 0 [MobilenetV1/MobilenetV1/Conv2d_13_depthwise/Relu6] - CONV_2D 72.069 5.888 5.888 0.148% 1.955% 0.000 0 [MobilenetV1/MobilenetV1/Conv2d_13_pointwise/Relu6] - AVERAGE_POOL_2D 77.958 0.036 0.036 0.001% 1.956% 0.000 0 [MobilenetV1/Logits/AvgPool_1a/AvgPool] - CONV_2D 77.994 1.445 1.445 0.036% 1.992% 0.000 0 [MobilenetV1/Logits/Conv2d_1c_1x1/BiasAdd] - RESHAPE 79.440 0.002 0.002 0.000% 1.992% 0.000 0 [MobilenetV1/Predictions/Reshape] - SOFTMAX 79.443 0.029 0.029 0.001% 1.993% 0.000 0 [MobilenetV1/Predictions/Softmax] - -============================== Top by Computation Time ============================== - [node type] [start] [first] [avg ms] [%] [cdf%] [mem KB] [times called] [Name] - CONV_2D 19.352 6.652 6.652 0.167% 0.167% 0.000 0 [MobilenetV1/MobilenetV1/Conv2d_3_pointwise/Relu6] - CONV_2D 6.421 6.107 6.107 0.153% 0.320% 0.000 0 [MobilenetV1/MobilenetV1/Conv2d_1_pointwise/Relu6] - CONV_2D 72.069 5.888 5.888 0.148% 0.468% 0.000 0 [MobilenetV1/MobilenetV1/Conv2d_13_pointwise/Relu6] - CONV_2D 30.694 5.800 5.800 0.145% 0.613% 0.000 0 [MobilenetV1/MobilenetV1/Conv2d_5_pointwise/Relu6] - CONV_2D 51.432 5.693 5.693 0.143% 0.756% 0.000 0 [MobilenetV1/MobilenetV1/Conv2d_9_pointwise/Relu6] - CONV_2D 57.493 5.472 5.472 0.137% 0.893% 0.000 0 [MobilenetV1/MobilenetV1/Conv2d_10_pointwise/Relu6] - CONV_2D 63.330 5.404 5.404 0.136% 1.029% 0.000 0 [MobilenetV1/MobilenetV1/Conv2d_11_pointwise/Relu6] - CONV_2D 45.752 5.322 5.322 0.133% 1.162% 0.000 0 [MobilenetV1/MobilenetV1/Conv2d_8_pointwise/Relu6] - CONV_2D 40.105 5.293 5.293 0.133% 1.295% 0.000 0 [MobilenetV1/MobilenetV1/Conv2d_7_pointwise/Relu6] - CONV_2D 0.000 4.269 4.269 0.107% 1.402% 0.000 0 [MobilenetV1/MobilenetV1/Conv2d_0/Relu6] - -Number of nodes executed: 31 -============================== Summary by node type ============================== - [Node type] [count] [avg ms] [avg %] [cdf %] [mem KB] [times called] - CONV_2D 15 1.406 89.270% 89.270% 0.000 0 - DEPTHWISE_CONV_2D 13 0.169 10.730% 100.000% 0.000 0 - SOFTMAX 1 0.000 0.000% 100.000% 0.000 0 - RESHAPE 1 0.000 0.000% 100.000% 0.000 0 - AVERAGE_POOL_2D 1 0.000 0.000% 100.000% 0.000 0 - -Timings (microseconds): count=50 first=79449 curr=81350 min=77385 max=88213 avg=79732 std=1929 -Memory (bytes): count=0 -31 nodes observed - - -Average inference timings in us: Warmup: 83235, Init: 38467, no stats: 79760.9 -``` diff --git a/tensorflow/contrib/lite/tools/benchmark/benchmark_model.cc b/tensorflow/contrib/lite/tools/benchmark/benchmark_model.cc deleted file mode 100644 index f86c0445b0525cd053c733b18bb7f1205d310d43..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/tools/benchmark/benchmark_model.cc +++ /dev/null @@ -1,168 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/contrib/lite/tools/benchmark/benchmark_model.h" - -#include - -#include -#include - -#include "tensorflow/contrib/lite/profiling/time.h" -#include "tensorflow/contrib/lite/tools/benchmark/logging.h" - -namespace { -void SleepForSeconds(double sleep_seconds) { - if (sleep_seconds <= 0.0) { - return; - } - // Convert the run_delay string into a timespec. - timespec req; - req.tv_sec = static_cast(sleep_seconds); - req.tv_nsec = (sleep_seconds - req.tv_sec) * 1000000000; - // If requested, sleep between runs for an arbitrary amount of time. - // This can be helpful to determine the effect of mobile processor - // scaling and thermal throttling. -#ifdef PLATFORM_WINDOWS - Sleep(sleep_seconds * 1000); -#else - nanosleep(&req, nullptr); -#endif -} - -} // namespace - -namespace tflite { -namespace benchmark { -using tensorflow::Stat; - -BenchmarkParams BenchmarkModel::DefaultParams() { - BenchmarkParams params; - params.AddParam("num_runs", BenchmarkParam::Create(50)); - params.AddParam("run_delay", BenchmarkParam::Create(-1.0f)); - params.AddParam("num_threads", BenchmarkParam::Create(1)); - params.AddParam("benchmark_name", BenchmarkParam::Create("")); - params.AddParam("output_prefix", BenchmarkParam::Create("")); - params.AddParam("warmup_runs", BenchmarkParam::Create(1)); - return params; -} - -BenchmarkModel::BenchmarkModel() : params_(DefaultParams()) {} - -void BenchmarkLoggingListener::OnBenchmarkEnd(const BenchmarkResults &results) { - auto inference_us = results.inference_time_us(); - auto init_us = results.startup_latency_us(); - auto warmup_us = results.warmup_time_us(); - TFLITE_LOG(INFO) << "Average inference timings in us: " - << "Warmup: " << warmup_us.avg() << ", " - << "Init: " << init_us << ", " - << "no stats: " << inference_us.avg(); -} - -std::vector BenchmarkModel::GetFlags() { - return { - CreateFlag("num_runs", ¶ms_, "number of runs"), - CreateFlag("run_delay", ¶ms_, "delay between runs in seconds"), - CreateFlag("num_threads", ¶ms_, "number of threads"), - CreateFlag("benchmark_name", ¶ms_, "benchmark name"), - CreateFlag("output_prefix", ¶ms_, - "benchmark output prefix"), - CreateFlag("warmup_runs", ¶ms_, - "how many runs to initialize model"), - }; -} - -void BenchmarkModel::LogParams() { - TFLITE_LOG(INFO) << "Num runs: [" << params_.Get("num_runs") << "]"; - TFLITE_LOG(INFO) << "Inter-run delay (seconds): [" - << params_.Get("run_delay") << "]"; - TFLITE_LOG(INFO) << "Num threads: [" << params_.Get("num_threads") - << "]"; - TFLITE_LOG(INFO) << "Benchmark name: [" - << params_.Get("benchmark_name") << "]"; - TFLITE_LOG(INFO) << "Output prefix: [" - << params_.Get("output_prefix") << "]"; - TFLITE_LOG(INFO) << "Warmup runs: [" << params_.Get("warmup_runs") - << "]"; -} - -void BenchmarkModel::PrepareInputsAndOutputs() {} - -Stat BenchmarkModel::Run(int num_times, RunType run_type) { - Stat run_stats; - TFLITE_LOG(INFO) << "Running benchmark for " << num_times << " iterations "; - for (int run = 0; run < num_times; run++) { - PrepareInputsAndOutputs(); - listeners_.OnSingleRunStart(run_type); - int64_t start_us = profiling::time::NowMicros(); - RunImpl(); - int64_t end_us = profiling::time::NowMicros(); - listeners_.OnSingleRunEnd(); - - run_stats.UpdateStat(end_us - start_us); - SleepForSeconds(params_.Get("run_delay")); - } - - std::stringstream stream; - run_stats.OutputToStream(&stream); - TFLITE_LOG(INFO) << stream.str() << std::endl; - - return run_stats; -} - -bool BenchmarkModel::ValidateParams() { return true; } - -void BenchmarkModel::Run(int argc, char **argv) { - if (!ParseFlags(argc, argv)) { - return; - } - Run(); -} - -void BenchmarkModel::Run() { - ValidateParams(); - LogParams(); - - listeners_.OnBenchmarkStart(params_); - int64_t initialization_start_us = profiling::time::NowMicros(); - Init(); - int64_t initialization_end_us = profiling::time::NowMicros(); - int64_t startup_latency_us = initialization_end_us - initialization_start_us; - TFLITE_LOG(INFO) << "Initialized session in " << startup_latency_us / 1e3 - << "ms"; - - uint64_t input_bytes = ComputeInputBytes(); - Stat warmup_time_us = - Run(params_.Get("warmup_runs"), WARMUP); - Stat inference_time_us = - Run(params_.Get("num_runs"), REGULAR); - listeners_.OnBenchmarkEnd( - {startup_latency_us, input_bytes, warmup_time_us, inference_time_us}); -} - -bool BenchmarkModel::ParseFlags(int argc, char **argv) { - auto flag_list = GetFlags(); - const bool parse_result = - Flags::Parse(&argc, const_cast(argv), flag_list); - if (!parse_result) { - std::string usage = Flags::Usage(argv[0], flag_list); - TFLITE_LOG(ERROR) << usage; - return false; - } - return true; -} - -} // namespace benchmark -} // namespace tflite diff --git a/tensorflow/contrib/lite/tools/benchmark/benchmark_test.cc b/tensorflow/contrib/lite/tools/benchmark/benchmark_test.cc deleted file mode 100644 index b697bb394db9b967dfaaff649517dcc23e85ccb0..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/tools/benchmark/benchmark_test.cc +++ /dev/null @@ -1,74 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include -#include -#include - -#include -#include -#include "tensorflow/contrib/lite/testing/util.h" -#include "tensorflow/contrib/lite/tools/benchmark/benchmark_tflite_model.h" -#include "tensorflow/contrib/lite/tools/benchmark/command_line_flags.h" - -namespace { -const std::string* g_model_path = nullptr; -} - -namespace tflite { -namespace benchmark { -namespace { - -BenchmarkParams CreateParams() { - BenchmarkParams params; - params.AddParam("num_runs", BenchmarkParam::Create(2)); - params.AddParam("run_delay", BenchmarkParam::Create(-1.0f)); - params.AddParam("num_threads", BenchmarkParam::Create(1)); - params.AddParam("benchmark_name", BenchmarkParam::Create("")); - params.AddParam("output_prefix", BenchmarkParam::Create("")); - params.AddParam("warmup_runs", BenchmarkParam::Create(1)); - params.AddParam("graph", BenchmarkParam::Create(*g_model_path)); - params.AddParam("input_layer", BenchmarkParam::Create("")); - params.AddParam("input_layer_shape", BenchmarkParam::Create("")); - params.AddParam("use_nnapi", BenchmarkParam::Create(false)); - return params; -} - -TEST(BenchmarkTest, DoesntCrash) { - ASSERT_THAT(g_model_path, testing::NotNull()); - - BenchmarkTfLiteModel benchmark(CreateParams()); - benchmark.Run(); -} - -} // namespace -} // namespace benchmark -} // namespace tflite - -int main(int argc, char** argv) { - std::string model_path; - std::vector flags = { - tflite::Flag::CreateFlag("graph", &model_path, "Path to model file.")}; - g_model_path = &model_path; - const bool parse_result = - tflite::Flags::Parse(&argc, const_cast(argv), flags); - if (!parse_result) { - std::cerr << tflite::Flags::Usage(argv[0], flags); - return 1; - } - - ::tflite::LogToStderr(); - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/contrib/lite/tools/benchmark/benchmark_tflite_model.h b/tensorflow/contrib/lite/tools/benchmark/benchmark_tflite_model.h deleted file mode 100644 index 25a302b2aaea400ea66e76dae3e6add71180e1cc..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/tools/benchmark/benchmark_tflite_model.h +++ /dev/null @@ -1,82 +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_LITE_TOOLS_BENCHMARK_BENCHMARK_TFLITE_MODEL_H_ -#define TENSORFLOW_CONTRIB_LITE_TOOLS_BENCHMARK_BENCHMARK_TFLITE_MODEL_H_ - -#include -#include -#include - -#include "tensorflow/contrib/lite/model.h" -#include "tensorflow/contrib/lite/profiling/profile_summarizer.h" -#include "tensorflow/contrib/lite/tools/benchmark/benchmark_model.h" - -namespace tflite { -namespace benchmark { - -// Dumps profiling events if profiling is enabled -class ProfilingListener : public BenchmarkListener { - public: - explicit ProfilingListener() : interpreter_(nullptr), has_profiles_(false) {} - - void SetInterpreter(Interpreter* interpreter); - - void OnSingleRunStart(RunType run_type) override; - - void OnSingleRunEnd() override; - - void OnBenchmarkEnd(const BenchmarkResults& results) override; - - private: - Interpreter* interpreter_; - profiling::Profiler profiler_; - profiling::ProfileSummarizer summarizer_; - bool has_profiles_; -}; - -// Benchmarks a TFLite model by running tflite interpreter. -class BenchmarkTfLiteModel : public BenchmarkModel { - public: - BenchmarkTfLiteModel(); - BenchmarkTfLiteModel(BenchmarkParams params); - virtual ~BenchmarkTfLiteModel() {} - - std::vector GetFlags() override; - void LogParams() override; - bool ValidateParams() override; - uint64_t ComputeInputBytes() override; - void Init() override; - void RunImpl() override; - - struct InputLayerInfo { - std::string name; - std::vector shape; - }; - - protected: - void PrepareInputsAndOutputs() override; - - private: - std::unique_ptr model; - std::unique_ptr interpreter; - std::vector inputs; - ProfilingListener profiling_listener_; -}; - -} // namespace benchmark -} // namespace tflite - -#endif // TENSORFLOW_CONTRIB_LITE_TOOLS_BENCHMARK_BENCHMARK_TFLITE_MODEL_H_ diff --git a/tensorflow/contrib/lite/tools/benchmark/ios/README.md b/tensorflow/contrib/lite/tools/benchmark/ios/README.md deleted file mode 100644 index 46144f7bf8e142b960d3fe1068686e366bb6c198..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/tools/benchmark/ios/README.md +++ /dev/null @@ -1,43 +0,0 @@ -# TFLite iOS benchmark app. - -## Description - -An iOS app to benchmark TFLite models. - -The app reads benchmark parameters from a JSON file named `benchmark_params.json` -in its `benchmark_data` directory. Any downloaded models for benchmarking should -also be placed in `benchmark_data` directory. - -The JSON file specifies the name of the model file and other benchmarking -parameters like inputs to the model, type of inputs, number of iterations, -number of threads. The default values in the JSON file are for the -Mobilenet_1.0_224 model -([paper](https://arxiv.org/pdf/1704.04861.pdf), -[tflite&pb](http://download.tensorflow.org/models/mobilenet_v1_2018_02_22/mobilenet_v1_1.0_224.tgz)) - -## To build/install/run - -- Follow instructions at -[iOS build for TFLite](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/lite/g3doc/ios.md) -to build TFLite. - -Running - -```bash -tensorflow/contrib/lite/build_ios_universal_lib.sh -``` -will also build `tensorflow/contrib/lite/gen/lib/benchmark-lib.a` . - -- Now copy the downloaded model file to `benchmark_data` directory. - -- Modify `benchmark_params.json` change the `input_layer`, `input_layer_shape` -and other benchmark parameters. - -- Change `Build Phases -> Copy Bundle Resources` and add the model file to the -resources that need to be copied. - -- Ensure that `Build Phases -> Link Binary With Library` contains the -`Accelerate framework` and `tensorflow/contrib/lite/gen/lib/benchmark-lib.a`. - -- Now try running the app. The app has a single button that runs the benchmark - on the model and displays results in a text view below. diff --git a/tensorflow/contrib/lite/tools/benchmark/ios/TFLiteBenchmark/TFLiteBenchmark.xcodeproj/project.pbxproj b/tensorflow/contrib/lite/tools/benchmark/ios/TFLiteBenchmark/TFLiteBenchmark.xcodeproj/project.pbxproj deleted file mode 100644 index b908f733d49b56a6b41ebea4185f1fe8c11edc60..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/tools/benchmark/ios/TFLiteBenchmark/TFLiteBenchmark.xcodeproj/project.pbxproj +++ /dev/null @@ -1,381 +0,0 @@ -// !$*UTF8*$! -{ - archiveVersion = 1; - classes = { - }; - objectVersion = 50; - objects = { - -/* Begin PBXBuildFile section */ - 6FE7579A20D59CE500F01636 /* benchmark_params.json in Resources */ = {isa = PBXBuildFile; fileRef = 6FE7579920D59CE500F01636 /* benchmark_params.json */; }; - 6FE7579D20D5A5E000F01636 /* benchmark-lib.a in Frameworks */ = {isa = PBXBuildFile; fileRef = 6FE7579C20D5A5E000F01636 /* benchmark-lib.a */; }; - 6FE7579F20D5A6A700F01636 /* Accelerate.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 6FE7579E20D5A6A700F01636 /* Accelerate.framework */; }; - 6FE757A120D5AB8100F01636 /* mobilenet_v1_1.0_224.tflite in Resources */ = {isa = PBXBuildFile; fileRef = 6FE757A020D5AB8000F01636 /* mobilenet_v1_1.0_224.tflite */; }; - 6FE93FFD20D592D8008C9FE4 /* AppDelegate.m in Sources */ = {isa = PBXBuildFile; fileRef = 6FE93FFC20D592D8008C9FE4 /* AppDelegate.m */; }; - 6FE9400020D592D8008C9FE4 /* BenchmarkViewController.mm in Sources */ = {isa = PBXBuildFile; fileRef = 6FE93FFF20D592D8008C9FE4 /* BenchmarkViewController.mm */; }; - 6FE9400320D592D8008C9FE4 /* Main.storyboard in Resources */ = {isa = PBXBuildFile; fileRef = 6FE9400120D592D8008C9FE4 /* Main.storyboard */; }; - 6FE9400520D592DA008C9FE4 /* Assets.xcassets in Resources */ = {isa = PBXBuildFile; fileRef = 6FE9400420D592DA008C9FE4 /* Assets.xcassets */; }; - 6FE9400B20D592DA008C9FE4 /* main.m in Sources */ = {isa = PBXBuildFile; fileRef = 6FE9400A20D592DA008C9FE4 /* main.m */; }; -/* End PBXBuildFile section */ - -/* Begin PBXFileReference section */ - 6FE7579920D59CE500F01636 /* benchmark_params.json */ = {isa = PBXFileReference; lastKnownFileType = text.json; path = benchmark_params.json; sourceTree = ""; }; - 6FE7579C20D5A5E000F01636 /* benchmark-lib.a */ = {isa = PBXFileReference; lastKnownFileType = archive.ar; name = "benchmark-lib.a"; path = "$SRCROOT/../../../../../../../tensorflow/contrib/lite/gen/lib/benchmark-lib.a"; sourceTree = ""; }; - 6FE7579E20D5A6A700F01636 /* Accelerate.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = Accelerate.framework; path = System/Library/Frameworks/Accelerate.framework; sourceTree = SDKROOT; }; - 6FE757A020D5AB8000F01636 /* mobilenet_v1_1.0_224.tflite */ = {isa = PBXFileReference; lastKnownFileType = file; path = mobilenet_v1_1.0_224.tflite; sourceTree = ""; }; - 6FE93FF820D592D8008C9FE4 /* TFLiteBenchmark.app */ = {isa = PBXFileReference; explicitFileType = wrapper.application; includeInIndex = 0; path = TFLiteBenchmark.app; sourceTree = BUILT_PRODUCTS_DIR; }; - 6FE93FFB20D592D8008C9FE4 /* AppDelegate.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = AppDelegate.h; sourceTree = ""; }; - 6FE93FFC20D592D8008C9FE4 /* AppDelegate.m */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.objc; path = AppDelegate.m; sourceTree = ""; }; - 6FE93FFE20D592D8008C9FE4 /* BenchmarkViewController.h */ = {isa = PBXFileReference; explicitFileType = sourcecode.cpp.h; path = BenchmarkViewController.h; sourceTree = ""; }; - 6FE93FFF20D592D8008C9FE4 /* BenchmarkViewController.mm */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.cpp.objcpp; path = BenchmarkViewController.mm; sourceTree = ""; }; - 6FE9400220D592D8008C9FE4 /* Base */ = {isa = PBXFileReference; lastKnownFileType = file.storyboard; name = Base; path = Base.lproj/Main.storyboard; sourceTree = ""; }; - 6FE9400420D592DA008C9FE4 /* Assets.xcassets */ = {isa = PBXFileReference; lastKnownFileType = folder.assetcatalog; path = Assets.xcassets; sourceTree = ""; }; - 6FE9400920D592DA008C9FE4 /* Info.plist */ = {isa = PBXFileReference; lastKnownFileType = text.plist.xml; path = Info.plist; sourceTree = ""; }; - 6FE9400A20D592DA008C9FE4 /* main.m */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.objc; path = main.m; sourceTree = ""; }; -/* End PBXFileReference section */ - -/* Begin PBXFrameworksBuildPhase section */ - 6FE93FF520D592D8008C9FE4 /* Frameworks */ = { - isa = PBXFrameworksBuildPhase; - buildActionMask = 2147483647; - files = ( - 6FE7579F20D5A6A700F01636 /* Accelerate.framework in Frameworks */, - 6FE7579D20D5A5E000F01636 /* benchmark-lib.a in Frameworks */, - ); - runOnlyForDeploymentPostprocessing = 0; - }; -/* End PBXFrameworksBuildPhase section */ - -/* Begin PBXGroup section */ - 6FE7579820D59C8B00F01636 /* benchmark_data */ = { - isa = PBXGroup; - children = ( - 6FE757A020D5AB8000F01636 /* mobilenet_v1_1.0_224.tflite */, - 6FE7579920D59CE500F01636 /* benchmark_params.json */, - ); - path = benchmark_data; - sourceTree = ""; - }; - 6FE7579B20D5A5E000F01636 /* Frameworks */ = { - isa = PBXGroup; - children = ( - 6FE7579E20D5A6A700F01636 /* Accelerate.framework */, - 6FE7579C20D5A5E000F01636 /* benchmark-lib.a */, - ); - name = Frameworks; - sourceTree = ""; - }; - 6FE93FEF20D592D8008C9FE4 = { - isa = PBXGroup; - children = ( - 6FE93FFA20D592D8008C9FE4 /* TFLiteBenchmark */, - 6FE93FF920D592D8008C9FE4 /* Products */, - 6FE7579B20D5A5E000F01636 /* Frameworks */, - ); - sourceTree = ""; - }; - 6FE93FF920D592D8008C9FE4 /* Products */ = { - isa = PBXGroup; - children = ( - 6FE93FF820D592D8008C9FE4 /* TFLiteBenchmark.app */, - ); - name = Products; - sourceTree = ""; - }; - 6FE93FFA20D592D8008C9FE4 /* TFLiteBenchmark */ = { - isa = PBXGroup; - children = ( - 6FE7579820D59C8B00F01636 /* benchmark_data */, - 6FE93FFB20D592D8008C9FE4 /* AppDelegate.h */, - 6FE93FFC20D592D8008C9FE4 /* AppDelegate.m */, - 6FE93FFE20D592D8008C9FE4 /* BenchmarkViewController.h */, - 6FE93FFF20D592D8008C9FE4 /* BenchmarkViewController.mm */, - 6FE9400120D592D8008C9FE4 /* Main.storyboard */, - 6FE9400420D592DA008C9FE4 /* Assets.xcassets */, - 6FE9400920D592DA008C9FE4 /* Info.plist */, - 6FE9400A20D592DA008C9FE4 /* main.m */, - ); - path = TFLiteBenchmark; - sourceTree = ""; - }; -/* End PBXGroup section */ - -/* Begin PBXNativeTarget section */ - 6FE93FF720D592D8008C9FE4 /* TFLiteBenchmark */ = { - isa = PBXNativeTarget; - buildConfigurationList = 6FE9400E20D592DA008C9FE4 /* Build configuration list for PBXNativeTarget "TFLiteBenchmark" */; - buildPhases = ( - 6FE93FF420D592D8008C9FE4 /* Sources */, - 6FE93FF520D592D8008C9FE4 /* Frameworks */, - 6FE93FF620D592D8008C9FE4 /* Resources */, - ); - buildRules = ( - ); - dependencies = ( - ); - name = TFLiteBenchmark; - productName = TFLiteBenchmark; - productReference = 6FE93FF820D592D8008C9FE4 /* TFLiteBenchmark.app */; - productType = "com.apple.product-type.application"; - }; -/* End PBXNativeTarget section */ - -/* Begin PBXProject section */ - 6FE93FF020D592D8008C9FE4 /* Project object */ = { - isa = PBXProject; - attributes = { - LastUpgradeCheck = 1000; - ORGANIZATIONNAME = Example; - TargetAttributes = { - 6FE93FF720D592D8008C9FE4 = { - CreatedOnToolsVersion = 10.0; - }; - }; - }; - buildConfigurationList = 6FE93FF320D592D8008C9FE4 /* Build configuration list for PBXProject "TFLiteBenchmark" */; - compatibilityVersion = "Xcode 9.3"; - developmentRegion = en; - hasScannedForEncodings = 0; - knownRegions = ( - en, - Base, - ); - mainGroup = 6FE93FEF20D592D8008C9FE4; - productRefGroup = 6FE93FF920D592D8008C9FE4 /* Products */; - projectDirPath = ""; - projectRoot = ""; - targets = ( - 6FE93FF720D592D8008C9FE4 /* TFLiteBenchmark */, - ); - }; -/* End PBXProject section */ - -/* Begin PBXResourcesBuildPhase section */ - 6FE93FF620D592D8008C9FE4 /* Resources */ = { - isa = PBXResourcesBuildPhase; - buildActionMask = 2147483647; - files = ( - 6FE757A120D5AB8100F01636 /* mobilenet_v1_1.0_224.tflite in Resources */, - 6FE9400520D592DA008C9FE4 /* Assets.xcassets in Resources */, - 6FE9400320D592D8008C9FE4 /* Main.storyboard in Resources */, - 6FE7579A20D59CE500F01636 /* benchmark_params.json in Resources */, - ); - runOnlyForDeploymentPostprocessing = 0; - }; -/* End PBXResourcesBuildPhase section */ - -/* Begin PBXSourcesBuildPhase section */ - 6FE93FF420D592D8008C9FE4 /* Sources */ = { - isa = PBXSourcesBuildPhase; - buildActionMask = 2147483647; - files = ( - 6FE9400020D592D8008C9FE4 /* BenchmarkViewController.mm in Sources */, - 6FE9400B20D592DA008C9FE4 /* main.m in Sources */, - 6FE93FFD20D592D8008C9FE4 /* AppDelegate.m in Sources */, - ); - runOnlyForDeploymentPostprocessing = 0; - }; -/* End PBXSourcesBuildPhase section */ - -/* Begin PBXVariantGroup section */ - 6FE9400120D592D8008C9FE4 /* Main.storyboard */ = { - isa = PBXVariantGroup; - children = ( - 6FE9400220D592D8008C9FE4 /* Base */, - ); - name = Main.storyboard; - sourceTree = ""; - }; -/* End PBXVariantGroup section */ - -/* Begin XCBuildConfiguration section */ - 6FE9400C20D592DA008C9FE4 /* Debug */ = { - isa = XCBuildConfiguration; - buildSettings = { - ALWAYS_SEARCH_USER_PATHS = NO; - CLANG_ANALYZER_NONNULL = YES; - CLANG_ANALYZER_NUMBER_OBJECT_CONVERSION = YES_AGGRESSIVE; - CLANG_CXX_LANGUAGE_STANDARD = "gnu++14"; - CLANG_CXX_LIBRARY = "libc++"; - CLANG_ENABLE_MODULES = YES; - CLANG_ENABLE_OBJC_ARC = YES; - CLANG_ENABLE_OBJC_WEAK = YES; - CLANG_WARN_BLOCK_CAPTURE_AUTORELEASING = YES; - CLANG_WARN_BOOL_CONVERSION = YES; - CLANG_WARN_COMMA = YES; - CLANG_WARN_CONSTANT_CONVERSION = YES; - CLANG_WARN_DEPRECATED_OBJC_IMPLEMENTATIONS = YES; - CLANG_WARN_DIRECT_OBJC_ISA_USAGE = YES_ERROR; - CLANG_WARN_DOCUMENTATION_COMMENTS = YES; - CLANG_WARN_EMPTY_BODY = YES; - CLANG_WARN_ENUM_CONVERSION = YES; - CLANG_WARN_INFINITE_RECURSION = YES; - CLANG_WARN_INT_CONVERSION = YES; - CLANG_WARN_NON_LITERAL_NULL_CONVERSION = YES; - CLANG_WARN_OBJC_IMPLICIT_RETAIN_SELF = YES; - CLANG_WARN_OBJC_LITERAL_CONVERSION = YES; - CLANG_WARN_OBJC_ROOT_CLASS = YES_ERROR; - CLANG_WARN_RANGE_LOOP_ANALYSIS = YES; - CLANG_WARN_STRICT_PROTOTYPES = YES; - CLANG_WARN_SUSPICIOUS_MOVE = YES; - CLANG_WARN_UNGUARDED_AVAILABILITY = YES_AGGRESSIVE; - CLANG_WARN_UNREACHABLE_CODE = YES; - CLANG_WARN__DUPLICATE_METHOD_MATCH = YES; - CODE_SIGN_IDENTITY = "iPhone Developer"; - COPY_PHASE_STRIP = NO; - DEBUG_INFORMATION_FORMAT = dwarf; - ENABLE_STRICT_OBJC_MSGSEND = YES; - ENABLE_TESTABILITY = YES; - GCC_C_LANGUAGE_STANDARD = gnu11; - GCC_DYNAMIC_NO_PIC = NO; - GCC_NO_COMMON_BLOCKS = YES; - GCC_OPTIMIZATION_LEVEL = 0; - GCC_PREPROCESSOR_DEFINITIONS = ( - "DEBUG=1", - "$(inherited)", - ); - GCC_WARN_64_TO_32_BIT_CONVERSION = YES; - GCC_WARN_ABOUT_RETURN_TYPE = YES_ERROR; - GCC_WARN_UNDECLARED_SELECTOR = YES; - GCC_WARN_UNINITIALIZED_AUTOS = YES_AGGRESSIVE; - GCC_WARN_UNUSED_FUNCTION = YES; - GCC_WARN_UNUSED_VARIABLE = YES; - IPHONEOS_DEPLOYMENT_TARGET = 11.0; - MTL_ENABLE_DEBUG_INFO = INCLUDE_SOURCE; - ONLY_ACTIVE_ARCH = YES; - OTHER_CFLAGS = ""; - OTHER_CPLUSPLUSFLAGS = "$(OTHER_CFLAGS)"; - SDKROOT = iphoneos; - }; - name = Debug; - }; - 6FE9400D20D592DA008C9FE4 /* Release */ = { - isa = XCBuildConfiguration; - buildSettings = { - ALWAYS_SEARCH_USER_PATHS = NO; - CLANG_ANALYZER_NONNULL = YES; - CLANG_ANALYZER_NUMBER_OBJECT_CONVERSION = YES_AGGRESSIVE; - CLANG_CXX_LANGUAGE_STANDARD = "gnu++14"; - CLANG_CXX_LIBRARY = "libc++"; - CLANG_ENABLE_MODULES = YES; - CLANG_ENABLE_OBJC_ARC = YES; - CLANG_ENABLE_OBJC_WEAK = YES; - CLANG_WARN_BLOCK_CAPTURE_AUTORELEASING = YES; - CLANG_WARN_BOOL_CONVERSION = YES; - CLANG_WARN_COMMA = YES; - CLANG_WARN_CONSTANT_CONVERSION = YES; - CLANG_WARN_DEPRECATED_OBJC_IMPLEMENTATIONS = YES; - CLANG_WARN_DIRECT_OBJC_ISA_USAGE = YES_ERROR; - CLANG_WARN_DOCUMENTATION_COMMENTS = YES; - CLANG_WARN_EMPTY_BODY = YES; - CLANG_WARN_ENUM_CONVERSION = YES; - CLANG_WARN_INFINITE_RECURSION = YES; - CLANG_WARN_INT_CONVERSION = YES; - CLANG_WARN_NON_LITERAL_NULL_CONVERSION = YES; - CLANG_WARN_OBJC_IMPLICIT_RETAIN_SELF = YES; - CLANG_WARN_OBJC_LITERAL_CONVERSION = YES; - CLANG_WARN_OBJC_ROOT_CLASS = YES_ERROR; - CLANG_WARN_RANGE_LOOP_ANALYSIS = YES; - CLANG_WARN_STRICT_PROTOTYPES = YES; - CLANG_WARN_SUSPICIOUS_MOVE = YES; - CLANG_WARN_UNGUARDED_AVAILABILITY = YES_AGGRESSIVE; - CLANG_WARN_UNREACHABLE_CODE = YES; - CLANG_WARN__DUPLICATE_METHOD_MATCH = YES; - CODE_SIGN_IDENTITY = "iPhone Developer"; - COPY_PHASE_STRIP = NO; - DEBUG_INFORMATION_FORMAT = "dwarf-with-dsym"; - ENABLE_NS_ASSERTIONS = NO; - ENABLE_STRICT_OBJC_MSGSEND = YES; - GCC_C_LANGUAGE_STANDARD = gnu11; - GCC_NO_COMMON_BLOCKS = YES; - GCC_WARN_64_TO_32_BIT_CONVERSION = YES; - GCC_WARN_ABOUT_RETURN_TYPE = YES_ERROR; - GCC_WARN_UNDECLARED_SELECTOR = YES; - GCC_WARN_UNINITIALIZED_AUTOS = YES_AGGRESSIVE; - GCC_WARN_UNUSED_FUNCTION = YES; - GCC_WARN_UNUSED_VARIABLE = YES; - IPHONEOS_DEPLOYMENT_TARGET = 11.0; - MTL_ENABLE_DEBUG_INFO = NO; - OTHER_CFLAGS = ""; - OTHER_CPLUSPLUSFLAGS = "$(OTHER_CFLAGS)"; - SDKROOT = iphoneos; - VALIDATE_PRODUCT = YES; - }; - name = Release; - }; - 6FE9400F20D592DA008C9FE4 /* Debug */ = { - isa = XCBuildConfiguration; - buildSettings = { - ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon; - CODE_SIGN_STYLE = Automatic; - "HEADER_SEARCH_PATHS[arch=*]" = ( - $SRCROOT/../../../../../../../, - $SRCROOT/../../../../../../../tensorflow/contrib/lite/downloads/eigen, - $SRCROOT/../../../../../../../tensorflow/contrib/lite/downloads/gemmlowp, - $SRCROOT/../../../../../../../tensorflow/contrib/lite/downloads/neon_2_sse, - $SRCROOT/../../../../../../../tensorflow/contrib/lite/downloads/farmhash/src, - $SRCROOT/../../../../../../../tensorflow/contrib/lite/downloads/flatbuffers/include, - ); - INFOPLIST_FILE = TFLiteBenchmark/Info.plist; - LD_RUNPATH_SEARCH_PATHS = ( - "$(inherited)", - "@executable_path/Frameworks", - ); - "LIBRARY_SEARCH_PATHS[arch=*]" = $SRCROOT/../../../../../../../tensorflow/contrib/lite/gen/lib; - PRODUCT_BUNDLE_IDENTIFIER = example.TFLiteBenchmark; - PRODUCT_NAME = "$(TARGET_NAME)"; - TARGETED_DEVICE_FAMILY = "1,2"; - "USER_HEADER_SEARCH_PATHS[arch=*]" = ""; - }; - name = Debug; - }; - 6FE9401020D592DA008C9FE4 /* Release */ = { - isa = XCBuildConfiguration; - buildSettings = { - ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon; - CODE_SIGN_STYLE = Automatic; - "HEADER_SEARCH_PATHS[arch=*]" = ( - $SRCROOT/../../../../../../../, - $SRCROOT/../../../../../../../tensorflow/contrib/lite/downloads/eigen, - $SRCROOT/../../../../../../../tensorflow/contrib/lite/downloads/gemmlowp, - $SRCROOT/../../../../../../../tensorflow/contrib/lite/downloads/neon_2_sse, - $SRCROOT/../../../../../../../tensorflow/contrib/lite/downloads/farmhash/src, - $SRCROOT/../../../../../../../tensorflow/contrib/lite/downloads/flatbuffers/include, - ); - INFOPLIST_FILE = TFLiteBenchmark/Info.plist; - LD_RUNPATH_SEARCH_PATHS = ( - "$(inherited)", - "@executable_path/Frameworks", - ); - "LIBRARY_SEARCH_PATHS[arch=*]" = $SRCROOT/../../../../../../../tensorflow/contrib/lite/gen/lib; - PRODUCT_BUNDLE_IDENTIFIER = example.TFLiteBenchmark; - PRODUCT_NAME = "$(TARGET_NAME)"; - TARGETED_DEVICE_FAMILY = "1,2"; - }; - name = Release; - }; -/* End XCBuildConfiguration section */ - -/* Begin XCConfigurationList section */ - 6FE93FF320D592D8008C9FE4 /* Build configuration list for PBXProject "TFLiteBenchmark" */ = { - isa = XCConfigurationList; - buildConfigurations = ( - 6FE9400C20D592DA008C9FE4 /* Debug */, - 6FE9400D20D592DA008C9FE4 /* Release */, - ); - defaultConfigurationIsVisible = 0; - defaultConfigurationName = Release; - }; - 6FE9400E20D592DA008C9FE4 /* Build configuration list for PBXNativeTarget "TFLiteBenchmark" */ = { - isa = XCConfigurationList; - buildConfigurations = ( - 6FE9400F20D592DA008C9FE4 /* Debug */, - 6FE9401020D592DA008C9FE4 /* Release */, - ); - defaultConfigurationIsVisible = 0; - defaultConfigurationName = Release; - }; -/* End XCConfigurationList section */ - }; - rootObject = 6FE93FF020D592D8008C9FE4 /* Project object */; -} diff --git a/tensorflow/contrib/lite/tools/make/Makefile b/tensorflow/contrib/lite/tools/make/Makefile deleted file mode 100644 index 16012a3fb16398003eb6cc934e6fe91318b2849a..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/tools/make/Makefile +++ /dev/null @@ -1,225 +0,0 @@ -# Find where we're running from, so we can store generated files here. -ifeq ($(origin MAKEFILE_DIR), undefined) - MAKEFILE_DIR := $(shell dirname $(realpath $(lastword $(MAKEFILE_LIST)))) -endif - -# Try to figure out the host system -HOST_OS := -ifeq ($(OS),Windows_NT) - HOST_OS = windows -else - UNAME_S := $(shell uname -s) - ifeq ($(UNAME_S),Linux) - HOST_OS := linux - endif - ifeq ($(UNAME_S),Darwin) - HOST_OS := osx - endif -endif - -HOST_ARCH := $(shell if [[ $(shell uname -m) =~ i[345678]86 ]]; then echo x86_32; else echo $(shell uname -m); fi) - -# Override these on the make command line to target a specific architecture. For example: -# make -f tensorflow/contrib/lite/Makefile TARGET=rpi TARGET_ARCH=armv7l -TARGET := $(HOST_OS) -TARGET_ARCH := $(HOST_ARCH) - -INCLUDES := \ --I. \ --I$(MAKEFILE_DIR)/../../../../../ \ --I$(MAKEFILE_DIR)/../../../../../../ \ --I$(MAKEFILE_DIR)/downloads/ \ --I$(MAKEFILE_DIR)/downloads/eigen \ --I$(MAKEFILE_DIR)/downloads/absl \ --I$(MAKEFILE_DIR)/downloads/gemmlowp \ --I$(MAKEFILE_DIR)/downloads/neon_2_sse \ --I$(MAKEFILE_DIR)/downloads/farmhash/src \ --I$(MAKEFILE_DIR)/downloads/flatbuffers/include \ --I$(OBJDIR) -# This is at the end so any globally-installed frameworks like protobuf don't -# override local versions in the source tree. -INCLUDES += -I/usr/local/include - -# These are the default libraries needed, but they can be added to or -# overridden by the platform-specific settings in target makefiles. -LIBS := \ --lstdc++ \ --lpthread \ --lm \ --lz - -# There are no rules for compiling objects for the host system (since we don't -# generate things like the protobuf compiler that require that), so all of -# these settings are for the target compiler. -CXXFLAGS := -O3 -DNDEBUG -CCFLAGS := ${CXXFLAGS} -CXXFLAGS += --std=c++11 -CFLAGS := -LDOPTS := -L/usr/local/lib -ARFLAGS := -r -TARGET_TOOLCHAIN_PREFIX := -CC_PREFIX := - -# This library is the main target for this makefile. It will contain a minimal -# runtime that can be linked in to other programs. -LIB_NAME := libtensorflow-lite.a - -# Benchmark static library and binary -BENCHMARK_LIB_NAME := benchmark-lib.a -BENCHMARK_BINARY_NAME := benchmark_model - -# A small example program that shows how to link against the library. -MINIMAL_SRCS := \ -tensorflow/contrib/lite/examples/minimal/minimal.cc - -# What sources we want to compile, must be kept in sync with the main Bazel -# build files. - -PROFILER_SRCS := \ - tensorflow/contrib/lite/profiling/time.cc -PROFILE_SUMMARIZER_SRCS := \ - tensorflow/contrib/lite/profiling/profile_summarizer.cc \ - tensorflow/core/util/stats_calculator.cc - -CORE_CC_ALL_SRCS := \ -$(wildcard tensorflow/contrib/lite/*.cc) \ -$(wildcard tensorflow/contrib/lite/*.c) \ -$(wildcard tensorflow/contrib/lite/c/*.c) \ -$(wildcard tensorflow/contrib/lite/core/api/*.cc) -ifneq ($(BUILD_TYPE),micro) -CORE_CC_ALL_SRCS += \ -$(wildcard tensorflow/contrib/lite/kernels/*.cc) \ -$(wildcard tensorflow/contrib/lite/kernels/internal/*.cc) \ -$(wildcard tensorflow/contrib/lite/kernels/internal/optimized/*.cc) \ -$(wildcard tensorflow/contrib/lite/kernels/internal/reference/*.cc) \ -$(PROFILER_SRCS) \ -$(wildcard tensorflow/contrib/lite/kernels/*.c) \ -$(wildcard tensorflow/contrib/lite/kernels/internal/*.c) \ -$(wildcard tensorflow/contrib/lite/kernels/internal/optimized/*.c) \ -$(wildcard tensorflow/contrib/lite/kernels/internal/reference/*.c) \ -$(wildcard tensorflow/contrib/lite/tools/make/downloads/farmhash/src/farmhash.cc) \ -$(wildcard tensorflow/contrib/lite/tools/make/downloads/fft2d/fftsg.c) -endif -# Remove any duplicates. -CORE_CC_ALL_SRCS := $(sort $(CORE_CC_ALL_SRCS)) -CORE_CC_EXCLUDE_SRCS := \ -$(wildcard tensorflow/contrib/lite/*test.cc) \ -$(wildcard tensorflow/contrib/lite/*/*test.cc) \ -$(wildcard tensorflow/contrib/lite/*/*/*test.cc) \ -$(wildcard tensorflow/contrib/lite/*/*/*/*test.cc) \ -$(wildcard tensorflow/contrib/lite/kernels/test_util.cc) \ -$(MINIMAL_SRCS) -ifeq ($(BUILD_TYPE),micro) -CORE_CC_EXCLUDE_SRCS += \ -tensorflow/contrib/lite/mmap_allocation.cc \ -tensorflow/contrib/lite/nnapi_delegate.cc -endif -# Filter out all the excluded files. -TF_LITE_CC_SRCS := $(filter-out $(CORE_CC_EXCLUDE_SRCS), $(CORE_CC_ALL_SRCS)) - -# Benchmark sources -BENCHMARK_SRCS_DIR := tensorflow/contrib/lite/tools/benchmark -BENCHMARK_ALL_SRCS := $(TFLITE_CC_SRCS) \ - $(wildcard $(BENCHMARK_SRCS_DIR)/*.cc) \ - $(PROFILE_SUMMARIZER_SRCS) - -BENCHMARK_SRCS := $(filter-out \ - $(wildcard $(BENCHMARK_SRCS_DIR)/*_test.cc), \ - $(BENCHMARK_ALL_SRCS)) - -# These target-specific makefiles should modify or replace options like -# CXXFLAGS or LIBS to work for a specific targetted architecture. All logic -# based on platforms or architectures should happen within these files, to -# keep this main makefile focused on the sources and dependencies. -include $(wildcard $(MAKEFILE_DIR)/targets/*_makefile.inc) - -ALL_SRCS := \ - $(MINIMAL_SRCS) \ - $(PROFILER_SRCS) \ - $(PROFILER_SUMMARY_SRCS) \ - $(TF_LITE_CC_SRCS) \ - $(BENCHMARK_SRCS) - -# Where compiled objects are stored. -GENDIR := $(MAKEFILE_DIR)/gen/$(TARGET)_$(TARGET_ARCH)/ -OBJDIR := $(GENDIR)obj/ -BINDIR := $(GENDIR)bin/ -LIBDIR := $(GENDIR)lib/ - -LIB_PATH := $(LIBDIR)$(LIB_NAME) -BENCHMARK_LIB := $(LIBDIR)$(BENCHMARK_LIB_NAME) -BENCHMARK_BINARY := $(BINDIR)$(BENCHMARK_BINARY_NAME) -MINIMAL_BINARY := $(BINDIR)minimal - -CXX := $(CC_PREFIX)${TARGET_TOOLCHAIN_PREFIX}g++ -CC := $(CC_PREFIX)${TARGET_TOOLCHAIN_PREFIX}gcc -AR := $(CC_PREFIX)${TARGET_TOOLCHAIN_PREFIX}ar - -MINIMAL_OBJS := $(addprefix $(OBJDIR), \ -$(patsubst %.cc,%.o,$(patsubst %.c,%.o,$(MINIMAL_SRCS)))) - -LIB_OBJS := $(addprefix $(OBJDIR), \ -$(patsubst %.cc,%.o,$(patsubst %.c,%.o,$(TF_LITE_CC_SRCS)))) - -BENCHMARK_OBJS := $(addprefix $(OBJDIR), \ -$(patsubst %.cc,%.o,$(patsubst %.c,%.o,$(BENCHMARK_SRCS)))) - -# For normal manually-created TensorFlow C++ source files. -$(OBJDIR)%.o: %.cc - @mkdir -p $(dir $@) - $(CXX) $(CXXFLAGS) $(INCLUDES) -c $< -o $@ -# For normal manually-created TensorFlow C++ source files. -$(OBJDIR)%.o: %.c - @mkdir -p $(dir $@) - $(CC) $(CCFLAGS) $(INCLUDES) -c $< -o $@ - -# The target that's compiled if there's no command-line arguments. -all: $(LIB_PATH) $(MINIMAL_BINARY) $(BENCHMARK_BINARY) - -# The target that's compiled for micro-controllers -micro: $(LIB_PATH) - -# Hack for generating schema file bypassing flatbuffer parsing -tensorflow/contrib/lite/schema/schema_generated.h: - @cp -u tensorflow/contrib/lite/schema/schema_generated.h.OPENSOURCE tensorflow/contrib/lite/schema/schema_generated.h - -# Gathers together all the objects we've compiled into a single '.a' archive. -$(LIB_PATH): tensorflow/contrib/lite/schema/schema_generated.h $(LIB_OBJS) - @mkdir -p $(dir $@) - $(AR) $(ARFLAGS) $(LIB_PATH) $(LIB_OBJS) - -$(MINIMAL_BINARY): $(MINIMAL_OBJS) $(LIB_PATH) - @mkdir -p $(dir $@) - $(CXX) $(CXXFLAGS) $(INCLUDES) \ - -o $(MINIMAL_BINARY) $(MINIMAL_OBJS) \ - $(LIBFLAGS) $(LIB_PATH) $(LDFLAGS) $(LIBS) - -$(BENCHMARK_LIB) : $(LIB_PATH) $(BENCHMARK_OBJS) - @mkdir -p $(dir $@) - $(AR) $(ARFLAGS) $(BENCHMARK_LIB) $(LIB_OBJS) $(BENCHMARK_OBJS) - -benchmark_lib: $(BENCHMARK_LIB) - -$(BENCHMARK_BINARY) : $(BENCHMARK_LIB) - @mkdir -p $(dir $@) - $(CXX) $(CXXFLAGS) $(INCLUDES) \ - -o $(BENCHMARK_BINARY) \ - $(LIBFLAGS) $(BENCHMARK_LIB) $(LDFLAGS) $(LIBS) - -benchmark: $(BENCHMARK_BINARY) - -# Gets rid of all generated files. -clean: - rm -rf $(MAKEFILE_DIR)/gen - -# Gets rid of target files only, leaving the host alone. Also leaves the lib -# directory untouched deliberately, so we can persist multiple architectures -# across builds for iOS and Android. -cleantarget: - rm -rf $(OBJDIR) - rm -rf $(BINDIR) - -$(DEPDIR)/%.d: ; -.PRECIOUS: $(DEPDIR)/%.d - --include $(patsubst %,$(DEPDIR)/%.d,$(basename $(ALL_SRCS))) diff --git a/tensorflow/contrib/lite/tools/make/build_ios_universal_lib.sh b/tensorflow/contrib/lite/tools/make/build_ios_universal_lib.sh deleted file mode 100755 index fe056945a652b04d078947f58bfe6ab60aa1f387..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/tools/make/build_ios_universal_lib.sh +++ /dev/null @@ -1,40 +0,0 @@ -#!/bin/bash -x -# 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. -# ============================================================================== - -set -e - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -cd "$SCRIPT_DIR/../../../../.." - -# Build library for supported architectures and packs them in a fat binary. -make_library() { - for arch in x86_64 armv7 armv7s arm64 - do - make -f tensorflow/contrib/lite/tools/make/Makefile TARGET=ios TARGET_ARCH=${arch} \ - -j 8 - done - mkdir -p tensorflow/contrib/lite/tools/make/gen/lib - lipo \ - tensorflow/contrib/lite/tools/make/gen/ios_x86_64/lib/${1} \ - tensorflow/contrib/lite/tools/make/gen/ios_armv7/lib/${1} \ - tensorflow/contrib/lite/tools/make/gen/ios_armv7s/lib/${1} \ - tensorflow/contrib/lite/tools/make/gen/ios_arm64/lib/${1} \ - -create \ - -output tensorflow/contrib/lite/tools/make/gen/lib/${1} -} - -make_library libtensorflow-lite.a -make_library benchmark-lib.a diff --git a/tensorflow/contrib/lite/tools/make/download_dependencies.sh b/tensorflow/contrib/lite/tools/make/download_dependencies.sh deleted file mode 100755 index 3570f9a38d3fdc435e5c0caeb04da39c422710f1..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/tools/make/download_dependencies.sh +++ /dev/null @@ -1,104 +0,0 @@ -#!/bin/bash -# 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. -# ============================================================================== - -set -e - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -cd "$SCRIPT_DIR/../../../../.." - -DOWNLOADS_DIR=tensorflow/contrib/lite/tools/make/downloads -BZL_FILE_PATH=tensorflow/workspace.bzl - -# Ensure it is being run from repo root -if [ ! -f $BZL_FILE_PATH ]; then - echo "Could not find ${BZL_FILE_PATH}": - echo "Likely you are not running this from the root directory of the repository."; - exit 1; -fi - -EIGEN_URL="$(grep -o 'http.*bitbucket.org/eigen/eigen/get/.*tar\.gz' "${BZL_FILE_PATH}" | grep -v mirror.bazel | head -n1)" -GEMMLOWP_URL="$(grep -o 'https://mirror.bazel.build/github.com/google/gemmlowp/.*zip' "${BZL_FILE_PATH}" | head -n1)" -GOOGLETEST_URL="https://github.com/google/googletest/archive/release-1.8.0.tar.gz" -ABSL_URL="$(grep -o 'https://github.com/abseil/abseil-cpp/.*tar.gz' "${BZL_FILE_PATH}" | head -n1)" -NEON_2_SSE_URL="https://github.com/intel/ARM_NEON_2_x86_SSE/archive/master.zip" -FARMHASH_URL="https://mirror.bazel.build/github.com/google/farmhash/archive/816a4ae622e964763ca0862d9dbd19324a1eaf45.tar.gz" -FLATBUFFERS_URL="https://github.com/google/flatbuffers/archive/1f5eae5d6a135ff6811724f6c57f911d1f46bb15.tar.gz" -FFT2D_URL="https://mirror.bazel.build/www.kurims.kyoto-u.ac.jp/~ooura/fft.tgz" - -# TODO(petewarden): Some new code in Eigen triggers a clang bug with iOS arm64, -# so work around it by patching the source. -replace_by_sed() { - local regex="${1}" - shift - # Detect the version of sed by the return value of "--version" flag. GNU-sed - # supports "--version" while BSD-sed doesn't. - if ! sed --version >/dev/null 2>&1; then - # BSD-sed. - sed -i '' -e "${regex}" "$@" - else - # GNU-sed. - sed -i -e "${regex}" "$@" - fi -} - -download_and_extract() { - local usage="Usage: download_and_extract URL DIR" - local url="${1:?${usage}}" - local dir="${2:?${usage}}" - echo "downloading ${url}" >&2 - mkdir -p "${dir}" - if [[ "${url}" == *gz ]]; then - curl -Ls "${url}" | tar -C "${dir}" --strip-components=1 -xz - elif [[ "${url}" == *zip ]]; then - tempdir=$(mktemp -d) - tempdir2=$(mktemp -d) - - curl -L ${url} > ${tempdir}/zipped.zip - unzip ${tempdir}/zipped.zip -d ${tempdir2} - - # If the zip file contains nested directories, extract the files from the - # inner directory. - if ls ${tempdir2}/*/* 1> /dev/null 2>&1; then - # unzip has no strip components, so unzip to a temp dir, and move the - # files we want from the tempdir to destination. - cp -R ${tempdir2}/*/* ${dir}/ - else - cp -R ${tempdir2}/* ${dir}/ - fi - rm -rf ${tempdir2} ${tempdir} - fi - - # Delete any potential BUILD files, which would interfere with Bazel builds. - find "${dir}" -type f -name '*BUILD' -delete -} - -download_and_extract "${EIGEN_URL}" "${DOWNLOADS_DIR}/eigen" -download_and_extract "${GEMMLOWP_URL}" "${DOWNLOADS_DIR}/gemmlowp" -download_and_extract "${GOOGLETEST_URL}" "${DOWNLOADS_DIR}/googletest" -download_and_extract "${ABSL_URL}" "${DOWNLOADS_DIR}/absl" -download_and_extract "${NEON_2_SSE_URL}" "${DOWNLOADS_DIR}/neon_2_sse" -download_and_extract "${FARMHASH_URL}" "${DOWNLOADS_DIR}/farmhash" -download_and_extract "${FLATBUFFERS_URL}" "${DOWNLOADS_DIR}/flatbuffers" -download_and_extract "${FFT2D_URL}" "${DOWNLOADS_DIR}/fft2d" - -replace_by_sed 's#static uint32x4_t p4ui_CONJ_XOR = vld1q_u32( conj_XOR_DATA );#static uint32x4_t p4ui_CONJ_XOR; // = vld1q_u32( conj_XOR_DATA ); - Removed by script#' \ - "${DOWNLOADS_DIR}/eigen/Eigen/src/Core/arch/NEON/Complex.h" -replace_by_sed 's#static uint32x2_t p2ui_CONJ_XOR = vld1_u32( conj_XOR_DATA );#static uint32x2_t p2ui_CONJ_XOR;// = vld1_u32( conj_XOR_DATA ); - Removed by scripts#' \ - "${DOWNLOADS_DIR}/eigen/Eigen/src/Core/arch/NEON/Complex.h" -replace_by_sed 's#static uint64x2_t p2ul_CONJ_XOR = vld1q_u64( p2ul_conj_XOR_DATA );#static uint64x2_t p2ul_CONJ_XOR;// = vld1q_u64( p2ul_conj_XOR_DATA ); - Removed by script#' \ - "${DOWNLOADS_DIR}/eigen/Eigen/src/Core/arch/NEON/Complex.h" - -echo "download_dependencies.sh completed successfully." >&2 diff --git a/tensorflow/contrib/lite/tools/optimize/BUILD b/tensorflow/contrib/lite/tools/optimize/BUILD deleted file mode 100644 index 51ccaedc23d0abfda83295879b007f2479d0c571..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/tools/optimize/BUILD +++ /dev/null @@ -1,25 +0,0 @@ -# TODO(suharshs): Write quantize_weights tests that use small exportable files. -# Then we can remove this file. -package( - default_visibility = ["//visibility:public"], -) - -licenses(["notice"]) # Apache 2.0 - -exports_files(["LICENSE"]) - -load("//tensorflow/contrib/lite:build_def.bzl", "tflite_copts") - -cc_library( - name = "quantize_weights", - srcs = ["quantize_weights.cc"], - hdrs = ["quantize_weights.h"], - deps = [ - "//tensorflow/contrib/lite:framework", - "//tensorflow/contrib/lite/kernels/internal:tensor_utils", - "//tensorflow/contrib/lite/schema:schema_fbs", - "//tensorflow/core:tflite_portable_logging", - "@com_google_absl//absl/memory", - "@flatbuffers", - ], -) diff --git a/tensorflow/contrib/lite/tools/verifier_test.cc b/tensorflow/contrib/lite/tools/verifier_test.cc deleted file mode 100644 index ad7d59ecb41a0c81a6a4d8edae5fa6b4b5a7bede..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/tools/verifier_test.cc +++ /dev/null @@ -1,287 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include -#include - -#include "flatbuffers/flatbuffers.h" -#include "flatbuffers/util.h" -#include -#include "tensorflow/contrib/lite/allocation.h" -#include "tensorflow/contrib/lite/error_reporter.h" -#include "tensorflow/contrib/lite/op_resolver.h" -#include "tensorflow/contrib/lite/schema/schema_generated.h" -#include "tensorflow/contrib/lite/testing/util.h" -#include "tensorflow/contrib/lite/tools/verifier.h" -#include "tensorflow/contrib/lite/version.h" -#include "tensorflow/core/framework/numeric_types.h" - -namespace tflite { - -using flatbuffers::FlatBufferBuilder; -using flatbuffers::Offset; - -// Build single subgraph model. -class TfLiteFlatbufferModelBuilder { - public: - TfLiteFlatbufferModelBuilder() { - buffers_.push_back( - CreateBuffer(builder_, builder_.CreateVector(std::vector{}))); - } - - TfLiteFlatbufferModelBuilder(const std::vector& builtin_ops, - const std::vector& custom_ops) { - buffers_.push_back( - CreateBuffer(builder_, builder_.CreateVector(std::vector{}))); - - for (const auto& iter : builtin_ops) { - resolver_.AddBuiltin(iter, &fake_op_); - } - for (const auto& iter : custom_ops) { - resolver_.AddCustom(iter.data(), &fake_op_); - } - } - - void AddTensor(const std::vector& shape, tflite::TensorType type, - const std::vector& buffer, const char* name) { - int buffer_index = 0; - if (!buffer.empty()) { - buffer_index = buffers_.size(); - buffers_.push_back(CreateBuffer(builder_, builder_.CreateVector(buffer))); - } - tensors_.push_back(CreateTensorDirect(builder_, &shape, type, buffer_index, - name, /*quantization=*/0)); - } - - void AddOperator(const std::vector& inputs, - const std::vector& outputs, - tflite::BuiltinOperator builtin_op, const char* custom_op) { - operator_codes_.push_back( - CreateOperatorCodeDirect(builder_, builtin_op, custom_op)); - operators_.push_back(CreateOperator( - builder_, operator_codes_.size() - 1, builder_.CreateVector(inputs), - builder_.CreateVector(outputs), BuiltinOptions_NONE, - /*builtin_options=*/0, - /*custom_options=*/0, tflite::CustomOptionsFormat_FLEXBUFFERS)); - } - - void FinishModel(const std::vector& inputs, - const std::vector& outputs) { - auto subgraph = std::vector>({CreateSubGraph( - builder_, builder_.CreateVector(tensors_), - builder_.CreateVector(inputs), builder_.CreateVector(outputs), - builder_.CreateVector(operators_), - builder_.CreateString("test_subgraph"))}); - auto result = CreateModel( - builder_, TFLITE_SCHEMA_VERSION, builder_.CreateVector(operator_codes_), - builder_.CreateVector(subgraph), builder_.CreateString("test_model"), - builder_.CreateVector(buffers_)); - tflite::FinishModelBuffer(builder_, result); - } - - bool Verify() { - return tflite::Verify(builder_.GetBufferPointer(), builder_.GetSize(), - resolver_, DefaultErrorReporter()); - } - - private: - FlatBufferBuilder builder_; - MutableOpResolver resolver_; - TfLiteRegistration fake_op_; - std::vector> operators_; - std::vector> operator_codes_; - std::vector> tensors_; - std::vector> buffers_; -}; - -TEST(VerifyModel, TestEmptyModel) { - FlatBufferBuilder builder; - auto model = CreateModel(builder, /*version=*/TFLITE_SCHEMA_VERSION, - /*operator_codes=*/0, /*subgraphs=*/0, - /*description=*/0, /*buffers=*/0); - ::tflite::FinishModelBuffer(builder, model); - - ASSERT_FALSE(Verify(builder.GetBufferPointer(), builder.GetSize(), - MutableOpResolver{}, DefaultErrorReporter())); -} - -TEST(VerifyModel, TestSimpleModel) { - TfLiteFlatbufferModelBuilder builder({}, {"test"}); - builder.AddOperator({0, 1}, {2}, BuiltinOperator_CUSTOM, "test"); - builder.AddTensor({2, 3}, TensorType_UINT8, {1, 2, 3, 4, 5, 6}, "input"); - builder.AddTensor( - {2}, TensorType_STRING, - {2, 0, 0, 0, 16, 0, 0, 0, 17, 0, 0, 0, 19, 0, 0, 0, 'A', 'B', 'C'}, - "data"); - builder.AddTensor({2, 3}, TensorType_INT32, {}, "output"); - builder.FinishModel({0, 1}, {2}); - ASSERT_TRUE(builder.Verify()); -} - -TEST(VerifyModel, TestCorruptedData) { - std::string model = "123"; - ASSERT_FALSE(Verify(model.data(), model.size(), MutableOpResolver{}, - /*error_reporter=*/nullptr)); -} - -TEST(VerifyModel, TestUnsupportedVersion) { - FlatBufferBuilder builder; - auto model = CreateModel(builder, /*version=*/1, /*operator_codes=*/0, - /*subgraphs=*/0, /*description=*/0, /*buffers=*/0); - ::tflite::FinishModelBuffer(builder, model); - ASSERT_FALSE(Verify(builder.GetBufferPointer(), builder.GetSize(), - MutableOpResolver{}, DefaultErrorReporter())); -} - -TEST(VerifyModel, TestRandomModificationIsNotAllowed) { - FlatBufferBuilder builder; - auto model = CreateModel(builder, /*version=*/TFLITE_SCHEMA_VERSION, - /*operator_codes=*/0, - /*subgraphs=*/0, /*description=*/0, /*buffers=*/0); - ::tflite::FinishModelBuffer(builder, model); - - std::string model_content(reinterpret_cast(builder.GetBufferPointer()), - builder.GetSize()); - for (int i = 0; i < model_content.size(); i++) { - model_content[i] = (model_content[i] + 137) % 255; - EXPECT_FALSE(Verify(model_content.data(), model_content.size(), - MutableOpResolver{}, DefaultErrorReporter())) - << "Fail at position: " << i; - } -} - -TEST(VerifyModel, TestIntTensorShapeIsGreaterThanBuffer) { - TfLiteFlatbufferModelBuilder builder; - builder.AddTensor({2, 3}, TensorType_UINT8, {1, 2, 3, 4}, "input"); - builder.FinishModel({}, {}); - ASSERT_FALSE(builder.Verify()); -} - -TEST(VerifyModel, TestIntTensorShapeIsSmallerThanBuffer) { - TfLiteFlatbufferModelBuilder builder; - builder.AddTensor({2, 1}, TensorType_UINT8, {1, 2, 3, 4}, "input"); - builder.FinishModel({}, {}); - ASSERT_FALSE(builder.Verify()); -} - -TEST(VerifyModel, TestIntTensorShapeOverflow) { - TfLiteFlatbufferModelBuilder builder; - builder.AddTensor({1024, 2048, 4096}, TensorType_UINT8, {1, 2, 3, 4}, - "input"); - builder.FinishModel({}, {}); - ASSERT_FALSE(builder.Verify()); -} - -TEST(VerifyModel, TensorBufferIsNotValid) { - FlatBufferBuilder builder; - std::vector shape = {2, 3}; - auto tensors = builder.CreateVector(std::vector>{ - CreateTensorDirect(builder, &shape, TensorType_INT32, /*buffer=*/2, - "input", /*quantization=*/0)}); - auto subgraph = std::vector>( - {CreateSubGraph(builder, tensors, /*inputs=*/0, /*outputs=*/0, - /*operators=*/0, builder.CreateString("Main"))}); - - auto buffers = builder.CreateVector(std::vector>{ - CreateBuffer(builder, builder.CreateVector( - std::vector{1, 2, 3, 4, 5, 6})), - }); - - auto model = CreateModel(builder, TFLITE_SCHEMA_VERSION, /*operator_codes=*/0, - builder.CreateVector(subgraph), - builder.CreateString("SmartReply"), buffers); - - ::tflite::FinishModelBuffer(builder, model); - ASSERT_FALSE(Verify(builder.GetBufferPointer(), builder.GetSize(), - MutableOpResolver{}, DefaultErrorReporter())); -} - -TEST(VerifyModel, StringTensorHasInvalidNumString) { - TfLiteFlatbufferModelBuilder builder; - builder.AddTensor( - {2}, TensorType_STRING, - {0x00, 0x00, 0x00, 0x20, 16, 0, 0, 0, 17, 0, 0, 0, 18, 0, 0, 0, 'A', 'B'}, - "input"); - builder.FinishModel({}, {}); - ASSERT_FALSE(builder.Verify()); -} - -TEST(VerifyModel, StringTensorOffsetTooSmall) { - TfLiteFlatbufferModelBuilder builder; - builder.AddTensor( - {2}, TensorType_STRING, - {2, 0, 0, 0, 12, 0, 0, 0, 17, 0, 0, 0, 18, 0, 0, 0, 'A', 'B'}, "input"); - builder.FinishModel({}, {}); - ASSERT_FALSE(builder.Verify()); -} - -TEST(VerifyModel, StringTensorOffsetOutOfRange) { - TfLiteFlatbufferModelBuilder builder; - builder.AddTensor( - {2}, TensorType_STRING, - {2, 0, 0, 0, 16, 0, 0, 0, 17, 0, 0, 0, 22, 0, 0, 0, 'A', 'B'}, "input"); - builder.FinishModel({}, {}); - ASSERT_FALSE(builder.Verify()); -} - -TEST(VerifyModel, StringTensorIsLargerThanRequired) { - TfLiteFlatbufferModelBuilder builder; - builder.AddTensor( - {2}, TensorType_STRING, - {2, 0, 0, 0, 16, 0, 0, 0, 17, 0, 0, 0, 18, 0, 0, 0, 'A', 'B', 'C'}, - "input"); - builder.FinishModel({}, {}); - ASSERT_FALSE(builder.Verify()); -} - -TEST(VerifyModel, AllOpsAreSupported) { - TfLiteFlatbufferModelBuilder builder({BuiltinOperator_ADD}, {"CustomOp"}); - builder.AddTensor({2, 3}, TensorType_UINT8, {1, 2, 3, 4}, "input1"); - builder.AddTensor({2, 3}, TensorType_UINT8, {1, 2, 3, 4}, "input2"); - builder.AddTensor({2, 3}, TensorType_UINT8, {}, "output"); - builder.AddOperator({0, 1}, {2}, BuiltinOperator_ADD, nullptr); - builder.AddOperator({0, 1}, {2}, BuiltinOperator_CUSTOM, "CustomOp"); - builder.FinishModel({}, {}); - ASSERT_FALSE(builder.Verify()); -} - -TEST(VerifyModel, UseUnsupportedBuiltinOps) { - TfLiteFlatbufferModelBuilder builder({BuiltinOperator_SUB}, {"CustomOp"}); - builder.AddTensor({2, 3}, TensorType_UINT8, {1, 2, 3, 4}, "input1"); - builder.AddTensor({2, 3}, TensorType_UINT8, {1, 2, 3, 4}, "input2"); - builder.AddTensor({2, 3}, TensorType_UINT8, {}, "output"); - builder.AddOperator({0, 1}, {2}, BuiltinOperator_ADD, nullptr); - builder.FinishModel({}, {}); - ASSERT_FALSE(builder.Verify()); -} - -TEST(VerifyModel, UseUnsupportedCustomOps) { - TfLiteFlatbufferModelBuilder builder({BuiltinOperator_ADD}, {"NewOp"}); - builder.AddTensor({2, 3}, TensorType_UINT8, {1, 2, 3, 4}, "input1"); - builder.AddTensor({2, 3}, TensorType_UINT8, {1, 2, 3, 4}, "input2"); - builder.AddTensor({2, 3}, TensorType_UINT8, {}, "output"); - builder.AddOperator({0, 1}, {2}, BuiltinOperator_CUSTOM, "Not supported"); - builder.FinishModel({}, {}); - ASSERT_FALSE(builder.Verify()); -} - -// TODO(yichengfan): make up malicious files to test with. - -} // namespace tflite - -int main(int argc, char** argv) { - ::tflite::LogToStderr(); - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/contrib/lite/util.cc b/tensorflow/contrib/lite/util.cc deleted file mode 100644 index 6aa35b52277910aea7ad0ea8753c2bad095b1f1f..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/util.cc +++ /dev/null @@ -1,58 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include "tensorflow/contrib/lite/util.h" - -#include - -namespace tflite { - -bool IsFlexOp(const char* custom_name) { - return custom_name && strncmp(custom_name, kFlexCustomCodePrefix, - strlen(kFlexCustomCodePrefix)) == 0; -} - -TfLiteIntArray* ConvertVectorToTfLiteIntArray(const std::vector& input) { - return ConvertArrayToTfLiteIntArray(input.size(), input.data()); -} - -TfLiteIntArray* ConvertArrayToTfLiteIntArray(const int rank, const int* dims) { - TfLiteIntArray* output = TfLiteIntArrayCreate(rank); - for (size_t i = 0; i < rank; i++) { - output->data[i] = dims[i]; - } - return output; -} - -bool EqualArrayAndTfLiteIntArray(const TfLiteIntArray* a, const int b_size, - const int* b) { - if (!a) return false; - if (a->size != b_size) return false; - for (int i = 0; i < a->size; ++i) { - if (a->data[i] != b[i]) return false; - } - return true; -} - -size_t CombineHashes(std::initializer_list hashes) { - size_t result = 0; - // Hash combiner used by TensorFlow core. - for (size_t hash : hashes) { - result = result ^ - (hash + 0x9e3779b97f4a7800ULL + (result << 10) + (result >> 4)); - } - return result; -} - -} // namespace tflite diff --git a/tensorflow/contrib/lite/util.h b/tensorflow/contrib/lite/util.h deleted file mode 100644 index 31292a6f8131f78f0939b4ebbcb46dfd9c3312df..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/util.h +++ /dev/null @@ -1,57 +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. -==============================================================================*/ - -// This file provides general C++ utility functions in TFLite. -// For example: Converting between `TfLiteIntArray`, `std::vector` and -// Flatbuffer vectors. These functions can't live in `context.h` since it's pure -// C. - -#ifndef TENSORFLOW_CONTRIB_LITE_UTIL_H_ -#define TENSORFLOW_CONTRIB_LITE_UTIL_H_ - -#include -#include "tensorflow/contrib/lite/c/c_api_internal.h" - -namespace tflite { - -// The prefix of Flex op custom code. -// This will be matched agains the `custom_code` field in `OperatorCode` -// Flatbuffer Table. -// WARNING: This is an experimental API and subject to change. -constexpr char kFlexCustomCodePrefix[] = "Flex"; - -// Checks whether the prefix of the custom name indicates the operation is an -// Flex operation. -bool IsFlexOp(const char* custom_name); - -// Converts a `std::vector` to a `TfLiteIntArray`. The caller takes ownership -// of the returned pointer. -TfLiteIntArray* ConvertVectorToTfLiteIntArray(const std::vector& input); - -// Converts an array (of the given size) to a `TfLiteIntArray`. The caller -// takes ownership of the returned pointer, and must make sure 'dims' has at -// least 'rank' elemnts. -TfLiteIntArray* ConvertArrayToTfLiteIntArray(const int rank, const int* dims); - -// Checks whether a `TfLiteIntArray` and an int array have matching elements. -// The caller must guarantee that 'b' has at least 'b_size' elements. -bool EqualArrayAndTfLiteIntArray(const TfLiteIntArray* a, const int b_size, - const int* b); - -size_t CombineHashes(std::initializer_list hashes); - -} // namespace tflite - -#endif // TENSORFLOW_CONTRIB_LITE_UTIL_H_ diff --git a/tensorflow/contrib/lite/util_test.cc b/tensorflow/contrib/lite/util_test.cc deleted file mode 100644 index 25f3aded7140ff1075a52d5d30e270ee049c7c88..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/lite/util_test.cc +++ /dev/null @@ -1,60 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include -#include -#include - -#include "tensorflow/contrib/lite/c/c_api_internal.h" -#include "tensorflow/contrib/lite/util.h" - -namespace tflite { -namespace { - -TEST(ConvertVectorToTfLiteIntArray, TestWithVector) { - std::vector input = {1, 2}; - TfLiteIntArray* output = ConvertVectorToTfLiteIntArray(input); - ASSERT_NE(output, nullptr); - EXPECT_EQ(output->size, 2); - EXPECT_EQ(output->data[0], 1); - EXPECT_EQ(output->data[1], 2); - TfLiteIntArrayFree(output); -} - -TEST(ConvertVectorToTfLiteIntArray, TestWithEmptyVector) { - std::vector input; - TfLiteIntArray* output = ConvertVectorToTfLiteIntArray(input); - ASSERT_NE(output, nullptr); - EXPECT_EQ(output->size, 0); - TfLiteIntArrayFree(output); -} - -TEST(UtilTest, IsFlexOp) { - EXPECT_TRUE(IsFlexOp("Flex")); - EXPECT_TRUE(IsFlexOp("FlexOp")); - EXPECT_FALSE(IsFlexOp("flex")); - EXPECT_FALSE(IsFlexOp("Fle")); - EXPECT_FALSE(IsFlexOp("OpFlex")); - EXPECT_FALSE(IsFlexOp(nullptr)); - EXPECT_FALSE(IsFlexOp("")); -} - -} // namespace -} // namespace tflite - -int main(int argc, char** argv) { - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/contrib/lookup/lookup_ops.py b/tensorflow/contrib/lookup/lookup_ops.py index 5abef822e82a1e9f818e54e32c2980a985d41ad8..229a72a780d5ccce8263444ffeae7700f6ac8613 100644 --- a/tensorflow/contrib/lookup/lookup_ops.py +++ b/tensorflow/contrib/lookup/lookup_ops.py @@ -42,7 +42,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.checkpointable import base as checkpointable from tensorflow.python.training.saver import BaseSaverBuilder from tensorflow.python.util.deprecation import deprecated @@ -92,7 +91,7 @@ def index_table_from_tensor(mapping, The bucket ID range is `[mapping size, mapping size + num_oov_buckets - 1]`. The underlying table must be initialized by calling - `tf.tables_initializer.run()` or `table.init.run()` once. + `session.run(tf.tables_initializer)` or `session.run(table.init)` once. Elements in `mapping` cannot have duplicates, otherwise when executing the table initializer op, it will throw a `FailedPreconditionError`. @@ -159,7 +158,7 @@ def string_to_index(tensor, mapping, default_value=-1, name=None): will throw a FailedPreconditionError. The underlying table must be initialized by calling - `tf.tables_initializer.run()` once. + `session.run(tf.tables_initializer)` once. For example: @@ -203,7 +202,7 @@ def index_to_string_table_from_tensor(mapping, default_value="UNK", name=None): (an out-of-vocabulary entry) is assigned the `default_value` The underlying table must be initialized by calling - `tf.tables_initializer.run()` or `table.init.run()` once. + `session.run(tf.tables_initializer)` or `session.run(table.init)` once. Elements in `mapping` cannot have duplicates, otherwise when executing the table initializer op, it will throw a `FailedPreconditionError`. @@ -258,7 +257,7 @@ def index_to_string(tensor, mapping, default_value="UNK", name=None): (an out-of-vocabulary entry) is assigned the `default_value` The underlying table must be initialized by calling - `tf.tables_initializer.run()` once. + `session.run(tf.tables_initializer)` once. For example: @@ -289,7 +288,7 @@ def index_to_string(tensor, mapping, default_value="UNK", name=None): return table.lookup(tensor) -class MutableHashTable(LookupInterface, checkpointable.CheckpointableBase): +class MutableHashTable(LookupInterface): """A generic mutable hash table implementation. Data can be inserted by calling the insert method and removed by calling the @@ -339,43 +338,56 @@ class MutableHashTable(LookupInterface, checkpointable.CheckpointableBase): 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 - executing_eagerly = context.executing_eagerly() - if executing_eagerly and shared_name is None: + 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(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 = checkpoint and shared_name is None + use_node_name_sharing = self._checkpoint and self._shared_name is None if self._default_value.get_shape().ndims == 0: - self._table_ref = gen_lookup_ops.mutable_hash_table_v2( - shared_name=shared_name, + table_ref = gen_lookup_ops.mutable_hash_table_v2( + shared_name=self._shared_name, use_node_name_sharing=use_node_name_sharing, - key_dtype=key_dtype, - value_dtype=value_dtype, - name=name) + key_dtype=self._key_dtype, + value_dtype=self._value_dtype, + name=self._name) else: - self._table_ref = gen_lookup_ops.mutable_hash_table_of_tensors_v2( - shared_name=shared_name, + 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=key_dtype, - value_dtype=value_dtype, + key_dtype=self._key_dtype, + value_dtype=self._value_dtype, value_shape=self._default_value.get_shape(), - name=name) - if executing_eagerly: - op_name = None + name=self._name) + + if context.executing_eagerly(): + self._table_name = None else: - op_name = self._table_ref.op.name.split("/")[-1] - super(MutableHashTable, self).__init__(key_dtype, value_dtype, - op_name) + self._table_name = table_ref.op.name.split("/")[-1] + return table_ref - if checkpoint: - saveable = MutableHashTable._Saveable(self, name) - ops.add_to_collection(ops.GraphKeys.SAVEABLE_OBJECTS, saveable) + @property + def name(self): + return self._table_name def size(self, name=None): """Compute the number of elements in this table. @@ -386,10 +398,11 @@ class MutableHashTable(LookupInterface, checkpointable.CheckpointableBase): Returns: A scalar tensor containing the number of elements in this table. """ - with ops.name_scope(name, "%s_Size" % self._name, - [self._table_ref]) as name: - with ops.colocate_with(self._table_ref): - return gen_lookup_ops.lookup_table_size_v2(self._table_ref, name=name) + 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. @@ -411,11 +424,12 @@ class MutableHashTable(LookupInterface, checkpointable.CheckpointableBase): 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._table_ref, keys, self._default_value)) as name: + 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._table_ref, keys, name=name) + self.resource_handle, keys, name=name) return op @@ -436,12 +450,13 @@ class MutableHashTable(LookupInterface, checkpointable.CheckpointableBase): Raises: TypeError: when `keys` do not match the table data types. """ - with ops.name_scope(name, "%s_lookup_table_find" % self._name, - (self._table_ref, keys, self._default_value)) as name: + 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._table_ref): + with ops.colocate_with(self.resource_handle): values = gen_lookup_ops.lookup_table_find_v2( - self._table_ref, keys, self._default_value, name=name) + self.resource_handle, keys, self._default_value, name=name) return values def insert(self, keys, values, name=None): @@ -461,14 +476,14 @@ class MutableHashTable(LookupInterface, checkpointable.CheckpointableBase): TypeError: when `keys` or `values` doesn't match the table data types. """ - with ops.name_scope(name, "%s_lookup_table_insert" % self._name, - [self._table_ref, keys, values]) as name: + 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._table_ref): + with ops.colocate_with(self.resource_handle): # pylint: disable=protected-access op = gen_lookup_ops.lookup_table_insert_v2( - self._table_ref, keys, values, name=name) + self.resource_handle, keys, values, name=name) return op def export(self, name=None): @@ -481,11 +496,11 @@ class MutableHashTable(LookupInterface, checkpointable.CheckpointableBase): 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._table_ref]) as name: - with ops.colocate_with(self._table_ref): + 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._table_ref, self._key_dtype, self._value_dtype, name=name) + self.resource_handle, self._key_dtype, self._value_dtype, name=name) return exported_keys, exported_values def _gather_saveables_for_checkpoint(self): @@ -507,12 +522,12 @@ class MutableHashTable(LookupInterface, checkpointable.CheckpointableBase): def restore(self, restored_tensors, restored_shapes): del restored_shapes # unused # pylint: disable=protected-access - with ops.colocate_with(self.op._table_ref): + with ops.colocate_with(self.op.resource_handle): return gen_lookup_ops.lookup_table_import_v2( - self.op._table_ref, restored_tensors[0], restored_tensors[1]) + self.op.resource_handle, restored_tensors[0], restored_tensors[1]) -class MutableDenseHashTable(LookupInterface, checkpointable.CheckpointableBase): +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 @@ -581,42 +596,55 @@ class MutableDenseHashTable(LookupInterface, checkpointable.CheckpointableBase): """ 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._value_shape = self._default_value.get_shape() + self._checkpoint = checkpoint + self._name = name - # 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 = checkpoint and shared_name is None - empty_key = ops.convert_to_tensor( + self._empty_key = ops.convert_to_tensor( empty_key, dtype=key_dtype, name="empty_key") - deleted_key = ops.convert_to_tensor( + self._deleted_key = ops.convert_to_tensor( deleted_key, dtype=key_dtype, name="deleted_key") - executing_eagerly = context.executing_eagerly() - if executing_eagerly and shared_name is None: + 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._table_ref = gen_lookup_ops.mutable_dense_hash_table_v2( - empty_key=empty_key, - deleted_key=deleted_key, - shared_name=shared_name, + 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, + shared_name=self._shared_name, use_node_name_sharing=use_node_name_sharing, - value_dtype=value_dtype, + value_dtype=self._value_dtype, value_shape=self._value_shape, - initial_num_buckets=initial_num_buckets, - name=name) - if executing_eagerly: - op_name = None + initial_num_buckets=self._initial_num_buckets, + name=self._name) + if context.executing_eagerly(): + self._table_name = None else: - op_name = self._table_ref.op.name.split("/")[-1] - super(MutableDenseHashTable, self).__init__( - key_dtype, value_dtype, op_name) + self._table_name = table_ref.op.name.split("/")[-1] + return table_ref - if checkpoint: - saveable = MutableDenseHashTable._Saveable(self, name) - ops.add_to_collection(ops.GraphKeys.SAVEABLE_OBJECTS, saveable) + @property + def name(self): + return self._table_name def size(self, name=None): """Compute the number of elements in this table. @@ -627,10 +655,11 @@ class MutableDenseHashTable(LookupInterface, checkpointable.CheckpointableBase): Returns: A scalar tensor containing the number of elements in this table. """ - with ops.name_scope(name, "%s_Size" % self._name, - [self._table_ref]) as name: - with ops.colocate_with(self._table_ref): - return gen_lookup_ops.lookup_table_size_v2(self._table_ref, name=name) + 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. @@ -649,12 +678,12 @@ class MutableDenseHashTable(LookupInterface, checkpointable.CheckpointableBase): Raises: TypeError: when `keys` do not match the table data types. """ - with ops.name_scope(name, "%s_lookup_table_find" % self._name, - [self._table_ref, keys]) as name: + 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._table_ref): + with ops.colocate_with(self.resource_handle): values = gen_lookup_ops.lookup_table_find_v2( - self._table_ref, keys, self._default_value, name=name) + self.resource_handle, keys, self._default_value, name=name) return values @@ -675,14 +704,14 @@ class MutableDenseHashTable(LookupInterface, checkpointable.CheckpointableBase): TypeError: when `keys` or `values` doesn't match the table data types. """ - with ops.name_scope(name, "%s_lookup_table_insert" % self._name, - [self._table_ref, keys, values]) as name: + 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._table_ref): + with ops.colocate_with(self.resource_handle): op = gen_lookup_ops.lookup_table_insert_v2( - self._table_ref, keys, values, name=name) + self.resource_handle, keys, values, name=name) return op def remove(self, keys, name=None): @@ -705,11 +734,12 @@ class MutableDenseHashTable(LookupInterface, checkpointable.CheckpointableBase): 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._table_ref, keys, self._default_value)) as name: + 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._table_ref, keys, name=name) + self.resource_handle, keys, name=name) return op @@ -723,11 +753,11 @@ class MutableDenseHashTable(LookupInterface, checkpointable.CheckpointableBase): 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._table_ref]) as name: - with ops.colocate_with(self._table_ref): + 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._table_ref, self._key_dtype, self._value_dtype, name=name) + self.resource_handle, self._key_dtype, self._value_dtype, name=name) return exported_keys, exported_values @@ -751,6 +781,6 @@ class MutableDenseHashTable(LookupInterface, checkpointable.CheckpointableBase): def restore(self, restored_tensors, restored_shapes): del restored_shapes # unused # pylint: disable=protected-access - with ops.colocate_with(self.op._table_ref): + with ops.colocate_with(self.op.resource_handle): return gen_lookup_ops.lookup_table_import_v2( - self.op._table_ref, restored_tensors[0], restored_tensors[1]) + self.op.resource_handle, restored_tensors[0], restored_tensors[1]) diff --git a/tensorflow/contrib/lookup/lookup_ops_test.py b/tensorflow/contrib/lookup/lookup_ops_test.py index 35b0d1bc4447cdd3964dbde6899deca590942c84..9b2c2dd87cc8a92fbb6b45504939be3788b60839 100644 --- a/tensorflow/contrib/lookup/lookup_ops_test.py +++ b/tensorflow/contrib/lookup/lookup_ops_test.py @@ -25,6 +25,7 @@ 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 @@ -50,7 +51,7 @@ class HashTableOpTest(test.TestCase): values = constant_op.constant([0, 1, 2], dtypes.int64) table = lookup.HashTable( lookup.KeyValueTensorInitializer(keys, values), default_val) - table.init.run() + table.initializer.run() self.assertAllEqual(3, table.size().eval()) @@ -74,7 +75,7 @@ class HashTableOpTest(test.TestCase): values = constant_op.constant([0, 1, 2], dtypes.int64) table = lookup.HashTable( lookup.KeyValueTensorInitializer(keys, values), default_val) - table.init.run() + table.initializer.run() self.assertAllEqual(3, table.size().eval()) @@ -94,7 +95,7 @@ class HashTableOpTest(test.TestCase): lookup.KeyValueTensorInitializer( keys, values, value_dtype=dtypes.int64), default_val) - table.init.run() + table.initializer.run() self.assertAllEqual(3, table.size().eval()) @@ -111,7 +112,7 @@ class HashTableOpTest(test.TestCase): values = np.array([0, 1, 2], dtype=np.int64) table = lookup.HashTable( lookup.KeyValueTensorInitializer(keys, values), default_val) - table.init.run() + table.initializer.run() self.assertAllEqual(3, table.size().eval()) @@ -156,7 +157,7 @@ class HashTableOpTest(test.TestCase): values = constant_op.constant([0, 1, 2], dtypes.int64) table = lookup.HashTable( lookup.KeyValueTensorInitializer(keys, values), default_val) - table.init.run() + table.initializer.run() input_string = constant_op.constant(["brain", "salad", "tank"]) output = table.lookup(input_string) @@ -171,7 +172,7 @@ class HashTableOpTest(test.TestCase): values = constant_op.constant([0, 1, 2], dtypes.int64) table = lookup.HashTable( lookup.KeyValueTensorInitializer(keys, values), default_val) - table.init.run() + table.initializer.run() sp_indices = [[0, 0], [0, 1], [1, 0]] sp_shape = [2, 2] @@ -194,7 +195,7 @@ class HashTableOpTest(test.TestCase): values = constant_op.constant([0, 1, 2], dtypes.int64) table = lookup.HashTable( lookup.KeyValueTensorInitializer(keys, values), default_val) - table.init.run() + table.initializer.run() # Ref types do not produce a lookup signature mismatch. input_string_ref = variables.Variable("brain") @@ -238,10 +239,10 @@ class HashTableOpTest(test.TestCase): values = constant_op.constant([0, 1, 2], dtypes.int64) table = lookup.HashTable( lookup.KeyValueTensorInitializer(keys, values), default_val) - table.init.run() + table.initializer.run() with self.assertRaisesOpError("Table already initialized"): - table.init.run() + table.initializer.run() def testInitializationWithInvalidDimensions(self): with self.cached_session(): @@ -273,13 +274,13 @@ class HashTableOpTest(test.TestCase): # Init the table in the first session. with session1: - table.init.run() + table.initializer.run() self.assertAllEqual(3, table.size().eval()) # Init the table in the second session and verify that we do not get a # "Table already initialized" error. with session2: - table.init.run() + table.initializer.run() self.assertAllEqual(3, table.size().eval()) def testHashTableInt32String(self): @@ -289,7 +290,7 @@ class HashTableOpTest(test.TestCase): values = constant_op.constant(["brain", "salad", "surgery"]) table = lookup.HashTable( lookup.KeyValueTensorInitializer(keys, values), default_val) - table.init.run() + table.initializer.run() input_tensor = constant_op.constant([0, 1, -1]) output = table.lookup(input_tensor) @@ -1669,7 +1670,7 @@ class IndexTableFromFile(test.TestCase): table = lookup.index_table_from_file( vocabulary_file=vocabulary_file, vocab_size=4) self.assertRaisesRegexp(errors_impl.InvalidArgumentError, - "Invalid vocab_size", table.init.run) + "Invalid vocab_size", table.initializer.run) def test_index_table_from_file_with_vocab_size(self): vocabulary_file = self._createVocabFile("f2i_vocab8.txt") @@ -1717,14 +1718,14 @@ class KeyValueTensorInitializerTest(test.TestCase): init = lookup.KeyValueTensorInitializer( ("brain", "salad", "surgery"), (0, 1, 2), dtypes.string, dtypes.int64) table = lookup.HashTable(init, default_value=-1) - table.init.run() + table.initializer.run() def test_int64(self): with ops.Graph().as_default(), self.cached_session(): init = lookup.KeyValueTensorInitializer( (42, 1, -1000), (0, 1, 2), dtypes.int64, dtypes.int64) table = lookup.HashTable(init, default_value=-1) - table.init.run() + table.initializer.run() def test_int32(self): with ops.Graph().as_default(), self.cached_session(): @@ -1733,7 +1734,7 @@ class KeyValueTensorInitializerTest(test.TestCase): table = lookup.HashTable(init, default_value=-1) with self.assertRaisesRegexp( errors_impl.OpError, "No OpKernel was registered"): - table.init.run() + table.initializer.run() class IndexTableFromTensor(test.TestCase): @@ -2021,7 +2022,7 @@ class InitializeTableFromFileOpTest(test.TestCase): dtypes.int64, lookup.TextFileIndex.LINE_NUMBER), default_value) - self.evaluate(table.init) + self.evaluate(table.initializer) output = table.lookup(constant_op.constant(["brain", "salad", "tank"])) @@ -2040,7 +2041,7 @@ class InitializeTableFromFileOpTest(test.TestCase): dtypes.int64, lookup.TextFileIndex.LINE_NUMBER), default_value) - table.init.run() + table.initializer.run() output = table.lookup( constant_op.constant((42, 1, 11), dtype=dtypes.int64)) @@ -2059,7 +2060,7 @@ class InitializeTableFromFileOpTest(test.TestCase): lookup.TextFileInitializer(vocabulary_file, dtypes.int64, key_index, dtypes.string, value_index), default_value) - table.init.run() + table.initializer.run() input_values = constant_op.constant([0, 1, 2, 3], dtypes.int64) output = table.lookup(input_values) @@ -2081,7 +2082,7 @@ class InitializeTableFromFileOpTest(test.TestCase): lookup.TextFileInitializer(vocabulary_file, dtypes.string, key_index, dtypes.int64, value_index), default_value) - table.init.run() + table.initializer.run() input_string = constant_op.constant(["brain", "salad", "surgery"]) output = table.lookup(input_string) @@ -2103,7 +2104,7 @@ class InitializeTableFromFileOpTest(test.TestCase): key_index, dtypes.int64, value_index), default_value) with self.assertRaisesOpError("is not a valid"): - table.init.run() + table.initializer.run() def testInvalidDataType(self): vocabulary_file = self._createVocabFile("one_column_3.txt") @@ -2131,7 +2132,7 @@ class InitializeTableFromFileOpTest(test.TestCase): default_value) with self.assertRaisesOpError("Invalid number of columns"): - table.init.run() + table.initializer.run() def testInitializeSameTableWithMultipleNodes(self): vocabulary_file = self._createVocabFile("one_column_5.txt") @@ -2200,7 +2201,7 @@ class InitializeTableFromFileOpTest(test.TestCase): default_value) # Initialize from file. - table1.init.run() + table1.initializer.run() self.assertEquals(vocab_size, table1.size().eval()) vocabulary_file2 = self._createVocabFile("one_column7.txt") @@ -2215,7 +2216,7 @@ class InitializeTableFromFileOpTest(test.TestCase): vocab_size=vocab_size), default_value) with self.assertRaisesOpError("Invalid vocab_size"): - table2.init.run() + table2.initializer.run() vocab_size = 1 vocabulary_file3 = self._createVocabFile("one_column3.txt") @@ -2230,7 +2231,7 @@ class InitializeTableFromFileOpTest(test.TestCase): default_value) # Smaller vocab size reads only vocab_size records. - table3.init.run() + table3.initializer.run() self.assertEquals(vocab_size, table3.size().eval()) def testFeedVocabularyName(self): @@ -2248,11 +2249,11 @@ class InitializeTableFromFileOpTest(test.TestCase): # Initialize with non existing file (old_file.txt) should fail. # TODO(yleon): Update message, which might change per FileSystem. with self.assertRaisesOpError("old_file.txt"): - table.init.run() + table.initializer.run() # Initialize the model feeding the vocabulary file. filenames = ops.get_collection(ops.GraphKeys.ASSET_FILEPATHS) - table.init.run(feed_dict={filenames[0]: vocabulary_file}) + table.initializer.run(feed_dict={filenames[0]: vocabulary_file}) input_string = constant_op.constant(["brain", "salad", "tank"]) output = table.lookup(input_string) @@ -2294,7 +2295,7 @@ class InitializeTableFromFileOpTest(test.TestCase): vocab_file, vocab_size=vocab_size), default_value) - table.init.run() + table.initializer.run() input_values = constant_op.constant([0, 1, 2, 3], dtypes.int64) @@ -2311,7 +2312,7 @@ class InitializeTableFromFileOpTest(test.TestCase): lookup.TextFileIdTableInitializer( vocab_file, vocab_size=vocab_size), default_value) - table.init.run() + table.initializer.run() input_string = constant_op.constant(["brain", "salad", "surgery", "UNK"]) @@ -2329,7 +2330,7 @@ class InitializeTableFromFileOpTest(test.TestCase): lookup.TextFileIdTableInitializer( vocab_file, vocab_size=vocab_size, key_dtype=dtypes.int64), default_value) - table.init.run() + table.initializer.run() out = table.lookup( constant_op.constant((42, 1, -1000, 11), dtype=dtypes.int64)) @@ -2358,7 +2359,7 @@ class IdTableWithHashBucketsTest(test.TestCase): default_value), oov_buckets) - table.init.run() + table.initializer.run() input_string = constant_op.constant(["brain", "salad", "surgery", "UNK"]) @@ -2380,7 +2381,7 @@ class IdTableWithHashBucketsTest(test.TestCase): oov_buckets, key_dtype=dtypes.int32) - table.init.run() + table.initializer.run() values = constant_op.constant((42, 1, -1000, 11), dtype=dtypes.int32) @@ -2401,7 +2402,7 @@ class IdTableWithHashBucketsTest(test.TestCase): default_value), oov_buckets) - table.init.run() + table.initializer.run() values = constant_op.constant((42, 1, -1000, 11), dtype=dtypes.int64) @@ -2416,7 +2417,7 @@ class IdTableWithHashBucketsTest(test.TestCase): # Set a table that only uses hash buckets, for each input value returns # an id calculated by fingerprint("input") mod oov_buckets. table = lookup.IdTableWithHashBuckets(None, oov_buckets) - table.init.run() + table.initializer.run() values = constant_op.constant(("brain", "salad", "surgery")) @@ -2438,7 +2439,7 @@ class IdTableWithHashBucketsTest(test.TestCase): # an id calculated by fingerprint("input") mod oov_buckets. table = lookup.IdTableWithHashBuckets( None, oov_buckets, key_dtype=dtypes.int32) - table.init.run() + table.initializer.run() input_string = constant_op.constant([42, 1, -1000], dtype=dtypes.int32) @@ -2520,7 +2521,7 @@ class IdTableWithHashBucketsTest(test.TestCase): shared_name=shared_name), oov_buckets) - table1.init.run() + table1.initializer.run() input_string_1 = constant_op.constant( ["brain", "salad", "surgery", "UNK"]) @@ -2536,7 +2537,7 @@ class IdTableWithHashBucketsTest(test.TestCase): oov_buckets = 1 # Underlying lookup table already initialized in previous session. - # No need to call table2.init.run() + # No need to call table2.initializer.run() table2 = lookup.IdTableWithHashBuckets( lookup.HashTable( lookup.TextFileIdTableInitializer( @@ -2605,7 +2606,7 @@ class IdTableWithHashBucketsTest(test.TestCase): vocab_file, vocab_size=3), -1), 1) - table.init.run() + table.initializer.run() sp_ids = table.lookup(sp_features) @@ -2634,7 +2635,7 @@ class IdTableWithHashBucketsTest(test.TestCase): -1), 1, key_dtype=dtypes.int32) - table.init.run() + table.initializer.run() sp_ids = table.lookup(sp_features) @@ -2663,7 +2664,7 @@ class IdTableWithHashBucketsTest(test.TestCase): -1), 1, key_dtype=dtypes.int64) - table.init.run() + table.initializer.run() sp_ids = table.lookup(sp_features) @@ -2737,7 +2738,7 @@ class MutableHashTableBenchmark(test.Benchmark): def benchmark_many_repeated_scalar_insert_scalar(self): table = self._create_table() - c = counter.Counter().make_one_shot_iterator().get_next() + c = dataset_ops.make_one_shot_iterator(counter.Counter()).get_next() value = variables.Variable(1.0) insert = table.insert(c, value) size = table.size() @@ -2758,7 +2759,7 @@ class MutableHashTableBenchmark(test.Benchmark): def benchmark_many_repeated_batch_32_insert_scalar(self): table = self._create_table() - c = counter.Counter().make_one_shot_iterator().get_next() + 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() diff --git a/tensorflow/contrib/losses/python/losses/loss_ops.py b/tensorflow/contrib/losses/python/losses/loss_ops.py index 619294b51822bd9983eda777acae5cf0d253926d..709a042bbcefb89125f7e4cd14a0d7ecd2b53281 100644 --- a/tensorflow/contrib/losses/python/losses/loss_ops.py +++ b/tensorflow/contrib/losses/python/losses/loss_ops.py @@ -22,7 +22,6 @@ from __future__ import division from __future__ import print_function from tensorflow.contrib.framework.python.ops import add_arg_scope -from tensorflow.python.compat import compat from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import math_ops @@ -60,41 +59,12 @@ def _scale_losses(losses, weights): """ # First, compute the sum of the losses over all elements: start_index = max(0, weights.get_shape().ndims) - reduction_indices = list(range(start_index, losses.get_shape().ndims)) - reduced_losses = math_ops.reduce_sum( - losses, reduction_indices=reduction_indices) + axis = list(range(start_index, losses.get_shape().ndims)) + reduced_losses = math_ops.reduce_sum(losses, axis=axis) reduced_losses = math_ops.multiply(reduced_losses, weights) return math_ops.reduce_sum(reduced_losses) -def _safe_div(numerator, denominator, name="value"): - """Computes a safe divide which returns 0 if the denominator is zero. - - Note that the function contains an additional conditional check that is - necessary for avoiding situations where the loss is zero causing NaNs to - creep into the gradient computation. - - Args: - numerator: An arbitrary `Tensor`. - denominator: A `Tensor` whose shape matches `numerator` and whose values are - assumed to be non-negative. - name: An optional name for the returned op. - - Returns: - The element-wise value of the numerator divided by the denominator. - """ - if compat.forward_compatible(2018, 11, 1): - return math_ops.div_no_nan(numerator, denominator, name=name) - return array_ops.where( - math_ops.greater(denominator, 0), - math_ops.div(numerator, - array_ops.where( - math_ops.equal(denominator, 0), - array_ops.ones_like(denominator), denominator)), - array_ops.zeros_like(numerator), - name=name) - - def _safe_mean(losses, num_present): """Computes a safe mean of the losses. @@ -107,7 +77,7 @@ def _safe_mean(losses, num_present): then zero is returned. """ total_loss = math_ops.reduce_sum(losses) - return _safe_div(total_loss, num_present, name="value") + return math_ops.div_no_nan(total_loss, num_present, name="value") @deprecated("2016-12-30", "Use tf.losses.compute_weighted_loss instead.") @@ -187,10 +157,9 @@ def _num_present(losses, weights, per_batch=False): # First, count the number of nonzero weights: if weights.get_shape().ndims >= 1: - reduction_indices = list(range(1, weights.get_shape().ndims)) + axis = list(range(1, weights.get_shape().ndims)) num_nonzero_per_batch = math_ops.reduce_sum( - math_ops.to_float(math_ops.not_equal(weights, 0)), - reduction_indices=reduction_indices) + math_ops.to_float(math_ops.not_equal(weights, 0)), axis=axis) # Next, determine the number of elements that weights would broadcast to: broadcast_dims = array_ops.slice( @@ -606,20 +575,20 @@ def mean_pairwise_squared_error(predictions, if weights.get_shape().ndims is None: raise ValueError("weights.get_shape().ndims cannot be None") - reduction_indices = list(range(1, diffs.get_shape().ndims)) + axis = list(range(1, diffs.get_shape().ndims)) sum_squares_diff_per_batch = math_ops.reduce_sum( - math_ops.square(diffs), reduction_indices=reduction_indices) + math_ops.square(diffs), axis=axis) num_present_per_batch = _num_present(diffs, weights, per_batch=True) - term1 = 2.0 * _safe_div(sum_squares_diff_per_batch, - num_present_per_batch, - name="value") + term1 = 2.0 * math_ops.div_no_nan( + sum_squares_diff_per_batch, num_present_per_batch, name="value") - sum_diff = math_ops.reduce_sum(diffs, reduction_indices=reduction_indices) - term2 = 2.0 * _safe_div(math_ops.square(sum_diff), - math_ops.square(num_present_per_batch), - name="value") + sum_diff = math_ops.reduce_sum(diffs, axis=axis) + term2 = 2.0 * math_ops.div_no_nan( + math_ops.square(sum_diff), + math_ops.square(num_present_per_batch), + name="value") loss = _scale_losses(term1 - term2, weights) @@ -674,7 +643,7 @@ def cosine_distance(predictions, radial_diffs = math_ops.multiply(predictions, labels) losses = 1 - math_ops.reduce_sum( - radial_diffs, reduction_indices=[ + radial_diffs, axis=[ axis, ]) return compute_weighted_loss(losses, weights, scope=scope) diff --git a/tensorflow/contrib/makefile/Makefile b/tensorflow/contrib/makefile/Makefile index 36125c198e01c177f66b78931ac30e38e17fc409..7ea6e34cf50ed8e292f11314550d992c3dde34c0 100644 --- a/tensorflow/contrib/makefile/Makefile +++ b/tensorflow/contrib/makefile/Makefile @@ -208,6 +208,16 @@ endif # override local versions in the source tree. INCLUDES += -I/usr/local/include +# If `$(WITH_TFLITE_FLEX)` is `true`, this Makefile will build a library +# for TensorFlow Lite Flex runtime. +# Farmhash and Flatbuffer is required for TensorFlow Lite Flex runtime. +ifeq ($(WITH_TFLITE_FLEX), true) + HOST_INCLUDES += -I$(MAKEFILE_DIR)/downloads/farmhash/src + HOST_INCLUDES += -I$(MAKEFILE_DIR)/downloads/flatbuffers/include + INCLUDES += -I$(MAKEFILE_DIR)/downloads/farmhash/src + INCLUDES += -I$(MAKEFILE_DIR)/downloads/flatbuffers/include +endif + LIBS := \ $(TARGET_NSYNC_LIB) \ -lstdc++ \ @@ -283,7 +293,7 @@ ifeq ($(TARGET),ANDROID) else ANDROID_HOST_OS_ARCH := $(ANDROID_HOST_OS_ARCH)-$(HOST_ARCH) endif - + ifndef ANDROID_ARCH ANDROID_ARCH := armeabi-v7a endif @@ -330,7 +340,7 @@ ifeq ($(TARGET),ANDROID) BIN_PREFIX := x86_64-linux-android MARCH_OPTION := endif - + ifndef NDK_ROOT $(error "NDK_ROOT is not defined.") endif @@ -717,6 +727,57 @@ tensorflow/core/util/reporter.cc \ tensorflow/tools/benchmark/benchmark_model.cc \ tensorflow/tools/benchmark/benchmark_model_main.cc +# If `$(WITH_TFLITE_FLEX)` is `true`, this Makefile will build a library +# for TensorFlow Lite Flex runtime. +# Adding the following dependencies> +# * TensorFlow Eager Runtime. +# * TensorFlow Lite Runtime. +# * TensorFlow Lite Flex Delegate. +ifeq ($(WITH_TFLITE_FLEX), true) + EAGER_CC_ALL_SRCS += $(wildcard tensorflow/core/common_runtime/eager/*.cc) + EAGER_CC_EXCLUDE_SRCS := $(wildcard tensorflow/core/common_runtime/eager/*test.cc) + EAGER_CC_SRCS := $(filter-out $(EAGER_CC_EXCLUDE_SRCS), $(EAGER_CC_ALL_SRCS)) + TF_CC_SRCS += $(EAGER_CC_SRCS) + + TF_LITE_CORE_CC_ALL_SRCS := \ + $(wildcard tensorflow/lite/*.cc) \ + $(wildcard tensorflow/lite/*.c) \ + $(wildcard tensorflow/lite/c/*.c) \ + $(wildcard tensorflow/lite/core/api/*.cc) + + TF_LITE_CORE_CC_ALL_SRCS += \ + $(wildcard tensorflow/lite/kernels/*.cc) \ + $(wildcard tensorflow/lite/kernels/internal/*.cc) \ + $(wildcard tensorflow/lite/kernels/internal/optimized/*.cc) \ + $(wildcard tensorflow/lite/kernels/internal/reference/*.cc) \ + $(PROFILER_SRCS) \ + $(wildcard tensorflow/lite/kernels/*.c) \ + $(wildcard tensorflow/lite/kernels/internal/*.c) \ + $(wildcard tensorflow/lite/kernels/internal/optimized/*.c) \ + $(wildcard tensorflow/lite/kernels/internal/reference/*.c) \ + $(wildcard tensorflow/lite/delegates/flex/*.cc) + + # Hack. This shouldn't be here? + TF_LITE_CORE_CC_ALL_SRCS += \ + $(wildcard tensorflow/contrib/makefile/downloads/farmhash/src/farmhash.cc) \ + + # Remove any duplicates. + TF_LITE_CORE_CC_ALL_SRCS := $(sort $(TF_LITE_CORE_CC_ALL_SRCS)) + TF_LITE_CORE_CC_EXCLUDE_SRCS := \ + $(wildcard tensorflow/lite/*test.cc) \ + $(wildcard tensorflow/lite/*/*test.cc) \ + $(wildcard tensorflow/lite/*/*/*test.cc) \ + $(wildcard tensorflow/lite/*/*/*/*test.cc) \ + $(wildcard tensorflow/lite/kernels/test_util.cc) \ + $(wildcard tensorflow/lite/delegates/flex/test_util.cc) \ + $(wildcard tensorflow/lite/nnapi_delegate.cc) \ + $(wildcard tensorflow/lite/mmap_allocation_disabled.cc) + + # Filter out all the excluded files. + TF_LITE_CC_SRCS := $(filter-out $(TF_LITE_CORE_CC_EXCLUDE_SRCS), $(TF_LITE_CORE_CC_ALL_SRCS)) + TF_CC_SRCS += $(TF_LITE_CC_SRCS) +endif + ifdef HEXAGON_LIBS TF_CC_SRCS += \ tensorflow/cc/framework/scope.cc \ diff --git a/tensorflow/contrib/makefile/README.md b/tensorflow/contrib/makefile/README.md index 6c3b02e12b3082be8bfcc316c4c6122931eb5f76..1293e59cbcba86115e99b505b1f0672a01526462 100644 --- a/tensorflow/contrib/makefile/README.md +++ b/tensorflow/contrib/makefile/README.md @@ -142,7 +142,7 @@ First, download and install JetPack for Android version 3.2 or greater from [Nvi git clone https://github.com/tensorflow/tensorflow.git cd tensorflow JETPACK=$HOME/JetPack_Android_3.2 -TEGRA_LIBS="$JETPACK/cuDNN/aarch64/cuda/lib64/libcudnn.so $JETPACK/cuda-9.0/extras/CUPTI/lib64/libcupti.so $JETPACK/cuda/targets/aarch64-linux-androideabi/lib64/libcufft.so" +TEGRA_LIBS="$JETPACK/cuDNN/aarch64/cuda/lib64/libcudnn.so $JETPACK/cuda/extras/CUPTI/lib64/libcupti.so $JETPACK/cuda/targets/aarch64-linux-androideabi/lib64/libcufft.so" ``` #### Building all CUDA-enabled native binaries: diff --git a/tensorflow/contrib/makefile/build_all_android.sh b/tensorflow/contrib/makefile/build_all_android.sh index fb9e77ae1bcfc3404f1fdf90ab2697a4e79a9836..dc29694449729fe80d072aa06118ba0a8e64ca54 100755 --- a/tensorflow/contrib/makefile/build_all_android.sh +++ b/tensorflow/contrib/makefile/build_all_android.sh @@ -34,7 +34,7 @@ echo "********************************************************************" echo "TensorFlow Lite is the recommended library for mobile and embedded machine learning inference." echo "You are currently using an older version. Please switch over to TensorFlow Lite." echo "" -echo "Link to the code: https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite" +echo "Link to the code: https://github.com/tensorflow/tensorflow/tree/master/tensorflow/lite" echo "********************************************************************" echo "" diff --git a/tensorflow/contrib/makefile/build_all_ios.sh b/tensorflow/contrib/makefile/build_all_ios.sh index 1d4677ef4bd1e8811998d1464e63902544153a49..9a8059ce50041f21d00884896783a02e6285d55c 100755 --- a/tensorflow/contrib/makefile/build_all_ios.sh +++ b/tensorflow/contrib/makefile/build_all_ios.sh @@ -35,11 +35,10 @@ echo "********************************************************************" echo "TensorFlow Lite is the recommended library for mobile and embedded machine learning inference." echo "You are currently using an older version. Please switch over to TensorFlow Lite." echo "" -echo "Link to the code: https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/lite" +echo "Link to the code: https://github.com/tensorflow/tensorflow/tree/master/tensorflow/lite" echo "********************************************************************" echo "" -DEFAULT_ARCH="i386 x86_64 armv7 armv7s arm64" while getopts "a:g:T" opt_name; do case "$opt_name" in a) BUILD_ARCH="${OPTARG}";; @@ -138,7 +137,7 @@ if [[ ! -z "${BUILD_ARCH}" ]]; then fi # build the ios tensorflow libraries. -echo "Building TensorFlow with flags: ${TF_SCRIPT_FLAGS} -f ${TF_CC_FLAGS}" +echo "Building TensorFlow with command: ${TF_SCRIPT_FLAGS} -f ${TF_CC_FLAGS}" tensorflow/contrib/makefile/compile_ios_tensorflow.sh ${TF_SCRIPT_FLAGS} -f "${TF_CC_FLAGS}" # Creates a static universal library in diff --git a/tensorflow/contrib/makefile/build_all_ios_with_tflite.sh b/tensorflow/contrib/makefile/build_all_ios_with_tflite.sh new file mode 100755 index 0000000000000000000000000000000000000000..8d34911f154101d4c0f4a02e69842986056b4b63 --- /dev/null +++ b/tensorflow/contrib/makefile/build_all_ios_with_tflite.sh @@ -0,0 +1,40 @@ +#!/usr/bin/env bash +# 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. +# ============================================================================== + +# This shell script is used to build TensorFlow Lite Flex runtime for iOS. +# It compiles TensorFlow Lite and TensorFlow codebases together, and enable a +# route to use TensorFlow kernels in TensorFlow Lite. +# +# After the script is executed, the multi-architecture static libraries will be +# created under: `tensorflow/contrib/makefile/gen/lib/`. + +set -e + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +TOP_SRCDIR="${SCRIPT_DIR}/../../../" +cd ${TOP_SRCDIR} + +# Exporting `WITH_TFLITE_FLEX`. The flag will be propagated all the way +# down to Makefile. +export WITH_TFLITE_FLEX="true" +# Execute `build_all_ios.sh` and propagate all parameters. +tensorflow/contrib/makefile/build_all_ios.sh $* + +# Copy all the libraries required for TFLite Flex runtime together. +cd "${TOP_SRCDIR}/tensorflow/contrib/makefile" +cp 'downloads/nsync/builds/lipo.ios.c++11/nsync.a' 'gen/lib/' +cp 'gen/protobuf_ios/lib/libprotobuf.a' 'gen/lib/' +cp 'gen/lib/libtensorflow-core.a' 'gen/lib/libtensorflow-lite.a' diff --git a/tensorflow/contrib/makefile/download_dependencies.sh b/tensorflow/contrib/makefile/download_dependencies.sh index dc9b17a62783817ec9a2998c4d5548c0f05e073b..2a5232b476712a96f84be0f4725beb78bc138297 100755 --- a/tensorflow/contrib/makefile/download_dependencies.sh +++ b/tensorflow/contrib/makefile/download_dependencies.sh @@ -30,11 +30,13 @@ EIGEN_URL="$(grep -o 'http.*bitbucket.org/eigen/eigen/get/.*tar\.gz' "${BZL_FILE GEMMLOWP_URL="$(grep -o 'https://mirror.bazel.build/github.com/google/gemmlowp/.*zip' "${BZL_FILE_PATH}" | head -n1)" GOOGLETEST_URL="https://github.com/google/googletest/archive/release-1.8.0.tar.gz" NSYNC_URL="$(grep -o 'https://mirror.bazel.build/github.com/google/nsync/.*tar\.gz' "${BZL_FILE_PATH}" | head -n1)" -# Note: The Protobuf source in `tensorflow/workspace.bzl` in TensorFlow -# 1.10 branch does not work. `make distclean` fails and blocks the build -# process. For now we're hardcoding to the version which is used by -# TensorFlow 1.9. -PROTOBUF_URL="https://mirror.bazel.build/github.com/google/protobuf/archive/396336eb961b75f03b25824fe86cf6490fb75e3a.tar.gz" + +# Note: The protobuf repo needs to be cloned due to its submodules. +# These variables contain the GitHub repo and the sha, from `tensorflow/workspace.bzl`, +# from which to clone it from and checkout to. +readonly PROTOBUF_REPO="https://github.com/protocolbuffers/protobuf.git" +readonly PROTOBUF_TAG="$(grep -o 'https://github.com/protocolbuffers/protobuf/archive/.*tar\.gz' "${BZL_FILE_PATH}" | head -n1 | awk '{print substr($0, index($0, "archive") + 8, index($0, "tar") - index($0, "archive") - 9) }')" + # TODO (yongtang): Replace the following with 'https://mirror.bazel.build/github.com/google/re2/.*tar\.gz' once # the archive has been propagated in mirror.bazel.build. RE2_URL="$(grep -o 'https://github.com/google/re2/.*tar\.gz' "${BZL_FILE_PATH}" | head -n1)" @@ -43,6 +45,10 @@ DOUBLE_CONVERSION_URL="$(grep -o "https.*google/double-conversion.*\.zip" "${BZL ABSL_URL="$(grep -o 'https://github.com/abseil/abseil-cpp/.*tar.gz' "${BZL_FILE_PATH}" | head -n1)" CUB_URL="$(grep -o 'https.*cub/archive.*zip' "${BZL_FILE_PATH}" | grep -v mirror.bazel | head -n1)" +# Required for TensorFlow Lite Flex runtime. +FARMHASH_URL="https://mirror.bazel.build/github.com/google/farmhash/archive/816a4ae622e964763ca0862d9dbd19324a1eaf45.tar.gz" +FLATBUFFERS_URL="https://github.com/google/flatbuffers/archive/1f5eae5d6a135ff6811724f6c57f911d1f46bb15.tar.gz" + # TODO(petewarden): Some new code in Eigen triggers a clang bug with iOS arm64, # so work around it by patching the source. replace_by_sed() { @@ -87,17 +93,46 @@ download_and_extract() { find "${dir}" -type f -name '*BUILD' -delete } +function clone_repository() { + local repo_url="${1}" + local destination_directory="${2}" + local commit_sha="${3}" + + if [[ -d "${destination_directory}" ]]; then + rm -rf "${destination_directory}" + fi + + git clone "${repo_url}" "${destination_directory}" + + pushd "$(pwd)" 1>/dev/null + + cd "${destination_directory}" + + if [[ -n "${commit_sha}" ]]; then + git checkout "${PROTOBUF_TAG}" + fi + + git submodule update --init + + popd 1>/dev/null +} + download_and_extract "${EIGEN_URL}" "${DOWNLOADS_DIR}/eigen" download_and_extract "${GEMMLOWP_URL}" "${DOWNLOADS_DIR}/gemmlowp" download_and_extract "${GOOGLETEST_URL}" "${DOWNLOADS_DIR}/googletest" download_and_extract "${NSYNC_URL}" "${DOWNLOADS_DIR}/nsync" -download_and_extract "${PROTOBUF_URL}" "${DOWNLOADS_DIR}/protobuf" download_and_extract "${RE2_URL}" "${DOWNLOADS_DIR}/re2" download_and_extract "${FFT2D_URL}" "${DOWNLOADS_DIR}/fft2d" download_and_extract "${DOUBLE_CONVERSION_URL}" "${DOWNLOADS_DIR}/double_conversion" download_and_extract "${ABSL_URL}" "${DOWNLOADS_DIR}/absl" download_and_extract "${CUB_URL}" "${DOWNLOADS_DIR}/cub/external/cub_archive" +# Required for TensorFlow Lite Flex runtime. +download_and_extract "${FARMHASH_URL}" "${DOWNLOADS_DIR}/farmhash" +download_and_extract "${FLATBUFFERS_URL}" "${DOWNLOADS_DIR}/flatbuffers" + +clone_repository "${PROTOBUF_REPO}" "${DOWNLOADS_DIR}/protobuf" "${PROTOBUF_TAG}" + replace_by_sed 's#static uint32x4_t p4ui_CONJ_XOR = vld1q_u32( conj_XOR_DATA );#static uint32x4_t p4ui_CONJ_XOR; // = vld1q_u32( conj_XOR_DATA ); - Removed by script#' \ "${DOWNLOADS_DIR}/eigen/Eigen/src/Core/arch/NEON/Complex.h" replace_by_sed 's#static uint32x2_t p2ui_CONJ_XOR = vld1_u32( conj_XOR_DATA );#static uint32x2_t p2ui_CONJ_XOR;// = vld1_u32( conj_XOR_DATA ); - Removed by scripts#' \ diff --git a/tensorflow/contrib/makefile/proto_text_pb_cc_files.txt b/tensorflow/contrib/makefile/proto_text_pb_cc_files.txt index 0d8df93d1116afc1a56f598a9ea776010ae38fd0..87c73ec1ca610cac6d63468887bc350bada5910b 100644 --- a/tensorflow/contrib/makefile/proto_text_pb_cc_files.txt +++ b/tensorflow/contrib/makefile/proto_text_pb_cc_files.txt @@ -27,6 +27,7 @@ tensorflow/core/grappler/costs/op_performance_data.pb.cc tensorflow/core/lib/core/error_codes.pb.cc tensorflow/core/protobuf/cluster.pb.cc tensorflow/core/protobuf/config.pb.cc +tensorflow/core/protobuf/eager_service.pb.cc tensorflow/core/protobuf/debug.pb.cc tensorflow/core/protobuf/device_properties.pb.cc tensorflow/core/protobuf/meta_graph.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 d982df93193268e1eb68a892592a48fbfbc50155..4120ea52ec5255b1efce7a6ce6890fc79c1e4831 100644 --- a/tensorflow/contrib/makefile/proto_text_pb_h_files.txt +++ b/tensorflow/contrib/makefile/proto_text_pb_h_files.txt @@ -29,6 +29,7 @@ tensorflow/core/protobuf/cluster.pb.h tensorflow/core/protobuf/config.pb.h tensorflow/core/protobuf/debug.pb.h tensorflow/core/protobuf/device_properties.pb.h +tensorflow/core/protobuf/eager_service.pb.h tensorflow/core/protobuf/meta_graph.pb.h tensorflow/core/protobuf/named_tensor.pb.h tensorflow/core/protobuf/queue_runner.pb.h diff --git a/tensorflow/contrib/makefile/tf_op_files.txt b/tensorflow/contrib/makefile/tf_op_files.txt index 91af933cfff695c9426fbceebfeb7cbc5eaaf80d..2cd7d6d519a55423a96526b541845392d9ec6bc2 100644 --- a/tensorflow/contrib/makefile/tf_op_files.txt +++ b/tensorflow/contrib/makefile/tf_op_files.txt @@ -42,6 +42,7 @@ tensorflow/core/kernels/conv_grad_filter_ops.cc tensorflow/core/kernels/conv_grad_input_ops.cc tensorflow/core/kernels/conv_grad_ops.cc tensorflow/core/kernels/conv_ops.cc +tensorflow/core/kernels/conv_ops_3d.cc tensorflow/core/kernels/conv_ops_fused.cc tensorflow/core/kernels/conv_ops_using_gemm.cc tensorflow/core/kernels/crop_and_resize_op.cc @@ -52,6 +53,8 @@ tensorflow/core/kernels/cwise_op_add_2.cc tensorflow/core/kernels/cwise_op_bitwise_and.cc tensorflow/core/kernels/cwise_op_bitwise_or.cc tensorflow/core/kernels/cwise_op_bitwise_xor.cc +tensorflow/core/kernels/cwise_op_cos.cc +tensorflow/core/kernels/cwise_op_cosh.cc tensorflow/core/kernels/cwise_op_div.cc tensorflow/core/kernels/cwise_op_equal_to_1.cc tensorflow/core/kernels/cwise_op_equal_to_2.cc @@ -86,10 +89,13 @@ tensorflow/core/kernels/cwise_op_rsqrt.cc tensorflow/core/kernels/cwise_op_select.cc tensorflow/core/kernels/cwise_op_sigmoid.cc tensorflow/core/kernels/cwise_op_sign.cc +tensorflow/core/kernels/cwise_op_sin.cc +tensorflow/core/kernels/cwise_op_sinh.cc tensorflow/core/kernels/cwise_op_sqrt.cc tensorflow/core/kernels/cwise_op_square.cc tensorflow/core/kernels/cwise_op_squared_difference.cc tensorflow/core/kernels/cwise_op_sub.cc +tensorflow/core/kernels/cwise_op_tan.cc tensorflow/core/kernels/cwise_op_tanh.cc tensorflow/core/kernels/cwise_op_xdivy.cc tensorflow/core/kernels/cwise_op_xlogy.cc @@ -113,6 +119,7 @@ tensorflow/core/kernels/fake_quant_ops.cc tensorflow/core/kernels/fifo_queue.cc tensorflow/core/kernels/fifo_queue_op.cc tensorflow/core/kernels/fill_functor.cc +tensorflow/core/kernels/fft_ops.cc tensorflow/core/kernels/function_ops.cc tensorflow/core/kernels/fused_batch_norm_op.cc tensorflow/core/kernels/gather_functor.cc @@ -151,14 +158,15 @@ tensorflow/core/kernels/mirror_pad_op_cpu_impl_2.cc tensorflow/core/kernels/mirror_pad_op_cpu_impl_3.cc tensorflow/core/kernels/mirror_pad_op_cpu_impl_4.cc tensorflow/core/kernels/mirror_pad_op_cpu_impl_5.cc +tensorflow/core/kernels/multinomial_op.cc tensorflow/core/kernels/no_op.cc tensorflow/core/kernels/non_max_suppression_op.cc tensorflow/core/kernels/one_hot_op.cc -tensorflow/core/kernels/ops_util.cc tensorflow/core/kernels/pack_op.cc tensorflow/core/kernels/pad_op.cc tensorflow/core/kernels/padding_fifo_queue.cc tensorflow/core/kernels/padding_fifo_queue_op.cc +tensorflow/core/kernels/pooling_ops_3d.cc tensorflow/core/kernels/pooling_ops_common.cc tensorflow/core/kernels/population_count_op.cc tensorflow/core/kernels/quantization_utils.cc @@ -244,7 +252,9 @@ tensorflow/core/kernels/spectrogram_op.cc tensorflow/core/kernels/split_lib_cpu.cc tensorflow/core/kernels/split_op.cc tensorflow/core/kernels/split_v_op.cc +tensorflow/core/kernels/stack.cc tensorflow/core/kernels/stack_ops.cc +tensorflow/core/kernels/stateless_random_ops.cc tensorflow/core/kernels/strided_slice_op.cc tensorflow/core/kernels/strided_slice_op_inst_0.cc tensorflow/core/kernels/strided_slice_op_inst_1.cc diff --git a/tensorflow/contrib/makefile/tf_proto_files.txt b/tensorflow/contrib/makefile/tf_proto_files.txt index 8bec3e3e01f7913944cda8dafc5f1993e96570bd..2712e906d719e72dacb60e213205ad68895f905f 100644 --- a/tensorflow/contrib/makefile/tf_proto_files.txt +++ b/tensorflow/contrib/makefile/tf_proto_files.txt @@ -34,6 +34,7 @@ tensorflow/core/lib/core/error_codes.proto tensorflow/core/protobuf/cluster.proto tensorflow/core/protobuf/config.proto tensorflow/core/protobuf/debug.proto +tensorflow/core/protobuf/eager_service.proto tensorflow/core/protobuf/device_properties.proto tensorflow/core/protobuf/meta_graph.proto tensorflow/core/protobuf/named_tensor.proto diff --git a/tensorflow/contrib/metrics/python/metrics/classification.py b/tensorflow/contrib/metrics/python/metrics/classification.py index 7053907da05b487df73481e3ced269bb69b8deae..9aabc4bec3053871e3ff6cd3a88fd76d293f48cc 100644 --- a/tensorflow/contrib/metrics/python/metrics/classification.py +++ b/tensorflow/contrib/metrics/python/metrics/classification.py @@ -18,13 +18,13 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +from tensorflow.python.distribute import distribution_strategy_context from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import math_ops from tensorflow.python.ops import metrics_impl from tensorflow.python.ops import variable_scope -from tensorflow.python.training import distribution_strategy_context # TODO(nsilberman): move into metrics/python/ops/ @@ -167,15 +167,15 @@ def f1_score(labels, predictions, weights=None, num_thresholds=200, (precision_at_t + recall_at_t + epsilon)) return math_ops.reduce_max(f1_at_thresholds) - def f1_across_towers(_, values): + def f1_across_replicas(_, values): best_f1 = compute_best_f1_score(tp=values['tp'], fp=values['fp'], fn=values['fn'], name='value') if metrics_collections: ops.add_to_collections(metrics_collections, best_f1) return best_f1 - best_f1 = distribution_strategy_context.get_tower_context().merge_call( - f1_across_towers, values) + best_f1 = distribution_strategy_context.get_replica_context().merge_call( + f1_across_replicas, args=(values,)) update_op = compute_best_f1_score(tp=update_ops['tp'], fp=update_ops['fp'], fn=update_ops['fn'], name='update') diff --git a/tensorflow/contrib/metrics/python/metrics/classification_test.py b/tensorflow/contrib/metrics/python/metrics/classification_test.py index d6a670f97b32a29129cb9ea0cd71c5a2b7597a47..e789d2cb9dfbac7b1e145be48b3f707af3fd4e18 100644 --- a/tensorflow/contrib/metrics/python/metrics/classification_test.py +++ b/tensorflow/contrib/metrics/python/metrics/classification_test.py @@ -291,12 +291,11 @@ class F1ScoreTest(test.TestCase): labels = labels.astype(np.float32) predictions = predictions.astype(np.float32) - tf_predictions, tf_labels = (dataset_ops.Dataset - .from_tensor_slices((predictions, labels)) - .repeat() - .batch(batch_size) - .make_one_shot_iterator() - .get_next()) + tf_predictions, tf_labels = dataset_ops.make_one_shot_iterator( + dataset_ops.Dataset + .from_tensor_slices((predictions, labels)) + .repeat() + .batch(batch_size)).get_next() f1, f1_op = classification.f1_score(tf_labels, tf_predictions, num_thresholds=3) diff --git a/tensorflow/contrib/metrics/python/ops/metric_ops.py b/tensorflow/contrib/metrics/python/ops/metric_ops.py index d6932f6e4b603b1a76250ab622f5fe8eaea81bc9..7b432f8bd20989c6d95310bcaca88d44ce3e0d1f 100644 --- a/tensorflow/contrib/metrics/python/ops/metric_ops.py +++ b/tensorflow/contrib/metrics/python/ops/metric_ops.py @@ -24,7 +24,6 @@ from __future__ import print_function import collections as collections_lib -from tensorflow.python.compat import compat from tensorflow.python.eager import context from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops @@ -46,32 +45,6 @@ from tensorflow.python.util.deprecation import deprecated _EPSILON = 1e-7 -def _safe_div(numerator, denominator): - """Computes a safe divide which returns 0 if the denominator is zero. - - Note that the function contains an additional conditional check that is - necessary for avoiding situations where the loss is zero causing NaNs to - creep into the gradient computation. - - Args: - numerator: An arbitrary `Tensor`. - denominator: A `Tensor` whose shape matches `numerator` and whose values are - assumed to be non-negative. - - Returns: - The element-wise value of the numerator divided by the denominator. - """ - if compat.forward_compatible(2018, 11, 1): - return math_ops.div_no_nan(numerator, denominator) - return array_ops.where( - math_ops.greater(denominator, 0), - math_ops.div(numerator, - array_ops.where( - math_ops.equal(denominator, 0), - array_ops.ones_like(denominator), denominator)), - array_ops.zeros_like(numerator)) - - @deprecated(None, 'Please switch to tf.metrics.true_positives. Note that the ' 'order of the labels and predictions arguments has been switched.') def streaming_true_positives(predictions, @@ -3247,24 +3220,20 @@ def streaming_covariance(predictions, # We update the means by Delta=Error*BatchCount/(BatchCount+PrevCount) # batch_mean_prediction is E[x_B] in the update equation - batch_mean_prediction = _safe_div( - math_ops.reduce_sum(weighted_predictions), - batch_count) - delta_mean_prediction = _safe_div( - (batch_mean_prediction - mean_prediction) * batch_count, - update_count) + batch_mean_prediction = math_ops.div_no_nan( + math_ops.reduce_sum(weighted_predictions), batch_count) + delta_mean_prediction = math_ops.div_no_nan( + (batch_mean_prediction - mean_prediction) * batch_count, update_count) update_mean_prediction = state_ops.assign_add(mean_prediction, delta_mean_prediction) # prev_mean_prediction is E[x_A] in the update equation prev_mean_prediction = update_mean_prediction - delta_mean_prediction # batch_mean_label is E[y_B] in the update equation - batch_mean_label = _safe_div( - math_ops.reduce_sum(weighted_labels), - batch_count) - delta_mean_label = _safe_div( - (batch_mean_label - mean_label) * batch_count, - update_count) + batch_mean_label = math_ops.div_no_nan( + math_ops.reduce_sum(weighted_labels), batch_count) + delta_mean_label = math_ops.div_no_nan( + (batch_mean_label - mean_label) * batch_count, update_count) update_mean_label = state_ops.assign_add(mean_label, delta_mean_label) # prev_mean_label is E[y_A] in the update equation prev_mean_label = update_mean_label - delta_mean_label @@ -3447,7 +3416,7 @@ def streaming_mean_cosine_distance(predictions, predictions.get_shape().assert_is_compatible_with(labels.get_shape()) radial_diffs = math_ops.multiply(predictions, labels) radial_diffs = math_ops.reduce_sum( - radial_diffs, reduction_indices=[ + radial_diffs, axis=[ dim, ], keepdims=True) mean_distance, update_op = streaming_mean(radial_diffs, weights, None, None, @@ -3926,9 +3895,8 @@ def cohen_kappa(labels, po_sum = math_ops.reduce_sum(po) total = math_ops.reduce_sum(pe_row) pe_sum = math_ops.reduce_sum( - _safe_div( - math_ops.to_double(pe_row * pe_col), - math_ops.to_double(total))) + math_ops.div_no_nan( + math_ops.to_double(pe_row * pe_col), math_ops.to_double(total))) po_sum, pe_sum, total = (math_ops.to_double(po_sum), math_ops.to_double(pe_sum), math_ops.to_double(total)) diff --git a/tensorflow/contrib/mixed_precision/python/loss_scale_manager_test.py b/tensorflow/contrib/mixed_precision/python/loss_scale_manager_test.py index 1b0383d24c0c472b4875d15c3650e37dfd2439e1..c922d0cd11fda3c51a51ceccf69798df7ce75f26 100644 --- a/tensorflow/contrib/mixed_precision/python/loss_scale_manager_test.py +++ b/tensorflow/contrib/mixed_precision/python/loss_scale_manager_test.py @@ -29,7 +29,7 @@ from tensorflow.python.platform import test def _GetExampleIter(inputs): dataset = dataset_ops.Dataset.from_tensor_slices(inputs) - return dataset.make_one_shot_iterator() + return dataset_ops.make_one_shot_iterator(dataset) class FixedLossScaleManagerTest(test.TestCase): diff --git a/tensorflow/contrib/mixed_precision/python/loss_scale_optimizer_test.py b/tensorflow/contrib/mixed_precision/python/loss_scale_optimizer_test.py index 9009df0eefec13146090ba5fc2096e71ba6eb89d..33f9a43e803ea845a25bba284e41e5a0e6228dad 100644 --- a/tensorflow/contrib/mixed_precision/python/loss_scale_optimizer_test.py +++ b/tensorflow/contrib/mixed_precision/python/loss_scale_optimizer_test.py @@ -132,7 +132,7 @@ class LossScaleOptimizerTest(test.TestCase): x = variable_scope.get_variable("x", initializer=1., dtype=dtypes.float32) dataset = dataset_ops.Dataset.from_tensor_slices([np.nan, np.inf, 0.1]) - itr = dataset.make_one_shot_iterator() + itr = dataset_ops.make_one_shot_iterator(dataset) lr = 1 opt = gd.GradientDescentOptimizer(lr) @@ -182,7 +182,7 @@ class LossScaleOptimizerTest(test.TestCase): x = variable_scope.get_variable("x", initializer=1., dtype=dtypes.float32) dataset = dataset_ops.Dataset.from_tensor_slices([np.nan, np.inf, 0.1]) - itr = dataset.make_one_shot_iterator() + itr = dataset_ops.make_one_shot_iterator(dataset) lr = 1 init_loss_scale = 8 diff --git a/tensorflow/contrib/model_pruning/README.md b/tensorflow/contrib/model_pruning/README.md index b313024e2852caf2385454771b289ad0162cc463..710a262f33872ada8d090d796f80dc06c2a27f84 100644 --- a/tensorflow/contrib/model_pruning/README.md +++ b/tensorflow/contrib/model_pruning/README.md @@ -51,9 +51,8 @@ The pruning library allows for specification of the following hyper parameters: | begin_pruning_step | integer | 0 | The global step at which to begin pruning | | end_pruning_step | integer | -1 | The global step at which to terminate pruning. Defaults to -1 implying that pruning continues till the training stops | | weight_sparsity_map | list of strings | [""] | list of weight variable name (or layer name):target sparsity pairs. Eg. [conv1:0.9,conv2/kernel:0.8]. For layers/weights not in this list, sparsity as specified by the target_sparsity hyperparameter is used. | -| threshold_decay | float | 0.9 | The decay factor to use for exponential decay of the thresholds | +| threshold_decay | float | 0.0 | The decay factor to use for exponential decay of the thresholds | | pruning_frequency | integer | 10 | How often should the masks be updated? (in # of global_steps) | -| nbins | integer | 256 | Number of bins to use for histogram computation. Note: When running on TPUs, a large (>1024) value for `nbins` may adversely affect the training time. | | block_height|integer | 1 | Number of rows in a block for block sparse matrices| | block_width |integer | 1 | Number of cols in a block for block sparse matrices| | block_pooling_function| string | AVG | The function to use to pool weight values in a block: average (AVG) or max (MAX)| diff --git a/tensorflow/contrib/model_pruning/python/layers/core_layers.py b/tensorflow/contrib/model_pruning/python/layers/core_layers.py index 764ab620bc2227ff5e8e3f473d689e0e133e83d4..1fa5c8cb485704a5fccc486e823bbc4050bf505a 100644 --- a/tensorflow/contrib/model_pruning/python/layers/core_layers.py +++ b/tensorflow/contrib/model_pruning/python/layers/core_layers.py @@ -21,6 +21,7 @@ from __future__ import print_function from tensorflow.python.framework import ops from tensorflow.python.framework import tensor_shape +from tensorflow.python.keras.engine import input_spec from tensorflow.python.layers import base from tensorflow.python.layers import utils from tensorflow.python.ops import array_ops @@ -119,15 +120,15 @@ class _MaskedConv(base.Layer): self.bias_initializer = bias_initializer self.kernel_regularizer = kernel_regularizer self.bias_regularizer = bias_regularizer - self.input_spec = base.InputSpec(ndim=self.rank + 2) + self.input_spec = input_spec.InputSpec(ndim=self.rank + 2) def build(self, input_shape): input_shape = tensor_shape.TensorShape(input_shape) channel_axis = 1 if self.data_format == 'channels_first' else -1 - if input_shape[channel_axis].value is None: + if tensor_shape.dimension_value(input_shape[channel_axis]) is None: raise ValueError('The channel dimension of the inputs ' 'should be defined. Found `None`.') - input_dim = input_shape[channel_axis].value + input_dim = tensor_shape.dimension_value(input_shape[channel_axis]) kernel_shape = self.kernel_size + (input_dim, self.filters) self.mask = self.add_variable( name='mask', @@ -171,7 +172,7 @@ class _MaskedConv(base.Layer): dtype=self.dtype) else: self.bias = None - self.input_spec = base.InputSpec( + self.input_spec = input_spec.InputSpec( ndim=self.rank + 2, axes={channel_axis: input_dim}) self.built = True @@ -393,19 +394,19 @@ class MaskedFullyConnected(base.Layer): self.bias_initializer = bias_initializer self.kernel_regularizer = kernel_regularizer self.bias_regularizer = bias_regularizer - self.input_spec = base.InputSpec(min_ndim=2) + self.input_spec = input_spec.InputSpec(min_ndim=2) def build(self, input_shape): input_shape = tensor_shape.TensorShape(input_shape) - if input_shape[-1].value is None: + if tensor_shape.dimension_value(input_shape[-1]) is None: raise ValueError('The last dimension of the inputs to `Dense` ' 'should be defined. Found `None`.') - self.input_spec = base.InputSpec( - min_ndim=2, axes={-1: input_shape[-1].value}) + self.input_spec = input_spec.InputSpec( + min_ndim=2, axes={-1: tensor_shape.dimension_value(input_shape[-1])}) self.kernel = self.add_variable( 'kernel', - shape=[input_shape[-1].value, self.units], + shape=[tensor_shape.dimension_value(input_shape[-1]), self.units], initializer=self.kernel_initializer, regularizer=self.kernel_regularizer, dtype=self.dtype, @@ -413,7 +414,7 @@ class MaskedFullyConnected(base.Layer): self.mask = self.add_variable( name='mask', - shape=[input_shape[-1].value, self.units], + shape=[tensor_shape.dimension_value(input_shape[-1]), self.units], initializer=init_ops.ones_initializer(), trainable=False, dtype=self.dtype) @@ -470,7 +471,7 @@ class MaskedFullyConnected(base.Layer): def compute_output_shape(self, input_shape): input_shape = tensor_shape.TensorShape(input_shape) input_shape = input_shape.with_rank_at_least(2) - if input_shape[-1].value is None: + if tensor_shape.dimension_value(input_shape[-1]) is None: raise ValueError( 'The innermost dimension of input_shape must be defined, but saw: %s' % input_shape) diff --git a/tensorflow/contrib/model_pruning/python/layers/rnn_cells.py b/tensorflow/contrib/model_pruning/python/layers/rnn_cells.py index 5f6c6aea74f2965ccfe552a58cde290b5506ef12..2959019d6d8eac489e0cc5ece61ea59ce725d604 100644 --- a/tensorflow/contrib/model_pruning/python/layers/rnn_cells.py +++ b/tensorflow/contrib/model_pruning/python/layers/rnn_cells.py @@ -94,7 +94,7 @@ class MaskedBasicLSTMCell(tf_rnn.BasicLSTMCell): self.built = False - input_depth = inputs_shape[1].value + input_depth = inputs_shape.dims[1].value h_depth = self._num_units self._mask = self.add_variable( name="mask", @@ -243,7 +243,7 @@ class MaskedLSTMCell(tf_rnn.LSTMCell): self.built = False - input_depth = inputs_shape[1].value + input_depth = inputs_shape.dims[1].value h_depth = self._num_units self._mask = self.add_variable( name="mask", @@ -304,7 +304,7 @@ class MaskedLSTMCell(tf_rnn.LSTMCell): c_prev = array_ops.slice(state, [0, 0], [-1, self._num_units]) m_prev = array_ops.slice(state, [0, self._num_units], [-1, num_proj]) - input_size = inputs.get_shape().with_rank(2)[1] + 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]") diff --git a/tensorflow/contrib/model_pruning/python/pruning.py b/tensorflow/contrib/model_pruning/python/pruning.py index d2b811641764df05c66654dfcb044fa7e78853a5..9966f7cf798d206fffbaeb4d16b6500a90d113e4 100644 --- a/tensorflow/contrib/model_pruning/python/pruning.py +++ b/tensorflow/contrib/model_pruning/python/pruning.py @@ -204,7 +204,7 @@ def get_pruning_hparams(): begin_pruning_step=0, end_pruning_step=-1, weight_sparsity_map=[''], - threshold_decay=0.9, + threshold_decay=0.0, pruning_frequency=10, nbins=256, block_height=1, @@ -214,7 +214,7 @@ def get_pruning_hparams(): target_sparsity=0.5, sparsity_function_begin_step=0, sparsity_function_end_step=100, - sparsity_function_exponent=3, + sparsity_function_exponent=3.0, use_tpu=False) @@ -397,28 +397,26 @@ class Pruning(object): raise ValueError('Sparsity variable undefined') sparsity = self._get_sparsity(weights.op.name) - with ops.name_scope(weights.op.name + '_pruning_ops'): abs_weights = math_ops.abs(weights) - max_value = math_ops.reduce_max(abs_weights) - cdf_fn = pruning_utils.compute_cdf_from_histogram - if self._spec.use_tpu: - cdf_fn = pruning_utils.compute_cdf - - norm_cdf = cdf_fn(abs_weights, [0.0, max_value], nbins=self._spec.nbins) - current_threshold = math_ops.multiply( - math_ops.div( - math_ops.reduce_sum( - math_ops.cast( - math_ops.less(norm_cdf, sparsity), dtypes.float32)), - float(self._spec.nbins)), max_value) - + k = math_ops.cast( + math_ops.round( + math_ops.cast(array_ops.size(abs_weights), dtypes.float32) * + (1 - sparsity)), dtypes.int32) + # Sort the entire array + values, _ = nn_ops.top_k( + array_ops.reshape(abs_weights, [-1]), k=array_ops.size(abs_weights)) + # Grab the (k-1) th value + current_threshold = array_ops.gather(values, k - 1) smoothed_threshold = math_ops.add_n([ math_ops.multiply(current_threshold, 1 - self._spec.threshold_decay), math_ops.multiply(threshold, self._spec.threshold_decay) ]) + new_mask = math_ops.cast( - math_ops.greater(abs_weights, smoothed_threshold), dtypes.float32) + math_ops.greater_equal(abs_weights, smoothed_threshold), + dtypes.float32) + return smoothed_threshold, new_mask def _maybe_update_block_mask(self, weights, threshold): @@ -456,13 +454,14 @@ class Pruning(object): pool_window = [self._block_dim[0], self._block_dim[1]] pool_fn = pruning_utils.factorized_pool - + squeeze_axis = None if not self._spec.use_tpu: pool_fn = nn_ops.pool abs_weights = array_ops.reshape( abs_weights, [1, abs_weights.get_shape()[0], abs_weights.get_shape()[1], 1]) + squeeze_axis = [0, 3] pooled_weights = pool_fn( abs_weights, @@ -473,7 +472,7 @@ class Pruning(object): name=weights.op.name + '_pooled') if pooled_weights.get_shape().ndims != 2: - pooled_weights = array_ops.squeeze(pooled_weights) + pooled_weights = array_ops.squeeze(pooled_weights, axis=squeeze_axis) smoothed_threshold, new_mask = self._update_mask(pooled_weights, threshold) diff --git a/tensorflow/contrib/model_pruning/python/pruning_test.py b/tensorflow/contrib/model_pruning/python/pruning_test.py index 1b6da5ce2b4ebb3ea3b204c4ed12bed8db951447..835614d8822147dadb029107ae0e917cc955eef0 100644 --- a/tensorflow/contrib/model_pruning/python/pruning_test.py +++ b/tensorflow/contrib/model_pruning/python/pruning_test.py @@ -102,7 +102,7 @@ class PruningTest(test.TestCase): weights = variables.VariableV1( math_ops.linspace(1.0, 100.0, 100), name="weights") masked_weights = pruning.apply_mask(weights) - sparsity = variables.VariableV1(0.5, name="sparsity") + sparsity = variables.VariableV1(0.95, name="sparsity") p = pruning.Pruning(sparsity=sparsity) p._spec.threshold_decay = 0.0 mask_update_op = p.mask_update_op() @@ -111,7 +111,7 @@ class PruningTest(test.TestCase): self.assertAllEqual(np.count_nonzero(masked_weights_val), 100) session.run(mask_update_op) masked_weights_val = masked_weights.eval() - self.assertAllEqual(np.count_nonzero(masked_weights_val), 50) + self.assertAllEqual(np.count_nonzero(masked_weights_val), 5) def _blockMasking(self, hparams, weights, expected_mask): diff --git a/tensorflow/contrib/model_pruning/python/pruning_utils.py b/tensorflow/contrib/model_pruning/python/pruning_utils.py index 91b0bb7f6003c047e4dcf342695f433edbc11614..8f2ba036469bd02328a831a3d1de2ffbd10f5004 100644 --- a/tensorflow/contrib/model_pruning/python/pruning_utils.py +++ b/tensorflow/contrib/model_pruning/python/pruning_utils.py @@ -25,16 +25,12 @@ 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 clip_ops from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import init_ops -from tensorflow.python.ops import math_ops from tensorflow.python.ops import nn_ops from tensorflow.python.ops import state_ops from tensorflow.python.ops import variable_scope -_NBINS = 256 - def weight_mask_variable(var, scope): """Create a mask for the weights. @@ -165,130 +161,6 @@ def expand_tensor(tensor, block_dims): return expanded_tensor -def _histogram(values, value_range, nbins=100, dtype=dtypes.int32, name=None): - """Return histogram of values. - - Given the tensor `values`, this operation returns a rank 1 histogram counting - the number of entries in `values` that fell into every bin. The bins are - equal width and determined by the arguments `value_range` and `nbins`. - - Args: - values: Numeric `Tensor`. - value_range: Shape [2] `Tensor` of same `dtype` as `values`. - values <= value_range[0] will be mapped to hist[0], - values >= value_range[1] will be mapped to hist[-1]. - nbins: Scalar `int32 Tensor`. Number of histogram bins. - dtype: dtype for returned histogram. - name: A name for this operation (defaults to 'histogram'). - - Returns: - A 1-D `Tensor` holding histogram of values. - - """ - with ops.name_scope(name, 'histogram', [values, value_range, nbins]) as scope: - values = ops.convert_to_tensor(values, name='values') - values = array_ops.reshape(values, [-1]) - value_range = ops.convert_to_tensor(value_range, name='value_range') - nbins_float = np.float32(nbins) - - # Map tensor values that fall within value_range to [0, 1]. - scaled_values = math_ops.truediv( - values - value_range[0], - value_range[1] - value_range[0], - name='scaled_values') - - # map tensor values within the open interval value_range to {0,.., nbins-1}, - # values outside the open interval will be zero or less, or nbins or more. - indices = math_ops.floor(nbins_float * scaled_values, name='indices') - - # Clip edge cases (e.g. value = value_range[1]) or "outliers." - indices = math_ops.cast( - clip_ops.clip_by_value(indices, 0, nbins_float - 1), dtypes.int32) - - return math_ops.unsorted_segment_sum( - array_ops.ones_like(indices, dtype=dtype), indices, nbins, name=scope) - - -def compute_cdf_from_histogram(values, value_range, **kwargs): - """Returns the normalized cumulative distribution of the given values tensor. - - Computes the histogram and uses tf.cumsum to arrive at cdf - - Args: - values: Numeric `Tensor`. - value_range: Shape [2] `Tensor` of same `dtype` as `values`. - **kwargs: keyword arguments: nbins, name - - Returns: - A 1-D `Tensor` holding normalized cdf of values. - - """ - nbins = kwargs.get('nbins', _NBINS) - name = kwargs.get('name', None) - with ops.name_scope(name, 'cdf', [values, value_range, nbins]): - histogram = _histogram( - values, value_range, dtype=dtypes.float32, nbins=nbins) - cdf = math_ops.cumsum(histogram) - return math_ops.div(cdf, math_ops.reduce_max(cdf)) - - -def compute_cdf(values, value_range, **kwargs): - """Returns the normalized cumulative distribution of the given values tensor. - - Uses tf.while_loop to directly compute the cdf of the values. - - Args: - values: Numeric `Tensor`. - value_range: Shape [2] `Tensor` of same `dtype` as `values` - **kwargs: keyword arguments: nbins, name - - Returns: - A 1-D `Tensor` holding normalized cdf of values. - - """ - nbins = kwargs.get('nbins', _NBINS) - name = kwargs.get('name', None) - with ops.name_scope(name, 'cdf', [values, value_range, nbins]): - values = ops.convert_to_tensor(values, name='values') - value_range = ops.convert_to_tensor(value_range, name='value_range') - nbins_float = np.float32(nbins) - - # Map tensor values that fall within value_range to [0, 1]. - scaled_values = math_ops.truediv( - values - value_range[0], - value_range[1] - value_range[0], - name='scaled_values') - - # map tensor values within the open interval value_range to {0,.., nbins-1}, - # values outside the open interval will be zero or less, or nbins or more. - indices = math_ops.floor(nbins_float * scaled_values, name='indices') - - # Clip edge cases (e.g. value = value_range[1]) or "outliers." - indices = math_ops.cast( - clip_ops.clip_by_value(indices, 0, nbins_float - 1), dtypes.int32) - - cdf = array_ops.zeros(nbins) - i = constant_op.constant(0) - - def loop_cond(loop_count, _): - return math_ops.less(loop_count, nbins) - - def loop_body(loop_count, cdf): - temp = math_ops.reduce_sum( - math_ops.cast( - math_ops.less_equal(indices, loop_count), dtypes.float32)) - cdf = math_ops.add( - cdf, - array_ops.one_hot( - loop_count, depth=nbins, on_value=temp, off_value=0.0)) - return [loop_count + 1, cdf] - - _, cdf = control_flow_ops.while_loop( - loop_cond, loop_body, [i, cdf], maximum_iterations=nbins) - - return math_ops.div(cdf, math_ops.reduce_max(cdf)) - - def factorized_pool(input_tensor, window_shape, pooling_type, @@ -336,7 +208,7 @@ def factorized_pool(input_tensor, padding=padding) return array_ops.squeeze( - array_ops.transpose(width_pooling, perm=[0, 1, 3, 2])) + array_ops.transpose(width_pooling, perm=[0, 1, 3, 2]), axis=[0, 1]) def determine_partitioned_axis(partitioned_variable): diff --git a/tensorflow/contrib/model_pruning/python/pruning_utils_test.py b/tensorflow/contrib/model_pruning/python/pruning_utils_test.py index 0aca843497611552d922715514118cac003c29b2..b85bc413155d53cd6d53e98dae0ad626531f61eb 100644 --- a/tensorflow/contrib/model_pruning/python/pruning_utils_test.py +++ b/tensorflow/contrib/model_pruning/python/pruning_utils_test.py @@ -19,13 +19,9 @@ from __future__ import division from __future__ import print_function from absl.testing import parameterized -import numpy as np from tensorflow.contrib.model_pruning.python import pruning_utils -from tensorflow.python.framework import constant_op 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_ops from tensorflow.python.ops import random_ops from tensorflow.python.ops import variable_scope @@ -33,60 +29,29 @@ from tensorflow.python.ops import variables from tensorflow.python.platform import test -class PruningUtilsTest(test.TestCase): - - def _compare_cdf(self, values): - abs_values = math_ops.abs(values) - max_value = math_ops.reduce_max(abs_values) - with self.cached_session(): - variables.global_variables_initializer().run() - cdf_from_histogram = pruning_utils.compute_cdf_from_histogram( - abs_values, [0.0, max_value], nbins=pruning_utils._NBINS) - cdf = pruning_utils.compute_cdf(abs_values, [0.0, max_value]) - self.assertAllEqual(cdf.eval(), cdf_from_histogram.eval()) - - def testHistogram(self): - width = 10 - height = 10 - nbins = 100 - expected_histogram = np.full(nbins, 1.0) - init = init_ops.constant_initializer(np.linspace(0.0, 1.0, width * height)) - weights = variable_scope.get_variable( - "weights", [width, height], initializer=init) - histogram = pruning_utils._histogram( - weights, [0, 1.0], nbins, dtype=np.float32) - with self.cached_session(): - variables.global_variables_initializer().run() - computed_histogram = histogram.eval() - self.assertAllEqual(expected_histogram, computed_histogram) - - def testCDF(self): - nbins = 5 - weights = constant_op.constant([-1, 0, 1, 1.5, 2, 3, 4, 5, 10, 100]) - abs_weights = math_ops.abs(weights) - norm_cdf = pruning_utils.compute_cdf_from_histogram( - abs_weights, [0.0, 5.0], nbins=nbins) - expected_cdf = np.array([0.1, 0.4, 0.5, 0.6, 1.0], dtype=np.float32) - with self.cached_session() as sess: - variables.global_variables_initializer().run() - norm_cdf_val = sess.run(norm_cdf) - self.assertAllEqual(len(norm_cdf_val), nbins) - self.assertAllEqual(expected_cdf, norm_cdf_val) - - def testCDFEquivalence2D(self): - width = 100 - height = 100 - weights = variable_scope.get_variable("weights", shape=[width, height]) - self._compare_cdf(weights) - - def testCDFEquivalence4D(self): - weights = variable_scope.get_variable("weights", shape=[5, 5, 128, 128]) - self._compare_cdf(weights) - - @parameterized.named_parameters( - ("1x1", [1, 1]), ("4x4", [4, 4]), ("6x6", [6, 6]), ("1x4", [1, 4]), - ("4x1", [4, 1]), ("1x8", [1, 8]), ("8x1", [8, 1])) + ("Input_32x32_block_1x1", [32, 32], [1, 1]), + # block size 6x6 + ("Input_3x3_block_6x6", [3, 3], [6, 6]), + ("Input_32x32_block_6x6", [32, 32], [6, 6]), + ("Input_2x32_block_6x6", [2, 32], [6, 6]), + ("Input_32x2_block_6x6", [32, 2], [6, 6]), + ("Input_30x30_block_6x6", [30, 30], [6, 6]), + # block size 4x4 + ("Input_32x32_block_4x4", [32, 32], [4, 4]), + ("Input_2x32_block_4x4", [2, 32], [4, 4]), + ("Input_32x2_block_4x4", [32, 2], [4, 4]), + ("Input_30x30_block_4x4", [30, 30], [4, 4]), + # block size 1x4 + ("Input_32x32_block_1x4", [32, 32], [1, 4]), + ("Input_2x32_block_1x4", [2, 32], [1, 4]), + ("Input_32x2_block_1x4", [32, 2], [1, 4]), + ("Input_30x30_block_1x4", [30, 30], [1, 4]), + # block size 4x1 + ("Input_32x32_block_4x1", [32, 32], [4, 1]), + ("Input_2x32_block_4x1", [2, 32], [4, 1]), + ("Input_32x2_block_4x1", [32, 2], [4, 1]), + ("Input_30x30_block_4x1", [30, 30], [4, 1])) class PruningUtilsParameterizedTest(test.TestCase, parameterized.TestCase): def _compare_pooling_methods(self, weights, pooling_kwargs): @@ -97,9 +62,11 @@ class PruningUtilsParameterizedTest(test.TestCase, parameterized.TestCase): array_ops.reshape( weights, [1, weights.get_shape()[0], - weights.get_shape()[1], 1]), **pooling_kwargs)) + weights.get_shape()[1], 1]), **pooling_kwargs), + axis=[0, 3]) pooled_weights_factorized_pool = pruning_utils.factorized_pool( weights, **pooling_kwargs) + self.assertAllClose(pooled_weights_tf.eval(), pooled_weights_factorized_pool.eval()) @@ -113,8 +80,8 @@ class PruningUtilsParameterizedTest(test.TestCase, parameterized.TestCase): [expanded_tensor, kronecker_product]) self.assertAllEqual(expanded_tensor_val, kronecker_product_val) - def testFactorizedAvgPool(self, window_shape): - weights = variable_scope.get_variable("weights", shape=[1024, 2048]) + def testFactorizedAvgPool(self, input_shape, window_shape): + weights = variable_scope.get_variable("weights", shape=input_shape) pooling_kwargs = { "window_shape": window_shape, "pooling_type": "AVG", @@ -123,8 +90,8 @@ class PruningUtilsParameterizedTest(test.TestCase, parameterized.TestCase): } self._compare_pooling_methods(weights, pooling_kwargs) - def testFactorizedMaxPool(self, window_shape): - weights = variable_scope.get_variable("weights", shape=[1024, 2048]) + def testFactorizedMaxPool(self, input_shape, window_shape): + weights = variable_scope.get_variable("weights", shape=input_shape) pooling_kwargs = { "window_shape": window_shape, "pooling_type": "MAX", @@ -133,8 +100,8 @@ class PruningUtilsParameterizedTest(test.TestCase, parameterized.TestCase): } self._compare_pooling_methods(weights, pooling_kwargs) - def testExpandTensor(self, block_dim): - weights = random_ops.random_normal(shape=[1024, 512]) + def testExpandTensor(self, input_shape, block_dim): + weights = random_ops.random_normal(shape=input_shape) self._compare_expand_tensor_with_kronecker_product(weights, block_dim) 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/nccl/BUILD b/tensorflow/contrib/nccl/BUILD deleted file mode 100644 index 9a9d4802608dc94bf70082e1585de3890a5dbabf..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/nccl/BUILD +++ /dev/null @@ -1,177 +0,0 @@ -# Description: -# Wrap NVIDIA (https://github.com/NVIDIA/nccl) NCCL with tensorflow ops. -# APIs are meant to change over time. - -package(default_visibility = ["//tensorflow:__subpackages__"]) - -licenses(["notice"]) # Apache 2.0 - -exports_files(["LICENSE"]) - -load( - "//tensorflow:tensorflow.bzl", - "tf_custom_op_library", - "tf_gen_op_libs", - "tf_gen_op_wrapper_py", -) -load("//tensorflow:tensorflow.bzl", "tf_cuda_cc_test") -load("//tensorflow:tensorflow.bzl", "cuda_py_test") -load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda") -load("//tensorflow:tensorflow.bzl", "tf_kernel_library") -load("//tensorflow:tensorflow.bzl", "tf_custom_op_py_library") -load("//tensorflow:tensorflow.bzl", "if_not_windows_cuda") - -tf_custom_op_library( - name = "python/ops/_nccl_ops.so", - srcs = [ - "ops/nccl_ops.cc", - ], - gpu_srcs = if_not_windows_cuda([ - "kernels/nccl_manager.cc", - "kernels/nccl_manager.h", - "kernels/nccl_ops.cc", - ]), - deps = [] + if_cuda([ - "@local_config_nccl//:nccl", - "//tensorflow/core:gpu_headers_lib", - "//tensorflow/core:protos_all_cc", - "//tensorflow/core:protos_all_proto_text", - ]), -) - -tf_cuda_cc_test( - name = "nccl_manager_test", - size = "medium", - srcs = if_cuda( - [ - "kernels/nccl_manager.cc", - "kernels/nccl_manager.h", - "kernels/nccl_manager_test.cc", - ], - [], - ), - # Disabled on jenkins until errors finding nvmlShutdown are found. - tags = [ - "manual", - "multi_gpu", - "no_oss", - "noguitar", - "notap", - ], - deps = - if_cuda([ - "@local_config_nccl//:nccl", - "//tensorflow/core:cuda", - "//tensorflow/core:test", - "//tensorflow/core:test_main", - "//tensorflow/core:testlib", - ]), -) - -tf_kernel_library( - name = "nccl_kernels", - srcs = if_cuda([ - "kernels/nccl_manager.cc", - "kernels/nccl_manager.h", - "kernels/nccl_ops.cc", - "kernels/nccl_rewrite.cc", - ]), - deps = if_cuda([ - "@local_config_nccl//:nccl", - "//tensorflow/core:core_cpu", - "//tensorflow/core:framework", - "//tensorflow/core:gpu_headers_lib", - "//tensorflow/core:lib", - "//tensorflow/core:stream_executor", - ]), - alwayslink = 1, -) - -tf_gen_op_libs( - op_lib_names = ["nccl_ops"], - deps = [ - "//tensorflow/core:lib", - ], -) - -tf_gen_op_wrapper_py( - name = "nccl_ops", - deps = [":nccl_ops_op_lib"], -) - -# Test only nccl ops lib without dso to test behavior when NCCL lib is not -# installed. See nccl_dependency_test for more details. -# -# Users should use the public nccl_py lib that also adds the dso. -tf_custom_op_py_library( - name = "nccl_ops_lib_without_dso", - srcs = [ - "__init__.py", - "python/ops/nccl_ops.py", - ], - kernels = if_cuda([":nccl_kernels"]) + [ - ":nccl_ops_op_lib", - ], - deps = [ - ":nccl_ops", - "//tensorflow/contrib/util:util_py", - "//tensorflow/python:device", - "//tensorflow/python:framework_ops", - "//tensorflow/python:platform", - "//tensorflow/python:util", - "//tensorflow/python/eager:context", - ], -) - -tf_custom_op_py_library( - name = "nccl_py", - dso = [":python/ops/_nccl_ops.so"], - visibility = ["//visibility:public"], - deps = [ - ":nccl_ops_lib_without_dso", - ], -) - -cuda_py_test( - name = "nccl_ops_test", - size = "small", - srcs = ["python/ops/nccl_ops_test.py"], - additional_deps = [ - ":nccl_py", - "//tensorflow/python:array_ops", - "//tensorflow/python:client_testlib", - "//tensorflow/python:framework_for_generated_wrappers", - "//tensorflow/python:framework_test_lib", - "//tensorflow/python:platform_test", - ], - # Disabled on jenkins until errors finding nvmlShutdown are found. - tags = [ - "manual", - "multi_gpu", - "no_oss", - "noguitar", - "notap", - ], -) - -cuda_py_test( - name = "nccl_dependency_test", - size = "small", - srcs = ["python/ops/nccl_dependency_test.py"], - additional_deps = [ - ":nccl_ops_lib_without_dso", - "//tensorflow/python:constant_op", - "//tensorflow/python:errors", - "//tensorflow/python:framework_ops", - "//tensorflow/python:util", - "//tensorflow/python:client_testlib", - "//tensorflow/python:platform_test", - ], - # Disable this test internally as static linking is used internally and only - # run for OSS to verify that NCCL is an optional dynamic dependency. - tags = [ - "manual", - "noguitar", - "notap", - ], -) diff --git a/tensorflow/contrib/nccl/__init__.py b/tensorflow/contrib/nccl/__init__.py deleted file mode 100644 index 4a682cb70369e1ae6edec67730618ab6a1ba6f47..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/nccl/__init__.py +++ /dev/null @@ -1,38 +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. -# ============================================================================== -"""Functions for using NVIDIA nccl collective ops. - -@@all_max -@@all_min -@@all_prod -@@all_sum -@@reduce_sum -@@broadcast - -""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -from tensorflow.contrib.nccl.python.ops.nccl_ops import all_max -from tensorflow.contrib.nccl.python.ops.nccl_ops import all_min -from tensorflow.contrib.nccl.python.ops.nccl_ops import all_prod -from tensorflow.contrib.nccl.python.ops.nccl_ops import all_sum -from tensorflow.contrib.nccl.python.ops.nccl_ops import broadcast -from tensorflow.contrib.nccl.python.ops.nccl_ops import reduce_sum - -from tensorflow.python.util.all_util import remove_undocumented -remove_undocumented(__name__) diff --git a/tensorflow/contrib/nccl/kernels/nccl_manager.cc b/tensorflow/contrib/nccl/kernels/nccl_manager.cc deleted file mode 100644 index 99fecf96517935bf3bde3636df83b4a9a4e1c779..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/nccl/kernels/nccl_manager.cc +++ /dev/null @@ -1,529 +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/contrib/nccl/kernels/nccl_manager.h" - -#include - -#ifdef GOOGLE_CUDA - -#include "tensorflow/core/lib/core/threadpool.h" -#include "tensorflow/core/platform/cuda.h" -#include "tensorflow/core/platform/env.h" - -namespace tensorflow { - -using se::cuda::ScopedActivateExecutorContext; - -// Contains data for a single stream used for nccl communication; this includes -// a background thread that calls NcclManager::LoopKernelLaunches. -struct NcclManager::NcclStream { - public: - NcclStream() {} - ~NcclStream() { - mutex_lock l(mu); - shutdown_requested = true; - cv.notify_all(); - } - - se::StreamExecutor* executor = nullptr; - - // The stream on which to run the nccl collective. - // This is a different stream than the tensorflow compute stream. - std::unique_ptr stream; - - // See NcclManager::LoopKernelLaunches for information on these. - std::unique_ptr thread; - mutex mu; - condition_variable cv; - // Has collective,rank pairs. - std::deque> pending_launches_ GUARDED_BY(mu); - bool shutdown_requested GUARDED_BY(mu) = false; -}; - -struct NcclManager::CommunicatorMember { - public: - CommunicatorMember() {} - ~CommunicatorMember() { - if (nccl_comm != nullptr) ncclCommDestroy(nccl_comm); - } - ncclComm_t nccl_comm; - - // Owned by NcclManager::device_to_comm_streams_. - NcclStream* nccl_stream = nullptr; -}; - -struct NcclManager::Communicator { - public: - explicit Communicator(std::vector members) - : num_devices(members.size()), members(std::move(members)) {} - - const int num_devices; - const std::vector members; // indexed by rank. -}; - -namespace { -ncclDataType_t ToNcclType(DataType t) { - switch (t) { - case DT_HALF: - return ncclHalf; - case DT_FLOAT: - return ncclFloat; - case DT_DOUBLE: - return ncclDouble; - case DT_INT32: - return ncclInt; - case DT_INT64: - return ncclInt64; - default: - return ncclFloat; - } -} -} // namespace - -// A participant in a Collective. See below. -struct NcclManager::Participant { - Participant(const Tensor* in_t, Tensor* out_t, EventMgr* event_mgr, - se::Stream* tensor_stream, se::StreamExecutor* executor, - int gpu_device_id, NcclManager::DoneCallback done_callback) - : in_t(in_t), - out_t(out_t), - event_mgr(event_mgr), - tensor_stream(tensor_stream), - executor(executor), - gpu_device_id(gpu_device_id), - done_callback(std::move(done_callback)) { - DCHECK(executor != nullptr); - DCHECK(event_mgr != nullptr); - DCHECK(tensor_stream != nullptr); - } - // Owned by the caller, who must keep it live until is called. - // Is NULL for participants that only receive data. - const Tensor* in_t; - - // Owned by the caller, who must keep it live until is called. - // Is NULL for participants that only send data. - Tensor* out_t; - - // Owned by the caller, who must keep it live until is called. - EventMgr* const event_mgr; - - // Owned by the caller, who must keep it live until is called. - se::Stream* const tensor_stream; - - // Matches the executor in CommunicatorMember::stream. Expected to be live for - // process lifetime. - se::StreamExecutor* const executor = nullptr; - - const int gpu_device_id; - - NcclManager::DoneCallback done_callback; - - bool root = false; -}; - -// A Collective tracks a single communicator operation (e.g., a single -// AllReduce call). -struct NcclManager::Collective { - Collective(DataType data_type_in, CollectiveType type_in, - ncclRedOp_t reduction_op_in, int num_devices) - : data_type(data_type_in), - type(type_in), - reduction_op(reduction_op_in), - remaining_participants(num_devices) { - participants.reserve(num_devices); - } - - const DataType data_type; - const CollectiveType type; - const ncclRedOp_t reduction_op; // applies when is a reduction. - - Communicator* communicator = nullptr; - - // All collective participants. - // - // Adding values in this vector is guarded by the mutex of the containing - // NcclManager. - std::vector> participants; - - // For collective types that have a root (e.g. the root of broadcast is the - // sender), this is the rank of the root. - int root_rank = -1; - - // How many participants have been registered so far. The Collective is - // eligible for running with == participants.size(). - // - // Guarded by the mutex of the containing Communicator. - int available_participants = 0; - - mutable std::atomic_int_fast32_t remaining_participants; -}; - -NcclManager::NcclManager() {} -NcclManager::~NcclManager() {} -NcclManager* NcclManager::instance() { - static NcclManager* instance = new NcclManager(); - return instance; -} - -NcclManager::Communicator* NcclManager::GetCommunicator( - NcclManager::Collective* collective) { - // Sort by executor to make ordering of executors deterministic. - std::sort(collective->participants.begin(), collective->participants.end(), - [](const std::unique_ptr& a, - const std::unique_ptr& b) { - return a->executor < b->executor; - }); - const int num_devices = collective->participants.size(); - - mutex_lock l(mu_); - - // Scan to find an existing communicator that provides nccl communication - // between the executors used by the participants in the collective. For - // example, if a collective is for GPUs 0, 1, and 2 then this will scan - // to find the communicator for GPUs 0, 1, and 2. - // - // Note that each executor identifies a context on one device, so this is the - // same as getting the communicator connecting the devices in the collective. - // A device can be in different communicators as well - for example, a - // communicator for GPUs 0 and 1 is separate from one for GPUs 0, 1, and 2. - // - // Since it's expected that a small number of distinct communicators will - // be needed, communicators_ is not garbage collected currently. - // - // Launching of kernels must be serialized so that, given collectives A and B, - // and an order of them (e.g., A before B), then for each comm_stream - // involved, the kernel for A is launched before the kernel for B. This is - // guaranteed currently be a global mutex controlling additions of the kernels - // to per-stream launch queues. The launch queues are processed by - // LoopKernelLaunches. - for (auto& comm : communicators_) { - if (comm->num_devices == num_devices) { - int i; - for (i = 0; i < num_devices; ++i) { - if (comm->members[i].nccl_stream->executor != - collective->participants[i]->executor) { - break; - } - } - if (i == num_devices) return comm.get(); - } - } - - auto* env = Env::Default(); - std::set used_streams; - - // Create and initialize a new communicator. - // Note that this is done under the lock; performance is not expected to - // matter as this happens a very small number of times. - std::vector members(num_devices); - std::vector devices(num_devices); - for (int i = 0; i < num_devices; ++i) { - auto* executor = collective->participants[i]->executor; - - // Find a communication stream to use for the device. - auto& streams = device_to_comm_streams_[executor]; - NcclStream* nccl_stream = nullptr; - for (const auto& s : streams) { - if (used_streams.insert(s.get()).second) { - nccl_stream = s.get(); - break; - } - } - if (nccl_stream == nullptr) { - nccl_stream = new NcclStream(); - nccl_stream->executor = executor; - nccl_stream->stream.reset(new se::Stream(executor)); - nccl_stream->stream->Init(); - - streams.emplace_back(nccl_stream); - used_streams.insert(nccl_stream); - - nccl_stream->thread.reset(env->StartThread( - ThreadOptions(), "nccl_kernel_launch", - [this, nccl_stream] { LoopKernelLaunches(nccl_stream); })); - } - - members[i].nccl_stream = nccl_stream; - devices[i] = collective->participants[i]->gpu_device_id; - } - - int device_count = num_devices; -#if NCCL_MAJOR >= 2 - // NCCL2 prevents InitAll for more communicators than devices (but doesn't - // check that device ids are unique). Work around it by initializing each - // rank individually. - cudaGetDeviceCount(&device_count); -#endif - std::vector nccl_comms(num_devices); - if (num_devices <= device_count) { - auto result = - ncclCommInitAll(nccl_comms.data(), num_devices, devices.data()); - CHECK_EQ(result, ncclSuccess) << ncclGetErrorString(result); - } else { - int savedDevice = 0; - CHECK_EQ(cudaGetDevice(&savedDevice), cudaSuccess); - ncclUniqueId commId; - ncclGetUniqueId(&commId); -#if NCCL_MAJOR >= 2 - CHECK_EQ(ncclGroupStart(), ncclSuccess); -#endif - for (int rank = 0; rank < num_devices; ++rank) { - cudaSetDevice(devices[rank]); - auto result = - ncclCommInitRank(nccl_comms.data() + rank, num_devices, commId, rank); - CHECK_EQ(result, ncclSuccess) << ncclGetErrorString(result); - } -#if NCCL_MAJOR >= 2 - CHECK_EQ(ncclGroupEnd(), ncclSuccess); -#endif - cudaSetDevice(savedDevice); - } - for (int rank = 0; rank < num_devices; ++rank) { - members[rank].nccl_comm = nccl_comms[rank]; - } - communicators_.emplace_back(new Communicator(std::move(members))); - return communicators_.back().get(); -} - -void NcclManager::AddToAllReduce(int num_devices, const string& key, - ncclRedOp_t reduction_op, - se::StreamExecutor* executor, - int gpu_device_id, EventMgr* event_mgr, - se::Stream* tensor_stream, const Tensor* in_t, - Tensor* out_t, - const DoneCallback& done_callback) { - std::unique_ptr participant( - new Participant(in_t, out_t, event_mgr, tensor_stream, executor, - gpu_device_id, done_callback)); - AddParticipant(num_devices, key, std::move(participant), in_t->dtype(), - kAllReduce, reduction_op); -} - -void NcclManager::AddBroadcastSend(int num_devices, const string& key, - se::StreamExecutor* executor, - int gpu_device_id, EventMgr* event_mgr, - se::Stream* tensor_stream, - const Tensor* in_t, - DoneCallback done_callback) { - std::unique_ptr participant( - new Participant(in_t, nullptr /* out_t */, event_mgr, tensor_stream, - executor, gpu_device_id, std::move(done_callback))); - participant->root = true; - AddParticipant(num_devices, key, std::move(participant), in_t->dtype(), - kBroadcast, ncclSum /* unused */); -} - -void NcclManager::AddBroadcastRecv(int num_devices, const string& key, - se::StreamExecutor* executor, - int gpu_device_id, EventMgr* event_mgr, - se::Stream* tensor_stream, Tensor* out_t, - DoneCallback done_callback) { - std::unique_ptr participant( - new Participant(nullptr /* in_t */, out_t, event_mgr, tensor_stream, - executor, gpu_device_id, std::move(done_callback))); - AddParticipant(num_devices, key, std::move(participant), out_t->dtype(), - kBroadcast, ncclSum /* unused */); -} - -void NcclManager::AddReduceSend(int num_devices, const string& key, - ncclRedOp_t reduction_op, - se::StreamExecutor* executor, int gpu_device_id, - EventMgr* event_mgr, se::Stream* tensor_stream, - const Tensor* in_t, - DoneCallback done_callback) { - std::unique_ptr participant( - new Participant(in_t, nullptr /* out_t */, event_mgr, tensor_stream, - executor, gpu_device_id, std::move(done_callback))); - AddParticipant(num_devices, key, std::move(participant), in_t->dtype(), - kReduce, reduction_op); -} - -void NcclManager::AddReduceRecv(int num_devices, const string& key, - ncclRedOp_t reduction_op, - se::StreamExecutor* executor, int gpu_device_id, - EventMgr* event_mgr, se::Stream* tensor_stream, - const Tensor* in_t, Tensor* out_t, - DoneCallback done_callback) { - std::unique_ptr participant( - new Participant(in_t, out_t, event_mgr, tensor_stream, executor, - gpu_device_id, std::move(done_callback))); - participant->root = true; - AddParticipant(num_devices, key, std::move(participant), in_t->dtype(), - kReduce, reduction_op); -} - -void NcclManager::AddParticipant(int num_devices, const string& key, - std::unique_ptr participant, - DataType data_type, - CollectiveType collective_type, - ncclRedOp_t reduction_op) { - Collective* to_run = nullptr; - { - mutex_lock l(mu_); - auto& collective_ptr = collectives_[key]; - if (collective_ptr == nullptr) { - collective_ptr.reset(new Collective(data_type, collective_type, - reduction_op, num_devices)); - } - Collective* collective = collective_ptr.get(); - DCHECK_EQ(collective->type, collective_type); - DCHECK_LT(collective->participants.size(), num_devices); - collective->participants.emplace_back(std::move(participant)); - ++collective->available_participants; - - if (collective->available_participants == num_devices) { - to_run = collective; - - // Ownership is going to be transferred to RunCollective. - collective_ptr.release(); - collectives_.erase(key); - } - } - - if (to_run != nullptr) { - RunCollective(key, to_run); - } -} - -void NcclManager::RunCollective(const string& key, Collective* collective) { - static mutex collective_mu(LINKER_INITIALIZED); - - auto* communicator = GetCommunicator(collective); - collective->communicator = communicator; - const int size = communicator->num_devices; - - for (int rank = 0; rank < size; ++rank) { - Participant* p = collective->participants[rank].get(); - NcclStream* nccl_stream = communicator->members[rank].nccl_stream; - CHECK(nccl_stream != nullptr); - - if (p->in_t != nullptr) { - // Wait to ensure that the kernel that produces the data in the input - // tensor has finished running before the nccl kernel runs on the - // communication stream. - nccl_stream->stream->ThenWaitFor(p->tensor_stream); - } - if (p->root) { - CHECK_EQ(collective->root_rank, -1); - collective->root_rank = rank; - } - } - - if (collective->type == kBroadcast) { - CHECK_NE(collective->root_rank, -1); - } - - { - // Allow only one collective at a time to queue kernels for launching. This - // is to prevent collectives from deadlocking each other. - // Note that it would be possible to run multiple collectives at once, if - // they have non-intersecting sets of devices. - mutex_lock l(collective_mu); - for (int rank = 0; rank < size; ++rank) { - NcclStream* nccl_stream = communicator->members[rank].nccl_stream; - mutex_lock l(nccl_stream->mu); - nccl_stream->pending_launches_.push_front( - std::make_pair(collective, rank)); - nccl_stream->cv.notify_all(); - } - } -} - -void NcclManager::LoopKernelLaunches(NcclStream* nccl_stream) { - se::Stream* comm_stream = nccl_stream->stream.get(); - ScopedActivateExecutorContext scoped_context(nccl_stream->executor); - const cudaStream_t* cu_stream = reinterpret_cast( - comm_stream->implementation()->GpuStreamMemberHack()); - - while (true) { - // Find collective to run. - std::pair next_launch; - { - mutex_lock l(nccl_stream->mu); - while (nccl_stream->pending_launches_.empty()) { - if (nccl_stream->shutdown_requested) { - // No work and shutdown requested, exit. - return; - } - nccl_stream->cv.wait(l); - } - next_launch = nccl_stream->pending_launches_.back(); - nccl_stream->pending_launches_.pop_back(); - } - Collective* collective = next_launch.first; - int rank = next_launch.second; - - // Launch the nccl kernel. - ncclDataType_t data_type = ToNcclType(collective->data_type); - Participant* p = collective->participants[rank].get(); - - auto nccl_comm = collective->communicator->members[rank].nccl_comm; - ncclResult_t nccl_result = ncclSuccess; - switch (collective->type) { - case kAllReduce: { - const void* sendbuff = p->in_t->tensor_data().data(); - void* recvbuff = const_cast(p->out_t->tensor_data().data()); - - nccl_result = - ncclAllReduce(sendbuff, recvbuff, p->in_t->NumElements(), data_type, - collective->reduction_op, nccl_comm, *cu_stream); - break; - } - case kBroadcast: { - const Tensor* buf_t = p->in_t ? p->in_t : p->out_t; - void* buf = const_cast(buf_t->tensor_data().data()); - nccl_result = ncclBcast(buf, buf_t->NumElements(), data_type, - collective->root_rank, nccl_comm, *cu_stream); - break; - } - case kReduce: { - const void* sendbuff = p->in_t->tensor_data().data(); - void* recvbuff = p->out_t - ? const_cast(p->out_t->tensor_data().data()) - : nullptr; - nccl_result = ncclReduce(sendbuff, recvbuff, p->in_t->NumElements(), - data_type, collective->reduction_op, - collective->root_rank, nccl_comm, *cu_stream); - break; - } - } - - // Run the done_callback when the nccl kernel finishes running. - auto done_callback = [collective, rank, nccl_result]() { - if (nccl_result == ncclSuccess) { - collective->participants[rank]->done_callback(Status::OK()); - } else { - // Propagate the error, but note that if other members of the collective - // did launch their kernels, then they are hanging. - collective->participants[rank]->done_callback(errors::Unknown( - "Error invoking NCCL: ", ncclGetErrorString(nccl_result))); - } - - // TODO(cwhipkey): use RefCounted after figuring out how to use in a - // custom op library. - // See tensorflow/core/lib/core/refcount.h for details on this locking. - if (collective->remaining_participants.load(std::memory_order_acquire) == - 1 || - collective->remaining_participants.fetch_sub(1) == 1) { - delete collective; - } - }; - p->event_mgr->ThenExecute(comm_stream, done_callback); - } -} - -} // namespace tensorflow - -#endif // GOOGLE_CUDA diff --git a/tensorflow/contrib/nccl/kernels/nccl_manager.h b/tensorflow/contrib/nccl/kernels/nccl_manager.h deleted file mode 100644 index 7d158cc98026678edafa0845df92038b449a9225..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/nccl/kernels/nccl_manager.h +++ /dev/null @@ -1,138 +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. -==============================================================================*/ -#ifndef TENSORFLOW_CONTRIB_NCCL_KERNELS_NCCL_MANAGER_H_ -#define TENSORFLOW_CONTRIB_NCCL_KERNELS_NCCL_MANAGER_H_ - -#ifdef GOOGLE_CUDA - -#include -#include - -// TODO(rmlarsen): Get rid of this workaround. "gpu_assert" is defined when -// setting EIGEN_USE_THREADS. But when defining EIGEN_USE_THREADS here, -// incAtomic and other CUDA specific symbols are no longer recognized. -#ifndef gpu_assert -#define gpu_assert(x) -#endif - -#include "third_party/nccl/nccl.h" -#include "tensorflow/core/common_runtime/gpu/gpu_event_mgr.h" -#include "tensorflow/core/framework/tensor.h" -#include "tensorflow/core/platform/mutex.h" -#include "tensorflow/core/platform/stream_executor.h" - -namespace tensorflow { - -// The communicator is used to make the asynchronous communicator calls and to -// manage the per-device streams used for communication. -// -// See nccl_ops.cc for example usage, including description of memory -// management and stream synchronization. -class NcclManager { - public: - typedef std::function DoneCallback; - NcclManager(); - ~NcclManager(); - - static NcclManager* instance(); - - // Add one participant to an all-reduce, sending in data from and - // receiving the result of the all-reduce in . The device for this - // participant is managed by , and its events are polled by - // . - // - // This is an asynchronous call. When is called, has - // been set to the all-reduce result (note: the stream may not yet have been - // synced). - // - // is the stream that should be waited on to ensure 's - // data is available on the GPU for the communication stream to access. It - // is also the stream that will use the produced data; is - // not called until the next kernel launched on would see the data. - void AddToAllReduce(int num_devices, const string& key, - ncclRedOp_t reduction_op, se::StreamExecutor* executor, - int gpu_device_id, EventMgr* event_mgr, - se::Stream* tensor_stream, const Tensor* in_t, - Tensor* out_t, const DoneCallback& done_callback); - - // AddBroadcastSend and AddBroadcastRecv combine to sent data from one sender - // to all receivers. - void AddBroadcastSend(int num_devices, const string& key, - se::StreamExecutor* executor, int gpu_device_id, - EventMgr* event_mgr, se::Stream* tensor_stream, - const Tensor* in_t, DoneCallback done_callback); - void AddBroadcastRecv(int num_devices, const string& key, - se::StreamExecutor* executor, int gpu_device_id, - EventMgr* event_mgr, se::Stream* tensor_stream, - Tensor* out_t, DoneCallback done_callback); - - // AddReduceSend and AddReduceRecv combine to sent data from all senders - // to one receiver. - void AddReduceSend(int num_devices, const string& key, - ncclRedOp_t reduction_op, se::StreamExecutor* executor, - int gpu_device_id, EventMgr* event_mgr, - se::Stream* tensor_stream, const Tensor* in_t, - DoneCallback done_callback); - void AddReduceRecv(int num_devices, const string& key, - ncclRedOp_t reduction_op, se::StreamExecutor* executor, - int gpu_device_id, EventMgr* event_mgr, - se::Stream* tensor_stream, const Tensor* in_t, - Tensor* out_t, DoneCallback done_callback); - - private: - enum CollectiveType { - kAllReduce = 1, - kBroadcast = 2, - kReduce = 3, - }; - struct Collective; - struct Communicator; - struct CommunicatorMember; - struct NcclStream; - struct Participant; - - Communicator* GetCommunicator(Collective* collective); - - void AddParticipant(int num_devices, const string& key, - std::unique_ptr participant, - DataType data_type, CollectiveType collective_type, - ncclRedOp_t reduction_op); - - // Run . This calls takes ownership of . - void RunCollective(const string& key, Collective* collective); - void LoopKernelLaunches(NcclStream* stream); - - mutex mu_; - - // Maps key to collectives currently being assembled or run. - std::unordered_map> collectives_ - GUARDED_BY(mu_); - - // Maps a device to the communication streams that make up its collective. - // This is used to share the stream across different communicators that - // include the same device. - std::map>> - device_to_comm_streams_ GUARDED_BY(mu_); - - std::vector> communicators_; - - TF_DISALLOW_COPY_AND_ASSIGN(NcclManager); -}; - -} // namespace tensorflow - -#endif // GOOGLE_CUDA - -#endif // TENSORFLOW_CONTRIB_NCCL_KERNELS_NCCL_MANAGER_H_ diff --git a/tensorflow/contrib/nccl/kernels/nccl_manager_test.cc b/tensorflow/contrib/nccl/kernels/nccl_manager_test.cc deleted file mode 100644 index 5144f7c38c8650ebfced1dfcc9378263ebaad8c0..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/nccl/kernels/nccl_manager_test.cc +++ /dev/null @@ -1,307 +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. -==============================================================================*/ - -#ifdef GOOGLE_CUDA - -#include -#include -#include - -#include "tensorflow/contrib/nccl/kernels/nccl_manager.h" -#include "tensorflow/core/common_runtime/device_factory.h" -#include "tensorflow/core/common_runtime/gpu/gpu_device.h" -#include "tensorflow/core/framework/tensor_testutil.h" -#include "tensorflow/core/lib/core/status_test_util.h" -#include "tensorflow/core/platform/test.h" - -namespace tensorflow { - -static std::vector GetGPUDevices() { - std::vector devices; - SessionOptions session_options; - session_options.config.mutable_gpu_options() - ->set_per_process_gpu_memory_fraction(0.1); - session_options.env = Env::Default(); - Status s = DeviceFactory::GetFactory(DEVICE_GPU) - ->AddDevices(session_options, "", &devices); - TF_CHECK_OK(s); - std::vector gpus; - for (Device* d : devices) { - if (d->device_type() == "GPU") { - gpus.push_back(static_cast(d)); - } else { - delete d; - } - } - return gpus; -} - -template -class NcclManagerTest : public ::testing::Test { - public: - // A single all-reduce to apply. - struct TestCase { - string key; - std::vector ins; - std::vector outs; - Tensor expected; - - mutex mu; - Status final_status; - int num_completed = 0; - }; - - static void SetUpTestCase() { - setenv("NCCL_DEBUG", "INFO", 1 /* replace */); - devices_ = new std::vector(GetGPUDevices()); - CHECK(!devices_->empty()); - LOG(ERROR) << "Running test with " << devices_->size() << " gpus"; - } - - static void TearDownTestCase() { - for (auto device : *devices_) delete device; - delete devices_; - } - - TestCase* MakeTestCase(int num_ranks, ncclRedOp_t reduction_op, - TensorShape shape, float value_offset) { - TestCase* test_case = new TestCase(); - test_case->expected = Tensor(data_type_, shape); - if (reduction_op == ncclProd) { - test::FillFn(&test_case->expected, - [](int) { return static_cast(1); }); - } else if (reduction_op == ncclSum) { - test::FillFn(&test_case->expected, - [](int) { return static_cast(0); }); - } else if (reduction_op == ncclMax) { - test::FillFn(&test_case->expected, [](int) { return -max_; }); - } else if (reduction_op == ncclMin) { - test::FillFn(&test_case->expected, [](int) { return max_; }); - } else { - LOG(FATAL) << "Invalid reduction_op " << reduction_op; - } - - float value_scale = 0.01; // Small scale to avoid fp16 overflow. - for (int rank = 0; rank < num_ranks; ++rank) { - auto* device = GetDevice(rank); - auto* stream = device->tensorflow_gpu_device_info()->stream; - - Tensor in_cpu(data_type_, shape); - test::FillFn(&in_cpu, [&](int index) { - return static_cast((index + 1) * value_scale + value_offset); - }); - for (int j = 0; j < shape.num_elements(); ++j) { - auto in_val = in_cpu.flat()(j); - auto out_expr = test_case->expected.template flat(); - if (reduction_op == ncclProd) { - out_expr(j) = out_expr(j) * in_val; - } else if (reduction_op == ncclSum) { - out_expr(j) = out_expr(j) + in_val; - } else if (reduction_op == ncclMax) { - if (in_val > out_expr(j)) { - out_expr(j) = in_val; - } - } else if (reduction_op == ncclMin) { - if (in_val < out_expr(j)) { - out_expr(j) = in_val; - } - } - } - - value_scale *= 10; - test_case->ins.emplace_back(GpuAllocator(device), data_type_, shape); - test_case->outs.emplace_back(GpuAllocator(device), data_type_, shape); - - const Tensor& in_gpu = test_case->ins.back(); - auto in_gpu_mem = AsDeviceMemory(in_gpu.flat().data()); - stream->ThenMemcpy(&in_gpu_mem, in_cpu.flat().data(), - in_cpu.TotalBytes()); - } - return test_case; - } - - void VerifyResults(const string& case_label, TestCase* test_case) { - // Wait for the done callback to be called. - { - test_case->mu.lock(); - while (test_case->num_completed != test_case->outs.size()) { - test_case->mu.unlock(); - Env::Default()->SleepForMicroseconds(10); - test_case->mu.lock(); - } - test_case->mu.unlock(); - } - // Copy memory to host and verify. - for (int rank = 0; rank < test_case->outs.size(); ++rank) { - auto* device = GetDevice(rank); - auto* stream = device->tensorflow_gpu_device_info()->stream; - const Tensor& out_gpu = test_case->outs[rank]; - Tensor out_cpu(data_type_, out_gpu.shape()); - auto out_gpu_mem = AsDeviceMemory(out_gpu.flat().data()); - stream->ThenMemcpy(out_cpu.flat().data(), out_gpu_mem, - out_cpu.TotalBytes()); - SE_ASSERT_OK(stream->BlockHostUntilDone()); - test::ExpectTensorNear(test_case->expected, out_cpu, 0.01); - } - } - - NcclManager::DoneCallback CreateDoneCallback(TestCase* test_case) { - return [this, test_case](Status s) { - mutex_lock l(test_case->mu); - ++test_case->num_completed; - test_case->final_status.Update(s); - }; - } - - static BaseGPUDevice* GetDevice(size_t rank) { - return devices_->at(rank % devices_->size()); - } - - private: - static Allocator* GpuAllocator(BaseGPUDevice* device) { - return device->GetAllocator(AllocatorAttributes()); - } - - static se::DeviceMemory AsDeviceMemory(const Scalar* cuda_memory) { - se::DeviceMemoryBase wrapped(const_cast(cuda_memory)); - se::DeviceMemory typed(wrapped); - return typed; - } - - private: - static std::vector* devices_; - static const DataType data_type_; - static const Scalar max_; -}; - -template -std::vector* NcclManagerTest::devices_ = nullptr; -template -const DataType NcclManagerTest::data_type_ = - DataTypeToEnum::value; -template -const Scalar NcclManagerTest::max_ = - Eigen::NumTraits::highest(); - -// Instantiate tests for float and half. -using TypeList = ::testing::Types; -TYPED_TEST_CASE(NcclManagerTest, TypeList); - -// Test basic sum reduction. -TYPED_TEST(NcclManagerTest, BasicSumReduction) { - const int num_ranks = 3; - - for (int op = 0; op < 4; ++op) { - ncclRedOp_t reduction_op = static_cast(op); - std::unique_ptr test_case( - this->MakeTestCase(num_ranks, reduction_op, TensorShape({2, 3}), 0.0f)); - for (int rank = 0; rank < num_ranks; ++rank) { - auto* device = this->GetDevice(rank); - auto* event_mgr = device->tensorflow_gpu_device_info()->event_mgr; - auto* stream = device->tensorflow_gpu_device_info()->stream; - NcclManager::instance()->AddToAllReduce( - num_ranks, "allreduce", reduction_op, device->executor(), - device->gpu_id(), event_mgr, stream, &test_case->ins[rank], - &test_case->outs[rank], this->CreateDoneCallback(test_case.get())); - } - - LOG(ERROR) << "Verifying results"; - this->VerifyResults("test_case", test_case.get()); - } -} - -// Same as the Basic test, but with multiple threads launching parts of many -// reductions. -// -// Testing the multi-rank execution is currently reduced as it can hang when run -// with num_ranks > devices->size(), for some GPUs (e.g. K20m). -// To test the higher settings, increase num_ranks, -// num_collectives_per_iteration and time_limit_micros. -TYPED_TEST(NcclManagerTest, MultipleCallers) { - const int num_ranks = 1; // 2; - const int num_collectives_per_iteration = 1; // 1000; - const int num_threads = 3; - const int time_limit_micros = 1; // 60 * 30 * 1000 * 1000; - - int64 start = Env::Default()->NowMicros(); - srand(Env::Default()->NowMicros()); - - for (;;) { - std::vector> case_and_rank; - std::vector> test_cases; - for (int i = 0; i < num_collectives_per_iteration; ++i) { - test_cases.emplace_back(this->MakeTestCase( - num_ranks, ncclSum, TensorShape({100, i % 5 + 1, i % 3 + 1}), - 1.1f * i)); - for (int j = 0; j < num_ranks; ++j) { - case_and_rank.emplace_back(i, j); - } - } - - for (int rank = 0; rank < num_ranks; ++rank) { - auto* device = this->GetDevice(rank); - auto* stream = device->tensorflow_gpu_device_info()->stream; - SE_ASSERT_OK(stream->BlockHostUntilDone()); - } - - std::shuffle(case_and_rank.begin(), case_and_rank.end(), - std::mt19937(std::random_device()())); - - mutex mu; // guards case_and_rank. - std::unique_ptr pool( - new thread::ThreadPool(Env::Default(), "test", num_threads)); - const int to_schedule = case_and_rank.size(); - for (int i = 0; i < to_schedule; ++i) { - auto fn = [&]() { - int rank; - int test_num; - { - mutex_lock l(mu); - test_num = case_and_rank.back().first; - rank = case_and_rank.back().second; - case_and_rank.pop_back(); - } - auto* device = this->GetDevice(rank); - auto* event_mgr = device->tensorflow_gpu_device_info()->event_mgr; - auto* stream = device->tensorflow_gpu_device_info()->stream; - typename TestFixture::TestCase* test_case = test_cases[test_num].get(); - NcclManager::instance()->AddToAllReduce( - num_ranks, strings::StrCat("allreduce", test_num), ncclSum, - device->executor(), device->gpu_id(), event_mgr, stream, - &test_case->ins[rank], &test_case->outs[rank], - this->CreateDoneCallback(test_case)); - }; - pool->Schedule(fn); - } - pool.reset(); // wait for all work to be scheduled. - - LOG(ERROR) << "Verifying results for " << num_collectives_per_iteration - << " collectives"; - for (int i = 0; i < test_cases.size(); ++i) { - this->VerifyResults(strings::StrCat("collective", i), - test_cases[i].get()); - } - - int64 delta = Env::Default()->NowMicros() - start; - if (delta > time_limit_micros) { - LOG(ERROR) << "Ran for " << delta << " quitting"; - break; - } - } -} - -} // namespace tensorflow - -#endif // GOOGLE_CUDA diff --git a/tensorflow/contrib/nccl/kernels/nccl_ops.cc b/tensorflow/contrib/nccl/kernels/nccl_ops.cc deleted file mode 100644 index c2b76caef38a4af248387b65701b8f8936e8431f..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/nccl/kernels/nccl_ops.cc +++ /dev/null @@ -1,246 +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. -==============================================================================*/ - -#if GOOGLE_CUDA - -#include - -#include "third_party/nccl/nccl.h" -#include "tensorflow/contrib/nccl/kernels/nccl_manager.h" -#include "tensorflow/core/framework/op_kernel.h" - -namespace tensorflow { -namespace { - -// Base class for all communicator ops that use nccl. -// -// About memory management and stream syncing: -// 1. The nccl communicator has a stream for each rank. -// 2. For input tensors to the communicator, the compute stream is passed to the -// NcclManager which will do a needed -// communicator_stream.ThenWaitFor(input_tensor_stream). -// 3. The done_callback of the async kernel is not called by the -// NcclManager until after the communicator kernel is complete. This -// is enough to a) keep the input tensor data valid for the lifetime of the -// collective; and b) ensure the data in the output tensor is available -// when the async op kernel's done callback is called. -class NcclAsyncOpBase : public AsyncOpKernel { - public: - explicit NcclAsyncOpBase(OpKernelConstruction* c) : AsyncOpKernel(c) { - OP_REQUIRES_OK(c, c->GetAttr("num_devices", &num_devices_)); - OP_REQUIRES_OK(c, c->GetAttr("shared_name", &collective_prefix_)); - } - - string GetCollectiveKey(OpKernelContext* c) { - return strings::StrCat(collective_prefix_, ";", c->step_id(), ";", - c->frame_iter().frame_id, ":", - c->frame_iter().iter_id); - } - - int num_devices() const { return num_devices_; } - - private: - int num_devices_; - string collective_prefix_; - - TF_DISALLOW_COPY_AND_ASSIGN(NcclAsyncOpBase); -}; - -class NcclReduceOpBase : public NcclAsyncOpBase { - public: - explicit NcclReduceOpBase(OpKernelConstruction* c) : NcclAsyncOpBase(c) { - string reduction; - OP_REQUIRES_OK(c, c->GetAttr("reduction", &reduction)); - if (reduction == "min") { - reduction_op_ = ncclMin; - } else if (reduction == "max") { - reduction_op_ = ncclMax; - } else if (reduction == "sum") { - reduction_op_ = ncclSum; - } else if (reduction == "prod") { - reduction_op_ = ncclProd; - } else { - OP_REQUIRES_OK(c, - errors::InvalidArgument("Invalid reduction: ", reduction)); - } - } - - ncclRedOp_t reduction_op() const { return reduction_op_; } - - private: - ncclRedOp_t reduction_op_; -}; - -// To execute a single all-reduce, this kernel is called once for each of the -// devices in the communicator. -class NcclAllReduceOpKernel : public NcclReduceOpBase { - public: - explicit NcclAllReduceOpKernel(OpKernelConstruction* c) - : NcclReduceOpBase(c) {} - - void ComputeAsync(OpKernelContext* c, DoneCallback done) override { - const Tensor* in_t = &c->input(0); - Tensor* out_t; - OP_REQUIRES_OK_ASYNC(c, c->allocate_output(0, in_t->shape(), &out_t), done); - - auto actual_done = [c, done](Status s) { - OP_REQUIRES_OK_ASYNC(c, s, done); - done(); - }; - - auto* compute_stream = c->op_device_context()->stream(); - auto* gpu_info = c->device()->tensorflow_gpu_device_info(); - NcclManager::instance()->AddToAllReduce( - num_devices(), GetCollectiveKey(c), reduction_op(), - compute_stream->parent(), gpu_info->gpu_id, gpu_info->event_mgr, - compute_stream, in_t, out_t, std::move(actual_done)); - } -}; -REGISTER_KERNEL_BUILDER(Name("NcclAllReduce").Device(DEVICE_GPU), - NcclAllReduceOpKernel); - -// To execute a single reduce, this kernel is called once for all but one of the -// devices in the communicator, and NcclReduceRecvKernel is called once for -// the remaining device. -class NcclReduceSendKernel : public NcclReduceOpBase { - public: - explicit NcclReduceSendKernel(OpKernelConstruction* c) - : NcclReduceOpBase(c) {} - - void ComputeAsync(OpKernelContext* c, DoneCallback done) override { - auto actual_done = [c, done](Status s) { - OP_REQUIRES_OK_ASYNC(c, s, done); - done(); - }; - - auto* compute_stream = c->op_device_context()->stream(); - auto* gpu_info = c->device()->tensorflow_gpu_device_info(); - NcclManager::instance()->AddReduceSend( - num_devices(), GetCollectiveKey(c), reduction_op(), - compute_stream->parent(), gpu_info->gpu_id, gpu_info->event_mgr, - compute_stream, &c->input(0), std::move(actual_done)); - } -}; -REGISTER_KERNEL_BUILDER(Name("_NcclReduceSend").Device(DEVICE_GPU), - NcclReduceSendKernel); - -// To execute a single reduce, this kernel is called once for one devices, and -// NcclReduceSendKernel is called for all other devices in the -// communicator. -class NcclReduceRecvKernel : public NcclReduceOpBase { - public: - explicit NcclReduceRecvKernel(OpKernelConstruction* c) - : NcclReduceOpBase(c) {} - - void ComputeAsync(OpKernelContext* c, DoneCallback done) override { - const Tensor& in_t = c->input(0); - Tensor* out_t; - OP_REQUIRES_OK_ASYNC(c, c->allocate_output(0, in_t.shape(), &out_t), done); - - auto actual_done = [c, done](Status s) { - OP_REQUIRES_OK_ASYNC(c, s, done); - done(); - }; - - auto* compute_stream = c->op_device_context()->stream(); - auto* gpu_info = c->device()->tensorflow_gpu_device_info(); - NcclManager::instance()->AddReduceRecv( - num_devices(), GetCollectiveKey(c), reduction_op(), - compute_stream->parent(), gpu_info->gpu_id, gpu_info->event_mgr, - compute_stream, &in_t, out_t, std::move(actual_done)); - } - - private: - ncclRedOp_t reduction_op_; -}; -REGISTER_KERNEL_BUILDER(Name("_NcclReduceRecv").Device(DEVICE_GPU), - NcclReduceRecvKernel); - -// To execute a single broadcast, this kernel is called once for one device, and -// NcclBroadcastRecvKernel is called for all other devices in the -// communicator. -class NcclBroadcastSendKernel : public NcclAsyncOpBase { - public: - explicit NcclBroadcastSendKernel(OpKernelConstruction* c) - : NcclAsyncOpBase(c) {} - - void ComputeAsync(OpKernelContext* c, DoneCallback done) override { - auto actual_done = [c, done](Status s) { - OP_REQUIRES_OK_ASYNC(c, s, done); - done(); - }; - - auto* compute_stream = c->op_device_context()->stream(); - auto* gpu_info = c->device()->tensorflow_gpu_device_info(); - NcclManager::instance()->AddBroadcastSend( - num_devices(), GetCollectiveKey(c), compute_stream->parent(), - gpu_info->gpu_id, gpu_info->event_mgr, compute_stream, &c->input(0), - std::move(actual_done)); - } -}; -REGISTER_KERNEL_BUILDER(Name("_NcclBroadcastSend").Device(DEVICE_GPU), - NcclBroadcastSendKernel); - -// To execute a single broadcast, this kernel is called once for all but one of -// the devices in the communicator, and NcclBroadcastSendKernel is called -// once for the remaining device. -class NcclBroadcastRecvKernel : public NcclAsyncOpBase { - public: - explicit NcclBroadcastRecvKernel(OpKernelConstruction* c) - : NcclAsyncOpBase(c) {} - - void ComputeAsync(OpKernelContext* c, DoneCallback done) override { - const Tensor& shape_t = c->input(0); - TensorShape shape; - OP_REQUIRES_OK_ASYNC( - c, TensorShapeUtils::MakeShape(shape_t.vec(), &shape), done); - Tensor* out_t; - OP_REQUIRES_OK_ASYNC(c, c->allocate_output(0, shape, &out_t), done); - - auto actual_done = [c, done](Status s) { - OP_REQUIRES_OK_ASYNC(c, s, done); - done(); - }; - - auto* compute_stream = c->op_device_context()->stream(); - auto* gpu_info = c->device()->tensorflow_gpu_device_info(); - NcclManager::instance()->AddBroadcastRecv( - num_devices(), GetCollectiveKey(c), compute_stream->parent(), - gpu_info->gpu_id, gpu_info->event_mgr, compute_stream, out_t, - std::move(actual_done)); - } -}; -REGISTER_KERNEL_BUILDER( - Name("_NcclBroadcastRecv").Device(DEVICE_GPU).HostMemory("shape"), - NcclBroadcastRecvKernel); - -// Define stub kernels for the ops that get replaced post placement. -class NcclStubKernel : public AsyncOpKernel { - public: - explicit NcclStubKernel(OpKernelConstruction* c) : AsyncOpKernel(c) {} - void ComputeAsync(OpKernelContext* c, DoneCallback done) override { - c->SetStatus(errors::Unimplemented( - "This op should be replaced during graph optimization.")); - done(); - } -}; -REGISTER_KERNEL_BUILDER(Name("NcclBroadcast").Device(DEVICE_GPU), - NcclStubKernel); -REGISTER_KERNEL_BUILDER(Name("NcclReduce").Device(DEVICE_GPU), NcclStubKernel); - -} // namespace -} // namespace tensorflow - -#endif // GOOGLE_CUDA diff --git a/tensorflow/contrib/nccl/ops/nccl_ops.cc b/tensorflow/contrib/nccl/ops/nccl_ops.cc deleted file mode 100644 index a353a34b80add119fcdc8bc4230eddf5a77b30e8..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/nccl/ops/nccl_ops.cc +++ /dev/null @@ -1,185 +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 { - -using shape_inference::InferenceContext; -using shape_inference::ShapeHandle; - -REGISTER_OP("NcclAllReduce") - .Input("input: T") - .Output("data: T") - .Attr("reduction: {'min', 'max', 'prod', 'sum'}") - .Attr("T: {half, float, float64, int32, int64}") - .Attr("num_devices: int") - .Attr("shared_name: string") - .SetIsStateful() - .SetShapeFn(shape_inference::UnchangedShape) - .Doc(R"doc( -Outputs a tensor containing the reduction across all input tensors passed to ops -within the same `shared_name. - -The graph should be constructed so if one op runs with shared_name value `c`, -then `num_devices` ops will run with shared_name value `c`. Failure to do so -will cause the graph execution to fail to complete. - -input: the input to the reduction -data: the value of the reduction across all `num_devices` devices. -reduction: the reduction operation to perform. -num_devices: The number of devices participating in this reduction. -shared_name: Identifier that shared between ops of the same reduction. -)doc"); - -// Note: This op has no kernel implementation, but is replaced by -// _NcclReduceSend and _NcclReduceRecv during graph optimization stage. -REGISTER_OP("NcclReduce") - .Input("input: num_devices * T") - .Output("data: T") - .Attr("reduction: {'min', 'max', 'prod', 'sum'}") - .Attr("T: {half, float, float64, int32, int64}") - .Attr("num_devices: int") - .SetIsStateful() - .SetShapeFn(shape_inference::UnchangedShape) - .Doc(R"doc( -Reduces `input` from `num_devices` using `reduction` to a single device. - -The graph should be constructed so that all inputs have a valid device -assignment, and the op itself is assigned one of these devices. - -input: The input to the reduction. -data: the value of the reduction across all `num_devices` devices. -reduction: the reduction operation to perform. - )doc"); - -REGISTER_OP("_NcclReduceSend") - .Input("input: T") - .Attr("reduction: {'min', 'max', 'prod', 'sum'}") - .Attr("T: {half, float, float64, int32, int64}") - .Attr("num_devices: int") - .Attr("shared_name: string") - .SetIsStateful() - .SetShapeFn(shape_inference::NoOutputs) - .Doc(R"doc( -Replacement node for NcclReduce. - -Reduces `input` to the NcclReduceRecv op registered in the same `shared_name`. -The graph should be constructed so that 'num_devices-1' devices run -`_NcclReduceSend` and one device runs _NcclReduceRecv op with shared_name value -`c`. Failure to do so will cause the graph execution to fail to complete. - -input: The input to the reduction. -reduction: the reduction operation to perform. -num_devices: The number of devices participating in this reduction. -shared_name: Identifier that is shared between ops of the same reduce. - )doc"); - -REGISTER_OP("_NcclReduceRecv") - .Input("input: T") - .Output("data: T") - .Attr("reduction: {'min', 'max', 'prod', 'sum'}") - .Attr("T: {half, float, float64, int32, int64}") - .Attr("num_devices: int") - .Attr("shared_name: string") - .SetIsStateful() - .SetShapeFn(shape_inference::UnchangedShape) - .Doc(R"doc( -Replacement node for NcclReduce. - -Reduces 'input' from this op and the NcclReduceSend ops registered in the same -`shared_name`. -The graph should be constructed so that 'num_devices-1' devices run -`_NcclReduceSend` and one device runs _NcclReduceRecv op with shared_name value -`c`. Failure to do so will cause the graph execution to fail to complete. - -input: The input to the reduction. -data: The reduced data received from this op and the NcclReduceSend op. -reduction: the reduction operation to perform. -num_devices: The number of devices participating in this reduction. -shared_name: Identifier that is shared between ops of the same reduce. - )doc"); - -// Note: This op has no kernel implementation, but is replaced by -// _NcclBroadcastSend and _NcclBroadcastRecv during graph optimization stage. -REGISTER_OP("NcclBroadcast") - .Input("input: T") - .Output("output: T") - .Attr("T: {half, float, float64, int32, int64}") - .Attr("shape: shape") - .SetIsStateful() - .SetShapeFn(shape_inference::UnchangedShape) - .Doc(R"doc( -Sends `input` to all devices that are connected to the output. - -The graph should be constructed so that all ops connected to the output have a -valid device assignment, and the op itself is assigned one of these devices. - -input: The input to the broadcast. -output: The same as input. -shape: The shape of the input tensor. - )doc"); - -REGISTER_OP("_NcclBroadcastSend") - .Input("input: T") - .Attr("T: {half, float, float64, int32, int64}") - .Attr("num_devices: int") - .Attr("shared_name: string") - .SetIsStateful() - .SetShapeFn(shape_inference::NoOutputs) - .Doc(R"doc( -Replacement node for NcclBroadcast. - -Sends `input` to the _NcclBroadcastRecv ops registered in the same -`shared_name`. -The graph should be constructed so that one device runs `_NcclBroadcastSend` and -`num_devices-1` devices run _NcclBroadcastRecv ops with shared_name value `c`. -Failure to do so will cause the graph execution to fail to complete. - -input: The input to the broadcast. -num_devices: The number of devices participating in this reduction. -shared_name: Identifier that is shared between ops of the same broadcast. - )doc"); - -REGISTER_OP("_NcclBroadcastRecv") - .Input("shape: int32") - .Output("output: T") - .Attr("T: {half, float, float64, int32, int64}") - .Attr("num_devices: int") - .Attr("shared_name: string") - .SetIsStateful() - .SetShapeFn([](InferenceContext* c) { - ShapeHandle out; - TF_RETURN_IF_ERROR(c->MakeShapeFromShapeTensor(0, &out)); - c->set_output(0, out); - return Status::OK(); - }) - .Doc(R"doc( -Replacement node for NcclBroadcast. - -Sends data of shape `shape` from the _NcclBroadcastSend op registered in the -same `shared_name`. -The graph should be constructed so that one device runs `_NcclBroadcastSend` and -`num_devices-1` devices run _NcclBroadcastRecv ops with shared_name value `c`. -Failure to do so will cause the graph execution to fail to complete. - -shape: The shape of the output. -output: The broadcast data received from the NcclBroadcastSend op. -num_devices: The number of devices participating in this reduction. -shared_name: Identifier that is shared between ops of the same broadcast. - )doc"); - -} // namespace tensorflow diff --git a/tensorflow/contrib/nccl/python/ops/nccl_dependency_test.py b/tensorflow/contrib/nccl/python/ops/nccl_dependency_test.py deleted file mode 100644 index c766080dbee7c9a6f4383ef6fa8cade7bba158af..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/nccl/python/ops/nccl_dependency_test.py +++ /dev/null @@ -1,59 +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. -# ============================================================================== -"""Dependency test for nccl to test behavior when NCCL is not installed.""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -from tensorflow.contrib import nccl -from tensorflow.python.framework import constant_op -from tensorflow.python.framework import errors_impl -from tensorflow.python.framework import ops -from tensorflow.python.platform import test -from tensorflow.python.util import tf_inspect - - -class NcclDependencyTest(test.TestCase): - """Verifies that importing nccl ops lib does not fail even if NCCL is not - installed but nccl ops throws an exception on use if NCCL is not installed. - """ - - def test_nccl_ops(self): - """Tests behavior of nccl ops when NCCL is not installed.""" - - public_methods = [ - m[0] - for m in tf_inspect.getmembers(nccl, tf_inspect.isfunction) - if not m[0].startswith('_') - ] - for method_name in public_methods: - with ops.device('/device:CPU:0'): - tensor = constant_op.constant(1) - - if method_name == 'broadcast': - arg = tensor - else: - arg = [tensor] - - nccl_op = getattr(nccl, method_name) - with ops.device('/device:CPU:0'): - with self.assertRaisesRegexp(errors_impl.NotFoundError, - r'cannot open shared object file'): - nccl_op(arg) - - -if __name__ == '__main__': - test.main() diff --git a/tensorflow/contrib/opt/BUILD b/tensorflow/contrib/opt/BUILD index f4ac70eb1a720c2acc3ef942f269228156749cba..12320d9e456ae93cbf95639a0c9e0c7f414f3518 100644 --- a/tensorflow/contrib/opt/BUILD +++ b/tensorflow/contrib/opt/BUILD @@ -14,6 +14,7 @@ py_library( name = "opt_py", srcs = [ "__init__.py", + "python/training/adam_gs_optimizer.py", "python/training/adamax.py", "python/training/addsign.py", "python/training/agn_optimizer.py", @@ -22,6 +23,7 @@ py_library( "python/training/external_optimizer.py", "python/training/ggt.py", "python/training/lars_optimizer.py", + "python/training/lazy_adam_gs_optimizer.py", "python/training/lazy_adam_optimizer.py", "python/training/matrix_functions.py", "python/training/model_average_optimizer.py", @@ -60,6 +62,21 @@ py_library( ], ) +py_test( + name = "adam_gs_optimizer_test", + srcs = ["python/training/adam_gs_optimizer_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":opt_py", + "//tensorflow/python:array_ops", + "//tensorflow/python:client_testlib", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:math_ops", + "//tensorflow/python:training", + "//third_party/py/numpy", + ], +) + py_test( name = "adamax_test", srcs = ["python/training/adamax_test.py"], @@ -148,6 +165,25 @@ py_test( ], ) +py_test( + name = "lazy_adam_gs_optimizer_test", + srcs = ["python/training/lazy_adam_gs_optimizer_test.py"], + srcs_version = "PY2AND3", + deps = [ + ":opt_py", + "//tensorflow/python:array_ops", + "//tensorflow/python:client_testlib", + "//tensorflow/python:constant_op", + "//tensorflow/python:dtypes", + "//tensorflow/python:framework_ops", + "//tensorflow/python:math_ops", + "//tensorflow/python:resource_variable_ops", + "//tensorflow/python:variables", + "//third_party/py/numpy", + "@absl_py//absl/testing:parameterized", + ], +) + py_test( name = "lazy_adam_optimizer_test", srcs = ["python/training/lazy_adam_optimizer_test.py"], @@ -283,6 +319,9 @@ tf_py_test( "//tensorflow/python:framework_for_generated_wrappers", "//third_party/py/numpy", ], + tags = [ + "oss_serial", + ], ) tf_py_test( diff --git a/tensorflow/contrib/opt/__init__.py b/tensorflow/contrib/opt/__init__.py index c7ea68efa9a13a471bba3f41d0600855793b20a2..e8fc52342ceabb47da97ca0f3c8a01e419a221a1 100644 --- a/tensorflow/contrib/opt/__init__.py +++ b/tensorflow/contrib/opt/__init__.py @@ -19,6 +19,7 @@ from __future__ import division from __future__ import print_function # pylint: disable=wildcard-import +from tensorflow.contrib.opt.python.training.adam_gs_optimizer import * from tensorflow.contrib.opt.python.training.adamax import * from tensorflow.contrib.opt.python.training.addsign import * from tensorflow.contrib.opt.python.training.agn_optimizer import * @@ -28,6 +29,7 @@ from tensorflow.contrib.opt.python.training.external_optimizer import * from tensorflow.contrib.opt.python.training.lars_optimizer import * from tensorflow.contrib.opt.python.training.ggt import * from tensorflow.contrib.opt.python.training.lazy_adam_optimizer import * +from tensorflow.contrib.opt.python.training.lazy_adam_gs_optimizer import * from tensorflow.contrib.opt.python.training.model_average_optimizer import * from tensorflow.contrib.opt.python.training.moving_average_optimizer import * from tensorflow.contrib.opt.python.training.multitask_optimizer_wrapper import * @@ -44,12 +46,14 @@ from tensorflow.python.util.all_util import remove_undocumented _allowed_symbols = [ 'AdaMaxOptimizer', + 'AdamGSOptimizer', 'PowerSignOptimizer', 'AddSignOptimizer', 'DelayCompensatedGradientDescentOptimizer', 'DropStaleGradientOptimizer', 'ExternalOptimizerInterface', 'LARSOptimizer', + 'LazyAdamGSOptimizer', 'LazyAdamOptimizer', 'NadamOptimizer', 'MovingAverageOptimizer', diff --git a/tensorflow/contrib/opt/python/training/adam_gs_optimizer.py b/tensorflow/contrib/opt/python/training/adam_gs_optimizer.py new file mode 100644 index 0000000000000000000000000000000000000000..c5c9fc74deaf0171a33d0eb1b5c6f60b3aa5e533 --- /dev/null +++ b/tensorflow/contrib/opt/python/training/adam_gs_optimizer.py @@ -0,0 +1,214 @@ +# 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. +# ============================================================================== + +"""Adam rewrite to use global step for computing beta1 & beta2 accumulation.""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +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.ops import math_ops +from tensorflow.python.ops import resource_variable_ops +from tensorflow.python.ops import state_ops +from tensorflow.python.training import optimizer +from tensorflow.python.training import training_ops +from tensorflow.python.util.tf_export import tf_export + + +@tf_export("train.AdamOptimizer") +class AdamGSOptimizer(optimizer.Optimizer): + """Optimizer that implements the Adam algorithm. + + See [Kingma et al., 2014](http://arxiv.org/abs/1412.6980) + ([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"): + 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 + optimizer keep its own independent beta1 and beta2 accumulators as non-slot + variables. + + Initialization: + + $$m_0 := 0 \text{(Initialize initial 1st moment vector)}$$ + $$v_0 := 0 \text{(Initialize initial 2nd moment vector)}$$ + $$t := 0 \text{(Initialize timestep)}$$ + + The update rule for `variable` with gradient `g` uses an optimization + described at the end of section2 of the paper: + + $$t := t + 1$$ + $$lr_t := \text{learning\_rate} * \sqrt{1 - beta_2^t} / (1 - beta_1^t)$$ + + $$m_t := beta_1 * m_{t-1} + (1 - beta_1) * g$$ + $$v_t := beta_2 * v_{t-1} + (1 - beta_2) * g * g$$ + $$variable := variable - lr_t * m_t / (\sqrt{v_t} + \epsilon)$$ + + The default value of 1e-8 for epsilon might not be a good default in + general. For example, when training an Inception network on ImageNet a + current good choice is 1.0 or 0.1. Note that since AdamOptimizer uses the + formulation just before Section 2.1 of the Kingma and Ba paper rather than + the formulation in Algorithm 1, the "epsilon" referred to here is "epsilon + hat" in the paper. + + The sparse implementation of this algorithm (used when the gradient is an + IndexedSlices object, typically because of `tf.gather` or an embedding + lookup in the forward pass) does apply momentum to variable slices even if + they were not used in the forward pass (meaning they have a gradient equal + to zero). Momentum decay (beta1) is also applied to the entire momentum + accumulator. This means that the sparse behavior is equivalent to the dense + behavior (in contrast to some momentum implementations which ignore momentum + unless a variable slice was actually used). + + 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. + 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 + """ + super(AdamGSOptimizer, self).__init__(use_locking, name) + self._lr = learning_rate + self._beta1 = beta1 + self._beta2 = beta2 + self._epsilon = epsilon + self._global_step = global_step + self._global_step_on_worker = None + + # Tensor versions of the constructor arguments, created in _prepare(). + self._lr_t = None + self._beta1_t = None + 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)) + + def _create_slots(self, var_list): + # Create slots for the first and second moments. + for v in var_list: + self._zeros_slot(v, "m", self._name) + self._zeros_slot(v, "v", self._name) + + def _prepare(self): + lr = self._call_if_callable(self._lr) + beta1 = self._call_if_callable(self._beta1) + beta2 = self._call_if_callable(self._beta2) + epsilon = self._call_if_callable(self._epsilon) + + self._lr_t = ops.convert_to_tensor(lr, name="learning_rate") + self._beta1_t = ops.convert_to_tensor(beta1, name="beta1") + self._beta2_t = ops.convert_to_tensor(beta2, name="beta2") + self._epsilon_t = ops.convert_to_tensor(epsilon, name="epsilon") + + # Performance optimization so that worker creates a copy of the global step + # to avoid overloading the parameter server holding the global step. + self._global_step_on_worker = math_ops.cast( + array_ops.identity(self._global_step) + 1, dtypes.float32) + + def _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.apply_adam( + 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 + + 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, + 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) + + def _apply_sparse_shared(self, grad, var, indices, scatter_add): + beta1_power, beta2_power = self._get_beta_accumulators() + beta1_power = math_ops.cast(beta1_power, var.dtype.base_dtype) + beta2_power = math_ops.cast(beta2_power, var.dtype.base_dtype) + lr_t = math_ops.cast(self._lr_t, var.dtype.base_dtype) + beta1_t = math_ops.cast(self._beta1_t, var.dtype.base_dtype) + beta2_t = math_ops.cast(self._beta2_t, var.dtype.base_dtype) + epsilon_t = math_ops.cast(self._epsilon_t, var.dtype.base_dtype) + lr = (lr_t * math_ops.sqrt(1 - beta2_power) / (1 - beta1_power)) + # 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) + 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) + v = self.get_slot(var, "v") + v_scaled_g_values = (grad * grad) * (1 - beta2_t) + v_t = state_ops.assign(v, v * beta2_t, use_locking=self._use_locking) + 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) + 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, + lambda x, i, v: state_ops.scatter_add( # pylint: disable=g-long-lambda + 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)]): + return x.value() + + def _resource_apply_sparse(self, grad, var, indices): + return self._apply_sparse_shared( + grad, var, indices, self._resource_scatter_add) diff --git a/tensorflow/contrib/opt/python/training/adam_gs_optimizer_test.py b/tensorflow/contrib/opt/python/training/adam_gs_optimizer_test.py new file mode 100644 index 0000000000000000000000000000000000000000..c68c965aef3729bebe7d0e0dd707c344321d9e3f --- /dev/null +++ b/tensorflow/contrib/opt/python/training/adam_gs_optimizer_test.py @@ -0,0 +1,382 @@ +# 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 AdamGS.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.contrib.opt.python.training import adam_gs_optimizer +from tensorflow.python.client import session +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 +from tensorflow.python.framework import test_util +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import math_ops +from tensorflow.python.ops import resource_variable_ops +from tensorflow.python.ops import variables +from tensorflow.python.platform import test + + +def adam_update_numpy(param, + g_t, + t, + m, + v, + alpha=0.001, + beta1=0.9, + beta2=0.999, + epsilon=1e-8): + alpha_t = alpha * np.sqrt(1 - beta2**t) / (1 - beta1**t) + + m_t = beta1 * m + (1 - beta1) * g_t + v_t = beta2 * v + (1 - beta2) * g_t * g_t + + param_t = param - alpha_t * m_t / (np.sqrt(v_t) + epsilon) + return param_t, m_t, v_t + + +class AdamGSOptimizerTest(test.TestCase): + + def doTestSparse(self, use_resource=False): + for dtype in [dtypes.half, dtypes.float32, dtypes.float64]: + with self.cached_session(): + # Initialize variables for numpy implementation. + m0, v0, m1, v1 = 0.0, 0.0, 0.0, 0.0 + var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype) + grads0_np = np.array([0.1, 0.1], dtype=dtype.as_numpy_dtype) + var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype) + grads1_np = np.array([0.01, 0.01], dtype=dtype.as_numpy_dtype) + + if use_resource: + global_step = resource_variable_ops.ResourceVariable( + array_ops.zeros([], dtypes.int64)) + var0 = resource_variable_ops.ResourceVariable(var0_np) + var1 = resource_variable_ops.ResourceVariable(var1_np) + else: + global_step = variables.Variable(array_ops.zeros([], dtypes.int64)) + var0 = variables.Variable(var0_np) + var1 = variables.Variable(var1_np) + grads0_np_indices = np.array([0, 1], dtype=np.int32) + grads0 = ops.IndexedSlices( + constant_op.constant(grads0_np), + constant_op.constant(grads0_np_indices), constant_op.constant([2])) + grads1_np_indices = np.array([0, 1], dtype=np.int32) + grads1 = ops.IndexedSlices( + constant_op.constant(grads1_np), + constant_op.constant(grads1_np_indices), constant_op.constant([2])) + opt = adam_gs_optimizer.AdamGSOptimizer(global_step=global_step) + update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]), + global_step=global_step) + variables.global_variables_initializer().run() + + # Fetch params to validate initial values + self.assertAllClose([1.0, 2.0], self.evaluate(var0)) + self.assertAllClose([3.0, 4.0], self.evaluate(var1)) + + beta1_power, beta2_power = opt._get_beta_accumulators() + + # Run 3 steps of Adam + for t in range(1, 4): + self.assertAllCloseAccordingToType(0.9**t, self.evaluate(beta1_power)) + self.assertAllCloseAccordingToType(0.999**t, + self.evaluate(beta2_power)) + update.run() + + var0_np, m0, v0 = adam_update_numpy(var0_np, grads0_np, t, m0, v0) + var1_np, m1, v1 = adam_update_numpy(var1_np, grads1_np, t, m1, v1) + + # Validate updated params + self.assertAllCloseAccordingToType(var0_np, self.evaluate(var0)) + self.assertAllCloseAccordingToType(var1_np, self.evaluate(var1)) + + def testSparse(self): + self.doTestSparse(use_resource=False) + + def testResourceSparse(self): + self.doTestSparse(use_resource=True) + + def testSparseDevicePlacement(self): + for index_dtype in [dtypes.int32, dtypes.int64]: + with self.cached_session(force_gpu=test.is_gpu_available()): + # If a GPU is available, tests that all optimizer ops can be placed on + # it (i.e. they have GPU kernels). + var = variables.Variable([[1.0], [2.0]]) + indices = constant_op.constant([0, 1], dtype=index_dtype) + gathered_sum = math_ops.reduce_sum(array_ops.gather(var, indices)) + optimizer = adam_gs_optimizer.AdamGSOptimizer(3.0) + minimize_op = optimizer.minimize(gathered_sum) + variables.global_variables_initializer().run() + minimize_op.run() + + def testSparseRepeatedIndices(self): + for dtype in [dtypes.half, dtypes.float32, dtypes.float64]: + with self.cached_session(): + repeated_index_global_step = variables.Variable( + array_ops.zeros([], dtypes.int64)) + aggregated_global_step = variables.Variable( + array_ops.zeros([], dtypes.int64)) + repeated_index_update_var = variables.Variable( + [[1.0], [2.0]], dtype=dtype) + aggregated_update_var = variables.Variable( + [[1.0], [2.0]], dtype=dtype) + grad_repeated_index = ops.IndexedSlices( + constant_op.constant( + [0.1, 0.1], shape=[2, 1], dtype=dtype), + constant_op.constant([1, 1]), + constant_op.constant([2, 1])) + grad_aggregated = ops.IndexedSlices( + constant_op.constant( + [0.2], shape=[1, 1], dtype=dtype), + constant_op.constant([1]), + constant_op.constant([2, 1])) + repeated_update = adam_gs_optimizer.AdamGSOptimizer( + global_step=repeated_index_global_step).apply_gradients( + [(grad_repeated_index, repeated_index_update_var)], + global_step=repeated_index_global_step) + aggregated_update = adam_gs_optimizer.AdamGSOptimizer( + global_step=aggregated_global_step).apply_gradients( + [(grad_aggregated, aggregated_update_var)], + global_step=aggregated_global_step) + variables.global_variables_initializer().run() + self.assertAllClose(aggregated_update_var.eval(), + self.evaluate(repeated_index_update_var)) + for _ in range(3): + repeated_update.run() + aggregated_update.run() + self.assertAllClose(aggregated_update_var.eval(), + self.evaluate(repeated_index_update_var)) + + def doTestBasic(self, use_resource=False, use_callable_params=False): + for i, dtype in enumerate([dtypes.half, dtypes.float32, dtypes.float64]): + with self.session(graph=ops.Graph()): + # Initialize variables for numpy implementation. + m0, v0, m1, v1 = 0.0, 0.0, 0.0, 0.0 + var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype) + grads0_np = np.array([0.1, 0.1], dtype=dtype.as_numpy_dtype) + var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype) + grads1_np = np.array([0.01, 0.01], dtype=dtype.as_numpy_dtype) + + if use_resource: + global_step = resource_variable_ops.ResourceVariable( + array_ops.zeros([], dtypes.int64), name="global_step_%d" % i) + var0 = resource_variable_ops.ResourceVariable( + var0_np, name="var0_%d" % i) + var1 = resource_variable_ops.ResourceVariable( + var1_np, name="var1_%d" % i) + else: + global_step = variables.Variable(array_ops.zeros([], dtypes.int64)) + var0 = variables.Variable(var0_np) + var1 = variables.Variable(var1_np) + grads0 = constant_op.constant(grads0_np) + grads1 = constant_op.constant(grads1_np) + + learning_rate = lambda: 0.001 + beta1 = lambda: 0.9 + beta2 = lambda: 0.999 + epsilon = lambda: 1e-8 + if not use_callable_params: + learning_rate = learning_rate() + beta1 = beta1() + beta2 = beta2() + epsilon = epsilon() + + opt = adam_gs_optimizer.AdamGSOptimizer(global_step=global_step, + learning_rate=learning_rate) + update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]), + global_step=global_step) + opt_variables = opt.variables() + beta1_power, beta2_power = opt._get_beta_accumulators() + self.assertTrue(beta1_power is not None) + self.assertTrue(beta2_power is not None) + self.assertNotIn(beta1_power, opt_variables) + self.assertNotIn(beta2_power, opt_variables) + + if not context.executing_eagerly(): + with ops.Graph().as_default(): + # Shouldn't return non-slot variables from other graphs. + self.assertEqual(0, len(opt.variables())) + self.evaluate(variables.global_variables_initializer()) + # Fetch params to validate initial values + self.assertAllClose([1.0, 2.0], self.evaluate(var0)) + self.assertAllClose([3.0, 4.0], self.evaluate(var1)) + + # Run 3 steps of Adam + for t in range(1, 4): + if not context.executing_eagerly(): + self.evaluate(update) + self.assertAllCloseAccordingToType( + 0.9**(t + 1), self.evaluate(beta1_power)) + self.assertAllCloseAccordingToType( + 0.999**(t + 1), self.evaluate(beta2_power)) + else: + if t > 1: + opt.apply_gradients(zip([grads0, grads1], [var0, var1]), + global_step=global_step) + beta1_power, beta2_power = opt._get_beta_accumulators() + self.assertAllCloseAccordingToType( + 0.9**t, self.evaluate(beta1_power)) + self.assertAllCloseAccordingToType( + 0.999**t, self.evaluate(beta2_power)) + + var0_np, m0, v0 = adam_update_numpy(var0_np, grads0_np, t, m0, v0) + var1_np, m1, v1 = adam_update_numpy(var1_np, grads1_np, t, m1, v1) + + # Validate updated params + self.assertAllCloseAccordingToType(var0_np, self.evaluate(var0)) + self.assertAllCloseAccordingToType(var1_np, self.evaluate(var1)) + if use_resource: + self.assertEqual("var0_%d/Adam:0" % (i,), + opt.get_slot(var=var0, name="m").name) + + def testBasic(self): + with self.cached_session(): + self.doTestBasic(use_resource=False) + + @test_util.run_in_graph_and_eager_modes(reset_test=True) + def testResourceBasic(self): + self.doTestBasic(use_resource=True) + + def testBasicCallableParams(self): + with context.eager_mode(): + self.doTestBasic(use_resource=True, use_callable_params=True) + + def testTensorLearningRate(self): + for dtype in [dtypes.half, dtypes.float32, dtypes.float64]: + with self.cached_session(): + global_step = variables.Variable(array_ops.zeros([], dtypes.int64)) + # Initialize variables for numpy implementation. + m0, v0, m1, v1 = 0.0, 0.0, 0.0, 0.0 + var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype) + grads0_np = np.array([0.1, 0.1], dtype=dtype.as_numpy_dtype) + var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype) + grads1_np = np.array([0.01, 0.01], dtype=dtype.as_numpy_dtype) + + var0 = variables.Variable(var0_np) + var1 = variables.Variable(var1_np) + grads0 = constant_op.constant(grads0_np) + grads1 = constant_op.constant(grads1_np) + opt = adam_gs_optimizer.AdamGSOptimizer( + global_step=global_step, learning_rate=constant_op.constant(0.001)) + update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]), + global_step=global_step) + variables.global_variables_initializer().run() + + # Fetch params to validate initial values + self.assertAllClose([1.0, 2.0], self.evaluate(var0)) + self.assertAllClose([3.0, 4.0], self.evaluate(var1)) + + beta1_power, beta2_power = opt._get_beta_accumulators() + + # Run 3 steps of Adam + for t in range(1, 4): + self.assertAllCloseAccordingToType(0.9**t, self.evaluate(beta1_power)) + self.assertAllCloseAccordingToType(0.999**t, + self.evaluate(beta2_power)) + update.run() + + var0_np, m0, v0 = adam_update_numpy(var0_np, grads0_np, t, m0, v0) + var1_np, m1, v1 = adam_update_numpy(var1_np, grads1_np, t, m1, v1) + + # Validate updated params + self.assertAllCloseAccordingToType(var0_np, self.evaluate(var0)) + self.assertAllCloseAccordingToType(var1_np, self.evaluate(var1)) + + def testSharing(self): + for dtype in [dtypes.half, dtypes.float32, dtypes.float64]: + with self.cached_session(): + global_step = variables.Variable(array_ops.zeros([], dtypes.int64)) + # Initialize variables for numpy implementation. + m0, v0, m1, v1 = 0.0, 0.0, 0.0, 0.0 + var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype) + grads0_np = np.array([0.1, 0.1], dtype=dtype.as_numpy_dtype) + var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype) + grads1_np = np.array([0.01, 0.01], dtype=dtype.as_numpy_dtype) + + var0 = variables.Variable(var0_np) + var1 = variables.Variable(var1_np) + grads0 = constant_op.constant(grads0_np) + grads1 = constant_op.constant(grads1_np) + opt = adam_gs_optimizer.AdamGSOptimizer(global_step=global_step) + update1 = opt.apply_gradients(zip([grads0, grads1], [var0, var1]), + global_step=global_step) + update2 = opt.apply_gradients(zip([grads0, grads1], [var0, var1]), + global_step=global_step) + variables.global_variables_initializer().run() + + beta1_power, beta2_power = opt._get_beta_accumulators() + + # Fetch params to validate initial values + self.assertAllClose([1.0, 2.0], self.evaluate(var0)) + self.assertAllClose([3.0, 4.0], self.evaluate(var1)) + + # Run 3 steps of intertwined Adam1 and Adam2. + for t in range(1, 4): + self.assertAllCloseAccordingToType(0.9**t, self.evaluate(beta1_power)) + self.assertAllCloseAccordingToType(0.999**t, + self.evaluate(beta2_power)) + if t % 2 == 0: + update1.run() + else: + update2.run() + + var0_np, m0, v0 = adam_update_numpy(var0_np, grads0_np, t, m0, v0) + var1_np, m1, v1 = adam_update_numpy(var1_np, grads1_np, t, m1, v1) + + # Validate updated params + self.assertAllCloseAccordingToType(var0_np, self.evaluate(var0)) + self.assertAllCloseAccordingToType(var1_np, self.evaluate(var1)) + + def testTwoSessions(self): + optimizer = adam_gs_optimizer.AdamGSOptimizer() + + with context.eager_mode(): + var0 = variables.Variable(np.array([1.0, 2.0]), name="v0") + grads0 = constant_op.constant(np.array([0.1, 0.1])) + optimizer.apply_gradients([(grads0, var0)]) + + g = ops.Graph() + with g.as_default(): + with session.Session(): + var0 = variables.Variable(np.array([1.0, 2.0]), name="v0") + grads0 = constant_op.constant(np.array([0.1, 0.1])) + optimizer.apply_gradients([(grads0, var0)]) + + gg = ops.Graph() + with gg.as_default(): + with session.Session(): + var0 = variables.Variable(np.array([1.0, 2.0]), name="v0") + grads0 = constant_op.constant(np.array([0.1, 0.1])) + + # If the optimizer saves any state not keyed by graph the following line + # fails. + optimizer.apply_gradients([(grads0, var0)]) + + def testSlotsUniqueEager(self): + with context.eager_mode(): + v1 = resource_variable_ops.ResourceVariable(1.) + v2 = resource_variable_ops.ResourceVariable(1.) + opt = adam_gs_optimizer.AdamGSOptimizer(1.) + opt.minimize(lambda: v1 + v2) + # There should be two unique slot variables for v1 and v2 respectively. + self.assertEqual(4, len(set(opt.variables()))) + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/opt/python/training/elastic_average_optimizer.py b/tensorflow/contrib/opt/python/training/elastic_average_optimizer.py index 6c203e5519e6a66d20e2509eca3c74eb66bf32c7..fa1a7aaff0aa59a6a64b1f0bf836a273926d785d 100644 --- a/tensorflow/contrib/opt/python/training/elastic_average_optimizer.py +++ b/tensorflow/contrib/opt/python/training/elastic_average_optimizer.py @@ -30,6 +30,7 @@ from tensorflow.python.ops import variables from tensorflow.python.training import optimizer from tensorflow.python.training import saver from tensorflow.python.training import session_run_hook +from tensorflow.python.training.saving import saveable_object_util LOCAL_VARIABLE_NAME = 'local_center_variable' GLOBAL_VARIABLE_NAME = 'global_center_variable' @@ -424,7 +425,7 @@ class ElasticAverageOptimizer(optimizer.Optimizer): if var_list is None: var_list = variables.trainable_variables() if not isinstance(var_list, dict): - var_list = saver.BaseSaverBuilder.OpListToDict(var_list) + var_list = saveable_object_util.op_list_to_dict(var_list) swapped_var_list = {} for key, var in var_list.items(): @@ -464,4 +465,4 @@ class _ElasticAverageOptimizerHook(session_run_hook.SessionRunHook): def after_create_session(self, session, coord): """Run initialization ops""" - session.run(self._variable_init_op) \ No newline at end of file + session.run(self._variable_init_op) diff --git a/tensorflow/contrib/opt/python/training/external_optimizer.py b/tensorflow/contrib/opt/python/training/external_optimizer.py index 82ebca7f20306e5658c8321716e39f9c7f8b8970..e5e52f7dc3a70892322c65ac968c14a9c3115df4 100644 --- a/tensorflow/contrib/opt/python/training/external_optimizer.py +++ b/tensorflow/contrib/opt/python/training/external_optimizer.py @@ -429,7 +429,7 @@ def _accumulate(list_): def _get_shape_tuple(tensor): - return tuple(dim.value for dim in tensor.get_shape()) + return tuple(tensor.get_shape().as_list()) def _prod(array): diff --git a/tensorflow/contrib/opt/python/training/ggt.py b/tensorflow/contrib/opt/python/training/ggt.py index cae952d8f50acbc3a176697fb3989db6c9ac3e9b..6dc17fe5a5210fa1700e1382016e40fa0a792917 100644 --- a/tensorflow/contrib/opt/python/training/ggt.py +++ b/tensorflow/contrib/opt/python/training/ggt.py @@ -21,6 +21,7 @@ import collections import numpy as np from tensorflow.contrib.optimizer_v2 import optimizer_v2 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 linalg_ops @@ -120,7 +121,7 @@ class GGTOptimizer(optimizer_v2.OptimizerV2): # Construct ordered dictionary for variable dimensions, sorted by name. shape_dict = {} for v in var_list: - shape_dict[v.name] = np.prod(v.get_shape()).value + shape_dict[v.name] = tensor_shape.dimension_value(np.prod(v.get_shape())) self.shape_dict = collections.OrderedDict( sorted(shape_dict.items(), key=lambda t: t[0])) diff --git a/tensorflow/contrib/opt/python/training/lars_optimizer.py b/tensorflow/contrib/opt/python/training/lars_optimizer.py index a8dafd9a4cb9c669400f74b545b3c165bd49b2a2..bc18177b6d0b1d3f4fc58236bbc3d445fb73d80d 100644 --- a/tensorflow/contrib/opt/python/training/lars_optimizer.py +++ b/tensorflow/contrib/opt/python/training/lars_optimizer.py @@ -18,6 +18,7 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import linalg_ops from tensorflow.python.ops import math_ops @@ -162,3 +163,14 @@ class LARSOptimizer(optimizer.Optimizer): math_ops.cast(self._momentum_tensor, grad.dtype), use_locking=self._use_locking, use_nesterov=self._use_nesterov) + + def _prepare(self): + learning_rate = self._learning_rate + if callable(learning_rate): + learning_rate = learning_rate() + self._learning_rate_tensor = ops.convert_to_tensor( + learning_rate, name="learning_rate") + momentum = self._momentum + if callable(momentum): + momentum = momentum() + self._momentum_tensor = ops.convert_to_tensor(momentum, name="momentum") diff --git a/tensorflow/contrib/opt/python/training/lazy_adam_gs_optimizer.py b/tensorflow/contrib/opt/python/training/lazy_adam_gs_optimizer.py new file mode 100644 index 0000000000000000000000000000000000000000..8827007e4d7f6722398a8e36bd626377842d92ef --- /dev/null +++ b/tensorflow/contrib/opt/python/training/lazy_adam_gs_optimizer.py @@ -0,0 +1,114 @@ +# 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. +# ============================================================================== + +"""LazyAdam rewrite to use global step for computing beta1 & beta2 accumulation. +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.opt.python.training import adam_gs_optimizer +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 resource_variable_ops +from tensorflow.python.ops import state_ops + + +class LazyAdamGSOptimizer(adam_gs_optimizer.AdamGSOptimizer): + """Variant of the Adam optimizer that handles sparse updates more efficiently. + + Branched from tf.contrib.opt.LazyAdamGSOptimizer. The only difference is to + pass global step for computing beta1 and beta2 accumulators, instead of having + optimizer keep its own independent beta1 and beta2 accumulators as non-slot + variables. + + The original Adam algorithm maintains two moving-average accumulators for + each trainable variable; the accumulators are updated at every step. + This class provides lazier handling of gradient updates for sparse variables. + It only updates moving-average accumulators for sparse variable indices that + appear in the current batch, rather than updating the accumulators for all + indices. Compared with the original Adam optimizer, it can provide large + improvements in model training throughput for some applications. However, it + provides slightly different semantics than the original Adam algorithm, and + may lead to different empirical results. + """ + + def _apply_sparse(self, grad, var): + beta1_power, beta2_power = self._get_beta_accumulators() + beta1_power = math_ops.cast(beta1_power, var.dtype.base_dtype) + beta2_power = math_ops.cast(beta2_power, var.dtype.base_dtype) + lr_t = math_ops.cast(self._lr_t, var.dtype.base_dtype) + beta1_t = math_ops.cast(self._beta1_t, var.dtype.base_dtype) + beta2_t = math_ops.cast(self._beta2_t, var.dtype.base_dtype) + epsilon_t = math_ops.cast(self._epsilon_t, var.dtype.base_dtype) + lr = (lr_t * math_ops.sqrt(1 - beta2_power) / (1 - beta1_power)) + + # \\(m := beta1 * m + (1 - beta1) * g_t\\) + m = self.get_slot(var, "m") + m_t = state_ops.scatter_update(m, grad.indices, + beta1_t * array_ops.gather(m, grad.indices) + + (1 - beta1_t) * grad.values, + use_locking=self._use_locking) + + # \\(v := beta2 * v + (1 - beta2) * (g_t * g_t)\\) + v = self.get_slot(var, "v") + v_t = state_ops.scatter_update(v, grad.indices, + beta2_t * array_ops.gather(v, grad.indices) + + (1 - beta2_t) * math_ops.square(grad.values), + use_locking=self._use_locking) + + # \\(variable -= learning_rate * m_t / (epsilon_t + sqrt(v_t))\\) + m_t_slice = array_ops.gather(m_t, grad.indices) + v_t_slice = array_ops.gather(v_t, grad.indices) + denominator_slice = math_ops.sqrt(v_t_slice) + epsilon_t + var_update = state_ops.scatter_sub(var, grad.indices, + lr * m_t_slice / denominator_slice, + use_locking=self._use_locking) + return control_flow_ops.group(var_update, m_t, v_t) + + def _resource_apply_sparse(self, grad, var, indices): + beta1_power, beta2_power = self._get_beta_accumulators() + beta1_power = math_ops.cast(beta1_power, var.dtype.base_dtype) + beta2_power = math_ops.cast(beta2_power, var.dtype.base_dtype) + lr_t = math_ops.cast(self._lr_t, var.dtype.base_dtype) + beta1_t = math_ops.cast(self._beta1_t, var.dtype.base_dtype) + beta2_t = math_ops.cast(self._beta2_t, var.dtype.base_dtype) + epsilon_t = math_ops.cast(self._epsilon_t, var.dtype.base_dtype) + lr = (lr_t * math_ops.sqrt(1 - beta2_power) / (1 - beta1_power)) + + # \\(m := beta1 * m + (1 - beta1) * g_t\\) + m = self.get_slot(var, "m") + m_t_slice = beta1_t * array_ops.gather(m, indices) + (1 - beta1_t) * grad + m_update_op = resource_variable_ops.resource_scatter_update(m.handle, + indices, + m_t_slice) + + # \\(v := beta2 * v + (1 - beta2) * (g_t * g_t)\\) + v = self.get_slot(var, "v") + v_t_slice = (beta2_t * array_ops.gather(v, indices) + + (1 - beta2_t) * math_ops.square(grad)) + v_update_op = resource_variable_ops.resource_scatter_update(v.handle, + indices, + v_t_slice) + + # \\(variable -= learning_rate * m_t / (epsilon_t + sqrt(v_t))\\) + var_slice = lr * m_t_slice / (math_ops.sqrt(v_t_slice) + epsilon_t) + var_update_op = resource_variable_ops.resource_scatter_sub(var.handle, + indices, + var_slice) + + return control_flow_ops.group(var_update_op, m_update_op, v_update_op) diff --git a/tensorflow/contrib/opt/python/training/lazy_adam_gs_optimizer_test.py b/tensorflow/contrib/opt/python/training/lazy_adam_gs_optimizer_test.py new file mode 100644 index 0000000000000000000000000000000000000000..bdc9a02a546c8399172d0c5b58941b4d80179955 --- /dev/null +++ b/tensorflow/contrib/opt/python/training/lazy_adam_gs_optimizer_test.py @@ -0,0 +1,402 @@ +# 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 LazyAdamGSOptimizer.""" + +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.opt.python.training import lazy_adam_gs_optimizer +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 +from tensorflow.python.framework import test_util +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import math_ops +from tensorflow.python.ops import resource_variable_ops +from tensorflow.python.ops import variables +from tensorflow.python.platform import test + + +def adam_update_numpy(param, + g_t, + t, + m, + v, + alpha=0.001, + beta1=0.9, + beta2=0.999, + epsilon=1e-8): + alpha_t = alpha * np.sqrt(1 - beta2**t) / (1 - beta1**t) + + m_t = beta1 * m + (1 - beta1) * g_t + v_t = beta2 * v + (1 - beta2) * g_t * g_t + + param_t = param - alpha_t * m_t / (np.sqrt(v_t) + epsilon) + return param_t, m_t, v_t + + +class LazyAdamGSOptimizerTest(test.TestCase, parameterized.TestCase): + + @parameterized.parameters([False, True]) + def testSparse(self, use_resource): + for dtype in [dtypes.half, dtypes.float32, dtypes.float64]: + with self.cached_session(): + # Initialize variables for numpy implementation. + m0, v0, m1, v1 = 0.0, 0.0, 0.0, 0.0 + var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype) + grads0_np = np.array([0.1, 0.1], dtype=dtype.as_numpy_dtype) + var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype) + grads1_np = np.array([0.01, 0.01], dtype=dtype.as_numpy_dtype) + + if use_resource: + global_step = resource_variable_ops.ResourceVariable( + array_ops.zeros([], dtypes.int64)) + var0 = resource_variable_ops.ResourceVariable(var0_np) + var1 = resource_variable_ops.ResourceVariable(var1_np) + else: + global_step = variables.Variable(array_ops.zeros([], dtypes.int64)) + var0 = variables.Variable(var0_np) + var1 = variables.Variable(var1_np) + + grads0_np_indices = np.array([0, 1], dtype=np.int32) + grads0 = ops.IndexedSlices( + constant_op.constant(grads0_np), + constant_op.constant(grads0_np_indices), constant_op.constant([2])) + grads1_np_indices = np.array([0, 1], dtype=np.int32) + grads1 = ops.IndexedSlices( + constant_op.constant(grads1_np), + constant_op.constant(grads1_np_indices), constant_op.constant([2])) + opt = lazy_adam_gs_optimizer.LazyAdamGSOptimizer( + global_step=global_step) + update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]), + global_step=global_step) + variables.global_variables_initializer().run() + + # Fetch params to validate initial values + self.assertAllClose([1.0, 2.0], var0.eval()) + self.assertAllClose([3.0, 4.0], var1.eval()) + + beta1_power, beta2_power = opt._get_beta_accumulators() + + # Run 3 steps of Adam + for t in range(1, 4): + self.assertAllCloseAccordingToType(0.9**t, beta1_power.eval()) + self.assertAllCloseAccordingToType(0.999**t, beta2_power.eval()) + update.run() + + var0_np, m0, v0 = adam_update_numpy(var0_np, grads0_np, t, m0, v0) + var1_np, m1, v1 = adam_update_numpy(var1_np, grads1_np, t, m1, v1) + + # Validate updated params + self.assertAllCloseAccordingToType(var0_np, var0.eval()) + self.assertAllCloseAccordingToType(var1_np, var1.eval()) + + @parameterized.parameters([False, True]) + def testSparseDevicePlacement(self, use_resource): + for index_dtype in [dtypes.int32, dtypes.int64]: + with self.cached_session(force_gpu=test.is_gpu_available()): + # If a GPU is available, tests that all optimizer ops can be placed on + # it (i.e. they have GPU kernels). + if use_resource: + global_step = resource_variable_ops.ResourceVariable( + array_ops.zeros([], dtypes.int64)) + var = resource_variable_ops.ResourceVariable([[1.0], [2.0]]) + else: + global_step = variables.Variable(array_ops.zeros([], dtypes.int64)) + var = variables.Variable([[1.0], [2.0]]) + + indices = constant_op.constant([0, 1], dtype=index_dtype) + gathered_sum = math_ops.reduce_sum(array_ops.gather(var, indices)) + optimizer = lazy_adam_gs_optimizer.LazyAdamGSOptimizer( + global_step=global_step, learning_rate=3.0) + minimize_op = optimizer.minimize(gathered_sum, global_step=global_step) + variables.global_variables_initializer().run() + minimize_op.run() + + @parameterized.parameters([False, True]) + def testSparseRepeatedIndices(self, use_resource): + for dtype in [dtypes.half, dtypes.float32, dtypes.float64]: + with self.cached_session(): + if use_resource: + repeated_index_global_step = resource_variable_ops.ResourceVariable( + array_ops.zeros([], dtypes.int64)) + aggregated_global_step = resource_variable_ops.ResourceVariable( + array_ops.zeros([], dtypes.int64)) + repeated_index_update_var = resource_variable_ops.ResourceVariable( + [[1.0], [2.0]], dtype=dtype) + aggregated_update_var = resource_variable_ops.ResourceVariable( + [[1.0], [2.0]], dtype=dtype) + else: + repeated_index_global_step = variables.Variable( + array_ops.zeros([], dtypes.int64)) + aggregated_global_step = variables.Variable( + array_ops.zeros([], dtypes.int64)) + repeated_index_update_var = variables.Variable( + [[1.0], [2.0]], dtype=dtype) + aggregated_update_var = variables.Variable( + [[1.0], [2.0]], dtype=dtype) + + grad_repeated_index = ops.IndexedSlices( + constant_op.constant( + [0.1, 0.1], shape=[2, 1], dtype=dtype), + constant_op.constant([1, 1]), + constant_op.constant([2, 1])) + grad_aggregated = ops.IndexedSlices( + constant_op.constant( + [0.2], shape=[1, 1], dtype=dtype), + constant_op.constant([1]), + constant_op.constant([2, 1])) + repeated_update_opt = lazy_adam_gs_optimizer.LazyAdamGSOptimizer( + global_step=repeated_index_global_step) + repeated_update = repeated_update_opt.apply_gradients( + [(grad_repeated_index, repeated_index_update_var)], + global_step=repeated_index_global_step) + aggregated_update_opt = lazy_adam_gs_optimizer.LazyAdamGSOptimizer( + global_step=aggregated_global_step) + aggregated_update = aggregated_update_opt.apply_gradients( + [(grad_aggregated, aggregated_update_var)], + global_step=aggregated_global_step) + variables.global_variables_initializer().run() + self.assertAllClose(aggregated_update_var.eval(), + repeated_index_update_var.eval()) + for _ in range(3): + repeated_update.run() + aggregated_update.run() + self.assertAllClose(aggregated_update_var.eval(), + repeated_index_update_var.eval()) + + def doTestBasic(self, use_resource=False, use_callable_params=False): + for i, dtype in enumerate([dtypes.half, dtypes.float32, dtypes.float64]): + with self.session(graph=ops.Graph()): + # Initialize variables for numpy implementation. + m0, v0, m1, v1 = 0.0, 0.0, 0.0, 0.0 + var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype) + grads0_np = np.array([0.1, 0.1], dtype=dtype.as_numpy_dtype) + var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype) + grads1_np = np.array([0.01, 0.01], dtype=dtype.as_numpy_dtype) + + if use_resource: + global_step = resource_variable_ops.ResourceVariable( + array_ops.zeros([], dtypes.int64), name="global_step_%d" % i) + var0 = resource_variable_ops.ResourceVariable( + var0_np, name="var0_%d" % i) + var1 = resource_variable_ops.ResourceVariable( + var1_np, name="var1_%d" % i) + else: + global_step = variables.Variable(array_ops.zeros([], dtypes.int64)) + var0 = variables.Variable(var0_np) + var1 = variables.Variable(var1_np) + grads0 = constant_op.constant(grads0_np) + grads1 = constant_op.constant(grads1_np) + + learning_rate = lambda: 0.001 + beta1 = lambda: 0.9 + beta2 = lambda: 0.999 + epsilon = lambda: 1e-8 + if not use_callable_params: + learning_rate = learning_rate() + beta1 = beta1() + beta2 = beta2() + epsilon = epsilon() + + opt = lazy_adam_gs_optimizer.LazyAdamGSOptimizer( + global_step=global_step, learning_rate=learning_rate) + update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]), + global_step=global_step) + opt_variables = opt.variables() + beta1_power, beta2_power = opt._get_beta_accumulators() + self.assertIsNotNone(beta1_power) + self.assertIsNotNone(beta2_power is not None) + self.assertNotIn(beta1_power, opt_variables) + self.assertNotIn(beta2_power, opt_variables) + + if not context.executing_eagerly(): + with ops.Graph().as_default(): + # Shouldn't return non-slot variables from other graphs. + self.assertEqual(0, len(opt.variables())) + self.evaluate(variables.global_variables_initializer()) + # Fetch params to validate initial values + self.assertAllClose([1.0, 2.0], self.evaluate(var0)) + self.assertAllClose([3.0, 4.0], self.evaluate(var1)) + + # Run 3 steps of Adam + for t in range(1, 4): + if not context.executing_eagerly(): + self.evaluate(update) + self.assertAllCloseAccordingToType( + 0.9**(t + 1), self.evaluate(beta1_power)) + self.assertAllCloseAccordingToType( + 0.999**(t + 1), self.evaluate(beta2_power)) + else: + if t > 1: + opt.apply_gradients(zip([grads0, grads1], [var0, var1]), + global_step=global_step) + beta1_power, beta2_power = opt._get_beta_accumulators() + self.assertAllCloseAccordingToType( + 0.9**t, self.evaluate(beta1_power)) + self.assertAllCloseAccordingToType( + 0.999**t, self.evaluate(beta2_power)) + + var0_np, m0, v0 = adam_update_numpy(var0_np, grads0_np, t, m0, v0) + var1_np, m1, v1 = adam_update_numpy(var1_np, grads1_np, t, m1, v1) + + # Validate updated params + self.assertAllCloseAccordingToType(var0_np, self.evaluate(var0)) + self.assertAllCloseAccordingToType(var1_np, self.evaluate(var1)) + if use_resource: + self.assertEqual("var0_%d/Adam:0" % (i,), + opt.get_slot(var=var0, name="m").name) + + def testBasic(self): + with self.cached_session(): + self.doTestBasic(use_resource=False) + + @test_util.run_in_graph_and_eager_modes(reset_test=True) + def testResourceBasic(self): + self.doTestBasic(use_resource=True) + + def testBasicCallableParams(self): + with context.eager_mode(): + self.doTestBasic(use_resource=True, use_callable_params=True) + + def testTensorLearningRate(self): + for dtype in [dtypes.half, dtypes.float32, dtypes.float64]: + with self.cached_session(): + global_step = variables.Variable(array_ops.zeros([], dtypes.int64)) + # Initialize variables for numpy implementation. + m0, v0, m1, v1 = 0.0, 0.0, 0.0, 0.0 + var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype) + grads0_np = np.array([0.1, 0.1], dtype=dtype.as_numpy_dtype) + var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype) + grads1_np = np.array([0.01, 0.01], dtype=dtype.as_numpy_dtype) + + var0 = variables.Variable(var0_np) + var1 = variables.Variable(var1_np) + grads0 = constant_op.constant(grads0_np) + grads1 = constant_op.constant(grads1_np) + opt = lazy_adam_gs_optimizer.LazyAdamGSOptimizer( + global_step=global_step, learning_rate=constant_op.constant(0.001)) + update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]), + global_step=global_step) + variables.global_variables_initializer().run() + + # Fetch params to validate initial values + self.assertAllClose([1.0, 2.0], var0.eval()) + self.assertAllClose([3.0, 4.0], var1.eval()) + + beta1_power, beta2_power = opt._get_beta_accumulators() + + # Run 3 steps of Adam + for t in range(1, 4): + self.assertAllCloseAccordingToType(0.9**t, beta1_power.eval()) + self.assertAllCloseAccordingToType(0.999**t, beta2_power.eval()) + update.run() + + var0_np, m0, v0 = adam_update_numpy(var0_np, grads0_np, t, m0, v0) + var1_np, m1, v1 = adam_update_numpy(var1_np, grads1_np, t, m1, v1) + + # Validate updated params + self.assertAllCloseAccordingToType(var0_np, var0.eval()) + self.assertAllCloseAccordingToType(var1_np, var1.eval()) + + def testSharing(self): + for dtype in [dtypes.half, dtypes.float32, dtypes.float64]: + with self.cached_session(): + global_step = variables.Variable(array_ops.zeros([], dtypes.int64)) + # Initialize variables for numpy implementation. + m0, v0, m1, v1 = 0.0, 0.0, 0.0, 0.0 + var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype) + grads0_np = np.array([0.1, 0.1], dtype=dtype.as_numpy_dtype) + var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype) + grads1_np = np.array([0.01, 0.01], dtype=dtype.as_numpy_dtype) + + var0 = variables.Variable(var0_np) + var1 = variables.Variable(var1_np) + grads0 = constant_op.constant(grads0_np) + grads1 = constant_op.constant(grads1_np) + opt = lazy_adam_gs_optimizer.LazyAdamGSOptimizer( + global_step=global_step) + update1 = opt.apply_gradients(zip([grads0, grads1], [var0, var1]), + global_step=global_step) + update2 = opt.apply_gradients(zip([grads0, grads1], [var0, var1]), + global_step=global_step) + variables.global_variables_initializer().run() + + beta1_power, beta2_power = opt._get_beta_accumulators() + + # Fetch params to validate initial values + self.assertAllClose([1.0, 2.0], var0.eval()) + self.assertAllClose([3.0, 4.0], var1.eval()) + + # Run 3 steps of intertwined Adam1 and Adam2. + for t in range(1, 4): + self.assertAllCloseAccordingToType(0.9**t, beta1_power.eval()) + self.assertAllCloseAccordingToType(0.999**t, beta2_power.eval()) + if t % 2 == 0: + update1.run() + else: + update2.run() + + var0_np, m0, v0 = adam_update_numpy(var0_np, grads0_np, t, m0, v0) + var1_np, m1, v1 = adam_update_numpy(var1_np, grads1_np, t, m1, v1) + + # Validate updated params + self.assertAllCloseAccordingToType(var0_np, var0.eval()) + self.assertAllCloseAccordingToType(var1_np, var1.eval()) + + def testTwoSessions(self): + optimizer = lazy_adam_gs_optimizer.LazyAdamGSOptimizer() + + with context.eager_mode(): + var0 = variables.Variable(np.array([1.0, 2.0]), name="v0") + grads0 = constant_op.constant(np.array([0.1, 0.1])) + optimizer.apply_gradients([(grads0, var0)]) + + g = ops.Graph() + with g.as_default(): + with self.session(graph=g): + var0 = variables.Variable(np.array([1.0, 2.0]), name="v0") + grads0 = constant_op.constant(np.array([0.1, 0.1])) + optimizer.apply_gradients([(grads0, var0)]) + + gg = ops.Graph() + with gg.as_default(): + with self.session(graph=gg): + var0 = variables.Variable(np.array([1.0, 2.0]), name="v0") + grads0 = constant_op.constant(np.array([0.1, 0.1])) + + # If the optimizer saves any state not keyed by graph the following line + # fails. + optimizer.apply_gradients([(grads0, var0)]) + + def testSlotsUniqueEager(self): + with context.eager_mode(): + v1 = resource_variable_ops.ResourceVariable(1.) + v2 = resource_variable_ops.ResourceVariable(1.) + opt = lazy_adam_gs_optimizer.LazyAdamGSOptimizer(1.) + opt.minimize(lambda: v1 + v2) + # There should be two non-slot variables, and two unique slot variables + # for v1 and v2 respectively. + self.assertLen(set(opt.variables()), 4) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/opt/python/training/moving_average_optimizer.py b/tensorflow/contrib/opt/python/training/moving_average_optimizer.py index 9ce50bfe1054072b315adecb87f1ba729dfe0d83..bf3e5c51f78cc3ca3c7c77009c9cf428c4988953 100644 --- a/tensorflow/contrib/opt/python/training/moving_average_optimizer.py +++ b/tensorflow/contrib/opt/python/training/moving_average_optimizer.py @@ -26,6 +26,7 @@ from tensorflow.python.ops import variables from tensorflow.python.training import moving_averages from tensorflow.python.training import optimizer from tensorflow.python.training import saver +from tensorflow.python.training.saving import saveable_object_util class MovingAverageOptimizer(optimizer.Optimizer): @@ -106,6 +107,32 @@ class MovingAverageOptimizer(optimizer.Optimizer): self._swapped_variable_name_map[v_avg.op.name] = v.op.name return control_flow_ops.group(train_op, ma_op, name='train_with_avg') + def _find_swapped_variable(self, v_name_to_tensor, v_name, tensor): + """Returns name of swapped variable for given tensor. + + Args: + v_name_to_tensor: Mapping from variable names to tensors. + v_name: name of the variable for which swapped variable should be returned + tensor: Tensor which correspond to variable for which swapped variable + should be returned. + + Returns: + Tensor which correspond to swapped variable. + + Raises: + ValueError: If swapped variable could not be found in v_name_to_tensor. + """ + swapped_v_name = self._swapped_variable_name_map.get(v_name, None) + if swapped_v_name is None: + return tensor + else: + if swapped_v_name in v_name_to_tensor: + return v_name_to_tensor[swapped_v_name] + else: + raise ValueError( + ('Variable to swap %s is not part of variables to save. ' + 'This breaks MovingAverageOptimizer.') % swapped_v_name) + def swapping_saver(self, var_list=None, name='swapping_saver', **kwargs): """Create a saver swapping moving averages and variables. @@ -139,35 +166,35 @@ class MovingAverageOptimizer(optimizer.Optimizer): if var_list is None: var_list = variables.global_variables() if not isinstance(var_list, dict): - var_list = saver.BaseSaverBuilder.OpListToDict(var_list) - - # OpListToDict converts variables to tensors. We make sure we can get - # the unique variable name for normal and resource vaiables. - def get_v_name(tensor): - if tensor.op.type == 'ReadVariableOp': - return tensor.op.inputs[0].op.name - else: - return tensor.op.name + var_list = saveable_object_util.op_list_to_dict(var_list) v_name_to_tensor = {} - for tensor in six.itervalues(var_list): - v_name = get_v_name(tensor) - v_name_to_tensor[v_name] = tensor + for k, tensor_or_list in six.iteritems(var_list): + # For each partitioned variable OpListToDict returns list of constituent + # parts instead of single tensor. + if (isinstance(tensor_or_list, list) + or isinstance(tensor_or_list, variables.PartitionedVariable)): + for tensor in tensor_or_list: + v_name = tensor.op.name + v_name_to_tensor[v_name] = tensor + else: + v_name_to_tensor[k] = tensor_or_list # Now swap variables and moving averages swapped_var_list = {} - for k, tensor in six.iteritems(var_list): - v_name = get_v_name(tensor) - swapped_v_name = self._swapped_variable_name_map.get(v_name, None) - tensor_to_save = tensor - if swapped_v_name is not None: - if swapped_v_name in v_name_to_tensor: - tensor_to_save = v_name_to_tensor[swapped_v_name] - else: - raise ValueError( - ('Variable to swap %s is not part of variables to save. ' - 'This breaks MovingAverageOptimizer.') % swapped_v_name) - swapped_var_list[k] = tensor_to_save + for k, tensor_or_list in six.iteritems(var_list): + if isinstance(tensor_or_list, list): + tensor_list_to_save = [] + for tensor in tensor_or_list: + v_name = tensor.op.name + swapped_variable = self._find_swapped_variable(v_name_to_tensor, + v_name, + tensor) + tensor_list_to_save.append(swapped_variable) + swapped_var_list[k] = tensor_list_to_save + else: + swapped_var_list[k] = self._find_swapped_variable( + v_name_to_tensor, k, tensor_or_list) # Build the swapping saver. return saver.Saver(swapped_var_list, name=name, **kwargs) diff --git a/tensorflow/contrib/opt/python/training/moving_average_optimizer_test.py b/tensorflow/contrib/opt/python/training/moving_average_optimizer_test.py index f22e7245285a8b2716645f9789eb5997928a22d2..643403eea6f88bcb33aa96d6539bc9a45a109c6b 100644 --- a/tensorflow/contrib/opt/python/training/moving_average_optimizer_test.py +++ b/tensorflow/contrib/opt/python/training/moving_average_optimizer_test.py @@ -26,6 +26,8 @@ from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.ops import control_flow_ops +from tensorflow.python.ops import math_ops +from tensorflow.python.ops import partitioned_variables from tensorflow.python.ops import state_ops from tensorflow.python.ops import variable_scope from tensorflow.python.ops import variables @@ -43,97 +45,171 @@ class MovingAverageOptimizerTest(test.TestCase): # Test that MovingAverageOptimizer works with resource variables. self._helpTestRun(use_resource=True) - def _helpTestRun(self, use_resource=False): + def testRunUsePartitionedVars(self): + # Test that MovingAverageOptimizer works with partitioned variables. + self._helpTestRun(use_partitioned_vars=True) + + def testRunUseResourcePartitionedVars(self): + # Test that MovingAverageOptimizer works with resource and partitioned + # variables. + self._helpTestRun(use_partitioned_vars=True, use_resource=True) + + def _helpTestRun(self, use_resource=False, use_partitioned_vars=False): + # Partitioned variables are represented as a "collection" of partitions. + # To simplify the test and reuse as much code as possible we employ + # following test strategy for partitioned variables. + # + # In the case of non-partitioned variables test runs on variables with + # shape [2]. + # + # In the case of partitioned variables we use shape [4] with two partitions, + # thus each partition has shape [2]. + # For partitioned variables the test is run twice (for loop over + # variable_part_names), first time on the first partition of each variable, + # second time on the second partition of each variable. + variable_part_names = ['part_0', 'part_1'] if use_partitioned_vars else [''] for sequential_update in [True, False]: for dtype in [dtypes.half, dtypes.float32, dtypes.float64]: - with self.session(graph=ops.Graph()) as sess: - orig_val0 = [1.0, 2.0] - orig_val1 = [3.0, 4.0] - var0 = variable_scope.get_variable( - 'var0', - initializer=constant_op.constant(orig_val0, dtype=dtype), - use_resource=use_resource) - var1 = variable_scope.get_variable( - 'var1', - initializer=constant_op.constant(orig_val1, dtype=dtype), - use_resource=use_resource) - grads0 = constant_op.constant([0.1, 0.1], dtype=dtype) - grads1 = constant_op.constant([0.01, 0.01], dtype=dtype) - - opt = moving_average_optimizer.MovingAverageOptimizer( - gradient_descent.GradientDescentOptimizer(learning_rate=2.0), - average_decay=0.5, - sequential_update=sequential_update) - save_dir = tempfile.mkdtemp( - prefix=os.path.join(self.get_temp_dir(), 'run_1')) - save_path = os.path.join(save_dir, 'model') - update = opt.apply_gradients( - list(six.moves.zip([grads0, grads1], [var0, var1]))) - global_vars = variables.global_variables() - ema_var0 = [ - v for v in global_vars - if v.op.name == 'var0/ExponentialMovingAverage' - ][0] - ema_var1 = [ - v for v in global_vars - if v.op.name == 'var1/ExponentialMovingAverage' - ][0] - perturb = control_flow_ops.group([ - state_ops.assign_add(var0, [1.0, 1.0]), - state_ops.assign_add(var1, [2.0, 2.0]), - state_ops.assign_add(ema_var0, [3.0, 3.0]), - state_ops.assign_add(ema_var1, [4.0, 4.0]) - ]) - - # Test that saver with missing ema variables will fail. - with self.assertRaisesRegexp(ValueError, r'Variable to swap'): - opt.swapping_saver(var_list=[var0]) - - train_saver = opt.swapping_saver() - train_saver_subset = opt.swapping_saver(var_list=[var0, ema_var0]) - inference_saver = saver.Saver() - variables.global_variables_initializer().run() - # Step 1. - update.run() - self.assertAllCloseAccordingToType([0.8, 1.8], var0.eval()) - self.assertAllCloseAccordingToType([2.98, 3.98], var1.eval()) - if sequential_update: - self.assertAllCloseAccordingToType([0.9, 1.9], ema_var0.eval()) - self.assertAllCloseAccordingToType([2.99, 3.99], ema_var1.eval()) - # Test that the swapping saver save/restore operation is identity. - train_saver.save(sess, save_path) - train_saver.restore(sess, save_path) - self.assertAllCloseAccordingToType([0.8, 1.8], var0.eval()) - self.assertAllCloseAccordingToType([2.98, 3.98], var1.eval()) - if sequential_update: - self.assertAllCloseAccordingToType([0.9, 1.9], ema_var0.eval()) - self.assertAllCloseAccordingToType([2.99, 3.99], ema_var1.eval()) - # Test that the subset saver saves the EMA variable as well. - if sequential_update: - subset_save_path = save_path + '_subset' - train_saver_subset.save(sess, subset_save_path) - perturb.run() - self.assertAllCloseAccordingToType([1.8, 2.8], var0.eval()) - self.assertAllCloseAccordingToType([3.9, 4.9], ema_var0.eval()) - self.assertAllCloseAccordingToType([4.98, 5.98], var1.eval()) - self.assertAllCloseAccordingToType([6.99, 7.99], ema_var1.eval()) - # Restoring should only restore var0 and ema_var0. - train_saver_subset.restore(sess, subset_save_path) + for var_part_name in variable_part_names: + with self.session(graph=ops.Graph()) as sess: + orig_val0 = [1.0, 2.0] + orig_val1 = [3.0, 4.0] + grads0 = [0.1, 0.1] + grads1 = [0.01, 0.01] + if use_partitioned_vars: + # Use partitioned variables. + # Create partitioned and duplicate each value used as initial + # value of variables. + partitioner = partitioned_variables.fixed_size_partitioner( + num_shards=2) + orig_val0 = orig_val0 * 2 + orig_val1 = orig_val1 * 2 + grads0 = grads0 * 2 + grads1 = grads1 * 2 + else: + # Regular (non-partitioned) variables. + partitioner = None + var0 = variable_scope.get_variable( + 'var0', + initializer=constant_op.constant(orig_val0, dtype=dtype), + use_resource=use_resource, + partitioner=partitioner) + var1 = variable_scope.get_variable( + 'var1', + initializer=constant_op.constant(orig_val1, dtype=dtype), + use_resource=use_resource, + partitioner=partitioner) + # Make a fake loss, such that gradient(loss, var0) == grads0 + # and gradient(loss, var1) == grads1 + grads0 = constant_op.constant(grads0, dtype=dtype) + grads1 = constant_op.constant(grads1, dtype=dtype) + loss = (math_ops.reduce_sum(grads0 * var0) + + math_ops.reduce_sum(grads1 * var1)) + + opt = moving_average_optimizer.MovingAverageOptimizer( + gradient_descent.GradientDescentOptimizer(learning_rate=2.0), + average_decay=0.5, + sequential_update=sequential_update) + save_dir = tempfile.mkdtemp( + prefix=os.path.join(self.get_temp_dir(), 'run_1')) + save_path = os.path.join(save_dir, 'model') + + update = opt.minimize(loss) + + # Get variables and their EMAs. In case of partitioned variables + # get proper part of each variable. + def _get_variable(var_name, part_name, ema): + """Returns variable of it's moving average by name.""" + matches = [ + v for v in variables.global_variables() + if ((var_name in v.op.name) + and (part_name in v.op.name) + and (('ExponentialMovingAverage' in v.op.name) == ema)) + ] + self.assertEqual(len(matches), 1) + return matches[0] + var0 = _get_variable('var0', var_part_name, ema=False) + var1 = _get_variable('var1', var_part_name, ema=False) + ema_var0 = _get_variable('var0', var_part_name, ema=True) + ema_var1 = _get_variable('var1', var_part_name, ema=True) + + perturb = control_flow_ops.group([ + state_ops.assign_add(var0, [1.0, 1.0]), + state_ops.assign_add(var1, [2.0, 2.0]), + state_ops.assign_add(ema_var0, [3.0, 3.0]), + state_ops.assign_add(ema_var1, [4.0, 4.0]) + ]) + + # Test that saver with missing ema variables will fail. + with self.assertRaisesRegexp(ValueError, r'Variable to swap'): + opt.swapping_saver(var_list=[var0]) + + train_saver = opt.swapping_saver() + train_saver_subset = opt.swapping_saver(var_list=[var0, ema_var0]) + inference_saver = saver.Saver() + variables.global_variables_initializer().run() + # Step 1. + update.run() self.assertAllCloseAccordingToType([0.8, 1.8], var0.eval()) - self.assertAllCloseAccordingToType([0.9, 1.9], ema_var0.eval()) - self.assertAllCloseAccordingToType([4.98, 5.98], var1.eval()) - self.assertAllCloseAccordingToType([6.99, 7.99], ema_var1.eval()) - # Restore back to previous state. + self.assertAllCloseAccordingToType([2.98, 3.98], var1.eval()) + if sequential_update: + self.assertAllCloseAccordingToType([0.9, 1.9], ema_var0.eval()) + self.assertAllCloseAccordingToType([2.99, 3.99], ema_var1.eval()) + # Test that the swapping saver save/restore operation is identity. + train_saver.save(sess, save_path) train_saver.restore(sess, save_path) + self.assertAllCloseAccordingToType([0.8, 1.8], var0.eval()) + self.assertAllCloseAccordingToType([2.98, 3.98], var1.eval()) + if sequential_update: + self.assertAllCloseAccordingToType([0.9, 1.9], ema_var0.eval()) + self.assertAllCloseAccordingToType([2.99, 3.99], ema_var1.eval()) + # Test that the subset saver saves the EMA variable as well. + if sequential_update: + subset_save_path = save_path + '_subset' + train_saver_subset.save(sess, subset_save_path) + perturb.run() + self.assertAllCloseAccordingToType([1.8, 2.8], var0.eval()) + self.assertAllCloseAccordingToType([3.9, 4.9], ema_var0.eval()) + self.assertAllCloseAccordingToType([4.98, 5.98], var1.eval()) + self.assertAllCloseAccordingToType([6.99, 7.99], ema_var1.eval()) + # Restoring should only restore var0 and ema_var0. + train_saver_subset.restore(sess, subset_save_path) + self.assertAllCloseAccordingToType([0.8, 1.8], var0.eval()) + self.assertAllCloseAccordingToType([0.9, 1.9], ema_var0.eval()) + self.assertAllCloseAccordingToType([4.98, 5.98], var1.eval()) + self.assertAllCloseAccordingToType([6.99, 7.99], ema_var1.eval()) + # Restore back to previous state. + train_saver.restore(sess, save_path) - # If updates are parallel, this is not always true after the 1st step. - if sequential_update: + # If updates are parallel, + # this is not always true after the 1st step. + if sequential_update: + # Test that the normal saver will have the averaged variables. + # We test that the average values are between the original value + # and the most recent variable values (since they are an average + # of the two). + val0 = var0.eval() + val1 = var1.eval() + train_saver.save(sess, save_path) + inference_saver.restore(sess, save_path) + avg_val0 = var0.eval() + avg_val1 = var1.eval() + for i in six.moves.range(len(val0)): + self.assertLess(val0[i], avg_val0[i]) + self.assertLess(avg_val0[i], orig_val0[i]) + self.assertLess(val1[i], avg_val1[i]) + self.assertLess(avg_val1[i], orig_val1[i]) + train_saver.restore(sess, save_path) + # Step 2. + update.run() # Test that the normal saver will have the averaged variables. - # We test that the average values are between the original value - # and the most recent variable values (since they are an average - # of the two). + # We test that the average values are between the original value and + # the most recent variable values (since they are an average of the + # two). val0 = var0.eval() val1 = var1.eval() + self.assertAllCloseAccordingToType([0.6, 1.6], val0) + self.assertAllCloseAccordingToType([2.96, 3.96], val1) train_saver.save(sess, save_path) inference_saver.restore(sess, save_path) avg_val0 = var0.eval() @@ -143,26 +219,6 @@ class MovingAverageOptimizerTest(test.TestCase): self.assertLess(avg_val0[i], orig_val0[i]) self.assertLess(val1[i], avg_val1[i]) self.assertLess(avg_val1[i], orig_val1[i]) - train_saver.restore(sess, save_path) - # Step 2. - update.run() - # Test that the normal saver will have the averaged variables. - # We test that the average values are between the original value and - # the most recent variable values (since they are an average of the - # two). - val0 = var0.eval() - val1 = var1.eval() - self.assertAllCloseAccordingToType([0.6, 1.6], val0) - self.assertAllCloseAccordingToType([2.96, 3.96], val1) - train_saver.save(sess, save_path) - inference_saver.restore(sess, save_path) - avg_val0 = var0.eval() - avg_val1 = var1.eval() - for i in six.moves.range(len(val0)): - self.assertLess(val0[i], avg_val0[i]) - self.assertLess(avg_val0[i], orig_val0[i]) - self.assertLess(val1[i], avg_val1[i]) - self.assertLess(avg_val1[i], orig_val1[i]) def testFailWhenSaverCreatedBeforeInitialized(self): with self.cached_session(): diff --git a/tensorflow/contrib/opt/python/training/nadam_optimizer.py b/tensorflow/contrib/opt/python/training/nadam_optimizer.py index 44a8890cb107440b79cf8fbbdfcfda503b1c910f..960826407b66b4efa3c2693efb6d2e17c4b47b33 100644 --- a/tensorflow/contrib/opt/python/training/nadam_optimizer.py +++ b/tensorflow/contrib/opt/python/training/nadam_optimizer.py @@ -1,4 +1,4 @@ -# Copyright 2015 The TensorFlow Authors. All Rights Reserved. +# 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. @@ -18,6 +18,7 @@ from __future__ import division from __future__ import print_function 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 math_ops from tensorflow.python.ops import state_ops @@ -83,14 +84,14 @@ class NadamOptimizer(adam.AdamOptimizer): with ops.control_dependencies([m_t]): m_t = scatter_add(m, indices, m_scaled_g_values) # m_bar = (1 - beta1) * g_t + beta1 * m_t - m_bar = m_scaled_g_values + beta1_t * m_t + m_bar = m_scaled_g_values + beta1_t * array_ops.gather(m_t, indices) # v_t = beta2 * v + (1 - beta2) * (g_t * g_t) v = self.get_slot(var, "v") v_scaled_g_values = (grad * grad) * (1 - beta2_t) v_t = state_ops.assign(v, v * beta2_t, use_locking=self._use_locking) 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_bar / (v_sqrt + epsilon_t), use_locking=self._use_locking) + v_t_slice = array_ops.gather(v_t, indices) + v_sqrt = math_ops.sqrt(v_t_slice) + var_update = scatter_add(var, indices, -lr * m_bar / (v_sqrt + epsilon_t)) return control_flow_ops.group(*[var_update, m_bar, v_t]) diff --git a/tensorflow/contrib/opt/python/training/nadam_optimizer_test.py b/tensorflow/contrib/opt/python/training/nadam_optimizer_test.py index 85e05ce71cec6ef897cadb7d123e630febb3c064..a4372f64874e7591dbceac901fad6c941209bef9 100644 --- a/tensorflow/contrib/opt/python/training/nadam_optimizer_test.py +++ b/tensorflow/contrib/opt/python/training/nadam_optimizer_test.py @@ -52,14 +52,19 @@ def nadam_update_numpy(param, class NadamOptimizerTest(test.TestCase): def doTestSparse(self, use_resource=False): + # need to use a larger value of epsilon here so that + # np.sqrt(v_t) + epsilon doesn't get rounded to 0 when + # the dtype is half and np.sqrt(v_t) = 0, as is the case + # when the gradient is 0 + sparse_epsilon = 1e-7 for dtype in [dtypes.half, dtypes.float32, dtypes.float64]: with self.cached_session(): # Initialize variables for numpy implementation. m0, v0, m1, v1 = 0.0, 0.0, 0.0, 0.0 - var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype) - grads0_np = np.array([0.1, 0.1], dtype=dtype.as_numpy_dtype) - var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype) - grads1_np = np.array([0.01, 0.01], dtype=dtype.as_numpy_dtype) + var0_np = np.array([1.0, 1.0, 2.0], dtype=dtype.as_numpy_dtype) + grads0_np = np.array([0.1, 0, 0.1], dtype=dtype.as_numpy_dtype) + var1_np = np.array([3.0, 3.0, 4.0], dtype=dtype.as_numpy_dtype) + grads1_np = np.array([0.01, 0, 0.01], dtype=dtype.as_numpy_dtype) if use_resource: var0 = resource_variable_ops.ResourceVariable(var0_np) @@ -67,21 +72,21 @@ class NadamOptimizerTest(test.TestCase): else: var0 = variables.Variable(var0_np) var1 = variables.Variable(var1_np) - grads0_np_indices = np.array([0, 1], dtype=np.int32) + grads0_np_indices = np.array([0, 2], dtype=np.int32) grads0 = ops.IndexedSlices( - constant_op.constant(grads0_np), - constant_op.constant(grads0_np_indices), constant_op.constant([2])) - grads1_np_indices = np.array([0, 1], dtype=np.int32) + constant_op.constant(grads0_np[grads0_np_indices]), + constant_op.constant(grads0_np_indices), constant_op.constant([3])) + grads1_np_indices = np.array([0, 2], dtype=np.int32) grads1 = ops.IndexedSlices( - constant_op.constant(grads1_np), - constant_op.constant(grads1_np_indices), constant_op.constant([2])) - opt = nadam_optimizer.NadamOptimizer() + constant_op.constant(grads1_np[grads1_np_indices]), + constant_op.constant(grads1_np_indices), constant_op.constant([3])) + opt = nadam_optimizer.NadamOptimizer(epsilon=sparse_epsilon) update = opt.apply_gradients(zip([grads0, grads1], [var0, var1])) variables.global_variables_initializer().run() # Fetch params to validate initial values - self.assertAllClose([1.0, 2.0], var0.eval()) - self.assertAllClose([3.0, 4.0], var1.eval()) + self.assertAllClose([1.0, 1.0, 2.0], var0.eval()) + self.assertAllClose([3.0, 3.0, 4.0], var1.eval()) beta1_power, beta2_power = opt._get_beta_accumulators() @@ -91,8 +96,10 @@ class NadamOptimizerTest(test.TestCase): self.assertAllCloseAccordingToType(0.999**t, beta2_power.eval()) update.run() - var0_np, m0, v0 = nadam_update_numpy(var0_np, grads0_np, t, m0, v0) - var1_np, m1, v1 = nadam_update_numpy(var1_np, grads1_np, t, m1, v1) + var0_np, m0, v0 = nadam_update_numpy(var0_np, grads0_np, t, m0, v0, + epsilon=sparse_epsilon) + var1_np, m1, v1 = nadam_update_numpy(var1_np, grads1_np, t, m1, v1, + epsilon=sparse_epsilon) # Validate updated params self.assertAllCloseAccordingToType(var0_np, var0.eval()) diff --git a/tensorflow/contrib/opt/python/training/weight_decay_optimizers.py b/tensorflow/contrib/opt/python/training/weight_decay_optimizers.py index 200b0d200826a6212a236680327f4daf7d07831f..8b8065c678e11e8fc237e71cf1d392ced5c22ada 100644 --- a/tensorflow/contrib/opt/python/training/weight_decay_optimizers.py +++ b/tensorflow/contrib/opt/python/training/weight_decay_optimizers.py @@ -59,6 +59,23 @@ class DecoupledWeightDecayExtension(object): Note that this extension decays weights BEFORE applying the update based on the gradient, i.e. this extension only has the desired behaviour for optimizers which do not depend on the value of'var' in the update step! + + Note: when applying a decay to the learning rate, be sure to manually apply + the decay to the `weight_decay` as well. For example: + + ```python + schedule = tf.train.piecewise_constant(tf.train.get_global_step(), + [10000, 15000], [1e-0, 1e-1, 1e-2]) + lr = 1e-1 * schedule() + wd = lambda: 1e-4 * schedule() + + # ... + + optimizer = tf.contrib.opt.MomentumWOptimizer(learning_rate=lr, + weight_decay=wd, + momentum=0.9, + use_nesterov=True) + ``` """ def __init__(self, weight_decay, **kwargs): diff --git a/tensorflow/contrib/optimizer_v2/BUILD b/tensorflow/contrib/optimizer_v2/BUILD index 3ba3ee29ec79687df522eb330665a2ce80061682..6e401406308604970677003aeea0f15c64cc74b6 100644 --- a/tensorflow/contrib/optimizer_v2/BUILD +++ b/tensorflow/contrib/optimizer_v2/BUILD @@ -48,7 +48,6 @@ py_library( srcs_version = "PY2AND3", deps = [ "//tensorflow/python:control_flow_ops", - "//tensorflow/python:distribute", "//tensorflow/python:framework", "//tensorflow/python:math_ops", "//tensorflow/python:resource_variable_ops", @@ -56,6 +55,8 @@ py_library( "//tensorflow/python:training", "//tensorflow/python:variable_scope", "//tensorflow/python:variables", + "//tensorflow/python/distribute:distribute_lib", + "//tensorflow/python/distribute:reduce_util", ], ) 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 ae8a5d243b871cd39f4697731d166d0fa03c1e5b..0243927ce44aec626973744507e75b20a42253e9 100644 --- a/tensorflow/contrib/optimizer_v2/checkpointable_utils_test.py +++ b/tensorflow/contrib/optimizer_v2/checkpointable_utils_test.py @@ -24,6 +24,7 @@ import os import six +from tensorflow.contrib.optimizer_v2 import adam from tensorflow.python.client import session as session_lib from tensorflow.python.eager import backprop from tensorflow.python.eager import context @@ -34,7 +35,6 @@ from tensorflow.python.framework import ops from tensorflow.python.framework import test_util from tensorflow.python.keras.engine import training from tensorflow.python.keras.layers import core -from tensorflow.python.keras.optimizer_v2 import adam from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import init_ops from tensorflow.python.ops import resource_variable_ops @@ -48,7 +48,7 @@ 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__() @@ -98,7 +98,7 @@ class CheckpointingTests(test.TestCase): # A nuisance Model using the same optimizer. Its slot variables should not # go in the checkpoint, since it is never depended on. other_model = MyModel() - optimizer = adam.Adam(0.001) + optimizer = adam.AdamOptimizer(0.001) optimizer_step = training_util.get_or_create_global_step() root_checkpointable = util.Checkpoint( optimizer=optimizer, model=model, optimizer_step=optimizer_step) @@ -130,8 +130,8 @@ class CheckpointingTests(test.TestCase): # non-Layer dependency of the model "model/_non_layer/a_variable", # The optimizer creates two non-slot variables - "optimizer/beta_1_power", - "optimizer/beta_2_power", + "optimizer/beta1_power", + "optimizer/beta2_power", # Slot variables "model/_second/kernel/.OPTIMIZER_SLOT/optimizer/m", "model/_second/kernel/.OPTIMIZER_SLOT/optimizer/v", @@ -145,7 +145,6 @@ class CheckpointingTests(test.TestCase): name + suffix for name in expected_checkpoint_names] # The optimizer and Dense layers also save get_config() JSON expected_checkpoint_names.extend([ - "optimizer/.ATTRIBUTES/OBJECT_CONFIG_JSON", "model/_second/.ATTRIBUTES/OBJECT_CONFIG_JSON", "model/_named_dense/.ATTRIBUTES/OBJECT_CONFIG_JSON" ]) @@ -163,20 +162,20 @@ class CheckpointingTests(test.TestCase): "my_model/dense/kernel", named_variables["model/_named_dense/kernel" + suffix].full_name) self.assertEqual( - "beta_1_power", - named_variables["optimizer/beta_1_power" + suffix].full_name) + "beta1_power", + named_variables["optimizer/beta1_power" + suffix].full_name) self.assertEqual( - "beta_2_power", - named_variables["optimizer/beta_2_power" + suffix].full_name) + "beta2_power", + named_variables["optimizer/beta2_power" + suffix].full_name) # Spot check the generated protocol buffers. self.assertEqual("optimizer", serialized_graph.nodes[0].children[1].local_name) optimizer_node = serialized_graph.nodes[serialized_graph.nodes[0].children[ 1].node_id] - self.assertEqual("beta_1_power", optimizer_node.children[0].local_name) + self.assertEqual("beta1_power", optimizer_node.children[0].local_name) self.assertEqual( - "beta_1_power", serialized_graph.nodes[ - optimizer_node.children[0].node_id].attributes[0].full_name) + "beta1_power", serialized_graph.nodes[optimizer_node.children[0] + .node_id].attributes[0].full_name) self.assertEqual( "my_model/dense/kernel", serialized_graph.nodes[optimizer_node.slot_variables[0] @@ -208,7 +207,7 @@ class CheckpointingTests(test.TestCase): @test_util.run_in_graph_and_eager_modes def testSaveRestore(self): model = MyModel() - optimizer = adam.Adam(0.001) + optimizer = adam.AdamOptimizer(0.001) root_checkpointable = util.Checkpoint( optimizer=optimizer, model=model) input_value = constant_op.constant([[3.]]) @@ -240,12 +239,12 @@ class CheckpointingTests(test.TestCase): if not context.executing_eagerly(): return # Restore-on-create is only supported when executing eagerly on_create_model = MyModel() - on_create_optimizer = adam.Adam( + on_create_optimizer = adam.AdamOptimizer( 0.001, # Preserve beta_1_power and beta_2_power when appying gradients # so we can test that they've been restored correctly. - beta_1=1.0, - beta_2=1.0) + beta1=1.0, + beta2=1.0) on_create_root = util.Checkpoint( optimizer=on_create_optimizer, model=on_create_model) # Deferred restoration @@ -277,7 +276,7 @@ class CheckpointingTests(test.TestCase): checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") for training_continuation in range(3): model = MyModel() - optimizer = adam.Adam(0.001) + optimizer = adam.AdamOptimizer(0.001) root = util.Checkpoint( optimizer=optimizer, model=model, optimizer_step=training_util.get_or_create_global_step()) @@ -302,7 +301,7 @@ class CheckpointingTests(test.TestCase): for training_continuation in range(3): with ops.Graph().as_default(): model = MyModel() - optimizer = adam.Adam(0.001) + optimizer = adam.AdamOptimizer(0.001) root = util.Checkpoint( optimizer=optimizer, model=model, global_step=training_util.get_or_create_global_step()) @@ -340,7 +339,7 @@ class CheckpointingTests(test.TestCase): with ops.Graph().as_default(), self.test_session( graph=ops.get_default_graph()), test_util.device(use_gpu=True): model = MyModel() - optimizer = adam.Adam(0.001) + optimizer = adam.AdamOptimizer(0.001) root = util.Checkpoint( optimizer=optimizer, model=model, global_step=training_util.get_or_create_global_step()) @@ -374,7 +373,7 @@ class CheckpointingTests(test.TestCase): graph=ops.get_default_graph()), test_util.device(use_gpu=True): model = MyModel() # Don't actually train so we can test variable values - optimizer = adam.Adam(0.) + optimizer = adam.AdamOptimizer(0.) root = util.Checkpoint( optimizer=optimizer, model=model, global_step=training_util.get_or_create_global_step()) @@ -423,7 +422,7 @@ class CheckpointingTests(test.TestCase): with context.eager_mode(): model = Model() - optimizer = adam.Adam(learning_rate=0.05) + optimizer = adam.AdamOptimizer(learning_rate=0.05) checkpoint_directory = self.get_temp_dir() checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") checkpoint = util.Checkpoint( @@ -441,10 +440,10 @@ class CheckpointingTests(test.TestCase): def testDeferredSlotRestoration(self): checkpoint_directory = self.get_temp_dir() - root = tracking.Checkpointable() + root = tracking.AutoCheckpointable() root.var = util.add_variable( root, name="var", initializer=0.) - optimizer = adam.Adam(0.1) + optimizer = adam.AdamOptimizer(0.1) if context.executing_eagerly(): optimizer.minimize(root.var.read_value) else: @@ -464,7 +463,7 @@ class CheckpointingTests(test.TestCase): 14.)) slots_path = util.CheckpointableSaver(root).save( os.path.join(checkpoint_directory, "with_slots")) - new_root = tracking.Checkpointable() + new_root = tracking.AutoCheckpointable() # Load the slot-containing checkpoint (deferred), then immediately overwrite # the non-slot variable (also deferred). slot_status = util.CheckpointableSaver( @@ -478,8 +477,8 @@ class CheckpointingTests(test.TestCase): no_slot_status.assert_consumed() no_slot_status.run_restore_ops() self.assertEqual(12., self.evaluate(new_root.var)) - new_root.optimizer = adam.Adam(0.1) - with self.assertRaisesRegexp(AssertionError, "beta_1_power"): + new_root.optimizer = adam.AdamOptimizer(0.1) + with self.assertRaisesRegexp(AssertionError, "beta1_power"): slot_status.assert_consumed() self.assertEqual(12., self.evaluate(new_root.var)) if context.executing_eagerly(): @@ -509,9 +508,9 @@ 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 = tracking.AutoCheckpointable() obj.var = variable_scope.get_variable(name="v", initializer=0.) - obj.opt = adam.Adam(0.1) + obj.opt = adam.AdamOptimizer(0.1) obj.opt.minimize(obj.var.read_value()) self.evaluate(util.gather_initializers(obj)) saver = util.CheckpointableSaver(obj) @@ -527,9 +526,9 @@ 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 = tracking.AutoCheckpointable() obj.var = variable_scope.get_variable(name="v", initializer=0.) - obj.opt = adam.Adam(0.1) + obj.opt = adam.AdamOptimizer(0.1) obj.opt.minimize(obj.var.read_value()) self.evaluate(util.gather_initializers(obj)) saver = util.CheckpointableSaver(obj) @@ -543,7 +542,7 @@ class CheckpointingTests(test.TestCase): with context.graph_mode(): checkpoint_directory = self.get_temp_dir() checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt") - optimizer = adam.Adam(0.001) + optimizer = adam.AdamOptimizer(0.001) # Construct a model in one graph first_graph = ops.Graph() first_session = session_lib.Session(graph=first_graph) @@ -614,7 +613,7 @@ class TemplateTests(test.TestCase): save_template = template.make_template("s1", _templated) v1_save, _, v2_save = save_template() - optimizer = adam.Adam(0.0) + optimizer = adam.AdamOptimizer(0.0) save_root = util.Checkpoint( my_template=save_template, optimizer=optimizer) optimizer.minimize(v1_save.read_value) @@ -626,7 +625,7 @@ class TemplateTests(test.TestCase): save_path = save_root.save(checkpoint_prefix) load_template = template.make_template("s2", _templated) - load_optimizer = adam.Adam(0.0) + load_optimizer = adam.AdamOptimizer(0.0) load_root = util.Checkpoint( my_template=load_template, optimizer=load_optimizer) status = load_root.restore(save_path) @@ -646,7 +645,7 @@ class CheckpointCompatibilityTests(test.TestCase): def _initialized_model(self): input_value = constant_op.constant([[3.]]) model = MyModel() - optimizer = adam.Adam(0.001) + optimizer = adam.AdamOptimizer(0.001) optimizer_step = training_util.get_or_create_global_step() root_checkpointable = util.Checkpoint( optimizer=optimizer, model=model, optimizer_step=optimizer_step) diff --git a/tensorflow/contrib/optimizer_v2/optimizer_v2.py b/tensorflow/contrib/optimizer_v2/optimizer_v2.py index b00262e1f331fdf07a5c37e7286c5c93f51ec524..1323ed014c9e51e273491694fa44a8e36cc723d0 100644 --- a/tensorflow/contrib/optimizer_v2/optimizer_v2.py +++ b/tensorflow/contrib/optimizer_v2/optimizer_v2.py @@ -22,6 +22,11 @@ from __future__ import print_function 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 from tensorflow.python.eager import context from tensorflow.python.framework import dtypes @@ -32,14 +37,13 @@ from tensorflow.python.ops import math_ops from tensorflow.python.ops import resource_variable_ops from tensorflow.python.ops import variable_scope from tensorflow.python.ops import variables -from tensorflow.python.training import distribute as distribute_lib -from tensorflow.python.training import distribution_strategy_context from tensorflow.python.training import optimizer as optimizer_v1 from tensorflow.python.training import slot_creator from tensorflow.python.training.checkpointable import base as checkpointable from tensorflow.python.util import nest +@six.add_metaclass(abc.ABCMeta) class _OptimizableVariable(object): """Interface for abstracting over variables in the optimizers.""" @@ -443,7 +447,7 @@ class _OptimizerV2State(object): if v is None: if colocate_with is None: colocate_with = self._non_slot_devices - with self._distribution.colocate_vars_with(colocate_with): + with self._distribution.extended.colocate_vars_with(colocate_with): # TODO(josh11b): Use get_variable() except for the legacy Adam use case. v = variable_scope.variable(initial_value, name=name, trainable=False) self._non_slot_dict[name] = v @@ -631,16 +635,16 @@ class OptimizerV2(optimizer_v1.Optimizer): # Map from graph_key to state for that graph. We use the graph_key # since it works in both eager and graph mode, and gives the outer # graph inside functions. - tower_context = distribution_strategy_context.get_tower_context() - if tower_context is None: - # In a cross-tower context for a DistributionStrategy, which means - # only one Optimizer will be created, not one per tower. + replica_context = distribute_ctx.get_replica_context() + if replica_context is None: + # In a cross-replica context for a DistributionStrategy, which means + # only one Optimizer will be created, not one per replica. self._per_graph_state = {} else: - # We use get_tower_context().merge_call() to get a single dict + # We use get_replica_context().merge_call() to get a single dict # shared across all model replicas when running with a # DistributionStrategy. - self._per_graph_state = tower_context.merge_call(lambda _: {}) + self._per_graph_state = replica_context.merge_call(lambda _: {}) # Hyper parameters, and whether they should be re-evaluated every step. self._hyper = {} @@ -654,11 +658,10 @@ class OptimizerV2(optimizer_v1.Optimizer): var_list=None, gate_gradients=GATE_OP, aggregation_method=None, - colocate_gradients_with_ops=False, name=None, grad_loss=None, stop_gradients=None, - scale_loss_by_num_towers=None): + scale_loss_by_num_replicas=None): """Add operations to minimize `loss` by updating `var_list`. This method simply combines calls `compute_gradients()` and @@ -677,14 +680,12 @@ class OptimizerV2(optimizer_v1.Optimizer): `GATE_NONE`, `GATE_OP`, or `GATE_GRAPH`. aggregation_method: Specifies the method used to combine gradient terms. Valid values are defined in the class `AggregationMethod`. - colocate_gradients_with_ops: If True, try colocating gradients with the - corresponding op. name: Optional name for the returned operation. grad_loss: Optional. A `Tensor` holding the gradient computed for `loss`. stop_gradients: Optional. A Tensor or list of tensors not to differentiate through. - scale_loss_by_num_towers: Optional boolean. If true, scale the loss down - by the number of towers. By default, auto-detects whether this is + 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. Returns: @@ -701,8 +702,8 @@ class OptimizerV2(optimizer_v1.Optimizer): Minimization (and gradient computation) is done with respect to the elements of `var_list` if not None, else with respect to any trainable variables created during the execution of the `loss` function. - `gate_gradients`, `aggregation_method`, `colocate_gradients_with_ops` and - `grad_loss` are ignored when eager execution is enabled. + `gate_gradients`, `aggregation_method`, and `grad_loss` are ignored when + eager execution is enabled. @end_compatibility """ grads_and_vars = self.compute_gradients( @@ -710,10 +711,9 @@ class OptimizerV2(optimizer_v1.Optimizer): var_list=var_list, gate_gradients=gate_gradients, aggregation_method=aggregation_method, - colocate_gradients_with_ops=colocate_gradients_with_ops, grad_loss=grad_loss, stop_gradients=stop_gradients, - scale_loss_by_num_towers=scale_loss_by_num_towers) + scale_loss_by_num_replicas=scale_loss_by_num_replicas) vars_with_grad = [v for g, v in grads_and_vars if g is not None] if not vars_with_grad: @@ -730,10 +730,9 @@ class OptimizerV2(optimizer_v1.Optimizer): var_list=None, gate_gradients=GATE_OP, aggregation_method=None, - colocate_gradients_with_ops=False, grad_loss=None, stop_gradients=None, - scale_loss_by_num_towers=None): + scale_loss_by_num_replicas=None): """Compute gradients of `loss` for the variables in `var_list`. This is the first part of `minimize()`. It returns a list @@ -753,13 +752,11 @@ class OptimizerV2(optimizer_v1.Optimizer): `GATE_NONE`, `GATE_OP`, or `GATE_GRAPH`. aggregation_method: Specifies the method used to combine gradient terms. Valid values are defined in the class `AggregationMethod`. - colocate_gradients_with_ops: If True, try colocating gradients with the - corresponding op. grad_loss: Optional. A `Tensor` holding the gradient computed for `loss`. stop_gradients: Optional. A Tensor or list of tensors not to differentiate through. - scale_loss_by_num_towers: Optional boolean. If true, scale the loss down - by the number of towers. By default, auto-detects whether this is + 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. Returns: @@ -773,8 +770,8 @@ class OptimizerV2(optimizer_v1.Optimizer): not callable. @compatibility(eager) - When eager execution is enabled, `gate_gradients`, `aggregation_method`, - and `colocate_gradients_with_ops` are ignored. + When eager execution is enabled, `gate_gradients`, and `aggregation_method` + are ignored. @end_compatibility """ # TODO(josh11b): Test that we handle weight decay in a reasonable way. @@ -784,18 +781,10 @@ class OptimizerV2(optimizer_v1.Optimizer): tape.watch(var_list) loss_value = loss() - # Scale loss for number of towers (callable-loss case). In this case, + # 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. - if scale_loss_by_num_towers is None: - scale_loss_by_num_towers = ( - distribute_lib.get_loss_reduction() == variable_scope - .VariableAggregation.MEAN) - if scale_loss_by_num_towers: - num_towers = distribution_strategy_context.get_distribution_strategy( - ).num_towers - if num_towers > 1: - loss_value *= 1. / num_towers + loss_value = self._scale_loss(loss_value, scale_loss_by_num_replicas) if var_list is None: var_list = tape.watched_variables() @@ -805,16 +794,8 @@ class OptimizerV2(optimizer_v1.Optimizer): raise RuntimeError("`loss` passed to Optimizer.compute_gradients should " "be a function when eager execution is enabled.") - # Scale loss for number of towers (non-callable-loss case). - if scale_loss_by_num_towers is None: - scale_loss_by_num_towers = ( - distribute_lib.get_loss_reduction() == variable_scope - .VariableAggregation.MEAN) - if scale_loss_by_num_towers: - num_towers = distribution_strategy_context.get_distribution_strategy( - ).num_towers - if num_towers > 1: - loss *= 1. / num_towers + # Scale loss for number of replicas (non-callable-loss case). + loss = self._scale_loss(loss, scale_loss_by_num_replicas) if gate_gradients not in [ optimizer_v1.Optimizer.GATE_NONE, optimizer_v1.Optimizer.GATE_OP, @@ -845,7 +826,6 @@ class OptimizerV2(optimizer_v1.Optimizer): grad_ys=grad_loss, gate_gradients=(gate_gradients == optimizer_v1.Optimizer.GATE_OP), aggregation_method=aggregation_method, - colocate_gradients_with_ops=colocate_gradients_with_ops, stop_gradients=stop_gradients) if gate_gradients == optimizer_v1.Optimizer.GATE_GRAPH: grads = control_flow_ops.tuple(grads) @@ -856,6 +836,18 @@ class OptimizerV2(optimizer_v1.Optimizer): ]) return grads_and_vars + @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_strategy().num_replicas_in_sync + if num_replicas > 1: + loss_value *= 1. / num_replicas + return loss_value + def apply_gradients(self, grads_and_vars, global_step=None, name=None): """Apply gradients to variables. @@ -890,8 +882,9 @@ class OptimizerV2(optimizer_v1.Optimizer): if not filtered: raise ValueError("No gradients provided for any variable: %s." % ([str(v) for _, v in grads_and_vars],)) - return distribution_strategy_context.get_tower_context().merge_call( - self._distributed_apply, filtered, global_step=global_step, name=name) + return distribute_ctx.get_replica_context().merge_call( + self._distributed_apply, args=(filtered,), + kwargs={"global_step": global_step, "name": name}) def _get_or_create_state(self, var_list=None): """Either looks up or creates `_OptimizerV2State`. @@ -926,8 +919,8 @@ class OptimizerV2(optimizer_v1.Optimizer): def _distributed_apply(self, distribution, grads_and_vars, global_step, name): """`apply_gradients` for use with a `DistributionStrategy`.""" - reduced_grads = distribution.batch_reduce( - variable_scope.VariableAggregation.SUM, grads_and_vars) + reduced_grads = distribution.extended.batch_reduce_to( + ds_reduce_util.ReduceOp.SUM, grads_and_vars) var_list = [v for _, v in grads_and_vars] grads_and_vars = zip(reduced_grads, var_list) @@ -943,7 +936,7 @@ class OptimizerV2(optimizer_v1.Optimizer): with ops.name_scope(name, self._name) as name: per_graph_state = self._get_or_create_state(var_list=unwrapped_var_list) # Include the current value of any dynamic hyper parameters in `state`. - non_slot_devices = distribution.non_slot_devices(var_list) + non_slot_devices = distribution.extended.non_slot_devices(var_list) state = per_graph_state._copy_with_dynamic_hyper( # pylint: disable=protected-access self._hyper, distribution, non_slot_devices) @@ -988,7 +981,8 @@ class OptimizerV2(optimizer_v1.Optimizer): # Use the processors to update the variables. update_ops = [] for grad, var in grads_and_vars: - update_ops.extend(distribution.update(var, update, grad, grouped=False)) + update_ops.extend(distribution.extended.update( + var, update, args=(grad,), group=False)) # Give the child class a chance to do something after applying # gradients @@ -1000,12 +994,12 @@ class OptimizerV2(optimizer_v1.Optimizer): update_ops = control_flow_ops.group(update_ops) with ops.control_dependencies([update_ops]): - finish_updates = distribution.update_non_slot( - non_slot_devices, finish, grouped=False) - # We said grouped=False, which means finish_updates is always a list. - # It will be [None] when finish() returns None. - if finish_updates == [None]: - finish_updates = [update_ops] + finish_updates = distribution.extended.update_non_slot( + non_slot_devices, finish, group=False) + # We said group=False, which means finish_updates is always a tuple. + # It will be (None,) when finish() returns None. + if finish_updates == (None,): + finish_updates = (update_ops,) # Update `global_step` (if any). if global_step is None: @@ -1016,8 +1010,8 @@ class OptimizerV2(optimizer_v1.Optimizer): def update_global_step(global_step, name): return global_step.assign_add(1, read_value=False, name=name) - apply_updates = distribution.update(global_step, update_global_step, - name) + apply_updates = distribution.extended.update( + global_step, update_global_step, args=(name,)) # Add the training op to the TRAIN_OP graph collection in graph mode. if not eager_execution: diff --git a/tensorflow/contrib/predictor/BUILD b/tensorflow/contrib/predictor/BUILD index d50b52b8ff1ce8188ab52c6968d716378efd9daa..53a3bc63e1d770b451846c45370fdee9ffa72d70 100644 --- a/tensorflow/contrib/predictor/BUILD +++ b/tensorflow/contrib/predictor/BUILD @@ -42,6 +42,7 @@ py_library( name = "saved_model_predictor", srcs = ["saved_model_predictor.py"], srcs_version = "PY2AND3", + visibility = ["//learning/brain/contrib/learn/tpu:__subpackages__"], deps = [ ":base_predictor", "//tensorflow/contrib/saved_model:saved_model_py", diff --git a/tensorflow/contrib/quantize/BUILD b/tensorflow/contrib/quantize/BUILD index 94a2d9672dba74d19cb0801aa8680e921c238c97..b67e68ea96a15f94e62050c92405eec4fe4be70f 100644 --- a/tensorflow/contrib/quantize/BUILD +++ b/tensorflow/contrib/quantize/BUILD @@ -202,9 +202,14 @@ 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 = [ + "nomsan", + ], deps = [ ":fold_batch_norms", ":quantize", diff --git a/tensorflow/contrib/quantize/README.md b/tensorflow/contrib/quantize/README.md index 0ab19c91bb036ad24beee3d99624e788d086a9a5..5b8da92491fb747c5a37dcfe03bcb21b5b903560 100644 --- a/tensorflow/contrib/quantize/README.md +++ b/tensorflow/contrib/quantize/README.md @@ -28,7 +28,7 @@ Since it's difficult to add these fake quantization operations to all the required locations in the model, there's a function available that rewrites the training graph. To create a fake quantized training graph: -``` +```python # Build forward pass of model. loss = tf.losses.get_total_loss() @@ -51,7 +51,7 @@ The rewritten *eval graph* is non-trivially different from the *training graph* since the quantization ops affect the batch normalization step. Because of this, we've added a separate rewrite for the *eval graph*: -``` +```python # Build eval model logits = tf.nn.softmax_cross_entropy_with_logits_v2(...) @@ -110,7 +110,7 @@ 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:
@@ -145,7 +145,7 @@ Mobilenet-v2, and Inception-v3) using this tool:
Our pre-trained models are available in the -TensorFlow Lite model repository. The code used to generate +TensorFlow Lite model repository. The code used to generate these models is available. diff --git a/tensorflow/contrib/quantize/python/graph_matcher.py b/tensorflow/contrib/quantize/python/graph_matcher.py index aa3ca991c060b208ec71ae27e1ddc75df8a2c723..cfbf5bf30f9ba224afdef0c849e33fe7915cf583 100644 --- a/tensorflow/contrib/quantize/python/graph_matcher.py +++ b/tensorflow/contrib/quantize/python/graph_matcher.py @@ -21,7 +21,10 @@ from __future__ import print_function import abc import itertools +import six + +@six.add_metaclass(abc.ABCMeta) class Pattern(object): """The parent class of all patterns (e.g. OpTypePattern and OneofPattern).""" diff --git a/tensorflow/contrib/quantize/python/quant_ops.py b/tensorflow/contrib/quantize/python/quant_ops.py index 6f659347fba019288361dd0420f2ade6dc1bebaf..8619708cdaecd78bcc7de0e8e0cbf2baa11bf6a2 100644 --- a/tensorflow/contrib/quantize/python/quant_ops.py +++ b/tensorflow/contrib/quantize/python/quant_ops.py @@ -138,7 +138,7 @@ def LastValueQuantize(inputs, if per_channel: if input_dim >= 2: batch_min = math_ops.reduce_min( - inputs, reduction_indices=reduce_dims, name='BatchMin') + inputs, axis=reduce_dims, name='BatchMin') else: batch_min = inputs else: @@ -147,7 +147,7 @@ def LastValueQuantize(inputs, if per_channel: if input_dim >= 2: batch_max = math_ops.reduce_max( - inputs, reduction_indices=reduce_dims, name='BatchMax') + inputs, axis=reduce_dims, name='BatchMax') else: batch_max = inputs else: @@ -263,7 +263,7 @@ def MovingAvgQuantize(inputs, if per_channel: if input_dim >= 2: batch_min = math_ops.reduce_min( - inputs, reduction_indices=reduce_dims, name='BatchMin') + inputs, axis=reduce_dims, name='BatchMin') else: batch_min = inputs else: @@ -272,7 +272,7 @@ def MovingAvgQuantize(inputs, if per_channel: if input_dim >= 2: batch_max = math_ops.reduce_max( - inputs, reduction_indices=reduce_dims, name='BatchMax') + inputs, axis=reduce_dims, name='BatchMax') else: batch_max = inputs else: diff --git a/tensorflow/contrib/quantize/python/quantize.py b/tensorflow/contrib/quantize/python/quantize.py index 92ca3f203954414159954f7f5d220f95b17967d0..7c973fe597181b822e617db1f85a08f1b678e26f 100644 --- a/tensorflow/contrib/quantize/python/quantize.py +++ b/tensorflow/contrib/quantize/python/quantize.py @@ -91,48 +91,50 @@ def Quantize(graph, # If `scope` is given, only quantize it if the consumer of weights # (the layer op) is in the right scope. - _InsertQuantOp( - context, - 'weights_quant', - layer_match.weight_tensor.op, - input_to_ops_map.ConsumerOperations(layer_match.weight_tensor.op), - is_training, - moving_avg=False, - ema_decay=ema_decay, - quant_delay=quant_delay, - narrow_range=True, - vars_collection=vars_collection, - bits=weight_bits, - symmetric=symmetric, - consumer_scope=scope) + if layer_match.weight_tensor is not None: + _InsertQuantOp( + context, + 'weights_quant', + layer_match.weight_tensor.op, + input_to_ops_map.ConsumerOperations(layer_match.weight_tensor.op), + is_training, + moving_avg=False, + ema_decay=ema_decay, + quant_delay=quant_delay, + narrow_range=True, + vars_collection=vars_collection, + bits=weight_bits, + symmetric=symmetric, + consumer_scope=scope) # Quantize the activations. - consumer_ops = input_to_ops_map.ConsumerOperations( - layer_match.activation_op) - add_context = context - if layer_match.bypass_op: - pattern_match_result = re.search(r'^(.*)/([^/]+)', context) - if pattern_match_result is not None: - add_context = pattern_match_result.group(1) - else: - add_context = '' - # If `scope` is given, only quantize it if the producer of weights - # (usually it's the layer op) is in the right scope. - _InsertQuantOp( - add_context, - 'act_quant', - layer_match.activation_op, - consumer_ops, - is_training, - moving_avg=True, - ema_decay=ema_decay, - quant_delay=quant_delay, - vars_collection=vars_collection, - bits=activation_bits, - symmetric=symmetric, - init_min=0.0, - producer_scope=scope) - quantized_ops.add(layer_match.activation_op) + if layer_match.activation_op is not None: + consumer_ops = input_to_ops_map.ConsumerOperations( + layer_match.activation_op) + add_context = context + if layer_match.bypass_op: + pattern_match_result = re.search(r'^(.*)/([^/]+)', context) + if pattern_match_result is not None: + add_context = pattern_match_result.group(1) + else: + add_context = '' + # If `scope` is given, only quantize it if the producer of weights + # (usually it's the layer op) is in the right scope. + _InsertQuantOp( + add_context, + 'act_quant', + layer_match.activation_op, + consumer_ops, + is_training, + moving_avg=True, + ema_decay=ema_decay, + quant_delay=quant_delay, + vars_collection=vars_collection, + bits=activation_bits, + symmetric=symmetric, + init_min=0.0, + producer_scope=scope) + quantized_ops.add(layer_match.activation_op) # Quantize the inputs and output to the bypass (if it exists). The input to # the bypass is the bias add, and the output is the activation. @@ -158,7 +160,7 @@ def Quantize(graph, # shouldn't quantize it, since the activation will be Fused into the # Add at inference time. consumers = input_to_ops_map.ConsumerOperations(layer_match.bypass_op) - if any([consumer.type in _ACTIVATION_TYPES for consumer in consumers]): + if any(consumer.type in _ACTIVATION_TYPES for consumer in consumers): logging.info('Skipping %s, because its followed by an activation.', layer_match.bypass_op.name) else: @@ -193,7 +195,7 @@ def Quantize(graph, # Add at inference time. consumers = input_to_ops_map.ConsumerOperations( layer_match.post_activation_bypass_op) - if any([consumer.type in _RELU_TYPES for consumer in consumers]): + if any(consumer.type in _RELU_TYPES for consumer in consumers): logging.info('Skipping %s, because its followed by an activation.', layer_match.post_activation_bypass_op.name) else: @@ -547,6 +549,8 @@ def _FindLayersToQuantize(graph): for match_result in sep_conv_matcher.match_graph(graph): layer_op = match_result.get_op(layer_pattern) weight_tensor = match_result.get_tensor(weight_identity_pattern) + if weight_tensor is None: + weight_tensor = match_result.get_tensor(weight_resource_var_pattern) activation_op = match_result.get_op(layer_pattern) if layer_op not in matched_layer_set: matched_layer_set.add(layer_op) @@ -681,7 +685,7 @@ def _InsertQuantOp(context, [1; 2^bits - 1] or wide range [0; 2^bits - 1]. producer_scope: The restriction of producer scope. If not None, the new op will be inserted only when the producer is in this scope. - consumer_scope: The restriction of producer scope. If not None, the new op + consumer_scope: The restriction of consumer scope. If not None, the new op will be inserted only when all the consumers are in this scope. Raises: ValueError: When producer operation is not directly connected to the diff --git a/tensorflow/contrib/quantize/python/quantize_test.py b/tensorflow/contrib/quantize/python/quantize_test.py index 212d902a3c64791adb50e7b3fa4a487f41b5bfbd..5681a213fe5eafb0814088ed34cc2253767c1d7e 100644 --- a/tensorflow/contrib/quantize/python/quantize_test.py +++ b/tensorflow/contrib/quantize/python/quantize_test.py @@ -18,6 +18,7 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +from tensorflow.contrib.framework.python.ops import variables from tensorflow.contrib.layers.python.layers import layers from tensorflow.contrib.quantize.python import quantize from tensorflow.python.framework import ops @@ -26,6 +27,7 @@ from tensorflow.python.ops import array_ops from tensorflow.python.ops import control_flow_ops from tensorflow.python.ops import init_ops from tensorflow.python.ops import math_ops +from tensorflow.python.ops import nn from tensorflow.python.ops import nn_ops from tensorflow.python.ops import partitioned_variables from tensorflow.python.ops import variable_scope @@ -525,6 +527,43 @@ class QuantizeTest(test_util.TensorFlowTestCase): self.assertTrue( 'FakeQuantWithMinMaxVars' in [i.op.type for i in reshape.op.inputs]) + def testSeparableConvWithResourceVar(self): + graph = ops.Graph() + with graph.as_default(): + with variable_scope.variable_scope('', use_resource=True): + batch_size, height, width, depth = 5, 128, 128, 3 + input1 = array_ops.zeros((batch_size, height, width, depth)) + kernel_size, depth_multiplier = 3, 1 + depthwise_shape = [kernel_size, kernel_size, depth, depth_multiplier] + depthwise_weights = variables.model_variable( + 'depthwise_weights', shape=depthwise_shape) + strides = [1, 1, 1, 1] + with variable_scope.variable_scope('depthwise_conv_1'): + conv1 = nn.depthwise_conv2d( + input1, depthwise_weights, strides, padding='SAME') + with variable_scope.variable_scope('depthwise_conv_2'): + conv2 = nn.depthwise_conv2d( + conv1, depthwise_weights, strides, padding='SAME') + math_ops.add(conv2, input1, name='add') + + quantize.Quantize(graph, True) + + # Test that the weights and activations of all convs have been quantized. + quant_node_name = 'FakeQuantWithMinMaxVars' + weights_quant = graph.get_operation_by_name( + 'depthwise_conv_1/weights_quant/' + quant_node_name) + self.assertEqual(weights_quant.type, quant_node_name) + act_quant = graph.get_operation_by_name('depthwise_conv_1/act_quant/' + + quant_node_name) + self.assertEqual(act_quant.type, quant_node_name) + + weights_quant = graph.get_operation_by_name( + 'depthwise_conv_2/weights_quant/' + quant_node_name) + self.assertEqual(weights_quant.type, quant_node_name) + act_quant = graph.get_operation_by_name('depthwise_conv_2/act_quant/' + + quant_node_name) + self.assertEqual(act_quant.type, quant_node_name) + def _WeightInit(self, stddev): """Returns truncated normal variable initializer. diff --git a/tensorflow/contrib/rate/BUILD b/tensorflow/contrib/rate/BUILD index c461a7145e27c4238161cec989448be807acd543..76db9aecf615d0a94f65cd7ea799db245828db1c 100644 --- a/tensorflow/contrib/rate/BUILD +++ b/tensorflow/contrib/rate/BUILD @@ -34,6 +34,11 @@ py_test( name = "rate_test", size = "small", srcs = ["rate_test.py"], + tags = [ + "manual", # TODO(b/120555555) + "no_oss", # TODO(b/120555555) + "notap", # TODO(b/120555555) + ], deps = [ ":rate", "//tensorflow/python:array_ops", diff --git a/tensorflow/contrib/receptive_field/README.md b/tensorflow/contrib/receptive_field/README.md index 79b015a9163f5727caa40b54579c71e57621c92f..d1c41e4c0a11028765c9fc0dc345cb29453baa31 100644 --- a/tensorflow/contrib/receptive_field/README.md +++ b/tensorflow/contrib/receptive_field/README.md @@ -185,5 +185,4 @@ Effective padding (vertical) = 1482 ## Authors -André Araujo (github id: andrefaraujo) and Mark Sandler (github id: -marksandler) +André Araujo (@andrefaraujo) and Mark Sandler (@marksandler) diff --git a/tensorflow/contrib/receptive_field/python/util/examples/compute_rf.py b/tensorflow/contrib/receptive_field/python/util/examples/compute_rf.py index d6fdd12bbe37fb0e0cb12f1d0adc3fce29b19e8a..72f98ccc32e945b48b5f1b570bcca323a5b5f48a 100644 --- a/tensorflow/contrib/receptive_field/python/util/examples/compute_rf.py +++ b/tensorflow/contrib/receptive_field/python/util/examples/compute_rf.py @@ -12,10 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Computes Receptive Field (RF) information given a graph protobuf. - -For an example of usage, see accompanying file compute_rf.sh -""" +"""Computes Receptive Field (RF) information given a graph protobuf.""" from __future__ import absolute_import from __future__ import division diff --git a/tensorflow/contrib/receptive_field/python/util/examples/rf_benchmark.py b/tensorflow/contrib/receptive_field/python/util/examples/rf_benchmark.py index a298b4d49038468299b58140758c69675368e855..325929a5937ac60a6134fae064e7633a4c57473d 100644 --- a/tensorflow/contrib/receptive_field/python/util/examples/rf_benchmark.py +++ b/tensorflow/contrib/receptive_field/python/util/examples/rf_benchmark.py @@ -16,8 +16,6 @@ The receptive field (and related parameters) for the different models are printed to stdout, and may also optionally be written to a CSV file. - -For an example of usage, see rf_benchmark.sh """ from __future__ import absolute_import @@ -262,11 +260,11 @@ def _model_rf(graphdef, information will be computed. model_type: Type of model to be used, used only for printing purposes. csv_writer: A CSV writer for RF parameters, which is used if it is not None. - input_resolution: Input resolution to use when computing RF - parameters. This is important for the case where padding can only be - defined if the input resolution is known, which may happen if using SAME - padding. This is assumed the resolution for both height and width. If - None, we consider the resolution is unknown. + input_resolution: Input resolution to use when computing RF parameters. This + is important for the case where padding can only be defined if the input + resolution is known, which may happen if using SAME padding. This is + assumed the resolution for both height and width. If None, we consider the + resolution is unknown. """ for desired_end_point_key in desired_end_point_keys: print('- %s:' % desired_end_point_key) @@ -283,10 +281,10 @@ def _model_rf(graphdef, if (receptive_field_x == receptive_field_y) and ( effective_stride_x == effective_stride_y) and ( effective_padding_x == effective_padding_y): - print('Receptive field size = %5s, effective stride = %5s, effective ' - 'padding = %5s' % (str(receptive_field_x), - str(effective_stride_x), - str(effective_padding_x))) + print( + 'Receptive field size = %5s, effective stride = %5s, effective ' + 'padding = %5s' % (str(receptive_field_x), str(effective_stride_x), + str(effective_padding_x))) else: print('Receptive field size: horizontal = %5s, vertical = %5s. ' 'Effective stride: horizontal = %5s, vertical = %5s. Effective ' @@ -362,9 +360,8 @@ def _process_model_rf(model_type='resnet_v1_50', defined if the input resolution is known, which may happen if using SAME padding. The entries in the list are assumed the resolution for both height and width. If one of the elements in the list is None, we consider - it to mean that the resolution is unknown. If the list itself is None, - we use the default list [None, 224, 321]. - + it to mean that the resolution is unknown. If the list itself is None, we + use the default list [None, 224, 321]. """ # Process default value for this list. if input_resolutions is None: @@ -477,8 +474,8 @@ def _mobilenet_v1_rf(csv_writer=None): csv_writer: A CSV writer for RF parameters, which is used if it is not None. """ for model_type in _SUPPORTED_MOBILENETV1_VARIANTS: - with slim.arg_scope( - [slim.batch_norm, slim.dropout], is_training=False) as arg_sc: + with slim.arg_scope([slim.batch_norm, slim.dropout], + is_training=False) as arg_sc: _process_model_rf(model_type, csv_writer, arg_sc) 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/receptive_field/python/util/receptive_field.py b/tensorflow/contrib/receptive_field/python/util/receptive_field.py index b9bd2f09761ab10a62d37e8e2580b93b9b8a4453..9127c772c75279d9c8eacc5a17680beba9247d01 100644 --- a/tensorflow/contrib/receptive_field/python/util/receptive_field.py +++ b/tensorflow/contrib/receptive_field/python/util/receptive_field.py @@ -12,12 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Functions to compute receptive field of a fully-convolutional network. - -Please refer to the following g3doc for detailed explanation on how this -computation is performed, and why it is important: -g3doc/photos/vision/features/delf/g3doc/rf_computation.md -""" +"""Functions to compute receptive field of a fully-convolutional network.""" from __future__ import absolute_import from __future__ import division @@ -96,8 +91,8 @@ class ReceptiveField(object): Args: y: An array of feature coordinates with shape `(..., d)`, where `d` is the number of dimensions of the coordinates. - axis: The dimensions for which to compute the input center coordinates. - If `None` (the default), compute the input center coordinates for all + axis: The dimensions for which to compute the input center coordinates. If + `None` (the default), compute the input center coordinates for all dimensions. Returns: @@ -127,8 +122,8 @@ class ReceptiveField(object): Args: x: An array of input center coordinates with shape `(..., d)`, where `d` is the number of dimensions of the coordinates. - axis: The dimensions for which to compute the feature coordinates. - If `None` (the default), compute the feature coordinates for all + axis: The dimensions for which to compute the feature coordinates. If + `None` (the default), compute the feature coordinates for all dimensions. Returns: @@ -274,14 +269,15 @@ def compute_receptive_field_from_graph_def(graph_def, continue # Get params for this layer. - (kernel_size_x, kernel_size_y, stride_x, stride_y, padding_x, - padding_y, _, _) = parse_layer_parameters.get_layer_params( + (kernel_size_x, kernel_size_y, stride_x, stride_y, padding_x, padding_y, + _, _) = parse_layer_parameters.get_layer_params( node, name_to_node, node_info[node.name].input_size) - logging.vlog(3, "kernel_size_x = %s, kernel_size_y = %s, " - "stride_x = %s, stride_y = %s, " - "padding_x = %s, padding_y = %s, input size = %s" % - (kernel_size_x, kernel_size_y, stride_x, stride_y, padding_x, - padding_y, node_info[node.name].input_size)) + logging.vlog( + 3, "kernel_size_x = %s, kernel_size_y = %s, " + "stride_x = %s, stride_y = %s, " + "padding_x = %s, padding_y = %s, input size = %s" % + (kernel_size_x, kernel_size_y, stride_x, stride_y, padding_x, + padding_y, node_info[node.name].input_size)) if padding_x is None or padding_y is None: undefined_padding = True @@ -352,15 +348,15 @@ def compute_receptive_field_from_graph_def(graph_def, raise ValueError( "Graph is not aligned since effective stride from different " "paths is different in vertical direction") - if (rf_sizes_x[inp_name] - 1 - ) / 2 - effective_paddings_x[inp_name] != ( - rf_size_input_x - 1) / 2 - effective_padding_input_x: + if (rf_sizes_x[inp_name] - + 1) / 2 - effective_paddings_x[inp_name] != ( + rf_size_input_x - 1) / 2 - effective_padding_input_x: raise ValueError( "Graph is not aligned since center shift from different " "paths is different in horizontal direction") - if (rf_sizes_y[inp_name] - 1 - ) / 2 - effective_paddings_y[inp_name] != ( - rf_size_input_y - 1) / 2 - effective_padding_input_y: + if (rf_sizes_y[inp_name] - + 1) / 2 - effective_paddings_y[inp_name] != ( + rf_size_input_y - 1) / 2 - effective_padding_input_y: raise ValueError( "Graph is not aligned since center shift from different " "paths is different in vertical direction") diff --git a/tensorflow/contrib/resampler/BUILD b/tensorflow/contrib/resampler/BUILD index b3f32b8f34e7b956b44bc82322bba16ed6fe43c7..bbf109967595a73a0fc4bacaf34859b30c2376fc 100644 --- a/tensorflow/contrib/resampler/BUILD +++ b/tensorflow/contrib/resampler/BUILD @@ -13,6 +13,7 @@ load( ) load("//tensorflow:tensorflow.bzl", "cuda_py_test") load("//tensorflow:tensorflow.bzl", "tf_custom_op_py_library") +load("//tensorflow/compiler/tests:build_defs.bzl", "tf_xla_py_test") tf_custom_op_py_library( name = "resampler_py", @@ -52,7 +53,12 @@ tf_kernel_library( ":resampler_ops_op_lib", "//tensorflow/core:framework", "//tensorflow/core:lib", - ], + ] + select({ + "//tensorflow:with_xla_support": [ + "//tensorflow/compiler/tf2xla/kernels:resampler_ops", + ], + "//conditions:default": [], + }), alwayslink = 1, ) @@ -93,3 +99,26 @@ cuda_py_test( "//tensorflow/python:array_ops", ], ) + +tf_xla_py_test( + name = "resampler_ops_xla_test", + size = "small", + srcs = ["xla/resampler_ops_xla_test.py"], + disabled_backends = [ + # TODO(b/74459949) Support BatchDot in CPU backend. + "cpu", + "cpu_ondemand", + ], + # TODO(b/112295522): the OSS build will not likely work in the short to medium term, currently it is blocked by the fact that bazel does not allow py_library to depend on cc_library: https://github.com/bazelbuild/bazel/issues/701 which may not be resolvable. + tags = ["no_oss"], + deps = [ + "//tensorflow/compiler/tests:xla_test", + "//tensorflow/compiler/tf2xla/kernels:resampler_ops", + "//tensorflow/contrib/resampler:resampler_ops", + "//tensorflow/contrib/resampler:resampler_py", + "//tensorflow/python:array_ops", + "//tensorflow/python:client_testlib", + "//tensorflow/python:platform_test", + "//third_party/py/numpy", + ], +) diff --git a/tensorflow/contrib/resampler/ops/resampler_ops.cc b/tensorflow/contrib/resampler/ops/resampler_ops.cc index 5ab212032e50ace9545762bebda5679f68fbf77c..f785d4ee5fcd63212882ccf736bfc61c35d68545 100644 --- a/tensorflow/contrib/resampler/ops/resampler_ops.cc +++ b/tensorflow/contrib/resampler/ops/resampler_ops.cc @@ -25,7 +25,7 @@ REGISTER_OP("Resampler") .Input("data: T") .Input("warp: T") .Output("output: T") - .Attr("T: {half, float, double}") + .Attr("T: {half, bfloat16, float, double}") .SetShapeFn([](InferenceContext* c) { ShapeHandle data; ShapeHandle warp; @@ -48,7 +48,7 @@ REGISTER_OP("ResamplerGrad") .Input("grad_output: T") .Output("grad_data: T") .Output("grad_warp: T") - .Attr("T: {half, float, double}") + .Attr("T: {half, bfloat16, float, double}") .SetShapeFn([](InferenceContext* c) { c->set_output(0, c->input(0)); c->set_output(1, c->input(1)); diff --git a/tensorflow/contrib/resampler/python/ops/resampler_ops.py b/tensorflow/contrib/resampler/python/ops/resampler_ops.py index 8b632527f6b1fc08454c77deac181b4f9c4e5f5f..0ee224a47821b0530304aa448eaca0b1b9f59d02 100644 --- a/tensorflow/contrib/resampler/python/ops/resampler_ops.py +++ b/tensorflow/contrib/resampler/python/ops/resampler_ops.py @@ -39,7 +39,9 @@ def resampler(data, warp, name="resampler"): data_num_channels]` containing 2D data that will be resampled. warp: Tensor of minimum rank 2 containing the coordinates at which resampling will be performed. Since only bilinear interpolation is - currently supported, the last dimension of the `warp` tensor must be 2. + currently supported, the last dimension of the `warp` tensor must be 2, + representing the (x, y) coordinate where x is the index for width and y is + the index for height. name: Optional name of the op. Returns: diff --git a/tensorflow/contrib/resampler/xla/resampler_ops_xla_test.py b/tensorflow/contrib/resampler/xla/resampler_ops_xla_test.py new file mode 100644 index 0000000000000000000000000000000000000000..cec4c3c23305034d167a248a637425507750064e --- /dev/null +++ b/tensorflow/contrib/resampler/xla/resampler_ops_xla_test.py @@ -0,0 +1,241 @@ +# 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 resampler ops.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.compiler.tests import xla_test +from tensorflow.contrib import resampler +from tensorflow.contrib.resampler.ops import gen_resampler_ops +from tensorflow.python.ops import array_ops +from tensorflow.python.platform import test + + +class ResamplerOpsTest(xla_test.XLATestCase): + + def _assertForwardOpMatchesExpected(self, image_np, warp_np, expected): + with self.test_session() as sess, self.test_scope(): + input_image = array_ops.placeholder(image_np.dtype) + warp = array_ops.placeholder(warp_np.dtype) + resampled = resampler.resampler(input_image, warp, name='resampler') + out = sess.run(resampled, {input_image: image_np, warp: warp_np}) + + self.assertAllCloseAccordingToType( + expected, out, rtol=5e-3, half_rtol=1e-2, bfloat16_rtol=3e-2) + + def _assertBackwardOpMatchesExpected(self, input_np, warp_np, grad_output_np, + expected_grad_data, expected_grad_warp): + with self.cached_session() as sess, self.test_scope(): + input_image = array_ops.placeholder(input_np.dtype) + warp = array_ops.placeholder(warp_np.dtype) + grad_output = array_ops.placeholder(grad_output_np.dtype) + + grad_data, grad_warp = gen_resampler_ops.resampler_grad( + input_image, warp, grad_output) + + grad_data_tf, grad_warp_tf = sess.run([grad_data, grad_warp], { + input_image: input_np, + warp: warp_np, + grad_output: grad_output_np + }) + + self.assertAllCloseAccordingToType( + expected_grad_warp, grad_warp_tf, half_rtol=1e-2, bfloat16_rtol=3e-2) + self.assertAllCloseAccordingToType( + expected_grad_data, grad_data_tf, half_rtol=1e-2, bfloat16_rtol=3e-2) + + def testSimple(self): + for dtype in self.float_types: + input_shape = [1, 2, 2, 1] + input_data = [0, 5, 13, 54] + input_np = np.array(input_data, dtype=dtype).reshape(input_shape) + + warp_shape = [1, 2] + warp_data = [0.7, 0.6] + warp_np = np.array(warp_data, dtype=dtype).reshape(warp_shape) + expected = [[26.42]] + self._assertForwardOpMatchesExpected(input_np, warp_np, expected) + + grad_output = np.ones([1, 1], dtype=dtype) + + expected_grad_data = [[[[0.12], [0.27999997]], [[0.18000001], + [0.42000002]]]] + + expected_grad_warp = [[26.60000038, 38.20000076]] + + self._assertBackwardOpMatchesExpected(input_np, warp_np, grad_output, + expected_grad_data, + expected_grad_warp) + + def testMultiChannel(self): + for dtype in self.float_types: + input_shape = [1, 2, 2, 3] + input_rgb_data = [0, 5, 13, 54, 135, 226, 37, 8, 234, 90, 255, 1] + input_np = np.array(input_rgb_data, dtype=dtype).reshape(input_shape) + + warp_shape = [1, 2] + warp_data = [0.7, 0.6] + warp_np = np.array(warp_data, dtype=dtype).reshape(warp_shape) + expected = [[59.58000183, 146.94000244, 107.37999725]] + self._assertForwardOpMatchesExpected(input_np, warp_np, expected) + + grad_output = np.ones([1, 3], dtype=dtype) + + expected_grad_data = [[[[0.12, 0.12, 0.12], + [0.27999997, 0.27999997, 0.27999997]], + [[0.18000001, 0.18000001, 0.18000001], + [0.42000002, 0.42000002, 0.42000002]]]] + + expected_grad_warp = [[199, 30]] + + self._assertBackwardOpMatchesExpected(input_np, warp_np, grad_output, + expected_grad_data, + expected_grad_warp) + + def testBatch2Height3byWidth3RGB(self): + for dtype in self.float_types: + input_shape = [2, 3, 3, 3] + input_rgb_data = [ + 0, 5, 13, 54, 135, 226, 37, 8, 234, 90, 255, 1, 30, 105, 2, 40, 115, + 3, 50, 125, 4, 60, 135, 5, 70, 145, 6, 0, 5, 13, 54, 135, 226, 37, 8, + 234, 90, 255, 1, 30, 105, 2, 40, 115, 3, 50, 125, 4, 60, 135, 5, 70, + 145, 6 + ] + input_np = np.array(input_rgb_data, dtype=dtype).reshape(input_shape) + + # 2 batches and 2 samples for each batch. + warp_shape = [2, 2, 2] + warp_data = [0.7, 0.6, 1, 0.7, 0.9, 1.2, 1.3, 1.6] + warp_np = np.array(warp_data, dtype=dtype).reshape(warp_shape) + + expected_forward = [[[43.92, 128.4, 65.86], [37.2, 114., 69.2]], + [[40.6, 122.8, 2.5], [51., 126, 4.1]]] + + self._assertForwardOpMatchesExpected(input_np, warp_np, expected_forward) + + expected_grad_data = [[[[0.12, 0.12, 0.12], + [0.57999998, 0.57999998, 0.57999998], + [0., 0., 0.]], + [[0.18000001, 0.18000001, 0.18000001], + [1.12, 1.12, 1.12], [0., 0., 0.]], + [[0., 0., 0.], [0., 0., 0.], [0., 0., 0.]]], + [[[0., 0., 0.], [0., 0., 0.], [0., 0., 0.]], + [[0.08000001, 0.08000001, 0.08000001], + [0.99999988, 0.99999988, 0.99999988], + [0.11999997, 0.11999997, 0.11999997]], + [[0.02000001, 0.02000001, 0.02000001], + [0.60000008, 0.60000008, 0.60000008], + [0.17999998, 0.17999998, 0.17999998]]]] + expected_grad_warp = [[[33.39999008, -96.20000458], [-26.10000229, + -278.]], + [[-162.99998474, 39.99999619], [21., 63.]]] + + grad_output = np.ones([2, 2, 3], dtype=dtype) + self._assertBackwardOpMatchesExpected(input_np, warp_np, grad_output, + expected_grad_data, + expected_grad_warp) + + def testOutOfBoundWarps(self): + # (x, y) are both less than 0. + for dtype in self.float_types: + input_shape = [1, 2, 2, 1] + input_data = [10, 5, 13, 54] + input_np = np.array(input_data, dtype=dtype).reshape(input_shape) + + warp_shape = [1, 2, 2] + warp_data = [-1, -1, 0.7, 0.6] + warp_np = np.array(warp_data, dtype=dtype).reshape(warp_shape) + expected = [[[0.0], [27.62]]] + self._assertForwardOpMatchesExpected(input_np, warp_np, expected) + + expected_grad_data = [[[[0.12], [0.27999997]], [[0.18000001], + [0.42000002]]]] + expected_grad_warp = [[[0., 0.], [22.60000038, 35.20000076]]] + + grad_output = np.ones([1, 2, 1], dtype=dtype) + self._assertBackwardOpMatchesExpected(input_np, warp_np, grad_output, + expected_grad_data, + expected_grad_warp) + + # One of (x, y) is less than 0. + for dtype in self.float_types: + input_shape = [1, 2, 2, 1] + input_data = [10, 5, 13, 54] + input_np = np.array(input_data, dtype=dtype).reshape(input_shape) + + warp_shape = [1, 2, 2] + # -1 is out of bound for grad_warp. + warp_data = [-1, 0.1, 0.7, 0.6] + warp_np = np.array(warp_data, dtype=dtype).reshape(warp_shape) + expected = [[[0.0], [27.62]]] + self._assertForwardOpMatchesExpected(input_np, warp_np, expected) + + expected_grad_data = [[[[0.12], [0.27999997]], [[0.18000001], + [0.42000002]]]] + expected_grad_warp = [[[0., 0.], [22.60000038, 35.20000076]]] + + grad_output = np.ones([1, 2, 1], dtype=dtype) + self._assertBackwardOpMatchesExpected(input_np, warp_np, grad_output, + expected_grad_data, + expected_grad_warp) + + # Both of (x, y) are greater than image size. + for dtype in self.float_types: + input_shape = [1, 2, 2, 1] + input_data = [10, 5, 13, 54] + input_np = np.array(input_data, dtype=dtype).reshape(input_shape) + + warp_shape = [1, 2, 2] + # -0.1 is *inbound* for grad_warp and grad_data, 2.1 is out of bound. + warp_data = [-0.1, 0.1, 1.2, 2.1] + warp_np = np.array(warp_data, dtype=dtype).reshape(warp_shape) + expected = [[[0.0], [0.0]]] + self._assertForwardOpMatchesExpected(input_np, warp_np, expected) + + expected_grad_data = [[[[0.81], [0.0]], [[0.09], [0.0]]]] + expected_grad_warp = [[[10.30, 2.7], [0.0, 0.0]]] + + grad_output = np.ones([1, 2, 1], dtype=dtype) + self._assertBackwardOpMatchesExpected(input_np, warp_np, grad_output, + expected_grad_data, + expected_grad_warp) + + # One of (x, y) is greater than image size. + for dtype in self.float_types: + input_shape = [1, 2, 2, 1] + input_data = [10, 5, 13, 54] + input_np = np.array(input_data, dtype=dtype).reshape(input_shape) + + warp_shape = [1, 2, 2] + warp_data = [0.1, -0.1, 1.2, 0.1] + warp_np = np.array(warp_data, dtype=dtype).reshape(warp_shape) + expected = [[[0.0], [0.0]]] + self._assertForwardOpMatchesExpected(input_np, warp_np, expected) + + expected_grad_data = [[[[0.81], [0.81]], [[0.0], [0.08]]]] + expected_grad_warp = [[[-4.5, 9.5], [-9.9, 39.20]]] + + grad_output = np.ones([1, 2, 1], dtype=dtype) + self._assertBackwardOpMatchesExpected(input_np, warp_np, grad_output, + expected_grad_data, + expected_grad_warp) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/rnn/BUILD b/tensorflow/contrib/rnn/BUILD index 391df8cdb4b1c6cd0e22ff2e27527c58abd4c303..d65d80df8073ef70d591c4ae2af99132f1c318ef 100644 --- a/tensorflow/contrib/rnn/BUILD +++ b/tensorflow/contrib/rnn/BUILD @@ -118,6 +118,7 @@ cuda_py_tests( "//tensorflow/python:rnn_cell", "//tensorflow/python:variable_scope", "//tensorflow/python:variables", + "@absl_py//absl/testing:parameterized", ], ) @@ -196,6 +197,7 @@ cuda_py_tests( srcs = ["python/kernel_tests/lstm_ops_test.py"], additional_deps = [ ":rnn_py", + "@absl_py//absl/testing:parameterized", "//third_party/py/numpy", "//tensorflow/python:array_ops", "//tensorflow/python:client_testlib", @@ -225,7 +227,10 @@ tf_custom_op_library( "kernels/lstm_ops_gpu.cu.cc", "kernels/lstm_ops.h", ], - deps = ["//tensorflow/core/kernels:eigen_helpers"], + deps = [ + "//tensorflow/core/kernels:eigen_contraction_kernel", + "//tensorflow/core/kernels:eigen_helpers", + ], ) tf_gen_op_wrapper_py( @@ -247,7 +252,10 @@ tf_custom_op_library( "kernels/gru_ops_gpu.cu.cc", "kernels/gru_ops.h", ], - deps = ["//tensorflow/core/kernels:eigen_helpers"], + deps = [ + "//tensorflow/core/kernels:eigen_contraction_kernel", + "//tensorflow/core/kernels:eigen_helpers", + ], ) tf_gen_op_wrapper_py( @@ -344,6 +352,7 @@ tf_kernel_library( deps = [ "//tensorflow/core:framework", "//tensorflow/core:lib", + "//tensorflow/core/kernels:eigen_contraction_kernel", "//tensorflow/core/kernels:eigen_helpers", "//third_party/eigen3", ], @@ -379,6 +388,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", @@ -397,7 +413,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/__init__.py b/tensorflow/contrib/rnn/__init__.py index 026bf08ced33cf0d663cf0940e8bea3f3f2aca28..cbc8af5350276bf3398cf29a24554fd27e0621ee 100644 --- a/tensorflow/contrib/rnn/__init__.py +++ b/tensorflow/contrib/rnn/__init__.py @@ -14,8 +14,6 @@ # ============================================================================== """RNN Cells and additional RNN operations. -See [Contrib RNN](https://tensorflow.org/api_guides/python/contrib.rnn) guide. - @@RNNCell @@LayerRNNCell diff --git a/tensorflow/contrib/rnn/kernels/blas_gemm.h b/tensorflow/contrib/rnn/kernels/blas_gemm.h index 9535a76566748eaf8b4756ad0dc26218262ed990..12f3182a6a8878aa27ee143fa6405903e3fc4ef3 100644 --- a/tensorflow/contrib/rnn/kernels/blas_gemm.h +++ b/tensorflow/contrib/rnn/kernels/blas_gemm.h @@ -21,6 +21,10 @@ limitations under the License. #include "tensorflow/core/kernels/eigen_activations.h" #include "tensorflow/core/platform/types.h" +#if defined(TENSORFLOW_USE_CUSTOM_CONTRACTION_KERNEL) +#include "tensorflow/core/kernels/eigen_contraction_kernel.h" +#endif + namespace tensorflow { class OpKernelContext; namespace functor { @@ -32,15 +36,26 @@ struct TensorCuBlasGemm { const T* b, int ldb, float beta, T* c, int ldc); }; +template +struct gemm_compute_type { + typedef T type; +}; + +template <> +struct gemm_compute_type { + typedef float type; +}; + template struct TensorBlasGemm; template struct TensorBlasGemm { static void compute(OpKernelContext* ctx, const Device& d, bool transa, - bool transb, float alpha, + bool transb, typename gemm_compute_type::type alpha, typename TTypes::ConstMatrix a, - typename TTypes::ConstMatrix b, float beta, + typename TTypes::ConstMatrix b, + typename gemm_compute_type::type beta, typename TTypes::Matrix c) { int64 m = c.dimensions()[0]; int64 n = c.dimensions()[1]; @@ -55,19 +70,23 @@ struct TensorBlasGemm { template struct TensorBlasGemm { static void compute(OpKernelContext* ctx, const Device& d, bool transa, - bool transb, T alpha, typename TTypes::ConstMatrix a, - typename TTypes::ConstMatrix b, T beta, + bool transb, typename gemm_compute_type::type alpha, + typename TTypes::ConstMatrix a, + typename TTypes::ConstMatrix b, + typename gemm_compute_type::type beta, typename TTypes::Matrix c) { Eigen::array, 1> contract_pairs; contract_pairs[0] = Eigen::IndexPair(transa == false, transb == true); - if (alpha == T(1) && beta == T(0)) { + if (alpha == typename gemm_compute_type::type(1.f) && + beta == typename gemm_compute_type::type(0.f)) { c.device(d) = a.contract(b, contract_pairs); - } else if (alpha == T(1) && beta == T(1)) { + } else if (alpha == typename gemm_compute_type::type(1.f) && + beta == typename gemm_compute_type::type(1.f)) { c.device(d) += a.contract(b, contract_pairs); } else { - c.device(d) = c.constant(alpha) * a.contract(b, contract_pairs) + - c.constant(beta) * c; + c.device(d) = c.constant(T(alpha)) * a.contract(b, contract_pairs) + + c.constant(T(beta)) * c; } } }; diff --git a/tensorflow/contrib/rnn/kernels/gru_ops.h b/tensorflow/contrib/rnn/kernels/gru_ops.h index 3e2cb39e64bb3f0b22ea66c5601af36c5fb9b0fd..38be58fa104f8b30e4aede6d18330960fc30dcb5 100644 --- a/tensorflow/contrib/rnn/kernels/gru_ops.h +++ b/tensorflow/contrib/rnn/kernels/gru_ops.h @@ -88,7 +88,9 @@ struct GRUBlockCellFprop : public GRUCell { typename TTypes::ConstMatrix const_x_h_prev(x_h_prev.data(), x_h_prev.dimensions()); TensorBlasGemm::compute( - ctx, d, false, false, T(1), const_x_h_prev, w_ru, T(0), r_u_bar); + ctx, d, false, false, typename gemm_compute_type::type(1.f), + const_x_h_prev, w_ru, typename gemm_compute_type::type(0.f), + r_u_bar); // Creating a bias matrix for adding by broadcasting 'b_ru' Eigen::array broadcast_shape({batch_size_, 1}); @@ -107,7 +109,8 @@ struct GRUBlockCellFprop : public GRUCell { typename TTypes::ConstMatrix const_x_h_prevr(x_h_prevr.data(), x_h_prevr.dimensions()); TensorBlasGemm::compute( - ctx, d, false, false, T(1), const_x_h_prevr, w_c, T(0), c); + ctx, d, false, false, typename gemm_compute_type::type(1.f), + const_x_h_prevr, w_c, typename gemm_compute_type::type(0.f), c); Eigen::array b_c_shape({1, b_c.dimensions()[0]}); c.device(d) += (b_c.reshape(b_c_shape).broadcast(broadcast_shape)); @@ -148,9 +151,10 @@ struct GRUBlockCellBprop : public GRUCell { // [2nd_component_of_d_x d_h_prevr] = d_c_bar X w_c^T typename TTypes::ConstMatrix const_d_c_bar(d_c_bar.data(), d_c_bar.dimensions()); - TensorBlasGemm::compute(ctx, d, false, true, T(1), - const_d_c_bar, w_c, T(0), - d_x_comp2_and_h_prevr); + TensorBlasGemm::compute( + ctx, d, false, true, typename gemm_compute_type::type(1.f), + const_d_c_bar, w_c, typename gemm_compute_type::type(0.f), + d_x_comp2_and_h_prevr); d_hr.device(d) = d_x_comp2_and_h_prevr.slice(h_offsets(), h_extends()); d_r_bar.device(d) = (d_hr * h_prev * r) * (r.constant(T(1)) - r); @@ -164,7 +168,8 @@ struct GRUBlockCellBprop : public GRUCell { typename TTypes::ConstMatrix const_d_r_bar_u_bar( d_r_bar_u_bar.data(), d_r_bar_u_bar.dimensions()); TensorBlasGemm::compute( - ctx, d, false, true, T(1), const_d_r_bar_u_bar, w_ru, T(0), + ctx, d, false, true, typename gemm_compute_type::type(1.f), + const_d_r_bar_u_bar, w_ru, typename gemm_compute_type::type(0.f), d_x_comp1_and_h_prev_comp1); // d_x = d_x_comp1 + d_x_comp2 diff --git a/tensorflow/contrib/rnn/kernels/lstm_ops.cc b/tensorflow/contrib/rnn/kernels/lstm_ops.cc index ee08d306f84baaba8b774ce3fa1a04d5f9a4f6dd..d369bc12ae88dafb4e3ca0095a08bcc3ee09bf70 100644 --- a/tensorflow/contrib/rnn/kernels/lstm_ops.cc +++ b/tensorflow/contrib/rnn/kernels/lstm_ops.cc @@ -61,7 +61,8 @@ void LSTMBlockCellFpropWithEigen( // states1 = xh * w + b typename TTypes::ConstMatrix const_xh(xh.data(), xh.dimensions()); TensorBlasGemm::compute( - ctx, d, false, false, T(1), const_xh, w, T(0), icfo); + ctx, d, false, false, typename gemm_compute_type::type(1.f), const_xh, + w, typename gemm_compute_type::type(0.f), icfo); Eigen::array b_shape({1, b.dimensions()[0]}); Eigen::array broadcast_shape({cell.batch_size(), 1}); icfo.device(d) += b.reshape(b_shape).broadcast(broadcast_shape); @@ -87,11 +88,11 @@ void LSTMBlockCellFpropWithEigen( if (use_peephole) { auto f_peep = cs_prev * wcf.reshape(p_shape).broadcast(p_broadcast_shape); f.device(d) = (icfo.slice(cell.icfo_f_offsets(), cell.cell_extents()) + - f.constant(forget_bias) + f_peep) + f.constant(T(forget_bias)) + f_peep) .sigmoid(); } else { f.device(d) = (icfo.slice(cell.icfo_f_offsets(), cell.cell_extents()) + - f.constant(forget_bias)) + f.constant(T(forget_bias))) .sigmoid(); } @@ -100,7 +101,7 @@ void LSTMBlockCellFpropWithEigen( if (cell_clip > 0.0f) { cs.device(d) = - cs.binaryExpr(cs.constant(cell_clip), Eigen::scalar_clip_op()); + cs.binaryExpr(cs.constant(T(cell_clip)), Eigen::scalar_clip_op()); } // co = tanh(cs) @@ -225,6 +226,7 @@ void LSTMBlockCellBpropWithEigen( template struct LSTMBlockCellBprop; DEFINE_CPU_SPECS(float); +DEFINE_CPU_SPECS(Eigen::half); #undef DEFINE_CPU_SPECS } // namespace functor @@ -373,7 +375,7 @@ class LSTMBlockCellOp : public OpKernel { Name("LSTMBlockCell").Device(DEVICE_CPU).TypeConstraint("T"), \ LSTMBlockCellOp); REGISTER_KERNEL(float); -// REGISTER_KERNEL(double); +REGISTER_KERNEL(Eigen::half); #undef REGISTER_KERNEL #if GOOGLE_CUDA @@ -398,7 +400,6 @@ namespace functor { DECLARE_GPU_SPEC(float); DECLARE_GPU_SPEC(Eigen::half); -// DECLARE_GPU_SPEC(double); #undef DECLARE_GPU_SPEC } // end namespace functor @@ -661,7 +662,7 @@ class LSTMBlockCellGradOp : public OpKernel { Name("LSTMBlockCellGrad").Device(DEVICE_CPU).TypeConstraint("T"), \ LSTMBlockCellGradOp); REGISTER_KERNEL(float); -// REGISTER_KERNEL(double); +REGISTER_KERNEL(Eigen::half); #undef REGISTER_KERNEL #if GOOGLE_CUDA @@ -1008,7 +1009,7 @@ class BlockLSTMOp : public OpKernel { Name("BlockLSTM").Device(DEVICE_CPU).TypeConstraint("T"), \ BlockLSTMOp); REGISTER_KERNEL(float); -// REGISTER_KERNEL(double); +REGISTER_KERNEL(Eigen::half); #undef REGISTER_KERNEL #if GOOGLE_CUDA @@ -1283,7 +1284,7 @@ class BlockLSTMGradOp : public OpKernel { Name("BlockLSTMGrad").Device(DEVICE_CPU).TypeConstraint("T"), \ BlockLSTMGradOp); REGISTER_KERNEL(float); -// REGISTER_KERNEL(double); +REGISTER_KERNEL(Eigen::half); #undef REGISTER_KERNEL #if GOOGLE_CUDA diff --git a/tensorflow/contrib/rnn/kernels/lstm_ops_gpu.cu.cc b/tensorflow/contrib/rnn/kernels/lstm_ops_gpu.cu.cc index b664b0f45ee08648e4dc10e8244340df1615ad19..15ae95f13cffa5d1469d737b23f2a83b9e5a694f 100644 --- a/tensorflow/contrib/rnn/kernels/lstm_ops_gpu.cu.cc +++ b/tensorflow/contrib/rnn/kernels/lstm_ops_gpu.cu.cc @@ -141,7 +141,7 @@ __global__ void lstm_gates(const T* icfo, const T* b, const T* cs_prev, // const int gid = batch_id * cell_size * 4 + act_id; const int cid = batch_id * cell_size + act_id; - Eigen::internal::scalar_sigmoid_op sigmoid_op; + Eigen::internal::scalar_logistic_op sigmoid_op; Eigen::internal::scalar_tanh_op tanh_op; Eigen::scalar_clip_op clip_op; @@ -169,7 +169,7 @@ __global__ void lstm_gates(const T* icfo, const T* b, const T* cs_prev, f[cid] = f_local; T cs_local = i_local * ci_local + f_local * cs_prev[cid]; - if (cell_clip_t > strict_cast(0.0f)) { + if (cell_clip > 0.0f) { cs_local = clip_op(cs_local, cell_clip_t); } cs[cid] = cs_local; @@ -248,7 +248,8 @@ void LSTMBlockCellFpropWithCUDA( // states1 = xh * w typename TTypes::ConstMatrix const_xh(xh.data(), xh.dimensions()); TensorBlasGemm::compute( - ctx, d, false, false, 1.f, const_xh, w, 0.f, icfo); + ctx, d, false, false, typename gemm_compute_type::type(1.f), const_xh, + w, typename gemm_compute_type::type(0.f), icfo); // Add bias, apply non-linearities and gating. // 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 index 245fa68eaef43ca8bc18c6087460d946228b0c85..7bad4a60a149011d5b8d745f45359fd25473e54e 100644 --- a/tensorflow/contrib/rnn/python/kernel_tests/core_rnn_cell_test.py +++ b/tensorflow/contrib/rnn/python/kernel_tests/core_rnn_cell_test.py @@ -20,6 +20,7 @@ from __future__ import print_function import os +from absl.testing import parameterized import numpy as np from tensorflow.contrib import rnn as contrib_rnn @@ -31,6 +32,8 @@ 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 @@ -207,6 +210,35 @@ class RNNCellTest(test.TestCase): # 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_lib.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( @@ -805,12 +837,13 @@ class RNNCellTest(test.TestCase): self.assertAllClose(res[1], [[0.13248, 0.13248]]) -class DropoutWrapperTest(test.TestCase): +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( @@ -832,7 +865,7 @@ class DropoutWrapperTest(test.TestCase): constant([[0.1, 0.1, 0.1]] * batch_size, dtype=dtypes.float32) ] * 2) outputs, final_state = rnn.dynamic_rnn( - cell=rnn_cell_impl.DropoutWrapper( + cell=wrapper_type( rnn_cell_impl.LSTMCell(3), dtype=x.dtype, **kwargs), time_major=True, parallel_iterations=parallel_iterations, @@ -845,16 +878,34 @@ class DropoutWrapperTest(test.TestCase): self.assertEqual(res[1].h.shape, (batch_size, 3)) return res - def testWrappedCellProperty(self): + @parameterized.parameters( + [rnn_cell_impl.DropoutWrapper, rnn_cell_impl.DropoutWrapperV2]) + def testDropoutWrapperProperties(self, wrapper_type): cell = rnn_cell_impl.BasicRNNCell(10) - wrapper = rnn_cell_impl.DropoutWrapper(cell) + wrapper = wrapper_type(cell) # Github issue 15810 self.assertEqual(wrapper.wrapped_cell, cell) - - def testDropoutWrapperKeepAllConstantInput(self): + 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) + 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) @@ -864,10 +915,13 @@ class DropoutWrapperTest(test.TestCase): self.assertAllClose(true_full_output[1], res[1].h) self.assertAllClose(true_full_final_c, res[1].c) - def testDropoutWrapperKeepAll(self): + @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) + 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) @@ -877,7 +931,9 @@ class DropoutWrapperTest(test.TestCase): self.assertAllClose(true_full_output[1], res[1].h) self.assertAllClose(true_full_final_c, res[1].c) - def testDropoutWrapperWithSeed(self): + @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 @@ -889,7 +945,8 @@ class DropoutWrapperTest(test.TestCase): output_keep_prob=keep_some, state_keep_prob=keep_some, seed=10, - parallel_iterations=1) + 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() @@ -899,18 +956,22 @@ class DropoutWrapperTest(test.TestCase): output_keep_prob=keep_some, state_keep_prob=keep_some, seed=10, - parallel_iterations=1) + 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) - def testDropoutWrapperKeepNoOutput(self): + @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-10) + 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) + 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) @@ -920,15 +981,18 @@ class DropoutWrapperTest(test.TestCase): self.assertAllClose(true_full_output[1], res[1].h) self.assertAllClose(true_full_final_c, res[1].c) - def testDropoutWrapperKeepNoStateExceptLSTMCellMemory(self): + @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-10) + 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) + 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]]], @@ -941,9 +1005,11 @@ class DropoutWrapperTest(test.TestCase): # c state of an LSTMStateTuple is NEVER modified. self.assertAllClose(true_c_state, res[1].c) - def testDropoutWrapperKeepNoInput(self): + @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-10) + 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) @@ -953,12 +1019,15 @@ class DropoutWrapperTest(test.TestCase): res = self._testDropoutWrapper( input_keep_prob=keep_none, output_keep_prob=keep_all, - state_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) - def testDropoutWrapperRecurrentOutput(self): + @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( @@ -966,6 +1035,7 @@ class DropoutWrapperTest(test.TestCase): 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) @@ -974,13 +1044,16 @@ class DropoutWrapperTest(test.TestCase): for m in output_mask[1:]: self.assertAllClose(output_mask[0], m) - def testDropoutWrapperRecurrentStateInputAndOutput(self): + @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) @@ -1002,7 +1075,10 @@ class DropoutWrapperTest(test.TestCase): for batch_entry in state_h_mask: self.assertAllClose(batch_entry, state_h_mask[0]) - def testDropoutWrapperRecurrentStateInputAndOutputWithSeed(self): + @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) @@ -1011,6 +1087,7 @@ class DropoutWrapperTest(test.TestCase): 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, @@ -1024,6 +1101,7 @@ class DropoutWrapperTest(test.TestCase): 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, @@ -1050,6 +1128,60 @@ class DropoutWrapperTest(test.TestCase): 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: diff --git a/tensorflow/contrib/rnn/python/kernel_tests/core_rnn_test.py b/tensorflow/contrib/rnn/python/kernel_tests/core_rnn_test.py index 5cba54dd3df5bbb33380505bd5a073f069a3a590..ef372b947cedf71e9d44423f10cc43375b467cd9 100644 --- a/tensorflow/contrib/rnn/python/kernel_tests/core_rnn_test.py +++ b/tensorflow/contrib/rnn/python/kernel_tests/core_rnn_test.py @@ -227,7 +227,7 @@ class RNNTest(test.TestCase): def testDropout(self): cell = Plus1RNNCell() full_dropout_cell = rnn_cell.DropoutWrapper( - cell, input_keep_prob=1e-12, seed=0) + cell, input_keep_prob=1e-6, seed=0) (name, dep), = full_dropout_cell._checkpoint_dependencies self.assertIs(dep, cell) self.assertEqual("cell", name) diff --git a/tensorflow/contrib/rnn/python/kernel_tests/lstm_ops_test.py b/tensorflow/contrib/rnn/python/kernel_tests/lstm_ops_test.py index 9ce0b399ba173b67285e907a050c71af5d57068c..d5700d2a200f6cdac06183366c0d11ec3531235b 100644 --- a/tensorflow/contrib/rnn/python/kernel_tests/lstm_ops_test.py +++ b/tensorflow/contrib/rnn/python/kernel_tests/lstm_ops_test.py @@ -18,6 +18,7 @@ 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.rnn.python.kernel_tests import benchmarking @@ -27,6 +28,8 @@ 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 gen_array_ops +from tensorflow.python.ops import gen_bitwise_ops from tensorflow.python.ops import gradients_impl from tensorflow.python.ops import init_ops from tensorflow.python.ops import rnn @@ -38,7 +41,70 @@ from tensorflow.python.platform import test block_lstm = lstm_ops._block_lstm # pylint: disable=protected-access -def blocks_match(sess, use_peephole): +class _MaskedRandomUniformInitializer(init_ops.RandomUniform): + """Initializer for uniform dist tensors with trailing bits zeroed-out. + + Allow returning tensors with last few mantissa bits set to 0. This potentially + helps avoid getting into precision issues when testing low precision (float16) + computation. + """ + + def __init__(self, + minval=0, + maxval=None, + seed=None, + dtype=dtypes.float16, + num_valid_mantissa_bits=4): + """Constructor. + + Args: + minval: A python scalar or a scalar tensor. Lower bound of the range of + random values to generate. + maxval: A python scalar or a scalar tensor. Upper bound of the range of + random values to generate. Defaults to 1 for float types. + seed: A Python integer. Used to create random seeds. See + `tf.set_random_seed` for behavior. + dtype: The data type. Only supports tf.float16 for now. + num_valid_mantissa_bits: number of non-zero mantissa bits, default to 4. + + Raises: + ValueError: An error if `dtype` is not tf.float16. + """ + if dtype not in (dtypes.float16,): + raise ValueError("dtype: %s not supported" % dtype.name) + + super(_MaskedRandomUniformInitializer, self).__init__( + minval=minval, maxval=maxval, seed=seed, dtype=dtype) + self._num_mantissa_bits = 10 + self._num_valid_mantissa_bits = num_valid_mantissa_bits + + def __call__(self, shape, dtype=dtypes.float16, partition_info=None): + if dtype and dtype != dtypes.float16: + raise ValueError("dtype: %s not supported" % dtype.name) + res = super(_MaskedRandomUniformInitializer, self).__call__( + shape, dtype, partition_info) + # get uint16 view of the underlying buffer. + res = gen_array_ops.bitcast(res, dtypes.uint16) + + # mask the last `shift` mantissa bits. + shift = self._num_mantissa_bits - self._num_valid_mantissa_bits + mask = (0xffff >> shift) << shift + res = gen_bitwise_ops.bitwise_and(res, mask) + + # restore float16 view. + return gen_array_ops.bitcast(res, dtype) + + +def _get_initializer(init_bound, dtype, seed): + if dtype == dtypes.float16: + return _MaskedRandomUniformInitializer( + -init_bound, init_bound, dtype=dtype, seed=seed) + else: + return init_ops.random_uniform_initializer( + -init_bound, init_bound, dtype=dtype, seed=seed) + + +def blocks_match(sess, use_peephole, dtype=dtypes.float32, cell_clip=None): batch_size = 2 input_size = 3 cell_size = 4 @@ -47,36 +113,42 @@ def blocks_match(sess, use_peephole): inputs = [] for _ in range(sequence_length): inp = ops.convert_to_tensor( - np.random.randn(batch_size, input_size), dtype=dtypes.float32) + np.random.randn(batch_size, input_size), dtype=dtype) inputs.append(inp) stacked_inputs = array_ops.stack(inputs) - initializer = init_ops.random_uniform_initializer(-0.01, 0.01, seed=19890212) + init_bound = 1e-1 if dtype == dtypes.float16 else 1e-2 + initializer = _get_initializer(init_bound, dtype=dtype, seed=19890212) with variable_scope.variable_scope("test", initializer=initializer): # magic naming so that the cells pick up these variables and reuse them if use_peephole: wci = variable_scope.get_variable( - "rnn/lstm_cell/w_i_diag", shape=[cell_size], dtype=dtypes.float32) + "rnn/lstm_cell/w_i_diag", shape=[cell_size], dtype=dtype) wcf = variable_scope.get_variable( - "rnn/lstm_cell/w_f_diag", shape=[cell_size], dtype=dtypes.float32) + "rnn/lstm_cell/w_f_diag", shape=[cell_size], dtype=dtype) wco = variable_scope.get_variable( - "rnn/lstm_cell/w_o_diag", shape=[cell_size], dtype=dtypes.float32) + "rnn/lstm_cell/w_o_diag", shape=[cell_size], dtype=dtype) w = variable_scope.get_variable( "rnn/lstm_cell/kernel", shape=[input_size + cell_size, cell_size * 4], - dtype=dtypes.float32) + dtype=dtype) b = variable_scope.get_variable( "rnn/lstm_cell/bias", shape=[cell_size * 4], - dtype=dtypes.float32, + dtype=dtype, initializer=init_ops.zeros_initializer()) basic_cell = rnn_cell.LSTMCell( - cell_size, use_peepholes=use_peephole, state_is_tuple=True, reuse=True) + cell_size, + use_peepholes=use_peephole, + cell_clip=cell_clip, + dtype=dtype, + state_is_tuple=True, + reuse=True) basic_outputs_op, basic_state_op = rnn.static_rnn( - basic_cell, inputs, dtype=dtypes.float32) + basic_cell, inputs, dtype=dtype) if use_peephole: _, _, _, _, _, _, block_outputs_op = block_lstm( @@ -87,7 +159,7 @@ def blocks_match(sess, use_peephole): wci=wci, wcf=wcf, wco=wco, - cell_clip=0, + cell_clip=cell_clip, use_peephole=True) else: _, _, _, _, _, _, block_outputs_op = block_lstm( @@ -95,13 +167,15 @@ def blocks_match(sess, use_peephole): inputs, w, b, - cell_clip=0) + cell_clip=cell_clip) fused_cell = lstm_ops.LSTMBlockFusedCell( - cell_size, cell_clip=0, use_peephole=use_peephole, reuse=True, + cell_size, + cell_clip=cell_clip, + use_peephole=use_peephole, + reuse=True, name="rnn/lstm_cell") - fused_outputs_op, fused_state_op = fused_cell( - stacked_inputs, dtype=dtypes.float32) + fused_outputs_op, fused_state_op = fused_cell(stacked_inputs, dtype=dtype) sess.run([variables.global_variables_initializer()]) basic_outputs, basic_state = sess.run([basic_outputs_op, basic_state_op[0]]) @@ -127,7 +201,19 @@ def blocks_match(sess, use_peephole): block_wgrads, fused_wgrads) -class LSTMBlockCellTest(test.TestCase): +class LSTMBlockCellTest(test.TestCase, parameterized.TestCase): + + TEST_CASES = ({ + "testcase_name": "Fp32", + "dtype": dtypes.float32, + "rtol": 1e-6, + "atol": 1e-6 + }, { + "testcase_name": "Fp16", + "dtype": dtypes.float16, + "rtol": 8e-3, + "atol": 8e-4 + }) def testNoneDimsWithDynamicRNN(self): with self.session(use_gpu=True, graph=ops.Graph()) as sess: @@ -314,41 +400,43 @@ class LSTMBlockCellTest(test.TestCase): for basic, block in zip(basic_res, block_res): self.assertAllClose(basic, block) - def testLSTMBasicToBlock(self): - with self.session(use_gpu=True) as sess: + def LSTMBasicToBlockTestHelper(self, + dtype=dtypes.float32, + use_peephole=False, + cell_clip=None, + rtol=1e-6, + atol=1e-6): + with self.session(use_gpu=True, graph=ops.Graph()) as sess: (basic_state, fused_state, basic_outputs, block_outputs, fused_outputs, basic_grads, block_grads, fused_grads, basic_wgrads, block_wgrads, fused_wgrads) = blocks_match( - sess, use_peephole=False) + sess, use_peephole=use_peephole, dtype=dtype, cell_clip=cell_clip) - self.assertAllClose(basic_outputs, block_outputs) - self.assertAllClose(basic_grads, block_grads) + self.assertAllClose(basic_outputs, block_outputs, rtol=rtol, atol=atol) + self.assertAllClose(basic_grads, block_grads, rtol=rtol, atol=atol) for basic, block in zip(basic_wgrads, block_wgrads): - self.assertAllClose(basic, block, rtol=1e-6, atol=1e-6) + self.assertAllClose(basic, block, rtol=rtol, atol=atol) - self.assertAllClose(basic_outputs, fused_outputs) - self.assertAllClose(basic_state, fused_state) - self.assertAllClose(basic_grads, fused_grads) - for basic, fused in zip(block_wgrads, fused_wgrads): - self.assertAllClose(basic, fused, rtol=1e-6, atol=1e-6) + self.assertAllClose(basic_outputs, fused_outputs, rtol=rtol, atol=atol) + self.assertAllClose(basic_state, fused_state, rtol=rtol, atol=atol) + self.assertAllClose(basic_grads, fused_grads, rtol=rtol, atol=atol) + for basic, fused in zip(basic_wgrads, fused_wgrads): + self.assertAllClose(basic, fused, rtol=rtol, atol=atol) - def testLSTMBasicToBlockPeeping(self): - with self.session(use_gpu=True) as sess: - (basic_state, fused_state, basic_outputs, block_outputs, fused_outputs, - basic_grads, block_grads, fused_grads, basic_wgrads, block_wgrads, - fused_wgrads) = blocks_match( - sess, use_peephole=True) + @parameterized.named_parameters(*TEST_CASES) + def testLSTMBasicToBlock(self, dtype, rtol, atol): + self.LSTMBasicToBlockTestHelper( + dtype, use_peephole=False, rtol=rtol, atol=atol) - self.assertAllClose(basic_outputs, block_outputs) - self.assertAllClose(basic_grads, block_grads) - for basic, block in zip(basic_wgrads, block_wgrads): - self.assertAllClose(basic, block, rtol=1e-6, atol=1e-6) + @parameterized.named_parameters(*TEST_CASES) + def testLSTMBasicToBlockPeeping(self, dtype, rtol, atol): + self.LSTMBasicToBlockTestHelper( + dtype, use_peephole=True, rtol=rtol, atol=atol) - self.assertAllClose(basic_outputs, fused_outputs) - self.assertAllClose(basic_state, fused_state) - self.assertAllClose(basic_grads, fused_grads) - for basic, fused in zip(block_wgrads, fused_wgrads): - self.assertAllClose(basic, fused, rtol=1e-6, atol=1e-6) + @parameterized.named_parameters(*TEST_CASES) + def testLSTMBasicToBlockCellClip(self, dtype, rtol, atol): + self.LSTMBasicToBlockTestHelper( + dtype, use_peephole=True, cell_clip=0.5, rtol=rtol, atol=atol) def testLSTMFusedSequenceLengths(self): """Verify proper support for sequence lengths in LSTMBlockFusedCell.""" @@ -444,16 +532,21 @@ class BenchmarkLSTMBlock(test.Benchmark): "batch_size": [1, 8, 13, 32, 67, 128], "cell_size": [128, 250, 512, 650, 1024, 1350], "time_steps": [40], - "use_gpu": [True, False] + "use_gpu": [True, False], + "dtype": ["float32", "float16"], }): + dtype = dtypes.float32 if config["dtype"] == "float32" else dtypes.float16 with ops.Graph().as_default(): with benchmarking.device(use_gpu=config["use_gpu"]): inputs = variable_scope.get_variable( "x", - [config["time_steps"], config["batch_size"], config["cell_size"]]) - cell = lstm_ops.LSTMBlockCell(config["cell_size"]) - outputs = rnn.dynamic_rnn( - cell, inputs, time_major=True, dtype=dtypes.float32) + dtype=dtype, + shape=[ + config["time_steps"], config["batch_size"], + config["cell_size"] + ]) + cell = lstm_ops.LSTMBlockCell(config["cell_size"], dtype=dtype) + outputs = rnn.dynamic_rnn(cell, inputs, time_major=True, dtype=dtype) init_op = variables.global_variables_initializer() with session.Session() as sess: @@ -464,12 +557,14 @@ class BenchmarkLSTMBlock(test.Benchmark): # is set, this will produce a copy-paste-able CSV file. print(",".join( map(str, [ - config["batch_size"], config["cell_size"], config["cell_size"], - config["time_steps"], config["use_gpu"], wall_time + config["dtype"], config["batch_size"], config["cell_size"], + config["cell_size"], config["time_steps"], config["use_gpu"], + wall_time ]))) benchmark_name_template = "_".join([ - "LSTMBlockCell_fprop", "BS%(batch_size)i", "CS%(cell_size)i", - "IS%(cell_size)i", "TS%(time_steps)i", "gpu_%(use_gpu)s" + "LSTMBlockCell_fprop", "DT_%(dtype)s", "BS%(batch_size)i", + "CS%(cell_size)i", "IS%(cell_size)i", "TS%(time_steps)i", + "gpu_%(use_gpu)s" ]) self.report_benchmark( @@ -488,8 +583,10 @@ class BenchmarkLSTMBlock(test.Benchmark): "batch_size": [1, 8, 13, 32, 67, 128], "cell_size": [128, 250, 512, 650, 1024, 1350], "time_steps": [40], - "use_gpu": [True, False] + "use_gpu": [True, False], + "dtype": ["float32", "float16"], }): + dtype = dtypes.float32 if config["dtype"] == "float32" else dtypes.float16 with ops.Graph().as_default(): with benchmarking.device(use_gpu=config["use_gpu"]): time_steps = config["time_steps"] @@ -498,21 +595,21 @@ class BenchmarkLSTMBlock(test.Benchmark): inputs = variable_scope.get_variable( "x", [time_steps, batch_size, cell_size], trainable=False, - dtype=dtypes.float32) + dtype=dtype) with variable_scope.variable_scope( "rnn", reuse=variable_scope.AUTO_REUSE): w = variable_scope.get_variable( "rnn/lstm_cell/kernel", shape=[input_size + cell_size, cell_size * 4], - dtype=dtypes.float32) + dtype=dtype) b = variable_scope.get_variable( "rnn/lstm_cell/bias", shape=[cell_size * 4], - dtype=dtypes.float32, + dtype=dtype, initializer=init_ops.zeros_initializer()) - cell = lstm_ops.LSTMBlockCell(cell_size) + cell = lstm_ops.LSTMBlockCell(cell_size, dtype=dtype) outputs = rnn.dynamic_rnn( - cell, inputs, time_major=True, dtype=dtypes.float32) + cell, inputs, time_major=True, dtype=dtype) grads = gradients_impl.gradients(outputs, [inputs, w, b]) init_op = variables.global_variables_initializer() @@ -524,12 +621,13 @@ class BenchmarkLSTMBlock(test.Benchmark): # is set, this will produce a copy-paste-able CSV file. print(",".join( map(str, [ - batch_size, cell_size, cell_size, time_steps, config["use_gpu"], - wall_time + config["dtype"], batch_size, cell_size, cell_size, time_steps, + config["use_gpu"], wall_time ]))) benchmark_name_template = "_".join([ - "LSTMBlockCell_bprop", "BS%(batch_size)i", "CS%(cell_size)i", - "IS%(cell_size)i", "TS%(time_steps)i", "gpu_%(use_gpu)s" + "LSTMBlockCell_bprop", "DT_%(dtype)s", "BS%(batch_size)i", + "CS%(cell_size)i", "IS%(cell_size)i", "TS%(time_steps)i", + "gpu_%(use_gpu)s" ]) self.report_benchmark( 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..d7ee7fb8faacb0876218a983d68f007e1905c11e 100644 --- a/tensorflow/contrib/rnn/python/kernel_tests/rnn_cell_test.py +++ b/tensorflow/contrib/rnn/python/kernel_tests/rnn_cell_test.py @@ -29,7 +29,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 @@ -763,6 +765,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/core_rnn_cell.py b/tensorflow/contrib/rnn/python/ops/core_rnn_cell.py index 645f82624bf67b96ffc8520289b293b45f0e69e2..6da58747e8905fae3b85b085991487d74f80b734 100644 --- a/tensorflow/contrib/rnn/python/ops/core_rnn_cell.py +++ b/tensorflow/contrib/rnn/python/ops/core_rnn_cell.py @@ -83,11 +83,11 @@ class _Linear(object): for shape in shapes: if shape.ndims != 2: raise ValueError("linear is expecting 2D arguments: %s" % shapes) - if shape[1].value is None: + if shape.dims[1].value is None: raise ValueError("linear expects shape[1] to be provided for shape %s, " "but saw %s" % (shape, shape[1])) else: - total_arg_size += shape[1].value + total_arg_size += shape.dims[1].value dtype = [a.dtype for a in args][0] @@ -156,11 +156,11 @@ def _linear(args, for shape in shapes: if shape.ndims != 2: raise ValueError("linear is expecting 2D arguments: %s" % shapes) - if shape[1].value is None: + if shape.dims[1].value is None: raise ValueError("linear expects shape[1] to be provided for shape %s, " "but saw %s" % (shape, shape[1])) else: - total_arg_size += shape[1].value + total_arg_size += shape.dims[1].value dtype = [a.dtype for a in args][0] diff --git a/tensorflow/contrib/rnn/python/ops/fused_rnn_cell.py b/tensorflow/contrib/rnn/python/ops/fused_rnn_cell.py index b7393d8b9880715cb381e1050b5ea757e36f2372..f90fd40990a32de18d5650dde9ff361ace77821e 100644 --- a/tensorflow/contrib/rnn/python/ops/fused_rnn_cell.py +++ b/tensorflow/contrib/rnn/python/ops/fused_rnn_cell.py @@ -20,10 +20,13 @@ from __future__ import print_function import abc +import six + from tensorflow.python.ops import array_ops from tensorflow.python.ops import rnn +@six.add_metaclass(abc.ABCMeta) class FusedRNNCell(object): """Abstract object representing a fused RNN cell. @@ -38,8 +41,6 @@ class FusedRNNCell(object): Every `FusedRNNCell` must implement `__call__` with the following signature. """ - __metaclass__ = abc.ABCMeta - @abc.abstractmethod def __call__(self, inputs, diff --git a/tensorflow/contrib/rnn/python/ops/gru_ops.py b/tensorflow/contrib/rnn/python/ops/gru_ops.py index 81ca12317be484ba420b7bbfac822e91d6d38bff..251a933eaec826b08266123245d9aef8573d3e06 100644 --- a/tensorflow/contrib/rnn/python/ops/gru_ops.py +++ b/tensorflow/contrib/rnn/python/ops/gru_ops.py @@ -20,7 +20,8 @@ from __future__ import print_function from tensorflow.contrib.rnn.ops import gen_gru_ops from tensorflow.contrib.util import loader from tensorflow.python.framework import ops -from tensorflow.python.layers import base as base_layer +from tensorflow.python.framework import tensor_shape +from tensorflow.python.keras.engine import input_spec from tensorflow.python.ops import array_ops from tensorflow.python.ops import init_ops from tensorflow.python.ops import math_ops @@ -164,7 +165,7 @@ class GRUBlockCell(LayerRNNCell): num_units = cell_size self._cell_size = num_units # Inputs must be 2-dimensional. - self.input_spec = base_layer.InputSpec(ndim=2) + self.input_spec = input_spec.InputSpec(ndim=2) @property def state_size(self): @@ -176,7 +177,7 @@ class GRUBlockCell(LayerRNNCell): def build(self, input_shape): # Check if the input size exist. - input_size = input_shape[1].value + input_size = tensor_shape.dimension_value(input_shape[1]) if input_size is None: raise ValueError("Expecting input_size to be set.") @@ -221,7 +222,7 @@ class GRUBlockCellV2(GRUBlockCell): def build(self, input_shape): """GRU cell.""" - input_size = input_shape[1].value + input_size = tensor_shape.dimension_value(input_shape[1]) if input_size is None: raise ValueError("Expecting input_size to be set.") diff --git a/tensorflow/contrib/rnn/python/ops/lstm_ops.py b/tensorflow/contrib/rnn/python/ops/lstm_ops.py index f2975b98061da45895481438aa34d2a6f6901a3a..b043026bc556a8879b15b432829baf8136250c0e 100644 --- a/tensorflow/contrib/rnn/python/ops/lstm_ops.py +++ b/tensorflow/contrib/rnn/python/ops/lstm_ops.py @@ -19,10 +19,13 @@ from __future__ import print_function import abc +import six + from tensorflow.contrib.rnn.ops import gen_lstm_ops from tensorflow.contrib.util import loader from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops +from tensorflow.python.keras.engine import input_spec from tensorflow.python.layers import base as base_layer from tensorflow.python.ops import array_ops from tensorflow.python.ops import init_ops @@ -113,7 +116,7 @@ def _lstm_block_cell(x, ValueError: If cell_size is None. """ if wci is None: - cell_size = cs_prev.get_shape().with_rank(2)[1].value + cell_size = cs_prev.get_shape().with_rank(2).dims[1].value if cell_size is None: raise ValueError("cell_size from `cs_prev` should not be None.") wci = array_ops.constant(0, dtype=dtypes.float32, shape=[cell_size]) @@ -154,7 +157,7 @@ def _block_lstm(seq_len_max, Args: seq_len_max: A `Tensor` of type `int64`. - x: A list of at least 1 `Tensor` objects of the same type in: `float32`. + x: A list of at least 1 `Tensor` objects of the same type. w: A `Tensor`. Must have the same type as `x`. b: A `Tensor`. Must have the same type as `x`. cs_prev: A `Tensor`. Must have the same type as `x`. @@ -187,21 +190,22 @@ def _block_lstm(seq_len_max, Raises: ValueError: If `b` does not have a valid shape. """ - batch_size = x[0].get_shape().with_rank(2)[0].value - cell_size4 = b.get_shape().with_rank(1)[0].value + dtype = x[0].dtype + batch_size = x[0].get_shape().with_rank(2).dims[0].value + cell_size4 = b.get_shape().with_rank(1).dims[0].value if cell_size4 is None: raise ValueError("`b` shape must not be None.") cell_size = cell_size4 / 4 zero_state = None if cs_prev is None or h_prev is None: zero_state = array_ops.constant( - 0, dtype=dtypes.float32, shape=[batch_size, cell_size]) + 0, dtype=dtype, shape=[batch_size, cell_size]) if cs_prev is None: cs_prev = zero_state if h_prev is None: h_prev = zero_state if wci is None: - wci = array_ops.constant(0, dtype=dtypes.float32, shape=[cell_size]) + wci = array_ops.constant(0, dtype=dtype, shape=[cell_size]) wcf = wci wco = wci @@ -238,13 +242,13 @@ def _LSTMBlockCellGrad(op, *grad): (i, cs, f, o, ci, co, _) = op.outputs (_, cs_grad, _, _, _, _, h_grad) = grad - batch_size = x.get_shape().with_rank(2)[0].value + batch_size = x.get_shape().with_rank(2).dims[0].value if batch_size is None: batch_size = -1 - input_size = x.get_shape().with_rank(2)[1].value + input_size = x.get_shape().with_rank(2).dims[1].value if input_size is None: raise ValueError("input_size from `x` should not be None.") - cell_size = cs_prev.get_shape().with_rank(2)[1].value + cell_size = cs_prev.get_shape().with_rank(2).dims[1].value if cell_size is None: raise ValueError("cell_size from `cs_prev` should not be None.") @@ -382,7 +386,7 @@ class LSTMBlockCell(LayerRNNCell): "scope": "lstm_cell" } # Inputs must be 2-dimensional. - self.input_spec = base_layer.InputSpec(ndim=2) + self.input_spec = input_spec.InputSpec(ndim=2) @property def state_size(self): @@ -393,10 +397,10 @@ class LSTMBlockCell(LayerRNNCell): return self._num_units def build(self, inputs_shape): - if not inputs_shape[1].value: + if not inputs_shape.dims[1].value: raise ValueError( "Expecting inputs_shape[1] to be set: %s" % str(inputs_shape)) - input_size = inputs_shape[1].value + input_size = inputs_shape.dims[1].value self._kernel = self.add_variable( self._names["W"], [input_size + self._num_units, self._num_units * 4]) self._bias = self.add_variable( @@ -439,6 +443,7 @@ class LSTMBlockCell(LayerRNNCell): return h, new_state +@six.add_metaclass(abc.ABCMeta) class LSTMBlockWrapper(base_layer.Layer): """This is a helper class that provides housekeeping for LSTM cells. @@ -511,10 +516,10 @@ class LSTMBlockWrapper(base_layer.Layer): inputs_shape = inputs.get_shape().with_rank(3) if not inputs_shape[2]: raise ValueError("Expecting inputs_shape[2] to be set: %s" % inputs_shape) - batch_size = inputs_shape[1].value + batch_size = inputs_shape.dims[1].value if batch_size is None: batch_size = array_ops.shape(inputs)[1] - time_len = inputs_shape[0].value + time_len = inputs_shape.dims[0].value if time_len is None: time_len = array_ops.shape(inputs)[0] @@ -624,7 +629,7 @@ class LSTMBlockFusedCell(LSTMBlockWrapper): self._use_peephole = use_peephole # Inputs must be 3-dimensional. - self.input_spec = base_layer.InputSpec(ndim=3) + self.input_spec = input_spec.InputSpec(ndim=3) @property def num_units(self): @@ -632,7 +637,7 @@ class LSTMBlockFusedCell(LSTMBlockWrapper): return self._num_units def build(self, input_shape): - input_size = input_shape[2].value + input_size = input_shape.dims[2].value self._kernel = self.add_variable( "kernel", [input_size + self._num_units, self._num_units * 4]) self._bias = self.add_variable( @@ -674,7 +679,7 @@ class LSTMBlockFusedCell(LSTMBlockWrapper): """ inputs_shape = inputs.get_shape().with_rank(3) - time_len = inputs_shape[0].value + time_len = inputs_shape.dims[0].value if time_len is None: time_len = array_ops.shape(inputs)[0] diff --git a/tensorflow/contrib/rnn/python/ops/rnn_cell.py b/tensorflow/contrib/rnn/python/ops/rnn_cell.py index 0f693e915415598c6b6327e63a10deb60a7f6d27..482e547a16be85804beec88a91fa03b053d09b27 100644 --- a/tensorflow/contrib/rnn/python/ops/rnn_cell.py +++ b/tensorflow/contrib/rnn/python/ops/rnn_cell.py @@ -30,7 +30,7 @@ from tensorflow.python.framework import ops from tensorflow.python.framework import tensor_shape from tensorflow.python.keras import activations from tensorflow.python.keras import initializers -from tensorflow.python.layers import base as base_layer +from tensorflow.python.keras.engine import input_spec from tensorflow.python.ops import array_ops from tensorflow.python.ops import clip_ops from tensorflow.python.ops import gen_array_ops @@ -251,11 +251,13 @@ class CoupledInputForgetGateLSTMCell(rnn_cell_impl.RNNCell): m_prev = array_ops.slice(state, [0, self._num_units], [-1, num_proj]) dtype = inputs.dtype - input_size = inputs.get_shape().with_rank(2)[1] + 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]") concat_w = _get_concat_variable( - "W", [input_size.value + num_proj, 3 * self._num_units], dtype, + "W", + [input_size.value + num_proj, 3 * self._num_units], + dtype, self._num_unit_shards) b = vs.get_variable( @@ -429,7 +431,7 @@ class TimeFreqLSTMCell(rnn_cell_impl.RNNCell): # initialize the first freq state to be zero m_prev_freq = array_ops.zeros( - [inputs.shape[0].value or inputs.get_shape()[0], self._num_units], + [inputs.shape.dims[0].value or inputs.get_shape()[0], self._num_units], dtype) for fq in range(len(freq_inputs)): c_prev = array_ops.slice(state, [0, 2 * fq * self._num_units], @@ -480,7 +482,7 @@ class TimeFreqLSTMCell(rnn_cell_impl.RNNCell): Raises: ValueError: if input_size cannot be inferred from static shape inference. """ - input_size = input_feat.get_shape().with_rank(2)[-1].value + input_size = input_feat.get_shape().with_rank(2).dims[-1].value if input_size is None: raise ValueError("Cannot infer input_size from static shape inference.") num_feats = int( @@ -636,7 +638,8 @@ class GridLSTMCell(rnn_cell_impl.RNNCell): ValueError: if an input_size was specified and the provided inputs have a different dimension. """ - batch_size = inputs.shape[0].value or array_ops.shape(inputs)[0] + batch_size = tensor_shape.dimension_value( + inputs.shape[0]) or array_ops.shape(inputs)[0] freq_inputs = self._make_tf_features(inputs) m_out_lst = [] state_out_lst = [] @@ -886,7 +889,7 @@ class GridLSTMCell(rnn_cell_impl.RNNCell): Raises: ValueError: if input_size cannot be inferred from static shape inference. """ - input_size = input_feat.get_shape().with_rank(2)[-1].value + input_size = input_feat.get_shape().with_rank(2).dims[-1].value if input_size is None: raise ValueError("Cannot infer input_size from static shape inference.") if slice_offset > 0: @@ -910,7 +913,7 @@ class GridLSTMCell(rnn_cell_impl.RNNCell): if not self._start_freqindex_list: if len(self._num_frequency_blocks) != 1: raise ValueError("Length of num_frequency_blocks" - " is not 1, but instead is %d", + " is not 1, but instead is %d" % len(self._num_frequency_blocks)) num_feats = int( (input_size - self._feature_size) / (self._frequency_skip)) + 1 @@ -1058,7 +1061,8 @@ class BidirectionalGridLSTMCell(GridLSTMCell): ValueError: if an input_size was specified and the provided inputs have a different dimension. """ - batch_size = inputs.shape[0].value or array_ops.shape(inputs)[0] + batch_size = tensor_shape.dimension_value( + inputs.shape[0]) or array_ops.shape(inputs)[0] fwd_inputs = self._make_tf_features(inputs) if self._backward_slice_offset: bwd_inputs = self._make_tf_features(inputs, self._backward_slice_offset) @@ -1289,7 +1293,7 @@ class HighwayWrapper(rnn_cell_impl.RNNCell): return self._cell.zero_state(batch_size, dtype) def _highway(self, inp, out): - input_size = inp.get_shape().with_rank(2)[1].value + input_size = inp.get_shape().with_rank(2).dims[1].value carry_weight = vs.get_variable("carry_w", [input_size, input_size]) carry_bias = vs.get_variable( "carry_b", [input_size], @@ -1458,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: @@ -1471,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: @@ -1505,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. @@ -1531,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)[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]) @@ -1594,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 @@ -1674,7 +1690,7 @@ class UGRNNCell(rnn_cell_impl.RNNCell): """ sigmoid = math_ops.sigmoid - input_size = inputs.get_shape().with_rank(2)[1] + 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]") @@ -1785,7 +1801,7 @@ class IntersectionRNNCell(rnn_cell_impl.RNNCell): sigmoid = math_ops.sigmoid tanh = math_ops.tanh - input_size = inputs.get_shape().with_rank(2)[1] + 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]") @@ -2067,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`. @@ -2088,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 @@ -2168,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. @@ -2362,11 +2378,12 @@ class GLSTMCell(rnn_cell_impl.RNNCell): """ (c_prev, m_prev) = state - self._batch_size = inputs.shape[0].value or array_ops.shape(inputs)[0] + self._batch_size = tensor_shape.dimension_value( + inputs.shape[0]) or array_ops.shape(inputs)[0] # If the input size is statically-known, calculate and validate its group # size. Otherwise, use the output group size. - input_size = inputs.shape[1].value + input_size = tensor_shape.dimension_value(inputs.shape[1]) if input_size is None: raise ValueError("input size must be statically known") if input_size % self._number_of_groups != 0: @@ -2587,11 +2604,11 @@ class LayerNormLSTMCell(rnn_cell_impl.RNNCell): for shape in shapes: if shape.ndims != 2: raise ValueError("linear is expecting 2D arguments: %s" % shapes) - if shape[1].value is None: + if tensor_shape.dimension_value(shape[1]) is None: raise ValueError("linear expects shape[1] to be provided for shape %s, " "but saw %s" % (shape, shape[1])) else: - total_arg_size += shape[1].value + total_arg_size += tensor_shape.dimension_value(shape[1]) dtype = [a.dtype for a in args][0] @@ -2649,7 +2666,7 @@ class LayerNormLSTMCell(rnn_cell_impl.RNNCell): (c_prev, m_prev) = state dtype = inputs.dtype - input_size = inputs.get_shape().with_rank(2)[1] + 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]") scope = vs.get_variable_scope() @@ -2739,15 +2756,17 @@ 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 # Restrict inputs to be 2-dimensional matrices - self.input_spec = base_layer.InputSpec(ndim=2) + self.input_spec = input_spec.InputSpec(ndim=2) @property def state_size(self): @@ -2758,11 +2777,11 @@ class SRUCell(rnn_cell_impl.LayerRNNCell): return self._num_units def build(self, inputs_shape): - if inputs_shape[1].value is None: + if tensor_shape.dimension_value(inputs_shape[1]) is None: raise ValueError( "Expected inputs.shape[-1] to be known, saw shape: %s" % inputs_shape) - input_depth = inputs_shape[1].value + input_depth = tensor_shape.dimension_value(inputs_shape[1]) # pylint: disable=protected-access self._kernel = self.add_variable( @@ -2772,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 @@ -2935,11 +2954,11 @@ class WeightNormLSTMCell(rnn_cell_impl.RNNCell): for shape in shapes: if shape.ndims != 2: raise ValueError("linear is expecting 2D arguments: %s" % shapes) - if shape[1].value is None: + if tensor_shape.dimension_value(shape[1]) is None: raise ValueError("linear expects shape[1] to be provided for shape %s, " "but saw %s" % (shape, shape[1])) else: - total_arg_size += shape[1].value + total_arg_size += tensor_shape.dimension_value(shape[1]) dtype = [a.dtype for a in args][0] @@ -2955,7 +2974,7 @@ class WeightNormLSTMCell(rnn_cell_impl.RNNCell): st = 0 with ops.control_dependencies(None): for i in range(len(args)): - en = st + shapes[i][1].value + en = st + tensor_shape.dimension_value(shapes[i][1]) wn.append( self._normalize(weights[st:en, :], name="norm_{}".format(i))) st = en @@ -3009,7 +3028,7 @@ class WeightNormLSTMCell(rnn_cell_impl.RNNCell): sigmoid = math_ops.sigmoid c, h = state - input_size = inputs.get_shape().with_rank(2)[1] + 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]") @@ -3084,7 +3103,7 @@ class IndRNNCell(rnn_cell_impl.LayerRNNCell): super(IndRNNCell, self).__init__(_reuse=reuse, name=name, dtype=dtype) # Inputs must be 2-dimensional. - self.input_spec = base_layer.InputSpec(ndim=2) + self.input_spec = input_spec.InputSpec(ndim=2) self._num_units = num_units self._activation = activation or math_ops.tanh @@ -3098,11 +3117,11 @@ class IndRNNCell(rnn_cell_impl.LayerRNNCell): return self._num_units def build(self, inputs_shape): - if inputs_shape[1].value is None: + if tensor_shape.dimension_value(inputs_shape[1]) is None: raise ValueError( "Expected inputs.shape[-1] to be known, saw shape: %s" % inputs_shape) - input_depth = inputs_shape[1].value + input_depth = tensor_shape.dimension_value(inputs_shape[1]) # pylint: disable=protected-access self._kernel_w = self.add_variable( "%s_w" % rnn_cell_impl._WEIGHTS_VARIABLE_NAME, @@ -3178,7 +3197,7 @@ class IndyGRUCell(rnn_cell_impl.LayerRNNCell): super(IndyGRUCell, self).__init__(_reuse=reuse, name=name, dtype=dtype) # Inputs must be 2-dimensional. - self.input_spec = base_layer.InputSpec(ndim=2) + self.input_spec = input_spec.InputSpec(ndim=2) self._num_units = num_units self._activation = activation or math_ops.tanh @@ -3194,11 +3213,11 @@ class IndyGRUCell(rnn_cell_impl.LayerRNNCell): return self._num_units def build(self, inputs_shape): - if inputs_shape[1].value is None: + if tensor_shape.dimension_value(inputs_shape[1]) is None: raise ValueError( "Expected inputs.shape[-1] to be known, saw shape: %s" % inputs_shape) - input_depth = inputs_shape[1].value + input_depth = tensor_shape.dimension_value(inputs_shape[1]) # pylint: disable=protected-access self._gate_kernel_w = self.add_variable( "gates/%s_w" % rnn_cell_impl._WEIGHTS_VARIABLE_NAME, @@ -3318,7 +3337,7 @@ class IndyLSTMCell(rnn_cell_impl.LayerRNNCell): super(IndyLSTMCell, self).__init__(_reuse=reuse, name=name, dtype=dtype) # Inputs must be 2-dimensional. - self.input_spec = base_layer.InputSpec(ndim=2) + self.input_spec = input_spec.InputSpec(ndim=2) self._num_units = num_units self._forget_bias = forget_bias @@ -3335,11 +3354,11 @@ class IndyLSTMCell(rnn_cell_impl.LayerRNNCell): return self._num_units def build(self, inputs_shape): - if inputs_shape[1].value is None: + if tensor_shape.dimension_value(inputs_shape[1]) is None: raise ValueError( "Expected inputs.shape[-1] to be known, saw shape: %s" % inputs_shape) - input_depth = inputs_shape[1].value + input_depth = tensor_shape.dimension_value(inputs_shape[1]) # pylint: disable=protected-access self._kernel_w = self.add_variable( "%s_w" % rnn_cell_impl._WEIGHTS_VARIABLE_NAME, @@ -3439,7 +3458,7 @@ class MinimalRNNCell(rnn_cell_impl.LayerRNNCell): super(MinimalRNNCell, self).__init__(name=name, dtype=dtype, **kwargs) # Inputs must be 2-dimensional. - self.input_spec = base_layer.InputSpec(ndim=2) + self.input_spec = input_spec.InputSpec(ndim=2) self.units = units self.activation = activations.get(activation) @@ -3492,12 +3511,13 @@ class MinimalRNNCell(rnn_cell_impl.LayerRNNCell): static shape inference. """ input_size = inputs.get_shape()[1] - if input_size.value is None: + if tensor_shape.dimension_value(input_size) is None: raise ValueError("Could not infer input size from inputs.get_shape()[-1]") feedforward_weight, gate_weight = array_ops.split( value=self.kernel, - num_or_size_splits=[input_size.value, 2 * self.units], + num_or_size_splits=[tensor_shape.dimension_value(input_size), + 2 * self.units], axis=0) feedforward = math_ops.matmul(inputs, feedforward_weight) @@ -3552,7 +3572,7 @@ class CFNCell(rnn_cell_impl.LayerRNNCell): super(CFNCell, self).__init__(name=name, dtype=dtype, **kwargs) # Inputs must be 2-dimensional. - self.input_spec = base_layer.InputSpec(ndim=2) + self.input_spec = input_spec.InputSpec(ndim=2) self.units = units self.activation = activations.get(activation) @@ -3611,7 +3631,7 @@ class CFNCell(rnn_cell_impl.LayerRNNCell): static shape inference. """ input_size = inputs.get_shape()[-1] - if input_size.value is None: + if tensor_shape.dimension_value(input_size) is None: raise ValueError("Could not infer input size from inputs.get_shape()[-1]") # The variable names u, v, w, b are consistent with the notations in the 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/rpc/python/kernel_tests/rpc_op_test_base.py b/tensorflow/contrib/rpc/python/kernel_tests/rpc_op_test_base.py index 0d615923e04915a8429252317025ac8e79f9bb4e..d6148715be91c78e6e5a99fc0f3caa905b5c1a7d 100644 --- a/tensorflow/contrib/rpc/python/kernel_tests/rpc_op_test_base.py +++ b/tensorflow/contrib/rpc/python/kernel_tests/rpc_op_test_base.py @@ -176,7 +176,9 @@ class RpcOpTestBase(object): expected_message_values = np.where( status_code_values == errors.INVALID_ARGUMENT, I_WARNED_YOU.encode('ascii'), b'') - self.assertAllEqual(expected_message_values, status_message_values) + for msg, expected in zip(status_message_values, expected_message_values): + self.assertTrue(expected in msg, + '"%s" did not contain "%s"' % (msg, expected)) def testVecHostPortRpc(self): with self.cached_session() as sess: diff --git a/tensorflow/contrib/saved_model/BUILD b/tensorflow/contrib/saved_model/BUILD index 291ff83791c7cded2dccc4719bb12e84f00afa42..f0242a3b40fd566ec0f477d462426d5f550d1620 100644 --- a/tensorflow/contrib/saved_model/BUILD +++ b/tensorflow/contrib/saved_model/BUILD @@ -82,35 +82,8 @@ py_library( name = "keras_saved_model", srcs = ["python/saved_model/keras_saved_model.py"], srcs_version = "PY2AND3", - tags = ["no_windows"], 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 = ["notsan"], - 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 2c5c8c4afdc5778e3bb182d0a492d20e758baf14..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,297 +18,9 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import os +from tensorflow.python.keras import saving -from tensorflow.python.client import session -from tensorflow.python.estimator import keras as estimator_keras_util -from tensorflow.python.estimator import model_fn as model_fn_lib -from tensorflow.python.estimator.export import export as export_helpers -from tensorflow.python.framework import errors -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.models import model_from_json -from tensorflow.python.lib.io import file_io -from tensorflow.python.ops import variables -from tensorflow.python.platform import gfile -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 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 - -def save_keras_model( - model, saved_model_path, custom_objects=None, as_text=None): - """Save a `tf.keras.Model` into Tensorflow SavedModel format. - - `save_model` generates new files/folders under the `saved_model_path` folder: - 1) an asset folder containing the json string of the model's - configuration (topology). - 2) a checkpoint containing the model weights. - 3) 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. - - Model Requirements: - - Model must be a sequential model or functional model. Subclassed models can - not be saved via this function, unless you provide an implementation for - get_config() and from_config(). - - All variables must be saveable by the model. In general, this condition is - met through the use of layers defined in the keras library. However, - there is currently a bug with variables created in Lambda layer functions - not being saved correctly (see - https://github.com/keras-team/keras/issues/9740). - - 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. - - Args: - model: A `tf.keras.Model` to be saved. - 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. - - Returns: - String path to the SavedModel folder, a subdirectory of `saved_model_path`. - - Raises: - NotImplementedError: If the passed in model is a subclassed model. - """ - if not model._is_graph_network: - raise NotImplementedError - - export_dir = export_helpers.get_timestamped_export_dir(saved_model_path) - temp_export_dir = export_helpers.get_temp_export_dir(export_dir) - - builder = saved_model_builder.SavedModelBuilder(temp_export_dir) - - # 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_json_and_variables(model, temp_export_dir) - - # 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} - - has_saved_vars = False - if model.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) - - gfile.Rename(temp_export_dir, export_dir) - return export_dir - - -def _export_model_json_and_variables(model, saved_model_path): - """Save model variables and json structure into SavedModel subdirectories.""" - # Save 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) - - # Save 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 _get_var_list(model): - """Return list of all checkpointed saveable objects in the model.""" - return checkpointable_utils.named_saveables(model) - - -def _export_mode( - mode, has_saved_vars, builder, model, custom_objects, checkpoint_path): - """Export a model, and optionally save 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. - - 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) - - # Clone the model into blank graph. This will create placeholders for inputs - # and targets. - clone = models_lib.clone_and_build_model( - model, 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. - if mode == model_fn_lib.ModeKeys.TRAIN: - clone._make_train_function() - builder._add_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), - main_op=variables.local_variables_initializer()) - return None - - -def _create_signature_def_map(model, mode): - """Create 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)} - export_outputs = model_fn_lib.export_outputs_for_mode( - mode, - predictions=outputs_dict, - loss=model.total_loss if model.optimizer else None, - metrics=estimator_keras_util._convert_keras_metrics_to_estimator(model)) - 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): - """Assert model and clone contain the same checkpointable objects.""" - - def get_non_optimizer_objects(m, g): - """Gather set of model and optimizer checkpointable objects.""" - # Set default graph because optimizer.variables() returns optimizer - # variables defined in the default graph. - with g.as_default(): - all_objects = set(checkpointable_utils.list_objects(m)) - optimizer_and_variables = set() - for obj in all_objects: - if isinstance(obj, optimizers.TFOptimizer): - optimizer_and_variables.update(checkpointable_utils.list_objects(obj)) - optimizer_and_variables.update(set(obj.optimizer.variables())) - return all_objects - optimizer_and_variables - - model_objects = get_non_optimizer_objects(model, model_graph) - clone_objects = get_non_optimizer_objects(clone, clone_graph) - - if len(model_objects) != len(clone_objects): - raise errors.InternalError( - None, None, - 'Model and clone must use the same variables.' - '\n\tModel variables: %s\n\t Clone variables: %s' - % (model_objects, clone_objects)) - - -def load_keras_model(saved_model_path): - """Load 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. - - 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 060c5045235ced50adf38222a0152a1700a252e8..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/saved_model/python/saved_model/keras_saved_model_test.py +++ /dev/null @@ -1,430 +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 errors -from tensorflow.python.framework import ops -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 constants -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) - 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): - 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): - 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 load_model(sess, path, mode): - tags = model_fn_lib.EXPORT_TAG_MAP[mode] - sig_def_key = (signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY - if mode == model_fn_lib.ModeKeys.PREDICT else 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 - - -@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( - (functional_model, True, training_module.AdadeltaOptimizer(), True), - (functional_model, True, training_module.AdadeltaOptimizer(), False), - (functional_model, False, None, False), - (sequential_model, True, training_module.AdadeltaOptimizer(), True), - (sequential_model, True, training_module.AdadeltaOptimizer(), False), - (sequential_model, False, None, 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()): - 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) - - eval_results = sess.run(outputs, {inputs[input_name]: input_arr, - inputs[target_name]: target_arr}) - - self.assertEqual(int(train_before_export), - sess.run(training_module.get_global_step())) - self.assertAllClose(ref_loss, eval_results['loss'], atol=1e-05) - self.assertAllClose( - ref_mae, eval_results['metrics/mae/update_op'], atol=1e-05) - self.assertAllClose( - ref_predict, eval_results['predictions/' + output_name], 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 = 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/mae/update_op', outputs) - self.assertIn('metrics/mae/value', outputs) - self.assertIn('predictions/' + output_name, outputs) - - # Train for a step - train_op = ops.get_collection(constants.TRAIN_OP_KEY) - 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) - - with self.assertRaisesRegexp( - errors.InternalError, 'Model and clone must use the same variables.'): - keras_saved_model._assert_same_non_optimizer_objects( - model, model_graph, clone, clone_graph) - - -if __name__ == '__main__': - test.main() diff --git a/tensorflow/contrib/seq2seq/BUILD b/tensorflow/contrib/seq2seq/BUILD index 18b56cd21942e28cb0dc3210df0bb04d55c1e16f..7d5ba90ded215a59dbded751efd497f142a95e61 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", ], @@ -215,3 +213,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 c1e36b2ea3677f742f7c699b616def0c0147e063..d815f81f847ad79ddcc6c6ecf5c050598e185d8d 100644 --- a/tensorflow/contrib/seq2seq/python/kernel_tests/attention_wrapper_test.py +++ b/tensorflow/contrib/seq2seq/python/kernel_tests/attention_wrapper_test.py @@ -35,6 +35,7 @@ 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 random_ops +from tensorflow.python.ops import rnn from tensorflow.python.ops import rnn_cell from tensorflow.python.ops import variables from tensorflow.python.ops import variable_scope as vs @@ -154,13 +155,13 @@ class AttentionWrapperTest(test.TestCase): 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]) + 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])[-1].value + [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) @@ -992,5 +993,67 @@ class AttentionWrapperTest(test.TestCase): expected_final_alignment_history=expected_final_alignment_history, name='testMultiAttention') + def testCustomizedAttention(self): + batch_size = 2 + max_time = 3 + num_units = 2 + memory = constant_op.constant([[[1., 1.], [2., 2.], [3., 3.]], + [[4., 4.], [5., 5.], [6., 6.]]]) + memory_sequence_length = constant_op.constant([3, 2]) + attention_mechanism = wrapper.BahdanauAttention(num_units, memory, + memory_sequence_length) + + # Sets all returned values to be all ones. + def _customized_attention(unused_attention_mechanism, unused_cell_output, + unused_attention_state, unused_attention_layer): + """Customized attention. + + Returns: + attention: `Tensor` of shape [batch_size, num_units], attention output. + alignments: `Tensor` of shape [batch_size, max_time], sigma value for + each input memory (prob. function of input keys). + next_attention_state: A `Tensor` representing the next state for the + attention. + """ + attention = array_ops.ones([batch_size, num_units]) + alignments = array_ops.ones([batch_size, max_time]) + next_attention_state = alignments + return attention, alignments, next_attention_state + + attention_cell = wrapper.AttentionWrapper( + rnn_cell.LSTMCell(2), + attention_mechanism, + attention_layer_size=None, # don't use attention layer. + output_attention=False, + alignment_history=(), + attention_fn=_customized_attention, + name='attention') + self.assertEqual(num_units, attention_cell.output_size) + + initial_state = attention_cell.zero_state( + batch_size=2, dtype=dtypes.float32) + source_input_emb = array_ops.ones([2, 3, 2]) + source_input_length = constant_op.constant([3, 2]) + + # 'state' is a tuple of + # (cell_state, h, attention, alignments, alignment_history, attention_state) + output, state = rnn.dynamic_rnn( + attention_cell, + inputs=source_input_emb, + sequence_length=source_input_length, + initial_state=initial_state, + dtype=dtypes.float32) + + with self.session() as sess: + sess.run(variables.global_variables_initializer()) + output_value, state_value = sess.run([output, state], feed_dict={}) + self.assertAllEqual(np.array([2, 3, 2]), output_value.shape) + self.assertAllClose(np.array([[1., 1.], [1., 1.]]), state_value.attention) + self.assertAllClose( + np.array([[1., 1., 1.], [1., 1., 1.]]), state_value.alignments) + self.assertAllClose( + np.array([[1., 1., 1.], [1., 1., 1.]]), state_value.attention_state) + + if __name__ == '__main__': test.main() 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..7ff04e1780c4c44df14d6e87c5afdbf533ca5c90 --- /dev/null +++ b/tensorflow/contrib/seq2seq/python/kernel_tests/attention_wrapper_v2_test.py @@ -0,0 +1,94 @@ +# 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 + +from absl.testing import parameterized +import numpy as np + +from tensorflow.contrib.seq2seq.python.ops import attention_wrapper as wrapper +from tensorflow.python.framework import ops +from tensorflow.python.framework import test_util +from tensorflow.python.ops import variables +from tensorflow.python.platform import test + + +@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 = ops.convert_to_tensor( + np.random.random((self.batch, self.timestep, self.memory_size)), + dtype=np.float32) + self.query = ops.convert_to_tensor( + np.random.random((self.batch, self.units)), dtype=np.float32) + self.state = ops.convert_to_tensor( + np.random.random((self.batch, self.timestep)), dtype=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) + attention_score = attention([self.query, self.state, self.memory]) + 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) + 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) + + score = attention([self.query, self.state, self.memory]) + 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)) + +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/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 0ba32cd3bf8a374f5f55bdc6b2325b03443cd545..ae3e7f1b5d8c9f06b5defbaee9cad3810e58abd4 100644 --- a/tensorflow/contrib/seq2seq/python/ops/attention_wrapper.py +++ b/tensorflow/contrib/seq2seq/python/ops/attention_wrapper.py @@ -28,6 +28,7 @@ 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 layers 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,76 +73,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 = ( - memory_sequence_length.shape[0].value - 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 = m.shape[0].value 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. @@ -204,20 +135,23 @@ 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 else self._values) self._batch_size = ( - self._keys.shape[0].value or array_ops.shape(self._keys)[0]) - self._alignments_size = (self._keys.shape[1].value or - array_ops.shape(self._keys)[1]) + 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]) @property def memory_layer(self): @@ -284,6 +218,207 @@ 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 only support Keras functional API since it takes multiple + input tensors, which is not available in sequential model. + """ + + def __init__(self, + probability_fn, + query_layer=None, + memory_layer=None, + **kwargs): + """Construct base AttentionMechanism class. + + Args: + 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`. + **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 + + def build(self, input_shape): + # The layer suppose to take 3 inputs, [query, state, memory]. + query_input_shape, _, memory_input_shape = input_shape + if self.query_layer is not None: + self.query_layer.build(query_input_shape) + if self.memory_layer is not None: + self.memory_layer.build(memory_input_shape) + # dtype of the layer is known at this moment, create the score_mask_value if + # needed. + self.score_mask_value = dtypes.as_dtype(self.dtype).as_numpy_dtype(-np.inf) + self.built = True + + def _setup_memory(self, memory, 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_mask: 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.") + with ops.name_scope( + self.name, "BaseAttentionMechanismInit", nest.flatten(memory)): + self.values = _prepare_memory( + memory, memory_mask=memory_mask, + check_inner_dims_defined=self._check_inner_dims_defined) + 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: + self.probability_fn = lambda score, prev: ( # pylint:disable=g-long-lambda + self.probability_fn(_maybe_mask_score( + score, self.score_mask_value, memory_mask=memory_mask), prev)) + self._memory_initialized = True + + def call(self, inputs, mask=None, **kwargs): + """Base method to calculate the attention score. + + Args: + inputs: a list of tensor that contains `query`, `state`, and `memory`. + `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. + This tensor should be shaped `[batch_size, max_time, feature]`. + 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. + **kwargs: Dict, other keyword arguments for the call method. + """ + query, state, memory, memory_mask = self._process_inputs(inputs, mask) + if not self._memory_initialized: + self._setup_memory(memory, memory_mask=memory_mask) + return self.calculate_attention(query, state) + + def calculate_attention(self, query, state): + raise NotImplementedError( + "calculate_attention need to be implemented by subclasses.") + + 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(), + } + base_config = super(_BaseAttentionMechanismV2, self).get_config() + return dict(list(base_config.items()) + list(config.items())) + + def _process_inputs(self, inputs, mask): + if len(inputs) != 3: + raise ValueError( + "Expect to have 3 inputs for attention, got %d" % len(inputs)) + query, state, memory = inputs + return query, state, memory, mask + + 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 + + def _luong_score(query, keys, scale): """Implements Luong-style (multiplicative) scoring function. @@ -302,7 +437,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. @@ -318,7 +453,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. @@ -336,12 +470,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 @@ -352,8 +482,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. @@ -427,13 +557,125 @@ 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, + scale=False, + probability_fn="softmax", + dtype=None, + name="LuongAttention", + **kwargs): + """Construct the AttentionMechanism mechanism. + + Args: + units: The depth of the attention mechanism. + 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) + super(LuongAttentionV2, self).__init__( + query_layer=None, + memory_layer=memory_layer, + probability_fn=wrapped_probability_fn, + name=name, + dtype=dtype, + **kwargs) + self.units = units + self.scale = scale + + def build(self, input_shape): + super(LuongAttentionV2, self).build(input_shape) + if self.scale: + self.scale_weight = self.add_weight( + "attention_g", initializer=init_ops.ones_initializer, shape=()) + else: + self.scale_weight = None + 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, @@ -451,40 +693,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 = keys.shape[2].value 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): @@ -575,12 +805,152 @@ 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, + normalize=False, + probability_fn="softmax", + dtype=None, + name="BahdanauAttention", + **kwargs): + """Construct the Attention mechanism. + + Args: + units: The depth of the query mechanism. + 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`. + 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) + super(BahdanauAttentionV2, self).__init__( + query_layer=query_layer, + memory_layer=memory_layer, + probability_fn=wrapped_probability_fn, + name=name, + dtype=dtype, + **kwargs) + self.units = units + self.normalize = normalize + + def build(self, input_shape): + super(BahdanauAttentionV2, self).build(input_shape) + self.attention_v = self.add_weight( + "attention_v", [self.units], dtype=self.dtype) + if self.normalize: + 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()) + else: + self.attention_g = None + self.attention_b = None + 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, + } + 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. @@ -649,8 +1019,9 @@ def monotonic_attention(p_choose_i, previous_attention, mode): previous_attention = ops.convert_to_tensor( previous_attention, name="previous_attention") if mode == "recursive": - # Use .shape[0].value when it's not None, or fall back on symbolic shape - batch_size = p_choose_i.shape[0].value or array_ops.shape(p_choose_i)[0] + # Use .shape[0] when it's not None, or fall back on symbolic shape + batch_size = tensor_shape.dimension_value( + p_choose_i.shape[0]) or array_ops.shape(p_choose_i)[0] # Compute [1, 1 - p_choose_i[0], 1 - p_choose_i[1], ..., 1 - p_choose_i[-2]] shifted_1mp_choose_i = array_ops.concat( [array_ops.ones((batch_size, 1)), 1 - p_choose_i[:, :-1]], 1) @@ -762,6 +1133,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. @@ -856,7 +1255,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) @@ -866,6 +1280,146 @@ 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, + normalize=False, + sigmoid_noise=0., + sigmoid_noise_seed=None, + score_bias_init=0., + mode="parallel", + dtype=None, + name="BahdanauMonotonicAttention", + **kwargs): + """Construct the Attention mechanism. + + Args: + units: The depth of the query mechanism. + 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. + 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) + super(BahdanauMonotonicAttentionV2, self).__init__( + query_layer=query_layer, + memory_layer=memory_layer, + probability_fn=wrapped_probability_fn, + name=name, + dtype=dtype, + **kwargs) + 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 + + def build(self, input_shape): + super(BahdanauMonotonicAttentionV2, self).build(input_shape) + self.attention_v = self.add_weight( + "attention_v", [self.units], dtype=self.dtype) + 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 not self.normalize: + self.attention_g = None + self.attention_b = None + else: + 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, + } + 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. @@ -956,7 +1510,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) @@ -966,6 +1525,129 @@ 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, + 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. + 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) + super(LuongMonotonicAttentionV2, self).__init__( + query_layer=None, + memory_layer=memory_layer, + probability_fn=wrapped_probability_fn, + name=name, + dtype=dtype, + **kwargs) + 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 + + def build(self, input_shape): + super(LuongMonotonicAttentionV2, self).build(input_shape) + if self.scale: + self.attention_g = self.add_weight( + "attention_g", initializer=init_ops.ones_initializer, shape=()) + else: + self.attention_g = 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", @@ -1022,6 +1704,97 @@ 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: + seq_len_mask = None + seq_len_batch_size = None + 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) + seq_len_batch_size = ( + tensor_shape.dimension_value(memory_sequence_length.shape[0]) + or array_ops.shape(memory_sequence_length)[0]) + else: + # For memory_mask is not None + seq_len_mask = memory_mask + seq_len_batch_size = ( + tensor_shape.dimension_value(memory_mask.shape[0]) + or array_ops.shape(memory_mask)[0]) + 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) + m_batch_size = tensor_shape.dimension_value( + m.shape[0]) or array_ops.shape(m)[0] + if seq_len_batch_size 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=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. @@ -1035,8 +1808,8 @@ def hardmax(logits, name=None): """ with ops.name_scope(name, "Hardmax", [logits]): logits = ops.convert_to_tensor(logits, name="logits") - if logits.get_shape()[-1].value is not None: - depth = logits.get_shape()[-1].value + if tensor_shape.dimension_value(logits.get_shape()[-1]) is not None: + depth = tensor_shape.dimension_value(logits.get_shape()[-1]) else: depth = array_ops.shape(logits)[-1] return array_ops.one_hot( @@ -1084,7 +1857,8 @@ class AttentionWrapper(rnn_cell_impl.RNNCell): output_attention=True, initial_cell_state=None, name=None, - attention_layer=None): + attention_layer=None, + attention_fn=None): """Construct the `AttentionWrapper`. **NOTE** If you are using the `BeamSearchDecoder` with a cell wrapped in @@ -1128,7 +1902,9 @@ class AttentionWrapper(rnn_cell_impl.RNNCell): feed the context and cell output into the attention layer to generate attention at each time step. If attention_mechanism is a list, attention_layer_size must be a list of the same length. If - attention_layer is set, this must be None. + attention_layer is set, this must be None. If attention_fn is set, + it must guaranteed that the outputs of attention_fn also meet the + above requirements. alignment_history: Python boolean, whether to store alignment history from all time steps in the final output state (currently stored as a time major `TensorArray` on which you must call `stack()`). @@ -1154,6 +1930,12 @@ class AttentionWrapper(rnn_cell_impl.RNNCell): the context as attention at each time step. If attention_mechanism is a list, attention_layer must be a list of the same length. If attention_layers_size is set, this must be None. + attention_fn: An optional callable function that allows users to provide + their own customized attention function, which takes input + (attention_mechanism, cell_output, attention_state, attention_layer) and + outputs (attention, alignments, next_attention_state). If provided, + the attention_layer_size should be the size of the outputs of + attention_fn. Raises: TypeError: `attention_layer_size` is not None and (`attention_mechanism` @@ -1224,17 +2006,22 @@ class AttentionWrapper(rnn_cell_impl.RNNCell): "layer per attention_mechanism, saw: %d vs %d" % (len(self._attention_layers), len(attention_mechanisms))) self._attention_layer_size = sum( - layer.compute_output_shape( + tensor_shape.dimension_value(layer.compute_output_shape( [None, - cell.output_size + mechanism.values.shape[-1].value])[-1].value + cell.output_size + tensor_shape.dimension_value( + mechanism.values.shape[-1])])[-1]) for layer, mechanism in zip( self._attention_layers, attention_mechanisms)) else: self._attention_layers = None self._attention_layer_size = sum( - attention_mechanism.values.get_shape()[-1].value + tensor_shape.dimension_value(attention_mechanism.values.shape[-1]) for attention_mechanism in attention_mechanisms) + if attention_fn is None: + attention_fn = _compute_attention + self._attention_fn = attention_fn + self._cell = cell self._attention_mechanisms = attention_mechanisms self._cell_input_fn = cell_input_fn @@ -1246,7 +2033,7 @@ class AttentionWrapper(rnn_cell_impl.RNNCell): else: final_state_tensor = nest.flatten(initial_cell_state)[-1] state_batch_size = ( - final_state_tensor.shape[0].value + tensor_shape.dimension_value(final_state_tensor.shape[0]) or array_ops.shape(final_state_tensor)[0]) error_message = ( "When constructing AttentionWrapper %s: " % self._base_name + @@ -1412,7 +2199,8 @@ class AttentionWrapper(rnn_cell_impl.RNNCell): cell_output, next_cell_state = self._cell(cell_inputs, cell_state) cell_batch_size = ( - cell_output.shape[0].value or array_ops.shape(cell_output)[0]) + tensor_shape.dimension_value(cell_output.shape[0]) or + array_ops.shape(cell_output)[0]) error_message = ( "When applying AttentionWrapper %s: " % self.name + "Non-matching batch sizes between the memory " @@ -1437,7 +2225,7 @@ class AttentionWrapper(rnn_cell_impl.RNNCell): all_attention_states = [] maybe_all_histories = [] for i, attention_mechanism in enumerate(self._attention_mechanisms): - attention, alignments, next_attention_state = _compute_attention( + attention, alignments, next_attention_state = self._attention_fn( attention_mechanism, cell_output, previous_attention_state[i], self._attention_layers[i] if self._attention_layers else None) alignment_history = previous_alignment_history[i].write( diff --git a/tensorflow/contrib/seq2seq/python/ops/beam_search_decoder.py b/tensorflow/contrib/seq2seq/python/ops/beam_search_decoder.py index 605e3143fd2459d098ee967568e9f2fa0073d0c5..8f8f057702951094758b277ce060955f3dc6e99d 100644 --- a/tensorflow/contrib/seq2seq/python/ops/beam_search_decoder.py +++ b/tensorflow/contrib/seq2seq/python/ops/beam_search_decoder.py @@ -85,7 +85,8 @@ def _tile_batch(t, multiplier): tiling = [1] * (t.shape.ndims + 1) tiling[1] = multiplier tiled_static_batch_size = ( - t.shape[0].value * multiplier if t.shape[0].value is not None else None) + t.shape.dims[0].value * multiplier + if t.shape.dims[0].value is not None else None) tiled = array_ops.tile(array_ops.expand_dims(t, 1), tiling) tiled = array_ops.reshape(tiled, array_ops.concat( @@ -138,9 +139,9 @@ def gather_tree_from_array(t, parent_ids, sequence_length): A `Tensor` which is a stacked `TensorArray` of the same size and type as `t` and where beams are sorted in each `Tensor` according to `parent_ids`. """ - max_time = parent_ids.shape[0].value or array_ops.shape(parent_ids)[0] - batch_size = parent_ids.shape[1].value or array_ops.shape(parent_ids)[1] - beam_width = parent_ids.shape[2].value or array_ops.shape(parent_ids)[2] + max_time = parent_ids.shape.dims[0].value or array_ops.shape(parent_ids)[0] + batch_size = parent_ids.shape.dims[1].value or array_ops.shape(parent_ids)[1] + beam_width = parent_ids.shape.dims[2].value or array_ops.shape(parent_ids)[2] # Generate beam ids that will be reordered by gather_tree. beam_ids = array_ops.expand_dims( @@ -191,9 +192,9 @@ def _check_static_batch_beam_maybe(shape, batch_size, beam_width): reshaped to [batch_size, beam_size, -1]. """ reshaped_shape = tensor_shape.TensorShape([batch_size, beam_width, None]) - if (batch_size is not None and shape[0].value is not None + if (batch_size is not None and shape.dims[0].value is not None and (shape[0] != batch_size * beam_width - or (shape.ndims >= 2 and shape[1].value is not None + or (shape.ndims >= 2 and shape.dims[1].value is not None and (shape[0] != batch_size or shape[1] != beam_width)))): tf_logging.warn("TensorArray reordering expects elements to be " "reshapable to %s which is incompatible with the " @@ -722,7 +723,7 @@ def _beam_search_step(time, logits, next_cell_state, beam_state, batch_size, total_probs = array_ops.expand_dims(beam_state.log_probs, 2) + step_log_probs # Calculate the continuation lengths by adding to all continuing beams. - vocab_size = logits.shape[-1].value or array_ops.shape(logits)[-1] + vocab_size = logits.shape.dims[-1].value or array_ops.shape(logits)[-1] lengths_to_add = array_ops.one_hot( indices=array_ops.fill([batch_size, beam_width], end_token), depth=vocab_size, @@ -920,6 +921,7 @@ def _get_scores(log_probs, sequence_lengths, length_penalty_weight, """ length_penalty_ = _length_penalty( sequence_lengths=sequence_lengths, penalty_factor=length_penalty_weight) + length_penalty_ = math_ops.cast(length_penalty_, dtype=log_probs.dtype) scores = log_probs / length_penalty_ coverage_penalty_weight = 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/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/signal/BUILD b/tensorflow/contrib/signal/BUILD index 6bd58c4d322c04d4d14d04678e24a05c0f876208..5e4f130b31483204a111e2f778fa5d0fc4526fea 100644 --- a/tensorflow/contrib/signal/BUILD +++ b/tensorflow/contrib/signal/BUILD @@ -4,129 +4,11 @@ licenses(["notice"]) # Apache 2.0 exports_files(["LICENSE"]) -load("//tensorflow:tensorflow.bzl", "cuda_py_tests") -load("//tensorflow:tensorflow.bzl", "py_test") # @unused - py_library( name = "signal_py", - srcs = ["__init__.py"] + glob(["python/ops/*.py"]), - srcs_version = "PY2AND3", - deps = [ - "//tensorflow/python:array_ops", - "//tensorflow/python:constant_op", - "//tensorflow/python:control_flow_ops", - "//tensorflow/python:dtypes", - "//tensorflow/python:framework_ops", - "//tensorflow/python:math_ops", - "//tensorflow/python:spectral_ops", - "//tensorflow/python:tensor_util", - "//tensorflow/python:util", - "//third_party/py/numpy", - ], -) - -py_library( - name = "test_util", - srcs = ["python/kernel_tests/test_util.py"], + srcs = ["__init__.py"], srcs_version = "PY2AND3", deps = [ - "//tensorflow/core:protos_all_py", - "//tensorflow/python:tf_optimizer", - "//tensorflow/python:training", - ], -) - -cuda_py_tests( - name = "mel_ops_test", - srcs = ["python/kernel_tests/mel_ops_test.py"], - additional_deps = [ - ":signal_py", - ":test_util", - "//third_party/py/numpy", - "//tensorflow/python:client_testlib", - ], -) - -cuda_py_tests( - name = "mfcc_ops_test", - srcs = ["python/kernel_tests/mfcc_ops_test.py"], - additional_deps = [ - ":signal_py", - "//third_party/py/numpy", - "//tensorflow/python:client_testlib", - "//tensorflow/python:framework", - "//tensorflow/python:framework_for_generated_wrappers", - "//tensorflow/python:framework_test_lib", - "//tensorflow/python:spectral_ops_test_util", - ], -) - -cuda_py_tests( - name = "reconstruction_ops_test", - srcs = ["python/kernel_tests/reconstruction_ops_test.py"], - additional_deps = [ - ":signal_py", - "//third_party/py/numpy", - "//tensorflow/python:array_ops", - "//tensorflow/python:gradients", - "//tensorflow/python:math_ops", - "//tensorflow/python:client_testlib", - "//tensorflow/python:framework", - "//tensorflow/python:framework_for_generated_wrappers", - "//tensorflow/python:framework_test_lib", - "//tensorflow/python:platform_test", - ], -) - -cuda_py_tests( - name = "shape_ops_test", - srcs = ["python/kernel_tests/shape_ops_test.py"], - additional_deps = [ - ":signal_py", - ":test_util", - "//third_party/py/numpy", - "//tensorflow/python:array_ops", - "//tensorflow/python:math_ops", - "//tensorflow/python:client_testlib", - "//tensorflow/python:framework", - "//tensorflow/python:framework_for_generated_wrappers", - "//tensorflow/python:framework_test_lib", - "//tensorflow/python:platform_test", - ], -) - -cuda_py_tests( - name = "spectral_ops_test", - size = "large", - srcs = ["python/kernel_tests/spectral_ops_test.py"], - additional_deps = [ - ":signal_py", - "//third_party/py/numpy", - "//tensorflow/python:array_ops", - "//tensorflow/python:gradients", - "//tensorflow/python:math_ops", - "//tensorflow/python:random_ops", - "//tensorflow/python:client_testlib", - "//tensorflow/python:framework", - "//tensorflow/python:framework_for_generated_wrappers", - "//tensorflow/python:framework_test_lib", - "//tensorflow/python:platform_test", - "//tensorflow/python:spectral_ops_test_util", - ], - tags = ["nomac"], -) - -cuda_py_tests( - name = "window_ops_test", - srcs = ["python/kernel_tests/window_ops_test.py"], - additional_deps = [ - ":signal_py", - ":test_util", - "//third_party/py/numpy", - "//tensorflow/python:client_testlib", - "//tensorflow/python:framework", - "//tensorflow/python:framework_for_generated_wrappers", - "//tensorflow/python:framework_test_lib", - "//tensorflow/python:platform_test", + "//tensorflow/python/ops/signal", ], ) diff --git a/tensorflow/contrib/signal/__init__.py b/tensorflow/contrib/signal/__init__.py index d088e744346aac0aa8675b95d7b792379fc7b019..d01f5ccf51c132082a419ec7db49045ef8bab725 100644 --- a/tensorflow/contrib/signal/__init__.py +++ b/tensorflow/contrib/signal/__init__.py @@ -14,6 +14,9 @@ # ============================================================================== """Signal processing operations. +`tf.contrib.signal` has been renamed to `tf.signal`. `tf.contrib.signal` will be +removed in TensorFlow 2.0. + See the [Contrib Signal](https://tensorflow.org/api_guides/python/contrib.signal) guide. @@ -39,18 +42,20 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -from tensorflow.contrib.signal.python.ops.mel_ops import linear_to_mel_weight_matrix -from tensorflow.contrib.signal.python.ops.mfcc_ops import mfccs_from_log_mel_spectrograms -from tensorflow.contrib.signal.python.ops.reconstruction_ops import overlap_and_add -from tensorflow.contrib.signal.python.ops.shape_ops import frame +from tensorflow.python.ops.signal.mel_ops import linear_to_mel_weight_matrix +from tensorflow.python.ops.signal.mfcc_ops import mfccs_from_log_mel_spectrograms +from tensorflow.python.ops.signal.reconstruction_ops import overlap_and_add +from tensorflow.python.ops.signal.shape_ops import frame +from tensorflow.python.ops.signal.spectral_ops import inverse_stft +from tensorflow.python.ops.signal.spectral_ops import inverse_stft_window_fn +from tensorflow.python.ops.signal.spectral_ops import stft +from tensorflow.python.ops.signal.window_ops import hamming_window +from tensorflow.python.ops.signal.window_ops import hann_window + +from tensorflow.python.util.all_util import remove_undocumented + # `frame` used to be named `frames`, which is a noun and not a verb. # Keep an alias to `frames` for backwards compatibility. -from tensorflow.contrib.signal.python.ops.shape_ops import frame as frames -from tensorflow.contrib.signal.python.ops.spectral_ops import inverse_stft -from tensorflow.contrib.signal.python.ops.spectral_ops import inverse_stft_window_fn -from tensorflow.contrib.signal.python.ops.spectral_ops import stft -from tensorflow.contrib.signal.python.ops.window_ops import hamming_window -from tensorflow.contrib.signal.python.ops.window_ops import hann_window +frames = frame -from tensorflow.python.util.all_util import remove_undocumented remove_undocumented(__name__) diff --git a/tensorflow/contrib/signal/python/__init__.py b/tensorflow/contrib/signal/python/__init__.py deleted file mode 100644 index e672d1146c53a813613c9076c0cb6056f7081441..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/signal/python/__init__.py +++ /dev/null @@ -1,19 +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. -# ============================================================================== -"""Signal ops.""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function diff --git a/tensorflow/contrib/signal/python/kernel_tests/reconstruction_ops_test.py b/tensorflow/contrib/signal/python/kernel_tests/reconstruction_ops_test.py deleted file mode 100644 index c476cd4e00d621bbf694b27236d82f18e8699c8c..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/signal/python/kernel_tests/reconstruction_ops_test.py +++ /dev/null @@ -1,192 +0,0 @@ -# Copyright 2017 The TensorFlow Authors. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== -"""Tests for reconstruction_ops.""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import numpy as np - -from tensorflow.contrib.signal.python.ops import reconstruction_ops -from tensorflow.python.framework import constant_op -from tensorflow.python.framework import dtypes -from tensorflow.python.ops import array_ops -from tensorflow.python.ops import gradients_impl -from tensorflow.python.ops import math_ops -from tensorflow.python.platform import test - - -class ReconstructionOpsTest(test.TestCase): - - def __init__(self, *args, **kwargs): - super(ReconstructionOpsTest, self).__init__(*args, **kwargs) - self.batch_size = 3 - self.frames = 3 - self.samples = 5 - - self.bases = np.array(range(2, 5)) - exponents = np.array(range(self.frames * self.samples)) - powers = np.power(self.bases[:, np.newaxis], exponents[np.newaxis, :]) - - self.powers = np.reshape(powers, [self.batch_size, self.frames, - self.samples]) - self.frame_hop = 2 - - # Hand computed example using powers of unique numbers: this is easily - # verified. - self.expected_string = ["1", "10", "100100", "1001000", "10010010000", - "100100000000", "1001000000000", "10000000000000", - "100000000000000"] - - def test_all_ones(self): - signal = constant_op.constant(np.ones((3, 5)), dtype=dtypes.int64) - reconstruction = reconstruction_ops.overlap_and_add(signal, 2) - - with self.session(use_gpu=True) as sess: - output = sess.run(reconstruction) - - expected_output = np.array([1, 1, 2, 2, 3, 2, 2, 1, 1]) - - self.assertAllClose(output, expected_output) - - def test_simple(self): - def make_input(frame_length, num_frames=3): - """Generate a tensor of num_frames frames of frame_length.""" - return np.reshape(np.arange(1, num_frames * frame_length + 1), - (-1, frame_length)) - - # List of (signal, expected_result, frame_hop). - configurations = [ - # All hop lengths on a frame length of 2. - (make_input(2), [1, 5, 9, 6], 1), - (make_input(2), [1, 2, 3, 4, 5, 6], 2), - - # All hop lengths on a frame length of 3. - (make_input(3), [1, 6, 15, 14, 9], 1), - (make_input(3), [1, 2, 7, 5, 13, 8, 9], 2), - (make_input(3), [1, 2, 3, 4, 5, 6, 7, 8, 9], 3), - - # All hop lengths on a frame length of 4. - (make_input(4), [1, 7, 18, 21, 19, 12], 1), - (make_input(4), [1, 2, 8, 10, 16, 18, 11, 12], 2), - (make_input(4), [1, 2, 3, 9, 6, 7, 17, 10, 11, 12], 3), - (make_input(4), [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], 4), - ] - - with self.session(use_gpu=True): - for signal, expected, frame_hop in configurations: - reconstruction = reconstruction_ops.overlap_and_add( - np.array(signal), frame_hop).eval() - expected_output = np.array(expected) - self.assertAllClose(reconstruction, expected_output) - - def test_powers(self): - signal = constant_op.constant(np.squeeze(self.powers[0, :, :]), - dtype=dtypes.int64) - reconstruction = reconstruction_ops.overlap_and_add(signal, self.frame_hop) - - with self.session(use_gpu=True) as sess: - output = sess.run(reconstruction) - string_output = [np.base_repr(x, self.bases[0]) for x in output] - - self.assertEqual(string_output, self.expected_string) - - def test_batch(self): - signal = constant_op.constant(self.powers, dtype=dtypes.int64) - reconstruction = reconstruction_ops.overlap_and_add(signal, self.frame_hop) - - with self.session(use_gpu=True) as sess: - output = sess.run(reconstruction) - - accumulator = True - for i in range(self.batch_size): - string_output = [np.base_repr(x, self.bases[i]) for x in output[i, :]] - accumulator = accumulator and (string_output == self.expected_string) - - self.assertTrue(accumulator) - - def test_one_element_batch(self): - input_matrix = np.squeeze(self.powers[0, :, :]) - input_matrix = input_matrix[np.newaxis, :, :].astype(float) - signal = constant_op.constant(input_matrix, dtype=dtypes.float32) - reconstruction = reconstruction_ops.overlap_and_add(signal, self.frame_hop) - - with self.session(use_gpu=True) as sess: - output = sess.run(reconstruction) - - string_output = [np.base_repr(int(x), self.bases[0]) for x in - np.squeeze(output)] - - self.assertEqual(output.shape, (1, 9)) - self.assertEqual(string_output, self.expected_string) - - def test_gradient(self): - configurations = [ - ((1, 128), 1), - ((5, 35), 17), - ((10, 128), 128), - ((2, 10, 128), 127), - ((2, 2, 10, 128), 126), - ((2, 2, 2, 10, 128), 125), - ] - - with self.session(use_gpu=True) as sess: - for shape, frame_hop in configurations: - signal = array_ops.zeros(shape) - reconstruction = reconstruction_ops.overlap_and_add(signal, frame_hop) - loss = math_ops.reduce_sum(reconstruction) - # Increasing any sample in the input frames by one will increase the sum - # of all the samples in the reconstruction by 1, so the gradient should - # be all ones, no matter the shape or hop. - gradient = sess.run(gradients_impl.gradients([loss], [signal])[0]) - self.assertTrue((gradient == 1.0).all()) - - def test_gradient_batch(self): - with self.session(use_gpu=True) as sess: - signal = array_ops.zeros((2, 10, 10)) - frame_hop = 10 - reconstruction = reconstruction_ops.overlap_and_add(signal, frame_hop) - - # Multiply the first batch-item's reconstruction by zeros. This will block - # gradient from flowing into the first batch item from the loss. Multiply - # the second batch item by the integers from 0 to 99. Since there is zero - # overlap, the gradient for this batch item will be 0-99 shaped as (10, - # 10). - reconstruction *= array_ops.stack( - [array_ops.zeros((100,)), math_ops.to_float(math_ops.range(100))]) - loss = math_ops.reduce_sum(reconstruction) - - # Verify that only the second batch item receives gradient. - gradient = sess.run(gradients_impl.gradients([loss], [signal])[0]) - expected_gradient = np.stack([ - np.zeros((10, 10)), - np.reshape(np.arange(100).astype(np.float32), (10, 10))]) - self.assertAllEqual(expected_gradient, gradient) - - def test_gradient_numerical(self): - with self.session(use_gpu=True): - shape = (2, 10, 10) - framed_signal = array_ops.zeros(shape) - frame_hop = 10 - reconstruction = reconstruction_ops.overlap_and_add( - framed_signal, frame_hop) - error = test.compute_gradient_error( - framed_signal, shape, reconstruction, [2, 100]) - self.assertLess(error, 2e-5) - - -if __name__ == "__main__": - test.main() diff --git a/tensorflow/contrib/signal/python/kernel_tests/test_util.py b/tensorflow/contrib/signal/python/kernel_tests/test_util.py deleted file mode 100644 index b4422a49887378187a2be46275d4dabf1fbd40a1..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/signal/python/kernel_tests/test_util.py +++ /dev/null @@ -1,47 +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. -# ============================================================================== -"""Test utilities for tf.contrib.signal.""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -from tensorflow.core.protobuf import rewriter_config_pb2 -from tensorflow.python.grappler import tf_optimizer -from tensorflow.python.training import saver - - -def grappler_optimize(graph, fetches=None, rewriter_config=None): - """Tries to optimize the provided graph using grappler. - - Args: - graph: A `tf.Graph` instance containing the graph to optimize. - fetches: An optional list of `Tensor`s to fetch (i.e. not optimize away). - Grappler uses the 'train_op' collection to look for fetches, so if not - provided this collection should be non-empty. - rewriter_config: An optional `tf.RewriterConfig` to use when rewriting the - graph. - - Returns: - A `tf.GraphDef` containing the rewritten graph. - """ - if rewriter_config is None: - rewriter_config = rewriter_config_pb2.RewriterConfig() - rewriter_config.min_graph_nodes = -1 - if fetches is not None: - for fetch in fetches: - graph.add_to_collection('train_op', fetch) - metagraph = saver.export_meta_graph(graph_def=graph.as_graph_def()) - return tf_optimizer.OptimizeGraph(rewriter_config, metagraph) diff --git a/tensorflow/contrib/signal/python/ops/reconstruction_ops.py b/tensorflow/contrib/signal/python/ops/reconstruction_ops.py deleted file mode 100644 index 4db8dc2ca090534f2cda66bd55c30dfa389b860a..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/signal/python/ops/reconstruction_ops.py +++ /dev/null @@ -1,150 +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. -# ============================================================================== -"""Signal reconstruction via overlapped addition of frames.""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -from tensorflow.contrib.signal.python.ops import shape_ops -from tensorflow.contrib.signal.python.ops import util_ops -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 math_ops - - -def _shuffle_to_front(input_tensor, k): - """Shuffles the last `k` indices of `input_tensor` to the front. - - Transposes `input_tensor` to have the last `k` indices at the front. The input - may have arbitrary rank and unknown shape. - - Args: - input_tensor: A `Tensor` of arbitrary rank and unknown shape. - k: A scalar `Tensor` specifying how many indices to shuffle. - - Returns: - A transposed version of `input_tensor` with `k` indices shuffled to the - front. - - Raises: - ValueError: If `input_tensor` is not at least rank `k` or `k` is not scalar. - """ - k = ops.convert_to_tensor(k, name="k") - k.shape.with_rank(0) - k_static = tensor_util.constant_value(k) - if k_static is not None: - input_tensor.shape.with_rank_at_least(k_static) - - rank = array_ops.rank(input_tensor) - outer_indices, inner_indices = array_ops.split(math_ops.range(rank), - [rank - k, k]) - permutation = array_ops.concat([inner_indices, outer_indices], 0) - - return array_ops.transpose(input_tensor, perm=permutation) - - -def overlap_and_add(signal, frame_step, name=None): - """Reconstructs a signal from a framed representation. - - Adds potentially overlapping frames of a signal with shape - `[..., frames, frame_length]`, offsetting subsequent frames by `frame_step`. - The resulting tensor has shape `[..., output_size]` where - - output_size = (frames - 1) * frame_step + frame_length - - Args: - signal: A [..., frames, frame_length] `Tensor`. All dimensions may be - unknown, and rank must be at least 2. - frame_step: An integer or scalar `Tensor` denoting overlap offsets. Must be - less than or equal to `frame_length`. - name: An optional name for the operation. - - Returns: - A `Tensor` with shape `[..., output_size]` containing the overlap-added - frames of `signal`'s inner-most two dimensions. - - Raises: - ValueError: If `signal`'s rank is less than 2, `frame_step` is not a scalar - integer or `frame_step` is greater than `frame_length`. - """ - with ops.name_scope(name, "overlap_and_add", [signal, frame_step]): - signal = ops.convert_to_tensor(signal, name="signal") - signal.shape.with_rank_at_least(2) - frame_step = ops.convert_to_tensor(frame_step, name="frame_step") - frame_step.shape.assert_has_rank(0) - if not frame_step.dtype.is_integer: - raise ValueError("frame_step must be an integer. Got %s" % - frame_step.dtype) - - signal_shape = array_ops.shape(signal) - - # All dimensions that are not part of the overlap-and-add. Can be empty for - # rank 2 inputs. - outer_dimensions = signal_shape[:-2] - - # If frame_length and frame_step are known at graph construction time, check - # frame_step is less than or equal to frame_length. - frame_step_static = tensor_util.constant_value(frame_step) - if (frame_step_static is not None and signal.shape.ndims is not None and - signal.shape[-1].value is not None): - if frame_step_static > signal.shape[-1].value: - raise ValueError( - "frame_step (%d) must be less than or equal to " - "frame_length (%d)" % ( - frame_step_static, signal.shape[-1].value)) - # If frame_length is equal to frame_step, there's no overlap so just - # reshape the tensor. - if frame_step_static == signal.shape[-1].value: - return array_ops.reshape(signal, array_ops.concat( - [outer_dimensions, [-1]], 0)) - - signal_rank = array_ops.rank(signal) - frames = signal_shape[-2] - frame_length = signal_shape[-1] - - subframe_length = util_ops.gcd(frame_length, frame_step) - subframe_step = frame_step // subframe_length - subframes_per_frame = frame_length // subframe_length - output_size = frame_step * (frames - 1) + frame_length - output_subframes = output_size // subframe_length - - # To avoid overlap-adding sample-by-sample, we overlap-add at the "subframe" - # level, where a subframe is gcd(frame_length, frame_step). Reshape signal - # from [..., frames, frame_length] into [..., subframes, subframe_length]. - subframe_shape = array_ops.concat( - [outer_dimensions, [-1, subframe_length]], 0) - subframe_signal = array_ops.reshape(signal, subframe_shape) - - # Now we shuffle the last [subframes, subframe_length] dimensions to the - # front. - # TODO(rjryan): Add an axis argument to unsorted_segment_sum so we can - # avoid this pair of transposes. - subframe_signal = _shuffle_to_front(subframe_signal, 2) - - # Use unsorted_segment_sum to add overlapping subframes together. - segment_ids = array_ops.reshape(shape_ops.frame( - math_ops.range(output_subframes), subframes_per_frame, subframe_step, - pad_end=False), [-1]) - result = math_ops.unsorted_segment_sum(subframe_signal, segment_ids, - num_segments=output_subframes) - - # result is a [subframes, subframe_length, ...outer_dimensions] tensor. We - # return a [...outer_dimensions, output_size] tensor with a transpose and - # reshape. - result_shape = array_ops.concat([outer_dimensions, [output_size]], 0) - return array_ops.reshape(_shuffle_to_front(result, signal_rank - 2), - result_shape) diff --git a/tensorflow/contrib/signal/python/ops/spectral_ops.py b/tensorflow/contrib/signal/python/ops/spectral_ops.py deleted file mode 100644 index a8b5deff6ca3a4a756d31b904e577f08f6155fd7..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/signal/python/ops/spectral_ops.py +++ /dev/null @@ -1,287 +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. -# ============================================================================== -"""Spectral operations (e.g. Short-time Fourier Transform).""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import functools - -import numpy as np - -from tensorflow.contrib.signal.python.ops import reconstruction_ops -from tensorflow.contrib.signal.python.ops import shape_ops -from tensorflow.contrib.signal.python.ops import window_ops -from tensorflow.python.framework import constant_op -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 math_ops -from tensorflow.python.ops import spectral_ops - - -def stft(signals, frame_length, frame_step, fft_length=None, - window_fn=functools.partial(window_ops.hann_window, periodic=True), - pad_end=False, name=None): - """Computes the [Short-time Fourier Transform][stft] of `signals`. - - Implemented with GPU-compatible ops and supports gradients. - - Args: - signals: A `[..., samples]` `float32` `Tensor` of real-valued signals. - frame_length: An integer scalar `Tensor`. The window length in samples. - frame_step: An integer scalar `Tensor`. The number of samples to step. - fft_length: An integer scalar `Tensor`. The size of the FFT to apply. - If not provided, uses the smallest power of 2 enclosing `frame_length`. - window_fn: A callable that takes a window length and a `dtype` keyword - argument and returns a `[window_length]` `Tensor` of samples in the - provided datatype. If set to `None`, no windowing is used. - pad_end: Whether to pad the end of `signals` with zeros when the provided - frame length and step produces a frame that lies partially past its end. - name: An optional name for the operation. - - Returns: - A `[..., frames, fft_unique_bins]` `Tensor` of `complex64` STFT values where - `fft_unique_bins` is `fft_length // 2 + 1` (the unique components of the - FFT). - - Raises: - ValueError: If `signals` is not at least rank 1, `frame_length` is - not scalar, or `frame_step` is not scalar. - - [stft]: https://en.wikipedia.org/wiki/Short-time_Fourier_transform - """ - with ops.name_scope(name, 'stft', [signals, frame_length, - frame_step]): - signals = ops.convert_to_tensor(signals, name='signals') - signals.shape.with_rank_at_least(1) - frame_length = ops.convert_to_tensor(frame_length, name='frame_length') - frame_length.shape.assert_has_rank(0) - frame_step = ops.convert_to_tensor(frame_step, name='frame_step') - frame_step.shape.assert_has_rank(0) - - if fft_length is None: - fft_length = _enclosing_power_of_two(frame_length) - else: - fft_length = ops.convert_to_tensor(fft_length, name='fft_length') - - framed_signals = shape_ops.frame( - signals, frame_length, frame_step, pad_end=pad_end) - - # Optionally window the framed signals. - if window_fn is not None: - window = window_fn(frame_length, dtype=framed_signals.dtype) - framed_signals *= window - - # spectral_ops.rfft produces the (fft_length/2 + 1) unique components of the - # FFT of the real windowed signals in framed_signals. - return spectral_ops.rfft(framed_signals, [fft_length]) - - -def inverse_stft_window_fn(frame_step, - forward_window_fn=functools.partial( - window_ops.hann_window, periodic=True), - name=None): - """Generates a window function that can be used in `inverse_stft`. - - Constructs a window that is equal to the forward window with a further - pointwise amplitude correction. `inverse_stft_window_fn` is equivalent to - `forward_window_fn` in the case where it would produce an exact inverse. - - See examples in `inverse_stft` documentation for usage. - - Args: - frame_step: An integer scalar `Tensor`. The number of samples to step. - forward_window_fn: window_fn used in the forward transform, `stft`. - name: An optional name for the operation. - - Returns: - A callable that takes a window length and a `dtype` keyword argument and - returns a `[window_length]` `Tensor` of samples in the provided datatype. - The returned window is suitable for reconstructing original waveform in - inverse_stft. - """ - with ops.name_scope(name, 'inverse_stft_window_fn', [forward_window_fn]): - frame_step = ops.convert_to_tensor(frame_step, name='frame_step') - frame_step.shape.assert_has_rank(0) - - def inverse_stft_window_fn_inner(frame_length, dtype): - """Computes a window that can be used in `inverse_stft`. - - Args: - frame_length: An integer scalar `Tensor`. The window length in samples. - dtype: Data type of waveform passed to `stft`. - - Returns: - A window suitable for reconstructing original waveform in `inverse_stft`. - - Raises: - ValueError: If `frame_length` is not scalar, `forward_window_fn` is not a - callable that takes a window length and a `dtype` keyword argument and - returns a `[window_length]` `Tensor` of samples in the provided datatype - `frame_step` is not scalar, or `frame_step` is not scalar. - """ - with ops.name_scope(name, 'inverse_stft_window_fn', [forward_window_fn]): - frame_length = ops.convert_to_tensor(frame_length, name='frame_length') - frame_length.shape.assert_has_rank(0) - - # Use equation 7 from Griffin + Lim. - forward_window = forward_window_fn(frame_length, dtype=dtype) - denom = math_ops.square(forward_window) - overlaps = -(-frame_length // frame_step) # Ceiling division. - denom = array_ops.pad(denom, [(0, overlaps * frame_step - frame_length)]) - denom = array_ops.reshape(denom, [overlaps, frame_step]) - denom = math_ops.reduce_sum(denom, 0, keepdims=True) - denom = array_ops.tile(denom, [overlaps, 1]) - denom = array_ops.reshape(denom, [overlaps * frame_step]) - - return forward_window / denom[:frame_length] - return inverse_stft_window_fn_inner - - -def inverse_stft(stfts, - frame_length, - frame_step, - fft_length=None, - window_fn=functools.partial(window_ops.hann_window, - periodic=True), - name=None): - """Computes the inverse [Short-time Fourier Transform][stft] of `stfts`. - - To reconstruct an original waveform, a complimentary window function should - be used in inverse_stft. Such a window function can be constructed with - tf.contrib.signal.inverse_stft_window_fn. - - Example: - - ```python - frame_length = 400 - frame_step = 160 - waveform = tf.placeholder(dtype=tf.float32, shape=[1000]) - stft = tf.contrib.signal.stft(waveform, frame_length, frame_step) - inverse_stft = tf.contrib.signal.inverse_stft( - stft, frame_length, frame_step, - window_fn=tf.contrib.signal.inverse_stft_window_fn(frame_step)) - ``` - - if a custom window_fn is used in stft, it must be passed to - inverse_stft_window_fn: - - ```python - frame_length = 400 - frame_step = 160 - window_fn = functools.partial(window_ops.hamming_window, periodic=True), - waveform = tf.placeholder(dtype=tf.float32, shape=[1000]) - stft = tf.contrib.signal.stft( - waveform, frame_length, frame_step, window_fn=window_fn) - inverse_stft = tf.contrib.signal.inverse_stft( - stft, frame_length, frame_step, - window_fn=tf.contrib.signal.inverse_stft_window_fn( - frame_step, forward_window_fn=window_fn)) - ``` - - Implemented with GPU-compatible ops and supports gradients. - - Args: - stfts: A `complex64` `[..., frames, fft_unique_bins]` `Tensor` of STFT bins - representing a batch of `fft_length`-point STFTs where `fft_unique_bins` - is `fft_length // 2 + 1` - frame_length: An integer scalar `Tensor`. The window length in samples. - frame_step: An integer scalar `Tensor`. The number of samples to step. - fft_length: An integer scalar `Tensor`. The size of the FFT that produced - `stfts`. If not provided, uses the smallest power of 2 enclosing - `frame_length`. - window_fn: A callable that takes a window length and a `dtype` keyword - argument and returns a `[window_length]` `Tensor` of samples in the - provided datatype. If set to `None`, no windowing is used. - name: An optional name for the operation. - - Returns: - A `[..., samples]` `Tensor` of `float32` signals representing the inverse - STFT for each input STFT in `stfts`. - - Raises: - ValueError: If `stfts` is not at least rank 2, `frame_length` is not scalar, - `frame_step` is not scalar, or `fft_length` is not scalar. - - [stft]: https://en.wikipedia.org/wiki/Short-time_Fourier_transform - """ - with ops.name_scope(name, 'inverse_stft', [stfts]): - stfts = ops.convert_to_tensor(stfts, name='stfts') - stfts.shape.with_rank_at_least(2) - frame_length = ops.convert_to_tensor(frame_length, name='frame_length') - frame_length.shape.assert_has_rank(0) - frame_step = ops.convert_to_tensor(frame_step, name='frame_step') - frame_step.shape.assert_has_rank(0) - if fft_length is None: - fft_length = _enclosing_power_of_two(frame_length) - else: - fft_length = ops.convert_to_tensor(fft_length, name='fft_length') - fft_length.shape.assert_has_rank(0) - - real_frames = spectral_ops.irfft(stfts, [fft_length]) - - # frame_length may be larger or smaller than fft_length, so we pad or - # truncate real_frames to frame_length. - frame_length_static = tensor_util.constant_value(frame_length) - # If we don't know the shape of real_frames's inner dimension, pad and - # truncate to frame_length. - if (frame_length_static is None or - real_frames.shape.ndims is None or - real_frames.shape[-1].value is None): - real_frames = real_frames[..., :frame_length] - real_frames_rank = array_ops.rank(real_frames) - real_frames_shape = array_ops.shape(real_frames) - paddings = array_ops.concat( - [array_ops.zeros([real_frames_rank - 1, 2], - dtype=frame_length.dtype), - [[0, math_ops.maximum(0, frame_length - real_frames_shape[-1])]]], 0) - real_frames = array_ops.pad(real_frames, paddings) - # We know real_frames's last dimension and frame_length statically. If they - # are different, then pad or truncate real_frames to frame_length. - elif real_frames.shape[-1].value > frame_length_static: - real_frames = real_frames[..., :frame_length_static] - elif real_frames.shape[-1].value < frame_length_static: - pad_amount = frame_length_static - real_frames.shape[-1].value - real_frames = array_ops.pad(real_frames, - [[0, 0]] * (real_frames.shape.ndims - 1) + - [[0, pad_amount]]) - - # The above code pads the inner dimension of real_frames to frame_length, - # but it does so in a way that may not be shape-inference friendly. - # Restore shape information if we are able to. - if frame_length_static is not None and real_frames.shape.ndims is not None: - real_frames.set_shape([None] * (real_frames.shape.ndims - 1) + - [frame_length_static]) - - # Optionally window and overlap-add the inner 2 dimensions of real_frames - # into a single [samples] dimension. - if window_fn is not None: - window = window_fn(frame_length, dtype=stfts.dtype.real_dtype) - real_frames *= window - return reconstruction_ops.overlap_and_add(real_frames, frame_step) - - -def _enclosing_power_of_two(value): - """Return 2**N for integer N such that 2**N >= value.""" - value_static = tensor_util.constant_value(value) - if value_static is not None: - return constant_op.constant( - int(2**np.ceil(np.log(value_static) / np.log(2.0))), value.dtype) - return math_ops.cast( - math_ops.pow(2.0, math_ops.ceil( - math_ops.log(math_ops.to_float(value)) / math_ops.log(2.0))), - value.dtype) diff --git a/tensorflow/contrib/slim/python/slim/data/data_decoder.py b/tensorflow/contrib/slim/python/slim/data/data_decoder.py index 5a32be6c5a329068a65655b0b7be020fcd22ea18..46d33597e429123f37bfa0f3c3b8b2dc5098fe7d 100644 --- a/tensorflow/contrib/slim/python/slim/data/data_decoder.py +++ b/tensorflow/contrib/slim/python/slim/data/data_decoder.py @@ -39,12 +39,13 @@ from __future__ import print_function import abc +import six + +@six.add_metaclass(abc.ABCMeta) class DataDecoder(object): """An abstract class which is used to decode data for a provider.""" - __metaclass__ = abc.ABCMeta - @abc.abstractmethod def decode(self, data, items): """Decodes the data to returns the tensors specified by the list of items. diff --git a/tensorflow/contrib/slim/python/slim/data/data_provider.py b/tensorflow/contrib/slim/python/slim/data/data_provider.py index a49c0969d96bf7eef0200269e168941f9b8433a5..3252b4fe8470f5f9733c67c2cf44ad888ae3d2c7 100644 --- a/tensorflow/contrib/slim/python/slim/data/data_provider.py +++ b/tensorflow/contrib/slim/python/slim/data/data_provider.py @@ -38,7 +38,10 @@ from __future__ import print_function import abc +import six + +@six.add_metaclass(abc.ABCMeta) class DataProvider(object): """Maps a list of requested data items to tensors from a data source. @@ -46,7 +49,6 @@ class DataProvider(object): method which returns arbitrary types of data. No assumption is made about the source of the data nor the mechanism for providing it. """ - __metaclass__ = abc.ABCMeta def __init__(self, items_to_tensors, num_samples): """Constructs the Data Provider. diff --git a/tensorflow/contrib/slim/python/slim/data/tfexample_decoder.py b/tensorflow/contrib/slim/python/slim/data/tfexample_decoder.py index a6ce45c20365d9893895101476c9711065bfc511..1b2b6acacca838f95cb758ae88f79263993ca69e 100644 --- a/tensorflow/contrib/slim/python/slim/data/tfexample_decoder.py +++ b/tensorflow/contrib/slim/python/slim/data/tfexample_decoder.py @@ -25,6 +25,8 @@ from __future__ import print_function import abc +import six + from tensorflow.contrib.slim.python.slim.data import data_decoder from tensorflow.python.framework import dtypes from tensorflow.python.framework import sparse_tensor @@ -37,6 +39,7 @@ from tensorflow.python.ops import parsing_ops from tensorflow.python.ops import sparse_ops +@six.add_metaclass(abc.ABCMeta) class ItemHandler(object): """Specifies the item-to-Features mapping for tf.parse_example. @@ -45,8 +48,6 @@ class ItemHandler(object): parsing. """ - __metaclass__ = abc.ABCMeta - def __init__(self, keys): """Constructs the handler with the name of the tf.Feature keys to use. diff --git a/tensorflow/contrib/solvers/python/kernel_tests/lanczos_test.py b/tensorflow/contrib/solvers/python/kernel_tests/lanczos_test.py index 8fcd7aeef6a6964902666a4f3c17e05b0c7b52ee..f31bdbd399c9de4f2f5d557b75b1ece6d64a765e 100644 --- a/tensorflow/contrib/solvers/python/kernel_tests/lanczos_test.py +++ b/tensorflow/contrib/solvers/python/kernel_tests/lanczos_test.py @@ -19,6 +19,7 @@ from __future__ import print_function import numpy as np +from tensorflow.python import tf2 from tensorflow.contrib.solvers.python.ops import lanczos from tensorflow.contrib.solvers.python.ops import util from tensorflow.python.framework import constant_op @@ -80,7 +81,8 @@ if __name__ == "__main__": for shape in [[4, 4], [7, 4], [5, 8]]: for orthogonalize in True, False: for steps in range(1, min(shape) + 1): - for use_static_shape in True, False: + # TF2 does not support placeholders so we skip it + for use_static_shape in set([True, tf2.enabled()]): arg_string = "%s_%s_%s_%s_staticshape_%s" % ( dtype.__name__, "_".join(map(str, shape)), orthogonalize, steps, use_static_shape) diff --git a/tensorflow/contrib/solvers/python/kernel_tests/least_squares_test.py b/tensorflow/contrib/solvers/python/kernel_tests/least_squares_test.py index 2a9100903aae5689919a6b25fcb18ff192f250b3..841a41a2339824ab8ca15f4bdd74be697cd6fe9f 100644 --- a/tensorflow/contrib/solvers/python/kernel_tests/least_squares_test.py +++ b/tensorflow/contrib/solvers/python/kernel_tests/least_squares_test.py @@ -19,6 +19,7 @@ from __future__ import print_function import numpy as np +from tensorflow.python import tf2 from tensorflow.contrib.solvers.python.ops import least_squares from tensorflow.contrib.solvers.python.ops import util from tensorflow.python.framework import constant_op @@ -76,7 +77,8 @@ def _get_least_squares_tests(dtype_, use_static_shape_, shape_): if __name__ == "__main__": for dtype in np.float32, np.float64: for shape in [[4, 4], [8, 5], [3, 7]]: - for use_static_shape in True, False: + # TF2 does not support placeholders under eager so we skip it + for use_static_shape in set([True, tf2.enabled()]): arg_string = "%s_%s_staticshape_%s" % (dtype.__name__, "_".join(map(str, shape)), use_static_shape) diff --git a/tensorflow/contrib/solvers/python/kernel_tests/linear_equations_test.py b/tensorflow/contrib/solvers/python/kernel_tests/linear_equations_test.py index a0e6eb87bc06fb1303a7eb86fa6760458f20a9b9..10807f7a80617e56abeb6d13ce419a49a2269aac 100644 --- a/tensorflow/contrib/solvers/python/kernel_tests/linear_equations_test.py +++ b/tensorflow/contrib/solvers/python/kernel_tests/linear_equations_test.py @@ -19,6 +19,7 @@ from __future__ import print_function import numpy as np +from tensorflow.python import tf2 from tensorflow.contrib.solvers.python.ops import linear_equations from tensorflow.contrib.solvers.python.ops import util from tensorflow.python.framework import constant_op @@ -113,7 +114,8 @@ def _get_linear_equations_tests(dtype_, use_static_shape_, shape_): if __name__ == "__main__": for dtype in np.float32, np.float64: for size in 1, 4, 10: - for use_static_shape in True, False: + # TF2 does not support placeholders under eager so we skip it + for use_static_shape in set([True, tf2.enabled()]): shape = [size, size] arg_string = "%s_%s_staticshape_%s" % (dtype.__name__, size, use_static_shape) 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/summary/summary_ops_test.py b/tensorflow/contrib/summary/summary_ops_test.py index 4d1807130c57039976dfa57c27bb0d4807e75212..10e4556dacbc17ec02c2bd698389b04d517d7076 100644 --- a/tensorflow/contrib/summary/summary_ops_test.py +++ b/tensorflow/contrib/summary/summary_ops_test.py @@ -152,6 +152,27 @@ class EagerFileTest(test_util.TensorFlowTestCase): self.assertEqual(len(events), 2) self.assertEqual(events[1].summary.value[0].tag, 'scalar') + def testRecordEveryNGlobalSteps(self): + step = training_util.get_or_create_global_step() + logdir = tempfile.mkdtemp() + + def run_step(): + summary_ops.scalar('scalar', i, step=step) + step.assign_add(1) + + with summary_ops.create_file_writer( + logdir).as_default(), summary_ops.record_summaries_every_n_global_steps( + 2, step): + for i in range(10): + run_step() + # And another 10 steps as a graph function. + run_step_fn = function.defun(run_step) + for i in range(10): + run_step_fn() + + events = summary_test_util.events_from_logdir(logdir) + self.assertEqual(len(events), 11) + def testMaxQueue(self): logs = tempfile.mkdtemp() with summary_ops.create_file_writer( @@ -279,12 +300,9 @@ class EagerDbTest(summary_test_util.SummaryDbTest): def testDbURIOpen(self): tmpdb_path = os.path.join(self.get_temp_dir(), 'tmpDbURITest.sqlite') - tmpdb_uri = six.moves.urllib_parse.urljoin("file:", tmpdb_path) - tmpdb_writer = summary_ops.create_db_writer( - tmpdb_uri, - "experimentA", - "run1", - "user1") + tmpdb_uri = six.moves.urllib_parse.urljoin('file:', tmpdb_path) + tmpdb_writer = summary_ops.create_db_writer(tmpdb_uri, 'experimentA', + 'run1', 'user1') with summary_ops.always_record_summaries(): with tmpdb_writer.as_default(): summary_ops.scalar('t1', 2.0) 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/tensor_forest/python/ops/model_ops.py b/tensorflow/contrib/tensor_forest/python/ops/model_ops.py index 596c59ead3460aa63eeff44d5a11a4a8c5cde0da..290c16fe3966791ea78986539750caf938a37322 100644 --- a/tensorflow/contrib/tensor_forest/python/ops/model_ops.py +++ b/tensorflow/contrib/tensor_forest/python/ops/model_ops.py @@ -17,6 +17,8 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import functools + from tensorflow.contrib.tensor_forest.python.ops import gen_model_ops # pylint: disable=unused-import @@ -28,10 +30,12 @@ from tensorflow.contrib.tensor_forest.python.ops.gen_model_ops import update_mod # pylint: enable=unused-import from tensorflow.contrib.util import loader +from tensorflow.python.eager import context from tensorflow.python.framework import ops from tensorflow.python.ops import resources from tensorflow.python.platform import resource_loader from tensorflow.python.training import saver +from tensorflow.python.training.checkpointable import tracking _model_ops = loader.load_op_library( @@ -88,6 +92,59 @@ class TreeVariableSavable(saver.BaseSaverBuilder.SaveableObject): params=self.params.serialized_params_proto) +class TreeVariable(tracking.TrackableResource): + """A tree model.""" + + def __init__(self, params, tree_config, stats_handle, name, container=None): + self._params = params + self._tree_config = tree_config + self._stats_handle = stats_handle + self._name = name + self._container = container + self._init_op = None + super(TreeVariable, self).__init__() + self._resource_handle = self.create_resource() + + def create_resource(self): + if context.executing_eagerly(): + # TODO(allenl): This will leak memory due to kernel caching by the + # shared_name attribute value (but is better than the alternative of + # sharing everything by default when executing eagerly; hopefully creating + # tables in a loop is uncommon). + shared_name = "tree_variable_%d" % (ops.uid(),) + else: + shared_name = self._name + return gen_model_ops.decision_tree_resource_handle_op( + self._container, shared_name=shared_name, name=self._name) + + def initialize(self): + return gen_model_ops.create_tree_variable( + self.resource_handle, + self._tree_config, + params=self._params.serialized_params_proto) + + @property + def initializer(self): + if self._init_op is None: + self._init_op = self.initialize() + return self._init_op + + def is_initialized(self): + return gen_model_ops.tree_is_initialized_op(self.resource_handle) + + def _gather_saveables_for_checkpoint(self): + """For object-based checkpointing.""" + return { + "tree_variable": + functools.partial( + TreeVariableSavable, + params=self._params, + tree_handle=self.resource_handle, + stats_handle=self._stats_handle, + create_op=self._init_op) + } + + def tree_variable(params, tree_config, stats_handle, name, container=None): r"""Creates a tree model and returns a handle to it. @@ -102,18 +159,13 @@ def tree_variable(params, tree_config, stats_handle, name, container=None): A `Tensor` of type mutable `string`. The handle to the tree. """ with ops.name_scope(name, "TreeVariable") as name: - resource_handle = gen_model_ops.decision_tree_resource_handle_op( - container, shared_name=name, name=name) - - create_op = gen_model_ops.create_tree_variable( - resource_handle, - tree_config, - params=params.serialized_params_proto) - is_initialized_op = gen_model_ops.tree_is_initialized_op(resource_handle) + tree_var = TreeVariable(params, tree_config, stats_handle, name, container) + resource_handle = tree_var.resource_handle + create_op = tree_var.initializer + is_initialized_op = tree_var.is_initialized() # Adds the variable to the savable list. - saveable = TreeVariableSavable(params, resource_handle, stats_handle, - create_op, - resource_handle.name) + saveable = tree_var._gather_saveables_for_checkpoint()["tree_variable"]( # pylint: disable=protected-access + name=resource_handle.name) ops.add_to_collection(ops.GraphKeys.SAVEABLE_OBJECTS, saveable) resources.register_resource(resource_handle, create_op, is_initialized_op) return resource_handle diff --git a/tensorflow/contrib/tensor_forest/python/ops/stats_ops.py b/tensorflow/contrib/tensor_forest/python/ops/stats_ops.py index 44d486edecc4e4f7ba8a9b6d680178298813621b..9184198cd4c8fd2a7609714d094d5ef2b6868658 100644 --- a/tensorflow/contrib/tensor_forest/python/ops/stats_ops.py +++ b/tensorflow/contrib/tensor_forest/python/ops/stats_ops.py @@ -17,6 +17,8 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import functools + from tensorflow.contrib.tensor_forest.python.ops import gen_stats_ops # pylint: disable=unused-import from tensorflow.contrib.tensor_forest.python.ops.gen_stats_ops import finalize_tree @@ -25,10 +27,12 @@ from tensorflow.contrib.tensor_forest.python.ops.gen_stats_ops import process_in # pylint: enable=unused-import from tensorflow.contrib.util import loader +from tensorflow.python.eager import context from tensorflow.python.framework import ops from tensorflow.python.ops import resources from tensorflow.python.platform import resource_loader from tensorflow.python.training import saver +from tensorflow.python.training.checkpointable import tracking _stats_ops = loader.load_op_library( @@ -84,8 +88,58 @@ class FertileStatsVariableSavable(saver.BaseSaverBuilder.SaveableObject): params=self.params.serialized_params_proto) -def fertile_stats_variable(params, stats_config, name, - container=None): +class FertileStatsVariable(tracking.TrackableResource): + """A Fertile stats variable.""" + + def __init__(self, params, stats_config, name, container=None): + self._params = params + self._stats_config = stats_config + self._name = name + self._container = container + self._init_op = None + super(FertileStatsVariable, self).__init__() + self._resource_handle = self.create_resource() + + def create_resource(self): + if context.executing_eagerly(): + # TODO(allenl): This will leak memory due to kernel caching by the + # shared_name attribute value (but is better than the alternative of + # sharing everything by default when executing eagerly; hopefully creating + # tables in a loop is uncommon). + shared_name = "fertile_stats_variable_%d" % (ops.uid(),) + else: + shared_name = self._name + return gen_stats_ops.fertile_stats_resource_handle_op( + self._container, shared_name=shared_name, name=self._name) + + def initialize(self): + return gen_stats_ops.create_fertile_stats_variable( + self.resource_handle, + self._stats_config, + params=self._params.serialized_params_proto) + + @property + def initializer(self): + if self._init_op is None: + self._init_op = self.initialize() + return self._init_op + + def is_initialized(self): + return gen_stats_ops.fertile_stats_is_initialized_op(self.resource_handle) + + def _gather_saveables_for_checkpoint(self): + """For object-based checkpointing.""" + return { + "fertile_stats_variable": + functools.partial( + FertileStatsVariableSavable, + params=self._params, + stats_handle=self.resource_handle, + create_op=self.initializer) + } + + +def fertile_stats_variable(params, stats_config, name, container=None): r"""Creates a stats object and returns a handle to it. Args: @@ -98,17 +152,15 @@ def fertile_stats_variable(params, stats_config, name, A `Tensor` of type mutable `string`. The handle to the stats. """ with ops.name_scope(name, "FertileStatsVariable") as name: - resource_handle = gen_stats_ops.fertile_stats_resource_handle_op( - container, shared_name=name, name=name) - - create_op = gen_stats_ops.create_fertile_stats_variable( - resource_handle, stats_config, - params=params.serialized_params_proto) - is_initialized_op = gen_stats_ops.fertile_stats_is_initialized_op( - resource_handle) + fertile_stats_var = FertileStatsVariable(params, stats_config, name, + container) + resource_handle = fertile_stats_var.resource_handle + create_op = fertile_stats_var.initializer + is_initialized_op = fertile_stats_var.is_initialized() # Adds the variable to the savable list. - saveable = FertileStatsVariableSavable(params, resource_handle, create_op, - resource_handle.name) + saveable = ( + fertile_stats_var._gather_saveables_for_checkpoint()[ # pylint: disable=protected-access + "fertile_stats_variable"](name=resource_handle.name)) ops.add_to_collection(ops.GraphKeys.SAVEABLE_OBJECTS, saveable) resources.register_resource(resource_handle, create_op, is_initialized_op) return resource_handle diff --git a/tensorflow/contrib/tensorboard/db/BUILD b/tensorflow/contrib/tensorboard/db/BUILD deleted file mode 100644 index 6507546ee9f81108add181a9c83064c9860005e2..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/tensorboard/db/BUILD +++ /dev/null @@ -1,138 +0,0 @@ -# Description: -# TensorBoard database code. - -package(default_visibility = ["//tensorflow:internal"]) - -licenses(["notice"]) # Apache 2.0 - -load( - "//tensorflow:tensorflow.bzl", - "tf_cc_binary", - "tf_cc_test", - "tf_copts", -) - -cc_library( - name = "schema", - srcs = ["schema.cc"], - hdrs = ["schema.h"], - copts = tf_copts(), - deps = [ - "//tensorflow/core:lib", - "//tensorflow/core/lib/db:sqlite", - ], -) - -tf_cc_test( - name = "schema_test", - size = "small", - srcs = ["schema_test.cc"], - deps = [ - ":schema", - "//tensorflow/core:test", - "//tensorflow/core:test_main", - ], -) - -cc_library( - name = "summary_db_writer", - srcs = ["summary_db_writer.cc"], - hdrs = ["summary_db_writer.h"], - copts = tf_copts(), - deps = [ - ":summary_converter", - "//tensorflow/core:framework", - "//tensorflow/core:lib", - "//tensorflow/core:lib_internal", - "//tensorflow/core:protos_all_cc", - "//tensorflow/core/kernels:summary_interface", - "//tensorflow/core/lib/db:sqlite", - ], -) - -tf_cc_test( - name = "summary_db_writer_test", - size = "small", - srcs = ["summary_db_writer_test.cc"], - deps = [ - ":schema", - ":summary_db_writer", - "//tensorflow/core:lib", - "//tensorflow/core:protos_all_cc", - "//tensorflow/core:test", - "//tensorflow/core:test_main", - "//tensorflow/core/lib/db:sqlite", - ], -) - -cc_library( - name = "summary_file_writer", - srcs = ["summary_file_writer.cc"], - hdrs = ["summary_file_writer.h"], - copts = tf_copts(), - deps = [ - ":summary_converter", - "//tensorflow/core:framework", - "//tensorflow/core:lib", - "//tensorflow/core:lib_internal", - "//tensorflow/core:protos_all_cc", - "//tensorflow/core:ptr_util", - "//tensorflow/core/kernels:summary_interface", - ], -) - -tf_cc_test( - name = "summary_file_writer_test", - size = "medium", # file i/o - timeout = "short", - srcs = ["summary_file_writer_test.cc"], - deps = [ - ":summary_file_writer", - "//tensorflow/core:lib", - "//tensorflow/core:lib_internal", - "//tensorflow/core:protos_all_cc", - "//tensorflow/core:test", - "//tensorflow/core:test_main", - ], -) - -cc_library( - name = "summary_converter", - srcs = ["summary_converter.cc"], - hdrs = ["summary_converter.h"], - copts = tf_copts(), - visibility = ["//visibility:private"], - deps = [ - "//tensorflow/core:framework", - "//tensorflow/core:lib", - "//tensorflow/core:lib_internal", - "//tensorflow/core:png_internal", - "//tensorflow/core:protos_all_cc", - ], -) - -tf_cc_binary( - name = "loader", - srcs = ["loader.cc"], - linkstatic = 1, - deps = [ - ":schema", - ":summary_db_writer", - "//tensorflow/core:framework", - "//tensorflow/core:framework_internal", - "//tensorflow/core:lib", - "//tensorflow/core:protos_all_cc", - "//tensorflow/core/lib/db:sqlite", - ], -) - -tf_cc_binary( - name = "vacuum", - srcs = ["vacuum.cc"], - linkstatic = 1, - deps = [ - "//tensorflow/core:framework_internal", - "//tensorflow/core:lib", - "//tensorflow/core/lib/db:sqlite", - ], -) diff --git a/tensorflow/contrib/tensorrt/BUILD b/tensorflow/contrib/tensorrt/BUILD index c96ca302d9e3f9c71a2907098255f4a3793ad74b..67461450f8ae53739f619622de8751b654dbc082 100644 --- a/tensorflow/contrib/tensorrt/BUILD +++ b/tensorflow/contrib/tensorrt/BUILD @@ -11,56 +11,21 @@ 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", ) -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", - ]), -) +exports_files(glob([ + "test/testdata/*", +])) tf_cuda_library( name = "trt_shape_function", @@ -68,90 +33,13 @@ tf_cuda_library( 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 = [ @@ -171,8 +59,8 @@ py_library( name = "trt_ops_py", srcs_version = "PY2AND3", deps = [ - ":trt_engine_op", - ":trt_engine_op_loader", + "//tensorflow/compiler/tf2tensorrt:trt_ops", + "//tensorflow/compiler/tf2tensorrt:trt_ops_loader", ], ) @@ -201,238 +89,13 @@ tf_py_wrap_cc( "//tensorflow/python:platform/base.i", ], deps = [ - ":test_utils", - ":trt_conversion", - ":trt_engine_op_kernel", + "//tensorflow/compiler/tf2tensorrt:test_utils", + "//tensorflow/compiler/tf2tensorrt:trt_conversion", + "//tensorflow/compiler/tf2tensorrt:trt_op_kernels", "//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", - ], - 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", - ]), -) - -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( - 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", - "//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:lib", - "//tensorflow/core:protos_all_cc", - "//tensorflow/core:test", - "//tensorflow/core:test_main", - ] + 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/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"], @@ -478,12 +141,20 @@ cuda_py_tests( "test/binary_tensor_weight_broadcast_test.py", "test/concatenation_test.py", "test/const_broadcast_test.py", + "test/conv2d_test.py", + "test/dynamic_input_shapes_test.py", + "test/identity_output_test.py", + "test/int32_test.py", + "test/lru_cache_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/topk_test.py", + "test/unary_test.py", "test/vgg_block_nchw_test.py", "test/vgg_block_test.py", ], @@ -499,41 +170,44 @@ cuda_py_tests( ], ) -cuda_py_tests( - name = "tf_trt_integration_test_no_oss", - srcs = [ - "test/unary_test.py", - ], +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/117274186): re-enable in OSS after crash fixed - "no_pip", # TODO(b/117274186): re-enable in OSS after crash fixed + "no_pip", + "no_tap", # It is not able to download the mnist data. "no_windows", "nomac", ], ) -cc_library( - name = "utils", - srcs = ["convert/utils.cc"], - hdrs = ["convert/utils.h"], - copts = tf_copts(), - deps = [ - "//tensorflow/core:lib", - ], +# The following rules forward the libraries that were moved in order to not +# break other internal targets. + +alias( + name = "trt_conversion", + actual = "//tensorflow/compiler/tf2tensorrt:trt_conversion", ) -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_op_kernels", + actual = "//tensorflow/compiler/tf2tensorrt:trt_op_kernels", +) + +alias( + name = "trt_engine_op_op_lib", + actual = "//tensorflow/compiler/tf2tensorrt:trt_engine_op_op_lib", ) diff --git a/tensorflow/contrib/tensorrt/README.md b/tensorflow/contrib/tensorrt/README.md index caf8b6db0dc0a220d593f9c0afc9464ca51a1e05..a9c2ad78a3db409e6e8669c48c4df37c8db19c4b 100644 --- a/tensorflow/contrib/tensorrt/README.md +++ b/tensorflow/contrib/tensorrt/README.md @@ -1,8 +1,46 @@ -# Using TensorRT in TensorFlow +# Using TensorRT in TensorFlow (TF-TRT) -This module provides necessary bindings and introduces TRT_engine_op operator -that wraps a subgraph in TensorRT. This is still a work in progress but should -be useable with most common graphs. +This module provides necessary bindings and introduces `TRTEngineOp` operator +that wraps a subgraph in TensorRT. This module is under active development. + +## Installing TF-TRT + +Currently TensorFlow nightly builds include TF-TRT by default, which means you +don't need to install TF-TRT separately. You can pull the latest TF containers +from docker hub or install the latest TF pip package to get access to the latest +TF-TRT. + +If you want to use TF-TRT on NVIDIA Jetson platform, you can find the download +links for the relevant TensorFlow pip packages here: +https://docs.nvidia.com/deeplearning/dgx/index.html#installing-frameworks-for-jetson + +## Installing TensorRT + +In order to make use of TF-TRT, you will need a local installation of TensorRT. +Installation instructions for compatibility with TensorFlow are provided on the +[TensorFlow GPU support](https://www.tensorflow.org/install/gpu) guide. + +## Examples + +You can find example scripts for running inference on deep learning models in +this repository: https://github.com/tensorflow/tensorrt + +We have used these examples to verify the accuracy and performance of TF-TRT. +For more information see +[Verified Models](https://docs.nvidia.com/deeplearning/dgx/integrate-tf-trt/index.html#verified-models). + +## Documentation + +[TF-TRT documentation](https://docs.nvidia.com/deeplearning/dgx/integrate-tf-trt/index.html) +gives an overview of the supported functionalities, provides tutorials and +verified models, explains best practices with troubleshooting guides. + +## Tests + +TF-TRT includes both Python tests and C++ unit tests. Most of Python tests are +located in the test directory and they can be executed using `bazel test` or +directly with the Python command. Most of the C++ unit tests are used to test +the conversion functions that convert each TF op to a number of TensorRT layers. ## Compilation @@ -18,12 +56,3 @@ bazel build --config=cuda --config=opt //tensorflow/tools/pip_package:build_pip_ bazel-bin/tensorflow/tools/pip_package/build_pip_package /tmp/ ``` -After the installation of tensorflow package, TensorRT transformation will be -available. An example use can be found in test/test_tftrt.py script - -## Installing TensorRT 3.0.4 - -In order to make use of TensorRT integration, you will need a local installation -of TensorRT 3.0.4 from the [NVIDIA Developer website](https://developer.nvidia.com/tensorrt). -Installation instructions for compatibility with TensorFlow are provided on the -[TensorFlow GPU support](https://www.tensorflow.org/install/gpu) guide. diff --git a/tensorflow/contrib/tensorrt/convert/convert_graph.h b/tensorflow/contrib/tensorrt/convert/convert_graph.h deleted file mode 100644 index 3525202369841fd0b76583cdd26de2247fcdfff3..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/tensorrt/convert/convert_graph.h +++ /dev/null @@ -1,100 +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_CONVERT_CONVERT_GRAPH_H_ -#define TENSORFLOW_CONTRIB_TENSORRT_CONVERT_CONVERT_GRAPH_H_ - -#include - -#include "tensorflow/contrib/tensorrt/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" -#include "tensorflow/core/lib/core/status.h" -#include "tensorflow/core/platform/types.h" - -#if GOOGLE_CUDA -#if GOOGLE_TENSORRT - -namespace tensorflow { -namespace tensorrt { -namespace convert { - -struct ConversionParams { - ConversionParams() - : input_graph_def(nullptr), - max_batch_size(1), - max_workspace_size_bytes(1 << 30), - output_graph_def(nullptr), - precision_mode(1), - minimum_segment_size(3), - graph_properties(nullptr), - cluster(nullptr), - is_dyn_op(false), - fixed_input_size(true), - max_cached_engines(1) {} - const tensorflow::GraphDef* input_graph_def; - const std::vector* output_names; - size_t max_batch_size; - size_t max_workspace_size_bytes; - tensorflow::GraphDef* output_graph_def; - int precision_mode; - int minimum_segment_size; - const tensorflow::grappler::GraphProperties* graph_properties; - const tensorflow::grappler::Cluster* cluster; - bool is_dyn_op; // Whether to create engine on conversion or execution time - bool fixed_input_size; // Assume non-batch ranks of input tensors are fixed - int max_cached_engines; // maximum number of cached engines - std::vector cached_engine_batches; // list of cached engines -}; - -// This method extracts calibration information from the resource managers -// and puts them in to engine nodedefs. -tensorflow::Status ConvertCalibGraphToInferGraph( - const tensorflow::GraphDef& graph_def, tensorflow::GraphDef* new_graph_def, - bool is_dyn_op); - -// - max_batch_size: maximum batch size which can be used for inference for -// optimization targets inference run with max batch size. -// - max_workspace_size_bytes: The upper bound of memory allowance for engine -// building. -tensorflow::Status ConvertGraphDefToTensorRT( - const tensorflow::GraphDef& graph_def, - const std::vector& output_names, size_t max_batch_size, - size_t max_workspace_size_bytes, tensorflow::GraphDef* new_graph_def, - int precision_mode = 1, int minimum_segment_size = 3, - bool is_dyn_op = false, int max_cached_engines = 1, - std::vector cached_engine_batches = {}); - -// Method to call from optimization pass -tensorflow::Status ConvertAfterShapes(ConversionParams& params); - -// Return compile time TensorRT library version information. -std::vector GetLinkedTensorRTVersion(); - -// Return runtime time TensorRT library version information. -std::vector GetLoadedTensorRTVersion(); - -// Helper method for the conversion, expose for testing. -std::pair GetDeviceAndAllocator( - const ConversionParams& params, const EngineInfo& engine); - -} // namespace convert -} // namespace tensorrt -} // namespace tensorflow - -#endif // GOOGLE_TENSORRT -#endif // GOOGLE_CUDA - -#endif // TENSORFLOW_CONTRIB_TENSORRT_CONVERT_CONVERT_GRAPH_H_ diff --git a/tensorflow/contrib/tensorrt/convert/convert_graph_test.cc b/tensorflow/contrib/tensorrt/convert/convert_graph_test.cc deleted file mode 100644 index 8146bed4b0541ca86fee5f9402f2d606cd012047..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/tensorrt/convert/convert_graph_test.cc +++ /dev/null @@ -1,140 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/contrib/tensorrt/convert/convert_graph.h" - -#include "tensorflow/contrib/tensorrt/convert/convert_nodes.h" -#include "tensorflow/core/common_runtime/device_mgr.h" -#include "tensorflow/core/common_runtime/device_set.h" -#include "tensorflow/core/grappler/clusters/cluster.h" -#include "tensorflow/core/grappler/grappler_item.h" -#include "tensorflow/core/lib/core/status.h" -#include "tensorflow/core/lib/core/status_test_util.h" -#include "tensorflow/core/platform/test.h" -#include "tensorflow/core/protobuf/config.pb.h" // NOLINT -#include "tensorflow/core/public/session.h" - -#if GOOGLE_CUDA -#if GOOGLE_TENSORRT - -namespace tensorflow { -namespace tensorrt { -namespace convert { - -class FakeCluster : public grappler::Cluster { - public: - FakeCluster() : Cluster(0) {} - - void SetDeviceSet(const DeviceSet* device_set) { device_set_ = device_set; } - - const DeviceSet* GetDeviceSet() const override { return device_set_; } - - string type() const override { return ""; } - Status Provision() override { return Status::OK(); } - Status Initialize(const grappler::GrapplerItem& item) override { - return Status::OK(); - } - Status Run(const GraphDef& graph_def, - const std::vector>& feed, - const std::vector& fetch, - RunMetadata* metadata) override { - return Status::OK(); - } - - private: - const DeviceSet* device_set_; -}; - -TEST(ConvertGraphTest, GetDeviceAndAllocator) { - ConversionParams params; - EngineInfo engine_info; - { - // params.cluster is not set, and no gpu device is available. - auto result = GetDeviceAndAllocator(params, engine_info); - EXPECT_EQ(-1, result.first); - EXPECT_EQ(nullptr, result.second); - } - - // Create a session with two (virtual) gpu device. - SessionOptions options; - ConfigProto* config = &options.config; - GPUOptions* gpu_options = config->mutable_gpu_options(); - auto virtual_devices = - gpu_options->mutable_experimental()->add_virtual_devices(); - virtual_devices->add_memory_limit_mb(200); - virtual_devices->add_memory_limit_mb(200); - std::unique_ptr session(NewSession(options)); - - { - // params.cluster is not set, should find and return first gpu id and - // corresponding allocator. - auto result = GetDeviceAndAllocator(params, engine_info); - EXPECT_EQ(0, result.first); - EXPECT_NE(nullptr, result.second); - EXPECT_EQ("GPU_0_bfc", result.second->Name()); - } - - FakeCluster cluster; - params.cluster = &cluster; - { - // params.cluster->GetDeviceSet() returns null, should find and return first - // gpu id and corresponding allocator. - auto result = GetDeviceAndAllocator(params, engine_info); - EXPECT_EQ(0, result.first); - EXPECT_NE(nullptr, result.second); - EXPECT_EQ("GPU_0_bfc", result.second->Name()); - } - - // Build the DeviceSet. - DeviceSet device_set; - const DeviceMgr* device_mgr = nullptr; - TF_ASSERT_OK(session->LocalDeviceManager(&device_mgr)); - for (auto d : device_mgr->ListDevices()) { - device_set.AddDevice(d); - } - cluster.SetDeviceSet(&device_set); - { - // engine_info.device is not set, should find and return first gpu id and - // corresponding allocator. - auto result = GetDeviceAndAllocator(params, engine_info); - EXPECT_EQ(0, result.first); - EXPECT_NE(nullptr, result.second); - EXPECT_EQ("GPU_0_bfc", result.second->Name()); - } - - engine_info.device = "/GPU:1"; - { - // Set to use second device. - auto result = GetDeviceAndAllocator(params, engine_info); - EXPECT_EQ(0, result.first); - EXPECT_NE(nullptr, result.second); - EXPECT_EQ("GPU_1_bfc", result.second->Name()); - } - - engine_info.device = "/GPU:3"; - { - // Set to use nonexistent device. - auto result = GetDeviceAndAllocator(params, engine_info); - EXPECT_EQ(-1, result.first); - EXPECT_EQ(nullptr, result.second); - } -} - -} // namespace convert -} // namespace tensorrt -} // namespace tensorflow - -#endif // GOOGLE_TENSORRT -#endif // GOOGLE_CUDA diff --git a/tensorflow/contrib/tensorrt/convert/convert_nodes.cc b/tensorflow/contrib/tensorrt/convert/convert_nodes.cc deleted file mode 100644 index eb2dffe185aff767db0d13ac18f0bb644a16bde5..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/tensorrt/convert/convert_nodes.cc +++ /dev/null @@ -1,2766 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/contrib/tensorrt/convert/convert_nodes.h" - -#include -#include -#include -#include -#include -#include -#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 "tensorflow/core/framework/node_def.pb.h" // NOLINT -#include "tensorflow/core/framework/node_def_builder.h" -#include "tensorflow/core/framework/tensor.pb.h" // NOLINT -#include "tensorflow/core/framework/tensor_shape.pb.h" // NOLINT -#include "tensorflow/core/framework/types.h" -#include "tensorflow/core/graph/algorithm.h" -#include "tensorflow/core/graph/graph.h" -#include "tensorflow/core/graph/graph_constructor.h" -#include "tensorflow/core/lib/core/errors.h" -#include "tensorflow/core/lib/core/status.h" -#include "tensorflow/core/lib/strings/numbers.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/tensor_coding.h" -#include "tensorflow/core/platform/types.h" - -#if GOOGLE_CUDA -#if GOOGLE_TENSORRT -#include "tensorrt/include/NvInfer.h" - -// Check if the types are equal. Cast to int first so that failure log message -// would work! -#define TFTRT_CHECK_EQ_TYPE(val1, val2) CHECK_EQ((int)val1, (int)val2) - -#define TFTRT_INTERNAL_ERROR_AT_NODE(node) \ - do { \ - return tensorflow::errors::Internal( \ - "TFTRT::", __FUNCTION__, "failed to add TRT layer, at: ", node); \ - } while (0) - -#define TFTRT_RETURN_ERROR_IF_FALSE(status, node) \ - do { \ - if (status == false) { \ - TFTRT_INTERNAL_ERROR_AT_NODE(node); \ - } \ - } while (0) - -#define TFTRT_RETURN_ERROR_IF_NULLPTR(ptr, node) \ - do { \ - if (ptr == nullptr) { \ - TFTRT_INTERNAL_ERROR_AT_NODE(node); \ - } \ - } while (0) - -namespace tensorflow { -namespace tensorrt { -// TODO(aaroey): put these constants into some class. -const char* const kInputPHName = "TensorRTInputPH_"; -const char* const kOutputPHName = "TensorRTOutputPH_"; - -namespace convert { -using ::tensorflow::str_util::Split; -using ::tensorflow::strings::StrAppend; -using ::tensorflow::strings::StrCat; - -inline tensorflow::Status ConvertDType(tensorflow::DataType tf_dtype, - nvinfer1::DataType* trt_dtype) { - switch (tf_dtype) { - case tensorflow::DataType::DT_FLOAT: - *trt_dtype = nvinfer1::DataType::kFLOAT; - break; - // TODO(aaroey): this should be DT_QINT8 which is not a well supported type. - case tensorflow::DataType::DT_INT8: - *trt_dtype = nvinfer1::DataType::kINT8; - break; - case tensorflow::DataType::DT_HALF: - *trt_dtype = nvinfer1::DataType::kHALF; - break; - case tensorflow::DataType::DT_INT32: - *trt_dtype = nvinfer1::DataType::kINT32; - break; - default: - return tensorflow::errors::InvalidArgument( - "Unsupported data type ", tensorflow::DataTypeString(tf_dtype)); - } - return tensorflow::Status::OK(); -} - -void GetInputProperties(const grappler::GraphProperties& graph_properties, - const Node* outside_node, const int out_port, - PartialTensorShape* shape, - tensorflow::DataType* dtype) { - if (graph_properties.HasOutputProperties(outside_node->name())) { - auto output_params = - graph_properties.GetOutputProperties(outside_node->name()); - auto out_shape = output_params.at(out_port); - *dtype = out_shape.dtype(); - *shape = out_shape.shape(); - } else { - VLOG(0) << "Unknown output shape" << outside_node->name(); - *dtype = outside_node->output_type(out_port); - } -} - -void GetOutputProperties(const grappler::GraphProperties& graph_properties, - const Node* outside_node, const int in_port, - PartialTensorShape* shape, - tensorflow::DataType* dtype) { - if (graph_properties.HasInputProperties(outside_node->name())) { - auto input_params = - graph_properties.GetInputProperties(outside_node->name()); - auto in_shape = input_params.at(in_port); - *dtype = in_shape.dtype(); - *shape = in_shape.shape(); - } else { - *dtype = outside_node->input_type(in_port); - } -} - -tensorflow::Status ValidateInputProperties(const PartialTensorShape& shape, - const tensorflow::DataType dtype, - nvinfer1::DataType* trt_dtype) { - // TODO(aaroey): some of these checks also apply to IsTensorRTCandidate(), so - // put them there instead. - TF_RETURN_IF_ERROR(ConvertDType(dtype, trt_dtype)); - if (shape.dims() < 0) { - return tensorflow::errors::InvalidArgument("Input tensor rank is unknown."); - } - if (shape.dims() > 9) { - return tensorflow::errors::OutOfRange( - "Input tensor rank is greater than 8."); - } - for (int d = 1; d < shape.dims(); ++d) { - if (shape.dim_size(d) < 0) { - return tensorflow::errors::InvalidArgument( - "Input tensor with shape ", shape.DebugString(), - " has an unknown non-batch dimemension at dim ", d); - } - } - return Status::OK(); -} - -string DebugString(const nvinfer1::Dims& dims) { - string out = StrCat("nvinfer1::Dims(nbDims=", dims.nbDims, ", d="); - for (int i = 0; i < dims.nbDims; ++i) { - StrAppend(&out, dims.d[i], ","); - } - StrAppend(&out, ")"); - return out; -} - -string DebugString(const nvinfer1::ITensor& tensor) { - return StrCat("nvinfer1::ITensor(@", reinterpret_cast(&tensor), - ", shape=", DebugString(tensor.getDimensions()), ")"); -} - -// Return whether or not the broadcast is feasible; -bool TensorRTGetBroadcastShape(const nvinfer1::Dims& operand_l, - const bool operand_l_is_tensor, - const nvinfer1::Dims& operand_r, - const bool operand_r_is_tensor, - nvinfer1::Dims* operand_l_new_shape, - nvinfer1::Dims* operand_r_new_shape) { - // *************************************************************************** - // TensorRT Elementwise op supports broadcast but requires both tensor to be - // of Identical rank - // - // We consider case of: - // 1. operand_l to be a Tensor & operand_r to be a Const; - // 2. operand_l to be a Tensor & operand_r to be a Tensor; - // note: const op const (constant folding) should fallback to TensorFlow - // - // broadcast scheme: - // T: 1 3 5 (tensor would not have batch dimension) - // W: 1 1 3 1 (weight would have all explicit dimensions) - // i. fill in explicit dimensions - // -> T: -1 1 3 5 (we put a -1 for batch dimension) - // -> W: 1 1 3 1 - // ii. compare broadcast feasibility - // - // We cannot support the following since TensorRT does not allow manipulation - // on batch dimension, we cannot generate output with proper shape - // T: 3 5 1 - // W: 1 1 1 1 3 5 1 - // -> T: 1 1 1 -1 3 5 1 - // -> W: 1 1 1 1 3 5 1 - // *************************************************************************** - const int max_nb_dims = nvinfer1::Dims::MAX_DIMS + 1; - const size_t element_size = sizeof(operand_l.d[0]); - - // fill in dimensions - int l_s[max_nb_dims]; - std::fill(l_s, l_s + max_nb_dims, 1); - int l_d = operand_l_is_tensor ? operand_l.nbDims + 1 : operand_l.nbDims; - int r_s[max_nb_dims]; - std::fill(r_s, r_s + max_nb_dims, 1); - int r_d = operand_r_is_tensor ? operand_r.nbDims + 1 : operand_r.nbDims; - - int max_d = std::max(l_d, r_d); - std::memcpy(l_s + max_d - operand_l.nbDims, operand_l.d, - operand_l.nbDims * element_size); - std::memcpy(r_s + max_d - operand_r.nbDims, operand_r.d, - operand_r.nbDims * element_size); - - // set -1 for batch dimension, since batch size is not supposed to be - // broadcasted - if (operand_l_is_tensor) { - if (max_d != l_d) { // if broadcast beyond batch dimension, fail - return false; - } - l_s[0] = -1; - } - if (operand_r_is_tensor) { - if (max_d != r_d) { // if broadcast beyond batch dimension, fail - return false; - } - r_s[0] = -1; - } - - // compare broadcast feasibility - for (int i = max_d - 1; i >= 0; i--) { - if ((l_s[i] != r_s[i]) && (l_s[i] != 1) && (r_s[i] != 1)) { - return false; - } - } - - // output new TensorRT Dimension (stripping the batch dimension) - operand_l_new_shape->nbDims = max_d - 1; - std::memcpy(operand_l_new_shape->d, l_s + 1, (max_d - 1) * element_size); - operand_r_new_shape->nbDims = max_d - 1; - std::memcpy(operand_r_new_shape->d, r_s + 1, (max_d - 1) * element_size); - - return true; -} - -inline bool DimsEqual(const nvinfer1::Dims& dim_l, - const nvinfer1::Dims& dim_r) { - if (dim_l.nbDims != dim_r.nbDims) { - return false; - } - for (int i = 0; i < dim_l.nbDims; i++) { - if (dim_l.d[i] != dim_r.d[i]) { - return false; - } - } - return true; -} - -inline nvinfer1::Dims GetTrtDimsForTensor(const tensorflow::Tensor& tensor) { - nvinfer1::Dims dims; - dims.nbDims = tensor.dims(); - for (int i = 0; i < dims.nbDims; i++) { - dims.d[i] = tensor.dim_size(i); - } - return dims; -} - -// Returns total number of elements in dims. Returning 0 means either some dim -// is 0 or the number of dims is 0. -// Note that for TF scalar constant, we always convert to dims [1]. -int64_t TrtDimsNumElements(const nvinfer1::Dims& dims) { - if (dims.nbDims == 0) return 0; - int64_t count = 1; - for (int d = 0; d < dims.nbDims; ++d) { - count *= dims.d[d]; - } - return count; -} - -static std::vector> CreateSamePadding( - const nvinfer1::DimsHW& stride, const nvinfer1::DimsHW& kernel, - const std::vector& input_dims) { - std::vector> padding(input_dims.size()); - CHECK_EQ(stride.nbDims, input_dims.size()); // TODO(jie): N+C? NC+? - - for (size_t i = 0; i < input_dims.size(); ++i) { - // Formula to calculate the padding - int p = ((input_dims[i] - 1) / stride.d[i]) * stride.d[i] + kernel.d[i] - - input_dims[i]; - p = (p > 0) ? p : 0; - - // Right precedence padding, like in TensorFlow - int left = p / 2; - int right = p - left; - - VLOG(2) << "PADDING_" << i << " pre: " << left << ", post: " << right - << "paras: " << input_dims[i] << ", " << stride.d[i] << ", " - << "kernel: " << kernel.d[i]; - padding[i] = {left, right}; - } - return padding; -} - -string GetCommonNameScope(const string& op_name_a, const string& op_name_b) { - size_t last_scope_separator = 0; - const size_t min_size = std::min(op_name_a.size(), op_name_b.size()); - for (size_t i = 0; i < min_size; ++i) { - if (op_name_a[i] != op_name_b[i]) break; - if (op_name_a[i] == '/') last_scope_separator = i + 1; - } - return op_name_a.substr(0, last_scope_separator); -} - -TRT_ShapedWeights::TRT_ShapedWeights(tensorflow::DataType type, - const void* values, nvinfer1::Dims shape) - : shape_(shape), type_(type), values_(CHECK_NOTNULL(values)) {} - -TRT_ShapedWeights::TRT_ShapedWeights(tensorflow::DataType type) - : shape_(), type_(type), values_(nullptr) { - shape_.nbDims = 0; -} - -TRT_ShapedWeights::TRT_ShapedWeights(const TRT_ShapedWeights& rhs) - : shape_(rhs.shape_), type_(rhs.type_), values_(rhs.values_) {} - -int64_t TRT_ShapedWeights::count() const { return TrtDimsNumElements(shape_); } - -nvinfer1::Weights TRT_ShapedWeights::GetWeightsForTRT() const { - nvinfer1::DataType trt_type(nvinfer1::DataType::kFLOAT); - TF_CHECK_OK(ConvertDType(type_, &trt_type)); - return nvinfer1::Weights{trt_type, values_, values_ == nullptr ? 0 : count()}; -} - -size_t TRT_ShapedWeights::size_bytes() const { - return this->count() * tensorflow::DataTypeSize(this->type_); -} - -string TRT_ShapedWeights::DebugString() const { - return StrCat("TRT_ShapedWeights(shape=", convert::DebugString(shape_), - ", type=", type_, - ", values=", reinterpret_cast(values_), ")"); -} - -TRT_TensorOrWeights::TRT_TensorOrWeights(nvinfer1::ITensor* tensor) - : tensor_(tensor), weights_(DT_FLOAT), is_tensor_(true) {} - -TRT_TensorOrWeights::TRT_TensorOrWeights(const TRT_ShapedWeights& weights) - : tensor_(nullptr), weights_(weights), is_tensor_(false) {} - -TRT_TensorOrWeights::TRT_TensorOrWeights(const TRT_TensorOrWeights& rhs) - : tensor_(rhs.tensor_), - weights_(rhs.weights_), - is_tensor_(rhs.is_tensor_) {} - -nvinfer1::Dims TRT_TensorOrWeights::shape() const { - if (is_tensor()) { - return tensor()->getDimensions(); - } else { - return weights().shape_; - } -} - -string TRT_TensorOrWeights::DebugString() const { - string output = "TRT_TensorOrWeights(type="; - if (is_tensor()) { - StrAppend(&output, "tensor @", reinterpret_cast(tensor_), - ", shape=", convert::DebugString(tensor_->getDimensions())); - } else { - StrAppend(&output, "weights=", weights_.DebugString()); - } - StrAppend(&output, ")"); - return output; -} - -class TFAttrs { - public: - explicit TFAttrs(const tensorflow::NodeDef& tf_node) { - for (const auto& attr : tf_node.attr()) { - attrs_.insert({attr.first, &attr.second}); - } - } - - bool count(const string& key) const { return attrs_.count(key); } - - tensorflow::AttrValue const* at(const string& key) const { - if (!attrs_.count(key)) { - LOG(FATAL) << "Attribute not found: " << key; - } - return attrs_.at(key); - } - - template - T get(const string& key) const; - - template - T get(const string& key, const T& default_value) const { - return attrs_.count(key) ? this->get(key) : default_value; - } - - std::vector GetAllAttrKeys() const { - std::vector attr_list; - for (const auto& attr_item : attrs_) { - attr_list.emplace_back(attr_item.first); - } - return attr_list; - } - - private: - typedef std::map AttrMap; - AttrMap attrs_; -}; - -template <> -string TFAttrs::get(const string& key) const { - return this->at(key)->s(); -} - -template <> -std::vector TFAttrs::get>(const string& key) const { - auto attr = this->at(key)->list().i(); - return std::vector(attr.begin(), attr.end()); -} - -template <> -std::vector TFAttrs::get>(const string& key) const { - auto attr = this->at(key)->list().f(); - return std::vector(attr.begin(), attr.end()); -} - -template <> -nvinfer1::DataType TFAttrs::get(const string& key) const { - nvinfer1::DataType trt_dtype(nvinfer1::DataType::kFLOAT); - TF_CHECK_OK(ConvertDType(this->at(key)->type(), &trt_dtype)); - return trt_dtype; -} - -template <> -tensorflow::DataType TFAttrs::get( - const string& key) const { - return this->at(key)->type(); -} - -template <> -float TFAttrs::get(const string& key) const { - return this->at(key)->f(); -} - -template <> -bool TFAttrs::get(const string& key) const { - return this->at(key)->b(); -} - -// TODO(jie): reorder4 & reorder2 should be merged? -// TODO(aaroey): fix the order of parameters. -template -void Reorder4(const nvinfer1::DimsNCHW& shape, const T* idata, - const nvinfer1::DimsNCHW& istrides, T* odata, - const nvinfer1::DimsNCHW& ostrides) { - for (int n = 0; n < shape.n(); ++n) { - for (int c = 0; c < shape.c(); ++c) { - for (int h = 0; h < shape.h(); ++h) { - for (int w = 0; w < shape.w(); ++w) { - odata[n * ostrides.n() + c * ostrides.c() + h * ostrides.h() + - w * ostrides.w()] = idata[n * istrides.n() + c * istrides.c() + - h * istrides.h() + w * istrides.w()]; - } - } - } - } -} - -template -void Reorder2(const nvinfer1::DimsHW& shape, const T* idata, - const nvinfer1::DimsHW& istrides, T* odata, - const nvinfer1::DimsHW& ostrides) { - for (int h = 0; h < shape.h(); ++h) { - for (int w = 0; w < shape.w(); ++w) { - odata[h * ostrides.h() + w * ostrides.w()] = - idata[h * istrides.h() + w * istrides.w()]; - } - } -} - -// TODO(jie): fallback to tensorflow!! -void ReorderCKtoKC(const TRT_ShapedWeights& iweights, - TRT_ShapedWeights* oweights) { - const int c = iweights.shape_.d[0]; - const int k = iweights.shape_.d[1]; - oweights->shape_.d[0] = k; - oweights->shape_.d[1] = c; - const nvinfer1::DimsHW istrides = {1, k}; - const nvinfer1::DimsHW ostrides = {c, 1}; - switch (iweights.type_) { - case tensorflow::DataType::DT_FLOAT: { - Reorder2({k, c}, static_cast(iweights.GetValues()), - istrides, - // TODO(aaroey): get rid of all the const_cast like this. - static_cast(const_cast(oweights->GetValues())), - ostrides); - break; - } - case tensorflow::DataType::DT_HALF: { - Reorder2( - {k, c}, static_cast(iweights.GetValues()), - istrides, - static_cast(const_cast(oweights->GetValues())), - ostrides); - break; - } - default: - LOG(FATAL) << "Unsupported type in reorder expected fp32 or fp16 but got " - << DataTypeString(iweights.type_); - } -} - -void ReorderRSCKToKCRS(const TRT_ShapedWeights& iweights, - TRT_ShapedWeights* oweights, const int num_groups) { - CHECK_EQ(iweights.type_, oweights->type_); - CHECK_EQ(iweights.size_bytes(), oweights->size_bytes()); - // K indexes over output channels, C over input channels, and R and S over the - // height and width of the convolution - const int r = iweights.shape_.d[0]; - const int s = iweights.shape_.d[1]; - // TRT requires GKcRS, while TF depthwise has RSCK where c=1, C=G - const int c = iweights.shape_.d[2] / num_groups; - const int k = iweights.shape_.d[3] * num_groups; - VLOG(2) << "num_groups: " << num_groups - << "c" << iweights.shape_.d[2] << " then " << c - << "k" << iweights.shape_.d[3] << " then " << k - << "r" << iweights.shape_.d[0] << " then " << r - << "s" << iweights.shape_.d[1] << " then " << s; - oweights->shape_.d[0] = k / num_groups; - oweights->shape_.d[1] = c * num_groups; - oweights->shape_.d[2] = r; - oweights->shape_.d[3] = s; - const nvinfer1::DimsNCHW istrides = {1, k, s * k * c, c * k}; - const nvinfer1::DimsNCHW ostrides = {c * r * s, r * s, s, 1}; - switch (iweights.type_) { - case tensorflow::DataType::DT_FLOAT: { - Reorder4({k, c, r, s}, static_cast(iweights.GetValues()), - istrides, - static_cast(const_cast(oweights->GetValues())), - ostrides); - break; - } - case tensorflow::DataType::DT_HALF: { - Reorder4( - {k, c, r, s}, static_cast(iweights.GetValues()), - istrides, - static_cast(const_cast(oweights->GetValues())), - ostrides); - break; - } - - default: - LOG(FATAL) << "Unsupported type, expected fp32 or fp16 but got " - << DataTypeString(iweights.type_); - } -} - -Converter::Converter(nvinfer1::INetworkDefinition* trt_network, bool fp16, - int max_batch_size) - : trt_network_(trt_network), fp16_(fp16), max_batch_size_(max_batch_size) { - this->RegisterOpConverters(); -} - -TRT_ShapedWeights Converter::GetTempWeights(tensorflow::DataType type, - const nvinfer1::Dims& dims) { - const int64_t size_bytes = - TrtDimsNumElements(dims) * tensorflow::DataTypeSize(type); - // TODO(jie): check weights size_bytes. 0 means type error - weight_store_.store_.push_back(std::vector(size_bytes)); - TRT_ShapedWeights weights(type, weight_store_.store_.back().data(), dims); - return weights; -} - -tensorflow::Status Converter::ConvertNode(const tensorflow::NodeDef& node_def) { - std::vector inputs; - TF_RETURN_IF_ERROR(this->GetInputs(node_def, &inputs)); - const string& op = node_def.op(); - std::vector outputs; - if (PluginFactoryTensorRT::GetInstance()->IsPlugin(op)) { - TF_RETURN_IF_ERROR(plugin_converter_(*this, node_def, inputs, &outputs)); - } else { - if (!op_registry_.count(op)) { - return tensorflow::errors::Unimplemented( - "No converter registered for op: " + op); - } - OpConverter op_converter = op_registry_.at(op); - TF_RETURN_IF_ERROR(op_converter(*this, node_def, inputs, &outputs)); - } - 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); - // We need to check the name before setting it. For Identity op where the - // output is the input, if its input is one of the engine input, setting - // the name here will overwrite engine input bindings which will cause - // runtime error. - if (output.is_tensor()) { - const char* tensor_name = output.tensor()->getName(); - if (tensor_name == nullptr || std::strlen(tensor_name) == 0) { - output.tensor()->setName(output_name.c_str()); - } - } - VLOG(2) << "Adding out tensor " << output_name << ": " - << output.DebugString(); - if (!trt_tensors_.insert({output_name, output}).second) { - return tensorflow::errors::AlreadyExists( - "Output tensor already exists for op: " + op); - } - } - return tensorflow::Status::OK(); -} - -TRT_TensorOrWeights Converter::GetTensorOrWeights(const string& name) { - if (!trt_tensors_.count(name)) { - return TRT_TensorOrWeights(nullptr); - } - return trt_tensors_.at(name); -} - -Status Converter::AddInputTensor(const string& name, - nvinfer1::ITensor* tensor) { - if (!trt_tensors_.insert({name, TRT_TensorOrWeights(tensor)}).second) { - return errors::AlreadyExists("Input tensor already exists for op: ", name); - } - return Status::OK(); -} - -Status Converter::TransposeTensor(nvinfer1::ITensor* input_tensor, - const std::vector& order_with_batch_dim, - const nvinfer1::ITensor** output_tensor) { - const auto dims = input_tensor->getDimensions(); - - if (order_with_batch_dim.size() - 1 != size_t(dims.nbDims)) { - return tensorflow::errors::InvalidArgument( - "Rank of perm for transpose does not match with that of the input."); - } - if (order_with_batch_dim[0] != 0) { - return tensorflow::errors::Unimplemented( - "Transpose at batch dimension is not supported."); - } - - nvinfer1::IShuffleLayer* layer = this->network()->addShuffle(*input_tensor); - TFTRT_RETURN_ERROR_IF_NULLPTR(layer, "TF-TRT Internal Transpose"); - - nvinfer1::Permutation permutation; - for (int32_t i = 0; i < dims.nbDims; ++i) { - permutation.order[i] = order_with_batch_dim[i + 1] - 1; - } - layer->setFirstTranspose(permutation); - - nvinfer1::Dims reshape_dims; - reshape_dims.nbDims = dims.nbDims; - for (int32_t i = 0; i < reshape_dims.nbDims; ++i) { - reshape_dims.d[i] = 0; - // TODO(aaroey): why not transposing the types as well? - reshape_dims.type[i] = dims.type[i]; - } - layer->setReshapeDimensions(reshape_dims); - - *output_tensor = layer->getOutput(0); - return tensorflow::Status::OK(); -} - -Status Converter::PrepareTensorForShape(const TRT_TensorOrWeights& input, - const nvinfer1::Dims& dims, - const nvinfer1::ITensor** tensor) { - // If -1 is not used for one of the dims, we can check if the shapes are - // compatible. - bool can_check_shapes = true; - for (int i = 0; i < dims.nbDims; i++) { - if (dims.d[i] == -1) { - can_check_shapes = false; - break; - } - } - if (can_check_shapes && - TrtDimsNumElements(input.shape()) != TrtDimsNumElements(dims)) { - return tensorflow::errors::InvalidArgument( - "Reshape shapes are not compatible."); - } - - if (input.is_tensor()) { - if (DimsEqual(input.shape(), dims)) { - *tensor = input.tensor(); - } else { - nvinfer1::IShuffleLayer* layer = this->network()->addShuffle( - *const_cast(input.tensor())); - TFTRT_RETURN_ERROR_IF_NULLPTR(layer, "TF-TRT Internal Reshape"); - layer->setReshapeDimensions(dims); - *tensor = layer->getOutput(0); - } - } else { - nvinfer1::IConstantLayer* layer = - this->network()->addConstant(dims, input.weights()); - TFTRT_RETURN_ERROR_IF_NULLPTR(layer, "TF-TRT Internal Reshape"); - *tensor = layer->getOutput(0); - } - return tensorflow::Status::OK(); -} - -Status Converter::GetInputs(const tensorflow::NodeDef& node_def, - std::vector* inputs) const { - for (auto const& input_name : node_def.input()) { - /************************************************************************* - * TODO(jie): handle case 1) here. - * Normalizes the inputs and extracts associated metadata: - * 1) Inputs can contain a colon followed by a suffix of characters. - * That suffix may be a single number (e.g. inputName:1) or several - * word characters separated from a number by a colon - * (e.g. inputName:foo:1). The - * latter case is used to denote inputs and outputs of functions. - * 2) Control dependency inputs contain caret at the beginning and we - * remove this and annotate the edge as a control dependency. - ************************************************************************/ - // skip control nodes - if (input_name[0] == '^') continue; - string name = input_name; - auto last = name.find_last_of(':'); - // TODO(aaroey): use TensorId - if (last != string::npos && last + 2 == name.size() && - name[last + 1] == '0') { - name.erase(last); - } - - if (trt_tensors_.count(name)) { - TRT_TensorOrWeights input = trt_tensors_.at(name); - inputs->push_back(input); - VLOG(2) << "Retrieved input " << name << ": " << input.DebugString(); - } else { - // TODO(aaroey): this should not happen, make it a CHECK. - // TODO(aaroey): use StrCat for pattern like this. - string msg("Node "); - StrAppend(&msg, node_def.name(), " should have an input named '", name, - "' but it is not available"); - LOG(ERROR) << msg; - return tensorflow::errors::InvalidArgument(msg); - } - } - return tensorflow::Status::OK(); -} - -TRT_ShapedWeights ConvertFP32ToFP16(Converter& ctx, - const TRT_ShapedWeights& weights_src) { - auto dtype_new = tensorflow::DataType::DT_HALF; - TRT_ShapedWeights weights = ctx.GetTempWeights(dtype_new, weights_src.shape_); - const float* src = static_cast(weights_src.GetValues()); - Eigen::half* dst = const_cast( - static_cast(weights.GetValues())); - for (int64_t i = 0; i < weights_src.count(); i++) { - dst[i] = Eigen::half_impl::float_to_half_rtne(src[i]); - } - return weights; -} - -// **************************************************************************** -// Constant folding functions -// TODO(jie): once optimizer kicks in, we should have done constant folding -// there. -// ***************************************************************************** -struct LambdaFactory { - enum class OP_CATEGORY : int { RSQRT = 0, NEG, ADD, MUL, SUB, RECIP }; - OP_CATEGORY op; - - template - std::function unary() { - switch (op) { - case OP_CATEGORY::RSQRT: { - VLOG(2) << "RSQRT GETS DONE"; - return [](T t) -> T { return 1.0 / sqrt(t); }; - } - case OP_CATEGORY::NEG: - return [](T t) -> T { return -t; }; - case OP_CATEGORY::RECIP: - return [](T t) -> T { return 1.0 / t; }; - default: - VLOG(2) << "Not supported op for unary: " << static_cast(op); - return nullptr; - } - } - - template - std::function binary() { - switch (op) { - case OP_CATEGORY::ADD: - return [](T l, T r) -> T { return l + r; }; - case OP_CATEGORY::SUB: - return [](T l, T r) -> T { return l - r; }; - case OP_CATEGORY::MUL: - return [](T l, T r) -> T { return l * r; }; - default: - LOG(WARNING) << "Not supported op for binary: " << static_cast(op); - } - return [](T l, T r) -> T { - LOG(FATAL) << "Unsupported op type "; - return l; - }; - } - - template - std::function broadcast_r(T val) { - VLOG(2) << "LAMBDA VAL : " << val; - switch (op) { - case OP_CATEGORY::ADD: - return [val](T l) -> T { - VLOG(2) << "LAMBDA VAL : " << val; - return l + val; - }; - case OP_CATEGORY::SUB: - return [val](T l) -> T { - VLOG(2) << "LAMBDA VAL : " << val; - return l - val; - }; - case OP_CATEGORY::MUL: - return [val](T l) -> T { - VLOG(2) << "LAMBDA VAL : " << val; - return l * val; - }; - default: - LOG(WARNING) << "Not supported op for binary: " << static_cast(op); - } - return [val](T l) -> T { - LOG(FATAL) << "Unsupported op type "; - return l; - }; - } - - template - std::function broadcast_l(T val) { - VLOG(2) << "LAMBDA VAL : " << val; - switch (op) { - case OP_CATEGORY::ADD: - return [val](T l) -> T { - VLOG(2) << "LAMBDA VAL : " << val; - return val + l; - }; - case OP_CATEGORY::SUB: - return [val](T l) -> T { - VLOG(2) << "LAMBDA VAL : " << val; - return val - l; - }; - case OP_CATEGORY::MUL: - return [val](T l) -> T { - VLOG(2) << "LAMBDA VAL : " << val; - return val * l; - }; - default: - LOG(ERROR) << "Not supported op for binary: " << static_cast(op); - } - return [val](T l) -> T { - LOG(FATAL) << "Unsupported op type "; - return l; - }; - } -}; - -template <> -std::function LambdaFactory::unary() { - switch (op) { - case OP_CATEGORY::RSQRT: { - VLOG(2) << "RSQRT GETS DONE"; - return [](Eigen::half t) -> Eigen::half { - return Eigen::half(1.0 / sqrt(static_cast(t))); - }; - } - case OP_CATEGORY::NEG: - return [](Eigen::half t) -> Eigen::half { return -t; }; - // TODO(aaroey): can we support RECIP? - default: - VLOG(2) << "Not supported op for unary: " << static_cast(op); - return nullptr; - } -} - -tensorflow::Status UnaryCompute(const TRT_ShapedWeights& iweights, - TRT_ShapedWeights* oweights, - LambdaFactory unary_op) { - CHECK_EQ(iweights.type_, oweights->type_); - switch (iweights.type_) { - case tensorflow::DataType::DT_FLOAT: { - auto inp = static_cast(iweights.GetValues()); - auto oup = static_cast(const_cast(oweights->GetValues())); - std::transform(inp, inp + iweights.count(), oup, unary_op.unary()); - break; - } - case tensorflow::DataType::DT_HALF: { - auto inp = static_cast(iweights.GetValues()); - auto oup = - static_cast(const_cast(oweights->GetValues())); - std::transform(inp, inp + iweights.count(), oup, - unary_op.unary()); - break; - } - default: - return tensorflow::errors::Unimplemented( - "Data type not supported: " + - tensorflow::DataTypeString(iweights.type_)); - } - return tensorflow::Status::OK(); -} - -// TODO(jie): broadcast is needed yet not implemented. -// Only implemented channel wise for the time being -tensorflow::Status BinaryTensorOpWeight( - Converter& ctx, const tensorflow::NodeDef& node_def, - const nvinfer1::ITensor* tensor, TRT_ShapedWeights weights, - bool swapped_inputs, std::vector* outputs) { - // tensor is the left operand while weights is the right operand; - // when swapped_inputs set to true, those two are swapped. - // TODO(aaroey): use a set. - if (node_def.op() != "Sub" && node_def.op() != "Add" && - node_def.op() != "Mul" && node_def.op() != "Div" && - node_def.op() != "RealDiv") { - return tensorflow::errors::Unimplemented( - "op not supported: " + node_def.op() + ", at: " + node_def.name()); - } - - // Check type consistency - nvinfer1::DataType ttype; - TF_RETURN_IF_ERROR(ConvertDType(weights.type_, &ttype)); - - // Check scale mode - auto dims_w = weights.shape_; - auto dims_t = tensor->getDimensions(); - - // TODO(jie): addScale checks for input tensor dimension - if (dims_t.nbDims != 3) { - return tensorflow::errors::InvalidArgument( - "addScale requires tensor with rank 3, " + node_def.name()); - } - - // default to element-wise - auto scale_mode = nvinfer1::ScaleMode::kELEMENTWISE; - - // TODO(jie): maybe use a permutation instead to support more cases; - bool permutation_flag = false; - - if (weights.count() == 1) { - VLOG(2) << "UNIFORM"; - scale_mode = nvinfer1::ScaleMode::kUNIFORM; - } else { - // no broadcasting on Batch dimension; - VLOG(2) << "WEIGHTS DIM: " << dims_w.nbDims - << " tensor DIM: " << dims_t.nbDims; - if (dims_w.nbDims == dims_t.nbDims + 1) { - if (dims_w.d[0] == 1) { - for (int i = 1; i < dims_w.nbDims; i++) { - dims_w.d[i - 1] = dims_w.d[i]; - } - dims_w.nbDims--; - } else { - return tensorflow::errors::InvalidArgument( - "Binary op cannot operate on batch, " + node_def.name()); - } - } - - if (dims_w.nbDims == dims_t.nbDims && dims_w.d[0] == dims_t.d[0]) { - scale_mode = nvinfer1::ScaleMode::kELEMENTWISE; - // default is element; - for (int i = 1; i < dims_w.nbDims; i++) { - if (dims_w.d[i] != dims_t.d[i]) { - // if dimension does not match, switch back to channel; - VLOG(2) << "channel"; - scale_mode = nvinfer1::ScaleMode::kCHANNEL; - break; - } - } - // if channel as candidate, validate it - if (scale_mode == nvinfer1::ScaleMode::kCHANNEL) { - for (int i = 1; i < dims_w.nbDims; i++) { - if (dims_w.d[i] != 1) - return tensorflow::errors::InvalidArgument( - "Weight shape not compatible at, " + node_def.name()); - } - } else { - VLOG(2) << "elementwise"; - } - } else if (dims_w.nbDims == 1 && - dims_w.d[0] == dims_t.d[dims_t.nbDims - 1]) { - // channel wise and broadcast required; - permutation_flag = true; - scale_mode = nvinfer1::ScaleMode::kCHANNEL; - } else { - return tensorflow::errors::InvalidArgument( - "Weight shape not compatible at, " + node_def.name()); - } - } - - // transpose last dimension - std::vector permutation(dims_t.nbDims + 1); - if (permutation_flag) { - if (scale_mode == nvinfer1::ScaleMode::kCHANNEL && dims_t.nbDims > 1) { - // we swap the last dimension into channel for trt. - // because of tensorflow default broadcasting rules. - for (int i = 0; i < static_cast(permutation.size()); i++) { - permutation[i] = i; - } - permutation[1] = dims_t.nbDims; - permutation[dims_t.nbDims] = 1; - TF_RETURN_IF_ERROR(ctx.TransposeTensor( - const_cast(tensor), permutation, &tensor)); - } else { - return tensorflow::errors::InvalidArgument( - "Transpose cannot be applied, " + node_def.name()); - } - } - - if (ctx.IsFP16()) { - weights = ConvertFP32ToFP16(ctx, weights); - } - - // prepare weights - TRT_ShapedWeights shift_weights(weights.type_); - TRT_ShapedWeights scale_weights(weights.type_); - TRT_ShapedWeights power_weights(weights.type_); - - // Maybe I should do a switch - if (node_def.op() == "Sub") { - if (swapped_inputs) { - shift_weights = weights; - nvinfer1::IUnaryLayer* layer = - ctx.network()->addUnary(*const_cast(tensor), - nvinfer1::UnaryOperation::kNEG); - TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); - tensor = layer->getOutput(0); - } else { - TRT_ShapedWeights neg_weights = ctx.GetTempWeightsLike(weights); - LambdaFactory unary_op; - unary_op.op = LambdaFactory::OP_CATEGORY::NEG; - TF_RETURN_IF_ERROR(UnaryCompute(weights, &neg_weights, unary_op)); - shift_weights = neg_weights; - } - } else if (node_def.op() == "Div" || node_def.op() == "RealDiv") { - if (swapped_inputs) { - scale_weights = weights; - nvinfer1::IUnaryLayer* layer = - ctx.network()->addUnary(*const_cast(tensor), - nvinfer1::UnaryOperation::kRECIP); - TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); - tensor = layer->getOutput(0); - } else { - TRT_ShapedWeights recip_weights = ctx.GetTempWeightsLike(weights); - LambdaFactory unary_op; - unary_op.op = LambdaFactory::OP_CATEGORY::RECIP; - TF_RETURN_IF_ERROR(UnaryCompute(weights, &recip_weights, unary_op)); - scale_weights = recip_weights; - } - } else if (node_def.op() == "Mul") { - scale_weights = weights; - } else if (node_def.op() == "Add") { - shift_weights = weights; - } else { - return tensorflow::errors::Unimplemented("Binary op not supported: " + - node_def.op()); - } - - nvinfer1::IScaleLayer* layer = ctx.network()->addScale( - *const_cast(tensor), scale_mode, shift_weights, - scale_weights, power_weights); - TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); - - const nvinfer1::ITensor* output_tensor = layer->getOutput(0); - // transpose back dimension - if (permutation_flag) { - TF_RETURN_IF_ERROR( - ctx.TransposeTensor(const_cast(output_tensor), - permutation, &output_tensor)); - } - - // Pass the output - outputs->push_back( - TRT_TensorOrWeights(const_cast(output_tensor))); - return tensorflow::Status::OK(); -} - -enum class ConvolutionType { DEFAULT, DEPTHWISE_CONV }; - -tensorflow::Status ConvertConv2DHelper( - Converter& ctx, const tensorflow::NodeDef& node_def, - const std::vector& inputs, - std::vector* outputs, int group) { - 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") { - TF_RETURN_IF_ERROR(ctx.TransposeTensor( - const_cast(tensor), {0, 3, 1, 2}, &tensor)); - h_index = 1; - w_index = 2; - // TODO(jie): transpose it - } - - // tensor after transpose (NCHW) - 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; - - 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()); - } - if (ctx.IsFP16()) { - weights_rsck = ConvertFP32ToFP16(ctx, inputs.at(1).weights()); - } - - TRT_ShapedWeights weights = ctx.GetTempWeightsLike(weights_rsck); - ReorderRSCKToKCRS(weights_rsck, &weights, num_groups); - TRT_ShapedWeights biases(weights.type_); - const int noutput = weights.shape_.d[0] * 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]); - - 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])}); - } 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()); - auto pad_layer = ctx.network()->addPadding( - *const_cast(tensor), - nvinfer1::DimsHW(padding[0].first, padding[1].first), - nvinfer1::DimsHW(padding[0].second, padding[1].second)); - TFTRT_RETURN_ERROR_IF_NULLPTR(pad_layer, node_def.name()); - padding = {{0, 0}, {0, 0}}; - tensor = pad_layer->getOutput(0); - VLOG(2) << "TENSOR after: " << DebugString(tensor->getDimensions()); - } - - nvinfer1::IConvolutionLayer* layer = - ctx.network()->addConvolution(*const_cast(tensor), - noutput, kernel_size, weights, biases); - 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); - 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! - TF_RETURN_IF_ERROR( - ctx.TransposeTensor(const_cast(output_tensor), - {0, 2, 3, 1}, &output_tensor)); - } - outputs->push_back( - TRT_TensorOrWeights(const_cast(output_tensor))); - return tensorflow::Status::OK(); -} - -tensorflow::Status ConvertConv2DHelper( - Converter& ctx, const tensorflow::NodeDef& node_def, - const std::vector& inputs, - std::vector* outputs, ConvolutionType type) { - switch (type) { - case ConvolutionType::DEFAULT: - return ConvertConv2DHelper(ctx, node_def, inputs, outputs, 1); - case ConvolutionType::DEPTHWISE_CONV: - return ConvertConv2DHelper(ctx, node_def, inputs, outputs, 0); - } - return tensorflow::errors::Unimplemented("unsupported convolution type at, " + - node_def.name()); -} - -tensorflow::Status BinaryTensorOpTensor( - Converter& ctx, const tensorflow::NodeDef& node_def, - const TRT_TensorOrWeights& operand_l, const TRT_TensorOrWeights& operand_r, - std::vector* outputs) { - static const std::unordered_map ops{ - {"Add", nvinfer1::ElementWiseOperation::kSUM}, - {"Mul", nvinfer1::ElementWiseOperation::kPROD}, - {"Sub", nvinfer1::ElementWiseOperation::kSUB}, - {"Div", nvinfer1::ElementWiseOperation::kDIV}, - {"RealDiv", nvinfer1::ElementWiseOperation::kDIV}, - {"Minimum", nvinfer1::ElementWiseOperation::kMIN}, - {"Maximum", nvinfer1::ElementWiseOperation::kMAX}, - }; - - const nvinfer1::ITensor* tensor_l; - const nvinfer1::ITensor* tensor_r; - - nvinfer1::Dims dim_l; - nvinfer1::Dims dim_r; - - if (!TensorRTGetBroadcastShape(operand_l.shape(), operand_l.is_tensor(), - operand_r.shape(), operand_r.is_tensor(), - &dim_l, &dim_r)) { - return tensorflow::errors::InvalidArgument( - "Binary op broadcast scheme not supported by TensorRT op: " + - node_def.op() + ", at: " + node_def.name()); - } - - TF_RETURN_IF_ERROR(ctx.PrepareTensorForShape(operand_l, dim_l, &tensor_l)); - TF_RETURN_IF_ERROR(ctx.PrepareTensorForShape(operand_r, dim_r, &tensor_r)); - - // get trt type & shape - TFAttrs attrs(node_def); - // maybe this part has to be moved into the block of rsqrt later - nvinfer1::DataType dtype = attrs.get("T"); - - // check type consistency - TFTRT_CHECK_EQ_TYPE(tensor_l->getType(), dtype); - TFTRT_CHECK_EQ_TYPE(tensor_r->getType(), dtype); - auto op_pair = ops.find(node_def.op()); - if (op_pair == ops.end()) { - return tensorflow::errors::Unimplemented( - "binary op: ", node_def.op(), " not supported at: ", node_def.name()); - } - - nvinfer1::IElementWiseLayer* layer = ctx.network()->addElementWise( - // TODO(aaroey): will tensor_l/tensor_r get modified? - *const_cast(tensor_l), - *const_cast(tensor_r), op_pair->second); - TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); - - nvinfer1::ITensor* output_tensor = layer->getOutput(0); - - // pass the output - outputs->push_back(TRT_TensorOrWeights(output_tensor)); - return tensorflow::Status::OK(); -} - -tensorflow::Status ConvertPlugin(Converter& ctx, - const tensorflow::NodeDef& node_def, - const std::vector& inputs, - std::vector* outputs) { - // prepare input - std::vector all_inputs; - for (auto input : inputs) { - all_inputs.emplace_back(const_cast(input.tensor())); - } - - // plugin is owned by PluginFactory - // TODO(jie): destroy plugins later (resource management) - PluginTensorRT* plugin = - PluginFactoryTensorRT::GetInstance()->CreatePlugin(node_def.op()); - - // passing attributes - // TODO(jie): support more general attribute - TFAttrs attrs(node_def); - auto attr_key_vector = attrs.GetAllAttrKeys(); - for (auto attr_key : attr_key_vector) { - // TODO(jie): support only list of float for toy example here. - auto data = attrs.get>(attr_key); - size_t size_data = data.size() * sizeof(float); - if (!plugin->SetAttribute(attr_key, static_cast(data.data()), - size_data)) { - return tensorflow::errors::InvalidArgument("plugin SetAttribute failed"); - } - } - - nvinfer1::IPluginLayer* layer = ctx.network()->addPlugin( - &all_inputs[0], static_cast(inputs.size()), *plugin); - - for (int i = 0; i < layer->getNbOutputs(); i++) { - nvinfer1::ITensor* output_tensor = layer->getOutput(i); - outputs->push_back(TRT_TensorOrWeights(output_tensor)); - } - return tensorflow::Status::OK(); -} - -tensorflow::Status ConvertTranspose( - Converter& ctx, const tensorflow::NodeDef& node_def, - const std::vector& inputs, - std::vector* outputs) { - 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()); - } - nvinfer1::ITensor* input_tensor = - const_cast(inputs.at(0).tensor()); - - TRT_ShapedWeights weights = inputs.at(1).weights(); - const int* weights_ptr = - static_cast(const_cast(weights.GetValues())); - std::vector perm(weights.count()); - for (int i = 0; i < weights.count(); i++) { - perm[i] = weights_ptr[i]; - } - - const nvinfer1::ITensor* output_tensor = nullptr; - TF_RETURN_IF_ERROR(ctx.TransposeTensor(input_tensor, perm, &output_tensor)); - outputs->push_back( - TRT_TensorOrWeights(const_cast(output_tensor))); - return tensorflow::Status::OK(); -} - -tensorflow::Status ConvertReshape( - Converter& ctx, const tensorflow::NodeDef& node_def, - const std::vector& inputs, - std::vector* outputs) { - if (inputs.size() != 2 || !inputs.at(1).is_weights()) { - return tensorflow::errors::InvalidArgument( - "Input expects weights for shape, at ", node_def.name()); - } - - TRT_ShapedWeights weights = inputs.at(1).weights(); - if (weights.count() == 0) { - return tensorflow::errors::Unimplemented( - "Reshape to shape=[] is not supported, at ", node_def.name()); - } - - // Get new_dims - const int* weights_ptr = - static_cast(const_cast(weights.GetValues())); - nvinfer1::Dims new_dims; - // Ignore first (batch) dimension because TRT abstracts batch away - new_dims.nbDims = weights.count() - 1; - for (int i = 1; i < weights.count(); i++) { - new_dims.d[i - 1] = weights_ptr[i]; - } - - // Check that batch dimension doesn't change - const nvinfer1::Dims input_dims = inputs.at(0).shape(); - if (weights_ptr[0] == -1) { - // Product of input shape should equal product of new_dims - if (TrtDimsNumElements(input_dims) != TrtDimsNumElements(new_dims)) { - return tensorflow::errors::Unimplemented( - "Reshape on batch dimension is not supported, at ", node_def.name()); - } - } else if (weights_ptr[0] != ctx.GetMaxBatchSize()) { - return tensorflow::errors::Unimplemented( - "Reshape on batch dimension is not supported, at ", node_def.name()); - } - - const nvinfer1::ITensor* output_tensor = nullptr; - TF_RETURN_IF_ERROR( - ctx.PrepareTensorForShape(inputs.at(0), new_dims, &output_tensor)); - outputs->push_back( - TRT_TensorOrWeights(const_cast(output_tensor))); - return tensorflow::Status::OK(); -} - -tensorflow::Status ConvertConv2D(Converter& ctx, - const tensorflow::NodeDef& node_def, - const std::vector& inputs, - std::vector* outputs) { - return ConvertConv2DHelper(ctx, node_def, inputs, outputs, - ConvolutionType::DEFAULT); -} - -tensorflow::Status ConvertConv2DDepthwise( - Converter& ctx, const tensorflow::NodeDef& node_def, - const std::vector& inputs, - std::vector* outputs) { - return ConvertConv2DHelper(ctx, node_def, inputs, outputs, - ConvolutionType::DEPTHWISE_CONV); -} - -tensorflow::Status ConvertPool(Converter& ctx, - const tensorflow::NodeDef& node_def, - const std::vector& inputs, - std::vector* outputs) { - const nvinfer1::ITensor* tensor = inputs.at(0).tensor(); - TFAttrs attrs(node_def); - - int h_index = 2; - int w_index = 3; - const auto data_format = attrs.get("data_format"); - if (data_format == "NHWC") { - h_index = 1; - w_index = 2; - TF_RETURN_IF_ERROR(ctx.TransposeTensor( - const_cast(tensor), {0, 3, 1, 2}, &tensor)); - } - - nvinfer1::PoolingType type; - if (node_def.op() == "MaxPool") { - type = nvinfer1::PoolingType::kMAX; - } else if (node_def.op() == "AvgPool") { - type = nvinfer1::PoolingType::kAVERAGE; - } else { - return tensorflow::errors::Unimplemented("Unsupported pool type: ", - node_def.op()); - } - - const auto tf_stride = attrs.get>("strides"); - const nvinfer1::DimsHW stride(tf_stride[h_index], tf_stride[w_index]); - - const auto tf_kernel = attrs.get>("ksize"); - const nvinfer1::DimsHW ksize(tf_kernel[h_index], tf_kernel[w_index]); - - auto tensor_dim = tensor->getDimensions(); - std::vector> padding; - const string padding_type = attrs.get("padding"); - if (padding_type == "SAME") { - // This is NCHW tensor with no batch dimension. - // 1 -> h - // 2 -> w - padding = CreateSamePadding( - stride, ksize, - {static_cast(tensor_dim.d[1]), static_cast(tensor_dim.d[2])}); - } else if (padding_type == "VALID") { - padding = {{0, 0}, {0, 0}}; - } else { - return tensorflow::errors::Unimplemented("Unsupported padding type: ", - padding_type); - } - - if (padding[0].first != padding[0].second || - padding[1].first != padding[1].second) { - VLOG(2) << "Padding!!!: " << padding[0].first << padding[0].second - << padding[1].first << padding[1].second; - auto pad_layer = ctx.network()->addPadding( - *const_cast(tensor), - nvinfer1::DimsHW(padding[0].first, padding[1].first), - nvinfer1::DimsHW(padding[0].second, padding[1].second)); - TFTRT_RETURN_ERROR_IF_NULLPTR(pad_layer, node_def.name()); - padding = {{0, 0}, {0, 0}}; - tensor = pad_layer->getOutput(0); - } - - nvinfer1::IPoolingLayer* layer = ctx.network()->addPooling( - *const_cast(tensor), type, ksize); - 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()); - const nvinfer1::ITensor* output_tensor = layer->getOutput(0); - - if (data_format == "NHWC") { - TF_RETURN_IF_ERROR( - ctx.TransposeTensor(const_cast(output_tensor), - {0, 2, 3, 1}, &output_tensor)); - } - outputs->push_back( - TRT_TensorOrWeights(const_cast(output_tensor))); - return tensorflow::Status::OK(); -} - -tensorflow::Status ConvertActivation( - Converter& ctx, const tensorflow::NodeDef& node_def, - const std::vector& inputs, - std::vector* outputs) { - const nvinfer1::ITensor* tensor = inputs.at(0).tensor(); - nvinfer1::IActivationLayer* layer = ctx.network()->addActivation( - *const_cast(tensor), nvinfer1::ActivationType::kRELU); - TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); - nvinfer1::ITensor* output_tensor = layer->getOutput(0); - outputs->push_back(TRT_TensorOrWeights(output_tensor)); - return tensorflow::Status::OK(); -} - -tensorflow::Status ConvertScale(Converter& ctx, - const tensorflow::NodeDef& node_def, - const std::vector& inputs, - std::vector* outputs) { - if (inputs.size() != 2 || !inputs.at(0).is_tensor() || - !inputs.at(1).is_weights()) { - return tensorflow::errors::Unimplemented( - "ConvertScale only supports tensorweight: ", node_def.name()); - } - - const nvinfer1::ITensor* tensor = inputs.at(0).tensor(); - TRT_ShapedWeights weights = inputs.at(1).weights(); - if (ctx.IsFP16()) { - weights = ConvertFP32ToFP16(ctx, inputs.at(1).weights()); - } - - TRT_ShapedWeights empty_weights(weights.type_); - TFAttrs attrs(node_def); - - const auto data_format = attrs.get("data_format"); - int channel_index; - const auto dims = tensor->getDimensions(); - if (data_format == "NHWC") { - // 1). NHWC is really N+C - channel_index = dims.nbDims - 1; // batch dimension is implicit here! - } else { - // 2). NCHW is really N+CHW - channel_index = 0; // batch dimension is implicit here! - } - - nvinfer1::Permutation permutation; - for (int32_t i = 0; i < dims.nbDims; ++i) { - permutation.order[i] = i; - } - - if (channel_index >= 0) { - permutation.order[0] = channel_index; - permutation.order[channel_index] = 0; - } else { - return tensorflow::errors::Unimplemented( - "TFTRT::BiasAdd cannot apply on batch dimension, at ", node_def.name()); - } - - // TensorRT addScale requires input to be of rank 3, we need to apply - // transpose as well as reshape - if (channel_index != 0 || dims.nbDims != 3) { - nvinfer1::IShuffleLayer* shuffle_layer = - ctx.network()->addShuffle(*const_cast(tensor)); - TFTRT_RETURN_ERROR_IF_NULLPTR(shuffle_layer, node_def.name()); - nvinfer1::Dims reshape_dims; - reshape_dims.nbDims = 3; - reshape_dims.d[0] = 0; // 0 copy from the input - reshape_dims.d[1] = dims.nbDims >= 2 ? 0 : 1; // 0 copy from the input - reshape_dims.d[2] = dims.nbDims >= 3 ? -1 : 1; // -1 infer from the rest - if (channel_index != 0) { - // maybe we do not need this check. concerned about TRT optimization - shuffle_layer->setFirstTranspose(permutation); - } - shuffle_layer->setReshapeDimensions(reshape_dims); - tensor = shuffle_layer->getOutput(0); - } - - nvinfer1::ScaleMode mode = nvinfer1::ScaleMode::kCHANNEL; - if (weights.shape_.d[0] == 1) { - mode = nvinfer1::ScaleMode::kUNIFORM; - } - - nvinfer1::IScaleLayer* layer = - ctx.network()->addScale(*const_cast(tensor), mode, - weights, empty_weights, empty_weights); - TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); - - nvinfer1::ITensor* output_tensor = layer->getOutput(0); - - // restore transpose & reshape - if (channel_index != 0 || dims.nbDims != 3) { - nvinfer1::IShuffleLayer* shuffle_layer = ctx.network()->addShuffle( - *const_cast(output_tensor)); - TFTRT_RETURN_ERROR_IF_NULLPTR(shuffle_layer, node_def.name()); - nvinfer1::Dims reshape_dims = dims; - int tmp = reshape_dims.d[channel_index]; - reshape_dims.d[channel_index] = reshape_dims.d[0]; - reshape_dims.d[0] = tmp; - shuffle_layer->setReshapeDimensions(reshape_dims); - if (channel_index != 0) { - shuffle_layer->setSecondTranspose(permutation); - } - output_tensor = shuffle_layer->getOutput(0); - } - - outputs->push_back(TRT_TensorOrWeights(output_tensor)); - return tensorflow::Status::OK(); -} - -Status GetTensorDimsWithProtoShape(const Tensor& tensor, - int tensor_proto_array_len, - 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->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, Converter* converter, - TRT_ShapedWeights* weights) { - nvinfer1::Dims weight_dims; - TF_RETURN_IF_ERROR(GetTensorDimsWithProtoShape(tensor, tensor_proto_array_len, - &weight_dims)); - const int64_t size_bytes = - tensorflow::DataTypeSize(dtype) * TrtDimsNumElements(weight_dims); - converter->weight_store()->store_.push_back(std::vector(size_bytes)); - void* dst = - static_cast(&(converter->weight_store()->store_.back()[0])); - if (tensor_proto_array_len == 1) { - std::fill_n((CType*)dst, TrtDimsNumElements(weight_dims), - *tensor_proto_array); - } else { - memcpy(dst, tensor_proto_array, size_bytes); - } - *weights = TRT_ShapedWeights(dtype, dst, weight_dims); - return Status::OK(); -} - -tensorflow::Status ConvertConst(Converter& ctx, - const tensorflow::NodeDef& node_def, - const std::vector& inputs, - std::vector* outputs) { - if (!inputs.empty()) { - return errors::InvalidArgument( - "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(); - tensorflow::Tensor tensor; - if (!tensor.FromProto(tensor_proto)) { - return tensorflow::errors::Internal("Cannot parse weight tensor proto: ", - 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(), &ctx, &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(), &ctx, &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 = ctx.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()); - } - // Pass the output. - outputs->push_back(TRT_TensorOrWeights(weights)); - return tensorflow::Status::OK(); -} - -tensorflow::Status ConvertIdentity( - Converter& ctx, const tensorflow::NodeDef& node_def, - const std::vector& inputs, - std::vector* outputs) { - outputs->push_back(inputs.at(0)); - return tensorflow::Status::OK(); -} - -tensorflow::Status ConvertBinary(Converter& ctx, - const tensorflow::NodeDef& node_def, - const std::vector& inputs, - std::vector* outputs) { - if (inputs.size() != 2) { - return tensorflow::errors::FailedPrecondition( - "Binary ops require two tensor input, at ", node_def.name()); - } - - // Constant folding should have been done by TensorFlow - - if (inputs.at(0).is_weights() && inputs.at(1).is_weights()) { - return tensorflow::errors::Unimplemented( - "Constant folding is falled back to TensorFlow, binary op received " - "both input as constant at: ", - node_def.name()); - } - - // Try to convert into Scale layer first (for better performance) - // Since scale layer supports restricted broadcast policy and op types, we - // allow failure and try to handle it through Elementwise op - // (BinaryTensorOpTensor) - Status status = tensorflow::Status::OK(); - if (inputs.at(0).is_tensor() && inputs.at(1).is_weights()) { - status = BinaryTensorOpWeight(ctx, node_def, inputs.at(0).tensor(), - inputs.at(1).weights(), false, outputs); - } else if (inputs.at(0).is_weights() && inputs.at(1).is_tensor()) { - status = BinaryTensorOpWeight(ctx, node_def, inputs.at(1).tensor(), - inputs.at(0).weights(), true, outputs); - } - if ((inputs.at(0).is_tensor() && inputs.at(1).is_tensor()) || !status.ok()) { - status = BinaryTensorOpTensor(ctx, node_def, inputs.at(0), inputs.at(1), - outputs); - } - return status; -} - -tensorflow::Status ConvertUnary(Converter& ctx, - const tensorflow::NodeDef& node_def, - const std::vector& inputs, - std::vector* outputs) { - static const std::unordered_map ops{ - {"Neg", nvinfer1::UnaryOperation::kNEG}, - {"Exp", nvinfer1::UnaryOperation::kEXP}, - {"Log", nvinfer1::UnaryOperation::kLOG}, - {"Sqrt", nvinfer1::UnaryOperation::kSQRT}, - {"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()); - } - - // TODO(jie): check type - const nvinfer1::ITensor* tensor = nullptr; - TF_RETURN_IF_ERROR( - ctx.PrepareTensorForShape(inputs.at(0), inputs.at(0).shape(), &tensor)); - - nvinfer1::IUnaryLayer* layer; - if (node_def.op() == "Rsqrt") { - layer = ctx.network()->addUnary(*const_cast(tensor), - nvinfer1::UnaryOperation::kSQRT); - TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); - tensor = layer->getOutput(0); - layer = ctx.network()->addUnary(*const_cast(tensor), - nvinfer1::UnaryOperation::kRECIP); - } else if (ops.count(node_def.op()) != 0) { - layer = ctx.network()->addUnary(*const_cast(tensor), - ops.at(node_def.op())); - } else { - return tensorflow::errors::InvalidArgument( - "Binary op: ", node_def.op(), " not supported, at ", node_def.name()); - } - - TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); - nvinfer1::ITensor* output_tensor = layer->getOutput(0); - outputs->push_back( - TRT_TensorOrWeights(const_cast(output_tensor))); - return tensorflow::Status::OK(); -} - -tensorflow::Status ConvertReduce(Converter& ctx, - const tensorflow::NodeDef& node_def, - const std::vector& inputs, - std::vector* outputs) { - 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()); - } - - const nvinfer1::ITensor* tensor = inputs.at(0).tensor(); - TRT_ShapedWeights index_list = inputs.at(1).weights(); - - TFAttrs attrs(node_def); - auto index_type = attrs.get("Tidx"); - - // Only expect to handle INT32 as attributes for now - if (index_type != tensorflow::DataType::DT_INT32) { - return tensorflow::errors::Unimplemented("Tidx supports only DT_INT32"); - } - - int axes = 0; - if (index_list.count() == 0) { - return tensorflow::errors::InvalidArgument( - "TRT cannot support reduce on all (batch) dimensions, at", - node_def.name()); - } else { - auto index_list_data = - static_cast(const_cast(index_list.GetValues())); - for (int i = 0; i < index_list.count(); i++) { - int axis = index_list_data[i]; - if (axis < 0) axis += tensor->getDimensions().nbDims + 1; - if (axis == 0) { - return tensorflow::errors::InvalidArgument( - "TRT cannot reduce at batch dimension, at", node_def.name()); - } - axes |= (1 << (axis - 1)); - } - } - - nvinfer1::ReduceOperation reduce_operation; - if (node_def.op() == "Sum") { - reduce_operation = nvinfer1::ReduceOperation::kSUM; - } else if (node_def.op() == "Prod") { - reduce_operation = nvinfer1::ReduceOperation::kPROD; - } else if (node_def.op() == "Max") { - reduce_operation = nvinfer1::ReduceOperation::kMAX; - } else if (node_def.op() == "Min") { - reduce_operation = nvinfer1::ReduceOperation::kMIN; - } else if (node_def.op() == "Mean") { - reduce_operation = nvinfer1::ReduceOperation::kAVG; - } else { - return tensorflow::errors::Unimplemented("Op not supported ", node_def.op(), - " , at ", node_def.name()); - } - - const auto keep_dims = attrs.get("keep_dims"); - nvinfer1::ILayer* layer = - ctx.network()->addReduce(*const_cast(tensor), - reduce_operation, axes, keep_dims); - TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); - - outputs->push_back(TRT_TensorOrWeights(layer->getOutput(0))); - return tensorflow::Status::OK(); -} - -tensorflow::Status ConvertPad(Converter& ctx, - const tensorflow::NodeDef& node_def, - const std::vector& inputs, - std::vector* outputs) { - // 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()); - } - - // Implement tensor binaryOp weight [channel wise] for now; - const nvinfer1::ITensor* tensor = inputs.at(0).tensor(); - const auto dims = tensor->getDimensions(); - // Restore implicit batch dimension - const int nb_dims = dims.nbDims + 1; - - TRT_ShapedWeights pads = inputs.at(1).weights(); - - TFAttrs attrs(node_def); - // Padding type here is done through TF type - // so I can leverage their EnumToDataType for my cast - auto padding_type = attrs.get("Tpaddings"); - // TODO(jie): handle data type conversion for TRT? - - if (pads.shape_.d[0] != nb_dims || pads.shape_.d[1] != 2) { - return tensorflow::errors::InvalidArgument( - "Pad only supports explicit padding on 4 dimensional tensor, at ", - node_def.name()); - } - - // Only expect to handle INT32 as attributes for now - if (padding_type != tensorflow::DataType::DT_INT32) { - return tensorflow::errors::Unimplemented( - "Tpaddings supports only DT_INT32"); - } - auto pad_data = static_cast(const_cast(pads.GetValues())); - - std::vector pad_index; - for (int i = 0; i < nb_dims; i++) { - if (pad_data[2 * i] != 0 || pad_data[2 * i + 1] != 0) { - pad_index.push_back(i); - } - } - - // No padding at all, we should exit - if (pad_index.size() == 0) { - outputs->push_back(inputs.at(0)); - return tensorflow::Status::OK(); - } - - // Only supports padding on less than 2 axis GIE-2579 - if (pad_index.size() > 2) { - return tensorflow::errors::InvalidArgument( - "Padding layer does not support padding on > 2"); - } - - // Padding on batch dimension is not supported - if (pad_index[0] == 0) { - return tensorflow::errors::InvalidArgument( - "Padding layer does not support padding on batch dimension"); - } - - // Not doing the legit thing here. ignoring padding on dim 1 and 3; - // TODO(jie): implement pad as uff parser - if (pad_index.size() == 2 && pad_index[0] == 0 && pad_index[1] == 3) { - return tensorflow::errors::Unimplemented( - "Padding layer does not support padding on dimension 1 and 3 yet"); - } - - bool legit_pad = true; - nvinfer1::DimsHW pre_padding(0, 0); - nvinfer1::DimsHW post_padding(0, 0); - - std::vector permuted_pad_index(pad_index); - if (pad_index[0] == 1) { - legit_pad = false; - TF_RETURN_IF_ERROR(ctx.TransposeTensor( - const_cast(tensor), {0, 3, 2, 1}, &tensor)); - permuted_pad_index[0] = 3; - } - - for (size_t i = 0; i < pad_index.size(); i++) { - int index = pad_index[i]; - if (permuted_pad_index[i] == 2) { - pre_padding.h() = pad_data[index * 2]; - post_padding.h() = pad_data[index * 2 + 1]; - } else if (permuted_pad_index[i] == 3) { - pre_padding.w() = pad_data[index * 2]; - post_padding.w() = pad_data[index * 2 + 1]; - } - } - - nvinfer1::IPaddingLayer* layer = ctx.network()->addPadding( - *const_cast(tensor), pre_padding, post_padding); - TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); - const nvinfer1::ITensor* output_tensor = layer->getOutput(0); - - if (!legit_pad) { - TF_RETURN_IF_ERROR( - ctx.TransposeTensor(const_cast(output_tensor), - {0, 3, 2, 1}, &output_tensor)); - } - - outputs->push_back( - TRT_TensorOrWeights(const_cast(output_tensor))); - return tensorflow::Status::OK(); -} - -tensorflow::Status ConvertConcat(Converter& ctx, - const tensorflow::NodeDef& node_def, - const std::vector& inputs, - std::vector* outputs) { - // not including the last input (axis) here - int input_size = static_cast(inputs.size()) - 1; - - if (!inputs.at(0).is_tensor()) { - return tensorflow::errors::InvalidArgument( - "Concat in TRT support only Tensor input, at ", node_def.name()); - } - - // We are retrieving the axis - TRT_ShapedWeights axis = inputs.at(input_size).weights(); - - TFAttrs attrs(node_def); - auto index_type = attrs.get("Tidx"); - - // TODO(jie): handle data type - // Only expect to handle INT32 as index attributes for now - if (index_type != tensorflow::DataType::DT_INT32) - return tensorflow::errors::Unimplemented("Tidx supports only DT_INT32, at ", - node_def.name()); - - int index = *(static_cast(const_cast(axis.GetValues()))); - - // TODO(jie): early termination with no-op (attr_size==1) - - auto dim = inputs.at(0).tensor()->getDimensions(); - // dimension check - if (index > dim.nbDims + 1) { - return tensorflow::errors::InvalidArgument( - "Concatenate on axis out of dimension range, at ", node_def.name()); - } - if (index == 0) { - return tensorflow::errors::InvalidArgument( - "Concatenate on batch dimension not supported, at ", node_def.name()); - } - if (index < 0) { - index = dim.nbDims + index + 1; - } - - std::vector inputs_vec; - // Shap chack (all input tensor should have same shape) - // starting from 0 since we are probably also doing transpose here; - for (int i = 0; i < input_size; i++) { - auto tensor_i = inputs.at(i).tensor(); - auto dim_i = tensor_i->getDimensions(); - if (dim_i.nbDims != dim.nbDims) { - return tensorflow::errors::InvalidArgument( - "Concatenate receives inputs with inconsistent dimensions, at ", - node_def.name()); - } - for (int j = 0; j < dim.nbDims; j++) { - // check dimension consistency on non-concatenate axis - if (j != index - 1 && dim_i.d[j] != dim.d[j]) { - return tensorflow::errors::InvalidArgument( - "Concatenate receives inputs with inconsistent shape, at", - node_def.name()); - } - } - - inputs_vec.push_back(tensor_i); - } - - // nvinfer1::ITensor const* tensor = inputs.at(0).tensor(); - nvinfer1::IConcatenationLayer* layer = ctx.network()->addConcatenation( - const_cast(inputs_vec.data()), - inputs_vec.size()); - TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); - layer->setAxis(index - 1); - nvinfer1::ITensor* output_tensor = layer->getOutput(0); - outputs->push_back(TRT_TensorOrWeights(output_tensor)); - return tensorflow::Status::OK(); -} - -tensorflow::Status ConvertFusedBatchNorm( - Converter& ctx, const tensorflow::NodeDef& node_def, - const std::vector& inputs, - std::vector* outputs) { - TFAttrs attrs(node_def); - float epsilon = attrs.get("epsilon"); - auto data_format = attrs.get("data_format"); - if (data_format != "NCHW") { - return tensorflow::errors::Unimplemented( - "only data_format=NCHW is supported, at " + node_def.name()); - } - bool is_training = attrs.get("is_training"); - if (is_training) { - return tensorflow::errors::Unimplemented( - "only is_training=false is supported, at " + node_def.name()); - } - nvinfer1::ITensor const* tensor = inputs.at(0).tensor(); - - // Check parameter types - auto parameter_type = inputs.at(1).weights().type_; - if ((parameter_type != tensorflow::DataType::DT_FLOAT) && - (parameter_type != tensorflow::DataType::DT_HALF)) { - return tensorflow::errors::Unimplemented( - "only float32 or float16 weight data type is supported, for node " + - node_def.name() + " got " + tensorflow::DataTypeString(parameter_type)); - } - for (int i = 1; i < 5; i++) { - if (inputs.at(i).weights().type_ != parameter_type) { - return tensorflow::errors::Unimplemented( - "Inconsistent parameter type for batchnormis not supported, at: " + - node_def.name()); - } - } - - 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()); - } - TRT_ShapedWeights* ptr_shape_weights = nullptr; - for (int i = 1; i < 5; i++) { - if (inputs.at(i).weights().count() == nweight) { - ptr_shape_weights = - const_cast(&(inputs.at(i).weights())); - } else if (inputs.at(i).weights().count() != 1) { - return tensorflow::errors::InvalidArgument( - "Inconsistent batchnorm parameter count, at: " + node_def.name()); - } - } - // We could technically have two weights with different shape. - // that requires two addScale op, arguably less performant - TRT_ShapedWeights combined_scale_weights = - ctx.GetTempWeightsLike(*ptr_shape_weights); - TRT_ShapedWeights combined_offset_weights = - ctx.GetTempWeightsLike(*ptr_shape_weights); - - const Eigen::half* cast_vals_array[4]; - const float* vals_array[4]; - for (int j = 0; j < 4; j++) { - cast_vals_array[j] = - static_cast(inputs.at(j + 1).weights().GetValues()); - vals_array[j] = - static_cast(inputs.at(j + 1).weights().GetValues()); - } - Eigen::half* cast_combined_scale_vals = const_cast( - static_cast(combined_scale_weights.GetValues())); - Eigen::half* cast_combined_offset_vals = const_cast( - static_cast(combined_offset_weights.GetValues())); - float* combined_scale_vals = const_cast( - static_cast(combined_scale_weights.GetValues())); - float* combined_offset_vals = const_cast( - static_cast(combined_offset_weights.GetValues())); - - for (size_t i = 0; i < nweight; ++i) { - float batchnorm_data[4]; - for (int j = 0; j < 4; j++) { - if (inputs.at(j + 1).weights().count() != 1) { - if (parameter_type == tensorflow::DT_FLOAT) { - batchnorm_data[j] = vals_array[j][i]; - } else if (parameter_type == tensorflow::DT_HALF) { - batchnorm_data[j] = - Eigen::half_impl::half_to_float(cast_vals_array[j][i]); - } - } else { - if (parameter_type == tensorflow::DT_FLOAT) { - batchnorm_data[j] = vals_array[j][0]; - } else if (parameter_type == tensorflow::DT_HALF) { - batchnorm_data[j] = - Eigen::half_impl::half_to_float(cast_vals_array[j][0]); - } - } - } - float scale = batchnorm_data[0]; - float offset = batchnorm_data[1]; - float mean = batchnorm_data[2]; - float variance = batchnorm_data[3]; - float combined_scale_val = scale / sqrtf(variance + epsilon); - float combined_offset_val = offset - mean * combined_scale_val; - if (parameter_type == tensorflow::DT_FLOAT) { - combined_scale_vals[i] = combined_scale_val; - combined_offset_vals[i] = combined_offset_val; - } else if (parameter_type == tensorflow::DT_HALF) { - cast_combined_scale_vals[i] = Eigen::half(combined_scale_val); - cast_combined_offset_vals[i] = Eigen::half(combined_offset_val); - } - } - - nvinfer1::ScaleMode mode = nweight == 1 ? nvinfer1::ScaleMode::kUNIFORM - : nvinfer1::ScaleMode::kCHANNEL; - nvinfer1::IScaleLayer* layer = - ctx.network()->addScale(*const_cast(tensor), mode, - combined_offset_weights.GetWeightsForTRT(), - combined_scale_weights.GetWeightsForTRT(), - dummy_power_weights.GetWeightsForTRT()); - TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); - nvinfer1::ITensor* output_tensor = layer->getOutput(0); - outputs->push_back(TRT_TensorOrWeights(output_tensor)); - return tensorflow::Status::OK(); -} - -tensorflow::Status ConvertMatMulHelper( - Converter& ctx, TRT_TensorOrWeights tensor_input, - TRT_ShapedWeights weights_raw, bool transpose_weight, string node_name, - std::vector* outputs) { - nvinfer1::ITensor* output_tensor; - if (!tensor_input.is_tensor()) { - return tensorflow::errors::InvalidArgument("Input 0 expects tensor"); - } - const nvinfer1::ITensor* tensor = tensor_input.tensor(); - - TRT_ShapedWeights weights(weights_raw.type_); - if (transpose_weight) { - weights = weights_raw; - } else { - TRT_ShapedWeights weights_ck = weights_raw; - weights = ctx.GetTempWeightsLike(weights_ck); - ReorderCKtoKC(weights_raw, &weights); - } - TRT_ShapedWeights biases(weights.type_); - - int noutput = weights.shape_.d[0]; - - auto input_dim = tensor->getDimensions(); - while (input_dim.nbDims != 3) { - input_dim.d[input_dim.nbDims++] = 1; - } - TF_RETURN_IF_ERROR( - ctx.PrepareTensorForShape(tensor_input, input_dim, &tensor)); - - nvinfer1::IFullyConnectedLayer* layer = ctx.network()->addFullyConnected( - *const_cast(tensor), noutput, weights, biases); - TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_name); - output_tensor = layer->getOutput(0); - - const nvinfer1::ITensor* temp_tensor = nullptr; - auto output_dim = output_tensor->getDimensions(); - output_dim.nbDims = 1; - TF_RETURN_IF_ERROR(ctx.PrepareTensorForShape( - TRT_TensorOrWeights(output_tensor), output_dim, &temp_tensor)); - output_tensor = const_cast(temp_tensor); - outputs->push_back(TRT_TensorOrWeights(output_tensor)); - return tensorflow::Status::OK(); -} - -// inputs are both two dimensional (tensorflow::ops::MatMul) -tensorflow::Status ConvertMatMul(Converter& ctx, - const tensorflow::NodeDef& node_def, - const std::vector& inputs, - std::vector* outputs) { - if (!inputs.at(0).is_tensor()) { - return tensorflow::errors::InvalidArgument("Input 0 expects tensor, 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) { - return tensorflow::errors::Unimplemented( - "data type is not supported, for node " + node_def.name() + " got " + - tensorflow::DataTypeString(tf_dtype)); - } - bool transpose_a = attrs.get("transpose_a"); - bool transpose_b = attrs.get("transpose_b"); - - // FullyConnected: - if (transpose_a) { - return tensorflow::errors::Internal( - "Transpose_a is not supported for TensorRT FullyConnected (op: " + - node_def.op() + "), at: " + node_def.name()); - } - if (inputs.at(1).is_tensor()) { - return tensorflow::errors::Internal( - "Operand 1 must be constant for TensorRT FullyConnected (op: " + - node_def.op() + "), at: " + node_def.name()); - } - return ConvertMatMulHelper(ctx, inputs.at(0), inputs.at(1).weights(), - transpose_b, node_def.name(), outputs); -} - -tensorflow::Status ConvertBatchMatMul( - Converter& ctx, const tensorflow::NodeDef& node_def, - const std::vector& inputs, - std::vector* outputs) { - 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) { - return tensorflow::errors::Unimplemented( - "data type is not supported, for node " + node_def.name() + " got " + - tensorflow::DataTypeString(tf_dtype)); - } - - bool transpose_a = attrs.get("adj_x"); - bool transpose_b = attrs.get("adj_y"); - - auto dims = inputs.at(0).shape(); - if (dims.nbDims == 1) { // NC * CK is only supported through fully connected - if (transpose_a == false && inputs.at(0).is_tensor() && - inputs.at(1).is_weights()) { - return ConvertMatMulHelper(ctx, inputs.at(0), inputs.at(1).weights(), - transpose_b, node_def.name(), outputs); - } else { - return tensorflow::errors::InvalidArgument( - "Invalid configuration for MatMul, at: " + node_def.name()); - } - } - - const nvinfer1::ITensor* tensor_l; - const nvinfer1::ITensor* tensor_r; - auto dims_l = inputs.at(0).shape(); - auto dims_r = inputs.at(1).shape(); - if (inputs.at(0).is_weights()) { - if (inputs.at(0).shape().d[0] != 1) { - return tensorflow::errors::InvalidArgument( - "Input 0 as weight assumes broadcast across batch for MatMul, at: " + - node_def.name()); - } else { - for (int i = 0; i < dims_l.nbDims - 1; i++) { - dims_l.d[i] = dims_l.d[i + 1]; - } - dims_l.nbDims--; - } - } - if (inputs.at(1).is_weights()) { - if (inputs.at(1).shape().d[0] != 1) { - return tensorflow::errors::InvalidArgument( - "Input 1 as weight assumes broadcast across batch for MatMul, at: " + - node_def.name()); - } else { - for (int i = 0; i < dims_r.nbDims - 1; i++) { - dims_r.d[i] = dims_r.d[i + 1]; - } - dims_r.nbDims--; - } - } - TF_RETURN_IF_ERROR( - ctx.PrepareTensorForShape(inputs.at(0), dims_l, &tensor_l)); - TF_RETURN_IF_ERROR( - ctx.PrepareTensorForShape(inputs.at(1), dims_r, &tensor_r)); - - nvinfer1::IMatrixMultiplyLayer* layer = ctx.network()->addMatrixMultiply( - *const_cast(tensor_l), transpose_a, - *const_cast(tensor_r), transpose_b); - TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); - nvinfer1::ITensor* output_tensor = layer->getOutput(0); - outputs->push_back(TRT_TensorOrWeights(output_tensor)); - return tensorflow::Status::OK(); -} - -tensorflow::Status ConvertSoftmax( - Converter& ctx, const tensorflow::NodeDef& node_def, - const std::vector& inputs, - std::vector* outputs) { - const nvinfer1::ITensor* tensor = inputs.at(0).tensor(); - - int nbDims = tensor->getDimensions().nbDims; - if (nbDims == 0) { - return tensorflow::errors::InvalidArgument( - "TensorRT Softmax cannot apply on batch dimension, at" + - node_def.name()); - } - nvinfer1::ISoftMaxLayer* layer = - ctx.network()->addSoftMax(*const_cast(tensor)); - TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); - // Tensorflow SoftMax assumes applying softmax on the last dimension. - layer->setAxes(1 << (nbDims - 1)); - - nvinfer1::ITensor* output_tensor = layer->getOutput(0); - outputs->push_back(TRT_TensorOrWeights(output_tensor)); - return tensorflow::Status::OK(); -} - -tensorflow::Status ConvertTopK(Converter& ctx, - const tensorflow::NodeDef& node_def, - const std::vector& inputs, - std::vector* outputs) { - 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()); - } - - 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()); - } - - nvinfer1::ITopKLayer* layer = ctx.network()->addTopK( - *const_cast(tensor), op, k, reducedAxes); - TFTRT_RETURN_ERROR_IF_NULLPTR(layer, node_def.name()); - - nvinfer1::ITensor* output_value_tensor = layer->getOutput(0); - nvinfer1::ITensor* output_indices_tensor = layer->getOutput(1); - outputs->push_back(TRT_TensorOrWeights(output_value_tensor)); - outputs->push_back(TRT_TensorOrWeights(output_indices_tensor)); - return tensorflow::Status::OK(); -} - -void Converter::RegisterOpConverters() { - // vgg_16 slim implementation - op_registry_["Conv2D"] = ConvertConv2D; - op_registry_["DepthwiseConv2dNative"] = ConvertConv2DDepthwise; - op_registry_["Relu"] = ConvertActivation; - op_registry_["MaxPool"] = ConvertPool; - op_registry_["AvgPool"] = ConvertPool; - op_registry_["BiasAdd"] = ConvertScale; - op_registry_["Const"] = ConvertConst; - // TODO(ben,jie): this is a temp hack. - op_registry_["Identity"] = ConvertIdentity; // Identity should be removed - op_registry_["Snapshot"] = ConvertIdentity; // Snapshot should be removed - - // resnet_50_v1 slim implementation - op_registry_["Add"] = ConvertBinary; - op_registry_["Mul"] = ConvertBinary; - op_registry_["Sub"] = ConvertBinary; - op_registry_["Pad"] = ConvertPad; - - op_registry_["ConcatV2"] = ConvertConcat; - op_registry_["FusedBatchNorm"] = ConvertFusedBatchNorm; - op_registry_["FusedBatchNormV2"] = ConvertFusedBatchNorm; - - op_registry_["Div"] = ConvertBinary; - op_registry_["RealDiv"] = ConvertBinary; - - op_registry_["Rsqrt"] = ConvertUnary; - op_registry_["Reciprocal"] = ConvertUnary; - op_registry_["Exp"] = ConvertUnary; - op_registry_["Log"] = ConvertUnary; - op_registry_["Sqrt"] = ConvertUnary; - op_registry_["Abs"] = ConvertUnary; - op_registry_["Neg"] = ConvertUnary; - - op_registry_["Transpose"] = ConvertTranspose; - op_registry_["Reshape"] = ConvertReshape; - - op_registry_["Sum"] = ConvertReduce; - op_registry_["Prod"] = ConvertReduce; - op_registry_["Max"] = ConvertReduce; - op_registry_["Min"] = ConvertReduce; - op_registry_["Mean"] = ConvertReduce; - op_registry_["Maximum"] = ConvertBinary; - op_registry_["Minimum"] = ConvertBinary; - op_registry_["Softmax"] = ConvertSoftmax; - op_registry_["MatMul"] = ConvertMatMul; - 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 std::vector& input_shapes, - Logger* logger, nvinfer1::IGpuAllocator* allocator, - TRTInt8Calibrator* calibrator, - TrtUniquePtrType* engine, - bool* convert_successfully) { - engine->reset(); - if (convert_successfully) *convert_successfully = false; - - // Create the builder. - TrtUniquePtrType builder( - nvinfer1::createInferBuilder(*logger)); - builder->setMaxBatchSize(max_batch_size); - builder->setMaxWorkspaceSize(max_workspace_size_bytes); - builder->setGpuAllocator(allocator); - if (precision_mode == FP16MODE) { - builder->setHalf2Mode(true); - } else if (precision_mode == INT8MODE) { - builder->setInt8Mode(true); - builder->setInt8Calibrator(calibrator); - } - - // Create the network. - auto trt_network = - TrtUniquePtrType(builder->createNetwork()); - if (!trt_network) { - return tensorflow::errors::Internal( - "Failed to create TensorRT network object"); - } - - // Build the network - VLOG(1) << "Starting engine conversion "; - Converter converter(trt_network.get(), precision_mode == FP16MODE, - max_batch_size); - 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")) { - int32 slot_number = -1; - if (!tensorflow::strings::safe_strto32( - node_name.c_str() + strlen(kInputPHName), &slot_number)) { - return tensorflow::errors::InvalidArgument( - "Failed to parse slot number from ", node_name); - } - nvinfer1::DataType dtype; - auto shape = input_shapes.at(slot_number); - auto status = ValidateInputProperties( - shape, node_def.attr().at("dtype").type(), &dtype); - if (!status.ok()) { - const string error_message = - StrCat("Validation failed for ", node_name, " and input slot ", - slot_number, ": ", status.error_message()); - LOG(WARNING) << error_message; - return Status(status.code(), error_message); - } - - nvinfer1::Dims input_dim; - for (int i = 1; i < shape.dims(); i++) { - input_dim.d[i - 1] = shape.dim_size(i); - } - input_dim.nbDims = shape.dims() - 1; - nvinfer1::ITensor* input_tensor = - converter.network()->addInput(node_name.c_str(), dtype, input_dim); - if (!input_tensor) { - return tensorflow::errors::InvalidArgument( - "Failed to create Input layer tensor ", node_name, - " rank=", shape.dims() - 1); - } - VLOG(2) << "Adding engine input tensor " << node_name << " with shape " - << DebugString(input_dim); - TF_RETURN_IF_ERROR(converter.AddInputTensor(node_name, input_tensor)); - } else if (tensorflow::str_util::StartsWith(node_name, kOutputPHName) && - (node_def.op() == "Identity")) { - int32 slot_number = -1; - if (!tensorflow::strings::safe_strto32( - node_name.c_str() + strlen(kOutputPHName), &slot_number)) { - return tensorflow::errors::InvalidArgument( - "Failed to parse slot number from ", node_name); - } - if (output_tensors.size() <= slot_number) { - output_tensors.resize(slot_number + 1); - } - output_tensors.at(slot_number) = {node_def.input(0), node_name}; - } else { - VLOG(2) << "Converting node: " << node_def.name() << " , " - << node_def.op(); - TF_RETURN_IF_ERROR(converter.ConvertNode(node_def)); - } - } - for (const auto& output : output_tensors) { - auto tensor_or_weights = converter.GetTensorOrWeights(output.first); - if (!tensor_or_weights.is_tensor()) { - return tensorflow::errors::InvalidArgument( - "Output node '" + output.first + "' is weights not tensor"); - } - nvinfer1::ITensor* tensor = tensor_or_weights.tensor(); - tensor->setName(output.second.c_str()); - if (!tensor) { - return tensorflow::errors::NotFound("Output tensor not found: " + - output.first); - } - VLOG(1) << "Marking output tensor " << output.first << ", as output tensor " - << output.second; - - converter.network()->markOutput(*tensor); - } - if (convert_successfully) *convert_successfully = true; - - // Build the engine. - VLOG(1) << "Starting engine creation"; - engine->reset(builder->buildCudaEngine(*converter.network())); - if (engine->get() == nullptr) { - return tensorflow::errors::Internal("Failed to build TensorRT engine"); - } - VLOG(1) << "Finished conversion"; - return tensorflow::Status::OK(); -} - -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 - std::vector* connections, - tensorflow::GraphDef* segment_def, string* common_scope) { - std::set marker_nodes; - // Update connection shapes/data types and add corresponding input/output - // nodes in the segment graphdef. - for (size_t i = 0; i < connections->size(); ++i) { - auto& connection = connections->at(i); - if (connection.is_control_edge()) continue; - auto outside_node = graph->FindNodeId(connection.outside_id); - if (!outside_node) { - // This should never happen, unless the original graph is problematic. - return tensorflow::errors::NotFound( - "Cannot find node with id ", connection.outside_id, " in the graph."); - } - // Updates the shape and data types of input/output connections. - tensorflow::DataType dtype; - tensorflow::PartialTensorShape partial_shape; - if (connection.is_input_edge) { - GetInputProperties(graph_properties, - graph->FindNodeId(connection.outside_id), - connection.outside_port, &partial_shape, &dtype); - connection.outside_shape = partial_shape; - } else { - GetOutputProperties(graph_properties, - graph->FindNodeId(connection.outside_id), - connection.outside_port, &partial_shape, &dtype); - connection.inside_shape = partial_shape; - } - connection.connection_type = dtype; - - // Add dummy input/output nodes to the segment graphdef. - if (connection.is_input_edge) { - const string node_name = StrCat(kInputPHName, connection.port_number); - if (marker_nodes.count(node_name)) { - VLOG(1) << "Reusing input " << node_name << " for the edge " - << connection.outside_node_name << ":" - << connection.outside_port << " -> " - << connection.inside_node_name << ":" << connection.inside_port; - continue; - } - marker_nodes.insert(node_name); - auto seg_node = segment_def->add_node(); - tensorflow::NodeDefBuilder builder(node_name, "Placeholder"); - auto status = builder.Attr("shape", partial_shape) - .Attr("dtype", dtype) - .Finalize(seg_node); - VLOG(1) << "Constructing input " << node_name << " for the edge " - << connection.outside_node_name << ":" << connection.outside_port - << " -> " << connection.inside_node_name << ":" - << connection.inside_port; - } else { - const string node_name = StrCat(kOutputPHName, connection.port_number); - if (marker_nodes.count(node_name)) { - VLOG(1) << "Reusing output " << node_name << " for the edge " - << connection.inside_node_name << ":" << connection.inside_port - << " -> " << connection.outside_node_name << ":" - << connection.outside_port; - continue; - } - 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); - VLOG(1) << "Constructing output " << node_name << " for the edge " - << connection.inside_node_name << ":" << connection.inside_port - << " -> " << connection.outside_node_name << ":" - << connection.outside_port; - } - } // for each connection. - - 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); - local_scope = GetCommonNameScope(local_scope, node->name()); - old_to_new_id_map[node_id] = segment_def->node_size(); - auto snode = segment_def->add_node(); - snode->CopyFrom(node->def()); - VLOG(2) << "Copying " << snode->name() << " to subgraph"; - } - // Update the inputs of the new input nodes to point to placeholder nodes. - for (int i = 0; i < connections->size(); ++i) { - auto& connection = connections->at(i); - if (connection.is_control_edge() || !connection.is_input_edge) continue; - auto snode = - segment_def->mutable_node(old_to_new_id_map[connection.inside_id]); - const string placeholder_name = - StrCat(kInputPHName, connection.port_number); - VLOG(1) << "Updating " << snode->name() << ":" << connection.inside_port - << " from " << snode->input(connection.inside_port) << " to " - << placeholder_name; - snode->set_input(connection.inside_port, placeholder_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); - const int input_size = snode->input_size(); - int input_idx = 0; - int actual_input_idx = 0; - while (input_idx < input_size) { - 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)) { - if (input.second == Graph::kControlSlot) { - VLOG(1) << "... removing control inputs " << input.first - << " from subgraph."; - ++input_idx; - continue; - } else { - return tensorflow::errors::InvalidArgument( - "Found non control input outside the segment that is not an " - "engine connection to ", - snode->name(), ": ", input.first); - } - } - if (actual_input_idx != input_idx) { - snode->set_input(actual_input_idx, snode->input(input_idx)); - } - ++input_idx; - ++actual_input_idx; - } - for (int remove = input_size - actual_input_idx; remove > 0; --remove) { - snode->mutable_input()->RemoveLast(); - } - } - *common_scope = local_scope; - VLOG(0) << "Segment @scope '" << local_scope << "', converted to graph"; - return tensorflow::Status::OK(); -} - -bool InputEdgeValidator::operator()(const tensorflow::Edge* in_edge) const { - if (in_edge->IsControlEdge()) return true; - PartialTensorShape shape; - tensorflow::DataType dtype; - GetInputProperties(graph_properties_, in_edge->src(), in_edge->src_output(), - &shape, &dtype); - nvinfer1::DataType trt_dtype; - Status status = ValidateInputProperties(shape, dtype, &trt_dtype); - if (!status.ok()) { - VLOG(1) << "--> Need to remove input node " << in_edge->dst()->name() - << ": " << status; - return false; - } - - - if (in_edge->src()->type_string() != "Const" && - // Single dimensional input tensor is not supported since the first - // dimension is treated as batch dimension. - shape.dims() < 2) { - VLOG(1) << "--> Need to remove input node " << in_edge->dst()->name() - << " which has an input at port " << in_edge->dst_input() << " with" - << " #dim<2" - << " and is not a const: " << shape; - return false; - } - return true; -} - -bool OutputEdgeValidator::operator()(const tensorflow::Edge* out_edge) const { - if (out_edge->IsControlEdge()) return true; - if (out_edge->src()->type_string() == "Const") { - VLOG(1) << "--> Need to remove output node " << out_edge->src()->name() - << " which is a Const."; - return false; - } - return true; -} - -} // namespace convert -} // namespace tensorrt -} // namespace tensorflow - -#endif // GOOGLE_TENSORRT -#endif // GOOGLE_CUDA diff --git a/tensorflow/contrib/tensorrt/convert/convert_nodes.h b/tensorflow/contrib/tensorrt/convert/convert_nodes.h deleted file mode 100644 index a9b015695ffc453a066b8ab535c10b38f0a3d23c..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/tensorrt/convert/convert_nodes.h +++ /dev/null @@ -1,339 +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_CONVERT_CONVERT_NODES_H_ -#define TENSORFLOW_CONTRIB_TENSORRT_CONVERT_CONVERT_NODES_H_ - -#include -#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/contrib/tensorrt/resources/trt_resources.h" -#include "tensorflow/core/framework/graph.pb.h" -#include "tensorflow/core/graph/graph.h" -#include "tensorflow/core/grappler/costs/graph_properties.h" -#include "tensorflow/core/lib/core/status.h" - -#if GOOGLE_CUDA -#if GOOGLE_TENSORRT -#include "tensorrt/include/NvInfer.h" - -namespace tensorflow { -namespace tensorrt { -extern const char* const kInputPHName; -extern const char* const kOutputPHName; - -namespace convert { - -struct EngineConnection { - // Constructs a non-control edge. - EngineConnection(const string& outside, int out_id, int out_port, - const string& inside, int in_id, int in_port, - bool input_edge, int port) - : outside_node_name(outside), - outside_id(out_id), - outside_port(out_port), - inside_node_name(inside), - inside_id(in_id), - inside_port(in_port), - is_input_edge(input_edge), - port_number(port) {} - - // Constructs a control edge. - EngineConnection(const string& outside, int out_id, const string& inside, - int in_id, bool input_edge) - : outside_node_name(outside), - outside_id(out_id), - outside_port(Graph::kControlSlot), - inside_node_name(inside), - inside_id(in_id), - inside_port(Graph::kControlSlot), - is_input_edge(input_edge), - port_number(Graph::kControlSlot) {} - - bool is_control_edge() const { return port_number == Graph::kControlSlot; } - - const string outside_node_name; - const int outside_id; - const int outside_port; - tensorflow::PartialTensorShape outside_shape; // Only set for input edge. - - const string inside_node_name; - const int inside_id; - const int inside_port; - tensorflow::PartialTensorShape inside_shape; // Only set for output edge. - - tensorflow::DataType connection_type; - const bool is_input_edge; - - // The port number of the TRT node connected with this edge. - const int port_number; -}; - -struct EngineInfo { - EngineInfo() - : engine_type(EngineType::TRTStatic), - max_workspace_size_bytes(0), - precision_mode(FP32MODE) {} - - string engine_name; - string device; - tensorflow::GraphDef segment_graph_def; - - // Non-control input connections inside this vector are sorted in a way such - // that, the segment nodes connecting to them are topological sorted. - // In addition, for non-control connections, there must be no duplicates. - std::vector connections; - - enum class EngineType { TRTStatic = 0, TRTDynamic = 1 }; - EngineType engine_type; - int64 max_workspace_size_bytes; - int maximum_cached_engines; - std::vector cached_engine_batches; - int precision_mode; -}; - -// Constructs a graphdef from the segment in the given graph. Adds placeholder -// nodes for input edges (InputPH_*) and identity nodes for output edges -// (OutputPH_*). This function needs to be called before TensorRT nodes -// inserted in order to correctly get sizes from the original graph. -// -// - subgraph_node_names: the node names of the subgraph. -// - subgraph_node_ids: the node ids of the subgraph, must be sorted in -// topological order. -// - segment_def: the output GraphDef, whose non-input/output nodedefs will be -// sorted in topological order. -// -// TODO(aaroey): add tests to validate these properties. -tensorflow::Status ConvertSegmentToGraphDef( - const tensorflow::Graph* graph, - const tensorflow::grappler::GraphProperties& graph_properties, - const std::set& subgraph_node_names, - const std::vector& subgraph_node_ids, - std::vector* connections, - tensorflow::GraphDef* segment_def, string* common_scope); - -// Converts given subgraph to a TRT engine saved in 'engine'. Returns ok iff -// 'builder' successfully build the engine. If the result is not ok, 'engine' -// will be set to nullptr -// Once returned, 'builder' is not needed any more and can be safely detroyed. -// -// - convert_successfully: indicates whether the converson to TensorRT network -// is successful. This is different than successfully building the engine: -// building can still fail afterwards. -tensorflow::Status ConvertGraphDefToEngine( - const tensorflow::GraphDef& gdef, int precision_mode, int max_batch_size, - size_t max_workspace_size_bytes, - const std::vector& input_shapes, - Logger* logger, nvinfer1::IGpuAllocator* allocator, - TRTInt8Calibrator* calibrator, - TrtUniquePtrType* engine, - bool* convert_successfully); - -// Helper class for the segmenter to determine whether an input edge to the TRT -// segment is valid. -class InputEdgeValidator { - public: - InputEdgeValidator(const grappler::GraphProperties& graph_properties) - : graph_properties_(graph_properties) {} - - // Return true if the specified edge is eligible to be an input edge of the - // TRT segment. - bool operator()(const tensorflow::Edge* in_edge) const; - - private: - const grappler::GraphProperties& graph_properties_; -}; - -// Helper class for the segmenter to determine whether an output edge from the -// TRT segment is valid. -class OutputEdgeValidator { - public: - // Return true if the specified edge is eligible to be an output edge of the - // TRT segment. - bool operator()(const tensorflow::Edge* out_edge) const; -}; - -//////////////////////////////////////////////////////////////////////////////// -// Classes/functions below are exposed for testing purposes only. -//////////////////////////////////////////////////////////////////////////////// - -string DebugString(const nvinfer1::Dims& dims); -string DebugString(const nvinfer1::ITensor& tensor); -int64_t TrtDimsNumElements(const nvinfer1::Dims& dims); - -// Class to convert TF weight to TRT weight. -class TRT_ShapedWeights { - public: - TRT_ShapedWeights(tensorflow::DataType type, const void* values, - nvinfer1::Dims shape); - - explicit TRT_ShapedWeights(tensorflow::DataType type); - - // TODO(aaroey): use rvalue reference. - TRT_ShapedWeights(const TRT_ShapedWeights& rhs); - - nvinfer1::Weights GetWeightsForTRT() const; - - const void* GetValues() const { return values_; } - - int64_t count() const; - - size_t size_bytes() const; - - // Default converter - operator nvinfer1::Weights() const { return GetWeightsForTRT(); } - - string DebugString() const; - - // TODO(aaroey): make these private. - nvinfer1::Dims shape_; // Note: shape.type[] is not used. - tensorflow::DataType type_; - - private: - // TODO(aaroey): this should not be const as it's always from TRTWeightStore. - const void* values_; - - friend bool operator==(const TRT_ShapedWeights& lhs, - const TRT_ShapedWeights& rhs); -}; - -class TRT_TensorOrWeights { - public: - explicit TRT_TensorOrWeights(nvinfer1::ITensor* tensor); - - explicit TRT_TensorOrWeights(const TRT_ShapedWeights& weights); - - // TODO(aaroey): use rvalue reference. - TRT_TensorOrWeights(const TRT_TensorOrWeights& rhs); - - bool is_tensor() const { return is_tensor_; } - bool is_weights() const { return !is_tensor_; } - - nvinfer1::ITensor* tensor() { - CHECK(is_tensor()); - return tensor_; - } - - const nvinfer1::ITensor* tensor() const { - CHECK(is_tensor()); - return tensor_; - } - - TRT_ShapedWeights& weights() { - CHECK(is_weights()); - return weights_; - } - - const TRT_ShapedWeights& weights() const { - CHECK(is_weights()); - return weights_; - } - - // TODO(aaroey): rename to dims() to be consistent. - nvinfer1::Dims shape() const; - - string DebugString() const; - - private: - nvinfer1::ITensor* tensor_; - TRT_ShapedWeights weights_; - const bool is_tensor_; -}; - -// Class to convert TF nodes to TRT network. -class Converter { - public: - Converter(nvinfer1::INetworkDefinition* trt_network, bool fp16, - int max_batch_size); - - virtual ~Converter() {} - - nvinfer1::INetworkDefinition* network() { return trt_network_; } - - TRTWeightStore* weight_store() { return &weight_store_; } - - bool IsFP16() const { return fp16_; } - - int GetMaxBatchSize() const { return max_batch_size_; } - - TRT_ShapedWeights GetTempWeights(tensorflow::DataType type, - const nvinfer1::Dims& dims); - - TRT_ShapedWeights GetTempWeightsLike(const TRT_ShapedWeights& weights) { - return GetTempWeights(weights.type_, weights.shape_); - } - - Status ConvertNode(const tensorflow::NodeDef& node_def); - - TRT_TensorOrWeights GetTensorOrWeights(const string& name); - - Status AddInputTensor(const string& name, nvinfer1::ITensor* tensor); - - Status TransposeTensor(nvinfer1::ITensor* input_tensor, - const std::vector& order_with_batch_dim, - const nvinfer1::ITensor** output_tensor); - - // Converts input into tensor with shape specified by dims. - Status PrepareTensorForShape(const TRT_TensorOrWeights& input, - const nvinfer1::Dims& dims, - const nvinfer1::ITensor** tensor); - - // Expose for testing purposes. - Status GetInputs(const tensorflow::NodeDef& node_def, - std::vector* inputs) const; - - private: - using OpConverter = - std::function&, - std::vector*)>; - - void RegisterOpConverters(); - - std::unordered_map op_registry_; - - std::unordered_map trt_tensors_; - - OpConverter plugin_converter_; - - nvinfer1::INetworkDefinition* trt_network_; - - // TODO(aaroey): inline the definition of TRTWeightStore here, and add APIs to - // operate the stored weights instead of operating it directly. - TRTWeightStore weight_store_; - - bool fp16_; - - int max_batch_size_; - - friend class ConverterForTest; -}; - -} // namespace convert -} // namespace tensorrt -} // namespace tensorflow - -#endif // GOOGLE_TENSORRT -#endif // GOOGLE_CUDA - -#endif // TENSORFLOW_CONTRIB_TENSORRT_CONVERT_CONVERT_NODES_H_ diff --git a/tensorflow/contrib/tensorrt/convert/convert_nodes_test.cc b/tensorflow/contrib/tensorrt/convert/convert_nodes_test.cc deleted file mode 100644 index d2407fa6216ee78f4c8dbfff3613c68328a6b4bd..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/tensorrt/convert/convert_nodes_test.cc +++ /dev/null @@ -1,646 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/contrib/tensorrt/convert/convert_nodes.h" - -#include -#include -#include - -#include -#include -#include "tensorflow/contrib/tensorrt/log/trt_logger.h" -#include "tensorflow/contrib/tensorrt/plugin/trt_plugin_factory.h" -#include "tensorflow/core/framework/node_def.pb.h" // NOLINT -#include "tensorflow/core/framework/tensor.h" -#include "tensorflow/core/framework/tensor.pb.h" // NOLINT -#include "tensorflow/core/framework/tensor_shape.h" -#include "tensorflow/core/framework/tensor_testutil.h" -#include "tensorflow/core/lib/core/status.h" -#include "tensorflow/core/lib/core/status_test_util.h" -#include "tensorflow/core/platform/test.h" - -#if GOOGLE_CUDA -#if GOOGLE_TENSORRT -#include "cuda/include/cuda.h" -#include "cuda/include/cuda_runtime_api.h" -#include "tensorrt/include/NvInfer.h" - -namespace tensorflow { -namespace tensorrt { -namespace convert { - -using ::testing::ElementsAre; - -void ExpectStatus(Status status, error::Code code, const char* substr) { - EXPECT_EQ(code, status.code()) << status; - EXPECT_THAT(status.error_message(), ::testing::HasSubstr(substr)) << status; -} - -nvinfer1::Dims GetTestDims(const std::vector& d) { - nvinfer1::Dims dims; - dims.nbDims = d.size(); - for (int i = 0; i < d.size(); ++i) { - dims.d[i] = d[i]; - } - return dims; -} - -// Fake ITensor implementation for testing purposes. -class FakeITensor : public nvinfer1::ITensor { - public: - FakeITensor() {} - - FakeITensor(const nvinfer1::Dims& dims, const string& name = "") - : name_(name), dims_(dims) {} - - FakeITensor(const string& name, const std::vector& dims) - : name_(name), dims_(GetTestDims(dims)) {} - - void setName(const char* name) override { name_ = name; } - - const char* getName() const override { return name_.c_str(); } - - void setDimensions(nvinfer1::Dims dimensions) override { dims_ = dimensions; } - - nvinfer1::Dims getDimensions() const override { return dims_; } - - void setType(nvinfer1::DataType type) override { type_ = type; } - - nvinfer1::DataType getType() const override { return type_; } - - bool isNetworkInput() const override { return false; } - - bool isNetworkOutput() const override { return false; } - - void setBroadcastAcrossBatch(bool broadcastAcrossBatch) override {} - - bool getBroadcastAcrossBatch() const override { return false; } - - nvinfer1::TensorLocation getLocation() const override { return location_; } - - void setLocation(nvinfer1::TensorLocation location) override { - location_ = location; - } - -#if NV_TENSORRT_MAJOR >= 5 - bool setDynamicRange(float min, float max) override {} -#endif - - private: - string name_; - nvinfer1::Dims dims_; - nvinfer1::DataType type_; - nvinfer1::TensorLocation location_; -}; - -bool Equals(const nvinfer1::Dims& lhs, const nvinfer1::Dims& rhs) { - if (lhs.nbDims != rhs.nbDims) return false; - for (int i = 0; i < lhs.nbDims; ++i) { - if (lhs.d[i] != rhs.d[i]) return false; - // We don't check the types in the tests. - } - return true; -} - -bool operator==(const TRT_ShapedWeights& lhs, const TRT_ShapedWeights& rhs) { - return Equals(lhs.shape_, rhs.shape_) && lhs.type_ == rhs.type_ && - lhs.values_ == rhs.values_; -} - -TEST(TRT_ShapedWeights_Test, Basic) { - { - float raw_weights[10]; - TRT_ShapedWeights weights(DT_FLOAT, raw_weights, GetTestDims({2, 5})); - - nvinfer1::Weights trt_weights = weights.GetWeightsForTRT(); - EXPECT_EQ(nvinfer1::DataType::kFLOAT, trt_weights.type); - EXPECT_EQ(static_cast(raw_weights), trt_weights.values); - EXPECT_EQ(10, trt_weights.count); - - EXPECT_EQ(static_cast(raw_weights), weights.GetValues()); - EXPECT_EQ(10, weights.count()); - EXPECT_EQ(40, weights.size_bytes()); - } - { - int32 raw_weights = 0; - TRT_ShapedWeights weights(DT_INT32, &raw_weights, GetTestDims({1, 1, 1})); - - nvinfer1::Weights trt_weights = weights.GetWeightsForTRT(); - EXPECT_EQ(nvinfer1::DataType::kINT32, trt_weights.type); - EXPECT_EQ(static_cast(&raw_weights), trt_weights.values); - EXPECT_EQ(1, trt_weights.count); - - EXPECT_EQ(static_cast(&raw_weights), weights.GetValues()); - EXPECT_EQ(1, weights.count()); - EXPECT_EQ(4, weights.size_bytes()); - } - { - TRT_ShapedWeights weights(DT_FLOAT); - - nvinfer1::Weights trt_weights = weights.GetWeightsForTRT(); - EXPECT_EQ(nvinfer1::DataType::kFLOAT, trt_weights.type); - EXPECT_EQ(nullptr, trt_weights.values); - EXPECT_EQ(0, trt_weights.count); - - EXPECT_EQ(nullptr, weights.GetValues()); - EXPECT_EQ(0, weights.count()); - EXPECT_EQ(0, weights.size_bytes()); - } -} - -TEST(TRT_TensorOrWeights_Test, Basic) { - { - nvinfer1::Dims dims; - dims.nbDims = 1; - dims.d[0] = 1; - FakeITensor itensor(dims); - - TRT_TensorOrWeights tw(&itensor); - EXPECT_EQ(true, tw.is_tensor()); - EXPECT_EQ(false, tw.is_weights()); - EXPECT_EQ(&itensor, tw.tensor()); - EXPECT_TRUE(Equals(dims, tw.shape())) - << "- expected: " << DebugString(dims) - << "\n vs\n- actual: " << DebugString(tw.shape()); - } - { - TRT_ShapedWeights weights(DT_FLOAT); - TRT_TensorOrWeights tw(weights); - EXPECT_EQ(false, tw.is_tensor()); - EXPECT_EQ(true, tw.is_weights()); - EXPECT_EQ(weights, tw.weights()); - - nvinfer1::Dims dims; - dims.nbDims = 0; - EXPECT_TRUE(Equals(dims, tw.shape())) - << "- expected: " << DebugString(dims) - << "\n vs\n- actual: " << DebugString(tw.shape()); - } -} - -class ConverterForTest : public Converter { - public: - ConverterForTest() - : Converter(nullptr, /*fp16=*/false, /*max_batch_size=*/1) { - QCHECK_EQ(0, cudaStreamCreate(&stream_)); - Reset(); - } - - ~ConverterForTest() override { QCHECK_EQ(0, cudaStreamDestroy(stream_)); } - - // Helper methods for testing purposes. - - void AddOpConverter(const string& op_name, OpConverter op_converter) { - op_registry_[op_name] = op_converter; - } - - void AddTensorOrWeights(const string& name, TRT_TensorOrWeights tw) { - ASSERT_TRUE(trt_tensors_.insert({name, tw}).second); - } - - void Reset() { - // Clear the tensor map. - trt_tensors_.clear(); - // Reset the INetworkDefinition. - engine_.reset(nullptr); - network_.reset(nullptr); - builder_.reset(nullptr); - builder_.reset(nvinfer1::createInferBuilder(logger_)); - network_.reset(builder_->createNetwork()); - trt_network_ = network_.get(); - } - - void BuildAndRun(const char* input_name, const std::vector& input_data, - const char* output_name, std::vector* output_data) { - // Mark the output tensor as TRT engine output. - TRT_TensorOrWeights tensor = GetTensorOrWeights(output_name); - tensor.tensor()->setName(output_name); - network()->markOutput(*tensor.tensor()); - - // Build the TRT engine. - QCHECK_EQ(nullptr, engine_.get()); - engine_.reset(builder_->buildCudaEngine(*network())); - CHECK_NOTNULL(engine_.get()); - - // Execute the TRT engine. - const int input_size = input_data.size() * sizeof(float); - const int output_size = output_data->size() * sizeof(float); - const int input_index = engine_->getBindingIndex(input_name); - const int output_index = engine_->getBindingIndex(output_name); - - ASSERT_EQ(engine_->getNbBindings(), 2); - void* buffers[2]; - ASSERT_EQ(0, cudaMalloc(&buffers[input_index], input_size)); - ASSERT_EQ(0, cudaMalloc(&buffers[output_index], output_size)); - ASSERT_EQ(0, cudaMemcpyAsync(buffers[input_index], input_data.data(), - input_size, cudaMemcpyHostToDevice, stream_)); - TrtUniquePtrType execution_context( - engine_->createExecutionContext()); - execution_context->enqueue(1, buffers, stream_, nullptr); - ASSERT_EQ(0, cudaMemcpyAsync(output_data->data(), buffers[output_index], - output_size, cudaMemcpyDeviceToHost, stream_)); - cudaStreamSynchronize(stream_); - ASSERT_EQ(0, cudaFree(buffers[input_index])); - ASSERT_EQ(0, cudaFree(buffers[output_index])); - } - - private: - Logger logger_; - TrtUniquePtrType builder_; - TrtUniquePtrType network_; - TrtUniquePtrType engine_; - cudaStream_t stream_; -}; - -class ConverterTest : public ::testing::Test { - protected: - nvinfer1::ITensor* AddTestTensor(const char* name, - const std::vector& dims) { - nvinfer1::ITensor* tensor = converter_.network()->addInput( - name, nvinfer1::DataType::kFLOAT, GetTestDims(dims)); - converter_.AddTensorOrWeights(name, TRT_TensorOrWeights{tensor}); - return tensor; - } - - template - TRT_ShapedWeights AddTestWeights(const char* name, const DataType dtype, - const std::vector& dims, - const std::vector& values) { - const nvinfer1::Dims trt_dims = GetTestDims(dims); - const int64_t num_elements = TrtDimsNumElements(trt_dims); - QCHECK_EQ(num_elements, values.size()) - << num_elements << " vs " << values.size(); - TRT_ShapedWeights weights(dtype); - if (num_elements) { - const int64_t size_bytes = DataTypeSize(dtype) * num_elements; - QCHECK_EQ(size_bytes, sizeof(CType) * values.size()) - << size_bytes << " vs " << sizeof(CType) * values.size(); - converter_.weight_store()->store_.push_back( - std::vector(size_bytes)); - void* dst = - static_cast(converter_.weight_store()->store_.back().data()); - memcpy(dst, values.data(), size_bytes); - weights = TRT_ShapedWeights(dtype, dst, trt_dims); - } - converter_.AddTensorOrWeights(name, TRT_TensorOrWeights{weights}); - return weights; - } - - NodeDef MakeNodeDef(const string& name, const string& op, - const std::vector& inputs) { - NodeDef node_def; - node_def.set_name(name); - node_def.set_op(op); - for (const string& input : inputs) { - node_def.add_input(input); - } - return node_def; - } - - ConverterForTest converter_; -}; - -TEST_F(ConverterTest, GetTempWeights) { - TRT_ShapedWeights weights = - converter_.GetTempWeights(DT_FLOAT, GetTestDims({2, 3})); - - nvinfer1::Weights trt_weights = weights.GetWeightsForTRT(); - EXPECT_EQ(nvinfer1::DataType::kFLOAT, trt_weights.type); - EXPECT_NE(nullptr, trt_weights.values); - EXPECT_EQ(6, trt_weights.count); - - EXPECT_NE(nullptr, weights.GetValues()); - EXPECT_EQ(6, weights.count()); - EXPECT_EQ(24, weights.size_bytes()); - - // TODO(aaroey): test the case where shape element count is 0. -} - -TEST_F(ConverterTest, GetInputs) { - NodeDef node_def; - node_def.add_input("^control_input"); - node_def.add_input("input"); - node_def.add_input("input:0"); - node_def.add_input("input:1"); - node_def.add_input("weird_input:2:3:4:0"); - - FakeITensor input, input_1, input_2; - TF_EXPECT_OK(converter_.AddInputTensor("input", &input)); - TF_EXPECT_OK(converter_.AddInputTensor("input:1", &input_1)); - TF_EXPECT_OK(converter_.AddInputTensor("weird_input:2:3:4", &input_2)); - - std::vector inputs; - TF_EXPECT_OK(converter_.GetInputs(node_def, &inputs)); - EXPECT_EQ(4, inputs.size()); - EXPECT_EQ(&input, inputs[0].tensor()); - EXPECT_EQ(&input, inputs[1].tensor()); - EXPECT_EQ(&input_1, inputs[2].tensor()); - EXPECT_EQ(&input_2, inputs[3].tensor()); -} - -TEST_F(ConverterTest, ConvertNode) { - FakeITensor output_tensors[2]; - auto op_converter = [&output_tensors]( - Converter& ctx, const NodeDef& node_def, - const std::vector& inputs, - std::vector* outputs) -> Status { - nvinfer1::Dims dims = inputs[0].tensor()->getDimensions(); - for (int i = 0; i < 2; ++i) { - dims.d[0] += 1; - output_tensors[i].setDimensions(dims); - outputs->push_back(TRT_TensorOrWeights(&output_tensors[i])); - } - return Status::OK(); - }; - converter_.AddOpConverter("MyOp", op_converter); - - FakeITensor input_tensor("my_input", {12345}); - TF_EXPECT_OK(converter_.AddInputTensor("my_input", &input_tensor)); - - NodeDef node_def = MakeNodeDef("my_op", "MyOp", {"my_input"}); - TF_EXPECT_OK(converter_.ConvertNode(node_def)); - - TRT_TensorOrWeights actual_output_1 = converter_.GetTensorOrWeights("my_op"); - EXPECT_EQ(&output_tensors[0], actual_output_1.tensor()); - EXPECT_EQ(12346, actual_output_1.tensor()->getDimensions().d[0]); - - TRT_TensorOrWeights actual_output_2 = - converter_.GetTensorOrWeights("my_op:1"); - EXPECT_EQ(&output_tensors[1], actual_output_2.tensor()); - EXPECT_EQ(12347, actual_output_2.tensor()->getDimensions().d[0]); -} - -TEST_F(ConverterTest, TransposeTensor) { - nvinfer1::ITensor* input_tensor = AddTestTensor("", {2, 3, 5}); - const nvinfer1::ITensor* output_tensor = nullptr; - - // Rank doesn't match. - ExpectStatus( - converter_.TransposeTensor(input_tensor, {0, 1}, &output_tensor), - error::INVALID_ARGUMENT, - "Rank of perm for transpose does not match with that of the input"); - - // Transpose at batch dimension. - ExpectStatus( - converter_.TransposeTensor(input_tensor, {1, 0, 2, 3}, &output_tensor), - error::UNIMPLEMENTED, "Transpose at batch dimension is not supported."); - - // OK. - TF_EXPECT_OK( - converter_.TransposeTensor(input_tensor, {0, 3, 1, 2}, &output_tensor)); - EXPECT_TRUE(Equals(GetTestDims({5, 2, 3}), output_tensor->getDimensions())) - << DebugString(*output_tensor); -} - -TEST_F(ConverterTest, PrepareTensorForShape_Tensor) { - nvinfer1::ITensor* input_tensor = AddTestTensor("", {2, 3, 5}); - TRT_TensorOrWeights tw(input_tensor); - const nvinfer1::ITensor* output_tensor = nullptr; - - // Shape size doesn't match. - ExpectStatus(converter_.PrepareTensorForShape(tw, GetTestDims({2, 3, 6}), - &output_tensor), - error::INVALID_ARGUMENT, "Reshape shapes are not compatible."); - - // TODO(aaroey): we should check the case where uninferred dimensions are not - // an exact divisor of input dim ensions, e.g. for dims {-1, 7}. - - // Infer shape, ok. - TF_EXPECT_OK(converter_.PrepareTensorForShape(tw, GetTestDims({-1, 2}), - &output_tensor)); - EXPECT_TRUE(Equals(GetTestDims({15, 2}), output_tensor->getDimensions())) - << DebugString(*output_tensor); - - // Regular shape. - TF_EXPECT_OK(converter_.PrepareTensorForShape(tw, GetTestDims({10, 3}), - &output_tensor)); - EXPECT_TRUE(Equals(GetTestDims({10, 3}), output_tensor->getDimensions())) - << DebugString(*output_tensor); -} - -#if NV_TENSORRT_MAJOR > 3 -TEST_F(ConverterTest, PrepareTensorForShape_Weights) { - TRT_ShapedWeights weights = - converter_.GetTempWeights(DT_FLOAT, GetTestDims({2, 3, 5})); - TRT_TensorOrWeights tw(weights); - const nvinfer1::ITensor* output_tensor = nullptr; - TF_EXPECT_OK(converter_.PrepareTensorForShape(tw, GetTestDims({10, 3}), - &output_tensor)); - EXPECT_TRUE(Equals(GetTestDims({10, 3}), output_tensor->getDimensions())) - << DebugString(*output_tensor); -} -#endif - -template -void TestConvertConst(ConverterForTest* converter) { - NodeDef node_def; - node_def.set_name("my_const"); - node_def.set_op("Const"); - - auto reset_and_test = [&node_def, converter]( - const Tensor& tensor, const bool as_tensor_content, - const std::vector& expected_dims, - const std::vector& expected_value) { - converter->Reset(); - - auto& attr = *node_def.mutable_attr(); - if (as_tensor_content) { - tensor.AsProtoTensorContent(attr["value"].mutable_tensor()); - } else { - tensor.AsProtoField(attr["value"].mutable_tensor()); - } - TF_EXPECT_OK(converter->ConvertNode(node_def)); - TRT_TensorOrWeights output = converter->GetTensorOrWeights("my_const"); - EXPECT_TRUE(Equals(GetTestDims(expected_dims), output.weights().shape_)) - << output.DebugString(); - ASSERT_EQ(expected_value.size(), output.weights().count()) - << output.DebugString(); - const OutputCType* actual_values = - static_cast(output.weights().GetValues()); - for (int i = 0; i < expected_value.size(); ++i) { - EXPECT_EQ(expected_value[i], actual_values[i]); - } - }; - - auto& attr = *node_def.mutable_attr(); - attr["dtype"].set_type(dtype); - { - // 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. - reset_and_test(t, false, {}, {}); - } - { - Tensor t = ::tensorflow::test::AsScalar(12); - reset_and_test(t, false, {1}, {12}); - reset_and_test(t, true, {1}, {12}); - } - { - Tensor t = ::tensorflow::test::AsTensor({1, 2}); - reset_and_test(t, false, {2}, {1, 2}); - reset_and_test(t, true, {2}, {1, 2}); - } - { - Tensor t = ::tensorflow::test::AsTensor({1, 2, 3, 4, 5, 6}, - TensorShape({2, 3})); - 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}); - } -} - -TEST_F(ConverterTest, ConvertConst) { - { - converter_.Reset(); - NodeDef node_def = MakeNodeDef("my_const", "Const", {"input"}); - AddTestTensor("input", {1}); - ExpectStatus( - converter_.ConvertNode(node_def), error::INVALID_ARGUMENT, - "Constant node is expected to have empty input list: my_const"); - } - { - converter_.Reset(); - NodeDef node_def = MakeNodeDef("my_const", "Const", {}); - (*node_def.mutable_attr())["dtype"].set_type(DT_DOUBLE); - ExpectStatus(converter_.ConvertNode(node_def), error::INVALID_ARGUMENT, - "Unsupported data type"); - } - - TestConvertConst(&converter_); - TestConvertConst(&converter_); -#if NV_TENSORRT_MAJOR > 3 - TestConvertConst(&converter_); -#endif -} - -TEST_F(ConverterTest, ConvertTranspose) { - { - // Input list is empty, should fail. - NodeDef node_def = MakeNodeDef("my_transpose", "Transpose", {}); - ExpectStatus(converter_.ConvertNode(node_def), error::INVALID_ARGUMENT, - "Input expects tensor and weights, at my_transpose"); - } - NodeDef node_def = - MakeNodeDef("my_transpose", "Transpose", {"input", "weights"}); - { - // Permutation is a tensor, should fail. - converter_.Reset(); - AddTestTensor("input", {1, 2, 3}); - AddTestTensor("weights", {3}); - ExpectStatus(converter_.ConvertNode(node_def), error::INVALID_ARGUMENT, - "Input expects tensor and weights, at my_transpose"); - } - { - // Transpose at batch dimension, should fail. - converter_.Reset(); - AddTestTensor("input", {1, 2, 3}); - AddTestWeights("weights", DT_INT32, {4}, {1, 0, 2, 3}); - ExpectStatus(converter_.ConvertNode(node_def), error::UNIMPLEMENTED, - "Transpose at batch dimension is not supported"); - } - { - // Permutation rank doesn't match, should fail. - converter_.Reset(); - AddTestTensor("input", {1, 2, 3}); - AddTestWeights("weights", DT_INT32, {3}, {0, 1, 2}); - ExpectStatus( - converter_.ConvertNode(node_def), error::INVALID_ARGUMENT, - "Rank of perm for transpose does not match with that of the input."); - } - { - // Ok. - converter_.Reset(); - AddTestTensor("input", {1, 2, 3}); - AddTestWeights("weights", DT_INT32, {4}, {0, 3, 1, 2}); - TF_EXPECT_OK(converter_.ConvertNode(node_def)); - TRT_TensorOrWeights output = converter_.GetTensorOrWeights("my_transpose"); - EXPECT_TRUE(output.is_tensor()); - EXPECT_TRUE( - Equals(GetTestDims({3, 1, 2}), output.tensor()->getDimensions())) - << output.DebugString(); - - std::vector output_data(6); - converter_.BuildAndRun("input", {1, 2, 3, 4, 5, 6}, "my_transpose", - &output_data); - EXPECT_THAT(output_data, ElementsAre(1, 4, 2, 5, 3, 6)); - } -} - -TEST_F(ConverterTest, ConvertReshape) { - { - // Input list is empty, should fail. - NodeDef node_def = MakeNodeDef("my_reshape", "Reshape", {}); - ExpectStatus(converter_.ConvertNode(node_def), error::INVALID_ARGUMENT, - "Input expects weights for shape, at my_reshape"); - } - NodeDef node_def = MakeNodeDef("my_reshape", "Reshape", {"input", "weights"}); - { - // Shape is a tensor, should fail. - converter_.Reset(); - AddTestTensor("input", {1, 2, 3}); - AddTestTensor("weights", {3}); - ExpectStatus(converter_.ConvertNode(node_def), error::INVALID_ARGUMENT, - "Input expects weights for shape, at my_reshape"); - } - { - // Reshape to scalar, should fail. - converter_.Reset(); - AddTestTensor("input", {1, 2, 3}); - AddTestWeights("weights", DT_INT32, {}, {}); - ExpectStatus(converter_.ConvertNode(node_def), error::UNIMPLEMENTED, - "Reshape to shape=[] is not supported, at my_reshape"); - } - { - // Reshape at batch dimension, should fail. - converter_.Reset(); - AddTestTensor("input", {1, 2, 3}); - AddTestWeights("weights", DT_INT32, {4}, {-1, 1, 1, 2}); - ExpectStatus(converter_.ConvertNode(node_def), error::UNIMPLEMENTED, - "Reshape on batch dimension is not supported, at my_reshape"); - } - { - // Reshape at batch dimension, should fail. - converter_.Reset(); - AddTestTensor("input", {1, 2, 3}); - AddTestWeights("weights", DT_INT32, {4}, {3, 1, 1, 2}); - ExpectStatus(converter_.ConvertNode(node_def), error::UNIMPLEMENTED, - "Reshape on batch dimension is not supported, at my_reshape"); - } - // Reshape on non batch dimensions, ok. - for (int batch_dim : {-1, 1}) { - converter_.Reset(); - AddTestTensor("input", {1, 2, 3}); - AddTestWeights("weights", DT_INT32, {4}, {batch_dim, 1, 3, 2}); - TF_EXPECT_OK(converter_.ConvertNode(node_def)); - TRT_TensorOrWeights output = converter_.GetTensorOrWeights("my_reshape"); - EXPECT_TRUE(output.is_tensor()); - EXPECT_TRUE( - Equals(GetTestDims({1, 3, 2}), output.tensor()->getDimensions())) - << output.DebugString(); - - std::vector output_data(6); - converter_.BuildAndRun("input", {1, 2, 3, 4, 5, 6}, "my_reshape", - &output_data); - EXPECT_THAT(output_data, ElementsAre(1, 2, 3, 4, 5, 6)); - } -} - -} // namespace convert -} // namespace tensorrt -} // namespace tensorflow - -#endif // GOOGLE_TENSORRT -#endif // GOOGLE_CUDA diff --git a/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.cc b/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.cc deleted file mode 100644 index ff4fba58bfccd7d9c4d744daa3646c3ee14190ad..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/tensorrt/convert/trt_optimization_pass.cc +++ /dev/null @@ -1,291 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "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/core/grappler/clusters/cluster.h" -#include "tensorflow/core/grappler/grappler_item.h" -#include "tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.h" -#include "tensorflow/core/lib/strings/numbers.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/stacktrace.h" - -#if GOOGLE_CUDA -#if GOOGLE_TENSORRT -namespace tensorflow { -namespace tensorrt { -namespace convert { -// TODO(sami): Remove VLOG messages once the code matures -using tensorflow::str_util::Uppercase; -using tensorflow::strings::StrAppend; -using tensorflow::strings::StrCat; - -tensorflow::Status TRTOptimizationPass::Init( - const tensorflow::RewriterConfig_CustomGraphOptimizer* config) { - VLOG(1) << "Called INIT for " << name_ << " with config = " << config; - if (config == nullptr) { - return tensorflow::Status::OK(); - } - const auto params = config->parameter_map(); - if (params.count("minimum_segment_size")) { - minimum_segment_size_ = params.at("minimum_segment_size").i(); - } - if (params.count("max_batch_size")) { - maximum_batch_size_ = params.at("max_batch_size").i(); - } - if (params.count("is_dynamic_op")) { - is_dynamic_op_ = params.at("is_dynamic_op").b(); - } - if (params.count("cached_engine_batches")) { - auto batch_vec = params.at("cached_engine_batches").list(); - batches_.reserve(batch_vec.i_size()); - for (const auto i : batch_vec.i()) { - batches_.push_back(i); - } - } - if (params.count("maximum_cached_engines")) { - max_cached_batches_ = params.at("maximum_cached_engines").i(); - } - if (params.count("max_workspace_size_bytes")) { - max_workspace_size_bytes_ = params.at("max_workspace_size_bytes").i(); - } - if (params.count("precision_mode")) { - TF_RETURN_IF_ERROR(GetPrecisionMode( - Uppercase(params.at("precision_mode").s()), &precision_mode_)); - } - return tensorflow::Status::OK(); -} - -void TRTOptimizationPass::PrintDebugInfo( - tensorflow::grappler::Cluster* cluster, - const tensorflow::grappler::GrapplerItem& item) { - VLOG(1) << "Cluster = " << cluster; - string offset(" "); - string offset2 = StrCat(offset, offset); - string offset3 = StrCat(offset2, offset); - string offset4 = StrCat(offset2, offset2); - if (cluster) { - VLOG(1) << offset << "type = " << cluster->type(); - VLOG(1) << offset << "num warmup steps = " << cluster->NumWarmupSteps(); - const auto dev_names = cluster->GetDeviceNames(); - if (dev_names.size()) { - VLOG(1) << offset << " Device names:"; - for (const auto s : dev_names) { - VLOG(1) << offset2 << s; - } - } - std::unordered_map peak_mem; - auto status = cluster->GetPeakMemoryUsage(&peak_mem); - if (status == tensorflow::Status::OK()) { - VLOG(1) << offset << "Peak Memory Usage :"; - for (auto s : peak_mem) { - VLOG(1) << offset2 << s.first << " = " << s.second; - } - } - - const auto dev_props = cluster->GetDevices(); - if (dev_props.size()) { - VLOG(1) << offset << "Device properties:"; - for (auto k : dev_props) { - VLOG(1) << offset2 << k.first; - const auto& dt = k.second; - VLOG(1) << offset3 << "type = " << dt.type(); - VLOG(1) << offset3 << "vendor = " << dt.vendor(); - VLOG(1) << offset3 << "model = " << dt.model(); - VLOG(1) << offset3 << "frequency = " << dt.frequency(); - VLOG(1) << offset3 << "num cores = " << dt.num_cores(); - VLOG(1) << offset3 << "num registers = " << dt.num_registers(); - VLOG(1) << offset3 << "L1 cache size = " << dt.l1_cache_size(); - VLOG(1) << offset3 << "L2 cache size = " << dt.l2_cache_size(); - VLOG(1) << offset3 << "L3 cache size = " << dt.l3_cache_size(); - VLOG(1) << offset3 << "SHMem per SMP = " - << dt.shared_memory_size_per_multiprocessor(); - VLOG(1) << offset3 << "memory size = " << dt.memory_size(); - VLOG(1) << offset3 << "bandwidth = " << dt.bandwidth(); - if (dt.environment_size()) { - VLOG(1) << offset3 << "environment :"; - for (const auto e : dt.environment()) { - VLOG(1) << offset4 << e.first << " = " << e.second; - } - } - } - } - } - VLOG(1) << "item: " << item.id; - if (item.feed.size()) { - VLOG(1) << offset << "Feeds :"; - for (const auto& f : item.feed) { - const auto& shape = f.second.shape(); - VLOG(1) << offset2 << f.first << " = shaped " << shape.DebugString(); - } - } else { - VLOG(1) << offset << "No Feeds"; - } - if (item.fetch.size()) { - VLOG(1) << offset << "Fetches :"; - for (const auto& f : item.fetch) { - VLOG(1) << offset2 << f; - } - } else { - VLOG(1) << offset << "No Fetches"; - } - - if (item.init_ops.size()) { - VLOG(1) << offset << "init ops :"; - for (const auto& f : item.init_ops) { - VLOG(1) << offset2 << f; - } - } else { - VLOG(1) << offset << "No init ops"; - } - VLOG(1) << "Save Op = " << item.save_op; - VLOG(1) << "Restore Op = " << item.restore_op; - VLOG(1) << "save_restore_loc_tensor = " << item.save_restore_loc_tensor; - if (item.keep_ops.size()) { - VLOG(1) << offset << "keep ops :"; - for (const auto& f : item.keep_ops) { - VLOG(1) << offset2 << f; - } - } else { - VLOG(1) << offset << "No keep ops"; - } - VLOG(3) << item.graph.DebugString(); - for (const auto dev : cluster->GetDeviceSet()->devices()) { - const auto& pname = dev->parsed_name(); - VLOG(1) << "Device name= " << dev->name() - << " parsedname job= " << pname.job << " id= " << pname.id - << " has_id: " << pname.has_id << " has_job: " << pname.has_job - << "has_type: " << pname.has_type << " type =" << pname.type; - } -} - -tensorflow::Status TRTOptimizationPass::Optimize( - tensorflow::grappler::Cluster* cluster, - const tensorflow::grappler::GrapplerItem& item, GraphDef* optimized_graph) { - VLOG(1) << "Called TRTOptimization Pass " << name_; - // This is a hack to workaround optimizer issue. MetaOptimizer calls - // optimization passes on function objects as well, we should not modify - // generated funcdefs! This is fragile but we don't have any other option - // until framework fixes it. - if (item.id != "tf_graph") { - LOG(WARNING) << name_ - << " is probably called on funcdef! This optimizer must *NOT* " - "be called on function objects."; - *optimized_graph = item.graph; - return tensorflow::Status::OK(); - } - if (VLOG_IS_ON(1)) { - VLOG(2) << CurrentStackTrace(); - PrintDebugInfo(cluster, item); - } - int max_dim = -1; - if (item.feed.size()) { - for (const auto& f : item.feed) { - const auto& shape = f.second.shape(); - if (shape.dims() > 0) { - if (shape.dim_size(0) > max_dim) max_dim = shape.dim_size(0); - } - } - } - if (maximum_batch_size_ < 0) { // automatic batch size from input - if (max_dim > 0) { - maximum_batch_size_ = max_dim; - VLOG(1) << "Setting maximum batch size to " << max_dim; - } else { - maximum_batch_size_ = 128; - LOG(WARNING) << "Maximum batch size is not set" - " and can't be deduced from inputs setting it to" - << maximum_batch_size_ - << ". Suggest configuring it from configuration parameters"; - } - } else { - if (max_dim > maximum_batch_size_) { - LOG(WARNING) << "Configured batch size " << maximum_batch_size_ - << " is less than input batch size " << max_dim - << " adjusting maximum batch size to match input batch size"; - } - } - tensorflow::grappler::GraphProperties static_graph_properties(item); - TF_RETURN_IF_ERROR(static_graph_properties.InferStatically(true)); - tensorflow::tensorrt::convert::ConversionParams cp; - - std::vector nodes_to_preserve; - for (const auto& n : item.NodesToPreserve()) { - auto tokens = str_util::Split(n, ":"); - string s = tokens.at(0); - for (int i = 1; i < tokens.size() - 1; ++i) { - StrAppend(&s, ":", tokens.at(i)); - } - int dumm_port = -1; - // 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)) { - StrAppend(&s, ":", tokens.back()); - } - nodes_to_preserve.push_back(s); - } - cp.input_graph_def = &item.graph; - cp.output_names = &nodes_to_preserve; - cp.max_batch_size = maximum_batch_size_; - cp.max_workspace_size_bytes = max_workspace_size_bytes_; - cp.output_graph_def = optimized_graph; - cp.precision_mode = precision_mode_; - cp.minimum_segment_size = minimum_segment_size_; - cp.graph_properties = &static_graph_properties; - cp.cluster = cluster; - cp.is_dyn_op = is_dynamic_op_; - cp.cached_engine_batches = batches_; - cp.max_cached_engines = max_cached_batches_; - auto status = tensorflow::tensorrt::convert::ConvertAfterShapes(cp); - VLOG(2) << optimized_graph->DebugString(); - VLOG(1) << "Returning from " << name_; - return status; -} - -void TRTOptimizationPass::Feedback( - tensorflow::grappler::Cluster* cluster, - const tensorflow::grappler::GrapplerItem& item, - const GraphDef& optimized_graph, double result) {} - -} // namespace convert -} // namespace tensorrt -} // namespace tensorflow - -class VerboseCustomGraphOptimizerRegistrar - : public tensorflow::grappler::CustomGraphOptimizerRegistrar { - public: - VerboseCustomGraphOptimizerRegistrar( - const tensorflow::grappler::CustomGraphOptimizerRegistry::Creator& cr, - const tensorflow::string& name) - : tensorflow::grappler::CustomGraphOptimizerRegistrar(cr, name) { - VLOG(1) << "Constructing a CustomOptimizationPass registration object for " - << name; - } -}; - -static VerboseCustomGraphOptimizerRegistrar TRTOptimizationPass_Registrar( - []() { - VLOG(1) - << "Instantiating CustomOptimizationPass object TensorRTOptimizer"; - return new tensorflow::tensorrt::convert::TRTOptimizationPass( - "TensorRTOptimizer"); - }, - ("TensorRTOptimizer")); - -#endif -#endif diff --git a/tensorflow/contrib/tensorrt/convert/utils.cc b/tensorflow/contrib/tensorrt/convert/utils.cc deleted file mode 100644 index e7a1febb8c076891596741fe30721e7acca15a73..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/tensorrt/convert/utils.cc +++ /dev/null @@ -1,69 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/contrib/tensorrt/convert/utils.h" - -#include "tensorflow/core/lib/core/errors.h" -#include "tensorflow/core/lib/core/status.h" - -namespace tensorflow { -namespace tensorrt { - -bool IsGoogleTensorRTEnabled() { - // TODO(laigd): consider also checking if tensorrt shared libraries are - // accessible. We can then direct users to this function to make sure they can - // safely write code that uses tensorrt conditionally. E.g. if it does not - // check for for tensorrt, and user mistakenly uses tensorrt, they will just - // crash and burn. -#if GOOGLE_CUDA && GOOGLE_TENSORRT - return true; -#else - return false; -#endif -} - -Status GetPrecisionModeName(const int precision_mode, string* name) { - switch (precision_mode) { - case FP32MODE: - *name = "FP32"; - break; - case FP16MODE: - *name = "FP16"; - break; - case INT8MODE: - *name = "INT8"; - break; - default: - return tensorflow::errors::OutOfRange("Unknown precision mode"); - } - return Status::OK(); -} - -Status GetPrecisionMode(const string& name, int* precision_mode) { - if (name == "FP32") { - *precision_mode = FP32MODE; - } else if (name == "FP16") { - *precision_mode = FP16MODE; - } else if (name == "INT8") { - *precision_mode = INT8MODE; - } else { - return tensorflow::errors::InvalidArgument("Invalid precision mode name: ", - name); - } - return Status::OK(); -} - -} // namespace tensorrt -} // namespace tensorflow diff --git a/tensorflow/contrib/tensorrt/convert/utils.h b/tensorflow/contrib/tensorrt/convert/utils.h deleted file mode 100644 index 0592f31462af2b20f3a13fe5119e89c2ba42dd8a..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/tensorrt/convert/utils.h +++ /dev/null @@ -1,50 +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_CONVERT_UTILS_H_ -#define TENSORFLOW_CONTRIB_TENSORRT_CONVERT_UTILS_H_ - -#include - -#include "tensorflow/core/lib/core/status.h" - -namespace tensorflow { -namespace tensorrt { - -template -struct TrtDestroyer { - void operator()(T* t) { - if (t) t->destroy(); - } -}; - -template -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; - -Status GetPrecisionModeName(const int precision_mode, string* name); - -Status GetPrecisionMode(const string& name, int* precision_mode); - -} // namespace tensorrt -} // namespace tensorflow - -#endif // TENSORFLOW_CONTRIB_TENSORRT_CONVERT_UTILS_H_ 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.cc b/tensorflow/contrib/tensorrt/kernels/trt_engine_op.cc deleted file mode 100644 index 019446813a56de6316a04c1738ae13d03e8f4713..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/tensorrt/kernels/trt_engine_op.cc +++ /dev/null @@ -1,606 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ -#include "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 "tensorflow/core/framework/graph_to_functiondef.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/stream_executor.h" -#include "tensorflow/core/platform/types.h" - -#if GOOGLE_CUDA -#if GOOGLE_TENSORRT -#include "cuda/include/cuda_runtime_api.h" - -namespace tensorflow { -namespace tensorrt { -static Logger logger; -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. -class AsyncHelper : public tensorflow::core::RefCounted { - public: - AsyncHelper(AsyncOpKernel::DoneCallback done) { done_ = done; } - ~AsyncHelper() override { done_(); } - - private: - AsyncOpKernel::DoneCallback done_; -}; - -#define TYPECASE(dt, X, Y) \ - case dt: { \ - return (void*)X->flat::Type>().data(); \ - } - -void* GetTensorAddress(const Tensor* tensor_ptr) { - auto tensor_type = tensor_ptr->dtype(); - switch (tensor_type) { - TYPECASE(tensorflow::DT_FLOAT, tensor_ptr, dest_ptr); - TYPECASE(tensorflow::DT_HALF, tensor_ptr, dest_ptr); - TYPECASE(tensorflow::DT_INT8, tensor_ptr, dest_ptr); - default: { - LOG(ERROR) << "Unsupported Data type " - << tensorflow::DataTypeString(tensor_type); - return nullptr; - } - } -} - -tensorflow::Status TRTEngineOp::ConstructFunctionHandle(OpKernelContext* ctx) { - VLOG(1) << "Constructing function handle"; - auto lib = ctx->function_library(); - if (lib == nullptr) { - return tensorflow::errors::Internal("Context function library is null"); - } - auto fdef = lib->GetFunctionLibraryDefinition()->Find(funcdef_name_); - if (fdef == nullptr) { - return tensorflow::errors::Internal("Native FunctionDef ", funcdef_name_, - " can't be found in function library"); - } - tensorflow::FunctionLibraryRuntime::InstantiateOptions inst_ops; - inst_ops.overlay_lib = nullptr; - inst_ops.state_handle = ""; - inst_ops.target = ctx->device()->name(); - native_func_ = 0; - auto status = lib->Instantiate(funcdef_name_, AttrSlice(&fdef->attr()), - inst_ops, &native_func_); - if (!status.ok()) { - LOG(ERROR) << " Instantiating native function " << funcdef_name_ - << " failed!"; - } - return status; -} - -TRTEngineOp::TRTEngineOp(OpKernelConstruction* context) - : AsyncOpKernel(context) { - // read serialized_engine - OP_REQUIRES_OK(context, - context->GetAttr("serialized_segment", &serialized_segment_)); - OP_REQUIRES_OK(context, - context->GetAttr("workspace_size_bytes", &workspace_size_)); - OP_REQUIRES_OK(context, context->GetAttr("static_engine", &static_engine_)); - if (!static_engine_) { - if (!segment_graph_.ParseFromString(serialized_segment_)) { - LOG(ERROR) << "Parsing segment graph failed!"; - context->SetStatus(tensorflow::errors::InvalidArgument( - "Failed to parse segment graphdef!")); - return; - } - serialized_segment_.resize(0); - } - VLOG(1) << "Constructing " << name(); - string precision_string; - OP_REQUIRES_OK(context, - context->GetAttr("precision_mode", &precision_string)); - string calibration_data; - OP_REQUIRES_OK(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_)); - calibration_mode_ = - (precision_mode_ == INT8MODE && calibration_data.size() == 0); - if (calibration_data.size()) { - 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()); - if (VLOG_IS_ON(1)) { - string s("Engine Batches= "); - for (auto i : cached_engine_batches_) { - StrAppend(&s, i, " "); - } - VLOG(1) << s; - } -} - -void TRTEngineOp::ExecuteNativeSegment(OpKernelContext* ctx, - AsyncHelper* helper) { - if (!calibration_mode_) { - VLOG(1) << "Executing native engine"; - } - std::vector inputs; - std::vector* outputs = new std::vector(); - if (native_func_ == tensorflow::kInvalidHandle) { - auto status = ConstructFunctionHandle(ctx); - if (!status.ok()) { - LOG(ERROR) << "Couldn't construct function handle " << funcdef_name_; - ctx->SetStatus(status); - return; - } - } - auto lib = ctx->function_library(); - tensorflow::FunctionLibraryRuntime::Options opts; - opts.step_id = ctx->step_id(); - opts.rendezvous = ctx->rendezvous(); - opts.cancellation_manager = ctx->cancellation_manager(); - opts.runner = ctx->runner(); - for (int i = 0; i < ctx->num_inputs(); i++) { - inputs.push_back(ctx->input(i)); - } - helper->Ref(); // Increment count for calculating native graph - VLOG(1) << "Executing native segment " << name(); - lib->Run(opts, native_func_, inputs, outputs, - [this, ctx, outputs, helper](const tensorflow::Status& s) { - tensorflow::core::ScopedUnref sc(helper); - VLOG(1) << "Native Segment completed"; - if (!s.ok()) { - ctx->SetStatus(s); - return; - } - for (size_t t = 0; t < outputs->size(); ++t) { - ctx->set_output(t, outputs->at(t)); - } - test::AddTestValue(StrCat(this->name(), ":ExecuteNativeSegment"), - "done"); - delete outputs; - }); -} - -void TRTEngineOp::ExecuteCalibration(OpKernelContext* ctx, - AsyncHelper* helper) { - 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"); - 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; - } - int num_inputs = ctx->num_inputs(); - // Pass input data to calibrator - std::unordered_map input_data; - for (int i = 0; i < num_inputs; i++) { - const Tensor& t = ctx->input(i); - void* data_address = GetTensorAddress(&t); - if (data_address == nullptr) { - ctx->SetStatus(tensorflow::errors::InvalidArgument( - "Unsupported data type encountered in input ", i)); - return; - } - // Check the allocated buffer is sufficient for input - const auto device_tensor = dev_tensors_.at(i).AccessTensor(ctx); - CHECK_EQ(t.TotalBytes(), device_tensor->TotalBytes()); - input_data.emplace(StrCat(kInputPHName, i), data_address); - } - VLOG(2) << "Filled map for sending"; - // copied from cuda_kernel_helper since it seems only valid in *.cu.cc files - const cudaStream_t* stream = CHECK_NOTNULL( - reinterpret_cast(ctx->op_device_context() - ->stream() - ->implementation() - ->GpuStreamMemberHack())); - calib_res->calibrator_->setBatch(input_data, *stream); - test::AddTestValue(StrCat(name(), ":ExecuteCalibration"), "done"); - VLOG(2) << "Passed calibration data"; - 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; - } - } - return smallest_engine; -} - -void TRTEngineOp::ComputeAsync(OpKernelContext* ctx, - AsyncOpKernel::DoneCallback done) { - auto helper = new AsyncHelper(done); - tensorflow::core::ScopedUnref sc(helper); - if (calibration_mode_) { - 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; - } - - 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 - << " 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()); - if (retry) { - LOG(WARNING) << "Failed to execute engine, " - << "retrying with native segment for " << name(); - ExecuteNativeSegment(ctx, helper); - return; - } -} - -bool TRTEngineOp::ExecuteTrtEngine( - OpKernelContext* ctx, const int num_batch, - nvinfer1::ICudaEngine* trt_engine_ptr, - nvinfer1::IExecutionContext* trt_execution_context_ptr) { - const bool kRetry = true; - 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 size_t binding_index = - trt_engine_ptr->getBindingIndex(input_name.c_str()); - if (binding_index == -1) { - LOG(ERROR) << "Input node not found, at " << input_name; - return kRetry; - } - - const Tensor& input_tensor = ctx->input(i); - const TensorShape& input_shape = input_tensor.shape(); - if (num_batch != input_shape.dim_size(0)) { - LOG(ERROR) << "Input data has inconsistent batch size: " << num_batch - << " vs " << input_shape.dim_size(0); - return kRetry; - } - auto dtype = trt_engine_ptr->getBindingDataType(binding_index); - switch (dtype) { - case nvinfer1::DataType::kFLOAT: - buffers[binding_index] = (void*)(input_tensor.flat().data()); - break; - case nvinfer1::DataType::kHALF: - LOG(ERROR) << "FP16 inputs are not supported yet!"; - return kRetry; - case nvinfer1::DataType::kINT8: - LOG(ERROR) << "INT8 inputs are not supported yet!"; - return kRetry; - case nvinfer1::DataType::kINT32: - buffers[binding_index] = (void*)(input_tensor.flat().data()); - break; - default: - LOG(ERROR) << "Unknown TRT data type: " << int(dtype); - return kRetry; - } - } - - for (int i = 0; i < ctx->num_outputs(); i++) { - // Create an output tensor - const string output_name = StrCat(kOutputPHName, i); - const size_t binding_index = - trt_engine_ptr->getBindingIndex(output_name.c_str()); - Tensor* output_tensor = nullptr; - - TensorShape output_shape; - if (binding_index != -1) { - auto dims = trt_engine_ptr->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]; - auto status = TensorShapeUtils::MakeShape( - trt_shape.data(), trt_shape.size(), &output_shape); - if (!status.ok()) { - LOG(ERROR) << "Failed to get output shape: " << status; - return kRetry; - } - } else { - LOG(ERROR) << "Output node not found, at " << output_name; - return kRetry; - } - auto status = ctx->allocate_output(i, output_shape, &output_tensor); - if (!status.ok()) { - LOG(ERROR) << "Allocating output failed with " << status; - ctx->SetStatus(status); - // Do not retry since we cannot allocate the same output twice. - // TODO(aaroey): ideally we should retry, fix this. - return !kRetry; - } - auto dtype = trt_engine_ptr->getBindingDataType(binding_index); - switch (dtype) { - case nvinfer1::DataType::kFLOAT: - buffers[binding_index] = - reinterpret_cast(output_tensor->flat().data()); - break; - case nvinfer1::DataType::kHALF: - LOG(WARNING) << "half size is not supported yet!"; - return kRetry; - case nvinfer1::DataType::kINT8: - LOG(WARNING) << "int8 is not supported yet!"; - return kRetry; - case nvinfer1::DataType::kINT32: - buffers[binding_index] = - reinterpret_cast(output_tensor->flat().data()); - break; - default: - LOG(WARNING) << "Unknown TRT data type: " << static_cast(dtype); - return kRetry; - } - } - // Copied from cuda_kernel_helper since it seems only valid in *.cu.cc files - const cudaStream_t* stream = CHECK_NOTNULL( - reinterpret_cast(ctx->op_device_context() - ->stream() - ->implementation() - ->GpuStreamMemberHack())); - - // TODO(jie): trt enqueue does not return error - auto ret = trt_execution_context_ptr->enqueue(num_batch, &buffers[0], *stream, - nullptr); - if (!ret) { - LOG(WARNING) << "Failed to enqueue batch for TRT engine: " << name(); - return kRetry; - } - test::AddTestValue(StrCat(name(), ":ExecuteTrtEngine"), "done"); - // Synchronization will be done by TF. - 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(); - } - 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; - } - 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_); - - if (static_engine_) { - if (engine_map_.size()) { - if (engine_map_.begin()->first >= batch_size) { - return engine_map_.begin()->second; - } - return null_pair; - } - 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(), - serialized_segment_.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())}; - // Runtime is safe to delete after engine creation - serialized_segment_.clear(); - if (max_batch_size < batch_size) { - return null_pair; - } - return engine_map_.at(max_batch_size); - } // 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()); - } - TrtUniquePtrType engine; - bool convert_successfully = false; - VLOG(0) << name() << " Constructing a new engine with batch size " - << batch_size; - // 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, &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}; - } - LOG(WARNING) << "Engine creation for batch size " << batch_size - << " failed " << status; - return null_pair; - } - VLOG(1) << "Conversion is done"; - TrtUniquePtrType exec_context( - engine->createExecutionContext()); - engine_map_[batch_size] = {std::move(engine), std::move(exec_context)}; - } - return engine_map_.at(batch_size); -} - -tensorflow::Status TRTEngineOp::AllocateCalibrationResources( - OpKernelContext* ctx, TRTCalibrationResource** cr) { - auto cres = new TRTCalibrationResource(); - *cr = cres; - // Get the allocator. - auto alloc = ctx->device()->GetAllocator(tensorflow::AllocatorAttributes()); - if (!alloc) { - LOG(WARNING) << "Can't get device allocator will not be able to " - "allocate memory from TensorFlow memory pool"; - cres->allocator_.reset(new TRTCudaAllocator); - } else { - cres->allocator_.reset(new TRTDeviceAllocator(alloc)); - } - // Get the input shapes. - 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); - VLOG(1) << " Constructing calibrator"; - for (int i = 0; i < num_inputs; i++) { - // allocate workspace on device for inputs - const tensorflow::Tensor& t = ctx->input(i); - 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)); - 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( - StrCat(kInputPHName, i), - std::pair(device_address, device_tensor->TotalBytes())); - } - cres->calibrator_.reset( - new TRTInt8Calibrator(device_buffers_, batch_size, name())); - const string label(name()); - auto segment_graph = &segment_graph_; - const int platform_gpu_id = - ctx->device()->tensorflow_gpu_device_info()->gpu_id; - if (platform_gpu_id < 0) { - LOG(ERROR) << "Can't get gpu_device_info from context->device()"; - return tensorflow::errors::InvalidArgument( - "Context->device doesn't contain device info!"); - } - const int64 workspace_size_bytes = workspace_size_; - cres->thr_.reset(new std::thread([cres, label, segment_graph, shapes, - platform_gpu_id, workspace_size_bytes]() { - VLOG(0) << "Starting calibration thread on device " << platform_gpu_id - << ", Calibration Resource @ " << cres; - auto err = cudaSetDevice(platform_gpu_id); - if (err != cudaSuccess) { - // TODO(aaroey): should return error here. - LOG(ERROR) << "Couldn't set cuda device to " << platform_gpu_id - << " in calibration thread"; - } - // ConvertGraphDefToEngine() will try to build the engine. This thread - // will loop inside buildCudaEngine() consuming the calibration data - // that is set by the TF op, and drive the builder until calibrator returns - // false. Engine is discarded after calibration table is generated - // - // 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_, - /*convert_successfully=*/nullptr); - if (!s.ok()) { - LOG(ERROR) << "Calibration failed: " << s; - cres->calibrator_->setDone(); // Ignore further pushes - } - VLOG(1) << "Calibration loop terminated " << label; - })); - VLOG(1) << "initialized calibrator resource"; - return tensorflow::Status::OK(); -} - -REGISTER_KERNEL_BUILDER(Name("TRTEngineOp").Device(DEVICE_GPU), TRTEngineOp); - -} // namespace tensorrt -} // namespace tensorflow - -#endif // GOOGLE_TENSORRT -#endif // GOOGLE_CUDA 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 8fe06758914261035c90a6fda3f114a63a8ac93a..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/tensorrt/kernels/trt_engine_op.h +++ /dev/null @@ -1,141 +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_; -}; - -} // 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/ops/trt_engine_op.cc b/tensorflow/contrib/tensorrt/ops/trt_engine_op.cc deleted file mode 100644 index e0c7b6272379a20e3dacb6cd7c3b39de735d844d..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/tensorrt/ops/trt_engine_op.cc +++ /dev/null @@ -1,57 +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. -==============================================================================*/ - -#if GOOGLE_CUDA -#if GOOGLE_TENSORRT - -#include "tensorflow/core/framework/op.h" -#include "tensorflow/core/framework/op_kernel.h" -#include "tensorflow/core/framework/shape_inference.h" -#include "tensorflow/core/framework/tensor_shape.h" - -namespace tensorflow { - -namespace shape_inference { -extern Status TRTEngineOpShapeInference(InferenceContext* c); -} - -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("static_engine: bool = true") - .Attr("fixed_input_size: bool = true") - .Attr("cached_engine_batches: list(int) = []") - .Attr("max_cached_engines_count: int = 1") - .Attr("workspace_size_bytes: int") - .Attr("precision_mode: {'FP32', 'FP16', 'INT8', 'INT8CALIB'}") - .Attr("calibration_data: string = ''") - .Input("in_tensor: InT") - .Output("out_tensor: OutT"); -// TODO(jie): TF requires concrete output shape for concrete input shapes. -// This is tricky for batch dimension, since we cannot ensure which input -// would carry the correct batch dimension (for the current stage of the -// implementation, we do require all input tensor to carry the same batch -// size, but this could change in the future). Hence we disable shape -// inference function as a workaround. -// .SetShapeFn(shape_inference::TRTEngineOpShapeInference); - -} // namespace tensorflow - -#endif // GOOGLE_TENSORRT -#endif // GOOGLE_CUDA diff --git a/tensorflow/contrib/tensorrt/python/__init__.py b/tensorflow/contrib/tensorrt/python/__init__.py index 7cdfe2b1a612be2eec473d806d0eb44b611ca68a..75490aecfbe84810520c82597d127a36d36de3ee 100644 --- a/tensorflow/contrib/tensorrt/python/__init__.py +++ b/tensorflow/contrib/tensorrt/python/__init__.py @@ -19,7 +19,7 @@ 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.compiler.tf2tensorrt.python.ops import trt_ops 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 diff --git a/tensorflow/contrib/tensorrt/python/ops/trt_engine_op.py b/tensorflow/contrib/tensorrt/python/ops/trt_engine_op.py deleted file mode 100644 index 31a313182be9a2fca7457a539670dbc911ccabb1..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/tensorrt/python/ops/trt_engine_op.py +++ /dev/null @@ -1,34 +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. -# ============================================================================= -"""Exposes the Python wrapper of TRTEngineOp.""" - -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import platform - -if platform.system() != "Windows": - # pylint: disable=wildcard-import,unused-import,g-import-not-at-top - from tensorflow.contrib.tensorrt.ops.gen_trt_engine_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")) -else: - raise RuntimeError("Windows platforms are not supported") diff --git a/tensorflow/contrib/tensorrt/python/trt_convert.py b/tensorflow/contrib/tensorrt/python/trt_convert.py index 99890d910e717dadbaa7786973e0df2fcb7e1ab7..49d72232aa0cfba3f5bf533de04f4d50e65275fd 100644 --- a/tensorflow/contrib/tensorrt/python/trt_convert.py +++ b/tensorflow/contrib/tensorrt/python/trt_convert.py @@ -45,12 +45,19 @@ 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") + +def _to_bytes(s): + """Encode s if it is a sequence of chars.""" + if isinstance(s, _six.text_type): + return s.encode("utf-8", errors="surrogateescape") + return s + + +def _to_string(s): + """Decode s if it is a sequence of bytes.""" + if isinstance(s, _six.binary_type): + return s.decode("utf-8") + return s class TrtPrecisionMode(object): @@ -63,19 +70,20 @@ class TrtPrecisionMode(object): return [TrtPrecisionMode.FP32, TrtPrecisionMode.FP16, TrtPrecisionMode.INT8] -def 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): +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_batches=None, + use_calibration=True): """Returns a RewriterConfig proto for TRT transformation. Args: - rewriter_config: a RewriterConfig proto to append the TensorRTOptimizer to. - If None, it will create one with default settings. + 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' @@ -89,12 +97,21 @@ def tensorrt_rewriter_config(rewriter_config=None, 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 + cached_engine_batches: 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 + should be <= 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. @@ -107,9 +124,16 @@ def tensorrt_rewriter_config(rewriter_config=None, 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: - rewriter_config = rewriter_config_pb2.RewriterConfig() - rewriter_config.optimizers.extend(["constfold", "layout"]) + # 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." @@ -117,7 +141,7 @@ def tensorrt_rewriter_config(rewriter_config=None, precision_mode, TrtPrecisionMode.supported_precision_modes)) - optimizer = rewriter_config.custom_optimizers.add() + 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 @@ -126,15 +150,16 @@ def tensorrt_rewriter_config(rewriter_config=None, "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 " + if cached_engine_batches: + if not isinstance(cached_engine_batches, list): + raise TypeError("cached_engine_batches should be a list.") + if len(cached_engine_batches) > maximum_cached_engines: + raise ValueError("cached_engine_batches should not contain more than " "maximum_cached_engines items.") optimizer.parameter_map["cached_engine_batches"].list.i.extend( - cached_engine_batch_sizes) - return rewriter_config + cached_engine_batches) + optimizer.parameter_map["use_calibration"].b = use_calibration + return rewriter_config_with_trt def create_inference_graph(input_graph_def, @@ -145,8 +170,8 @@ def create_inference_graph(input_graph_def, minimum_segment_size=3, is_dynamic_op=False, maximum_cached_engines=1, - cached_engine_batch_sizes=None, - rewriter_config=None, + cached_engine_batches=None, + use_calibration=True, input_saved_model_dir=None, input_saved_model_tags=None, output_saved_model_dir=None, @@ -172,14 +197,21 @@ def create_inference_graph(input_graph_def, 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 + cached_engine_batches: 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 + should be <= 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. - rewriter_config: a RewriterConfig proto to append the TensorRTOptimizer to. - If None, it will create one with default settings. + 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. @@ -187,8 +219,9 @@ def create_inference_graph(input_graph_def, 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. If not specified, - a default ConfigProto will be used. + 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 @@ -318,21 +351,30 @@ def create_inference_graph(input_graph_def, grappler_meta_graph_def.collection_def["train_op"].CopyFrom( output_collection) - # Create RewriterConfig. - rewriter_config = tensorrt_rewriter_config( + # 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) + cached_engine_batches, use_calibration) + session_config_with_trt.graph_options.rewrite_options.CopyFrom( + rewriter_config_with_trt) # Run Grappler. transformed_graph_def = tf_optimizer.OptimizeGraph( - rewriter_config, grappler_meta_graph_def, graph_id=b"tf_graph") + 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, diff --git a/tensorflow/contrib/tensorrt/python/trt_convert_test.py b/tensorflow/contrib/tensorrt/python/trt_convert_test.py index 530adafcb3fe67051051129a84847e33a30dd85f..abd822c7b71b4d7cca59482bdb51a922a28d480c 100644 --- a/tensorflow/contrib/tensorrt/python/trt_convert_test.py +++ b/tensorflow/contrib/tensorrt/python/trt_convert_test.py @@ -20,12 +20,13 @@ from __future__ import print_function import os -from tensorflow.contrib.tensorrt.python import trt_convert # pylint: disable=unused-import -from tensorflow.contrib.tensorrt.python.ops import trt_engine_op +from tensorflow.compiler.tf2tensorrt.python.ops import trt_ops # pylint: enable=unused-import +from tensorflow.contrib.tensorrt.python import trt_convert from tensorflow.core.framework import graph_pb2 from tensorflow.core.protobuf import config_pb2 +from tensorflow.core.protobuf import rewriter_config_pb2 from tensorflow.python.framework import dtypes from tensorflow.python.framework import graph_util from tensorflow.python.framework import importer @@ -46,9 +47,9 @@ from tensorflow.python.tools import saved_model_utils class TrtConvertTest(test_util.TensorFlowTestCase): """Class to test Tensorflow-TensorRT integration python API.""" - def testTensorrtRewriterConfig(self): - """Test case for trt_convert.tensorrt_rewriter_config().""" - rewriter_cfg = trt_convert.tensorrt_rewriter_config( + def testGetTensorrtRewriterConfig(self): + """Test case for trt_convert.get_tensorrt_rewriter_config().""" + rewriter_cfg = trt_convert.get_tensorrt_rewriter_config( rewriter_config=None, max_batch_size=128, max_workspace_size_bytes=1234, @@ -56,8 +57,11 @@ class TrtConvertTest(test_util.TensorFlowTestCase): minimum_segment_size=10, is_dynamic_op=True, maximum_cached_engines=2, - cached_engine_batch_sizes=[1, 128]) - self.assertEqual(["constfold", "layout"], rewriter_cfg.optimizers) + cached_engine_batches=[1, 128]) + self.assertEqual(["constfold", "layout", "constfold"], + rewriter_cfg.optimizers) + self.assertEqual(rewriter_config_pb2.RewriterConfig.ONE, + rewriter_cfg.meta_optimizer_iterations) trt_optimizer = None for optimizer in rewriter_cfg.custom_optimizers: if optimizer.name == "TensorRTOptimizer": @@ -80,8 +84,7 @@ class TrtConvertTest(test_util.TensorFlowTestCase): trt_optimizer.parameter_map["precision_mode"].s) self.assertEqual(2, trt_optimizer.parameter_map["maximum_cached_engines"].i) self.assertEqual( - [1, 128], - trt_optimizer.parameter_map["cached_engine_batches"].list.i) + [1, 128], trt_optimizer.parameter_map["cached_engine_batches"].list.i) def _GetConfigProto(self): """Get ConfigProto for session creation.""" @@ -158,7 +161,7 @@ class TrtConvertTest(test_util.TensorFlowTestCase): node_name_to_op = {node.name: node.op for node in graph_def.node} self.assertEqual({ "input": "Placeholder", - "my_trt_op_0": "TRTEngineOp", + "TRTEngineOp_0": "TRTEngineOp", "output": "Identity" }, node_name_to_op) @@ -184,11 +187,12 @@ class TrtConvertTest(test_util.TensorFlowTestCase): self.assertAllEqual([[[4.0]]] * batch_size, result) execute_engine_test_value = ("done" if expect_engine_is_run else "") execute_native_segment_test_value = ("" if expect_engine_is_run else "done") - self.assertEqual(execute_engine_test_value, - trt_convert.get_test_value("my_trt_op_0:ExecuteTrtEngine")) + self.assertEqual( + execute_engine_test_value, + trt_convert.get_test_value("TRTEngineOp_0:ExecuteTrtEngine")) self.assertEqual( execute_native_segment_test_value, - trt_convert.get_test_value("my_trt_op_0:ExecuteNativeSegment")) + trt_convert.get_test_value("TRTEngineOp_0:ExecuteNativeSegment")) def testCreateInferenceGraph_MinimumSegmentSize(self): if not trt_convert.is_tensorrt_enabled(): @@ -234,8 +238,8 @@ class TrtConvertTest(test_util.TensorFlowTestCase): # Run with batch size 2, a new engine is created and cached. self._TestRun(sess, 2, True) # Run with batch size 3, since the number of cached engines has reached - # the max, it should fall back to TF function. - self._TestRun(sess, 3, False) + # the max, it should evict an old engine and create a new one. + self._TestRun(sess, 3, True) # Test the output SavedModel with ops.Graph().as_default(): @@ -246,8 +250,8 @@ class TrtConvertTest(test_util.TensorFlowTestCase): # Run with batch size 2, a new engine is created and cached. self._TestRun(sess, 2, True) # Run with batch size 3, since the number of cached engines has reached - # the max, it should fall back to TF function. - self._TestRun(sess, 3, False) + # the max, it should evict an old engine and create a new one. + self._TestRun(sess, 3, True) def testCreateInferenceGraph_StaticOp(self): if not trt_convert.is_tensorrt_enabled(): diff --git a/tensorflow/contrib/tensorrt/resources/trt_resources.h b/tensorflow/contrib/tensorrt/resources/trt_resources.h deleted file mode 100644 index d7d56cb95e033ea55bd3aa385a707e7a7cfc557b..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/tensorrt/resources/trt_resources.h +++ /dev/null @@ -1,100 +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() { - 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 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_; -}; - -class TRTWeightStore { - public: - TRTWeightStore() {} - - virtual ~TRTWeightStore() { VLOG(1) << "Destroying store" << DebugString(); } - - string DebugString() { - std::stringstream oss; - size_t len_bytes = 0; - for (const auto& v : store_) { - len_bytes += v.size() * sizeof(uint8_t); - } - oss << " Number of entries = " << store_.size() << std::endl - << " Total number of bytes = " - << store_.size() * sizeof(std::vector) + len_bytes - << std::endl; - return oss.str(); - } - - std::list> store_; -}; - -} // 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/base_test.py b/tensorflow/contrib/tensorrt/test/base_test.py index 18096e0ff1ec6b9872346d8a84ac93c542cfb643..17e0b6f4d2c4bbaf56ef143b78c543c9e130b458 100644 --- a/tensorflow/contrib/tensorrt/test/base_test.py +++ b/tensorflow/contrib/tensorrt/test/base_test.py @@ -56,8 +56,9 @@ class SimpleSingleEngineTest(trt_test.TfTrtIntegrationTestBase): strides=[1, 2, 2, 1], padding="SAME", name="conv") - bias = constant_op.constant( - [4., 1.5, 2., 3., 5., 7.], name="bias", dtype=dtype) + bias = constant_op.constant([4., 1.5, 2., 3., 5., 7.], + name="bias", + dtype=dtype) added = nn.bias_add(conv, bias, name="bias_add") relu = nn.relu(added, "relu") identity = array_ops.identity(relu, "identity") @@ -67,17 +68,18 @@ class SimpleSingleEngineTest(trt_test.TfTrtIntegrationTestBase): return trt_test.TfTrtIntegrationTestParams( gdef=g.as_graph_def(), input_names=[input_name], - input_dims=[input_dims], + input_dims=[[input_dims]], output_names=[output_name], - expected_output_dims=[(100, 6, 6, 6)]) + expected_output_dims=[[[100, 6, 6, 6]]]) def ExpectedEnginesToBuild(self, run_params): """Return the expected engines to build.""" - # TODO(aaroey): LayoutOptimizer adds additional nodes to the graph which - # breaks the connection check, fix it. - # - my_trt_op_0 should have ["weights", "conv", "bias", "bias_add", - # "relu", "identity", "max_pool"] - return ["my_trt_op_0"] + return { + "TRTEngineOp_0": [ + "weights", "conv", "bias", "bias_add", "relu", "identity", + "max_pool" + ] + } class SimpleMultiEnginesTest(trt_test.TfTrtIntegrationTestBase): @@ -92,7 +94,7 @@ class SimpleMultiEnginesTest(trt_test.TfTrtIntegrationTestBase): g = ops.Graph() with g.as_default(): inp = array_ops.placeholder( - dtype=dtype, shape=[None] + input_dims[1:], name=input_name) + dtype=dtype, shape=input_dims, name=input_name) with g.device("/GPU:0"): conv_filter = constant_op.constant( [[[[1., 0.5, 4., 6., 0.5, 1.], [1., 0.5, 1., 1., 0.5, 1.]]]], @@ -105,10 +107,10 @@ class SimpleMultiEnginesTest(trt_test.TfTrtIntegrationTestBase): padding="SAME", name="conv") c1 = constant_op.constant( - np.random.randn(input_dims[0], 12, 12, 6), dtype=dtype, name="c1") + np.random.randn(12, 12, 6), dtype=dtype, name="c1") p = math_ops.mul(conv, c1, name="mul") c2 = constant_op.constant( - np.random.randn(input_dims[0], 12, 12, 6), dtype=dtype, name="c2") + np.random.randn(12, 12, 6), dtype=dtype, name="c2") q = math_ops.div(conv, c2, name="div") edge = self.trt_incompatible_op(q, name="incompatible") @@ -123,28 +125,27 @@ class SimpleMultiEnginesTest(trt_test.TfTrtIntegrationTestBase): return trt_test.TfTrtIntegrationTestParams( gdef=g.as_graph_def(), input_names=[input_name], - input_dims=[input_dims], + input_dims=[[input_dims]], output_names=[output_name], - expected_output_dims=[(100, 12, 12, 6)]) + expected_output_dims=[[[100, 12, 12, 6]]]) def ExpectedEnginesToBuild(self, run_params): """Return the expected engines to build.""" - # TODO(aaroey): LayoutOptimizer adds additional nodes to the graph which - # breaks the connection check, fix it. - # - my_trt_op_0 should have ["mul", "sub", "div1", "mul1", "add1", - # "add", "sub1"]; - # - my_trt_op_1 should have ["weights","conv", "div"] - return ["my_trt_op_0", "my_trt_op_1"] + return { + "TRTEngineOp_0": [ + "add", "add1", "c1", "div1", "mul", "mul1", "sub", "sub1" + ], + "TRTEngineOp_1": ["c2", "conv", "div", "weights"] + } - def ShouldRunTest(self, run_params): - # TODO(aaroey): LayoutOptimizer adds Transpose(Const, Const) to the graph - # which breaks the conversion. We should fix it as: - # - Detect the invalid NodeDef earlier before adding them to segment - # - Let it able to change the RewriterConfig when calling - # create_inference_graph(). - # It will be good to add debugging feature for Grappler to print the graph - # after running each optimizer. - return False + def GetConversionParams(self, run_params): + """Return a ConversionParams for test.""" + return super( + SimpleMultiEnginesTest, self + ).GetConversionParams(run_params)._replace( + # Disable layout optimizer, since it'll add Transpose(Const, Const) to + # the graph and breaks the conversion check. + rewriter_config=trt_test.OptimizerDisabledRewriterConfig()) class PartiallyConvertedTestA(trt_test.TfTrtIntegrationTestBase): @@ -153,7 +154,7 @@ class PartiallyConvertedTestA(trt_test.TfTrtIntegrationTestBase): """Setup method.""" super(PartiallyConvertedTestA, self).setUp() # Let it fail to build the second engine. - trt_convert.add_test_value("my_trt_op_1:CreateTRTNode", "fail") + trt_convert.add_test_value("TRTEngineOp_1:CreateTRTNode", "fail") def GetParams(self): """Create a graph containing two segment.""" @@ -182,22 +183,24 @@ class PartiallyConvertedTestA(trt_test.TfTrtIntegrationTestBase): return trt_test.TfTrtIntegrationTestParams( gdef=g.as_graph_def(), input_names=[input_name], - input_dims=[input_dims], + input_dims=[[input_dims]], output_names=[output_name], - expected_output_dims=[tuple(input_dims)]) + expected_output_dims=[[input_dims]]) def ExpectedEnginesToBuild(self, run_params): """Return the expected engines to build.""" return { # Only the first engine is built. - "my_trt_op_0": ["c0", "c1", "add0", "add1", "mul0", "mul1"] + "TRTEngineOp_0": ["c0", "c1", "add0", "add1", "mul0", "mul1"] } def ShouldRunTest(self, run_params): """Whether to run the test.""" # Disable the test in fp16 mode since multiple matmul and add ops together # can cause overflow. - return run_params.precision_mode != "FP16" + return ((run_params.precision_mode != "FP16") and + not (trt_test.IsQuantizationMode(run_params.precision_mode) and + not run_params.use_calibration)) class PartiallyConvertedTestB(PartiallyConvertedTestA): @@ -207,13 +210,13 @@ class PartiallyConvertedTestB(PartiallyConvertedTestA): super(PartiallyConvertedTestB, self).setUp() # Let it fail to build the first engine. trt_convert.clear_test_values("") - trt_convert.add_test_value("my_trt_op_0:CreateTRTNode", "fail") + trt_convert.add_test_value("TRTEngineOp_0:CreateTRTNode", "fail") def ExpectedEnginesToBuild(self, run_params): """Return the expected engines to build.""" return { # Only the second engine is built. - "my_trt_op_1": ["c2", "c3", "add2", "add3", "mul2", "mul3"] + "TRTEngineOp_1": ["c2", "c3", "add2", "add3", "mul2", "mul3"] } @@ -250,15 +253,15 @@ class ConstInputTest(trt_test.TfTrtIntegrationTestBase): return trt_test.TfTrtIntegrationTestParams( gdef=g.as_graph_def(), input_names=[input_name], - input_dims=[input_dims], + input_dims=[[input_dims]], output_names=[output_name], - expected_output_dims=[tuple(input_dims)]) + expected_output_dims=[[input_dims]]) def ExpectedEnginesToBuild(self, run_params): """Return the expected engines to build.""" return { - "my_trt_op_0": ["add", "add1", "mul"], - "my_trt_op_1": ["add2", "add3", "mul1"] + "TRTEngineOp_0": ["add", "add1", "mul"], + "TRTEngineOp_1": ["add2", "add3", "mul1"] } @@ -283,13 +286,13 @@ class ConstDataInputSingleEngineTest(trt_test.TfTrtIntegrationTestBase): return trt_test.TfTrtIntegrationTestParams( gdef=g.as_graph_def(), input_names=[input_name], - input_dims=[input_dims], + input_dims=[[input_dims]], output_names=[output_name], - expected_output_dims=[tuple(input_dims)]) + expected_output_dims=[[input_dims]]) def ExpectedEnginesToBuild(self, run_params): """Return the expected engines to build.""" - return {"my_trt_op_0": ["c", "add", "add1", "mul"]} + return {"TRTEngineOp_0": ["c", "add", "add1", "mul"]} class ConstDataInputMultipleEnginesTest(trt_test.TfTrtIntegrationTestBase): @@ -317,19 +320,19 @@ class ConstDataInputMultipleEnginesTest(trt_test.TfTrtIntegrationTestBase): return trt_test.TfTrtIntegrationTestParams( gdef=g.as_graph_def(), input_names=[input_name], - input_dims=[input_dims], + input_dims=[[input_dims]], output_names=[output_name], - expected_output_dims=[tuple(input_dims)]) + expected_output_dims=[[input_dims]]) def ExpectedEnginesToBuild(self, run_params): """Return the expected engines to build.""" return { - "my_trt_op_0": ["add2", "add3", "mul1"], + "TRTEngineOp_0": ["add2", "add3", "mul1"], # Why segment ["add", "add1", "mul"] was assigned segment id 1 # instead of 0: the parent node of this segment is actually const # node 'c', but it's removed later since it's const output of the # segment which is not allowed. - "my_trt_op_1": ["add", "add1", "mul"] + "TRTEngineOp_1": ["add", "add1", "mul"] } @@ -366,15 +369,15 @@ class ControlDependencyTest(trt_test.TfTrtIntegrationTestBase): return trt_test.TfTrtIntegrationTestParams( gdef=g.as_graph_def(), input_names=[input_name], - input_dims=[input_dims], + input_dims=[[input_dims]], output_names=[output_name], - expected_output_dims=[tuple(input_dims)]) + expected_output_dims=[[input_dims]]) def ExpectedEnginesToBuild(self, run_params): """Return the expected engines to build.""" return { - "my_trt_op_0": ["c1", "add", "add1", "mul"], - "my_trt_op_1": ["c2", "add2", "add3", "mul1"] + "TRTEngineOp_0": ["c1", "add", "add1", "mul"], + "TRTEngineOp_1": ["c2", "add2", "add3", "mul1"] } diff --git a/tensorflow/contrib/tensorrt/test/batch_matmul_test.py b/tensorflow/contrib/tensorrt/test/batch_matmul_test.py index 4b8880817876143dc753cfacdb79d4ad50347fe0..46e3407d9669085a9737bacbeec1a20765ef88cc 100644 --- a/tensorflow/contrib/tensorrt/test/batch_matmul_test.py +++ b/tensorflow/contrib/tensorrt/test/batch_matmul_test.py @@ -71,42 +71,20 @@ class BatchMatMulTest(trt_test.TfTrtIntegrationTestBase): return trt_test.TfTrtIntegrationTestParams( gdef=g.as_graph_def(), input_names=[input_name, w1_name, w2_name], - input_dims=[input_dims, w1_dims, w2_dims], + input_dims=[[input_dims, w1_dims, w2_dims]], output_names=[output_name], - expected_output_dims=[(12, 5, 8, 7)]) + expected_output_dims=[[[12, 5, 8, 7]]]) def ExpectedEnginesToBuild(self, run_params): """Return the expected engines to build.""" if (run_params.dynamic_engine and not trt_test.IsQuantizationMode(run_params.precision_mode)): - return ["my_trt_op_0", "my_trt_op_1"] - return ["my_trt_op_1"] + return ["TRTEngineOp_0", "TRTEngineOp_1"] + return ["TRTEngineOp_1"] def ExpectedEnginesToRun(self, run_params): """Return the expected engines to run.""" - return ["my_trt_op_1"] - - def ShouldRunTest(self, run_params): - """Whether to run the test.""" - # TODO(aaroey): Trt library will fail like: - # - # ../builder/cudnnBuilder2.cpp:685: - # virtual std::vector> - # nvinfer1::builder::Node::getSupportedFormats( - # const nvinfer1::query::Ports&, - # const nvinfer1::cudnn::HardwareContext&, - # nvinfer1::builder::Format::Type, - # const nvinfer1::builder::FormatTypeHack&) const: - # Assertion `sf' failed. - # - # To reproduce, run: - # bazel test -c opt --copt=-mavx \ - # --test_arg=BatchMatMulTest.testTfTrt_ToolConversion_INT8_DynamicEngine \ - # tensorflow/contrib/tensorrt:batch_matmul_test - # - # Investigate and fix it. - return not trt_test.IsQuantizationMode(run_params.precision_mode) + return ["TRTEngineOp_1"] if __name__ == "__main__": diff --git a/tensorflow/contrib/tensorrt/test/biasadd_matmul_test.py b/tensorflow/contrib/tensorrt/test/biasadd_matmul_test.py index 7d006b73d5363117d4b943f09a7a28b1c0a08b95..ca23629efacba1df27ffb466d24b189d6074a917 100644 --- a/tensorflow/contrib/tensorrt/test/biasadd_matmul_test.py +++ b/tensorflow/contrib/tensorrt/test/biasadd_matmul_test.py @@ -41,6 +41,7 @@ class BiasaddMatMulTest(trt_test.TfTrtIntegrationTestBase): input_name = "input" input_matrix_rows = 4 input_matrix_columns = 144 + # Note that tf.nn.bias_add supports up to 5 dimensions. input_dims = [input_matrix_rows, input_matrix_columns] output_name = "output" g = ops.Graph() @@ -74,18 +75,18 @@ class BiasaddMatMulTest(trt_test.TfTrtIntegrationTestBase): x5 = nn.bias_add(x5, b) x5 = gen_array_ops.reshape(x5, [4, -1]) - x6 = gen_array_ops.reshape(x, [4, 12, 12]) - b = self._ConstOp((12,)) + x6 = gen_array_ops.reshape(x, [4, 24, 6]) + b = self._ConstOp((6,)) x6 = nn.bias_add(x6, b, data_format="NHWC") x6 = gen_array_ops.reshape(x6, [4, -1]) - x7 = gen_array_ops.reshape(x, [4, 12, 3, 4]) - b = self._ConstOp((4,)) + x7 = gen_array_ops.reshape(x, [4, 12, 4, 3]) + b = self._ConstOp((3,)) x7 = nn.bias_add(x7, b, data_format="NHWC") x7 = gen_array_ops.reshape(x7, [4, -1]) - x8 = gen_array_ops.reshape(x, [4, 12, 3, 2, 2]) - b = self._ConstOp((2,)) + x8 = gen_array_ops.reshape(x, [4, 4, 3, 2, 6]) + b = self._ConstOp((6,)) x8 = nn.bias_add(x8, b, data_format="NHWC") x8 = gen_array_ops.reshape(x8, [4, -1]) @@ -94,13 +95,13 @@ class BiasaddMatMulTest(trt_test.TfTrtIntegrationTestBase): x9 = nn.bias_add(x9, b, data_format="NCHW") x9 = gen_array_ops.reshape(x9, [4, -1]) - x10 = gen_array_ops.reshape(x, [4, 12, 3, 4]) - b = self._ConstOp((12,)) + x10 = gen_array_ops.reshape(x, [4, 3, 4, 12]) + b = self._ConstOp((3,)) x10 = nn.bias_add(x10, b, data_format="NCHW") x10 = gen_array_ops.reshape(x10, [4, -1]) - x11 = gen_array_ops.reshape(x, [4, 12, 12]) - b = self._ConstOp((12,)) + x11 = gen_array_ops.reshape(x, [4, 6, 24]) + b = self._ConstOp((6,)) x11 = nn.bias_add(x11, b, data_format="NCHW") x11 = gen_array_ops.reshape(x11, [4, -1]) @@ -110,44 +111,24 @@ class BiasaddMatMulTest(trt_test.TfTrtIntegrationTestBase): return trt_test.TfTrtIntegrationTestParams( gdef=g.as_graph_def(), input_names=[input_name], - input_dims=[input_dims], + input_dims=[[input_dims]], output_names=[output_name], - expected_output_dims=[(4, 6680)]) + expected_output_dims=[[[4, 6680]]]) def GetConversionParams(self, run_params): """Return a ConversionParams for test.""" - return super(BiasaddMatMulTest, - self).GetConversionParams(run_params)._replace( - max_batch_size=4, maximum_cached_engines=2) - - def _ValidEngines(self): - """Engines expected to build and run.""" - return ["my_trt_op_0"] - - def _InvalidEngines(self): - """Engines that will cause conversion error at building time.""" - return ["my_trt_op_1", "my_trt_op_2"] + conversion_params = super(BiasaddMatMulTest, + self).GetConversionParams(run_params) + return conversion_params._replace( + max_batch_size=4, + maximum_cached_engines=1, + # Disable layout optimizer, since it will convert BiasAdd with NHWC + # format to NCHW format under four dimentional input. + rewriter_config=trt_test.OptimizerDisabledRewriterConfig()) def ExpectedEnginesToBuild(self, run_params): """Return the expected engines to build.""" - # In dynamic engine mode the engines are built in execution time, not in - # conversion time, so build errors occurs later. Here three of the engines - # will be failed to built but the corresponding engine op are still created. - # TODO(aaroey, jjsjann123): fix this. - if (run_params.dynamic_engine and - not trt_test.IsQuantizationMode(run_params.precision_mode)): - return self._ValidEngines() + self._InvalidEngines() - return self._ValidEngines() - - def ExpectedEnginesToRun(self, run_params): - """Return the expected engines to run.""" - return self._ValidEngines() - - def ShouldRunTest(self, run_params): - """Whether to run the test.""" - # TODO(aaroey): Trt 4.0 forbids conversion for tensors with rank <3 in int8 - # mode, which is a bug. Re-enable this when trt library is fixed. - return not trt_test.IsQuantizationMode(run_params.precision_mode) + return ["TRTEngineOp_0"] if __name__ == "__main__": diff --git a/tensorflow/contrib/tensorrt/test/binary_tensor_weight_broadcast_test.py b/tensorflow/contrib/tensorrt/test/binary_tensor_weight_broadcast_test.py index b53cb3c091ea477ef0974d9d14d82c587a431152..846fd009db07b151e1eca794e9a8a38ff834a465 100644 --- a/tensorflow/contrib/tensorrt/test/binary_tensor_weight_broadcast_test.py +++ b/tensorflow/contrib/tensorrt/test/binary_tensor_weight_broadcast_test.py @@ -26,7 +26,6 @@ from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import gen_array_ops -from tensorflow.python.ops import math_ops from tensorflow.python.platform import test @@ -56,21 +55,21 @@ class BinaryTensorWeightBroadcastTest(trt_test.TfTrtIntegrationTestBase): ]: a = self._ConstOp(weights_shape) f = x + a - x = math_ops.sigmoid(f) + x = self.trt_incompatible_op(f) a = self._ConstOp(weights_shape) f = a + x - x = math_ops.sigmoid(f) + x = self.trt_incompatible_op(f) gen_array_ops.reshape(x, [5, -1], name=output_name) return trt_test.TfTrtIntegrationTestParams( gdef=g.as_graph_def(), input_names=[input_name], - input_dims=[input_dims], + input_dims=[[input_dims]], output_names=[output_name], - expected_output_dims=[(5, 23040)]) + expected_output_dims=[[[5, 23040]]]) def ExpectedEnginesToBuild(self, run_params): """Return the expected engines to build.""" - return ["my_trt_op_%d" % i for i in range(16)] + return ["TRTEngineOp_%d" % i for i in range(16)] if __name__ == "__main__": diff --git a/tensorflow/contrib/tensorrt/test/concatenation_test.py b/tensorflow/contrib/tensorrt/test/concatenation_test.py index 465cb022964df046bf03a481bb1c6b65750aa883..5d8742ae356c091ba831bbd48741dee34cd39d08 100644 --- a/tensorflow/contrib/tensorrt/test/concatenation_test.py +++ b/tensorflow/contrib/tensorrt/test/concatenation_test.py @@ -73,13 +73,13 @@ class ConcatenationTest(trt_test.TfTrtIntegrationTestBase): return trt_test.TfTrtIntegrationTestParams( gdef=g.as_graph_def(), input_names=[input_name], - input_dims=[input_dims], + input_dims=[[input_dims]], output_names=[output_name], - expected_output_dims=[(2, 126)]) + expected_output_dims=[[[2, 126]]]) def ExpectedEnginesToBuild(self, run_params): """Return the expected engines to build.""" - return ["my_trt_op_0"] + return ["TRTEngineOp_0"] if __name__ == "__main__": diff --git a/tensorflow/contrib/tensorrt/test/const_broadcast_test.py b/tensorflow/contrib/tensorrt/test/const_broadcast_test.py index e32f0478661caaab5386339c819b524656baf066..9137d0072f66321d8420b7caac6acc329541123f 100644 --- a/tensorflow/contrib/tensorrt/test/const_broadcast_test.py +++ b/tensorflow/contrib/tensorrt/test/const_broadcast_test.py @@ -58,13 +58,13 @@ class ConstBroadcastTest(trt_test.TfTrtIntegrationTestBase): return trt_test.TfTrtIntegrationTestParams( gdef=g.as_graph_def(), input_names=[input_name], - input_dims=[input_dims], + input_dims=[[input_dims]], output_names=[output_name], - expected_output_dims=[(5, 12, 12, 1)]) + expected_output_dims=[[[5, 12, 12, 1]]]) def ExpectedEnginesToBuild(self, run_params): """Return the expected engines to build.""" - return ['my_trt_op_0'] + return ['TRTEngineOp_0'] def ExpectedAbsoluteTolerance(self, run_params): """The absolute tolerance to compare floating point results.""" diff --git a/tensorflow/contrib/tensorrt/test/conv2d_test.py b/tensorflow/contrib/tensorrt/test/conv2d_test.py new file mode 100644 index 0000000000000000000000000000000000000000..e7993b4620931736cd872bfffb4ebe177555fcd2 --- /dev/null +++ b/tensorflow/contrib/tensorrt/test/conv2d_test.py @@ -0,0 +1,191 @@ +# 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. +# ============================================================================== +"""Model script to test TF-TensorRT integration.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.contrib.tensorrt.test import tf_trt_integration_test_base as trt_test +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 gen_nn_ops +from tensorflow.python.platform import test + + +def conv2d_layer(inputs, + filters, + kernel_size, + strides=(1, 1), + padding="valid", + data_format="channels_last", + dilation_rate=(1, 1), + name=None): + dtype = inputs.dtype + c_axis = -1 if data_format == "channels_last" else 1 + nchan = inputs.shape[c_axis] + weights_shape = (kernel_size[0], kernel_size[1], nchan, filters) + weights = constant_op.constant(np.random.randn(*weights_shape), dtype=dtype) + padding = padding.upper() + if data_format == "channels_last": + strides = [1] + list(strides) + [1] + dilations = [1] + list(dilation_rate) + [1] + data_format = "NHWC" + else: + strides = [1, 1] + list(strides) + dilations = [1, 1] + list(dilation_rate) + data_format = "NCHW" + return gen_nn_ops.conv2d( + inputs, + weights, + strides=strides, + padding=padding, + dilations=dilations, + data_format=data_format) + + +def div_round_up(n, d): + return (n - 1) // d + 1 + + +def build_graph(input_dims, + dtype, + num_filters, + data_format, + kernel_sizes, + dilation_rates, + padding="same"): + g = ops.Graph() + with g.as_default(): + inp = array_ops.placeholder( + dtype=dtype, shape=[None] + input_dims[1:], name="input") + with g.device("/GPU:0"): + results = [] + for kernel_size in kernel_sizes: + for dilation_rate in dilation_rates: + result = conv2d_layer(inp, num_filters, kernel_size, (1, 1), padding, + data_format, dilation_rate) + results.append(result) + output = sum(results) + output = array_ops.identity(output, name="output") + return g + + +class Conv2DNCHWTest(trt_test.TfTrtIntegrationTestBase): + + def GetParams(self): + """Testing conversion of Conv2D (data_format=NCHW) in TF-TRT conversion.""" + np.random.seed(1234) + input_dims = [13, 3, 7, 11] + g = build_graph( + input_dims=input_dims, + dtype=dtypes.float32, + num_filters=5, + data_format="channels_first", + kernel_sizes=[(3, 3), (3, 2)], + dilation_rates=[(1, 1), (2, 3)]) + return trt_test.TfTrtIntegrationTestParams( + gdef=g.as_graph_def(), + input_names=["input"], + input_dims=[[input_dims]], + output_names=["output"], + expected_output_dims=[[[13, 5, 7, 11]]]) + + def ExpectedEnginesToBuild(self, run_params): + """Return the expected engines to build.""" + return ["TRTEngineOp_0"] + + +class Conv2DNHWCTest(trt_test.TfTrtIntegrationTestBase): + + def GetParams(self): + """Testing conversion of Conv2D (data_format=NCHW) in TF-TRT conversion.""" + np.random.seed(1234) + input_dims = [13, 7, 11, 3] + g = build_graph( + input_dims=input_dims, + dtype=dtypes.float32, + num_filters=5, + data_format="channels_last", + kernel_sizes=[(3, 3), (3, 2)], + dilation_rates=[(1, 1), (2, 3)]) + return trt_test.TfTrtIntegrationTestParams( + gdef=g.as_graph_def(), + input_names=["input"], + input_dims=[[input_dims]], + output_names=["output"], + expected_output_dims=[[[13, 7, 11, 5]]]) + + def ExpectedEnginesToBuild(self, run_params): + """Return the expected engines to build.""" + return ["TRTEngineOp_0"] + + +class Conv2DStridedNCHWTest(trt_test.TfTrtIntegrationTestBase): + + def GetParams(self): + """Testing conversion of strided Conv2D (data_format=NCHW) in TF-TRT + + conversion. + """ + np.random.seed(1234) + dtype = dtypes.float32 + input_name = "input" + n, c, h, w = 13, 3, 7, 11 + num_filters = 5 + input_dims = [n, c, h, w] + output_name = "output" + g = ops.Graph() + with g.as_default(): + inp = array_ops.placeholder( + dtype=dtype, shape=[None] + input_dims[1:], name=input_name) + with g.device("/GPU:0"): + output = inp + output = conv2d_layer( + output, + num_filters, (3, 2), + strides=(2, 2), + padding="same", + data_format="channels_first") + h = div_round_up(h, 2) + w = div_round_up(w, 2) + output = conv2d_layer( + output, + num_filters, (3, 3), + strides=(2, 2), + dilation_rate=(2, 3), + padding="same", + data_format="channels_first") + h = div_round_up(h, 2) + w = div_round_up(w, 2) + output = array_ops.identity(output, name=output_name) + return trt_test.TfTrtIntegrationTestParams( + gdef=g.as_graph_def(), + input_names=[input_name], + input_dims=[[input_dims]], + output_names=[output_name], + expected_output_dims=[[[n, num_filters, h, w]]]) + + def ExpectedEnginesToBuild(self, run_params): + """Return the expected engines to build.""" + return ["TRTEngineOp_0"] + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/tensorrt/test/dynamic_input_shapes_test.py b/tensorflow/contrib/tensorrt/test/dynamic_input_shapes_test.py new file mode 100644 index 0000000000000000000000000000000000000000..cc28cd6087997359e81ffaa6dc8bd958109cc565 --- /dev/null +++ b/tensorflow/contrib/tensorrt/test/dynamic_input_shapes_test.py @@ -0,0 +1,107 @@ +# 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-TRT INT8 conversion without calibration on Mnist model.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.contrib.tensorrt.test import tf_trt_integration_test_base as trt_test +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 nn +from tensorflow.python.platform import test + + +class DynamicInputShapesTest(trt_test.TfTrtIntegrationTestBase): + + def GetParams(self): + # TODO(laigd): we should test the following cases: + # - batch size is not changed, other dims are changing + # - batch size is decreasing, other dims are identical + # - batch size is decreasing, other dims are changing + # - batch size is increasing, other dims are identical + # - batch size is increasing, other dims are changing + input_dims = [[[1, 5, 5, 1]], [[10, 5, 5, 1]], [[3, 5, 5, 1]], + [[1, 5, 5, 1]], [[1, 3, 1, 1]], [[2, 9, 9, 1]], + [[1, 224, 224, 1]], [[1, 128, 224, 1]]] + expected_output_dims = input_dims + + g = ops.Graph() + with g.as_default(): + x = array_ops.placeholder( + shape=(None, None, None, 1), dtype=dtypes.float32, name="input") + conv_filter1 = constant_op.constant( + np.ones([3, 3, 1, 8]), name="weights1", dtype=dtypes.float32) + bias1 = constant_op.constant(np.random.randn(8), dtype=dtypes.float32) + x = nn.conv2d( + input=x, + filter=conv_filter1, + strides=[1, 1, 1, 1], + padding="SAME", + name="conv") + x = nn.bias_add(x, bias1) + x = nn.relu(x) + conv_filter2 = constant_op.constant( + np.ones([3, 3, 8, 1]), name="weights2", dtype=dtypes.float32) + bias2 = constant_op.constant(np.random.randn(1), dtype=dtypes.float32) + x = nn.conv2d( + input=x, + filter=conv_filter2, + strides=[1, 1, 1, 1], + padding="SAME", + name="conv") + x = nn.bias_add(x, bias2) + x = array_ops.identity(x, name="output") + + return trt_test.TfTrtIntegrationTestParams( + gdef=g.as_graph_def(), + input_names=["input"], + input_dims=input_dims, + output_names=["output"], + expected_output_dims=expected_output_dims) + + def GetConversionParams(self, run_params): + """Return a ConversionParams for test.""" + conversion_params = super(DynamicInputShapesTest, + self).GetConversionParams(run_params) + return conversion_params._replace( + maximum_cached_engines=10, + # Disable layout optimizer, since it will convert BiasAdd with NHWC + # format to NCHW format under four dimentional input. + rewriter_config=trt_test.OptimizerDisabledRewriterConfig()) + + def ExpectedEnginesToBuild(self, run_params): + return ["TRTEngineOp_0"] + + def ShouldRunTest(self, run_params): + return (run_params.dynamic_engine and + not trt_test.IsQuantizationMode(run_params.precision_mode)) + + def ExpectedAbsoluteTolerance(self, run_params): + """The absolute tolerance to compare floating point results.""" + return 1.e-03 if run_params.precision_mode == "FP32" else 1.e-01 + + def ExpectedRelativeTolerance(self, run_params): + """The relative tolerance to compare floating point results.""" + return 1.e-03 if run_params.precision_mode == "FP32" else 1.e-01 + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/tensorrt/test/identity_output_test.py b/tensorflow/contrib/tensorrt/test/identity_output_test.py new file mode 100644 index 0000000000000000000000000000000000000000..b568eeda945d997a832b7f71a5bfd8c42e127e65 --- /dev/null +++ b/tensorflow/contrib/tensorrt/test/identity_output_test.py @@ -0,0 +1,74 @@ +# 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. +# ============================================================================== +"""This test checks a situation where the same tensor is considered as an output + +multiple times because it has been duplicated by 2+ indentity ops. Previously, +the tensor would be renamed multiple times, overwriting the output binding name +which resulted in a runtime error when the binding would not be found. +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.contrib.tensorrt.test import tf_trt_integration_test_base as trt_test +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 math_ops +from tensorflow.python.platform import test + + +class IdentityTest(trt_test.TfTrtIntegrationTestBase): + + def _ConstOp(self, shape): + return constant_op.constant(np.random.randn(*shape), dtype=dtypes.float32) + + def GetParams(self): + """Testing engine with the same tensor repeated as output via identity.""" + input_name = 'input' + input_dims = [100, 32] + g = ops.Graph() + with g.as_default(): + x = array_ops.placeholder( + dtype=dtypes.float32, shape=input_dims, name=input_name) + + b = self._ConstOp((32, 4)) + x1 = math_ops.matmul(x, b) + b = self._ConstOp((1, 4)) + x1 = x1 + b + + out1 = array_ops.identity(x1, name='output1') + out2 = array_ops.identity(x1, name='output2') + iden1 = array_ops.identity(x1) + out3 = array_ops.identity(iden1, name='output3') + + return trt_test.TfTrtIntegrationTestParams( + gdef=g.as_graph_def(), + input_names=[input_name], + input_dims=[[input_dims]], + output_names=['output1', 'output2', 'output3'], + expected_output_dims=[[[100, 4]] * 3]) + + def ExpectedEnginesToBuild(self, run_params): + """Return the expected engines to build.""" + return ['TRTEngineOp_0'] + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/tensorrt/test/int32_test.py b/tensorflow/contrib/tensorrt/test/int32_test.py new file mode 100644 index 0000000000000000000000000000000000000000..8cf538703880b130322a7dd504094c7a298e6522 --- /dev/null +++ b/tensorflow/contrib/tensorrt/test/int32_test.py @@ -0,0 +1,82 @@ +# 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. +# ============================================================================== +"""Test conversion of graphs involving INT32 tensors and operations.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.contrib.tensorrt.test import tf_trt_integration_test_base as trt_test +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 math_ops +from tensorflow.python.ops import nn +from tensorflow.python.platform import test + + +class ExcludeUnsupportedInt32Test(trt_test.TfTrtIntegrationTestBase): + + def _ConstOp(self, shape, dtype): + return constant_op.constant(np.random.randn(*shape), dtype=dtype) + + def GetParams(self): + """Test exclusion of ops which are not supported in INT32 mode by TF-TRT""" + input_name = 'input' + output_name = 'output' + input_dims = [100, 4] + dtype = dtypes.int32 + g = ops.Graph() + with g.as_default(): + x = array_ops.placeholder(dtype=dtype, shape=input_dims, name=input_name) + b = self._ConstOp((4, 10), dtype) + x = math_ops.matmul(x, b) + b = self._ConstOp((10,), dtype) + x = nn.bias_add(x, b) + x = array_ops.identity(x, name=output_name) + return trt_test.TfTrtIntegrationTestParams( + gdef=g.as_graph_def(), + input_names=[input_name], + input_dims=[[input_dims]], + output_names=[output_name], + expected_output_dims=[[[100, 10]]]) + + def GetConversionParams(self, run_params): + """Return a ConversionParams for test.""" + conversion_params = super(ExcludeUnsupportedInt32Test, + self).GetConversionParams(run_params) + return conversion_params._replace( + max_batch_size=100, + maximum_cached_engines=1, + # Disable layout optimizer, since it will convert BiasAdd with NHWC + # format to NCHW format under four dimentional input. + rewriter_config=trt_test.OptimizerDisabledRewriterConfig()) + + def ExpectedEnginesToBuild(self, run_params): + """Return the expected engines to build.""" + return [] + + def ShouldRunTest(self, run_params): + """Whether to run the test.""" + # TODO(aaroey): Trt 4.0 forbids conversion for tensors with rank <3 in int8 + # mode, which is a bug. Re-enable this when trt library is fixed. + return not trt_test.IsQuantizationMode(run_params.precision_mode) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/tensorrt/test/lru_cache_test.py b/tensorflow/contrib/tensorrt/test/lru_cache_test.py new file mode 100644 index 0000000000000000000000000000000000000000..7702413e6cee667796b7fbf4121c6e0d9118d35c --- /dev/null +++ b/tensorflow/contrib/tensorrt/test/lru_cache_test.py @@ -0,0 +1,78 @@ +# 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. +# ============================================================================== +"""Test LRUCache by running different input batch sizes on same network.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.contrib.tensorrt.test import tf_trt_integration_test_base as trt_test +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 math_ops +from tensorflow.python.ops import nn +from tensorflow.python.platform import test + + +class LRUCacheTest(trt_test.TfTrtIntegrationTestBase): + + def GetParams(self): + dtype = dtypes.float32 + input_name = "input" + input_dims = [[[1, 10, 10, 2]], [[2, 10, 10, 2]], [[4, 10, 10, 2]], + [[2, 10, 10, 2]]] + expected_output_dims = [[[1, 10, 10, 1]], [[2, 10, 10, 1]], [[4, 10, 10, + 1]], + [[2, 10, 10, 1]]] + output_name = "output" + g = ops.Graph() + with g.as_default(): + x = array_ops.placeholder( + dtype=dtype, shape=[None, 10, 10, 2], name=input_name) + conv_filter = constant_op.constant( + np.random.randn(3, 3, 2, 1), dtype=dtypes.float32) + x = nn.conv2d( + input=x, + filter=conv_filter, + strides=[1, 1, 1, 1], + padding="SAME", + name="conv") + bias = constant_op.constant( + np.random.randn(1, 10, 10, 1), dtype=dtypes.float32) + x = math_ops.add(x, bias) + x = nn.relu(x) + x = array_ops.identity(x, name="output") + return trt_test.TfTrtIntegrationTestParams( + gdef=g.as_graph_def(), + input_names=[input_name], + input_dims=input_dims, + output_names=[output_name], + expected_output_dims=expected_output_dims) + + def ExpectedEnginesToBuild(self, run_params): + """Return the expected engines to build.""" + return ["TRTEngineOp_0"] + + def ShouldRunTest(self, run_params): + return (run_params.dynamic_engine and + not trt_test.IsQuantizationMode(run_params.precision_mode)) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/tensorrt/test/manual_test.py b/tensorflow/contrib/tensorrt/test/manual_test.py index 1187c759b4b5483cbf5afe136401abe86d6ef989..aad7b9f30728cbb3f4ec5fa730c5dbe46fe9fc3f 100644 --- a/tensorflow/contrib/tensorrt/test/manual_test.py +++ b/tensorflow/contrib/tensorrt/test/manual_test.py @@ -67,9 +67,9 @@ class ManualTest(trt_test.TfTrtIntegrationTestBase): return trt_test.TfTrtIntegrationTestParams( gdef=gdef, input_names=params_map['input_names'], - input_dims=params_map['input_dims'], + input_dims=[params_map['input_dims']], output_names=params_map['output_names'], - expected_output_dims=params_map['expected_output_dims']) + expected_output_dims=[params_map['expected_output_dims']]) def GetConversionParams(self, run_params): """Return a ConversionParams for test.""" diff --git a/tensorflow/contrib/tensorrt/test/memory_alignment_test.py b/tensorflow/contrib/tensorrt/test/memory_alignment_test.py index bc7c90081ff38a832b523948db10c02de7acefc2..cc64329bbd53eaaebf7929e48bbfa8d8beeeadff 100644 --- a/tensorflow/contrib/tensorrt/test/memory_alignment_test.py +++ b/tensorflow/contrib/tensorrt/test/memory_alignment_test.py @@ -62,13 +62,13 @@ class MemoryAlignmentTest(trt_test.TfTrtIntegrationTestBase): return trt_test.TfTrtIntegrationTestParams( gdef=g.as_graph_def(), input_names=[input_name], - input_dims=[input_dims], + input_dims=[[input_dims]], output_names=[output_name], - expected_output_dims=[(2, 15, 15, 10)]) + expected_output_dims=[[[2, 15, 15, 10]]]) def ExpectedEnginesToBuild(self, run_params): """Return the expected engines to build.""" - return ["my_trt_op_0"] + return ["TRTEngineOp_0"] def ExpectedAbsoluteTolerance(self, run_params): """The absolute tolerance to compare floating point results.""" diff --git a/tensorflow/contrib/tensorrt/test/multi_connection_neighbor_engine_test.py b/tensorflow/contrib/tensorrt/test/multi_connection_neighbor_engine_test.py index 11be4feaf7bf8ce6c8bd16f1546dc17450c342f1..a14bb0396ece74c8de73008d2007bce5c763b0ed 100644 --- a/tensorflow/contrib/tensorrt/test/multi_connection_neighbor_engine_test.py +++ b/tensorflow/contrib/tensorrt/test/multi_connection_neighbor_engine_test.py @@ -25,8 +25,6 @@ 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 gen_math_ops -from tensorflow.python.ops import math_ops from tensorflow.python.ops import nn from tensorflow.python.platform import test @@ -60,14 +58,14 @@ class MultiConnectionNeighborEngineTest(trt_test.TfTrtIntegrationTestBase): b = constant_op.constant( np.random.normal(5.0, 1.0, [1, 4, 1, 1]), name="bias", dtype=dtype) q = conv - b - edge = math_ops.sigmoid(q) + edge = self.trt_incompatible_op(q) b = constant_op.constant( np.random.normal(5.0, 1.0, [1, 4, 1, 1]), name="bias", dtype=dtype) d = b + conv - edge3 = math_ops.sigmoid(d) + edge3 = self.trt_incompatible_op(d) - edge1 = gen_math_ops.tan(conv) + edge1 = self.trt_incompatible_op(conv) t = t - edge1 q = q + edge t = t + q @@ -77,13 +75,13 @@ class MultiConnectionNeighborEngineTest(trt_test.TfTrtIntegrationTestBase): return trt_test.TfTrtIntegrationTestParams( gdef=g.as_graph_def(), input_names=[input_name], - input_dims=[input_dims], + input_dims=[[input_dims]], output_names=[output_name], - expected_output_dims=[(2, 4, 5, 4)]) + expected_output_dims=[[[2, 4, 5, 4]]]) def ExpectedEnginesToBuild(self, run_params): """Return the expected engines to build.""" - return ["my_trt_op_0", "my_trt_op_1"] + return ["TRTEngineOp_0", "TRTEngineOp_1"] if __name__ == "__main__": diff --git a/tensorflow/contrib/tensorrt/test/neighboring_engine_test.py b/tensorflow/contrib/tensorrt/test/neighboring_engine_test.py index eddeafa38bc71743ac6c9d8e5e8db76f28ca7bf4..06a86bbb8df4c11a471475c040b74099a6fe2361 100644 --- a/tensorflow/contrib/tensorrt/test/neighboring_engine_test.py +++ b/tensorflow/contrib/tensorrt/test/neighboring_engine_test.py @@ -59,15 +59,15 @@ class NeighboringEngineTest(trt_test.TfTrtIntegrationTestBase): return trt_test.TfTrtIntegrationTestParams( gdef=g.as_graph_def(), input_names=[input_name], - input_dims=[input_dims], + input_dims=[[input_dims]], output_names=[output_name], - expected_output_dims=[(2, 4, 5, 4)]) + expected_output_dims=[[[2, 4, 5, 4]]]) def ExpectedEnginesToBuild(self, run_params): """Return the expected engines to build.""" return { - "my_trt_op_0": ["bias", "mul", "sub"], - "my_trt_op_1": ["weights", "conv"] + "TRTEngineOp_0": ["bias", "mul", "sub"], + "TRTEngineOp_1": ["weights", "conv"] } diff --git a/tensorflow/contrib/tensorrt/test/quantization_mnist_test.py b/tensorflow/contrib/tensorrt/test/quantization_mnist_test.py new file mode 100644 index 0000000000000000000000000000000000000000..d68211a7ee344f3d07d01e308ee60246a61816f6 --- /dev/null +++ b/tensorflow/contrib/tensorrt/test/quantization_mnist_test.py @@ -0,0 +1,293 @@ +# 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-TRT INT8 conversion without calibration on Mnist model.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +# pylint: disable=unused-import +from tensorflow.compiler.tf2tensorrt.python.ops import trt_ops +# pylint: enable=unused-import +from tensorflow.contrib.tensorrt.python import trt_convert +from tensorflow.core.protobuf import config_pb2 +from tensorflow.python import data +from tensorflow.python import keras +from tensorflow.python.estimator.estimator import Estimator +from tensorflow.python.estimator.model_fn import EstimatorSpec +from tensorflow.python.estimator.model_fn import ModeKeys +from tensorflow.python.estimator.run_config import RunConfig +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import graph_util +from tensorflow.python.framework import importer +from tensorflow.python.framework import ops +from tensorflow.python.framework import test_util +from tensorflow.python.keras.datasets import mnist +from tensorflow.python.layers import layers +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import gen_array_ops +from tensorflow.python.ops import math_ops +from tensorflow.python.ops import metrics +from tensorflow.python.ops import nn +from tensorflow.python.ops import variable_scope +from tensorflow.python.ops.losses import losses +from tensorflow.python.platform import test +from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.summary import summary +from tensorflow.python.training import saver +from tensorflow.python.training.adam import AdamOptimizer +from tensorflow.python.training.checkpoint_management import latest_checkpoint +from tensorflow.python.training.training_util import get_global_step + +INPUT_NODE_NAME = 'input' +OUTPUT_NODE_NAME = 'output' + + +class QuantizationAwareTrainingMNISTTest(test_util.TensorFlowTestCase): + + def _BuildGraph(self, x): + + def _Quantize(x, r): + x = gen_array_ops.quantize_and_dequantize_v2(x, -r, r) + return x + + def _DenseLayer(x, num_inputs, num_outputs, quantization_range, name): + """Dense layer with quantized outputs. + + Args: + x: input to the dense layer + num_inputs: number of input columns of x + num_outputs: number of output columns + quantization_range: the min/max range for quantization + name: name of the variable scope + + Returns: + The output of the layer. + """ + with variable_scope.variable_scope(name): + kernel = variable_scope.get_variable( + 'kernel', + shape=[num_inputs, num_outputs], + dtype=dtypes.float32, + initializer=keras.initializers.glorot_uniform()) + bias = variable_scope.get_variable( + 'bias', + shape=[num_outputs], + dtype=dtypes.float32, + initializer=keras.initializers.zeros()) + x = math_ops.matmul(x, kernel) + x = _Quantize(x, quantization_range) + x = nn.bias_add(x, bias) + x = _Quantize(x, quantization_range) + return x + + x = _Quantize(x, 1) + # Conv + Bias + Relu6 + x = layers.conv2d(x, filters=32, kernel_size=3, use_bias=True) + x = nn.relu6(x) + # Conv + Bias + Relu6 + x = layers.conv2d(x, filters=64, kernel_size=3, use_bias=True) + x = nn.relu6(x) + # Reduce + x = math_ops.reduce_mean(x, [1, 2]) + x = _Quantize(x, 6) + # FC1 + x = _DenseLayer(x, 64, 512, 6, name='dense') + x = nn.relu6(x) + # FC2 + x = _DenseLayer(x, 512, 10, 25, name='dense_1') + x = array_ops.identity(x, name=OUTPUT_NODE_NAME) + return x + + def _GetGraphDef(self, use_trt, max_batch_size, model_dir): + """Get the frozen mnist GraphDef. + + Args: + use_trt: whether use TF-TRT to convert the graph. + max_batch_size: the max batch size to apply during TF-TRT conversion. + model_dir: the model directory to load the checkpoints. + + Returns: + The frozen mnist GraphDef. + """ + graph = ops.Graph() + with self.session(graph=graph) as sess: + with graph.device('/GPU:0'): + x = array_ops.placeholder( + shape=(None, 28, 28, 1), dtype=dtypes.float32, name=INPUT_NODE_NAME) + self._BuildGraph(x) + # Load weights + mnist_saver = saver.Saver() + checkpoint_file = latest_checkpoint(model_dir) + mnist_saver.restore(sess, checkpoint_file) + # Freeze + graph_def = graph_util.convert_variables_to_constants( + sess, sess.graph_def, output_node_names=[OUTPUT_NODE_NAME]) + # Convert with TF-TRT + if use_trt: + logging.info('Number of nodes before TF-TRT conversion: %d', + len(graph_def.node)) + graph_def = trt_convert.create_inference_graph( + graph_def, + outputs=[OUTPUT_NODE_NAME], + max_batch_size=max_batch_size, + precision_mode='INT8', + # There is a 2GB GPU memory limit for each test, so we set + # max_workspace_size_bytes to 256MB to leave enough room for TF + # runtime to allocate GPU memory. + max_workspace_size_bytes=1 << 28, + minimum_segment_size=2, + use_calibration=False, + ) + logging.info('Number of nodes after TF-TRT conversion: %d', + len(graph_def.node)) + num_engines = len( + [1 for n in graph_def.node if str(n.op) == 'TRTEngineOp']) + self.assertEqual(1, num_engines) + return graph_def + + def _Run(self, is_training, use_trt, batch_size, num_epochs, model_dir): + """Train or evaluate the model. + + Args: + is_training: whether to train or evaluate the model. In training mode, + quantization will be simulated where the quantize_and_dequantize_v2 are + placed. + use_trt: if true, use TRT INT8 mode for evaluation, which will perform + real quantization. Otherwise use native TensorFlow which will perform + simulated quantization. Ignored if is_training is True. + batch_size: batch size. + num_epochs: how many epochs to train. Ignored if is_training is False. + model_dir: where to save or load checkpoint. + + Returns: + The Estimator evaluation result. + """ + # Get dataset + train_data, test_data = mnist.load_data() + + def _PreprocessFn(x, y): + x = math_ops.cast(x, dtypes.float32) + x = array_ops.expand_dims(x, axis=2) + x = 2.0 * (x / 255.0) - 1.0 + y = math_ops.cast(y, dtypes.int32) + return x, y + + def _EvalInputFn(): + mnist_x, mnist_y = test_data + dataset = data.Dataset.from_tensor_slices((mnist_x, mnist_y)) + dataset = dataset.apply( + data.experimental.map_and_batch( + map_func=_PreprocessFn, + batch_size=batch_size, + num_parallel_calls=8)) + dataset = dataset.repeat(count=1) + iterator = dataset.make_one_shot_iterator() + features, labels = iterator.get_next() + return features, labels + + def _TrainInputFn(): + mnist_x, mnist_y = train_data + dataset = data.Dataset.from_tensor_slices((mnist_x, mnist_y)) + dataset = dataset.shuffle(2 * len(mnist_x)) + dataset = dataset.apply( + data.experimental.map_and_batch( + map_func=_PreprocessFn, + batch_size=batch_size, + num_parallel_calls=8)) + dataset = dataset.repeat(count=num_epochs) + iterator = dataset.make_one_shot_iterator() + features, labels = iterator.get_next() + return features, labels + + def _ModelFn(features, labels, mode): + if is_training: + logits_out = self._BuildGraph(features) + else: + graph_def = self._GetGraphDef(use_trt, batch_size, model_dir) + logits_out = importer.import_graph_def( + graph_def, + input_map={INPUT_NODE_NAME: features}, + return_elements=[OUTPUT_NODE_NAME + ':0'], + name='')[0] + + loss = losses.sparse_softmax_cross_entropy( + labels=labels, logits=logits_out) + summary.scalar('loss', loss) + + classes_out = math_ops.argmax(logits_out, axis=1, name='classes_out') + accuracy = metrics.accuracy( + labels=labels, predictions=classes_out, name='acc_op') + summary.scalar('accuracy', accuracy[1]) + + if mode == ModeKeys.EVAL: + return EstimatorSpec( + mode, loss=loss, eval_metric_ops={'accuracy': accuracy}) + elif mode == ModeKeys.TRAIN: + optimizer = AdamOptimizer(learning_rate=1e-2) + train_op = optimizer.minimize(loss, global_step=get_global_step()) + return EstimatorSpec(mode, loss=loss, train_op=train_op) + + config_proto = config_pb2.ConfigProto() + config_proto.gpu_options.allow_growth = True + estimator = Estimator( + model_fn=_ModelFn, + model_dir=model_dir if is_training else None, + config=RunConfig(session_config=config_proto)) + + if is_training: + estimator.train(_TrainInputFn) + results = estimator.evaluate(_EvalInputFn) + logging.info('accuracy: %s', str(results['accuracy'])) + return results + + # To generate the checkpoint, set a different model_dir and call self._Run() + # by setting is_training=True and num_epochs=1000, e.g.: + # model_dir = '/tmp/quantization_mnist' + # self._Run( + # is_training=True, + # use_trt=False, + # batch_size=128, + # num_epochs=100, + # model_dir=model_dir) + def testEval(self): + if not trt_convert.is_tensorrt_enabled(): + return + model_dir = test.test_src_dir_path('contrib/tensorrt/test/testdata') + + accuracy_tf_native = self._Run( + is_training=False, + use_trt=False, + batch_size=128, + num_epochs=None, + model_dir=model_dir)['accuracy'] + logging.info('accuracy_tf_native: %f', accuracy_tf_native) + self.assertAllClose(0.9662, accuracy_tf_native, rtol=1e-3, atol=1e-3) + + if trt_convert.get_linked_tensorrt_version()[0] < 5: + return + + accuracy_tf_trt = self._Run( + is_training=False, + use_trt=True, + batch_size=128, + num_epochs=None, + model_dir=model_dir)['accuracy'] + logging.info('accuracy_tf_trt: %f', accuracy_tf_trt) + self.assertAllClose(0.9675, accuracy_tf_trt, rtol=1e-3, atol=1e-3) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/tensorrt/test/quantization_test.py b/tensorflow/contrib/tensorrt/test/quantization_test.py new file mode 100644 index 0000000000000000000000000000000000000000..ce1b25ebf3c52ac5710dea654134925bb5b6ceca --- /dev/null +++ b/tensorflow/contrib/tensorrt/test/quantization_test.py @@ -0,0 +1,144 @@ +# 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. +# ============================================================================== +"""Model script to test TF-TensorRT integration.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.contrib.tensorrt.python import trt_convert +from tensorflow.contrib.tensorrt.test import tf_trt_integration_test_base as trt_test +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 gen_array_ops +from tensorflow.python.ops import math_ops +from tensorflow.python.platform import test + + +def _GetParams(add_quantization_nodes, dtype=dtypes.float32): + input_name = "input" + input_dims = [8, 8] + output_name = "output" + + def _Quantize(x, r): + if add_quantization_nodes: + x = gen_array_ops.fake_quant_with_min_max_vars(x, -r, r) + return x + + g = ops.Graph() + with g.as_default(): + x = array_ops.placeholder( + dtype=dtype, shape=[None] + input_dims[1:], name=input_name) + x = _Quantize(x, 10.0) + x = x + 5 + x = _Quantize(x, 15.0) + x = x - 5 + x = _Quantize(x, 10.0) + x = x * 0.1 + x = _Quantize(x, 1.0) + w = constant_op.constant(np.ones((8, 1)), dtype=dtypes.float32) + x = math_ops.matmul(x, w) + x = _Quantize(x, 10.0) + x = array_ops.identity(x, name=output_name) + + return trt_test.TfTrtIntegrationTestParams( + gdef=g.as_graph_def(), + input_names=[input_name], + input_dims=[[input_dims]], + output_names=[output_name], + expected_output_dims=[[[8, 1]]]) + + +class QuantizationMissingAllRangesTest(trt_test.TfTrtIntegrationTestBase): + + def GetParams(self): + """Create a graph containing single segment with no quantization ranges.""" + return _GetParams(add_quantization_nodes=False) + + def ShouldRunTest(self, run_params): + if trt_convert.get_linked_tensorrt_version()[0] < 5: + return False + # Only test static engine mode, with or without calibration. + return (trt_test.IsQuantizationMode(run_params.precision_mode) and + not run_params.use_optimizer and not run_params.dynamic_engine) + + def ExpectedEnginesToBuild(self, run_params): + """Return the expected engines to build.""" + if run_params.use_calibration: + # In static engine mode with calibration, it should build a calibration + # engine. + return ["TRTEngineOp_0"] + # In static engine mode without calibration, the engine building will fail + # since no quantization ranges are set, which results in no TRT nodes. + return [] + + +class QuantizationWithRangesTest(trt_test.TfTrtIntegrationTestBase): + + def GetParams(self): + """Create a graph containing single segment with no quantization ranges.""" + return _GetParams(add_quantization_nodes=True) + + def ShouldRunTest(self, run_params): + if trt_convert.get_linked_tensorrt_version()[0] < 5: + return False + # Test static/dynamic engine with/without calibration. + return (trt_test.IsQuantizationMode(run_params.precision_mode) and + not run_params.use_optimizer) + + def ExpectedEnginesToBuild(self, run_params): + """Return the expected engines to build.""" + return ["TRTEngineOp_0"] + + def ExpectedAbsoluteTolerance(self, run_params): + """The absolute tolerance to compare floating point results.""" + return 1.e-05 if run_params.precision_mode == "FP32" else 1.e-01 + + def ExpectedRelativeTolerance(self, run_params): + """The relative tolerance to compare floating point results.""" + return 1.e-05 if run_params.precision_mode == "FP32" else 1.e-01 + + +class NonQuantizedPrecisionsWithRangesTest(trt_test.TfTrtIntegrationTestBase): + + def GetParams(self): + """Create a graph containing single segment with no quantization ranges.""" + return _GetParams(add_quantization_nodes=True) + + def ShouldRunTest(self, run_params): + # Only test FP32/FP16 mode. + return not trt_test.IsQuantizationMode(run_params.precision_mode) + + def ExpectedEnginesToBuild(self, run_params): + """Return the expected engines to build.""" + # The fake quant ops are not supported in FP32/FP16 mode, and will split the + # graph into three TRT segments. + return ["TRTEngineOp_0", "TRTEngineOp_1", "TRTEngineOp_2", "TRTEngineOp_3"] + + def ExpectedAbsoluteTolerance(self, run_params): + """The absolute tolerance to compare floating point results.""" + return 1.e-05 if run_params.precision_mode == "FP32" else 1.e-01 + + def ExpectedRelativeTolerance(self, run_params): + """The relative tolerance to compare floating point results.""" + return 1.e-05 if run_params.precision_mode == "FP32" else 1.e-01 + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/tensorrt/test/rank_two_test.py b/tensorflow/contrib/tensorrt/test/rank_two_test.py index 74a4a059257ffde4c86df1f18b3ce35c3790ec7a..97159bb008068efbbcdb0fc6844890a42a08f46c 100644 --- a/tensorflow/contrib/tensorrt/test/rank_two_test.py +++ b/tensorflow/contrib/tensorrt/test/rank_two_test.py @@ -51,8 +51,10 @@ class RankTwoTest(trt_test.TfTrtIntegrationTestBase): c = constant_op.constant(3.0, name="c%d_3" % i) q = math_ops.add(q, c, name="add%d_3" % i) if i == 0: + axis = constant_op.constant(-1, dtype=dtypes.int32, name="axis") for j in range(2): - q = array_ops.expand_dims(q, -1, name="expand%d_%d" % (i, j)) + q = array_ops.expand_dims(q, axis, name="expand%d_%d" % (i, j)) + q = self.trt_incompatible_op(q) q = gen_math_ops.reciprocal(q, name="reciprocal%d" % i) outputs.append(q) # Combine both paths @@ -61,29 +63,23 @@ class RankTwoTest(trt_test.TfTrtIntegrationTestBase): return trt_test.TfTrtIntegrationTestParams( gdef=g.as_graph_def(), input_names=input_names, - input_dims=input_dims, + input_dims=[input_dims], output_names=[output_name], - expected_output_dims=[tuple(input_dims[1])]) + expected_output_dims=[[input_dims[1]]]) def ExpectedEnginesToBuild(self, run_params): """Return the expected engines to build.""" return { - "my_trt_op_0": [ + "TRTEngineOp_0": [ "add0_1", "add0_2", "add0_3", "c0_1", "c0_2", "c0_3", "abs0_1", - "abs0_2" + "abs0_2", "expand0_0", "expand0_1", "axis" ], - "my_trt_op_1": [ + "TRTEngineOp_1": [ "add", "add1_1", "add1_2", "add1_3", "c1_1", "c1_2", "c1_3", "abs1_1", "abs1_2", "reciprocal0", "reciprocal1" ], } - def ShouldRunTest(self, run_params): - """Whether to run the test.""" - # TODO(aaroey): Trt 4.0 forbids conversion for tensors with rank <3 in int8 - # mode, which is a bug. Re-enable this when trt library is fixed. - return not trt_test.IsQuantizationMode(run_params.precision_mode) - if __name__ == "__main__": test.main() diff --git a/tensorflow/contrib/tensorrt/test/reshape_transpose_test.py b/tensorflow/contrib/tensorrt/test/reshape_transpose_test.py index 3cf7dadb1f472247250687a0879df431bd81ee94..7fb2cbde07c4987d925e9abc915ede52119ec6df 100644 --- a/tensorflow/contrib/tensorrt/test/reshape_transpose_test.py +++ b/tensorflow/contrib/tensorrt/test/reshape_transpose_test.py @@ -72,15 +72,15 @@ class ReshapeTest(trt_test.TfTrtIntegrationTestBase): return trt_test.TfTrtIntegrationTestParams( gdef=g.as_graph_def(), input_names=[input_name], - input_dims=[input_dims], + input_dims=[[input_dims]], output_names=[output_name], - expected_output_dims=[tuple(input_dims)]) + expected_output_dims=[[input_dims]]) def ExpectedEnginesToBuild(self, run_params): """Return the expected engines to build.""" return { - "my_trt_op_3": ["reshape-%d" % i for i in range(7)] + - ["reshape-%d/shape" % i for i in range(7)] + "TRTEngineOp_0": ["reshape-%d" % i for i in range(7)] + + ["reshape-%d/shape" % i for i in range(7)] } def ShouldRunTest(self, run_params): @@ -117,7 +117,7 @@ class TransposeTest(trt_test.TfTrtIntegrationTestBase): # Note: by default Grappler will run the TRT optimizer twice. At the # first time it will group the two transpose ops below to same segment # then fail the conversion due to the expected batch dimension problem. - # At the second time, since the input of bridge op is my_trt_op_0, it + # At the second time, since the input of bridge op is TRTEngineOp_0, it # will fail to do shape inference which then cause conversion to fail. # TODO(laigd): support shape inference, make TRT optimizer run only # once, and fix this. @@ -129,14 +129,14 @@ class TransposeTest(trt_test.TfTrtIntegrationTestBase): return trt_test.TfTrtIntegrationTestParams( gdef=g.as_graph_def(), input_names=[input_name], - input_dims=[input_dims], + input_dims=[[input_dims]], output_names=[output_name], - expected_output_dims=[(24, 100, 2, 24)]) + expected_output_dims=[[[24, 100, 2, 24]]]) def ExpectedEnginesToBuild(self, run_params): """Return the expected engines to build.""" return { - "my_trt_op_0": [ + "TRTEngineOp_0": [ "transpose-1", "transpose-1/perm", "transposeback", "transposeback/perm" ] diff --git a/tensorflow/contrib/tensorrt/test/test_tftrt.py b/tensorflow/contrib/tensorrt/test/test_tftrt.py index d26f26008635733c6c364a98b72b88c1e552f5fe..090aa8bdb0487973e186631af3b4edac48096a5f 100644 --- a/tensorflow/contrib/tensorrt/test/test_tftrt.py +++ b/tensorflow/contrib/tensorrt/test/test_tftrt.py @@ -191,7 +191,7 @@ def user(multi_engine, minimum_segment_size=2, # minimum number of nodes in an engine is_dynamic_op=False, maximum_cached_engines=1, - cached_engine_batch_sizes=[]) + cached_engine_batches=[]) o1 = run_graph(orig_graph, dummy_input) o2 = run_graph(trt_graph, dummy_input) o3 = run_graph(trt_graph, dummy_input) @@ -206,7 +206,7 @@ def user(multi_engine, minimum_segment_size=2, # minimum number of nodes in an engine is_dynamic_op=False, maximum_cached_engines=1, - cached_engine_batch_sizes=[]) + cached_engine_batches=[]) int8_calib_gdef = trt.create_inference_graph( input_graph_def=orig_graph, outputs=["output"], @@ -216,7 +216,7 @@ def user(multi_engine, minimum_segment_size=2, # minimum number of nodes in an engine is_dynamic_op=False, maximum_cached_engines=1, - cached_engine_batch_sizes=[]) + cached_engine_batches=[]) 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) diff --git a/tensorflow/contrib/tensorrt/test/testdata/checkpoint b/tensorflow/contrib/tensorrt/test/testdata/checkpoint new file mode 100644 index 0000000000000000000000000000000000000000..a603e1aec91adab04fd9801ba05a2ee9adfbb6e8 --- /dev/null +++ b/tensorflow/contrib/tensorrt/test/testdata/checkpoint @@ -0,0 +1,3 @@ +model_checkpoint_path: "model.ckpt-46900" +all_model_checkpoint_paths: "model.ckpt-0" +all_model_checkpoint_paths: "model.ckpt-46900" diff --git a/tensorflow/contrib/tensorrt/test/testdata/model.ckpt-46900.data-00000-of-00001 b/tensorflow/contrib/tensorrt/test/testdata/model.ckpt-46900.data-00000-of-00001 new file mode 100644 index 0000000000000000000000000000000000000000..88a998f184b275121e1e76eb51d2310da149f10a Binary files /dev/null and b/tensorflow/contrib/tensorrt/test/testdata/model.ckpt-46900.data-00000-of-00001 differ diff --git a/tensorflow/contrib/tensorrt/test/testdata/model.ckpt-46900.index b/tensorflow/contrib/tensorrt/test/testdata/model.ckpt-46900.index new file mode 100644 index 0000000000000000000000000000000000000000..537976571337508ab1798d33646c51d62a146ecc Binary files /dev/null and b/tensorflow/contrib/tensorrt/test/testdata/model.ckpt-46900.index differ diff --git a/tensorflow/contrib/tensorrt/test/tf_trt_integration_test_base.py b/tensorflow/contrib/tensorrt/test/tf_trt_integration_test_base.py index a725d0651c92fe18bcfd284cffd40cdfec2e6c69..9a00cdb11a0f98d9b9be0d8e88a79cf45a8a7e3a 100644 --- a/tensorflow/contrib/tensorrt/test/tf_trt_integration_test_base.py +++ b/tensorflow/contrib/tensorrt/test/tf_trt_integration_test_base.py @@ -25,11 +25,12 @@ import warnings import numpy as np import six -from tensorflow.contrib.tensorrt.python import trt_convert # pylint: disable=unused-import -from tensorflow.contrib.tensorrt.python.ops import trt_engine_op +from tensorflow.compiler.tf2tensorrt.python.ops import trt_ops # pylint: enable=unused-import +from tensorflow.contrib.tensorrt.python import trt_convert from tensorflow.core.protobuf import config_pb2 +from tensorflow.core.protobuf import rewriter_config_pb2 from tensorflow.python.framework import dtypes from tensorflow.python.framework import graph_io from tensorflow.python.framework import importer @@ -38,18 +39,29 @@ from tensorflow.python.framework import test_util from tensorflow.python.ops import math_ops from tensorflow.python.platform import tf_logging as logging -TfTrtIntegrationTestParams = namedtuple("TfTrtIntegrationTestParams", [ - "gdef", "input_names", "input_dims", "output_names", "expected_output_dims" +TfTrtIntegrationTestParams = namedtuple( + "TfTrtIntegrationTestParams", + [ + "gdef", + # A list of names of the input placeholder nodes. + "input_names", + # A list of list of output shapes of the input placeholder nodes. + "input_dims", + # A list of names of the output identity nodes. + "output_names", + # A list of list of expected output shapes of the output identity nodes. + "expected_output_dims" + ]) + +RunParams = namedtuple("RunParams", [ + "use_optimizer", "precision_mode", "dynamic_engine", "test_name", + "use_calibration" ]) -RunParams = namedtuple( - "RunParams", - ["use_optimizer", "precision_mode", "dynamic_engine", "test_name"]) - ConversionParams = namedtuple("ConversionParams", [ "max_batch_size", "max_workspace_size_bytes", "precision_mode", "minimum_segment_size", "is_dynamic_op", "maximum_cached_engines", - "cached_engine_batch_sizes", "rewriter_config" + "cached_engine_batches", "rewriter_config", "use_calibration" ]) PRECISION_MODES = ["FP32", "FP16", "INT8"] @@ -65,6 +77,34 @@ class GraphState(object): INFERENCE = 2 +def OptimizerDisabledRewriterConfig(): + """Returns a RewriterConfig with all default Grappler optimizers disabled.""" + rewriter_config = rewriter_config_pb2.RewriterConfig() + + # Turn off all default Grappler optimizers. + off = rewriter_config_pb2.RewriterConfig.OFF + rewriter_config.layout_optimizer = off + rewriter_config.constant_folding = off + rewriter_config.shape_optimization = off + rewriter_config.remapping = off + rewriter_config.arithmetic_optimization = off + rewriter_config.dependency_optimization = off + rewriter_config.loop_optimization = off + rewriter_config.function_optimization = off + rewriter_config.debug_stripper = off + rewriter_config.disable_model_pruning = True + rewriter_config.scoped_allocator_optimization = off + rewriter_config.memory_optimization = ( + rewriter_config_pb2.RewriterConfig.NO_MEM_OPT) + rewriter_config.pin_to_host_optimization = off + rewriter_config.auto_parallel.enable = False + + # Run only once for each enabled optimizer. + rewriter_config.meta_optimizer_iterations = ( + rewriter_config_pb2.RewriterConfig.ONE) + return rewriter_config + + class TfTrtIntegrationTestBase(test_util.TensorFlowTestCase): """Class to test Tensorflow-TensorRT integration.""" @@ -129,21 +169,33 @@ class TfTrtIntegrationTestBase(test_util.TensorFlowTestCase): def GetConversionParams(self, run_params): """Return a ConversionParams for test.""" + batch_list = [] + for dims_list in self._GetParamsCached().input_dims: + assert dims_list + # Each list of shapes should have same batch size. + input_batches = [dims[0] for dims in dims_list] + assert max(input_batches) == min(input_batches) + batch_list.append(input_batches[0]) return ConversionParams( - max_batch_size=max([ - dims[0] for dims in self._GetParamsCached().input_dims if len(dims) - ]), + # We use the minimum of all the batch sizes, so when multiple different + # input shapes are provided it'll always create new engines in the + # cache, and we can therefore test the cache behavior. + max_batch_size=min(batch_list), max_workspace_size_bytes=1 << 25, precision_mode=run_params.precision_mode, minimum_segment_size=2, is_dynamic_op=run_params.dynamic_engine, maximum_cached_engines=1, - cached_engine_batch_sizes=None, - rewriter_config=None) + cached_engine_batches=None, + rewriter_config=None, + use_calibration=run_params.use_calibration) def ShouldRunTest(self, run_params): """Whether to run the test.""" - return True + # This setting combination requires quantization nodes to be present in + # order to build the engine. + return not (IsQuantizationMode(run_params.precision_mode) and + not run_params.use_calibration) def VerifyRunForEngine(self, engine_name, graph_state, expect_run=True): """Verify the state of a particular engine after sess.run().""" @@ -194,34 +246,38 @@ class TfTrtIntegrationTestBase(test_util.TensorFlowTestCase): def _PrepareRun(self, graph_state): """Set up necessary testing environment before calling sess.run().""" # Clear test values added by TRTEngineOp. - trt_convert.clear_test_values("my_trt_op_.*:ExecuteTrtEngine") - trt_convert.clear_test_values("my_trt_op_.*:ExecuteCalibration") - trt_convert.clear_test_values("my_trt_op_.*:ExecuteNativeSegment") + trt_convert.clear_test_values("TRTEngineOp_.*:ExecuteTrtEngine") + trt_convert.clear_test_values("TRTEngineOp_.*:ExecuteCalibration") + trt_convert.clear_test_values("TRTEngineOp_.*:ExecuteNativeSegment") + + def _GetGPUOptions(self): + gpu_options = config_pb2.GPUOptions() + gpu_options.allow_growth = True + return gpu_options def _GetConfigProto(self, run_params, graph_state): """Get config proto based on specific settings.""" + conversion_params = self.GetConversionParams(run_params) if graph_state != GraphState.ORIGINAL and run_params.use_optimizer: - conversion_params = self.GetConversionParams(run_params) - rewriter_cfg = trt_convert.tensorrt_rewriter_config( + rewriter_cfg = trt_convert.get_tensorrt_rewriter_config( conversion_params.rewriter_config, conversion_params.max_batch_size, conversion_params.max_workspace_size_bytes, conversion_params.precision_mode, conversion_params.minimum_segment_size, conversion_params.is_dynamic_op, conversion_params.maximum_cached_engines, - conversion_params.cached_engine_batch_sizes) + conversion_params.cached_engine_batches, + conversion_params.use_calibration) graph_options = config_pb2.GraphOptions(rewrite_options=rewriter_cfg) else: graph_options = config_pb2.GraphOptions() - - gpu_options = config_pb2.GPUOptions() - gpu_options.allow_growth = True - if trt_convert.get_linked_tensorrt_version()[0] == 3: - gpu_options.per_process_gpu_memory_fraction = 0.50 + if conversion_params.rewriter_config is not None: + graph_options.rewrite_options.CopyFrom( + conversion_params.rewriter_config) config = config_pb2.ConfigProto( - gpu_options=gpu_options, graph_options=graph_options) + gpu_options=self._GetGPUOptions(), graph_options=graph_options) return config def _ExpectTestValue(self, engine_name, method, expected_value): @@ -245,13 +301,16 @@ class TfTrtIntegrationTestBase(test_util.TensorFlowTestCase): def _RunGraph(self, run_params, gdef, - input_data, + inputs_data, config, graph_state, num_runs=2): """Run given graphdef multiple times.""" params = self._GetParamsCached() - assert len(params.input_names) == len(input_data) + for current_input_data in inputs_data: + assert len(params.input_names) == len(current_input_data) + + vals = [] g = ops.Graph() with g.as_default(): io_ops = importer.import_graph_def( @@ -259,38 +318,48 @@ class TfTrtIntegrationTestBase(test_util.TensorFlowTestCase): return_elements=params.input_names + params.output_names, name="") inputs = [op.outputs[0] for op in io_ops[:len(params.input_names)]] - assert len(inputs) == len(input_data) + for current_input_data in inputs_data: + assert len(inputs) == len(current_input_data) outputs = [op.outputs[0] for op in io_ops[len(params.input_names):]] - with self.test_session( - graph=g, config=config, use_gpu=True, force_gpu=True) as sess: - val = None - # Defaults to 2 runs to verify result across multiple runs is same. - for _ in range(num_runs): - self._PrepareRun(graph_state) - new_val = sess.run( - outputs, {inputs[i]: input_data[i] for i in range(len(inputs))}) - output_len = len(params.expected_output_dims) - self.assertEqual(output_len, len(new_val)) - for i in range(output_len): - self.assertEqual(params.expected_output_dims[i], new_val[i].shape) - if val is not None: - self.assertAllClose(val, new_val, atol=1.e-06, rtol=1.e-06) - val = new_val - self.VerifyRun(run_params, graph_state) - return val + with self.test_session( + graph=g, config=config, use_gpu=True, force_gpu=True) as sess: + # Run for each input(s) shape + for shape_index in range(len(inputs_data)): + val = None + # Defaults to 2 runs to verify result across multiple runs is same. + for _ in range(num_runs): + self._PrepareRun(graph_state) + new_val = sess.run(outputs, { + inputs[i]: inputs_data[shape_index][i] + for i in range(len(inputs)) + }) + output_len = len(params.expected_output_dims[shape_index]) + self.assertEqual(output_len, len(new_val)) + for i in range(output_len): + self.assertEqual( + list(params.expected_output_dims[shape_index][i]), + list(new_val[i].shape)) + if val is not None: + self.assertAllClose(val, new_val, atol=1.e-06, rtol=1.e-06) + val = new_val + self.VerifyRun(run_params, graph_state) + vals.append(val) + return vals # Use real data that is representative of the inference dataset # for calibration. For this test script it is random data. - def _RunCalibration(self, run_params, gdef, input_data, config): + def _RunCalibration(self, run_params, gdef, inputs_data, config): """Run calibration on given graph.""" return self._RunGraph( - run_params, gdef, input_data, config, GraphState.CALIBRATE, num_runs=5) + run_params, gdef, inputs_data, config, GraphState.CALIBRATE, num_runs=5) - def _GetTrtGraphDef(self, run_params, gdef): + def _GetTrtGraphDef(self, run_params, graph_state, gdef): """Return trt converted graphdef.""" params = self._GetParamsCached() conversion_params = self.GetConversionParams(run_params) logging.info(conversion_params) + + config_for_trt = self._GetConfigProto(run_params, graph_state) return trt_convert.create_inference_graph( input_graph_def=gdef, outputs=params.input_names + params.output_names, @@ -300,8 +369,9 @@ class TfTrtIntegrationTestBase(test_util.TensorFlowTestCase): minimum_segment_size=conversion_params.minimum_segment_size, is_dynamic_op=conversion_params.is_dynamic_op, maximum_cached_engines=conversion_params.maximum_cached_engines, - cached_engine_batch_sizes=conversion_params.cached_engine_batch_sizes, - rewriter_config=conversion_params.rewriter_config) + cached_engine_batches=conversion_params.cached_engine_batches, + use_calibration=conversion_params.use_calibration, + session_config=config_for_trt) def _WriteGraph(self, run_params, gdef, graph_state): if graph_state == GraphState.ORIGINAL: @@ -400,10 +470,12 @@ class TfTrtIntegrationTestBase(test_util.TensorFlowTestCase): is_dynamic_engine = not node.attr["static_engine"].b self.assertEqual(run_params.dynamic_engine, is_dynamic_engine, node.name) + self.assertEqual(node.attr["use_calibration"].b, + run_params.use_calibration, node.name) has_calibration_data = len(node.attr["calibration_data"].s) if (IsQuantizationMode(run_params.precision_mode) and - graph_state == GraphState.INFERENCE): + run_params.use_calibration and graph_state == GraphState.INFERENCE): self.assertTrue(has_calibration_data, node.name) else: self.assertFalse(has_calibration_data, node.name) @@ -431,35 +503,47 @@ class TfTrtIntegrationTestBase(test_util.TensorFlowTestCase): dtypes.as_dtype(node.attr["dtype"].type).as_numpy_dtype()) assert len(params.input_names) == len(input_dtypes) - input_data = [] - for i in range(len(params.input_names)): - dtype = input_dtypes[params.input_names[i]] - # Multiply the input by some constant to avoid all zeros input for integer - # types. - scale = 10.0 if np.issubdtype(dtype, np.integer) else 1.0 - dims = params.input_dims[i] - input_data.append((scale * np.random.random_sample(dims)).astype(dtype)) + inputs_data = [] + for inp in params.input_dims: + current_input_data = [] + for i in range(len(params.input_names)): + dtype = input_dtypes[params.input_names[i]] + # Multiply the input by some constant to avoid all zeros input for + # integer types. + scale = 10.0 if np.issubdtype(dtype, np.integer) else 1.0 + dims = inp[i] + # TODO(laigd): add debug options. E.g. we can set the input data to be + # continuous natural numbers: + # seq = np.arange(np.prod(dims)) + # seq.resize(dims) + # input_data.append(scale * seq.astype(dtype)) + current_input_data.append( + (scale * np.random.random_sample(dims)).astype(dtype)) + inputs_data.append(current_input_data) + self._VerifyGraphDef(run_params, input_gdef, GraphState.ORIGINAL) # Get reference result without running trt. config_no_trt = self._GetConfigProto(run_params, GraphState.ORIGINAL) logging.info("Running original graph w/o trt, config:\n%s", str(config_no_trt)) - ref_result = self._RunGraph(run_params, input_gdef, input_data, + ref_result = self._RunGraph(run_params, input_gdef, inputs_data, config_no_trt, GraphState.ORIGINAL) # Run calibration if necessary. - if IsQuantizationMode(run_params.precision_mode): + if (IsQuantizationMode(run_params.precision_mode) and + run_params.use_calibration): calib_config = self._GetConfigProto(run_params, GraphState.CALIBRATE) logging.info("Running calibration graph, config:\n%s", str(calib_config)) if run_params.use_optimizer: - result = self._RunCalibration(run_params, input_gdef, input_data, + result = self._RunCalibration(run_params, input_gdef, inputs_data, calib_config) else: - calib_gdef = self._GetTrtGraphDef(run_params, input_gdef) + calib_gdef = self._GetTrtGraphDef(run_params, GraphState.CALIBRATE, + input_gdef) self._VerifyGraphDef(run_params, calib_gdef, GraphState.CALIBRATE) - result = self._RunCalibration(run_params, calib_gdef, input_data, + result = self._RunCalibration(run_params, calib_gdef, inputs_data, calib_config) infer_gdef = trt_convert.calib_graph_to_infer_graph( calib_gdef, run_params.dynamic_engine) @@ -478,10 +562,11 @@ class TfTrtIntegrationTestBase(test_util.TensorFlowTestCase): logging.info("Running final inference graph, config:\n%s", str(infer_config)) if not run_params.use_optimizer: - infer_gdef = self._GetTrtGraphDef(run_params, infer_gdef) + infer_gdef = self._GetTrtGraphDef(run_params, GraphState.INFERENCE, + infer_gdef) self._VerifyGraphDef(run_params, infer_gdef, GraphState.INFERENCE) - result = self._RunGraph(run_params, infer_gdef, input_data, infer_config, + result = self._RunGraph(run_params, infer_gdef, inputs_data, infer_config, GraphState.INFERENCE) self.assertAllClose( ref_result, @@ -519,27 +604,38 @@ def _AddTests(test_class): use_optimizer_options = [False, True] dynamic_engine_options = [False, True] - for (use_optimizer, precision_mode, dynamic_engine) in itertools.product( - use_optimizer_options, PRECISION_MODES, dynamic_engine_options): + use_calibration_options = [False, True] + opts = itertools.product(use_optimizer_options, PRECISION_MODES, + dynamic_engine_options, use_calibration_options) + for (use_optimizer, precision_mode, dynamic_engine, use_calibration) in opts: if IsQuantizationMode(precision_mode): if use_optimizer: # TODO(aaroey): if use_optimizer is True we need to get the inference # graphdef using custom python wrapper class, which is not currently # supported yet. continue - if not dynamic_engine: + if use_calibration and not dynamic_engine: + # Static engine with use_calibration=False will be static, so we want to + # test that. If use_calibration=True, only dynamic op is supported. # TODO(aaroey): construction of static calibration engine is not # supported yet. continue + else: + if use_calibration: + # Don't calibrate in FP32 or FP16 mode + continue conversion = "OptimizerConversion" if use_optimizer else "ToolConversion" - engine_type = ("DynamicEngine" if dynamic_engine else "StaticEngine") - test_name = "%s_%s_%s" % (conversion, precision_mode, engine_type) + engine_type = "DynamicEngine" if dynamic_engine else "StaticEngine" + calibration_type = "UseCalibration" if use_calibration else "NoCalibration" + test_name = "%s_%s_%s_%s" % (conversion, engine_type, precision_mode, + calibration_type) run_params = RunParams( use_optimizer=use_optimizer, precision_mode=precision_mode, dynamic_engine=dynamic_engine, - test_name=test_name) + test_name=test_name, + use_calibration=use_calibration) setattr(test_class, "testTfTrt_" + test_name, _GetTest(run_params)) diff --git a/tensorflow/contrib/tensorrt/test/topk_test.py b/tensorflow/contrib/tensorrt/test/topk_test.py new file mode 100644 index 0000000000000000000000000000000000000000..633a51982b9a6acf1926033628793c1edbd2d118 --- /dev/null +++ b/tensorflow/contrib/tensorrt/test/topk_test.py @@ -0,0 +1,58 @@ +# 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. +# ============================================================================== +"""Model script to test TF-TensorRT integration.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.tensorrt.test import tf_trt_integration_test_base as trt_test +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import ops +from tensorflow.python.framework import constant_op +from tensorflow.python.ops import nn_ops +from tensorflow.python.ops import array_ops +from tensorflow.python.platform import test + + +class TopKTest(trt_test.TfTrtIntegrationTestBase): + + def GetParams(self): + """Testing Top-K in TF-TRT conversion.""" + dtype = dtypes.float32 + input_name = "input" + input_dims = [100, 100] + k = 5 + g = ops.Graph() + with g.as_default(): + x = array_ops.placeholder(dtype=dtype, shape=input_dims, name=input_name) + k_tensor = constant_op.constant(k, dtype=dtypes.int32, name="Const") + values, indices = nn_ops.top_k(x, k_tensor, name="TopK") + values = array_ops.identity(values, name="output_values") + indices = array_ops.identity(indices, name="output_indices") + return trt_test.TfTrtIntegrationTestParams( + gdef=g.as_graph_def(), + input_names=[input_name], + input_dims=[[input_dims]], + output_names=["output_values", "output_indices"], + expected_output_dims=[[[100, k], [100, k]]]) + + def ExpectedEnginesToBuild(self, run_params): + """Return the expected engines to build.""" + return {"TRTEngineOp_0": ["Const", "TopK"]} + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/contrib/tensorrt/test/unary_test.py b/tensorflow/contrib/tensorrt/test/unary_test.py index 8736bfb6449b3c25a411ec081ad58b1f8be84617..497ea2848aae42a61db4f8f5a5c973525d5892d9 100644 --- a/tensorflow/contrib/tensorrt/test/unary_test.py +++ b/tensorflow/contrib/tensorrt/test/unary_test.py @@ -100,16 +100,13 @@ class UnaryTest(trt_test.TfTrtIntegrationTestBase): return trt_test.TfTrtIntegrationTestParams( gdef=g.as_graph_def(), input_names=[input_name, input2_name], - input_dims=[input_dims, input2_dims], + input_dims=[[input_dims, input2_dims]], output_names=[output_name], - expected_output_dims=[(12, 5, 8, 12)]) + expected_output_dims=[[[12, 5, 8, 12]]]) def ExpectedEnginesToBuild(self, run_params): """Return the expected engines to build.""" - return [ - "my_trt_op_0", "my_trt_op_1", "my_trt_op_2", "my_trt_op_3", - "my_trt_op_4" - ] + return ["TRTEngineOp_0"] if __name__ == "__main__": diff --git a/tensorflow/contrib/tensorrt/test/utils.cc b/tensorflow/contrib/tensorrt/test/utils.cc deleted file mode 100644 index 276308b3a0a6ce864969afb0179c6a3f00d6b70b..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/tensorrt/test/utils.cc +++ /dev/null @@ -1,101 +0,0 @@ -/* Copyright 2018 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/contrib/tensorrt/test/utils.h" - -#include -#include - -#include "re2/re2.h" -#include "tensorflow/core/platform/macros.h" - -namespace tensorflow { -namespace tensorrt { -namespace test { - -// TODO(aaroey): make this class thread-safe. -class TestValueManager { - public: - static TestValueManager* singleton() { - static TestValueManager* manager = new TestValueManager(); - return manager; - } - - void Enable() { - VLOG(1) << "Enabling test value"; - enabled_ = true; - } - - void Add(const string& label, const string& value) { - if (TF_PREDICT_FALSE(enabled_)) { - QCHECK_NE("", value); - VLOG(1) << "Adding test value: " << label << " -> " << value; - values_.insert({label, value}); - } - } - - string Get(const string& label) { - if (TF_PREDICT_FALSE(enabled_)) { - VLOG(1) << "Getting test value by " << label; - auto itr = values_.find(label); - if (itr == values_.end()) return ""; - return itr->second; - } - return ""; - } - - void Clear(const string& pattern) { - if (TF_PREDICT_FALSE(enabled_)) { - VLOG(1) << "Clearing test values"; - if (pattern.empty()) { - values_.clear(); - return; - } - std::vector keys_to_clear; - for (const auto& kv : values_) { - if (RE2::FullMatch(kv.first, pattern)) { - keys_to_clear.push_back(kv.first); - } - } - for (const string& key : keys_to_clear) { - values_.erase(key); - } - } - } - - private: - TestValueManager() : enabled_(false) {} - - bool enabled_; - std::unordered_map values_; -}; - -void EnableTestValue() { TestValueManager::singleton()->Enable(); } - -void ClearTestValues(const string& pattern) { - TestValueManager::singleton()->Clear(pattern); -} - -void AddTestValue(const string& label, const string& value) { - TestValueManager::singleton()->Add(label, value); -} - -string GetTestValue(const string& label) { - return TestValueManager::singleton()->Get(label); -} - -} // namespace test -} // namespace tensorrt -} // namespace tensorflow diff --git a/tensorflow/contrib/tensorrt/test/utils.h b/tensorflow/contrib/tensorrt/test/utils.h deleted file mode 100644 index 4bb4120206cfaae70107e55d1818e3af2f02717a..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/tensorrt/test/utils.h +++ /dev/null @@ -1,44 +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_TEST_UTILS_H_ -#define TENSORFLOW_CONTRIB_TENSORRT_TEST_UTILS_H_ - -#include "tensorflow/core/lib/core/errors.h" -#include "tensorflow/core/lib/core/status.h" - -namespace tensorflow { -namespace tensorrt { -namespace test { - -// Helper methods to inject values used by testing tools. -void EnableTestValue(); -void ClearTestValues(const string& pattern); -void AddTestValue(const string& label, const string& value); -string GetTestValue(const string& label); - -#define TRT_RETURN_IF_TEST_VALUE(label, value_to_return) \ - do { \ - if (::tensorflow::tensorrt::test::GetTestValue(label) == \ - value_to_return) { \ - return errors::Internal("Injected manually"); \ - } \ - } while (0) - -} // namespace test -} // namespace tensorrt -} // namespace tensorflow - -#endif // TENSORFLOW_CONTRIB_TENSORRT_TEST_UTILS_H_ diff --git a/tensorflow/contrib/tensorrt/test/vgg_block_nchw_test.py b/tensorflow/contrib/tensorrt/test/vgg_block_nchw_test.py index b0271a04b364864b841c2ec9fe53aac74611b2c3..b5fed73d2d75271e2c5c533670923d42f233e80b 100644 --- a/tensorflow/contrib/tensorrt/test/vgg_block_nchw_test.py +++ b/tensorflow/contrib/tensorrt/test/vgg_block_nchw_test.py @@ -70,13 +70,13 @@ class VGGBlockNCHWTest(trt_test.TfTrtIntegrationTestBase): return trt_test.TfTrtIntegrationTestParams( gdef=g.as_graph_def(), input_names=[input_name], - input_dims=[input_dims], + input_dims=[[input_dims]], output_names=[output_name], - expected_output_dims=[(5, 6, 2, 2)]) + expected_output_dims=[[[5, 6, 2, 2]]]) def ExpectedEnginesToBuild(self, run_params): """Return the expected engines to build.""" - return ["my_trt_op_0"] + return ["TRTEngineOp_0"] if __name__ == "__main__": diff --git a/tensorflow/contrib/tensorrt/test/vgg_block_test.py b/tensorflow/contrib/tensorrt/test/vgg_block_test.py index d7c165784bfe14bb5faffd266770328237a3eb80..307128f1a89c46d63e851b6a7cd2d6abe7e39ff8 100644 --- a/tensorflow/contrib/tensorrt/test/vgg_block_test.py +++ b/tensorflow/contrib/tensorrt/test/vgg_block_test.py @@ -61,13 +61,13 @@ class VGGBlockTest(trt_test.TfTrtIntegrationTestBase): return trt_test.TfTrtIntegrationTestParams( gdef=g.as_graph_def(), input_names=[input_name], - input_dims=[input_dims], + input_dims=[[input_dims]], output_names=[output_name], - expected_output_dims=[(5, 2, 2, 6)]) + expected_output_dims=[[[5, 2, 2, 6]]]) def ExpectedEnginesToBuild(self, run_params): """Return the expected engines to build.""" - return ["my_trt_op_0"] + return ["TRTEngineOp_0"] if __name__ == "__main__": diff --git a/tensorflow/contrib/tensorrt/trt_conversion.i b/tensorflow/contrib/tensorrt/trt_conversion.i index 6ea15fb8eff13663625420288a37ba002d57fa47..c12895c730047898f366bf651c798c1f1c5b93f7 100644 --- a/tensorflow/contrib/tensorrt/trt_conversion.i +++ b/tensorflow/contrib/tensorrt/trt_conversion.i @@ -99,9 +99,9 @@ _LIST_OUTPUT_TYPEMAP(int, PyLong_FromLong); #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/util/stat_summarizer.h" -#include "tensorflow/contrib/tensorrt/convert/convert_graph.h" -#include "tensorflow/contrib/tensorrt/convert/utils.h" -#include "tensorflow/contrib/tensorrt/test/utils.h" +#include "tensorflow/compiler/tf2tensorrt/convert/convert_graph.h" +#include "tensorflow/compiler/tf2tensorrt/convert/utils.h" +#include "tensorflow/compiler/tf2tensorrt/utils/test_utils.h" %} %ignoreall diff --git a/tensorflow/contrib/text/python/ops/skip_gram_ops_test.py b/tensorflow/contrib/text/python/ops/skip_gram_ops_test.py index 832d34d60d0553ae54a54d2f41eb2e27370535f6..49260f272eeb27bcaeb7cf314969d067811f2582 100644 --- a/tensorflow/contrib/text/python/ops/skip_gram_ops_test.py +++ b/tensorflow/contrib/text/python/ops/skip_gram_ops_test.py @@ -339,7 +339,7 @@ class SkipGramOpsTest(test.TestCase): lookup.KeyValueTensorInitializer(keys, values), -1) with self.cached_session(): - vocab_freq_table.init.run() + vocab_freq_table.initializer.run() # No vocab_freq_table specified - output should be the same as input. no_table_output = skip_gram_ops._filter_input( @@ -396,7 +396,7 @@ class SkipGramOpsTest(test.TestCase): lookup.KeyValueTensorInitializer(keys, values), -1) with self.cached_session(): - vocab_freq_table.init.run() + vocab_freq_table.initializer.run() output = skip_gram_ops._filter_input( input_tensor=input_tensor, vocab_freq_table=vocab_freq_table, diff --git a/tensorflow/contrib/tfprof/README.md b/tensorflow/contrib/tfprof/README.md index 7faf2b9b24acfd71f0ffa6d4a8477a34ff3ed321..f40e76f554e8815aac96344d8cb0b911bafdd712 100644 --- a/tensorflow/contrib/tfprof/README.md +++ b/tensorflow/contrib/tfprof/README.md @@ -1,24 +1,23 @@ # tfprof: TensorFlow Profiler and Beyond -

Please use `tf.profiler.xxx` instead of `tf.contrib.tfprof.xxx`

-

Full Document in tensorflow/core/profiler/README.md

+

Full Document in +tensorflow/core/profiler/README.md

### Features -* Profile model architectures - * parameters, tensor shapes, float operations, device placement, etc. -* Profile model performance - * execution time, memory consumption - * Profile multiple steps. -* Auto profile and advise. - * accelerator utilization check - * expensive operation check - * operation configuration check - * distributed runtime check (Not OSS) +* Profile model architectures + * parameters, tensor shapes, float operations, device placement, etc. +* Profile model performance + * execution time, memory consumption + * Profile multiple steps. +* Auto profile and advise. + * accelerator utilization check + * expensive operation check + * operation configuration check + * distributed runtime check (Not OSS) ### Interfaces -* Python API -* Command Line -* Visualization -* C++ API (Not public, contact us if needed.) +* Python API +* Command Line +* Visualization 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/lstm.py b/tensorflow/contrib/timeseries/examples/lstm.py index b1c7475442c58b9a190c818b752760a4fb4fe6f0..a7c7230dd2fed6cf14bb272df4e9bfb8f3699493 100644 --- a/tensorflow/contrib/timeseries/examples/lstm.py +++ b/tensorflow/contrib/timeseries/examples/lstm.py @@ -254,8 +254,8 @@ def train_and_predict( if export_directory is None: export_directory = tempfile.mkdtemp() input_receiver_fn = estimator.build_raw_serving_input_receiver_fn() - export_location = estimator.export_savedmodel( - export_directory, input_receiver_fn) + export_location = estimator.export_saved_model(export_directory, + input_receiver_fn) # Warm up and predict using the SavedModel with tf.Graph().as_default(): with tf.Session() as session: diff --git a/tensorflow/contrib/timeseries/examples/multivariate.py b/tensorflow/contrib/timeseries/examples/multivariate.py index e81cb18ad7b928a6fd2a748ea6b258c49cf722ae..6b60542e2200615ca004722627fa743ca9729b3b 100644 --- a/tensorflow/contrib/timeseries/examples/multivariate.py +++ b/tensorflow/contrib/timeseries/examples/multivariate.py @@ -66,8 +66,8 @@ def multivariate_train_and_sample( if export_directory is None: export_directory = tempfile.mkdtemp() input_receiver_fn = estimator.build_raw_serving_input_receiver_fn() - export_location = estimator.export_savedmodel( - export_directory, input_receiver_fn) + export_location = estimator.export_saved_model(export_directory, + input_receiver_fn) with tf.Graph().as_default(): numpy.random.seed(1) # Make the example a bit more deterministic with tf.Session() as session: 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 c230919168b937b26c68e141e15f0762ad70f3e6..2a22295197dc225cefbedf2736adeea5491a9fc2 100644 --- a/tensorflow/contrib/timeseries/python/timeseries/BUILD +++ b/tensorflow/contrib/timeseries/python/timeseries/BUILD @@ -104,8 +104,10 @@ py_test( srcs = [ "estimators_test.py", ], + shard_count = 3, srcs_version = "PY2AND3", tags = [ + "no_mac", "no_pip_gpu", # b/63391119 "nomsan", # Takes too long to run. "notsan", # b/67865658 @@ -279,6 +281,7 @@ py_library( "input_pipeline.py", ], srcs_version = "PY2AND3", + visibility = ["//visibility:public"], deps = [ ":feature_keys", ":model_utils", @@ -359,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 9bbe87e301678c7acb57846555e3f97273c8d806..a8d5e1a49dd4313f58f2f515bc3f292ecce5cbd4 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 @@ -102,12 +101,12 @@ class FlatPredictionModel(training.Model): [batch size, output window size, num_features], where num_features is the same as the constructor argument. """ - if input_window_features.shape[1].value == 0: + if input_window_features.shape.dims[1].value == 0: # TODO(allenl): Make reshape()'s static shape information work on # zero-size Tensors? Currently this special case is required because # otherwise the Dense layers get unknown last dimensions. activation = self._output_flatten(output_window_features) - elif output_window_features.shape[2].value == 0: + elif output_window_features.shape.dims[2].value == 0: activation = self._input_flatten(input_window_features) else: activation = array_ops.concat( @@ -438,7 +437,7 @@ class ARModel(model.TimeSeriesModel): output_window_features = array_ops.zeros( [batch_size, self.output_window_size, 0], dtype=self.dtype) - static_batch_size = times.get_shape()[0].value + static_batch_size = times.get_shape().dims[0].value input_window_features.set_shape( [static_batch_size, self.input_window_size, input_feature_size]) output_window_features.set_shape( @@ -462,8 +461,8 @@ 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)) @@ -772,7 +771,7 @@ class ARModel(model.TimeSeriesModel): # windows matching self.window_size (as with training), but this looping # allows easy plotting of "in-sample" predictions. times.get_shape().assert_has_rank(2) - static_window_size = times.get_shape()[1].value + static_window_size = times.get_shape().dims[1].value if (static_window_size is not None and static_window_size < self.window_size): raise ValueError( @@ -965,16 +964,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 +977,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/estimators.py b/tensorflow/contrib/timeseries/python/timeseries/estimators.py index af68aa03cf6583dc474eda6cda2e648fa1c3d08d..146ed9f27134e3e2a6c74627b6b78e53d65155f0 100644 --- a/tensorflow/contrib/timeseries/python/timeseries/estimators.py +++ b/tensorflow/contrib/timeseries/python/timeseries/estimators.py @@ -32,7 +32,7 @@ from tensorflow.contrib.timeseries.python.timeseries.state_space_models.filterin from tensorflow.python.estimator import estimator_lib from tensorflow.python.estimator.canned import optimizers from tensorflow.python.estimator.export import export_lib -from tensorflow.python.feature_column import feature_column +from tensorflow.python.feature_column import feature_column_lib as feature_column from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.framework import tensor_shape diff --git a/tensorflow/contrib/timeseries/python/timeseries/estimators_test.py b/tensorflow/contrib/timeseries/python/timeseries/estimators_test.py index 6ec7184c6839ebc4375432fa54c9b880729b9550..7d780559f976516823611f3fe0ded056e4be088c 100644 --- a/tensorflow/contrib/timeseries/python/timeseries/estimators_test.py +++ b/tensorflow/contrib/timeseries/python/timeseries/estimators_test.py @@ -30,7 +30,7 @@ from tensorflow.contrib.timeseries.python.timeseries import saved_model_utils from tensorflow.python.client import session from tensorflow.python.estimator import estimator_lib -from tensorflow.python.feature_column import feature_column +from tensorflow.python.feature_column import feature_column_lib as feature_column from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.platform import test @@ -98,8 +98,8 @@ class TimeSeriesRegressorTest(test.TestCase): ) = list(second_estimator.predict(input_fn=predict_input_fn)) self.assertAllEqual([10, 1], estimator_predictions["mean"].shape) input_receiver_fn = first_estimator.build_raw_serving_input_receiver_fn() - export_location = first_estimator.export_savedmodel(self.get_temp_dir(), - input_receiver_fn) + export_location = first_estimator.export_saved_model( + self.get_temp_dir(), input_receiver_fn) with ops.Graph().as_default(): with session.Session() as sess: signatures = loader.load(sess, [tag_constants.SERVING], export_location) diff --git a/tensorflow/contrib/timeseries/python/timeseries/head_test.py b/tensorflow/contrib/timeseries/python/timeseries/head_test.py index 04d17bc123eae5128315657948dc29e59cd2c941..8f692d94da45bfaed6c72cf75d525346865aea34 100644 --- a/tensorflow/contrib/timeseries/python/timeseries/head_test.py +++ b/tensorflow/contrib/timeseries/python/timeseries/head_test.py @@ -38,7 +38,7 @@ from tensorflow.core.example import example_pb2 from tensorflow.python.client import session as session_lib from tensorflow.python.estimator import estimator_lib -from tensorflow.python.feature_column import feature_column +from tensorflow.python.feature_column import feature_column_lib as feature_column from tensorflow.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops @@ -402,8 +402,8 @@ class OneShotTests(parameterized.TestCase): self.assertIn("average_loss", result) self.assertNotIn(feature_keys.State.STATE_TUPLE, result) input_receiver_fn = estimator.build_raw_serving_input_receiver_fn() - export_location = estimator.export_savedmodel(_new_temp_dir(), - input_receiver_fn) + export_location = estimator.export_saved_model(_new_temp_dir(), + input_receiver_fn) graph = ops.Graph() with graph.as_default(): with session_lib.Session() as session: @@ -438,7 +438,7 @@ class OneShotTests(parameterized.TestCase): output = session.run(fetches, feed_dict=feeds) self.assertEqual((2, 15, 5), output["mean"].shape) # Build a parsing input function, then make a tf.Example for it to parse. - export_location = estimator.export_savedmodel( + export_location = estimator.export_saved_model( _new_temp_dir(), estimator.build_one_shot_parsing_serving_input_receiver_fn( filtering_length=20, prediction_length=15)) diff --git a/tensorflow/contrib/timeseries/python/timeseries/math_utils.py b/tensorflow/contrib/timeseries/python/timeseries/math_utils.py index 9c585fe6a7537d105dba57818b5b33f559bfa6bc..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. @@ -262,8 +290,8 @@ def batch_times_matrix(batch, matrix, adj_x=False, adj_y=False): assert matrix.get_shape().ndims == 2 if adj_x: batch = array_ops.transpose(batch, [0, 2, 1]) - batch_dimension = batch.get_shape()[0].value - first_dimension = batch.get_shape()[1].value + batch_dimension = batch.get_shape().dims[0].value + first_dimension = batch.get_shape().dims[1].value tensor_batch_shape = array_ops.shape(batch) if batch_dimension is None: batch_dimension = tensor_batch_shape[0] @@ -802,7 +830,7 @@ class InputStatisticsFromMiniBatch(object): array_ops.shape(times)[1] - 1, self._dtype)) # Co-locate updates with their variables to minimize race conditions when # updating statistics. - with ops.colocate_with(auxiliary_variables.max_time_seen): + with ops.device(auxiliary_variables.max_time_seen.device): # There is a race condition if this value is being updated from multiple # workers. However, it should eventually reach the correct value if the # last chunk is presented enough times. @@ -810,16 +838,16 @@ class InputStatisticsFromMiniBatch(object): auxiliary_variables.max_time_seen, gen_math_ops.maximum(auxiliary_variables.max_time_seen, math_ops.reduce_max(times))) - with ops.colocate_with(auxiliary_variables.chunk_count): + with ops.device(auxiliary_variables.chunk_count.device): chunk_count_assign = state_ops.assign_add(auxiliary_variables.chunk_count, array_ops.shape( times, out_type=dtypes.int64)[0]) - with ops.colocate_with(auxiliary_variables.inter_observation_duration_sum): + with ops.device(auxiliary_variables.inter_observation_duration_sum.device): inter_observation_duration_assign = state_ops.assign_add( auxiliary_variables.inter_observation_duration_sum, math_ops.reduce_sum(batch_inter_observation_duration)) - with ops.colocate_with(auxiliary_variables.example_count): + with ops.device(auxiliary_variables.example_count.device): example_count_assign = state_ops.assign_add( auxiliary_variables.example_count, array_ops.size(times, out_type=dtypes.int64)) @@ -829,11 +857,11 @@ class InputStatisticsFromMiniBatch(object): # the series are then members of fewer chunks. For series which are much # longer than the chunk size (the usual/expected case), this effect becomes # irrelevant. - with ops.colocate_with(auxiliary_variables.overall_feature_sum): + with ops.device(auxiliary_variables.overall_feature_sum.device): overall_feature_sum_assign = state_ops.assign_add( auxiliary_variables.overall_feature_sum, math_ops.reduce_sum(values, axis=[0, 1])) - with ops.colocate_with(auxiliary_variables.overall_feature_sum_of_squares): + with ops.device(auxiliary_variables.overall_feature_sum_of_squares.device): overall_feature_sum_of_squares_assign = state_ops.assign_add( auxiliary_variables.overall_feature_sum_of_squares, math_ops.reduce_sum(values**2, axis=[0, 1])) @@ -869,7 +897,7 @@ class InputStatisticsFromMiniBatch(object): state_ops.assign(statistics.series_start_moments.mean, mean), state_ops.assign(statistics.series_start_moments.variance, variance)) - with ops.colocate_with(statistics.start_time): + with ops.device(statistics.start_time.device): series_start_update = control_flow_ops.cond( # Update moments whenever we even match the lowest time seen so far, # to ensure that series start statistics are eventually updated to diff --git a/tensorflow/contrib/timeseries/python/timeseries/model.py b/tensorflow/contrib/timeseries/python/timeseries/model.py index 7644764a7459db3951fe9a2790389713dd412a8f..a8cd4287e0003de300b7114cf3f88d21d3239e6e 100644 --- a/tensorflow/contrib/timeseries/python/timeseries/model.py +++ b/tensorflow/contrib/timeseries/python/timeseries/model.py @@ -21,13 +21,16 @@ from __future__ import print_function import abc import collections +import six + from tensorflow.contrib.timeseries.python.timeseries import math_utils from tensorflow.contrib.timeseries.python.timeseries.feature_keys import PredictionFeatures from tensorflow.contrib.timeseries.python.timeseries.feature_keys import TrainEvalFeatures -from tensorflow.python.feature_column import feature_column +from tensorflow.python.feature_column import feature_column_lib as feature_column 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 math_ops @@ -52,11 +55,10 @@ ModelOutputs = collections.namedtuple( # pylint: disable=invalid-name ]) +@six.add_metaclass(abc.ABCMeta) class TimeSeriesModel(object): """Base class for creating generative time series models.""" - __metaclass__ = abc.ABCMeta - def __init__(self, num_features, exogenous_feature_columns=None, @@ -712,7 +714,7 @@ class SequentialTimeSeriesModel(TimeSeriesModel): `outputs` and computed in state_update_fn. """ times = ops.convert_to_tensor(times, dtype=dtypes.int64) - window_static_shape = times.get_shape()[1].value + window_static_shape = tensor_shape.dimension_value(times.shape[1]) if self._static_unrolling_window_size_threshold is None: static_unroll = False else: @@ -789,7 +791,7 @@ class SequentialTimeSeriesModel(TimeSeriesModel): [_window_size_tensor_array(self.dtype) for _ in outputs]] if static_unroll: arguments = initial_loop_arguments - for step_number in range(times.get_shape()[1].value): + for step_number in range(tensor_shape.dimension_value(times.shape[1])): arguments = _state_update_step( array_ops.constant(step_number, dtypes.int32), *arguments[1:], reuse=(step_number > 0)) # Variable sharing between steps diff --git a/tensorflow/contrib/timeseries/python/timeseries/state_space_models/BUILD b/tensorflow/contrib/timeseries/python/timeseries/state_space_models/BUILD index 3c07a74ed8af9e3ab70408f9b43cb62b6bd4c7f2..cf5e749042afd83f927a3d22edfd3a9538ab2ffd 100644 --- a/tensorflow/contrib/timeseries/python/timeseries/state_space_models/BUILD +++ b/tensorflow/contrib/timeseries/python/timeseries/state_space_models/BUILD @@ -40,7 +40,10 @@ py_test( timeout = "long", # Moderate but for asan srcs = ["state_space_model_test.py"], srcs_version = "PY2AND3", - tags = ["no_windows"], # TODO: needs investigation on Windows + tags = [ + "no_mac", + "no_windows", # TODO: needs investigation on Windows + ], deps = [ ":state_space_model", "//tensorflow/contrib/layers:layers_py", @@ -75,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", @@ -232,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 7fa538a16ecd7dcf39beeb001992fd7927cee70b..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 @@ -20,7 +20,7 @@ from __future__ import print_function import abc -from tensorflow.contrib import distributions +import six from tensorflow.contrib.timeseries.python.timeseries import math_utils @@ -32,11 +32,10 @@ from tensorflow.python.ops import math_ops from tensorflow.python.util import nest +@six.add_metaclass(abc.ABCMeta) class FilteringStepPostprocessor(object): """Base class for processors that are applied after each filter step.""" - __metaclass__ = abc.ABCMeta - @abc.abstractmethod def process_filtering_step(self, current_times, current_values, predicted_state, filtered_state, outputs): @@ -90,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/timeseries/python/timeseries/state_space_models/state_space_model.py b/tensorflow/contrib/timeseries/python/timeseries/state_space_models/state_space_model.py index d04c72100749fd0c96ac74cccbaeb93b7fcd5db4..2ecc7eafdaf1e3dc3a76f99f995e39261e333da7 100644 --- a/tensorflow/contrib/timeseries/python/timeseries/state_space_models/state_space_model.py +++ b/tensorflow/contrib/timeseries/python/timeseries/state_space_models/state_space_model.py @@ -35,6 +35,7 @@ from tensorflow.python.estimator import estimator_lib from tensorflow.python.framework import constant_op 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 gen_math_ops @@ -510,7 +511,7 @@ class StateSpaceModel(model.SequentialTimeSeriesModel): estimated_state, estimated_state_covariance, previous_times = state state_transition = ops.convert_to_tensor( self.get_state_transition(), dtype=self.dtype) - state_dimension = state_transition.get_shape()[0].value + state_dimension = tensor_shape.dimension_value(state_transition.shape[0]) # Learning the observation model would be redundant since we transform # `exogenous_values` to the state space via a linear transformation anyway. observation_model = linalg_ops.eye( @@ -572,8 +573,9 @@ class StateSpaceModel(model.SequentialTimeSeriesModel): start_mean, start_covariance, previous_times = state with variable_scope.variable_scope("exogenous_noise_increasing_mean"): mean_addition = layers.fully_connected( - exogenous_values, start_mean.get_shape()[1].value, activation_fn=None) - state_dimension = start_covariance.get_shape()[1].value + exogenous_values, + tensor_shape.dimension_value(start_mean.shape[1]), activation_fn=None) + state_dimension = tensor_shape.dimension_value(start_covariance.shape[1]) with variable_scope.variable_scope("exogenous_noise_increasing_covariance"): covariance_addition = ( math_utils.transform_to_covariance_matrices( @@ -712,7 +714,7 @@ class StateSpaceModel(model.SequentialTimeSeriesModel): """ with variable_scope.variable_scope(self._variable_scope): state_dimension = ops.convert_to_tensor( - self.get_state_transition()).get_shape()[0].value + self.get_state_transition()).get_shape().dims[0].value if self._configuration.trainable_start_state: base_covariance = math_utils.variable_covariance_matrix( state_dimension, "prior_state_var", @@ -742,7 +744,7 @@ class StateSpaceModel(model.SequentialTimeSeriesModel): with variable_scope.variable_scope(self._variable_scope): state_transition = ops.convert_to_tensor( self.get_state_transition(), dtype=self.dtype) - state_dimension = state_transition.get_shape()[0].value + state_dimension = state_transition.get_shape().dims[0].value return variable_scope.get_variable( name="prior_state_mean", shape=[state_dimension], @@ -920,7 +922,7 @@ class StateSpaceModel(model.SequentialTimeSeriesModel): self, minimum_initial_variance=1e-5): state_noise_transform = ops.convert_to_tensor( self.get_noise_transform(), dtype=self.dtype) - state_noise_dimension = state_noise_transform.get_shape()[1].value + state_noise_dimension = state_noise_transform.get_shape().dims[1].value if self._input_statistics is not None: feature_variance = self._scale_variance( self._input_statistics.series_start_moments.variance) diff --git a/tensorflow/contrib/timeseries/python/timeseries/state_space_models/state_space_model_test.py b/tensorflow/contrib/timeseries/python/timeseries/state_space_models/state_space_model_test.py index 80126ac786e7fc41d334076fc19bea7d091e19ab..26857ba235e4fe13904ca4f1334f4662a795f8a8 100644 --- a/tensorflow/contrib/timeseries/python/timeseries/state_space_models/state_space_model_test.py +++ b/tensorflow/contrib/timeseries/python/timeseries/state_space_models/state_space_model_test.py @@ -185,9 +185,8 @@ class StateSpaceEquivalenceTests(test.TestCase): "exogenous": [-1., -2., -3., -4.] })) estimator.train(combined_input_fn, steps=1) - export_location = estimator.export_savedmodel( - self.get_temp_dir(), - estimator.build_raw_serving_input_receiver_fn()) + export_location = estimator.export_saved_model( + self.get_temp_dir(), estimator.build_raw_serving_input_receiver_fn()) with ops.Graph().as_default() as graph: random_model.initialize_graph() with self.session(graph=graph) as session: diff --git a/tensorflow/contrib/timeseries/python/timeseries/state_space_models/test_utils.py b/tensorflow/contrib/timeseries/python/timeseries/state_space_models/test_utils.py index 5f127700d99f1a9cf2549e2fdb57ce6090440ac7..f7f5024b34218ceb04d13bb351f6d2d302069bce 100644 --- a/tensorflow/contrib/timeseries/python/timeseries/state_space_models/test_utils.py +++ b/tensorflow/contrib/timeseries/python/timeseries/state_space_models/test_utils.py @@ -24,6 +24,7 @@ from tensorflow.contrib.timeseries.python.timeseries import math_utils 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 linalg_ops from tensorflow.python.ops import math_ops @@ -35,7 +36,7 @@ def transition_power_test_template(test_case, model, num_steps): transition_matrix = ops.convert_to_tensor( model.get_state_transition(), dtype=model.dtype) step_number = array_ops.placeholder(shape=[], dtype=dtypes.int64) - state_dimension = transition_matrix.get_shape()[0].value + state_dimension = tensor_shape.dimension_value(transition_matrix.shape[0]) previous_matrix = array_ops.placeholder( shape=[state_dimension, state_dimension], dtype=transition_matrix.dtype) true_single_step_update = math_ops.matmul(previous_matrix, @@ -63,8 +64,8 @@ def noise_accumulator_test_template(test_case, model, num_steps): model.get_state_transition(), dtype=model.dtype) noise_transform = ops.convert_to_tensor( model.get_noise_transform(), dtype=model.dtype) - state_dimension = transition_matrix.get_shape()[0].value - state_noise_dimension = noise_transform.get_shape()[1].value + state_dimension = tensor_shape.dimension_value(transition_matrix.shape[0]) + state_noise_dimension = tensor_shape.dimension_value(noise_transform.shape[1]) gen_noise_addition = math_utils.sign_magnitude_positive_definite( raw=random_ops.random_normal( shape=[state_noise_dimension, state_noise_dimension], diff --git a/tensorflow/contrib/timeseries/python/timeseries/state_space_models/varma.py b/tensorflow/contrib/timeseries/python/timeseries/state_space_models/varma.py index 6746dd7b433466c473402e0e8374377093a73492..ddee749b498121ee62c13bde59680269bc497d23 100644 --- a/tensorflow/contrib/timeseries/python/timeseries/state_space_models/varma.py +++ b/tensorflow/contrib/timeseries/python/timeseries/state_space_models/varma.py @@ -52,6 +52,7 @@ from tensorflow.contrib.timeseries.python.timeseries import math_utils from tensorflow.contrib.timeseries.python.timeseries.state_space_models import state_space_model 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 init_ops from tensorflow.python.ops import linalg_ops @@ -191,7 +192,8 @@ class VARMA(state_space_model.StateSpaceModel): initial_transition_noise_scale = 0. state_noise_transform = ops.convert_to_tensor( self.get_noise_transform(), dtype=self.dtype) - state_noise_dimension = state_noise_transform.get_shape()[1].value + state_noise_dimension = tensor_shape.dimension_value( + state_noise_transform.shape[1]) return math_utils.variable_covariance_matrix( state_noise_dimension, "state_transition_noise", dtype=self.dtype, diff --git a/tensorflow/contrib/tpu/BUILD b/tensorflow/contrib/tpu/BUILD index 4855a379b4497f93243809669de05ebfab1886b1..faf11afb21e7159dada285314bcafeeb66c69627 100644 --- a/tensorflow/contrib/tpu/BUILD +++ b/tensorflow/contrib/tpu/BUILD @@ -1,23 +1,25 @@ # Description: Operations defined for Cloud TPUs -licenses(["notice"]) # Apache 2.0 - load( "//tensorflow:tensorflow.bzl", "tf_custom_op_library", "tf_gen_op_libs", "tf_gen_op_wrapper_py", + "tf_py_test", ) load("//tensorflow:tensorflow.bzl", "tf_custom_op_py_library") -load("//tensorflow:tensorflow.bzl", "tf_py_test") + +licenses(["notice"]) # Apache 2.0 package( default_visibility = [ "//cloud/vmm/testing/tests/tpu:__subpackages__", + "//knowledge/cerebra/sense/im2query:__subpackages__", "//learning/brain:__subpackages__", "//learning/deepmind:__subpackages__", "//medical/pathology:__subpackages__", "//tensorflow:__subpackages__", + "//vr/perception:__subpackages__", ], ) @@ -68,15 +70,19 @@ py_library( srcs_version = "PY2AND3", deps = [ ":async_checkpoint", + ":functional", ":tpu_lib", + ":tpu_ordinal_selector_py", "//tensorflow/contrib/training:training_py", "//tensorflow/core:protos_all_py", "//tensorflow/python:array_ops", "//tensorflow/python:control_flow_ops", "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:function", "//tensorflow/python:init_ops", "//tensorflow/python:math_ops", "//tensorflow/python:platform", + "//tensorflow/python:session", "//tensorflow/python:state_ops", "//tensorflow/python:summary", "//tensorflow/python:summary_ops_v2", @@ -99,6 +105,8 @@ tf_gen_op_libs( "replication_ops", "tpu_configuration_ops", "tpu_embedding_ops", + "tpu_ordinal_selector_op", + "functional_ops", ], deps = [ "//tensorflow/contrib/tpu/proto:tpu_embedding_configuration_proto_cc", @@ -150,6 +158,52 @@ tf_gen_op_wrapper_py( ], ) +tf_custom_op_library( + name = "python/ops/_tpu_ordinal_selector.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"], + kernels = [ + ":tpu_ordinal_selector_op_op_lib", + ], + srcs_version = "PY2AND3", + visibility = ["//visibility:public"], + deps = [ + ":tpu_ordinal_selector_op", + ], +) + +tf_gen_op_wrapper_py( + name = "tpu_ordinal_selector_op", + deps = [ + ":tpu_ordinal_selector_op_op_lib", + ], +) + +tf_gen_op_wrapper_py( + name = "gen_functional_ops", + out = "python/tpu/gen_functional_ops.py", + hidden = [ + "TPUPartitionedCall", + ], + deps = [":functional_ops_op_lib"], +) + +py_library( + name = "functional", + srcs = ["python/tpu/functional.py"], + visibility = [ + "//visibility:public", + ], + deps = [ + ":gen_functional_ops", + ], +) + py_library( name = "profiler", srcs = ["python/profiler/__init__.py"], @@ -190,6 +244,7 @@ py_library( ], srcs_version = "PY2AND3", deps = [ + ":feature_column", ":keras_support", # split out to avoid cycle with tpu_strategy ":tpu_embedding", ":tpu_estimator", @@ -208,11 +263,12 @@ py_library( "//cloud/vmm/testing/tests/tpu:__subpackages__", "//learning/brain:__subpackages__", "//tensorflow:__subpackages__", - "//third_party/cloud_tpu/models/keras:__subpackages__", + "//third_party/cloud_tpu/models/keras_colab:__subpackages__", + "//third_party/cloud_tpu/models/resnet50_keras:__subpackages__", ], deps = [ ":tpu_lib", - "//tensorflow/contrib/cluster_resolver:tpu_cluster_resolver_py", + "//tensorflow/contrib/cluster_resolver:cluster_resolver_py", "//tensorflow/contrib/distribute", "//tensorflow/contrib/framework:framework_py", "//tensorflow/contrib/tpu/proto:compilation_result_proto_py", @@ -243,6 +299,7 @@ py_library( "python/tpu/bfloat16.py", "python/tpu/device_assignment.py", "python/tpu/session_support.py", + "python/tpu/tensor_tracer.py", "python/tpu/topology.py", "python/tpu/tpu.py", "python/tpu/tpu_feed.py", @@ -259,9 +316,10 @@ py_library( ":tpu_py", "//tensorflow/compiler/xla/experimental/xla_sharding", "//tensorflow/compiler/xla/python_api:xla_shape", - "//tensorflow/contrib/cluster_resolver:tpu_cluster_resolver_py", + "//tensorflow/contrib/cluster_resolver:cluster_resolver_py", "//tensorflow/contrib/compiler:xla", "//tensorflow/contrib/tpu/proto:compilation_result_proto_py", + "//tensorflow/contrib/tpu/proto:dynamic_padding_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", @@ -301,13 +359,15 @@ py_library( tf_py_test( name = "datasets_test", + size = "medium", srcs = ["python/tpu/datasets_test.py"], additional_deps = [ "//tensorflow/python:client_testlib", ":datasets", ], - flaky = 1, # TODO(b/117363808): fails 1/1000 OSS runs grpc_enabled = True, + shard_count = 4, + tags = ["no_oss"], ) tf_py_test( @@ -394,7 +454,8 @@ py_library( srcs = ["python/tpu/tpu_embedding.py"], srcs_version = "PY2AND3", deps = [ - "//tensorflow/contrib/tpu:tpu_ops", + ":tpu_lib", + ":tpu_ops", "//tensorflow/contrib/tpu/proto:tpu_embedding_configuration_proto_py", "//tensorflow/python:array_ops", "//tensorflow/python:framework_for_generated_wrappers", @@ -406,3 +467,37 @@ py_library( "@six_archive//:six", ], ) + +py_library( + name = "feature_column", + srcs = ["python/tpu/feature_column.py"], + deps = [ + ":tpu_lib", + "//tensorflow/python:framework_ops", + "//tensorflow/python:init_ops", + "//tensorflow/python:variable_scope", + "//tensorflow/python/feature_column", + "//tensorflow/python/feature_column:feature_column_py", + ], +) + +tf_py_test( + name = "feature_column_test", + srcs = [ + "python/tpu/feature_column_test.py", + ], + additional_deps = [ + ":feature_column", + "//third_party/py/numpy", + "//tensorflow/python:dtypes", + "//tensorflow/python:framework_ops", + "//tensorflow/python:lookup_ops", + "//tensorflow/python:parsing_ops", + "//tensorflow/python:session", + "//tensorflow/python:sparse_tensor", + "//tensorflow/python:variables", + "//tensorflow/python/feature_column", + "//tensorflow/python/feature_column:feature_column_py", + ], + main = "python/tpu/feature_column_test.py", +) diff --git a/tensorflow/contrib/tpu/ops/functional_ops.cc b/tensorflow/contrib/tpu/ops/functional_ops.cc new file mode 100644 index 0000000000000000000000000000000000000000..aa81e8b24b5e303f5de5d2938b9474fc6b7af6c9 --- /dev/null +++ b/tensorflow/contrib/tpu/ops/functional_ops.cc @@ -0,0 +1,31 @@ +/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "tensorflow/core/framework/common_shape_fns.h" +#include "tensorflow/core/framework/op.h" +#include "tensorflow/core/framework/shape_inference.h" + +namespace tensorflow { + +REGISTER_OP("TPUPartitionedCall") + .Input("args: Tin") + .Input("device_ordinal: int32") + .Output("output: Tout") + .Attr("Tin: list(type) >= 0") + .Attr("Tout: list(type) >= 0") + .Attr("f: func") + .SetShapeFn(shape_inference::UnknownShape); + +} // namespace tensorflow diff --git a/tensorflow/contrib/tpu/ops/infeed_ops.cc b/tensorflow/contrib/tpu/ops/infeed_ops.cc index efc546f9a6077de9cac5a5acefa3fc7206547fc6..2ed16c2a2270a5399059d7e07f5903e11098bbf9 100644 --- a/tensorflow/contrib/tpu/ops/infeed_ops.cc +++ b/tensorflow/contrib/tpu/ops/infeed_ops.cc @@ -40,6 +40,7 @@ REGISTER_OP("InfeedEnqueue") .Input("input: dtype") .Attr("dtype: type") .Attr("shape: shape = {}") + .Attr("layout: list(int) = []") .Attr("device_ordinal: int = -1") .SetShapeFn(shape_inference::NoOutputs) .SetIsStateful() @@ -49,6 +50,9 @@ An op which feeds a single Tensor value into the computation. input: A tensor that will be provided using the infeed mechanism. dtype: The type of elements in the tensor. shape: The shape of the tensor. +layout: A vector holding the requested layout in minor-to-major sequence. +If a layout attribute is passed, but its values are all -1, the layout will +be computed by the infeed operation. device_ordinal: The TPU device to use. This should be -1 when the Op is running on a TPU device, and >= 0 when the Op is running on the CPU device. @@ -58,6 +62,7 @@ REGISTER_OP("InfeedEnqueueTuple") .Input("inputs: dtypes") .Attr("dtypes: list(type)") .Attr("shapes: list(shape)") + .Attr("layouts: list(int) = []") .Attr("device_ordinal: int = -1") .SetShapeFn(shape_inference::NoOutputs) .SetIsStateful() @@ -67,6 +72,10 @@ An op which feeds multiple Tensor values into the computation as an XLA tuple. inputs: A list of tensors that will be provided using the infeed mechanism. dtypes: The element types of each element in `inputs`. shapes: The shapes of each tensor in `inputs`. +layouts: A vector holding the requested layout in minor-to-major sequence for +all the tuple shapes, in the order the shapes appear in the "shapes" input. +The layout elements for a sub-shape can be set to -1, in which case the +corresponding layout will be computed by the infeed operation. device_ordinal: The TPU device to use. This should be -1 when the Op is running on a TPU device, and >= 0 when the Op is running on the CPU device. 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 0ef29bdf734467aa9dee5c157bc8d8a7e0a85f13..676aed0b7b651494eda80ff2d7c7c31097529590 100644 --- a/tensorflow/contrib/tpu/ops/tpu_embedding_ops.cc +++ b/tensorflow/contrib/tpu/ops/tpu_embedding_ops.cc @@ -37,18 +37,18 @@ namespace tensorflow { // pieces of the TF Graph. // 1. Pass this TPUEmbeddingConfiguration to tpu.initialize_system() as the // tpu_embedding_config parameter. -// 2. Use the TPUEmbeddingLoad Op to initialize the embedding tables in TPU +// 2. Use the LoadTPUEmbedding Ops to initialize the embedding tables in TPU // memories, sharded across the memories attached to each Host. -// 3. Use TPUEmbeddingEnqueueSparseBatch to provide the TPU with embedding +// 3. Use EnqueueTPUEmbeddingSparseBatch to provide the TPU with embedding // indices and aggregation weights. -// 4. TPUEmbeddingReceiveActivations returns a list of Tensors, containing the +// 4. RecvTPUEmbeddingActivations returns a list of Tensors, containing the // activations from each table specified in the configuration. // 5. TPUEmbeddingActivations, when used with appropriate Python libraries, // enables the automatic differentiation of models that use embeddings. -// 6. TPUEmbeddingSendGradients takes a list of Tensors (of the same shapes +// 6. SendTPUEmbeddingGradients takes a list of Tensors (of the same shapes // as those returned by TPUEmbeddingReceiveActivations) containing gradients // to use in updating the embedding tables. -// 7. Before saving a checkpoint, use the TPUEmbeddingRetrieve Op to update +// 7. Before saving a checkpoint, use the RetrieveTPUEmbedding Ops to update // the Graph's embedding table Variables from the updated tables in the // TPU memories. // @@ -455,20 +455,21 @@ REGISTER_OP("SendTPUEmbeddingGradients") return Status::OK(); }) .Doc(R"doc( -An op that performs gradient updates of embedding tables. - -The TensorList argument has the same length and shapes as the return value of -TPUEmbeddingReceiveActivations, but contains gradients of the model's loss -with respect to the embedding activations. The embedding tables are updated -from these gradients via the optimizer specified in the configuration given -to tpu.initialize_system. +An op that performs gradient updates of embedding tables using the specified +learning rates. inputs: A TensorList of gradients with which to update embedding tables. - It contains one tensor per embedding table in the model. -learning_rates: A list 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. + 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 optimizer specified in the TPU embedding + 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. + 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. config: Serialized TPUEmbeddingConfiguration proto. )doc"); diff --git a/tensorflow/contrib/tpu/ops/tpu_ordinal_selector_op.cc b/tensorflow/contrib/tpu/ops/tpu_ordinal_selector_op.cc new file mode 100644 index 0000000000000000000000000000000000000000..54e6b20f7f388b67a96ac8acfe814a4202b56a18 --- /dev/null +++ b/tensorflow/contrib/tpu/ops/tpu_ordinal_selector_op.cc @@ -0,0 +1,39 @@ +/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "tensorflow/core/framework/op.h" +#include "tensorflow/core/framework/shape_inference.h" + +namespace tensorflow { + +REGISTER_OP("TPUOrdinalSelector") + .Output("device_ordinals: int32") + .SetIsStateful() + .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) { + c->set_output(0, + c->Vector(shape_inference::InferenceContext::kUnknownDim)); + return Status::OK(); + }) + .Doc(R"doc( +A TPU core selector Op. + +This Op produces a set of TPU cores (for warm-up) or a single TPU core +(for regular inference) to execute the TPU program on. The output is +consumed by TPUPartitionedCall. + +device_ordinals: A vector 1 or more TPU cores. +)doc"); + +} // namespace tensorflow diff --git a/tensorflow/contrib/tpu/profiler/BUILD b/tensorflow/contrib/tpu/profiler/BUILD index 38d1c3049ef7185f2f9f448361029d066678cdae..541fbf33a302a4d850422885fdbbc438bd6b9b7b 100644 --- a/tensorflow/contrib/tpu/profiler/BUILD +++ b/tensorflow/contrib/tpu/profiler/BUILD @@ -94,13 +94,6 @@ tf_proto_library( visibility = ["//visibility:public"], ) -tf_proto_library( - name = "tf_op_stats_proto", - srcs = ["tf_op_stats.proto"], - cc_api_version = 2, - visibility = ["//visibility:public"], -) - tf_proto_library( name = "tpu_profiler_analysis_proto", srcs = ["tpu_profiler_analysis.proto"], diff --git a/tensorflow/contrib/tpu/profiler/pip_package/cloud_tpu_profiler/main.py b/tensorflow/contrib/tpu/profiler/pip_package/cloud_tpu_profiler/main.py index 63641e00c5dbf4b4e635ecfea8bef98c7d0b7075..a081c4354a779d37140338793e66844c3fcf7a12 100644 --- a/tensorflow/contrib/tpu/profiler/pip_package/cloud_tpu_profiler/main.py +++ b/tensorflow/contrib/tpu/profiler/pip_package/cloud_tpu_profiler/main.py @@ -90,12 +90,12 @@ def main(unused_argv=None): tf_version = tf.__version__ print('TensorFlow version %s detected' % tf_version) - if FLAGS.service_addr is None and FLAGS.tpu is None: + if not FLAGS.service_addr and not FLAGS.tpu: sys.exit('You must specify either --service_addr or --tpu.') tpu_cluster_resolver = None - if FLAGS.service_addr is not None: - if FLAGS.tpu is not None: + if FLAGS.service_addr: + if FLAGS.tpu: tf.logging.warn('Both --service_addr and --tpu are set. Ignoring ' '--tpu and using --service_addr.') service_addr = FLAGS.service_addr 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/tf_op_stats.proto b/tensorflow/contrib/tpu/profiler/tf_op_stats.proto deleted file mode 100644 index 1e66801efd4b2a997ed85289b9b1690bb5d07737..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/tpu/profiler/tf_op_stats.proto +++ /dev/null @@ -1,261 +0,0 @@ -// This proto describes the format of tensorflow operation level stats for -// profiling (in tensorboard) purpose. - -syntax = "proto2"; - -package tensorflow.tpu; - -// Result proto for OpMetrics. -message OpMetricsResult { - // True if this OP is executed on the device; False if it is executed on the - // host. - optional bool on_device = 1; - reserved 2; // was uint32 id. - // Name of this OP. - optional string name = 3; - // Rank of this OP. - optional uint64 rank = 4; - // The starting time in cycles of the last instance of this OP executed. - optional double last_starttime_in_cycles = 5; - // The ending time in cycles of the last instance of this OP executed. - optional double last_endtime_in_cycles = 6; - // If this OP (say A), is an immediate child of another OP (say B), this field - // stores the sum of duration in microseconds of A inside B. If A appears more - // than once in B, the duration of all A's appearances will be added together. - // This sum will be reset after the self-time of B is calculated so that it - // can be reused for a new parent OP. - optional double sum_of_duration_in_us_as_children = 7; - // Number of instances that this OP occurred. - optional uint64 occurrences = 8; - // Total time in microseconds spent in this OP (accumulated - // over all of its occurrences). - optional double total_time_in_us = 9; - // Total self time in microseconds spent in this OP - // (accumulated over all of its occurrences). - optional double total_self_time_in_us = 10; - // The total self time as a fraction of sum of all OP's - // total self time on the host. - optional double host_total_self_time_as_fraction_of_all_op_time = 11; - // Cumulative total self time in fraction on the host. - optional double host_cumulative_total_self_time_as_fraction_of_all_op_time = - 12; - // The total self time as a fraction of sum of all OP's - // total self time on the device. - optional double device_total_self_time_as_fraction_of_all_op_time = 13; - // Cumulative total self time in fraction on the device. - optional double device_cumulative_total_self_time_as_fraction_of_all_op_time = - 14; - // Total number of FLOPs incurred by this OP. - optional double total_flops = 15; - // Total number of bytes accessed by this OP. - optional double total_bytes_accessed = 16; - // Total time in microseconds that special hw unit 1 is occupied by this OP. - optional double unit1_occupancy_in_us = 17; - // Total time in microseconds that special hw unit 2 is occupied by this OP. - optional double unit2_occupancy_in_us = 18; - // Total memory stall time in microseconds. - optional double total_memory_stall_in_us = 19; -} - -// Result proto for OpMetricsDb. -message OpMetricsDbResult { - // A bunch of OpMetricsResults. - repeated OpMetricsResult metrics_db = 1; - // The total host infeed-enqueue duration in picoseconds. - optional uint64 total_host_infeed_enq_duration_ps = 2; - // The total of the difference between the start times of two - // consecutive infeed-enqueues (per host) in picoseconds. - optional uint64 total_host_infeed_enq_start_timestamp_ps_diff = 3; - // The total device time in microseconds. - optional double total_device_time_in_us = 4; - // The total host time in microseconds. - optional double total_host_time_in_us = 5; -} - -// Result proto for StepInfo. -message StepInfoResult { - // The (micro) step number. - optional uint32 step_num = 1; - // The step duration in picoseconds. - optional uint64 duration_ps = 2; - // The infeed duration in picoseconds. - optional uint64 infeed_duration_ps = 3; - // The outfeed duration in picoseconds. - optional uint64 host_outfeed_ps = 8; - // The start time of this step in picoseconds. - optional uint64 begin_ps = 4; - // The waiting time within this step in picoseconds. - optional uint64 wait_duration_ps = 5; - // The unit b outfeed duration in picoseconds. - optional uint64 unit_b_outfeed_ps = 9; - // The time spent on cross-replica-sum in picoseconds. - optional uint64 crs_duration_ps = 6; - // Percentage of unit b time spent on infeed. - optional double unit_b_infeed_percent = 7; -} - -// Result proto for a sequence of steps. -message StepSequenceResult { - // A sequence of StepInfoResults. - repeated StepInfoResult step_sequence = 1; -} - -// Result proto for a StepDatabase. -message StepDatabaseResult { - // A map from core_id to StepSequenceResult. - map step_sequence_per_core = 1; -} - -// Result proto for looping-related metrics. -message LoopingResult { - // The total iteration time in nanoseconds. - optional double iteration_time_ns = 1; - // The total number of iterations. - optional int32 num_iterations = 2; - // The total computation time in nanoseconds. - optional double computation_time_ns = 3; - // The total number of computations. - optional int32 num_computations = 4; -} - -// Result proto for HloExtraInfo. -message HloExtraInfoResult { - // Category of the HLO op given by the compiler. - optional string category = 1; - // The long name of the HLO that includes the dimensions. - optional string long_name = 2; - // The per-TPU-core batch size inferred from this HLO. - optional int64 per_core_batch_size = 3; -} - -// Result proto for HloExtraInfoMap. -message HloExtraInfoMapResult { - // A map from HLO name to HloExtraInfo. - map hlo_extrainfo_map = 1; -} - -// Result proto for host-independent job information. -message HostIndependentJobInfoResult { - // The change-list number of this build. - optional int64 change_list = 1; - // The time of this build. - optional int64 build_time = 2; - // The target of this build. - optional string build_target = 3; -} - -// Result proto for host-dependent job information. -message HostDependentJobInfoResult { - // This ID of the host where the job was run on. - optional string host_id = 1; - // The command line used to run the job. - optional string command_line = 2; - // The start time of the job on this host. - optional int64 start_time = 3; -} - -// Result proto for RunEnvironment (the run environment of a profiling session). -message RunEnvironmentResult { - // Number of hosts used. - optional int32 host_count = 1; - // The type of TPU used. - optional string tpu_type = 2; - // The number of TPU cores used. - optional int32 tpu_core_count = 3; - // The per-TPU-core batch size. - optional int32 per_core_batch_size = 4; - // Host-independent job information. - optional HostIndependentJobInfoResult host_independent_job_info = 5; - // Host-dependent job information. - repeated HostDependentJobInfoResult host_dependent_job_info = 6; - // The number of replicas, corresponds to input parallelism. - // If there is no model parallelism, replica_count = tpu_core_count - optional int32 replica_count = 7; - // The number of cores used for a single replica, e.g. model parallelism. - // If there is no model parallelism, then num_cores_per_replica = 1 - optional int32 num_cores_per_replica = 8; -} - -// The types of host operations that are tracked. -enum HostOp { - // Invalid host op. - kINVALIDHostOp = 0; - // Each of host op type has two parts: - // (1) the stage where the op happens and (2) the op name. - // stage = Input Data Producer, op = Get Next Batch. - kInputDataProducerGetNextBatch = 1; - // stage = Input Data Producer, op = Session Run. - kInputDataProducerSessionRun = 2; - // stage = Input Data Producer, op = Forward Batch. - kInputDataProducerForwardBatch = 3; - // stage = Infeed Thread, op = Get Next Batch. - kInfeedThreadGetNextBatch = 4; - // stage = Infeed Thread, op = Session Run. - kInfeedThreadSessionRun = 5; - // stage = Infeed Thread, op = Forward Batch. - kInfeedThreadForwardBatch = 6; - // stage = Outfeed Thread, op = Get Next Batch. - kOutfeedThreadGetNextBatch = 7; - // stage = Outfeed Thread, op = Session Run. - kOutfeedThreadSessionRun = 8; - // stage = Outfeed Thread, op = Forward Batch. - kOutfeedThreadForwardBatch = 9; -} - -// Result proto for the host ops per TPU step. -message HostOpsPerTpuStep { - // Whether the data in this message is valid. - optional bool valid = 1 [default = false]; - // The current TPU step number. - optional uint32 tpu_step_num = 2; - // The beginning time of the current TPU step on the device in picoseconds. - optional uint64 tpu_step_begin_ps = 3; - // The ending time of the current TPU step on the device in picoseconds. - optional uint64 tpu_step_end_ps = 4; - // For each possible host operation, maps to the difference between the TPU - // step number that the host op targets and the current TPU step number. - // The key is HostOp, value is the step difference. - map step_diffs = 5; -} - -message HostOpsDetailsPerCore { - // Map from core id to HostOpsPerTpuStep. - map core_map = 1; -} - -message HostOpsDetailsPerHost { - // Map from hostname to a map from core id to HostOpsPerTpuStep. - map host_map = 1; -} - -// Result proto for the host ops for all TPU steps. -message HostOpsResult { - reserved 1; // (was repeated HostOpsPerTpuStep host_op_sequence) - // A sequence of records with one for each TPU step. Each record - // is a map from hostname to a map from core id to HostOpsPerTpuStep. - repeated HostOpsDetailsPerHost hostops_details = 2; -} - -// Result proto for TfStatsHelper. -message TfOpStats { - // The result for the TF-metric database. - optional OpMetricsDbResult tf_metrics_db = 1; - // The result for the HLO-metric database. - optional OpMetricsDbResult hlo_metrics_db = 2; - // The result for the step database. - optional StepDatabaseResult step_db = 3; - // The result for the looping-related metrics. - optional LoopingResult looping = 4; - // The result for the HloExtraInfoMap. - optional HloExtraInfoMapResult hlo_extrainfo_map = 5; - // Overall matrix unit utilization in percentage. - optional double matrix_unit_utilization_percent = 6; - // The run environment of this profiling session. - optional RunEnvironmentResult run_environment = 7; - // The result for the host operations. - optional HostOpsResult host_ops = 8; - // A map from core ID to name. - map core_id_to_name_map = 9; - // The result for hw unit b stats. - optional bytes unit_b_stats = 10; -} 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/proto/BUILD b/tensorflow/contrib/tpu/proto/BUILD index c20cab844cfaf083be2702a29ac2a152c7b72c2a..ea98ee25c89e1b7bef39276bae5c98bf382dbd7f 100644 --- a/tensorflow/contrib/tpu/proto/BUILD +++ b/tensorflow/contrib/tpu/proto/BUILD @@ -49,6 +49,15 @@ tf_proto_library( visibility = ["//visibility:public"], ) +tf_proto_library( + name = "dynamic_padding_proto", + srcs = [ + "dynamic_padding.proto", + ], + cc_api_version = 2, + visibility = ["//visibility:public"], +) + tf_proto_library_py( name = "compilation_result_proto", srcs = [ diff --git a/tensorflow/contrib/tpu/proto/dynamic_padding.proto b/tensorflow/contrib/tpu/proto/dynamic_padding.proto new file mode 100644 index 0000000000000000000000000000000000000000..c9ebf181169a583d774ef77ca0b8c243ce733615 --- /dev/null +++ b/tensorflow/contrib/tpu/proto/dynamic_padding.proto @@ -0,0 +1,19 @@ +syntax = "proto3"; + +option cc_enable_arenas = true; + +package tensorflow.tpu; + +// A mapping between the dynamic shape dimension of an input and the arg that +// represents the real shape. +message PaddingMap { + // Input arg index with dynamic shapes. + int32 arg_index = 1; + + // The dynamic shape dimension index. + int32 shape_index = 2; + + // The arg index that dynamic dimension maps to, which represents the value + // of the real shape. + int32 padding_arg_index = 3; +} diff --git a/tensorflow/contrib/tpu/proto/optimization_parameters.proto b/tensorflow/contrib/tpu/proto/optimization_parameters.proto index c2e3be03db0e4cca1a664f9e79aa9107384de312..bc50c613f3d2a09f9e51353fab4938055549a4cd 100644 --- a/tensorflow/contrib/tpu/proto/optimization_parameters.proto +++ b/tensorflow/contrib/tpu/proto/optimization_parameters.proto @@ -9,9 +9,38 @@ message ClippingLimits { google.protobuf.FloatValue upper = 2; // +inf if not set } -// Get the learning rate from the parameters of the SendTPUEmbeddingGradients -// op. +// Dynamic learning rate specification in the TPUEmbeddingConfiguration. The +// actual learning rates are provided as a scalar input list to the +// SendTPUEmbeddingGradients Op indexed by their tag specified through the +// following proto. message DynamicLearningRate { + // For tables where learning rates are dynamically computed and communicated + // to the TPU embedding program, a tag must be specified for the learning + // rate. + // + // The tag must be a non-negative integer. The total number of unique tags + // must be less than or equal to the number of tables in the TPU embedding + // configuration (a table does not specify any tag if it uses a constant + // learning rate, and specifies exactly one tag if it uses dynamic learning + // rates). + // + // All tags in the range [0, number_of_unique_tags) must be present in the TPU + // embedding configuration, i.e. a tag cannot be skipped if a different tag + // numerically greater than it is used in the configuration. + // + // If multiple tables specify the same tag, they *MUST* have + // the same dynamic learning rate, for example, their dynamic learning rate + // could be computed by the same TensorFlow sub-graph. The partitioning of the + // embedding layer would be more optimal if the number_of_unique_tags is as + // *LOW* as possible, i.e., if many tables share the same tag. + // + // The learning_rate input of the SendTPUEmbeddingGradients op is used to + // communicate dynamic learning rates to the TPU embedding program. + // The learning_rate input is a list of scalars where the size of the list is + // equal to the number of unique tags. The learning rate associated with a + // particular tag is specified by populating its corresponding index in the + // list of learning_rate scalars. + int32 tag = 1; } // Source of learning rate to use. @@ -154,6 +183,14 @@ message OptimizationParameters { // updates; not present means no limits are applied. ClippingLimits gradient_clipping_limits = 7; + // Amount of weight decay to apply; see weight_decay_optimizers.py for + // details. Almost all optimizers are supported with this option (MDL Adagrad + // Light does not work, and SGD does not behave as expected if it is enabled). + // Although there is no check, users who want weight decay will probably also + // want to enable gradient accumulation as well so that the decay will happen + // once per minibatch. + float weight_decay_factor = 16; + // Whether to use gradient accumulation (do two passes over the input // gradients: one to accumulate them into a temporary array and another to // apply them using the actual optimization algorithm). This feature is @@ -178,7 +215,8 @@ message OptimizationParameters { } // Specification of an optimization algorithm's state variables (both the main -// value vector and any extra accumulators, etc.). +// value vector and any extra accumulators, etc.). This proto is only used +// internally by the TPU software and is not exposed directly to the TF model. message StateVariableSpecification { // Parameter name for the state variable. string name = 1; @@ -186,6 +224,20 @@ message StateVariableSpecification { // A normal state variable that should be saved and restored in checkpoints // and used as an input or output to non-debug TensorFlow ops. message UserDefined { + // For padding embedding rows, this field specifies the initial value to be + // used. Separate initial values need to be specified for the embeddings and + // any extra accumulators. The initial values should be specified so as to + // maintain two invariants during model training: + // (1) The embedding vector multiplied by zero returns a vector containing + // all zeros. To maintain this invariant, the embedding values should + // never be NaNs or +-infinity. + // (2) Repeatedly applying the optimizer using a gradient vector of all + // zeros does not cause the embeddings or slot variables to become NaNs + // or +-infinity. + // The padding row is looked up when no embedding IDs are present for a + // feature. The semantics of embedding lookup dictate that the output must + // be zero under this scenario. + double padding_initial_value = 1; } // A state variable that should be filled with a constant and normally hidden diff --git a/tensorflow/contrib/tpu/python/ops/tpu_ops.py b/tensorflow/contrib/tpu/python/ops/tpu_ops.py index 968adccf2b82e80bd008e54d3af614bd74852795..9260e7b8a800c3bf160923af95867d44342000a3 100644 --- a/tensorflow/contrib/tpu/python/ops/tpu_ops.py +++ b/tensorflow/contrib/tpu/python/ops/tpu_ops.py @@ -137,10 +137,18 @@ if platform.system() != "Windows": """ return gen_tpu_ops.collective_permute(x, source_target_pairs, name=name) + @ops.RegisterGradient("CollectivePermute") + def _collective_permute_grad(op, grad): + # The gradient of a collective permute operation is also a collective + # permute, but with source/target pairs reversed. The gradient with respect + # to input argument `source_target_pairs` is `None`. + source_target_pairs = op.inputs[1][:, ::-1] + return [gen_tpu_ops.collective_permute(grad, source_target_pairs), None] + @ops.RegisterGradient("CrossReplicaSum") def _cross_replica_sum_grad(op, grad): # The gradient of a cross replica sum is also a cross-replica sum. - # The graident with respect to group_assignment is None. + # The gradient with respect to group_assignment is None. return [gen_tpu_ops.cross_replica_sum(grad, op.inputs[1]), None] # This extra type checking exists to give a more helpful error message in @@ -209,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/contrib/tpu/proto/ + 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: @@ -329,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/tpu/python/tpu/async_checkpoint.py b/tensorflow/contrib/tpu/python/tpu/async_checkpoint.py index 78253d83fc4dcc0df153d4441373b755b7349560..1b09ce173a64ba3f93ec019c8fd65dc4710f0fcf 100644 --- a/tensorflow/contrib/tpu/python/tpu/async_checkpoint.py +++ b/tensorflow/contrib/tpu/python/tpu/async_checkpoint.py @@ -80,6 +80,8 @@ class AsyncCheckpointSaverHook(basic_session_run_hooks.CheckpointSaverHook): self._summary_writer = None self._global_step_tensor = None + self._last_checkpoint_step = None + def _set_steps_per_run(self, steps_per_run): self._steps_per_run = steps_per_run @@ -102,7 +104,8 @@ class AsyncCheckpointSaverHook(basic_session_run_hooks.CheckpointSaverHook): training_util.write_graph( ops.get_default_graph().as_graph_def(add_shapes=True), self._checkpoint_dir, "graph.pbtxt") - self._write_graph_thread = threading.Thread(target=_write_graph_fn) + self._write_graph_thread = threading.Thread(target=_write_graph_fn, + args=[self]) self._write_graph_thread.start() saver_def = self._get_saver().saver_def if self._get_saver() else None @@ -136,8 +139,7 @@ class AsyncCheckpointSaverHook(basic_session_run_hooks.CheckpointSaverHook): last_step = session.run(self._global_step_tensor) - # Save the last checkpoint synchronously if needed. - if last_step != self._timer.last_triggered_step(): + if self._last_checkpoint_step != last_step: self._save(session, last_step, asynchronous=False) for l in self._listeners: @@ -163,15 +165,17 @@ class AsyncCheckpointSaverHook(basic_session_run_hooks.CheckpointSaverHook): SessionLog( status=SessionLog.CHECKPOINT, checkpoint_path=self._save_path), step) + + for l in self._listeners: + l.after_save(session, step) + end_time = time.time() logging.info("Checkpoint actual writing time: (%.3f sec)", end_time - start_time) logging.info("Checkpoint finished for %d into %s.", step, self._save_path) - for l in self._listeners: - l.before_save(session, step) - if not asynchronous: + self._last_checkpoint_step = step _save_fn() return @@ -181,6 +185,7 @@ class AsyncCheckpointSaverHook(basic_session_run_hooks.CheckpointSaverHook): logging.info("Saver thread still in progress, skipping checkpoint.") return + self._last_checkpoint_step = step self._save_thread = threading.Thread(target=_save_fn) self._save_thread.start() diff --git a/tensorflow/contrib/tpu/python/tpu/datasets.py b/tensorflow/contrib/tpu/python/tpu/datasets.py index c694e9c1bca10d9930492c29dd1c3cbc7f7f5d04..bc0cd41d210ac6f8de1b20ebf744ee1e1dd04137 100644 --- a/tensorflow/contrib/tpu/python/tpu/datasets.py +++ b/tensorflow/contrib/tpu/python/tpu/datasets.py @@ -133,7 +133,7 @@ def StreamingFilesDataset(files, with ops.device('/job:%s' % file_reader_job): if isinstance(files, str): source_dataset = dataset_ops.Dataset.list_files(files) - elif isinstance(files, dataset_ops.Dataset): + elif isinstance(files, dataset_ops.DatasetV2): source_dataset = files else: raise ValueError('files was not a string or a dataset: %s' % files) @@ -142,21 +142,18 @@ def StreamingFilesDataset(files, source_dataset = source_dataset.shuffle( buffer_size=filename_shuffle_buffer_size) - # NOTE: We perform the `repeat` on the source dataset, because the output - # dataset does not currently have enough information to recreate an iterator - # over the source dataset when it reaches the end. - source_dataset = source_dataset.repeat(num_epochs) - source_dataset = source_dataset.apply( interleave_ops.parallel_interleave( reader_fn, cycle_length=num_parallel_reads, sloppy=sloppy)) + source_dataset = source_dataset.repeat(num_epochs) + if batch_transfer_size: source_dataset = source_dataset.batch(batch_transfer_size) source_dataset = source_dataset.prefetch(1) - source_iterator = source_dataset.make_one_shot_iterator() + source_iterator = dataset_ops.make_one_shot_iterator(source_dataset) source_handle = source_iterator.string_handle() @function.Defun(dtypes.string) diff --git a/tensorflow/contrib/tpu/python/tpu/datasets_test.py b/tensorflow/contrib/tpu/python/tpu/datasets_test.py index b58d05eac56f3586e183333f7c1a3867ee57456c..8a94f527bb6dffa48e71e6500ae5e9e9589fbf5c 100644 --- a/tensorflow/contrib/tpu/python/tpu/datasets_test.py +++ b/tensorflow/contrib/tpu/python/tpu/datasets_test.py @@ -27,6 +27,7 @@ from tensorflow.python.client import session from tensorflow.python.data.ops import dataset_ops from tensorflow.python.data.ops import readers from tensorflow.python.framework import dtypes +from tensorflow.python.framework import ops from tensorflow.python.framework import tensor_shape from tensorflow.python.lib.io import python_io from tensorflow.python.platform import test @@ -55,6 +56,7 @@ class DatasetsTest(test.TestCase): session_config = config_pb2.ConfigProto(cluster_def=self._cluster_def) self._sess = session.Session(self._worker.target, config=session_config) + self._worker_device = '/job:' + worker_job.name def testTextLineDataset(self): all_contents = [] @@ -70,7 +72,8 @@ class DatasetsTest(test.TestCase): dataset = datasets.StreamingFilesDataset( os.path.join(self.get_temp_dir(), 'text_line.*.txt'), filetype='text') - iterator = dataset.make_initializable_iterator() + with ops.device(self._worker_device): + iterator = dataset_ops.make_initializable_iterator(dataset) self._sess.run(iterator.initializer) get_next = iterator.get_next() @@ -94,7 +97,8 @@ class DatasetsTest(test.TestCase): dataset = datasets.StreamingFilesDataset( os.path.join(self.get_temp_dir(), 'tf_record*'), filetype='tfrecord') - iterator = dataset.make_initializable_iterator() + with ops.device(self._worker_device): + iterator = dataset_ops.make_initializable_iterator(dataset) self._sess.run(iterator.initializer) get_next = iterator.get_next() @@ -121,7 +125,8 @@ class DatasetsTest(test.TestCase): dataset = datasets.StreamingFilesDataset(filenames, filetype='tfrecord') - iterator = dataset.make_initializable_iterator() + with ops.device(self._worker_device): + iterator = dataset_ops.make_initializable_iterator(dataset) self._sess.run(iterator.initializer) get_next = iterator.get_next() @@ -154,7 +159,8 @@ class DatasetsTest(test.TestCase): os.path.join(self.get_temp_dir(), 'fixed_length*'), filetype=FixedLengthFile) - iterator = dataset.make_initializable_iterator() + with ops.device(self._worker_device): + iterator = dataset_ops.make_initializable_iterator(dataset) self._sess.run(iterator.initializer) get_next = iterator.get_next() @@ -177,7 +183,8 @@ class DatasetsTest(test.TestCase): dataset = datasets.StreamingFilesDataset( dataset_ops.Dataset.range(10), filetype=gen_dataset) - iterator = dataset.make_initializable_iterator() + with ops.device(self._worker_device): + iterator = dataset_ops.make_initializable_iterator(dataset) self._sess.run(iterator.initializer) get_next = iterator.get_next() diff --git a/tensorflow/contrib/tpu/python/tpu/device_assignment.py b/tensorflow/contrib/tpu/python/tpu/device_assignment.py index b9e2a4287a97d2f2d8cb2fb73a12d3f24f090007..3313dc749c2c7606101b2dc96614df2d052dfed1 100644 --- a/tensorflow/contrib/tpu/python/tpu/device_assignment.py +++ b/tensorflow/contrib/tpu/python/tpu/device_assignment.py @@ -25,20 +25,32 @@ from six.moves import xrange # pylint: disable=redefined-builtin from tensorflow.contrib.tpu.python.tpu.topology import Topology -def _tpu_device_name(job, task, device): - """Returns the device name for the TPU `device` on `task` of `job`.""" - if job is None: - return "/task:%d/device:TPU:%d" % (task, device) - else: - return "/job:%s/task:%d/device:TPU:%d" % (job, task, device) +SINGLE_CORE_ASSIGNMENT = [[[0, 0, 0]]] -def _tpu_host_device_name(job, task): - """Returns the device name for the CPU device on `task` of `job`.""" - if job is None: - return "/task:%d/device:CPU:0" % task - else: - return "/job:%s/task:%d/device:CPU:0" % (job, task) +def _compute_task_and_cores_to_replicas(core_assignment, topology): + """Computes a nested dict which maps task and logical core to replicas.""" + task_and_cores_to_replicas = {} + for replica in xrange(core_assignment.shape[0]): + for logical_core in xrange(core_assignment.shape[1]): + coordinates = core_assignment[replica, logical_core, :] + task_id = topology.task_ordinal_at_coordinates(coordinates) + if task_id not in task_and_cores_to_replicas: + task_and_cores_to_replicas[task_id] = {} + if logical_core not in task_and_cores_to_replicas[task_id]: + task_and_cores_to_replicas[task_id][logical_core] = set() + + task_and_cores_to_replicas[task_id][logical_core].add(replica) + + task_to_sorted_replica_id = {} + + for task, core_to_replicas in task_and_cores_to_replicas.items(): + core_to_sorted_replicas = {} + for core, replicas in core_to_replicas.items(): + core_to_sorted_replicas[core] = sorted(replicas) + + task_to_sorted_replica_id[task] = core_to_sorted_replicas + return task_to_sorted_replica_id class DeviceAssignment(object): @@ -68,10 +80,7 @@ class DeviceAssignment(object): core_assignment = np.asarray(core_assignment, dtype=np.int32) self._topology = topology - self._topology_tasks, self._topology_devices = ( - self._invert_topology(topology)) - topology_rank = self._topology_tasks.ndim if core_assignment.ndim != 3: raise ValueError("core_assignment must be a rank 3 numpy array, " "got shape {}".format(core_assignment.shape)) @@ -79,52 +88,15 @@ class DeviceAssignment(object): self._num_replicas = core_assignment.shape[0] self._num_cores_per_replica = core_assignment.shape[1] - if core_assignment.shape[-1] != topology_rank: + if core_assignment.shape[-1] != topology.mesh_rank: raise ValueError( "minor dimension of core_assignment must have size equal to topology " - "rank ({}), got shape {}".format(topology_rank, + "rank ({}), got shape {}".format(topology.mesh_rank, core_assignment.shape)) self._core_assignment = core_assignment - self._task_and_cores_to_replicas = self._compute_task_and_cores_to_replicas( - self._core_assignment, self._topology_tasks) - - def _invert_topology(self, topology): - """Inverts a [task,device,axis] topology to [x,y,z] -> task/device maps.""" - mesh_shape = topology.mesh_shape - tasks = np.full(list(mesh_shape), -1, dtype=np.int32) - devices = np.full(list(mesh_shape), -1, dtype=np.int32) - for task in xrange(topology.device_coordinates.shape[0]): - for device in xrange(topology.device_coordinates.shape[1]): - x, y, z = topology.device_coordinates[task, device, :] - tasks[x, y, z] = task - devices[x, y, z] = device - return tasks, devices - - def _compute_task_and_cores_to_replicas(self, core_assignment, - topology_tasks): - """Computes a nested dict which maps task and logical core to replicas.""" - task_and_cores_to_replicas = {} - for replica in xrange(core_assignment.shape[0]): - for logical_core in xrange(core_assignment.shape[1]): - x, y, z = core_assignment[replica, logical_core, :] - task_id = topology_tasks[x, y, z] - if task_id not in task_and_cores_to_replicas: - task_and_cores_to_replicas[task_id] = {} - if logical_core not in task_and_cores_to_replicas[task_id]: - task_and_cores_to_replicas[task_id][logical_core] = set() - - task_and_cores_to_replicas[task_id][logical_core].add(replica) - - task_to_sorted_replica_id = {} - - for task, core_to_replicas in task_and_cores_to_replicas.items(): - core_to_sorted_replicas = {} - for core, replicas in core_to_replicas.items(): - core_to_sorted_replicas[core] = sorted(replicas) - - task_to_sorted_replica_id[task] = core_to_sorted_replicas - return task_to_sorted_replica_id + self._task_and_cores_to_replicas = _compute_task_and_cores_to_replicas( + self._core_assignment, topology) @property def topology(self): @@ -179,18 +151,17 @@ class DeviceAssignment(object): def tpu_ordinal(self, replica=0, logical_core=0): """Returns the ordinal of the TPU device assigned to a logical core.""" coordinates = self._coordinates(replica, logical_core) - return self._topology_devices[coordinates] + return self._topology.tpu_device_ordinal_at_coordinates(coordinates) def host_device(self, replica=0, logical_core=0, job=None): """Returns the CPU device attached to a logical core.""" coordinates = self._coordinates(replica, logical_core) - return _tpu_host_device_name(job, self._topology_tasks[coordinates]) + return self._topology.cpu_device_name_at_coordinates(coordinates, job=job) def tpu_device(self, replica=0, logical_core=0, job=None): """Returns the name of the TPU device assigned to a logical core.""" coordinates = self._coordinates(replica, logical_core) - return _tpu_device_name(job, self._topology_tasks[coordinates], - self._topology_devices[coordinates]) + return self._topology.tpu_device_name_at_coordinates(coordinates, job=job) def device_assignment(topology, diff --git a/tensorflow/contrib/tpu/python/tpu/feature_column.py b/tensorflow/contrib/tpu/python/tpu/feature_column.py new file mode 100644 index 0000000000000000000000000000000000000000..8edf131bc24fd003806263570b63ee8514c49896 --- /dev/null +++ b/tensorflow/contrib/tpu/python/tpu/feature_column.py @@ -0,0 +1,429 @@ +# 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. +# =================================================================== +"""TPU Feature Column Library.""" +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import math + +from tensorflow.contrib.tpu.python.tpu import tpu +from tensorflow.contrib.tpu.python.tpu import tpu_function +from tensorflow.python.feature_column import feature_column as fc +from tensorflow.python.feature_column import feature_column_lib as fc_lib +from tensorflow.python.framework import ops +from tensorflow.python.ops import init_ops +from tensorflow.python.ops import variable_scope +# pylint: disable=protected-access + + +_TPU_FC_TO_SCOPE = '_tpu_feature_column_scope' +_SUPPORTED_CATEGORICAL_COLUMNS = (fc._IdentityCategoricalColumn, + fc._VocabularyFileCategoricalColumn, + fc._VocabularyListCategoricalColumn, + fc._WeightedCategoricalColumn, + fc_lib.IdentityCategoricalColumn, + fc_lib.VocabularyFileCategoricalColumn, + fc_lib.VocabularyListCategoricalColumn, + fc_lib.WeightedCategoricalColumn) + + +def embedding_column(categorical_column, + dimension, + combiner='mean', + initializer=None): + """TPU embedding_column for `tf.feature_column.embedding_column`. + + Note that the interface for TPU embedding_column is different from the non-TPU + version. The following args available for the non-TPU version are NOT + supported: ckpt_to_load_from, tensor_name_in_ckp, max_norm and trainable. + + Args: + categorical_column: A categorical_column returned from + categorical_column_with_identity, weighted_categorical_column, + categorical_column_with_vocabulary_list or + categorical_column_with_vocabulary_file. + dimension: An integer specifying dimension of the embedding, must be > 0. + combiner: A string specifying how to reduce if there are multiple entries + in a single row. For more information, see + `tf.feature_column.embedding_column`. + initializer: A variable initializer function to be used in embedding + variable initialization. If not specified, defaults to + `tf.truncated_normal_initializer` with mean `0.0` and standard deviation + `1/sqrt(dimension)`. + + Returns: + A _TPUEmbeddingColumn. + + Raises: + ValueError: if `dimension` not > 0. + ValueError: if `initializer` is specified but not callable. + """ + if not isinstance(categorical_column, _SUPPORTED_CATEGORICAL_COLUMNS): + raise TypeError( + 'categorical_column for tpu ' + ' embedding_column must be type %s, got %s.' % (' or '.join([ + cc.__name__ for cc in _SUPPORTED_CATEGORICAL_COLUMNS + ]), type(categorical_column))) + if (dimension is None) or (dimension < 1): + raise ValueError('Invalid dimension {}.'.format(dimension)) + + if (initializer is not None) and (not callable(initializer)): + raise ValueError('initializer must be callable if specified. ' + 'Embedding of column_name: {}'.format( + categorical_column.name)) + if initializer is None: + initializer = init_ops.truncated_normal_initializer( + mean=0.0, stddev=1 / math.sqrt(dimension)) + + embedding_shape = categorical_column._num_buckets, dimension # pylint: disable=protected-access + + def _creator(weight_collections, scope): + embedding_column_layer = fc._EmbeddingColumnLayer( + embedding_shape=embedding_shape, + initializer=initializer, + weight_collections=weight_collections, + trainable=True, + name='embedding_column_layer') + return embedding_column_layer(None, scope=scope) # pylint: disable=not-callable + + column = _TPUEmbeddingColumn( + categorical_column=categorical_column, + dimension=dimension, + combiner=combiner, + layer_creator=_creator, + ckpt_to_load_from=None, + tensor_name_in_ckpt=None, + max_norm=None, + trainable=True) + # For Embedding column, the initializer is hidden inside the creator Fn, which + # is not accessiable later. So, we attach it to a speicial field. Also note + # that non-TPU Embedding column and non-TPU shared Embedding column handle the + # initializer differently. See shared_embedding_columns for details. + column._tpu_initializer = initializer + return column + + +def shared_embedding_columns(categorical_columns, + dimension, + combiner='mean', + initializer=None, + shared_embedding_collection_name=None): + """List of dense columns that convert from sparse, categorical input.""" + for categorical_column in categorical_columns: + if not isinstance(categorical_column, _SUPPORTED_CATEGORICAL_COLUMNS): + raise TypeError( + 'categorical_column for tpu ' + ' shared_embedding_columns must be type %s, got %s.' % (' or '.join([ + cc.__name__ for cc in _SUPPORTED_CATEGORICAL_COLUMNS + ]), type(categorical_column))) + columns = fc_lib.shared_embedding_columns( + categorical_columns, + dimension, + combiner=combiner, + initializer=initializer, + shared_embedding_collection_name=shared_embedding_collection_name, + ckpt_to_load_from=None, + tensor_name_in_ckpt=None, + max_norm=None, + trainable=True) + + # Use the initializer and shared_embedding_collection_name to create TPU + # version + initializer = columns[0].initializer + shared_embedding_collection_name = columns[0].shared_embedding_collection_name + tpu_columns = [] + + # Create the state (_SharedEmbeddingColumnLayer) here. + for categorical_column in categorical_columns: + column = _TPUSharedEmbeddingColumn( + categorical_column=categorical_column, + dimension=dimension, + combiner=combiner, + initializer=initializer, + shared_embedding_collection_name=shared_embedding_collection_name, + ckpt_to_load_from=None, + tensor_name_in_ckpt=None, + max_norm=None, + trainable=True) + tpu_columns.append(column) + + return tpu_columns + + +class _TPUBaseEmbeddingColumn(object): + """Base class for TPU Embedding Column.""" + + def __init__(self, categorical_column): + self._tpu_categorical_column = categorical_column + + def get_combiner(self): + """Returns the embedding combiner.""" + raise NotImplementedError('not implemented') + + def get_embedding_table_size(self): + """Returns the embedding table size, tuple of vocab size and dimension.""" + raise NotImplementedError('not implemented') + + def get_feature_key_name(self): + """Returns the feature key name in the features dict.""" + raise NotImplementedError('not impl') + + def get_weight_key_name(self): + """Return the key name for weights.""" + raise NotImplementedError('not impl') + + def get_embedding_var_name(self): + """Returns the embedding variable name. + + Feature key name and embedding variable name are usually one-to-one mapping. + But for shared embedding columns, it is many-to-one mapping. + """ + raise NotImplementedError('not impl') + + def get_initializer(self): + """Returns the initializer.""" + raise NotImplementedError('not impl') + + def is_categorical_column_weighted(self): + """Check if the categorical column of the embedding column is weighted.""" + raise NotImplementedError('not impl') + + +class _TPUEmbeddingColumn(_TPUBaseEmbeddingColumn, fc._EmbeddingColumn): + """Core Embedding Column.""" + + def __new__(cls, + categorical_column, + dimension, + combiner='mean', + layer_creator=None, + ckpt_to_load_from=None, + tensor_name_in_ckpt=None, + max_norm=None, + trainable=True): + # Note, args ckpt_to_load_from, tensor_name_in_ckpt, max_norm and trainable + # are not supported on TPU. They are solely for matching the signature of + # __new__ of parent class fc._EmbeddingColumn. + return fc._EmbeddingColumn.__new__( + cls, + categorical_column, + dimension, + combiner=combiner, + layer_creator=layer_creator, + ckpt_to_load_from=ckpt_to_load_from, + tensor_name_in_ckpt=tensor_name_in_ckpt, + max_norm=max_norm, + trainable=trainable) + + def __init__(self, + categorical_column, + dimension, + combiner='mean', + layer_creator=None, + ckpt_to_load_from=None, + tensor_name_in_ckpt=None, + max_norm=None, + trainable=True): + _TPUBaseEmbeddingColumn.__init__(self, categorical_column) + self._key = None + + def get_combiner(self): + return self.combiner + + def get_embedding_table_size(self): + """Returns num_ids and width.""" + return (self.categorical_column._num_buckets, self.dimension) + + def get_feature_key_name(self): + """get_feature_key_name.""" + if self.is_categorical_column_weighted(): + return self.categorical_column.categorical_column.name + return self.categorical_column.name + + def get_weight_key_name(self): + """get_weight_key_name.""" + if self.is_categorical_column_weighted(): + return self.categorical_column.weight_feature_key + return None + + def get_embedding_var_name(self): + """get_embedding_var_name.""" + return self.categorical_column.name + + def get_initializer(self): + return self._tpu_initializer + + def is_categorical_column_weighted(self): + """Check if the categorical column of the embedding column is weighted.""" + if isinstance( + self.categorical_column, + ( + fc._WeightedCategoricalColumn, # pylint: disable=protected-access + fc_lib.WeightedCategoricalColumn)): + return True + return False + + def _get_dense_tensor(self, inputs, weight_collections=None, trainable=None): + if tpu.under_tpu_inference_context(): + def host_computation(): + return fc._EmbeddingColumn._get_dense_tensor( + self, inputs, weight_collections, trainable) + return tpu.outside_compilation(host_computation) + + if _is_running_on_cpu(): + return fc._EmbeddingColumn._get_dense_tensor( + self, inputs, weight_collections, trainable) + + # TPU mode + # Get the embeddings from the LazyBuilder. + tensor = inputs.get(self.get_feature_key_name()) + + # Add to collection for _create_tpu_embedding_variables_and_ops + _record_variable_scope_and_name(self.get_embedding_var_name(), + 'embedding_weights') + + return tensor + + +class _TPUSharedEmbeddingColumn(_TPUBaseEmbeddingColumn, + fc._SharedEmbeddingColumn): + """Core Shared Embedding Column.""" + + def __new__(cls, + categorical_column, + dimension, + combiner='mean', + initializer=None, + shared_embedding_collection_name=None, + ckpt_to_load_from=None, + tensor_name_in_ckpt=None, + max_norm=None, + trainable=True): + return fc._SharedEmbeddingColumn.__new__( + cls, + categorical_column, + dimension, + combiner=combiner, + initializer=initializer, + shared_embedding_collection_name=shared_embedding_collection_name, + ckpt_to_load_from=ckpt_to_load_from, + tensor_name_in_ckpt=tensor_name_in_ckpt, + max_norm=max_norm, + trainable=trainable) + + def __init__(self, + categorical_column, + dimension, + combiner='mean', + initializer=None, + shared_embedding_collection_name=None, + ckpt_to_load_from=None, + tensor_name_in_ckpt=None, + max_norm=None, + trainable=True): + + _TPUBaseEmbeddingColumn.__init__(self, categorical_column) + self._key = None + + def get_combiner(self): + return self.combiner + + def get_embedding_table_size(self): + """Returns num_ids and width.""" + return (self.categorical_column._num_buckets, self.dimension) + + def get_feature_key_name(self): + """get_feature_key_name.""" + if self.is_categorical_column_weighted(): + return self.categorical_column.categorical_column.name + return self.categorical_column.name + + def get_weight_key_name(self): + """get_weight_key_name.""" + if self.is_categorical_column_weighted(): + return self.categorical_column.weight_feature_key + return None + + def get_embedding_var_name(self): + """get_embedding_var_name.""" + return self.shared_embedding_collection_name + + def get_initializer(self): + return self.initializer + + def is_categorical_column_weighted(self): + """Check if the categorical column of the embedding column is weighted.""" + if isinstance( + self.categorical_column, + ( + fc._WeightedCategoricalColumn, # pylint: disable=protected-access + fc_lib.WeightedCategoricalColumn)): + return True + return False + + def _get_dense_tensor(self, inputs, weight_collections=None, trainable=None): + if tpu.under_tpu_inference_context(): + def host_computation(): + return fc._SharedEmbeddingColumn._get_dense_tensor( + self, inputs, weight_collections, trainable) + return tpu.outside_compilation(host_computation) + + if _is_running_on_cpu(): + return fc._SharedEmbeddingColumn._get_dense_tensor( + self, inputs, weight_collections, trainable) + + # TPU mode + # Get the embeddings from the LazyBuilder. + tensor = inputs.get(self.get_feature_key_name()) + + # Add to collection for _create_tpu_embedding_variables_and_ops + _record_variable_scope_and_name( + self.get_embedding_var_name(), + 'embedding_weights', + is_shared_embedding=True) + return tensor + + +def _record_variable_scope_and_name(embedding_var_name, + embedding_var_name_in_fc, + is_shared_embedding=False): + """Add embedding variable name and scope to collection.""" + g = ops.get_default_graph() + collection = g.get_collection_ref(_TPU_FC_TO_SCOPE) + if not collection: + collection.append({}) + + var_def_dict = collection[0] + + captured_scope = None + + if is_shared_embedding and (embedding_var_name in var_def_dict): + 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, ' + '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, + embedding_var_name_in_fc) + + +def _is_running_on_cpu(): + """Returns True if the current context is CPU model.""" + return tpu_function.get_tpu_context().number_of_shards is None diff --git a/tensorflow/contrib/tpu/python/tpu/feature_column_test.py b/tensorflow/contrib/tpu/python/tpu/feature_column_test.py new file mode 100644 index 0000000000000000000000000000000000000000..75164cce4c261cc541dd6b01ee22699d286d9621 --- /dev/null +++ b/tensorflow/contrib/tpu/python/tpu/feature_column_test.py @@ -0,0 +1,286 @@ +# 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 contrib.tpu.python.tpu.feature_column.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.contrib.tpu.python.tpu import feature_column as tpu_fc +from tensorflow.python.client import session +from tensorflow.python.feature_column import feature_column as fc +from tensorflow.python.feature_column import feature_column_lib as fc_lib +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import ops +from tensorflow.python.framework import sparse_tensor +from tensorflow.python.ops import lookup_ops +from tensorflow.python.ops import parsing_ops +from tensorflow.python.ops import variables as variables_lib +from tensorflow.python.platform import test + + +def _initialized_session(): + sess = session.Session() + sess.run(variables_lib.global_variables_initializer()) + sess.run(lookup_ops.tables_initializer()) + return sess + + +class EmbeddingColumnTest(test.TestCase): + + def test_defaults(self): + categorical_column = fc_lib.categorical_column_with_identity( + key='aaa', num_buckets=3) + embedding_dimension = 2 + embedding_column = tpu_fc.embedding_column( + categorical_column, dimension=embedding_dimension) + self.assertIs(categorical_column, embedding_column.categorical_column) + self.assertEqual(embedding_dimension, embedding_column.dimension) + self.assertEqual('mean', embedding_column.combiner) + self.assertEqual('aaa_embedding', embedding_column.name) + self.assertEqual('aaa_embedding', embedding_column._var_scope_name) + self.assertEqual((embedding_dimension,), embedding_column._variable_shape) + self.assertEqual({ + 'aaa': parsing_ops.VarLenFeature(dtypes.int64) + }, embedding_column._parse_example_spec) + + def test_all_constructor_args(self): + categorical_column = fc_lib.categorical_column_with_identity( + key='aaa', num_buckets=3) + embedding_dimension = 2 + embedding_column = tpu_fc.embedding_column( + categorical_column, + dimension=embedding_dimension, + combiner='my_combiner', + initializer=lambda: 'my_initializer') + self.assertIs(categorical_column, embedding_column.categorical_column) + self.assertEqual(embedding_dimension, embedding_column.dimension) + self.assertEqual('my_combiner', embedding_column.combiner) + self.assertEqual('aaa_embedding', embedding_column.name) + self.assertEqual('aaa_embedding', embedding_column._var_scope_name) + self.assertEqual((embedding_dimension,), embedding_column._variable_shape) + self.assertEqual({ + 'aaa': parsing_ops.VarLenFeature(dtypes.int64) + }, embedding_column._parse_example_spec) + + def test_get_dense_tensor(self): + # Inputs. + vocabulary_size = 3 + sparse_input = sparse_tensor.SparseTensorValue( + # example 0, ids [2] + # example 1, ids [0, 1] + # example 2, ids [] + # example 3, ids [1] + indices=((0, 0), (1, 0), (1, 4), (3, 0)), + values=(2, 0, 1, 1), + dense_shape=(4, 5)) + + # Embedding variable. + embedding_dimension = 2 + embedding_values = ( + (1., 2.), # id 0 + (3., 5.), # id 1 + (7., 11.) # id 2 + ) + + def _initializer(shape, dtype, partition_info): + self.assertAllEqual((vocabulary_size, embedding_dimension), shape) + self.assertEqual(dtypes.float32, dtype) + self.assertIsNone(partition_info) + return embedding_values + + # Expected lookup result, using combiner='mean'. + expected_lookups = ( + # example 0, ids [2], embedding = [7, 11] + (7., 11.), + # example 1, ids [0, 1], embedding = mean([1, 2] + [3, 5]) = [2, 3.5] + (2., 3.5), + # example 2, ids [], embedding = [0, 0] + (0., 0.), + # example 3, ids [1], embedding = [3, 5] + (3., 5.), + ) + + # Build columns. + categorical_column = fc_lib.categorical_column_with_identity( + key='aaa', num_buckets=vocabulary_size) + embedding_column = tpu_fc.embedding_column( + categorical_column, + dimension=embedding_dimension, + initializer=_initializer) + + # Provide sparse input and get dense result. + embedding_lookup = embedding_column._get_dense_tensor( + fc._LazyBuilder({ + 'aaa': sparse_input + })) + + # Assert expected embedding variable and lookups. + global_vars = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES) + self.assertItemsEqual(('embedding_weights:0',), + tuple([v.name for v in global_vars])) + with _initialized_session(): + self.assertAllEqual(embedding_values, global_vars[0].eval()) + self.assertAllEqual(expected_lookups, embedding_lookup.eval()) + + +class SharedEmbeddingColumnTest(test.TestCase): + + def test_defaults(self): + categorical_column_a = fc_lib.categorical_column_with_identity( + key='aaa', num_buckets=3) + categorical_column_b = fc_lib.categorical_column_with_identity( + key='bbb', num_buckets=3) + embedding_dimension = 2 + embedding_column_b, embedding_column_a = tpu_fc.shared_embedding_columns( + [categorical_column_b, categorical_column_a], + dimension=embedding_dimension) + self.assertIs(categorical_column_a, embedding_column_a.categorical_column) + self.assertIs(categorical_column_b, embedding_column_b.categorical_column) + self.assertEqual(embedding_dimension, embedding_column_a.dimension) + self.assertEqual(embedding_dimension, embedding_column_b.dimension) + self.assertEqual('mean', embedding_column_a.combiner) + self.assertEqual('mean', embedding_column_b.combiner) + self.assertIsNotNone(embedding_column_a.initializer) + self.assertIsNotNone(embedding_column_b.initializer) + self.assertEqual('aaa_bbb_shared_embedding', + embedding_column_a.shared_embedding_collection_name) + self.assertEqual('aaa_bbb_shared_embedding', + embedding_column_b.shared_embedding_collection_name) + self.assertEqual('aaa_shared_embedding', embedding_column_a.name) + self.assertEqual('bbb_shared_embedding', embedding_column_b.name) + self.assertEqual('aaa_bbb_shared_embedding', + embedding_column_a._var_scope_name) + self.assertEqual('aaa_bbb_shared_embedding', + embedding_column_b._var_scope_name) + self.assertEqual((embedding_dimension,), embedding_column_a._variable_shape) + self.assertEqual((embedding_dimension,), embedding_column_b._variable_shape) + self.assertEqual({ + 'aaa': parsing_ops.VarLenFeature(dtypes.int64) + }, embedding_column_a._parse_example_spec) + self.assertEqual({ + 'bbb': parsing_ops.VarLenFeature(dtypes.int64) + }, embedding_column_b._parse_example_spec) + + def test_all_constructor_args(self): + categorical_column_a = fc_lib.categorical_column_with_identity( + key='aaa', num_buckets=3) + categorical_column_b = fc_lib.categorical_column_with_identity( + key='bbb', num_buckets=3) + embedding_dimension = 2 + embedding_column_a, embedding_column_b = tpu_fc.shared_embedding_columns( + [categorical_column_a, categorical_column_b], + dimension=embedding_dimension, + combiner='my_combiner', + initializer=lambda: 'my_initializer', + shared_embedding_collection_name='var_scope_name') + self.assertIs(categorical_column_a, embedding_column_a.categorical_column) + self.assertIs(categorical_column_b, embedding_column_b.categorical_column) + self.assertEqual(embedding_dimension, embedding_column_a.dimension) + self.assertEqual(embedding_dimension, embedding_column_b.dimension) + self.assertEqual('my_combiner', embedding_column_a.combiner) + self.assertEqual('my_combiner', embedding_column_b.combiner) + self.assertEqual('my_initializer', embedding_column_a.initializer()) + self.assertEqual('my_initializer', embedding_column_b.initializer()) + self.assertEqual('var_scope_name', + embedding_column_a.shared_embedding_collection_name) + self.assertEqual('var_scope_name', + embedding_column_b.shared_embedding_collection_name) + self.assertEqual('aaa_shared_embedding', embedding_column_a.name) + self.assertEqual('bbb_shared_embedding', embedding_column_b.name) + self.assertEqual('var_scope_name', embedding_column_a._var_scope_name) + self.assertEqual('var_scope_name', embedding_column_b._var_scope_name) + self.assertEqual((embedding_dimension,), embedding_column_a._variable_shape) + self.assertEqual((embedding_dimension,), embedding_column_b._variable_shape) + self.assertEqual({ + 'aaa': parsing_ops.VarLenFeature(dtypes.int64) + }, embedding_column_a._parse_example_spec) + self.assertEqual({ + 'bbb': parsing_ops.VarLenFeature(dtypes.int64) + }, embedding_column_b._parse_example_spec) + + def test_get_dense_tensor(self): + # Inputs. + vocabulary_size = 3 + # -1 values are ignored. + input_a = np.array([ + [2, -1, -1], # example 0, ids [2] + [0, 1, -1] + ]) # example 1, ids [0, 1] + input_b = np.array([ + [0, -1, -1], # example 0, ids [0] + [-1, -1, -1] + ]) # example 1, ids [] + input_features = {'aaa': input_a, 'bbb': input_b} + + # Embedding variable. + embedding_dimension = 2 + embedding_values = ( + (1., 2.), # id 0 + (3., 5.), # id 1 + (7., 11.) # id 2 + ) + + def _initializer(shape, dtype, partition_info): + self.assertAllEqual((vocabulary_size, embedding_dimension), shape) + self.assertEqual(dtypes.float32, dtype) + self.assertIsNone(partition_info) + return embedding_values + + # Expected lookup result, using combiner='mean'. + expected_lookups_a = ( + # example 0: + (7., 11.), # ids [2], embedding = [7, 11] + # example 1: + (2., 3.5), # ids [0, 1], embedding = mean([1, 2] + [3, 5]) = [2, 3.5] + ) + expected_lookups_b = ( + # example 0: + (1., 2.), # ids [0], embedding = [1, 2] + # example 1: + (0., 0.), # ids [], embedding = [0, 0] + ) + + # Build columns. + categorical_column_a = fc_lib.categorical_column_with_identity( + key='aaa', num_buckets=vocabulary_size) + categorical_column_b = fc_lib.categorical_column_with_identity( + key='bbb', num_buckets=vocabulary_size) + embedding_column_a, embedding_column_b = tpu_fc.shared_embedding_columns( + [categorical_column_a, categorical_column_b], + dimension=embedding_dimension, + initializer=_initializer) + + # Provide sparse input and get dense result. + embedding_lookup_a = embedding_column_a._get_dense_tensor( + fc._LazyBuilder(input_features)) + embedding_lookup_b = embedding_column_b._get_dense_tensor( + fc._LazyBuilder(input_features)) + + # Assert expected embedding variable and lookups. + global_vars = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES) + self.assertItemsEqual(('embedding_weights:0',), + tuple([v.name for v in global_vars])) + embedding_var = global_vars[0] + with _initialized_session(): + self.assertAllEqual(embedding_values, embedding_var.eval()) + self.assertAllEqual(expected_lookups_a, embedding_lookup_a.eval()) + self.assertAllEqual(expected_lookups_b, embedding_lookup_b.eval()) + + +if __name__ == '__main__': + test.main() diff --git a/tensorflow/contrib/tpu/python/tpu/functional.py b/tensorflow/contrib/tpu/python/tpu/functional.py new file mode 100644 index 0000000000000000000000000000000000000000..1ec9b5b33d007eb2eaa557438f32ea69053261c6 --- /dev/null +++ b/tensorflow/contrib/tpu/python/tpu/functional.py @@ -0,0 +1,25 @@ +# 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. +# ============================================================================= +"""Functional operations.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.contrib.tpu.python.tpu import gen_functional_ops + + +TPUPartitionedCall = gen_functional_ops._tpu_partitioned_call # pylint: disable=invalid-name,protected-access + diff --git a/tensorflow/contrib/tpu/python/tpu/keras_support.py b/tensorflow/contrib/tpu/python/tpu/keras_support.py index 083b65a8dacd983e3f0abede708067cc8194d1c0..37fe9af8c4b154a2e20a957f6ca5d97df3d413be 100644 --- a/tensorflow/contrib/tpu/python/tpu/keras_support.py +++ b/tensorflow/contrib/tpu/python/tpu/keras_support.py @@ -81,6 +81,7 @@ from tensorflow.python.keras import metrics as metrics_module from tensorflow.python.keras import models from tensorflow.python.keras import optimizers as keras_optimizers from tensorflow.python.keras.engine import base_layer +from tensorflow.python.keras.engine import base_layer_utils from tensorflow.python.keras.engine import training_arrays from tensorflow.python.keras.engine import training_utils from tensorflow.python.keras.layers import embeddings @@ -97,14 +98,25 @@ from tensorflow.python.platform import tf_logging as logging # TODO(b/114775106): temporary shim to optionally initialize the TPU # This increases the odds our session is initialized, but shouldn't be needed. +_TEST_REWRITE_OP = None + + def _maybe_initialize_tpu(session): """Initialize the TPU if it has not already been initialized.""" + global _TEST_REWRITE_OP try: + # Try to use cached version to avoid another ground of graph optimization. + test_rewrite_op = _TEST_REWRITE_OP + if (test_rewrite_op is None or + test_rewrite_op[0].graph != ops.get_default_graph()): + + def test_op(): + return constant_op.constant(1) + constant_op.constant(1) - def test_op(): - return constant_op.constant(1) + constant_op.constant(1) + test_rewrite_op = tpu.rewrite(test_op) + _TEST_REWRITE_OP = test_rewrite_op - session.run(tpu.rewrite(test_op)) + session.run(test_rewrite_op) except errors.FailedPreconditionError as _: session.run(tpu.initialize_system()) @@ -121,7 +133,7 @@ def _tpu_session_context(): An error occurred connecting or initializing your TPU. The session has been reset. re-run keras_to_tpu_model to create a new session. -""" + e) +""" + str(e)) def setup_tpu_session(cluster_resolver): @@ -427,7 +439,7 @@ class TPURewriteContext(object): self._default_placeholder = array_ops.placeholder self._default_name_scope = ops.name_scope - self._default_make_variable = base_layer.make_variable + self._default_make_variable = base_layer_utils.make_variable self._default_random_normal = random_ops.random_normal self._default_qr = gen_linalg_ops.qr @@ -475,14 +487,14 @@ class TPURewriteContext(object): gen_linalg_ops.qr = qr ops.name_scope = _name_scope - base_layer.make_variable = variable_scope.get_variable + base_layer_utils.make_variable = variable_scope.get_variable logging.info('Overriding default placeholder.') return def __exit__(self, exc_type, exc_val, exc_tb): array_ops.placeholder = self._default_placeholder ops.name_scope = self._default_name_scope - base_layer.make_variable = self._default_make_variable + base_layer_utils.make_variable = self._default_make_variable random_ops.random_normal = self._default_random_normal gen_linalg_ops.qr = self._default_qr @@ -532,6 +544,7 @@ class TPUInfeedInstance(object): pass +@six.add_metaclass(abc.ABCMeta) class TPUInfeedManager(object): """TPUInfeedManager manages the data infeeding of data to a TPU computation. @@ -716,7 +729,7 @@ class TPUDatasetInfeedManager(TPUInfeedManager): dummy_x_shape[0] *= tpu_assignment.num_towers dummy_y_shape = dataset.output_shapes[1].as_list() dummy_y_shape[0] *= tpu_assignment.num_towers - self._iterator = dataset.make_initializable_iterator() + self._iterator = dataset_ops.make_initializable_iterator(dataset) K.get_session().run(self._iterator.initializer) self._get_next_ops = [] @@ -757,7 +770,7 @@ class TPUDatasetInfeedManager(TPUInfeedManager): def _verify_dataset_shape(self, dataset): """Verifies a dataset is of an appropriate shape for TPUs.""" - if not isinstance(dataset, dataset_ops.Dataset): + if not isinstance(dataset, dataset_ops.DatasetV2): raise ValueError('The function passed as the `x` parameter did not ' 'return a `tf.data.Dataset`.') if not isinstance(dataset.output_classes, tuple): @@ -966,7 +979,7 @@ class TPUFunction(object): # When running on more than one core, concatenate outputs at the end # of processing. In backprop stage, the gradients will be - # calculdated according to the local inputs as gradient of + # calculated according to the local inputs as gradient of # cross-replica-concat being zero for any outputs other than those # from mlocal core so the loss calculation is identical. num_towers = self.model._tpu_assignment.num_towers @@ -993,14 +1006,17 @@ class TPUFunction(object): for tensor in tpu_targets ] - if is_training or is_test: + if is_training or is_test: + with variable_scope.variable_scope( + 'metrics', reuse=variable_scope.AUTO_REUSE): self._cloned_model.compile( optimizer=_replicated_optimizer(self._cloned_optimizer), loss=self.model.loss, loss_weights=self.model.loss_weights, - metrics=metrics_module.clone_metrics(self.model.metrics), + metrics=metrics_module.clone_metrics( + self.model._compile_metrics), weighted_metrics=metrics_module.clone_metrics( - self.model.weighted_metrics), + self.model._compile_weighted_metrics), target_tensors=tpu_targets, ) @@ -1012,29 +1028,29 @@ class TPUFunction(object): # the Momentum optimizer) when _make_train_function is invoked. with keras_tpu_variables.replicated_variable_for_optimizer( self._tpu_assignment.num_towers): - self._cloned_model._make_train_function() + self._cloned_model._make_fit_function() else: - self._cloned_model._make_train_function() + self._cloned_model._make_fit_function() self._outfeed_spec = [ tensor_spec.TensorSpec(tensor.shape, tensor.dtype, tensor.name) - for tensor in self._cloned_model.train_function.outputs + for tensor in self._cloned_model._fit_function.outputs ] return [ - self._cloned_model.train_function.updates_op, + self._cloned_model._fit_function.updates_op, tpu_ops.outfeed_enqueue_tuple( - self._cloned_model.train_function.outputs, + self._cloned_model._fit_function.outputs, name='outfeed-enqueue-train') ] elif is_test: - self._cloned_model._make_test_function() + self._cloned_model._make_eval_function() self._outfeed_spec = [ tensor_spec.TensorSpec(tensor.shape, tensor.dtype, tensor.name) - for tensor in self._cloned_model.test_function.outputs + for tensor in self._cloned_model._eval_function.outputs ] return [ tpu_ops.outfeed_enqueue_tuple( - self._cloned_model.test_function.outputs, + self._cloned_model._eval_function.outputs, name='outfeed-enqueue-test') ] elif is_predict: @@ -1060,7 +1076,7 @@ class TPUFunction(object): # `execute op` replicates `_model_fn` `num_replicas` times, with each shard # running on a different logical core. compile_op, execute_op = tpu.split_compile_and_replicate( - _model_fn, inputs=[[]] * self._tpu_assignment.num_towers) + _model_fn, inputs=[[] for _ in range(self._tpu_assignment.num_towers)]) # Generate CPU side operations to enqueue features/labels and dequeue # outputs from the model call. @@ -1170,13 +1186,9 @@ class TPUFunction(object): # pipelined loop. return None, None - if (self.model.uses_learning_phase and - not isinstance(K.learning_phase(), int)): + if isinstance(inputs[-1], int): # Remove the learning_phase flag at the end. We currently hard code the # learning_phase in TPUFunction. - assert isinstance(inputs[-1], int), ( - 'Expect the final element be learning_phase flag. Got {}'.format( - inputs[-1])) inputs = inputs[:-1] if (self.execution_mode == model_fn_lib.ModeKeys.TRAIN or @@ -1205,7 +1217,7 @@ class TPUFunction(object): """ # TODO(xiejw): Decide how to reduce outputs, or discard all but first. if self.execution_mode == model_fn_lib.ModeKeys.PREDICT: - outputs = [[]] * len(self._outfeed_spec) + outputs = [[] for _ in range(len(self._outfeed_spec))] outputs_per_replica = len(self._outfeed_spec) for i in range(self._tpu_assignment.num_towers): @@ -1361,9 +1373,16 @@ class KerasTPUModel(models.Model): # not hashable. self._numpy_to_infeed_manager_list = [] + # Add distribution specific arguments since we don't call the Model init. + self._distribution_strategy = None + self._compile_distribution = None + self.predict_function = None self.test_function = None self.train_function = None + self._fit_function = None + self._eval_function = None + self._stateful_metric_functions = [] cluster_resolver = strategy._tpu_cluster_resolver self._tpu_name_or_address = cluster_resolver.get_master() @@ -1378,13 +1397,22 @@ class KerasTPUModel(models.Model): self.compile( self._cpu_model.optimizer, self._cpu_model.loss, - self._cpu_model.metrics, + self._cpu_model._compile_metrics, self._cpu_model.loss_weights, self._cpu_model.sample_weight_mode, - self._cpu_model.weighted_metrics, + self._cpu_model._compile_weighted_metrics, self._cpu_model.target_tensors, ) + # This flag must be disabled upon model mutation, such as changing the model + # layers or recompiling the model to use a different optimizer. New function + # definitions are generated whenever this flag is disabled, ensuring that + # internal graph functions are always using the current model structure. + # + # Requires declaration here because this constructor skips the + # Model constructor. + self._built_graph_functions = False + def get_config(self): return { 'cpu_model': self._cpu_model, @@ -1442,7 +1470,7 @@ class KerasTPUModel(models.Model): assert not self._numpy_to_infeed_manager_list # Ensure empty. infeed_managers = [] # Managers to clean up at the end of the fit call. - if isinstance(x, dataset_ops.Dataset): + if isinstance(x, dataset_ops.DatasetV2): # TODO(b/111413240): Support taking a tf.data.Dataset directly. raise ValueError( 'Taking a Dataset directly is not yet supported. Please ' @@ -1468,7 +1496,7 @@ class KerasTPUModel(models.Model): y = infeed_manager.dummy_y infeed_managers.append((x, infeed_manager)) - if isinstance(validation_data, dataset_ops.Dataset): + if isinstance(validation_data, dataset_ops.DatasetV2): # TODO(b/111413240): Support taking a tf.data.Dataset directly. raise ValueError( 'Taking a Dataset directly is not yet supported. Please ' @@ -1516,11 +1544,18 @@ class KerasTPUModel(models.Model): verbose=1, sample_weight=None, steps=None): - assert not self._numpy_to_infeed_manager_list # Ensure empty. + original_numpy_to_infeed_manager_list = [] + if self._numpy_to_infeed_manager_list: + # evaluate call may be executed as callbacks during the training. In this + # case, _numpy_to_infeed_manager_list is not empty, so save it for + # recovery at the end of evaluate call. + original_numpy_to_infeed_manager_list = self._numpy_to_infeed_manager_list + self._numpy_to_infeed_manager_list = [] with _tpu_session_context(): - infeed_managers = [] # Managers to clean up at the end of the fit call. - if isinstance(x, dataset_ops.Dataset): + # Managers to clean up at the end of the evaluate call. + infeed_managers = [] + if isinstance(x, dataset_ops.DatasetV2): # TODO(b/111413240): Support taking a tf.data.Dataset directly. raise ValueError( 'Taking a Dataset directly is not yet supported. Please ' @@ -1549,7 +1584,8 @@ class KerasTPUModel(models.Model): return super(KerasTPUModel, self).evaluate(x, y, batch_size, verbose, sample_weight, steps) finally: - self._numpy_to_infeed_manager_list = [] + self._numpy_to_infeed_manager_list = ( + original_numpy_to_infeed_manager_list) def _pipeline_fit(self, x, y, batch_size, epochs, verbose, callbacks, validation_split, validation_data, shuffle, class_weight, @@ -1618,7 +1654,7 @@ class KerasTPUModel(models.Model): self._make_train_function() sample_weights = sample_weights or [] val_sample_weights = val_sample_weights or [] - if self.uses_learning_phase and not isinstance(K.learning_phase(), int): + if not isinstance(K.learning_phase(), int): ins = inputs + targets + sample_weights + [1] else: ins = inputs + targets + sample_weights @@ -1644,14 +1680,10 @@ class KerasTPUModel(models.Model): callbacks, self, do_validation=do_validation, - val_inputs=val_inputs, - val_targets=val_targets, - val_sample_weights=val_sample_weights, batch_size=batch_size, epochs=epochs, steps_per_epoch=steps_per_epoch, samples=num_training_samples, - validation_steps=validation_steps, verbose=verbose, count_mode=count_mode) @@ -1668,7 +1700,7 @@ class KerasTPUModel(models.Model): callbacks.on_train_begin() for epoch in range(initial_epoch, epochs): # Reset stateful metrics - for m in self.stateful_metric_functions: + for m in self.metrics: m.reset_states() # Update callbacks callbacks.on_epoch_begin(epoch) @@ -1891,7 +1923,7 @@ class KerasTPUModel(models.Model): if validation_data: if (isinstance(validation_data, iterator_ops.Iterator) or isinstance(validation_data, iterator_ops.EagerIterator) or - isinstance(validation_data, dataset_ops.Dataset)): + isinstance(validation_data, dataset_ops.DatasetV2)): raise ValueError('KerasTPUModel cannot handle a Dataset or Iterator ' 'for validation_data. Please instead pass a function ' 'that returns a `tf.data.Dataset`.') @@ -1965,10 +1997,21 @@ class KerasTPUModel(models.Model): def optimizer(self, optimizer): self._optimizer = optimizer + @property + def metrics(self): + if self._tpu_model: + return self._tpu_model.metrics + return self._stateful_metric_functions + + @metrics.setter + def metrics(self, metrics): + self._stateful_metric_functions = metrics + def _make_train_function(self): if not self.train_function: self.train_function = TPUFunction( - self, model_fn_lib.ModeKeys.TRAIN, + self, + model_fn_lib.ModeKeys.TRAIN, tpu_assignment=self._tpu_assignment) return self.train_function @@ -1979,6 +2022,21 @@ class KerasTPUModel(models.Model): self, model_fn_lib.ModeKeys.EVAL, tpu_assignment=self._tpu_assignment) return self.test_function + def _make_fit_function(self): + if not self._fit_function: + self._fit_function = TPUFunction( + self, + model_fn_lib.ModeKeys.TRAIN, + tpu_assignment=self._tpu_assignment) + + return self._fit_function + + def _make_eval_function(self): + if not self._eval_function: + self._eval_function = TPUFunction( + self, model_fn_lib.ModeKeys.EVAL, tpu_assignment=self._tpu_assignment) + return self._eval_function + def _make_predict_function(self): if not self.predict_function: self.predict_function = TPUFunction( @@ -2015,6 +2073,8 @@ class KerasTPUModel(models.Model): # tpu_model may not be compiled, e.g., loading weights and then predict. return for k, v in six.iteritems(cpu_optimizer_config): + if k == 'name': + continue opt_var = getattr(self._tpu_model.optimizer, k) if isinstance(opt_var, variables.Variable): logging.info('CPU -> TPU %s: %s {%s}', k, v, K.get_value(opt_var)) @@ -2043,6 +2103,8 @@ class KerasTPUModel(models.Model): self._cpu_model.set_weights(tpu_weights) for k, v in six.iteritems(tpu_optimizer_config): logging.info('TPU -> CPU %s: %s', k, v) + if k == 'name': + continue opt_var = getattr(self.cpu_optimizer, k) if isinstance(opt_var, variables.Variable): K.get_session().run(opt_var.assign(v)) @@ -2172,10 +2234,10 @@ def tpu_model(model, strategy=None): cpu_model.compile( _clone_optimizer(model.optimizer, optimizer_config), model.loss, - metrics_module.clone_metrics(model.metrics), + metrics_module.clone_metrics(model._compile_metrics), model.loss_weights, model.sample_weight_mode, - metrics_module.clone_metrics(model.weighted_metrics), + metrics_module.clone_metrics(model._compile_weighted_metrics), ) if model_weights: diff --git a/tensorflow/contrib/tpu/python/tpu/keras_tpu_variables.py b/tensorflow/contrib/tpu/python/tpu/keras_tpu_variables.py index 28d3a938510a450ccba0d921663d848e2adec72f..de425626c813784ef657d17eac0c7bb77599a155 100644 --- a/tensorflow/contrib/tpu/python/tpu/keras_tpu_variables.py +++ b/tensorflow/contrib/tpu/python/tpu/keras_tpu_variables.py @@ -69,6 +69,7 @@ class ReplicatedVariable(object): def __init__(self, name, variables): self._name = name self._primary_var = variables[0] + self._common_name = self._primary_var.name.split(":")[0] self._vars = variables self._cached_value = None self._dtype = variables[0].dtype @@ -217,6 +218,10 @@ class ReplicatedVariable(object): def get(self): return self._primary_var + @property + def _in_graph_mode(self): + return self._primary_var._in_graph_mode # pylint: disable=protected-access + def _should_act_as_resource_variable(self): """Pass resource_variable_ops.is_resource_variable check.""" pass diff --git a/tensorflow/contrib/tpu/python/tpu/session_support.py b/tensorflow/contrib/tpu/python/tpu/session_support.py index a95275487899c4770ef99b620a7671eec2bb81eb..f5735cecc38b7033f21fc4d4105cfead233379fa 100644 --- a/tensorflow/contrib/tpu/python/tpu/session_support.py +++ b/tensorflow/contrib/tpu/python/tpu/session_support.py @@ -43,12 +43,19 @@ class CoordinatorShutdownException(Exception): pass +def _clone_session(session, graph=None): + return session_lib.Session( + target=session.sess_str, + config=session._config, # pylint: disable=protected-access + graph=graph if graph else session.graph) + + def _make_heartbeat_op(session, device, request_ph): """Return a heartbeat op or None if heartbeats are not supported by device.""" try: # Test if we can connect in a isolated graph + session with ops.Graph().as_default(): - with session_lib.Session(target=session.sess_str) as temp_session: + with _clone_session(session) as temp_session: with ops.device(device): heartbeat_op = tpu_ops.worker_heartbeat('') options = config_pb2.RunOptions(timeout_in_ms=5000) @@ -178,7 +185,8 @@ def all_worker_devices(session): """Return a list of devices for each worker in the system.""" devices = session.list_devices() return [ - device.name for device in devices + device.name + for device in devices if ':CPU:' in device.name and 'coordinator' not in device.name ] @@ -220,6 +228,7 @@ class WatchdogManager(threading.Thread): self.ping_interval = ping_interval self.shutdown_timeout = shutdown_timeout self.daemon = True + self._config = session._config # pylint: disable=protected-access self._target = session.sess_str self._running = False self._devices = devices @@ -234,6 +243,7 @@ class WatchdogManager(threading.Thread): self._session = session_lib.Session( target=self._target, graph=self._graph, + config=self._config, ) if self._devices is None: @@ -246,12 +256,14 @@ class WatchdogManager(threading.Thread): self._worker_manager.configure( event_pb2.WorkerHeartbeatRequest( watchdog_config=event_pb2.WatchdogConfig( - timeout_ms=self.shutdown_timeout * 1000,))) + timeout_ms=self.shutdown_timeout * 1000,), + shutdown_mode=event_pb2.WAIT_FOR_COORDINATOR)) def configure_and_run(self): - logging.info('Enabling watchdog timer with %d second timeout ' - 'and %d second ping interval.', - self.shutdown_timeout, self.ping_interval) + logging.info( + 'Enabling watchdog timer with %d second timeout ' + 'and %d second ping interval.', self.shutdown_timeout, + self.ping_interval) self._reset_manager() self._running = True self.start() @@ -260,7 +272,8 @@ class WatchdogManager(threading.Thread): logging.info('Stopping worker watchdog.') self._worker_manager.configure( event_pb2.WorkerHeartbeatRequest( - watchdog_config=event_pb2.WatchdogConfig(timeout_ms=-1,))) + watchdog_config=event_pb2.WatchdogConfig(timeout_ms=-1,), + shutdown_mode=event_pb2.NOT_CONFIGURED)) self._running = False self.join() @@ -334,8 +347,7 @@ class GracefulShutdownHook(session_run_hook.SessionRunHook): with self._graph.as_default(): logging.info('Installing graceful shutdown hook.') - self._session = session_lib.Session( - target=training_session.sess_str, graph=self._graph) + self._session = _clone_session(training_session, self._graph) self._workers = WorkerHeartbeatManager.from_devices( self._session, all_worker_devices(self._session)) self._heartbeat_supported = self._workers.num_workers() > 0 diff --git a/tensorflow/contrib/tpu/python/tpu/tensor_tracer.py b/tensorflow/contrib/tpu/python/tpu/tensor_tracer.py new file mode 100644 index 0000000000000000000000000000000000000000..bf492e78a15acc92017663a286e8c8f0b2045339 --- /dev/null +++ b/tensorflow/contrib/tpu/python/tpu/tensor_tracer.py @@ -0,0 +1,1147 @@ +# 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. +# ======================================================================== +"""A utility to trace tensor values on TPU.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import os +import os.path +import re +import sys + +from tensorflow.contrib.tpu.python.ops import tpu_ops +from tensorflow.contrib.tpu.python.tpu import tpu +from tensorflow.python.framework import constant_op +from tensorflow.python.framework import dtypes +from tensorflow.python.framework import ops +from tensorflow.python.framework import tensor_util +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import control_flow_util +from tensorflow.python.ops import gen_math_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.platform import gfile +from tensorflow.python.platform import tf_logging as logging + +_TRACER_LOG_PREFIX = ' [>>>TT>>>]' +_DEVICE_TYPE_TPU = 'tpu' +_DEVICE_TYPE_CPU = 'cpu' +_TRACE_MODE_NAN_INF = 'nan-inf' +_TRACE_MODE_PART_TENSOR = 'part-tensor' +_TRACE_MODE_PART_TENSOR_SIZE = 3 +_TRACE_MODE_FULL_TENSOR = 'full-tensor' +_TRACE_MODE_NORM = 'norm' +_TRACE_MODE_MAX_ABS = 'max-abs' +_SUBMODE_BRIEF = 'brief' +_SUBMODE_DETAILED = 'detailed' +_REASON_OUTSIDE_OP_RANGE = 'not-traced-outside-op-range' +_REASON_UNSAFE_OP = 'not-traced-unsafe-op' +_REASON_UNSAFE_SCALAR = 'not-traced-unsafe-scalar' +_REASON_LESS_INTERESTING_OP = 'not-traced-less-interesting-op' +_REASON_DEVICE_MISMATCH = 'not-traced-device-mismatch' +_REASON_DYNAMIC_SHAPE = 'not-traced-dynamic-shape' +_REASON_SCALAR_GET_TRACED = 'traced-scalar' +_REASON_TENSOR_GET_TRACED = 'traced-tensor' +_REASON_USER_INCLUDED = 'traced-user-included' +_REASON_USER_EXCLUDED = 'not-traced-user-excluded' +_REASON_NOT_EXECUTED = 'not-traced-not-in-exec-path' +_REASON_NON_NUMERIC_TENSOR = 'not-traced-non-numeric-tensor' +_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_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_OPS = 'number-of-ops:' +_FIELD_NAME_NUM_TENSORS = 'number-of-tensors:' +_FIELD_NAME_TOPOLOGICAL_SORT_SUCCEED = 'topological-sort-succeed:' +_FLAGS_ENV_VAR = 'TENSOR_TRACER_FLAGS' +_FLAG_SINGLE_QUOTE_PAT = re.compile(r"\s*--([^=]+)='([^']*)'") +_FLAG_DOUBLE_QUOTE_PAT = re.compile(r'\s*--([^=]+)="([^"]*)"') +_FLAG_NO_QUOTE_PAT = re.compile(r'\s*--([^=]+)=(\S*)') +_FLAG_NO_EQUAL_PAT = re.compile(r'\s*--([^=]+)\s*') +_FLAG_NAME_ENABLE = 'enable' +_FLAG_NAME_TRACE_MODE = 'trace_mode' +_FLAG_NAME_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_USE_TEST_UNDECLARED_OUTPUTS_DIR = 'use_test_undeclared_outputs_dir' +_FLAG_NAME_OP_RANGE = 'op_range' +_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' + + +def tensor_tracepoint(tensor, checkpoint_name): + """Adds a checkpoint with the given checkpoint name for the given tensor. + + The tensor will be added to the list of tensors that will be traced by the + tensor tracer. + + Args: + tensor: the tensor object for which the tracing is requested. + checkpoint_name: a string name for the checkpoint. This name has to be a + unique name if used within model comparison. The tensors that have the same + checkpoint identifier is compared in model comparison. + Returns: + The provided tensor. + """ + + tensor.graph.get_collection(_TENSOR_TRACER_COLLECTION) + tensor.graph.add_to_collection(_TENSOR_TRACER_COLLECTION, + (tensor, checkpoint_name)) + return tensor + + +def keras_layer_tracepoint(layer, checkpoint_name): + """An interface for adding the tensor outputs of a keras layer. + + Encapsulates tensor_tracepoint. + + Args: + layer: A keras layer. + checkpoint_name: a string name for the checkpoint. This name has to be a + unique name if used within model comparison. The tensors that have the same + checkpoint identifier is compared in model comparison. + + Returns: + The provided layer. + """ + try: + outputs = layer.output + if tensor_util.is_tensor(outputs): + tensor_tracepoint(outputs, '%s' % (checkpoint_name)) + else: + idx = 0 + for output_tensor in outputs: + if tensor_util.is_tensor(outputs): + tensor_tracepoint(output_tensor, '%s_%d' % (checkpoint_name, idx)) + idx += 1 + except AttributeError: + pass + except RuntimeError: + pass + return layer + + +class TensorTracer(object): + """A software construct for tracing tensor values in a TF graph on TPU. + + This utility is disabled by default. It can be enabled by setting + the TENSOR_TRACER_FLAGS env variable as: + export TENSOR_TRACER_FLAGS="--enable=1" + If it is enabled, it will trace the output tensor values of + selected Ops in the graph. It has two outputs: (1) the traces and (2) + a report. The traces are dumped to a specified local file on the TPU + host. The report is printed to the log.info of the TPU job. + By passing options via the env variable, users can change: + (1) the trace mode (e.g., detecting NaN/Inf, printing partial or + full tensor values) + (2) which Ops to be traced (via op.name or op.type) + (3) output trace file path. + """ + + @staticmethod + def _match_next_flag(flags, pos): + """Returns the match for the next TensorTracer flag. + + Args: + flags: a string that contains the flags. + pos: where in flags to start the search. + + Returns: + A pair where the first element is the regular-expression + match found and the second element indicates if the match + has a value. + """ + + match = _FLAG_DOUBLE_QUOTE_PAT.match(flags, pos) + if match: + return match, True + match = _FLAG_SINGLE_QUOTE_PAT.match(flags, pos) + if match: + return match, True + match = _FLAG_NO_QUOTE_PAT.match(flags, pos) + if match: + return match, True + match = _FLAG_NO_EQUAL_PAT.match(flags, pos) + if match: + # The flag is found but is not given a value. + return match, False + # The flag is not found. + return None, False + + @staticmethod + def validate_flag_names(): + """Validates if the TensorTrace flags passed are valid.""" + valid_flag_names = [_FLAG_NAME_ENABLE, _FLAG_NAME_TRACE_MODE, + _FLAG_NAME_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_USE_TEST_UNDECLARED_OUTPUTS_DIR, + _FLAG_NAME_INCLUDE_LESS_INTERESTING_OPS, + _FLAG_NAME_OP_RANGE] + tensor_tracer_flags = os.environ.get(_FLAGS_ENV_VAR) + if not tensor_tracer_flags: + return + pos = 0 + while True: + match, _ = TensorTracer._match_next_flag(tensor_tracer_flags, pos) + if not match: + break + flag_name = match.group(1) + if flag_name not in valid_flag_names: + raise ValueError( + 'The flag name "%s" passed via the environment variable "%s" ' + 'is invalid. Valid flag names are:' + '\n%s'%(flag_name, _FLAGS_ENV_VAR, valid_flag_names)) + pos = match.end() + + @staticmethod + def print_flag_values(): + """Prints all TensorTracer flags passed via environment variables.""" + + tensor_tracer_flags = os.environ.get(_FLAGS_ENV_VAR) + if not tensor_tracer_flags: + return 'Env variable "%s" is not set'%_FLAGS_ENV_VAR + result = 'Env variable "%s" is set to "%s"\n'%(_FLAGS_ENV_VAR, + tensor_tracer_flags) + result += 'Individual flag value:\n' + pos = 0 + while True: + match, has_value = TensorTracer._match_next_flag( + tensor_tracer_flags, pos) + if not match: + break + flag_name = match.group(1) + if has_value: + flag_value = match.group(2) + else: + flag_value = None + result += ' %s: %s\n'%(flag_name, flag_value) + pos = match.end() + result += '\n' + return result + + @staticmethod + def get_flag_value(wanted_flag_name): + """Returns the value of a TensorTracer flags. + + Args: + wanted_flag_name: the name the the flag we are looking for. + + Returns: + A pair where the first element indicates if the flag is + found and the second element is the value of the flag. + + Raises: + RuntimeError: If supposedly deadcode is reached. + """ + + tensor_tracer_flags = os.getenv(_FLAGS_ENV_VAR) + if not tensor_tracer_flags: + return False, None + pos = 0 + while True: + match, has_value = TensorTracer._match_next_flag( + tensor_tracer_flags, pos) + if not match: + return False, None + flag_name = match.group(1) + if has_value: + flag_value = match.group(2) + else: + flag_value = None + if flag_name == wanted_flag_name: + return True, flag_value + pos = match.end() + raise RuntimeError('Should not reach here.') + + @staticmethod + def flag_value_to_re_list(flag_name): + """Converts list of strings to compiled RE.""" + + re_list = [] + found, flag_value = TensorTracer.get_flag_value(flag_name) + if not found or not flag_value: + return re_list + list_of_values = flag_value.split() + for v in list_of_values: + r = re.compile(v) + re_list.append(r) + return re_list + + @staticmethod + def _is_flag_on(flag_name): + """Returns True if the given flag is on.""" + + found, flag_value = TensorTracer.get_flag_value(flag_name) + if not found: + return False + if flag_value is None: + return True + # Depends on the flag value. + flag_value = flag_value.lower() + enabled = flag_value in ['1', 't', 'true', 'y', 'yes'] + return enabled + + @staticmethod + def is_enabled(): + """Returns True if TensorTracer is enabled.""" + + return TensorTracer._is_flag_on(_FLAG_NAME_ENABLE) + + @staticmethod + def use_test_undeclared_outputs_dir(): + """Decides the output directory of the report and trace files. + + Args: + None. + + Returns: + True if the output files should be written to the + test-undeclared-outputs-directory defined via an + env variable. + """ + + return TensorTracer._is_flag_on( + _FLAG_NAME_USE_TEST_UNDECLARED_OUTPUTS_DIR) + + + @staticmethod + def check_device_type(device_type): + """Checks if the given device type is valid.""" + + if device_type not in [_DEVICE_TYPE_TPU, _DEVICE_TYPE_CPU]: + raise ValueError('Invalid device_type "%s"'%device_type) + + @staticmethod + def check_trace_mode(trace_mode): + """Checks if the given trace mode is valid.""" + + valid_trace_modes = [_TRACE_MODE_NAN_INF, _TRACE_MODE_PART_TENSOR, + _TRACE_MODE_FULL_TENSOR, _TRACE_MODE_NORM, + _TRACE_MODE_MAX_ABS] + if trace_mode not in valid_trace_modes: + raise ValueError('Invalid trace mode "%s" given to the Tensor_Tracer.' + 'Valid trace modes are: %s'%(trace_mode, + valid_trace_modes)) + + @staticmethod + def check_submode(submode): + """Checks if the given submode is valid.""" + + if not submode: + return + valid_submodes = [_SUBMODE_DETAILED, _SUBMODE_BRIEF] + if submode not in valid_submodes: + raise ValueError('Invalid submode "%s" given to the Tensor_Tracer.' + 'Valid submodes are: %s'%(submode, + valid_submodes)) + + @staticmethod + def unsafe_op(op): + """Returns True if this op is not safe to be traced.""" + + if control_flow_util.IsInCond(op): + return True + # Reasons for not including following op types: + # Assign: cause incorrect result with CPU tracing. + if op.type in ['Assign']: + return True + return False + + @staticmethod + def device_mismatch(device_type, op): + if device_type == _DEVICE_TYPE_TPU: + # pylint: disable=protected-access + return tpu._TPU_REPLICATE_ATTR not in op.node_def.attr + # pylint: enable=protected-access + return False + + @staticmethod + def unsafe_scalar_trace(op): + """Return true if scalar output tensor from Op is not safe to be traced.""" + + # Tracing the following causes cycle in the graph on TPU. + if op.type in ['LoopCond', 'Enter', 'Merge', 'Const', + 'Switch', 'Less', 'ReadVariableOp']: + return True + # Tracing the following will cause casting-issue + # with the norm tracing mode or other compilation issues on CPU. + if op.type in ['VarHandleOp', 'IteratorToStringHandle', + 'IteratorGetNext', 'OneShotIterator', + 'IteratorV2', 'MakeIterator', + 'BatchDatasetV2', 'MapDataset', + 'FixedLengthRecordDataset', 'TakeDataset', 'ZipDataset', + 'Placeholder', 'PlaceholderWithDefault', 'StridedSlice']: + return True + return False + + @staticmethod + def less_interesting_op(op): + """Returns True if the given Op is not an interesting one to be traced.""" + + found, _ = TensorTracer.get_flag_value( + _FLAG_NAME_INCLUDE_LESS_INTERESTING_OPS) + if found: + # users force to include all ops. + return False + # Following ops are highly unlikey to cause bugs. + return op.type in ['Const', 'Identity', 'Cast', 'Shape'] + + @staticmethod + def reason(op_idx, details): + """Returns reason why the Op at op_idx is traced or not.""" + + return '%d %s'%(op_idx, details) + + @staticmethod + def topological_sort(g): + """Performs topological sort on the given graph. + + Args: + g: the graph. + + Returns: + A pair where the first element indicates if the topological + sort succeeded (True if there is no cycle found; False if a + cycle is found) and the second element is either the sorted + list of nodes or the cycle of nodes found. + """ + + def visit(op, cycle, permanently_marked_ops, + temporarily_marked_ops, sorted_ops): + """Recursively visits all Ops in a graph. + + Args: + op: the current Op being visited. + cycle: a cycle of Ops found. + permanently_marked_ops: the set of Ops that were already visited. + temporarily_marked_ops: the set of Ops that we have visited during + the current descent. + sorted_ops: the list of Ops sorted in topological order. + """ + + if cycle: + return + if op in permanently_marked_ops: + return + if op in temporarily_marked_ops: + cycle = temporarily_marked_ops + return + temporarily_marked_ops.add(op) + for i in range(len(op.outputs)): + out_tensor = op.outputs[i] + for consumer_op in out_tensor.consumers(): + visit(consumer_op, cycle, permanently_marked_ops, + temporarily_marked_ops, sorted_ops) + # pylint: disable=protected-access + for ctrl_output_op in op._control_outputs: + # pylint: enable=protected-access + visit(ctrl_output_op, cycle, permanently_marked_ops, + temporarily_marked_ops, sorted_ops) + temporarily_marked_ops.remove(op) + permanently_marked_ops.add(op) + sorted_ops.insert(0, op) + + graph_cycle = set([]) + sorted_ops = [] + permanently_marked_ops = set([]) + temporarily_marked_ops = set([]) + unsorted_ops = g.get_operations() + for op in unsorted_ops: + visit(op, graph_cycle, permanently_marked_ops, + temporarily_marked_ops, sorted_ops) + if graph_cycle: + return (False, graph_cycle) + else: + assert len(unsorted_ops) == len(sorted_ops) + return (True, sorted_ops) + + @staticmethod + def _make_op_and_tensor_maps(op_list): + """Creates various maps and lists from op_list. + + Args: + op_list: a list of Ops + + Returns: + opname_idx_map: a map from Op's name to its index in op_list. + tensor_list: a list of output tensors of the Ops in op_list. + tensorname_idx_map: a map from output tensor name to its index + in tensor_list. + """ + + opname_idx_map = {} + tensor_list = [] + tensorname_idx_map = {} + for op_id, op in enumerate(op_list): + if op.name in opname_idx_map: + raise ValueError('Duplicated Op name: %s'%op.name) + opname_idx_map[op.name] = op_id + for output_tensor in op.outputs: + if output_tensor.name not in tensorname_idx_map: + tensor_list.append(output_tensor) + tensorname_idx_map[output_tensor.name] = len(tensor_list)-1 + return (opname_idx_map, tensor_list, tensorname_idx_map) + + def __init__(self): + """Initializes a TensorTracer. + + Sets the various member fields from the flags (if given) or the defaults. + """ + self._version = 'use-outside-compilation' + self._device_type = None + TensorTracer.validate_flag_names() + found, self._trace_mode = TensorTracer.get_flag_value(_FLAG_NAME_TRACE_MODE) + if not found or not self._trace_mode: + self._trace_mode = _TRACE_MODE_NAN_INF + TensorTracer.check_trace_mode(self._trace_mode) + found, self._submode = TensorTracer.get_flag_value(_FLAG_NAME_SUBMODE) + if not found or not self._submode: + self._submode = _SUBMODE_DETAILED + TensorTracer.check_submode(self._submode) + self._part_tensor_size = _TRACE_MODE_PART_TENSOR_SIZE + self._instrument_records = {} + self._set_trace_file_path() + self._set_report_file() + self._set_op_range() + self._set_excluded_opnames() + self._set_excluded_optypes() + self._set_included_opnames() + self._set_included_optypes() + self._num_replicas = None + self._replica_id = None + + def _add_replica_id_to_graph(self, num_replicas, result_tensor): + """Adds nodes for computing the replica ID to the graph.""" + + if not num_replicas: + 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 \ + 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) + + def _set_report_file(self): + """Sets the path of the output report file.""" + + found, self._report_file_path = TensorTracer.get_flag_value( + _FLAG_NAME_REPORT_FILE) + if found and self._report_file_path \ + and TensorTracer.use_test_undeclared_outputs_dir(): + if os.path.isabs(self._report_file_path): + raise ValueError('If use_test_undeclared_outputs_dir is set,' + 'report_file_path cannot be an absolute path (%s)' + %self._report_file_path) + outputs_dir = os.environ.get(_TEST_UNDECLARED_OUTPUTS_DIR_ENV_VAR) + self._report_file_path = os.path.join(outputs_dir, + self._report_file_path) + if not self._report_file_path: + self._report_file = None + return + try: + self._report_file = gfile.Open(self._report_file_path, 'w') + except IOError as e: + raise e + + def _close_report_file(self): + if self._report_file: + self._report_file.close() + + def _set_op_range(self): + """Sets the index range of the Ops that we will consider tracing.""" + + found, op_range = TensorTracer.get_flag_value(_FLAG_NAME_OP_RANGE) + if not found or not op_range: + self._op_range = (-1, -1) # this means including all ops. + return + match = _OP_RANGE_PAT.match(op_range) + if not match: + self._op_range = (-1, -1) # this means including all ops. + return + self._op_range = (int(match.group(1)), int(match.group(2))) + + def _inside_op_range(self, idx): + """Return True if the given index is inside the selected range.""" + + if idx < self._op_range[0]: + return False + return self._op_range[1] < 0 or idx <= self._op_range[1] + + def _set_excluded_opnames(self): + self._excluded_opname_re_list = TensorTracer.flag_value_to_re_list( + _FLAG_NAME_EXCLUDED_OPNAMES) + + def _set_excluded_optypes(self): + self._excluded_optype_re_list = TensorTracer.flag_value_to_re_list( + _FLAG_NAME_EXCLUDED_OPTYPES) + + def _set_included_opnames(self): + self._included_opname_re_list = TensorTracer.flag_value_to_re_list( + _FLAG_NAME_INCLUDED_OPNAMES) + + def _set_included_optypes(self): + self._included_optype_re_list = TensorTracer.flag_value_to_re_list( + _FLAG_NAME_INCLUDED_OPTYPES) + + def _is_user_included_op(self, op): + for opname_re in self._included_opname_re_list: + if opname_re.match(op.name): + return True + for optype_re in self._included_optype_re_list: + if optype_re.match(op.type): + return True + return False + + def _is_user_excluded_op(self, op): + for opname_re in self._excluded_opname_re_list: + if opname_re.match(op.name): + return True + for optype_re in self._excluded_optype_re_list: + if optype_re.match(op.type): + return True + return False + + def _write_report(self, content): + """Writes the given content to the report.""" + + line = '%s %s'%(_TRACER_LOG_PREFIX, content) + if self._report_file: + self._report_file.write(line) + else: + logging.info(line) + + def _write_config_section(self): + """Writes the config section of the report.""" + + self._write_report('%s %s\n'%(_MARKER_SECTION_BEGIN, _SECTION_NAME_CONFIG)) + self._write_report('%s %s\n'%(_FIELD_NAME_VERSION, self._version)) + self._write_report('%s %s\n'%(_FIELD_NAME_DEVICE, self._device_type)) + self._write_report('%s %s\n'%(_FIELD_NAME_TRACE_MODE, self._trace_mode)) + self._write_report('%s %s\n'%(_FIELD_NAME_SUBMODE, self._submode)) + self._write_report('%s %s\n'%(_FIELD_NAME_NUM_REPLICAS, self._num_replicas)) + self._write_report('%s %s\n'%(_MARKER_SECTION_END, _SECTION_NAME_CONFIG)) + + def _write_reason_section(self): + """Writes the reason section of the report.""" + + self._write_report('%s %s\n'%(_MARKER_SECTION_BEGIN, _SECTION_NAME_REASON)) + for key in sorted(self._instrument_records): + self._write_report('"%s" %s\n'%(key, self._instrument_records[key])) + self._write_report('%s %s\n'%(_MARKER_SECTION_END, _SECTION_NAME_REASON)) + + def _write_op_list_section(self, op_list): + """Writes the Op-list section of the report.""" + + self._write_report('%s %s\n'%(_MARKER_SECTION_BEGIN, _SECTION_NAME_OP_LIST)) + self._write_report('%s %d\n'%(_FIELD_NAME_NUM_OPS, len(op_list))) + for i in range(0, len(op_list)): + op = op_list[i] + line = '%d "%s" %s'%(i, op.name, op.type) + for out_tensor in op.outputs: + if out_tensor.name not in self._tensorname_idx_map: + raise ValueError( + 'out_tensor %s is not in tensorname_idx_map'%out_tensor.name) + line += ' %d'%self._tensorname_idx_map[out_tensor.name] + line += '\n' + self._write_report(line) + self._write_report('%s %s\n'%(_MARKER_SECTION_END, _SECTION_NAME_OP_LIST)) + + def _write_tensor_list_section(self, tensor_list, opname_idx_map): + """Writes the tensor-list section of the report.""" + + self._write_report('%s %s\n'%(_MARKER_SECTION_BEGIN, + _SECTION_NAME_TENSOR_LIST)) + self._write_report('%s %d\n'%(_FIELD_NAME_NUM_TENSORS, len(tensor_list))) + for i in range(0, len(tensor_list)): + tensor = tensor_list[i] + line = '%d "%s"'%(i, tensor.name) + for consumer_op in tensor.consumers(): + if consumer_op.name not in opname_idx_map: + raise ValueError( + 'consumer_op %s is not in opname_idx_map'%consumer_op.name) + line += ' %d'%opname_idx_map[consumer_op.name] + line += '\n' + self._write_report(line) + self._write_report('%s %s\n'%(_MARKER_SECTION_END, + _SECTION_NAME_TENSOR_LIST)) + + def _write_graph_section(self, succeed, sorted_or_cycle): + """Writes the graph section of the report.""" + + self._write_report('%s %s\n'%(_MARKER_SECTION_BEGIN, _SECTION_NAME_GRAPH)) + self._write_report('%s %s\n'%(_FIELD_NAME_TOPOLOGICAL_SORT_SUCCEED, + succeed)) + l = list(sorted_or_cycle) + for i in range(0, len(l)): + self._write_report('%d "%s"\n'%(i, l[i].name)) + self._write_report('%s %s\n'%(_MARKER_SECTION_END, _SECTION_NAME_GRAPH)) + + def _preprocess_traced_tensor(self, tensor): + """Computes NAN/Norm/Max on TPUs before sending to CPU. + + Args: + tensor: The tensor to be traced. + Returns: + A tensor that should be input to the trace_function. + Raises: + RuntimeError: If the trace mode is invalid. + """ + + def _detect_nan_inf(tensor): + """Trace function for detecting any NaN/Inf in the tensor.""" + + if tensor.dtype.is_floating: + 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) + # The shape has to be 1. Set it if it does not have the information. + output_tensor = array_ops.reshape(output_tensor, [1]) + return output_tensor + + def _show_norm(tensor): + tensor = math_ops.cast(tensor, dtypes.float32) + output_tensor = linalg_ops.norm(tensor) + # The shape has to be 1. Set it if it does not have the information. + output_tensor = array_ops.reshape(output_tensor, [1]) + return output_tensor + + def _show_max_abs(tensor): + tensor = math_ops.cast(tensor, dtypes.float32) + output_tensor = math_ops.reduce_max(math_ops.abs(tensor)) + zero = constant_op.constant(0, dtypes.float32) + output_tensor = gen_math_ops.maximum(zero, output_tensor) + # The shape has to be 1. Set it if it does not have the information. + output_tensor = array_ops.reshape(output_tensor, [1]) + return output_tensor + + if self._trace_mode == _TRACE_MODE_NAN_INF: + return _detect_nan_inf(tensor) + if self._trace_mode == _TRACE_MODE_PART_TENSOR: + return tensor + if self._trace_mode == _TRACE_MODE_FULL_TENSOR: + return tensor + if self._trace_mode == _TRACE_MODE_NORM: + return _show_norm(tensor) + if self._trace_mode == _TRACE_MODE_MAX_ABS: + return _show_max_abs(tensor) + raise RuntimeError( + 'Tensor trace fun for %s is not yet implemented' % self._trace_mode) + + def _make_tensor_trace_fun(self, tensor_name): + """Makes the tensor tracing function called by outside compilation. + + Args: + tensor_name: name of the tensor being traced. + + Returns: + A function to be passed as the first argument to outside compilation. + + Raises: + RuntimeError: If the trace mode is invalid. + """ + + def _print_tensor(tensor_name, num_elements, tensor, output_tensor): + """Prints a tensor value to a file. + + Args: + tensor_name: name of the tensor being traced. + num_elements: number of elements to print (-1 means print all). + tensor: the tensor needs to be returned. + output_tensor: the tensor needs to be printed. + + Returns: + The same tensor passed via the "tensor" argument. + + Raises: + ValueError: If tensor_name is not already in + self._tensorname_idx_map. + """ + + if self._submode == _SUBMODE_BRIEF: + if tensor_name not in self._tensorname_idx_map: + raise ValueError( + 'Tensor name %s is not in the tensorname_idx_map'%tensor_name) + msg = '%d'%self._tensorname_idx_map[tensor_name] + else: + msg = '"%s"'%tensor_name + + if self._trace_file_path: + output_stream = _OUTPUT_STREAM_ESCAPE + self._trace_file_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 _show_part_tensor(tensor): + """Trace function for printing part of the tensor.""" + + return _print_tensor(tensor_name, self._part_tensor_size, + tensor, tensor) + + def _show_full_tensor(tensor): + """Trace function for printing the entire tensor.""" + + return _print_tensor(tensor_name, -1, tensor, tensor) + + if self._trace_mode == _TRACE_MODE_PART_TENSOR: + return _show_part_tensor + # The input tensor has a shape of "[1]" for _TRACE_MODE_NAN_INF, + # _TRACE_MODE_NORM, and _TRACE_MODE_MAX_ABS, as related computations are + # performed within TPUs and only their results are transferred to CPU. + # Simply, print the full tensor for these trace modes. + if self._trace_mode in [ + _TRACE_MODE_NAN_INF, _TRACE_MODE_NORM, _TRACE_MODE_FULL_TENSOR, + _TRACE_MODE_MAX_ABS + ]: + return _show_full_tensor + + raise RuntimeError('Tensor trace fun for %s is not yet implemented' + %self._trace_mode) + + def _skip_op(self, op_id, op, user_included, user_excluded, + in_exec_path=True): + """Returns True if we should not trace Op.""" + + if 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 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.unsafe_op(op): + self._instrument_records[op.name] = TensorTracer.reason( + op_id, _REASON_UNSAFE_OP) + return True + if TensorTracer.device_mismatch(self._device_type, op): + self._instrument_records[op.name] = TensorTracer.reason( + op_id, _REASON_DEVICE_MISMATCH) + return True + if TensorTracer.less_interesting_op(op): + self._instrument_records[op.name] = TensorTracer.reason( + op_id, _REASON_LESS_INTERESTING_OP) + return True + return False + + def _skip_tensor(self, op_id, out_tensor, user_included, + user_excluded): + """Returns True if we should not trace out_tensor.""" + + # Skips a tensor if the tensor has a non-numeric type. + # Note: we cannot use check_ops.is_numeric_tensor(out_tensor) + # because it also excludes tensors with dtypes, bool, and + # float32_ref, which we actually want to trace. + non_numeric_tensor_types = set([dtypes.variant, dtypes.resource, + dtypes.string]) + if out_tensor.dtype in non_numeric_tensor_types: + self._instrument_records[out_tensor.name] = TensorTracer.reason( + op_id, _REASON_NON_NUMERIC_TENSOR) + return True + + if user_included: + self._instrument_records[out_tensor.name] = TensorTracer.reason( + op_id, _REASON_USER_INCLUDED) + return False + if user_excluded: + self._instrument_records[out_tensor.name] = TensorTracer.reason( + op_id, _REASON_USER_EXCLUDED) + return True + if not out_tensor.get_shape().is_fully_defined(): + # If trace mode is nan-inf, norm or max, then the tensor will be reduced + # to a scalar before the outside compilation call. + if self._trace_mode in [ + _TRACE_MODE_NAN_INF, _TRACE_MODE_NORM, _TRACE_MODE_MAX_ABS + ]: + self._instrument_records[out_tensor.name] = TensorTracer.reason( + op_id, _REASON_TENSOR_GET_TRACED) + return False + else: + self._instrument_records[out_tensor.name] = TensorTracer.reason( + op_id, _REASON_DYNAMIC_SHAPE) + return True + rank = len(out_tensor.shape) + if rank < 1: + # scalar + if TensorTracer.unsafe_scalar_trace(out_tensor.op): + self._instrument_records[out_tensor.name] = TensorTracer.reason( + op_id, _REASON_UNSAFE_SCALAR) + return True + else: + self._instrument_records[out_tensor.name] = TensorTracer.reason( + op_id, _REASON_SCALAR_GET_TRACED) + return False + else: + # tensor + self._instrument_records[out_tensor.name] = TensorTracer.reason( + op_id, _REASON_TENSOR_GET_TRACED) + return False + + def _filter_execution_path_operations(self, operations, fetches): + """Returns the set of ops in the execution path to compute given fetches.""" + # If no fetch provided, then return all operations. + if fetches is None: + return set(operations) + # Convert to list, if a single element is provided. + if not isinstance(fetches, (list, tuple)): + fetches = [fetches] + # If a tensor is given as fetch, convert it to op. + op_fetches = [] + for fetch in fetches: + if isinstance(fetch, ops.Operation): + op_fetches.append(fetch) + elif isinstance(fetch, ops.Tensor): + op_fetches.append(fetch.op) + else: + raise RuntimeError('Given fetch:%s is neither a tensor nor an op.' + %fetch) + + execution_path_operations = set(op_fetches) + traverse_stack = list(op_fetches) + while True: + if not traverse_stack: + break + head_op = traverse_stack.pop() + input_ops = [tensor_input.op for tensor_input in head_op.inputs] + input_ops.extend(head_op.control_inputs) + + for input_op in input_ops: + if input_op not in execution_path_operations: + execution_path_operations.add(input_op) + traverse_stack.append(input_op) + return execution_path_operations + + def _pre_tracing(self, graph): + """Work needs to be done prior to TPU or CPU tracing.""" + + 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) + # 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) + + def _post_tracing(self, succeed, sorted_or_cycle): + """Work needs to be done after TPU or CPU tracing.""" + + self._write_reason_section() + self._write_graph_section(succeed, sorted_or_cycle) + self._close_report_file() + + def _get_checkpoints(self, graph): + """Returns the list of Ops that produce the tensors traced with API. + + Args: + graph: the graph of Ops. + + Returns: + A set of operation names which should be traced. + """ + + self._write_report('%s %s\n'%(_MARKER_SECTION_BEGIN, + _TENSOR_TRACER_CHECKPOINT)) + checkpoint_operations = set() + tensor_tracer_variables = graph.get_collection(_TENSOR_TRACER_COLLECTION) + for (tensor, checkpoint_name) in tensor_tracer_variables: + self._write_report('%s %s\n'%(tensor.name, checkpoint_name)) + checkpoint_operations.add(tensor.op.name) + self._write_report('%s %s\n'%(_MARKER_SECTION_END, + _TENSOR_TRACER_CHECKPOINT)) + return checkpoint_operations + + def trace_tpu(self, graph, result_tensor, num_replicas=None, fetches=None): + """Traces the tensors generated by TPU Ops in a TF graph. + + 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. + fetches: the list of fetches given to session.run, used to determine the + ops in execution path. If None, the whole graph will be traced. + + 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. + """ + + def _cast_unsupported_dtypes(tensor): + """Casts tensor to a supported type.""" + + if tensor.dtype.__eq__(dtypes.int64): + # outside-compilation doesn't support int64 input yet. + return math_ops.cast(tensor, dtypes.int32) + if tensor.dtype.__eq__(dtypes.bfloat16) or tensor.dtype.__eq__( + dtypes.float16): + # Since host can't handle bf16, convert tensor to f32. + return math_ops.cast(tensor, dtypes.float32) + return tensor + + 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) + # 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) + 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) + 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 + # 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 + processed_out_tensor = self._preprocess_traced_tensor(out_tensor) + processed_out_tensor = _cast_unsupported_dtypes(processed_out_tensor) + trace_op = tpu.outside_compilation( + self._make_tensor_trace_fun(tensor_name), processed_out_tensor) + if consumers: + for consumer_op in consumers: + # pylint: disable=protected-access + consumer_op._add_control_input(trace_op) + # pylint: enable=protected-access + else: + # if there is no consumer, we will add the control dependence later + # when we add the control dependency to the output operations. + tracing_ops.append(trace_op) + self._post_tracing(succeed, sorted_or_cycle) + return (result_tensor_copy, tracing_ops) + + def trace_cpu(self, graph): + """Traces the tensors generated by CPU Ops in a TF graph. + + Args: + graph: the graph of Ops executed on the CPU. + + 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). + """ + + self._device_type = _DEVICE_TYPE_CPU + TensorTracer.check_device_type(self._device_type) + self._num_replicas = 1 + self._replica_id = 0 + (operations, succeed, sorted_or_cycle) = self._pre_tracing(graph) + tracing_calls = {} + 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 + 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 + processed_out_tensor = self._preprocess_traced_tensor(out_tensor) + trace_fun = self._make_tensor_trace_fun(out_tensor.name) + trace_call = (trace_fun, [processed_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 ab89c6aa8ca3a5e62bdfab991ee1a9335397c070..6ae718cc2c9716587849aeee8abcd0a1de82a9ae 100644 --- a/tensorflow/contrib/tpu/python/tpu/topology.py +++ b/tensorflow/contrib/tpu/python/tpu/topology.py @@ -19,10 +19,27 @@ from __future__ import division 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 +def _tpu_device_name(job, task, device): + """Returns the device name for the TPU `device` on `task` of `job`.""" + if job is None: + return "/task:%d/device:TPU:%d" % (task, device) + else: + return "/job:%s/task:%d/device:TPU:%d" % (job, task, device) + + +def _tpu_host_device_name(job, task): + """Returns the device name for the CPU device on `task` of `job`.""" + if job is None: + return "/task:%d/device:CPU:0" % task + else: + return "/job:%s/task:%d/device:CPU:0" % (job, task) + + class Topology(object): """Describes a set of TPU devices. @@ -71,6 +88,8 @@ class Topology(object): raise ValueError("`device_coordinates` must be a rank 3 int32 array " "with minor dimension equal to the mesh shape rank") + self._topology_tasks, self._topology_devices = self._invert_topology() + def _parse_topology(self, serialized): """Parses a serialized `TopologyProto` into `self`.""" proto = topology_pb2.TopologyProto() @@ -106,6 +125,17 @@ class Topology(object): len(proto.mesh_shape))) self._device_coordinates = coords + def _invert_topology(self): + """Inverts a [task,device,axis] topology to [x,y,z] -> task/device maps.""" + tasks = np.full(list(self.mesh_shape), -1, dtype=np.int32) + devices = np.full(list(self.mesh_shape), -1, dtype=np.int32) + for task in xrange(self.device_coordinates.shape[0]): + for device in xrange(self.device_coordinates.shape[1]): + x, y, z = self.device_coordinates[task, device, :] + tasks[x, y, z] = task + devices[x, y, z] = device + return tasks, devices + @property def mesh_shape(self): """A rank 1 int32 array describing the shape of the TPU topology.""" @@ -130,6 +160,43 @@ class Topology(object): """ return self._device_coordinates + def task_ordinal_at_coordinates(self, device_coordinates): + """Returns the TensorFlow task number attached to `device_coordinates`. + + Args: + device_coordinates: An integer sequence describing a device's physical + coordinates in the TPU fabric. + + Returns: + Returns the TensorFlow task number that contains the TPU device with those + physical coordinates. + """ + return self._topology_tasks[tuple(device_coordinates)] + + def tpu_device_ordinal_at_coordinates(self, device_coordinates): + """Returns the TensorFlow device number at `device_coordinates`. + + Args: + device_coordinates: An integer sequence describing a device's physical + coordinates in the TPU fabric. + + Returns: + Returns the TensorFlow device number within the task corresponding to + attached to the device with those physical coordinates. + """ + return self._topology_devices[tuple(device_coordinates)] + + def cpu_device_name_at_coordinates(self, device_coordinates, job=None): + """Returns the CPU device attached to a logical core.""" + return _tpu_host_device_name( + job, self._topology_tasks[tuple(device_coordinates)]) + + def tpu_device_name_at_coordinates(self, device_coordinates, job=None): + """Returns the name of the TPU device assigned to a logical core.""" + return _tpu_device_name(job, + self._topology_tasks[tuple(device_coordinates)], + self._topology_devices[tuple(device_coordinates)]) + @property def num_tasks(self): """Returns the number of TensorFlow tasks in the TPU slice.""" diff --git a/tensorflow/contrib/tpu/python/tpu/topology_test.py b/tensorflow/contrib/tpu/python/tpu/topology_test.py index e67fdb263aa48a37f65c3623365ebcf8f98bebd4..fafe3254d84551d3d7ed8a9d3346849411714f97 100644 --- a/tensorflow/contrib/tpu/python/tpu/topology_test.py +++ b/tensorflow/contrib/tpu/python/tpu/topology_test.py @@ -27,7 +27,7 @@ from tensorflow.python.platform import test class TopologyTest(test.TestCase): def testSerialization(self): - """Test if the class is able to generate serialzied string.""" + """Tests if the class is able to generate serialized strings.""" original_topology = topology.Topology( mesh_shape=[1, 1, 2], device_coordinates=[[[0, 0, 0], [0, 0, 1]]], diff --git a/tensorflow/contrib/tpu/python/tpu/tpu.py b/tensorflow/contrib/tpu/python/tpu/tpu.py index a51dcc8020a423632667d9828ac4ca4a5644a4c2..ebbccea02c70f06ac3e1231a359f2df4ebad3142 100644 --- a/tensorflow/contrib/tpu/python/tpu/tpu.py +++ b/tensorflow/contrib/tpu/python/tpu/tpu.py @@ -19,23 +19,29 @@ 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.proto import dynamic_padding_pb2 as dynamic_padding 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.python.compat import compat as api_compat from tensorflow.python.framework import device as pydev +from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors from tensorflow.python.framework import ops +from tensorflow.python.framework import tensor_shape from tensorflow.python.ops import array_ops from tensorflow.python.ops import control_flow_ops +from tensorflow.python.ops import math_ops from tensorflow.python.ops import variable_scope from tensorflow.python.platform import tf_logging as logging from tensorflow.python.util import compat +from tensorflow.python.util import nest # Operations that indicate some error in the users graph, e.g. a placeholder @@ -372,14 +378,11 @@ 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. - 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() + external_control_inputs = [ + array_ops.identity(x.outputs[0]).op + for x in external_control_inputs + if x.outputs + ] # pylint: disable=protected-access op._add_control_inputs(external_control_inputs) # pylint: enable=protected-access @@ -483,14 +486,19 @@ 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: computation: A Python function that builds the computation to replicate. inputs: A list of lists of input tensors or `None` (equivalent to `[[]]`), indexed by `[replica_num][input_num]`. All replicas must - have the same number of inputs. + have the same number of inputs. Each input can be a nested structure + containing values that are convertible to tensors. Note that passing an + N-dimension list of compatible values will result in a N-dimention list of + scalar tensors rather than a single Rank-N tensors. If you need different + behavior, convert part of inputs to tensors with `tf.convert_to_tensor`. infeed_queue: If not `None`, the `InfeedQueue` from which to append a tuple of arguments as inputs to computation. device_assignment: If not `None`, a `DeviceAssignment` describing the @@ -500,15 +508,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, @@ -516,7 +634,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 @@ -529,7 +648,11 @@ def split_compile_and_replicate(computation, computation: A Python function that builds the computation to replicate. inputs: A list of lists of input tensors or `None` (equivalent to `[[]]`), indexed by `[replica_num][input_num]`. All replicas must - have the same number of inputs. + have the same number of inputs. Each input can be a nested structure + containing values that are convertible to tensors. Note that passing an + N-dimension list of compatible values will result in a N-dimention list of + scalar tensors rather than a single Rank-N tensors. If you need different + behavior, convert part of inputs to tensors with `tf.convert_to_tensor`. infeed_queue: If not `None`, the `InfeedQueue` from which to append a tuple of arguments as inputs to computation. device_assignment: If not `None`, a `DeviceAssignment` describing the @@ -542,6 +665,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]`. @@ -549,6 +681,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 @@ -583,24 +719,32 @@ def split_compile_and_replicate(computation, if num_replicas == 0: return [] + # Checks all replicas have the same structure. + for i in xrange(1, num_replicas): + nest.assert_same_structure(inputs[0], inputs[i]) + + # Flatten inputs. + flat_inputs = [ + nest.flatten(per_replica_input) for per_replica_input in inputs + ] # Converts inputs to Tensors. - inputs = [[ops.convert_to_tensor(x) for x in inp] for inp in inputs] + flat_inputs = [[ops.convert_to_tensor(x) for x in inp] for inp in flat_inputs] # Verifies that all replicas have matching numbers and types of inputs - input_types = [x.dtype for x in inputs[0]] - input_arity = len(input_types) + flat_input_types = [x.dtype for x in flat_inputs[0]] + input_arity = len(inputs[0]) + flat_input_arity = len(flat_input_types) for i in range(num_replicas): if len(inputs[i]) != input_arity: raise ValueError("Replicas must have the same number of inputs. " "Replica 0 had {} inputs, replica {} had {} " "inputs.".format(input_arity, i, len(inputs[i]))) - types = [x.dtype for x in inputs[i]] - if types != input_types: - raise ValueError( - "Replicas must have matching input types. Replica 0 had " - "input types {}, replica {} had input types {}".format( - input_types, i, types)) + types = [x.dtype for x in flat_inputs[i]] + if types != flat_input_types: + raise ValueError("Replicas must have matching input types. Replica 0 had " + "input types {}, replica {} had input types {}".format( + flat_input_types, i, types)) arg_error = xla.check_function_argument_count( computation, input_arity, infeed_queue) @@ -619,13 +763,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, input_arity): - replicas = [inputs[replica][i] for replica in xrange(num_replicas)] - computation_inputs.append( + flat_replicated_inputs = [] + for i in range(0, len(flat_inputs[0])): + replicas = [flat_inputs[replica][i] for replica in xrange(num_replicas)] + flat_replicated_inputs.append( tpu_ops.tpu_replicated_input(replicas, name="input{}".format(i))) cluster_name = graph.unique_name("cluster") @@ -645,10 +810,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 flat_replicated_inputs: + # pylint: disable=protected-access + # Add an attribute to the identity node so that they could be removed in + # encapsulate TPU computation pass if unused. However we don't remove + # inputs when dynamic padding is enabled. + # TODO(rxsang): Use other ways except argument index in padding_map so + # outside compilation can work with dynamic padding correctly. + if maximum_shapes is None: + i.op._set_attr("_tpu_input_identity", + attr_value_pb2.AttrValue(b=True)) + # pylint: enable=protected-access + + # Unflatten the computation inputs to match original input structure. + computation_inputs = nest.pack_sequence_as( + structure=inputs[0], + flat_sequence=flat_replicated_inputs[:flat_input_arity]) # If there is an infeed queue, adds the dequeued values to the # computation's inputs. @@ -690,47 +871,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)): - new_output_tensors.append(array_ops.identity(t)) - output_tensors = new_output_tensors context.ExitResult(output_tensors) finally: context.report_unsupported_operations() @@ -742,11 +888,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() @@ -756,39 +897,157 @@ 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 shard(computation, - inputs=None, - num_shards=1, - input_shard_axes=None, - outputs_from_all_shards=True, - output_shard_axes=None, - infeed_queue=None, - device_assignment=None, - name=None): +def _postprocess_non_flat_outputs(outputs): + """Validates non-flat outputs, add backs device assignments and other attrs. + + Args: + outputs: Output from `computation` inside `tpu.rewrite`. + + Returns: + Tensors extracted from outputs and an empty list because Operations are not + allowed in non-flat outputs.. + """ + + # Flatten output items. + flat_outputs = nest.flatten(outputs) + + # Convert all non-Operation outputs to Tensors. + for i, o in enumerate(flat_outputs): + if isinstance(o, ops.Operation): + raise ValueError( + "tpu.rewrite does not support Operation as return value in non-flat " + "output structure. You can set returned Operations as control " + "dependencies of returned Tensors so Operations are triggered when " + 'Tensors are evaluated. Operation found: "%s"' % o.name) + + try: + o = ops.convert_to_tensor(o) + except Exception as e: + raise ValueError( + "TPU function return values must all either be Operations or " + 'convertible to Tensors. Got error: "%s"' % str(e)) + + # Wraps outputs in Identity ops. Otherwise a replicated input copied + # straight to an output would bypass the replicate(). This would be bad + # because the TPUReplicatedInput/TPUReplicatedOutput operator would not + # be rewritten away, leading to a runtime error. + # TODO(phawkins): extend the rewrite to elide these nodes instead. + with ops.device(core(0)): + o = array_ops.identity(o) + # pylint: disable=protected-access + o.op._set_attr("_tpu_output_identity", attr_value_pb2.AttrValue(b=True)) + # pylint: enable=protected-access + flat_outputs[i] = array_ops.identity(o) + + # All flat_outputs are Tensors, and no Operations. + return flat_outputs, [] + + +def split_compile_and_shard(computation, + inputs=None, + num_shards=1, + input_shard_axes=None, + outputs_from_all_shards=True, + output_shard_axes=None, + infeed_queue=None, + device_assignment=None, + name=None): """Shards `computation` for parallel execution. `inputs` must be a list of Tensors or None (equivalent to an empty list), each @@ -804,9 +1063,6 @@ def 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. @@ -842,12 +1098,14 @@ def shard(computation, is equal to the number of cores in the TPU system. name: (Deprecated) Does nothing. Returns: - A list of output tensors. + A tuple of (compile op, [output tensors]). Raises: ValueError: If num_shards <= 0 ValueError: If len(input_shard_axes) != len(inputs) ValueError: If len(output_shard_axes) != len(outputs from `computation`) """ + # TODO(phawkins): consider adding support for broadcasting Tensors passed as + # inputs. if num_shards <= 0: raise ValueError("num_shards must be a positive integer.") @@ -877,7 +1135,7 @@ def shard(computation, else: transposed_inputs = [[]] * num_shards - outputs = replicate( + compile_op, outputs = split_compile_and_replicate( computation, transposed_inputs, infeed_queue=infeed_queue, @@ -894,7 +1152,7 @@ def shard(computation, # one so it can be used as a control dependency or fetch node. # TODO(b/36647078) remove disable when pylint bug is fixed. # pylint: disable=indexing-exception - return [outputs[0]] + return compile_op, [outputs[0]] # pylint: enable=indexing-exception # TODO(b/36647078) remove disable when pylint bug is fixed. @@ -928,7 +1186,87 @@ def shard(computation, # TODO(phawkins): use a smarter policy, e.g., round-robin across shards. results.append(x[0]) - return results + return compile_op, results + + +def shard(computation, + inputs=None, + num_shards=1, + input_shard_axes=None, + outputs_from_all_shards=True, + output_shard_axes=None, + infeed_queue=None, + device_assignment=None, + name=None): + """Shards `computation` for parallel execution. + + `inputs` must be a list of Tensors or None (equivalent to an empty list), each + of which has a corresponding split axis (from `input_shard_axes`). Each input + is split into `num_shards` pieces along the corresponding axis, and + computation is applied to each shard in parallel. + + Tensors are broadcast to all shards if they are lexically captured by + `computation`. e.g., + + x = tf.constant(7) + def computation(): + return x + 3 + ... = shard(computation, ...) + + TODO(phawkins): consider adding support for broadcasting Tensors passed + as inputs. + + If `outputs_from_all_shards` is true, the outputs from all shards of + `computation` are concatenated back together along their `output_shards_axes`. + Otherwise, each output is taken from an arbitrary shard. + + Inputs and outputs of the computation must be at least rank-1 Tensors. + + Args: + computation: A Python function that builds a computation to apply to each + shard of the input. + inputs: A list of input tensors or None (equivalent to an empty list). Each + input tensor has a corresponding shard axes, given by `input_shard_axes`, + which must have size divisible by `num_shards`. + num_shards: The number of shards. + input_shard_axes: A list of dimensions along which to shard `inputs`, or + `None`. `None` means "shard all inputs along dimension 0". If not `None`, + there must be one dimension per input. + outputs_from_all_shards: Boolean or list of boolean. For each output, if + `True`, outputs from all shards are concatenated along the corresponding + `output_shard_axes` entry. Otherwise, each output is taken + from an arbitrary shard. If the argument is a boolean, the argument's + value is used for each output. + output_shard_axes: A list of dimensions along which to concatenate the + outputs of `computation`, or `None`. `None` means "concatenate all outputs + along dimension 0". If not `None`, there must be one dimension per output. + Ignored if `outputs_from_all_shards` is False. + infeed_queue: If not `None`, the `InfeedQueue` to use to augment the inputs + of `computation`. + device_assignment: If not `None`, a `DeviceAssignment` describing the + mapping between logical cores in the computation with physical cores in + the TPU topology. Uses a default device assignment if `None`. The + `DeviceAssignment` may be omitted if each shard of the computation uses + only one core, and there is either only one shard, or the number of shards + is equal to the number of cores in the TPU system. + name: (Deprecated) Does nothing. + Returns: + A list of output tensors. + Raises: + ValueError: If num_shards <= 0 + ValueError: If len(input_shard_axes) != len(inputs) + ValueError: If len(output_shard_axes) != len(outputs from `computation`) + """ + return split_compile_and_shard( + computation, + inputs=inputs, + num_shards=num_shards, + input_shard_axes=input_shard_axes, + outputs_from_all_shards=outputs_from_all_shards, + output_shard_axes=output_shard_axes, + infeed_queue=infeed_queue, + device_assignment=device_assignment, + name=name)[1] def batch_parallel(computation, @@ -1004,9 +1342,14 @@ def rewrite(computation, `rewrite` is a list of tensors corresponding to the tensors from the output of `computation`. - All `Operation`s returned from `computation` will be executed when - evaluating any of the returned output tensors. + All `Operation`s constructed during `computation` will be executed when + evaluating any of the returned output tensors, not just the ones returned. inputs: A list of input tensors or `None` (equivalent to an empty list). + Each input can be a nested structure containing values that are + convertible to tensors. Note that passing an N-dimension list of + compatible values will result in a N-dimention list of scalar tensors + rather than a single Rank-N tensors. If you need different behavior, + convert part of inputs to tensors with `tf.convert_to_tensor`. infeed_queue: If not `None`, the `InfeedQueue` from which to append a tuple of arguments as inputs to `computation`. device_assignment: if not `None`, a `DeviceAssignment` describing the @@ -1015,11 +1358,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( @@ -1114,7 +1461,7 @@ def validate_inference_rewrite_for_variables(graph): Raises: RuntimeError: if validation failed. """ - if not any([x.type == "GuaranteeConst" for x in graph.get_operations()]): + if not any(x.type == "GuaranteeConst" for x in graph.get_operations()): raise RuntimeError( "No GuaranteeConst ops found in the graph after running " "tpu.rewrite_for_inference(...). Please check that you are using " diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_context.py b/tensorflow/contrib/tpu/python/tpu/tpu_context.py index da6bdf67d686fba09d66386de982b57aa28d4dd4..9ff0c55187e42d83e3c2c962ceca2efafa4c55e9 100644 --- a/tensorflow/contrib/tpu/python/tpu/tpu_context.py +++ b/tensorflow/contrib/tpu/python/tpu/tpu_context.py @@ -41,7 +41,7 @@ _NUM_CORES_TO_COMPUTATION_SHAPE = { class TPUContext(object): - """The context of current input_fn invocation.""" + """A context that holds the current configuration of the TPU computation.""" def __init__(self, internal_ctx, @@ -208,7 +208,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: diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_embedding.py b/tensorflow/contrib/tpu/python/tpu/tpu_embedding.py index cd9bfbcdce0931a8fa84c22d959c780235c7a9f9..1a909a3ac6fae79070a7762b94bfa138f93a5fb5 100644 --- a/tensorflow/contrib/tpu/python/tpu/tpu_embedding.py +++ b/tensorflow/contrib/tpu/python/tpu/tpu_embedding.py @@ -28,6 +28,7 @@ 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.python.framework import dtypes from tensorflow.python.framework import ops from tensorflow.python.framework import sparse_tensor @@ -43,19 +44,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( @@ -158,6 +146,8 @@ class AdamParameters(_OptimizationParameters): beta1=0.9, beta2=0.999, epsilon=1e-08, + lazy_adam=True, + sum_inside_sqrt=True, use_gradient_accumulation=False, pipeline_execution_with_tensor_core=True): """Optimization parameters for Adam. @@ -169,10 +159,14 @@ class AdamParameters(_OptimizationParameters): beta2: A float value. The exponential decay rate for the 2nd moment estimates. epsilon: A small constant for numerical stability. + lazy_adam: Use lazy Adam instead of Adam. Lazy Adam trains faster. + Please see `optimization_parameters.proto` for details. + sum_inside_sqrt: This improves training speed. Please see + `optimization_parameters.proto` for details. use_gradient_accumulation: setting this to `True` makes embedding - gradients calculation more accurate but slower. Please see - `optimization_parameters.proto` for details. - for details. + gradients calculation more accurate but slower. Please see + `optimization_parameters.proto` for details. + for details. pipeline_execution_with_tensor_core: setting this to `True` makes training faster, but trained model will be different if step N and step N+1 involve the same set of embedding ID. Please see @@ -184,6 +178,8 @@ class AdamParameters(_OptimizationParameters): self.beta1 = beta1 self.beta2 = beta2 self.epsilon = epsilon + self.lazy_adam = lazy_adam + self.sum_inside_sqrt = sum_inside_sqrt class StochasticGradientDescentParameters(_OptimizationParameters): @@ -276,6 +272,8 @@ class TPUEmbedding(object): # could have a field to indicate that the feature should not be used to # update embedding table (cr/204852758, cr/204940540). Also, this can support # different combiners for different features within the same table. + # TODO(shizhiw, b/118512626): Remove `batch_size` from `__init__` and move it + # to `FeatureConfig`? # TODO(shizhiw): will it be cleaner to make `table_to_config_dict` and # `feature_to_table_dict` lists of `TableSpec` and `FeatureSpec` respectively? @@ -291,10 +289,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: @@ -305,12 +302,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 - not be `None` in inference. - tpu_embedding_test: A `bool`. Only used for testing. + be `None` in inference. Raises: ValueError: if any input is invalid. @@ -327,15 +323,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,7 +377,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`. @@ -884,6 +882,10 @@ class _AdamHandler(_OptimizerHandler): self._optimization_parameters.beta2) table_descriptor.optimization_parameters.adam.epsilon = ( self._optimization_parameters.epsilon) + table_descriptor.optimization_parameters.adam.use_non_lazy_adam = ( + not self._optimization_parameters.lazy_adam) + table_descriptor.optimization_parameters.adam.use_sum_inside_sqrt = ( + self._optimization_parameters.sum_inside_sqrt) def create_variables_and_ops(self, table, variable_name, num_hosts, table_config, table_variables, @@ -1055,17 +1057,14 @@ def _create_partitioned_variables(name, 'As TPU embedding is not optimized for small tables, ' 'please consider other ways for this embedding lookup.') - slicing = [num_hosts, 1] - - # TODO(shizhiw): deprecated, use tf.get_variable()? - return partitioned_variables.create_partitioned_variables( - name=name, - slicing=slicing, + return list(variable_scope.get_variable( + name, shape=(vocabulary_size, embedding_dimension), + partitioner=partitioned_variables.fixed_size_partitioner(num_hosts), dtype=dtypes.float32, initializer=initializer, collections=collections, - trainable=False) + trainable=False)) @ops.RegisterGradient('TPUEmbeddingActivations') diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py b/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py index 29aa0d656807619b431b7844a159be2d684aa462..41d0562513daba36798a88d6d5b5330e54d88d15 100644 --- a/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py +++ b/tensorflow/contrib/tpu/python/tpu/tpu_estimator.py @@ -31,9 +31,13 @@ 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.tpu import error_handling +from tensorflow.contrib.tpu.python.tpu import functional as tpu_functional from tensorflow.contrib.tpu.python.tpu import session_support +from tensorflow.contrib.tpu.python.tpu import tensor_tracer from tensorflow.contrib.tpu.python.tpu import tpu from tensorflow.contrib.tpu.python.tpu import tpu_config from tensorflow.contrib.tpu.python.tpu import tpu_context @@ -44,15 +48,16 @@ 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.python.client import session as tf_session from tensorflow.python.data.ops import dataset_ops from tensorflow.python.data.util import nest as data_nest from tensorflow.python.estimator import estimator as estimator_lib from tensorflow.python.estimator import model_fn as model_fn_lib -from tensorflow.python.estimator import util as estimator_util from tensorflow.python.estimator.export import export_output as export_output_lib from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes from tensorflow.python.framework import errors +from tensorflow.python.framework import function from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import check_ops @@ -88,6 +93,7 @@ _ONE_GIGABYTE = 1024 * 1024 * 1024 _TPU_ENQUEUE_OPS = '_tpu_enqueue_ops' _TPU_TRAIN_OP = '_tpu_train_op' _REWRITE_FOR_INFERENCE_MODE = '_rewrite_for_inference' +_KEY_WHEN_PREDICTIONS_IS_A_TENSOR = '_key_when_predictions_is_a_tensor' # Ideally _USE_TPU_KEY should be reserved as well. However there are already # models that make use of this key, thus it can not be reserved now to prevent @@ -109,6 +115,25 @@ ops.register_proto_function( from_proto=resource_variable_ops._from_proto_fn) # pylint: disable=protected-access +def _is_iterable(obj): + """A Python 2 and 3 compatible util to check whether `obj` is iterable.""" + try: + iter(obj) + return True + except TypeError: + return False + + +class CatchInvalidHostcallFunctions(control_flow_ops.XLAControlFlowContext): + + def AddOp(self, op): + if op.type in [ + 'AudioSummary', 'AudioSummaryV2', 'HistogramSummary', 'ImageSummary', + 'MergeSummary', 'ScalarSummary', 'TensorSummary', 'TensorSummaryV2' + ]: + raise ValueError('Use tf.contrib.summary inside of host_calls.') + + def _create_global_step(graph): graph = graph or ops.get_default_graph() if training.get_global_step(graph) is not None: @@ -289,9 +314,9 @@ class TPUEstimatorSpec(model_fn_lib._TPUEstimatorSpec): # pylint: disable=prote host_calls['host_call'] = host_call _OutfeedHostCall.validate(host_calls) - training_hooks = list(training_hooks or []) - evaluation_hooks = list(evaluation_hooks or []) - prediction_hooks = list(prediction_hooks or []) + training_hooks = tuple(training_hooks or []) + evaluation_hooks = tuple(evaluation_hooks or []) + prediction_hooks = tuple(prediction_hooks or []) for hook in training_hooks + evaluation_hooks + prediction_hooks: if not isinstance(hook, session_run_hook.SessionRunHook): @@ -326,7 +351,17 @@ 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'])] - hooks = list(hooks or []) + if tensor_tracer.TensorTracer.is_enabled(): + 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)] + hooks = tuple(hooks or []) scaffold = self.scaffold_fn() if self.scaffold_fn else None return model_fn_lib.EstimatorSpec( mode=self.mode, @@ -402,13 +437,17 @@ class TPUInfeedOutfeedSessionHook(session_run_hook.SessionRunHook): ctx, enqueue_ops, dequeue_ops, + tpu_compile_op, run_infeed_loop_on_coordinator=True, - rendezvous=None): + rendezvous=None, + master=None, + session_config=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._run_infeed_loop_on_coordinator = run_infeed_loop_on_coordinator self._initial_infeed_sleep_secs = ( ctx.config.tpu_config.initial_infeed_sleep_secs) @@ -416,15 +455,15 @@ class TPUInfeedOutfeedSessionHook(session_run_hook.SessionRunHook): self._feed_error = None self._finished = False self._should_initialize_tpu = True + 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._init_ops = [tpu.initialize_system(job=self._master_job)] self._finalize_ops = [tpu.shutdown_system(job=self._master_job)] else: - self._init_ops = [] self._finalize_ops = [] summary_writer_init_ops = contrib_summary.summary_writer_initializer_op() @@ -465,12 +504,31 @@ class TPUInfeedOutfeedSessionHook(session_run_hook.SessionRunHook): def _create_infeed_controller(self, name, target, args): return _OpQueueContext(name=name, target=target, args=args) + def _assertCompilationSucceeded(self, result, coord): + proto = tpu_compilation_result.CompilationResultProto() + proto.ParseFromString(result) + if proto.status_error_message: + logging.error('Compilation failed: {}'.format(proto.status_error_message)) + coord.request_stop() + else: + logging.info('Compilation succeeded') + def after_create_session(self, session, coord): - logging.info('Init TPU system') - start = time.time() + if self._should_initialize_tpu: + logging.info('Init TPU system') + start = time.time() + with ops.Graph().as_default(): + with tf_session.Session( + self._master, config=self._session_config) as sess: + sess.run(tpu.initialize_system(job=self._master_job)) + logging.info('Initialized TPU in %d seconds', time.time() - start) + session.run(self._init_ops, options=config_pb2.RunOptions(timeout_in_ms=5 * 60 * 1000)) - logging.info('Initialized TPU in %d seconds', time.time() - start) + + if os.environ.get('TPU_SPLIT_COMPILE_AND_EXECUTE', '') == '1': + logging.info('Compiling user program: this may take a while...') + self._assertCompilationSucceeded(session.run(self._tpu_compile_op), coord) self._infeed_controller = self._create_infeed_controller( name='InfeedController', target=self._run_infeed, args=(session,)) @@ -512,13 +570,17 @@ class TPUInfeedOutfeedSessionHook(session_run_hook.SessionRunHook): class TPUInfeedOutfeedSessionHookForPrediction(TPUInfeedOutfeedSessionHook): - def __init__(self, ctx, enqueue_ops, dequeue_ops, rendezvous=None): + def __init__(self, ctx, enqueue_ops, dequeue_ops, tpu_compile_op, + rendezvous=None, master=None, session_config=None): super(TPUInfeedOutfeedSessionHookForPrediction, self).__init__( ctx, enqueue_ops, dequeue_ops, + tpu_compile_op=tpu_compile_op, run_infeed_loop_on_coordinator=False, - rendezvous=rendezvous) + rendezvous=rendezvous, + master=master, + session_config=session_config) def _create_infeed_controller(self, name, target, args): return _OpSignalOnceQueueContext(name=name, target=target, args=args) @@ -1229,7 +1291,7 @@ class _InputPipeline(object): # first one. self._infeed_queue = infeed_queues[0] return enqueue_ops, [ - estimator_util.MultiHostDatasetInitializerHook(all_dataset_initializers) + util_lib.MultiHostDatasetInitializerHook(all_dataset_initializers) ], run_infeed_loop_on_coordinator def _validate_input_pipeline(self): @@ -1255,6 +1317,44 @@ class _InputPipeline(object): logging.warn(err_msg) +def call_computation(computation, + experimental_exported_model_uses_all_cores=True): + """Call computation. + + computation uses a single-core for TPU inference. If + `experimental_exported_model_uses_all_cores` is `True`, this function will + round-robin + computation among all TPU cores visible to the host; otherwise, it will use + a single core. + + Args: + computation: A Python function that takes no inputs and builds computation + graph. If `computation` returns m outputs, this function will return a + list of m Tensors. + experimental_exported_model_uses_all_cores: Whether to round-robin among all + cores visible to the host, or to use a single core. + + Returns: + A list of output tensors. + """ + if experimental_exported_model_uses_all_cores: + # Using `TPUPartitionedCall` makes it possible to target a different + # TPU core with every `Session.run()` call. Note that the entire inference + # graph executes on a single core, and that invocations of this graph + # will round-robin among the cores attached to a host. + @function.Defun() + def tpu_subgraph(): + return computation() + + return tpu_functional.TPUPartitionedCall( + args=tpu_subgraph.captured_inputs, + device_ordinal=gen_tpu_ordinal_selector_op.tpu_ordinal_selector(), + Tout=[o.type for o in tpu_subgraph.definition.signature.output_arg], + f=tpu_subgraph) + else: + return computation() + + class _ModelFnWrapper(object): """A `model_fn` wrapper. @@ -1318,9 +1418,16 @@ 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, + fetches=[loss, train_op]) + # We must run train_op to update the variables prior to running the # outfeed. - with ops.control_dependencies([train_op]): + with ops.control_dependencies([train_op]+tracing_ops): host_call_outfeed_ops = [] if (isinstance(estimator_spec, model_fn_lib._TPUEstimatorSpec) # pylint: disable=protected-access and estimator_spec.host_call is not None): @@ -1448,7 +1555,7 @@ class _ModelFnWrapper(object): raise TypeError('TPUEstimatorSpec.predictions must be dict of Tensors.') for (key, tensor) in predictions.items(): - if tensor.shape[0].value is None: + if tensor.shape.dims[0].value is None: raise ValueError( 'The tensor with key ({}) in TPUEstimatorSpec.predictions has ' 'dynamic shape (should be static). Tensor: {}'.format(key, tensor)) @@ -1618,7 +1725,7 @@ class _OutfeedHostCall(object): 'Exception while calling %s: %s. It is likely the tensors ' '(%s[1]) do not match the ' 'function\'s arguments', name, e, name) - raise e + raise return ret def record(self, host_calls): @@ -1705,6 +1812,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: @@ -1712,6 +1823,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. # @@ -1724,24 +1843,39 @@ class _OutfeedHostCall(object): raise RuntimeError( 'All tensors outfed from TPU should preserve batch size ' 'dimension, but got scalar {}'.format(dequeue_ops[i][0])) - # TODO(xiejw): Allow users to specify the axis for batch size - # dimension. - dequeue_ops[i] = array_ops.concat(dequeue_ops[i], axis=0) + # TODO(xiejw): Make the specification of the outfeed combinaton + # function more explicit and well-documented. We may want to give the + # user the option of concatenating along any axis. + if (self._ctx.config.tpu_config.per_host_input_for_training is + tpu_config.InputPipelineConfig.BROADCAST): + # If the infeed is in BROADCAST mode (each core recieving the same + # input), then we assume that the cores also produce identical + # copies of the same output, and we simply take the output from + # the first core. This mode is used by Mesh-TensorFlow. + with ops.control_dependencies(dequeue_ops[i]): + dequeue_ops[i] = array_ops.identity(dequeue_ops[i][0]) + else: + # Assume that the input has been batch-split and that axis 0 of the + # output tensors represents the batch size. Concatenate along + # the axis 0 to re-combine the batch. + dequeue_ops[i] = array_ops.concat(dequeue_ops[i], axis=0) if self._tensor_keys[name] is not None: # The user-provided eval_metrics[1] is a dict. dequeue_ops = dict(zip(self._tensor_keys[name], dequeue_ops)) try: - ret[name] = 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 ' '(%s[1]) do not match the ' 'function\'s arguments', name, e, name) - raise e + 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 @@ -2033,7 +2167,10 @@ class TPUEstimator(estimator_lib.Estimator): batch_axis=None, eval_on_tpu=True, export_to_tpu=True, - warm_start_from=None): + export_to_cpu=True, + warm_start_from=None, + experimental_exported_model_uses_all_cores=False, + experimental_export_device_assignment=False): """Constructs an `TPUEstimator` instance. Args: @@ -2076,12 +2213,26 @@ 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 a `WarmStartSettings`, then all variables are warm-started, and it is assumed that vocabularies and Tensor names are unchanged. + experimental_exported_model_uses_all_cores: Whether to round-robin among + all cores visible to the host which is serving the saved model, or to + use a single core. This is a temporary flag to enable using all TPU + cores for inference with TPUPartitionedCall(). Once outside compilation + is supported in TPUPartitionedCall(), this flag will be enabled by + default. + experimental_export_device_assignment: Whether to include the device + assignment in the exported model. Doing so is useful in case of model + parallel inference but will tie the exported model to the TPU topology + used to export the model. Raises: ValueError: `params` has reserved keys already. @@ -2145,7 +2296,17 @@ class TPUEstimator(estimator_lib.Estimator): self._config, train_batch_size, eval_batch_size, predict_batch_size, use_tpu, eval_on_tpu) + 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 = {} @@ -2154,43 +2315,48 @@ class TPUEstimator(estimator_lib.Estimator): builder, input_receiver_fn_map, checkpoint_path, - strip_default_attrs, save_variables=True, mode=model_fn_lib.ModeKeys.PREDICT, export_tags=None, check_variables=True): if self._export_to_tpu and mode != model_fn_lib.ModeKeys.PREDICT: - 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, - strip_default_attrs, - 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 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: + 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, - strip_default_attrs, - 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: @@ -2205,6 +2371,88 @@ class TPUEstimator(estimator_lib.Estimator): raise ValueError('mode must be {}; ' 'got {}.'.format(_REWRITE_FOR_INFERENCE_MODE, mode)) + computation, capture = self._build_computation_for_inference( + features, labels, mode, config) + tensors = call_computation( + computation, + experimental_exported_model_uses_all_cores=self + ._experimental_exported_model_uses_all_cores) + estimator_spec, export_outputs_dict, predictions_dict, none_indices = ( + capture.get()) + predictions_list = tensors[:len(predictions_dict)] + export_outputs_list_without_none = tensors[len(predictions_dict):] + + # Reinsert `None`s which we've taken out in + # `_build_computation_for_inference()`. + export_outputs_list = [] + while none_indices or export_outputs_list_without_none: + if none_indices and none_indices[0] == len(export_outputs_list): + export_outputs_list.append(None) + none_indices.pop(0) + else: + export_outputs_list.append(export_outputs_list_without_none.pop(0)) + + # Reconstruct `export_outputs` with updated tensors. + new_export_outputs_dict = nest.pack_sequence_as(export_outputs_dict, + export_outputs_list) + export_outputs = estimator_spec.export_outputs + new_export_outputs = collections.OrderedDict( + (k, _clone_export_output_with_tensors(export_outputs[k], v)) + for k, v in six.iteritems(new_export_outputs_dict)) + # Reconstruct `predictions` with updated tensors. + new_predictions = nest.pack_sequence_as(predictions_dict, predictions_list) + if (len(new_predictions) == 1 and + _KEY_WHEN_PREDICTIONS_IS_A_TENSOR in new_predictions): + new_predictions = new_predictions[_KEY_WHEN_PREDICTIONS_IS_A_TENSOR] + + return estimator_spec._replace( + export_outputs=new_export_outputs, predictions=new_predictions) + + def _build_computation_for_inference(self, features, labels, mode, config): + capture = _CapturedObject() + + def computation(): + """Computation to be passed to `TPUPartitionedCall()`.""" + tpu_computation, tpu_capture = self._build_tpu_computation_for_inference( + features, labels, mode, config) + + if self._experimental_export_device_assignment: + # Export the device assignment as part of the model. This is useful for + # model parallel usecases where the model relies on the mapping between + # logical and physical devices. + with self._ctx.with_mode(mode) as ctx: + device_assignment = ctx.device_assignment + else: + device_assignment = None + tensors_on_cpu = tpu.rewrite_for_inference( + tpu_computation, device_assignment=device_assignment) + (estimator_spec, export_outputs_dict, export_outputs_list, + predictions_dict) = ( + tpu_capture.get()) + predictions_list = tensors_on_cpu[:len(predictions_dict)] + export_outputs_tpu_on_cpu_list = tensors_on_cpu[len(predictions_dict):] + + # Reconstruct tensors used in export_outputs, with TPU tensors replaced + # with their CPU counterpart returned from `rewrite_for_inference()`. + # `function.Defun()` does not like `None`s in return values, so we leave + # `None`s out but record their positions for later reconstruction. + export_outputs_list_without_none = [] + none_indices = [] + for i, t in enumerate(export_outputs_list): + if t is None: + none_indices.append(i) + else: + export_outputs_list_without_none.append( + export_outputs_tpu_on_cpu_list.pop(0)) + + capture.capture((estimator_spec, export_outputs_dict, predictions_dict, + none_indices)) + return predictions_list + export_outputs_list_without_none + + return computation, capture + + def _build_tpu_computation_for_inference(self, features, labels, mode, + config): capture = _CapturedObject() def computation(): @@ -2225,47 +2473,30 @@ class TPUEstimator(estimator_lib.Estimator): # We pick the TPU tensors out from `export_output` and later return them # from `computation` for rewriting. - tensors_dict = collections.OrderedDict( + export_outputs_dict = collections.OrderedDict( (k, _export_output_to_tensors(v)) for k, v in six.iteritems(estimator_spec.export_outputs)) - tensors = nest.flatten(tensors_dict) - tpu_tensors = [t for t in tensors if _is_tpu_tensor(t)] - - # We cannot return anything other than `tpu_tensors` here so we capture - # the rest for later use. - capture.capture((estimator_spec, tensors_dict, tensors)) - return tpu_tensors - - tpu_tensors_on_cpu = tpu.rewrite_for_inference(computation) - estimator_spec, tensors_dict, tensors = capture.get() - - # Reconstruct `tensors`, but with `tpu_tensors` replaced with - # `tpu_tensors_on_cpu`. - new_tensors = [] - for t in tensors: - if _is_tpu_tensor(t): - new_tensors.append(tpu_tensors_on_cpu.pop(0)) - elif t is None: - new_tensors.append(None) + export_outputs_list = nest.flatten(export_outputs_dict) + export_outputs_tpu_list = [ + t for t in export_outputs_list if t is not None + ] + + if isinstance(estimator_spec.predictions, dict): + predictions_dict = collections.OrderedDict( + (k, v) for k, v in six.iteritems(estimator_spec.predictions)) else: - # Only fetching `tpu_tensors_on_cpu` does not trigger - # TPU computation and blocks, so we add the control dependency here. - control_inputs = ( - tpu_tensors_on_cpu if isinstance(tpu_tensors_on_cpu, - (list, tuple)) else - (tpu_tensors_on_cpu,)) - with ops.control_dependencies(control_inputs): - new_tensors.append(array_ops.identity(t)) - - # Reconstruct `tensors_dict`. - new_tensors_dict = nest.pack_sequence_as(tensors_dict, new_tensors) - # Reconstruct `export_outputs`. - export_outputs = estimator_spec.export_outputs - new_export_outputs = collections.OrderedDict( - (k, _clone_export_output_with_tensors(export_outputs[k], v)) - for k, v in six.iteritems(new_tensors_dict)) + predictions_dict = { + _KEY_WHEN_PREDICTIONS_IS_A_TENSOR: estimator_spec.predictions + } + predictions_list = nest.flatten(predictions_dict) - return estimator_spec._replace(export_outputs=new_export_outputs) + # We cannot return everything we want through the return values, so + # capture the rest here for later use. + capture.capture((estimator_spec, export_outputs_dict, export_outputs_list, + predictions_dict)) + return predictions_list + export_outputs_tpu_list + + return computation, capture def _create_global_step(self, graph): """Creates a global step suitable for TPUs. @@ -2483,7 +2714,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): @@ -2512,7 +2747,7 @@ class TPUEstimator(estimator_lib.Estimator): graph.add_to_collection(_TPU_ENQUEUE_OPS, enqueue_op) if mode == model_fn_lib.ModeKeys.TRAIN: - loss, host_call, scaffold, training_hooks = ( + compile_op, loss, host_call, scaffold, training_hooks = ( _train_on_tpu_system(ctx, model_fn_wrapper, dequeue_fn)) host_ops = host_call.create_tpu_hostcall() if host_ops is None: @@ -2547,9 +2782,12 @@ class TPUEstimator(estimator_lib.Estimator): ctx, enqueue_ops, host_ops, + tpu_compile_op=compile_op, run_infeed_loop_on_coordinator=( run_infeed_loop_on_coordinator), rendezvous=self._rendezvous[mode], + master=self._config.master, + session_config=self._session_config, ), InstallSignalHandlerHook() ]) @@ -2602,8 +2840,8 @@ class TPUEstimator(estimator_lib.Estimator): scaffold=scaffold) if mode == model_fn_lib.ModeKeys.EVAL: - total_loss, host_calls, scaffold, eval_hooks = _eval_on_tpu_system( - ctx, model_fn_wrapper, dequeue_fn) + compile_op, total_loss, host_calls, scaffold, eval_hooks = ( + _eval_on_tpu_system(ctx, model_fn_wrapper, dequeue_fn)) iterations_per_loop_var = _create_or_get_iterations_per_loop() mean_loss = math_ops.div( total_loss, @@ -2650,10 +2888,13 @@ class TPUEstimator(estimator_lib.Estimator): ctx, enqueue_ops, eval_update_ops + host_ops, + tpu_compile_op=compile_op, run_infeed_loop_on_coordinator=( run_infeed_loop_on_coordinator), - rendezvous=self._rendezvous[mode]), - ] + input_hooks + rendezvous=self._rendezvous[mode], + master=self._config.evaluation_master, + session_config=self._session_config, + )] + input_hooks if eval_hooks: hooks.extend(eval_hooks) @@ -2668,7 +2909,7 @@ class TPUEstimator(estimator_lib.Estimator): # Predict assert mode == model_fn_lib.ModeKeys.PREDICT - (dummy_predict_op, host_calls, + (compile_op, dummy_predict_op, host_calls, scaffold, prediction_hooks) = _predict_on_tpu_system( ctx, model_fn_wrapper, dequeue_fn) with ops.control_dependencies([dummy_predict_op]): @@ -2724,7 +2965,10 @@ class TPUEstimator(estimator_lib.Estimator): hooks = [ _StoppingPredictHook(scalar_stopping_signal), TPUInfeedOutfeedSessionHookForPrediction( - ctx, enqueue_ops, host_ops, rendezvous=self._rendezvous[mode]), + ctx, enqueue_ops, host_ops, rendezvous=self._rendezvous[mode], + tpu_compile_op=compile_op, + master=self._config.master, + session_config=self._session_config), ] + input_hooks if prediction_hooks: @@ -2739,17 +2983,6 @@ class TPUEstimator(estimator_lib.Estimator): return _model_fn -def _is_tpu_tensor(tensor): - if not isinstance(tensor, ops.Tensor): - return False - try: - tensor.op.get_attr(tpu._OUTSIDE_COMPILATION_ATTR) # pylint: disable=protected-access - except ValueError: - return True - else: - return False - - def _export_output_to_tensors(export_output): """Get a list of `Tensors` used in `export_output`. @@ -2769,7 +3002,7 @@ def _export_output_to_tensors(export_output): elif isinstance(export_output, export_output_lib.RegressionOutput): return [export_output.value] elif isinstance(export_output, export_output_lib.PredictOutput): - return export_output.outputs.values() + return list(export_output.outputs.values()) else: raise ValueError( '`export_output` must be have type `ClassificationOutput`, ' @@ -2821,15 +3054,16 @@ def _eval_on_tpu_system(ctx, model_fn_wrapper, dequeue_fn): return training_loop.repeat(iterations_per_loop_var, single_tpu_eval_step, [_ZERO_LOSS]) - (loss,) = tpu.shard( + (compile_op, loss,) = tpu.split_compile_and_shard( multi_tpu_eval_steps_on_single_shard, inputs=[], num_shards=ctx.num_replicas, outputs_from_all_shards=False, device_assignment=ctx.device_assignment) + loss = loss[0] scaffold = _get_scaffold(captured_scaffold_fn) - return loss, host_calls, scaffold, captured_eval_hooks.get() + return compile_op, loss, host_calls, scaffold, captured_eval_hooks.get() def _train_on_tpu_system(ctx, model_fn_wrapper, dequeue_fn): @@ -2844,15 +3078,16 @@ def _train_on_tpu_system(ctx, model_fn_wrapper, dequeue_fn): return training_loop.repeat(iterations_per_loop_var, single_tpu_train_step, [_INITIAL_LOSS]) - (loss,) = tpu.shard( + (compile_op, loss,) = tpu.split_compile_and_shard( multi_tpu_train_steps_on_single_shard, inputs=[], num_shards=ctx.num_replicas, outputs_from_all_shards=False, device_assignment=ctx.device_assignment) + loss = loss[0] scaffold = _get_scaffold(captured_scaffold_fn) - return loss, host_call, scaffold, captured_training_hooks.get() + return compile_op, loss, host_call, scaffold, captured_training_hooks.get() def _predict_on_tpu_system(ctx, model_fn_wrapper, dequeue_fn): @@ -2872,15 +3107,17 @@ def _predict_on_tpu_system(ctx, model_fn_wrapper, dequeue_fn): cond, single_tpu_predict_step, inputs=inputs, name=b'loop') return outputs - (dummy_predict_op,) = tpu.shard( + (compile_op, dummy_predict_op,) = tpu.split_compile_and_shard( multi_tpu_predict_steps_on_single_shard, inputs=[], num_shards=ctx.num_replicas, outputs_from_all_shards=False, device_assignment=ctx.device_assignment) + dummy_predict_op = dummy_predict_op[0] scaffold = _get_scaffold(captured_scaffold_fn) - return dummy_predict_op, host_calls, scaffold, captured_predict_hooks.get() + return (compile_op, dummy_predict_op, host_calls, scaffold, + captured_predict_hooks.get()) def _wrap_computation_in_while_loop(device, op_fn): @@ -2999,6 +3236,12 @@ class _CapturingContext(control_flow_ops.ControlFlowContext): control_flow_ops.ControlFlowContext.__init__(self) self._message = message + def to_control_flow_context_def(self, context_def, export_scope=None): + # pylint: disable=useless-super-delegation + # NOTE(slebedev): the method is required by `ControlFlowContext`. + super(_CapturingContext, self).to_control_flow_context_def( + context_def, export_scope) + def AddOp(self, op): # pylint: disable=invalid-name for c in op.inputs: if tpu._TPU_REPLICATE_ATTR in c.op.node_def.attr: # pylint: disable=protected-access @@ -3039,7 +3282,7 @@ class _Inputs(object): @staticmethod def from_input_fn(return_values): """Returns an `_Inputs` instance according to `input_fn` return value.""" - if isinstance(return_values, dataset_ops.Dataset): + if isinstance(return_values, dataset_ops.DatasetV2): dataset = return_values return _Inputs(dataset=dataset) @@ -3064,7 +3307,7 @@ class _Inputs(object): The initializer must be run before calling `features_and_labels`. """ - self._iterator = self._dataset.make_initializable_iterator() + self._iterator = dataset_ops.make_initializable_iterator(self._dataset) return self._iterator.initializer def features_and_labels(self): diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_estimator_signals_test.py b/tensorflow/contrib/tpu/python/tpu/tpu_estimator_signals_test.py index 3786e52b949dfac8c1587d1ea3041b625f00183f..e3ea983abfd24d03c964fbc647b56262e15e0a96 100644 --- a/tensorflow/contrib/tpu/python/tpu/tpu_estimator_signals_test.py +++ b/tensorflow/contrib/tpu/python/tpu/tpu_estimator_signals_test.py @@ -21,8 +21,8 @@ from __future__ import print_function import numpy as np from tensorflow.contrib.tpu.python.tpu import tpu_estimator -from tensorflow.python import data as dataset_lib from tensorflow.python.client import session +from tensorflow.python.data.ops import dataset_ops from tensorflow.python.framework import errors from tensorflow.python.framework import ops from tensorflow.python.platform import test @@ -34,10 +34,10 @@ def make_input_fn(num_samples): def input_fn(params): batch_size = params['batch_size'] - da1 = dataset_lib.Dataset.from_tensor_slices(a) - da2 = dataset_lib.Dataset.from_tensor_slices(b) + da1 = dataset_ops.Dataset.from_tensor_slices(a) + da2 = dataset_ops.Dataset.from_tensor_slices(b) - dataset = dataset_lib.Dataset.zip((da1, da2)) + dataset = dataset_ops.Dataset.zip((da1, da2)) dataset = dataset.map(lambda fa, fb: {'a': fa, 'b': fb}) dataset = dataset.batch(batch_size) return dataset @@ -50,10 +50,10 @@ def make_input_fn_with_labels(num_samples): def input_fn(params): batch_size = params['batch_size'] - da1 = dataset_lib.Dataset.from_tensor_slices(a) - da2 = dataset_lib.Dataset.from_tensor_slices(b) + da1 = dataset_ops.Dataset.from_tensor_slices(a) + da2 = dataset_ops.Dataset.from_tensor_slices(b) - dataset = dataset_lib.Dataset.zip((da1, da2)) + dataset = dataset_ops.Dataset.zip((da1, da2)) dataset = dataset.map(lambda fa, fb: ({'a': fa}, fb)) dataset = dataset.batch(batch_size) return dataset @@ -71,7 +71,7 @@ class TPUEstimatorStoppingSignalsTest(test.TestCase): with ops.Graph().as_default(): dataset = input_fn(params) - features = dataset.make_one_shot_iterator().get_next() + features = dataset_ops.make_one_shot_iterator(dataset).get_next() # With tf.data.Dataset.batch, the batch is None, i.e., dynamic shape. self.assertIsNone(features['a'].shape.as_list()[0]) diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_feed.py b/tensorflow/contrib/tpu/python/tpu/tpu_feed.py index e75a09492ec12b95bad32b221a8e78a1b79f3a6b..d5957b7e8ec40b40c7af8822378cee6134ef0d0f 100644 --- a/tensorflow/contrib/tpu/python/tpu/tpu_feed.py +++ b/tensorflow/contrib/tpu/python/tpu/tpu_feed.py @@ -26,7 +26,6 @@ import numpy as np from six.moves import xrange # pylint: disable=redefined-builtin from tensorflow.compiler.xla.experimental.xla_sharding import xla_sharding -from tensorflow.compiler.xla.python_api import xla_shape from tensorflow.contrib.tpu.python.ops import tpu_ops from tensorflow.contrib.tpu.python.tpu import tpu from tensorflow.contrib.tpu.python.tpu import tpu_sharding @@ -92,8 +91,7 @@ class InfeedQueue(object): else: raise ValueError( "number of tuple elements cannot be inferred from InfeedQueue " - "constructor" - ) + "constructor") if number_of_tuple_elements <= 0: raise ValueError("number_of_tuple_elements %d must be > 0" % number_of_tuple_elements) @@ -293,9 +291,8 @@ class InfeedQueue(object): self.number_of_tuple_elements """ if len(input_tensors) != self.number_of_tuple_elements: - raise ValueError( - "input_tensors is %s, but should be a list of %d Tensors", ( - str(input_tensors), self.number_of_tuple_elements)) + raise ValueError("input_tensors is %s, but should be a list of %d Tensors" + % (str(input_tensors), self.number_of_tuple_elements)) self.set_tuple_shapes([t.shape for t in input_tensors]) self.set_tuple_types([t.dtype for t in input_tensors]) @@ -451,8 +448,8 @@ class InfeedQueue(object): for i in xrange(1, self.number_of_tuple_elements): if devices[0] != devices[i]: raise ValueError( - "input devices for shard %d are %s, but should all be the same", - index, str(devices)) + "input devices for shard %d are %s, but should all be the same" % + (index, str(devices))) with ops.colocate_with(inputs[0]): return tpu_ops.infeed_enqueue_tuple( inputs=inputs, @@ -792,18 +789,14 @@ class _PartitionedInfeedQueue(InfeedQueue): Args: tensor: Input tensor for partitioning. - dims: A list of integer describes how to partition the input tensor. + dims: 1-D np.array of the list of integer describes how to partition the + input tensor. Raises: ValueError: If the tensor can't be partitioned by dims or the num_cores_per_replica doesn't match the number of partitions(dims.prod()). """ - if dims is None: - return - - dims = np.array(dims) - if (dims < 1).any(): raise ValueError("All input partition dims must be >= 1.") @@ -823,11 +816,6 @@ class _PartitionedInfeedQueue(InfeedQueue): "partition dims = {}).".format(tensor.shape.as_list(), dims)) tensor.shape.assert_is_fully_defined() - if (np.array(tensor.shape.as_list()) % dims != 0).any(): - raise ValueError( - "All input partition dims must divide exactly into the `Tensor` " - "shape (tensor shape = {}, input partition dims = {}).".format( - tensor.shape.as_list(), dims)) def _partition_or_replicate_on_host(self, tensor, dims): """Partitions or replicates the input tensor. @@ -840,16 +828,39 @@ class _PartitionedInfeedQueue(InfeedQueue): Returns: An iterator of `Tensor`s or a list of partioned tensors. """ - self._check_input_partition_dims(tensor, dims) if dims is None: return itertools.repeat(tensor) - else: - output = [tensor] - for axis, dim in enumerate(dims): - if dim > 1: - output = [array_ops.split(x, dim, axis=axis) for x in output] - output = nest.flatten(output) - return output + dims = np.array(dims) + self._check_input_partition_dims(tensor, dims) + output = [tensor] + shape_list = np.array(tensor.shape.as_list()) + quotients, remainders = np.divmod(shape_list, dims) + for axis, (quotient, remainder, dim, original_size) in enumerate( + zip(quotients, remainders, dims, shape_list)): + if dim <= 1: + continue + if remainder > 0: + # For each dimension, when it cannot be evenly partitioned, XLA assumes + # tensors are partitioned in a greedy manner by using + # ceil_ratio(size/dim) first. E.g. 2D tensor with shape (5, 14) and dims + # are (2, 4). Since 5 % 2 = 1 and 14 % 4 = 2, [5, 14] => + # [[(3, 4), (3, 4), (2, 4), (2, 2)], + # [(2, 4), (2, 4), (2, 4), (2, 2)]] + ceil_ratio = quotient + 1 + num_full_slots, left_over = np.divmod(original_size, ceil_ratio) + num_or_size_splits = [ceil_ratio] * num_full_slots + [left_over] + if len(num_or_size_splits) < dim: + num_or_size_splits += [0] * (dim - len(num_or_size_splits)) + new_output = [] + for x in output: + new_output.append( + array_ops.split( + x, num_or_size_splits=num_or_size_splits, axis=axis)) + output = new_output + else: + output = [array_ops.split(x, dim, axis=axis) for x in output] + output = nest.flatten(output) + return output def _tag_sharding_attribute_for_dequeued_tensor(self, tensor, dims): """Tags appropriate XLA sharding attribute to the dequeued tensor. @@ -866,13 +877,9 @@ class _PartitionedInfeedQueue(InfeedQueue): elif np.prod(dims) == 1: return xla_sharding.assign_device(tensor, 0) else: - tile_shape = np.array(tensor.shape.as_list()) // dims tile_assignment = np.arange(np.prod(dims)).reshape(dims) return xla_sharding.tile( tensor=tensor, - tile_shape=xla_shape.CreateShapeFromDtypeAndTuple( - dtype=np.dtype(tensor.dtype.as_numpy_dtype), - shape_tuple=tile_shape), tile_assignment=tile_assignment) def _tag_sharding_attribute_for_dequeued_tensors(self, dequeues, dims): diff --git a/tensorflow/contrib/tpu/python/tpu/tpu_system_metadata.py b/tensorflow/contrib/tpu/python/tpu/tpu_system_metadata.py index ec682e5829c4df536a043334b74200f0b6259df3..d66ecfcf4a56b8da1c2d2f518bebe4baa76b315e 100644 --- a/tensorflow/contrib/tpu/python/tpu/tpu_system_metadata.py +++ b/tensorflow/contrib/tpu/python/tpu/tpu_system_metadata.py @@ -52,6 +52,7 @@ def _query_tpu_system_metadata(master_address, cluster_def=None, devices = [] device_dict = collections.defaultdict(list) + # TODO(b/120564445): Replace with standard library for retries. retry_count = 1 while True: logging.info('Querying Tensorflow master (%s) for TPU system metadata.', diff --git a/tensorflow/contrib/tpu/python/tpu/training_loop.py b/tensorflow/contrib/tpu/python/tpu/training_loop.py index b6c350ecd7588221b0e7bc979ed1be3b911c8cfd..0187b4bec6ecc55943bf48b9268a74e18ea5b488 100644 --- a/tensorflow/contrib/tpu/python/tpu/training_loop.py +++ b/tensorflow/contrib/tpu/python/tpu/training_loop.py @@ -166,8 +166,8 @@ def while_loop(condition, body, inputs=None, infeed_queue=None, name=None): # control dependencies from any side-effecting operations. if input_arity == 0: inputs = [array_ops.constant(0)] - return control_flow_ops.while_loop(condition_wrapper, body_wrapper, inputs, - name="") + return control_flow_ops.while_loop( + condition_wrapper, body_wrapper, inputs, name="", parallel_iterations=1) def repeat(n, body, inputs=None, infeed_queue=None, name=None): diff --git a/tensorflow/contrib/tpu/python/tpu/util.py b/tensorflow/contrib/tpu/python/tpu/util.py index b8ea307d8900cf1b6d1e6e808d0b9ede26f86490..dfb8ce1d1821da05c853bb0d10b1db3a857ccb1b 100644 --- a/tensorflow/contrib/tpu/python/tpu/util.py +++ b/tensorflow/contrib/tpu/python/tpu/util.py @@ -19,8 +19,11 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import time import six +from tensorflow.python.platform import tf_logging as logging +from tensorflow.python.training import training def check_positive_integer(value, name): """Checks whether `value` is a positive integer.""" @@ -29,3 +32,20 @@ def check_positive_integer(value, name): if value <= 0: raise ValueError('{} must be positive, got {}'.format(name, value)) + + +# TODO(b/118302029) Remove this copy of MultiHostDatasetInitializerHook after we +# release a tensorflow_estimator with MultiHostDatasetInitializerHook in +# python/estimator/util.py. +class MultiHostDatasetInitializerHook(training.SessionRunHook): + """Creates a SessionRunHook that initializes all passed iterators.""" + + def __init__(self, dataset_initializers): + self._initializers = dataset_initializers + + def after_create_session(self, session, coord): + del coord + start = time.time() + session.run(self._initializers) + logging.info('Initialized dataset iterators in %d seconds', + time.time() - start) diff --git a/tensorflow/contrib/tpu/tpu_estimator.md b/tensorflow/contrib/tpu/tpu_estimator.md index b6514e19dc92fe4c7cdcdb6582a7c0ad5ad573d5..552febd80bd35b37a95cdaaf8d5923278311ac8e 100644 --- a/tensorflow/contrib/tpu/tpu_estimator.md +++ b/tensorflow/contrib/tpu/tpu_estimator.md @@ -89,12 +89,9 @@ handle training: dataset = tf.data.TFRecordDataset( filename, buffer_size=FLAGS.dataset_reader_buffer_size) - dataset = dataset.map(parser).cache().repeat().batch(batch_size) - images, labels = dataset.make_one_shot_iterator().get_next() - # set_shape to give inputs statically known shapes. - images.set_shape([batch_size, 28 * 28]) - labels.set_shape([batch_size]) - return images, labels + dataset = dataset.map(parser).cache().repeat().batch( + batch_size, drop_remainder=True) + return dataset return input_fn 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..d98e0b7a5ed52c00a8cf2b1a1bbc53f1b1cd28c7 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 +// //third_party/tensorflow/contrib/tpu/proto/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/training/BUILD b/tensorflow/contrib/training/BUILD index 00295f57f60858db5234ce28cc643ea9eee44daa..f6427ae05a20f253edf030eff0f860361616042b 100644 --- a/tensorflow/contrib/training/BUILD +++ b/tensorflow/contrib/training/BUILD @@ -26,7 +26,6 @@ py_library( "python/training/resample.py", "python/training/sampling_ops.py", "python/training/sequence_queueing_state_saver.py", - "python/training/tensor_queue_dataset.py", "python/training/training.py", "python/training/tuner.py", ], @@ -287,28 +286,6 @@ py_test( ], ) -py_test( - name = "tensor_queue_dataset_test", - size = "large", - srcs = ["python/training/tensor_queue_dataset_test.py"], - srcs_version = "PY2AND3", - tags = ["notsan"], - deps = [ - ":training_py", - "//tensorflow/python:client_testlib", - "//tensorflow/python:framework_for_generated_wrappers", - "//tensorflow/python:gradients", - "//tensorflow/python:math_ops", - "//tensorflow/python:platform", - "//tensorflow/python:random_seed", - "//tensorflow/python:training", - "//tensorflow/python:variables", - "//tensorflow/python/data", - "//tensorflow/python/data/experimental/kernel_tests/serialization:dataset_serialization_test_base", - "//third_party/py/numpy", - ], -) - tf_proto_library( name = "protos_all", srcs = glob(["**/*.proto"]), diff --git a/tensorflow/contrib/training/__init__.py b/tensorflow/contrib/training/__init__.py index 3547e71184ec2b99163ea4247c01d24487811b47..87ce57ef060a0eb9383248255713421c14988416 100644 --- a/tensorflow/contrib/training/__init__.py +++ b/tensorflow/contrib/training/__init__.py @@ -59,8 +59,6 @@ from tensorflow.contrib.training.python.training.hparam import * from tensorflow.contrib.training.python.training.resample import * from tensorflow.contrib.training.python.training.sampling_ops import * from tensorflow.contrib.training.python.training.sequence_queueing_state_saver import * -from tensorflow.contrib.training.python.training.tensor_queue_dataset import enqueue_in_queue_dataset -from tensorflow.contrib.training.python.training.tensor_queue_dataset import prepend_from_queue_and_padded_batch_dataset from tensorflow.contrib.training.python.training.training import add_gradients_summaries from tensorflow.contrib.training.python.training.training import clip_gradient_norms from tensorflow.contrib.training.python.training.training import clip_gradient_norms_fn @@ -79,7 +77,6 @@ _allowed_symbols = [ 'FeedingQueueRunner', 'get_or_create_eval_step', 'StopAfterNEvalsHook', 'SummaryAtEndHook', 'wait_for_new_checkpoint', 'add_gradients_summaries', 'clip_gradient_norms', 'clip_gradient_norms_fn', 'create_train_op', - 'multiply_gradients', 'enqueue_in_queue_dataset', - 'prepend_from_queue_and_padded_batch_dataset', 'train'] + 'multiply_gradients', 'train'] remove_undocumented(__name__, _allowed_symbols) diff --git a/tensorflow/contrib/training/python/training/hparam.py b/tensorflow/contrib/training/python/training/hparam.py index 3beb7bfe3048a8f0294f7e9149b5a07b5fcc7d17..27f0d9b2e38c433d4fb4573285ecb8c9946112e8 100644 --- a/tensorflow/contrib/training/python/training/hparam.py +++ b/tensorflow/contrib/training/python/training/hparam.py @@ -187,7 +187,7 @@ def _cast_to_type_if_compatible(name, param_type, value): return param_type(value) -def parse_values(values, type_map): +def parse_values(values, type_map, ignore_unknown=False): """Parses hyperparameter values from a string into a python map. `values` is a string containing comma-separated `name=value` pairs. @@ -233,6 +233,9 @@ def parse_values(values, type_map): type T if either V has type T, or V is a list of elements of type T. Hence, for a multidimensional parameter 'x' taking float values, 'x=[0.1,0.2]' will parse successfully if type_map['x'] = float. + ignore_unknown: Bool. Whether values that are missing a type in type_map + should be ignored. If set to True, a ValueError will not be raised for + unknown hyperparameter type. Returns: A python map mapping each name to either: @@ -260,6 +263,8 @@ def parse_values(values, type_map): m_dict = m.groupdict() name = m_dict['name'] if name not in type_map: + if ignore_unknown: + continue raise ValueError('Unknown hyperparameter type for %s' % name) type_ = type_map[name] @@ -494,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] @@ -512,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. """ @@ -520,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(): @@ -543,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. @@ -552,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(): @@ -591,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. @@ -600,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/hparam_test.py b/tensorflow/contrib/training/python/training/hparam_test.py index 660c97f25e8458c345c8914bcaf98f37d047e50e..a990e04711ce68bd928a508484f0d6f657dd2f8c 100644 --- a/tensorflow/contrib/training/python/training/hparam_test.py +++ b/tensorflow/contrib/training/python/training/hparam_test.py @@ -216,6 +216,14 @@ class HParamsTest(test.TestCase): self.assertTrue(isinstance(parse_dict['arr'], dict)) self.assertDictEqual(parse_dict['arr'], {1: 10}) + def testParseValuesWithIndexAssigment1_IgnoreUnknown(self): + """Assignment to an index position.""" + parse_dict = hparam.parse_values( + 'arr[1]=10,b=5', {'arr': int}, ignore_unknown=True) + self.assertEqual(len(parse_dict), 1) + self.assertTrue(isinstance(parse_dict['arr'], dict)) + self.assertDictEqual(parse_dict['arr'], {1: 10}) + def testParseValuesWithIndexAssigment2(self): """Assignment to multiple index positions.""" parse_dict = hparam.parse_values('arr[0]=10,arr[5]=20', {'arr': int}) @@ -223,6 +231,14 @@ class HParamsTest(test.TestCase): self.assertTrue(isinstance(parse_dict['arr'], dict)) self.assertDictEqual(parse_dict['arr'], {0: 10, 5: 20}) + def testParseValuesWithIndexAssigment2_IgnoreUnknown(self): + """Assignment to multiple index positions.""" + parse_dict = hparam.parse_values( + 'arr[0]=10,arr[5]=20,foo=bar', {'arr': int}, ignore_unknown=True) + self.assertEqual(len(parse_dict), 1) + self.assertTrue(isinstance(parse_dict['arr'], dict)) + self.assertDictEqual(parse_dict['arr'], {0: 10, 5: 20}) + def testParseValuesWithIndexAssigment3(self): """Assignment to index positions in multiple names.""" parse_dict = hparam.parse_values('arr[0]=10,arr[1]=20,L[5]=100,L[10]=200', @@ -234,6 +250,17 @@ class HParamsTest(test.TestCase): self.assertTrue(isinstance(parse_dict['L'], dict)) self.assertDictEqual(parse_dict['L'], {5: 100, 10: 200}) + def testParseValuesWithIndexAssigment3_IgnoreUnknown(self): + """Assignment to index positions in multiple names.""" + parse_dict = hparam.parse_values( + 'arr[0]=10,C=5,arr[1]=20,B[0]=kkk,L[5]=100,L[10]=200', + {'arr': int, 'L': int}, ignore_unknown=True) + self.assertEqual(len(parse_dict), 2) + self.assertTrue(isinstance(parse_dict['arr'], dict)) + self.assertDictEqual(parse_dict['arr'], {0: 10, 1: 20}) + self.assertTrue(isinstance(parse_dict['L'], dict)) + self.assertDictEqual(parse_dict['L'], {5: 100, 10: 200}) + def testParseValuesWithIndexAssigment4(self): """Assignment of index positions and scalars.""" parse_dict = hparam.parse_values('x=10,arr[1]=20,y=30', @@ -246,6 +273,17 @@ class HParamsTest(test.TestCase): self.assertEqual(parse_dict['x'], 10) self.assertEqual(parse_dict['y'], 30) + def testParseValuesWithIndexAssigment4_IgnoreUnknown(self): + """Assignment of index positions and scalars.""" + parse_dict = hparam.parse_values( + 'x=10,foo[0]=bar,arr[1]=20,zzz=78,y=30', + {'x': int, 'y': int, 'arr': int}, ignore_unknown=True) + self.assertEqual(len(parse_dict), 3) + self.assertTrue(isinstance(parse_dict['arr'], dict)) + self.assertDictEqual(parse_dict['arr'], {1: 20}) + self.assertEqual(parse_dict['x'], 10) + self.assertEqual(parse_dict['y'], 30) + def testParseValuesWithIndexAssigment5(self): """Different variable types.""" parse_dict = hparam.parse_values('a[0]=5,b[1]=true,c[2]=abc,d[3]=3.14', { @@ -264,24 +302,55 @@ class HParamsTest(test.TestCase): self.assertTrue(isinstance(parse_dict['d'], dict)) self.assertDictEqual(parse_dict['d'], {3: 3.14}) + def testParseValuesWithIndexAssigment5_IgnoreUnknown(self): + """Different variable types.""" + parse_dict = hparam.parse_values( + 'a[0]=5,cc=4,b[1]=true,c[2]=abc,mm=2,d[3]=3.14', + {'a': int, 'b': bool, 'c': str, 'd': float}, + ignore_unknown=True) + self.assertEqual(set(parse_dict.keys()), {'a', 'b', 'c', 'd'}) + self.assertTrue(isinstance(parse_dict['a'], dict)) + self.assertDictEqual(parse_dict['a'], {0: 5}) + self.assertTrue(isinstance(parse_dict['b'], dict)) + self.assertDictEqual(parse_dict['b'], {1: True}) + self.assertTrue(isinstance(parse_dict['c'], dict)) + self.assertDictEqual(parse_dict['c'], {2: 'abc'}) + self.assertTrue(isinstance(parse_dict['d'], dict)) + self.assertDictEqual(parse_dict['d'], {3: 3.14}) + def testParseValuesWithBadIndexAssigment1(self): """Reject assignment of list to variable type.""" with self.assertRaisesRegexp(ValueError, r'Assignment of a list to a list index.'): hparam.parse_values('arr[1]=[1,2,3]', {'arr': int}) + def testParseValuesWithBadIndexAssigment1_IgnoreUnknown(self): + """Reject assignment of list to variable type.""" + with self.assertRaisesRegexp(ValueError, + r'Assignment of a list to a list index.'): + hparam.parse_values( + 'arr[1]=[1,2,3],c=8', {'arr': int}, ignore_unknown=True) + def testParseValuesWithBadIndexAssigment2(self): """Reject if type missing.""" with self.assertRaisesRegexp(ValueError, r'Unknown hyperparameter type for arr'): hparam.parse_values('arr[1]=5', {}) + def testParseValuesWithBadIndexAssigment2_IgnoreUnknown(self): + """Ignore missing type.""" + hparam.parse_values('arr[1]=5', {}, ignore_unknown=True) + def testParseValuesWithBadIndexAssigment3(self): """Reject type of the form name[index].""" with self.assertRaisesRegexp(ValueError, 'Unknown hyperparameter type for arr'): hparam.parse_values('arr[1]=1', {'arr[1]': int}) + def testParseValuesWithBadIndexAssigment3_IgnoreUnknown(self): + """Ignore type of the form name[index].""" + hparam.parse_values('arr[1]=1', {'arr[1]': int}, ignore_unknown=True) + def testWithReusedVariables(self): with self.assertRaisesRegexp(ValueError, 'Multiple assignments to variable \'x\''): diff --git a/tensorflow/contrib/training/python/training/sampling_ops.py b/tensorflow/contrib/training/python/training/sampling_ops.py index 7140f2a46d57f0f3b76ff4f1ea9d0d73808405c8..849b77d60954dc91726e261a0523943d704e5d21 100644 --- a/tensorflow/contrib/training/python/training/sampling_ops.py +++ b/tensorflow/contrib/training/python/training/sampling_ops.py @@ -19,6 +19,7 @@ from __future__ import print_function 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 check_ops from tensorflow.python.ops import control_flow_ops @@ -300,10 +301,11 @@ def _verify_data_inputs(tensor_list): """Verify that batched data inputs are well-formed.""" for tensor in tensor_list: # Data tensor should have a batch dimension. - tensor_shape = tensor.get_shape().with_rank_at_least(1) + shape = tensor.get_shape().with_rank_at_least(1) # Data batch dimensions must be compatible. - tensor_shape[0].assert_is_compatible_with(tensor_list[0].get_shape()[0]) + tensor_shape.dimension_at_index(shape, 0).assert_is_compatible_with( + tensor_list[0].get_shape()[0]) return tensor_list @@ -340,10 +342,11 @@ def _verify_input(tensor_list, labels, probs_list): for tensor in tensor_list: # Data tensor should have a batch dimension. - tensor_shape = tensor.get_shape().with_rank_at_least(1) + shape = tensor.get_shape().with_rank_at_least(1) # Data and label batch dimensions must be compatible. - tensor_shape[0].assert_is_compatible_with(labels.get_shape()[0]) + tensor_shape.dimension_at_index(shape, 0).assert_is_compatible_with( + labels.get_shape()[0]) # Data and labels must have the same, strictly positive batch size. Since we # can't assume we know the batch size at graph creation, add runtime checks. diff --git a/tensorflow/contrib/training/python/training/tensor_queue_dataset.py b/tensorflow/contrib/training/python/training/tensor_queue_dataset.py deleted file mode 100644 index 8896a95327a4cb609a9a78412afa68b316a3131e..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/training/python/training/tensor_queue_dataset.py +++ /dev/null @@ -1,201 +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. -# ============================================================================== -"""Python wrappers for Datasets and Iterators.""" -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -from tensorflow.python.data.ops import dataset_ops -from tensorflow.python.data.util import convert -from tensorflow.python.data.util import nest -from tensorflow.python.data.util import sparse -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.ops import gen_dataset_ops -from tensorflow.python.util import nest as tf_nest - - -class _PrependFromQueueAndPaddedBatchDataset(dataset_ops.UnaryDataset): - """A `Dataset` that prepends a queue to another `Dataset`. - - A vector of handles to the queue is returned as the first component of - the associated iterator. This vector can be passed to - `enqueue_in_queue_dataset` to add new elements to the queue. - """ - - def __init__(self, input_dataset, batch_size, padded_shapes, padding_values): - """Initialize `PrependFromQueueAndPaddedBatchDataset`.""" - super(_PrependFromQueueAndPaddedBatchDataset, self).__init__(input_dataset) - if sparse.any_sparse(input_dataset.output_classes): - raise TypeError( - "Batching of padded sparse tensors is not currently supported") - self._input_dataset = input_dataset - self._batch_size = ops.convert_to_tensor( - batch_size, dtype=dtypes.int64, name="batch_size") - if padded_shapes is None: - self._padded_shapes = nest.map_structure( - convert.partial_shape_to_tensor, input_dataset.output_shapes) - else: - self._padded_shapes = nest.map_structure_up_to( - input_dataset.output_shapes, convert.partial_shape_to_tensor, - padded_shapes) - # pylint: disable=protected-access - padding_values = ( - padding_values if padding_values is not None else - dataset_ops._default_padding(input_dataset)) - self._padding_values = nest.map_structure_up_to( - input_dataset.output_shapes, dataset_ops._padding_value_to_tensor, - padding_values, input_dataset.output_types) - # pylint: enable=protected-access - - def _as_variant_tensor(self): - # pylint: disable=protected-access - return gen_dataset_ops.prepend_from_queue_and_padded_batch_dataset( - self._input_dataset._as_variant_tensor(), - batch_size=self._batch_size, - padded_shapes=[ - ops.convert_to_tensor(s, dtype=dtypes.int64) - for s in nest.flatten(self._padded_shapes) - ], - padding_values=nest.flatten(self._padding_values), - output_shapes=nest.flatten( - sparse.as_dense_shapes(self.output_shapes, self.output_classes))) - # pylint: enable=protected-access - - @property - def output_classes(self): - return (ops.Tensor, self._input_dataset.output_classes) - - def _as_batch_shape(self, shape_like): - return tensor_shape.vector(None).concatenate( - tensor_util.constant_value_as_shape(shape_like)) - - @property - def output_shapes(self): - # First output is a variant representing the Queue - return (tensor_shape.vector(None), - nest.map_structure(self._as_batch_shape, self._padded_shapes)) - - @property - def output_types(self): - # First output is a variant representing the Queue - return (dtypes.variant, self._input_dataset.output_types) - - -def prepend_from_queue_and_padded_batch_dataset(batch_size, - padding_values=None, - padded_shapes=None): - """A transformation that prepends a queue to a `Dataset` and batches results. - - A vector of handles to the queue is returned as the first component of the - associated iterator. This vector can be passed to `enqueue_in_queue_dataset` - to add new elements to the queue. - - Below is an example of how this dataset might be used to split incoming - variable-length sequences into "head" and "rest" parts, where "rest" parts - are re-enqueued back into the dataset. A more realistic example would - perform some calculation on the "head" and modify some components of "rest" - with the result (before re-enqueueing). - - ```python - dataset = tf.data.Dataset.from_tensor_slices([2*x for x in range(10)]) - # Make a dataset of variable-length vectors and their lengths. - dataset = dataset.map(lambda count: (count, tf.ones((count,)))) - # Emit a queue we can prepend to, and counts/values as padded batch. - dataset = dataset.apply( - tf.contrib.training.prepend_from_queue_and_padded_batch_dataset( - batch_size=10)) - dataset = dataset.prefetch(1) - - iterator = dataset.make_one_shot_iterator() - queue, (count, padded_value) = iterator.get_next() - - # Split the padded_value into two pieces: head and rest - rest_indices = tf.squeeze(tf.where(count > 3), axis=1) - bound = tf.minimum(3, tf.reduce_max(count)) - value_head = padded_value[:, :bound] - count_rest = tf.gather(count - 3, rest_indices) - value_rest = tf.gather(padded_value[:, bound:], rest_indices) - queue_rest = tf.gather(queue, rest_indices) - enqueue_rest_op = tf.contrib.training.enqueue_in_queue_dataset( - queue_rest, (count_rest, value_rest)) - with tf.control_dependencies([enqueue_rest_op]): - calculation = fn(value_head) - - while True: # Will raise OutOfRange when finished with all pieces. - session.run(calculation) - ``` - - Args: - batch_size: `int64` scalar tensor. The batch size to use when performing - padded batching. - padding_values: (optional) Nested tuple of scalar tensors. If provided, - the structure and dtypes of padding_values should match that of - incoming dataset's `output_types`. - padded_shapes: (optional) Nested tuple of `int64` vector tensors. - If provided, the structure must match that of the incoming dataset's - `output_types`. If not provided, the incoming dataset's `output_shapes` - is used. Any unknown (`None` or `-1`) dimensions in the shapes are - treated as being unique per-batch: for each batch time, an unknown - dimension is replaced with the maximum given value of this dimension - across all tensors for the given component in the batch. - - Returns: - A `Dataset` transformation function, which can be passed to - `tf.data.Dataset.apply`. - """ - - def _apply_fn(dataset): - return _PrependFromQueueAndPaddedBatchDataset( - dataset, - batch_size=batch_size, - padding_values=padding_values, - padded_shapes=padded_shapes) - - return _apply_fn - - -def enqueue_in_queue_dataset(queue, components): - """Enqueue components into queue from `PrependFromQueueAndPaddedBatchDataset`. - - The components' dtypes and shapes must be compatible with the `output_shapes` - attribute of the `dataset` created by - `prepend_from_queue_and_padded_batch_dataset`. This operation supports both - non-batched and batched modes. - - For more details, see the example in the docstring for - `prepend_from_queue_and_padded_batch_dataset`. - - Args: - queue: `variant` scalar or vector tensor. - The tensor emitted by the first component of the iterator associated with - `prepend_from_queue_and_padded_batch_dataset`. If this is a scalar, - then the `components` input tensors should not have a prepended batch - dimension. - components: Nested tuple of tensors, each with a leading batch dimension - if `queue` is a vector. The structure, dtypes, and shapes - (excluding batch dimension) must match the nested tuples - `dataset.output_types[1]` and `dataset.output_shapes[1]` (the non-queue - output types and shapes) of the `dataset` emitted by - the original `prepend_from_queue_and_padded_batch_dataset` call. - - Returns: - An `Operation` that enqueues `components` into the dataset(s) associated - with entries of `queue`. - """ - return gen_dataset_ops.enqueue_in_queue_dataset( - queue=queue, components=tf_nest.flatten(components)) diff --git a/tensorflow/contrib/training/python/training/tensor_queue_dataset_test.py b/tensorflow/contrib/training/python/training/tensor_queue_dataset_test.py deleted file mode 100644 index c1657fec7bbe4a3227c3ea273b72176ac4066c50..0000000000000000000000000000000000000000 --- a/tensorflow/contrib/training/python/training/tensor_queue_dataset_test.py +++ /dev/null @@ -1,355 +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. -# ============================================================================== -"""Tests for TensorQueueDataset.""" -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import numpy as np - -from tensorflow.contrib.training.python.training import tensor_queue_dataset as tqd -from tensorflow.python.data.experimental.kernel_tests.serialization import dataset_serialization_test_base -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 ops -from tensorflow.python.ops import array_ops -from tensorflow.python.ops import math_ops -from tensorflow.python.ops import string_ops -from tensorflow.python.platform import test - - -class PrependFromQueueAndPaddedBatchDatasetTest(test.TestCase): - - def testNoEnqueue(self): - dataset = dataset_ops.Dataset.from_tensor_slices([0, 1, 2]) - dataset = dataset.apply( - tqd.prepend_from_queue_and_padded_batch_dataset(batch_size=1)) - self.assertEqual((dtypes.variant, dtypes.int32), dataset.output_types) - self.assertAllEqual(([None],) * 2, - [x.as_list() for x in dataset.output_shapes]) - iterator = dataset.make_one_shot_iterator() - _, value = iterator.get_next() - self.assertEqual([0], self.evaluate(value)) - self.assertEqual([1], self.evaluate(value)) - self.assertEqual([2], self.evaluate(value)) - with self.assertRaisesOpError("End of sequence"): - self.evaluate(value) - - def testBatchedNoEnqueue(self): - dataset = dataset_ops.Dataset.from_tensor_slices([0, 1, 2]) - dataset = dataset.apply( - tqd.prepend_from_queue_and_padded_batch_dataset(batch_size=2)) - iterator = dataset.make_one_shot_iterator() - _, value = iterator.get_next() - self.assertAllEqual([0, 1], self.evaluate(value)) - self.assertAllEqual([2], self.evaluate(value)) - with self.assertRaisesOpError("End of sequence"): - self.evaluate(value) - - def testBatchedWithBiggerPaddingNoEnqueue(self): - dataset = dataset_ops.Dataset.from_tensor_slices([[0], [1], [2]]) - dataset = dataset.apply( - tqd.prepend_from_queue_and_padded_batch_dataset( - batch_size=2, padded_shapes=[3])) - iterator = dataset.make_one_shot_iterator() - _, value = iterator.get_next() - self.assertAllEqual([[0, 0, 0], [1, 0, 0]], self.evaluate(value)) - self.assertAllEqual([[2, 0, 0]], self.evaluate(value)) - with self.assertRaisesOpError("End of sequence"): - self.evaluate(value) - - def testBatchedWithBiggerPaddingOneEnqueue(self): - dataset = dataset_ops.Dataset.from_tensor_slices([[0], [1], [2]]) - dataset = dataset.apply( - tqd.prepend_from_queue_and_padded_batch_dataset( - batch_size=1, padded_shapes=[3])) - iterator = dataset.make_one_shot_iterator() - queue_handle, value = iterator.get_next() - enqueue_negative = tqd.enqueue_in_queue_dataset(queue_handle, -value) - with self.cached_session() as sess: - self.assertAllEqual([[0, 0, 0]], sess.run(value)) - value_1, _ = sess.run([value, enqueue_negative]) - self.assertAllEqual([[1, 0, 0]], value_1) - value_2, _ = sess.run([value, enqueue_negative]) - self.assertAllEqual([[-1, 0, 0]], value_2) - value_3 = sess.run(value) - self.assertAllEqual([[1, 0, 0]], value_3) - value_4, _ = sess.run([value, enqueue_negative]) - self.assertAllEqual([[2, 0, 0]], value_4) - value_5 = sess.run(value) - self.assertAllEqual([[-2, 0, 0]], value_5) - with self.assertRaisesOpError("End of sequence"): - sess.run(value) - - def testOneEnqueue(self): - dataset = dataset_ops.Dataset.from_tensor_slices([0, 1, 2]) - dataset = dataset.apply( - tqd.prepend_from_queue_and_padded_batch_dataset(batch_size=1)) - iterator = dataset.make_one_shot_iterator() - queue_handle, value = iterator.get_next() - enqueue_negative = tqd.enqueue_in_queue_dataset(queue_handle, -value) - with self.cached_session() as sess: - self.assertEqual([0], sess.run(value)) - value_1, _ = sess.run([value, enqueue_negative]) - self.assertEqual([1], value_1) - value_2, _ = sess.run([value, enqueue_negative]) - self.assertEqual([-1], value_2) - value_3 = sess.run(value) - self.assertEqual([1], value_3) - value_4, _ = sess.run([value, enqueue_negative]) - self.assertEqual([2], value_4) - value_5 = sess.run(value) - self.assertEqual([-2], value_5) - with self.assertRaisesOpError("End of sequence"): - sess.run(value) - - def testBatchedOneEnqueue(self): - dataset = dataset_ops.Dataset.from_tensor_slices([0, 1, 2]) - dataset = dataset.apply( - tqd.prepend_from_queue_and_padded_batch_dataset(batch_size=2)) - iterator = dataset.make_one_shot_iterator() - queue_handle, value = iterator.get_next() - enqueue_negative = tqd.enqueue_in_queue_dataset(queue_handle, -value) - enqueue_zeroth = tqd.enqueue_in_queue_dataset([queue_handle[0]], - array_ops.expand_dims( - value[0], axis=0)) - with self.cached_session() as sess: - value_0, _ = sess.run([value, enqueue_negative]) - self.assertAllEqual([0, 1], value_0) - value_1, _ = sess.run([value, enqueue_zeroth]) - self.assertAllEqual([0, -1], value_1) - value_2, _ = sess.run([value, enqueue_negative]) - self.assertAllEqual([0, 2], value_2) - self.assertAllEqual([0, -2], sess.run(value)) - with self.assertRaisesOpError("End of sequence"): - sess.run(value) - - def testManyEnqueue(self): - dataset = dataset_ops.Dataset.from_tensor_slices([0, 1]) - dataset = dataset.apply( - tqd.prepend_from_queue_and_padded_batch_dataset(batch_size=1)) - iterator = dataset.make_one_shot_iterator() - queue_handle, value = iterator.get_next() - enqueue_many_more = [ - tqd.enqueue_in_queue_dataset(queue_handle, value + 100 + i) - for i in range(1000) - ] - with self.cached_session() as sess: - value_0, _ = sess.run((value, enqueue_many_more)) - self.assertEqual([0], value_0) - rest = [] - for _ in range(1000): - rest.append(sess.run(value)) - self.assertEquals([[100 + i] for i in range(1000)], sorted(rest)) - # Going back to the original input. - value_1, _ = sess.run((value, enqueue_many_more)) - self.assertEqual(1, value_1) - rest = [] - for _ in range(1000): - rest.append(sess.run(value)) - self.assertEquals([[100 + i + 1] for i in range(1000)], sorted(rest)) - with self.assertRaisesOpError("End of sequence"): - sess.run(value) - - def testEnqueueWithPrefetch(self): - dataset = dataset_ops.Dataset.from_tensor_slices([0]) - dataset = dataset.apply( - tqd.prepend_from_queue_and_padded_batch_dataset(batch_size=1)) - # Prefetching will request additional values before they are - # available to the queue. - dataset = dataset.prefetch(buffer_size=3) - iterator = dataset.make_one_shot_iterator() - queue_handle, value = iterator.get_next() - enqueue = tqd.enqueue_in_queue_dataset(queue_handle, value + 1) - with self.cached_session() as sess: - i = 0 - while i < 4: - received, _ = sess.run((value, enqueue)) - if received.size > 0: - self.assertAllEqual([i], received) - i += 1 - received_last = False - while True: - try: - received = sess.run(value) - if received.size > 0: - self.assertAllEqual([4], received) - received_last = True - except errors.OutOfRangeError: - break - self.assertTrue(received_last) - - def testDatasetWithPaddedShapeSmallerThanInputFails(self): - dataset = dataset_ops.Dataset.from_tensor_slices([[0, 0, 0]]).repeat(None) - dataset = dataset.apply( - tqd.prepend_from_queue_and_padded_batch_dataset( - batch_size=1, padded_shapes=[2])) - iterator = dataset.make_one_shot_iterator() - _, value = iterator.get_next() - with self.cached_session() as sess: - with self.assertRaisesOpError( - r"Incompatible input shapes at component 0 between " - r"input dataset this dataset: \[3\] vs. \[2\]"): - sess.run(value) - - def testEnqueueWithIncompatibleInputsFailsWithInformativeError(self): - dataset = dataset_ops.Dataset.from_tensor_slices([0]).repeat(None) - dataset = dataset.apply( - tqd.prepend_from_queue_and_padded_batch_dataset(batch_size=1)) - iterator = dataset.make_one_shot_iterator() - queue_handle, value = iterator.get_next() - - enqueue_bad_structure = tqd.enqueue_in_queue_dataset( - queue_handle, (value, value)) - enqueue_bad_dtype = tqd.enqueue_in_queue_dataset(queue_handle, - np.array( - [1.0], - dtype=np.float32)) - enqueue_bad_shape_no_batch_dim = tqd.enqueue_in_queue_dataset( - queue_handle, ([1],)) - enqueue_bad_shape = tqd.enqueue_in_queue_dataset(queue_handle, - np.array( - [[1]], dtype=np.int32)) - - with self.cached_session() as sess: - with self.assertRaisesOpError( - "mismatched number of tensors. Queue expects 1 tensors but " - "tried to insert 2"): - sess.run(enqueue_bad_structure) - with self.assertRaisesOpError(r"Expected component 0 to have batched " - r"shape \[1,...\], but saw shape: \[\]"): - sess.run(enqueue_bad_shape_no_batch_dim) - with self.assertRaisesOpError( - r"mismatched shapes at component 0. Attempted to insert tensor " - r"with shape \[1\] but queue expected shape: \[\]"): - sess.run(enqueue_bad_shape) - with self.assertRaisesOpError( - r"mismatched dtypes at component 0. Attempted to insert tensor " - r"of type float but queue expected type: int32"): - sess.run(enqueue_bad_dtype) - - def testEnqueueWithPaddedBatchFailsWithInformativeError(self): - dataset = dataset_ops.Dataset.from_tensor_slices([0, 1, 2]) - dataset = dataset.apply( - tqd.prepend_from_queue_and_padded_batch_dataset(batch_size=1)) - with self.assertRaisesRegexp( - TypeError, r"Unable to create padding for field of type 'variant'"): - dataset.padded_batch(batch_size=10, padded_shapes=[1]) - - def testOneEnqueueWithPadding(self): - dataset = dataset_ops.Dataset.from_tensor_slices([0, 2, 4, 6]) - # Make a dataset of variable-length vectors and their lengths. - dataset = dataset.map( - lambda c: (c, c * array_ops.ones((c,), dtype=c.dtype))) - # Emit a queue we can prepend to, and counts/values as padded - # batch. - dataset = dataset.apply( - tqd.prepend_from_queue_and_padded_batch_dataset(batch_size=3)) - - iterator = dataset.make_one_shot_iterator() - queue, (count, padded_value) = iterator.get_next() - - # Split the padded_value into two pieces: head and rest - rest_indices = array_ops.squeeze(array_ops.where(count > 2), axis=1) - bound = math_ops.minimum(2, math_ops.reduce_max(count)) - value_head = padded_value[:, :bound] - count_rest = array_ops.gather(count - 2, rest_indices) - value_rest = array_ops.gather(padded_value, rest_indices)[:, bound:] - queue_rest = array_ops.gather(queue, rest_indices) - enqueue_rest_op = tqd.enqueue_in_queue_dataset(queue_rest, - (count_rest, value_rest)) - with ops.control_dependencies([enqueue_rest_op]): - calc = array_ops.identity(value_head) - - with self.cached_session() as sess: - self.assertAllEqual([[0, 0], [2, 2], [4, 4]], sess.run(calc)) - self.assertAllEqual([[4, 4], [6, 6]], sess.run(calc)) - self.assertAllEqual([[6, 6]], sess.run(calc)) - self.assertAllEqual([[6, 6]], sess.run(calc)) - # Get some final batches due to prefetching. - for _ in range(3): - try: - self.assertAllEqual( - np.empty(shape=(0, 0), dtype=np.int32), sess.run(calc)) - except errors.OutOfRangeError as e: - self.assertTrue(str(e).startswith("End of sequence")) - - def testNonstandardPadding(self): - dataset = dataset_ops.Dataset.from_tensor_slices([0, 2, 4, 6]) - # Make a dataset of variable-length vectors and their lengths. - dataset = dataset.map( - lambda c: (c, c * array_ops.ones((c,), dtype=c.dtype))) - # Emit a queue we can prepend to, and counts/values as padded - # batch. - dataset = dataset.apply( - tqd.prepend_from_queue_and_padded_batch_dataset( - batch_size=3, padding_values=( - 0, - -1, - ))) - - iterator = dataset.make_one_shot_iterator() - _, (unused_count, padded_value) = iterator.get_next() - - with self.cached_session() as sess: - self.assertAllEqual([[-1, -1, -1, -1], [2, 2, -1, -1], [4, 4, 4, 4]], - sess.run(padded_value)) - self.assertAllEqual([[6] * 6], sess.run(padded_value)) - with self.assertRaisesOpError("End of sequence"): - sess.run(padded_value) - - -# TODO(ebrevdo): Figure out how to use run_core_tests to test state -# saving of an iterator that's had some tensors enqueued into its queue. -class PrependFromQueueAndPaddedBatchDatasetSerializationTest( - dataset_serialization_test_base.DatasetSerializationTestBase): - - def testPrependFromQueueAndPaddedBatch(self): - - def build_dataset(seq_lens): - return dataset_ops.Dataset.from_tensor_slices(seq_lens).map( - lambda x: array_ops.fill([x], x)).apply( - tqd.prepend_from_queue_and_padded_batch_dataset(batch_size=4)) - - seq_lens1 = np.random.randint(1, 20, size=(32,)).astype(np.int32) - seq_lens2 = np.random.randint(21, 40, size=(32,)).astype(np.int32) - self.run_core_tests(lambda: build_dataset(seq_lens1), - lambda: build_dataset(seq_lens2), 8) - - def testPrependFromQueueAndPaddedBatchNonDefaultPadding(self): - - def build_dataset(seq_lens): - - def fill_tuple(x): - filled = array_ops.fill([x], x) - return (filled, string_ops.as_string(filled)) - - padded_shape = [-1] - return dataset_ops.Dataset.from_tensor_slices(seq_lens).map( - fill_tuple).apply( - tqd.prepend_from_queue_and_padded_batch_dataset( - batch_size=4, - padded_shapes=(padded_shape, padded_shape), - padding_values=(-1, ""))) - - seq_lens1 = np.random.randint(1, 20, size=(32,)).astype(np.int32) - seq_lens2 = np.random.randint(21, 40, size=(32,)).astype(np.int32) - self.run_core_tests(lambda: build_dataset(seq_lens1), - lambda: build_dataset(seq_lens2), 8) - - -if __name__ == "__main__": - test.main() diff --git a/tensorflow/contrib/training/python/training/training.py b/tensorflow/contrib/training/python/training/training.py index c272a2ac144068cfb7355c2647eebf5bd0ce9d50..fc6e38ab4a5243cb7502f4ca42db03cbfd342a40 100644 --- a/tensorflow/contrib/training/python/training/training.py +++ b/tensorflow/contrib/training/python/training/training.py @@ -419,7 +419,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: diff --git a/tensorflow/contrib/training/python/training/tuner.py b/tensorflow/contrib/training/python/training/tuner.py index 8843632619f0881f888ca76c9de484f081786b19..ad647a61da7adb549ec21a940bf8682b722353b5 100644 --- a/tensorflow/contrib/training/python/training/tuner.py +++ b/tensorflow/contrib/training/python/training/tuner.py @@ -21,9 +21,12 @@ from __future__ import print_function import abc +import six + from tensorflow.contrib.framework.python.framework import experimental +@six.add_metaclass(abc.ABCMeta) class Tuner(object): """Tuner class is the interface for Experiment hyper-parameters tuning. @@ -42,8 +45,6 @@ class Tuner(object): learn_runner.tune(experiment_fn=_create_my_experiment, tuner) """ - __metaclass__ = abc.ABCMeta - @experimental @abc.abstractmethod def next_trial(self): 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/util/__init__.py b/tensorflow/contrib/util/__init__.py index 338acef63f244613cbd14a2da04c7ec4d811a0af..acc5a049aa87649e4f8bf3a00be605616ea7b630 100644 --- a/tensorflow/contrib/util/__init__.py +++ b/tensorflow/contrib/util/__init__.py @@ -15,8 +15,6 @@ """Utilities for dealing with Tensors. -See [Contrib Util](https://tensorflow.org/api_guides/python/contrib.util) guide. - @@constant_value @@make_tensor_proto @@make_ndarray diff --git a/tensorflow/contrib/verbs/rdma.cc b/tensorflow/contrib/verbs/rdma.cc index f7c979e86320d59ad033e2b8d7fcdff89ce0d133..9db80f6b5736d849d88e1e41ea467a5ff11844f5 100644 --- a/tensorflow/contrib/verbs/rdma.cc +++ b/tensorflow/contrib/verbs/rdma.cc @@ -30,7 +30,6 @@ limitations under the License. #include "tensorflow/core/distributed_runtime/rendezvous_mgr_interface.h" #include "tensorflow/core/distributed_runtime/rpc/grpc_util.h" #include "tensorflow/core/distributed_runtime/session_mgr.h" -#include "tensorflow/core/distributed_runtime/rpc/grpc_util.h" #include "tensorflow/core/framework/rendezvous.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/lib/core/status.h" @@ -1028,7 +1027,10 @@ Status RdmaTensorResponse::PrepareRecvTensor( return errors::Aborted( "RecvTensor expects a different device incarnation: ", parsed.src_incarnation, " vs. ", (*src_dev)->attributes().incarnation(), - ". Your worker job was probably restarted. Check your " + ". Your worker job (\"", + channel_->adapter_->worker_env_->session_mgr->LegacySession() + ->worker_name, + "\") was probably restarted. Check your " "worker job for the reason why it was restarted."); } diff --git a/tensorflow/contrib/verbs/rdma_mgr.cc b/tensorflow/contrib/verbs/rdma_mgr.cc index 2784bf124ceaacd8e01f0653287fa7f006d0d608..2f2375427862ad1e99a0e6bfc506382d200e9b1d 100644 --- a/tensorflow/contrib/verbs/rdma_mgr.cc +++ b/tensorflow/contrib/verbs/rdma_mgr.cc @@ -277,9 +277,18 @@ void RdmaMgr::InitAllocators() { ProcessState::singleton()->AddCPUFreeVisitor(free_visitor); #if GOOGLE_CUDA + GPUProcessState::singleton()->AddCUDAHostAllocVisitor(0, alloc_visitor); + GPUProcessState::singleton()->AddCUDAHostFreeVisitor(0, free_visitor); + if (IsGDRAvailable()) { // Note we don't free allocated GPU memory so there is no free visitor - int32_t bus_id = TryToReadNumaNode(rdma_adapter_->context_->device) + 1; + + // TODO: This is to fix the 'invalid use of member in static member function + // bug'. + // Waiting for better implementation. + // int32_t bus_id = TryToReadNumaNode(rdma_adapter_->context_->device) + // + 1; + int32_t bus_id = 0; SubAllocator::Visitor cuda_alloc_visitor = [](void* ptr, int gpu_id, size_t num_bytes) { @@ -288,9 +297,6 @@ void RdmaMgr::InitAllocators() { }; GPUProcessState::singleton()->AddGPUAllocVisitor(bus_id, cuda_alloc_visitor); - GPUProcessState::singleton()->AddCUDAHostAllocVisitor(bus_id, - alloc_visitor); - GPUProcessState::singleton()->AddCUDAHostFreeVisitor(bus_id, free_visitor); LOG(INFO) << "Instrumenting GPU allocator with bus_id " << bus_id; } #endif // GOOGLE_CUDA diff --git a/tensorflow/contrib/verbs/verbs_server_lib.cc b/tensorflow/contrib/verbs/verbs_server_lib.cc index 5b72b1604aca2e0c593978c6104322372788eb3c..19ef109f671ee57ce2aceb55110c50aa44352223 100644 --- a/tensorflow/contrib/verbs/verbs_server_lib.cc +++ b/tensorflow/contrib/verbs/verbs_server_lib.cc @@ -33,6 +33,8 @@ RendezvousMgrInterface* NewRdmaRendezvousMgr(const WorkerEnv* env) { return new RdmaRendezvousMgr(env); } +std::once_flag reg_mem_visitors_call; + } // namespace VerbsServer::VerbsServer(const ServerDef& server_def, Env* env) @@ -76,10 +78,6 @@ Status VerbsServer::ChannelCacheFactory(const ServerDef& server_def, return Status::OK(); } -namespace { -std::once_flag reg_mem_visitors_call; -} // namespace - Status VerbsServer::Init(ServiceInitFunction service_func, RendezvousMgrCreationFunction rendezvous_mgr_func) { std::call_once(reg_mem_visitors_call, []() { RdmaMgr::RegMemVisitors(); }); diff --git a/tensorflow/core/BUILD b/tensorflow/core/BUILD index 4f95f207ad8c3d9e283c24565dbd8b319f8119c7..7652f03a9f62e5918032c5284b453269ac33aba5 100644 --- a/tensorflow/core/BUILD +++ b/tensorflow/core/BUILD @@ -49,7 +49,7 @@ # filegroup ":android_proto_srcs" - Protos # filegroup ":android_srcs" - Core sources # cc_library ":android_tensorflow_lib" - Native library -# cc_library ":android_tensorflow_lib_selective_registration" - Native library +# cc_library ":android_tensorflow_lib_lite" - Native library, without ops, # supporting SELECTIVE_REGISTRATION feature. # portable_proto_library ":android_proto_lib" (Google-internal) # @@ -70,6 +70,9 @@ 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", @@ -95,6 +98,8 @@ load("//tensorflow:tensorflow.bzl", "tf_cc_test_gpu") load("//tensorflow:tensorflow.bzl", "tf_cc_tests_gpu") load("//tensorflow:tensorflow.bzl", "tf_cuda_cc_test") load("//tensorflow:tensorflow.bzl", "tf_version_info_genrule") +load("//tensorflow:tensorflow.bzl", "if_nccl") +load("//tensorflow:tensorflow.bzl", "tensorflow_opensource_extra_deps") load("//tensorflow:tensorflow.bzl", "tf_cuda_only_cc_test") # For platform specific build config @@ -130,7 +135,6 @@ load( "tf_kernel_tests_linkstatic", "tf_lib_proto_compiler_deps", "tf_lib_proto_parsing_deps", - "tf_nano_proto_library", "tf_platform_hdrs", "tf_platform_srcs", "tf_proto_library", @@ -177,7 +181,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", @@ -251,15 +254,6 @@ tf_jspb_proto_library( deps = [":protos_all_cc"], ) -tf_nano_proto_library( - name = "protos_all_nano_proto", - field_style = "accessors", - generate_equals = 1, - generate_intdefs = 1, - visibility = ["//visibility:public"], - deps = [":protos_all_cc"], -) - proto_library( name = "example_protos", srcs = [ @@ -309,6 +303,7 @@ filegroup( "platform/env_time.h", "platform/logging.h", "platform/macros.h", + "platform/platform_strings.h", "platform/types.h", ], visibility = ["//visibility:private"], @@ -335,6 +330,16 @@ cc_library( ], ) +cc_library( + name = "framework_bounds_check", + hdrs = ["framework/bounds_check.h"], + visibility = ["//tensorflow/core/kernels:friends"], + deps = [ + "//tensorflow/core:platform_base", + "//third_party/eigen3", + ], +) + filegroup( name = "platform_port_hdrs", srcs = [ @@ -441,6 +446,18 @@ cc_library( ] + tf_additional_human_readable_json_deps(), ) +cc_library( + name = "logger", + srcs = ["platform/logger.cc"], + hdrs = ["platform/logger.h"], + copts = tf_copts(), + visibility = ["//visibility:public"], + deps = [ + ":lib_proto_parsing", + "@protobuf_archive//:protobuf", + ], +) + filegroup( name = "platform_env_hdrs", srcs = [ @@ -476,7 +493,10 @@ cc_library( ":platform_env_internal_hdrs", ], copts = tf_copts(), - visibility = ["//tensorflow/core:__subpackages__"], + visibility = [ + "//tensorflow/c:__subpackages__", + "//tensorflow/core:__subpackages__", + ], deps = [ ":error_codes_proto_cc", ":lib", @@ -518,6 +538,19 @@ cc_library( ], ) +cc_library( + name = "platform_strings", + srcs = tf_platform_srcs([ + "platform/platform_strings.cc", + "platform/platform_strings_computed.h", + ]), + hdrs = [ + "platform/platform_strings.h", + ], + visibility = ["//tensorflow/core:__subpackages__"], + deps = [":lib"], +) + filegroup( name = "platform_other_hdrs", srcs = [ @@ -646,7 +679,6 @@ cc_library( "lib/core/arena.h", "lib/core/bitmap.h", "lib/core/bits.h", - "lib/core/casts.h", "lib/core/coding.h", "lib/core/errors.h", "lib/core/notification.h", @@ -785,10 +817,13 @@ cc_library( }), visibility = ["//visibility:public"], deps = [ + ":function_ops_op_lib", + ":functional_ops_op_lib", ":lib", ":lib_internal", ":protos_all_cc", "//tensorflow/core/platform/default/build_config:gtest", + "//tensorflow/core/kernels:required", ] + tf_additional_test_deps(), ) @@ -822,6 +857,7 @@ tf_cuda_library( hdrs = [ "example/feature_util.h", "framework/allocator.h", + "framework/bounds_check.h", "framework/variant.h", "framework/variant_encode_decode.h", "framework/variant_op_registry.h", @@ -837,6 +873,7 @@ tf_cuda_library( "framework/dataset_stateful_op_whitelist.h", "framework/device_base.h", "framework/function.h", + "framework/function_handle_cache.h", "framework/graph_def_util.h", "framework/graph_to_functiondef.h", "framework/kernel_def_builder.h", @@ -852,6 +889,7 @@ tf_cuda_library( "framework/op_def_builder.h", "framework/op_def_util.h", "framework/op_kernel.h", + "framework/ops_util.h", "framework/partial_tensor_shape.h", "framework/queue_interface.h", "framework/reader_interface.h", @@ -879,6 +917,7 @@ tf_cuda_library( "util/bcast.h", "util/cuda_kernel_helper.h", "util/device_name_utils.h", + "util/dump_graph.h", "util/events_writer.h", "util/example_proto_fast_parsing.h", "util/example_proto_helper.h", @@ -896,6 +935,7 @@ tf_cuda_library( "util/stream_executor_util.h", "util/strided_slice_op.h", "util/tensor_format.h", + "util/tensor_ops_util.h", "util/tensor_slice_reader.h", "util/tensor_slice_reader_cache.h", "util/tensor_slice_writer.h", @@ -979,6 +1019,7 @@ cc_library( ":lib", ":lib_internal", ":protos_all_cc", + "//tensorflow/core/util/proto:proto_utils", ], ) @@ -1033,8 +1074,10 @@ tf_gen_op_libs( "batch_ops", "bitwise_ops", "boosted_trees_ops", + "tensor_forest_ops", "candidate_sampling_ops", "checkpoint_ops", + "clustering_ops", "collective_ops", "control_flow_ops", "ctc_ops", @@ -1053,11 +1096,14 @@ tf_gen_op_libs( "logging_ops", "manip_ops", "math_ops", + "mkl_nn_ops", + "nccl_ops", "nn_ops", "no_op", "parsing_ops", "random_grad", "random_ops", + "stateful_random_ops", "remote_fused_graph_ops", "rpc_ops", "scoped_allocator_ops", @@ -1078,7 +1124,11 @@ tf_gen_op_libs( op_lib_names = [ "string_ops", ], - deps = ["@com_google_absl//absl/strings"], + deps = [ + ":lib_internal", + ":lib_proto_parsing", + "@com_google_absl//absl/strings", + ], ) tf_gen_op_libs( @@ -1158,6 +1208,7 @@ cc_library( name = "ragged_ops", deps = [ ":ragged_array_ops_op_lib", + ":ragged_conversion_ops_op_lib", ":ragged_math_ops_op_lib", ], ) @@ -1165,6 +1216,7 @@ cc_library( tf_gen_op_libs( op_lib_names = [ "ragged_array_ops", + "ragged_conversion_ops", "ragged_math_ops", ], ) @@ -1178,8 +1230,10 @@ cc_library( ":batch_ops_op_lib", ":bitwise_ops_op_lib", ":boosted_trees_ops_op_lib", + ":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", @@ -1199,11 +1253,13 @@ cc_library( ":lookup_ops_op_lib", ":manip_ops_op_lib", ":math_ops_op_lib", + ":nccl_ops_op_lib", ":nn_ops_op_lib", ":no_op_op_lib", ":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", @@ -1221,7 +1277,7 @@ cc_library( ":training_ops_op_lib", ":user_ops_op_lib", ":word2vec_ops", - ] + tf_additional_cloud_op_deps(), + ] + if_mkl([":mkl_nn_ops_op_lib"]) + tf_additional_cloud_op_deps(), alwayslink = 1, ) @@ -1277,7 +1333,9 @@ cc_library( ":framework", ":lib", ":nn_ops_op_lib", - ], + ] + if_mkl([ + ":mkl_nn_ops_op_lib", + ]), alwayslink = 1, ) @@ -1320,7 +1378,7 @@ cc_library( # This includes implementations of all kernels built into TensorFlow. cc_library( - name = "all_kernels_statically_linked", + name = "all_kernels_impl", visibility = ["//visibility:private"], deps = [ "//tensorflow/core/kernels:array", @@ -1328,14 +1386,15 @@ cc_library( "//tensorflow/core/kernels:batch_kernels", "//tensorflow/core/kernels:bincount_op", "//tensorflow/core/kernels:boosted_trees_ops", + "//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", "//tensorflow/core/kernels:cudnn_rnn_kernels", "//tensorflow/core/kernels:data_flow", - "//tensorflow/core/kernels:dataset_ops", "//tensorflow/core/kernels:decode_proto_op", "//tensorflow/core/kernels:encode_proto_op", "//tensorflow/core/kernels:fake_quant_ops", @@ -1346,18 +1405,20 @@ cc_library( "//tensorflow/core/kernels:image", "//tensorflow/core/kernels:io", "//tensorflow/core/kernels:linalg", - "//tensorflow/core/kernels:list_kernels", "//tensorflow/core/kernels:lookup", "//tensorflow/core/kernels:logging", "//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", @@ -1399,6 +1460,8 @@ cc_library( ]) + if_cuda([ "//tensorflow/core/grappler/optimizers:gpu_swapping_kernels", "//tensorflow/core/grappler/optimizers:gpu_swapping_ops", + ]) + if_nccl([ + "//tensorflow/core/kernels:nccl_kernels", ]), ) @@ -1407,8 +1470,13 @@ cc_library( visibility = ["//visibility:public"], deps = if_dynamic_kernels( [], - otherwise = [":all_kernels_statically_linked"], - ), + otherwise = [":all_kernels_impl"], + ) + [ + # TODO(gunan): Work on the API between these and rest of TF and make + # these also dynamically loading. + "//tensorflow/core/kernels:dataset_ops", # Depends on grappler + "//tensorflow/core/kernels:list_kernels", # Depends on variant_op_registry.h + ], ) tf_cuda_library( @@ -1422,7 +1490,8 @@ tf_cuda_library( ":example_parser_configuration", ":gpu_runtime", ":lib", - ], + ":ops", + ] + tensorflow_opensource_extra_deps(), ) cc_library( @@ -1474,12 +1543,16 @@ cc_library( ":test", ":testlib_ops", "//tensorflow/cc:scope", - "//tensorflow/core/kernels:cast_op", - "//tensorflow/core/kernels:constant_op", "//tensorflow/core/kernels:ops_testutil", "//tensorflow/core/kernels:ops_util", - "//tensorflow/core/kernels:random_ops", - ], + ] + if_dynamic_kernels( + [], + otherwise = [ + "//tensorflow/core/kernels:cast_op", + "//tensorflow/core/kernels:constant_op", + "//tensorflow/core/kernels:random_ops", + ], + ), ) cc_library( @@ -1551,6 +1624,9 @@ filegroup( "**/*main.cc", "debug/**/*", "framework/op_gen_*", + "framework/node_def_util.*", + "framework/op_kernel.*", + "framework/dataset.*", "lib/jpeg/**/*", "lib/png/**/*", "lib/gif/**/*", @@ -1558,6 +1634,7 @@ filegroup( "util/stats_calculator.*", "util/reporter.*", "platform/**/cuda_libdevice_path.*", + "platform/**/logger.cc", "platform/default/test_benchmark.*", "platform/cuda.h", "platform/google/**/*", @@ -1592,6 +1669,9 @@ filegroup( "common_runtime/**/*.cc", "graph/**/*.h", "graph/**/*.cc", + "framework/node_def_util.*", + "framework/op_kernel.*", + "framework/dataset.*", ], exclude = [ "**/*test.*", @@ -1620,6 +1700,9 @@ filegroup( # operators, use :android_tensorflow_lib if you want full operator # support. # +# If you just need TensorFlow types, e.g. Tensors, use +# :android_tensorflow_lib_lite_no_runtime. +# # Compiles to a trivial library on non-Android to prevent irrelevant # build errors. If not building this as part of an android_binary, # a command such as the following must be used: @@ -1630,7 +1713,33 @@ filegroup( cc_library( name = "android_tensorflow_lib_lite", srcs = if_android(["//tensorflow/core:android_srcs"]), - copts = tf_copts(android_optimization_level_override = None), + copts = tf_copts(android_optimization_level_override = None) + [ + "-DSUPPORT_SELECTIVE_REGISTRATION", + ], + linkopts = ["-lz"], + tags = [ + "manual", + "notap", + ], + visibility = ["//visibility:public"], + deps = [ + ":mobile_additional_lib_deps", + ":protos_all_cc_impl", + ":stats_calculator_portable", + "//third_party/eigen3", + "@double_conversion//:double-conversion", + "@nsync//:nsync_cpp", + "@protobuf_archive//:protobuf", + ], + alwayslink = 1, +) + +cc_library( + name = "android_tensorflow_lib_lite_nortti", + srcs = if_android(["//tensorflow/core:android_srcs"]), + copts = tf_copts(android_optimization_level_override = None) + [ + "-DSUPPORT_SELECTIVE_REGISTRATION", + ] + tf_opts_nortti_if_android(), linkopts = ["-lz"], tags = [ "manual", @@ -1652,6 +1761,8 @@ 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", ], ) @@ -1737,56 +1848,6 @@ cc_library( alwayslink = 1, ) -# Android library for use with the SELECTIVE_REGISTRATION feature. -# Does not contain operators. In contrast to android_tensorflow_lib_lite, -# this links in framework support for all types, relying on selective -# registration of ops to prune code size. -cc_library( - name = "android_tensorflow_lib_selective_registration", - srcs = if_android(["//tensorflow/core:android_srcs"]), - copts = tf_copts(android_optimization_level_override = None) + [ - "-DSUPPORT_SELECTIVE_REGISTRATION", - ], - linkopts = if_android(["-lz"]), - tags = [ - "manual", - "notap", - ], - visibility = ["//visibility:public"], - deps = [ - ":protos_all_cc_impl", - "//third_party/eigen3", - "@double_conversion//:double-conversion", - "@nsync//:nsync_cpp", - "@protobuf_archive//:protobuf", - ], - alwayslink = 1, -) - -# Android library for use with the SELECTIVE_REGISTRATION feature with -# no proto_rtti. -cc_library( - name = "android_tensorflow_lib_selective_registration_nortti", - srcs = if_android(["//tensorflow/core:android_srcs"]), - copts = tf_copts(android_optimization_level_override = None) + tf_opts_nortti_if_android() + [ - "-DSUPPORT_SELECTIVE_REGISTRATION", - ], - linkopts = if_android(["-lz"]), - tags = [ - "manual", - "notap", - ], - visibility = ["//visibility:public"], - deps = [ - ":protos_all_cc_impl", - "//third_party/eigen3", - "@double_conversion//:double-conversion", - "@nsync//:nsync_cpp", - "@protobuf_archive//:protobuf", - ], - alwayslink = 1, -) - filegroup( name = "android_op_registrations_and_gradients", srcs = glob( @@ -1960,6 +2021,13 @@ tf_pyclif_proto_library( visibility = ["//visibility:public"], ) +tf_pyclif_proto_library( + name = "framework/step_stats_pyclif", + proto_lib = ":protos_all_cc", + proto_srcfile = "framework/step_stats.proto", + visibility = ["//visibility:public"], +) + tf_pyclif_proto_library( name = "framework/types_pyclif", proto_lib = ":protos_all_cc", @@ -2025,9 +2093,7 @@ tf_proto_library_cc( srcs = ["protobuf/master.proto"], cc_api_version = 2, protodeps = tf_additional_all_protos(), - visibility = [ - "//tensorflow:internal", - ], + visibility = ["//tensorflow:internal"], ) tf_proto_library_cc( @@ -2139,6 +2205,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", @@ -2153,6 +2220,7 @@ cc_library( "lib/**/*.cc", "platform/*.cc", "platform/profile_utils/**/*.cc", + ] + [ "framework/resource_handle.cc", "util/env_var.cc", ], @@ -2166,6 +2234,7 @@ cc_library( "platform/**/env_time.cc", "platform/**/cuda_libdevice_path.cc", "platform/**/device_tracer.cc", + "platform/**/logger.cc", "platform/**/logging.cc", "platform/**/human_readable_json.cc", "platform/abi.cc", @@ -2178,6 +2247,7 @@ cc_library( "platform/**/stream_executor.h", "platform/**/env_time.cc", "platform/**/device_tracer.cc", + "platform/**/logger.cc", "platform/**/logging.cc", "platform/**/human_readable_json.cc", "platform/abi.cc", @@ -2260,7 +2330,6 @@ cc_library( srcs = ["lib/png/png_io.cc"], hdrs = [ "lib/bfloat16/bfloat16.h", - "lib/core/casts.h", "lib/core/stringpiece.h", "lib/png/png_io.h", "platform/byte_order.h", @@ -2283,6 +2352,7 @@ cc_library( ":lib", ":lib_internal", "//tensorflow/core/platform/default/build_config:png", + "@com_google_absl//absl/base", "@com_google_absl//absl/strings", "@zlib_archive//:zlib", ], @@ -2290,7 +2360,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", @@ -2299,7 +2374,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 = [ @@ -2380,7 +2455,6 @@ cc_library( ]), hdrs = [ "lib/bfloat16/bfloat16.h", - "lib/core/casts.h", "lib/core/stringpiece.h", "lib/png/png_io.h", "platform/byte_order.h", @@ -2500,6 +2574,7 @@ FRAMEWORK_INTERNAL_PRIVATE_HEADERS = [ }) FRAMEWORK_INTERNAL_PUBLIC_HEADERS = [ + "framework/model.h", # only needed for tests "framework/op_segment.h", "framework/rendezvous.h", # only needed for tests "framework/resource_var.h", @@ -2619,6 +2694,9 @@ tf_cuda_library( ":protos_all_cc", ":stats_calculator_portable", ":version_lib", + "@com_google_absl//absl/base", + "@com_google_absl//absl/container:flat_hash_set", + "@com_google_absl//absl/strings", "//tensorflow/core/platform/default/build_config:platformlib", "//tensorflow/core/kernels:bounds_check", "//third_party/eigen3", @@ -2706,6 +2784,7 @@ cc_library( # in this library. GRAPH_HDRS = [ "graph/algorithm.h", + "graph/collective_order.h", "graph/colors.h", "graph/control_flow.h", "graph/costmodel.h", @@ -2732,6 +2811,7 @@ tf_cuda_library( name = "graph", srcs = [ "graph/algorithm.cc", + "graph/collective_order.cc", "graph/colors.cc", "graph/control_flow.cc", "graph/costmodel.cc", @@ -2749,6 +2829,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", ], ) @@ -2763,12 +2846,16 @@ CORE_CPU_BASE_HDRS = GRAPH_HDRS + [ "framework/versions.h", "common_runtime/process_function_library_runtime.h", "common_runtime/function.h", + "common_runtime/scoped_allocator.h", + "common_runtime/scoped_allocator_mgr.h", ] tf_cuda_library( name = "core_cpu_base", srcs = [ "common_runtime/eval_const_tensor.cc", + "common_runtime/scoped_allocator.cc", + "common_runtime/scoped_allocator_mgr.cc", "common_runtime/shape_refiner.cc", "common_runtime/shape_refiner.h", "framework/versions.h", @@ -2795,7 +2882,6 @@ tf_cuda_library( ":functional_ops_op_lib", "//tensorflow/core/kernels:bounds_check", "//tensorflow/core/kernels:required", - ":core_cpu_impl", ]), alwayslink = 1, ) @@ -2829,6 +2915,7 @@ CORE_CPU_LIB_HEADERS = CORE_CPU_BASE_HDRS + [ "common_runtime/mkl_cpu_allocator.h", "common_runtime/optimization_registry.h", "common_runtime/pending_counts.h", + "common_runtime/partitioning_utils.h", "common_runtime/placer.h", "common_runtime/process_util.h", "common_runtime/profile_handler.h", @@ -2836,8 +2923,6 @@ CORE_CPU_LIB_HEADERS = CORE_CPU_BASE_HDRS + [ "common_runtime/rendezvous_mgr.h", "common_runtime/rendezvous_util.h", "common_runtime/ring_reducer.h", - "common_runtime/scoped_allocator.h", - "common_runtime/scoped_allocator_mgr.h", "common_runtime/session_factory.h", "common_runtime/single_threaded_cpu_device.h", "common_runtime/stats_publisher_interface.h", @@ -2885,6 +2970,7 @@ tf_cuda_library( "common_runtime/mkl_cpu_allocator.cc", "common_runtime/optimization_registry.cc", "common_runtime/parallel_concat_optimizer.cc", + "common_runtime/partitioning_utils.cc", "common_runtime/placer.cc", "common_runtime/pool_allocator.cc", "common_runtime/process_function_library_runtime.cc", @@ -2894,8 +2980,6 @@ tf_cuda_library( "common_runtime/rendezvous_mgr.cc", "common_runtime/rendezvous_util.cc", "common_runtime/ring_reducer.cc", - "common_runtime/scoped_allocator.cc", - "common_runtime/scoped_allocator_mgr.cc", "common_runtime/session.cc", "common_runtime/session_factory.cc", "common_runtime/session_options.cc", @@ -2922,8 +3006,10 @@ tf_cuda_library( ":lib_internal", ":proto_text", ":protos_all_cc", + "@com_google_absl//absl/memory", + "@com_google_absl//absl/strings", "//third_party/eigen3", - "//tensorflow/core/grappler:grappler_item", + "//tensorflow/core/grappler/utils:functions", ] + mkl_deps(), alwayslink = 1, ) @@ -2949,13 +3035,9 @@ tf_cuda_library( copts = tf_copts(), deps = [ ":framework", - ":framework_internal", - ":function_ops_op_lib", - ":functional_grad", - ":functional_ops_op_lib", ":graph", ":lib", - ":lib_internal", + ":metrics", ":proto_text", ":protos_all_cc", "//tensorflow/core/grappler:grappler_item", @@ -2963,8 +3045,13 @@ tf_cuda_library( "//tensorflow/core/grappler/clusters:virtual_cluster", "//tensorflow/core/grappler/optimizers:meta_optimizer", "//third_party/eigen3", + ] + mkl_deps() + tf_additional_core_deps() + if_static([ + ":core_cpu_impl", + ":function_ops_op_lib", + ":functional_grad", + ":functional_ops_op_lib", "//tensorflow/core/kernels:required", - ] + mkl_deps() + tf_additional_core_deps() + if_static([":core_cpu_impl"]), + ]), alwayslink = 1, ) @@ -2981,6 +3068,15 @@ cc_library( deps = [":lib_internal"], ) +tf_cuda_library( + name = "metrics", + srcs = ["common_runtime/metrics.cc"], + hdrs = ["common_runtime/metrics.h"], + deps = [ + ":lib", + ], +) + tf_cuda_library( name = "direct_session_internal", srcs = ["common_runtime/direct_session.cc"], @@ -2997,6 +3093,7 @@ tf_cuda_library( ":graph", ":lib", ":lib_internal", + ":metrics", ":proto_text", ":protos_all_cc", "//tensorflow/core/debug:debug_graph_utils", @@ -3032,7 +3129,9 @@ tf_cuda_library( ], copts = tf_copts(), cuda_deps = if_cuda_is_configured(tf_additional_cupti_wrapper_deps() + tf_additional_device_tracer_cuda_deps()), - visibility = ["//visibility:private"], + visibility = [ + "//tensorflow:internal", + ], deps = [ ":core_cpu_internal", ":lib", @@ -3301,7 +3400,6 @@ tf_cc_tests( size = "small", srcs = [ "lib/core/arena_test.cc", - "lib/core/bit_cast_test.cc", "lib/core/bitmap_test.cc", "lib/core/blocking_counter_test.cc", "lib/core/coding_test.cc", @@ -3359,6 +3457,7 @@ tf_cc_tests( "platform/profile_utils/cpu_utils_test.cc", "platform/stacktrace_handler_test.cc", "platform/subprocess_test.cc", + "platform/vmodule_benchmark_test.cc", ], deps = [ ":lib", @@ -3372,6 +3471,20 @@ tf_cc_tests( ], ) +tf_cc_test( + name = "vmodule_test", + srcs = ["platform/vmodule_test.cc"], + tags = ["optonly"], + deps = [ + ":lib", + ":lib_internal", + ":lib_test_internal", + ":protos_all_cc", + ":test", + "//third_party/eigen3", + ], +) + tf_cc_test( name = "lib_random_random_distributions_test", srcs = ["lib/random/random_distributions_test.cc"], @@ -3387,6 +3500,16 @@ tf_cc_test( ], ) +tf_cc_test( + name = "platform_strings_test", + size = "small", + srcs = ["platform/platform_strings_test.cc"], + deps = [ + ":lib", + ":platform_strings", + ], +) + tf_cc_test( name = "platform_env_test", size = "small", @@ -3402,6 +3525,29 @@ tf_cc_test( ], ) +tf_cc_test( + name = "platform_fake_python_env_test", + size = "small", + srcs = ["platform/fake_python_env_test.cc"], + args = [ + "/some/path/to/pythontest.runfiles/org_tensorflow/stuff/to/run.py", + ], + tags = [ + "local", + "no_windows", + "nogpu", + "nomac", + "notap", + ], + deps = [ + ":lib", + ":lib_internal", + ":lib_test_internal", + ":test", + ":test_main", + ], +) + tf_cc_test( name = "platform_abi_test", size = "small", @@ -3507,6 +3653,7 @@ tf_cc_test( ":lib_internal", ":test", ":test_main", + "@com_google_absl//absl/base", ], ) @@ -3574,7 +3721,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", @@ -3595,6 +3741,7 @@ tf_cc_tests( "framework/kernel_def_builder_test.cc", "framework/kernel_def_util_test.cc", "framework/memory_types_test.cc", + "framework/model_test.cc", "framework/node_def_builder_test.cc", "framework/node_def_util_test.cc", "framework/op_compatibility_test.cc", @@ -3632,6 +3779,7 @@ tf_cc_tests( "util/bcast_test.cc", "util/command_line_flags_test.cc", "util/device_name_utils_test.cc", + "util/dump_graph_test.cc", "util/equal_graph_def_test.cc", "util/events_writer_test.cc", "util/example_proto_fast_parsing_test.cc", @@ -3680,6 +3828,7 @@ tf_cc_tests( "//tensorflow/cc:while_loop", "//tensorflow/core/kernels:ops_util", "//third_party/eigen3", + "@com_google_absl//absl/base", ], ) @@ -3687,6 +3836,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({ @@ -3736,6 +3886,27 @@ tf_cc_test( ], ) +tf_cc_tests( + name = "collective_order_test", + size = "small", + srcs = [ + "graph/collective_order_test.cc", + ], + deps = [ + ":core", + ":core_cpu", + ":core_cpu_internal", + ":framework", + ":framework_internal", + ":lib", + ":lib_internal", + ":ops", + ":protos_all_cc", + ":test", + "@com_google_googletest//:gtest_main", + ], +) + tf_cc_tests_gpu( name = "ring_reducer_test", size = "medium", @@ -3761,6 +3932,7 @@ tf_cc_tests_gpu( ":test", ":test_main", ":testlib", + "@com_google_absl//absl/memory", ], ) @@ -3789,6 +3961,7 @@ tf_cc_tests_gpu( ":test", ":test_main", ":testlib", + "@com_google_absl//absl/memory", ], ) @@ -3950,20 +4123,6 @@ tf_cuda_cc_test( ], ) -tf_cc_test_gpu( - name = "cuda_libdevice_path_test", - size = "small", - srcs = ["platform/cuda_libdevice_path_test.cc"], - linkstatic = tf_kernel_tests_linkstatic(), - tags = tf_cuda_tests_tags(), - deps = [ - ":cuda_libdevice_path", - ":lib", - ":test", - ":test_main", - ], -) - tf_cuda_only_cc_test( name = "util_cuda_kernel_helper_test", srcs = [ @@ -4062,6 +4221,7 @@ tf_cc_test( "//tensorflow/core/kernels:identity_op", "//tensorflow/core/kernels:immutable_constant_op", "//tensorflow/core/kernels:matmul_op", + "//tensorflow/core/kernels:topk_op", "//third_party/eigen3", ], ) @@ -4096,7 +4256,7 @@ tf_cc_test( ], ) -tf_cc_test( +tf_cuda_cc_test( name = "common_runtime_process_function_library_runtime_test", size = "small", srcs = ["common_runtime/process_function_library_runtime_test.cc"], @@ -4105,6 +4265,7 @@ tf_cc_test( ":core_cpu", ":core_cpu_internal", ":framework", + ":framework_internal", ":lib", ":test", ":test_main", @@ -4113,6 +4274,7 @@ tf_cc_test( "//tensorflow/core/kernels:cast_op", "//tensorflow/core/kernels:cwise_op", "//tensorflow/core/kernels:function_ops", + "//tensorflow/core/kernels:resource_variable_ops", ], ) @@ -4154,6 +4316,27 @@ tf_cc_test( ], ) +tf_cc_test( + name = "common_runtime_partitioning_utils_test", + size = "small", + srcs = ["common_runtime/partitioning_utils_test.cc"], + deps = [ + ":core_cpu", + ":core_cpu_internal", + ":framework", + ":lib", + ":ops", + ":test", + ":test_main", + ":testlib", + "//tensorflow/cc:cc_ops", + "//tensorflow/cc:cc_ops_internal", + "//tensorflow/cc:function_ops", + "//tensorflow/core/kernels:function_ops", + "//tensorflow/core/kernels:identity_op", + ], +) + tf_cuda_cc_test( name = "common_runtime_direct_session_test", size = "small", @@ -4346,13 +4529,17 @@ tf_cc_test( "//tensorflow/cc:cc_ops_internal", "//tensorflow/cc:function_ops", "//tensorflow/cc:functional_ops", + "//tensorflow/cc:sendrecv_ops", "//tensorflow/core/kernels:cast_op", "//tensorflow/core/kernels:cwise_op", "//tensorflow/core/kernels:function_ops", "//tensorflow/core/kernels:matmul_op", + "//tensorflow/core/kernels:partitioned_function_ops", "//tensorflow/core/kernels:random_ops", "//tensorflow/core/kernels:shape_ops", "//third_party/eigen3", + "@com_google_absl//absl/memory", + "@com_google_absl//absl/strings", ], ) @@ -4814,7 +5001,7 @@ filegroup( cc_library( name = "cuda_libdevice_path", - srcs = ["platform/cuda_libdevice_path.cc"] + tf_additional_libdevice_srcs(), + srcs = tf_additional_libdevice_srcs(), hdrs = ["platform/cuda_libdevice_path.h"], copts = tf_copts(), data = tf_additional_libdevice_data(), @@ -4831,6 +5018,7 @@ transitive_hdrs( "//tensorflow/core:core_cpu", "//tensorflow/core:framework", "//tensorflow/core:lib", + "//tensorflow/core:platform_strings", "//tensorflow/core:protos_all_cc", "//tensorflow/core:stream_executor", ], diff --git a/tensorflow/core/api_def/BUILD b/tensorflow/core/api_def/BUILD index 06b797e32edc046bab498f8d775040d57ef62ce9..f610facd75d8cfa42845714f87498bd7afff58e2 100644 --- a/tensorflow/core/api_def/BUILD +++ b/tensorflow/core/api_def/BUILD @@ -17,6 +17,10 @@ load( "tf_cc_binary", "tf_cc_test", ) +load( + "//third_party/mkl:build_defs.bzl", + "if_mkl", +) filegroup( name = "base_api_def", @@ -40,6 +44,7 @@ cc_library( name = "excluded_ops_lib", srcs = ["excluded_ops.cc"], hdrs = ["excluded_ops.h"], + copts = if_mkl(["-DINTEL_MKL=1"]), ) cc_library( diff --git a/tensorflow/core/api_def/api_test.cc b/tensorflow/core/api_def/api_test.cc index 51812caeb2979270c913adee4fba2ce02f9c4d0e..7405e2ace72d1c08cf87cc0040e617379e18149b 100644 --- a/tensorflow/core/api_def/api_test.cc +++ b/tensorflow/core/api_def/api_test.cc @@ -35,7 +35,6 @@ limitations under the License. #include "tensorflow/core/platform/protobuf.h" #include "tensorflow/core/platform/test.h" #include "tensorflow/core/platform/types.h" -#include "tensorflow/core/util/command_line_flags.h" namespace tensorflow { namespace { @@ -176,6 +175,22 @@ void TestDeprecatedAttributesSetCorrectly( } } } + +void TestDeprecationVersionSetCorrectly( + const std::unordered_map& api_defs_map) { + for (const auto& name_and_api_def : api_defs_map) { + const auto& name = name_and_api_def.first; + const auto& api_def = name_and_api_def.second; + if (api_def.deprecation_version() != 0) { + ASSERT_TRUE(api_def.deprecation_version() > 0) + << "Found ApiDef with negative deprecation_version"; + ASSERT_FALSE(api_def.deprecation_message().empty()) + << "ApiDef that includes deprecation_version > 0 must also specify " + << "a deprecation_message. Op " << name + << " has deprecation_version > 0 but deprecation_message is not set."; + } + } +} } // namespace class BaseApiTest : public ::testing::Test { @@ -268,6 +283,12 @@ TEST_F(BaseApiTest, DeprecationSetCorrectly) { TestDeprecatedAttributesSetCorrectly(api_defs_map_); } +// Checks that deprecation_version is set for entire op only if +// deprecation_message is set. +TEST_F(BaseApiTest, DeprecationVersionSetCorrectly) { + TestDeprecationVersionSetCorrectly(api_defs_map_); +} + class PythonApiTest : public ::testing::Test { protected: PythonApiTest() { @@ -309,4 +330,10 @@ TEST_F(PythonApiTest, DeprecationSetCorrectly) { TestDeprecatedAttributesSetCorrectly(api_defs_map_); } +// Checks that deprecation_version is set for entire op only if +// deprecation_message is set. +TEST_F(PythonApiTest, DeprecationVersionSetCorrectly) { + TestDeprecationVersionSetCorrectly(api_defs_map_); +} + } // namespace tensorflow diff --git a/tensorflow/core/api_def/base_api/api_def_BatchDataset.pbtxt b/tensorflow/core/api_def/base_api/api_def_BatchDataset.pbtxt index 639d962874d083472e6df13550e107026fd2d0a1..32def912f83e420eab58a3071f573ae81139a298 100644 --- a/tensorflow/core/api_def/base_api/api_def_BatchDataset.pbtxt +++ b/tensorflow/core/api_def/base_api/api_def_BatchDataset.pbtxt @@ -1,5 +1,6 @@ op { graph_op_name: "BatchDataset" + visibility: HIDDEN in_arg { name: "batch_size" description: <

+ +In Python, this scatter operation would look like this: + +```python + indices = tf.constant([[4], [3], [1], [7]]) + updates = tf.constant([9, 10, 11, 12]) + tensor = tf.ones([8], dtype=tf.int32) + updated = tf.tensor_scatter_update(tensor, indices, updates) + with tf.Session() as sess: + print(sess.run(scatter)) +``` + +The resulting tensor would look like this: + + [1, 11, 1, 10, 9, 1, 1, 12] + +We can also, insert entire slices of a higher rank tensor all at once. For +example, if we wanted to insert two slices in the first dimension of a +rank-3 tensor with two matrices of new values. + +In Python, this scatter operation would look like this: + +```python + indices = tf.constant([[0], [2]]) + updates = tf.constant([[[5, 5, 5, 5], [6, 6, 6, 6], + [7, 7, 7, 7], [8, 8, 8, 8]], + [[5, 5, 5, 5], [6, 6, 6, 6], + [7, 7, 7, 7], [8, 8, 8, 8]]]) + tensor = tf.ones([4, 4, 4]) + updated = tf.tensor_scatter_update(tensor, indices, updates) + with tf.Session() as sess: + print(sess.run(scatter)) +``` + +The resulting tensor would look like this: + + [[[5, 5, 5, 5], [6, 6, 6, 6], [7, 7, 7, 7], [8, 8, 8, 8]], + [[1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1]], + [[5, 5, 5, 5], [6, 6, 6, 6], [7, 7, 7, 7], [8, 8, 8, 8]], + [[1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1]]] + +Note that on CPU, if an out of bound index is found, an error is returned. +On GPU, if an out of bound index is found, the index is ignored. +END +} diff --git a/tensorflow/core/api_def/base_api/api_def_TensorSliceDataset.pbtxt b/tensorflow/core/api_def/base_api/api_def_TensorSliceDataset.pbtxt index a26a98fd7f3a6564309efd28dff8c2bc93d7a67f..30cb803b26bf836a7b02cc3fb6875175046eab94 100644 --- a/tensorflow/core/api_def/base_api/api_def_TensorSliceDataset.pbtxt +++ b/tensorflow/core/api_def/base_api/api_def_TensorSliceDataset.pbtxt @@ -1,4 +1,5 @@ op { graph_op_name: "TensorSliceDataset" + visibility: HIDDEN summary: "Creates a dataset that emits each dim-0 slice of `components` once." } diff --git a/tensorflow/core/api_def/base_api/api_def_TextLineDataset.pbtxt b/tensorflow/core/api_def/base_api/api_def_TextLineDataset.pbtxt index 6b630509964ed56ecaf401b10a46c5e53cd46528..31ef3e3335e2812156fc3d1af2c5c1724fa52310 100644 --- a/tensorflow/core/api_def/base_api/api_def_TextLineDataset.pbtxt +++ b/tensorflow/core/api_def/base_api/api_def_TextLineDataset.pbtxt @@ -1,5 +1,6 @@ op { graph_op_name: "TextLineDataset" + visibility: HIDDEN in_arg { name: "filenames" description: <